diff --git a/.ai/skills/datafusion-ffi/SKILL.md b/.ai/skills/datafusion-ffi/SKILL.md new file mode 100644 index 0000000000000..ba02d22c09b30 --- /dev/null +++ b/.ai/skills/datafusion-ffi/SKILL.md @@ -0,0 +1,364 @@ +--- +name: datafusion-ffi +description: Patterns and review checklist for the `datafusion-ffi` crate. Use whenever the user adds, edits, or reviews code under `datafusion/ffi/` — new `FFI_X` wrappers, `ForeignX` implementations, codec changes, or expanding an existing wrapper to cover more of a trait's surface. Also use when reviewing PRs that touch this crate. +--- + +# DataFusion FFI Skill + +This crate exposes a stable C ABI for DataFusion traits so that independently-compiled libraries (different Rust versions, plugins, `datafusion-python`, etc.) can interoperate at runtime. Stability and correctness here are load-bearing: a missed pattern can cause segfaults, leaks, or silently dropped trait behavior on the consumer side. + +Read the crate's `README.md` first if you have not — it establishes the vocabulary (`FFI_X` / `ForeignX`, `library_marker_id`, `release`, `TaskContextProvider`, stabby vs `#[repr(C)]`). + +## When to use + +Trigger this skill any time the work touches `datafusion/ffi/`: + +- Adding a new `FFI_` + `Foreign` pair +- Adding a method to an existing `FFI_X` struct +- Reviewing a PR that touches this crate +- Changing the codec / proto serialization layer +- Bumping the wrapped DataFusion trait surface (e.g. a new default method appeared upstream) + +## Hard rules + +1. **No `datafusion` dependency.** `datafusion-ffi` must not depend on the umbrella `datafusion` crate. Use the leaf crates (`datafusion-common`, `datafusion-expr`, `datafusion-catalog`, `datafusion-physical-plan`, etc.). `datafusion` is fine in `[dev-dependencies]`. +2. **`#[repr(C)]` on every `FFI_X` struct**, not `#[stabby::stabby]`. Stabby is used for `SString`/`SVec` only. Reasons documented in the README (build time, Arrow types lack `IStable`). +3. **`unsafe extern "C"` on every function-pointer field — including `version`.** The one exception is the `library_marker_id` field, which is plain `extern "C" fn() -> usize`. Plain (safe) `extern "C"` also applies to the standalone function defs in `src/lib.rs` — `pub extern "C" fn version()` and `pub extern "C" fn get_library_marker_id()` — which coerce into the `unsafe extern "C"` `version` field slot at construction. +4. **Match `Send`/`Sync` to the wrapped trait.** Raw `*mut c_void` makes every `FFI_X` `!Send + !Sync` by default — `unsafe impl` whichever bounds the consumer-facing trait requires. Most DataFusion traits (`TableProvider`, `ExecutionPlan`, all UDFs, codecs) need both. `Send`-only: `RecordBatchStream`, `Accumulator`, `GroupsAccumulator`, `PartitionEvaluator` (mutable / stream APIs). The matching `ForeignX` always carries the same bounds — pick consistently. +5. **`#![deny(clippy::clone_on_ref_ptr)]`** is on at the crate root. Use `Arc::clone(&x)`, never `x.clone()` on `Arc`. +6. **Run before pushing:** `cargo fmt --all`, `cargo clippy -p datafusion-ffi --all-targets --all-features -- -D warnings`, `cargo test -p datafusion-ffi`. +7. **`api change` label required.** Any PR that modifies an `FFI_X` struct layout (adds/removes/reorders fields, changes a function-pointer signature, adds a variant to an FFI enum, or changes the `version` extern) must carry the `api change` GitHub label. Layout changes break ABI for already-compiled consumer libraries. The label is the project-wide convention for highlighting breaking public-API changes in release notes — see `docs/source/contributor-guide/api-health.md` §"What to do when making breaking API changes?" (step 1 names the label explicitly). Downstream users (e.g. `datafusion-python`, plugin authors) read the labelled notes to know they must recompile against the new DataFusion major. The `version()` extern in `src/lib.rs` returns the major of workspace `CARGO_PKG_VERSION`; consumers compare it at load time and can refuse mismatched producers. Apply via `gh pr edit --add-label "api change"` (label name contains a space — must be quoted). When reviewing such a PR, block merge until label present. +8. **No FFI struct changes in patch releases.** Patch releases ship from branches matching `^branch-\d+$` (e.g. `branch-53`, `branch-52`). FFI struct layout changes (anything that would earn rule 7's `api change` label) **must not** target a release branch and must not be back-ported. Patch releases are ABI-stable by contract — a consumer compiled against `53.1.0` must keep working against `53.1.1`. Before reviewing/approving an FFI PR, check the PR's base branch: if it matches the regex above, or the PR description / labels indicate patch / back-port, reject the FFI struct change and ask the author to retarget `main`. Bugfixes that do not alter struct layout (e.g. fixing a function-pointer body) are fine to back-port. Quick check: `gh pr view --json baseRefName,labels --jq '.baseRefName'` then match against `^branch-\d+$`. Do **not** glob-match `branch-*` — back-port / cherry-pick working branches (e.g. `branch-53-cherry-pick-1`) also share that prefix but are not release branches; only the strict `branch-` form is the freeze target. + +## The standard wrapper shape + +A new `FFI_X` for trait `X` must follow this template. Use `FFI_CatalogProvider` (`src/catalog_provider.rs`) as the canonical reference — it shows the full shape (codec field, nested FFI types, `FFI_Option`/`FFI_Result` returns, Arc-backed `PrivateData`) without async or capability-flag noise. `FFI_TableProvider` (`src/table_provider.rs`) covers async (`scan`, `FFI_SessionRef`, `FfiFuture`) and the one `Option` capability flag (`supports_filters_pushdown`). + +### 1. The `FFI_X` struct + +```rust +#[repr(C)] +#[derive(Debug)] +pub struct FFI_X { + some_method: unsafe extern "C" fn(this: &Self, ...) -> FFI_Result<...>, + optional_method: Option FFI_Result<...>>, + pub logical_codec: FFI_LogicalExtensionCodec, + + clone: unsafe extern "C" fn(&Self) -> Self, + release: unsafe extern "C" fn(&mut Self), + pub version: unsafe extern "C" fn() -> u64, + + private_data: *mut c_void, + pub library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_X {} +unsafe impl Sync for FFI_X {} +``` + +Field rules: + +- **One `unsafe extern "C" fn` per trait method.** Always populate — `Arc` dispatch picks override-or-default at call time, so the producer side gets the right answer without the consumer needing to know. See § "Method coverage". +- **`Option` is the capability-flag exception**, not a template. Crate uses it exactly once: `FFI_TableProvider::supports_filters_pushdown`. See § "Method coverage". +- **Codec field** (`FFI_LogicalExtensionCodec` / `FFI_PhysicalExtensionCodec`) only if the trait moves `Expr`s / `LogicalPlan`s / `ExecutionPlan`s across the boundary. +- **Method function pointers are private by default.** Mark `pub` only if a downstream library needs to invoke them directly (rare — typically only `version`, `library_marker_id`, embedded codecs are `pub`). +- **`version: super::version` is mandatory.** Consumers gate compatibility on it. +- **`library_marker_id: crate::get_library_marker_id` is mandatory *when the wrapper uses the standard `ForeignX` adapter pattern*.** Two flavors exist: + - **Arc-backed (immutable / shareable traits — `TableProvider`, `ExecutionPlan`, all UDFs, codecs):** consumer-side `From<&FFI_X> for Arc` consults the marker to choose `Arc::clone(inner)` vs `Arc::new(ForeignX(...))`. + - **Box-backed (mutable / move-only traits — `Accumulator`, `GroupsAccumulator`, `PartitionEvaluator`):** consumer-side `From for Box` consults the marker to take the inner `Box` directly vs `Box::new(ForeignX(...))`. Producer-side `From> for FFI_X` *also* uses an `is::()` downcast as an additional re-wrap bypass; the marker check covers the reverse direction. + + The field is dead ABI surface — and must be omitted with a one-line module-doc rationale — only when **neither** flavor applies: no `ForeignX` adapter, no reverse `From for {Arc,Box}`, and the trait is impl'd directly on `FFI_X`. Canonical example: `FFI_RecordBatchStream` (`impl RecordBatchStream for FFI_RecordBatchStream` at `record_batch_stream.rs:149`, no `ForeignRecordBatchStream`, no reverse `From`). + + Before flagging a missing `library_marker_id` as a gap, run all three greps on the wrapper file: `Foreign`, `From<&?FFI_X> for Arc<`, `From for Box<`. Hit on any → marker is required and its absence is a real gap. Zero hits on all three → marker is intentionally not needed, not a gap. +- **`Send`/`Sync` bounds match the wrapped trait** (rule 4). Most wrappers want both; `Send`-only for streams / mutable traits. + +### 2. `PrivateData` shape + +Default — for read-only, shareable traits — use `Arc`: + +```rust +struct XPrivateData { + inner: Arc, + runtime: Option, // include when any async method exists +} +``` + +For traits that require `&mut self` (e.g. `Accumulator`, `GroupsAccumulator`, `PartitionEvaluator`), use `Box`: + +```rust +struct XPrivateData { + inner: Box, + runtime: Option, // include when any async method exists +} +``` + +A `Box`-backed `FFI_X` **cannot implement `Clone`**; document this and skip the `clone` function pointer, or hand-write a release path that distinguishes producer vs consumer side. Canonical example: `FFI_Accumulator` in `src/udaf/accumulator.rs`. See also `FFI_GroupsAccumulator` (`src/udaf/groups_accumulator.rs`) and `FFI_PartitionEvaluator` (`src/udwf/partition_evaluator.rs`). + +### 3. Function-pointer wrappers + +Naming convention: `_fn_wrapper`. + +```rust +unsafe extern "C" fn some_method_fn_wrapper(this: &FFI_X, ...) -> FFI_Result { + // 1. Recover inner via this.inner() + // 2. Translate FFI types → native types + // 3. Call native method + // 4. Translate native Result → FFI_Result via sresult_return! or .into() +} +``` + +### 4. `clone` / `release` / `Drop` + +```rust +unsafe extern "C" fn clone_fn_wrapper(this: &FFI_X) -> FFI_X { /* re-Box new private_data, copy fn ptrs */ } + +unsafe extern "C" fn release_fn_wrapper(this: &mut FFI_X) { + unsafe { + debug_assert!(!this.private_data.is_null()); + drop(Box::from_raw(this.private_data as *mut XPrivateData)); + this.private_data = std::ptr::null_mut(); + } +} + +impl Drop for FFI_X { fn drop(&mut self) { unsafe { (self.release)(self) } } } +impl Clone for FFI_X { fn clone(&self) -> Self { unsafe { (self.clone)(self) } } } +``` + +`release` must null `private_data` so a double-free debug-asserts loudly. + +### 5. Constructor split + +```rust +impl FFI_X { + pub fn new(inner: Arc, runtime: Option, + task_ctx_provider: impl Into, + logical_codec: Option>) -> Self { + // build FFI_LogicalExtensionCodec from defaults, then forward + Self::new_with_ffi_codec(inner, runtime, ffi_codec) + } + + pub fn new_with_ffi_codec(inner: Arc, runtime: Option, + logical_codec: FFI_LogicalExtensionCodec) -> Self { + // Round-trip downcast: if inner is already a ForeignX, return its FFI directly. + if let Some(foreign) = inner.downcast_ref::() { + return foreign.0.clone(); + } + // …allocate XPrivateData and populate fn ptrs… + } +} +``` + +The round-trip downcast is **mandatory** — without it, repeated FFI hops nest `ForeignX(FFI_X(ForeignX(...)))` and you pay the boundary cost every layer. + +### 6. The `Foreign` consumer + +```rust +#[derive(Debug)] +pub struct ForeignX(pub FFI_X); +unsafe impl Send for ForeignX {} +unsafe impl Sync for ForeignX {} + +impl From<&FFI_X> for Arc { + fn from(p: &FFI_X) -> Self { + if (p.library_marker_id)() == crate::get_library_marker_id() { + Arc::clone(unsafe { p.inner() }) + } else { + Arc::new(ForeignX(p.clone())) + } + } +} + +impl X for ForeignX { /* call each fn pointer, translate types back */ } +``` + +The marker-id check is **mandatory** for every `From<&FFI_X> for Arc`. Skipping it breaks the local-bypass optimization and forces the producer's data through serialization. + +### 7. Tests + +The crate has **two distinct test surfaces** and a new wrapper usually needs entries in both. They are not interchangeable; they catch different classes of bug. + +#### a. In-process unit tests (`#[cfg(test)] mod tests` inside `src/.rs`) + +Run on every `cargo test -p datafusion-ffi`. Producer and consumer live in the same compilation unit, so `library_marker_id` returns the same value on both sides. To force the foreign path you must override the marker: + +```rust +let mut ffi_x = FFI_X::new(provider, …); +ffi_x.library_marker_id = crate::mock_foreign_marker_id; // forces the ForeignX branch +let arc: Arc = (&ffi_x).into(); +assert!(arc.downcast_ref::().is_some()); +``` + +Every wrapper must include at minimum: + +- A **local-bypass test** — build `FFI_X` from a concrete native type, convert to `Arc`, `downcast_ref::()` must succeed. +- A **forced-foreign test** — set `library_marker_id = crate::mock_foreign_marker_id`, convert, `downcast_ref::()` must succeed, then exercise every method end-to-end. + +Templates: `test_ffi_table_provider_local_bypass` and `test_round_trip_ffi_table_provider_scan` in `src/table_provider.rs`. + +What unit tests catch: Rust-level correctness (translation logic, lifetime bugs, leaks under valgrind/miri, Send/Sync, error propagation, codec round-trips). What they **cannot** catch: real ABI bugs. Both producer and consumer share `#[repr(C)]` layout because they are the exact same struct definition in memory. + +#### b. Cross-library integration tests (`tests/ffi_*.rs`, gated by the `integration-tests` feature) + +The crate is published as `crate-type = ["cdylib", "rlib"]`. The integration tests in `datafusion/ffi/tests/` use `libloading` to `dlopen` the crate's own `cdylib` and call `datafusion_ffi_get_module` — a `#[unsafe(no_mangle)] extern "C"` entry point defined in `src/tests/mod.rs` and gated by `#[cfg(feature = "integration-tests")]`. The test executable links against the rlib (consumer side); the dlopen'd cdylib is the producer side. Even though both are built from the same source, they are independent compilation outputs going through the actual FFI symbol path. + +Run with: + +```bash +cargo test -p datafusion-ffi --features integration-tests +``` + +To add coverage for a new wrapper: + +1. **Add a constructor** in `src/tests/.rs` (or a new file there). Return a populated `FFI_X` from a known-good native type. +2. **Wire it into `ForeignLibraryModule`** in `src/tests/mod.rs`: add a field of type `extern "C" fn(...) -> FFI_X` and populate it in `datafusion_ffi_get_module`. This struct is the cross-library contract. Adding a field is itself an ABI change for the test module; integration tests will rebuild the cdylib automatically. +3. **Add the test** in `tests/ffi_.rs` under `#[cfg(feature = "integration-tests")] mod tests { … }`. Call `datafusion_ffi::tests::utils::get_module()` to load the cdylib, invoke your constructor through the returned `ForeignLibraryModule`, convert into `Arc`, and exercise every method. + +When adding a method to an existing wrapper, reuse an existing fixture and constructor when a small trait-method override can cover it. This applies both to omitted default methods and methods newly added to the trait. Add a dedicated test type or `ForeignLibraryModule` field only when the existing fixtures cannot cover the method. + +What integration tests catch that unit tests cannot: + +- **Real ABI layout bugs.** Two builds means the consumer's view of `FFI_X` is reconstructed from declaration, not aliased to the producer's memory. Mismatched alignment, padding, niche optimization, or accidentally non-`#[repr(C)]` types surface here. +- **Symbol visibility / `no_mangle`** issues. +- **`library_marker_id` correctness without mocking** — the two libraries genuinely have different statics, so the foreign branch is taken for real. +- **Drop / leak ordering** when the producer side is in a `dlopen`'d image. + +#### Which tests does my change need? + +| Change | Unit | Integration | +| --------------------------------------------------------------------- | ---- | ----------- | +| New `FFI_X` wrapper | Yes | Yes | +| New method on existing `FFI_X` | Yes | Yes if the method takes/returns a non-trivial FFI type. See note below. | +| Bugfix to a wrapper body, no signature change | Yes | Only if reproducing the bug requires cross-library symbol lookup or `dlopen` semantics | +| Layout change (`#[repr(C)]` field add/remove/reorder, fn-ptr sig) | Yes | **Mandatory** — this is exactly the bug class integration tests exist for | +| New `From for FFI_X` or codec change | Yes | Yes if the codec is exercised by the cross-library round-trip | + +**"Non-trivial FFI type" for the table above** — anything other than: + +- Primitives (`u8`/`u64`/`bool`/`usize`, etc.) and `#[repr(u8)]` FFI enums (`FFI_TableType`, `Volatility`, `InsertOp`, `TableProviderFilterPushDown`). +- A `stabby::string::String` (`SString`) returned by value, with no other args or returns. + +Concrete skippable example: `fn name(&self) -> SString` reading a field already validated by another method. Concrete *non*-skippable examples: anything returning `SVec`, `FFI_Option`, `FFI_Result`, `WrappedSchema`, `WrappedArray`, an `FfiFuture`, an `FFI_*` sub-struct, or any `*mut`/`*const` pointer. These exercise alignment / padding / niche-opt across the ABI boundary and need the two-build coverage. When unsure, write the integration test; the cost is one constructor + ~20 lines. + +If you skip the integration test for a layout change, you have effectively shipped untested ABI. + +## Method coverage — the silent-default gap + +**This is the area where the crate currently has real holes.** When the wrapped trait has methods with *default implementations*, those defaults are typically the trait's "no-op / unsupported" answer (`None`, `false`, `Unsupported`, `not_impl_err!()`). If the producer overrides a default but the FFI struct does not carry a function pointer for it, the consumer's `Foreign` falls back to the trait default — **silently losing the override**. The producer thinks it implemented `delete_from`; the consumer behaves as if it never did. + +### Rule + +When adding or auditing an `FFI_X`, **enumerate every method on the wrapped trait, including defaulted ones**, and for each one decide: + +| Category | Action | +| ------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------- | +| Required method (no default) | Mandatory `unsafe extern "C" fn` field. | +| Defaulted, plausible override (statistics, distribution, ordering, simplify, DML, …) | Mandatory `unsafe extern "C" fn` field — same as a required method. The wrapper body calls `inner.method(...)` and `Arc` dispatch picks override-or-default for free. | +| Defaulted, deprecated or vestigial | Document the skip in a `// FFI omitted: …` comment. | +| Defaulted, derived purely from other methods already plumbed | Skip — but call out the derivation in a comment. | + +**Do not use `Option` just because the underlying trait has a default.** `Arc` erases override-vs-default info, so the producer side cannot know whether to populate the slot. Always plumb the fn pointer; the wrapper body invokes the trait method and dynamic dispatch does the right thing. + +`Option` is used exactly once in the crate today: `FFI_TableProvider::supports_filters_pushdown`, gated by the `can_support_pushdown_filters: bool` argument to `FFI_TableProvider::new`. It is an exception, not a template. Reach for it only when (a) the producer's constructor takes an explicit capability flag and (b) skipping the FFI call is meaningfully cheaper than letting the trait default run on the producer side. Otherwise plumb the fn pointer unconditionally. + +### Known gaps to close + +Tracked gaps live on GitHub under the [`ffi`](https://github.com/apache/datafusion/issues?q=is%3Aissue+is%3Aopen+label%3Affi) label — that query is the source of truth and stays current as issues are filed or closed. Treat new PRs in those areas as opportunities to fix the listed methods; treat new wrappers as required to avoid creating more. Each open issue names the specific wrapper, the missing methods, and the severity. + +Quick CLI list: + +```bash +gh issue list --repo apache/datafusion --label ffi --state open --limit 50 +``` + +Common severity classes seen on the label today: + +- **DML / optimizer-relevant defaults silently lost** — e.g. `delete_from`/`update`/`truncate` on table providers, distribution / ordering / pushdown on execution plans, `value_from_stats` on aggregates. These demote producer capability on the foreign side. +- **SQL surface area silently absent** — naming hooks (`display_name`, `schema_name`, `documentation`), null-handling and within-group clauses on UDAFs, etc. +- **Performance regressions, not correctness** — e.g. `memoize` on partition evaluators. +- **Open design questions** — none currently tracked. (Historical entry: whether `FFI_RecordBatchStream` needs `library_marker_id` — resolved no, it impls the trait directly on `FFI_X` with no `Foreign` adapter and no reverse `From`, so the marker has no consultation site. See `library_marker_id` rule in §"Method coverage".) + +When in doubt, open the label query; do not assume the list above is exhaustive. Wrappers without an open issue are **not** certified complete — re-enumerate the trait surface (see "How to audit a wrapper's coverage" below) whenever you audit one; upstream trait drift can introduce new defaulted methods at any time and silently re-open the silent-override-loss bug class. + +Conversely, an open issue under the `ffi` label is a **claim of a gap, not proof of one.** Past audits have filed false positives by enumerating gaps from memory rather than from current source (e.g. #22335 claimed `size` missing on `FFI_GroupsAccumulator` when it had been plumbed since PR #14775). Before acting on a listed gap — opening a fix PR, re-citing it in a new audit, or extending the list — run the dual-grep audit below and confirm the field is actually absent. If the issue is a false positive, close it as `not planned`, link the `file:line` of the existing plumbing, and remove the corresponding bullet from this skill. + +When *closing* a gap, add the fn pointer unconditionally — the wrapper body calls the trait method on the inner `Arc` and Rust's dynamic dispatch picks the producer's override or falls back to the default. Use `Option` only when the new method also gains a corresponding capability flag on the producer's `new()`. Either way the layout changes, so the PR is an ABI break: mark `api change`, do not back-port to `branch-`, and the workspace major bump in the next release makes the `version()` extern surface the change to consumers at load time. + +### How to audit a wrapper's coverage + +Every audit — opening an issue, filing a fix PR, or re-confirming a listed gap — must compare two sides drawn from current source. Never enumerate either side from memory or from a prior audit; trait surface and FFI struct both drift. + +**Side A — trait defaults (what could go missing):** + +```bash +# Find the trait definition +grep -rn "pub trait X" datafusion/ --include='*.rs' + +# Inspect for `fn method(...) { default_body }` — the body marks it as a default +``` + +**Side B — FFI wrapper coverage (what is already plumbed):** + +```bash +# List every fn-pointer field on the FFI struct +grep -nE 'pub [a-z_]+: (unsafe )?extern "C" fn' datafusion/ffi/src/.../X.rs +``` + +Diff Side A against Side B. Any claim of a gap — in an issue body, audit summary, or PR description — must cite `file:line` for **both** the trait default and the FFI struct line where the field is (or is not). An issue body with only one side cited is incomplete and likely a false positive; reject it pending a re-grep. + +If a method's body is non-trivial, the consumer-side default is non-trivial too. Decide explicitly whether the FFI should let the consumer recompute the same default, or whether the producer's override is what should travel. + +## Type-bridging conventions + +- **Strings / vecs** crossing the boundary: `stabby::string::String as SString`, `stabby::vec::Vec as SVec`. Native conversion is `Vec::into_iter().collect::>()` and back. +- **Optional / fallible** values: use this crate's `FFI_Option` and `FFI_Result` (`src/ffi_option.rs`), *not* stabby's, because ours do not require `T: IStable`. +- **Schema / arrays**: `WrappedSchema` (`src/arrow_wrappers.rs`) wraps `FFI_ArrowSchema`. Never expose `FFI_ArrowSchema` directly. +- **Logical `Expr` / `LogicalPlan`**: serialize via `datafusion-proto` using the embedded `FFI_LogicalExtensionCodec`. Same for physical plans → `FFI_PhysicalExtensionCodec`. +- **Enums** (`Volatility`, `TableType`, `InsertOp`, `TableProviderFilterPushDown`): `#[repr(u8)]`, with `From for FFI_X` and `From<&FFI_X> for Native`. Always write a round-trip unit test that exercises every variant. +- **Errors**: every `FFI_X` method that can fail returns `FFI_Result`. Use the `sresult!`, `sresult_return!`, `df_result!` macros from `src/util.rs` — do not roll your own. +- **Infallible trait methods**: if an FFI call can fail but the native trait cannot return the error, log the transport error before returning the trait's fallback (`None`, `false`, or a default). Never discard it with `.ok()`, `.unwrap_or_default()`, or equivalent. +- **Owned FFI returns**: consume an owned `FFI_X` with `From` instead of converting through `&FFI_X` and cloning across the boundary. Keep the borrowed conversion's local marker fast path. + +## Async, sessions, and task context + +- Any async method becomes `unsafe extern "C" fn(...) -> FfiFuture>`. The wrapper body uses `async move { ... }.into_ffi()`. Store `Option` in `PrivateData` so the producer side can re-enter its own runtime if needed. +- Methods taking `&dyn Session` cross the boundary as `FFI_SessionRef`. On the consumer side, try `session.as_local()` first; only construct a `ForeignSession::try_from(&session)` if that returns `None`. See `scan_fn_wrapper` in `src/table_provider.rs`. +- Anything that needs to deserialize an `Expr` / `LogicalPlan` needs a `TaskContext`. Threading a fresh `TaskContext` per call is wrong because new UDFs may have been registered since construction. Use `FFI_TaskContextProvider` (`src/execution/task_ctx_provider.rs`), which holds a `Weak` ref to a `TaskContextProvider`. If the weak ref is dead at call time, return a clear error — do not panic. + +## Memory model checklist for every new `FFI_X` + +- [ ] `private_data` is `Box::into_raw`-ed exactly once at construction. +- [ ] Every constructor path (including `clone_fn_wrapper`) allocates a fresh `Box` for its own `private_data`. +- [ ] `release_fn_wrapper` `Box::from_raw`s it and nulls the pointer. +- [ ] `Drop` calls `release`. +- [ ] No method touches `private_data` directly outside the producer side. Consumer-side methods on `ForeignX` use only the function pointers. +- [ ] `library_marker_id` and `version` are populated in **every** constructor (including `clone`). +- [ ] No method dereferences a pointer it did not check for nullness (debug-assert at minimum). + +## Quick PR-review checklist + +When reviewing a PR that touches `datafusion/ffi/`: + +1. **Trait coverage.** Pull up the underlying trait. List its methods. Confirm every non-defaulted method has a function pointer. For each *defaulted* method, ask whether a real-world producer would override it — if yes, the PR must either plumb it through as a plain `unsafe extern "C" fn` (let dynamic dispatch on `Arc` pick override-or-default) or explicitly justify the omission in a comment. Reserve `Option` for the capability-flag pattern described in §"Method coverage" — do not use it just because the trait has a default. +2. **Layout fields.** `clone`, `release`, `version`, `private_data` all present? `library_marker_id` present **iff** wrapper has a `ForeignX` adapter and a reverse `From<&?FFI_X> for {Arc,Box}` consultation site (Arc-backed for shareable traits, Box-backed for `&mut self` traits); if the trait is impl'd directly on `FFI_X` (no `ForeignX`, no reverse `From`), `library_marker_id` is dead surface and must be omitted with a one-line rationale. +3. **Marker-id bypass** in `From<&FFI_X>` (where applicable per rule 2)? +4. **Round-trip downcast** in the constructor? +5. **`Drop` calling `release`**, and `release` nulling the pointer? +6. **`Send`/`Sync` unsafe impls** match the wrapped trait's bounds (rule 4), and `FFI_X` + `ForeignX` agree? +7. **Stabby types** for strings/vecs; crate-local `FFI_Option`/`FFI_Result` for optional/fallible payloads? +8. **Async** uses `FfiFuture` + `.into_ffi()`, never blocking? +9. **Codec** present on any method that ships an `Expr` / plan? +10. **Unit tests** include both local-bypass and `mock_foreign_marker_id` forced-foreign cases? **Integration tests** in `tests/ffi_*.rs` exist for any wrapper that takes/returns non-trivial FFI types, and for *every* layout change? `cargo test -p datafusion-ffi --features integration-tests` must pass before merging an ABI-affecting PR. +11. **No `datafusion` runtime dep** crept into `Cargo.toml`? +12. **`Arc::clone(&x)`** everywhere — no implicit `x.clone()` on `Arc` (the lint will reject it but worth pre-flagging). +13. **`cargo clippy --all-targets --all-features -- -D warnings`** clean on the crate? +14. **`api change` label** on the PR if any `FFI_X` struct layout / fn-ptr signature / FFI enum / `version` extern changed? Block merge until applied. +15. **Base branch check.** If FFI struct layout changed, base branch must be `main`, never a release branch matching `^branch-\d+$` (e.g. `branch-53`). Reject back-ports of ABI-breaking changes to patch-release branches. Verify with `gh pr view --json baseRefName,labels --jq '.baseRefName'`; match strictly against `^branch-\d+$` — cherry-pick working branches like `branch-53-cherry-pick-1` also start with `branch-` and must not false-positive. + +## References + +- Crate README: `datafusion/ffi/README.md` — vocabulary + memory-model rationale. +- Canonical wrapper to model after: `src/catalog_provider.rs`. Async + capability-flag variants: `src/table_provider.rs`. +- Mutable-trait variant: `src/udaf/accumulator.rs` (`Box`). +- Optional-method pattern: `FFI_TableProvider::supports_filters_pushdown`. +- Codec wiring: `src/proto/logical_extension_codec.rs`, `src/proto/physical_extension_codec.rs`. +- Examples crate: `datafusion-examples/examples/ffi` (end-to-end producer + consumer). diff --git a/.ai/skills/pr_review/SKILL.md b/.ai/skills/pr_review/SKILL.md new file mode 100644 index 0000000000000..1de9b0726a4bf --- /dev/null +++ b/.ai/skills/pr_review/SKILL.md @@ -0,0 +1,35 @@ +--- +name: pr_review +description: Review Apache DataFusion pull requests following the project's PR review guide. Use whenever asked to review a DataFusion PR or PR URL, and whenever creating a PR, to check the changes against the same criteria before submitting. +--- + +# DataFusion PR Review + +This skill describes the mechanics for doing PR reviews from the command line. + +When creating a PR, skip the "Collect PR context" step and instead check the +changes against each area of the +[PR review guide](../../../docs/source/contributor-guide/pr_review.md) before +submitting. + +## Collect PR context + +- Check out the PR locally: `gh pr checkout ` (ask first if the + working tree has other work in progress). +- Fetch the PR description, comments, and reviews: + `gh pr view --json title,body,comments,reviews` +- Fetch CI status: `gh pr checks `. + +## Compute the diff + +```bash +# find the remote that points at apache/datafusion (e.g. `apache`, `upstream`, or `origin`) +UPSTREAM=$(git remote -v | grep -m1 'apache/datafusion' | cut -f1) +git fetch $UPSTREAM main +git diff $(git merge-base HEAD $UPSTREAM/main) +``` + +## Review checklist + +Work through each area from the +[PR review guide](../../../docs/source/contributor-guide/pr_review.md). diff --git a/.asf.yaml b/.asf.yaml index ee337fad7c136..0c04b12c4f9c0 100644 --- a/.asf.yaml +++ b/.asf.yaml @@ -78,6 +78,7 @@ github: - "cargo test (macos-aarch64)" - "Verify Vendored Code" - "Check cargo fmt" + - "Check GitHub Actions install tooling" - "clippy" - "check Cargo.toml formatting" - "check configs.md and ***_functions.md is up-to-date" @@ -97,6 +98,12 @@ github: branch-52: required_pull_request_reviews: required_approving_review_count: 1 + branch-53: + required_pull_request_reviews: + required_approving_review_count: 1 + branch-54: + required_pull_request_reviews: + required_approving_review_count: 1 pull_requests: # enable updating head branches of pull requests allow_update_branch: true @@ -108,4 +115,3 @@ github: # https://datafusion.apache.org/ publish: whoami: asf-site - diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index 955e59d74d08b..62449afbbbe1f 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -9,7 +9,7 @@ body: description: Please describe what you are trying to do. placeholder: > A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] - (This section helps Arrow developers understand the context and *why* for this feature, in addition to the *what*) + (This section helps DataFusion developers understand the context and *why* for this feature, in addition to the *what*) - type: textarea attributes: label: Describe the solution you'd like diff --git a/.github/actions/setup-rust-runtime/action.yaml b/.github/actions/setup-rust-runtime/action.yaml index e0341de93b83d..ad8fbaccc07e7 100644 --- a/.github/actions/setup-rust-runtime/action.yaml +++ b/.github/actions/setup-rust-runtime/action.yaml @@ -31,3 +31,14 @@ runs: run: | echo "RUST_BACKTRACE=1" >> $GITHUB_ENV echo "RUSTFLAGS=-C debuginfo=line-tables-only -C incremental=false" >> $GITHUB_ENV + # Work around intermittent "[16] Error in the HTTP2 framing layer" + # failures from curl when cargo fetches crates from crates.io. + # Disabling HTTP/2 multiplexing forces cargo to serialize requests, + # and raising retries makes transient network hiccups self-heal. + # + # Reference: + # https://doc.rust-lang.org/cargo/reference/config.html?#httpmultiplexing + # https://doc.rust-lang.org/cargo/reference/config.html?#netretry + echo "CARGO_HTTP_MULTIPLEXING=false" >> $GITHUB_ENV + echo "CARGO_NET_RETRY=10" >> $GITHUB_ENV + echo "CARGO_HTTP_RETRY=10" >> $GITHUB_ENV diff --git a/.github/dependabot.yml b/.github/dependabot.yml index 2cd4bdfdd7923..12ddff783b4d2 100644 --- a/.github/dependabot.yml +++ b/.github/dependabot.yml @@ -68,6 +68,10 @@ updates: interval: "weekly" open-pull-requests-limit: 10 labels: [auto-dependencies] + groups: + codeql-actions: + patterns: + - "github/codeql-action/*" - package-ecosystem: "pip" directory: "/docs" schedule: diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index 907d90523978c..01a44953e83b6 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -11,12 +11,19 @@ We generally require a GitHub issue to be filed for all bug fixes and enhancemen ## What changes are included in this PR? ## Are these changes tested? @@ -33,8 +40,6 @@ If tests are not included in your PR, please explain why (for example, are they - diff --git a/.github/workflows/audit.yml b/.github/workflows/audit.yml index 320255b595afa..875eeffdf053d 100644 --- a/.github/workflows/audit.yml +++ b/.github/workflows/audit.yml @@ -43,12 +43,14 @@ jobs: security_audit: runs-on: ubuntu-latest steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install cargo-audit - uses: taiki-e/install-action@213ccc1a076163c093f914550b94feb90fab916d # v2.79.2 + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 with: tool: cargo-audit - name: Run audit check # Note: you can ignore specific RUSTSEC issues using the `--ignore` flag ,for example: # run: cargo audit --ignore RUSTSEC-2026-0001 - run: cargo audit + # TODO: remove once object_store upgrades to quick-xml >= 0.41.0 + # https://github.com/apache/datafusion/issues/23297 + run: cargo audit --ignore RUSTSEC-2026-0194 --ignore RUSTSEC-2026-0195 diff --git a/.github/workflows/breaking_changes_detector.yml b/.github/workflows/breaking_changes_detector.yml index 158cc17e94d0e..3a279a27f54d9 100644 --- a/.github/workflows/breaking_changes_detector.yml +++ b/.github/workflows/breaking_changes_detector.yml @@ -55,7 +55,7 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 @@ -89,7 +89,7 @@ jobs: - name: Install cargo-semver-checks if: steps.changed_crates.outputs.packages != '' - uses: taiki-e/install-action@213ccc1a076163c093f914550b94feb90fab916d # v2.79.2 + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 with: tool: cargo-semver-checks diff --git a/.github/workflows/breaking_changes_detector_comment.yml b/.github/workflows/breaking_changes_detector_comment.yml index 579c61cb9d5c7..f3a3400d00f9c 100644 --- a/.github/workflows/breaking_changes_detector_comment.yml +++ b/.github/workflows/breaking_changes_detector_comment.yml @@ -104,39 +104,66 @@ jobs: echo "${DELIM}" } >> "$GITHUB_OUTPUT" - # The marker `` is what makes the comment - # "sticky": maintain-one-comment uses it to find and replace (or - # delete) the existing comment instead of stacking new ones. + + # Find any existing sticky comment by its hidden marker so we can update + # or delete it instead of stacking new ones. + - name: Find existing sticky comment + id: find + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + PR_NUMBER: ${{ steps.read.outputs.pr_number }} + run: | + COMMENT_ID=$(gh api --paginate "repos/${REPO}/issues/${PR_NUMBER}/comments" \ + --jq '.[] | select(.body | contains("")) | .id' \ + | head -n1) + echo "comment_id=${COMMENT_ID}" >> "$GITHUB_OUTPUT" + + # update the existing comment found above, or create a new one. The hidden + # marker `` stays in the body so the next run + # finds it again. LOGS is interpolated via a shell parameter expansion, + # whose result bash does not re-scan, so untrusted log content cannot + # inject further commands. - name: Upsert sticky comment if: steps.read.outputs.result != 'success' - uses: actions-cool/maintain-one-comment@909842216bc8e8658364c572ec52100f4c2cc50a # v3.3.0 - with: - token: ${{ secrets.GITHUB_TOKEN }} - number: ${{ steps.read.outputs.pr_number }} - body-include: '' - body: | - - Thank you for opening this pull request! - - Reviewer note: [cargo-semver-checks](https://github.com/obi1kenobi/cargo-semver-checks) reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). - -
- Details - - ``` - ${{ steps.read.outputs.logs }} - ``` - -
+ env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + PR_NUMBER: ${{ steps.read.outputs.pr_number }} + COMMENT_ID: ${{ steps.find.outputs.comment_id }} + LOGS: ${{ steps.read.outputs.logs }} + run: | + set -euo pipefail + BODY=" + Thank you for opening this pull request! + + Reviewer note: [cargo-semver-checks](https://github.com/obi1kenobi/cargo-semver-checks) reported the current version number is not SemVer-compatible with the changes in this pull request (compared against the base branch). + +
+ Details + + \`\`\` + ${LOGS} + \`\`\` + +
" + + # Use --raw-field (not --field): always sends the value as a literal string. while --field would treat a leading `@` as a file to read + # (even though the body does not start with user input we are being cautious) + if [ -n "$COMMENT_ID" ]; then + gh api "repos/${REPO}/issues/comments/${COMMENT_ID}" --method PATCH --raw-field body="$BODY" + else + gh api "repos/${REPO}/issues/${PR_NUMBER}/comments" --method POST --raw-field body="$BODY" + fi + # Clear a stale comment once the breaking change is resolved. - name: Delete sticky comment - if: steps.read.outputs.result == 'success' - uses: actions-cool/maintain-one-comment@909842216bc8e8658364c572ec52100f4c2cc50a # v3.3.0 - with: - token: ${{ secrets.GITHUB_TOKEN }} - number: ${{ steps.read.outputs.pr_number }} - body-include: '' - delete: true + if: steps.read.outputs.result == 'success' && steps.find.outputs.comment_id != '' + env: + GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} + REPO: ${{ github.repository }} + COMMENT_ID: ${{ steps.find.outputs.comment_id }} + run: gh api -X DELETE "repos/${REPO}/issues/comments/${COMMENT_ID}" - name: Add "auto detected api change" label if: steps.read.outputs.result != 'success' diff --git a/.github/workflows/codeql.yml b/.github/workflows/codeql.yml index 4716a8c5bcded..9c2de289177c2 100644 --- a/.github/workflows/codeql.yml +++ b/.github/workflows/codeql.yml @@ -40,16 +40,16 @@ jobs: steps: - name: Checkout repository - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false - name: Initialize CodeQL - uses: github/codeql-action/init@9e0d7b8d25671d64c341c19c0152d693099fb5ba # v4 + uses: github/codeql-action/init@5595ccaf912efad79be6eef63a5619ff05969be3 # v4 with: languages: actions - name: Perform CodeQL Analysis - uses: github/codeql-action/analyze@9e0d7b8d25671d64c341c19c0152d693099fb5ba # v4 + uses: github/codeql-action/analyze@5595ccaf912efad79be6eef63a5619ff05969be3 # v4 with: category: "/language:actions" diff --git a/.github/workflows/dependencies.yml b/.github/workflows/dependencies.yml index 2f3a127ef98c4..f76ced5296871 100644 --- a/.github/workflows/dependencies.yml +++ b/.github/workflows/dependencies.yml @@ -25,6 +25,7 @@ on: push: branches-ignore: - 'gh-readonly-queue/**' + - 'dependabot/**' pull_request: merge_group: # manual trigger @@ -41,7 +42,7 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -50,9 +51,9 @@ jobs: with: rust-version: stable - name: Check dependencies + working-directory: dev/depcheck run: | - cd dev/depcheck - cargo run + cargo run --locked detect-unused-dependencies: name: Detect Unused Dependencies @@ -60,8 +61,10 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install cargo-machete - run: cargo install cargo-machete --version ^0.9 --locked + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 + with: + tool: cargo-machete@0.9 - name: Detect unused dependencies run: cargo machete --with-metadata diff --git a/.github/workflows/dev.yml b/.github/workflows/dev.yml index 0cb71cc14e1ab..a8c1b8c07af96 100644 --- a/.github/workflows/dev.yml +++ b/.github/workflows/dev.yml @@ -20,6 +20,7 @@ on: push: branches-ignore: - 'gh-readonly-queue/**' + - 'dependabot/**' pull_request: merge_group: @@ -35,10 +36,11 @@ jobs: runs-on: ubuntu-latest name: Check License Header steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Install HawkEye - # This CI job is bound by installation time, use `--profile dev` to speed it up - run: cargo install hawkeye --version 6.2.0 --locked --profile dev + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 + with: + tool: hawkeye@6.2.0 - name: Run license header check run: ci/scripts/license_header.sh @@ -46,8 +48,8 @@ jobs: name: Use prettier to check formatting of documents runs-on: ubuntu-slim steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - - uses: actions/setup-node@6044e13b5dc448c55e2357c09f80417699197238 # v6.2.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 with: node-version: "20" - name: Prettier check @@ -58,13 +60,13 @@ jobs: name: Check Markdown Links runs-on: ubuntu-latest steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Load tool versions run: | source ci/scripts/utils/tool_versions.sh echo "LYCHEE_VERSION=${LYCHEE_VERSION}" >> "$GITHUB_ENV" - name: Install lychee - uses: taiki-e/install-action@213ccc1a076163c093f914550b94feb90fab916d # v2.79.2 + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 with: tool: lychee@${{ env.LYCHEE_VERSION }} - name: Run markdown link check @@ -74,7 +76,7 @@ jobs: name: Validate required_status_checks in .asf.yaml runs-on: ubuntu-latest steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - run: pip install pyyaml - run: python3 ci/scripts/check_asf_yaml_status_checks.py @@ -82,13 +84,15 @@ jobs: name: Spell Check with Typos runs-on: ubuntu-latest steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: persist-credentials: false # Version fixed on purpose. It uses heuristics to detect typos, so upgrading # it may cause checks to fail more often. # We can upgrade it manually once a while. - - name: Install typos-cli - run: cargo install typos-cli --locked --version 1.37.0 + - name: Install typos + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 + with: + tool: typos@1.37.0 - name: Run typos check run: ci/scripts/typos_check.sh diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml index f0fbea566af69..05f48d885fb5b 100644 --- a/.github/workflows/docs.yaml +++ b/.github/workflows/docs.yaml @@ -34,25 +34,28 @@ jobs: runs-on: ubuntu-latest steps: - name: Checkout docs sources - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Checkout asf-site branch - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: asf-site path: asf-site - name: Setup uv - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 - name: Install dependencies run: uv sync --package datafusion-docs - - name: Install dependency graph tooling + - name: Install Graphviz run: | set -x sudo apt-get update sudo apt-get install -y graphviz - cargo install cargo-depgraph --version ^1.6 --locked + - name: Install cargo-depgraph + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 + with: + tool: cargo-depgraph@1.6 - name: Build docs run: | diff --git a/.github/workflows/docs_pr.yaml b/.github/workflows/docs_pr.yaml index 4b8d25b0611eb..d476d40d065b8 100644 --- a/.github/workflows/docs_pr.yaml +++ b/.github/workflows/docs_pr.yaml @@ -26,6 +26,9 @@ on: push: paths: - "docs/**" + branches: + - main + - branch-* pull_request: paths: - "docs/**" @@ -42,20 +45,23 @@ jobs: name: Test doc build runs-on: ubuntu-latest steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 - name: Setup uv - uses: astral-sh/setup-uv@08807647e7069bb48b6ef5acd8ec9567f424441b # v8.1.0 + uses: astral-sh/setup-uv@c771a70e6277c0a99b617c7a806ffedaca235ff9 # v9.0.0 - name: Install doc dependencies run: uv sync --package datafusion-docs - - name: Install dependency graph tooling + - name: Install Graphviz run: | set -x sudo apt-get update sudo apt-get install -y graphviz - cargo install cargo-depgraph --version ^1.6 --locked + - name: Install cargo-depgraph + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 + with: + tool: cargo-depgraph@1.6 - name: Build docs html and check for warnings run: | set -x diff --git a/.github/workflows/extended.yml b/.github/workflows/extended.yml index a143cb49fd35b..a6e303e3d6ff4 100644 --- a/.github/workflows/extended.yml +++ b/.github/workflows/extended.yml @@ -63,23 +63,24 @@ jobs: runs-on: ${{ vars.USE_RUNS_ON == 'true' && format('runs-on={0},family=m8a+m7a+c8a,cpu=32,image=ubuntu24-full-x64,extras=s3-cache,disk=large,tag=datafusion', github.run_id) || 'ubuntu-latest' }} # note: do not use amd/rust container to preserve disk space steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ github.event.inputs.pr_head_sha }} # will be empty if triggered by push submodules: true fetch-depth: 1 - name: Free Disk Space (Ubuntu) uses: jlumbroso/free-disk-space@54081f138730dfa15788a46383842cd2f914a1be # v1.3.1 - - name: Install Rust - run: | - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y - source $HOME/.cargo/env - rustup toolchain install - - name: Install Protobuf Compiler - run: | - sudo apt-get update - sudo apt-get install -y protobuf-compiler + - parallel: + - name: Install Rust + run: | + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y + source $HOME/.cargo/env + rustup toolchain install + - name: Install Protobuf Compiler + run: | + sudo apt-get update + sudo apt-get install -y protobuf-compiler # For debugging, test binaries can be large. - name: Show available disk space run: | @@ -98,10 +99,11 @@ jobs: --tests \ --bins \ --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption - - name: Verify Working Directory Clean - run: git diff --exit-code - - name: Cleanup - run: cargo clean + - parallel: + - name: Verify Working Directory Clean + run: git diff --exit-code + - name: Cleanup + run: cargo clean # Check answers are correct when hash values collide hash-collisions: @@ -110,8 +112,8 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: ref: ${{ github.event.inputs.pr_head_sha }} # will be empty if triggered by push submodules: true @@ -132,16 +134,17 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - with: - ref: ${{ github.event.inputs.pr_head_sha }} # will be empty if triggered by push - submodules: true - fetch-depth: 1 - # Don't use setup-builder to avoid configuring RUST_BACKTRACE which is expensive - - name: Install protobuf compiler - run: | - apt-get update && apt-get install -y protobuf-compiler + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - parallel: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + with: + ref: ${{ github.event.inputs.pr_head_sha }} # will be empty if triggered by push + submodules: true + fetch-depth: 1 + # Don't use setup-builder to avoid configuring RUST_BACKTRACE which is expensive + - name: Install protobuf compiler + run: | + apt-get update && apt-get install -y protobuf-compiler - name: Run sqllogictest run: | - cargo test --features backtrace,parquet_encryption --profile ci-optimized --test sqllogictests -- --include-sqlite \ No newline at end of file + cargo test --features backtrace,parquet_encryption --profile ci-optimized --test sqllogictests -- --include-sqlite diff --git a/.github/workflows/labeler.yml b/.github/workflows/labeler.yml index a3714a4a7c8fe..d47bd76c0caa2 100644 --- a/.github/workflows/labeler.yml +++ b/.github/workflows/labeler.yml @@ -44,7 +44,7 @@ jobs: github.event_name == 'pull_request_target' && (github.event.action == 'opened' || github.event.action == 'synchronize') - uses: actions/labeler@f27b608878404679385c85cfa523b85ccb86e213 # v6.1.0 + uses: actions/labeler@bf12e9b00b37c5c0ca2b87b79b2daf7891dbda13 # v7.0.0 with: repo-token: ${{ secrets.GITHUB_TOKEN }} configuration-path: .github/workflows/labeler/labeler-config.yml diff --git a/.github/workflows/large_files.yml b/.github/workflows/large_files.yml index 5a127e443fcb7..2648988a7d3dd 100644 --- a/.github/workflows/large_files.yml +++ b/.github/workflows/large_files.yml @@ -32,7 +32,7 @@ jobs: check-files: runs-on: ubuntu-slim steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: fetch-depth: 0 - name: Check size of new Git objects diff --git a/.github/workflows/rust.yml b/.github/workflows/rust.yml index 5ff1f6467bbf1..eaa9b21a1b343 100644 --- a/.github/workflows/rust.yml +++ b/.github/workflows/rust.yml @@ -27,6 +27,7 @@ on: push: branches-ignore: - 'gh-readonly-queue/**' + - 'dependabot/**' paths-ignore: - "docs/**" - "**.md" @@ -50,14 +51,14 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: rust-version: stable - name: Rust Dependency Cache - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: shared-key: "amd-ci-check" # this job uses it's own cache becase check has a separate cache and we need it to be fast as it blocks other jobs save-if: ${{ github.ref_name == 'main' }} @@ -79,7 +80,7 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -104,13 +105,13 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: rust-version: stable - name: Rust Dependency Cache - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: save-if: false # set in linux-test shared-key: "amd-ci" @@ -141,8 +142,8 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -161,6 +162,26 @@ jobs: - name: Check datafusion-proto (avro) run: cargo check --profile ci --no-default-features -p datafusion-proto --features=avro + # Check datafusion-ffi features + # + # Ensure via `cargo check` that the crate can be built with a + # subset of the features packages enabled. + linux-datafusion-ffi-features: + name: cargo check datafusion-ffi features + needs: linux-build-lib + runs-on: ${{ vars.USE_RUNS_ON == 'true' && format('runs-on={0},family=m8a+m7a+c8a,cpu=16,image=ubuntu24-full-x64,extras=s3-cache,disk=large,tag=datafusion', github.run_id) || 'ubuntu-latest' }} + container: + image: amd64/rust + steps: + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Setup Rust toolchain + uses: ./.github/actions/setup-builder + with: + rust-version: stable + - name: Check datafusion-ffi (no-default-features) + run: cargo check --profile ci --no-default-features -p datafusion-ffi + # Check datafusion crate features # @@ -173,14 +194,14 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: rust-version: stable - name: Rust Dependency Cache - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: save-if: false # set in linux-test shared-key: "amd-ci" @@ -239,7 +260,7 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -276,8 +297,8 @@ jobs: volumes: - /usr/local:/host/usr/local steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -285,16 +306,22 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: stable + - name: Install llvm-tools-preview + run: rustup component add llvm-tools-preview + - name: Install cargo-llvm-cov + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 + with: + tool: cargo-llvm-cov - name: Rust Dependency Cache - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: - save-if: ${{ github.ref_name == 'main' }} - shared-key: "amd-ci" + save-if: ${{ github.ref_name == 'main' }} + shared-key: "amd-ci" - name: Run tests (excluding doctests and datafusion-cli) env: RUST_BACKTRACE: 1 run: | - cargo test \ + cargo llvm-cov \ --profile ci \ --exclude datafusion-examples \ --exclude ffi_example_table_provider \ @@ -303,18 +330,27 @@ jobs: --lib \ --tests \ --bins \ - --features serde,avro,json,backtrace,integration-tests,parquet_encryption,substrait - - name: Verify Working Directory Clean - run: git diff --exit-code - # Check no temporary directories created during test. - # `false/` folder is excuded for rust cache. - - name: Verify Working Directory Clean (No Untracked Files) - run: | - STATUS="$(git status --porcelain | sed -e '/^?? false\/$/d' -e '/^?? false$/d')" - if [ -n "$STATUS" ]; then - echo "$STATUS" - exit 1 - fi + --features serde,avro,json,backtrace,integration-tests,parquet_encryption,substrait \ + --codecov \ + --output-path target/codecov.json + - parallel: + - name: Verify Working Directory Clean + run: git diff --exit-code + # Check no temporary directories created during test. + # `false/` folder is excuded for rust cache. + - name: Verify Working Directory Clean (No Untracked Files) + run: | + STATUS="$(git status --porcelain | sed -e '/^?? false\/$/d' -e '/^?? false$/d')" + if [ -n "$STATUS" ]; then + echo "$STATUS" + exit 1 + fi + - name: Upload coverage to codecov.io + uses: codecov/codecov-action@fb8b3582c8e4def4969c97caa2f19720cb33a72f # v7.0.0 + with: + files: target/codecov.json + fail_ci_if_error: false + token: ${{ secrets.CODECOV_TOKEN }} # datafusion-cli tests linux-test-datafusion-cli: @@ -322,15 +358,15 @@ jobs: needs: linux-build-lib runs-on: ${{ vars.USE_RUNS_ON == 'true' && format('runs-on={0},family=m8a+m7a+c8a,cpu=16,image=ubuntu24-full-x64,extras=s3-cache,disk=large,tag=datafusion', github.run_id) || 'ubuntu-latest' }} steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 - name: Setup Rust toolchain run: rustup toolchain install stable - name: Rust Dependency Cache - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: save-if: false # set in linux-test shared-key: "amd-ci" @@ -354,8 +390,8 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -364,7 +400,7 @@ jobs: with: rust-version: stable - name: Rust Dependency Cache - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 with: save-if: ${{ github.ref_name == 'main' }} shared-key: "amd-ci-linux-test-example" @@ -385,8 +421,8 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -407,8 +443,8 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -420,18 +456,19 @@ jobs: name: build and run with wasm-pack runs-on: ubuntu-24.04 steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 - - name: Setup for wasm32 - run: | - rustup target add wasm32-unknown-unknown - - name: Install dependencies - run: | - sudo apt-get update -qq - sudo apt-get install -y -qq clang - - name: Setup wasm-pack - uses: taiki-e/install-action@213ccc1a076163c093f914550b94feb90fab916d # v2.79.2 - with: - tool: wasm-pack + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - parallel: + - name: Setup for wasm32 + run: | + rustup target add wasm32-unknown-unknown + - name: Install dependencies + run: | + sudo apt-get update -qq + sudo apt-get install -y -qq clang + - name: Setup wasm-pack + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 + with: + tool: wasm-pack - name: Run tests with headless mode working-directory: ./datafusion/wasmtest run: | @@ -448,23 +485,24 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 - - name: Setup Rust toolchain - uses: ./.github/actions/setup-builder - with: - rust-version: stable - - name: Generate benchmark data and expected query results - run: | - mkdir -p datafusion/sqllogictest/test_files/tpch/data - git clone https://github.com/databricks/tpch-dbgen.git - cd tpch-dbgen - make - ./dbgen -f -s 0.1 - mv *.tbl ../datafusion/sqllogictest/test_files/tpch/data + - parallel: + - name: Setup Rust toolchain + uses: ./.github/actions/setup-builder + with: + rust-version: stable + - name: Generate benchmark data and expected query results + run: | + mkdir -p datafusion/sqllogictest/test_files/tpch/data + git clone https://github.com/databricks/tpch-dbgen.git + cd tpch-dbgen + make + ./dbgen -f -s 0.1 + mv *.tbl ../datafusion/sqllogictest/test_files/tpch/data - name: Verify that benchmark queries return expected results run: | # increase stack size to fix stack overflow @@ -496,8 +534,8 @@ jobs: --health-timeout 5s --health-retries 5 steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -521,8 +559,8 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -565,7 +603,7 @@ jobs: name: cargo test (macos-aarch64) runs-on: macos-15 steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -581,7 +619,7 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: @@ -598,57 +636,23 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder with: rust-version: stable - name: Run run: | - echo '' > datafusion/proto/src/generated/datafusion.rs ci/scripts/rust_fmt.sh - # Coverage job disabled due to - # https://github.com/apache/datafusion/issues/3678 - - # coverage: - # name: coverage - # runs-on: ubuntu-latest - # steps: - # - uses: actions/checkout@v4 - # with: - # submodules: true - # - name: Install protobuf compiler - # shell: bash - # run: | - # mkdir -p $HOME/d/protoc - # cd $HOME/d/protoc - # export PROTO_ZIP="protoc-21.4-linux-x86_64.zip" - # curl -LO https://github.com/protocolbuffers/protobuf/releases/download/v21.4/$PROTO_ZIP - # unzip $PROTO_ZIP - # export PATH=$PATH:$HOME/d/protoc/bin - # protoc --version - # - name: Setup Rust toolchain - # run: | - # rustup toolchain install stable - # rustup default stable - # rustup component add rustfmt clippy - # - name: Cache Cargo - # uses: actions/cache@v4 - # with: - # path: /home/runner/.cargo - # # this key is not equal because the user is different than on a container (runner vs github) - # key: cargo-coverage-cache3- - # - name: Run coverage - # run: | - # export PATH=$PATH:$HOME/d/protoc/bin - # rustup toolchain install stable - # rustup default stable - # cargo install --version 0.20.1 cargo-tarpaulin - # cargo tarpaulin --all --out Xml - # - name: Report coverage - # continue-on-error: true - # run: bash <(curl -s https://codecov.io/bash) + check-workflow-tool-installs: + name: Check GitHub Actions install tooling + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 + - name: Check workflow tool installs + run: ci/scripts/check_no_cargo_install_in_workflows.sh + clippy: name: clippy @@ -657,8 +661,8 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -666,13 +670,14 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: stable - - name: Install Clippy - run: rustup component add clippy - - name: Rust Dependency Cache - uses: Swatinem/rust-cache@c19371144df3bb44fab255c43d04cbc2ab54d1c4 # v2.9.1 - with: - save-if: ${{ github.ref_name == 'main' }} - shared-key: "amd-ci-clippy" + - parallel: + - name: Install Clippy + run: rustup component add clippy + - name: Rust Dependency Cache + uses: Swatinem/rust-cache@6323deb102c322ba6fcbdcafc7e3dddab59af2b6 # v2.9.2 + with: + save-if: ${{ github.ref_name == 'main' }} + shared-key: "amd-ci-clippy" - name: Run clippy run: ci/scripts/rust_clippy.sh @@ -683,7 +688,7 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -692,7 +697,9 @@ jobs: with: rust-version: stable - name: Install taplo - run: cargo +stable install taplo-cli --version ^0.9 --locked + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 + with: + tool: taplo-cli@0.9 # if you encounter an error, try running 'taplo format' to fix the formatting automatically. - name: Check Cargo.toml formatting run: taplo format --check @@ -704,8 +711,8 @@ jobs: container: image: amd64/rust steps: - - uses: runs-on/action@d141ef83eb66d096ce8afc767e09115a65c63b60 # v2.1.2 - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: runs-on/action@46910bf61b41721b0579f237e186afb35477007a # v2.3.0 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 @@ -713,7 +720,7 @@ jobs: uses: ./.github/actions/setup-builder with: rust-version: stable - - uses: actions/setup-node@6044e13b5dc448c55e2357c09f80417699197238 # v6.2.0 + - uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 with: node-version: "20" - name: Check if configs.md has been modified @@ -740,20 +747,21 @@ jobs: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 with: submodules: true fetch-depth: 1 - - name: Mark repository as safe for git - # Required for git commands inside container (avoids "dubious ownership" error) - run: git config --global --add safe.directory "$GITHUB_WORKSPACE" + - parallel: + - name: Mark repository as safe for git + # Required for git commands inside container (avoids "dubious ownership" error) + run: git config --global --add safe.directory "$GITHUB_WORKSPACE" - - name: Set up Node.js (required for prettier) - # doc_prettier_check.sh uses npx to run prettier for Markdown formatting - uses: actions/setup-node@v6 - with: - node-version: '18' + - name: Set up Node.js (required for prettier) + # doc_prettier_check.sh uses npx to run prettier for Markdown formatting + uses: actions/setup-node@820762786026740c76f36085b0efc47a31fe5020 # v7.0.0 + with: + node-version: '18' - name: Run examples docs check script run: | @@ -770,11 +778,11 @@ jobs: container: image: amd64/rust steps: - - uses: actions/checkout@de0fac2e4500dabe0009e67214ff5f5447ce83dd # v6.0.2 + - uses: actions/checkout@3d3c42e5aac5ba805825da76410c181273ba90b1 # v7.0.1 - name: Setup Rust toolchain uses: ./.github/actions/setup-builder - name: Install cargo-msrv - uses: taiki-e/install-action@213ccc1a076163c093f914550b94feb90fab916d # v2.79.2 + uses: taiki-e/install-action@6c6fd71fe4fb72c3697d269963d0e15df8adedad # v2.85.10 with: tool: cargo-msrv diff --git a/.github/workflows/stale.yml b/.github/workflows/stale.yml index 8627b3bf044ff..81188559d89f0 100644 --- a/.github/workflows/stale.yml +++ b/.github/workflows/stale.yml @@ -28,7 +28,7 @@ jobs: issues: write pull-requests: write steps: - - uses: actions/stale@b5d41d4e1d5dceea10e7104786b73624c18a190f # v10.2.0 + - uses: actions/stale@4391f3da665fdf50b6810c1a66712fb9ba21aa93 # v11.0.0 with: stale-pr-message: "Thank you for your contribution. Unfortunately, this pull request is stale because it has been open 60 days with no activity. Please remove the stale label or comment or this will be closed in 7 days." days-before-pr-stale: 60 diff --git a/.gitignore b/.gitignore index c1f9677e47366..2bcc0950d01b3 100644 --- a/.gitignore +++ b/.gitignore @@ -73,9 +73,6 @@ datafusion/core/benches/data/* filtered_rat.txt rat.txt -# data generated by examples -datafusion-examples/examples/datafusion-examples/ - # Samply profile data profile.json.gz diff --git a/AGENTS.md b/AGENTS.md index 9dff7f6f1ffd1..8fdf314ed4b6c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -5,6 +5,7 @@ - [Quick Start Setup](docs/source/contributor-guide/development_environment.md#quick-start) - [Testing Quick Start](docs/source/contributor-guide/testing.md#testing-quick-start) - [Before Submitting a PR](docs/source/contributor-guide/index.md#before-submitting-a-pr) +- [Reviewing Pull Requests](docs/source/contributor-guide/pr_review.md) - [Contributor Guide](docs/source/contributor-guide/index.md) - [Architecture Guide](docs/source/contributor-guide/architecture.md) @@ -37,5 +38,24 @@ When creating a PR, you MUST follow the [PR template](.github/pull_request_templ ## Testing -See the [Testing Quick Start](docs/source/contributor-guide/testing.md#testing-quick-start) -for the recommended pre-PR test commands. +If documentation files changed then run +```bash +./ci/scripts/doc_prettier_check.sh --write --allow-dirty +``` + +Otherwise, run extended tests +```bash +RUST_BACKTRACE=1 cargo test --profile ci \ + --exclude datafusion-examples --exclude datafusion-benchmarks --exclude datafusion-cli \ + --workspace --lib --tests --bins \ + --features avro,json,backtrace,extended_tests,recursive_protection,parquet_encryption +``` + +For modified code identify local benchmarks(if any) and run them against `main`. See [Benchmarks](benchmarks/README.md). + +## Agent Skills + +Repository-specific agent skills live under `.ai/skills/`. Each subdirectory is +a single skill with a `SKILL.md` (YAML frontmatter + body). Check that +directory for applicable skills before working on a task; new skills go in +`.ai/skills//SKILL.md`. diff --git a/Cargo.lock b/Cargo.lock index 4d5b15075ecef..62c04d332b98e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -128,9 +128,9 @@ dependencies = [ [[package]] name = "anyhow" -version = "1.0.102" +version = "1.0.103" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" [[package]] name = "ar_archive_writer" @@ -141,6 +141,15 @@ dependencies = [ "object", ] +[[package]] +name = "arc-swap" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +dependencies = [ + "rustversion", +] + [[package]] name = "arrayref" version = "0.3.9" @@ -155,9 +164,9 @@ checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" [[package]] name = "arrow" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "378530e55cd479eda3c14eb345310799717e6f76d0c332041e8487022166b471" +checksum = "61d285d16bce7d0be61912f7928342b673067b6b7d7ef6cc179258ba7de1fecf" dependencies = [ "arrow-arith", "arrow-array", @@ -178,9 +187,9 @@ dependencies = [ [[package]] name = "arrow-arith" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a0ab212d2c1886e802f51c5212d78ebbcbb0bec980fff9dadc1eb8d45cd0b738" +checksum = "757ef1836251e88222542a7da2623bc1c9cb9e20afefa6db2c41e79991cd91d4" dependencies = [ "arrow-array", "arrow-buffer", @@ -192,9 +201,9 @@ dependencies = [ [[package]] name = "arrow-array" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cfd33d3e92f207444098c75b42de99d329562be0cf686b307b097cc52b4e999e" +checksum = "bc9a4a4b2b5ecd0e04df03471661cb61f28bed3c7fd50994715129b01b2edb97" dependencies = [ "ahash", "arrow-buffer", @@ -204,6 +213,7 @@ dependencies = [ "chrono-tz", "half", "hashbrown 0.17.1", + "libc", "num-complex", "num-integer", "num-traits", @@ -211,9 +221,9 @@ dependencies = [ [[package]] name = "arrow-avro" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "049230728cd6e093088c8d231b4beede184e35cad7777c1505c0d5a8571f4376" +checksum = "9fb45cd6bd2b25c0965793b83200eaca82214273a8030fbbc2d783e4c7c65a61" dependencies = [ "arrow-array", "arrow-buffer", @@ -235,21 +245,21 @@ dependencies = [ [[package]] name = "arrow-buffer" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0c6cd424c2693bcdbc150d843dc9d4d137dd2de4782ce6df491ad11a3a0416c0" +checksum = "c12b576ef18c1deb80925a248b25ad84f419198d791b8e293fc6aaa60441fe90" dependencies = [ "bytes", "half", - "num-bigint", + "num-bigint 0.5.1", "num-traits", ] [[package]] name = "arrow-cast" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c5aefb56a2c02e9e2b30746241058b85f8983f0fcff2ba0c6d09006e1cded7f" +checksum = "68338a9096a5dc9bc11927c58c43a8526d96bf6abd2012ef6c0c9f505991cc79" dependencies = [ "arrow-array", "arrow-buffer", @@ -258,7 +268,7 @@ dependencies = [ "arrow-schema", "arrow-select", "atoi", - "base64 0.22.1", + "base64 0.23.1", "chrono", "comfy-table", "half", @@ -269,9 +279,9 @@ dependencies = [ [[package]] name = "arrow-csv" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e94e8cf7e517657a52b91ea1263acf38c4ca62a84655d72458a3359b12ab97de" +checksum = "25011b52b346407d497ef0030e12b45e4f2d0cc279efc09c4f3d09106db30e36" dependencies = [ "arrow-array", "arrow-cast", @@ -284,9 +294,9 @@ dependencies = [ [[package]] name = "arrow-data" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3c88210023a2bfee1896af366309a3028fc3bcbd6515fa29a7990ee1baa08ee0" +checksum = "723fe4aeed7604e00b9883a465af4ff0a0e6c44c03e41a68c3d1cbc403e0e44d" dependencies = [ "arrow-buffer", "arrow-schema", @@ -297,9 +307,9 @@ dependencies = [ [[package]] name = "arrow-flight" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "28abfe8bf9f124e5fc83b334af4fa58f8d0323ad25312ccb2d1da50178415704" +checksum = "2bebfacc9d71f0728f6774164e4d4254b5e504d2b46812d0512d8290ec119a64" dependencies = [ "arrow-arith", "arrow-array", @@ -312,11 +322,10 @@ dependencies = [ "arrow-schema", "arrow-select", "arrow-string", - "base64 0.22.1", + "base64 0.23.1", "bytes", "futures", "once_cell", - "paste", "prost", "prost-types", "tonic", @@ -325,9 +334,9 @@ dependencies = [ [[package]] name = "arrow-ipc" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "238438f0834483703d88896db6fe5a7138b2230debc31b34c0336c2996e3c64f" +checksum = "149437b14371f5b9ec60f5ddc751483ae99d7a7072653c0075e5e469156eea7b" dependencies = [ "arrow-array", "arrow-buffer", @@ -341,9 +350,9 @@ dependencies = [ [[package]] name = "arrow-json" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "205ca2119e6d679d5c133c6f30e68f027738d95ed948cf77677ea69c7800036b" +checksum = "f18b9123ccfec418a663f821c9a034af339711678c11ffe00d3ec07da5ff9f7e" dependencies = [ "arrow-array", "arrow-buffer", @@ -366,9 +375,9 @@ dependencies = [ [[package]] name = "arrow-ord" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1bffd8fd2579286a5d63bac898159873e5094a79009940bcb42bbfce4f19f1d0" +checksum = "e6c08dff0686cf23ca4f562803f191ccbeb726dbae6309cd4b4aaf65e0f2c979" dependencies = [ "arrow-array", "arrow-buffer", @@ -379,9 +388,9 @@ dependencies = [ [[package]] name = "arrow-row" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bab5994731204603c73ba69267616c50f80780774c6bb0476f1f830625115e0c" +checksum = "bbec439386df71ad570e6758a946111322b9e9dc8db83b5527321f0b4c9119c2" dependencies = [ "arrow-array", "arrow-buffer", @@ -392,9 +401,9 @@ dependencies = [ [[package]] name = "arrow-schema" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f633dbfdf39c039ada1bf9e34c694816eb71fbb7dc78f613993b7245e078a1ed" +checksum = "e6fed2ca0d1eade57e811cbe73b98ad50cc08a1183e13b2d2aa43a7df593f40e" dependencies = [ "bitflags", "serde", @@ -404,9 +413,9 @@ dependencies = [ [[package]] name = "arrow-select" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8cd065c54172ac787cf3f2f8d4107e0d3fdc26edba76fdf4f4cc170258942222" +checksum = "466b19cf75130b891dc1b23a84b343c714c62c64c9c62e365c76aa0ff90a53fb" dependencies = [ "ahash", "arrow-array", @@ -418,9 +427,9 @@ dependencies = [ [[package]] name = "arrow-string" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "29dd7cda3ab9692f43a2e4acc444d760cc17b12bb6d8232ddf64e9bab7c06b42" +checksum = "c838a25bb3691e919e0f617616ac51a4ff8517a952e29ca133cf0c22b2ce65b1" dependencies = [ "arrow-array", "arrow-buffer", @@ -451,9 +460,9 @@ dependencies = [ [[package]] name = "async-compression" -version = "0.4.42" +version = "0.4.43" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e79b3f8a79cccc2898f31920fc69f304859b3bd567490f75ebf51ae1c792a9ac" +checksum = "3976abdc8fe7d1133d43d304afd42abdf5bc3e1319d263d223bde07b5efc4be8" dependencies = [ "compression-codecs", "compression-core", @@ -463,9 +472,9 @@ dependencies = [ [[package]] name = "async-ffi" -version = "0.5.0" +version = "0.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f4de21c0feef7e5a556e51af767c953f0501f7f300ba785cc99c47bdc8081a50" +checksum = "39cd9de47399986d5b216c6bef9434dfff1689ab61ba8d1e2720dc5fe5c84083" [[package]] name = "async-recursion" @@ -475,7 +484,7 @@ checksum = "3b43422f69d8ff38f95f1b2bb76517c91589a924d1559a0e935d7c8ce0274c11" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -497,18 +506,18 @@ checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "async-trait" -version = "0.1.89" +version = "0.1.92" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +checksum = "82f6aeea286b8eb4dd3431a1be1b59d290ace00f5bfd8e2a159bc2a05e2c1667" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.3", ] [[package]] @@ -534,9 +543,9 @@ checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" [[package]] name = "aws-config" -version = "1.8.16" +version = "1.10.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "50f156acdd2cf55f5aa53ee416c4ac851cf1222694506c0b1f78c85695e9ca9d" +checksum = "1b180a3c8b55960db3426d8964b8745e652466a1a49fe1a2eda828046d30b5e4" dependencies = [ "aws-credential-types", "aws-runtime", @@ -548,6 +557,7 @@ dependencies = [ "aws-smithy-json", "aws-smithy-runtime", "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", "aws-types", "bytes", @@ -564,9 +574,9 @@ dependencies = [ [[package]] name = "aws-credential-types" -version = "1.2.14" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8f20799b373a1be121fe3005fba0c2090af9411573878f224df44b42727fcaf7" +checksum = "e93964ffdaf57857f544be3666a5f57570bb699e934700f11b49708f61bb556e" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", @@ -598,9 +608,9 @@ dependencies = [ [[package]] name = "aws-runtime" -version = "1.7.3" +version = "1.9.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dcd93c82209ac7413532388067dce79be5a8780c1786e5fae3df22e4dee2864" +checksum = "c9007227e10b5fed2f3e0a2beff489211e2b5604c400b7a9d5d81ca9d64c24bb" dependencies = [ "aws-credential-types", "aws-sigv4", @@ -623,10 +633,11 @@ dependencies = [ [[package]] name = "aws-sdk-sso" -version = "1.98.0" +version = "1.105.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d69c77aafa20460c68b6b3213c84f6423b6e76dbf89accd3e1789a686ffd9489" +checksum = "6ffd0fbe7873cb548a7aa60f9573c268fff94155397fd4f14dc9f1ecaaab8516" dependencies = [ + "arc-swap", "aws-credential-types", "aws-runtime", "aws-smithy-async", @@ -635,6 +646,7 @@ dependencies = [ "aws-smithy-observability", "aws-smithy-runtime", "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", "aws-types", "bytes", @@ -647,10 +659,11 @@ dependencies = [ [[package]] name = "aws-sdk-ssooidc" -version = "1.100.0" +version = "1.107.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1c7e7b09346d5ca22a2a08267555843a6a0127fb20d8964cb6ecfb8fdb190225" +checksum = "175763eb222a46377df7aa257a3bca980ab3e96703fefc8f4d0b8da6ad2e254c" dependencies = [ + "arc-swap", "aws-credential-types", "aws-runtime", "aws-smithy-async", @@ -659,6 +672,7 @@ dependencies = [ "aws-smithy-observability", "aws-smithy-runtime", "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", "aws-types", "bytes", @@ -671,10 +685,11 @@ dependencies = [ [[package]] name = "aws-sdk-sts" -version = "1.103.0" +version = "1.110.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c2249b81a2e73a8027c41c378463a81ec39b8510f184f2caab87de912af0f49b" +checksum = "dd8b14781dfbff48984017d57167b6ea0b6471c6920ec52b44a2677c7feb3c13" dependencies = [ + "arc-swap", "aws-credential-types", "aws-runtime", "aws-smithy-async", @@ -684,6 +699,7 @@ dependencies = [ "aws-smithy-query", "aws-smithy-runtime", "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", "aws-smithy-xml", "aws-types", @@ -696,9 +712,9 @@ dependencies = [ [[package]] name = "aws-sigv4" -version = "1.4.3" +version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68dc0b907359b120170613b5c09ccc61304eac3998ff6274b97d93ee6490115a" +checksum = "723c2234ad7511ceef63eab016b7ba6ff7c55590fefb96fa8467af014a07309f" dependencies = [ "aws-credential-types", "aws-smithy-http", @@ -718,9 +734,9 @@ dependencies = [ [[package]] name = "aws-smithy-async" -version = "1.2.14" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ffcaf626bdda484571968400c326a244598634dc75fd451325a54ad1a59acfc" +checksum = "f02e407fb3b54891734224b9ffac8a71fdd35f542500fa1af95754a6b2beb316" dependencies = [ "futures-util", "pin-project-lite", @@ -729,9 +745,9 @@ dependencies = [ [[package]] name = "aws-smithy-http" -version = "0.63.6" +version = "0.64.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ba1ab2dc1c2c3749ead27180d333c42f11be8b0e934058fb4b2258ee8dbe5231" +checksum = "37843d9add67c3aff5856f409c6dc315d3cdff60f9c0cb5b670dab1e9920306d" dependencies = [ "aws-smithy-runtime-api", "aws-smithy-types", @@ -750,9 +766,9 @@ dependencies = [ [[package]] name = "aws-smithy-http-client" -version = "1.1.12" +version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6a2f165a7feee6f263028b899d0a181987f4fa7179a6411a32a439fba7c5f769" +checksum = "635d23afda0a6ab48d666c4d447c4873e8d1e83518a2be2093122397e50b838e" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api", @@ -774,43 +790,49 @@ dependencies = [ [[package]] name = "aws-smithy-json" -version = "0.62.5" +version = "0.63.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9648b0bb82a2eedd844052c6ad2a1a822d1f8e3adee5fbf668366717e428856a" +checksum = "3dc65a121adb4b33729919fcfa14fa36fb33c1555a8f06bb0e2188dbfdc1d9ef" dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", ] [[package]] name = "aws-smithy-observability" -version = "0.2.6" +version = "0.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a06c2315d173edbf1920da8ba3a7189695827002e4c0fc961973ab1c54abca9c" +checksum = "8e86338c869539a581bf161247762a6e87f92c5c075060057b5ed6d06632ed0c" dependencies = [ "aws-smithy-runtime-api", ] [[package]] name = "aws-smithy-query" -version = "0.60.15" +version = "0.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1a56d79744fb3edb5d722ef79d86081e121d3b9422cb209eb03aea6aa4f21ebd" +checksum = "512346c7212ab7436df2d77a16d976a468ae44a418835511d2a69269810aaf62" dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", + "aws-smithy-xml", "urlencoding", ] [[package]] name = "aws-smithy-runtime" -version = "1.11.1" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0504b1ab12debb5959e5165ee5fe97dd387e7aa7ea6a477bfd7635dfe769a4f5" +checksum = "07505b34e8f4b3591a4fa69e9792b52289b95488dbbc68c3c0075b7bedb245e1" dependencies = [ "aws-smithy-async", "aws-smithy-http", "aws-smithy-http-client", "aws-smithy-observability", "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", "bytes", "fastrand", @@ -827,9 +849,9 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api" -version = "1.12.0" +version = "1.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b71a13df6ada0aafbf21a73bdfcdf9324cfa9df77d96b8446045be3cde61b42e" +checksum = "3b98f2e1fd67ec06618f9c291e5e495a468e60519e44c9c1979cd0521f3affdb" dependencies = [ "aws-smithy-async", "aws-smithy-runtime-api-macros", @@ -845,20 +867,31 @@ dependencies = [ [[package]] name = "aws-smithy-runtime-api-macros" -version = "1.0.0" +version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8d7396fd9500589e62e460e987ecb671bad374934e55ec3b5f498cc7a8a8a7b7" +checksum = "221eaa237ddf1ca79b60d1372aad77e47f9c0ea5b3ce5099da8c61d027dc77b3" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", +] + +[[package]] +name = "aws-smithy-schema" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d56e0a4e53127a632224e43633b0fe045fa9e1e3cfc68b9830f1115e103f910" +dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-types", + "http 1.4.0", ] [[package]] name = "aws-smithy-types" -version = "1.4.7" +version = "1.6.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9d73dbfbaa8e4bc57b9045137680b958d274823509a360abfd8e1d514d40c95c" +checksum = "d6dc683efb34b9e755675b37fedbe0103141e5b6df7bdc9eb6967756a8c167d8" dependencies = [ "base64-simd", "bytes", @@ -879,22 +912,26 @@ dependencies = [ [[package]] name = "aws-smithy-xml" -version = "0.60.15" +version = "0.62.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ce02add1aa3677d022f8adf81dcbe3046a95f17a1b1e8979c145cd21d3d22b3" +checksum = "ce84f71c72fee2cbbadde6e7d082f5fb466e3a84733855295fa7aafd1b31b7d8" dependencies = [ + "aws-smithy-runtime-api", + "aws-smithy-schema", + "aws-smithy-types", "xmlparser", ] [[package]] name = "aws-types" -version = "1.3.15" +version = "1.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2f4bbcaa9304ea40902d3d5f42a0428d1bd895a2b0f6999436fb279ffddc58ac" +checksum = "eec1cd5469f328c782dc3e33d4153cf118a54e33cbb3356d60d16f89883e1f94" dependencies = [ "aws-credential-types", "aws-smithy-async", "aws-smithy-runtime-api", + "aws-smithy-schema", "aws-smithy-types", "rustc_version", "tracing", @@ -955,6 +992,12 @@ version = "0.22.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" +[[package]] +name = "base64" +version = "0.23.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac07cdecf99051d9a5238b80f35af32cdeba5b336e55d957b318b50137e18da5" + [[package]] name = "base64-simd" version = "0.8.0" @@ -973,7 +1016,7 @@ checksum = "4d6867f1565b3aad85681f1015055b087fcfd840d6aeee6eee7f2da317603695" dependencies = [ "autocfg", "libm", - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", ] @@ -995,9 +1038,9 @@ dependencies = [ [[package]] name = "blake3" -version = "1.8.5" +version = "1.8.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +checksum = "76ae7bad254120e9e4c63bafc385310756f90c484eac0e36b8317cf09cb92a77" dependencies = [ "arrayref", "arrayvec", @@ -1120,6 +1163,15 @@ dependencies = [ "alloc-stdlib", ] +[[package]] +name = "bs58" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf88ba1141d185c399bee5288d850d63b8369520c1eafc32a0430b5b6c287bf4" +dependencies = [ + "tinyvec", +] + [[package]] name = "bstr" version = "1.12.1" @@ -1144,9 +1196,9 @@ checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" [[package]] name = "bytes" -version = "1.11.1" +version = "1.12.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" +checksum = "fc652a48c352aef3ea3aed32080501cf3ef6ed5da78602a020c991775b0aff04" [[package]] name = "bytes-utils" @@ -1210,9 +1262,9 @@ dependencies = [ [[package]] name = "chrono" -version = "0.4.44" +version = "0.4.45" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" +checksum = "1aa79e62e7697b8e29b513a68abacf485adcd1fe8284a4316c5ae868e6633327" dependencies = [ "iana-time-zone", "js-sys", @@ -1261,9 +1313,9 @@ dependencies = [ [[package]] name = "clap" -version = "4.6.1" +version = "4.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +checksum = "473c7e07f409a8d772161724aa8db6a765a2532a70f9667eeb7b49d3d02fbdca" dependencies = [ "clap_builder", "clap_derive", @@ -1271,9 +1323,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.6.0" +version = "4.6.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +checksum = "7b48fea5a88e9ae728a2dcbedbfc0e730f7d60da42e1cb049a83c9fb8b789889" dependencies = [ "anstream", "anstyle", @@ -1283,14 +1335,14 @@ dependencies = [ [[package]] name = "clap_derive" -version = "4.6.1" +version = "4.6.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f2ce8604710f6733aa641a2b3731eaa1e8b3d9973d5e3565da11800813f997a9" +checksum = "d012d2b9d65aca7f18f4d9878a045bc17899bba951561ba5ec3c2ba1eed9a061" dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.3", ] [[package]] @@ -1319,9 +1371,9 @@ dependencies = [ [[package]] name = "cmov" -version = "0.5.3" +version = "0.5.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3f88a43d011fc4a6876cb7344703e297c71dda42494fee094d5f7c76bf13f746" +checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" [[package]] name = "colorchoice" @@ -1362,9 +1414,9 @@ checksum = "cc14f565cf027a105f7a44ccf9e5b424348421a1d8952a8fc9d499d313107789" [[package]] name = "console" -version = "0.16.3" +version = "0.16.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" +checksum = "4fe5f465a4f6fee88fad41b85d990f84c835335e85b5d9e6e63e0d06d28cba7c" dependencies = [ "encode_unicode", "libc", @@ -1521,9 +1573,9 @@ dependencies = [ [[package]] name = "crossbeam-epoch" -version = "0.9.18" +version = "0.9.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +checksum = "2d6914041f254d6e9176c01941b21115dcfb7089e55135a35411081bd106ef3f" dependencies = [ "crossbeam-utils", ] @@ -1582,9 +1634,9 @@ dependencies = [ [[package]] name = "ctor" -version = "1.0.5" +version = "1.0.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "378f0974ae2468eaf63aa036dbe9c926b0dc7ea64c156f2ea618bc2f75b934f0" +checksum = "2d83cb7e7a873830708d6b02a78cd36a592c6fa14bf267b68725103b85c0d77f" dependencies = [ "link-section", "linktime-proc-macro", @@ -1625,7 +1677,7 @@ dependencies = [ "proc-macro2", "quote", "strsim", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -1636,14 +1688,14 @@ checksum = "ac3984ec7bd6cfa798e62b4a642426a5be0e68f9401cfc2a01e3fa9ea2fcdb8d" dependencies = [ "darling_core", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "dashmap" -version = "6.1.0" +version = "6.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" dependencies = [ "cfg-if", "crossbeam-utils", @@ -1655,7 +1707,7 @@ dependencies = [ [[package]] name = "datafusion" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-schema", @@ -1700,9 +1752,10 @@ dependencies = [ "flate2", "futures", "glob", + "half", "indexmap 2.14.0", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "liblzma", "log", "nix", @@ -1728,9 +1781,8 @@ dependencies = [ [[package]] name = "datafusion-benchmarks" -version = "53.1.0" +version = "55.0.0" dependencies = [ - "anstream", "arrow", "async-trait", "bytes", @@ -1738,6 +1790,7 @@ dependencies = [ "criterion", "datafusion", "datafusion-common", + "datafusion-common-runtime", "datafusion-proto", "env_logger", "futures", @@ -1759,7 +1812,7 @@ dependencies = [ [[package]] name = "datafusion-catalog" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", @@ -1773,7 +1826,7 @@ dependencies = [ "datafusion-physical-plan", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools 0.15.0", "log", "object_store", "parking_lot", @@ -1782,7 +1835,7 @@ dependencies = [ [[package]] name = "datafusion-catalog-listing" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", @@ -1798,14 +1851,15 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "futures", - "itertools 0.14.0", + "itertools 0.15.0", "log", "object_store", + "percent-encoding", ] [[package]] name = "datafusion-cli" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", @@ -1837,7 +1891,7 @@ dependencies = [ [[package]] name = "datafusion-common" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-ipc", @@ -1850,9 +1904,10 @@ dependencies = [ "hex", "indexmap 2.14.0", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "libc", "log", + "num-traits", "object_store", "parquet", "rand 0.9.4", @@ -1865,7 +1920,7 @@ dependencies = [ [[package]] name = "datafusion-common-runtime" -version = "53.1.0" +version = "55.0.0" dependencies = [ "futures", "log", @@ -1874,7 +1929,7 @@ dependencies = [ [[package]] name = "datafusion-datasource" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-compression", @@ -1887,16 +1942,18 @@ dependencies = [ "datafusion-common-runtime", "datafusion-execution", "datafusion-expr", + "datafusion-functions", "datafusion-physical-expr", "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "flate2", "futures", "glob", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "liblzma", "log", "object_store", @@ -1911,7 +1968,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-arrow" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-ipc", @@ -1925,16 +1982,17 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools 0.15.0", "object_store", "tokio", ] [[package]] name = "datafusion-datasource-avro" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-avro", @@ -1944,6 +2002,7 @@ dependencies = [ "datafusion-datasource", "datafusion-physical-expr-adapter", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1951,7 +2010,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-csv" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", @@ -1963,6 +2022,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1972,7 +2032,7 @@ dependencies = [ [[package]] name = "datafusion-datasource-json" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", @@ -1984,6 +2044,7 @@ dependencies = [ "datafusion-expr", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-session", "futures", "object_store", @@ -1994,9 +2055,10 @@ dependencies = [ [[package]] name = "datafusion-datasource-parquet" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", + "arrow-schema", "async-trait", "bytes", "chrono", @@ -2013,10 +2075,11 @@ dependencies = [ "datafusion-physical-expr-adapter", "datafusion-physical-expr-common", "datafusion-physical-plan", + "datafusion-proto-models", "datafusion-pruning", "datafusion-session", "futures", - "itertools 0.14.0", + "itertools 0.15.0", "log", "object_store", "parking_lot", @@ -2027,17 +2090,17 @@ dependencies = [ [[package]] name = "datafusion-doc" -version = "53.1.0" +version = "55.0.0" [[package]] name = "datafusion-examples" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-flight", "arrow-schema", "async-trait", - "base64 0.22.1", + "base64 0.23.1", "bytes", "dashmap", "datafusion", @@ -2072,11 +2135,12 @@ dependencies = [ [[package]] name = "datafusion-execution" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-buffer", "async-trait", + "bytes", "chrono", "dashmap", "datafusion-common", @@ -2088,14 +2152,17 @@ dependencies = [ "object_store", "parking_lot", "parquet", + "pin-project-lite", "rand 0.9.4", "tempfile", + "tokio", + "tokio-util", "url", ] [[package]] name = "datafusion-expr" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-schema", @@ -2108,10 +2175,12 @@ dependencies = [ "datafusion-functions-aggregate-common", "datafusion-functions-window-common", "datafusion-physical-expr-common", + "datafusion-proto-common", + "datafusion-proto-models", "env_logger", "indexmap 2.14.0", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "recursive", "serde_json", "sqlparser", @@ -2119,18 +2188,18 @@ dependencies = [ [[package]] name = "datafusion-expr-common" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "datafusion-common", "indexmap 2.14.0", "insta", - "itertools 0.14.0", + "itertools 0.15.0", ] [[package]] name = "datafusion-ffi" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-schema", @@ -2167,11 +2236,11 @@ dependencies = [ [[package]] name = "datafusion-functions" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-buffer", - "base64 0.22.1", + "base64 0.23.1", "blake2", "blake3", "chrono", @@ -2187,7 +2256,7 @@ dependencies = [ "datafusion-physical-expr-common", "env_logger", "hex", - "itertools 0.14.0", + "itertools 0.15.0", "log", "md-5 0.11.0", "memchr", @@ -2201,7 +2270,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "criterion", @@ -2213,8 +2282,8 @@ dependencies = [ "datafusion-macros", "datafusion-physical-expr", "datafusion-physical-expr-common", - "foldhash 0.2.0", "half", + "hashbrown 0.17.1", "log", "num-traits", "rand 0.9.4", @@ -2222,7 +2291,7 @@ dependencies = [ [[package]] name = "datafusion-functions-aggregate-common" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "criterion", @@ -2234,7 +2303,7 @@ dependencies = [ [[package]] name = "datafusion-functions-nested" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-ord", @@ -2251,7 +2320,7 @@ dependencies = [ "datafusion-physical-expr", "datafusion-physical-expr-common", "hashbrown 0.17.1", - "itertools 0.14.0", + "itertools 0.15.0", "itoa", "log", "memchr", @@ -2260,7 +2329,7 @@ dependencies = [ [[package]] name = "datafusion-functions-table" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", @@ -2274,7 +2343,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "criterion", @@ -2290,7 +2359,7 @@ dependencies = [ [[package]] name = "datafusion-functions-window-common" -version = "53.1.0" +version = "55.0.0" dependencies = [ "datafusion-common", "datafusion-physical-expr-common", @@ -2298,16 +2367,16 @@ dependencies = [ [[package]] name = "datafusion-macros" -version = "53.1.0" +version = "55.0.0" dependencies = [ "datafusion-doc", "quote", - "syn 2.0.117", + "syn 3.0.3", ] [[package]] name = "datafusion-optimizer" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", @@ -2325,7 +2394,7 @@ dependencies = [ "env_logger", "indexmap 2.14.0", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "log", "recursive", "regex", @@ -2334,7 +2403,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "criterion", @@ -2344,11 +2413,12 @@ dependencies = [ "datafusion-functions", "datafusion-functions-aggregate-common", "datafusion-physical-expr-common", + "datafusion-proto-models", "half", "hashbrown 0.17.1", "indexmap 2.14.0", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "parking_lot", "petgraph", "rand 0.9.4", @@ -2359,7 +2429,7 @@ dependencies = [ [[package]] name = "datafusion-physical-expr-adapter" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "datafusion-common", @@ -2367,21 +2437,22 @@ dependencies = [ "datafusion-functions", "datafusion-physical-expr", "datafusion-physical-expr-common", - "itertools 0.14.0", + "itertools 0.15.0", ] [[package]] name = "datafusion-physical-expr-common" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "chrono", "criterion", "datafusion-common", "datafusion-expr-common", + "datafusion-proto-models", "hashbrown 0.17.1", "indexmap 2.14.0", - "itertools 0.14.0", + "itertools 0.15.0", "parking_lot", "pin-project", "rand 0.9.4", @@ -2389,7 +2460,7 @@ dependencies = [ [[package]] name = "datafusion-physical-optimizer" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "datafusion-common", @@ -2402,15 +2473,16 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-pruning", + "datafusion-session", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "recursive", "tokio", ] [[package]] name = "datafusion-physical-plan" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "arrow-data", @@ -2418,6 +2490,7 @@ dependencies = [ "arrow-ord", "arrow-schema", "async-trait", + "bytes", "criterion", "datafusion-common", "datafusion-common-runtime", @@ -2430,12 +2503,14 @@ dependencies = [ "datafusion-functions-window-common", "datafusion-physical-expr", "datafusion-physical-expr-common", + "datafusion-proto-common", + "datafusion-proto-models", "futures", "half", "hashbrown 0.17.1", "indexmap 2.14.0", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "log", "num-traits", "parking_lot", @@ -2443,16 +2518,16 @@ dependencies = [ "rand 0.9.4", "rstest", "rstest_reuse", + "serde_json", "tokio", ] [[package]] name = "datafusion-proto" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", - "chrono", "datafusion", "datafusion-catalog", "datafusion-catalog-listing", @@ -2473,19 +2548,19 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "datafusion-proto-common", + "datafusion-proto-models", "doc-comment", "object_store", - "pbjson 0.9.0", "pretty_assertions", "prost", - "serde", + "recursive", "serde_json", "tokio", ] [[package]] name = "datafusion-proto-common" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "datafusion-common", @@ -2495,9 +2570,20 @@ dependencies = [ "serde", ] +[[package]] +name = "datafusion-proto-models" +version = "55.0.0" +dependencies = [ + "datafusion-common", + "datafusion-proto-common", + "pbjson 0.9.0", + "prost", + "serde", +] + [[package]] name = "datafusion-pruning" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "datafusion-common", @@ -2509,14 +2595,15 @@ dependencies = [ "datafusion-physical-expr-common", "datafusion-physical-plan", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "log", ] [[package]] name = "datafusion-session" -version = "53.1.0" +version = "55.0.0" dependencies = [ + "arrow-schema", "async-trait", "datafusion-common", "datafusion-execution", @@ -2527,7 +2614,7 @@ dependencies = [ [[package]] name = "datafusion-spark" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "bigdecimal", @@ -2550,13 +2637,14 @@ dependencies = [ "serde_json", "sha1 0.11.0", "sha2", + "tokio", "twox-hash", "url", ] [[package]] name = "datafusion-sql" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "bigdecimal", @@ -2571,17 +2659,18 @@ dependencies = [ "env_logger", "indexmap 2.14.0", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "log", "recursive", "regex", "rstest", "sqlparser", + "stacker", ] [[package]] name = "datafusion-sqllogictest" -version = "53.1.0" +version = "55.0.0" dependencies = [ "arrow", "async-trait", @@ -2596,7 +2685,7 @@ dependencies = [ "futures", "half", "indicatif", - "itertools 0.14.0", + "itertools 0.15.0", "log", "object_store", "postgres-types", @@ -2613,7 +2702,7 @@ dependencies = [ [[package]] name = "datafusion-substrait" -version = "53.1.0" +version = "55.0.0" dependencies = [ "async-recursion", "async-trait", @@ -2622,7 +2711,7 @@ dependencies = [ "datafusion-functions-aggregate", "half", "insta", - "itertools 0.14.0", + "itertools 0.15.0", "object_store", "pbjson-types", "prost", @@ -2634,7 +2723,7 @@ dependencies = [ [[package]] name = "datafusion-wasmtest" -version = "53.1.0" +version = "55.0.0" dependencies = [ "bytes", "chrono", @@ -2733,7 +2822,7 @@ checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -2774,7 +2863,7 @@ dependencies = [ "enum-ordinalize", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -2812,14 +2901,14 @@ checksum = "8ca9601fb2d62598ee17836250842873a413586e5d7ed88b356e38ddbb0ec631" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "env_filter" -version = "1.0.1" +version = "2.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32e90c2accc4b07a8456ea0debdc2e7587bdd890680d71173a15d4ae604f6eef" +checksum = "900d271a03799a1ee8d1ca9b19893b48ca674a9284fefcfb85f05e74ed314217" dependencies = [ "log", "regex", @@ -2827,9 +2916,9 @@ dependencies = [ [[package]] name = "env_logger" -version = "0.11.10" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0621c04f2196ac3f488dd583365b9c09be011a4ab8b9f37248ffcc8f6198b56a" +checksum = "de671bd27a75a797dc9ae289ba1e77276e75e2026408aab65185384e2d5cd3f6" dependencies = [ "anstream", "anstyle", @@ -3069,7 +3158,7 @@ checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -3176,9 +3265,9 @@ dependencies = [ [[package]] name = "glob" -version = "0.3.3" +version = "0.3.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" +checksum = "e4eba85ea1d0a966a983acd07deee566e67395d2d96b6fb39e62b5a833f1eb0b" [[package]] name = "globset" @@ -3649,9 +3738,9 @@ dependencies = [ [[package]] name = "indicatif" -version = "0.18.4" +version = "0.18.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" +checksum = "9433806cd6b4ec1aba79c021c7e4c58fb4c3b9977c085062e611ac929998fb0c" dependencies = [ "console", "portable-atomic", @@ -3662,9 +3751,9 @@ dependencies = [ [[package]] name = "insta" -version = "1.47.2" +version = "1.48.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b4a6248eb93a4401ed2f37dfe8ea592d3cf05b7cf4f8efa867b6895af7e094e" +checksum = "86f0f8fee8c926415c58d6ae43a08523a26faccb2323f5e6b644fe7dd4ef6b82" dependencies = [ "console", "globset", @@ -3672,27 +3761,22 @@ dependencies = [ "regex", "serde", "similar", + "strip-ansi-escapes", "tempfile", "walkdir", ] [[package]] name = "insta-cmd" -version = "0.6.0" +version = "0.7.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ffeeefa927925cced49ccb01bf3e57c9d4cd132df21e576eb9415baeab2d3de6" +checksum = "bffdf4af1db390cf0401535d7c1303cd079a074d28d8473b026fdb6559c41403" dependencies = [ "insta", "serde", "serde_json", ] -[[package]] -name = "integer-encoding" -version = "3.0.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8bb03732005da905c88227371639bf1ad885cc712789c011c31c5fb3ab3ccf02" - [[package]] name = "ipnet" version = "2.12.0" @@ -3733,6 +3817,15 @@ dependencies = [ "either", ] +[[package]] +name = "itertools" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b4baf93f58d4425749ca49a51c50ebab072c5df6994d08fed93541c331481dc" +dependencies = [ + "either", +] + [[package]] name = "itoa" version = "1.0.18" @@ -3760,7 +3853,7 @@ checksum = "e000de030ff8022ea1da3f466fbb0f3a809f5e51ed31f6dd931c35181ad8e6d7" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -3862,9 +3955,9 @@ checksum = "b3a6a8c165077efc8f3a971534c50ea6a1a18b329ef4a66e897a7e3a1494565f" [[package]] name = "libc" -version = "0.2.186" +version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" +checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" [[package]] name = "libloading" @@ -3878,9 +3971,9 @@ dependencies = [ [[package]] name = "liblzma" -version = "0.4.6" +version = "0.4.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b6033b77c21d1f56deeae8014eb9fbe7bdf1765185a6c508b5ca82eeaed7f899" +checksum = "45aec2360b3933207e27908049d8e4df4e476b58180afb1e56b2a4fb72efe4ba" dependencies = [ "liblzma-sys", ] @@ -3904,9 +3997,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" -version = "0.1.47" +version = "0.1.49" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d1eacfa31c33ec25e873c136ba5669f00f9866d0688bea7be4d3f7e43067df6" +checksum = "6a45a52f43e1c16f667ccfe4dd8c85b7f7c204fd5e3bf46c5b0db9a5c3c0b8e9" dependencies = [ "cc", "cty", @@ -3938,15 +4031,15 @@ dependencies = [ [[package]] name = "link-section" -version = "0.16.1" +version = "0.19.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5c8600ca3dbe044f07955b443ff606c50f45295b863289bbe7d0844d50cf11e4" +checksum = "5ee1a0d6e252afe82e7bc2db42fba60e02ddf3b1accaf8cb21d96e34ba61f3d4" [[package]] name = "linktime-proc-macro" -version = "0.1.0" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a44cd706ff0d503ee32b2071166510ca27e281228de10cd3aa8d35ff94560f81" +checksum = "348d0075b1fc163b26d72a7f75fc5141daf2fd1bdf128d873cbaf6785d495bdf" [[package]] name = "linux-raw-sys" @@ -3971,9 +4064,9 @@ dependencies = [ [[package]] name = "log" -version = "0.4.29" +version = "0.4.33" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +checksum = "0ceec5bc11778974d1bcb055b18002eba7f4b3518b6a0081b3af5f21666da9ad" [[package]] name = "lru-slab" @@ -3983,9 +4076,9 @@ checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" [[package]] name = "lz4_flex" -version = "0.13.0" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "db9a0d582c2874f68138a16ce1867e0ffde6c0bb0a0df85e1f36d04146db488a" +checksum = "ecbdfe44b1bd960b68170b417450a628c43f7cf56bb3c5317e61cb230ee7f226" dependencies = [ "twox-hash", ] @@ -4018,15 +4111,15 @@ dependencies = [ [[package]] name = "memchr" -version = "2.8.0" +version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" +checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" [[package]] name = "mimalloc" -version = "0.1.50" +version = "0.1.52" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3627c4272df786b9260cabaa46aec1d59c93ede723d4c3ef646c503816b0640" +checksum = "2d4139bb28d14ad1facf21d5eb8825051b326e172d216b39f6d31df53cc97862" dependencies = [ "libmimalloc-sys", ] @@ -4128,7 +4221,7 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-complex", "num-integer", "num-iter", @@ -4146,6 +4239,16 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93e7820bc0a80a0238e650327316f929ba18d5be054b647490a3a6a339f3e7c0" +dependencies = [ + "num-integer", + "num-traits", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -4187,7 +4290,7 @@ version = "0.4.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824" dependencies = [ - "num-bigint", + "num-bigint 0.4.6", "num-integer", "num-traits", ] @@ -4346,15 +4449,6 @@ version = "0.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" -[[package]] -name = "ordered-float" -version = "2.10.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" -dependencies = [ - "num-traits", -] - [[package]] name = "outref" version = "0.5.2" @@ -4402,9 +4496,9 @@ dependencies = [ [[package]] name = "parquet" -version = "58.3.0" +version = "59.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5dafa7d01085b62a47dd0c1829550a0a36710ea9c4fe358a05a85477cec8a908" +checksum = "7065842956a20c2a536924ce8e4d9955f7422451511b9eb7500d7bfe5077e59c" dependencies = [ "ahash", "arrow-array", @@ -4413,7 +4507,7 @@ dependencies = [ "arrow-ipc", "arrow-schema", "arrow-select", - "base64 0.22.1", + "base64 0.23.1", "brotli", "bytes", "chrono", @@ -4422,16 +4516,14 @@ dependencies = [ "half", "hashbrown 0.17.1", "lz4_flex", - "num-bigint", + "num-bigint 0.5.1", "num-integer", "num-traits", "object_store", - "paste", "ring", "seq-macro", "simdutf8", "snap", - "thrift", "tokio", "twox-hash", "zstd", @@ -4459,15 +4551,9 @@ dependencies = [ "regex", "regex-syntax", "structmeta", - "syn 2.0.117", + "syn 2.0.119", ] -[[package]] -name = "paste" -version = "1.0.15" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" - [[package]] name = "pbjson" version = "0.8.0" @@ -4584,22 +4670,22 @@ dependencies = [ [[package]] name = "pin-project" -version = "1.1.12" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cbf0d9e68100b3a7989b4901972f265cd542e560a3a8a724e1e20322f4d06ce9" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.1.12" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a990e22f43e84855daf260dded30524ef4a9021cc7541c26540500a50b624389" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -4671,21 +4757,21 @@ dependencies = [ [[package]] name = "postgres-derive" -version = "0.4.8" +version = "0.4.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ca1dad89d9ffdbf78502fde418eeede499b87772d88be780478f7f76dc8d471f" +checksum = "4d9d9089bb0ce62f4b5d52a0be0f4acfb35738b979380670d3dea85fe38d6ddd" dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "postgres-protocol" -version = "0.6.11" +version = "0.6.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56201207dac53e2f38e848e31b4b91616a6bb6e0c7205b77718994a7f49e70fc" +checksum = "08808e3c483c46e999108051c78334f473d5adb59d78bb80a1268c7e6aa6c514" dependencies = [ "base64 0.22.1", "byteorder", @@ -4701,9 +4787,9 @@ dependencies = [ [[package]] name = "postgres-types" -version = "0.2.13" +version = "0.2.14" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8dc729a129e682e8d24170cd30ae1aa01b336b096cbb56df6d534ffec133d186" +checksum = "851ca9db4932932d69f3ea811b1abe63087a0f740a47692619dd40d4899b68be" dependencies = [ "bytes", "chrono", @@ -4753,7 +4839,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" dependencies = [ "proc-macro2", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -4776,9 +4862,9 @@ dependencies = [ [[package]] name = "prost" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d2ea70524a2f82d518bce41317d0fae74151505651af45faf1ffbd6fd33f0568" +checksum = "528ac67416ff8646872a3c02cad9cc4ee5dc9f9540c9b10771855c95cb2e5ae1" dependencies = [ "bytes", "prost-derive", @@ -4786,9 +4872,9 @@ dependencies = [ [[package]] name = "prost-build" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "343d3bd7056eda839b03204e68deff7d1b13aba7af2b2fd16890697274262ee7" +checksum = "03da047801ff44bb6a4d407d4860c05fd70bb81714e6b2f3812603d5b145b042" dependencies = [ "heck", "itertools 0.14.0", @@ -4799,28 +4885,28 @@ dependencies = [ "prost", "prost-types", "regex", - "syn 2.0.117", + "syn 2.0.119", "tempfile", ] [[package]] name = "prost-derive" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27c6023962132f4b30eb4c172c91ce92d933da334c59c23cddee82358ddafb0b" +checksum = "b570b25f7617e43d59005d0990ccb79e950a423952cea19671b7a876da390adf" dependencies = [ "anyhow", "itertools 0.14.0", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "prost-types" -version = "0.14.3" +version = "0.14.4" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8991c4cbdb8bc5b11f0b074ffe286c30e523de90fee5ba8132f1399f23cb3dd7" +checksum = "f94967dc7688f3054c7fac87473ffae4cc4c3904800e2d9f5b857246d8963b0a" dependencies = [ "prost", ] @@ -4856,9 +4942,9 @@ dependencies = [ [[package]] name = "quinn" -version = "0.11.9" +version = "0.11.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" +checksum = "0c1a41e437b6bbd489372cd4971de128e85c855f56c57f283d20ff016cf7c0a8" dependencies = [ "bytes", "cfg_aliases", @@ -4876,9 +4962,9 @@ dependencies = [ [[package]] name = "quinn-proto" -version = "0.11.14" +version = "0.11.15" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +checksum = "4fcb935c5bec503c2f0e306bdd3e58bb9029dcb14fa8d9ac76e3a5256ac0763e" dependencies = [ "bytes", "getrandom 0.3.4", @@ -4911,9 +4997,9 @@ dependencies = [ [[package]] name = "quote" -version = "1.0.45" +version = "1.0.47" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +checksum = "1fbf4db142a473a8d80c26bbf18454ed458bf8d26c8219c331daecfdbd079001" dependencies = [ "proc-macro2", ] @@ -5063,7 +5149,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "76009fbe0614077fc1a2ce255e3a1881a2e3a3527097d5dc6d8212c585e7e38b" dependencies = [ "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5112,14 +5198,14 @@ checksum = "b7186006dcb21920990093f30e3dea63b7d6e977bf1256be20c3563a5db070da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "regex" -version = "1.12.3" +version = "1.13.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +checksum = "f020237b6c8eed93db2e2cb53c00c60a8e1bc73da7d073199a1180401450218d" dependencies = [ "aho-corasick", "memchr", @@ -5129,9 +5215,9 @@ dependencies = [ [[package]] name = "regex-automata" -version = "0.4.14" +version = "0.4.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +checksum = "8fcfdb36bda0c880c5931cdc7a2bcdc8ba4556847b9d912bca70bc94708711ad" dependencies = [ "aho-corasick", "memchr", @@ -5146,9 +5232,9 @@ checksum = "cab834c73d247e67f4fae452806d17d3c7501756d98c8808d7c9c7aa7d18f973" [[package]] name = "regex-syntax" -version = "0.8.10" +version = "0.8.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" [[package]] name = "regress" @@ -5247,7 +5333,7 @@ dependencies = [ "regex", "relative-path", "rustc_version", - "syn 2.0.117", + "syn 2.0.119", "unicode-ident", ] @@ -5259,7 +5345,7 @@ checksum = "b3a8fb4672e840a587a66fc577a5491375df51ddb88f2a2c2a792598c326fe14" dependencies = [ "quote", "rand 0.8.6", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5348,9 +5434,9 @@ checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" [[package]] name = "rustyline" -version = "18.0.0" +version = "18.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4a990b25f351b25139ddc7f21ee3f6f56f86d6846b74ac8fad3a719a287cd4a0" +checksum = "53f6a737db68eb1a8ccff86b584b2fc13eca6a7bb6f78ebc7c529547e3ab9684" dependencies = [ "bitflags", "cfg-if", @@ -5436,7 +5522,7 @@ dependencies = [ "proc-macro2", "quote", "serde_derive_internals", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5486,9 +5572,9 @@ checksum = "1bc711410fbe7399f390ca1c3b60ad0f53f80e95c5eb935e52268a0e2cd49acc" [[package]] name = "serde" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", "serde_derive", @@ -5496,22 +5582,22 @@ dependencies = [ [[package]] name = "serde_core" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +checksum = "67dca2c9c51e58a4791a4b1ed58308b39c64224d349a935ab5039aa360942a48" dependencies = [ "serde_derive", ] [[package]] name = "serde_derive" -version = "1.0.228" +version = "1.0.229" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +checksum = "e7a5d71263a5a7d47b41f6b3f06ba276f10cc18b0931f1799f710578e2309348" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.3", ] [[package]] @@ -5522,14 +5608,14 @@ checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "serde_json" -version = "1.0.149" +version = "1.0.151" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +checksum = "c841b55ecdae098c80dcae9cf767f6f8a0c2cdb3416bbef72181df4d0fe73f14" dependencies = [ "indexmap 2.14.0", "itoa", @@ -5547,7 +5633,7 @@ checksum = "175ee3e80ae9982737ca543e96133087cbd9a485eecc3bc4de9c1a37b47ea59c" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5568,7 +5654,7 @@ dependencies = [ "proc-macro2", "quote", "serde", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5585,11 +5671,12 @@ dependencies = [ [[package]] name = "serde_with" -version = "3.18.0" +version = "3.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "dd5414fad8e6907dbdd5bc441a50ae8d6e26151a03b1de04d89a5576de61d01f" +checksum = "76a5c54c7310e7b8b9577c286d7e399ddd876c3e12b3ed917a8aabc4b96e9e8c" dependencies = [ "base64 0.22.1", + "bs58", "chrono", "hex", "indexmap 1.9.3", @@ -5604,14 +5691,14 @@ dependencies = [ [[package]] name = "serde_with_macros" -version = "3.18.0" +version = "3.21.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d3db8978e608f1fe7357e211969fd9abdcae80bac1ba7a3369bb7eb6b404eb65" +checksum = "84d57bc0c8b9a17920c178daa6bb924850d54a9c97ab45194bb8c17ad66bb660" dependencies = [ "darling", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5805,14 +5892,14 @@ checksum = "a6dd45d8fc1c79299bfbb7190e42ccbbdf6a5f52e4a6ad98d92357ea965bd289" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "stabby" -version = "72.1.1" +version = "72.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "976399a0c48ea769ef7f5dc303bb88240ab8d84008647a6b2303eced3dab3945" +checksum = "3d53d2428934c46277fafd2d41e39357595aa1e47954c75db2b14ed90632f3cc" dependencies = [ "rustversion", "stabby-abi", @@ -5820,9 +5907,9 @@ dependencies = [ [[package]] name = "stabby-abi" -version = "72.1.1" +version = "72.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f7b54832a9a1f92a0e55e74a5c0332744426edc515bb3fbad82f10b874a87f0d" +checksum = "f375eae680bb54203ee5e47d4cd2ae7b79c0a79ed90919279f38f500ad53f190" dependencies = [ "rustc_version", "rustversion", @@ -5832,15 +5919,14 @@ dependencies = [ [[package]] name = "stabby-macros" -version = "72.1.1" +version = "72.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a768b1e51e4dbfa4fa52ae5c01241c0a41e2938fdffbb84add0c8238092f9091" +checksum = "ea664671a576c5f7e32fee291ac123d82af5e92b0689beb3555347c00c76eef1" dependencies = [ "proc-macro-crate", "proc-macro2", "quote", - "rand 0.8.6", - "syn 1.0.109", + "syn 2.0.119", ] [[package]] @@ -5851,9 +5937,9 @@ checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" [[package]] name = "stacker" -version = "0.1.24" +version = "0.1.25" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "640c8cdd92b6b12f5bcb1803ca3bbf5ab96e5e6b6b96b9ab77dabe9e880b3190" +checksum = "707f49d46706bacf8a2b00d51dace3f9de527c13eec3778f570c411f89e69967" dependencies = [ "cc", "cfg-if", @@ -5873,6 +5959,15 @@ dependencies = [ "unicode-properties", ] +[[package]] +name = "strip-ansi-escapes" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a8f8038e7e7969abb3f1b7c2a811225e9296da208539e0f79c5251d6cac0025" +dependencies = [ + "vte", +] + [[package]] name = "strsim" version = "0.11.1" @@ -5888,7 +5983,7 @@ dependencies = [ "proc-macro2", "quote", "structmeta-derive", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5899,7 +5994,7 @@ checksum = "152a0b65a590ff6c3da95cabe2353ee04e6167c896b28e3b14478c2636c922fc" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5917,7 +6012,7 @@ dependencies = [ "heck", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -5952,7 +6047,7 @@ dependencies = [ "serde", "serde_json", "serde_yaml", - "syn 2.0.117", + "syn 2.0.119", "typify", "walkdir", ] @@ -5965,9 +6060,9 @@ checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" [[package]] name = "syn" -version = "1.0.109" +version = "2.0.119" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "72b64191b275b66ffe2469e8af2c1cfe3bafa67b529ead792a6d0160888b4237" +checksum = "872831b642d1a07999a962a351ed35b955ea2cfc8f3862091e2a240a84f17297" dependencies = [ "proc-macro2", "quote", @@ -5976,9 +6071,9 @@ dependencies = [ [[package]] name = "syn" -version = "2.0.117" +version = "3.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +checksum = "53e9bae58849f64dfa4f5d5ae372c8341f7305f82a3868709269343628b659a3" dependencies = [ "proc-macro2", "quote", @@ -6002,14 +6097,14 @@ checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "sysinfo" -version = "0.39.2" +version = "0.39.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "14311e7e9a03114cd4b65eedd54e8fed2945e17f08586ae97ef53bc0669f9581" +checksum = "d2071df9448915b71c4fe6d25deaf1c22f12bd234f01540b77312bb8e41361e6" dependencies = [ "libc", "memchr", @@ -6086,22 +6181,22 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.18" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +checksum = "ec86235f5fcc2a73650310756d2ac5b138a5780bbbdfae3eeccec992c435ba4f" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.18" +version = "2.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +checksum = "bc04cd3e1236dd4a98afca4569f2deb3f120e5422a4023be2cb683f8486292af" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 3.0.3", ] [[package]] @@ -6113,17 +6208,6 @@ dependencies = [ "cfg-if", ] -[[package]] -name = "thrift" -version = "0.17.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7e54bc85fc7faa8bc175c4bab5b92ba8d9a3ce893d0e9f42cc455c8ab16a9e09" -dependencies = [ - "byteorder", - "integer-encoding", - "ordered-float", -] - [[package]] name = "time" version = "0.3.47" @@ -6201,9 +6285,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" [[package]] name = "tokio" -version = "1.52.3" +version = "1.53.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" +checksum = "202caea871b69668250d242070849eb495be178ed697a3e98aebce5bc81a0bed" dependencies = [ "bytes", "libc", @@ -6224,14 +6308,14 @@ checksum = "385a6cb71ab9ab790c5fe8d67f1645e6c450a7ce006a33de03daa956cf70a496" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] name = "tokio-postgres" -version = "0.7.17" +version = "0.7.18" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4dd8df5ef180f6364759a6f00f7aadda4fbbac86cdee37480826a6ff9f3574ce" +checksum = "a528f7d280f6d5b9cd149635c8705b0dd049754bc67d81d31fa25169a93809d3" dependencies = [ "async-trait", "byteorder", @@ -6265,9 +6349,9 @@ dependencies = [ [[package]] name = "tokio-stream" -version = "0.1.18" +version = "0.1.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +checksum = "a3d06f0b082ba57c26b79407372e57cf2a1e28124f78e9479fe80322cf53420b" dependencies = [ "futures-core", "pin-project-lite", @@ -6277,22 +6361,23 @@ dependencies = [ [[package]] name = "tokio-util" -version = "0.7.18" +version = "0.7.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +checksum = "494815d09bf52b5548659851081238f0ca39ff638363907596da739561c62c52" dependencies = [ "bytes", "futures-core", "futures-sink", + "libc", "pin-project-lite", "tokio", ] [[package]] name = "toml" -version = "1.1.2+spec-1.1.0" +version = "1.1.3+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "81f3d15e84cbcd896376e6730314d59fb5a87f31e4b038454184435cd57defee" +checksum = "53c96ecdfa941c8fc4fcaed14f99ada8ebed502eef533015095a07e3301d4c3c" dependencies = [ "indexmap 2.14.0", "serde_core", @@ -6335,9 +6420,9 @@ dependencies = [ [[package]] name = "toml_writer" -version = "1.1.1+spec-1.1.0" +version = "1.1.2+spec-1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" +checksum = "7d56353a2a665ad0f41a421187180aab746c8c325620617ad883a99a1cbe66d2" [[package]] name = "tonic" @@ -6447,7 +6532,7 @@ checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -6493,11 +6578,11 @@ checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" [[package]] name = "twox-hash" -version = "2.1.2" +version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" +checksum = "8464ec13c3691491391d9fce00f6416c9a48e46972f72d7865688be2080192c9" dependencies = [ - "rand 0.9.4", + "rand 0.10.1", ] [[package]] @@ -6531,7 +6616,7 @@ dependencies = [ "semver", "serde", "serde_json", - "syn 2.0.117", + "syn 2.0.119", "thiserror", "unicode-ident", ] @@ -6549,7 +6634,7 @@ dependencies = [ "serde", "serde_json", "serde_tokenstream", - "syn 2.0.117", + "syn 2.0.119", "typify-impl", ] @@ -6688,9 +6773,9 @@ checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" [[package]] name = "uuid" -version = "1.23.1" +version = "1.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ddd74a9687298c6858e9b88ec8935ec45d22e8fd5e6394fa1bd4e99a87789c76" +checksum = "bf3923a6f5c4c6382e0b653c4117f48d631ea17f38ed86e2a828e6f7412f5239" dependencies = [ "getrandom 0.4.2", "js-sys", @@ -6715,6 +6800,15 @@ version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5c3082ca00d5a5ef149bb8b555a72ae84c9c59f7250f013ac822ac2e49b19c64" +[[package]] +name = "vte" +version = "0.14.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "231fdcd7ef3037e8330d8e17e61011a2c244126acc0a982f4040ac3f9f0bc077" +dependencies = [ + "memchr", +] + [[package]] name = "walkdir" version = "2.5.0" @@ -6818,7 +6912,7 @@ dependencies = [ "bumpalo", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", "wasm-bindgen-shared", ] @@ -6861,7 +6955,7 @@ checksum = "caf0ca1bd612b988616bac1ab34c4e4290ef18f7148a1d8b7f31c150080e9295" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -7034,7 +7128,7 @@ checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -7045,7 +7139,7 @@ checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -7292,7 +7386,7 @@ dependencies = [ "heck", "indexmap 2.14.0", "prettyplease", - "syn 2.0.117", + "syn 2.0.119", "wasm-metadata", "wit-bindgen-core", "wit-component", @@ -7308,7 +7402,7 @@ dependencies = [ "prettyplease", "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", "wit-bindgen-core", "wit-bindgen-rust", ] @@ -7397,7 +7491,7 @@ checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", "synstructure", ] @@ -7418,7 +7512,7 @@ checksum = "70e3cd084b1788766f53af483dd21f93881ff30d7320490ec3ef7526d203bad4" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] @@ -7438,7 +7532,7 @@ checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", "synstructure", ] @@ -7478,7 +7572,7 @@ checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" dependencies = [ "proc-macro2", "quote", - "syn 2.0.117", + "syn 2.0.119", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index 78c271d524fb8..4526fa0c58934 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -47,9 +47,10 @@ members = [ "datafusion/pruning", "datafusion/physical-plan", "datafusion/proto", - "datafusion/proto/gen", "datafusion/proto-common", "datafusion/proto-common/gen", + "datafusion/proto-models", + "datafusion/proto-models/gen", "datafusion/session", "datafusion/spark", "datafusion/sql", @@ -77,9 +78,9 @@ license = "Apache-2.0" readme = "README.md" repository = "https://github.com/apache/datafusion" # Define Minimum Supported Rust Version (MSRV) -rust-version = "1.88.0" +rust-version = "1.94.0" # Define DataFusion version -version = "53.1.0" +version = "55.0.0" [workspace.dependencies] # We turn off default-features for some dependencies here so the workspaces which inherit them can @@ -88,75 +89,76 @@ version = "53.1.0" # # See for more details: https://github.com/rust-lang/cargo/issues/11329 apache-avro = { version = "0.21", default-features = false } -arrow = { version = "58.3.0", features = [ +arrow = { version = "59.2.0", features = [ "prettyprint", "chrono-tz", ] } -arrow-avro = { version = "58.3.0", default-features = false, features = [ +arrow-avro = { version = "59.2.0", default-features = false, features = [ "deflate", "snappy", "zstd", "bzip2", "xz", ] } -arrow-buffer = { version = "58.3.0", default-features = false } -arrow-data = { version = "58.3.0", default-features = false } -arrow-flight = { version = "58.3.0", features = [ +arrow-buffer = { version = "59.2.0", default-features = false } +arrow-data = { version = "59.2.0", default-features = false } +arrow-flight = { version = "59.2.0", features = [ "flight-sql-experimental", ] } # Both codecs are required here to make sure that code paths like # file-spilling have access to all compression codecs. -arrow-ipc = { version = "58.3.0", default-features = false, features = [ +arrow-ipc = { version = "59.2.0", default-features = false, features = [ "lz4", "zstd", ] } -arrow-ord = { version = "58.3.0", default-features = false } -arrow-schema = { version = "58.3.0", default-features = false } +arrow-ord = { version = "59.2.0", default-features = false } +arrow-schema = { version = "59.2.0", default-features = false } async-trait = "0.1.89" bigdecimal = "0.4.8" bytes = "1.11" bzip2 = "0.6.1" -chrono = { version = "0.4.44", default-features = false } +chrono = { version = "0.4.45", default-features = false } criterion = "0.8" -ctor = "1.0.5" -dashmap = "6.0.1" -datafusion = { path = "datafusion/core", version = "53.1.0", default-features = false } -datafusion-catalog = { path = "datafusion/catalog", version = "53.1.0" } -datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "53.1.0" } -datafusion-common = { path = "datafusion/common", version = "53.1.0", default-features = false } -datafusion-common-runtime = { path = "datafusion/common-runtime", version = "53.1.0" } -datafusion-datasource = { path = "datafusion/datasource", version = "53.1.0", default-features = false } -datafusion-datasource-arrow = { path = "datafusion/datasource-arrow", version = "53.1.0", default-features = false } -datafusion-datasource-avro = { path = "datafusion/datasource-avro", version = "53.1.0", default-features = false } -datafusion-datasource-csv = { path = "datafusion/datasource-csv", version = "53.1.0", default-features = false } -datafusion-datasource-json = { path = "datafusion/datasource-json", version = "53.1.0", default-features = false } -datafusion-datasource-parquet = { path = "datafusion/datasource-parquet", version = "53.1.0", default-features = false } -datafusion-doc = { path = "datafusion/doc", version = "53.1.0" } -datafusion-execution = { path = "datafusion/execution", version = "53.1.0", default-features = false } -datafusion-expr = { path = "datafusion/expr", version = "53.1.0", default-features = false } -datafusion-expr-common = { path = "datafusion/expr-common", version = "53.1.0" } -datafusion-ffi = { path = "datafusion/ffi", version = "53.1.0" } -datafusion-functions = { path = "datafusion/functions", version = "53.1.0" } -datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "53.1.0" } -datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "53.1.0" } -datafusion-functions-nested = { path = "datafusion/functions-nested", version = "53.1.0", default-features = false } -datafusion-functions-table = { path = "datafusion/functions-table", version = "53.1.0" } -datafusion-functions-window = { path = "datafusion/functions-window", version = "53.1.0" } -datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "53.1.0" } -datafusion-macros = { path = "datafusion/macros", version = "53.1.0" } -datafusion-optimizer = { path = "datafusion/optimizer", version = "53.1.0", default-features = false } -datafusion-physical-expr = { path = "datafusion/physical-expr", version = "53.1.0", default-features = false } -datafusion-physical-expr-adapter = { path = "datafusion/physical-expr-adapter", version = "53.1.0", default-features = false } -datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "53.1.0", default-features = false } -datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "53.1.0" } -datafusion-physical-plan = { path = "datafusion/physical-plan", version = "53.1.0" } -datafusion-proto = { path = "datafusion/proto", version = "53.1.0" } -datafusion-proto-common = { path = "datafusion/proto-common", version = "53.1.0" } -datafusion-pruning = { path = "datafusion/pruning", version = "53.1.0" } -datafusion-session = { path = "datafusion/session", version = "53.1.0" } -datafusion-spark = { path = "datafusion/spark", version = "53.1.0" } -datafusion-sql = { path = "datafusion/sql", version = "53.1.0" } -datafusion-substrait = { path = "datafusion/substrait", version = "53.1.0" } +ctor = "1.0.7" +dashmap = "6.2.1" +datafusion = { path = "datafusion/core", version = "55.0.0", default-features = false } +datafusion-catalog = { path = "datafusion/catalog", version = "55.0.0" } +datafusion-catalog-listing = { path = "datafusion/catalog-listing", version = "55.0.0" } +datafusion-common = { path = "datafusion/common", version = "55.0.0", default-features = false } +datafusion-common-runtime = { path = "datafusion/common-runtime", version = "55.0.0" } +datafusion-datasource = { path = "datafusion/datasource", version = "55.0.0", default-features = false } +datafusion-datasource-arrow = { path = "datafusion/datasource-arrow", version = "55.0.0", default-features = false } +datafusion-datasource-avro = { path = "datafusion/datasource-avro", version = "55.0.0", default-features = false } +datafusion-datasource-csv = { path = "datafusion/datasource-csv", version = "55.0.0", default-features = false } +datafusion-datasource-json = { path = "datafusion/datasource-json", version = "55.0.0", default-features = false } +datafusion-datasource-parquet = { path = "datafusion/datasource-parquet", version = "55.0.0", default-features = false } +datafusion-doc = { path = "datafusion/doc", version = "55.0.0" } +datafusion-execution = { path = "datafusion/execution", version = "55.0.0", default-features = false } +datafusion-expr = { path = "datafusion/expr", version = "55.0.0", default-features = false } +datafusion-expr-common = { path = "datafusion/expr-common", version = "55.0.0" } +datafusion-ffi = { path = "datafusion/ffi", version = "55.0.0" } +datafusion-functions = { path = "datafusion/functions", version = "55.0.0" } +datafusion-functions-aggregate = { path = "datafusion/functions-aggregate", version = "55.0.0" } +datafusion-functions-aggregate-common = { path = "datafusion/functions-aggregate-common", version = "55.0.0" } +datafusion-functions-nested = { path = "datafusion/functions-nested", version = "55.0.0", default-features = false } +datafusion-functions-table = { path = "datafusion/functions-table", version = "55.0.0" } +datafusion-functions-window = { path = "datafusion/functions-window", version = "55.0.0" } +datafusion-functions-window-common = { path = "datafusion/functions-window-common", version = "55.0.0" } +datafusion-macros = { path = "datafusion/macros", version = "55.0.0" } +datafusion-optimizer = { path = "datafusion/optimizer", version = "55.0.0", default-features = false } +datafusion-physical-expr = { path = "datafusion/physical-expr", version = "55.0.0", default-features = false } +datafusion-physical-expr-adapter = { path = "datafusion/physical-expr-adapter", version = "55.0.0", default-features = false } +datafusion-physical-expr-common = { path = "datafusion/physical-expr-common", version = "55.0.0", default-features = false } +datafusion-physical-optimizer = { path = "datafusion/physical-optimizer", version = "55.0.0" } +datafusion-physical-plan = { path = "datafusion/physical-plan", version = "55.0.0" } +datafusion-proto = { path = "datafusion/proto", version = "55.0.0", default-features = false } +datafusion-proto-common = { path = "datafusion/proto-common", version = "55.0.0" } +datafusion-proto-models = { path = "datafusion/proto-models", version = "55.0.0" } +datafusion-pruning = { path = "datafusion/pruning", version = "55.0.0" } +datafusion-session = { path = "datafusion/session", version = "55.0.0" } +datafusion-spark = { path = "datafusion/spark", version = "55.0.0" } +datafusion-sql = { path = "datafusion/sql", version = "55.0.0" } +datafusion-substrait = { path = "datafusion/substrait", version = "55.0.0" } doc-comment = "0.3" env_logger = "0.11" @@ -168,22 +170,24 @@ hashbrown = { version = "0.17.1" } hex = { version = "0.4.3" } indexmap = "2.14.0" insta = { version = "1.47.2", features = ["glob", "filters"] } -itertools = "0.14" +itertools = "0.15" itoa = "1.0" liblzma = { version = "0.4.6", features = ["static"] } log = "^0.4" -memchr = "2.8.0" +memchr = "2.8.1" num-traits = { version = "0.2" } object_store = { version = "0.13.2", default-features = false } parking_lot = "0.12" -parquet = { version = "58.3.0", default-features = false, features = [ +parquet = { version = "59.2.0", default-features = false, features = [ "arrow", "async", "object_store", ] } pbjson = { version = "0.9.0" } pbjson-types = "0.9" +percent-encoding = "2.3" pin-project = "1" +pin-project-lite = "^0.2.7" # Should match arrow-flight's version of prost. prost = "0.14.1" rand = "0.9" @@ -193,6 +197,7 @@ rstest = "0.26.1" serde_json = "1" sha2 = "^0.11.0" sqlparser = { version = "0.62.0", default-features = false, features = ["std", "visitor"] } +stacker = "0.1.24" strum = "0.28.0" strum_macros = "0.28.0" tempfile = "3" @@ -204,25 +209,27 @@ url = "2.5.7" uuid = "1.23" zstd = { version = "0.13", default-features = false } +# Keep this list sorted alphabetically. [workspace.lints.clippy] +# https://github.com/apache/datafusion/issues/18881 +allow_attributes = "warn" +assigning_clones = "warn" +inefficient_to_string = "warn" # Detects large stack-allocated futures that may cause stack overflow crashes (see threshold in clippy.toml) large_futures = "warn" -used_underscore_binding = "warn" -or_fun_call = "warn" -unnecessary_lazy_evaluations = "warn" -uninlined_format_args = "warn" -inefficient_to_string = "warn" # https://github.com/apache/datafusion/issues/18503 needless_pass_by_value = "warn" -# https://github.com/apache/datafusion/issues/18881 -allow_attributes = "warn" -assigning_clones = "warn" +or_fun_call = "warn" +uninlined_format_args = "warn" +unnecessary_lazy_evaluations = "warn" +unused_async = "warn" +used_underscore_binding = "warn" [workspace.lints.rust] unexpected_cfgs = { level = "warn", check-cfg = [ 'cfg(datafusion_coop, values("tokio", "tokio_fallback", "per_stream"))', - "cfg(tarpaulin)", - "cfg(tarpaulin_include)", + "cfg(coverage)", + "cfg(coverage_nightly)", ] } unused_qualifications = "deny" diff --git a/README.md b/README.md index 5297b68e2179f..73c4409ef9b54 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ [![Discord chat][discord-badge]][discord-url] [![Linkedin][linkedin-badge]][linkedin-url] ![Crates.io MSRV][msrv-badge] +[![Codecov][codecov-badge]][codecov-url] [crates-badge]: https://img.shields.io/crates/v/datafusion.svg [crates-url]: https://crates.io/crates/datafusion @@ -40,11 +41,13 @@ [commit-activity-badge]: https://img.shields.io/github/commit-activity/m/apache/datafusion [open-issues-badge]: https://img.shields.io/github/issues-raw/apache/datafusion [open-issues-url]: https://github.com/apache/datafusion/issues -[pending-pr-badge]: https://img.shields.io/github/issues-search/apache/datafusion?query=is%3Apr+is%3Aopen+draft%3Afalse+review%3Arequired+status%3Asuccess&label=Pending%20PRs&logo=github -[pending-pr-url]: https://github.com/apache/datafusion/pulls?q=is%3Apr+is%3Aopen+draft%3Afalse+review%3Arequired+status%3Asuccess+sort%3Aupdated-desc +[pending-pr-badge]: https://img.shields.io/github/issues-search/apache/datafusion?query=is%3Apr+is%3Aopen+draft%3Afalse+review%3Arequired&label=Pending%20PRs&logo=github +[pending-pr-url]: https://github.com/apache/datafusion/pulls?q=is%3Apr+is%3Aopen+draft%3Afalse+review%3Arequired+sort%3Aupdated-desc [linkedin-badge]: https://img.shields.io/badge/Follow-Linkedin-blue [linkedin-url]: https://www.linkedin.com/company/apache-datafusion/ [msrv-badge]: https://img.shields.io/crates/msrv/datafusion?label=Min%20Rust%20Version +[codecov-badge]: https://codecov.io/github/apache/datafusion/graph/badge.svg +[codecov-url]: https://app.codecov.io/github/apache/datafusion/tree/main [Website](https://datafusion.apache.org/) | [API Docs](https://docs.rs/datafusion/latest/datafusion/) | @@ -67,6 +70,8 @@ See [use cases] for examples. The following related subprojects target end users queries. - [DataFusion Comet](https://github.com/apache/datafusion-comet/) is an accelerator for Apache Spark based on DataFusion. +- [DataFusion Ballista](https://github.com/apache/datafusion-ballista/) is a distributed query execution engine + that scales DataFusion across a cluster of nodes. "Out of the box," DataFusion offers [SQL](https://datafusion.apache.org/user-guide/sql/index.html) and [DataFrame](https://datafusion.apache.org/user-guide/dataframe.html) APIs, excellent [performance], @@ -106,8 +111,14 @@ It lets you start quickly from a fully working engine, and then customize those Please see the [contributor guide] and [communication] pages for more information. +We discuss our [roadmap] via GitHub issues and invite you +to join the conversation. The current discussion is the +[DataFusion 2026 Q3-Q4 Roadmap Discussion]. + [contributor guide]: https://datafusion.apache.org/contributor-guide [communication]: https://datafusion.apache.org/contributor-guide/communication.html +[roadmap]: https://datafusion.apache.org/contributor-guide/roadmap.html +[datafusion 2026 q3-q4 roadmap discussion]: https://github.com/apache/datafusion/issues/22882 ## Crate features diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 97e0d901b95f9..62eea98439ee1 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -40,14 +40,14 @@ snmalloc = ["snmalloc-rs"] mimalloc_extended = ["libmimalloc-sys/extended"] [dependencies] -anstream = "1.0" arrow = { workspace = true } async-trait = "0.1" bytes = { workspace = true } -clap = { version = "4.6.1", features = ["derive", "env", "color"] } +clap = { version = "4.6.0", features = ["derive", "env", "string"] } criterion = { workspace = true, features = ["html_reports"] } datafusion = { workspace = true, default-features = true } datafusion-common = { workspace = true, default-features = true } +datafusion-common-runtime = { workspace = true } env_logger = { workspace = true } futures = { workspace = true } libmimalloc-sys = { version = "0.1", optional = true } @@ -62,10 +62,10 @@ serde_json = { workspace = true } snmalloc-rs = { version = "0.7", optional = true } tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot"] } tokio-util = { version = "0.7.17" } -toml = "1.1" +toml = "1.1.3" [dev-dependencies] -datafusion-proto = { workspace = true } +datafusion-proto = { workspace = true, features = ["parquet"] } tempfile = { workspace = true } [[bench]] diff --git a/benchmarks/README.md b/benchmarks/README.md index d143de662e47d..f357ff4da58ce 100644 --- a/benchmarks/README.md +++ b/benchmarks/README.md @@ -483,6 +483,14 @@ Your benchmark should create and use an instance of `BenchmarkRun` defined in `b - Call its `start_new_case` method with a string that will appear in the "Query" column of the compare output. - Use `write_iter` to record elapsed times for the behavior you're benchmarking. +- Call `set_memory_pool` with the `RuntimeEnv`'s memory pool (`ctx.runtime_env().memory_pool`), + and again for each new runtime if your benchmark builds one per query. Each case then reports a + `pool_peak_bytes` field: the peak `MemoryPool` reservation reached while running it, which is the + largest value across that case's iterations. The field is omitted when the benchmark runs without + `--memory-limit`, since no pool is installed to record. Comparing it against the peak RSS printed + by `print_memory_stats` shows how much of the run's memory the pool actually accounted for; the + pool only tracks the "large" allocations that scale with input size, so the two are expected to + differ. - When all cases are done, call the `BenchmarkRun`'s `maybe_write_json` method, giving it the value of the `--output` structopt field on `RunOpt`. @@ -496,6 +504,61 @@ The ClickBench[1] benchmarks are widely cited in the industry and focus on grouping / aggregation / filtering. This runner uses the scripts and queries from [2]. +The runner applies two ClickBench-specific setup steps automatically: + +- ClickBench stores `EventDate` as `UInt16` days since `1970-01-01`. + The runner registers the parquet data as `hits_raw`, then creates a + `hits` view that casts `EventDate` through `INTEGER` to `DATE` for the + benchmark queries. +- The source partitioned ClickBench dataset stores string columns without + the `string` Parquet logical type annotation. For partitioned runs, the + runner enables the parquet `binary_as_string` option so those columns + are read as strings. + +If you set up ClickBench manually through SQL, register the single-file +dataset as follows: + +```sql +CREATE EXTERNAL TABLE hits_raw +STORED AS PARQUET +LOCATION 'benchmarks/data/hits.parquet'; +``` + +For the partitioned dataset, register the directory and enable +`binary_as_string`: + +```sql +CREATE EXTERNAL TABLE hits_raw +STORED AS PARQUET +LOCATION 'benchmarks/data/hits_partitioned' +OPTIONS ('binary_as_string' 'true'); +``` + +After registering either dataset as `hits_raw`, create the `hits` view with +the required `EventDate` conversion: + +```sql +CREATE VIEW hits AS +SELECT * EXCEPT ("EventDate"), + CAST(CAST("EventDate" AS INTEGER) AS DATE) AS "EventDate" +FROM hits_raw; +``` + +From the repository root, download data and run the default ClickBench +queries against the single parquet file: + +```shell +./benchmarks/bench.sh data clickbench_1 +./benchmarks/bench.sh run clickbench_1 +``` + +Or run against the partitioned dataset: + +```shell +./benchmarks/bench.sh data clickbench_partitioned +./benchmarks/bench.sh run clickbench_partitioned +``` + [1]: https://github.com/ClickHouse/ClickBench [2]: https://github.com/ClickHouse/ClickBench/tree/main/datafusion @@ -558,7 +621,15 @@ Test performance of end-to-end sort SQL queries. (While the `Sort` benchmark foc Sort integration benchmark runs whole table sort queries on TPCH `lineitem` table, with different characteristics. For example, different number of sort keys, different sort key cardinality, different number of payload columns, etc. -If the TPCH tables have been converted as sorted on their first column (see [Sorted Conversion](#sorted-conversion)), you can use the `--sorted` flag to indicate that the input data is pre-sorted, allowing DataFusion to leverage that order during query execution. +The `--sorted` flag does not sort or rewrite the input files. It declares that the `lineitem` Parquet input is already sorted ascending by its first column (`l_orderkey`). DataFusion can then leverage that ordering during query execution. + +To generate the expected TPC-H SF=1 Parquet input for this benchmark, run: + +```bash +./bench.sh data tpch +``` + +For the `lineitem` table used by `sort-tpch`, this uses `tpchgen-cli` to generate Parquet data that is already ordered by `l_orderkey`. If you use a different input directory, only pass `--sorted` when the `lineitem` files already have that ordering. Additionally, an optional `--limit` flag is available for the sort benchmark. When specified, this flag appends a `LIMIT n` clause to the SQL query, effectively converting the query into a TopK query. Combining the `--sorted` and `--limit` options enables benchmarking of TopK queries on pre-sorted inputs. @@ -578,7 +649,7 @@ See [`sort_tpch.rs`](src/sort_tpch.rs) for more details. cargo run --release --bin dfbench -- sort-tpch -p './datafusion/benchmarks/data/tpch_sf1' -o '/tmp/sort_tpch.json' --query 2 ``` -3. Run all queries as TopK queries on presorted data: +3. Run all queries as TopK queries on already sorted data: ```bash cargo run --release --bin dfbench -- sort-tpch --sorted --limit 10 -p './datafusion/benchmarks/data/tpch_sf1' -o '/tmp/sort_tpch.json' @@ -598,6 +669,14 @@ In addition, topk_tpch is available from the bench.sh script: ./bench.sh run topk_tpch ``` +To benchmark TopK queries on TPC-H `lineitem` input ordered by `l_orderkey`, use: + +```bash +./bench.sh run topk_sorted_tpch +``` + +This runs `dfbench sort-tpch --sorted --limit 100` through the benchmark script, using `--sorted` to declare the existing `l_orderkey` ordering. + ## IMDB Run Join Order Benchmark (JOB) on IMDB dataset. @@ -874,7 +953,7 @@ Several queries are included to test hash joins under various workloads. ## Sort Merge Join -This benchmark focuses on the performance of queries with sort merge joins joins, minimizing other overheads such as scanning data sources or evaluating predicates. +This benchmark focuses on the performance of queries with sort merge joins, minimizing other overheads such as scanning data sources or evaluating predicates. Several queries are included to test sort merge joins under various workloads. diff --git a/benchmarks/bench.sh b/benchmarks/bench.sh index 29957f25e370d..52b78c844a73a 100755 --- a/benchmarks/bench.sh +++ b/benchmarks/bench.sh @@ -99,9 +99,15 @@ tpcds: TPCDS inspired benchmark on Scale Factor (SF) 1 (~1GB), sort_tpch: Benchmark of sorting speed for end-to-end sort queries on TPC-H dataset (SF=1) sort_tpch10: Benchmark of sorting speed for end-to-end sort queries on TPC-H dataset (SF=10) topk_tpch: Benchmark of top-k (sorting with limit) queries on TPC-H dataset (SF=1) +topk_sorted_tpch: Benchmark of top-k queries on TPC-H lineitem ordered by l_orderkey (SF=1) +push_down_topk: Benchmark of ORDER BY ... LIMIT over outer joins on TPC-H dataset (SF=1) — exercises pushing TopK through a join external_aggr: External aggregation benchmark on TPC-H dataset (SF=1) wide_schema: Small-projection queries on a wide synthetic dataset (1024 cols × 256 files) — measures per-file metadata overhead (runs both 'wide' and 'narrow' subgroups: narrow is an internal baseline; the wide-vs-narrow ratio is the signal) +predicate_eval: Conjunctive (AND) filter-evaluation micro-benchmarks; each subgroup is a different predicate pattern, to test how an + adaptive predicate-ordering system behaves across them (see https://github.com/apache/datafusion/issues/11262) + (subgroups via BENCH_SUBGROUP: costsel, cost, selectivity, cardinality, width, scale, neutral, correlation, drift) + (toggle a system under test with its native DATAFUSION_* env var; size data with PRED_ROWS, string width with PRED_FILL) # ClickBench Benchmarks clickbench_1: ClickBench queries against a single parquet file @@ -245,6 +251,10 @@ main() { wide_schema) data_wide_schema ;; + predicate_eval) + # Data is generated inline by the suite's load SQL. + echo "predicate_eval: no external data to generate" + ;; tpcds) data_tpcds ;; @@ -337,7 +347,11 @@ main() { # same data as for tpch10 data_tpch "10" "parquet" ;; - topk_tpch) + topk_tpch|topk_sorted_tpch) + # same data as for tpch + data_tpch "1" "parquet" + ;; + push_down_topk) # same data as for tpch data_tpch "1" "parquet" ;; @@ -458,6 +472,9 @@ main() { wide_schema) run_wide_schema ;; + predicate_eval) + run_predicate_eval + ;; tpcds) run_tpcds ;; @@ -561,6 +578,12 @@ main() { topk_tpch) run_topk_tpch ;; + topk_sorted_tpch) + run_topk_sorted_tpch + ;; + push_down_topk) + run_push_down_topk + ;; nlj) run_nlj ;; @@ -767,17 +790,60 @@ data_wide_schema() { run_wide_schema() { echo "Running wide_schema benchmark (wide subgroup)..." debug_run env BENCH_NAME=wide_schema BENCH_SUBGROUP=wide \ + DATA_DIR="${DATA_DIR}" \ SIMULATE_LATENCY="${SIMULATE_LATENCY}" \ ${QUERY:+BENCH_QUERY="${QUERY}"} \ bash -c "$SQL_CARGO_COMMAND" echo "Running wide_schema benchmark (narrow baseline subgroup)..." debug_run env BENCH_NAME=wide_schema BENCH_SUBGROUP=narrow \ + DATA_DIR="${DATA_DIR}" \ SIMULATE_LATENCY="${SIMULATE_LATENCY}" \ ${QUERY:+BENCH_QUERY="${QUERY}"} \ bash -c "$SQL_CARGO_COMMAND" } +# Runs the push_down_topk benchmark (ORDER BY ... LIMIT over outer joins). +# Reuses the TPC-H parquet data, so it needs `./bench.sh data tpch` (or +# `data push_down_topk`) first. +run_push_down_topk() { + echo "Running push_down_topk benchmark..." + + debug_run env BENCH_NAME=push_down_topk \ + BENCH_SIZE="1" \ + DATA_DIR="${DATA_DIR}" \ + SIMULATE_LATENCY="${SIMULATE_LATENCY}" \ + ${QUERY:+BENCH_QUERY="${QUERY}"} \ + bash -c "$SQL_CARGO_COMMAND" +} + +# Runs the predicate_eval benchmark suite: conjunctive (AND) filter-evaluation +# micro-benchmarks where each subgroup is a different predicate pattern, used to +# test how an adaptive predicate-ordering system behaves across them (see +# https://github.com/apache/datafusion/issues/11262). Data is generated inline +# by the suite's load SQL, so there is no data step. +# +# By default the suite measures DataFusion's built-in left-deep AND short-circuit +# and sets no engine config of its own. To evaluate a system under test, export +# its native DATAFUSION_* config before invoking bench.sh -- the harness reads +# SessionConfig::from_env, and that environment is inherited here, e.g. +# DATAFUSION_EXECUTION_ADAPTIVE_FILTER_REORDERING=true ./bench.sh run predicate_eval +# Suite-specific knobs (string-substituted into the load SQL, not engine config): +# BENCH_SUBGROUP run one subgroup (costsel, cost, selectivity, cardinality, +# width, scale, neutral, correlation, drift) +# PRED_ROWS synthetic row count (default 1_000_000; the scale subgroup +# overrides this per query) +# PRED_FILL filler chars per marker = string-column width knob +run_predicate_eval() { + echo "Running predicate_eval benchmark (subgroup=${BENCH_SUBGROUP:-all}, rows=${PRED_ROWS:-1000000})..." + debug_run env BENCH_NAME=predicate_eval \ + ${BENCH_SUBGROUP:+BENCH_SUBGROUP="${BENCH_SUBGROUP}"} \ + PRED_ROWS="${PRED_ROWS:-1000000}" \ + ${PRED_FILL:+PRED_FILL="${PRED_FILL}"} \ + ${QUERY:+BENCH_QUERY="${QUERY}"} \ + bash -c "$SQL_CARGO_COMMAND" +} + # Runs the tpch in memory (needs tpch parquet data) run_tpch_mem() { SCALE_FACTOR=$1 @@ -1444,6 +1510,16 @@ run_topk_tpch() { $CARGO_COMMAND --bin dfbench -- sort-tpch --iterations 5 --path "${TPCH_DIR}" -o "${RESULTS_FILE}" --limit 100 ${QUERY_ARG} ${LATENCY_ARG} } +# Runs the sorted sort tpch integration benchmark with limit 100 (topk) +run_topk_sorted_tpch() { + TPCH_DIR="${DATA_DIR}/tpch_sf1" + RESULTS_FILE="${RESULTS_DIR}/run_topk_sorted_tpch.json" + echo "RESULTS_FILE: ${RESULTS_FILE}" + echo "Running sorted topk tpch benchmark..." + + $CARGO_COMMAND --bin dfbench -- sort-tpch --iterations 5 --path "${TPCH_DIR}" -o "${RESULTS_FILE}" --sorted --limit 100 ${QUERY_ARG} ${LATENCY_ARG} +} + # Runs the nlj benchmark run_nlj() { RESULTS_FILE="${RESULTS_DIR}/nlj.json" diff --git a/benchmarks/benches/sql.rs b/benchmarks/benches/sql.rs index 73302b4763818..9240a19470db9 100644 --- a/benchmarks/benches/sql.rs +++ b/benchmarks/benches/sql.rs @@ -22,26 +22,13 @@ //! Cargo, for example: `BENCH_NAME=tpch cargo bench --bench sql`. use clap::Parser; -use criterion::{Criterion, SamplingMode, criterion_group, criterion_main}; -use datafusion::error::Result; -use datafusion::prelude::SessionContext; -use datafusion_benchmarks::sql_benchmark::SqlBenchmark; -use datafusion_benchmarks::util::{CommonOpt, print_memory_stats}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_benchmarks::sql_benchmark_runner::{ + BenchmarkFilter, SqlRunConfig, default_criterion_replacements, + default_sql_benchmark_directory, run_criterion_benchmarks_impl, +}; +use datafusion_benchmarks::util::CommonOpt; use datafusion_common::instant::Instant; -use log::{debug, info}; -use std::collections::BTreeMap; -use std::fs; -use std::sync::LazyLock; -use tokio::runtime::Runtime; - -static SQL_BENCHMARK_DIRECTORY: LazyLock = LazyLock::new(|| { - format!( - "{}{}{}", - env!("CARGO_MANIFEST_DIR"), - std::path::MAIN_SEPARATOR, - "sql_benchmarks" - ) -}); #[cfg(feature = "snmalloc")] #[global_allocator] @@ -82,7 +69,7 @@ struct EnvParser { subgroup: Option, #[arg(env = "BENCH_QUERY")] - query: Option, + query: Option, } pub fn sql(c: &mut Criterion) { @@ -90,229 +77,30 @@ pub fn sql(c: &mut Criterion) { let start = Instant::now(); let args = EnvParser::parse(); - let rt = make_tokio_runtime(); + let config = SqlRunConfig { + common: args.options, + filter: BenchmarkFilter { + name: args.name, + subgroup: args.subgroup, + query: args.query, + }, + replacements: default_criterion_replacements(), + query_filename: None, + persist_results: args.persist_results, + validate_results: args.validate, + output: None, + }; println!("Loading benchmarks..."); - let benchmarks = rt.block_on(async { - let ctx = make_ctx(&args).expect("SessionContext creation failed"); - - load_benchmarks(&args, &ctx, &SQL_BENCHMARK_DIRECTORY) - .await - .unwrap_or_else(|err| panic!("failed load benchmarks: {err:?}")) - }); + run_criterion_benchmarks_impl(&default_sql_benchmark_directory(), &config, c) + .unwrap_or_else(|err| panic!("failed to run SQL benchmarks: {err:?}")); println!( - "Loaded benchmarks in {} ms ...", + "Completed benchmarks in {} ms ...", start.elapsed().as_millis() ); - - for (group, benchmarks) in benchmarks { - let mut group = c.benchmark_group(group); - group.sample_size(10); - group.sampling_mode(SamplingMode::Flat); - - for mut benchmark in benchmarks { - // create a context - let ctx = make_ctx(&args).expect("SessionContext creation failed"); - - // initialize the benchmark. This parses the benchmark file and does any pre-execution - // work such as loading data into tables - rt.block_on(async { - benchmark - .initialize(&ctx) - .await - .expect("initialization failed"); - - // run assertions - benchmark.assert(&ctx).await.expect("assertion failed"); - }); - - let mut name = benchmark.name().to_string(); - if !benchmark.subgroup().is_empty() { - name.push('_'); - name.push_str(benchmark.subgroup()); - } - - if args.persist_results { - handle_persist(&rt, &ctx, &name, &mut benchmark); - } else if args.validate { - handle_verify(&rt, &ctx, &name, &mut benchmark); - } else { - info!("Running benchmark {name} ..."); - - let name = name.clone(); - group.bench_function(name.clone(), |b| { - b.iter(|| handle_run(&rt, &ctx, &args, &mut benchmark, &name)) - }); - - print_memory_stats(); - - info!("Benchmark {name} completed"); - } - - // run cleanup - rt.block_on(async { - benchmark.cleanup(&ctx).await.expect("Cleanup failed"); - }); - } - - group.finish(); - } -} - -fn handle_run( - rt: &Runtime, - ctx: &SessionContext, - args: &EnvParser, - benchmark: &mut SqlBenchmark, - name: &str, -) { - rt.block_on(async { - benchmark - .run(ctx, args.validate) - .await - .unwrap_or_else(|err| panic!("Failed to run benchmark {name}: {err:?}")) - }); -} - -fn handle_persist( - rt: &Runtime, - ctx: &SessionContext, - name: &str, - benchmark: &mut SqlBenchmark, -) { - info!("Running benchmark {name} prior to persisting results ..."); - - rt.block_on(async { - info!("Persisting benchmark {name} ..."); - - benchmark - .persist(ctx) - .await - .expect("Failed to persist results"); - }); - - info!("Persisted benchmark {name} successfully"); -} - -fn handle_verify( - rt: &Runtime, - ctx: &SessionContext, - name: &str, - benchmark: &mut SqlBenchmark, -) { - info!("Verifying benchmark {name} results ..."); - - rt.block_on(async { - benchmark - .run(ctx, true) - .await - .unwrap_or_else(|err| panic!("Failed to run benchmark {name}: {err:?}")); - benchmark - .verify(ctx) - .await - .unwrap_or_else(|err| panic!("Verification failed: {err:?}")); - }); - - info!("Verified benchmark {name} results successfully"); } criterion_group!(benches, sql); criterion_main!(benches); - -fn make_tokio_runtime() -> Runtime { - tokio::runtime::Builder::new_multi_thread() - .enable_all() - .build() - .unwrap() -} - -fn make_ctx(args: &EnvParser) -> Result { - let config = args.options.config()?; - let rt = args.options.build_runtime()?; - - Ok(SessionContext::new_with_config_rt(config, rt)) -} - -/// Recursively walks the directory tree starting at `path` and -/// calls the call back function for every file encountered. -pub fn list_files(path: &str, callback: &mut F) -where - F: FnMut(&str), -{ - let mut entries: Vec = - fs::read_dir(path).unwrap().filter_map(Result::ok).collect(); - entries.sort_by_key(|entry| entry.path()); - - for dir_entry in entries { - let path = dir_entry.path(); - if path.is_dir() { - // Recurse into the sub‑directory - list_files(&path.to_string_lossy(), callback); - } else { - // For files, invoke the callback with the full path as a string - let full_str = path.to_string_lossy(); - callback(&full_str); - } - } -} - -/// Loads all benchmark files in the `sql_benchmarks` directory. -/// For each file ending with `.benchmark` it creates a new -/// `SqlBenchmark` instance. -async fn load_benchmarks( - args: &EnvParser, - ctx: &SessionContext, - path: &str, -) -> Result>> { - let mut benches = BTreeMap::new(); - let mut paths = Vec::new(); - - list_files(path, &mut |path: &str| { - if path.ends_with(".benchmark") { - paths.push(path.to_string()); - } - }); - - for path in paths { - debug!("Loading benchmark from {path}"); - - let benchmark = SqlBenchmark::new(ctx, &path, &*SQL_BENCHMARK_DIRECTORY).await?; - let entries = benches - .entry(benchmark.group().to_string()) - .or_insert(vec![]); - - entries.push(benchmark); - } - - benches = filter_benchmarks(args, benches); - benches.iter_mut().for_each(|(_, benchmarks)| { - benchmarks.sort_by(|b1, b2| b1.name().cmp(b2.name())) - }); - - Ok(benches) -} - -fn filter_benchmarks( - args: &EnvParser, - benchmarks: BTreeMap>, -) -> BTreeMap> { - match &args.name { - Some(bench_name) => benchmarks - .into_iter() - .filter(|(key, _val)| key.eq_ignore_ascii_case(bench_name)) - .map(|(key, mut val)| { - if let Some(subgroup) = &args.subgroup { - val.retain(|bench| bench.subgroup().eq_ignore_ascii_case(subgroup)); - } - if let Some(query_number) = &args.query { - let padded = format!("Q{query_number:0>2}"); - val.retain(|bench| bench.name().eq_ignore_ascii_case(&padded)); - } - (key, val) - }) - .collect(), - None => benchmarks, - } -} diff --git a/benchmarks/queries/clickbench/queries/q27.sql b/benchmarks/queries/clickbench/queries/q27.sql index ba234d34f8877..dbd6aeaf8128a 100644 --- a/benchmarks/queries/clickbench/queries/q27.sql +++ b/benchmarks/queries/clickbench/queries/q27.sql @@ -1,4 +1,5 @@ -- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591 -- set datafusion.execution.parquet.binary_as_string = true -SELECT "CounterID", AVG(length("URL")) AS l, COUNT(*) AS c FROM hits WHERE "URL" <> '' GROUP BY "CounterID" HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; +-- DataFusion length(...) counts characters; use octet_length(...) for ClickBench byte-length semantics. +SELECT "CounterID", AVG(octet_length("URL")) AS l, COUNT(*) AS c FROM hits WHERE "URL" <> '' GROUP BY "CounterID" HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; diff --git a/benchmarks/queries/clickbench/queries/q28.sql b/benchmarks/queries/clickbench/queries/q28.sql index 6a3bd037bece7..6d00194b74929 100644 --- a/benchmarks/queries/clickbench/queries/q28.sql +++ b/benchmarks/queries/clickbench/queries/q28.sql @@ -1,4 +1,5 @@ -- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591 -- set datafusion.execution.parquet.binary_as_string = true -SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; +-- DataFusion length(...) counts characters; use octet_length(...) for ClickBench byte-length semantics. +SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(octet_length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; diff --git a/benchmarks/queries/h2o/window.sql b/benchmarks/queries/h2o/window.sql index fa16a3de32ca5..ece2c75abd205 100644 --- a/benchmarks/queries/h2o/window.sql +++ b/benchmarks/queries/h2o/window.sql @@ -117,3 +117,132 @@ SELECT id2, largest2_v2 FROM ( ROW_NUMBER() OVER (PARTITION BY id2 ORDER BY v2 DESC) AS order_v2 FROM large WHERE v2 IS NOT NULL ) sub_query WHERE order_v2 <= 2; + +-- Window Top-N partition cardinality sweep (id3 % N gives N distinct partitions). +-- These exercise PartitionedTopKExec across cardinalities to validate it stays +-- competitive with the SortExec+Filter baseline as partition count grows. +-- Window Top-N: 100 partitions +SELECT pk, largest2_v2 FROM ( + SELECT id3 % 100 AS pk, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id3 % 100 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; + +-- Window Top-N: 1,000 partitions +SELECT pk, largest2_v2 FROM ( + SELECT id3 % 1000 AS pk, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id3 % 1000 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; + +-- Window Top-N: 10,000 partitions +SELECT pk, largest2_v2 FROM ( + SELECT id3 % 10000 AS pk, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id3 % 10000 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; + +-- Window Top-N: 100,000 partitions +SELECT pk, largest2_v2 FROM ( + SELECT id3 % 100000 AS pk, v2 AS largest2_v2, + ROW_NUMBER() OVER (PARTITION BY id3 % 100000 ORDER BY v2 DESC) AS order_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE order_v2 <= 2; + +-- Window Top-N (RANK top-2 per partition, ~100 partitions) +-- The RANK queries below mirror the ROW_NUMBER cardinality sweep +-- above and add heavy-ties variants. RANK semantics retain boundary +-- ties (`WHERE rk <= K` may keep more than K rows per partition), so +-- this exercises PartitionedTopKRank's ties-Vec path. +SELECT pk, largest_v2 FROM ( + SELECT (id3 % 100) AS pk, v2 AS largest_v2, + RANK() OVER (PARTITION BY (id3 % 100) ORDER BY v2 DESC) AS rk_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE rk_v2 <= 2; + +-- Window Top-N (RANK top-2 per partition, ~1K partitions) +SELECT pkey, largest_v2 FROM ( + SELECT (id3 % 1000) AS pkey, v2 AS largest_v2, + RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY v2 DESC) AS rk_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE rk_v2 <= 2; + +-- Window Top-N (RANK top-2 per partition, ~1K partitions, heavy ties) +-- v2 % 10 forces 10 distinct OBY values, so most rows tie at the boundary +-- and exercise PartitionedTopKRank's ties-Vec path. +SELECT pkey, largest_v2 FROM ( + SELECT (id3 % 1000) AS pkey, v2 AS largest_v2, + RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY (v2 % 10) DESC) AS rk_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE rk_v2 <= 2; + +-- Window Top-N (RANK top-2 per partition, ~10K partitions, low ties) +SELECT id2, largest_v2 FROM ( + SELECT id2, v2 AS largest_v2, + RANK() OVER (PARTITION BY id2 ORDER BY v2 DESC) AS rk_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE rk_v2 <= 2; + +-- Window Top-N (RANK top-2 per partition, ~10K partitions, heavy ties) +SELECT id2, largest_v2 FROM ( + SELECT id2, v2 AS largest_v2, + RANK() OVER (PARTITION BY id2 ORDER BY (v2 % 10) DESC) AS rk_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE rk_v2 <= 2; + +-- Window Top-N (RANK top-2 per partition, ~100K partitions) +SELECT pk, largest_v2 FROM ( + SELECT (id3 % 100000) AS pk, v2 AS largest_v2, + RANK() OVER (PARTITION BY (id3 % 100000) ORDER BY v2 DESC) AS rk_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE rk_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~100 partitions) +-- The DENSE_RANK queries below mirror the RANK cardinality sweep above. +-- DENSE_RANK semantics keep every row whose ORDER BY value is among the +-- K distinct-greatest values in the partition, so total kept per partition +-- is unbounded in rows-per-distinct-value — exercises PartitionedTopKDenseRank's +-- HashMap-of-groups path. +SELECT pk, largest_v2 FROM ( + SELECT (id3 % 100) AS pk, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY (id3 % 100) ORDER BY v2 DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~1K partitions) +SELECT pkey, largest_v2 FROM ( + SELECT (id3 % 1000) AS pkey, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY v2 DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~1K partitions, heavy ties) +-- v2 % 10 forces 10 distinct OBY values; most rows share the top-2 distinct +-- values so appends dominate — exercises the "Case A" append-to-existing-Vec +-- fast path in PartitionedTopKDenseRank. +SELECT pkey, largest_v2 FROM ( + SELECT (id3 % 1000) AS pkey, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY (id3 % 1000) ORDER BY (v2 % 10) DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~10K partitions, low ties) +SELECT id2, largest_v2 FROM ( + SELECT id2, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY id2 ORDER BY v2 DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~10K partitions, heavy ties) +SELECT id2, largest_v2 FROM ( + SELECT id2, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY id2 ORDER BY (v2 % 10) DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; + +-- Window Top-N (DENSE_RANK top-2 per partition, ~100K partitions) +SELECT pk, largest_v2 FROM ( + SELECT (id3 % 100000) AS pk, v2 AS largest_v2, + DENSE_RANK() OVER (PARTITION BY (id3 % 100000) ORDER BY v2 DESC) AS dr_v2 + FROM large WHERE v2 IS NOT NULL +) sub_query WHERE dr_v2 <= 2; diff --git a/benchmarks/sql_benchmarks/README.md b/benchmarks/sql_benchmarks/README.md index 1705cf0d2f58b..dfb09e0a3a4a2 100644 --- a/benchmarks/sql_benchmarks/README.md +++ b/benchmarks/sql_benchmarks/README.md @@ -36,29 +36,103 @@ in the community: | `hj` | Hash join benchmark | | `imdb` | IMDb benchmark | | `nlj` | Nested‑loop join benchmark | +| `push_down_topk` | `ORDER BY ... LIMIT` over outer joins (TPC-H data); exercises pushing a TopK through a join | | `smj` | Sort‑merge join benchmark | | `sort tpch` | Sorting benchmarks against the TPC-H lineitem table | | `taxi` | NYC taxi dataset benchmark | | `tpcds` | TPC‑DS queries | | `tpch` | TPC‑H queries | | `wide_schema` | Small-projection queries on a wide (1024-col, 256-file) synthetic dataset; runs `wide` + `narrow` subgroups for comparison | +| `predicate_eval` | Conjunctive (AND) filter-evaluation micro-benchmarks; each subgroup is a different predicate pattern, to test how an adaptive predicate-ordering system behaves across them ([#11262](https://github.com/apache/datafusion/issues/11262)). Subgroups (`--subgroup`): `costsel`, `cost`, `selectivity`, `cardinality`, `width`, `scale`, `neutral`, `correlation`, `drift`. Configure the system under test through its DataFusion settings. | # Running Benchmarks -The easiest way to run a benchmark is to use the `bench.sh` shell script (up one level from this document) -as it takes care of configuring any required environment variables and can populate any required data files. -However, it is possible to directly run a sql benchmark using the `cargo bench` command. For example: +Use `benchmark_runner` to run SQL benchmarks. It reads each suite's `.suite` +file and exposes the suite's configuration as command-line options. Use the +`bench.sh` shell script one level above this directory to download or generate +required data files. ```shell -BENCH_NAME=tpch cargo bench --bench sql +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch ``` +## SQL benchmark runner + +The `benchmark_runner` binary discovers suites from this directory and exposes +suite-specific options alongside the common benchmark options. The suite name +must come before all options. + +```bash +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- --list +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch --help +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch --query 15 --format csv +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- clickbench --partitioning partitioned --dry-run +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch --query 1 --result-mode persist +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- tpch --query 1 --result-mode validate +``` + +Use `--path PATH` or `-p PATH` to override `DATA_DIR` for a suite that declares +that path replacement. Suites without a `DATA_DIR` replacement reject the +option. Suite-specific values follow this precedence: command-line option, +environment variable, then the default in the suite metadata. + +`--dry-run` prints the resolved suite, filters, run mode, common options, +suite-specific values, and path replacements as JSON. It reports source metadata +for suite options and path replacements. It validates the command but does not +load benchmark definitions, create a session, read datasets, execute SQL, or +write benchmark results. + +Use `--result-mode persist` to save query results or `--result-mode validate` to +compare them with saved results. The default, `--result-mode none`, does neither. +For compatibility with direct Criterion runs, the runner also reads +`BENCH_PERSIST_RESULTS` and `BENCH_VALIDATE`. Persistence takes precedence when +both variables are `true`. An explicit `--result-mode` overrides both variables. + +### Suite metadata + +Each discoverable suite has one TOML metadata file named +`/.suite`. The runner accepts these top-level fields: + +| Field | Required | Description | +|-------|----------|-------------| +| `description` | Yes | Non-empty text shown by `--list` and suite help. | +| `query_pattern` | No | Relative benchmark filename pattern. It must contain exactly one `{QUERY_ID}` or `{QUERY_ID_PADDED}` placeholder and defaults to `q{QUERY_ID_PADDED}.benchmark`. | +| `path_replacements` | No | Map of replacement names to paths. Relative paths resolve from the suite directory. `DATA_DIR` enables `--path/-p`. | +| `options` | No | Array of suite-specific option tables described below. | +| `examples` | No | Array of `command` and `description` pairs appended to suite help. Both values must contain text. | + +Each `[[options]]` table has these fields: + +| Field | Required | Description | +|-------|----------|-------------| +| `name` | Yes | Long option name without `--`; use lowercase ASCII letters, digits, and hyphens. | +| `short` | No | One ASCII letter or digit without `-`. | +| `env` | Yes | Environment variable that supplies the option value. | +| `default` | Yes | Value used when neither the command line nor the environment supplies one. | +| `values` | No | Accepted values. Omit the field to allow any value. Include `"..."` to allow the listed values plus any other value. Without `"..."`, the list is closed. | +| `help` | Yes | Non-empty text shown in suite help. | + +Option names, short names, and environment keys must be unique within a suite. +An option environment key cannot also appear in `path_replacements`. Suite +options cannot reuse the runner's global names: `help` (`-h`), `query` (`-q`), +`subgroup`, `iterations` (`-i`), `partitions` (`-n`), `batch-size` (`-s`), +`mem-pool-type`, `memory-limit`, `sort-spill-reservation-bytes`, `debug` (`-d`), +`simulate-latency`, `criterion`, `list`, `output` (`-o`), `save-baseline`, +`path` (`-p`), `result-mode`, or `dry-run`. + # Benchmark configuration -Sql benchmarks are configured via environment variables. Cargo's bench command and -[criterion](https://github.com/criterion-rs/criterion.rs) (the underlying benchmark framework) have an unfortunate -limitation in that custom command arguments cannot be passed into a benchmark. The alternative is to use environment -variables to pass in arguments which is what is used here. +`benchmark_runner` is the preferred interface for configuring and running SQL +benchmarks. Run ` --help` to see the common and suite-specific options: + +```shell +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- h2o --help +``` + +The runner maps suite options to the environment variables below for +compatibility with benchmark files and direct Criterion runs. Direct +`cargo bench --bench sql` invocations cannot accept custom arguments, so they +still use environment variables. The SQL benchmarking tool uses the following environment variables: @@ -74,10 +148,10 @@ The SQL benchmarking tool uses the following environment variables: | MEM_POOL_TYPE | The memory pool type to use, should be one of "fair" or "greedy". | | MEMORY_LIMIT | Memory limit (e.g. '100M', '1.5G'). If not specified, run all pre-defined memory limits for given query if there's any, otherwise run with no memory limit. | -Example – Run the H2O window benchmarks on the 'small' sized CSV data files: +Example: run the H2O window benchmarks on the small CSV data files: -``` bash -BENCH_NAME=h2o BENCH_SUBGROUP=window H2O_BENCH_SIZE=small H20_FILE_TYPE=csv cargo bench --bench sql +```shell +cargo run -p datafusion-benchmarks --release --bin benchmark_runner -- h2o --subgroup window --size small --format csv ``` Some benchmarks use custom environment variables as outlined below: @@ -94,24 +168,27 @@ Some benchmarks use custom environment variables as outlined below: | BENCH_SORTED | Used in the sort_tpch benchmark to indicate whether the lineitem table should be sorted. | false | | SORTED_BY | Used in the clickbench_sorted benchmark to indicate the column to sort by. | `EventTime` | | SORTED_ORDER | Used in the clickbench_sorted benchmark to indicate the sort order of the column. | `ASC` | +| PRED_ROWS | Used in the predicate_eval benchmark to size the synthetic table (the `scale` subgroup overrides this per query). | `1000000` | +| PRED_FILL | Used in the predicate_eval benchmark as the string-column width knob (filler chars per marker). | `30` | ## How it works -SQL benchmarks are run via cargo's bench command using [criterion](https://docs.rs/criterion/latest/criterion/) -for running and gathering statistics of each sql being benchmarked. +The runner executes SQL benchmarks with its basic runner by default. Pass +`--criterion` to gather statistics with +[Criterion](https://docs.rs/criterion/latest/criterion/). Each individual benchmark is represented by a `.benchmark` file that contains a number of directives instructing the tool on how to load data, run initializations, run assertions, run the benchmark, optionally persist and validate results, and finally run any cleanup if required. -Variables are supported in two forms: +Benchmark files support replacement variables in two forms: -* string substitution based on environment variables (with default values if unset): \${ENV_VAR} and +* string substitution with an optional default: \${ENV_VAR} and \${ENV_VAR:-default}. -* if / else based on whether an environment variable is true or not +* if / else based on whether a replacement value is true or not (\${ENV_VAR:-default|true value|false value}). In this form only the value `true` (case-insensitive) selects the - true branch; any other set value selects the false branch. If ENV_VAR is unset, the valud of `default` is used to -* select the branch. + true branch; any other supplied value selects the false branch. If the value is absent, the parser uses `default` to + select the branch. Comments in files are supported with lines starting with # or --. @@ -153,8 +230,8 @@ The above showcases the use of defaults for variables: `${NAME:-default}` The name of the benchmark. This will be used as part of the display name used by criterion.

Example:
name Q${QUERY_NUMBER_PADDED}
-The `name` directive also makes the value available to benchmark-file replacements as `BENCH_NAME`. This is separate -from the `BENCH_NAME` environment variable used to select which benchmark group to run. +The `name` directive also makes the value available to benchmark-file replacements as `BENCH_NAME`. This value is +separate from the suite name passed to `benchmark_runner`. @@ -217,8 +294,8 @@ The run directive called during execution of the benchmark. If a path to a file the run directive that path will be parsed and any sql statements in that file will be executed during the benchmark run. If no path is specified the next line is required to be the sql statement to execute.

Multiple statements are allowed within a single run directive, however a benchmark file may contain only one run directive. When -running with `BENCH_PERSIST_RESULTS` or `BENCH_VALIDATE`, only the last `SELECT` or `WITH` statement from that run -directive will be used for comparison.

The run directive (including any following sql statement) must be +when persisting or validating results, only the last `SELECT` or `WITH` statement from that run directive will be used +for comparison.

The run directive (including any following sql statement) must be followed by a blank line.

Example:
run sql_benchmarks/imdb/queries/${QUERY_NUMBER_PADDED}.sql
diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q00.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q00.benchmark new file mode 100644 index 0000000000000..0ea18a72733b9 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q00.benchmark @@ -0,0 +1,17 @@ +name Q00 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT COUNT(*) +FROM hits; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q00.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..1512ef10b5d7e --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q01.benchmark @@ -0,0 +1,18 @@ +name Q01 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT COUNT(*) +FROM hits +WHERE "AdvEngineID" <> 0; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q01.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..3bc1a4ec4acbd --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q02.benchmark @@ -0,0 +1,17 @@ +name Q02 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT SUM("AdvEngineID"), COUNT(*), AVG("ResolutionWidth") +FROM hits; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q02.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..c545a27d8c46d --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q03.benchmark @@ -0,0 +1,17 @@ +name Q03 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT AVG("UserID") +FROM hits; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q03.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..5ae8ad3b8ffed --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q04.benchmark @@ -0,0 +1,17 @@ +name Q04 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT COUNT(DISTINCT "UserID") +FROM hits; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q04.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..dd2f654698d50 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q05.benchmark @@ -0,0 +1,17 @@ +name Q05 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT COUNT(DISTINCT "SearchPhrase") +FROM hits; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q05.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..1b5e105a1acc2 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q06.benchmark @@ -0,0 +1,17 @@ +name Q06 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT MIN("EventDate"), MAX("EventDate") +FROM hits; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q06.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..882f5ad3d8327 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q07.benchmark @@ -0,0 +1,20 @@ +name Q07 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "AdvEngineID", COUNT(*) +FROM hits +WHERE "AdvEngineID" <> 0 +GROUP BY "AdvEngineID" +ORDER BY COUNT(*) DESC; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q07.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..525a39cf47ad8 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q08.benchmark @@ -0,0 +1,16 @@ +name Q08 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "RegionID", COUNT(DISTINCT "UserID") AS u FROM hits GROUP BY "RegionID" ORDER BY u DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q08.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..3f58fc3c95f4b --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q09.benchmark @@ -0,0 +1,16 @@ +name Q09 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "RegionID", SUM("AdvEngineID"), COUNT(*) AS c, AVG("ResolutionWidth"), COUNT(DISTINCT "UserID") FROM hits GROUP BY "RegionID" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q09.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..4a3506d74728d --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q10.benchmark @@ -0,0 +1,16 @@ +name Q10 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhoneModel" ORDER BY u DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q10.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q11.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q11.benchmark new file mode 100644 index 0000000000000..b18f1946782f0 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q11.benchmark @@ -0,0 +1,16 @@ +name Q11 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "MobilePhone", "MobilePhoneModel", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "MobilePhoneModel" <> '' GROUP BY "MobilePhone", "MobilePhoneModel" ORDER BY u DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q11.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q12.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q12.benchmark new file mode 100644 index 0000000000000..0586305e3d4f0 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q12.benchmark @@ -0,0 +1,16 @@ +name Q12 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchPhrase", COUNT(*) AS c FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q12.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q13.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q13.benchmark new file mode 100644 index 0000000000000..b36449e05bd4c --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q13.benchmark @@ -0,0 +1,16 @@ +name Q13 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchPhrase", COUNT(DISTINCT "UserID") AS u FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY u DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q13.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q14.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q14.benchmark new file mode 100644 index 0000000000000..2b7c3b196f22e --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q14.benchmark @@ -0,0 +1,16 @@ +name Q14 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchEngineID", "SearchPhrase", COUNT(*) AS c FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchEngineID", "SearchPhrase" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q14.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q15.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q15.benchmark new file mode 100644 index 0000000000000..8e8be046446a7 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q15.benchmark @@ -0,0 +1,16 @@ +name Q15 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "UserID", COUNT(*) FROM hits GROUP BY "UserID" ORDER BY COUNT(*) DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q15.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q16.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q16.benchmark new file mode 100644 index 0000000000000..93fb630d73699 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q16.benchmark @@ -0,0 +1,16 @@ +name Q16 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "UserID", "SearchPhrase", COUNT(*) FROM hits GROUP BY "UserID", "SearchPhrase" ORDER BY COUNT(*) DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q16.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q17.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q17.benchmark new file mode 100644 index 0000000000000..60725ae005997 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q17.benchmark @@ -0,0 +1,16 @@ +name Q17 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "UserID", "SearchPhrase", COUNT(*) FROM hits GROUP BY "UserID", "SearchPhrase" LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q17.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q18.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q18.benchmark new file mode 100644 index 0000000000000..1f5bad2a029f5 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q18.benchmark @@ -0,0 +1,16 @@ +name Q18 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "UserID", extract(minute FROM to_timestamp_seconds("EventTime")) AS m, "SearchPhrase", COUNT(*) FROM hits GROUP BY "UserID", m, "SearchPhrase" ORDER BY COUNT(*) DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q18.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q19.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q19.benchmark new file mode 100644 index 0000000000000..7bd760aaff9fe --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q19.benchmark @@ -0,0 +1,18 @@ +name Q19 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "UserID" +FROM hits +WHERE "UserID" = 435090932899640449; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q19.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q20.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q20.benchmark new file mode 100644 index 0000000000000..6ec6c5c0a61ef --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q20.benchmark @@ -0,0 +1,18 @@ +name Q20 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT COUNT(*) +FROM hits +WHERE "URL" LIKE '%google%'; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q20.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q21.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q21.benchmark new file mode 100644 index 0000000000000..a1123e9391983 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q21.benchmark @@ -0,0 +1,16 @@ +name Q21 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchPhrase", MIN("URL"), COUNT(*) AS c FROM hits WHERE "URL" LIKE '%google%' AND "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q21.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q22.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q22.benchmark new file mode 100644 index 0000000000000..9df61823b3107 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q22.benchmark @@ -0,0 +1,16 @@ +name Q22 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchPhrase", MIN("URL"), MIN("Title"), COUNT(*) AS c, COUNT(DISTINCT "UserID") FROM hits WHERE "Title" LIKE '%Google%' AND "URL" NOT LIKE '%.google.%' AND "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q22.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q23.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q23.benchmark new file mode 100644 index 0000000000000..aa742cb56bfc7 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q23.benchmark @@ -0,0 +1,16 @@ +name Q23 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT * FROM hits WHERE "URL" LIKE '%google%' ORDER BY "EventTime" LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q23.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q24.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q24.benchmark new file mode 100644 index 0000000000000..4b30c5fef3c72 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q24.benchmark @@ -0,0 +1,16 @@ +name Q24 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "EventTime" LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q24.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q25.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q25.benchmark new file mode 100644 index 0000000000000..5a8a425703662 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q25.benchmark @@ -0,0 +1,16 @@ +name Q25 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "SearchPhrase" LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q25.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q26.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q26.benchmark new file mode 100644 index 0000000000000..b87f59a847adf --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q26.benchmark @@ -0,0 +1,16 @@ +name Q26 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "EventTime", "SearchPhrase" LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q26.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q27.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q27.benchmark new file mode 100644 index 0000000000000..84e43c2272d57 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q27.benchmark @@ -0,0 +1,17 @@ +name Q27 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +-- DataFusion length(...) counts characters; use octet_length(...) for ClickBench byte-length semantics. +SELECT "CounterID", AVG(octet_length("URL")) AS l, COUNT(*) AS c FROM hits WHERE "URL" <> '' GROUP BY "CounterID" HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q27.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q28.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q28.benchmark new file mode 100644 index 0000000000000..02cbfb20c09f1 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q28.benchmark @@ -0,0 +1,17 @@ +name Q28 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +-- DataFusion length(...) counts characters; use octet_length(...) for ClickBench byte-length semantics. +SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(octet_length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q28.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q29.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q29.benchmark new file mode 100644 index 0000000000000..a76d6c6f2d4b0 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q29.benchmark @@ -0,0 +1,106 @@ +name Q29 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT SUM("ResolutionWidth"), + SUM("ResolutionWidth" + 1), + SUM("ResolutionWidth" + 2), + SUM("ResolutionWidth" + 3), + SUM("ResolutionWidth" + 4), + SUM("ResolutionWidth" + 5), + SUM("ResolutionWidth" + 6), + SUM("ResolutionWidth" + 7), + SUM("ResolutionWidth" + 8), + SUM("ResolutionWidth" + 9), + SUM("ResolutionWidth" + 10), + SUM("ResolutionWidth" + 11), + SUM("ResolutionWidth" + 12), + SUM("ResolutionWidth" + 13), + SUM("ResolutionWidth" + 14), + SUM("ResolutionWidth" + 15), + SUM("ResolutionWidth" + 16), + SUM("ResolutionWidth" + 17), + SUM("ResolutionWidth" + 18), + SUM("ResolutionWidth" + 19), + SUM("ResolutionWidth" + 20), + SUM("ResolutionWidth" + 21), + SUM("ResolutionWidth" + 22), + SUM("ResolutionWidth" + 23), + SUM("ResolutionWidth" + 24), + SUM("ResolutionWidth" + 25), + SUM("ResolutionWidth" + 26), + SUM("ResolutionWidth" + 27), + SUM("ResolutionWidth" + 28), + SUM("ResolutionWidth" + 29), + SUM("ResolutionWidth" + 30), + SUM("ResolutionWidth" + 31), + SUM("ResolutionWidth" + 32), + SUM("ResolutionWidth" + 33), + SUM("ResolutionWidth" + 34), + SUM("ResolutionWidth" + 35), + SUM("ResolutionWidth" + 36), + SUM("ResolutionWidth" + 37), + SUM("ResolutionWidth" + 38), + SUM("ResolutionWidth" + 39), + SUM("ResolutionWidth" + 40), + SUM("ResolutionWidth" + 41), + SUM("ResolutionWidth" + 42), + SUM("ResolutionWidth" + 43), + SUM("ResolutionWidth" + 44), + SUM("ResolutionWidth" + 45), + SUM("ResolutionWidth" + 46), + SUM("ResolutionWidth" + 47), + SUM("ResolutionWidth" + 48), + SUM("ResolutionWidth" + 49), + SUM("ResolutionWidth" + 50), + SUM("ResolutionWidth" + 51), + SUM("ResolutionWidth" + 52), + SUM("ResolutionWidth" + 53), + SUM("ResolutionWidth" + 54), + SUM("ResolutionWidth" + 55), + SUM("ResolutionWidth" + 56), + SUM("ResolutionWidth" + 57), + SUM("ResolutionWidth" + 58), + SUM("ResolutionWidth" + 59), + SUM("ResolutionWidth" + 60), + SUM("ResolutionWidth" + 61), + SUM("ResolutionWidth" + 62), + SUM("ResolutionWidth" + 63), + SUM("ResolutionWidth" + 64), + SUM("ResolutionWidth" + 65), + SUM("ResolutionWidth" + 66), + SUM("ResolutionWidth" + 67), + SUM("ResolutionWidth" + 68), + SUM("ResolutionWidth" + 69), + SUM("ResolutionWidth" + 70), + SUM("ResolutionWidth" + 71), + SUM("ResolutionWidth" + 72), + SUM("ResolutionWidth" + 73), + SUM("ResolutionWidth" + 74), + SUM("ResolutionWidth" + 75), + SUM("ResolutionWidth" + 76), + SUM("ResolutionWidth" + 77), + SUM("ResolutionWidth" + 78), + SUM("ResolutionWidth" + 79), + SUM("ResolutionWidth" + 80), + SUM("ResolutionWidth" + 81), + SUM("ResolutionWidth" + 82), + SUM("ResolutionWidth" + 83), + SUM("ResolutionWidth" + 84), + SUM("ResolutionWidth" + 85), + SUM("ResolutionWidth" + 86), + SUM("ResolutionWidth" + 87), + SUM("ResolutionWidth" + 88), + SUM("ResolutionWidth" + 89) +FROM hits; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q29.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q30.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q30.benchmark new file mode 100644 index 0000000000000..740a6724cca38 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q30.benchmark @@ -0,0 +1,16 @@ +name Q30 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SearchEngineID", "ClientIP", COUNT(*) AS c, SUM("IsRefresh"), AVG("ResolutionWidth") FROM hits WHERE "SearchPhrase" <> '' GROUP BY "SearchEngineID", "ClientIP" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q30.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q31.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q31.benchmark new file mode 100644 index 0000000000000..91035bcf6916f --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q31.benchmark @@ -0,0 +1,16 @@ +name Q31 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "WatchID", "ClientIP", COUNT(*) AS c, SUM("IsRefresh"), AVG("ResolutionWidth") FROM hits WHERE "SearchPhrase" <> '' GROUP BY "WatchID", "ClientIP" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q31.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q32.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q32.benchmark new file mode 100644 index 0000000000000..15a58676098fd --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q32.benchmark @@ -0,0 +1,16 @@ +name Q32 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "WatchID", "ClientIP", COUNT(*) AS c, SUM("IsRefresh"), AVG("ResolutionWidth") FROM hits GROUP BY "WatchID", "ClientIP" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q32.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q33.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q33.benchmark new file mode 100644 index 0000000000000..2742a609c306f --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q33.benchmark @@ -0,0 +1,16 @@ +name Q33 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "URL", COUNT(*) AS c FROM hits GROUP BY "URL" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q33.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q34.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q34.benchmark new file mode 100644 index 0000000000000..6b8c2beb9c7aa --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q34.benchmark @@ -0,0 +1,16 @@ +name Q34 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT 1, "URL", COUNT(*) AS c FROM hits GROUP BY 1, "URL" ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q34.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q35.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q35.benchmark new file mode 100644 index 0000000000000..75a6b210996ab --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q35.benchmark @@ -0,0 +1,16 @@ +name Q35 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "ClientIP", "ClientIP" - 1, "ClientIP" - 2, "ClientIP" - 3, COUNT(*) AS c FROM hits GROUP BY "ClientIP", "ClientIP" - 1, "ClientIP" - 2, "ClientIP" - 3 ORDER BY c DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q35.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q36.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q36.benchmark new file mode 100644 index 0000000000000..95f7c4b03b203 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q36.benchmark @@ -0,0 +1,16 @@ +name Q36 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "URL", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "DontCountHits" = 0 AND "IsRefresh" = 0 AND "URL" <> '' GROUP BY "URL" ORDER BY PageViews DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q36.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q37.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q37.benchmark new file mode 100644 index 0000000000000..dd8ef38bb2ed2 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q37.benchmark @@ -0,0 +1,16 @@ +name Q37 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "Title", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "DontCountHits" = 0 AND "IsRefresh" = 0 AND "Title" <> '' GROUP BY "Title" ORDER BY PageViews DESC LIMIT 10; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q37.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q38.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q38.benchmark new file mode 100644 index 0000000000000..93d4d1722d3f2 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q38.benchmark @@ -0,0 +1,16 @@ +name Q38 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "URL", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 AND "IsLink" <> 0 AND "IsDownload" = 0 GROUP BY "URL" ORDER BY PageViews DESC LIMIT 10 OFFSET 1000; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q38.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q39.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q39.benchmark new file mode 100644 index 0000000000000..443e2120fca92 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q39.benchmark @@ -0,0 +1,16 @@ +name Q39 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "TraficSourceID", "SearchEngineID", "AdvEngineID", CASE WHEN ("SearchEngineID" = 0 AND "AdvEngineID" = 0) THEN "Referer" ELSE '' END AS Src, "URL" AS Dst, COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 GROUP BY "TraficSourceID", "SearchEngineID", "AdvEngineID", Src, Dst ORDER BY PageViews DESC LIMIT 10 OFFSET 1000; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q39.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q40.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q40.benchmark new file mode 100644 index 0000000000000..b3358dc1661d7 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q40.benchmark @@ -0,0 +1,16 @@ +name Q40 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "URLHash", "EventDate", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 AND "TraficSourceID" IN (-1, 6) AND "RefererHash" = 3594120000172545465 GROUP BY "URLHash", "EventDate" ORDER BY PageViews DESC LIMIT 10 OFFSET 100; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q40.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q41.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q41.benchmark new file mode 100644 index 0000000000000..0cbafea4682ca --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q41.benchmark @@ -0,0 +1,16 @@ +name Q41 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "WindowClientWidth", "WindowClientHeight", COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-01' AND "EventDate" <= '2013-07-31' AND "IsRefresh" = 0 AND "DontCountHits" = 0 AND "URLHash" = 2868770270353813622 GROUP BY "WindowClientWidth", "WindowClientHeight" ORDER BY PageViews DESC LIMIT 10 OFFSET 10000; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q41.csv diff --git a/benchmarks/sql_benchmarks/clickbench/benchmarks/q42.benchmark b/benchmarks/sql_benchmarks/clickbench/benchmarks/q42.benchmark new file mode 100644 index 0000000000000..7822062fd0150 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/benchmarks/q42.benchmark @@ -0,0 +1,16 @@ +name Q42 +group clickbench + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT DATE_TRUNC('minute', to_timestamp_seconds("EventTime")) AS M, COUNT(*) AS PageViews FROM hits WHERE "CounterID" = 62 AND "EventDate" >= '2013-07-14' AND "EventDate" <= '2013-07-15' AND "IsRefresh" = 0 AND "DontCountHits" = 0 GROUP BY DATE_TRUNC('minute', to_timestamp_seconds("EventTime")) ORDER BY DATE_TRUNC('minute', M) LIMIT 10 OFFSET 1000; + +result sql_benchmarks/clickbench/results/${CLICKBENCH_TYPE:-single}/q42.csv diff --git a/benchmarks/sql_benchmarks/clickbench/clickbench.suite b/benchmarks/sql_benchmarks/clickbench/clickbench.suite new file mode 100644 index 0000000000000..74d8a5cc2ae56 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/clickbench.suite @@ -0,0 +1,25 @@ +description = "ClickBench analytics queries over the hits dataset" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "partitioning" +env = "CLICKBENCH_TYPE" +default = "single" +values = ["single", "partitioned"] +help = "Selects the single-file or partitioned ClickBench dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench" +description = "Run all ClickBench queries against the single-file dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench --query 7" +description = "Run ClickBench query 7." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench --partitioning partitioned" +description = "Run all ClickBench queries against the partitioned dataset." diff --git a/benchmarks/sql_benchmarks/clickbench/init/load-partitioned.sql b/benchmarks/sql_benchmarks/clickbench/init/load-partitioned.sql new file mode 100644 index 0000000000000..2e4a39625c304 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/init/load-partitioned.sql @@ -0,0 +1,3 @@ +CREATE EXTERNAL TABLE hits_raw STORED AS PARQUET LOCATION '${DATA_DIR:-data}/hits_partitioned/'; + +CREATE VIEW hits AS SELECT * EXCEPT ("EventDate"), CAST(CAST("EventDate" AS INTEGER) AS DATE) AS "EventDate" FROM hits_raw \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/clickbench/init/load-single.sql b/benchmarks/sql_benchmarks/clickbench/init/load-single.sql new file mode 100644 index 0000000000000..3bba41744371d --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/init/load-single.sql @@ -0,0 +1,3 @@ +CREATE EXTERNAL TABLE hits_raw STORED AS PARQUET LOCATION '${DATA_DIR:-data}/hits.parquet'; + +CREATE VIEW hits AS SELECT * EXCEPT ("EventDate"), CAST(CAST("EventDate" AS INTEGER) AS DATE) AS "EventDate" FROM hits_raw \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/clickbench/init/set_config.sql b/benchmarks/sql_benchmarks/clickbench/init/set_config.sql new file mode 100644 index 0000000000000..ee2ac0b3c9529 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench/init/set_config.sql @@ -0,0 +1,5 @@ +# ClickBench partitioned dataset was written by an ancient version of PyArrow that +# wrote strings with the wrong logical type. To read it correctly, we must +# automatically convert binary to string. + +SET datafusion.execution.parquet.binary_as_string = true; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q00.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q00.benchmark new file mode 100644 index 0000000000000..10f58f493e5ef --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q00.benchmark @@ -0,0 +1,18 @@ +name Q00 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT COUNT(DISTINCT "SearchPhrase"), COUNT(DISTINCT "MobilePhone"), COUNT(DISTINCT "MobilePhoneModel") +FROM hits; + +result sql_benchmarks/clickbench_extended/results/q00.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..cfaaec1037fc3 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q01.benchmark @@ -0,0 +1,18 @@ +name Q01 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT COUNT(DISTINCT "HitColor"), COUNT(DISTINCT "BrowserCountry"), COUNT(DISTINCT "BrowserLanguage") +FROM hits; + +result sql_benchmarks/clickbench_extended/results/q01.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..711919c35fce6 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q02.benchmark @@ -0,0 +1,17 @@ +name Q02 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "BrowserCountry", COUNT(DISTINCT "SocialNetwork"), COUNT(DISTINCT "HitColor"), COUNT(DISTINCT "BrowserLanguage"), COUNT(DISTINCT "SocialAction") FROM hits GROUP BY 1 ORDER BY 2 DESC LIMIT 10; + +result sql_benchmarks/clickbench_extended/results/q02.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..a1fe3aa0f3f34 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q03.benchmark @@ -0,0 +1,17 @@ +name Q03 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "SocialSourceNetworkID", "RegionID", COUNT(*), AVG("Age"), AVG("ParamPrice"), STDDEV("ParamPrice") as s, VAR("ParamPrice") FROM hits GROUP BY "SocialSourceNetworkID", "RegionID" HAVING s IS NOT NULL ORDER BY s DESC LIMIT 10; + +result sql_benchmarks/clickbench_extended/results/q03.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..f08525fb91960 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q04.benchmark @@ -0,0 +1,17 @@ +name Q04 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "ClientIP", "WatchID", COUNT(*) c, MIN("ResponseStartTiming") tmin, MEDIAN("ResponseStartTiming") tmed, MAX("ResponseStartTiming") tmax FROM hits WHERE "JavaEnable" = 0 GROUP BY "ClientIP", "WatchID" HAVING c > 1 ORDER BY tmed DESC LIMIT 10; + +result sql_benchmarks/clickbench_extended/results/q04.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..8e594d24afc19 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q05.benchmark @@ -0,0 +1,17 @@ +name Q05 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "ClientIP", "WatchID", COUNT(*) c, MIN("ResponseStartTiming") tmin, APPROX_PERCENTILE_CONT(0.95) WITHIN GROUP (ORDER BY "ResponseStartTiming") tp95, MAX("ResponseStartTiming") tmax FROM 'hits' WHERE "JavaEnable" = 0 GROUP BY "ClientIP", "WatchID" HAVING c > 1 ORDER BY tp95 DESC LIMIT 10; + +result sql_benchmarks/clickbench_extended/results/q05.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..4ae1a8efb629c --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q06.benchmark @@ -0,0 +1,17 @@ +name Q06 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT COUNT(*) AS ShareCount FROM hits WHERE "IsMobile" = 1 AND "MobilePhoneModel" LIKE 'iPhone%' AND "SocialAction" = 'share' AND "SocialSourceNetworkID" IN (5, 12) AND "ClientTimeZone" BETWEEN -5 AND 5 AND regexp_match("Referer", '\/campaign\/(spring|summer)_promo') IS NOT NULL AND CASE WHEN split_part(split_part(CAST("URL" AS STRING), 'resolution=', 2), '&', 1) ~ '^\d+$' THEN split_part(split_part(CAST("URL" AS STRING), 'resolution=', 2), '&', 1)::INT ELSE 0 END > 1920 AND levenshtein(CAST("UTMSource" AS STRING), CAST("UTMCampaign" AS STRING)) < 3; + +result sql_benchmarks/clickbench_extended/results/q06.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..9bf7e1052958b --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q07.benchmark @@ -0,0 +1,17 @@ +name Q07 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT "WatchID", MIN("ResolutionWidth") as wmin, MAX("ResolutionWidth") as wmax, SUM("IsRefresh") as srefresh FROM hits GROUP BY "WatchID" ORDER BY "WatchID" DESC LIMIT 10; + +result sql_benchmarks/clickbench_extended/results/q07.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..709a3b74e870b --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q08.benchmark @@ -0,0 +1,19 @@ +name Q08 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +-- Must set for ClickBench hits_partitioned dataset. See https://github.com/apache/datafusion/issues/16591 +-- set datafusion.execution.parquet.binary_as_string = true +SELECT "RegionID", "UserAgent", "OS", AVG(to_timestamp("ResponseEndTiming")-to_timestamp("ResponseStartTiming")) as avg_response_time, AVG(to_timestamp("ResponseEndTiming")-to_timestamp("ConnectTiming")) as avg_latency FROM hits GROUP BY "RegionID", "UserAgent", "OS" ORDER BY avg_latency DESC limit 10; + +result sql_benchmarks/clickbench_extended/results/q08.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..8405941975e6c --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q09.benchmark @@ -0,0 +1,21 @@ +name Q09 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT MAX(len) FROM ( + SELECT LENGTH(FIRST_VALUE("URL" ORDER BY "EventTime")) as len + FROM hits + GROUP BY "UserID" +); + +result sql_benchmarks/clickbench_extended/results/q09.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..a1a1210d99fcc --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q10.benchmark @@ -0,0 +1,21 @@ +name Q10 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT MAX(len) FROM ( + SELECT LENGTH(FIRST_VALUE("URL" ORDER BY "EventTime")) as len + FROM hits + GROUP BY "OS" +); + +result sql_benchmarks/clickbench_extended/results/q10.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q11.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q11.benchmark new file mode 100644 index 0000000000000..60a482eaee9b0 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q11.benchmark @@ -0,0 +1,21 @@ +name Q11 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT MAX(fv) FROM ( + SELECT FIRST_VALUE("WatchID" ORDER BY "EventTime") as fv + FROM hits + GROUP BY "UserID" +); + +result sql_benchmarks/clickbench_extended/results/q11.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q12.benchmark b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q12.benchmark new file mode 100644 index 0000000000000..81b69296beb46 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/benchmarks/q12.benchmark @@ -0,0 +1,21 @@ +name Q12 +group clickbench_extended +subgroup ${CLICKBENCH_TYPE:-single} + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench/init/load-${CLICKBENCH_TYPE:-single}.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT MAX(fv) FROM ( + SELECT FIRST_VALUE("WatchID" ORDER BY "EventTime") as fv + FROM hits + GROUP BY "OS" +); + +result sql_benchmarks/clickbench_extended/results/q12.csv diff --git a/benchmarks/sql_benchmarks/clickbench_extended/clickbench_extended.suite b/benchmarks/sql_benchmarks/clickbench_extended/clickbench_extended.suite new file mode 100644 index 0000000000000..dfca00b4a03db --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_extended/clickbench_extended.suite @@ -0,0 +1,21 @@ +description = "Extended ClickBench queries over the hits dataset" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "partitioning" +env = "CLICKBENCH_TYPE" +default = "single" +values = ["single", "partitioned"] +help = "Selects the single-file or partitioned ClickBench dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench_extended" +description = "Run all extended ClickBench queries against the single-file dataset." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench_extended --query 4 --partitioning partitioned" +description = "Run extended ClickBench query 4 against the partitioned dataset." diff --git a/benchmarks/sql_benchmarks/clickbench_sorted/benchmarks/q00.benchmark b/benchmarks/sql_benchmarks/clickbench_sorted/benchmarks/q00.benchmark new file mode 100644 index 0000000000000..5c95a91b6addb --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_sorted/benchmarks/q00.benchmark @@ -0,0 +1,17 @@ +name Q00 +group clickbench_sorted + +init sql_benchmarks/clickbench/init/set_config.sql + +load sql_benchmarks/clickbench_sorted/init/load.sql + +assert I +SELECT COUNT(*) > 0 from hits; +---- +true + +run +SELECT * FROM hits ORDER BY "EventTime" DESC limit 10; + +result sql_benchmarks/clickbench_sorted/results/q00.csv + diff --git a/benchmarks/sql_benchmarks/clickbench_sorted/clickbench_sorted.suite b/benchmarks/sql_benchmarks/clickbench_sorted/clickbench_sorted.suite new file mode 100644 index 0000000000000..5c8a0909e3f55 --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_sorted/clickbench_sorted.suite @@ -0,0 +1,28 @@ +description = "ClickBench query over a pre-sorted hits dataset" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "sort-column" +env = "SORTED_BY" +default = "EventTime" +values = ["EventTime", "..."] +help = "Selects the column used to sort the ClickBench data." + +[[options]] +name = "sort-order" +env = "SORTED_ORDER" +default = "ASC" +values = ["ASC", "DESC"] +help = "Selects the sort direction for the ClickBench data." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench_sorted" +description = "Run the sorted ClickBench query ordered by EventTime ascending." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- clickbench_sorted --sort-column UserID --sort-order DESC" +description = "Run the query over data sorted by UserID descending." diff --git a/benchmarks/sql_benchmarks/clickbench_sorted/init/load.sql b/benchmarks/sql_benchmarks/clickbench_sorted/init/load.sql new file mode 100644 index 0000000000000..fa3c379c7b05b --- /dev/null +++ b/benchmarks/sql_benchmarks/clickbench_sorted/init/load.sql @@ -0,0 +1,8 @@ +-- Run benchmark with prefer_existing_sort configuration +-- This allows DataFusion to optimize away redundant sorts while maintaining parallelism + +set datafusion.optimizer.prefer_existing_sort=true; + +CREATE EXTERNAL TABLE hits_raw STORED AS PARQUET LOCATION '${DATA_DIR:-data}/hits_sorted.parquet' WITH ORDER ("${SORTED_BY:-EventTime}" ${SORTED_ORDER:-ASC}); + +CREATE VIEW hits AS SELECT * EXCEPT ("EventDate"), CAST(CAST("EventDate" AS INTEGER) AS DATE) AS "EventDate" FROM hits_raw; diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q01.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q01.benchmark new file mode 100644 index 0000000000000..e499243c55002 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q01.benchmark @@ -0,0 +1,20 @@ +subgroup groupby + +name Q01 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id1, SUM(v1) AS v1 +FROM x +GROUP BY id1; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q01.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q02.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q02.benchmark new file mode 100644 index 0000000000000..a1477574384ce --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q02.benchmark @@ -0,0 +1,20 @@ +subgroup groupby + +name Q02 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id1, id2, SUM(v1) AS v1 +FROM x +GROUP BY id1, id2; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q02.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q03.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q03.benchmark new file mode 100644 index 0000000000000..4368bc9f46217 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q03.benchmark @@ -0,0 +1,20 @@ +subgroup groupby + +name Q03 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id3, SUM(v1) AS v1, AVG(v3) AS v3 +FROM x +GROUP BY id3; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q03.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q04.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q04.benchmark new file mode 100644 index 0000000000000..1813613e5d2b5 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q04.benchmark @@ -0,0 +1,20 @@ +subgroup groupby + +name Q04 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id4, AVG(v1) AS v1, AVG(v2) AS v2, AVG(v3) AS v3 +FROM x +GROUP BY id4; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q04.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q05.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q05.benchmark new file mode 100644 index 0000000000000..dc6b4a4feaa90 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q05.benchmark @@ -0,0 +1,20 @@ +subgroup groupby + +name Q05 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id6, SUM(v1) AS v1, SUM(v2) AS v2, SUM(v3) AS v3 +FROM x +GROUP BY id6; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q05.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q06.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q06.benchmark new file mode 100644 index 0000000000000..67eaeafaf804b --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q06.benchmark @@ -0,0 +1,20 @@ +subgroup groupby + +name Q06 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id4, id5, MEDIAN(v3) AS median_v3, STDDEV(v3) AS sd_v3 +FROM x +GROUP BY id4, id5; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q06.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q07.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q07.benchmark new file mode 100644 index 0000000000000..faa125eb9ec55 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q07.benchmark @@ -0,0 +1,20 @@ +subgroup groupby + +name Q07 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id3, MAX(v1) - MIN(v2) AS range_v1_v2 +FROM x +GROUP BY id3; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q07.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q08.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q08.benchmark new file mode 100644 index 0000000000000..54d46080c678d --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q08.benchmark @@ -0,0 +1,22 @@ +subgroup groupby + +name Q08 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id6, largest2_v3 FROM + ( + SELECT id6, v3 AS largest2_v3, ROW_NUMBER() OVER (PARTITION BY id6 ORDER BY v3 DESC) AS order_v3 + FROM x WHERE v3 IS NOT NULL + ) sub_query WHERE order_v3 <= 2; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q08.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q09.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q09.benchmark new file mode 100644 index 0000000000000..133ddef5c296e --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q09.benchmark @@ -0,0 +1,20 @@ +subgroup groupby + +name Q09 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id2, id4, POWER(CORR(v1, v2), 2) AS r2 +FROM x +GROUP BY id2, id4; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q09.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q10.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q10.benchmark new file mode 100644 index 0000000000000..a302ef2408260 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/groupby/q10.benchmark @@ -0,0 +1,18 @@ +subgroup groupby + +name Q10 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} groupby ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_groupby_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT id1, id2, id3, id4, id5, id6, SUM(v3) AS v3, COUNT(*) AS count FROM x GROUP BY id1, id2, id3, id4, id5, id6; + +result sql_benchmarks/h2o/results/groupby/${H2O_BENCH_SIZE:-small}/q10.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/join/q01.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q01.benchmark new file mode 100644 index 0000000000000..4271ba8e43efc --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q01.benchmark @@ -0,0 +1,28 @@ +subgroup join + +name Q01 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} join ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_join_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT x.id1, + x.id2, + x.id3, + x.id4 as xid4, + small.id4 as smallid4, + x.id5, + x.id6, + x.v1, + small.v2 +FROM x +INNER JOIN small ON x.id1 = small.id1; + +result sql_benchmarks/h2o/results/join/${H2O_BENCH_SIZE:-small}/q01.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/join/q02.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q02.benchmark new file mode 100644 index 0000000000000..48369c4a58197 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q02.benchmark @@ -0,0 +1,30 @@ +subgroup join + +name Q02 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} join ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_join_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT x.id1 as xid1, + medium.id1 as mediumid1, + x.id2, + x.id3, + x.id4 as xid4, + medium.id4 as mediumid4, + x.id5 as xid5, + medium.id5 as mediumid5, + x.id6, + x.v1, + medium.v2 +FROM x +INNER JOIN medium ON x.id2 = medium.id2; + +result sql_benchmarks/h2o/results/join/${H2O_BENCH_SIZE:-small}/q02.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/join/q03.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q03.benchmark new file mode 100644 index 0000000000000..abf7296b2128e --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q03.benchmark @@ -0,0 +1,30 @@ +subgroup join + +name Q03 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} join ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_join_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT x.id1 as xid1, + medium.id1 as mediumid1, + x.id2, + x.id3, + x.id4 as xid4, + medium.id4 as mediumid4, + x.id5 as xid5, + medium.id5 as mediumid5, + x.id6, + x.v1, + medium.v2 +FROM x +LEFT JOIN medium ON x.id2 = medium.id2; + +result sql_benchmarks/h2o/results/join/${H2O_BENCH_SIZE:-small}/q03.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/join/q04.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q04.benchmark new file mode 100644 index 0000000000000..8884fb38e8a5d --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q04.benchmark @@ -0,0 +1,30 @@ +subgroup join + +name Q04 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} join ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_join_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT x.id1 as xid1, + medium.id1 as mediumid1, + x.id2, + x.id3, + x.id4 as xid4, + medium.id4 as mediumid4, + x.id5 as xid5, + medium.id5 as mediumid5, + x.id6, + x.v1, + medium.v2 +FROM x +JOIN medium ON x.id5 = medium.id5; + +result sql_benchmarks/h2o/results/join/${H2O_BENCH_SIZE:-small}/q04.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/join/q05.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q05.benchmark new file mode 100644 index 0000000000000..19a5c47bfe093 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/join/q05.benchmark @@ -0,0 +1,32 @@ +subgroup join + +name Q05 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} join ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_join_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +SELECT x.id1 as xid1, + large.id1 as largeid1, + x.id2 as xid2, + large.id2 as largeid2, + x.id3, + x.id4 as xid4, + large.id4 as largeid4, + x.id5 as xid5, + large.id5 as largeid5, + x.id6 as xid6, + large.id6 as largeid6, + x.v1, + large.v2 +FROM x +JOIN large ON x.id3 = large.id3; + +result sql_benchmarks/h2o/results/join/${H2O_BENCH_SIZE:-small}/q05.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q01.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q01.benchmark new file mode 100644 index 0000000000000..4a95f07cc18d7 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q01.benchmark @@ -0,0 +1,25 @@ +subgroup window + +name Q01 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- Basic Window +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER () AS window_basic +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q01.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q02.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q02.benchmark new file mode 100644 index 0000000000000..8bc4b605b8d33 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q02.benchmark @@ -0,0 +1,26 @@ +subgroup window + +name Q02 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- Sorted Window +SELECT + id1, + id2, + id3, + v2, + first_value(v2) OVER (ORDER BY id3) AS first_order_by, + row_number() OVER (ORDER BY id3) AS row_number_order_by +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q02.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q03.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q03.benchmark new file mode 100644 index 0000000000000..53f5cff31837b --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q03.benchmark @@ -0,0 +1,27 @@ +subgroup window + +name Q03 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- PARTITION BY +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (PARTITION BY id1) AS sum_by_id1, + sum(v2) OVER (PARTITION BY id2) AS sum_by_id2, + sum(v2) OVER (PARTITION BY id3) AS sum_by_id3 +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q03.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q04.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q04.benchmark new file mode 100644 index 0000000000000..0111235ebb8e9 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q04.benchmark @@ -0,0 +1,25 @@ +subgroup window + +name Q04 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- PARTITION BY ORDER BY +SELECT + id1, + id2, + id3, + v2, + first_value(v2) OVER (PARTITION BY id2 ORDER BY id3) AS first_by_id2_ordered_by_id3 +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q04.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q05.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q05.benchmark new file mode 100644 index 0000000000000..e75827f8aafb9 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q05.benchmark @@ -0,0 +1,26 @@ +subgroup window + +name Q05 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- Lead and Lag +SELECT + id1, + id2, + id3, + v2, + first_value(v2) OVER (ORDER BY id3 ROWS BETWEEN 1 PRECEDING AND 1 PRECEDING) AS my_lag, + first_value(v2) OVER (ORDER BY id3 ROWS BETWEEN 1 FOLLOWING AND 1 FOLLOWING) AS my_lead +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q05.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q06.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q06.benchmark new file mode 100644 index 0000000000000..8d58e41f4ef21 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q06.benchmark @@ -0,0 +1,25 @@ +subgroup window + +name Q06 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- Moving Averages +SELECT + id1, + id2, + id3, + v2, + avg(v2) OVER (ORDER BY id3 ROWS BETWEEN 100 PRECEDING AND CURRENT ROW) AS my_moving_average +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q06.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q07.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q07.benchmark new file mode 100644 index 0000000000000..f0d8abdbc4aa3 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q07.benchmark @@ -0,0 +1,25 @@ +subgroup window + +name Q07 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- Rolling Sum +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (ORDER BY id3 ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS my_rolling_sum +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q07.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q08.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q08.benchmark new file mode 100644 index 0000000000000..599c331d45be8 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q08.benchmark @@ -0,0 +1,25 @@ +subgroup window + +name Q08 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- RANGE BETWEEN +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (ORDER BY v2 RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) AS my_range_between +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q08.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q09.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q09.benchmark new file mode 100644 index 0000000000000..286e3dfd0f9cc --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q09.benchmark @@ -0,0 +1,26 @@ +subgroup window + +name Q09 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- First PARTITION BY ROWS BETWEEN +SELECT + id1, + id2, + id3, + v2, + first_value(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 1 PRECEDING AND 1 PRECEDING) AS my_lag_by_id2, + first_value(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 1 FOLLOWING AND 1 FOLLOWING) AS my_lead_by_id2 +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q09.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q10.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q10.benchmark new file mode 100644 index 0000000000000..92cfffe233464 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q10.benchmark @@ -0,0 +1,25 @@ +subgroup window + +name Q10 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- Moving Averages PARTITION BY +SELECT + id1, + id2, + id3, + v2, + avg(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN 100 PRECEDING AND CURRENT ROW) AS my_moving_average_by_id2 +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q10.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q11.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q11.benchmark new file mode 100644 index 0000000000000..509b0931d8ce7 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q11.benchmark @@ -0,0 +1,25 @@ +subgroup window + +name Q11 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- Rolling Sum PARTITION BY +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (PARTITION BY id2 ORDER BY id3 ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) AS my_rolling_sum_by_id2 +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q11.csv diff --git a/benchmarks/sql_benchmarks/h2o/benchmarks/window/q12.benchmark b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q12.benchmark new file mode 100644 index 0000000000000..c63777d704503 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/benchmarks/window/q12.benchmark @@ -0,0 +1,25 @@ +subgroup window + +name Q12 +group h2o + +echo Loading ${H2O_BENCH_SIZE:-small} window ${H2O_FILE_TYPE:-csv} h2o data + +load sql_benchmarks/h2o/init/load_window_${H2O_BENCH_SIZE:-small}_${H2O_FILE_TYPE:-csv}.sql + +assert I +SELECT COUNT(*) > 0 FROM x +---- +true + +run +-- RANGE BETWEEN PARTITION BY +SELECT + id1, + id2, + id3, + v2, + sum(v2) OVER (PARTITION BY id2 ORDER BY v2 RANGE BETWEEN 3 PRECEDING AND CURRENT ROW) AS my_range_between_by_id2 +FROM x; + +result sql_benchmarks/h2o/results/window/${H2O_BENCH_SIZE:-small}/q12.csv diff --git a/benchmarks/sql_benchmarks/h2o/h2o.suite b/benchmarks/sql_benchmarks/h2o/h2o.suite new file mode 100644 index 0000000000000..27d83285ba026 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/h2o.suite @@ -0,0 +1,33 @@ +description = "H2O group-by, join, and window SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "size" +env = "H2O_BENCH_SIZE" +default = "small" +values = ["small", "medium", "big"] +help = "Selects the H2O dataset size." + +[[options]] +name = "format" +short = "f" +env = "H2O_FILE_TYPE" +default = "csv" +values = ["csv", "parquet"] +help = "Selects the H2O data format." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- h2o" +description = "Run all H2O queries with the small CSV datasets." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- h2o --query 3 --subgroup window" +description = "Run H2O window query 3." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- h2o --subgroup join --size medium -f parquet" +description = "Run the H2O join queries with the medium Parquet dataset." diff --git a/benchmarks/sql_benchmarks/h2o/init/load_groupby_big_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_groupby_big_csv.sql new file mode 100644 index 0000000000000..a930b58aad89d --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_groupby_big_csv.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/G1_1e9_1e9_100_0.csv'; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/load_groupby_big_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_groupby_big_parquet.sql new file mode 100644 index 0000000000000..16a561ee99b2a --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_groupby_big_parquet.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/G1_1e9_1e9_100_0.parquet'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_groupby_medium_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_groupby_medium_csv.sql new file mode 100644 index 0000000000000..8992ed251f547 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_groupby_medium_csv.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/G1_1e8_1e8_100_0.csv'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_groupby_medium_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_groupby_medium_parquet.sql new file mode 100644 index 0000000000000..dcf77c6defb75 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_groupby_medium_parquet.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/G1_1e8_1e8_100_0.parquet'; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/load_groupby_small_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_groupby_small_csv.sql new file mode 100644 index 0000000000000..9c353b406ad05 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_groupby_small_csv.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/G1_1e7_1e7_100_0.csv'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_groupby_small_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_groupby_small_parquet.sql new file mode 100644 index 0000000000000..4e9eb51f74ab5 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_groupby_small_parquet.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/G1_1e7_1e7_100_0.parquet'; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/h2o/init/load_join_big_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_join_big_csv.sql new file mode 100644 index 0000000000000..6c67c9fd56074 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_join_big_csv.sql @@ -0,0 +1,7 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_NA_0.csv'; + +CREATE EXTERNAL TABLE small STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_1e3_0.csv'; + +CREATE EXTERNAL TABLE medium STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_1e6_0.csv'; + +CREATE EXTERNAL TABLE large STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_1e9_NA.csv'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_join_big_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_join_big_parquet.sql new file mode 100644 index 0000000000000..f84f17199617f --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_join_big_parquet.sql @@ -0,0 +1,7 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_NA_0.parquet'; + +CREATE EXTERNAL TABLE small STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_1e3_0.parquet'; + +CREATE EXTERNAL TABLE medium STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_1e6_0.parquet'; + +CREATE EXTERNAL TABLE large STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_1e9_NA.parquet'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_join_medium_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_join_medium_csv.sql new file mode 100644 index 0000000000000..fcb3916a9751b --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_join_medium_csv.sql @@ -0,0 +1,7 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_NA_0.csv'; + +CREATE EXTERNAL TABLE small STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_1e2_0.csv'; + +CREATE EXTERNAL TABLE medium STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_1e5_0.csv'; + +CREATE EXTERNAL TABLE large STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_1e8_NA.csv'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_join_medium_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_join_medium_parquet.sql new file mode 100644 index 0000000000000..175a38b44ab12 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_join_medium_parquet.sql @@ -0,0 +1,7 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_NA_0.parquet'; + +CREATE EXTERNAL TABLE small STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_1e2_0.parquet'; + +CREATE EXTERNAL TABLE medium STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_1e5_0.parquet'; + +CREATE EXTERNAL TABLE large STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_1e8_NA.parquet'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_join_small_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_join_small_csv.sql new file mode 100644 index 0000000000000..0867e6c06bb5a --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_join_small_csv.sql @@ -0,0 +1,7 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_NA_0.csv'; + +CREATE EXTERNAL TABLE small STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_1e1_0.csv'; + +CREATE EXTERNAL TABLE medium STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_1e4_0.csv'; + +CREATE EXTERNAL TABLE large STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_1e7_NA.csv'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_join_small_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_join_small_parquet.sql new file mode 100644 index 0000000000000..c32a6e24b6b44 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_join_small_parquet.sql @@ -0,0 +1,7 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_NA_0.parquet'; + +CREATE EXTERNAL TABLE small STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_1e1_0.parquet'; + +CREATE EXTERNAL TABLE medium STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_1e4_0.parquet'; + +CREATE EXTERNAL TABLE large STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_1e7_NA.parquet'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_window_big_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_window_big_csv.sql new file mode 100644 index 0000000000000..e712ef1458d5f --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_window_big_csv.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_1e9_NA.csv'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_window_big_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_window_big_parquet.sql new file mode 100644 index 0000000000000..33d58870a150d --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_window_big_parquet.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e9_1e9_NA.parquet'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_window_medium_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_window_medium_csv.sql new file mode 100644 index 0000000000000..9331c17a4ec89 --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_window_medium_csv.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_1e8_NA.csv'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_window_medium_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_window_medium_parquet.sql new file mode 100644 index 0000000000000..8d4f290b0cdad --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_window_medium_parquet.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e8_1e8_NA.parquet'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_window_small_csv.sql b/benchmarks/sql_benchmarks/h2o/init/load_window_small_csv.sql new file mode 100644 index 0000000000000..1d0ea6992ff0d --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_window_small_csv.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS CSV LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_1e7_NA.csv'; diff --git a/benchmarks/sql_benchmarks/h2o/init/load_window_small_parquet.sql b/benchmarks/sql_benchmarks/h2o/init/load_window_small_parquet.sql new file mode 100644 index 0000000000000..8c16e0b22a7bd --- /dev/null +++ b/benchmarks/sql_benchmarks/h2o/init/load_window_small_parquet.sql @@ -0,0 +1 @@ +CREATE EXTERNAL TABLE x STORED AS PARQUET LOCATION '${DATA_DIR:-data}/h2o/J1_1e7_1e7_NA.parquet'; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..7b9e201c4d6c9 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q01.benchmark @@ -0,0 +1,22 @@ +name Q01 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q1: Very Small Build Side (Dense) +-- Build Side: nation (25 rows) | Probe Side: customer (1.5M rows) +-- density: 1.0, +-- prob_hit: 1.0, +-- build_size: "25", +-- probe_size: "1.5M", +SELECT n_nationkey +FROM nation + JOIN customer ON c_nationkey = n_nationkey; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..2f0ba5ffc5a8c --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q02.benchmark @@ -0,0 +1,24 @@ +name Q02 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q2: Very Small Build Side (Sparse, range < 1024) +-- Build Side: nation (25 rows, range 961) | Probe Side: customer (1.5M rows) +-- density: 0.026, +-- prob_hit: 1.0, +-- build_size: "25", +-- probe_size: "1.5M", +SELECT l.k +FROM (SELECT c_nationkey * 40 as k + FROM customer) l + JOIN (SELECT n_nationkey * 40 as k + FROM nation) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..412b96eb0c54f --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q03.benchmark @@ -0,0 +1,21 @@ +name Q03 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q3: 100% Density, 100% Hit rate +-- density: 1.0, +-- prob_hit: 1.0, +-- build_size: "100K", +-- probe_size: "60M", +SELECT s_suppkey +FROM supplier + JOIN lineitem ON s_suppkey = l_suppkey; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..8cf41b76079b8 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q04.benchmark @@ -0,0 +1,26 @@ +name Q04 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q4: 100% Density, 10% Hit rate +-- density: 1.0, +-- prob_hit: 0.1, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k + FROM ( + SELECT CASE WHEN l_suppkey % 10 = 0 THEN l_suppkey ELSE l_suppkey + 1000000 END as k + FROM lineitem + ) l + JOIN ( + SELECT s_suppkey as k FROM supplier + ) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..7d985a2f8c15c --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q05.benchmark @@ -0,0 +1,23 @@ +name Q05 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q5: 75% Density, 100% Hit rate +-- density: 0.75, +-- prob_hit: 1.0, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k +FROM (SELECT l_suppkey * 4 / 3 as k + FROM lineitem) l + JOIN (SELECT s_suppkey * 4 / 3 as k + FROM supplier) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..5fd1ebf602e37 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q06.benchmark @@ -0,0 +1,30 @@ +name Q06 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q6: 75% Density, 10% Hit rate +-- density: 0.75, +-- prob_hit: 0.1, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k + FROM ( + SELECT CASE + WHEN l_suppkey % 10 = 0 THEN l_suppkey * 4 / 3 + WHEN l_suppkey % 10 < 9 THEN (l_suppkey * 4 / 3 / 4) * 4 + 3 + ELSE l_suppkey * 4 / 3 + 1000000 + END as k + FROM lineitem + ) l + JOIN ( + SELECT s_suppkey * 4 / 3 as k FROM supplier + ) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..a0be4a484af85 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q07.benchmark @@ -0,0 +1,23 @@ +name Q07 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q7: 50% Density, 100% Hit rate +-- density: 0.5, +-- prob_hit: 1.0, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k +FROM (SELECT l_suppkey * 2 as k + FROM lineitem) l + JOIN (SELECT s_suppkey * 2 as k + FROM supplier) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..aa4ac0039fec5 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q08.benchmark @@ -0,0 +1,30 @@ +name Q08 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q8: 50% Density, 10% Hit rate +-- density: 0.5, +-- prob_hit: 0.1, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k + FROM ( + SELECT CASE + WHEN l_suppkey % 10 = 0 THEN l_suppkey * 2 + WHEN l_suppkey % 10 < 9 THEN l_suppkey * 2 + 1 + ELSE l_suppkey * 2 + 1000000 + END as k + FROM lineitem + ) l + JOIN ( + SELECT s_suppkey * 2 as k FROM supplier + ) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..bed67f360cc09 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q09.benchmark @@ -0,0 +1,23 @@ +name Q09 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q9: 20% Density, 100% Hit rate +-- density: 0.2, +-- prob_hit: 1.0, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k +FROM (SELECT l_suppkey * 5 as k + FROM lineitem) l + JOIN (SELECT s_suppkey * 5 as k + FROM supplier) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..857881326b911 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q10.benchmark @@ -0,0 +1,30 @@ +name Q10 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q10: 20% Density, 10% Hit rate +-- density: 0.2, +-- prob_hit: 0.1, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k + FROM ( + SELECT CASE + WHEN l_suppkey % 10 = 0 THEN l_suppkey * 5 + WHEN l_suppkey % 10 < 9 THEN l_suppkey * 5 + 1 + ELSE l_suppkey * 5 + 1000000 + END as k + FROM lineitem + ) l + JOIN ( + SELECT s_suppkey * 5 as k FROM supplier + ) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q11.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q11.benchmark new file mode 100644 index 0000000000000..de241a3601710 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q11.benchmark @@ -0,0 +1,23 @@ +name Q11 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q11: 10% Density, 100% Hit rate +-- density: 0.1, +-- prob_hit: 1.0, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k +FROM (SELECT l_suppkey * 10 as k + FROM lineitem) l + JOIN (SELECT s_suppkey * 10 as k + FROM supplier) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q12.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q12.benchmark new file mode 100644 index 0000000000000..f83e8e94a1f7a --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q12.benchmark @@ -0,0 +1,30 @@ +name Q12 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q12: 10% Density, 10% Hit rate +-- density: 0.1, +-- prob_hit: 0.1, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k + FROM ( + SELECT CASE + WHEN l_suppkey % 10 = 0 THEN l_suppkey * 10 + WHEN l_suppkey % 10 < 9 THEN l_suppkey * 10 + 1 + ELSE l_suppkey * 10 + 1000000 + END as k + FROM lineitem + ) l + JOIN ( + SELECT s_suppkey * 10 as k FROM supplier + ) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q13.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q13.benchmark new file mode 100644 index 0000000000000..d60d3543d0821 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q13.benchmark @@ -0,0 +1,23 @@ +name Q13 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q13: 1% Density, 100% Hit rate +-- density: 0.01, +-- prob_hit: 1.0, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k +FROM (SELECT l_suppkey * 100 as k + FROM lineitem) l + JOIN (SELECT s_suppkey * 100 as k + FROM supplier) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q14.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q14.benchmark new file mode 100644 index 0000000000000..0997bb0a431b5 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q14.benchmark @@ -0,0 +1,31 @@ +name Q14 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q14: 1% Density, 10% Hit rate +-- density: 0.01, +-- prob_hit: 0.1, +-- build_size: "100K", +-- probe_size: "60M", +SELECT l.k + FROM ( + SELECT CASE + WHEN l_suppkey % 10 = 0 THEN l_suppkey * 100 + WHEN l_suppkey % 10 < 9 THEN l_suppkey * 100 + 1 + ELSE l_suppkey * 100 + 11000000 -- oob + END as k + FROM lineitem + ) l + JOIN ( + SELECT s_suppkey * 100 as k FROM supplier + ) s ON l.k = s.k;; + diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q15.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q15.benchmark new file mode 100644 index 0000000000000..413baa420d47e --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q15.benchmark @@ -0,0 +1,33 @@ +name Q15 +group hj + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q15: 20% Density, 10% Hit rate, 20% Duplicates in Build Side +-- density: 0.2, +-- prob_hit: 0.1, +-- build_size: "100K_(20%_dups)", +-- probe_size: "60M", +SELECT l.k + FROM ( + SELECT CASE + WHEN l_suppkey % 10 = 0 THEN ((l_suppkey % 80000) + 1) * 25 / 4 + ELSE ((l_suppkey % 80000) + 1) * 25 / 4 + 1 + END as k + FROM lineitem + ) l + JOIN ( + SELECT CASE + WHEN s_suppkey <= 80000 THEN (s_suppkey * 25) / 4 + ELSE ((s_suppkey - 80000) * 25) / 4 + END as k + FROM supplier + ) s ON l.k = s.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q16.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q16.benchmark new file mode 100644 index 0000000000000..3bc097088e739 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q16.benchmark @@ -0,0 +1,21 @@ +name Q16 +group hj + +init sql_benchmarks/hj/init/set_config.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q16: RightSemi, Small build (25 rows), 100% Hit rate +-- Build Side: nation (25 rows) | Probe Side: customer (1.5M rows) +SELECT c.k +FROM (SELECT CAST(n_nationkey AS INT) as k FROM nation) n +RIGHT SEMI JOIN (SELECT CAST(c_nationkey AS INT) as k FROM customer) c +ON n.k = c.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q17.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q17.benchmark new file mode 100644 index 0000000000000..75604de89aea1 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q17.benchmark @@ -0,0 +1,21 @@ +name Q17 +group hj + +init sql_benchmarks/hj/init/set_config.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM lineitem +---- +true + +expect_plan HashJoinExec + +run +-- Q17: RightSemi, Medium build (100K rows), 100% Hit rate +-- Build Side: supplier (100K rows) | Probe Side: lineitem (60M rows) +SELECT l.k +FROM (SELECT CAST(s_suppkey AS INT) as k FROM supplier) s +RIGHT SEMI JOIN (SELECT CAST(l_suppkey AS INT) as k FROM lineitem) l +ON s.k = l.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q18.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q18.benchmark new file mode 100644 index 0000000000000..e9af18ba7d95f --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q18.benchmark @@ -0,0 +1,24 @@ +name Q18 +group hj + +init sql_benchmarks/hj/init/set_config.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM lineitem +---- +true + +expect_plan HashJoinExec + +run +-- Q18: RightSemi, Medium build (100K rows), 10% Hit rate +-- Build Side: supplier (100K rows) | Probe Side: lineitem (60M rows) +SELECT l.k +FROM (SELECT CAST(s_suppkey AS INT) as k FROM supplier) s +RIGHT SEMI JOIN ( + SELECT CAST(CASE WHEN l_suppkey % 10 = 0 THEN l_suppkey ELSE l_suppkey + 1000000 END AS INT) as k + FROM lineitem +) l +ON s.k = l.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q19.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q19.benchmark new file mode 100644 index 0000000000000..fc70b7bc060c1 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q19.benchmark @@ -0,0 +1,21 @@ +name Q19 +group hj + +init sql_benchmarks/hj/init/set_config.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM customer +---- +true + +expect_plan HashJoinExec + +run +-- Q19: RightAnti, Small build (25 rows), 100% Hit rate (no output) +-- Build Side: nation (25 rows) | Probe Side: customer (1.5M rows) +SELECT c.k +FROM (SELECT CAST(n_nationkey AS INT) as k FROM nation) n +RIGHT ANTI JOIN (SELECT CAST(c_nationkey AS INT) as k FROM customer) c +ON n.k = c.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q20.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q20.benchmark new file mode 100644 index 0000000000000..4fb421f1c0ff8 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q20.benchmark @@ -0,0 +1,21 @@ +name Q20 +group hj + +init sql_benchmarks/hj/init/set_config.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM lineitem +---- +true + +expect_plan HashJoinExec + +run +-- Q20: RightAnti, Medium build (100K rows), 100% Hit rate (no output) +-- Build Side: supplier (100K rows) | Probe Side: lineitem (60M rows) +SELECT l.k +FROM (SELECT CAST(s_suppkey AS INT) as k FROM supplier) s +RIGHT ANTI JOIN (SELECT CAST(l_suppkey AS INT) as k FROM lineitem) l +ON s.k = l.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q21.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q21.benchmark new file mode 100644 index 0000000000000..dae927178f868 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q21.benchmark @@ -0,0 +1,24 @@ +name Q21 +group hj + +init sql_benchmarks/hj/init/set_config.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM lineitem +---- +true + +expect_plan HashJoinExec + +run +-- Q21: RightAnti, Medium build (100K rows), 10% Hit rate (90% output) +-- Build Side: supplier (100K rows) | Probe Side: lineitem (60M rows) +SELECT l.k +FROM (SELECT CAST(s_suppkey AS INT) as k FROM supplier) s +RIGHT ANTI JOIN ( + SELECT CAST(CASE WHEN l_suppkey % 10 = 0 THEN l_suppkey ELSE l_suppkey + 1000000 END AS INT) as k + FROM lineitem +) l +ON s.k = l.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q22.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q22.benchmark new file mode 100644 index 0000000000000..868b88a5aaed2 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q22.benchmark @@ -0,0 +1,28 @@ +name Q22 +group hj + +init sql_benchmarks/hj/init/set_config.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM lineitem +---- +true + +expect_plan HashJoinExec + +run +-- Q22: RightSemi, Medium build (100K rows), ~1% Hit rate, fanout ~100 +-- Build Side: supplier (100K rows) collapsed onto 1K distinct keys +-- Probe Side: lineitem (60M rows) +SELECT l.k +FROM ( + SELECT CAST(((s_suppkey - 1) % 1000) + 1 AS INT) as k + FROM supplier +) s +RIGHT SEMI JOIN ( + SELECT CAST(l_suppkey AS INT) as k + FROM lineitem +) l +ON s.k = l.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q23.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q23.benchmark new file mode 100644 index 0000000000000..7aa8acc87e93b --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q23.benchmark @@ -0,0 +1,31 @@ +name Q23 +group hj + +init sql_benchmarks/hj/init/set_config_no_stats.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM lineitem +---- +true + +expect_plan HashJoinExec + +run +-- Q23: high-fanout string-key inner join. +-- Build ~32K rows / ~415 distinct keys (fanout ~78), probe ~2.3M rows +-- (all carrying the dominant key), output ~176M pairs. Long keys (~28 chars) +-- make per-pair key comparison expensive; count(*) isolates the match path. +-- Thresholds zeroed to force Partitioned mode (simulates absent row-count stats). +SELECT count(*) +FROM ( + SELECT 'high_fanout_string_join_key_' || CAST((s_suppkey % 415) + 1 AS VARCHAR) as k + FROM supplier + WHERE s_suppkey <= 32340 +) s +JOIN ( + SELECT 'high_fanout_string_join_key_1' as k + FROM lineitem + WHERE l_orderkey % 265 = 0 +) l ON s.k = l.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q24.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q24.benchmark new file mode 100644 index 0000000000000..2ea60f0f87009 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q24.benchmark @@ -0,0 +1,31 @@ +name Q24 +group hj + +init sql_benchmarks/hj/init/set_config_no_stats.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM lineitem +---- +true + +expect_plan HashJoinExec + +run +-- Q24: single-hot-bucket long string-key inner join. +-- Build rows all share one long string key, so each matching probe row fans +-- out to the whole build side. count(*) focuses the benchmark on hash match +-- and equality filtering without buffering joined rows. +-- Thresholds zeroed to force Partitioned mode (simulates absent row-count stats). +SELECT count(*) +FROM ( + SELECT 'single_hot_bucket_string_join_key' as k + FROM supplier + WHERE s_suppkey <= 3000 +) s +JOIN ( + SELECT 'single_hot_bucket_string_join_key' as k + FROM lineitem + WHERE l_orderkey % 3000 = 0 +) l ON s.k = l.k; diff --git a/benchmarks/sql_benchmarks/hj/benchmarks/q25.benchmark b/benchmarks/sql_benchmarks/hj/benchmarks/q25.benchmark new file mode 100644 index 0000000000000..b29d6b959a853 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/benchmarks/q25.benchmark @@ -0,0 +1,33 @@ +name Q25 +group hj + +init sql_benchmarks/hj/init/set_config_no_stats.sql + +load sql_benchmarks/hj/init/load.sql + +assert I +SELECT count(*) > 0 FROM lineitem +---- +true + +expect_plan HashJoinExec + +run +-- Q25: skewed high-fanout multi-column string-key inner join. +-- This tracks candidate-pair filtering for composite keys: the first key is +-- skewed and the second long string key must also be checked before emitting +-- each match. count(*) isolates the match path. +-- Thresholds zeroed to force Partitioned mode (simulates absent row-count stats). +SELECT count(*) +FROM ( + SELECT CAST((s_suppkey % 256) + 1 AS INT) as k1, + 'multi_column_high_fanout_key' as k2 + FROM supplier + WHERE s_suppkey <= 20000 +) s +JOIN ( + SELECT CAST(1 AS INT) as k1, + 'multi_column_high_fanout_key' as k2 + FROM lineitem + WHERE l_orderkey % 250 = 0 +) l ON s.k1 = l.k1 AND s.k2 = l.k2; diff --git a/benchmarks/sql_benchmarks/hj/hj.suite b/benchmarks/sql_benchmarks/hj/hj.suite new file mode 100644 index 0000000000000..31ff12a046c59 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/hj.suite @@ -0,0 +1,21 @@ +description = "Hash join SQL benchmarks derived from TPC-H" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "..."] +help = "Selects the TPC-H scale factor used by the hash join benchmarks." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- hj" +description = "Run all hash join queries at scale factor 1." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- hj --query 16 --scale-factor 10" +description = "Run hash join query 16 at scale factor 10." diff --git a/benchmarks/sql_benchmarks/hj/init/load.sql b/benchmarks/sql_benchmarks/hj/init/load.sql new file mode 100644 index 0000000000000..174dac5fbaed5 --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/init/load.sql @@ -0,0 +1,7 @@ +CREATE EXTERNAL TABLE nation STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/nation/nation.1.parquet'; + +CREATE EXTERNAL TABLE supplier STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/supplier/supplier.1.parquet'; + +CREATE EXTERNAL TABLE customer STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/customer/customer.1.parquet'; + +CREATE EXTERNAL TABLE lineitem STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/lineitem/lineitem.1.parquet'; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/hj/init/set_config.sql b/benchmarks/sql_benchmarks/hj/init/set_config.sql new file mode 100644 index 0000000000000..39a3ce259b0ae --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/init/set_config.sql @@ -0,0 +1 @@ +set datafusion.optimizer.join_reordering = false; diff --git a/benchmarks/sql_benchmarks/hj/init/set_config_no_stats.sql b/benchmarks/sql_benchmarks/hj/init/set_config_no_stats.sql new file mode 100644 index 0000000000000..547cbe80bc49f --- /dev/null +++ b/benchmarks/sql_benchmarks/hj/init/set_config_no_stats.sql @@ -0,0 +1,3 @@ +set datafusion.optimizer.join_reordering = false; +set datafusion.optimizer.hash_join_single_partition_threshold = 0; +set datafusion.optimizer.hash_join_single_partition_threshold_rows = 0; diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/01a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/01a.benchmark new file mode 100644 index 0000000000000..1641b348b861a --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/01a.benchmark @@ -0,0 +1,35 @@ +name Q01a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mc.note) AS production_note, + MIN(t.title) AS movie_title, + MIN(t.production_year) AS movie_year +FROM company_type AS ct, + info_type AS it, + movie_companies AS mc, + movie_info_idx AS mi_idx, + title AS t +WHERE ct.kind = 'production companies' + AND it.info = 'top 250 rank' + AND mc.note NOT LIKE '%(as Metro-Goldwyn-Mayer Pictures)%' + AND (mc.note LIKE '%(co-production)%' + OR mc.note LIKE '%(presents)%') + AND ct.id = mc.company_type_id + AND t.id = mc.movie_id + AND t.id = mi_idx.movie_id + AND mc.movie_id = mi_idx.movie_id + AND it.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/01a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/01b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/01b.benchmark new file mode 100644 index 0000000000000..e8515ab3a88e0 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/01b.benchmark @@ -0,0 +1,34 @@ +name Q01b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mc.note) AS production_note, + MIN(t.title) AS movie_title, + MIN(t.production_year) AS movie_year +FROM company_type AS ct, + info_type AS it, + movie_companies AS mc, + movie_info_idx AS mi_idx, + title AS t +WHERE ct.kind = 'production companies' + AND it.info = 'bottom 10 rank' + AND mc.note NOT LIKE '%(as Metro-Goldwyn-Mayer Pictures)%' + AND t.production_year BETWEEN 2005 AND 2010 + AND ct.id = mc.company_type_id + AND t.id = mc.movie_id + AND t.id = mi_idx.movie_id + AND mc.movie_id = mi_idx.movie_id + AND it.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/01b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/01c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/01c.benchmark new file mode 100644 index 0000000000000..fb9711a34fd80 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/01c.benchmark @@ -0,0 +1,35 @@ +name Q01c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mc.note) AS production_note, + MIN(t.title) AS movie_title, + MIN(t.production_year) AS movie_year +FROM company_type AS ct, + info_type AS it, + movie_companies AS mc, + movie_info_idx AS mi_idx, + title AS t +WHERE ct.kind = 'production companies' + AND it.info = 'top 250 rank' + AND mc.note NOT LIKE '%(as Metro-Goldwyn-Mayer Pictures)%' + AND (mc.note LIKE '%(co-production)%') + AND t.production_year >2010 + AND ct.id = mc.company_type_id + AND t.id = mc.movie_id + AND t.id = mi_idx.movie_id + AND mc.movie_id = mi_idx.movie_id + AND it.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/01c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/01d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/01d.benchmark new file mode 100644 index 0000000000000..00dff7d071994 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/01d.benchmark @@ -0,0 +1,34 @@ +name Q01d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mc.note) AS production_note, + MIN(t.title) AS movie_title, + MIN(t.production_year) AS movie_year +FROM company_type AS ct, + info_type AS it, + movie_companies AS mc, + movie_info_idx AS mi_idx, + title AS t +WHERE ct.kind = 'production companies' + AND it.info = 'bottom 10 rank' + AND mc.note NOT LIKE '%(as Metro-Goldwyn-Mayer Pictures)%' + AND t.production_year >2000 + AND ct.id = mc.company_type_id + AND t.id = mc.movie_id + AND t.id = mi_idx.movie_id + AND mc.movie_id = mi_idx.movie_id + AND it.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/01d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/02a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/02a.benchmark new file mode 100644 index 0000000000000..d3455b56a4e17 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/02a.benchmark @@ -0,0 +1,30 @@ +name Q02a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS movie_title +FROM company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + title AS t +WHERE cn.country_code ='[de]' + AND k.keyword ='character-name-in-title' + AND cn.id = mc.company_id + AND mc.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/02a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/02b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/02b.benchmark new file mode 100644 index 0000000000000..b6cf22600adcf --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/02b.benchmark @@ -0,0 +1,30 @@ +name Q02b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS movie_title +FROM company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + title AS t +WHERE cn.country_code ='[nl]' + AND k.keyword ='character-name-in-title' + AND cn.id = mc.company_id + AND mc.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/02b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/02c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/02c.benchmark new file mode 100644 index 0000000000000..b020e9e3cdd87 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/02c.benchmark @@ -0,0 +1,30 @@ +name Q02c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS movie_title +FROM company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + title AS t +WHERE cn.country_code ='[sm]' + AND k.keyword ='character-name-in-title' + AND cn.id = mc.company_id + AND mc.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/02c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/02d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/02d.benchmark new file mode 100644 index 0000000000000..08355454213d4 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/02d.benchmark @@ -0,0 +1,30 @@ +name Q02d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS movie_title +FROM company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + title AS t +WHERE cn.country_code ='[us]' + AND k.keyword ='character-name-in-title' + AND cn.id = mc.company_id + AND mc.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/02d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/03a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/03a.benchmark new file mode 100644 index 0000000000000..22112a2894832 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/03a.benchmark @@ -0,0 +1,36 @@ +name Q03a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS movie_title +FROM keyword AS k, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE k.keyword LIKE '%sequel%' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German') + AND t.production_year > 2005 + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND mk.movie_id = mi.movie_id + AND k.id = mk.keyword_id; + +result sql_benchmarks/imdb/results/03a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/03b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/03b.benchmark new file mode 100644 index 0000000000000..ab24455fd0f4d --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/03b.benchmark @@ -0,0 +1,29 @@ +name Q03b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS movie_title +FROM keyword AS k, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE k.keyword LIKE '%sequel%' + AND mi.info IN ('Bulgaria') + AND t.production_year > 2010 + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND mk.movie_id = mi.movie_id + AND k.id = mk.keyword_id; + +result sql_benchmarks/imdb/results/03b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/03c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/03c.benchmark new file mode 100644 index 0000000000000..65cfe87df168f --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/03c.benchmark @@ -0,0 +1,38 @@ +name Q03c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS movie_title +FROM keyword AS k, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE k.keyword LIKE '%sequel%' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND t.production_year > 1990 + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND mk.movie_id = mi.movie_id + AND k.id = mk.keyword_id; + +result sql_benchmarks/imdb/results/03c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/04a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/04a.benchmark new file mode 100644 index 0000000000000..ff5992501de70 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/04a.benchmark @@ -0,0 +1,33 @@ +name Q04a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi_idx.info) AS rating, + MIN(t.title) AS movie_title +FROM info_type AS it, + keyword AS k, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE it.info ='rating' + AND k.keyword LIKE '%sequel%' + AND mi_idx.info > '5.0' + AND t.production_year > 2005 + AND t.id = mi_idx.movie_id + AND t.id = mk.movie_id + AND mk.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/04a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/04b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/04b.benchmark new file mode 100644 index 0000000000000..fbcbf42aedd42 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/04b.benchmark @@ -0,0 +1,33 @@ +name Q04b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi_idx.info) AS rating, + MIN(t.title) AS movie_title +FROM info_type AS it, + keyword AS k, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE it.info ='rating' + AND k.keyword LIKE '%sequel%' + AND mi_idx.info > '9.0' + AND t.production_year > 2010 + AND t.id = mi_idx.movie_id + AND t.id = mk.movie_id + AND mk.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/04b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/04c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/04c.benchmark new file mode 100644 index 0000000000000..cc0791f6fc993 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/04c.benchmark @@ -0,0 +1,33 @@ +name Q04c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi_idx.info) AS rating, + MIN(t.title) AS movie_title +FROM info_type AS it, + keyword AS k, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE it.info ='rating' + AND k.keyword LIKE '%sequel%' + AND mi_idx.info > '2.0' + AND t.production_year > 1990 + AND t.id = mi_idx.movie_id + AND t.id = mk.movie_id + AND mk.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/04c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/05a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/05a.benchmark new file mode 100644 index 0000000000000..04ea2cb309113 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/05a.benchmark @@ -0,0 +1,40 @@ +name Q05a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS typical_european_movie +FROM company_type AS ct, + info_type AS it, + movie_companies AS mc, + movie_info AS mi, + title AS t +WHERE ct.kind = 'production companies' + AND mc.note LIKE '%(theatrical)%' + AND mc.note LIKE '%(France)%' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German') + AND t.production_year > 2005 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND mc.movie_id = mi.movie_id + AND ct.id = mc.company_type_id + AND it.id = mi.info_type_id; + +result sql_benchmarks/imdb/results/05a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/05b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/05b.benchmark new file mode 100644 index 0000000000000..d2a8011bd86f9 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/05b.benchmark @@ -0,0 +1,35 @@ +name Q05b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS american_vhs_movie +FROM company_type AS ct, + info_type AS it, + movie_companies AS mc, + movie_info AS mi, + title AS t +WHERE ct.kind = 'production companies' + AND mc.note LIKE '%(VHS)%' + AND mc.note LIKE '%(USA)%' + AND mc.note LIKE '%(1994)%' + AND mi.info IN ('USA', + 'America') + AND t.production_year > 2010 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND mc.movie_id = mi.movie_id + AND ct.id = mc.company_type_id + AND it.id = mi.info_type_id; + +result sql_benchmarks/imdb/results/05b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/05c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/05c.benchmark new file mode 100644 index 0000000000000..7467bf826da8c --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/05c.benchmark @@ -0,0 +1,42 @@ +name Q05c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS american_movie +FROM company_type AS ct, + info_type AS it, + movie_companies AS mc, + movie_info AS mi, + title AS t +WHERE ct.kind = 'production companies' + AND mc.note NOT LIKE '%(TV)%' + AND mc.note LIKE '%(USA)%' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND t.production_year > 1990 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND mc.movie_id = mi.movie_id + AND ct.id = mc.company_type_id + AND it.id = mi.info_type_id; + +result sql_benchmarks/imdb/results/05c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/06a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/06a.benchmark new file mode 100644 index 0000000000000..cadd66c86abe1 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/06a.benchmark @@ -0,0 +1,33 @@ +name Q06a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(k.keyword) AS movie_keyword, + MIN(n.name) AS actor_name, + MIN(t.title) AS marvel_movie +FROM cast_info AS ci, + keyword AS k, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword = 'marvel-cinematic-universe' + AND n.name LIKE '%Downey%Robert%' + AND t.production_year > 2010 + AND k.id = mk.keyword_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mk.movie_id + AND n.id = ci.person_id; + +result sql_benchmarks/imdb/results/06a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/06b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/06b.benchmark new file mode 100644 index 0000000000000..08d310baea120 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/06b.benchmark @@ -0,0 +1,40 @@ +name Q06b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(k.keyword) AS movie_keyword, + MIN(n.name) AS actor_name, + MIN(t.title) AS hero_movie +FROM cast_info AS ci, + keyword AS k, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword IN ('superhero', + 'sequel', + 'second-part', + 'marvel-comics', + 'based-on-comic', + 'tv-special', + 'fight', + 'violence') + AND n.name LIKE '%Downey%Robert%' + AND t.production_year > 2014 + AND k.id = mk.keyword_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mk.movie_id + AND n.id = ci.person_id; + +result sql_benchmarks/imdb/results/06b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/06c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/06c.benchmark new file mode 100644 index 0000000000000..125b48c5a3f1d --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/06c.benchmark @@ -0,0 +1,33 @@ +name Q06c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(k.keyword) AS movie_keyword, + MIN(n.name) AS actor_name, + MIN(t.title) AS marvel_movie +FROM cast_info AS ci, + keyword AS k, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword = 'marvel-cinematic-universe' + AND n.name LIKE '%Downey%Robert%' + AND t.production_year > 2014 + AND k.id = mk.keyword_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mk.movie_id + AND n.id = ci.person_id; + +result sql_benchmarks/imdb/results/06c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/06d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/06d.benchmark new file mode 100644 index 0000000000000..0ce0c10b6b032 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/06d.benchmark @@ -0,0 +1,40 @@ +name Q06d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(k.keyword) AS movie_keyword, + MIN(n.name) AS actor_name, + MIN(t.title) AS hero_movie +FROM cast_info AS ci, + keyword AS k, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword IN ('superhero', + 'sequel', + 'second-part', + 'marvel-comics', + 'based-on-comic', + 'tv-special', + 'fight', + 'violence') + AND n.name LIKE '%Downey%Robert%' + AND t.production_year > 2000 + AND k.id = mk.keyword_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mk.movie_id + AND n.id = ci.person_id; + +result sql_benchmarks/imdb/results/06d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/06e.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/06e.benchmark new file mode 100644 index 0000000000000..d6eb6b8a7f0f0 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/06e.benchmark @@ -0,0 +1,33 @@ +name Q06e +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(k.keyword) AS movie_keyword, + MIN(n.name) AS actor_name, + MIN(t.title) AS marvel_movie +FROM cast_info AS ci, + keyword AS k, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword = 'marvel-cinematic-universe' + AND n.name LIKE '%Downey%Robert%' + AND t.production_year > 2000 + AND k.id = mk.keyword_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mk.movie_id + AND n.id = ci.person_id; + +result sql_benchmarks/imdb/results/06e.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/06f.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/06f.benchmark new file mode 100644 index 0000000000000..8387633632e3c --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/06f.benchmark @@ -0,0 +1,39 @@ +name Q06f +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(k.keyword) AS movie_keyword, + MIN(n.name) AS actor_name, + MIN(t.title) AS hero_movie +FROM cast_info AS ci, + keyword AS k, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword IN ('superhero', + 'sequel', + 'second-part', + 'marvel-comics', + 'based-on-comic', + 'tv-special', + 'fight', + 'violence') + AND t.production_year > 2000 + AND k.id = mk.keyword_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mk.movie_id + AND n.id = ci.person_id; + +result sql_benchmarks/imdb/results/06f.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/07a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/07a.benchmark new file mode 100644 index 0000000000000..1ad5388cc28be --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/07a.benchmark @@ -0,0 +1,47 @@ +name Q07a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS of_person, + MIN(t.title) AS biography_movie +FROM aka_name AS an, + cast_info AS ci, + info_type AS it, + link_type AS lt, + movie_link AS ml, + name AS n, + person_info AS pi, + title AS t +WHERE an.name LIKE '%a%' + AND it.info ='mini biography' + AND lt.link ='features' + AND n.name_pcode_cf BETWEEN 'A' AND 'F' + AND (n.gender='m' + OR (n.gender = 'f' + AND n.name LIKE 'B%')) + AND pi.note ='Volker Boehm' + AND t.production_year BETWEEN 1980 AND 1995 + AND n.id = an.person_id + AND n.id = pi.person_id + AND ci.person_id = n.id + AND t.id = ci.movie_id + AND ml.linked_movie_id = t.id + AND lt.id = ml.link_type_id + AND it.id = pi.info_type_id + AND pi.person_id = an.person_id + AND pi.person_id = ci.person_id + AND an.person_id = ci.person_id + AND ci.movie_id = ml.linked_movie_id; + +result sql_benchmarks/imdb/results/07a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/07b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/07b.benchmark new file mode 100644 index 0000000000000..bfc2e107a99df --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/07b.benchmark @@ -0,0 +1,45 @@ +name Q07b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS of_person, + MIN(t.title) AS biography_movie +FROM aka_name AS an, + cast_info AS ci, + info_type AS it, + link_type AS lt, + movie_link AS ml, + name AS n, + person_info AS pi, + title AS t +WHERE an.name LIKE '%a%' + AND it.info ='mini biography' + AND lt.link ='features' + AND n.name_pcode_cf LIKE 'D%' + AND n.gender='m' + AND pi.note ='Volker Boehm' + AND t.production_year BETWEEN 1980 AND 1984 + AND n.id = an.person_id + AND n.id = pi.person_id + AND ci.person_id = n.id + AND t.id = ci.movie_id + AND ml.linked_movie_id = t.id + AND lt.id = ml.link_type_id + AND it.id = pi.info_type_id + AND pi.person_id = an.person_id + AND pi.person_id = ci.person_id + AND an.person_id = ci.person_id + AND ci.movie_id = ml.linked_movie_id; + +result sql_benchmarks/imdb/results/07b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/07c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/07c.benchmark new file mode 100644 index 0000000000000..449df56c14d89 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/07c.benchmark @@ -0,0 +1,52 @@ +name Q07c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS cast_member_name, + MIN(pi.info) AS cast_member_info +FROM aka_name AS an, + cast_info AS ci, + info_type AS it, + link_type AS lt, + movie_link AS ml, + name AS n, + person_info AS pi, + title AS t +WHERE an.name IS NOT NULL + AND (an.name LIKE '%a%' + OR an.name LIKE 'A%') + AND it.info ='mini biography' + AND lt.link IN ('references', + 'referenced in', + 'features', + 'featured in') + AND n.name_pcode_cf BETWEEN 'A' AND 'F' + AND (n.gender='m' + OR (n.gender = 'f' + AND n.name LIKE 'A%')) + AND pi.note IS NOT NULL + AND t.production_year BETWEEN 1980 AND 2010 + AND n.id = an.person_id + AND n.id = pi.person_id + AND ci.person_id = n.id + AND t.id = ci.movie_id + AND ml.linked_movie_id = t.id + AND lt.id = ml.link_type_id + AND it.id = pi.info_type_id + AND pi.person_id = an.person_id + AND pi.person_id = ci.person_id + AND an.person_id = ci.person_id + AND ci.movie_id = ml.linked_movie_id; + +result sql_benchmarks/imdb/results/07c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/08a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/08a.benchmark new file mode 100644 index 0000000000000..72914b32c326b --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/08a.benchmark @@ -0,0 +1,41 @@ +name Q08a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an1.name) AS actress_pseudonym, + MIN(t.title) AS japanese_movie_dubbed +FROM aka_name AS an1, + cast_info AS ci, + company_name AS cn, + movie_companies AS mc, + name AS n1, + role_type AS rt, + title AS t +WHERE ci.note ='(voice: English version)' + AND cn.country_code ='[jp]' + AND mc.note LIKE '%(Japan)%' + AND mc.note NOT LIKE '%(USA)%' + AND n1.name LIKE '%Yo%' + AND n1.name NOT LIKE '%Yu%' + AND rt.role ='actress' + AND an1.person_id = n1.id + AND n1.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.role_id = rt.id + AND an1.person_id = ci.person_id + AND ci.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/08a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/08b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/08b.benchmark new file mode 100644 index 0000000000000..a66486d16de24 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/08b.benchmark @@ -0,0 +1,46 @@ +name Q08b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS acress_pseudonym, + MIN(t.title) AS japanese_anime_movie +FROM aka_name AS an, + cast_info AS ci, + company_name AS cn, + movie_companies AS mc, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note ='(voice: English version)' + AND cn.country_code ='[jp]' + AND mc.note LIKE '%(Japan)%' + AND mc.note NOT LIKE '%(USA)%' + AND (mc.note LIKE '%(2006)%' + OR mc.note LIKE '%(2007)%') + AND n.name LIKE '%Yo%' + AND n.name NOT LIKE '%Yu%' + AND rt.role ='actress' + AND t.production_year BETWEEN 2006 AND 2007 + AND (t.title LIKE 'One Piece%' + OR t.title LIKE 'Dragon Ball Z%') + AND an.person_id = n.id + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.role_id = rt.id + AND an.person_id = ci.person_id + AND ci.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/08b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/08c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/08c.benchmark new file mode 100644 index 0000000000000..116a9c9f60bd3 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/08c.benchmark @@ -0,0 +1,36 @@ +name Q08c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(a1.name) AS writer_pseudo_name, + MIN(t.title) AS movie_title +FROM aka_name AS a1, + cast_info AS ci, + company_name AS cn, + movie_companies AS mc, + name AS n1, + role_type AS rt, + title AS t +WHERE cn.country_code ='[us]' + AND rt.role ='writer' + AND a1.person_id = n1.id + AND n1.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.role_id = rt.id + AND a1.person_id = ci.person_id + AND ci.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/08c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/08d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/08d.benchmark new file mode 100644 index 0000000000000..def2f26b3db4c --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/08d.benchmark @@ -0,0 +1,36 @@ +name Q08d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an1.name) AS costume_designer_pseudo, + MIN(t.title) AS movie_with_costumes +FROM aka_name AS an1, + cast_info AS ci, + company_name AS cn, + movie_companies AS mc, + name AS n1, + role_type AS rt, + title AS t +WHERE cn.country_code ='[us]' + AND rt.role ='costume designer' + AND an1.person_id = n1.id + AND n1.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.role_id = rt.id + AND an1.person_id = ci.person_id + AND ci.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/08d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/09a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/09a.benchmark new file mode 100644 index 0000000000000..7cb040bc6dbca --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/09a.benchmark @@ -0,0 +1,49 @@ +name Q09a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS alternative_name, + MIN(chn.name) AS character_name, + MIN(t.title) AS movie +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + movie_companies AS mc, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND mc.note IS NOT NULL + AND (mc.note LIKE '%(USA)%' + OR mc.note LIKE '%(worldwide)%') + AND n.gender ='f' + AND n.name LIKE '%Ang%' + AND rt.role ='actress' + AND t.production_year BETWEEN 2005 AND 2015 + AND ci.movie_id = t.id + AND t.id = mc.movie_id + AND ci.movie_id = mc.movie_id + AND mc.company_id = cn.id + AND ci.role_id = rt.id + AND n.id = ci.person_id + AND chn.id = ci.person_role_id + AND an.person_id = n.id + AND an.person_id = ci.person_id; + +result sql_benchmarks/imdb/results/09a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/09b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/09b.benchmark new file mode 100644 index 0000000000000..a3b7f1e200225 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/09b.benchmark @@ -0,0 +1,47 @@ +name Q09b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS alternative_name, + MIN(chn.name) AS voiced_character, + MIN(n.name) AS voicing_actress, + MIN(t.title) AS american_movie +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + movie_companies AS mc, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note = '(voice)' + AND cn.country_code ='[us]' + AND mc.note LIKE '%(200%)%' + AND (mc.note LIKE '%(USA)%' + OR mc.note LIKE '%(worldwide)%') + AND n.gender ='f' + AND n.name LIKE '%Angel%' + AND rt.role ='actress' + AND t.production_year BETWEEN 2007 AND 2010 + AND ci.movie_id = t.id + AND t.id = mc.movie_id + AND ci.movie_id = mc.movie_id + AND mc.company_id = cn.id + AND ci.role_id = rt.id + AND n.id = ci.person_id + AND chn.id = ci.person_role_id + AND an.person_id = n.id + AND an.person_id = ci.person_id; + +result sql_benchmarks/imdb/results/09b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/09c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/09c.benchmark new file mode 100644 index 0000000000000..1588622de447f --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/09c.benchmark @@ -0,0 +1,46 @@ +name Q09c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS alternative_name, + MIN(chn.name) AS voiced_character_name, + MIN(n.name) AS voicing_actress, + MIN(t.title) AS american_movie +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + movie_companies AS mc, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND n.gender ='f' + AND n.name LIKE '%An%' + AND rt.role ='actress' + AND ci.movie_id = t.id + AND t.id = mc.movie_id + AND ci.movie_id = mc.movie_id + AND mc.company_id = cn.id + AND ci.role_id = rt.id + AND n.id = ci.person_id + AND chn.id = ci.person_role_id + AND an.person_id = n.id + AND an.person_id = ci.person_id; + +result sql_benchmarks/imdb/results/09c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/09d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/09d.benchmark new file mode 100644 index 0000000000000..959a61c3b6d21 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/09d.benchmark @@ -0,0 +1,45 @@ +name Q09d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS alternative_name, + MIN(chn.name) AS voiced_char_name, + MIN(n.name) AS voicing_actress, + MIN(t.title) AS american_movie +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + movie_companies AS mc, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND n.gender ='f' + AND rt.role ='actress' + AND ci.movie_id = t.id + AND t.id = mc.movie_id + AND ci.movie_id = mc.movie_id + AND mc.company_id = cn.id + AND ci.role_id = rt.id + AND n.id = ci.person_id + AND chn.id = ci.person_role_id + AND an.person_id = n.id + AND an.person_id = ci.person_id; + +result sql_benchmarks/imdb/results/09d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/10a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/10a.benchmark new file mode 100644 index 0000000000000..ba58639156680 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/10a.benchmark @@ -0,0 +1,38 @@ +name Q10a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS uncredited_voiced_character, + MIN(t.title) AS russian_movie +FROM char_name AS chn, + cast_info AS ci, + company_name AS cn, + company_type AS ct, + movie_companies AS mc, + role_type AS rt, + title AS t +WHERE ci.note LIKE '%(voice)%' + AND ci.note LIKE '%(uncredited)%' + AND cn.country_code = '[ru]' + AND rt.role = 'actor' + AND t.production_year > 2005 + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mc.movie_id + AND chn.id = ci.person_role_id + AND rt.id = ci.role_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id; + +result sql_benchmarks/imdb/results/10a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/10b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/10b.benchmark new file mode 100644 index 0000000000000..1947b640b3c86 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/10b.benchmark @@ -0,0 +1,37 @@ +name Q10b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS character, + MIN(t.title) AS russian_mov_with_actor_producer +FROM char_name AS chn, + cast_info AS ci, + company_name AS cn, + company_type AS ct, + movie_companies AS mc, + role_type AS rt, + title AS t +WHERE ci.note LIKE '%(producer)%' + AND cn.country_code = '[ru]' + AND rt.role = 'actor' + AND t.production_year > 2010 + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mc.movie_id + AND chn.id = ci.person_role_id + AND rt.id = ci.role_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id; + +result sql_benchmarks/imdb/results/10b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/10c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/10c.benchmark new file mode 100644 index 0000000000000..2fb881324b620 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/10c.benchmark @@ -0,0 +1,36 @@ +name Q10c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS character, + MIN(t.title) AS movie_with_american_producer +FROM char_name AS chn, + cast_info AS ci, + company_name AS cn, + company_type AS ct, + movie_companies AS mc, + role_type AS rt, + title AS t +WHERE ci.note LIKE '%(producer)%' + AND cn.country_code = '[us]' + AND t.production_year > 1990 + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mc.movie_id + AND chn.id = ci.person_role_id + AND rt.id = ci.role_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id; + +result sql_benchmarks/imdb/results/10c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/11a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/11a.benchmark new file mode 100644 index 0000000000000..d24bc35146bec --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/11a.benchmark @@ -0,0 +1,46 @@ +name Q11a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS from_company, + MIN(lt.link) AS movie_link_type, + MIN(t.title) AS non_polish_sequel_movie +FROM company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cn.country_code !='[pl]' + AND (cn.name LIKE '%Film%' + OR cn.name LIKE '%Warner%') + AND ct.kind ='production companies' + AND k.keyword ='sequel' + AND lt.link LIKE '%follow%' + AND mc.note IS NULL + AND t.production_year BETWEEN 1950 AND 2000 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/11a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/11b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/11b.benchmark new file mode 100644 index 0000000000000..e2dd4cafba597 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/11b.benchmark @@ -0,0 +1,47 @@ +name Q11b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS from_company, + MIN(lt.link) AS movie_link_type, + MIN(t.title) AS sequel_movie +FROM company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cn.country_code !='[pl]' + AND (cn.name LIKE '%Film%' + OR cn.name LIKE '%Warner%') + AND ct.kind ='production companies' + AND k.keyword ='sequel' + AND lt.link LIKE '%follows%' + AND mc.note IS NULL + AND t.production_year = 1998 + AND t.title LIKE '%Money%' + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/11b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/11c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/11c.benchmark new file mode 100644 index 0000000000000..9fde2824afc8a --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/11c.benchmark @@ -0,0 +1,48 @@ +name Q11c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS from_company, + MIN(mc.note) AS production_note, + MIN(t.title) AS movie_based_on_book +FROM company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cn.country_code !='[pl]' + AND (cn.name LIKE '20th Century Fox%' + OR cn.name LIKE 'Twentieth Century Fox%') + AND ct.kind != 'production companies' + AND ct.kind IS NOT NULL + AND k.keyword IN ('sequel', + 'revenge', + 'based-on-novel') + AND mc.note IS NOT NULL + AND t.production_year > 1950 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/11c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/11d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/11d.benchmark new file mode 100644 index 0000000000000..c66a6d5ee04da --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/11d.benchmark @@ -0,0 +1,46 @@ +name Q11d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS from_company, + MIN(mc.note) AS production_note, + MIN(t.title) AS movie_based_on_book +FROM company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cn.country_code !='[pl]' + AND ct.kind != 'production companies' + AND ct.kind IS NOT NULL + AND k.keyword IN ('sequel', + 'revenge', + 'based-on-novel') + AND mc.note IS NOT NULL + AND t.production_year > 1950 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/11d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/12a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/12a.benchmark new file mode 100644 index 0000000000000..53cb5fe7705c9 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/12a.benchmark @@ -0,0 +1,46 @@ +name Q12a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS drama_horror_movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + title AS t +WHERE cn.country_code = '[us]' + AND ct.kind = 'production companies' + AND it1.info = 'genres' + AND it2.info = 'rating' + AND mi.info IN ('Drama', + 'Horror') + AND mi_idx.info > '8.0' + AND t.production_year BETWEEN 2005 AND 2008 + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND mi.info_type_id = it1.id + AND mi_idx.info_type_id = it2.id + AND t.id = mc.movie_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id + AND mc.movie_id = mi.movie_id + AND mc.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id; + +result sql_benchmarks/imdb/results/12a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/12b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/12b.benchmark new file mode 100644 index 0000000000000..02d76f9192ec0 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/12b.benchmark @@ -0,0 +1,46 @@ +name Q12b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS budget, + MIN(t.title) AS unsuccsessful_movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + title AS t +WHERE cn.country_code ='[us]' + AND ct.kind IS NOT NULL + AND (ct.kind ='production companies' + OR ct.kind = 'distributors') + AND it1.info ='budget' + AND it2.info ='bottom 10 rank' + AND t.production_year >2000 + AND (t.title LIKE 'Birdemic%' + OR t.title LIKE '%Movie%') + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND mi.info_type_id = it1.id + AND mi_idx.info_type_id = it2.id + AND t.id = mc.movie_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id + AND mc.movie_id = mi.movie_id + AND mc.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id; + +result sql_benchmarks/imdb/results/12b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/12c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/12c.benchmark new file mode 100644 index 0000000000000..f104486194943 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/12c.benchmark @@ -0,0 +1,48 @@ +name Q12c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS mainstream_movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + title AS t +WHERE cn.country_code = '[us]' + AND ct.kind = 'production companies' + AND it1.info = 'genres' + AND it2.info = 'rating' + AND mi.info IN ('Drama', + 'Horror', + 'Western', + 'Family') + AND mi_idx.info > '7.0' + AND t.production_year BETWEEN 2000 AND 2010 + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND mi.info_type_id = it1.id + AND mi_idx.info_type_id = it2.id + AND t.id = mc.movie_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id + AND mc.movie_id = mi.movie_id + AND mc.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id; + +result sql_benchmarks/imdb/results/12c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/13a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/13a.benchmark new file mode 100644 index 0000000000000..60f65978022cf --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/13a.benchmark @@ -0,0 +1,45 @@ +name Q13a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS release_date, + MIN(miidx.info) AS rating, + MIN(t.title) AS german_movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it, + info_type AS it2, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS miidx, + title AS t +WHERE cn.country_code ='[de]' + AND ct.kind ='production companies' + AND it.info ='rating' + AND it2.info ='release dates' + AND kt.kind ='movie' + AND mi.movie_id = t.id + AND it2.id = mi.info_type_id + AND kt.id = t.kind_id + AND mc.movie_id = t.id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id + AND miidx.movie_id = t.id + AND it.id = miidx.info_type_id + AND mi.movie_id = miidx.movie_id + AND mi.movie_id = mc.movie_id + AND miidx.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/13a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/13b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/13b.benchmark new file mode 100644 index 0000000000000..fbd016322bdac --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/13b.benchmark @@ -0,0 +1,48 @@ +name Q13b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS producing_company, + MIN(miidx.info) AS rating, + MIN(t.title) AS movie_about_winning +FROM company_name AS cn, + company_type AS ct, + info_type AS it, + info_type AS it2, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS miidx, + title AS t +WHERE cn.country_code ='[us]' + AND ct.kind ='production companies' + AND it.info ='rating' + AND it2.info ='release dates' + AND kt.kind ='movie' + AND t.title != '' + AND (t.title LIKE '%Champion%' + OR t.title LIKE '%Loser%') + AND mi.movie_id = t.id + AND it2.id = mi.info_type_id + AND kt.id = t.kind_id + AND mc.movie_id = t.id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id + AND miidx.movie_id = t.id + AND it.id = miidx.info_type_id + AND mi.movie_id = miidx.movie_id + AND mi.movie_id = mc.movie_id + AND miidx.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/13b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/13c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/13c.benchmark new file mode 100644 index 0000000000000..b053b9eba9543 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/13c.benchmark @@ -0,0 +1,48 @@ +name Q13c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS producing_company, + MIN(miidx.info) AS rating, + MIN(t.title) AS movie_about_winning +FROM company_name AS cn, + company_type AS ct, + info_type AS it, + info_type AS it2, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS miidx, + title AS t +WHERE cn.country_code ='[us]' + AND ct.kind ='production companies' + AND it.info ='rating' + AND it2.info ='release dates' + AND kt.kind ='movie' + AND t.title != '' + AND (t.title LIKE 'Champion%' + OR t.title LIKE 'Loser%') + AND mi.movie_id = t.id + AND it2.id = mi.info_type_id + AND kt.id = t.kind_id + AND mc.movie_id = t.id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id + AND miidx.movie_id = t.id + AND it.id = miidx.info_type_id + AND mi.movie_id = miidx.movie_id + AND mi.movie_id = mc.movie_id + AND miidx.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/13c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/13d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/13d.benchmark new file mode 100644 index 0000000000000..f9807dcb2cde1 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/13d.benchmark @@ -0,0 +1,45 @@ +name Q13d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS producing_company, + MIN(miidx.info) AS rating, + MIN(t.title) AS movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it, + info_type AS it2, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS miidx, + title AS t +WHERE cn.country_code ='[us]' + AND ct.kind ='production companies' + AND it.info ='rating' + AND it2.info ='release dates' + AND kt.kind ='movie' + AND mi.movie_id = t.id + AND it2.id = mi.info_type_id + AND kt.id = t.kind_id + AND mc.movie_id = t.id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id + AND miidx.movie_id = t.id + AND it.id = miidx.info_type_id + AND mi.movie_id = miidx.movie_id + AND mi.movie_id = mc.movie_id + AND miidx.movie_id = mc.movie_id; + +result sql_benchmarks/imdb/results/13d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/14a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/14a.benchmark new file mode 100644 index 0000000000000..7b9fc0fb20796 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/14a.benchmark @@ -0,0 +1,56 @@ +name Q14a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi_idx.info) AS rating, + MIN(t.title) AS northern_dark_movie +FROM info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind = 'movie' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND mi_idx.info < '8.5' + AND t.production_year > 2010 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/14a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/14b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/14b.benchmark new file mode 100644 index 0000000000000..b843cbd341a25 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/14b.benchmark @@ -0,0 +1,57 @@ +name Q14b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi_idx.info) AS rating, + MIN(t.title) AS western_dark_production +FROM info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title') + AND kt.kind = 'movie' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND mi_idx.info > '6.0' + AND t.production_year > 2010 + AND (t.title LIKE '%murder%' + OR t.title LIKE '%Murder%' + OR t.title LIKE '%Mord%') + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/14b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/14c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/14c.benchmark new file mode 100644 index 0000000000000..2ea8cb2d3843a --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/14c.benchmark @@ -0,0 +1,58 @@ +name Q14c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi_idx.info) AS rating, + MIN(t.title) AS north_european_dark_production +FROM info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IS NOT NULL + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind IN ('movie', + 'episode') + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Danish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND mi_idx.info < '8.5' + AND t.production_year > 2005 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/14c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/15a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/15a.benchmark new file mode 100644 index 0000000000000..47999ab30df78 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/15a.benchmark @@ -0,0 +1,49 @@ +name Q15a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS release_date, + MIN(t.title) AS internet_movie +FROM aka_title AS at_, + company_name AS cn, + company_type AS ct, + info_type AS it1, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE cn.country_code = '[us]' + AND it1.info = 'release dates' + AND mc.note LIKE '%(200%)%' + AND mc.note LIKE '%(worldwide)%' + AND mi.note LIKE '%internet%' + AND mi.info LIKE 'USA:% 200%' + AND t.production_year > 2000 + AND t.id = at_.movie_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = at_.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = at_.movie_id + AND mc.movie_id = at_.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id; + +result sql_benchmarks/imdb/results/15a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/15b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/15b.benchmark new file mode 100644 index 0000000000000..ec90b379fe5d2 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/15b.benchmark @@ -0,0 +1,50 @@ +name Q15b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS release_date, + MIN(t.title) AS youtube_movie +FROM aka_title AS at_, + company_name AS cn, + company_type AS ct, + info_type AS it1, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE cn.country_code = '[us]' + AND cn.name = 'YouTube' + AND it1.info = 'release dates' + AND mc.note LIKE '%(200%)%' + AND mc.note LIKE '%(worldwide)%' + AND mi.note LIKE '%internet%' + AND mi.info LIKE 'USA:% 200%' + AND t.production_year BETWEEN 2005 AND 2010 + AND t.id = at_.movie_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = at_.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = at_.movie_id + AND mc.movie_id = at_.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id; + +result sql_benchmarks/imdb/results/15b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/15c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/15c.benchmark new file mode 100644 index 0000000000000..a9e134520389f --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/15c.benchmark @@ -0,0 +1,49 @@ +name Q15c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS release_date, + MIN(t.title) AS modern_american_internet_movie +FROM aka_title AS at_, + company_name AS cn, + company_type AS ct, + info_type AS it1, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE cn.country_code = '[us]' + AND it1.info = 'release dates' + AND mi.note LIKE '%internet%' + AND mi.info IS NOT NULL + AND (mi.info LIKE 'USA:% 199%' + OR mi.info LIKE 'USA:% 200%') + AND t.production_year > 1990 + AND t.id = at_.movie_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = at_.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = at_.movie_id + AND mc.movie_id = at_.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id; + +result sql_benchmarks/imdb/results/15c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/15d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/15d.benchmark new file mode 100644 index 0000000000000..7f51437509651 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/15d.benchmark @@ -0,0 +1,46 @@ +name Q15d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(at_.title) AS aka_title, + MIN(t.title) AS internet_movie_title +FROM aka_title AS at_, + company_name AS cn, + company_type AS ct, + info_type AS it1, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE cn.country_code = '[us]' + AND it1.info = 'release dates' + AND mi.note LIKE '%internet%' + AND t.production_year > 1990 + AND t.id = at_.movie_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = at_.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = at_.movie_id + AND mc.movie_id = at_.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id; + +result sql_benchmarks/imdb/results/15d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/16a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/16a.benchmark new file mode 100644 index 0000000000000..dd440026a5f91 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/16a.benchmark @@ -0,0 +1,42 @@ +name Q16a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS cool_actor_pseudonym, + MIN(t.title) AS series_named_after_char +FROM aka_name AS an, + cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cn.country_code ='[us]' + AND k.keyword ='character-name-in-title' + AND t.episode_nr >= 50 + AND t.episode_nr < 100 + AND an.person_id = n.id + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND an.person_id = ci.person_id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/16a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/16b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/16b.benchmark new file mode 100644 index 0000000000000..7fade8228fc1a --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/16b.benchmark @@ -0,0 +1,40 @@ +name Q16b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS cool_actor_pseudonym, + MIN(t.title) AS series_named_after_char +FROM aka_name AS an, + cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cn.country_code ='[us]' + AND k.keyword ='character-name-in-title' + AND an.person_id = n.id + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND an.person_id = ci.person_id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/16b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/16c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/16c.benchmark new file mode 100644 index 0000000000000..d1ea1f6f04b14 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/16c.benchmark @@ -0,0 +1,41 @@ +name Q16c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS cool_actor_pseudonym, + MIN(t.title) AS series_named_after_char +FROM aka_name AS an, + cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cn.country_code ='[us]' + AND k.keyword ='character-name-in-title' + AND t.episode_nr < 100 + AND an.person_id = n.id + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND an.person_id = ci.person_id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/16c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/16d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/16d.benchmark new file mode 100644 index 0000000000000..7622fc980d632 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/16d.benchmark @@ -0,0 +1,42 @@ +name Q16d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(an.name) AS cool_actor_pseudonym, + MIN(t.title) AS series_named_after_char +FROM aka_name AS an, + cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cn.country_code ='[us]' + AND k.keyword ='character-name-in-title' + AND t.episode_nr >= 5 + AND t.episode_nr < 100 + AND an.person_id = n.id + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND an.person_id = ci.person_id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/16d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/17a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/17a.benchmark new file mode 100644 index 0000000000000..3bf51dc255dde --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/17a.benchmark @@ -0,0 +1,38 @@ +name Q17a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS member_in_charnamed_american_movie, + MIN(n.name) AS a1 +FROM cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cn.country_code ='[us]' + AND k.keyword ='character-name-in-title' + AND n.name LIKE 'B%' + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/17a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/17b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/17b.benchmark new file mode 100644 index 0000000000000..abe492623a76e --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/17b.benchmark @@ -0,0 +1,37 @@ +name Q17b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS member_in_charnamed_movie, + MIN(n.name) AS a1 +FROM cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword ='character-name-in-title' + AND n.name LIKE 'Z%' + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/17b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/17c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/17c.benchmark new file mode 100644 index 0000000000000..83561d72f194e --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/17c.benchmark @@ -0,0 +1,37 @@ +name Q17c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS member_in_charnamed_movie, + MIN(n.name) AS a1 +FROM cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword ='character-name-in-title' + AND n.name LIKE 'X%' + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/17c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/17d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/17d.benchmark new file mode 100644 index 0000000000000..d7df85a5b68fb --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/17d.benchmark @@ -0,0 +1,36 @@ +name Q17d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS member_in_charnamed_movie +FROM cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword ='character-name-in-title' + AND n.name LIKE '%Bert%' + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/17d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/17e.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/17e.benchmark new file mode 100644 index 0000000000000..b05b5e1cd1a2c --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/17e.benchmark @@ -0,0 +1,36 @@ +name Q17e +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS member_in_charnamed_movie +FROM cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cn.country_code ='[us]' + AND k.keyword ='character-name-in-title' + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/17e.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/17f.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/17f.benchmark new file mode 100644 index 0000000000000..4feef0a7f8ac8 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/17f.benchmark @@ -0,0 +1,36 @@ +name Q17f +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS member_in_charnamed_movie +FROM cast_info AS ci, + company_name AS cn, + keyword AS k, + movie_companies AS mc, + movie_keyword AS mk, + name AS n, + title AS t +WHERE k.keyword ='character-name-in-title' + AND n.name LIKE '%B%' + AND n.id = ci.person_id + AND ci.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_id = cn.id + AND ci.movie_id = mc.movie_id + AND ci.movie_id = mk.movie_id + AND mc.movie_id = mk.movie_id; + +result sql_benchmarks/imdb/results/17f.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/18a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/18a.benchmark new file mode 100644 index 0000000000000..c3e5309e00e50 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/18a.benchmark @@ -0,0 +1,42 @@ +name Q18a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(t.title) AS movie_title +FROM cast_info AS ci, + info_type AS it1, + info_type AS it2, + movie_info AS mi, + movie_info_idx AS mi_idx, + name AS n, + title AS t +WHERE ci.note IN ('(producer)', + '(executive producer)') + AND it1.info = 'budget' + AND it2.info = 'votes' + AND n.gender = 'm' + AND n.name LIKE '%Tim%' + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/18a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/18b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/18b.benchmark new file mode 100644 index 0000000000000..d527cb39858ed --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/18b.benchmark @@ -0,0 +1,50 @@ +name Q18b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(t.title) AS movie_title +FROM cast_info AS ci, + info_type AS it1, + info_type AS it2, + movie_info AS mi, + movie_info_idx AS mi_idx, + name AS n, + title AS t +WHERE ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND it1.info = 'genres' + AND it2.info = 'rating' + AND mi.info IN ('Horror', + 'Thriller') + AND mi.note IS NULL + AND mi_idx.info > '8.0' + AND n.gender IS NOT NULL + AND n.gender = 'f' + AND t.production_year BETWEEN 2008 AND 2014 + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/18b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/18c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/18c.benchmark new file mode 100644 index 0000000000000..30aeff6153497 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/18c.benchmark @@ -0,0 +1,50 @@ +name Q18c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(t.title) AS movie_title +FROM cast_info AS ci, + info_type AS it1, + info_type AS it2, + movie_info AS mi, + movie_info_idx AS mi_idx, + name AS n, + title AS t +WHERE ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND it1.info = 'genres' + AND it2.info = 'votes' + AND mi.info IN ('Horror', + 'Action', + 'Sci-Fi', + 'Thriller', + 'Crime', + 'War') + AND n.gender = 'm' + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND mi.movie_id = mi_idx.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/18c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/19a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/19a.benchmark new file mode 100644 index 0000000000000..eef6a7cdecf4e --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/19a.benchmark @@ -0,0 +1,58 @@ +name Q19a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS voicing_actress, + MIN(t.title) AS voiced_movie +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + movie_companies AS mc, + movie_info AS mi, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND it.info = 'release dates' + AND mc.note IS NOT NULL + AND (mc.note LIKE '%(USA)%' + OR mc.note LIKE '%(worldwide)%') + AND mi.info IS NOT NULL + AND (mi.info LIKE 'Japan:%200%' + OR mi.info LIKE 'USA:%200%') + AND n.gender ='f' + AND n.name LIKE '%Ang%' + AND rt.role ='actress' + AND t.production_year BETWEEN 2005 AND 2009 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mi.movie_id = ci.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id; + +result sql_benchmarks/imdb/results/19a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/19b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/19b.benchmark new file mode 100644 index 0000000000000..49a29d2646c75 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/19b.benchmark @@ -0,0 +1,56 @@ +name Q19b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS voicing_actress, + MIN(t.title) AS kung_fu_panda +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + movie_companies AS mc, + movie_info AS mi, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note = '(voice)' + AND cn.country_code ='[us]' + AND it.info = 'release dates' + AND mc.note LIKE '%(200%)%' + AND (mc.note LIKE '%(USA)%' + OR mc.note LIKE '%(worldwide)%') + AND mi.info IS NOT NULL + AND (mi.info LIKE 'Japan:%2007%' + OR mi.info LIKE 'USA:%2008%') + AND n.gender ='f' + AND n.name LIKE '%Angel%' + AND rt.role ='actress' + AND t.production_year BETWEEN 2007 AND 2008 + AND t.title LIKE '%Kung%Fu%Panda%' + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mi.movie_id = ci.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id; + +result sql_benchmarks/imdb/results/19b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/19c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/19c.benchmark new file mode 100644 index 0000000000000..1c13abf5fbbe1 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/19c.benchmark @@ -0,0 +1,55 @@ +name Q19c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS voicing_actress, + MIN(t.title) AS jap_engl_voiced_movie +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + movie_companies AS mc, + movie_info AS mi, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND it.info = 'release dates' + AND mi.info IS NOT NULL + AND (mi.info LIKE 'Japan:%200%' + OR mi.info LIKE 'USA:%200%') + AND n.gender ='f' + AND n.name LIKE '%An%' + AND rt.role ='actress' + AND t.production_year > 2000 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mi.movie_id = ci.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id; + +result sql_benchmarks/imdb/results/19c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/19d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/19d.benchmark new file mode 100644 index 0000000000000..34dfd2ef43a64 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/19d.benchmark @@ -0,0 +1,51 @@ +name Q19d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS voicing_actress, + MIN(t.title) AS jap_engl_voiced_movie +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + movie_companies AS mc, + movie_info AS mi, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND it.info = 'release dates' + AND n.gender ='f' + AND rt.role ='actress' + AND t.production_year > 2000 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mi.movie_id = ci.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id; + +result sql_benchmarks/imdb/results/19d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/20a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/20a.benchmark new file mode 100644 index 0000000000000..3a30479008616 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/20a.benchmark @@ -0,0 +1,55 @@ +name Q20a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS complete_downey_ironman_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + keyword AS k, + kind_type AS kt, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind LIKE '%complete%' + AND chn.name NOT LIKE '%Sherlock%' + AND (chn.name LIKE '%Tony%Stark%' + OR chn.name LIKE '%Iron%Man%') + AND k.keyword IN ('superhero', + 'sequel', + 'second-part', + 'marvel-comics', + 'based-on-comic', + 'tv-special', + 'fight', + 'violence') + AND kt.kind = 'movie' + AND t.production_year > 1950 + AND kt.id = t.kind_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = ci.movie_id + AND mk.movie_id = cc.movie_id + AND ci.movie_id = cc.movie_id + AND chn.id = ci.person_role_id + AND n.id = ci.person_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/20a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/20b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/20b.benchmark new file mode 100644 index 0000000000000..000eef0c3481e --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/20b.benchmark @@ -0,0 +1,56 @@ +name Q20b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(t.title) AS complete_downey_ironman_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + keyword AS k, + kind_type AS kt, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind LIKE '%complete%' + AND chn.name NOT LIKE '%Sherlock%' + AND (chn.name LIKE '%Tony%Stark%' + OR chn.name LIKE '%Iron%Man%') + AND k.keyword IN ('superhero', + 'sequel', + 'second-part', + 'marvel-comics', + 'based-on-comic', + 'tv-special', + 'fight', + 'violence') + AND kt.kind = 'movie' + AND n.name LIKE '%Downey%Robert%' + AND t.production_year > 2000 + AND kt.id = t.kind_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = ci.movie_id + AND mk.movie_id = cc.movie_id + AND ci.movie_id = cc.movie_id + AND chn.id = ci.person_role_id + AND n.id = ci.person_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/20b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/20c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/20c.benchmark new file mode 100644 index 0000000000000..4fa02af954c47 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/20c.benchmark @@ -0,0 +1,58 @@ +name Q20c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(n.name) AS cast_member, + MIN(t.title) AS complete_dynamic_hero_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + keyword AS k, + kind_type AS kt, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind LIKE '%complete%' + AND chn.name IS NOT NULL + AND (chn.name LIKE '%man%' + OR chn.name LIKE '%Man%') + AND k.keyword IN ('superhero', + 'marvel-comics', + 'based-on-comic', + 'tv-special', + 'fight', + 'violence', + 'magnet', + 'web', + 'claw', + 'laser') + AND kt.kind = 'movie' + AND t.production_year > 2000 + AND kt.id = t.kind_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = ci.movie_id + AND mk.movie_id = cc.movie_id + AND ci.movie_id = cc.movie_id + AND chn.id = ci.person_role_id + AND n.id = ci.person_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/20c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/21a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/21a.benchmark new file mode 100644 index 0000000000000..45713d402c719 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/21a.benchmark @@ -0,0 +1,59 @@ +name Q21a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS company_name, + MIN(lt.link) AS link_type, + MIN(t.title) AS western_follow_up +FROM company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cn.country_code !='[pl]' + AND (cn.name LIKE '%Film%' + OR cn.name LIKE '%Warner%') + AND ct.kind ='production companies' + AND k.keyword ='sequel' + AND lt.link LIKE '%follow%' + AND mc.note IS NULL + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German') + AND t.production_year BETWEEN 1950 AND 2000 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND mi.movie_id = t.id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND ml.movie_id = mi.movie_id + AND mk.movie_id = mi.movie_id + AND mc.movie_id = mi.movie_id; + +result sql_benchmarks/imdb/results/21a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/21b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/21b.benchmark new file mode 100644 index 0000000000000..9fc4a1acd88ef --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/21b.benchmark @@ -0,0 +1,53 @@ +name Q21b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS company_name, + MIN(lt.link) AS link_type, + MIN(t.title) AS german_follow_up +FROM company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cn.country_code !='[pl]' + AND (cn.name LIKE '%Film%' + OR cn.name LIKE '%Warner%') + AND ct.kind ='production companies' + AND k.keyword ='sequel' + AND lt.link LIKE '%follow%' + AND mc.note IS NULL + AND mi.info IN ('Germany', + 'German') + AND t.production_year BETWEEN 2000 AND 2010 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND mi.movie_id = t.id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND ml.movie_id = mi.movie_id + AND mk.movie_id = mi.movie_id + AND mc.movie_id = mi.movie_id; + +result sql_benchmarks/imdb/results/21b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/21c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/21c.benchmark new file mode 100644 index 0000000000000..9143fc3fc642f --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/21c.benchmark @@ -0,0 +1,60 @@ +name Q21c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS company_name, + MIN(lt.link) AS link_type, + MIN(t.title) AS western_follow_up +FROM company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cn.country_code !='[pl]' + AND (cn.name LIKE '%Film%' + OR cn.name LIKE '%Warner%') + AND ct.kind ='production companies' + AND k.keyword ='sequel' + AND lt.link LIKE '%follow%' + AND mc.note IS NULL + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German', + 'English') + AND t.production_year BETWEEN 1950 AND 2010 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND mi.movie_id = t.id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND ml.movie_id = mi.movie_id + AND mk.movie_id = mi.movie_id + AND mc.movie_id = mi.movie_id; + +result sql_benchmarks/imdb/results/21c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/22a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/22a.benchmark new file mode 100644 index 0000000000000..053bb3a0885bd --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/22a.benchmark @@ -0,0 +1,64 @@ +name Q22a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS western_violent_movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE cn.country_code != '[us]' + AND it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind IN ('movie', + 'episode') + AND mc.note NOT LIKE '%(USA)%' + AND mc.note LIKE '%(200%)%' + AND mi.info IN ('Germany', + 'German', + 'USA', + 'American') + AND mi_idx.info < '7.0' + AND t.production_year > 2008 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND t.id = mc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mk.movie_id = mc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mc.movie_id + AND mc.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id; + +result sql_benchmarks/imdb/results/22a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/22b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/22b.benchmark new file mode 100644 index 0000000000000..5e3c9011e6dae --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/22b.benchmark @@ -0,0 +1,64 @@ +name Q22b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS western_violent_movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE cn.country_code != '[us]' + AND it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind IN ('movie', + 'episode') + AND mc.note NOT LIKE '%(USA)%' + AND mc.note LIKE '%(200%)%' + AND mi.info IN ('Germany', + 'German', + 'USA', + 'American') + AND mi_idx.info < '7.0' + AND t.production_year > 2009 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND t.id = mc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mk.movie_id = mc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mc.movie_id + AND mc.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id; + +result sql_benchmarks/imdb/results/22b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/22c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/22c.benchmark new file mode 100644 index 0000000000000..72f9eae846548 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/22c.benchmark @@ -0,0 +1,70 @@ +name Q22c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS western_violent_movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE cn.country_code != '[us]' + AND it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind IN ('movie', + 'episode') + AND mc.note NOT LIKE '%(USA)%' + AND mc.note LIKE '%(200%)%' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Danish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND mi_idx.info < '8.5' + AND t.production_year > 2005 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND t.id = mc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mk.movie_id = mc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mc.movie_id + AND mc.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id; + +result sql_benchmarks/imdb/results/22c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/22d.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/22d.benchmark new file mode 100644 index 0000000000000..c7906c6f0b757 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/22d.benchmark @@ -0,0 +1,68 @@ +name Q22d +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS western_violent_movie +FROM company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE cn.country_code != '[us]' + AND it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind IN ('movie', + 'episode') + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Danish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND mi_idx.info < '8.5' + AND t.production_year > 2005 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND t.id = mc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mk.movie_id = mc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mc.movie_id + AND mc.movie_id = mi_idx.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id; + +result sql_benchmarks/imdb/results/22d.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/23a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/23a.benchmark new file mode 100644 index 0000000000000..6922670965b0c --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/23a.benchmark @@ -0,0 +1,55 @@ +name Q23a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(kt.kind) AS movie_kind, + MIN(t.title) AS complete_us_internet_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + company_name AS cn, + company_type AS ct, + info_type AS it1, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE cct1.kind = 'complete+verified' + AND cn.country_code = '[us]' + AND it1.info = 'release dates' + AND kt.kind IN ('movie') + AND mi.note LIKE '%internet%' + AND mi.info IS NOT NULL + AND (mi.info LIKE 'USA:% 199%' + OR mi.info LIKE 'USA:% 200%') + AND t.production_year > 2000 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = cc.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = cc.movie_id + AND mc.movie_id = cc.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id + AND cct1.id = cc.status_id; + +result sql_benchmarks/imdb/results/23a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/23b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/23b.benchmark new file mode 100644 index 0000000000000..800d1a4d6f9f5 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/23b.benchmark @@ -0,0 +1,57 @@ +name Q23b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(kt.kind) AS movie_kind, + MIN(t.title) AS complete_nerdy_internet_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + company_name AS cn, + company_type AS ct, + info_type AS it1, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE cct1.kind = 'complete+verified' + AND cn.country_code = '[us]' + AND it1.info = 'release dates' + AND k.keyword IN ('nerd', + 'loner', + 'alienation', + 'dignity') + AND kt.kind IN ('movie') + AND mi.note LIKE '%internet%' + AND mi.info LIKE 'USA:% 200%' + AND t.production_year > 2000 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = cc.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = cc.movie_id + AND mc.movie_id = cc.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id + AND cct1.id = cc.status_id; + +result sql_benchmarks/imdb/results/23b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/23c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/23c.benchmark new file mode 100644 index 0000000000000..7ef7b698ce737 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/23c.benchmark @@ -0,0 +1,58 @@ +name Q23c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(kt.kind) AS movie_kind, + MIN(t.title) AS complete_us_internet_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + company_name AS cn, + company_type AS ct, + info_type AS it1, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + title AS t +WHERE cct1.kind = 'complete+verified' + AND cn.country_code = '[us]' + AND it1.info = 'release dates' + AND kt.kind IN ('movie', + 'tv movie', + 'video movie', + 'video game') + AND mi.note LIKE '%internet%' + AND mi.info IS NOT NULL + AND (mi.info LIKE 'USA:% 199%' + OR mi.info LIKE 'USA:% 200%') + AND t.production_year > 1990 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = cc.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = cc.movie_id + AND mc.movie_id = cc.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND cn.id = mc.company_id + AND ct.id = mc.company_type_id + AND cct1.id = cc.status_id; + +result sql_benchmarks/imdb/results/23c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/24a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/24a.benchmark new file mode 100644 index 0000000000000..085b9104e512d --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/24a.benchmark @@ -0,0 +1,66 @@ +name Q24a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS voiced_char_name, + MIN(n.name) AS voicing_actress_name, + MIN(t.title) AS voiced_action_movie_jap_eng +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND it.info = 'release dates' + AND k.keyword IN ('hero', + 'martial-arts', + 'hand-to-hand-combat') + AND mi.info IS NOT NULL + AND (mi.info LIKE 'Japan:%201%' + OR mi.info LIKE 'USA:%201%') + AND n.gender ='f' + AND n.name LIKE '%An%' + AND rt.role ='actress' + AND t.production_year > 2010 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mc.movie_id = mk.movie_id + AND mi.movie_id = ci.movie_id + AND mi.movie_id = mk.movie_id + AND ci.movie_id = mk.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id + AND k.id = mk.keyword_id; + +result sql_benchmarks/imdb/results/24a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/24b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/24b.benchmark new file mode 100644 index 0000000000000..bdb50db40cead --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/24b.benchmark @@ -0,0 +1,69 @@ +name Q24b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS voiced_char_name, + MIN(n.name) AS voicing_actress_name, + MIN(t.title) AS kung_fu_panda +FROM aka_name AS an, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + name AS n, + role_type AS rt, + title AS t +WHERE ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND cn.name = 'DreamWorks Animation' + AND it.info = 'release dates' + AND k.keyword IN ('hero', + 'martial-arts', + 'hand-to-hand-combat', + 'computer-animated-movie') + AND mi.info IS NOT NULL + AND (mi.info LIKE 'Japan:%201%' + OR mi.info LIKE 'USA:%201%') + AND n.gender ='f' + AND n.name LIKE '%An%' + AND rt.role ='actress' + AND t.production_year > 2010 + AND t.title LIKE 'Kung Fu Panda%' + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mc.movie_id = mk.movie_id + AND mi.movie_id = ci.movie_id + AND mi.movie_id = mk.movie_id + AND ci.movie_id = mk.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id + AND k.id = mk.keyword_id; + +result sql_benchmarks/imdb/results/24b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/25a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/25a.benchmark new file mode 100644 index 0000000000000..4994f4492de15 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/25a.benchmark @@ -0,0 +1,58 @@ +name Q25a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS male_writer, + MIN(t.title) AS violent_movie_title +FROM cast_info AS ci, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'blood', + 'gore', + 'death', + 'female-nudity') + AND mi.info = 'Horror' + AND n.gender = 'm' + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi_idx.movie_id = mk.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id; + +result sql_benchmarks/imdb/results/25a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/25b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/25b.benchmark new file mode 100644 index 0000000000000..56acb0a881368 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/25b.benchmark @@ -0,0 +1,60 @@ +name Q25b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS male_writer, + MIN(t.title) AS violent_movie_title +FROM cast_info AS ci, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'blood', + 'gore', + 'death', + 'female-nudity') + AND mi.info = 'Horror' + AND n.gender = 'm' + AND t.production_year > 2010 + AND t.title LIKE 'Vampire%' + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi_idx.movie_id = mk.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id; + +result sql_benchmarks/imdb/results/25b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/25c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/25c.benchmark new file mode 100644 index 0000000000000..113b75c77bc24 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/25c.benchmark @@ -0,0 +1,65 @@ +name Q25c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS male_writer, + MIN(t.title) AS violent_movie_title +FROM cast_info AS ci, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'violence', + 'blood', + 'gore', + 'death', + 'female-nudity', + 'hospital') + AND mi.info IN ('Horror', + 'Action', + 'Sci-Fi', + 'Thriller', + 'Crime', + 'War') + AND n.gender = 'm' + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi_idx.movie_id = mk.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id; + +result sql_benchmarks/imdb/results/25c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/26a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/26a.benchmark new file mode 100644 index 0000000000000..17ebb36029223 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/26a.benchmark @@ -0,0 +1,69 @@ +name Q26a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS character_name, + MIN(mi_idx.info) AS rating, + MIN(n.name) AS playing_actor, + MIN(t.title) AS complete_hero_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind LIKE '%complete%' + AND chn.name IS NOT NULL + AND (chn.name LIKE '%man%' + OR chn.name LIKE '%Man%') + AND it2.info = 'rating' + AND k.keyword IN ('superhero', + 'marvel-comics', + 'based-on-comic', + 'tv-special', + 'fight', + 'violence', + 'magnet', + 'web', + 'claw', + 'laser') + AND kt.kind = 'movie' + AND mi_idx.info > '7.0' + AND t.production_year > 2000 + AND kt.id = t.kind_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND t.id = cc.movie_id + AND t.id = mi_idx.movie_id + AND mk.movie_id = ci.movie_id + AND mk.movie_id = cc.movie_id + AND mk.movie_id = mi_idx.movie_id + AND ci.movie_id = cc.movie_id + AND ci.movie_id = mi_idx.movie_id + AND cc.movie_id = mi_idx.movie_id + AND chn.id = ci.person_role_id + AND n.id = ci.person_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/26a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/26b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/26b.benchmark new file mode 100644 index 0000000000000..bc03fd914c08a --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/26b.benchmark @@ -0,0 +1,62 @@ +name Q26b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS character_name, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS complete_hero_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind LIKE '%complete%' + AND chn.name IS NOT NULL + AND (chn.name LIKE '%man%' + OR chn.name LIKE '%Man%') + AND it2.info = 'rating' + AND k.keyword IN ('superhero', + 'marvel-comics', + 'based-on-comic', + 'fight') + AND kt.kind = 'movie' + AND mi_idx.info > '8.0' + AND t.production_year > 2005 + AND kt.id = t.kind_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND t.id = cc.movie_id + AND t.id = mi_idx.movie_id + AND mk.movie_id = ci.movie_id + AND mk.movie_id = cc.movie_id + AND mk.movie_id = mi_idx.movie_id + AND ci.movie_id = cc.movie_id + AND ci.movie_id = mi_idx.movie_id + AND cc.movie_id = mi_idx.movie_id + AND chn.id = ci.person_role_id + AND n.id = ci.person_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/26b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/26c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/26c.benchmark new file mode 100644 index 0000000000000..b07e738425c9c --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/26c.benchmark @@ -0,0 +1,67 @@ +name Q26c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS character_name, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS complete_hero_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind LIKE '%complete%' + AND chn.name IS NOT NULL + AND (chn.name LIKE '%man%' + OR chn.name LIKE '%Man%') + AND it2.info = 'rating' + AND k.keyword IN ('superhero', + 'marvel-comics', + 'based-on-comic', + 'tv-special', + 'fight', + 'violence', + 'magnet', + 'web', + 'claw', + 'laser') + AND kt.kind = 'movie' + AND t.production_year > 2000 + AND kt.id = t.kind_id + AND t.id = mk.movie_id + AND t.id = ci.movie_id + AND t.id = cc.movie_id + AND t.id = mi_idx.movie_id + AND mk.movie_id = ci.movie_id + AND mk.movie_id = cc.movie_id + AND mk.movie_id = mi_idx.movie_id + AND ci.movie_id = cc.movie_id + AND ci.movie_id = mi_idx.movie_id + AND cc.movie_id = mi_idx.movie_id + AND chn.id = ci.person_role_id + AND n.id = ci.person_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id + AND it2.id = mi_idx.info_type_id; + +result sql_benchmarks/imdb/results/26c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/27a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/27a.benchmark new file mode 100644 index 0000000000000..ff0cfebf81050 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/27a.benchmark @@ -0,0 +1,68 @@ +name Q27a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS producing_company, + MIN(lt.link) AS link_type, + MIN(t.title) AS complete_western_sequel +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cct1.kind IN ('cast', + 'crew') + AND cct2.kind = 'complete' + AND cn.country_code !='[pl]' + AND (cn.name LIKE '%Film%' + OR cn.name LIKE '%Warner%') + AND ct.kind ='production companies' + AND k.keyword ='sequel' + AND lt.link LIKE '%follow%' + AND mc.note IS NULL + AND mi.info IN ('Sweden', + 'Germany', + 'Swedish', + 'German') + AND t.production_year BETWEEN 1950 AND 2000 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND mi.movie_id = t.id + AND t.id = cc.movie_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND ml.movie_id = mi.movie_id + AND mk.movie_id = mi.movie_id + AND mc.movie_id = mi.movie_id + AND ml.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND mc.movie_id = cc.movie_id + AND mi.movie_id = cc.movie_id; + +result sql_benchmarks/imdb/results/27a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/27b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/27b.benchmark new file mode 100644 index 0000000000000..bf0fa5b69ec52 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/27b.benchmark @@ -0,0 +1,68 @@ +name Q27b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS producing_company, + MIN(lt.link) AS link_type, + MIN(t.title) AS complete_western_sequel +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cct1.kind IN ('cast', + 'crew') + AND cct2.kind = 'complete' + AND cn.country_code !='[pl]' + AND (cn.name LIKE '%Film%' + OR cn.name LIKE '%Warner%') + AND ct.kind ='production companies' + AND k.keyword ='sequel' + AND lt.link LIKE '%follow%' + AND mc.note IS NULL + AND mi.info IN ('Sweden', + 'Germany', + 'Swedish', + 'German') + AND t.production_year = 1998 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND mi.movie_id = t.id + AND t.id = cc.movie_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND ml.movie_id = mi.movie_id + AND mk.movie_id = mi.movie_id + AND mc.movie_id = mi.movie_id + AND ml.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND mc.movie_id = cc.movie_id + AND mi.movie_id = cc.movie_id; + +result sql_benchmarks/imdb/results/27b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/27c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/27c.benchmark new file mode 100644 index 0000000000000..fd7444531277e --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/27c.benchmark @@ -0,0 +1,72 @@ +name Q27c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS producing_company, + MIN(lt.link) AS link_type, + MIN(t.title) AS complete_western_sequel +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + company_name AS cn, + company_type AS ct, + keyword AS k, + link_type AS lt, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + movie_link AS ml, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind LIKE 'complete%' + AND cn.country_code !='[pl]' + AND (cn.name LIKE '%Film%' + OR cn.name LIKE '%Warner%') + AND ct.kind ='production companies' + AND k.keyword ='sequel' + AND lt.link LIKE '%follow%' + AND mc.note IS NULL + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Denish', + 'Norwegian', + 'German', + 'English') + AND t.production_year BETWEEN 1950 AND 2010 + AND lt.id = ml.link_type_id + AND ml.movie_id = t.id + AND t.id = mk.movie_id + AND mk.keyword_id = k.id + AND t.id = mc.movie_id + AND mc.company_type_id = ct.id + AND mc.company_id = cn.id + AND mi.movie_id = t.id + AND t.id = cc.movie_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id + AND ml.movie_id = mk.movie_id + AND ml.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND ml.movie_id = mi.movie_id + AND mk.movie_id = mi.movie_id + AND mc.movie_id = mi.movie_id + AND ml.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND mc.movie_id = cc.movie_id + AND mi.movie_id = cc.movie_id; + +result sql_benchmarks/imdb/results/27c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/28a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/28a.benchmark new file mode 100644 index 0000000000000..1fd17967b12f4 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/28a.benchmark @@ -0,0 +1,82 @@ +name Q28a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS complete_euro_dark_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE cct1.kind = 'crew' + AND cct2.kind != 'complete+verified' + AND cn.country_code != '[us]' + AND it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind IN ('movie', + 'episode') + AND mc.note NOT LIKE '%(USA)%' + AND mc.note LIKE '%(200%)%' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Danish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND mi_idx.info < '8.5' + AND t.production_year > 2000 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND t.id = mc.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = cc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = cc.movie_id + AND mc.movie_id = mi_idx.movie_id + AND mc.movie_id = cc.movie_id + AND mi_idx.movie_id = cc.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/28a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/28b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/28b.benchmark new file mode 100644 index 0000000000000..0b68663f7fba6 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/28b.benchmark @@ -0,0 +1,76 @@ +name Q28b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS complete_euro_dark_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE cct1.kind = 'crew' + AND cct2.kind != 'complete+verified' + AND cn.country_code != '[us]' + AND it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind IN ('movie', + 'episode') + AND mc.note NOT LIKE '%(USA)%' + AND mc.note LIKE '%(200%)%' + AND mi.info IN ('Sweden', + 'Germany', + 'Swedish', + 'German') + AND mi_idx.info > '6.5' + AND t.production_year > 2005 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND t.id = mc.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = cc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = cc.movie_id + AND mc.movie_id = mi_idx.movie_id + AND mc.movie_id = cc.movie_id + AND mi_idx.movie_id = cc.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/28b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/28c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/28c.benchmark new file mode 100644 index 0000000000000..b64d407a67d51 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/28c.benchmark @@ -0,0 +1,82 @@ +name Q28c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn.name) AS movie_company, + MIN(mi_idx.info) AS rating, + MIN(t.title) AS complete_euro_dark_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + company_name AS cn, + company_type AS ct, + info_type AS it1, + info_type AS it2, + keyword AS k, + kind_type AS kt, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind = 'complete' + AND cn.country_code != '[us]' + AND it1.info = 'countries' + AND it2.info = 'rating' + AND k.keyword IN ('murder', + 'murder-in-title', + 'blood', + 'violence') + AND kt.kind IN ('movie', + 'episode') + AND mc.note NOT LIKE '%(USA)%' + AND mc.note LIKE '%(200%)%' + AND mi.info IN ('Sweden', + 'Norway', + 'Germany', + 'Denmark', + 'Swedish', + 'Danish', + 'Norwegian', + 'German', + 'USA', + 'American') + AND mi_idx.info < '8.5' + AND t.production_year > 2005 + AND kt.id = t.kind_id + AND t.id = mi.movie_id + AND t.id = mk.movie_id + AND t.id = mi_idx.movie_id + AND t.id = mc.movie_id + AND t.id = cc.movie_id + AND mk.movie_id = mi.movie_id + AND mk.movie_id = mi_idx.movie_id + AND mk.movie_id = mc.movie_id + AND mk.movie_id = cc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mc.movie_id + AND mi.movie_id = cc.movie_id + AND mc.movie_id = mi_idx.movie_id + AND mc.movie_id = cc.movie_id + AND mi_idx.movie_id = cc.movie_id + AND k.id = mk.keyword_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND ct.id = mc.company_type_id + AND cn.id = mc.company_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/28c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/29a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/29a.benchmark new file mode 100644 index 0000000000000..40affdb3d557c --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/29a.benchmark @@ -0,0 +1,83 @@ +name Q29a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS voiced_char, + MIN(n.name) AS voicing_actress, + MIN(t.title) AS voiced_animation +FROM aka_name AS an, + complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + info_type AS it3, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + name AS n, + person_info AS pi, + role_type AS rt, + title AS t +WHERE cct1.kind ='cast' + AND cct2.kind ='complete+verified' + AND chn.name = 'Queen' + AND ci.note IN ('(voice)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND it.info = 'release dates' + AND it3.info = 'trivia' + AND k.keyword = 'computer-animation' + AND mi.info IS NOT NULL + AND (mi.info LIKE 'Japan:%200%' + OR mi.info LIKE 'USA:%200%') + AND n.gender ='f' + AND n.name LIKE '%An%' + AND rt.role ='actress' + AND t.title = 'Shrek 2' + AND t.production_year BETWEEN 2000 AND 2010 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = cc.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mc.movie_id = mk.movie_id + AND mc.movie_id = cc.movie_id + AND mi.movie_id = ci.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = cc.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id + AND n.id = pi.person_id + AND ci.person_id = pi.person_id + AND it3.id = pi.info_type_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/29a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/29b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/29b.benchmark new file mode 100644 index 0000000000000..9d43c7151071c --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/29b.benchmark @@ -0,0 +1,81 @@ +name Q29b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS voiced_char, + MIN(n.name) AS voicing_actress, + MIN(t.title) AS voiced_animation +FROM aka_name AS an, + complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + info_type AS it3, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + name AS n, + person_info AS pi, + role_type AS rt, + title AS t +WHERE cct1.kind ='cast' + AND cct2.kind ='complete+verified' + AND chn.name = 'Queen' + AND ci.note IN ('(voice)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND it.info = 'release dates' + AND it3.info = 'height' + AND k.keyword = 'computer-animation' + AND mi.info LIKE 'USA:%200%' + AND n.gender ='f' + AND n.name LIKE '%An%' + AND rt.role ='actress' + AND t.title = 'Shrek 2' + AND t.production_year BETWEEN 2000 AND 2005 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = cc.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mc.movie_id = mk.movie_id + AND mc.movie_id = cc.movie_id + AND mi.movie_id = ci.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = cc.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id + AND n.id = pi.person_id + AND ci.person_id = pi.person_id + AND it3.id = pi.info_type_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/29b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/29c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/29c.benchmark new file mode 100644 index 0000000000000..9d0cbdc14cc02 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/29c.benchmark @@ -0,0 +1,82 @@ +name Q29c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(chn.name) AS voiced_char, + MIN(n.name) AS voicing_actress, + MIN(t.title) AS voiced_animation +FROM aka_name AS an, + complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + char_name AS chn, + cast_info AS ci, + company_name AS cn, + info_type AS it, + info_type AS it3, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_keyword AS mk, + name AS n, + person_info AS pi, + role_type AS rt, + title AS t +WHERE cct1.kind ='cast' + AND cct2.kind ='complete+verified' + AND ci.note IN ('(voice)', + '(voice: Japanese version)', + '(voice) (uncredited)', + '(voice: English version)') + AND cn.country_code ='[us]' + AND it.info = 'release dates' + AND it3.info = 'trivia' + AND k.keyword = 'computer-animation' + AND mi.info IS NOT NULL + AND (mi.info LIKE 'Japan:%200%' + OR mi.info LIKE 'USA:%200%') + AND n.gender ='f' + AND n.name LIKE '%An%' + AND rt.role ='actress' + AND t.production_year BETWEEN 2000 AND 2010 + AND t.id = mi.movie_id + AND t.id = mc.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = cc.movie_id + AND mc.movie_id = ci.movie_id + AND mc.movie_id = mi.movie_id + AND mc.movie_id = mk.movie_id + AND mc.movie_id = cc.movie_id + AND mi.movie_id = ci.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = cc.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND cn.id = mc.company_id + AND it.id = mi.info_type_id + AND n.id = ci.person_id + AND rt.id = ci.role_id + AND n.id = an.person_id + AND ci.person_id = an.person_id + AND chn.id = ci.person_role_id + AND n.id = pi.person_id + AND ci.person_id = pi.person_id + AND it3.id = pi.info_type_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/29c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/30a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/30a.benchmark new file mode 100644 index 0000000000000..747c43e60ccdb --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/30a.benchmark @@ -0,0 +1,75 @@ +name Q30a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS writer, + MIN(t.title) AS complete_violent_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + cast_info AS ci, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind IN ('cast', + 'crew') + AND cct2.kind ='complete+verified' + AND ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'violence', + 'blood', + 'gore', + 'death', + 'female-nudity', + 'hospital') + AND mi.info IN ('Horror', + 'Thriller') + AND n.gender = 'm' + AND t.production_year > 2000 + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = cc.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = cc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = cc.movie_id + AND mi_idx.movie_id = mk.movie_id + AND mi_idx.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/30a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/30b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/30b.benchmark new file mode 100644 index 0000000000000..6f29177a91e71 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/30b.benchmark @@ -0,0 +1,78 @@ +name Q30b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS writer, + MIN(t.title) AS complete_gore_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + cast_info AS ci, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind IN ('cast', + 'crew') + AND cct2.kind ='complete+verified' + AND ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'violence', + 'blood', + 'gore', + 'death', + 'female-nudity', + 'hospital') + AND mi.info IN ('Horror', + 'Thriller') + AND n.gender = 'm' + AND t.production_year > 2000 + AND (t.title LIKE '%Freddy%' + OR t.title LIKE '%Jason%' + OR t.title LIKE 'Saw%') + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = cc.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = cc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = cc.movie_id + AND mi_idx.movie_id = mk.movie_id + AND mi_idx.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/30b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/30c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/30c.benchmark new file mode 100644 index 0000000000000..78cbbdf2f4543 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/30c.benchmark @@ -0,0 +1,77 @@ +name Q30c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS writer, + MIN(t.title) AS complete_violent_movie +FROM complete_cast AS cc, + comp_cast_type AS cct1, + comp_cast_type AS cct2, + cast_info AS ci, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE cct1.kind = 'cast' + AND cct2.kind ='complete+verified' + AND ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'violence', + 'blood', + 'gore', + 'death', + 'female-nudity', + 'hospital') + AND mi.info IN ('Horror', + 'Action', + 'Sci-Fi', + 'Thriller', + 'Crime', + 'War') + AND n.gender = 'm' + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = cc.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = cc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = cc.movie_id + AND mi_idx.movie_id = mk.movie_id + AND mi_idx.movie_id = cc.movie_id + AND mk.movie_id = cc.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id + AND cct1.id = cc.subject_id + AND cct2.id = cc.status_id; + +result sql_benchmarks/imdb/results/30c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/31a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/31a.benchmark new file mode 100644 index 0000000000000..8f3b5a1567da6 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/31a.benchmark @@ -0,0 +1,70 @@ +name Q31a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS writer, + MIN(t.title) AS violent_liongate_movie +FROM cast_info AS ci, + company_name AS cn, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND cn.name LIKE 'Lionsgate%' + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'violence', + 'blood', + 'gore', + 'death', + 'female-nudity', + 'hospital') + AND mi.info IN ('Horror', + 'Thriller') + AND n.gender = 'm' + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = mc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = mc.movie_id + AND mi_idx.movie_id = mk.movie_id + AND mi_idx.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id + AND cn.id = mc.company_id; + +result sql_benchmarks/imdb/results/31a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/31b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/31b.benchmark new file mode 100644 index 0000000000000..7395f37089c14 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/31b.benchmark @@ -0,0 +1,75 @@ +name Q31b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS writer, + MIN(t.title) AS violent_liongate_movie +FROM cast_info AS ci, + company_name AS cn, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND cn.name LIKE 'Lionsgate%' + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'violence', + 'blood', + 'gore', + 'death', + 'female-nudity', + 'hospital') + AND mc.note LIKE '%(Blu-ray)%' + AND mi.info IN ('Horror', + 'Thriller') + AND n.gender = 'm' + AND t.production_year > 2000 + AND (t.title LIKE '%Freddy%' + OR t.title LIKE '%Jason%' + OR t.title LIKE 'Saw%') + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = mc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = mc.movie_id + AND mi_idx.movie_id = mk.movie_id + AND mi_idx.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id + AND cn.id = mc.company_id; + +result sql_benchmarks/imdb/results/31b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/31c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/31c.benchmark new file mode 100644 index 0000000000000..ca1efcf21385d --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/31c.benchmark @@ -0,0 +1,73 @@ +name Q31c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(mi.info) AS movie_budget, + MIN(mi_idx.info) AS movie_votes, + MIN(n.name) AS writer, + MIN(t.title) AS violent_liongate_movie +FROM cast_info AS ci, + company_name AS cn, + info_type AS it1, + info_type AS it2, + keyword AS k, + movie_companies AS mc, + movie_info AS mi, + movie_info_idx AS mi_idx, + movie_keyword AS mk, + name AS n, + title AS t +WHERE ci.note IN ('(writer)', + '(head writer)', + '(written by)', + '(story)', + '(story editor)') + AND cn.name LIKE 'Lionsgate%' + AND it1.info = 'genres' + AND it2.info = 'votes' + AND k.keyword IN ('murder', + 'violence', + 'blood', + 'gore', + 'death', + 'female-nudity', + 'hospital') + AND mi.info IN ('Horror', + 'Action', + 'Sci-Fi', + 'Thriller', + 'Crime', + 'War') + AND t.id = mi.movie_id + AND t.id = mi_idx.movie_id + AND t.id = ci.movie_id + AND t.id = mk.movie_id + AND t.id = mc.movie_id + AND ci.movie_id = mi.movie_id + AND ci.movie_id = mi_idx.movie_id + AND ci.movie_id = mk.movie_id + AND ci.movie_id = mc.movie_id + AND mi.movie_id = mi_idx.movie_id + AND mi.movie_id = mk.movie_id + AND mi.movie_id = mc.movie_id + AND mi_idx.movie_id = mk.movie_id + AND mi_idx.movie_id = mc.movie_id + AND mk.movie_id = mc.movie_id + AND n.id = ci.person_id + AND it1.id = mi.info_type_id + AND it2.id = mi_idx.info_type_id + AND k.id = mk.keyword_id + AND cn.id = mc.company_id; + +result sql_benchmarks/imdb/results/31c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/32a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/32a.benchmark new file mode 100644 index 0000000000000..54380bc0c2852 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/32a.benchmark @@ -0,0 +1,33 @@ +name Q32a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(lt.link) AS link_type, + MIN(t1.title) AS first_movie, + MIN(t2.title) AS second_movie +FROM keyword AS k, + link_type AS lt, + movie_keyword AS mk, + movie_link AS ml, + title AS t1, + title AS t2 +WHERE k.keyword ='10,000-mile-club' + AND mk.keyword_id = k.id + AND t1.id = mk.movie_id + AND ml.movie_id = t1.id + AND ml.linked_movie_id = t2.id + AND lt.id = ml.link_type_id + AND mk.movie_id = t1.id; + +result sql_benchmarks/imdb/results/32a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/32b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/32b.benchmark new file mode 100644 index 0000000000000..7f6582efd272a --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/32b.benchmark @@ -0,0 +1,33 @@ +name Q32b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(lt.link) AS link_type, + MIN(t1.title) AS first_movie, + MIN(t2.title) AS second_movie +FROM keyword AS k, + link_type AS lt, + movie_keyword AS mk, + movie_link AS ml, + title AS t1, + title AS t2 +WHERE k.keyword ='character-name-in-title' + AND mk.keyword_id = k.id + AND t1.id = mk.movie_id + AND ml.movie_id = t1.id + AND ml.linked_movie_id = t2.id + AND lt.id = ml.link_type_id + AND mk.movie_id = t1.id; + +result sql_benchmarks/imdb/results/32b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/33a.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/33a.benchmark new file mode 100644 index 0000000000000..f62e614f899b1 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/33a.benchmark @@ -0,0 +1,66 @@ +name Q33a +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn1.name) AS first_company, + MIN(cn2.name) AS second_company, + MIN(mi_idx1.info) AS first_rating, + MIN(mi_idx2.info) AS second_rating, + MIN(t1.title) AS first_movie, + MIN(t2.title) AS second_movie +FROM company_name AS cn1, + company_name AS cn2, + info_type AS it1, + info_type AS it2, + kind_type AS kt1, + kind_type AS kt2, + link_type AS lt, + movie_companies AS mc1, + movie_companies AS mc2, + movie_info_idx AS mi_idx1, + movie_info_idx AS mi_idx2, + movie_link AS ml, + title AS t1, + title AS t2 +WHERE cn1.country_code = '[us]' + AND it1.info = 'rating' + AND it2.info = 'rating' + AND kt1.kind IN ('tv series') + AND kt2.kind IN ('tv series') + AND lt.link IN ('sequel', + 'follows', + 'followed by') + AND mi_idx2.info < '3.0' + AND t2.production_year BETWEEN 2005 AND 2008 + AND lt.id = ml.link_type_id + AND t1.id = ml.movie_id + AND t2.id = ml.linked_movie_id + AND it1.id = mi_idx1.info_type_id + AND t1.id = mi_idx1.movie_id + AND kt1.id = t1.kind_id + AND cn1.id = mc1.company_id + AND t1.id = mc1.movie_id + AND ml.movie_id = mi_idx1.movie_id + AND ml.movie_id = mc1.movie_id + AND mi_idx1.movie_id = mc1.movie_id + AND it2.id = mi_idx2.info_type_id + AND t2.id = mi_idx2.movie_id + AND kt2.id = t2.kind_id + AND cn2.id = mc2.company_id + AND t2.id = mc2.movie_id + AND ml.linked_movie_id = mi_idx2.movie_id + AND ml.linked_movie_id = mc2.movie_id + AND mi_idx2.movie_id = mc2.movie_id; + +result sql_benchmarks/imdb/results/33a.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/33b.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/33b.benchmark new file mode 100644 index 0000000000000..01f21763de5c1 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/33b.benchmark @@ -0,0 +1,64 @@ +name Q33b +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn1.name) AS first_company, + MIN(cn2.name) AS second_company, + MIN(mi_idx1.info) AS first_rating, + MIN(mi_idx2.info) AS second_rating, + MIN(t1.title) AS first_movie, + MIN(t2.title) AS second_movie +FROM company_name AS cn1, + company_name AS cn2, + info_type AS it1, + info_type AS it2, + kind_type AS kt1, + kind_type AS kt2, + link_type AS lt, + movie_companies AS mc1, + movie_companies AS mc2, + movie_info_idx AS mi_idx1, + movie_info_idx AS mi_idx2, + movie_link AS ml, + title AS t1, + title AS t2 +WHERE cn1.country_code = '[nl]' + AND it1.info = 'rating' + AND it2.info = 'rating' + AND kt1.kind IN ('tv series') + AND kt2.kind IN ('tv series') + AND lt.link LIKE '%follow%' + AND mi_idx2.info < '3.0' + AND t2.production_year = 2007 + AND lt.id = ml.link_type_id + AND t1.id = ml.movie_id + AND t2.id = ml.linked_movie_id + AND it1.id = mi_idx1.info_type_id + AND t1.id = mi_idx1.movie_id + AND kt1.id = t1.kind_id + AND cn1.id = mc1.company_id + AND t1.id = mc1.movie_id + AND ml.movie_id = mi_idx1.movie_id + AND ml.movie_id = mc1.movie_id + AND mi_idx1.movie_id = mc1.movie_id + AND it2.id = mi_idx2.info_type_id + AND t2.id = mi_idx2.movie_id + AND kt2.id = t2.kind_id + AND cn2.id = mc2.company_id + AND t2.id = mc2.movie_id + AND ml.linked_movie_id = mi_idx2.movie_id + AND ml.linked_movie_id = mc2.movie_id + AND mi_idx2.movie_id = mc2.movie_id; + +result sql_benchmarks/imdb/results/33b.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/benchmarks/33c.benchmark b/benchmarks/sql_benchmarks/imdb/benchmarks/33c.benchmark new file mode 100644 index 0000000000000..a0b7abed6cdbc --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/benchmarks/33c.benchmark @@ -0,0 +1,68 @@ +name Q33c +group imdb + +echo Loading imdb tables + +load sql_benchmarks/imdb/init/load_${IMDB_FILE_TYPE:-parquet}.sql + +assert I +SELECT COUNT(*) > 0 from title; +---- +true + +run +SELECT MIN(cn1.name) AS first_company, + MIN(cn2.name) AS second_company, + MIN(mi_idx1.info) AS first_rating, + MIN(mi_idx2.info) AS second_rating, + MIN(t1.title) AS first_movie, + MIN(t2.title) AS second_movie +FROM company_name AS cn1, + company_name AS cn2, + info_type AS it1, + info_type AS it2, + kind_type AS kt1, + kind_type AS kt2, + link_type AS lt, + movie_companies AS mc1, + movie_companies AS mc2, + movie_info_idx AS mi_idx1, + movie_info_idx AS mi_idx2, + movie_link AS ml, + title AS t1, + title AS t2 +WHERE cn1.country_code != '[us]' + AND it1.info = 'rating' + AND it2.info = 'rating' + AND kt1.kind IN ('tv series', + 'episode') + AND kt2.kind IN ('tv series', + 'episode') + AND lt.link IN ('sequel', + 'follows', + 'followed by') + AND mi_idx2.info < '3.5' + AND t2.production_year BETWEEN 2000 AND 2010 + AND lt.id = ml.link_type_id + AND t1.id = ml.movie_id + AND t2.id = ml.linked_movie_id + AND it1.id = mi_idx1.info_type_id + AND t1.id = mi_idx1.movie_id + AND kt1.id = t1.kind_id + AND cn1.id = mc1.company_id + AND t1.id = mc1.movie_id + AND ml.movie_id = mi_idx1.movie_id + AND ml.movie_id = mc1.movie_id + AND mi_idx1.movie_id = mc1.movie_id + AND it2.id = mi_idx2.info_type_id + AND t2.id = mi_idx2.movie_id + AND kt2.id = t2.kind_id + AND cn2.id = mc2.company_id + AND t2.id = mc2.movie_id + AND ml.linked_movie_id = mi_idx2.movie_id + AND ml.linked_movie_id = mc2.movie_id + AND mi_idx2.movie_id = mc2.movie_id; + +result sql_benchmarks/imdb/results/33c.csv + +cleanup sql_benchmarks/imdb/init/cleanup.sql \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/imdb.suite b/benchmarks/sql_benchmarks/imdb/imdb.suite new file mode 100644 index 0000000000000..7422b06bbc345 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/imdb.suite @@ -0,0 +1,26 @@ +description = "Join Order Benchmark queries over the IMDb dataset" + +query_pattern = "{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "format" +short = "f" +env = "IMDB_FILE_TYPE" +default = "parquet" +values = ["parquet", "csv"] +help = "Selects the IMDb data format." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- imdb" +description = "Run all IMDb queries against Parquet data." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- imdb --query 01a" +description = "Run IMDb query 01a." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- imdb --query 01a -f csv" +description = "Run IMDb query 01a against CSV data." diff --git a/benchmarks/sql_benchmarks/imdb/init/cleanup.sql b/benchmarks/sql_benchmarks/imdb/init/cleanup.sql new file mode 100644 index 0000000000000..5ec8696caaa50 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/init/cleanup.sql @@ -0,0 +1,41 @@ +DROP TABLE IF EXISTS aka_name; + +DROP TABLE IF EXISTS aka_title; + +DROP TABLE IF EXISTS cast_info; + +DROP TABLE IF EXISTS char_name; + +DROP TABLE IF EXISTS comp_cast_type; + +DROP TABLE IF EXISTS company_name; + +DROP TABLE IF EXISTS company_type; + +DROP TABLE IF EXISTS complete_cast; + +DROP TABLE IF EXISTS info_type; + +DROP TABLE IF EXISTS keyword; + +DROP TABLE IF EXISTS kind_type; + +DROP TABLE IF EXISTS link_type; + +DROP TABLE IF EXISTS movie_companies; + +DROP TABLE IF EXISTS movie_info; + +DROP TABLE IF EXISTS movie_info_idx; + +DROP TABLE IF EXISTS movie_keyword; + +DROP TABLE IF EXISTS movie_link; + +DROP TABLE IF EXISTS name; + +DROP TABLE IF EXISTS person_info; + +DROP TABLE IF EXISTS role_type; + +DROP TABLE IF EXISTS title; \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/imdb/init/load_csv.sql b/benchmarks/sql_benchmarks/imdb/init/load_csv.sql new file mode 100644 index 0000000000000..02e8867388aa1 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/init/load_csv.sql @@ -0,0 +1,170 @@ +CREATE EXTERNAL TABLE aka_name ( + id integer unsigned NOT NULL, + person_id integer NOT NULL, + name varchar(218) NOT NULL, + imdb_index varchar(12), + name_pcode_cf varchar(5), + name_pcode_nf varchar(5), + surname_pcode varchar(5), + md5sum varchar(32) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/aka_name.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE aka_title ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + title varchar(553) NOT NULL, + imdb_index varchar(12), + kind_id integer NOT NULL, + production_year integer, + phonetic_code varchar(5), + episode_of_id integer, + season_nr integer, + episode_nr integer, + note varchar(72), + md5sum varchar(32) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/aka_title.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE cast_info ( + id integer unsigned NOT NULL, + person_id integer NOT NULL, + movie_id integer NOT NULL, + person_role_id integer, + note varchar(992), + nr_order integer, + role_id integer NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/cast_info.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE char_name ( + id integer unsigned NOT NULL, + name varchar(478) NOT NULL, + imdb_index varchar(12), + imdb_id integer, + name_pcode_nf varchar(5), + surname_pcode varchar(5), + md5sum varchar(32) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/char_name.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE comp_cast_type ( + id integer unsigned NOT NULL, + kind varchar(32) NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/comp_cast_type.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE company_name ( + id integer unsigned NOT NULL, + name varchar(200) NOT NULL, + country_code varchar(255), + imdb_id integer, + name_pcode_nf varchar(5), + name_pcode_sf varchar(5), + md5sum varchar(32) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/company_name.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE company_type ( + id integer unsigned NOT NULL, + kind varchar(32) NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/company_type.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE complete_cast ( + id integer unsigned NOT NULL, + movie_id integer, + subject_id integer NOT NULL, + status_id integer NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/complete_cast.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE info_type ( + id integer unsigned NOT NULL, + info varchar(32) NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/info_type.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE keyword ( + id integer unsigned NOT NULL, + keyword varchar(74) NOT NULL, + phonetic_code varchar(5) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/keyword.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE kind_type ( + id integer unsigned NOT NULL, + kind varchar(15) NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/kind_type.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE link_type ( + id integer unsigned NOT NULL, + link varchar(32) NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/link_type.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE movie_companies ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + company_id integer NOT NULL, + company_type_id integer NOT NULL, + note varchar(208) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/movie_companies.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE movie_info ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + info_type_id integer NOT NULL, + info varchar(8000) NOT NULL, + note varchar(387) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/movie_info.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE movie_info_idx ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + info_type_id integer NOT NULL, + info varchar(10) NOT NULL, + note varchar(1) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/movie_info_idx.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE movie_keyword ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + keyword_id integer NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/movie_keyword.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE movie_link ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + linked_movie_id integer NOT NULL, + link_type_id integer NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/movie_link.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE name ( + id integer unsigned NOT NULL, + name varchar(106) NOT NULL, + imdb_index varchar(12), + imdb_id integer, + gender varchar(1), + name_pcode_cf varchar(5), + name_pcode_nf varchar(5), + surname_pcode varchar(5), + md5sum varchar(32) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/name.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE person_info ( + id integer unsigned NOT NULL, + person_id integer NOT NULL, + info_type_id integer NOT NULL, + info text NOT NULL, + note varchar(430) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/person_info.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE role_type ( + id integer unsigned NOT NULL, + role varchar(32) NOT NULL +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/role_type.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); + +CREATE EXTERNAL TABLE title ( + id integer unsigned NOT NULL, + title varchar(334) NOT NULL, + imdb_index varchar(12), + kind_id integer NOT NULL, + production_year integer, + imdb_id integer, + phonetic_code varchar(5), + episode_of_id integer, + season_nr integer, + episode_nr integer, + series_years varchar(49), + md5sum varchar(32) +) STORED AS CSV LOCATION '${DATA_DIR:-data}/imdb/title.csv' OPTIONS ('has_header' 'false', 'format.delimiter' ',', 'format.escape' '\'); diff --git a/benchmarks/sql_benchmarks/imdb/init/load_parquet.sql b/benchmarks/sql_benchmarks/imdb/init/load_parquet.sql new file mode 100644 index 0000000000000..1c1d28b2436d5 --- /dev/null +++ b/benchmarks/sql_benchmarks/imdb/init/load_parquet.sql @@ -0,0 +1,170 @@ +CREATE EXTERNAL TABLE aka_name ( + id integer unsigned NOT NULL, + person_id integer NOT NULL, + name varchar(218) NOT NULL, + imdb_index varchar(12), + name_pcode_cf varchar(5), + name_pcode_nf varchar(5), + surname_pcode varchar(5), + md5sum varchar(32) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/aka_name.parquet'; + +CREATE EXTERNAL TABLE aka_title ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + title varchar(553) NOT NULL, + imdb_index varchar(12), + kind_id integer NOT NULL, + production_year integer, + phonetic_code varchar(5), + episode_of_id integer, + season_nr integer, + episode_nr integer, + note varchar(72), + md5sum varchar(32) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/aka_title.parquet'; + +CREATE EXTERNAL TABLE cast_info ( + id integer unsigned NOT NULL, + person_id integer NOT NULL, + movie_id integer NOT NULL, + person_role_id integer, + note varchar(992), + nr_order integer, + role_id integer NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/cast_info.parquet'; + +CREATE EXTERNAL TABLE char_name ( + id integer unsigned NOT NULL, + name varchar(478) NOT NULL, + imdb_index varchar(12), + imdb_id integer, + name_pcode_nf varchar(5), + surname_pcode varchar(5), + md5sum varchar(32) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/char_name.parquet'; + +CREATE EXTERNAL TABLE comp_cast_type ( + id integer unsigned NOT NULL, + kind varchar(32) NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/comp_cast_type.parquet'; + +CREATE EXTERNAL TABLE company_name ( + id integer unsigned NOT NULL, + name varchar(200) NOT NULL, + country_code varchar(255), + imdb_id integer, + name_pcode_nf varchar(5), + name_pcode_sf varchar(5), + md5sum varchar(32) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/company_name.parquet'; + +CREATE EXTERNAL TABLE company_type ( + id integer unsigned NOT NULL, + kind varchar(32) NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/company_type.parquet'; + +CREATE EXTERNAL TABLE complete_cast ( + id integer unsigned NOT NULL, + movie_id integer, + subject_id integer NOT NULL, + status_id integer NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/complete_cast.parquet'; + +CREATE EXTERNAL TABLE info_type ( + id integer unsigned NOT NULL, + info varchar(32) NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/info_type.parquet'; + +CREATE EXTERNAL TABLE keyword ( + id integer unsigned NOT NULL, + keyword varchar(74) NOT NULL, + phonetic_code varchar(5) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/keyword.parquet'; + +CREATE EXTERNAL TABLE kind_type ( + id integer unsigned NOT NULL, + kind varchar(15) NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/kind_type.parquet'; + +CREATE EXTERNAL TABLE link_type ( + id integer unsigned NOT NULL, + link varchar(32) NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/link_type.parquet'; + +CREATE EXTERNAL TABLE movie_companies ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + company_id integer NOT NULL, + company_type_id integer NOT NULL, + note varchar(208) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/movie_companies.parquet'; + +CREATE EXTERNAL TABLE movie_info ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + info_type_id integer NOT NULL, + info varchar(8000) NOT NULL, + note varchar(387) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/movie_info.parquet'; + +CREATE EXTERNAL TABLE movie_info_idx ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + info_type_id integer NOT NULL, + info varchar(10) NOT NULL, + note varchar(1) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/movie_info_idx.parquet'; + +CREATE EXTERNAL TABLE movie_keyword ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + keyword_id integer NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/movie_keyword.parquet'; + +CREATE EXTERNAL TABLE movie_link ( + id integer unsigned NOT NULL, + movie_id integer NOT NULL, + linked_movie_id integer NOT NULL, + link_type_id integer NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/movie_link.parquet'; + +CREATE EXTERNAL TABLE name ( + id integer unsigned NOT NULL, + name varchar(106) NOT NULL, + imdb_index varchar(12), + imdb_id integer, + gender varchar(1), + name_pcode_cf varchar(5), + name_pcode_nf varchar(5), + surname_pcode varchar(5), + md5sum varchar(32) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/name.parquet'; + +CREATE EXTERNAL TABLE person_info ( + id integer unsigned NOT NULL, + person_id integer NOT NULL, + info_type_id integer NOT NULL, + info text NOT NULL, + note varchar(430) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/person_info.parquet'; + +CREATE EXTERNAL TABLE role_type ( + id integer unsigned NOT NULL, + role varchar(32) NOT NULL +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/role_type.parquet'; + +CREATE EXTERNAL TABLE title ( + id integer unsigned NOT NULL, + title varchar(334) NOT NULL, + imdb_index varchar(12), + kind_id integer NOT NULL, + production_year integer, + imdb_id integer, + phonetic_code varchar(5), + episode_of_id integer, + season_nr integer, + episode_nr integer, + series_years varchar(49), + md5sum varchar(32) +) STORED AS PARQUET LOCATION '${DATA_DIR:-data}/imdb/title.parquet'; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..a3d65c01d3ac2 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q01.benchmark @@ -0,0 +1,12 @@ + +name Q01 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q1: INNER 10K x 10K | LOW 0.1% +SELECT * +FROM range(10000) AS t1 + JOIN range(10000) AS t2 + ON (t1.value + t2.value) % 1000 = 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..6e81be67f16f1 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q02.benchmark @@ -0,0 +1,12 @@ + +name Q02 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q2: INNER 10K x 10K | Medium 20% +SELECT * +FROM range(10000) AS t1 + JOIN range(10000) AS t2 + ON (t1.value + t2.value) % 5 = 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..e561fd7c47030 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q03.benchmark @@ -0,0 +1,12 @@ + +name Q03 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q3: INNER 10K x 10K | High 90% +SELECT * +FROM range(10000) AS t1 + JOIN range(10000) AS t2 + ON (t1.value + t2.value) % 10 <> 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..0dac2d78a50b2 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q04.benchmark @@ -0,0 +1,12 @@ + +name Q04 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q4: INNER 30K x 30K | Medium 20% +SELECT * +FROM range(30000) AS t1 + JOIN range(30000) AS t2 + ON (t1.value + t2.value) % 5 = 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..714c9ded43b72 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q05.benchmark @@ -0,0 +1,12 @@ + +name Q05 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q5: INNER 10K x 200K | LOW 0.1% (small to large) +SELECT * +FROM range(10000) AS t1 + JOIN range(200000) AS t2 + ON (t1.value + t2.value) % 1000 = 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..cb40e71b9db38 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q06.benchmark @@ -0,0 +1,12 @@ + +name Q06 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q6: INNER 200K x 10K | LOW 0.1% (large to small) +SELECT * +FROM range(200000) AS t1 + JOIN range(10000) AS t2 + ON (t1.value + t2.value) % 1000 = 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..71c29bb123c9b --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q07.benchmark @@ -0,0 +1,12 @@ + +name Q07 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q7: RIGHT OUTER 10K x 200K | LOW 0.1% +SELECT * +FROM range(10000) AS t1 + RIGHT JOIN range(200000) AS t2 + ON (t1.value + t2.value) % 1000 = 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..b7e1abc67bf55 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q08.benchmark @@ -0,0 +1,12 @@ + +name Q08 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q8: LEFT OUTER 200K x 10K | LOW 0.1% +SELECT * +FROM range(200000) AS t1 + LEFT JOIN range(10000) AS t2 + ON (t1.value + t2.value) % 1000 = 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..c505717008686 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q09.benchmark @@ -0,0 +1,12 @@ + +name Q09 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q9: FULL OUTER 30K x 30K | LOW 0.1% +SELECT * +FROM range(30000) AS t1 + FULL JOIN range(30000) AS t2 + ON (t1.value + t2.value) % 1000 = 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..f71fa2ebea1b7 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q10.benchmark @@ -0,0 +1,12 @@ + +name Q10 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q10: FULL OUTER 30K x 30K | High 90% +SELECT * +FROM range(30000) AS t1 + FULL JOIN range(30000) AS t2 + ON (t1.value + t2.value) % 10 <> 0; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q11.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q11.benchmark new file mode 100644 index 0000000000000..f54ea79f9d3af --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q11.benchmark @@ -0,0 +1,12 @@ + +name Q11 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q11: INNER 30K x 30K | MEDIUM 50% | cheap predicate +SELECT * +FROM range(30000) AS t1 + INNER JOIN range(30000) AS t2 + ON (t1.value > t2.value); diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q12.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q12.benchmark new file mode 100644 index 0000000000000..9010716a858a1 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q12.benchmark @@ -0,0 +1,12 @@ + +name Q12 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q12: FULL OUTER 30K x 30K | MEDIUM 50% | cheap predicate +SELECT * +FROM range(30000) AS t1 + FULL JOIN range(30000) AS t2 + ON (t1.value > t2.value); diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q13.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q13.benchmark new file mode 100644 index 0000000000000..6e9069bfabb6e --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q13.benchmark @@ -0,0 +1,12 @@ + +name Q13 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q13: LEFT SEMI 30K x 30K | HIGH 99.9% +SELECT t1.* +FROM range(30000) AS t1 + LEFT SEMI JOIN range(30000) AS t2 +ON t1.value < t2.value; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q14.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q14.benchmark new file mode 100644 index 0000000000000..85d95de966094 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q14.benchmark @@ -0,0 +1,12 @@ + +name Q14 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q14: LEFT ANTI 30K x 30K | LOW 0.003% +SELECT t1.* +FROM range(30000) AS t1 + LEFT ANTI JOIN range(30000) AS t2 +ON t1.value < t2.value; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q15.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q15.benchmark new file mode 100644 index 0000000000000..7d9e2adbe7da2 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q15.benchmark @@ -0,0 +1,12 @@ + +name Q15 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q15: RIGHT SEMI 30K x 30K | HIGH 99.9% +SELECT t1.* +FROM range(30000) AS t2 + RIGHT SEMI JOIN range(30000) AS t1 +ON t2.value < t1.value; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q16.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q16.benchmark new file mode 100644 index 0000000000000..c9237f88a5dc5 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q16.benchmark @@ -0,0 +1,12 @@ + +name Q16 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q16: RIGHT ANTI 30K x 30K | LOW 0.003% +SELECT t1.* +FROM range(30000) AS t2 + RIGHT ANTI JOIN range(30000) AS t1 +ON t2.value < t1.value; diff --git a/benchmarks/sql_benchmarks/nlj/benchmarks/q17.benchmark b/benchmarks/sql_benchmarks/nlj/benchmarks/q17.benchmark new file mode 100644 index 0000000000000..f4243a52dbb20 --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/benchmarks/q17.benchmark @@ -0,0 +1,14 @@ + +name Q17 +group nlj + +expect_plan NestedLoopJoinExec + +run +-- Q17: LEFT MARK | HIGH 99.9% +SELECT * +FROM range(30000) AS t2(k2) +WHERE k2 > 0 + OR EXISTS (SELECT 1 + FROM range(30000) AS t1(k1) + WHERE t2.k2 > t1.k1); diff --git a/benchmarks/sql_benchmarks/nlj/nlj.suite b/benchmarks/sql_benchmarks/nlj/nlj.suite new file mode 100644 index 0000000000000..21b4cb298cd8e --- /dev/null +++ b/benchmarks/sql_benchmarks/nlj/nlj.suite @@ -0,0 +1,11 @@ +description = "Nested-loop join SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- nlj" +description = "Run all nested-loop join queries." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- nlj --query 7" +description = "Run nested-loop join query 7." diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q30.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q30.benchmark new file mode 100644 index 0000000000000..760ea2ca902a4 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q30.benchmark @@ -0,0 +1,7 @@ +subgroup cardinality + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=cardinality +QPAD=30 +DATASET=ints +NAME=cardinality_q30_k2 diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q31.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q31.benchmark new file mode 100644 index 0000000000000..74f22715d1eb6 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q31.benchmark @@ -0,0 +1,7 @@ +subgroup cardinality + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=cardinality +QPAD=31 +DATASET=ints +NAME=cardinality_q31_k4 diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q32.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q32.benchmark new file mode 100644 index 0000000000000..b6b69c3852361 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q32.benchmark @@ -0,0 +1,7 @@ +subgroup cardinality + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=cardinality +QPAD=32 +DATASET=ints +NAME=cardinality_q32_k8 diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q33.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q33.benchmark new file mode 100644 index 0000000000000..1260e68137860 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cardinality/q33.benchmark @@ -0,0 +1,7 @@ +subgroup cardinality + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=cardinality +QPAD=33 +DATASET=ints +NAME=cardinality_q33_k16 diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q70.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q70.benchmark new file mode 100644 index 0000000000000..ef20f7dc495b8 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q70.benchmark @@ -0,0 +1,7 @@ +subgroup correlation + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=correlation +QPAD=70 +DATASET=corr +NAME=correlation_q70_independent diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q71.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q71.benchmark new file mode 100644 index 0000000000000..8875f6c44e359 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q71.benchmark @@ -0,0 +1,7 @@ +subgroup correlation + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=correlation +QPAD=71 +DATASET=corr +NAME=correlation_q71_positive diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q72.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q72.benchmark new file mode 100644 index 0000000000000..8109f1439aedb --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q72.benchmark @@ -0,0 +1,7 @@ +subgroup correlation + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=correlation +QPAD=72 +DATASET=corr +NAME=correlation_q72_anti diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q73.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q73.benchmark new file mode 100644 index 0000000000000..cc3f7bcf54901 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/correlation/q73.benchmark @@ -0,0 +1,7 @@ +subgroup correlation + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=correlation +QPAD=73 +DATASET=corrproxy +NAME=correlation_q73_redundant_proxy diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cost/q10.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cost/q10.benchmark new file mode 100644 index 0000000000000..9b864b859457d --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cost/q10.benchmark @@ -0,0 +1,7 @@ +subgroup cost + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=cost +QPAD=10 +DATASET=mixed +NAME=cost_q10_expensive_first diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cost/q11.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cost/q11.benchmark new file mode 100644 index 0000000000000..296ea443b3fec --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/cost/q11.benchmark @@ -0,0 +1,7 @@ +subgroup cost + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=cost +QPAD=11 +DATASET=mixed +NAME=cost_q11_cheap_first diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q01.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q01.benchmark new file mode 100644 index 0000000000000..abedd1d580831 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q01.benchmark @@ -0,0 +1,7 @@ +subgroup costsel + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=costsel +QPAD=01 +DATASET=markers +NAME=costsel_q01_regexp_selective_last diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q02.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q02.benchmark new file mode 100644 index 0000000000000..f50aab66427ec --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q02.benchmark @@ -0,0 +1,7 @@ +subgroup costsel + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=costsel +QPAD=02 +DATASET=markers +NAME=costsel_q02_regexp_selective_first diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q03.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q03.benchmark new file mode 100644 index 0000000000000..10c4ce184eb34 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/costsel/q03.benchmark @@ -0,0 +1,7 @@ +subgroup costsel + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=costsel +QPAD=03 +DATASET=mixed +NAME=costsel_q03_cheap_unselective_then_expensive_selective diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/drift/q80.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/drift/q80.benchmark new file mode 100644 index 0000000000000..970adc53f8017 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/drift/q80.benchmark @@ -0,0 +1,7 @@ +subgroup drift + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=drift +QPAD=80 +DATASET=drift +NAME=drift_q80_a_then_b diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/drift/q81.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/drift/q81.benchmark new file mode 100644 index 0000000000000..93cde75ffef87 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/drift/q81.benchmark @@ -0,0 +1,7 @@ +subgroup drift + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=drift +QPAD=81 +DATASET=drift +NAME=drift_q81_b_then_a diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/neutral/q60.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/neutral/q60.benchmark new file mode 100644 index 0000000000000..039fee622b48b --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/neutral/q60.benchmark @@ -0,0 +1,7 @@ +subgroup neutral + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=neutral +QPAD=60 +DATASET=ints +NAME=neutral_q60_cheap_uniform diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/neutral/q61.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/neutral/q61.benchmark new file mode 100644 index 0000000000000..edaf89b471c5f --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/neutral/q61.benchmark @@ -0,0 +1,7 @@ +subgroup neutral + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=neutral +QPAD=61 +DATASET=markers +NAME=neutral_q61_expensive_uniform diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q50.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q50.benchmark new file mode 100644 index 0000000000000..0bef31e14f402 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q50.benchmark @@ -0,0 +1,8 @@ +subgroup scale + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=scale +QPAD=50 +DATASET=mixed +PRED_ROWS=5000 +NAME=scale_q50_5k diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q51.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q51.benchmark new file mode 100644 index 0000000000000..8f1315fb113b1 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q51.benchmark @@ -0,0 +1,8 @@ +subgroup scale + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=scale +QPAD=51 +DATASET=mixed +PRED_ROWS=100000 +NAME=scale_q51_100k diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q52.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q52.benchmark new file mode 100644 index 0000000000000..7ddbfc19b443d --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q52.benchmark @@ -0,0 +1,8 @@ +subgroup scale + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=scale +QPAD=52 +DATASET=mixed +PRED_ROWS=5000000 +NAME=scale_q52_5m diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q53.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q53.benchmark new file mode 100644 index 0000000000000..6cea5c44a108b --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/scale/q53.benchmark @@ -0,0 +1,8 @@ +subgroup scale + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=scale +QPAD=53 +DATASET=mixed +PRED_ROWS=50000000 +NAME=scale_q53_50m diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/selectivity/q20.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/selectivity/q20.benchmark new file mode 100644 index 0000000000000..077a62650d2f0 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/selectivity/q20.benchmark @@ -0,0 +1,7 @@ +subgroup selectivity + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=selectivity +QPAD=20 +DATASET=ints +NAME=selectivity_q20_unselective_first diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/selectivity/q21.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/selectivity/q21.benchmark new file mode 100644 index 0000000000000..24fc6ef4cd62f --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/selectivity/q21.benchmark @@ -0,0 +1,7 @@ +subgroup selectivity + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=selectivity +QPAD=21 +DATASET=ints +NAME=selectivity_q21_selective_first diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q40.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q40.benchmark new file mode 100644 index 0000000000000..df66cf16a37ec --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q40.benchmark @@ -0,0 +1,8 @@ +subgroup width + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=width +QPAD=40 +DATASET=markers +PRED_FILL=2 +NAME=width_q40_narrow diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q41.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q41.benchmark new file mode 100644 index 0000000000000..c260dc9985a0c --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q41.benchmark @@ -0,0 +1,8 @@ +subgroup width + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=width +QPAD=41 +DATASET=markers +PRED_FILL=30 +NAME=width_q41_wide diff --git a/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q42.benchmark b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q42.benchmark new file mode 100644 index 0000000000000..988ff59c70fe5 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/benchmarks/width/q42.benchmark @@ -0,0 +1,8 @@ +subgroup width + +template sql_benchmarks/predicate_eval/predicate_eval.benchmark.template +SUBGROUP=width +QPAD=42 +DATASET=markers +PRED_FILL=170 +NAME=width_q42_xwide diff --git a/benchmarks/sql_benchmarks/predicate_eval/init/cleanup.sql b/benchmarks/sql_benchmarks/predicate_eval/init/cleanup.sql new file mode 100644 index 0000000000000..48f076a9fa652 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/init/cleanup.sql @@ -0,0 +1 @@ +DROP TABLE IF EXISTS t; diff --git a/benchmarks/sql_benchmarks/predicate_eval/load/corr.sql b/benchmarks/sql_benchmarks/predicate_eval/load/corr.sql new file mode 100644 index 0000000000000..2d7ceb73e608d --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/load/corr.sql @@ -0,0 +1,19 @@ +-- Correlation dataset: a base column plus derived columns that control the +-- *conditional* selectivity of one predicate given another (its selectivity +-- among the rows that already passed the other). +-- +-- x uniform [0,100) +-- x_pos = x (perfectly positively correlated: `x 0) +-- 'bbb' present in ~86% of rows (value % 7 <> 0) +-- 'ccc' present in ~80% of rows (value % 5 <> 0) +-- 'ddd' present in ~75% of rows (value % 4 <> 0) +-- 'rare' present in ~0.1% of rows (value % 1009 = 5) <- the selective one +-- +-- PRED_FILL sets the filler width per marker (the string-column width knob: ~6*PRED_FILL +-- chars per row), and PRED_ROWS sizes the table. +CREATE TABLE t AS +SELECT + repeat('q', ${PRED_FILL:-30}) + || CASE WHEN value % 10 <> 0 THEN 'aaa' ELSE 'zzz' END + || repeat('q', ${PRED_FILL:-30}) + || CASE WHEN value % 7 <> 0 THEN 'bbb' ELSE 'zzz' END + || repeat('q', ${PRED_FILL:-30}) + || CASE WHEN value % 5 <> 0 THEN 'ccc' ELSE 'zzz' END + || repeat('q', ${PRED_FILL:-30}) + || CASE WHEN value % 4 <> 0 THEN 'ddd' ELSE 'zzz' END + || repeat('q', ${PRED_FILL:-30}) + || CASE WHEN value % 1009 = 5 THEN 'rare' ELSE 'zzzz' END + || repeat('q', ${PRED_FILL:-30}) AS s +FROM generate_series(1, ${PRED_ROWS:-1000000}); diff --git a/benchmarks/sql_benchmarks/predicate_eval/load/mixed.sql b/benchmarks/sql_benchmarks/predicate_eval/load/mixed.sql new file mode 100644 index 0000000000000..a51c1040daca6 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/load/mixed.sql @@ -0,0 +1,26 @@ +-- Mixed-cost dataset: cheap integer columns (`cN < k` ~ k% selectivity) +-- alongside one wide string column carrying three markers matched by expensive +-- `regexp_like`: +-- +-- 'rare' present in ~0.1% of rows (value % 1009 = 5) +-- 'ten' present in ~10% of rows (value % 10 = 0) +-- 'aaa' present in ~90% of rows (value % 10 <> 0) +-- +-- This lets a single table mix cheap integer compares with expensive regexp +-- scans at independently chosen selectivities (e.g. a cheap, unselective compare +-- next to an expensive, selective regexp). PRED_FILL is the string-width knob; +-- PRED_ROWS sizes the table. +CREATE TABLE t AS +SELECT + (value * 1) % 100 AS c0, + (value * 3) % 100 AS c1, + (value * 7) % 100 AS c2, + (value * 9) % 100 AS c3, + repeat('q', ${PRED_FILL:-30}) + || CASE WHEN value % 1009 = 5 THEN 'rare' ELSE 'zzzz' END + || repeat('q', ${PRED_FILL:-30}) + || CASE WHEN value % 10 = 0 THEN 'ten' ELSE 'zzz' END + || repeat('q', ${PRED_FILL:-30}) + || CASE WHEN value % 10 <> 0 THEN 'aaa' ELSE 'zzz' END + || repeat('q', ${PRED_FILL:-30}) AS s +FROM generate_series(1, ${PRED_ROWS:-1000000}); diff --git a/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.benchmark.template b/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.benchmark.template new file mode 100644 index 0000000000000..0030a7e946eca --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.benchmark.template @@ -0,0 +1,34 @@ +# Shared template for every predicate_eval benchmark. Each qNN.benchmark sets +# its `subgroup` directive and then includes this template with parameters: +# SUBGROUP subgroup name, also the query sub-directory (e.g. costsel) +# QPAD zero-padded query id, also the query file stem (e.g. 01) +# DATASET load script stem under load/ (e.g. markers) +# NAME criterion display name (e.g. costsel_q01_regexp_selective_last) +# Optional (consumed by the load scripts via ${...:-default}): +# PRED_ROWS synthetic row count (default 1_000_000) +# PRED_FILL filler chars per marker = string-column width knob (default 30) +# +# The run SQL lives in queries/${SUBGROUP}/q${QPAD}.sql so the WHERE clause is +# readable on its own. The table is always named `t`, so the assert and cleanup +# are uniform across datasets. +# +# The suite is implementation-agnostic and sets no engine config of its own: it +# measures DataFusion's built-in left-deep `AND` short-circuit by default. To +# evaluate a predicate-ordering system under test, set its native config via the +# environment (the bench harness builds its SessionContext with +# SessionConfig::from_env), e.g. +# DATAFUSION_EXECUTION_ADAPTIVE_FILTER_REORDERING=true + +load sql_benchmarks/predicate_eval/load/${DATASET}.sql + +name ${NAME} +group predicate_eval + +assert I +SELECT count(*) > 0 FROM t; +---- +true + +run sql_benchmarks/predicate_eval/queries/${SUBGROUP}/q${QPAD}.sql + +cleanup sql_benchmarks/predicate_eval/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.suite b/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.suite new file mode 100644 index 0000000000000..af1a326cd8c51 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/predicate_eval.suite @@ -0,0 +1,31 @@ +description = "Conjunctive filter evaluation micro-benchmarks covering predicate cost, selectivity, cardinality, width, scale, correlation, and drift" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[options]] +name = "rows" +short = "r" +env = "PRED_ROWS" +default = "1000000" +values = ["1000000", "..."] +help = "Sets the number of rows in generated predicate-evaluation datasets." + +[[options]] +name = "fill" +short = "f" +env = "PRED_FILL" +default = "30" +values = ["2", "30", "170", "..."] +help = "Sets the filler width for generated string columns." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- predicate_eval" +description = "Run all predicate-evaluation subgroups with default data sizes." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- predicate_eval --query 20 --subgroup selectivity" +description = "Run selectivity query 20." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- predicate_eval --subgroup width -r 500000 -f 170" +description = "Run the width subgroup with 500,000 extra-wide rows." diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q30.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q30.sql new file mode 100644 index 0000000000000..3be840e917383 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q30.sql @@ -0,0 +1,6 @@ +-- Hidden: cheap integer compares; `c1 < 5` matches ~5%, the `c0 < 90` family +-- ~90%. k = 2 here. q30..q33 sweep k = 2/4/8/16 with one ~5% predicate written +-- last among ~90% ones. +SELECT count(*) FROM t +WHERE c0 < 90 + AND c1 < 5; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q31.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q31.sql new file mode 100644 index 0000000000000..4ba84f8124be9 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q31.sql @@ -0,0 +1,6 @@ +-- k = 4: three ~90% compares followed by one ~5% compare. See q30. +SELECT count(*) FROM t +WHERE c0 < 90 + AND c1 < 90 + AND c2 < 90 + AND c3 < 5; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q32.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q32.sql new file mode 100644 index 0000000000000..d9e920cc62574 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q32.sql @@ -0,0 +1,10 @@ +-- k = 8: seven ~90% compares followed by one ~5% compare. See q30. +SELECT count(*) FROM t +WHERE c0 < 90 + AND c1 < 90 + AND c2 < 90 + AND c3 < 90 + AND c4 < 90 + AND c5 < 90 + AND c6 < 90 + AND c7 < 5; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q33.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q33.sql new file mode 100644 index 0000000000000..2408427ab7632 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/cardinality/q33.sql @@ -0,0 +1,18 @@ +-- k = 16: fifteen ~90% compares followed by one ~5% compare. See q30. +SELECT count(*) FROM t +WHERE c0 < 90 + AND c1 < 90 + AND c2 < 90 + AND c3 < 90 + AND c4 < 90 + AND c5 < 90 + AND c6 < 90 + AND c7 < 90 + AND c8 < 90 + AND c9 < 90 + AND c10 < 90 + AND c11 < 90 + AND c12 < 90 + AND c13 < 90 + AND c14 < 90 + AND c15 < 5; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q70.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q70.sql new file mode 100644 index 0000000000000..86e33534c705d --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q70.sql @@ -0,0 +1,6 @@ +-- Hidden: `x` and `ind` are independent, each ~20%, so the conjunction matches +-- ~4% and the second predicate is just as selective among the first's survivors +-- as on its own. Baseline for the correlation sweep. cf. q71, q72. +SELECT count(*) FROM t +WHERE x < 20 + AND ind < 20; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q71.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q71.sql new file mode 100644 index 0000000000000..eda61cc289e92 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q71.sql @@ -0,0 +1,6 @@ +-- Hidden: `x_pos` is a copy of `x`, so `x < 20 AND x_pos < 20` still matches +-- ~20% (not the ~4% independence would imply) -- the second predicate removes +-- none of the first's survivors. cf. q70. +SELECT count(*) FROM t +WHERE x < 20 + AND x_pos < 20; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q72.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q72.sql new file mode 100644 index 0000000000000..ff987524da6ed --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q72.sql @@ -0,0 +1,6 @@ +-- Hidden: `x_anti` is `99 - x`, so `x < 50 AND x_anti < 50` is empty -- the +-- second predicate removes all of the first's survivors, though each matches +-- ~50% alone. cf. q70. +SELECT count(*) FROM t +WHERE x < 50 + AND x_anti < 50; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q73.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q73.sql new file mode 100644 index 0000000000000..5e1e822e92eca --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/correlation/q73.sql @@ -0,0 +1,14 @@ +-- Hidden: `c0 = 1` is a perfect proxy for the s1/s2/s3 regexes -- after the +-- cheap proxy, each of those keeps every survivor while the equally selective +-- (~30%) s4 regex still discards ~70%. The optimal order is [c0, s4, s1/s2/s3] +-- (one informative regex on 30% of rows, the three redundant ones on 9%), but +-- the four regexes are marginally identical -- same width, same marker offset, +-- same cost, same selectivity -- so ranking them takes their *joint* +-- distribution with the proxy. Written with the redundant regexes first, +-- grouped with their proxy, as an author naturally would. +SELECT count(*) FROM t +WHERE c0 = 1 + AND regexp_like(s1, 'a.a') + AND regexp_like(s2, 'c.c') + AND regexp_like(s3, 'd.d') + AND regexp_like(s4, 'b.b'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/cost/q10.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/cost/q10.sql new file mode 100644 index 0000000000000..b089ebc7a192a --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/cost/q10.sql @@ -0,0 +1,6 @@ +-- Hidden: both predicates match ~10%, but `regexp_like(s, 'ten')` scans the +-- string (expensive) while `c0 < 10` is a cheap compare. Equal selectivity, +-- unequal cost; expensive one written first. cf. q11 (opposite order). +SELECT count(*) FROM t +WHERE regexp_like(s, 'ten') + AND c0 < 10; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/cost/q11.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/cost/q11.sql new file mode 100644 index 0000000000000..82d748c93b3b2 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/cost/q11.sql @@ -0,0 +1,5 @@ +-- Same two predicates as q10 (both ~10%; regexp expensive, compare cheap), +-- opposite written order. cf. q10. +SELECT count(*) FROM t +WHERE c0 < 10 + AND regexp_like(s, 'ten'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q01.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q01.sql new file mode 100644 index 0000000000000..bc029ed5d8297 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q01.sql @@ -0,0 +1,10 @@ +-- Hidden in the data: the five markers have very different selectivities -- +-- 'aaa' ~90%, 'bbb' ~86%, 'ccc' ~80%, 'ddd' ~75%, 'rare' ~0.1% -- while every +-- regexp_like costs about the same. 'rare' (most selective) is written last. +-- cf. q02 (most selective written first). +SELECT count(*) FROM t +WHERE regexp_like(s, 'aaa') + AND regexp_like(s, 'bbb') + AND regexp_like(s, 'ccc') + AND regexp_like(s, 'ddd') + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q02.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q02.sql new file mode 100644 index 0000000000000..7f7fc61831ff0 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q02.sql @@ -0,0 +1,8 @@ +-- Same predicates and hidden selectivities as q01 ('rare' ~0.1% is the +-- selective one, the rest 75-90%), but with 'rare' written first. cf. q01. +SELECT count(*) FROM t +WHERE regexp_like(s, 'rare') + AND regexp_like(s, 'aaa') + AND regexp_like(s, 'bbb') + AND regexp_like(s, 'ccc') + AND regexp_like(s, 'ddd'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q03.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q03.sql new file mode 100644 index 0000000000000..a583a498b211c --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/costsel/q03.sql @@ -0,0 +1,6 @@ +-- Hidden: `c0 < 90` matches ~90% (cheap integer compare); `regexp_like(s, +-- 'rare')` matches ~0.1% (scans the wide string). The cheaper predicate is the +-- less selective one. +SELECT count(*) FROM t +WHERE c0 < 90 + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/drift/q80.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/drift/q80.sql new file mode 100644 index 0000000000000..b8cb61e85a478 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/drift/q80.sql @@ -0,0 +1,7 @@ +-- The non-obvious property: selectivity changes across the scan. Rows arrive in +-- `seq` order; `a_sel = 0` matches ~0.1% in the first 10% of rows and ~50% +-- after, `b_sel = 0` is the mirror -- so which predicate is more selective flips +-- partway through. cf. q81 (opposite order). +SELECT count(*) FROM t +WHERE a_sel = 0 + AND b_sel = 0; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/drift/q81.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/drift/q81.sql new file mode 100644 index 0000000000000..d65ef475cc0e0 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/drift/q81.sql @@ -0,0 +1,5 @@ +-- Same drifting predicates as q80 (a_sel/b_sel flip which is more selective +-- partway through the scan), opposite written order. cf. q80. +SELECT count(*) FROM t +WHERE b_sel = 0 + AND a_sel = 0; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/neutral/q60.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/neutral/q60.sql new file mode 100644 index 0000000000000..b217f56953272 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/neutral/q60.sql @@ -0,0 +1,7 @@ +-- Hidden: four integer compares of equal cost, each ~50% selective. Nothing is +-- selective and the costs are equal, so the predicates are interchangeable. +SELECT count(*) FROM t +WHERE c0 < 50 + AND c1 < 50 + AND c2 < 50 + AND c3 < 50; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/neutral/q61.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/neutral/q61.sql new file mode 100644 index 0000000000000..7029a3d9f8f7d --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/neutral/q61.sql @@ -0,0 +1,8 @@ +-- Hidden: four regexp scans of about equal cost, all unselective ('aaa' ~90%, +-- 'bbb' ~86%, 'ccc' ~80%, 'ddd' ~75%). Like q60 the predicates are +-- interchangeable, but here each one is expensive. +SELECT count(*) FROM t +WHERE regexp_like(s, 'aaa') + AND regexp_like(s, 'bbb') + AND regexp_like(s, 'ccc') + AND regexp_like(s, 'ddd'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q50.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q50.sql new file mode 100644 index 0000000000000..03a0f1c0db285 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q50.sql @@ -0,0 +1,6 @@ +-- Same predicates as costsel/q03 (`c0 < 90` ~90% cheap, `regexp_like(s, 'rare')` +-- ~0.1% expensive). q50..q53 sweep table size; here PRED_ROWS=5_000, roughly a +-- single batch. +SELECT count(*) FROM t +WHERE c0 < 90 + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q51.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q51.sql new file mode 100644 index 0000000000000..28174a5df4f44 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q51.sql @@ -0,0 +1,4 @@ +-- q50 at PRED_ROWS=100_000 (~12 batches). See q50. +SELECT count(*) FROM t +WHERE c0 < 90 + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q52.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q52.sql new file mode 100644 index 0000000000000..74938c4634f78 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q52.sql @@ -0,0 +1,4 @@ +-- q50 at PRED_ROWS=5_000_000 (~610 batches). See q50. +SELECT count(*) FROM t +WHERE c0 < 90 + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q53.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q53.sql new file mode 100644 index 0000000000000..8edb4d4a057d2 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/scale/q53.sql @@ -0,0 +1,4 @@ +-- q50 at PRED_ROWS=50_000_000 (~6100 batches); builds a ~9 GB table. See q50. +SELECT count(*) FROM t +WHERE c0 < 90 + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/selectivity/q20.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/selectivity/q20.sql new file mode 100644 index 0000000000000..3638f757a720d --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/selectivity/q20.sql @@ -0,0 +1,6 @@ +-- Hidden: two equally cheap integer compares of unequal selectivity -- `c4 < 95` +-- matches ~95%, `c0 < 5` matches ~5%. Less selective one written first. +-- cf. q21 (opposite order). +SELECT count(*) FROM t +WHERE c4 < 95 + AND c0 < 5; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/selectivity/q21.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/selectivity/q21.sql new file mode 100644 index 0000000000000..5181faf38784f --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/selectivity/q21.sql @@ -0,0 +1,5 @@ +-- Same two equally-cheap compares as q20 (`c4 < 95` ~95%, `c0 < 5` ~5%), +-- opposite written order. cf. q20. +SELECT count(*) FROM t +WHERE c0 < 5 + AND c4 < 95; diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/width/q40.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/width/q40.sql new file mode 100644 index 0000000000000..1b3df3e937eb3 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/width/q40.sql @@ -0,0 +1,9 @@ +-- Same predicate set and hidden selectivities as costsel/q01 ('rare' ~0.1%, the +-- rest 75-90%); only the string-column width differs across q40/q41/q42. Narrow: +-- PRED_FILL=2, ~12 chars/row. +SELECT count(*) FROM t +WHERE regexp_like(s, 'aaa') + AND regexp_like(s, 'bbb') + AND regexp_like(s, 'ccc') + AND regexp_like(s, 'ddd') + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/width/q41.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/width/q41.sql new file mode 100644 index 0000000000000..a03b576d9c959 --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/width/q41.sql @@ -0,0 +1,7 @@ +-- q40 with wide strings: PRED_FILL=30, ~186 chars/row. See q40. +SELECT count(*) FROM t +WHERE regexp_like(s, 'aaa') + AND regexp_like(s, 'bbb') + AND regexp_like(s, 'ccc') + AND regexp_like(s, 'ddd') + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/predicate_eval/queries/width/q42.sql b/benchmarks/sql_benchmarks/predicate_eval/queries/width/q42.sql new file mode 100644 index 0000000000000..cb55d828e32ab --- /dev/null +++ b/benchmarks/sql_benchmarks/predicate_eval/queries/width/q42.sql @@ -0,0 +1,7 @@ +-- q40 with extra-wide strings: PRED_FILL=170, ~1KB/row. See q40. +SELECT count(*) FROM t +WHERE regexp_like(s, 'aaa') + AND regexp_like(s, 'bbb') + AND regexp_like(s, 'ccc') + AND regexp_like(s, 'ddd') + AND regexp_like(s, 'rare'); diff --git a/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..a7ec837319af3 --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q01.benchmark @@ -0,0 +1,22 @@ +-- LEFT JOIN, ORDER BY a column from the preserved (left) side, small LIMIT. +-- Canonical push_down_topk_through_join case: the TopK can be duplicated +-- below the join over the customer scan so only the top 10 rows (by +-- c_acctbal) are joined against orders. + +name Q01 +group push_down_topk + +load sql_benchmarks/push_down_topk/init/load.sql + +assert I +SELECT COUNT(*) > 0 from customer; +---- +true + +run +SELECT c_custkey, c_acctbal +FROM customer LEFT JOIN orders ON c_custkey = o_custkey +ORDER BY c_acctbal +LIMIT 10; + +cleanup sql_benchmarks/push_down_topk/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..eb70ef34fe739 --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q02.benchmark @@ -0,0 +1,21 @@ +-- RIGHT JOIN, ORDER BY a column from the preserved (right) side. +-- Symmetric to Q01: the TopK is pushed below the join over the orders +-- scan (the right/preserved side). + +name Q02 +group push_down_topk + +load sql_benchmarks/push_down_topk/init/load.sql + +assert I +SELECT COUNT(*) > 0 from orders; +---- +true + +run +SELECT o_orderkey, o_totalprice +FROM customer RIGHT JOIN orders ON c_custkey = o_custkey +ORDER BY o_totalprice +LIMIT 10; + +cleanup sql_benchmarks/push_down_topk/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..503cc45710e91 --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q03.benchmark @@ -0,0 +1,21 @@ +-- LEFT JOIN, multi-column ORDER BY (both columns from the preserved side). +-- All sort exprs must come from the preserved side for the rule to fire; +-- this checks that multi-column sorts are still pushed. + +name Q03 +group push_down_topk + +load sql_benchmarks/push_down_topk/init/load.sql + +assert I +SELECT COUNT(*) > 0 from customer; +---- +true + +run +SELECT c_custkey, c_acctbal, c_nationkey +FROM customer LEFT JOIN orders ON c_custkey = o_custkey +ORDER BY c_acctbal, c_nationkey +LIMIT 100; + +cleanup sql_benchmarks/push_down_topk/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..143455721127c --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q04.benchmark @@ -0,0 +1,21 @@ +-- CROSS JOIN, ORDER BY a column from one side. +-- Cross joins preserve every row from both sides; the rule pushes the +-- TopK below the join over the side referenced by ORDER BY. + +name Q04 +group push_down_topk + +load sql_benchmarks/push_down_topk/init/load.sql + +assert I +SELECT COUNT(*) > 0 from customer; +---- +true + +run +SELECT c_custkey, c_acctbal +FROM customer CROSS JOIN nation +ORDER BY c_acctbal +LIMIT 10; + +cleanup sql_benchmarks/push_down_topk/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..74ed2ec592bc6 --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/benchmarks/q05.benchmark @@ -0,0 +1,23 @@ +-- Negative case: ORDER BY references the probe (non-preserved) side. +-- The rule MUST NOT fire here -- orders is the right side of a LEFT JOIN +-- so it isn't preserved (rows can be NULL when there's no match), and +-- pushing a TopK onto orders would change semantics. Included so the +-- bench captures the no-pushdown path alongside the positive cases. + +name Q05 +group push_down_topk + +load sql_benchmarks/push_down_topk/init/load.sql + +assert I +SELECT COUNT(*) > 0 from orders; +---- +true + +run +SELECT c_custkey, o_totalprice +FROM customer LEFT JOIN orders ON c_custkey = o_custkey +ORDER BY o_totalprice +LIMIT 10; + +cleanup sql_benchmarks/push_down_topk/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/push_down_topk/init/cleanup.sql b/benchmarks/sql_benchmarks/push_down_topk/init/cleanup.sql new file mode 100644 index 0000000000000..9e271dba06ff0 --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/init/cleanup.sql @@ -0,0 +1,5 @@ +DROP TABLE IF EXISTS customer; + +DROP TABLE IF EXISTS orders; + +DROP TABLE IF EXISTS nation; diff --git a/benchmarks/sql_benchmarks/push_down_topk/init/load.sql b/benchmarks/sql_benchmarks/push_down_topk/init/load.sql new file mode 100644 index 0000000000000..f5f5bb641d4e5 --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/init/load.sql @@ -0,0 +1,5 @@ +CREATE EXTERNAL TABLE customer STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/customer/customer.1.parquet'; + +CREATE EXTERNAL TABLE orders STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/orders/orders.1.parquet'; + +CREATE EXTERNAL TABLE nation STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/nation/nation.1.parquet'; diff --git a/benchmarks/sql_benchmarks/push_down_topk/push_down_topk.suite b/benchmarks/sql_benchmarks/push_down_topk/push_down_topk.suite new file mode 100644 index 0000000000000..a70139c7669ca --- /dev/null +++ b/benchmarks/sql_benchmarks/push_down_topk/push_down_topk.suite @@ -0,0 +1,21 @@ +description = "TopK pushdown benchmarks for ORDER BY LIMIT over TPC-H joins" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "..."] +help = "Selects the TPC-H scale factor used by the TopK benchmarks." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- push_down_topk" +description = "Run all TopK pushdown queries at scale factor 1." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- push_down_topk --query 3 --scale-factor 10" +description = "Run TopK pushdown query 3 at scale factor 10." diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..f1d44a6fb3c16 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q01.benchmark @@ -0,0 +1,21 @@ +name Q01 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q1: INNER 1M x 1M | 1:1 +WITH t1_sorted AS ( + SELECT value as key FROM range(1000000) ORDER BY value + ), + t2_sorted AS ( +SELECT value as key FROM range(1000000) ORDER BY value + ) +SELECT t1_sorted.key as k1, t2_sorted.key as k2 +FROM t1_sorted JOIN t2_sorted ON t1_sorted.key = t2_sorted.key + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..cd30f53256407 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q02.benchmark @@ -0,0 +1,25 @@ +name Q02 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q2: INNER 1M x 10M | 1:10 +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted JOIN t2_sorted ON t1_sorted.key = t2_sorted.key + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..02bab2a6850bd --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q03.benchmark @@ -0,0 +1,25 @@ +name Q03 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q3: INNER 1M x 1M | 1:100 +WITH t1_sorted AS ( + SELECT value % 10000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 10000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted JOIN t2_sorted ON t1_sorted.key = t2_sorted.key + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..2442906a3f2f7 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q04.benchmark @@ -0,0 +1,26 @@ +name Q04 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q4: INNER 1M x 10M | 1:10 | 1% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted JOIN t2_sorted ON t1_sorted.key = t2_sorted.key +WHERE t2_sorted.data % 100 = 0 + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..1735d0e0e65ec --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q05.benchmark @@ -0,0 +1,26 @@ +name Q05 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q5: INNER 1M x 1M | 1:100 | 10% +WITH t1_sorted AS ( + SELECT value % 10000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 10000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted JOIN t2_sorted ON t1_sorted.key = t2_sorted.key +WHERE t1_sorted.data <> t2_sorted.data AND t2_sorted.data % 10 = 0 + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..8c18ee164f2a5 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q06.benchmark @@ -0,0 +1,25 @@ +name Q06 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q6: LEFT 1M x 10M | 1:10 +WITH t1_sorted AS ( + SELECT value % 105000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted LEFT JOIN t2_sorted ON t1_sorted.key = t2_sorted.key + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..20619b0948707 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q07.benchmark @@ -0,0 +1,26 @@ +name Q07 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q7: LEFT 1M x 10M | 1:10 | 50% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted LEFT JOIN t2_sorted ON t1_sorted.key = t2_sorted.key +WHERE t2_sorted.data IS NULL OR t2_sorted.data % 2 = 0 + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..7597f2012ab47 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q08.benchmark @@ -0,0 +1,26 @@ +name Q08 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q8: FULL 1M x 1M | 1:10 +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 125000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ) +SELECT t1_sorted.key as k1, t1_sorted.data as d1, + t2_sorted.key as k2, t2_sorted.data as d2 +FROM t1_sorted FULL JOIN t2_sorted ON t1_sorted.key = t2_sorted.key + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..ca0565c6a69a9 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q09.benchmark @@ -0,0 +1,29 @@ +name Q09 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q9: FULL 1M x 10M | 1:10 | 10% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key as k1, t1_sorted.data as d1, + t2_sorted.key as k2, t2_sorted.data as d2 +FROM t1_sorted FULL JOIN t2_sorted ON t1_sorted.key = t2_sorted.key +WHERE (t1_sorted.data IS NULL OR t2_sorted.data IS NULL + OR t1_sorted.data <> t2_sorted.data) + AND (t1_sorted.data IS NULL OR t1_sorted.data % 10 = 0) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..cbba0610c590d --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q10.benchmark @@ -0,0 +1,29 @@ +name Q10 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q10: LEFT SEMI 1M x 10M | 1:10 +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key +FROM range(10000000) +ORDER BY key + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q11.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q11.benchmark new file mode 100644 index 0000000000000..2431b2646ec9c --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q11.benchmark @@ -0,0 +1,31 @@ +name Q11 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q11: LEFT SEMI 1M x 10M | 1:10 | 1% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key + AND t2_sorted.data <> t1_sorted.data + AND t2_sorted.data % 100 = 0 +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q12.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q12.benchmark new file mode 100644 index 0000000000000..79e0a8e8a51cb --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q12.benchmark @@ -0,0 +1,31 @@ +name Q12 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q12: LEFT SEMI 1M x 10M | 1:10 | 50% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key + AND t2_sorted.data <> t1_sorted.data + AND t2_sorted.data % 2 = 0 +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q13.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q13.benchmark new file mode 100644 index 0000000000000..7e13e687434ab --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q13.benchmark @@ -0,0 +1,31 @@ +name Q13 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q13: LEFT SEMI 1M x 10M | 1:10 | 90% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key + AND t2_sorted.data <> t1_sorted.data + AND t2_sorted.data % 10 <> 0 +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q14.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q14.benchmark new file mode 100644 index 0000000000000..a56a0d5863aec --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q14.benchmark @@ -0,0 +1,29 @@ +name Q14 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q14: LEFT ANTI 1M x 10M | 1:10 +WITH t1_sorted AS ( + SELECT value % 105000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key +FROM range(10000000) +ORDER BY key + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE NOT EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q15.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q15.benchmark new file mode 100644 index 0000000000000..bd64d74422f99 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q15.benchmark @@ -0,0 +1,29 @@ +name Q15 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q15: LEFT ANTI 1M x 10M | 1:10 | partial match +WITH t1_sorted AS ( + SELECT value % 120000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key +FROM range(10000000) +ORDER BY key + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE NOT EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q16.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q16.benchmark new file mode 100644 index 0000000000000..282d6374ebd27 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q16.benchmark @@ -0,0 +1,29 @@ +name Q16 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q16: LEFT ANTI 1M x 1M | 1:1 | stress +WITH t1_sorted AS ( + SELECT value % 110000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key +FROM range(1000000) +ORDER BY key + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE NOT EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q17.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q17.benchmark new file mode 100644 index 0000000000000..7f1c9a0ae2485 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q17.benchmark @@ -0,0 +1,26 @@ +name Q17 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q17: INNER 1M x 50M | 1:50 | 5% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(50000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted JOIN t2_sorted ON t1_sorted.key = t2_sorted.key +WHERE t2_sorted.data <> t1_sorted.data AND t2_sorted.data % 20 = 0 + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q18.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q18.benchmark new file mode 100644 index 0000000000000..fac7edd19b0b9 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q18.benchmark @@ -0,0 +1,31 @@ +name Q18 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q18: LEFT SEMI 1M x 50M | 1:50 | 2% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(50000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key + AND t2_sorted.data <> t1_sorted.data + AND t2_sorted.data % 50 = 0 +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q19.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q19.benchmark new file mode 100644 index 0000000000000..a867bb6bff4e2 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q19.benchmark @@ -0,0 +1,29 @@ +name Q19 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q19: LEFT ANTI 1M x 50M | 1:50 | partial match +WITH t1_sorted AS ( + SELECT value % 150000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key +FROM range(50000000) +ORDER BY key + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE NOT EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q20.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q20.benchmark new file mode 100644 index 0000000000000..317c6290f7964 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q20.benchmark @@ -0,0 +1,26 @@ +name Q20 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q20: INNER 1M x 10M | 1:100 + GROUP BY +WITH t1_sorted AS ( + SELECT value % 10000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 10000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, count(*) as cnt +FROM t1_sorted JOIN t2_sorted ON t1_sorted.key = t2_sorted.key +GROUP BY t1_sorted.key + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q21.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q21.benchmark new file mode 100644 index 0000000000000..3fe460ea000f2 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q21.benchmark @@ -0,0 +1,25 @@ +name Q21 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q21: INNER 10M x 10M | unique keys (1:1) | 50% join filter +WITH t1_sorted AS ( + SELECT value as key, value as data +FROM range(10000000) ORDER BY value + ), + t2_sorted AS ( +SELECT value as key, value as data +FROM range(10000000) ORDER BY value + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted JOIN t2_sorted + ON t1_sorted.key = t2_sorted.key + AND t1_sorted.data + t2_sorted.data < 10000000 + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q22.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q22.benchmark new file mode 100644 index 0000000000000..fe0063de5761e --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q22.benchmark @@ -0,0 +1,25 @@ +name Q22 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q22: LEFT 10M x 10M | unique keys (1:1) | 50% join filter +WITH t1_sorted AS ( + SELECT value as key, value as data +FROM range(10000000) ORDER BY value + ), + t2_sorted AS ( +SELECT value as key, value as data +FROM range(10000000) ORDER BY value + ) +SELECT t1_sorted.key, t1_sorted.data as d1, t2_sorted.data as d2 +FROM t1_sorted LEFT JOIN t2_sorted + ON t1_sorted.key = t2_sorted.key + AND t1_sorted.data + t2_sorted.data < 10000000 + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q23.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q23.benchmark new file mode 100644 index 0000000000000..592effd993d3b --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q23.benchmark @@ -0,0 +1,26 @@ +name Q23 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q23: FULL 10M x 10M | unique keys (1:1) | 50% join filter +WITH t1_sorted AS ( + SELECT value as key, value as data +FROM range(10000000) ORDER BY value + ), + t2_sorted AS ( +SELECT value as key, value as data +FROM range(10000000) ORDER BY value + ) +SELECT t1_sorted.key as k1, t1_sorted.data as d1, + t2_sorted.key as k2, t2_sorted.data as d2 +FROM t1_sorted FULL JOIN t2_sorted + ON t1_sorted.key = t2_sorted.key + AND t1_sorted.data + t2_sorted.data < 10000000 + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/benchmarks/q24.benchmark b/benchmarks/sql_benchmarks/smj/benchmarks/q24.benchmark new file mode 100644 index 0000000000000..893eb1fb78733 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/benchmarks/q24.benchmark @@ -0,0 +1,32 @@ +name Q24 +group smj + +init +set datafusion.optimizer.prefer_hash_join=false; + +expect_plan SortMergeJoinExec + +run +-- Q24: LEFT MARK 1M x 10M | 1:10 | 1% +WITH t1_sorted AS ( + SELECT value % 100000 as key, value as data +FROM range(1000000) +ORDER BY key, data + ), + t2_sorted AS ( +SELECT value % 100000 as key, value as data +FROM range(10000000) +ORDER BY key, data + ) +SELECT t1_sorted.key, t1_sorted.data +FROM t1_sorted +WHERE t1_sorted.data < 0 + OR EXISTS ( + SELECT 1 FROM t2_sorted + WHERE t2_sorted.key = t1_sorted.key + AND t2_sorted.data <> t1_sorted.data + AND t2_sorted.data % 100 = 0 +) + +cleanup +reset datafusion.optimizer.prefer_hash_join; diff --git a/benchmarks/sql_benchmarks/smj/smj.suite b/benchmarks/sql_benchmarks/smj/smj.suite new file mode 100644 index 0000000000000..44db22ffe20f5 --- /dev/null +++ b/benchmarks/sql_benchmarks/smj/smj.suite @@ -0,0 +1,11 @@ +description = "Sort-merge join SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- smj" +description = "Run all sort-merge join queries." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- smj --query 12" +description = "Run sort-merge join query 12." diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..b6f1a37e3d03f --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q01.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q01 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q1: 1 sort key (type: INTEGER, cardinality: 7) + 1 payload column +SELECT l_linenumber, l_partkey +FROM lineitem +ORDER BY l_linenumber +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q01.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..1238beb00583a --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q02.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q02 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q2: 1 sort key (type: BIGINT, cardinality: 1.5M) + 1 payload column +SELECT l_orderkey, l_partkey +FROM lineitem +ORDER BY l_orderkey +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q02.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..aadbe86c61602 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q03.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q03 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q3: 1 sort key (type: VARCHAR, cardinality: 4.5M) + 1 payload column +SELECT l_comment, l_partkey +FROM lineitem +ORDER BY l_comment +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q03.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..8119a6c51be33 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q04.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q04 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q4: 2 sort keys {(BIGINT, 1.5M), (INTEGER, 7)} + 1 payload column +SELECT l_orderkey, l_linenumber, l_partkey +FROM lineitem +ORDER BY l_orderkey, l_linenumber +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q04.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..5ee9e610cc3bf --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q05.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q05 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q5: 3 sort keys {(INTEGER, 7), (BIGINT, 10k), (BIGINT, 1.5M)} + no payload column +SELECT l_linenumber, l_suppkey, l_orderkey +FROM lineitem +ORDER BY l_linenumber, l_suppkey, l_orderkey +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q05.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..54ce6fa44341d --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q06.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q06 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q6: 3 sort keys {(INTEGER, 7), (BIGINT, 10k), (BIGINT, 1.5M)} + 1 payload column +SELECT l_linenumber, l_suppkey, l_orderkey, l_partkey +FROM lineitem +ORDER BY l_linenumber, l_suppkey, l_orderkey +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q06.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..8932810cc1f97 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q07.benchmark @@ -0,0 +1,58 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q07 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q7: 3 sort keys {(INTEGER, 7), (BIGINT, 10k), (BIGINT, 1.5M)} + 12 all other columns +SELECT l_linenumber, + l_suppkey, + l_orderkey, + l_partkey, + l_quantity, + l_extendedprice, + l_discount, + l_tax, + l_returnflag, + l_linestatus, + l_shipdate, + l_commitdate, + l_receiptdate, + l_shipinstruct, + l_shipmode +FROM lineitem +ORDER BY l_linenumber, l_suppkey, l_orderkey +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q07.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..f09e6e9f72f21 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q08.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q08 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q8: 4 sort keys {(BIGINT, 1.5M), (BIGINT, 10k), (INTEGER, 7), (VARCHAR, 4.5M)} + no payload column +SELECT l_orderkey, l_suppkey, l_linenumber, l_comment +FROM lineitem +ORDER BY l_orderkey, l_suppkey, l_linenumber, l_comment +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q08.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..5e7a2ea63747a --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q09.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q09 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q9: 4 sort keys {(BIGINT, 1.5M), (BIGINT, 10k), (INTEGER, 7), (VARCHAR, 4.5M)} + 1 payload column +SELECT l_orderkey, l_suppkey, l_linenumber, l_comment, l_partkey +FROM lineitem +ORDER BY l_orderkey, l_suppkey, l_linenumber, l_comment +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q09.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..535393526147e --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q10.benchmark @@ -0,0 +1,59 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q10 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q10: 4 sort keys {(BIGINT, 1.5M), (BIGINT, 10k), (INTEGER, 7), (VARCHAR, 4.5M)} + 12 all other columns +SELECT l_orderkey, + l_suppkey, + l_linenumber, + l_comment, + l_partkey, + l_quantity, + l_extendedprice, + l_discount, + l_tax, + l_returnflag, + l_linestatus, + l_shipdate, + l_commitdate, + l_receiptdate, + l_shipinstruct, + l_shipmode +FROM lineitem +ORDER BY l_orderkey, l_suppkey, l_linenumber, l_comment +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q10.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q11.benchmark b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q11.benchmark new file mode 100644 index 0000000000000..efce2005f3beb --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/benchmarks/q11.benchmark @@ -0,0 +1,44 @@ +echo Loading tpch items sorted: ${BENCH_SORTED:-false} + +# +# Sort queries with different characteristics: +# - Sort key with fixed length or variable length (VARCHAR) +# - Sort key with different cardinality +# - Different number of sort keys +# - Different number of payload columns (thin: 1 additional column other +# than sort keys; wide: all columns except sort keys) +# +# DataSet is `lineitem` table in TPCH dataset (16 columns, 6M rows for +# scale factor 1.0, cardinality is counted from SF1 dataset) +# +# Key Columns: +# - Column `l_linenumber`, type: `INTEGER`, cardinality: 7 +# - Column `l_suppkey`, type: `BIGINT`, cardinality: 10k +# - Column `l_orderkey`, type: `BIGINT`, cardinality: 1.5M +# - Column `l_comment`, type: `VARCHAR`, cardinality: 4.5M (len is ~26 chars) +# +# Payload Columns: +# - Thin variant: `l_partkey` column with `BIGINT` type (1 column) +# - Wide variant: all columns except for possible key columns (12 columns) + +name Q11 +group sort_tpch +subgroup sf${BENCH_SIZE:-1} + +echo Loading sort_tpch sf ${BENCH_SIZE:-1} data + +load sql_benchmarks/sort_tpch/init/load.sql + +assert I +SELECT COUNT(*) > 0 from lineitem; +---- +true + +run +-- Q11: 1 sort key (type: VARCHAR, cardinality: 4.5M) + 1 payload column +SELECT l_shipmode, l_comment, l_partkey +FROM lineitem +ORDER BY l_shipmode +${LIMIT:-false|LIMIT 100| } + +result sql_benchmarks/sort_tpch/results/sf${BENCH_SIZE:-1}/q11.csv diff --git a/benchmarks/sql_benchmarks/sort_tpch/init/load.sql b/benchmarks/sql_benchmarks/sort_tpch/init/load.sql new file mode 100644 index 0000000000000..395d8da009d21 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/init/load.sql @@ -0,0 +1,3 @@ +CREATE EXTERNAL TABLE lineitem_raw STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpch_sf${BENCH_SIZE:-1}/lineitem/lineitem.1.parquet'; + +CREATE TABLE lineitem as (SELECT * FROM lineitem_raw${BENCH_SORTED:-false| order by l_orderkey asc| }); \ No newline at end of file diff --git a/benchmarks/sql_benchmarks/sort_tpch/sort_tpch.suite b/benchmarks/sql_benchmarks/sort_tpch/sort_tpch.suite new file mode 100644 index 0000000000000..38ee9c132b284 --- /dev/null +++ b/benchmarks/sql_benchmarks/sort_tpch/sort_tpch.suite @@ -0,0 +1,28 @@ +description = "Sorting benchmarks over the TPC-H lineitem table" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "..."] +help = "Selects the TPC-H scale factor." + +[[options]] +name = "sorted" +env = "BENCH_SORTED" +default = "false" +values = ["false", "true"] +help = "Controls whether the lineitem table is loaded in l_orderkey order." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- sort_tpch" +description = "Run all TPC-H sorting queries at scale factor 1." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- sort_tpch --query 4 --sorted true" +description = "Run sorting query 4 over pre-sorted lineitem data." diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q01.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q01.benchmark new file mode 100644 index 0000000000000..f3a7cb1b45d8e --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q01.benchmark @@ -0,0 +1,16 @@ +name Q01 +group tpcds + + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/1.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/1.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q02.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q02.benchmark new file mode 100644 index 0000000000000..4b68a1829dcc7 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q02.benchmark @@ -0,0 +1,15 @@ +name Q02 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/2.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/2.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q03.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q03.benchmark new file mode 100644 index 0000000000000..74c8e78584821 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q03.benchmark @@ -0,0 +1,15 @@ +name Q03 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/3.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/3.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q04.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q04.benchmark new file mode 100644 index 0000000000000..dbd2a8879763b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q04.benchmark @@ -0,0 +1,16 @@ +name Q04 +group tpcds + + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/4.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/4.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q05.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q05.benchmark new file mode 100644 index 0000000000000..49bbfbda237c0 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q05.benchmark @@ -0,0 +1,15 @@ +name Q05 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/5.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/5.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q06.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q06.benchmark new file mode 100644 index 0000000000000..b2349eda3639e --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q06.benchmark @@ -0,0 +1,15 @@ +name Q06 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/6.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/6.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q07.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q07.benchmark new file mode 100644 index 0000000000000..1e354cf1b0ab5 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q07.benchmark @@ -0,0 +1,15 @@ +name Q07 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/7.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/7.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q08.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q08.benchmark new file mode 100644 index 0000000000000..b9a511dbe0d31 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q08.benchmark @@ -0,0 +1,15 @@ +name Q08 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/8.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/8.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q09.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q09.benchmark new file mode 100644 index 0000000000000..15cf4226acf9e --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q09.benchmark @@ -0,0 +1,15 @@ +name Q09 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/9.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/9.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q10.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q10.benchmark new file mode 100644 index 0000000000000..401bd3dea294b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q10.benchmark @@ -0,0 +1,15 @@ +name Q10 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/10.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/10.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q11.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q11.benchmark new file mode 100644 index 0000000000000..f54ba637bed31 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q11.benchmark @@ -0,0 +1,15 @@ +name Q11 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/11.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/11.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q12.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q12.benchmark new file mode 100644 index 0000000000000..16d3530dd676b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q12.benchmark @@ -0,0 +1,15 @@ +name Q12 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/12.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/12.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q13.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q13.benchmark new file mode 100644 index 0000000000000..7ef0d003d09d8 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q13.benchmark @@ -0,0 +1,15 @@ +name Q13 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/13.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/13.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q14.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q14.benchmark new file mode 100644 index 0000000000000..748e11083588b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q14.benchmark @@ -0,0 +1,15 @@ +name Q14 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/14.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/14.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q15.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q15.benchmark new file mode 100644 index 0000000000000..3a5c1d6c34207 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q15.benchmark @@ -0,0 +1,15 @@ +name Q15 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/15.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/15.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q16.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q16.benchmark new file mode 100644 index 0000000000000..e9cf989f4fdb9 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q16.benchmark @@ -0,0 +1,15 @@ +name Q16 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/16.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/16.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q17.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q17.benchmark new file mode 100644 index 0000000000000..5a9eb9ab11aef --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q17.benchmark @@ -0,0 +1,15 @@ +name Q17 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/17.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/17.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q18.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q18.benchmark new file mode 100644 index 0000000000000..eca4093edccd9 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q18.benchmark @@ -0,0 +1,15 @@ +name Q18 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/18.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/18.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q19.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q19.benchmark new file mode 100644 index 0000000000000..385524636da71 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q19.benchmark @@ -0,0 +1,15 @@ +name Q19 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/19.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/19.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q20.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q20.benchmark new file mode 100644 index 0000000000000..f05d81638250b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q20.benchmark @@ -0,0 +1,15 @@ +name Q20 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/20.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/20.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q21.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q21.benchmark new file mode 100644 index 0000000000000..98ed74677a082 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q21.benchmark @@ -0,0 +1,15 @@ +name Q21 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/21.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/21.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q22.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q22.benchmark new file mode 100644 index 0000000000000..e1eccfc852987 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q22.benchmark @@ -0,0 +1,15 @@ +name Q22 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/22.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/22.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q23.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q23.benchmark new file mode 100644 index 0000000000000..47153714ac5ad --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q23.benchmark @@ -0,0 +1,15 @@ +name Q23 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/23.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/23.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q24.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q24.benchmark new file mode 100644 index 0000000000000..05540a4606336 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q24.benchmark @@ -0,0 +1,15 @@ +name Q24 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/24.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/24.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q25.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q25.benchmark new file mode 100644 index 0000000000000..0c0a0f8ce8b55 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q25.benchmark @@ -0,0 +1,15 @@ +name Q25 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/25.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/25.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q26.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q26.benchmark new file mode 100644 index 0000000000000..8481c3b660ec9 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q26.benchmark @@ -0,0 +1,15 @@ +name Q26 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/26.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/26.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q27.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q27.benchmark new file mode 100644 index 0000000000000..2357a8aae87b1 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q27.benchmark @@ -0,0 +1,15 @@ +name Q27 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/27.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/27.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q28.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q28.benchmark new file mode 100644 index 0000000000000..f4cbce1430eee --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q28.benchmark @@ -0,0 +1,15 @@ +name Q28 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/28.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/28.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q29.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q29.benchmark new file mode 100644 index 0000000000000..77b9b058128df --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q29.benchmark @@ -0,0 +1,15 @@ +name Q29 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/29.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/29.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q30.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q30.benchmark new file mode 100644 index 0000000000000..e7c144674d96d --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q30.benchmark @@ -0,0 +1,15 @@ +name Q30 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/30.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/30.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q31.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q31.benchmark new file mode 100644 index 0000000000000..84f2b0ba5ed69 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q31.benchmark @@ -0,0 +1,15 @@ +name Q31 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/31.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/31.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q32.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q32.benchmark new file mode 100644 index 0000000000000..42b995ac6608c --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q32.benchmark @@ -0,0 +1,15 @@ +name Q32 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/32.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/32.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q33.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q33.benchmark new file mode 100644 index 0000000000000..ad3c21990c088 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q33.benchmark @@ -0,0 +1,15 @@ +name Q33 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/33.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/33.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q34.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q34.benchmark new file mode 100644 index 0000000000000..af74cd4abdd72 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q34.benchmark @@ -0,0 +1,15 @@ +name Q34 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/34.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/34.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q35.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q35.benchmark new file mode 100644 index 0000000000000..62ae55e6b8444 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q35.benchmark @@ -0,0 +1,15 @@ +name Q35 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/35.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/35.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q36.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q36.benchmark new file mode 100644 index 0000000000000..c1d1ee9ebc1c4 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q36.benchmark @@ -0,0 +1,15 @@ +name Q36 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/36.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/36.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q37.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q37.benchmark new file mode 100644 index 0000000000000..47dfb9229353a --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q37.benchmark @@ -0,0 +1,15 @@ +name Q37 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/37.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/37.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q38.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q38.benchmark new file mode 100644 index 0000000000000..14616a6abf631 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q38.benchmark @@ -0,0 +1,15 @@ +name Q38 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/38.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/38.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q39.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q39.benchmark new file mode 100644 index 0000000000000..12c02d8135568 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q39.benchmark @@ -0,0 +1,15 @@ +name Q39 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/39.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/39.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q40.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q40.benchmark new file mode 100644 index 0000000000000..c5e787bbb0a1f --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q40.benchmark @@ -0,0 +1,15 @@ +name Q40 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/40.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/40.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q41.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q41.benchmark new file mode 100644 index 0000000000000..bc1daf4f55cc2 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q41.benchmark @@ -0,0 +1,15 @@ +name Q41 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/41.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/41.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q42.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q42.benchmark new file mode 100644 index 0000000000000..1054b6223cad3 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q42.benchmark @@ -0,0 +1,15 @@ +name Q42 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/42.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/42.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q43.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q43.benchmark new file mode 100644 index 0000000000000..902ebcbe0d357 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q43.benchmark @@ -0,0 +1,15 @@ +name Q43 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/43.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/43.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q44.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q44.benchmark new file mode 100644 index 0000000000000..620b7c4f9e05f --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q44.benchmark @@ -0,0 +1,15 @@ +name Q44 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/44.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/44.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q45.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q45.benchmark new file mode 100644 index 0000000000000..7d5cc3d05dbb7 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q45.benchmark @@ -0,0 +1,15 @@ +name Q45 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/45.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/45.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q46.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q46.benchmark new file mode 100644 index 0000000000000..398921d09178a --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q46.benchmark @@ -0,0 +1,15 @@ +name Q46 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/46.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/46.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q47.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q47.benchmark new file mode 100644 index 0000000000000..7a14ecf23cdab --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q47.benchmark @@ -0,0 +1,15 @@ +name Q47 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/47.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/47.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q48.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q48.benchmark new file mode 100644 index 0000000000000..c60972b34d3cc --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q48.benchmark @@ -0,0 +1,15 @@ +name Q48 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/48.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/48.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q49.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q49.benchmark new file mode 100644 index 0000000000000..ebfdce644c7e8 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q49.benchmark @@ -0,0 +1,15 @@ +name Q49 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/49.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/49.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q50.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q50.benchmark new file mode 100644 index 0000000000000..bf8056b7d178c --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q50.benchmark @@ -0,0 +1,15 @@ +name Q50 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/50.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/50.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q51.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q51.benchmark new file mode 100644 index 0000000000000..90982ca601b36 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q51.benchmark @@ -0,0 +1,15 @@ +name Q51 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/51.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/51.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q52.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q52.benchmark new file mode 100644 index 0000000000000..ed9d4ea86be60 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q52.benchmark @@ -0,0 +1,15 @@ +name Q52 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/52.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/52.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q53.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q53.benchmark new file mode 100644 index 0000000000000..b77eac22c97f3 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q53.benchmark @@ -0,0 +1,15 @@ +name Q53 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/53.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/53.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q54.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q54.benchmark new file mode 100644 index 0000000000000..83bb72e103cd3 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q54.benchmark @@ -0,0 +1,15 @@ +name Q54 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/54.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/54.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q55.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q55.benchmark new file mode 100644 index 0000000000000..41ce8b54e4d87 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q55.benchmark @@ -0,0 +1,15 @@ +name Q55 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/55.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/55.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q56.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q56.benchmark new file mode 100644 index 0000000000000..5fead000b5761 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q56.benchmark @@ -0,0 +1,15 @@ +name Q56 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/56.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/56.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q57.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q57.benchmark new file mode 100644 index 0000000000000..78368b6057266 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q57.benchmark @@ -0,0 +1,15 @@ +name Q57 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/57.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/57.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q58.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q58.benchmark new file mode 100644 index 0000000000000..6d3e80b4bfbba --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q58.benchmark @@ -0,0 +1,15 @@ +name Q58 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/58.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/58.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q59.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q59.benchmark new file mode 100644 index 0000000000000..33bcd35d3fc45 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q59.benchmark @@ -0,0 +1,15 @@ +name Q59 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/59.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/59.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q60.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q60.benchmark new file mode 100644 index 0000000000000..766e1eb101f50 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q60.benchmark @@ -0,0 +1,15 @@ +name Q60 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/60.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/60.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q61.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q61.benchmark new file mode 100644 index 0000000000000..0c41e1ef0ca27 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q61.benchmark @@ -0,0 +1,15 @@ +name Q61 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/61.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/61.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q62.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q62.benchmark new file mode 100644 index 0000000000000..e097f807d27aa --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q62.benchmark @@ -0,0 +1,15 @@ +name Q62 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/62.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/62.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q63.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q63.benchmark new file mode 100644 index 0000000000000..b2cee6313e3b7 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q63.benchmark @@ -0,0 +1,15 @@ +name Q63 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/63.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/63.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q64.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q64.benchmark new file mode 100644 index 0000000000000..5116830e58f16 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q64.benchmark @@ -0,0 +1,15 @@ +name Q64 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/64.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/64.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q65.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q65.benchmark new file mode 100644 index 0000000000000..eb33f20a55835 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q65.benchmark @@ -0,0 +1,15 @@ +name Q65 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/65.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/65.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q66.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q66.benchmark new file mode 100644 index 0000000000000..f9bedd5474c40 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q66.benchmark @@ -0,0 +1,15 @@ +name Q66 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/66.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/66.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q67.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q67.benchmark new file mode 100644 index 0000000000000..1d387a7fb66ed --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q67.benchmark @@ -0,0 +1,15 @@ +name Q67 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/67.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/67.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q68.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q68.benchmark new file mode 100644 index 0000000000000..e5303d6543ba8 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q68.benchmark @@ -0,0 +1,15 @@ +name Q68 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/68.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/68.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q69.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q69.benchmark new file mode 100644 index 0000000000000..bc0e043af2dc1 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q69.benchmark @@ -0,0 +1,15 @@ +name Q69 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/69.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/69.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q70.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q70.benchmark new file mode 100644 index 0000000000000..345a17db0b31b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q70.benchmark @@ -0,0 +1,15 @@ +name Q70 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/70.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/70.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q71.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q71.benchmark new file mode 100644 index 0000000000000..4dd7d2f90e0c5 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q71.benchmark @@ -0,0 +1,15 @@ +name Q71 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/71.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/71.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q72.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q72.benchmark new file mode 100644 index 0000000000000..15faf7eb4c496 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q72.benchmark @@ -0,0 +1,15 @@ +name Q72 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/72.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/72.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q73.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q73.benchmark new file mode 100644 index 0000000000000..5579742338649 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q73.benchmark @@ -0,0 +1,15 @@ +name Q73 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/73.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/73.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q74.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q74.benchmark new file mode 100644 index 0000000000000..a40113a16f089 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q74.benchmark @@ -0,0 +1,15 @@ +name Q74 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/74.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/74.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q75.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q75.benchmark new file mode 100644 index 0000000000000..6e11461c2f6bb --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q75.benchmark @@ -0,0 +1,15 @@ +name Q75 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/75.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/75.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q76.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q76.benchmark new file mode 100644 index 0000000000000..281d65129051b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q76.benchmark @@ -0,0 +1,15 @@ +name Q76 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/76.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/76.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q77.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q77.benchmark new file mode 100644 index 0000000000000..42d3518239c2a --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q77.benchmark @@ -0,0 +1,15 @@ +name Q77 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/77.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/77.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q78.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q78.benchmark new file mode 100644 index 0000000000000..03ea2b583b9b1 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q78.benchmark @@ -0,0 +1,15 @@ +name Q78 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/78.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/78.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q79.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q79.benchmark new file mode 100644 index 0000000000000..151ee27b39cd0 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q79.benchmark @@ -0,0 +1,15 @@ +name Q79 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/79.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/79.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q80.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q80.benchmark new file mode 100644 index 0000000000000..9f6809fa36066 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q80.benchmark @@ -0,0 +1,15 @@ +name Q80 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/80.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/80.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q81.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q81.benchmark new file mode 100644 index 0000000000000..bd5bfec578253 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q81.benchmark @@ -0,0 +1,15 @@ +name Q81 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/81.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/81.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q82.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q82.benchmark new file mode 100644 index 0000000000000..7fce855ac696d --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q82.benchmark @@ -0,0 +1,15 @@ +name Q82 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/82.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/82.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q83.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q83.benchmark new file mode 100644 index 0000000000000..c39cf514a9749 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q83.benchmark @@ -0,0 +1,15 @@ +name Q83 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/83.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/83.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q84.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q84.benchmark new file mode 100644 index 0000000000000..8debc4f705cc8 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q84.benchmark @@ -0,0 +1,15 @@ +name Q84 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/84.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/84.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q85.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q85.benchmark new file mode 100644 index 0000000000000..050e1efcdc75b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q85.benchmark @@ -0,0 +1,15 @@ +name Q85 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/85.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/85.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q86.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q86.benchmark new file mode 100644 index 0000000000000..53b65089e3a6e --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q86.benchmark @@ -0,0 +1,15 @@ +name Q86 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/86.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/86.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q87.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q87.benchmark new file mode 100644 index 0000000000000..71021946c82c6 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q87.benchmark @@ -0,0 +1,15 @@ +name Q87 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/87.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/87.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q88.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q88.benchmark new file mode 100644 index 0000000000000..49e07041e4b5a --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q88.benchmark @@ -0,0 +1,15 @@ +name Q88 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/88.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/88.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q89.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q89.benchmark new file mode 100644 index 0000000000000..2256006201502 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q89.benchmark @@ -0,0 +1,15 @@ +name Q89 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/89.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/89.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q90.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q90.benchmark new file mode 100644 index 0000000000000..d95d5ff3e2f6c --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q90.benchmark @@ -0,0 +1,15 @@ +name Q90 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/90.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/90.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q91.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q91.benchmark new file mode 100644 index 0000000000000..e82bfc5d4c9b9 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q91.benchmark @@ -0,0 +1,15 @@ +name Q91 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/91.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/91.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q92.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q92.benchmark new file mode 100644 index 0000000000000..bc7b9236bf4ea --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q92.benchmark @@ -0,0 +1,15 @@ +name Q92 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/92.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/92.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q93.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q93.benchmark new file mode 100644 index 0000000000000..0b9645f9cbf2b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q93.benchmark @@ -0,0 +1,15 @@ +name Q93 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/93.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/93.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q94.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q94.benchmark new file mode 100644 index 0000000000000..f5932537fd31b --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q94.benchmark @@ -0,0 +1,15 @@ +name Q94 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/94.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/94.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q95.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q95.benchmark new file mode 100644 index 0000000000000..3eda91c9ba1e0 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q95.benchmark @@ -0,0 +1,15 @@ +name Q95 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/95.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/95.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q96.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q96.benchmark new file mode 100644 index 0000000000000..caef1b71556ce --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q96.benchmark @@ -0,0 +1,15 @@ +name Q96 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/96.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/96.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q97.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q97.benchmark new file mode 100644 index 0000000000000..c81446698bfe9 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q97.benchmark @@ -0,0 +1,15 @@ +name Q97 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/97.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/97.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q98.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q98.benchmark new file mode 100644 index 0000000000000..b598baa846d77 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q98.benchmark @@ -0,0 +1,15 @@ +name Q98 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/98.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/98.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/benchmarks/q99.benchmark b/benchmarks/sql_benchmarks/tpcds/benchmarks/q99.benchmark new file mode 100644 index 0000000000000..d017d447bcaa0 --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/benchmarks/q99.benchmark @@ -0,0 +1,15 @@ +name Q99 +group tpcds + +load sql_benchmarks/tpcds/init/load.sql + +assert I +SELECT COUNT(*) > 0 from web_site; +---- +true + +run ../datafusion/core/tests/tpc-ds/99.sql + +result sql_benchmarks/tpcds/results/sf${BENCH_SIZE:-1}/99.csv + +cleanup sql_benchmarks/tpcds/init/cleanup.sql diff --git a/benchmarks/sql_benchmarks/tpcds/init/cleanup.sql b/benchmarks/sql_benchmarks/tpcds/init/cleanup.sql new file mode 100644 index 0000000000000..2a6ed79c5196d --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/init/cleanup.sql @@ -0,0 +1,47 @@ +DROP TABLE IF EXISTS call_center; + +DROP TABLE IF EXISTS catalog_page; + +DROP TABLE IF EXISTS catalog_returns; + +DROP TABLE IF EXISTS catalog_sales; + +DROP TABLE IF EXISTS customer; + +DROP TABLE IF EXISTS customer_address; + +DROP TABLE IF EXISTS customer_demographics; + +DROP TABLE IF EXISTS date_dim; + +DROP TABLE IF EXISTS household_demographics; + +DROP TABLE IF EXISTS income_band; + +DROP TABLE IF EXISTS inventory; + +DROP TABLE IF EXISTS item; + +DROP TABLE IF EXISTS promotion; + +DROP TABLE IF EXISTS reason; + +DROP TABLE IF EXISTS ship_mode; + +DROP TABLE IF EXISTS store; + +DROP TABLE IF EXISTS store_returns; + +DROP TABLE IF EXISTS store_sales; + +DROP TABLE IF EXISTS time_dim; + +DROP TABLE IF EXISTS warehouse; + +DROP TABLE IF EXISTS web_page; + +DROP TABLE IF EXISTS web_returns; + +DROP TABLE IF EXISTS web_sales; + +DROP TABLE IF EXISTS web_site; diff --git a/benchmarks/sql_benchmarks/tpcds/init/load.sql b/benchmarks/sql_benchmarks/tpcds/init/load.sql new file mode 100644 index 0000000000000..6b89199646f7f --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/init/load.sql @@ -0,0 +1,47 @@ +CREATE EXTERNAL TABLE call_center STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/call_center.parquet'; + +CREATE EXTERNAL TABLE catalog_page STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/catalog_page.parquet'; + +CREATE EXTERNAL TABLE catalog_returns STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/catalog_returns.parquet'; + +CREATE EXTERNAL TABLE catalog_sales STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/catalog_sales.parquet'; + +CREATE EXTERNAL TABLE customer STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/customer.parquet'; + +CREATE EXTERNAL TABLE customer_address STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/customer_address.parquet'; + +CREATE EXTERNAL TABLE customer_demographics STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/customer_demographics.parquet'; + +CREATE EXTERNAL TABLE date_dim STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/date_dim.parquet'; + +CREATE EXTERNAL TABLE household_demographics STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/household_demographics.parquet'; + +CREATE EXTERNAL TABLE income_band STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/income_band.parquet'; + +CREATE EXTERNAL TABLE inventory STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/inventory.parquet'; + +CREATE EXTERNAL TABLE item STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/item.parquet'; + +CREATE EXTERNAL TABLE promotion STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/promotion.parquet'; + +CREATE EXTERNAL TABLE reason STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/reason.parquet'; + +CREATE EXTERNAL TABLE ship_mode STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/ship_mode.parquet'; + +CREATE EXTERNAL TABLE store STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/store.parquet'; + +CREATE EXTERNAL TABLE store_returns STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/store_returns.parquet'; + +CREATE EXTERNAL TABLE store_sales STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/store_sales.parquet'; + +CREATE EXTERNAL TABLE time_dim STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/time_dim.parquet'; + +CREATE EXTERNAL TABLE warehouse STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/warehouse.parquet'; + +CREATE EXTERNAL TABLE web_page STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/web_page.parquet'; + +CREATE EXTERNAL TABLE web_returns STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/web_returns.parquet'; + +CREATE EXTERNAL TABLE web_sales STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/web_sales.parquet'; + +CREATE EXTERNAL TABLE web_site STORED AS PARQUET LOCATION '${DATA_DIR:-data}/tpcds_sf${BENCH_SIZE:-1}/web_site.parquet'; diff --git a/benchmarks/sql_benchmarks/tpcds/tpcds.suite b/benchmarks/sql_benchmarks/tpcds/tpcds.suite new file mode 100644 index 0000000000000..7261c3d4dfc6d --- /dev/null +++ b/benchmarks/sql_benchmarks/tpcds/tpcds.suite @@ -0,0 +1,21 @@ +description = "TPC-DS SQL benchmarks" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[options]] +name = "scale-factor" +env = "BENCH_SIZE" +default = "1" +values = ["1", "10", "100", "..."] +help = "Selects the TPC-DS scale factor." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpcds" +description = "Run all TPC-DS queries at scale factor 1." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpcds --query 42 --scale-factor 10" +description = "Run TPC-DS query 42 at scale factor 10." diff --git a/benchmarks/sql_benchmarks/tpch/tpch.suite b/benchmarks/sql_benchmarks/tpch/tpch.suite index 317b32e57f45f..0330cc0f32584 100644 --- a/benchmarks/sql_benchmarks/tpch/tpch.suite +++ b/benchmarks/sql_benchmarks/tpch/tpch.suite @@ -1,16 +1,47 @@ -name = "tpch" description = "TPC-H SQL benchmarks" +# Query patterns control how numeric QUERY_ID values map to .benchmark files +# during discovery and command resolution. Use exactly one query-id token: +# - {QUERY_ID_PADDED}: two-digit ids, such as q01.benchmark +# - {QUERY_ID}: unpadded ids, such as query-1.benchmark +# If omitted, this defaults to q{QUERY_ID_PADDED}.benchmark. +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +# Path replacements define path-like variables used while parsing benchmark +# files. Relative paths are resolved from this suite file's directory and then +# passed to SqlBenchmark's replacement mapping, so benchmark SQL can refer to +# values such as ${DATA_DIR}. For timed runs, the runner's --path/-p option +# overrides DATA_DIR. +[path_replacements] +DATA_DIR = "../../data" + [[options]] name = "format" short = "f" +env = "TPCH_FILE_TYPE" default = "parquet" values = ["parquet", "csv", "mem"] help = "Selects the TPC-H data format." [[options]] name = "scale-factor" -short = "sf" +env = "BENCH_SIZE" default = "1" -values = ["1", "10"] +values = ["1", "10", "..."] help = "Selects the TPC-H scale factor." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpch" +description = "Run all TPC-H queries with the default parquet SF1 configuration." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpch --query 15" +description = "Run TPC-H query 15 with the default parquet SF1 configuration." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpch --query 15 -f csv" +description = "Run TPC-H query 15 against CSV data." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- tpch --query 15 --scale-factor 10" +description = "Run TPC-H query 15 at scale factor 10." diff --git a/benchmarks/sql_benchmarks/wide_schema/init/load.sql b/benchmarks/sql_benchmarks/wide_schema/init/load.sql index 4fbcda1d5817e..72486106770aa 100644 --- a/benchmarks/sql_benchmarks/wide_schema/init/load.sql +++ b/benchmarks/sql_benchmarks/wide_schema/init/load.sql @@ -3,4 +3,4 @@ -- BENCH_SUBGROUP=wide → 1024-col synthetic dataset (the actual benchmark) -- BENCH_SUBGROUP=narrow → 8-col baseline (companion only — meaningful -- only when compared to the wide numbers) -CREATE EXTERNAL TABLE events STORED AS PARQUET LOCATION 'data/wide_schema/${BENCH_SUBGROUP:-wide}/'; +CREATE EXTERNAL TABLE events STORED AS PARQUET LOCATION '${DATA_DIR:-data}/wide_schema/${BENCH_SUBGROUP:-wide}/'; diff --git a/benchmarks/sql_benchmarks/wide_schema/wide_schema.suite b/benchmarks/sql_benchmarks/wide_schema/wide_schema.suite new file mode 100644 index 0000000000000..275f15e102677 --- /dev/null +++ b/benchmarks/sql_benchmarks/wide_schema/wide_schema.suite @@ -0,0 +1,14 @@ +description = "Projection benchmarks over synthetic wide and narrow schemas" + +query_pattern = "q{QUERY_ID_PADDED}.benchmark" + +[path_replacements] +DATA_DIR = "../../data" + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- wide_schema" +description = "Run all wide-schema projection queries." + +[[examples]] +command = "cargo run --release --bin benchmark_runner -- wide_schema --query 2 --subgroup narrow" +description = "Run query 2 with the narrow schema." diff --git a/benchmarks/src/benchmark_runner/cli.rs b/benchmarks/src/benchmark_runner/cli.rs deleted file mode 100644 index 30fb71b5a30cc..0000000000000 --- a/benchmarks/src/benchmark_runner/cli.rs +++ /dev/null @@ -1,81 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! CLI construction and argument conversion for `benchmark_runner`. -//! -//! This module owns the clap command tree for the initial runner surface: -//! top-level help and suite listing. - -use clap::builder::styling::{AnsiColor, Styles}; -use clap::{ArgMatches, Command}; -use datafusion_common::{Result, exec_datafusion_err}; - -const HELP_STYLES: Styles = Styles::styled() - .header(AnsiColor::Green.on_default().bold()) - .usage(AnsiColor::Green.on_default().bold()) - .literal(AnsiColor::Cyan.on_default().bold()) - .placeholder(AnsiColor::Cyan.on_default()); - -#[derive(Debug)] -pub enum RunnerCommand { - Help, - List, -} - -/// Builds the command tree for help and suite listing. -pub fn build_cli() -> Command { - Command::new("benchmark_runner") - .about("Inspect DataFusion SQL benchmark suites.") - .styles(HELP_STYLES) - .subcommand_required(false) - .arg_required_else_help(false) - .disable_help_subcommand(true) - .subcommand(Command::new("help").about("Print help")) - .subcommand(Command::new("list").about("List SQL benchmark suites")) -} - -/// Converts clap matches into a typed command. -pub(crate) fn command_from_matches(matches: &ArgMatches) -> Result { - match matches.subcommand() { - None | Some(("help", _)) => Ok(RunnerCommand::Help), - Some(("list", _)) => Ok(RunnerCommand::List), - Some((name, _)) => Err(exec_datafusion_err!("Unknown command '{name}'")), - } -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn list_rejects_unrecognized_options() { - let matches = - build_cli().try_get_matches_from(["benchmark_runner", "list", "--format"]); - - assert!(matches.is_err(), "{matches:?}"); - } - - #[test] - fn help_mentions_list_command() { - let err = build_cli() - .try_get_matches_from(["benchmark_runner", "--help"]) - .unwrap_err(); - let help = err.to_string(); - - assert!(help.contains("list")); - } -} diff --git a/benchmarks/src/benchmark_runner/mod.rs b/benchmarks/src/benchmark_runner/mod.rs deleted file mode 100644 index 458e5f974c152..0000000000000 --- a/benchmarks/src/benchmark_runner/mod.rs +++ /dev/null @@ -1,104 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Command-line inspection for SQL benchmark suites. -//! -//! This module backs the `benchmark_runner` binary. The initial command -//! surface lists discovered SQL benchmark suites from `.suite` files and -//! prints the top-level help. -//! -//! Common invocations: -//! -//! ```text -//! cargo run --bin benchmark_runner -- --help -//! cargo run --release --bin benchmark_runner -- list -//! ``` -//! -//! The public entry point is [`run_cli`]. The submodules are kept private so -//! the command-line flow remains the single supported API: -//! -//! - `cli` builds the clap command tree and parses the selected command. -//! - `suite` loads `.suite` metadata and discovers benchmark query files. -//! - `output` formats colored `list` command output. - -mod cli; -mod output; -mod suite; - -use crate::benchmark_runner::cli::{RunnerCommand, build_cli, command_from_matches}; -use crate::benchmark_runner::output::format_suite_list_styled; -use crate::benchmark_runner::suite::SuiteRegistry; -use datafusion::error::Result; -use datafusion_common::DataFusionError; -use std::io::Write; -use std::path::PathBuf; - -/// Runs the benchmark runner command-line flow for the provided argument list. -/// -/// This discovers suite metadata, parses the help/list command, and dispatches -/// to the selected implementation. -pub fn run_cli(args: I) -> Result<()> -where - I: IntoIterator, - T: Clone + Into, -{ - let benchmark_dir = default_benchmark_dir(); - let registry = SuiteRegistry::discover(&benchmark_dir)?; - let mut cli = build_cli(); - let matches = match cli.try_get_matches_from_mut(args) { - Ok(matches) => matches, - Err(e) if e.kind() == clap::error::ErrorKind::DisplayHelp => { - e.print()?; - return Ok(()); - } - Err(e) => return Err(DataFusionError::External(Box::new(e))), - }; - let command = command_from_matches(&matches)?; - - match command { - RunnerCommand::Help => { - cli.print_long_help()?; - println!(); - } - RunnerCommand::List => { - print_styled(&format_suite_list_styled(®istry)?)?; - } - } - - Ok(()) -} - -/// Writes already styled output through `anstream` so ANSI color handling -/// matches clap help output on supported terminals. -fn print_styled(output: &str) -> Result<()> { - let mut stdout = anstream::stdout(); - - write!(&mut stdout, "{output}") - .map_err(|e| DataFusionError::External(Box::new(e)))?; - Ok(()) -} - -/// Resolves the SQL benchmark root from either the repository root or the -/// benchmarks crate manifest directory. -fn default_benchmark_dir() -> PathBuf { - let repo_root_path = PathBuf::from("benchmarks/sql_benchmarks"); - if repo_root_path.exists() { - repo_root_path - } else { - PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("sql_benchmarks") - } -} diff --git a/benchmarks/src/benchmark_runner/output.rs b/benchmarks/src/benchmark_runner/output.rs deleted file mode 100644 index 9eb975311e29f..0000000000000 --- a/benchmarks/src/benchmark_runner/output.rs +++ /dev/null @@ -1,173 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Formatting helpers for human-readable benchmark runner output. -//! -//! The runner intentionally uses the same colored style for `list` output as -//! clap uses for help text. - -use crate::benchmark_runner::suite::{SuiteConfig, SuiteOption, SuiteRegistry}; -use clap::builder::styling::{AnsiColor, Style}; -use datafusion_common::Result; -use std::fmt::{Display, Write as _}; - -/// Formats the `list` command output with suite summaries, query-id hints, and -/// configurable suite options. -pub fn format_suite_list_styled(registry: &SuiteRegistry) -> Result { - let mut output = String::new(); - - for suite in registry.suites() { - write_suite_list_entry(&mut output, suite)?; - } - - Ok(output) -} - -/// Writes one suite entry for the `list` command. -fn write_suite_list_entry(output: &mut String, suite: &SuiteConfig) -> Result<()> { - writeln!(output, "{}", header(&suite.name))?; - writeln!(output, " {}: {}", label("description"), suite.description)?; - - let queries = suite.discover_queries()?; - - if let (Some(first), Some(last)) = (queries.first(), queries.last()) { - writeln!( - output, - " {}: {}-{} discovered under {} as {}", - label("query ids"), - value(first.id), - value(last.id), - suite.query_search_root().display(), - literal("qNN.benchmark") - )?; - } - writeln!(output, " {}:", label("options"))?; - - for option in &suite.options { - write_suite_list_option(output, option)?; - } - - Ok(()) -} - -/// Writes one suite option summary for the `list` command. -fn write_suite_list_option(output: &mut String, option: &SuiteOption) -> Result<()> { - let values = option - .values - .iter() - .map(|v| { - if v == &option.default { - format!("{} ({})", value(v), label("default")) - } else { - value(v) - } - }) - .collect::>() - .join(", "); - - writeln!( - output, - " {} {} {}", - literal(option_display(option)), - placeholder(""), - values - )?; - - Ok(()) -} - -fn option_display(option: &SuiteOption) -> String { - match &option.short { - Some(short) => format!("-{short}, --{}", option.name), - None => format!("--{}", option.name), - } -} - -fn header(text: impl Display) -> String { - styled(AnsiColor::Green.on_default().bold(), text) -} - -fn literal(text: impl Display) -> String { - styled(AnsiColor::Cyan.on_default().bold(), text) -} - -fn placeholder(text: impl Display) -> String { - styled(AnsiColor::Cyan.on_default(), text) -} - -fn value(text: impl Display) -> String { - styled(AnsiColor::Green.on_default(), text) -} - -fn label(text: impl Display) -> String { - styled(Style::new().bold(), text) -} - -fn styled(style: Style, text: impl Display) -> String { - format!("{style}{text}{style:#}") -} - -#[cfg(test)] -mod tests { - use super::*; - use std::path::PathBuf; - - fn manifest_path(path: &str) -> PathBuf { - PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path) - } - - fn strip_ansi(input: &str) -> String { - let mut output = String::new(); - let mut chars = input.chars(); - while let Some(c) = chars.next() { - if c == '\x1b' { - for c in chars.by_ref() { - if c == 'm' { - break; - } - } - } else { - output.push(c); - } - } - output - } - - #[test] - fn list_output_mentions_tpch_options() { - let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); - let output = strip_ansi(&format_suite_list_styled(®istry).unwrap()); - - assert!(output.contains("tpch\n description: TPC-H SQL benchmarks")); - assert!(output.contains("-f, --format parquet (default), csv, mem")); - assert!(output.contains("-sf, --scale-factor 1 (default), 10")); - } - - #[test] - fn styled_list_output_includes_ansi_sequences() { - let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); - let output = format_suite_list_styled(®istry).unwrap(); - - assert!(output.contains("\u{1b}[")); - assert!(output.contains("tpch")); - assert!(output.contains("-f")); - assert!(output.contains("--format")); - assert!(output.contains("-sf")); - assert!(output.contains("--scale-factor")); - assert!(output.contains("")); - } -} diff --git a/benchmarks/src/benchmark_runner/suite.rs b/benchmarks/src/benchmark_runner/suite.rs deleted file mode 100644 index fe5b3339b1643..0000000000000 --- a/benchmarks/src/benchmark_runner/suite.rs +++ /dev/null @@ -1,500 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Suite-file loading and validation. -//! -//! A suite is described by a text `.suite` file that declares which options -//! the runner should display. Query discovery recursively scans from the suite -//! file's directory for `qNN.benchmark` files. Discovered queries are cached -//! lazily because they are reused by listing during a single CLI run. - -use datafusion_common::{DataFusionError, Result}; -use serde::{Deserialize, Serialize}; -use std::cell::OnceCell; -use std::collections::HashSet; -use std::fs::{self, DirEntry}; -use std::path::{Path, PathBuf}; - -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct SuiteQuery { - /// Numeric query id parsed from a `qNN.benchmark` file. - pub id: usize, - /// File name as it appears on disk, for example `q01.benchmark`. - pub file_name: String, - /// Full path to the benchmark file. - pub path: PathBuf, -} - -/// Parsed `.suite` file plus runtime metadata derived from its location. -/// -/// The serialized fields define the text configuration format. The skipped -/// fields are populated from the suite file path and used for discovery and -/// caching during a single runner invocation. -#[derive(Debug, Deserialize, Serialize)] -pub struct SuiteConfig { - /// Suite selector used on the command line, such as `tpch`. - pub name: String, - /// Human-readable suite description shown by `list`. - pub description: String, - /// Suite-specific options shown by `list`. - #[serde(default)] - pub options: Vec, - /// Path to the `.suite` file that produced this config. - #[serde(skip)] - pub suite_path: PathBuf, - /// Directory containing the `.suite` file. - #[serde(skip)] - pub suite_dir: PathBuf, - /// Lazily discovered benchmark query files for this suite. - #[serde(skip)] - pub(crate) query_cache: OnceCell>, -} - -impl Clone for SuiteConfig { - fn clone(&self) -> Self { - Self { - name: self.name.clone(), - description: self.description.clone(), - options: self.options.clone(), - suite_path: self.suite_path.clone(), - suite_dir: self.suite_dir.clone(), - query_cache: OnceCell::new(), - } - } -} - -#[derive(Debug, Clone, Deserialize, Serialize)] -pub struct SuiteOption { - /// Long option name, without the leading `--`. - pub name: String, - /// Optional short alias, without the leading `-`. - #[serde(default)] - pub short: Option, - /// Default option value. - pub default: String, - /// Allowed option values. - pub values: Vec, - /// Help text shown in command output. - pub help: String, -} - -/// Discovered suite metadata, sorted by suite name. -#[derive(Debug, Clone)] -pub struct SuiteRegistry { - suites: Vec, -} - -impl SuiteConfig { - /// Loads, parses, and validates one `.suite` file. - /// - /// The suite file path and containing directory are stored on the returned - /// config so later discovery can be resolved relative to the suite file. - pub fn from_file(path: impl AsRef) -> Result { - let suite_path = path.as_ref().to_path_buf(); - let contents = fs::read_to_string(&suite_path).map_err(|e| { - DataFusionError::External( - format!("failed to read suite file {}: {e}", suite_path.display()).into(), - ) - })?; - let mut suite: Self = toml::from_str(&contents).map_err(|e| { - DataFusionError::External( - format!("failed to parse suite file {}: {e}", suite_path.display()) - .into(), - ) - })?; - - suite.suite_dir = suite_path - .parent() - .map(Path::to_path_buf) - .unwrap_or_else(|| PathBuf::from(".")); - suite.suite_path = suite_path; - suite.validate()?; - - Ok(suite) - } - - /// Returns the directory recursively searched for query benchmark files. - pub fn query_search_root(&self) -> &Path { - &self.suite_dir - } - - /// Discovers and caches the suite's benchmark query files. - /// - /// Query files are found by recursively scanning from the suite directory, - /// accepting only `qNN.benchmark` files, sorting by numeric query id, and - /// rejecting duplicate ids. - pub fn discover_queries(&self) -> Result> { - if let Some(queries) = self.query_cache.get() { - return Ok(queries.clone()); - } - - let queries = self.scan_queries()?; - let _ = self.query_cache.set(queries.clone()); - Ok(queries) - } - - /// Performs uncached query discovery and duplicate-id validation. - fn scan_queries(&self) -> Result> { - let mut queries = Vec::new(); - - self.scan_query_dir(self.query_search_root(), &mut queries)?; - queries.sort_by(|left, right| { - left.id - .cmp(&right.id) - .then_with(|| left.path.cmp(&right.path)) - }); - - for pair in queries.windows(2) { - let [left, right] = pair else { - continue; - }; - if left.id == right.id { - return Err(DataFusionError::Configuration(format!( - "duplicate QUERY_ID {} in suite '{}': {} and {}", - left.id, - self.name, - left.path.display(), - right.path.display() - ))); - } - } - - Ok(queries) - } - - /// Recursively scans a directory and appends valid query benchmark files to - /// the provided collection. - fn scan_query_dir(&self, dir: &Path, queries: &mut Vec) -> Result<()> { - let mut entries = read_dir_entries(dir, "benchmark query directory")?; - entries.sort_by_key(|entry| entry.file_name()); - - for entry in entries { - let path = entry.path(); - let file_type = entry.file_type().map_err(|e| { - DataFusionError::External( - format!( - "failed to read benchmark query entry type {}: {e}", - path.display() - ) - .into(), - ) - })?; - - if file_type.is_dir() { - self.scan_query_dir(&path, queries)?; - continue; - } - - if path - .extension() - .is_none_or(|extension| extension != "benchmark") - { - continue; - } - - let file_name = entry.file_name().to_string_lossy().into_owned(); - if let Some(id) = parse_query_file_name(&file_name) { - queries.push(SuiteQuery { - id, - file_name, - path, - }); - } - } - - Ok(()) - } - - /// Validates suite metadata that cannot be enforced by TOML deserialization. - fn validate(&self) -> Result<()> { - self.validate_suite_fields()?; - self.validate_options() - } - - /// Validates required suite-level fields. - fn validate_suite_fields(&self) -> Result<()> { - if self.name.trim().is_empty() { - return Err(DataFusionError::Configuration( - "suite name cannot be empty".to_string(), - )); - } - - Ok(()) - } - - /// Validates suite-defined option declarations. - fn validate_options(&self) -> Result<()> { - let mut option_names = HashSet::new(); - for option in &self.options { - if !option_names.insert(option.name.as_str()) { - return Err(DataFusionError::Configuration(format!( - "duplicate option name '{}'", - option.name - ))); - } - - option.validate()?; - } - - Ok(()) - } -} - -impl SuiteOption { - /// Validates one suite-defined option. - fn validate(&self) -> Result<()> { - if self.name.trim().is_empty() { - return Err(DataFusionError::Configuration( - "option name cannot be empty".to_string(), - )); - } - - if !is_valid_cli_option_name(&self.name) { - return Err(DataFusionError::Configuration(format!( - "invalid option name '{}'; expected lowercase ASCII letters, digits, and hyphens", - self.name - ))); - } - - self.validate_short_alias()?; - - if self.help.trim().is_empty() { - return Err(DataFusionError::Configuration(format!( - "help for option '{}' cannot be empty", - self.name - ))); - } - - if self.values.is_empty() { - return Err(DataFusionError::Configuration(format!( - "values for option '{}' cannot be empty", - self.name - ))); - } - - let mut values = HashSet::new(); - for value in &self.values { - if !values.insert(value.as_str()) { - return Err(DataFusionError::Configuration(format!( - "duplicate value '{}' for option '{}'", - value, self.name - ))); - } - } - - if !self.values.contains(&self.default) { - return Err(DataFusionError::Configuration(format!( - "default value '{}' for option '{}' must be present in values", - self.default, self.name - ))); - } - - Ok(()) - } - - /// Validates the optional short alias for one suite-defined option. - fn validate_short_alias(&self) -> Result<()> { - let Some(short) = &self.short else { - return Ok(()); - }; - - if short.trim().is_empty() { - return Err(DataFusionError::Configuration(format!( - "short alias for option '{}' cannot be empty", - self.name - ))); - } - - if !is_valid_cli_option_name(short) { - return Err(DataFusionError::Configuration(format!( - "invalid short alias '{short}' for option '{}'; expected lowercase ASCII letters, digits, and hyphens", - self.name - ))); - } - - Ok(()) - } -} - -impl SuiteRegistry { - /// Discovers all suite files below the SQL benchmark root. - /// - /// Each direct child directory is searched for `.suite` files. Suite names - /// must be unique across the registry so command selectors are - /// unambiguous. - pub fn discover(root: impl AsRef) -> Result { - let root = root.as_ref(); - let mut suites = Vec::new(); - let mut suite_names = HashSet::new(); - - for entry in read_dir_entries(root, "benchmark suite root")? { - if !entry - .file_type() - .map_err(|e| { - DataFusionError::External( - format!( - "failed to read benchmark suite entry type {}: {e}", - entry.path().display() - ) - .into(), - ) - })? - .is_dir() - { - continue; - } - - let mut suite_files = Vec::new(); - let suite_dir = entry.path(); - for suite_entry in read_dir_entries(&suite_dir, "benchmark suite directory")? - { - let path = suite_entry.path(); - if path - .extension() - .is_some_and(|extension| extension == "suite") - { - suite_files.push(path); - } - } - suite_files.sort(); - - for suite_file in suite_files { - let suite = SuiteConfig::from_file(suite_file)?; - if !suite_names.insert(suite.name.clone()) { - return Err(DataFusionError::Configuration(format!( - "duplicate suite name '{}'", - suite.name - ))); - } - suites.push(suite); - } - } - - suites.sort_by(|left, right| left.name.cmp(&right.name)); - - Ok(Self { suites }) - } - - /// Returns discovered suites sorted by selector name. - pub fn suites(&self) -> &[SuiteConfig] { - &self.suites - } -} - -fn parse_query_file_name(file_name: &str) -> Option { - let query_id = file_name.strip_prefix('q')?.strip_suffix(".benchmark")?; - - if query_id.len() < 2 || !query_id.chars().all(|c| c.is_ascii_digit()) { - return None; - } - - query_id.parse().ok() -} - -fn is_valid_cli_option_name(name: &str) -> bool { - name.chars() - .all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-') -} - -fn read_dir_entries(dir: &Path, label: &str) -> Result> { - let entries = fs::read_dir(dir).map_err(|e| { - DataFusionError::External( - format!("failed to read {label} {}: {e}", dir.display()).into(), - ) - })?; - - entries - .collect::>>() - .map_err(|e| DataFusionError::External(Box::new(e))) -} - -#[cfg(test)] -mod tests { - use super::*; - use std::fs; - - fn manifest_path(path: &str) -> PathBuf { - PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path) - } - - #[test] - fn discovers_tpch_suite_file() { - let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); - - assert_eq!(registry.suites().len(), 1); - assert_eq!(registry.suites()[0].name, "tpch"); - assert_eq!(registry.suites()[0].options.len(), 2); - } - - #[test] - fn discovers_query_ids_in_numeric_order() { - let registry = SuiteRegistry::discover(manifest_path("sql_benchmarks")).unwrap(); - let suite = ®istry.suites()[0]; - let queries = suite.discover_queries().unwrap(); - let ids = queries.iter().map(|query| query.id).collect::>(); - - assert_eq!(ids.first(), Some(&1)); - assert_eq!(ids.last(), Some(&22)); - assert!(!ids.contains(&0)); - } - - #[test] - fn rejects_duplicate_suite_names() { - let dir = tempfile::tempdir().unwrap(); - let one = dir.path().join("one"); - let two = dir.path().join("two"); - fs::create_dir_all(&one).unwrap(); - fs::create_dir_all(&two).unwrap(); - fs::write( - one.join("suite.suite"), - "name = \"dup\"\ndescription = \"one\"\n", - ) - .unwrap(); - fs::write( - two.join("suite.suite"), - "name = \"dup\"\ndescription = \"two\"\n", - ) - .unwrap(); - - let err = SuiteRegistry::discover(dir.path()).unwrap_err(); - assert!(err.to_string().contains("duplicate suite name")); - } - - #[test] - fn rejects_invalid_option_metadata() { - let dir = tempfile::tempdir().unwrap(); - let suite_dir = dir.path().join("suite"); - fs::create_dir_all(&suite_dir).unwrap(); - fs::write( - suite_dir.join("bad.suite"), - r#" -name = "bad" -description = "bad suite" - -[[options]] -name = "BAD" -default = "one" -values = ["one"] -help = "bad" -"#, - ) - .unwrap(); - - let err = SuiteRegistry::discover(dir.path()).unwrap_err(); - assert!(err.to_string().contains("invalid option name")); - } -} diff --git a/benchmarks/src/bin/benchmark_runner.rs b/benchmarks/src/bin/benchmark_runner.rs index 0efd169947f3c..c1700c42ba2aa 100644 --- a/benchmarks/src/bin/benchmark_runner.rs +++ b/benchmarks/src/bin/benchmark_runner.rs @@ -15,7 +15,33 @@ // specific language governing permissions and limitations // under the License. -use datafusion_benchmarks::benchmark_runner::run_cli; +//! DataFusion SQL benchmark runner. + +use clap::{ + Arg, ArgAction, ArgMatches, Command, CommandFactory, FromArgMatches, Parser, + ValueEnum, +}; +use criterion::Criterion; +use datafusion::error::Result; +use datafusion::prelude::SessionContext; +use datafusion_benchmarks::sql_benchmark::SqlBenchmark; +use datafusion_benchmarks::sql_benchmark_runner::{ + BenchmarkFilter, SqlRunConfig, default_sql_benchmark_directory, ensure_selection, + filter_benchmarks, finish_benchmark, load_benchmark_definitions_for_query, make_ctx, + prepare_benchmark, run_criterion_benchmarks_impl, +}; +use datafusion_benchmarks::sql_benchmark_suite::{ + ReservedOptions, SuiteExample, SuiteMetadata, discover_suites, +}; +use datafusion_benchmarks::util::{BenchmarkRun, CommonOpt, print_memory_stats}; +use datafusion_common::instant::Instant; +use datafusion_common::{DataFusionError, exec_datafusion_err}; +use datafusion_common_runtime::SpawnedTask; +use serde::{Serialize, Serializer}; +use std::collections::{BTreeMap, BTreeSet}; +use std::ffi::{OsStr, OsString}; +use std::io::IsTerminal; +use std::path::Path; #[cfg(feature = "snmalloc")] #[global_allocator] @@ -27,10 +53,2137 @@ static ALLOC: snmalloc_rs::SnMalloc = snmalloc_rs::SnMalloc; #[global_allocator] static ALLOC: mimalloc::MiMalloc = mimalloc::MiMalloc; -fn main() { +#[tokio::main] +async fn main() { env_logger::init(); - if let Err(e) = run_cli(std::env::args()) { - eprintln!("{e}"); + if let Err(error) = run_cli().await { + eprintln!("Error: {error}"); std::process::exit(1); } } + +#[derive(Debug)] +enum CliAction { + List, + Simple(SqlRunConfig), + Criterion { + config: SqlRunConfig, + save_baseline: Option, + }, + DryRun(DryRunOutput), +} + +#[derive(Debug, Serialize)] +struct ResolvedSuiteValue { + value: String, + #[serde(serialize_with = "serialize_value_source")] + source: datafusion_benchmarks::sql_benchmark_suite::ValueSource, + environment: String, +} + +#[derive(Debug, Serialize)] +struct ResolvedPathValue { + value: String, + #[serde(serialize_with = "serialize_value_source")] + source: datafusion_benchmarks::sql_benchmark_suite::ValueSource, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)] +#[serde(rename_all = "snake_case")] +enum RunMode { + Simple, + Criterion, +} + +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, Serialize, ValueEnum)] +#[serde(rename_all = "snake_case")] +enum ResultMode { + #[default] + None, + Persist, + Validate, +} + +impl ResultMode { + fn config_flags(self) -> (bool, bool) { + match self { + Self::None => (false, false), + Self::Persist => (true, false), + Self::Validate => (false, true), + } + } +} + +#[derive(Debug, Serialize)] +struct DryRunCommonOptions { + iterations: usize, + partitions: Option, + batch_size: Option, +} + +#[derive(Debug, Serialize)] +struct DryRunOutput { + suite: String, + query: Option, + subgroup: Option, + mode: RunMode, + result_mode: ResultMode, + common_options: DryRunCommonOptions, + suite_options: BTreeMap, + path_replacements: BTreeMap, +} + +#[derive(Debug, Parser)] +#[command( + name = "benchmark_runner", + about = "Run DataFusion SQL benchmarks", + styles = criterion_like_styles(), +)] +struct Cli { + #[arg(value_name = "BENCHMARK", help = "SQL benchmark group to run")] + benchmark: Option, + + #[arg(short = 'q', long = "query", env = "BENCH_QUERY")] + query: Option, + + #[arg(long = "subgroup", env = "BENCH_SUBGROUP")] + subgroup: Option, + + #[command(flatten)] + common: CommonOpt, + + #[arg( + long = "criterion", + action = ArgAction::SetTrue, + help = "Run benchmarks with Criterion" + )] + criterion: bool, + + #[arg( + long = "list", + action = ArgAction::SetTrue, + help = "List available SQL benchmark groups" + )] + list: bool, + + #[arg( + short = 'o', + long = "output", + help = "Write simple runner results as JSON to this path" + )] + output: Option, + + #[arg( + long = "save-baseline", + value_name = "BASELINE", + help = "Save Criterion measurements to the named baseline" + )] + save_baseline: Option, + + #[arg(short = 'p', long = "path", value_name = "PATH")] + path: Option, + + #[arg( + long = "result-mode", + value_enum, + value_name = "MODE", + help = "Handle expected results: none, persist, or validate" + )] + result_mode: Option, + + #[arg(long = "dry-run", action = ArgAction::SetTrue)] + dry_run: bool, +} + +/// Parses CLI arguments, runs the selected action, and prints any output. +async fn run_cli() -> Result<()> { + let benchmark_dir = default_sql_benchmark_directory(); + let output = run_cli_from(std::env::args_os(), &benchmark_dir).await?; + + if !output.is_empty() { + println!("{output}"); + } + + Ok(()) +} + +fn serialize_value_source( + source: &datafusion_benchmarks::sql_benchmark_suite::ValueSource, + serializer: S, +) -> std::result::Result +where + S: Serializer, +{ + let value = match source { + datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine => { + "command_line" + } + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Environment => { + "environment" + } + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default => "default", + }; + serializer.serialize_str(value) +} + +fn clap_display_output(error: &DataFusionError) -> Option { + let DataFusionError::External(error) = error else { + return None; + }; + let error = error.downcast_ref::()?; + matches!( + error.kind(), + clap::error::ErrorKind::DisplayHelp | clap::error::ErrorKind::DisplayVersion + ) + .then(|| error.to_string()) +} + +async fn run_cli_from(args: I, benchmark_dir: &Path) -> Result +where + I: IntoIterator, + T: Into + Clone, +{ + match parse_cli_from(args, benchmark_dir) { + Ok(action) => run_cli_action(action, benchmark_dir).await, + Err(error) => clap_display_output(&error).ok_or(error), + } +} + +fn format_examples(examples: &[SuiteExample]) -> String { + if examples.is_empty() { + return String::new(); + } + + let mut output = String::from("Examples:\n"); + for example in examples { + output.push_str(" "); + output.push_str(example.command()); + output.push_str("\n "); + output.push_str(example.description()); + output.push('\n'); + } + output +} + +fn build_cli(suite: Option<&SuiteMetadata>) -> Command { + let mut command = Cli::command(); + + if let Some(suite) = suite { + command = command.about(suite.description().to_string()); + + for option in suite.options() { + let mut arg = Arg::new(option.name().to_string()) + .long(option.name().to_string()) + .help(option.help().to_string()) + .env(option.env().to_string()) + .default_value(option.default().to_string()); + + if let Some(short) = option.short() { + arg = arg.short(short); + } + if let Some(values) = option + .values() + .filter(|values| !values.iter().any(|value| value == "...")) + { + arg = arg.value_parser(values.to_vec()); + } + + command = command.arg(arg); + } + + let examples = format_examples(suite.examples()); + + if !examples.is_empty() { + command = command.after_help(examples); + } + } + + command +} + +fn reserved_options() -> (BTreeSet, BTreeSet) { + let command = Cli::command(); + let long = command + .get_arguments() + .filter_map(|arg| arg.get_long().map(ToOwned::to_owned)) + .collect(); + let short = command + .get_arguments() + .filter_map(|arg| arg.get_short()) + .collect(); + (long, short) +} + +fn suite_metadata(benchmark_dir: &Path) -> Result> { + let (long, short) = reserved_options(); + discover_suites( + benchmark_dir, + &ReservedOptions { + long: &long, + short: &short, + }, + ) +} + +fn format_suite_list(suites: &[SuiteMetadata]) -> String { + let mut output = String::from("SQL benchmarks:\n"); + for suite in suites { + let query_word = if suite.benchmark_count() == 1 { + "query " + } else { + "queries " + }; + output.push_str(&format!( + " {:<24} {} {query_word}{}\n", + suite.name(), + suite.benchmark_count(), + suite.description() + )); + } + output.trim_end().to_string() +} + +fn locate_suite_arg(args: &[OsString]) -> Result> { + let Some(argument) = args.get(1) else { + return Ok(None); + }; + if argument == OsStr::new("--help") + || argument == OsStr::new("-h") + || argument == OsStr::new("--list") + || argument == OsStr::new("--dry-run") + { + return Ok(None); + } + let suite = argument.to_str().ok_or_else(|| { + DataFusionError::External("suite name is not valid Unicode".into()) + })?; + if suite.starts_with('-') { + return Err(exec_datafusion_err!( + "suite must be the first argument; options must follow the suite" + )); + } + Ok(Some(suite)) +} + +fn try_parse_cli_from(args: I, benchmark_dir: &Path) -> Result +where + I: IntoIterator, + T: Into + Clone, +{ + let args = args.into_iter().map(Into::into).collect::>(); + let suite = locate_suite_arg(&args)?; + let (long, short) = reserved_options(); + let reserved = ReservedOptions { + long: &long, + short: &short, + }; + let suite = suite + .map(|name| { + if !benchmark_dir.join(name).is_dir() { + let available = discover_suites(benchmark_dir, &reserved)?; + return Err(exec_datafusion_err!( + "unknown benchmark '{name}'\n\n{}", + format_suite_list(&available) + )); + } + SuiteMetadata::load(benchmark_dir, name, &reserved) + }) + .transpose()?; + let matches = build_cli(suite.as_ref()) + .try_get_matches_from(args) + .map_err(|error| DataFusionError::External(Box::new(error)))?; + + cli_action_from_matches(&matches, suite.as_ref()) +} + +fn parse_cli_from(args: I, benchmark_dir: &Path) -> Result +where + I: IntoIterator, + T: Into + Clone, +{ + try_parse_cli_from(args, benchmark_dir) +} + +/// Converts parsed arguments into an executable action and validates mode options. +fn cli_action_from_matches( + matches: &ArgMatches, + suite: Option<&SuiteMetadata>, +) -> Result { + let cli = Cli::from_arg_matches(matches) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + + if cli.dry_run && cli.list { + return Err(exec_datafusion_err!("--list cannot be used with --dry-run")); + } + if cli.dry_run && cli.benchmark.is_none() { + return Err(exec_datafusion_err!("--dry-run requires a benchmark suite")); + } + if cli.list || cli.benchmark.is_none() { + return Ok(CliAction::List); + } + if cli.criterion && cli.output.is_some() { + return Err(exec_datafusion_err!( + "--output cannot be used with --criterion" + )); + } + if !cli.criterion && cli.save_baseline.is_some() { + return Err(exec_datafusion_err!( + "--save-baseline cannot be used without --criterion" + )); + } + if !cli.criterion && cli.common.iterations == 0 { + return Err(exec_datafusion_err!("iterations must be greater than zero")); + } + + // we need to know if iterations was set on the command line, not the default value + let iterations_from_cli = matches.value_source("iterations") + == Some(clap::parser::ValueSource::CommandLine); + + if cli.criterion && iterations_from_cli { + return Err(exec_datafusion_err!( + "--iterations cannot be used with --criterion" + )); + } + + let suite = + suite.ok_or_else(|| exec_datafusion_err!("benchmark suite is required"))?; + + if cli.path.is_some() && !suite.path_replacements().contains_key("DATA_DIR") { + return Err(exec_datafusion_err!( + "--path cannot be used because suite '{}' does not declare DATA_DIR", + suite.name() + )); + } + + let result_mode = resolve_result_mode(cli.result_mode)?; + let (persist_results, validate_results) = result_mode.config_flags(); + let mut config = SqlRunConfig { + common: cli.common, + filter: BenchmarkFilter { + name: cli.benchmark, + subgroup: cli.subgroup, + query: cli.query, + }, + replacements: Default::default(), + query_filename: None, + persist_results, + validate_results, + output: cli.output, + }; + let suite_options: BTreeMap = suite + .options() + .iter() + .map(|option| { + let value = matches + .get_one::(option.name()) + .expect("suite options always have defaults") + .clone(); + let source = match matches.value_source(option.name()) { + Some(clap::parser::ValueSource::CommandLine) => { + datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine + } + Some(clap::parser::ValueSource::EnvVariable) => { + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Environment + } + Some(clap::parser::ValueSource::DefaultValue) => { + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default + } + vs => unreachable!("unexpected suite option source: {vs:?}"), + }; + ( + option.name().to_string(), + ResolvedSuiteValue { + value, + source, + environment: option.env().to_string(), + }, + ) + }) + .collect(); + let path_replacements: BTreeMap = suite + .path_replacements() + .iter() + .map(|(key, default)| { + let (value, source) = if key == "DATA_DIR" { + cli.path.as_ref().map_or_else( + || { + ( + default.display().to_string(), + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default, + ) + }, + |path| { + ( + path.display().to_string(), + datafusion_benchmarks::sql_benchmark_suite::ValueSource::CommandLine, + ) + }, + ) + } else { + ( + default.display().to_string(), + datafusion_benchmarks::sql_benchmark_suite::ValueSource::Default, + ) + }; + ( + key.to_ascii_lowercase(), + ResolvedPathValue { value, source }, + ) + }) + .collect(); + + config.replacements = suite_options + .values() + .map(|resolved| { + ( + resolved.environment.to_ascii_lowercase(), + resolved.value.clone(), + ) + }) + .chain( + path_replacements + .iter() + .map(|(key, resolved)| (key.clone(), resolved.value.clone())), + ) + .collect(); + config.query_filename = config + .filter + .query + .as_deref() + .map(|query| suite.query_filename(query)) + .transpose()?; + + if cli.dry_run { + let mode = if cli.criterion { + RunMode::Criterion + } else { + RunMode::Simple + }; + return Ok(CliAction::DryRun(DryRunOutput { + suite: suite.name().to_string(), + query: config.filter.query.clone(), + subgroup: config.filter.subgroup.clone(), + mode, + result_mode, + common_options: DryRunCommonOptions { + iterations: config.common.iterations, + partitions: config.common.partitions, + batch_size: config.common.batch_size, + }, + suite_options, + path_replacements, + })); + } + + if cli.criterion { + Ok(CliAction::Criterion { + config, + save_baseline: cli.save_baseline, + }) + } else { + Ok(CliAction::Simple(config)) + } +} + +/// Executes a parsed CLI action and returns any text that should be printed. +async fn run_cli_action(action: CliAction, benchmark_dir: &Path) -> Result { + match action { + CliAction::List => Ok(format_suite_list(&suite_metadata(benchmark_dir)?)), + CliAction::Simple(config) => { + run_simple_benchmarks(benchmark_dir, config).await?; + Ok(String::new()) + } + CliAction::Criterion { + config, + save_baseline, + } => { + if config.output.is_some() { + return Err(exec_datafusion_err!( + "--output cannot be used with --criterion" + )); + } + let benchmark_dir = benchmark_dir.to_path_buf(); + + SpawnedTask::spawn_blocking(move || { + run_criterion_benchmarks( + &benchmark_dir, + &config, + save_baseline.as_deref(), + ) + }) + .await + .map_err(|e| DataFusionError::External(Box::new(e)))??; + + Ok(String::new()) + } + CliAction::DryRun(output) => serde_json::to_string_pretty(&output) + .map_err(|error| DataFusionError::External(Box::new(error))), + } +} + +fn resolve_result_mode(explicit: Option) -> Result { + if let Some(mode) = explicit { + return Ok(mode); + } + + let persist = parse_compat_bool("BENCH_PERSIST_RESULTS")?; + let validate = parse_compat_bool("BENCH_VALIDATE")?; + + Ok(if persist { + ResultMode::Persist + } else if validate { + ResultMode::Validate + } else { + ResultMode::None + }) +} + +fn parse_compat_bool(name: &str) -> Result { + let Some(value) = std::env::var_os(name) else { + return Ok(false); + }; + let value = value + .into_string() + .map_err(|_| exec_datafusion_err!("{name} contains invalid UTF-8"))?; + + value.parse::().map_err(|_| { + exec_datafusion_err!("invalid value '{value}' for {name}; expected true or false") + }) +} + +/// Builds the default Criterion runner and optionally records a named baseline. +fn run_criterion_benchmarks( + benchmark_dir: &Path, + config: &SqlRunConfig, + save_baseline: Option<&str>, +) -> Result<()> { + let mut criterion = Criterion::default() + .sample_size(10) + .with_output_color(std::io::stdout().is_terminal()); + + if let Some(save_baseline) = save_baseline { + criterion = criterion.save_baseline(save_baseline.to_string()); + } + + run_criterion_benchmarks_impl(benchmark_dir, config, &mut criterion)?; + criterion.final_summary(); + + Ok(()) +} + +/// Runs selected benchmarks with fixed iteration counts and optional JSON output. +pub async fn run_simple_benchmarks( + benchmark_dir: &Path, + config: SqlRunConfig, +) -> Result<()> { + if config.common.iterations == 0 { + return Err(exec_datafusion_err!("iterations must be greater than zero")); + } + + let listing_ctx = make_ctx(&config.common)?; + let all_benchmarks = load_benchmark_definitions_for_query( + &config.filter, + &listing_ctx, + benchmark_dir, + &config.replacements, + config.query_filename.as_deref(), + ) + .await?; + let selected = filter_benchmarks(&config.filter, all_benchmarks.clone()); + let mut run = BenchmarkRun::new(); + + ensure_selection(&config.filter, &all_benchmarks, &selected)?; + + for (_group, benchmarks) in selected { + for mut benchmark in benchmarks { + let ctx = make_ctx(&config.common)?; + let result = + run_simple_benchmark(&ctx, &mut benchmark, &config, &mut run).await; + let cleanup_result = benchmark.cleanup(&ctx).await; + + finish_benchmark(result, cleanup_result)?; + } + } + + run.maybe_write_json(config.output.as_ref())?; + + Ok(()) +} + +/// Runs one benchmark case, recording each timed iteration. +async fn run_simple_benchmark( + ctx: &SessionContext, + benchmark: &mut SqlBenchmark, + config: &SqlRunConfig, + run: &mut BenchmarkRun, +) -> Result<()> { + prepare_benchmark(ctx, benchmark, config).await?; + + let case_name = benchmark_case_name(benchmark); + + // Each case gets its own `SessionContext`, so hand over its pool before the + // case starts. + run.set_memory_pool(&ctx.runtime_env().memory_pool); + run.start_new_case(&case_name); + + for iteration in 0..config.common.iterations { + let start = Instant::now(); + let row_count = benchmark.run(ctx, false).await?; + let elapsed = start.elapsed(); + let ms = elapsed.as_secs_f64() * 1000.0; + + println!("{case_name} iteration {iteration}: {ms:.1} ms, {row_count} rows"); + + run.write_iter(elapsed, row_count); + } + + print_memory_stats(&*ctx.runtime_env().memory_pool); + + Ok(()) +} + +fn benchmark_case_name(benchmark: &SqlBenchmark) -> String { + let mut name = format!("{}/{}", benchmark.group(), benchmark.name()); + + if !benchmark.subgroup().is_empty() { + name.push('/'); + name.push_str(benchmark.subgroup()); + } + + name +} + +fn criterion_like_styles() -> clap::builder::Styles { + use clap::builder::styling::AnsiColor; + + clap::builder::Styles::styled() + .header(AnsiColor::Green.on_default().bold()) + .usage(AnsiColor::Green.on_default().bold()) + .literal(AnsiColor::Cyan.on_default().bold()) + .placeholder(AnsiColor::Cyan.on_default()) +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_benchmarks::sql_benchmark_runner::{ + load_benchmark_definitions, sort_benchmarks, unknown_benchmark_error, + }; + use datafusion_benchmarks::sql_benchmark_suite::ValueSource; + use std::collections::HashMap; + use std::ffi::OsString; + use std::fs; + use std::path::{Path, PathBuf}; + use std::sync::{Mutex, MutexGuard}; + + static ENV_MUTEX: Mutex<()> = Mutex::new(()); + + struct ScopedEnv { + previous: Vec<(&'static str, Option)>, + _lock: MutexGuard<'static, ()>, + } + + impl ScopedEnv { + fn set(name: &'static str, value: impl Into) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let previous = std::env::var_os(name); + // SAFETY: ENV_MUTEX serializes changes made through ScopedEnv in this + // test module; it does not synchronize environment access elsewhere. + unsafe { std::env::set_var(name, value.into()) }; + Self { + previous: vec![(name, previous)], + _lock: lock, + } + } + + fn remove(name: &'static str) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let previous = std::env::var_os(name); + // SAFETY: ENV_MUTEX serializes changes made through ScopedEnv in this + // test module; it does not synchronize environment access elsewhere. + unsafe { std::env::remove_var(name) }; + Self { + previous: vec![(name, previous)], + _lock: lock, + } + } + + fn set_many(changes: [(&'static str, Option<&str>); N]) -> Self { + let lock = ENV_MUTEX.lock().unwrap_or_else(|error| error.into_inner()); + let mut previous = Vec::with_capacity(N); + for (name, value) in changes { + previous.push((name, std::env::var_os(name))); + // SAFETY: this guard holds ENV_MUTEX until it restores all entries. + unsafe { + match value { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + } + Self { + previous, + _lock: lock, + } + } + } + + impl Drop for ScopedEnv { + fn drop(&mut self) { + // SAFETY: this guard holds ENV_MUTEX until after all entries are restored. + unsafe { + for (name, previous) in self.previous.drain(..).rev() { + match previous { + Some(value) => std::env::set_var(name, value), + None => std::env::remove_var(name), + } + } + } + } + } + + /// Loads benchmark definitions, applies CLI-style filters, and sorts each group. + async fn load_benchmarks( + filter: &BenchmarkFilter, + ctx: &SessionContext, + benchmark_dir: &Path, + ) -> Result>> { + let benches = + load_benchmark_definitions(filter, ctx, benchmark_dir, &Default::default()) + .await?; + let mut benches = filter_benchmarks(filter, benches); + + sort_benchmarks(&mut benches); + + Ok(benches) + } + + fn write_benchmark(root: &Path, relative_path: &str, contents: &str) -> PathBuf { + let path = root.join(relative_path); + + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(&path, contents).unwrap(); + + path + } + + fn write_suite(root: &Path, name: &str, description: &str) -> PathBuf { + let path = root.join(name).join(format!("{name}.suite")); + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(&path, format!("description = {description:?}\n")).unwrap(); + path + } + + fn common(iterations: usize) -> CommonOpt { + CommonOpt { + iterations, + partitions: None, + batch_size: None, + mem_pool_type: "fair".to_string(), + memory_limit: None, + sort_spill_reservation_bytes: None, + debug: false, + simulate_latency: false, + } + } + + async fn run_cli_with_dir(args: I, benchmark_dir: &Path) -> Result + where + I: IntoIterator, + T: Into + Clone, + { + run_cli_from(args, benchmark_dir).await + } + + fn suite_root() -> tempfile::TempDir { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + fs::write( + temp.path().join("alpha/alpha.suite"), + r#"description = "Alpha benchmark" + +[path_replacements] +DATA_DIR = "data" + +[[options]] +name = "format" +short = "f" +env = "ALPHA_FORMAT" +default = "parquet" +values = ["parquet", "csv"] +help = "Alpha input format" + +[[examples]] +command = "benchmark_runner alpha -q 1 -f csv" +description = "Run query one against CSV data." +"#, + ) + .unwrap(); + temp + } + + #[test] + fn suite_help_contains_metadata() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let error = + try_parse_cli_from(["benchmark_runner", "alpha", "--help"], temp.path()) + .unwrap_err(); + let help = error.to_string(); + + assert!(help.contains("Alpha benchmark"), "{help}"); + assert!(help.contains("--format"), "{help}"); + assert!(help.contains("ALPHA_FORMAT"), "{help}"); + assert!( + help.contains("benchmark_runner alpha -q 1 -f csv"), + "{help}" + ); + } + + #[test] + fn accepts_interleaved_named_options() { + let temp = suite_root(); + let action = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--partitions", + "2", + "--format", + "csv", + "--query", + "5", + ], + temp.path(), + ) + .unwrap(); + let CliAction::Simple(config) = action else { + panic!("expected simple run") + }; + + assert_eq!(config.common.partitions, Some(2)); + assert_eq!(config.filter.query.as_deref(), Some("5")); + assert_eq!(config.query_filename.as_deref(), Some("q05.benchmark")); + assert_eq!(config.replacements["alpha_format"], "csv"); + assert_eq!( + config.replacements["data_dir"], + temp.path().join("alpha/data").display().to_string() + ); + } + + #[test] + fn dry_run_uses_suite_default() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let action = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + + assert_eq!(output.suite_options["format"].value, "parquet"); + assert_eq!(output.suite_options["format"].source, ValueSource::Default); + } + + #[test] + fn result_mode_defaults_to_none() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let CliAction::DryRun(output) = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap() + else { + panic!("expected dry run"); + }; + assert_eq!(output.result_mode, ResultMode::None); + } + + #[tokio::test] + async fn result_mode_persist_writes_expected_results() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let result_path = temp.path().join("expected.csv"); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + &format!( + "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", + result_path.display() + ), + ); + + run_cli_from( + [ + "benchmark_runner", + "alpha", + "--query", + "1", + "--result-mode", + "persist", + ], + temp.path(), + ) + .await + .unwrap(); + + let persisted = fs::read_to_string(result_path).unwrap(); + assert!(persisted.contains("value"), "{persisted}"); + assert!(persisted.contains('1'), "{persisted}"); + } + + #[tokio::test] + async fn result_mode_validate_accepts_expected_results() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let result_path = temp.path().join("expected.csv"); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + &format!( + "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", + result_path.display() + ), + ); + fs::write(&result_path, "value\n1\n").unwrap(); + + run_cli_from( + [ + "benchmark_runner", + "alpha", + "--query", + "1", + "--result-mode", + "validate", + ], + temp.path(), + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn result_mode_validate_reports_mismatched_results() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let result_path = temp.path().join("expected.csv"); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + &format!( + "name Q01\ngroup alpha\n\nresult {}\n\nrun\nSELECT 1 AS value\n", + result_path.display() + ), + ); + fs::write(&result_path, "value\n2\n").unwrap(); + + let error = run_cli_from( + [ + "benchmark_runner", + "alpha", + "--query", + "1", + "--result-mode", + "validate", + ], + temp.path(), + ) + .await + .unwrap_err(); + assert!(error.to_string().contains("expected value"), "{error}"); + } + + #[test] + fn explicit_result_modes_populate_config() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", Some("invalid")), + ("BENCH_VALIDATE", Some("invalid")), + ]); + let temp = suite_root(); + for (value, expected, persist, validate) in [ + ("none", ResultMode::None, false, false), + ("persist", ResultMode::Persist, true, false), + ("validate", ResultMode::Validate, false, true), + ] { + let CliAction::DryRun(output) = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--result-mode", + value, + "--dry-run", + ], + temp.path(), + ) + .unwrap() else { + panic!("expected dry run"); + }; + assert_eq!(output.result_mode, expected); + + let CliAction::Simple(config) = parse_cli_from( + ["benchmark_runner", "alpha", "--result-mode", value], + temp.path(), + ) + .unwrap() else { + panic!("expected simple run"); + }; + assert_eq!(config.persist_results, persist); + assert_eq!(config.validate_results, validate); + } + } + + #[test] + fn compatibility_environment_resolves_result_mode() { + for (persist, validate, expected) in [ + (Some("true"), None, ResultMode::Persist), + (None, Some("true"), ResultMode::Validate), + (Some("true"), Some("true"), ResultMode::Persist), + (Some("false"), Some("false"), ResultMode::None), + ] { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", persist), + ("BENCH_VALIDATE", validate), + ]); + let temp = suite_root(); + let CliAction::DryRun(output) = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap() + else { + panic!("expected dry run"); + }; + assert_eq!(output.result_mode, expected); + } + } + + #[test] + fn invalid_result_mode_environment_is_rejected_without_cli_override() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", Some("invalid")), + ("BENCH_VALIDATE", Some("false")), + ]); + let temp = suite_root(); + let error = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap_err(); + assert!( + error.to_string().contains("BENCH_PERSIST_RESULTS"), + "{error}" + ); + } + + #[test] + fn invalid_result_mode_cli_value_lists_allowed_values() { + let _env = ScopedEnv::set_many([ + ("BENCH_PERSIST_RESULTS", None), + ("BENCH_VALIDATE", None), + ]); + let temp = suite_root(); + let error = parse_cli_from( + ["benchmark_runner", "alpha", "--result-mode", "invalid"], + temp.path(), + ) + .unwrap_err(); + let message = error.to_string(); + for allowed in ["none", "persist", "validate"] { + assert!(message.contains(allowed), "{message}"); + } + } + + #[test] + fn environment_beats_suite_default() { + let _env = ScopedEnv::set("ALPHA_FORMAT", "csv"); + let temp = suite_root(); + let action = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + + assert_eq!(output.suite_options["format"].value, "csv"); + assert_eq!( + output.suite_options["format"].source, + ValueSource::Environment + ); + } + + #[test] + fn cli_equals_syntax_beats_environment_and_default() { + let _env = ScopedEnv::set("ALPHA_FORMAT", "parquet"); + let temp = suite_root(); + let action = parse_cli_from( + ["benchmark_runner", "alpha", "--format=csv", "--dry-run"], + temp.path(), + ) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + + assert_eq!(output.suite_options["format"].value, "csv"); + assert_eq!( + output.suite_options["format"].source, + ValueSource::CommandLine + ); + } + + #[test] + fn empty_environment_value_is_validated() { + let _env = ScopedEnv::set("ALPHA_FORMAT", ""); + let temp = suite_root(); + let error = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap_err(); + + let message = error.to_string(); + assert!(message.contains("a value is required"), "{message}"); + assert!(message.contains("parquet, csv"), "{message}"); + } + + #[cfg(unix)] + #[test] + fn non_unicode_environment_value_is_rejected_by_clap() { + use std::os::unix::ffi::OsStringExt; + + let _env = ScopedEnv::set("ALPHA_FORMAT", OsString::from_vec(vec![0xff])); + let temp = suite_root(); + let error = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap_err(); + + assert!(error.to_string().contains("invalid UTF-8"), "{error}"); + } + + #[test] + fn attached_short_cli_value_beats_invalid_environment() { + let _env = ScopedEnv::set("ALPHA_FORMAT", "invalid"); + let temp = suite_root(); + let action = parse_cli_from( + ["benchmark_runner", "alpha", "-fcsv", "--dry-run"], + temp.path(), + ) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + + assert_eq!(output.suite_options["format"].value, "csv"); + assert_eq!( + output.suite_options["format"].source, + ValueSource::CommandLine + ); + } + + #[test] + fn invalid_suite_environment_does_not_block_help() { + let _env = ScopedEnv::set("ALPHA_FORMAT", "invalid"); + let temp = suite_root(); + let error = + try_parse_cli_from(["benchmark_runner", "alpha", "--help"], temp.path()) + .unwrap_err(); + let help = error.to_string(); + + assert!(help.contains("Alpha benchmark"), "{help}"); + assert!(help.contains("--format"), "{help}"); + } + + #[test] + fn dry_run_rejects_list_instead_of_listing() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let error = parse_cli_from( + ["benchmark_runner", "alpha", "--list", "--dry-run"], + temp.path(), + ) + .unwrap_err(); + + assert!(error.to_string().contains("--list"), "{error}"); + assert!(error.to_string().contains("--dry-run"), "{error}"); + } + + #[test] + fn dry_run_requires_suite() { + let temp = suite_root(); + let error = + parse_cli_from(["benchmark_runner", "--dry-run"], temp.path()).unwrap_err(); + + assert!(error.to_string().contains("--dry-run"), "{error}"); + assert!(error.to_string().contains("suite"), "{error}"); + } + + #[test] + fn dry_run_resolves_default_and_overridden_path() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let default_action = + parse_cli_from(["benchmark_runner", "alpha", "--dry-run"], temp.path()) + .unwrap(); + let CliAction::DryRun(default_output) = default_action else { + panic!("expected dry run") + }; + assert_eq!( + default_output.path_replacements["data_dir"].source, + ValueSource::Default + ); + + let override_action = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--path", + "/tmp/alpha-data", + "--dry-run", + ], + temp.path(), + ) + .unwrap(); + let CliAction::DryRun(override_output) = override_action else { + panic!("expected dry run") + }; + assert_eq!( + override_output.path_replacements["data_dir"].value, + "/tmp/alpha-data" + ); + assert_eq!( + override_output.path_replacements["data_dir"].source, + ValueSource::CommandLine + ); + } + + #[test] + fn path_is_rejected_without_data_dir_replacement() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + fs::write( + temp.path().join("alpha/alpha.suite"), + fs::read_to_string(temp.path().join("alpha/alpha.suite")) + .unwrap() + .replace("[path_replacements]\nDATA_DIR = \"data\"\n\n", ""), + ) + .unwrap(); + + let error = + parse_cli_from(["benchmark_runner", "alpha", "--path", "data"], temp.path()) + .unwrap_err(); + assert!(error.to_string().contains("--path"), "{error}"); + assert!(error.to_string().contains("DATA_DIR"), "{error}"); + } + + #[test] + fn criterion_dry_run_reports_mode_and_keeps_cross_checks() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let action = parse_cli_from( + ["benchmark_runner", "alpha", "--criterion", "--dry-run"], + temp.path(), + ) + .unwrap(); + let CliAction::DryRun(output) = action else { + panic!("expected dry run") + }; + assert_eq!(output.mode, RunMode::Criterion); + + let error = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion", + "--iterations", + "2", + "--dry-run", + ], + temp.path(), + ) + .unwrap_err(); + assert!(error.to_string().contains("--iterations"), "{error}"); + } + + #[test] + fn dry_run_rejects_invalid_query_form() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let error = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--query", + "../secret", + "--dry-run", + ], + temp.path(), + ) + .unwrap_err(); + + assert!( + error.to_string().contains("invalid query identifier"), + "{error}" + ); + } + + #[test] + fn uppercase_q_dry_run_uses_same_query_filename_as_lowercase_q() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + + assert!(matches!( + parse_cli_from( + ["benchmark_runner", "alpha", "--query", "Q1", "--dry-run",], + temp.path(), + ) + .unwrap(), + CliAction::DryRun(_) + )); + + let filename = |query| { + let CliAction::Simple(config) = parse_cli_from( + ["benchmark_runner", "alpha", "--query", query], + temp.path(), + ) + .unwrap() else { + panic!("expected simple run") + }; + config.query_filename + }; + + assert_eq!(filename("Q1"), filename("q1")); + } + + #[tokio::test] + async fn dry_run_returns_deterministic_json_without_reading_benchmarks() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + fs::write( + temp.path().join("alpha/benchmarks/q01.benchmark"), + "this benchmark is intentionally invalid", + ) + .unwrap(); + + let output = run_cli_with_dir( + [ + "benchmark_runner", + "alpha", + "--query", + "7", + "--partitions", + "2", + "--dry-run", + ], + temp.path(), + ) + .await + .unwrap(); + let json: serde_json::Value = serde_json::from_str(&output).unwrap(); + + assert_eq!(json["suite"], "alpha"); + assert_eq!(json["query"], "7"); + assert_eq!(json["mode"], "simple"); + assert_eq!(json["common_options"]["partitions"], 2); + assert_eq!(json["suite_options"]["format"]["source"], "default"); + } + + #[test] + fn cli_lists_when_benchmark_is_omitted() { + let temp = suite_root(); + let action = parse_cli_from(["benchmark_runner"], temp.path()).unwrap(); + + assert!(matches!(action, CliAction::List)); + } + + #[test] + fn cli_lists_with_explicit_list_flag() { + let temp = suite_root(); + let action = parse_cli_from(["benchmark_runner", "--list"], temp.path()).unwrap(); + + assert!(matches!(action, CliAction::List)); + } + + #[test] + fn cli_defaults_to_basic_runner() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let action = + parse_cli_from(["benchmark_runner", "alpha", "--query", "1"], temp.path()) + .unwrap(); + let CliAction::Simple(config) = action else { + panic!("expected basic runner"); + }; + + assert_eq!(config.filter.name.as_deref(), Some("alpha")); + assert_eq!(config.filter.query.as_deref(), Some("1")); + } + + #[test] + fn cli_reads_query_from_env() { + let _env = ScopedEnv::set("BENCH_QUERY", "8"); + let temp = suite_root(); + let action = parse_cli_from( + ["benchmark_runner", "alpha", "--format", "parquet"], + temp.path(), + ); + let action = action.unwrap(); + let CliAction::Simple(config) = action else { + panic!("expected basic runner"); + }; + + assert_eq!(config.filter.name.as_deref(), Some("alpha")); + assert_eq!(config.filter.query.as_deref(), Some("8")); + } + + #[test] + fn cli_accepts_criterion_runner() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let action = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion", + "--save-baseline", + "main", + ], + temp.path(), + ) + .unwrap(); + + let CliAction::Criterion { + config, + save_baseline, + } = action + else { + panic!("expected criterion runner"); + }; + + assert_eq!(config.filter.name.as_deref(), Some("alpha")); + assert_eq!(save_baseline.as_deref(), Some("main")); + } + + #[test] + fn cli_rejects_output_with_criterion() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion", + "--output", + "results.json", + ], + temp.path(), + ) + .unwrap_err(); + + assert!(err.to_string().contains("--output")); + assert!(err.to_string().contains("--criterion")); + } + + #[test] + fn cli_rejects_save_baseline_without_criterion() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + ["benchmark_runner", "alpha", "--save-baseline", "main"], + temp.path(), + ) + .unwrap_err(); + + assert!(err.to_string().contains("--save-baseline")); + assert!(err.to_string().contains("--criterion")); + } + + #[test] + fn cli_rejects_iterations_with_criterion() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + [ + "benchmark_runner", + "alpha", + "--criterion", + "--iterations", + "3", + ], + temp.path(), + ) + .unwrap_err(); + + assert!(err.to_string().contains("--iterations")); + assert!(err.to_string().contains("--criterion")); + } + + #[test] + fn cli_rejects_zero_basic_iterations() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let err = parse_cli_from( + ["benchmark_runner", "alpha", "--iterations", "0"], + temp.path(), + ) + .unwrap_err(); + + assert!(err.to_string().contains("iterations")); + } + + #[tokio::test] + async fn run_cli_lists_when_no_benchmark_is_supplied() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + write_suite(temp.path(), "alpha", "Alpha workload"); + + let output = run_cli_with_dir(["benchmark_runner"], temp.path()) + .await + .unwrap(); + + assert!(output.contains("SQL benchmarks:")); + assert!(output.contains("alpha")); + } + + #[tokio::test] + async fn run_cli_lists_with_explicit_list_flag() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + write_suite(temp.path(), "alpha", "Alpha workload"); + + let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) + .await + .unwrap(); + + assert!(output.contains("SQL benchmarks:")); + assert!(output.contains("alpha")); + } + + #[tokio::test] + async fn run_cli_top_level_help_is_successful_output() { + let temp = suite_root(); + let output = run_cli_with_dir(["benchmark_runner", "--help"], temp.path()) + .await + .unwrap(); + + assert!(output.contains("Run DataFusion SQL benchmarks"), "{output}"); + assert!(output.contains("Usage:"), "{output}"); + } + + #[tokio::test] + async fn run_cli_suite_help_is_successful_output() { + let _env = ScopedEnv::remove("ALPHA_FORMAT"); + let temp = suite_root(); + let output = + run_cli_with_dir(["benchmark_runner", "alpha", "--help"], temp.path()) + .await + .unwrap(); + + assert!(output.contains("Alpha benchmark"), "{output}"); + assert!(output.contains("--format"), "{output}"); + } + + #[tokio::test] + async fn run_cli_real_parse_error_remains_an_error() { + let temp = suite_root(); + let error = run_cli_with_dir( + ["benchmark_runner", "alpha", "--not-an-option"], + temp.path(), + ) + .await + .unwrap_err(); + + assert!(error.to_string().contains("unexpected argument"), "{error}"); + } + + #[tokio::test] + async fn run_cli_reports_unknown_benchmark_with_list() { + let temp = suite_root(); + + let err = run_cli_with_dir(["benchmark_runner", "missing"], temp.path()) + .await + .unwrap_err(); + let message = err.to_string(); + + assert!(message.contains("unknown benchmark 'missing'"), "{message}"); + assert!(message.contains("alpha"), "{message}"); + } + + #[test] + fn criterion_runner_saves_named_baseline() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + + let output = tempfile::tempdir().unwrap(); + let mut criterion = Criterion::default() + .sample_size(10) + .warm_up_time(std::time::Duration::from_millis(1)) + .measurement_time(std::time::Duration::from_millis(10)) + .without_plots() + .output_directory(output.path()) + .save_baseline("acceptance".to_string()); + let config = SqlRunConfig { + common: common(3), + filter: BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("1".to_string()), + }, + replacements: HashMap::new(), + query_filename: None, + persist_results: false, + validate_results: false, + output: None, + }; + + run_criterion_benchmarks_impl(temp.path(), &config, &mut criterion).unwrap(); + criterion.final_summary(); + + assert!( + output + .path() + .join("alpha") + .join("Q01") + .join("acceptance") + .join("estimates.json") + .exists() + ); + } + + #[tokio::test] + async fn simple_runner_reports_unknown_query_for_known_benchmark() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + + let config = SqlRunConfig { + common: common(1), + filter: BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("9".to_string()), + }, + replacements: HashMap::new(), + query_filename: None, + persist_results: false, + validate_results: false, + output: None, + }; + let err = run_simple_benchmarks(temp.path(), config) + .await + .unwrap_err(); + let message = err.to_string(); + + assert!( + message.contains("no SQL benchmark query matched benchmark 'alpha'"), + "{message}" + ); + assert!(message.contains("query '9'"), "{message}"); + assert!(message.contains("normalized: 'Q09'"), "{message}"); + assert!(message.contains("Available alpha queries:"), "{message}"); + assert!(message.contains("Q01"), "{message}"); + assert!(!message.contains("unknown benchmark 'alpha'"), "{message}"); + } + + #[tokio::test] + async fn simple_runner_reports_unknown_subgroup_for_known_benchmark() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\nsubgroup wide\n\nrun\nSELECT 1\n", + ); + + let config = SqlRunConfig { + common: common(1), + filter: BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: Some("narrow".to_string()), + query: None, + }, + replacements: HashMap::new(), + query_filename: None, + persist_results: false, + validate_results: false, + output: None, + }; + let err = run_simple_benchmarks(temp.path(), config) + .await + .unwrap_err(); + let message = err.to_string(); + + assert!( + message.contains( + "no SQL benchmark subgroup matched benchmark 'alpha' with subgroup 'narrow'" + ), + "{message}" + ); + assert!(message.contains("Available alpha subgroups:"), "{message}"); + assert!(message.contains("wide"), "{message}"); + assert!(!message.contains("unknown benchmark 'alpha'"), "{message}"); + } + + #[tokio::test] + async fn basic_runner_executes_iterations_and_writes_json() { + let temp = tempfile::tempdir().unwrap(); + let output = temp.path().join("results.json"); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT * FROM (VALUES (1), (2)) AS t(v)\n", + ); + + let config = SqlRunConfig { + common: common(2), + filter: BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("1".to_string()), + }, + replacements: HashMap::new(), + query_filename: None, + persist_results: false, + validate_results: false, + output: Some(output.clone()), + }; + + run_simple_benchmarks(temp.path(), config).await.unwrap(); + + let json = fs::read_to_string(output).unwrap(); + + assert!(json.contains("\"query\": \"alpha/Q01\"")); + assert!(json.contains("\"row_count\": 2")); + assert_eq!(json.matches("\"row_count\": 2").count(), 2); + } + + #[tokio::test] + async fn basic_runner_reports_run_and_cleanup_failures() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT * FROM missing_run_table\n\ncleanup\nSELECT * FROM missing_cleanup_table\n", + ); + + let config = SqlRunConfig { + common: common(1), + filter: BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("1".to_string()), + }, + replacements: HashMap::new(), + query_filename: None, + persist_results: false, + validate_results: false, + output: None, + }; + let err = run_simple_benchmarks(temp.path(), config) + .await + .unwrap_err(); + let message = err.to_string(); + + assert!(message.contains("missing_run_table"), "{message}"); + assert!(message.contains("cleanup also failed"), "{message}"); + assert!(message.contains("missing_cleanup_table"), "{message}"); + } + + #[tokio::test] + async fn discovery_lists_groups_from_directories() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + write_benchmark( + temp.path(), + "beta/benchmarks/q02.benchmark", + "name Q02\n\nrun\nSELECT 2\n", + ); + let ctx = SessionContext::new(); + let benches = load_benchmarks(&BenchmarkFilter::default(), &ctx, temp.path()) + .await + .unwrap(); + + assert_eq!(benches["alpha"].len(), 1); + assert_eq!(benches["beta"].len(), 1); + } + + #[tokio::test] + async fn discovery_filters_benchmark_subgroup_and_query() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\nsubgroup wide\n\nrun\nSELECT 1\n", + ); + write_benchmark( + temp.path(), + "alpha/benchmarks/q02.benchmark", + "name Q02\nsubgroup narrow\n\nrun\nSELECT 2\n", + ); + + let ctx = SessionContext::new(); + let filter = BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: Some("wide".to_string()), + query: Some("1".to_string()), + }; + let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); + + assert_eq!(benches.len(), 1); + assert_eq!(benches["alpha"].len(), 1); + assert_eq!(benches["alpha"][0].name(), "Q01"); + } + + #[tokio::test] + async fn cli_subgroup_filter_is_used_for_benchmark_replacements() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "wide_schema/benchmarks/q01.benchmark", + "name Q01\nsubgroup ${BENCH_SUBGROUP:-wide}\n\nrun\nSELECT '${BENCH_SUBGROUP:-wide}'\n", + ); + + let ctx = SessionContext::new(); + let filter = BenchmarkFilter { + name: Some("wide_schema".to_string()), + subgroup: Some("narrow".to_string()), + query: None, + }; + let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); + + assert_eq!(benches["wide_schema"].len(), 1); + assert_eq!(benches["wide_schema"][0].subgroup(), "narrow"); + } + + #[tokio::test] + async fn benchmark_replacements_use_explicit_data_dir() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "clickbench/benchmarks/q01.benchmark", + "name Q01\nsubgroup ${DATA_DIR:-data}\n\nrun\nSELECT 1\n", + ); + + let ctx = SessionContext::new(); + let expected = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("data") + .to_string_lossy() + .into_owned(); + let replacements = HashMap::from([("data_dir".to_string(), expected.clone())]); + let benches = load_benchmark_definitions( + &BenchmarkFilter::default(), + &ctx, + temp.path(), + &replacements, + ) + .await + .unwrap(); + + assert_eq!(benches["clickbench"][0].subgroup(), expected); + } + + #[tokio::test] + async fn query_filter_matches_starts_with_when_exact_match_is_absent() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01a.benchmark", + "name Q01a\n\nrun\nSELECT 1\n", + ); + + let ctx = SessionContext::new(); + let filter = BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("1".to_string()), + }; + let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); + + assert_eq!(benches["alpha"].len(), 1); + assert_eq!(benches["alpha"][0].name(), "Q01a"); + } + + #[tokio::test] + async fn query_filter_matches_token_start_when_exact_match_is_absent() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "predicate_eval/benchmarks/costsel/q01.benchmark", + "name costsel_q01_regexp_selective_last\n\nrun\nSELECT 1\n", + ); + + let ctx = SessionContext::new(); + let filter = BenchmarkFilter { + name: Some("predicate_eval".to_string()), + subgroup: None, + query: Some("1".to_string()), + }; + let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); + + assert_eq!(benches["predicate_eval"].len(), 1); + assert_eq!( + benches["predicate_eval"][0].name(), + "costsel_q01_regexp_selective_last" + ); + } + + #[tokio::test] + async fn query_filter_prefers_starts_with_match_over_token_match() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/token.benchmark", + "name costsel_q01_regexp_selective_last\n\nrun\nSELECT 1\n", + ); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01a.benchmark", + "name Q01a\n\nrun\nSELECT 2\n", + ); + + let ctx = SessionContext::new(); + let filter = BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("1".to_string()), + }; + let benches = load_benchmarks(&filter, &ctx, temp.path()).await.unwrap(); + + assert_eq!(benches["alpha"].len(), 1); + assert_eq!(benches["alpha"][0].name(), "Q01a"); + } + + #[tokio::test] + async fn list_output_is_sorted_and_includes_counts_and_descriptions() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "beta/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + write_benchmark( + temp.path(), + "beta/benchmarks/q02.benchmark", + "name Q02\n\nrun\nSELECT 2\n", + ); + write_suite(temp.path(), "alpha", "Alpha workload"); + write_suite(temp.path(), "beta", "Beta workload"); + + let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) + .await + .unwrap(); + + assert_eq!( + output, + "SQL benchmarks:\n alpha 1 query Alpha workload\n beta 2 queries Beta workload" + ); + } + + #[tokio::test] + async fn list_does_not_parse_benchmark_sql() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "not valid benchmark syntax", + ); + write_suite(temp.path(), "alpha", "Alpha workload"); + + let output = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) + .await + .unwrap(); + + assert_eq!( + output, + "SQL benchmarks:\n alpha 1 query Alpha workload" + ); + } + + #[tokio::test] + async fn list_malformed_metadata_names_its_file() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + let metadata_path = temp.path().join("alpha/alpha.suite"); + fs::write(&metadata_path, "not valid metadata").unwrap(); + + let error = run_cli_with_dir(["benchmark_runner", "--list"], temp.path()) + .await + .unwrap_err(); + + assert!( + error + .to_string() + .contains(&metadata_path.display().to_string()), + "{error}" + ); + } + + #[tokio::test] + async fn unknown_benchmark_error_includes_available_benchmarks() { + let temp = tempfile::tempdir().unwrap(); + + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nrun\nSELECT 1\n", + ); + let ctx = SessionContext::new(); + let benches = load_benchmarks(&BenchmarkFilter::default(), &ctx, temp.path()) + .await + .unwrap(); + let message = unknown_benchmark_error("missing", &benches).to_string(); + + assert!(message.contains("unknown benchmark 'missing'"), "{message}"); + assert!(message.contains("alpha"), "{message}"); + } +} diff --git a/benchmarks/src/bin/external_aggr.rs b/benchmarks/src/bin/external_aggr.rs index a6e322c7fabc0..c19554eb33583 100644 --- a/benchmarks/src/bin/external_aggr.rs +++ b/benchmarks/src/bin/external_aggr.rs @@ -34,7 +34,7 @@ use datafusion::datasource::listing::{ use datafusion::datasource::{MemTable, TableProvider}; use datafusion::error::Result; use datafusion::execution::SessionStateBuilder; -use datafusion::execution::memory_pool::FairSpillPool; +use datafusion::execution::memory_pool::{FairSpillPool, PeakRecordingPool}; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::physical_plan::{collect, displayable}; @@ -169,7 +169,7 @@ impl ExternalAggrConfig { )); let query_results = self - .benchmark_query(query_id, mem_limit, mem_pool_type) + .benchmark_query(query_id, mem_limit, mem_pool_type, &mut benchmark_run) .await?; for iter in query_results { benchmark_run.write_iter(iter.elapsed, iter.row_count); @@ -182,11 +182,15 @@ impl ExternalAggrConfig { } /// Benchmark query `query_id` in `AGGR_QUERIES` + /// + /// `benchmark_run` is handed this query's runtime, which is built here + /// because each query runs under its own memory limit. async fn benchmark_query( &self, query_id: usize, mem_limit: u64, mem_pool_type: &str, + benchmark_run: &mut BenchmarkRun, ) -> Result> { let query_name = format!("Q{query_id}({})", human_readable_size(mem_limit as usize)); @@ -198,6 +202,12 @@ impl ExternalAggrConfig { return exec_err!("Invalid memory pool type: {}", mem_pool_type); } }; + // This benchmark builds its pool directly rather than going through + // `CommonOpt::runtime_env_builder`, so it has to install the recorder + // itself to report a peak. + let memory_pool: Arc = + Arc::new(PeakRecordingPool::new(memory_pool)); + benchmark_run.set_memory_pool(&memory_pool); let runtime_env = RuntimeEnvBuilder::new() .with_memory_pool(memory_pool) .build_arc()?; @@ -318,9 +328,7 @@ impl ExternalAggrConfig { ); let extension = DEFAULT_PARQUET_EXTENSION; - let options = ListingOptions::new(format) - .with_file_extension(extension) - .with_collect_stat(state.config().collect_statistics()); + let options = ListingOptions::new(format).with_file_extension(extension); let table_path = ListingTableUrl::parse(path)?; let config = ListingTableConfig::new(table_path).with_listing_options(options); diff --git a/benchmarks/src/cancellation.rs b/benchmarks/src/cancellation.rs index d3da1b0e83623..5f7fdcc43d99d 100644 --- a/benchmarks/src/cancellation.rs +++ b/benchmarks/src/cancellation.rs @@ -37,11 +37,11 @@ use datafusion::prelude::*; use datafusion_common::instant::Instant; use futures::TryStreamExt; use object_store::ObjectStore; +use object_store::buffered::BufWriter; use parquet::arrow::AsyncArrowWriter; -use parquet::arrow::async_writer::ParquetObjectWriter; use rand::Rng; use rand::distr::Alphanumeric; -use rand::rngs::ThreadRng; +use rand::prelude::*; use tokio::runtime::Runtime; use tokio_util::sync::CancellationToken; @@ -215,7 +215,8 @@ async fn find_or_generate_files( if files_on_disk.is_empty() { println!("No data files found, generating (this will take a bit)"); - generate_data(data_dir.as_ref(), num_files, num_rows_per_file).await?; + let mut rng = StdRng::seed_from_u64(0); + generate_data(&mut rng, data_dir.as_ref(), num_files, num_rows_per_file).await?; println!("Done generating files"); let files_on_disk = find_files_on_disk(data_dir)?; @@ -269,6 +270,7 @@ async fn load_data( } async fn generate_data( + rng: &mut StdRng, data_dir: impl AsRef, num_files: usize, num_rows_per_file: usize, @@ -295,12 +297,12 @@ async fn generate_data( for file_num in 1..=num_files { println!("Generating file {file_num} of {num_files}"); let data = columns.iter().map(|(column_name, column_type)| { - let column = random_data(column_type, num_rows_per_file); + let column = random_data(rng, column_type, num_rows_per_file); (column_name, column) }); let to_write = RecordBatch::try_from_iter(data).unwrap(); let path = object_store::path::Path::from(format!("{file_num}.parquet").as_str()); - let object_store_writer = ParquetObjectWriter::new(Arc::clone(&store) as _, path); + let object_store_writer = BufWriter::new(Arc::clone(&store) as _, path); let mut writer = AsyncArrowWriter::try_new(object_store_writer, to_write.schema(), None)?; @@ -311,13 +313,12 @@ async fn generate_data( Ok(()) } -fn random_data(column_type: &DataType, rows: usize) -> Arc { - let mut rng = rand::rng(); - let values = (0..rows).map(|_| random_value(&mut rng, column_type)); +fn random_data(rng: &mut StdRng, column_type: &DataType, rows: usize) -> Arc { + let values = (0..rows).map(|_| random_value(rng, column_type)); ScalarValue::iter_to_array(values).unwrap() } -fn random_value(rng: &mut ThreadRng, column_type: &DataType) -> ScalarValue { +fn random_value(rng: &mut StdRng, column_type: &DataType) -> ScalarValue { match column_type { DataType::Float64 => ScalarValue::Float64(Some(rng.random())), DataType::Boolean => ScalarValue::Boolean(Some(rng.random())), diff --git a/benchmarks/src/clickbench.rs b/benchmarks/src/clickbench.rs index 70aaeb7d2d192..a2e65aa5618a9 100644 --- a/benchmarks/src/clickbench.rs +++ b/benchmarks/src/clickbench.rs @@ -213,6 +213,7 @@ impl RunOpt { self.register_hits(&ctx).await?; let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query_path = get_query_path(&self.queries_path, query_id); let Some(sql) = get_query_sql(&query_path)? else { @@ -278,7 +279,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } diff --git a/benchmarks/src/dict.rs b/benchmarks/src/dict.rs index f8451715ea81e..e04b5f816adcc 100644 --- a/benchmarks/src/dict.rs +++ b/benchmarks/src/dict.rs @@ -333,6 +333,7 @@ impl RunOpt { let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query = &DICTIONARY_QUERIES[query_id - 1]; diff --git a/benchmarks/src/h2o.rs b/benchmarks/src/h2o.rs index 8b6e04932cb39..feb4bf2fa11ce 100644 --- a/benchmarks/src/h2o.rs +++ b/benchmarks/src/h2o.rs @@ -109,6 +109,7 @@ impl RunOpt { let iterations = self.common.iterations; let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { benchmark_run.start_new_case(&format!("Query {query_id}")); let sql = queries.get_query(query_id)?; @@ -131,7 +132,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); if self.common.debug { ctx.sql(sql) diff --git a/benchmarks/src/hj.rs b/benchmarks/src/hj.rs index 301fe0d599cd6..4f97b24d0f02c 100644 --- a/benchmarks/src/hj.rs +++ b/benchmarks/src/hj.rs @@ -25,8 +25,6 @@ use std::path::PathBuf; use futures::StreamExt; -// TODO: Add existence joins - /// Run the Hash Join benchmark /// /// This micro-benchmark focuses on the performance characteristics of Hash Joins. @@ -59,6 +57,7 @@ struct HashJoinQuery { prob_hit: f64, build_size: &'static str, probe_size: &'static str, + isolate_partitioned_join: bool, } /// Inline SQL queries for Hash Join benchmarks @@ -71,6 +70,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 1.0, build_size: "25", probe_size: "1.5M", + isolate_partitioned_join: false, }, // Q2: Very Small Build Side (Sparse, range < 1024) // Build Side: nation (25 rows, range 961) | Probe Side: customer (1.5M rows) @@ -87,6 +87,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 1.0, build_size: "25", probe_size: "1.5M", + isolate_partitioned_join: false, }, // Q3: 100% Density, 100% Hit rate HashJoinQuery { @@ -95,6 +96,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 1.0, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q4: 100% Density, 10% Hit rate HashJoinQuery { @@ -110,6 +112,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 0.1, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q5: 75% Density, 100% Hit rate HashJoinQuery { @@ -125,6 +128,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 1.0, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q6: 75% Density, 10% Hit rate HashJoinQuery { @@ -144,6 +148,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 0.1, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q7: 50% Density, 100% Hit rate HashJoinQuery { @@ -159,6 +164,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 1.0, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q8: 50% Density, 10% Hit rate HashJoinQuery { @@ -178,6 +184,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 0.1, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q9: 20% Density, 100% Hit rate HashJoinQuery { @@ -193,6 +200,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 1.0, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q10: 20% Density, 10% Hit rate HashJoinQuery { @@ -212,6 +220,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 0.1, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q11: 10% Density, 100% Hit rate HashJoinQuery { @@ -227,6 +236,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 1.0, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q12: 10% Density, 10% Hit rate HashJoinQuery { @@ -246,6 +256,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 0.1, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q13: 1% Density, 100% Hit rate HashJoinQuery { @@ -261,6 +272,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 1.0, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q14: 1% Density, 10% Hit rate HashJoinQuery { @@ -280,6 +292,7 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 0.1, build_size: "100K", probe_size: "60M", + isolate_partitioned_join: false, }, // Q15: 20% Density, 10% Hit rate, 20% Duplicates in Build Side HashJoinQuery { @@ -302,6 +315,208 @@ const HASH_QUERIES: &[HashJoinQuery] = &[ prob_hit: 0.1, build_size: "100K_(20%_dups)", probe_size: "60M", + isolate_partitioned_join: false, + }, + // RightSemi Join benchmarks with Int32 keys + // + // Fanout (average build rows matched per probe row, as measured by running + // the equivalent INNER JOIN under `EXPLAIN ANALYZE` and reading the + // `HashJoinExec` metrics): 1 for Q16-Q18. Build keys here are primary + // keys (`n_nationkey`, `s_suppkey`), so each probe row matches at most + // one build row. `prob_hit` controls what fraction of probe rows find + // that one match. + // + // Fanout still matters because semi joins short-circuit after the first + // match. Coverage of fanout > 1 (build-side duplicates) is left for a + // follow-up. + // + // Q16: RightSemi, Small build (25 rows), 100% Hit rate + // Build Side: nation (25 rows) | Probe Side: customer (1.5M rows) + HashJoinQuery { + sql: r###"SELECT c.k + FROM (SELECT CAST(n_nationkey AS INT) as k FROM nation) n + RIGHT SEMI JOIN (SELECT CAST(c_nationkey AS INT) as k FROM customer) c + ON n.k = c.k"###, + density: 1.0, + prob_hit: 1.0, + build_size: "25", + probe_size: "1.5M_RightSemi", + isolate_partitioned_join: false, + }, + // Q17: RightSemi, Medium build (100K rows), 100% Hit rate + // Build Side: supplier (100K rows) | Probe Side: lineitem (60M rows) + HashJoinQuery { + sql: r###"SELECT l.k + FROM (SELECT CAST(s_suppkey AS INT) as k FROM supplier) s + RIGHT SEMI JOIN (SELECT CAST(l_suppkey AS INT) as k FROM lineitem) l + ON s.k = l.k"###, + density: 1.0, + prob_hit: 1.0, + build_size: "100K", + probe_size: "60M_RightSemi", + isolate_partitioned_join: false, + }, + // Q18: RightSemi, Medium build (100K rows), 10% Hit rate + // Build Side: supplier (100K rows) | Probe Side: lineitem (60M rows) + HashJoinQuery { + sql: r###"SELECT l.k + FROM (SELECT CAST(s_suppkey AS INT) as k FROM supplier) s + RIGHT SEMI JOIN ( + SELECT CAST(CASE WHEN l_suppkey % 10 = 0 THEN l_suppkey ELSE l_suppkey + 1000000 END AS INT) as k + FROM lineitem + ) l + ON s.k = l.k"###, + density: 1.0, + prob_hit: 0.1, + build_size: "100K", + probe_size: "60M_RightSemi", + isolate_partitioned_join: false, + }, + // RightAnti Join benchmarks with Int32 keys + // + // Fanout (average build rows matched per probe row, as measured by running + // the equivalent INNER JOIN under `EXPLAIN ANALYZE` and reading the + // `HashJoinExec` metrics): 1 for Q19-Q21. Build keys here are primary + // keys (`n_nationkey`, `s_suppkey`), so each probe row matches at most + // one build row. `prob_hit` controls what fraction of probe rows find + // that one match (and are therefore filtered *out* by anti). + // + // Fanout still matters because anti joins short-circuit after the first + // match. Coverage of fanout > 1 (build-side duplicates) is left for a + // follow-up. + // + // Q19: RightAnti, Small build (25 rows), 100% Hit rate (no output) + // Build Side: nation (25 rows) | Probe Side: customer (1.5M rows) + HashJoinQuery { + sql: r###"SELECT c.k + FROM (SELECT CAST(n_nationkey AS INT) as k FROM nation) n + RIGHT ANTI JOIN (SELECT CAST(c_nationkey AS INT) as k FROM customer) c + ON n.k = c.k"###, + density: 1.0, + prob_hit: 1.0, + build_size: "25", + probe_size: "1.5M_RightAnti", + isolate_partitioned_join: false, + }, + // Q20: RightAnti, Medium build (100K rows), 100% Hit rate (no output) + // Build Side: supplier (100K rows) | Probe Side: lineitem (60M rows) + HashJoinQuery { + sql: r###"SELECT l.k + FROM (SELECT CAST(s_suppkey AS INT) as k FROM supplier) s + RIGHT ANTI JOIN (SELECT CAST(l_suppkey AS INT) as k FROM lineitem) l + ON s.k = l.k"###, + density: 1.0, + prob_hit: 1.0, + build_size: "100K", + probe_size: "60M_RightAnti", + isolate_partitioned_join: false, + }, + // Q21: RightAnti, Medium build (100K rows), 10% Hit rate (90% output) + // Build Side: supplier (100K rows) | Probe Side: lineitem (60M rows) + HashJoinQuery { + sql: r###"SELECT l.k + FROM (SELECT CAST(s_suppkey AS INT) as k FROM supplier) s + RIGHT ANTI JOIN ( + SELECT CAST(CASE WHEN l_suppkey % 10 = 0 THEN l_suppkey ELSE l_suppkey + 1000000 END AS INT) as k + FROM lineitem + ) l + ON s.k = l.k"###, + density: 1.0, + prob_hit: 0.1, + build_size: "100K", + probe_size: "60M_RightAnti", + isolate_partitioned_join: false, + }, + // Q22: RightSemi, Medium build (100K rows), ~1% Hit rate, fanout ~100 + // + // Build Side: supplier (100K rows) collapsed onto 1K distinct keys + // Probe Side: lineitem (60M rows). Each matching probe row produces many + // duplicate probe indices before RightSemi deduplication. + HashJoinQuery { + sql: r###"SELECT l.k + FROM ( + SELECT CAST(((s_suppkey - 1) % 1000) + 1 AS INT) as k + FROM supplier + ) s + RIGHT SEMI JOIN ( + SELECT CAST(l_suppkey AS INT) as k + FROM lineitem + ) l + ON s.k = l.k"###, + density: 1.0, + prob_hit: 0.01, + build_size: "100K_(fanout_100)", + probe_size: "60M_RightSemi", + isolate_partitioned_join: false, + }, + // Q23: skewed high-fanout string-key inner join. + // Build ~32K rows / ~415 distinct keys (fanout ~78), probe ~2.3M rows all + // carrying the same dominant key — one partition does nearly all the work. + // Long keys (~28 chars) make per-pair key comparison expensive; count(*) + // isolates the match path. + HashJoinQuery { + sql: r###"SELECT count(*) + FROM ( + SELECT 'high_fanout_string_join_key_' || CAST((s_suppkey % 415) + 1 AS VARCHAR) as k + FROM supplier + WHERE s_suppkey <= 32340 + ) s + JOIN ( + SELECT 'high_fanout_string_join_key_1' as k + FROM lineitem + WHERE l_orderkey % 265 = 0 + ) l ON s.k = l.k"###, + density: 1.0, + prob_hit: 1.0, + build_size: "32K_(fanout~78)", + probe_size: "2.3M_long_keys_count", + isolate_partitioned_join: true, + }, + // Q24: single-hot-bucket long string-key inner join. + // Build rows all share one long string key, so each matching probe row fans + // out to the whole build side. The output is counted to focus on the hash + // match/equality path without buffering the joined rows. + HashJoinQuery { + sql: r###"SELECT count(*) + FROM ( + SELECT 'single_hot_bucket_string_join_key' as k + FROM supplier + WHERE s_suppkey <= 3000 + ) s + JOIN ( + SELECT 'single_hot_bucket_string_join_key' as k + FROM lineitem + WHERE l_orderkey % 3000 = 0 + ) l ON s.k = l.k"###, + density: 1.0, + prob_hit: 1.0, + build_size: "3K_(single_hot_bucket)", + probe_size: "20K_long_keys_count", + isolate_partitioned_join: true, + }, + // Q25: skewed high-fanout multi-column string-key inner join. + // This tracks the same candidate-pair filtering path for composite join + // keys, where the first key is skewed and the second long string key must + // also be checked before emitting each match. + HashJoinQuery { + sql: r###"SELECT count(*) + FROM ( + SELECT CAST((s_suppkey % 256) + 1 AS INT) as k1, + 'multi_column_high_fanout_key' as k2 + FROM supplier + WHERE s_suppkey <= 20000 + ) s + JOIN ( + SELECT CAST(1 AS INT) as k1, + 'multi_column_high_fanout_key' as k2 + FROM lineitem + WHERE l_orderkey % 250 = 0 + ) l ON s.k1 = l.k1 AND s.k2 = l.k2"###, + density: 1.0, + prob_hit: 1.0, + build_size: "20K_(fanout~78_multi_key)", + probe_size: "240K_multi_key_count", + isolate_partitioned_join: true, }, ]; @@ -323,7 +538,9 @@ impl RunOpt { None => 1..=HASH_QUERIES.len(), }; - let config = self.common.config()?; + let mut config = self.common.config()?; + // Disable join reordering to ensure the optimizer doesn't swap join sides + config.options_mut().optimizer.join_reordering = false; let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); @@ -347,6 +564,7 @@ impl RunOpt { } let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query_index = query_id - 1; @@ -362,6 +580,20 @@ impl RunOpt { ); benchmark_run.start_new_case(&case_name); + // For Q23 force Partitioned mode: zero the CollectLeft thresholds + // so the planner cannot prove the build side is small (as happens + // when the datasource provides no row-count stats). + if query.isolate_partitioned_join { + ctx.sql( + "SET datafusion.optimizer.hash_join_single_partition_threshold = 0", + ) + .await?; + ctx.sql( + "SET datafusion.optimizer.hash_join_single_partition_threshold_rows = 0", + ) + .await?; + } + let query_run = self .benchmark_query(query.sql, &query_id.to_string(), &ctx) .await; diff --git a/benchmarks/src/imdb/convert.rs b/benchmarks/src/imdb/convert.rs index aaed186da4905..bd6b37b2a2b1c 100644 --- a/benchmarks/src/imdb/convert.rs +++ b/benchmarks/src/imdb/convert.rs @@ -82,7 +82,7 @@ impl ConvertOpt { println!( "Converting '{}' to {} files in directory '{}'", - &input_path, self.file_format, &output_path + input_path, self.file_format, output_path ); match self.file_format.as_str() { "csv" => { diff --git a/benchmarks/src/imdb/run.rs b/benchmarks/src/imdb/run.rs index 6d3b5c6bafb40..5822bbcb0d89e 100644 --- a/benchmarks/src/imdb/run.rs +++ b/benchmarks/src/imdb/run.rs @@ -295,7 +295,7 @@ impl RunOpt { let mut benchmark_run = BenchmarkRun::new(); for query_id in query_range { benchmark_run.start_new_case(&format!("Query {query_id}")); - let query_run = self.benchmark_query(query_id).await?; + let query_run = self.benchmark_query(query_id, &mut benchmark_run).await?; for iter in query_run { benchmark_run.write_iter(iter.elapsed, iter.row_count); } @@ -304,7 +304,13 @@ impl RunOpt { Ok(()) } - async fn benchmark_query(&self, query_id: usize) -> Result> { + /// `benchmark_run` is handed this query's runtime, which is built here so + /// each query gets a pool of its own. + async fn benchmark_query( + &self, + query_id: usize, + benchmark_run: &mut BenchmarkRun, + ) -> Result> { let mut config = self .common .config()? @@ -314,6 +320,7 @@ impl RunOpt { self.hash_join_buffering_capacity; let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); // register tables self.register_tables(&ctx).await?; @@ -348,14 +355,14 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } async fn register_tables(&self, ctx: &SessionContext) -> Result<()> { for table in IMDB_TABLES { - let table_provider = { self.get_table(ctx, table).await? }; + let table_provider = { self.get_table(ctx, table)? }; if self.mem_table { println!("Loading table '{table}' into memory"); @@ -416,7 +423,7 @@ impl RunOpt { Ok(result) } - async fn get_table( + fn get_table( &self, ctx: &SessionContext, table: &str, @@ -425,7 +432,6 @@ impl RunOpt { let table_format = self.file_format.as_str(); // Obtain a snapshot of the SessionState - let state = ctx.state(); let (format, path, extension): (Arc, String, &'static str) = match table_format { // dbgen creates .tbl ('|' delimited) files without header @@ -458,9 +464,7 @@ impl RunOpt { } }; - let options = ListingOptions::new(format) - .with_file_extension(extension) - .with_collect_stat(state.config().collect_statistics()); + let options = ListingOptions::new(format).with_file_extension(extension); let table_path = ListingTableUrl::parse(path)?; let config = ListingTableConfig::new(table_path).with_listing_options(options); diff --git a/benchmarks/src/lib.rs b/benchmarks/src/lib.rs index f41fd5ebed205..7d8b7044bbdd8 100644 --- a/benchmarks/src/lib.rs +++ b/benchmarks/src/lib.rs @@ -16,7 +16,6 @@ // under the License. //! DataFusion benchmark runner -pub mod benchmark_runner; pub mod cancellation; pub mod clickbench; pub mod dict; @@ -28,6 +27,8 @@ pub mod smj; pub mod sort_pushdown; pub mod sort_tpch; pub mod sql_benchmark; +pub mod sql_benchmark_runner; +pub mod sql_benchmark_suite; pub mod tpcds; pub mod tpch; pub mod util; diff --git a/benchmarks/src/nlj.rs b/benchmarks/src/nlj.rs index 361cc35ec200c..485ee069d1bba 100644 --- a/benchmarks/src/nlj.rs +++ b/benchmarks/src/nlj.rs @@ -211,6 +211,7 @@ impl RunOpt { let ctx = SessionContext::new_with_config_rt(config, rt); let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query_index = query_id - 1; // Convert 1-based to 0-based index diff --git a/benchmarks/src/smj.rs b/benchmarks/src/smj.rs index 3d173b7116e2b..9282f72c2fab6 100644 --- a/benchmarks/src/smj.rs +++ b/benchmarks/src/smj.rs @@ -550,6 +550,7 @@ impl RunOpt { let ctx = SessionContext::new_with_config_rt(config, rt); let mut benchmark_run = BenchmarkRun::new(); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); for query_id in query_range { let query_index = query_id - 1; // Convert 1-based to 0-based index diff --git a/benchmarks/src/sort_pushdown.rs b/benchmarks/src/sort_pushdown.rs index 8e34706ac140a..77f889e702e3d 100644 --- a/benchmarks/src/sort_pushdown.rs +++ b/benchmarks/src/sort_pushdown.rs @@ -137,7 +137,7 @@ impl RunOpt { for query_id in query_ids { benchmark_run.start_new_case(&format!("{query_id}")); - let query_results = self.benchmark_query(query_id).await; + let query_results = self.benchmark_query(query_id, &mut benchmark_run).await; match query_results { Ok(query_results) => { for iter in query_results { @@ -156,17 +156,25 @@ impl RunOpt { Ok(()) } - async fn benchmark_query(&self, query_id: usize) -> Result> { + /// `benchmark_run` is handed this query's runtime, which is built here so + /// each query gets a pool of its own. + async fn benchmark_query( + &self, + query_id: usize, + benchmark_run: &mut BenchmarkRun, + ) -> Result> { let sql = self.load_query(query_id)?; let config = self.common.config()?; let rt = self.common.build_runtime()?; let state = SessionStateBuilder::new() - .with_config(config) + // Always collect statistics for sort pushdown + .with_config(config.with_collect_statistics(true)) .with_runtime_env(rt) .with_default_features() .build(); let ctx = SessionContext::from(state); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); self.register_tables(&ctx).await?; @@ -190,7 +198,7 @@ impl RunOpt { let avg = millis.iter().sum::() / millis.len() as f64; println!("Query {query_id} avg time: {avg:.2} ms"); - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } @@ -255,9 +263,7 @@ impl RunOpt { ); let extension = DEFAULT_PARQUET_EXTENSION; - let options = ListingOptions::new(format) - .with_file_extension(extension) - .with_collect_stat(true); // Always collect statistics for sort pushdown + let options = ListingOptions::new(format).with_file_extension(extension); let table_path = ListingTableUrl::parse(path)?; let schema = options.infer_schema(&state, &table_path).await?; diff --git a/benchmarks/src/sort_tpch.rs b/benchmarks/src/sort_tpch.rs index 206911c45adde..d5f81c04a3ba4 100644 --- a/benchmarks/src/sort_tpch.rs +++ b/benchmarks/src/sort_tpch.rs @@ -64,8 +64,8 @@ pub struct RunOpt { #[arg(short = 'm', long = "mem-table")] mem_table: bool, - /// Mark the first column of each table as sorted in ascending order. - /// The tables should have been created with the `--sort` option for this to have any effect. + /// Declare that the first column of the input table is already sorted in ascending order. + /// This flag only attaches ordering metadata; it does not sort the input files. #[arg(short = 't', long = "sorted")] sorted: bool, @@ -187,7 +187,7 @@ impl RunOpt { for query_id in query_range { benchmark_run.start_new_case(&format!("{query_id}")); - let query_results = self.benchmark_query(query_id).await; + let query_results = self.benchmark_query(query_id, &mut benchmark_run).await; match query_results { Ok(query_results) => { for iter in query_results { @@ -207,7 +207,14 @@ impl RunOpt { } /// Benchmark query `query_id` in `SORT_QUERIES` - async fn benchmark_query(&self, query_id: usize) -> Result> { + /// + /// `benchmark_run` is handed this query's runtime, which is built here so + /// each query gets a pool of its own. + async fn benchmark_query( + &self, + query_id: usize, + benchmark_run: &mut BenchmarkRun, + ) -> Result> { let config = self.common.config()?; let rt = self.common.build_runtime()?; let state = SessionStateBuilder::new() @@ -216,6 +223,7 @@ impl RunOpt { .with_default_features() .build(); let ctx = SessionContext::from(state); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); // register tables self.register_tables(&ctx).await?; @@ -250,7 +258,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory usage stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } @@ -333,9 +341,7 @@ impl RunOpt { ); let extension = DEFAULT_PARQUET_EXTENSION; - let options = ListingOptions::new(format) - .with_file_extension(extension) - .with_collect_stat(state.config().collect_statistics()); + let options = ListingOptions::new(format).with_file_extension(extension); let table_path = ListingTableUrl::parse(path)?; let schema = options.infer_schema(&state, &table_path).await?; diff --git a/benchmarks/src/sql_benchmark.rs b/benchmarks/src/sql_benchmark.rs index fc6da24b8a9b2..f69012402a3c2 100644 --- a/benchmarks/src/sql_benchmark.rs +++ b/benchmarks/src/sql_benchmark.rs @@ -74,6 +74,20 @@ impl SqlBenchmark { ctx: &SessionContext, full_path: impl AsRef, benchmark_directory: impl AsRef, + ) -> Result { + Self::new_with_replacements(ctx, full_path, benchmark_directory, HashMap::new()) + .await + } + + /// Creates a benchmark using caller-provided template replacements. + /// + /// Caller values take precedence over environment variables during + /// `${...}` substitution; `BENCHMARK_DIR` is still set internally. + pub async fn new_with_replacements( + ctx: &SessionContext, + full_path: impl AsRef, + benchmark_directory: impl AsRef, + replacement_mapping: HashMap, ) -> Result { let full_path = full_path.as_ref(); let benchmark_directory = benchmark_directory.as_ref(); @@ -83,7 +97,7 @@ impl SqlBenchmark { group: group_name, subgroup: String::new(), benchmark_path: full_path.to_path_buf(), - replacement_mapping: HashMap::new(), + replacement_mapping, expect: vec![], queries: HashMap::new(), result_queries: vec![], @@ -201,7 +215,11 @@ impl SqlBenchmark { /// # Errors /// Returns an error if a `run` query fails or if expected plan strings /// are not found. - pub async fn run(&mut self, ctx: &SessionContext, save_results: bool) -> Result<()> { + pub async fn run( + &mut self, + ctx: &SessionContext, + save_results: bool, + ) -> Result { let run_queries = self .queries .get(&QueryDirective::Run) @@ -270,7 +288,7 @@ impl SqlBenchmark { // Store results for verification self.last_results = Some(result); - Ok(()) + Ok(result_count) } /// Calls run and persists results to disk as a CSV file. @@ -283,7 +301,7 @@ impl SqlBenchmark { /// Returns an error if no results are available or if writing to the /// target path fails. pub async fn persist(&mut self, ctx: &SessionContext) -> Result<()> { - self.run(ctx, true).await?; + let _ = self.run(ctx, true).await?; // Check if we have result queries to persist for if self.result_queries.is_empty() { @@ -837,10 +855,14 @@ impl BenchmarkDirective { )); } - debug!("Processing {} file: {}", splits[0], splits[1]); + let query_path = resolve_benchmark_file_path(splits[1]); + debug!("Processing {} file: {}", splits[0], query_path.display()); - let query_file = fs::read_to_string(splits[1]).map_err(|e| { - exec_datafusion_err!("Failed to read query file {}: {e}", splits[1]) + let query_file = fs::read_to_string(&query_path).map_err(|e| { + exec_datafusion_err!( + "Failed to read query file {}: {e}", + query_path.display() + ) })?; let query_file = query_file.replace("\r\n", "\n"); @@ -1121,7 +1143,8 @@ impl BenchmarkDirective { } // restart the load from the template file - Box::pin(bench.process_file(ctx, Path::new(splits[1]))).await + let path = resolve_benchmark_file_path(splits[1]); + Box::pin(bench.process_file(ctx, &path)).await } async fn process_include( @@ -1137,7 +1160,8 @@ impl BenchmarkDirective { )); } - Box::pin(bench.process_file(ctx, Path::new(splits[1]))).await + let path = resolve_benchmark_file_path(splits[1]); + Box::pin(bench.process_file(ctx, &path)).await } fn process_echo( @@ -1552,6 +1576,15 @@ fn make_array_formatter<'a>( } } +fn resolve_benchmark_file_path(path: &str) -> PathBuf { + let path = PathBuf::from(path); + if path.is_absolute() || path.exists() { + path + } else { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join(path) + } +} + #[cfg(test)] mod tests { use super::*; @@ -1635,6 +1668,42 @@ mod tests { replacements } + #[test] + fn resolves_sql_benchmarks_paths_from_manifest_directory() { + let path = + resolve_benchmark_file_path("sql_benchmarks/clickbench/init/set_config.sql"); + + assert!(path.exists(), "resolved path should exist: {path:?}"); + } + + #[tokio::test] + async fn run_returns_row_count_when_not_saving_results() { + let contents = "name Q01\n\nrun\nSELECT * FROM (VALUES (1), (2), (3)) AS t(v)\n"; + let mut benchmark = parse_benchmark(contents).await.unwrap(); + let ctx = SessionContext::new(); + + benchmark.initialize(&ctx).await.unwrap(); + let row_count = benchmark.run(&ctx, false).await.unwrap(); + + assert_eq!(row_count, 3); + } + + #[tokio::test] + async fn run_returns_row_count_when_saving_results() { + let contents = "name Q01\n\nrun\nSELECT * FROM (VALUES (1), (2)) AS t(v)\n"; + let mut benchmark = parse_benchmark(contents).await.unwrap(); + let ctx = SessionContext::new(); + + benchmark.initialize(&ctx).await.unwrap(); + let row_count = benchmark.run(&ctx, true).await.unwrap(); + + assert_eq!(row_count, 2); + assert_eq!( + formatted_last_results(&benchmark), + vec![vec!["1"], vec!["2"]] + ); + } + fn env_map(entries: &[(&str, &str)]) -> HashMap { entries .iter() diff --git a/benchmarks/src/sql_benchmark_runner.rs b/benchmarks/src/sql_benchmark_runner.rs new file mode 100644 index 0000000000000..b321881fbf364 --- /dev/null +++ b/benchmarks/src/sql_benchmark_runner.rs @@ -0,0 +1,769 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared SQL benchmark runner used by `benchmark_runner` and the Criterion +//! SQL benchmark harness. + +use crate::sql_benchmark::SqlBenchmark; +use crate::util::{CommonOpt, print_memory_stats}; +use criterion::{Criterion, SamplingMode}; +use datafusion::error::Result; +use datafusion::prelude::SessionContext; +use datafusion_common::{DataFusionError, exec_datafusion_err}; +use std::any::Any; +use std::collections::{BTreeMap, HashMap}; +use std::fs; +use std::panic::{AssertUnwindSafe, catch_unwind}; +use std::path::{Path, PathBuf}; +use tokio::runtime::Runtime; + +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct BenchmarkFilter { + pub name: Option, + pub subgroup: Option, + pub query: Option, +} + +#[derive(Debug, Clone)] +pub struct SqlRunConfig { + pub common: CommonOpt, + pub filter: BenchmarkFilter, + pub replacements: HashMap, + pub query_filename: Option, + pub persist_results: bool, + pub validate_results: bool, + pub output: Option, +} + +/// Runs the selected SQL benchmarks through a caller-provided Criterion instance. +pub fn run_criterion_benchmarks_impl( + benchmark_dir: &Path, + config: &SqlRunConfig, + criterion: &mut Criterion, +) -> Result<()> { + let rt = make_tokio_runtime()?; + let listing_ctx = make_ctx(&config.common)?; + let all_benchmarks = rt.block_on(load_benchmark_definitions_for_query( + &config.filter, + &listing_ctx, + benchmark_dir, + &config.replacements, + config.query_filename.as_deref(), + ))?; + let selected = filter_benchmarks(&config.filter, all_benchmarks.clone()); + + ensure_selection(&config.filter, &all_benchmarks, &selected)?; + + for (group_name, benchmarks) in selected { + let mut group = criterion.benchmark_group(group_name); + + group.sample_size(10); + group.sampling_mode(SamplingMode::Flat); + + for mut benchmark in benchmarks { + let ctx = make_ctx(&config.common)?; + let result = + run_criterion_benchmark(&rt, &ctx, &mut benchmark, config, &mut group); + let cleanup_result = rt.block_on(benchmark.cleanup(&ctx)); + + finish_benchmark(result, cleanup_result)?; + } + + group.finish(); + } + + Ok(()) +} + +/// Runs one benchmark case inside Criterion and converts benchmark panics to errors. +fn run_criterion_benchmark( + rt: &Runtime, + ctx: &SessionContext, + benchmark: &mut SqlBenchmark, + config: &SqlRunConfig, + group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>, +) -> Result<()> { + rt.block_on(prepare_benchmark(ctx, benchmark, config))?; + + let name = criterion_function_name(benchmark); + let result = catch_unwind(AssertUnwindSafe(|| { + group.bench_function(name.clone(), |b| { + b.iter(|| { + let _ = rt.block_on(async { + benchmark.run(ctx, false).await.unwrap_or_else(|err| { + panic!("Failed to run benchmark {name}: {err:?}") + }) + }); + }); + }); + })); + + match result { + Ok(()) => { + print_memory_stats(&*ctx.runtime_env().memory_pool); + Ok(()) + } + Err(payload) => Err(panic_payload_to_error(payload.as_ref())), + } +} + +/// Extracts a readable message from a panic payload. +fn panic_payload_to_error(payload: &(dyn Any + Send)) -> DataFusionError { + let message = if let Some(message) = payload.downcast_ref::() { + message.as_str() + } else if let Some(message) = payload.downcast_ref::<&str>() { + message + } else { + "unknown panic" + }; + + exec_datafusion_err!("criterion benchmark failed: {message}") +} + +pub fn default_sql_benchmark_directory() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("sql_benchmarks") +} + +/// Replacements used by the Criterion SQL benchmark harness. +pub fn default_criterion_replacements() -> HashMap { + criterion_replacements(std::env::var("DATA_DIR").ok()) +} + +fn criterion_replacements(data_dir: Option) -> HashMap { + HashMap::from([( + "data_dir".to_string(), + data_dir.unwrap_or_else(|| { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("data") + .to_string_lossy() + .into_owned() + }), + )]) +} + +fn make_tokio_runtime() -> Result { + tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build() + .map_err(|e| DataFusionError::External(Box::new(e))) +} + +pub fn make_ctx(common: &CommonOpt) -> Result { + let config = common.config()?; + let rt = common.build_runtime()?; + + Ok(SessionContext::new_with_config_rt(config, rt)) +} + +/// Discovers benchmark definition files in stable path order. +fn discover_benchmark_paths(path: &Path) -> Result> { + let mut paths = Vec::new(); + + collect_benchmark_paths(path, &mut paths)?; + paths.sort(); + + Ok(paths) +} + +/// Loads all benchmark definitions with replacements derived from the filter. +pub async fn load_benchmark_definitions( + filter: &BenchmarkFilter, + ctx: &SessionContext, + benchmark_dir: &Path, + replacements: &HashMap, +) -> Result>> { + load_benchmark_definitions_for_query(filter, ctx, benchmark_dir, replacements, None) + .await +} + +/// Loads benchmark definitions, optionally limiting discovery to one filename. +pub async fn load_benchmark_definitions_for_query( + filter: &BenchmarkFilter, + ctx: &SessionContext, + benchmark_dir: &Path, + replacements: &HashMap, + query_filename: Option<&str>, +) -> Result>> { + let mut benches = BTreeMap::new(); + let mut replacements = replacements.clone(); + let selected_suite_dir = filter + .name + .as_ref() + .map(|name| benchmark_dir.join(name.to_ascii_lowercase())) + .filter(|path| path.is_dir()); + let discovery_dir = selected_suite_dir.as_deref().unwrap_or(benchmark_dir); + if let Some(subgroup) = &filter.subgroup { + replacements.insert("bench_subgroup".to_string(), subgroup.to_string()); + } + + for path in discover_benchmark_paths(discovery_dir)? + .into_iter() + .filter(|path| { + query_filename.is_none_or(|filename| { + path.file_name() + .is_some_and(|candidate| candidate.eq_ignore_ascii_case(filename)) + }) + }) + { + let benchmark = SqlBenchmark::new_with_replacements( + ctx, + &path, + benchmark_dir, + replacements.clone(), + ) + .await?; + benches + .entry(benchmark.group().to_string()) + .or_insert_with(Vec::new) + .push(benchmark); + } + + sort_benchmarks(&mut benches); + + Ok(benches) +} + +pub fn sort_benchmarks(benchmarks: &mut BTreeMap>) { + benchmarks + .values_mut() + .for_each(|benchmarks| benchmarks.sort_by(|a, b| a.name().cmp(b.name()))); +} + +/// Applies benchmark, subgroup, and query filters to discovered benchmark groups. +pub fn filter_benchmarks( + filter: &BenchmarkFilter, + benchmarks: BTreeMap>, +) -> BTreeMap> { + match &filter.name { + Some(bench_name) => benchmarks + .into_iter() + .filter(|(key, _)| key.eq_ignore_ascii_case(bench_name)) + .map(|(key, mut value)| { + if let Some(subgroup) = &filter.subgroup { + value.retain(|bench| bench.subgroup().eq_ignore_ascii_case(subgroup)); + } + if let Some(query) = &filter.query { + retain_query_matches(&mut value, query); + } + (key, value) + }) + .filter(|(_, value)| !value.is_empty()) + .collect(), + None => benchmarks, + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +enum QueryMatchRank { + Exact, + StartsWith, + TokenStartsWith, + Contains, +} + +/// Retains the best benchmark name matches for a query selector like `1` or `Q01`. +/// +/// Exact matches keep all matching benchmarks; fallback matches keep one stable +/// best match to avoid running adjacent query variants unexpectedly. +fn retain_query_matches(benchmarks: &mut Vec, query: &str) { + let normalized = normalize_query(query); + let best_rank = benchmarks + .iter() + .filter_map(|bench| query_match_rank(bench.name(), &normalized)) + .min(); + let Some(best_rank) = best_rank else { + benchmarks.clear(); + return; + }; + + // if exact match retain all matches + if best_rank == QueryMatchRank::Exact { + benchmarks.retain(|bench| { + query_match_rank(bench.name(), &normalized) == Some(QueryMatchRank::Exact) + }); + return; + } + + let selected = benchmarks + .iter() + .filter(|bench| query_match_rank(bench.name(), &normalized) == Some(best_rank)) + .min_by(|left, right| { + left.name() + .cmp(right.name()) + .then_with(|| left.subgroup().cmp(right.subgroup())) + }) + .cloned(); + + benchmarks.clear(); + + if let Some(benchmark) = selected { + benchmarks.push(benchmark); + } +} + +/// Ranks query-name matches, preferring direct `Q01...` names before fallback +/// matches inside descriptive names such as `costsel_q01...`. +fn query_match_rank(name: &str, normalized_query: &str) -> Option { + let name = name.to_ascii_uppercase(); + let normalized_query = normalized_query.to_ascii_uppercase(); + + if name == normalized_query { + Some(QueryMatchRank::Exact) + } else if name.starts_with(&normalized_query) { + Some(QueryMatchRank::StartsWith) + } else if name + .split(|c: char| !c.is_ascii_alphanumeric()) + .any(|token| token.starts_with(&normalized_query)) + { + Some(QueryMatchRank::TokenStartsWith) + } else if name.contains(&normalized_query) { + Some(QueryMatchRank::Contains) + } else { + None + } +} + +/// Converts user query selectors into the SQL benchmark `QNN` naming form. +fn normalize_query(query: &str) -> String { + let query = query.trim_start_matches(['Q', 'q']); + let split = query + .find(|c: char| !c.is_ascii_digit()) + .unwrap_or(query.len()); + let (number, suffix) = query.split_at(split); + + format!("Q{number:0>2}{suffix}") +} + +pub fn format_benchmark_list(benchmarks: &BTreeMap>) -> String { + let mut output = String::from("SQL benchmarks:\n"); + + for (name, benchmarks) in benchmarks { + let query_word = if benchmarks.len() == 1 { + "query" + } else { + "queries" + }; + output.push_str(&format!(" {name:<24} {} {query_word}\n", benchmarks.len())); + } + + output.trim_end().to_string() +} + +/// Recursively collects `.benchmark` files below `path`. +fn collect_benchmark_paths(path: &Path, paths: &mut Vec) -> Result<()> { + let mut entries = fs::read_dir(path)? + .filter_map(std::result::Result::ok) + .collect::>(); + + entries.sort_by_key(|entry| entry.path()); + + for entry in entries { + let path = entry.path(); + if path.is_dir() { + collect_benchmark_paths(&path, paths)?; + } else if path.extension().is_some_and(|ext| ext == "benchmark") { + paths.push(path); + } + } + + Ok(()) +} + +pub fn unknown_benchmark_error( + requested: &str, + benchmarks: &BTreeMap>, +) -> DataFusionError { + exec_datafusion_err!( + "unknown benchmark '{requested}'\n\n{}", + format_benchmark_list(benchmarks) + ) +} + +fn unknown_subgroup_error( + benchmark_name: &str, + subgroup: &str, + benchmarks: &[SqlBenchmark], +) -> DataFusionError { + exec_datafusion_err!( + "no SQL benchmark subgroup matched benchmark '{benchmark_name}' with subgroup '{subgroup}'\n\n{}", + format_subgroup_list(benchmark_name, benchmarks) + ) +} + +fn unknown_query_error( + benchmark_name: &str, + query: &str, + subgroup: Option<&str>, + benchmarks: &[SqlBenchmark], +) -> DataFusionError { + let normalized = normalize_query(query); + + exec_datafusion_err!( + "no SQL benchmark query matched benchmark '{benchmark_name}' with query '{query}' (normalized: '{normalized}')\n\n{}", + format_query_list(benchmark_name, subgroup, benchmarks) + ) +} + +fn format_subgroup_list(benchmark_name: &str, benchmarks: &[SqlBenchmark]) -> String { + let mut entries = benchmarks + .iter() + .map(|bench| { + if bench.subgroup().is_empty() { + "".to_string() + } else { + bench.subgroup().to_string() + } + }) + .collect::>(); + + entries.sort(); + entries.dedup(); + + let mut output = format!("Available {benchmark_name} subgroups:\n"); + + if entries.is_empty() { + output.push_str(" "); + } else { + for entry in entries { + output.push_str(&format!(" {entry}\n")); + } + } + + output.trim_end().to_string() +} + +/// Formats available query names for an unknown-query error message. +fn format_query_list( + benchmark_name: &str, + subgroup: Option<&str>, + benchmarks: &[SqlBenchmark], +) -> String { + let mut entries = benchmarks + .iter() + .filter(|bench| { + subgroup + .is_none_or(|subgroup| bench.subgroup().eq_ignore_ascii_case(subgroup)) + }) + .map(|bench| { + if bench.subgroup().is_empty() { + bench.name().to_string() + } else { + format!("{}/{} ", bench.subgroup(), bench.name()) + } + }) + .take(10) + .collect::>(); + + entries.sort(); + entries.dedup(); + if entries.len() == 10 { + entries.push("...".to_string()); + } + + let mut output = match subgroup { + Some(subgroup) => { + format!("Available {benchmark_name} queries in subgroup '{subgroup}':\n") + } + None => format!("Available {benchmark_name} queries:\n"), + }; + + if entries.is_empty() { + output.push_str(" "); + } else { + for entry in entries { + output.push_str(&format!(" {entry}\n")); + } + } + + output.trim_end().to_string() +} + +/// Initializes a benchmark and performs any configured assertion or validation step. +pub async fn prepare_benchmark( + ctx: &SessionContext, + benchmark: &mut SqlBenchmark, + config: &SqlRunConfig, +) -> Result<()> { + benchmark.initialize(ctx).await?; + benchmark.assert(ctx).await?; + + if config.persist_results { + benchmark.persist(ctx).await?; + } else if config.validate_results { + let _ = benchmark.run(ctx, true).await?; + benchmark.verify(ctx).await?; + } + + Ok(()) +} + +/// Ensures filtering selected at least one benchmark and emits targeted errors. +pub fn ensure_selection( + filter: &BenchmarkFilter, + all_benchmarks: &BTreeMap>, + selected: &BTreeMap>, +) -> Result<()> { + if selected.is_empty() { + if let Some(name) = &filter.name { + if let Some((benchmark_name, benchmarks)) = all_benchmarks + .iter() + .find(|(key, _)| key.eq_ignore_ascii_case(name)) + { + if let Some(subgroup) = &filter.subgroup { + let has_subgroup = benchmarks + .iter() + .any(|bench| bench.subgroup().eq_ignore_ascii_case(subgroup)); + + if !has_subgroup { + return Err(unknown_subgroup_error( + benchmark_name, + subgroup, + benchmarks, + )); + } + } + + if let Some(query) = &filter.query { + return Err(unknown_query_error( + benchmark_name, + query, + filter.subgroup.as_deref(), + benchmarks, + )); + } + } + return Err(unknown_benchmark_error(name, all_benchmarks)); + } + return Err(exec_datafusion_err!("no SQL benchmarks discovered")); + } + + Ok(()) +} + +/// Combines benchmark and cleanup results without hiding cleanup failures. +pub fn finish_benchmark(result: Result<()>, cleanup_result: Result<()>) -> Result<()> { + match (result, cleanup_result) { + (Ok(()), Ok(())) => Ok(()), + (Ok(()), Err(cleanup_error)) => Err(cleanup_error), + (Err(error), Ok(())) => Err(error), + (Err(error), Err(cleanup_error)) => Err(exec_datafusion_err!( + "{error}; cleanup also failed: {cleanup_error}" + )), + } +} + +fn criterion_function_name(benchmark: &SqlBenchmark) -> String { + let mut name = benchmark.name().to_string(); + + if !benchmark.subgroup().is_empty() { + name.push('_'); + name.push_str(benchmark.subgroup()); + } + + name +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion::prelude::SessionContext; + use std::path::{Path, PathBuf}; + + fn write_benchmark(root: &Path, relative_path: &str, contents: &str) -> PathBuf { + let path = root.join(relative_path); + + fs::create_dir_all(path.parent().unwrap()).unwrap(); + fs::write(&path, contents).unwrap(); + + path + } + + #[tokio::test] + async fn caller_replacements_reach_parser() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q01.benchmark", + "name Q01\n\nload\nSELECT '${ALPHA_FORMAT}'\n\nrun\nSELECT 1\n", + ); + let replacements = + HashMap::from([("alpha_format".to_string(), "csv".to_string())]); + + let result = load_benchmark_definitions( + &BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("1".to_string()), + }, + &SessionContext::new(), + temp.path(), + &replacements, + ) + .await; + + assert!(result.is_ok()); + } + + #[tokio::test] + async fn query_filename_filters_paths_before_parsing() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "alpha/benchmarks/q07.benchmark", + "name Q07\n\nrun\nSELECT 7\n", + ); + write_benchmark( + temp.path(), + "alpha/benchmarks/q08.benchmark", + "this is not a benchmark definition", + ); + write_benchmark( + temp.path(), + "beta/benchmarks/q07.benchmark", + "this is not a benchmark definition", + ); + + let benches = load_benchmark_definitions_for_query( + &BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("7".to_string()), + }, + &SessionContext::new(), + temp.path(), + &HashMap::new(), + Some("q07.benchmark"), + ) + .await + .unwrap(); + + assert_eq!(benches["alpha"].len(), 1); + assert_eq!(benches["alpha"][0].name(), "Q07"); + } + + #[test] + fn criterion_replacements_use_benchmarks_data_directory() { + let expected = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("data") + .to_string_lossy() + .into_owned(); + + assert_eq!(criterion_replacements(None)["data_dir"], expected); + } + + #[test] + fn criterion_replacements_use_explicit_data_directory() { + let replacements = criterion_replacements(Some("/custom/data".to_string())); + + assert_eq!(replacements["data_dir"], "/custom/data"); + } + + #[tokio::test] + async fn query_filename_keeps_matches_in_multiple_subgroups() { + let temp = tempfile::tempdir().unwrap(); + for subgroup in ["aggregate", "window"] { + write_benchmark( + temp.path(), + &format!("alpha/benchmarks/{subgroup}/q03.benchmark"), + &format!("name Q03\nsubgroup {subgroup}\n\nrun\nSELECT 3\n"), + ); + } + + let filter = BenchmarkFilter { + name: Some("alpha".to_string()), + subgroup: None, + query: Some("3".to_string()), + }; + let benches = load_benchmark_definitions_for_query( + &filter, + &SessionContext::new(), + temp.path(), + &HashMap::new(), + Some("q03.benchmark"), + ) + .await + .unwrap(); + assert_eq!(filter_benchmarks(&filter, benches)["alpha"].len(), 2); + + let filter = BenchmarkFilter { + subgroup: Some("window".to_string()), + ..filter + }; + let benches = load_benchmark_definitions_for_query( + &filter, + &SessionContext::new(), + temp.path(), + &HashMap::new(), + Some("q03.benchmark"), + ) + .await + .unwrap(); + assert_eq!(filter_benchmarks(&filter, benches)["alpha"].len(), 1); + } + + #[tokio::test] + async fn query_filename_accepts_alphanumeric_pattern() { + let temp = tempfile::tempdir().unwrap(); + write_benchmark( + temp.path(), + "imdb/benchmarks/01a.benchmark", + "name Q01a\n\nrun\nSELECT 1\n", + ); + + let benches = load_benchmark_definitions_for_query( + &BenchmarkFilter { + name: Some("imdb".to_string()), + subgroup: None, + query: Some("1a".to_string()), + }, + &SessionContext::new(), + temp.path(), + &HashMap::new(), + Some("01a.benchmark"), + ) + .await + .unwrap(); + + assert_eq!(benches["imdb"][0].name(), "Q01a"); + } + + #[test] + fn normalizes_query_like_existing_sql_harness() { + assert_eq!(normalize_query("1"), "Q01"); + assert_eq!(normalize_query("01"), "Q01"); + assert_eq!(normalize_query("6a"), "Q06a"); + assert_eq!(normalize_query("Q06a"), "Q06a"); + } + + #[test] + fn criterion_names_match_existing_sql_harness() { + let temp = tempfile::tempdir().unwrap(); + let benchmark_path = write_benchmark( + temp.path(), + "tpch/benchmarks/q01.benchmark", + "name Q01\nsubgroup sf1\n\nrun\nSELECT 1\n", + ); + let ctx = SessionContext::new(); + let rt = make_tokio_runtime().unwrap(); + let benchmark = rt + .block_on(SqlBenchmark::new(&ctx, &benchmark_path, temp.path())) + .unwrap(); + + assert_eq!(benchmark.group(), "tpch"); + assert_eq!(criterion_function_name(&benchmark), "Q01_sf1"); + } +} diff --git a/benchmarks/src/sql_benchmark_suite.rs b/benchmarks/src/sql_benchmark_suite.rs new file mode 100644 index 0000000000000..aa7a3c5d52c8e --- /dev/null +++ b/benchmarks/src/sql_benchmark_suite.rs @@ -0,0 +1,849 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Metadata parsing and validation for SQL benchmark suites. + +use std::collections::{BTreeMap, BTreeSet}; +use std::fs; +use std::io; +use std::path::{Component, Path, PathBuf}; + +use datafusion_common::{DataFusionError, Result}; +use serde::Deserialize; + +const DEFAULT_QUERY_PATTERN: &str = "q{QUERY_ID_PADDED}.benchmark"; + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawSuite { + description: String, + query_pattern: Option, + #[serde(default)] + path_replacements: BTreeMap, + #[serde(default)] + options: Vec, + #[serde(default)] + examples: Vec, +} + +#[derive(Debug, Deserialize)] +#[serde(deny_unknown_fields)] +struct RawSuiteOption { + name: String, + short: Option, + env: String, + default: String, + values: Option>, + help: String, +} + +/// Validated metadata for one benchmark suite. +#[derive(Debug, Clone)] +pub struct SuiteMetadata { + name: String, + directory: PathBuf, + description: String, + query_pattern: String, + path_replacements: BTreeMap, + options: Vec, + examples: Vec, + benchmark_count: usize, +} + +/// A suite-specific command-line option. +#[derive(Debug, Clone)] +pub struct SuiteOption { + name: String, + short: Option, + env: String, + default: String, + values: Option>, + help: String, +} + +/// An example invocation from a suite metadata file. +#[derive(Debug, Clone, Deserialize)] +#[serde(deny_unknown_fields)] +pub struct SuiteExample { + command: String, + description: String, +} + +/// Global option names unavailable to suite-specific options. +pub struct ReservedOptions<'a> { + pub long: &'a BTreeSet, + pub short: &'a BTreeSet, +} + +/// Where a resolved option value originated. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ValueSource { + CommandLine, + Environment, + Default, +} + +/// An option value together with its origin. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct ResolvedValue { + pub value: String, + pub source: ValueSource, +} + +fn metadata_error(message: impl Into) -> DataFusionError { + DataFusionError::Configuration(message.into()) +} + +impl SuiteMetadata { + /// Loads and validates `//.suite`. + pub fn load(root: &Path, name: &str, reserved: &ReservedOptions) -> Result { + let directory = root.join(name); + let metadata_path = directory.join(format!("{name}.suite")); + let contents = fs::read_to_string(&metadata_path)?; + let raw: RawSuite = toml::from_str(&contents).map_err(|error| { + metadata_error(format!("{}: {error}", metadata_path.display())) + })?; + Self::from_raw(name, directory, raw, reserved) + } + + fn from_raw( + name: &str, + directory: PathBuf, + raw: RawSuite, + reserved: &ReservedOptions, + ) -> Result { + if raw.description.trim().is_empty() { + return Err(metadata_error("suite description must not be empty")); + } + + let query_pattern = raw + .query_pattern + .clone() + .unwrap_or_else(|| DEFAULT_QUERY_PATTERN.to_string()); + + validate_query_pattern(&query_pattern)?; + + let mut long_names = BTreeSet::new(); + let mut short_names = BTreeSet::new(); + let mut env_names = BTreeSet::new(); + let mut options = Vec::with_capacity(raw.options.len()); + + for option in raw.options { + Self::validate_option( + reserved, + &mut long_names, + &mut short_names, + &mut env_names, + &raw.path_replacements, + &option, + )?; + + let suite_option = SuiteOption { + name: option.name, + short: option.short.as_deref().map(parse_short).transpose()?, + env: option.env, + default: option.default, + values: option.values, + help: option.help, + }; + + if !suite_option.accepts(&suite_option.default) { + return Err(metadata_error(format!( + "default value '{}' is not accepted by option '{}'", + suite_option.default, suite_option.name + ))); + } + + options.push(suite_option); + } + + for example in &raw.examples { + if example.command.trim().is_empty() { + return Err(metadata_error("example command must not be empty")); + } + if example.description.trim().is_empty() { + return Err(metadata_error("example description must not be empty")); + } + } + + let path_replacements = raw + .path_replacements + .into_iter() + .map(|(key, path)| { + let path = PathBuf::from(path); + let path = if path.is_relative() { + directory.join(path) + } else { + path + }; + (key, path) + }) + .collect(); + let benchmark_count = count_benchmarks(&directory)?; + + Ok(Self { + name: name.to_string(), + directory, + description: raw.description, + query_pattern, + path_replacements, + options, + examples: raw.examples, + benchmark_count, + }) + } + + fn validate_option( + reserved: &ReservedOptions, + long_names: &mut BTreeSet, + short_names: &mut BTreeSet, + env_names: &mut BTreeSet, + path_replacements: &BTreeMap, + option: &RawSuiteOption, + ) -> Result<()> { + if !valid_long_name(&option.name) { + return Err(metadata_error(format!( + "invalid option name '{}'", + option.name + ))); + } + if reserved.long.contains(&option.name) || !long_names.insert(option.name.clone()) + { + return Err(metadata_error(format!( + "option name '{}' is reserved or duplicated", + option.name + ))); + } + let short = option.short.as_deref().map(parse_short).transpose()?; + if let Some(short) = short + && (reserved.short.contains(&short) || !short_names.insert(short)) + { + return Err(metadata_error(format!( + "option short name '{short}' is reserved or duplicated" + ))); + } + if !env_names.insert(option.env.clone()) { + return Err(metadata_error(format!( + "option environment key '{}' is duplicated", + option.env + ))); + } + if path_replacements.contains_key(&option.env) { + return Err(metadata_error(format!( + "environment key '{}' is used by both an option and a path replacement", + option.env + ))); + } + if option.help.trim().is_empty() { + return Err(metadata_error(format!( + "help for option '{}' must not be empty", + option.name + ))); + } + + Ok(()) + } + + pub fn name(&self) -> &str { + &self.name + } + + pub fn directory(&self) -> &Path { + &self.directory + } + + pub fn description(&self) -> &str { + &self.description + } + + pub fn query_pattern(&self) -> &str { + &self.query_pattern + } + + pub fn path_replacements(&self) -> &BTreeMap { + &self.path_replacements + } + + pub fn options(&self) -> &[SuiteOption] { + &self.options + } + + pub fn examples(&self) -> &[SuiteExample] { + &self.examples + } + + pub fn benchmark_count(&self) -> usize { + self.benchmark_count + } + + /// Formats a query identifier using this suite's query pattern. + pub fn query_filename(&self, query: &str) -> Result { + let query = query.strip_prefix(['q', 'Q']).unwrap_or(query); + let digit_count = query.bytes().take_while(u8::is_ascii_digit).count(); + + if digit_count == 0 || !query.bytes().all(|byte| byte.is_ascii_alphanumeric()) { + return Err(metadata_error(format!( + "invalid query identifier '{query}'" + ))); + } + + let (digits, suffix) = query.split_at(digit_count); + let replacement = if self.query_pattern.contains("{QUERY_ID_PADDED}") { + let digits = digits.trim_start_matches('0'); + let digits = if digits.is_empty() { "0" } else { digits }; + format!("{digits:0>2}{suffix}") + } else { + query.to_string() + }; + + Ok(self + .query_pattern + .replace("{QUERY_ID_PADDED}", &replacement) + .replace("{QUERY_ID}", &replacement)) + } +} + +impl SuiteOption { + pub fn name(&self) -> &str { + &self.name + } + + pub fn short(&self) -> Option { + self.short + } + + pub fn env(&self) -> &str { + &self.env + } + + pub fn default(&self) -> &str { + &self.default + } + + pub fn values(&self) -> Option<&[String]> { + self.values.as_deref() + } + + pub fn help(&self) -> &str { + &self.help + } + + /// Whether `value` belongs to this option's configured value set. + pub fn accepts(&self, value: &str) -> bool { + self.values.as_ref().is_none_or(|values| { + values + .iter() + .any(|allowed| allowed == value || allowed == "...") + }) + } +} + +impl SuiteExample { + pub fn command(&self) -> &str { + &self.command + } + + pub fn description(&self) -> &str { + &self.description + } +} + +/// Finds and loads suite metadata immediately below `root`, sorted by name. +pub fn discover_suites( + root: &Path, + reserved: &ReservedOptions, +) -> Result> { + let mut suites = Vec::new(); + + for entry in collect_sorted_entries(fs::read_dir(root)?)? { + if !entry.file_type()?.is_dir() { + continue; + } + + let name = entry.file_name().to_string_lossy().into_owned(); + let expected = entry.path().join(format!("{name}.suite")); + let suite_files = collect_sorted_entries(fs::read_dir(entry.path())?)? + .into_iter() + .filter(|entry| entry.path().extension().is_some_and(|ext| ext == "suite")) + .collect::>(); + + if suite_files.is_empty() { + continue; + } + if suite_files.len() != 1 || suite_files[0].path() != expected { + return Err(metadata_error(format!( + "suite metadata filename must match directory name '{name}'" + ))); + } + + suites.push(SuiteMetadata::load(root, &name, reserved)?); + } + + suites.sort_by(|left, right| left.name.cmp(&right.name)); + + Ok(suites) +} + +fn valid_long_name(name: &str) -> bool { + name.bytes() + .next() + .is_some_and(|byte| byte.is_ascii_alphanumeric()) + && name.bytes().all(|byte| { + byte.is_ascii_lowercase() || byte.is_ascii_digit() || byte == b'-' + }) +} + +fn parse_short(short: &str) -> Result { + let mut chars = short.chars(); + let value = chars.next().filter(char::is_ascii_alphanumeric); + + match (value, chars.next()) { + (Some(value), None) => Ok(value), + _ => Err(metadata_error(format!( + "invalid option short name '{short}': expected one ASCII alphanumeric character" + ))), + } +} + +fn validate_query_pattern(pattern: &str) -> Result<()> { + let path = Path::new(pattern); + if path.is_absolute() { + return Err(metadata_error("query pattern must not be absolute")); + } + if path + .components() + .any(|component| component == Component::ParentDir) + { + return Err(metadata_error( + "query pattern must not contain a parent component", + )); + } + + let placeholders = pattern.matches("{QUERY_ID}").count() + + pattern.matches("{QUERY_ID_PADDED}").count(); + if placeholders != 1 { + return Err(metadata_error( + "query pattern must contain exactly one query identifier placeholder", + )); + } + + Ok(()) +} + +fn count_benchmarks(directory: &Path) -> Result { + let mut count = 0; + for entry in collect_sorted_entries(fs::read_dir(directory)?)? { + if entry.file_type()?.is_dir() { + count += count_benchmarks(&entry.path())?; + } else if entry + .path() + .extension() + .is_some_and(|ext| ext == "benchmark") + { + count += 1; + } + } + + Ok(count) +} + +fn collect_sorted_entries( + entries: impl IntoIterator>, +) -> io::Result> { + let mut entries = entries.into_iter().collect::>>()?; + entries.sort_by_key(|entry| entry.file_name()); + + Ok(entries) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::sql_benchmark_runner::default_sql_benchmark_directory; + use std::collections::BTreeSet; + use std::fs; + use std::io; + use std::path::Path; + + fn reserved() -> ReservedOptions<'static> { + let long = Box::leak(Box::new(BTreeSet::from([ + "help".to_string(), + "query".to_string(), + ]))); + let short = Box::leak(Box::new(BTreeSet::from(['h', 'q']))); + ReservedOptions { long, short } + } + + fn write_suite(root: &Path, name: &str, metadata: &str) { + let directory = root.join(name); + fs::create_dir_all(&directory).unwrap(); + fs::write(directory.join(format!("{name}.suite")), metadata).unwrap(); + } + + fn minimal(extra: &str) -> String { + format!("description = \"Benchmark\"\n{extra}") + } + + #[test] + fn loads_complete_suite() { + let temp = tempfile::tempdir().unwrap(); + write_suite( + temp.path(), + "alpha", + r#" +description = "Alpha benchmark" +query_pattern = "q{QUERY_ID_PADDED}.benchmark" +[path_replacements] +DATA_DIR = "../../data" +[[options]] +name = "format" +short = "f" +env = "ALPHA_FORMAT" +default = "parquet" +values = ["parquet", "csv"] +help = "Select the file format." +[[examples]] +command = "benchmark_runner alpha -q 1 -f csv" +description = "Run query 1 against CSV." +"#, + ); + + let suite = SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap(); + assert_eq!(suite.name(), "alpha"); + assert_eq!(suite.description(), "Alpha benchmark"); + assert_eq!(suite.options()[0].short(), Some('f')); + assert!(suite.options()[0].accepts("csv")); + assert!(!suite.options()[0].accepts("json")); + assert_eq!( + suite.path_replacements()["DATA_DIR"], + temp.path().join("alpha/../../data") + ); + assert_eq!(suite.examples().len(), 1); + } + + #[test] + fn rejects_unknown_field() { + let temp = tempfile::tempdir().unwrap(); + write_suite( + temp.path(), + "alpha", + "description = \"Alpha\"\ndescripton = \"bad\"\n", + ); + let error = SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap_err(); + assert!(error.to_string().contains("descripton")); + assert!(error.to_string().contains("alpha.suite")); + } + + #[test] + fn validates_value_sets() { + let closed = suite_option(Some(vec!["csv", "parquet"])); + assert!(closed.accepts("csv")); + assert!(!closed.accepts("json")); + assert!(suite_option(Some(vec!["1", "10", "..."])).accepts("100")); + assert!(suite_option(None).accepts("anything")); + } + + fn suite_option(values: Option>) -> SuiteOption { + SuiteOption { + name: "format".to_string(), + short: Some('f'), + env: "FORMAT".to_string(), + default: "csv".to_string(), + values: values.map(|values| values.into_iter().map(str::to_string).collect()), + help: "Format".to_string(), + } + } + + #[test] + fn rejects_invalid_metadata() { + let cases = [ + ("empty description", "description = \" \"\n", "description"), + ( + "invalid long", + &minimal( + "[[options]]\nname = \"Bad_name\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "Bad_name", + ), + ( + "long starts hyphen", + &minimal( + "[[options]]\nname = \"-bad\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "-bad", + ), + ( + "long reserved", + &minimal( + "[[options]]\nname = \"query\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "query", + ), + ( + "short long", + &minimal( + "[[options]]\nname = \"format\"\nshort = \"ff\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "ff", + ), + ( + "short invalid", + &minimal( + "[[options]]\nname = \"format\"\nshort = \"-\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "short", + ), + ( + "short reserved", + &minimal( + "[[options]]\nname = \"format\"\nshort = \"q\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + "q", + ), + ( + "empty help", + &minimal( + "[[options]]\nname = \"format\"\nenv = \"ENV\"\ndefault = \"x\"\nhelp = \" \"\n", + ), + "help", + ), + ( + "bad default", + &minimal( + "[[options]]\nname = \"format\"\nenv = \"ENV\"\ndefault = \"json\"\nvalues = [\"csv\"]\nhelp = \"help\"\n", + ), + "json", + ), + ( + "absolute pattern", + "description = \"Benchmark\"\nquery_pattern = \"/q{QUERY_ID}.benchmark\"\n", + "absolute", + ), + ( + "parent pattern", + "description = \"Benchmark\"\nquery_pattern = \"../q{QUERY_ID}.benchmark\"\n", + "parent", + ), + ( + "no placeholder", + "description = \"Benchmark\"\nquery_pattern = \"q.benchmark\"\n", + "placeholder", + ), + ( + "two placeholders", + "description = \"Benchmark\"\nquery_pattern = \"{QUERY_ID}-{QUERY_ID_PADDED}.benchmark\"\n", + "exactly one", + ), + ( + "empty example command", + &minimal("[[examples]]\ncommand = \" \"\ndescription = \"example\"\n"), + "command", + ), + ( + "empty example description", + &minimal("[[examples]]\ncommand = \"runner\"\ndescription = \" \"\n"), + "description", + ), + ]; + + for (name, metadata, expected) in cases { + let temp = tempfile::tempdir().unwrap(); + write_suite(temp.path(), "alpha", metadata); + let error = + SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap_err(); + assert!(error.to_string().contains(expected), "{name}: {error}"); + } + } + + #[test] + fn rejects_duplicate_and_colliding_options() { + let fields = [("name", "format"), ("short", "f"), ("env", "FORMAT")]; + for (field, value) in fields { + let temp = tempfile::tempdir().unwrap(); + write_suite( + temp.path(), + "alpha", + &minimal(&format!( + r#" +[[options]] +name = "format" +short = "f" +env = "FORMAT" +default = "x" +help = "help" +[[options]] +name = "{name}" +short = "{short}" +env = "{env}" +default = "x" +help = "help" +"#, + name = if field == "name" { value } else { "other" }, + short = if field == "short" { value } else { "o" }, + env = if field == "env" { value } else { "OTHER" } + )), + ); + let error = + SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap_err(); + assert!(error.to_string().contains(value), "{field}: {error}"); + } + + let temp = tempfile::tempdir().unwrap(); + write_suite( + temp.path(), + "alpha", + &minimal( + "[path_replacements]\nFORMAT = \"data\"\n[[options]]\nname = \"format\"\nenv = \"FORMAT\"\ndefault = \"x\"\nhelp = \"help\"\n", + ), + ); + let error = SuiteMetadata::load(temp.path(), "alpha", &reserved()).unwrap_err(); + assert!(error.to_string().contains("FORMAT")); + } + + #[test] + fn discovers_sorted_suites_and_counts_benchmarks() { + let temp = tempfile::tempdir().unwrap(); + write_suite(temp.path(), "zeta", "description = \"Zeta\"\n"); + write_suite(temp.path(), "alpha", "description = \"Alpha\"\n"); + fs::create_dir_all(temp.path().join("alpha/nested")).unwrap(); + fs::write(temp.path().join("alpha/q01.benchmark"), "").unwrap(); + fs::write(temp.path().join("alpha/nested/q02.benchmark"), "").unwrap(); + fs::write(temp.path().join("alpha/ignored.sql"), "").unwrap(); + let suites = discover_suites(temp.path(), &reserved()).unwrap(); + assert_eq!( + suites.iter().map(SuiteMetadata::name).collect::>(), + ["alpha", "zeta"] + ); + assert_eq!(suites[0].benchmark_count(), 2); + } + + #[test] + fn checked_in_suites_cover_benchmark_directories() { + let root = default_sql_benchmark_directory(); + for entry in fs::read_dir(&root).unwrap() { + let entry = entry.unwrap(); + if !entry.file_type().unwrap().is_dir() { + continue; + } + let directory = entry.path(); + let has_benchmark = count_benchmarks(&directory).unwrap() > 0; + if has_benchmark { + let name = entry.file_name().to_string_lossy().into_owned(); + assert!( + directory.join(format!("{name}.suite")).is_file(), + "benchmark directory {name} is missing {name}.suite" + ); + } + } + + let long = BTreeSet::from([ + "batch-size".to_string(), + "debug".to_string(), + "iterations".to_string(), + "output".to_string(), + "partitions".to_string(), + "path".to_string(), + "query".to_string(), + ]); + let short = BTreeSet::from(['q', 'i', 'n', 's', 'd', 'p', 'o']); + let suites = discover_suites( + &root, + &ReservedOptions { + long: &long, + short: &short, + }, + ) + .unwrap(); + let by_name = suites + .iter() + .map(|suite| (suite.name(), suite)) + .collect::>(); + + assert_eq!( + by_name["imdb"].query_filename("1a").unwrap(), + "01a.benchmark" + ); + assert_eq!( + by_name["imdb"].query_filename("01a").unwrap(), + "01a.benchmark" + ); + assert_eq!(by_name["clickbench"].options()[0].name(), "partitioning"); + assert!( + by_name["tpch"] + .options() + .iter() + .all(|option| option.short() != Some('s')) + ); + } + + #[test] + fn rejects_mismatched_suite_filename() { + let temp = tempfile::tempdir().unwrap(); + fs::create_dir_all(temp.path().join("wrong")).unwrap(); + fs::write( + temp.path().join("wrong/other.suite"), + "description = \"Wrong\"", + ) + .unwrap(); + + let error = discover_suites(temp.path(), &reserved()).unwrap_err(); + assert!(error.to_string().contains("wrong")); + } + + #[test] + fn propagates_directory_entry_errors() { + let entries = std::iter::once(Err::(io::Error::new( + io::ErrorKind::PermissionDenied, + "entry denied", + ))); + + let error = collect_sorted_entries(entries).unwrap_err(); + assert_eq!(error.kind(), io::ErrorKind::PermissionDenied); + assert_eq!(error.to_string(), "entry denied"); + } + + #[test] + fn formats_query_filenames() { + let temp = tempfile::tempdir().unwrap(); + write_suite(temp.path(), "padded", "description = \"Padded\"\n"); + write_suite( + temp.path(), + "plain", + "description = \"Plain\"\nquery_pattern = \"{QUERY_ID}.benchmark\"\n", + ); + let padded = SuiteMetadata::load(temp.path(), "padded", &reserved()).unwrap(); + let plain = SuiteMetadata::load(temp.path(), "plain", &reserved()).unwrap(); + + assert_eq!(padded.query_filename("7").unwrap(), "q07.benchmark"); + assert_eq!(padded.query_filename("07").unwrap(), "q07.benchmark"); + assert_eq!( + padded.query_filename("Q1").unwrap(), + padded.query_filename("q1").unwrap() + ); + assert_eq!( + plain.query_filename("Q01a").unwrap(), + plain.query_filename("q01a").unwrap() + ); + assert_eq!( + padded.query_filename("184467440737095516160").unwrap(), + "q184467440737095516160.benchmark" + ); + assert_eq!(plain.query_filename("01a").unwrap(), "01a.benchmark"); + assert!(plain.query_filename("abc").is_err()); + assert!(plain.query_filename("1-a").is_err()); + } +} diff --git a/benchmarks/src/tpcds/run.rs b/benchmarks/src/tpcds/run.rs index 58821340034da..3eaaf172c0f16 100644 --- a/benchmarks/src/tpcds/run.rs +++ b/benchmarks/src/tpcds/run.rs @@ -21,6 +21,7 @@ use std::sync::Arc; use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats}; +use arrow::datatypes::Schema; use arrow::record_batch::RecordBatch; use arrow::util::pretty::{self, pretty_format_batches}; use datafusion::datasource::file_format::parquet::ParquetFormat; @@ -34,7 +35,7 @@ use datafusion::physical_plan::{collect, displayable}; use datafusion::prelude::*; use datafusion_common::instant::Instant; use datafusion_common::utils::get_available_parallelism; -use datafusion_common::{DEFAULT_PARQUET_EXTENSION, plan_err}; +use datafusion_common::{Constraint, Constraints, DEFAULT_PARQUET_EXTENSION, plan_err}; use clap::Args; use log::info; @@ -71,6 +72,61 @@ pub const TPCDS_TABLES: &[&str] = &[ "web_site", ]; +static TPCDS_PRIMARY_KEYS: &[(&str, &[&str])] = &[ + ("call_center", &["cc_call_center_sk"]), + ("catalog_page", &["cp_catalog_page_sk"]), + ("catalog_returns", &["cr_item_sk", "cr_order_number"]), + ("catalog_sales", &["cs_item_sk", "cs_order_number"]), + ("customer", &["c_customer_sk"]), + ("customer_address", &["ca_address_sk"]), + ("customer_demographics", &["cd_demo_sk"]), + ("date_dim", &["d_date_sk"]), + ("household_demographics", &["hd_demo_sk"]), + ("income_band", &["ib_income_band_sk"]), + ( + "inventory", + &["inv_date_sk", "inv_item_sk", "inv_warehouse_sk"], + ), + ("item", &["i_item_sk"]), + ("promotion", &["p_promo_sk"]), + ("reason", &["r_reason_sk"]), + ("ship_mode", &["sm_ship_mode_sk"]), + ("store", &["s_store_sk"]), + ("store_returns", &["sr_item_sk", "sr_ticket_number"]), + ("store_sales", &["ss_item_sk", "ss_ticket_number"]), + ("time_dim", &["t_time_sk"]), + ("warehouse", &["w_warehouse_sk"]), + ("web_page", &["wp_web_page_sk"]), + ("web_returns", &["wr_item_sk", "wr_order_number"]), + ("web_sales", &["ws_item_sk", "ws_order_number"]), + ("web_site", &["web_site_sk"]), +]; + +/// Get the constraints for a TPC-DS table. Only primary keys are returned; +/// TPC-DS also defines foreign keys, but those are currently unsupported. +fn table_constraints(table: &str, schema: &Schema) -> Constraints { + let columns = TPCDS_PRIMARY_KEYS + .iter() + .find(|(name, _)| *name == table) + .map(|(_, columns)| *columns) + .unwrap_or_else(|| unimplemented!("unknown TPC-DS table: {table}")); + + Constraints::new_unverified(vec![primary_key(schema, columns)]) +} + +fn primary_key(schema: &Schema, column_names: &[&str]) -> Constraint { + let indices = column_names + .iter() + .map(|column_name| { + schema.index_of(column_name).unwrap_or_else(|_| { + panic!("primary key column '{column_name}' not found in schema") + }) + }) + .collect(); + + Constraint::PrimaryKey(indices) +} + /// Get the SQL statements from the specified query file pub fn get_query_sql(base_query_path: &str, query: usize) -> Result> { if query > 0 && query < 100 { @@ -170,6 +226,7 @@ impl RunOpt { self.hash_join_buffering_capacity; let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); // register tables self.register_tables(&ctx).await?; @@ -234,7 +291,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } @@ -308,7 +365,6 @@ impl RunOpt { table: &str, ) -> Result> { let path = self.path.to_str().unwrap(); - let target_partitions = self.partitions(); // Obtain a snapshot of the SessionState let state = ctx.state(); @@ -324,10 +380,10 @@ impl RunOpt { let table_path = ListingTableUrl::parse(path)?; let options = ListingOptions::new(Arc::new(format)) - .with_file_extension(DEFAULT_PARQUET_EXTENSION) - .with_target_partitions(target_partitions) - .with_collect_stat(state.config().collect_statistics()); + .with_file_extension(DEFAULT_PARQUET_EXTENSION); + let schema = options.infer_schema(&state, &table_path).await?; + let constraints = table_constraints(table, schema.as_ref()); if self.common.debug { println!( @@ -347,9 +403,11 @@ impl RunOpt { .with_listing_options(options) .with_schema(schema); - Ok(Arc::new(ListingTable::try_new(config)?.with_cache( - ctx.runtime_env().cache_manager.get_file_statistic_cache(), - ))) + let provider = ListingTable::try_new(config)? + .with_constraints(constraints) + .with_cache(ctx.runtime_env().cache_manager.get_file_statistic_cache()); + + Ok(Arc::new(provider)) } fn iterations(&self) -> usize { diff --git a/benchmarks/src/tpch/mod.rs b/benchmarks/src/tpch/mod.rs index 08cedc0e5b4c3..9f3226ed5a8f6 100644 --- a/benchmarks/src/tpch/mod.rs +++ b/benchmarks/src/tpch/mod.rs @@ -20,7 +20,7 @@ use arrow::datatypes::SchemaBuilder; use datafusion::{ arrow::datatypes::{DataType, Field, Schema}, - common::plan_err, + common::{Constraint, Constraints, plan_err}, error::Result, }; use std::fs; @@ -138,6 +138,42 @@ pub fn get_tpch_table_schema(table: &str) -> Schema { } } +static TPCH_PRIMARY_KEYS: &[(&str, &[&str])] = &[ + ("region", &["r_regionkey"]), + ("nation", &["n_nationkey"]), + ("part", &["p_partkey"]), + ("supplier", &["s_suppkey"]), + ("partsupp", &["ps_partkey", "ps_suppkey"]), + ("customer", &["c_custkey"]), + ("orders", &["o_orderkey"]), + ("lineitem", &["l_orderkey", "l_linenumber"]), +]; + +/// Get the constraints for a TPC-H table. Only primary keys are returned; TPC-H +/// also defines foreign keys, but those are currently unsupported. +fn table_constraints(table: &str, schema: &Schema) -> Constraints { + let columns = TPCH_PRIMARY_KEYS + .iter() + .find(|(name, _)| *name == table) + .map(|(_, columns)| *columns) + .unwrap_or_else(|| unimplemented!("unknown TPC-H table: {table}")); + + Constraints::new_unverified(vec![primary_key(schema, columns)]) +} + +fn primary_key(schema: &Schema, column_names: &[&str]) -> Constraint { + let indices = column_names + .iter() + .map(|column_name| { + schema.index_of(column_name).unwrap_or_else(|_| { + panic!("primary key column '{column_name}' not found in schema") + }) + }) + .collect(); + + Constraint::PrimaryKey(indices) +} + /// Get the SQL statements from the specified query file pub fn get_query_sql(query: usize) -> Result> { get_query_sql_for_scale_factor(query, 1.0) diff --git a/benchmarks/src/tpch/run.rs b/benchmarks/src/tpch/run.rs index 75983ee141d93..47edfbac4b5a7 100644 --- a/benchmarks/src/tpch/run.rs +++ b/benchmarks/src/tpch/run.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use super::{ TPCH_QUERY_END_ID, TPCH_QUERY_START_ID, TPCH_TABLES, get_query_sql_for_scale_factor, - get_tbl_tpch_table_schema, get_tpch_table_schema, + get_tbl_tpch_table_schema, get_tpch_table_schema, table_constraints, }; use crate::util::{BenchmarkRun, CommonOpt, QueryResult, print_memory_stats}; @@ -137,6 +137,7 @@ impl RunOpt { self.hash_join_buffering_capacity; let rt = self.common.build_runtime()?; let ctx = SessionContext::new_with_config_rt(config, rt); + benchmark_run.set_memory_pool(&ctx.runtime_env().memory_pool); // register tables self.register_tables(&ctx).await?; let scale_factor = self.scale_factor()?; @@ -208,7 +209,7 @@ impl RunOpt { println!("Query {query_id} avg time: {avg:.2} ms"); // Print memory stats using mimalloc (only when compiled with --features mimalloc_extended) - print_memory_stats(); + print_memory_stats(&*ctx.runtime_env().memory_pool); Ok(query_results) } @@ -283,7 +284,6 @@ impl RunOpt { ) -> Result> { let path = self.path.to_str().unwrap(); let table_format = self.file_format.as_str(); - let target_partitions = self.partitions(); // Obtain a snapshot of the SessionState let state = ctx.state(); @@ -320,16 +320,16 @@ impl RunOpt { }; let table_path = ListingTableUrl::parse(path)?; - let options = ListingOptions::new(format) - .with_file_extension(extension) - .with_target_partitions(target_partitions) - .with_collect_stat(state.config().collect_statistics()); + let options = ListingOptions::new(format).with_file_extension(extension); + let schema = match table_format { "parquet" => options.infer_schema(&state, &table_path).await?, "tbl" => Arc::new(get_tbl_tpch_table_schema(table)), "csv" => Arc::new(get_tpch_table_schema(table)), _ => unreachable!(), }; + let constraints = table_constraints(table, schema.as_ref()); + let options = if self.sorted { let key_column_name = schema.fields()[0].name(); options @@ -342,9 +342,11 @@ impl RunOpt { .with_listing_options(options) .with_schema(schema); - Ok(Arc::new(ListingTable::try_new(config)?.with_cache( - ctx.runtime_env().cache_manager.get_file_statistic_cache(), - ))) + let provider = ListingTable::try_new(config)? + .with_constraints(constraints) + .with_cache(ctx.runtime_env().cache_manager.get_file_statistic_cache()); + + Ok(Arc::new(provider)) } fn iterations(&self) -> usize { diff --git a/benchmarks/src/util/memory.rs b/benchmarks/src/util/memory.rs index 11b96ef227756..f0339c9cf0c95 100644 --- a/benchmarks/src/util/memory.rs +++ b/benchmarks/src/util/memory.rs @@ -15,8 +15,32 @@ // specific language governing permissions and limitations // under the License. -/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api -pub fn print_memory_stats() { +use datafusion::execution::memory_pool::{MemoryPool, PeakRecordingPool}; + +/// Print Peak RSS, Peak Commit, Page Faults based on mimalloc api, followed by +/// the peak reservation of `memory_pool` when a memory limit was configured. +pub fn print_memory_stats(memory_pool: &dyn MemoryPool) { + print_allocator_stats(); + print_pool_stats(memory_pool); +} + +/// Print the peak reservation `memory_pool` has seen. +/// +/// Prints nothing when the benchmark ran without a memory limit, since no +/// [`PeakRecordingPool`] was installed to record. Comparing this against the +/// peak RSS above shows how much of a run's memory the pool actually accounted +/// for — DataFusion only tracks the "large" allocations that scale with input +/// size, so the two are expected to differ. +fn print_pool_stats(memory_pool: &dyn MemoryPool) { + if let Some(recorder) = PeakRecordingPool::from_pool(memory_pool) { + println!( + "Peak pool reserved: {}", + datafusion_common::human_readable_size(recorder.max_reserved()) + ); + } +} + +fn print_allocator_stats() { #[cfg(all(feature = "mimalloc", feature = "mimalloc_extended"))] { use datafusion_common::human_readable_size; diff --git a/benchmarks/src/util/options.rs b/benchmarks/src/util/options.rs index a3e6d2a4c5538..4a1c14674a1d0 100644 --- a/benchmarks/src/util/options.rs +++ b/benchmarks/src/util/options.rs @@ -21,7 +21,10 @@ use clap::Args; use datafusion::{ execution::{ disk_manager::DiskManagerBuilder, - memory_pool::{FairSpillPool, GreedyMemoryPool, MemoryPool, TrackConsumersPool}, + memory_pool::{ + FairSpillPool, GreedyMemoryPool, MemoryPool, PeakRecordingPool, + TrackConsumersPool, + }, object_store::ObjectStoreUrl, runtime_env::{RuntimeEnv, RuntimeEnvBuilder}, }, @@ -125,6 +128,9 @@ impl CommonOpt { ))); } }; + // Record the peak reservation so benchmarks can report it next to + // peak RSS. Purely observational: every call is delegated. + let pool: Arc = Arc::new(PeakRecordingPool::new(pool)); rt_builder = rt_builder .with_memory_pool(pool) .with_disk_manager_builder(DiskManagerBuilder::default()); diff --git a/benchmarks/src/util/run.rs b/benchmarks/src/util/run.rs index df17674e62961..772d421bc7bf4 100644 --- a/benchmarks/src/util/run.rs +++ b/benchmarks/src/util/run.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use datafusion::execution::memory_pool::{MemoryPool, PeakRecordingPool}; use datafusion::{DATAFUSION_VERSION, error::Result}; use datafusion_common::utils::get_available_parallelism; use serde::{Serialize, Serializer}; @@ -22,6 +23,7 @@ use serde_json::Value; use std::{ collections::HashMap, path::Path, + sync::Arc, time::{Duration, SystemTime}, }; @@ -91,6 +93,16 @@ pub struct BenchQuery { #[serde(serialize_with = "serialize_start_time")] start_time: SystemTime, success: bool, + /// Peak [`MemoryPool`] reservation observed while running this query, in + /// bytes. Recorded for failed queries too, since a query that ran out of + /// memory is one whose peak is worth seeing. + /// + /// `None` (and omitted from the JSON) only when the benchmark ran without a + /// memory limit, since there is then no pool to record. + /// + /// [`MemoryPool`]: datafusion::execution::memory_pool::MemoryPool + #[serde(skip_serializing_if = "Option::is_none")] + pool_peak_bytes: Option, } /// Internal representation of a single benchmark query iteration result. pub struct QueryResult { @@ -102,6 +114,10 @@ pub struct BenchmarkRun { context: RunContext, queries: Vec, current_case: Option, + /// The pool queries run against, when one was handed over with + /// [`BenchmarkRun::set_memory_pool`]. Only read through + /// [`BenchmarkRun::peak_recorder`]. + memory_pool: Option>, } impl Default for BenchmarkRun { @@ -117,15 +133,44 @@ impl BenchmarkRun { context: RunContext::new(), queries: vec![], current_case: None, + memory_pool: None, } } + + /// Report the peak reservation of `memory_pool` alongside each query. + /// + /// Call this with the pool of the [`RuntimeEnv`] the queries run against. + /// Has no effect unless a [`PeakRecordingPool`] is installed, which + /// [`CommonOpt::runtime_env_builder`] does whenever a memory limit is + /// configured; without one `pool_peak_bytes` is omitted from the results. + /// + /// Benchmarks that build a runtime per query should call this each time, so + /// each query reports against the pool it actually ran on. + /// + /// [`RuntimeEnv`]: datafusion::execution::runtime_env::RuntimeEnv + /// [`CommonOpt::runtime_env_builder`]: super::CommonOpt::runtime_env_builder + pub fn set_memory_pool(&mut self, memory_pool: &Arc) { + self.memory_pool = Some(Arc::clone(memory_pool)); + } + + /// The recorder in front of the pool set by [`Self::set_memory_pool`]. + fn peak_recorder(&self) -> Option<&PeakRecordingPool> { + PeakRecordingPool::from_pool(self.memory_pool.as_deref()?) + } + /// begin a new case. iterations added after this will be included in the new case pub fn start_new_case(&mut self, id: &str) { + // Give this query its own memory pool reading rather than inheriting + // the high-water mark of the queries that ran before it. + if let Some(recorder) = self.peak_recorder() { + recorder.reset_peak(); + } self.queries.push(BenchQuery { query: id.to_owned(), iterations: vec![], start_time: SystemTime::now(), success: true, + pool_peak_bytes: None, }); if let Some(c) = self.current_case.as_mut() { *c += 1; @@ -135,10 +180,14 @@ impl BenchmarkRun { } /// Write a new iteration to the current case pub fn write_iter(&mut self, elapsed: Duration, row_count: usize) { + // The peak is not reset between iterations, so this ends up holding the + // largest reservation seen across all of them. + let pool_peak_bytes = self.peak_recorder().map(PeakRecordingPool::peak_reserved); if let Some(idx) = self.current_case { self.queries[idx] .iterations - .push(QueryIter { elapsed, row_count }) + .push(QueryIter { elapsed, row_count }); + self.queries[idx].pool_peak_bytes = pool_peak_bytes; } else { panic!("no cases existed yet"); } @@ -159,8 +208,12 @@ impl BenchmarkRun { /// Mark current query pub fn mark_failed(&mut self) { + // A query that failed under a memory limit wrote no iteration, so this + // is the only chance to record what it had reserved when it gave up. + let pool_peak_bytes = self.peak_recorder().map(PeakRecordingPool::peak_reserved); if let Some(idx) = self.current_case { self.queries[idx].success = false; + self.queries[idx].pool_peak_bytes = pool_peak_bytes; } else { unreachable!("Cannot mark failure: no current case"); } @@ -182,3 +235,87 @@ impl BenchmarkRun { Ok(()) } } + +#[cfg(test)] +mod tests { + use datafusion::execution::memory_pool::{GreedyMemoryPool, MemoryConsumer}; + + use super::*; + + fn recording_pool(limit: usize) -> Arc { + Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new( + limit, + )))) + } + + #[test] + fn each_case_reports_its_own_peak() { + let pool = recording_pool(1024); + let mut run = BenchmarkRun::new(); + run.set_memory_pool(&pool); + + run.start_new_case("q1"); + let reservation = MemoryConsumer::new("q1").register(&pool); + reservation.try_grow(600).unwrap(); + run.write_iter(Duration::from_millis(1), 1); + drop(reservation); + + // The second case must not inherit the first case's high-water mark. + run.start_new_case("q2"); + let reservation = MemoryConsumer::new("q2").register(&pool); + reservation.try_grow(100).unwrap(); + run.write_iter(Duration::from_millis(1), 1); + + assert_eq!(run.queries[0].pool_peak_bytes, Some(600)); + assert_eq!(run.queries[1].pool_peak_bytes, Some(100)); + } + + #[test] + fn a_later_pool_replaces_an_earlier_one() { + let first = recording_pool(1024); + let mut run = BenchmarkRun::new(); + run.set_memory_pool(&first); + MemoryConsumer::new("q1") + .register(&first) + .try_grow(600) + .unwrap(); + + // Benchmarks that build a runtime per query hand over the new pool + // before the next case; the reading follows it. + let second = recording_pool(1024); + run.set_memory_pool(&second); + run.start_new_case("q2"); + MemoryConsumer::new("q2") + .register(&second) + .try_grow(100) + .unwrap(); + run.write_iter(Duration::from_millis(1), 1); + + assert_eq!(run.queries[0].pool_peak_bytes, Some(100)); + } + + #[test] + fn a_failed_query_still_reports_its_peak() { + let pool = recording_pool(1024); + let mut run = BenchmarkRun::new(); + run.set_memory_pool(&pool); + + run.start_new_case("q1"); + let reservation = MemoryConsumer::new("q1").register(&pool); + reservation.try_grow(600).unwrap(); + // No `write_iter`: the query failed before completing an iteration. + run.mark_failed(); + + assert_eq!(run.queries[0].pool_peak_bytes, Some(600)); + } + + #[test] + fn the_peak_is_omitted_without_a_recording_pool() { + let mut run = BenchmarkRun::new(); + run.start_new_case("q1"); + run.write_iter(Duration::from_millis(1), 1); + + assert_eq!(run.queries[0].pool_peak_bytes, None); + assert!(!run.to_json().contains("pool_peak_bytes")); + } +} diff --git a/ci/scripts/check_no_cargo_install_in_workflows.sh b/ci/scripts/check_no_cargo_install_in_workflows.sh new file mode 100755 index 0000000000000..aa84b2cf8f366 --- /dev/null +++ b/ci/scripts/check_no_cargo_install_in_workflows.sh @@ -0,0 +1,30 @@ +#!/usr/bin/env bash +# +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +set -euo pipefail + +SCRIPT_NAME="$(basename "${BASH_SOURCE[0]}")" +WORKFLOWS_DIR=".github/workflows" + +if grep -R -E -w -n --include='*.yml' --include='*.yaml' -- 'cargo.*install' "${WORKFLOWS_DIR}"; then + echo "[${SCRIPT_NAME}] Found workflow Rust tool installs that should use taiki-e/install-action instead." >&2 + exit 1 +fi + +echo "[${SCRIPT_NAME}] GitHub Actions workflow tool installs look good." diff --git a/clippy.toml b/clippy.toml index ea3609b574c06..7b781d9b6605f 100644 --- a/clippy.toml +++ b/clippy.toml @@ -1,6 +1,7 @@ disallowed-methods = [ { path = "tokio::task::spawn", reason = "To provide cancel-safety, use `SpawnedTask::spawn` instead (https://github.com/apache/datafusion/issues/6513)" }, { path = "tokio::task::spawn_blocking", reason = "To provide cancel-safety, use `SpawnedTask::spawn_blocking` instead (https://github.com/apache/datafusion/issues/6513)" }, + { path = "std::vec::Vec::reserve", reason = "Use `Vec::try_reserve` so allocation failures can be reported instead of panicking", replacement = "try_reserve" }, ] disallowed-types = [ diff --git a/datafusion-cli/Cargo.toml b/datafusion-cli/Cargo.toml index baf8e2c297fd2..62eedafe798d4 100644 --- a/datafusion-cli/Cargo.toml +++ b/datafusion-cli/Cargo.toml @@ -37,7 +37,7 @@ backtrace = ["datafusion/backtrace"] [dependencies] arrow = { workspace = true } async-trait = { workspace = true } -aws-config = "1.8.16" +aws-config = "1.8.18" aws-credential-types = "1.2.13" chrono = { workspace = true } clap = { version = "4.5.60", features = ["cargo", "derive"] } @@ -75,7 +75,7 @@ workspace = true [dev-dependencies] ctor = { workspace = true } insta = { workspace = true } -insta-cmd = "0.6.0" +insta-cmd = "0.7.0" rstest = { workspace = true } testcontainers-modules = { workspace = true, features = ["minio"] } # Makes sure `test_display_pg_json` behaves in a consistent way regardless of diff --git a/datafusion-cli/src/command.rs b/datafusion-cli/src/command.rs index 8aaa8025d1c3a..e847f7fdb501b 100644 --- a/datafusion-cli/src/command.rs +++ b/datafusion-cli/src/command.rs @@ -259,7 +259,7 @@ impl FromStr for OutputFormat { } impl OutputFormat { - pub async fn execute(&self, print_options: &mut PrintOptions) -> Result<()> { + pub fn execute(&self, print_options: &mut PrintOptions) -> Result<()> { match self { Self::ChangeFormat(format) => { if let Ok(format) = format.parse::() { diff --git a/datafusion-cli/src/exec.rs b/datafusion-cli/src/exec.rs index 09347d6d7dc2c..fc230d5362346 100644 --- a/datafusion-cli/src/exec.rs +++ b/datafusion-cli/src/exec.rs @@ -23,12 +23,12 @@ use crate::print_format::PrintFormat; use crate::{ command::{Command, OutputFormat}, helper::CliHelper, - object_storage::get_object_store, + object_storage::{get_object_store, stdin::StdinUtils}, print_options::{MaxRows, PrintOptions}, }; use datafusion::common::instant::Instant; use datafusion::common::{plan_datafusion_err, plan_err}; -use datafusion::config::ConfigFileType; +use datafusion::config::{ConfigFileType, Dialect}; use datafusion::datasource::listing::ListingTableUrl; use datafusion::error::{DataFusionError, Result}; use datafusion::execution::memory_pool::MemoryConsumer; @@ -148,7 +148,7 @@ pub async fn exec_from_repl( Command::OutputFormat(subcommand) => { if let Some(subcommand) = subcommand { if let Ok(command) = subcommand.parse::() { - if let Err(e) = command.execute(print_options).await { + if let Err(e) = command.execute(print_options) { eprintln!("{e}") } } else { @@ -223,9 +223,8 @@ pub(super) async fn exec_and_print( let dialect = &options.sql_parser.dialect; let dialect = dialect_from_str(dialect).ok_or_else(|| { plan_datafusion_err!( - "Unsupported SQL dialect: {dialect}. Available dialects: \ - Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, \ - MsSQL, ClickHouse, BigQuery, Ansi, DuckDB, Databricks." + "Unsupported SQL dialect: {dialect}. Available dialects: {}.", + Dialect::available() ) })?; @@ -418,17 +417,23 @@ async fn create_plan( // Note that cmd is a mutable reference so that create_external_table function can remove all // datafusion-cli specific options before passing through to datafusion. Otherwise, datafusion // will raise Configuration errors. - if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan { + if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &mut plan { // To support custom formats, treat error as None let format = config_file_type_from_str(&cmd.file_type); - register_object_store_and_config_extensions( - ctx, - &cmd.location, - &cmd.options, - format, - resolve_region, - ) - .await?; + + // Expose stdin (e.g. `cat data.csv | datafusion-cli`) as a `stdin://` + // object store, registered like any other scheme in `get_object_store`. + for location in &mut cmd.locations { + *location = StdinUtils::rewrite_location(location, format.as_ref()); + register_object_store_and_config_extensions( + ctx, + location, + &cmd.options, + format.clone(), + resolve_region, + ) + .await?; + } } if let LogicalPlan::Copy(copy_to) = &mut plan { @@ -531,14 +536,16 @@ mod tests { if let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = &plan { let format = config_file_type_from_str(&cmd.file_type); - register_object_store_and_config_extensions( - &ctx, - &cmd.location, - &cmd.options, - format, - false, - ) - .await?; + for location in &cmd.locations { + register_object_store_and_config_extensions( + &ctx, + location, + &cmd.options, + format.clone(), + false, + ) + .await?; + } } else { return plan_err!("LogicalPlan is not a CreateExternalTable"); } @@ -613,9 +620,8 @@ mod tests { let dialect = &task_ctx.session_config().options().sql_parser.dialect; let dialect = dialect_from_str(dialect).ok_or_else(|| { plan_datafusion_err!( - "Unsupported SQL dialect: {dialect}. Available dialects: \ - Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, \ - MsSQL, ClickHouse, BigQuery, Ansi, DuckDB, Databricks." + "Unsupported SQL dialect: {dialect}. Available dialects: {}.", + Dialect::available() ) })?; for location in locations { diff --git a/datafusion-cli/src/functions.rs b/datafusion-cli/src/functions.rs index df066992fb979..164af2559d2f6 100644 --- a/datafusion-cli/src/functions.rs +++ b/datafusion-cli/src/functions.rs @@ -42,6 +42,7 @@ use datafusion::physical_plan::ExecutionPlan; use datafusion::scalar::ScalarValue; use async_trait::async_trait; +use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use parquet::basic::ConvertedType; use parquet::data_type::{ByteArray, FixedLenByteArray}; use parquet::file::reader::FileReader; @@ -546,15 +547,17 @@ impl TableFunctionImpl for MetadataCacheFunc { for (path, entry) in cached_entries { path_arr.push(path.to_string()); file_modified_arr - .push(Some(entry.object_meta.last_modified.timestamp_millis())); - file_size_bytes_arr.push(entry.object_meta.size); - e_tag_arr.push(entry.object_meta.e_tag); - version_arr.push(entry.object_meta.version); + .push(Some(entry.value.meta.last_modified.timestamp_millis())); + file_size_bytes_arr.push(entry.value.meta.size); + e_tag_arr.push(entry.value.meta.e_tag); + version_arr.push(entry.value.meta.version); metadata_size_bytes.push(entry.size_bytes as u64); hits_arr.push(entry.hits as u64); let mut extra = entry - .extra + .value + .file_metadata + .extra_info() .iter() .map(|(k, v)| format!("{k}={v}")) .collect::>(); @@ -646,6 +649,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { Field::new("num_columns", DataType::UInt64, false), Field::new("table_size_bytes", DataType::Utf8, false), Field::new("statistics_size_bytes", DataType::UInt64, false), + Field::new("hits", DataType::UInt64, false), ])); // construct record batch from metadata @@ -659,6 +663,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { let mut num_columns_arr = vec![]; let mut table_size_bytes_arr = vec![]; let mut statistics_size_bytes_arr = vec![]; + let mut hits_arr = vec![]; if let Some(file_statistics_cache) = self.cache_manager.get_file_statistic_cache() { @@ -667,14 +672,23 @@ impl TableFunctionImpl for StatisticsCacheFunc { table_arr .push(path.table.map_or_else(|| "".to_string(), |t| t.to_string())); file_modified_arr - .push(Some(entry.object_meta.last_modified.timestamp_millis())); - file_size_bytes_arr.push(entry.object_meta.size); - e_tag_arr.push(entry.object_meta.e_tag); - version_arr.push(entry.object_meta.version); - num_rows_arr.push(entry.num_rows.to_string()); - num_columns_arr.push(entry.num_columns as u64); - table_size_bytes_arr.push(entry.table_size_bytes.to_string()); - statistics_size_bytes_arr.push(entry.statistics_size_bytes as u64); + .push(Some(entry.value.meta.last_modified.timestamp_millis())); + file_size_bytes_arr.push(entry.value.meta.size); + e_tag_arr.push(entry.value.meta.e_tag); + version_arr.push(entry.value.meta.version); + num_rows_arr.push(entry.value.statistics.num_rows.to_string()); + num_columns_arr + .push(entry.value.statistics.column_statistics.len() as u64); + table_size_bytes_arr + .push(entry.value.statistics.total_byte_size.to_string()); + statistics_size_bytes_arr.push( + entry + .value + .statistics + .heap_size(&mut DFHeapSizeCtx::default()) + as u64, + ); + hits_arr.push(entry.hits as u64); } } @@ -691,6 +705,7 @@ impl TableFunctionImpl for StatisticsCacheFunc { Arc::new(UInt64Array::from(num_columns_arr)), Arc::new(StringArray::from(table_size_bytes_arr)), Arc::new(UInt64Array::from(statistics_size_bytes_arr)), + Arc::new(UInt64Array::from(hits_arr)), ], )?; @@ -798,6 +813,7 @@ impl TableFunctionImpl for ListFilesCacheFunc { DataType::List(Arc::new(metadata_field.clone())), true, ), + Field::new("hits", DataType::UInt64, false), ])); let mut table_arr = vec![]; @@ -811,6 +827,7 @@ impl TableFunctionImpl for ListFilesCacheFunc { let mut etag_arr = vec![]; let mut version_arr = vec![]; let mut offsets: Vec = vec![0]; + let mut hits_arr = vec![]; if let Some(list_files_cache) = self.cache_manager.get_list_files_cache() { let now = Instant::now(); @@ -827,15 +844,16 @@ impl TableFunctionImpl for ListFilesCacheFunc { .map(|t| t.duration_since(now).as_millis() as i64), ); - for meta in entry.metas.files.iter() { + for meta in entry.value.files.iter() { file_path_arr.push(meta.location.to_string()); file_modified_arr.push(meta.last_modified.timestamp_millis()); file_size_bytes_arr.push(meta.size); etag_arr.push(meta.e_tag.clone()); version_arr.push(meta.version.clone()); } - current_offset += entry.metas.files.len() as i32; + current_offset += entry.value.files.len() as i32; offsets.push(current_offset); + hits_arr.push(entry.hits as u64); } } @@ -867,6 +885,7 @@ impl TableFunctionImpl for ListFilesCacheFunc { Arc::new(struct_arr), None, )), + Arc::new(UInt64Array::from(hits_arr)), ], )?; diff --git a/datafusion-cli/src/main.rs b/datafusion-cli/src/main.rs index 935bf0a9744dd..20a2537d7c10c 100644 --- a/datafusion-cli/src/main.rs +++ b/datafusion-cli/src/main.rs @@ -37,6 +37,7 @@ use datafusion_cli::functions::{ use datafusion_cli::object_storage::instrumented::{ InstrumentedObjectStoreMode, InstrumentedObjectStoreRegistry, }; +use datafusion_cli::object_storage::{StdinCarriesCommands, is_stdin_location}; use datafusion_cli::{ DATAFUSION_CLI_VERSION, exec, pool_type::PoolType, @@ -158,6 +159,23 @@ struct Args { object_store_profiling: InstrumentedObjectStoreMode, } +impl Args { + /// Without -c/-f the CLI enters the REPL, which reads its SQL from + /// stdin — interactively or piped. + fn repl_mode(&self) -> bool { + self.command.is_empty() && self.file.is_empty() + } + + /// Whether the CLI consumes stdin for its own SQL input. This covers the + /// REPL (no -c/-f, reading SQL interactively or piped) as well as an + /// explicit `-f /dev/stdin` (or the other stdin pseudo-paths), where the + /// SQL file *is* stdin. In either case stdin is already spoken for and + /// cannot also back a `LOCATION '/dev/stdin'` table. + fn reads_sql_from_stdin(&self) -> bool { + self.repl_mode() || self.file.iter().any(|f| is_stdin_location(f)) + } +} + #[tokio::main] /// Calls [`main_inner`], then handles printing errors and returning the correct exit code pub async fn main() -> ExitCode { @@ -268,6 +286,7 @@ async fn main_inner() -> Result<()> { instrumented_registry: Arc::clone(&instrumented_registry), }; + let repl_mode = args.repl_mode(); let commands = args.command; let files = args.file; let rc = match args.rc { @@ -285,7 +304,7 @@ async fn main_inner() -> Result<()> { } }; - if commands.is_empty() && files.is_empty() { + if repl_mode { if !rc.is_empty() { exec::exec_from_files(&ctx, rc, &print_options).await?; } @@ -316,7 +335,8 @@ fn get_session_config(args: &Args) -> Result { if batch_size == 0 { return config_err!("batch_size must be greater than 0"); } - config_options.execution.batch_size = batch_size; + config_options.execution.batch_size = + datafusion_common::config::ConfigNonZeroUsize::try_new(batch_size)?; }; // use easier to understand "tree" mode by default @@ -330,8 +350,16 @@ fn get_session_config(args: &Args) -> Result { config_options.format.null = String::from("NULL"); } - let session_config = + let mut session_config = SessionConfig::from(config_options).with_information_schema(true); + + if args.reads_sql_from_stdin() { + // When stdin carries the session's SQL — the REPL (including any rc + // file run before it) or an explicit `-f /dev/stdin` — it cannot also + // serve as a data source for `LOCATION '/dev/stdin'`. + session_config = session_config.with_extension(Arc::new(StdinCarriesCommands)); + } + Ok(session_config) } @@ -441,9 +469,10 @@ mod tests { use std::time::Duration; use super::*; + use datafusion::execution::cache::default_cache::DefaultCache; use datafusion::{ common::test_util::batches_to_string, - execution::cache::{DefaultListFilesCache, cache_manager::CacheManagerConfig}, + execution::cache::cache_manager::CacheManagerConfig, prelude::{ParquetReadOptions, col, lit, split_part}, }; use insta::assert_snapshot; @@ -613,9 +642,9 @@ mod tests { +-----------------------------------+-----------------+---------------------+------+------------------+ | filename | file_size_bytes | metadata_size_bytes | hits | extra | +-----------------------------------+-----------------+---------------------+------+------------------+ - | alltypes_plain.parquet | 1851 | 8882 | 2 | page_index=false | - | alltypes_tiny_pages.parquet | 454233 | 269074 | 2 | page_index=true | - | lz4_raw_compressed_larger.parquet | 380836 | 1339 | 2 | page_index=false | + | alltypes_plain.parquet | 1851 | 8794 | 1 | page_index=false | + | alltypes_tiny_pages.parquet | 454233 | 268970 | 2 | page_index=true | + | lz4_raw_compressed_larger.parquet | 380836 | 1331 | 1 | page_index=false | +-----------------------------------+-----------------+---------------------+------+------------------+ "); @@ -644,9 +673,9 @@ mod tests { +-----------------------------------+-----------------+---------------------+------+------------------+ | filename | file_size_bytes | metadata_size_bytes | hits | extra | +-----------------------------------+-----------------+---------------------+------+------------------+ - | alltypes_plain.parquet | 1851 | 8882 | 5 | page_index=false | - | alltypes_tiny_pages.parquet | 454233 | 269074 | 2 | page_index=true | - | lz4_raw_compressed_larger.parquet | 380836 | 1339 | 3 | page_index=false | + | alltypes_plain.parquet | 1851 | 8794 | 4 | page_index=false | + | alltypes_tiny_pages.parquet | 454233 | 268970 | 2 | page_index=true | + | lz4_raw_compressed_larger.parquet | 380836 | 1331 | 2 | page_index=false | +-----------------------------------+-----------------+---------------------+------+------------------+ "); @@ -682,17 +711,48 @@ mod tests { .await?; } - let sql = "SELECT split_part(path, '/', -1) as filename, table, file_size_bytes, num_rows, num_columns, table_size_bytes from statistics_cache() order by filename"; + let sql = "SELECT split_part(path, '/', -1) as filename, table, file_size_bytes, num_rows, num_columns, hits, table_size_bytes from statistics_cache() order by filename"; let df = ctx.sql(sql).await?; let rbs = df.collect().await?; - assert_snapshot!(batches_to_string(&rbs),@r" - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+ - | filename | table | file_size_bytes | num_rows | num_columns | table_size_bytes | - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+ - | alltypes_plain.parquet | alltypes_plain | 1851 | Exact(8) | 11 | Absent | - | alltypes_tiny_pages.parquet | alltypes_tiny_pages | 454233 | Exact(7300) | 13 | Absent | - | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836 | Exact(10000) | 1 | Absent | - +-----------------------------------+---------------------------+-----------------+--------------+-------------+------------------+ + assert_snapshot!(batches_to_string(&rbs),@" + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + | filename | table | file_size_bytes | num_rows | num_columns | hits | table_size_bytes | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + | alltypes_plain.parquet | alltypes_plain | 1851 | Exact(8) | 11 | 0 | Absent | + | alltypes_tiny_pages.parquet | alltypes_tiny_pages | 454233 | Exact(7300) | 13 | 0 | Absent | + | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836 | Exact(10000) | 1 | 0 | Absent | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + "); + + // increase the number of hits + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from alltypes_plain") + .await? + .collect() + .await?; + ctx.sql("select * from lz4_raw_compressed_larger") + .await? + .collect() + .await?; + + let sql = "SELECT split_part(path, '/', -1) as filename, table, file_size_bytes, num_rows, num_columns, hits, table_size_bytes from statistics_cache() order by filename"; + let df = ctx.sql(sql).await?; + let rbs = df.collect().await?; + assert_snapshot!(batches_to_string(&rbs),@" + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + | filename | table | file_size_bytes | num_rows | num_columns | hits | table_size_bytes | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ + | alltypes_plain.parquet | alltypes_plain | 1851 | Exact(8) | 11 | 3 | Absent | + | alltypes_tiny_pages.parquet | alltypes_tiny_pages | 454233 | Exact(7300) | 13 | 0 | Absent | + | lz4_raw_compressed_larger.parquet | lz4_raw_compressed_larger | 380836 | Exact(10000) | 1 | 1 | Absent | + +-----------------------------------+---------------------------+-----------------+--------------+-------------+------+------------------+ "); Ok(()) @@ -700,7 +760,7 @@ mod tests { #[tokio::test] async fn test_list_files_cache() -> Result<(), DataFusionError> { - let list_files_cache = Arc::new(DefaultListFilesCache::new( + let list_files_cache = Arc::new(DefaultCache::new_with_ttl( 1024, Some(Duration::from_secs(1)), )); @@ -750,7 +810,7 @@ mod tests { .collect() .await?; - let sql = "SELECT metadata_size_bytes, expires_in, metadata_list FROM list_files_cache()"; + let sql = "SELECT metadata_size_bytes, expires_in, metadata_list, hits FROM list_files_cache()"; let df = ctx .sql(sql) .await? @@ -778,16 +838,17 @@ mod tests { "filename", "file_size_bytes", "etag", + "hits", ])? .sort(vec![col("filename").sort(true, false)])?; let rbs = df.collect().await?; assert_snapshot!(batches_to_string(&rbs),@r" - +---------------------+-----------+-----------------+------+ - | metadata_size_bytes | filename | file_size_bytes | etag | - +---------------------+-----------+-----------------+------+ - | 212 | 0.parquet | 3642 | 0 | - | 212 | 1.parquet | 3642 | 1 | - +---------------------+-----------+-----------------+------+ + +---------------------+-----------+-----------------+------+------+ + | metadata_size_bytes | filename | file_size_bytes | etag | hits | + +---------------------+-----------+-----------------+------+------+ + | 212 | 0.parquet | 3642 | 0 | 2 | + | 212 | 1.parquet | 3642 | 1 | 2 | + +---------------------+-----------+-----------------+------+------+ "); Ok(()) diff --git a/datafusion-cli/src/object_storage.rs b/datafusion-cli/src/object_storage.rs index 34787838929f1..5e6337e303f6f 100644 --- a/datafusion-cli/src/object_storage.rs +++ b/datafusion-cli/src/object_storage.rs @@ -16,6 +16,9 @@ // under the License. pub mod instrumented; +pub(crate) mod stdin; + +pub use stdin::{StdinCarriesCommands, is_stdin_location}; use async_trait::async_trait; use aws_config::BehaviorVersion; @@ -53,6 +56,10 @@ use object_store::aws::resolve_bucket_region; // Provide a local mock when running tests so we don't make network calls #[cfg(test)] +#[expect( + clippy::unused_async, + reason = "matches object_store::aws::resolve_bucket_region" +)] async fn resolve_bucket_region( _bucket: &str, _client_options: &ClientOptions, @@ -173,7 +180,10 @@ struct CredentialsFromConfig { impl CredentialsFromConfig { /// Attempt find AWS S3 credentials via the AWS SDK pub async fn try_new() -> Result { - let config = aws_config::defaults(BehaviorVersion::latest()).load().await; + // Loading the SDK config produces a large future, so box it to avoid + // potentially triggering the `large_futures` clippy lint. + let config = + Box::pin(aws_config::defaults(BehaviorVersion::latest()).load()).await; let region = config.region().map(|r| r.to_string()); let credentials = config @@ -564,6 +574,9 @@ pub(crate) async fn get_object_store( .with_url(url.origin().ascii_serialization()) .build()?, ), + _ if scheme == stdin::StdinUtils::SCHEME => { + stdin::StdinUtils::get_or_create(state, url).await? + } _ => { // For other types, try to get from `object_store_registry`: state @@ -594,7 +607,7 @@ mod tests { #[tokio::test] async fn s3_object_store_builder_default() -> Result<()> { - if let Err(DataFusionError::Execution(e)) = check_aws_envs().await { + if let Err(DataFusionError::Execution(e)) = check_aws_envs() { // Skip test if AWS envs are not set eprintln!("{e}"); return Ok(()); @@ -759,7 +772,7 @@ mod tests { #[tokio::test] async fn s3_object_store_builder_resolves_region_when_none_provided() -> Result<()> { - if let Err(DataFusionError::Execution(e)) = check_aws_envs().await { + if let Err(DataFusionError::Execution(e)) = check_aws_envs() { // Skip test if AWS envs are not set eprintln!("{e}"); return Ok(()); @@ -792,7 +805,7 @@ mod tests { #[tokio::test] async fn s3_object_store_builder_overrides_region_when_resolve_region_enabled() -> Result<()> { - if let Err(DataFusionError::Execution(e)) = check_aws_envs().await { + if let Err(DataFusionError::Execution(e)) = check_aws_envs() { // Skip test if AWS envs are not set eprintln!("{e}"); return Ok(()); @@ -903,7 +916,7 @@ mod tests { table_options } - async fn check_aws_envs() -> Result<()> { + fn check_aws_envs() -> Result<()> { let aws_envs = [ "AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY", diff --git a/datafusion-cli/src/object_storage/stdin.rs b/datafusion-cli/src/object_storage/stdin.rs new file mode 100644 index 0000000000000..9e63068f4f32f --- /dev/null +++ b/datafusion-cli/src/object_storage/stdin.rs @@ -0,0 +1,388 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Exposes the process's standard input as a `stdin://` object store so that +//! piped data (e.g. `cat data.csv | datafusion-cli`) can be queried via +//! `CREATE EXTERNAL TABLE ... LOCATION '/dev/stdin'`. + +use std::io::{IsTerminal, Read}; +use std::sync::Arc; + +use datafusion::common::exec_datafusion_err; +use datafusion::config::ConfigFileType; +use datafusion::error::Result; +use datafusion::execution::context::SessionState; +use futures::TryStreamExt; + +use object_store::memory::InMemory; +use object_store::path::Path as ObjectStorePath; +use object_store::{ObjectStore, ObjectStoreExt}; +use url::Url; + +/// Marker [`SessionConfig`] extension recording that the session reads its SQL +/// commands from stdin (the interactive or piped REPL). stdin cannot then also +/// serve as a data source: reading it for table data would silently consume +/// the remaining SQL statements. +/// +/// [`SessionConfig`]: datafusion::execution::context::SessionConfig +#[derive(Debug)] +pub struct StdinCarriesCommands; + +/// Filesystem paths that refer to the process's standard input. +/// +/// These are intentionally limited to the well known pseudo-files exposed by +/// the operating system so that ordinary files are never accidentally treated +/// as stdin. +const STDIN_LOCATIONS: [&str; 3] = ["/dev/stdin", "/dev/fd/0", "/proc/self/fd/0"]; + +/// Returns `true` if `path` refers to the process's standard input. +/// +/// Re-exported as [`crate::object_storage::is_stdin_location`] so the CLI entry +/// point can detect when it reads its SQL from stdin via `-f /dev/stdin` and +/// avoid also offering stdin as a `LOCATION '/dev/stdin'` data source. +pub fn is_stdin_location(path: &str) -> bool { + STDIN_LOCATIONS.contains(&path) +} + +/// Utilities for exposing the process's standard input as an object store. +/// +/// stdin is surfaced as a `stdin://` object store and dispatched alongside the +/// other schemes (`s3`, `gs`, `http`, ...) so that reading piped data flows +/// through the normal object-store/listing code path, conceptually similar to +/// DuckDB's `PipeFileSystem`. +pub(crate) struct StdinUtils; + +impl StdinUtils { + /// The URL scheme used to expose stdin as an object store, mirroring how + /// `s3`, `gs`, `http`, etc. are addressed. + pub(crate) const SCHEME: &'static str = "stdin"; + + /// Rewrites the well known stdin pseudo-paths (e.g. `/dev/stdin`) to a + /// canonical `stdin://` URL so that reading from standard input flows + /// through the same object-store/listing code path as any other scheme. + /// Non-stdin locations are returned unchanged. + /// + /// The listing layer filters candidate files by extension, so the canonical + /// object is named with the extension matching the declared `STORED AS` + /// format. The name thereby also records which format stdin was consumed + /// as: a later stdin-backed table declaring a different format resolves to + /// a path the buffered store does not contain and is rejected by + /// [`Self::get_or_create`]. + pub(crate) fn rewrite_location( + location: &str, + format: Option<&ConfigFileType>, + ) -> String { + if !is_stdin_location(location) { + return location.to_string(); + } + + let object_name = match format { + Some(ConfigFileType::CSV) => "stdin.csv", + Some(ConfigFileType::JSON) => "stdin.json", + Some(ConfigFileType::PARQUET) => "stdin.parquet", + _ => "stdin", + }; + format!("{}:///{object_name}", Self::SCHEME) + } + + /// Returns the object store backing the `stdin://` scheme, buffering all of + /// standard input when the store is first constructed and reusing that + /// buffer for any subsequent `stdin://` table created in the same session. + /// + /// stdin is a one-shot stream: it can only be read once. The object store + /// registry keys by scheme/authority, so every `stdin://` URL maps to the + /// same store. Without this guard, a second `CREATE EXTERNAL TABLE ... + /// LOCATION '/dev/stdin'` would re-read (now-EOF) stdin, build an empty + /// store, and overwrite the populated one, silently emptying the earlier + /// table. Reusing the already-registered store avoids that. + /// + /// A later stdin-backed table declaring a different `STORED AS` format + /// resolves to an object the store does not contain (the object name + /// records the format stdin was consumed as) and is rejected with a clear + /// error — both reading the buffer as another format and re-reading stdin + /// would be silently wrong. + pub(crate) async fn get_or_create( + state: &SessionState, + url: &Url, + ) -> Result> { + let Ok(existing) = state.runtime_env().object_store_registry.get_store(url) + else { + return Self::object_store(state, url).await; + }; + + let path = ObjectStorePath::from_url_path(url.path())?; + if existing.head(&path).await.is_err() { + let buffered = existing + .list(None) + .try_next() + .await + .ok() + .flatten() + .map(|meta| format!(" as '{}'", meta.location)) + .unwrap_or_default(); + return Err(exec_datafusion_err!( + "stdin was already read{buffered} by an earlier statement; all \ + tables backed by stdin in a session must declare the same \ + STORED AS format" + )); + } + Ok(existing) + } + + /// Builds the object store backing the `stdin://` scheme by reading all of + /// standard input into memory. + /// + /// A pipe (e.g. `cat data.csv | datafusion-cli`) is not seekable and reports + /// a size of `0`, so it cannot be read directly by the file based formats + /// (CSV requires seeking, Parquet needs the footer at the end of the file). + /// Buffering the whole input up front sidesteps these limitations and lets + /// the data be read like any other object, including being scanned more than + /// once. + async fn object_store( + state: &SessionState, + url: &Url, + ) -> Result> { + if state + .config() + .get_extension::() + .is_some() + { + return Err(exec_datafusion_err!( + "stdin is already being read for SQL commands, so it cannot \ + also supply table data; pass the query with -c/--command or \ + -f/--file so that stdin carries the data, e.g. \ + `cat data.csv | datafusion-cli -f query.sql`" + )); + } + if std::io::stdin().is_terminal() { + return Err(exec_datafusion_err!( + "stdin is connected to a terminal, not piped data; pipe the \ + input in, e.g. `cat data.csv | datafusion-cli -f query.sql`" + )); + } + + let mut buffer = Vec::new(); + std::io::stdin() + .lock() + .read_to_end(&mut buffer) + .map_err(|e| exec_datafusion_err!("Failed to read from stdin: {e}"))?; + Self::in_memory_object_store(url, buffer).await + } + + /// Stores `data` at the path referenced by `url` in a fresh [`InMemory`] + /// store. + async fn in_memory_object_store( + url: &Url, + data: Vec, + ) -> Result> { + let store = InMemory::new(); + store + .put(&ObjectStorePath::from_url_path(url.path())?, data.into()) + .await?; + Ok(Arc::new(store)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use datafusion::prelude::{SessionConfig, SessionContext}; + + #[test] + fn rewrites_stdin_locations() { + // stdin pseudo-paths are rewritten to a `stdin://` URL carrying the + // extension that matches the declared format. + assert_eq!( + StdinUtils::rewrite_location("/dev/stdin", Some(&ConfigFileType::CSV)), + "stdin:///stdin.csv" + ); + assert_eq!( + StdinUtils::rewrite_location("/dev/fd/0", Some(&ConfigFileType::JSON)), + "stdin:///stdin.json" + ); + assert_eq!( + StdinUtils::rewrite_location( + "/proc/self/fd/0", + Some(&ConfigFileType::PARQUET) + ), + "stdin:///stdin.parquet" + ); + assert_eq!( + StdinUtils::rewrite_location("/dev/stdin", None), + "stdin:///stdin" + ); + + // Ordinary locations are left untouched. + for location in ["/dev/stdout", "data/stdin.csv", "stdin", "s3://b/f.csv"] { + assert_eq!( + StdinUtils::rewrite_location(location, Some(&ConfigFileType::CSV)), + location + ); + } + } + + /// Buffers `data` into the `stdin://` object store and reads it back through + /// a `CREATE EXTERNAL TABLE`, returning the number of rows in the table. + /// + /// This exercises the full path used for `/dev/stdin` short of the actual + /// stdin read, which cannot be driven from a unit test. + async fn count_stdin_rows( + data: Vec, + stored_as: &str, + format: Option, + options: &str, + ) -> Result { + let location = StdinUtils::rewrite_location("/dev/stdin", format.as_ref()); + let url = Url::parse(&location).unwrap(); + let store = StdinUtils::in_memory_object_store(&url, data).await?; + + let ctx = SessionContext::new(); + ctx.register_object_store(&url, store); + ctx.sql(&format!( + "CREATE EXTERNAL TABLE t STORED AS {stored_as} LOCATION '{location}' {options}" + )) + .await? + .collect() + .await?; + + ctx.sql("SELECT * FROM t").await?.count().await + } + + #[tokio::test] + async fn reuses_buffered_stdin_store() -> Result<()> { + // stdin can only be read once, so a second `stdin://` table must reuse + // the store buffered by the first instead of re-reading (now-empty) + // stdin and overwriting it. + // + // The very first read happens inside `get_or_create` -> `object_store`, + // which consumes the real process stdin and so cannot be driven from a + // unit test. Seed the registry with the store that first read would have + // produced (as the first `CREATE EXTERNAL TABLE` does), then drive the + // lookup through `get_or_create` and assert it hands back that exact + // store rather than rebuilding it. + let url = Url::parse("stdin:///stdin.csv").unwrap(); + let path = ObjectStorePath::from_url_path(url.path())?; + let buffered: Arc = Arc::new(InMemory::new()); + buffered.put(&path, b"a\n1\n2\n".to_vec().into()).await?; + + let ctx = SessionContext::new(); + ctx.register_object_store(&url, Arc::clone(&buffered)); + + let reused = StdinUtils::get_or_create(&ctx.state(), &url).await?; + assert!( + Arc::ptr_eq(&buffered, &reused), + "get_or_create must reuse the registered stdin store, not rebuild it" + ); + let bytes = reused.get(&path).await?.bytes().await?; + assert_eq!(bytes.as_ref(), b"a\n1\n2\n"); + Ok(()) + } + + #[tokio::test] + async fn rejects_second_stdin_table_with_different_format() -> Result<()> { + // The buffered object's name records the format stdin was consumed + // as; a later stdin table declaring a different format must fail with + // a clear error rather than a downstream "not found" (or silently + // misreading the bytes as another format). + let csv_url = Url::parse("stdin:///stdin.csv").unwrap(); + let store = + StdinUtils::in_memory_object_store(&csv_url, b"a\n1\n".to_vec()).await?; + + let ctx = SessionContext::new(); + ctx.register_object_store(&csv_url, store); + + let json_url = Url::parse("stdin:///stdin.json").unwrap(); + let err = StdinUtils::get_or_create(&ctx.state(), &json_url) + .await + .unwrap_err() + .to_string(); + assert!( + err.contains("must declare the same STORED AS format") + && err.contains("stdin.csv"), + "unexpected error: {err}" + ); + Ok(()) + } + + #[tokio::test] + async fn errors_when_stdin_carries_commands() { + // Once the REPL owns stdin for SQL commands, building the stdin store + // must fail with a clear error instead of swallowing the remaining + // statements as table data. + let config = SessionConfig::new().with_extension(Arc::new(StdinCarriesCommands)); + let ctx = SessionContext::new_with_config(config); + + let url = Url::parse("stdin:///stdin.csv").unwrap(); + let err = StdinUtils::get_or_create(&ctx.state(), &url) + .await + .unwrap_err(); + assert!( + err.to_string().contains("SQL commands"), + "unexpected error: {err}" + ); + } + + #[tokio::test] + async fn stdin_object_store_reads_csv() -> Result<()> { + let data = b"a,b\n1,foo\n2,bar\n".to_vec(); + let rows = count_stdin_rows( + data, + "CSV", + Some(ConfigFileType::CSV), + "OPTIONS ('format.has_header' 'true')", + ) + .await?; + assert_eq!(rows, 2); + Ok(()) + } + + #[tokio::test] + async fn stdin_object_store_reads_json() -> Result<()> { + let data = b"{\"a\": 1, \"b\": \"foo\"}\n{\"a\": 2, \"b\": \"bar\"}\n".to_vec(); + let rows = count_stdin_rows(data, "JSON", Some(ConfigFileType::JSON), "").await?; + assert_eq!(rows, 2); + Ok(()) + } + + #[tokio::test] + async fn stdin_object_store_reads_parquet() -> Result<()> { + use datafusion::arrow::array::Int32Array; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::arrow::record_batch::RecordBatch; + use parquet::arrow::ArrowWriter; + + // Parquet requires random access to the footer, which a real pipe cannot + // provide; the in-memory buffer makes this work. + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![1, 2, 3]))], + ) + .unwrap(); + + let mut data = Vec::new(); + let mut writer = ArrowWriter::try_new(&mut data, schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let rows = + count_stdin_rows(data, "PARQUET", Some(ConfigFileType::PARQUET), "").await?; + assert_eq!(rows, 3); + Ok(()) + } +} diff --git a/datafusion-cli/tests/cli_integration.rs b/datafusion-cli/tests/cli_integration.rs index 4849ac9e9a5e2..4dc244445a2eb 100644 --- a/datafusion-cli/tests/cli_integration.rs +++ b/datafusion-cli/tests/cli_integration.rs @@ -173,6 +173,222 @@ fn cli_quick_test<'a>( assert_cmd_snapshot!(cmd); } +/// Read data piped into the CLI via the `/dev/stdin` pseudo-path. +/// +/// Unix-only: `/dev/stdin` does not exist on Windows. This drives the real +/// binary through an actual pipe, exercising the stdin read that the in-process +/// unit tests cannot. +#[cfg(unix)] +#[test] +fn test_cli_read_from_stdin() { + let stdout = run_cli_with_stdin( + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin' \ + OPTIONS ('format.has_header' 'true'); \ + SELECT b, count(*) AS c FROM t GROUP BY b ORDER BY b;", + b"a,b\n1,foo\n2,bar\n3,foo\n", + ); + + assert!( + stdout.contains("| foo | 2 |") && stdout.contains("| bar | 1 |"), + "unexpected output:\n{stdout}" + ); +} + +/// stdin is a one-shot stream, so a second `/dev/stdin` table in the same +/// session must reuse the buffered input rather than re-reading (now-empty) +/// stdin and silently emptying the first table. +#[cfg(unix)] +#[test] +fn test_cli_read_from_stdin_twice_reuses_buffer() { + let stdout = run_cli_with_stdin( + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin' \ + OPTIONS ('format.has_header' 'true'); \ + CREATE EXTERNAL TABLE t2 STORED AS CSV LOCATION '/dev/stdin' \ + OPTIONS ('format.has_header' 'true'); \ + SELECT count(*) AS t_count FROM t; \ + SELECT count(*) AS t2_count FROM t2;", + b"a,b\n1,foo\n2,bar\n", + ); + + // Both tables must still see the two buffered rows. + let counts: Vec<&str> = stdout + .lines() + .filter(|line| line.trim_start().starts_with("| 2 ")) + .collect(); + assert_eq!( + counts.len(), + 2, + "expected both stdin tables to report 2 rows, got:\n{stdout}" + ); +} + +/// A later `/dev/stdin` table declaring a different `STORED AS` format must be +/// rejected with a clear error: stdin is one-shot, its bytes were already +/// buffered under the first table's format, and silently reading them as +/// another format would be wrong. +#[cfg(unix)] +#[test] +fn test_cli_read_from_stdin_mixed_formats_rejected() { + use std::io::Write; + use std::process::Stdio; + + let mut child = cli() + .args([ + "-q", + "--command", + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin' \ + OPTIONS ('format.has_header' 'true'); \ + CREATE EXTERNAL TABLE t2 STORED AS JSON LOCATION '/dev/stdin';", + ]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("failed to spawn datafusion-cli"); + + child + .stdin + .take() + .unwrap() + .write_all(b"a,b\n1,foo\n2,bar\n") + .unwrap(); + + let output = child.wait_with_output().unwrap(); + // Fatal errors in `--command` mode are reported on stdout. + let stdout = String::from_utf8_lossy(&output.stdout); + + assert!( + !output.status.success(), + "expected the mismatched format to fail, stdout:\n{stdout}" + ); + assert!( + stdout.contains("must declare the same STORED AS format"), + "expected a clear mismatch error, got:\n{stdout}" + ); +} + +/// When the SQL itself arrives on stdin (the piped REPL, e.g. `cat script.sql +/// | datafusion-cli`), stdin cannot double as a data source: the statement +/// must fail with a clear error instead of silently consuming the rest of the +/// script as table data, and the remaining statements must still run. +#[cfg(unix)] +#[test] +fn test_cli_stdin_location_rejected_when_sql_comes_from_stdin() { + use std::io::Write; + use std::process::Stdio; + + let mut child = cli() + .arg("-q") + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("failed to spawn datafusion-cli"); + + child + .stdin + .take() + .unwrap() + .write_all( + b"CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin';\n\ + SELECT 123 + 456;\n", + ) + .unwrap(); + + let output = child.wait_with_output().unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + + assert!( + stderr.contains("SQL commands"), + "expected a clear error about stdin carrying SQL.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + // The statement after the failed CREATE must still execute rather than + // being consumed as table data. + assert!( + stdout.contains("579"), + "expected the following statement to still run.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); +} + +/// `-f /dev/stdin` reads the SQL script from stdin, exactly like the piped +/// REPL, so stdin still cannot double as a `LOCATION '/dev/stdin'` data source. +/// The offending statement must fail with the same clear error, and later +/// statements in the script must still run. +/// +/// `/dev/stdin` only passes the `-f` file check when stdin is a redirected +/// regular file (a pipe is not `is_file()`), so the binary is driven with a +/// temp script file as its stdin rather than a pipe. +#[cfg(unix)] +#[test] +fn test_cli_dash_f_stdin_location_rejected() { + use std::process::Stdio; + + let script = env::temp_dir().join(format!( + "datafusion_cli_dash_f_stdin_{}.sql", + std::process::id() + )); + fs::write( + &script, + b"CREATE EXTERNAL TABLE t STORED AS CSV LOCATION '/dev/stdin';\n\ + SELECT 123 + 456;\n", + ) + .unwrap(); + let stdin = fs::File::open(&script).unwrap(); + + let output = cli() + .args(["-q", "-f", "/dev/stdin"]) + .stdin(Stdio::from(stdin)) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .output() + .expect("failed to spawn datafusion-cli"); + + let _ = fs::remove_file(&script); + + let stdout = String::from_utf8_lossy(&output.stdout); + let stderr = String::from_utf8_lossy(&output.stderr); + + assert!( + stderr.contains("SQL commands"), + "expected a clear error about stdin carrying SQL.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + // The statement after the failed CREATE must still execute rather than + // being consumed as table data. + assert!( + stdout.contains("579"), + "expected the following statement to still run.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); +} + +/// Spawns the real `datafusion-cli` binary, pipes `stdin` into it, and returns +/// its stdout after asserting a successful exit. +#[cfg(unix)] +fn run_cli_with_stdin(command: &str, stdin: &[u8]) -> String { + use std::io::Write; + use std::process::Stdio; + + let mut child = cli() + .args(["-q", "--command", command]) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .spawn() + .expect("failed to spawn datafusion-cli"); + + child.stdin.take().unwrap().write_all(stdin).unwrap(); + + let output = child.wait_with_output().unwrap(); + let stdout = String::from_utf8_lossy(&output.stdout).into_owned(); + let stderr = String::from_utf8_lossy(&output.stderr); + + assert!( + output.status.success(), + "datafusion-cli failed.\nstdout:\n{stdout}\nstderr:\n{stderr}" + ); + stdout +} + #[test] fn cli_explain_environment_overrides() { let mut settings = make_settings(); diff --git a/datafusion-examples/Cargo.toml b/datafusion-examples/Cargo.toml index bb8a92dbe05e7..6d6d917ac46ec 100644 --- a/datafusion-examples/Cargo.toml +++ b/datafusion-examples/Cargo.toml @@ -50,17 +50,17 @@ async-trait = { workspace = true } bytes = { workspace = true } dashmap = { workspace = true } # note only use main datafusion crate for examples -base64 = "0.22.1" +base64 = "0.23.0" datafusion-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } -datafusion-proto = { workspace = true } +datafusion-proto = { workspace = true, features = ["parquet"] } datafusion-sql = { workspace = true } env_logger = { workspace = true } futures = { workspace = true } insta = { workspace = true } log = { workspace = true } mimalloc = { version = "0.1", default-features = false } -object_store = { workspace = true, features = ["aws", "http"] } +object_store = { workspace = true, features = ["aws", "fs", "http"] } prost = { workspace = true } rand = { workspace = true } serde = { version = "1", features = ["derive"] } diff --git a/datafusion-examples/README.md b/datafusion-examples/README.md index 073f269d4a35d..86cfffe1a80e8 100644 --- a/datafusion-examples/README.md +++ b/datafusion-examples/README.md @@ -88,19 +88,21 @@ cargo run --example dataframe -- dataframe #### Category: Single Process -| Subcommand | File Path | Description | -| ---------------------- | ----------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- | -| catalog | [`data_io/catalog.rs`](examples/data_io/catalog.rs) | Register tables into a custom catalog | -| in_memory_object_store | [`data_io/in_memory_object_store.rs`](examples/data_io/in_memory_object_store.rs) | Read CSV from an in-memory object store (pattern applies to JSON/Parquet) | -| json_shredding | [`data_io/json_shredding.rs`](examples/data_io/json_shredding.rs) | Implement filter rewriting for JSON shredding | -| parquet_adv_idx | [`data_io/parquet_advanced_index.rs`](examples/data_io/parquet_advanced_index.rs) | Create a secondary index across multiple parquet files | -| parquet_emb_idx | [`data_io/parquet_embedded_index.rs`](examples/data_io/parquet_embedded_index.rs) | Store a custom index inside Parquet files | -| parquet_enc | [`data_io/parquet_encrypted.rs`](examples/data_io/parquet_encrypted.rs) | Read & write encrypted Parquet files | -| parquet_enc_with_kms | [`data_io/parquet_encrypted_with_kms.rs`](examples/data_io/parquet_encrypted_with_kms.rs) | Encrypted Parquet I/O using a KMS-backed factory | -| parquet_exec_visitor | [`data_io/parquet_exec_visitor.rs`](examples/data_io/parquet_exec_visitor.rs) | Extract statistics by visiting an ExecutionPlan | -| parquet_idx | [`data_io/parquet_index.rs`](examples/data_io/parquet_index.rs) | Create a secondary index | -| query_http_csv | [`data_io/query_http_csv.rs`](examples/data_io/query_http_csv.rs) | Query CSV files via HTTP | -| remote_catalog | [`data_io/remote_catalog.rs`](examples/data_io/remote_catalog.rs) | Interact with a remote catalog | +| Subcommand | File Path | Description | +| ----------------------- | ----------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- | +| catalog | [`data_io/catalog.rs`](examples/data_io/catalog.rs) | Register tables into a custom catalog | +| in_memory_object_store | [`data_io/in_memory_object_store.rs`](examples/data_io/in_memory_object_store.rs) | Read CSV from an in-memory object store (pattern applies to JSON/Parquet) | +| json_shredding | [`data_io/json_shredding.rs`](examples/data_io/json_shredding.rs) | Implement filter rewriting for JSON shredding | +| object_store_spill | [`data_io/object_store_spill.rs`](examples/data_io/object_store_spill.rs) | Use ObjectStore-backed spill files | +| parquet_adv_idx | [`data_io/parquet_advanced_index.rs`](examples/data_io/parquet_advanced_index.rs) | Create a secondary index across multiple parquet files | +| parquet_emb_idx | [`data_io/parquet_embedded_index.rs`](examples/data_io/parquet_embedded_index.rs) | Store a custom index inside Parquet files | +| parquet_enc | [`data_io/parquet_encrypted.rs`](examples/data_io/parquet_encrypted.rs) | Read & write encrypted Parquet files | +| parquet_enc_with_kms | [`data_io/parquet_encrypted_with_kms.rs`](examples/data_io/parquet_encrypted_with_kms.rs) | Encrypted Parquet I/O using a KMS-backed factory | +| parquet_exec_visitor | [`data_io/parquet_exec_visitor.rs`](examples/data_io/parquet_exec_visitor.rs) | Extract statistics by visiting an ExecutionPlan | +| parquet_idx | [`data_io/parquet_index.rs`](examples/data_io/parquet_index.rs) | Create a secondary index | +| partitioned_file_schema | [`data_io/partitioned_file_schema.rs`](examples/data_io/partitioned_file_schema.rs) | Provide an explicit arrow schema for a PartitionedFile | +| query_http_csv | [`data_io/query_http_csv.rs`](examples/data_io/query_http_csv.rs) | Query CSV files via HTTP | +| remote_catalog | [`data_io/remote_catalog.rs`](examples/data_io/remote_catalog.rs) | Interact with a remote catalog | ## DataFrame Examples @@ -218,14 +220,15 @@ cargo run --example dataframe -- dataframe #### Category: Single Process -| Subcommand | File Path | Description | -| --------------- | ----------------------------------------------------------- | ----------------------------------------------- | -| adv_udaf | [`udf/advanced_udaf.rs`](examples/udf/advanced_udaf.rs) | Advanced User Defined Aggregate Function (UDAF) | -| adv_udf | [`udf/advanced_udf.rs`](examples/udf/advanced_udf.rs) | Advanced User Defined Scalar Function (UDF) | -| adv_udwf | [`udf/advanced_udwf.rs`](examples/udf/advanced_udwf.rs) | Advanced User Defined Window Function (UDWF) | -| async_udf | [`udf/async_udf.rs`](examples/udf/async_udf.rs) | Asynchronous User Defined Scalar Function | -| udaf | [`udf/simple_udaf.rs`](examples/udf/simple_udaf.rs) | Simple UDAF example | -| udf | [`udf/simple_udf.rs`](examples/udf/simple_udf.rs) | Simple UDF example | -| udtf | [`udf/simple_udtf.rs`](examples/udf/simple_udtf.rs) | Simple UDTF example | -| udwf | [`udf/simple_udwf.rs`](examples/udf/simple_udwf.rs) | Simple UDWF example | -| table_list_udtf | [`udf/table_list_udtf.rs`](examples/udf/table_list_udtf.rs) | Session-aware UDTF table list example | +| Subcommand | File Path | Description | +| --------------- | ----------------------------------------------------------------------- | ----------------------------------------------- | +| adv_udaf | [`udf/advanced_udaf.rs`](examples/udf/advanced_udaf.rs) | Advanced User Defined Aggregate Function (UDAF) | +| adv_udf | [`udf/advanced_udf.rs`](examples/udf/advanced_udf.rs) | Advanced User Defined Scalar Function (UDF) | +| adv_udwf | [`udf/advanced_udwf.rs`](examples/udf/advanced_udwf.rs) | Advanced User Defined Window Function (UDWF) | +| async_udf | [`udf/async_udf.rs`](examples/udf/async_udf.rs) | Asynchronous User Defined Scalar Function | +| struct_udaf | [`udf/struct_returning_udaf.rs`](examples/udf/struct_returning_udaf.rs) | Struct-returning UDAF with window metadata | +| udaf | [`udf/simple_udaf.rs`](examples/udf/simple_udaf.rs) | Simple UDAF example | +| udf | [`udf/simple_udf.rs`](examples/udf/simple_udf.rs) | Simple UDF example | +| udtf | [`udf/simple_udtf.rs`](examples/udf/simple_udtf.rs) | Simple UDTF example | +| udwf | [`udf/simple_udwf.rs`](examples/udf/simple_udwf.rs) | Simple UDWF example | +| table_list_udtf | [`udf/table_list_udtf.rs`](examples/udf/table_list_udtf.rs) | Session-aware UDTF table list example | diff --git a/datafusion-examples/examples/custom_data_source/adapter_serialization.rs b/datafusion-examples/examples/custom_data_source/adapter_serialization.rs index d82bd2097ce1d..f18b888f3eb56 100644 --- a/datafusion-examples/examples/custom_data_source/adapter_serialization.rs +++ b/datafusion-examples/examples/custom_data_source/adapter_serialization.rs @@ -62,7 +62,8 @@ use datafusion_proto::bytes::{ use datafusion_proto::physical_plan::from_proto::parse_physical_expr_with_converter; use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter; use datafusion_proto::physical_plan::{ - PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, + PhysicalExtensionCodec, PhysicalPlanDecodeContext, PhysicalPlanNodeExt, + PhysicalProtoConverterExtension, }; use datafusion_proto::protobuf::physical_plan_node::PhysicalPlanType; use datafusion_proto::protobuf::{ @@ -275,6 +276,7 @@ impl PhysicalExtensionCodec for AdapterPreservingCodec { buf: &[u8], inputs: &[Arc], _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { // Try to parse as our extension payload if let Ok(payload) = serde_json::from_slice::(buf) @@ -303,6 +305,7 @@ impl PhysicalExtensionCodec for AdapterPreservingCodec { &self, _node: Arc, _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { // We don't need this for the example - adapter wrapping happens in // `execution_plan_to_proto` instead. diff --git a/datafusion-examples/examples/custom_data_source/custom_datasource.rs b/datafusion-examples/examples/custom_data_source/custom_datasource.rs index 701a886d2a140..6f176f8b46609 100644 --- a/datafusion-examples/examples/custom_data_source/custom_datasource.rs +++ b/datafusion-examples/examples/custom_data_source/custom_datasource.rs @@ -26,6 +26,7 @@ use async_trait::async_trait; use datafusion::arrow::array::{UInt8Builder, UInt64Builder}; use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::arrow::record_batch::RecordBatch; +use datafusion::common::assert_batches_eq; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::datasource::{TableProvider, TableType, provider_as_source}; use datafusion::error::Result; @@ -35,8 +36,8 @@ use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::memory::MemoryStream; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - SendableRecordBatchStream, project_schema, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, project_schema, }; use datafusion::prelude::*; @@ -53,6 +54,33 @@ pub async fn custom_datasource() -> Result<()> { search_accounts(db.clone(), Some(col("bank_account").gt(lit(8000u64))), 1).await?; search_accounts(db.clone(), Some(col("bank_account").gt(lit(200u64))), 2).await?; + // exercise SQL paths that push down non-trivial projections: + // - `SELECT 1 ...` requests no source columns (projection: Some([])) + // - `SELECT COUNT(id) ...` requests a single column (projection: Some([0])) + let ctx = SessionContext::new(); + ctx.register_table("accounts", Arc::new(db))?; + let constant_batches = ctx + .sql("SELECT 1 AS a FROM accounts") + .await? + .collect() + .await?; + assert_batches_eq!( + [ + "+---+", "| a |", "+---+", "| 1 |", "| 1 |", "| 1 |", "+---+", + ], + &constant_batches + ); + + let count_batches = ctx + .sql("SELECT COUNT(id) AS cnt FROM accounts") + .await? + .collect() + .await?; + assert_batches_eq!( + ["+-----+", "| cnt |", "+-----+", "| 3 |", "+-----+",], + &count_batches + ); + Ok(()) } @@ -118,7 +146,7 @@ impl Debug for CustomDataSource { } impl CustomDataSource { - pub(crate) async fn create_physical_plan( + pub(crate) fn create_physical_plan( &self, projections: Option<&Vec>, schema: SchemaRef, @@ -180,13 +208,14 @@ impl TableProvider for CustomDataSource { _filters: &[Expr], _limit: Option, ) -> Result> { - return self.create_physical_plan(projection, self.schema()).await; + self.create_physical_plan(projection, self.schema()) } } #[derive(Debug, Clone)] struct CustomExec { db: CustomDataSource, + projection: Option>, projected_schema: SchemaRef, cache: Arc, } @@ -202,6 +231,7 @@ impl CustomExec { let cache = Self::compute_properties(projected_schema.clone()); Self { db, + projection: projections.cloned(), projected_schema, cache: Arc::new(cache), } @@ -238,13 +268,24 @@ impl ExecutionPlan for CustomExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -263,32 +304,35 @@ impl ExecutionPlan for CustomExec { account_array.append_value(user.bank_account); } + // Build a batch holding every column the table can produce, then let + // Arrow drop the columns the query didn't ask for. `RecordBatch::project` + // preserves the row count, which matters when the projection selects + // zero columns (e.g. `SELECT 1 FROM t`). + let full_batch = RecordBatch::try_new( + self.db.schema(), + vec![ + Arc::new(id_array.finish()), + Arc::new(account_array.finish()), + ], + )?; + let batch = match &self.projection { + Some(indices) => full_batch.project(indices)?, + None => full_batch, + }; + Ok(Box::pin(MemoryStream::try_new( - vec![RecordBatch::try_new( - self.projected_schema.clone(), - vec![ - Arc::new(id_array.finish()), - Arc::new(account_array.finish()), - ], - )?], - self.schema(), + vec![batch], + self.projected_schema.clone(), None, )?)) } fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion-examples/examples/data_io/main.rs b/datafusion-examples/examples/data_io/main.rs index 4656a83670aaf..041308463cda9 100644 --- a/datafusion-examples/examples/data_io/main.rs +++ b/datafusion-examples/examples/data_io/main.rs @@ -21,7 +21,7 @@ //! //! ## Usage //! ```bash -//! cargo run --example data_io -- [all|catalog|in_memory_object_store|json_shredding|parquet_adv_idx|parquet_emb_idx|parquet_enc_with_kms|parquet_enc|parquet_exec_visitor|parquet_idx|query_http_csv|remote_catalog] +//! cargo run --example data_io -- [all|catalog|in_memory_object_store|json_shredding|object_store_spill|parquet_adv_idx|parquet_emb_idx|parquet_enc_with_kms|parquet_enc|parquet_exec_visitor|parquet_idx|query_http_csv|remote_catalog] //! ``` //! //! Each subcommand runs a corresponding example: @@ -36,6 +36,9 @@ //! - `json_shredding` //! (file: json_shredding.rs, desc: Implement filter rewriting for JSON shredding) //! +//! - `object_store_spill` +//! (file: object_store_spill.rs, desc: Use ObjectStore-backed spill files) +//! //! - `parquet_adv_idx` //! (file: parquet_advanced_index.rs, desc: Create a secondary index across multiple parquet files) //! @@ -54,6 +57,9 @@ //! - `parquet_idx` //! (file: parquet_index.rs, desc: Create a secondary index) //! +//! - `partitioned_file_schema` +//! (file: partitioned_file_schema.rs, desc: Provide an explicit arrow schema for a PartitionedFile) +//! //! - `query_http_csv` //! (file: query_http_csv.rs, desc: Query CSV files via HTTP) //! @@ -63,12 +69,14 @@ mod catalog; mod in_memory_object_store; mod json_shredding; +mod object_store_spill; mod parquet_advanced_index; mod parquet_embedded_index; mod parquet_encrypted; mod parquet_encrypted_with_kms; mod parquet_exec_visitor; mod parquet_index; +mod partitioned_file_schema; mod query_http_csv; mod remote_catalog; @@ -83,12 +91,14 @@ enum ExampleKind { Catalog, InMemoryObjectStore, JsonShredding, + ObjectStoreSpill, ParquetAdvIdx, ParquetEmbIdx, ParquetEnc, ParquetEncWithKms, ParquetExecVisitor, ParquetIdx, + PartitionedFileSchema, QueryHttpCsv, RemoteCatalog, } @@ -113,6 +123,9 @@ impl ExampleKind { in_memory_object_store::in_memory_object_store().await? } ExampleKind::JsonShredding => json_shredding::json_shredding().await?, + ExampleKind::ObjectStoreSpill => { + object_store_spill::object_store_spill().await? + } ExampleKind::ParquetAdvIdx => { parquet_advanced_index::parquet_advanced_index().await? } @@ -127,6 +140,9 @@ impl ExampleKind { parquet_exec_visitor::parquet_exec_visitor().await? } ExampleKind::ParquetIdx => parquet_index::parquet_index().await?, + ExampleKind::PartitionedFileSchema => { + partitioned_file_schema::read_partitioned_file().await? + } ExampleKind::QueryHttpCsv => query_http_csv::query_http_csv().await?, ExampleKind::RemoteCatalog => remote_catalog::remote_catalog().await?, } diff --git a/datafusion-examples/examples/data_io/object_store_spill.rs b/datafusion-examples/examples/data_io/object_store_spill.rs new file mode 100644 index 0000000000000..d7d5392f66953 --- /dev/null +++ b/datafusion-examples/examples/data_io/object_store_spill.rs @@ -0,0 +1,273 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! See `main.rs` for how to run it. +//! +//! [`object_store_spill`] demonstrates how to use the [`TempFileFactory`] API to configure +//! DataFusion to spill intermediate results to remote storage when it exceeds +//! the configured memory limits. +//! +//! See [`datafusion::execution::memory_pool`] for more information on how +//! DataFusion decides when operators should spill, and [`SpillFile`] for the +//! spill file abstraction this example implements. +use std::future::Future; +use std::io::Write; +use std::path::Path as StdPath; +use std::pin::Pin; +use std::sync::Arc; +use std::sync::atomic::{AtomicU64, Ordering}; + +use bytes::Bytes; +use datafusion::common::Result; +use datafusion::execution::disk_manager::DiskManagerBuilder; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::execution::{SpillFile, SpillWriter, TempFileFactory}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::exec_err; +use futures::{Stream, StreamExt, TryStreamExt, stream}; +use object_store::local::LocalFileSystem; +use object_store::path::Path; +use object_store::{ObjectStore, ObjectStoreExt, PutPayload}; +use tempfile::tempdir; + +/// Demonstrates configuring DataFusion with spill files backed by an ObjectStore. +pub async fn object_store_spill() -> Result<()> { + // A real system would use S3, GCS, Azure, or some other ObjectStore for + // remote spills. This example uses a local-file-backed ObjectStore for + // simplicity. + let tmp_dir = tempdir()?; + let store: Arc = + Arc::new(LocalFileSystem::new_with_prefix(tmp_dir.path())?); + + // Create the custom TempFileFactory that creates spill files in the ObjectStore. + let temp_file_factory = Arc::new(ObjectStoreTempFileFactory::new(store)); + let disk_manager_builder = + DiskManagerBuilder::default().with_temp_file_factory(temp_file_factory.clone()); + let runtime = RuntimeEnvBuilder::new() + .with_disk_manager_builder(disk_manager_builder) // use the factory + // and set a small memory limit so the example spills + .with_memory_limit(1024 * 1024, 1.0) + .build_arc()?; + + // Configure a SessionContext for running queries; use a single partition + // and no sort spill reservation to make the example deterministic and keep + // the spill behavior easy to observe. + let config = SessionConfig::new() + .with_sort_spill_reservation_bytes(0) + .with_sort_in_place_threshold_bytes(0) + .with_target_partitions(1); + let ctx = SessionContext::new_with_config_rt(config, Arc::clone(&runtime)); + + // Run an SQL query that sorts a "large" amount of data. Given the + // SessionContext's low memory limit, the sort will spill. + let row_count = 10_000_000; + let mut stream = ctx + .sql(&format!( + "SELECT * FROM generate_series(1, {row_count}) AS t(v) ORDER BY v DESC" + )) + .await? + .execute_stream() + .await?; + + // Drive the query to completion, and verify output + let mut output_rows = 0; + while let Some(batch) = stream.next().await { + output_rows += batch?.num_rows(); + } + + assert_eq!(output_rows, row_count as usize); + assert!( + temp_file_factory.created_files() > 0, + "expected the custom TempFileFactory to be used for spilling" + ); + + Ok(()) +} + +/// Creates spill files backed by an [`ObjectStore`]. +/// +/// DataFusion calls this factory whenever an operator needs a new temporary +/// file for spilling. A remote deployment would use the same pattern with an +/// S3, GCS, Azure, or other remote ObjectStore implementation. +struct ObjectStoreTempFileFactory { + /// ObjectStore used for spill file reads and writes. + store: Arc, + /// Monotonic counter used to create unique object paths. + counter: AtomicU64, + /// Counts how many spill files DataFusion requested from this factory. + created_files: AtomicU64, +} + +impl ObjectStoreTempFileFactory { + /// Create a new spill file factory that stores spill data in `store`. + fn new(store: Arc) -> Self { + Self { + store, + counter: AtomicU64::new(0), + created_files: AtomicU64::new(0), + } + } + + /// Return the number of spill files created through this factory. + fn created_files(&self) -> u64 { + self.created_files.load(Ordering::Relaxed) + } +} + +impl TempFileFactory for ObjectStoreTempFileFactory { + /// Create one logical spill file backed by an ObjectStore path. + fn create_temp_file(&self, description: &str) -> Result> { + let id = self.counter.fetch_add(1, Ordering::Relaxed); + self.created_files.fetch_add(1, Ordering::Relaxed); + + // Convert a query-provided spill description into an ObjectStore-safe path component. + // + // For example, `"Sort Spill: partition 0"` becomes `"Sort_Spill__partition_0"`. + let cleaned_description: String = description + .chars() + .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' }) + .collect(); + let location = Path::from(format!("spill/{cleaned_description}-{id}.bin")); + + // Return a SpillFile implementation that reads and writes this ObjectStore path. + Ok(Arc::new(ObjectStoreSpillFile { + store: Arc::clone(&self.store), + location, + size: Arc::new(AtomicU64::new(0)), + })) + } +} + +/// Logical spill file stored at an ObjectStore path. +/// +/// DataFusion writes spill data by calling [`SpillFile::open_writer`] and reads +/// it back by calling [`SpillFile::read_stream`]. +struct ObjectStoreSpillFile { + /// ObjectStore containing the spill object. + store: Arc, + /// ObjectStore path for this spill object. + location: Path, + /// Last committed object size, updated when the writer finishes. + size: Arc, +} + +impl SpillFile for ObjectStoreSpillFile { + /// Return no local filesystem path because the spill file is accessed through ObjectStore. + fn path(&self) -> Option<&StdPath> { + None // Remote ObjectStores do not have a local OS path. + } + + /// Return the size of the uploaded object + fn size(&self) -> Option { + // Return the last committed size, which this example tracks after upload. + Some(self.size.load(Ordering::Relaxed)) + } + + /// Read the spill file contents as a byte stream. + fn read_stream(&self) -> Result> + Send>>> { + let store = Arc::clone(&self.store); + let location = self.location.clone(); + + // Use `stream::once` to defer the ObjectStore read until DataFusion + // polls the returned stream. + let result_stream = + async move { store.get(&location).await.map(|r| r.into_stream()) }; + let stream = stream::once(result_stream) + .try_flatten() + .map_err(Into::into); + + Ok(Box::pin(stream)) + } + + /// Open a synchronous writer for this spill file. + fn open_writer(&self) -> Result> { + // Create a writer that buffers bytes and uploads them on finish. + Ok(Box::new(ObjectStoreSpillWriter { + store: Arc::clone(&self.store), + location: self.location.clone(), + size: Arc::clone(&self.size), + buffer: Vec::new(), + })) + } +} + +/// Adapts DataFusion's [`SpillWriter`] API to ObjectStore. +/// +/// This simple example buffers bytes in memory and uploads them in +/// [`SpillWriter::finish`]. A production remote implementation should consider +/// multipart or streaming uploads. +struct ObjectStoreSpillWriter { + /// ObjectStore to read/write bytes to. + store: Arc, + /// ObjectStore path to upload to. + location: Path, + /// Shared size field on the corresponding [`ObjectStoreSpillFile`]. + size: Arc, + /// Buffered spill bytes waiting to be uploaded. + /// + /// This simple example buffers the spill and uploads it on finish. + /// Production remote stores should consider multipart or streaming uploads. + buffer: Vec, +} + +impl Write for ObjectStoreSpillWriter { + /// Append bytes to the in-memory buffer. + fn write(&mut self, buf: &[u8]) -> std::io::Result { + // Buffer bytes written through the synchronous Write API. + self.buffer.extend_from_slice(buf); + Ok(buf.len()) + } + + /// No-op because data is committed in [`SpillWriter::finish`]. + fn flush(&mut self) -> std::io::Result<()> { + Ok(()) + } +} + +impl SpillWriter for ObjectStoreSpillWriter { + /// Upload buffered bytes to ObjectStore and mark the spill file complete. + fn finish(&mut self) -> Result<()> { + // Move the buffered bytes into the upload future. + let store = Arc::clone(&self.store); + let location = self.location.clone(); + let data = std::mem::take(&mut self.buffer); + let size = data.len() as u64; + + // This simple example buffers the spill and uploads it on finish. + // Production remote stores should consider multipart or streaming uploads. + block_on_object_store(async move { + store + .put(&location, PutPayload::from_bytes(data.into())) + .await?; + Ok(()) + })?; + + self.size.store(size, Ordering::Relaxed); + Ok(()) + } +} + +/// Run an async ObjectStore operation. +/// +/// Adding a native async API is tracked in +fn block_on_object_store(future: impl Future>) -> Result { + if let Ok(handle) = tokio::runtime::Handle::try_current() { + tokio::task::block_in_place(|| handle.block_on(future)) + } else { + exec_err!("No current Tokio runtime available") + } +} diff --git a/datafusion-examples/examples/data_io/parquet_advanced_index.rs b/datafusion-examples/examples/data_io/parquet_advanced_index.rs index 9bdcda265ea7e..b6440eb3e2078 100644 --- a/datafusion-examples/examples/data_io/parquet_advanced_index.rs +++ b/datafusion-examples/examples/data_io/parquet_advanced_index.rs @@ -41,13 +41,14 @@ use datafusion::parquet::arrow::ArrowWriter; use datafusion::parquet::arrow::arrow_reader::{ ArrowReaderOptions, ParquetRecordBatchReaderBuilder, RowSelection, RowSelector, }; -use datafusion::parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; +use datafusion::parquet::arrow::async_reader::AsyncFileReader; +use datafusion::parquet::errors::ParquetError; use datafusion::parquet::file::metadata::{PageIndexPolicy, ParquetMetaData}; use datafusion::parquet::file::properties::{EnabledStatistics, WriterProperties}; use datafusion::parquet::schema::types::ColumnPath; use datafusion::physical_expr::PhysicalExpr; use datafusion::physical_expr::utils::{Guarantee, LiteralGuarantee}; -use datafusion::physical_optimizer::pruning::PruningPredicate; +use datafusion::physical_optimizer::pruning::PruningPredicateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; use datafusion::prelude::*; @@ -59,7 +60,7 @@ use bytes::Bytes; use datafusion::datasource::memory::DataSourceExec; use futures::FutureExt; use futures::future::BoxFuture; -use object_store::ObjectStore; +use object_store::{ObjectStore, ObjectStoreExt}; use tempfile::TempDir; use url::Url; @@ -155,6 +156,7 @@ use url::Url; /// ``` /// /// [`ListingTable`]: datafusion::datasource::listing::ListingTable +/// [`PruningPredicate`]: datafusion::physical_optimizer::pruning::PruningPredicate /// [Page Index](https://github.com/apache/parquet-format/blob/master/PageIndex.md) pub async fn parquet_advanced_index() -> Result<()> { // the object store is used to read the parquet files (in this case, it is @@ -300,8 +302,9 @@ impl IndexTableProvider { // In this example, we use the PruningPredicate's literal guarantees to // analyze the predicate. In a real system, using // `PruningPredicate::prune` would likely be easier to do. - let pruning_predicate = - PruningPredicate::try_new(Arc::clone(predicate), self.schema())?; + let pruning_predicate = PruningPredicateBuilder::new() + .with_file_schema(self.schema()) + .try_build(Arc::clone(predicate))?; // The PruningPredicate's guarantees must all be satisfied in order for // the predicate to possibly evaluate to true. @@ -552,12 +555,19 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { &self, _partition_index: usize, partitioned_file: PartitionedFile, - metadata_size_hint: Option, + _metadata_size_hint: Option, _metrics: &ExecutionPlanMetricsSet, ) -> Result> { // for this example we ignore the partition index and metrics // but in a real system you would likely use them to report details on // the performance of the reader. + // + // We also ignore the metadata size hint as this reader always serves + // metadata from the pre-populated `self.metadata` cache, so it never + // performs the footer fetch the hint is meant to optimize. A real + // implementation would likely pass the hint to + // `ParquetMetaDataReader::with_prefetch_hint` to reduce the number of + // IO requests needed to load the footer. let filename = partitioned_file .object_meta .location @@ -568,13 +578,7 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { .to_string(); let object_store = Arc::clone(&self.object_store); - let mut inner = - ParquetObjectReader::new(object_store, partitioned_file.object_meta.location) - .with_file_size(partitioned_file.object_meta.size); - - if let Some(hint) = metadata_size_hint { - inner = inner.with_footer_size_hint(hint) - }; + let location = partitioned_file.object_meta.location; let metadata = self .metadata @@ -583,16 +587,18 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { Ok(Box::new(ParquetReaderWithCache { filename, metadata: Arc::clone(metadata), - inner, + object_store, + location, })) } } -/// wrapper around a ParquetObjectReader that caches metadata +/// An [`AsyncFileReader`] that reads from an [`ObjectStore`] and caches metadata struct ParquetReaderWithCache { filename: String, metadata: Arc, - inner: ParquetObjectReader, + object_store: Arc, + location: object_store::path::Path, } impl AsyncFileReader for ParquetReaderWithCache { @@ -601,7 +607,15 @@ impl AsyncFileReader for ParquetReaderWithCache { range: Range, ) -> BoxFuture<'_, datafusion::parquet::errors::Result> { println!("get_bytes: {} Reading range {:?}", self.filename, range); - self.inner.get_bytes(range) + let object_store = Arc::clone(&self.object_store); + let location = self.location.clone(); + async move { + object_store + .get_range(&location, range) + .await + .map_err(|e| ParquetError::External(Box::new(e))) + } + .boxed() } fn get_byte_ranges( @@ -612,7 +626,15 @@ impl AsyncFileReader for ParquetReaderWithCache { "get_byte_ranges: {} Reading ranges {:?}", self.filename, ranges ); - self.inner.get_byte_ranges(ranges) + let object_store = Arc::clone(&self.object_store); + let location = self.location.clone(); + async move { + object_store + .get_ranges(&location, &ranges) + .await + .map_err(|e| ParquetError::External(Box::new(e))) + } + .boxed() } fn get_metadata( diff --git a/datafusion-examples/examples/data_io/parquet_embedded_index.rs b/datafusion-examples/examples/data_io/parquet_embedded_index.rs index 40b5b468ff5bf..a8a3c97fa11f8 100644 --- a/datafusion-examples/examples/data_io/parquet_embedded_index.rs +++ b/datafusion-examples/examples/data_io/parquet_embedded_index.rs @@ -87,7 +87,7 @@ //! 2. Read and deserialize the index. //! //! 3. Create a `TableProvider` that knows how to use the index to quickly find -//! the relevant files, row groups, data pages or rows based on on pushed down +//! the relevant files, row groups, data pages or rows based on pushed down //! filters. //! //! # FAQ: Why do other Parquet readers skip over the custom index? diff --git a/datafusion-examples/examples/data_io/parquet_index.rs b/datafusion-examples/examples/data_io/parquet_index.rs index 9be84d8249342..753d1b30fc0e8 100644 --- a/datafusion-examples/examples/data_io/parquet_index.rs +++ b/datafusion-examples/examples/data_io/parquet_index.rs @@ -42,7 +42,7 @@ use datafusion::parquet::arrow::{ ArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder, }; use datafusion::physical_expr::PhysicalExpr; -use datafusion::physical_optimizer::pruning::PruningPredicate; +use datafusion::physical_optimizer::pruning::PruningPredicateBuilder; use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::*; use std::collections::HashSet; @@ -274,7 +274,8 @@ impl TableProvider for IndexTableProvider { /// Simple in memory secondary index for a set of parquet files /// /// The index is represented as an arrow [`RecordBatch`] that can be passed -/// directly by the DataFusion [`PruningPredicate`] API +/// directly by the DataFusion +/// [`datafusion::physical_optimizer::pruning::PruningPredicate`] API /// /// The `RecordBatch` looks as follows. /// @@ -362,8 +363,9 @@ impl ParquetMetadataIndex { ) -> Result> { // Use the PruningPredicate API to determine which files can not // possibly have any relevant data. - let pruning_predicate = - PruningPredicate::try_new(predicate, self.schema().clone())?; + let pruning_predicate = PruningPredicateBuilder::new() + .with_file_schema(self.schema().clone()) + .try_build(predicate)?; // Now evaluate the pruning predicate into a boolean mask, one element per // file in the index. If the mask is true, the file may have rows that diff --git a/datafusion-examples/examples/data_io/partitioned_file_schema.rs b/datafusion-examples/examples/data_io/partitioned_file_schema.rs new file mode 100644 index 0000000000000..b423ebb6a38b0 --- /dev/null +++ b/datafusion-examples/examples/data_io/partitioned_file_schema.rs @@ -0,0 +1,147 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! See `main.rs` for how to run it. + +use arrow::array::{Int32Array, RecordBatch}; +use arrow_schema::{DataType, Field, Schema, SchemaRef}; +use datafusion::common::Result; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::datasource::physical_plan::{FileScanConfigBuilder, ParquetSource}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::parquet::arrow::ArrowWriter; +use datafusion::parquet::file::reader::Length; +use datafusion::physical_plan::ExecutionPlan; +use futures::StreamExt; +use std::fs::File; +use std::path::Path; +use std::sync::Arc; +use tempfile::TempDir; + +/// Demonstrates how to attach a per-file Arrow schema to a [`PartitionedFile`] +/// via [`PartitionedFile::with_arrow_schema`]. +/// +/// By default DataFusion infers a file's physical schema by reading its +/// metadata (e.g. the Parquet footer) when the scan begins. When the schema is +/// already known, it can be supplied up front so this inference step is +/// skipped, saving an I/O round trip and metadata parse per file. +/// +/// The example writes a small Parquet file with a single `Int32` column `a` and +/// reads it back three ways: +/// - without a schema, letting DataFusion infer it at query time; +/// - with the correct schema, skipping inference; +/// - with a deliberately mismatched schema (`a` typed as `Int64`), which +/// surfaces as an error since the provided schema does not match the data +/// actually stored in the file. +/// +/// Note that the schema passed to [`PartitionedFile::with_arrow_schema`] must +/// describe only the columns physically stored in the file and must not include +/// partition columns. +pub async fn read_partitioned_file() -> Result<()> { + let tmpdir = TempDir::new()?; + let file_path = tmpdir.path().join("partitioned-file"); + let file = File::create(file_path.as_path())?; + + let file_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch = RecordBatch::try_new( + file_schema.clone(), + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], + )?; + let mut writer = ArrowWriter::try_new(&file, file_schema.clone(), None)?; + writer.write(&batch)?; + writer.finish()?; + + let table_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + // Specify another field in the table which is missing from the file schema. + // Illustrates that the table schema does not need to match the PartitionedFile schema + // for a scan to succeed. + Field::new("b", DataType::Float64, true), + ])); + + // Infer file schema at query time. + { + let batch = + read_file(file_path.as_path(), file.len(), table_schema.clone(), None) + .await?; + println!("{batch:?}"); + } + + // Provide the correct file schema to skip inferring at query time. + { + let batch = read_file( + file_path.as_path(), + file.len(), + table_schema.clone(), + Some(file_schema.clone()), + ) + .await?; + println!("{batch:?}"); + } + + // A mismatching file schema returns an error. + { + let mismatching_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let error = read_file( + file_path.as_path(), + file.len(), + table_schema.clone(), + Some(mismatching_schema), + ) + .await + .unwrap_err(); + println!("Got schema error: {error:?}"); + } + + Ok(()) +} + +/// Scans a single Parquet file with the given `source_schema`, optionally +/// supplying the file's Arrow schema to skip schema inference. A `None` +/// `file_schema` lets DataFusion infer the schema from the file metadata at +/// query time. +async fn read_file( + file_path: &Path, + file_len: u64, + source_schema: SchemaRef, + file_schema: Option, +) -> Result { + let mut partitioned_file = + PartitionedFile::new(file_path.to_string_lossy(), file_len); + if let Some(schema) = file_schema { + partitioned_file = partitioned_file.with_arrow_schema(schema); + } + + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::local_filesystem(), + Arc::new(ParquetSource::new(source_schema)), + ) + .with_file(partitioned_file) + .build(); + + let exec = DataSourceExec::from_data_source(config); + let mut result = exec.execute(0, Arc::new(TaskContext::default()))?; + result.next().await.ok_or_else(|| { + datafusion::error::DataFusionError::Internal( + "execution produced no batches".into(), + ) + })? +} diff --git a/datafusion-examples/examples/data_io/remote_catalog.rs b/datafusion-examples/examples/data_io/remote_catalog.rs index 16814752b3ec2..a24ca2238181d 100644 --- a/datafusion-examples/examples/data_io/remote_catalog.rs +++ b/datafusion-examples/examples/data_io/remote_catalog.rs @@ -130,6 +130,7 @@ struct RemoteCatalogInterface {} impl RemoteCatalogInterface { /// Establish a connection to the remote catalog + #[expect(clippy::unused_async)] pub async fn connect() -> Result { // In a real implementation this method might connect to a remote // catalog, validate credentials, cache basic information, etc @@ -137,6 +138,7 @@ impl RemoteCatalogInterface { } /// Fetches information for a specific table + #[expect(clippy::unused_async)] pub async fn table_info(&self, name: &str) -> Result> { if name != "remote_table" { return Ok(None); @@ -155,6 +157,7 @@ impl RemoteCatalogInterface { } /// Fetches data for a table from a remote data source + #[expect(clippy::unused_async)] pub async fn read_data(&self, name: &str) -> Result { if name != "remote_table" { return plan_err!("Remote table not found: {}", name); diff --git a/datafusion-examples/examples/dataframe/cache_factory.rs b/datafusion-examples/examples/dataframe/cache_factory.rs index a92c3dc4ce26a..ffbce298b4f17 100644 --- a/datafusion-examples/examples/dataframe/cache_factory.rs +++ b/datafusion-examples/examples/dataframe/cache_factory.rs @@ -23,12 +23,14 @@ use std::sync::{Arc, RwLock}; use arrow::array::RecordBatch; use async_trait::async_trait; +use datafusion::catalog::Session; use datafusion::catalog::memory::MemorySourceConfig; use datafusion::common::DFSchemaRef; use datafusion::error::Result; use datafusion::execution::context::QueryPlanner; use datafusion::execution::session_state::CacheFactory; use datafusion::execution::{SessionState, SessionStateBuilder}; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; use datafusion::logical_expr::{ Extension, LogicalPlan, UserDefinedLogicalNode, UserDefinedLogicalNodeCore, }; @@ -145,7 +147,8 @@ impl ExtensionPlanner for CacheNodePlanner { node: &dyn UserDefinedLogicalNode, logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], - session_state: &SessionState, + session_state: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { if let Some(cache_node) = node.as_any().downcast_ref::() { assert_eq!(logical_inputs.len(), 1, "Inconsistent number of inputs"); @@ -198,7 +201,7 @@ impl QueryPlanner for CacheNodeQueryPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { let physical_planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( diff --git a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs index dc374c7e02fe5..89eff74e0d730 100644 --- a/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs +++ b/datafusion-examples/examples/execution_monitoring/memory_pool_execution_plan.rs @@ -26,9 +26,9 @@ //! - Handle memory pressure by spilling to disk //! - Release memory when done +use arrow::array::record_batch; use arrow::record_batch::RecordBatch; use arrow_schema::SchemaRef; -use datafusion::common::record_batch; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::common::{exec_datafusion_err, internal_err}; use datafusion::datasource::{DefaultTableSource, memory::MemTable}; @@ -39,7 +39,8 @@ use datafusion::execution::{SendableRecordBatchStream, TaskContext}; use datafusion::logical_expr::LogicalPlanBuilder; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, }; use datafusion::prelude::*; use futures::stream::{StreamExt, TryStreamExt}; @@ -237,9 +238,10 @@ impl ExecutionPlan for BufferingExecutionPlan { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _options: ReplaceChildrenOptions, ) -> Result> { if children.len() == 1 { Ok(Arc::new(BufferingExecutionPlan::new( @@ -251,6 +253,16 @@ impl ExecutionPlan for BufferingExecutionPlan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -295,17 +307,10 @@ impl ExecutionPlan for BufferingExecutionPlan { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.properties.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion-examples/examples/external_dependency/query_aws_s3.rs b/datafusion-examples/examples/external_dependency/query_aws_s3.rs index 63507bb3eed11..7dc2f76be4f0c 100644 --- a/datafusion-examples/examples/external_dependency/query_aws_s3.rs +++ b/datafusion-examples/examples/external_dependency/query_aws_s3.rs @@ -66,7 +66,7 @@ pub async fn query_aws_s3() -> Result<()> { // dynamic query by the file path let ctx = ctx.enable_url_table(); let df = ctx - .sql(format!(r#"SELECT * FROM '{}' LIMIT 10"#, &path).as_str()) + .sql(format!(r#"SELECT * FROM '{path}' LIMIT 10"#).as_str()) .await?; // print the results diff --git a/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs b/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs index 7894e97f3796d..29b04d0042547 100644 --- a/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs +++ b/datafusion-examples/examples/ffi/ffi_example_table_provider/src/lib.rs @@ -17,9 +17,10 @@ use std::sync::Arc; -use arrow::array::RecordBatch; +use arrow::array::{RecordBatch, record_batch}; +use arrow::datatypes as arrow_schema; use arrow::datatypes::{DataType, Field, Schema}; -use datafusion::{common::record_batch, datasource::MemTable}; +use datafusion::datasource::MemTable; use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; use datafusion_ffi::table_provider::FFI_TableProvider; use ffi_module_interface::TableProviderModule; diff --git a/datafusion-examples/examples/flight/server.rs b/datafusion-examples/examples/flight/server.rs index b73c81dd7d2c3..ac8908d7c820e 100644 --- a/datafusion-examples/examples/flight/server.rs +++ b/datafusion-examples/examples/flight/server.rs @@ -19,7 +19,7 @@ use std::sync::Arc; -use arrow::ipc::writer::{CompressionContext, DictionaryTracker, IpcDataGenerator}; +use arrow::ipc::writer::{DictionaryTracker, IpcDataGenerator, IpcWriteContext}; use arrow_flight::{ Action, ActionType, Criteria, Empty, FlightData, FlightDescriptor, FlightInfo, HandshakeRequest, HandshakeResponse, PutResult, SchemaResult, Ticket, @@ -112,7 +112,7 @@ impl FlightService for FlightServiceImpl { // add an initial FlightData message that sends schema let options = arrow::ipc::writer::IpcWriteOptions::default(); - let mut compression_context = CompressionContext::default(); + let mut compression_context = IpcWriteContext::default(); let schema_flight_data = SchemaAsIpc::new(&schema, &options); let mut flights = vec![FlightData::from(schema_flight_data)]; diff --git a/datafusion-examples/examples/proto/composed_extension_codec.rs b/datafusion-examples/examples/proto/composed_extension_codec.rs index ae9503dd87b19..51c1bc7c5518b 100644 --- a/datafusion-examples/examples/proto/composed_extension_codec.rs +++ b/datafusion-examples/examples/proto/composed_extension_codec.rs @@ -39,16 +39,18 @@ use datafusion::common::Result; use datafusion::common::internal_err; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::execution::TaskContext; +use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion::physical_plan::{DisplayAs, ExecutionPlan}; use datafusion::prelude::SessionContext; use datafusion_proto::physical_plan::{ AsExecutionPlan, ComposedPhysicalExtensionCodec, PhysicalExtensionCodec, + PhysicalProtoConverterExtension, }; use datafusion_proto::protobuf; /// Example of using multiple extension codecs for serialization / deserialization -pub async fn composed_extension_codec() -> Result<()> { - // build execution plan that has both types of nodes +pub fn composed_extension_codec() -> Result<()> { + // Build execution plan that has both types of nodes // // Note each node requires a different `PhysicalExtensionCodec` to decode let exec_plan = Arc::new(ParentExec { @@ -63,18 +65,18 @@ pub async fn composed_extension_codec() -> Result<()> { Arc::new(ChildPhysicalExtensionCodec {}), ]); - // serialize execution plan to proto + // Serialize execution plan to proto let proto: protobuf::PhysicalPlanNode = protobuf::PhysicalPlanNode::try_from_physical_plan( exec_plan.clone(), &composed_codec, )?; - // deserialize proto back to execution plan + // Deserialize proto back to execution plan let result_exec_plan: Arc = proto.try_into_physical_plan(&ctx.task_ctx(), &composed_codec)?; - // assert that the original and deserialized execution plans are equal + // Assert that the original and deserialized execution plans are equal assert_eq!(format!("{exec_plan:?}"), format!("{result_exec_plan:?}")); Ok(()) @@ -110,13 +112,24 @@ impl ExecutionPlan for ParentExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unreachable!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -128,7 +141,7 @@ impl ExecutionPlan for ParentExec { fn apply_expressions( &self, _f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) @@ -145,6 +158,7 @@ impl PhysicalExtensionCodec for ParentPhysicalExtensionCodec { buf: &[u8], inputs: &[Arc], _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { if buf == "ParentExec".as_bytes() { Ok(Arc::new(ParentExec { @@ -155,7 +169,12 @@ impl PhysicalExtensionCodec for ParentPhysicalExtensionCodec { } } - fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { if node.is::() { buf.extend_from_slice("ParentExec".as_bytes()); Ok(()) @@ -191,13 +210,24 @@ impl ExecutionPlan for ChildExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unreachable!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -209,7 +239,7 @@ impl ExecutionPlan for ChildExec { fn apply_expressions( &self, _f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) @@ -226,6 +256,7 @@ impl PhysicalExtensionCodec for ChildPhysicalExtensionCodec { buf: &[u8], _inputs: &[Arc], _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { if buf == "ChildExec".as_bytes() { Ok(Arc::new(ChildExec {})) @@ -234,7 +265,12 @@ impl PhysicalExtensionCodec for ChildPhysicalExtensionCodec { } } - fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { if node.is::() { buf.extend_from_slice("ChildExec".as_bytes()); Ok(()) diff --git a/datafusion-examples/examples/proto/expression_deduplication.rs b/datafusion-examples/examples/proto/expression_deduplication.rs index 26d246b2efca8..8ee59fa14d9cd 100644 --- a/datafusion-examples/examples/proto/expression_deduplication.rs +++ b/datafusion-examples/examples/proto/expression_deduplication.rs @@ -52,7 +52,7 @@ use datafusion_proto::physical_plan::from_proto::parse_physical_expr_with_conver use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter; use datafusion_proto::physical_plan::{ DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, PhysicalPlanDecodeContext, - PhysicalProtoConverterExtension, + PhysicalPlanNodeExt, PhysicalProtoConverterExtension, }; use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use prost::Message; @@ -72,7 +72,7 @@ use prost::Message; /// In real scenarios, expressions can be much more complex, e.g. a large InList /// expression could be megabytes in size, so deduplication can save significant memory /// in addition to more correctly representing the original plan structure. -pub async fn expression_deduplication() -> Result<()> { +pub fn expression_deduplication() -> Result<()> { println!("=== Expression Deduplication Example ===\n"); // Create a schema for our test expressions @@ -187,6 +187,7 @@ impl PhysicalExtensionCodec for CachingCodec { _buf: &[u8], _inputs: &[Arc], _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { datafusion::common::not_impl_err!("No custom extension nodes") } @@ -196,6 +197,7 @@ impl PhysicalExtensionCodec for CachingCodec { &self, _node: Arc, _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { datafusion::common::not_impl_err!("No custom extension nodes") } diff --git a/datafusion-examples/examples/proto/main.rs b/datafusion-examples/examples/proto/main.rs index 3f525b5d46afa..d534eda24ba64 100644 --- a/datafusion-examples/examples/proto/main.rs +++ b/datafusion-examples/examples/proto/main.rs @@ -64,10 +64,10 @@ impl ExampleKind { } } ExampleKind::ComposedExtensionCodec => { - composed_extension_codec::composed_extension_codec().await? + composed_extension_codec::composed_extension_codec()? } ExampleKind::ExpressionDeduplication => { - expression_deduplication::expression_deduplication().await? + expression_deduplication::expression_deduplication()? } } Ok(()) diff --git a/datafusion-examples/examples/query_planning/expr_api.rs b/datafusion-examples/examples/query_planning/expr_api.rs index c087019c687c5..08efff7777691 100644 --- a/datafusion-examples/examples/query_planning/expr_api.rs +++ b/datafusion-examples/examples/query_planning/expr_api.rs @@ -33,6 +33,7 @@ use datafusion::functions_aggregate::first_last::first_value_udaf; use datafusion::logical_expr::execution_props::ExecutionProps; use datafusion::logical_expr::expr::BinaryExpr; use datafusion::logical_expr::interval_arithmetic::Interval; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; use datafusion::logical_expr::simplify::SimplifyContext; use datafusion::logical_expr::{ColumnarValue, ExprFunctionExt, ExprSchemable, Operator}; use datafusion::optimizer::analyzer::type_coercion::TypeCoercionRewriter; @@ -57,7 +58,7 @@ use datafusion::prelude::*; /// 5. Analyze predicates for boundary ranges: [`range_analysis_demo`] /// 6. Get the types of the expressions: [`expression_type_demo`] /// 7. Apply type coercion to expressions: [`type_coercion_demo`] -pub async fn expr_api() -> Result<()> { +pub fn expr_api() -> Result<()> { // The easiest way to do create expressions is to use the // "fluent"-style API: let expr = col("a") + lit(5); @@ -541,8 +542,12 @@ fn type_coercion_demo() -> Result<()> { // Evaluation with an expression that has not been type coerced cannot succeed. let props = ExecutionProps::default(); - let physical_expr = - datafusion::physical_expr::create_physical_expr(&expr, &df_schema, &props)?; + let physical_expr = datafusion::physical_expr::create_physical_expr( + &expr, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + )?; let e = physical_expr.evaluate(&batch).unwrap_err(); assert!( e.find_root() @@ -566,6 +571,7 @@ fn type_coercion_demo() -> Result<()> { &coerced_expr, &df_schema, &props, + &PhysicalPlanningContext::default(), )?; assert!(physical_expr.evaluate(&batch).is_ok()); @@ -578,6 +584,7 @@ fn type_coercion_demo() -> Result<()> { &coerced_expr, &df_schema, &props, + &PhysicalPlanningContext::default(), )?; assert!(physical_expr.evaluate(&batch).is_ok()); @@ -606,6 +613,7 @@ fn type_coercion_demo() -> Result<()> { &coerced_expr, &df_schema, &props, + &PhysicalPlanningContext::default(), )?; assert!(physical_expr.evaluate(&batch).is_ok()); diff --git a/datafusion-examples/examples/query_planning/main.rs b/datafusion-examples/examples/query_planning/main.rs index d3f99aedceb3d..2e4310082c9dd 100644 --- a/datafusion-examples/examples/query_planning/main.rs +++ b/datafusion-examples/examples/query_planning/main.rs @@ -94,12 +94,12 @@ impl ExampleKind { } } ExampleKind::AnalyzerRule => analyzer_rule::analyzer_rule().await?, - ExampleKind::ExprApi => expr_api::expr_api().await?, + ExampleKind::ExprApi => expr_api::expr_api()?, ExampleKind::OptimizerRule => optimizer_rule::optimizer_rule().await?, ExampleKind::ParseSqlExpr => parse_sql_expr::parse_sql_expr().await?, ExampleKind::PlanToSql => plan_to_sql::plan_to_sql_examples().await?, ExampleKind::PlannerApi => planner_api::planner_api().await?, - ExampleKind::Pruning => pruning::pruning().await?, + ExampleKind::Pruning => pruning::pruning()?, ExampleKind::ThreadPools => thread_pools::thread_pools().await?, } Ok(()) diff --git a/datafusion-examples/examples/query_planning/pruning.rs b/datafusion-examples/examples/query_planning/pruning.rs index 7fdc4a7952d68..023058f825f64 100644 --- a/datafusion-examples/examples/query_planning/pruning.rs +++ b/datafusion-examples/examples/query_planning/pruning.rs @@ -26,8 +26,11 @@ use datafusion::common::pruning::PruningStatistics; use datafusion::common::{DFSchema, ScalarValue}; use datafusion::error::Result; use datafusion::execution::context::ExecutionProps; +use datafusion::logical_expr::physical_planning_context::PhysicalPlanningContext; use datafusion::physical_expr::create_physical_expr; -use datafusion::physical_optimizer::pruning::PruningPredicate; +use datafusion::physical_optimizer::pruning::{ + PruningPredicate, PruningPredicateBuilder, +}; use datafusion::prelude::*; /// This example shows how to use DataFusion's `PruningPredicate` to prove @@ -43,7 +46,7 @@ use datafusion::prelude::*; /// one might do as part of a higher level storage engine. See /// `parquet_index.rs` for an example that uses pruning in the context of an /// individual query. -pub async fn pruning() -> Result<()> { +pub fn pruning() -> Result<()> { // In this example, we'll use the PruningPredicate to determine if // the expression `x = 5 AND y = 10` can never be true based on statistics @@ -194,8 +197,17 @@ impl PruningStatistics for MyCatalog { fn create_pruning_predicate(expr: Expr, schema: &SchemaRef) -> PruningPredicate { let df_schema = DFSchema::try_from(Arc::clone(schema)).unwrap(); let props = ExecutionProps::new(); - let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap(); - PruningPredicate::try_new(physical_expr, Arc::clone(schema)).unwrap() + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &props, + &PhysicalPlanningContext::default(), + ) + .unwrap(); + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(schema)) + .try_build(physical_expr) + .unwrap() } fn i32_array<'a>(values: impl Iterator>) -> ArrayRef { diff --git a/datafusion-examples/examples/relation_planner/table_sample.rs b/datafusion-examples/examples/relation_planner/table_sample.rs index 42342e5f1a641..7a8f533ac3a9b 100644 --- a/datafusion-examples/examples/relation_planner/table_sample.rs +++ b/datafusion-examples/examples/relation_planner/table_sample.rs @@ -100,24 +100,30 @@ use futures::{ use rand::{Rng, SeedableRng, rngs::StdRng}; use tonic::async_trait; -use datafusion::optimizer::simplify_expressions::simplify_literal::parse_literal; use datafusion::{ + catalog::Session, execution::{ - RecordBatchStream, SendableRecordBatchStream, SessionState, SessionStateBuilder, - TaskContext, context::QueryPlanner, + RecordBatchStream, SendableRecordBatchStream, SessionStateBuilder, TaskContext, + context::QueryPlanner, }, physical_expr::EquivalenceProperties, physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildStats, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + StatisticsArgs, metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput}, }, physical_planner::{DefaultPhysicalPlanner, ExtensionPlanner, PhysicalPlanner}, prelude::*, }; +use datafusion::{ + optimizer::simplify_expressions::simplify_literal::parse_literal, + physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}, +}; use datafusion_common::{ DFSchemaRef, DataFusionError, Result, Statistics, internal_err, not_impl_err, plan_datafusion_err, plan_err, tree_node::TreeNodeRecursion, }; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ UserDefinedLogicalNode, UserDefinedLogicalNodeCore, logical_plan::{Extension, LogicalPlan, LogicalPlanBuilder}, @@ -563,7 +569,7 @@ impl QueryPlanner for TableSampleQueryPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( TableSampleExtensionPlanner, @@ -585,7 +591,8 @@ impl ExtensionPlanner for TableSampleExtensionPlanner { node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], - _session_state: &SessionState, + _session_state: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { let Some(sample_node) = node.as_any().downcast_ref::() else { @@ -694,9 +701,10 @@ impl ExecutionPlan for SampleExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self::try_new( children.swap_remove(0), @@ -706,6 +714,16 @@ impl ExecutionPlan for SampleExec { )?)) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -722,8 +740,16 @@ impl ExecutionPlan for SampleExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let mut stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?); + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let mut stats = input_stats[0].as_ref().clone(); let ratio = self.upper_bound - self.lower_bound; // Scale statistics by sampling ratio (inexact due to randomness) @@ -741,18 +767,11 @@ impl ExecutionPlan for SampleExec { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion-examples/examples/sql_ops/frontend.rs b/datafusion-examples/examples/sql_ops/frontend.rs index b34c720a78198..27eb97ee7ab25 100644 --- a/datafusion-examples/examples/sql_ops/frontend.rs +++ b/datafusion-examples/examples/sql_ops/frontend.rs @@ -154,7 +154,7 @@ impl ContextProvider for MyContextProvider { None } - fn get_higher_order_meta(&self, _name: &str) -> Option> { + fn get_higher_order_meta(&self, _name: &str) -> Option> { None } diff --git a/datafusion-examples/examples/udf/advanced_udaf.rs b/datafusion-examples/examples/udf/advanced_udaf.rs index f1651dbf28913..bca4c7edab2c5 100644 --- a/datafusion-examples/examples/udf/advanced_udaf.rs +++ b/datafusion-examples/examples/udf/advanced_udaf.rs @@ -23,8 +23,10 @@ use datafusion::{arrow::datatypes::DataType, logical_expr::Volatility}; use std::sync::Arc; use arrow::array::{ - ArrayRef, AsArray, Float32Array, PrimitiveArray, PrimitiveBuilder, UInt32Array, + Array, ArrayRef, AsArray, BooleanArray, Float32Array, PrimitiveArray, + PrimitiveBuilder, UInt32Array, }; +use arrow::buffer::NullBuffer; use arrow::datatypes::{ArrowNativeTypeOp, ArrowPrimitiveType, Float64Type, UInt32Type}; use arrow::record_batch::RecordBatch; use arrow_schema::FieldRef; @@ -237,7 +239,7 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&arrow::array::BooleanArray>, + opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "single argument to update_batch"); @@ -268,7 +270,6 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&arrow::array::BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 2, "two arguments to merge_batch"); @@ -280,7 +281,7 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator { self.null_state.accumulate( group_indices, partial_counts, - opt_filter, + None, total_num_groups, |group_index, partial_count| { self.counts[group_index] += partial_count; @@ -292,7 +293,7 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator { self.null_state.accumulate( group_indices, partial_prods, - opt_filter, + None, total_num_groups, |group_index, new_value: ::Native| { let prod = &mut self.prods[group_index]; @@ -360,6 +361,38 @@ impl GroupsAccumulator for GeometricMeanGroupsAccumulator { ]) } + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + assert_eq!(values.len(), 1, "single argument to convert_to_state"); + + let prods = values[0] + .as_primitive::() + .clone() + .with_data_type(self.prod_data_type.clone()); + let counts = UInt32Array::from_value(1, prods.len()); + + let filter_nulls = opt_filter.map(|filter| { + let validity = match filter.nulls() { + Some(nulls) => filter.values() & nulls.inner(), + None => filter.values().clone(), + }; + NullBuffer::new(validity) + }); + let nulls = NullBuffer::union(filter_nulls.as_ref(), prods.nulls()); + + let prods = + PrimitiveArray::::new(prods.values().clone(), nulls.clone()) + .with_data_type(self.prod_data_type.clone()); + let counts = UInt32Array::new(counts.values().clone(), nulls); + + Ok(vec![ + Arc::new(prods) as ArrayRef, + Arc::new(counts) as ArrayRef, + ]) + } fn size(&self) -> usize { self.counts.capacity() * size_of::() + self.prods.capacity() * size_of::() diff --git a/datafusion-examples/examples/udf/main.rs b/datafusion-examples/examples/udf/main.rs index 89f3fd801deec..0eff5f7a30a2c 100644 --- a/datafusion-examples/examples/udf/main.rs +++ b/datafusion-examples/examples/udf/main.rs @@ -39,6 +39,9 @@ //! - `async_udf` //! (file: async_udf.rs, desc: Asynchronous User Defined Scalar Function) //! +//! - `struct_udaf` +//! (file: struct_returning_udaf.rs, desc: Struct-returning UDAF with window metadata) +//! //! - `udaf` //! (file: simple_udaf.rs, desc: Simple UDAF example) //! @@ -62,6 +65,7 @@ mod simple_udaf; mod simple_udf; mod simple_udtf; mod simple_udwf; +mod struct_returning_udaf; mod table_list_udtf; use datafusion::error::{DataFusionError, Result}; @@ -76,6 +80,7 @@ enum ExampleKind { AdvUdf, AdvUdwf, AsyncUdf, + StructUdaf, Udf, Udaf, Udwf, @@ -102,6 +107,9 @@ impl ExampleKind { ExampleKind::AdvUdf => advanced_udf::advanced_udf().await?, ExampleKind::AdvUdwf => advanced_udwf::advanced_udwf().await?, ExampleKind::AsyncUdf => async_udf::async_udf().await?, + ExampleKind::StructUdaf => { + struct_returning_udaf::struct_returning_udaf().await? + } ExampleKind::Udaf => simple_udaf::simple_udaf().await?, ExampleKind::Udf => simple_udf::simple_udf().await?, ExampleKind::Udtf => simple_udtf::simple_udtf().await?, diff --git a/datafusion-examples/examples/udf/struct_returning_udaf.rs b/datafusion-examples/examples/udf/struct_returning_udaf.rs new file mode 100644 index 0000000000000..5bb32b9ef28a3 --- /dev/null +++ b/datafusion-examples/examples/udf/struct_returning_udaf.rs @@ -0,0 +1,280 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! See `main.rs` for how to run it. +//! +//! This example shows how an extension can return window metadata from an +//! aggregate by passing the relevant input columns directly to the aggregate. + +use std::sync::Arc; + +use arrow::array::{ + ArrayRef, Float64Array, StructArray, TimestampNanosecondArray, UInt64Array, +}; +use arrow::datatypes::{DataType, Field, Fields, Schema, TimeUnit}; +use arrow::record_batch::RecordBatch; +use datafusion::assert_batches_eq; +use datafusion::common::{cast::as_primitive_array, exec_err}; +use datafusion::datasource::MemTable; +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::{AccumulatorFactoryFunction, Volatility, create_udaf}; +use datafusion::physical_plan::Accumulator; +use datafusion::prelude::*; +use datafusion::scalar::ScalarValue; + +pub async fn struct_returning_udaf() -> Result<()> { + let ctx = create_context()?; + + register_augmented_avg(&ctx); + + // The `augmented_avg` aggregate returns both the average and metadata about + // the time window from which the average was computed. + let sql = " + SELECT + augmented_avg(time, value)['window_start'] AS window_start, + augmented_avg(time, value)['window_end'] AS window_end, + augmented_avg(time, value)['window_duration'] AS window_duration, + augmented_avg(time, value)['avg_value'] AS avg_value + FROM t + GROUP BY date_bin(INTERVAL '5 microseconds', time) + ORDER BY window_start + "; + + let results = ctx.sql(sql).await?.collect().await?; + let expected = [ + "+----------------------------+----------------------------+-----------------+-----------+", + "| window_start | window_end | window_duration | avg_value |", + "+----------------------------+----------------------------+-----------------+-----------+", + "| 1970-01-01T00:00:00.000001 | 1970-01-01T00:00:00.000002 | 1000 | 15.0 |", + "| 1970-01-01T00:00:00.000005 | 1970-01-01T00:00:00.000009 | 4000 | 3.0 |", + "+----------------------------+----------------------------+-----------------+-----------+", + ]; + assert_batches_eq!(expected, &results); + + println!("Struct-returning aggregate produced window metadata:"); + ctx.sql(sql).await?.show().await?; + + Ok(()) +} + +fn create_context() -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new( + "time", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + ), + Field::new("value", DataType::Float64, false), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(TimestampNanosecondArray::from(vec![ + 1000, 2000, 5000, 7000, 9000, + ])) as ArrayRef, + Arc::new(Float64Array::from(vec![10.0, 20.0, 1.0, 3.0, 5.0])), + ], + )?; + + let ctx = SessionContext::new(); + let provider = MemTable::try_new(schema, vec![vec![batch]])?; + ctx.register_table("t", Arc::new(provider))?; + Ok(ctx) +} + +fn register_augmented_avg(ctx: &SessionContext) { + let accumulator: AccumulatorFactoryFunction = + Arc::new(|_| Ok(Box::new(AugmentedAvg::new()))); + + let augmented_avg = create_udaf( + "augmented_avg", + vec![ + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Float64, + ], + Arc::new(AugmentedAvg::output_datatype()), + Volatility::Immutable, + accumulator, + Arc::new(AugmentedAvg::state_datatypes()), + ); + + ctx.register_udaf(augmented_avg); +} + +#[derive(Debug, Clone)] +struct AugmentedAvg { + window_start: Option, + window_end: Option, + sum: f64, + count: u64, +} + +impl AugmentedAvg { + fn new() -> Self { + Self { + window_start: None, + window_end: None, + sum: 0.0, + count: 0, + } + } + + fn fields() -> Fields { + vec![ + Field::new( + "window_start", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + ), + Field::new( + "window_end", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + ), + Field::new("window_duration", DataType::Int64, true), + Field::new("avg_value", DataType::Float64, true), + ] + .into() + } + + fn output_datatype() -> DataType { + DataType::Struct(Self::fields()) + } + + fn state_datatypes() -> Vec { + vec![ + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Timestamp(TimeUnit::Nanosecond, None), + DataType::Float64, + DataType::UInt64, + ] + } + + fn update_one(&mut self, time: i64, value: f64) { + self.window_start = Some(self.window_start.map_or(time, |start| start.min(time))); + self.window_end = Some(self.window_end.map_or(time, |end| end.max(time))); + self.sum += value; + self.count += 1; + } +} + +impl Accumulator for AugmentedAvg { + fn state(&mut self) -> Result> { + // DataFusion can merge partial aggregate results across execution + // stages, so all values needed to reconstruct the final struct are + // included in the state. + Ok(vec![ + ScalarValue::TimestampNanosecond(self.window_start, None), + ScalarValue::TimestampNanosecond(self.window_end, None), + ScalarValue::Float64(Some(self.sum)), + ScalarValue::UInt64(Some(self.count)), + ]) + } + + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let [times, values] = values else { + return exec_err!("augmented_avg expects time and value arrays"); + }; + let times = + as_primitive_array::(times)?; + let values = as_primitive_array::(values)?; + + // Track the window bounds and aggregate values directly from the input + // rows assigned to each group by `date_bin`. + for (time, value) in times.iter().zip(values.iter()) { + if let (Some(time), Some(value)) = (time, value) { + self.update_one(time, value); + } + } + + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let [starts, ends, sums, counts] = states else { + return exec_err!("augmented_avg expects four state arrays"); + }; + let starts = + as_primitive_array::(starts)?; + let ends = as_primitive_array::(ends)?; + let sums = as_primitive_array::(sums)?; + let counts = counts + .as_any() + .downcast_ref::() + .ok_or_else(|| { + DataFusionError::Execution("Expected UInt64Array".to_string()) + })?; + + // Combine partial states by preserving the earliest start, latest end, + // and additive average components. + for (((start, end), sum), count) in starts + .iter() + .zip(ends.iter()) + .zip(sums.iter()) + .zip(counts.iter()) + { + let Some(count) = count else { + continue; + }; + if count == 0 { + continue; + } + if let (Some(start), Some(end), Some(sum)) = (start, end, sum) { + self.window_start = Some( + self.window_start + .map_or(start, |current| current.min(start)), + ); + self.window_end = + Some(self.window_end.map_or(end, |current| current.max(end))); + self.sum += sum; + self.count += count; + } + } + + Ok(()) + } + + fn evaluate(&mut self) -> Result { + let duration = self + .window_start + .zip(self.window_end) + .map(|(start, end)| end - start); + let avg = (self.count > 0).then_some(self.sum / self.count as f64); + + // Return one Struct scalar whose fields can be projected from SQL with + // expressions like `augmented_avg(time, value)['window_start']`. + let struct_array = StructArray::try_new( + AugmentedAvg::fields(), + vec![ + Arc::new(TimestampNanosecondArray::from(vec![self.window_start])) + as ArrayRef, + Arc::new(TimestampNanosecondArray::from(vec![self.window_end])) + as ArrayRef, + Arc::new(arrow::array::Int64Array::from(vec![duration])) as ArrayRef, + Arc::new(Float64Array::from(vec![avg])) as ArrayRef, + ], + None, + )?; + + Ok(ScalarValue::Struct(Arc::new(struct_array))) + } + + fn size(&self) -> usize { + size_of_val(self) + } +} diff --git a/datafusion-testing b/datafusion-testing index 7833a65d5b08b..13bbae38776c2 160000 --- a/datafusion-testing +++ b/datafusion-testing @@ -1 +1 @@ -Subproject commit 7833a65d5b08be2ca484ea938f471cf01df54e18 +Subproject commit 13bbae38776c2bfbc1fab1be7e7220222d4284bf diff --git a/datafusion/catalog-listing/Cargo.toml b/datafusion/catalog-listing/Cargo.toml index 61b55397137df..abe58f45994be 100644 --- a/datafusion/catalog-listing/Cargo.toml +++ b/datafusion/catalog-listing/Cargo.toml @@ -46,6 +46,7 @@ futures = { workspace = true } itertools = { workspace = true } log = { workspace = true } object_store = { workspace = true } +percent-encoding = { workspace = true } [dev-dependencies] chrono = { workspace = true } diff --git a/datafusion/catalog-listing/src/config.rs b/datafusion/catalog-listing/src/config.rs index ca4d2abfcd737..2b83c8ec92b2c 100644 --- a/datafusion/catalog-listing/src/config.rs +++ b/datafusion/catalog-listing/src/config.rs @@ -152,8 +152,7 @@ impl ListingTableConfig { /// # use datafusion_datasource_parquet::file_format::ParquetFormat; /// # let table_paths = ListingTableUrl::parse("file:///path/to/data").unwrap(); /// let options = ListingOptions::new(Arc::new(ParquetFormat::default())) - /// .with_file_extension(".parquet") - /// .with_collect_stat(true); + /// .with_file_extension(".parquet"); /// /// let config = ListingTableConfig::new(table_paths).with_listing_options(options); /// // Configure file format and options diff --git a/datafusion/catalog-listing/src/helpers.rs b/datafusion/catalog-listing/src/helpers.rs index 0389b3cb17fe9..098f3d51ef911 100644 --- a/datafusion/catalog-listing/src/helpers.rs +++ b/datafusion/catalog-listing/src/helpers.rs @@ -17,7 +17,7 @@ //! Helper functions for the table implementation -use std::mem; +use std::borrow::Cow; use std::sync::Arc; use datafusion_catalog::Session; @@ -34,6 +34,7 @@ use arrow::{ record_batch::RecordBatch, }; use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use futures::stream::FuturesUnordered; use futures::{StreamExt, TryStreamExt, stream::BoxStream}; use log::{debug, trace}; @@ -44,6 +45,10 @@ use datafusion_expr::{Expr, Volatility}; use datafusion_physical_expr::create_physical_expr; use object_store::path::Path; use object_store::{ObjectMeta, ObjectStore}; +use percent_encoding::{AsciiSet, CONTROLS, percent_decode_str, utf8_percent_encode}; + +const PARTITION_VALUE_ENCODE_SET: &AsciiSet = + &CONTROLS.add(b' ').add(b'%').add(b'/').add(b'?').add(b'#'); /// Check whether the given expression can be resolved using only the columns `col_names`. /// This means that if this function returns true: @@ -137,41 +142,6 @@ pub fn expr_applicable_for_cols(col_names: &[&str], expr: &Expr) -> bool { /// The maximum number of concurrent listing requests const CONCURRENCY_LIMIT: usize = 100; -/// Partition the list of files into `n` groups -#[deprecated(since = "47.0.0", note = "use `FileGroup::split_files` instead")] -pub fn split_files( - mut partitioned_files: Vec, - n: usize, -) -> Vec> { - if partitioned_files.is_empty() { - return vec![]; - } - - // ObjectStore::list does not guarantee any consistent order and for some - // implementations such as LocalFileSystem, it may be inconsistent. Thus - // Sort files by path to ensure consistent plans when run more than once. - partitioned_files.sort_by(|a, b| a.path().cmp(b.path())); - - // effectively this is div with rounding up instead of truncating - let chunk_size = partitioned_files.len().div_ceil(n); - let mut chunks = Vec::with_capacity(n); - let mut current_chunk = Vec::with_capacity(chunk_size); - for file in partitioned_files.drain(..) { - current_chunk.push(file); - if current_chunk.len() == chunk_size { - let full_chunk = - mem::replace(&mut current_chunk, Vec::with_capacity(chunk_size)); - chunks.push(full_chunk); - } - } - - if !current_chunk.is_empty() { - chunks.push(current_chunk) - } - - chunks -} - #[derive(Debug)] pub struct Partition { /// The path to the partition, including the table prefix @@ -308,7 +278,16 @@ pub fn evaluate_partition_prefix<'a>( Some(PartitionValue::Single(val)) => { // if a partition only has a single literal value, then it can be added to the // prefix - parts.push(format!("{p}={val}")); + let encoded = encode_partition_value(val); + if encoded != val.as_str() { + // The same decoded value can be represented by both raw and + // percent-encoded partition directories. Prefix pruning is + // an optimization, so stop before this partition rather + // than listing only one spelling and potentially skipping + // valid rows. + break; + } + parts.push(format!("{p}={encoded}")); } _ => { // break on the first unconstrainted partition to create a common prefix @@ -325,7 +304,11 @@ pub fn evaluate_partition_prefix<'a>( } } -fn filter_partitions( +fn encode_partition_value(value: &str) -> Cow<'_, str> { + utf8_percent_encode(value, PARTITION_VALUE_ENCODE_SET).into() +} + +pub fn filter_partitioned_file( pf: PartitionedFile, filters: &[Expr], df_schema: &DFSchema, @@ -346,7 +329,12 @@ fn filter_partitions( let filter = utils::conjunction(filters.iter().cloned()).unwrap_or_else(|| lit(true)); let props = ExecutionProps::new(); - let expr = create_physical_expr(&filter, df_schema, &props)?; + let expr = create_physical_expr( + &filter, + df_schema, + &props, + &PhysicalPlanningContext::default(), + )?; // Since we're only operating on a single file, our batch and resulting "array" holds only one // value indicating if the input file matches the provided filters @@ -379,7 +367,7 @@ fn try_into_partitioned_file( .into_iter() .zip(partition_cols) .map(|(parsed, (_, datatype))| { - ScalarValue::try_from_string(parsed.to_string(), datatype) + ScalarValue::try_from_string(parsed.into_owned(), datatype) }) .collect::>>()?; @@ -447,7 +435,7 @@ pub async fn pruned_partition_list<'a>( )) }) .try_filter_map(move |pf| { - futures::future::ready(filter_partitions(pf, filters, &df_schema)) + futures::future::ready(filter_partitioned_file(pf, filters, &df_schema)) }) .boxed()) } @@ -459,6 +447,7 @@ fn object_meta_to_partitioned_file( ) -> Result> { Ok(Some(PartitionedFile { object_meta, + arrow_schema: None, partition_values: vec![], range: None, statistics: None, @@ -470,12 +459,15 @@ fn object_meta_to_partitioned_file( } /// Extract the partition values for the given `file_path` (in the given `table_path`) -/// associated to the partitions defined by `table_partition_cols` +/// associated to the partitions defined by `table_partition_cols`. +/// +/// Partition values are percent-decoded to match Hive-style object-store paths +/// that encode special characters in path segments. pub fn parse_partitions_for_path<'a, I>( table_path: &ListingTableUrl, file_path: &'a Path, table_partition_cols: I, -) -> Option> +) -> Option>> where I: IntoIterator, { @@ -484,7 +476,13 @@ where let mut part_values = vec![]; for (part, expected_partition) in subpath.zip(table_partition_cols) { match part.split_once('=') { - Some((name, val)) if name == expected_partition => part_values.push(val), + Some((name, val)) if name == expected_partition => { + // Preserve the original value if percent-decoding produces invalid UTF-8. + let decoded = percent_decode_str(val) + .decode_utf8() + .unwrap_or(Cow::Borrowed(val)); + part_values.push(decoded); + } _ => { debug!( "Ignoring file: file_path='{file_path}', table_path='{table_path}', part='{part}', partition_col='{expected_partition}'", @@ -560,7 +558,7 @@ mod tests { #[test] fn test_parse_partitions_for_path() { assert_eq!( - Some(vec![]), + Some(vec![] as Vec>), parse_partitions_for_path( &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), &Path::from("bucket/mytable/file.csv"), @@ -584,15 +582,51 @@ mod tests { ) ); assert_eq!( - Some(vec!["v1"]), + Some(vec![Cow::Borrowed("v1")]), parse_partitions_for_path( &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), &Path::from("bucket/mytable/mypartition=v1/file.csv"), vec!["mypartition"] ) ); + for (path, column, expected) in [ + ( + "bucket/mytable/mypartition=v%2F1/file.csv", + "mypartition", + "v/1", + ), + ( + "bucket/mytable/name=John%20Doe/file.csv", + "name", + "John Doe", + ), + ( + "bucket/mytable/mypartition=test%20dir%2Ffile/file.csv", + "mypartition", + "test dir/file", + ), + ( + "bucket/mytable/mypartition=%C3%A9/file.csv", + "mypartition", + "é", + ), + ( + "bucket/mytable/mypartition=%FF/file.csv", + "mypartition", + "%FF", + ), + ] { + assert_eq!( + Some(vec![Cow::Borrowed(expected)]), + parse_partitions_for_path( + &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), + &Path::parse(path).unwrap(), + vec![column] + ) + ); + } assert_eq!( - Some(vec!["v1"]), + Some(vec![Cow::Borrowed("v1")]), parse_partitions_for_path( &ListingTableUrl::parse("file:///bucket/mytable/").unwrap(), &Path::from("bucket/mytable/mypartition=v1/file.csv"), @@ -609,7 +643,7 @@ mod tests { ) ); assert_eq!( - Some(vec!["v1", "v2"]), + Some(vec![Cow::Borrowed("v1"), Cow::Borrowed("v2")]), parse_partitions_for_path( &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), &Path::from("bucket/mytable/mypartition=v1/otherpartition=v2/file.csv"), @@ -617,7 +651,7 @@ mod tests { ) ); assert_eq!( - Some(vec!["v1"]), + Some(vec![Cow::Borrowed("v1")]), parse_partitions_for_path( &ListingTableUrl::parse("file:///bucket/mytable").unwrap(), &Path::from("bucket/mytable/mypartition=v1/otherpartition=v2/file.csv"), @@ -649,6 +683,32 @@ mod tests { ); } + #[test] + fn test_try_into_partitioned_file_decodes_partition_value() { + let table_path = ListingTableUrl::parse("file:///bucket/mytable").unwrap(); + let partition_cols = vec![("category".to_string(), DataType::Utf8)]; + let meta = ObjectMeta { + location: Path::parse( + "bucket/mytable/category=Electronics%2FComputers/data.parquet", + ) + .unwrap(), + last_modified: chrono::Utc::now(), + size: 100, + e_tag: None, + version: None, + }; + + let result = + try_into_partitioned_file(meta, &partition_cols, &table_path).unwrap(); + assert!(result.is_some()); + let pf = result.unwrap(); + assert_eq!(pf.partition_values.len(), 1); + assert_eq!( + pf.partition_values[0], + ScalarValue::Utf8(Some("Electronics/Computers".to_string())) + ); + } + #[test] fn test_try_into_partitioned_file_root_file_skipped() { // File in root directory (not inside any partition path) should be @@ -803,6 +863,27 @@ mod tests { Some(Path::from("a=foo")), ); + assert_eq!( + evaluate_partition_prefix( + partitions, + &[col("a").eq(lit("Electronics/Computers"))], + ), + None, + ); + + assert_eq!( + evaluate_partition_prefix(partitions, &[col("a").eq(lit("John Doe"))]), + None, + ); + + assert_eq!( + evaluate_partition_prefix( + partitions, + &[col("a").eq(lit("foo")).and(col("b").eq(lit("John Doe")))], + ), + Some(Path::from("a=foo")), + ); + assert_eq!( evaluate_partition_prefix( partitions, diff --git a/datafusion/catalog-listing/src/options.rs b/datafusion/catalog-listing/src/options.rs index 0ab15e05abba1..44337e52a1e05 100644 --- a/datafusion/catalog-listing/src/options.rs +++ b/datafusion/catalog-listing/src/options.rs @@ -20,10 +20,10 @@ use datafusion_catalog::Session; use datafusion_common::plan_err; use datafusion_datasource::ListingTableUrl; use datafusion_datasource::file_format::FileFormat; -use datafusion_execution::config::SessionConfig; -use datafusion_expr::SortExpr; +use datafusion_expr::{Partitioning, SortExpr}; use futures::StreamExt; use futures::TryStreamExt; +use itertools::AllEqualValueError; use itertools::Itertools; use std::sync::Arc; @@ -38,13 +38,6 @@ pub struct ListingOptions { /// The expected partition column names in the folder structure. /// See [Self::with_table_partition_cols] for details pub table_partition_cols: Vec<(String, DataType)>, - /// Set true to try to guess statistics from the files. - /// This can add a lot of overhead as it will usually require files - /// to be opened and at least partially parsed. - pub collect_stat: bool, - /// Group files to avoid that the number of partitions exceeds - /// this limit - pub target_partitions: usize, /// Optional pre-known sort order(s). Must be `SortExpr`s. /// /// DataFusion may take advantage of this ordering to omit sorts @@ -61,6 +54,46 @@ pub struct ListingOptions { /// multiple equivalent orderings, the outer `Vec` will have a /// single element. pub file_sort_order: Vec>, + /// Declared output partitioning for scans from this table. + /// + /// Expressions are logical expressions over the full table schema. When set, + /// [`ListingTable`](crate::ListingTable) creates one file group per + /// declared output partition. When unset, file grouping uses the scan-time + /// [`SessionConfig::target_partitions`](datafusion_execution::config::SessionConfig::target_partitions). + /// + /// Files are listed in path order, split into whole-file groups across the + /// declared partition count, and then padded with trailing empty groups when + /// needed. DataFusion does not route files by partition values or validate + /// row placement, so callers must ensure file group `i` contains rows for + /// partition `i`. Layouts that require explicit file-to-partition assignment + /// are not supported. + /// + /// For example, range partitioning on column `a` with split points + /// `[10, 20, 30]` declares four output partitions. With three path-ordered + /// files, the trailing partition is preserved as empty: + /// + /// ```text + /// files in path order: f0, f1, f2 + /// + /// file groups: + /// partition 0: [f0] + /// partition 1: [f1] + /// partition 2: [f2] + /// partition 3: [] + /// ``` + /// + /// With five path-ordered files, a partition can contain multiple files: + /// + /// ```text + /// files in path order: f0, f1, f2, f3, f4 + /// + /// file groups: + /// partition 0: [f0, f1] + /// partition 1: [f2, f3] + /// partition 2: [f4] + /// partition 3: [] + /// ``` + pub output_partitioning: Option, } impl ListingOptions { @@ -68,30 +101,16 @@ impl ListingOptions { /// Default values: /// - use default file extension filter /// - no input partition to discover - /// - one target partition - /// - do not collect statistics pub fn new(format: Arc) -> Self { Self { file_extension: format.get_ext(), format, table_partition_cols: vec![], - collect_stat: false, - target_partitions: 1, file_sort_order: vec![], + output_partitioning: None, } } - /// Set options from [`SessionConfig`] and returns self. - /// - /// Currently this sets `target_partitions` and `collect_stat` - /// but if more options are added in the future that need to be coordinated - /// they will be synchronized through this method. - pub fn with_session_config_options(mut self, config: &SessionConfig) -> Self { - self = self.with_target_partitions(config.target_partitions()); - self = self.with_collect_stat(config.collect_statistics()); - self - } - /// Set file extension on [`ListingOptions`] and returns self. /// /// # Example @@ -136,6 +155,17 @@ impl ListingOptions { self } + /// Set declared output partitioning. + /// + /// See [`Self::output_partitioning`] for the contract. + pub fn with_output_partitioning( + mut self, + output_partitioning: Option, + ) -> Self { + self.output_partitioning = output_partitioning; + self + } + /// Set `table partition columns` on [`ListingOptions`] and returns self. /// /// "partition columns," used to support [Hive Partitioning], are @@ -205,40 +235,6 @@ impl ListingOptions { self } - /// Set stat collection on [`ListingOptions`] and returns self. - /// - /// ``` - /// # use std::sync::Arc; - /// # use datafusion_catalog_listing::ListingOptions; - /// # use datafusion_datasource_parquet::file_format::ParquetFormat; - /// - /// let listing_options = - /// ListingOptions::new(Arc::new(ParquetFormat::default())).with_collect_stat(true); - /// - /// assert_eq!(listing_options.collect_stat, true); - /// ``` - pub fn with_collect_stat(mut self, collect_stat: bool) -> Self { - self.collect_stat = collect_stat; - self - } - - /// Set number of target partitions on [`ListingOptions`] and returns self. - /// - /// ``` - /// # use std::sync::Arc; - /// # use datafusion_catalog_listing::ListingOptions; - /// # use datafusion_datasource_parquet::file_format::ParquetFormat; - /// - /// let listing_options = - /// ListingOptions::new(Arc::new(ParquetFormat::default())).with_target_partitions(8); - /// - /// assert_eq!(listing_options.target_partitions, 8); - /// ``` - pub fn with_target_partitions(mut self, target_partitions: usize) -> Self { - self.target_partitions = target_partitions; - self - } - /// Set file sort order on [`ListingOptions`] and returns self. /// /// ``` @@ -409,11 +405,10 @@ impl ListingOptions { match partition_keys.into_iter().all_equal_value() { Ok(v) => Ok(v), - Err(None) => Ok(vec![]), - Err(Some(diff)) => { - let mut sorted_diff = [diff.0, diff.1]; - sorted_diff.sort(); - plan_err!("Found mixed partition values on disk {:?}", sorted_diff) + Err(AllEqualValueError(None)) => Ok(vec![]), + Err(AllEqualValueError(Some(mut diff))) => { + diff.sort(); + plan_err!("Found mixed partition values on disk {:?}", diff) } } } diff --git a/datafusion/catalog-listing/src/table.rs b/datafusion/catalog-listing/src/table.rs index 7ee743a6abe71..b3328cc06303d 100644 --- a/datafusion/catalog-listing/src/table.rs +++ b/datafusion/catalog-listing/src/table.rs @@ -16,37 +16,46 @@ // under the License. use crate::config::SchemaSource; -use crate::helpers::{expr_applicable_for_cols, pruned_partition_list}; +use crate::helpers::{ + expr_applicable_for_cols, filter_partitioned_file, pruned_partition_list, +}; use crate::{ListingOptions, ListingTableConfig}; use arrow::datatypes::{Field, Schema, SchemaBuilder, SchemaRef}; use async_trait::async_trait; use datafusion_catalog::{ScanArgs, ScanResult, Session, TableProvider}; use datafusion_common::stats::Precision; use datafusion_common::{ - Constraints, SchemaExt, Statistics, internal_datafusion_err, plan_err, project_schema, + Constraints, DFSchema, SchemaExt, Statistics, internal_datafusion_err, plan_err, + project_schema, }; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_groups::FileGroup; -use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use datafusion_datasource::file_scan_config::{ + FileScanConfig, FileScanConfigBuilder, output_partitioning_from_partition_fields, +}; use datafusion_datasource::file_sink_config::{FileOutputMode, FileSinkConfig}; #[expect(deprecated)] use datafusion_datasource::schema_adapter::SchemaAdapterFactory; use datafusion_datasource::{ - ListingTableUrl, PartitionedFile, TableSchema, compute_all_files_statistics, + ListingTableUrl, PartitionedFile, TableSchemaBuilder, compute_all_files_statistics, +}; +use datafusion_execution::cache::cache_manager::{ + CachedFileMetadata, FileStatisticsCache, SchemaFingerprint, TableScopedPath, }; -use datafusion_execution::cache::TableScopedPath; -use datafusion_execution::cache::cache_manager::FileStatisticsCache; use datafusion_expr::dml::InsertOp; use datafusion_expr::execution_props::ExecutionProps; -use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; -use datafusion_physical_expr::create_lex_ordering; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion_expr::{ + Expr, Partitioning as LogicalPartitioning, TableProviderFilterPushDown, TableType, +}; +use datafusion_physical_expr::{create_lex_ordering, create_physical_partitioning}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::empty::EmptyExec; use futures::{Stream, StreamExt, TryStreamExt, future, stream}; use object_store::ObjectStore; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; /// Result of a file listing operation from [`ListingTable::list_files_for_scan`]. @@ -56,7 +65,7 @@ pub struct ListFilesResult { pub file_groups: Vec, /// Aggregated statistics for all files. pub statistics: Statistics, - /// Whether files are grouped by partition values (enables Hash partitioning). + /// Whether files are grouped by partition values. pub grouped_by_partition: bool, } @@ -142,15 +151,15 @@ pub struct ListFilesResult { /// # use datafusion_datasource_parquet::file_format::ParquetFormat;/// # /// # use datafusion_catalog::Session; /// async fn get_listing_table(session: &dyn Session) -> Result> { -/// let table_path = "/path/to/parquet"; +/// let table_path = "/path/to/parquet"; /// -/// // Parse the path -/// let table_path = ListingTableUrl::parse(table_path)?; +/// // Parse the path +/// let table_path = ListingTableUrl::parse(table_path)?; /// -/// // Create default parquet options -/// let file_format = ParquetFormat::new(); -/// let listing_options = ListingOptions::new(Arc::new(file_format)) -/// .with_file_extension(".parquet"); +/// // Create default parquet options +/// let file_format = ParquetFormat::new(); +/// let listing_options = ListingOptions::new(Arc::new(file_format)) +/// .with_file_extension(".parquet"); /// /// // Resolve the schema /// let resolved_schema = listing_options @@ -164,8 +173,8 @@ pub struct ListFilesResult { /// // Create a new TableProvider /// let provider = Arc::new(ListingTable::try_new(config)?); /// -/// # Ok(provider) -/// # } +/// Ok(provider) +/// } /// ``` #[derive(Debug, Clone)] pub struct ListingTable { @@ -186,13 +195,17 @@ pub struct ListingTable { /// The SQL definition for this table, if any definition: Option, /// Cache for collected file statistics - collected_statistics: Option>, + collected_statistics: Option>, /// Constraints applied to this table constraints: Constraints, /// Column default expressions for columns that are not physically present in the data files column_defaults: HashMap, /// Optional [`PhysicalExprAdapterFactory`] for creating physical expression adapters expr_adapter_factory: Option>, + /// Precomputed fingerprint of `file_schema` for file-statistics cache + /// validation. Constant for the table, so computed once here instead of per + /// file. + file_schema_fingerprint: Arc, } impl ListingTable { @@ -223,6 +236,9 @@ impl ListingTable { .with_metadata(file_schema.metadata().clone()), ); + let file_schema_fingerprint = + Arc::new(SchemaFingerprint::from_schema(&file_schema)); + let table = Self { table_paths: config.table_paths, file_schema, @@ -234,6 +250,7 @@ impl ListingTable { constraints: Constraints::default(), column_defaults: HashMap::new(), expr_adapter_factory: config.expr_adapter_factory, + file_schema_fingerprint, }; Ok(table) @@ -259,7 +276,7 @@ impl ListingTable { /// Setting a statistics cache on the `SessionContext` can avoid refetching statistics /// multiple times in the same session. /// - pub fn with_cache(mut self, cache: Option>) -> Self { + pub fn with_cache(mut self, cache: Option>) -> Self { self.collected_statistics = cache; self } @@ -321,14 +338,15 @@ impl ListingTable { /// Creates a file source for this table fn create_file_source(&self) -> Arc { - let table_schema = TableSchema::new( - Arc::clone(&self.file_schema), - self.options - .table_partition_cols - .iter() - .map(|(col, field)| Arc::new(Field::new(col, field.clone(), false))) - .collect(), - ); + let table_schema = TableSchemaBuilder::from(&self.file_schema) + .with_table_partition_cols( + self.options + .table_partition_cols + .iter() + .map(|(col, field)| Arc::new(Field::new(col, field.clone(), false))) + .collect::>(), + ) + .build(); self.options.format.file_source(table_schema) } @@ -433,6 +451,20 @@ fn derive_common_ordering_from_files(file_groups: &[FileGroup]) -> Option datafusion_common::Result { + let files = file_group + .into_inner() + .into_iter() + .map(|file| filter_partitioned_file(file, filters, df_schema)) + .filter_map(Result::transpose) + .collect::>>()?; + Ok(FileGroup::new(files)) +} + // Expressions can be used for partition pruning if they can be evaluated using // only the partition columns and there are partition columns. fn can_be_evaluated_for_partition_pruning( @@ -500,9 +532,19 @@ impl TableProvider for ListingTable { can_be_evaluated_for_partition_pruning(&table_partition_col_names, filter) }); - // We should not limit the number of partitioned files to scan if there are filters and limit - // at the same time. This is because the limit should be applied after the filters are applied. - let statistic_file_limit = if filters.is_empty() { limit } else { None }; + let declared_output_partitioning = self.options.output_partitioning.as_ref(); + + // We should not limit files before assigning declared output partitions + // or before applying non-partition filters. + let statistic_file_limit = + if filters.is_empty() && declared_output_partitioning.is_none() { + limit + } else { + None + }; + let file_group_count = declared_output_partitioning + .and_then(LogicalPartitioning::partition_count) + .unwrap_or_else(|| state.config().target_partitions()); let ListFilesResult { file_groups: mut partitioned_file_lists, @@ -522,17 +564,19 @@ impl TableProvider for ListingTable { state.execution_props(), &partitioned_file_lists, )?; - match state - .config_options() - .execution - .split_file_groups_by_statistics + let split_file_groups_by_statistics = declared_output_partitioning.is_none() + && state + .config_options() + .execution + .split_file_groups_by_statistics; + match split_file_groups_by_statistics .then(|| { output_ordering.first().map(|output_ordering| { FileScanConfig::split_groups_by_statistics_with_target_partitions( &self.table_schema, &partitioned_file_lists, output_ordering, - self.options.target_partitions, + file_group_count, ) }) }) @@ -540,7 +584,7 @@ impl TableProvider for ListingTable { { Some(Err(e)) => log::debug!("failed to split file groups by statistics: {e}"), Some(Ok(new_groups)) => { - if new_groups.len() <= self.options.target_partitions { + if new_groups.len() <= file_group_count { partitioned_file_lists = new_groups; } else { log::debug!( @@ -551,6 +595,51 @@ impl TableProvider for ListingTable { None => {} // no ordering required }; + let output_partitioning = if let Some(output_partitioning) = + declared_output_partitioning + { + let output_partitioning = match output_partitioning { + LogicalPartitioning::RoundRobinBatch(_) => { + return datafusion_common::not_impl_err!( + "RoundRobinBatch output partitioning is not supported for ListingTable" + ); + } + LogicalPartitioning::DistributeBy(_) => { + return datafusion_common::not_impl_err!( + "DistributeBy output partitioning is not supported for ListingTable" + ); + } + LogicalPartitioning::Hash(_, _) | LogicalPartitioning::Range(_) => { + let df_schema = DFSchema::try_from(Arc::clone(&self.table_schema))?; + create_physical_partitioning( + output_partitioning, + &df_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )? + } + }; + let partition_count = output_partitioning.partition_count(); + if partitioned_file_lists.len() != partition_count { + return plan_err!( + "ListingTable output_partitioning has {partition_count} partitions, but the scan has {} file groups", + partitioned_file_lists.len() + ); + } + Some(output_partitioning) + } else if partitioned_by_file_group { + // Files are grouped by partition column values: declare output + // partitioning on those columns so the optimizer can skip + // repartitioning for aggregates and joins on the partition columns. + output_partitioning_from_partition_fields( + &self.table_schema, + &table_partition_cols.clone().into(), + partitioned_file_lists.len(), + ) + } else { + None + }; + let Some(object_store_url) = self.table_paths.first().map(ListingTableUrl::object_store) else { @@ -560,24 +649,22 @@ impl TableProvider for ListingTable { }; let file_source = self.create_file_source(); + let scan_config = FileScanConfigBuilder::new(object_store_url, file_source) + .with_file_groups(partitioned_file_lists) + .with_constraints(self.constraints.clone()) + .with_statistics(statistics) + .with_projection_indices(projection)? + .with_limit(limit) + .with_output_ordering(output_ordering) + .with_output_partitioning(output_partitioning) + .with_expr_adapter(self.expr_adapter_factory.clone()) + .build(); // create the execution plan let plan = self .options .format - .create_physical_plan( - state, - FileScanConfigBuilder::new(object_store_url, file_source) - .with_file_groups(partitioned_file_lists) - .with_constraints(self.constraints.clone()) - .with_statistics(statistics) - .with_projection_indices(projection)? - .with_limit(limit) - .with_output_ordering(output_ordering) - .with_expr_adapter(self.expr_adapter_factory.clone()) - .with_partitioned_by_file_group(partitioned_by_file_group) - .build(), - ) + .create_physical_plan(state, scan_config) .await?; Ok(ScanResult::new(plan)) @@ -689,42 +776,63 @@ impl ListingTable { /// Get the list of files for a scan as well as the file level statistics. /// The list is grouped to let the execution plan know how the files should /// be distributed to different threads / executors. + /// + /// If [`ListingOptions::output_partitioning`] is set, returns one file + /// group per declared partition, including empty trailing groups. pub async fn list_files_for_scan<'a>( &'a self, ctx: &'a dyn Session, filters: &'a [Expr], limit: Option, ) -> datafusion_common::Result { - let store = if let Some(url) = self.table_paths.first() { - ctx.runtime_env().object_store(url)? + if let Some(output_partitioning) = self.options.output_partitioning.as_ref() { + self.list_files_for_declared_output_partitioning( + ctx, + output_partitioning, + filters, + ) + .await } else { - return Ok(ListFilesResult { - file_groups: vec![], - statistics: Statistics::new_unknown(&self.file_schema), - grouped_by_partition: false, - }); - }; + self.list_files_for_regular_scan(ctx, filters, limit).await + } + } + + async fn collect_files_for_scan<'a>( + &'a self, + ctx: &'a dyn Session, + store: &'a Arc, + listing_time_filters: &'a [Expr], + file_limit: Option, + ) -> datafusion_common::Result<(FileGroup, bool)> { // list files (with partitions) let file_list = future::try_join_all(self.table_paths.iter().map(|table_path| { pruned_partition_list( ctx, store.as_ref(), table_path, - filters, + listing_time_filters, &self.options.file_extension, &self.options.table_partition_cols, ) })) .await?; let meta_fetch_concurrency = - ctx.config_options().execution.meta_fetch_concurrency; - let file_list = stream::iter(file_list).flatten_unordered(meta_fetch_concurrency); + ctx.config_options().execution.meta_fetch_concurrency.get(); + // Table paths can overlap, for example when one path is a directory and + // another names a file inside it. A ListingTable uses one object store, + // so the object path uniquely identifies a file within this scan. + let mut seen_files = HashSet::new(); + let file_list = stream::iter(file_list) + .flatten_unordered(meta_fetch_concurrency) + .try_filter(move |file| { + future::ready(seen_files.insert(file.object_meta.location.clone())) + }); // collect the statistics and ordering if required by the config let files = file_list .map(|part_file| async { let part_file = part_file?; - let (statistics, ordering) = if self.options.collect_stat { - self.do_collect_statistics_and_ordering(ctx, &store, &part_file) + let (statistics, ordering) = if ctx.config().collect_statistics() { + self.do_collect_statistics_and_ordering(ctx, store, &part_file) .await? } else { (Arc::new(Statistics::new_unknown(&self.file_schema)), None) @@ -734,44 +842,146 @@ impl ListingTable { .with_ordering(ordering)) }) .boxed() - .buffer_unordered(ctx.config_options().execution.meta_fetch_concurrency); + .buffer_unordered( + ctx.config_options().execution.meta_fetch_concurrency.get(), + ); - let (file_group, inexact_stats) = - get_files_with_limit(files, limit, self.options.collect_stat).await?; + get_files_with_limit(files, file_limit, ctx.config().collect_statistics()).await + } + + async fn list_files_for_regular_scan<'a>( + &'a self, + ctx: &'a dyn Session, + filters: &'a [Expr], + limit: Option, + ) -> datafusion_common::Result { + let file_group_count = ctx.config().target_partitions(); + if file_group_count == 0 { + return plan_err!( + "ListingTable requires target_partitions to be greater than zero" + ); + } + + let store = if let Some(url) = self.table_paths.first() { + ctx.runtime_env().object_store(url)? + } else { + return Ok(ListFilesResult { + file_groups: vec![], + statistics: Statistics::new_unknown(&self.file_schema), + grouped_by_partition: false, + }); + }; + let (file_group, inexact_stats) = self + .collect_files_for_scan(ctx, &store, filters, limit) + .await?; // Threshold: 0 = disabled, N > 0 = enabled when distinct_keys >= N // // When enabled, files are grouped by their Hive partition column values, allowing - // FileScanConfig to declare Hash partitioning. This enables the optimizer to skip - // hash repartitioning for aggregates and joins on partition columns. + // FileScanConfig to declare output partitioning. This enables the optimizer to + // skip repartitioning for aggregates and joins on partition columns. let threshold = ctx.config_options().optimizer.preserve_file_partitions; - let (file_groups, grouped_by_partition) = if threshold > 0 - && !self.options.table_partition_cols.is_empty() - { - let grouped = - file_group.group_by_partition_values(self.options.target_partitions); - if grouped.len() >= threshold { - (grouped, true) + let (file_groups, grouped_by_partition) = + if threshold > 0 && !self.options.table_partition_cols.is_empty() { + let grouped = file_group.group_by_partition_values(file_group_count); + if grouped.len() >= threshold { + (grouped, true) + } else { + let all_files: Vec<_> = + grouped.into_iter().flat_map(|g| g.into_inner()).collect(); + ( + FileGroup::new(all_files).split_files(file_group_count), + false, + ) + } } else { - let all_files: Vec<_> = - grouped.into_iter().flat_map(|g| g.into_inner()).collect(); - ( - FileGroup::new(all_files).split_files(self.options.target_partitions), - false, - ) - } + (file_group.split_files(file_group_count), false) + }; + + self.list_files_result_from_groups( + ctx, + file_groups, + inexact_stats, + grouped_by_partition, + ) + } + + async fn list_files_for_declared_output_partitioning<'a>( + &'a self, + ctx: &'a dyn Session, + output_partitioning: &LogicalPartitioning, + filters: &'a [Expr], + ) -> datafusion_common::Result { + let Some(file_group_count) = output_partitioning.partition_count() else { + return datafusion_common::not_impl_err!( + "DistributeBy output partitioning is not supported for ListingTable" + ); + }; + if file_group_count == 0 { + return plan_err!( + "ListingTable output_partitioning requires at least one partition" + ); + } + + let store = if let Some(url) = self.table_paths.first() { + ctx.runtime_env().object_store(url)? } else { - ( - file_group.split_files(self.options.target_partitions), - false, - ) + return Ok(ListFilesResult { + file_groups: vec![], + statistics: Statistics::new_unknown(&self.file_schema), + grouped_by_partition: false, + }); }; + let (file_group, inexact_stats) = + self.collect_files_for_scan(ctx, &store, &[], None).await?; + let mut file_groups = file_group.split_files(file_group_count); + if !file_groups.is_empty() { + file_groups.resize_with(file_group_count, || FileGroup::new(vec![])); + } + let file_groups = + self.filter_declared_file_groups_by_partition_filters(file_groups, filters)?; + self.list_files_result_from_groups(ctx, file_groups, inexact_stats, false) + } + + fn filter_declared_file_groups_by_partition_filters( + &self, + file_groups: Vec, + filters: &[Expr], + ) -> datafusion_common::Result> { + if filters.is_empty() { + return Ok(file_groups); + } + + let df_schema = DFSchema::from_unqualified_fields( + self.options + .table_partition_cols + .iter() + .map(|(name, data_type)| Field::new(name, data_type.clone(), true)) + .collect(), + Default::default(), + )?; + + file_groups + .into_iter() + .map(|file_group| { + filter_file_group_by_partition_filters(file_group, filters, &df_schema) + }) + .collect::>>() + } + + fn list_files_result_from_groups( + &self, + ctx: &dyn Session, + file_groups: Vec, + inexact_stats: bool, + grouped_by_partition: bool, + ) -> datafusion_common::Result { let (file_groups, stats) = compute_all_files_statistics( file_groups, self.schema(), - self.options.collect_stat, + ctx.config().collect_statistics(), inexact_stats, )?; @@ -797,18 +1007,18 @@ impl ListingTable { store: &Arc, part_file: &PartitionedFile, ) -> datafusion_common::Result<(Arc, Option)> { - use datafusion_execution::cache::cache_manager::CachedFileMetadata; - let path = TableScopedPath { table: part_file.table_reference.clone(), path: part_file.object_meta.location.clone(), }; let meta = &part_file.object_meta; - // Check cache first - if we have valid cached statistics and ordering + // Check cache first. The key stays `{table, path}` for cheap lookups; + // the cached value carries the schema fingerprint to prevent reusing + // stats computed under a different file schema. if let Some(cache) = &self.collected_statistics && let Some(cached) = cache.get(&path) - && cached.is_valid_for(meta) + && cached.is_valid_for(meta, &self.file_schema_fingerprint) { // Return cached statistics and ordering return Ok((Arc::clone(&cached.statistics), cached.ordering.clone())); @@ -829,6 +1039,7 @@ impl ListingTable { &path, CachedFileMetadata::new( meta.clone(), + Arc::clone(&self.file_schema_fingerprint), Arc::clone(&statistics), file_meta.ordering.clone(), ), diff --git a/datafusion/catalog/src/catalog.rs b/datafusion/catalog/src/catalog.rs index 34cdf74440cb3..07da1293a781d 100644 --- a/datafusion/catalog/src/catalog.rs +++ b/datafusion/catalog/src/catalog.rs @@ -15,195 +15,8 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; -use std::fmt::Debug; -use std::sync::Arc; - -pub use crate::schema::SchemaProvider; -use datafusion_common::Result; -use datafusion_common::not_impl_err; - -/// Represents a catalog, comprising a number of named schemas. -/// -/// # Catalog Overview -/// -/// To plan and execute queries, DataFusion needs a "Catalog" that provides -/// metadata such as which schemas and tables exist, their columns and data -/// types, and how to access the data. -/// -/// The Catalog API consists: -/// * [`CatalogProviderList`]: a collection of `CatalogProvider`s -/// * [`CatalogProvider`]: a collection of `SchemaProvider`s (sometimes called a "database" in other systems) -/// * [`SchemaProvider`]: a collection of `TableProvider`s (often called a "schema" in other systems) -/// * [`TableProvider`]: individual tables -/// -/// # Implementing Catalogs -/// -/// To implement a catalog, you implement at least one of the [`CatalogProviderList`], -/// [`CatalogProvider`] and [`SchemaProvider`] traits and register them -/// appropriately in the `SessionContext`. -/// -/// DataFusion comes with a simple in-memory catalog implementation, -/// `MemoryCatalogProvider`, that is used by default and has no persistence. -/// DataFusion does not include more complex Catalog implementations because -/// catalog management is a key design choice for most data systems, and thus -/// it is unlikely that any general-purpose catalog implementation will work -/// well across many use cases. -/// -/// # Implementing "Remote" catalogs -/// -/// See [`remote_catalog`] for an end to end example of how to implement a -/// remote catalog. -/// -/// Sometimes catalog information is stored remotely and requires a network call -/// to retrieve. For example, the [Delta Lake] table format stores table -/// metadata in files on S3 that must be first downloaded to discover what -/// schemas and tables exist. -/// -/// [Delta Lake]: https://delta.io/ -/// [`remote_catalog`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/data_io/remote_catalog.rs -/// -/// The [`CatalogProvider`] can support this use case, but it takes some care. -/// The planning APIs in DataFusion are not `async` and thus network IO can not -/// be performed "lazily" / "on demand" during query planning. The rationale for -/// this design is that using remote procedure calls for all catalog accesses -/// required for query planning would likely result in multiple network calls -/// per plan, resulting in very poor planning performance. -/// -/// To implement [`CatalogProvider`] and [`SchemaProvider`] for remote catalogs, -/// you need to provide an in memory snapshot of the required metadata. Most -/// systems typically either already have this information cached locally or can -/// batch access to the remote catalog to retrieve multiple schemas and tables -/// in a single network call. -/// -/// Note that [`SchemaProvider::table`] **is** an `async` function in order to -/// simplify implementing simple [`SchemaProvider`]s. For many table formats it -/// is easy to list all available tables but there is additional non trivial -/// access required to read table details (e.g. statistics). -/// -/// The pattern that DataFusion itself uses to plan SQL queries is to walk over -/// the query to find all table references, performing required remote catalog -/// lookups in parallel, storing the results in a cached snapshot, and then plans -/// the query using that snapshot. -/// -/// # Example Catalog Implementations -/// -/// Here are some examples of how to implement custom catalogs: -/// -/// * [`datafusion-cli`]: [`DynamicFileCatalogProvider`] catalog provider -/// that treats files and directories on a filesystem as tables. -/// -/// * The [`catalog.rs`]: a simple directory based catalog. -/// -/// * [delta-rs]: [`UnityCatalogProvider`] implementation that can -/// read from Delta Lake tables -/// -/// [`datafusion-cli`]: https://datafusion.apache.org/user-guide/cli/index.html -/// [`DynamicFileCatalogProvider`]: https://github.com/apache/datafusion/blob/31b9b48b08592b7d293f46e75707aad7dadd7cbc/datafusion-cli/src/catalog.rs#L75 -/// [`catalog.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/data_io/catalog.rs -/// [delta-rs]: https://github.com/delta-io/delta-rs -/// [`UnityCatalogProvider`]: https://github.com/delta-io/delta-rs/blob/951436ecec476ce65b5ed3b58b50fb0846ca7b91/crates/deltalake-core/src/data_catalog/unity/datafusion.rs#L111-L123 -/// -/// [`TableProvider`]: crate::TableProvider -pub trait CatalogProvider: Any + Debug + Sync + Send { - /// Retrieves the list of available schema names in this catalog. - fn schema_names(&self) -> Vec; - - /// Retrieves a specific schema from the catalog by name, provided it exists. - fn schema(&self, name: &str) -> Option>; - - /// Adds a new schema to this catalog. - /// - /// If a schema of the same name existed before, it is replaced in - /// the catalog and returned. - /// - /// By default returns a "Not Implemented" error - fn register_schema( - &self, - name: &str, - schema: Arc, - ) -> Result>> { - // use variables to avoid unused variable warnings - let _ = name; - let _ = schema; - not_impl_err!("Registering new schemas is not supported") - } - - /// Removes a schema from this catalog. Implementations of this method should return - /// errors if the schema exists but cannot be dropped. For example, in DataFusion's - /// default in-memory catalog, `MemoryCatalogProvider`, a non-empty schema - /// will only be successfully dropped when `cascade` is true. - /// This is equivalent to how DROP SCHEMA works in PostgreSQL. - /// - /// Implementations of this method should return None if schema with `name` - /// does not exist. - /// - /// By default returns a "Not Implemented" error - fn deregister_schema( - &self, - _name: &str, - _cascade: bool, - ) -> Result>> { - not_impl_err!("Deregistering new schemas is not supported") - } -} - -impl dyn CatalogProvider { - /// Returns `true` if the catalog provider is of type `T`. - /// - /// Prefer this over `downcast_ref::().is_some()`. Works correctly when - /// called on `Arc` via auto-deref. - pub fn is(&self) -> bool { - (self as &dyn Any).is::() - } - - /// Attempts to downcast this catalog provider to a concrete type `T`, - /// returning `None` if the provider is not of that type. - /// - /// Works correctly when called on `Arc` via auto-deref, - /// unlike `(&arc as &dyn Any).downcast_ref::()` which would attempt to - /// downcast the `Arc` itself. - pub fn downcast_ref(&self) -> Option<&T> { - (self as &dyn Any).downcast_ref() - } -} - -/// Represent a list of named [`CatalogProvider`]s. -/// -/// Please see the documentation on [`CatalogProvider`] for details of -/// implementing a custom catalog. -pub trait CatalogProviderList: Any + Debug + Sync + Send { - /// Adds a new catalog to this catalog list - /// If a catalog of the same name existed before, it is replaced in the list and returned. - fn register_catalog( - &self, - name: String, - catalog: Arc, - ) -> Option>; - - /// Retrieves the list of available catalog names - fn catalog_names(&self) -> Vec; - - /// Retrieves a specific catalog by name, provided it exists. - fn catalog(&self, name: &str) -> Option>; -} - -impl dyn CatalogProviderList { - /// Returns `true` if the catalog provider list is of type `T`. - /// - /// Prefer this over `downcast_ref::().is_some()`. Works correctly when - /// called on `Arc` via auto-deref. - pub fn is(&self) -> bool { - (self as &dyn Any).is::() - } - - /// Attempts to downcast this catalog provider list to a concrete type `T`, - /// returning `None` if the provider list is not of that type. - /// - /// Works correctly when called on `Arc` via - /// auto-deref, unlike `(&arc as &dyn Any).downcast_ref::()` which would - /// attempt to downcast the `Arc` itself. - pub fn downcast_ref(&self) -> Option<&T> { - (self as &dyn Any).downcast_ref() - } -} +// Re-export from this module for backwards compatibility. +pub use datafusion_session::{CatalogProvider, CatalogProviderList}; +// Re-export so users can access this type through `datafusion_catalog` and +// `datafusion::catalog` without depending directly on `datafusion_session`. +pub use datafusion_session::EmptyCatalogProviderList; diff --git a/datafusion/catalog/src/cte_worktable.rs b/datafusion/catalog/src/cte_worktable.rs index dd313ebb4cbff..5ec688526c92b 100644 --- a/datafusion/catalog/src/cte_worktable.rs +++ b/datafusion/catalog/src/cte_worktable.rs @@ -36,14 +36,16 @@ use crate::{ScanArgs, ScanResult, Session, TableProvider}; pub struct CteWorkTable { /// The name of the CTE work table name: String, - /// This schema must be shared across both the static and recursive terms of a recursive query + /// Schema exposed by recursive self-references while planning the recursive term. + /// + /// This is a conservative work-table schema, not the final recursive query output + /// schema. For example, the SQL planner may mark fields nullable here so recursive + /// references do not inherit unsound anchor-term nullability assumptions. table_schema: SchemaRef, } impl CteWorkTable { - /// construct a new CteWorkTable with the given name and schema - /// This schema must match the schema of the recursive term of the query - /// Since the scan method will contain an physical plan that assumes this schema + /// Construct a new CteWorkTable with the given name and self-reference schema. pub fn new(name: &str, table_schema: SchemaRef) -> Self { Self { name: name.to_owned(), @@ -56,7 +58,7 @@ impl CteWorkTable { &self.name } - /// The schema of the recursive term of the query + /// The schema exposed by scans of the recursive self-reference. pub fn schema(&self) -> SchemaRef { Arc::clone(&self.table_schema) } diff --git a/datafusion/catalog/src/information_schema.rs b/datafusion/catalog/src/information_schema.rs index 34c677c3dd43e..d9ad7791af67c 100644 --- a/datafusion/catalog/src/information_schema.rs +++ b/datafusion/catalog/src/information_schema.rs @@ -20,6 +20,7 @@ //! [Information Schema]: https://en.wikipedia.org/wiki/Information_schema use crate::streaming::StreamingTable; +use crate::table::TableFunction; use crate::{CatalogProviderList, SchemaProvider, TableProvider}; use arrow::array::builder::{BooleanBuilder, UInt8Builder}; use arrow::{ @@ -81,14 +82,28 @@ impl InformationSchemaProvider { /// Creates a new [`InformationSchemaProvider`] for the provided `catalog_list` pub fn new(catalog_list: Arc) -> Self { Self { - config: InformationSchemaConfig { catalog_list }, + config: InformationSchemaConfig { + catalog_list, + table_functions: HashMap::new(), + }, } } + + /// Attach the session's table (UDTF) functions so that they appear in + /// `information_schema.routines` / `SHOW FUNCTIONS`. + pub fn with_table_functions( + mut self, + table_functions: HashMap>, + ) -> Self { + self.config.table_functions = table_functions; + self + } } #[derive(Clone, Debug)] struct InformationSchemaConfig { catalog_list: Arc, + table_functions: HashMap>, } impl InformationSchemaConfig { @@ -136,7 +151,7 @@ impl InformationSchemaConfig { Ok(()) } - async fn make_schemata(&self, builder: &mut InformationSchemataBuilder) { + fn make_schemata(&self, builder: &mut InformationSchemataBuilder) { for catalog_name in self.catalog_list.catalog_names() { let catalog = self.catalog_list.catalog(&catalog_name).unwrap(); @@ -301,6 +316,26 @@ impl InformationSchemaConfig { ) } } + + // Table functions (UDTFs) don't have scalar signatures; their return + // type is always a table, so emit a single row per UDTF with + // routine_type = "FUNCTION", function_type = "TABLE" and + // data_type = "TABLE". + for name in self.table_functions.keys() { + builder.add_routine( + catalog_name, + schema_name, + name, + "FUNCTION", + // No signature is available for UDTFs; report deterministic + // = false to stay conservative. + false, + Some(&"TABLE"), + "TABLE", + None::, + None::, + ) + } Ok(()) } @@ -400,6 +435,14 @@ impl InformationSchemaConfig { } } + // UDTFs deliberately do NOT appear in `information_schema.parameters`. + // A same-named scalar UDF (e.g. `generate_series` exists as both a + // scalar UDF in functions-nested and a UDTF in functions-table) would + // cross-join with a UDTF row keyed only by (name, rid) and produce + // spurious `TABLE`-typed variants of every scalar signature in + // SHOW FUNCTIONS. `show_functions_to_plan` sources UDTFs directly + // from `information_schema.routines` via a UNION branch instead. + Ok(()) } @@ -967,18 +1010,34 @@ struct InformationSchemata { config: InformationSchemaConfig, } +/// The Arrow schema of [`information_schema.schemata`] rows. +/// +/// Useful for downstream catalog implementations that want to declare a +/// `TableProvider` for `schemata` before populating any rows via +/// [`InformationSchemataBuilder`]. +/// +/// Columns and nullability match +/// . +/// +/// [`information_schema.schemata`]: https://www.postgresql.org/docs/current/infoschema-schemata.html +pub fn schemata_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("catalog_name", DataType::Utf8, false), + Field::new("schema_name", DataType::Utf8, false), + Field::new("schema_owner", DataType::Utf8, true), + Field::new("default_character_set_catalog", DataType::Utf8, true), + Field::new("default_character_set_schema", DataType::Utf8, true), + Field::new("default_character_set_name", DataType::Utf8, true), + Field::new("sql_path", DataType::Utf8, true), + ])) +} + impl InformationSchemata { fn new(config: InformationSchemaConfig) -> Self { - let schema = Arc::new(Schema::new(vec![ - Field::new("catalog_name", DataType::Utf8, false), - Field::new("schema_name", DataType::Utf8, false), - Field::new("schema_owner", DataType::Utf8, true), - Field::new("default_character_set_catalog", DataType::Utf8, true), - Field::new("default_character_set_schema", DataType::Utf8, true), - Field::new("default_character_set_name", DataType::Utf8, true), - Field::new("sql_path", DataType::Utf8, true), - ])); - Self { schema, config } + Self { + schema: schemata_schema(), + config, + } } fn builder(&self) -> InformationSchemataBuilder { @@ -995,7 +1054,16 @@ impl InformationSchemata { } } -struct InformationSchemataBuilder { +/// Builder that produces [`RecordBatch`] values matching the schema of +/// `information_schema.schemata` (see [`schemata_schema`]). +/// +/// Intended for downstream catalog implementations that need to emit +/// `schemata` rows from their own metadata source rather than going +/// through DataFusion's `InformationSchemaProvider`, which enumerates +/// schemas synchronously via `CatalogProviderList` and so is unsuitable +/// for catalog backends that resolve asynchronously. +#[derive(Debug)] +pub struct InformationSchemataBuilder { schema: SchemaRef, catalog_name: StringBuilder, schema_name: StringBuilder, @@ -1006,8 +1074,32 @@ struct InformationSchemataBuilder { sql_path: StringBuilder, } +impl Default for InformationSchemataBuilder { + fn default() -> Self { + Self::new() + } +} + impl InformationSchemataBuilder { - fn add_schemata( + /// Construct an empty builder. + pub fn new() -> Self { + Self { + schema: schemata_schema(), + catalog_name: StringBuilder::new(), + schema_name: StringBuilder::new(), + schema_owner: StringBuilder::new(), + default_character_set_catalog: StringBuilder::new(), + default_character_set_schema: StringBuilder::new(), + default_character_set_name: StringBuilder::new(), + sql_path: StringBuilder::new(), + } + } + + /// Append one row to the builder. `schema_owner` is the optional SQL + /// schema owner; the three `default_character_set_*` columns and + /// `sql_path` are written as null (DataFusion does not model those + /// concepts; see the PostgreSQL docs link on [`schemata_schema`]). + pub fn add_schemata( &mut self, catalog_name: &str, schema_name: &str, @@ -1019,15 +1111,19 @@ impl InformationSchemataBuilder { Some(owner) => self.schema_owner.append_value(owner), None => self.schema_owner.append_null(), } - // refer to https://www.postgresql.org/docs/current/infoschema-schemata.html, - // these rows apply to a feature that is not implemented in DataFusion self.default_character_set_catalog.append_null(); self.default_character_set_schema.append_null(); self.default_character_set_name.append_null(); self.sql_path.append_null(); } - fn finish(&mut self) -> RecordBatch { + /// Finalize the builder into a [`RecordBatch`]. + /// + /// Returns an error only if Arrow buffer construction fails, which + /// the builder's column-count and type invariants make unreachable + /// under normal use. The `Result` return type preserves room to add + /// validation in the future without a breaking API change. + pub fn finish(&mut self) -> Result { RecordBatch::try_new( Arc::clone(&self.schema), vec![ @@ -1040,7 +1136,7 @@ impl InformationSchemataBuilder { Arc::new(self.sql_path.finish()), ], ) - .unwrap() + .map_err(DataFusionError::from) } } @@ -1056,8 +1152,8 @@ impl PartitionStream for InformationSchemata { Arc::clone(&self.schema), // TODO: Stream this futures::stream::once(async move { - config.make_schemata(&mut builder).await; - Ok(builder.finish()) + config.make_schemata(&mut builder); + builder.finish() }), )) } @@ -1413,11 +1509,63 @@ impl PartitionStream for InformationSchemaParameters { mod tests { use super::*; use crate::CatalogProvider; + use arrow::array::Array; + + #[test] + fn schemata_builder_emits_canonical_schema_and_rows() { + // Construct via `Default` so the test exercises both `new()` (via + // the `Default` impl) and the public column-layout contract. + let mut builder = InformationSchemataBuilder::default(); + builder.add_schemata("cat", "schema_one", Some("alice")); + builder.add_schemata("cat", "schema_two", None); + let batch = builder.finish().expect("finish should not fail"); + + assert_eq!(batch.schema(), schemata_schema()); + assert_eq!(batch.num_rows(), 2); + + let col = |name: &str| { + batch + .column_by_name(name) + .unwrap_or_else(|| panic!("missing column {name}")) + }; + let string_col = |name: &str| { + col(name) + .as_any() + .downcast_ref::() + .unwrap_or_else(|| panic!("{name} should be a StringArray")) + }; + + let catalog = string_col("catalog_name"); + assert_eq!(catalog.value(0), "cat"); + assert_eq!(catalog.value(1), "cat"); + + let schema = string_col("schema_name"); + assert_eq!(schema.value(0), "schema_one"); + assert_eq!(schema.value(1), "schema_two"); + + let owner = string_col("schema_owner"); + assert_eq!(owner.value(0), "alice"); + assert!(owner.is_null(1)); + + // The three character-set columns and sql_path are unconditionally + // null — they exist for SQL-standard column-layout compatibility. + for name in [ + "default_character_set_catalog", + "default_character_set_schema", + "default_character_set_name", + "sql_path", + ] { + let c = string_col(name); + assert!(c.is_null(0), "{name} row 0 should be null"); + assert!(c.is_null(1), "{name} row 1 should be null"); + } + } #[tokio::test] async fn make_tables_uses_table_type() { let config = InformationSchemaConfig { catalog_list: Arc::new(Fixture), + table_functions: HashMap::new(), }; let mut builder = InformationSchemaTablesBuilder { catalog_names: StringBuilder::new(), diff --git a/datafusion/catalog/src/lib.rs b/datafusion/catalog/src/lib.rs index 33d54b7cb89d5..815bfe32fac72 100644 --- a/datafusion/catalog/src/lib.rs +++ b/datafusion/catalog/src/lib.rs @@ -25,7 +25,10 @@ #![cfg_attr(not(test), deny(clippy::clone_on_ref_ptr))] #![cfg_attr(test, allow(clippy::needless_pass_by_value))] -//! Interfaces and default implementations of catalogs and schemas. +//! Default implementations of catalogs and schemas. +//! +//! The catalog interfaces are defined in [`datafusion_session`] and re-exported +//! by this crate. //! //! Implementations //! * Information schema: [`information_schema`] @@ -57,8 +60,3 @@ pub use memory::{ }; pub use schema::*; pub use table::*; - -// For backwards compatibility, -mod session { - pub use datafusion_session::Session; -} diff --git a/datafusion/catalog/src/memory/table.rs b/datafusion/catalog/src/memory/table.rs index 8102c15079658..4cf96cb364be8 100644 --- a/datafusion/catalog/src/memory/table.rs +++ b/datafusion/catalog/src/memory/table.rs @@ -33,11 +33,11 @@ use arrow::record_batch::RecordBatch; use datafusion_common::error::Result; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Constraints, DFSchema, SchemaExt, not_impl_err, plan_err}; -use datafusion_common_runtime::JoinSet; use datafusion_datasource::memory::{MemSink, MemorySourceConfig}; use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::source::DataSourceExec; use datafusion_expr::dml::InsertOp; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{Expr, SortExpr, TableType}; use datafusion_physical_expr::{ LexOrdering, create_physical_expr, create_physical_sort_exprs, @@ -45,13 +45,12 @@ use datafusion_physical_expr::{ use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - PhysicalExpr, PlanProperties, common, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PhysicalExpr, PlanProperties, ReplaceChildrenOptions, collect_partitioned, }; use datafusion_session::Session; use async_trait::async_trait; -use futures::StreamExt; use log::debug; use parking_lot::Mutex; use tokio::sync::RwLock; @@ -146,68 +145,28 @@ impl MemTable { state: &dyn Session, ) -> Result { let schema = t.schema(); - let constraints = t.constraints(); - let exec = t.scan(state, None, &[], None).await?; - let partition_count = exec.output_partitioning().partition_count(); - - let mut join_set = JoinSet::new(); - - for part_idx in 0..partition_count { - let task = state.task_ctx(); - let exec = Arc::clone(&exec); - join_set.spawn(async move { - let stream = exec.execute(part_idx, task)?; - common::collect(stream).await - }); - } + let constraints = t.constraints().cloned().unwrap_or_default(); - let mut data: Vec> = - Vec::with_capacity(exec.output_partitioning().partition_count()); - - while let Some(result) = join_set.join_next().await { - match result { - Ok(res) => data.push(res?), - Err(e) => { - if e.is_panic() { - std::panic::resume_unwind(e.into_panic()); - } else { - unreachable!(); - } - } - } - } - - let mut exec = DataSourceExec::new(Arc::new(MemorySourceConfig::try_new( - &data, - Arc::clone(&schema), - None, - )?)); - if let Some(cons) = constraints { - exec = exec.with_constraints(cons.clone()); - } - - if let Some(num_partitions) = output_partitions { + let exec = t.scan(state, None, &[], None).await?; + let data = collect_partitioned(exec, state.task_ctx()).await?; + + // Optionally repartition the collected batches. + let data = if let Some(num_partitions) = output_partitions { + let source = DataSourceExec::new(Arc::new(MemorySourceConfig::try_new( + &data, + Arc::clone(&schema), + None, + )?)); let exec = RepartitionExec::try_new( - Arc::new(exec), + Arc::new(source), Partitioning::RoundRobinBatch(num_partitions), )?; + collect_partitioned(Arc::new(exec), state.task_ctx()).await? + } else { + data + }; - // execute and collect results - let mut output_partitions = vec![]; - for i in 0..exec.properties().output_partitioning().partition_count() { - // execute this *output* partition and collect all batches - let task_ctx = state.task_ctx(); - let mut stream = exec.execute(i, task_ctx)?; - let mut batches = vec![]; - while let Some(result) = stream.next().await { - batches.push(result?); - } - output_partitions.push(batches); - } - - return MemTable::try_new(Arc::clone(&schema), output_partitions); - } - MemTable::try_new(Arc::clone(&schema), data) + MemTable::try_new(schema, data).map(|table| table.with_constraints(constraints)) } } @@ -252,8 +211,12 @@ impl TableProvider for MemTable { let eqp = state.execution_props(); let mut file_sort_order = vec![]; for sort_exprs in sort_order.iter() { - let physical_exprs = - create_physical_sort_exprs(sort_exprs, &df_schema, eqp)?; + let physical_exprs = create_physical_sort_exprs( + sort_exprs, + &df_schema, + eqp, + &PhysicalPlanningContext::default(), + )?; file_sort_order.extend(LexOrdering::new(physical_exprs)); } source = source.try_with_sort_information(file_sort_order)?; @@ -399,8 +362,12 @@ impl TableProvider for MemTable { let physical_assignments: HashMap> = assignments .iter() .map(|(name, expr)| { - let physical_expr = - create_physical_expr(expr, &df_schema, state.execution_props())?; + let physical_expr = create_physical_expr( + expr, + &df_schema, + state.execution_props(), + &PhysicalPlanningContext::default(), + )?; Ok((name.clone(), physical_expr)) }) .collect::>()?; @@ -513,8 +480,12 @@ fn evaluate_filters_to_mask( let mut combined_mask: Option = None; for filter_expr in filters { - let physical_expr = - create_physical_expr(filter_expr, df_schema, execution_props)?; + let physical_expr = create_physical_expr( + filter_expr, + df_schema, + execution_props, + &PhysicalPlanningContext::default(), + )?; let result = physical_expr.evaluate(batch)?; let array = result.into_array(batch.num_rows())?; @@ -601,13 +572,24 @@ impl ExecutionPlan for DmlResultExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -630,7 +612,7 @@ impl ExecutionPlan for DmlResultExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/catalog/src/schema.rs b/datafusion/catalog/src/schema.rs index d99027593ccce..40b20caeb9bb9 100644 --- a/datafusion/catalog/src/schema.rs +++ b/datafusion/catalog/src/schema.rs @@ -15,93 +15,5 @@ // specific language governing permissions and limitations // under the License. -//! Describes the interface and built-in implementations of schemas, -//! representing collections of named tables. - -use async_trait::async_trait; -use datafusion_common::{DataFusionError, exec_err}; -use std::any::Any; -use std::fmt::Debug; -use std::sync::Arc; - -use crate::table::TableProvider; -use datafusion_common::Result; -use datafusion_expr::TableType; - -/// Represents a schema, comprising a number of named tables. -/// -/// Please see [`CatalogProvider`] for details of implementing a custom catalog. -/// -/// [`CatalogProvider`]: super::CatalogProvider -#[async_trait] -pub trait SchemaProvider: Any + Debug + Sync + Send { - /// Returns the owner of the Schema, default is None. This value is reported - /// as part of `information_tables.schemata - fn owner_name(&self) -> Option<&str> { - None - } - - /// Retrieves the list of available table names in this schema. - fn table_names(&self) -> Vec; - - /// Retrieves a specific table from the schema by name, if it exists, - /// otherwise returns `None`. - async fn table( - &self, - name: &str, - ) -> Result>, DataFusionError>; - - /// Retrieves the type of a specific table from the schema by name, if it exists, otherwise - /// returns `None`. Implementations for which this operation is cheap but [Self::table] is - /// expensive can override this to improve operations that only need the type, e.g. - /// `SELECT * FROM information_schema.tables`. - async fn table_type(&self, name: &str) -> Result> { - self.table(name).await.map(|o| o.map(|t| t.table_type())) - } - - /// If supported by the implementation, adds a new table named `name` to - /// this schema. - /// - /// If a table of the same name was already registered, returns "Table - /// already exists" error. - #[expect(unused_variables)] - fn register_table( - &self, - name: String, - table: Arc, - ) -> Result>> { - exec_err!("schema provider does not support registering tables") - } - - /// If supported by the implementation, removes the `name` table from this - /// schema and returns the previously registered [`TableProvider`], if any. - /// - /// If no `name` table exists, returns Ok(None). - #[expect(unused_variables)] - fn deregister_table(&self, name: &str) -> Result>> { - exec_err!("schema provider does not support deregistering tables") - } - - /// Returns true if table exist in the schema provider, false otherwise. - fn table_exist(&self, name: &str) -> bool; -} - -impl dyn SchemaProvider { - /// Returns `true` if the schema provider is of type `T`. - /// - /// Prefer this over `downcast_ref::().is_some()`. Works correctly when - /// called on `Arc` via auto-deref. - pub fn is(&self) -> bool { - (self as &dyn Any).is::() - } - - /// Attempts to downcast this schema provider to a concrete type `T`, - /// returning `None` if the provider is not of that type. - /// - /// Works correctly when called on `Arc` via auto-deref, - /// unlike `(&arc as &dyn Any).downcast_ref::()` which would attempt to - /// downcast the `Arc` itself. - pub fn downcast_ref(&self) -> Option<&T> { - (self as &dyn Any).downcast_ref() - } -} +// Re-export from this module for backwards compatibility. +pub use datafusion_session::SchemaProvider; diff --git a/datafusion/catalog/src/stream.rs b/datafusion/catalog/src/stream.rs index 8501ea65902e2..c8060456dd2a7 100644 --- a/datafusion/catalog/src/stream.rs +++ b/datafusion/catalog/src/stream.rs @@ -53,7 +53,15 @@ impl TableProviderFactory for StreamTableFactory { cmd: &CreateExternalTable, ) -> Result> { let schema: SchemaRef = Arc::clone(cmd.schema.inner()); - let location = cmd.location.clone(); + let location = match cmd.locations.as_slice() { + [single] => single.clone(), + _ => { + return config_err!( + "Stream tables support exactly one location; \ + use a listing table to read multiple files" + ); + } + }; let encoding = cmd.file_type.parse()?; let header = if let Ok(opt) = cmd .options diff --git a/datafusion/catalog/src/streaming.rs b/datafusion/catalog/src/streaming.rs index e609877c2b778..50f05355aa75e 100644 --- a/datafusion/catalog/src/streaming.rs +++ b/datafusion/catalog/src/streaming.rs @@ -22,9 +22,13 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use async_trait::async_trait; use datafusion_common::{DFSchema, Result, plan_err}; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{Expr, SortExpr, TableType}; use datafusion_physical_expr::equivalence::project_ordering; -use datafusion_physical_expr::{LexOrdering, create_physical_sort_exprs}; +use datafusion_physical_expr::projection::ProjectionMapping; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, Partitioning, create_physical_sort_exprs, +}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::streaming::{PartitionStream, StreamingTableExec}; use log::debug; @@ -38,6 +42,7 @@ pub struct StreamingTable { partitions: Vec>, infinite: bool, sort_order: Vec, + output_partitioning: Option, } impl StreamingTable { @@ -62,6 +67,7 @@ impl StreamingTable { partitions, infinite: false, sort_order: vec![], + output_partitioning: None, }) } @@ -76,6 +82,33 @@ impl StreamingTable { self.sort_order = sort_order; self } + + /// Declares the output partitioning of this streaming table. + /// + /// The partitioning expressions refer to the table schema before scan + /// projection. If a scan projection removes a partitioning expression, the + /// physical plan reports unknown partitioning. + pub fn with_output_partitioning(mut self, output_partitioning: Partitioning) -> Self { + self.output_partitioning = Some(output_partitioning); + self + } + + fn output_partitioning( + &self, + projection: Option<&Vec>, + ) -> Result { + let Some(output_partitioning) = &self.output_partitioning else { + return Ok(Partitioning::UnknownPartitioning(self.partitions.len())); + }; + let Some(projection) = projection else { + return Ok(output_partitioning.clone()); + }; + + let projection_mapping = + ProjectionMapping::from_indices(projection, &self.schema)?; + let eq_properties = EquivalenceProperties::new(Arc::clone(&self.schema)); + Ok(output_partitioning.project(&projection_mapping, &eq_properties)) + } } #[async_trait] @@ -99,8 +132,12 @@ impl TableProvider for StreamingTable { let df_schema = DFSchema::try_from(Arc::clone(&self.schema))?; let eqp = state.execution_props(); - let original_sort_exprs = - create_physical_sort_exprs(&self.sort_order, &df_schema, eqp)?; + let original_sort_exprs = create_physical_sort_exprs( + &self.sort_order, + &df_schema, + eqp, + &PhysicalPlanningContext::default(), + )?; if let Some(p) = projection { // When performing a projection, the output columns will not match @@ -119,13 +156,16 @@ impl TableProvider for StreamingTable { vec![] }; - Ok(Arc::new(StreamingTableExec::try_new( + let exec = StreamingTableExec::try_new( Arc::clone(&self.schema), self.partitions.clone(), projection, LexOrdering::new(physical_sort), self.infinite, limit, - )?)) + )? + .with_output_partitioning(self.output_partitioning(projection)?)?; + + Ok(Arc::new(exec)) } } diff --git a/datafusion/catalog/src/table.rs b/datafusion/catalog/src/table.rs index 5d1391bed1172..2a10efbdcce6a 100644 --- a/datafusion/catalog/src/table.rs +++ b/datafusion/catalog/src/table.rs @@ -15,598 +15,8 @@ // specific language governing permissions and limitations // under the License. -use std::any::Any; -use std::borrow::Cow; -use std::fmt::Debug; -use std::sync::Arc; - -use crate::session::Session; -use arrow::datatypes::SchemaRef; -use async_trait::async_trait; -use datafusion_common::{Constraints, Statistics, not_impl_err}; -use datafusion_common::{Result, internal_err}; -use datafusion_expr::Expr; - -use datafusion_expr::dml::InsertOp; -use datafusion_expr::{ - CreateExternalTable, LogicalPlan, TableProviderFilterPushDown, TableType, +// Re-export from this module for backwards compatibility. +pub use datafusion_session::{ + ScanArgs, ScanResult, TableFunction, TableFunctionArgs, TableFunctionImpl, + TableProvider, TableProviderFactory, }; -use datafusion_physical_plan::ExecutionPlan; - -/// A table which can be queried and modified. -/// -/// Please see [`CatalogProvider`] for details of implementing a custom catalog. -/// -/// [`TableProvider`] represents a source of data which can provide data as -/// Apache Arrow [`RecordBatch`]es. Implementations of this trait provide -/// important information for planning such as: -/// -/// 1. [`Self::schema`]: The schema (columns and their types) of the table -/// 2. [`Self::supports_filters_pushdown`]: Should filters be pushed into this scan -/// 2. [`Self::scan`]: An [`ExecutionPlan`] that can read data -/// -/// [`RecordBatch`]: https://docs.rs/arrow/latest/arrow/record_batch/struct.RecordBatch.html -/// [`CatalogProvider`]: super::CatalogProvider -#[async_trait] -pub trait TableProvider: Any + Debug + Sync + Send { - /// Get a reference to the schema for this table - fn schema(&self) -> SchemaRef; - - /// Get a reference to the constraints of the table. - /// Returns: - /// - `None` for tables that do not support constraints. - /// - `Some(&Constraints)` for tables supporting constraints. - /// Therefore, a `Some(&Constraints::empty())` return value indicates that - /// this table supports constraints, but there are no constraints. - fn constraints(&self) -> Option<&Constraints> { - None - } - - /// Get the type of this table for metadata/catalog purposes. - fn table_type(&self) -> TableType; - - /// Get the create statement used to create this table, if available. - fn get_table_definition(&self) -> Option<&str> { - None - } - - /// Get the [`LogicalPlan`] of this table, if available. - fn get_logical_plan(&'_ self) -> Option> { - None - } - - /// Get the default value for a column, if available. - fn get_column_default(&self, _column: &str) -> Option<&Expr> { - None - } - - /// Create an [`ExecutionPlan`] for scanning the table with optional - /// `projection`, `filter`, and `limit`, described below. - /// - /// The returned `ExecutionPlan` is responsible for scanning the datasource's - /// partitions in a streaming, parallelized fashion. - /// - /// # Projection - /// - /// If specified, only a subset of columns should be returned, in the order - /// specified. The projection is a set of indexes of the fields in - /// [`Self::schema`]. - /// - /// DataFusion provides the projection so the scan reads only the columns - /// actually used in the query, an optimization called "Projection - /// Pushdown". Some datasources, such as Parquet, can use this information - /// to go significantly faster when only a subset of columns is required. - /// - /// # Filters - /// - /// A list of boolean filter [`Expr`]s to evaluate *during* the scan, in the - /// manner specified by [`Self::supports_filters_pushdown`]. Only rows for - /// which *all* of the `Expr`s evaluate to `true` must be returned (that is, - /// the expressions are `AND`ed together). - /// - /// To enable filter pushdown, override - /// [`Self::supports_filters_pushdown`]. The default implementation does not - /// push down filters, and `filters` will be empty. - /// - /// DataFusion pushes filters into scans whenever possible ("Filter - /// Pushdown"). Depending on the data format and implementation, evaluating - /// predicates during the scan can significantly improve performance. - /// - /// ## Note: Some columns may appear *only* in Filters - /// - /// In some cases, a query may use a column only in a filter and the - /// projection will not contain all columns referenced by the filter - /// expressions. - /// - /// For example, given the query `SELECT t.a FROM t WHERE t.b > 5`, - /// - /// ```text - /// ┌────────────────────┐ - /// │ Projection(t.a) │ - /// └────────────────────┘ - /// ▲ - /// │ - /// │ - /// ┌────────────────────┐ Filter ┌────────────────────┐ Projection ┌────────────────────┐ - /// │ Filter(t.b > 5) │────Pushdown──▶ │ Projection(t.a) │ ───Pushdown───▶ │ Projection(t.a) │ - /// └────────────────────┘ └────────────────────┘ └────────────────────┘ - /// ▲ ▲ ▲ - /// │ │ │ - /// │ │ ┌────────────────────┐ - /// ┌────────────────────┐ ┌────────────────────┐ │ Scan │ - /// │ Scan │ │ Scan │ │ filter=(t.b > 5) │ - /// └────────────────────┘ │ filter=(t.b > 5) │ │ projection=(t.a) │ - /// └────────────────────┘ └────────────────────┘ - /// - /// Initial Plan If `TableProviderFilterPushDown` Projection pushdown notes that - /// returns true, filter pushdown the scan only needs t.a - /// pushes the filter into the scan - /// BUT internally evaluating the - /// predicate still requires t.b - /// ``` - /// - /// # Limit - /// - /// If `limit` is specified, the scan must produce *at least* this many - /// rows, though it may return more. Like Projection Pushdown and Filter - /// Pushdown, DataFusion pushes `LIMIT`s as far down in the plan as - /// possible. This is called "Limit Pushdown", and some sources can use the - /// information to improve performance. - /// - /// Note: If any pushed-down filters are `Inexact`, the `LIMIT` cannot be - /// pushed down. Inexact filters do not guarantee that every filtered row is - /// removed, so applying the limit could leave too few rows to return in the - /// final result. - /// - /// # Evaluation Order - /// - /// The logical evaluation order is `filters`, then `limit`, then - /// `projection`. - /// - /// Note that `limit` applies to the filtered result, not to the unfiltered - /// input, and `projection` affects only which columns are returned, not - /// which rows qualify. - /// - /// For example, if a scan receives: - /// - /// - `projection = [a]` - /// - `filters = [b > 5]` - /// - `limit = Some(3)` - /// - /// It must logically produce results equivalent to: - /// - /// ```text - /// PROJECTION a (LIMIT 3 (SCAN WHERE b > 5)) - /// ``` - /// - /// As noted above, columns referenced only by pushed-down filters may be - /// absent from `projection`. - async fn scan( - &self, - state: &dyn Session, - projection: Option<&Vec>, - filters: &[Expr], - limit: Option, - ) -> Result>; - - /// Create an [`ExecutionPlan`] for scanning the table using structured arguments. - /// - /// This method uses [`ScanArgs`] to pass scan parameters in a structured way - /// and returns a [`ScanResult`] containing the execution plan. - /// - /// Table providers can override this method to take advantage of additional - /// parameters like the upcoming `preferred_ordering` that may not be available through - /// other scan methods. - /// - /// # Arguments - /// * `state` - The session state containing configuration and context - /// * `args` - Structured scan arguments including projection, filters, limit, and ordering preferences - /// - /// # Returns - /// A [`ScanResult`] containing the [`ExecutionPlan`] for scanning the table - /// - /// See [`Self::scan`] for detailed documentation about projection, filters, and limits. - async fn scan_with_args<'a>( - &self, - state: &dyn Session, - args: ScanArgs<'a>, - ) -> Result { - let filters = args.filters().unwrap_or(&[]); - let projection = args.projection().map(|p| p.to_vec()); - let limit = args.limit(); - let plan = self - .scan(state, projection.as_ref(), filters, limit) - .await?; - Ok(plan.into()) - } - - /// Specify if DataFusion should provide filter expressions to the - /// TableProvider to apply *during* the scan. - /// - /// Some TableProviders can evaluate filters more efficiently than the - /// `Filter` operator in DataFusion, for example by using an index. - /// - /// # Parameters and Return Value - /// - /// The return `Vec` must have one element for each element of the `filters` - /// argument. The value of each element indicates if the TableProvider can - /// apply the corresponding filter during the scan. The position in the return - /// value corresponds to the expression in the `filters` parameter. - /// - /// If the length of the resulting `Vec` does not match the `filters` input - /// an error will be thrown. - /// - /// Each element in the resulting `Vec` is one of the following: - /// * [`Exact`] or [`Inexact`]: The TableProvider can apply the filter - /// during scan - /// * [`Unsupported`]: The TableProvider cannot apply the filter during scan - /// - /// By default, this function returns [`Unsupported`] for all filters, - /// meaning no filters will be provided to [`Self::scan`]. - /// - /// [`Unsupported`]: TableProviderFilterPushDown::Unsupported - /// [`Exact`]: TableProviderFilterPushDown::Exact - /// [`Inexact`]: TableProviderFilterPushDown::Inexact - /// # Example - /// - /// ```rust - /// # use std::any::Any; - /// # use std::sync::Arc; - /// # use arrow::datatypes::SchemaRef; - /// # use async_trait::async_trait; - /// # use datafusion_catalog::{TableProvider, Session}; - /// # use datafusion_common::Result; - /// # use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; - /// # use datafusion_physical_plan::ExecutionPlan; - /// // Define a struct that implements the TableProvider trait - /// #[derive(Debug)] - /// struct TestDataSource {} - /// - /// #[async_trait] - /// impl TableProvider for TestDataSource { - /// # fn schema(&self) -> SchemaRef { todo!() } - /// # fn table_type(&self) -> TableType { todo!() } - /// # async fn scan(&self, s: &dyn Session, p: Option<&Vec>, f: &[Expr], l: Option) -> Result> { - /// todo!() - /// # } - /// // Override the supports_filters_pushdown to evaluate which expressions - /// // to accept as pushdown predicates. - /// fn supports_filters_pushdown(&self, filters: &[&Expr]) -> Result> { - /// // Process each filter - /// let support: Vec<_> = filters.iter().map(|expr| { - /// match expr { - /// // This example only supports a between expr with a single column named "c1". - /// Expr::Between(between_expr) => { - /// between_expr.expr - /// .try_as_col() - /// .map(|column| { - /// if column.name == "c1" { - /// TableProviderFilterPushDown::Exact - /// } else { - /// TableProviderFilterPushDown::Unsupported - /// } - /// }) - /// // If there is no column in the expr set the filter to unsupported. - /// .unwrap_or(TableProviderFilterPushDown::Unsupported) - /// } - /// _ => { - /// // For all other cases return Unsupported. - /// TableProviderFilterPushDown::Unsupported - /// } - /// } - /// }).collect(); - /// Ok(support) - /// } - /// } - /// ``` - fn supports_filters_pushdown( - &self, - filters: &[&Expr], - ) -> Result> { - Ok(vec![ - TableProviderFilterPushDown::Unsupported; - filters.len() - ]) - } - - /// Get statistics for this table, if available - /// Although not presently used in mainline DataFusion, this allows implementation specific - /// behavior for downstream repositories, in conjunction with specialized optimizer rules to - /// perform operations such as re-ordering of joins. - fn statistics(&self) -> Option { - None - } - - /// Return an [`ExecutionPlan`] to insert data into this table, if - /// supported. - /// - /// The returned plan should return a single row in a UInt64 - /// column called "count" such as the following - /// - /// ```text - /// +-------+, - /// | count |, - /// +-------+, - /// | 6 |, - /// +-------+, - /// ``` - /// - /// # See Also - /// - /// See [`DataSinkExec`] for the common pattern of inserting a - /// streams of `RecordBatch`es as files to an ObjectStore. - /// - /// [`DataSinkExec`]: datafusion_datasource::sink::DataSinkExec - async fn insert_into( - &self, - _state: &dyn Session, - _input: Arc, - _insert_op: InsertOp, - ) -> Result> { - not_impl_err!("Insert into not implemented for this table") - } - - /// Delete rows matching the filter predicates. - /// - /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). - /// Empty `filters` deletes all rows. - async fn delete_from( - &self, - _state: &dyn Session, - _filters: Vec, - ) -> Result> { - not_impl_err!("DELETE not supported for {} table", self.table_type()) - } - - /// Update rows matching the filter predicates. - /// - /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). - /// Empty `filters` updates all rows. - async fn update( - &self, - _state: &dyn Session, - _assignments: Vec<(String, Expr)>, - _filters: Vec, - ) -> Result> { - not_impl_err!("UPDATE not supported for {} table", self.table_type()) - } - - /// Remove all rows from the table. - /// - /// Should return an [ExecutionPlan] producing a single row with count (UInt64), - /// representing the number of rows removed. - async fn truncate(&self, _state: &dyn Session) -> Result> { - not_impl_err!("TRUNCATE not supported for {} table", self.table_type()) - } -} - -impl dyn TableProvider { - /// Returns `true` if the table provider is of type `T`. - /// - /// Prefer this over `downcast_ref::().is_some()`. Works correctly when - /// called on `Arc` via auto-deref. - pub fn is(&self) -> bool { - (self as &dyn Any).is::() - } - - /// Attempts to downcast this table provider to a concrete type `T`, - /// returning `None` if the provider is not of that type. - /// - /// Works correctly when called on `Arc` via auto-deref, - /// unlike `(&arc as &dyn Any).downcast_ref::()` which would attempt to - /// downcast the `Arc` itself. - pub fn downcast_ref(&self) -> Option<&T> { - (self as &dyn Any).downcast_ref() - } -} - -/// Arguments for scanning a table with [`TableProvider::scan_with_args`]. -#[derive(Debug, Clone, Default)] -pub struct ScanArgs<'a> { - filters: Option<&'a [Expr]>, - projection: Option<&'a [usize]>, - limit: Option, -} - -impl<'a> ScanArgs<'a> { - /// Set the column projection for the scan. - /// - /// The projection is a list of column indices from [`TableProvider::schema`] - /// that should be included in the scan results. If `None`, all columns are included. - /// - /// # Arguments - /// * `projection` - Optional slice of column indices to project - pub fn with_projection(mut self, projection: Option<&'a [usize]>) -> Self { - self.projection = projection; - self - } - - /// Get the column projection for the scan. - /// - /// Returns a reference to the projection column indices, or `None` if - /// no projection was specified (meaning all columns should be included). - pub fn projection(&self) -> Option<&'a [usize]> { - self.projection - } - - /// Set the filter expressions for the scan. - /// - /// Filters are boolean expressions that should be evaluated during the scan - /// to reduce the number of rows returned. All expressions are combined with AND logic. - /// Whether filters are actually pushed down depends on [`TableProvider::supports_filters_pushdown`]. - /// - /// # Arguments - /// * `filters` - Optional slice of filter expressions - pub fn with_filters(mut self, filters: Option<&'a [Expr]>) -> Self { - self.filters = filters; - self - } - - /// Get the filter expressions for the scan. - /// - /// Returns a reference to the filter expressions, or `None` if no filters were specified. - pub fn filters(&self) -> Option<&'a [Expr]> { - self.filters - } - - /// Set the maximum number of rows to return from the scan. - /// - /// If specified, the scan should return at most this many rows. This is typically - /// used to optimize queries with `LIMIT` clauses. - /// - /// # Arguments - /// * `limit` - Optional maximum number of rows to return - pub fn with_limit(mut self, limit: Option) -> Self { - self.limit = limit; - self - } - - /// Get the maximum number of rows to return from the scan. - /// - /// Returns the row limit, or `None` if no limit was specified. - pub fn limit(&self) -> Option { - self.limit - } -} - -/// Result of a table scan operation from [`TableProvider::scan_with_args`]. -#[derive(Debug, Clone)] -pub struct ScanResult { - /// The ExecutionPlan to run. - plan: Arc, -} - -impl ScanResult { - /// Create a new `ScanResult` with the given execution plan. - /// - /// # Arguments - /// * `plan` - The execution plan that will perform the table scan - pub fn new(plan: Arc) -> Self { - Self { plan } - } - - /// Get a reference to the execution plan for this scan result. - /// - /// Returns a reference to the [`ExecutionPlan`] that will perform - /// the actual table scanning and data retrieval. - pub fn plan(&self) -> &Arc { - &self.plan - } - - /// Consume this ScanResult and return the execution plan. - /// - /// Returns the owned [`ExecutionPlan`] that will perform - /// the actual table scanning and data retrieval. - pub fn into_inner(self) -> Arc { - self.plan - } -} - -impl From> for ScanResult { - fn from(plan: Arc) -> Self { - Self::new(plan) - } -} - -/// A factory which creates [`TableProvider`]s at runtime given a URL. -/// -/// For example, this can be used to create a table "on the fly" -/// from a directory of files only when that name is referenced. -#[async_trait] -pub trait TableProviderFactory: Debug + Sync + Send { - /// Create a TableProvider with the given url - async fn create( - &self, - state: &dyn Session, - cmd: &CreateExternalTable, - ) -> Result>; -} - -/// Describes arguments provided to the table function call. -pub struct TableFunctionArgs<'e, 's> { - /// Call arguments. - exprs: &'e [Expr], - /// Session within which the function is called. - session: &'s dyn Session, -} - -impl<'e, 's> TableFunctionArgs<'e, 's> { - /// Make a new [`TableFunctionArgs`]. - pub fn new(exprs: &'e [Expr], session: &'s dyn Session) -> Self { - Self { exprs, session } - } - - /// Get expressions passed as the called function arguments. - pub fn exprs(&self) -> &'e [Expr] { - self.exprs - } - - /// Get a session where the table function is called. - pub fn session(&self) -> &'s dyn Session { - self.session - } -} - -/// A trait for table function implementations -pub trait TableFunctionImpl: Debug + Sync + Send + Any { - /// Create a table provider - #[deprecated( - since = "53.0.0", - note = "Implement `TableFunctionImpl::call_with_args` instead" - )] - fn call(&self, _exprs: &[Expr]) -> Result> { - internal_err!( - "TableFunctionImpl::call is not implemented. Implement TableFunctionImpl::call_with_args instead." - ) - } - - /// Create a table provider - fn call_with_args(&self, args: TableFunctionArgs) -> Result> { - #[expect(deprecated)] - self.call(args.exprs) - } -} - -/// A table that uses a function to generate data -#[derive(Clone, Debug)] -pub struct TableFunction { - /// Name of the table function - name: String, - /// Function implementation - fun: Arc, -} - -impl TableFunction { - /// Create a new table function - pub fn new(name: String, fun: Arc) -> Self { - Self { name, fun } - } - - /// Get the name of the table function - pub fn name(&self) -> &str { - &self.name - } - - /// Get the implementation of the table function - pub fn function(&self) -> &Arc { - &self.fun - } - - /// Get the function implementation and generate a table - #[deprecated( - since = "53.0.0", - note = "Use `TableFunction::create_table_provider_with_args` instead" - )] - pub fn create_table_provider(&self, args: &[Expr]) -> Result> { - #[expect(deprecated)] - self.fun.call(args) - } - - /// Get the function implementation and generate a table - pub fn create_table_provider_with_args( - &self, - args: TableFunctionArgs, - ) -> Result> { - self.fun.call_with_args(args) - } -} diff --git a/datafusion/catalog/src/view.rs b/datafusion/catalog/src/view.rs index 45084e65f23f2..723634c34c0f7 100644 --- a/datafusion/catalog/src/view.rs +++ b/datafusion/catalog/src/view.rs @@ -59,17 +59,6 @@ impl ViewTable { } } - #[deprecated( - since = "47.0.0", - note = "Use `ViewTable::new` instead and apply TypeCoercion to the logical plan if needed" - )] - pub fn try_new( - logical_plan: LogicalPlan, - definition: Option, - ) -> Result { - Ok(Self::new(logical_plan, definition)) - } - /// Get definition ref pub fn definition(&self) -> Option<&String> { self.definition.as_ref() diff --git a/datafusion/common-runtime/src/join_set.rs b/datafusion/common-runtime/src/join_set.rs index 1857a4111dbcb..3ac5912243817 100644 --- a/datafusion/common-runtime/src/join_set.rs +++ b/datafusion/common-runtime/src/join_set.rs @@ -21,10 +21,13 @@ use std::task::{Context, Poll}; use tokio::runtime::Handle; use tokio::task::{AbortHandle, Id, JoinError, LocalSet}; -/// A wrapper around Tokio's JoinSet that forwards all API calls while optionally +/// A wrapper around [Tokio's `JoinSet`] that forwards all API calls while optionally /// instrumenting spawned tasks and blocking closures with custom tracing behavior. -/// If no tracer is injected via `trace_utils::set_tracer`, tasks and closures are executed +/// If no tracer is injected via [`set_join_set_tracer`], tasks and closures are executed /// without any instrumentation. +/// +/// [Tokio's `JoinSet`]: tokio::task::JoinSet +/// [`set_join_set_tracer`]: crate::trace_utils::set_join_set_tracer #[derive(Debug)] pub struct JoinSet { inner: tokio::task::JoinSet, diff --git a/datafusion/common/Cargo.toml b/datafusion/common/Cargo.toml index 740d4e45b8d05..1eb23089a4021 100644 --- a/datafusion/common/Cargo.toml +++ b/datafusion/common/Cargo.toml @@ -77,6 +77,7 @@ indexmap = { workspace = true } itertools = { workspace = true } libc = "0.2.185" log = { workspace = true } +num-traits = { workspace = true } object_store = { workspace = true, optional = true } parquet = { workspace = true, optional = true, default-features = true } recursive = { workspace = true, optional = true } diff --git a/datafusion/common/src/column.rs b/datafusion/common/src/column.rs index c7f0b5a4f4881..f8893aa423fa1 100644 --- a/datafusion/common/src/column.rs +++ b/datafusion/common/src/column.rs @@ -271,7 +271,7 @@ impl Column { }) .map_err(|err| { let mut diagnostic = Diagnostic::new_error( - format!("column '{}' is ambiguous", &self.name), + format!("column '{}' is ambiguous", self.name), self.spans().first(), ); // TODO If [`DFSchema`] had spans, we could show the @@ -439,7 +439,7 @@ mod tests { &[], ) .expect_err("should've failed to find field"); - let expected = "Schema error: No field named z. \ + let expected = "Schema error: No field named z.\n\ Valid fields are t1.a, t1.b, t2.c, t2.d, t3.a, t3.b, t3.c, t3.d, t3.e."; assert_eq!(err.strip_backtrace(), expected); diff --git a/datafusion/common/src/config.rs b/datafusion/common/src/config.rs index e6d1ebbbbe746..f5742f09f9b08 100644 --- a/datafusion/common/src/config.rs +++ b/datafusion/common/src/config.rs @@ -21,7 +21,7 @@ use arrow_ipc::CompressionType; #[cfg(feature = "parquet_encryption")] use crate::encryption::{FileDecryptionProperties, FileEncryptionProperties}; -use crate::error::_config_err; +use crate::error::{_config_datafusion_err, _config_err}; use crate::format::{ExplainAnalyzeCategories, ExplainFormat, MetricType}; use crate::parquet_config::DFParquetWriterVersion; use crate::parsers::{CompressionTypeVariant, CsvQuoteStyle}; @@ -33,6 +33,7 @@ use std::any::Any; use std::collections::{BTreeMap, HashMap}; use std::error::Error; use std::fmt::{self, Display}; +use std::num::NonZeroUsize; use std::str::FromStr; #[cfg(feature = "parquet_encryption")] use std::sync::Arc; @@ -278,8 +279,8 @@ config_namespace! { /// are normalized automatically. pub enable_options_value_normalization: bool, warn = "`enable_options_value_normalization` is deprecated and ignored", default = false - /// Configure the SQL dialect used by DataFusion's parser; supported values include: Generic, - /// MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, MsSQL, ClickHouse, BigQuery, Ansi, DuckDB and Databricks. + /// Configure the SQL dialect used by DataFusion's parser. + /// The configuration reference lists the supported values from [`Dialect::available`]. pub dialect: Dialect, default = Dialect::Generic // no need to lowercase because `sqlparser::dialect_from_str`] is case-insensitive @@ -300,7 +301,7 @@ config_namespace! { pub collect_spans: bool, default = false /// Specifies the recursion depth limit when parsing complex SQL Queries - pub recursion_limit: usize, default = 50 + pub recursion_limit: ConfigNonZeroUsize, default = non_zero_usize_default(50) /// Specifies the default null ordering for query results. There are 4 options: /// - `nulls_max`: Nulls appear last in ascending order. @@ -323,44 +324,172 @@ config_namespace! { } } -/// This is the SQL dialect used by DataFusion's parser. -/// This mirrors [sqlparser::dialect::Dialect](https://docs.rs/sqlparser/latest/sqlparser/dialect/trait.Dialect.html) -/// trait in order to offer an easier API and avoid adding the `sqlparser` dependency -#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] -pub enum Dialect { - #[default] - Generic, - MySQL, - PostgreSQL, - Hive, - SQLite, - Snowflake, - Redshift, - MsSQL, - ClickHouse, - BigQuery, - Ansi, - DuckDB, - Databricks, +/// Metadata for a SQL dialect supported by DataFusion configuration. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +#[non_exhaustive] +pub struct DialectInfo { + pub dialect: Dialect, + pub canonical_name: &'static str, + pub display_name: &'static str, + pub aliases: &'static [&'static str], +} + +// Keep this key in sync with the `SqlParserOptions::dialect` config path. +const SQL_PARSER_DIALECT_CONFIG_KEY: &str = "datafusion.sql_parser.dialect"; + +macro_rules! dialect_display_list { + ($($display_name:literal),+ $(,)?) => { + dialect_display_list!(@acc [] $($display_name),+) + }; + (@acc [$($acc:tt)*] $last:literal) => { + concat!($($acc)* $last) + }; + (@acc [$($acc:tt)*] $next:literal, $($rest:literal),+) => { + dialect_display_list!(@acc [$($acc)* $next, ", ",] $($rest),+) + }; +} + +macro_rules! dialect_metadata { + ( + default: $default_variant:ident; + $( + $variant:ident { + canonical_name: $canonical_name:literal, + display_name: $display_name:literal, + aliases: [$($alias:literal),* $(,)?], + } + ),+ $(,)? + ) => { + /// This is the SQL dialect used by DataFusion's parser. + /// This mirrors [sqlparser::dialect::Dialect](https://docs.rs/sqlparser/latest/sqlparser/dialect/trait.Dialect.html) + /// trait in order to offer an easier API and avoid adding the `sqlparser` dependency + #[derive(Debug, Clone, Copy, PartialEq, Eq)] + pub enum Dialect { + $($variant,)+ + } + + impl Default for Dialect { + fn default() -> Self { + Self::$default_variant + } + } + + const DIALECT_INFOS: &[DialectInfo] = &[ + $( + DialectInfo { + dialect: Dialect::$variant, + canonical_name: $canonical_name, + display_name: $display_name, + aliases: &[$($alias),*], + }, + )+ + ]; + + const AVAILABLE_DIALECTS: &str = dialect_display_list!($($display_name),+); + const DIALECT_CONFIG_DESCRIPTION: &str = concat!( + "Configure the SQL dialect used by DataFusion's parser; supported values include: ", + dialect_display_list!($($display_name),+), + "." + ); + }; +} + +dialect_metadata! { + default: Generic; + Generic { + canonical_name: "generic", + display_name: "Generic", + aliases: [], + }, + MySQL { + canonical_name: "mysql", + display_name: "MySQL", + aliases: [], + }, + PostgreSQL { + canonical_name: "postgresql", + display_name: "PostgreSQL", + aliases: ["postgres"], + }, + Hive { + canonical_name: "hive", + display_name: "Hive", + aliases: [], + }, + SQLite { + canonical_name: "sqlite", + display_name: "SQLite", + aliases: [], + }, + Snowflake { + canonical_name: "snowflake", + display_name: "Snowflake", + aliases: [], + }, + Redshift { + canonical_name: "redshift", + display_name: "Redshift", + aliases: [], + }, + MsSQL { + canonical_name: "mssql", + display_name: "MsSQL", + aliases: [], + }, + ClickHouse { + canonical_name: "clickhouse", + display_name: "ClickHouse", + aliases: [], + }, + BigQuery { + canonical_name: "bigquery", + display_name: "BigQuery", + aliases: [], + }, + Ansi { + canonical_name: "ansi", + display_name: "Ansi", + aliases: [], + }, + DuckDB { + canonical_name: "duckdb", + display_name: "DuckDB", + aliases: [], + }, + Databricks { + canonical_name: "databricks", + display_name: "Databricks", + aliases: [], + }, + Spark { + canonical_name: "spark", + display_name: "Spark", + aliases: ["sparksql"], + }, +} + +impl Dialect { + /// Return metadata for all supported dialects. + pub fn metadata() -> &'static [DialectInfo] { + DIALECT_INFOS + } + + /// Return all supported dialect names, for use in error messages. + pub fn available() -> &'static str { + AVAILABLE_DIALECTS + } + + fn info(&self) -> &'static DialectInfo { + DIALECT_INFOS + .iter() + .find(|info| info.dialect == *self) + .expect("all Dialect variants are listed in DIALECT_INFOS") + } } impl AsRef for Dialect { fn as_ref(&self) -> &str { - match self { - Self::Generic => "generic", - Self::MySQL => "mysql", - Self::PostgreSQL => "postgresql", - Self::Hive => "hive", - Self::SQLite => "sqlite", - Self::Snowflake => "snowflake", - Self::Redshift => "redshift", - Self::MsSQL => "mssql", - Self::ClickHouse => "clickhouse", - Self::BigQuery => "bigquery", - Self::Ansi => "ansi", - Self::DuckDB => "duckdb", - Self::Databricks => "databricks", - } + self.info().canonical_name } } @@ -368,33 +497,31 @@ impl FromStr for Dialect { type Err = DataFusionError; fn from_str(s: &str) -> Result { - let value = match s.to_ascii_lowercase().as_str() { - "generic" => Self::Generic, - "mysql" => Self::MySQL, - "postgresql" | "postgres" => Self::PostgreSQL, - "hive" => Self::Hive, - "sqlite" => Self::SQLite, - "snowflake" => Self::Snowflake, - "redshift" => Self::Redshift, - "mssql" => Self::MsSQL, - "clickhouse" => Self::ClickHouse, - "bigquery" => Self::BigQuery, - "ansi" => Self::Ansi, - "duckdb" => Self::DuckDB, - "databricks" => Self::Databricks, - other => { - let error_message = format!( - "Invalid Dialect: {other}. Expected one of: Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, MsSQL, ClickHouse, BigQuery, Ansi, DuckDB, Databricks" - ); - return Err(DataFusionError::Configuration(error_message)); + for info in DIALECT_INFOS { + if info.canonical_name.eq_ignore_ascii_case(s) + || info + .aliases + .iter() + .any(|alias| alias.eq_ignore_ascii_case(s)) + { + return Ok(info.dialect); } - }; - Ok(value) + } + + Err(DataFusionError::Configuration(format!( + "Invalid Dialect: {s}. Expected one of: {}", + Self::available() + ))) } } impl ConfigField for Dialect { fn visit(&self, v: &mut V, key: &str, description: &'static str) { + let description = if key == SQL_PARSER_DIALECT_CONFIG_KEY { + DIALECT_CONFIG_DESCRIPTION + } else { + description + }; v.some(key, self, description) } @@ -456,6 +583,218 @@ impl Display for SpillCompression { } } +/// A `usize` configuration value that rejects zero when set from strings. +/// +/// Use this for options where zero is never a meaningful runtime value. +/// Invalid values return a configuration error through [`ConfigField`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ConfigNonZeroUsize(NonZeroUsize); + +/// Private helper for hard-coded defaults in `config_namespace!`, which cannot +/// use `?`. All external construction should use [`ConfigNonZeroUsize::try_new`]. +const fn non_zero_usize_default(value: usize) -> ConfigNonZeroUsize { + match NonZeroUsize::new(value) { + Some(value) => ConfigNonZeroUsize(value), + None => panic!("value must be greater than 0"), + } +} + +impl ConfigNonZeroUsize { + /// Creates a [`ConfigNonZeroUsize`], returning a configuration error if + /// `value` is zero. + pub fn try_new(value: usize) -> Result { + NonZeroUsize::new(value) + .map(Self) + .ok_or_else(|| _config_datafusion_err!("value must be greater than 0")) + } + + /// Returns the wrapped `usize`. + pub const fn get(self) -> usize { + self.0.get() + } +} + +impl From for usize { + fn from(value: ConfigNonZeroUsize) -> Self { + value.get() + } +} + +impl FromStr for ConfigNonZeroUsize { + type Err = DataFusionError; + + fn from_str(s: &str) -> Result { + Self::try_new(default_config_transform(s)?) + } +} + +impl ConfigField for ConfigNonZeroUsize { + fn visit(&self, v: &mut V, key: &str, description: &'static str) { + v.some(key, self, description) + } + + fn set(&mut self, key: &str, value: &str) -> Result<()> { + if !key.is_empty() { + return _config_err!( + "Config field batch_size is a scalar ConfigNonZeroUsize and does not have nested field \"{}\"", + key + ); + } + + *self = ConfigNonZeroUsize::from_str(value)?; + Ok(()) + } + + fn reset(&mut self, key: &str) -> Result<()> { + if key.is_empty() { + Ok(()) + } else { + _config_err!( + "Config field batch_size is a scalar ConfigNonZeroUsize and does not have nested field \"{}\"", + key + ) + } + } +} + +impl Display for ConfigNonZeroUsize { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.get()) + } +} + +/// A `usize` configuration value that rejects 0 and 1 when set from strings. +/// +/// Use this for options whose consumer divides the value in half to size an +/// internal buffer (e.g. a bounded channel capacity): values below 2 would +/// round down to a zero-capacity buffer and panic. Invalid values return a +/// configuration error through [`ConfigField`] instead. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)] +pub struct ConfigMinTwoUsize(usize); + +/// Private helper for hard-coded defaults in `config_namespace!`, which cannot +/// use `?`. All external construction should use [`ConfigMinTwoUsize::try_new`]. +const fn min_two_usize_default(value: usize) -> ConfigMinTwoUsize { + if value >= 2 { + ConfigMinTwoUsize(value) + } else { + panic!("value must be at least 2") + } +} + +impl ConfigMinTwoUsize { + /// Creates a [`ConfigMinTwoUsize`], returning a configuration error if + /// `value` is less than 2. + pub fn try_new(value: usize) -> Result { + if value >= 2 { + Ok(Self(value)) + } else { + _config_err!("value must be at least 2") + } + } + + /// Returns the wrapped `usize`. + pub const fn get(self) -> usize { + self.0 + } +} + +impl From for usize { + fn from(value: ConfigMinTwoUsize) -> Self { + value.get() + } +} + +impl FromStr for ConfigMinTwoUsize { + type Err = DataFusionError; + + fn from_str(s: &str) -> Result { + Self::try_new(default_config_transform(s)?) + } +} + +impl ConfigField for ConfigMinTwoUsize { + fn visit(&self, v: &mut V, key: &str, description: &'static str) { + v.some(key, self, description) + } + + fn set(&mut self, key: &str, value: &str) -> Result<()> { + if !key.is_empty() { + return _config_err!( + "Config field max_buffered_batches_per_output_file is a scalar ConfigMinTwoUsize and does not have nested field \"{}\"", + key + ); + } + + *self = ConfigMinTwoUsize::from_str(value)?; + Ok(()) + } + + fn reset(&mut self, key: &str) -> Result<()> { + if key.is_empty() { + Ok(()) + } else { + _config_err!( + "Config field max_buffered_batches_per_output_file is a scalar ConfigMinTwoUsize and does not have nested field \"{}\"", + key + ) + } + } +} + +impl Display for ConfigMinTwoUsize { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.get()) + } +} + +/// Policy for handling duplicate keys in Spark-compatible map-construction +/// functions (`map_from_arrays`, `map_from_entries`, `str_to_map`). Mirrors +/// Spark's [`spark.sql.mapKeyDedupPolicy`](https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961). +#[derive(Debug, Default, Clone, Copy, PartialEq, Eq)] +pub enum MapKeyDedupPolicy { + /// Raise `[DUPLICATED_MAP_KEY]` at runtime on any duplicate key. + #[default] + Exception, + /// Keep the last occurrence of each duplicate key. + LastWin, +} + +impl FromStr for MapKeyDedupPolicy { + type Err = DataFusionError; + + fn from_str(s: &str) -> Result { + match s.to_ascii_uppercase().as_str() { + "EXCEPTION" => Ok(Self::Exception), + "LAST_WIN" => Ok(Self::LastWin), + other => Err(DataFusionError::Configuration(format!( + "Invalid MapKeyDedupPolicy: {other}. Expected one of: EXCEPTION, LAST_WIN" + ))), + } + } +} + +impl ConfigField for MapKeyDedupPolicy { + fn visit(&self, v: &mut V, key: &str, description: &'static str) { + v.some(key, self, description) + } + + fn set(&mut self, _: &str, value: &str) -> Result<()> { + *self = MapKeyDedupPolicy::from_str(value)?; + Ok(()) + } +} + +impl Display for MapKeyDedupPolicy { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let str = match self { + Self::Exception => "EXCEPTION", + Self::LastWin => "LAST_WIN", + }; + write!(f, "{str}") + } +} + impl From for Option { fn from(c: SpillCompression) -> Self { match c { @@ -476,7 +815,7 @@ config_namespace! { /// Default batch size while creating new batches, it's especially useful for /// buffer-in-memory batches since creating tiny batches would result in too much /// metadata memory consumption - pub batch_size: usize, default = 8192 + pub batch_size: ConfigNonZeroUsize, default = non_zero_usize_default(8192) /// A perfect hash join (see `HashJoinExec` for more details) will be considered /// if the range of keys (max - min) on the build side is < this threshold. @@ -504,8 +843,7 @@ config_namespace! { pub coalesce_batches: bool, default = true /// Should DataFusion collect statistics when first creating a table. - /// Has no effect after the table is created. Applies to the default - /// `ListingTableProvider` in DataFusion. Defaults to true. + /// Has no effect after the table is created. Defaults to true. pub collect_statistics: bool, default = true /// Number of partitions for query execution. Increasing partitions can increase @@ -540,6 +878,17 @@ config_namespace! { /// the new schema verification step. pub skip_physical_aggregate_schema_check: bool, default = false + /// Temporary switch for aggregate stream implementations that are being + /// migrated from `GroupedHashAggregateStream`. + /// + /// When set to true, DataFusion tries the migrated implementations when + /// their preconditions are satisfied. When set to false, grouped + /// aggregation falls back to `GroupedHashAggregateStream`. This option + /// will be removed after the migration is finished. + /// + /// See for details. + pub enable_migration_aggregate: bool, default = true + /// Sets the compression codec used when spilling data to disk. /// /// Since datafusion writes spill files using the Arrow IPC Stream format, @@ -594,28 +943,36 @@ config_namespace! { /// may create spill files larger than the limit. /// /// Default: 128 MB - pub max_spill_file_size_bytes: usize, default = 128 * 1024 * 1024 + pub max_spill_file_size_bytes: ConfigNonZeroUsize, default = non_zero_usize_default(128 * 1024 * 1024) /// Number of files to read in parallel when inferring schema and statistics - pub meta_fetch_concurrency: usize, default = 32 + pub meta_fetch_concurrency: ConfigNonZeroUsize, default = non_zero_usize_default(32) /// Guarantees a minimum level of output files running in parallel. /// RecordBatches will be distributed in round robin fashion to each /// parallel writer. Each writer is closed and a new file opened once /// soft_max_rows_per_output_file is reached. - pub minimum_parallel_output_files: usize, default = 4 + pub minimum_parallel_output_files: ConfigNonZeroUsize, default = non_zero_usize_default(4) /// Target number of rows in output files when writing multiple. /// This is a soft max, so it can be exceeded slightly. There also /// will be one file smaller than the limit if the total /// number of rows written is not roughly divisible by the soft max - pub soft_max_rows_per_output_file: usize, default = 50000000 + pub soft_max_rows_per_output_file: ConfigNonZeroUsize, default = non_zero_usize_default(50000000) /// This is the maximum number of RecordBatches buffered /// for each output file being worked. Higher values can potentially /// give faster write performance at the cost of higher peak - /// memory consumption - pub max_buffered_batches_per_output_file: usize, default = 2 + /// memory consumption. + /// + /// This budget is split evenly between two independent points in the + /// write pipeline (see the demuxer diagram in #7791): how many files + /// can be in flight from the demuxer to a writer task, and how many + /// RecordBatches are buffered for a single file's writer. Must be at + /// least 2 so each half gets at least 1 unit of buffering - 0 or 1 + /// would leave one side with a zero-capacity channel and panic at + /// write time. + pub max_buffered_batches_per_output_file: ConfigMinTwoUsize, default = min_two_usize_default(2) /// Should sub directories be ignored when scanning directories for data /// files. Defaults to true (ignores subdirectories), consistent with @@ -639,6 +996,19 @@ config_namespace! { /// Should DataFusion keep the columns used for partition_by in the output RecordBatches pub keep_partition_by_columns: bool, default = false + /// When `true` (the default), DataFusion's built-in file scans + /// dynamically rebalance files across partitions at query execution + /// time: a partition that goes idle reads files (or byte-range morsels) + /// originally assigned to a sibling partition, which keeps all + /// partitions busy in a single process. + /// + /// Executors that depend on the plan-time partition assignment — such as + /// Ballista and datafusion-distributed, which run each partition as an + /// isolated task and never poll the siblings — should set this to + /// `false` so each partition reads only its own file group and no + /// runtime reassignment occurs. + pub enable_file_stream_work_stealing: bool, default = true + /// Aggregation ratio (number of distinct groups / number of input rows) /// threshold for skipping partial aggregation. If the value is greater /// then partial aggregation will skip aggregation for further input @@ -709,134 +1079,117 @@ config_namespace! { } } -/// Options for content-defined chunking (CDC) when writing parquet files. -/// See [`ParquetOptions::use_content_defined_chunking`]. -/// -/// Can be enabled with default options by setting -/// `use_content_defined_chunking` to `true`, or configured with sub-fields -/// like `use_content_defined_chunking.min_chunk_size`. -#[derive(Debug, Clone, PartialEq)] -pub struct CdcOptions { - /// Minimum chunk size in bytes. The rolling hash will not trigger a split - /// until this many bytes have been accumulated. Default is 256 KiB. - pub min_chunk_size: usize, - - /// Maximum chunk size in bytes. A split is forced when the accumulated - /// size exceeds this value. Default is 1 MiB. - pub max_chunk_size: usize, - - /// Normalization level. Increasing this improves deduplication ratio - /// but increases fragmentation. Recommended range is [-3, 3], default is 0. - pub norm_level: i32, -} - -// Note: `CdcOptions` intentionally does NOT implement `Default` so that the -// blanket `impl ConfigField for Option` does not -// apply. This allows the specific `impl ConfigField for Option` -// below to handle "true"/"false" for enabling/disabling CDC. -// Use `CdcOptions::default()` (the inherent method) instead of `Default::default()`. -impl CdcOptions { - /// Returns a new `CdcOptions` with default values. - #[expect(clippy::should_implement_trait)] - pub fn default() -> Self { +config_namespace! { + /// Options for content-defined chunking (CDC) when writing parquet files. + /// Mirrors `parquet::file::properties::CdcOptions`. + /// + /// Carried as a [`ParquetCdcOptions`] in [`ParquetOptions::content_defined_chunking`] + /// with an explicit `enabled` flag, so it can be toggled with dotted config + /// keys (`content_defined_chunking.enabled = true|false`) and the result is + /// independent of the order in which the keys are set. + pub struct ParquetCdcOptions { + /// (writing) EXPERIMENTAL: Enable content-defined chunking (CDC) when writing + /// parquet files. When enabled, parallel writing is automatically disabled + /// since the chunker state must persist across row groups. + pub enabled: bool, default = false + + /// Minimum chunk size in bytes. The rolling hash will not trigger a split + /// until this many bytes have been accumulated. Default is 256 KiB. + pub min_chunk_size: usize, default = 256 * 1024 + + /// Maximum chunk size in bytes. A split is forced when the accumulated + /// size exceeds this value. Default is 1 MiB. + pub max_chunk_size: usize, default = 1024 * 1024 + + /// Normalization level. Increasing this improves deduplication ratio + /// but increases fragmentation. Recommended range is [-3, 3], default is 0. + pub norm_level: i32, default = 0 + } +} + +impl ParquetCdcOptions { + /// Returns enabled CDC options with the default chunking parameters. + /// + /// Shorthand for `ParquetCdcOptions { enabled: true, ..Default::default() }`; + /// combine with struct-update syntax to override parameters, e.g. + /// `ParquetCdcOptions { min_chunk_size: 4096, ..ParquetCdcOptions::enabled() }`. + pub fn enabled() -> Self { Self { - min_chunk_size: 256 * 1024, - max_chunk_size: 1024 * 1024, - norm_level: 0, + enabled: true, + ..Default::default() } } + + /// Returns disabled CDC options (equivalent to [`ParquetCdcOptions::default`]). + pub fn disabled() -> Self { + Self::default() + } } -impl ConfigField for CdcOptions { - fn set(&mut self, key: &str, value: &str) -> Result<()> { - let (key, rem) = key.split_once('.').unwrap_or((key, "")); - match key { - "min_chunk_size" => self.min_chunk_size.set(rem, value), - "max_chunk_size" => self.max_chunk_size.set(rem, value), - "norm_level" => self.norm_level.set(rem, value), - _ => _config_err!("Config value \"{}\" not found on CdcOptions", key), +/// Target maximum size of a Parquet row group in bytes. +/// +/// Wraps a `usize` so the "must be greater than zero" constraint (arrow-rs +/// panics on a zero byte limit) is validated when the config is set, rather +/// than when the writer properties are built. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct MaxRowGroupBytes(usize); + +impl MaxRowGroupBytes { + /// Creates a `MaxRowGroupBytes`, rejecting zero. + pub fn try_new(value: usize) -> Result { + if value == 0 { + return Err(DataFusionError::Configuration( + "max_row_group_bytes must be greater than 0".to_string(), + )); } + Ok(Self(value)) } - fn visit(&self, v: &mut V, key_prefix: &str, _description: &'static str) { - let key = format!("{key_prefix}.min_chunk_size"); - self.min_chunk_size.visit(v, &key, "Minimum chunk size in bytes. The rolling hash will not trigger a split until this many bytes have been accumulated. Default is 256 KiB."); - let key = format!("{key_prefix}.max_chunk_size"); - self.max_chunk_size.visit(v, &key, "Maximum chunk size in bytes. A split is forced when the accumulated size exceeds this value. Default is 1 MiB."); - let key = format!("{key_prefix}.norm_level"); - self.norm_level.visit(v, &key, "Normalization level. Increasing this improves deduplication ratio but increases fragmentation. Recommended range is [-3, 3], default is 0."); + /// Returns the configured byte limit. + pub fn get(&self) -> usize { + self.0 } +} - fn reset(&mut self, key: &str) -> Result<()> { - let (key, rem) = key.split_once('.').unwrap_or((key, "")); - match key { - "min_chunk_size" => { - if rem.is_empty() { - self.min_chunk_size = CdcOptions::default().min_chunk_size; - Ok(()) - } else { - self.min_chunk_size.reset(rem) - } - } - "max_chunk_size" => { - if rem.is_empty() { - self.max_chunk_size = CdcOptions::default().max_chunk_size; - Ok(()) - } else { - self.max_chunk_size.reset(rem) - } - } - "norm_level" => { - if rem.is_empty() { - self.norm_level = CdcOptions::default().norm_level; - Ok(()) - } else { - self.norm_level.reset(rem) - } - } - _ => _config_err!("Config value \"{}\" not found on CdcOptions", key), - } +impl FromStr for MaxRowGroupBytes { + type Err = DataFusionError; + + fn from_str(s: &str) -> Result { + let value = s.parse::().map_err(|_| { + DataFusionError::Configuration(format!( + "Invalid max_row_group_bytes: '{s}'. Expected a positive integer." + )) + })?; + Self::try_new(value) + } +} + +impl Display for MaxRowGroupBytes { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.0) } } -/// `ConfigField` for `Option` — allows setting the option to -/// `"true"` (enable with defaults) or `"false"` (disable), in addition to -/// setting individual sub-fields like `min_chunk_size`. -impl ConfigField for Option { +/// `ConfigField` for `Option`. A custom impl (rather than the +/// blanket `Option` one) so an invalid value is rejected without leaving the +/// option in an invalid intermediate state on error. `MaxRowGroupBytes` +/// deliberately does not implement `Default`, so the blanket impl does not apply. +impl ConfigField for Option { fn visit(&self, v: &mut V, key: &str, description: &'static str) { match self { - Some(s) => s.visit(v, key, description), + Some(s) => v.some(key, s, description), None => v.none(key, description), } } - fn set(&mut self, key: &str, value: &str) -> Result<()> { - if key.is_empty() { - match value.to_ascii_lowercase().as_str() { - "true" => { - *self = Some(CdcOptions::default()); - Ok(()) - } - "false" => { - *self = None; - Ok(()) - } - _ => _config_err!( - "Expected 'true' or 'false' for use_content_defined_chunking, got '{value}'" - ), - } - } else { - self.get_or_insert_with(CdcOptions::default).set(key, value) - } + fn set(&mut self, _key: &str, value: &str) -> Result<()> { + *self = Some(MaxRowGroupBytes::from_str(value)?); + Ok(()) } - fn reset(&mut self, key: &str) -> Result<()> { - if key.is_empty() { - *self = None; - Ok(()) - } else { - self.get_or_insert_with(CdcOptions::default).reset(key) - } + fn reset(&mut self, _key: &str) -> Result<()> { + *self = None; + Ok(()) } } @@ -929,6 +1282,17 @@ config_namespace! { /// parquet reader setting. 0 means no caching. pub max_predicate_cache_size: Option, default = None + /// Maximum number of values in an `IN (...)` list for which pruning will + /// occur. Longer lists will not be used to prune files, row groups, or + /// data pages. + /// + /// Higher values help in cases such as filtering on a list of + /// ~25-100 identifiers, but also make the predicate more expensive to + /// evaluate. Set to 0 to disable `IN (...)` list pruning entirely. + /// + /// Defaults to 20. + pub max_in_list_size: usize, default = 20 + // The following options affect writing to parquet files // and map to parquet::file::properties::WriterProperties @@ -973,9 +1337,21 @@ config_namespace! { /// (writing) Target maximum number of rows in each row group (defaults to 1M /// rows). Writing larger row groups requires more memory to write, but - /// can get better compression and be faster to read. + /// can get better compression and be faster to read. When + /// `max_row_group_bytes` is also set, the writer flushes a row group when + /// either limit is reached, whichever comes first. pub max_row_group_size: usize, default = 1024 * 1024 + /// (writing) Target maximum size of each row group in bytes. When set, + /// the writer flushes whenever either this limit or `max_row_group_size` + /// is reached, whichever comes first. Useful for bounding writer memory + /// on wide schemas where a row-count limit can map to very different + /// byte sizes. Matches the behavior of `parquet.block.size` in + /// parquet-mr. If `None` (the default), only the row-count limit + /// applies. Currently only honored when `allow_single_file_parallelism` + /// is `false`; by default the parallel file writer ignores this limit. + pub max_row_group_bytes: Option, default = None + /// (writing) Sets "created by" property pub created_by: String, default = concat!("datafusion version ", env!("CARGO_PKG_VERSION")).into() @@ -1036,11 +1412,14 @@ config_namespace! { /// data frame. pub maximum_buffered_record_batches_per_stream: usize, default = 2 - /// (writing) EXPERIMENTAL: Enable content-defined chunking (CDC) when writing - /// parquet files. When `Some`, CDC is enabled with the given options; when `None` - /// (the default), CDC is disabled. When CDC is enabled, parallel writing is - /// automatically disabled since the chunker state must persist across row groups. - pub use_content_defined_chunking: Option, default = None + /// (writing) EXPERIMENTAL: Content-defined chunking (CDC) options when writing + /// parquet files. Disabled by default; toggle with + /// `content_defined_chunking.enabled = true|false`. The chunking parameters live + /// under the same prefix (e.g. `content_defined_chunking.min_chunk_size`). When + /// enabled, parallel writing is automatically disabled since the chunker state + /// must persist across row groups. Mirrors + /// `parquet::file::properties::WriterProperties::content_defined_chunking`. + pub content_defined_chunking: ParquetCdcOptions, default = Default::default() } } @@ -1124,6 +1503,22 @@ config_namespace! { /// into the file scan phase. pub enable_topk_dynamic_filter_pushdown: bool, default = true + /// When set to true, uncorrelated scalar subqueries are + /// left in the logical plan and executed by `ScalarSubqueryExec` during + /// physical execution. When set to false, all scalar subqueries + /// (including uncorrelated ones) are rewritten to left joins by the + /// `ScalarSubqueryToJoin` optimizer rule. + /// + /// Note disabling this option is not recommended. It restores + /// pre + /// behavior, which silently produces incorrect results for + /// multi-row subqueries and does not support scalar subqueries in + /// ORDER BY / JOIN ON / aggregate-function arguments. This option is + /// intended as a temporary escape hatch for distributed execution + /// frameworks and is planned to be removed in a future DataFusion + /// release. + pub enable_physical_uncorrelated_scalar_subquery: bool, default = true + /// When set to true, the optimizer will attempt to push down Join dynamic filters /// into the file scan phase. pub enable_join_dynamic_filter_pushdown: bool, default = true @@ -1151,8 +1546,13 @@ config_namespace! { /// in parallel using the provided `target_partitions` level pub repartition_aggregations: bool, default = true - /// Minimum total files size in bytes to perform file scan repartitioning. - pub repartition_file_min_size: usize, default = 10 * 1024 * 1024 + /// Minimum total file size in bytes for file-group byte-range + /// splitting to fire. Files (or merged file groups) smaller than this + /// stay as one partition. Lower values produce more, smaller + /// partitions — better at filling `target_partitions` worth of cores + /// when files are modestly sized, at the cost of slightly more + /// per-partition open / metadata-load overhead. + pub repartition_file_min_size: usize, default = 1024 * 1024 /// Should DataFusion repartition data using the join keys to execute joins in parallel /// using the provided `target_partitions` level @@ -1184,7 +1584,7 @@ config_namespace! { pub repartition_file_scans: bool, default = true /// Minimum number of distinct partition values required to group files by their - /// Hive partition column values (enabling Hash partitioning declaration). + /// Hive partition column values (enabling output partitioning declaration). /// /// How the option is used: /// - preserve_file_partitions=0: Disable it. @@ -1478,6 +1878,24 @@ impl<'a> TryFrom<&'a FormatOptions> for arrow::util::display::FormatOptions<'a> } } +config_namespace! { + /// Options controlling DataFusion's Spark-compatibility layer (functions + /// under `datafusion/spark`). Keys here mirror their `spark.sql.*` + /// equivalents in Apache Spark. + pub struct SparkOptions { + /// Policy for handling duplicate keys in Spark-compatible map-construction + /// functions (`map_from_arrays`, `map_from_entries`, `str_to_map`). + /// + /// Mirrors Spark's + /// [`spark.sql.mapKeyDedupPolicy`](https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961): + /// - `EXCEPTION` (default): raise `[DUPLICATED_MAP_KEY]` at runtime on any duplicate key. + /// - `LAST_WIN`: keep the last occurrence of each duplicate key. + /// + /// Values are case-insensitive. + pub map_key_dedup_policy: MapKeyDedupPolicy, default = MapKeyDedupPolicy::Exception + } +} + /// A key value pair, with a corresponding description #[derive(Debug, Clone, Hash, PartialEq, Eq)] pub struct ConfigEntry { @@ -1509,6 +1927,8 @@ pub struct ConfigOptions { pub extensions: Extensions, /// Formatting options when printing batches pub format: FormatOptions, + /// Spark-compatibility options (functions under `datafusion/spark`) + pub spark: SparkOptions, } impl ConfigField for ConfigOptions { @@ -1519,6 +1939,7 @@ impl ConfigField for ConfigOptions { self.explain.visit(v, "datafusion.explain", ""); self.sql_parser.visit(v, "datafusion.sql_parser", ""); self.format.visit(v, "datafusion.format", ""); + self.spark.visit(v, "datafusion.spark", ""); } fn set(&mut self, key: &str, value: &str) -> Result<()> { @@ -1531,6 +1952,7 @@ impl ConfigField for ConfigOptions { "explain" => self.explain.set(rem, value), "sql_parser" => self.sql_parser.set(rem, value), "format" => self.format.set(rem, value), + "spark" => self.spark.set(rem, value), _ => _config_err!("Config value \"{key}\" not found on ConfigOptions"), } } @@ -1570,6 +1992,7 @@ impl ConfigField for ConfigOptions { "explain" => self.explain.reset(rem), "sql_parser" => self.sql_parser.reset(rem), "format" => self.format.reset(rem), + "spark" => self.spark.reset(rem), other => _config_err!("Config value \"{other}\" not found on ConfigOptions"), } } @@ -1613,7 +2036,8 @@ impl ConfigOptions { } return Ok(()); } - return ConfigField::set(self, inner_key, value); + return ConfigField::set(self, inner_key, value) + .map_err(|e| e.context(format!("Error setting config {key}"))); } if !self.extensions.0.contains_key(prefix) @@ -3452,6 +3876,7 @@ impl Display for OutputFormat { #[cfg(test)] mod tests { #[cfg(feature = "parquet")] + use crate::assert_contains; use crate::config::TableParquetOptions; use crate::config::{ ConfigEntry, ConfigExtension, ConfigField, ConfigFileType, ExtensionOptions, @@ -4012,7 +4437,7 @@ mod tests { let err = config .set("datafusion.execution.parquet.writer_version", "3.0") .unwrap_err(); - assert_eq!( + assert_contains!( err.to_string(), "Invalid or Unsupported Configuration: Invalid parquet writer version: 3.0. Expected one of: 1.0, 2.0" ); @@ -4020,74 +4445,179 @@ mod tests { #[cfg(feature = "parquet")] #[test] - fn set_cdc_option_with_boolean_true() { + fn set_cdc_enabled_flag() { use crate::config::ConfigOptions; let mut config = ConfigOptions::default(); - assert!( - config - .execution - .parquet - .use_content_defined_chunking - .is_none() - ); + // CDC is disabled by default. + assert!(!config.execution.parquet.content_defined_chunking.enabled); - // Setting to "true" should enable CDC with default options + // `.enabled = true` enables CDC; parameters keep their defaults. config .set( - "datafusion.execution.parquet.use_content_defined_chunking", + "datafusion.execution.parquet.content_defined_chunking.enabled", "true", ) .unwrap(); - let cdc = config - .execution - .parquet - .use_content_defined_chunking - .as_ref() - .expect("CDC should be enabled"); + let cdc = &config.execution.parquet.content_defined_chunking; + assert!(cdc.enabled); assert_eq!(cdc.min_chunk_size, 256 * 1024); assert_eq!(cdc.max_chunk_size, 1024 * 1024); assert_eq!(cdc.norm_level, 0); - // Setting to "false" should disable CDC + // `.enabled = false` disables CDC. config .set( - "datafusion.execution.parquet.use_content_defined_chunking", + "datafusion.execution.parquet.content_defined_chunking.enabled", "false", ) .unwrap(); - assert!( - config - .execution - .parquet - .use_content_defined_chunking - .is_none() - ); + assert!(!config.execution.parquet.content_defined_chunking.enabled); } #[cfg(feature = "parquet")] #[test] - fn set_cdc_option_with_subfields() { + fn set_cdc_param_does_not_enable() { use crate::config::ConfigOptions; let mut config = ConfigOptions::default(); - // Setting sub-fields should also enable CDC + // Setting a parameter does NOT enable CDC (`enabled` is a distinct field, + // defaulting to false), and the result is independent of key order. config .set( - "datafusion.execution.parquet.use_content_defined_chunking.min_chunk_size", + "datafusion.execution.parquet.content_defined_chunking.min_chunk_size", "1024", ) .unwrap(); - let cdc = config - .execution - .parquet - .use_content_defined_chunking - .as_ref() - .expect("CDC should be enabled"); + let cdc = &config.execution.parquet.content_defined_chunking; + assert!(!cdc.enabled); assert_eq!(cdc.min_chunk_size, 1024); - // Other fields should be defaults assert_eq!(cdc.max_chunk_size, 1024 * 1024); assert_eq!(cdc.norm_level, 0); } + + #[test] + fn test_dialect_metadata_roundtrip() { + use crate::config::Dialect; + use std::str::FromStr; + + assert_eq!(Dialect::default(), Dialect::Generic); + assert!(!Dialect::metadata().is_empty()); + + for info in Dialect::metadata() { + let dialect = info.dialect; + + assert_eq!(Dialect::from_str(info.canonical_name).unwrap(), dialect); + assert_eq!( + Dialect::from_str(&info.canonical_name.to_ascii_uppercase()).unwrap(), + dialect + ); + assert_eq!(dialect.as_ref(), info.canonical_name); + assert_eq!(dialect.to_string(), info.canonical_name); + } + } + + #[test] + fn test_dialect_aliases() { + use crate::config::Dialect; + use std::str::FromStr; + + for info in Dialect::metadata() { + for alias in info.aliases { + assert_eq!(Dialect::from_str(alias).unwrap(), info.dialect); + assert_eq!( + Dialect::from_str(&alias.to_ascii_uppercase()).unwrap(), + info.dialect + ); + } + } + } + + #[test] + fn test_available_dialects_includes_each_display_name_once() { + use crate::config::Dialect; + use std::collections::BTreeSet; + + let available = Dialect::available(); + let listed: Vec<_> = available.split(", ").collect(); + let display_names: Vec<_> = Dialect::metadata() + .iter() + .map(|info| info.display_name) + .collect(); + let unique_display_names: BTreeSet<_> = display_names.iter().copied().collect(); + + assert_eq!(display_names.len(), unique_display_names.len()); + assert_eq!(listed, display_names); + } + + #[test] + fn test_dialect_config_description_uses_metadata() { + use crate::config::{ConfigOptions, Dialect, SQL_PARSER_DIALECT_CONFIG_KEY}; + + let description = ConfigOptions::default() + .entries() + .into_iter() + .find(|entry| entry.key == SQL_PARSER_DIALECT_CONFIG_KEY) + .unwrap() + .description; + + assert!(description.contains(Dialect::available())); + } + + #[test] + fn test_invalid_dialect_error_lists_available_dialects() { + use crate::config::Dialect; + use std::str::FromStr; + + let error = Dialect::from_str("notadialect").unwrap_err().to_string(); + + assert!(error.contains("Invalid Dialect: notadialect")); + assert!(error.contains(Dialect::available())); + } + + #[test] + fn max_row_group_bytes_rejects_zero() { + use crate::config::MaxRowGroupBytes; + use std::str::FromStr; + + assert!(MaxRowGroupBytes::try_new(0).is_err()); + assert!(MaxRowGroupBytes::from_str("0").is_err()); + assert!(MaxRowGroupBytes::from_str("not_a_number").is_err()); + assert_eq!(MaxRowGroupBytes::try_new(128).unwrap().get(), 128); + assert_eq!(MaxRowGroupBytes::from_str("128").unwrap().get(), 128); + } + + #[test] + fn parquet_max_row_group_bytes_config_set_rejects_zero() { + use crate::config::ConfigOptions; + + let mut options = ConfigOptions::new(); + options + .set("datafusion.execution.parquet.max_row_group_bytes", "1024") + .unwrap(); + assert_eq!( + options + .execution + .parquet + .max_row_group_bytes + .map(|v| v.get()), + Some(1024) + ); + + // Zero is rejected at set time, leaving the previous value unchanged. + assert!( + options + .set("datafusion.execution.parquet.max_row_group_bytes", "0") + .is_err() + ); + assert_eq!( + options + .execution + .parquet + .max_row_group_bytes + .map(|v| v.get()), + Some(1024) + ); + } } diff --git a/datafusion/common/src/dfschema.rs b/datafusion/common/src/dfschema.rs index e3da99163ed69..262f1dcf619d9 100644 --- a/datafusion/common/src/dfschema.rs +++ b/datafusion/common/src/dfschema.rs @@ -599,30 +599,6 @@ impl DFSchema { .all(|(dffield, arrowfield)| dffield.name() == arrowfield.name()) } - /// Check to see if fields in 2 Arrow schemas are compatible - #[deprecated(since = "47.0.0", note = "This method is no longer used")] - pub fn check_arrow_schema_type_compatible( - &self, - arrow_schema: &Schema, - ) -> Result<()> { - let self_arrow_schema = self.as_arrow(); - self_arrow_schema - .fields() - .iter() - .zip(arrow_schema.fields().iter()) - .try_for_each(|(l_field, r_field)| { - if !can_cast_types(r_field.data_type(), l_field.data_type()) { - _plan_err!("Column {} (type: {}) is not compatible with column {} (type: {})", - r_field.name(), - r_field.data_type(), - l_field.name(), - l_field.data_type()) - } else { - Ok(()) - } - }) - } - /// Returns true if the two schemas have the same qualified named /// fields with logically equivalent data types. Returns false otherwise. /// @@ -641,11 +617,6 @@ impl DFSchema { }) } - #[deprecated(since = "47.0.0", note = "Use has_equivalent_names_and_types` instead")] - pub fn equivalent_names_and_types(&self, other: &Self) -> bool { - self.has_equivalent_names_and_types(other).is_ok() - } - /// Returns Ok if the two schemas have the same qualified named /// fields with the compatible data types. /// @@ -700,6 +671,8 @@ impl DFSchema { /// logically equivalent. For example: /// - a Dictionary type is logically equal to a plain V type /// - a Dictionary is also logically equal to Dictionary + /// - a RunEndEncoded type is logically equal to a plain V type + /// - a RunEndEncoded is also logically equal to RunEndEncoded /// - Utf8 and Utf8View are logically equal pub fn datatype_is_logically_equal(dt1: &DataType, dt2: &DataType) -> bool { // check nested fields @@ -711,8 +684,17 @@ impl DFSchema { | (othertype, DataType::Dictionary(_, v1)) => { Self::datatype_is_logically_equal(v1.as_ref(), othertype) } + (DataType::RunEndEncoded(_, v1), DataType::RunEndEncoded(_, v2)) => { + Self::datatype_is_logically_equal(v1.data_type(), v2.data_type()) + } + (DataType::RunEndEncoded(_, v1), othertype) + | (othertype, DataType::RunEndEncoded(_, v1)) => { + Self::datatype_is_logically_equal(v1.data_type(), othertype) + } (DataType::List(f1), DataType::List(f2)) | (DataType::LargeList(f1), DataType::LargeList(f2)) + | (DataType::ListView(f1), DataType::ListView(f2)) + | (DataType::LargeListView(f1), DataType::LargeListView(f2)) | (DataType::FixedSizeList(f1, _), DataType::FixedSizeList(f2, _)) => { // Don't compare the names of the technical inner field // Usually "item" but that's not mandated @@ -771,8 +753,17 @@ impl DFSchema { Self::datatype_is_semantically_equal(k1.as_ref(), k2.as_ref()) && Self::datatype_is_semantically_equal(v1.as_ref(), v2.as_ref()) } + (DataType::RunEndEncoded(k1, v1), DataType::RunEndEncoded(k2, v2)) => { + Self::datatype_is_semantically_equal(k1.data_type(), k2.data_type()) + && Self::datatype_is_semantically_equal( + v1.data_type(), + v2.data_type(), + ) + } (DataType::List(f1), DataType::List(f2)) | (DataType::LargeList(f1), DataType::LargeList(f2)) + | (DataType::ListView(f1), DataType::ListView(f2)) + | (DataType::LargeListView(f1), DataType::LargeListView(f2)) | (DataType::FixedSizeList(f1, _), DataType::FixedSizeList(f2, _)) => { // Don't compare the names of the technical inner field // Usually "item" but that's not mandated @@ -1281,13 +1272,13 @@ pub trait SchemaExt { /// This is a specialized version of Eq that ignores differences /// in nullability and metadata. /// - /// It works the same as [`DFSchema::equivalent_names_and_types`]. + /// It works the same as [`DFSchema::has_equivalent_names_and_types`]. fn equivalent_names_and_types(&self, other: &Self) -> bool; /// Returns nothing if the two schemas have the same qualified named /// fields with logically equivalent data types. Returns internal error otherwise. /// - /// Use [DFSchema]::equivalent_names_and_types for stricter semantic type + /// Use [DFSchema]::has_equivalent_names_and_types for stricter semantic type /// equivalence checking. /// /// It is only used by insert into cases. @@ -1427,11 +1418,8 @@ mod tests { let schema = DFSchema::try_from_qualified_schema("t1", &test_schema_1())?; // lookup with unqualified name "t1.c0" let err = schema.index_of_column(&col).unwrap_err(); - let expected = "Schema error: No field named \"t1.c0\". \ - Column names are case sensitive. \ - You can use double quotes to refer to the \"\"t1.c0\"\" column \ - or set the datafusion.sql_parser.enable_ident_normalization configuration. \ - Did you mean 't1.c0'?."; + let expected = "Schema error: No field named \"t1.c0\". Did you mean 't1.c0'?\n\ + Valid fields are t1.c0, t1.c1."; assert_eq!(err.strip_backtrace(), expected); Ok(()) } @@ -1449,12 +1437,47 @@ mod tests { // lookup with unqualified name "t1.c0" let err = schema.index_of_column(&col).unwrap_err(); - let expected = "Schema error: No field named \"t1.c0\". \ + let expected = "Schema error: No field named \"t1.c0\".\n\ Valid fields are t1.\"CapitalColumn\", t1.\"field.with.period\"."; assert_eq!(err.strip_backtrace(), expected); Ok(()) } + #[test] + fn field_not_found_suggests_closest_field_name() -> Result<()> { + let schema = DFSchema::try_from(Schema::new(vec![ + Field::new("abzz", DataType::Boolean, true), + Field::new("abcd", DataType::Boolean, true), + ]))?; + + let err = schema.field_with_unqualified_name("abc").unwrap_err(); + let expected = "Schema error: No field named abc. Did you mean 'abcd'?\n\ + Valid fields are abzz, abcd."; + assert_eq!(err.strip_backtrace(), expected); + Ok(()) + } + + #[test] + fn field_not_found_suggests_case_sensitive_qualified_field() -> Result<()> { + let schema = DFSchema::try_from_qualified_schema( + "hits", + &Schema::new(vec![ + Field::new("WatchID", DataType::Boolean, true), + Field::new("URL", DataType::Boolean, true), + Field::new("URLHash", DataType::Boolean, true), + ]), + )?; + + let err = schema.field_with_unqualified_name("url").unwrap_err(); + let expected = "Schema error: No field named url. Did you mean 'hits.\"URL\"'?\n\ + Column names are case sensitive. \ + You can use double quotes to refer to the hits.\"URL\" column \ + or disable the datafusion.sql_parser.enable_ident_normalization configuration.\n\ + Valid fields are hits.\"WatchID\", hits.\"URL\", hits.\"URLHash\"."; + assert_eq!(err.strip_backtrace(), expected); + Ok(()) + } + #[test] fn from_unqualified_schema() -> Result<()> { let schema = DFSchema::try_from(test_schema_1())?; @@ -1751,12 +1774,20 @@ mod tests { &DataType::List(Field::new_list_field(DataType::Int8, true).into()), &DataType::List(Field::new("element", DataType::Int8, false).into()) )); + assert!(DFSchema::datatype_is_logically_equal( + &DataType::ListView(Field::new_list_field(DataType::Int8, true).into()), + &DataType::ListView(Field::new("element", DataType::Int8, false).into()) + )); // Fails if element type is different assert!(!DFSchema::datatype_is_logically_equal( &DataType::List(Field::new_list_field(DataType::Int8, true).into()), &DataType::List(Field::new_list_field(DataType::Int16, true).into()) )); + assert!(!DFSchema::datatype_is_logically_equal( + &DataType::ListView(Field::new_list_field(DataType::Int8, true).into()), + &DataType::ListView(Field::new_list_field(DataType::Int16, true).into()) + )); // Test maps let map_field = DataType::Map( @@ -1893,6 +1924,50 @@ mod tests { )); } + #[test] + fn test_datatype_is_logically_equivalent_to_ree() { + // RunEndEncoded is logically equal to its value type + assert!(DFSchema::datatype_is_logically_equal( + &DataType::Utf8, + &DataType::RunEndEncoded( + Field::new("run", DataType::Int32, false).into(), + Field::new("val", DataType::Utf8, true).into(), + ) + )); + + // Dictionary is logically equal to the logically equivalent value type + assert!(DFSchema::datatype_is_logically_equal( + &DataType::Utf8View, + &DataType::RunEndEncoded( + Field::new("run", DataType::Int32, false).into(), + Field::new("val", DataType::Utf8, true).into(), + ) + )); + + assert!(DFSchema::datatype_is_logically_equal( + &DataType::RunEndEncoded( + Field::new("run", DataType::Int32, false).into(), + Field::new( + "val", + DataType::List(Field::new("element", DataType::Utf8, false).into()), + true + ) + .into(), + ), + &DataType::RunEndEncoded( + Field::new("run", DataType::Int64, false).into(), + Field::new( + "val", + DataType::List( + Field::new("element", DataType::Utf8View, false).into() + ), + true + ) + .into(), + ), + )); + } + #[test] fn test_datatype_is_semantically_equal() { assert!(DFSchema::datatype_is_semantically_equal( @@ -1942,12 +2017,20 @@ mod tests { &DataType::List(Field::new_list_field(DataType::Int8, true).into()), &DataType::List(Field::new("element", DataType::Int8, false).into()) )); + assert!(DFSchema::datatype_is_semantically_equal( + &DataType::ListView(Field::new_list_field(DataType::Int8, true).into()), + &DataType::ListView(Field::new("element", DataType::Int8, false).into()) + )); // Fails if element type is different assert!(!DFSchema::datatype_is_semantically_equal( &DataType::List(Field::new_list_field(DataType::Int8, true).into()), &DataType::List(Field::new_list_field(DataType::Int16, true).into()) )); + assert!(!DFSchema::datatype_is_semantically_equal( + &DataType::ListView(Field::new_list_field(DataType::Int8, true).into()), + &DataType::ListView(Field::new_list_field(DataType::Int16, true).into()) + )); // Test maps let map_field = DataType::Map( @@ -2063,6 +2146,18 @@ mod tests { )); } + #[test] + fn test_datatype_is_not_semantically_equivalent_to_ree() { + // RunEndEncoded is not semantically equal to its value type + assert!(!DFSchema::datatype_is_semantically_equal( + &DataType::Utf8, + &DataType::RunEndEncoded( + Field::new("run", DataType::Int32, false).into(), + Field::new("val", DataType::Utf8, true).into(), + ) + )); + } + fn test_schema_2() -> Schema { Schema::new(vec![ Field::new("c100", DataType::Boolean, true), diff --git a/datafusion/common/src/error.rs b/datafusion/common/src/error.rs index c6c50371c26c1..02016387c0a96 100644 --- a/datafusion/common/src/error.rs +++ b/datafusion/common/src/error.rs @@ -45,7 +45,7 @@ use std::io; use std::result; use std::sync::Arc; -use crate::utils::datafusion_strsim::normalized_levenshtein; +use crate::utils::datafusion_strsim::{levenshtein, normalized_levenshtein}; use crate::utils::quote_identifier; use crate::{Column, DFSchema, Diagnostic, TableReference}; use arrow::error::ArrowError; @@ -198,6 +198,77 @@ pub enum SchemaError { }, } +fn case_insensitive_field_match<'a>( + field: &Column, + valid_fields: &'a [Column], +) -> Option<&'a Column> { + let field_name = field.name(); + let field_flat_name = field.flat_name(); + let field_name_lower = field_name.to_lowercase(); + let field_flat_name_lower = field_flat_name.to_lowercase(); + + valid_fields.iter().find(|valid_field| { + let valid_field_name = valid_field.name(); + let valid_field_flat_name = valid_field.flat_name(); + let valid_field_name_lower = valid_field_name.to_lowercase(); + let valid_field_flat_name_lower = valid_field_flat_name.to_lowercase(); + + let name_differs_only_by_case = + field_name_lower == valid_field_name_lower && field_name != valid_field_name; + let flat_name_differs_only_by_case = field_flat_name_lower + == valid_field_flat_name_lower + && field_flat_name != valid_field_flat_name; + + name_differs_only_by_case || flat_name_differs_only_by_case + }) +} + +/// Find the most similar field name based on edit distance. +/// Returns `None` if all candidate edit distances are too far away. +fn closest_valid_field<'a>( + field: &Column, + valid_fields: &'a [Column], +) -> Option<&'a Column> { + // Find the most similar valid field name. + let target_names = [ + field.name().to_lowercase(), + field.flat_name().to_lowercase(), + ]; + + let mut best_match: Option<(usize, usize, usize, &Column)> = None; + for (index, valid_field) in valid_fields.iter().enumerate() { + let valid_names = [ + valid_field.name().to_lowercase(), + valid_field.flat_name().to_lowercase(), + ]; + for target in &target_names { + for valid_name in &valid_names { + let distance = levenshtein(target, valid_name); + let max_len = target.chars().count().max(valid_name.chars().count()); + // If there are no shared characters, or we would have to edit + // more than half of the longer name, don't suggest a potential match. + if max_len == 0 || distance * 2 > max_len { + continue; + } + + let should_replace = best_match.is_none_or( + |(best_distance, best_max_len, best_index, _)| { + distance < best_distance + || distance == best_distance + && (max_len > best_max_len + || max_len == best_max_len && index < best_index) + }, + ); + if should_replace { + best_match = Some((distance, max_len, index, valid_field)); + } + } + } + } + + best_match.map(|(_, _, _, valid_field)| valid_field) +} + impl Display for SchemaError { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { @@ -205,44 +276,39 @@ impl Display for SchemaError { field, valid_fields, } => { + let closest_field = closest_valid_field(field, valid_fields); + let case_sensitive_match = + case_insensitive_field_match(field, valid_fields); + write!(f, "No field named {}", field.quoted_flat_name())?; - let lower_valid_fields = valid_fields - .iter() - .map(|column| column.flat_name().to_lowercase()) - .collect::>(); - - let valid_fields_names = valid_fields - .iter() - .map(|column| column.flat_name()) - .collect::>(); - if lower_valid_fields.contains(&field.flat_name().to_lowercase()) { + if let Some(matched) = closest_field { + write!(f, ". Did you mean '{}'?", matched.quoted_flat_name())?; + } else { + write!(f, ".")?; + } + + if let Some(case_sensitive_match) = case_sensitive_match { write!( f, - ". Column names are case sensitive. You can use double quotes to refer to the \"{}\" column \ - or set the datafusion.sql_parser.enable_ident_normalization configuration", - field.quoted_flat_name() + "\nColumn names are case sensitive. You can use double quotes to refer to the {} column \ + or disable the datafusion.sql_parser.enable_ident_normalization configuration.", + case_sensitive_match.quoted_flat_name() )?; } - let field_name = field.name(); - if let Some(matched) = valid_fields_names - .iter() - .filter(|str| normalized_levenshtein(str, field_name) >= 0.5) - .collect::>() - .first() - { - write!(f, ". Did you mean '{matched}'?")?; - } else if !valid_fields.is_empty() { + + if !valid_fields.is_empty() { write!( f, - ". Valid fields are {}", + "\nValid fields are {}.", valid_fields .iter() .map(|field| field.quoted_flat_name()) .collect::>() .join(", ") - )?; + ) + } else { + Ok(()) } - write!(f, ".") } Self::DuplicateQualifiedField { qualifier, name } => { write!( @@ -621,14 +687,11 @@ impl DataFusionError { return Some(diagnostics); } - if let Some(source) = self - .head - .source() - .and_then(|source| source.downcast_ref::()) { + let source = self.head.source().and_then(|source| { + source.downcast_ref::() + })?; self.head = source; - } else { - return None; } } } @@ -757,10 +820,10 @@ impl DataFusionErrorBuilder { macro_rules! unwrap_or_internal_err { ($Value: ident) => { $Value.ok_or_else(|| { - $crate::DataFusionError::Internal(format!( + $crate::error::_internal_datafusion_err!( "{} should not be None", stringify!($Value) - )) + ) })? }; } @@ -778,19 +841,19 @@ macro_rules! unwrap_or_internal_err { macro_rules! assert_or_internal_err { ($cond:expr) => { if !$cond { - return Err($crate::DataFusionError::Internal(format!( + return Err($crate::error::_internal_datafusion_err!( "Assertion failed: {}", stringify!($cond) - ))); + )); } }; ($cond:expr, $($arg:tt)+) => { if !$cond { - return Err($crate::DataFusionError::Internal(format!( + return Err($crate::error::_internal_datafusion_err!( "Assertion failed: {}: {}", stringify!($cond), format!($($arg)+) - ))); + )); } }; } @@ -810,27 +873,27 @@ macro_rules! assert_eq_or_internal_err { let left_val = &$left; let right_val = &$right; if left_val != right_val { - return Err($crate::DataFusionError::Internal(format!( + return Err($crate::error::_internal_datafusion_err!( "Assertion failed: {} == {} (left: {:?}, right: {:?})", stringify!($left), stringify!($right), left_val, right_val - ))); + )); } }}; ($left:expr, $right:expr, $($arg:tt)+) => {{ let left_val = &$left; let right_val = &$right; if left_val != right_val { - return Err($crate::DataFusionError::Internal(format!( + return Err($crate::error::_internal_datafusion_err!( "Assertion failed: {} == {} (left: {:?}, right: {:?}): {}", stringify!($left), stringify!($right), left_val, right_val, format!($($arg)+) - ))); + )); } }}; } @@ -850,27 +913,27 @@ macro_rules! assert_ne_or_internal_err { let left_val = &$left; let right_val = &$right; if left_val == right_val { - return Err($crate::DataFusionError::Internal(format!( + return Err($crate::error::_internal_datafusion_err!( "Assertion failed: {} != {} (left: {:?}, right: {:?})", stringify!($left), stringify!($right), left_val, right_val - ))); + )); } }}; ($left:expr, $right:expr, $($arg:tt)+) => {{ let left_val = &$left; let right_val = &$right; if left_val == right_val { - return Err($crate::DataFusionError::Internal(format!( + return Err($crate::error::_internal_datafusion_err!( "Assertion failed: {} != {} (left: {:?}, right: {:?}): {}", stringify!($left), stringify!($right), left_val, right_val, format!($($arg)+) - ))); + )); } }}; } @@ -1138,7 +1201,6 @@ mod test { use std::sync::Arc; use arrow::error::ArrowError; - use insta::assert_snapshot; fn ok_result() -> Result<()> { Ok(()) @@ -1157,14 +1219,8 @@ mod test { ok_result() } - let err = check().unwrap_err(); - assert_snapshot!( - err.to_string(), - @r" - Internal error: Assertion failed: 1 == 2 (left: 1, right: 2): expected equality. - This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/datafusion/issues - " - ); + let err = check().unwrap_err().strip_backtrace(); + assert!(err.starts_with("Internal error: Assertion failed: 1 == 2 (left: 1, right: 2): expected equality")); } #[test] @@ -1180,14 +1236,8 @@ mod test { ok_result() } - let err = check().unwrap_err(); - assert_snapshot!( - err.to_string(), - @r" - Internal error: Assertion failed: 3 != 3 (left: 3, right: 3): values must differ. - This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/datafusion/issues - " - ); + let err = check().unwrap_err().strip_backtrace(); + assert!(err.starts_with("Internal error: Assertion failed: 3 != 3 (left: 3, right: 3): values must differ")); } #[test] @@ -1204,14 +1254,8 @@ mod test { ok_result() } - let err = check().unwrap_err(); - assert_snapshot!( - err.to_string(), - @r" - Internal error: Assertion failed: false. - This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/datafusion/issues - " - ); + let err = check().unwrap_err().strip_backtrace(); + assert!(err.starts_with("Internal error: Assertion failed: false")); } #[test] @@ -1221,13 +1265,9 @@ mod test { ok_result() } - let err = check().unwrap_err(); - assert_snapshot!( - err.to_string(), - @r" - Internal error: Assertion failed: false: custom message. - This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/datafusion/issues - " + let err = check().unwrap_err().strip_backtrace(); + assert!( + err.starts_with("Internal error: Assertion failed: false: custom message") ); } @@ -1238,14 +1278,8 @@ mod test { ok_result() } - let err = check().unwrap_err(); - assert_snapshot!( - err.to_string(), - @r" - Internal error: Assertion failed: false: custom 42. - This issue was likely caused by a bug in DataFusion's code. Please help us to resolve this by filing a bug report in our issue tracker: https://github.com/apache/datafusion/issues - " - ); + let err = check().unwrap_err().strip_backtrace(); + assert!(err.starts_with("Internal error: Assertion failed: false: custom 42")); } #[test] @@ -1273,23 +1307,69 @@ mod test { // To pass the test the environment variable RUST_BACKTRACE should be set to 1 to enforce backtrace #[cfg(feature = "backtrace")] - #[test] - fn test_enabled_backtrace() { + fn ensure_rust_backtrace_enabled() { match std::env::var("RUST_BACKTRACE") { Ok(val) if val == "1" => {} _ => panic!("Environment variable RUST_BACKTRACE must be set to 1"), }; + } + + // To pass the test the environment variable RUST_BACKTRACE should be set to 1 to enforce backtrace + #[cfg(feature = "backtrace")] + #[test] + fn test_enabled_backtrace() { + ensure_rust_backtrace_enabled(); let res: Result<(), DataFusionError> = plan_err!("Err"); - let err = res.unwrap_err().to_string(); - assert!(err.contains(DataFusionError::BACK_TRACE_SEP)); - assert_eq!( - err.split(DataFusionError::BACK_TRACE_SEP) - .collect::>() - .first() - .unwrap(), - &"Error during planning: Err" + assert_error_have_message_and_backtrace( + &res.unwrap_err(), + "Error during planning: Err", + ); + } + + #[cfg(not(feature = "backtrace"))] + #[test] + fn test_disabled_backtrace() { + let res: Result<(), DataFusionError> = plan_err!("Err"); + assert_err_without_backtrace_and_equal( + &res.unwrap_err(), + "Error during planning: Err", + ); + } + + #[cfg(not(feature = "backtrace"))] + fn assert_err_without_backtrace_and_equal( + err: &DataFusionError, + expected_message: &str, + ) { + let err = err.to_string(); + assert!(!err.contains(DataFusionError::BACK_TRACE_SEP)); + assert_eq!(err, expected_message); + } + + #[cfg(not(feature = "backtrace"))] + fn assert_internal_err_without_backtrace_and_equal( + err: &DataFusionError, + expected_message: &str, + ) { + let expected_message_before_backtrace = format!( + "{expected_message}.\nThis issue was likely caused by a bug in DataFusion's code. \ + Please help us to resolve this by filing a bug report in our issue tracker: \ + https://github.com/apache/datafusion/issues" + ); + assert_err_without_backtrace_and_equal( + err, + expected_message_before_backtrace.as_str(), ); + } + + #[cfg(feature = "backtrace")] + fn assert_error_have_message_and_backtrace( + err: &DataFusionError, + message_before_backtrace: &str, + ) { + let err = err.to_string(); + assert!(err.contains(DataFusionError::BACK_TRACE_SEP)); assert!( !err.split(DataFusionError::BACK_TRACE_SEP) .collect::>() @@ -1297,15 +1377,272 @@ mod test { .unwrap() .is_empty() ); + assert_eq!( + err.split(DataFusionError::BACK_TRACE_SEP) + .collect::>() + .first() + .copied() + .unwrap(), + message_before_backtrace, + "full error is: {err}" + ); } + #[cfg(feature = "backtrace")] + #[test] + fn test_enabled_backtrace_for_unwrap_or_internal_err() { + ensure_rust_backtrace_enabled(); + + fn get_error() -> Result<(), DataFusionError> { + let item = None::<()>; + unwrap_or_internal_err!(item); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_error_have_message_and_backtrace( + &res.unwrap_err(), + "Internal error: item should not be None", + ); + } + + // To pass the test the environment variable RUST_BACKTRACE should be set to 1 to enforce backtrace #[cfg(not(feature = "backtrace"))] #[test] - fn test_disabled_backtrace() { - let res: Result<(), DataFusionError> = plan_err!("Err"); - let res = res.unwrap_err().to_string(); - assert!(!res.contains(DataFusionError::BACK_TRACE_SEP)); - assert_eq!(res, "Error during planning: Err"); + fn test_disabled_backtrace_for_unwrap_or_internal_err() { + fn get_error() -> Result<(), DataFusionError> { + let item = None::<()>; + unwrap_or_internal_err!(item); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_internal_err_without_backtrace_and_equal( + &res.unwrap_err(), + "Internal error: item should not be None", + ); + } + + #[cfg(feature = "backtrace")] + #[test] + fn test_enabled_backtrace_for_assert_or_internal_err_without_args() { + ensure_rust_backtrace_enabled(); + + fn get_error() -> Result<(), DataFusionError> { + assert_or_internal_err!(false); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_error_have_message_and_backtrace( + &res.unwrap_err(), + "Internal error: Assertion failed: false", + ); + } + + #[cfg(feature = "backtrace")] + #[test] + fn test_enabled_backtrace_for_assert_or_internal_err_with_args() { + ensure_rust_backtrace_enabled(); + + fn get_error() -> Result<(), DataFusionError> { + assert_or_internal_err!(false, "my cool context"); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_error_have_message_and_backtrace( + &res.unwrap_err(), + "Internal error: Assertion failed: false: my cool context", + ); + } + + #[cfg(not(feature = "backtrace"))] + #[test] + fn test_disabled_backtrace_for_assert_or_internal_err_without_args() { + fn get_error() -> Result<(), DataFusionError> { + assert_or_internal_err!(false); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_internal_err_without_backtrace_and_equal( + &res.unwrap_err(), + "Internal error: Assertion failed: false", + ); + } + + #[cfg(not(feature = "backtrace"))] + #[test] + fn test_disabled_backtrace_for_assert_or_internal_err_with_args() { + fn get_error() -> Result<(), DataFusionError> { + assert_or_internal_err!(false, "my cool context"); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_internal_err_without_backtrace_and_equal( + &res.unwrap_err(), + "Internal error: Assertion failed: false: my cool context", + ); + } + + #[cfg(feature = "backtrace")] + #[test] + fn test_enabled_backtrace_for_assert_eq_or_internal_err_without_args() { + ensure_rust_backtrace_enabled(); + + fn get_error() -> Result<(), DataFusionError> { + let arg1 = 1; + let arg2 = 2; + assert_eq_or_internal_err!(arg1, arg2); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_error_have_message_and_backtrace( + &res.unwrap_err(), + "Internal error: Assertion failed: arg1 == arg2 (left: 1, right: 2)", + ); + } + + #[cfg(feature = "backtrace")] + #[test] + fn test_enabled_backtrace_for_assert_eq_or_internal_err_with_args() { + ensure_rust_backtrace_enabled(); + + fn get_error() -> Result<(), DataFusionError> { + let arg1 = 1; + let arg2 = 2; + assert_eq_or_internal_err!(arg1, arg2, "my cool context"); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_error_have_message_and_backtrace( + &res.unwrap_err(), + "Internal error: Assertion failed: arg1 == arg2 (left: 1, right: 2): my cool context", + ); + } + + #[cfg(not(feature = "backtrace"))] + #[test] + fn test_disabled_backtrace_for_assert_eq_or_internal_err_without_args() { + fn get_error() -> Result<(), DataFusionError> { + let arg1 = 1; + let arg2 = 2; + assert_eq_or_internal_err!(arg1, arg2); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_internal_err_without_backtrace_and_equal( + &res.unwrap_err(), + "Internal error: Assertion failed: arg1 == arg2 (left: 1, right: 2)", + ); + } + + #[cfg(not(feature = "backtrace"))] + #[test] + fn test_disabled_backtrace_for_assert_eq_or_internal_err_with_args() { + fn get_error() -> Result<(), DataFusionError> { + let arg1 = 1; + let arg2 = 2; + assert_eq_or_internal_err!(arg1, arg2, "my cool context"); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_internal_err_without_backtrace_and_equal( + &res.unwrap_err(), + "Internal error: Assertion failed: arg1 == arg2 (left: 1, right: 2): my cool context", + ); + } + + #[cfg(feature = "backtrace")] + #[test] + fn test_enabled_backtrace_for_assert_ne_or_internal_err_without_args() { + ensure_rust_backtrace_enabled(); + + fn get_error() -> Result<(), DataFusionError> { + let arg1 = 1; + let arg2 = 1; + assert_ne_or_internal_err!(arg1, arg2); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_error_have_message_and_backtrace( + &res.unwrap_err(), + "Internal error: Assertion failed: arg1 != arg2 (left: 1, right: 1)", + ); + } + + #[cfg(feature = "backtrace")] + #[test] + fn test_enabled_backtrace_for_assert_ne_or_internal_err_with_args() { + ensure_rust_backtrace_enabled(); + + fn get_error() -> Result<(), DataFusionError> { + let arg1 = 1; + let arg2 = 1; + assert_ne_or_internal_err!(arg1, arg2, "my cool context"); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_error_have_message_and_backtrace( + &res.unwrap_err(), + "Internal error: Assertion failed: arg1 != arg2 (left: 1, right: 1): my cool context", + ); + } + + #[cfg(not(feature = "backtrace"))] + #[test] + fn test_disabled_backtrace_for_assert_ne_or_internal_err_without_args() { + fn get_error() -> Result<(), DataFusionError> { + let arg1 = 1; + let arg2 = 1; + assert_ne_or_internal_err!(arg1, arg2); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_internal_err_without_backtrace_and_equal( + &res.unwrap_err(), + "Internal error: Assertion failed: arg1 != arg2 (left: 1, right: 1)", + ); + } + + #[cfg(not(feature = "backtrace"))] + #[test] + fn test_disabled_backtrace_for_assert_ne_or_internal_err_with_args() { + fn get_error() -> Result<(), DataFusionError> { + let arg1 = 1; + let arg2 = 1; + assert_ne_or_internal_err!(arg1, arg2, "my cool context"); + + unreachable!("should return error"); + } + + let res: Result<(), DataFusionError> = get_error(); + assert_internal_err_without_backtrace_and_equal( + &res.unwrap_err(), + "Internal error: Assertion failed: arg1 != arg2 (left: 1, right: 1): my cool context", + ); } #[test] diff --git a/datafusion/common/src/file_options/mod.rs b/datafusion/common/src/file_options/mod.rs index 5d2abd23172ed..97b4a44f03223 100644 --- a/datafusion/common/src/file_options/mod.rs +++ b/datafusion/common/src/file_options/mod.rs @@ -114,14 +114,14 @@ mod tests { properties .bloom_filter_properties(&ColumnPath::from("")) .expect("expected bloom properties!") - .fpp, + .fpp(), 0.123 ); assert_eq!( properties .bloom_filter_properties(&ColumnPath::from("")) .expect("expected bloom properties!") - .ndv, + .ndv(), 123 ); @@ -242,7 +242,7 @@ mod tests { properties .bloom_filter_properties(&col1) .expect("expected bloom properties!") - .fpp, + .fpp(), 0.123 ); @@ -250,7 +250,7 @@ mod tests { properties .bloom_filter_properties(&col2_nested) .expect("expected bloom properties!") - .fpp, + .fpp(), 0.456 ); @@ -258,7 +258,7 @@ mod tests { properties .bloom_filter_properties(&col1) .expect("expected bloom properties!") - .ndv, + .ndv(), 123 ); @@ -266,7 +266,7 @@ mod tests { properties .bloom_filter_properties(&col2_nested) .expect("expected bloom properties!") - .ndv, + .ndv(), 456 ); diff --git a/datafusion/common/src/file_options/parquet_writer.rs b/datafusion/common/src/file_options/parquet_writer.rs index 3f827fbfa75a0..c539245764d45 100644 --- a/datafusion/common/src/file_options/parquet_writer.rs +++ b/datafusion/common/src/file_options/parquet_writer.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use crate::{ _internal_datafusion_err, DataFusionError, Result, - config::{ParquetOptions, TableParquetOptions}, + config::{ParquetCdcOptions, ParquetOptions, TableParquetOptions}, }; use arrow::datatypes::Schema; @@ -157,8 +157,8 @@ impl TryFrom<&TableParquetOptions> for WriterPropertiesBuilder { } if let Some(bloom_filter_ndv) = options.bloom_filter_ndv { - builder = - builder.set_column_bloom_filter_ndv(path.clone(), bloom_filter_ndv); + builder = builder + .set_column_bloom_filter_max_ndv(path.clone(), bloom_filter_ndv); } } @@ -166,6 +166,42 @@ impl TryFrom<&TableParquetOptions> for WriterPropertiesBuilder { } } +/// Convert DataFusion's [`ParquetCdcOptions`] into parquet-rs's `Option`. +/// +/// parquet-rs has no `enabled` flag; CDC is on when the option is `Some`. So a +/// disabled [`ParquetCdcOptions`] maps to `None`, and an enabled one to `Some` +/// with the chunking parameters. +impl From<&ParquetCdcOptions> for Option { + fn from(value: &ParquetCdcOptions) -> Self { + value + .enabled + .then_some(parquet::file::properties::CdcOptions { + min_chunk_size: value.min_chunk_size, + max_chunk_size: value.max_chunk_size, + norm_level: value.norm_level, + }) + } +} + +/// Convert parquet-rs's `Option<&CdcOptions>` back into DataFusion's +/// [`ParquetCdcOptions`]. +/// +/// The presence of parquet-rs options means CDC was enabled, so `Some` maps to +/// `enabled: true`; `None` yields the disabled default. +impl From> for ParquetCdcOptions { + fn from(value: Option<&parquet::file::properties::CdcOptions>) -> Self { + match value { + Some(cdc) => ParquetCdcOptions { + enabled: true, + min_chunk_size: cdc.min_chunk_size, + max_chunk_size: cdc.max_chunk_size, + norm_level: cdc.norm_level, + }, + None => ParquetCdcOptions::default(), + } + } +} + impl ParquetOptions { /// Convert the global session options, [`ParquetOptions`], into a single write action's [`WriterPropertiesBuilder`]. /// @@ -183,6 +219,7 @@ impl ParquetOptions { dictionary_page_size_limit, statistics_enabled, max_row_group_size, + max_row_group_bytes, created_by, column_index_truncate_length, statistics_truncate_length, @@ -191,7 +228,7 @@ impl ParquetOptions { bloom_filter_on_write, bloom_filter_fpp, bloom_filter_ndv, - use_content_defined_chunking, + content_defined_chunking, // not in WriterProperties enable_page_index: _, @@ -211,6 +248,7 @@ impl ParquetOptions { coerce_int96_tz: _, // not used for writer props skip_arrow_metadata: _, max_predicate_cache_size: _, + max_in_list_size: _, } = self; let mut builder = WriterProperties::builder() @@ -225,6 +263,7 @@ impl ParquetOptions { .unwrap_or(DEFAULT_STATISTICS_ENABLED), ) .set_max_row_group_row_count(Some(*max_row_group_size)) + .set_max_row_group_bytes(max_row_group_bytes.as_ref().map(|v| v.get())) .set_created_by(created_by.clone()) .set_column_index_truncate_length(*column_index_truncate_length) .set_statistics_truncate_length(*statistics_truncate_length) @@ -235,7 +274,7 @@ impl ParquetOptions { builder = builder.set_bloom_filter_fpp(*bloom_filter_fpp); }; if let Some(bloom_filter_ndv) = bloom_filter_ndv { - builder = builder.set_bloom_filter_ndv(*bloom_filter_ndv); + builder = builder.set_bloom_filter_max_ndv(*bloom_filter_ndv); }; if let Some(dictionary_enabled) = dictionary_enabled { builder = builder.set_dictionary_enabled(*dictionary_enabled); @@ -249,26 +288,7 @@ impl ParquetOptions { if let Some(encoding) = encoding { builder = builder.set_encoding(parse_encoding_string(encoding)?); } - if let Some(cdc) = use_content_defined_chunking { - if cdc.min_chunk_size == 0 { - return Err(DataFusionError::Configuration( - "CDC min_chunk_size must be greater than 0".to_string(), - )); - } - if cdc.max_chunk_size <= cdc.min_chunk_size { - return Err(DataFusionError::Configuration(format!( - "CDC max_chunk_size ({}) must be greater than min_chunk_size ({})", - cdc.max_chunk_size, cdc.min_chunk_size - ))); - } - builder = builder.set_content_defined_chunking(Some( - parquet::file::properties::CdcOptions { - min_chunk_size: cdc.min_chunk_size, - max_chunk_size: cdc.max_chunk_size, - norm_level: cdc.norm_level, - }, - )); - } + builder = builder.set_content_defined_chunking(content_defined_chunking.into()); Ok(builder) } @@ -411,7 +431,8 @@ mod tests { #[cfg(feature = "parquet_encryption")] use crate::config::ConfigFileEncryptionProperties; use crate::config::{ - CdcOptions, ParquetColumnOptions, ParquetEncryptionOptions, ParquetOptions, + MaxRowGroupBytes, ParquetCdcOptions, ParquetColumnOptions, + ParquetEncryptionOptions, ParquetOptions, }; use crate::parquet_config::DFParquetWriterVersion; use parquet::basic::Compression; @@ -453,9 +474,10 @@ mod tests { writer_version, compression: Some("zstd(22)".into()), dictionary_enabled: Some(!defaults.dictionary_enabled.unwrap_or(false)), - dictionary_page_size_limit: 42, + dictionary_page_size_limit: 43, statistics_enabled: Some("chunk".into()), max_row_group_size: 42, + max_row_group_bytes: Some(MaxRowGroupBytes::try_new(42).unwrap()), created_by: "wordy".into(), column_index_truncate_length: Some(42), statistics_truncate_length: Some(42), @@ -468,6 +490,7 @@ mod tests { // not in WriterProperties, but itemizing here to not skip newly added props enable_page_index: defaults.enable_page_index, pruning: defaults.pruning, + max_in_list_size: defaults.max_in_list_size, skip_metadata: defaults.skip_metadata, metadata_size_hint: defaults.metadata_size_hint, pushdown_filters: defaults.pushdown_filters, @@ -485,7 +508,7 @@ mod tests { coerce_int96: None, coerce_int96_tz: None, max_predicate_cache_size: defaults.max_predicate_cache_size, - use_content_defined_chunking: defaults.use_content_defined_chunking.clone(), + content_defined_chunking: defaults.content_defined_chunking.clone(), } } @@ -513,8 +536,8 @@ mod tests { } .into(), ), - bloom_filter_fpp: bloom_filter_default_props.map(|p| p.fpp), - bloom_filter_ndv: bloom_filter_default_props.map(|p| p.ndv), + bloom_filter_fpp: bloom_filter_default_props.map(|p| p.fpp()), + bloom_filter_ndv: bloom_filter_default_props.map(|p| p.ndv()), } } @@ -558,13 +581,16 @@ mod tests { TableParquetOptions { global: ParquetOptions { // global options - data_pagesize_limit: props.dictionary_page_size_limit(), + data_pagesize_limit: props.data_page_size_limit(), write_batch_size: props.write_batch_size(), writer_version: props.writer_version().into(), dictionary_page_size_limit: props.dictionary_page_size_limit(), max_row_group_size: props .max_row_group_row_count() .unwrap_or(DEFAULT_MAX_ROW_GROUP_ROW_COUNT), + max_row_group_bytes: props + .max_row_group_bytes() + .and_then(|v| MaxRowGroupBytes::try_new(v).ok()), created_by: props.created_by().to_string(), column_index_truncate_length: props.column_index_truncate_length(), statistics_truncate_length: props.statistics_truncate_length(), @@ -584,6 +610,7 @@ mod tests { // not in WriterProperties enable_page_index: global_options_defaults.enable_page_index, pruning: global_options_defaults.pruning, + max_in_list_size: global_options_defaults.max_in_list_size, skip_metadata: global_options_defaults.skip_metadata, metadata_size_hint: global_options_defaults.metadata_size_hint, pushdown_filters: global_options_defaults.pushdown_filters, @@ -603,13 +630,7 @@ mod tests { skip_arrow_metadata: global_options_defaults.skip_arrow_metadata, coerce_int96: None, coerce_int96_tz: None, - use_content_defined_chunking: props.content_defined_chunking().map(|c| { - CdcOptions { - min_chunk_size: c.min_chunk_size, - max_chunk_size: c.max_chunk_size, - norm_level: c.norm_level, - } - }), + content_defined_chunking: props.content_defined_chunking().into(), }, column_specific_options, key_value_metadata, @@ -812,10 +833,12 @@ mod tests { ); assert_eq!( default_writer_props.bloom_filter_properties(&"default".into()), - Some(&BloomFilterProperties { - fpp: 0.42, - ndv: DEFAULT_BLOOM_FILTER_NDV - }), + Some( + &BloomFilterProperties::builder() + .with_fpp(0.42) + .with_max_ndv(DEFAULT_BLOOM_FILTER_NDV) + .build() + ), "should have only the fpp set, and the ndv at default", ); } @@ -823,11 +846,12 @@ mod tests { #[test] fn test_cdc_enabled_with_custom_options() { let mut opts = TableParquetOptions::default(); - opts.global.use_content_defined_chunking = Some(CdcOptions { + opts.global.content_defined_chunking = ParquetCdcOptions { + enabled: true, min_chunk_size: 128 * 1024, max_chunk_size: 512 * 1024, norm_level: 2, - }); + }; opts.arrow_schema(&Arc::new(Schema::empty())); let props = WriterPropertiesBuilder::try_from(&opts).unwrap().build(); @@ -846,46 +870,61 @@ mod tests { assert!(props.content_defined_chunking().is_none()); } + #[test] + fn test_cdc_params_ignored_when_disabled() { + // Parameters are customized but `enabled` is false, so CDC stays off. + let mut opts = TableParquetOptions::default(); + opts.global.content_defined_chunking = ParquetCdcOptions { + enabled: false, + min_chunk_size: 128 * 1024, + max_chunk_size: 512 * 1024, + norm_level: 2, + }; + opts.arrow_schema(&Arc::new(Schema::empty())); + + let props = WriterPropertiesBuilder::try_from(&opts).unwrap().build(); + assert!(props.content_defined_chunking().is_none()); + } + #[test] fn test_cdc_round_trip_through_writer_props() { let mut opts = TableParquetOptions::default(); - opts.global.use_content_defined_chunking = Some(CdcOptions { + opts.global.content_defined_chunking = ParquetCdcOptions { + enabled: true, min_chunk_size: 64 * 1024, max_chunk_size: 2 * 1024 * 1024, norm_level: -1, - }); + }; opts.arrow_schema(&Arc::new(Schema::empty())); let props = WriterPropertiesBuilder::try_from(&opts).unwrap().build(); let recovered = session_config_from_writer_props(&props); - let cdc = recovered.global.use_content_defined_chunking.unwrap(); + let cdc = recovered.global.content_defined_chunking; + assert!(cdc.enabled); assert_eq!(cdc.min_chunk_size, 64 * 1024); assert_eq!(cdc.max_chunk_size, 2 * 1024 * 1024); assert_eq!(cdc.norm_level, -1); } #[test] - fn test_cdc_validation_zero_min_chunk_size() { + fn test_max_row_group_bytes_disabled_by_default() { let mut opts = TableParquetOptions::default(); - opts.global.use_content_defined_chunking = Some(CdcOptions { - min_chunk_size: 0, - ..CdcOptions::default() - }); opts.arrow_schema(&Arc::new(Schema::empty())); - assert!(WriterPropertiesBuilder::try_from(&opts).is_err()); + + let props = WriterPropertiesBuilder::try_from(&opts).unwrap().build(); + assert_eq!(props.max_row_group_bytes(), None); } #[test] - fn test_cdc_validation_max_not_greater_than_min() { + fn test_max_row_group_bytes_propagated_to_writer_props() { let mut opts = TableParquetOptions::default(); - opts.global.use_content_defined_chunking = Some(CdcOptions { - min_chunk_size: 512 * 1024, - max_chunk_size: 256 * 1024, - ..CdcOptions::default() - }); + opts.global.max_row_group_bytes = + Some(MaxRowGroupBytes::try_new(64 * 1024 * 1024).unwrap()); opts.arrow_schema(&Arc::new(Schema::empty())); - assert!(WriterPropertiesBuilder::try_from(&opts).is_err()); + + let props = WriterPropertiesBuilder::try_from(&opts).unwrap().build(); + assert_eq!(props.max_row_group_bytes(), Some(64 * 1024 * 1024)); } #[test] @@ -903,7 +942,7 @@ mod tests { // the WriterProperties::default, with only ndv set let default_writer_props = WriterProperties::builder() .set_bloom_filter_enabled(true) - .set_bloom_filter_ndv(42) + .set_bloom_filter_max_ndv(42) .build(); assert_eq!( @@ -913,10 +952,12 @@ mod tests { ); assert_eq!( default_writer_props.bloom_filter_properties(&"default".into()), - Some(&BloomFilterProperties { - fpp: DEFAULT_BLOOM_FILTER_FPP, - ndv: 42 - }), + Some( + &BloomFilterProperties::builder() + .with_fpp(DEFAULT_BLOOM_FILTER_FPP) + .with_max_ndv(42) + .build() + ), "should have only the ndv set, and the fpp at default", ); } diff --git a/datafusion/common/src/format.rs b/datafusion/common/src/format.rs index a6bd42be691a9..ea88eca4a65bc 100644 --- a/datafusion/common/src/format.rs +++ b/datafusion/common/src/format.rs @@ -23,6 +23,8 @@ use arrow::util::display::{DurationFormat, FormatOptions}; use crate::config::{ConfigField, Visit}; use crate::error::{DataFusionError, Result}; +#[cfg(feature = "sql")] +use sqlparser::ast::{Expr, UtilityOption, Value, ValueWithSpan}; /// The default [`FormatOptions`] to use within DataFusion /// Also see [`crate::config::FormatOptions`] @@ -430,3 +432,470 @@ impl ConfigField for ExplainAnalyzeCategories { Ok(()) } } + +/// Normalized options for a single `EXPLAIN` statement. +/// +/// This collects the knobs that can be set per-statement from either the +/// legacy keyword form (`EXPLAIN ANALYZE VERBOSE FORMAT tree ...`) or the +/// Postgres-style `EXPLAIN (option [arg], ...) ...` form supported on +/// dialects whose +/// [`Dialect::supports_explain_with_utility_options`](https://docs.rs/sqlparser/latest/sqlparser/dialect/trait.Dialect.html#method.supports_explain_with_utility_options) +/// returns `true`. +/// +/// Fields that are `None` / `false` mean "not set at the statement level" — +/// the physical planner falls back to the corresponding session config +/// value. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct ExplainStatementOptions { + /// Whether to actually execute the plan and gather metrics. + /// + /// Corresponds to the `ANALYZE` keyword or the `ANALYZE` option. + pub analyze: bool, + /// Whether to include extra detail in the output. + /// + /// Corresponds to the `VERBOSE` keyword or the `VERBOSE` option. + pub verbose: bool, + /// Output format for the plan. When `None`, the session-config + /// default (`datafusion.explain.format`) is used. + pub format: Option, + /// Override for [`MetricType`] (summary / dev) when running + /// `EXPLAIN ANALYZE`. + pub analyze_level: Option, + /// Override for [`ExplainAnalyzeCategories`] (rows / bytes / timing + /// / uncategorized) when running `EXPLAIN ANALYZE`. + pub analyze_categories: Option, + /// Override for `datafusion.explain.show_statistics`. + pub show_statistics: Option, +} + +#[cfg(feature = "sql")] +impl ExplainStatementOptions { + /// Parse a list of [`UtilityOption`] values (produced by sqlparser's + /// `parse_utility_options`) into a normalized [`ExplainStatementOptions`]. + /// + /// Argument grammar accepted: + /// - `OPTION` — bare, implies `TRUE` for boolean options. + /// - `OPTION TRUE` / `OPTION FALSE` + /// - `OPTION ON` / `OPTION OFF` + /// - `OPTION 1` / `OPTION 0` + /// - `OPTION ` or `OPTION ''` for format / level / metrics. + /// + /// Options recognized by DataFusion are: `ANALYZE`, `VERBOSE`, `FORMAT`, + /// `METRICS`, `LEVEL`, `TIMING`, `SUMMARY`, `COSTS`. + /// + /// Postgres-only options (`BUFFERS`, `WAL`, `SETTINGS`, `GENERIC_PLAN`, + /// `MEMORY`) return a helpful "not supported" error. Any other option + /// name produces an `unknown EXPLAIN option` error. + pub fn from_utility_options(opts: &[UtilityOption]) -> Result { + let mut out = ExplainStatementOptions::default(); + // Track whether METRICS was explicitly set so TIMING can merge + // into it rather than overwrite. + let mut metrics_explicit = false; + + for opt in opts { + let name = opt.name.value.to_ascii_lowercase(); + match name.as_str() { + "analyze" => { + out.analyze = parse_bool_arg(&opt.arg, &name)?; + } + "verbose" => { + out.verbose = parse_bool_arg(&opt.arg, &name)?; + } + "format" => { + let s = parse_ident_or_string_arg(&opt.arg, &name)?; + out.format = Some(ExplainFormat::from_str(&s)?); + } + "metrics" => { + let s = parse_ident_or_string_arg(&opt.arg, &name)?; + out.analyze_categories = + Some(ExplainAnalyzeCategories::from_str(&s)?); + metrics_explicit = true; + } + "level" => { + let s = parse_ident_or_string_arg(&opt.arg, &name)?; + out.analyze_level = Some(MetricType::from_str(&s)?); + } + "timing" => { + let enable = parse_bool_arg(&opt.arg, &name)?; + out.analyze_categories = Some(adjust_timing( + out.analyze_categories.take(), + enable, + metrics_explicit, + )); + } + "summary" => { + let summary = parse_bool_arg(&opt.arg, &name)?; + out.analyze_level = Some(if summary { + MetricType::Summary + } else { + MetricType::Dev + }); + } + "costs" => { + out.show_statistics = Some(parse_bool_arg(&opt.arg, &name)?); + } + // Postgres options DataFusion does not model. Give a helpful + // pointer rather than silently accepting them. + "buffers" | "wal" | "settings" | "generic_plan" | "memory" => { + let upper = name.to_ascii_uppercase(); + return Err(DataFusionError::NotImplemented(format!( + "EXPLAIN option {upper} is not supported by DataFusion; \ + see METRICS for category filtering" + ))); + } + _ => { + return Err(DataFusionError::Plan(format!( + "unknown EXPLAIN option: {}", + opt.name.value + ))); + } + } + } + + Ok(out) + } +} + +/// Parse a boolean argument for an EXPLAIN option. +/// +/// `None` (bare option, e.g. `ANALYZE`) is treated as `true`. Accepts +/// identifiers `TRUE`/`FALSE`/`ON`/`OFF` (case-insensitive) and the numeric +/// literals `0` / `1`. +#[cfg(feature = "sql")] +fn parse_bool_arg(arg: &Option, name: &str) -> Result { + let Some(expr) = arg else { + return Ok(true); + }; + match expr { + Expr::Identifier(ident) => match ident.value.to_ascii_lowercase().as_str() { + "true" | "on" => Ok(true), + "false" | "off" => Ok(false), + other => Err(DataFusionError::Plan(format!( + "expected boolean for EXPLAIN option {name}, got '{other}'" + ))), + }, + Expr::Value(ValueWithSpan { value, .. }) => match value { + Value::Boolean(b) => Ok(*b), + Value::Number(n, _) => match n.as_str() { + "0" => Ok(false), + "1" => Ok(true), + other => Err(DataFusionError::Plan(format!( + "expected boolean (0 or 1) for EXPLAIN option {name}, got '{other}'" + ))), + }, + Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => { + match s.to_ascii_lowercase().as_str() { + "true" | "on" | "1" => Ok(true), + "false" | "off" | "0" => Ok(false), + other => Err(DataFusionError::Plan(format!( + "expected boolean for EXPLAIN option {name}, got '{other}'" + ))), + } + } + other => Err(DataFusionError::Plan(format!( + "expected boolean for EXPLAIN option {name}, got '{other}'" + ))), + }, + other => Err(DataFusionError::Plan(format!( + "expected boolean for EXPLAIN option {name}, got '{other}'" + ))), + } +} + +/// Parse an identifier-or-string argument (used for `FORMAT`, `METRICS`, +/// `LEVEL`). +#[cfg(feature = "sql")] +fn parse_ident_or_string_arg(arg: &Option, name: &str) -> Result { + let expr = arg.as_ref().ok_or_else(|| { + DataFusionError::Plan(format!( + "EXPLAIN option {} requires an argument", + name.to_ascii_uppercase() + )) + })?; + match expr { + Expr::Identifier(ident) => Ok(ident.value.clone()), + Expr::Value(ValueWithSpan { value, .. }) => match value { + Value::SingleQuotedString(s) | Value::DoubleQuotedString(s) => Ok(s.clone()), + other => Err(DataFusionError::Plan(format!( + "expected identifier or string for EXPLAIN option {name}, got '{other}'" + ))), + }, + other => Err(DataFusionError::Plan(format!( + "expected identifier or string for EXPLAIN option {name}, got '{other}'" + ))), + } +} + +/// Merge a `TIMING on/off` option into an existing `METRICS` selection. +/// +/// If METRICS was already specified, we only add/remove the Timing category +/// within that selection. If METRICS was not specified, TIMING effectively +/// means "Only(Timing)" when on, or "show everything except timing" when off. +#[cfg(feature = "sql")] +fn adjust_timing( + current: Option, + enable: bool, + metrics_explicit: bool, +) -> ExplainAnalyzeCategories { + // METRICS was not specified — TIMING alone shapes the selection. + if !metrics_explicit { + return if enable { + ExplainAnalyzeCategories::All + } else { + ExplainAnalyzeCategories::Only(vec![ + MetricCategory::Rows, + MetricCategory::Bytes, + MetricCategory::Uncategorized, + ]) + }; + } + + // METRICS was specified explicitly earlier — merge into its list. When + // METRICS was explicit, `current` is always `Some(_)`; fall back to All + // to be safe. + match current.unwrap_or(ExplainAnalyzeCategories::All) { + ExplainAnalyzeCategories::All if enable => ExplainAnalyzeCategories::All, + ExplainAnalyzeCategories::All => { + // Everything except timing: rows, bytes, uncategorized. + ExplainAnalyzeCategories::Only(vec![ + MetricCategory::Rows, + MetricCategory::Bytes, + MetricCategory::Uncategorized, + ]) + } + ExplainAnalyzeCategories::Only(mut cats) if enable => { + if !cats.contains(&MetricCategory::Timing) { + cats.push(MetricCategory::Timing); + } + ExplainAnalyzeCategories::Only(cats) + } + ExplainAnalyzeCategories::Only(cats) => ExplainAnalyzeCategories::Only( + cats.into_iter() + .filter(|c| *c != MetricCategory::Timing) + .collect(), + ), + } +} + +#[cfg(all(test, feature = "sql"))] +mod explain_options_tests { + use super::*; + use sqlparser::ast::Ident; + use sqlparser::tokenizer::Span; + + fn bare(name: &str) -> UtilityOption { + UtilityOption { + name: Ident { + value: name.to_string(), + quote_style: None, + span: Span::empty(), + }, + arg: None, + } + } + + fn with_ident_arg(name: &str, arg: &str) -> UtilityOption { + UtilityOption { + name: Ident { + value: name.to_string(), + quote_style: None, + span: Span::empty(), + }, + arg: Some(Expr::Identifier(Ident { + value: arg.to_string(), + quote_style: None, + span: Span::empty(), + })), + } + } + + fn with_string_arg(name: &str, arg: &str) -> UtilityOption { + UtilityOption { + name: Ident { + value: name.to_string(), + quote_style: None, + span: Span::empty(), + }, + arg: Some(Expr::Value(ValueWithSpan { + value: Value::SingleQuotedString(arg.to_string()), + span: Span::empty(), + })), + } + } + + fn with_bool_arg(name: &str, b: bool) -> UtilityOption { + UtilityOption { + name: Ident { + value: name.to_string(), + quote_style: None, + span: Span::empty(), + }, + arg: Some(Expr::Value(ValueWithSpan { + value: Value::Boolean(b), + span: Span::empty(), + })), + } + } + + fn with_number_arg(name: &str, n: &str) -> UtilityOption { + UtilityOption { + name: Ident { + value: name.to_string(), + quote_style: None, + span: Span::empty(), + }, + arg: Some(Expr::Value(ValueWithSpan { + value: Value::Number(n.to_string(), false), + span: Span::empty(), + })), + } + } + + #[test] + fn bare_analyze_and_verbose() { + let opts = ExplainStatementOptions::from_utility_options(&[ + bare("ANALYZE"), + bare("VERBOSE"), + ]) + .unwrap(); + assert!(opts.analyze); + assert!(opts.verbose); + assert!(opts.format.is_none()); + } + + #[test] + fn format_from_ident_and_string() { + let opts = ExplainStatementOptions::from_utility_options(&[with_ident_arg( + "FORMAT", "tree", + )]) + .unwrap(); + assert_eq!(opts.format, Some(ExplainFormat::Tree)); + + let opts = ExplainStatementOptions::from_utility_options(&[with_string_arg( + "FORMAT", "pgjson", + )]) + .unwrap(); + assert_eq!(opts.format, Some(ExplainFormat::PostgresJSON)); + } + + #[test] + fn metrics_and_level() { + let opts = ExplainStatementOptions::from_utility_options(&[ + with_string_arg("METRICS", "rows,bytes"), + with_ident_arg("LEVEL", "dev"), + ]) + .unwrap(); + assert_eq!( + opts.analyze_categories, + Some(ExplainAnalyzeCategories::Only(vec![ + MetricCategory::Rows, + MetricCategory::Bytes, + ])) + ); + assert_eq!(opts.analyze_level, Some(MetricType::Dev)); + } + + #[test] + fn on_off_numeric_bool() { + let opts = ExplainStatementOptions::from_utility_options(&[ + with_ident_arg("ANALYZE", "ON"), + with_ident_arg("VERBOSE", "off"), + with_bool_arg("COSTS", true), + ]) + .unwrap(); + assert!(opts.analyze); + assert!(!opts.verbose); + assert_eq!(opts.show_statistics, Some(true)); + + let opts = ExplainStatementOptions::from_utility_options(&[ + with_number_arg("ANALYZE", "1"), + with_number_arg("VERBOSE", "0"), + ]) + .unwrap(); + assert!(opts.analyze); + assert!(!opts.verbose); + } + + #[test] + fn summary_sugar_sets_level() { + let opts = ExplainStatementOptions::from_utility_options(&[with_ident_arg( + "SUMMARY", "ON", + )]) + .unwrap(); + assert_eq!(opts.analyze_level, Some(MetricType::Summary)); + + let opts = ExplainStatementOptions::from_utility_options(&[with_bool_arg( + "SUMMARY", false, + )]) + .unwrap(); + assert_eq!(opts.analyze_level, Some(MetricType::Dev)); + } + + #[test] + fn timing_merges_with_metrics() { + // METRICS then TIMING off → timing is removed from the list + let opts = ExplainStatementOptions::from_utility_options(&[ + with_string_arg("METRICS", "rows,timing"), + with_bool_arg("TIMING", false), + ]) + .unwrap(); + assert_eq!( + opts.analyze_categories, + Some(ExplainAnalyzeCategories::Only(vec![MetricCategory::Rows])) + ); + + // METRICS 'rows' then TIMING on → timing is appended + let opts = ExplainStatementOptions::from_utility_options(&[ + with_string_arg("METRICS", "rows"), + with_bool_arg("TIMING", true), + ]) + .unwrap(); + assert_eq!( + opts.analyze_categories, + Some(ExplainAnalyzeCategories::Only(vec![ + MetricCategory::Rows, + MetricCategory::Timing, + ])) + ); + } + + #[test] + fn timing_alone() { + let opts = ExplainStatementOptions::from_utility_options(&[with_bool_arg( + "TIMING", false, + )]) + .unwrap(); + assert_eq!( + opts.analyze_categories, + Some(ExplainAnalyzeCategories::Only(vec![ + MetricCategory::Rows, + MetricCategory::Bytes, + MetricCategory::Uncategorized, + ])) + ); + } + + #[test] + fn unknown_option_rejected() { + let err = + ExplainStatementOptions::from_utility_options(&[bare("FOO")]).unwrap_err(); + assert!( + err.to_string().contains("unknown EXPLAIN option: FOO"), + "got: {err}" + ); + } + + #[test] + fn postgres_only_options_rejected() { + for pg_only in ["BUFFERS", "WAL", "SETTINGS", "GENERIC_PLAN", "MEMORY"] { + let err = ExplainStatementOptions::from_utility_options(&[bare(pg_only)]) + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains(pg_only), + "msg did not include {pg_only}: {msg}" + ); + assert!(msg.contains("not supported"), "msg: {msg}"); + } + } +} diff --git a/datafusion/common/src/functional_dependencies.rs b/datafusion/common/src/functional_dependencies.rs index 24ca33c0c2c90..8b15c49c565f1 100644 --- a/datafusion/common/src/functional_dependencies.rs +++ b/datafusion/common/src/functional_dependencies.rs @@ -151,8 +151,10 @@ pub struct FunctionalDependence { /// Describes functional dependency mode. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum Dependency { - Single, // A determinant key may occur only once. - Multi, // A determinant key may occur multiple times (in multiple rows). + /// A determinant key may occur only once. + Single, + /// A determinant key may occur multiple times (in multiple rows). + Multi, } impl FunctionalDependence { diff --git a/datafusion/common/src/hash_utils.rs b/datafusion/common/src/hash_utils.rs index fcc2e919b6cc2..cfe57999689b1 100644 --- a/datafusion/common/src/hash_utils.rs +++ b/datafusion/common/src/hash_utils.rs @@ -31,8 +31,54 @@ use itertools::Itertools; use std::collections::HashMap; use std::hash::{BuildHasher, Hash, Hasher}; -/// The hash random state used throughout DataFusion for hashing. +/// [`RandomState`] is optimized for speed and suitable for hash tables and +/// bloom filters. [`QualityRandomState`] is optimized for statistical quality +/// and suitable for algorithms such as HyperLogLog. The tradeoff is that the +/// fast variant gives up some statistical quality, while the quality variant +/// is slightly slower. +/// +/// See: pub type RandomState = FixedState; +pub type QualityRandomState = foldhash::quality::FixedState; + +/// Fixed quality hash state used by HyperLogLog sketches. +/// +/// The seed is part of the HLL wire/storage semantics: serialized sketches only +/// remain mergeable if every producer uses the same hash state. +pub const HLL_RANDOM_STATE: QualityRandomState = QualityRandomState::with_seed(0); + +/// Hash state used by [`create_hashes`]. +/// +/// Multi-column hashing folds the previous column hash into a fresh hasher +/// before hashing the next column. This trait keeps that seeded hasher in the +/// same foldhash tier as the top-level hash state. +pub trait HashState: BuildHasher { + type SeededState: BuildHasher; + + fn seeded_state(&self, seed: u64) -> Self::SeededState; +} + +impl HashState for FixedState { + type SeededState = foldhash::fast::SeedableRandomState; + + fn seeded_state(&self, seed: u64) -> Self::SeededState { + foldhash::fast::SeedableRandomState::with_seed( + seed, + foldhash::SharedSeed::global_fixed(), + ) + } +} + +impl HashState for foldhash::quality::FixedState { + type SeededState = foldhash::quality::SeedableRandomState; + + fn seeded_state(&self, seed: u64) -> Self::SeededState { + foldhash::quality::SeedableRandomState::with_seed( + seed, + foldhash::SharedSeed::global_fixed(), + ) + } +} #[cfg(not(feature = "force_hash_collisions"))] use crate::cast::{ @@ -45,6 +91,8 @@ use crate::error::Result; use crate::error::{_internal_datafusion_err, _internal_err}; use std::cell::RefCell; +mod build_hasher; + // Combines two hashes into one hash #[inline] pub fn combine_hashes(l: u64, r: u64) -> u64 { @@ -99,7 +147,7 @@ thread_local! { /// ``` pub fn with_hashes( arrays: I, - random_state: &RandomState, + random_state: &impl HashState, callback: F, ) -> Result where @@ -140,9 +188,32 @@ where }).map_err(|_| _internal_datafusion_err!("with_hashes cannot access thread-local storage during or after thread destruction"))? } +/// Creates hashes for the given arrays using a thread-local buffer and a custom +/// hash builder, then calls the provided callback with the computed hashes. +/// +/// Hash compatibility with [`with_hashes`] follows the rules documented on +/// [`create_hashes_with_hasher`]. +pub fn with_hashes_with_hasher( + arrays: I, + hash_builder: &S, + callback: F, +) -> Result +where + I: IntoIterator, + T: AsDynArray, + F: FnOnce(&[u64]) -> Result, + S: BuildHasher, +{ + build_hasher::with_hashes_with_hasher(arrays, hash_builder, callback) +} + #[cfg(not(feature = "force_hash_collisions"))] -fn hash_null(random_state: &RandomState, hashes_buffer: &'_ mut [u64], mul_col: bool) { - if mul_col { +fn hash_null( + random_state: &S, + hashes_buffer: &'_ mut [u64], + multi_col: bool, +) { + if multi_col { hashes_buffer.iter_mut().for_each(|hash| { // stable hash for null value *hash = combine_hashes(random_state.hash_one(1), *hash); @@ -155,13 +226,13 @@ fn hash_null(random_state: &RandomState, hashes_buffer: &'_ mut [u64], mul_col: } pub trait HashValue { - fn hash_one(&self, state: &RandomState) -> u64; + fn hash_one(&self, state: &S) -> u64; /// Write this value into an existing hasher (same data as `hash_one`). fn hash_write(&self, hasher: &mut impl Hasher); } impl HashValue for &T { - fn hash_one(&self, state: &RandomState) -> u64 { + fn hash_one(&self, state: &S) -> u64 { T::hash_one(self, state) } fn hash_write(&self, hasher: &mut impl Hasher) { @@ -172,7 +243,7 @@ impl HashValue for &T { macro_rules! hash_value { ($($t:ty),+) => { $(impl HashValue for $t { - fn hash_one(&self, state: &RandomState) -> u64 { + fn hash_one(&self, state: &S) -> u64 { state.hash_one(self) } fn hash_write(&self, hasher: &mut impl Hasher) { @@ -187,27 +258,45 @@ hash_value!(bool, str, [u8], IntervalDayTime, IntervalMonthDayNano); macro_rules! hash_float_value { ($(($t:ty, $i:ty)),+) => { $(impl HashValue for $t { - fn hash_one(&self, state: &RandomState) -> u64 { - state.hash_one(<$i>::from_ne_bytes(self.to_ne_bytes())) + fn hash_one(&self, state: &S) -> u64 { + // +0.0 and -0.0 differ only in the sign bit but compare equal + // under IEEE 754; normalize -0.0 → +0.0 so Hash agrees with Eq. + let bits = <$i>::from_ne_bytes(self.to_ne_bytes()); + let bits = if bits << 1 == 0 { 0 } else { bits }; + state.hash_one(bits) } fn hash_write(&self, hasher: &mut impl Hasher) { - hasher.write(&self.to_ne_bytes()) + let bits = <$i>::from_ne_bytes(self.to_ne_bytes()); + let bits: $i = if bits << 1 == 0 { 0 } else { bits }; + hasher.write(&bits.to_ne_bytes()) } })+ }; } hash_float_value!((half::f16, u16), (f32, u32), (f64, u64)); -/// Create a `SeedableRandomState` whose per-hasher seed incorporates `seed`. -/// This folds the previous hash into the hasher's initial state so only the -/// new value needs to pass through the hash function — same cost as `hash_one`. #[cfg(not(feature = "force_hash_collisions"))] -#[inline] -fn seeded_state(seed: u64) -> foldhash::fast::SeedableRandomState { - foldhash::fast::SeedableRandomState::with_seed( - seed, - foldhash::SharedSeed::global_fixed(), - ) +trait ChildHashing { + fn create_hashes(&self, arrays: I, hashes_buffer: &mut [u64]) -> Result<()> + where + I: IntoIterator, + T: AsDynArray; +} + +#[cfg(not(feature = "force_hash_collisions"))] +struct HashStateChildHashing<'a, S> { + hash_state: &'a S, +} + +#[cfg(not(feature = "force_hash_collisions"))] +impl ChildHashing for HashStateChildHashing<'_, S> { + fn create_hashes(&self, arrays: I, hashes_buffer: &mut [u64]) -> Result<()> + where + I: IntoIterator, + T: AsDynArray, + { + create_hashes(arrays, self.hash_state, hashes_buffer).map(|_| ()) + } } /// Builds hash values of PrimitiveArray and writes them into `hashes_buffer` @@ -216,7 +305,7 @@ fn seeded_state(seed: u64) -> foldhash::fast::SeedableRandomState { #[cfg(not(feature = "force_hash_collisions"))] fn hash_array_primitive( array: &PrimitiveArray, - random_state: &RandomState, + random_state: &impl HashState, hashes_buffer: &mut [u64], rehash: bool, ) where @@ -231,7 +320,7 @@ fn hash_array_primitive( if array.null_count() == 0 { if rehash { for (hash, &value) in hashes_buffer.iter_mut().zip(array.values().iter()) { - let mut hasher = seeded_state(*hash).build_hasher(); + let mut hasher = random_state.seeded_state(*hash).build_hasher(); value.hash_write(&mut hasher); *hash = hasher.finish(); } @@ -243,7 +332,7 @@ fn hash_array_primitive( } else if rehash { for i in array.nulls().unwrap().valid_indices() { let value = unsafe { array.value_unchecked(i) }; - let mut hasher = seeded_state(hashes_buffer[i]).build_hasher(); + let mut hasher = random_state.seeded_state(hashes_buffer[i]).build_hasher(); value.hash_write(&mut hasher); hashes_buffer[i] = hasher.finish(); } @@ -261,7 +350,7 @@ fn hash_array_primitive( #[cfg(not(feature = "force_hash_collisions"))] fn hash_array( array: &T, - random_state: &RandomState, + random_state: &impl HashState, hashes_buffer: &mut [u64], rehash: bool, ) where @@ -316,7 +405,7 @@ fn hash_string_view_array_inner< const REHASH: bool, >( array: &GenericByteViewArray, - random_state: &RandomState, + random_state: &impl HashState, hashes_buffer: &mut [u64], ) { assert_eq!( @@ -345,7 +434,7 @@ fn hash_string_view_array_inner< // all views are inlined, no need to access external buffers if !HAS_BUFFERS || view_len <= 12 { if REHASH { - let mut hasher = seeded_state(*hash).build_hasher(); + let mut hasher = random_state.seeded_state(*hash).build_hasher(); v.hash_write(&mut hasher); *hash = hasher.finish(); } else { @@ -356,7 +445,7 @@ fn hash_string_view_array_inner< // view is not inlined, so we need to hash the bytes as well let value = view_bytes(view_len, v); if REHASH { - let mut hasher = seeded_state(*hash).build_hasher(); + let mut hasher = random_state.seeded_state(*hash).build_hasher(); value.hash_write(&mut hasher); *hash = hasher.finish(); } else { @@ -371,7 +460,7 @@ fn hash_string_view_array_inner< #[cfg(not(feature = "force_hash_collisions"))] fn hash_generic_byte_view_array( array: &GenericByteViewArray, - random_state: &RandomState, + random_state: &impl HashState, hashes_buffer: &mut [u64], rehash: bool, ) { @@ -390,7 +479,7 @@ fn hash_generic_byte_view_array( } (false, false, true) => { for (hash, &view) in hashes_buffer.iter_mut().zip(array.views().iter()) { - let mut hasher = seeded_state(*hash).build_hasher(); + let mut hasher = random_state.seeded_state(*hash).build_hasher(); view.hash_write(&mut hasher); *hash = hasher.finish(); } @@ -428,31 +517,25 @@ fn hash_generic_byte_view_array( } } -/// Hash dictionary array with compile-time specialization for null handling. +/// Scatter precomputed dictionary value hashes to key positions. /// -/// Uses const generics to eliminate runtim branching in the hot loop: +/// Uses const generics to eliminate runtime branching in the hot loop: /// - `HAS_NULL_KEYS`: Whether to check for null dictionary keys /// - `HAS_NULL_VALUES`: Whether to check for null dictionary values /// - `MULTI_COL`: Whether to combine with existing hash (true) or initialize (false) #[cfg(not(feature = "force_hash_collisions"))] #[inline(never)] -fn hash_dictionary_inner< +fn hash_dictionary_scatter< K: ArrowDictionaryKeyType, const HAS_NULL_KEYS: bool, const HAS_NULL_VALUES: bool, const MULTI_COL: bool, >( array: &DictionaryArray, - random_state: &RandomState, + dict_hashes: &[u64], hashes_buffer: &mut [u64], -) -> Result<()> { - // Hash each dictionary value once, and then use that computed - // hash for each key value to avoid a potentially expensive - // redundant hashing for large dictionary elements (e.g. strings) +) { let dict_values = array.values(); - let mut dict_hashes = vec![0; dict_values.len()]; - create_hashes([dict_values], random_state, &mut dict_hashes)?; - if HAS_NULL_KEYS { for (hash, key) in hashes_buffer.iter_mut().zip(array.keys().iter()) { if let Some(key) = key { @@ -478,70 +561,98 @@ fn hash_dictionary_inner< } } } - Ok(()) } -/// Hash the values in a dictionary array #[cfg(not(feature = "force_hash_collisions"))] -fn hash_dictionary( +fn dispatch_dictionary_scatter( array: &DictionaryArray, - random_state: &RandomState, + dict_hashes: &[u64], hashes_buffer: &mut [u64], multi_col: bool, -) -> Result<()> { +) { let has_null_keys = array.keys().null_count() != 0; let has_null_values = array.values().null_count() != 0; - // Dispatcher based on null presence and multi-column mode - // Should reduce branching within hot loops match (has_null_keys, has_null_values, multi_col) { - (false, false, false) => hash_dictionary_inner::( + (false, false, false) => hash_dictionary_scatter::( array, - random_state, + dict_hashes, hashes_buffer, ), - (false, false, true) => hash_dictionary_inner::( + (false, false, true) => hash_dictionary_scatter::( array, - random_state, + dict_hashes, hashes_buffer, ), - (false, true, false) => hash_dictionary_inner::( + (false, true, false) => hash_dictionary_scatter::( array, - random_state, + dict_hashes, hashes_buffer, ), - (false, true, true) => hash_dictionary_inner::( + (false, true, true) => hash_dictionary_scatter::( array, - random_state, + dict_hashes, hashes_buffer, ), - (true, false, false) => hash_dictionary_inner::( + (true, false, false) => hash_dictionary_scatter::( array, - random_state, + dict_hashes, hashes_buffer, ), - (true, false, true) => hash_dictionary_inner::( + (true, false, true) => hash_dictionary_scatter::( array, - random_state, + dict_hashes, hashes_buffer, ), - (true, true, false) => hash_dictionary_inner::( + (true, true, false) => hash_dictionary_scatter::( array, - random_state, + dict_hashes, hashes_buffer, ), - (true, true, true) => hash_dictionary_inner::( + (true, true, true) => hash_dictionary_scatter::( array, - random_state, + dict_hashes, hashes_buffer, ), } } +/// Hash the values in a dictionary array. +#[cfg(not(feature = "force_hash_collisions"))] +fn hash_dictionary( + array: &DictionaryArray, + random_state: &impl HashState, + hashes_buffer: &mut [u64], + multi_col: bool, +) -> Result<()> { + // Hash each dictionary value once, and then use that computed + // hash for each key value to avoid a potentially expensive + // redundant hashing for large dictionary elements (e.g. strings) + let dict_values = array.values(); + let mut dict_hashes = vec![0; dict_values.len()]; + create_hashes([dict_values], random_state, &mut dict_hashes)?; + dispatch_dictionary_scatter(array, &dict_hashes, hashes_buffer, multi_col); + Ok(()) +} + +#[cfg(not(feature = "force_hash_collisions"))] +fn hash_dictionary_with_child_hashing( + array: &DictionaryArray, + child_hashing: &impl ChildHashing, + hashes_buffer: &mut [u64], + multi_col: bool, +) -> Result<()> { + let dict_values = array.values(); + let mut dict_hashes = vec![0; dict_values.len()]; + child_hashing.create_hashes([dict_values], &mut dict_hashes)?; + dispatch_dictionary_scatter(array, &dict_hashes, hashes_buffer, multi_col); + Ok(()) +} + #[cfg(not(feature = "force_hash_collisions"))] fn hash_struct_array( array: &StructArray, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], ) -> Result<()> { let nulls = array.nulls(); @@ -549,7 +660,7 @@ fn hash_struct_array( // Create hashes for each row that combines the hashes over all the column at that row. let mut values_hashes = vec![0u64; row_len]; - create_hashes(array.columns(), random_state, &mut values_hashes)?; + child_hashing.create_hashes(array.columns(), &mut values_hashes)?; // Separate paths to avoid allocating Vec when there are no nulls if let Some(nulls) = nulls { @@ -571,7 +682,7 @@ fn hash_struct_array( #[cfg(not(feature = "force_hash_collisions"))] fn hash_map_array( array: &MapArray, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], ) -> Result<()> { let nulls = array.nulls(); @@ -590,7 +701,7 @@ fn hash_map_array( .iter() .map(|col| col.slice(first_offset, entries_len)) .collect(); - create_hashes(&sliced_columns, random_state, &mut values_hashes)?; + child_hashing.create_hashes(&sliced_columns, &mut values_hashes)?; // Combine the hashes for entries on each row with each other and previous hash for that row // Adjust indices by first_offset since values_hashes is sliced starting from first_offset @@ -622,7 +733,7 @@ fn hash_map_array( #[cfg(not(feature = "force_hash_collisions"))] fn hash_list_array( array: &GenericListArray, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], ) -> Result<()> where @@ -633,11 +744,10 @@ where let last_offset = array.value_offsets().last().cloned().unwrap_or_default(); let value_bytes_len = (last_offset - first_offset).as_usize(); let mut values_hashes = vec![0u64; value_bytes_len]; - create_hashes( + child_hashing.create_hashes( [array .values() .slice(first_offset.as_usize(), value_bytes_len)], - random_state, &mut values_hashes, )?; @@ -673,7 +783,7 @@ where #[cfg(not(feature = "force_hash_collisions"))] fn hash_list_view_array( array: &GenericListViewArray, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], ) -> Result<()> where @@ -684,7 +794,7 @@ where let sizes = array.value_sizes(); let nulls = array.nulls(); let mut values_hashes = vec![0u64; values.len()]; - create_hashes([values], random_state, &mut values_hashes)?; + child_hashing.create_hashes([values], &mut values_hashes)?; if let Some(nulls) = nulls { for (i, (offset, size)) in offsets.iter().zip(sizes.iter()).enumerate() { if nulls.is_valid(i) { @@ -712,7 +822,7 @@ where #[cfg(not(feature = "force_hash_collisions"))] fn hash_union_array( array: &UnionArray, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], ) -> Result<()> { let DataType::Union(union_fields, _mode) = array.data_type() else { @@ -722,12 +832,12 @@ fn hash_union_array( if array.is_dense() { // Dense union: children only contain values of their type, so they're already compact. // Use the default hashing approach which is efficient for dense unions. - hash_union_array_default(array, union_fields, random_state, hashes_buffer) + hash_union_array_default(array, union_fields, child_hashing, hashes_buffer) } else { // Sparse union: each child has the same length as the union array. // Optimization: only hash the elements that are actually referenced by type_ids, // instead of hashing all K*N elements (where K = num types, N = array length). - hash_sparse_union_array(array, union_fields, random_state, hashes_buffer) + hash_sparse_union_array(array, union_fields, child_hashing, hashes_buffer) } } @@ -744,7 +854,7 @@ fn hash_union_array( fn hash_union_array_default( array: &UnionArray, union_fields: &UnionFields, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], ) -> Result<()> { let mut child_hashes: HashMap> = @@ -754,7 +864,7 @@ fn hash_union_array_default( for (type_id, _field) in union_fields.iter() { let child = array.child(type_id); let mut child_hash_buffer = vec![0; child.len()]; - create_hashes([child], random_state, &mut child_hash_buffer)?; + child_hashing.create_hashes([child], &mut child_hash_buffer)?; child_hashes.insert(type_id, child_hash_buffer); } @@ -785,7 +895,7 @@ fn hash_union_array_default( fn hash_sparse_union_array( array: &UnionArray, union_fields: &UnionFields, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], ) -> Result<()> { use std::collections::HashMap; @@ -796,7 +906,7 @@ fn hash_sparse_union_array( return hash_union_array_default( array, union_fields, - random_state, + child_hashing, hashes_buffer, ); } @@ -824,7 +934,7 @@ fn hash_sparse_union_array( // Hash the filtered array let mut filtered_hashes = vec![0u64; filtered.len()]; - create_hashes([&filtered], random_state, &mut filtered_hashes)?; + child_hashing.create_hashes([&filtered], &mut filtered_hashes)?; // Scatter hashes back to correct positions for (hash, &idx) in filtered_hashes.iter().zip(indices.iter()) { @@ -840,14 +950,14 @@ fn hash_sparse_union_array( #[cfg(not(feature = "force_hash_collisions"))] fn hash_fixed_list_array( array: &FixedSizeListArray, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], ) -> Result<()> { let values = array.values(); let value_length = array.value_length() as usize; let nulls = array.nulls(); let mut values_hashes = vec![0u64; values.len()]; - create_hashes([values], random_state, &mut values_hashes)?; + child_hashing.create_hashes([values], &mut values_hashes)?; if let Some(nulls) = nulls { for i in 0..array.len() { if nulls.is_valid(i) { @@ -875,11 +985,12 @@ fn hash_fixed_list_array( #[cfg(not(feature = "force_hash_collisions"))] fn hash_run_array_inner< R: RunEndIndexType, + C: ChildHashing + ?Sized, const HAS_NULL_VALUES: bool, const REHASH: bool, >( array: &RunArray, - random_state: &RandomState, + child_hashing: &C, hashes_buffer: &mut [u64], ) -> Result<()> { // We find the relevant runs that cover potentially sliced arrays, so we can only hash those @@ -906,11 +1017,8 @@ fn hash_run_array_inner< end_physical_index - start_physical_index, ); let mut values_hashes = vec![0u64; sliced_values.len()]; - create_hashes( - std::slice::from_ref(&sliced_values), - random_state, - &mut values_hashes, - )?; + child_hashing + .create_hashes(std::slice::from_ref(&sliced_values), &mut values_hashes)?; let mut start_in_slice = 0; for (adjusted_physical_index, &absolute_run_end) in run_ends_values @@ -946,24 +1054,26 @@ fn hash_run_array_inner< #[cfg(not(feature = "force_hash_collisions"))] fn hash_run_array( array: &RunArray, - random_state: &RandomState, + child_hashing: &impl ChildHashing, hashes_buffer: &mut [u64], rehash: bool, ) -> Result<()> { let has_null_values = array.values().null_count() != 0; match (has_null_values, rehash) { - (false, false) => { - hash_run_array_inner::(array, random_state, hashes_buffer) - } + (false, false) => hash_run_array_inner::( + array, + child_hashing, + hashes_buffer, + ), (false, true) => { - hash_run_array_inner::(array, random_state, hashes_buffer) + hash_run_array_inner::(array, child_hashing, hashes_buffer) } (true, false) => { - hash_run_array_inner::(array, random_state, hashes_buffer) + hash_run_array_inner::(array, child_hashing, hashes_buffer) } (true, true) => { - hash_run_array_inner::(array, random_state, hashes_buffer) + hash_run_array_inner::(array, child_hashing, hashes_buffer) } } } @@ -973,7 +1083,7 @@ fn hash_run_array( #[cfg(not(feature = "force_hash_collisions"))] fn hash_single_array( array: &dyn Array, - random_state: &RandomState, + random_state: &impl HashState, hashes_buffer: &mut [u64], rehash: bool, ) -> Result<()> { @@ -997,38 +1107,67 @@ fn hash_single_array( } DataType::Struct(_) => { let array = as_struct_array(array)?; - hash_struct_array(array, random_state, hashes_buffer)?; + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_struct_array(array, &child_hashing, hashes_buffer)?; } DataType::List(_) => { let array = as_list_array(array)?; - hash_list_array(array, random_state, hashes_buffer)?; + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_list_array(array, &child_hashing, hashes_buffer)?; } DataType::LargeList(_) => { let array = as_large_list_array(array)?; - hash_list_array(array, random_state, hashes_buffer)?; + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_list_array(array, &child_hashing, hashes_buffer)?; } DataType::ListView(_) => { let array = as_list_view_array(array)?; - hash_list_view_array(array, random_state, hashes_buffer)?; + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_list_view_array(array, &child_hashing, hashes_buffer)?; } DataType::LargeListView(_) => { let array = as_large_list_view_array(array)?; - hash_list_view_array(array, random_state, hashes_buffer)?; + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_list_view_array(array, &child_hashing, hashes_buffer)?; } DataType::Map(_, _) => { let array = as_map_array(array)?; - hash_map_array(array, random_state, hashes_buffer)?; + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_map_array(array, &child_hashing, hashes_buffer)?; } DataType::FixedSizeList(_,_) => { let array = as_fixed_size_list_array(array)?; - hash_fixed_list_array(array, random_state, hashes_buffer)?; + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_fixed_list_array(array, &child_hashing, hashes_buffer)?; } DataType::Union(_, _) => { let array = as_union_array(array)?; - hash_union_array(array, random_state, hashes_buffer)?; + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_union_array(array, &child_hashing, hashes_buffer)?; } DataType::RunEndEncoded(_, _) => downcast_run_array! { - array => hash_run_array(array, random_state, hashes_buffer, rehash)?, + array => { + let child_hashing = HashStateChildHashing { + hash_state: random_state, + }; + hash_run_array(array, &child_hashing, hashes_buffer, rehash)? + }, _ => unreachable!() } _ => { @@ -1046,7 +1185,7 @@ fn hash_single_array( #[cfg(feature = "force_hash_collisions")] fn hash_single_array( _array: &dyn Array, - _random_state: &RandomState, + _random_state: &impl HashState, hashes_buffer: &mut [u64], _rehash: bool, ) -> Result<()> { @@ -1099,7 +1238,7 @@ impl AsDynArray for &ArrayRef { /// `hashes_buffer` should be pre-sized appropriately. pub fn create_hashes<'a, I, T>( arrays: I, - random_state: &RandomState, + random_state: &impl HashState, hashes_buffer: &'a mut [u64], ) -> Result<&'a mut [u64]> where @@ -1114,8 +1253,36 @@ where Ok(hashes_buffer) } +/// Creates hash values for every row using a caller-provided hash builder. +/// +/// The number of rows to hash is determined by `hashes_buffer.len()`. +/// `hashes_buffer` should be pre-sized appropriately. +/// +/// # Hash compatibility +/// +/// Hash values are not guaranteed to be bit-for-bit identical to those from +/// [`create_hashes`], even when `hash_builder` also implements [`HashState`]. +/// The optimized [`HashState`] path seeds the hasher from the previous hash +/// when rehashing some primitive and byte-view values, whereas this function +/// combines independently computed hashes. Use one API consistently if hashes +/// are persisted or exchanged. +pub fn create_hashes_with_hasher<'a, I, T, S>( + arrays: I, + hash_builder: &S, + hashes_buffer: &'a mut [u64], +) -> Result<&'a mut [u64]> +where + I: IntoIterator, + T: AsDynArray, + S: BuildHasher, +{ + build_hasher::create_hashes_with_hasher(arrays, hash_builder, hashes_buffer) +} + #[cfg(test)] mod tests { + #[cfg(not(feature = "force_hash_collisions"))] + use std::hash::{BuildHasherDefault, Hasher}; use std::sync::Arc; use arrow::array::*; @@ -1124,6 +1291,23 @@ mod tests { use super::*; + #[cfg(not(feature = "force_hash_collisions"))] + #[derive(Default)] + struct TestHasher(u64); + + #[cfg(not(feature = "force_hash_collisions"))] + impl Hasher for TestHasher { + fn finish(&self) -> u64 { + self.0 + } + + fn write(&mut self, bytes: &[u8]) { + for byte in bytes { + self.0 = self.0.wrapping_mul(37).wrapping_add(u64::from(*byte)); + } + } + } + #[test] fn create_hashes_for_decimal_array() -> Result<()> { let array = vec![1, 2, 3, 4] @@ -1360,6 +1544,206 @@ mod tests { Ok(()) } + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_create_hashes_with_custom_hasher() { + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 1, 4])); + let hash_builder = BuildHasherDefault::::default(); + + let mut custom_hashes = vec![0; array.len()]; + create_hashes_with_hasher([&array], &hash_builder, &mut custom_hashes).unwrap(); + + let random_state = RandomState::with_seed(0); + let mut default_hashes = vec![0; array.len()]; + create_hashes([&array], &random_state, &mut default_hashes).unwrap(); + + assert_eq!(custom_hashes[0], custom_hashes[2]); + assert_ne!(custom_hashes[0], custom_hashes[1]); + assert_ne!(custom_hashes, default_hashes); + } + + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_create_hashes_with_custom_hasher_normalizes_negative_zero() { + let array: ArrayRef = Arc::new(Float64Array::from(vec![0.0, -0.0])); + let hash_builder = BuildHasherDefault::::default(); + let mut hashes = vec![0; array.len()]; + + create_hashes_with_hasher([&array], &hash_builder, &mut hashes).unwrap(); + + assert_eq!(hashes[0], hashes[1]); + } + + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_create_hashes_dictionary_with_custom_hasher() { + let strings = [Some("foo"), None, Some("bar"), Some("foo"), None]; + let string_array: ArrayRef = + Arc::new(strings.iter().cloned().collect::()); + let dict_array: ArrayRef = Arc::new( + strings + .iter() + .cloned() + .collect::>(), + ); + let hash_builder = BuildHasherDefault::::default(); + + let mut string_hashes = vec![0; strings.len()]; + create_hashes_with_hasher([&string_array], &hash_builder, &mut string_hashes) + .unwrap(); + + let mut dict_hashes = vec![0; strings.len()]; + create_hashes_with_hasher([&dict_array], &hash_builder, &mut dict_hashes) + .unwrap(); + + assert_eq!(string_hashes, dict_hashes); + } + + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_create_hashes_struct_with_custom_hasher() { + let struct_array = StructArray::from(vec![ + ( + Arc::new(Field::new("int", DataType::Int32, false)), + Arc::new(Int32Array::from(vec![1, 2, 1, 3])) as ArrayRef, + ), + ( + Arc::new(Field::new("string", DataType::Utf8, false)), + Arc::new(StringArray::from(vec!["alpha", "beta", "alpha", "alpha"])) + as ArrayRef, + ), + ]); + let hash_builder = BuildHasherDefault::::default(); + + let mut child_hashes = vec![0; struct_array.len()]; + create_hashes_with_hasher( + struct_array.columns(), + &hash_builder, + &mut child_hashes, + ) + .unwrap(); + let expected_hashes = child_hashes + .into_iter() + .map(|hash| combine_hashes(0, hash)) + .collect::>(); + + let array: ArrayRef = Arc::new(struct_array); + let mut actual_hashes = vec![0; array.len()]; + create_hashes_with_hasher([&array], &hash_builder, &mut actual_hashes).unwrap(); + + assert_eq!(actual_hashes, expected_hashes); + assert_eq!(actual_hashes[0], actual_hashes[2]); + assert_ne!(actual_hashes[0], actual_hashes[3]); + } + + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_create_hashes_long_utf8_view_with_custom_hasher() { + let values = vec![ + Some("this string is longer than twelve bytes"), + None, + Some("another string longer than twelve bytes"), + Some("this string is longer than twelve bytes"), + ]; + let view_array = StringViewArray::from(values.clone()); + assert!(!view_array.data_buffers().is_empty()); + let view_array: ArrayRef = Arc::new(view_array); + let hash_builder = BuildHasherDefault::::default(); + + let mut view_hashes = vec![0; view_array.len()]; + create_hashes_with_hasher([&view_array], &hash_builder, &mut view_hashes) + .unwrap(); + let expected_hashes = values + .iter() + .map(|value| { + value + .map(|value| hash_builder.hash_one(value.as_bytes())) + .unwrap_or_default() + }) + .collect::>(); + assert_eq!(view_hashes, expected_hashes); + + let prefix_array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 1])); + let mut expected_hashes = vec![0; prefix_array.len()]; + create_hashes_with_hasher([&prefix_array], &hash_builder, &mut expected_hashes) + .unwrap(); + for (hash, value) in expected_hashes.iter_mut().zip(&values) { + if let Some(value) = value { + *hash = combine_hashes(hash_builder.hash_one(value.as_bytes()), *hash); + } + } + + let mut view_hashes = vec![0; view_array.len()]; + create_hashes_with_hasher( + [&prefix_array, &view_array], + &hash_builder, + &mut view_hashes, + ) + .unwrap(); + assert_eq!(view_hashes, expected_hashes); + } + + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_single_column_leaf_hashes_match_with_same_hasher() { + let arrays: Vec = vec![ + Arc::new(Int32Array::from(vec![Some(1), None, Some(-1)])), + Arc::new(Float64Array::from(vec![Some(0.0), Some(-0.0), None])), + Arc::new(StringArray::from(vec![Some("foo"), None, Some("bar")])), + Arc::new(BinaryArray::from(vec![ + Some(&b"short"[..]), + None, + Some(&b"longer than twelve bytes"[..]), + ])), + Arc::new(StringViewArray::from(vec![ + Some("short"), + None, + Some("longer than twelve bytes"), + ])), + ]; + let random_state = RandomState::with_seed(0); + + for array in arrays { + let mut default_hashes = vec![0; array.len()]; + create_hashes([&array], &random_state, &mut default_hashes).unwrap(); + + let mut custom_hashes = vec![0; array.len()]; + create_hashes_with_hasher([&array], &random_state, &mut custom_hashes) + .unwrap(); + + assert_eq!( + custom_hashes, + default_hashes, + "single-column parity failed for {}", + array.data_type() + ); + } + } + + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_with_hashes_with_custom_hasher() { + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let hash_builder = BuildHasherDefault::::default(); + + let mut expected_hashes = vec![0; int_array.len()]; + create_hashes_with_hasher( + [&int_array, &str_array], + &hash_builder, + &mut expected_hashes, + ) + .unwrap(); + + let actual_hashes = + with_hashes_with_hasher([&int_array, &str_array], &hash_builder, |hashes| { + Ok(hashes.to_vec()) + }) + .unwrap(); + + assert_eq!(actual_hashes, expected_hashes); + } + #[test] // Tests actual values of hashes, which are different if forcing collisions #[cfg(not(feature = "force_hash_collisions"))] @@ -1808,6 +2192,31 @@ mod tests { assert_eq!(hashes1, hashes2); } + #[test] + #[cfg(not(feature = "force_hash_collisions"))] + fn test_create_hashes_with_quality_hash_state() { + let int_array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + let str_array: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); + let quality_state = foldhash::quality::FixedState::with_seed(0); + + let mut one_col_hashes = vec![0; int_array.len()]; + create_hashes([&int_array], &quality_state, &mut one_col_hashes).unwrap(); + let expected_hashes: Vec<_> = [1i32, 2, 3, 4] + .iter() + .map(|value| quality_state.hash_one(value)) + .collect(); + assert_eq!(one_col_hashes, expected_hashes); + + let mut two_col_hashes = vec![0; int_array.len()]; + create_hashes( + [&int_array, &str_array], + &quality_state, + &mut two_col_hashes, + ) + .unwrap(); + assert_ne!(two_col_hashes, one_col_hashes); + } + #[test] fn test_with_hashes() { let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); diff --git a/datafusion/common/src/hash_utils/build_hasher.rs b/datafusion/common/src/hash_utils/build_hasher.rs new file mode 100644 index 0000000000000..12258beb11403 --- /dev/null +++ b/datafusion/common/src/hash_utils/build_hasher.rs @@ -0,0 +1,494 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use super::{AsDynArray, HASH_BUFFER, MAX_BUFFER_SIZE}; +#[cfg(not(feature = "force_hash_collisions"))] +use super::{ + ChildHashing, combine_hashes, hash_dictionary_with_child_hashing, + hash_fixed_list_array, hash_list_array, hash_list_view_array, hash_map_array, + hash_run_array, hash_struct_array, hash_union_array, +}; +#[cfg(not(feature = "force_hash_collisions"))] +use crate::cast::{ + as_binary_view_array, as_boolean_array, as_fixed_size_list_array, + as_generic_binary_array, as_large_list_array, as_large_list_view_array, + as_list_array, as_list_view_array, as_map_array, as_string_array, + as_string_view_array, as_struct_array, as_union_array, +}; +use crate::error::Result; +use crate::error::{_internal_datafusion_err, _internal_err}; +#[cfg(feature = "force_hash_collisions")] +use arrow::array::Array; +#[cfg(not(feature = "force_hash_collisions"))] +use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; +#[cfg(not(feature = "force_hash_collisions"))] +use arrow::array::*; +#[cfg(not(feature = "force_hash_collisions"))] +use arrow::datatypes::*; +#[cfg(not(feature = "force_hash_collisions"))] +use arrow::{downcast_dictionary_array, downcast_primitive_array}; +use std::hash::BuildHasher; + +pub(super) fn with_hashes_with_hasher( + arrays: I, + hash_builder: &S, + callback: F, +) -> Result +where + I: IntoIterator, + T: AsDynArray, + F: FnOnce(&[u64]) -> Result, + S: BuildHasher, +{ + let mut iter = arrays.into_iter().peekable(); + + let required_size = match iter.peek() { + Some(arr) => arr.as_dyn_array().len(), + None => { + return _internal_err!("with_hashes_with_hasher requires at least one array"); + } + }; + + HASH_BUFFER.try_with(|cell| { + let mut buffer = cell.try_borrow_mut().map_err(|_| { + _internal_datafusion_err!( + "with_hashes_with_hasher cannot be called reentrantly on the same thread" + ) + })?; + + buffer.clear(); + buffer.resize(required_size, 0); + + create_hashes_with_hasher_impl(iter, hash_builder, &mut buffer[..required_size])?; + + let result = callback(&buffer[..required_size])?; + + if buffer.capacity() > MAX_BUFFER_SIZE { + buffer.truncate(MAX_BUFFER_SIZE); + buffer.shrink_to_fit(); + } + + Ok(result) + }).map_err(|_| { + _internal_datafusion_err!( + "with_hashes_with_hasher cannot access thread-local storage during or after thread destruction" + ) + })? +} + +pub(super) fn create_hashes_with_hasher<'a, I, T, S>( + arrays: I, + hash_builder: &S, + hashes_buffer: &'a mut [u64], +) -> Result<&'a mut [u64]> +where + I: IntoIterator, + T: AsDynArray, + S: BuildHasher, +{ + create_hashes_with_hasher_impl(arrays, hash_builder, hashes_buffer) +} + +fn create_hashes_with_hasher_impl<'a, I, T, S>( + arrays: I, + hash_builder: &S, + hashes_buffer: &'a mut [u64], +) -> Result<&'a mut [u64]> +where + I: IntoIterator, + T: AsDynArray, + S: BuildHasher, +{ + for (i, array) in arrays.into_iter().enumerate() { + let rehash = i >= 1; + hash_single_array_with_hasher( + array.as_dyn_array(), + hash_builder, + hashes_buffer, + rehash, + )?; + } + Ok(hashes_buffer) +} + +#[cfg(not(feature = "force_hash_collisions"))] +struct BuildHasherChildHashing<'a, S> { + hash_builder: &'a S, +} + +#[cfg(not(feature = "force_hash_collisions"))] +impl ChildHashing for BuildHasherChildHashing<'_, S> { + fn create_hashes(&self, arrays: I, hashes_buffer: &mut [u64]) -> Result<()> + where + I: IntoIterator, + T: AsDynArray, + { + create_hashes_with_hasher_impl(arrays, self.hash_builder, hashes_buffer) + .map(|_| ()) + } +} + +#[cfg(not(feature = "force_hash_collisions"))] +trait BuildHasherHashValue { + fn hash_one_with_hasher(&self, state: &S) -> u64; +} + +#[cfg(not(feature = "force_hash_collisions"))] +impl BuildHasherHashValue for &T { + fn hash_one_with_hasher(&self, state: &S) -> u64 { + T::hash_one_with_hasher(self, state) + } +} + +macro_rules! build_hasher_hash_value { + ($($t:ty),+) => { + $(#[cfg(not(feature = "force_hash_collisions"))] + impl BuildHasherHashValue for $t { + fn hash_one_with_hasher(&self, state: &S) -> u64 { + state.hash_one(self) + } + })+ + }; +} +build_hasher_hash_value!(i8, i16, i32, i64, i128, i256, u8, u16, u32, u64, u128); +build_hasher_hash_value!(bool, str, [u8], IntervalDayTime, IntervalMonthDayNano); + +macro_rules! build_hasher_hash_float_value { + ($(($t:ty, $i:ty)),+) => { + $(#[cfg(not(feature = "force_hash_collisions"))] + impl BuildHasherHashValue for $t { + fn hash_one_with_hasher(&self, state: &S) -> u64 { + let bits = <$i>::from_ne_bytes(self.to_ne_bytes()); + let bits = if bits << 1 == 0 { 0 } else { bits }; + state.hash_one(bits) + } + })+ + }; +} +build_hasher_hash_float_value!((half::f16, u16), (f32, u32), (f64, u64)); + +#[cfg(not(feature = "force_hash_collisions"))] +fn hash_null_with_hasher( + hash_builder: &S, + hashes_buffer: &mut [u64], + multi_col: bool, +) { + if hashes_buffer.is_empty() { + return; + } + + let null_hash = hash_builder.hash_one(1); + if multi_col { + hashes_buffer.iter_mut().for_each(|hash| { + *hash = combine_hashes(null_hash, *hash); + }) + } else { + hashes_buffer.fill(null_hash); + } +} + +#[cfg(not(feature = "force_hash_collisions"))] +fn hash_array_primitive_with_hasher( + array: &PrimitiveArray, + hash_builder: &S, + hashes_buffer: &mut [u64], + rehash: bool, +) where + T: ArrowPrimitiveType, + S: BuildHasher, +{ + assert_eq!( + hashes_buffer.len(), + array.len(), + "hashes_buffer and array should be of equal length" + ); + + if array.null_count() == 0 { + if rehash { + for (hash, &value) in hashes_buffer.iter_mut().zip(array.values().iter()) { + *hash = combine_hashes(value.hash_one_with_hasher(hash_builder), *hash); + } + } else { + for (hash, &value) in hashes_buffer.iter_mut().zip(array.values().iter()) { + *hash = value.hash_one_with_hasher(hash_builder); + } + } + } else if rehash { + for i in array.nulls().unwrap().valid_indices() { + let value = unsafe { array.value_unchecked(i) }; + hashes_buffer[i] = combine_hashes( + value.hash_one_with_hasher(hash_builder), + hashes_buffer[i], + ); + } + } else { + for i in array.nulls().unwrap().valid_indices() { + let value = unsafe { array.value_unchecked(i) }; + hashes_buffer[i] = value.hash_one_with_hasher(hash_builder); + } + } +} + +#[cfg(not(feature = "force_hash_collisions"))] +fn hash_array_with_hasher( + array: &T, + hash_builder: &S, + hashes_buffer: &mut [u64], + rehash: bool, +) where + T: ArrayAccessor, + T::Item: BuildHasherHashValue, + S: BuildHasher, +{ + assert_eq!( + hashes_buffer.len(), + array.len(), + "hashes_buffer and array should be of equal length" + ); + + if array.null_count() == 0 { + if rehash { + for (i, hash) in hashes_buffer.iter_mut().enumerate() { + let value = unsafe { array.value_unchecked(i) }; + *hash = combine_hashes(value.hash_one_with_hasher(hash_builder), *hash); + } + } else { + for (i, hash) in hashes_buffer.iter_mut().enumerate() { + let value = unsafe { array.value_unchecked(i) }; + *hash = value.hash_one_with_hasher(hash_builder); + } + } + } else if rehash { + for i in array.nulls().unwrap().valid_indices() { + let value = unsafe { array.value_unchecked(i) }; + hashes_buffer[i] = combine_hashes( + value.hash_one_with_hasher(hash_builder), + hashes_buffer[i], + ); + } + } else { + for i in array.nulls().unwrap().valid_indices() { + let value = unsafe { array.value_unchecked(i) }; + hashes_buffer[i] = value.hash_one_with_hasher(hash_builder); + } + } +} + +#[cfg(not(feature = "force_hash_collisions"))] +#[inline(never)] +fn hash_string_view_array_inner_with_hasher< + T: ByteViewType, + S: BuildHasher, + const HAS_NULLS: bool, + const HAS_BUFFERS: bool, + const REHASH: bool, +>( + array: &GenericByteViewArray, + hash_builder: &S, + hashes_buffer: &mut [u64], +) { + assert_eq!( + hashes_buffer.len(), + array.len(), + "hashes_buffer and array should be of equal length" + ); + + let buffers = array.data_buffers(); + let view_bytes = |view_len: u32, view: u128| { + let view = ByteView::from(view); + let offset = view.offset as usize; + unsafe { + let data = buffers.get_unchecked(view.buffer_index as usize); + data.get_unchecked(offset..offset + view_len as usize) + } + }; + + let hashes_and_views = hashes_buffer.iter_mut().zip(array.views().iter()); + for (i, (hash, &v)) in hashes_and_views.enumerate() { + if HAS_NULLS && array.is_null(i) { + continue; + } + let view_len = v as u32; + if !HAS_BUFFERS || view_len <= 12 { + if REHASH { + *hash = combine_hashes(v.hash_one_with_hasher(hash_builder), *hash); + } else { + *hash = v.hash_one_with_hasher(hash_builder); + } + continue; + } + let value = view_bytes(view_len, v); + if REHASH { + *hash = combine_hashes(value.hash_one_with_hasher(hash_builder), *hash); + } else { + *hash = value.hash_one_with_hasher(hash_builder); + } + } +} + +#[cfg(not(feature = "force_hash_collisions"))] +fn hash_generic_byte_view_array_with_hasher( + array: &GenericByteViewArray, + hash_builder: &S, + hashes_buffer: &mut [u64], + rehash: bool, +) { + match ( + array.null_count() != 0, + !array.data_buffers().is_empty(), + rehash, + ) { + (false, false, false) => { + for (hash, &view) in hashes_buffer.iter_mut().zip(array.views().iter()) { + *hash = view.hash_one_with_hasher(hash_builder); + } + } + (false, false, true) => { + for (hash, &view) in hashes_buffer.iter_mut().zip(array.views().iter()) { + *hash = combine_hashes(view.hash_one_with_hasher(hash_builder), *hash); + } + } + (false, true, false) => { + hash_string_view_array_inner_with_hasher::( + array, + hash_builder, + hashes_buffer, + ) + } + (false, true, true) => { + hash_string_view_array_inner_with_hasher::( + array, + hash_builder, + hashes_buffer, + ) + } + (true, false, false) => { + hash_string_view_array_inner_with_hasher::( + array, + hash_builder, + hashes_buffer, + ) + } + (true, false, true) => { + hash_string_view_array_inner_with_hasher::( + array, + hash_builder, + hashes_buffer, + ) + } + (true, true, false) => { + hash_string_view_array_inner_with_hasher::( + array, + hash_builder, + hashes_buffer, + ) + } + (true, true, true) => { + hash_string_view_array_inner_with_hasher::( + array, + hash_builder, + hashes_buffer, + ) + } + } +} + +#[cfg(not(feature = "force_hash_collisions"))] +fn hash_single_array_with_hasher( + array: &dyn Array, + hash_builder: &S, + hashes_buffer: &mut [u64], + rehash: bool, +) -> Result<()> { + let child_hashing = BuildHasherChildHashing { hash_builder }; + + downcast_primitive_array! { + array => hash_array_primitive_with_hasher(array, hash_builder, hashes_buffer, rehash), + DataType::Null => hash_null_with_hasher(hash_builder, hashes_buffer, rehash), + DataType::Boolean => hash_array_with_hasher(&as_boolean_array(array)?, hash_builder, hashes_buffer, rehash), + DataType::Utf8 => hash_array_with_hasher(&as_string_array(array)?, hash_builder, hashes_buffer, rehash), + DataType::Utf8View => hash_generic_byte_view_array_with_hasher(as_string_view_array(array)?, hash_builder, hashes_buffer, rehash), + DataType::LargeUtf8 => hash_array_with_hasher(&as_largestring_array(array), hash_builder, hashes_buffer, rehash), + DataType::Binary => hash_array_with_hasher(&as_generic_binary_array::(array)?, hash_builder, hashes_buffer, rehash), + DataType::BinaryView => hash_generic_byte_view_array_with_hasher(as_binary_view_array(array)?, hash_builder, hashes_buffer, rehash), + DataType::LargeBinary => hash_array_with_hasher(&as_generic_binary_array::(array)?, hash_builder, hashes_buffer, rehash), + DataType::FixedSizeBinary(_) => { + let array: &FixedSizeBinaryArray = array.as_any().downcast_ref().unwrap(); + hash_array_with_hasher(&array, hash_builder, hashes_buffer, rehash) + } + DataType::Dictionary(_, _) => downcast_dictionary_array! { + array => hash_dictionary_with_child_hashing(array, &child_hashing, hashes_buffer, rehash)?, + _ => unreachable!() + } + DataType::Struct(_) => { + let array = as_struct_array(array)?; + hash_struct_array(array, &child_hashing, hashes_buffer)?; + } + DataType::List(_) => { + let array = as_list_array(array)?; + hash_list_array(array, &child_hashing, hashes_buffer)?; + } + DataType::LargeList(_) => { + let array = as_large_list_array(array)?; + hash_list_array(array, &child_hashing, hashes_buffer)?; + } + DataType::ListView(_) => { + let array = as_list_view_array(array)?; + hash_list_view_array(array, &child_hashing, hashes_buffer)?; + } + DataType::LargeListView(_) => { + let array = as_large_list_view_array(array)?; + hash_list_view_array(array, &child_hashing, hashes_buffer)?; + } + DataType::Map(_, _) => { + let array = as_map_array(array)?; + hash_map_array(array, &child_hashing, hashes_buffer)?; + } + DataType::FixedSizeList(_,_) => { + let array = as_fixed_size_list_array(array)?; + hash_fixed_list_array(array, &child_hashing, hashes_buffer)?; + } + DataType::Union(_, _) => { + let array = as_union_array(array)?; + hash_union_array(array, &child_hashing, hashes_buffer)?; + } + DataType::RunEndEncoded(_, _) => downcast_run_array! { + array => hash_run_array(array, &child_hashing, hashes_buffer, rehash)?, + _ => unreachable!() + } + _ => { + return _internal_err!( + "Unsupported data type in hasher: {}", + array.data_type() + ); + } + } + Ok(()) +} + +#[cfg(feature = "force_hash_collisions")] +fn hash_single_array_with_hasher( + _array: &dyn Array, + _hash_builder: &S, + hashes_buffer: &mut [u64], + _rehash: bool, +) -> Result<()> { + for hash in hashes_buffer.iter_mut() { + *hash = 0; + } + Ok(()) +} diff --git a/datafusion/common/src/heap_size.rs b/datafusion/common/src/heap_size.rs index edb64709d5aa4..037f807fce9d8 100644 --- a/datafusion/common/src/heap_size.rs +++ b/datafusion/common/src/heap_size.rs @@ -15,6 +15,30 @@ // specific language governing permissions and limitations // under the License. +//! Estimating the heap-allocated memory owned by a value. +//! +//! The [`DFHeapSize`] trait reports the number of bytes a value owns on the +//! heap, **excluding** the stack size of the value itself. +//! +//! Implementations need to use [`DFHeapSizeCtx`] that is pushed through every +//! nested call. The context records which allocations have already been measured +//! so they are only counted once. +//! +//! # Example +//! +//! ``` +//! use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; +//! use std::sync::Arc; +//! +//! let shared: Arc = Arc::new("hello".to_string()); +//! let alias = Arc::clone(&shared); +//! +//! let mut ctx = DFHeapSizeCtx::default(); +//! // The shared allocation is counted once even when reached twice. +//! let total = shared.heap_size(&mut ctx) + alias.heap_size(&mut ctx); +//! assert_eq!(total, shared.heap_size(&mut DFHeapSizeCtx::default())); +//! ``` + use crate::stats::Precision; use crate::{ColumnStatistics, ScalarValue, Statistics, TableReference}; use arrow::array::{ @@ -32,12 +56,15 @@ use std::collections::HashMap; use std::fmt::Debug; use std::sync::Arc; -/// This is a temporary solution until and -/// are resolved. -/// Trait for calculating the size of various containers +/// Trait for computing how many bytes a value has allocated on the heap. +/// +/// Implementations need to use [`DFHeapSizeCtx`] that is pushed through every +/// nested call. The context records which allocations have already been measured +/// so they are only counted once. +/// pub trait DFHeapSize { - /// Return the size of any bytes allocated on the heap by this object, - /// including heap memory in those structures + /// Return the number of bytes this value has allocated on the heap, + /// including heap memory owned transitively by nested values. /// /// Note that the size of the type itself is not included in the result -- /// instead, that size is added by the caller (e.g. container). @@ -49,6 +76,12 @@ pub struct DFHeapSizeCtx { seen: HashSet, } +impl DFHeapSizeCtx { + fn count_allocation_once(&mut self, ptr: usize) -> bool { + self.seen.insert(ptr) + } +} + impl DFHeapSize for Statistics { fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { self.num_rows.heap_size(ctx) @@ -254,11 +287,21 @@ impl DFHeapSize for HashMap { } } +fn arc_ptr(arc: &Arc) -> usize { + Arc::as_ptr(arc) as usize +} + +/// For unsized types, `Arc::as_ptr` returns the data address + metadata - we only need the thin address +/// Casting through `*const i32` gets us the thin pointer +fn arc_unsized_ptr(arc: &Arc) -> usize { + Arc::as_ptr(arc) as *const i32 as usize +} + impl DFHeapSize for Arc { fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { - let ptr = Arc::as_ptr(self) as usize; + let ptr = arc_ptr(self); - if !ctx.seen.insert(ptr) { + if !ctx.count_allocation_once(ptr) { return 0; } @@ -269,9 +312,9 @@ impl DFHeapSize for Arc { impl DFHeapSize for Arc { fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { - let ptr = Arc::as_ptr(self) as *const i32 as usize; + let ptr = arc_unsized_ptr(self); - if !ctx.seen.insert(ptr) { + if !ctx.count_allocation_once(ptr) { return 0; } @@ -282,9 +325,9 @@ impl DFHeapSize for Arc { impl DFHeapSize for Arc { fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { - let ptr = Arc::as_ptr(self) as *const i32 as usize; + let ptr = arc_unsized_ptr(self); - if !ctx.seen.insert(ptr) { + if !ctx.count_allocation_once(ptr) { return 0; } @@ -299,47 +342,6 @@ impl DFHeapSize for Fields { } } -impl DFHeapSize for StructArray { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - self.get_array_memory_size() - } -} - -impl DFHeapSize for LargeListArray { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - self.get_array_memory_size() - } -} - -impl DFHeapSize for LargeListViewArray { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - self.get_array_memory_size() - } -} - -impl DFHeapSize for ListArray { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - self.get_array_memory_size() - } -} - -impl DFHeapSize for ListViewArray { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - self.get_array_memory_size() - } -} - -impl DFHeapSize for FixedSizeListArray { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - self.get_array_memory_size() - } -} -impl DFHeapSize for MapArray { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - self.get_array_memory_size() - } -} - impl DFHeapSize for Box { fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { size_of::() + self.as_ref().heap_size(ctx) @@ -362,6 +364,17 @@ where } } +impl DFHeapSize for (A, B, C) +where + A: DFHeapSize, + B: DFHeapSize, + C: DFHeapSize, +{ + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.0.heap_size(ctx) + self.1.heap_size(ctx) + self.2.heap_size(ctx) + } +} + impl DFHeapSize for String { fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { self.capacity() @@ -370,6 +383,7 @@ impl DFHeapSize for String { impl DFHeapSize for str { fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + // Internal accounting helper for owners like Arc self.len() } } @@ -382,24 +396,6 @@ impl DFHeapSize for UnionFields { } } -impl DFHeapSize for UnionMode { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for TimeUnit { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for IntervalUnit { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - impl DFHeapSize for Field { fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { self.name().heap_size(ctx) @@ -424,103 +420,72 @@ impl DFHeapSize for IntervalDayTime { } } -impl DFHeapSize for DateTime { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for bool { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} -impl DFHeapSize for u8 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for u16 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for u32 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for u64 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for i8 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for i16 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for i32 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} -impl DFHeapSize for i64 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for i128 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for i256 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for f16 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} - -impl DFHeapSize for f32 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} -impl DFHeapSize for f64 { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } +/// Implement [`DFHeapSize`] for types that own no heap allocations. +macro_rules! impl_zero_heap_size { + ($($t:ty),+ $(,)?) => { + $( + impl DFHeapSize for $t { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + 0 // no heap allocations + } + } + )+ + }; +} + +impl_zero_heap_size!( + bool, + u8, + u16, + u32, + u64, + usize, + i8, + i16, + i32, + i64, + i128, + i256, + f16, + f32, + f64, + UnionMode, + TimeUnit, + IntervalUnit, + DateTime, +); + +/// Implement [`DFHeapSize`] for Arrow arrays types. +macro_rules! impl_array_heap_size { + ($($t:ty),+ $(,)?) => { + $( + impl DFHeapSize for $t { + fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { + self.get_array_memory_size() + } + } + )+ + }; } -impl DFHeapSize for usize { - fn heap_size(&self, _: &mut DFHeapSizeCtx) -> usize { - 0 // no heap allocations - } -} +impl_array_heap_size!( + StructArray, + LargeListArray, + LargeListViewArray, + ListArray, + ListViewArray, + FixedSizeListArray, + MapArray, +); #[cfg(test)] mod tests { use super::*; + fn size(v: &T) -> usize { + v.heap_size(&mut DFHeapSizeCtx::default()) + } + #[test] fn test_heap_size_arc_avoid_double_accounting() { let a1 = Arc::new(vec![1, 2, 3]); @@ -558,4 +523,243 @@ mod tests { assert_eq!(heap_size, heap_size_with_clones); } + + #[test] + fn test_arc_dyn() { + let a1: Arc = Arc::new(String::from("hello")); + let baseline = size(&a1); + + let a2 = Arc::clone(&a1); + let mut ctx = DFHeapSizeCtx::default(); + let with_clones = a1.heap_size(&mut ctx) + a2.heap_size(&mut ctx); + assert_eq!(baseline, with_clones); + } + + #[test] + fn test_primitives() { + assert_eq!(size(&true), 0); + assert_eq!(size(&0u8), 0); + assert_eq!(size(&0u16), 0); + assert_eq!(size(&0u32), 0); + assert_eq!(size(&0u64), 0); + assert_eq!(size(&0usize), 0); + assert_eq!(size(&0i8), 0); + assert_eq!(size(&0i16), 0); + assert_eq!(size(&0i32), 0); + assert_eq!(size(&0i64), 0); + assert_eq!(size(&0i128), 0); + assert_eq!(size(&i256::ZERO), 0); + assert_eq!(size(&0f32), 0); + assert_eq!(size(&0f64), 0); + assert_eq!(size(&f16::from_f32(0.0)), 0); + } + + #[test] + fn test_heap_size_union_mode() { + assert_eq!(size(&UnionMode::Sparse), 0); + assert_eq!(size(&UnionMode::Dense), 0); + } + + #[test] + fn test_heap_size_time_units() { + assert_eq!(size(&TimeUnit::Second), 0); + assert_eq!(size(&IntervalUnit::YearMonth), 0); + assert_eq!(size(&DateTime::::UNIX_EPOCH), 0); + assert_eq!(size(&Utc::now()), 0); + } + + #[test] + fn test_string() { + let mut s = String::with_capacity(32); + s.push_str("hello"); + assert_eq!(size(&s), 32); + + let empty = String::new(); + assert_eq!(size(&empty), 0); + } + + #[test] + fn test_owned_str() { + let a: Arc = Arc::from("Hello"); + assert!(size(&a) > 0); + } + + #[test] + fn test_option() { + let some: Option = Some(String::from("hi")); + assert_eq!(size(&some), some.as_ref().unwrap().capacity()); + + let none: Option = None; + assert_eq!(size(&none), 0); + } + + #[test] + fn test_vec() { + let v: Vec = vec![1, 2, 3]; + assert!(size(&v) > 0); + + let strings = vec![String::from("ab"), String::from("cdef")]; + assert!(size(&strings) > 0); + + let empty: Vec = Vec::new(); + assert_eq!(size(&empty), 0); + } + + #[test] + fn test_box() { + let b: Box = Box::new(42); + assert!(size(&b) > 0); + + let b: Box = Box::new(String::from("hello")); + assert!(size(&b) > 0); + } + + #[test] + fn test_tuple() { + let zero = (1i32, 2i64); + assert_eq!(size(&zero), 0); + + let t = (String::from("hello"), String::from("world")); + assert!(size(&t) > 0); + } + + #[test] + fn test_hashmap() { + let m: HashMap = HashMap::new(); + assert_eq!(size(&m), 0); + + let mut m: HashMap = HashMap::new(); + m.insert("key".into(), "value".into()); + + assert!(size(&m) > 0); + } + + #[test] + fn test_precision() { + let exact: Precision = Precision::Exact(42); + assert_eq!(size(&exact), 0); + + let inexact: Precision = Precision::Inexact(99); + assert_eq!(size(&inexact), 0); + + let absent: Precision = Precision::Absent; + assert_eq!(size(&absent), 0); + } + + #[test] + fn test_scalar_values() { + assert_eq!(size(&ScalarValue::Null), 0); + assert_eq!(size(&ScalarValue::Int32(Some(42))), 0); + assert_eq!(size(&ScalarValue::Boolean(Some(true))), 0); + assert_eq!(size(&ScalarValue::Float64(None)), 0); + + let sv = ScalarValue::Utf8(Some(String::from("hello"))); + assert_eq!(size(&sv), "hello".len()); + + let sv = ScalarValue::Utf8(None); + assert_eq!(size(&sv), 0); + } + + #[test] + fn test_data_type_primitives() { + assert_eq!(size(&DataType::Int32), 0); + assert_eq!(size(&DataType::Utf8), 0); + assert_eq!(size(&DataType::Boolean), 0); + assert_eq!(size(&DataType::Null), 0); + } + + #[test] + fn test_data_type_with_field() { + let list = DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + assert!(size(&list) > 0); + } + + #[test] + fn test_table_references() { + let tr = TableReference::bare("users"); + // Arc overhead (two usize counts) plus the bytes of "users". + assert!(size(&tr) > 0); + let tr = TableReference::full("cat", "schema", "users"); + assert!(size(&tr) > 0); + } + + #[test] + fn test_column_statistics() { + let mut col = ColumnStatistics::new_unknown(); + col.max_value = Precision::Exact(ScalarValue::Utf8(Some("hello".into()))); + col.min_value = Precision::Exact(ScalarValue::Utf8(Some("ab".into()))); + assert_eq!(size(&col), "hello".len() + "ab".len()); + + let mut col = ColumnStatistics::new_unknown(); + col.max_value = Precision::Exact(ScalarValue::Utf8(Some("hello".into()))); + let stats = Statistics { + num_rows: Precision::Exact(10), + total_byte_size: Precision::Absent, + column_statistics: vec![col], + }; + assert!(size(&stats) > 0); + } + + #[test] + fn test_field() { + let field = Field::new("temperature", DataType::Float64, true); + assert!(size(&field) > 0); + } + + #[test] + fn test_list_array() { + use arrow::array::types::Int32Type; + + let array = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![Some(4)]), + ]); + assert_eq!(size(&array), array.get_array_memory_size()); + assert!(size(&array) > 0); + + let large = + LargeListArray::from_iter_primitive::(vec![Some(vec![ + Some(1), + Some(2), + ])]); + assert_eq!(size(&large), large.get_array_memory_size()); + assert!(size(&large) > 0); + } + + #[test] + fn test_struct_array() { + use arrow::array::Int32Array; + + let array = StructArray::from(vec![( + Arc::new(Field::new("a", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![1, 2, 3])) as _, + )]); + assert_eq!(size(&array), array.get_array_memory_size()); + assert!(size(&array) > 0); + } + + #[test] + fn test_fixed_size_list_array() { + use arrow::array::Int32Array; + + let values = Arc::new(Int32Array::from(vec![1, 2, 3, 4])); + let field = Arc::new(Field::new("item", DataType::Int32, true)); + let array = FixedSizeListArray::new(field, 2, values, None); + assert_eq!(size(&array), array.get_array_memory_size()); + assert!(size(&array) > 0); + } + + #[test] + fn test_map_array() { + use arrow::array::{Int32Builder, MapBuilder, StringBuilder}; + + let mut builder = + MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + builder.keys().append_value("key"); + builder.values().append_value(1); + builder.append(true).unwrap(); + let array = builder.finish(); + assert_eq!(size(&array), array.get_array_memory_size()); + assert!(size(&array) > 0); + } } diff --git a/datafusion/common/src/join_type.rs b/datafusion/common/src/join_type.rs index d517844db48b4..c77a1475ed227 100644 --- a/datafusion/common/src/join_type.rs +++ b/datafusion/common/src/join_type.rs @@ -156,6 +156,24 @@ impl JoinType { | JoinType::RightSemi ) } + + /// Returns true when an empty build-side map necessarily produces an empty + /// result for this join type, even if the build side still contains rows. + /// + /// Every output row of these join types requires a matching build row, so + /// when the map has no matchable keys the result is empty regardless of the + /// probe side. Note this is a subset of + /// [`Self::empty_build_side_produces_empty_result`]: an empty build side + /// yields an empty map, but the map can also be empty when every build row + /// has a NULL join key under [`NullEquality::NullEqualsNothing`]. + /// + /// [`NullEquality::NullEqualsNothing`]: crate::NullEquality::NullEqualsNothing + pub fn empty_map_produces_empty_result(self) -> bool { + matches!( + self, + JoinType::Inner | JoinType::LeftSemi | JoinType::RightSemi + ) + } } impl Display for JoinType { diff --git a/datafusion/common/src/lib.rs b/datafusion/common/src/lib.rs index e865c548bb554..2eebfe4963057 100644 --- a/datafusion/common/src/lib.rs +++ b/datafusion/common/src/lib.rs @@ -30,6 +30,7 @@ mod dfschema; mod functional_dependencies; mod join_type; mod param_value; +mod partitioning; mod schema_reference; mod table_reference; mod unnest; @@ -92,12 +93,13 @@ pub use join_type::{JoinConstraint, JoinSide, JoinType}; pub use nested_struct::cast_column; pub use null_equality::NullEquality; pub use param_value::ParamValues; +pub use partitioning::{SplitPoint, validate_range_split_points}; pub use scalar::{ScalarType, ScalarValue}; pub use schema_reference::SchemaReference; pub use spans::{Location, Span, Spans}; pub use stats::{ColumnStatistics, Statistics}; pub use table_reference::{ResolvedTableReference, TableReference}; -pub use unnest::{RecursionUnnestOption, UnnestOptions}; +pub use unnest::{NullHandling, RecursionUnnestOption, UnnestOptions}; pub use utils::project_schema; // These are hidden from docs purely to avoid polluting the public view of what this crate exports. diff --git a/datafusion/common/src/nested_struct.rs b/datafusion/common/src/nested_struct.rs index cdd6215d08e2f..e915b91b911cc 100644 --- a/datafusion/common/src/nested_struct.rs +++ b/datafusion/common/src/nested_struct.rs @@ -18,9 +18,10 @@ use crate::error::{_plan_err, Result}; use arrow::{ array::{ - Array, ArrayRef, DictionaryArray, GenericListArray, GenericListViewArray, - StructArray, downcast_integer, new_null_array, + Array, ArrayRef, AsArray, DictionaryArray, FixedSizeListArray, GenericListArray, + GenericListViewArray, StructArray, downcast_integer, make_array, new_null_array, }, + buffer::NullBuffer, compute::{CastOptions, can_cast_types, cast_with_options}, datatypes::{DataType, DataType::Struct, Field, FieldRef}, }; @@ -58,9 +59,7 @@ fn cast_struct_column( target_fields: &[Arc], cast_options: &CastOptions, ) -> Result { - if source_col.data_type() == &DataType::Null - || (!source_col.is_empty() && source_col.null_count() == source_col.len()) - { + if source_col.data_type() == &DataType::Null { return Ok(new_null_array( &Struct(target_fields.to_vec().into()), source_col.len(), @@ -70,6 +69,14 @@ fn cast_struct_column( if let Some(source_struct) = source_col.as_any().downcast_ref::() { let source_fields = source_struct.fields(); validate_struct_compatibility(source_fields, target_fields)?; + + if !source_col.is_empty() && source_col.null_count() == source_col.len() { + return Ok(new_null_array( + &Struct(target_fields.to_vec().into()), + source_col.len(), + )); + } + let mut fields: Vec> = Vec::with_capacity(target_fields.len()); let mut arrays: Vec = Vec::with_capacity(target_fields.len()); let num_rows = source_col.len(); @@ -183,6 +190,15 @@ pub fn cast_column( (DataType::LargeList(_), DataType::LargeList(target_inner)) => { cast_list_column::(source_col, target_inner, cast_options) } + ( + DataType::FixedSizeList(_, source_list_size), + DataType::FixedSizeList(target_inner, target_list_size), + ) if source_list_size == target_list_size => cast_fixed_size_list_column( + source_col, + target_inner, + *target_list_size, + cast_options, + ), (DataType::ListView(_), DataType::ListView(target_inner)) => { cast_list_view_column::(source_col, target_inner, cast_options) } @@ -208,15 +224,7 @@ fn cast_list_column( target_inner_field: &FieldRef, cast_options: &CastOptions, ) -> Result { - let source_list = source_col - .as_any() - .downcast_ref::>() - .ok_or_else(|| { - crate::error::DataFusionError::Plan(format!( - "Expected list array but got {}", - source_col.data_type() - )) - })?; + let source_list = source_col.as_list::(); let cast_values = cast_column( source_list.values(), @@ -238,15 +246,7 @@ fn cast_list_view_column( target_inner_field: &FieldRef, cast_options: &CastOptions, ) -> Result { - let source_list = source_col - .as_any() - .downcast_ref::>() - .ok_or_else(|| { - crate::error::DataFusionError::Plan(format!( - "Expected list view array but got {}", - source_col.data_type() - )) - })?; + let source_list = source_col.as_list_view::(); let cast_values = cast_column( source_list.values(), @@ -264,6 +264,82 @@ fn cast_list_view_column( Ok(Arc::new(result)) } +fn cast_fixed_size_list_column( + source_col: &ArrayRef, + target_inner_field: &FieldRef, + target_list_size: i32, + cast_options: &CastOptions, +) -> Result { + let source_list = source_col.as_fixed_size_list(); + + let source_values = source_list.values(); + let target_type = target_inner_field.data_type(); + + let cast_values = match cast_column(source_values, target_type, cast_options) { + Ok(cast_values) => cast_values, + Err(error) => match cast_fixed_size_list_values_with_parent_nulls( + source_values, + target_type, + cast_options, + source_list.nulls(), + target_list_size, + ) { + Some(masked_cast) => masked_cast?, + None => return Err(error), + }, + }; + + Ok(Arc::new(FixedSizeListArray::try_new( + Arc::clone(target_inner_field), + target_list_size, + cast_values, + source_list.nulls().cloned(), + )?)) +} + +fn cast_fixed_size_list_values_with_parent_nulls( + source_values: &ArrayRef, + target_type: &DataType, + cast_options: &CastOptions, + parent_nulls: Option<&NullBuffer>, + list_size: i32, +) -> Option> { + let parent_nulls = parent_nulls.filter(|nulls| nulls.null_count() > 0)?; + + // FixedSizeList stores child slots for null parent lists. Those child + // values are semantically hidden, but recursive casts still inspect them. + let hidden_child_nulls = parent_nulls.expand(list_size as usize); + let masked_values = mask_array_values(source_values, &hidden_child_nulls); + Some(masked_values.and_then(|values| cast_column(&values, target_type, cast_options))) +} + +fn mask_array_values( + values: &ArrayRef, + additional_nulls: &NullBuffer, +) -> Result { + let nulls = NullBuffer::union(values.nulls(), Some(additional_nulls)); + + if let Some(struct_array) = values.as_any().downcast_ref::() { + let struct_nulls = nulls + .as_ref() + .expect("additional nulls always produce nulls"); + let arrays = struct_array + .columns() + .iter() + .map(|child| mask_array_values(child, struct_nulls)) + .collect::>>()?; + return Ok(Arc::new(StructArray::new( + struct_array.fields().clone(), + arrays, + nulls, + ))); + } + + Ok(make_array( + values.to_data().into_builder().nulls(nulls).build()?, + )) +} + fn cast_dictionary_column( source_col: &ArrayRef, source_key_type: &DataType, @@ -425,6 +501,12 @@ pub fn validate_data_type_compatibility( (Struct(source_nested), Struct(target_nested)) => { validate_struct_compatibility(source_nested, target_nested)?; } + ( + DataType::FixedSizeList(s, source_list_size), + DataType::FixedSizeList(t, target_list_size), + ) if source_list_size == target_list_size => { + validate_field_compatibility(s, t)?; + } (DataType::List(s), DataType::List(t)) | (DataType::LargeList(s), DataType::LargeList(t)) | (DataType::ListView(s), DataType::ListView(t)) @@ -460,8 +542,8 @@ pub fn validate_data_type_compatibility( /// name-based nested struct casting logic, rather than Arrow's standard cast. /// /// This is the case when both types are struct types, or both are the same -/// container type (List, LargeList, ListView, LargeListView, Dictionary) wrapping -/// types that recursively contain structs. +/// container type (List, LargeList, equal-width FixedSizeList, ListView, +/// LargeListView, Dictionary) wrapping types that recursively contain structs. /// /// Use this predicate at both planning time (to decide whether to apply struct /// compatibility validation) and execution time (to decide whether to route @@ -472,6 +554,12 @@ pub fn requires_nested_struct_cast( ) -> bool { match (source_type, target_type) { (Struct(_), Struct(_)) => true, + ( + DataType::FixedSizeList(s, source_list_size), + DataType::FixedSizeList(t, target_list_size), + ) if source_list_size == target_list_size => { + requires_nested_struct_cast(s.data_type(), t.data_type()) + } (DataType::List(s), DataType::List(t)) | (DataType::LargeList(s), DataType::LargeList(t)) | (DataType::ListView(s), DataType::ListView(t)) @@ -508,8 +596,9 @@ mod tests { use crate::{assert_contains, format::DEFAULT_CAST_OPTIONS}; use arrow::{ array::{ - BinaryArray, Int32Array, Int32Builder, Int64Array, ListArray, ListViewArray, - MapArray, MapBuilder, NullArray, StringArray, StringBuilder, + BinaryArray, FixedSizeListArray, Int32Array, Int32Builder, Int64Array, + ListArray, ListViewArray, MapArray, MapBuilder, NullArray, StringArray, + StringBuilder, }, buffer::{NullBuffer, ScalarBuffer}, datatypes::{DataType, Field, FieldRef, Int32Type}, @@ -1307,6 +1396,275 @@ mod tests { assert!(b_col.iter().all(|v| v.is_none())); } + fn fixed_size_list_struct_field(fields: Vec<(&str, DataType)>) -> FieldRef { + arc_field( + "item", + struct_type( + fields + .into_iter() + .map(|(name, data_type)| field(name, data_type)) + .collect(), + ), + ) + } + + fn create_fixed_size_list_test_fields( + source_struct_fields: Vec<(&str, DataType)>, + target_struct_fields: Vec<(&str, DataType)>, + ) -> (FieldRef, FieldRef) { + ( + fixed_size_list_struct_field(source_struct_fields), + fixed_size_list_struct_field(target_struct_fields), + ) + } + + fn fixed_size_list_struct_values( + array: &ArrayRef, + ) -> (&FixedSizeListArray, &StructArray) { + let list = array.as_any().downcast_ref::().unwrap(); + let values = list + .values() + .as_any() + .downcast_ref::() + .unwrap(); + (list, values) + } + + #[test] + fn test_cast_fixed_size_list_struct() { + let struct_arr = StructArray::from(vec![( + arc_field("a", DataType::Int32), + Arc::new(Int32Array::from(vec![1, 2, 3, 4])) as ArrayRef, + )]); + + let (source_field, target_field) = create_fixed_size_list_test_fields( + vec![("a", DataType::Int32)], + vec![("a", DataType::Int64), ("b", DataType::Utf8)], + ); + let source_col: ArrayRef = Arc::new(FixedSizeListArray::new( + source_field, + 2, + Arc::new(struct_arr), + Some(NullBuffer::from(vec![true, false])), + )); + let target_type = DataType::FixedSizeList(target_field, 2); + + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + let (result_list, struct_values) = fixed_size_list_struct_values(&result); + assert_eq!(result_list.len(), 2); + assert!(result_list.is_valid(0)); + assert!(result_list.is_null(1)); + let a_col = get_column_as!(&struct_values, "a", Int64Array); + assert_eq!(a_col.values(), &[1, 2, 3, 4]); + let b_col = get_column_as!(&struct_values, "b", StringArray); + assert!(b_col.iter().all(|v| v.is_none())); + } + + #[test] + fn test_validate_fixed_size_list_struct_compatibility() { + let (source_field, target_field) = create_fixed_size_list_test_fields( + vec![("a", DataType::Int32)], + vec![("a", DataType::Int64), ("b", DataType::Utf8)], + ); + let source = DataType::FixedSizeList(source_field, 2); + let target = DataType::FixedSizeList(target_field, 2); + + assert!(requires_nested_struct_cast(&source, &target)); + assert!(validate_data_type_compatibility("col", &source, &target).is_ok()); + } + + #[test] + fn test_validate_fixed_size_list_struct_missing_non_nullable_field_rejected() { + let (source_field, _) = create_fixed_size_list_test_fields( + vec![("a", DataType::Int32)], + vec![("a", DataType::Int64), ("b", DataType::Utf8)], + ); + let source = DataType::FixedSizeList(source_field, 2); + let target = DataType::FixedSizeList( + arc_field( + "item", + struct_type(vec![ + field("a", DataType::Int32), + non_null_field("b", DataType::Utf8), + ]), + ), + 2, + ); + + let error = validate_data_type_compatibility("col", &source, &target) + .unwrap_err() + .to_string(); + assert_contains!( + error, + "target field 'b' is non-nullable but missing from source" + ); + } + + #[test] + fn test_fixed_size_list_struct_size_mismatch_rejected() { + let source_field = fixed_size_list_struct_field(vec![("a", DataType::Int32)]); + let target_field = Arc::clone(&source_field); + let source_type = DataType::FixedSizeList(Arc::clone(&source_field), 2); + let target_type = DataType::FixedSizeList(target_field, 3); + + let validation_error = + validate_data_type_compatibility("col", &source_type, &target_type) + .unwrap_err() + .to_string(); + assert_contains!(validation_error, "Cannot cast struct field 'col'"); + + let struct_arr = StructArray::from(vec![( + arc_field("a", DataType::Int32), + Arc::new(Int32Array::from(vec![1, 2])) as ArrayRef, + )]); + let source_col: ArrayRef = Arc::new(FixedSizeListArray::new( + source_field, + 2, + Arc::new(struct_arr), + None, + )); + + let runtime_error = cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS) + .unwrap_err() + .to_string(); + assert_contains!( + runtime_error, + "cannot cast fixed-size-list to fixed-size-list with different size" + ); + } + + #[test] + fn test_cast_fixed_size_list_struct_all_null() { + let (source_field, target_field) = create_fixed_size_list_test_fields( + vec![("a", DataType::Int32)], + vec![("a", DataType::Int64), ("b", DataType::Utf8)], + ); + let source_col: ArrayRef = + Arc::new(FixedSizeListArray::new_null(source_field, 2, 2)); + let target_type = DataType::FixedSizeList(target_field, 2); + + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + let (result_list, struct_values) = fixed_size_list_struct_values(&result); + assert_eq!(result_list.null_count(), 2); + let a_col = get_column_as!(&struct_values, "a", Int64Array); + let b_col = get_column_as!(&struct_values, "b", StringArray); + assert!(a_col.iter().all(|v| v.is_none())); + assert!(b_col.iter().all(|v| v.is_none())); + } + + #[test] + fn test_fixed_size_list_struct_planner_runtime_parity_on_incompatible_type() { + let source_field = + arc_field("item", struct_type(vec![field("a", DataType::Binary)])); + let target_field = + arc_field("item", struct_type(vec![field("a", DataType::Int32)])); + let source_type = DataType::FixedSizeList(Arc::clone(&source_field), 2); + let target_type = DataType::FixedSizeList(target_field, 2); + let validation_error = + validate_data_type_compatibility("col", &source_type, &target_type) + .unwrap_err() + .to_string(); + assert_contains!(validation_error, "Cannot cast struct field 'a'"); + + let struct_arr = StructArray::from(vec![( + arc_field("a", DataType::Binary), + Arc::new(BinaryArray::from(vec![ + Some(b"x".as_ref()), + Some(b"y".as_ref()), + ])) as ArrayRef, + )]); + let source_col: ArrayRef = Arc::new(FixedSizeListArray::new( + source_field, + 2, + Arc::new(struct_arr), + None, + )); + + let runtime_error = cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS) + .unwrap_err() + .to_string(); + assert_contains!(runtime_error, "Cannot cast struct field 'a'"); + } + + #[test] + fn test_cast_fixed_size_list_struct_missing_non_nullable_field_runtime_rejected() { + let source_field = + arc_field("item", struct_type(vec![field("a", DataType::Int32)])); + let target_field = arc_field( + "item", + struct_type(vec![ + field("a", DataType::Int32), + non_null_field("b", DataType::Utf8), + ]), + ); + let source_col: ArrayRef = + Arc::new(FixedSizeListArray::new_null(source_field, 2, 1)); + let target_type = DataType::FixedSizeList(target_field, 2); + + let error = cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS) + .unwrap_err() + .to_string(); + assert_contains!( + error, + "target field 'b' is non-nullable but missing from source" + ); + } + + #[test] + fn test_cast_fixed_size_list_returns_error_for_non_nullable_child() { + let source_field = Arc::new(Field::new("item", DataType::Int32, true)); + let target_field = Arc::new(Field::new("item", DataType::Int32, false)); + let source_col: ArrayRef = Arc::new(FixedSizeListArray::new( + source_field, + 2, + Arc::new(Int32Array::from(vec![None, Some(1)])), + None, + )); + let target_type = DataType::FixedSizeList(target_field, 2); + + let error = cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS) + .unwrap_err() + .to_string(); + assert_contains!(error, "Found unmasked nulls for non-nullable"); + } + + #[test] + fn test_cast_sliced_fixed_size_list_struct_ignores_hidden_child_values() { + let source_field = + arc_field("item", struct_type(vec![field("a", DataType::Utf8)])); + let target_field = + arc_field("item", struct_type(vec![field("a", DataType::Int32)])); + let struct_arr = StructArray::from(vec![( + arc_field("a", DataType::Utf8), + Arc::new(StringArray::from(vec![ + "0", "0", "not_int", "also_bad", "1", "2", + ])) as ArrayRef, + )]); + let source_col: ArrayRef = Arc::new( + FixedSizeListArray::new( + source_field, + 2, + Arc::new(struct_arr), + Some(NullBuffer::from(vec![true, false, true])), + ) + .slice(1, 2), + ); + let target_type = DataType::FixedSizeList(target_field, 2); + + let result = + cast_column(&source_col, &target_type, &DEFAULT_CAST_OPTIONS).unwrap(); + let (result_list, struct_values) = fixed_size_list_struct_values(&result); + assert!(result_list.is_null(0)); + assert!(result_list.is_valid(1)); + let a_col = get_column_as!(&struct_values, "a", Int32Array); + assert!(a_col.is_null(0)); + assert!(a_col.is_null(1)); + assert_eq!(a_col.value(2), 1); + assert_eq!(a_col.value(3), 2); + } + #[test] fn test_requires_nested_struct_cast() { let s1 = struct_type(vec![field("a", DataType::Int32)]); @@ -1322,8 +1680,12 @@ mod tests { &DataType::Dictionary(Box::new(DataType::Int32), Box::new(s2.clone())), )); assert!(requires_nested_struct_cast( - &DataType::ListView(arc_field("item", s1)), - &DataType::ListView(arc_field("item", s2)), + &DataType::ListView(arc_field("item", s1.clone())), + &DataType::ListView(arc_field("item", s2.clone())), + )); + assert!(requires_nested_struct_cast( + &DataType::FixedSizeList(arc_field("item", s1), 2), + &DataType::FixedSizeList(arc_field("item", s2), 2), )); // Non-struct types should return false. @@ -1335,5 +1697,9 @@ mod tests { &DataType::List(arc_field("item", DataType::Int32)), &DataType::List(arc_field("item", DataType::Int64)), )); + assert!(!requires_nested_struct_cast( + &DataType::FixedSizeList(arc_field("item", DataType::Int32), 2), + &DataType::FixedSizeList(arc_field("item", DataType::Int64), 2), + )); } } diff --git a/datafusion/common/src/partitioning.rs b/datafusion/common/src/partitioning.rs new file mode 100644 index 0000000000000..8a7212c2e3089 --- /dev/null +++ b/datafusion/common/src/partitioning.rs @@ -0,0 +1,104 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::utils::compare_rows; +use crate::{Result, ScalarValue, error::_plan_err}; +use arrow::compute::SortOptions; +use std::cmp::Ordering; +use std::fmt::{self, Display}; + +/// A boundary between adjacent range partitions. +/// +/// A split point is a tuple with one [`ScalarValue`] per partitioning +/// expression. Split points are interpreted lexicographically according to the +/// ordering of the range partitioning that owns them. +/// +/// `N` split points define `N + 1` partitions: +/// +/// ```text +/// partition 0: key < split_points[0] +/// partition 1: split_points[0] <= key < split_points[1] +/// ... +/// partition N - 1: split_points[N - 2] <= key < split_points[N - 1] +/// partition N: split_points[N - 1] <= key +/// ``` +/// +/// Values equal to split point `i` belong to partition `i + 1`, so interior +/// partitions are lower-inclusive and upper-exclusive. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct SplitPoint { + values: Vec, +} + +impl SplitPoint { + /// Creates a new split point from its tuple values. + pub fn new(values: Vec) -> Self { + Self { values } + } + + /// Returns the tuple values for this split point. + pub fn values(&self) -> &[ScalarValue] { + &self.values + } +} + +impl Display for SplitPoint { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let values = self + .values + .iter() + .map(ToString::to_string) + .collect::>() + .join(", "); + write!(f, "({values})") + } +} + +/// Validates that split points match the ordering width and are strictly +/// ordered according to the provided sort options. +pub fn validate_range_split_points( + split_points: &[SplitPoint], + sort_options: &[SortOptions], +) -> Result<()> { + let width = sort_options.len(); + for (idx, split_point) in split_points.iter().enumerate() { + let split_point_width = split_point.values().len(); + if split_point_width != width { + return _plan_err!( + "Range partitioning split point {idx} has width {split_point_width}, but ordering has width {width}" + ); + } + } + + for (idx, split_points) in split_points.windows(2).enumerate() { + if compare_rows( + split_points[0].values(), + split_points[1].values(), + sort_options, + )? != Ordering::Less + { + return _plan_err!( + "Range partitioning split points must be strictly ordered: split point {idx} ({}) must be less than split point {} ({})", + split_points[0], + idx + 1, + split_points[1] + ); + } + } + + Ok(()) +} diff --git a/datafusion/common/src/pruning.rs b/datafusion/common/src/pruning.rs index ebae23f0723a1..a36ac9f795b95 100644 --- a/datafusion/common/src/pruning.rs +++ b/datafusion/common/src/pruning.rs @@ -305,7 +305,7 @@ impl PruningStatistics for PartitionPruningStatistics { /// that has statistics of its columns. /// /// It is up to the caller to decide what each container represents. For -/// example, they can come from a file (e.g. [`PartitionedFile`]) or a set of of +/// example, they can come from a file (e.g. [`PartitionedFile`]) or a set of /// files (e.g. [`FileGroup`]) /// /// [`PartitionedFile`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.PartitionedFile.html diff --git a/datafusion/common/src/scalar/consts.rs b/datafusion/common/src/scalar/consts.rs index 599c2523cd2c7..df12265a3723c 100644 --- a/datafusion/common/src/scalar/consts.rs +++ b/datafusion/common/src/scalar/consts.rs @@ -17,6 +17,9 @@ // Constants defined for scalar construction. +use arrow::datatypes::{Decimal32Type, Decimal64Type, Decimal128Type, DecimalType}; +use arrow::datatypes::{Decimal256Type, i256}; + // Next F16 value above π (upper bound) pub(super) const PI_UPPER_F16: half::f16 = half::f16::from_bits(0x4249); @@ -54,3 +57,63 @@ pub(super) const NEGATIVE_FRAC_PI_2_LOWER_F32: f32 = // Next f64 value below -π/2 (lower bound) pub(super) const NEGATIVE_FRAC_PI_2_LOWER_F64: f64 = (-std::f64::consts::FRAC_PI_2).next_down(); + +// Generate lookup table for 1 values of decimals (1, 10, 100, etc.) +macro_rules! decimal_ones_lut { + () => {{ + let mut values = [1; _]; + let mut i = 1; + while i < values.len() { + values[i] = values[i - 1] * 10; + i += 1; + } + values + }}; +} + +// 1, 10, 100 values meant to be indexed by scale. We omit handling for MAX_SCALE +// itself (we don't go to MAX_SCALE + 1) since we can't represent a 1 value at +// that scale. +pub(super) const DECIMAL32_ONES: [i32; Decimal32Type::MAX_SCALE as usize] = + decimal_ones_lut!(); +pub(super) const DECIMAL64_ONES: [i64; Decimal64Type::MAX_SCALE as usize] = + decimal_ones_lut!(); +pub(super) const DECIMAL128_ONES: [i128; Decimal128Type::MAX_SCALE as usize] = + decimal_ones_lut!(); +pub(super) const DECIMAL256_ONES: [i256; Decimal256Type::MAX_SCALE as usize] = { + // This code was generated by codex and frankly I don't know how it works, + // but the test below verifies it outputs the correct values so ¯\_(ツ)_/¯ + // + // This is mainly a shortcut for not needing to manually list out each value + // anyway. + // + // TODO: simplify this after https://github.com/apache/arrow-rs/pull/10363 + // lands upstream + let mut values = [i256::ONE; _]; + let mut i = 1; + while i < values.len() { + let (low, high) = values[i - 1].to_parts(); + let low_product = (low as u64 as u128) * 10; + let high_product = (low >> 64) * 10 + (low_product >> 64); + let low = ((high_product as u64 as u128) << 64) | low_product as u64 as u128; + let carry = (high_product >> 64) as i128; + values[i] = i256::from_parts(low, high * 10 + carry); + i += 1; + } + values +}; + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_ensure_correct_decimal256_ones() { + for (scale, val) in DECIMAL256_ONES.iter().enumerate() { + let zeros = "0".repeat(scale); + let num = "1".to_string() + &zeros; + let num = i256::from_string(&num).unwrap(); + assert_eq!(num, *val, "{scale}"); + } + } +} diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs index 644ed1085d742..cb0442392ad21 100644 --- a/datafusion/common/src/scalar/mod.rs +++ b/datafusion/common/src/scalar/mod.rs @@ -23,7 +23,7 @@ mod struct_builder; use std::borrow::Borrow; use std::cmp::Ordering; -use std::collections::{HashSet, VecDeque}; +use std::collections::{HashMap, HashSet, VecDeque}; use std::convert::Infallible; use std::fmt; use std::fmt::Write; @@ -54,6 +54,9 @@ use crate::cast::{ use crate::error::{_exec_err, _internal_err, _not_impl_err, DataFusionError, Result}; use crate::format::DEFAULT_CAST_OPTIONS; use crate::hash_utils::create_hashes; +use crate::scalar::consts::{ + DECIMAL32_ONES, DECIMAL64_ONES, DECIMAL128_ONES, DECIMAL256_ONES, +}; use crate::utils::SingleRowListArrayBuilder; use crate::{_internal_datafusion_err, arrow_datafusion_err}; use arrow::array::{ @@ -83,15 +86,20 @@ use arrow::datatypes::{ Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType, Field, FieldRef, Float32Type, Int8Type, Int16Type, Int32Type, Int64Type, IntervalDayTime, IntervalDayTimeType, IntervalMonthDayNano, IntervalMonthDayNanoType, IntervalUnit, - IntervalYearMonthType, RunEndIndexType, TimeUnit, TimestampMicrosecondType, - TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, - UInt16Type, UInt32Type, UInt64Type, UnionFields, UnionMode, i256, - validate_decimal_precision_and_scale, + IntervalYearMonthType, MAX_DECIMAL32_FOR_EACH_PRECISION, + MAX_DECIMAL64_FOR_EACH_PRECISION, MAX_DECIMAL128_FOR_EACH_PRECISION, + MAX_DECIMAL256_FOR_EACH_PRECISION, MIN_DECIMAL32_FOR_EACH_PRECISION, + MIN_DECIMAL64_FOR_EACH_PRECISION, MIN_DECIMAL128_FOR_EACH_PRECISION, + MIN_DECIMAL256_FOR_EACH_PRECISION, RunEndIndexType, TimeUnit, + TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, + TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, UnionFields, + UnionMode, i256, validate_decimal_precision_and_scale, }; use arrow::util::display::{ArrayFormatter, FormatOptions, array_value_to_string}; use cache::{get_or_create_cached_key_array, get_or_create_cached_null_array}; use chrono::{Duration, NaiveDate}; use half::f16; +use num_traits::ToPrimitive; pub use struct_builder::ScalarStructBuilder; const SECONDS_PER_DAY: i64 = 86_400; @@ -148,6 +156,30 @@ pub fn date_to_timestamp_multiplier( } } +/// Returns the multiplier that converts the input timestamp representation into +/// the desired timestamp unit, if the conversion requires a multiplication that +/// can overflow an `i64`. +pub fn timestamp_to_timestamp_multiplier( + source_type: &DataType, + target_type: &DataType, +) -> Option { + let (DataType::Timestamp(source_unit, _), DataType::Timestamp(target_unit, _)) = + (source_type, target_type) + else { + return None; + }; + + match (source_unit, target_unit) { + (TimeUnit::Second, TimeUnit::Millisecond) => Some(1_000), + (TimeUnit::Second, TimeUnit::Microsecond) => Some(1_000_000), + (TimeUnit::Second, TimeUnit::Nanosecond) => Some(1_000_000_000), + (TimeUnit::Millisecond, TimeUnit::Microsecond) => Some(1_000), + (TimeUnit::Millisecond, TimeUnit::Nanosecond) => Some(1_000_000), + (TimeUnit::Microsecond, TimeUnit::Nanosecond) => Some(1_000), + _ => None, + } +} + /// Ensures the provided value can be represented as a timestamp with the given /// multiplier. Returns an [`DataFusionError::Execution`] when the converted /// value would overflow the timestamp range. @@ -1779,48 +1811,56 @@ impl ScalarValue { *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match 10_i32.checked_pow(*scale as u32) { - Some(value) => { - ScalarValue::Decimal32(Some(value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + *precision != *scale as u8, + "Can't represent one at scale {} with precision {}", + *scale, + *precision + ); + let one = DECIMAL32_ONES[*scale as usize]; + ScalarValue::Decimal32(Some(one), *precision, *scale) } DataType::Decimal64(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i64::from(10).checked_pow(*scale as u32) { - Some(value) => { - ScalarValue::Decimal64(Some(value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + *precision != *scale as u8, + "Can't represent one at scale {} with precision {}", + *scale, + *precision + ); + let one = DECIMAL64_ONES[*scale as usize]; + ScalarValue::Decimal64(Some(one), *precision, *scale) } DataType::Decimal128(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i128::from(10).checked_pow(*scale as u32) { - Some(value) => { - ScalarValue::Decimal128(Some(value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + *precision != *scale as u8, + "Can't represent one at scale {} with precision {}", + *scale, + *precision + ); + let one = DECIMAL128_ONES[*scale as usize]; + ScalarValue::Decimal128(Some(one), *precision, *scale) } DataType::Decimal256(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i256::from(10).checked_pow(*scale as u32) { - Some(value) => { - ScalarValue::Decimal256(Some(value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + *precision != *scale as u8, + "Can't represent one at scale {} with precision {}", + *scale, + *precision + ); + let one = DECIMAL256_ONES[*scale as usize]; + ScalarValue::Decimal256(Some(one), *precision, *scale) } _ => { return _not_impl_err!( @@ -1833,10 +1873,10 @@ impl ScalarValue { /// Create a negative one value in the given type. pub fn new_negative_one(datatype: &DataType) -> Result { Ok(match datatype { - DataType::Int8 | DataType::UInt8 => ScalarValue::Int8(Some(-1)), - DataType::Int16 | DataType::UInt16 => ScalarValue::Int16(Some(-1)), - DataType::Int32 | DataType::UInt32 => ScalarValue::Int32(Some(-1)), - DataType::Int64 | DataType::UInt64 => ScalarValue::Int64(Some(-1)), + DataType::Int8 => ScalarValue::Int8(Some(-1)), + DataType::Int16 => ScalarValue::Int16(Some(-1)), + DataType::Int32 => ScalarValue::Int32(Some(-1)), + DataType::Int64 => ScalarValue::Int64(Some(-1)), DataType::Float16 => ScalarValue::Float16(Some(f16::NEG_ONE)), DataType::Float32 => ScalarValue::Float32(Some(-1.0)), DataType::Float64 => ScalarValue::Float64(Some(-1.0)), @@ -1845,48 +1885,56 @@ impl ScalarValue { *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match 10_i32.checked_pow(*scale as u32) { - Some(value) => { - ScalarValue::Decimal32(Some(-value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + *precision != *scale as u8, + "Can't represent negative one at scale {} with precision {}", + *scale, + *precision + ); + let one = DECIMAL32_ONES[*scale as usize]; + ScalarValue::Decimal32(Some(-one), *precision, *scale) } DataType::Decimal64(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i64::from(10).checked_pow(*scale as u32) { - Some(value) => { - ScalarValue::Decimal64(Some(-value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + *precision != *scale as u8, + "Can't represent negative one at scale {} with precision {}", + *scale, + *precision + ); + let one = DECIMAL64_ONES[*scale as usize]; + ScalarValue::Decimal64(Some(-one), *precision, *scale) } DataType::Decimal128(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i128::from(10).checked_pow(*scale as u32) { - Some(value) => { - ScalarValue::Decimal128(Some(-value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + *precision != *scale as u8, + "Can't represent negative one at scale {} with precision {}", + *scale, + *precision + ); + let one = DECIMAL128_ONES[*scale as usize]; + ScalarValue::Decimal128(Some(-one), *precision, *scale) } DataType::Decimal256(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i256::from(10).checked_pow(*scale as u32) { - Some(value) => { - ScalarValue::Decimal256(Some(-value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + *precision != *scale as u8, + "Can't represent one at scale {} with precision {}", + *scale, + *precision + ); + let one = DECIMAL256_ONES[*scale as usize]; + ScalarValue::Decimal256(Some(-one), *precision, *scale) } _ => { return _not_impl_err!( @@ -1914,48 +1962,64 @@ impl ScalarValue { *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match 10_i32.checked_pow((*scale + 1) as u32) { - Some(value) => { - ScalarValue::Decimal32(Some(value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + (*precision - *scale as u8) > 1, + "Can't represent ten at scale {} with precision {}", + *scale, + *precision + ); + // +1 safe since we validate above that scale must be less than + // the max possible scale + let ten = DECIMAL32_ONES[*scale as usize + 1]; + ScalarValue::Decimal32(Some(ten), *precision, *scale) } DataType::Decimal64(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i64::from(10).checked_pow((*scale + 1) as u32) { - Some(value) => { - ScalarValue::Decimal64(Some(value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + (*precision - *scale as u8) > 1, + "Can't represent ten at scale {} with precision {}", + *scale, + *precision + ); + // +1 safe since we validate above that scale must be less than + // the max possible scale + let ten = DECIMAL64_ONES[*scale as usize + 1]; + ScalarValue::Decimal64(Some(ten), *precision, *scale) } DataType::Decimal128(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i128::from(10).checked_pow((*scale + 1) as u32) { - Some(value) => { - ScalarValue::Decimal128(Some(value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + (*precision - *scale as u8) > 1, + "Can't represent ten at scale {} with precision {}", + *scale, + *precision + ); + // +1 safe since we validate above that scale must be less than + // the max possible scale + let ten = DECIMAL128_ONES[*scale as usize + 1]; + ScalarValue::Decimal128(Some(ten), *precision, *scale) } DataType::Decimal256(precision, scale) => { Self::validate_decimal_or_internal_err::( *precision, *scale, )?; assert_or_internal_err!(*scale >= 0, "Negative scale is not supported"); - match i256::from(10).checked_pow((*scale + 1) as u32) { - Some(value) => { - ScalarValue::Decimal256(Some(value), *precision, *scale) - } - None => return _internal_err!("Unsupported scale {scale}"), - } + assert_or_internal_err!( + (*precision - *scale as u8) > 1, + "Can't represent ten at scale {} with precision {}", + *scale, + *precision + ); + // +1 safe since we validate above that scale must be less than + // the max possible scale + let ten = DECIMAL256_ONES[*scale as usize + 1]; + ScalarValue::Decimal256(Some(ten), *precision, *scale) } _ => { return _not_impl_err!( @@ -2274,7 +2338,18 @@ impl ScalarValue { | ScalarValue::Int64(None) | ScalarValue::Float16(None) | ScalarValue::Float32(None) - | ScalarValue::Float64(None) => Ok(self.clone()), + | ScalarValue::Float64(None) + | ScalarValue::IntervalYearMonth(None) + | ScalarValue::IntervalDayTime(None) + | ScalarValue::IntervalMonthDayNano(None) + | ScalarValue::Decimal32(None, _, _) + | ScalarValue::Decimal64(None, _, _) + | ScalarValue::Decimal128(None, _, _) + | ScalarValue::Decimal256(None, _, _) + | ScalarValue::TimestampSecond(None, _) + | ScalarValue::TimestampMillisecond(None, _) + | ScalarValue::TimestampMicrosecond(None, _) + | ScalarValue::TimestampNanosecond(None, _) => Ok(self.clone()), ScalarValue::Float16(Some(v)) => Ok(ScalarValue::Float16(Some(-v))), ScalarValue::Float64(Some(v)) => Ok(ScalarValue::Float64(Some(-v))), ScalarValue::Float32(Some(v)) => Ok(ScalarValue::Float32(Some(-v))), @@ -2561,63 +2636,107 @@ impl ScalarValue { /// distance is greater than [`usize::MAX`]. If the type is a float, then the distance will be /// rounded to the nearest integer. /// - /// /// Note: the datatype itself must support subtraction. pub fn distance(&self, other: &ScalarValue) -> Option { + self.distance_u64(other) + .and_then(|d| usize::try_from(d).ok()) + } + + /// Helper to convert a rounded float distance to u64, returning None if it exceeds u64::MAX, is negative, or is not finite. + fn rounded_float_distance_u64(diff: f64) -> Option { + if diff.is_finite() && diff >= 0.0 && diff < u64::MAX as f64 { + Some(diff as u64) + } else { + None + } + } + + /// Absolute distance between two numeric values (of the same type). This method will return + /// None if either one of the arguments are null. It might also return None if the resulting + /// distance is greater than [`u64::MAX`]. If the type is a float, then the distance will be + /// rounded to the nearest integer. + /// + /// Note: the datatype itself must support subtraction. + pub fn distance_u64(&self, other: &ScalarValue) -> Option { match (self, other) { - (Self::Int8(Some(l)), Self::Int8(Some(r))) => Some(l.abs_diff(*r) as _), - (Self::Int16(Some(l)), Self::Int16(Some(r))) => Some(l.abs_diff(*r) as _), - (Self::Int32(Some(l)), Self::Int32(Some(r))) => Some(l.abs_diff(*r) as _), - (Self::Int64(Some(l)), Self::Int64(Some(r))) => Some(l.abs_diff(*r) as _), - (Self::UInt8(Some(l)), Self::UInt8(Some(r))) => Some(l.abs_diff(*r) as _), - (Self::UInt16(Some(l)), Self::UInt16(Some(r))) => Some(l.abs_diff(*r) as _), - (Self::UInt32(Some(l)), Self::UInt32(Some(r))) => Some(l.abs_diff(*r) as _), - (Self::UInt64(Some(l)), Self::UInt64(Some(r))) => Some(l.abs_diff(*r) as _), + (Self::Int8(Some(l)), Self::Int8(Some(r))) => Some(l.abs_diff(*r) as u64), + (Self::Int16(Some(l)), Self::Int16(Some(r))) => Some(l.abs_diff(*r) as u64), + (Self::Int32(Some(l)), Self::Int32(Some(r))) => Some(l.abs_diff(*r) as u64), + (Self::Int64(Some(l)), Self::Int64(Some(r))) => Some(l.abs_diff(*r)), + (Self::UInt8(Some(l)), Self::UInt8(Some(r))) => Some(l.abs_diff(*r) as u64), + (Self::UInt16(Some(l)), Self::UInt16(Some(r))) => Some(l.abs_diff(*r) as u64), + (Self::UInt32(Some(l)), Self::UInt32(Some(r))) => Some(l.abs_diff(*r) as u64), + (Self::UInt64(Some(l)), Self::UInt64(Some(r))) => Some(l.abs_diff(*r)), // TODO: we might want to look into supporting ceil/floor here for floats. (Self::Float16(Some(l)), Self::Float16(Some(r))) => { - Some((f16::to_f32(*l) - f16::to_f32(*r)).abs().round() as _) + let diff = (f16::to_f32(*l) - f16::to_f32(*r)).abs().round(); + Self::rounded_float_distance_u64(diff as f64) } (Self::Float32(Some(l)), Self::Float32(Some(r))) => { - Some((l - r).abs().round() as _) + let diff = (l - r).abs().round(); + Self::rounded_float_distance_u64(diff as f64) } (Self::Float64(Some(l)), Self::Float64(Some(r))) => { - Some((l - r).abs().round() as _) + let diff = (l - r).abs().round(); + Self::rounded_float_distance_u64(diff) } - (Self::Date32(Some(l)), Self::Date32(Some(r))) => Some(l.abs_diff(*r) as _), - (Self::Date64(Some(l)), Self::Date64(Some(r))) => Some(l.abs_diff(*r) as _), + (Self::Date32(Some(l)), Self::Date32(Some(r))) => Some(l.abs_diff(*r) as u64), + (Self::Date64(Some(l)), Self::Date64(Some(r))) => Some(l.abs_diff(*r)), // Timestamp values are stored as epoch ticks regardless of timezone // annotation, so the distance is tz-independent (tz is display metadata). (Self::TimestampSecond(Some(l), _), Self::TimestampSecond(Some(r), _)) => { - Some(l.abs_diff(*r) as _) + Some(l.abs_diff(*r)) } ( Self::TimestampMillisecond(Some(l), _), Self::TimestampMillisecond(Some(r), _), - ) => Some(l.abs_diff(*r) as _), + ) => Some(l.abs_diff(*r)), ( Self::TimestampMicrosecond(Some(l), _), Self::TimestampMicrosecond(Some(r), _), - ) => Some(l.abs_diff(*r) as _), + ) => Some(l.abs_diff(*r)), ( Self::TimestampNanosecond(Some(l), _), Self::TimestampNanosecond(Some(r), _), - ) => Some(l.abs_diff(*r) as _), + ) => Some(l.abs_diff(*r)), ( - Self::Decimal128(Some(l), lprecision, lscale), - Self::Decimal128(Some(r), rprecision, rscale), + Self::Decimal32(Some(l), _, lscale), + Self::Decimal32(Some(r), _, rscale), ) => { - if lprecision == rprecision && lscale == rscale { - l.checked_sub(*r)?.checked_abs()?.to_usize() + // In order to be aligned with PartialOrd we only + // check for equal scale, ignoring precision + if lscale == rscale { + Some(l.abs_diff(*r) as u64) } else { None } } ( - Self::Decimal256(Some(l), lprecision, lscale), - Self::Decimal256(Some(r), rprecision, rscale), + Self::Decimal64(Some(l), _, lscale), + Self::Decimal64(Some(r), _, rscale), ) => { - if lprecision == rprecision && lscale == rscale { - l.checked_sub(*r)?.checked_abs()?.to_usize() + if lscale == rscale { + Some(l.abs_diff(*r)) + } else { + None + } + } + ( + Self::Decimal128(Some(l), _, lscale), + Self::Decimal128(Some(r), _, rscale), + ) => { + if lscale == rscale { + l.checked_sub(*r)?.checked_abs()?.to_u64() + } else { + None + } + } + ( + Self::Decimal256(Some(l), _, lscale), + Self::Decimal256(Some(r), _, rscale), + ) => { + if lscale == rscale { + l.checked_sub(*r)?.checked_abs()?.to_u64() } else { None } @@ -3190,7 +3309,8 @@ impl ScalarValue { let values = if values.is_empty() { new_empty_array(data_type) } else { - Self::iter_to_array(values.iter().cloned()).unwrap() + let arr = Self::iter_to_array(values.iter().cloned()).unwrap(); + cast_with_options(&arr, data_type, &DEFAULT_CAST_OPTIONS).unwrap() }; Arc::new( SingleRowListArrayBuilder::new(values) @@ -3252,7 +3372,8 @@ impl ScalarValue { let values = if values.len() == 0 { new_empty_array(data_type) } else { - Self::iter_to_array(values).unwrap() + let arr = Self::iter_to_array(values).unwrap(); + cast_with_options(&arr, data_type, &DEFAULT_CAST_OPTIONS).unwrap() }; Arc::new( SingleRowListArrayBuilder::new(values) @@ -3295,7 +3416,8 @@ impl ScalarValue { let values = if values.is_empty() { new_empty_array(data_type) } else { - Self::iter_to_array(values.iter().cloned()).unwrap() + let arr = Self::iter_to_array(values.iter().cloned()).unwrap(); + cast_with_options(&arr, data_type, &DEFAULT_CAST_OPTIONS).unwrap() }; Arc::new(SingleRowListArrayBuilder::new(values).build_large_list_array()) } @@ -4212,22 +4334,70 @@ impl ScalarValue { Some(v.as_ref().map(|v| v.as_str())) } - /// Try to cast this value to a ScalarValue of type `data_type` + /// Cast this value to a `ScalarValue` of type `target_type` using the + /// default [`CastOptions`]. + /// + /// This is a general-purpose cast with the same semantics as the Arrow + /// [`cast_with_options`] kernel and can therefore **lose information** -- + /// for example casting the floating point value `123.45` to the integer + /// `123`. + /// + /// Returns an error for casts the Arrow kernel cannot perform. + /// + /// # See Also + /// - [`try_cast_literal_to_type`]: for a *value-preserving* cast + /// + /// [`try_cast_literal_to_type`]: https://docs.rs/datafusion/latest/datafusion/logical_expr_common/casts/fn.try_cast_literal_to_type.html pub fn cast_to(&self, target_type: &DataType) -> Result { self.cast_to_with_options(target_type, &DEFAULT_CAST_OPTIONS) } - /// Try to cast this value to a ScalarValue of type `data_type` with [`CastOptions`] + /// Cast this value to type `target_type` with the given [`CastOptions`]. + /// + /// # See Also + /// - [`ScalarValue::cast_to`] for more details. + /// - [`try_cast_literal_to_type`]: for a *value-preserving* cast + /// + /// [`try_cast_literal_to_type`]: https://docs.rs/datafusion/latest/datafusion/logical_expr_common/casts/fn.try_cast_literal_to_type.html pub fn cast_to_with_options( &self, target_type: &DataType, cast_options: &CastOptions<'static>, ) -> Result { let source_type = self.data_type(); + + // Fast path: an identical target type needs no conversion at all. + if &source_type == target_type { + return Ok(self.clone()); + } + + // Fast path: conversions among the string types (`Utf8`, `LargeUtf8`, + // `Utf8View`) are value-preserving, so we can rewrap the string + // directly instead of building a single-row array and invoking the + // arrow cast kernel. + if source_type.is_string() && target_type.is_string() { + // `self` is one of the string types, so `try_as_str` returns `Some` + let value = self.try_as_str().flatten().map(|s| s.to_string()); + return Ok(match target_type { + DataType::Utf8 => ScalarValue::Utf8(value), + DataType::LargeUtf8 => ScalarValue::LargeUtf8(value), + DataType::Utf8View => ScalarValue::Utf8View(value), + _ => unreachable!("matched a string target type above"), + }); + } + if let Some(multiplier) = date_to_timestamp_multiplier(&source_type, target_type) - && let Some(value) = self.date_scalar_value_as_i64() + .or_else(|| timestamp_to_timestamp_multiplier(&source_type, target_type)) + && let Some(value) = self.temporal_scalar_value_as_i64() { - ensure_timestamp_in_bounds(value, multiplier, &source_type, target_type)?; + match ensure_timestamp_in_bounds(value, multiplier, &source_type, target_type) + { + Ok(()) => {} + Err(_) if cast_options.safe => { + return ScalarValue::try_new_null(target_type); + } + Err(e) => return Err(e), + } } let scalar_array = self.to_array()?; @@ -4247,10 +4417,14 @@ impl ScalarValue { ScalarValue::try_from_array(&cast_arr, 0) } - fn date_scalar_value_as_i64(&self) -> Option { + fn temporal_scalar_value_as_i64(&self) -> Option { match self { ScalarValue::Date32(Some(value)) => Some(i64::from(*value)), ScalarValue::Date64(Some(value)) => Some(*value), + ScalarValue::TimestampSecond(Some(value), _) + | ScalarValue::TimestampMillisecond(Some(value), _) + | ScalarValue::TimestampMicrosecond(Some(value), _) + | ScalarValue::TimestampNanosecond(Some(value), _) => Some(*value), _ => None, } } @@ -4713,12 +4887,32 @@ impl ScalarValue { .sum::() } + /// Estimates [size](Self::size) of [`HashMap`] keyed by [`ScalarValue`] in bytes. + /// + /// Includes the size of the [`HashMap`] container itself. Heap payload of + /// `V` is not accounted for; callers storing heap-backed values should + /// supplement this estimate. + #[allow(clippy::allow_attributes, clippy::mutable_key_type)] // ScalarValue has interior mutability but is intentionally used as hash key + pub fn size_of_hashmap(map: &HashMap) -> usize { + size_of_val(map) + + ((size_of::() + size_of::()) * map.capacity()) + + map.keys().map(|k| k.size() - size_of_val(k)).sum::() + } + /// Compacts the allocation referenced by `self` to the minimum, copying the data if /// necessary. /// /// This can be relevant when `self` is a list or contains a list as a nested value, as /// a single list holds an Arc to its entire original array buffer. pub fn compact(&mut self) { + // copy_array_data + compact_view_buffers + downcast back, all in one step. + macro_rules! compact_array { + ($arr:expr, $from_type:ty, $($as_method:tt)+) => { + *Arc::make_mut($arr) = ScalarValue::compact_view_buffers( + Arc::new(<$from_type>::from(copy_array_data(&$arr.to_data()))) as ArrayRef, + ).$($as_method)+.clone() + }; + } match self { ScalarValue::Null | ScalarValue::Boolean(_) @@ -4762,33 +4956,20 @@ impl ScalarValue { | ScalarValue::LargeBinary(_) | ScalarValue::BinaryView(_) => (), ScalarValue::FixedSizeList(arr) => { - let array = copy_array_data(&arr.to_data()); - *Arc::make_mut(arr) = FixedSizeListArray::from(array); - } - ScalarValue::List(arr) => { - let array = copy_array_data(&arr.to_data()); - *Arc::make_mut(arr) = ListArray::from(array); + compact_array!(arr, FixedSizeListArray, as_fixed_size_list()) } + ScalarValue::List(arr) => compact_array!(arr, ListArray, as_list::()), ScalarValue::LargeList(arr) => { - let array = copy_array_data(&arr.to_data()); - *Arc::make_mut(arr) = LargeListArray::from(array) + compact_array!(arr, LargeListArray, as_list::()) } ScalarValue::ListView(arr) => { - let array = copy_array_data(&arr.to_data()); - *Arc::make_mut(arr) = ListViewArray::from(array); + compact_array!(arr, ListViewArray, as_list_view::()) } ScalarValue::LargeListView(arr) => { - let array = copy_array_data(&arr.to_data()); - *Arc::make_mut(arr) = LargeListViewArray::from(array) - } - ScalarValue::Struct(arr) => { - let array = copy_array_data(&arr.to_data()); - *Arc::make_mut(arr) = StructArray::from(array); - } - ScalarValue::Map(arr) => { - let array = copy_array_data(&arr.to_data()); - *Arc::make_mut(arr) = MapArray::from(array); + compact_array!(arr, LargeListViewArray, as_list_view::()) } + ScalarValue::Struct(arr) => compact_array!(arr, StructArray, as_struct()), + ScalarValue::Map(arr) => compact_array!(arr, MapArray, as_map()), ScalarValue::Union(val, _, _) => { if let Some((_, value)) = val.as_mut() { value.compact(); @@ -4809,6 +4990,95 @@ impl ScalarValue { self } + /// Recursively compacts the backing buffers of any [`StringViewArray`] or + /// [`BinaryViewArray`] nested within `array`. + /// + /// View-typed arrays keep an `Arc` reference to their original backing + /// buffers, so a single scalar extracted from a large batch still retains + /// the entire buffer. Calling [`.gc()`][StringViewArray::gc] copies only + /// the bytes that are actually referenced by the surviving views, releasing + /// the rest. + /// + /// Container types (`List`, `LargeList`, `FixedSizeList`, `ListView`, + /// `LargeListView`, `Struct`, `Map`) are handled by recursing into their + /// child / values arrays and reconstructing the parent with the compacted + /// children. All other types are returned unchanged. + fn compact_view_buffers(array: ArrayRef) -> ArrayRef { + // Macro for the i32/i64-offset list pair (List / LargeList). + macro_rules! gc_list { + ($field:expr, $offset_type:ty, $array_type:ty) => {{ + let list = array.as_list::<$offset_type>(); + Arc::new(<$array_type>::new( + Arc::clone($field), + list.offsets().clone(), + ScalarValue::compact_view_buffers(Arc::clone(list.values())), + list.nulls().cloned(), + )) as ArrayRef + }}; + } + // Macro for the i32/i64-offset list-view pair (ListView / LargeListView). + macro_rules! gc_list_view { + ($field:expr, $offset_type:ty, $array_type:ty) => {{ + let list = array.as_list_view::<$offset_type>(); + Arc::new(<$array_type>::new( + Arc::clone($field), + list.offsets().clone(), + list.sizes().clone(), + ScalarValue::compact_view_buffers(Arc::clone(list.values())), + list.nulls().cloned(), + )) as ArrayRef + }}; + } + + match array.data_type() { + DataType::Utf8View => Arc::new(array.as_string_view().gc()), + DataType::BinaryView => Arc::new(array.as_binary_view().gc()), + DataType::Struct(_) => { + let s = array.as_struct(); + let columns = s + .columns() + .iter() + .map(|c| ScalarValue::compact_view_buffers(Arc::clone(c))) + .collect(); + Arc::new(StructArray::new( + s.fields().clone(), + columns, + s.nulls().cloned(), + )) + } + DataType::List(field) => gc_list!(field, i32, ListArray), + DataType::LargeList(field) => gc_list!(field, i64, LargeListArray), + DataType::FixedSizeList(field, size) => { + let list = array.as_fixed_size_list(); + Arc::new(FixedSizeListArray::new( + Arc::clone(field), + *size, + ScalarValue::compact_view_buffers(Arc::clone(list.values())), + list.nulls().cloned(), + )) + } + DataType::ListView(field) => gc_list_view!(field, i32, ListViewArray), + DataType::LargeListView(field) => { + gc_list_view!(field, i64, LargeListViewArray) + } + DataType::Map(field, ordered) => { + let map = array.as_map(); + let entries = ScalarValue::compact_view_buffers(Arc::new( + map.entries().clone(), + ) + as ArrayRef); + Arc::new(MapArray::new( + Arc::clone(field), + map.offsets().clone(), + entries.as_struct().clone(), + map.nulls().cloned(), + *ordered, + )) + } + _ => array, + } + } + /// Returns the minimum value for the given numeric `DataType`. /// /// This function returns the smallest representable value for numeric @@ -4836,28 +5106,21 @@ impl ScalarValue { DataType::Float16 => Some(ScalarValue::Float16(Some(f16::NEG_INFINITY))), DataType::Float32 => Some(ScalarValue::Float32(Some(f32::NEG_INFINITY))), DataType::Float64 => Some(ScalarValue::Float64(Some(f64::NEG_INFINITY))), + DataType::Decimal32(precision, scale) => { + let min = MIN_DECIMAL32_FOR_EACH_PRECISION[*precision as usize]; + Some(ScalarValue::Decimal32(Some(min), *precision, *scale)) + } + DataType::Decimal64(precision, scale) => { + let min = MIN_DECIMAL64_FOR_EACH_PRECISION[*precision as usize]; + Some(ScalarValue::Decimal64(Some(min), *precision, *scale)) + } DataType::Decimal128(precision, scale) => { - // For decimal, min is -10^(precision-scale) + 10^(-scale) - // But for simplicity, we use the minimum i128 value that fits the precision - let max_digits = 10_i128.pow(*precision as u32) - 1; - Some(ScalarValue::Decimal128( - Some(-max_digits), - *precision, - *scale, - )) + let min = MIN_DECIMAL128_FOR_EACH_PRECISION[*precision as usize]; + Some(ScalarValue::Decimal128(Some(min), *precision, *scale)) } DataType::Decimal256(precision, scale) => { - // Similar to Decimal128 but with i256 - // For now, use a large negative value - let max_digits = i256::from_i128(10_i128) - .checked_pow(*precision as u32) - .and_then(|v| v.checked_sub(i256::from_i128(1))) - .unwrap_or(i256::MAX); - Some(ScalarValue::Decimal256( - Some(max_digits.neg_wrapping()), - *precision, - *scale, - )) + let min = MIN_DECIMAL256_FOR_EACH_PRECISION[*precision as usize]; + Some(ScalarValue::Decimal256(Some(min), *precision, *scale)) } DataType::Date32 => Some(ScalarValue::Date32(Some(i32::MIN))), DataType::Date64 => Some(ScalarValue::Date64(Some(i64::MIN))), @@ -4932,27 +5195,21 @@ impl ScalarValue { DataType::Float16 => Some(ScalarValue::Float16(Some(f16::INFINITY))), DataType::Float32 => Some(ScalarValue::Float32(Some(f32::INFINITY))), DataType::Float64 => Some(ScalarValue::Float64(Some(f64::INFINITY))), + DataType::Decimal32(precision, scale) => { + let max = MAX_DECIMAL32_FOR_EACH_PRECISION[*precision as usize]; + Some(ScalarValue::Decimal32(Some(max), *precision, *scale)) + } + DataType::Decimal64(precision, scale) => { + let max = MAX_DECIMAL64_FOR_EACH_PRECISION[*precision as usize]; + Some(ScalarValue::Decimal64(Some(max), *precision, *scale)) + } DataType::Decimal128(precision, scale) => { - // For decimal, max is 10^(precision-scale) - 10^(-scale) - // But for simplicity, we use the maximum i128 value that fits the precision - let max_digits = 10_i128.pow(*precision as u32) - 1; - Some(ScalarValue::Decimal128( - Some(max_digits), - *precision, - *scale, - )) + let max = MAX_DECIMAL128_FOR_EACH_PRECISION[*precision as usize]; + Some(ScalarValue::Decimal128(Some(max), *precision, *scale)) } DataType::Decimal256(precision, scale) => { - // Similar to Decimal128 but with i256 - let max_digits = i256::from_i128(10_i128) - .checked_pow(*precision as u32) - .and_then(|v| v.checked_sub(i256::from_i128(1))) - .unwrap_or(i256::MAX); - Some(ScalarValue::Decimal256( - Some(max_digits), - *precision, - *scale, - )) + let max = MAX_DECIMAL256_FOR_EACH_PRECISION[*precision as usize]; + Some(ScalarValue::Decimal256(Some(max), *precision, *scale)) } DataType::Date32 => Some(ScalarValue::Date32(Some(i32::MAX))), DataType::Date64 => Some(ScalarValue::Date64(Some(i64::MAX))), @@ -5048,7 +5305,8 @@ impl ScalarValue { /// as necessary. pub fn copy_array_data(src_data: &ArrayData) -> ArrayData { let mut copy = MutableArrayData::new(vec![&src_data], true, src_data.len()); - copy.extend(0, 0, src_data.len()); + copy.try_extend(0, 0, src_data.len()) + .expect("copy_array_data failed due to offset overflow"); copy.freeze() } @@ -5259,6 +5517,35 @@ macro_rules! format_option { }}; } +macro_rules! format_decimal { + ($F:expr, $TYPE:ty, $VALUE:expr, $PRECISION:expr, $SCALE:expr) => {{ + match $VALUE { + Some(value) => write!( + $F, + "{}", + <$TYPE>::format_decimal(*value, *$PRECISION, *$SCALE) + ), + None => write!($F, "NULL"), + } + }}; +} + +macro_rules! format_decimal_debug { + ($F:expr, $TYPE_NAME:literal, $TYPE:ty, $VALUE:expr, $PRECISION:expr, $SCALE:expr) => {{ + match $VALUE { + Some(value) => write!( + $F, + "{}({},{},{})", + $TYPE_NAME, + <$TYPE>::format_decimal(*value, *$PRECISION, *$SCALE), + $PRECISION, + $SCALE + ), + None => write!($F, "{}(NULL,{},{})", $TYPE_NAME, $PRECISION, $SCALE), + } + }}; +} + // Implement Display trait for ScalarValue // // # Panics @@ -5268,16 +5555,16 @@ impl fmt::Display for ScalarValue { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { ScalarValue::Decimal32(v, p, s) => { - write!(f, "{v:?},{p:?},{s:?}")?; + format_decimal!(f, Decimal32Type, v, p, s)? } ScalarValue::Decimal64(v, p, s) => { - write!(f, "{v:?},{p:?},{s:?}")?; + format_decimal!(f, Decimal64Type, v, p, s)? } ScalarValue::Decimal128(v, p, s) => { - write!(f, "{v:?},{p:?},{s:?}")?; + format_decimal!(f, Decimal128Type, v, p, s)? } ScalarValue::Decimal256(v, p, s) => { - write!(f, "{v:?},{p:?},{s:?}")?; + format_decimal!(f, Decimal256Type, v, p, s)? } ScalarValue::Boolean(e) => format_option!(f, e)?, ScalarValue::Float16(e) => format_option!(f, e)?, @@ -5451,8 +5738,9 @@ fn fmt_list(arr: &dyn Array, f: &mut fmt::Formatter) -> fmt::Result { write!(f, "{value_formatter}") } -/// writes a byte array to formatter. `[1, 2, 3]` ==> `"1,2,3"` -fn fmt_binary(data: &[u8], f: &mut fmt::Formatter) -> fmt::Result { +/// Writes a byte array for ScalarValue Debug formatting. +/// `[1, 2, 3]` -> `"1,2,3"` +fn fmt_binary_debug(data: &[u8], f: &mut fmt::Formatter) -> fmt::Result { let mut iter = data.iter(); if let Some(b) = iter.next() { write!(f, "{b}")?; @@ -5466,10 +5754,46 @@ fn fmt_binary(data: &[u8], f: &mut fmt::Formatter) -> fmt::Result { impl fmt::Debug for ScalarValue { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { - ScalarValue::Decimal32(_, _, _) => write!(f, "Decimal32({self})"), - ScalarValue::Decimal64(_, _, _) => write!(f, "Decimal64({self})"), - ScalarValue::Decimal128(_, _, _) => write!(f, "Decimal128({self})"), - ScalarValue::Decimal256(_, _, _) => write!(f, "Decimal256({self})"), + ScalarValue::Decimal32(value, precision, scale) => { + format_decimal_debug!( + f, + "Decimal32", + Decimal32Type, + value, + precision, + scale + ) + } + ScalarValue::Decimal64(value, precision, scale) => { + format_decimal_debug!( + f, + "Decimal64", + Decimal64Type, + value, + precision, + scale + ) + } + ScalarValue::Decimal128(value, precision, scale) => { + format_decimal_debug!( + f, + "Decimal128", + Decimal128Type, + value, + precision, + scale + ) + } + ScalarValue::Decimal256(value, precision, scale) => { + format_decimal_debug!( + f, + "Decimal256", + Decimal256Type, + value, + precision, + scale + ) + } ScalarValue::Boolean(_) => write!(f, "Boolean({self})"), ScalarValue::Float16(_) => write!(f, "Float16({self})"), ScalarValue::Float32(_) => write!(f, "Float32({self})"), @@ -5503,13 +5827,13 @@ impl fmt::Debug for ScalarValue { ScalarValue::Binary(None) => write!(f, "Binary({self})"), ScalarValue::Binary(Some(b)) => { write!(f, "Binary(\"")?; - fmt_binary(b.as_slice(), f)?; + fmt_binary_debug(b.as_slice(), f)?; write!(f, "\")") } ScalarValue::BinaryView(None) => write!(f, "BinaryView({self})"), ScalarValue::BinaryView(Some(b)) => { write!(f, "BinaryView(\"")?; - fmt_binary(b.as_slice(), f)?; + fmt_binary_debug(b.as_slice(), f)?; write!(f, "\")") } ScalarValue::FixedSizeBinary(size, None) => { @@ -5517,13 +5841,13 @@ impl fmt::Debug for ScalarValue { } ScalarValue::FixedSizeBinary(size, Some(b)) => { write!(f, "FixedSizeBinary({size}, \"")?; - fmt_binary(b.as_slice(), f)?; + fmt_binary_debug(b.as_slice(), f)?; write!(f, "\")") } ScalarValue::LargeBinary(None) => write!(f, "LargeBinary({self})"), ScalarValue::LargeBinary(Some(b)) => { write!(f, "LargeBinary(\"")?; - fmt_binary(b.as_slice(), f)?; + fmt_binary_debug(b.as_slice(), f)?; write!(f, "\")") } ScalarValue::FixedSizeList(_) => write!(f, "FixedSizeList({self})"), @@ -8652,6 +8976,80 @@ mod tests { ScalarValue::from("larger than 12 bytes string"), DataType::Utf8View, ); + + // Cases also covered by `try_cast_literal_to_type` in datafusion-expr-common + + // identity casts (exercise the no-conversion fast path in `cast_to`) + check_scalar_cast(ScalarValue::Int32(Some(5)), DataType::Int32); + check_scalar_cast(ScalarValue::from("foo"), DataType::Utf8); + check_scalar_cast(ScalarValue::Utf8(None), DataType::Utf8); + check_scalar_cast( + ScalarValue::Dictionary( + Box::new(DataType::Int32), + Box::new(ScalarValue::from("foo")), + ), + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)), + ); + + // integer widening / narrowing (in range) + check_scalar_cast(ScalarValue::Int32(Some(123)), DataType::Int64); + check_scalar_cast(ScalarValue::Int64(Some(123)), DataType::Int32); + check_scalar_cast(ScalarValue::UInt32(Some(123)), DataType::Int64); + check_scalar_cast(ScalarValue::Int32(Some(123)), DataType::UInt64); + + // integer <-> decimal + check_scalar_cast(ScalarValue::Int32(Some(123)), DataType::Decimal128(10, 0)); + check_scalar_cast(ScalarValue::Decimal128(Some(123), 3, 0), DataType::Int64); + // decimal rescale + check_scalar_cast( + ScalarValue::Decimal128(Some(12300), 5, 2), + DataType::Decimal128(8, 5), + ); + + // timestamp unit conversion + check_scalar_cast( + ScalarValue::TimestampNanosecond(Some(123456), None), + DataType::Timestamp(TimeUnit::Microsecond, None), + ); + // timestamp timezone conversion + check_scalar_cast( + ScalarValue::TimestampSecond(Some(12345), None), + DataType::Timestamp(TimeUnit::Second, Some("+00:00".into())), + ); + // int64 <-> timestamp + check_scalar_cast( + ScalarValue::Int64(Some(12345)), + DataType::Timestamp(TimeUnit::Nanosecond, None), + ); + check_scalar_cast( + ScalarValue::TimestampSecond(Some(12345), Some("+00:00".into())), + DataType::Int64, + ); + + // additional string conversions + check_scalar_cast(ScalarValue::from("foo"), DataType::LargeUtf8); + check_scalar_cast(ScalarValue::LargeUtf8(Some("foo".into())), DataType::Utf8); + check_scalar_cast( + ScalarValue::LargeUtf8(Some("foo".into())), + DataType::Utf8View, + ); + check_scalar_cast(ScalarValue::Utf8View(Some("foo".into())), DataType::Utf8); + + // dictionary unwrap + check_scalar_cast( + ScalarValue::Dictionary( + Box::new(DataType::Int32), + Box::new(ScalarValue::from("foo")), + ), + DataType::Utf8, + ); + + // binary -> fixed size binary + check_scalar_cast( + ScalarValue::Binary(Some(vec![1, 2, 3])), + DataType::FixedSizeBinary(3), + ); + check_scalar_cast( { let element_field = @@ -8742,6 +9140,16 @@ mod tests { let cast_scalar = ScalarValue::try_from_array(&cast_array, 0).unwrap(); assert_eq!(cast_scalar.data_type(), desired_type); + // `ScalarValue::cast_to` (which has array-free fast paths) must produce + // exactly the same result as casting through the arrow kernel above. + let cast_to_scalar = scalar + .cast_to(&desired_type) + .expect("Failed to cast_to scalar"); + assert_eq!( + cast_to_scalar, cast_scalar, + "cast_to({scalar:?} -> {desired_type:?}) disagreed with the arrow cast kernel" + ); + // Some time later the "cast" scalar is turned back into an array: let array = cast_scalar .to_array_of_size(10) @@ -9122,8 +9530,8 @@ mod tests { ), ]; for (lhs, rhs, expected) in cases.iter() { - let distance = lhs.distance(rhs).unwrap(); - assert_eq!(distance, *expected); + let distance = lhs.distance_u64(rhs).unwrap(); + assert_eq!(distance, *expected as u64); } } @@ -9140,7 +9548,7 @@ mod tests { ), ]; for (lhs, rhs) in cases.iter() { - let distance = lhs.distance(rhs); + let distance = lhs.distance_u64(rhs); assert!(distance.is_none(), "{lhs} vs {rhs}"); } } @@ -9186,13 +9594,9 @@ mod tests { ScalarValue::Decimal128(Some(123), 5, 5), ScalarValue::Decimal128(Some(120), 5, 3), ), - ( - ScalarValue::Decimal128(Some(123), 5, 5), - ScalarValue::Decimal128(Some(120), 3, 5), - ), ( ScalarValue::Decimal256(Some(123.into()), 5, 5), - ScalarValue::Decimal256(Some(120.into()), 3, 5), + ScalarValue::Decimal256(Some(120.into()), 5, 3), ), // Distance 2 * 2^50 is larger than usize ( @@ -9214,11 +9618,124 @@ mod tests { ), ]; for (lhs, rhs) in cases { - let distance = lhs.distance(&rhs); + let distance = lhs.distance_u64(&rhs); assert!(distance.is_none()); } } + #[test] + fn test_scalar_distance_u64_boundaries() { + // 1. Full-domain integer ranges + // i64::MIN to i64::MAX -> distance is u64::MAX + let lhs = ScalarValue::Int64(Some(i64::MIN)); + let rhs = ScalarValue::Int64(Some(i64::MAX)); + assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); + assert_eq!(rhs.distance_u64(&lhs), Some(u64::MAX)); + + // u64::MIN to u64::MAX -> distance is u64::MAX + let lhs = ScalarValue::UInt64(Some(u64::MIN)); + let rhs = ScalarValue::UInt64(Some(u64::MAX)); + assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); + assert_eq!(rhs.distance_u64(&lhs), Some(u64::MAX)); + + // 2. Decimal128 overflow edges (around u64::MAX) + // distance equal to u64::MAX fits + let lhs = ScalarValue::Decimal128(Some(0), 20, 0); + let rhs = ScalarValue::Decimal128(Some(u64::MAX as i128), 20, 0); + assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); + + // distance greater than u64::MAX overflows + let lhs = ScalarValue::Decimal128(Some(0), 20, 0); + let rhs = ScalarValue::Decimal128(Some(u64::MAX as i128 + 1), 20, 0); + assert_eq!(lhs.distance_u64(&rhs), None); + + // 3. Decimal256 overflow edges (around u64::MAX) + // distance equal to u64::MAX fits + let lhs = ScalarValue::Decimal256(Some(i256::from_parts(0, 0)), 20, 0); + let rhs = + ScalarValue::Decimal256(Some(i256::from_parts(u64::MAX as u128, 0)), 20, 0); + assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); + + // distance greater than u64::MAX overflows + let lhs = ScalarValue::Decimal256(Some(i256::from_parts(0, 0)), 20, 0); + let rhs = ScalarValue::Decimal256( + Some(i256::from_parts(u64::MAX as u128 + 1, 0)), + 20, + 0, + ); + assert_eq!(lhs.distance_u64(&rhs), None); + + // 4. Float64 overflow edges (around u64::MAX) + let lhs = ScalarValue::Float64(Some(0.0)); + let val: f64 = 18446744073709500000.0; + let rhs = ScalarValue::Float64(Some(val)); + assert_eq!(lhs.distance_u64(&rhs), Some(18446744073709500416)); + + // float value > u64::MAX overflows + let rhs = ScalarValue::Float64(Some(1.9e19)); + assert_eq!(lhs.distance_u64(&rhs), None); + + // exact 2^64 boundary (18446744073709551616.0) is greater than u64::MAX, so it should return None + let exact_2_64_f64 = ScalarValue::Float64(Some(18446744073709551616.0)); + assert_eq!(lhs.distance_u64(&exact_2_64_f64), None); + + // exact 2^64 boundary as Float32 should also return None + let lhs_f32 = ScalarValue::Float32(Some(0.0)); + let exact_2_64_f32 = ScalarValue::Float32(Some(18446744073709551616.0)); + assert_eq!(lhs_f32.distance_u64(&exact_2_64_f32), None); + + // largest float32 value below 2^64 (2^64 - 2^41 = 18446741874686296064.0) should fit + let below_2_64_f32 = ScalarValue::Float32(Some(18446741874686296064.0)); + assert_eq!( + lhs_f32.distance_u64(&below_2_64_f32), + Some(18446741874686296064) + ); + + // Inf, NegInf, NaN + let inf = ScalarValue::Float64(Some(f64::INFINITY)); + let neg_inf = ScalarValue::Float64(Some(f64::NEG_INFINITY)); + let nan = ScalarValue::Float64(Some(f64::NAN)); + assert_eq!(lhs.distance_u64(&inf), None); + assert_eq!(lhs.distance_u64(&neg_inf), None); + assert_eq!(lhs.distance_u64(&nan), None); + + let inf_f32 = ScalarValue::Float32(Some(f32::INFINITY)); + let neg_inf_f32 = ScalarValue::Float32(Some(f32::NEG_INFINITY)); + let nan_f32 = ScalarValue::Float32(Some(f32::NAN)); + assert_eq!(lhs_f32.distance_u64(&inf_f32), None); + assert_eq!(lhs_f32.distance_u64(&neg_inf_f32), None); + assert_eq!(lhs_f32.distance_u64(&nan_f32), None); + + let lhs_f16 = ScalarValue::Float16(Some(f16::ZERO)); + let inf_f16 = ScalarValue::Float16(Some(f16::INFINITY)); + let neg_inf_f16 = ScalarValue::Float16(Some(f16::NEG_INFINITY)); + let nan_f16 = ScalarValue::Float16(Some(f16::NAN)); + assert_eq!(lhs_f16.distance_u64(&inf_f16), None); + assert_eq!(lhs_f16.distance_u64(&neg_inf_f16), None); + assert_eq!(lhs_f16.distance_u64(&nan_f16), None); + + // 5. Date and Timestamp boundaries + // Date32: i32::MIN to i32::MAX + let lhs = ScalarValue::Date32(Some(i32::MIN)); + let rhs = ScalarValue::Date32(Some(i32::MAX)); + assert_eq!(lhs.distance_u64(&rhs), Some(u32::MAX as u64)); + + // TimestampSecond: i64::MIN to i64::MAX + let lhs = ScalarValue::TimestampSecond(Some(i64::MIN), None); + let rhs = ScalarValue::TimestampSecond(Some(i64::MAX), None); + assert_eq!(lhs.distance_u64(&rhs), Some(u64::MAX)); + + // 6. Decimal scale matching (ignoring precision) + let lhs = ScalarValue::Decimal128(Some(100), 10, 2); + let rhs = ScalarValue::Decimal128(Some(150), 15, 2); + assert_eq!(lhs.distance_u64(&rhs), Some(50)); + assert_eq!(rhs.distance_u64(&lhs), Some(50)); + + let lhs = ScalarValue::Decimal128(Some(100), 10, 2); + let rhs = ScalarValue::Decimal128(Some(150), 10, 3); + assert_eq!(lhs.distance_u64(&rhs), None); + } + #[test] fn test_scalar_interval_negate() { let cases = [ @@ -9516,6 +10033,36 @@ mod tests { ); } + #[test] + fn test_decimal_display_and_debug() { + let decimal32 = ScalarValue::Decimal32(Some(123), 3, 2); + assert_eq!(decimal32.to_string(), "1.23"); + assert_eq!(format!("{decimal32:?}"), "Decimal32(1.23,3,2)"); + + let decimal64 = ScalarValue::Decimal64(Some(-12345), 5, 3); + assert_eq!(decimal64.to_string(), "-12.345"); + assert_eq!(format!("{decimal64:?}"), "Decimal64(-12.345,5,3)"); + + let decimal128 = ScalarValue::Decimal128(Some(1), 1, 1); + assert_eq!(decimal128.to_string(), "0.1"); + assert_eq!(format!("{decimal128:?}"), "Decimal128(0.1,1,1)"); + + let decimal128_trailing_zero = ScalarValue::Decimal128(Some(120), 3, 2); + assert_eq!(decimal128_trailing_zero.to_string(), "1.20"); + assert_eq!( + format!("{decimal128_trailing_zero:?}"), + "Decimal128(1.20,3,2)" + ); + + let decimal256 = ScalarValue::Decimal256(Some(i256::from(100123)), 28, 3); + assert_eq!(decimal256.to_string(), "100.123"); + assert_eq!(format!("{decimal256:?}"), "Decimal256(100.123,28,3)"); + + let null_decimal = ScalarValue::Decimal128(None, 10, 2); + assert_eq!(null_decimal.to_string(), "NULL"); + assert_eq!(format!("{null_decimal:?}"), "Decimal128(NULL,10,2)"); + } + #[test] fn test_struct_display_null() { let fields = vec![Field::new("a", DataType::Int32, false)]; @@ -9845,6 +10392,55 @@ mod tests { ); } + #[test] + fn safe_cast_date_to_timestamp_overflow_returns_null() { + let scalar = ScalarValue::Date32(Some(i32::MAX)); + let safe_options = CastOptions { + safe: true, + ..DEFAULT_CAST_OPTIONS + }; + + let casted = scalar + .cast_to_with_options( + &DataType::Timestamp(TimeUnit::Nanosecond, None), + &safe_options, + ) + .expect("expected safe cast to return null"); + + assert_eq!(casted, ScalarValue::TimestampNanosecond(None, None)); + } + + #[test] + fn cast_timestamp_to_timestamp_overflow_returns_error() { + let scalar = ScalarValue::TimestampSecond(Some(i64::MAX), None); + let err = scalar + .cast_to(&DataType::Timestamp(TimeUnit::Nanosecond, None)) + .expect_err("expected cast to fail"); + assert!( + err.to_string() + .contains("converted value exceeds the representable i64 range"), + "unexpected error: {err}" + ); + } + + #[test] + fn safe_cast_timestamp_to_timestamp_overflow_returns_null() { + let scalar = ScalarValue::TimestampSecond(Some(i64::MAX), None); + let safe_options = CastOptions { + safe: true, + ..DEFAULT_CAST_OPTIONS + }; + + let casted = scalar + .cast_to_with_options( + &DataType::Timestamp(TimeUnit::Nanosecond, None), + &safe_options, + ) + .expect("expected safe cast to return null"); + + assert_eq!(casted, ScalarValue::TimestampNanosecond(None, None)); + } + #[test] fn null_dictionary_scalar_produces_null_dictionary_array() { let dictionary_scalar = ScalarValue::Dictionary( @@ -10708,4 +11304,310 @@ mod tests { ] ); } + + // ── compact / compact_view_buffers ─────────────────────────────────────── + + /// Builds a `StringViewArray` with `n` strings that are all longer than + /// 12 bytes so they are stored in backing buffers rather than inline. + fn make_long_strings(n: usize) -> StringViewArray { + let mut b = StringViewBuilder::new(); + for i in 0..n { + b.append_value(format!("long_string_value_pad_{i:04}")); + } + b.finish() + } + + /// Total bytes across all backing buffers of a `StringViewArray`. + fn utf8view_buffer_bytes(a: &StringViewArray) -> usize { + a.data_buffers().iter().map(|b| b.len()).sum() + } + + #[test] + fn test_compact_list_utf8view() { + const N: usize = 50; + let strings = make_long_strings(N); + let one_len = strings.value(0).len(); + assert!(utf8view_buffer_bytes(&strings) >= N * one_len); + + let single_row_list_array = + SingleRowListArrayBuilder::new(Arc::new(strings.slice(0, 1)) as ArrayRef) + .build_list_array(); + let mut scalar = ScalarValue::List(Arc::new(single_row_list_array)); + scalar.compact(); + + let ScalarValue::List(arr) = &scalar else { + panic!("expected List") + }; + assert_eq!( + utf8view_buffer_bytes(arr.values().as_string_view()), + one_len + ); + assert_eq!(arr.values().as_string_view().value(0), strings.value(0)); + } + + #[test] + fn test_compact_large_list_utf8view() { + const N: usize = 50; + let strings = make_long_strings(N); + let one_len = strings.value(0).len(); + assert!(utf8view_buffer_bytes(&strings) >= N * one_len); + + let single_row_list_array = + SingleRowListArrayBuilder::new(Arc::new(strings.slice(0, 1)) as ArrayRef) + .build_large_list_array(); + let mut scalar = ScalarValue::LargeList(Arc::new(single_row_list_array)); + scalar.compact(); + + let ScalarValue::LargeList(arr) = &scalar else { + panic!("expected LargeList") + }; + assert_eq!( + utf8view_buffer_bytes(arr.values().as_string_view()), + one_len + ); + assert_eq!(arr.values().as_string_view().value(0), strings.value(0)); + } + + #[test] + fn test_compact_fixed_size_list_utf8view() { + const N: usize = 50; + let strings = make_long_strings(N); + let one_len = strings.value(0).len(); + assert!(utf8view_buffer_bytes(&strings) >= N * one_len); + + let single_row_list_array = + SingleRowListArrayBuilder::new(Arc::new(strings.slice(0, 1)) as ArrayRef) + .build_fixed_size_list_array(1); + let mut scalar = ScalarValue::FixedSizeList(Arc::new(single_row_list_array)); + scalar.compact(); + + let ScalarValue::FixedSizeList(arr) = &scalar else { + panic!("expected FixedSizeList") + }; + assert_eq!( + utf8view_buffer_bytes(arr.values().as_string_view()), + one_len + ); + assert_eq!(arr.values().as_string_view().value(0), strings.value(0)); + } + + #[test] + fn test_compact_list_view_utf8view() { + const N: usize = 50; + let strings = make_long_strings(N); + let one_len = strings.value(0).len(); + assert!(utf8view_buffer_bytes(&strings) >= N * one_len); + + let single_row_list_array = + SingleRowListArrayBuilder::new(Arc::new(strings.slice(0, 1)) as ArrayRef) + .build_list_view_array(); + let mut scalar = ScalarValue::ListView(Arc::new(single_row_list_array)); + scalar.compact(); + + let ScalarValue::ListView(arr) = &scalar else { + panic!("expected ListView") + }; + assert_eq!( + utf8view_buffer_bytes(arr.values().as_string_view()), + one_len + ); + assert_eq!(arr.values().as_string_view().value(0), strings.value(0)); + } + + #[test] + fn test_compact_large_list_view_utf8view() { + const N: usize = 50; + let strings = make_long_strings(N); + let one_len = strings.value(0).len(); + assert!(utf8view_buffer_bytes(&strings) >= N * one_len); + + let single_row_list_array = + SingleRowListArrayBuilder::new(Arc::new(strings.slice(0, 1)) as ArrayRef) + .build_large_list_view_array(); + let mut scalar = ScalarValue::LargeListView(Arc::new(single_row_list_array)); + scalar.compact(); + + let ScalarValue::LargeListView(arr) = &scalar else { + panic!("expected LargeListView") + }; + assert_eq!( + utf8view_buffer_bytes(arr.values().as_string_view()), + one_len + ); + assert_eq!(arr.values().as_string_view().value(0), strings.value(0)); + } + + #[test] + fn test_compact_struct_utf8view() { + const N: usize = 50; + let strings = make_long_strings(N); + let one_len = strings.value(0).len(); + + let field = Arc::new(Field::new("name", DataType::Utf8View, true)); + let struct_arr = StructArray::new( + Fields::from(vec![Arc::clone(&field)]), + vec![Arc::new(strings.slice(0, 1)) as ArrayRef], + None, + ); + + let mut scalar = ScalarValue::Struct(Arc::new(struct_arr)); + scalar.compact(); + + let ScalarValue::Struct(arr) = &scalar else { + panic!("expected Struct") + }; + let col = arr.column(0).as_string_view(); + assert_eq!(utf8view_buffer_bytes(col), one_len); + assert_eq!(col.value(0), strings.value(0)); + } + + #[test] + fn test_compact_map_utf8view() { + const N: usize = 50; + let strings = make_long_strings(N); + let one_len = strings.value(0).len(); + + let key_field = Arc::new(Field::new("key", DataType::Utf8View, false)); + let val_field = Arc::new(Field::new("value", DataType::Int32, true)); + let entries = StructArray::new( + Fields::from(vec![Arc::clone(&key_field), Arc::clone(&val_field)]), + vec![ + Arc::new(strings.slice(0, 1)) as ArrayRef, + Arc::new(Int32Array::from(vec![1i32])) as ArrayRef, + ], + None, + ); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![key_field, val_field])), + false, + )); + let map = MapArray::new( + entries_field, + OffsetBuffer::new(vec![0i32, 1].into()), + entries, + None, + false, + ); + + let mut scalar = ScalarValue::Map(Arc::new(map)); + scalar.compact(); + + let ScalarValue::Map(arr) = &scalar else { + panic!("expected Map") + }; + let keys = arr.entries().column(0).as_string_view(); + assert_eq!(utf8view_buffer_bytes(keys), one_len); + assert_eq!(keys.value(0), strings.value(0)); + } + + #[test] + fn test_zero_size_fsl() { + let s = ScalarValue::new_default(&DataType::FixedSizeList( + Field::new("a", DataType::Int32, true).into(), + 0, + )) + .unwrap(); + assert_eq!(s.to_string(), "[]"); + } + + #[test] + fn test_decimal_value_bounds() { + fn run_tests() { + // 0.1111, 0.2222, etc. + let max_scale = D::TYPE_CONSTRUCTOR(D::MAX_PRECISION, D::MAX_SCALE); + // 1.111, 2.222, etc. + let max_scale_less_one = + D::TYPE_CONSTRUCTOR(D::MAX_PRECISION, D::MAX_SCALE - 1); + // 11.11, 22.22, etc. + let max_scale_less_two = + D::TYPE_CONSTRUCTOR(D::MAX_PRECISION, D::MAX_SCALE - 2); + + // Invalid (can't represent the value) + assert!(ScalarValue::new_one(&max_scale).is_err()); + assert!(ScalarValue::new_negative_one(&max_scale).is_err()); + assert!(ScalarValue::new_ten(&max_scale).is_err()); + assert!(ScalarValue::new_ten(&max_scale_less_one).is_err()); + + // Valid + let one = ScalarValue::Int32(Some(1)); + let neg_one = ScalarValue::Int32(Some(-1)); + let ten = ScalarValue::Int32(Some(10)); + + let num = ScalarValue::new_one(&max_scale_less_one).unwrap(); + assert_eq!(num.cast_to(&DataType::Int32).unwrap(), one); + let num = ScalarValue::new_negative_one(&max_scale_less_one).unwrap(); + assert_eq!(num.cast_to(&DataType::Int32).unwrap(), neg_one); + let num = ScalarValue::new_ten(&max_scale_less_two).unwrap(); + assert_eq!(num.cast_to(&DataType::Int32).unwrap(), ten); + } + + run_tests::(); + run_tests::(); + run_tests::(); + run_tests::(); + } + + #[test] + fn test_new_list_nested_nullability_mismatch_issue_24022() { + // requested element type: Struct(n: Int32 nullable=true) + let requested_element_type = + DataType::Struct(Fields::from(vec![Field::new("n", DataType::Int32, true)])); + + // inferred from concrete values: Struct(n: Int32 nullable=false) + let inferred_field = Field::new("n", DataType::Int32, false); + + let value = ScalarValue::Struct(Arc::new(StructArray::from(vec![( + Arc::new(inferred_field), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]))); + + let expected_struct_array = StructArray::from(vec![( + Arc::new(Field::new("n", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]); + let expected_array = Arc::new(expected_struct_array) as ArrayRef; + + // Test new_list + let list = ScalarValue::new_list( + std::slice::from_ref(&value), + &requested_element_type, + true, + ); + assert_eq!( + list.data_type(), + &DataType::List(Arc::new(Field::new_list_field( + requested_element_type.clone(), + true + ))) + ); + assert_eq!(&list.value(0), &expected_array); + + // Test new_list_from_iter + let list_from_iter = ScalarValue::new_list_from_iter( + std::iter::once(value.clone()), + &requested_element_type, + true, + ); + assert_eq!( + list_from_iter.data_type(), + &DataType::List(Arc::new(Field::new_list_field( + requested_element_type.clone(), + true + ))) + ); + assert_eq!(&list_from_iter.value(0), &expected_array); + + // Test new_large_list + let large_list = ScalarValue::new_large_list(&[value], &requested_element_type); + assert_eq!( + large_list.data_type(), + &DataType::LargeList(Arc::new(Field::new( + "item", + requested_element_type.clone(), + true + ))) + ); + assert_eq!(&large_list.value(0), &expected_array); + } } diff --git a/datafusion/common/src/stats.rs b/datafusion/common/src/stats.rs index 320fd43751025..1a226c369884f 100644 --- a/datafusion/common/src/stats.rs +++ b/datafusion/common/src/stats.rs @@ -195,8 +195,12 @@ impl Precision { /// Return the estimate of applying a filter with estimated selectivity /// `selectivity` to this Precision. A selectivity of `1.0` means that all /// rows are selected. A selectivity of `0.5` means half the rows are - /// selected. Will always return inexact statistics. + /// selected. An exact zero is preserved, since filtering an empty input + /// cannot produce rows; any other known value is demoted to inexact. pub fn with_estimated_selectivity(self, selectivity: f64) -> Self { + if self == Precision::Exact(0) { + return self; + } self.map(|v| ((v as f64 * selectivity).ceil()) as usize) .to_inexact() } @@ -318,6 +322,12 @@ impl Precision { } } +impl From> for Precision { + fn from(option: Option) -> Self { + option.map_or(Precision::Absent, Precision::Exact) + } +} + impl Debug for Precision { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { @@ -417,7 +427,9 @@ impl Statistics { } /// Calculates `total_byte_size` based on the schema and `num_rows`. - /// If any of the columns has non-primitive width, `total_byte_size` is set to inexact. + /// If any of the columns has non-primitive width, or `num_rows` is unknown, + /// the previous `total_byte_size` is kept but downgraded to inexact rather + /// than discarded. pub fn calculate_total_byte_size(&mut self, schema: &Schema) { let mut row_size = Some(0); for field in schema.fields() { @@ -431,11 +443,11 @@ impl Statistics { } } } - match row_size { - None => { + match (row_size, &self.num_rows) { + (None, _) | (Some(_), Precision::Absent) => { self.total_byte_size = self.total_byte_size.to_inexact(); } - Some(size) => { + (Some(size), _) => { self.total_byte_size = self.num_rows.multiply(&Precision::Exact(size)); } } @@ -539,6 +551,10 @@ impl Statistics { skip: usize, n_partitions: usize, ) -> Result { + if fetch.is_none() && skip == 0 { + return Ok(self); + } + let fetch_val = fetch.unwrap_or(usize::MAX); // Get the ratio of rows after / rows before on a per-partition basis @@ -592,18 +608,18 @@ impl Statistics { .. } => check_num_rows(fetch.and_then(|v| v.checked_mul(n_partitions)), false), }; - let ratio: f64 = match (num_rows_before, self.num_rows) { + let ratio: Option = match (num_rows_before, self.num_rows) { ( Precision::Exact(nr_before) | Precision::Inexact(nr_before), Precision::Exact(nr_after) | Precision::Inexact(nr_after), ) => { if nr_before == 0 { - 0.0 + Some(0.0) } else { - nr_after as f64 / nr_before as f64 + Some(nr_after as f64 / nr_before as f64) } } - _ => 0.0, + _ => None, }; self.column_statistics = self .column_statistics @@ -611,11 +627,11 @@ impl Statistics { .map(|cs| { let mut cs = cs.to_inexact(); // Scale byte_size by the row ratio - cs.byte_size = match cs.byte_size { - Precision::Exact(n) | Precision::Inexact(n) => { + cs.byte_size = match (cs.byte_size, ratio) { + (Precision::Exact(n) | Precision::Inexact(n), Some(ratio)) => { Precision::Inexact((n as f64 * ratio) as usize) } - Precision::Absent => Precision::Absent, + _ => Precision::Absent, }; // NDV can never exceed the number of rows if let Some(&rows) = self.num_rows.get_value() { @@ -637,11 +653,11 @@ impl Statistics { Some(sum) => Precision::Inexact(sum), None => { // Fall back to scaling original total_byte_size if not all columns have byte_size - match &self.total_byte_size { - Precision::Exact(n) | Precision::Inexact(n) => { + match (&self.total_byte_size, ratio) { + (Precision::Exact(n) | Precision::Inexact(n), Some(ratio)) => { Precision::Inexact((*n as f64 * ratio) as usize) } - Precision::Absent => Precision::Absent, + _ => Precision::Absent, } } }; @@ -823,8 +839,8 @@ pub fn estimate_ndv_with_overlap( let right_min = right.min_value.get_value()?; let right_max = right.max_value.get_value()?; - let range_left = left_max.distance(left_min)?; - let range_right = right_max.distance(right_min)?; + let range_left = left_max.distance_u64(left_min)?; + let range_right = right_max.distance_u64(right_min)?; // Constant columns (range == 0) can't use the proportional overlap // formula below, so check interval overlap directly instead. @@ -853,7 +869,7 @@ pub fn estimate_ndv_with_overlap( return Some(ndv_left + ndv_right); } - let overlap_range = overlap_max.distance(overlap_min)? as f64; + let overlap_range = overlap_max.distance_u64(overlap_min)? as f64; let overlap_left = overlap_range / range_left as f64; let overlap_right = overlap_range / range_right as f64; @@ -1192,6 +1208,44 @@ mod tests { assert_eq!(absent_precision.get_value(), None); } + #[test] + fn test_with_estimated_selectivity() { + // Filtering an empty input cannot produce rows, so the zero stays exact. + assert_eq!( + Precision::Exact(0).with_estimated_selectivity(0.5), + Precision::Exact(0) + ); + assert_eq!( + Precision::Exact(0).with_estimated_selectivity(1.0), + Precision::Exact(0) + ); + + // Any other known value is scaled and demoted, since the selectivity is + // itself an estimate. + assert_eq!( + Precision::Exact(100).with_estimated_selectivity(0.5), + Precision::Inexact(50) + ); + assert_eq!( + Precision::Exact(100).with_estimated_selectivity(1.0), + Precision::Inexact(100) + ); + assert_eq!( + Precision::Exact(3).with_estimated_selectivity(0.5), + Precision::Inexact(2) + ); + + // An inexact zero is an estimate, not a proof, and stays inexact. + assert_eq!( + Precision::Inexact(0).with_estimated_selectivity(0.5), + Precision::Inexact(0) + ); + assert_eq!( + Precision::::Absent.with_estimated_selectivity(0.5), + Precision::Absent + ); + } + #[test] fn test_map() { let exact_precision = Precision::Exact(42); @@ -2370,6 +2424,38 @@ mod tests { assert_eq!(result.total_byte_size, Precision::Exact(800)); } + #[test] + fn test_with_fetch_no_limit_preserves_absent_num_rows() { + let original_stats = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Exact(800), + column_statistics: vec![col_stats_i64(10)], + }; + + let result = original_stats.clone().with_fetch(None, 0, 1).unwrap(); + + assert_eq!(result, original_stats); + } + + #[test] + fn test_with_fetch_absent_num_rows_does_not_zero_byte_size() { + let original_stats = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Exact(800), + column_statistics: vec![col_stats_i64(10)], + }; + + let result = original_stats.with_fetch(Some(1), 0, 1).unwrap(); + + assert_eq!(result.num_rows, Precision::Inexact(1)); + assert_eq!(result.total_byte_size, Precision::Absent); + assert_eq!(result.column_statistics[0].byte_size, Precision::Absent); + assert_eq!( + result.column_statistics[0].distinct_count, + Precision::Inexact(1) + ); + } + #[test] fn test_with_fetch_with_skip() { // Test with both skip and fetch @@ -3261,4 +3347,33 @@ mod tests { precision_add_for_sum_in_place(&mut lhs, &Precision::Absent); assert_eq!(lhs, Precision::Absent); } + + #[test] + fn test_calculate_total_byte_size() { + let primitive_schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let non_primitive_schema = + Schema::new(vec![Field::new("a", DataType::Utf8, false)]); + + // All-primitive schema with a known row count computes an exact size. + let mut stats = Statistics::new_unknown(&primitive_schema); + stats.num_rows = Precision::Exact(10); + stats.calculate_total_byte_size(&primitive_schema); + assert_eq!(stats.total_byte_size, Precision::Exact(40)); + + // All-primitive schema with an unknown row count keeps a previously + // known `total_byte_size`, downgraded to inexact, instead of + // discarding it to `Absent`. + let mut stats = Statistics::new_unknown(&primitive_schema); + stats.total_byte_size = Precision::Exact(1234); + stats.calculate_total_byte_size(&primitive_schema); + assert_eq!(stats.total_byte_size, Precision::Inexact(1234)); + + // Non-primitive schema always downgrades any existing + // `total_byte_size` to inexact, regardless of `num_rows`. + let mut stats = Statistics::new_unknown(&non_primitive_schema); + stats.num_rows = Precision::Exact(10); + stats.total_byte_size = Precision::Exact(999); + stats.calculate_total_byte_size(&non_primitive_schema); + assert_eq!(stats.total_byte_size, Precision::Inexact(999)); + } } diff --git a/datafusion/common/src/test_util.rs b/datafusion/common/src/test_util.rs index f060704944233..c0353ff408d70 100644 --- a/datafusion/common/src/test_util.rs +++ b/datafusion/common/src/test_util.rs @@ -174,7 +174,7 @@ macro_rules! assert_contains { } /// A macro to assert that one string is NOT contained within another with -/// a nice error message if they are are. +/// a nice error message if they are. /// /// Usage: `assert_not_contains!(actual, unexpected)` /// @@ -364,15 +364,20 @@ macro_rules! create_array { /// Creates a record batch from literal slice of values, suitable for rapid /// testing and development. /// +/// **Deprecated**: prefer the upstream macro from `arrow`, +/// [`arrow::array::record_batch`], which now supports both the literal slice +/// form shown below and a variable/expression form. +/// /// Example: /// ``` -/// use datafusion_common::record_batch; +/// use arrow::array::record_batch; /// let batch = record_batch!( /// ("a", Int32, vec![1, 2, 3]), /// ("b", Float64, vec![Some(4.0), None, Some(5.0)]), /// ("c", Utf8, vec!["alpha", "beta", "gamma"]) /// ); /// ``` +#[deprecated(since = "55.0.0", note = "Use `arrow::array::record_batch` instead")] #[macro_export] macro_rules! record_batch { ($(($name: expr, $type: ident, $values: expr)),*) => { @@ -616,7 +621,29 @@ pub mod array_conversion { } } - //#TODO add impl for f16 + impl IntoArrayRef for Vec { + fn into_array_ref(self) -> ArrayRef { + create_array!(Float16, self) + } + } + + impl IntoArrayRef for Vec> { + fn into_array_ref(self) -> ArrayRef { + create_array!(Float16, self) + } + } + + impl IntoArrayRef for &[half::f16] { + fn into_array_ref(self) -> ArrayRef { + create_array!(Float16, self.to_vec()) + } + } + + impl IntoArrayRef for &[Option] { + fn into_array_ref(self) -> ArrayRef { + create_array!(Float16, self.to_vec()) + } + } impl IntoArrayRef for Vec { fn into_array_ref(self) -> ArrayRef { @@ -776,6 +803,10 @@ mod tests { } #[test] + #[expect( + deprecated, + reason = "testing the deprecated record_batch! macro itself" + )] fn test_create_record_batch() -> Result<()> { use arrow::array::Array; diff --git a/datafusion/common/src/unnest.rs b/datafusion/common/src/unnest.rs index db48edd061605..58aed390ace78 100644 --- a/datafusion/common/src/unnest.rs +++ b/datafusion/common/src/unnest.rs @@ -19,23 +19,38 @@ use crate::Column; +/// How [`UnnestOptions`] handles `NULL` and empty list values in the input column. +/// +/// The variants enumerate the three observable behaviors so that callers do +/// not have to compose multiple boolean flags to express what they want. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Hash)] +pub enum NullHandling { + /// Drop rows where the input list is `NULL` or empty. Matches the + /// default behavior of systems such as DuckDB and ClickHouse. + Drop, + /// Preserve `NULL` input rows as a single output row containing `NULL`. + /// Empty lists still produce zero output rows. This is the default and + /// matches DataFusion's historical `preserve_nulls = true` behavior. + #[default] + Preserve, + /// Like [`Self::Preserve`], and additionally treat an empty list + /// identically to a `NULL` list, producing a single output row + /// containing `NULL`. + PreserveAndExpandEmpty, +} + /// Options for unnesting a column that contains a list type, /// replicating values in the other, non nested rows. /// /// Conceptually this operation is like joining each row with all the /// values in the list column. /// -/// If `preserve_nulls` is false, nulls and empty lists -/// from the input column are not carried through to the output. This -/// is the default behavior for other systems such as ClickHouse and -/// DuckDB -/// -/// If `preserve_nulls` is true (the default), nulls from the input -/// column are carried through to the output. +/// The behavior with `NULL` and empty input lists is controlled by +/// [`NullHandling`]. See its variants for full details. /// /// # Examples /// -/// ## `Unnest(c1)`, preserve_nulls: false +/// ## `Unnest(c1)`, null_handling: NullHandling::Drop /// ```text /// ┌─────────┐ ┌─────┐ ┌─────────┐ ┌─────┐ /// │ {1, 2} │ │ A │ Unnest │ 1 │ │ A │ @@ -49,7 +64,7 @@ use crate::Column; /// c1 c2 /// ``` /// -/// ## `Unnest(c1)`, preserve_nulls: true +/// ## `Unnest(c1)`, null_handling: NullHandling::Preserve /// ```text /// ┌─────────┐ ┌─────┐ ┌─────────┐ ┌─────┐ /// │ {1, 2} │ │ A │ Unnest │ 1 │ │ A │ @@ -63,13 +78,30 @@ use crate::Column; /// c1 c2 c1 c2 /// ``` /// +/// ## `Unnest(c1)`, null_handling: NullHandling::PreserveAndExpandEmpty +/// ```text +/// ┌─────────┐ ┌─────┐ ┌─────────┐ ┌─────┐ +/// │ {1, 2} │ │ A │ Unnest │ 1 │ │ A │ +/// ├─────────┤ ├─────┤ ├─────────┤ ├─────┤ +/// │ null │ │ B │ │ 2 │ │ A │ +/// ├─────────┤ ├─────┤ ────────────▶ ├─────────┤ ├─────┤ +/// │ {} │ │ D │ │ null │ │ B │ +/// ├─────────┤ ├─────┤ ├─────────┤ ├─────┤ +/// │ {3} │ │ E │ │ null │ │ D │ +/// └─────────┘ └─────┘ ├─────────┤ ├─────┤ +/// c1 c2 │ 3 │ │ E │ +/// └─────────┘ └─────┘ +/// c1 c2 +/// ``` +/// /// `recursions` instruct how a column should be unnested (e.g unnesting a column multiple /// time, with depth = 1 and depth = 2). Any unnested column not being mentioned inside this /// options is inferred to be unnested with depth = 1 #[derive(Debug, Clone, PartialEq, PartialOrd, Hash, Eq)] pub struct UnnestOptions { - /// Should nulls in the input be preserved? Defaults to true - pub preserve_nulls: bool, + /// How to handle `NULL` and empty list values in the input column. + /// Defaults to [`NullHandling::Preserve`]. + pub null_handling: NullHandling, /// If specific columns need to be unnested multiple times (e.g at different depth), /// declare them here. Any unnested columns not being mentioned inside this option /// will be unnested with depth = 1 @@ -88,8 +120,7 @@ pub struct RecursionUnnestOption { impl Default for UnnestOptions { fn default() -> Self { Self { - // default to true to maintain backwards compatible behavior - preserve_nulls: true, + null_handling: NullHandling::Preserve, recursions: vec![], } } @@ -101,13 +132,41 @@ impl UnnestOptions { Default::default() } - /// Set the behavior with nulls in the input as described on - /// [`Self`] - pub fn with_preserve_nulls(mut self, preserve_nulls: bool) -> Self { - self.preserve_nulls = preserve_nulls; + /// Set the [`NullHandling`] mode used when unnesting `NULL` or empty + /// input lists. + pub fn with_null_handling(mut self, null_handling: NullHandling) -> Self { + self.null_handling = null_handling; self } + /// Backward-compatible setter that maps the previous boolean + /// `preserve_nulls` flag onto [`NullHandling`]. + /// + /// `true` maps to [`NullHandling::Preserve`]; `false` maps to + /// [`NullHandling::Drop`]. To opt into the new empty-list-preserving + /// mode, call [`Self::with_null_handling`] directly with + /// [`NullHandling::PreserveAndExpandEmpty`]. + pub fn with_preserve_nulls(self, preserve_nulls: bool) -> Self { + let null_handling = if preserve_nulls { + NullHandling::Preserve + } else { + NullHandling::Drop + }; + self.with_null_handling(null_handling) + } + + /// Returns true if `NULL` input rows produce a single output row + /// containing `NULL`. + pub fn preserve_nulls(&self) -> bool { + !matches!(self.null_handling, NullHandling::Drop) + } + + /// Returns true if empty input lists should produce a single + /// output row containing `NULL`. + pub fn expand_empty_as_null(&self) -> bool { + matches!(self.null_handling, NullHandling::PreserveAndExpandEmpty) + } + /// Set the recursions for the unnest operation pub fn with_recursions(mut self, recursion: RecursionUnnestOption) -> Self { self.recursions.push(recursion); diff --git a/datafusion/common/src/utils/hex.rs b/datafusion/common/src/utils/hex.rs new file mode 100644 index 0000000000000..872d54f40c6f7 --- /dev/null +++ b/datafusion/common/src/utils/hex.rs @@ -0,0 +1,397 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Hex encoding of bytes and integers. +//! +//! [`encode_bytes`] and [`encode_bytes_into`] encode a byte slice into an +//! owned `String` or an appended `Vec`, respectively; [`encode_bytes_to_slice`] +//! writes into a caller-provided, pre-sized buffer. [`encode_u64`] encodes an +//! integer, trimming leading zeros. All four take a [`HexCase`] to choose +//! between lowercase and uppercase digits. + +use arrow::datatypes::ArrowNativeType; + +use crate::Result; +use crate::error::_internal_err; + +/// Case of the emitted hex digits. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HexCase { + /// Digits `0123456789abcdef`. + Lower, + /// Digits `0123456789ABCDEF`. + Upper, +} + +const LOWER_DIGITS: &[u8; 16] = b"0123456789abcdef"; +const UPPER_DIGITS: &[u8; 16] = b"0123456789ABCDEF"; + +/// Maps a full byte to its two hex digits, so encoding advances a whole byte +/// per iteration instead of a nibble. +const LOOKUP_LOWER: [[u8; 2]; 256] = build_lookup(LOWER_DIGITS); +const LOOKUP_UPPER: [[u8; 2]; 256] = build_lookup(UPPER_DIGITS); + +const fn build_lookup(digits: &[u8; 16]) -> [[u8; 2]; 256] { + let mut table = [[0u8; 2]; 256]; + let mut i = 0; + while i < 256 { + table[i][0] = digits[i >> 4]; + table[i][1] = digits[i & 0xF]; + i += 1; + } + table +} + +impl HexCase { + #[inline] + const fn lookup(self) -> &'static [[u8; 2]; 256] { + match self { + HexCase::Lower => &LOOKUP_LOWER, + HexCase::Upper => &LOOKUP_UPPER, + } + } + + #[inline] + const fn digits(self) -> &'static [u8; 16] { + match self { + HexCase::Lower => LOWER_DIGITS, + HexCase::Upper => UPPER_DIGITS, + } + } +} + +/// Trait for converting integer types to hexadecimal in a buffer +pub trait ToHex: ArrowNativeType { + /// Writes the hex representation into `buf` and returns the written + /// subslice. Digits are right-aligned with leading zeros trimmed. + fn write_hex(self, case: HexCase, buf: &mut [u8; 16]) -> &[u8]; +} + +macro_rules! impl_to_hex_signed { + ($ty:ty) => { + impl ToHex for $ty { + #[inline(always)] + fn write_hex(self, case: HexCase, buf: &mut [u8; 16]) -> &[u8] { + encode_u64(self as i64 as u64, case, buf) + } + } + }; +} + +macro_rules! impl_to_hex_unsigned { + ($ty:ty) => { + impl ToHex for $ty { + #[inline(always)] + fn write_hex(self, case: HexCase, buf: &mut [u8; 16]) -> &[u8] { + encode_u64(self as u64, case, buf) + } + } + }; +} + +impl_to_hex_signed!(i8); +impl_to_hex_signed!(i16); +impl_to_hex_signed!(i32); +impl_to_hex_signed!(i64); +impl_to_hex_unsigned!(u8); +impl_to_hex_unsigned!(u16); +impl_to_hex_unsigned!(u32); +impl_to_hex_unsigned!(u64); + +/// Appends the hex encoding of `bytes` to `out`. +/// +/// Allocates only through `out`'s own growth. Callers that must bound or guard +/// that growth should reserve capacity in `out` before calling. +#[inline(always)] +pub fn encode_bytes_into(bytes: &[u8], case: HexCase, out: &mut Vec) { + let lookup = case.lookup(); + for &byte in bytes { + out.extend_from_slice(&lookup[byte as usize]); + } +} + +/// Writes the hex encoding of `bytes` into `out`. +/// +/// This is for callers that already own a pre-sized buffer (for example a +/// slice of a larger, pre-allocated output array) and want to write directly +/// into it rather than appending to a `Vec`. +/// +/// Returns an internal error if `out` is not exactly `2 * bytes.len()` bytes +/// long, without filling any of the `out` buffer. +/// +/// # Example +/// +/// ``` +/// use datafusion_common::utils::hex::{HexCase, encode_bytes_to_slice}; +/// +/// let mut out = [0u8; 8]; +/// encode_bytes_to_slice(&[0xde, 0xad, 0xbe, 0xef], HexCase::Lower, &mut out)?; +/// assert_eq!(&out, b"deadbeef"); +/// # Ok::<(), datafusion_common::DataFusionError>(()) +/// ``` +#[inline(always)] +pub fn encode_bytes_to_slice(bytes: &[u8], case: HexCase, out: &mut [u8]) -> Result<()> { + let expected = bytes.len() * 2; + if out.len() != expected { + return _internal_err!( + "hex output buffer is {} bytes, expected {expected}", + out.len() + ); + } + let lookup = case.lookup(); + for (&b, chunk) in bytes.iter().zip(out.chunks_exact_mut(2)) { + chunk.copy_from_slice(&lookup[b as usize]); + } + Ok(()) +} + +/// Returns the hex encoding of `bytes` as an owned `String`. +/// +/// # Example +/// +/// ``` +/// use datafusion_common::utils::hex::{HexCase, encode_bytes}; +/// +/// assert_eq!(encode_bytes(&[0xde, 0xad, 0xbe, 0xef], HexCase::Lower), "deadbeef"); +/// assert_eq!(encode_bytes(&[0xde, 0xad, 0xbe, 0xef], HexCase::Upper), "DEADBEEF"); +/// ``` +#[inline] +pub fn encode_bytes(bytes: &[u8], case: HexCase) -> String { + let mut out = Vec::with_capacity(bytes.len() * 2); + encode_bytes_into(bytes, case, &mut out); + // SAFETY: `out` holds only ASCII hex digits, which are valid UTF-8. + unsafe { String::from_utf8_unchecked(out) } +} + +/// Writes `v` as hex into `buf` and returns the written subslice. +/// +/// Digits are written right-aligned with leading zeros trimmed, so the result +/// borrows the tail of `buf`. Zero encodes as `"0"`. +/// +/// Signed values should be cast with `as u64`, which yields the two's +/// complement representation that both `to_hex` and Spark's `hex` produce for +/// negative input. +/// +/// # Example +/// +/// The caller owns the buffer and can reuse it across calls; each call +/// returns a fresh subslice of it, borrowed for as long as `buf` is: +/// +/// ``` +/// use datafusion_common::utils::hex::{HexCase, encode_u64}; +/// +/// let mut buf = [0u8; 16]; +/// assert_eq!(encode_u64(0xAB, HexCase::Lower, &mut buf), b"ab"); +/// assert_eq!(encode_u64(0, HexCase::Lower, &mut buf), b"0"); +/// ``` +#[inline(always)] +pub fn encode_u64(v: u64, case: HexCase, buf: &mut [u8; 16]) -> &[u8] { + let start = write_digits(v, case, buf); + &buf[start..] +} + +/// Writes the digits of `v` right-aligned in `buf`, returning the index of the +/// first digit. +/// +/// Split out from [`encode_u64`] so the mutable borrow of `buf` ends before the +/// returned slice reborrows it. +#[inline(always)] +fn write_digits(v: u64, case: HexCase, buf: &mut [u8; 16]) -> usize { + if v == 0 { + buf[15] = b'0'; + return 15; + } + + // Consume two nibbles (one full byte) per iteration. + let lookup = case.lookup(); + let mut pos = 16; + let mut rest = v; + while rest >= 0x10 { + pos -= 2; + let pair = lookup[(rest & 0xFF) as usize]; + buf[pos] = pair[0]; + buf[pos + 1] = pair[1]; + rest >>= 8; + } + if rest > 0 { + // A single high nibble (0x1..=0xF) remains. + pos -= 1; + buf[pos] = case.digits()[rest as usize]; + } + + pos +} + +#[cfg(test)] +mod tests { + use super::*; + + fn hex_u64(v: u64, case: HexCase) -> String { + let mut buf = [0u8; 16]; + String::from_utf8(encode_u64(v, case, &mut buf).to_vec()).unwrap() + } + + #[test] + fn encode_u64_zero() { + assert_eq!(hex_u64(0, HexCase::Lower), "0"); + assert_eq!(hex_u64(0, HexCase::Upper), "0"); + } + + #[test] + fn encode_u64_single_nibble() { + for v in 1..=0xFu64 { + assert_eq!(hex_u64(v, HexCase::Lower), format!("{v:x}")); + assert_eq!(hex_u64(v, HexCase::Upper), format!("{v:X}")); + } + } + + #[test] + fn encode_u64_digit_count_boundaries() { + // Straddle each odd/even digit-count boundary: the two-nibbles-per + // iteration loop plus the trailing single-nibble fixup. + for v in [ + 0x10u64, + 0xFF, + 0x100, + 0xFFF, + 0x1000, + 0xFFFFF, + 0xFFFF_FFFF, + 0x1_0000_0000, + ] { + assert_eq!(hex_u64(v, HexCase::Lower), format!("{v:x}")); + assert_eq!(hex_u64(v, HexCase::Upper), format!("{v:X}")); + } + } + + #[test] + fn encode_u64_max() { + assert_eq!(hex_u64(u64::MAX, HexCase::Lower), "ffffffffffffffff"); + assert_eq!(hex_u64(u64::MAX, HexCase::Upper), "FFFFFFFFFFFFFFFF"); + } + + #[test] + fn encode_u64_signed_is_twos_complement() { + // Callers cast signed values with `as u64`; this is the behaviour both + // `to_hex` and Spark `hex` rely on for negative input. + assert_eq!(hex_u64(-1i64 as u64, HexCase::Lower), "ffffffffffffffff"); + assert_eq!(hex_u64(i64::MIN as u64, HexCase::Upper), "8000000000000000"); + } + + #[test] + fn encode_bytes_empty() { + assert_eq!(encode_bytes(&[], HexCase::Lower), ""); + assert_eq!(encode_bytes(&[], HexCase::Upper), ""); + } + + #[test] + fn encode_bytes_examples() { + assert_eq!(encode_bytes(&[0x00], HexCase::Lower), "00"); + assert_eq!(encode_bytes(&[0xAB], HexCase::Lower), "ab"); + assert_eq!(encode_bytes(&[0xAB], HexCase::Upper), "AB"); + assert_eq!( + encode_bytes(&[0xde, 0xad, 0xbe, 0xef], HexCase::Lower), + "deadbeef" + ); + assert_eq!( + encode_bytes(&[0xde, 0xad, 0xbe, 0xef], HexCase::Upper), + "DEADBEEF" + ); + } + + #[test] + fn encode_bytes_covers_every_byte_value() { + let bytes: Vec = (0..=255u8).collect(); + + let expected: String = bytes.iter().map(|b| format!("{b:02x}")).collect(); + assert_eq!(encode_bytes(&bytes, HexCase::Lower), expected); + + let expected: String = bytes.iter().map(|b| format!("{b:02X}")).collect(); + assert_eq!(encode_bytes(&bytes, HexCase::Upper), expected); + } + + #[test] + fn encode_bytes_into_appends_without_clearing() { + let mut out = b"prefix-".to_vec(); + encode_bytes_into(&[0x01, 0x02], HexCase::Lower, &mut out); + assert_eq!(out, b"prefix-0102"); + } + + #[test] + fn encode_u64_reused_buffer_leaks_no_stale_digits() { + let mut buf = [0u8; 16]; + assert_eq!( + encode_u64(u64::MAX, HexCase::Lower, &mut buf), + b"ffffffffffffffff" + ); + assert_eq!(encode_u64(0, HexCase::Lower, &mut buf), b"0"); + assert_eq!(encode_u64(0xAB, HexCase::Lower, &mut buf), b"ab"); + } + + #[test] + fn encode_bytes_to_slice_empty() -> Result<()> { + let mut out: [u8; 0] = []; + encode_bytes_to_slice(&[], HexCase::Lower, &mut out)?; + assert_eq!(out, [] as [u8; 0]); + Ok(()) + } + + #[test] + fn encode_bytes_to_slice_examples() -> Result<()> { + let mut out = [0u8; 8]; + encode_bytes_to_slice(&[0xde, 0xad, 0xbe, 0xef], HexCase::Lower, &mut out)?; + assert_eq!(&out, b"deadbeef"); + + let mut out = [0u8; 8]; + encode_bytes_to_slice(&[0xde, 0xad, 0xbe, 0xef], HexCase::Upper, &mut out)?; + assert_eq!(&out, b"DEADBEEF"); + Ok(()) + } + + #[test] + fn encode_bytes_to_slice_agrees_with_encode_bytes() -> Result<()> { + let bytes: Vec = (0..=255u8).collect(); + for case in [HexCase::Lower, HexCase::Upper] { + let mut out = vec![0u8; bytes.len() * 2]; + encode_bytes_to_slice(&bytes, case, &mut out)?; + assert_eq!(String::from_utf8(out).unwrap(), encode_bytes(&bytes, case)); + } + Ok(()) + } + + #[test] + fn encode_bytes_to_slice_rejects_wrong_length() { + // Too short: the old `debug_assert` let release builds silently drop + // the remaining input. + let mut short = [0u8; 6]; + let err = + encode_bytes_to_slice(&[0xde, 0xad, 0xbe, 0xef], HexCase::Lower, &mut short) + .unwrap_err(); + assert!( + err.message() + .contains("hex output buffer is 6 bytes, expected 8"), + "unexpected message: {err}" + ); + + // Too long: would have left stale bytes at the tail. + let mut long = [0u8; 10]; + assert!( + encode_bytes_to_slice(&[0xde, 0xad, 0xbe, 0xef], HexCase::Lower, &mut long) + .is_err() + ); + } +} diff --git a/datafusion/common/src/utils/memory.rs b/datafusion/common/src/utils/memory.rs index 78ec434d2b577..21c084119e120 100644 --- a/datafusion/common/src/utils/memory.rs +++ b/datafusion/common/src/utils/memory.rs @@ -21,7 +21,8 @@ use crate::error::_exec_datafusion_err; use crate::{HashSet, Result}; use arrow::array::ArrayData; use arrow::record_batch::RecordBatch; -use std::{mem::size_of, ptr::NonNull}; +use std::mem::size_of; +use std::num::NonZero; /// Estimates the memory size required for a hash table prior to allocation. /// @@ -131,34 +132,74 @@ pub fn estimate_memory_size(num_elements: usize, fixed_size: usize) -> Result /// `Buffer`. This method provides temporary fix until the issue is resolved: /// pub fn get_record_batch_memory_size(batch: &RecordBatch) -> usize { - // Store pointers to `Buffer`'s start memory address (instead of actual - // used data region's pointer represented by current `Array`) - let mut counted_buffers: HashSet> = HashSet::new(); - let mut total_size = 0; - - for array in batch.columns() { - let array_data = array.to_data(); - count_array_data_memory_size(&array_data, &mut counted_buffers, &mut total_size); + RecordBatchMemoryCounter::new().count_batch(batch) +} + +/// Tracks the memory used by a sequence of [`RecordBatch`]es that may share +/// underlying buffers, counting each buffer exactly once. +/// +/// Use this instead of [`get_record_batch_memory_size`] to account for the +/// total memory of a sequence of batches, e.g. when buffering the batches of +/// an input stream. Such batches can share buffers (for example, operators +/// like aggregates emit one large batch as multiple zero-copy slices), and +/// calling [`get_record_batch_memory_size`] per batch counts the shared +/// buffers once per batch, while this counter counts them exactly once. A +/// batch's buffers are kept alive by the batch even when only a sub-range is +/// referenced, so counting unique buffers in full reflects the memory the +/// batches actually retain. +#[derive(Debug, Default)] +pub struct RecordBatchMemoryCounter { + /// Start addresses of `Buffer`s that have already been counted (instead of + /// actual used data region's pointer represented by current `Array`) + counted_buffers: HashSet>, + /// Total memory of all unique buffers counted so far + memory_usage: usize, +} + +impl RecordBatchMemoryCounter { + pub fn new() -> Self { + Self::default() } - total_size + /// Count `batch`, returning the memory used by its buffers that have not + /// been counted before. + pub fn count_batch(&mut self, batch: &RecordBatch) -> usize { + let mut total_size = 0; + + for array in batch.columns() { + let array_data = array.to_data(); + count_array_data_memory_size( + &array_data, + &mut self.counted_buffers, + &mut total_size, + ); + } + + self.memory_usage += total_size; + total_size + } + + /// Total memory of the unique buffers of all batches counted so far. + pub fn memory_usage(&self) -> usize { + self.memory_usage + } } /// Count the memory usage of `array_data` and its children recursively. fn count_array_data_memory_size( array_data: &ArrayData, - counted_buffers: &mut HashSet>, + counted_buffers: &mut HashSet>, total_size: &mut usize, ) { // Count memory usage for `array_data` for buffer in array_data.buffers() { - if counted_buffers.insert(buffer.data_ptr()) { + if counted_buffers.insert(buffer.data_ptr().addr()) { *total_size += buffer.capacity(); } // Otherwise the buffer's memory is already counted } if let Some(null_buffer) = array_data.nulls() - && counted_buffers.insert(null_buffer.inner().inner().data_ptr()) + && counted_buffers.insert(null_buffer.inner().inner().data_ptr().addr()) { *total_size += null_buffer.inner().inner().capacity(); } @@ -295,6 +336,29 @@ mod record_batch_tests { assert_eq!(size_origin, size_sliced); } + #[test] + fn test_record_batch_memory_counter_buffer_shared_across_batches() { + let schema = Arc::new(Schema::new(vec![Field::new( + "ints", + DataType::Int32, + false, + )])); + + let int_array = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let batch = RecordBatch::try_new(schema, vec![Arc::new(int_array)]).unwrap(); + let slices = [batch.slice(0, 2), batch.slice(2, 2), batch.slice(4, 2)]; + + // Counting each slice individually counts the shared buffer once per slice + let summed: usize = slices.iter().map(get_record_batch_memory_size).sum(); + assert_eq!(summed, 3 * get_record_batch_memory_size(&batch)); + + // A counter shared across the batches counts it exactly once + let mut counter = RecordBatchMemoryCounter::new(); + let deduped: usize = slices.iter().map(|slice| counter.count_batch(slice)).sum(); + assert_eq!(deduped, get_record_batch_memory_size(&batch)); + assert_eq!(counter.memory_usage(), get_record_batch_memory_size(&batch)); + } + #[test] fn test_get_record_batch_memory_size_nested_array() { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/common/src/utils/mod.rs b/datafusion/common/src/utils/mod.rs index 0c667b17c3fd9..73772b319351c 100644 --- a/datafusion/common/src/utils/mod.rs +++ b/datafusion/common/src/utils/mod.rs @@ -19,6 +19,7 @@ pub(crate) mod aggregate; pub mod expr; +pub mod hex; pub mod memory; pub mod proxy; pub mod string_utils; @@ -395,6 +396,137 @@ pub fn longest_consecutive_prefix>( count } +/// Splits `vec` at index `n`, returning the first `n` elements and leaving the +/// remaining `vec.len() - n` elements in `vec`. +/// +/// Allocates for whichever side is smaller, so the new allocation is +/// `min(n, vec.len() - n)` rather than always `n` (as `vec.drain(0..n).collect()` +/// would). This matters when the split emits a prefix under memory pressure, +/// where `n` can be close to `vec.len()`. +pub fn split_vec_min_alloc(vec: &mut Vec, n: usize) -> Vec { + if n * 2 <= vec.len() { + vec.drain(0..n).collect() + } else { + let remaining = vec.split_off(n); + std::mem::replace(vec, remaining) + } +} + +#[cfg(test)] +mod split_vec_min_alloc_tests { + use super::split_vec_min_alloc; + + #[test] + fn drain_branch() { + // n * 2 <= len -> drain+collect branch (allocates n elements) + let mut v = vec![1, 2, 3, 4, 5, 6]; + let first = split_vec_min_alloc(&mut v, 2); + assert_eq!(first, vec![1, 2]); + assert_eq!(v, vec![3, 4, 5, 6]); + } + + #[test] + fn split_off_branch() { + // remaining < n -> split_off+replace branch (allocates remaining elements) + let mut v = vec![1, 2, 3, 4, 5, 6]; + let first = split_vec_min_alloc(&mut v, 4); + assert_eq!(first, vec![1, 2, 3, 4]); + assert_eq!(v, vec![5, 6]); + } + + #[test] + fn exactly_half() { + // n * 2 == len -> drain branch (boundary) + let mut v = vec![1, 2, 3, 4]; + let first = split_vec_min_alloc(&mut v, 2); + assert_eq!(first, vec![1, 2]); + assert_eq!(v, vec![3, 4]); + } + + #[test] + fn take_all() { + let mut v = vec![1, 2, 3]; + let first = split_vec_min_alloc(&mut v, 3); + assert_eq!(first, vec![1, 2, 3]); + assert!(v.is_empty()); + } + + #[test] + fn take_none() { + let mut v = vec![1, 2, 3]; + let first = split_vec_min_alloc(&mut v, 0); + assert!(first.is_empty()); + assert_eq!(v, vec![1, 2, 3]); + } + + #[test] + fn emitted_prefix_does_not_realloc_on_push() { + // Demonstrates *why* the split-off branch must NOT call `shrink_to_fit`. + // + // Downstream callers (e.g. `multi_group_by/bytes.rs`, which does + // `first_n_offsets.push(offset_n)` right after the split) push onto the + // emitted prefix immediately. The split-off branch hands the original + // backing allocation to that prefix, so the prefix already has spare + // capacity for the very next push. + // + // If we shrank the prefix to fit, that next push would have to + // reallocate, and Vec's growth strategy would land it at a *larger* + // capacity than the original allocation we started with -- the opposite + // of the memory saving `shrink_to_fit` was meant to deliver. + + // A Vec with a known, deliberately large capacity. n*2 > len, so this + // takes the split-off branch. + let mut v: Vec = Vec::with_capacity(64); + v.extend(0..10); + let original_capacity = v.capacity(); + assert!(original_capacity >= 64); + + // Emit a prefix that is most of the Vec (n = 8, remaining = 2). + let mut prefix = split_vec_min_alloc(&mut v, 8); + assert_eq!(prefix, vec![0, 1, 2, 3, 4, 5, 6, 7]); + + // The split-off branch moved the original backing store into `prefix`, + // so it keeps the original (large) capacity -- no shrink happened. + assert_eq!( + prefix.capacity(), + original_capacity, + "split-off branch must hand the original allocation to the prefix" + ); + + // The caller's very next operation: push one element onto the prefix. + prefix.push(99); + + // Because the capacity was preserved, the push reused the existing + // allocation: post-push capacity is unchanged and still <= original. + // This is the realloc that `shrink_to_fit` would have forced. + assert_eq!( + prefix.capacity(), + original_capacity, + "push must reuse the preserved allocation (no realloc)" + ); + assert!(prefix.capacity() <= original_capacity); + + // Counter-demonstration: had we shrunk the prefix to fit (capacity 8), + // the same push would have reallocated. Vec doubles on growth, so the + // post-push capacity (16) ends up LARGER than where a length-8 prefix + // started -- and we paid a realloc for it. + let mut shrunk: Vec = prefix[..8].to_vec(); + shrunk.shrink_to_fit(); + let shrunk_capacity = shrink_then_push_capacity(&mut shrunk); + assert!( + shrunk_capacity > 8, + "shrink-to-fit then push reallocates to a larger capacity" + ); + } + + /// Helper for the counter-demonstration above: push one element and report + /// the resulting capacity. + fn shrink_then_push_capacity(v: &mut Vec) -> usize { + v.push(99); + v.capacity() + } +} + /// Creates single element [`ListArray`], [`LargeListArray`] and /// [`FixedSizeListArray`] from other arrays /// @@ -482,7 +614,8 @@ impl SingleRowListArrayBuilder { /// Build a single element [`FixedSizeListArray`] pub fn build_fixed_size_list_array(self, list_size: usize) -> FixedSizeListArray { let (field, arr) = self.into_field_and_arr(); - FixedSizeListArray::new(field, list_size as i32, arr, None) + FixedSizeListArray::try_new_with_length(field, list_size as i32, arr, None, 1) + .unwrap() } /// Build a single element [`FixedSizeListArray`] and wrap as [`ScalarValue::FixedSizeList`] @@ -1104,16 +1237,7 @@ pub fn adjust_offsets_for_slice( ) -> OffsetBuffer { let offsets = list.offsets(); - if let (Some(first), Some(last)) = (offsets.first(), offsets.last()) - && (!first.is_zero() || last.as_usize() != list.values().len()) - { - let offsets = offsets.iter().map(|offset| *offset - *first).collect(); - - //todo: use unsafe Offset::new_unchecked? - return OffsetBuffer::new(offsets); - } - - offsets.clone() + offsets.clone().subtract(offsets[0]) } /// For lists and large lists, truncates the sublist of null values @@ -1157,11 +1281,11 @@ fn truncate_list_nulls( let (valid_or_empty, _nulls) = valid_or_empty.into_parts(); for (start, end) in valid_or_empty.set_slices() { - mutable_array_data.extend( + mutable_array_data.try_extend( 0, offsets[start].as_usize(), offsets[end].as_usize(), - ); + )?; } let lengths = std::iter::zip(offsets.lengths(), nulls) @@ -1255,6 +1379,95 @@ fn fsl_values_row_number(list_size: i32, array_len: usize) -> Result Ok(PrimitiveArray::new(rows_number.into(), None)) } +/// Replace `-0.0` with `+0.0` in any `Float16`, `Float32`, or `Float64` array. +/// For non-float arrays returns the input unchanged. NaN payloads are +/// preserved. +/// +/// Arrow's comparison kernels (`arrow::compute::kernels::cmp::eq` etc.) and +/// row-encoding (`arrow::row::RowConverter`) use IEEE 754 totalOrder +/// semantics, which treats `-0.0` and `+0.0` as distinct. SQL semantics +/// (PostgreSQL / IEEE 754 equality) require them to compare equal, so +/// callers normalize before invoking those kernels. +/// +/// The common case - no `-0.0` present - is allocation-free: a single +/// read-only scan of the underlying buffer (auto-vectorizable to an +/// OR-reduction) decides whether to fall through to the rewriting path. +/// Only arrays that actually contain `-0.0` pay for a new buffer. +pub fn normalize_float_zero(array: &ArrayRef) -> ArrayRef { + use arrow::array::{Float16Array, Float32Array, Float64Array}; + use arrow::datatypes::{Float16Type, Float32Type, Float64Type}; + // -0.0 has only the sign bit set; no other finite or NaN value shares + // this bit pattern, so a strict-equality scan reliably gates the rewrite. + const NEG_ZERO_F16_BITS: u16 = half::f16::NEG_ZERO.to_bits(); + const NEG_ZERO_F32_BITS: u32 = (-0.0_f32).to_bits(); + const NEG_ZERO_F64_BITS: u64 = (-0.0_f64).to_bits(); + match array.data_type() { + DataType::Float32 => { + let arr: &Float32Array = array.as_primitive::(); + if !arr + .values() + .iter() + .any(|v| v.to_bits() == NEG_ZERO_F32_BITS) + { + return Arc::clone(array); + } + let normalized: Float32Array = + arr.unary(|v| if v.to_bits() << 1 == 0 { 0.0_f32 } else { v }); + Arc::new(normalized) + } + DataType::Float64 => { + let arr: &Float64Array = array.as_primitive::(); + if !arr + .values() + .iter() + .any(|v| v.to_bits() == NEG_ZERO_F64_BITS) + { + return Arc::clone(array); + } + let normalized: Float64Array = + arr.unary(|v| if v.to_bits() << 1 == 0 { 0.0_f64 } else { v }); + Arc::new(normalized) + } + DataType::Float16 => { + let arr: &Float16Array = array.as_primitive::(); + if !arr + .values() + .iter() + .any(|v| v.to_bits() == NEG_ZERO_F16_BITS) + { + return Arc::clone(array); + } + let normalized: Float16Array = arr.unary(|v| { + if v.to_bits() << 1 == 0 { + half::f16::from_bits(0) + } else { + v + } + }); + Arc::new(normalized) + } + _ => Arc::clone(array), + } +} + +/// Replace `-0.0` with `+0.0` in `Float16`, `Float32`, or `Float64` scalar +/// values. Other variants are returned unchanged. See [`normalize_float_zero`] +/// for context. +pub fn normalize_float_zero_scalar(scalar: ScalarValue) -> ScalarValue { + match scalar { + ScalarValue::Float32(Some(v)) if v.to_bits() << 1 == 0 => { + ScalarValue::Float32(Some(0.0)) + } + ScalarValue::Float64(Some(v)) if v.to_bits() << 1 == 0 => { + ScalarValue::Float64(Some(0.0)) + } + ScalarValue::Float16(Some(v)) if v.to_bits() << 1 == 0 => { + ScalarValue::Float16(Some(half::f16::from_bits(0))) + } + other => other, + } +} + #[cfg(test)] mod tests { use std::sync::Arc; @@ -1266,6 +1479,7 @@ mod tests { buffer::NullBuffer, datatypes::Int32Type, }; + #[cfg(feature = "sql")] use sqlparser::ast::Ident; #[test] diff --git a/datafusion/core/Cargo.toml b/datafusion/core/Cargo.toml index 55151caf2f8f0..0fe48ddf6a3e0 100644 --- a/datafusion/core/Cargo.toml +++ b/datafusion/core/Cargo.toml @@ -163,7 +163,7 @@ zstd = { workspace = true, optional = true } async-trait = { workspace = true } criterion = { workspace = true, features = ["async_tokio", "async_futures"] } ctor = { workspace = true } -dashmap = "6.1.0" +dashmap = "6.2.1" datafusion-doc = { workspace = true } datafusion-functions-window-common = { workspace = true } datafusion-macros = { workspace = true } @@ -173,13 +173,14 @@ bytes = { workspace = true } env_logger = { workspace = true } glob = { workspace = true } insta = { workspace = true } +half = { workspace = true } rand = { workspace = true, features = ["small_rng"] } rand_distr = "0.5" recursive = { workspace = true } regex = { workspace = true } rstest = { workspace = true } serde_json = { workspace = true } -sysinfo = "0.39.2" +sysinfo = "0.39.3" test-utils = { path = "../../test-utils" } tokio = { workspace = true, features = ["rt-multi-thread", "parking_lot", "fs"] } @@ -247,11 +248,26 @@ harness = false name = "parquet_struct_query" required-features = ["parquet"] +[[bench]] +harness = false +name = "parquet_nested_schema_pruning" +required-features = ["parquet"] + [[bench]] harness = false name = "parquet_struct_projection" required-features = ["parquet"] +[[bench]] +harness = false +name = "parquet_struct_shared_prefix_pushdown" +required-features = ["parquet"] + +[[bench]] +harness = false +name = "cse_projection_pushdown" +required-features = ["parquet"] + [[bench]] harness = false name = "range_and_generate_series" diff --git a/datafusion/core/benches/cse_projection_pushdown.rs b/datafusion/core/benches/cse_projection_pushdown.rs new file mode 100644 index 0000000000000..f5f9ec55e8912 --- /dev/null +++ b/datafusion/core/benches/cse_projection_pushdown.rs @@ -0,0 +1,184 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for the interaction between Common Subexpression Elimination +//! (CSE) and projection pushdown on parquet sources. +//! +//! Each query repeats a scalar function call several times, which the logical +//! CSE pass extracts into a single intermediate projection referenced by +//! column. These benchmarks measure the end-to-end cost of such queries, which +//! is dominated by how many times the extracted expression is ultimately +//! evaluated per row. + +use arrow::array::{Float64Array, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::instant::Instant; +use futures::stream::StreamExt; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use rand::prelude::*; +use rand::rng; +use std::sync::Arc; +use tempfile::NamedTempFile; + +const NUM_BATCHES: usize = 1024; +const BATCH_SIZE: usize = 1024; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Float64, false), + Field::new("c", DataType::Int64, false), + ])) +} + +fn generate_batch() -> RecordBatch { + let mut rng = rng(); + let len = BATCH_SIZE; + + let a: Float64Array = (0..len) + .map(|_| Some(rng.random_range(1.0..1000.0))) + .collect(); + let b: Float64Array = (0..len) + .map(|_| Some(rng.random_range(1.0..1000.0))) + .collect(); + let c: Int64Array = (0..len) + .map(|_| Some(rng.random_range(1i64..1000))) + .collect(); + + RecordBatch::try_new(schema(), vec![Arc::new(a), Arc::new(b), Arc::new(c)]).unwrap() +} + +fn generate_file() -> NamedTempFile { + let now = Instant::now(); + let mut named_file = tempfile::Builder::new() + .prefix("cse_projection_pushdown") + .suffix(".parquet") + .tempfile() + .unwrap(); + + println!("Generating parquet file - {}", named_file.path().display()); + + let props = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_max_row_group_row_count(Some(1024 * 1024)) + .build(); + + let mut writer = + ArrowWriter::try_new(&mut named_file, schema(), Some(props)).unwrap(); + + for _ in 0..NUM_BATCHES { + let batch = generate_batch(); + writer.write(&batch).unwrap(); + } + writer.close().unwrap(); + + println!( + "Generated parquet file in {} seconds", + now.elapsed().as_secs_f32() + ); + + named_file +} + +fn criterion_benchmark(c: &mut Criterion) { + let temp_file = generate_file(); + let file_path = temp_file.path().display().to_string(); + + let partitions = 4; + let config = SessionConfig::new().with_target_partitions(partitions); + let context = SessionContext::new_with_config(config); + + let local_rt = tokio::runtime::Builder::new_current_thread() + .build() + .unwrap(); + + let query_rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(partitions) + .build() + .unwrap(); + + local_rt + .block_on(context.register_parquet("t", file_path.as_str(), Default::default())) + .unwrap(); + + // Queries that repeat a scalar function call, which CSE extracts into a + // single intermediate projection referenced by column. + let queries = vec![ + // Same sqrt(a) appears 3 times. + ( + "repeated_sqrt", + "SELECT sqrt(a) + 1, sqrt(a) * 2, sqrt(a) / b FROM t", + ), + // power(a, 2) appears in multiple places. + ( + "repeated_power", + "SELECT power(a, 2) + b, power(a, 2) - b, power(a, 2) * c FROM t", + ), + // Deeper nesting: ln(abs(a)) repeated. + ( + "repeated_nested_fn", + "SELECT ln(abs(a)) + 1, ln(abs(a)) * b, ln(abs(a)) + c FROM t", + ), + // Mixed: some repeated, some unique. + ( + "mixed_repeated_unique", + "SELECT sqrt(a) + sqrt(a), abs(b), sqrt(a) * c FROM t", + ), + // A trivial function (abs) repeated. + ( + "repeated_cheap_abs", + "SELECT abs(a) + 1, abs(a) * 2, abs(a) / b FROM t", + ), + // Baseline: no repeated expressions (CSE does not fire). + ( + "no_repeated_exprs", + "SELECT sqrt(a), abs(b), power(a, 2) FROM t", + ), + ]; + + for (name, query) in queries { + c.bench_function(&format!("cse_pushdown: {name}"), |b| { + b.iter(|| { + let query = query.to_string(); + let context = context.clone(); + let (sender, mut receiver) = futures::channel::mpsc::unbounded(); + + query_rt.spawn(async move { + let query = context.sql(&query).await.unwrap(); + let mut stream = query.execute_stream().await.unwrap(); + + while let Some(next) = stream.next().await { + sender.unbounded_send(next).unwrap(); + } + }); + + local_rt.block_on(async { + while receiver.next().await.transpose().unwrap().is_some() {} + }) + }); + }); + } + + drop(temp_file); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/core/benches/filter_query_sql.rs b/datafusion/core/benches/filter_query_sql.rs index 3b80518d32dcd..6ddf6fa31820a 100644 --- a/datafusion/core/benches/filter_query_sql.rs +++ b/datafusion/core/benches/filter_query_sql.rs @@ -23,12 +23,11 @@ use arrow::{ use criterion::{Criterion, criterion_group, criterion_main}; use datafusion::prelude::SessionContext; use datafusion::{datasource::MemTable, error::Result}; -use futures::executor::block_on; use std::hint::black_box; use std::sync::Arc; use tokio::runtime::Runtime; -async fn query(ctx: &SessionContext, rt: &Runtime, sql: &str) { +fn query(ctx: &SessionContext, rt: &Runtime, sql: &str) { // execute the query let df = rt.block_on(ctx.sql(sql)).unwrap(); black_box(rt.block_on(df.collect()).unwrap()); @@ -71,28 +70,28 @@ fn criterion_benchmark(c: &mut Criterion) { c.bench_function("filter_array", |b| { let ctx = create_context(array_len, batch_size).unwrap(); - b.iter(|| block_on(query(&ctx, &rt, "select f32, f64 from t where f32 >= f64"))) + b.iter(|| query(&ctx, &rt, "select f32, f64 from t where f32 >= f64")) }); c.bench_function("filter_scalar", |b| { let ctx = create_context(array_len, batch_size).unwrap(); b.iter(|| { - block_on(query( + query( &ctx, &rt, "select f32, f64 from t where f32 >= 250 and f64 > 250", - )) + ) }) }); c.bench_function("filter_scalar in list", |b| { let ctx = create_context(array_len, batch_size).unwrap(); b.iter(|| { - block_on(query( + query( &ctx, &rt, "select f32, f64 from t where f32 in (10, 20, 30, 40)", - )) + ) }) }); } diff --git a/datafusion/core/benches/map_query_sql.rs b/datafusion/core/benches/map_query_sql.rs index 67904197bc257..6e7d584c6fce6 100644 --- a/datafusion/core/benches/map_query_sql.rs +++ b/datafusion/core/benches/map_query_sql.rs @@ -22,8 +22,7 @@ use std::sync::Arc; use arrow::array::{ArrayRef, Int32Array, RecordBatch}; use criterion::{Criterion, criterion_group, criterion_main}; use parking_lot::Mutex; -use rand::Rng; -use rand::prelude::ThreadRng; +use rand::prelude::*; use tokio::runtime::Runtime; use datafusion::prelude::SessionContext; @@ -33,7 +32,7 @@ use datafusion_functions_nested::map::map; mod data_utils; -fn build_keys(rng: &mut ThreadRng) -> Vec { +fn build_keys(rng: &mut StdRng) -> Vec { let mut keys = HashSet::with_capacity(1000); while keys.len() < 1000 { let key = rng.random_range(0..9999).to_string(); @@ -42,7 +41,7 @@ fn build_keys(rng: &mut ThreadRng) -> Vec { keys.into_iter().collect() } -fn build_values(rng: &mut ThreadRng) -> Vec { +fn build_values(rng: &mut StdRng) -> Vec { let mut values = vec![]; for _ in 0..1000 { values.push(rng.random_range(0..9999)); @@ -67,7 +66,7 @@ fn criterion_benchmark(c: &mut Criterion) { let rt = Runtime::new().unwrap(); let df = rt.block_on(ctx.lock().table("t")).unwrap(); - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let keys = build_keys(&mut rng); let values = build_values(&mut rng); let mut key_buffer = Vec::new(); diff --git a/datafusion/core/benches/parquet_nested_schema_pruning.rs b/datafusion/core/benches/parquet_nested_schema_pruning.rs new file mode 100644 index 0000000000000..de4f0a57a5c41 --- /dev/null +++ b/datafusion/core/benches/parquet_nested_schema_pruning.rs @@ -0,0 +1,445 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for schema-driven nested projection pruning in Parquet. +//! +//! A table's declared (logical) schema can be *narrower* than the physical +//! parquet type of a nested column — e.g. the table declares +//! `events: LIST>` while the file contains +//! `events: LIST>`. Engines like Spark +//! communicate nested projection pruning to the scan exactly this way +//! (a clipped read schema), so the reader should fetch and decode only the +//! leaves the declared schema names. +//! +//! Each dataset shape is measured three ways: +//! +//! 1. **narrow_schema**: wide file, narrow declared table schema — the +//! interesting case; ideally close to (3) +//! 2. **full_schema**: wide file, full table schema — the cost of reading +//! everything +//! 3. **physically_narrow**: a file that only contains the narrow columns — +//! the floor +//! +//! At setup the benchmark reads the parquet scan's `bytes_scanned` metric for +//! (1), (2) and (3) so the IO pattern is visible in addition to wall time, and +//! asserts that nested projection pruning keeps the narrow declared schema's +//! scan well below the full schema's, close to the physically-narrow floor +//! (see [`assert_scan_prunes`]). + +use arrow::array::{ + ArrayRef, Int32Array, Int64Array, ListArray, StringArray, StructArray, +}; +use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::datasource::listing::{ + ListingTable, ListingTableConfig, ListingTableConfigExt, +}; +use datafusion::physical_plan::metrics::MetricsSet; +use datafusion::physical_plan::{ExecutionPlan, collect}; +use datafusion::prelude::SessionContext; +use datafusion_datasource::ListingTableUrl; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use std::hint::black_box; +use std::sync::Arc; +use std::time::Duration; +use tempfile::NamedTempFile; +use tokio::runtime::Runtime; + +const NUM_BATCHES: usize = 2; +const ROWS_PER_BATCH: usize = 256; +const ROW_GROUP_ROW_COUNT: usize = 256; +const ELEMS_PER_ROW: usize = 3; +const NUM_PAD_FIELDS: usize = 8; +const PAD_LEN: usize = 2048; + +/// The narrow item fields: the subset of the struct the table declares. +fn narrow_item_fields() -> Fields { + Fields::from(vec![ + Field::new("x", DataType::Int64, true), + Field::new("y", DataType::Utf8, true), + ]) +} + +/// The wide item fields as written to the file: the narrow fields plus +/// `NUM_PAD_FIELDS` fat string fields the table schema does not mention. +/// +/// Derived from [`narrow_item_fields`] so the shared columns (`x`, `y`) match +/// by construction — same names, types and nullability — and only the extra +/// pad fields distinguish the two. +fn wide_item_fields() -> Fields { + let mut fields: Vec = narrow_item_fields() + .iter() + .map(|f| f.as_ref().clone()) + .collect(); + for i in 0..NUM_PAD_FIELDS { + fields.push(Field::new(format!("pad_{i}"), DataType::Utf8, false)); + } + Fields::from(fields) +} + +fn list_schema(item_fields: Fields) -> SchemaRef { + let item = Arc::new(Field::new("item", DataType::Struct(item_fields), true)); + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("events", DataType::List(item), true), + ])) +} + +fn struct_schema(item_fields: Fields) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(item_fields), true), + ])) +} + +/// Distinct pad values so dictionary encoding cannot collapse them. +fn pad_values(count: usize, seed: usize) -> ArrayRef { + let base = "x".repeat(PAD_LEN); + let values: Vec = (0..count) + .map(|i| format!("{:08}{base}", seed + i)) + .collect(); + Arc::new(StringArray::from(values)) +} + +/// Struct children for `count` elements, restricted to `fields`. +fn item_columns(fields: &Fields, count: usize, seed: usize) -> Vec { + fields + .iter() + .enumerate() + .map(|(i, field)| match field.name().as_str() { + "x" => Arc::new(Int64Array::from_iter_values( + (0..count).map(|j| (seed + j) as i64), + )) as ArrayRef, + "y" => Arc::new(StringArray::from_iter_values( + (0..count).map(|j| format!("y-{}", seed + j)), + )) as ArrayRef, + // `seed + i` keeps each pad column's values distinct from the + // others (and matches the additive seeding used above); a + // multiplier like `seed * (i + 1)` collapses to the same seed for + // every column when `seed == 0` (the first batch). + _ => pad_values(count, seed + i), + }) + .collect() +} + +fn list_batch(fields: &Fields, batch_id: usize) -> RecordBatch { + let num_elems = ROWS_PER_BATCH * ELEMS_PER_ROW; + let seed = batch_id * num_elems; + let struct_array = + StructArray::new(fields.clone(), item_columns(fields, num_elems, seed), None); + let item = Arc::new(Field::new("item", DataType::Struct(fields.clone()), true)); + let events = ListArray::new( + item, + OffsetBuffer::from_lengths(std::iter::repeat_n(ELEMS_PER_ROW, ROWS_PER_BATCH)), + Arc::new(struct_array), + None, + ); + let ids = Int32Array::from_iter_values( + (0..ROWS_PER_BATCH).map(|i| (batch_id * ROWS_PER_BATCH + i) as i32), + ); + RecordBatch::try_new( + list_schema(fields.clone()), + vec![Arc::new(ids), Arc::new(events)], + ) + .unwrap() +} + +fn struct_batch(fields: &Fields, batch_id: usize) -> RecordBatch { + let seed = batch_id * ROWS_PER_BATCH; + let struct_array = StructArray::new( + fields.clone(), + item_columns(fields, ROWS_PER_BATCH, seed), + None, + ); + let ids = + Int32Array::from_iter_values((0..ROWS_PER_BATCH).map(|i| (seed + i) as i32)); + RecordBatch::try_new( + struct_schema(fields.clone()), + vec![Arc::new(ids), Arc::new(struct_array)], + ) + .unwrap() +} + +fn generate_file( + schema: SchemaRef, + batch_fn: impl Fn(usize) -> RecordBatch, + prefix: &str, +) -> NamedTempFile { + let mut named_file = tempfile::Builder::new() + .prefix(prefix) + .suffix(".parquet") + .tempfile() + .unwrap(); + + let properties = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_dictionary_enabled(false) + .set_max_row_group_row_count(Some(ROW_GROUP_ROW_COUNT)) + .build(); + + let mut writer = + ArrowWriter::try_new(&mut named_file, schema, Some(properties)).unwrap(); + for batch_id in 0..NUM_BATCHES { + writer.write(&batch_fn(batch_id)).unwrap(); + } + let metadata = writer.close().unwrap(); + println!( + "Generated {} ({} rows, {} row groups, {} bytes)", + named_file.path().display(), + metadata.file_metadata().num_rows(), + metadata.row_groups().len(), + std::fs::metadata(named_file.path()).unwrap().len(), + ); + named_file +} + +/// Register `path` as `table`, declaring `table_schema` (which may be narrower +/// than the file's physical schema). +fn register_table( + ctx: &SessionContext, + rt: &Runtime, + table: &str, + path: &str, + table_schema: SchemaRef, +) { + let url = ListingTableUrl::parse(path).unwrap(); + let config = rt + .block_on(ListingTableConfig::new(url).infer_options(&ctx.state())) + .unwrap() + .with_schema(table_schema); + let provider = ListingTable::try_new(config).unwrap(); + ctx.register_table(table, Arc::new(provider)).unwrap(); +} + +fn query(ctx: &SessionContext, rt: &Runtime, sql: &str) { + let df = rt.block_on(ctx.sql(sql)).unwrap(); + black_box(rt.block_on(df.collect()).unwrap()); +} + +/// Recursively collect the metrics of every node in `plan` into `out`. +fn gather_metrics(plan: &Arc, out: &mut MetricsSet) { + if let Some(metrics) = plan.metrics() { + for metric in metrics.iter() { + out.push(Arc::clone(metric)); + } + } + for child in plan.children() { + gather_metrics(child, out); + } +} + +/// Execute `sql` and return the parquet scan's `bytes_scanned` metric, read +/// from the typed metrics API rather than scraped from display output (which +/// would silently break if the format ever changed). +fn scan_bytes(ctx: &SessionContext, rt: &Runtime, sql: &str) -> usize { + let df = rt.block_on(ctx.sql(sql)).unwrap(); + let plan = rt.block_on(df.create_physical_plan()).unwrap(); + // Fully drive the plan so the scan populates its metrics. + black_box( + rt.block_on(collect(Arc::clone(&plan), ctx.task_ctx())) + .unwrap(), + ); + + let mut metrics = MetricsSet::new(); + gather_metrics(&plan, &mut metrics); + metrics + .aggregate_by_name() + .sum_by_name("bytes_scanned") + .map(|v| v.as_usize()) + .expect("parquet scan should report a bytes_scanned metric") +} + +/// Report and assert the `bytes_scanned` improvement for one dataset shape. +/// +/// `narrow` selects from a wide file through a narrow declared schema, `full` +/// through the full schema, and `floor` from a physically-narrow file. +/// Nested projection pruning clips the narrow read to the declared leaves, so +/// `narrow` should read substantially less than `full`, close to `floor`, +/// the cost of a file that never had the extra leaves to begin with. +fn assert_scan_prunes( + ctx: &SessionContext, + rt: &Runtime, + label: &str, + narrow_sql: &str, + full_sql: &str, + floor_sql: &str, +) { + let narrow = scan_bytes(ctx, rt, narrow_sql); + let full = scan_bytes(ctx, rt, full_sql); + let floor = scan_bytes(ctx, rt, floor_sql); + println!( + "{label}: bytes_scanned narrow_schema={narrow} full_schema={full} \ + physically_narrow={floor}" + ); + assert!( + narrow * 2 < full, + "{label}: expected the narrow declared schema to read less than half \ + of the full schema's {full} bytes (physically-narrow floor is \ + {floor} bytes), but it read {narrow}" + ); +} + +struct Fixture { + ctx: SessionContext, + rt: Runtime, + _files: Vec, +} + +/// Tables: +/// `_narrow_schema`: wide file, narrow declared schema +/// `_full_schema`: wide file, full declared schema +/// `_physically_narrow`: narrow file, narrow declared schema +fn setup( + name: &str, + schema_fn: fn(Fields) -> SchemaRef, + batch_fn: fn(&Fields, usize) -> RecordBatch, +) -> Fixture { + let rt = Runtime::new().unwrap(); + let ctx = SessionContext::new(); + + let wide = wide_item_fields(); + let narrow = narrow_item_fields(); + + let wide_file = generate_file(schema_fn(wide.clone()), |i| batch_fn(&wide, i), name); + let narrow_file = generate_file( + schema_fn(narrow.clone()), + |i| batch_fn(&narrow, i), + &format!("{name}_narrow"), + ); + let wide_path = wide_file.path().display().to_string(); + let narrow_path = narrow_file.path().display().to_string(); + + register_table( + &ctx, + &rt, + &format!("{name}_narrow_schema"), + &wide_path, + schema_fn(narrow.clone()), + ); + register_table( + &ctx, + &rt, + &format!("{name}_full_schema"), + &wide_path, + schema_fn(wide.clone()), + ); + register_table( + &ctx, + &rt, + &format!("{name}_physically_narrow"), + &narrow_path, + schema_fn(narrow.clone()), + ); + + Fixture { + ctx, + rt, + _files: vec![wide_file, narrow_file], + } +} + +fn list_struct_benchmarks(c: &mut Criterion) { + let f = setup("list_struct", list_schema, list_batch); + let (ctx, rt) = (&f.ctx, &f.rt); + + assert_scan_prunes( + ctx, + rt, + "list_struct", + "SELECT events FROM list_struct_narrow_schema", + "SELECT events FROM list_struct_full_schema", + "SELECT events FROM list_struct_physically_narrow", + ); + + let mut group = c.benchmark_group("list_struct"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + + // wide file, narrow declared schema: should only read the narrow leaves + group.bench_function("select_events_narrow_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT events FROM list_struct_narrow_schema")) + }); + + // wide file, full schema: the cost of reading everything + group.bench_function("select_events_full_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT events FROM list_struct_full_schema")) + }); + + // narrow file: the floor + group.bench_function("select_events_physically_narrow", |b| { + b.iter(|| query(ctx, rt, "SELECT events FROM list_struct_physically_narrow")) + }); + + // aggregation over one narrow leaf through unnest + group.bench_function("sum_x_narrow_schema", |b| { + b.iter(|| { + query( + ctx, + rt, + "SELECT SUM(e['x']) FROM (SELECT UNNEST(events) AS e FROM list_struct_narrow_schema)", + ) + }) + }); + + group.finish(); +} + +fn top_level_struct_benchmarks(c: &mut Criterion) { + let f = setup("struct", struct_schema, struct_batch); + let (ctx, rt) = (&f.ctx, &f.rt); + + assert_scan_prunes( + ctx, + rt, + "top_level_struct", + "SELECT s FROM struct_narrow_schema", + "SELECT s FROM struct_full_schema", + "SELECT s FROM struct_physically_narrow", + ); + + let mut group = c.benchmark_group("top_level_struct"); + group.sample_size(10); + group.warm_up_time(Duration::from_secs(1)); + group.measurement_time(Duration::from_secs(3)); + + group.bench_function("select_struct_narrow_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT s FROM struct_narrow_schema")) + }); + + group.bench_function("select_struct_full_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT s FROM struct_full_schema")) + }); + + group.bench_function("select_struct_physically_narrow", |b| { + b.iter(|| query(ctx, rt, "SELECT s FROM struct_physically_narrow")) + }); + + // get_field on a schema-narrowed struct column: the expression-level + // pruning path interacting with the schema-level narrowing + group.bench_function("sum_x_narrow_schema", |b| { + b.iter(|| query(ctx, rt, "SELECT SUM(s['x']) FROM struct_narrow_schema")) + }); + + group.finish(); +} + +criterion_group!(benches, list_struct_benchmarks, top_level_struct_benchmarks); +criterion_main!(benches); diff --git a/datafusion/core/benches/parquet_query_sql.rs b/datafusion/core/benches/parquet_query_sql.rs index f099137973592..2e7794bfd19b4 100644 --- a/datafusion/core/benches/parquet_query_sql.rs +++ b/datafusion/core/benches/parquet_query_sql.rs @@ -32,7 +32,6 @@ use parquet::file::properties::{WriterProperties, WriterVersion}; use rand::distr::Alphanumeric; use rand::distr::uniform::SampleUniform; use rand::prelude::*; -use rand::rng; use std::fs::File; use std::io::Read; use std::ops::Range; @@ -69,36 +68,36 @@ fn schema() -> SchemaRef { ])) } -fn generate_batch() -> RecordBatch { +fn generate_batch(rng: &mut StdRng) -> RecordBatch { let schema = schema(); let len = WRITE_RECORD_BATCH_SIZE; RecordBatch::try_new( schema, vec![ - generate_string_dictionary("prefix", 10, len, 1.0), - generate_string_dictionary("prefix", 10, len, 0.5), - generate_string_dictionary("prefix", 100, len, 1.0), - generate_string_dictionary("prefix", 100, len, 0.5), - generate_string_dictionary("prefix", 1000, len, 1.0), - generate_string_dictionary("prefix", 1000, len, 0.5), - generate_strings(0..100, len, 1.0), - generate_strings(0..100, len, 0.5), - generate_primitive::(len, 1.0, -2000..2000), - generate_primitive::(len, 0.5, -2000..2000), - generate_primitive::(len, 1.0, -1000.0..1000.0), - generate_primitive::(len, 0.5, -1000.0..1000.0), + generate_string_dictionary(rng, "prefix", 10, len, 1.0), + generate_string_dictionary(rng, "prefix", 10, len, 0.5), + generate_string_dictionary(rng, "prefix", 100, len, 1.0), + generate_string_dictionary(rng, "prefix", 100, len, 0.5), + generate_string_dictionary(rng, "prefix", 1000, len, 1.0), + generate_string_dictionary(rng, "prefix", 1000, len, 0.5), + generate_strings(rng, 0..100, len, 1.0), + generate_strings(rng, 0..100, len, 0.5), + generate_primitive::(rng, len, 1.0, -2000..2000), + generate_primitive::(rng, len, 0.5, -2000..2000), + generate_primitive::(rng, len, 1.0, -1000.0..1000.0), + generate_primitive::(rng, len, 0.5, -1000.0..1000.0), ], ) .unwrap() } fn generate_string_dictionary( + rng: &mut StdRng, prefix: &str, cardinality: usize, len: usize, valid_percent: f64, ) -> ArrayRef { - let mut rng = rng(); let strings: Vec<_> = (0..cardinality).map(|x| format!("{prefix}#{x}")).collect(); Arc::new(DictionaryArray::::from_iter((0..len).map( @@ -110,11 +109,11 @@ fn generate_string_dictionary( } fn generate_strings( + rng: &mut StdRng, string_length_range: Range, len: usize, valid_percent: f64, ) -> ArrayRef { - let mut rng = rng(); Arc::new(StringArray::from_iter((0..len).map(|_| { rng.random_bool(valid_percent).then(|| { let string_len = rng.random_range(string_length_range.clone()); @@ -126,6 +125,7 @@ fn generate_strings( } fn generate_primitive( + rng: &mut StdRng, len: usize, valid_percent: f64, range: Range, @@ -134,7 +134,6 @@ where T: ArrowPrimitiveType, T::Native: SampleUniform, { - let mut rng = rng(); Arc::new(PrimitiveArray::::from_iter((0..len).map(|_| { rng.random_bool(valid_percent) .then(|| rng.random_range(range.clone())) @@ -160,8 +159,9 @@ fn generate_file() -> NamedTempFile { let mut writer = ArrowWriter::try_new(&mut named_file, schema, Some(properties)).unwrap(); + let mut rng = StdRng::seed_from_u64(0); for _ in 0..NUM_BATCHES { - let batch = generate_batch(); + let batch = generate_batch(&mut rng); writer.write(&batch).unwrap(); } diff --git a/datafusion/core/benches/parquet_struct_query.rs b/datafusion/core/benches/parquet_struct_query.rs index e7e91f0dd0e1e..b7132973c1bff 100644 --- a/datafusion/core/benches/parquet_struct_query.rs +++ b/datafusion/core/benches/parquet_struct_query.rs @@ -27,7 +27,6 @@ use parquet::arrow::ArrowWriter; use parquet::file::properties::{WriterProperties, WriterVersion}; use rand::distr::Alphanumeric; use rand::prelude::*; -use rand::rng; use std::hint::black_box; use std::ops::Range; use std::path::Path; @@ -59,8 +58,7 @@ fn schema() -> SchemaRef { ])) } -fn generate_strings(len: usize) -> ArrayRef { - let mut rng = rng(); +fn generate_strings(rng: &mut StdRng, len: usize) -> ArrayRef { Arc::new(StringArray::from_iter((0..len).map(|_| { let string_len = rng.random_range(STRING_LENGTH_RANGE.clone()); Some( @@ -71,7 +69,7 @@ fn generate_strings(len: usize) -> ArrayRef { }))) } -fn generate_batch(batch_id: usize) -> RecordBatch { +fn generate_batch(rng: &mut StdRng, batch_id: usize) -> RecordBatch { let schema = schema(); let len = WRITE_RECORD_BATCH_SIZE; @@ -84,7 +82,7 @@ fn generate_batch(batch_id: usize) -> RecordBatch { let struct_id_array = Arc::new(Int32Array::from(id_values)); // Generate random strings for struct value field - let value_array = generate_strings(len); + let value_array = generate_strings(rng, len); // Construct StructArray let struct_array = StructArray::from(vec![ @@ -120,8 +118,9 @@ fn generate_file() -> NamedTempFile { let mut writer = ArrowWriter::try_new(&mut named_file, schema, Some(properties)).unwrap(); + let mut rng = StdRng::seed_from_u64(0); for batch_id in 0..NUM_BATCHES { - let batch = generate_batch(batch_id); + let batch = generate_batch(&mut rng, batch_id); writer.write(&batch).unwrap(); } diff --git a/datafusion/core/benches/parquet_struct_shared_prefix_pushdown.rs b/datafusion/core/benches/parquet_struct_shared_prefix_pushdown.rs new file mode 100644 index 0000000000000..f08c2b554717c --- /dev/null +++ b/datafusion/core/benches/parquet_struct_shared_prefix_pushdown.rs @@ -0,0 +1,329 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for row-filter pushdown with predicates that reach several struct +//! leaves under a common prefix. +//! +//! The existing `parquet_struct_query` bench exercises single-field struct +//! predicates only; `parquet_struct_projection` has no `WHERE` clause. Neither +//! drives the row-filter planner with multiple accesses under the same struct +//! root. +//! +//! Two properties of the planner shape these queries, and both are easy to get +//! wrong: +//! +//! * `execution.parquet.pushdown_filters` must be enabled. It defaults to +//! `false`, in which case no row filter is built and every case below +//! degenerates to a plain scan. +//! * The predicate must be a *single* conjunct. `build_row_filter` calls +//! `split_conjunction` before building filter candidates, and each candidate +//! collects its own access paths, so `s['a'] = 5 AND s['b'] = 5` becomes two +//! independent single-access candidates and never reaches multi-access +//! planning. The cases below use the `(s['a'] + s['b']) = 10` form so every +//! access lands in one candidate. +//! +//! Nested access is written `s['inner']['x']`, which the planner represents as a +//! single flattened `get_field(s, 'inner', 'x')`. That form is pushdown-eligible; +//! a chained `get_field(get_field(s, 'inner'), 'x')` is not. +//! +//! Dataset schema: +//! +//! ```sql +//! CREATE TABLE t ( +//! id INT, +//! s STRUCT< +//! a INT, b INT, c INT, d INT, e INT, +//! inner STRUCT +//! > +//! ); +//! ``` +//! +//! All struct leaves mirror the top-level `id`, so a sum of `n` leaves equals +//! `n * id` and every predicate is satisfied by exactly one row (`id = 5`). +//! Holding the match count fixed keeps the cases comparable, and each is +//! asserted to return that single row. + +use arrow::array::{ArrayRef, Int32Array, StructArray}; +use arrow::datatypes::{DataType, Field, Fields, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion::prelude::{SessionConfig, SessionContext}; +use datafusion_common::instant::Instant; +use parquet::arrow::ArrowWriter; +use parquet::file::properties::{WriterProperties, WriterVersion}; +use std::hint::black_box; +use std::path::Path; +use std::sync::Arc; +use tempfile::NamedTempFile; +use tokio::runtime::Runtime; + +/// The number of batches to write +const NUM_BATCHES: usize = 128; +/// The number of rows in each record batch to write +const WRITE_RECORD_BATCH_SIZE: usize = 4096; +/// The number of rows in a row group +const ROW_GROUP_ROW_COUNT: usize = 65536; +/// The number of row groups expected +const EXPECTED_ROW_GROUPS: usize = 8; +/// Number of rows every predicate is expected to match. +const EXPECTED_MATCHES: usize = 1; + +fn inner_struct_fields() -> Fields { + Fields::from(vec![ + Field::new("x", DataType::Int32, false), + Field::new("y", DataType::Int32, false), + Field::new("z", DataType::Int32, false), + ]) +} + +fn struct_fields() -> Fields { + Fields::from(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + Field::new("d", DataType::Int32, false), + Field::new("e", DataType::Int32, false), + Field::new("inner", DataType::Struct(inner_struct_fields()), false), + ]) +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(struct_fields()), false), + ])) +} + +fn generate_batch(batch_id: usize) -> RecordBatch { + let schema = schema(); + let len = WRITE_RECORD_BATCH_SIZE; + + // Sequential IDs give distinct per-row values so a predicate like + // `s['a'] = 5` matches exactly one row, mirroring parquet_struct_query. + let base_id = (batch_id * len) as i32; + let id_values: Vec = (0..len).map(|i| base_id + i as i32).collect(); + let id_array = Arc::new(Int32Array::from(id_values.clone())); + + let leaf = || Arc::new(Int32Array::from(id_values.clone())) as ArrayRef; + + let inner_struct = StructArray::from(vec![ + (Arc::new(Field::new("x", DataType::Int32, false)), leaf()), + (Arc::new(Field::new("y", DataType::Int32, false)), leaf()), + (Arc::new(Field::new("z", DataType::Int32, false)), leaf()), + ]); + + let struct_array = StructArray::from(vec![ + (Arc::new(Field::new("a", DataType::Int32, false)), leaf()), + (Arc::new(Field::new("b", DataType::Int32, false)), leaf()), + (Arc::new(Field::new("c", DataType::Int32, false)), leaf()), + (Arc::new(Field::new("d", DataType::Int32, false)), leaf()), + (Arc::new(Field::new("e", DataType::Int32, false)), leaf()), + ( + Arc::new(Field::new( + "inner", + DataType::Struct(inner_struct_fields()), + false, + )), + Arc::new(inner_struct) as ArrayRef, + ), + ]); + + RecordBatch::try_new(schema, vec![id_array, Arc::new(struct_array)]).unwrap() +} + +fn generate_file() -> NamedTempFile { + let now = Instant::now(); + let mut named_file = tempfile::Builder::new() + .prefix("parquet_struct_shared_prefix_pushdown") + .suffix(".parquet") + .tempfile() + .unwrap(); + + println!("Generating parquet file - {}", named_file.path().display()); + let schema = schema(); + + let properties = WriterProperties::builder() + .set_writer_version(WriterVersion::PARQUET_2_0) + .set_max_row_group_row_count(Some(ROW_GROUP_ROW_COUNT)) + .build(); + + let mut writer = + ArrowWriter::try_new(&mut named_file, schema, Some(properties)).unwrap(); + + for batch_id in 0..NUM_BATCHES { + let batch = generate_batch(batch_id); + writer.write(&batch).unwrap(); + } + + let metadata = writer.close().unwrap(); + let file_metadata = metadata.file_metadata(); + let expected_rows = WRITE_RECORD_BATCH_SIZE * NUM_BATCHES; + assert_eq!( + file_metadata.num_rows() as usize, + expected_rows, + "Expected {expected_rows} rows but got {}", + file_metadata.num_rows() + ); + assert_eq!( + metadata.row_groups().len(), + EXPECTED_ROW_GROUPS, + "Expected {EXPECTED_ROW_GROUPS} row groups but got {}", + metadata.row_groups().len() + ); + + println!( + "Generated parquet file with {} rows and {} row groups in {:.2}s", + file_metadata.num_rows(), + metadata.row_groups().len(), + now.elapsed().as_secs_f32() + ); + + named_file +} + +fn create_context(file_path: &str, rt: &Runtime) -> SessionContext { + let mut config = SessionConfig::new(); + // Row-filter pushdown is off by default. Without it no row filter is built, + // and these benchmarks would time a plain scan for every predicate shape. + config.options_mut().execution.parquet.pushdown_filters = true; + + let ctx = SessionContext::new_with_config(config); + rt.block_on(ctx.register_parquet("t", file_path, Default::default())) + .unwrap(); + ctx +} + +fn query(ctx: &SessionContext, rt: &Runtime, sql: &str) { + let ctx = ctx.clone(); + let sql = sql.to_string(); + let df = rt.block_on(ctx.sql(&sql)).unwrap(); + black_box(rt.block_on(df.collect()).unwrap()); +} + +/// Fails unless `sql` actually pushes a row filter into the Parquet decoder and +/// matches [`EXPECTED_MATCHES`] rows. +/// +/// Guards the two silent-failure modes: a disabled `pushdown_filters` and a +/// predicate shape that turns out not to be pushdown-eligible. Either would +/// leave the benchmark timing a plain scan instead of row-filter pushdown. +/// +/// Metrics are read off the executed plan rather than scraped from +/// `EXPLAIN ANALYZE` text, so the check does not depend on output formatting. +fn assert_pushdown_active(ctx: &SessionContext, rt: &Runtime, name: &str, sql: &str) { + let (rows, pruned) = rt + .block_on(async { + let plan = ctx.sql(sql).await?.create_physical_plan().await?; + let batches = + datafusion::physical_plan::collect(Arc::clone(&plan), ctx.task_ctx()) + .await?; + let rows = batches.iter().map(|b| b.num_rows()).sum::(); + Ok::<_, datafusion_common::DataFusionError>((rows, rows_pruned(&plan))) + }) + .unwrap(); + + assert_eq!( + rows, EXPECTED_MATCHES, + "`{name}` matched {rows} rows, expected {EXPECTED_MATCHES}" + ); + assert!( + pruned > 0, + "`{name}` pruned no rows via the Parquet row filter, so it does not \ + exercise row-filter pushdown (is `pushdown_filters` enabled, and is \ + the predicate a single pushdown-eligible conjunct?)" + ); +} + +/// Total `pushdown_rows_pruned` reported anywhere in the executed plan. +fn rows_pruned(plan: &Arc) -> usize { + let mut total = plan + .metrics() + .and_then(|metrics| metrics.sum_by_name("pushdown_rows_pruned")) + .map(|value| value.as_usize()) + .unwrap_or(0); + + for child in plan.children() { + total += rows_pruned(child); + } + + total +} + +fn criterion_benchmark(c: &mut Criterion) { + let (file_path, temp_file) = match std::env::var("PARQUET_FILE") { + Ok(file) => (file, None), + Err(_) => { + let temp_file = generate_file(); + (temp_file.path().display().to_string(), Some(temp_file)) + } + }; + + assert!(Path::new(&file_path).exists(), "path not found"); + println!("Using parquet file {file_path}"); + + let rt = Runtime::new().unwrap(); + let ctx = create_context(&file_path, &rt); + + // Baseline: one access on a single struct leaf. + let sql = "select id from t where s['a'] = 5"; + assert_pushdown_active(&ctx, &rt, "1_access", sql); + c.bench_function("1_access", |b| b.iter(|| query(&ctx, &rt, sql))); + + // Two accesses inside one conjunct, sharing the struct root `s`. + let sql = "select id from t where (s['a'] + s['b']) = 10"; + assert_pushdown_active(&ctx, &rt, "2_access_shared_root", sql); + c.bench_function("2_access_shared_root", |b| b.iter(|| query(&ctx, &rt, sql))); + + // Three accesses sharing the struct root `s`. + let sql = "select id from t where (s['a'] + s['b'] + s['c']) = 15"; + assert_pushdown_active(&ctx, &rt, "3_access_shared_root", sql); + c.bench_function("3_access_shared_root", |b| b.iter(|| query(&ctx, &rt, sql))); + + // Five accesses sharing the struct root `s`, amplifying planning cost. + let sql = "select id from t \ + where (s['a'] + s['b'] + s['c'] + s['d'] + s['e']) = 25"; + assert_pushdown_active(&ctx, &rt, "5_access_shared_root", sql); + c.bench_function("5_access_shared_root", |b| b.iter(|| query(&ctx, &rt, sql))); + + // Two accesses sharing the deeper prefix `s.inner`. + let sql = "select id from t where (s['inner']['x'] + s['inner']['y']) = 10"; + assert_pushdown_active(&ctx, &rt, "2_access_shared_nested_prefix", sql); + c.bench_function("2_access_shared_nested_prefix", |b| { + b.iter(|| query(&ctx, &rt, sql)) + }); + + // Three accesses sharing the deeper prefix `s.inner`. + let sql = "select id from t \ + where (s['inner']['x'] + s['inner']['y'] + s['inner']['z']) = 15"; + assert_pushdown_active(&ctx, &rt, "3_access_shared_nested_prefix", sql); + c.bench_function("3_access_shared_nested_prefix", |b| { + b.iter(|| query(&ctx, &rt, sql)) + }); + + // Mix: two accesses on `s` leaves and two on `s.inner` leaves. + let sql = "select id from t \ + where (s['a'] + s['b'] + s['inner']['x'] + s['inner']['y']) = 20"; + assert_pushdown_active(&ctx, &rt, "mixed_depth_shared_prefix", sql); + c.bench_function("mixed_depth_shared_prefix", |b| { + b.iter(|| query(&ctx, &rt, sql)) + }); + + // Temporary file must outlive the benchmarks, it is deleted when dropped + drop(temp_file); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/core/benches/sort.rs b/datafusion/core/benches/sort.rs index 7544f7ae26d43..ac4be5b8b2c9f 100644 --- a/datafusion/core/benches/sort.rs +++ b/datafusion/core/benches/sort.rs @@ -66,12 +66,10 @@ //! ~10% duplicates rows) //! ``` -use std::sync::Arc; - -use arrow::array::StringViewArray; +use arrow::array::{ArrayRef, StringViewArray, StringViewBuilder}; use arrow::{ - array::{DictionaryArray, Float64Array, Int64Array, StringArray}, - datatypes::{Int32Type, Schema}, + array::{Array, DictionaryArray, Float64Array, Int64Array, StringArray}, + datatypes::{Field, Int32Type, Schema}, record_batch::RecordBatch, }; use datafusion::physical_plan::sorts::sort::SortExec; @@ -87,11 +85,16 @@ use datafusion::{ use datafusion_datasource::memory::MemorySourceConfig; use datafusion_physical_expr::{PhysicalSortExpr, expressions::col}; use datafusion_physical_expr_common::sort_expr::LexOrdering; +use std::sync::Arc; +use std::time::Duration; /// Benchmarks for SortPreservingMerge stream use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_execution::config::SessionConfig; use futures::StreamExt; +use itertools::Itertools; use rand::rngs::StdRng; +use rand::seq::SliceRandom; use rand::{Rng, SeedableRng}; use tokio::runtime::Runtime; @@ -103,10 +106,55 @@ const NUM_STREAMS: usize = 8; const BATCH_SIZE: usize = 1024; /// Input sizes to benchmark. The small size (100K) exercises the -/// in-memory concat-and-sort path; the large size (10M) exercises +/// in-memory concat-and-sort path; the large size (1M) exercises /// the sort-then-merge path with high fan-in. const INPUT_SIZES: &[(u64, &str)] = &[(100_000, "100k"), (1_000_000, "1M")]; +/// Number of extra (non-sort-key) payload columns to carry alongside the sort +/// keys in the axis benchmarks. Measures the cost of reordering wide batches. +const EXTRA_COLUMN_COUNTS: &[usize] = &[0, 5, 20, 100]; + +/// Input ordering profiles for the SortExec axis benchmarks. +#[derive(Clone, Copy, Debug)] +enum DataProfile { + Sorted, + Unsorted, + /// Fully sorted, then 10% of rows swapped to random positions. + NearlySorted, +} + +impl DataProfile { + /// Arrange `v` (whose initial order is irrelevant) into this profile. + fn apply(self, mut v: Vec) -> Vec { + let mut rng = StdRng::seed_from_u64(99); + match self { + DataProfile::Sorted => v.sort_unstable(), + DataProfile::Unsorted => v.shuffle(&mut rng), + DataProfile::NearlySorted => { + v.sort_unstable(); + let n = v.len(); + + // 10% is globally misplaced + for _ in 0..n / 10 { + v.swap(rng.random_range(0..n), rng.random_range(0..n)); + } + } + } + v + } +} + +/// Sort-key cardinality, i.e. how much the key values overlap across rows and +/// partitions. Only affects the sort keys, not the extra payload columns. +#[derive(Clone, Copy, Debug)] +enum Cardinality { + /// Heavy overlap: i64 in `0..input_size` (~1/3 duplicates), 100 distinct + /// strings repeated across all rows. + Low, + /// Minimal overlap: full-range i64 and random strings (~no duplicates). + High, +} + type PartitionedBatches = Vec>; type StreamGenerator = Box PartitionedBatches>; @@ -178,25 +226,25 @@ fn criterion_benchmark(c: &mut Criterion) { for (name, f) in &cases { c.bench_function(&format!("merge sorted {name} {size_label}"), |b| { let data = f(true); - let case = BenchCase::merge_sorted(&data); + let case = BenchCase::merge_sorted(BATCH_SIZE, &data); b.iter(move || case.run()) }); c.bench_function(&format!("sort merge {name} {size_label}"), |b| { let data = f(false); - let case = BenchCase::sort_merge(&data); + let case = BenchCase::sort_merge(BATCH_SIZE, &data); b.iter(move || case.run()) }); c.bench_function(&format!("sort {name} {size_label}"), |b| { let data = f(false); - let case = BenchCase::sort(&data); + let case = BenchCase::sort(BATCH_SIZE, &data); b.iter(move || case.run()) }); c.bench_function(&format!("sort partitioned {name} {size_label}"), |b| { let data = f(false); - let case = BenchCase::sort_partitioned(&data); + let case = BenchCase::sort_partitioned(BATCH_SIZE, &data); b.iter(move || case.run()) }); } @@ -215,9 +263,11 @@ struct BenchCase { impl BenchCase { /// Prepare to run a benchmark that merges the specified /// pre-sorted partitions (streams) together using all keys - fn merge_sorted(partitions: &[Vec]) -> Self { + fn merge_sorted(batch_size: usize, partitions: &[Vec]) -> Self { let runtime = tokio::runtime::Builder::new_multi_thread().build().unwrap(); - let session_ctx = SessionContext::new(); + let session_ctx = SessionContext::new_with_config( + SessionConfig::new().with_batch_size(batch_size), + ); let task_ctx = session_ctx.task_ctx(); let schema = partitions[0][0].schema(); @@ -234,9 +284,11 @@ impl BenchCase { } /// Test SortExec in "partitioned" mode followed by a SortPreservingMerge - fn sort_merge(partitions: &[Vec]) -> Self { + fn sort_merge(batch_size: usize, partitions: &[Vec]) -> Self { let runtime = tokio::runtime::Builder::new_multi_thread().build().unwrap(); - let session_ctx = SessionContext::new(); + let session_ctx = SessionContext::new_with_config( + SessionConfig::new().with_batch_size(batch_size), + ); let task_ctx = session_ctx.task_ctx(); let schema = partitions[0][0].schema(); @@ -255,9 +307,11 @@ impl BenchCase { /// Test SortExec in "partitioned" mode which sorts the input streams /// individually into some number of output streams - fn sort(partitions: &[Vec]) -> Self { + fn sort(batch_size: usize, partitions: &[Vec]) -> Self { let runtime = tokio::runtime::Builder::new_multi_thread().build().unwrap(); - let session_ctx = SessionContext::new(); + let session_ctx = SessionContext::new_with_config( + SessionConfig::new().with_batch_size(batch_size), + ); let task_ctx = session_ctx.task_ctx(); let schema = partitions[0][0].schema(); @@ -276,9 +330,11 @@ impl BenchCase { /// Test SortExec in "partitioned" mode which sorts the input streams /// individually into some number of output streams - fn sort_partitioned(partitions: &[Vec]) -> Self { + fn sort_partitioned(batch_size: usize, partitions: &[Vec]) -> Self { let runtime = tokio::runtime::Builder::new_multi_thread().build().unwrap(); - let session_ctx = SessionContext::new(); + let session_ctx = SessionContext::new_with_config( + SessionConfig::new().with_batch_size(batch_size), + ); let task_ctx = session_ctx.task_ctx(); let schema = partitions[0][0].schema(); @@ -312,11 +368,15 @@ impl BenchCase { } } -/// Make sort exprs for each column in `schema` +const EXTRA_COLUMN_NAME_PREFIX: &str = "extra_"; + +/// Make sort exprs for each column in `schema`, skipping non-sort payload +/// columns added by [`with_extra_columns`]. fn make_sort_exprs(schema: &Schema) -> LexOrdering { let sort_exprs = schema .fields() .iter() + .filter(|f| !f.name().starts_with(EXTRA_COLUMN_NAME_PREFIX)) .map(|f| PhysicalSortExpr::new_default(col(f.name(), schema).unwrap())); LexOrdering::new(sort_exprs).unwrap() } @@ -328,10 +388,19 @@ fn i64_streams(sorted: bool, input_size: u64) -> PartitionedBatches { values.sort_unstable(); } - split_tuples(values, |v| { - let array = Int64Array::from(v); - RecordBatch::try_from_iter(vec![("i64", Arc::new(array) as _)]).unwrap() - }) + split_tuples(values, build_i64_batch) +} + +/// Build a single-column i64 [`RecordBatch`]. +fn build_i64_batch(v: Vec) -> RecordBatch { + let array = Int64Array::from(v); + RecordBatch::try_from_iter(vec![("i64", Arc::new(array) as _)]).unwrap() +} + +/// Build a single-column utf8 view [`RecordBatch`] under the given column name. +fn build_utf8_view_batch(name: &str, v: Vec>>) -> RecordBatch { + let array: StringViewArray = v.into_iter().collect(); + RecordBatch::try_from_iter(vec![(name, Arc::new(array) as _)]).unwrap() } /// Create streams of f64 (where approximately 1/3 values are repeated) @@ -369,10 +438,7 @@ fn utf8_view_low_cardinality_streams( if sorted { values.sort_unstable(); } - split_tuples(values, |v| { - let array: StringViewArray = v.into_iter().collect(); - RecordBatch::try_from_iter(vec![("utf_view_low", Arc::new(array) as _)]).unwrap() - }) + split_tuples(values, |v| build_utf8_view_batch("utf_view_low", v)) } /// Create streams of high cardinality (~ no duplicates) utf8_view values @@ -384,10 +450,7 @@ fn utf8_view_high_cardinality_streams( if sorted { values.sort_unstable(); } - split_tuples(values, |v| { - let array: StringViewArray = v.into_iter().collect(); - RecordBatch::try_from_iter(vec![("utf_view_high", Arc::new(array) as _)]).unwrap() - }) + split_tuples(values, |v| build_utf8_view_batch("utf_view_high", v)) } /// Create streams of high cardinality (~ no duplicates) utf8 values @@ -485,25 +548,32 @@ fn mixed_tuple_streams(sorted: bool, input_size: u64) -> PartitionedBatches { tuples.sort_unstable(); } - split_tuples(tuples, |tuples| { - let (tuples, i64_values): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); - let (tuples, utf8_low2): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); - let (f64_values, utf8_low1): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); - - let f64_values: Float64Array = f64_values.into_iter().map(|v| v as f64).collect(); - - let utf8_low1: StringArray = utf8_low1.into_iter().collect(); - let utf8_low2: StringArray = utf8_low2.into_iter().collect(); - let i64_values: Int64Array = i64_values.into_iter().collect(); + split_tuples(tuples, build_mixed_tuple_batch) +} - RecordBatch::try_from_iter(vec![ - ("f64", Arc::new(f64_values) as _), - ("utf_low1", Arc::new(utf8_low1) as _), - ("utf_low2", Arc::new(utf8_low2) as _), - ("i64", Arc::new(i64_values) as _), - ]) - .unwrap() - }) +/// The tuple shape used by the `mixed tuple` case: (i64, utf8_low, utf8_low, i64) +type MixedTuple = (((i64, Option>), Option>), i64); + +/// Build a (f64, utf8_low, utf8_low, i64) batch from [`MixedTuple`]s +/// (the leading i64 becomes the f64 column). +fn build_mixed_tuple_batch(tuples: Vec) -> RecordBatch { + let (tuples, i64_values): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); + let (tuples, utf8_low2): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); + let (f64_values, utf8_low1): (Vec<_>, Vec<_>) = tuples.into_iter().unzip(); + + let f64_values: Float64Array = f64_values.into_iter().map(|v| v as f64).collect(); + + let utf8_low1: StringArray = utf8_low1.into_iter().collect(); + let utf8_low2: StringArray = utf8_low2.into_iter().collect(); + let i64_values: Int64Array = i64_values.into_iter().collect(); + + RecordBatch::try_from_iter(vec![ + ("f64", Arc::new(f64_values) as _), + ("utf_low1", Arc::new(utf8_low1) as _), + ("utf_low2", Arc::new(utf8_low2) as _), + ("i64", Arc::new(i64_values) as _), + ]) + .unwrap() } /// Create a batch of (f64, utf8_view_low, utf8_view_low, i64) @@ -681,10 +751,10 @@ impl DataGenerator { } /// Create sorted values of high cardinality (~ no duplicates) utf8 values - fn utf8_high_cardinality_values(&mut self) -> Vec> { + fn utf8_high_cardinality_values(&mut self) -> Vec>> { // make random strings let mut input = (0..self.input_size) - .map(|_| Some(self.random_string())) + .map(|_| Some(self.random_string().into())) .collect::>(); input.sort_unstable(); @@ -699,6 +769,26 @@ impl DataGenerator { .map(char::from) .collect::() } + + /// i64 values with the given cardinality (initial order is irrelevant since + /// callers reorder via [`DataProfile::apply`]). + fn i64_values_by(&mut self, card: Cardinality) -> Vec { + match card { + Cardinality::Low => self.i64_values(), + // Full i64 range -> effectively unique (minimal overlap) + Cardinality::High => { + (0..self.input_size).map(|_| self.rng.random()).collect() + } + } + } + + /// utf8 values with the given cardinality. + fn utf8_values_by(&mut self, card: Cardinality) -> Vec>> { + match card { + Cardinality::Low => self.utf8_low_cardinality_values(), + Cardinality::High => self.utf8_high_cardinality_values(), + } + } } /// Splits the `input` tuples randomly into batches of `BATCH_SIZE` distributed across @@ -733,5 +823,247 @@ where .collect() } -criterion_group!(benches, criterion_benchmark); +fn create_single_partition( + input: Vec, + f: F, + batch_size: usize, +) -> Vec +where + F: Fn(Vec) -> RecordBatch, +{ + input + .into_iter() + .chunks(batch_size) + .into_iter() + .map(|x| f(x.collect_vec())) + .collect() +} + +/// Read a duration (seconds, may be fractional) from `var`. panics if set to a value that isn't a number. +fn env_duration(var: &str) -> Option { + let s = std::env::var(var).ok()?; + + let secs = s + .parse::() + .unwrap_or_else(|e| panic!("invalid {var}={s:?}: {e}")); + + Some(Duration::from_secs_f64(secs)) +} + +/// Read a `usize` from `var`. panics if set to a value that isn't an integer. +fn env_usize(var: &str) -> Option { + let s = std::env::var(var).ok()?; + + Some( + s.parse::() + .unwrap_or_else(|e| panic!("invalid {var}={s:?}: {e}")), + ) +} + +type AxisGenerator = Box Vec>; + +/// Benchmarks `SortExec` (at the 1M input size) on single partition across the following axes: +/// 1. Sort columns +/// - single column with a specialized impl (primitive or byte(view)) +/// - multiple columns, which will use fallback impl +/// 2. Number of columns in the record batch - more columns mean more data to +/// copy while reordering and more memory to hold +/// 3. Value cardinality - whether the sort-key values overlap or not +/// 4. Input ordering - already sorted / unsorted / nearly sorted +fn sort_axis_benchmark(c: &mut Criterion) { + let input_size = 1_000_000u64; + let size_label = "1M"; + + const AXIS_BATCH_SIZE: usize = 8192; + + let cases: Vec<(&str, AxisGenerator)> = vec![ + ( + "i64", + Box::new(move |p, card, extra| { + i64_axis(p, card, extra, input_size, AXIS_BATCH_SIZE) + }), + ), + ( + "utf8 view", + Box::new(move |p, card, extra| { + utf8_view_axis(p, card, extra, input_size, AXIS_BATCH_SIZE) + }), + ), + ( + "mixed tuple", + Box::new(move |p, card, extra| { + mixed_tuple_axis(p, card, extra, input_size, AXIS_BATCH_SIZE) + }), + ), + ]; + + let mut group = c.benchmark_group("sort_axis"); + + if let Some(sample_size) = env_usize("SORT_AXIS_SAMPLE_SIZE") { + group.sample_size(sample_size); + } + + if let Some(warm_up_time) = env_duration("SORT_AXIS_WARMUP_SECS") { + group.warm_up_time(warm_up_time); + } + + if let Some(measurement_time) = env_duration("SORT_AXIS_MEASUREMENT_SECS") { + group.measurement_time(measurement_time); + } + + for (name, f) in &cases { + for card in [Cardinality::Low, Cardinality::High] { + for &extra in EXTRA_COLUMN_COUNTS { + for profile in [ + DataProfile::Sorted, + DataProfile::Unsorted, + DataProfile::NearlySorted, + ] { + group.bench_function( + format!( + "sort {name} {size_label} {card:?} cardinality {profile:?} +{extra}cols", + ), + |b| { + let data = f(profile, card, extra); + let case = BenchCase::sort_partitioned(AXIS_BATCH_SIZE, &[data]); + b.iter(move || case.run()) + }, + ); + } + } + } + } + + group.finish(); +} + +/// Single-column i64 batches +fn i64_axis( + profile: DataProfile, + card: Cardinality, + extra: usize, + input_size: u64, + batch_size: usize, +) -> Vec { + let values = profile.apply(DataGenerator::new(input_size).i64_values_by(card)); + let batches = create_single_partition(values, build_i64_batch, batch_size); + with_extra_columns(batches, extra) +} + +/// Single-column utf8 view batches +fn utf8_view_axis( + profile: DataProfile, + card: Cardinality, + extra: usize, + input_size: u64, + batch_size: usize, +) -> Vec { + let values = profile.apply(DataGenerator::new(input_size).utf8_values_by(card)); + let batches = create_single_partition( + values, + |v| build_utf8_view_batch("utf_view", v), + batch_size, + ); + with_extra_columns(batches, extra) +} + +/// Multi-column (f64, utf8, utf8, i64) batches. +fn mixed_tuple_axis( + profile: DataProfile, + card: Cardinality, + extra: usize, + input_size: u64, + batch_size: usize, +) -> Vec { + let mut data_gen = DataGenerator::new(input_size); + let tuples: Vec = data_gen + .i64_values_by(card) + .into_iter() + .zip(data_gen.utf8_values_by(card)) + .zip(data_gen.utf8_values_by(card)) + .zip(data_gen.i64_values_by(card)) + .collect(); + let batches = create_single_partition( + profile.apply(tuples), + build_mixed_tuple_batch, + batch_size, + ); + with_extra_columns(batches, extra) +} + +/// Append `n` extra non-sort-key payload columns to every batch, split across i64, string, string view and dictionary +fn with_extra_columns(batches: Vec, n: usize) -> Vec { + if n == 0 { + return batches; + } + let mut rng = StdRng::seed_from_u64(7); + + type Generator = Box ArrayRef>; + + let generators: Vec = vec![ + Box::new(|data_gen: &mut DataGenerator| { + let arr = Int64Array::from_iter_values(data_gen.i64_values()); + + Arc::new(arr) + }), + Box::new(|data_gen: &mut DataGenerator| { + let values = data_gen.utf8_low_cardinality_values(); + let arr: StringArray = values.iter().map(|item| item.as_deref()).collect(); + + Arc::new(arr) + }), + Box::new(|data_gen: &mut DataGenerator| { + let values = data_gen.utf8_low_cardinality_values(); + let mut builder = + StringViewBuilder::with_capacity(values.len()).with_deduplicate_strings(); + for v in values { + builder.append_option(v.as_deref()); + } + + let arr = builder.finish(); + + Arc::new(arr) + }), + Box::new(|data_gen: &mut DataGenerator| { + let values = data_gen.utf8_low_cardinality_values(); + + let arr: DictionaryArray = + values.iter().map(|item| item.as_deref()).collect(); + + Arc::new(arr) + }), + ]; + + let generator_index = (0..n) + .map(|_| rng.random_range(0..generators.len())) + .collect::>(); + + let mut generator = DataGenerator { input_size: 1, rng }; + + batches + .into_iter() + .map(|batch| { + let num_rows = batch.num_rows(); + let mut fields = batch.schema().fields().iter().cloned().collect::>(); + let mut columns = batch.columns().to_vec(); + generator.input_size = num_rows as u64; + + for (col_index, gen_index) in generator_index.iter().enumerate() { + let gen_fn = &generators[*gen_index]; + + let array = gen_fn(&mut generator); + fields.push(Arc::new(Field::new( + format!("{EXTRA_COLUMN_NAME_PREFIX}{col_index}"), + array.data_type().clone(), + array.logical_null_count() > 0, + ))); + columns.push(array); + } + + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + }) + .collect() +} + +criterion_group!(benches, criterion_benchmark, sort_axis_benchmark); criterion_main!(benches); diff --git a/datafusion/core/benches/sql_planner.rs b/datafusion/core/benches/sql_planner.rs index fcc8da30fedd9..5fae803708edc 100644 --- a/datafusion/core/benches/sql_planner.rs +++ b/datafusion/core/benches/sql_planner.rs @@ -41,11 +41,37 @@ const BENCHMARKS_PATH_1: &str = "../../benchmarks/"; const BENCHMARKS_PATH_2: &str = "./benchmarks/"; const CLICKBENCH_DATA_PATH: &str = "data/hits_partitioned/"; -/// Create a logical plan from the specified sql +/// Create a logical plan from the specified sql (parse + analyze only, NO optimization) fn logical_plan(ctx: &SessionContext, rt: &Runtime, sql: &str) { black_box(rt.block_on(ctx.sql(sql)).unwrap()); } +/// Parse SQL and run the analyzer to get an analyzed (but unoptimized) LogicalPlan. +/// This is the input to the optimizer. +fn analyzed_plan( + ctx: &SessionContext, + rt: &Runtime, + sql: &str, +) -> datafusion_expr::LogicalPlan { + let state = ctx.state(); + let plan = rt.block_on(state.create_logical_plan(sql)).unwrap(); + state + .analyzer() + .execute_and_check(plan, state.config().options(), |_, _| {}) + .unwrap() +} + +/// Run ONLY the optimizer on a pre-analyzed plan. Measures optimizer cost in isolation. +fn optimize_plan(ctx: &SessionContext, plan: &datafusion_expr::LogicalPlan) { + let state = ctx.state(); + black_box( + state + .optimizer() + .optimize(plan.clone(), &state, |_, _| {}) + .unwrap(), + ); +} + /// Create a physical ExecutionPlan (by way of logical plan) fn physical_plan(ctx: &SessionContext, rt: &Runtime, sql: &str) { black_box(rt.block_on(async { @@ -107,15 +133,23 @@ fn create_context() -> SessionContext { /// Register the table definitions as a MemTable with the context and return the /// context -#[expect(clippy::needless_pass_by_value)] fn register_defs(ctx: SessionContext, defs: Vec) -> SessionContext { - defs.iter().for_each(|TableDef { name, schema }| { + for TableDef { + name, + schema, + constraints, + } in defs + { ctx.register_table( - name, - Arc::new(MemTable::try_new(Arc::new(schema.clone()), vec![vec![]]).unwrap()), + &name, + Arc::new( + MemTable::try_new(Arc::new(schema), vec![vec![]]) + .unwrap() + .with_constraints(constraints), + ), ) .unwrap(); - }); + } ctx } @@ -646,6 +680,433 @@ fn criterion_benchmark(c: &mut Criterion) { c.bench_function("with_param_values_many_columns", |b| { benchmark_with_param_values_many_columns(&ctx, &rt, b); }); + + // ========================================================================== + // Optimizer-focused benchmarks + // These benchmarks are designed to stress the logical optimizer with + // varying plan sizes, expression counts, and node type distributions. + // ========================================================================== + + // --- Deep join trees (many plan nodes, few expressions) --- + // Tests optimizer traversal cost as plan node count grows. + // Each join adds ~3 nodes (Join, TableScan, CrossJoin/Filter). + + // Register additional tables for join benchmarks + for i in 3..=16 { + ctx.register_table(format!("j{i}"), create_table_provider("x", 10)) + .unwrap(); + } + + c.bench_function("logical_join_chain_4", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT j3.x0 FROM j3 \ + JOIN j4 ON j3.x0 = j4.x0 \ + JOIN j5 ON j4.x0 = j5.x0 \ + JOIN j6 ON j5.x0 = j6.x0", + ) + }) + }); + + c.bench_function("logical_join_chain_8", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT j3.x0 FROM j3 \ + JOIN j4 ON j3.x0 = j4.x0 \ + JOIN j5 ON j4.x0 = j5.x0 \ + JOIN j6 ON j5.x0 = j6.x0 \ + JOIN j7 ON j6.x0 = j7.x0 \ + JOIN j8 ON j7.x0 = j8.x0 \ + JOIN j9 ON j8.x0 = j9.x0 \ + JOIN j10 ON j9.x0 = j10.x0", + ) + }) + }); + + c.bench_function("logical_join_chain_16", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT j3.x0 FROM j3 \ + JOIN j4 ON j3.x0 = j4.x0 \ + JOIN j5 ON j4.x0 = j5.x0 \ + JOIN j6 ON j5.x0 = j6.x0 \ + JOIN j7 ON j6.x0 = j7.x0 \ + JOIN j8 ON j7.x0 = j8.x0 \ + JOIN j9 ON j8.x0 = j9.x0 \ + JOIN j10 ON j9.x0 = j10.x0 \ + JOIN j11 ON j10.x0 = j11.x0 \ + JOIN j12 ON j11.x0 = j12.x0 \ + JOIN j13 ON j12.x0 = j13.x0 \ + JOIN j14 ON j13.x0 = j14.x0 \ + JOIN j15 ON j14.x0 = j15.x0 \ + JOIN j16 ON j15.x0 = j16.x0 \ + JOIN j3 AS j3b ON j16.x0 = j3b.x0 \ + JOIN j4 AS j4b ON j3b.x0 = j4b.x0", + ) + }) + }); + + // --- Wide expressions (few plan nodes, many expressions) --- + // Tests expression processing overhead in optimizer rules like + // SimplifyExpressions, CommonSubexprEliminate, OptimizeProjections. + + // Many WHERE clauses (filter expressions) + { + let predicates: Vec = (0..50).map(|i| format!("a{i} > 0")).collect(); + let query = format!("SELECT a0 FROM t1 WHERE {}", predicates.join(" AND ")); + c.bench_function("logical_wide_filter_50_predicates", |b| { + b.iter(|| logical_plan(&ctx, &rt, &query)) + }); + } + + { + let predicates: Vec = (0..200).map(|i| format!("a{i} > 0")).collect(); + let query = format!("SELECT a0 FROM t1 WHERE {}", predicates.join(" AND ")); + c.bench_function("logical_wide_filter_200_predicates", |b| { + b.iter(|| logical_plan(&ctx, &rt, &query)) + }); + } + + // Many aggregate expressions + { + let aggs: Vec = + (0..50).map(|i| format!("SUM(a{i}), AVG(a{i})")).collect(); + let query = format!("SELECT {} FROM t1", aggs.join(", ")); + c.bench_function("logical_wide_aggregate_100_exprs", |b| { + b.iter(|| logical_plan(&ctx, &rt, &query)) + }); + } + + // Many CASE WHEN expressions (complex expressions) + { + let cases: Vec = (0..50) + .map(|i| { + format!("CASE WHEN a{i} > 0 THEN a{i} * 2 ELSE a{i} + 1 END AS r{i}") + }) + .collect(); + let query = format!("SELECT {} FROM t1", cases.join(", ")); + c.bench_function("logical_wide_case_50_exprs", |b| { + b.iter(|| logical_plan(&ctx, &rt, &query)) + }); + } + + // --- Mixed: deep plan + wide expressions --- + // This is the worst case for optimizer: many nodes AND many expressions. + + c.bench_function("logical_join_4_with_agg_and_filter", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT j3.x0, SUM(j4.x1), AVG(j5.x2), COUNT(j6.x3), \ + MIN(j3.x4), MAX(j4.x5) \ + FROM j3 \ + JOIN j4 ON j3.x0 = j4.x0 \ + JOIN j5 ON j4.x0 = j5.x0 \ + JOIN j6 ON j5.x0 = j6.x0 \ + WHERE j3.x1 > 0 AND j4.x2 < 100 AND j5.x3 != j6.x4 \ + GROUP BY j3.x0 \ + HAVING SUM(j4.x1) > 10 \ + ORDER BY j3.x0", + ) + }) + }); + + c.bench_function("logical_join_8_with_agg_sort_limit", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT j3.x0, j4.x1, j5.x2, \ + SUM(j6.x3), AVG(j7.x4), COUNT(j8.x5), \ + MIN(j9.x6), MAX(j10.x7) \ + FROM j3 \ + JOIN j4 ON j3.x0 = j4.x0 \ + JOIN j5 ON j4.x0 = j5.x0 \ + JOIN j6 ON j5.x0 = j6.x0 \ + JOIN j7 ON j6.x0 = j7.x0 \ + JOIN j8 ON j7.x0 = j8.x0 \ + JOIN j9 ON j8.x0 = j9.x0 \ + JOIN j10 ON j9.x0 = j10.x0 \ + WHERE j3.x1 > 0 AND j5.x2 < 100 \ + GROUP BY j3.x0, j4.x1, j5.x2 \ + ORDER BY j3.x0 DESC \ + LIMIT 100", + ) + }) + }); + + // --- Subqueries (trigger decorrelation rules) --- + // Tests rules like DecorrelatePredicateSubquery, ScalarSubqueryToJoin. + + c.bench_function("logical_correlated_subquery_exists", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT a0, a1 FROM t1 \ + WHERE EXISTS (SELECT 1 FROM t2 WHERE t2.b0 = t1.a0)", + ) + }) + }); + + c.bench_function("logical_correlated_subquery_in", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT a0, a1 FROM t1 \ + WHERE a0 IN (SELECT b0 FROM t2 WHERE t2.b1 = t1.a1)", + ) + }) + }); + + c.bench_function("logical_scalar_subquery", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT a0, (SELECT MAX(b1) FROM t2 WHERE t2.b0 = t1.a0) AS max_b \ + FROM t1", + ) + }) + }); + + c.bench_function("logical_multiple_subqueries", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT a0, a1 FROM t1 \ + WHERE a0 IN (SELECT b0 FROM t2 WHERE b1 > 0) \ + AND EXISTS (SELECT 1 FROM t2 WHERE t2.b0 = t1.a0 AND t2.b1 < 100) \ + AND a1 > (SELECT AVG(b1) FROM t2)", + ) + }) + }); + + // --- UNION queries (test OptimizeUnions, PropagateEmptyRelation) --- + + c.bench_function("logical_union_4_branches", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT a0, a1 FROM t1 WHERE a0 > 0 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 10 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 20 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 30", + ) + }) + }); + + c.bench_function("logical_union_8_branches", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "SELECT a0, a1 FROM t1 WHERE a0 > 0 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 10 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 20 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 30 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 40 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 50 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 60 \ + UNION ALL SELECT a0, a1 FROM t1 WHERE a0 > 70", + ) + }) + }); + + // --- DISTINCT (test ReplaceDistinctWithAggregate) --- + + c.bench_function("logical_distinct_many_columns", |b| { + let cols: Vec = (0..50).map(|i| format!("a{i}")).collect(); + let query = format!("SELECT DISTINCT {} FROM t1", cols.join(", ")); + b.iter(|| logical_plan(&ctx, &rt, &query)) + }); + + // --- Nested views / CTEs (deeper plan trees) --- + + c.bench_function("logical_nested_cte_4_levels", |b| { + b.iter(|| { + logical_plan( + &ctx, + &rt, + "WITH \ + cte1 AS (SELECT a0, a1, a2 FROM t1 WHERE a0 > 0), \ + cte2 AS (SELECT a0, a1 FROM cte1 WHERE a1 > 0), \ + cte3 AS (SELECT a0 FROM cte2 WHERE a0 < 100), \ + cte4 AS (SELECT a0, COUNT(*) AS cnt FROM cte3 GROUP BY a0) \ + SELECT * FROM cte4 ORDER BY a0 LIMIT 10", + ) + }) + }); + + // --- TPC-H logical plans (uncommented from existing code) --- + // These test real-world query patterns with moderate plan complexity. + + c.bench_function("logical_plan_tpch_all", |b| { + b.iter(|| { + for sql in &all_tpch_sql_queries { + logical_plan(&tpch_ctx, &rt, sql) + } + }) + }); + + c.bench_function("logical_plan_tpcds_all", |b| { + b.iter(|| { + for sql in &all_tpcds_sql_queries { + logical_plan(&tpcds_ctx, &rt, sql) + } + }) + }); + + // ========================================================================== + // Optimizer-only benchmarks + // These measure ONLY the optimizer, not SQL parsing or analysis. + // Plans are pre-parsed and pre-analyzed in setup, then only optimization + // is measured in the benchmark loop. + // ========================================================================== + + // Simple select (baseline: few nodes, few expressions) + { + let plan = analyzed_plan(&ctx, &rt, "SELECT c1 FROM t700"); + c.bench_function("optimizer_select_one_from_700", |b| { + b.iter(|| optimize_plan(&ctx, &plan)) + }); + } + + // Wide select (many expressions, few nodes) + { + let plan = analyzed_plan(&ctx, &rt, "SELECT * FROM t1000"); + c.bench_function("optimizer_select_all_from_1000", |b| { + b.iter(|| optimize_plan(&ctx, &plan)) + }); + } + + // Deep join chains (many nodes, few expressions) + { + let plan = analyzed_plan( + &ctx, + &rt, + "SELECT j3.x0 FROM j3 \ + JOIN j4 ON j3.x0 = j4.x0 \ + JOIN j5 ON j4.x0 = j5.x0 \ + JOIN j6 ON j5.x0 = j6.x0", + ); + c.bench_function("optimizer_join_chain_4", |b| { + b.iter(|| optimize_plan(&ctx, &plan)) + }); + } + + { + let plan = analyzed_plan( + &ctx, + &rt, + "SELECT j3.x0 FROM j3 \ + JOIN j4 ON j3.x0 = j4.x0 \ + JOIN j5 ON j4.x0 = j5.x0 \ + JOIN j6 ON j5.x0 = j6.x0 \ + JOIN j7 ON j6.x0 = j7.x0 \ + JOIN j8 ON j7.x0 = j8.x0 \ + JOIN j9 ON j8.x0 = j9.x0 \ + JOIN j10 ON j9.x0 = j10.x0", + ); + c.bench_function("optimizer_join_chain_8", |b| { + b.iter(|| optimize_plan(&ctx, &plan)) + }); + } + + // Wide filter (many expressions) + { + let predicates: Vec = (0..200).map(|i| format!("a{i} > 0")).collect(); + let query = format!("SELECT a0 FROM t1 WHERE {}", predicates.join(" AND ")); + let plan = analyzed_plan(&ctx, &rt, &query); + c.bench_function("optimizer_wide_filter_200", |b| { + b.iter(|| optimize_plan(&ctx, &plan)) + }); + } + + // Wide aggregate (many expressions) + { + let aggs: Vec = + (0..50).map(|i| format!("SUM(a{i}), AVG(a{i})")).collect(); + let query = format!("SELECT {} FROM t1", aggs.join(", ")); + let plan = analyzed_plan(&ctx, &rt, &query); + c.bench_function("optimizer_wide_aggregate_100", |b| { + b.iter(|| optimize_plan(&ctx, &plan)) + }); + } + + // Subquery (tests decorrelation rules) + { + let plan = analyzed_plan( + &ctx, + &rt, + "SELECT a0, a1 FROM t1 \ + WHERE EXISTS (SELECT 1 FROM t2 WHERE t2.b0 = t1.a0)", + ); + c.bench_function("optimizer_correlated_exists", |b| { + b.iter(|| optimize_plan(&ctx, &plan)) + }); + } + + // Mixed: joins + aggregates + filter + { + let plan = analyzed_plan( + &ctx, + &rt, + "SELECT j3.x0, SUM(j4.x1), AVG(j5.x2), COUNT(j6.x3), \ + MIN(j3.x4), MAX(j4.x5) \ + FROM j3 \ + JOIN j4 ON j3.x0 = j4.x0 \ + JOIN j5 ON j4.x0 = j5.x0 \ + JOIN j6 ON j5.x0 = j6.x0 \ + WHERE j3.x1 > 0 AND j4.x2 < 100 AND j5.x3 != j6.x4 \ + GROUP BY j3.x0 \ + HAVING SUM(j4.x1) > 10 \ + ORDER BY j3.x0", + ); + c.bench_function("optimizer_join_4_with_agg_filter", |b| { + b.iter(|| optimize_plan(&ctx, &plan)) + }); + } + + // TPC-H all queries (optimizer only) + { + let plans: Vec<_> = all_tpch_sql_queries + .iter() + .map(|sql| analyzed_plan(&tpch_ctx, &rt, sql)) + .collect(); + c.bench_function("optimizer_tpch_all", |b| { + b.iter(|| { + for plan in &plans { + optimize_plan(&tpch_ctx, plan) + } + }) + }); + } + + // TPC-DS all queries (optimizer only) + { + let plans: Vec<_> = all_tpcds_sql_queries + .iter() + .map(|sql| analyzed_plan(&tpcds_ctx, &rt, sql)) + .collect(); + c.bench_function("optimizer_tpcds_all", |b| { + b.iter(|| { + for plan in &plans { + optimize_plan(&tpcds_ctx, plan) + } + }) + }); + } } criterion_group!(benches, criterion_benchmark); diff --git a/datafusion/core/benches/sql_planner_extended.rs b/datafusion/core/benches/sql_planner_extended.rs index d4955313c79c3..5bea9860c4be7 100644 --- a/datafusion/core/benches/sql_planner_extended.rs +++ b/datafusion/core/benches/sql_planner_extended.rs @@ -324,6 +324,57 @@ fn build_non_case_left_join_df_with_push_down_filter( rt.block_on(async { ctx.sql(&query).await.unwrap() }) } +/// Join + wide-OR filter + N chained CTEs, each adding one column defined +/// by a depth-K nested CASE ladder over the same input column. Exercises +/// the physical `ProjectionPushdown` rule on long projection chains. +fn build_chained_case_projection_query( + chained_steps: usize, + case_depth: usize, + or_width: usize, +) -> String { + let mut q = String::new(); + q.push_str("WITH s0 AS (\n SELECT l.c0, l.c1 FROM t l LEFT JOIN t r ON l.c0 = r.c0"); + if or_width > 0 { + q.push_str("\n WHERE ("); + for i in 0..or_width { + if i > 0 { + q.push_str(" OR "); + } + let _ = write!(&mut q, "l.c1 = '{i}'"); + } + q.push(')'); + } + q.push_str("\n)"); + + for n in 1..=chained_steps { + q.push_str(",\n"); + let _ = write!(&mut q, "s{n} AS (SELECT *, "); + for d in 0..case_depth { + let _ = write!(&mut q, "CASE WHEN c0 = '{d}' THEN 'label' ELSE "); + } + q.push_str("c0"); + for _ in 0..case_depth { + q.push_str(" END"); + } + let _ = write!(&mut q, " AS d{n} FROM s{prev})", prev = n - 1); + } + + let _ = write!(&mut q, "\nSELECT * FROM s{chained_steps}"); + q +} + +fn build_chained_case_projection_df( + rt: &Runtime, + chained_steps: usize, + case_depth: usize, + or_width: usize, +) -> DataFrame { + let ctx = SessionContext::new(); + register_string_table(&ctx, 100, 1000); + let query = build_chained_case_projection_query(chained_steps, case_depth, or_width); + rt.block_on(async { ctx.sql(&query).await.unwrap() }) +} + fn criterion_benchmark(c: &mut Criterion) { let baseline_ctx = SessionContext::new(); let case_heavy_ctx = SessionContext::new(); @@ -335,12 +386,16 @@ fn criterion_benchmark(c: &mut Criterion) { let df = build_test_data_frame(&baseline_ctx, &rt); let case_heavy_left_join_df = build_case_heavy_left_join_df(&case_heavy_ctx, &rt); - c.bench_function("logical_plan_optimize", |b| { + // really slow :( + let mut group = c.benchmark_group("sample_size_5"); + group.sample_size(5); + group.bench_function("logical_plan_optimize", |b| { b.iter(|| { let df_clone = df.clone(); black_box(rt.block_on(async { df_clone.into_optimized_plan().unwrap() })); }) }); + group.finish(); c.bench_function("logical_plan_optimize_hotspot_case_heavy_left_join", |b| { b.iter(|| { @@ -460,6 +515,16 @@ fn criterion_benchmark(c: &mut Criterion) { } } control_group.finish(); + + let chained_df = build_chained_case_projection_df(&rt, 80, 23, 30); + c.bench_function("physical_plan_chained_case_projection_hotspot", |b| { + b.iter(|| { + let df_clone = chained_df.clone(); + black_box( + rt.block_on(async { df_clone.create_physical_plan().await.unwrap() }), + ); + }) + }); } criterion_group!(benches, criterion_benchmark); diff --git a/datafusion/core/benches/sql_query_with_io.rs b/datafusion/core/benches/sql_query_with_io.rs index fc8caf31acd11..c6600e197374b 100644 --- a/datafusion/core/benches/sql_query_with_io.rs +++ b/datafusion/core/benches/sql_query_with_io.rs @@ -124,8 +124,10 @@ async fn setup_context(object_store: Arc) -> SessionContext { let table_name = table_name(table_id); let file_format = ParquetFormat::default().with_enable_pruning(true); let options = ListingOptions::new(Arc::new(file_format)) - .with_table_partition_cols(vec![(String::from("partition"), DataType::UInt8)]) - .with_target_partitions(THREADS); + .with_table_partition_cols(vec![( + String::from("partition"), + DataType::UInt8, + )]); // make sure we actually find the data let path = format!("data://my_store/{table_name}/"); diff --git a/datafusion/core/benches/struct_query_sql.rs b/datafusion/core/benches/struct_query_sql.rs index 96434fc379ea6..848d5a3c3e5de 100644 --- a/datafusion/core/benches/struct_query_sql.rs +++ b/datafusion/core/benches/struct_query_sql.rs @@ -23,12 +23,11 @@ use arrow::{ use criterion::{Criterion, criterion_group, criterion_main}; use datafusion::prelude::SessionContext; use datafusion::{datasource::MemTable, error::Result}; -use futures::executor::block_on; use std::hint::black_box; use std::sync::Arc; use tokio::runtime::Runtime; -async fn query(ctx: &SessionContext, rt: &Runtime, sql: &str) { +fn query(ctx: &SessionContext, rt: &Runtime, sql: &str) { // execute the query let df = rt.block_on(ctx.sql(sql)).unwrap(); black_box(rt.block_on(df.collect()).unwrap()); @@ -71,7 +70,7 @@ fn criterion_benchmark(c: &mut Criterion) { let rt = Runtime::new().unwrap(); c.bench_function("struct", |b| { - b.iter(|| block_on(query(&ctx, &rt, "select struct(f32, f64) from t"))) + b.iter(|| query(&ctx, &rt, "select struct(f32, f64) from t")) }); } diff --git a/datafusion/core/benches/topk_aggregate.rs b/datafusion/core/benches/topk_aggregate.rs index c78b1ea494407..d8ca0d58b8d21 100644 --- a/datafusion/core/benches/topk_aggregate.rs +++ b/datafusion/core/benches/topk_aggregate.rs @@ -74,7 +74,7 @@ fn test_distinct_schema() -> SchemaRef { Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])) } -async fn create_context( +fn create_context( partition_cnt: i32, sample_cnt: i32, asc: bool, @@ -94,7 +94,7 @@ async fn create_context( Ok(ctx) } -async fn create_context_distinct( +fn create_context_distinct( partition_cnt: i32, sample_cnt: i32, use_topk: bool, @@ -306,12 +306,8 @@ fn assert_utf8_utf8view_match( asc: bool, use_topk: bool, ) { - let ctx_utf8 = rt - .block_on(create_context(partitions, samples, asc, use_topk, false)) - .unwrap(); - let ctx_view = rt - .block_on(create_context(partitions, samples, asc, use_topk, true)) - .unwrap(); + let ctx_utf8 = create_context(partitions, samples, asc, use_topk, false).unwrap(); + let ctx_view = create_context(partitions, samples, asc, use_topk, true).unwrap(); let batches_utf8 = rt .block_on(aggregate_string(ctx_utf8, limit, use_topk)) .unwrap(); @@ -390,15 +386,9 @@ fn criterion_benchmark(c: &mut Criterion) { .name_tpl .replace("{rows}", &total_rows.to_string()) .replace("{limit}", &limit.to_string()); - let ctx = rt - .block_on(create_context( - partitions, - samples, - case.asc, - case.use_topk, - case.use_view, - )) - .unwrap(); + let ctx = + create_context(partitions, samples, case.asc, case.use_topk, case.use_view) + .unwrap(); c.bench_function(&name, |b| { b.iter(|| run(&rt, ctx.clone(), limit, case.use_topk, case.asc)) }); @@ -462,15 +452,9 @@ fn criterion_benchmark(c: &mut Criterion) { } else { format!("string aggregate {total_rows} {scenario} rows [{type_label}]") }; - let ctx = rt - .block_on(create_context( - partitions, - samples, - case.asc, - case.use_topk, - case.use_view, - )) - .unwrap(); + let ctx = + create_context(partitions, samples, case.asc, case.use_topk, case.use_view) + .unwrap(); c.bench_function(&name, |b| { b.iter(|| run_string(&rt, ctx.clone(), limit, case.use_topk)) }); @@ -478,11 +462,7 @@ fn criterion_benchmark(c: &mut Criterion) { // DISTINCT benchmarks for use_topk in [false, true] { - let ctx = rt.block_on(async { - create_context_distinct(partitions, samples, use_topk) - .await - .unwrap() - }); + let ctx = create_context_distinct(partitions, samples, use_topk).unwrap(); let topk_label = if use_topk { "TopK" } else { "no TopK" }; for asc in [false, true] { let dir = if asc { "asc" } else { "desc" }; diff --git a/datafusion/core/src/bin/print_functions_docs.rs b/datafusion/core/src/bin/print_functions_docs.rs index c34865a32d532..86f433ac8e12c 100644 --- a/datafusion/core/src/bin/print_functions_docs.rs +++ b/datafusion/core/src/bin/print_functions_docs.rs @@ -287,7 +287,7 @@ impl DocProvider for WindowUDF { } } -impl DocProvider for Arc { +impl DocProvider for Arc { fn get_name(&self) -> String { self.name().to_string() } diff --git a/datafusion/core/src/dataframe/mod.rs b/datafusion/core/src/dataframe/mod.rs index 0f38988c69405..325ae91d27bbf 100644 --- a/datafusion/core/src/dataframe/mod.rs +++ b/datafusion/core/src/dataframe/mod.rs @@ -48,6 +48,7 @@ use std::sync::Arc; use arrow::array::{Array, ArrayRef, Int64Array, StringArray}; use arrow::compute::{cast, concat}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::util::display::{ArrayFormatter, FormatOptions}; use arrow_schema::FieldRef; use datafusion_common::config::{CsvOptions, JsonOptions}; use datafusion_common::{ @@ -57,13 +58,11 @@ use datafusion_common::{ }; use datafusion_expr::select_expr::SelectExpr; use datafusion_expr::{ - ExplainOption, SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE, case, - dml::InsertOp, - expr::{Alias, ScalarFunction}, - is_null, lit, - utils::COUNT_STAR_EXPANSION, + ExplainOption, ScalarUDF, SortExpr, TableProviderFilterPushDown, UNNAMED_TABLE, case, + dml::InsertOp, is_null, lit, utils::COUNT_STAR_EXPANSION, }; use datafusion_functions::core::coalesce; +use datafusion_functions::math::nanvl; use datafusion_functions_aggregate::expr_fn::{ avg, count, max, median, min, stddev, sum, }; @@ -979,8 +978,13 @@ impl DataFrame { /// Return a new `DataFrame` that has statistics for a DataFrame. /// - /// Only summarizes numeric datatypes at the moment and returns nulls for - /// non numeric datatypes. The output format is modeled after pandas + /// The summary contains the `count`, `null_count`, `mean`, `std`, `min`, + /// `max`, and `median` of each column. `count` and `null_count` are + /// computed for every column; `min` and `max` for every column except + /// `Boolean`; and `mean`, `std`, and `median` only for numeric columns + /// (other columns report `null` for these). `min`/`max` of binary columns + /// (`Binary`, `LargeBinary`, `BinaryView`, `FixedSizeBinary`) are rendered + /// as lowercase hex. The output format is modeled after pandas /// /// # Example /// ``` @@ -1074,9 +1078,7 @@ impl DataFrame { vec![], original_schema_fields .clone() - .filter(|f| { - !matches!(f.data_type(), DataType::Binary | DataType::Boolean) - }) + .filter(|f| !matches!(f.data_type(), DataType::Boolean)) .map(|f| min(ident(f.name())).alias(f.name())) .collect::>(), ), @@ -1085,9 +1087,7 @@ impl DataFrame { vec![], original_schema_fields .clone() - .filter(|f| { - !matches!(f.data_type(), DataType::Binary | DataType::Boolean) - }) + .filter(|f| !matches!(f.data_type(), DataType::Boolean)) .map(|f| max(ident(f.name())).alias(f.name())) .collect::>(), ), @@ -1126,6 +1126,22 @@ impl DataFrame { Arc::new(StringArray::from(vec!["null"])) } else if field.data_type().is_numeric() { cast(column, &DataType::Float64)? + } else if field.data_type().is_binary() { + let formatter = ArrayFormatter::try_new( + column.as_ref(), + &FormatOptions::default(), + )?; + let values: Vec> = (0..column.len()) + .map(|i| { + if column.is_null(i) { + None + } else { + let value = formatter.value(i); + Some(value.to_string()) + } + }) + .collect(); + Arc::new(StringArray::from(values)) } else { cast(column, &DataType::Utf8)? } @@ -1133,7 +1149,8 @@ impl DataFrame { _ => Arc::new(StringArray::from(vec!["null"])), } } - //Handling error when only boolean/binary column, and in other cases + // Handles the case where all columns were filtered out + // (e.g. only boolean columns for mean/std/min/max/median) Err(err) if err.to_string().contains( "Error during planning: \ @@ -1517,7 +1534,7 @@ impl DataFrame { /// # } pub async fn to_string(self) -> Result { let options = self.session_state.config().options().format.clone(); - let arrow_options: arrow::util::display::FormatOptions = (&options).try_into()?; + let arrow_options: FormatOptions = (&options).try_into()?; let registry = self.session_state.extension_type_registry(); let formatter_factory = DFArrayFormatterFactory::new(Arc::clone(registry)); @@ -2422,7 +2439,7 @@ impl DataFrame { } /// Fill null values in specified columns with a given value - /// If no columns are specified (empty vector), applies to all columns + /// If no columns are specified (empty slice), applies to all columns /// Only fills if the value can be cast to the column's type /// /// # Arguments @@ -2441,17 +2458,70 @@ impl DataFrame { /// .read_csv("tests/data/example.csv", CsvReadOptions::new()) /// .await?; /// // Fill nulls in only columns "a" and "c": - /// let df = df.fill_null(ScalarValue::from(0), vec!["a".to_owned(), "c".to_owned()])?; + /// let df = df.fill_null(&ScalarValue::from(0), &["a", "c"])?; /// // Fill nulls across all columns: - /// let df = df.fill_null(ScalarValue::from(0), vec![])?; + /// let df = df.fill_null(&ScalarValue::from(0), &[])?; + /// # Ok(()) + /// # } + /// ``` + pub fn fill_null(&self, value: &ScalarValue, columns: &[&str]) -> Result { + self.fill_columns(value, columns, &coalesce(), |_| true) + } + + // Helper to find columns from names + fn find_columns(&self, names: &[impl AsRef]) -> Result> { + let schema = self.logical_plan().schema(); + names + .iter() + .map(|name| { + let name = name.as_ref(); + schema + .field_with_name(None, name) + .cloned() + .map_err(|_| plan_datafusion_err!("Column '{}' not found", name)) + }) + .collect() + } + + /// Fill NaN values in specified floating-point columns with a given value + /// If no columns are specified (empty slice), applies to all columns + /// Only floating-point columns are affected; other columns are left unchanged + /// Only fills if the value can be cast to the column's type + /// + /// # Arguments + /// * `value` - Value to fill NaNs with + /// * `columns` - List of column names to fill. If empty, fills all columns. + /// + /// # Example + /// ``` + /// # use datafusion::prelude::*; + /// # use datafusion::error::Result; + /// # use datafusion_common::ScalarValue; + /// # #[tokio::main] + /// # async fn main() -> Result<()> { + /// let ctx = SessionContext::new(); + /// let df = ctx + /// .read_csv("tests/data/example.csv", CsvReadOptions::new()) + /// .await?; + /// // Fill NaN in only columns "a" and "c": + /// let df = df.fill_nan(&ScalarValue::from(0.0), &["a", "c"])?; + /// // Fill NaN across all columns: + /// let df = df.fill_nan(&ScalarValue::from(0.0), &[])?; /// # Ok(()) /// # } /// ``` - #[expect(clippy::needless_pass_by_value)] - pub fn fill_null( + pub fn fill_nan(&self, value: &ScalarValue, columns: &[&str]) -> Result { + self.fill_columns(value, columns, &nanvl(), |field| { + field.data_type().is_floating() + }) + } + + fn fill_columns( &self, - value: ScalarValue, - columns: Vec, + value: &ScalarValue, + columns: &[impl AsRef], + func: &Arc, + applies: impl Fn(&FieldRef) -> bool, ) -> Result { let cols = if columns.is_empty() { self.logical_plan() @@ -2461,28 +2531,21 @@ impl DataFrame { .map(Arc::clone) .collect() } else { - self.find_columns(&columns)? + self.find_columns(columns)? }; - // Create projections for each column let projections = self .logical_plan() .schema() .fields() .iter() .map(|field| { - if cols.contains(field) { + if cols.contains(field) && applies(field) { // Try to cast fill value to column type. If the cast fails, fallback to the original column. match value.clone().cast_to(field.data_type()) { - Ok(fill_value) => Expr::Alias(Alias { - expr: Box::new(Expr::ScalarFunction(ScalarFunction { - func: coalesce(), - args: vec![col(field.name()), lit(fill_value)], - })), - relation: None, - name: field.name().to_string(), - metadata: None, - }), + Ok(fill_value) => func + .call(vec![col(field.name()), lit(fill_value)]) + .alias(field.name()), Err(_) => col(field.name()), } } else { @@ -2494,20 +2557,6 @@ impl DataFrame { self.clone().select(projections) } - // Helper to find columns from names - fn find_columns(&self, names: &[String]) -> Result> { - let schema = self.logical_plan().schema(); - names - .iter() - .map(|name| { - schema - .field_with_name(None, name) - .cloned() - .map_err(|_| plan_datafusion_err!("Column '{}' not found", name)) - }) - .collect() - } - /// Find qualified columns for this dataframe from names /// /// # Arguments diff --git a/datafusion/core/src/datasource/file_format/csv.rs b/datafusion/core/src/datasource/file_format/csv.rs index a068b4f5c0413..90d7eb3b41388 100644 --- a/datafusion/core/src/datasource/file_format/csv.rs +++ b/datafusion/core/src/datasource/file_format/csv.rs @@ -45,6 +45,7 @@ mod tests { use datafusion_datasource::file_format::FileFormat; use datafusion_datasource::write::BatchSerializer; use datafusion_expr::{col, lit}; + use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::{ExecutionPlan, collect}; use arrow::array::{ @@ -215,9 +216,16 @@ mod tests { assert_eq!(tt_batches, 50 /* 100/2 */); // test metadata - assert_eq!(exec.partition_statistics(None)?.num_rows, Precision::Absent); assert_eq!( - exec.partition_statistics(None)?.total_byte_size, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .num_rows, + Precision::Absent + ); + assert_eq!( + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .total_byte_size, Precision::Absent ); @@ -583,8 +591,7 @@ mod tests { //convert compressed_stream to decoded_stream let decoded_stream = compressed_csv - .read_to_delimited_chunks_from_stream(compressed_stream.unwrap()) - .await; + .read_to_delimited_chunks_from_stream(compressed_stream.unwrap()); let (schema, records_read) = compressed_csv .infer_schema_from_stream(&session_state, records_to_read, decoded_stream) .await?; @@ -697,7 +704,7 @@ mod tests { ) -> Result { let df = ctx.sql(&format!("EXPLAIN {sql}")).await?; let result = df.collect().await?; - let plan = format!("{}", &pretty_format_batches(&result)?); + let plan = format!("{}", pretty_format_batches(&result)?); let re = Regex::new(r"DataSourceExec: file_groups=\{(\d+) group").unwrap(); diff --git a/datafusion/core/src/datasource/file_format/json.rs b/datafusion/core/src/datasource/file_format/json.rs index 5b3e22705620e..1f6f27242e723 100644 --- a/datafusion/core/src/datasource/file_format/json.rs +++ b/datafusion/core/src/datasource/file_format/json.rs @@ -36,6 +36,7 @@ mod tests { BatchDeserializer, DecoderDeserializer, DeserializerOutput, }; use datafusion_datasource::file_format::FileFormat; + use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::{ExecutionPlan, collect}; use arrow::compute::concat_batches; @@ -117,9 +118,16 @@ mod tests { assert_eq!(tt_batches, 6 /* 12/2 */); // test metadata - assert_eq!(exec.partition_statistics(None)?.num_rows, Precision::Absent); assert_eq!( - exec.partition_statistics(None)?.total_byte_size, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .num_rows, + Precision::Absent + ); + assert_eq!( + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .total_byte_size, Precision::Absent ); @@ -225,7 +233,7 @@ mod tests { .collect() .await?; - let plan = format!("{}", &pretty::pretty_format_batches(&result)?); + let plan = format!("{}", pretty::pretty_format_batches(&result)?); let re = Regex::new(r"file_groups=\{(\d+) group").unwrap(); diff --git a/datafusion/core/src/datasource/file_format/mod.rs b/datafusion/core/src/datasource/file_format/mod.rs index b04238ebc9b37..c46b472bd6404 100644 --- a/datafusion/core/src/datasource/file_format/mod.rs +++ b/datafusion/core/src/datasource/file_format/mod.rs @@ -67,7 +67,7 @@ pub(crate) mod test_util { .await? }; - let table_schema = TableSchema::new(file_schema.clone(), vec![]); + let table_schema = TableSchema::from(&file_schema); let statistics = format .infer_stats(state, &store, file_schema.clone(), &meta) diff --git a/datafusion/core/src/datasource/file_format/options.rs b/datafusion/core/src/datasource/file_format/options.rs index bd0ac36087381..8ef780ea1e972 100644 --- a/datafusion/core/src/datasource/file_format/options.rs +++ b/datafusion/core/src/datasource/file_format/options.rs @@ -34,6 +34,7 @@ use crate::error::Result; use crate::execution::context::{SessionConfig, SessionState}; use arrow::datatypes::{DataType, Schema, SchemaRef}; +use datafusion_catalog_listing::SchemaSource; use datafusion_common::config::{ConfigFileDecryptionProperties, TableOptions}; use datafusion_common::{ DEFAULT_ARROW_EXTENSION, DEFAULT_AVRO_EXTENSION, DEFAULT_CSV_EXTENSION, @@ -595,6 +596,11 @@ pub trait ReadOptions<'a> { table_path: ListingTableUrl, ) -> Result; + /// Returns whether the read schema was inferred or specified. + fn schema_source(&self) -> SchemaSource { + SchemaSource::Specified + } + /// helper function to reduce repetitive code. Infers the schema from sources if not provided. Infinite data sources not supported through this function. async fn _get_resolved_schema( &'a self, @@ -620,7 +626,7 @@ pub trait ReadOptions<'a> { impl ReadOptions<'_> for CsvReadOptions<'_> { fn to_listing_options( &self, - config: &SessionConfig, + _config: &SessionConfig, table_options: TableOptions, ) -> ListingOptions { let file_format = CsvFormat::default() @@ -639,7 +645,6 @@ impl ReadOptions<'_> for CsvReadOptions<'_> { ListingOptions::new(Arc::new(file_format)) .with_file_extension(self.file_extension) - .with_session_config_options(config) .with_table_partition_cols(self.table_partition_cols.clone()) .with_file_sort_order(self.file_sort_order.clone()) } @@ -653,6 +658,10 @@ impl ReadOptions<'_> for CsvReadOptions<'_> { self._get_resolved_schema(config, state, table_path, self.schema) .await } + + fn schema_source(&self) -> SchemaSource { + schema_source_from_option(self.schema) + } } #[cfg(feature = "parquet")] @@ -660,7 +669,7 @@ impl ReadOptions<'_> for CsvReadOptions<'_> { impl ReadOptions<'_> for ParquetReadOptions<'_> { fn to_listing_options( &self, - config: &SessionConfig, + _config: &SessionConfig, table_options: TableOptions, ) -> ListingOptions { let mut options = table_options.parquet; @@ -685,7 +694,6 @@ impl ReadOptions<'_> for ParquetReadOptions<'_> { .with_file_extension(self.file_extension) .with_table_partition_cols(self.table_partition_cols.clone()) .with_file_sort_order(self.file_sort_order.clone()) - .with_session_config_options(config) } async fn get_resolved_schema( @@ -697,13 +705,17 @@ impl ReadOptions<'_> for ParquetReadOptions<'_> { self._get_resolved_schema(config, state, table_path, self.schema) .await } + + fn schema_source(&self) -> SchemaSource { + schema_source_from_option(self.schema) + } } #[async_trait] impl ReadOptions<'_> for JsonReadOptions<'_> { fn to_listing_options( &self, - config: &SessionConfig, + _config: &SessionConfig, table_options: TableOptions, ) -> ListingOptions { let file_format = JsonFormat::default() @@ -714,7 +726,6 @@ impl ReadOptions<'_> for JsonReadOptions<'_> { ListingOptions::new(Arc::new(file_format)) .with_file_extension(self.file_extension) - .with_session_config_options(config) .with_table_partition_cols(self.table_partition_cols.clone()) .with_file_sort_order(self.file_sort_order.clone()) } @@ -728,6 +739,10 @@ impl ReadOptions<'_> for JsonReadOptions<'_> { self._get_resolved_schema(config, state, table_path, self.schema) .await } + + fn schema_source(&self) -> SchemaSource { + schema_source_from_option(self.schema) + } } #[cfg(feature = "avro")] @@ -735,14 +750,13 @@ impl ReadOptions<'_> for JsonReadOptions<'_> { impl ReadOptions<'_> for AvroReadOptions<'_> { fn to_listing_options( &self, - config: &SessionConfig, + _config: &SessionConfig, _table_options: TableOptions, ) -> ListingOptions { let file_format = AvroFormat; ListingOptions::new(Arc::new(file_format)) .with_file_extension(self.file_extension) - .with_session_config_options(config) .with_table_partition_cols(self.table_partition_cols.clone()) } @@ -755,20 +769,23 @@ impl ReadOptions<'_> for AvroReadOptions<'_> { self._get_resolved_schema(config, state, table_path, self.schema) .await } + + fn schema_source(&self) -> SchemaSource { + schema_source_from_option(self.schema) + } } #[async_trait] impl ReadOptions<'_> for ArrowReadOptions<'_> { fn to_listing_options( &self, - config: &SessionConfig, + _config: &SessionConfig, _table_options: TableOptions, ) -> ListingOptions { let file_format = ArrowFormat; ListingOptions::new(Arc::new(file_format)) .with_file_extension(self.file_extension) - .with_session_config_options(config) .with_table_partition_cols(self.table_partition_cols.clone()) } @@ -781,4 +798,16 @@ impl ReadOptions<'_> for ArrowReadOptions<'_> { self._get_resolved_schema(config, state, table_path, self.schema) .await } + + fn schema_source(&self) -> SchemaSource { + schema_source_from_option(self.schema) + } +} + +fn schema_source_from_option(schema: Option<&Schema>) -> SchemaSource { + if schema.is_some() { + SchemaSource::Specified + } else { + SchemaSource::Inferred + } } diff --git a/datafusion/core/src/datasource/file_format/parquet.rs b/datafusion/core/src/datasource/file_format/parquet.rs index c977deab32aa4..0f5db4a057d76 100644 --- a/datafusion/core/src/datasource/file_format/parquet.rs +++ b/datafusion/core/src/datasource/file_format/parquet.rs @@ -141,6 +141,7 @@ mod tests { use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::dml::InsertOp; + use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ExecutionPlan, collect}; @@ -715,12 +716,16 @@ mod tests { // test metadata assert_eq!( - exec.partition_statistics(None)?.num_rows, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .num_rows, Precision::Exact(8) ); // TODO correct byte size: https://github.com/apache/datafusion/issues/14936 assert_eq!( - exec.partition_statistics(None)?.total_byte_size, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .total_byte_size, Precision::Absent, ); @@ -764,11 +769,15 @@ mod tests { // note: even if the limit is set, the executor rounds up to the batch size assert_eq!( - exec.partition_statistics(None)?.num_rows, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .num_rows, Precision::Exact(8) ); assert_eq!( - exec.partition_statistics(None)?.total_byte_size, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .total_byte_size, Precision::Absent, ); let batches = collect(exec, task_ctx).await?; diff --git a/datafusion/core/src/datasource/listing/table.rs b/datafusion/core/src/datasource/listing/table.rs index d14ec1f56dce2..56a5d5779596c 100644 --- a/datafusion/core/src/datasource/listing/table.rs +++ b/datafusion/core/src/datasource/listing/table.rs @@ -87,15 +87,13 @@ impl ListingTableConfigExt for ListingTableConfig { let listing_file_extension = if let Some(compression_type) = maybe_compression_type { - format!("{}.{}", &file_extension, &compression_type) + format!("{file_extension}.{compression_type}") } else { file_extension }; - let listing_options = ListingOptions::new(file_format) - .with_file_extension(listing_file_extension) - .with_target_partitions(state.config().target_partitions()) - .with_collect_stat(state.config().collect_statistics()); + let listing_options = + ListingOptions::new(file_format).with_file_extension(listing_file_extension); Ok(self.with_listing_options(listing_options)) } @@ -125,7 +123,7 @@ mod tests { }, }; use arrow::{compute::SortOptions, record_batch::RecordBatch}; - use arrow_schema::{DataType, Field, Schema, SchemaRef}; + use arrow_schema::{DataType, Field, Schema, SchemaRef, TimeUnit}; use datafusion_catalog::TableProvider; use datafusion_catalog_listing::{ ListingOptions, ListingTable, ListingTableConfig, SchemaSource, @@ -139,12 +137,18 @@ mod tests { use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_format::FileFormat; use datafusion_expr::dml::InsertOp; - use datafusion_expr::{BinaryExpr, LogicalPlanBuilder, Operator}; + use datafusion_expr::{ + BinaryExpr, LogicalPlanBuilder, Operator, Partitioning as LogicalPartitioning, + RangePartitioning as LogicalRangePartitioning, + }; use datafusion_physical_expr::PhysicalSortExpr; - use datafusion_physical_expr::expressions::binary; + use datafusion_physical_expr::expressions::{Column, binary}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::empty::EmptyExec; - use datafusion_physical_plan::{ExecutionPlanProperties, collect}; + use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; + use datafusion_physical_plan::{ + ExecutionPlanProperties, Partitioning, RangePartitioning, SplitPoint, collect, + }; use std::collections::HashMap; use std::io::Write; use std::sync::Arc; @@ -177,6 +181,21 @@ mod tests { .collect() } + fn listing_table_with_files( + ctx: &SessionContext, + files: &[&str], + table_path: &str, + options: ListingOptions, + schema: Schema, + ) -> Result { + register_test_store(ctx, &files.iter().map(|f| (*f, 10)).collect::>()); + + let config = ListingTableConfig::new(ListingTableUrl::parse(table_path)?) + .with_listing_options(options) + .with_schema(Arc::new(schema)); + ListingTable::try_new(config) + } + #[tokio::test] async fn test_schema_source_tracking_comprehensive() -> Result<()> { let ctx = SessionContext::new(); @@ -247,11 +266,15 @@ mod tests { // test metadata assert_eq!( - exec.partition_statistics(None)?.num_rows, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .num_rows, Precision::Exact(8) ); assert_eq!( - exec.partition_statistics(None)?.total_byte_size, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new())? + .total_byte_size, Precision::Absent, ); @@ -372,7 +395,9 @@ mod tests { #[tokio::test] async fn read_empty_table() -> Result<()> { - let ctx = SessionContext::new(); + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(4), + ); let path = String::from("table/p1=v1/file.json"); register_test_store(&ctx, &[(&path, 100)]); @@ -381,8 +406,7 @@ mod tests { let opt = ListingOptions::new(Arc::new(format)) .with_file_extension(ext) - .with_table_partition_cols(vec![(String::from("p1"), DataType::Utf8)]) - .with_target_partitions(4); + .with_table_partition_cols(vec![(String::from("p1"), DataType::Utf8)]); let table_path = ListingTableUrl::parse("test:///table/")?; let file_schema = @@ -438,12 +462,13 @@ mod tests { output_partitioning: usize, file_ext: Option<&str>, ) -> Result<()> { - let ctx = SessionContext::new(); + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(target_partitions), + ); register_test_store(&ctx, &files.iter().map(|f| (*f, 10)).collect::>()); let opt = ListingOptions::new(Arc::new(JsonFormat::default())) - .with_file_extension_opt(file_ext) - .with_target_partitions(target_partitions); + .with_file_extension_opt(file_ext); let schema = Schema::new(vec![Field::new("a", DataType::Boolean, false)]); @@ -470,12 +495,13 @@ mod tests { output_partitioning: usize, file_ext: Option<&str>, ) -> Result<()> { - let ctx = SessionContext::new(); + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(target_partitions), + ); register_test_store(&ctx, &files.iter().map(|f| (*f, 10)).collect::>()); let opt = ListingOptions::new(Arc::new(JsonFormat::default())) - .with_file_extension_opt(file_ext) - .with_target_partitions(target_partitions); + .with_file_extension_opt(file_ext); let schema = Schema::new(vec![Field::new("a", DataType::Boolean, false)]); @@ -505,7 +531,9 @@ mod tests { output_partitioning: usize, file_ext: Option<&str>, ) -> Result<()> { - let ctx = SessionContext::new(); + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(target_partitions), + ); let (store, _) = make_test_store_and_state( &files.iter().map(|f| (*f, 10)).collect::>(), ); @@ -514,7 +542,8 @@ mod tests { .state() .config_options() .execution - .meta_fetch_concurrency; + .meta_fetch_concurrency + .get(); let expected_concurrency = files.len().min(meta_fetch_concurrency); let head_concurrency_store = ensure_head_concurrency(store, expected_concurrency); @@ -523,9 +552,7 @@ mod tests { let format = JsonFormat::default(); - let opt = ListingOptions::new(Arc::new(format)) - .with_file_extension_opt(file_ext) - .with_target_partitions(target_partitions); + let opt = ListingOptions::new(Arc::new(format)).with_file_extension_opt(file_ext); let schema = Schema::new(vec![Field::new("a", DataType::Boolean, false)]); @@ -1286,6 +1313,236 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_list_files_uses_declared_output_partitioning_count() -> Result<()> { + let files = ["bucket/key-prefix/file0", "bucket/key-prefix/file1"]; + + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(1), + ); + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(LogicalPartitioning::Range( + LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::from(10i32)]), + SplitPoint::new(vec![ScalarValue::from(20i32)]), + SplitPoint::new(vec![ScalarValue::from(30i32)]), + ], + )?, + ))); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new("a", DataType::Int32, false)]), + )?; + + let result = table.list_files_for_scan(&ctx.state(), &[], None).await?; + let group_sizes = result + .file_groups + .iter() + .map(|group| group.len()) + .collect::>(); + + assert_eq!(group_sizes, vec![1, 1, 0, 0]); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_normalizes_split_point_types() -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(123_000_000_000), + None, + )])], + )?); + let expected_output_partitioning = + Partitioning::Range(RangePartitioning::try_new( + LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions::default(), + )]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::TimestampSecond( + Some(123), + None, + )])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new( + "a", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]), + )?; + + let scan = table.scan(&ctx.state(), None, &[], None).await?; + assert_eq!(scan.output_partitioning(), &expected_output_partitioning); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_rejects_invalid_split_point_type() + -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "not-an-int".to_string(), + ))])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new("a", DataType::Int32, false)]), + )?; + + let err = table.scan(&ctx.state(), None, &[], None).await.unwrap_err(); + assert_contains!( + err.to_string(), + "Range output partitioning split point 0 value 0 with type Utf8 cannot be represented exactly as ordering expression type Int32" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_range_output_partitioning_rejects_lossy_timestamp_split_point() + -> Result<()> { + let files = ["bucket/key-prefix/file0"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::TimestampNanosecond( + Some(123_456), + None, + )])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_output_partitioning(Some(output_partitioning)); + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/key-prefix/", + opt, + Schema::new(vec![Field::new( + "a", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]), + )?; + + let err = table.scan(&ctx.state(), None, &[], None).await.unwrap_err(); + assert_contains!( + err.to_string(), + "Range output partitioning split point 0 value 0 with type Timestamp(ns) cannot be represented exactly as ordering expression type Timestamp(s)" + ); + + Ok(()) + } + + #[tokio::test] + async fn test_partition_filter_preserves_declared_output_partitioning() -> Result<()> + { + let files = ["bucket/test/pid=1/file1", "bucket/test/pid=2/file2"]; + + let ctx = SessionContext::new(); + let output_partitioning = + LogicalPartitioning::Range(LogicalRangePartitioning::try_new( + vec![col("pid").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::from(2i32)])], + )?); + let expected_output_partitioning = + Partitioning::Range(RangePartitioning::try_new( + LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("pid", 1)), + SortOptions::default(), + )]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::from(2i32)])], + )?); + + let opt = ListingOptions::new(Arc::new(JsonFormat::default())) + .with_file_extension_opt(Some("")) + .with_table_partition_cols(vec![("pid".to_string(), DataType::Int32)]) + .with_output_partitioning(Some(output_partitioning.clone())); + + let table = listing_table_with_files( + &ctx, + &files, + "test:///bucket/test/", + opt, + Schema::new(vec![Field::new("a", DataType::Boolean, false)]), + )?; + + let unfiltered = table.scan(&ctx.state(), None, &[], None).await?; + assert_eq!( + unfiltered.output_partitioning(), + &expected_output_partitioning + ); + + let filter = Expr::eq(col("pid"), lit(2_i32)); + let file_groups = table + .list_files_for_scan(&ctx.state(), std::slice::from_ref(&filter), None) + .await? + .file_groups + .into_iter() + .map(|group| { + group + .into_inner() + .into_iter() + .map(|file| file.path().to_string()) + .collect::>() + }) + .collect::>(); + assert_eq!( + file_groups, + vec![ + Vec::::new(), + vec!["bucket/test/pid=2/file2".to_string()] + ] + ); + + let filtered = table.scan(&ctx.state(), None, &[filter], None).await?; + assert_eq!( + filtered.output_partitioning(), + &expected_output_partitioning + ); + + Ok(()) + } + #[tokio::test] async fn test_listing_table_prunes_extra_files_in_hive() -> Result<()> { let files = [ @@ -1296,7 +1553,9 @@ mod tests { "bucket/test/other/file5", ]; - let ctx = SessionContext::new(); + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(1), + ); register_test_store(&ctx, &files.iter().map(|f| (*f, 10)).collect::>()); let opt = ListingOptions::new(Arc::new(JsonFormat::default())) @@ -1341,10 +1600,12 @@ mod tests { let filename = format!("{}/{}", testdata, "alltypes_plain.parquet"); let table_path = ListingTableUrl::parse(filename)?; - let ctx = SessionContext::new(); - let state = ctx.state(); + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(1), + ); + let mut state = ctx.state(); - // Test 1: Default behavior - stats not collected + // Test 1: Default behavior - stats collected let opt_default = ListingOptions::new(Arc::new(ParquetFormat::default())); let schema_default = opt_default.infer_schema(&state, &table_path).await?; let config_default = ListingTableConfig::new(table_path.clone()) @@ -1355,52 +1616,66 @@ mod tests { let exec_default = table_default.scan(&state, None, &[], None).await?; assert_eq!( - exec_default.partition_statistics(None)?.num_rows, - Precision::Absent + StatisticsContext::new() + .compute(exec_default.as_ref(), &StatisticsArgs::new())? + .num_rows, + Precision::Exact(8) ); // TODO correct byte size: https://github.com/apache/datafusion/issues/14936 assert_eq!( - exec_default.partition_statistics(None)?.total_byte_size, + StatisticsContext::new() + .compute(exec_default.as_ref(), &StatisticsArgs::new())? + .total_byte_size, Precision::Absent ); - // Test 2: Explicitly disable stats - let opt_disabled = ListingOptions::new(Arc::new(ParquetFormat::default())) - .with_collect_stat(false); - let schema_disabled = opt_disabled.infer_schema(&state, &table_path).await?; + // Test 2: Explicitly disable stats via session config + let cfg = state.config_mut(); + cfg.options_mut().execution.collect_statistics = false; + let opt = ListingOptions::new(Arc::new(ParquetFormat::default())); + let schema_disabled = opt.infer_schema(&state, &table_path).await?; let config_disabled = ListingTableConfig::new(table_path.clone()) - .with_listing_options(opt_disabled) + .with_listing_options(opt) .with_schema(schema_disabled); let table_disabled = ListingTable::try_new(config_disabled)?; let exec_disabled = table_disabled.scan(&state, None, &[], None).await?; assert_eq!( - exec_disabled.partition_statistics(None)?.num_rows, + StatisticsContext::new() + .compute(exec_disabled.as_ref(), &StatisticsArgs::new())? + .num_rows, Precision::Absent ); assert_eq!( - exec_disabled.partition_statistics(None)?.total_byte_size, + StatisticsContext::new() + .compute(exec_disabled.as_ref(), &StatisticsArgs::new())? + .total_byte_size, Precision::Absent ); - // Test 3: Explicitly enable stats - let opt_enabled = ListingOptions::new(Arc::new(ParquetFormat::default())) - .with_collect_stat(true); - let schema_enabled = opt_enabled.infer_schema(&state, &table_path).await?; + // Test 3: Re-enable stats via session config + let cfg = state.config_mut(); + cfg.options_mut().execution.collect_statistics = true; + let opt = ListingOptions::new(Arc::new(ParquetFormat::default())); + let schema_enabled = opt.infer_schema(&state, &table_path).await?; let config_enabled = ListingTableConfig::new(table_path) - .with_listing_options(opt_enabled) + .with_listing_options(opt) .with_schema(schema_enabled); let table_enabled = ListingTable::try_new(config_enabled)?; let exec_enabled = table_enabled.scan(&state, None, &[], None).await?; assert_eq!( - exec_enabled.partition_statistics(None)?.num_rows, + StatisticsContext::new() + .compute(exec_enabled.as_ref(), &StatisticsArgs::new())? + .num_rows, Precision::Exact(8) ); // TODO correct byte size: https://github.com/apache/datafusion/issues/14936 assert_eq!( - exec_enabled.partition_statistics(None)?.total_byte_size, + StatisticsContext::new() + .compute(exec_enabled.as_ref(), &StatisticsArgs::new())? + .total_byte_size, Precision::Absent, ); @@ -1456,14 +1731,16 @@ mod tests { #[tokio::test] async fn test_basic_table_scan() -> Result<()> { - let ctx = SessionContext::new(); + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_collect_statistics(false), + ); // Test basic table creation and scanning let path = "table/file.json"; register_test_store(&ctx, &[(path, 10)]); let format = JsonFormat::default(); - let opt = ListingOptions::new(Arc::new(format)).with_collect_stat(false); + let opt = ListingOptions::new(Arc::new(format)); let schema = Schema::new(vec![Field::new("a", DataType::Boolean, false)]); let table_path = ListingTableUrl::parse("test:///table/")?; diff --git a/datafusion/core/src/datasource/listing_table_factory.rs b/datafusion/core/src/datasource/listing_table_factory.rs index a1eb7ffb64b7d..68f6743189447 100644 --- a/datafusion/core/src/datasource/listing_table_factory.rs +++ b/datafusion/core/src/datasource/listing_table_factory.rs @@ -27,9 +27,11 @@ use crate::datasource::listing::{ }; use crate::execution::context::SessionState; -use arrow::datatypes::DataType; +use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::{Result, config_datafusion_err}; -use datafusion_common::{ToDFSchema, arrow_datafusion_err, plan_err}; +use datafusion_common::{ + ToDFSchema, arrow_datafusion_err, internal_datafusion_err, plan_err, +}; use datafusion_expr::CreateExternalTable; use async_trait::async_trait; @@ -71,20 +73,66 @@ impl TableProviderFactory for ListingTableFactory { ))? .create(session_state, &cmd.options)?; - let mut table_path = - ListingTableUrl::parse(&cmd.location)?.with_table_ref(cmd.name.clone()); - let file_extension = match table_path.is_collection() { - // Setting the extension to be empty instead of allowing the default extension seems - // odd, but was done to ensure existing behavior isn't modified. It seems like this - // could be refactored to either use the default extension or set the fully expected - // extension when compression is included (e.g. ".csv.gz") - true => "", - false => &get_extension(cmd.location.as_str()), + let table_paths = cmd + .locations + .iter() + .map(|location| { + Ok(ListingTableUrl::parse(location)?.with_table_ref(cmd.name.clone())) + }) + .collect::>>()?; + let Some(first_path) = table_paths.first() else { + return plan_err!("CREATE EXTERNAL TABLE requires at least one location"); }; - let mut options = ListingOptions::new(file_format) - .with_session_config_options(session_state.config()) - .with_file_extension(file_extension); + let mut seen_paths = HashSet::with_capacity(table_paths.len()); + if let Some(duplicate) = table_paths.iter().find(|path| !seen_paths.insert(*path)) + { + return plan_err!( + "Duplicate location '{}' in CREATE EXTERNAL TABLE", + duplicate.as_str() + ); + } + + // `ListingTable` resolves a single object store (from the first location) + // and scans every location with it, so locations spanning different + // object stores would silently read the wrong data. Reading across + // object stores is intentionally not supported (see + // https://github.com/apache/datafusion/issues/16303); reject it here with + // a clear error rather than producing incorrect results at scan time. + let object_store_url = first_path.object_store(); + if let Some(other) = table_paths + .iter() + .find(|path| path.object_store() != object_store_url) + { + return plan_err!( + "All locations of a CREATE EXTERNAL TABLE must be on the same \ + object store, but found '{}' and '{}'", + object_store_url.as_str(), + other.object_store().as_str() + ); + } + + // With a single location the historical extension handling is kept. With + // more than one location the files may have different extensions, so the + // extension filter is left empty and the explicit paths/globs are used + // as provided. + let file_extension = if table_paths.len() == 1 { + match first_path.is_collection() { + // Setting the extension to be empty instead of allowing the default extension seems + // odd, but was done to ensure existing behavior isn't modified. It seems like this + // could be refactored to either use the default extension or set the fully expected + // extension when compression is included (e.g. ".csv.gz") + true => String::new(), + false => get_extension(&cmd.locations[0]), + } + } else { + String::new() + }; + let mut options = + ListingOptions::new(file_format).with_file_extension(file_extension); + + // Partition columns are derived from the first location; all locations + // are expected to share the same partitioning. let (provided_schema, table_partition_cols) = if cmd.schema.fields().is_empty() { let infer_parts = session_state .config_options() @@ -92,7 +140,7 @@ impl TableProviderFactory for ListingTableFactory { .listing_table_factory_infer_partitions; let part_cols = if cmd.table_partition_cols.is_empty() && infer_parts { options - .infer_partitions(session_state, &table_path) + .infer_partitions(session_state, first_path) .await? .into_iter() } else { @@ -142,36 +190,75 @@ impl TableProviderFactory for ListingTableFactory { options = options.with_table_partition_cols(table_partition_cols); - options - .validate_partitions(session_state, &table_path) - .await?; + // Validate partitions against every location before any glob rewriting. + for table_path in &table_paths { + options + .validate_partitions(session_state, table_path) + .await?; + } - let resolved_schema = match provided_schema { + let (resolved_table_paths, resolved_schema) = match provided_schema { // We will need to check the table columns against the schema // this is done so that we can do an ORDER BY for external table creation // specifically for parquet file format. // See: https://github.com/apache/datafusion/issues/7317 None => { - // if the folder then rewrite a file path as 'path/*.parquet' - // to only read the files the reader can understand - if table_path.is_folder() && table_path.get_glob().is_none() { - // Since there are no files yet to infer an actual extension, - // derive the pattern based on compression type. - // So for gzipped CSV the pattern is `*.csv.gz` - let glob = match options.format.compression_type() { - Some(compression) => { - match options.format.get_ext_with_compression(&compression) { - // Use glob based on `FileFormat` extension - Ok(ext) => format!("*.{ext}"), - // Fallback to `file_type`, if not supported by `FileFormat` - Err(_) => format!("*.{}", cmd.file_type.to_lowercase()), + let mut resolved_paths = Vec::with_capacity(table_paths.len()); + let mut inferred_schema: Option<(String, SchemaRef)> = None; + for mut table_path in table_paths { + // if the folder then rewrite a file path as 'path/*.parquet' + // to only read the files the reader can understand + if table_path.is_folder() && table_path.get_glob().is_none() { + // Since there are no files yet to infer an actual extension, + // derive the pattern based on compression type. + // So for gzipped CSV the pattern is `*.csv.gz` + let glob = match options.format.compression_type() { + Some(compression) => { + match options + .format + .get_ext_with_compression(&compression) + { + // Use glob based on `FileFormat` extension + Ok(ext) => format!("*.{ext}"), + // Fallback to `file_type`, if not supported by `FileFormat` + Err(_) => { + format!("*.{}", cmd.file_type.to_lowercase()) + } + } } + None => format!("*.{}", cmd.file_type.to_lowercase()), + }; + table_path = table_path.with_glob(glob.as_ref())?; + } + let schema = options.infer_schema(session_state, &table_path).await?; + // All locations must resolve to the same fields. Schema + // and field metadata may differ between files without + // changing the fields read by the table. + let location = table_path.to_string(); + match &inferred_schema { + None => inferred_schema = Some((location, schema)), + Some((existing_location, existing)) + if !schemas_have_same_fields(existing, &schema) => + { + return plan_err!( + "All locations of a CREATE EXTERNAL TABLE must have the \ + same schema, but schema inferred from '{}' differs from \ + schema inferred from '{}'", + location, + existing_location + ); } - None => format!("*.{}", cmd.file_type.to_lowercase()), - }; - table_path = table_path.with_glob(glob.as_ref())?; + Some(_) => {} + } + resolved_paths.push(table_path); } - let schema = options.infer_schema(session_state, &table_path).await?; + // `table_paths` was guaranteed non-empty above, so the loop ran + // at least once and `inferred_schema` is always `Some` here. + let (_, schema) = inferred_schema.ok_or_else(|| { + internal_datafusion_err!( + "no schema could be inferred from the provided locations" + ) + })?; let df_schema = Arc::clone(&schema).to_dfschema()?; let column_refs: HashSet<_> = cmd .order_exprs @@ -186,11 +273,11 @@ impl TableProviderFactory for ListingTableFactory { } } - schema + (resolved_paths, schema) } - Some(s) => s, + Some(s) => (table_paths, s), }; - let config = ListingTableConfig::new(table_path) + let config = ListingTableConfig::new_with_multi_paths(resolved_table_paths) .with_listing_options(options.with_file_sort_order(cmd.order_exprs.clone())) .with_schema(resolved_schema); let provider = ListingTable::try_new(config)? @@ -222,6 +309,19 @@ fn get_extension(path: &str) -> String { } } +fn schemas_have_same_fields(left: &SchemaRef, right: &SchemaRef) -> bool { + left.fields().len() == right.fields().len() + && left + .fields() + .iter() + .zip(right.fields()) + .all(|(left, right)| { + left.name() == right.name() + && left.data_type() == right.data_type() + && left.is_nullable() == right.is_nullable() + }) +} + #[cfg(test)] mod tests { use super::*; @@ -229,21 +329,60 @@ mod tests { datasource::file_format::csv::CsvFormat, execution::context::SessionContext, test_util::parquet_test_data, }; - use datafusion_execution::cache::CacheAccessor; - use datafusion_execution::cache::cache_manager::CacheManagerConfig; - use datafusion_execution::cache::file_statistics_cache::DefaultFileStatisticsCache; + use arrow::datatypes::{Field, Schema}; + use datafusion_execution::cache::cache_manager::{ + CacheManagerConfig, DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, + }; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use glob::Pattern; use std::collections::HashMap; use std::fs; use std::fs::File; - use std::path::PathBuf; + use std::path::{Path, PathBuf}; use datafusion_common::parsers::CompressionTypeVariant; use datafusion_common::{DFSchema, TableReference}; + use datafusion_execution::cache::Cache; + use datafusion_execution::cache::default_cache::DefaultCache; use datafusion_expr::registry::ExtensionTypeRegistryRef; + fn factory_and_state() -> (ListingTableFactory, SessionState) { + let factory = ListingTableFactory::new(); + let context = SessionContext::new(); + let state = context.state(); + (factory, state) + } + + fn write_csv(path: &Path, contents: &str) { + fs::write(path, contents).unwrap(); + } + + fn csv_cmd_with_locations(paths: &[&Path]) -> CreateExternalTable { + let locations = paths + .iter() + .map(|path| path.to_str().unwrap().to_string()) + .collect::>(); + + CreateExternalTable::builder( + TableReference::bare("foo"), + locations[0].clone(), + "csv", + Arc::new(DFSchema::empty()), + ) + .with_locations(locations) + .with_options(HashMap::from([("format.has_header".into(), "true".into())])) + .build() + } + + fn assert_error_contains(error: impl std::fmt::Display, expected: &str) { + let error = error.to_string(); + assert!( + error.contains(expected), + "expected error to contain '{expected}', got: {error}" + ); + } + #[tokio::test] async fn test_create_using_non_std_file_ext() { let csv_file = tempfile::Builder::new() @@ -474,6 +613,151 @@ mod tests { assert!(listing_options.table_partition_cols.is_empty()); } + #[tokio::test] + async fn test_create_with_multiple_locations() { + let dir = tempfile::tempdir().unwrap(); + let file_a = dir.path().join("file_a.csv"); + let file_b = dir.path().join("file_b.csv"); + write_csv(&file_a, "c1,c2\n1,a\n2,b\n"); + write_csv(&file_b, "c1,c2\n3,c\n"); + + let (factory, state) = factory_and_state(); + let cmd = csv_cmd_with_locations(&[&file_a, &file_b]); + + let table_provider = factory.create(&state, &cmd).await.unwrap(); + let listing_table = table_provider.downcast_ref::().unwrap(); + + // Both locations are registered as table paths + assert_eq!(2, listing_table.table_paths().len()); + + // Schema is inferred from the files and shared across both locations + let field_names: Vec<_> = listing_table + .schema() + .fields() + .iter() + .map(|f| f.name().clone()) + .collect(); + assert_eq!(field_names, vec!["c1".to_string(), "c2".to_string()]); + } + + #[tokio::test] + async fn test_create_with_duplicate_locations_errors() { + let dir = tempfile::tempdir().unwrap(); + let file = dir.path().join("file.csv"); + write_csv(&file, "c1,c2\n1,a\n"); + + let (factory, state) = factory_and_state(); + let cmd = csv_cmd_with_locations(&[&file, &file]); + let err = factory.create(&state, &cmd).await.unwrap_err(); + assert_error_contains(err, "Duplicate location"); + } + + #[tokio::test] + async fn test_create_with_overlapping_locations_reads_each_file_once() { + let dir = tempfile::tempdir().unwrap(); + let file_a = dir.path().join("file_a.csv"); + let file_b = dir.path().join("file_b.csv"); + write_csv(&file_a, "c1,c2\n1,a\n"); + write_csv(&file_b, "c1,c2\n2,b\n"); + + let (factory, state) = factory_and_state(); + let cmd = csv_cmd_with_locations(&[dir.path(), file_a.as_path()]); + let table_provider = factory.create(&state, &cmd).await.unwrap(); + let listing_table = table_provider.downcast_ref::().unwrap(); + + let listed_files = listing_table + .list_files_for_scan(&state, &[], None) + .await + .unwrap() + .file_groups + .iter() + .map(|group| group.len()) + .sum::(); + assert_eq!(listed_files, 2); + } + + #[tokio::test] + async fn test_create_with_multiple_locations_mismatched_schema_errors() { + let dir = tempfile::tempdir().unwrap(); + let file_a = dir.path().join("file_a.csv"); + let file_b = dir.path().join("file_b.csv"); + write_csv(&file_a, "c1,c2\n1,a\n"); + // Different column names -> different inferred schema + write_csv(&file_b, "x1,x2\n1,a\n"); + + let (factory, state) = factory_and_state(); + let cmd = csv_cmd_with_locations(&[&file_a, &file_b]); + let err = factory.create(&state, &cmd).await.unwrap_err(); + assert_error_contains(err, "same schema"); + } + + #[test] + fn test_schema_comparison_ignores_schema_metadata() { + let fields = + vec![ + Field::new("c1", DataType::Int32, true).with_metadata(HashMap::from([( + "field_source".to_string(), + "a".to_string(), + )])), + ]; + let schema_a = Arc::new(Schema::new_with_metadata( + fields.clone(), + HashMap::from([("source".to_string(), "a".to_string())]), + )); + let schema_b = + Arc::new(Schema::new_with_metadata( + vec![Field::new("c1", DataType::Int32, true).with_metadata( + HashMap::from([("field_source".to_string(), "b".to_string())]), + )], + HashMap::from([("source".to_string(), "b".to_string())]), + )); + let schema_c = + Arc::new(Schema::new(vec![Field::new("c2", DataType::Int32, true)])); + + assert_ne!(schema_a, schema_b); + assert!(schemas_have_same_fields(&schema_a, &schema_b)); + assert!(!schemas_have_same_fields(&schema_a, &schema_c)); + } + + #[tokio::test] + async fn test_create_with_no_locations_errors() { + let (factory, state) = factory_and_state(); + + let cmd = CreateExternalTable::builder( + TableReference::bare("foo"), + "unused", + "csv", + Arc::new(DFSchema::empty()), + ) + .with_locations(vec![]) + .build(); + + let err = factory.create(&state, &cmd).await.unwrap_err(); + assert_error_contains(err, "at least one location"); + } + + #[tokio::test] + async fn test_create_with_locations_on_different_stores_errors() { + let (factory, state) = factory_and_state(); + + // Two locations on different object stores (different buckets) are not + // supported: ListingTable would scan both against the first store. + let cmd = CreateExternalTable::builder( + TableReference::bare("foo"), + "s3://bucket_a/file.parquet", + "parquet", + Arc::new(DFSchema::empty()), + ) + .with_locations(vec![ + "s3://bucket_a/file.parquet".to_string(), + "s3://bucket_b/file.parquet".to_string(), + ]) + .build(); + + let err = factory.create(&state, &cmd).await.unwrap_err(); + assert_error_contains(err, "same object store"); + } + #[tokio::test] async fn test_statistics_cache_prewarming() { let factory = ListingTableFactory::new(); @@ -484,7 +768,8 @@ mod tests { .to_string(); // Test with collect_statistics enabled - let file_statistics_cache = Arc::new(DefaultFileStatisticsCache::default()); + let file_statistics_cache = + Arc::new(DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT)); let cache_config = CacheManagerConfig::default() .with_file_statistics_cache(Some(file_statistics_cache.clone())); let runtime = RuntimeEnvBuilder::new() @@ -514,7 +799,8 @@ mod tests { ); // Test with collect_statistics disabled - let file_statistics_cache = Arc::new(DefaultFileStatisticsCache::default()); + let file_statistics_cache = + Arc::new(DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT)); let cache_config = CacheManagerConfig::default() .with_file_statistics_cache(Some(file_statistics_cache.clone())); let runtime = RuntimeEnvBuilder::new() @@ -552,6 +838,7 @@ mod tests { use datafusion_execution::config::SessionConfig; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; + use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList}; use std::any::Any; use std::collections::HashMap; @@ -567,6 +854,9 @@ mod tests { fn config(&self) -> &SessionConfig { unimplemented!() } + fn catalog_list(&self) -> Arc { + Arc::new(EmptyCatalogProviderList) + } async fn create_physical_plan( &self, _logical_plan: &datafusion_expr::LogicalPlan, @@ -587,7 +877,7 @@ mod tests { } fn higher_order_functions( &self, - ) -> &HashMap> { + ) -> &HashMap> { unimplemented!() } fn aggregate_functions( diff --git a/datafusion/core/src/datasource/memory_test.rs b/datafusion/core/src/datasource/memory_test.rs index c7721cafb02ea..d7311c1d9c960 100644 --- a/datafusion/core/src/datasource/memory_test.rs +++ b/datafusion/core/src/datasource/memory_test.rs @@ -28,7 +28,7 @@ mod tests { use arrow::record_batch::RecordBatch; use arrow_schema::SchemaRef; use datafusion_catalog::TableProvider; - use datafusion_common::{DataFusionError, Result}; + use datafusion_common::{Constraint, Constraints, DataFusionError, Result}; use datafusion_expr::LogicalPlanBuilder; use datafusion_expr::dml::InsertOp; use futures::StreamExt; @@ -103,6 +103,57 @@ mod tests { Ok(()) } + /// Builds a single-batch [`MemTable`] over an `(a, b)` schema, optionally + /// attaching the given constraints. + fn source_table(constraints: Option) -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![4, 5, 6])), + ], + )?; + let table = MemTable::try_new(schema, vec![vec![batch]])?; + Ok(match constraints { + Some(constraints) => table.with_constraints(constraints), + None => table, + }) + } + + #[tokio::test] + async fn test_load_preserves_constraints() -> Result<()> { + let session_ctx = SessionContext::new(); + let constraints = + Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]); + + // Single partition + let source = Arc::new(source_table(Some(constraints.clone()))?); + let loaded = MemTable::load(source, None, &session_ctx.state()).await?; + assert_eq!(loaded.constraints(), Some(&constraints)); + + // Multiple partitions + let source = Arc::new(source_table(Some(constraints.clone()))?); + let loaded = MemTable::load(source, Some(2), &session_ctx.state()).await?; + assert_eq!(loaded.constraints(), Some(&constraints)); + + Ok(()) + } + + #[tokio::test] + async fn test_load_without_constraints() -> Result<()> { + let session_ctx = SessionContext::new(); + + let source = Arc::new(source_table(None)?); + let loaded = MemTable::load(source, None, &session_ctx.state()).await?; + assert_eq!(loaded.constraints(), Some(&Constraints::default())); + + Ok(()) + } + #[tokio::test] async fn test_invalid_projection() -> Result<()> { let session_ctx = SessionContext::new(); diff --git a/datafusion/core/src/datasource/physical_plan/avro.rs b/datafusion/core/src/datasource/physical_plan/avro.rs index 2954a47403299..c9ee2cc407783 100644 --- a/datafusion/core/src/datasource/physical_plan/avro.rs +++ b/datafusion/core/src/datasource/physical_plan/avro.rs @@ -34,7 +34,7 @@ mod tests { use datafusion_common::{Result, ScalarValue, test_util}; use datafusion_datasource::file_format::FileFormat; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; - use datafusion_datasource::{PartitionedFile, TableSchema}; + use datafusion_datasource::{PartitionedFile, TableSchemaBuilder}; use datafusion_datasource_avro::AvroFormat; use datafusion_datasource_avro::source::AvroSource; use datafusion_execution::object_store::ObjectStoreUrl; @@ -223,10 +223,13 @@ mod tests { partitioned_file.partition_values = vec![ScalarValue::from("2021-10-26")]; let projection = Some(vec![0, 1, file_schema.fields().len(), 2]); - let table_schema = TableSchema::new( - file_schema.clone(), - vec![Arc::new(Field::new("date", DataType::Utf8, false))], - ); + let table_schema = TableSchemaBuilder::from(file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "date", + DataType::Utf8, + false, + ))]) + .build(); let source = Arc::new(AvroSource::new(table_schema.clone())); let conf = FileScanConfigBuilder::new(object_store_url, source) // select specific columns of the files as well as the partitioning diff --git a/datafusion/core/src/datasource/physical_plan/csv.rs b/datafusion/core/src/datasource/physical_plan/csv.rs index 82c47b6c7281c..56642d583e414 100644 --- a/datafusion/core/src/datasource/physical_plan/csv.rs +++ b/datafusion/core/src/datasource/physical_plan/csv.rs @@ -122,7 +122,7 @@ mod tests { quote: b'"', ..Default::default() }; - let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema)); + let table_schema = TableSchema::from(&file_schema); let source = Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options)); let config = @@ -194,7 +194,7 @@ mod tests { quote: b'"', ..Default::default() }; - let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema)); + let table_schema = TableSchema::from(&file_schema); let source = Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options)); let config = @@ -265,7 +265,7 @@ mod tests { quote: b'"', ..Default::default() }; - let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema)); + let table_schema = TableSchema::from(&file_schema); let source = Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options)); let config = @@ -335,7 +335,7 @@ mod tests { quote: b'"', ..Default::default() }; - let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema)); + let table_schema = TableSchema::from(&file_schema); let source = Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options)); let config = @@ -371,7 +371,7 @@ mod tests { file_compression_type: FileCompressionType, ) -> Result<()> { use datafusion_common::ScalarValue; - use datafusion_datasource::TableSchema; + use datafusion_datasource::TableSchemaBuilder; let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); @@ -400,10 +400,13 @@ mod tests { quote: b'"', ..Default::default() }; - let table_schema = TableSchema::new( - Arc::clone(&file_schema), - vec![Arc::new(Field::new("date", DataType::Utf8, false))], - ); + let table_schema = TableSchemaBuilder::from(&file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "date", + DataType::Utf8, + false, + ))]) + .build(); let source = Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options)); let config = @@ -508,7 +511,7 @@ mod tests { quote: b'"', ..Default::default() }; - let table_schema = TableSchema::from_file_schema(Arc::clone(&file_schema)); + let table_schema = TableSchema::from(&file_schema); let source = Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options)); let config = diff --git a/datafusion/core/src/datasource/physical_plan/parquet.rs b/datafusion/core/src/datasource/physical_plan/parquet.rs index 6f38df46e3d2e..d562bbe8490f4 100644 --- a/datafusion/core/src/datasource/physical_plan/parquet.rs +++ b/datafusion/core/src/datasource/physical_plan/parquet.rs @@ -54,7 +54,7 @@ mod tests { use datafusion_datasource::source::DataSourceExec; use datafusion_datasource::file::FileSource; - use datafusion_datasource::{PartitionedFile, TableSchema}; + use datafusion_datasource::{PartitionedFile, TableSchemaBuilder}; use datafusion_datasource_parquet::source::ParquetSource; use datafusion_datasource_parquet::{ DefaultParquetFileReaderFactory, ParquetFileReaderFactory, ParquetFormat, @@ -62,10 +62,10 @@ mod tests { use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::{Expr, col, lit, when}; use datafusion_physical_expr::planner::logical2physical; - use datafusion_physical_plan::analyze::AnalyzeExec; + use datafusion_physical_plan::analyze::AnalyzeExecBuilder; use datafusion_physical_plan::collect; use datafusion_physical_plan::metrics::{ - ExecutionPlanMetricsSet, MetricType, MetricValue, MetricsSet, + ExecutionPlanMetricsSet, MetricValue, MetricsSet, }; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; @@ -231,21 +231,22 @@ mod tests { let parquet_exec = self.build_parquet_exec(file_group.clone(), Arc::clone(&parquet_source)); - let analyze_exec = Arc::new(AnalyzeExec::new( - false, - false, - vec![MetricType::Summary, MetricType::Dev], - None, - // use a new ParquetSource to avoid sharing execution metrics - self.build_parquet_exec( - file_group.clone(), - self.build_file_source(Arc::clone(table_schema)), - ), - Arc::new(Schema::new(vec![ - Field::new("plan_type", DataType::Utf8, true), - Field::new("plan", DataType::Utf8, true), - ])), - )); + let analyze_exec = Arc::new( + AnalyzeExecBuilder::new( + false, + false, + // use a new ParquetSource to avoid sharing execution metrics + self.build_parquet_exec( + file_group.clone(), + self.build_file_source(Arc::clone(table_schema)), + ), + Arc::new(Schema::new(vec![ + Field::new("plan_type", DataType::Utf8, true), + Field::new("plan", DataType::Utf8, true), + ])), + ) + .build(), + ); let session_ctx = SessionContext::new(); let task_ctx = session_ctx.task_ctx(); @@ -1642,9 +1643,8 @@ mod tests { ), ]); - let table_schema = TableSchema::new( - Arc::clone(&schema), - vec![ + let table_schema = TableSchemaBuilder::from(&schema) + .with_table_partition_cols(vec![ Arc::new(Field::new("year", DataType::Utf8, false)), Arc::new(Field::new("month", DataType::UInt8, false)), Arc::new(Field::new( @@ -1655,8 +1655,8 @@ mod tests { ), false, )), - ], - ); + ]) + .build(); let source = Arc::new(ParquetSource::new(table_schema.clone())); let config = FileScanConfigBuilder::new(object_store_url, source) .with_file(partitioned_file) diff --git a/datafusion/core/src/execution/context/mod.rs b/datafusion/core/src/execution/context/mod.rs index 67dbe6b7402ed..5b287f103abdd 100644 --- a/datafusion/core/src/execution/context/mod.rs +++ b/datafusion/core/src/execution/context/mod.rs @@ -64,6 +64,7 @@ use datafusion_catalog::memory::MemorySchemaProvider; use datafusion_catalog::{ DynamicFileCatalog, TableFunction, TableFunctionImpl, UrlTableFactory, }; +use datafusion_catalog_listing::SchemaSource; use datafusion_common::config::{ConfigField, ConfigOptions}; use datafusion_common::metadata::ScalarAndMetadata; use datafusion_common::{ @@ -75,12 +76,12 @@ use datafusion_common::{ }; pub use datafusion_execution::TaskContext; use datafusion_execution::cache::cache_manager::{ - DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, DEFAULT_LIST_FILES_CACHE_TTL, - DEFAULT_METADATA_CACHE_LIMIT, + DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, + DEFAULT_LIST_FILES_CACHE_TTL, DEFAULT_METADATA_CACHE_LIMIT, }; pub use datafusion_execution::config::SessionConfig; use datafusion_execution::disk_manager::{ - DEFAULT_MAX_TEMP_DIRECTORY_SIZE, DiskManagerBuilder, + DEFAULT_MAX_SPILL_MERGE_FAN_IN, DEFAULT_MAX_TEMP_DIRECTORY_SIZE, DiskManagerBuilder, }; use datafusion_execution::registry::SerializerRegistry; use datafusion_expr::HigherOrderUDF; @@ -102,7 +103,6 @@ use datafusion_session::SessionStore; use async_trait::async_trait; use chrono::{DateTime, Utc}; -use datafusion_execution::cache::file_statistics_cache::DEFAULT_FILE_STATISTICS_MEMORY_LIMIT; use object_store::ObjectStore; use parking_lot::RwLock; use url::Url; @@ -166,7 +166,7 @@ where /// * Create a [`DataFrame`] from a CSV or Parquet data source. /// * Register a CSV or Parquet data source as a table that can be referenced from a SQL query. /// * Register a custom data source that can be referenced from a SQL query. -/// * Execution a SQL query +/// * Execute a SQL query /// /// # Example: DataFrame API /// @@ -686,8 +686,8 @@ impl SessionContext { pub async fn execute_logical_plan(&self, plan: LogicalPlan) -> Result { match plan { LogicalPlan::Ddl(ddl) => { - // Box::pin avoids allocating the stack space within this function's frame - // for every one of these individual async functions, decreasing the risk of + // Box async DDL handlers to avoid reserving space for all of their + // futures in this function's state machine, decreasing the risk of // stack overflows. match ddl { DdlStatement::CreateExternalTable(cmd) => { @@ -702,32 +702,26 @@ impl SessionContext { Box::pin(self.create_view(cmd)).await } DdlStatement::CreateCatalogSchema(cmd) => { - Box::pin(self.create_catalog_schema(cmd)).await - } - DdlStatement::CreateCatalog(cmd) => { - Box::pin(self.create_catalog(cmd)).await + self.create_catalog_schema(cmd) } + DdlStatement::CreateCatalog(cmd) => self.create_catalog(cmd), DdlStatement::DropTable(cmd) => Box::pin(self.drop_table(cmd)).await, DdlStatement::DropView(cmd) => Box::pin(self.drop_view(cmd)).await, - DdlStatement::DropCatalogSchema(cmd) => { - Box::pin(self.drop_schema(cmd)).await - } + DdlStatement::DropCatalogSchema(cmd) => self.drop_schema(cmd), DdlStatement::CreateFunction(cmd) => { - Box::pin(self.create_function(cmd)).await - } - DdlStatement::DropFunction(cmd) => { - Box::pin(self.drop_function(cmd)).await + Box::pin(self.create_function(*cmd)).await } + DdlStatement::DropFunction(cmd) => self.drop_function(&cmd), ddl => Ok(DataFrame::new(self.state(), LogicalPlan::Ddl(ddl))), } } // TODO what about the other statements (like TransactionStart and TransactionEnd) LogicalPlan::Statement(Statement::SetVariable(stmt)) => { - self.set_variable(stmt).await?; + self.set_variable(stmt)?; self.return_empty_dataframe() } LogicalPlan::Statement(Statement::ResetVariable(stmt)) => { - self.reset_variable(stmt).await?; + self.reset_variable(stmt)?; self.return_empty_dataframe() } LogicalPlan::Statement(Statement::Prepare(Prepare { @@ -986,7 +980,7 @@ impl SessionContext { Ok(()) } - async fn create_catalog_schema(&self, cmd: CreateCatalogSchema) -> Result { + fn create_catalog_schema(&self, cmd: CreateCatalogSchema) -> Result { let CreateCatalogSchema { schema_name, if_not_exists, @@ -1027,7 +1021,7 @@ impl SessionContext { } } - async fn create_catalog(&self, cmd: CreateCatalog) -> Result { + fn create_catalog(&self, cmd: CreateCatalog) -> Result { let CreateCatalog { catalog_name, if_not_exists, @@ -1077,7 +1071,7 @@ impl SessionContext { } } - async fn drop_schema(&self, cmd: DropCatalogSchema) -> Result { + fn drop_schema(&self, cmd: DropCatalogSchema) -> Result { let DropCatalogSchema { name, if_exists: allow_missing, @@ -1112,7 +1106,7 @@ impl SessionContext { exec_err!("Schema '{schema_ref}' doesn't exist.") } - async fn set_variable(&self, stmt: SetVariable) -> Result<()> { + fn set_variable(&self, stmt: SetVariable) -> Result<()> { let SetVariable { variable, value, .. } = stmt; @@ -1147,7 +1141,7 @@ impl SessionContext { Ok(()) } - async fn reset_variable(&self, stmt: ResetVariable) -> Result<()> { + fn reset_variable(&self, stmt: ResetVariable) -> Result<()> { let variable = stmt.variable; if variable.starts_with("datafusion.runtime.") { return self.reset_runtime_variable(&variable); @@ -1207,6 +1201,14 @@ impl SessionContext { let limit = Self::parse_capacity_limit(variable, value)?; builder.with_file_statistics_cache_limit(limit) } + "max_spill_merge_fan_in" => { + let fan_in = value.parse::().map_err(|e| { + DataFusionError::Plan(format!( + "Failed to parse non-negative integer from '{variable}', value '{value}': {e}" + )) + })?; + builder.with_max_spill_merge_fan_in(fan_in) + } _ => return plan_err!("Unknown runtime configuration: {variable}"), // Remember to update `reset_runtime_variable()` when adding new options }; @@ -1251,6 +1253,10 @@ impl SessionContext { DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, ); } + "max_spill_merge_fan_in" => { + builder = + builder.with_max_spill_merge_fan_in(DEFAULT_MAX_SPILL_MERGE_FAN_IN); + } _ => return plan_err!("Unknown runtime configuration: {variable}"), }; *state = SessionStateBuilder::from(state.clone()) @@ -1284,7 +1290,11 @@ impl SessionContext { if limit.trim().is_empty() { return Err(plan_datafusion_err!("Empty limit value found!")); } - let (number, unit) = limit.split_at(limit.len() - 1); + let (unit_start, unit) = limit + .char_indices() + .next_back() + .ok_or_else(|| plan_datafusion_err!("Empty limit value found!"))?; + let number = &limit[..unit_start]; let number: f64 = number.parse().map_err(|_| { plan_datafusion_err!("Failed to parse number from memory limit '{limit}'") })?; @@ -1295,9 +1305,9 @@ impl SessionContext { } match unit { - "K" => Ok((number * 1024.0) as usize), - "M" => Ok((number * 1024.0 * 1024.0) as usize), - "G" => Ok((number * 1024.0 * 1024.0 * 1024.0) as usize), + 'K' => Ok((number * 1024.0) as usize), + 'M' => Ok((number * 1024.0 * 1024.0) as usize), + 'G' => Ok((number * 1024.0 * 1024.0 * 1024.0) as usize), _ => plan_err!("Unsupported unit '{unit}' in memory limit '{limit}'"), } } @@ -1327,7 +1337,10 @@ impl SessionContext { if limit == "0" { return Ok(0); } - let (number, unit) = limit.split_at(limit.len() - 1); + let (unit_start, unit) = limit.char_indices().next_back().ok_or_else(|| { + plan_datafusion_err!("Empty limit value found for '{config_name}'") + })?; + let number = &limit[..unit_start]; let number: f64 = number.parse().map_err(|_| { plan_datafusion_err!( "Failed to parse number from '{config_name}', limit '{limit}'" @@ -1340,9 +1353,9 @@ impl SessionContext { } match unit { - "K" => Ok((number * 1024.0) as usize), - "M" => Ok((number * 1024.0 * 1024.0) as usize), - "G" => Ok((number * 1024.0 * 1024.0 * 1024.0) as usize), + 'K' => Ok((number * 1024.0) as usize), + 'M' => Ok((number * 1024.0 * 1024.0) as usize), + 'G' => Ok((number * 1024.0 * 1024.0 * 1024.0) as usize), _ => plan_err!( "Unsupported unit '{unit}' in '{config_name}', limit '{limit}'. \ Unit must be one of: 'K', 'M', 'G'" @@ -1361,14 +1374,20 @@ impl SessionContext { let mut seconds = None; for duration in duration.split_inclusive(&['m', 's']) { - let (number, unit) = duration.split_at(duration.len() - 1); + let (unit_start, unit) = + duration.char_indices().next_back().ok_or_else(|| { + plan_datafusion_err!( + "Duration should not be empty or blank for '{config_name}'" + ) + })?; + let number = &duration[..unit_start]; let number: u64 = number.parse().map_err(|_| { plan_datafusion_err!("Failed to parse number from duration '{duration}' for '{config_name}'") })?; match unit { - "m" if minutes.is_none() && seconds.is_none() => minutes = Some(number), - "s" if seconds.is_none() => seconds = Some(number), + 'm' if minutes.is_none() && seconds.is_none() => minutes = Some(number), + 's' if seconds.is_none() => seconds = Some(number), other => plan_err!( "Invalid duration unit: '{other}'. The unit must be either 'm' (minutes), or 's' (seconds), and be in the correct order for '{config_name}'" )?, @@ -1448,7 +1467,7 @@ impl SessionContext { && table_provider.table_type() == table_type { schema.deregister_table(&table)?; - self.invalidate_caches(&Some(table_ref.clone()), table_type)?; + self.invalidate_caches(&table_ref, table_type)?; return Ok(true); } Ok(false) @@ -1456,7 +1475,7 @@ impl SessionContext { fn invalidate_caches( &self, - table_ref: &Option, + table_ref: &TableReference, table_type: TableType, ) -> Result<()> { if table_type == TableType::Base { @@ -1505,7 +1524,7 @@ impl SessionContext { self.return_empty_dataframe() } - async fn drop_function(&self, stmt: DropFunction) -> Result { + fn drop_function(&self, stmt: &DropFunction) -> Result { // we don't know function type at this point // decision has been made to drop all functions let mut dropped = false; @@ -1626,7 +1645,7 @@ impl SessionContext { /// - `SELECT "my_HIGHER_ORDER_FUNC"(x)` will look for a function named `"my_HIGHER_ORDER_FUNC"` /// /// Any functions registered with the function name or its aliases will be overwritten with this new function - pub fn register_higher_order_function(&self, f: Arc) { + pub fn register_higher_order_function(&self, f: Arc) { let mut state = self.state.write(); state.register_higher_order_function(f).ok(); } @@ -1725,12 +1744,20 @@ impl SessionContext { } } - let resolved_schema = options - .get_resolved_schema(&session_config, self.state(), table_paths[0].clone()) - .await?; + let schema_table_path = table_paths[0].clone(); let config = ListingTableConfig::new_with_multi_paths(table_paths) - .with_listing_options(listing_options) - .with_schema(resolved_schema); + .with_listing_options(listing_options); + let config = match options.schema_source() { + SchemaSource::Inferred | SchemaSource::Unset => { + config.infer_schema(&self.state()).await? + } + SchemaSource::Specified => { + let resolved_schema = options + .get_resolved_schema(&session_config, self.state(), schema_table_path) + .await?; + config.with_schema(resolved_schema) + } + }; let provider = ListingTable::try_new(config)? .with_cache(self.runtime_env().cache_manager.get_file_statistic_cache()); self.read_table(Arc::new(provider)) @@ -1942,7 +1969,7 @@ impl SessionContext { .deregister_table(&table); if let Ok(Some(ref table_provider)) = result { - self.invalidate_caches(&Some(table_ref), table_provider.table_type())?; + self.invalidate_caches(&table_ref, table_provider.table_type())?; } result @@ -2063,7 +2090,7 @@ impl FunctionRegistry for SessionContext { self.state.read().udf(name) } - fn higher_order_function(&self, name: &str) -> Result> { + fn higher_order_function(&self, name: &str) -> Result> { self.state.read().higher_order_function(name) } @@ -2081,8 +2108,8 @@ impl FunctionRegistry for SessionContext { fn register_higher_order_function( &mut self, - function: Arc, - ) -> Result>> { + function: Arc, + ) -> Result>> { self.state.write().register_higher_order_function(function) } @@ -2153,16 +2180,9 @@ impl From for SessionStateBuilder { } } +// Re-export from this module for backwards compatibility. /// A planner used to add extensions to DataFusion logical and physical plans. -#[async_trait] -pub trait QueryPlanner: Debug { - /// Given a [`LogicalPlan`], create an [`ExecutionPlan`] suitable for execution - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &SessionState, - ) -> Result>; -} +pub use datafusion_session::{QueryPlanner, UnsupportedQueryPlanner}; /// Interface for handling `CREATE FUNCTION` statements and interacting with /// [SessionState] to create and register functions ([`ScalarUDF`], @@ -2221,7 +2241,7 @@ pub enum RegisterFunction { /// Window user defined function Window(Arc), /// Higher-order user defined function - HigherOrder(Arc), + HigherOrder(Arc), /// Table user defined function Table(String, Arc), } @@ -2349,6 +2369,7 @@ mod tests { use arrow_schema::FieldRef; use datafusion_common::DataFusionError; use datafusion_common::datatype::DataTypeExt; + use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use std::error::Error; use std::path::PathBuf; @@ -2361,6 +2382,7 @@ mod tests { use crate::physical_planner::PhysicalPlanner; use async_trait::async_trait; use datafusion_expr::planner::TypePlanner; + use datafusion_session::Session; use sqlparser::ast; use tempfile::TempDir; @@ -2530,7 +2552,7 @@ mod tests { let ctx = SessionContext::new_with_state(session_state).enable_url_table(); let result = plan_and_collect( &ctx, - format!("select c_name from '{}' limit 3;", &url).as_str(), + format!("select c_name from '{url}' limit 3;").as_str(), ) .await?; @@ -2807,7 +2829,7 @@ mod tests { async fn create_physical_plan( &self, _logical_plan: &LogicalPlan, - _session_state: &SessionState, + _session_state: &dyn Session, ) -> Result> { not_impl_err!("query not supported") } @@ -2816,7 +2838,8 @@ mod tests { &self, _expr: &Expr, _input_dfschema: &DFSchema, - _session_state: &SessionState, + _session_state: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, ) -> Result> { unimplemented!() } @@ -2830,7 +2853,7 @@ mod tests { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { let physical_planner = MyPhysicalPlanner {}; physical_planner @@ -2948,7 +2971,7 @@ mod tests { // Invalid durations for duration in [ "0s", "0m", "1s0m", "1s1m", "XYZ", "1h", "XYZm2s", "", " ", "-1m", "1m 1s", - "1m1s ", " 1m1s", + "1m1s ", " 1m1s", "1\u{b5}", ] { let have = SessionContext::parse_duration(LIST_FILES_CACHE_TTL, duration); assert!(have.is_err()); @@ -3044,6 +3067,7 @@ mod tests { "G", "1024B", "invalid_size", + "1\u{b5}", ] { #[expect(deprecated)] let have = SessionContext::parse_memory_limit(limit); @@ -3079,6 +3103,7 @@ mod tests { "G", "1024B", "invalid_size", + "1\u{b5}", ] { let have = SessionContext::parse_capacity_limit(MEMORY_LIMIT, limit); assert!(have.is_err()); diff --git a/datafusion/core/src/execution/session_state.rs b/datafusion/core/src/execution/session_state.rs index de5e6b97c1af9..bfd38faacf816 100644 --- a/datafusion/core/src/execution/session_state.rs +++ b/datafusion/core/src/execution/session_state.rs @@ -30,7 +30,9 @@ use crate::datasource::provider_as_source; use crate::execution::SessionStateDefaults; use crate::execution::context::{EmptySerializerRegistry, FunctionFactory, QueryPlanner}; use crate::physical_planner::{DefaultPhysicalPlanner, PhysicalPlanner}; -use arrow_schema::{DataType, FieldRef}; +#[cfg(feature = "sql")] +use arrow_schema::DataType; +use arrow_schema::FieldRef; use datafusion_catalog::MemoryCatalogProviderList; use datafusion_catalog::information_schema::{ INFORMATION_SCHEMA, InformationSchemaProvider, @@ -53,6 +55,7 @@ use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::TableSource; use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::expr_rewriter::FunctionRewrite; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::planner::ExprPlanner; #[cfg(feature = "sql")] use datafusion_expr::planner::{RelationPlanner, TypePlanner}; @@ -70,12 +73,10 @@ use datafusion_optimizer::{ }; use datafusion_physical_expr::create_physical_expr; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; -use datafusion_physical_optimizer::PhysicalOptimizerContext; -use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::optimizer::PhysicalOptimizer; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::operator_statistics::StatisticsRegistry; -use datafusion_session::Session; +use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session}; #[cfg(feature = "sql")] use datafusion_sql::{ parser::{DFParserBuilder, Statement}, @@ -162,7 +163,7 @@ pub struct SessionState { /// Scalar functions that are registered with the context scalar_functions: HashMap>, /// Higher order functions that are registered with the context - higher_order_functions: HashMap>, + higher_order_functions: HashMap>, /// Aggregate functions registered in the context aggregate_functions: HashMap>, /// Window functions registered in the context @@ -267,6 +268,29 @@ impl Session for SessionState { self.config() } + fn catalog_list(&self) -> Arc { + Arc::clone(self.catalog_list()) + } + + fn query_planner(&self) -> Arc { + // Disambiguate: `SessionState` has an inherent `query_planner` (returning + // `&Arc<...>`) with the same name as this trait method. The qualified path + // calls the inherent one; a bare `self.query_planner()` would recurse. + Arc::clone(SessionState::query_planner(self)) + } + + fn optimize(&self, plan: &LogicalPlan) -> datafusion_common::Result { + SessionState::optimize(self, plan) + } + + fn physical_optimizers(&self) -> &[Arc] { + SessionState::physical_optimizers(self) + } + + fn statistics_registry(&self) -> Option<&StatisticsRegistry> { + SessionState::statistics_registry(self) + } + async fn create_physical_plan( &self, logical_plan: &LogicalPlan, @@ -286,7 +310,7 @@ impl Session for SessionState { &self.scalar_functions } - fn higher_order_functions(&self) -> &HashMap> { + fn higher_order_functions(&self) -> &HashMap> { &self.higher_order_functions } @@ -347,9 +371,10 @@ impl SessionState { let resolved_ref = self.resolve_table_ref(table_ref); if self.config.information_schema() && *resolved_ref.schema == *INFORMATION_SCHEMA { - return Ok(Arc::new(InformationSchemaProvider::new(Arc::clone( - &self.catalog_list, - )))); + return Ok(Arc::new( + InformationSchemaProvider::new(Arc::clone(&self.catalog_list)) + .with_table_functions(self.table_functions.clone()), + )); } self.catalog_list @@ -437,13 +462,12 @@ impl SessionState { ) -> datafusion_common::Result { let dialect = dialect_from_str(dialect).ok_or_else(|| { plan_datafusion_err!( - "Unsupported SQL dialect: {dialect}. Available dialects: \ - Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, \ - MsSQL, ClickHouse, BigQuery, Ansi, DuckDB, Databricks." + "Unsupported SQL dialect: {dialect}. Available dialects: {}.", + Dialect::available() ) })?; - let recursion_limit = self.config.options().sql_parser.recursion_limit; + let recursion_limit = self.config.options().sql_parser.recursion_limit.get(); let mut statements = DFParserBuilder::new(sql) .with_dialect(dialect.as_ref()) @@ -486,13 +510,12 @@ impl SessionState { ) -> datafusion_common::Result { let dialect = dialect_from_str(dialect).ok_or_else(|| { plan_datafusion_err!( - "Unsupported SQL dialect: {dialect}. Available dialects: \ - Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, \ - MsSQL, ClickHouse, BigQuery, Ansi, DuckDB, Databricks." + "Unsupported SQL dialect: {dialect}. Available dialects: {}.", + Dialect::available() ) })?; - let recursion_limit = self.config.options().sql_parser.recursion_limit; + let recursion_limit = self.config.options().sql_parser.recursion_limit.get(); let expr = DFParserBuilder::new(sql) .with_dialect(dialect.as_ref()) .with_recursion_limit(recursion_limit) @@ -691,6 +714,7 @@ impl SessionState { stringified_plans, schema: Arc::clone(&e.schema), logical_optimization_succeeded: false, + show_statistics: e.show_statistics, })); } Err(e) => return Err(e), @@ -728,6 +752,7 @@ impl SessionState { stringified_plans, schema: Arc::clone(&e.schema), logical_optimization_succeeded, + show_statistics: e.show_statistics, })) } else { let analyzed_plan = self.analyzer.execute_and_check( @@ -796,7 +821,12 @@ impl SessionState { .transform_up(|expr| rewrite.rewrite(expr, df_schema, config_options))? .data; } - create_physical_expr(&expr, df_schema, self.execution_props()) + create_physical_expr( + &expr, + df_schema, + self.execution_props(), + &PhysicalPlanningContext::default(), + ) } /// Return the session ID @@ -934,7 +964,7 @@ impl SessionState { } /// Return reference to higher_order_functions - pub fn higher_order_functions(&self) -> &HashMap> { + pub fn higher_order_functions(&self) -> &HashMap> { &self.higher_order_functions } @@ -1034,7 +1064,7 @@ pub struct SessionStateBuilder { catalog_list: Option>, table_functions: Option>>, scalar_functions: Option>>, - higher_order_functions: Option>>, + higher_order_functions: Option>>, aggregate_functions: Option>>, window_functions: Option>>, extension_types: Option, @@ -1371,7 +1401,7 @@ impl SessionStateBuilder { /// Set the map of [`HigherOrderUDF`]s pub fn with_higher_order_functions( mut self, - higher_order_functions: Vec>, + higher_order_functions: Vec>, ) -> Self { self.higher_order_functions = Some(higher_order_functions); self @@ -1791,9 +1821,7 @@ impl SessionStateBuilder { } /// Returns the current scalar_functions value - pub fn higher_order_functions( - &mut self, - ) -> &mut Option>> { + pub fn higher_order_functions(&mut self) -> &mut Option>> { &mut self.higher_order_functions } @@ -2016,7 +2044,7 @@ impl ContextProvider for SessionContextProvider<'_> { self.state.scalar_functions().get(name).cloned() } - fn get_higher_order_meta(&self, name: &str) -> Option> { + fn get_higher_order_meta(&self, name: &str) -> Option> { self.state.higher_order_functions().get(name).cloned() } @@ -2106,7 +2134,7 @@ impl FunctionRegistry for SessionState { fn higher_order_function( &self, name: &str, - ) -> datafusion_common::Result> { + ) -> datafusion_common::Result> { self.higher_order_functions .get(name) .cloned() @@ -2142,8 +2170,8 @@ impl FunctionRegistry for SessionState { fn register_higher_order_function( &mut self, - function: Arc, - ) -> datafusion_common::Result>> { + function: Arc, + ) -> datafusion_common::Result>> { function.aliases().iter().for_each(|alias| { self.higher_order_functions .insert(alias.clone(), Arc::clone(&function)); @@ -2191,7 +2219,7 @@ impl FunctionRegistry for SessionState { fn deregister_higher_order_function( &mut self, name: &str, - ) -> datafusion_common::Result>> { + ) -> datafusion_common::Result>> { let function = self.higher_order_functions.remove(name); if let Some(function) = &function { for alias in function.aliases() { @@ -2311,7 +2339,7 @@ impl QueryPlanner for DefaultQueryPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> datafusion_common::Result> { let planner = DefaultPhysicalPlanner::default(); planner @@ -2354,23 +2382,34 @@ mod tests { use crate::logical_expr::{AggregateUDF, ScalarUDF, TableSource, WindowUDF}; use crate::physical_plan::ExecutionPlan; use crate::sql::planner::ContextProvider; - use crate::sql::{ResolvedTableReference, TableReference}; use arrow::array::{ArrayRef, Int32Array, RecordBatch, StringArray}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_catalog::MemoryCatalogProviderList; - use datafusion_common::DFSchema; - use datafusion_common::Result; use datafusion_common::config::Dialect; + use datafusion_common::{DFSchema, ResolvedTableReference, Result, TableReference}; use datafusion_execution::config::SessionConfig; use datafusion_expr::Expr; use datafusion_expr::HigherOrderUDF; use datafusion_optimizer::Optimizer; use datafusion_optimizer::optimizer::OptimizerRule; use datafusion_physical_plan::display::DisplayableExecutionPlan; + use datafusion_session::Session; use datafusion_sql::planner::{PlannerContext, SqlToRel}; use std::collections::HashMap; use std::sync::Arc; + #[test] + #[cfg(feature = "sql")] + fn test_configured_dialect_names_are_accepted_by_sqlparser() { + for info in Dialect::metadata() { + assert!( + sqlparser::dialect::dialect_from_str(info.canonical_name).is_some(), + "sqlparser should accept configured dialect {}", + info.canonical_name + ); + } + } + #[test] #[cfg(feature = "sql")] fn test_session_state_with_default_features() { @@ -2446,6 +2485,9 @@ mod tests { let session_state = SessionStateBuilder::new() .with_catalog_list(Arc::new(MemoryCatalogProviderList::new())) .build(); + let session_catalogs = Session::catalog_list(&session_state); + assert!(Arc::ptr_eq(&session_catalogs, session_state.catalog_list())); + let table_ref = session_state.resolve_table_ref("employee").to_string(); session_state .schema_for_ref(&table_ref)? @@ -2677,7 +2719,7 @@ mod tests { self.state.scalar_functions().get(name).cloned() } - fn get_higher_order_meta(&self, name: &str) -> Option> { + fn get_higher_order_meta(&self, name: &str) -> Option> { self.state.higher_order_functions().get(name).cloned() } diff --git a/datafusion/core/src/execution/session_state_defaults.rs b/datafusion/core/src/execution/session_state_defaults.rs index 5e85c1bbc5e9e..584879cb197b5 100644 --- a/datafusion/core/src/execution/session_state_defaults.rs +++ b/datafusion/core/src/execution/session_state_defaults.rs @@ -114,7 +114,7 @@ impl SessionStateDefaults { } /// returns the list of default [`HigherOrderUDF`]s - pub fn default_higher_order_functions() -> Vec> { + pub fn default_higher_order_functions() -> Vec> { #[cfg(feature = "nested_expressions")] return functions_nested::all_default_higher_order_functions(); diff --git a/datafusion/core/src/optimizer_rule_reference.md b/datafusion/core/src/optimizer_rule_reference.md index 1f9f37f530557..1367ed0843c59 100644 --- a/datafusion/core/src/optimizer_rule_reference.md +++ b/datafusion/core/src/optimizer_rule_reference.md @@ -75,20 +75,19 @@ in multiple phases. | 3 | `join_selection` | - | Chooses join implementation, build side, and partition mode from statistics and stream properties. | | 4 | `LimitedDistinctAggregation` | - | Pushes limit hints into grouped distinct-style aggregations when only a small result is needed. | | 5 | `FilterPushdown` | pre-optimization phase | Pushes supported physical filters down toward data sources before distribution and sorting are enforced. | -| 6 | `EnforceDistribution` | - | Adds repartitioning only where needed to satisfy physical distribution requirements. | -| 7 | `CombinePartialFinalAggregate` | - | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant. | -| 8 | `EnforceSorting` | - | Adds or removes local sorts to satisfy required input orderings. | +| 6 | `WindowTopN` | - | Replaces eligible row-number window and filter patterns with per-partition TopK execution. | +| 7 | `EnsureRequirements` | - | Enforces both distribution and sorting requirements in a single idempotent rule. | +| 8 | `CombinePartialFinalAggregate` | - | Collapses adjacent partial and final aggregates when the distributed shape makes them redundant. | | 9 | `OptimizeAggregateOrder` | - | Updates aggregate expressions to use the best ordering once sort requirements are known. | -| 10 | `WindowTopN` | - | Replaces eligible row-number window and filter patterns with per-partition TopK execution. | -| 11 | `ProjectionPushdown` | early pass | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure. | -| 12 | `OutputRequirements` | remove phase | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done. | -| 13 | `LimitAggregation` | - | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets. | -| 14 | `LimitPushPastWindows` | - | Pushes fetch limits through bounded window operators when doing so keeps the result correct. | -| 15 | `HashJoinBuffering` | - | Adds buffering on the probe side of hash joins so probing can start before build completion. | -| 16 | `LimitPushdown` | - | Moves physical limits into child operators or fetch-enabled variants to cut data early. | -| 17 | `TopKRepartition` | - | Pushes TopK below hash repartition when the partition key is a prefix of the sort key. | -| 18 | `ProjectionPushdown` | late pass | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities. | -| 19 | `PushdownSort` | - | Pushes sort requirements into data sources that can already return sorted output. | -| 20 | `EnsureCooperative` | - | Wraps non-cooperative plan parts so long-running tasks yield fairly. | -| 21 | `FilterPushdown(Post)` | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving. | -| 22 | `SanityCheckPlan` | - | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. | +| 10 | `ProjectionPushdown` | early pass | Pushes projections toward inputs before later physical rewrites add more limit and TopK structure. | +| 11 | `OutputRequirements` | remove phase | Removes the temporary output-requirement helper nodes after requirement-sensitive planning is done. | +| 12 | `LimitAggregation` | - | Passes a limit hint into eligible aggregations so they can keep fewer accumulator buckets. | +| 13 | `LimitPushPastWindows` | - | Pushes fetch limits through bounded window operators when doing so keeps the result correct. | +| 14 | `HashJoinBuffering` | - | Adds buffering on the probe side of hash joins so probing can start before build completion. | +| 15 | `LimitPushdown` | - | Moves physical limits into child operators or fetch-enabled variants to cut data early. | +| 16 | `TopKRepartition` | - | Pushes TopK below hash repartition when the partition key is a prefix of the sort key. | +| 17 | `ProjectionPushdown` | late pass | Runs projection pushdown again after limit and TopK rewrites expose new pruning opportunities. | +| 18 | `PushdownSort` | - | Pushes sort requirements into data sources that can already return sorted output. | +| 19 | `EnsureCooperative` | - | Wraps non-cooperative plan parts so long-running tasks yield fairly. | +| 20 | `FilterPushdown(Post)` | post-optimization phase | Pushes dynamic filters at the end of optimization, after plan references stop moving. | +| 21 | `SanityCheckPlan` | - | Validates that the final physical plan meets ordering, distribution, and infinite-input safety requirements. | diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs index d225cff1deafc..3c1e7b50780a5 100644 --- a/datafusion/core/src/physical_planner.rs +++ b/datafusion/core/src/physical_planner.rs @@ -26,16 +26,15 @@ use crate::datasource::listing::ListingTableUrl; use crate::datasource::physical_plan::{FileOutputMode, FileSinkConfig}; use crate::datasource::{DefaultTableSource, source_as_provider}; use crate::error::{DataFusionError, Result}; -use crate::execution::context::{ExecutionProps, SessionState}; +use crate::execution::context::ExecutionProps; use crate::logical_expr::utils::generate_sort_key; use crate::logical_expr::{ Aggregate, EmptyRelation, Join, Projection, Sort, TableScan, Unnest, Values, Window, }; -use crate::logical_expr::{ - Expr, LogicalPlan, Partitioning as LogicalPartitioning, PlanType, Repartition, - UserDefinedLogicalNode, +use crate::logical_expr::{Expr, LogicalPlan, PlanType, Repartition}; +use crate::physical_expr::{ + create_physical_expr, create_physical_exprs, create_physical_partitioning, }; -use crate::physical_expr::{create_physical_expr, create_physical_exprs}; use crate::physical_plan::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use crate::physical_plan::analyze::AnalyzeExec; use crate::physical_plan::explain::ExplainExec; @@ -52,8 +51,8 @@ use crate::physical_plan::union::UnionExec; use crate::physical_plan::unnest::UnnestExec; use crate::physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; use crate::physical_plan::{ - ExecutionPlan, ExecutionPlanProperties, InputOrderMode, Partitioning, PhysicalExpr, - WindowExpr, displayable, windows, + ExecutionPlan, ExecutionPlanProperties, InputOrderMode, PhysicalExpr, WindowExpr, + displayable, windows, }; use crate::schema_equivalence::schema_satisfied_by; @@ -79,7 +78,6 @@ use datafusion_common::{ use datafusion_datasource::file_groups::FileGroup; use datafusion_datasource::memory::MemorySourceConfig; use datafusion_expr::dml::{CopyTo, InsertOp}; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_expr::expr::{ Alias, GroupingSet, NullTreatment, WindowFunction, WindowFunctionParams, physical_name, @@ -87,6 +85,9 @@ use datafusion_expr::expr::{ use datafusion_expr::expr_rewriter::unnormalize_cols; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::logical_plan::builder::wrap_projection_for_join_if_necessary; +use datafusion_expr::physical_planning_context::{ + PhysicalPlanningContext, ScalarSubqueryResults, SubqueryIndex, +}; use datafusion_expr::utils::{expr_to_columns, split_conjunction}; use datafusion_expr::{ Analyze, BinaryExpr, DescribeTable, DmlStatement, Explain, ExplainFormat, Extension, @@ -100,7 +101,6 @@ use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::{ LexOrdering, PhysicalSortExpr, create_physical_sort_exprs, }; -use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::execution_plan::InvariantLevel; use datafusion_physical_plan::joins::PiecewiseMergeJoinExec; @@ -108,6 +108,7 @@ use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::recursive_query::RecursiveQueryExec; use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; use datafusion_physical_plan::unnest::ListUnnest; +use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule, Session}; use async_trait::async_trait; use datafusion_physical_plan::async_func::{AsyncFuncExec, AsyncMapper}; @@ -117,144 +118,31 @@ use itertools::{Itertools, multiunzip}; use log::debug; use tokio::sync::Mutex; -/// Physical query planner that converts a `LogicalPlan` to an -/// `ExecutionPlan` suitable for execution. -#[async_trait] -pub trait PhysicalPlanner: Send + Sync { - /// Create a physical plan from a logical plan - async fn create_physical_plan( - &self, - logical_plan: &LogicalPlan, - session_state: &SessionState, - ) -> Result>; +// Re-export from this module for backwards compatibility. +pub use datafusion_session::{ExtensionPlanner, PhysicalPlanner}; - /// Create a physical expression from a logical expression - /// suitable for evaluation - /// - /// `expr`: the expression to convert - /// - /// `input_dfschema`: the logical plan schema for evaluating `expr` - fn create_physical_expr( - &self, - expr: &Expr, - input_dfschema: &DFSchema, - session_state: &SessionState, - ) -> Result>; +struct SessionOptimizerContext<'a> { + session: &'a dyn Session, } -/// This trait exposes the ability to plan an [`ExecutionPlan`] out of a [`LogicalPlan`]. -#[async_trait] -pub trait ExtensionPlanner { - /// Create a physical plan for a [`UserDefinedLogicalNode`]. - /// - /// `input_dfschema`: the logical plan schema for the inputs to this node - /// - /// Returns an error when the planner knows how to plan the concrete - /// implementation of `node` but errors while doing so. - /// - /// Returns `None` when the planner does not know how to plan the - /// `node` and wants to delegate the planning to another - /// [`ExtensionPlanner`]. - async fn plan_extension( - &self, - planner: &dyn PhysicalPlanner, - node: &dyn UserDefinedLogicalNode, - logical_inputs: &[&LogicalPlan], - physical_inputs: &[Arc], - session_state: &SessionState, - ) -> Result>>; - - /// Create a physical plan for a [`LogicalPlan::TableScan`]. - /// - /// This is useful for planning valid [`TableSource`]s that are not [`TableProvider`]s. - /// - /// Returns: - /// * `Ok(Some(plan))` if the planner knows how to plan the `scan` - /// * `Ok(None)` if the planner does not know how to plan the `scan` and wants to delegate the planning to another [`ExtensionPlanner`] - /// * `Err` if the planner knows how to plan the `scan` but errors while doing so - /// - /// # Example - /// - /// ```rust,ignore - /// use std::sync::Arc; - /// use datafusion::physical_plan::ExecutionPlan; - /// use datafusion::logical_expr::TableScan; - /// use datafusion::execution::context::SessionState; - /// use datafusion::error::Result; - /// use datafusion_physical_planner::{ExtensionPlanner, PhysicalPlanner}; - /// use async_trait::async_trait; - /// - /// // Your custom table source type - /// struct MyCustomTableSource { /* ... */ } - /// - /// // Your custom execution plan - /// struct MyCustomExec { /* ... */ } - /// - /// struct MyExtensionPlanner; - /// - /// #[async_trait] - /// impl ExtensionPlanner for MyExtensionPlanner { - /// async fn plan_extension( - /// &self, - /// _planner: &dyn PhysicalPlanner, - /// _node: &dyn UserDefinedLogicalNode, - /// _logical_inputs: &[&LogicalPlan], - /// _physical_inputs: &[Arc], - /// _session_state: &SessionState, - /// ) -> Result>> { - /// Ok(None) - /// } - /// - /// async fn plan_table_scan( - /// &self, - /// _planner: &dyn PhysicalPlanner, - /// scan: &TableScan, - /// _session_state: &SessionState, - /// ) -> Result>> { - /// // Check if this is your custom table source - /// if scan.source.is::() { - /// // Create a custom execution plan for your table source - /// let exec = MyCustomExec::new( - /// scan.table_name.clone(), - /// Arc::clone(scan.projected_schema.inner()), - /// ); - /// Ok(Some(Arc::new(exec))) - /// } else { - /// // Return None to let other extension planners handle it - /// Ok(None) - /// } - /// } - /// } - /// ``` - /// - /// [`TableSource`]: datafusion_expr::TableSource - /// [`TableProvider`]: datafusion_catalog::TableProvider - async fn plan_table_scan( +impl PhysicalOptimizerContext for SessionOptimizerContext<'_> { + fn config_options(&self) -> &datafusion_common::config::ConfigOptions { + self.session.config_options() + } + + fn statistics_registry( &self, - _planner: &dyn PhysicalPlanner, - _scan: &TableScan, - _session_state: &SessionState, - ) -> Result>> { - Ok(None) + ) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry> { + self.session.statistics_registry() } } /// Default single node physical query planner that converts a /// `LogicalPlan` to an `ExecutionPlan` suitable for execution. /// -/// This planner will first flatten the `LogicalPlan` tree via a -/// depth first approach, which allows it to identify the leaves -/// of the tree. -/// -/// Tasks are spawned from these leaves and traverse back up the -/// tree towards the root, converting each `LogicalPlan` node it -/// reaches into their equivalent `ExecutionPlan` node. When these -/// tasks reach a common node, they will terminate until the last -/// task reaches the node which will then continue building up the -/// tree. -/// -/// Up to [`planning_concurrency`] tasks are buffered at once to -/// execute concurrently. +/// This planner first flattens the `LogicalPlan` tree with a depth-first +/// traversal. It then builds the physical plan from the leaves to the root. +/// Up to [`planning_concurrency`] tasks execute concurrently. /// /// [`planning_concurrency`]: crate::config::ExecutionOptions::planning_concurrency #[derive(Default)] @@ -268,7 +156,7 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { if let Some(plan) = self .handle_explain_or_analyze(logical_plan, session_state) @@ -293,9 +181,15 @@ impl PhysicalPlanner for DefaultPhysicalPlanner { &self, expr: &Expr, input_dfschema: &DFSchema, - session_state: &SessionState, + session_state: &dyn Session, + planning_ctx: &PhysicalPlanningContext, ) -> Result> { - create_physical_expr(expr, input_dfschema, session_state.execution_props()) + create_physical_expr( + expr, + input_dfschema, + session_state.execution_props(), + planning_ctx, + ) } } @@ -372,13 +266,13 @@ impl DefaultPhysicalPlanner { &self, logical_schema: &DFSchemaRef, physical_plan: &Arc, - context: &str, + context_factory: impl FnOnce() -> String, ) -> Result<()> { if !logical_schema.matches_arrow_schema(&physical_plan.schema()) { return plan_err!( "{} created an ExecutionPlan with mismatched schema. \ LogicalPlan schema: {:?}, ExecutionPlan schema: {:?}", - context, + context_factory(), logical_schema, physical_plan.schema() ); @@ -416,9 +310,9 @@ impl DefaultPhysicalPlanner { /// collected, planned as separate physical plans, and each assigned an /// index in a shared [`ScalarSubqueryResults`] container that will hold its /// result at execution time. The index map and shared results container are - /// registered in [`ExecutionProps`] so that [`create_physical_expr`] can - /// convert `Expr::ScalarSubquery` into [`ScalarSubqueryExpr`] nodes that - /// read from that container. + /// stored in a [`PhysicalPlanningContext`] and passed explicitly to + /// [`create_physical_expr`] so it can convert `Expr::ScalarSubquery` into + /// [`ScalarSubqueryExpr`] nodes that read from that container. /// /// The resulting physical plan is wrapped in a [`ScalarSubqueryExec`] node /// that executes those subquery plans before any data flows through the @@ -434,37 +328,48 @@ impl DefaultPhysicalPlanner { fn create_initial_plan<'a>( &'a self, logical_plan: &'a LogicalPlan, - session_state: &'a SessionState, + session_state: &'a dyn Session, ) -> futures::future::BoxFuture<'a, Result>> { Box::pin(async move { - let all_subqueries = Self::collect_scalar_subqueries(logical_plan); + // When `enable_physical_uncorrelated_scalar_subquery` is disabled, the + // `ScalarSubqueryToJoin` optimizer rule rewrites all uncorrelated + // scalar subqueries to joins, so none should reach this point. + // Skip collection in that case to avoid creating a no-op + // `ScalarSubqueryExec` wrapper. + let all_subqueries = if session_state + .config_options() + .optimizer + .enable_physical_uncorrelated_scalar_subquery + { + Self::collect_scalar_subqueries(logical_plan) + } else { + Vec::new() + }; let (links, index_map) = self .plan_scalar_subqueries(all_subqueries, session_state) .await?; if links.is_empty() { return self - .create_initial_plan_inner(logical_plan, session_state) + .create_initial_plan_inner( + logical_plan, + session_state, + &PhysicalPlanningContext::default(), + ) .await; } - // Create the shared `ScalarSubqueryResults` container and register - // it in `ExecutionProps` so that `create_physical_expr` can resolve - // `Expr::ScalarSubquery` into `ScalarSubqueryExpr` nodes. We clone - // the `SessionState` so these are available throughout physical - // planning without mutating the caller's state. - // - // Ideally, the subquery state would live in a dedicated planning - // context rather than in `ExecutionProps`. It's here because - // `create_physical_expr` only receives `&ExecutionProps`. + // Build a `PhysicalPlanningContext` that carries the index map and + // shared results container into calls that create physical expressions. + // The context is threaded explicitly through physical planning rather + // than being stashed in `ExecutionProps`, so the planner does not need + // a mutable `SessionState` and each recursively planned subtree receives + // the correct context. let results = ScalarSubqueryResults::new(links.len()); - let mut owned = session_state.clone(); - owned.execution_props_mut().subquery_indexes = index_map; - owned.execution_props_mut().subquery_results = results.clone(); - let session_state = Cow::Owned(owned); + let planning_ctx = PhysicalPlanningContext::new(index_map, results.clone()); let plan = self - .create_initial_plan_inner(logical_plan, &session_state) + .create_initial_plan_inner(logical_plan, session_state, &planning_ctx) .await?; Ok(Arc::new(ScalarSubqueryExec::new(plan, links, results))) }) @@ -475,7 +380,8 @@ impl DefaultPhysicalPlanner { async fn create_initial_plan_inner( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, + planning_ctx: &PhysicalPlanningContext, ) -> Result> { // DFS the tree to flatten it into a Vec. // This will allow us to build the Physical Plan from the leaves up @@ -526,9 +432,9 @@ impl DefaultPhysicalPlanner { let max_concurrency = planning_concurrency.min(flat_tree_leaf_indices.len()); // Spawning tasks which will traverse leaf up to the root. - let tasks = flat_tree_leaf_indices - .into_iter() - .map(|index| self.task_helper(index, Arc::clone(&flat_tree), session_state)); + let tasks = flat_tree_leaf_indices.into_iter().map(|index| { + self.task_helper(index, Arc::clone(&flat_tree), session_state, planning_ctx) + }); let mut outputs = futures::stream::iter(tasks) .buffer_unordered(max_concurrency) .try_collect::>() @@ -555,7 +461,8 @@ impl DefaultPhysicalPlanner { &'a self, leaf_starter_index: usize, flat_tree: Arc>>, - session_state: &'a SessionState, + session_state: &'a dyn Session, + planning_ctx: &'a PhysicalPlanningContext, ) -> Result>> { // We always start with a leaf, so can ignore status and pass empty children let mut node = flat_tree.get(leaf_starter_index).ok_or_else(|| { @@ -567,6 +474,7 @@ impl DefaultPhysicalPlanner { .map_logical_node_to_physical( node.node, session_state, + planning_ctx, ChildrenContainer::None, ) .await?; @@ -584,6 +492,7 @@ impl DefaultPhysicalPlanner { .map_logical_node_to_physical( node.node, session_state, + planning_ctx, ChildrenContainer::One(plan), ) .await?; @@ -620,7 +529,12 @@ impl DefaultPhysicalPlanner { let children = children.into_iter().map(|epc| epc.plan).collect(); let children = ChildrenContainer::Multiple(children); plan = self - .map_logical_node_to_physical(node.node, session_state, children) + .map_logical_node_to_physical( + node.node, + session_state, + planning_ctx, + children, + ) .await?; } } @@ -634,7 +548,8 @@ impl DefaultPhysicalPlanner { async fn map_logical_node_to_physical( &self, node: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, + planning_ctx: &PhysicalPlanningContext, children: ChildrenContainer, ) -> Result> { let execution_props = session_state.execution_props(); @@ -647,6 +562,7 @@ impl DefaultPhysicalPlanner { filters, fetch, projected_schema, + statistics_requests, .. } = scan; @@ -656,10 +572,13 @@ impl DefaultPhysicalPlanner { // referred to in the query let filters = unnormalize_cols(filters.iter().cloned()); let filters_vec = filters.into_iter().collect::>(); + let stats_requests = + statistics_requests.iter().cloned().collect::>(); let opts = ScanArgs::default() .with_projection(projection.as_deref()) .with_filters(Some(&filters_vec)) - .with_limit(*fetch); + .with_limit(*fetch) + .with_statistics_requests(&stats_requests); let res = source.scan_with_args(session_state, opts).await?; Arc::clone(res.plan()) } else { @@ -669,8 +588,9 @@ impl DefaultPhysicalPlanner { break; } - maybe_plan = - planner.plan_table_scan(self, scan, session_state).await?; + maybe_plan = planner + .plan_table_scan(self, scan, session_state, planning_ctx) + .await?; } let plan = match maybe_plan { @@ -682,9 +602,11 @@ impl DefaultPhysicalPlanner { ); } }; - let context = - format!("Extension planner for table scan {}", scan.table_name); - self.ensure_schema_matches(projected_schema, &plan, &context)?; + + self.ensure_schema_matches(projected_schema, &plan, || { + format!("Extension planner for table scan {}", scan.table_name) + })?; + plan } } @@ -694,7 +616,12 @@ impl DefaultPhysicalPlanner { .map(|row| { row.iter() .map(|expr| { - create_physical_expr(expr, schema, execution_props) + create_physical_expr( + expr, + schema, + execution_props, + planning_ctx, + ) }) .collect::>>>() }) @@ -907,6 +834,35 @@ impl DefaultPhysicalPlanner { ); } } + LogicalPlan::Dml(DmlStatement { + table_name, + target, + op: WriteOp::MergeInto(merge_op), + input, + .. + }) => { + let provider = source_as_provider(target).map_err(|e| { + e.context(format!("MERGE INTO operation on table '{table_name}'")) + })?; + let input_exec = children.one()?; + let target_schema = DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?; + let merge_schema = Arc::new(target_schema.join(input.schema())?); + provider + .merge_into( + session_state, + input_exec, + merge_schema, + merge_op.on.clone(), + merge_op.clauses.clone(), + ) + .await + .map_err(|e| { + e.context(format!("MERGE INTO operation on table '{table_name}'")) + })? + } LogicalPlan::Window(Window { window_expr, .. }) => { assert_or_internal_err!( !window_expr.is_empty(), @@ -953,7 +909,14 @@ impl DefaultPhysicalPlanner { let logical_schema = node.schema(); let window_expr = window_expr .iter() - .map(|e| create_window_expr(e, logical_schema, execution_props)) + .map(|e| { + create_window_expr( + e, + logical_schema, + execution_props, + planning_ctx, + ) + }) .collect::>>()?; let can_repartition = session_state.config().target_partitions() > 1 @@ -1059,6 +1022,7 @@ impl DefaultPhysicalPlanner { logical_input_schema, &physical_input_schema, execution_props, + planning_ctx, )?; let agg_filter = aggr_expr @@ -1069,6 +1033,7 @@ impl DefaultPhysicalPlanner { logical_input_schema, &physical_input_schema, execution_props, + planning_ctx, ) .build() .map(lowered_aggregate_to_tuple) @@ -1166,17 +1131,23 @@ impl DefaultPhysicalPlanner { LogicalPlan::Projection(Projection { input, expr, .. }) => self .create_project_physical_exec_with_props( execution_props, + planning_ctx, children.one()?, input, expr, + node.schema(), )?, LogicalPlan::Filter(Filter { predicate, input, .. }) => { let physical_input = children.one()?; let input_dfschema = input.schema(); - let runtime_expr = - create_physical_expr(predicate, input_dfschema, execution_props)?; + let runtime_expr = create_physical_expr( + predicate, + input_dfschema, + execution_props, + planning_ctx, + )?; let input_schema = input.schema(); let filter = match self.try_plan_async_exprs( @@ -1234,25 +1205,12 @@ impl DefaultPhysicalPlanner { }) => { let physical_input = children.one()?; let input_dfschema = input.as_ref().schema(); - let physical_partitioning = match partitioning_scheme { - LogicalPartitioning::RoundRobinBatch(n) => { - Partitioning::RoundRobinBatch(*n) - } - LogicalPartitioning::Hash(expr, n) => { - let runtime_expr = expr - .iter() - .map(|e| { - create_physical_expr(e, input_dfschema, execution_props) - }) - .collect::>>()?; - Partitioning::Hash(runtime_expr, *n) - } - LogicalPartitioning::DistributeBy(_) => { - return not_impl_err!( - "Physical plan does not support DistributeBy partitioning" - ); - } - }; + let physical_partitioning = create_physical_partitioning( + partitioning_scheme, + input_dfschema, + execution_props, + planning_ctx, + )?; Arc::new(RepartitionExec::try_new( physical_input, physical_partitioning, @@ -1263,8 +1221,12 @@ impl DefaultPhysicalPlanner { }) => { let physical_input = children.one()?; let input_dfschema = input.as_ref().schema(); - let sort_exprs = - create_physical_sort_exprs(expr, input_dfschema, execution_props)?; + let sort_exprs = create_physical_sort_exprs( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; let Some(ordering) = LexOrdering::new(sort_exprs) else { return internal_err!( "SortExec requires at least one sort expression" @@ -1393,9 +1355,11 @@ impl DefaultPhysicalPlanner { LogicalPlan::Projection(Projection { input, expr, .. }), ) => self.create_project_physical_exec_with_props( execution_props, + planning_ctx, physical_left, input, expr, + left.schema(), )?, _ => physical_left, }; @@ -1406,9 +1370,11 @@ impl DefaultPhysicalPlanner { LogicalPlan::Projection(Projection { input, expr, .. }), ) => self.create_project_physical_exec_with_props( execution_props, + planning_ctx, physical_right, input, expr, + right.schema(), )?, _ => physical_right, }; @@ -1473,9 +1439,18 @@ impl DefaultPhysicalPlanner { let join_on = keys .iter() .map(|(l, r)| { - let l = create_physical_expr(l, left_df_schema, execution_props)?; - let r = - create_physical_expr(r, right_df_schema, execution_props)?; + let l = create_physical_expr( + l, + left_df_schema, + execution_props, + planning_ctx, + )?; + let r = create_physical_expr( + r, + right_df_schema, + execution_props, + planning_ctx, + )?; Ok((l, r)) }) .collect::>()?; @@ -1576,6 +1551,7 @@ impl DefaultPhysicalPlanner { expr, &filter_df_schema, execution_props, + planning_ctx, )?; let column_indices = join_utils::JoinFilter::build_column_indices( left_field_indices, @@ -1696,11 +1672,13 @@ impl DefaultPhysicalPlanner { lhs_logical, left_df_schema, execution_props, + planning_ctx, )?; let on_right = create_physical_expr( rhs_logical, right_df_schema, execution_props, + planning_ctx, )?; Arc::new(PiecewiseMergeJoinExec::try_new( @@ -1724,6 +1702,11 @@ impl DefaultPhysicalPlanner { } else if session_state.config().target_partitions() > 1 && session_state.config().repartition_joins() && !prefer_hash_join + && !*null_aware + // Null-aware joins (e.g. `NOT IN` with a nullable subquery) must + // use the CollectLeft HashJoin below: SortMergeJoinExec does not + // implement null-aware anti-join semantics and would return wrong + // results when the right side contains a null join key. { // Use SortMergeJoin if hash join is not preferred let join_on_len = join_on.len(); @@ -1772,20 +1755,26 @@ impl DefaultPhysicalPlanner { if let Some((input, expr)) = new_project { self.create_project_physical_exec_with_props( execution_props, + planning_ctx, join, input, expr, + new_logical.schema(), )? } else { join } } LogicalPlan::RecursiveQuery(RecursiveQuery { - name, is_distinct, .. + name, + is_distinct, + schema, + .. }) => { let [static_term, recursive_term] = children.two()?; Arc::new(RecursiveQueryExec::try_new( name.clone(), + Arc::clone(schema.inner()), static_term, recursive_term, *is_distinct, @@ -1810,6 +1799,7 @@ impl DefaultPhysicalPlanner { &logical_input, &children, session_state, + planning_ctx, ) .await?; } @@ -1822,8 +1812,10 @@ impl DefaultPhysicalPlanner { ), }?; - let context = format!("Extension planner for {node:?}"); - self.ensure_schema_matches(node.schema(), &plan, &context)?; + self.ensure_schema_matches(node.schema(), &plan, || { + format!("Extension planner for {node:?}") + })?; + plan } @@ -1870,6 +1862,7 @@ impl DefaultPhysicalPlanner { input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result { if group_expr.len() == 1 { match &group_expr[0] { @@ -1879,6 +1872,7 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + planning_ctx, ) } Expr::GroupingSet(GroupingSet::Cube(exprs)) => create_cube_physical_expr( @@ -1886,6 +1880,7 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + planning_ctx, ), Expr::GroupingSet(GroupingSet::Rollup(exprs)) => { create_rollup_physical_expr( @@ -1893,10 +1888,16 @@ impl DefaultPhysicalPlanner { input_dfschema, input_schema, execution_props, + planning_ctx, ) } expr => Ok(PhysicalGroupBy::new_single(vec![tuple_err(( - create_physical_expr(expr, input_dfschema, execution_props), + create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + ), physical_name(expr), ))?])), } @@ -1910,7 +1911,12 @@ impl DefaultPhysicalPlanner { .iter() .map(|e| { tuple_err(( - create_physical_expr(e, input_dfschema, execution_props), + create_physical_expr( + e, + input_dfschema, + execution_props, + planning_ctx, + ), physical_name(e), )) }) @@ -1935,6 +1941,7 @@ fn merge_grouping_set_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result { let num_groups = grouping_sets.len(); let mut all_exprs: Vec = vec![]; @@ -1949,6 +1956,7 @@ fn merge_grouping_set_physical_expr( expr, input_dfschema, execution_props, + planning_ctx, )?); null_exprs.push(get_null_physical_expr_pair( @@ -1956,6 +1964,7 @@ fn merge_grouping_set_physical_expr( input_dfschema, input_schema, execution_props, + planning_ctx, )?); } } @@ -1986,6 +1995,7 @@ fn create_cube_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result { let num_of_exprs = exprs.len(); let num_groups = num_of_exprs * num_of_exprs; @@ -2001,12 +2011,14 @@ fn create_cube_physical_expr( input_dfschema, input_schema, execution_props, + planning_ctx, )?); all_exprs.push(get_physical_expr_pair( expr, input_dfschema, execution_props, + planning_ctx, )?) } @@ -2032,6 +2044,7 @@ fn create_rollup_physical_expr( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result { let num_of_exprs = exprs.len(); @@ -2048,12 +2061,14 @@ fn create_rollup_physical_expr( input_dfschema, input_schema, execution_props, + planning_ctx, )?); all_exprs.push(get_physical_expr_pair( expr, input_dfschema, execution_props, + planning_ctx, )?) } @@ -2080,8 +2095,10 @@ fn get_null_physical_expr_pair( input_dfschema: &DFSchema, input_schema: &Schema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<(Arc, String)> { - let physical_expr = create_physical_expr(expr, input_dfschema, execution_props)?; + let physical_expr = + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?; let physical_name = physical_name(&expr.clone())?; let data_type = physical_expr.data_type(input_schema)?; @@ -2151,8 +2168,10 @@ fn get_physical_expr_pair( expr: &Expr, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result<(Arc, String)> { - let physical_expr = create_physical_expr(expr, input_dfschema, execution_props)?; + let physical_expr = + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?; let physical_name = physical_name(expr)?; Ok((physical_expr, physical_name)) } @@ -2405,11 +2424,14 @@ pub fn is_window_frame_bound_valid(window_frame: &WindowFrame) -> bool { } /// Create a window expression with a name from a logical expression +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_window_expr_with_name( e: &Expr, name: impl Into, logical_schema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result> { let name = name.into(); let physical_schema = Arc::clone(logical_schema.inner()); @@ -2428,12 +2450,24 @@ pub fn create_window_expr_with_name( filter, }, } = window_fun.as_ref(); - let physical_args = - create_physical_exprs(args, logical_schema, execution_props)?; - let partition_by = - create_physical_exprs(partition_by, logical_schema, execution_props)?; - let order_by = - create_physical_sort_exprs(order_by, logical_schema, execution_props)?; + let physical_args = create_physical_exprs( + args, + logical_schema, + execution_props, + planning_ctx, + )?; + let partition_by = create_physical_exprs( + partition_by, + logical_schema, + execution_props, + planning_ctx, + )?; + let order_by = create_physical_sort_exprs( + order_by, + logical_schema, + execution_props, + planning_ctx, + )?; if !is_window_frame_bound_valid(window_frame) { return plan_err!( @@ -2448,7 +2482,9 @@ pub fn create_window_expr_with_name( == NullTreatment::IgnoreNulls; let physical_filter = filter .as_ref() - .map(|f| create_physical_expr(f, logical_schema, execution_props)) + .map(|f| { + create_physical_expr(f, logical_schema, execution_props, planning_ctx) + }) .transpose()?; windows::create_window_expr( @@ -2469,10 +2505,13 @@ pub fn create_window_expr_with_name( } /// Create a window expression from a logical expression or an alias +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_window_expr( e: &Expr, logical_schema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result> { // unpack aliased logical expressions, e.g. "sum(col) over () as total" let (name, e) = match e { @@ -2482,7 +2521,7 @@ pub fn create_window_expr( ), _ => (e.schema_name().to_string(), e.clone()), }; - create_window_expr_with_name(&e, name, logical_schema, execution_props) + create_window_expr_with_name(&e, name, logical_schema, execution_props, planning_ctx) } type AggregateExprWithOptionalArgs = ( @@ -2503,11 +2542,13 @@ pub fn create_aggregate_expr_with_name_and_maybe_filter( physical_input_schema: &Schema, execution_props: &ExecutionProps, ) -> Result { + let planning_ctx = PhysicalPlanningContext::default(); let mut builder = LoweredAggregateBuilder::new( e, logical_input_schema, physical_input_schema, execution_props, + &planning_ctx, ) .with_human_display(human_display); @@ -2542,11 +2583,13 @@ pub fn create_aggregate_expr_and_maybe_filter( _ => (None, String::default(), e.clone()), }; + let planning_ctx = PhysicalPlanningContext::default(); let mut builder = LoweredAggregateBuilder::new( &e, logical_input_schema, physical_input_schema, execution_props, + &planning_ctx, ) .with_human_display(human_display); @@ -2572,7 +2615,7 @@ impl DefaultPhysicalPlanner { async fn handle_explain_or_analyze( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result>> { let execution_plan = match logical_plan { LogicalPlan::Explain(e) => self.handle_explain(e, session_state).await?, @@ -2586,13 +2629,15 @@ impl DefaultPhysicalPlanner { async fn handle_explain( &self, e: &Explain, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { use PlanType::*; let mut stringified_plans = vec![]; let config = &session_state.config_options().explain; let explain_format = &e.explain_format; + // Statement-level override wins over session config for show_statistics. + let show_statistics = e.show_statistics.unwrap_or(config.show_statistics); if !e.logical_optimization_succeeded { return Ok(Arc::new(ExplainExec::new( @@ -2665,7 +2710,7 @@ impl DefaultPhysicalPlanner { stringified_plans.push(StringifiedPlan::new( InitialPhysicalPlan, displayable(input.as_ref()) - .set_show_statistics(config.show_statistics) + .set_show_statistics(show_statistics) .set_show_schema(config.show_schema) .indent(e.verbose) .to_string(), @@ -2674,7 +2719,7 @@ impl DefaultPhysicalPlanner { // Show statistics + schema in verbose output even if not // explicitly requested if e.verbose { - if !config.show_statistics { + if !show_statistics { stringified_plans.push(StringifiedPlan::new( InitialPhysicalPlanWithStats, displayable(input.as_ref()) @@ -2703,7 +2748,7 @@ impl DefaultPhysicalPlanner { stringified_plans.push(StringifiedPlan::new( plan_type, displayable(plan) - .set_show_statistics(config.show_statistics) + .set_show_statistics(show_statistics) .set_show_schema(config.show_schema) .indent(e.verbose) .to_string(), @@ -2716,7 +2761,7 @@ impl DefaultPhysicalPlanner { stringified_plans.push(StringifiedPlan::new( FinalPhysicalPlan, displayable(input.as_ref()) - .set_show_statistics(config.show_statistics) + .set_show_statistics(show_statistics) .set_show_schema(config.show_schema) .indent(e.verbose) .to_string(), @@ -2725,7 +2770,7 @@ impl DefaultPhysicalPlanner { // Show statistics + schema in verbose output even if not // explicitly requested if e.verbose { - if !config.show_statistics { + if !show_statistics { stringified_plans.push(StringifiedPlan::new( FinalPhysicalPlanWithStats, displayable(input.as_ref()) @@ -2774,30 +2819,34 @@ impl DefaultPhysicalPlanner { async fn handle_analyze( &self, a: &Analyze, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { let input = self.create_physical_plan(&a.input, session_state).await?; let schema = Arc::clone(a.schema.inner()); let show_statistics = session_state.config_options().explain.show_statistics; - let analyze_level = session_state.config_options().explain.analyze_level; + // Statement-level overrides take precedence over the session config. + let analyze_level = a + .analyze_level + .unwrap_or_else(|| session_state.config_options().explain.analyze_level); let metric_types = analyze_level.included_types(); - let analyze_categories = session_state - .config_options() - .explain - .analyze_categories - .clone(); + let analyze_categories = a.analyze_categories.clone().unwrap_or_else(|| { + session_state + .config_options() + .explain + .analyze_categories + .clone() + }); let metric_categories = match analyze_categories { ExplainAnalyzeCategories::All => None, ExplainAnalyzeCategories::Only(cats) => Some(cats), }; - Ok(Arc::new(AnalyzeExec::new( - a.verbose, - show_statistics, - metric_types, - metric_categories, - input, - schema, - ))) + Ok(Arc::new( + AnalyzeExec::builder(a.verbose, show_statistics, input, schema) + .with_metric_types(metric_types) + .with_metric_categories(metric_categories) + .with_format(a.format.clone()) + .build(), + )) } /// Optimize a physical plan by applying each physical optimizer, @@ -2806,7 +2855,7 @@ impl DefaultPhysicalPlanner { pub fn optimize_physical_plan( &self, plan: Arc, - session_state: &SessionState, + session_state: &dyn Session, mut observer: F, ) -> Result> where @@ -2827,10 +2876,13 @@ impl DefaultPhysicalPlanner { InvariantChecker(InvariantLevel::Always).check(&plan)?; let mut new_plan = Arc::clone(&plan); + let optimizer_context = SessionOptimizerContext { + session: session_state, + }; for optimizer in optimizers { let before_schema = new_plan.schema(); new_plan = optimizer - .optimize_with_context(new_plan, session_state) + .optimize_with_context(new_plan, &optimizer_context) .map_err(|e| { DataFusionError::Context(optimizer.name().to_string(), Box::new(e)) })?; @@ -2910,7 +2962,7 @@ impl DefaultPhysicalPlanner { async fn plan_scalar_subqueries( &self, subqueries: Vec, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result<(Vec, DFHashMap)> { let mut links = Vec::with_capacity(subqueries.len()); let mut index_map = DFHashMap::with_capacity(subqueries.len()); @@ -2935,9 +2987,11 @@ impl DefaultPhysicalPlanner { fn create_project_physical_exec_with_props( &self, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, input_exec: Arc, input: &Arc, expr: &[Expr], + output_schema: &DFSchema, ) -> Result> { let input_logical_schema = input.as_ref().schema(); let input_physical_schema = input_exec.schema(); @@ -2972,8 +3026,12 @@ impl DefaultPhysicalPlanner { physical_name(e) }; - let physical_expr = - create_physical_expr(e, input_logical_schema, execution_props); + let physical_expr = create_physical_expr( + e, + input_logical_schema, + execution_props, + planning_ctx, + ); tuple_err((physical_expr, physical_name)) }) @@ -2991,7 +3049,11 @@ impl DefaultPhysicalPlanner { .into_iter() .map(|(expr, alias)| ProjectionExpr { expr, alias }) .collect(); - Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input_exec)?)) + Ok(Arc::new(ProjectionExec::try_new_with_schema_metadata( + proj_exprs, + input_exec, + output_schema.as_arrow(), + )?)) } PlanAsyncExpr::Async( async_map, @@ -3003,8 +3065,11 @@ impl DefaultPhysicalPlanner { .into_iter() .map(|(expr, alias)| ProjectionExpr { expr, alias }) .collect(); - let new_proj_exec = - ProjectionExec::try_new(proj_exprs, Arc::new(async_exec))?; + let new_proj_exec = ProjectionExec::try_new_with_schema_metadata( + proj_exprs, + Arc::new(async_exec), + output_schema.as_arrow(), + )?; Ok(Arc::new(new_proj_exec)) } _ => internal_err!("Unexpected PlanAsyncExpressions variant"), @@ -3104,15 +3169,14 @@ impl<'a> OptimizationInvariantChecker<'a> { previous_schema: &Arc, ) -> Result<()> { // if the rule is not permitted to change the schema, confirm that it did not change. - if self.rule.schema_check() - && !is_allowed_schema_change(previous_schema.as_ref(), plan.schema().as_ref()) - { - internal_err!( - "PhysicalOptimizer rule '{}' failed. Schema mismatch. Expected original schema: {}, got new schema: {}", - self.rule.name(), - previous_schema, - plan.schema() - )? + if self.rule.schema_check() { + is_allowed_schema_change(previous_schema.as_ref(), plan.schema().as_ref()) + .map_err(|e| { + e.context(format!( + "PhysicalOptimizer rule '{}' failed. Schema mismatch.", + self.rule.name(), + )) + })? } // check invariants per each ExecutionPlan node @@ -3131,28 +3195,45 @@ impl<'a> OptimizationInvariantChecker<'a> { /// This change is allowed because for any field the non-nullable domain `F` is a strict subset /// of the nullable domain `F ∪ { NULL }`. A physical schema that guarantees a stricter subset /// of values will not violate any assumptions made based on the less strict schema. -fn is_allowed_schema_change(old: &Schema, new: &Schema) -> bool { +fn is_allowed_schema_change(old: &Schema, new: &Schema) -> Result<()> { if new.metadata != old.metadata { - return false; + return internal_err!( + "Schema metadata mismatch: Expected original metadata: {:?}, got metadata: {:?}", + old.metadata, + new.metadata + ); } if new.fields.len() != old.fields.len() { - return false; + return internal_err!( + "Schema field mismatch: Expected original field count: {}, got field count: {}", + old.fields.len(), + new.fields.len() + ); } let new_fields = new.fields.iter().map(|f| f.as_ref()); let old_fields = old.fields.iter().map(|f| f.as_ref()); old_fields .zip(new_fields) - .all(|(old, new)| is_allowed_field_change(old, new)) + .try_for_each(|(old, new)| is_allowed_field_change(old, new)) } -fn is_allowed_field_change(old_field: &Field, new_field: &Field) -> bool { - new_field.name() == old_field.name() +fn is_allowed_field_change(old_field: &Field, new_field: &Field) -> Result<()> { + if new_field.name() == old_field.name() && new_field.data_type() == old_field.data_type() && new_field.metadata() == old_field.metadata() && (new_field.is_nullable() == old_field.is_nullable() || !new_field.is_nullable()) + { + Ok(()) + } else { + internal_err!( + "Schema field unallowed change: old field: {:?}, new field: {:?}", + old_field, + new_field + ) + } } impl<'n> TreeNodeVisitor<'n> for OptimizationInvariantChecker<'_> { @@ -3201,40 +3282,202 @@ mod tests { use std::fmt::{self, Debug}; use std::mem::size_of_val; use std::ops::{BitAnd, Not}; + use std::sync::atomic::{AtomicBool, Ordering as AtomicOrdering}; use super::*; use crate::datasource::MemTable; use crate::datasource::file_format::options::CsvReadOptions; use crate::physical_plan::{ - DisplayAs, DisplayFormatType, PlanProperties, SendableRecordBatchStream, - expressions, + DisplayAs, DisplayFormatType, Partitioning, PlanProperties, + SendableRecordBatchStream, expressions, }; use crate::prelude::{SessionConfig, SessionContext}; use crate::test_util::{scan_empty, scan_empty_with_partitions}; + use crate::execution::context::SessionState; use crate::execution::session_state::SessionStateBuilder; + use crate::logical_expr::UserDefinedLogicalNode; use arrow::array::{ArrayRef, DictionaryArray, Int32Array}; use arrow::datatypes::{DataType, Field, Int32Type}; use arrow_schema::{FieldRef, SchemaRef}; - use datafusion_common::config::ConfigOptions; + use datafusion_catalog::CatalogProviderList; + use datafusion_common::config::{ConfigOptions, TableOptions}; use datafusion_common::{ - DFSchemaRef, ScalarValue, TableReference, ToDFSchema as _, assert_batches_eq, - assert_contains, + DFSchemaRef, ScalarValue, SplitPoint, TableReference, ToDFSchema as _, + assert_batches_eq, assert_contains, }; use datafusion_execution::TaskContext; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_expr::builder::subquery_alias; + use datafusion_expr::dml::MergeIntoClause; use datafusion_expr::expr::AggregateFunctionParams; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; + use datafusion_expr::registry::ExtensionTypeRegistryRef; use datafusion_expr::{ - Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, LogicalPlanBuilder, - Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, - WindowFunctionDefinition, col, lit, + Accumulator, AggregateUDF, AggregateUDFImpl, ExprFunctionExt, HigherOrderUDF, + LogicalPlanBuilder, Partitioning as LogicalPartitioning, RangePartitioning, + ScalarUDF, Signature, TableSource, UserDefinedLogicalNodeCore, Volatility, + WindowFunctionDefinition, WindowUDF, col, lit, scalar_subquery, }; use datafusion_functions_aggregate::count::{count_all, count_udaf}; use datafusion_functions_aggregate::expr_fn::sum; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; + use datafusion_session::QueryPlanner; + + #[derive(Debug)] + struct ContextCheckingRule { + invoked: Arc, + } + + impl PhysicalOptimizerRule for ContextCheckingRule { + fn optimize( + &self, + plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + Ok(plan) + } + + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, + ) -> Result> { + assert!(context.statistics_registry().is_some()); + self.invoked.store(true, AtomicOrdering::Relaxed); + Ok(plan) + } + + fn name(&self) -> &str { + "context_checking_rule" + } + + fn schema_check(&self) -> bool { + true + } + } + + #[derive(Debug)] + struct TestQueryPlanner { + invoked: Arc, + } + + #[async_trait] + impl QueryPlanner for TestQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + self.invoked.store(true, AtomicOrdering::Relaxed); + DefaultPhysicalPlanner::default() + .create_physical_plan(logical_plan, session) + .await + } + } + + struct TestSession { + inner: SessionState, + query_planner: Arc, + } + + #[async_trait] + impl Session for TestSession { + fn session_id(&self) -> &str { + self.inner.session_id() + } + + fn config(&self) -> &SessionConfig { + self.inner.config() + } + + fn catalog_list(&self) -> Arc { + Arc::clone(self.inner.catalog_list()) + } + + fn query_planner(&self) -> Arc { + Arc::clone(&self.query_planner) + } + + fn optimize(&self, plan: &LogicalPlan) -> Result { + self.inner.optimize(plan) + } + + fn physical_optimizers(&self) -> &[Arc] { + self.inner.physical_optimizers() + } + + fn statistics_registry( + &self, + ) -> Option<&datafusion_physical_plan::operator_statistics::StatisticsRegistry> + { + self.inner.statistics_registry() + } + + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + ) -> Result> { + let logical_plan = self.optimize(logical_plan)?; + self.query_planner() + .create_physical_plan(&logical_plan, self) + .await + } + + fn create_physical_expr( + &self, + expr: Expr, + df_schema: &DFSchema, + ) -> Result> { + Session::create_physical_expr(&self.inner, expr, df_schema) + } + + fn scalar_functions(&self) -> &HashMap> { + Session::scalar_functions(&self.inner) + } + + fn higher_order_functions(&self) -> &HashMap> { + Session::higher_order_functions(&self.inner) + } + + fn aggregate_functions(&self) -> &HashMap> { + Session::aggregate_functions(&self.inner) + } + + fn window_functions(&self) -> &HashMap> { + Session::window_functions(&self.inner) + } + + fn extension_type_registry(&self) -> &ExtensionTypeRegistryRef { + Session::extension_type_registry(&self.inner) + } + + fn runtime_env(&self) -> &Arc { + self.inner.runtime_env() + } + + fn execution_props(&self) -> &ExecutionProps { + self.inner.execution_props() + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn table_options(&self) -> &TableOptions { + self.inner.table_options() + } + + fn table_options_mut(&mut self) -> &mut TableOptions { + self.inner.table_options_mut() + } + + fn task_ctx(&self) -> Arc { + self.inner.task_ctx() + } + } fn make_session_state() -> SessionState { let runtime = Arc::new(RuntimeEnv::default()); @@ -3247,6 +3490,85 @@ mod tests { .build() } + #[derive(Debug)] + struct CaptureMergeProvider { + schema: SchemaRef, + captured: Mutex>, + } + + #[async_trait] + impl TableProvider for CaptureMergeProvider { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + _projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Ok(Arc::new(EmptyExec::new(Arc::clone(&self.schema)))) + } + + async fn merge_into( + &self, + state: &dyn Session, + source: Arc, + merge_schema: DFSchemaRef, + on: Expr, + clauses: Vec, + ) -> Result> { + let physical_on = state.create_physical_expr(on, &merge_schema)?; + *self.captured.lock().await = + Some((merge_schema, format!("{physical_on:?}"), clauses.len())); + Ok(source) + } + } + + #[tokio::test] + async fn merge_into_provider_receives_combined_logical_schema() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let target = Arc::new(CaptureMergeProvider { + schema: Arc::clone(&schema), + captured: Mutex::new(None), + }); + let source = Arc::new(MemTable::try_new(Arc::clone(&schema), vec![vec![]])?); + let ctx = SessionContext::new(); + ctx.register_table("target", target.clone())?; + ctx.register_table("source", source)?; + + ctx.sql( + "MERGE INTO target AS t USING source AS s ON t.id = s.id \ + WHEN MATCHED AND t.id > s.id THEN DELETE", + ) + .await? + .create_physical_plan() + .await?; + + let captured = target.captured.lock().await; + let (merge_schema, physical_on, clause_count) = + captured.as_ref().expect("merge_into should be called"); + assert_eq!(*clause_count, 1); + assert_eq!( + merge_schema.index_of_column(&Column::new(Some("target"), "id"))?, + 0 + ); + assert_eq!( + merge_schema.index_of_column(&Column::new(Some("s"), "id"))?, + 1 + ); + assert_contains!(physical_on, "index: 0"); + assert_contains!(physical_on, "index: 1"); + Ok(()) + } + async fn plan(logical_plan: &LogicalPlan) -> Result> { let session_state = make_session_state(); // optimize the logical plan @@ -3257,6 +3579,35 @@ mod tests { .await } + #[tokio::test] + async fn plans_with_non_session_state_implementation() -> Result<()> { + let invoked = Arc::new(AtomicBool::new(false)); + let inner = SessionStateBuilder::new() + .with_default_features() + .with_physical_optimizer_rules(vec![Arc::new(ContextCheckingRule { + invoked: Arc::clone(&invoked), + })]) + .with_statistics_registry( + datafusion_physical_plan::operator_statistics::StatisticsRegistry::new(), + ) + .build(); + let query_planner_invoked = Arc::new(AtomicBool::new(false)); + let session = TestSession { + inner, + query_planner: Arc::new(TestQueryPlanner { + invoked: Arc::clone(&query_planner_invoked), + }), + }; + assert!(session.as_any().downcast_ref::().is_none()); + + let logical_plan = LogicalPlanBuilder::empty(false).build()?; + let physical_plan = session.create_physical_plan(&logical_plan).await?; + assert!(physical_plan.is::()); + assert!(query_planner_invoked.load(AtomicOrdering::Relaxed)); + assert!(invoked.load(AtomicOrdering::Relaxed)); + Ok(()) + } + async fn aggregate_explain(logical_plan: &LogicalPlan) -> Result { let physical_plan = plan(logical_plan).await?; Ok(displayable(physical_plan.as_ref()).indent(true).to_string()) @@ -3277,6 +3628,46 @@ mod tests { Field::new(name, DataType::Int64, nullable) } + #[tokio::test] + async fn logical_range_repartition_plans_output_partitioning() -> Result<()> { + let batch = RecordBatch::try_from_iter(vec![( + "a", + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )])?; + let table = Arc::new(MemTable::try_new(batch.schema(), vec![vec![batch]])?); + let source = Arc::new(DefaultTableSource::new(table)); + let logical_plan = LogicalPlanBuilder::scan("test", source, None)? + .repartition(LogicalPartitioning::Range(RangePartitioning::try_new( + vec![col("a").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?))? + .build()?; + + let planner = DefaultPhysicalPlanner::default(); + let physical_plan = planner + .create_initial_plan(&logical_plan, &make_session_state()) + .await?; + let repartition = physical_plan + .as_ref() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!( + "expected RepartitionExec, got {}", + physical_plan.name() + ) + })?; + let Partitioning::Range(range) = repartition.partitioning() else { + return internal_err!( + "expected Range target partitioning, got {:?}", + repartition.partitioning() + ); + }; + assert_eq!(range.partition_count(), 2); + assert_eq!(physical_plan.output_partitioning().partition_count(), 2); + + Ok(()) + } + #[test] fn test_create_window_expr_unwraps_alias_with_metadata() -> Result<()> { use std::collections::HashMap; @@ -3299,13 +3690,54 @@ mod tests { )) .alias_with_metadata("window_alias", Some(metadata)); - let window_expr = - create_window_expr(&expr, &logical_schema, &ExecutionProps::new())?; + let window_expr = create_window_expr( + &expr, + &logical_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; assert_eq!(window_expr.name(), "window_alias"); Ok(()) } + #[tokio::test] + async fn test_projection_preserves_field_metadata_for_aggregate() -> Result<()> { + use datafusion_common::metadata::FieldMetadata; + use datafusion_expr::expr::AggregateFunction; + use datafusion_functions_aggregate::min_max::max_udaf; + + let schema = Schema::new(vec![Field::new("value", DataType::Utf8, false)]); + let input = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(schema.to_dfschema()?), + }); + let metadata = + FieldMetadata::from(HashMap::from([("foo".to_string(), "bar".to_string())])); + let projection = LogicalPlan::Projection(Projection::try_new( + vec![col("value").alias_with_metadata("value", Some(metadata))], + Arc::new(input), + )?); + let aggregate = LogicalPlan::Aggregate(Aggregate::try_new( + Arc::new(projection), + vec![], + vec![Expr::AggregateFunction(AggregateFunction::new_udf( + max_udaf(), + vec![col("value")], + false, + None, + vec![], + None, + ))], + )?); + + DefaultPhysicalPlanner::default() + .create_physical_plan(&aggregate, &SessionContext::new().state()) + .await?; + + Ok(()) + } + #[derive(Debug, Default)] struct NullAccumulator; @@ -3555,6 +3987,7 @@ mod tests { logical_input_schema, physical_input_schema, session_state.execution_props(), + &PhysicalPlanningContext::default(), ); insta::assert_debug_snapshot!(cube, @r#" @@ -3686,6 +4119,7 @@ mod tests { logical_input_schema, physical_input_schema, session_state.execution_props(), + &PhysicalPlanningContext::default(), ); insta::assert_debug_snapshot!(rollup, @r#" @@ -3790,6 +4224,7 @@ mod tests { &col("a").not(), &dfschema, &make_session_state(), + &PhysicalPlanningContext::default(), )?; let expected = expressions::not(expressions::col("a", &schema)?)?; @@ -3817,6 +4252,36 @@ mod tests { Ok(()) } + #[tokio::test] + async fn scalar_subquery_in_extension_expr_plans() -> Result<()> { + let subquery = LogicalPlanBuilder::empty(true) + .project(vec![lit(42_i32)])? + .build()?; + let logical_plan = LogicalPlan::Extension(Extension { + node: Arc::new(NoOpExtensionNode { + expressions: vec![scalar_subquery(Arc::new(subquery))], + ..Default::default() + }), + }); + let planner = DefaultPhysicalPlanner::with_extension_planners(vec![Arc::new( + ExpressionExtensionPlanner, + )]); + let session = TestSession { + inner: make_session_state(), + query_planner: Arc::new(TestQueryPlanner { + invoked: Arc::new(AtomicBool::new(false)), + }), + }; + assert!(session.as_any().downcast_ref::().is_none()); + + let plan = planner + .create_physical_plan(&logical_plan, &session) + .await?; + + assert_contains!(format!("{plan:?}"), "ScalarSubqueryExec"); + Ok(()) + } + #[tokio::test] async fn error_during_extension_planning() { let session_state = make_session_state(); @@ -4308,6 +4773,7 @@ mod tests { stringified_plans, schema: schema.to_dfschema_ref().unwrap(), logical_optimization_succeeded: false, + show_statistics: None, }; let plan = planner .handle_explain(&explain, &ctx.state()) @@ -4336,7 +4802,8 @@ mod tests { _node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc], - _session_state: &SessionState, + _session_state: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { internal_err!("BOOM") } @@ -4345,6 +4812,7 @@ mod tests { #[derive(PartialEq, Eq, Hash)] struct NoOpExtensionNode { schema: DFSchemaRef, + expressions: Vec, } impl Default for NoOpExtensionNode { @@ -4357,6 +4825,7 @@ mod tests { ) .unwrap(), ), + expressions: vec![], } } } @@ -4389,7 +4858,7 @@ mod tests { } fn expressions(&self) -> Vec { - vec![] + self.expressions.clone() } fn fmt_for_explain(&self, f: &mut fmt::Formatter) -> fmt::Result { @@ -4398,10 +4867,13 @@ mod tests { fn with_exprs_and_inputs( &self, - _exprs: Vec, + exprs: Vec, _inputs: Vec, ) -> Result { - unimplemented!("NoOp"); + Ok(Self { + schema: Arc::clone(&self.schema), + expressions: exprs, + }) } fn supports_limit_pushdown(&self) -> bool { @@ -4461,11 +4933,26 @@ mod tests { vec![] } + fn replace_children( + self: Arc, + children: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + if children.is_empty() { + Ok(self) + } else { + exec_err!("NoOpExecutionPlan does not support children") + } + } + fn with_new_children( self: Arc, - _children: Vec>, + children: Vec>, ) -> Result> { - unimplemented!("NoOpExecutionPlan::with_new_children"); + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn execute( @@ -4478,16 +4965,36 @@ mod tests { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } + Ok(TreeNodeRecursion::Continue) + } + } + + struct ExpressionExtensionPlanner; + + #[async_trait] + impl ExtensionPlanner for ExpressionExtensionPlanner { + async fn plan_extension( + &self, + planner: &dyn PhysicalPlanner, + node: &dyn UserDefinedLogicalNode, + _logical_inputs: &[&LogicalPlan], + _physical_inputs: &[Arc], + session_state: &dyn Session, + planning_ctx: &PhysicalPlanningContext, + ) -> Result>> { + for expr in node.expressions() { + planner.create_physical_expr( + &expr, + node.schema(), + session_state, + planning_ctx, + )?; } - Ok(tnr) + Ok(Some(Arc::new(NoOpExecutionPlan::new(Arc::clone( + node.schema().inner(), + ))))) } } @@ -4504,7 +5011,8 @@ mod tests { _node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc], - _session_state: &SessionState, + _session_state: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { Ok(Some(Arc::new(NoOpExecutionPlan::new(SchemaRef::new( Schema::new(vec![Field::new("b", DataType::Int32, false)]), @@ -4606,12 +5114,22 @@ digraph { fn name(&self) -> &str { "always ok" } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self(children))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } @@ -4630,7 +5148,7 @@ digraph { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -4661,12 +5179,22 @@ digraph { fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn children(&self) -> Vec<&Arc> { unimplemented!() } @@ -4682,7 +5210,7 @@ digraph { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -4720,10 +5248,16 @@ digraph { // ok plan let ok_node: Arc = Arc::new(OkExtensionNode(vec![])); let child = Arc::clone(&ok_node); - let ok_plan = Arc::clone(&ok_node).with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&child)])?, - Arc::clone(&child), - ])?; + let ok_plan = Arc::clone(&ok_node).replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&child)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?, + Arc::clone(&child), + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; // Test: check should pass with same schema let equal_schema = ok_plan.schema(); @@ -4735,33 +5269,45 @@ digraph { let expected_err = OptimizationInvariantChecker::new(&rule) .check(&ok_plan, &different_schema) .unwrap_err(); - assert!(expected_err.to_string().contains("PhysicalOptimizer rule 'OptimizerRuleWithSchemaCheck' failed. Schema mismatch. Expected original schema")); + assert!(expected_err.to_string().contains("PhysicalOptimizer rule 'OptimizerRuleWithSchemaCheck' failed. Schema mismatch.")); + + // The recursive `check_invariants` walk only runs under `debug_assertions` + // (see `OptimizationInvariantChecker::check`). In release builds the walk is + // skipped, so the checker returns `Ok` rather than surfacing the node's error. // Test: should fail when extension node fails it's own invariant check let failing_node: Arc = Arc::new(InvariantFailsExtensionNode); - let expected_err = OptimizationInvariantChecker::new(&rule) - .check(&failing_node, &ok_plan.schema()) - .unwrap_err(); - assert!( - expected_err.to_string().contains( + let result = OptimizationInvariantChecker::new(&rule) + .check(&failing_node, &ok_plan.schema()); + if cfg!(debug_assertions) { + assert!(result.unwrap_err().to_string().contains( "extension node failed it's user-defined always-invariant check" - ) - ); + )); + } else { + assert!(result.is_ok()); + } // Test: should fail when descendent extension node fails let failing_node: Arc = Arc::new(InvariantFailsExtensionNode); - let invalid_plan = ok_node.with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&failing_node)])?, - Arc::clone(&child), - ])?; - let expected_err = OptimizationInvariantChecker::new(&rule) - .check(&invalid_plan, &ok_plan.schema()) - .unwrap_err(); - assert!( - expected_err.to_string().contains( + let invalid_plan = ok_node.replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&failing_node)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?, + Arc::clone(&child), + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let result = OptimizationInvariantChecker::new(&rule) + .check(&invalid_plan, &ok_plan.schema()); + if cfg!(debug_assertions) { + assert!(result.unwrap_err().to_string().contains( "extension node failed it's user-defined always-invariant check" - ) - ); + )); + } else { + assert!(result.is_ok()); + } Ok(()) } @@ -4785,12 +5331,22 @@ digraph { fn schema(&self) -> SchemaRef { Arc::new(Schema::empty()) } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn children(&self) -> Vec<&Arc> { vec![] } @@ -4806,7 +5362,7 @@ digraph { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -4837,10 +5393,16 @@ digraph { let failing_node: Arc = Arc::new(ExecutableInvariantFails); let ok_node: Arc = Arc::new(OkExtensionNode(vec![])); let child = Arc::clone(&ok_node); - let plan = ok_node.with_new_children(vec![ - Arc::clone(&child).with_new_children(vec![Arc::clone(&failing_node)])?, - Arc::clone(&child), - ])?; + let plan = ok_node.replace_children( + vec![ + Arc::clone(&child).replace_children( + vec![Arc::clone(&failing_node)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?, + Arc::clone(&child), + ], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; let expected_err = InvariantChecker(InvariantLevel::Executable) .check(&plan) .unwrap_err(); @@ -4962,9 +5524,8 @@ digraph { } #[tokio::test] - // When schemas match, planning proceeds past the schema_satisfied_by check. - // It then panics on unimplemented error in NoOpExecutionPlan. - #[should_panic(expected = "NoOpExecutionPlan")] + // When schemas match, planning proceeds past the schema_satisfied_by check + // and succeeds. async fn test_aggregate_schema_check_passes() { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); @@ -5150,7 +5711,8 @@ digraph { _node: &dyn UserDefinedLogicalNode, _logical_inputs: &[&LogicalPlan], _physical_inputs: &[Arc], - _session_state: &SessionState, + _session_state: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { Ok(None) } @@ -5159,7 +5721,8 @@ digraph { &self, _planner: &dyn PhysicalPlanner, scan: &TableScan, - _session_state: &SessionState, + _session_state: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { if scan.source.is::() { Ok(Some(Arc::new(EmptyExec::new(Arc::clone( diff --git a/datafusion/core/src/test/mod.rs b/datafusion/core/src/test/mod.rs index 717182f1d3d5b..f46a5a0749065 100644 --- a/datafusion/core/src/test/mod.rs +++ b/datafusion/core/src/test/mod.rs @@ -103,7 +103,7 @@ pub fn scan_partitioned_csv( quote: b'"', ..Default::default() }; - let table_schema = TableSchema::from_file_schema(schema); + let table_schema = TableSchema::from(schema); let source = Arc::new(CsvSource::new(table_schema.clone()).with_csv_options(options)); let config = FileScanConfigBuilder::from(partitioned_csv_config(file_groups, source)?) diff --git a/datafusion/core/src/test_util/mod.rs b/datafusion/core/src/test_util/mod.rs index aad659eacbe55..d70c0d186d007 100644 --- a/datafusion/core/src/test_util/mod.rs +++ b/datafusion/core/src/test_util/mod.rs @@ -45,7 +45,7 @@ use crate::execution::{SendableRecordBatchStream, SessionState, SessionStateBuil use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_catalog::Session; -use datafusion_common::{DFSchemaRef, TableReference}; +use datafusion_common::{DFSchemaRef, TableReference, plan_err}; use datafusion_expr::{ CreateExternalTable, Expr, LogicalPlan, SortExpr, TableType, UserDefinedLogicalNodeCore, @@ -187,8 +187,12 @@ impl TableProviderFactory for TestTableFactory { _: &dyn Session, cmd: &CreateExternalTable, ) -> Result> { + let Some(location) = cmd.locations.first() else { + return plan_err!("TestTableFactory requires at least one location"); + }; + Ok(Arc::new(TestTableProvider { - url: cmd.location.to_string(), + url: location.clone(), schema: Arc::clone(cmd.schema.inner()), })) } diff --git a/datafusion/core/src/test_util/parquet.rs b/datafusion/core/src/test_util/parquet.rs index c53495421307b..e25fe746695cf 100644 --- a/datafusion/core/src/test_util/parquet.rs +++ b/datafusion/core/src/test_util/parquet.rs @@ -29,6 +29,7 @@ use crate::datasource::object_store::ObjectStoreUrl; use crate::datasource::physical_plan::ParquetSource; use crate::error::Result; use crate::logical_expr::execution_props::ExecutionProps; +use crate::logical_expr::physical_planning_context::PhysicalPlanningContext; use crate::logical_expr::simplify::SimplifyContext; use crate::optimizer::simplify_expressions::ExprSimplifier; use crate::physical_expr::create_physical_expr; @@ -149,7 +150,7 @@ impl TestParquetFile { /// ``` /// /// Otherwise if `maybe_filter` is None, return just a `DataSourceExec` - pub async fn create_scan( + pub fn create_scan( &self, ctx: &SessionContext, maybe_filter: Option, @@ -172,8 +173,12 @@ impl TestParquetFile { if let Some(filter) = maybe_filter { let simplifier = ExprSimplifier::new(context); let filter = simplifier.coerce(filter, &df_schema).unwrap(); - let physical_filter_expr = - create_physical_expr(&filter, &df_schema, &ExecutionProps::default())?; + let physical_filter_expr = create_physical_expr( + &filter, + &df_schema, + &ExecutionProps::default(), + &PhysicalPlanningContext::default(), + )?; let source = Arc::new( ParquetSource::new(Arc::clone(&self.schema)) diff --git a/datafusion/core/tests/config_from_env.rs b/datafusion/core/tests/config_from_env.rs index 6375d4e25d8eb..15a047cbbda51 100644 --- a/datafusion/core/tests/config_from_env.rs +++ b/datafusion/core/tests/config_from_env.rs @@ -16,6 +16,7 @@ // under the License. use datafusion::config::ConfigOptions; +use datafusion_common::assert_contains; use std::env; #[test] @@ -34,7 +35,7 @@ fn from_env() { // invalid testing env::set_var(env_key, "ttruee"); let err = ConfigOptions::from_env().unwrap_err().strip_backtrace(); - assert_eq!( + assert_contains!( err, "Error parsing 'ttruee' as bool\ncaused by\nExternal error: provided string was not `true` or `false`" ); @@ -45,18 +46,18 @@ fn from_env() { // for valid testing env::set_var(env_key, "4096"); let config = ConfigOptions::from_env().unwrap(); - assert_eq!(config.execution.batch_size, 4096); + assert_eq!(config.execution.batch_size.get(), 4096); // for invalid testing env::set_var(env_key, "abc"); let err = ConfigOptions::from_env().unwrap_err().strip_backtrace(); - assert_eq!( + assert_contains!( err, "Error parsing 'abc' as usize\ncaused by\nExternal error: invalid digit found in string" ); env::remove_var(env_key); let config = ConfigOptions::from_env().unwrap(); - assert_eq!(config.execution.batch_size, 8192); // set to its default value + assert_eq!(config.execution.batch_size.get(), 8192); // set to its default value } } diff --git a/datafusion/core/tests/core_integration.rs b/datafusion/core/tests/core_integration.rs index f85538b5c3405..9b350e5529bcf 100644 --- a/datafusion/core/tests/core_integration.rs +++ b/datafusion/core/tests/core_integration.rs @@ -63,6 +63,9 @@ mod tracing; /// Run all tests that are found in the `extension_types` directory mod extension_types; +/// Helper functions for tests. +mod helper; + #[cfg(test)] #[ctor::ctor(unsafe)] fn init() { diff --git a/datafusion/core/tests/custom_sources_cases/mod.rs b/datafusion/core/tests/custom_sources_cases/mod.rs index cef75b444f6fe..7abbcd6e9578c 100644 --- a/datafusion/core/tests/custom_sources_cases/mod.rs +++ b/datafusion/core/tests/custom_sources_cases/mod.rs @@ -40,9 +40,12 @@ use datafusion_common::project_schema; use datafusion_common::stats::Precision; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_physical_expr::EquivalenceProperties; -use datafusion_physical_plan::PlanProperties; +use datafusion_physical_plan::StatisticsArgs; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, PlanProperties, ReplaceChildrenOptions, +}; use async_trait::async_trait; use futures::stream::Stream; @@ -164,13 +167,24 @@ impl ExecutionPlan for CustomExecutionPlan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -179,8 +193,12 @@ impl ExecutionPlan for CustomExecutionPlan { Ok(Box::pin(TestCustomRecordBatchStream { nb_batch: 1 })) } - fn partition_statistics(&self, partition: Option) -> Result> { - if partition.is_some() { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + if args.partition().is_some() { return Ok(Arc::new(Statistics::new_unknown(&self.schema()))); } let batch = TEST_CUSTOM_RECORD_BATCH!().unwrap(); @@ -208,18 +226,11 @@ impl ExecutionPlan for CustomExecutionPlan { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs index e52c559ec79ef..a8f7f09ad016b 100644 --- a/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs +++ b/datafusion/core/tests/custom_sources_cases/provider_filter_pushdown.rs @@ -41,6 +41,7 @@ use datafusion_expr::expr::{BinaryExpr, Cast}; use datafusion_functions_aggregate::expr_fn::count; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use async_trait::async_trait; @@ -117,9 +118,10 @@ impl ExecutionPlan for CustomPlan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { // CustomPlan has no children if children.is_empty() { @@ -129,6 +131,16 @@ impl ExecutionPlan for CustomPlan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -152,18 +164,11 @@ impl ExecutionPlan for CustomPlan { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion/core/tests/custom_sources_cases/statistics.rs b/datafusion/core/tests/custom_sources_cases/statistics.rs index 01c4deac5ccd3..6213be8e2d24f 100644 --- a/datafusion/core/tests/custom_sources_cases/statistics.rs +++ b/datafusion/core/tests/custom_sources_cases/statistics.rs @@ -37,6 +37,9 @@ use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{project_schema, stats::Precision}; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, StatisticsArgs, StatisticsContext, +}; use async_trait::async_trait; @@ -159,13 +162,24 @@ impl ExecutionPlan for StatisticsValidation { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -174,8 +188,12 @@ impl ExecutionPlan for StatisticsValidation { unimplemented!("This plan only serves for testing statistics") } - fn partition_statistics(&self, partition: Option) -> Result> { - if partition.is_some() { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + if args.partition().is_some() { Ok(Arc::new(Statistics::new_unknown(&self.schema))) } else { Ok(Arc::new(self.stats.clone())) @@ -184,18 +202,11 @@ impl ExecutionPlan for StatisticsValidation { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } @@ -247,7 +258,11 @@ async fn sql_basic() -> Result<()> { let physical_plan = df.create_physical_plan().await.unwrap(); // the statistics should be those of the source - assert_eq!(stats, *physical_plan.partition_statistics(None)?); + assert_eq!( + stats, + *StatisticsContext::new() + .compute(physical_plan.as_ref(), &StatisticsArgs::new())? + ); Ok(()) } @@ -263,7 +278,8 @@ async fn sql_filter() -> Result<()> { .unwrap(); let physical_plan = df.create_physical_plan().await.unwrap(); - let stats = physical_plan.partition_statistics(None)?; + let stats = StatisticsContext::new() + .compute(physical_plan.as_ref(), &StatisticsArgs::new())?; assert_eq!(stats.num_rows, Precision::Inexact(7)); Ok(()) @@ -278,7 +294,8 @@ async fn sql_limit() -> Result<()> { let physical_plan = df.create_physical_plan().await.unwrap(); // when the limit is smaller than the original number of lines we mark the statistics as inexact // and cap NDV at the new row count - let limit_stats = physical_plan.partition_statistics(None)?; + let limit_stats = StatisticsContext::new() + .compute(physical_plan.as_ref(), &StatisticsArgs::new())?; assert_eq!(limit_stats.num_rows, Precision::Exact(5)); // c1: NDV=2 stays at 2 (already below limit of 5) assert_eq!( @@ -297,7 +314,11 @@ async fn sql_limit() -> Result<()> { .unwrap(); let physical_plan = df.create_physical_plan().await.unwrap(); // when the limit is larger than the original number of lines, statistics remain unchanged - assert_eq!(stats, *physical_plan.partition_statistics(None)?); + assert_eq!( + stats, + *StatisticsContext::new() + .compute(physical_plan.as_ref(), &StatisticsArgs::new())? + ); Ok(()) } @@ -314,7 +335,8 @@ async fn sql_window() -> Result<()> { let physical_plan = df.create_physical_plan().await.unwrap(); - let result = physical_plan.partition_statistics(None)?; + let result = StatisticsContext::new() + .compute(physical_plan.as_ref(), &StatisticsArgs::new())?; assert_eq!(stats.num_rows, result.num_rows); let col_stats = &result.column_statistics; diff --git a/datafusion/core/tests/data/int_to_float_cast_precision.csv b/datafusion/core/tests/data/int_to_float_cast_precision.csv new file mode 100644 index 0000000000000..187d7affca616 --- /dev/null +++ b/datafusion/core/tests/data/int_to_float_cast_precision.csv @@ -0,0 +1,3 @@ +k,v +1,16777217 +2,16777216 diff --git a/datafusion/core/tests/dataframe/describe.rs b/datafusion/core/tests/dataframe/describe.rs index 9aa8a49c97ae3..056bc21a69186 100644 --- a/datafusion/core/tests/dataframe/describe.rs +++ b/datafusion/core/tests/dataframe/describe.rs @@ -15,6 +15,12 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + +use arrow::array::{ + BinaryArray, BinaryViewArray, FixedSizeBinaryArray, LargeBinaryArray, RecordBatch, +}; +use arrow::datatypes::{DataType, Field, Schema}; use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion_common::test_util::batches_to_string; use datafusion_common::{Result, test_util::parquet_test_data}; @@ -112,6 +118,59 @@ async fn describe_null() -> Result<()> { Ok(()) } +#[tokio::test] +async fn describe_binary_columns() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("bin", DataType::Binary, true), + Field::new("lbin", DataType::LargeBinary, true), + Field::new("vbin", DataType::BinaryView, true), + Field::new("fbin", DataType::FixedSizeBinary(2), true), + ])); + + let bin: BinaryArray = vec![Some([0x00u8, 0x01]), Some([0xff, 0xee]), None] + .into_iter() + .collect(); + let lbin: LargeBinaryArray = vec![Some([0x00u8, 0x01]), Some([0xff, 0xee]), None] + .into_iter() + .collect(); + let vbin: BinaryViewArray = vec![Some([0x00u8, 0x01]), Some([0xff, 0xee]), None] + .into_iter() + .collect(); + let fbin = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + [Some([0x00u8, 0x01]), Some([0xff, 0xee]), None].into_iter(), + 2, + )?; + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(bin), + Arc::new(lbin), + Arc::new(vbin), + Arc::new(fbin), + ], + )?; + let ctx = SessionContext::new(); + ctx.register_batch("t", batch)?; + let result = ctx.table("t").await?.describe().await?.collect().await?; + + assert_snapshot!(batches_to_string(&result), + @r" + +------------+------+------+------+------+ + | describe | bin | lbin | vbin | fbin | + +------------+------+------+------+------+ + | count | 2 | 2 | 2 | 2 | + | null_count | 1 | 1 | 1 | 1 | + | mean | null | null | null | null | + | std | null | null | null | null | + | min | 0001 | 0001 | 0001 | 0001 | + | max | ffee | ffee | ffee | ffee | + | median | null | null | null | null | + +------------+------+------+------+------+ +"); + + Ok(()) +} + /// Return a SessionContext with parquet file registered async fn parquet_context() -> SessionContext { let ctx = SessionContext::new(); diff --git a/datafusion/core/tests/dataframe/mod.rs b/datafusion/core/tests/dataframe/mod.rs index e55a373adab9f..4966c7aa9ae73 100644 --- a/datafusion/core/tests/dataframe/mod.rs +++ b/datafusion/core/tests/dataframe/mod.rs @@ -21,9 +21,10 @@ mod describe; use arrow::array::{ Array, ArrayRef, BooleanArray, DictionaryArray, FixedSizeListArray, - FixedSizeListBuilder, Float32Array, Float64Array, Int8Array, Int32Array, - Int32Builder, LargeListArray, ListArray, ListBuilder, RecordBatch, StringArray, - StringBuilder, StructBuilder, UInt32Array, UInt32Builder, UnionArray, record_batch, + FixedSizeListBuilder, Float16Array, Float32Array, Float64Array, Int8Array, + Int16Array, Int32Array, Int32Builder, Int64Array, LargeListArray, ListArray, + ListBuilder, RecordBatch, StringArray, StringBuilder, StructBuilder, UInt8Array, + UInt16Array, UInt32Array, UInt32Builder, UInt64Array, UnionArray, record_batch, }; use arrow::buffer::ScalarBuffer; use arrow::datatypes::{ @@ -39,6 +40,7 @@ use datafusion_functions_aggregate::expr_fn::{ array_agg, avg, avg_distinct, count, count_distinct, max, median, min, sum, sum_distinct, }; +use datafusion_functions_nested::expr_fn::{array_filter, array_transform, make_array}; use datafusion_functions_nested::make_array::make_array_udf; use datafusion_functions_window::expr_fn::{first_value, lead, row_number}; use insta::assert_snapshot; @@ -78,8 +80,8 @@ use datafusion_expr::{ CreateMemoryTable, CreateView, DdlStatement, Expr, ExprFunctionExt, ExprSchemable, LogicalPlan, LogicalPlanBuilder, ScalarFunctionImplementation, SortExpr, TableType, WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, cast, col, - create_udf, exists, in_subquery, lit, out_ref_col, placeholder, scalar_subquery, - when, wildcard, + create_udf, exists, in_subquery, lambda, lambda_var, lit, out_ref_col, placeholder, + scalar_subquery, when, wildcard, }; use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::aggregate::AggregateExprBuilder; @@ -90,7 +92,9 @@ use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; use datafusion_physical_plan::empty::EmptyExec; -use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties, displayable}; +use datafusion_physical_plan::{ + ExecutionPlan, ExecutionPlanProperties, collect, displayable, +}; use datafusion::error::Result as DataFusionResult; use datafusion::execution::options::JsonReadOptions; @@ -840,7 +844,7 @@ async fn test_aggregate_with_pk() -> Result<()> { let aggr_expr = vec![]; let df = df.aggregate(group_expr, aggr_expr)?; - // Since id and name are functionally dependant, we can use name among + // Since id and name are functionally dependent, we can use name among // expression even if it is not part of the group by expression and can // select "name" column even though it wasn't explicitly grouped let df = df.select(vec![col("id"), col("name")])?; @@ -895,7 +899,7 @@ async fn test_aggregate_with_pk2() -> Result<()> { " ); - // Since id and name are functionally dependant, we can use name among expression + // Since id and name are functionally dependent, we can use name among expression // even if it is not part of the group by expression. let df_results = df.collect().await?; @@ -943,7 +947,7 @@ async fn test_aggregate_with_pk3() -> Result<()> { " ); - // Since id and name are functionally dependant, we can use name among expression + // Since id and name are functionally dependent, we can use name among expression // even if it is not part of the group by expression. let df_results = df.collect().await?; @@ -1141,7 +1145,13 @@ async fn test_aggregate_name_collision() -> Result<()> { // The select expr has the same display_name as the group_expr, // but since they are different expressions, it should fail. .expect_err("Expected error"); - assert_snapshot!(df.strip_backtrace(), @r#"Schema error: No field named aggregate_test_100.c2. Valid fields are "aggregate_test_100.c2 + aggregate_test_100.c3"."#); + assert_snapshot!( + df.strip_backtrace(), + @r#" +Schema error: No field named aggregate_test_100.c2. +Valid fields are "aggregate_test_100.c2 + aggregate_test_100.c3". +"# + ); Ok(()) } @@ -1204,7 +1214,7 @@ async fn window_using_aggregates() -> Result<()> { +-------------+----------+-----------------+---------------+--------+-----+------+----+------+ | first_value | last_val | approx_distinct | approx_median | median | max | min | c2 | c3 | +-------------+----------+-----------------+---------------+--------+-----+------+----+------+ - | | | | | | | | 1 | -85 | + | | | 0 | | | | | 1 | -85 | | -85 | -101 | 14 | -12.0 | -12.0 | 83 | -101 | 4 | -54 | | -85 | -101 | 17 | -25.0 | -25.0 | 83 | -101 | 5 | -31 | | -85 | -12 | 10 | -32.75 | -34.0 | 83 | -85 | 3 | 13 | @@ -3001,22 +3011,22 @@ async fn test_count_wildcard_on_sort() -> Result<()> { assert_snapshot!( pretty_format_batches(&sql_results).unwrap(), @r" - +---------------+------------------------------------------------------------------------------------+ - | plan_type | plan | - +---------------+------------------------------------------------------------------------------------+ - | logical_plan | Sort: count(*) ASC NULLS LAST | - | | Projection: t1.b, count(Int64(1)) AS count(*) | - | | Aggregate: groupBy=[[t1.b]], aggr=[[count(Int64(1))]] | - | | TableScan: t1 projection=[b] | - | physical_plan | SortPreservingMergeExec: [count(*)@1 ASC NULLS LAST] | - | | SortExec: expr=[count(*)@1 ASC NULLS LAST], preserve_partitioning=[true] | - | | ProjectionExec: expr=[b@0 as b, count(Int64(1))@1 as count(*)] | - | | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[count(Int64(1))] | - | | RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=1 | - | | AggregateExec: mode=Partial, gby=[b@0 as b], aggr=[count(Int64(1))] | - | | DataSourceExec: partitions=1, partition_sizes=[1] | - | | | - +---------------+------------------------------------------------------------------------------------+ + +---------------+-------------------------------------------------------------------------------------+ + | plan_type | plan | + +---------------+-------------------------------------------------------------------------------------+ + | logical_plan | Sort: count(*) ASC NULLS LAST | + | | Projection: t1.b, count(Int64(1)) AS count(*) | + | | Aggregate: groupBy=[[t1.b]], aggr=[[count(Int64(1))]] | + | | TableScan: t1 projection=[b] | + | physical_plan | SortPreservingMergeExec: [count(*)@1 ASC NULLS LAST] | + | | ProjectionExec: expr=[b@0 as b, count(Int64(1))@1 as count(*)] | + | | SortExec: expr=[count(Int64(1))@1 ASC NULLS LAST], preserve_partitioning=[true] | + | | AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[count(Int64(1))] | + | | RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=1 | + | | AggregateExec: mode=Partial, gby=[b@0 as b], aggr=[count(Int64(1))] | + | | DataSourceExec: partitions=1, partition_sizes=[1] | + | | | + +---------------+-------------------------------------------------------------------------------------+ " ); @@ -3345,7 +3355,11 @@ async fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl( // To be able to remove user specific paths from the plan, for stable assertions let testdata_clean = Path::new(&testdata).canonicalize()?.display().to_string(); - let testdata_clean = testdata_clean.strip_prefix("/").unwrap_or(&testdata_clean); + let testdata_clean = testdata_clean.replace("\\", "/"); + let testdata_clean = testdata_clean + .strip_prefix("//?/") + .or_else(|| testdata_clean.strip_prefix("/")) + .unwrap_or(&testdata_clean); // Use displayable() rather than explain().collect() to avoid table formatting issues. We need // to replace machine-specific paths with variable lengths, which breaks table alignment and @@ -4358,6 +4372,7 @@ async fn unnest_column_nulls() -> Result<()> { let options = UnnestOptions::new().with_preserve_nulls(false); let results = df + .clone() .unnest_columns_with_options(&["list"], options)? .collect() .await?; @@ -4374,6 +4389,156 @@ async fn unnest_column_nulls() -> Result<()> { " ); + // Outer-unnest semantics: NULL and empty lists both produce a single + // output row containing NULL. + let options = UnnestOptions::new() + .with_null_handling(datafusion_common::NullHandling::PreserveAndExpandEmpty); + let results = df + .unnest_columns_with_options(&["list"], options)? + .collect() + .await?; + assert_snapshot!( + batches_to_string(&results), + @r" + +------+----+ + | list | id | + +------+----+ + | 1 | A | + | 2 | A | + | | B | + | | C | + | 3 | D | + +------+----+ + " + ); + + Ok(()) +} + +/// Outer-unnest on a list-of-struct column. Verifies that +/// (a) struct elements unnest into flattened sub-columns and +/// (b) NULL and empty lists both still produce a single output row whose +/// struct sub-columns are all NULL. +#[tokio::test] +async fn unnest_outer_list_of_struct() -> Result<()> { + use arrow::array::{Int32Array, StructArray}; + + // Per-row sub-list lengths: 2, 1, 0 (empty), 0 (null) + let names = StringArray::from(vec!["alice", "bob", "carol"]); + let ages = Int32Array::from(vec![30, 40, 50]); + let struct_values = StructArray::from(vec![ + ( + Arc::new(Field::new("name", DataType::Utf8, true)), + Arc::new(names) as ArrayRef, + ), + ( + Arc::new(Field::new("age", DataType::Int32, true)), + Arc::new(ages) as ArrayRef, + ), + ]); + let struct_field = + Arc::new(Field::new("item", struct_values.data_type().clone(), true)); + let offsets = arrow::buffer::OffsetBuffer::::from_lengths([2, 1, 0, 0]); + let validity = arrow::buffer::NullBuffer::from(vec![true, true, true, false]); + let people = ListArray::new( + struct_field, + offsets, + Arc::new(struct_values), + Some(validity), + ); + let group = Int32Array::from(vec![1, 2, 3, 4]); + + let batch = RecordBatch::try_from_iter(vec![ + ("people", Arc::new(people) as ArrayRef), + ("group", Arc::new(group) as ArrayRef), + ])?; + + let ctx = SessionContext::new(); + ctx.register_batch("teams", batch)?; + let df = ctx.table("teams").await?; + + let options = UnnestOptions::new() + .with_null_handling(datafusion_common::NullHandling::PreserveAndExpandEmpty); + let results = df + // Unnest the list, then expand the resulting struct rows into columns. + .unnest_columns_with_options(&["people"], options.clone())? + .unnest_columns_with_options(&["people"], options)? + .collect() + .await?; + assert_snapshot!( + batches_to_string(&results), + @r" + +-------------+------------+-------+ + | people.name | people.age | group | + +-------------+------------+-------+ + | alice | 30 | 1 | + | bob | 40 | 1 | + | carol | 50 | 2 | + | | | 3 | + | | | 4 | + +-------------+------------+-------+ + " + ); + + Ok(()) +} + +/// Outer-unnest applied to a `FixedSizeList` column. For fixed-size lists, +/// every non-null row has the fixed length, so "empty" never occurs — +/// `PreserveAndExpandEmpty` should behave identically to `Preserve` here. +/// The test pins that equivalence so we notice if it ever diverges. +#[tokio::test] +async fn unnest_outer_fixed_size_list() -> Result<()> { + let batch = get_fixed_list_batch()?; + let ctx = SessionContext::new(); + ctx.register_batch("shapes", batch)?; + let df = ctx.table("shapes").await?; + + let preserve_results = df + .clone() + .unnest_columns_with_options( + &["tags"], + UnnestOptions::new().with_preserve_nulls(true), + )? + .collect() + .await?; + let outer_results = df + .unnest_columns_with_options( + &["tags"], + UnnestOptions::new().with_null_handling( + datafusion_common::NullHandling::PreserveAndExpandEmpty, + ), + )? + .collect() + .await?; + assert_eq!( + batches_to_sort_string(&preserve_results), + batches_to_sort_string(&outer_results), + "FixedSizeList has no empty case, so PreserveAndExpandEmpty must \ + match Preserve exactly" + ); + + // And the snapshot itself, to make the expected shape explicit. + assert_snapshot!( + batches_to_sort_string(&outer_results), + @r" + +----------+-------+ + | shape_id | tags | + +----------+-------+ + | 1 | | + | 2 | tag21 | + | 2 | tag22 | + | 3 | tag31 | + | 3 | tag32 | + | 4 | | + | 5 | tag51 | + | 5 | tag52 | + | 6 | tag61 | + | 6 | tag62 | + +----------+-------+ + " + ); + Ok(()) } @@ -6305,7 +6470,10 @@ async fn test_alias_nested() -> Result<()> { let select2 = df.select(vec![col("alias1.a")]); assert_snapshot!( select2.unwrap_err().strip_backtrace(), - @"Schema error: No field named alias1.a. Valid fields are alias2.a, alias2.b, alias2.one." + @r#" +Schema error: No field named alias1.a. Did you mean 'alias2.a'? +Valid fields are alias2.a, alias2.b, alias2.one. +"# ); Ok(()) } @@ -6459,11 +6627,8 @@ async fn test_fill_null() -> Result<()> { // Use fill_null to replace nulls on each column. let df_filled = df - .fill_null(ScalarValue::Int32(Some(0)), vec!["a".to_string()])? - .fill_null( - ScalarValue::Utf8(Some("default".to_string())), - vec!["b".to_string()], - )?; + .fill_null(&ScalarValue::Int32(Some(0)), &["a"])? + .fill_null(&ScalarValue::Utf8(Some("default".to_string())), &["b"])?; let results = df_filled.collect().await?; assert_snapshot!( @@ -6489,8 +6654,7 @@ async fn test_fill_null_all_columns() -> Result<()> { // Use fill_null to replace nulls on all columns. // Only column "b" will be replaced since ScalarValue::Utf8(Some("default".to_string())) // can be cast to Utf8. - let df_filled = - df.fill_null(ScalarValue::Utf8(Some("default".to_string())), vec![])?; + let df_filled = df.fill_null(&ScalarValue::Utf8(Some("default".to_string())), &[])?; let results = df_filled.clone().collect().await?; @@ -6508,7 +6672,7 @@ async fn test_fill_null_all_columns() -> Result<()> { ); // Fill column "a" null values with a value that cannot be cast to Int32. - let df_filled = df_filled.fill_null(ScalarValue::Int32(Some(0)), vec![])?; + let df_filled = df_filled.fill_null(&ScalarValue::Int32(Some(0)), &[])?; let results = df_filled.collect().await?; assert_snapshot!( @@ -6526,6 +6690,173 @@ async fn test_fill_null_all_columns() -> Result<()> { Ok(()) } +async fn create_nan_table() -> Result { + // create a DataFrame with a NaN value in a float column "a" and a + // non-float column "b" that must stay untouched by fill_nan. + // "+-----+---+", + // "| a | b |", + // "+-----+---+", + // "| 1.0 | 1 |", + // "| NaN | 2 |", + // "| 3.0 | 3 |", + // "+-----+---+", + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, true), + Field::new("b", DataType::Int32, true), + ])); + let a_values = Float64Array::from(vec![Some(1.0), Some(f64::NAN), Some(3.0)]); + let b_values = Int32Array::from(vec![Some(1), Some(2), Some(3)]); + let batch = RecordBatch::try_new( + schema.clone(), + vec![Arc::new(a_values), Arc::new(b_values)], + )?; + + let ctx = SessionContext::new(); + let table = MemTable::try_new(schema.clone(), vec![vec![batch]])?; + ctx.register_table("t_nan", Arc::new(table))?; + let df = ctx.table("t_nan").await?; + Ok(df) +} + +#[tokio::test] +async fn test_fill_nan() -> Result<()> { + let df = create_nan_table().await?; + + // Fill NaNs in the float column "a" with 0.0. + let df_filled = df.fill_nan(&ScalarValue::Float64(Some(0.0)), &["a"])?; + + let results = df_filled.collect().await?; + assert_snapshot!( + batches_to_sort_string(&results), + @r" + +-----+---+ + | a | b | + +-----+---+ + | 0.0 | 2 | + | 1.0 | 1 | + | 3.0 | 3 | + +-----+---+ + " + ); + + Ok(()) +} + +#[tokio::test] +async fn test_fill_nan_all_columns() -> Result<()> { + let df = create_nan_table().await?; + + // Fill NaNs across all columns. Only the float column "a" is affected; + // the non-float column "b" is left unchanged since NaN only exists for + // floating-point types. + let df_filled = df.fill_nan(&ScalarValue::Float64(Some(0.0)), &[])?; + + let results = df_filled.collect().await?; + assert_snapshot!( + batches_to_sort_string(&results), + @r" + +-----+---+ + | a | b | + +-----+---+ + | 0.0 | 2 | + | 1.0 | 1 | + | 3.0 | 3 | + +-----+---+ + " + ); + Ok(()) +} + +#[tokio::test] +async fn test_fill_nan_non_float_column() -> Result<()> { + let df = create_nan_table().await?; + + // Explicitly naming a non-float column is a no-op, not an error: NaN does + // not exist for Int32, so column "b" (and the un-targeted "a") are unchanged. + let df_filled = df.fill_nan(&ScalarValue::Float64(Some(0.0)), &["b"])?; + + let results = df_filled.collect().await?; + assert_snapshot!( + batches_to_sort_string(&results), + @r" + +-----+---+ + | a | b | + +-----+---+ + | 1.0 | 1 | + | 3.0 | 3 | + | NaN | 2 | + +-----+---+ + " + ); + + Ok(()) +} + +#[tokio::test] +async fn test_fill_nan_unknown_column() -> Result<()> { + let df = create_nan_table().await?; + + // A column name that is not in the schema is propagated as an error. + let err = df + .fill_nan(&ScalarValue::Float64(Some(0.0)), &["does_not_exist"]) + .unwrap_err(); + + assert_snapshot!(err.strip_backtrace(), @"Error during planning: Column 'does_not_exist' not found"); + + Ok(()) +} + +#[tokio::test] +async fn test_fill_nan_casts_fill_value() -> Result<()> { + let df = create_nan_table().await?; + + // Int32(0) is not the column's type (Float64) but can be cast to it, so the + // NaN is replaced with 0.0. Exercises the cross-type cast path — the other + // positive tests pass a Float64 value, which skips the actual cast. + let df_filled = df.fill_nan(&ScalarValue::Int32(Some(0)), &["a"])?; + + let results = df_filled.collect().await?; + assert_snapshot!( + batches_to_sort_string(&results), + @r" + +-----+---+ + | a | b | + +-----+---+ + | 0.0 | 2 | + | 1.0 | 1 | + | 3.0 | 3 | + +-----+---+ + " + ); + + Ok(()) +} + +#[tokio::test] +async fn test_fill_nan_uncastable_value() -> Result<()> { + let df = create_nan_table().await?; + + // The float column "a" is targeted, but "abc" cannot be cast to Float64, so + // the fill is skipped and column "a" keeps its original NaN value. + let df_filled = df.fill_nan(&ScalarValue::Utf8(Some("abc".to_string())), &["a"])?; + + let results = df_filled.collect().await?; + assert_snapshot!( + batches_to_sort_string(&results), + @r" + +-----+---+ + | a | b | + +-----+---+ + | 1.0 | 1 | + | 3.0 | 3 | + | NaN | 2 | + +-----+---+ + " + ); + + Ok(()) +} + #[tokio::test] async fn test_insert_into_casting_support() -> Result<()> { // Testing case1: @@ -6611,24 +6942,80 @@ async fn test_insert_into_casting_support() -> Result<()> { #[tokio::test] async fn test_dataframe_from_columns() -> Result<()> { - let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); - let b: ArrayRef = Arc::new(BooleanArray::from(vec![true, true, false])); - let c: ArrayRef = Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); - let df = DataFrame::from_columns(vec![("a", a), ("b", b), ("c", c)])?; + let bools: ArrayRef = Arc::new(BooleanArray::from(vec![true, false, true])); + let i8s: ArrayRef = Arc::new(Int8Array::from(vec![-1, 0, 1])); + let i16s: ArrayRef = Arc::new(Int16Array::from(vec![-1, 0, 1])); + let i32s: ArrayRef = Arc::new(Int32Array::from(vec![-1, 0, 1])); + let i64s: ArrayRef = Arc::new(Int64Array::from(vec![-1, 0, 1])); + + let u8s: ArrayRef = Arc::new(UInt8Array::from(vec![0, 1, 2])); + let u16s: ArrayRef = Arc::new(UInt16Array::from(vec![0, 1, 2])); + let u32s: ArrayRef = Arc::new(UInt32Array::from(vec![0, 1, 2])); + let u64s: ArrayRef = Arc::new(UInt64Array::from(vec![0, 1, 2])); + + let f16s: ArrayRef = Arc::new(Float16Array::from(vec![ + half::f16::from_f64(1.0), + half::f16::from_f64(2.0), + half::f16::from_f64(3.0), + ])); + let f32s: ArrayRef = Arc::new(Float32Array::from(vec![1.0, 2.0, 3.0])); + let f64s: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])); + + let strings: ArrayRef = + Arc::new(StringArray::from(vec![Some("foo"), Some("bar"), None])); + + let df = DataFrame::from_columns(vec![ + ("bool", bools), + ("i8", i8s), + ("i16", i16s), + ("i32", i32s), + ("i64", i64s), + ("u8", u8s), + ("u16", u16s), + ("u32", u32s), + ("u64", u64s), + ("f16", f16s), + ("f32", f32s), + ("f64", f64s), + ("str", strings), + ])?; - assert_eq!(df.schema().fields().len(), 3); + assert_eq!(df.schema().fields().len(), 13); assert_eq!(df.clone().count().await?, 3); - let rows = df.sort(vec![col("a").sort(true, true)])?; + let expected_types = [ + ("bool", DataType::Boolean), + ("i8", DataType::Int8), + ("i16", DataType::Int16), + ("i32", DataType::Int32), + ("i64", DataType::Int64), + ("u8", DataType::UInt8), + ("u16", DataType::UInt16), + ("u32", DataType::UInt32), + ("u64", DataType::UInt64), + ("f16", DataType::Float16), + ("f32", DataType::Float32), + ("f64", DataType::Float64), + ("str", DataType::Utf8), + ]; + + let schema = df.schema(); + + for (name, data_type) in expected_types { + assert_eq!(schema.field_with_name(None, name)?.data_type(), &data_type); + } + + let rows = df.sort(vec![col("i32").sort(true, true)])?; + assert_batches_eq!( &[ - "+---+-------+-----+", - "| a | b | c |", - "+---+-------+-----+", - "| 1 | true | foo |", - "| 2 | true | bar |", - "| 3 | false | |", - "+---+-------+-----+", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| bool | i8 | i16 | i32 | i64 | u8 | u16 | u32 | u64 | f16 | f32 | f64 | str |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| true | -1 | -1 | -1 | -1 | 0 | 0 | 0 | 0 | 1 | 1.0 | 1.0 | foo |", + "| false | 0 | 0 | 0 | 0 | 1 | 1 | 1 | 1 | 2 | 2.0 | 2.0 | bar |", + "| true | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 3 | 3.0 | 3.0 | |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", ], &rows.collect().await? ); @@ -6638,25 +7025,86 @@ async fn test_dataframe_from_columns() -> Result<()> { #[tokio::test] async fn test_dataframe_macro() -> Result<()> { + let bools = [true, false, true]; + let i8s = [-1_i8, 0, 1]; + let i16s = [-1_i16, 0, 1]; + let i32s = [-1_i32, 0, 1]; + let i64s = [-1_i64, 0, 1]; + + let u8s = [0_u8, 1, 2]; + let u16s = [0_u16, 1, 2]; + let u32s = [0_u32, 1, 2]; + let u64s = [0_u64, 1, 2]; + + let f16s = [ + half::f16::from_f64(1.0), + half::f16::from_f64(2.0), + half::f16::from_f64(3.0), + ]; + let f32s = [1.0_f32, 2.0, 3.0]; + let f64s = [1.0_f64, 2.0, 3.0]; + + let strings = ["foo", "bar", "baz"]; + let df = dataframe!( - "a" => [1, 2, 3], - "b" => [true, true, false], - "c" => [Some("foo"), Some("bar"), None] + // Vec + "bool" => bools.to_vec(), + "i8" => i8s.to_vec(), + "i16" => i16s.to_vec(), + "i32" => i32s.to_vec(), + + // Vec> + "i64" => vec![Some(i64s[0]), None, Some(i64s[2])], + "u8" => vec![Some(u8s[0]), None, Some(u8s[2])], + "u16" => vec![Some(u16s[0]), None, Some(u16s[2])], + + // &[T] + "u32" => &u32s, + "u64" => &u64s, + "f16" => &f16s, + + // &[Option] + "f32" => &[Some(f32s[0]), None, Some(f32s[2])], + "f64" => &[Some(f64s[0]), None, Some(f64s[2])], + "str" => &[Some(strings[0]), None, Some(strings[2])], )?; - assert_eq!(df.schema().fields().len(), 3); + assert_eq!(df.schema().fields().len(), 13); assert_eq!(df.clone().count().await?, 3); - let rows = df.sort(vec![col("a").sort(true, true)])?; + let expected_types = [ + ("bool", DataType::Boolean), + ("i8", DataType::Int8), + ("i16", DataType::Int16), + ("i32", DataType::Int32), + ("i64", DataType::Int64), + ("u8", DataType::UInt8), + ("u16", DataType::UInt16), + ("u32", DataType::UInt32), + ("u64", DataType::UInt64), + ("f16", DataType::Float16), + ("f32", DataType::Float32), + ("f64", DataType::Float64), + ("str", DataType::Utf8), + ]; + + let schema = df.schema(); + + for (name, data_type) in expected_types { + assert_eq!(schema.field_with_name(None, name)?.data_type(), &data_type); + } + + let rows = df.sort(vec![col("i32").sort(true, true)])?; + assert_batches_eq!( &[ - "+---+-------+-----+", - "| a | b | c |", - "+---+-------+-----+", - "| 1 | true | foo |", - "| 2 | true | bar |", - "| 3 | false | |", - "+---+-------+-----+", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| bool | i8 | i16 | i32 | i64 | u8 | u16 | u32 | u64 | f16 | f32 | f64 | str |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", + "| true | -1 | -1 | -1 | -1 | 0 | 0 | 0 | 0 | 1 | 1.0 | 1.0 | foo |", + "| false | 0 | 0 | 0 | | | | 1 | 1 | 2 | | | |", + "| true | 1 | 1 | 1 | 1 | 2 | 2 | 2 | 2 | 3 | 3.0 | 3.0 | baz |", + "+-------+----+-----+-----+-----+----+-----+-----+-----+-----+-----+-----+-----+", ], &rows.collect().await? ); @@ -6745,7 +7193,8 @@ async fn test_copy_to_preserves_order() -> Result<()> { DataSinkExec: sink=CsvSink(file_groups=[]) SortExec: expr=[column1@0 DESC], preserve_partitioning=[false] DataSourceExec: partitions=1, partition_sizes=[1] - DataSourceExec: partitions=1, partition_sizes=[1] + ProjectionExec: expr=[CAST(column1@0 AS UInt64) as count] + DataSourceExec: partitions=1, partition_sizes=[1] " ); Ok(()) @@ -7062,3 +7511,45 @@ async fn test_grouping_with_alias() -> Result<()> { Ok(()) } + +#[tokio::test] +async fn test_unresolved_lambda_variable() -> Result<()> { + let plan = table_with_mixed_lists() + .await? + .with_column( + "c", + array_transform( + make_array(vec![col("list")]), + lambda( + ["x"], + array_filter( + lambda_var("x"), + lambda(["y"], lambda_var("y").gt_eq(lit(2))), + ), + ), + ), + )? + .select_columns(&["list", "c"])? + .into_unoptimized_plan() + .resolve_lambda_variables()? + .data; + + let session = SessionContext::new(); + let exec = session.state().create_physical_plan(&plan).await?; + let context = session.task_ctx(); + let results = collect(exec, context).await?; + + let expected = [ + "+-----------+----------+", + "| list | c |", + "+-----------+----------+", + "| [1, 2, 3] | [[2, 3]] |", + "| | [] |", + "| [] | [[]] |", + "| | [] |", + "+-----------+----------+", + ]; + assert_batches_eq!(expected, &results); + + Ok(()) +} diff --git a/datafusion/core/tests/datasource/object_store_access.rs b/datafusion/core/tests/datasource/object_store_access.rs index 83b84f6f9284e..2503de862e06a 100644 --- a/datafusion/core/tests/datasource/object_store_access.rs +++ b/datafusion/core/tests/datasource/object_store_access.rs @@ -231,23 +231,13 @@ async fn query_multi_csv_file() { ); } -/// Test that a CSV file split into byte ranges via repartitioning exercises -/// range-based object store access. +/// Test that a CSV file split into byte ranges via repartitioning produces +/// exactly one GET request per byte range — no extra requests for boundary seeking. /// /// With a single file and `target_partitions=3`, the repartitioner produces -/// exactly 3 ranges. For each range, `calculate_range` calls -/// `find_first_newline` via a GET for every non-file boundary it touches -/// (the start boundary if `start > 0`, the end boundary if `end < file_size`), -/// plus one GET for the actual data — so 2 GETs for the first range (end scan -/// + data), 3 for the middle range (start scan + end scan + data), and 2 for -/// the last range (start scan + data) = 7 data GETs total. Additionally, -/// adjacent ranges share a boundary position, so each shared boundary is scanned -/// twice — once as the left range's end and again as the right range's start — -/// producing the duplicate GETs visible in the snapshot. Add the 1 HEAD for -/// file-size metadata = **8 total**. -/// -/// This differs from the JSON reader which uses [`AlignedBoundaryStream`] and -/// needs only 1 GET per range. +/// exactly 3 ranges. Each range is served by a single [`AlignedBoundaryStream`] +/// which issues exactly one bounded `get_opts` call, so there are 3 data GETs +/// plus 1 HEAD (to determine file size) = **4 total**. /// /// This test documents the current request pattern to catch regressions. #[tokio::test] @@ -275,15 +265,11 @@ async fn query_csv_file_with_byte_range_partitions() { +---------+-------+-------+ ------- Object Store Request Summary ------- RequestCountingObjectStore() - Total Requests: 8 + Total Requests: 4 - GET (opts) path=csv_range_table.csv head=true + - GET (opts) path=csv_range_table.csv range=0-129 - GET (opts) path=csv_range_table.csv range=42-129 - - GET (opts) path=csv_range_table.csv range=0-49 - - GET (opts) path=csv_range_table.csv range=42-129 - - GET (opts) path=csv_range_table.csv range=85-129 - - GET (opts) path=csv_range_table.csv range=49-89 - GET (opts) path=csv_range_table.csv range=85-129 - - GET (opts) path=csv_range_table.csv range=89-129 " ); } @@ -904,7 +890,7 @@ async fn query_single_parquet_file_with_single_predicate() { RequestCountingObjectStore() Total Requests: 2 - GET (opts) path=parquet_table.parquet head=true - - GET (ranges) path=parquet_table.parquet ranges=1064-1481,1481-1594,1594-2011,2011-2124 + - GET (ranges) path=parquet_table.parquet ranges=1064-1594,1594-2124 " ); } @@ -928,8 +914,8 @@ async fn query_single_parquet_file_multi_row_groups_multiple_predicates() { RequestCountingObjectStore() Total Requests: 3 - GET (opts) path=parquet_table.parquet head=true - - GET (ranges) path=parquet_table.parquet ranges=4-421,421-534,534-951,951-1064 - - GET (ranges) path=parquet_table.parquet ranges=1064-1481,1481-1594,1594-2011,2011-2124 + - GET (ranges) path=parquet_table.parquet ranges=4-534,534-1064 + - GET (ranges) path=parquet_table.parquet ranges=1064-1594,1594-2124 " ); } diff --git a/datafusion/core/tests/execution/datasource_split.rs b/datafusion/core/tests/execution/datasource_split.rs index 370249cd8044e..171e8736496a3 100644 --- a/datafusion/core/tests/execution/datasource_split.rs +++ b/datafusion/core/tests/execution/datasource_split.rs @@ -61,6 +61,7 @@ async fn datasource_splits_large_batches() -> datafusion_common::Result<()> { .options() .execution .batch_size + .get() ); let total: usize = batches.iter().map(|b| b.num_rows()).sum(); assert_eq!(total, batch_size); @@ -70,7 +71,7 @@ async fn datasource_splits_large_batches() -> datafusion_common::Result<()> { #[tokio::test] async fn datasource_exact_batch_size_no_split() -> datafusion_common::Result<()> { let session_config = datafusion_execution::config::SessionConfig::new(); - let configured_batch_size = session_config.options().execution.batch_size; + let configured_batch_size = session_config.options().execution.batch_size.get(); let batches = create_and_collect_batches(configured_batch_size).await?; diff --git a/datafusion/core/tests/expr_api/simplification.rs b/datafusion/core/tests/expr_api/simplification.rs index 245aba66849ce..e9a975239a481 100644 --- a/datafusion/core/tests/expr_api/simplification.rs +++ b/datafusion/core/tests/expr_api/simplification.rs @@ -639,22 +639,29 @@ fn test_simplify_power() { // Power(c3, 0) ===> 1 { let expr = power(col("c3_non_null"), lit(0)); - let expected = lit(1i64); + let expected = lit(1.0f64); test_simplify(expr, expected) } - // Power(c3, 1) ===> c3 + // Power(c3, 1) ===> cast(c3 AS Float64) { let expr = power(col("c3_non_null"), lit(1)); - let expected = col("c3_non_null"); + let expected = + Expr::Cast(Cast::new(Box::new(col("c3_non_null")), DataType::Float64)); test_simplify(expr, expected) } - // Power(c3, Log(c3, c4)) ===> c4 + // Power(c3, Log(c3, c4)) ===> cast(c4 AS Float64) + // The simplifier rewrites `power(b, log(b, x))` to `x`, but the + // rewritten expression must keep the same type as the original + // `power` call. `power` returns Float64, so the UInt32 c4 has to be cast + // to Float64 to preserve the output schema the optimizer already + // committed to. { let expr = power( col("c3_non_null"), log(col("c3_non_null"), col("c4_non_null")), ); - let expected = col("c4_non_null"); + let expected = + Expr::Cast(Cast::new(Box::new(col("c4_non_null")), DataType::Float64)); test_simplify(expr, expected) } // Power(c3, c4) ===> Power(c3, c4) diff --git a/datafusion/core/tests/extension_types/pretty_printing.rs b/datafusion/core/tests/extension_types/pretty_printing.rs index c0796887b8b6e..f097b5bec97fc 100644 --- a/datafusion/core/tests/extension_types/pretty_printing.rs +++ b/datafusion/core/tests/extension_types/pretty_printing.rs @@ -40,10 +40,16 @@ async fn create_test_table() -> Result { // define data. let batch = RecordBatch::try_new( schema, - vec![Arc::new(FixedSizeBinaryArray::from(vec![ - &[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], - &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 5, 6], - ]))], + vec![Arc::new( + FixedSizeBinaryArray::try_from_iter( + vec![ + &[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], + &[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 5, 6], + ] + .into_iter(), + ) + .unwrap(), + )], )?; let state = SessionStateBuilder::default() diff --git a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs index 4726e7c4aca5c..f9e7f2e10e789 100644 --- a/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs @@ -350,7 +350,12 @@ async fn run_aggregate_test(input1: Vec, group_by_columns: Vec<&str schema.clone(), ) .unwrap(), - ) as Arc; + ); + assert_ne!( + aggregate_exec_running.input_order_mode(), + &InputOrderMode::Linear, + "running aggregate should observe ordered input for group_by: {group_by:?}" + ); let aggregate_exec_usual = Arc::new( AggregateExec::try_new( @@ -362,7 +367,7 @@ async fn run_aggregate_test(input1: Vec, group_by_columns: Vec<&str schema.clone(), ) .unwrap(), - ) as Arc; + ); let task_ctx = ctx.task_ctx(); let collected_usual = collect(aggregate_exec_usual.clone(), task_ctx.clone()) diff --git a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs index fe31098622c58..3579c6af844bb 100644 --- a/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs +++ b/datafusion/core/tests/fuzz_cases/aggregation_fuzzer/context_generator.rs @@ -41,6 +41,7 @@ use crate::fuzz_cases::aggregation_fuzzer::data_generator::Dataset; /// - `batch_size` /// - `target_partitions` /// - `skip_partial parameters` +/// - `enable_migration_aggregate` /// - hint `sorted` or not /// - `spilling` or not (TODO, I think a special `MemoryPool` may be needed /// to support this) @@ -96,11 +97,13 @@ impl SessionContextGenerator { let batch_size = self.max_batch_size; let target_partitions = 1; let skip_partial_params = SkipPartialParams::ensure_not_trigger(); + let enable_migration_aggregate = false; let builder = GeneratedSessionContextBuilder { batch_size, target_partitions, skip_partial_params, + enable_migration_aggregate, sort_hint: false, table_name: self.table_name.clone(), table_provider: Arc::new(provider), @@ -120,6 +123,7 @@ impl SessionContextGenerator { // - `batch_size`, from range: [1, `total_rows_num`] // - `target_partitions`, from range: [1, cpu_num] // - `skip_partial`, trigger or not trigger currently for simplicity + // - `enable_migration_aggregate`, true or false // - `sorted`, if found a sorted dataset, will or will not push down this information // - `spilling`(TODO) let batch_size = rng.random_range(1..=self.max_batch_size); @@ -131,6 +135,8 @@ impl SessionContextGenerator { let skip_partial_params = self.candidate_skip_partial_params[skip_partial_params_idx]; + let enable_migration_aggregate = rng.random_bool(0.5); + let (provider, sort_hint) = if rng.random_bool(0.5) && !self.dataset.sort_keys.is_empty() { // Sort keys exist and random to push down @@ -150,6 +156,7 @@ impl SessionContextGenerator { target_partitions, sort_hint, skip_partial_params, + enable_migration_aggregate, table_name: self.table_name.clone(), table_provider: Arc::new(provider), }; @@ -173,6 +180,7 @@ struct GeneratedSessionContextBuilder { target_partitions: usize, sort_hint: bool, skip_partial_params: SkipPartialParams, + enable_migration_aggregate: bool, table_name: String, table_provider: Arc, } @@ -197,6 +205,10 @@ impl GeneratedSessionContextBuilder { "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", &ScalarValue::Float64(Some(self.skip_partial_params.ratio_threshold)), ); + session_config = session_config.set_bool( + "datafusion.execution.enable_migration_aggregate", + self.enable_migration_aggregate, + ); let ctx = SessionContext::new_with_config(session_config); ctx.register_table(self.table_name, self.table_provider)?; @@ -206,6 +218,7 @@ impl GeneratedSessionContextBuilder { target_partitions: self.target_partitions, sort_hint: self.sort_hint, skip_partial_params: self.skip_partial_params, + enable_migration_aggregate: self.enable_migration_aggregate, }; Ok(SessionContextWithParams { ctx, params }) @@ -220,6 +233,7 @@ pub struct SessionContextParams { target_partitions: usize, sort_hint: bool, skip_partial_params: SkipPartialParams, + enable_migration_aggregate: bool, } /// Partial skipping parameters diff --git a/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs b/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs index a57095066ee12..60b09976355e9 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/ordering.rs @@ -16,9 +16,9 @@ // under the License. use crate::fuzz_cases::equivalence::utils::{ - TestScalarUDF, create_random_schema, create_test_params, create_test_schema_2, - generate_table_for_eq_properties, generate_table_for_orderings, - is_table_same_after_sort, + TestScalarUDF, contains_overflowable_arithmetic, create_random_schema, + create_test_params, create_test_schema_2, generate_table_for_eq_properties, + generate_table_for_orderings, is_table_same_after_sort, }; use arrow::compute::SortOptions; use datafusion_common::Result; @@ -144,14 +144,27 @@ fn test_ordering_satisfy_with_equivalence_complex_random() -> Result<()> { let err_msg = format!( "Error in test case requirement:{ordering:?}, expected: {expected:?}, eq_properties: {eq_properties}", ); - // Check whether ordering_satisfy API result and - // experimental result matches. - - assert_eq!( - eq_properties.ordering_satisfy(ordering)?, - (expected | false), - "{err_msg}" + // A rejection turns inconclusive only from the first `+`/`-` + // key onwards, since possible overflow makes an ordering + // underivable even when the sample happens to be sorted. A + // table sorted by the full ordering is sorted by every prefix + // of it, so a rejected arithmetic-free prefix still proves + // the rejection is genuine. + let conclusive_prefix = LexOrdering::new( + ordering + .iter() + .take_while(|sort_expr| { + !contains_overflowable_arithmetic(&sort_expr.expr) + }) + .cloned(), ); + if eq_properties.ordering_satisfy(ordering)? { + assert!(expected, "{err_msg}"); + } else if let Some(prefix) = conclusive_prefix + && !eq_properties.ordering_satisfy(prefix)? + { + assert!(!expected, "{err_msg}"); + } } } } diff --git a/datafusion/core/tests/fuzz_cases/equivalence/projection.rs b/datafusion/core/tests/fuzz_cases/equivalence/projection.rs index 2f67e211ce915..9593e1cf11565 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/projection.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/projection.rs @@ -16,8 +16,8 @@ // under the License. use crate::fuzz_cases::equivalence::utils::{ - TestScalarUDF, apply_projection, create_random_schema, - generate_table_for_eq_properties, is_table_same_after_sort, + TestScalarUDF, apply_projection, contains_overflowable_arithmetic, + create_random_schema, generate_table_for_eq_properties, is_table_same_after_sort, }; use arrow::compute::SortOptions; use datafusion_common::Result; @@ -179,13 +179,29 @@ fn ordering_satisfy_after_projection_random() -> Result<()> { let err_msg = format!( "Error in test case requirement:{ordering:?}, expected: {expected:?}, eq_properties: {eq_properties}, projected_eq: {projected_eq}, projection_mapping: {projection_mapping:?}" ); - // Check whether ordering_satisfy API result and - // experimental result matches. - assert_eq!( - projected_eq.ordering_satisfy(ordering)?, - expected, - "{err_msg}" + // Same reasoning as in `ordering.rs`: only keys from + // the first `+`/`-` source onwards are inconclusive, + // so assert on the longest prefix without one. + let conclusive_prefix = LexOrdering::new( + ordering + .iter() + .take_while(|sort_expr| { + !projection_mapping.iter().any(|(source, targets)| { + targets + .iter() + .any(|(target, _)| target.eq(&sort_expr.expr)) + && contains_overflowable_arithmetic(source) + }) + }) + .cloned(), ); + if projected_eq.ordering_satisfy(ordering)? { + assert!(expected, "{err_msg}"); + } else if let Some(prefix) = conclusive_prefix + && !projected_eq.ordering_satisfy(prefix)? + { + assert!(!expected, "{err_msg}"); + } } } } diff --git a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs index 8350cafb215cb..ca73db3ae99ec 100644 --- a/datafusion/core/tests/fuzz_cases/equivalence/utils.rs +++ b/datafusion/core/tests/fuzz_cases/equivalence/utils.rs @@ -21,11 +21,12 @@ use std::sync::Arc; use arrow::array::{ArrayRef, Float32Array, Float64Array, RecordBatch, UInt32Array}; use arrow::compute::{SortColumn, SortOptions, lexsort_to_indices, take_record_batch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion_common::tree_node::TreeNode; use datafusion_common::utils::{compare_rows, get_row_at_idx}; use datafusion_common::{Result, exec_err, internal_datafusion_err, plan_err}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, Operator, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; use datafusion_physical_expr::equivalence::{ EquivalenceClass, ProjectionMapping, convert_to_orderings, @@ -33,7 +34,7 @@ use datafusion_physical_expr::equivalence::{ use datafusion_physical_expr::{ConstExpr, EquivalenceProperties}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; -use datafusion_physical_plan::expressions::{Column, col}; +use datafusion_physical_plan::expressions::{BinaryExpr, Column, col}; use itertools::izip; use rand::prelude::*; @@ -209,6 +210,20 @@ fn add_equal_conditions_test() -> Result<()> { Ok(()) } +/// Returns `true` if `expr` contains a `+` or `-` anywhere in its tree. +/// +/// The equivalence framework conservatively discards orderings derived from +/// `+`/`-` expressions, because wrapping overflow can break them over the +/// type's full domain even when a finite batch happens to remain sorted. +pub fn contains_overflowable_arithmetic(expr: &Arc) -> bool { + expr.exists(|e| { + Ok(e.downcast_ref::().is_some_and(|binary| { + matches!(binary.op(), Operator::Plus | Operator::Minus) + })) + }) + .unwrap() +} + /// Checks if the table (RecordBatch) remains unchanged when sorted according to the provided `required_ordering`. /// /// The function works by adding a unique column of ascending integers to the original table. This column ensures diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs index fdb2934817bc5..81c7c9f83928e 100644 --- a/datafusion/core/tests/fuzz_cases/join_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/join_fuzz.rs @@ -1008,14 +1008,12 @@ impl JoinFuzzTestCase { if join_tests.contains(&HjSmj) { let err_msg_row_cnt = format!( - "HashJoinExec and SortMergeJoinExec produced different row counts, batch_size: {}", - &batch_size + "HashJoinExec and SortMergeJoinExec produced different row counts, batch_size: {batch_size}" ); assert_eq!(hj_rows, smj_rows, "{}", err_msg_row_cnt.as_str()); let err_msg_contents = format!( - "SortMergeJoinExec and HashJoinExec produced different results, batch_size: {}", - &batch_size + "SortMergeJoinExec and HashJoinExec produced different results, batch_size: {batch_size}" ); // row level compare if any of joins returns the result // the reason is different formatting when there is no rows @@ -1070,10 +1068,10 @@ impl JoinFuzzTestCase { let mut file = std::fs::File::create(&file_path).unwrap(); println!( "{}: Saving batch idx {} rows {} to parquet {}", - &out_name, + out_name, idx, batch.num_rows(), - &file_path + file_path ); let mut writer = parquet::arrow::ArrowWriter::try_new( &mut file, diff --git a/datafusion/core/tests/fuzz_cases/once_exec.rs b/datafusion/core/tests/fuzz_cases/once_exec.rs index 403e377a690e2..c1db9a110d863 100644 --- a/datafusion/core/tests/fuzz_cases/once_exec.rs +++ b/datafusion/core/tests/fuzz_cases/once_exec.rs @@ -16,13 +16,14 @@ // under the License. use arrow_schema::SchemaRef; -use datafusion_common::internal_datafusion_err; use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, internal_datafusion_err}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, }; use std::fmt::{Debug, Formatter}; use std::sync::{Arc, Mutex}; @@ -87,19 +88,30 @@ impl ExecutionPlan for OnceExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, - ) -> datafusion_common::Result> { + _: ReplaceChildrenOptions, + ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Returns a stream which yields data fn execute( &self, partition: usize, _context: Arc, - ) -> datafusion_common::Result { + ) -> Result { assert_eq!(partition, 0); let stream = self.stream.lock().unwrap().take(); @@ -109,17 +121,10 @@ impl ExecutionPlan for OnceExec { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, - ) -> datafusion_common::Result, - ) -> datafusion_common::Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + _f: &mut dyn FnMut( + &Arc, + ) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion/core/tests/fuzz_cases/pruning.rs b/datafusion/core/tests/fuzz_cases/pruning.rs index 8ce5207f91190..7624c97cf47f7 100644 --- a/datafusion/core/tests/fuzz_cases/pruning.rs +++ b/datafusion/core/tests/fuzz_cases/pruning.rs @@ -249,12 +249,7 @@ impl Utf8Test { for (idx, truncation_length) in [Some(1), Some(2), None].iter().enumerate() { // parquet files only support 32767 row groups per file, so chunk up into multiple files so we don't error if running on a large number of row groups for (rg_idx, row_groups) in row_groups.chunks(32766).enumerate() { - let buf = write_parquet_file( - *truncation_length, - Arc::clone(&schema), - row_groups.to_vec(), - ) - .await; + let buf = write_parquet_file(*truncation_length, &schema, row_groups); let filename = format!("test_fuzz_utf8_{idx}_{rg_idx}.parquet"); let size = buf.len(); let path = Path::from(filename); @@ -314,10 +309,10 @@ async fn execute_with_predicate( values } -async fn write_parquet_file( +fn write_parquet_file( truncation_length: Option, - schema: Arc, - row_groups: Vec>, + schema: &Arc, + row_groups: &[Vec], ) -> Bytes { let mut buf = BytesMut::new().writer(); let props = WriterProperties::builder() @@ -326,11 +321,11 @@ async fn write_parquet_file( let props = props.build(); { let mut writer = - ArrowWriter::try_new(&mut buf, schema.clone(), Some(props)).unwrap(); - for rg_values in row_groups.iter() { + ArrowWriter::try_new(&mut buf, Arc::clone(schema), Some(props)).unwrap(); + for rg_values in row_groups { let arr = StringArray::from_iter_values(rg_values.iter()); let batch = - RecordBatch::try_new(schema.clone(), vec![Arc::new(arr)]).unwrap(); + RecordBatch::try_new(Arc::clone(schema), vec![Arc::new(arr)]).unwrap(); writer.write(&batch).unwrap(); writer.flush().unwrap(); // finishes the current row group and starts a new one } diff --git a/datafusion/core/tests/fuzz_cases/sort_fuzz.rs b/datafusion/core/tests/fuzz_cases/sort_fuzz.rs index 0d8a066d432dd..675854ddb54b1 100644 --- a/datafusion/core/tests/fuzz_cases/sort_fuzz.rs +++ b/datafusion/core/tests/fuzz_cases/sort_fuzz.rs @@ -40,7 +40,6 @@ use test_utils::{batches_to_vec, partitions_to_sorted_vec}; const KB: usize = 1 << 10; #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] async fn test_sort_10k_mem() { for (batch_size, should_spill) in [(5, false), (20000, true), (500000, true)] { let (input, collected) = SortTest::new() @@ -58,7 +57,6 @@ async fn test_sort_10k_mem() { } #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] async fn test_sort_100k_mem() { for (batch_size, should_spill) in [(5, false), (10000, false), (20000, true), (1000000, true)] @@ -78,7 +76,6 @@ async fn test_sort_100k_mem() { } #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] async fn test_sort_strings_100k_mem() { for (batch_size, should_spill) in [(5, false), (1000, false), (10000, true), (20000, true)] @@ -116,7 +113,6 @@ async fn test_sort_strings_100k_mem() { } #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] async fn test_sort_multi_columns_100k_mem() { for (batch_size, should_spill) in [(5, false), (1000, false), (10000, true), (20000, true)] diff --git a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs index d401557e966d6..f82d0165f2fdb 100644 --- a/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs +++ b/datafusion/core/tests/fuzz_cases/spilling_fuzz_in_memory_constrained_env.rs @@ -23,6 +23,7 @@ use std::sync::Arc; use crate::fuzz_cases::aggregate_fuzz::assert_spill_count_metric; use crate::fuzz_cases::once_exec::OnceExec; use arrow::array::UInt64Array; +use arrow::row::{RowConverter, SortField}; use arrow::{array::StringArray, compute::SortOptions, record_batch::RecordBatch}; use arrow_schema::{DataType, Field, Schema}; use datafusion::common::Result; @@ -39,13 +40,25 @@ use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_functions_aggregate::array_agg::array_agg_udaf; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::{Column, col}; +use datafusion_physical_expr_common::metrics::MetricsSet; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; +use datafusion_physical_plan::metrics::MetricValue; +use datafusion_physical_plan::spill::get_record_batch_memory_size; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use futures::StreamExt; +use arrow::array::Int32Array; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_execution::memory_pool::{ + MemoryPool, PeakRecordingPool, UnboundedMemoryPool, +}; +use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; +use datafusion_physical_plan::spill::SpillManager; + #[tokio::test] async fn test_sort_with_limited_memory() -> Result<()> { let record_batch_size = 8192; @@ -69,16 +82,18 @@ async fn test_sort_with_limited_memory() -> Result<()> { // Basic test with a lot of groups that cannot all fit in memory and 1 record batch // from each spill file is too much memory - let spill_count = run_sort_test_with_limited_memory(RunTestWithLimitedMemoryArgs { + let metrics = run_sort_test_with_limited_memory(RunTestWithLimitedMemoryArgs { pool_size, task_ctx: Arc::new(task_ctx), number_of_record_batches: 100, get_size_of_record_batch_to_generate: Box::pin(move |_| record_batch_size), memory_behavior: Default::default(), + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; - let total_spill_files_size = spill_count * record_batch_size; + let total_spill_files_size = + metrics.spill_count().unwrap_or_default() * record_batch_size; assert!( total_spill_files_size > pool_size, "Total spill files size {total_spill_files_size} should be greater than pool size {pool_size}", @@ -119,6 +134,7 @@ async fn test_sort_with_limited_memory_and_different_sizes_of_record_batch() -> } }), memory_behavior: Default::default(), + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; @@ -157,6 +173,7 @@ async fn test_sort_with_limited_memory_and_different_sizes_of_record_batch_and_c } }), memory_behavior: MemoryBehavior::TakeAllMemoryAndReleaseEveryNthBatch(10), + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; @@ -195,6 +212,7 @@ async fn test_sort_with_limited_memory_and_different_sizes_of_record_batch_and_t } }), memory_behavior: MemoryBehavior::TakeAllMemoryAtTheBeginning, + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; @@ -227,9 +245,303 @@ async fn test_sort_with_limited_memory_and_large_record_batch() -> Result<()> { number_of_record_batches: 100, get_size_of_record_batch_to_generate: Box::pin(move |_| pool_size / 6), memory_behavior: Default::default(), + assert_all_output_batches_roughly_match_batch_size_conf: true, + }) + .await?; + + Ok(()) +} + +#[tokio::test] +async fn test_sort_with_limited_memory_and_oversized_record_batch() -> Result<()> { + let record_batch_size = 8192; + let pool_size = 2 * MB as usize; + let task_ctx = { + let memory_pool = Arc::new(FairSpillPool::new(pool_size)); + TaskContext::default() + .with_session_config( + SessionConfig::new() + .with_batch_size(record_batch_size) + .with_sort_spill_reservation_bytes(1), + ) + .with_runtime(Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(memory_pool) + .build()?, + )) + }; + + let number_of_record_batches = 100; + + // Each spilled run's largest batch is so big that two merge streams cannot be + // reserved at once even at the smallest read-buffer size (`2 * (2 * batch) > + // pool`), yet a single stream still fits (`2 * batch < pool`). Reducing the + // buffer size therefore cannot help, the multi-level merge has to re-spill a + // run with a smaller batch size to make progress instead of failing with + // `ResourcesExhausted`. + let metrics = run_sort_test_with_limited_memory(RunTestWithLimitedMemoryArgs { + pool_size, + task_ctx: Arc::new(task_ctx), + number_of_record_batches, + get_size_of_record_batch_to_generate: Box::pin(move |_| pool_size / 3), + memory_behavior: Default::default(), + + assert_all_output_batches_roughly_match_batch_size_conf: false, }) .await?; + let output_batches = get_output_batches_from_metrics(&metrics); + + // minimum 2 batches more + assert!( + output_batches >= number_of_record_batches + 2, + "output_batches {output_batches} should be greater than number_of_record_batches ({number_of_record_batches}) + 2" + ); + + Ok(()) +} + +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(true, false, false).await +} + +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(false, false, false).await +} + +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_multi_column() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(true, true, false).await +} + +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_multi_column() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(false, true, false).await +} + +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_round_robin_tied_values() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(true, false, true).await +} + +#[tokio::test] +async fn test_sort_preserving_merge_peak_memory_with_spilled_input_no_round_robin_tied_values() +-> Result<()> { + run_sort_preserving_merge_peak_memory_with_spilled_input(false, false, true).await +} + +/// Intended to measure the maximum number of record batches held in memory by +/// the SortPreservingMergeStream in a convoluted way by measuring the peak +/// memory reservation. Relevant for merging spilled streams, where the produced +/// record batches suffer from the following issue: +/// https://github.com/apache/arrow-rs/issues/6363 +/// +/// After an IPC roundtrip, all columns in a [`RecordBatch`] share a single +/// parent buffer. It causes the memory reservation to be inflated, but the +/// bigger issue is the increase in the peak allocated memory caused by +/// prev_cursors in SortPreservingMergeExec. The increase is caused by the fact +/// that the FieldCursor inside prev_cursors holds a reference for the entire +/// Buffer allocated for the input record batch, preventing it from being +/// dropped and thus increasing the number of concomitent input record batches +/// living during the merging phase +async fn run_sort_preserving_merge_peak_memory_with_spilled_input( + round_robin: bool, + multi_column_sort: bool, + tied_values: bool, +) -> Result<()> { + let num_batches = 10usize; + let num_rows_per_batch = 100usize; + // payload is ~100x larger than the sort key (i32 = 4 bytes, string ≈ 400 bytes) + let large_string = "x".repeat(400); + + let schema = Arc::new(Schema::new(vec![ + Field::new("sort_key", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + ])); + + // Unbounded env used only for spilling the input; the merge runs under its + // own pool below. + let spill_env = Arc::new(RuntimeEnvBuilder::new().build()?); + + let mut partition_batches: Vec> = Vec::new(); + + for stream_idx in 0..2usize { + // Each stream covers a non-overlapping key range so both are individually + // sorted: stream 0 → [0, 1000), stream 1 → [1000, 2000). When + // `tied_values` is set, every row of every batch in both streams + // instead carries the same sort key, so every comparison between the + // two streams is a tie. + let batches: Vec = (0..num_batches) + .map(|b| { + // Interleave streams: stream 0 → even slots [0,200,400,...], + // stream 1 → odd slots [100,300,500,...] so the merge + // alternates between them on every batch. + let base = ((b * 2 + stream_idx) * num_rows_per_batch) as i32; + let sort_col: Int32Array = if tied_values { + std::iter::repeat_n(0, num_rows_per_batch).collect() + } else { + (base..base + num_rows_per_batch as i32).collect() + }; + let payload_col: StringArray = + std::iter::repeat_n(large_string.as_str(), num_rows_per_batch) + .map(Some) + .collect(); + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sort_col), Arc::new(payload_col)], + ) + .unwrap() + }) + .collect(); + + // Spill to disk then read back: each RecordBatch is now IPC-backed, + // meaning all columns share a single parent buffer. As a result, + // get_buffer_memory_size() on the sort_key column returns the full + // parent-buffer capacity (≈ batch size of both columns combined) rather + // than just the key data (num_rows * 4 bytes). + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let manager = + SpillManager::new(Arc::clone(&spill_env), metrics, Arc::clone(&schema)); + let spill_file = manager + .spill_record_batch_and_finish(&batches, "stream")? + .expect("non-empty input should produce a spill file"); + + let mut stream = manager.read_spill_as_stream(spill_file, None)?; + let mut ipc_batches: Vec = Vec::new(); + while let Some(batch) = stream.next().await { + ipc_batches.push(batch?); + } + partition_batches.push(ipc_batches); + } + + let ipc_batch_size = get_record_batch_memory_size(&partition_batches[0][0]); + + // Build a 2-partition plan from the IPC-recovered batches. + let input = + MemorySourceConfig::try_new_exec(&partition_batches, Arc::clone(&schema), None)?; + + let sort_key_expr = PhysicalSortExpr { + expr: col("sort_key", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }; + // `payload` has the same value in every row, so adding it as a secondary + // sort key doesn't change the resulting order — it only forces the merge + // onto the row-oriented (`RowValues`/`RowCursorStream`) comparison path + // used whenever more than one sort expression is present. + let mut sort_exprs = vec![sort_key_expr]; + if multi_column_sort { + sort_exprs.push(PhysicalSortExpr { + expr: col("payload", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }); + } + + // When sorting by more than one column, the merge switches to the + // row-oriented `RowValues`/`RowCursorStream` path + // + // `RowCursorStream` also tracks one *shared* (not per-partition) + // reservation sized to `converter.size()` (`stream.rs`: + // `self.reservation.try_resize(self.converter.size())`) — the + // `RowConverter`'s own fixed internal state, separate from the `Rows` + // it produces per batch. + let (row_batch_size, converter_size) = if multi_column_sort { + let sort_fields = sort_exprs + .iter() + .map(|s| { + let data_type = s.expr.data_type(&schema)?; + Ok(SortField::new_with_options(data_type, s.options)) + }) + .collect::>>()?; + let converter = RowConverter::new(sort_fields)?; + let cols = sort_exprs + .iter() + .map(|s| { + s.expr + .evaluate(&partition_batches[0][0])? + .into_array(partition_batches[0][0].num_rows()) + }) + .collect::>>()?; + let rows = converter.convert_columns(&cols)?; + (rows.size(), converter.size()) + } else { + (0, 0) + }; + + let merge = Arc::new( + SortPreservingMergeExec::new(LexOrdering::new(sort_exprs).unwrap(), input) + .with_round_robin_repartition(round_robin), + ); + + // PeakRecordingPool records peak reserved bytes as a running high-water mark + // (via grow/shrink deltas), independent of any per-consumer registration + // bookkeeping - unlike TrackConsumersPool, whose tracked-consumer entry (and + // its peak) gets discarded the moment the consumer unregisters, which now + // happens mid-poll (inside the drain loop below) rather than when the + // caller eventually drops the returned stream. + let tracking_pool = Arc::new(PeakRecordingPool::new(Arc::new( + UnboundedMemoryPool::default(), + ))); + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::clone(&tracking_pool) as Arc) + .build()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().with_batch_size(num_rows_per_batch)) + .with_runtime(Arc::new(runtime)), + ); + + let mut output = merge.execute(0, task_ctx)?; + let mut total_rows = 0usize; + while let Some(batch) = output.next().await { + total_rows += batch?.num_rows(); + } + assert_eq!(total_rows, 2 * num_batches * num_rows_per_batch); + + let peak_bytes = tracking_pool.peak_reserved(); + + // in the single column case, the cursor takes up an ipc_batch_size worth of memory due to the + // IPC roundtrip issue + // for the multi-column case, we've calculated row_batch_size above + let cursor_unit = if multi_column_sort { + row_batch_size + } else { + ipc_batch_size + }; + + // BatchBuilder needs to hold 3 Record batches simultaneously to merge two + // streams (because a stream can cross a record batch boundary) + // there is also one cursor needed per stream + let mut max_peak = 3 * ipc_batch_size + 2 * cursor_unit + converter_size; + + // with round robin enabled, 2 extra cursors live in memory + // see https://github.com/apache/datafusion/issues/23604 + if round_robin { + max_peak += 2 * cursor_unit; + }; + + assert!( + peak_bytes > 0, + "peak reservation {peak_bytes} should be greater than 0" + ); + assert!( + peak_bytes <= max_peak, + "peak reservation {peak_bytes} bytes exceeds max_peak ({max_peak} bytes); \ + round_robin={round_robin}, multi_column_sort={multi_column_sort}", + ); + Ok(()) } @@ -240,6 +552,9 @@ struct RunTestWithLimitedMemoryArgs { get_size_of_record_batch_to_generate: Pin usize + Send + 'static>>, memory_behavior: MemoryBehavior, + + /// When true we would `assert_eq(the number of output_rows metric / output_batches metric == task_ctx.batch_size)` + assert_all_output_batches_roughly_match_batch_size_conf: bool, } #[derive(Default)] @@ -252,7 +567,7 @@ enum MemoryBehavior { async fn run_sort_test_with_limited_memory( mut args: RunTestWithLimitedMemoryArgs, -) -> Result { +) -> Result { let get_size_of_record_batch_to_generate = std::mem::replace( &mut args.get_size_of_record_batch_to_generate, Box::pin(move |_| unreachable!("should not be called after take")), @@ -312,7 +627,23 @@ async fn run_sort_test_with_limited_memory( let result = sort_exec.execute(0, Arc::clone(&args.task_ctx))?; - run_test(args, sort_exec, result).await + let number_of_record_batches = args.number_of_record_batches; + let assert_output_batch_size = + args.assert_all_output_batches_roughly_match_batch_size_conf; + + let metrics = run_test(args, sort_exec, result).await?; + + assert_baseline_metrics_for_non_empty_output( + &metrics, + number_of_record_batches * record_batch_size as usize, + if assert_output_batch_size { + Some(record_batch_size as usize) + } else { + None + }, + ); + + Ok(metrics) } fn grow_memory_as_much_as_possible( @@ -346,17 +677,19 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory() -> Result<() // Basic test with a lot of groups that cannot all fit in memory and 1 record batch // from each spill file is too much memory - let spill_count = + let metrics = run_test_aggregate_with_high_cardinality(RunTestWithLimitedMemoryArgs { pool_size, task_ctx: Arc::new(task_ctx), number_of_record_batches: 100, get_size_of_record_batch_to_generate: Box::pin(move |_| record_batch_size), memory_behavior: Default::default(), + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; - let total_spill_files_size = spill_count * record_batch_size; + let total_spill_files_size = + metrics.spill_count().unwrap_or_default() * record_batch_size; assert!( total_spill_files_size > pool_size, "Total spill files size {total_spill_files_size} should be greater than pool size {pool_size}", @@ -393,6 +726,7 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_ } }), memory_behavior: Default::default(), + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; @@ -427,6 +761,7 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_ } }), memory_behavior: MemoryBehavior::TakeAllMemoryAndReleaseEveryNthBatch(10), + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; @@ -461,6 +796,7 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_different_ } }), memory_behavior: MemoryBehavior::TakeAllMemoryAtTheBeginning, + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; @@ -490,6 +826,7 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_large_reco number_of_record_batches: 100, get_size_of_record_batch_to_generate: Box::pin(move |_| pool_size / 6), memory_behavior: Default::default(), + assert_all_output_batches_roughly_match_batch_size_conf: true, }) .await?; @@ -498,7 +835,7 @@ async fn test_aggregate_with_high_cardinality_with_limited_memory_and_large_reco async fn run_test_aggregate_with_high_cardinality( mut args: RunTestWithLimitedMemoryArgs, -) -> Result { +) -> Result { let get_size_of_record_batch_to_generate = std::mem::replace( &mut args.get_size_of_record_batch_to_generate, Box::pin(move |_| unreachable!("should not be called after take")), @@ -587,12 +924,13 @@ async fn run_test( args: RunTestWithLimitedMemoryArgs, plan: Arc, result_stream: SendableRecordBatchStream, -) -> Result { +) -> Result { let number_of_record_batches = args.number_of_record_batches; consume_stream_and_simulate_other_running_memory_consumers(args, result_stream) .await?; + let metrics = plan.metrics().expect("must have metrics"); let spill_count = assert_spill_count_metric(true, plan); assert!( @@ -600,7 +938,7 @@ async fn run_test( "Expected spill, but did not, number of record batches: {number_of_record_batches}", ); - Ok(spill_count) + Ok(metrics) } /// Consume the stream and change the amount of memory used while consuming it based on the [`MemoryBehavior`] provided @@ -656,3 +994,56 @@ async fn consume_stream_and_simulate_other_running_memory_consumers( Ok(()) } + +/// Assert baseline metrics are as expected or around that +/// +/// `output_batch_size` should be `None` when you expect to not get batched at the same size +/// `Some(session conf batch size)` for the rest +fn assert_baseline_metrics_for_non_empty_output( + metrics: &MetricsSet, + expected_output_rows: usize, + output_batch_size: Option, +) { + let end_time = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::EndTimestamp(end) => Some(end), + _ => None, + }) + .expect("Must have end time metric since it exists in the baseline"); + + assert_ne!(end_time.value(), None); + + assert_eq!(metrics.output_rows(), Some(expected_output_rows)); + + let output_bytes = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::OutputBytes(total) => Some(total), + _ => None, + }) + .expect("Must have output_bytes metric since it exists in the baseline"); + + assert_ne!(output_bytes.value(), 0_usize); + + let output_batches = get_output_batches_from_metrics(metrics); + + if let Some(output_batch_size) = output_batch_size { + assert_eq!( + output_batches, + expected_output_rows.div_ceil(output_batch_size) + ); + } else { + assert_ne!(output_batches, 0,); + } +} + +fn get_output_batches_from_metrics(metrics: &MetricsSet) -> usize { + metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::OutputBatches(total) => Some(total.value()), + _ => None, + }) + .expect("Must have output_batches metric since it exists in the baseline") +} diff --git a/datafusion/core/tests/helper/mod.rs b/datafusion/core/tests/helper/mod.rs new file mode 100644 index 0000000000000..809f8ca087547 --- /dev/null +++ b/datafusion/core/tests/helper/mod.rs @@ -0,0 +1,23 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared helpers for the `core_integration` test crate. +//! +//! Keep cross-cutting test utilities here when they are used by multiple test +//! modules under `core/tests`. Placing them in this submodule avoids creating +//! an additional Cargo integration test target for each helper file. +pub(crate) mod plan_metrics; diff --git a/datafusion/core/tests/helper/plan_metrics.rs b/datafusion/core/tests/helper/plan_metrics.rs new file mode 100644 index 0000000000000..12d3eaba1ad96 --- /dev/null +++ b/datafusion/core/tests/helper/plan_metrics.rs @@ -0,0 +1,54 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Helpers for aggregating execution metrics across a physical plan tree. +//! +//! `ExecutionPlan::metrics()` returns metrics for a single plan node only; it +//! does not include metrics from child operators. These helpers recursively walk +//! the plan tree so tests can assert on metrics that may move between operators +//! after optimizer rewrites, such as pushing a `SortExec` below a +//! `ProjectionExec`. + +use datafusion_physical_plan::ExecutionPlan; + +/// Returns the total number of spill events recorded by `plan` and all of its +/// descendants. +/// +/// Missing `spill_count` metrics are treated as zero. +pub fn plan_spill_count(plan: &dyn ExecutionPlan) -> usize { + let own = plan.metrics().and_then(|m| m.spill_count()).unwrap_or(0); + + own + plan + .children() + .into_iter() + .map(|child| plan_spill_count(child.as_ref())) + .sum::() +} + +/// Returns the total number of spilled bytes recorded by `plan` and all of its +/// descendants. +/// +/// Missing `spilled_bytes` metrics are treated as zero. +pub fn plan_spilled_bytes(plan: &dyn ExecutionPlan) -> usize { + let own = plan.metrics().and_then(|m| m.spilled_bytes()).unwrap_or(0); + + own + plan + .children() + .into_iter() + .map(|child| plan_spilled_bytes(child.as_ref())) + .sum::() +} diff --git a/datafusion/core/tests/macro_hygiene/mod.rs b/datafusion/core/tests/macro_hygiene/mod.rs index 9fd60cd1f06f3..144062278cc10 100644 --- a/datafusion/core/tests/macro_hygiene/mod.rs +++ b/datafusion/core/tests/macro_hygiene/mod.rs @@ -41,6 +41,10 @@ mod plan_datafusion_err { } mod record_batch { + #![expect( + deprecated, + reason = "exercising hygiene of the deprecated `datafusion_common::record_batch!` while it is still exported" + )] // NO other imports! use datafusion_common::record_batch; diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs index 32df6c5d62937..83ebb266c8257 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/mod.rs @@ -18,5 +18,6 @@ //! Validates query's actual memory usage is consistent with the specified memory //! limit. +mod smj_mem_validation; mod sort_mem_validation; mod utils; diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs new file mode 100644 index 0000000000000..3af642fffe101 --- /dev/null +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/smj_mem_validation.rs @@ -0,0 +1,105 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Memory-limit validation tests for sort-merge join queries. +//! +//! These tests run in separate processes to accurately measure memory usage. + +use datafusion::prelude::SessionConfig; + +use crate::memory_limit::memory_limit_validation::utils; + +/// Ensures the planner selected a sort-merge join. +const SMJ_OPERATOR_NAME: &str = "SortMergeJoinExec"; + +/// Configure a two-partition sort-merge join and reduce the sort reservation so +/// the join can spill under the tested memory limits. +fn smj_session_config() -> SessionConfig { + SessionConfig::new() + .with_target_partitions(2) + .with_sort_spill_reservation_bytes(1024 * 1024) + .set_bool("datafusion.optimizer.prefer_hash_join", false) +} + +/// Build a join with one large buffered key group and scalar output. +fn smj_sum_query(series_len: usize) -> String { + format!( + "SELECT sum(rr.v) FROM generate_series(0, 0) AS l(k) \ + JOIN (SELECT i % 1 AS k, i AS v FROM generate_series(1, {series_len}) AS r(i)) rr \ + ON l.k = rr.k" + ) +} + +#[test] +fn smj_with_mem_limit_1_runner() { + utils::spawn_test_process("smj_mem_validation", "smj_with_mem_limit_1"); +} + +#[test] +fn smj_with_mem_limit_2_runner() { + utils::spawn_test_process("smj_mem_validation", "smj_with_mem_limit_2"); +} + +#[test] +fn smj_no_mem_limit_runner() { + utils::spawn_test_process("smj_mem_validation", "smj_no_mem_limit"); +} + +/// Verify a 40 MB pool forces spilling within the RSS allowance. +#[tokio::test] +async fn smj_with_mem_limit_1() { + utils::validate_query_with_memory_limits_and_config( + 40_000_000 * 4, + Some(40_000_000), + &smj_sum_query(5_000_000), + &smj_sum_query(500_000), + smj_session_config(), + Some(SMJ_OPERATOR_NAME), + Some(true), + ) + .await; +} + +/// Verify a 16 MB pool forces spilling. The 5M join keys (~40 MB) stay resident +/// independently of the pool limit, so this case needs a larger RSS allowance. +#[tokio::test] +async fn smj_with_mem_limit_2() { + utils::validate_query_with_memory_limits_and_config( + 16_000_000 * 12, + Some(16_000_000), + &smj_sum_query(5_000_000), + &smj_sum_query(500_000), + smj_session_config(), + Some(SMJ_OPERATOR_NAME), + Some(true), + ) + .await; +} + +#[tokio::test] +async fn smj_no_mem_limit() { + utils::validate_query_with_memory_limits_and_config( + 40_000_000 * 5, + None, + &smj_sum_query(5_000_000), + &smj_sum_query(500_000), + smj_session_config(), + Some(SMJ_OPERATOR_NAME), + Some(false), + ) + .await; +} diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs index bf04123fff7fa..b55a3039ec9d4 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/sort_mem_validation.rs @@ -21,7 +21,6 @@ //! This file is organized as: //! - Test runners that spawn individual test processes //! - Test cases that contain the actual validation logic -use std::{process::Command, str}; use crate::memory_limit::memory_limit_validation::utils; @@ -32,67 +31,40 @@ use crate::memory_limit::memory_limit_validation::utils; #[test] fn memory_limit_validation_runner_works_runner() { - spawn_test_process("memory_limit_validation_runner_works"); + utils::spawn_test_process( + "sort_mem_validation", + "memory_limit_validation_runner_works", + ); } #[test] fn sort_no_mem_limit_runner() { - spawn_test_process("sort_no_mem_limit"); + utils::spawn_test_process("sort_mem_validation", "sort_no_mem_limit"); } #[test] fn sort_with_mem_limit_1_runner() { - spawn_test_process("sort_with_mem_limit_1"); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_1"); } #[test] fn sort_with_mem_limit_2_runner() { - spawn_test_process("sort_with_mem_limit_2"); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2"); } #[test] fn sort_with_mem_limit_3_runner() { - spawn_test_process("sort_with_mem_limit_3"); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_3"); } #[test] fn sort_with_mem_limit_2_cols_1_runner() { - spawn_test_process("sort_with_mem_limit_2_cols_1"); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2_cols_1"); } #[test] fn sort_with_mem_limit_2_cols_2_runner() { - spawn_test_process("sort_with_mem_limit_2_cols_2"); -} - -/// Helper function that executes a test in a separate process with the required -/// environment variable set. Re-invokes the current test binary directly, -/// avoiding cargo overhead and recompilation. -fn spawn_test_process(test: &str) { - let test_path = - format!("memory_limit::memory_limit_validation::sort_mem_validation::{test}"); - - let exe = std::env::current_exe().expect("Failed to get test binary path"); - - let output = Command::new(exe) - .arg(&test_path) - .arg("--exact") - .arg("--nocapture") - .env("DATAFUSION_TEST_MEM_LIMIT_VALIDATION", "1") - .output() - .expect("Failed to execute test command"); - - let stdout = str::from_utf8(&output.stdout).unwrap_or(""); - let stderr = str::from_utf8(&output.stderr).unwrap_or(""); - - assert!( - output.status.success(), - "Test '{}' failed with status: {}\nstdout:\n{}\nstderr:\n{}", - test, - output.status, - stdout, - stderr - ); + utils::spawn_test_process("sort_mem_validation", "sort_with_mem_limit_2_cols_2"); } // =========================================================================== diff --git a/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs b/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs index 2c9fae20c8606..788b8f4942ee4 100644 --- a/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs +++ b/datafusion/core/tests/memory_limit/memory_limit_validation/utils.rs @@ -16,11 +16,14 @@ // under the License. use datafusion_common_runtime::SpawnedTask; +use std::process::Command; +use std::str; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; use sysinfo::{ProcessRefreshKind, ProcessesToUpdate, System}; use tokio::time::{Duration, interval}; +use datafusion::physical_plan::{ExecutionPlan, collect, displayable}; use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_common::human_readable_size; use datafusion_execution::{memory_pool::FairSpillPool, runtime_env::RuntimeEnvBuilder}; @@ -98,6 +101,42 @@ where (result, peak_rss) } +/// Helper function that executes a test in a separate process with the required +/// environment variable set. Re-invokes the current test binary directly, +/// avoiding cargo overhead and recompilation. +pub fn spawn_test_process(module: &str, test: &str) { + let test_path = format!("memory_limit::memory_limit_validation::{module}::{test}"); + let exe = std::env::current_exe().expect("Failed to get test binary path"); + let output = Command::new(exe) + .arg(&test_path) + .arg("--exact") + .arg("--nocapture") + .env("DATAFUSION_TEST_MEM_LIMIT_VALIDATION", "1") + .output() + .expect("Failed to execute test command"); + + let stdout = str::from_utf8(&output.stdout).unwrap_or(""); + let stderr = str::from_utf8(&output.stderr).unwrap_or(""); + assert!( + output.status.success(), + "Test '{test}' failed with status: {}\nstdout:\n{stdout}\nstderr:\n{stderr}", + output.status, + ); +} + +fn operator_spill_count(plan: &dyn ExecutionPlan, operator_name: &str) -> usize { + let own = if plan.name() == operator_name { + plan.metrics().and_then(|m| m.spill_count()).unwrap_or(0) + } else { + 0 + }; + own + plan + .children() + .into_iter() + .map(|child| operator_spill_count(child.as_ref(), operator_name)) + .sum::() +} + /// Query runner that validates the memory usage of the query. /// /// Note this function is supposed to run in a separate process for accurate memory @@ -132,6 +171,30 @@ pub async fn validate_query_with_memory_limits( mem_limit_bytes: Option, query: &str, baseline_query: &str, +) { + let session_config = SessionConfig::new().with_target_partitions(4); // Make sure the configuration is the same if test is running on different machines + validate_query_with_memory_limits_and_config( + expected_mem_bytes, + mem_limit_bytes, + query, + baseline_query, + session_config, + None, + None, + ) + .await; +} + +/// Validate memory usage with a custom session configuration and optional +/// operator and spill assertions. +pub async fn validate_query_with_memory_limits_and_config( + expected_mem_bytes: i64, + mem_limit_bytes: Option, + query: &str, + baseline_query: &str, + session_config: SessionConfig, + expected_operator_name: Option<&str>, + expected_operator_spill: Option, ) { if std::env::var("DATAFUSION_TEST_MEM_LIMIT_VALIDATION").is_err() { println!("Skipping test because DATAFUSION_TEST_MEM_LIMIT_VALIDATION is not set"); @@ -151,18 +214,50 @@ pub async fn validate_query_with_memory_limits( None => runtime_builder.build_arc().unwrap(), }; - let session_config = SessionConfig::new().with_target_partitions(4); // Make sure the configuration is the same if test is running on different machines - let ctx = SessionContext::new_with_config_rt(session_config, runtime); let df = ctx.sql(query).await.unwrap(); + let physical_plan = df.create_physical_plan().await.unwrap(); + + if let Some(expected) = expected_operator_name { + let plan_display = displayable(physical_plan.as_ref()).indent(true).to_string(); + assert!( + plan_display.contains(expected), + "expected physical plan to contain `{expected}`, but got:\n{plan_display}", + ); + } + // Run a query with 10% data to estimate the constant overhead - let df_small = ctx.sql(baseline_query).await.unwrap(); + let baseline_plan = ctx + .sql(baseline_query) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + let baseline_task_ctx = ctx.task_ctx(); + let (_, baseline_max_rss) = measure_max_rss(|| async move { + collect(baseline_plan, baseline_task_ctx).await.unwrap() + }) + .await; - let (_, baseline_max_rss) = - measure_max_rss(|| async { df_small.collect().await.unwrap() }).await; + let execution_plan = Arc::clone(&physical_plan); + let execution_task_ctx = ctx.task_ctx(); + let (_, max_rss) = measure_max_rss(|| async move { + collect(execution_plan, execution_task_ctx).await.unwrap() + }) + .await; - let (_, max_rss) = measure_max_rss(|| async { df.collect().await.unwrap() }).await; + if let (Some(operator), Some(expect_spill)) = + (expected_operator_name, expected_operator_spill) + { + let spill_count = operator_spill_count(physical_plan.as_ref(), operator); + assert_eq!( + spill_count > 0, + expect_spill, + "unexpected spill_count={spill_count} for {operator}", + ); + } println!( "Memory before: {}, Memory after: {}", diff --git a/datafusion/core/tests/memory_limit/mod.rs b/datafusion/core/tests/memory_limit/mod.rs index 64861f237074e..84d7e9c4508b5 100644 --- a/datafusion/core/tests/memory_limit/mod.rs +++ b/datafusion/core/tests/memory_limit/mod.rs @@ -62,6 +62,8 @@ use async_trait::async_trait; use futures::StreamExt; use tokio::fs::File; +use crate::helper::plan_metrics::{plan_spill_count, plan_spilled_bytes}; + #[cfg(test)] #[ctor::ctor(unsafe)] fn init() { @@ -99,7 +101,8 @@ async fn group_by_row_hash() { TestCase::new() .with_query("select count(*) from t GROUP BY response_bytes") .with_expected_errors(vec![ - "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n GroupedHashAggregateStream" + "Resources exhausted: Additional allocation failed", + "for FinalHashAggregateStream[0]", ]) .with_memory_limit(2_000) .run() @@ -112,7 +115,8 @@ async fn group_by_hash() { // group by dict column .with_query("select count(*) from t GROUP BY service, host, pod, container") .with_expected_errors(vec![ - "Resources exhausted: Additional allocation failed", "with top memory consumers (across reservations) as:\n GroupedHashAggregateStream" + "Resources exhausted: Additional allocation failed", + "for PartialHashAggregateStream[0]", ]) .with_memory_limit(1_000) .run() @@ -423,7 +427,7 @@ async fn oom_grouped_hash_aggregate() { .with_query("SELECT COUNT(*), SUM(request_bytes) FROM t GROUP BY host") .with_expected_errors(vec![ "Failed to allocate additional", - "GroupedHashAggregateStream[0] (count(1), sum(t.request_bytes))", + "for PartialHashAggregateStream[0]", ]) .with_memory_limit(1_000) .run() @@ -546,16 +550,73 @@ async fn test_external_sort_zero_merge_reservation() { let _result = collect(stream).await; // Ensures the query spilled during execution - let metrics = physical_plan.metrics().unwrap(); - let spill_count = metrics.spill_count().unwrap(); + let spill_count = plan_spill_count(physical_plan.as_ref()); assert!(spill_count > 0); } +/// End-to-end (SQL-level) reproducer for the skewed-batch multi-level merge bug. +/// +/// The workload is a sort over wide rows under a tight memory budget. Each spilled +/// run's largest record batch is so wide that two merge streams cannot both be +/// reserved at once (`~4 * max_batch > pool`), yet a single stream still fits +/// (`~2 * max_batch < pool`). Reducing the read-ahead buffer therefore cannot help. +/// +/// Before the fix the multi-level merge gave up here with `ResourcesExhausted`; now +/// it re-spills the blocking run with a smaller batch size and the query completes. +/// +/// This complements the low-level unit tests in `multi_level_merge.rs`: it drives the +/// whole sort -> spill -> multi-level-merge pipeline from a SQL query, so the coverage +/// survives refactors of the merge internals. +#[tokio::test] +async fn test_sort_skewed_batches_spill() { + let pool_size = 2 * 1024 * 1024; // 2MB + let runtime = RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(FairSpillPool::new(pool_size))) + .build_arc() + .unwrap(); + + let config = SessionConfig::new() + .with_sort_spill_reservation_bytes(1) + .with_batch_size(8192) + .with_target_partitions(1); + + let ctx = SessionContext::new_with_config_rt(config, runtime); + + // Each row carries a ~100-byte string payload, so a full 8192-row batch is + // ~0.9MB. Reserving two such streams needs ~4 * 0.9MB > 2MB and cannot fit, + // while a single stream (~1.8MB) still fits - exactly the skew the fix handles. + // Sorting by the narrow `v` key forces the wide payload to be carried through + // the spill/merge path. + let row_count = 131072; + let query = "SELECT v, repeat('a', 100) AS payload \ + FROM generate_series(1, 131072) AS t(v) \ + ORDER BY v DESC"; + let df = ctx.sql(query).await.unwrap(); + + let physical_plan = df.create_physical_plan().await.unwrap(); + let task_ctx = Arc::new(TaskContext::from(&ctx.state())); + let stream = physical_plan.execute(0, task_ctx).unwrap(); + let batches = collect(stream) + .await + .expect("skewed sort should re-spill and complete, not exhaust memory"); + + // Every input row must come out of the merge. + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, row_count); + + // The query must actually spill, otherwise it never reaches the merge path + // this test is meant to cover. + assert!( + plan_spill_count(physical_plan.as_ref()) > 0, + "expected the sort to spill to disk" + ); +} + // Tests for disk limit (`max_temp_directory_size` in `DiskManager`) // ------------------------------------------------------------------ // Create a new `SessionContext` with specified disk limit, memory pool limit, and spill compression codec -async fn setup_context( +fn setup_context( disk_limit: u64, memory_pool_limit: usize, spill_compression: SpillCompression, @@ -596,7 +657,7 @@ async fn setup_context( #[tokio::test] async fn test_disk_spill_limit_reached() -> Result<()> { let spill_compression = SpillCompression::Uncompressed; - let ctx = setup_context(1024 * 1024, 1024 * 1024, spill_compression).await?; // 1MB disk limit, 1MB memory limit + let ctx = setup_context(1024 * 1024, 1024 * 1024, spill_compression)?; // 1MB disk limit, 1MB memory limit let df = ctx .sql("select * from generate_series(1, 1000000000000) as t1(v1) order by v1 desc") @@ -624,7 +685,7 @@ async fn test_disk_spill_limit_reached() -> Result<()> { async fn test_disk_spill_limit_not_reached() -> Result<()> { let disk_spill_limit = 1024 * 1024; // 1MB let spill_compression = SpillCompression::Uncompressed; - let ctx = setup_context(disk_spill_limit, 128 * 1024, spill_compression).await?; // 1MB disk limit, 128KB memory limit + let ctx = setup_context(disk_spill_limit, 128 * 1024, spill_compression)?; // 1MB disk limit, 128KB memory limit let df = ctx .sql("select * from generate_series(1, 10000) as t1(v1) order by v1 desc") @@ -637,8 +698,8 @@ async fn test_disk_spill_limit_not_reached() -> Result<()> { .await .expect("Query execution failed"); - let spill_count = plan.metrics().unwrap().spill_count().unwrap(); - let spilled_bytes = plan.metrics().unwrap().spilled_bytes().unwrap(); + let spill_count = plan_spill_count(plan.as_ref()); + let spilled_bytes = plan_spilled_bytes(plan.as_ref()); println!("spill count {spill_count}, spill bytes {spilled_bytes}"); assert!(spill_count > 0); @@ -660,7 +721,7 @@ async fn test_disk_spill_limit_not_reached() -> Result<()> { async fn test_spill_file_compressed_with_zstd() -> Result<()> { let disk_spill_limit = 1024 * 1024; // 1MB let spill_compression = SpillCompression::Zstd; - let ctx = setup_context(disk_spill_limit, 128 * 1024, spill_compression).await?; // 1MB disk limit, 128KB memory limit, zstd + let ctx = setup_context(disk_spill_limit, 128 * 1024, spill_compression)?; // 1MB disk limit, 128KB memory limit, zstd let df = ctx .sql("select * from generate_series(1, 100000) as t1(v1) order by v1 desc") @@ -673,8 +734,8 @@ async fn test_spill_file_compressed_with_zstd() -> Result<()> { .await .expect("Query execution failed"); - let spill_count = plan.metrics().unwrap().spill_count().unwrap(); - let spilled_bytes = plan.metrics().unwrap().spilled_bytes().unwrap(); + let spill_count = plan_spill_count(plan.as_ref()); + let spilled_bytes = plan_spilled_bytes(plan.as_ref()); println!("spill count {spill_count}"); assert!(spill_count > 0); @@ -696,7 +757,7 @@ async fn test_spill_file_compressed_with_zstd() -> Result<()> { async fn test_spill_file_compressed_with_lz4_frame() -> Result<()> { let disk_spill_limit = 1024 * 1024; // 1MB let spill_compression = SpillCompression::Lz4Frame; - let ctx = setup_context(disk_spill_limit, 128 * 1024, spill_compression).await?; // 1MB disk limit, 128KB memory limit, lz4_frame + let ctx = setup_context(disk_spill_limit, 128 * 1024, spill_compression)?; // 1MB disk limit, 128KB memory limit, lz4_frame let df = ctx .sql("select * from generate_series(1, 100000) as t1(v1) order by v1 desc") @@ -709,8 +770,8 @@ async fn test_spill_file_compressed_with_lz4_frame() -> Result<()> { .await .expect("Query execution failed"); - let spill_count = plan.metrics().unwrap().spill_count().unwrap(); - let spilled_bytes = plan.metrics().unwrap().spilled_bytes().unwrap(); + let spill_count = plan_spill_count(plan.as_ref()); + let spilled_bytes = plan_spilled_bytes(plan.as_ref()); println!("spill count {spill_count}"); assert!(spill_count > 0); diff --git a/datafusion/core/tests/memory_limit/union_nullable_spill.rs b/datafusion/core/tests/memory_limit/union_nullable_spill.rs index c5ef2387d3cdc..d04273bc7fdb1 100644 --- a/datafusion/core/tests/memory_limit/union_nullable_spill.rs +++ b/datafusion/core/tests/memory_limit/union_nullable_spill.rs @@ -103,10 +103,15 @@ fn build_task_ctx(pool_size: usize) -> Arc { /// have mismatched nullability (one child's `val` is non-nullable, the other's /// is nullable with NULLs). A tiny FairSpillPool forces all batches to spill. /// -/// UnionExec returns child streams without schema coercion, so batches from -/// different children carry different per-field nullability into the shared -/// SpillPool. The IPC writer must use the SpillManager's canonical (nullable) -/// schema — not the first batch's schema — so readback batches are valid. +/// `UnionExec` now re-stamps every child batch with its own declared (nullable) +/// schema before they reach `RepartitionExec` (see +/// ), so this no longer +/// exercises mismatched-nullability batches arriving at the SpillManager via +/// `UnionExec` specifically. It's kept as a regression test for the +/// SpillManager fix itself: the IPC writer must use the SpillManager's +/// canonical schema -- not the first batch's schema -- so readback batches +/// stay valid for any caller that does hand it batches with differing +/// nullability. See . /// /// Otherwise, sort_batch will panic with /// `Column 'val' is declared as non-nullable but contains null values` diff --git a/datafusion/core/tests/optimizer/mod.rs b/datafusion/core/tests/optimizer/mod.rs index c8208ef3efa90..0bfe1fac68795 100644 --- a/datafusion/core/tests/optimizer/mod.rs +++ b/datafusion/core/tests/optimizer/mod.rs @@ -216,7 +216,7 @@ impl ContextProvider for MyContextProvider { self.udfs.get(name).cloned() } - fn get_higher_order_meta(&self, _name: &str) -> Option> { + fn get_higher_order_meta(&self, _name: &str) -> Option> { None } diff --git a/datafusion/core/tests/parquet/content_defined_chunking.rs b/datafusion/core/tests/parquet/content_defined_chunking.rs index 6a98ded1bd4cf..bd89b502bd272 100644 --- a/datafusion/core/tests/parquet/content_defined_chunking.rs +++ b/datafusion/core/tests/parquet/content_defined_chunking.rs @@ -25,7 +25,7 @@ use arrow::array::{AsArray, Int32Array, StringArray}; use arrow::datatypes::{DataType, Field, Int32Type, Int64Type, Schema}; use arrow::record_batch::RecordBatch; use datafusion::prelude::{ParquetReadOptions, SessionContext}; -use datafusion_common::config::{CdcOptions, TableParquetOptions}; +use datafusion_common::config::{ParquetCdcOptions, TableParquetOptions}; use parquet::arrow::ArrowWriter; use parquet::arrow::arrow_reader::ArrowReaderMetadata; use parquet::file::properties::WriterProperties; @@ -97,7 +97,7 @@ async fn cdc_data_round_trip() { let batch = make_test_batch(5000); let mut opts = TableParquetOptions::default(); - opts.global.use_content_defined_chunking = Some(CdcOptions::default()); + opts.global.content_defined_chunking = ParquetCdcOptions::enabled(); let props = writer_props(&mut opts, &batch.schema()); let tmp = write_parquet_file(&batch, props); @@ -145,11 +145,12 @@ async fn cdc_affects_page_boundaries() { // Write WITH CDC using small chunk sizes to maximize effect let mut cdc_opts = TableParquetOptions::default(); - cdc_opts.global.use_content_defined_chunking = Some(CdcOptions { + cdc_opts.global.content_defined_chunking = ParquetCdcOptions { + enabled: true, min_chunk_size: 512, max_chunk_size: 2048, norm_level: 0, - }); + }; let cdc_file = write_parquet_file(&batch, writer_props(&mut cdc_opts, &batch.schema())); let cdc_meta = read_metadata(&cdc_file); diff --git a/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs b/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs new file mode 100644 index 0000000000000..5ee42b30674bf --- /dev/null +++ b/datafusion/core/tests/parquet/dynamic_row_group_pruning.rs @@ -0,0 +1,718 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end test for **runtime row-group pruning** driven by a TopK +//! `SortExec`'s `DynamicFilterPhysicalExpr`. +//! +//! A 5-row-group parquet file is constructed with disjoint statistics on +//! the sort column (`v`): row group `i` contains values +//! `[i*100, (i+1)*100)`. The query `ORDER BY v DESC LIMIT 5` fills the +//! TopK heap from the row group with the largest values; the threshold +//! then proves the remaining row groups cannot contribute. The runtime +//! `RowGroupPruner` in the parquet scan must observe the tightened +//! threshold and increment `row_groups_pruned_dynamic_filter`. +//! +//! We assert a property (`pruned >= 1`) rather than an exact count +//! because batch-arrival timing affects how soon the TopK heap fills, +//! and we don't want this test to become flaky. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use arrow_schema::{DataType, Field, Schema}; + +use datafusion::prelude::SessionConfig; + +use crate::parquet::Unit::RowGroup; +use crate::parquet::{ContextWithParquet, Scenario}; + +/// Build five `RecordBatch`es whose `v` column ranges are disjoint: +/// batch `i` carries `v` values `[i*100, (i+1)*100)`. When written with +/// `max_row_group_row_count = 100` each batch lands in its own row group. +fn build_five_disjoint_batches(schema: &Arc) -> Vec { + (0..5i64) + .map(|rg| { + let base = rg * 100; + let values: Vec = (base..base + 100).collect(); + let col: ArrayRef = Arc::new(Int64Array::from(values)); + RecordBatch::try_new(Arc::clone(schema), vec![col]).unwrap() + }) + .collect() +} + +/// Build five `RecordBatch`es in *descending* value order: batch 0 holds +/// `v ∈ [400, 500)`, batch 4 holds `v ∈ [0, 100)`. The physical row-group +/// order on disk therefore does **not** match the order a `ORDER BY v ASC` +/// query wants — sort-pushdown's `reorder_by_statistics` must rearrange +/// the access plan so the scan reads RG 4 first, then RG 3, etc. +fn build_five_disjoint_batches_desc(schema: &Arc) -> Vec { + (0..5i64) + .map(|rg| { + let base = (4 - rg) * 100; + let values: Vec = (base..base + 100).collect(); + let col: ArrayRef = Arc::new(Int64Array::from(values)); + RecordBatch::try_new(Arc::clone(schema), vec![col]).unwrap() + }) + .collect() +} + +/// `ORDER BY v DESC LIMIT 5` against a 5-RG file with disjoint per-RG +/// stats must trigger runtime RG pruning: the first RG read fills the +/// heap, and the tightened threshold proves every other RG unreachable. +#[tokio::test] +async fn dynamic_rg_pruning_metric_fires_for_topk_descending_limit() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batches = build_five_disjoint_batches(&schema); + + // `with_custom_data` honors the custom schema + batches and ignores + // `Scenario`. `Unit::RowGroup(100)` enables `pushdown_filters`, which + // is required for the TopK dynamic filter to reach the parquet scan. + let mut ctx = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroup(100), + Arc::clone(&schema), + batches, + ) + .await; + + let output = ctx.query("SELECT v FROM t ORDER BY v DESC LIMIT 5").await; + + assert_eq!(output.result_rows, 5, "query must return LIMIT rows",); + + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert!( + pruned >= 1, + "dynamic RG pruner must skip at least one row group; \ + pruned={pruned}\n{}", + output.description(), + ); +} + +/// Regression for the rg_plan / `reorder_by_statistics` ordering bug. +/// +/// When `sort_order_for_reorder` is set on the parquet scan, +/// `prepare_access_plan` calls +/// [`PreparedAccessPlan::reorder_by_statistics`], which rearranges +/// `row_group_indexes` so the decoder reads row groups in stats-optimal +/// order (smallest-min first for ASC, etc.). The stream's per-RG plan +/// (`rg_plan`) — which the runtime pruner walks one entry at a time — +/// **must use this reordered list**, not the access plan's natural +/// (index-ascending) order. Otherwise the pruner would consult the +/// metadata of RG K while the decoder is actually about to yield RG K', +/// silently producing wrong results. +/// +/// This test makes the failure visible: +/// +/// - File is written with RGs in *descending* `v` order (RG 0 has the +/// largest values, RG 4 has the smallest). +/// - Query is `ORDER BY v ASC LIMIT 5`, so sort-pushdown reorders the +/// access plan to read RG 4 first, then RG 3, etc. +/// - The smallest five values (which form the entire correct LIMIT +/// answer) live in RG 4 alone. After they are emitted, the TopK +/// threshold tightens enough that the per-RG pruner skips every other +/// RG. +/// +/// Without the fix, `rg_plan` would be `[0, 1, 2, 3, 4]` while the +/// decoder reads `[4, 3, 2, 1, 0]`. The first yielded reader (for RG 4 +/// in the decoder) would be tracked as if it were RG 0, the pruner +/// would check RG 1's stats (id range 300..400) against a threshold +/// already tightened to `v < 5`, prune RG 1 (because nothing in +/// 300..400 can satisfy `v < 5`), and then the rebuild via +/// `into_builder` would scan a row group whose data does not match its +/// expected metadata. The query would return fewer than five rows or +/// the wrong rows. +#[tokio::test] +async fn dynamic_rg_pruning_handles_sort_pushdown_reorder() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batches = build_five_disjoint_batches_desc(&schema); + + let mut ctx = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroup(100), + Arc::clone(&schema), + batches, + ) + .await; + + let output = ctx.query("SELECT v FROM t ORDER BY v ASC LIMIT 5").await; + + // Correctness — the five smallest values in the file are 0..=4. + // If `rg_plan` is misaligned with the decoder's read order, the + // pruner consults the wrong RG's stats and the result row count or + // values would drift. + assert_eq!(output.result_rows, 5, "query must return LIMIT rows"); + let formatted = output.pretty_results(); + for v in 0..=4i64 { + assert!( + formatted.contains(&format!("| {v} ")), + "output must contain the smallest value {v}; got:\n{formatted}", + ); + } + + // Behavior — the per-RG pruner must engage. We don't pin the exact + // count (batch-arrival timing affects how soon the heap fills); we + // only require that at least one row group is skipped at runtime. + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert!( + pruned >= 1, + "with `sort_order_for_reorder` active and a tight TopK, the \ + runtime pruner must skip at least one row group; pruned={pruned}\n{}", + output.description(), + ); +} + +/// A query without ORDER BY does not produce a TopK and therefore no +/// `DynamicFilterPhysicalExpr` reaches the scan. The runtime pruner must +/// stay quiet — the metric should be 0. +#[tokio::test] +async fn dynamic_rg_pruning_metric_quiet_without_topk() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batches = build_five_disjoint_batches(&schema); + + let mut ctx = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroup(100), + Arc::clone(&schema), + batches, + ) + .await; + + // Plain `SELECT *` — no sort, no limit, no dynamic filter. + let output = ctx.query("SELECT v FROM t").await; + assert_eq!(output.result_rows, 500); + + let pruned = output.row_groups_pruned_dynamic_filter().unwrap_or(0); + assert_eq!( + pruned, + 0, + "without TopK there is no dynamic filter, so the runtime pruner \ + must not fire; pruned={pruned}\n{}", + output.description(), + ); +} + +/// Regression for "into_builder called mid-row-group" — surfaced by +/// ClickBench Q24 / Q26 (`SELECT … WHERE x <> '' ORDER BY ts LIMIT 10`). +/// +/// The push-decoder state machine re-enters Step 2 on every iteration of +/// the `transition` loop, including iterations where Step 3 returned +/// `NeedsData` and pushed byte ranges but has not yet produced a reader +/// for the upcoming row group. At those moments the decoder is in +/// `ReadingRowGroup` state but `is_at_row_group_boundary()` is `false`, +/// and the runtime row-group pruner's `into_builder()` rebuild path +/// errored out with: +/// +/// ```text +/// Parquet error: into_builder called mid-row-group; +/// check is_at_row_group_boundary() first +/// ``` +/// +/// The fix in `push_decoder.rs::Step 2` gates the prune-and-rebuild on +/// `is_at_row_group_boundary()`. This test reproduces the trigger: a +/// many-RG file (so the pruner has work to do) plus an `ORDER BY` query +/// whose TopK threshold tightens enough to make the pruner want to +/// rebuild more than once during the scan. Before the fix the query +/// returned an `Execution` / `Parquet` error; after the fix it returns +/// the expected ten rows and the pruner fires. +#[tokio::test] +async fn dynamic_rg_pruner_does_not_call_into_builder_mid_row_group() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + // 20 disjoint row groups of 50 values each. With 20 RGs the pruner + // gets multiple boundaries to attempt rebuilds, so any path that + // calls `into_builder` outside a boundary is hit reliably. + let batches: Vec = (0..20i64) + .map(|rg| { + let base = rg * 50; + let values: Vec = (base..base + 50).collect(); + let col: ArrayRef = Arc::new(Int64Array::from(values)); + RecordBatch::try_new(Arc::clone(&schema), vec![col]).unwrap() + }) + .collect(); + + let mut ctx = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroup(50), + Arc::clone(&schema), + batches, + ) + .await; + + let output = ctx.query("SELECT v FROM t ORDER BY v ASC LIMIT 10").await; + + // Correctness: smallest ten values are 0..=9. + assert_eq!(output.result_rows, 10, "query must return LIMIT rows"); + let formatted = output.pretty_results(); + for v in 0..=9i64 { + assert!( + formatted.contains(&format!("| {v} ")), + "output must contain smallest value {v}; got:\n{formatted}", + ); + } + + // Behavior: with 20 disjoint RGs and a tight TopK, the dynamic + // pruner must skip a meaningful share of them. We don't pin the + // exact count — what matters is that the scan *completed* without + // the mid-row-group rebuild error. + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert!( + pruned >= 1, + "dynamic RG pruner must skip at least one row group; \ + pruned={pruned}\n{}", + output.description(), + ); +} + +/// Build five sorted `RecordBatch`es with 1000 values each so that, when +/// the writer is configured with `row_per_group=1000` and +/// `data_page_row_count_limit=100`, every row group ends up with **ten +/// data pages** of 100 rows each. RG `i` covers `[i*1000, (i+1)*1000)`, +/// monotonically ascending — page index will then have tight per-page +/// `min`/`max` and can prune at sub-RG granularity. +fn build_five_thousand_row_rgs(schema: &Arc) -> Vec { + (0..5i64) + .map(|rg| { + let base = rg * 1000; + let values: Vec = (base..base + 1000).collect(); + let col: ArrayRef = Arc::new(Int64Array::from(values)); + RecordBatch::try_new(Arc::clone(schema), vec![col]).unwrap() + }) + .collect() +} + +/// Regression test for : +/// when a page-index `RowSelection` is live, the runtime dynamic row-group +/// pruner is intentionally **not built**, so its `into_builder` rebuild can +/// never drop a row group without slicing the carried selection (which would +/// silently return wrong rows). Correctness is bought at the cost of the +/// dynamic-pruning optimization for this scan. +/// +/// The behavior asserted below (pruner disabled → +/// `row_groups_pruned_dynamic_filter == 0`) is expected to change once the +/// proper upstream fix lands, which keeps both mechanisms: +/// (tracked on the +/// DataFusion side in ). +/// +/// Layout: 5 RGs × 1000 rows, with `data_page_row_count_limit=100` so +/// each RG has 10 pages of 100 rows. +/// +/// Query: `SELECT v FROM t WHERE v >= 500 ORDER BY v DESC LIMIT 5`. +/// - `v >= 500` engages the page index: in RG 0 (values 0..1000) the +/// first 5 pages (values 0..500) are pruned, the last 5 (500..1000) +/// are scanned. RGs 1..4 keep all their pages (every page has +/// `max >= 500`). The decoder receives a `RowSelection` that masks +/// out those first 5 pages of RG 0 — its presence is what suppresses +/// the runtime pruner. +/// - `ORDER BY v DESC LIMIT 5` would let the tightened TopK threshold +/// (≥ 4995) prune RGs 0..3, but because a row selection is present the +/// runtime pruner is never created, so `row_groups_pruned_dynamic_filter` +/// stays 0. Results are still correct and page-index pruning still runs. +#[tokio::test] +async fn dynamic_rg_pruning_disabled_when_page_index_row_selection_present() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batches = build_five_thousand_row_rgs(&schema); + + // `RowGroupAndPage(1000, 100)` enables both `pushdown_filters` and + // page-index pruning, and writes a parquet file with 1000-row RGs + // partitioned into 100-row pages. + let mut ctx = ContextWithParquet::with_custom_data( + Scenario::Int, + crate::parquet::Unit::RowGroupAndPage(1000, 100), + Arc::clone(&schema), + batches, + ) + .await; + + let output = ctx + .query("SELECT v FROM t WHERE v >= 500 ORDER BY v DESC LIMIT 5") + .await; + + // Correctness — top-5 values descending are 4995..=4999 (all in RG 4). + assert_eq!(output.result_rows, 5, "query must return LIMIT rows"); + let formatted = output.pretty_results(); + for v in 4995..=4999i64 { + assert!( + formatted.contains(&format!("| {v} ")), + "output must contain top-5 descending value {v}; got:\n{formatted}", + ); + } + + // Page-index pruning still engages: RG 0's first 5 pages are entirely + // < 500. #24355 only suppresses the *runtime* row-group pruner, not + // page-index pruning, so this must remain non-zero. + let pages_pruned = output.metric_value("page_index_pages_pruned").unwrap_or(0); + assert!( + pages_pruned >= 5, + "page index must prune at least 5 pages (RG 0 pages 0..5 for v < 500); \ + pruned={pages_pruned}\n{}", + output.description(), + ); + + // The runtime dynamic pruner must be disabled while a page-index row + // selection is live (#24355): with no pruner there is no rebuild that + // could misapply the carried selection. Before the fix the pruner ran + // and this metric was >= 1. + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert_eq!( + pruned, + 0, + "runtime row-group pruning must be skipped when a page-index row \ + selection is present; pruned={pruned}\n{}", + output.description(), + ); +} + +/// Co-existence test: a `WHERE` clause that gets pushed into the parquet +/// `RowFilter` plus a `TopK` that drives the dynamic RG pruner. +/// +/// `v % 2 = 0` cannot be statically pruned and is not page-index-amenable +/// either, so it must run per-row inside the parquet decoder as a +/// `RowFilter`. `ORDER BY v DESC LIMIT 3` then fills the TopK heap and +/// tightens the threshold, triggering runtime RG pruning. The decoder +/// rebuild that happens via +/// `into_builder().with_row_groups(remaining).build()` must preserve the +/// installed `RowFilter` (and any `RowSelection` derived from page-index +/// pruning) across the rebuild — if it didn't, either: +/// +/// - The post-prune RGs would silently drop their per-row filtering and +/// the result would contain odd values, OR +/// - The rebuilt decoder would re-emit rows the original was about to +/// yield, double-counting against the limit. +/// +/// This test catches both regressions: it pins both the exact result rows +/// (top three even values descending: 498, 496, 494) and asserts the +/// dynamic pruner fired at least once. +#[tokio::test] +async fn dynamic_rg_pruning_coexists_with_row_filter() { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let batches = build_five_disjoint_batches(&schema); + + let mut ctx = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroup(100), + Arc::clone(&schema), + batches, + ) + .await; + + // `v % 2 = 0` survives stats pruning (every RG straddles even / odd), + // so the predicate is pushed into the decoder as a `RowFilter` and + // evaluated per row. The TopK on top still tightens the threshold and + // engages the runtime RG pruner. + let output = ctx + .query("SELECT v FROM t WHERE v % 2 = 0 ORDER BY v DESC LIMIT 3") + .await; + + assert_eq!(output.result_rows, 3, "query must return LIMIT rows"); + let formatted = output.pretty_results(); + for v in [498i64, 496, 494] { + assert!( + formatted.contains(&format!("| {v} ")), + "output must contain top-3 even descending value {v}; got:\n{formatted}", + ); + } + + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert!( + pruned >= 1, + "with WHERE v % 2 = 0 + TopK the runtime pruner must still skip at \ + least one row group; pruned={pruned}\n{}", + output.description(), + ); +} + +/// Build five two-column `RecordBatch`es: `a` is physically clustered +/// (batch `i` carries `a ∈ [i*100, (i+1)*100)`, disjoint per-RG stats) +/// and `b` is a per-batch shuffle (identical `[0, 100)` range in every +/// RG, useless for pruning). +fn build_two_col_leading_clustered(schema: &Arc) -> Vec { + (0..5i64) + .map(|rg| { + let base = rg * 100; + let a: Vec = (base..base + 100).collect(); + // pseudo-shuffled b, same value set in every RG + let b: Vec = (0..100).map(|i| (i * 37) % 100).collect(); + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(a)) as ArrayRef, + Arc::new(Int64Array::from(b)) as ArrayRef, + ], + ) + .unwrap() + }) + .collect() +} + +/// Build five two-column `RecordBatch`es where the *leading* sort key +/// ties everywhere (`a = 1` in every row / RG) and the *secondary* key +/// is clustered but stored in DESC disk order: batch 0 carries +/// `b ∈ [400, 500)`, batch 4 carries `b ∈ [0, 100)`. +/// +/// An `ORDER BY a, b LIMIT k` query wants the rows in batch 4 first; +/// reading disk order decodes every RG with a monotonically *improving* +/// threshold that never proves a later RG unwinnable. +fn build_two_col_leading_tied_desc(schema: &Arc) -> Vec { + (0..5i64) + .map(|rg| { + let base = (4 - rg) * 100; + let a: Vec = vec![1; 100]; + let b: Vec = (base..base + 100).collect(); + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(a)) as ArrayRef, + Arc::new(Int64Array::from(b)) as ArrayRef, + ], + ) + .unwrap() + }) + .collect() +} + +fn two_col_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])) +} + +/// A multi-column `ORDER BY a, b LIMIT k` must still engage the runtime +/// RG pruner through the *leading* disjunct of the lexicographic dynamic +/// filter (`a < x OR (a = x AND b < y)`): once the heap fills from the +/// first (best) row group, `min(a) > x` alone proves later RGs +/// unwinnable regardless of `b`. +#[tokio::test] +async fn dynamic_rg_pruning_fires_for_multi_column_sort_leading_clustered() { + let schema = two_col_schema(); + let batches = build_two_col_leading_clustered(&schema); + + let mut ctx = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroup(100), + Arc::clone(&schema), + batches, + ) + .await; + + let output = ctx + .query("SELECT a, b FROM t ORDER BY a ASC, b ASC LIMIT 5") + .await; + + assert_eq!(output.result_rows, 5, "query must return LIMIT rows"); + + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert!( + pruned >= 1, + "multi-column TopK must prune via the leading column's disjunct; \ + pruned={pruned}\n{}", + output.description(), + ); +} + +/// When the leading sort key ties across all row groups, pruning (and +/// reading the right RG first) must fall to the *secondary* key: RG +/// stats give `min(a) = max(a) = 1` everywhere, so the lex dynamic +/// filter reduces to `a = 1 AND b < y` — prunable via `min(b)`. +/// +/// The disk order is adversarial (secondary key DESC), so without +/// multi-column stats reorder the scan reads the worst RG first and the +/// threshold never proves later RGs unwinnable. With multi-column +/// reorder the best RG is read first and every other RG is pruned. +#[tokio::test] +async fn dynamic_rg_pruning_fires_for_multi_column_sort_leading_tied() { + let schema = two_col_schema(); + let batches = build_two_col_leading_tied_desc(&schema); + + let mut ctx = ContextWithParquet::with_custom_data( + Scenario::Int, + RowGroup(100), + Arc::clone(&schema), + batches, + ) + .await; + + let output = ctx + .query("SELECT a, b FROM t ORDER BY a ASC, b ASC LIMIT 5") + .await; + + assert_eq!(output.result_rows, 5, "query must return LIMIT rows"); + // The leading key `a = 1` is tied everywhere, so correctness rests + // entirely on the secondary key: the five smallest `b` values must come + // back, in ascending secondary order. Assert the exact result rows + // (full two-column text, in order) rather than just probing for each + // `b` — a bare `| {b} ` match would be satisfied by the leading `a = 1` + // column even if that `b` were missing or misordered. + let formatted = output.pretty_results(); + let data_rows: Vec<&str> = formatted + .lines() + .filter(|line| line.starts_with("| 1 |")) + .collect(); + assert_eq!( + data_rows, + vec![ + "| 1 | 0 |", + "| 1 | 1 |", + "| 1 | 2 |", + "| 1 | 3 |", + "| 1 | 4 |", + ], + "output must be exactly (a=1, b=0..=4) in ascending secondary order; got:\n{formatted}", + ); + + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert!( + pruned >= 1, + "with the leading key tied everywhere, the secondary key must \ + drive RG reorder + pruning; pruned={pruned}\n{}", + output.description(), + ); +} + +/// Build the #24352 fixture: four 2048-row row groups where the filter column +/// (`search_phrase`) differs from the sort column (`event_time`), and one row +/// group (the second) has an empty post-predicate selection invisible to +/// statistics — its only small `event_time` (50) sits on the row whose +/// `search_phrase` is `''`. +/// +/// RG 0: event_time = i*1000 (i in 0..2048) +/// RG 1: i=2048 -> (50, ''), else (20000+i, 'p'||i) (i in 2048..4096) +/// RG 2: event_time = 100 + (i-4096) (i in 4096..6144) +/// RG 3: event_time = 5000 + (i-6144) (i in 6144..8192) +fn build_q26_batches(schema: &Arc) -> Vec { + (0..4i64) + .map(|rg| { + let mut event_time = Vec::with_capacity(2048); + let mut search_phrase: Vec = Vec::with_capacity(2048); + for j in 0..2048i64 { + let i = rg * 2048 + j; + let (et, sp) = if i < 2048 { + (i * 1000, format!("p{i}")) + } else if i < 4096 { + if i == 2048 { + (50, String::new()) + } else { + (20000 + i, format!("p{i}")) + } + } else if i < 6144 { + (100 + (i - 4096), format!("p{i}")) + } else { + (5000 + (i - 6144), format!("p{i}")) + }; + event_time.push(et); + search_phrase.push(sp); + } + RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int64Array::from(event_time)) as ArrayRef, + Arc::new(StringArray::from(search_phrase)) as ArrayRef, + ], + ) + .unwrap() + }) + .collect() +} + +/// Regression for #24352: with `pushdown_filters` + TopK dynamic filter, a row +/// group whose post-predicate selection is empty is silently finished by +/// arrow-rs without handing back a reader. Before `rg_plan` was synced to the +/// decoder frontier (`peek_next_row_group`), it trailed the decoder by one, so +/// a later runtime prune rebuilt the decoder from a stale plan and re-read an +/// already-delivered row group — the duplicate rows displaced the true top-k. +#[tokio::test] +async fn topk_pushdown_does_not_reread_delivered_row_group() { + let schema = Arc::new(Schema::new(vec![ + Field::new("event_time", DataType::Int64, false), + Field::new("search_phrase", DataType::Utf8, false), + ])); + let batches = build_q26_batches(&schema); + + // `RowGroup(2048)` writes one row group per 2048-row batch (4 RGs) and + // enables `pushdown_filters`, required for the dynamic filter to reach the + // parquet scan. Page-index reading is disabled: this test exercises the + // #24352 empty-row-group / rg_plan-sync path, which is row-filter-driven and + // does not need the page index. With the page index on, `search_phrase <> ''` + // produces an intra-row-group `RowSelection`, and #24355 disables the runtime + // pruner whenever a row selection is present — which would stop this test + // from exercising the dynamic pruner at all. + let mut config = SessionConfig::new(); + config.options_mut().execution.parquet.enable_page_index = false; + let mut ctx = ContextWithParquet::with_config( + Scenario::Int, + RowGroup(2048), + config, + Some(Arc::clone(&schema)), + Some(batches), + ) + .await; + + let output = ctx + .query( + "SELECT search_phrase FROM t \ + WHERE search_phrase <> '' ORDER BY event_time LIMIT 10", + ) + .await; + + // `search_phrase` is unique per row, so any repeated value is the same + // source row emitted twice. The correct answer is the 10 smallest- + // `event_time` non-empty phrases, matching DuckDB / pushdown-off. + assert_eq!(output.result_rows, 10, "{}", output.description()); + + // The test must actually exercise the runtime prune/rebuild path that + // caused #24352 (not just a happy-path scan), otherwise a future default or + // optimizer change could let it pass without the bug's precondition. Assert + // the dynamic filter pruned at least one row group. + let pruned = output + .row_groups_pruned_dynamic_filter() + .expect("`row_groups_pruned_dynamic_filter` metric must be registered"); + assert!( + pruned >= 1, + "test must exercise dynamic RG pruning (the #24352 path); pruned={pruned}\n{}", + output.description(), + ); + + let formatted = output.pretty_results(); + for p in [ + "p0", "p4096", "p4097", "p4098", "p4099", "p4100", "p4101", "p4102", "p4103", + "p4104", + ] { + assert!( + formatted.contains(&format!("| {p} ")), + "missing {p} from top-k; got:\n{formatted}", + ); + } + // The bug emitted p4096 twice (and dropped p4101..=p4104); assert no dup. + assert_eq!( + formatted.matches("| p4096 ").count(), + 1, + "p4096 emitted more than once — rg_plan/decoder desync; got:\n{formatted}", + ); +} diff --git a/datafusion/core/tests/parquet/expr_adapter.rs b/datafusion/core/tests/parquet/expr_adapter.rs index fd70d74a9140c..535828fa29c2f 100644 --- a/datafusion/core/tests/parquet/expr_adapter.rs +++ b/datafusion/core/tests/parquet/expr_adapter.rs @@ -18,8 +18,8 @@ use std::sync::Arc; use arrow::array::{ - Array, ArrayRef, BooleanArray, Int32Array, Int64Array, LargeListArray, ListArray, - RecordBatch, StringArray, StructArray, record_batch, + Array, ArrayRef, BooleanArray, FixedSizeListArray, Int32Array, Int64Array, + LargeListArray, ListArray, RecordBatch, StringArray, StructArray, record_batch, }; use arrow::buffer::OffsetBuffer; use arrow::compute::concat_batches; @@ -60,13 +60,19 @@ async fn write_parquet(batch: RecordBatch, store: Arc, path: &s enum NestedListKind { List, LargeList, + FixedSizeList, } +const FIXED_SIZE_LIST_LEN: usize = 2; + impl NestedListKind { fn field_data_type(self, item_field: Arc) -> DataType { match self { Self::List => DataType::List(item_field), Self::LargeList => DataType::LargeList(item_field), + Self::FixedSizeList => { + DataType::FixedSizeList(item_field, FIXED_SIZE_LIST_LEN as i32) + } } } @@ -89,6 +95,19 @@ impl NestedListKind { values, None, )), + Self::FixedSizeList => { + assert_eq!( + lengths.as_slice(), + &[FIXED_SIZE_LIST_LEN], + "FixedSizeList fixtures must contain exactly {FIXED_SIZE_LIST_LEN} elements per row" + ); + Arc::new(FixedSizeListArray::new( + item_field, + FIXED_SIZE_LIST_LEN as i32, + values, + None, + )) + } } } @@ -96,6 +115,7 @@ impl NestedListKind { match self { Self::List => "list", Self::LargeList => "large_list", + Self::FixedSizeList => "fixed_size_list", } } } @@ -277,7 +297,8 @@ fn nested_list_table_schema( } // Helper to extract message values from a nested list column. -// Returns the values at indices 0 and 1 from either a ListArray or LargeListArray. +// Returns the values at indices 0 and 1 from either a ListArray, LargeListArray, +// or FixedSizeListArray. fn extract_nested_list_values( kind: NestedListKind, column: &ArrayRef, @@ -297,7 +318,50 @@ fn extract_nested_list_values( .expect("messages should be a LargeListArray"); (list.value(0), list.value(1)) } + NestedListKind::FixedSizeList => { + let list = column + .as_any() + .downcast_ref::() + .expect("messages should be a FixedSizeListArray"); + (list.value(0), list.value(1)) + } + } +} + +fn evolved_messages(kind: NestedListKind) -> Vec> { + let mut messages = vec![NestedMessageRow { + id: 30, + name: "gamma", + chain: Some("eth"), + ignored: Some(99), + }]; + if matches!(kind, NestedListKind::FixedSizeList) { + messages.push(NestedMessageRow { + id: 40, + name: "delta", + chain: Some("doge"), + ignored: Some(100), + }); + } + messages +} + +fn error_messages(kind: NestedListKind) -> Vec> { + let mut messages = vec![NestedMessageRow { + id: 10, + name: "alpha", + chain: Some("eth"), + ignored: None, + }]; + if matches!(kind, NestedListKind::FixedSizeList) { + messages.push(NestedMessageRow { + id: 20, + name: "beta", + chain: Some("doge"), + ignored: None, + }); } + messages } // Helper to set up a nested list test fixture. @@ -352,15 +416,11 @@ async fn assert_nested_list_struct_schema_evolution(kind: NestedListKind) -> Res ); // new.parquet shape: messages item struct adds nullable `chain` and extra `ignored`. + let new_messages = evolved_messages(kind); let new_batch = nested_messages_batch( kind, 2, - &[NestedMessageRow { - id: 30, - name: "gamma", - chain: Some("eth"), - ignored: Some(99), - }], + &new_messages, &message_fields(DataType::Utf8, true, true, true), ); @@ -429,7 +489,12 @@ async fn assert_nested_list_struct_schema_evolution(kind: NestedListKind) -> Res .as_any() .downcast_ref::() .unwrap(); - assert_eq!(new_chain.iter().collect::>(), vec![Some("eth")]); + let expected_new_chain = if matches!(kind, NestedListKind::FixedSizeList) { + vec![Some("eth"), Some("doge")] + } else { + vec![Some("eth")] + }; + assert_eq!(new_chain.iter().collect::>(), expected_new_chain); let projected = ctx .sql( @@ -863,12 +928,12 @@ async fn test_struct_schema_evolution_projection_and_filter() -> Result<()> { Ok(()) } -/// Macro to generate paired test functions for List and LargeList variants. -/// Expands to two `#[tokio::test]` functions with the specified names. -macro_rules! test_struct_schema_evolution_pair { +/// Macro to generate schema evolution tests for list-like variants. +macro_rules! test_struct_schema_evolution_variants { ( list: $list_test:ident, large_list: $large_list_test:ident, + fixed_size_list: $fixed_size_list_test:ident, fn: $assertion_fn:path $(, args: $($arg:expr),+)? ) => { #[tokio::test] @@ -880,10 +945,16 @@ macro_rules! test_struct_schema_evolution_pair { async fn $large_list_test() { $assertion_fn(NestedListKind::LargeList $(, $($arg),+)?).await; } + + #[tokio::test] + async fn $fixed_size_list_test() { + $assertion_fn(NestedListKind::FixedSizeList $(, $($arg),+)?).await; + } }; ( list: $list_test:ident, large_list: $large_list_test:ident, + fixed_size_list: $fixed_size_list_test:ident, fn_result: $assertion_fn:path ) => { #[tokio::test] @@ -895,31 +966,34 @@ macro_rules! test_struct_schema_evolution_pair { async fn $large_list_test() -> Result<()> { $assertion_fn(NestedListKind::LargeList).await } + + #[tokio::test] + async fn $fixed_size_list_test() -> Result<()> { + $assertion_fn(NestedListKind::FixedSizeList).await + } }; } -test_struct_schema_evolution_pair!( +test_struct_schema_evolution_variants!( list: test_list_struct_schema_evolution_end_to_end, large_list: test_large_list_struct_schema_evolution_end_to_end, + fixed_size_list: test_fixed_size_list_struct_schema_evolution_end_to_end, fn_result: assert_nested_list_struct_schema_evolution ); async fn assert_nested_list_struct_schema_evolution_errors( kind: NestedListKind, + source_includes_chain: bool, chain_type: DataType, chain_nullable: bool, expected_error: &str, ) { + let messages = error_messages(kind); let batch = nested_messages_batch( kind, 1, - &[NestedMessageRow { - id: 10, - name: "alpha", - chain: Some("eth"), - ignored: None, - }], - &message_fields(DataType::Utf8, true, true, false), + &messages, + &message_fields(DataType::Utf8, true, source_includes_chain, false), ); let table_schema = @@ -949,6 +1023,7 @@ async fn assert_nested_list_struct_schema_evolution_errors( async fn assert_non_nullable_missing_chain_field_fails(kind: NestedListKind) { assert_nested_list_struct_schema_evolution_errors( kind, + false, DataType::Utf8, false, "non-nullable", @@ -959,6 +1034,7 @@ async fn assert_non_nullable_missing_chain_field_fails(kind: NestedListKind) { async fn assert_incompatible_chain_field_fails(kind: NestedListKind) { assert_nested_list_struct_schema_evolution_errors( kind, + true, incompatible_chain_type(), true, "Cannot cast struct field 'chain'", @@ -970,15 +1046,17 @@ fn incompatible_chain_type() -> DataType { DataType::Struct(vec![Arc::new(Field::new("value", DataType::Utf8, true))].into()) } -test_struct_schema_evolution_pair!( +test_struct_schema_evolution_variants!( list: test_list_struct_schema_evolution_non_nullable_missing_field_fails, large_list: test_large_list_struct_schema_evolution_non_nullable_missing_field_fails, + fixed_size_list: test_fixed_size_list_struct_schema_evolution_non_nullable_missing_field_fails, fn: assert_non_nullable_missing_chain_field_fails ); -test_struct_schema_evolution_pair!( +test_struct_schema_evolution_variants!( list: test_list_struct_schema_evolution_incompatible_field_fails, large_list: test_large_list_struct_schema_evolution_incompatible_field_fails, + fixed_size_list: test_fixed_size_list_struct_schema_evolution_incompatible_field_fails, fn: assert_incompatible_chain_field_fails ); diff --git a/datafusion/core/tests/parquet/external_access_plan.rs b/datafusion/core/tests/parquet/external_access_plan.rs index 31be6fd979fd6..8fd9689ae3a8d 100644 --- a/datafusion/core/tests/parquet/external_access_plan.rs +++ b/datafusion/core/tests/parquet/external_access_plan.rs @@ -29,8 +29,10 @@ use datafusion::common::Result; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::physical_plan::ParquetSource; use datafusion::prelude::SessionContext; -use datafusion_common::{DFSchema, assert_contains}; -use datafusion_datasource_parquet::{ParquetAccessPlan, RowGroupAccess}; +use datafusion_common::{DFSchema, assert_batches_eq, assert_contains}; +use datafusion_datasource_parquet::{ + ParquetAccessPlan, ParquetRowSelection, RowGroupAccess, +}; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::{Expr, col, lit}; use datafusion_physical_plan::ExecutionPlan; @@ -152,6 +154,94 @@ async fn skip_scan() { } } +#[tokio::test] +async fn row_selection_extension() { + // The file has 2 row groups of 5 rows each (10 rows total). Attach a + // file-level `ParquetRowSelection` to the `PartitionedFile` and verify it + // survives the path from `PartitionedFile` into the parquet opener/reader. + + // select a single row in the first row group + let parquet_metrics = TestFull { + access_plan: None, + row_selection: Some(ParquetRowSelection::new(RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(1), + RowSelector::skip(7), + ]))), + expected_rows: 1, + expected_output: Some(&[ + "+------+------------+", + "| utf8 | large_utf8 |", + "+------+------------+", + "| c | c |", + "+------+------------+", + ]), + predicate: None, + } + .run() + .await + .unwrap(); + + // only the first row group is read, so some bytes are scanned + let bytes_scanned = metric_value(&parquet_metrics, "bytes_scanned").unwrap(); + assert_ne!(bytes_scanned, 0, "metrics : {parquet_metrics:#?}",); +} + +#[tokio::test] +async fn row_selection_extension_spanning_row_groups() { + // A selection whose selectors straddle the row group boundary (row 4 is the + // last row of group 0, rows 5-6 are the first rows of group 1). + let parquet_metrics = TestFull { + access_plan: None, + row_selection: Some(ParquetRowSelection::new(RowSelection::from(vec![ + RowSelector::skip(4), + RowSelector::select(3), + RowSelector::skip(3), + ]))), + expected_rows: 3, + expected_output: Some(&[ + "+------+------------+", + "| utf8 | large_utf8 |", + "+------+------------+", + "| | |", + "| e | e |", + "| f | f |", + "+------+------------+", + ]), + predicate: None, + } + .run() + .await + .unwrap(); + + let bytes_scanned = metric_value(&parquet_metrics, "bytes_scanned").unwrap(); + assert_ne!(bytes_scanned, 0, "metrics : {parquet_metrics:#?}",); +} + +#[tokio::test] +async fn bad_row_selection_extension() { + // selection specifies fewer rows than the file actually contains + let err = TestFull { + access_plan: None, + row_selection: Some(ParquetRowSelection::new(RowSelection::from(vec![ + RowSelector::skip(2), + RowSelector::select(1), + ]))), + expected_rows: 10000, + expected_output: None, + predicate: None, + } + .run() + .await + .unwrap_err(); + let err_string = err.to_string(); + assert_contains!(&err_string, "Invalid Parquet RowSelection"); + assert_contains!( + &err_string, + "File has 10 rows, but selection specifies 3 rows." + ); +} + #[tokio::test] async fn plan_and_filter() { // show that row group pruning is applied even when an initial plan is supplied @@ -170,7 +260,9 @@ async fn plan_and_filter() { // initial let parquet_metrics = TestFull { access_plan, + row_selection: None, expected_rows: 0, + expected_output: None, predicate: Some(predicate), } .run() @@ -227,7 +319,9 @@ async fn bad_row_groups() { RowGroupAccess::Skip, RowGroupAccess::Scan, ])), + row_selection: None, expected_rows: 0, + expected_output: None, predicate: None, } .run() @@ -249,8 +343,10 @@ async fn bad_selection() { ])), RowGroupAccess::Skip, ])), + row_selection: None, // expects that we hit an error, this should not be run expected_rows: 10000, + expected_output: None, predicate: None, } .run() @@ -300,7 +396,9 @@ impl Test { } = self; TestFull { access_plan, + row_selection: None, expected_rows, + expected_output: None, predicate: None, } .run() @@ -317,7 +415,9 @@ impl Test { /// 4. Returns the statistics from running the plan struct TestFull { access_plan: Option, + row_selection: Option, expected_rows: usize, + expected_output: Option<&'static [&'static str]>, predicate: Option, } @@ -327,7 +427,9 @@ impl TestFull { let Self { access_plan, + row_selection, expected_rows, + expected_output, predicate, } = self; @@ -352,6 +454,11 @@ impl TestFull { partitioned_file = partitioned_file.with_extension(access_plan); } + // add the file-level row selection, if any, as an extension + if let Some(row_selection) = row_selection { + partitioned_file = partitioned_file.with_extension(row_selection); + } + // Create a DataSourceExec to read the file let object_store_url = ObjectStoreUrl::local_filesystem(); // add the predicate, if requested @@ -380,6 +487,9 @@ impl TestFull { "results: \n{}", pretty_format_batches(&results).unwrap() ); + if let Some(expected_output) = expected_output { + assert_batches_eq!(expected_output, &results); + } std::fs::remove_file(file_name).unwrap(); diff --git a/datafusion/core/tests/parquet/file_statistics.rs b/datafusion/core/tests/parquet/file_statistics.rs index 3e3b90a348b04..f6d733ec69720 100644 --- a/datafusion/core/tests/parquet/file_statistics.rs +++ b/datafusion/core/tests/parquet/file_statistics.rs @@ -18,6 +18,7 @@ use std::fs; use std::sync::Arc; +use arrow::datatypes::{DataType, Field, Schema}; use datafusion::datasource::TableProvider; use datafusion::datasource::file_format::parquet::ParquetFormat; use datafusion::datasource::listing::{ @@ -26,14 +27,14 @@ use datafusion::datasource::listing::{ use datafusion::datasource::source::DataSourceExec; use datafusion::execution::context::SessionState; use datafusion::execution::session_state::SessionStateBuilder; -use datafusion::prelude::SessionContext; -use datafusion_common::DFSchema; +use datafusion::prelude::{ParquetReadOptions, SessionContext}; use datafusion_common::stats::Precision; -use datafusion_execution::cache::DefaultListFilesCache; +use datafusion_common::{DFSchema, TableReference}; use datafusion_execution::cache::cache_manager::{ - CacheManagerConfig, FileStatisticsCache, + CacheManagerConfig, DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, + DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, FileStatisticsCache, ListFilesCache, }; -use datafusion_execution::cache::file_statistics_cache::DefaultFileStatisticsCache; +use datafusion_execution::cache::default_cache::DefaultCache; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::{Expr, col, lit}; @@ -44,6 +45,7 @@ use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::filter_pushdown::FilterPushdown; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::filter::FilterExec; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use tempfile::tempdir; #[tokio::test] @@ -52,18 +54,21 @@ async fn check_stats_precision_with_filter_pushdown() { let filename = format!("{}/{}", testdata, "alltypes_plain.parquet"); let table_path = ListingTableUrl::parse(filename).unwrap(); - let opt = - ListingOptions::new(Arc::new(ParquetFormat::default())).with_collect_stat(true); + let opt = ListingOptions::new(Arc::new(ParquetFormat::default())); let table = get_listing_table(&table_path, None, &opt).await; let (_, _, state) = get_cache_runtime_state(); let mut options: ConfigOptions = state.config().options().as_ref().clone(); options.execution.parquet.pushdown_filters = true; + options.execution.collect_statistics = true; // Scan without filter, stats are exact let exec = table.scan(&state, None, &[], None).await.unwrap(); assert_eq!( - exec.partition_statistics(None).unwrap().num_rows, + StatisticsContext::new() + .compute(exec.as_ref(), &StatisticsArgs::new()) + .unwrap() + .num_rows, Precision::Exact(8), "Stats without filter should be exact" ); @@ -95,7 +100,10 @@ async fn check_stats_precision_with_filter_pushdown() { ); // Scan with filter pushdown, stats are inexact assert_eq!( - optimized_exec.partition_statistics(None).unwrap().num_rows, + StatisticsContext::new() + .compute(optimized_exec.as_ref(), &StatisticsArgs::new()) + .unwrap() + .num_rows, Precision::Inexact(8), "Stats after filter pushdown should be inexact" ); @@ -105,15 +113,20 @@ async fn check_stats_precision_with_filter_pushdown() { async fn load_table_stats_with_session_level_cache() { let testdata = datafusion::test_util::parquet_test_data(); let filename = format!("{}/{}", testdata, "alltypes_plain.parquet"); - let table_path = ListingTableUrl::parse(filename).unwrap(); + let table_path = ListingTableUrl::parse(filename) + .unwrap() + .with_table_ref(TableReference::bare("alltypes_plain")); - let (cache1, _, state1) = get_cache_runtime_state(); + let (cache1, _, mut state1) = get_cache_runtime_state(); + let cfg_1 = state1.config_mut(); + cfg_1.options_mut().execution.collect_statistics = true; // Create a separate DefaultFileStatisticsCache - let (cache2, _, state2) = get_cache_runtime_state(); + let (cache2, _, mut state2) = get_cache_runtime_state(); + let cfg_2 = state2.config_mut(); + cfg_2.options_mut().execution.collect_statistics = true; - let opt = - ListingOptions::new(Arc::new(ParquetFormat::default())).with_collect_stat(true); + let opt = ListingOptions::new(Arc::new(ParquetFormat::default())); let table1 = get_listing_table(&table_path, Some(cache1), &opt).await; let table2 = get_listing_table(&table_path, Some(cache2), &opt).await; @@ -123,11 +136,17 @@ async fn load_table_stats_with_session_level_cache() { let exec1 = table1.scan(&state1, None, &[], None).await.unwrap(); assert_eq!( - exec1.partition_statistics(None).unwrap().num_rows, + StatisticsContext::new() + .compute(exec1.as_ref(), &StatisticsArgs::new()) + .unwrap() + .num_rows, Precision::Exact(8) ); assert_eq!( - exec1.partition_statistics(None).unwrap().total_byte_size, + StatisticsContext::new() + .compute(exec1.as_ref(), &StatisticsArgs::new()) + .unwrap() + .total_byte_size, // Byte size is absent because we cannot estimate the output size // of the Arrow data since there are variable length columns. Precision::Absent, @@ -139,11 +158,17 @@ async fn load_table_stats_with_session_level_cache() { assert_eq!(get_static_cache_size(&state2), 0); let exec2 = table2.scan(&state2, None, &[], None).await.unwrap(); assert_eq!( - exec2.partition_statistics(None).unwrap().num_rows, + StatisticsContext::new() + .compute(exec2.as_ref(), &StatisticsArgs::new()) + .unwrap() + .num_rows, Precision::Exact(8) ); assert_eq!( - exec2.partition_statistics(None).unwrap().total_byte_size, + StatisticsContext::new() + .compute(exec2.as_ref(), &StatisticsArgs::new()) + .unwrap() + .total_byte_size, // Absent because the data contains variable length columns Precision::Absent, ); @@ -154,11 +179,17 @@ async fn load_table_stats_with_session_level_cache() { assert_eq!(get_static_cache_size(&state1), 1); let exec3 = table1.scan(&state1, None, &[], None).await.unwrap(); assert_eq!( - exec3.partition_statistics(None).unwrap().num_rows, + StatisticsContext::new() + .compute(exec3.as_ref(), &StatisticsArgs::new()) + .unwrap() + .num_rows, Precision::Exact(8) ); assert_eq!( - exec3.partition_statistics(None).unwrap().total_byte_size, + StatisticsContext::new() + .compute(exec3.as_ref(), &StatisticsArgs::new()) + .unwrap() + .total_byte_size, // Absent because the data contains variable length columns Precision::Absent, ); @@ -166,6 +197,92 @@ async fn load_table_stats_with_session_level_cache() { assert_eq!(get_static_cache_size(&state1), 1); } +#[tokio::test] +async fn anonymous_parquet_stats_cache_with_explicit_wider_schema() { + let temp_dir = tempdir().unwrap(); + let parquet_path = temp_dir.path().join("data.parquet"); + let parquet_path = parquet_path.to_string_lossy().to_string(); + + let ctx = SessionContext::new_with_config( + SessionConfig::new().with_collect_statistics(true), + ); + let cache = ctx + .runtime_env() + .cache_manager + .get_file_statistic_cache() + .unwrap(); + + ctx.sql(&format!( + "COPY ( + SELECT 1::BIGINT AS id, 1000::BIGINT AS population + ) TO '{parquet_path}' STORED AS PARQUET" + )) + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!(cache.len(), 0); + + ctx.read_parquet(&parquet_path, ParquetReadOptions::default()) + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!(cache.len(), 1); + + let wider_schema = Schema::new(vec![ + Field::new("id", DataType::Int64, true), + Field::new("population", DataType::Int64, true), + Field::new("extra", DataType::Int64, true), + ]); + + let plan = ctx + .read_parquet( + &parquet_path, + ParquetReadOptions::default().schema(&wider_schema), + ) + .await + .unwrap() + .select_columns(&["id", "extra"]) + .unwrap() + .create_physical_plan() + .await + .unwrap(); + + let stats = StatisticsContext::new() + .compute(plan.as_ref(), &StatisticsArgs::new()) + .unwrap(); + assert_eq!(stats.column_statistics.len(), 2); + assert_eq!(stats.column_statistics[1].null_count, Precision::Exact(1)); + + // #23072: the cache now validates file_schema, so the wider read no + // longer bypasses the cache (as in #22950), but it overwrites the existing + // `{table, path}` entry instead of adding a schema-specific key. + assert_eq!(cache.len(), 1); + + // Repeat the wider read: same path + same file_schema -> reuse (no new + // entry) and a cache hit. Under #22950's bypass this read could never reuse. + ctx.read_parquet( + &parquet_path, + ParquetReadOptions::default().schema(&wider_schema), + ) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + assert_eq!(cache.len(), 1); + let hits: usize = cache.list_entries().values().map(|e| e.hits).sum(); + assert_eq!( + hits, 1, + "expected a cache hit on the repeat read, got {hits}" + ); +} + #[tokio::test] async fn list_files_with_session_level_cache() { let p_name = "alltypes_plain.parquet"; @@ -238,7 +355,7 @@ async fn list_files_with_session_level_cache() { async fn get_listing_table( table_path: &ListingTableUrl, - static_cache: Option>, + static_cache: Option>, opt: &ListingOptions, ) -> ListingTable { let schema = opt @@ -256,14 +373,13 @@ async fn get_listing_table( .with_cache(static_cache) } -fn get_cache_runtime_state() -> ( - Arc, - Arc, - SessionState, -) { +fn get_cache_runtime_state() +-> (Arc, Arc, SessionState) { let cache_config = CacheManagerConfig::default(); - let file_static_cache = Arc::new(DefaultFileStatisticsCache::default()); - let list_file_cache = Arc::new(DefaultListFilesCache::default()); + let file_static_cache = + Arc::new(DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT)); + let list_file_cache = + Arc::new(DefaultCache::new(DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT)); let cache_config = cache_config .with_file_statistics_cache(Some(file_static_cache.clone())) diff --git a/datafusion/core/tests/parquet/filter_pushdown.rs b/datafusion/core/tests/parquet/filter_pushdown.rs index e6266b2c088d7..dabb2f35b24b1 100644 --- a/datafusion/core/tests/parquet/filter_pushdown.rs +++ b/datafusion/core/tests/parquet/filter_pushdown.rs @@ -515,7 +515,6 @@ impl<'a> TestCase<'a> { let exec = self .test_parquet_file .create_scan(&ctx, Some(filter.clone())) - .await .unwrap(); let result = collect(exec.clone(), ctx.task_ctx()).await.unwrap(); @@ -648,7 +647,8 @@ async fn predicate_cache_stats_issue_19561() -> datafusion_common::Result<()> { let mut config = SessionConfig::new(); config.options_mut().execution.parquet.pushdown_filters = true; // force to get multiple batches to trigger repeated metric compound bug - config.options_mut().execution.batch_size = 1; + config.options_mut().execution.batch_size = + datafusion_common::config::ConfigNonZeroUsize::try_new(1)?; let ctx = SessionContext::new_with_config(config); // The cache is on by default, and used when filter pushdown is enabled PredicateCacheTest { diff --git a/datafusion/core/tests/parquet/mod.rs b/datafusion/core/tests/parquet/mod.rs index 0e936a79ebe9f..7066a4147c017 100644 --- a/datafusion/core/tests/parquet/mod.rs +++ b/datafusion/core/tests/parquet/mod.rs @@ -46,6 +46,7 @@ use tempfile::NamedTempFile; mod content_defined_chunking; mod custom_reader; +mod dynamic_row_group_pruning; #[cfg(feature = "parquet_encryption")] mod encryption; mod expr_adapter; @@ -99,6 +100,10 @@ enum Unit { RowGroup(usize), // pass max row per page in parquet writer Page(usize), + // pass max row per row_group AND max row per page. Use when a test + // needs both multi-RG layout AND multiple pages within each RG so the + // page index can prune at sub-RG granularity. + RowGroupAndPage(usize, usize), } /// Test fixture that has an execution context that has an external @@ -147,6 +152,12 @@ struct TestOutput { } impl TestOutput { + /// Pretty-printed result batches, useful for asserting concrete row + /// values in regression tests. + fn pretty_results(&self) -> &str { + &self.pretty_results + } + /// retrieve the value of the named metric, if any fn metric_value(&self, metric_name: &str) -> Option { if let Some(pm) = self.pruning_metric(metric_name) { @@ -259,6 +270,13 @@ impl TestOutput { .map(|pm| pm.total_pruned()) } + /// The number of row groups pruned at runtime by the dynamic + /// row-group pruner (e.g. driven by a TopK `SortExec` threshold + /// pushed down via `DynamicFilterPhysicalExpr`). + fn row_groups_pruned_dynamic_filter(&self) -> Option { + self.metric_value("row_groups_pruned_dynamic_filter") + } + fn description(&self) -> String { format!( "Input:\n{}\nQuery:\n{}\nOutput:\n{}\nMetrics:\n{}", @@ -305,12 +323,29 @@ impl ContextWithParquet { Unit::RowGroup(row_per_group) => { config = config.with_parquet_bloom_filter_pruning(true); config.options_mut().execution.parquet.pushdown_filters = true; - make_test_file_rg(scenario, row_per_group, custom_schema, custom_batches) - .await + make_test_file_rg( + scenario, + row_per_group, + None, + custom_schema, + custom_batches, + ) } Unit::Page(row_per_page) => { config = config.with_parquet_page_index_pruning(true); - make_test_file_page(scenario, row_per_page).await + make_test_file_page(scenario, row_per_page) + } + Unit::RowGroupAndPage(row_per_group, row_per_page) => { + config = config.with_parquet_bloom_filter_pruning(true); + config = config.with_parquet_page_index_pruning(true); + config.options_mut().execution.parquet.pushdown_filters = true; + make_test_file_rg( + scenario, + row_per_group, + Some(row_per_page), + custom_schema, + custom_batches, + ) } }; let parquet_path = file.path().to_string_lossy(); @@ -726,11 +761,11 @@ fn make_bytearray_batch( let name: StringArray = std::iter::repeat_n(Some(name), num_rows).collect(); let service_string: StringArray = string_values.iter().map(Some).collect(); let service_binary: BinaryArray = binary_values.iter().map(Some).collect(); - let service_fixedsize: FixedSizeBinaryArray = fixedsize_values - .iter() - .map(|value| Some(value.as_slice())) - .collect::>() - .into(); + let service_fixedsize = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + fixedsize_values.iter().map(|value| Some(value.as_slice())), + 3, + ) + .unwrap(); let service_large_binary: LargeBinaryArray = large_binary_values.iter().map(Some).collect(); @@ -1136,9 +1171,10 @@ fn create_data_batch(scenario: Scenario) -> Vec { } /// Create a test parquet file with various data types -async fn make_test_file_rg( +fn make_test_file_rg( scenario: Scenario, row_per_group: usize, + row_per_page: Option, custom_schema: Option, custom_batches: Option>, ) -> NamedTempFile { @@ -1148,11 +1184,19 @@ async fn make_test_file_rg( .tempfile() .expect("tempfile creation"); - let props = WriterProperties::builder() + let mut props_builder = WriterProperties::builder() .set_max_row_group_row_count(Some(row_per_group)) .set_bloom_filter_enabled(true) - .set_statistics_enabled(EnabledStatistics::Page) - .build(); + .set_statistics_enabled(EnabledStatistics::Page); + if let Some(rpp) = row_per_page { + // Bound rows per page so the page index can prune at sub-RG + // granularity. `write_batch_size` must also be set so the writer + // does not buffer the whole RG into one page. + props_builder = props_builder + .set_data_page_row_count_limit(rpp) + .set_write_batch_size(rpp); + } + let props = props_builder.build(); let (batches, schema) = if let (Some(schema), Some(batches)) = (custom_schema, custom_batches) { @@ -1173,7 +1217,7 @@ async fn make_test_file_rg( output_file } -async fn make_test_file_page(scenario: Scenario, row_per_page: usize) -> NamedTempFile { +fn make_test_file_page(scenario: Scenario, row_per_page: usize) -> NamedTempFile { let mut output_file = tempfile::Builder::new() .prefix("parquet_page_pruning") .suffix(".parquet") diff --git a/datafusion/core/tests/parquet/ordering.rs b/datafusion/core/tests/parquet/ordering.rs index faecb4ca6a861..1bdad7f593846 100644 --- a/datafusion/core/tests/parquet/ordering.rs +++ b/datafusion/core/tests/parquet/ordering.rs @@ -101,3 +101,83 @@ async fn test_create_table_with_order_writes_sorting_columns() -> Result<()> { Ok(()) } + +/// Test that partition columns are removed and remaining column indices are +/// remapped when writing sorting_columns to Parquet metadata. +#[tokio::test] +async fn test_partitioned_table_remaps_sorting_columns() -> Result<()> { + use parquet::file::reader::FileReader; + use parquet::file::serialized_reader::SerializedFileReader; + use std::fs::File; + + let ctx = SessionContext::new(); + let tmp_dir = tempdir()?; + let table_path = tmp_dir.path().join("sorted_partitioned_table"); + std::fs::create_dir_all(&table_path)?; + + let create_table_sql = format!( + "CREATE EXTERNAL TABLE sorted_partitioned_data (a INT, b VARCHAR, part VARCHAR) \ + STORED AS PARQUET \ + LOCATION '{}' \ + PARTITIONED BY (part) \ + WITH ORDER (part ASC NULLS FIRST, a ASC NULLS FIRST, b DESC NULLS LAST)", + table_path.display() + ); + ctx.sql(&create_table_sql).await?; + + ctx.sql( + "INSERT INTO sorted_partitioned_data VALUES \ + (2, 'c', 'x'), (1, 'a', 'x'), (1, 'b', 'x')", + ) + .await? + .collect() + .await?; + + let parquet_file = find_parquet_file(&table_path)? + .expect("expected a Parquet file in the partition directory"); + + let file = File::open(parquet_file)?; + let reader = SerializedFileReader::new(file)?; + let metadata = reader.metadata(); + let parquet_schema = metadata.file_metadata().schema_descr(); + assert_eq!(parquet_schema.num_columns(), 2); + assert_eq!(parquet_schema.column(0).name(), "a"); + assert_eq!(parquet_schema.column(1).name(), "b"); + + let sorting = metadata + .row_group(0) + .sorting_columns() + .expect("expected sorting_columns in row group metadata"); + assert_eq!(sorting.len(), 2); + + assert_eq!(sorting[0].column_idx, 0); + assert!(!sorting[0].descending); + assert!(sorting[0].nulls_first); + + assert_eq!(sorting[1].column_idx, 1); + assert!(sorting[1].descending); + assert!(!sorting[1].nulls_first); + + Ok(()) +} + +fn find_parquet_file( + path: &std::path::Path, +) -> std::io::Result> { + for entry in std::fs::read_dir(path)? { + let entry = entry?; + let path = entry.path(); + if path.is_dir() { + if let Some(path) = find_parquet_file(&path)? { + return Ok(Some(path)); + } + } else if path + .extension() + .is_some_and(|extension| extension == "parquet") + { + return Ok(Some(path)); + } + } + + Ok(None) +} diff --git a/datafusion/core/tests/parquet/page_pruning.rs b/datafusion/core/tests/parquet/page_pruning.rs index a41803191ad05..372a7a601d492 100644 --- a/datafusion/core/tests/parquet/page_pruning.rs +++ b/datafusion/core/tests/parquet/page_pruning.rs @@ -38,6 +38,7 @@ use datafusion_expr::{Expr, col, lit}; use datafusion_physical_expr::create_physical_expr; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use futures::StreamExt; use object_store::ObjectMeta; use object_store::path::Path; @@ -74,7 +75,13 @@ async fn get_parquet_exec( let df_schema = schema.clone().to_dfschema().unwrap(); let execution_props = ExecutionProps::new(); - let predicate = create_physical_expr(&filter, &df_schema, &execution_props).unwrap(); + let predicate = create_physical_expr( + &filter, + &df_schema, + &execution_props, + &PhysicalPlanningContext::default(), + ) + .unwrap(); let source = Arc::new( ParquetSource::new(schema.clone()) diff --git a/datafusion/core/tests/parquet/row_group_pruning.rs b/datafusion/core/tests/parquet/row_group_pruning.rs index 441d1af3e96fd..0721715921909 100644 --- a/datafusion/core/tests/parquet/row_group_pruning.rs +++ b/datafusion/core/tests/parquet/row_group_pruning.rs @@ -2078,3 +2078,26 @@ async fn test_limit_pruning_exceeds_fully_matched() -> datafusion_common::error: .await; Ok(()) } + +#[tokio::test] +async fn prune_like_prefix() { + // UTF8 scenario: 2 row groups (5 rows each) + // RG1: ["a","b","c","d",NULL] => min="a", max="d" + // RG2: ["e","f","g","h","i"] => min="e", max="i" + // + // LIKE 'a%' => build_like_match produces: "a" <= max AND min <= "a" (actually min < "b") + // RG1: "a" <= "d" ✓, "a" < "b" ✓ => matched + // RG2: "a" <= "i" ✓, "e" < "b" ✗ => pruned + RowGroupPruningTest::new() + .with_scenario(Scenario::UTF8) + .with_query("SELECT * FROM t WHERE utf8 LIKE 'a%'") + .with_expected_errors(Some(0)) + .with_matched_by_stats(Some(1)) + .with_pruned_by_stats(Some(1)) + .with_pruned_files(Some(0)) + .with_matched_by_bloom_filter(Some(1)) + .with_pruned_by_bloom_filter(Some(0)) + .with_expected_rows(1) // only "a" matches LIKE 'a%' + .test_row_group_prune() + .await; +} diff --git a/datafusion/core/tests/parquet/schema_coercion.rs b/datafusion/core/tests/parquet/schema_coercion.rs index 6f7e2e328d0c3..be45ab38dabad 100644 --- a/datafusion/core/tests/parquet/schema_coercion.rs +++ b/datafusion/core/tests/parquet/schema_coercion.rs @@ -53,7 +53,7 @@ async fn multi_parquet_coercion() { // batch2: c2(int64), c3(float32) let batch2 = RecordBatch::try_from_iter(vec![("c2", c2), ("c3", c3)]).unwrap(); - let (meta, _files) = store_parquet(vec![batch1, batch2]).await.unwrap(); + let (meta, _files) = store_parquet(vec![batch1, batch2]).unwrap(); let file_group = meta.into_iter().map(Into::into).collect(); // cast c1 to utf8, c2 to int32, c3 to float64 @@ -107,7 +107,7 @@ async fn multi_parquet_coercion_projection() { let batch2 = RecordBatch::try_from_iter(vec![("c2", c2), ("c1", c1s), ("c3", c3)]).unwrap(); - let (meta, _files) = store_parquet(vec![batch1, batch2]).await.unwrap(); + let (meta, _files) = store_parquet(vec![batch1, batch2]).unwrap(); let file_group = meta.into_iter().map(Into::into).collect(); // cast c1 to utf8, c2 to int32, c3 to float64 @@ -146,7 +146,7 @@ async fn multi_parquet_coercion_projection() { } /// Writes `batches` to a temporary parquet file -pub async fn store_parquet( +pub fn store_parquet( batches: Vec, ) -> Result<(Vec, Vec)> { // Each batch writes to their own file diff --git a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs index 808e163b08369..2d22b60856ca5 100644 --- a/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/aggregate_statistics.rs @@ -29,16 +29,17 @@ use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::ParquetSource; use datafusion::datasource::source::DataSourceExec; use datafusion::prelude::{SessionConfig, SessionContext}; -use datafusion_common::assert_batches_eq; use datafusion_common::cast::as_int64_array; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, Result, Statistics}; +use datafusion_common::{ScalarValue, assert_batches_eq}; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_execution::TaskContext; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::Operator; use datafusion_functions_aggregate::count::count_udaf; +use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::{self, cast}; use datafusion_physical_optimizer::PhysicalOptimizerRule; @@ -553,3 +554,305 @@ async fn test_count_distinct_optimization() -> Result<()> { Ok(()) } + +/// Regression test for https://github.com/apache/datafusion/issues/22554 +/// +/// TopK aggregation for DISTINCT queries was unconditionally dropping NULL +/// group keys, producing wrong results with NULLS FIRST / NULLS LAST ordering. +#[tokio::test] +async fn topk_distinct_preserves_nulls() -> Result<()> { + let ctx = SessionContext::new_with_config(SessionConfig::new()); + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("v", DataType::Utf8, true)])), + vec![Arc::new(StringArray::from(vec![None, Some(""), Some("a")]))], + )?; + let table = MemTable::try_new(batch.schema(), vec![vec![batch]])?; + ctx.register_table("t", Arc::new(table))?; + + // ASC NULLS FIRST LIMIT 1 → NULL should come first + let result = ctx + .sql("SELECT DISTINCT v FROM t ORDER BY v ASC NULLS FIRST LIMIT 1") + .await? + .collect() + .await?; + assert_batches_eq!(&["+---+", "| v |", "+---+", "| |", "+---+"], &result); + assert!(result[0].column(0).is_null(0), "first row should be NULL"); + + // ASC NULLS FIRST LIMIT 2 → NULL, then empty string + let result = ctx + .sql("SELECT DISTINCT v FROM t ORDER BY v ASC NULLS FIRST LIMIT 2") + .await? + .collect() + .await?; + assert_eq!(result[0].num_rows(), 2); + assert!(result[0].column(0).is_null(0)); + assert!(!result[0].column(0).is_null(1)); + + // ASC NULLS LAST LIMIT 1 → empty string (smallest non-null) + let result = ctx + .sql("SELECT DISTINCT v FROM t ORDER BY v ASC NULLS LAST LIMIT 1") + .await? + .collect() + .await?; + assert!( + !result[0].column(0).is_null(0), + "first row should NOT be NULL" + ); + + // Full result with NULLS LAST should include NULL at end + let result = ctx + .sql("SELECT DISTINCT v FROM t ORDER BY v ASC NULLS LAST LIMIT 3") + .await? + .collect() + .await?; + assert_eq!(result[0].num_rows(), 3); + assert!(result[0].column(0).is_null(2), "last row should be NULL"); + + // Integer column + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, true)])), + vec![Arc::new(Int64Array::from(vec![None, Some(3), Some(1)]))], + )?; + let table = MemTable::try_new(batch.schema(), vec![vec![batch]])?; + ctx.register_table("t_int", Arc::new(table))?; + + let result = ctx + .sql("SELECT DISTINCT v FROM t_int ORDER BY v ASC NULLS FIRST LIMIT 1") + .await? + .collect() + .await?; + assert!( + result[0].column(0).is_null(0), + "integer NULL should be first" + ); + + let result = ctx + .sql("SELECT DISTINCT v FROM t_int ORDER BY v DESC NULLS LAST LIMIT 2") + .await? + .collect() + .await?; + assert_eq!(result[0].num_rows(), 2); + assert!(!result[0].column(0).is_null(0)); + assert!(!result[0].column(0).is_null(1)); + + Ok(()) +} + +#[tokio::test] +async fn test_sum_from_statistics() -> Result<()> { + enum SumArg { + ColumnA, + ColumnB, + CastColumnA(DataType), + Binary, + } + + struct TestCase { + name: &'static str, + data_type: DataType, + sum_value_a: Precision, + sum_value_b: Precision, + sum_arg: SumArg, + is_distinct: bool, + expected_value: Option, + } + + for case in [ + TestCase { + name: "exact statistics", + data_type: DataType::Int64, + sum_value_a: Precision::Exact(ScalarValue::Int64(Some(10))), + sum_value_b: Precision::Absent, + sum_arg: SumArg::ColumnA, + is_distinct: false, + expected_value: Some(ScalarValue::Int64(Some(10))), + }, + TestCase { + name: "second column statistics", + data_type: DataType::Int64, + sum_value_a: Precision::Exact(ScalarValue::Int64(Some(10))), + sum_value_b: Precision::Exact(ScalarValue::Int64(Some(42))), + sum_arg: SumArg::ColumnB, + is_distinct: false, + expected_value: Some(ScalarValue::Int64(Some(42))), + }, + TestCase { + name: "casted int32 column statistics", + data_type: DataType::Int32, + sum_value_a: Precision::Exact(ScalarValue::Int32(Some(10))), + sum_value_b: Precision::Absent, + sum_arg: SumArg::CastColumnA(DataType::Int64), + is_distinct: false, + expected_value: Some(ScalarValue::Int64(Some(10))), + }, + TestCase { + name: "decimal statistics uses aggregate return type", + data_type: DataType::Decimal128(5, 2), + sum_value_a: Precision::Exact(ScalarValue::Decimal128(Some(12345), 5, 2)), + sum_value_b: Precision::Absent, + sum_arg: SumArg::ColumnA, + is_distinct: false, + expected_value: Some(ScalarValue::Decimal128(Some(12345), 15, 2)), + }, + TestCase { + name: "inexact statistics", + data_type: DataType::Int64, + sum_value_a: Precision::Inexact(ScalarValue::Int64(Some(10))), + sum_value_b: Precision::Absent, + sum_arg: SumArg::ColumnA, + is_distinct: false, + expected_value: None, + }, + TestCase { + name: "absent statistics", + data_type: DataType::Int64, + sum_value_a: Precision::Absent, + sum_value_b: Precision::Absent, + sum_arg: SumArg::ColumnA, + is_distinct: false, + expected_value: None, + }, + TestCase { + name: "null statistics", + data_type: DataType::Int64, + sum_value_a: Precision::Exact(ScalarValue::Int64(None)), + sum_value_b: Precision::Absent, + sum_arg: SumArg::ColumnA, + is_distinct: false, + expected_value: None, + }, + TestCase { + name: "binary expr", + data_type: DataType::Int64, + sum_value_a: Precision::Exact(ScalarValue::Int64(Some(10))), + sum_value_b: Precision::Exact(ScalarValue::Int64(Some(42))), + sum_arg: SumArg::Binary, + is_distinct: false, + expected_value: None, + }, + TestCase { + name: "distinct sum", + data_type: DataType::Int64, + sum_value_a: Precision::Exact(ScalarValue::Int64(Some(10))), + sum_value_b: Precision::Absent, + sum_arg: SumArg::ColumnA, + is_distinct: true, + expected_value: None, + }, + ] { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", case.data_type.clone(), true), + Field::new("b", case.data_type.clone(), true), + ])); + + let statistics = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: vec![ + ColumnStatistics { + sum_value: case.sum_value_a, + ..Default::default() + }, + ColumnStatistics { + sum_value: case.sum_value_b, + ..Default::default() + }, + ], + }; + + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::new(ParquetSource::new(Arc::clone(&schema))), + ) + .with_file(PartitionedFile::new("x".to_string(), 100)) + .with_statistics(statistics) + .build(); + + let source: Arc = DataSourceExec::from_data_source(config); + let schema = source.schema(); + + let (agg_args, alias): (Vec>, _) = + match case.sum_arg { + SumArg::ColumnA => (vec![expressions::col("a", &schema)?], "SUM(a)"), + SumArg::ColumnB => (vec![expressions::col("b", &schema)?], "SUM(b)"), + SumArg::CastColumnA(cast_type) => ( + vec![cast(expressions::col("a", &schema)?, &schema, cast_type)?], + "SUM(CAST(a))", + ), + SumArg::Binary => ( + vec![expressions::binary( + expressions::col("a", &schema)?, + Operator::Plus, + expressions::col("b", &schema)?, + &schema, + )?], + "SUM(a + b)", + ), + }; + + let sum_expr_builder = AggregateExprBuilder::new(sum_udaf(), agg_args) + .schema(Arc::clone(&schema)) + .alias(alias); + let sum_expr_builder = if case.is_distinct { + sum_expr_builder.distinct() + } else { + sum_expr_builder + }; + let sum_expr = sum_expr_builder.build()?; + + let partial_agg = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::default(), + vec![Arc::new(sum_expr.clone())], + vec![None], + source, + Arc::clone(&schema), + )?; + + let final_agg = AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::default(), + vec![Arc::new(sum_expr)], + vec![None], + Arc::new(partial_agg), + Arc::clone(&schema), + )?; + + let conf = ConfigOptions::new(); + let optimized = + AggregateStatistics::new().optimize(Arc::new(final_agg), &conf)?; + + if let Some(expected_value) = case.expected_value { + assert!( + optimized.is::(), + "'{}': expected ProjectionExec", + case.name + ); + + let task_ctx = Arc::new(TaskContext::default()); + let result = common::collect(optimized.execute(0, task_ctx)?).await?; + assert_eq!(result.len(), 1, "'{}': expected 1 batch", case.name); + assert_eq!( + result[0].schema().field(0).data_type(), + &expected_value.data_type(), + "'{}': unexpected data type", + case.name + ); + assert_eq!( + ScalarValue::try_from_array(result[0].column(0), 0)?, + expected_value, + "'{}': unexpected value", + case.name + ); + } else { + assert!( + optimized.is::(), + "'{}': expected AggregateExec (not optimized)", + case.name + ); + } + } + + Ok(()) +} diff --git a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs index 78bb02ab1108b..2dbacf1d898ac 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_distribution.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_distribution.rs @@ -20,10 +20,10 @@ use std::ops::Deref; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - check_integrity, coalesce_partitions_exec, parquet_exec_with_sort, - parquet_exec_with_stats, repartition_exec, schema, sort_exec, - sort_exec_with_preserve_partitioning, sort_merge_join_exec, - sort_preserving_merge_exec, union_exec, + RequirementsTestExec, bounded_window_exec_with_can_repartition, check_integrity, + coalesce_partitions_exec, parquet_exec_with_sort, parquet_exec_with_stats, + repartition_exec, schema, sort_exec, sort_exec_with_preserve_partitioning, + sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; use arrow::array::{RecordBatch, UInt8Array, UInt64Array}; @@ -55,13 +55,15 @@ use datafusion_physical_expr_common::sort_expr::{ }; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::enforce_distribution::*; -use datafusion_physical_optimizer::enforce_sorting::EnforceSorting; +use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; use datafusion_physical_optimizer::output_requirements::OutputRequirements; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; -use datafusion_physical_expr::Distribution; +use datafusion_physical_expr::{ + Distribution, Partitioning, RangePartitioning, SplitPoint, +}; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::execution_plan::ExecutionPlan; use datafusion_physical_plan::expressions::col; @@ -72,7 +74,8 @@ use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlanProperties, PlanProperties, displayable, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, + PlanProperties, ReplaceChildrenOptions, displayable, }; use insta::Settings; @@ -191,9 +194,10 @@ impl ExecutionPlan for SortRequiredExec { vec![Some(OrderingRequirements::from(self.expr.clone()))] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_eq!(children.len(), 1); let child = children.pop().unwrap(); @@ -203,18 +207,21 @@ impl ExecutionPlan for SortRequiredExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } fn execute( @@ -279,8 +286,12 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { vec![&self.input] } - fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + ]) } fn maintains_input_order(&self) -> Vec { @@ -291,18 +302,29 @@ impl ExecutionPlan for SinglePartitionMaintainsOrderExec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_eq!(children.len(), 1); let child = children.pop().unwrap(); Ok(Arc::new(Self::new(child))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -348,6 +370,47 @@ fn parquet_exec_multiple_sorted( DataSourceExec::from_data_source(config) } +fn parquet_exec_with_output_partitioning( + output_partitioning: Partitioning, +) -> Arc { + let file_groups = (0..output_partitioning.partition_count()) + .map(|partition| { + FileGroup::new(vec![PartitionedFile::new(format!("p{partition}"), 100)]) + }) + .collect::>(); + + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::new(ParquetSource::new(schema())), + ) + .with_file_groups(file_groups) + .with_output_partitioning(Some(output_partitioning)) + .build(); + + DataSourceExec::from_data_source(config) +} + +fn range_partitioning( + column: &str, + split_values: impl IntoIterator, + options: SortOptions, +) -> Result { + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col(column, &schema())?, + options, + }] + .into(); + let split_points = split_values + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int64(Some(value))])) + .collect::>(); + + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) +} + fn csv_exec() -> Arc { csv_exec_with_sort(vec![]) } @@ -593,7 +656,9 @@ fn test_suite_default_config_options() -> ConfigOptions { config.execution.target_partitions = 10; // Use a small batch size, to trigger RoundRobin in tests - config.execution.batch_size = 1; + config.execution.batch_size = + datafusion_common::config::ConfigNonZeroUsize::try_new(1) + .expect("test batch size must be greater than zero"); config } @@ -695,17 +760,14 @@ impl TestConfig { // TODO: End state payloads will be checked here. } - for run in optimizers_to_run { - optimized = match run { - Run::Distribution => { - let optimizer = EnforceDistribution::new(); - optimizer.optimize(optimized, &self.config)? - } - Run::Sorting => { - let optimizer = EnforceSorting::new(); - optimizer.optimize(optimized, &self.config)? - } - }; + // With `EnsureRequirements`, distribution and sorting enforcement are + // composed into a single idempotent pass, so the historical sequence + // of `Run::Distribution` / `Run::Sorting` collapses to repeated calls + // of the same rule. The sequences are preserved so existing test + // assertions (which encode legacy run orders) remain stable. + for _run in optimizers_to_run { + let optimizer = EnsureRequirements::new(); + optimized = optimizer.optimize(optimized, &self.config)?; } // Remove the ancillary output requirements operator when done: @@ -724,6 +786,509 @@ impl TestConfig { } } +#[derive(Debug, Clone, Copy)] +enum ExpectedPlan { + Reuse, + Hash, +} + +#[test] +fn range_satisfaction_config_matrix() -> Result<()> { + const INPUT_PARTITIONS: usize = 4; + const MET: usize = INPUT_PARTITIONS; + const NOT_MET: usize = INPUT_PARTITIONS + 1; + const DISABLED: usize = 0; + const EQUAL: usize = INPUT_PARTITIONS; + const GREATER: usize = INPUT_PARTITIONS + 1; + use ExpectedPlan::{Hash, Reuse}; + + let config_cases = [ + // subset preserve target exact subset incompatible + (NOT_MET, DISABLED, EQUAL, [Reuse, Hash, Hash]), + (NOT_MET, DISABLED, GREATER, [Hash, Hash, Hash]), + (NOT_MET, NOT_MET, EQUAL, [Reuse, Hash, Hash]), + (NOT_MET, NOT_MET, GREATER, [Hash, Hash, Hash]), + (NOT_MET, MET, EQUAL, [Reuse, Hash, Hash]), + (NOT_MET, MET, GREATER, [Reuse, Reuse, Hash]), + (MET, DISABLED, EQUAL, [Reuse, Reuse, Hash]), + (MET, DISABLED, GREATER, [Reuse, Reuse, Hash]), + (MET, NOT_MET, EQUAL, [Reuse, Reuse, Hash]), + (MET, NOT_MET, GREATER, [Reuse, Reuse, Hash]), + (MET, MET, EQUAL, [Reuse, Reuse, Hash]), + (MET, MET, GREATER, [Reuse, Reuse, Hash]), + ]; + for (subset_threshold, preserve_file_partitions, target_partitions, expected) in + config_cases + { + let key_cases = [ + ("exact", vec![col("a", &schema())?], expected[0]), + ( + "subset", + vec![col("a", &schema())?, col("b", &schema())?], + expected[1], + ), + ("incompatible", vec![col("b", &schema())?], expected[2]), + ]; + for (key_match, partition_keys, expected_plan) in key_cases { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let requirement = RequirementsTestExec::new(input) + .with_required_input_distribution(Distribution::KeyPartitioned( + partition_keys, + )) + .into_arc(); + + let mut config = + TestConfig::default().with_query_execution_partitions(target_partitions); + config.config.optimizer.subset_repartition_threshold = subset_threshold; + config.config.optimizer.preserve_file_partitions = preserve_file_partitions; + + let plan = config.to_plan(requirement, &DISTRIB_DISTRIB_SORT); + let plan = displayable(plan.as_ref()).indent(true).to_string(); + let repartitions = plan + .lines() + .filter(|line| line.contains("RepartitionExec:")) + .collect::>(); + + let matches_expected = match expected_plan { + Reuse => repartitions.is_empty(), + Hash => matches!( + repartitions.as_slice(), + [repartition] if repartition.contains("partitioning=Hash") + ), + }; + assert!( + matches_expected, + "unexpected optimized plan for key_match={key_match}, \ + subset_threshold={subset_threshold}, \ + preserve_file_partitions={preserve_file_partitions}, \ + target_partitions={target_partitions}:\n{plan}" + ); + } + } + + Ok(()) +} + +#[test] +fn range_aggregate_reuses_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let aggregate = + aggregate_exec_with_alias(input, vec![("a".to_string(), "a".to_string())]); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(aggregate, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[] + AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_grouping_set_aggregate_rehashes_with_grouping_id() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20], + SortOptions::default(), + )?); + let input_schema = input.schema(); + let group_by = PhysicalGroupBy::new( + vec![ + (col("a", &input_schema)?, "a".to_string()), + (col("b", &input_schema)?, "b".to_string()), + ], + vec![ + (lit(ScalarValue::Int64(None)), "a".to_string()), + (lit(ScalarValue::Int64(None)), "b".to_string()), + ], + vec![vec![false, true], vec![false, false]], + true, + ); + let partial = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + input, + Arc::clone(&input_schema), + )?); + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + group_by.as_final(), + vec![], + vec![], + Arc::clone(&partial) as _, + partial.schema(), + )?); + + let plan = TestConfig::default() + .with_query_execution_partitions(3) + .to_plan(aggregate, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b, __grouping_id@2 as __grouping_id], aggr=[] + RepartitionExec: partitioning=Hash([a@0, b@1, __grouping_id@2], 3), input_partitions=3 + AggregateExec: mode=Partial, gby=[(a@0 as a, NULL as b), (a@0 as a, b@1 as b)], aggr=[] + DataSourceExec: file_groups={3 groups: [[p0], [p1], [p2]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20)], 3), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_inner_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 30, 40], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Inner); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a1@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + RepartitionExec: partitioning=Hash([a1@0], 4), input_partitions=4 + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_right_mark_hash_join_reuses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::RightMark); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=RightMark, on=[(a@0, a@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_right_semi_hash_join_rehashes_incompatible_sort_options() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [20], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [20], + SortOptions { + descending: true, + nulls_first: true, + }, + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::RightSemi); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(a@0, a@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 + DataSourceExec: file_groups={2 groups: [[p0], [p1]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(20)], 2), file_type=parquet + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=2 + DataSourceExec: file_groups={2 groups: [[p0], [p1]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 DESC], [(20)], 2), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_window_reuses_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let window = bounded_window_exec_with_can_repartition( + "a", + vec![], + &[col("a", &schema())?], + input, + true, + ); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(window, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + "# + ); + + Ok(()) +} + +#[test] +fn range_window_rehashes_incompatible_range_partitioning() -> Result<()> { + let input = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let window = bounded_window_exec_with_can_repartition( + "b", + vec![], + &[col("b", &schema())?], + input, + true, + ); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(window, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([b@1], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + "# + ); + + Ok(()) +} + +#[test] +fn range_full_hash_join_reuses_compatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Full); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_full_hash_join_rehashes_incompatible_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = projection_exec_with_alias( + parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 30, 40], + SortOptions::default(), + )?), + vec![ + ("a".to_string(), "a1".to_string()), + ("b".to_string(), "b1".to_string()), + ], + ); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a1", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::Full); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=Full, on=[(a@0, a1@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + RepartitionExec: partitioning=Hash([a1@0], 4), input_partitions=4 + ProjectionExec: expr=[a@0 as a1, b@1 as b1] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (30), (40)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_left_mark_hash_join_reuses_range_partitioning() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftMark); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(a@0, a@0)] + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + +#[test] +fn range_left_anti_hash_join_rehashes_incompatible_null_options() -> Result<()> { + let left = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions::default(), + )?); + let right = parquet_exec_with_output_partitioning(range_partitioning( + "a", + [10, 20, 30], + SortOptions { + descending: false, + nulls_first: false, + }, + )?); + let join_on = vec![( + Arc::new(Column::new_with_schema("a", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a", &right.schema())?) as _, + )]; + let join = hash_join_exec(left, right, &join_on, &JoinType::LeftAnti); + + let plan = TestConfig::default() + .with_query_execution_partitions(4) + .to_plan(join, &DISTRIB_DISTRIB_SORT); + + assert_plan!( + plan, + @r" + HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(a@0, a@0)] + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), file_type=parquet + RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=4 + DataSourceExec: file_groups={4 groups: [[p0], [p1], [p2], [p3]]}, projection=[a, b, c, d, e], output_partitioning=Range([a@0 ASC NULLS LAST], [(10), (20), (30)], 4), file_type=parquet + " + ); + + Ok(()) +} + #[test] fn multi_hash_joins() -> Result<()> { let left = parquet_exec(); @@ -1597,15 +2162,15 @@ fn multi_smj_joins() -> Result<()> { assert_plan!(plan_distrib, @r" SortMergeJoinExec: join_type=..., on=[(a@0, c@2)] SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)] - SortExec: expr=[a@0 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1 + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[b1@1 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1 - ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true + ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[c@2 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1 + RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[c@2 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet "); } @@ -1621,18 +2186,18 @@ fn multi_smj_joins() -> Result<()> { _ => { assert_plan!(plan_distrib, @r" SortMergeJoinExec: join_type=..., on=[(a@0, c@2)] - SortExec: expr=[a@0 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10 + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC + SortExec: expr=[a@0 ASC], preserve_partitioning=[true] SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)] - SortExec: expr=[a@0 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1 + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[b1@1 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1 - ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true + ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[c@2 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1 + RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[c@2 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet "); } @@ -1651,8 +2216,8 @@ fn multi_smj_joins() -> Result<()> { SortExec: expr=[a@0 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[b1@1 ASC], preserve_partitioning=[false] - ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true SortExec: expr=[c@2 ASC], preserve_partitioning=[false] @@ -1672,17 +2237,16 @@ fn multi_smj_joins() -> Result<()> { // TODO(wiedld): show different test result if enforce distribution first. assert_plan!(plan_sort, @r" SortMergeJoinExec: join_type=..., on=[(a@0, c@2)] - RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[a@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)] - RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC + SortExec: expr=[a@0 ASC], preserve_partitioning=[true] + SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true + ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[b1@1 ASC], preserve_partitioning=[false] - ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true SortExec: expr=[c@2 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet @@ -1709,15 +2273,15 @@ fn multi_smj_joins() -> Result<()> { assert_plan!(plan_distrib, @r" SortMergeJoinExec: join_type=..., on=[(b1@6, c@2)] SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)] - SortExec: expr=[a@0 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1 + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[b1@1 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1 - ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true + ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[c@2 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1 + RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[c@2 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet "); } @@ -1726,18 +2290,18 @@ fn multi_smj_joins() -> Result<()> { // TODO(wiedld): show different test result if enforce sorting first. assert_plan!(plan_distrib, @r" SortMergeJoinExec: join_type=..., on=[(b1@6, c@2)] - SortExec: expr=[b1@6 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10 + RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@6 ASC + SortExec: expr=[b1@6 ASC], preserve_partitioning=[true] SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)] - SortExec: expr=[a@0 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1 + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[b1@1 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1 - ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true + ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[c@2 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1 + RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[c@2 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet "); } @@ -1758,8 +2322,8 @@ fn multi_smj_joins() -> Result<()> { SortExec: expr=[a@0 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[b1@1 ASC], preserve_partitioning=[false] - ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true SortExec: expr=[c@2 ASC], preserve_partitioning=[false] @@ -1771,17 +2335,16 @@ fn multi_smj_joins() -> Result<()> { // TODO(wiedld): show different test result if enforce distribution first. assert_plan!(plan_sort, @r" SortMergeJoinExec: join_type=..., on=[(b1@6, c@2)] - RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[b1@6 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)] - RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + RepartitionExec: partitioning=Hash([b1@6], 10), input_partitions=10, preserve_order=true, sort_exprs=b1@6 ASC + SortExec: expr=[b1@6 ASC], preserve_partitioning=[true] + SortMergeJoinExec: join_type=..., on=[(a@0, b1@1)] + RepartitionExec: partitioning=Hash([a@0], 10), input_partitions=1, maintains_sort_order=true + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true + ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - RepartitionExec: partitioning=Hash([b1@1], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[b1@1 ASC], preserve_partitioning=[false] - ProjectionExec: expr=[a@0 as a1, b@1 as b1, c@2 as c1, d@3 as d1, e@4 as e1] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true SortExec: expr=[c@2 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet @@ -1854,16 +2417,16 @@ fn smj_join_key_ordering() -> Result<()> { let plan_distrib = test_config.to_plan(join.clone(), &DISTRIB_DISTRIB_SORT); assert_plan!(plan_distrib, @r" SortMergeJoinExec: join_type=Inner, on=[(b3@1, b2@1), (a3@0, a2@0)] - SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[true] - ProjectionExec: expr=[a1@0 as a3, b1@1 as b3] - ProjectionExec: expr=[a1@1 as a1, b1@0 as b1] + ProjectionExec: expr=[a1@0 as a3, b1@1 as b3] + ProjectionExec: expr=[a1@1 as a1, b1@0 as b1] + SortExec: expr=[b1@0 ASC, a1@1 ASC], preserve_partitioning=[true] AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[] RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10 AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[] RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[true] - ProjectionExec: expr=[a@1 as a2, b@0 as b2] + ProjectionExec: expr=[a@1 as a2, b@0 as b2] + SortExec: expr=[b@0 ASC, a@1 ASC], preserve_partitioning=[true] AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[] RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10 AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[] @@ -1875,25 +2438,21 @@ fn smj_join_key_ordering() -> Result<()> { let plan_sort = test_config.to_plan(join, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" SortMergeJoinExec: join_type=Inner, on=[(b3@1, b2@1), (a3@0, a2@0)] - RepartitionExec: partitioning=Hash([b3@1, a3@0], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[b3@1 ASC, a3@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - ProjectionExec: expr=[a1@0 as a3, b1@1 as b3] - ProjectionExec: expr=[a1@1 as a1, b1@0 as b1] - AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[] - RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10 - AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[] - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet - RepartitionExec: partitioning=Hash([b2@1, a2@0], 10), input_partitions=1, maintains_sort_order=true - SortExec: expr=[b2@1 ASC, a2@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - ProjectionExec: expr=[a@1 as a2, b@0 as b2] - AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[] - RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10 - AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[] - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + ProjectionExec: expr=[a1@0 as a3, b1@1 as b3] + ProjectionExec: expr=[a1@1 as a1, b1@0 as b1] + SortExec: expr=[b1@0 ASC, a1@1 ASC], preserve_partitioning=[true] + AggregateExec: mode=FinalPartitioned, gby=[b1@0 as b1, a1@1 as a1], aggr=[] + RepartitionExec: partitioning=Hash([b1@0, a1@1], 10), input_partitions=10 + AggregateExec: mode=Partial, gby=[b@1 as b1, a@0 as a1], aggr=[] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + ProjectionExec: expr=[a@1 as a2, b@0 as b2] + SortExec: expr=[b@0 ASC, a@1 ASC], preserve_partitioning=[true] + AggregateExec: mode=FinalPartitioned, gby=[b@0 as b, a@1 as a], aggr=[] + RepartitionExec: partitioning=Hash([b@0, a@1], 10), input_partitions=10 + AggregateExec: mode=Partial, gby=[b@1 as b, a@0 as a], aggr=[] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet "); Ok(()) @@ -1937,9 +2496,8 @@ fn merge_does_not_need_sort() -> Result<()> { let plan_sort = test_config.to_plan(exec, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" - SortExec: expr=[a@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + SortPreservingMergeExec: [a@0 ASC] + DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet "); Ok(()) @@ -2269,9 +2827,8 @@ fn repartition_ignores_sort_preserving_merge() -> Result<()> { let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" - SortExec: expr=[c@2 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet + SortPreservingMergeExec: [c@2 ASC] + DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet "); Ok(()) @@ -2309,11 +2866,10 @@ fn repartition_ignores_sort_preserving_merge_with_union() -> Result<()> { let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" - SortExec: expr=[c@2 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - UnionExec - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet + SortPreservingMergeExec: [c@2 ASC] + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet "); Ok(()) @@ -2435,8 +2991,8 @@ fn repartition_transitively_with_projection() -> Result<()> { let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" - SortExec: expr=[sum@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec + SortPreservingMergeExec: [sum@0 ASC] + SortExec: expr=[sum@0 ASC], preserve_partitioning=[true] ProjectionExec: expr=[a@0 + b@1 as sum] RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet @@ -2509,8 +3065,8 @@ fn repartition_transitively_past_sort_with_projection() -> Result<()> { let plan_distrib = test_config.to_plan(plan.clone(), &DISTRIB_DISTRIB_SORT); assert_plan!(plan_distrib, @r" - SortExec: expr=[c@2 ASC], preserve_partitioning=[false] - ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c] + ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c] + SortExec: expr=[c@2 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet "); // Since this projection is trivial, increasing parallelism is not beneficial @@ -2547,8 +3103,8 @@ fn repartition_transitively_past_sort_with_filter() -> Result<()> { let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" - SortExec: expr=[a@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec + SortPreservingMergeExec: [a@0 ASC] + SortExec: expr=[a@0 ASC], preserve_partitioning=[true] FilterExec: c@2 = 0 RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet @@ -2584,8 +3140,8 @@ fn repartition_transitively_past_sort_with_projection_and_filter() -> Result<()> assert_plan!(plan_distrib, @r" SortPreservingMergeExec: [a@0 ASC] - SortExec: expr=[a@0 ASC], preserve_partitioning=[true] - ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c] + ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c] + SortExec: expr=[a@0 ASC], preserve_partitioning=[true] FilterExec: c@2 = 0 RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet @@ -2598,9 +3154,9 @@ fn repartition_transitively_past_sort_with_projection_and_filter() -> Result<()> let plan_sort = test_config.to_plan(plan, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" - SortExec: expr=[a@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c] + SortPreservingMergeExec: [a@0 ASC] + ProjectionExec: expr=[a@0 as a, b@1 as b, c@2 as c] + SortExec: expr=[a@0 ASC], preserve_partitioning=[true] FilterExec: c@2 = 0 RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet @@ -3117,11 +3673,10 @@ fn parallelization_sort_preserving_merge_with_union() -> Result<()> { let plan_parquet_sort = test_config.to_plan(plan_parquet, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_parquet_sort, @r" - SortExec: expr=[c@2 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - UnionExec - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet + SortPreservingMergeExec: [c@2 ASC] + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet "); // no SPM // has coalesce @@ -3138,11 +3693,10 @@ fn parallelization_sort_preserving_merge_with_union() -> Result<()> { let plan_csv_sort = test_config.to_plan(plan_csv.clone(), &SORT_DISTRIB_DISTRIB); assert_plan!(plan_csv_sort, @r" - SortExec: expr=[c@2 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - UnionExec - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false + SortPreservingMergeExec: [c@2 ASC] + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=csv, has_header=false "); // no SPM // has coalesce @@ -3474,8 +4028,8 @@ fn do_not_preserve_ordering_through_repartition() -> Result<()> { let plan_sort = test_config.to_plan(physical_plan, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" - SortExec: expr=[a@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec + SortPreservingMergeExec: [a@0 ASC] + SortExec: expr=[a@0 ASC], preserve_partitioning=[true] FilterExec: c@2 = 0 RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2 DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet @@ -3545,12 +4099,11 @@ fn do_not_preserve_ordering_through_repartition2() -> Result<()> { let plan_sort = test_config.to_plan(physical_plan, &SORT_DISTRIB_DISTRIB); assert_plan!(plan_sort, @r" - SortExec: expr=[a@0 ASC], preserve_partitioning=[false] - CoalescePartitionsExec - SortExec: expr=[a@0 ASC], preserve_partitioning=[true] - FilterExec: c@2 = 0 - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2 - DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet + SortPreservingMergeExec: [a@0 ASC] + SortExec: expr=[a@0 ASC], preserve_partitioning=[true] + FilterExec: c@2 = 0 + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2 + DataSourceExec: file_groups={2 groups: [[x], [y]]}, projection=[a, b, c, d, e], output_ordering=[c@2 ASC], file_type=parquet "); Ok(()) @@ -3603,14 +4156,15 @@ fn do_not_put_sort_when_input_is_invalid() -> Result<()> { config.execution.target_partitions = 10; config.optimizer.enable_round_robin_repartition = true; config.optimizer.prefer_existing_sort = false; - let dist_plan = EnforceDistribution::new().optimize(physical_plan, &config)?; + let dist_plan = EnsureRequirements::new().optimize(physical_plan, &config)?; // Since at the start of the rule ordering requirement is not satisfied // EnforceDistribution rule doesn't satisfy this requirement either. assert_plan!(dist_plan, @r" SortRequiredExec: [a@0 ASC] - FilterExec: c@2 = 0 - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + SortExec: expr=[a@0 ASC], preserve_partitioning=[true] + FilterExec: c@2 = 0 + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet "); Ok(()) @@ -3639,7 +4193,7 @@ fn put_sort_when_input_is_valid() -> Result<()> { config.execution.target_partitions = 10; config.optimizer.enable_round_robin_repartition = true; config.optimizer.prefer_existing_sort = false; - let dist_plan = EnforceDistribution::new().optimize(physical_plan, &config)?; + let dist_plan = EnsureRequirements::new().optimize(physical_plan, &config)?; // Since at the start of the rule ordering requirement is satisfied // EnforceDistribution rule satisfy this requirement also. assert_plan!(dist_plan, @r" @@ -3792,8 +4346,8 @@ async fn test_distribute_sort_parquet() -> Result<()> { test_config.to_plan(physical_plan.clone(), &[Run::Distribution]); assert_plan!(plan_distribution, @r" - SortExec: expr=[c@2 ASC], preserve_partitioning=[false] - CoalescePartitionsExec + SortPreservingMergeExec: [c@2 ASC] + SortExec: expr=[c@2 ASC], preserve_partitioning=[true] DataSourceExec: file_groups={10 groups: [[x:0..8192000], [x:8192000..16384000], [x:16384000..24576000], [x:24576000..32768000], [x:32768000..40960000], [x:40960000..49152000], [x:49152000..57344000], [x:57344000..65536000], [x:65536000..73728000], [x:73728000..81920000]]}, projection=[a, b, c, d, e], file_type=parquet "); @@ -4008,3 +4562,38 @@ fn adjust_input_keys_ordering_no_transform_for_filter_scan() -> Result<()> { ); Ok(()) } + +/// Verifies the `ensure_distribution` fast path: when no child of a node is +/// replaced (no `RepartitionExec` or `SortExec` injection is required), +/// the rule must reuse the input `Arc` unchanged instead +/// of calling `with_new_children`. For a deep `ProjectionExec` chain over a +/// single-partition scan with `target_partitions = 1`, every node hits this +/// fast path, so the root returned by `ensure_distribution` must be the +/// same `Arc` as the input. +/// +/// Regression test for the optimization that avoids +/// `ProjectionExec::with_new_children` (which recomputes schema, equivalence +/// properties, output ordering, and partitioning) on the common point-query +/// plan shape. +#[test] +fn ensure_distribution_reuses_plan_arc_when_no_redistribution_needed() -> Result<()> { + let scan = parquet_exec(); + let proj1 = projection_exec_with_alias( + scan, + vec![ + ("a".to_string(), "a".to_string()), + ("b".to_string(), "b".to_string()), + ], + ); + let proj2 = + projection_exec_with_alias(proj1, vec![("a".to_string(), "a".to_string())]); + let plan: Arc = proj2; + + let result = ensure_distribution_helper(Arc::clone(&plan), 1, false)?; + + assert!( + Arc::ptr_eq(&result, &plan), + "ensure_distribution must reuse the input Arc when no children require redistribution" + ); + Ok(()) +} diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs index 40bcdbbd6efef..a8162f137ed0a 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting.rs @@ -33,8 +33,8 @@ use arrow::compute::{SortOptions}; use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::config::{ConfigOptions, CsvOptions}; use datafusion_common::tree_node::{TreeNode, TransformedResult}; -use datafusion_common::{create_array, Result, TableReference}; -use datafusion_datasource::file_scan_config::FileScanConfigBuilder; +use datafusion_common::{create_array, DataFusionError, NullEquality, Result, TableReference}; +use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::source::DataSourceExec; use datafusion_expr_common::operator::Operator; use datafusion_expr::{JoinType, SortExpr}; @@ -42,28 +42,35 @@ use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_expr_common::sort_expr::{ LexOrdering, PhysicalSortExpr, PhysicalSortRequirement, OrderingRequirements }; -use datafusion_physical_expr::{Distribution, Partitioning}; +use datafusion_physical_expr::{Distribution, Partitioning, PhysicalExpr}; use datafusion_physical_expr::expressions::{col, BinaryExpr, Column, NotExpr}; +use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::sorts::sort::SortExec; -use datafusion_physical_plan::{displayable, get_plan_string, ExecutionPlan}; +use datafusion_physical_plan::{displayable, get_plan_string, ExecutionPlan, ExecutionPlanProperties}; use datafusion::datasource::physical_plan::CsvSource; use datafusion::datasource::listing::PartitionedFile; -use datafusion_physical_optimizer::enforce_sorting::{EnforceSorting, PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, parallelize_sorts, ensure_sorting}; +use datafusion_physical_optimizer::enforce_sorting::{PlanWithCorrespondingCoalescePartitions, PlanWithCorrespondingSort, parallelize_sorts, ensure_sorting}; +use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; use datafusion_physical_optimizer::enforce_sorting::replace_with_order_preserving_variants::{replace_with_order_preserving_variants, OrderPreservationContext}; use datafusion_physical_optimizer::enforce_sorting::sort_pushdown::{SortPushDown, assign_initial_requirements, pushdown_sorts}; -use datafusion_physical_optimizer::enforce_distribution::EnforceDistribution; +use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; +use datafusion_physical_optimizer::limit_pushdown::LimitPushdown; +use datafusion_physical_optimizer::projection_pushdown::ProjectionPushdown; use datafusion_physical_optimizer::output_requirements::OutputRequirementExec; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion::prelude::*; -use arrow::array::{record_batch, ArrayRef, Int32Array, RecordBatch}; +use arrow::array::{record_batch, Array, ArrayRef, Int32Array, RecordBatch}; use arrow::datatypes::{Field}; use arrow_schema::Schema; use datafusion_execution::TaskContext; use datafusion_catalog::streaming::StreamingTable; +use datafusion_expr_common::columnar_value::ColumnarValue; +use datafusion_physical_expr::projection::ProjectionExpr; +use datafusion_physical_plan::projection::ProjectionExec; use futures::StreamExt; use insta::{Settings, assert_snapshot}; @@ -117,11 +124,19 @@ impl EnforceSortingTest { pub(crate) fn run(&self) -> String { let mut config = ConfigOptions::new(); config.optimizer.repartition_sorts = self.repartition_sorts; - - // This file has 4 rules that use tree node, apply these rules as in the - // EnforceSorting::optimize implementation - // After these operations tree nodes should be in a consistent state. - // This code block makes sure that these rules doesn't violate tree node integrity. + // Pin target_partitions so snapshots stay deterministic across + // machines with different CPU counts. Now that the underlying + // optimizer is `EnsureRequirements` (which performs distribution + // enforcement), the partition count appears in `Hash([…], N)` + // nodes in the output plan; without pinning, snapshots taken on + // an N-core machine fail on an M-core machine. 10 matches the + // existing convention in `enforce_distribution.rs`. + config.execution.target_partitions = 10; + + // This file has 4 sub-rules that use tree node; apply them in the same + // order EnsureRequirements does internally. After these operations the + // tree nodes should be in a consistent state; this block exists to make + // sure those sub-rules don't violate tree node integrity. { let plan_requirements = PlanWithCorrespondingSort::new_default(Arc::clone(&self.plan)); @@ -175,9 +190,9 @@ impl EnforceSortingTest { let input_plan_string = displayable(self.plan.as_ref()).indent(true).to_string(); // Run the actual optimizer - let optimized_physical_plan = EnforceSorting::new() + let optimized_physical_plan = EnsureRequirements::new() .optimize(Arc::clone(&self.plan), &config) - .expect("enforce_sorting failed"); + .expect("ensure_requirements failed"); // Get string representation of the plan let optimized_plan_string = displayable(optimized_physical_plan.as_ref()) @@ -218,9 +233,54 @@ async fn test_remove_unnecessary_sort5() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet Optimized Plan: - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(col_a@0, c@2)] - DataSourceExec: partitions=1, partition_sizes=[0] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + SortPreservingMergeExec: [a@2 ASC] + HashJoinExec: mode=Partitioned, join_type=Inner, on=[(col_a@0, c@2)] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + "); + Ok(()) +} + +#[tokio::test] +async fn test_hash_join_interleaved_projection_preserves_parent_sort() -> Result<()> { + let left_schema = create_test_schema()?; + let right_schema = create_test_schema2()?; + let left = parquet_exec(left_schema.clone()); + let right = parquet_exec(right_schema.clone()); + let on = vec![( + Arc::new(Column::new_with_schema("nullable_col", &left_schema)?) as _, + Arc::new(Column::new_with_schema("col_a", &right_schema)?) as _, + )]; + let join = Arc::new(HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Right, + // Interleave a right-side column before a left-side column. + Some(vec![2, 0]), + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?); + let ordering = [sort_expr("nullable_col", &join.schema())].into(); + let physical_plan = sort_exec(ordering, join); + + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 10; + let optimized_plan = + EnsureRequirements::new().optimize(Arc::clone(&physical_plan), &config)?; + let optimized_plan = SanityCheckPlan::new().optimize(optimized_plan, &config)?; + + assert_snapshot!(displayable(optimized_plan.as_ref()).indent(true), @r" + SortPreservingMergeExec: [nullable_col@1 ASC] + SortExec: expr=[nullable_col@1 ASC], preserve_partitioning=[true] + HashJoinExec: mode=CollectLeft, join_type=Right, on=[(nullable_col@0, col_a@0)], projection=[col_a@2, nullable_col@0] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet "); Ok(()) } @@ -255,14 +315,11 @@ async fn test_do_not_remove_sort_with_limit() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet Optimized Plan: - SortPreservingMergeExec: [nullable_col@0 ASC, non_nullable_col@1 ASC] - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=2 - UnionExec - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet - LocalLimitExec: fetch=100 - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet + LocalLimitExec: fetch=100 + SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); // We should keep the bottom `SortExec`. Ok(()) @@ -282,12 +339,18 @@ async fn test_union_inputs_sorted() -> Result<()> { let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true); assert_snapshot!(test.run(), @r" - Input / Optimized Plan: + Input Plan: SortPreservingMergeExec: [nullable_col@0 ASC] UnionExec DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + + Optimized Plan: + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); // should not add a sort at the output of the union, input plan should not be changed @@ -313,12 +376,18 @@ async fn test_union_inputs_different_sorted() -> Result<()> { let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true); assert_snapshot!(test.run(), @r" - Input / Optimized Plan: + Input Plan: SortPreservingMergeExec: [nullable_col@0 ASC] UnionExec DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC], file_type=parquet SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + + Optimized Plan: + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC], file_type=parquet + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); // should not add a sort at the output of the union, input plan should not be changed @@ -353,23 +422,20 @@ async fn test_union_inputs_different_sorted2() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet Optimized Plan: - SortPreservingMergeExec: [nullable_col@0 ASC, non_nullable_col@1 ASC] - UnionExec - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); Ok(()) } -#[tokio::test] +#[test] // Test with `repartition_sorts` enabled to preserve pre-sorted partitions and avoid resorting -async fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_with_repartition_sorts_true() +fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_with_repartition_sorts_true() -> Result<()> { assert_snapshot!( - union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl(true).await?, + union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl(true)?, @r" Input Plan: OutputRequirementExec: order_by=[(nullable_col@0, asc)], dist_by=SinglePartition @@ -390,12 +456,12 @@ async fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_with_reparti Ok(()) } -#[tokio::test] +#[test] // Test with `repartition_sorts` disabled, causing a full resort of the data -async fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_with_repartition_sorts_false() +fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_with_repartition_sorts_false() -> Result<()> { assert_snapshot!( - union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl(false).await?, + union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl(false)?, @r" Input Plan: OutputRequirementExec: order_by=[(nullable_col@0, asc)], dist_by=SinglePartition @@ -416,7 +482,7 @@ async fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_with_reparti Ok(()) } -async fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl( +fn union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl( repartition_sorts: bool, ) -> Result { let schema = create_test_schema()?; @@ -485,13 +551,11 @@ async fn test_union_inputs_different_sorted3() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet Optimized Plan: - SortPreservingMergeExec: [nullable_col@0 ASC] - UnionExec - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); // should adjust sorting in the first input of the union such that it is not unnecessarily fine Ok(()) @@ -529,14 +593,12 @@ async fn test_union_inputs_different_sorted4() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet Optimized Plan: - SortPreservingMergeExec: [nullable_col@0 ASC, non_nullable_col@1 ASC] - UnionExec - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + UnionExec + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); Ok(()) @@ -583,12 +645,9 @@ async fn test_union_inputs_different_sorted5() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet Optimized Plan: - SortPreservingMergeExec: [nullable_col@0 ASC] - UnionExec - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); Ok(()) @@ -631,14 +690,10 @@ async fn test_union_inputs_different_sorted6() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet Optimized Plan: - SortPreservingMergeExec: [nullable_col@0 ASC] - UnionExec - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); // Should adjust the requirement in the third input of the union so // that it is not unnecessarily fine. @@ -664,13 +719,20 @@ async fn test_union_inputs_different_sorted7() -> Result<()> { // Union has unnecessarily fine ordering below it. We should be able to replace them with absolutely necessary ordering. let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true); assert_snapshot!(test.run(), @r" - Input / Optimized Plan: + Input Plan: SortPreservingMergeExec: [nullable_col@0 ASC] UnionExec SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + + Optimized Plan: + UnionExec + SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); // Union preserves the inputs ordering, and we should not change any of the SortExecs under UnionExec @@ -807,9 +869,10 @@ async fn test_soft_hard_requirements_remove_soft_requirement_without_pushdowns() Optimized Plan: ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as count] - BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "#); // TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed // let expected_optimized = [ @@ -855,10 +918,11 @@ async fn test_soft_hard_requirements_remove_soft_requirement_without_pushdowns() Optimized Plan: BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] - ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col] - SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "#); // TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed // let expected_optimized = [ @@ -918,10 +982,11 @@ async fn test_soft_hard_requirements_multiple_soft_requirements() -> Result<()> Optimized Plan: BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] - ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col] - SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "#); // TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed // let expected_optimized = [ @@ -985,10 +1050,11 @@ async fn test_soft_hard_requirements_multiple_soft_requirements() -> Result<()> Optimized Plan: BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] - ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col] - SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "#); // TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed // let expected_optimized = [ @@ -1054,10 +1120,11 @@ async fn test_soft_hard_requirements_multiple_sorts() -> Result<()> { Optimized Plan: SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] - ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col] - SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + ProjectionExec: expr=[nullable_col@0 + non_nullable_col@1 as nullable_col] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "#); // TODO When sort pushdown respects to the alternatives, and removes soft SortExecs this should be changed // let expected_optimized = [ @@ -1165,7 +1232,7 @@ async fn test_window_multi_path_sort() -> Result<()> { // are not necessarily the same to be able to remove them. let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true); assert_snapshot!(test.run(), @r#" - Input Plan: + Input / Optimized Plan: BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] SortPreservingMergeExec: [nullable_col@0 DESC NULLS LAST] UnionExec @@ -1173,13 +1240,6 @@ async fn test_window_multi_path_sort() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC], file_type=parquet SortExec: expr=[nullable_col@0 DESC NULLS LAST], preserve_partitioning=[false] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet - - Optimized Plan: - WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Range, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] - SortPreservingMergeExec: [nullable_col@0 ASC] - UnionExec - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC, non_nullable_col@1 ASC], file_type=parquet - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC], file_type=parquet "#); Ok(()) @@ -1270,14 +1330,12 @@ async fn test_union_inputs_different_sorted_with_limit() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet Optimized Plan: - SortPreservingMergeExec: [nullable_col@0 ASC] - UnionExec - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - GlobalLimitExec: skip=0, fetch=100 - LocalLimitExec: fetch=100 - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 DESC NULLS LAST], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + GlobalLimitExec: skip=0, fetch=100 + LocalLimitExec: fetch=100 + SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 DESC NULLS LAST], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet "); Ok(()) @@ -1346,10 +1404,12 @@ async fn test_sort_merge_join_order_by_left() -> Result<()> { Optimized Plan: SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)] - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[col_a@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet "); } _ => { @@ -1362,11 +1422,12 @@ async fn test_sort_merge_join_order_by_left() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet Optimized Plan: - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] - SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)] - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)] + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[col_a@0 ASC], preserve_partitioning=[false] + SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet "); } @@ -1436,10 +1497,12 @@ async fn test_sort_merge_join_order_by_right() -> Result<()> { Optimized Plan: SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)] - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[col_a@0 ASC, col_b@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet "); } JoinType::RightAnti => { @@ -1453,10 +1516,12 @@ async fn test_sort_merge_join_order_by_right() -> Result<()> { Optimized Plan: SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)] - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[col_a@0 ASC, col_b@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet "); } _ => { @@ -1469,11 +1534,12 @@ async fn test_sort_merge_join_order_by_right() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet Optimized Plan: - SortExec: expr=[col_a@2 ASC, col_b@3 ASC], preserve_partitioning=[false] - SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)] - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + SortMergeJoinExec: join_type=..., on=[(nullable_col@0, col_a@0)] + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[col_a@0 ASC], preserve_partitioning=[false] + SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet "); } @@ -1518,11 +1584,12 @@ async fn test_sort_merge_join_complex_order_by() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet Optimized Plan: - SortExec: expr=[col_b@3 ASC, nullable_col@0 ASC], preserve_partitioning=[false] - SortMergeJoinExec: join_type=Inner, on=[(nullable_col@0, col_a@0)] - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + SortMergeJoinExec: join_type=Inner, on=[(nullable_col@0, col_a@0)] + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[col_a@0 ASC], preserve_partitioning=[false] + SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet "); // can not push down the sort requirements, need to add SortExec @@ -1546,10 +1613,12 @@ async fn test_sort_merge_join_complex_order_by() -> Result<()> { Optimized Plan: SortMergeJoinExec: join_type=Inner, on=[(nullable_col@0, col_a@0)] - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet - SortExec: expr=[col_a@0 ASC, col_b@1 ASC], preserve_partitioning=[false] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], file_type=parquet + SortExec: expr=[col_a@0 ASC], preserve_partitioning=[true] + RepartitionExec: partitioning=Hash([col_a@0], 10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[col_a, col_b], file_type=parquet "); // Can push down the sort requirements since col_a = nullable_col @@ -1628,10 +1697,7 @@ async fn test_with_lost_ordering_unbounded() -> Result<()> { StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC] Optimized Plan: - SortPreservingMergeExec: [a@0 ASC] - RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true - StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC] + StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC] "); let test_with_repartition_sorts = @@ -1646,10 +1712,7 @@ async fn test_with_lost_ordering_unbounded() -> Result<()> { StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC] Optimized Plan: - SortPreservingMergeExec: [a@0 ASC] - RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10, preserve_order=true, sort_exprs=a@0 ASC - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true - StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC] + StreamingTableExec: partition_sizes=1, projection=[a, b, c, d, e], infinite_source=true, output_ordering=[a@0 ASC] "); Ok(()) @@ -1663,12 +1726,15 @@ async fn test_with_lost_ordering_bounded() -> Result<()> { EnforceSortingTest::new(physical_plan.clone()).with_repartition_sorts(false); assert_snapshot!(test_no_repartition_sorts.run(), @r" - Input / Optimized Plan: + Input Plan: SortExec: expr=[a@0 ASC], preserve_partitioning=[false] CoalescePartitionsExec RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10 RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false + + Optimized Plan: + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false "); let test_with_repartition_sorts = @@ -1683,11 +1749,7 @@ async fn test_with_lost_ordering_bounded() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false Optimized Plan: - SortPreservingMergeExec: [a@0 ASC] - SortExec: expr=[a@0 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=Hash([c@2], 10), input_partitions=10 - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=csv, has_header=false "); Ok(()) @@ -1705,11 +1767,15 @@ async fn test_do_not_pushdown_through_spm() -> Result<()> { let test = EnforceSortingTest::new(physical_plan.clone()).with_repartition_sorts(true); assert_snapshot!(test.run(), @r" - Input / Optimized Plan: + Input Plan: SortExec: expr=[b@1 ASC], preserve_partitioning=[false] SortPreservingMergeExec: [a@0 ASC, b@1 ASC] RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false + + Optimized Plan: + SortExec: expr=[b@1 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false "); Ok(()) @@ -1741,10 +1807,8 @@ async fn test_pushdown_through_spm() -> Result<()> { DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false Optimized Plan: - SortPreservingMergeExec: [a@0 ASC, b@1 ASC] - SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1, maintains_sort_order=true - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false + SortExec: expr=[a@0 ASC, b@1 ASC, c@2 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC, b@1 ASC], file_type=csv, has_header=false "); Ok(()) } @@ -1773,11 +1837,8 @@ async fn test_window_multi_layer_requirement() -> Result<()> { Optimized Plan: BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - SortPreservingMergeExec: [a@0 ASC, b@1 ASC] - SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10 - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false + SortExec: expr=[a@0 ASC, b@1 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false "#); Ok(()) @@ -1900,8 +1961,7 @@ async fn test_add_required_sort() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) @@ -1967,9 +2027,8 @@ async fn test_remove_unnecessary_sort2() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=10 - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: partitions=1, partition_sizes=[0] + SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) @@ -2012,9 +2071,7 @@ async fn test_remove_unnecessary_sort3() -> Result<()> { Optimized Plan: AggregateExec: mode=Final, gby=[], aggr=[] - CoalescePartitionsExec - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) @@ -2064,10 +2121,8 @@ async fn test_remove_unnecessary_sort4() -> Result<()> { SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[true] FilterExec: NOT non_nullable_col@1 UnionExec - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: partitions=1, partition_sizes=[0] - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) @@ -2215,8 +2270,7 @@ async fn test_remove_unnecessary_spm1() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] - DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) @@ -2239,14 +2293,61 @@ async fn test_remove_unnecessary_spm2() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - LocalLimitExec: fetch=100 - SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false] - DataSourceExec: partitions=1, partition_sizes=[0] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) } +#[test] +fn test_spm_fetch_preserves_ordering_through_child_rewrite() -> Result<()> { + let schema = create_test_schema()?; + let ordering: LexOrdering = [sort_expr("non_nullable_col", &schema)].into(); + let source = parquet_exec_with_sort(Arc::clone(&schema), vec![ordering.clone()]); + let projection = projection_exec( + vec![ + (col("nullable_col", &schema)?, "nullable_col".to_string()), + ( + col("non_nullable_col", &schema)?, + "non_nullable_col".to_string(), + ), + ], + source, + )?; + let plan = sort_preserving_merge_exec_with_fetch(ordering.clone(), projection, 100); + + let optimized = PlanWithCorrespondingSort::new_default(plan) + .transform_up(ensure_sorting)? + .data; + let optimized = check_integrity(optimized)?.plan; + let limit = optimized + .downcast_ref::() + .expect("SPM fetch should become a local limit"); + assert_eq!(limit.fetch(), 100); + assert_eq!(limit.required_ordering().as_ref(), Some(&ordering)); + + let config = ConfigOptions::new(); + let optimized = ProjectionPushdown::new().optimize(optimized, &config)?; + let limit = optimized + .downcast_ref::() + .expect("projection rewrite should retain the local limit"); + assert_eq!(limit.required_ordering().as_ref(), Some(&ordering)); + assert!(limit.input().is::()); + + let optimized = LimitPushdown::new().optimize(optimized, &config)?; + let source = optimized + .downcast_ref::() + .expect("limit should be pushed into the parquet scan"); + let config = source + .data_source() + .downcast_ref::() + .expect("parquet scan should use FileScanConfig"); + assert_eq!(config.limit, Some(100)); + assert!(config.preserve_order); + + Ok(()) +} + #[tokio::test] async fn test_change_wrong_sorting() -> Result<()> { let schema = create_test_schema()?; @@ -2267,7 +2368,7 @@ async fn test_change_wrong_sorting() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - SortExec: expr=[nullable_col@0 ASC, non_nullable_col@1 ASC], preserve_partitioning=[false] + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] DataSourceExec: partitions=1, partition_sizes=[0] "); @@ -2295,7 +2396,7 @@ async fn test_change_wrong_sorting2() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - SortExec: expr=[non_nullable_col@1 ASC], preserve_partitioning=[false] + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] DataSourceExec: partitions=1, partition_sizes=[0] "); @@ -2360,22 +2461,16 @@ async fn test_commutativity() -> Result<()> { "#); let config = ConfigOptions::new(); - let rules = vec![ - Arc::new(EnforceDistribution::new()) as Arc, - Arc::new(EnforceSorting::new()) as Arc, - ]; - let mut first_plan = orig_plan.clone(); - for rule in rules { - first_plan = rule.optimize(first_plan, &config)?; - } - - let rules = vec![ - Arc::new(EnforceSorting::new()) as Arc, - Arc::new(EnforceDistribution::new()) as Arc, - Arc::new(EnforceSorting::new()) as Arc, - ]; + // Idempotency check: under the previous design this verified that + // `[EnforceDistribution, EnforceSorting]` produced the same plan as + // `[EnforceSorting, EnforceDistribution, EnforceSorting]`. With the + // merged `EnsureRequirements` rule the property collapses to + // "running EnsureRequirements N times is the same as running it once", + // which is the idempotency guarantee the merged rule provides. + let rule = EnsureRequirements::new(); + let first_plan = rule.optimize(orig_plan.clone(), &config)?; let mut second_plan = orig_plan.clone(); - for rule in rules { + for _ in 0..3 { second_plan = rule.optimize(second_plan, &config)?; } @@ -2414,10 +2509,8 @@ async fn test_coalesce_propagate() -> Result<()> { DataSourceExec: partitions=1, partition_sizes=[0] Optimized Plan: - SortPreservingMergeExec: [nullable_col@0 ASC] - SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[true] - RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 - DataSourceExec: partitions=1, partition_sizes=[0] + SortExec: expr=[nullable_col@0 ASC], preserve_partitioning=[false] + DataSourceExec: partitions=1, partition_sizes=[0] "); Ok(()) @@ -2850,3 +2943,517 @@ async fn test_sort_with_streaming_table() -> Result<()> { Ok(()) } + +/// Regression: `parallelize_sorts` must not relocate a per-partition `SortExec` +/// below an order-preserving `ProjectionExec` that *reorders* columns without +/// remapping the sort-key column indices. +/// +/// Builds the minimal physical plan that reproduces the bug (as it looks after +/// distribution enforcement, before sorting enforcement): +/// +/// ```text +/// SortExec(fetch=4) [score@1 DESC, a@0 ASC] (global, single-partition) +/// CoalescePartitionsExec +/// ProjectionExec [a@0, score@2 as score, b@1 as value] <- reorder: score 2 -> 1 +/// ProjectionExec [a@0, b@1, c+d as score] <- computes score at index 2 +/// SortExec(fetch=1000) [c+d DESC] preserve_partitioning=[true] <- inner ordering +/// RepartitionExec(RoundRobinBatch) <- multi-partition +/// DataSourceExec +/// ``` +/// +/// `parallelize_sorts` turns the `CoalescePartitionsExec` + global `SortExec` +/// into a `SortPreservingMergeExec` + per-partition `SortExec`, sinking the +/// per-partition sort *below* the reordering projection. The sort key +/// `score@1` is valid in the projection's output schema, but in the child +/// schema `[a, b, score]` index 1 is `b` and `score` is at index 2. If the +/// index is not remapped, the relocated `SortExec` references the wrong column +/// and `SanityCheckPlan` rejects the plan with +/// "does not satisfy order requirements ... Child-0 order: []". +fn reorder_projection_physical_plan() -> Result> { + let schema = create_test_schema3()?; // [a, b, c, d, e] + let source = parquet_exec(schema.clone()); + let repartitioned = repartition_exec(source); // RoundRobinBatch -> multi-partition + + // The score-source expression `c + d`. The inner sort below orders by this + // expression, and the lower projection aliases the *same* expression to + // `score`, so the projection output is already ordered by `score`. This + // existing ordering is what drives sort enforcement to relocate the outer + // sort below the reorder projection. + let score_expr = Arc::new(BinaryExpr::new( + col("c", &schema)?, + Operator::Plus, + col("d", &schema)?, + )) as Arc; + + // Inner per-partition, fetch-bearing sort on `c + d`. + let inner_ordering: LexOrdering = [PhysicalSortExpr::new( + Arc::clone(&score_expr), + SortOptions { + descending: true, + nulls_first: false, + }, + )] + .into(); + let inner_sort = Arc::new( + SortExec::new(inner_ordering, repartitioned) + .with_fetch(Some(1000)) + .with_preserve_partitioning(true), + ); + + // Lower projection: compute `score` (= c + d) as the last column. Output + // schema: [a, b, score]; output is ordered by `score`. + let lower = projection_exec( + vec![ + (col("a", &schema)?, "a".to_string()), + (col("b", &schema)?, "b".to_string()), + (Arc::clone(&score_expr), "score".to_string()), + ], + inner_sort, + )?; + + // Upper projection: reorder so `score` moves from input index 2 to output + // index 1, and rename `b` to `value`. Output schema: [a, score, value]. + let lower_schema = lower.schema(); + let upper = projection_exec( + vec![ + (col("a", &lower_schema)?, "a".to_string()), + (col("score", &lower_schema)?, "score".to_string()), + (col("b", &lower_schema)?, "value".to_string()), + ], + lower, + )?; + + // Global, fetch-bearing sort on the renamed column + a tiebreaker, expressed + // in the upper projection's output schema (score@1 DESC NULLS LAST, a@0 ASC). + let upper_schema = upper.schema(); + let ordering: LexOrdering = [ + sort_expr_options( + "score", + &upper_schema, + SortOptions { + descending: true, + nulls_first: false, + }, + ), + sort_expr("a", &upper_schema), + ] + .into(); + let coalesced = coalesce_partitions_exec(upper); + Ok(sort_exec_with_fetch(ordering, Some(4), coalesced)) +} + +#[tokio::test] +async fn test_parallelize_sorts_remaps_index_through_reordering_projection() -> Result<()> +{ + let physical_plan = reorder_projection_physical_plan()?; + + // `EnsureRequirements` (with sort repartitioning enabled) runs the sort + // enforcement pass, including `parallelize_sorts`. + let mut config = ConfigOptions::new(); + config.optimizer.repartition_sorts = true; + let optimized = EnsureRequirements::new().optimize(physical_plan, &config)?; + + // The optimized plan must be physically valid. Before the fix this fails: + // the per-partition `SortExec` was relocated below the reordering projection + // but kept the key `score@1` (valid only in the projection output), while its + // child schema `[a, b, score]` has `score` at index 2 — so `SanityCheckPlan` + // reports `does not satisfy order requirements: [...]. Child-0 order: []`. + SanityCheckPlan::new() + .optimize(Arc::clone(&optimized), &ConfigOptions::default()) + .unwrap_or_else(|e| { + panic!( + "Sort enforcement produced a plan that fails SanityCheckPlan \ + (stale sort-key index after relocating the SortExec below a \ + reordering ProjectionExec): {e}\n\nPlan:\n{}", + displayable(optimized.as_ref()).indent(true) + ) + }); + + Ok(()) +} + +#[tokio::test] +async fn test_push_sort_through_reordered_projection_to_union() -> Result<()> { + let schema = create_test_schema3()?; + let ordering: LexOrdering = [sort_expr("a", &schema)].into(); + + let sorted_source = parquet_exec_with_sort(schema.clone(), vec![ordering.clone()]); + let unsorted_source = sort_exec(ordering.clone(), parquet_exec(schema.clone())); + let union = union_exec(vec![sorted_source, unsorted_source]); + + let projection = projection_exec( + vec![ + (col("c", &schema)?, "c".to_string()), + (col("b", &schema)?, "b".to_string()), + (col("a", &schema)?, "a".to_string()), + ], + union, + )?; + + let physical_plan = + sort_exec([sort_expr("a", &projection.schema())].into(), projection); + + let test = EnforceSortingTest::new(physical_plan).with_repartition_sorts(true); + assert_snapshot!(test.run(), @r" + Input Plan: + SortExec: expr=[a@2 ASC], preserve_partitioning=[false] + ProjectionExec: expr=[c@2 as c, b@1 as b, a@0 as a] + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [a@2 ASC] + ProjectionExec: expr=[c@2 as c, b@1 as b, a@0 as a] + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], output_ordering=[a@0 ASC], file_type=parquet + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + Ok(()) +} + +#[tokio::test] +async fn test_push_sort_through_alias_reordered_projection() -> Result<()> { + let schema = create_test_schema3()?; + let source = parquet_exec(schema.clone()); + let projection = projection_exec( + vec![ + (col("c", &schema)?, "c_alias".to_string()), + (col("b", &schema)?, "b_alias".to_string()), + (col("a", &schema)?, "a_alias".to_string()), + ], + source, + )?; + + let physical_plan = sort_exec( + [sort_expr("a_alias", &projection.schema())].into(), + projection, + ); + + let test = EnforceSortingTest::new(physical_plan); + assert_snapshot!(test.run(), @r" + Input Plan: + SortExec: expr=[a_alias@2 ASC], preserve_partitioning=[false] + ProjectionExec: expr=[c@2 as c_alias, b@1 as b_alias, a@0 as a_alias] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + + Optimized Plan: + ProjectionExec: expr=[c@2 as c_alias, b@1 as b_alias, a@0 as a_alias] + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + Ok(()) +} + +#[tokio::test] +async fn test_does_not_push_sort_through_computed_projection() -> Result<()> { + let schema = create_test_schema3()?; + let source = parquet_exec(schema.clone()); + let computed_expr = Arc::new(BinaryExpr::new( + col("a", &schema)?, + Operator::Plus, + col("b", &schema)?, + )) as Arc; + let projection = projection_exec( + vec![ + (computed_expr, "sort_key".to_string()), + (col("c", &schema)?, "c".to_string()), + ], + source, + )?; + + let physical_plan = sort_exec( + [sort_expr("sort_key", &projection.schema())].into(), + projection, + ); + + let test = EnforceSortingTest::new(physical_plan); + assert_snapshot!(test.run(), @" + Input Plan: + SortExec: expr=[sort_key@0 ASC], preserve_partitioning=[false] + ProjectionExec: expr=[a@0 + b@1 as sort_key, c@2 as c] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + + Optimized Plan: + SortExec: expr=[sort_key@0 ASC], preserve_partitioning=[false] + CoalescePartitionsExec + ProjectionExec: expr=[a@0 + b@1 as sort_key, c@2 as c] + RepartitionExec: partitioning=RoundRobinBatch(10), input_partitions=1 + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + Ok(()) +} + +#[tokio::test] +async fn test_push_fetch_sort_through_alias_reordered_projection() -> Result<()> { + let schema = create_test_schema3()?; + let source = parquet_exec(schema.clone()); + let projection = projection_exec( + vec![ + (col("c", &schema)?, "c_alias".to_string()), + (col("b", &schema)?, "b_alias".to_string()), + (col("a", &schema)?, "a_alias".to_string()), + ], + source, + )?; + + let physical_plan = sort_exec_with_fetch( + [sort_expr("a_alias", &projection.schema())].into(), + Some(3), + projection, + ); + + let test = EnforceSortingTest::new(physical_plan); + assert_snapshot!(test.run(), @r" + Input Plan: + SortExec: TopK(fetch=3), expr=[a_alias@2 ASC], preserve_partitioning=[false] + ProjectionExec: expr=[c@2 as c_alias, b@1 as b_alias, a@0 as a_alias] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + + Optimized Plan: + ProjectionExec: expr=[c@2 as c_alias, b@1 as b_alias, a@0 as a_alias] + SortExec: TopK(fetch=3), expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + Ok(()) +} + +#[tokio::test] +async fn test_push_sort_through_reordered_projection_remaps_multiple_keys_and_options() +-> Result<()> { + let schema = create_test_schema3()?; + let source = parquet_exec(schema.clone()); + + let projection = projection_exec( + vec![ + (col("c", &schema)?, "c_alias".to_string()), + (col("b", &schema)?, "b_alias".to_string()), + (col("a", &schema)?, "a_alias".to_string()), + ], + source, + )?; + + let projection_schema = projection.schema(); + let ordering: LexOrdering = [ + sort_expr_options( + "c_alias", + &projection_schema, + SortOptions { + descending: true, + nulls_first: false, + }, + ), + sort_expr("a_alias", &projection_schema), + ] + .into(); + + let physical_plan = sort_exec(ordering, projection); + + let test = EnforceSortingTest::new(physical_plan); + assert_snapshot!(test.run(), @r" + Input Plan: + SortExec: expr=[c_alias@0 DESC NULLS LAST, a_alias@2 ASC], preserve_partitioning=[false] + ProjectionExec: expr=[c@2 as c_alias, b@1 as b_alias, a@0 as a_alias] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + + Optimized Plan: + ProjectionExec: expr=[c@2 as c_alias, b@1 as b_alias, a@0 as a_alias] + SortExec: expr=[c@2 DESC NULLS LAST, a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + Ok(()) +} + +#[tokio::test] +async fn test_does_not_push_fetch_sort_through_projection_over_union() -> Result<()> { + let schema = create_test_schema3()?; + let union = union_exec(vec![ + parquet_exec(schema.clone()), + parquet_exec(schema.clone()), + ]); + + let projection = projection_exec( + vec![ + (col("c", &schema)?, "c".to_string()), + (col("b", &schema)?, "b".to_string()), + (col("a", &schema)?, "a".to_string()), + ], + union, + )?; + + let physical_plan = sort_exec_with_fetch( + [sort_expr("a", &projection.schema())].into(), + Some(4), + projection, + ); + + let test = EnforceSortingTest::new(physical_plan); + assert_snapshot!(test.run(), @r" + Input Plan: + SortExec: TopK(fetch=4), expr=[a@2 ASC], preserve_partitioning=[false] + ProjectionExec: expr=[c@2 as c, b@1 as b, a@0 as a] + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + + Optimized Plan: + SortExec: TopK(fetch=4), expr=[a@2 ASC], preserve_partitioning=[false] + CoalescePartitionsExec + ProjectionExec: expr=[c@2 as c, b@1 as b, a@0 as a] + UnionExec + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); + + Ok(()) +} + +/// A pass-through wrapper around a column: just assert that column does not contain any nulls +#[derive(Debug, Eq)] +struct AssertNotNull { + inner: Arc, +} + +impl AssertNotNull { + fn new(inner: Arc) -> Arc { + Arc::new(Self { inner }) + } +} + +impl PartialEq for AssertNotNull { + fn eq(&self, other: &Self) -> bool { + self.inner.eq(&other.inner) + } +} + +impl std::hash::Hash for AssertNotNull { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + } +} + +impl std::fmt::Display for AssertNotNull { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "assert_not_null({})", self.inner) + } +} + +impl PhysicalExpr for AssertNotNull { + fn data_type(&self, input_schema: &Schema) -> Result { + self.inner.data_type(input_schema) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let child = self.inner.evaluate(batch)?; + match child { + ColumnarValue::Array(a) if a.logical_null_count() > 0 => Err( + DataFusionError::Internal("AssertNotNull evaluated to null".to_string()), + ), + ColumnarValue::Scalar(s) if s.is_null() => Err(DataFusionError::Internal( + "AssertNotNull evaluated to null".to_string(), + )), + child => Ok(child), + } + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(AssertNotNull { + inner: Arc::clone(&children[0]), + })) + } + + fn get_properties( + &self, + children: &[datafusion_expr::sort_properties::ExprProperties], + ) -> Result { + Ok(children[0].clone()) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "assert_not_null({})", self.inner) + } +} + +#[tokio::test] +async fn test_passthrough_wrapper_projection_keeps_ordering() -> Result<()> { + fn sort_expr(name: &str, schema: &Schema) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: col(name, schema).unwrap(), + options: Default::default(), + } + } + + pub fn projection_exec( + expr: Vec<(Arc, String)>, + input: Arc, + ) -> Result> { + let proj_exprs: Vec = expr + .into_iter() + .map(|(expr, alias)| ProjectionExpr { expr, alias }) + .collect(); + Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) + } + + let batch = record_batch!( + ("a", Utf8, ["x", "y"]), + ("b", Utf8, ["1", "2"]), + ("c", Utf8, ["1", "2"]) + )?; + let schema = batch.schema(); + let source = Arc::new(DataSourceExec::new(Arc::new( + datafusion::datasource::memory::MemorySourceConfig::try_new( + &[vec![batch]], + schema.clone(), + None, + )? + .try_with_sort_information(vec![ + LexOrdering::new([ + sort_expr("a", &schema), + sort_expr("b", &schema), + sort_expr("c", &schema), + ]) + .unwrap(), + ])?, + ))) as Arc; + + let projection = projection_exec( + vec![ + (AssertNotNull::new(col("a", &schema)?), "a".to_string()), + (AssertNotNull::new(col("b", &schema)?), "b".to_string()), + (AssertNotNull::new(col("c", &schema)?), "c".to_string()), + ], + source, + )?; + + let ordering = LexOrdering::new([ + sort_expr("a", &projection.schema()), + sort_expr("b", &projection.schema()), + sort_expr("c", &projection.schema()), + ]) + .unwrap(); + + let sort_satisfied = projection + .equivalence_properties() + .ordering_satisfy(ordering.clone())?; + + let plan_str = displayable(projection.as_ref()).indent(true).to_string(); + assert!( + sort_satisfied, + "sort should be satisfied, ordering: {ordering}\nplan:\n{plan_str}" + ); + + Ok(()) +} diff --git a/datafusion/core/tests/physical_optimizer/enforce_sorting_monotonicity.rs b/datafusion/core/tests/physical_optimizer/enforce_sorting_monotonicity.rs index de7611ff211a5..99e25a6c82595 100644 --- a/datafusion/core/tests/physical_optimizer/enforce_sorting_monotonicity.rs +++ b/datafusion/core/tests/physical_optimizer/enforce_sorting_monotonicity.rs @@ -433,8 +433,10 @@ fn test_window_partial_constant_and_set_monotonicity_8() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST] + WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -457,8 +459,10 @@ fn test_window_partial_constant_and_set_monotonicity_9() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 DESC NULLS LAST] + WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -475,10 +479,17 @@ fn test_window_partial_constant_and_set_monotonicity_10() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST] + SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -495,10 +506,17 @@ fn test_window_partial_constant_and_set_monotonicity_11() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST] + SortExec: expr=[avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -520,10 +538,17 @@ fn test_window_partial_constant_and_set_monotonicity_12() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [non_nullable_col@1 ASC NULLS LAST, count@2 ASC NULLS LAST] + SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -541,10 +566,17 @@ fn test_window_partial_constant_and_set_monotonicity_13() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [non_nullable_col@1 ASC NULLS LAST, max@2 DESC NULLS LAST] + SortExec: expr=[non_nullable_col@1 ASC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -562,10 +594,17 @@ fn test_window_partial_constant_and_set_monotonicity_14() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[min@2 DESC NULLS LAST, non_nullable_col@1 ASC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [min@2 DESC NULLS LAST, non_nullable_col@1 ASC NULLS LAST] + SortExec: expr=[min@2 DESC NULLS LAST, non_nullable_col@1 ASC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -583,10 +622,17 @@ fn test_window_partial_constant_and_set_monotonicity_15() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST] + SortExec: expr=[avg@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -806,8 +852,10 @@ fn test_window_partial_constant_and_set_monotonicity_24() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, count@2 DESC NULLS LAST] + WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -825,10 +873,17 @@ fn test_window_partial_constant_and_set_monotonicity_25() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -845,10 +900,17 @@ fn test_window_partial_constant_and_set_monotonicity_26() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[min@2 DESC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [min@2 DESC NULLS LAST] + SortExec: expr=[min@2 DESC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "#); } @@ -865,10 +927,17 @@ fn test_window_partial_constant_and_set_monotonicity_27() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[avg@2 DESC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [avg@2 DESC NULLS LAST] + SortExec: expr=[avg@2 DESC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "#); } @@ -891,10 +960,17 @@ fn test_window_partial_constant_and_set_monotonicity_28() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[count@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [count@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST] + SortExec: expr=[count@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[count: Ok(Field { name: "count", data_type: Int64 }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -918,8 +994,10 @@ fn test_window_partial_constant_and_set_monotonicity_29() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 DESC] + WindowAggExec: wdw=[max: Ok(Field { name: "max", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "#) } @@ -935,10 +1013,17 @@ fn test_window_partial_constant_and_set_monotonicity_30() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[min@2 DESC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [min@2 DESC NULLS LAST] + SortExec: expr=[min@2 DESC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[min: Ok(Field { name: "min", data_type: Int32, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "#); } @@ -955,10 +1040,17 @@ fn test_window_partial_constant_and_set_monotonicity_31() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST], preserve_partitioning=[false] WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST], preserve_partitioning=[true] + WindowAggExec: wdw=[avg: Ok(Field { name: "avg", data_type: Float64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: CurrentRow, end_bound: Following(UInt64(NULL)), is_causal: false }] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1179,8 +1271,10 @@ fn test_window_partial_constant_and_set_monotonicity_40() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST] + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1198,10 +1292,17 @@ fn test_window_partial_constant_and_set_monotonicity_41() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[max@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [max@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST] + SortExec: expr=[max@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1219,10 +1320,17 @@ fn test_window_partial_constant_and_set_monotonicity_42() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST] + SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1240,10 +1348,17 @@ fn test_window_partial_constant_and_set_monotonicity_43() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST, avg@2 ASC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1265,10 +1380,17 @@ fn test_window_partial_constant_and_set_monotonicity_44() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[count@2 ASC], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [count@2 ASC] + SortExec: expr=[count@2 ASC], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1286,10 +1408,17 @@ fn test_window_partial_constant_and_set_monotonicity_45() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 DESC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 DESC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1313,8 +1442,10 @@ fn test_window_partial_constant_and_set_monotonicity_46() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, min@2 DESC NULLS LAST] + BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1337,8 +1468,10 @@ fn test_window_partial_constant_and_set_monotonicity_47() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST] + BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1553,8 +1686,10 @@ fn test_window_partial_constant_and_set_monotonicity_56() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [count@2 ASC NULLS LAST, nullable_col@0 ASC NULLS LAST] + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1572,10 +1707,17 @@ fn test_window_partial_constant_and_set_monotonicity_57() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1593,10 +1735,17 @@ fn test_window_partial_constant_and_set_monotonicity_58() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST] + SortExec: expr=[min@2 DESC NULLS LAST, nullable_col@0 ASC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1613,10 +1762,17 @@ fn test_window_partial_constant_and_set_monotonicity_59() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[avg@2 ASC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [avg@2 ASC NULLS LAST] + SortExec: expr=[avg@2 ASC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1639,10 +1795,17 @@ fn test_window_partial_constant_and_set_monotonicity_60() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST, count@2 ASC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1660,10 +1823,17 @@ fn test_window_partial_constant_and_set_monotonicity_61() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, max@2 ASC] + SortExec: expr=[nullable_col@0 ASC NULLS LAST, max@2 ASC], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[max: Field { "max": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1681,10 +1851,17 @@ fn test_window_partial_constant_and_set_monotonicity_62() { ], }.run(), @ r#" - Input / Optimized Plan: + Input Plan: SortExec: expr=[nullable_col@0 ASC NULLS LAST, min@2 DESC NULLS LAST], preserve_partitioning=[false] BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + + Optimized Plan: + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST, min@2 DESC NULLS LAST] + SortExec: expr=[nullable_col@0 ASC NULLS LAST, min@2 DESC NULLS LAST], preserve_partitioning=[true] + BoundedWindowAggExec: wdw=[min: Field { "min": nullable Int32 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } @@ -1707,8 +1884,10 @@ fn test_window_partial_constant_and_set_monotonicity_63() { DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet Optimized Plan: - BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] - DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet + SortPreservingMergeExec: [nullable_col@0 ASC NULLS LAST] + BoundedWindowAggExec: wdw=[avg: Field { "avg": nullable Float64 }, frame: ROWS BETWEEN 1 PRECEDING AND CURRENT ROW], mode=[Sorted] + RepartitionExec: partitioning=Hash([nullable_col@0], 10), input_partitions=1, maintains_sort_order=true + DataSourceExec: file_groups={1 group: [[x]]}, projection=[nullable_col, non_nullable_col], output_ordering=[nullable_col@0 ASC NULLS LAST], file_type=parquet "# ); } diff --git a/datafusion/core/tests/physical_optimizer/ensure_requirements.rs b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs new file mode 100644 index 0000000000000..86b60519da370 --- /dev/null +++ b/datafusion/core/tests/physical_optimizer/ensure_requirements.rs @@ -0,0 +1,1434 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Integration tests for `EnsureRequirements`. +//! +//! Ported verbatim from `datafusion/physical-optimizer/src/ensure_requirements/mod.rs` +//! so the tests live alongside the rest of the `physical_optimizer/` integration +//! suite and can use real `ExecutionPlan`s where convenient. + +use insta::assert_snapshot; + +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{TransformedResult, TreeNode, TreeNodeRecursion}; +use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_optimizer::ensure_requirements::EnsureRequirements; +use datafusion_physical_optimizer::ensure_requirements::enforce_sorting::{ + PlanWithCorrespondingCoalescePartitions, parallelize_sorts, +}; + +use std::sync::Arc; + +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion_common::Result; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr::{ + EquivalenceProperties, LexOrdering, PhysicalExpr, PhysicalSortExpr, +}; +use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::limit::GlobalLimitExec; +use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::union::UnionExec; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, +}; + +use datafusion_physical_optimizer::output_requirements::OutputRequirementExec; +use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; + +use datafusion_common::{JoinType, NullEquality}; +use datafusion_physical_expr::Distribution; +use datafusion_physical_expr_common::sort_expr::OrderingRequirements; +use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec}; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; + +const TEST_TARGET_PARTITIONS: usize = 8; + +/// Mock ExecutionPlan with configurable partition count and output ordering. +#[derive(Debug)] +struct MockMultiPartitionExec { + properties: Arc, +} + +impl MockMultiPartitionExec { + fn new(partition_count: usize) -> Self { + Self::with_partitioning(Partitioning::UnknownPartitioning(partition_count)) + } + + /// A source that is already partitioned on `a`, as an aggregate or a partitioned + /// join below the node under test would be. + fn hash_partitioned_on_a(partition_count: usize) -> Self { + Self::with_partitioning(Partitioning::Hash( + vec![Arc::new(Column::new("a", 0))], + partition_count, + )) + } + + fn with_partitioning(partitioning: Partitioning) -> Self { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])); + let mut eq = EquivalenceProperties::new(Arc::clone(&schema)); + if let Some(ordering) = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions { + descending: false, + nulls_first: false, + }, + )]) { + eq.add_orderings(vec![ordering.into_iter().collect::>()]); + } + let properties = PlanProperties::new( + eq, + partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + ); + Self { + properties: Arc::new(properties), + } + } +} + +impl DisplayAs for MockMultiPartitionExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "MockMultiPartitionExec") + } +} + +impl ExecutionPlan for MockMultiPartitionExec { + fn name(&self) -> &str { + "MockMultiPartitionExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } +} + +fn test_config() -> ConfigOptions { + let mut config = ConfigOptions::default(); + // Keep plan-shape tests deterministic across machines with different CPU counts. + config.execution.target_partitions = TEST_TARGET_PARTITIONS; + config +} + +/// Helper: run EnsureRequirements and verify SanityCheckPlan passes +fn optimize_and_sanity_check( + plan: Arc, +) -> Result> { + let config = test_config(); + let optimized = EnsureRequirements::new().optimize(plan, &config)?; + // SanityCheckPlan must pass + SanityCheckPlan::new().optimize(Arc::clone(&optimized), &config)?; + Ok(optimized) +} + +/// Helper: verify idempotency — running twice produces the same plan +fn assert_idempotent(plan: Arc) { + let config = test_config(); + let p1 = EnsureRequirements::new() + .optimize(plan, &config) + .expect("first optimize failed"); + let p2 = EnsureRequirements::new() + .optimize(Arc::clone(&p1), &config) + .expect("second optimize failed"); + + let s1 = plan_string(&p1); + let s2 = plan_string(&p2); + assert_eq!( + s1, s2, + "EnsureRequirements is NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}" + ); + + // Both must pass SanityCheckPlan + SanityCheckPlan::new() + .optimize(p1, &config) + .expect("SanityCheckPlan failed on first pass"); + SanityCheckPlan::new() + .optimize(p2, &config) + .expect("SanityCheckPlan failed on second pass"); +} + +/// Single-column `LexOrdering` on `(name, idx)` with the given options. +/// Most tests in this file want a one-column ordering on the canonical +/// `a@0` or `b@1` columns; this helper trims the 7-line per-test +/// boilerplate down to a single call. +fn sort_expr_on( + name: &str, + idx: usize, + descending: bool, + nulls_first: bool, +) -> LexOrdering { + LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::new(Column::new(name, idx)), + SortOptions { + descending, + nulls_first, + }, + )]) + .unwrap() +} + +/// Render an execution plan with `displayable(...).indent(true)`. +fn plan_string(plan: &Arc) -> String { + datafusion_physical_plan::displayable(plan.as_ref()) + .indent(true) + .to_string() +} + +/// Run `EnsureRequirements`, assert `SanityCheckPlan` passes, snapshot the +/// resulting plan with `insta`, and verify idempotency by running the rule a +/// second time and checking the plan is unchanged. Use this for plan-shape +/// tests so a single call covers "correct plan + sanity + idempotent" and +/// updating an intentional plan change is a single `cargo insta accept`. +macro_rules! assert_ensure_requirements_plan { + ($plan:expr, @ $snapshot:literal $(,)?) => {{ + let config = test_config(); + let p1 = EnsureRequirements::new() + .optimize($plan, &config) + .expect("EnsureRequirements::optimize failed (pass 1)"); + SanityCheckPlan::new() + .optimize(Arc::clone(&p1), &config) + .expect("SanityCheckPlan failed (pass 1)"); + let p1_str = plan_string(&p1); + insta::assert_snapshot!(p1_str, @ $snapshot); + + // Idempotency: a second pass must produce the same plan. + let p2 = EnsureRequirements::new() + .optimize(p1, &config) + .expect("EnsureRequirements::optimize failed (pass 2)"); + let p2_str = plan_string(&p2); + assert_eq!( + p1_str, p2_str, + "EnsureRequirements is NOT idempotent!\nPass 1:\n{p1_str}\nPass 2:\n{p2_str}", + ); + SanityCheckPlan::new() + .optimize(p2, &config) + .expect("SanityCheckPlan failed (pass 2)"); + }}; +} + +/// Union with mixed partition counts + sort + limit. +#[test] +fn test_union_mixed_partitions_sort_limit() { + let live = Arc::new(MockMultiPartitionExec::new(32)); + let historical = Arc::new(MockMultiPartitionExec::new(1)); + + let union = UnionExec::try_new(vec![live as _, historical as _]).unwrap(); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, union)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(21))); + + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=0, fetch=21 + SortPreservingMergeExec: [a@0 DESC] + UnionExec + SortExec: expr=[a@0 DESC], preserve_partitioning=[true] + MockMultiPartitionExec + SortExec: expr=[a@0 DESC], preserve_partitioning=[false] + MockMultiPartitionExec + "); +} + +/// Idempotency: union with mixed partitions +#[test] +fn test_idempotent_union_mixed_partitions() { + let live = Arc::new(MockMultiPartitionExec::new(8)); + let hist = Arc::new(MockMultiPartitionExec::new(1)); + let union = UnionExec::try_new(vec![live as _, hist as _]).unwrap(); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, union)); + let limit = Arc::new(GlobalLimitExec::new(sort, 0, Some(5))); + + assert_idempotent(limit); +} + +// ======================================================================== +// Projection + multi-partition tests (pushdown_sorts trigger path) +// ======================================================================== + +/// ProjectionExec over multi-partition + sort DESC + limit. +/// This is the topology where pushdown_sorts pushes sort through projection +/// onto the multi-partition source. The optimizer must still produce a valid plan. +#[test] +fn test_projection_over_multi_partition_sort_limit() { + let source = Arc::new(MockMultiPartitionExec::new(16)); + // Identity projection + let proj_exprs: Vec<(Arc, String)> = vec![ + (Arc::new(Column::new("a", 0)), "a".to_string()), + (Arc::new(Column::new("b", 1)), "b".to_string()), + ]; + let projection = Arc::new(ProjectionExec::try_new(proj_exprs, source as _).unwrap()); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, projection)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(21))); + + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=0, fetch=21 + SortPreservingMergeExec: [a@0 DESC] + ProjectionExec: expr=[a@0 as a, b@1 as b] + SortExec: expr=[a@0 DESC], preserve_partitioning=[true] + MockMultiPartitionExec + "); +} + +// ======================================================================== +// Single partition tests (no unnecessary operators) +// ======================================================================== + +/// Single partition source + sort + limit should NOT add SortPreservingMergeExec. +#[test] +fn test_single_partition_no_unnecessary_spm() { + let source = Arc::new(MockMultiPartitionExec::new(1)); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, source)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(10))); + + // Snapshot asserts the plan-shape property: no `SortPreservingMergeExec` + // is added on a single-partition source. + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=0, fetch=10 + SortExec: expr=[a@0 DESC], preserve_partitioning=[false] + MockMultiPartitionExec + "); +} + +/// Source already has correct ordering → should not add SortExec. +#[test] +fn test_sort_already_satisfied_no_extra_sort() { + let source = Arc::new(MockMultiPartitionExec::new(1)); + + // Sort ASC matches MockMultiPartitionExec's output ordering (a ASC) + let sort_expr = sort_expr_on("a", 0, false, false); + + let sort = Arc::new(SortExec::new(sort_expr, source)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(10))); + + // Snapshot asserts the plan-shape property: no `SortExec` is added + // when the source already satisfies the ordering. + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=0, fetch=10 + MockMultiPartitionExec + "); +} + +// ======================================================================== +// Various partition counts (stress test) +// ======================================================================== + +/// Test with different partition counts: 2, 4, 8, 16, 32, 64 +#[test] +fn test_various_partition_counts_all_pass_sanity_check() { + for n in [2, 4, 8, 16, 32, 64] { + let source = Arc::new(MockMultiPartitionExec::new(n)); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, source)); + let limit: Arc = + Arc::new(GlobalLimitExec::new(sort, 0, Some(21))); + + let result = optimize_and_sanity_check(limit); + assert!( + result.is_ok(), + "SanityCheckPlan failed for {n} partitions: {:?}", + result.err() + ); + } +} + +// ======================================================================== +// CoalescePartitionsExec tests +// ======================================================================== + +/// CoalescePartitionsExec + sort should produce valid plan +#[test] +fn test_coalesce_then_sort_limit() { + let source = Arc::new(MockMultiPartitionExec::new(8)); + let coalesce: Arc = Arc::new(CoalescePartitionsExec::new(source)); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, coalesce)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(10))); + + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=0, fetch=10 + SortPreservingMergeExec: [a@0 DESC] + SortExec: expr=[a@0 DESC], preserve_partitioning=[true] + MockMultiPartitionExec + "); +} + +// ======================================================================== +// Filter + multi-partition tests +// ======================================================================== + +/// FilterExec over multi-partition + sort + limit +#[test] +fn test_filter_over_multi_partition_sort_limit() { + use datafusion_common::ScalarValue; + use datafusion_physical_expr::expressions::Literal; + use datafusion_physical_plan::filter::FilterExec; + + let source = Arc::new(MockMultiPartitionExec::new(16)); + + // Simple always-true filter + let predicate = Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))); + let filter = Arc::new(FilterExec::try_new(predicate, source as _).unwrap()); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, filter)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(21))); + + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=0, fetch=21 + SortPreservingMergeExec: [a@0 DESC] + SortExec: expr=[a@0 DESC], preserve_partitioning=[true] + FilterExec: true + MockMultiPartitionExec + "); +} + +// ======================================================================== +// RepartitionExec tests +// ======================================================================== + +/// Existing RepartitionExec + sort + limit must remain valid +#[test] +fn test_repartition_sort_limit_idempotent() { + let source = Arc::new(MockMultiPartitionExec::new(1)); + let repartition = Arc::new( + RepartitionExec::try_new(source as _, Partitioning::RoundRobinBatch(8)).unwrap(), + ); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, repartition)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(10))); + + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=0, fetch=10 + SortExec: expr=[a@0 DESC], preserve_partitioning=[false] + MockMultiPartitionExec + "); +} + +// ======================================================================== +// Skip + Fetch (offset + limit) tests +// ======================================================================== + +/// GlobalLimitExec with skip=5, fetch=10 must produce valid plan +#[test] +fn test_skip_and_fetch_multi_partition() { + let source = Arc::new(MockMultiPartitionExec::new(16)); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, source)); + // skip=5, fetch=10 + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 5, Some(10))); + + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=5, fetch=10 + SortPreservingMergeExec: [a@0 DESC] + SortExec: expr=[a@0 DESC], preserve_partitioning=[true] + MockMultiPartitionExec + "); +} + +// ======================================================================== +// Multiple sort columns test +// ======================================================================== + +/// Sort on (a DESC, b ASC) with multi-partition +#[test] +fn test_multi_column_sort_multi_partition() { + let source = Arc::new(MockMultiPartitionExec::new(32)); + + let sort_expr = LexOrdering::new(vec![ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions { + descending: true, + nulls_first: true, + }, + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 1)), + SortOptions { + descending: false, + nulls_first: false, + }, + ), + ]) + .unwrap(); + + let sort = Arc::new(SortExec::new(sort_expr, source)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(21))); + + assert_ensure_requirements_plan!(limit, @r" + GlobalLimitExec: skip=0, fetch=21 + SortPreservingMergeExec: [a@0 DESC, b@1 ASC NULLS LAST] + SortExec: expr=[a@0 DESC, b@1 ASC NULLS LAST], preserve_partitioning=[true] + MockMultiPartitionExec + "); +} + +// ======================================================================== +// pushdown_sorts distribution-awareness regression tests. +// These cover the specific bug where pushdown_sorts pushed a SortExec +// through an intermediate node onto a multi-partition source, setting +// preserve_partitioning=true without inserting SortPreservingMergeExec. +// ======================================================================== + +/// Regression: `OutputRequirementExec(SinglePartition)` wrapping a +/// multi-partition source. `ensure_sorting` must insert a +/// `SortPreservingMergeExec` to satisfy the `SinglePartition` requirement. +/// +/// The final `ensure_distribution` pass catches the distribution +/// violation from `pushdown_sorts`, producing a valid plan (via +/// `CoalescePartitionsExec` or `SortPreservingMergeExec`). +#[test] +fn test_output_requirement_single_partition_over_multi_partition_source() { + let source = Arc::new(MockMultiPartitionExec::new(10)); + + let sort_expr = sort_expr_on("a", 0, true, true); + + // OutputRequirementExec with SinglePartition + ordering requirement + let output_req: Arc = Arc::new(OutputRequirementExec::new( + source, + Some(OrderingRequirements::from(sort_expr)), + Distribution::SinglePartition, + Some(21), + )); + + // SinglePartition must be satisfied (via SPM or Coalesce+Sort) — snapshot + // documents which one the optimizer chooses. + assert_ensure_requirements_plan!(output_req, @r" + OutputRequirementExec: order_by=[(a@0, desc)], dist_by=SinglePartition + SortPreservingMergeExec: [a@0 DESC], fetch=21 + SortExec: TopK(fetch=21), expr=[a@0 DESC], preserve_partitioning=[true] + MockMultiPartitionExec + "); +} + +/// Regression: `pushdown_sorts` pushes a sort through a `ProjectionExec` +/// onto a multi-partition source. The result must include a +/// `SortPreservingMergeExec` (or equivalent) when the parent requires +/// `SinglePartition`. +/// +/// Without the distribution-aware pushdown the standalone +/// `pushdown_sorts` traversal would not propagate distribution; the +/// final `ensure_distribution` pass then catches the violation and +/// inserts a `CoalescePartitionsExec` to satisfy `SinglePartition`. +#[test] +fn test_sort_pushdown_through_projection_adds_spm() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(10)); + + // Identity projection + let proj_exprs: Vec<(Arc, String)> = vec![ + (Arc::new(Column::new("a", 0)), "a".to_string()), + (Arc::new(Column::new("b", 1)), "b".to_string()), + ]; + let projection: Arc = + Arc::new(ProjectionExec::try_new(proj_exprs, source).unwrap()); + + let sort_expr = sort_expr_on("a", 0, true, true); + + // OutputRequirementExec(SinglePartition) → ProjectionExec → multi-partition source + let output_req: Arc = Arc::new(OutputRequirementExec::new( + projection, + Some(OrderingRequirements::from(sort_expr)), + Distribution::SinglePartition, + Some(21), + )); + + // SinglePartition must be satisfied. The final ensure_distribution pass + // adds CoalescePartitionsExec or SortPreservingMergeExec as needed — the + // snapshot documents which. + assert_ensure_requirements_plan!(output_req, @r" + OutputRequirementExec: order_by=[(a@0, desc)], dist_by=SinglePartition + SortPreservingMergeExec: [a@0 DESC], fetch=21 + ProjectionExec: expr=[a@0 as a, b@1 as b] + SortExec: TopK(fetch=21), expr=[a@0 DESC], preserve_partitioning=[true] + MockMultiPartitionExec + "); +} + +// ======================================================================== +// Idempotency tests for distribution-fix scenarios +// These verify that the pushdown_sorts distribution fix actually +// makes EnsureRequirements idempotent for the bug-triggering topologies. +// ======================================================================== + +/// Idempotency for the `OutputRequirementExec(SinglePartition)` + +/// multi-partition source scenario. Running twice must produce the same plan. +#[test] +fn test_idempotent_output_requirement_single_partition() { + let source = Arc::new(MockMultiPartitionExec::new(10)); + let sort_expr = sort_expr_on("a", 0, true, true); + + let output_req: Arc = Arc::new(OutputRequirementExec::new( + source, + Some(OrderingRequirements::from(sort_expr)), + Distribution::SinglePartition, + Some(21), + )); + + assert_idempotent(output_req); +} + +/// Idempotency for the `OutputRequirementExec(SinglePartition)` → +/// `ProjectionExec` → multi-partition source scenario. +#[test] +fn test_idempotent_projection_over_multi_partition_with_single_partition_requirement() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(10)); + let proj_exprs: Vec<(Arc, String)> = vec![ + (Arc::new(Column::new("a", 0)), "a".to_string()), + (Arc::new(Column::new("b", 1)), "b".to_string()), + ]; + let projection: Arc = + Arc::new(ProjectionExec::try_new(proj_exprs, source).unwrap()); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let output_req: Arc = Arc::new(OutputRequirementExec::new( + projection, + Some(OrderingRequirements::from(sort_expr)), + Distribution::SinglePartition, + Some(21), + )); + + assert_idempotent(output_req); +} + +/// Regression for #21973 (the issue this PR fixes). +/// +/// Topology: `SPM → Sort(preserve=true) → multi-partition`. Under the old +/// two-rule pipeline `EnforceSorting::pushdown_sorts` could mutate +/// `preserve_partitioning` after `EnforceDistribution` had settled +/// distribution, so pass 2 could regress this parallel plan into a serial +/// one. This is the exact topology that caused +/// [`test_pushdown_through_spm`](../enforce_sorting.rs) to fail before this +/// PR; `EnsureRequirements` must keep the SPM-over-parallel-sort shape +/// stable across passes. +/// +/// Also acts as the "no extra SPM when already optimal" check — the +/// input plan already contains exactly one `SortPreservingMergeExec`, +/// so the optimised plan must too (we used to have a separate test for +/// this property, but on this input it is implied by idempotency +/// combined with the SPM-count assertion). +#[test] +fn test_issue_21973_idempotent_spm_sort_multi_partition() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(10)); + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new( + SortExec::new(sort_expr.clone(), source).with_preserve_partitioning(true), + ); + let spm: Arc = + Arc::new(SortPreservingMergeExec::new(sort_expr, sort)); + let limit: Arc = Arc::new(GlobalLimitExec::new(spm, 0, Some(21))); + + // No-extra-SPM property: count SortPreservingMergeExec occurrences in + // the first optimisation pass — must be ≤ 1 (the original SPM survives, + // none are added). + let config = test_config(); + let optimized = EnsureRequirements::new() + .optimize(Arc::clone(&limit), &config) + .expect("optimize failed"); + let plan_str = plan_string(&optimized); + let spm_count = plan_str.matches("SortPreservingMergeExec").count(); + assert!( + spm_count <= 1, + "Extra SortPreservingMergeExec added ({spm_count} found):\n{plan_str}" + ); + + assert_idempotent(limit); +} + +/// Regression for #21973: the `parallelize_sorts` rewrite path. +/// +/// Input: `Sort(DESC) ← CoalescePartitionsExec ← multi-partition`. The first +/// pass must rewrite this into a parallel plan +/// `SortPreservingMergeExec ← Sort(preserve=true) ← multi-partition`, and +/// subsequent passes must keep that parallel shape. Under the old two-rule +/// pipeline `pushdown_sorts` could regress this back into a serial sort. +/// +/// Runs 3 passes (one more than the standard idempotency check) and asserts +/// the parallel plan structure survives each one. +#[test] +fn test_issue_21973_parallel_sort_survives_multiple_passes() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(8)); + let coalesce: Arc = Arc::new(CoalescePartitionsExec::new(source)); + + let sort_expr = sort_expr_on("a", 0, true, true); + let sort: Arc = Arc::new(SortExec::new(sort_expr, coalesce)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(10))); + + let config = test_config(); + + let p1 = EnsureRequirements::new() + .optimize(Arc::clone(&limit), &config) + .expect("pass 1"); + let s1 = plan_string(&p1); + + let p2 = EnsureRequirements::new() + .optimize(Arc::clone(&p1), &config) + .expect("pass 2"); + let s2 = plan_string(&p2); + + let p3 = EnsureRequirements::new() + .optimize(Arc::clone(&p2), &config) + .expect("pass 3"); + let s3 = plan_string(&p3); + + // The parallel-sort shape must appear after pass 1 and survive every + // subsequent pass. Specifically: SortPreservingMergeExec on top of a + // Sort with preserve_partitioning=true (no CoalescePartitionsExec). + for (i, plan_str) in [&s1, &s2, &s3].iter().enumerate() { + assert!( + plan_str.contains("SortPreservingMergeExec"), + "pass {} regressed to serial: missing SortPreservingMergeExec:\n{plan_str}", + i + 1 + ); + assert!( + plan_str.contains("preserve_partitioning=[true]"), + "pass {} regressed to serial: Sort lost preserve_partitioning=true (#21973):\n{plan_str}", + i + 1 + ); + assert!( + !plan_str.contains("CoalescePartitionsExec"), + "pass {} regressed to serial: CoalescePartitionsExec re-introduced (#21973):\n{plan_str}", + i + 1 + ); + } + + assert_eq!( + s1, s2, + "not idempotent between pass 1 and 2 (#21973):\n{s1}\nvs\n{s2}" + ); + assert_eq!( + s2, s3, + "not idempotent between pass 2 and 3 (#21973):\n{s2}\nvs\n{s3}" + ); + + // All passes must produce sanity-checkable plans. + for (i, p) in [p1, p2, p3].into_iter().enumerate() { + SanityCheckPlan::new() + .optimize(p, &config) + .unwrap_or_else(|e| { + panic!("SanityCheckPlan failed on pass {}: {e:?}", i + 1) + }); + } +} + +/// Idempotency: Sort → Aggregate → Sort → Aggregate pattern (#18989). +/// This tests the multi-aggregate topology that caused SanityCheckPlan +/// failures in upstream issue #18989. +#[test] +fn test_idempotent_sort_aggregate_sort_aggregate() { + use datafusion_physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, + }; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])); + + let source: Arc = Arc::new(MockMultiPartitionExec::new(4)); + + // Partial aggregate + let group_by = PhysicalGroupBy::new_single(vec![( + Arc::new(Column::new("a", 0)) as _, + "a".to_string(), + )]); + let partial_agg: Arc = Arc::new( + AggregateExec::try_new( + AggregateMode::Partial, + group_by, + vec![], + vec![], + source, + Arc::clone(&schema), + ) + .unwrap(), + ); + + let sort_expr = sort_expr_on("a", 0, false, false); + + let sort: Arc = Arc::new(SortExec::new(sort_expr, partial_agg)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(10))); + + assert_idempotent(limit); +} + +/// Stress test: idempotency with ALL partition counts from 1 to 64 +#[test] +fn test_idempotent_all_partition_counts_1_to_64() { + for n in 1..=64 { + let source = Arc::new(MockMultiPartitionExec::new(n)); + let sort_expr = sort_expr_on("a", 0, true, true); + let sort = Arc::new(SortExec::new(sort_expr, source)); + let limit: Arc = + Arc::new(GlobalLimitExec::new(sort, 0, Some(10))); + assert_idempotent(limit); + } +} + +/// Regression for #14150: the standalone distribution enforcement path +/// lost `fetch` when applied twice. Verify `EnsureRequirements` +/// preserves fetch across multiple passes. +#[test] +fn test_issue_14150_fetch_survives_multiple_passes() { + // Simulate: SELECT * FROM multi_partition_table ORDER BY a LIMIT 5 + // with target_partitions > 1 (triggers RepartitionExec) + let source: Arc = Arc::new(MockMultiPartitionExec::new(1)); + let repartition = Arc::new( + RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(4)).unwrap(), + ); + + let sort_expr = sort_expr_on("a", 0, false, true); + + let sort = Arc::new(SortExec::new(sort_expr, repartition as _).with_fetch(Some(5))); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(5))); + + let config = test_config(); + + // Pass 1 + let p1 = EnsureRequirements::new() + .optimize(Arc::clone(&limit), &config) + .unwrap(); + let s1 = plan_string(&p1); + + // Pass 2 + let p2 = EnsureRequirements::new() + .optimize(Arc::clone(&p1), &config) + .unwrap(); + let s2 = plan_string(&p2); + + // Pass 3 + let p3 = EnsureRequirements::new() + .optimize(Arc::clone(&p2), &config) + .unwrap(); + let s3 = plan_string(&p3); + + // Fetch must survive all passes + assert!(s1.contains("fetch=5"), "fetch=5 lost after pass 1:\n{s1}"); + assert!( + s2.contains("fetch=5"), + "fetch=5 lost after pass 2 (#14150 regression):\n{s2}" + ); + assert!(s3.contains("fetch=5"), "fetch=5 lost after pass 3:\n{s3}"); + + // Plans must be identical (idempotent) + assert_eq!(s1, s2, "Plan changed between pass 1 and 2:\n{s1}\nvs\n{s2}"); + assert_eq!(s2, s3, "Plan changed between pass 2 and 3:\n{s2}\nvs\n{s3}"); +} + +/// Sharper #14150 reproduce: input plan already contains a +/// `SortPreservingMergeExec` with an explicit `fetch`, sitting directly +/// above a `SortExec(fetch=…)` on a multi-partition source. This is the +/// exact shape that originally triggered the bug — the old +/// `EnforceDistribution::optimize` path would call +/// `remove_dist_changing_operators()` on this SPM, strip it, and then +/// `add_merge_on_top()` re-create an SPM **without** copying the saved +/// `fetch`. Pass 2 saw an SPM with no fetch and #14150 silently bit. +/// +/// `EnsureRequirements` preserves the `fetch` value across every pass. +/// Note: it may legitimately deduplicate the `fetch` field between +/// adjacent operators (e.g. push it onto the surrounding +/// `GlobalLimitExec` and drop it from the SPM), so this test asserts +/// the #14150 property — \"`fetch=5` must appear somewhere in the +/// plan after every pass\" — rather than byte-identical idempotency +/// (which is covered by `test_issue_14150_fetch_survives_multiple_passes` +/// on the more realistic input shape where the SPM is inserted by the +/// optimizer itself). +#[test] +fn test_issue_14150_fetch_survives_with_input_spm() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(4)); + + let sort_expr = sort_expr_on("a", 0, false, true); + + // Sort with fetch=5 (TopK). + let sort = Arc::new( + SortExec::new(sort_expr.clone(), Arc::clone(&source)).with_fetch(Some(5)), + ); + + // SPM with fetch=5 above the sort — this is what `EnforceDistribution` + // used to strip and re-add without `fetch`. + let spm: Arc = + Arc::new(SortPreservingMergeExec::new(sort_expr, sort).with_fetch(Some(5))); + + let limit: Arc = Arc::new(GlobalLimitExec::new(spm, 0, Some(5))); + + let config = test_config(); + + let p1 = EnsureRequirements::new() + .optimize(Arc::clone(&limit), &config) + .unwrap(); + let s1 = plan_string(&p1); + + let p2 = EnsureRequirements::new() + .optimize(Arc::clone(&p1), &config) + .unwrap(); + let s2 = plan_string(&p2); + + // The #14150 property: `fetch=5` must survive both passes (the + // historical bug was that pass 2 dropped it when the SPM got + // re-created in `add_merge_on_top`). + assert!(s1.contains("fetch=5"), "fetch=5 lost after pass 1:\n{s1}"); + assert!( + s2.contains("fetch=5"), + "fetch=5 lost after pass 2 (#14150 regression):\n{s2}" + ); +} + +// ======================================================================== +// Mock operator with configurable distribution / ordering requirements +// (used by window-function and distribution tests below) +// ======================================================================== + +/// Mock operator requiring specific distribution and/or ordering from its +/// single child. Simulates operators like `BoundedWindowAggExec` that +/// demand hash-partitioning + ordering without pulling in complex window +/// expression machinery. +#[derive(Debug)] +struct MockReqExec { + input: Arc, + dist: Distribution, + ord: Option, + properties: Arc, +} + +impl MockReqExec { + fn new( + input: Arc, + dist: Distribution, + ord: Option, + ) -> Self { + let properties = Arc::new(PlanProperties::new( + input.equivalence_properties().clone(), + input.output_partitioning().clone(), + input.pipeline_behavior(), + input.boundedness(), + )); + Self { + input, + dist, + ord, + properties, + } + } +} + +impl DisplayAs for MockReqExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "MockReqExec") + } +} + +impl ExecutionPlan for MockReqExec { + fn name(&self) -> &str { + "MockReqExec" + } + fn properties(&self) -> &Arc { + &self.properties + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + self.dist.clone(), + ]) + } + fn required_input_ordering(&self) -> Vec> { + vec![ + self.ord + .as_ref() + .map(|o| OrderingRequirements::from(o.clone())), + ] + } + fn maintains_input_order(&self) -> Vec { + vec![true] + } + fn replace_children( + self: Arc, + mut children: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + assert_eq!(children.len(), 1); + Ok(Arc::new(MockReqExec::new( + children.pop().expect("1 child"), + self.dist.clone(), + self.ord.clone(), + ))) + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( + &self, + _p: usize, + _c: Arc, + ) -> Result { + unimplemented!() + } +} + +// ======================================================================== +// Additional idempotency tests covering remaining sub-passes +// ======================================================================== + +/// Idempotency: plan that triggers `parallelize_sorts`. +/// CoalescePartitionsExec → SortExec(preserve=false) → multi-partition +/// source. After optimization Sort+SPM should be parallel. Running twice +/// must produce the same plan. +#[test] +fn test_idempotent_parallelize_sorts() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(8)); + let coalesce: Arc = Arc::new(CoalescePartitionsExec::new(source)); + + let sort_expr = sort_expr_on("a", 0, true, true); + + // Sort without preserve_partitioning on top of coalesced input + let sort: Arc = Arc::new(SortExec::new(sort_expr, coalesce)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(10))); + + // First pass should parallelize the sort (Sort+SPM replaces Coalesce+Sort) + let config = test_config(); + let p1 = EnsureRequirements::new() + .optimize(Arc::clone(&limit), &config) + .expect("first optimize failed"); + let p2 = EnsureRequirements::new() + .optimize(Arc::clone(&p1), &config) + .expect("second optimize failed"); + + let s1 = plan_string(&p1); + let s2 = plan_string(&p2); + assert_eq!( + s1, s2, + "parallelize_sorts NOT idempotent!\nFirst:\n{s1}\nSecond:\n{s2}" + ); + + SanityCheckPlan::new() + .optimize(p1, &config) + .expect("SanityCheckPlan failed on first pass"); + SanityCheckPlan::new() + .optimize(p2, &config) + .expect("SanityCheckPlan failed on second pass"); +} + +/// Idempotency: SortMergeJoinExec with two multi-partition inputs + ORDER BY +/// + LIMIT. Tests that join key reordering + sort enforcement is stable. +#[test] +fn test_idempotent_sort_merge_join() { + let left: Arc = Arc::new(MockMultiPartitionExec::new(4)); + let right: Arc = Arc::new(MockMultiPartitionExec::new(4)); + + let on = vec![( + Arc::new(Column::new("a", 0)) as Arc, + Arc::new(Column::new("a", 0)) as Arc, + )]; + + let join: Arc = Arc::new( + SortMergeJoinExec::try_new( + left, + right, + on, + None, + JoinType::Inner, + vec![SortOptions { + descending: false, + nulls_first: false, + }], + NullEquality::NullEqualsNothing, + ) + .expect("SortMergeJoinExec creation failed"), + ); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, join)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(50))); + + assert_idempotent(limit); +} + +/// Idempotency: window-function-like operator over multi-partition source. +/// Uses MockReqExec with hash distribution + ordering to simulate +/// BoundedWindowAggExec requirements. Tests that window partitioning + +/// sort requirements are stable across optimizer passes. +#[test] +fn test_idempotent_window_over_multi_partition() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(8)); + + // Window function requires Hash(a) distribution + ordering [a ASC, b ASC] + let ord = LexOrdering::new(vec![ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions { + descending: false, + nulls_first: false, + }, + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 1)), + SortOptions { + descending: false, + nulls_first: false, + }, + ), + ]) + .unwrap(); + + let dist = Distribution::KeyPartitioned(vec![Arc::new(Column::new("a", 0))]); + let window_like: Arc = + Arc::new(MockReqExec::new(source, dist, Some(ord))); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, window_like)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(20))); + + assert_idempotent(limit); +} + +/// Idempotency: multiple levels of sort + limit. +/// GlobalLimitExec → SortExec → ProjectionExec → GlobalLimitExec → SortExec +/// → multi-partition source. Tests deeply nested sort/limit stability. +#[test] +fn test_idempotent_nested_subqueries_sort_limit() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(8)); + + // Inner sort (a DESC) + inner limit — use DESC to avoid matching + // MockMultiPartitionExec's built-in ASC ordering, which would cause + // the optimizer to eliminate the sort differently across passes. + let inner_sort: Arc = + Arc::new(SortExec::new(sort_expr_on("a", 0, true, true), source)); + let inner_limit: Arc = + Arc::new(GlobalLimitExec::new(inner_sort, 0, Some(100))); + + // Identity projection + let proj_exprs: Vec<(Arc, String)> = vec![ + (Arc::new(Column::new("a", 0)), "a".to_string()), + (Arc::new(Column::new("b", 1)), "b".to_string()), + ]; + let projection: Arc = + Arc::new(ProjectionExec::try_new(proj_exprs, inner_limit).unwrap()); + + // Outer sort (a DESC) + outer limit + let outer_sort: Arc = + Arc::new(SortExec::new(sort_expr_on("a", 0, true, true), projection)); + let outer_limit: Arc = + Arc::new(GlobalLimitExec::new(outer_sort, 0, Some(10))); + + assert_idempotent(outer_limit); +} + +/// Idempotency: RepartitionExec(Hash) + sort + limit. +/// Tests that hash distribution + ordering enforcement is stable. +#[test] +fn test_idempotent_repartition_hash_sort_limit() { + let source: Arc = Arc::new(MockMultiPartitionExec::new(8)); + + let hash_exprs: Vec> = vec![Arc::new(Column::new("a", 0))]; + let repartition = Arc::new( + RepartitionExec::try_new(source, Partitioning::Hash(hash_exprs, 4)).unwrap(), + ); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort: Arc = Arc::new(SortExec::new(sort_expr, repartition)); + let limit: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(15))); + + assert_idempotent(limit); +} + +/// EnsureRequirements applied twice on a HashJoinExec plan must produce +/// identical plans. Tests that hash distribution enforcement is stable. +#[test] +fn test_enforce_distribution_idempotent_hash_join() { + let left: Arc = Arc::new(MockMultiPartitionExec::new(4)); + let right: Arc = Arc::new(MockMultiPartitionExec::new(4)); + + let on = vec![( + Arc::new(Column::new("a", 0)) as Arc, + Arc::new(Column::new("a", 0)) as Arc, + )]; + + let join: Arc = Arc::new( + HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .expect("HashJoinExec creation failed"), + ); + + let config = test_config(); + let p1 = EnsureRequirements::new() + .optimize(Arc::clone(&join), &config) + .expect("first EnsureRequirements pass failed"); + let p2 = EnsureRequirements::new() + .optimize(Arc::clone(&p1), &config) + .expect("second EnsureRequirements pass failed"); + + let s1 = plan_string(&p1); + let s2 = plan_string(&p2); + + assert_eq!( + s1, s2, + "EnsureRequirements not idempotent for HashJoinExec!\nPass 1:\n{s1}\nPass 2:\n{s2}" + ); +} + +/// Idempotency on a complex plan: +/// `GlobalLimitExec → SortExec → ProjectionExec → UnionExec(multi, single)`. +/// The union + projection + sort topology has historically been a fertile +/// ground for non-idempotent behaviour, so we keep it as a separate +/// idempotency test — `assert_idempotent` already proves `f(f(x)) == f(x)`, +/// which for a deterministic optimiser is equivalent to stability across +/// any finite number of passes (the previous 10x sweep was overkill). +#[test] +fn test_idempotent_union_projection_sort() { + let live: Arc = Arc::new(MockMultiPartitionExec::new(16)); + let hist: Arc = Arc::new(MockMultiPartitionExec::new(1)); + let union: Arc = UnionExec::try_new(vec![live, hist]).unwrap(); + + // Identity projection + let proj_exprs: Vec<(Arc, String)> = vec![ + (Arc::new(Column::new("a", 0)), "a".to_string()), + (Arc::new(Column::new("b", 1)), "b".to_string()), + ]; + let projection: Arc = + Arc::new(ProjectionExec::try_new(proj_exprs, union).unwrap()); + + let sort_expr = sort_expr_on("a", 0, true, true); + + let sort = Arc::new(SortExec::new(sort_expr, projection)); + let plan: Arc = Arc::new(GlobalLimitExec::new(sort, 0, Some(21))); + + assert_idempotent(plan); +} + +/// Builds the plan shape that phase 3a (`parallelize_sorts`) sees in the reproducer, +/// i.e. the output of the distribution + sorting phases, not a freshly planned tree: +/// +/// ```text +/// CoalescePartitionsExec <- the node `parallelize_sorts` rewrites +/// HashJoinExec: mode=CollectLeft +/// CoalescePartitionsExec <- satisfies `SinglePartition` on the build side +/// +/// RepartitionExec: RoundRobinBatch +/// CoalescePartitionsExec <- links the join into the coalesce cascade +/// MockMultiPartitionExec +/// ``` +/// +/// Both coalesces below the join matter. The probe-side one is what makes +/// `update_coalesce_ctx_children` mark the join as connected — it only skips children that +/// require `SinglePartition`, and the probe side does not — so the walk descends into the +/// join. The build-side one is the one that must survive. +fn collect_left_plan_before_parallelize_sorts( + build: Arc, + join_type: JoinType, +) -> Result> { + let build: Arc = Arc::new(CoalescePartitionsExec::new(build)); + let probe: Arc = Arc::new(RepartitionExec::try_new( + Arc::new(CoalescePartitionsExec::new(Arc::new( + MockMultiPartitionExec::new(4), + ))), + Partitioning::RoundRobinBatch(TEST_TARGET_PARTITIONS), + )?); + + let on = vec![( + Arc::new(Column::new("a", 0)) as Arc, + Arc::new(Column::new("a", 0)) as Arc, + )]; + let join: Arc = Arc::new(HashJoinExec::try_new( + build, + probe, + on, + None, + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?); + + Ok(Arc::new(CoalescePartitionsExec::new(join))) +} + +/// Runs phase 3a of `EnsureRequirements` (`parallelize_sorts`) on its own, the same way +/// the rule drives it, and checks the result with `SanityCheckPlan`. +/// +/// The phase is driven directly rather than through `EnsureRequirements::optimize` because +/// the earlier phases would rebuild the plan shape above into something that never reaches +/// the code path under test. +fn parallelize_sorts_and_sanity_check( + plan: Arc, +) -> Result> { + let ctx = PlanWithCorrespondingCoalescePartitions::new_default(plan); + let rewritten = ctx.transform_up(parallelize_sorts).data()?.plan; + SanityCheckPlan::new().optimize(Arc::clone(&rewritten), &test_config())?; + Ok(rewritten) +} + +/// A `CollectLeft` `HashJoinExec` requires `Distribution::SinglePartition` on its build +/// (left) child, so the distribution phase puts a `CoalescePartitionsExec` on top of a +/// multi-partition build side. The sort-parallelization phase must not take that coalesce +/// back out again. +/// +/// It used to, because `remove_bottleneck_in_subplan` removed a coalesce found at +/// `children[0]` positionally, without consulting the parent's distribution requirement for +/// that child. The result was a build side left multi-partition with nothing to re-enforce +/// distribution afterwards, which `SanityCheckPlan` rejected with "does not satisfy +/// distribution requirements: SinglePartition". +#[test] +fn test_collect_left_join_keeps_build_side_coalesce() -> Result<()> { + let plan = collect_left_plan_before_parallelize_sorts( + Arc::new(MockMultiPartitionExec::new(4)), + JoinType::Left, + )?; + + let rewritten = parallelize_sorts_and_sanity_check(plan)?; + + // The build-side coalesce is retained; the probe-side one is still removed, which is + // the parallelization this phase exists for. + assert_snapshot!(plan_string(&rewritten), @r" + CoalescePartitionsExec + HashJoinExec: mode=CollectLeft, join_type=Left, on=[(a@0, a@0)] + CoalescePartitionsExec + MockMultiPartitionExec + RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4 + MockMultiPartitionExec + "); + + Ok(()) +} + +/// The same removal, with a build side that is already hash-partitioned on the join key +/// rather than `UnknownPartitioning`. This is the shape a `JoinSelection` input swap leaves +/// behind (a `CollectLeft` join reported as `join_type=Right`) when the build subtree is the +/// output of an aggregate or a partitioned join: the build side satisfies the join's *hash* +/// requirement but still not `SinglePartition`, so the coalesce is just as load-bearing. +#[test] +fn test_collect_left_join_keeps_hash_partitioned_build_side_coalesce() -> Result<()> { + let plan = collect_left_plan_before_parallelize_sorts( + Arc::new(MockMultiPartitionExec::hash_partitioned_on_a( + TEST_TARGET_PARTITIONS, + )), + JoinType::Right, + )?; + + let rewritten = parallelize_sorts_and_sanity_check(plan)?; + + assert_snapshot!(plan_string(&rewritten), @r" + CoalescePartitionsExec + HashJoinExec: mode=CollectLeft, join_type=Right, on=[(a@0, a@0)] + CoalescePartitionsExec + MockMultiPartitionExec + RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4 + MockMultiPartitionExec + "); + + Ok(()) +} diff --git a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs index f56b8c6d70624..a98c1b7bcf98b 100644 --- a/datafusion/core/tests/physical_optimizer/filter_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/filter_pushdown.rs @@ -18,7 +18,7 @@ use std::sync::{Arc, LazyLock}; use arrow::{ - array::record_batch, + array::{RecordBatch, record_batch}, datatypes::{DataType, Field, Schema, SchemaRef}, util::pretty::pretty_format_batches, }; @@ -34,17 +34,29 @@ use datafusion::{ scalar::ScalarValue, }; use datafusion_catalog::memory::DataSourceExec; -use datafusion_common::config::ConfigOptions; +use datafusion_common::{ + JoinType, + config::ConfigOptions, + tree_node::{TreeNode, TreeNodeRecursion}, +}; use datafusion_datasource::{ PartitionedFile, file_groups::FileGroup, file_scan_config::FileScanConfigBuilder, }; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_expr::ScalarUDF; use datafusion_functions::math::random::RandomFunc; -use datafusion_functions_aggregate::{count::count_udaf, min_max::min_udaf}; -use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, expressions::col}; +use datafusion_functions_aggregate::{ + count::count_udaf, + min_max::{max_udaf, min_udaf}, +}; +use datafusion_physical_expr::{ + LexOrdering, PhysicalSortExpr, + expressions::{DynamicFilterPhysicalExpr, col}, + utils::conjunction, +}; use datafusion_physical_expr::{ - Partitioning, ScalarFunctionExpr, aggregate::AggregateExprBuilder, + Partitioning, RangePartitioning, ScalarFunctionExpr, SplitPoint, + aggregate::AggregateExprBuilder, }; use datafusion_physical_optimizer::{ PhysicalOptimizerRule, filter_pushdown::FilterPushdown, @@ -55,6 +67,7 @@ use datafusion_physical_plan::{ coalesce_partitions::CoalescePartitionsExec, collect, filter::{FilterExec, FilterExecBuilder}, + joins::{HashJoinExec, PartitionMode}, projection::ProjectionExec, repartition::RepartitionExec, sorts::sort::SortExec, @@ -175,9 +188,6 @@ fn test_pushdown_into_scan_with_config_options() { // distinction this test exercises is not reachable via SQL. #[tokio::test] async fn test_static_filter_pushdown_through_hash_join() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - // Create build side with limited values let build_batches = vec![ record_batch!( @@ -738,6 +748,65 @@ fn test_pushdown_through_aggregates_on_grouping_columns() { ); } +#[test] +fn test_pushdown_through_aggregates_preserves_parent_filter_order() { + // AggregateExec may push filters on grouping columns to its input, but must + // keep filters on aggregate outputs above itself. The parent-filter result + // order must match the incoming filter order, otherwise an unsupported + // aggregate-output filter can be reported as pushed down and removed. + let scan = TestScanBuilder::new(schema()).with_support(true).build(); + + let aggregate_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("a", &schema()).unwrap()]) + .schema(schema()) + .alias("cnt") + .build() + .map(Arc::new) + .unwrap(), + ]; + let group_by = PhysicalGroupBy::new_single(vec![ + (col("a", &schema()).unwrap(), "a".to_string()), + (col("b", &schema()).unwrap(), "b".to_string()), + ]); + let aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Final, + group_by, + aggregate_expr, + vec![None], + scan, + schema(), + ) + .unwrap(), + ); + + let aggregate_schema = aggregate.schema(); + let aggregate_output_filter = col_lit_predicate( + "cnt", + ScalarValue::Int64(Some(1)), + aggregate_schema.as_ref(), + ); + let grouping_key_filter = col_lit_predicate("b", "bar", aggregate_schema.as_ref()); + let predicate = conjunction(vec![aggregate_output_filter, grouping_key_filter]); + let plan = Arc::new(FilterExec::try_new(predicate, aggregate).unwrap()); + + insta::assert_snapshot!( + OptimizationTest::new(plan, FilterPushdown::new(), true), + @r" + OptimizationTest: + input: + - FilterExec: cnt@2 = 1 AND b@1 = bar + - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt] + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true + output: + Ok: + - FilterExec: cnt@2 = 1 + - AggregateExec: mode=Final, gby=[a@0 as a, b@1 as b], aggr=[cnt], ordering_mode=PartiallySorted([1]) + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=b@1 = bar + " + ); +} + /// Test various combinations of handling of child pushdown results /// in an ExecutionPlan in combination with support/not support in a DataSource. #[test] @@ -874,15 +943,73 @@ async fn test_topk_filter_passes_through_coalesce_partitions() { ); } +fn hashjoin_pushdown_scans() -> ( + SchemaRef, + Arc, + SchemaRef, + Arc, +) { + let build_side_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Utf8, false), + Field::new("c", DataType::Float64, false), + ])); + let build_scan = TestScanBuilder::new(Arc::clone(&build_side_schema)) + .with_support(true) + .with_batches(vec![ + record_batch!( + ("a", Utf8, ["aa", "ab"]), + ("b", Utf8, ["ba", "bb"]), + ("c", Float64, [1.0, 2.0]) + ) + .unwrap(), + ]) + .build(); + + let probe_side_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Utf8, false), + Field::new("e", DataType::Float64, false), + ])); + let probe_scan = TestScanBuilder::new(Arc::clone(&probe_side_schema)) + .with_support(true) + .with_batches(vec![ + record_batch!( + ("a", Utf8, ["aa", "ab", "ac", "ad"]), + ("b", Utf8, ["ba", "bb", "bc", "bd"]), + ("e", Float64, [1.0, 2.0, 3.0, 4.0]) + ) + .unwrap(), + ]) + .build(); + + (build_side_schema, build_scan, probe_side_schema, probe_scan) +} + +async fn optimize_and_collect_pushdown_plan( + plan: Arc, + config: ConfigOptions, +) -> (Arc, Vec) { + let plan = FilterPushdown::new_post_optimization() + .optimize(plan, &config) + .unwrap(); + let session_ctx = + SessionContext::new_with_config(SessionConfig::from(config).with_batch_size(10)); + session_ctx.register_object_store( + ObjectStoreUrl::parse("test://").unwrap().as_ref(), + Arc::new(InMemory::new()), + ); + let task_ctx = session_ctx.state().task_ctx(); + let batches = collect(Arc::clone(&plan), task_ctx).await.unwrap(); + (plan, batches) +} + // Not portable to sqllogictest: this test pins `PartitionMode::Partitioned` // by hand-wiring `RepartitionExec(Hash, 12)` on both join sides. A SQL // INNER JOIN over small parquet inputs plans as `CollectLeft`, so the // per-partition CASE filter this test exercises is not reachable via SQL. #[tokio::test] async fn test_hashjoin_dynamic_filter_pushdown_partitioned() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - // Rough sketch of the MRE we're trying to recreate: // COPY (select i as k from generate_series(1, 10000000) as t(i)) // TO 'test_files/scratch/push_down_filter/t1.parquet' @@ -923,43 +1050,8 @@ async fn test_hashjoin_dynamic_filter_pushdown_partitioned() { // | | | // +---------------+------------------------------------------------------------+ - // Create build side with limited values - let build_batches = vec![ - record_batch!( - ("a", Utf8, ["aa", "ab"]), - ("b", Utf8, ["ba", "bb"]), - ("c", Float64, [1.0, 2.0]) // Extra column not used in join - ) - .unwrap(), - ]; - let build_side_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("c", DataType::Float64, false), - ])); - let build_scan = TestScanBuilder::new(Arc::clone(&build_side_schema)) - .with_support(true) - .with_batches(build_batches) - .build(); - - // Create probe side with more values - let probe_batches = vec![ - record_batch!( - ("a", Utf8, ["aa", "ab", "ac", "ad"]), - ("b", Utf8, ["ba", "bb", "bc", "bd"]), - ("e", Float64, [1.0, 2.0, 3.0, 4.0]) // Extra column not used in join - ) - .unwrap(), - ]; - let probe_side_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("e", DataType::Float64, false), - ])); - let probe_scan = TestScanBuilder::new(Arc::clone(&probe_side_schema)) - .with_support(true) - .with_batches(probe_batches) - .build(); + let (build_side_schema, build_scan, probe_side_schema, probe_scan) = + hashjoin_pushdown_scans(); // Create RepartitionExec nodes for both sides with hash partitioning on join keys let partition_count = 12; @@ -1057,20 +1149,7 @@ async fn test_hashjoin_dynamic_filter_pushdown_partitioned() { let mut config = ConfigOptions::default(); config.execution.parquet.pushdown_filters = true; config.optimizer.enable_dynamic_filter_pushdown = true; - let plan = FilterPushdown::new_post_optimization() - .optimize(plan, &config) - .unwrap(); - let config = SessionConfig::new().with_batch_size(10); - let session_ctx = SessionContext::new_with_config(config); - session_ctx.register_object_store( - ObjectStoreUrl::parse("test://").unwrap().as_ref(), - Arc::new(InMemory::new()), - ); - let state = session_ctx.state(); - let task_ctx = state.task_ctx(); - let batches = collect(Arc::clone(&plan), Arc::clone(&task_ctx)) - .await - .unwrap(); + let (plan, batches) = optimize_and_collect_pushdown_plan(plan, config).await; // Now check what our filter looks like #[cfg(not(feature = "force_hash_collisions"))] @@ -1127,53 +1206,214 @@ async fn test_hashjoin_dynamic_filter_pushdown_partitioned() { ); } -// Not portable to sqllogictest: this test specifically pins a -// `RepartitionExec(Hash, 12)` between `HashJoinExec(CollectLeft)` and the -// probe-side scan to verify the dynamic filter link survives that boundary -// (regression for #17451). The same CollectLeft filter content and -// pushdown counters are already covered by the simpler slt port -// (push_down_filter_parquet.slt::test_hashjoin_dynamic_filter_pushdown). #[tokio::test] -async fn test_hashjoin_dynamic_filter_pushdown_collect_left() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; +async fn test_hashjoin_dynamic_filter_pushdown_range_partitioned() { + // Rough sketch of the Range-partitioned MRE we're trying to recreate. The + // test hand-wires identical Range repartitioning: + // + // EXPLAIN + // SELECT * + // FROM build + // JOIN probe + // ON build.a = probe.a AND build.b = probe.b; + // + // +---------------+------------------------------------------------------------+ + // | plan_type | plan | + // +---------------+------------------------------------------------------------+ + // | physical_plan | ┌───────────────────────────┐ | + // | | │ HashJoinExec │ | + // | | │ -------------------- ├──────────────┐ | + // | | │ on: (a = a), (b = b) │ │ | + // | | └─────────────┬─────────────┘ │ | + // | | ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ | + // | | │ RepartitionExec ││ RepartitionExec │ | + // | | │ -------------------- ││ -------------------- │ | + // | | │ partition_count(in->out): ││ partition_count(in->out): │ | + // | | │ 1 -> 2 ││ 1 -> 2 │ | + // | | │ ││ │ | + // | | │ partitioning_scheme: ││ partitioning_scheme: │ | + // | | │ Range([a ASC, b ASC], 2) ││ Range([a ASC, b ASC], 2) │ | + // | | │ split: (aa, bb) ││ split: (aa, bb) │ | + // | | └─────────────┬─────────────┘└─────────────┬─────────────┘ | + // | | ┌─────────────┴─────────────┐┌─────────────┴─────────────┐ | + // | | │ DataSourceExec (build) ││ DataSourceExec (probe) │ | + // | | │ -------------------- ││ -------------------- │ | + // | | │ rows: (aa,ba), (ab,bb) ││ rows: (aa,ba) ... (ad,bd) │ | + // | | │ ││ predicate: DynamicFilter │ | + // | | │ ││ range CASE -> filter_0/1 │ | + // | | └───────────────────────────┘└───────────────────────────┘ | + // | | | + // +---------------+------------------------------------------------------------+ - let build_batches = vec![ - record_batch!( - ("a", Utf8, ["aa", "ab"]), - ("b", Utf8, ["ba", "bb"]), - ("c", Float64, [1.0, 2.0]) // Extra column not used in join + let (build_side_schema, build_scan, probe_side_schema, probe_scan) = + hashjoin_pushdown_scans(); + + let split_points = vec![SplitPoint::new(vec![ + ScalarValue::Utf8(Some("aa".to_string())), + ScalarValue::Utf8(Some("bb".to_string())), + ])]; + + // Build side: DataSource -> RepartitionExec (Range) + let build_range_ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new( + col("a", &build_side_schema).unwrap(), + SortOptions::default(), + ), + PhysicalSortExpr::new( + col("b", &build_side_schema).unwrap(), + SortOptions::default(), + ), + ]) + .unwrap(); + let build_repartition = Arc::new( + RepartitionExec::try_new( + build_scan, + Partitioning::Range( + RangePartitioning::try_new(build_range_ordering, split_points.clone()) + .unwrap(), + ), ) .unwrap(), - ]; - let build_side_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("c", DataType::Float64, false), - ])); - let build_scan = TestScanBuilder::new(Arc::clone(&build_side_schema)) - .with_support(true) - .with_batches(build_batches) - .build(); + ); - // Create probe side with more values - let probe_batches = vec![ - record_batch!( - ("a", Utf8, ["aa", "ab", "ac", "ad"]), - ("b", Utf8, ["ba", "bb", "bc", "bd"]), - ("e", Float64, [1.0, 2.0, 3.0, 4.0]) // Extra column not used in join + // Probe side: DataSource -> RepartitionExec (Range) + let probe_range_ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new( + col("a", &probe_side_schema).unwrap(), + SortOptions::default(), + ), + PhysicalSortExpr::new( + col("b", &probe_side_schema).unwrap(), + SortOptions::default(), + ), + ]) + .unwrap(); + let probe_repartition = Arc::new( + RepartitionExec::try_new( + Arc::clone(&probe_scan), + Partitioning::Range( + RangePartitioning::try_new(probe_range_ordering, split_points).unwrap(), + ), ) .unwrap(), + ); + + // Create HashJoinExec with partitioned inputs + let on = vec![ + ( + col("a", &build_side_schema).unwrap(), + col("a", &probe_side_schema).unwrap(), + ), + ( + col("b", &build_side_schema).unwrap(), + col("b", &probe_side_schema).unwrap(), + ), ]; - let probe_side_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("e", DataType::Float64, false), - ])); - let probe_scan = TestScanBuilder::new(Arc::clone(&probe_side_schema)) - .with_support(true) - .with_batches(probe_batches) - .build(); + let hash_join = Arc::new( + HashJoinExec::try_new( + build_repartition, + probe_repartition, + on, + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + datafusion_common::NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + + // Top-level CoalescePartitionsExec + let cp = Arc::new(CoalescePartitionsExec::new(hash_join)) as Arc; + // Add a sort for deterministic output + let plan = Arc::new(SortExec::new( + LexOrdering::new(vec![PhysicalSortExpr::new( + col("a", &probe_side_schema).unwrap(), + SortOptions::new(true, false), // descending, nulls_first + )]) + .unwrap(), + cp, + )) as Arc; + + // expect the predicate to be pushed down into the probe side DataSource + insta::assert_snapshot!( + OptimizationTest::new(Arc::clone(&plan), FilterPushdown::new_post_optimization(), true), + @r" + OptimizationTest: + input: + - SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false] + - CoalescePartitionsExec + - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0), (b@1, b@1)] + - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1 + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true + - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1 + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true + output: + Ok: + - SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false] + - CoalescePartitionsExec + - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0), (b@1, b@1)] + - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1 + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true + - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1 + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ empty ] + " + ); + + // Actually apply the optimization to the plan and execute to see the filter in action + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + config.optimizer.enable_dynamic_filter_pushdown = true; + config.optimizer.preserve_file_partitions = 1; + let (plan, batches) = optimize_and_collect_pushdown_plan(plan, config).await; + + // Now check what our filter looks like + insta::assert_snapshot!( + format!("{}", format_plan_for_test(&plan)), + @r" + - SortExec: expr=[a@0 DESC NULLS LAST], preserve_partitioning=[false] + - CoalescePartitionsExec + - HashJoinExec: mode=Partitioned, join_type=Inner, on=[(a@0, a@0), (b@1, b@1)] + - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1 + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true + - RepartitionExec: partitioning=Range([a@0 ASC, b@1 ASC], [(aa, bb)], 2), input_partitions=1 + - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, e], file_type=test, pushdown_supported=true, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN a@0 >= aa AND a@0 <= aa AND b@1 >= ba AND b@1 <= ba AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}]) ELSE a@0 >= ab AND a@0 <= ab AND b@1 >= bb AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:ab,c1:bb}]) END ] + " + ); + + let result = format!("{}", pretty_format_batches(&batches).unwrap()); + + let probe_scan_metrics = probe_scan.metrics().unwrap(); + + // The probe side had 4 rows, but after applying the dynamic filter only 2 rows should remain. + // The number of output rows from the probe side scan should stay consistent across executions. + // Issue: https://github.com/apache/datafusion/issues/17451 + assert_eq!(probe_scan_metrics.output_rows().unwrap(), 2); + + insta::assert_snapshot!( + result, + @r" + +----+----+-----+----+----+-----+ + | a | b | c | a | b | e | + +----+----+-----+----+----+-----+ + | ab | bb | 2.0 | ab | bb | 2.0 | + | aa | ba | 1.0 | aa | ba | 1.0 | + +----+----+-----+----+----+-----+ + ", + ); +} + +// Not portable to sqllogictest: this test specifically pins a +// `RepartitionExec(Hash, 12)` between `HashJoinExec(CollectLeft)` and the +// probe-side scan to verify the dynamic filter link survives that boundary +// (regression for #17451). The same CollectLeft filter content and +// pushdown counters are already covered by the simpler slt port +// (push_down_filter_parquet.slt::test_hashjoin_dynamic_filter_pushdown). +#[tokio::test] +async fn test_hashjoin_dynamic_filter_pushdown_collect_left() { + let (build_side_schema, build_scan, probe_side_schema, probe_scan) = + hashjoin_pushdown_scans(); // Create RepartitionExec nodes for both sides with hash partitioning on join keys let partition_count = 12; @@ -1255,20 +1495,7 @@ async fn test_hashjoin_dynamic_filter_pushdown_collect_left() { let mut config = ConfigOptions::default(); config.execution.parquet.pushdown_filters = true; config.optimizer.enable_dynamic_filter_pushdown = true; - let plan = FilterPushdown::new_post_optimization() - .optimize(plan, &config) - .unwrap(); - let config = SessionConfig::new().with_batch_size(10); - let session_ctx = SessionContext::new_with_config(config); - session_ctx.register_object_store( - ObjectStoreUrl::parse("test://").unwrap().as_ref(), - Arc::new(InMemory::new()), - ); - let state = session_ctx.state(); - let task_ctx = state.task_ctx(); - let batches = collect(Arc::clone(&plan), Arc::clone(&task_ctx)) - .await - .unwrap(); + let (plan, batches) = optimize_and_collect_pushdown_plan(plan, config).await; // Now check what our filter looks like insta::assert_snapshot!( @@ -1307,9 +1534,6 @@ async fn test_hashjoin_dynamic_filter_pushdown_collect_left() { #[test] fn test_hashjoin_parent_filter_pushdown_same_column_names() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - let build_side_schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Utf8, false), Field::new("build_val", DataType::Utf8, false), @@ -1376,9 +1600,6 @@ fn test_hashjoin_parent_filter_pushdown_same_column_names() { #[test] fn test_hashjoin_parent_filter_pushdown_mark_join() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - let left_schema = Arc::new(Schema::new(vec![ Field::new("id", DataType::Utf8, false), Field::new("val", DataType::Utf8, false), @@ -1442,15 +1663,10 @@ fn test_hashjoin_parent_filter_pushdown_mark_join() { ); } -/// Test that filters on join key columns are pushed to both sides of semi/anti joins. -/// For LeftSemi/LeftAnti, the output only contains left columns, but filters on -/// join key columns can also be pushed to the right (non-preserved) side because -/// the equijoin condition guarantees the key values match. +/// Semi-join key filters can be pushed to both sides, but anti-join filters must +/// only rely on the output side to preserve their semantics. #[test] fn test_hashjoin_parent_filter_pushdown_semi_anti_join() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - let left_schema = Arc::new(Schema::new(vec![ Field::new("k", DataType::Utf8, false), Field::new("v", DataType::Utf8, false), @@ -1474,9 +1690,9 @@ fn test_hashjoin_parent_filter_pushdown_semi_anti_join() { let join = Arc::new( HashJoinExec::try_new( - left_scan, - right_scan, - on, + left_scan, + Arc::clone(&right_scan), + on.clone(), None, &JoinType::LeftSemi, None, @@ -1515,6 +1731,24 @@ fn test_hashjoin_parent_filter_pushdown_semi_anti_join() { - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[k, w], file_type=test, pushdown_supported=true, predicate=k@0 = x " ); + + let join = Arc::new( + HashJoinExec::try_new( + TestScanBuilder::new(Arc::clone(&left_schema)).build(), + right_scan, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::Partitioned, + datafusion_common::NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + let predicate = Arc::new(Literal::new(ScalarValue::Boolean(Some(false)))); + let plan = Arc::new(FilterExec::try_new(predicate, join).unwrap()); + assert_parent_filter_remains(plan); } #[test] @@ -1753,6 +1987,16 @@ fn col_lit_predicate( )) } +fn assert_parent_filter_remains(plan: Arc) { + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + let optimized = FilterPushdown::new().optimize(plan, &config).unwrap(); + assert!( + optimized.downcast_ref::().is_some(), + "parent filter must remain" + ); +} + // ==== Aggregate Dynamic Filter tests ==== // // The end-to-end min/max dynamic filter cases (simple/min/max/mixed/all-nulls) @@ -1997,13 +2241,65 @@ fn test_pushdown_grouping_sets_filter_on_common_column() { ); } +#[tokio::test] +async fn test_no_pushdown_through_global_aggregate_with_name_collision() { + let input_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let scan = TestScanBuilder::new(Arc::clone(&input_schema)) + .with_support(true) + .with_batches(vec![record_batch!(("a", Int64, [1, 20])).unwrap()]) + .build(); + let aggregate_expr = vec![ + AggregateExprBuilder::new(max_udaf(), vec![col("a", &input_schema).unwrap()]) + .schema(Arc::clone(&input_schema)) + .alias("a") + .build() + .map(Arc::new) + .unwrap(), + ]; + let aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new_single(vec![]), + aggregate_expr, + vec![None], + scan, + input_schema, + ) + .unwrap(), + ); + + // This is a physical filter above the aggregate, not a SQL WHERE clause. + // Pushing it through would evaluate input `a` instead of MAX(a). + let predicate = Arc::new(BinaryExpr::new( + col("a", aggregate.schema().as_ref()).unwrap(), + Operator::Lt, + Arc::new(Literal::new(ScalarValue::Int64(Some(10)))), + )); + let plan = Arc::new(FilterExec::try_new(predicate, aggregate).unwrap()); + + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + let optimized = FilterPushdown::new().optimize(plan, &config).unwrap(); + assert!(optimized.downcast_ref::().is_some()); + + let session_ctx = SessionContext::new(); + session_ctx.register_object_store( + ObjectStoreUrl::parse("test://").unwrap().as_ref(), + Arc::new(InMemory::new()), + ); + let batches = collect(optimized, session_ctx.state().task_ctx()) + .await + .unwrap(); + assert!( + batches.is_empty(), + "MAX(a) = 20 must be filtered out instead of applying a < 10 to input rows" + ); +} + #[test] -fn test_pushdown_with_empty_group_by() { - // Test that filters can be pushed down when GROUP BY is empty (no grouping columns) - // SELECT count(*) as cnt FROM table WHERE a = 'foo' - // There are no grouping columns, so the filter should still push down +fn test_no_pushdown_constant_false_through_global_aggregate() { let scan = TestScanBuilder::new(schema()).with_support(true).build(); - let aggregate_expr = vec![ AggregateExprBuilder::new(count_udaf(), vec![col("c", &schema()).unwrap()]) .schema(schema()) @@ -2012,41 +2308,58 @@ fn test_pushdown_with_empty_group_by() { .map(Arc::new) .unwrap(), ]; + let aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new_single(vec![]), + aggregate_expr, + vec![None], + scan, + schema(), + ) + .unwrap(), + ); + let predicate = Arc::new(Literal::new(ScalarValue::Boolean(Some(false)))); + let plan = Arc::new(FilterExec::try_new(predicate, aggregate).unwrap()); - // Empty GROUP BY - no grouping columns - let group_by = PhysicalGroupBy::new_single(vec![]); + assert_parent_filter_remains(plan); +} +#[test] +fn test_no_pushdown_constant_false_through_empty_grouping_set() { + let scan = TestScanBuilder::new(schema()).with_support(true).build(); + let group_by = PhysicalGroupBy::new( + vec![(col("a", &schema()).unwrap(), "a".to_string())], + vec![( + Arc::new(Literal::new(ScalarValue::Utf8(None))), + "a".to_string(), + )], + vec![vec![true]], + true, + ); + let aggregate_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("c", &schema()).unwrap()]) + .schema(schema()) + .alias("cnt") + .build() + .map(Arc::new) + .unwrap(), + ]; let aggregate = Arc::new( AggregateExec::try_new( AggregateMode::Final, group_by, - aggregate_expr.clone(), + aggregate_expr, vec![None], scan, schema(), ) .unwrap(), ); - - // Filter on 'a' - let predicate = col_lit_predicate("a", "foo", &schema()); + let predicate = Arc::new(Literal::new(ScalarValue::Boolean(Some(false)))); let plan = Arc::new(FilterExec::try_new(predicate, aggregate).unwrap()); - // The filter should be pushed down even with empty GROUP BY - insta::assert_snapshot!( - OptimizationTest::new(plan, FilterPushdown::new(), true), - @r" - OptimizationTest: - input: - - FilterExec: a@0 = foo - - AggregateExec: mode=Final, gby=[], aggr=[cnt] - - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true - output: - Ok: - - AggregateExec: mode=Final, gby=[], aggr=[cnt] - - DataSourceExec: file_groups={1 group: [[test.parquet]]}, projection=[a, b, c], file_type=test, pushdown_supported=true, predicate=a@0 = foo - " - ); + assert_parent_filter_remains(plan); } #[test] @@ -2309,9 +2622,6 @@ fn test_pushdown_with_computed_grouping_key() { // on a hand-wired plan, which does trigger the `false` path. #[tokio::test] async fn test_hashjoin_dynamic_filter_all_partitions_empty() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - // Test scenario where all build-side partitions are empty // This validates the code path that sets the filter to `false` when no rows can match @@ -2444,46 +2754,8 @@ async fn test_hashjoin_dynamic_filter_all_partitions_empty() { // PartitionMode::Partitioned, which SQL never picks for small parquet inputs. #[tokio::test] async fn test_hashjoin_hash_table_pushdown_partitioned() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - - // Create build side with limited values - let build_batches = vec![ - record_batch!( - ("a", Utf8, ["aa", "ab"]), - ("b", Utf8, ["ba", "bb"]), - ("c", Float64, [1.0, 2.0]) // Extra column not used in join - ) - .unwrap(), - ]; - let build_side_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("c", DataType::Float64, false), - ])); - let build_scan = TestScanBuilder::new(Arc::clone(&build_side_schema)) - .with_support(true) - .with_batches(build_batches) - .build(); - - // Create probe side with more values - let probe_batches = vec![ - record_batch!( - ("a", Utf8, ["aa", "ab", "ac", "ad"]), - ("b", Utf8, ["ba", "bb", "bc", "bd"]), - ("e", Float64, [1.0, 2.0, 3.0, 4.0]) // Extra column not used in join - ) - .unwrap(), - ]; - let probe_side_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("e", DataType::Float64, false), - ])); - let probe_scan = TestScanBuilder::new(Arc::clone(&probe_side_schema)) - .with_support(true) - .with_batches(probe_batches) - .build(); + let (build_side_schema, build_scan, probe_side_schema, probe_scan) = + hashjoin_pushdown_scans(); // Create RepartitionExec nodes for both sides with hash partitioning on join keys let partition_count = 12; @@ -2553,24 +2825,11 @@ async fn test_hashjoin_hash_table_pushdown_partitioned() { )) as Arc; // Apply the optimization with config setting that forces HashTable strategy - let session_config = SessionConfig::default() - .with_batch_size(10) - .set_usize("datafusion.optimizer.hash_join_inlist_pushdown_max_size", 1) - .set_bool("datafusion.execution.parquet.pushdown_filters", true) - .set_bool("datafusion.optimizer.enable_dynamic_filter_pushdown", true); - let plan = FilterPushdown::new_post_optimization() - .optimize(plan, session_config.options()) - .unwrap(); - let session_ctx = SessionContext::new_with_config(session_config); - session_ctx.register_object_store( - ObjectStoreUrl::parse("test://").unwrap().as_ref(), - Arc::new(InMemory::new()), - ); - let state = session_ctx.state(); - let task_ctx = state.task_ctx(); - let batches = collect(Arc::clone(&plan), Arc::clone(&task_ctx)) - .await - .unwrap(); + let mut config = ConfigOptions::default(); + config.optimizer.hash_join_inlist_pushdown_max_size = 1; + config.execution.parquet.pushdown_filters = true; + config.optimizer.enable_dynamic_filter_pushdown = true; + let (plan, batches) = optimize_and_collect_pushdown_plan(plan, config).await; // Verify that hash_lookup is used instead of IN (SET) let plan_str = format_plan_for_test(&plan).to_string(); @@ -2610,45 +2869,8 @@ async fn test_hashjoin_hash_table_pushdown_partitioned() { // IN (SET) invariant is captured in the slt port. #[tokio::test] async fn test_hashjoin_hash_table_pushdown_collect_left() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - - let build_batches = vec![ - record_batch!( - ("a", Utf8, ["aa", "ab"]), - ("b", Utf8, ["ba", "bb"]), - ("c", Float64, [1.0, 2.0]) // Extra column not used in join - ) - .unwrap(), - ]; - let build_side_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("c", DataType::Float64, false), - ])); - let build_scan = TestScanBuilder::new(Arc::clone(&build_side_schema)) - .with_support(true) - .with_batches(build_batches) - .build(); - - // Create probe side with more values - let probe_batches = vec![ - record_batch!( - ("a", Utf8, ["aa", "ab", "ac", "ad"]), - ("b", Utf8, ["ba", "bb", "bc", "bd"]), - ("e", Float64, [1.0, 2.0, 3.0, 4.0]) // Extra column not used in join - ) - .unwrap(), - ]; - let probe_side_schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("e", DataType::Float64, false), - ])); - let probe_scan = TestScanBuilder::new(Arc::clone(&probe_side_schema)) - .with_support(true) - .with_batches(probe_batches) - .build(); + let (build_side_schema, build_scan, probe_side_schema, probe_scan) = + hashjoin_pushdown_scans(); // Create RepartitionExec nodes for both sides with hash partitioning on join keys let partition_count = 12; @@ -2704,24 +2926,11 @@ async fn test_hashjoin_hash_table_pushdown_collect_left() { )) as Arc; // Apply the optimization with config setting that forces HashTable strategy - let session_config = SessionConfig::default() - .with_batch_size(10) - .set_usize("datafusion.optimizer.hash_join_inlist_pushdown_max_size", 1) - .set_bool("datafusion.execution.parquet.pushdown_filters", true) - .set_bool("datafusion.optimizer.enable_dynamic_filter_pushdown", true); - let plan = FilterPushdown::new_post_optimization() - .optimize(plan, session_config.options()) - .unwrap(); - let session_ctx = SessionContext::new_with_config(session_config); - session_ctx.register_object_store( - ObjectStoreUrl::parse("test://").unwrap().as_ref(), - Arc::new(InMemory::new()), - ); - let state = session_ctx.state(); - let task_ctx = state.task_ctx(); - let batches = collect(Arc::clone(&plan), Arc::clone(&task_ctx)) - .await - .unwrap(); + let mut config = ConfigOptions::default(); + config.optimizer.hash_join_inlist_pushdown_max_size = 1; + config.execution.parquet.pushdown_filters = true; + config.optimizer.enable_dynamic_filter_pushdown = true; + let (plan, batches) = optimize_and_collect_pushdown_plan(plan, config).await; // Verify that hash_lookup is used instead of IN (SET) let plan_str = format_plan_for_test(&plan).to_string(); @@ -2755,19 +2964,30 @@ async fn test_hashjoin_hash_table_pushdown_collect_left() { ); } -// Not portable to sqllogictest: asserts on `HashJoinExec::dynamic_filter_for_test().is_used()` -// which is a debug-only API. The observable behavior (probe-side scan -// receiving the dynamic filter when the data source supports it) is -// already covered by the simpler CollectLeft port in push_down_filter_parquet.slt; -// the with_support(false) branch has no SQL analog (parquet always supports -// pushdown). -#[tokio::test] -async fn test_hashjoin_dynamic_filter_pushdown_is_used() { - use datafusion_common::JoinType; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; +// Not portable to sqllogictest: verifies whether the optimized probe-side plan +// retains the HashJoinExec's dynamic filter expression. The with_support(false) +// branch has no SQL analog because parquet supports filter pushdown. +#[test] +fn test_hashjoin_dynamic_filter_pushdown_is_used() { + fn contains_expression_id(plan: &Arc, expression_id: u64) -> bool { + let mut found = false; + plan.apply(|node| { + node.apply_expressions(&mut |root| { + root.apply(|expr| { + if expr.expression_id() == Some(expression_id) { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + }) + }) + .unwrap(); + found + } - // Test both cases: probe side with and without filter pushdown support - for (probe_supports_pushdown, expected_is_used) in [(false, false), (true, true)] { + for (probe_supports_pushdown, expected_consumer) in [(false, false), (true, true)] { let build_side_schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Utf8, false), Field::new("b", DataType::Utf8, false), @@ -2820,31 +3040,26 @@ async fn test_hashjoin_dynamic_filter_pushdown_is_used() { .unwrap(), ) as Arc; - // Apply filter pushdown optimization let mut config = ConfigOptions::default(); config.execution.parquet.pushdown_filters = true; config.optimizer.enable_dynamic_filter_pushdown = true; let plan = FilterPushdown::new_post_optimization() .optimize(plan, &config) .unwrap(); - - // Get the HashJoinExec to check the dynamic filter let hash_join = plan .downcast_ref::() .expect("Plan should be HashJoinExec"); + let dynamic_filters = hash_join.dynamic_expressions_produced(); + let expression_id = dynamic_filters + .first() + .expect("Dynamic filter should be created") + .expression_id() + .expect("Dynamic filters always have an expression ID"); - // Verify that a dynamic filter was created - let dynamic_filter = hash_join - .dynamic_filter_expr() - .expect("Dynamic filter should be created"); - - // Verify that is_used() returns the expected value based on probe side support. - // When probe_supports_pushdown=false: no consumer holds a reference (is_used=false) - // When probe_supports_pushdown=true: probe side holds a reference (is_used=true) assert_eq!( - dynamic_filter.is_used(), - expected_is_used, - "is_used() should return {expected_is_used} when probe side support is {probe_supports_pushdown}" + contains_expression_id(hash_join.right(), expression_id), + expected_consumer, + "probe consumer should be {expected_consumer} when pushdown support is {probe_supports_pushdown}" ); } } @@ -2949,11 +3164,6 @@ async fn test_filter_with_projection_pushdown() { /// counting nodes. Neither API is observable from SQL. #[tokio::test] async fn test_discover_dynamic_filters_via_expressions_api() { - use datafusion_common::JoinType; - use datafusion_common::tree_node::TreeNodeRecursion; - use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; - use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode}; - fn count_dynamic_filters(plan: &Arc) -> usize { let mut count = 0; @@ -3047,6 +3257,72 @@ async fn test_discover_dynamic_filters_via_expressions_api() { ); } +#[test] +fn test_discover_dynamic_expression_producers() { + fn producer_count(plan: &Arc) -> usize { + let mut count = 0; + plan.apply(|node| { + count += node.dynamic_expressions_produced().len(); + Ok(TreeNodeRecursion::Continue) + }) + .expect("plan traversal should succeed"); + count + } + + let build_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Int32, false), + ])); + let build_scan = TestScanBuilder::new(Arc::clone(&build_schema)) + .with_support(true) + .with_batches(vec![ + record_batch!(("a", Utf8, ["foo", "bar"]), ("b", Int32, [1, 2])).unwrap(), + ]) + .build(); + + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("c", DataType::Float64, false), + ])); + let probe_scan = TestScanBuilder::new(Arc::clone(&probe_schema)) + .with_support(true) + .with_batches(vec![ + record_batch!( + ("a", Utf8, ["foo", "bar", "baz", "qux"]), + ("c", Float64, [1.0, 2.0, 3.0, 4.0]) + ) + .unwrap(), + ]) + .build(); + + let plan = Arc::new( + HashJoinExec::try_new( + build_scan, + probe_scan, + vec![( + col("a", &build_schema).unwrap(), + col("a", &probe_schema).unwrap(), + )], + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + datafusion_common::NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ) as Arc; + assert_eq!(producer_count(&plan), 0); + + let mut config = ConfigOptions::default(); + config.optimizer.enable_dynamic_filter_pushdown = true; + config.execution.parquet.pushdown_filters = true; + let optimized_plan = FilterPushdown::new_post_optimization() + .optimize(plan, &config) + .unwrap(); + assert_eq!(producer_count(&optimized_plan), 1); +} + // ==== Filter pushdown through SortExec tests ==== /// FilterExec above a plain SortExec (no fetch) should be pushed below it. @@ -3374,3 +3650,41 @@ fn test_filter_pushdown_through_sort_with_projection() { " ); } + +/// `FilterPushdown::new_post_optimization()` must be idempotent. When applied +/// to a HashJoinExec, the rule installs a dynamic filter on the probe-side +/// scan; before the fix in `HashJoinExec::gather_filters_for_pushdown`, every +/// invocation created a *new* `DynamicFilterPhysicalExpr` and ANDed it onto +/// the probe side's existing predicate, producing +/// `DynamicFilter AND DynamicFilter AND ...` after N passes. +/// +/// AQE (datafusion-ballista#1359) re-runs the optimizer chain after every +/// completed stage, so this would compound indefinitely without the guard. +#[test] +fn post_phase_is_idempotent_on_hash_join() { + use crate::physical_optimizer::test_utils::{hash_join_exec, parquet_exec, schema}; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_optimizer::filter_pushdown::FilterPushdown; + use datafusion_physical_plan::get_plan_string; + use datafusion_physical_plan::joins::utils::JoinOn; + + let s = schema(); + let left = parquet_exec(Arc::clone(&s)); + let right = parquet_exec(Arc::clone(&s)); + let join_on: JoinOn = vec![( + Arc::new(Column::new_with_schema("a", &left.schema()).unwrap()), + Arc::new(Column::new_with_schema("a", &right.schema()).unwrap()), + )]; + let plan = hash_join_exec(left, right, join_on, None, &JoinType::Inner).unwrap(); + + let config = ConfigOptions::new(); + let rule = FilterPushdown::new_post_optimization(); + let once = rule.optimize(plan, &config).unwrap(); + let twice = rule.optimize(Arc::clone(&once), &config).unwrap(); + + assert_eq!( + get_plan_string(&once), + get_plan_string(&twice), + "second invocation of FilterPushdown::new_post_optimization mutated the plan", + ); +} diff --git a/datafusion/core/tests/physical_optimizer/join_selection.rs b/datafusion/core/tests/physical_optimizer/join_selection.rs index 050baa9e792e9..265279a8ca1e4 100644 --- a/datafusion/core/tests/physical_optimizer/join_selection.rs +++ b/datafusion/core/tests/physical_optimizer/join_selection.rs @@ -36,16 +36,21 @@ use datafusion_physical_expr::expressions::col; use datafusion_physical_expr::expressions::{BinaryExpr, Column, NegativeExpr}; use datafusion_physical_expr::intervals::utils::check_support; use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; +use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::join_selection::JoinSelection; -use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::displayable; use datafusion_physical_plan::joins::utils::ColumnIndex; use datafusion_physical_plan::joins::utils::JoinFilter; use datafusion_physical_plan::joins::{HashJoinExec, NestedLoopJoinExec, PartitionMode}; use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, ExecutionPlanProperties, ReplaceChildrenOptions, +}; +use datafusion_physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, StatisticsArgs, + StatisticsContext, execution_plan::{Boundedness, EmissionType}, }; @@ -249,23 +254,99 @@ async fn test_join_with_swap() { .expect("The type of the plan should not be changed"); assert_eq!( - swapped_join - .left() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(8192) ); assert_eq!( - swapped_join - .right() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(2097152) ); } +#[tokio::test] +async fn test_join_with_swap_to_sort_preserving_merge_fetch_side() { + let (big, _) = create_big_and_small(); + let top1_input = Arc::new(StatisticsExec::new( + big_statistics(), + Schema::new(vec![Field::new("top_col", DataType::Int32, false)]), + )); + let top1 = Arc::new( + SortPreservingMergeExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new( + "top_col", 0, + )))] + .into(), + top1_input, + ) + .with_fetch(Some(1)), + ); + + let join = Arc::new( + HashJoinExec::try_new( + Arc::clone(&big), + top1, + vec![( + Arc::new(Column::new_with_schema("big_col", &big.schema()).unwrap()), + Arc::new(Column::new("top_col", 0)), + )], + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(), + ); + + let optimized_join = JoinSelection::new() + .optimize(join, &ConfigOptions::new()) + .unwrap(); + let optimized_join = optimized_join + .downcast_ref::() + .map(|projection| projection.input()) + .unwrap_or(&optimized_join); + let swapped_join = optimized_join + .downcast_ref::() + .expect("optimized plan should contain a hash join"); + + let left_spm = swapped_join + .left() + .downcast_ref::() + .expect("SPM fetch side should become the left/build input"); + assert_eq!(left_spm.fetch(), Some(1)); + let statistics_context = StatisticsContext::new(); + assert_eq!( + statistics_context + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) + .unwrap() + .num_rows, + Precision::Inexact(1) + ); + let left_byte_size = statistics_context + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) + .unwrap() + .total_byte_size; + let right_byte_size = big_statistics().total_byte_size; + assert!( + left_byte_size.get_value() < right_byte_size.get_value(), + "SPM fetch side should be estimated smaller than the big side" + ); + assert_eq!( + statistics_context + .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) + .unwrap() + .num_rows, + big_statistics().num_rows + ); +} + #[tokio::test] async fn test_left_join_no_swap() { let (big, small) = create_big_and_small(); @@ -297,17 +378,15 @@ async fn test_left_join_no_swap() { .expect("The type of the plan should not be changed"); assert_eq!( - swapped_join - .left() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(8192) ); assert_eq!( - swapped_join - .right() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(2097152) @@ -348,17 +427,15 @@ async fn test_join_with_swap_semi() { assert_eq!(swapped_join.schema().fields().len(), 1); assert_eq!( - swapped_join - .left() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(8192) ); assert_eq!( - swapped_join - .right() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(2097152) @@ -401,17 +478,15 @@ async fn test_join_with_swap_mark() { assert_eq!(swapped_join.schema().fields().len(), 2); assert_eq!( - swapped_join - .left() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(8192) ); assert_eq!( - swapped_join - .right() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(2097152) @@ -529,17 +604,15 @@ async fn test_join_no_swap() { .expect("The type of the plan should not be changed"); assert_eq!( - swapped_join - .left() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(8192) ); assert_eq!( - swapped_join - .right() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(2097152) @@ -604,17 +677,15 @@ async fn test_nl_join_with_swap(join_type: JoinType) { ); assert_eq!( - swapped_join - .left() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(8192) ); assert_eq!( - swapped_join - .right() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(2097152) @@ -677,17 +748,15 @@ async fn test_nl_join_with_swap_no_proj(join_type: JoinType) { ); assert_eq!( - swapped_join - .left() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.left().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(8192) ); assert_eq!( - swapped_join - .right() - .partition_statistics(None) + StatisticsContext::new() + .compute(swapped_join.right().as_ref(), &StatisticsArgs::new()) .unwrap() .total_byte_size, Precision::Inexact(2097152) @@ -1041,13 +1110,24 @@ impl ExecutionPlan for UnboundedExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -1062,16 +1142,9 @@ impl ExecutionPlan for UnboundedExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } @@ -1152,13 +1225,24 @@ impl ExecutionPlan for StatisticsExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -1167,8 +1251,12 @@ impl ExecutionPlan for StatisticsExec { unimplemented!("This plan only serves for testing statistics") } - fn partition_statistics(&self, partition: Option) -> Result> { - Ok(Arc::new(if partition.is_some() { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::new(if args.partition().is_some() { Statistics::new_unknown(&self.schema) } else { self.stats.clone() @@ -1177,16 +1265,9 @@ impl ExecutionPlan for StatisticsExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } @@ -1229,8 +1310,8 @@ struct TestCase { expecting_swap: bool, } -#[tokio::test] -async fn test_join_with_swap_full() -> Result<()> { +#[test] +fn test_join_with_swap_full() -> Result<()> { // NOTE: Currently, some initial conditions are not viable after join order selection. // For example, full join always comes in partitioned mode. See the warning in // function "swap". If this changes in the future, we should update these tests. @@ -1277,13 +1358,13 @@ async fn test_join_with_swap_full() -> Result<()> { }, ]; for case in cases.into_iter() { - test_join_with_maybe_swap_unbounded_case(case).await? + test_join_with_maybe_swap_unbounded_case(case)? } Ok(()) } -#[tokio::test] -async fn test_cases_without_collect_left_check() -> Result<()> { +#[test] +fn test_cases_without_collect_left_check() -> Result<()> { let mut cases = vec![]; let join_types = vec![JoinType::LeftSemi, JoinType::Inner]; for join_type in join_types { @@ -1370,13 +1451,13 @@ async fn test_cases_without_collect_left_check() -> Result<()> { } for case in cases.into_iter() { - test_join_with_maybe_swap_unbounded_case(case).await? + test_join_with_maybe_swap_unbounded_case(case)? } Ok(()) } -#[tokio::test] -async fn test_not_support_collect_left() -> Result<()> { +#[test] +fn test_not_support_collect_left() -> Result<()> { let mut cases = vec![]; // After [JoinSelection] optimization, these join types cannot run in CollectLeft mode except // [JoinType::LeftSemi] @@ -1425,13 +1506,13 @@ async fn test_not_support_collect_left() -> Result<()> { } for case in cases.into_iter() { - test_join_with_maybe_swap_unbounded_case(case).await? + test_join_with_maybe_swap_unbounded_case(case)? } Ok(()) } -#[tokio::test] -async fn test_not_supporting_swaps_possible_collect_left() -> Result<()> { +#[test] +fn test_not_supporting_swaps_possible_collect_left() -> Result<()> { let mut cases = vec![]; let the_ones_not_support_collect_left = vec![JoinType::Right, JoinType::RightAnti, JoinType::RightSemi]; @@ -1525,12 +1606,12 @@ async fn test_not_supporting_swaps_possible_collect_left() -> Result<()> { } for case in cases.into_iter() { - test_join_with_maybe_swap_unbounded_case(case).await? + test_join_with_maybe_swap_unbounded_case(case)? } Ok(()) } -async fn test_join_with_maybe_swap_unbounded_case(t: TestCase) -> Result<()> { +fn test_join_with_maybe_swap_unbounded_case(t: TestCase) -> Result<()> { let left_unbounded = t.initial_sources_unbounded.0 == SourceType::Unbounded; let right_unbounded = t.initial_sources_unbounded.1 == SourceType::Unbounded; let left_exec = Arc::new(UnboundedExec::new( diff --git a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs index 572ae83540892..b8ebc80348134 100644 --- a/datafusion/core/tests/physical_optimizer/limit_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/limit_pushdown.rs @@ -714,3 +714,123 @@ fn no_limit_preserves_plan_identity() -> Result<()> { Ok(()) } + +#[test] +fn outer_offset_does_not_leak_through_sort_into_inner_limit() -> Result<()> { + // Regression test for https://github.com/apache/datafusion/issues/22489 + // + // When an outer OFFSET is separated from an inner LIMIT by a SortExec + // with different sort keys, the outer skip must not reduce the inner + // fetch. Before the fix, combine_limit merged them, producing + // GlobalLimitExec(skip=1, fetch=7) instead of preserving the inner + // LIMIT 8. + // + // Plan structure: + // GlobalLimitExec: skip=1, fetch=None (outer OFFSET 1) + // SortExec: [c1 DESC] (outer sort — different key) + // GlobalLimitExec: skip=0, fetch=8 (inner LIMIT 8) + // SortExec: [c2 ASC] (inner sort — different key) + // EmptyExec + let schema = create_schema(); + let empty = empty_exec(Arc::clone(&schema)); + + let inner_ordering: LexOrdering = [PhysicalSortExpr { + expr: col("c2", &schema)?, + options: SortOptions::default(), + }] + .into(); + let inner_sort = sort_exec(inner_ordering, empty); + let inner_limit = global_limit_exec(inner_sort, 0, Some(8)); + + let outer_ordering: LexOrdering = [PhysicalSortExpr { + expr: col("c1", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }] + .into(); + let outer_sort = sort_exec(outer_ordering, inner_limit); + let outer_limit = global_limit_exec(outer_sort, 1, None); + + let initial = format_plan(&outer_limit); + insta::assert_snapshot!( + initial, + @r" + GlobalLimitExec: skip=1, fetch=None + SortExec: expr=[c1@0 DESC NULLS LAST], preserve_partitioning=[false] + GlobalLimitExec: skip=0, fetch=8 + SortExec: expr=[c2@1 ASC], preserve_partitioning=[false] + EmptyExec + " + ); + + let after_optimize = + LimitPushdown::new().optimize(outer_limit, &ConfigOptions::new())?; + let optimized = format_plan(&after_optimize); + insta::assert_snapshot!( + optimized, + @r" + GlobalLimitExec: skip=1, fetch=None + SortExec: expr=[c1@0 DESC NULLS LAST], preserve_partitioning=[false] + SortExec: TopK(fetch=8), expr=[c2@1 ASC], preserve_partitioning=[false] + EmptyExec + " + ); + + Ok(()) +} + +#[test] +fn outer_offset_with_same_sort_key_still_pushes_limit() -> Result<()> { + // Companion to outer_offset_does_not_leak_through_sort_into_inner_limit: + // when both sorts use the *same* key, the inner LIMIT should still be + // pushed into the SortExec as TopK. + // + // Plan structure: + // GlobalLimitExec: skip=1, fetch=None (outer OFFSET 1) + // SortExec: [c1 ASC] (outer sort — same key) + // GlobalLimitExec: skip=0, fetch=8 (inner LIMIT 8) + // SortExec: [c1 ASC] (inner sort — same key) + // EmptyExec + let schema = create_schema(); + let empty = empty_exec(Arc::clone(&schema)); + + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col("c1", &schema)?, + options: SortOptions::default(), + }] + .into(); + + let inner_sort = sort_exec(ordering.clone(), empty); + let inner_limit = global_limit_exec(inner_sort, 0, Some(8)); + let outer_sort = sort_exec(ordering, inner_limit); + let outer_limit = global_limit_exec(outer_sort, 1, None); + + let initial = format_plan(&outer_limit); + insta::assert_snapshot!( + initial, + @r" + GlobalLimitExec: skip=1, fetch=None + SortExec: expr=[c1@0 ASC], preserve_partitioning=[false] + GlobalLimitExec: skip=0, fetch=8 + SortExec: expr=[c1@0 ASC], preserve_partitioning=[false] + EmptyExec + " + ); + + let after_optimize = + LimitPushdown::new().optimize(outer_limit, &ConfigOptions::new())?; + let optimized = format_plan(&after_optimize); + insta::assert_snapshot!( + optimized, + @r" + GlobalLimitExec: skip=1, fetch=None + SortExec: expr=[c1@0 ASC], preserve_partitioning=[false] + SortExec: TopK(fetch=8), expr=[c1@0 ASC], preserve_partitioning=[false] + EmptyExec + " + ); + + Ok(()) +} diff --git a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs b/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs index c523b4a752a82..323dfd6183306 100644 --- a/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs +++ b/datafusion/core/tests/physical_optimizer/limited_distinct_aggregation.rs @@ -36,7 +36,7 @@ use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::{ ExecutionPlan, aggregates::{AggregateExec, AggregateMode}, - collect, + collect, displayable, limit::{GlobalLimitExec, LocalLimitExec}, }; @@ -104,6 +104,121 @@ async fn test_partial_final() -> Result<()> { Ok(()) } +// Ensure operator respect the soft limit and stops early: `AggregateExec`'s +// `output_rows` metric should be smaller than then total distinct group count. +#[tokio::test] +async fn limited_distinct_aggregate_stream_respects_soft_limit() -> Result<()> { + // Snapshot for an aggregate operator node from `EXPLAIN ANALYZE`. + // + // Example: In an `EXPLAIN ANALYZE` output + // ```txt + // AggregateExec: mode=partial, limit=10, metrics=[output_rows=100, ...] + // ``` + // we get: + // ```txt + // AggregateRuntimeMetric { + // mode: Partial, + // limit: Some(10), + // output_rows: 100, + // } + // ``` + #[derive(Debug)] + struct AggregateRuntimeMetric { + mode: AggregateMode, + limit: Option, + output_rows: usize, + } + + fn collect_aggregate_runtime_metrics( + plan: &Arc, + metrics: &mut Vec, + ) { + if let Some(agg) = plan.downcast_ref::() { + let output_rows = agg + .metrics() + .and_then(|metrics| metrics.aggregate_by_name().output_rows()) + .expect("AggregateExec should record output_rows after execution"); + + metrics.push(AggregateRuntimeMetric { + mode: *agg.mode(), + limit: agg.limit_options().map(|config| config.limit()), + output_rows, + }); + } + + for child in plan.children() { + collect_aggregate_runtime_metrics(child, metrics); + } + } + + fn aggregate_runtime_metrics( + plan: &Arc, + ) -> Vec { + let mut metrics = vec![]; + collect_aggregate_runtime_metrics(plan, &mut metrics); + metrics + } + + let cfg = SessionConfig::new() + .with_target_partitions(2) + .with_batch_size(10) + .set_bool("datafusion.execution.enable_migration_aggregate", true); + let ctx = SessionContext::new_with_config(cfg); + + let dataframe = ctx + .sql( + "SELECT DISTINCT value % 100000 AS v \ + FROM generate_series(1000000) \ + LIMIT 10", + ) + .await?; + let plan = dataframe.create_physical_plan().await?; + let formatted_plan = displayable(plan.as_ref()).indent(false).to_string(); + assert!( + formatted_plan.contains("AggregateExec: mode=Partial"), + "expected a partial aggregate in plan:\n{formatted_plan}" + ); + assert!( + formatted_plan.contains("AggregateExec: mode=FinalPartitioned"), + "expected a final partitioned aggregate in plan:\n{formatted_plan}" + ); + + let batches = collect(Arc::clone(&plan), ctx.task_ctx()).await?; + assert_eq!( + batches.iter().map(|batch| batch.num_rows()).sum::(), + 10 + ); + + let metrics = aggregate_runtime_metrics(&plan); + let partial = metrics + .iter() + .find(|metric| metric.mode == AggregateMode::Partial) + .expect("expected partial aggregate metrics"); + let final_aggregate = metrics + .iter() + .find(|metric| { + matches!( + metric.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) + }) + .expect("expected final aggregate metrics"); + + assert_eq!(partial.limit, Some(10)); + assert_eq!(final_aggregate.limit, Some(10)); + + assert!( + partial.output_rows <= 100, + "partial aggregate should stop before emitting all distinct groups: {metrics:?}" + ); + assert!( + final_aggregate.output_rows <= 100, + "final aggregate should stop before emitting all distinct groups: {metrics:?}" + ); + + Ok(()) +} + #[tokio::test] async fn test_single_local() -> Result<()> { let source = mock_data()?; diff --git a/datafusion/core/tests/physical_optimizer/mod.rs b/datafusion/core/tests/physical_optimizer/mod.rs index b7ba661d2343a..f3b2884dab188 100644 --- a/datafusion/core/tests/physical_optimizer/mod.rs +++ b/datafusion/core/tests/physical_optimizer/mod.rs @@ -24,11 +24,13 @@ mod combine_partial_final_agg; mod enforce_distribution; mod enforce_sorting; mod enforce_sorting_monotonicity; +mod ensure_requirements; mod filter_pushdown; mod join_selection; #[expect(clippy::needless_pass_by_value)] mod limit_pushdown; mod limited_distinct_aggregation; +mod output_requirements; mod partition_statistics; mod projection_pushdown; mod pushdown_sort; diff --git a/datafusion/core/tests/physical_optimizer/output_requirements.rs b/datafusion/core/tests/physical_optimizer/output_requirements.rs new file mode 100644 index 0000000000000..79b47dc4418a7 --- /dev/null +++ b/datafusion/core/tests/physical_optimizer/output_requirements.rs @@ -0,0 +1,209 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use crate::physical_optimizer::test_utils::{parquet_exec, schema, sort_exec, sort_expr}; + +use arrow::array::{cast::AsArray, record_batch, types::Int32Type}; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::datasource::source::DataSourceExec; +use datafusion::prelude::SessionContext; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_optimizer::optimizer::PhysicalOptimizer; +use datafusion_physical_optimizer::output_requirements::OutputRequirements; +use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; +use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::{ExecutionPlan, collect, displayable, get_plan_string}; + +/// `OutputRequirements::new_add_mode()` must be idempotent: re-applying it to +/// its own output must not stack additional `OutputRequirementExec` wrappers. +/// +/// AQE (datafusion-ballista#1359) re-runs the optimizer chain after every +/// completed stage; without this guarantee, every replan adds another wrapper. +#[test] +fn add_mode_is_idempotent_on_bare_scan() { + // Exercises the path where `require_top_ordering_helper` returns + // `is_changed = false` and the rule adds a default (empty-requirement) + // wrapper. + assert_add_mode_idempotent(parquet_exec(schema())); +} + +#[test] +fn add_mode_is_idempotent_on_sorted_plan() { + // Exercises the path where the helper recognizes a top-level `SortExec` + // and produces a wrapper carrying that ordering requirement + // (`is_changed = true` branch). + let s = schema(); + let ordering: LexOrdering = [sort_expr("a", &s)].into(); + let plan = sort_exec(ordering, parquet_exec(Arc::clone(&s))); + assert_add_mode_idempotent(plan); +} + +#[test] +fn add_mode_is_idempotent_on_scalar_subquery() { + // Exercises the below-root case: the wrapper carrying the ordering lands + // under the `ScalarSubqueryExec`, so the root guard in `require_top_ordering` + // does not fire on the second pass. Without treating the existing wrapper as + // already-handled, the second pass would stamp a redundant empty wrapper on + // top of the subquery. + let s = schema(); + let ordering: LexOrdering = [sort_expr("a", &s)].into(); + let sort = sort_exec(ordering, parquet_exec(Arc::clone(&s))); + + let subqueries = vec![ScalarSubqueryLink { + plan: parquet_exec(Arc::clone(&s)), + index: SubqueryIndex::new(0), + }]; + let plan = Arc::new(ScalarSubqueryExec::new( + sort, + subqueries, + ScalarSubqueryResults::new(1), + )) as Arc; + + assert_add_mode_idempotent(plan); +} + +fn assert_add_mode_idempotent(plan: Arc) { + let config = ConfigOptions::new(); + let rule = OutputRequirements::new_add_mode(); + + let once = rule + .optimize(plan, &config) + .expect("first add-mode optimize pass should succeed"); + let twice = rule + .optimize(Arc::clone(&once), &config) + .expect("second add-mode optimize pass should succeed"); + + assert_eq!( + get_plan_string(&once), + get_plan_string(&twice), + "second invocation of OutputRequirements::new_add_mode mutated the plan", + ); +} + +/// For a `ScalarSubqueryExec` root, `require_top_ordering_helper` descends +/// through the main input (child 0) and wraps the global `SortExec` with an +/// `OutputRequirementExec` carrying its ordering, leaving the subquery child +/// untouched. Without this, the multi-child root is skipped and the query's +/// global ORDER BY requirement is lost. +#[test] +fn require_top_ordering_descends_through_scalar_subquery() { + let s = schema(); + let ordering: LexOrdering = [sort_expr("a", &s)].into(); + let sort = sort_exec(ordering, parquet_exec(Arc::clone(&s))); + + // A subquery child makes `children.len() == 2`, exercising the multi-child path. + let subqueries = vec![ScalarSubqueryLink { + plan: parquet_exec(Arc::clone(&s)), + index: SubqueryIndex::new(0), + }]; + let plan = Arc::new(ScalarSubqueryExec::new( + sort, + subqueries, + ScalarSubqueryResults::new(1), + )) as Arc; + + let optimized = OutputRequirements::new_add_mode() + .optimize(plan, &ConfigOptions::new()) + .expect("add-mode optimize should succeed"); + + insta::assert_snapshot!( + displayable(optimized.as_ref()).indent(true).to_string(), + @r" + ScalarSubqueryExec: subqueries=1 + OutputRequirementExec: order_by=[(a@0, asc)], dist_by=SinglePartition + SortExec: expr=[a@0 ASC], preserve_partitioning=[false] + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet + "); +} + +/// A `ScalarSubqueryExec` plan root must preserve its main input's global +/// ordering end to end. +/// +/// The main input is a `SortPreservingMergeExec` over a two-partition ordered +/// source — the shape federated/custom planners hand to the optimizer: an +/// order-preserving merge with no `SortExec` above it. `OutputRequirements` +/// records the global ORDER BY under the multi-child subquery root, the rest of +/// the pipeline keeps the merge, and executing the optimized plan returns the +/// rows in global order regardless of how the source is partitioned. +#[tokio::test] +async fn scalar_subquery_root_preserves_global_ordering_end_to_end() { + // Two partitions, each already sorted on `a`. Global order requires a sort-preserving merge; + // a plain concatenation would interleave them as 1, 3, 5, 7, 2, 4, 6, 8. + let p1 = record_batch!(("a", Int32, [1, 3, 5, 7])).expect("build partition 1 batch"); + let p2 = record_batch!(("a", Int32, [2, 4, 6, 8])).expect("build partition 2 batch"); + let schema = p1.schema(); + let ordering: LexOrdering = [sort_expr("a", &schema)].into(); + let source = DataSourceExec::from_data_source( + MemorySourceConfig::try_new(&[vec![p1], vec![p2]], Arc::clone(&schema), None) + .expect("build memory source config") + .try_with_sort_information(vec![ordering.clone()]) + .expect("attach sort information to source"), + ); + // The main plan establishes the query's global ordering via an `SortPreservingMergeExec` over the two sorted partitions. + let main_input = Arc::new(SortPreservingMergeExec::new(ordering, source)); + + // Dummy subquery that returns a single row + let sq_batch = record_batch!(("v", Int32, [42])).expect("build subquery batch"); + let subquery = MemorySourceConfig::try_new_exec( + &[vec![sq_batch.clone()]], + sq_batch.schema(), + None, + ) + .expect("build subquery exec"); + + let plan = Arc::new(ScalarSubqueryExec::new( + main_input, + vec![ScalarSubqueryLink { + plan: subquery, + index: SubqueryIndex::new(0), + }], + ScalarSubqueryResults::new(1), + )) as Arc; + + // Run the full default physical optimizer pipeline. + let mut config = ConfigOptions::new(); + config.execution.target_partitions = 4; + let mut optimized = plan; + for rule in PhysicalOptimizer::new().rules { + optimized = rule + .optimize(optimized, &config) + .unwrap_or_else(|e| panic!("optimizer rule {} failed: {e}", rule.name())); + } + + // The executed rows come back in global order: the two sorted partitions + // are merged into 1, 2, 3, 4, 5, 6, 7, 8. + let batches = collect(optimized, SessionContext::new().task_ctx()) + .await + .expect("execute optimized plan"); + let values: Vec = batches + .iter() + .flat_map(|b| { + b.column(0) + .as_primitive::() + .values() + .iter() + .copied() + }) + .collect(); + assert_eq!(values, vec![1, 2, 3, 4, 5, 6, 7, 8]); +} diff --git a/datafusion/core/tests/physical_optimizer/partition_statistics.rs b/datafusion/core/tests/physical_optimizer/partition_statistics.rs index d06e506abfebf..6cabcdb710393 100644 --- a/datafusion/core/tests/physical_optimizer/partition_statistics.rs +++ b/datafusion/core/tests/physical_optimizer/partition_statistics.rs @@ -55,6 +55,7 @@ mod test { use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; + use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::union::{InterleaveExec, UnionExec}; use datafusion_physical_plan::windows::{WindowAggExec, create_window_expr}; use datafusion_physical_plan::{ @@ -238,7 +239,12 @@ mod test { async fn test_statistics_by_partition_of_data_source() -> Result<()> { let scan = create_scan_exec_with_statistics(None, Some(2)).await; let statistics = (0..scan.output_partitioning().partition_count()) - .map(|idx| scan.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + scan.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; // Partition 1: ids [3,4], dates [2025-03-01, 2025-03-02] let expected_statistic_partition_1 = create_partition_statistics( @@ -282,7 +288,12 @@ mod test { let projection: Arc = Arc::new(ProjectionExec::try_new(exprs, scan)?); let statistics = (0..projection.output_partitioning().partition_count()) - .map(|idx| projection.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + projection.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; // Projection only includes id column, not the date partition column let expected_statistic_partition_1 = @@ -314,7 +325,12 @@ mod test { let sort = SortExec::new(ordering.clone().into(), scan_1); let sort_exec: Arc = Arc::new(sort); let statistics = (0..sort_exec.output_partitioning().partition_count()) - .map(|idx| sort_exec.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + sort_exec.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; // All 4 files merged: ids [1-4], dates [2025-03-01, 2025-03-04] let expected_statistic_partition = create_partition_statistics( @@ -353,7 +369,12 @@ mod test { Some((DATE_2025_03_03, DATE_2025_03_04)), ); let statistics = (0..sort_exec.output_partitioning().partition_count()) - .map(|idx| sort_exec.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + sort_exec.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 2); assert_eq!(*statistics[0], expected_statistic_partition_1); @@ -380,7 +401,8 @@ mod test { )?; let filter: Arc = Arc::new(FilterExec::try_new(predicate, scan)?); - let full_statistics = filter.partition_statistics(None)?; + let full_statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; let expected_full_statistic = Statistics { num_rows: Precision::Inexact(0), total_byte_size: Precision::Inexact(0), @@ -406,7 +428,12 @@ mod test { assert_eq!(*full_statistics, expected_full_statistic); let statistics = (0..filter.output_partitioning().partition_count()) - .map(|idx| filter.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + filter.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 2); let expected_partition_statistic = Statistics { @@ -442,7 +469,12 @@ mod test { let union_exec: Arc = UnionExec::try_new(vec![scan.clone(), scan])?; let statistics = (0..union_exec.output_partitioning().partition_count()) - .map(|idx| union_exec.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + union_exec.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; // Check that we have 4 partitions (2 from each scan) assert_eq!(statistics.len(), 4); @@ -505,7 +537,12 @@ mod test { // Verify the result of partition statistics let stats = (0..interleave.output_partitioning().partition_count()) - .map(|idx| interleave.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + interleave.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(stats.len(), 2); @@ -551,15 +588,20 @@ mod test { let cross_join: Arc = Arc::new(CrossJoinExec::new(left_scan, right_scan)); let statistics = (0..cross_join.output_partitioning().partition_count()) - .map(|idx| cross_join.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + cross_join.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; // Check that we have 2 partitions assert_eq!(statistics.len(), 2); // Cross join output schema: [left.id, left.date, right.id] - // Cross join doesn't propagate Column's byte_size let expected_statistic_partition_1 = Statistics { num_rows: Precision::Exact(8), - total_byte_size: Precision::Exact(512), + total_byte_size: Precision::Exact(96), + // Cross join doesn't propagate Column's byte_size column_statistics: vec![ // column 0: left.id (Int32, file column from t1) ColumnStatistics { @@ -593,7 +635,7 @@ mod test { }; let expected_statistic_partition_2 = Statistics { num_rows: Precision::Exact(8), - total_byte_size: Precision::Exact(512), + total_byte_size: Precision::Exact(96), column_statistics: vec![ // column 0: left.id (Int32, file column from t1) ColumnStatistics { @@ -658,45 +700,48 @@ mod test { // Test partition_statistics(None) - returns overall statistics // For RightSemi join, output columns come from right side only - let full_statistics = nested_loop_join.partition_statistics(None)?; + let full_statistics = StatisticsContext::new() + .compute(nested_loop_join.as_ref(), &StatisticsArgs::new())?; // With empty join columns, estimate_join_statistics returns Inexact row count // based on the outer side (right side for RightSemi) - let mut expected_full_statistics = create_partition_statistics( + let expected_full_statistics = create_partition_statistics( 4, 32, 1, 4, Some((DATE_2025_03_01, DATE_2025_03_04)), - ); - expected_full_statistics.num_rows = Precision::Inexact(4); - expected_full_statistics.total_byte_size = Precision::Absent; + ) + .to_inexact(); assert_eq!(*full_statistics, expected_full_statistics); // Test partition_statistics(Some(idx)) - returns partition-specific statistics // Partition 1: ids [3,4], dates [2025-03-01, 2025-03-02] - let mut expected_statistic_partition_1 = create_partition_statistics( + let expected_statistic_partition_1 = create_partition_statistics( 2, 16, 3, 4, Some((DATE_2025_03_01, DATE_2025_03_02)), - ); - expected_statistic_partition_1.num_rows = Precision::Inexact(2); - expected_statistic_partition_1.total_byte_size = Precision::Absent; + ) + .to_inexact(); // Partition 2: ids [1,2], dates [2025-03-03, 2025-03-04] - let mut expected_statistic_partition_2 = create_partition_statistics( + let expected_statistic_partition_2 = create_partition_statistics( 2, 16, 1, 2, Some((DATE_2025_03_03, DATE_2025_03_04)), - ); - expected_statistic_partition_2.num_rows = Precision::Inexact(2); - expected_statistic_partition_2.total_byte_size = Precision::Absent; + ) + .to_inexact(); let statistics = (0..nested_loop_join.output_partitioning().partition_count()) - .map(|idx| nested_loop_join.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + nested_loop_join.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 2); assert_eq!(*statistics[0], expected_statistic_partition_1); @@ -726,7 +771,12 @@ mod test { Some((DATE_2025_03_01, DATE_2025_03_04)), ); let statistics = (0..coalesce_partitions.output_partitioning().partition_count()) - .map(|idx| coalesce_partitions.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + coalesce_partitions.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 1); assert_eq!(*statistics[0], expected_statistic_partition); @@ -743,7 +793,12 @@ mod test { let local_limit: Arc = Arc::new(LocalLimitExec::new(scan.clone(), 1)); let statistics = (0..local_limit.output_partitioning().partition_count()) - .map(|idx| local_limit.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + local_limit.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 2); let mut expected_0 = Statistics::clone(&statistics[0]); @@ -770,7 +825,12 @@ mod test { let global_limit: Arc = Arc::new(GlobalLimitExec::new(scan.clone(), 0, Some(2))); let statistics = (0..global_limit.output_partitioning().partition_count()) - .map(|idx| global_limit.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + global_limit.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 1); // GlobalLimit takes from first partition: ids [3,4], dates [2025-03-01, 2025-03-02] @@ -829,7 +889,10 @@ mod test { @"AggregateExec: mode=Partial, gby=[id@0 as id, 1 + id@0 as expr], aggr=[COUNT(c)]" ); - let p0_statistics = aggregate_exec_partial.partition_statistics(Some(0))?; + let p0_statistics = StatisticsContext::new().compute( + aggregate_exec_partial.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + )?; // Aggregate doesn't propagate num_rows and ColumnStatistics byte_size from input let expected_p0_statistics = Statistics { @@ -868,7 +931,10 @@ mod test { ], }; - let p1_statistics = aggregate_exec_partial.partition_statistics(Some(1))?; + let p1_statistics = StatisticsContext::new().compute( + aggregate_exec_partial.as_ref(), + &StatisticsArgs::new().with_partition(Some(1)), + )?; assert_eq!(*p1_statistics, expected_p1_statistics); validate_statistics_with_data( @@ -890,10 +956,16 @@ mod test { aggregate_exec_partial.schema(), )?); - let p0_statistics = agg_final.partition_statistics(Some(0))?; + let p0_statistics = StatisticsContext::new().compute( + agg_final.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + )?; assert_eq!(*p0_statistics, expected_p0_statistics); - let p1_statistics = agg_final.partition_statistics(Some(1))?; + let p1_statistics = StatisticsContext::new().compute( + agg_final.as_ref(), + &StatisticsArgs::new().with_partition(Some(1)), + )?; assert_eq!(*p1_statistics, expected_p1_statistics); validate_statistics_with_data( @@ -938,8 +1010,20 @@ mod test { ], }; - assert_eq!(empty_stat, *agg_partial.partition_statistics(Some(0))?); - assert_eq!(empty_stat, *agg_partial.partition_statistics(Some(1))?); + assert_eq!( + empty_stat, + *StatisticsContext::new().compute( + agg_partial.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)) + )? + ); + assert_eq!( + empty_stat, + *StatisticsContext::new().compute( + agg_partial.as_ref(), + &StatisticsArgs::new().with_partition(Some(1)) + )? + ); validate_statistics_with_data( agg_partial.clone(), vec![ExpectedStatistics::Empty, ExpectedStatistics::Empty], @@ -965,8 +1049,20 @@ mod test { agg_partial.schema(), )?); - assert_eq!(empty_stat, *agg_final.partition_statistics(Some(0))?); - assert_eq!(empty_stat, *agg_final.partition_statistics(Some(1))?); + assert_eq!( + empty_stat, + *StatisticsContext::new().compute( + agg_final.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)) + )? + ); + assert_eq!( + empty_stat, + *StatisticsContext::new().compute( + agg_final.as_ref(), + &StatisticsArgs::new().with_partition(Some(1)) + )? + ); validate_statistics_with_data( agg_final, @@ -985,6 +1081,50 @@ mod test { scan_schema.clone(), )?); + let expect_partial_stat = Statistics { + num_rows: Precision::Exact(1), + total_byte_size: Precision::Absent, + column_statistics: vec![ColumnStatistics::new_unknown()], + }; + assert_eq!( + expect_partial_stat, + *StatisticsContext::new().compute( + agg_partial.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)) + )? + ); + assert_eq!( + expect_partial_stat, + *StatisticsContext::new().compute( + agg_partial.as_ref(), + &StatisticsArgs::new().with_partition(Some(1)) + )? + ); + + let expect_partial_overall_stat = Statistics { + num_rows: Precision::Exact(2), + total_byte_size: Precision::Absent, + column_statistics: vec![ColumnStatistics::new_unknown()], + }; + assert_eq!( + expect_partial_overall_stat, + *StatisticsContext::new() + .compute(agg_partial.as_ref(), &StatisticsArgs::new())? + ); + + // Verify that the partial aggregate emits one accumulator-state row per + // output partition, even when the corresponding input partitions are empty. + let partitions = execute_stream_partitioned( + agg_partial.clone(), + Arc::new(TaskContext::default()), + )?; + assert_eq!(2, partitions.len()); + for partition_stream in partitions { + let result: Vec = partition_stream.try_collect().await?; + let rows = result.iter().map(|batch| batch.num_rows()).sum::(); + assert_eq!(1, rows); + } + let coalesce = Arc::new(CoalescePartitionsExec::new(agg_partial.clone())); let agg_final = Arc::new(AggregateExec::try_new( @@ -1002,7 +1142,13 @@ mod test { column_statistics: vec![ColumnStatistics::new_unknown()], }; - assert_eq!(expect_stat, *agg_final.partition_statistics(Some(0))?); + assert_eq!( + expect_stat, + *StatisticsContext::new().compute( + agg_final.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)) + )? + ); // Verify that the aggregate final result has exactly one partition with one row let mut partitions = execute_stream_partitioned( @@ -1030,7 +1176,10 @@ mod test { let mut all_batches = vec![]; for (i, partition_stream) in partitions.into_iter().enumerate() { let batches: Vec = partition_stream.try_collect().await?; - let actual = plan.partition_statistics(Some(i))?; + let actual = StatisticsContext::new().compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(i)), + )?; let expected = compute_record_batch_statistics( std::slice::from_ref(&batches), &schema, @@ -1040,7 +1189,8 @@ mod test { all_batches.push(batches); } - let actual = plan.partition_statistics(None)?; + let actual = + StatisticsContext::new().compute(plan.as_ref(), &StatisticsArgs::new())?; let expected = compute_record_batch_statistics(&all_batches, &schema, None); assert_eq!(*actual, expected); @@ -1057,7 +1207,12 @@ mod test { )?); let statistics = (0..repartition.partitioning().partition_count()) - .map(|idx| repartition.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + repartition.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 3); @@ -1108,13 +1263,16 @@ mod test { Partitioning::RoundRobinBatch(2), )?); - let result = repartition.partition_statistics(Some(2)); + let result = StatisticsContext::new().compute( + repartition.as_ref(), + &StatisticsArgs::new().with_partition(Some(2)), + ); assert!(result.is_err()); let error = result.unwrap_err(); assert!( error .to_string() - .contains("RepartitionExec invalid partition 2 (expected less than 2)") + .contains("Invalid partition index: 2, the partition count is 2") ); let partitions = execute_stream_partitioned( @@ -1137,8 +1295,19 @@ mod test { Partitioning::RoundRobinBatch(0), )?); - let result = repartition.partition_statistics(Some(0))?; - assert_eq!(*result, Statistics::new_unknown(&scan_schema)); + // Requesting a specific partition of a zero-partition plan is out of + // range, so the context rejects it. + let result = StatisticsContext::new().compute( + repartition.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ); + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("Invalid partition index: 0, the partition count is 0") + ); // Verify that the result has exactly 0 partitions let partitions = execute_stream_partitioned( @@ -1164,7 +1333,12 @@ mod test { // Verify the result of partition statistics of repartition let stats = (0..repartition.partitioning().partition_count()) - .map(|idx| repartition.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + repartition.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(stats.len(), 2); @@ -1222,7 +1396,12 @@ mod test { // Verify partition statistics are properly propagated (not unknown) let statistics = (0..window_agg.output_partitioning().partition_count()) - .map(|idx| window_agg.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + window_agg.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 2); @@ -1308,7 +1487,10 @@ mod test { // Try to test with single partition let empty_single = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let stats = empty_single.partition_statistics(Some(0))?; + let stats = StatisticsContext::new().compute( + empty_single.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + )?; assert_eq!(stats.num_rows, Precision::Exact(0)); assert_eq!(stats.total_byte_size, Precision::Exact(0)); assert_eq!(stats.column_statistics.len(), 2); @@ -1323,7 +1505,8 @@ mod test { assert_eq!(col_stat.byte_size, Precision::Exact(0)); } - let overall_stats = empty_single.partition_statistics(None)?; + let overall_stats = StatisticsContext::new() + .compute(empty_single.as_ref(), &StatisticsArgs::new())?; assert_eq!(stats, overall_stats); validate_statistics_with_data(empty_single, vec![ExpectedStatistics::Empty], 0) @@ -1334,7 +1517,12 @@ mod test { Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(3)); let statistics = (0..empty_multi.output_partitioning().partition_count()) - .map(|idx| empty_multi.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + empty_multi.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; assert_eq!(statistics.len(), 3); @@ -1394,7 +1582,12 @@ mod test { // Test partition statistics for CollectLeft mode let statistics = (0..collect_left_join.output_partitioning().partition_count()) - .map(|idx| collect_left_join.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + collect_left_join.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; // Check that we have the expected number of partitions @@ -1470,7 +1663,12 @@ mod test { // Test partition statistics for Partitioned mode let statistics = (0..partitioned_join.output_partitioning().partition_count()) - .map(|idx| partitioned_join.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + partitioned_join.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; // Check that we have the expected number of partitions @@ -1544,7 +1742,12 @@ mod test { // Test partition statistics for Auto mode let statistics = (0..auto_join.output_partitioning().partition_count()) - .map(|idx| auto_join.partition_statistics(Some(idx))) + .map(|idx| { + StatisticsContext::new().compute( + auto_join.as_ref(), + &StatisticsArgs::new().with_partition(Some(idx)), + ) + }) .collect::>>()?; // Check that we have the expected number of partitions diff --git a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs index 6f88e01059fc9..3a8d82f111145 100644 --- a/datafusion/core/tests/physical_optimizer/projection_pushdown.rs +++ b/datafusion/core/tests/physical_optimizer/projection_pushdown.rs @@ -25,7 +25,7 @@ use datafusion::datasource::physical_plan::CsvSource; use datafusion::datasource::source::DataSourceExec; use datafusion_common::config::{ConfigOptions, CsvOptions}; use datafusion_common::{JoinSide, JoinType, NullEquality, Result, ScalarValue}; -use datafusion_datasource::TableSchema; +use datafusion_datasource::TableSchemaBuilder; use datafusion_datasource::file_scan_config::FileScanConfigBuilder; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; @@ -724,7 +724,7 @@ fn test_output_req_after_projection() -> Result<()> { ] .into(), )), - Distribution::HashPartitioned(vec![ + Distribution::KeyPartitioned(vec![ Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1)), ]), @@ -746,7 +746,7 @@ fn test_output_req_after_projection() -> Result<()> { actual, @r" ProjectionExec: expr=[c@2 as c, a@0 as new_a, b@1 as b] - OutputRequirementExec: order_by=[(b@1, asc), (c@2 + a@0, asc)], dist_by=HashPartitioned[[a@0, b@1]]) + OutputRequirementExec: order_by=[(b@1, asc), (c@2 + a@0, asc)], dist_by=KeyPartitioned[[a@0, b@1]]) DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=csv, has_header=false " ); @@ -762,7 +762,7 @@ fn test_output_req_after_projection() -> Result<()> { assert_snapshot!( actual, @r" - OutputRequirementExec: order_by=[(b@2, asc), (c@0 + new_a@1, asc)], dist_by=HashPartitioned[[new_a@1, b@2]]) + OutputRequirementExec: order_by=[(b@2, asc), (c@0 + new_a@1, asc)], dist_by=KeyPartitioned[[new_a@1, b@2]]) DataSourceExec: file_groups={1 group: [[x]]}, projection=[c, a@0 as new_a, b], file_type=csv, has_header=false " ); @@ -797,10 +797,12 @@ fn test_output_req_after_projection() -> Result<()> { Arc::new(Column::new("new_a", 1)), Arc::new(Column::new("b", 2)), ]; - if let Distribution::HashPartitioned(vec) = after_optimize + if let Distribution::KeyPartitioned(vec) = after_optimize .downcast_ref::() .unwrap() - .required_input_distribution()[0] + .input_distribution_requirements() + .child_distribution(0) + .unwrap() .clone() { assert!( @@ -809,7 +811,7 @@ fn test_output_req_after_projection() -> Result<()> { .all(|(actual, expected)| actual.eq(&expected)) ); } else { - panic!("Expected HashPartitioned distribution!"); + panic!("Expected KeyPartitioned distribution!"); }; Ok(()) @@ -1574,10 +1576,13 @@ fn partitioned_data_source() -> Arc { quote: b'"', ..Default::default() }; - let table_schema = TableSchema::new( - Arc::clone(&file_schema), - vec![Arc::new(Field::new("partition_col", DataType::Utf8, true))], - ); + let table_schema = TableSchemaBuilder::from(&file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "partition_col", + DataType::Utf8, + true, + ))]) + .build(); let config = FileScanConfigBuilder::new( ObjectStoreUrl::parse("test:///").unwrap(), Arc::new(CsvSource::new(table_schema).with_csv_options(options)), diff --git a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs index 6a5833f43d42c..b72563a942ae3 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_sort.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_sort.rs @@ -24,18 +24,26 @@ //! 4. Early termination is enabled for TopK queries //! 5. Prefix matching works correctly +use arrow::array::{ArrayRef, Int64Array}; +use arrow::datatypes::{DataType, Field, Schema}; +use arrow::record_batch::RecordBatch; +use datafusion::prelude::SessionContext; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{Result, assert_batches_eq}; use datafusion_physical_expr::expressions; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::pushdown_sort::PushdownSort; +use datafusion_physical_plan::collect; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - OptimizationTest, TestScan, coalesce_partitions_exec, parquet_exec, - parquet_exec_with_sort, projection_exec, projection_exec_with_alias, + OptimizationTest, TestScan, coalesce_partitions_exec, inexact_memory_exec, + parquet_exec, parquet_exec_with_sort, projection_exec, projection_exec_with_alias, repartition_exec, schema, simple_projection_exec, sort_exec, sort_exec_with_fetch, - sort_expr, sort_expr_named, test_scan_with_ordering, + sort_exec_with_fetch_and_preserve_partitioning, sort_expr, sort_expr_named, + test_scan_with_ordering, }; #[test] @@ -119,6 +127,91 @@ fn test_sort_with_limit_phase1() { ); } +#[test] +fn test_standalone_inexact_partitioned_topk_adds_global_merge() { + // Inexact pushdown keeps the SortExec. If that SortExec is a + // partition-preserving TopK, it still needs a final merge across partitions + // to preserve the global ORDER BY ... LIMIT semantics. + let schema = schema(); + let a = sort_expr("a", &schema); + let source = Arc::new(TestScan::new(schema.clone(), vec![]).with_partition_count(2)); + + let ordering = LexOrdering::new(vec![a]).unwrap(); + let plan = sort_exec_with_fetch_and_preserve_partitioning(ordering, Some(10), source); + + insta::assert_snapshot!( + OptimizationTest::new(plan, PushdownSort::new(), true), + @r" + OptimizationTest: + input: + - SortExec: TopK(fetch=10), expr=[a@0 ASC], preserve_partitioning=[true] + - TestScan + output: + Ok: + - SortPreservingMergeExec: [a@0 ASC], fetch=10 + - SortExec: TopK(fetch=10), expr=[a@0 ASC], preserve_partitioning=[true] + - TestScan: requested_ordering=[a@0 ASC] + " + ); +} + +#[test] +fn test_standalone_inexact_single_partition_topk_no_global_merge() { + let schema = schema(); + let a = sort_expr("a", &schema); + let source = Arc::new(TestScan::new(schema.clone(), vec![]).with_partition_count(1)); + + let ordering = LexOrdering::new(vec![a]).unwrap(); + let plan = sort_exec_with_fetch_and_preserve_partitioning(ordering, Some(10), source); + + insta::assert_snapshot!( + OptimizationTest::new(plan, PushdownSort::new(), true), + @r" + OptimizationTest: + input: + - SortExec: TopK(fetch=10), expr=[a@0 ASC], preserve_partitioning=[true] + - TestScan + output: + Ok: + - SortExec: TopK(fetch=10), expr=[a@0 ASC], preserve_partitioning=[true] + - TestScan: requested_ordering=[a@0 ASC] + " + ); +} + +#[tokio::test] +async fn test_standalone_inexact_partitioned_topk_returns_global_limit() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let partition_0 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![1, 100])) as ArrayRef], + )?; + let partition_1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vec![2, 3])) as ArrayRef], + )?; + let source = inexact_memory_exec( + &[vec![partition_0], vec![partition_1]], + Arc::clone(&schema), + )?; + + let ordering = LexOrdering::new(vec![sort_expr("a", &schema)]).unwrap(); + let plan = sort_exec_with_fetch_and_preserve_partitioning(ordering, Some(3), source); + + let mut config = ConfigOptions::new(); + config.optimizer.enable_sort_pushdown = true; + let optimized = PushdownSort::new().optimize(plan, &config)?; + + let ctx = SessionContext::new(); + let batches = collect(optimized, ctx.task_ctx()).await?; + + let expected = [ + "+---+", "| a |", "+---+", "| 1 |", "| 2 |", "| 3 |", "+---+", + ]; + assert_batches_eq!(expected, &batches); + Ok(()) +} + #[test] fn test_sort_multiple_columns_phase1() { // Phase 1: Sort on multiple columns - reverse multi-column ordering diff --git a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs index 8b659e757aa2a..0c2286527dbc5 100644 --- a/datafusion/core/tests/physical_optimizer/pushdown_utils.rs +++ b/datafusion/core/tests/physical_optimizer/pushdown_utils.rs @@ -30,6 +30,7 @@ use datafusion_physical_expr_common::physical_expr::fmt_sql; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_plan::filter::batch_filter; use datafusion_physical_plan::filter_pushdown::{FilterPushdownPhase, PushedDown}; +use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use datafusion_physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, displayable, filter::FilterExec, @@ -112,7 +113,7 @@ pub struct TestSource { impl TestSource { pub fn new(schema: SchemaRef, support: bool, batches: Vec) -> Self { - let table_schema = datafusion_datasource::TableSchema::new(schema, vec![]); + let table_schema = datafusion_datasource::TableSchema::from(schema); Self { support, metrics: ExecutionPlanMetricsSet::new(), @@ -238,21 +239,17 @@ impl FileSource for TestSource { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit predicate (filter) expression if present - if let Some(predicate) = &self.predicate { - f(predicate.as_ref())?; - } - - // Visit projection expressions if present - if let Some(projection) = &self.projection { - for proj_expr in projection { - f(proj_expr.expr.as_ref())?; - } - } - - Ok(TreeNodeRecursion::Continue) + datafusion_physical_plan::apply_expression_roots( + self.predicate.iter().chain( + self.projection + .iter() + .flatten() + .map(|proj_expr| &proj_expr.expr), + ), + f, + ) } } @@ -493,9 +490,10 @@ impl ExecutionPlan for TestNode { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert!(children.len() == 1); Ok(Arc::new(TestNode::new( @@ -505,6 +503,16 @@ impl ExecutionPlan for TestNode { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -572,10 +580,8 @@ impl ExecutionPlan for TestNode { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit the predicate expression - f(self.predicate.as_ref())?; - Ok(TreeNodeRecursion::Continue) + datafusion_physical_plan::apply_expression_roots([&self.predicate], f) } } diff --git a/datafusion/core/tests/physical_optimizer/sanity_checker.rs b/datafusion/core/tests/physical_optimizer/sanity_checker.rs index 217570846d56e..184125dcbe180 100644 --- a/datafusion/core/tests/physical_optimizer/sanity_checker.rs +++ b/datafusion/core/tests/physical_optimizer/sanity_checker.rs @@ -19,8 +19,9 @@ use insta::assert_snapshot; use std::sync::Arc; use crate::physical_optimizer::test_utils::{ - bounded_window_exec, global_limit_exec, local_limit_exec, memory_exec, - projection_exec, repartition_exec, sort_exec, sort_expr, sort_expr_options, + bounded_window_exec, bounded_window_exec_with_can_repartition, global_limit_exec, + hash_join_exec, local_limit_exec, memory_exec, projection_exec, repartition_exec, + sort_exec, sort_exec_with_preserve_partitioning, sort_expr, sort_expr_options, sort_merge_join_exec, sort_preserving_merge_exec, union_exec, }; @@ -29,12 +30,13 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use datafusion::datasource::stream::{FileStreamProvider, StreamConfig, StreamTable}; use datafusion::prelude::{CsvReadOptions, SessionContext}; use datafusion_common::config::ConfigOptions; -use datafusion_common::{JoinType, Result, ScalarValue}; -use datafusion_physical_expr::Partitioning; +use datafusion_common::{JoinType, NullEquality, Result, ScalarValue}; use datafusion_physical_expr::expressions::{Literal, col}; +use datafusion_physical_expr::{Partitioning, RangePartitioning, SplitPoint}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_optimizer::PhysicalOptimizerRule; use datafusion_physical_optimizer::sanity_checker::SanityCheckPlan; +use datafusion_physical_plan::joins::{StreamJoinPartitionMode, SymmetricHashJoinExec}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::{ExecutionPlan, displayable}; @@ -400,6 +402,146 @@ fn assert_sanity_check(plan: &Arc, is_sane: bool) { ); } +fn range_partitioned_exec( + schema: &SchemaRef, + key: &str, + split_points: impl IntoIterator, +) -> Result> { + let split_points = split_points + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [sort_expr(key, schema)].into(), + split_points, + )?); + RepartitionExec::try_new(memory_exec(schema), partitioning) + .map(|exec| Arc::new(exec) as Arc) +} + +#[test] +fn test_partitioned_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + join_on.clone(), + None, + &JoinType::Inner, + )?; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + join_on, + None, + &JoinType::Inner, + )?; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + +#[test] +fn test_partitioned_right_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + join_on.clone(), + None, + &JoinType::Right, + )?; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = hash_join_exec( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + join_on, + None, + &JoinType::Right, + )?; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + +#[test] +fn test_sort_merge_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + let ordering: LexOrdering = [sort_expr("a", &schema)].into(); + + let compatible_join = sort_merge_join_exec( + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&compatible_join, true); + + let incompatible_join = sort_merge_join_exec( + sort_exec_with_preserve_partitioning( + ordering.clone(), + range_partitioned_exec(&schema, "a", [10])?, + ), + sort_exec_with_preserve_partitioning( + ordering, + range_partitioned_exec(&schema, "a", [20])?, + ), + &join_on, + &JoinType::Inner, + ); + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + +#[test] +fn test_symmetric_hash_join_requires_co_partitioned_children() -> Result<()> { + let schema = create_test_schema2(); + let join_on = vec![(col("a", &schema)?, col("a", &schema)?)]; + + let compatible_join = Arc::new(SymmetricHashJoinExec::try_new( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [10])?, + join_on.clone(), + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?) as Arc; + assert_sanity_check(&compatible_join, true); + + let incompatible_join = Arc::new(SymmetricHashJoinExec::try_new( + range_partitioned_exec(&schema, "a", [10])?, + range_partitioned_exec(&schema, "a", [20])?, + join_on, + None, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + None, + None, + StreamJoinPartitionMode::Partitioned, + )?) as Arc; + assert_sanity_check(&incompatible_join, false); + + Ok(()) +} + #[tokio::test] /// Tests that plan is valid when the sort requirements are satisfied. async fn test_bounded_window_agg_sort_requirement() -> Result<()> { @@ -458,6 +600,76 @@ async fn test_bounded_window_agg_no_sort_requirement() -> Result<()> { Ok(()) } +#[tokio::test] +/// Tests that a window over a compatible range-partitioned input satisfies +/// the window's key distribution requirement without a hash repartition. +async fn test_bounded_window_agg_range_partitioning() -> Result<()> { + let schema = create_test_schema2(); + let source = range_partitioned_exec(&schema, "a", [10, 20, 30])?; + let ordering: LexOrdering = [sort_expr_options( + "a", + &schema, + SortOptions { + descending: false, + nulls_first: false, + }, + )] + .into(); + let partition_by = vec![col("a", &schema)?]; + let sort = sort_exec_with_preserve_partitioning(ordering, source); + let bw = + bounded_window_exec_with_can_repartition("a", vec![], &partition_by, sort, true); + let plan_str = displayable(bw.as_ref()).indent(true).to_string(); + let actual = plan_str.trim(); + assert_snapshot!( + actual, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[a@0 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + "# + ); + assert_sanity_check(&bw, true); + Ok(()) +} + +#[tokio::test] +/// Tests that a window over an incompatible range-partitioned input fails +/// the window's key distribution requirement. +async fn test_bounded_window_agg_incompatible_range_partitioning() -> Result<()> { + let schema = create_test_schema2(); + let source = range_partitioned_exec(&schema, "a", [10, 20, 30])?; + let ordering: LexOrdering = [sort_expr_options( + "b", + &schema, + SortOptions { + descending: false, + nulls_first: false, + }, + )] + .into(); + let partition_by = vec![col("b", &schema)?]; + let sort = sort_exec_with_preserve_partitioning(ordering, source); + let bw = + bounded_window_exec_with_can_repartition("b", vec![], &partition_by, sort, true); + let plan_str = displayable(bw.as_ref()).indent(true).to_string(); + let actual = plan_str.trim(); + assert_snapshot!( + actual, + @r#" + BoundedWindowAggExec: wdw=[count: Field { "count": Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[true] + RepartitionExec: partitioning=Range([a@0 ASC], [(10), (20), (30)], 4), input_partitions=1 + DataSourceExec: partitions=1, partition_sizes=[0] + "# + ); + // Range([a]) does not colocate `b` values, so the window's key + // distribution requirement is not satisfied. + assert_sanity_check(&bw, false); + Ok(()) +} + #[tokio::test] /// A valid when a single partition requirement /// is satisfied. diff --git a/datafusion/core/tests/physical_optimizer/test_utils.rs b/datafusion/core/tests/physical_optimizer/test_utils.rs index 6814ab2358ffc..0835497f3451e 100644 --- a/datafusion/core/tests/physical_optimizer/test_utils.rs +++ b/datafusion/core/tests/physical_optimizer/test_utils.rs @@ -27,7 +27,7 @@ use arrow::record_batch::RecordBatch; use datafusion::datasource::listing::PartitionedFile; use datafusion::datasource::memory::MemorySourceConfig; use datafusion::datasource::physical_plan::ParquetSource; -use datafusion::datasource::source::DataSourceExec; +use datafusion::datasource::source::{DataSource, DataSourceExec}; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::tree_node::{ @@ -42,9 +42,10 @@ use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::{WindowFrame, WindowFunctionDefinition}; use datafusion_functions_aggregate::count::count_udaf; -use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; use datafusion_physical_expr::expressions::{self, col}; +use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::{Distribution, EquivalenceProperties}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{ LexOrdering, OrderingRequirements, PhysicalSortExpr, @@ -69,8 +70,9 @@ use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::union::UnionExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_window_expr}; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, InputOrderMode, Partitioning, - PlanProperties, SortOrderPushdownResult, displayable, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + InputDistributionRequirements, InputOrderMode, Partitioning, PlanProperties, + ReplaceChildrenOptions, SortOrderPushdownResult, StatisticsArgs, displayable, }; /// Create a non sorted parquet exec @@ -232,6 +234,16 @@ pub fn memory_exec(schema: &SchemaRef) -> Arc { MemorySourceConfig::try_new_exec(&[vec![]], Arc::clone(schema), None).unwrap() } +pub fn inexact_memory_exec( + partitions: &[Vec], + schema: SchemaRef, +) -> Result> { + let source = InexactMemorySource { + inner: MemorySourceConfig::try_new(partitions, schema, None)?, + }; + Ok(Arc::new(DataSourceExec::new(Arc::new(source)))) +} + pub fn hash_join_exec( left: Arc, right: Arc, @@ -265,6 +277,22 @@ pub fn bounded_window_exec_with_partition( sort_exprs: impl IntoIterator, partition_by: &[Arc], input: Arc, +) -> Arc { + bounded_window_exec_with_can_repartition( + col_name, + sort_exprs, + partition_by, + input, + false, + ) +} + +pub fn bounded_window_exec_with_can_repartition( + col_name: &str, + sort_exprs: impl IntoIterator, + partition_by: &[Arc], + input: Arc, + can_repartition: bool, ) -> Arc { let sort_exprs = sort_exprs.into_iter().collect::>(); let schema = input.schema(); @@ -287,7 +315,7 @@ pub fn bounded_window_exec_with_partition( vec![window_expr], Arc::clone(&input), InputOrderMode::Sorted, - false, + can_repartition, ) .unwrap(), ) @@ -375,6 +403,18 @@ pub fn sort_exec_with_preserve_partitioning( Arc::new(SortExec::new(ordering, input).with_preserve_partitioning(true)) } +pub fn sort_exec_with_fetch_and_preserve_partitioning( + ordering: LexOrdering, + fetch: Option, + input: Arc, +) -> Arc { + Arc::new( + SortExec::new(ordering, input) + .with_fetch(fetch) + .with_preserve_partitioning(true), + ) +} + pub fn sort_exec_with_fetch( ordering: LexOrdering, fetch: Option, @@ -398,6 +438,7 @@ pub fn projection_exec( #[derive(Debug)] pub struct RequirementsTestExec { required_input_ordering: Option, + required_input_distribution: Distribution, maintains_input_order: bool, input: Arc, } @@ -406,6 +447,7 @@ impl RequirementsTestExec { pub fn new(input: Arc) -> Self { Self { required_input_ordering: None, + required_input_distribution: Distribution::UnspecifiedDistribution, maintains_input_order: true, input, } @@ -420,6 +462,15 @@ impl RequirementsTestExec { self } + /// sets the required input distribution + pub fn with_required_input_distribution( + mut self, + required_input_distribution: Distribution, + ) -> Self { + self.required_input_distribution = required_input_distribution; + self + } + /// set the maintains_input_order flag pub fn with_maintains_input_order(mut self, maintains_input_order: bool) -> Self { self.maintains_input_order = maintains_input_order; @@ -463,6 +514,10 @@ impl ExecutionPlan for RequirementsTestExec { ] } + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + InputDistributionRequirements::new(vec![self.required_input_distribution.clone()]) + } + fn maintains_input_order(&self) -> Vec { vec![self.maintains_input_order] } @@ -471,17 +526,29 @@ impl ExecutionPlan for RequirementsTestExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_eq!(children.len(), 1); Ok(RequirementsTestExec::new(Arc::clone(&children[0])) .with_required_input_ordering(self.required_input_ordering.clone()) + .with_required_input_distribution(self.required_input_distribution.clone()) .with_maintains_input_order(self.maintains_input_order) .into_arc()) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -492,16 +559,9 @@ impl ExecutionPlan for RequirementsTestExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in required_input_ordering if present - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = &self.required_input_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } @@ -909,6 +969,18 @@ impl TestScan { self.supports_fetch = supports; self } + + /// Set the number of output partitions reported by this scan. + pub fn with_partition_count(mut self, partition_count: usize) -> Self { + let eq_properties = self.plan_properties.equivalence_properties().clone(); + self.plan_properties = Arc::new(PlanProperties::new( + eq_properties, + Partitioning::UnknownPartitioning(partition_count), + EmissionType::Incremental, + Boundedness::Bounded, + )); + self + } } impl DisplayAs for TestScan { @@ -964,9 +1036,10 @@ impl ExecutionPlan for TestScan { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { if children.is_empty() { Ok(self) @@ -975,6 +1048,16 @@ impl ExecutionPlan for TestScan { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -983,7 +1066,11 @@ impl ExecutionPlan for TestScan { internal_err!("TestScan is for testing optimizer only, not for execution") } - fn partition_statistics(&self, _partition: Option) -> Result> { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { Ok(Arc::new(Statistics::new_unknown(&self.schema))) } @@ -1038,24 +1125,9 @@ impl ExecutionPlan for TestScan { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in output_ordering - let mut tnr = TreeNodeRecursion::Continue; - for ordering in &self.output_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - - // Visit expressions in requested_ordering if present - if let Some(ordering) = &self.requested_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } @@ -1066,3 +1138,67 @@ pub fn test_scan_with_ordering( ) -> Arc { Arc::new(TestScan::with_ordering(schema, ordering)) } + +#[derive(Debug, Clone)] +struct InexactMemorySource { + inner: MemorySourceConfig, +} + +impl DataSource for InexactMemorySource { + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.inner.apply_expressions(f) + } + + fn open( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.open(partition, context) + } + + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + self.inner.fmt_as(t, f) + } + + fn output_partitioning(&self) -> Partitioning { + self.inner.output_partitioning() + } + + fn eq_properties(&self) -> EquivalenceProperties { + self.inner.eq_properties() + } + + fn partition_statistics(&self, partition: Option) -> Result> { + self.inner.partition_statistics(partition) + } + + fn with_fetch(&self, limit: Option) -> Option> { + let mut new_source = self.clone(); + new_source.inner = new_source.inner.with_limit(limit); + Some(Arc::new(new_source)) + } + + fn fetch(&self) -> Option { + self.inner.fetch() + } + + fn try_swapping_with_projection( + &self, + _projection: &ProjectionExprs, + ) -> Result>> { + Ok(None) + } + + fn try_pushdown_sort( + &self, + _order: &[PhysicalSortExpr], + ) -> Result>> { + Ok(SortOrderPushdownResult::Inexact { + inner: Arc::new(self.clone()), + }) + } +} diff --git a/datafusion/core/tests/physical_optimizer/window_topn.rs b/datafusion/core/tests/physical_optimizer/window_topn.rs index e3f73a85353cc..be78b77b32a2d 100644 --- a/datafusion/core/tests/physical_optimizer/window_topn.rs +++ b/datafusion/core/tests/physical_optimizer/window_topn.rs @@ -25,6 +25,7 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::Operator; use datafusion_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits}; +use datafusion_functions_window::rank::{dense_rank_udwf, rank_udwf}; use datafusion_functions_window::row_number::row_number_udwf; use datafusion_physical_expr::expressions::{BinaryExpr, Column, col, lit}; use datafusion_physical_expr::window::StandardWindowExpr; @@ -63,7 +64,9 @@ fn optimize_disabled(plan: Arc) -> Result = Arc::new(PlaceholderRowExec::new(Arc::clone(&s))); - // Sort by pk ASC, val ASC - let ordering = LexOrdering::new(vec![ - PhysicalSortExpr::new_default(col("pk", &s)?).asc(), - PhysicalSortExpr::new_default(col("val", &s)?).asc(), - ]) - .unwrap(); - - let sort: Arc = - Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true)); - // ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) let partition_by = vec![col("pk", &s)?]; let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()]; @@ -104,7 +97,7 @@ fn build_window_topn_plan( let window: Arc = Arc::new(BoundedWindowAggExec::try_new( vec![window_expr], - sort, + input, InputOrderMode::Sorted, true, )?); @@ -226,7 +219,7 @@ fn basic_row_number_rn_lteq_3() -> Result<()> { let optimized = optimize(plan)?; assert_snapshot!(plan_str(optimized.as_ref()), @r#" BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - PartitionedTopKExec: fetch=3, partition=[pk@0], order=[val@1 ASC] + PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC] PlaceholderRowExec "#); Ok(()) @@ -238,7 +231,7 @@ fn rn_lt_3_becomes_fetch_2() -> Result<()> { let optimized = optimize(plan)?; assert_snapshot!(plan_str(optimized.as_ref()), @r#" BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - PartitionedTopKExec: fetch=2, partition=[pk@0], order=[val@1 ASC] + PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@0], order=[val@1 ASC] PlaceholderRowExec "#); Ok(()) @@ -251,16 +244,6 @@ fn flipped_3_gteq_rn() -> Result<()> { let input: Arc = Arc::new(PlaceholderRowExec::new(Arc::clone(&s))); - let ordering = LexOrdering::new(vec![ - PhysicalSortExpr::new_default(col("pk", &s)?).asc(), - PhysicalSortExpr::new_default(col("val", &s)?).asc(), - ]) - .unwrap(); - - let sort: Arc = Arc::new( - SortExec::new(ordering.clone(), input).with_preserve_partitioning(true), - ); - let partition_by = vec![col("pk", &s)?]; let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()]; @@ -283,7 +266,7 @@ fn flipped_3_gteq_rn() -> Result<()> { let window: Arc = Arc::new(BoundedWindowAggExec::try_new( vec![window_expr], - sort, + input, InputOrderMode::Sorted, true, )?); @@ -300,7 +283,7 @@ fn flipped_3_gteq_rn() -> Result<()> { let optimized = optimize(plan)?; assert_snapshot!(plan_str(optimized.as_ref()), @r#" BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - PartitionedTopKExec: fetch=3, partition=[pk@0], order=[val@1 ASC] + PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC] PlaceholderRowExec "#); Ok(()) @@ -352,15 +335,6 @@ fn with_projection_between() -> Result<()> { let s = schema(); let input: Arc = Arc::new(PlaceholderRowExec::new(Arc::clone(&s))); - let ordering = LexOrdering::new(vec![ - PhysicalSortExpr::new_default(col("pk", &s)?).asc(), - PhysicalSortExpr::new_default(col("val", &s)?).asc(), - ]) - .unwrap(); - - let sort: Arc = - Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true)); - let partition_by = vec![col("pk", &s)?]; let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()]; @@ -383,7 +357,7 @@ fn with_projection_between() -> Result<()> { let window: Arc = Arc::new(BoundedWindowAggExec::try_new( vec![window_expr], - sort, + input, InputOrderMode::Sorted, true, )?); @@ -418,8 +392,184 @@ fn with_projection_between() -> Result<()> { assert_snapshot!(plan_str(optimized.as_ref()), @r#" ProjectionExec: expr=[pk@0 as pk, val@1 as val, row_number@2 as row_number] BoundedWindowAggExec: wdw=[row_number: Field { "row_number": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] - PartitionedTopKExec: fetch=3, partition=[pk@0], order=[val@1 ASC] + PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@0], order=[val@1 ASC] PlaceholderRowExec "#); Ok(()) } + +// ---------------------------------------------------------------------- +// RANK rule tests +// ---------------------------------------------------------------------- + +/// Build: FilterExec(rk op limit) → BoundedWindowAggExec( PBY pk OBY val) +/// +/// Matches the pre-`EnsureRequirements` plan shape (no `SortExec` under the window). +/// +/// `udwf_factory` selects the window UDWF (rank, dense_rank, ...) and +/// `udwf_name` is the column name produced by that UDWF (matters because +/// the rule resolves the filter column by index, but the snapshot prints +/// the name). +fn build_ranking_topn_plan( + udwf_factory: fn() -> Arc, + udwf_name: &str, + limit_value: i64, + op: Operator, +) -> Result> { + let s = schema(); + let input: Arc = Arc::new(PlaceholderRowExec::new(Arc::clone(&s))); + + let partition_by = vec![col("pk", &s)?]; + let order_by = vec![PhysicalSortExpr::new_default(col("val", &s)?).asc()]; + + let window_expr = Arc::new(StandardWindowExpr::new( + create_udwf_window_expr(&udwf_factory(), &[], &s, udwf_name.to_string(), false)?, + &partition_by, + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + )); + + let window: Arc = Arc::new(BoundedWindowAggExec::try_new( + vec![window_expr], + input, + InputOrderMode::Sorted, + true, + )?); + + let rk_col = Arc::new(Column::new(udwf_name, 2)); + let limit_lit = lit(ScalarValue::UInt64(Some(limit_value as u64))); + // Place column on whichever side matches the operator's expectation. + let predicate: Arc = match op { + Operator::LtEq | Operator::Lt => Arc::new(BinaryExpr::new(rk_col, op, limit_lit)), + Operator::GtEq | Operator::Gt => Arc::new(BinaryExpr::new(limit_lit, op, rk_col)), + _ => unreachable!("only =/> are supported by the rule"), + }; + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, window)?); + + Ok(filter) +} + +/// Build a RANK plan with NO ORDER BY: every row ties at rank 1 — degenerate. +fn build_rank_no_order_by_plan(limit_value: i64) -> Result> { + let s = schema(); + let input: Arc = Arc::new(PlaceholderRowExec::new(Arc::clone(&s))); + + let ordering = + LexOrdering::new(vec![PhysicalSortExpr::new_default(col("pk", &s)?).asc()]) + .unwrap(); + + let sort: Arc = + Arc::new(SortExec::new(ordering.clone(), input).with_preserve_partitioning(true)); + + let partition_by = vec![col("pk", &s)?]; + + let window_expr = Arc::new(StandardWindowExpr::new( + create_udwf_window_expr(&rank_udwf(), &[], &s, "rank".to_string(), false)?, + &partition_by, + &[], // empty ORDER BY + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + )); + + let window: Arc = Arc::new(BoundedWindowAggExec::try_new( + vec![window_expr], + sort, + InputOrderMode::Sorted, + true, + )?); + + let rk_col = Arc::new(Column::new("rank", 2)); + let limit_lit = lit(ScalarValue::UInt64(Some(limit_value as u64))); + let predicate = Arc::new(BinaryExpr::new(rk_col, Operator::LtEq, limit_lit)); + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, window)?); + + Ok(filter) +} + +#[test] +fn basic_rank_rk_lteq_3() -> Result<()> { + let plan = build_ranking_topn_plan(rank_udwf, "rank", 3, Operator::LtEq)?; + let optimized = optimize(plan)?; + assert_snapshot!(plan_str(optimized.as_ref()), @r#" + BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC] + PlaceholderRowExec + "#); + Ok(()) +} + +#[test] +fn rank_rk_lt_4_becomes_fetch_3() -> Result<()> { + let plan = build_ranking_topn_plan(rank_udwf, "rank", 4, Operator::Lt)?; + let optimized = optimize(plan)?; + assert_snapshot!(plan_str(optimized.as_ref()), @r#" + BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC] + PlaceholderRowExec + "#); + Ok(()) +} + +#[test] +fn rank_flipped_3_gteq_rk() -> Result<()> { + let plan = build_ranking_topn_plan(rank_udwf, "rank", 3, Operator::GtEq)?; + let optimized = optimize(plan)?; + assert_snapshot!(plan_str(optimized.as_ref()), @r#" + BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC] + PlaceholderRowExec + "#); + Ok(()) +} + +#[test] +fn rank_flipped_4_gt_rk_becomes_fetch_3() -> Result<()> { + let plan = build_ranking_topn_plan(rank_udwf, "rank", 4, Operator::Gt)?; + let optimized = optimize(plan)?; + assert_snapshot!(plan_str(optimized.as_ref()), @r#" + BoundedWindowAggExec: wdw=[rank: Field { "rank": UInt64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] + PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@0], order=[val@1 ASC] + PlaceholderRowExec + "#); + Ok(()) +} + +#[test] +fn rank_no_order_by_no_change() -> Result<()> { + // Without ORDER BY, every row ties at rank 1 — the optimization is + // degenerate (entire input would be retained, ties storage unbounded). + // The rule must skip. + let plan = build_rank_no_order_by_plan(3)?; + let before = plan_str(plan.as_ref()); + let optimized = optimize(plan)?; + let after = plan_str(optimized.as_ref()); + assert_eq!( + before, after, + "RANK with empty ORDER BY must not be rewritten" + ); + Ok(()) +} + +#[test] +fn dense_rank_no_change() -> Result<()> { + // DENSE_RANK is not yet supported by the rule. The plan must pass + // through unchanged. + let plan = build_ranking_topn_plan(dense_rank_udwf, "dense_rank", 3, Operator::LtEq)?; + let before = plan_str(plan.as_ref()); + let optimized = optimize(plan)?; + let after = plan_str(optimized.as_ref()); + assert_eq!( + before, after, + "DENSE_RANK is unsupported and must not be rewritten" + ); + Ok(()) +} diff --git a/datafusion/core/tests/sql/aggregates/dict_nulls.rs b/datafusion/core/tests/sql/aggregates/dict_nulls.rs index 8733b9e87b57a..c6c3f02829c43 100644 --- a/datafusion/core/tests/sql/aggregates/dict_nulls.rs +++ b/datafusion/core/tests/sql/aggregates/dict_nulls.rs @@ -292,7 +292,7 @@ async fn test_first_last_value_group_by_dict_nulls() -> Result<()> { /// Test MAX with dictionary columns containing null keys and values as specified in the SQL query #[tokio::test] async fn test_max_with_fuzz_table_dict_nulls() -> Result<()> { - let (ctx_single, ctx_multi) = setup_fuzz_test_contexts().await?; + let (ctx_single, ctx_multi) = setup_fuzz_test_contexts()?; // Execute the SQL query with MAX aggregations let sql = "SELECT @@ -333,7 +333,7 @@ async fn test_max_with_fuzz_table_dict_nulls() -> Result<()> { /// Test MIN with fuzz table containing dictionary columns with null keys and values and timestamp data (single and multiple partitions) #[tokio::test] async fn test_min_timestamp_with_fuzz_table_dict_nulls() -> Result<()> { - let (ctx_single, ctx_multi) = setup_fuzz_timestamp_test_contexts().await?; + let (ctx_single, ctx_multi) = setup_fuzz_timestamp_test_contexts()?; // Execute the SQL query with MIN aggregation on timestamp let sql = "SELECT @@ -373,7 +373,7 @@ async fn test_min_timestamp_with_fuzz_table_dict_nulls() -> Result<()> { /// Test COUNT and COUNT DISTINCT with fuzz table containing dictionary columns with null keys and values (single and multiple partitions) #[tokio::test] async fn test_count_distinct_with_fuzz_table_dict_nulls() -> Result<()> { - let (ctx_single, ctx_multi) = setup_fuzz_count_test_contexts().await?; + let (ctx_single, ctx_multi) = setup_fuzz_count_test_contexts()?; // Execute the SQL query with COUNT and COUNT DISTINCT aggregations let sql = "SELECT @@ -414,7 +414,7 @@ async fn test_count_distinct_with_fuzz_table_dict_nulls() -> Result<()> { /// Test MEDIAN and MEDIAN DISTINCT with fuzz table containing various numeric types and dictionary columns with null keys and values (single and multiple partitions) #[tokio::test] async fn test_median_distinct_with_fuzz_table_dict_nulls() -> Result<()> { - let (ctx_single, ctx_multi) = setup_fuzz_median_test_contexts().await?; + let (ctx_single, ctx_multi) = setup_fuzz_median_test_contexts()?; // Execute the SQL query with MEDIAN and MEDIAN DISTINCT aggregations let sql = "SELECT diff --git a/datafusion/core/tests/sql/aggregates/mod.rs b/datafusion/core/tests/sql/aggregates/mod.rs index ede40d5c4ceca..b209e91cc81e7 100644 --- a/datafusion/core/tests/sql/aggregates/mod.rs +++ b/datafusion/core/tests/sql/aggregates/mod.rs @@ -259,20 +259,20 @@ impl TestData { } /// Sets up test contexts for TestData with both single and multiple partitions -pub async fn setup_test_contexts( +pub fn setup_test_contexts( test_data: &TestData, ) -> Result<(SessionContext, SessionContext)> { // Single partition context - let ctx_single = create_context_with_partitions(test_data, 1).await?; + let ctx_single = create_context_with_partitions(test_data, 1)?; // Multiple partition context - let ctx_multi = create_context_with_partitions(test_data, 3).await?; + let ctx_multi = create_context_with_partitions(test_data, 3)?; Ok((ctx_single, ctx_multi)) } /// Creates a session context with the specified number of partitions and registers test data -pub async fn create_context_with_partitions( +pub fn create_context_with_partitions( test_data: &TestData, num_partitions: usize, ) -> Result { @@ -348,7 +348,7 @@ pub async fn run_snapshot_test( test_data: &TestData, sql: &str, ) -> Result> { - let (ctx_single, ctx_multi) = setup_test_contexts(test_data).await?; + let (ctx_single, ctx_multi) = setup_test_contexts(test_data)?; let results = test_query_consistency(&ctx_single, &ctx_multi, sql).await?; Ok(results) } @@ -430,20 +430,20 @@ impl FuzzTestData { } /// Sets up test contexts for fuzz table with both single and multiple partitions -pub async fn setup_fuzz_test_contexts() -> Result<(SessionContext, SessionContext)> { +pub fn setup_fuzz_test_contexts() -> Result<(SessionContext, SessionContext)> { let test_data = FuzzTestData::new(); // Single partition context - let ctx_single = create_fuzz_context_with_partitions(&test_data, 1).await?; + let ctx_single = create_fuzz_context_with_partitions(&test_data, 1)?; // Multiple partition context - let ctx_multi = create_fuzz_context_with_partitions(&test_data, 3).await?; + let ctx_multi = create_fuzz_context_with_partitions(&test_data, 3)?; Ok((ctx_single, ctx_multi)) } /// Creates a session context with fuzz table partitioned into specified number of partitions -pub async fn create_fuzz_context_with_partitions( +pub fn create_fuzz_context_with_partitions( test_data: &FuzzTestData, num_partitions: usize, ) -> Result { @@ -604,21 +604,20 @@ impl FuzzCountTestData { } /// Sets up test contexts for fuzz table with duration/binary columns and both single and multiple partitions -pub async fn setup_fuzz_count_test_contexts() -> Result<(SessionContext, SessionContext)> -{ +pub fn setup_fuzz_count_test_contexts() -> Result<(SessionContext, SessionContext)> { let test_data = FuzzCountTestData::new(); // Single partition context - let ctx_single = create_fuzz_count_context_with_partitions(&test_data, 1).await?; + let ctx_single = create_fuzz_count_context_with_partitions(&test_data, 1)?; // Multiple partition context - let ctx_multi = create_fuzz_count_context_with_partitions(&test_data, 3).await?; + let ctx_multi = create_fuzz_count_context_with_partitions(&test_data, 3)?; Ok((ctx_single, ctx_multi)) } /// Creates a session context with fuzz count table partitioned into specified number of partitions -pub async fn create_fuzz_count_context_with_partitions( +pub fn create_fuzz_count_context_with_partitions( test_data: &FuzzCountTestData, num_partitions: usize, ) -> Result { @@ -808,21 +807,20 @@ impl FuzzMedianTestData { } /// Sets up test contexts for fuzz table with numeric types for median testing and both single and multiple partitions -pub async fn setup_fuzz_median_test_contexts() -> Result<(SessionContext, SessionContext)> -{ +pub fn setup_fuzz_median_test_contexts() -> Result<(SessionContext, SessionContext)> { let test_data = FuzzMedianTestData::new(); // Single partition context - let ctx_single = create_fuzz_median_context_with_partitions(&test_data, 1).await?; + let ctx_single = create_fuzz_median_context_with_partitions(&test_data, 1)?; // Multiple partition context - let ctx_multi = create_fuzz_median_context_with_partitions(&test_data, 3).await?; + let ctx_multi = create_fuzz_median_context_with_partitions(&test_data, 3)?; Ok((ctx_single, ctx_multi)) } /// Creates a session context with fuzz median table partitioned into specified number of partitions -pub async fn create_fuzz_median_context_with_partitions( +pub fn create_fuzz_median_context_with_partitions( test_data: &FuzzMedianTestData, num_partitions: usize, ) -> Result { @@ -959,21 +957,20 @@ impl FuzzTimestampTestData { } /// Sets up test contexts for fuzz table with timestamps and both single and multiple partitions -pub async fn setup_fuzz_timestamp_test_contexts() --> Result<(SessionContext, SessionContext)> { +pub fn setup_fuzz_timestamp_test_contexts() -> Result<(SessionContext, SessionContext)> { let test_data = FuzzTimestampTestData::new(); // Single partition context - let ctx_single = create_fuzz_timestamp_context_with_partitions(&test_data, 1).await?; + let ctx_single = create_fuzz_timestamp_context_with_partitions(&test_data, 1)?; // Multiple partition context - let ctx_multi = create_fuzz_timestamp_context_with_partitions(&test_data, 3).await?; + let ctx_multi = create_fuzz_timestamp_context_with_partitions(&test_data, 3)?; Ok((ctx_single, ctx_multi)) } /// Creates a session context with fuzz timestamp table partitioned into specified number of partitions -pub async fn create_fuzz_timestamp_context_with_partitions( +pub fn create_fuzz_timestamp_context_with_partitions( test_data: &FuzzTimestampTestData, num_partitions: usize, ) -> Result { diff --git a/datafusion/core/tests/sql/explain_analyze.rs b/datafusion/core/tests/sql/explain_analyze.rs index b093563d9adda..4c8b8f9c01122 100644 --- a/datafusion/core/tests/sql/explain_analyze.rs +++ b/datafusion/core/tests/sql/explain_analyze.rs @@ -774,8 +774,8 @@ async fn test_physical_plan_display_indent() { actual, @r" SortPreservingMergeExec: [the_min@2 DESC], fetch=10 - SortExec: TopK(fetch=10), expr=[the_min@2 DESC], preserve_partitioning=[true] - ProjectionExec: expr=[c1@0 as c1, max(aggregate_test_100.c12)@1 as max(aggregate_test_100.c12), min(aggregate_test_100.c12)@2 as the_min] + ProjectionExec: expr=[c1@0 as c1, max(aggregate_test_100.c12)@1 as max(aggregate_test_100.c12), min(aggregate_test_100.c12)@2 as the_min] + SortExec: TopK(fetch=10), expr=[min(aggregate_test_100.c12)@2 DESC], preserve_partitioning=[true] AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[max(aggregate_test_100.c12), min(aggregate_test_100.c12)] RepartitionExec: partitioning=Hash([c1@0], 9000), input_partitions=9000 AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[max(aggregate_test_100.c12), min(aggregate_test_100.c12)] @@ -827,7 +827,7 @@ async fn test_physical_plan_display_indent_multi_children() { } #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] +#[cfg_attr(coverage, ignore)] async fn csv_explain_analyze() { // This test uses the execute function to run an actual plan under EXPLAIN ANALYZE let ctx = SessionContext::new(); @@ -849,7 +849,7 @@ async fn csv_explain_analyze() { } #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] +#[cfg_attr(coverage, ignore)] async fn csv_explain_analyze_order_by() { let ctx = SessionContext::new(); register_aggregate_csv_by_sql(&ctx).await; @@ -866,7 +866,7 @@ async fn csv_explain_analyze_order_by() { } #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] +#[cfg_attr(coverage, ignore)] async fn parquet_explain_analyze() { let ctx = SessionContext::new(); register_alltypes_parquet(&ctx).await; @@ -913,7 +913,7 @@ async fn parquet_explain_analyze() { // (e.g. nested/recursive expansion causing full schema to be scanned). // Keeping this test ensures we don't regress that behavior. #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] +#[cfg_attr(coverage, ignore)] async fn parquet_recursive_projection_pushdown() -> Result<()> { use parquet::arrow::arrow_writer::ArrowWriter; use parquet::file::properties::WriterProperties; @@ -1014,7 +1014,7 @@ async fn parquet_recursive_projection_pushdown() -> Result<()> { SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] RecursiveQueryExec: name=number_series, is_distinct=false CoalescePartitionsExec - ProjectionExec: expr=[id@0 as id, 1 as level] + ProjectionExec: expr=[CAST(id@0 AS Int64) as id, CAST(1 AS Int64) as level] FilterExec: id@0 = 1 RepartitionExec: partitioning=RoundRobinBatch(NUM_CORES), input_partitions=1 DataSourceExec: file_groups={1 group: [[TMP_DIR/hierarchy.parquet]]}, projection=[id], file_type=parquet, predicate=id@0 = 1, pruning_predicate=id_null_count@2 != row_count@3 AND id_min@0 <= 1 AND 1 <= id_max@1, required_guarantees=[id in (1)] @@ -1030,7 +1030,7 @@ async fn parquet_recursive_projection_pushdown() -> Result<()> { } #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] +#[cfg_attr(coverage, ignore)] async fn parquet_explain_analyze_verbose() { let ctx = SessionContext::new(); register_alltypes_parquet(&ctx).await; @@ -1047,7 +1047,7 @@ async fn parquet_explain_analyze_verbose() { } #[tokio::test] -#[cfg_attr(tarpaulin, ignore)] +#[cfg_attr(coverage, ignore)] async fn csv_explain_analyze_verbose() { // This test uses the execute function to run an actual plan under EXPLAIN VERBOSE ANALYZE let ctx = SessionContext::new(); @@ -1267,3 +1267,115 @@ async fn explain_analyze_categories() { ); } } + +/// Returns a [`SessionContext`] configured with the PostgreSQL dialect so +/// that `EXPLAIN (option, ...)` utility-option syntax is accepted. +fn session_ctx_with_pg_dialect() -> SessionContext { + use std::str::FromStr; + let mut config = SessionConfig::new(); + let options = config.options_mut(); + options.sql_parser.dialect = + datafusion::config::Dialect::from_str("PostgreSQL").unwrap(); + SessionContext::new_with_config(config) +} + +async fn collect_explain(ctx: &SessionContext, sql: &str) -> String { + let dataframe = ctx.sql(sql).await.unwrap(); + let batches = dataframe.collect().await.unwrap(); + arrow::util::pretty::pretty_format_batches(&batches) + .unwrap() + .to_string() +} + +/// Verifies that the Postgres-style `EXPLAIN (METRICS '...')` form produces +/// the same category filtering as `SET datafusion.explain.analyze_categories`. +#[tokio::test] +async fn explain_analyze_paren_metrics_filtering() { + let ctx = session_ctx_with_pg_dialect(); + let sql = "EXPLAIN (ANALYZE, METRICS 'rows') \ + SELECT * FROM generate_series(10) as t1(v1) ORDER BY v1 DESC"; + let plan = collect_explain(&ctx, sql).await; + assert!( + plan.contains("output_rows"), + "rows category should include output_rows:\n{plan}" + ); + assert!( + !plan.contains("elapsed_compute"), + "rows-only METRICS should exclude elapsed_compute:\n{plan}" + ); + assert!( + !plan.contains("output_bytes"), + "rows-only METRICS should exclude output_bytes:\n{plan}" + ); +} + +/// Verifies that a statement-level METRICS overrides the session config. +#[tokio::test] +async fn explain_analyze_paren_metrics_overrides_session_config() { + let ctx = session_ctx_with_pg_dialect(); + // Session default: show only `rows` via config. + { + let state = ctx.state_ref(); + let mut state = state.write(); + state.config_mut().options_mut().explain.analyze_categories = + ExplainAnalyzeCategories::Only(vec![MetricCategory::Rows]); + } + // Statement overrides with 'bytes' — we should see output_bytes but not + // output_rows (except row-count metrics with the `output_bytes` substring + // are avoided because the metric names are distinct). + let sql = "EXPLAIN (ANALYZE, METRICS 'bytes') \ + SELECT * FROM generate_series(10) as t1(v1) ORDER BY v1 DESC"; + let plan = collect_explain(&ctx, sql).await; + assert!( + plan.contains("output_bytes"), + "statement-level METRICS='bytes' should show output_bytes:\n{plan}" + ); + assert!( + !plan.contains("output_rows"), + "statement-level METRICS='bytes' should hide output_rows:\n{plan}" + ); +} + +/// Verifies that `EXPLAIN (ANALYZE, LEVEL summary)` only shows summary metrics, +/// overriding the session default of `dev`. +#[tokio::test] +async fn explain_analyze_paren_level_overrides_session_config() { + let ctx = session_ctx_with_pg_dialect(); + // Session default: Dev + { + let state = ctx.state_ref(); + let mut state = state.write(); + state.config_mut().options_mut().explain.analyze_level = MetricType::Dev; + } + let sql = "EXPLAIN (ANALYZE, LEVEL summary) \ + SELECT * FROM generate_series(10) as t1(v1) ORDER BY v1 DESC"; + let plan = collect_explain(&ctx, sql).await; + // `spill_count` is Dev-only; `output_rows` is Summary. + assert!( + plan.contains("output_rows"), + "summary should still show output_rows:\n{plan}" + ); + assert!( + !plan.contains("spill_count"), + "summary should hide Dev-only spill_count:\n{plan}" + ); +} + +/// Verifies that `EXPLAIN (ANALYZE, BUFFERS)` returns a helpful error. +#[tokio::test] +async fn explain_paren_buffers_rejected() { + let ctx = session_ctx_with_pg_dialect(); + let err = ctx + .sql("EXPLAIN (ANALYZE, BUFFERS) SELECT 1") + .await + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("BUFFERS"), + "error should mention BUFFERS: {msg}" + ); + assert!( + msg.contains("not supported"), + "error should say not supported: {msg}" + ); +} diff --git a/datafusion/core/tests/sql/mod.rs b/datafusion/core/tests/sql/mod.rs index 9a1dc5502ee60..afed2f82d57a8 100644 --- a/datafusion/core/tests/sql/mod.rs +++ b/datafusion/core/tests/sql/mod.rs @@ -70,6 +70,8 @@ mod path_partition; mod runtime_config; pub mod select; mod sql_api; +mod union_comparison; +mod union_nullable; mod unparser; async fn register_aggregate_csv_by_sql(ctx: &SessionContext) { diff --git a/datafusion/core/tests/sql/path_partition.rs b/datafusion/core/tests/sql/path_partition.rs index 2eff1c262f855..82a15eb401fc4 100644 --- a/datafusion/core/tests/sql/path_partition.rs +++ b/datafusion/core/tests/sql/path_partition.rs @@ -38,6 +38,7 @@ use datafusion_common::ScalarValue; use datafusion_common::stats::Precision; use datafusion_common::test_util::batches_to_sort_string; use datafusion_execution::config::SessionConfig; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use async_trait::async_trait; use bytes::Bytes; @@ -461,8 +462,8 @@ async fn parquet_statistics() -> Result<()> { let schema = physical_plan.schema(); assert_eq!(schema.fields().len(), 4); - let stat_cols = physical_plan - .partition_statistics(None)? + let stat_cols = StatisticsContext::new() + .compute(physical_plan.as_ref(), &StatisticsArgs::new())? .column_statistics .clone(); assert_eq!(stat_cols.len(), 4); @@ -488,8 +489,8 @@ async fn parquet_statistics() -> Result<()> { let schema = physical_plan.schema(); assert_eq!(schema.fields().len(), 2); - let stat_cols = physical_plan - .partition_statistics(None)? + let stat_cols = StatisticsContext::new() + .compute(physical_plan.as_ref(), &StatisticsArgs::new())? .column_statistics .clone(); assert_eq!(stat_cols.len(), 2); @@ -609,8 +610,7 @@ async fn create_partitioned_alltypes_parquet_table( .iter() .map(|x| (x.0.to_owned(), x.1.clone())) .collect::>(), - ) - .with_session_config_options(&ctx.copied_config()); + ); let table_path = ListingTableUrl::parse(table_path).unwrap(); let store_path = diff --git a/datafusion/core/tests/sql/runtime_config.rs b/datafusion/core/tests/sql/runtime_config.rs index 2db1e1ce12f72..b0e4bccf30aba 100644 --- a/datafusion/core/tests/sql/runtime_config.rs +++ b/datafusion/core/tests/sql/runtime_config.rs @@ -23,12 +23,16 @@ use std::time::Duration; use datafusion::execution::context::SessionContext; use datafusion::execution::context::TaskContext; use datafusion::prelude::SessionConfig; -use datafusion_execution::cache::DefaultListFilesCache; -use datafusion_execution::cache::cache_manager::CacheManagerConfig; -use datafusion_execution::cache::file_statistics_cache::DefaultFileStatisticsCache; +use datafusion_execution::cache::cache_manager::{ + CacheManagerConfig, DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, + DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, +}; +use datafusion_execution::cache::default_cache::DefaultCache; use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_physical_plan::common::collect; +use crate::helper::plan_metrics::plan_spill_count; + #[tokio::test] async fn test_memory_limit_with_spill() { let ctx = SessionContext::new(); @@ -55,8 +59,7 @@ async fn test_memory_limit_with_spill() { let stream = plan.execute(0, task_ctx).unwrap(); let _results = collect(stream).await; - let metrics = plan.metrics().unwrap(); - let spill_count = metrics.spill_count().unwrap(); + let spill_count = plan_spill_count(plan.as_ref()); assert!(spill_count > 0, "Expected spills but none occurred"); } @@ -85,8 +88,7 @@ async fn test_no_spill_with_adequate_memory() { let stream = plan.execute(0, task_ctx).unwrap(); let _results = collect(stream).await; - let metrics = plan.metrics().unwrap(); - let spill_count = metrics.spill_count().unwrap(); + let spill_count = plan_spill_count(plan.as_ref()); assert_eq!(spill_count, 0, "Expected no spills but some occurred"); } @@ -113,7 +115,7 @@ async fn test_multiple_configs() { assert!(result.is_ok(), "Should not fail due to memory limit"); let state = ctx.state(); - let batch_size = state.config().options().execution.batch_size; + let batch_size = state.config().options().execution.batch_size.get(); assert_eq!(batch_size, 2048); } @@ -225,6 +227,34 @@ async fn test_max_temp_directory_size_enforcement() { ); } +#[tokio::test] +async fn test_max_spill_merge_fan_in_runtime_config() { + let ctx = SessionContext::new(); + + ctx.sql("SET datafusion.runtime.max_spill_merge_fan_in = '8'") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(ctx.runtime_env().disk_manager.max_spill_merge_fan_in(), 8); + + ctx.sql("RESET datafusion.runtime.max_spill_merge_fan_in") + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(ctx.runtime_env().disk_manager.max_spill_merge_fan_in(), 0); + + let error = ctx + .sql("SET datafusion.runtime.max_spill_merge_fan_in = '-1'") + .await + .unwrap_err() + .to_string(); + assert!(error.contains("Failed to parse non-negative integer")); +} + #[tokio::test] async fn test_test_metadata_cache_limit() { let ctx = SessionContext::new(); @@ -260,7 +290,8 @@ async fn test_test_metadata_cache_limit() { #[tokio::test] async fn test_list_files_cache_limit() { - let list_files_cache = Arc::new(DefaultListFilesCache::default()); + let list_files_cache = + Arc::new(DefaultCache::new(DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT)); let rt = RuntimeEnvBuilder::new() .with_cache_manager( @@ -303,7 +334,8 @@ async fn test_list_files_cache_limit() { #[tokio::test] async fn test_list_files_cache_ttl() { - let list_files_cache = Arc::new(DefaultListFilesCache::default()); + let list_files_cache = + Arc::new(DefaultCache::new(DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT)); let rt = RuntimeEnvBuilder::new() .with_cache_manager( @@ -347,7 +379,8 @@ async fn test_list_files_cache_ttl() { #[tokio::test] async fn test_file_statistics_cache_limit() { - let file_statistics_cache = Arc::new(DefaultFileStatisticsCache::default()); + let file_statistics_cache = + Arc::new(DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT)); let rt = RuntimeEnvBuilder::new() .with_cache_manager( diff --git a/datafusion/core/tests/sql/sql_api.rs b/datafusion/core/tests/sql/sql_api.rs index 290aa737d2742..ca18406a8e40d 100644 --- a/datafusion/core/tests/sql/sql_api.rs +++ b/datafusion/core/tests/sql/sql_api.rs @@ -208,18 +208,112 @@ async fn ddl_can_not_be_planned_by_session_state() { ); } +async fn merge_into_context() -> SessionContext { + let ctx = SessionContext::new(); + ctx.sql("CREATE TABLE target (id INT)").await.unwrap(); + ctx.sql("CREATE TABLE source (id INT)").await.unwrap(); + ctx +} + +async fn assert_merge_sql_error(ctx: &SessionContext, sql: &str, expected: &str) { + let err = ctx.sql(sql).await.unwrap_err(); + assert_contains!(err.strip_backtrace(), expected); +} + +async fn assert_merge_physical_error(ctx: &SessionContext, sql: &str, expected: &str) { + let err = ctx + .sql(sql) + .await + .unwrap() + .create_physical_plan() + .await + .unwrap_err(); + assert_contains!(err.strip_backtrace(), expected); +} + #[tokio::test] -async fn invalid_wrapped_negation_fails_during_optimization() { +async fn merge_into_rejects_source_alias_colliding_with_target_name() { + // Canonicalizing `t.id` to `target.id` must not collapse it onto a source + // that also uses `target` as its qualifier. + let ctx = merge_into_context().await; + + for target_ref in ["target", "public.target", "datafusion.public.target"] { + assert_merge_sql_error( + &ctx, + &format!( + "MERGE INTO {target_ref} AS t USING source AS target \ + ON t.id = target.id WHEN MATCHED THEN DELETE" + ), + &format!( + "MERGE source may not use the target table name '{target_ref}' \ + as a qualifier" + ), + ) + .await; + } +} + +#[tokio::test] +async fn merge_into_rejects_subqueries_correlated_to_target_alias() { + let ctx = merge_into_context().await; + assert_merge_sql_error( + &ctx, + "MERGE INTO target AS t USING source AS s \ + ON EXISTS (SELECT 1 FROM source AS x WHERE x.id = t.id) \ + WHEN MATCHED THEN DELETE", + "MERGE subqueries correlated to target alias 't' are not supported", + ) + .await; + + // Source-correlated and uncorrelated subqueries remain supported through + // logical optimization. + for sql in [ + "MERGE INTO target AS t USING source AS s \ + ON EXISTS (SELECT 1 FROM source AS x WHERE x.id = s.id) \ + WHEN MATCHED THEN DELETE", + "MERGE INTO target AS t USING source AS s \ + ON t.id = ANY (SELECT id FROM source) \ + WHEN MATCHED THEN DELETE", + ] { + assert_merge_physical_error(&ctx, sql, "MERGE INTO not supported for Base table") + .await; + } +} + +#[tokio::test] +async fn merge_into_requires_boolean_conditions() { + let ctx = merge_into_context().await; + + for (sql, expected) in [ + ( + "MERGE INTO target USING source ON 1 WHEN MATCHED THEN DELETE", + "MERGE ON condition must be boolean type, but got Int64", + ), + ( + "MERGE INTO target USING source ON true \ + WHEN MATCHED AND 1 THEN DELETE", + "MERGE WHEN condition must be boolean type, but got Int64", + ), + ( + "MERGE INTO target USING source ON NULL \ + WHEN MATCHED AND NULL THEN DELETE", + "MERGE INTO not supported for Base table", + ), + ] { + assert_merge_physical_error(&ctx, sql, expected).await; + } +} + +#[tokio::test] +async fn invalid_wrapped_negation_fails_during_planning() { let ctx = SessionContext::new(); let err = ctx .sql("SELECT * FROM (SELECT 1) WHERE ((-'a') IS NULL)") .await - .unwrap() - .into_optimized_plan() .unwrap_err(); assert_contains!( err.strip_backtrace(), - "Negation only supports numeric, interval and timestamp types" + "Unary operator '-' only supports signed numeric, interval and timestamp types" ); } diff --git a/datafusion/core/tests/sql/union_comparison.rs b/datafusion/core/tests/sql/union_comparison.rs new file mode 100644 index 0000000000000..87c8c4b8f5bf9 --- /dev/null +++ b/datafusion/core/tests/sql/union_comparison.rs @@ -0,0 +1,505 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +/* +tests for union type comparison coercion. + +when comparing a union type with an "opaque" (non-union) scalar type, the +coercion rule picks the scalar type if any union variant can be cast to it. +the actual extraction at execution time is delegated to arrow's +`cast(Union -> T)`, which selects the source variant using three passes: + +1. exact match: a variant whose type equals the target +2. same type family: e.g. Utf8 / LargeUtf8 / Utf8View are interchangeable, + so Utf8 is preferred over Int32 when the target is Utf8View +3. castable: the first variant (by type_id order) where can_cast_types is true + +rows whose active variant is not the selected one become NULL. + +current limitations exercised by these tests: +- numeric literals default to Int64, so a comparison against `42` won't pick + the Int32 variant exactly +- when multiple variants are equally good in pass 3, the smaller type_id wins +*/ + +use arrow::array::*; +use arrow::buffer::ScalarBuffer; +use arrow::compute::can_cast_types; +use arrow::datatypes::{DataType, Field, Schema, UnionFields, UnionMode}; +use datafusion::assert_batches_eq; +use datafusion::prelude::*; +use datafusion_common::Result; +use std::sync::Arc; + +// create a Union(Int32, Utf8) sparse union array +fn create_sparse_union_array(values: Vec) -> UnionArray { + let union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(); + + let mut int_values = Vec::new(); + let mut str_values = Vec::new(); + let mut type_ids = Vec::new(); + + for value in values { + match value { + UnionValue::Int(v) => { + int_values.push(v); + str_values.push(None); + type_ids.push(0); + } + UnionValue::Str(v) => { + int_values.push(None); + str_values.push(v); + type_ids.push(1); + } + } + } + + let int_array = Int32Array::from(int_values); + let str_array = StringArray::from(str_values); + let type_ids = ScalarBuffer::::from(type_ids); + + UnionArray::try_new( + union_fields, + type_ids, + None, + vec![Arc::new(int_array) as Arc, Arc::new(str_array)], + ) + .unwrap() +} + +#[derive(Debug)] +enum UnionValue { + Int(Option), + Str(Option<&'static str>), +} + +// arrow's cast layer now supports Union -> T whenever any variant can be cast +// to T. this is what the union coercion rule in DataFusion relies on at +// execution time, so we pin the expectation here. +#[test] +fn test_arrow_union_cast_support() { + let union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(); + let union_type = DataType::Union(union_fields, UnionMode::Sparse); + + assert!(can_cast_types(&union_type, &DataType::Int64)); + assert!(can_cast_types(&union_type, &DataType::Int32)); + assert!(can_cast_types(&union_type, &DataType::Utf8)); +} + +#[tokio::test] +async fn test_union_eq_int32() -> Result<()> { + let union_array = create_sparse_union_array(vec![ + UnionValue::Int(Some(67)), + UnionValue::Str(Some("hello")), + UnionValue::Int(Some(123)), + ]); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "val", + DataType::Union( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(), + UnionMode::Sparse, + ), + true, + ), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(union_array), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("test", batch)?; + + let df = ctx + .sql("SELECT id FROM test WHERE val = CAST(67 AS INT)") + .await?; + let results = df.collect().await?; + + let expected = ["+----+", "| id |", "+----+", "| 1 |", "+----+"]; + assert_batches_eq!(expected, &results); + + Ok(()) +} + +#[tokio::test] +async fn test_union_eq_string() -> Result<()> { + let union_array = create_sparse_union_array(vec![ + UnionValue::Int(Some(67)), + UnionValue::Str(Some("hello")), + UnionValue::Str(Some("world")), + ]); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "val", + DataType::Union( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(), + UnionMode::Sparse, + ), + true, + ), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(union_array), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("test", batch)?; + + let df = ctx.sql("SELECT id FROM test WHERE val = 'hello'").await?; + let results = df.collect().await?; + + let expected = ["+----+", "| id |", "+----+", "| 2 |", "+----+"]; + assert_batches_eq!(expected, &results); + + Ok(()) +} + +#[tokio::test] +async fn test_union_comparison_operators() -> Result<()> { + let union_array = create_sparse_union_array(vec![ + UnionValue::Int(Some(10)), + UnionValue::Int(Some(20)), + UnionValue::Int(Some(30)), + UnionValue::Str(Some("foo")), + ]); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "val", + DataType::Union( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(), + UnionMode::Sparse, + ), + true, + ), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(union_array), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("test", batch)?; + + // test > - cast literals to Int32 + let df = ctx + .sql("SELECT id FROM test WHERE val > CAST(15 AS INT)") + .await?; + let results = df.collect().await?; + let expected = ["+----+", "| id |", "+----+", "| 2 |", "| 3 |", "+----+"]; + assert_batches_eq!(expected, &results); + + // test < + let df = ctx + .sql("SELECT id FROM test WHERE val < CAST(15 AS INT)") + .await?; + let results = df.collect().await?; + let expected = ["+----+", "| id |", "+----+", "| 1 |", "+----+"]; + assert_batches_eq!(expected, &results); + + // test != + let df = ctx + .sql("SELECT id FROM test WHERE val != CAST(20 AS INT)") + .await?; + let results = df.collect().await?; + let expected = ["+----+", "| id |", "+----+", "| 1 |", "| 3 |", "+----+"]; + assert_batches_eq!(expected, &results); + + Ok(()) +} + +#[tokio::test] +async fn test_union_with_null_values() -> Result<()> { + let union_array = create_sparse_union_array(vec![ + UnionValue::Int(Some(10)), + UnionValue::Int(None), // null int + UnionValue::Str(Some("foo")), + UnionValue::Str(None), // null string + ]); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "val", + DataType::Union( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(), + UnionMode::Sparse, + ), + true, + ), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(union_array), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("test", batch)?; + + let df = ctx + .sql("SELECT id FROM test WHERE val = CAST(10 AS INT)") + .await?; + let results = df.collect().await?; + let expected = ["+----+", "| id |", "+----+", "| 1 |", "+----+"]; + assert_batches_eq!(expected, &results); + + let df = ctx.sql("SELECT id FROM test WHERE val IS NULL").await?; + let results = df.collect().await?; + + // row 2 has null int and row 4 has null string + // both should appear as null after cast + let expected = ["+----+", "| id |", "+----+", "| 2 |", "| 4 |", "+----+"]; + assert_batches_eq!(expected, &results); + + Ok(()) +} + +#[tokio::test] +async fn test_union_non_matching_variants_are_null() -> Result<()> { + let union_array = create_sparse_union_array(vec![ + UnionValue::Int(Some(10)), + UnionValue::Str(Some("hello")), + UnionValue::Int(Some(30)), + ]); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "val", + DataType::Union( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(), + UnionMode::Sparse, + ), + true, + ), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(union_array), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("test", batch)?; + + // When casting to Int32, the string variant becomes NULL + let df = ctx + .sql("SELECT id, CAST(val AS INT) as val_int FROM test") + .await?; + let results = df.collect().await?; + + let expected = [ + "+----+---------+", + "| id | val_int |", + "+----+---------+", + "| 1 | 10 |", + "| 2 | |", // null because it's a string + "| 3 | 30 |", + "+----+---------+", + ]; + assert_batches_eq!(expected, &results); + + Ok(()) +} + +// tests cast-compatible variant matching +// when comparing Union(Int32, Utf8) with Int64, it finds the Int32 variant and casts it +#[tokio::test] +async fn test_union_cast_compatible_variant() -> Result<()> { + let union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(); + + let union_array = create_sparse_union_array(vec![ + UnionValue::Int(Some(10)), + UnionValue::Str(Some("hello")), + ]); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "val", + DataType::Union(union_fields, UnionMode::Sparse), + true, + ), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(union_array), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("test", batch)?; + + // Int32 variant can be cast to Int64, so this should work + let df = ctx + .sql("SELECT id FROM test WHERE val = CAST(10 AS BIGINT)") + .await?; + let results = df.collect().await?; + + let expected = ["+----+", "| id |", "+----+", "| 1 |", "+----+"]; + assert_batches_eq!(expected, &results); + + Ok(()) +} + +// equality between two identical Union types: the coercion rule keeps the +// common Union type and arrow-ord handles the comparison directly. row 1 has +// the same active variant + value in both columns, row 2 has the same active +// variant but different values, so only row 1 should match. +#[tokio::test] +async fn test_union_eq_same_union() -> Result<()> { + let union_fields = UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("int", DataType::Int32, true), + Field::new("str", DataType::Utf8, true), + ], + ) + .unwrap(); + + let union_array1 = create_sparse_union_array(vec![ + UnionValue::Int(Some(10)), + UnionValue::Str(Some("hello")), + ]); + + let union_array2 = create_sparse_union_array(vec![ + UnionValue::Int(Some(10)), + UnionValue::Str(Some("world")), + ]); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new( + "val1", + DataType::Union(union_fields.clone(), UnionMode::Sparse), + true, + ), + Field::new( + "val2", + DataType::Union(union_fields, UnionMode::Sparse), + true, + ), + ])); + + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(union_array1), + Arc::new(union_array2), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("test", batch)?; + + let df = ctx + .sql("SELECT id FROM test WHERE val1 = val2") + .await + .unwrap(); + + let batches = df.collect().await?; + let expected = ["+----+", "| id |", "+----+", "| 1 |", "+----+"]; + assert_batches_eq!(expected, &batches); + + Ok(()) +} diff --git a/datafusion/core/tests/sql/union_nullable.rs b/datafusion/core/tests/sql/union_nullable.rs new file mode 100644 index 0000000000000..d2dc66336621a --- /dev/null +++ b/datafusion/core/tests/sql/union_nullable.rs @@ -0,0 +1,204 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Regression tests asserting that every batch yielded by a `UNION ALL` +//! reports the union's own declared schema, even when the same column is +//! `NOT NULL` on one leg and nullable on another. See +//! . + +use std::sync::Arc; + +use arrow::array::{Int64Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema}; +use datafusion::prelude::*; +use datafusion_common::Result; + +/// Builds two single-partition tables that agree on `id`/`status` types but +/// disagree on whether `status` is nullable, then runs `UNION ALL` over them. +async fn union_all_mismatched_nullable( + left_nullable: bool, + right_nullable: bool, +) -> Result { + let ctx = SessionContext::new(); + + let schema_a = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("status", DataType::Utf8, left_nullable), + ])); + let batch_a = RecordBatch::try_new( + Arc::clone(&schema_a), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(StringArray::from(vec!["ok", "ok"])), + ], + )?; + ctx.register_batch("table_a", batch_a)?; + + let schema_b = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("status", DataType::Utf8, right_nullable), + ])); + let status_values: Vec> = if right_nullable { + vec![Some("done"), None] + } else { + vec![Some("done"), Some("also-done")] + }; + let batch_b = RecordBatch::try_new( + Arc::clone(&schema_b), + vec![ + Arc::new(Int64Array::from(vec![3, 4])), + Arc::new(StringArray::from(status_values)), + ], + )?; + ctx.register_batch("table_b", batch_b)?; + + ctx.sql( + "SELECT id, status FROM table_a \ + UNION ALL \ + SELECT id, status FROM table_b", + ) + .await +} + +/// The schema DataFusion actually commits to for a query: the logical plan +/// after the `Analyzer` (which includes the `UNION` nullability/type +/// coercion this test targets) and `Optimizer` have run. `DataFrame::schema` +/// alone is not enough here -- it reflects the raw, pre-`Analyzer` plan (see +/// `SessionState::create_logical_plan`), which for a `UNION` still has the +/// first leg's un-coerced type. +fn analyzed_schema(df: &DataFrame) -> Result { + Ok(df + .clone() + .into_optimized_plan()? + .schema() + .as_arrow() + .clone()) +} + +/// Every `RecordBatch` actually produced by a `UNION ALL` must match the +/// query's analyzed output schema field-for-field -- including +/// nullability -- no matter which leg it came from. +async fn assert_every_batch_matches_declared_schema(df: DataFrame) -> Result<()> { + let declared_schema = analyzed_schema(&df)?; + + let batches = df.collect().await?; + assert!(!batches.is_empty()); + for batch in &batches { + assert_eq!( + batch.schema().as_ref(), + &declared_schema, + "a UNION ALL leg produced a RecordBatch whose schema disagrees \ + with the union's declared output schema (commonly a dropped \ + nullable flag) -- this is what downstream consumers that check \ + schema equality across batches (e.g. pyarrow) reject with \ + `ArrowInvalid: Schema at index N was different`" + ); + } + Ok(()) +} + +#[tokio::test] +async fn union_all_same_type_left_not_null_right_nullable() -> Result<()> { + let df = union_all_mismatched_nullable(false, true).await?; + assert!( + analyzed_schema(&df)? + .field_with_name("status")? + .is_nullable() + ); + assert_every_batch_matches_declared_schema(df).await +} + +#[tokio::test] +async fn union_all_same_type_left_nullable_right_not_null() -> Result<()> { + let df = union_all_mismatched_nullable(true, false).await?; + assert!( + analyzed_schema(&df)? + .field_with_name("status")? + .is_nullable() + ); + assert_every_batch_matches_declared_schema(df).await +} + +#[tokio::test] +async fn union_all_same_type_both_not_null_stays_not_null() -> Result<()> { + let df = union_all_mismatched_nullable(false, false).await?; + let declared_schema = analyzed_schema(&df)?; + assert!( + !declared_schema.field_with_name("status")?.is_nullable(), + "status should remain NOT NULL when neither leg is nullable" + ); + assert_every_batch_matches_declared_schema(df).await +} + +/// Same bug, but the coercion also has to widen the *type* (Int32 -> Int64) +/// on one leg. The leg that already matched the target type still needed +/// its nullability reconciled at execution time, independent of whichever +/// legs needed a `CAST`. +#[tokio::test] +async fn union_all_widening_cast_also_fixes_nullable() -> Result<()> { + use arrow::array::Int32Array; + + let ctx = SessionContext::new(); + + let schema_a = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("val", DataType::Int32, false), + ])); + let batch_a = RecordBatch::try_new( + Arc::clone(&schema_a), + vec![ + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![10, 20])), + ], + )?; + ctx.register_batch("table_a", batch_a)?; + + let schema_b = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int64, false), + Field::new("val", DataType::Int64, true), + ])); + let batch_b = RecordBatch::try_new( + Arc::clone(&schema_b), + vec![ + Arc::new(Int64Array::from(vec![3, 4])), + Arc::new(Int64Array::from(vec![Some(30), None])), + ], + )?; + ctx.register_batch("table_b", batch_b)?; + + let df = ctx + .sql( + "SELECT id, val FROM table_a \ + UNION ALL \ + SELECT id, val FROM table_b", + ) + .await?; + + let declared_schema = analyzed_schema(&df)?; + assert_eq!( + declared_schema.field_with_name("val")?.data_type(), + &DataType::Int64 + ); + assert!(declared_schema.field_with_name("val")?.is_nullable()); + + let batches = df.collect().await?; + assert!(!batches.is_empty()); + for batch in &batches { + assert_eq!(batch.schema().as_ref(), &declared_schema); + } + Ok(()) +} diff --git a/datafusion/core/tests/sql/unparser.rs b/datafusion/core/tests/sql/unparser.rs index d6ca872e198c3..355a58fd6f45b 100644 --- a/datafusion/core/tests/sql/unparser.rs +++ b/datafusion/core/tests/sql/unparser.rs @@ -37,14 +37,22 @@ use std::fs::ReadDir; use std::future::Future; +use std::sync::Arc; use arrow::array::RecordBatch; +use arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::Result; +use datafusion::datasource::empty::EmptyTable; +use datafusion::optimizer::{ + OptimizerRule, single_distinct_to_groupby::SingleDistinctToGroupBy, +}; use datafusion::prelude::{ParquetReadOptions, SessionContext}; +use datafusion_catalog::memory::MemorySchemaProvider; +use datafusion_catalog::{CatalogProvider, MemoryCatalogProvider, SchemaProvider}; use datafusion_common::Column; use datafusion_expr::Expr; use datafusion_sql::unparser::Unparser; -use datafusion_sql::unparser::dialect::DefaultDialect; +use datafusion_sql::unparser::dialect::{DefaultDialect, DuckDBDialect}; use itertools::Itertools; use recursive::{set_minimum_stack_size, set_stack_allocation_size}; @@ -218,6 +226,527 @@ async fn sort_batches( df.collect().await } +const ISSUE_22961_QUERY: &str = r#" +SELECT * FROM +( +SELECT + item_id, + order_id, + product_id, + quantity, + unit_price, + quantity * unit_price AS line_total + FROM + "warehouse"."main"."order_items" +) oi +JOIN ( + SELECT + order_id, + customer_id, + order_date, + lower(STATUS) AS STATUS, + lower(channel) AS channel, + coalesce(discount_pct, 0) AS discount_pct, + coalesce(shipping_cost, 0) AS shipping_cost, + STATUS IN ('completed', 'shipped') AS is_fulfilled + FROM + "warehouse"."main"."orders" +) o USING (order_id) +JOIN ( + SELECT + p.product_id, + p.category_id, + p.sku, + p.name AS product_name, + p.price, + p.cost, + p.weight_kg, + p.is_active, + p.stock_qty, + round(p.price - p.cost, 2) AS gross_margin, + round((p.price - p.cost) / nullif(p.price, 0), 4) AS margin_pct, + c.name AS category_name + FROM + "warehouse"."main"."products" p + LEFT JOIN "warehouse"."main"."categories" c USING (category_id) +) p USING (product_id) +"#; + +fn issue_22961_context() -> Result { + let ctx = SessionContext::new(); + + let schema_provider = Arc::new(MemorySchemaProvider::new()); + schema_provider.register_table( + "order_items".to_string(), + Arc::new(EmptyTable::new(Arc::new(Schema::new(vec![ + Field::new("item_id", DataType::Int32, false), + Field::new("order_id", DataType::Int32, true), + Field::new("product_id", DataType::Int32, true), + Field::new("quantity", DataType::Int32, true), + Field::new("unit_price", DataType::Decimal128(10, 2), true), + ])))), + )?; + schema_provider.register_table( + "orders".to_string(), + Arc::new(EmptyTable::new(Arc::new(Schema::new(vec![ + Field::new("order_id", DataType::Int32, false), + Field::new("customer_id", DataType::Int32, true), + Field::new("order_date", DataType::Date32, true), + Field::new("status", DataType::Utf8, true), + Field::new("channel", DataType::Utf8, true), + Field::new("discount_pct", DataType::Decimal128(5, 2), true), + Field::new("shipping_cost", DataType::Decimal128(8, 2), true), + ])))), + )?; + schema_provider.register_table( + "products".to_string(), + Arc::new(EmptyTable::new(Arc::new(Schema::new(vec![ + Field::new("product_id", DataType::Int32, false), + Field::new("category_id", DataType::Int32, true), + Field::new("sku", DataType::Utf8, true), + Field::new("name", DataType::Utf8, true), + Field::new("price", DataType::Decimal128(10, 2), true), + Field::new("cost", DataType::Decimal128(10, 2), true), + Field::new("weight_kg", DataType::Decimal128(6, 3), true), + Field::new("is_active", DataType::Boolean, true), + Field::new("stock_qty", DataType::Int32, true), + ])))), + )?; + schema_provider.register_table( + "categories".to_string(), + Arc::new(EmptyTable::new(Arc::new(Schema::new(vec![ + Field::new("category_id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + Field::new("parent_id", DataType::Int32, true), + Field::new("display_rank", DataType::Int32, true), + ])))), + )?; + + let catalog = Arc::new(MemoryCatalogProvider::new()); + catalog.register_schema("main", schema_provider)?; + ctx.register_catalog("warehouse", catalog); + + Ok(ctx) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_preserves_derived_table_scope() -> Result<()> { + let ctx = issue_22961_context()?; + let plan = ctx.sql(ISSUE_22961_QUERY).await?.into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert!(!sql.contains(r#""o"."__common_expr_1""#)); + assert!(!sql.contains(r#""o"."__common_expr_2""#)); + assert!(sql.contains( + r#"ON "oi"."order_id" = "o"."order_id" INNER JOIN (SELECT "p"."product_id""# + )); + + Ok(()) +} + +// https://github.com/apache/datafusion/issues/23138 +// +// CSE on `coalesce(discount_pct, 0)` factors a shared CAST into an extra inner +// projection, so `SubqueryAlias: o` ends up over two stacked projections. When +// the unparser renders that as nested derived tables it must qualify the +// pass-through `order_id` with a name in scope at each level -- it must not +// rebase it to the outer subquery alias `o`, which is not visible inside the +// inner derived table. +const ISSUE_23138_QUERY: &str = r#" +SELECT * FROM +( + SELECT order_id FROM "warehouse"."main"."order_items" +) oi +JOIN ( + SELECT order_id, coalesce(discount_pct, 0) AS discount_pct_2 + FROM "warehouse"."main"."orders" +) o USING (order_id) +"#; + +#[tokio::test] +async fn optimized_duckdb_unparse_qualifies_nested_passthrough_column() -> Result<()> { + let ctx = issue_22961_context()?; + let plan = ctx.sql(ISSUE_23138_QUERY).await?.into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + // The intermediate derived table has no `o` in scope, so the pass-through + // `order_id` must be unqualified there, not rebased to the subquery alias + // `o` (which is only the base-table alias one level deeper). The bug emitted + // `"o"."order_id"` inside that derived table; the fix emits a bare column. + let expected = concat!( + r#"SELECT "o"."order_id", "o"."discount_pct_2" "#, + r#"FROM "warehouse"."main"."order_items" AS "oi" "#, + r#"INNER JOIN (SELECT "order_id", "#, + r#"CASE WHEN "__common_expr_1" IS NOT NULL "#, + r#"THEN "__common_expr_1" ELSE 0.00 END AS "discount_pct_2" "#, + r#"FROM (SELECT CAST("o"."discount_pct" AS DECIMAL(22,2)) "#, + r#"AS "__common_expr_1", "o"."order_id" "#, + r#"FROM "warehouse"."main"."orders" AS "o")) AS "o" "#, + r#"ON "oi"."order_id" = "o"."order_id""#, + ); + assert_eq!(sql, expected); + + assert!( + sql.contains(r#"(SELECT "order_id", CASE WHEN"#), + "pass-through order_id should be unqualified in derived table: {sql}" + ); + assert!( + !sql.contains(r#"(SELECT "o"."order_id", CASE WHEN"#), + "derived table must not reference out-of-scope alias o: {sql}" + ); + + Ok(()) +} + +// https://github.com/apache/datafusion/issues/23317 +// +// `SingleDistinctToGroupBy` rewrites single DISTINCT aggregates into a +// two-phase aggregate plan. The inner Aggregate defines intermediate fields +// such as `group_alias_0`, `alias1`, and `alias2`. The unparser must preserve +// that inner Aggregate as a derived table before the outer Aggregate +// references those fields. +// +// Without `SingleDistinctToGroupBy`, the Aggregate still sits over an unnamed +// derived Projection. In that SQL scope, base table aliases `cs` and `c` are no +// longer visible, so aggregate expressions must refer to the derived table's +// output columns unqualified. +const ISSUE_23317_QUERY: &str = r#" +WITH cohort AS ( + SELECT + signup_year, + sum(customers) AS customers, + sum(revenue) AS revenue + FROM + ( + SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers, + round(sum(cs.total_revenue), 2) AS revenue + FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) + GROUP BY + 1 + ) + GROUP BY + signup_year +) +SELECT + * +FROM + cohort +"#; + +const ISSUE_23317_HAVING_QUERY: &str = r#" +SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers +FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) +GROUP BY + 1 +HAVING + count(DISTINCT cs.customer_id) > 0 +"#; + +const ISSUE_23317_QUALIFY_QUERY: &str = r#" +SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers, + row_number() OVER (ORDER BY date_part('year', c.signup_date)) AS rn +FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) +GROUP BY + 1 +QUALIFY + rn = 1 AND count(DISTINCT cs.customer_id) > 0 +"#; + +// https://github.com/apache/datafusion/issues/23668 +// +// Extends the #23317 aggregate-scope fix to the window and ORDER BY clauses. +// Reuses issue_23317_context() (same derived-projection shape). + +// Window sorting by an aggregate, over a derived-projection input. Already +// correct today; this locks the OVER clause against keeping the out-of-scope +// `cs` qualifier across the refactor. +const ISSUE_23668_WINDOW_QUERY: &str = r#" +SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers, + row_number() OVER (ORDER BY count(DISTINCT cs.customer_id) DESC) AS rn +FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) +GROUP BY + 1 +"#; + +// ORDER BY an aggregate that is NOT selected, so it can't use a select alias +// and is unprojected through the Aggregate. It must be normalized like the +// SELECT list, not keep the out-of-scope `cs` qualifier. +const ISSUE_23668_ORDER_BY_QUERY: &str = r#" +SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers +FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) +GROUP BY + 1 +ORDER BY + round(sum(cs.total_revenue), 2) DESC +"#; + +// ORDER BY a selected aggregate keeps a top-level Sort (the direct `Sort` arm, +// vs the projection-absorbed one above). It resolves to the select alias, so +// this covers routing only -- the normalization in that arm isn't reachable +// from SQL (an unselected aggregate takes the absorbed path above instead). +const ISSUE_23668_TOP_LEVEL_SORT_QUERY: &str = r#" +SELECT + date_part('year', c.signup_date) AS signup_year, + count(DISTINCT cs.customer_id) AS customers +FROM + "warehouse"."main"."sales" cs + JOIN "warehouse"."main"."customers" c USING (customer_id) +GROUP BY + 1 +ORDER BY + customers DESC +"#; + +fn issue_23317_context() -> Result { + let ctx = SessionContext::new(); + + let schema_provider = Arc::new(MemorySchemaProvider::new()); + schema_provider.register_table( + "customers".to_string(), + Arc::new(EmptyTable::new(Arc::new(Schema::new(vec![ + Field::new("customer_id", DataType::Int32, false), + Field::new("signup_date", DataType::Date32, true), + ])))), + )?; + schema_provider.register_table( + "sales".to_string(), + Arc::new(EmptyTable::new(Arc::new(Schema::new(vec![ + Field::new("customer_id", DataType::Int32, false), + Field::new("total_revenue", DataType::Decimal128(12, 2), true), + ])))), + )?; + + let catalog = Arc::new(MemoryCatalogProvider::new()); + catalog.register_schema("main", schema_provider)?; + ctx.register_catalog("warehouse", catalog); + + Ok(ctx) +} + +async fn assert_issue_23317_unparsed_sql_plans( + ctx: &SessionContext, + sql: &str, +) -> Result<()> { + ctx.sql(sql).await?.into_optimized_plan()?; + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_preserves_nested_aggregate_scope() -> Result<()> { + let ctx = issue_23317_context()?; + let plan = ctx.sql(ISSUE_23317_QUERY).await?.into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains(concat!( + r#"FROM (SELECT sum("total_revenue") AS "alias2", "#, + r#"date_part('year', "signup_date") AS "group_alias_0", "#, + r#""customer_id" AS "alias1" "# + )), + "inner aggregate should define the aliases before the outer aggregate uses them: {sql}", + ); + assert!( + !sql.contains(r#"date_part('year', "c"."signup_date") AS "group_alias_0""#), + "inner aggregate must not reference out-of-scope alias c: {sql}", + ); + + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_unqualifies_aggregate_input_projection() -> Result<()> { + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx.sql(ISSUE_23317_QUERY).await?.into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains( + r#"SELECT date_part('year', "signup_date") AS "signup_year", count(DISTINCT "customer_id") AS "customers", round(sum("total_revenue"), 2) AS "revenue" FROM ("# + ), + "aggregate expressions should resolve against the derived projection output: {sql}", + ); + assert!( + !sql.contains(r#"date_part('year', "c"."signup_date") AS "signup_year""#), + "derived aggregate must not reference out-of-scope alias c: {sql}", + ); + + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_having_unqualifies_agg_input() -> Result<()> { + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx + .sql(ISSUE_23317_HAVING_QUERY) + .await? + .into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains(r#"HAVING (count(DISTINCT "customer_id") > 0)"#), + "HAVING aggregate should resolve against the derived projection output: {sql}", + ); + assert!( + !sql.contains(r#"count(DISTINCT "cs"."customer_id")"#), + "HAVING must not reference out-of-scope alias cs: {sql}", + ); + + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_qualify_unqualifies_agg_input() -> Result<()> { + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx + .sql(ISSUE_23317_QUALIFY_QUERY) + .await? + .into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains("QUALIFY"), + "expected QUALIFY clause in unparsed SQL: {sql}", + ); + assert!( + sql.contains(r#"count(DISTINCT "customer_id") > 0"#), + "QUALIFY aggregate should resolve against the derived projection output: {sql}", + ); + assert!( + !sql.contains(r#"count(DISTINCT "cs"."customer_id")"#), + "QUALIFY must not reference out-of-scope alias cs: {sql}", + ); + + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_window_over_agg_unqualifies_input() -> Result<()> { + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx + .sql(ISSUE_23668_WINDOW_QUERY) + .await? + .into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains(r#"OVER (ORDER BY count(DISTINCT "customer_id")"#), + "window ORDER BY aggregate should resolve against the derived projection output: {sql}", + ); + assert!( + !sql.contains(r#"count(DISTINCT "cs"."customer_id")"#), + "window OVER clause must not reference out-of-scope alias cs: {sql}", + ); + + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_order_by_unqualifies_agg_input() -> Result<()> { + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx + .sql(ISSUE_23668_ORDER_BY_QUERY) + .await? + .into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains(r#"ORDER BY round(sum("total_revenue"), 2)"#), + "ORDER BY aggregate should resolve against the derived projection output: {sql}", + ); + assert!( + !sql.contains(r#"sum("cs"."total_revenue")"#), + "ORDER BY must not reference out-of-scope alias cs: {sql}", + ); + + Ok(()) +} + +#[tokio::test] +async fn optimized_duckdb_unparse_top_level_sort_over_agg_uses_select_alias() -> Result<()> +{ + let ctx = issue_23317_context()?; + assert!(ctx.remove_optimizer_rule(SingleDistinctToGroupBy::new().name())); + + let plan = ctx + .sql(ISSUE_23668_TOP_LEVEL_SORT_QUERY) + .await? + .into_optimized_plan()?; + let dialect = DuckDBDialect::new(); + let unparser = Unparser::new(&dialect); + let sql = unparser.plan_to_sql(&plan)?.to_string(); + + assert_issue_23317_unparsed_sql_plans(&ctx, &sql).await?; + + assert!( + sql.contains(r#"ORDER BY "customers""#), + "top-level ORDER BY should resolve to the select alias: {sql}", + ); + assert!( + !sql.contains(r#""cs"."customer_id") AS "customers""#), + "aggregate output must not reference out-of-scope alias cs: {sql}", + ); + + Ok(()) +} + /// The outcome of running a single roundtrip test. /// /// A successful test produces [`TestCaseResult::Success`]. diff --git a/datafusion/core/tests/tpcds_planning.rs b/datafusion/core/tests/tpcds_planning.rs index 3ad74962bc2c0..c1c3265e521d6 100644 --- a/datafusion/core/tests/tpcds_planning.rs +++ b/datafusion/core/tests/tpcds_planning.rs @@ -1036,10 +1036,10 @@ async fn regression_test(query_no: u8, create_physical: bool) -> Result<()> { for table in &tables { ctx.register_table( table.name.as_str(), - Arc::new(MemTable::try_new( - Arc::new(table.schema.clone()), - vec![vec![]], - )?), + Arc::new( + MemTable::try_new(Arc::new(table.schema.clone()), vec![vec![]])? + .with_constraints(table.constraints.clone()), + ), )?; } diff --git a/datafusion/core/tests/user_defined/insert_operation.rs b/datafusion/core/tests/user_defined/insert_operation.rs index 326c767d97610..0eefcdb551a65 100644 --- a/datafusion/core/tests/user_defined/insert_operation.rs +++ b/datafusion/core/tests/user_defined/insert_operation.rs @@ -28,7 +28,9 @@ use datafusion_common::config::Dialect; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_expr::{Expr, TableType, dml::InsertOp}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_plan::execution_plan::SchedulingType; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, execution_plan::SchedulingType, +}; use datafusion_physical_plan::{ DisplayAs, ExecutionPlan, PlanProperties, execution_plan::{Boundedness, EmissionType}, @@ -162,14 +164,25 @@ impl ExecutionPlan for TestInsertExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert!(children.is_empty()); Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -180,18 +193,11 @@ impl ExecutionPlan for TestInsertExec { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.plan_properties.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion/core/tests/user_defined/mod.rs b/datafusion/core/tests/user_defined/mod.rs index bc9949f5d681c..4dad3ec4577d9 100644 --- a/datafusion/core/tests/user_defined/mod.rs +++ b/datafusion/core/tests/user_defined/mod.rs @@ -41,3 +41,7 @@ mod relation_planner; /// Tests for insert operations mod insert_operation; + +/// Tests for `StatisticsRequest`s flowing from a custom optimizer rule +/// through the physical planner into a custom `TableProvider`. +mod statistics_requests; diff --git a/datafusion/core/tests/user_defined/statistics_requests.rs b/datafusion/core/tests/user_defined/statistics_requests.rs new file mode 100644 index 0000000000000..64c6676c23746 --- /dev/null +++ b/datafusion/core/tests/user_defined/statistics_requests.rs @@ -0,0 +1,221 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End-to-end test that a *custom* optimizer rule can annotate a +//! `TableScan` with `StatisticsRequest`s and have them reach a *custom* +//! `TableProvider`'s `scan_with_args`. +//! +//! DataFusion ships no rule that populates `TableScan::statistics_requests` +//! and no provider that consumes `ScanArgs::statistics_requests`. This test +//! plays both roles, demonstrating that the request-side hooks are +//! sufficient to build the whole feature outside of DataFusion. + +use std::sync::{Arc, Mutex}; + +use arrow::array::{Int64Array, RecordBatch}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use async_trait::async_trait; +use datafusion::catalog::{ScanArgs, ScanResult, Session, TableProvider}; +use datafusion::common::tree_node::Transformed; +use datafusion::common::{Column, Result}; +use datafusion::datasource::TableType; +use datafusion::datasource::memory::MemorySourceConfig; +use datafusion::execution::context::SessionContext; +use datafusion::execution::session_state::SessionStateBuilder; +use datafusion::logical_expr::statistics::StatisticsRequest; +use datafusion::logical_expr::{Expr, LogicalPlan}; +use datafusion::optimizer::{ApplyOrder, OptimizerConfig, OptimizerRule}; +use datafusion::physical_plan::ExecutionPlan; + +/// A custom optimizer rule that annotates every `TableScan` with a +/// `RowCount` request plus a `Min` request for each of its columns. +/// +/// This stands in for whatever request-derivation logic an external +/// implementer would write (e.g. Min/Max for sort keys, DistinctCount for +/// join keys). Here it is intentionally trivial and deterministic. +#[derive(Debug)] +struct RequestColumnStatistics; + +impl OptimizerRule for RequestColumnStatistics { + fn name(&self) -> &str { + "test_request_column_statistics" + } + + fn apply_order(&self) -> Option { + Some(ApplyOrder::TopDown) + } + + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> Result> { + let LogicalPlan::TableScan(mut scan) = plan else { + return Ok(Transformed::no(plan)); + }; + // Insert into the scan's existing request set. `BTreeSet::insert` + // reports whether the value was new, so the rule is idempotent — and + // composes with other rules' requests for free: re-inserting an + // existing request is a no-op, and we report `Transformed::yes` only + // when something was actually added, so the optimizer reaches a + // fixpoint without a manual "already visited" guard. + let mut changed = scan.statistics_requests.insert(StatisticsRequest::RowCount); + for field in scan.projected_schema.fields() { + let req = + StatisticsRequest::Min(Arc::new(Column::new_unqualified(field.name()))); + changed |= scan.statistics_requests.insert(req); + } + Ok(if changed { + Transformed::yes(LogicalPlan::TableScan(scan)) + } else { + Transformed::no(LogicalPlan::TableScan(scan)) + }) + } +} + +/// A `TableProvider` that records the `statistics_requests` it was asked +/// for, so the test can assert what reached it. +#[derive(Debug)] +struct RecordingTable { + schema: SchemaRef, + batch: RecordBatch, + last_requests: Arc>>, +} + +#[async_trait] +impl TableProvider for RecordingTable { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec>, + _filters: &[Expr], + _limit: Option, + ) -> Result> { + Ok(MemorySourceConfig::try_new_exec( + &[vec![self.batch.clone()]], + Arc::clone(&self.schema), + projection.cloned(), + )?) + } + + async fn scan_with_args<'a>( + &self, + state: &dyn Session, + args: ScanArgs<'a>, + ) -> Result { + // Record what reached us, then delegate to `scan`. + *self.last_requests.lock().unwrap() = args.statistics_requests().to_vec(); + let plan = self + .scan( + state, + args.projection().map(|p| p.to_vec()).as_ref(), + args.filters().unwrap_or(&[]), + args.limit(), + ) + .await?; + Ok(ScanResult::new(plan)) + } +} + +fn make_table() -> (Arc, Arc>>) { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int64Array::from(vec![1, 2, 3])), + Arc::new(Int64Array::from(vec![10, 20, 30])), + ], + ) + .unwrap(); + let last_requests = Arc::new(Mutex::new(Vec::new())); + let provider = Arc::new(RecordingTable { + schema, + batch, + last_requests: Arc::clone(&last_requests), + }); + (provider, last_requests) +} + +#[tokio::test] +async fn custom_rule_requests_reach_custom_provider() -> Result<()> { + let (provider, last_requests) = make_table(); + + let state = SessionStateBuilder::new() + .with_default_features() + .with_optimizer_rule(Arc::new(RequestColumnStatistics)) + .build(); + let ctx = SessionContext::new_with_state(state); + ctx.register_table("t", provider)?; + + ctx.sql("SELECT a, b FROM t").await?.collect().await?; + + let got = last_requests.lock().unwrap().clone(); + assert_eq!( + got.len(), + 3, + "expected RowCount + Min(a) + Min(b), got {got:?}" + ); + assert!( + got.contains(&StatisticsRequest::RowCount), + "expected RowCount, got {got:?}" + ); + assert!( + got.contains(&StatisticsRequest::Min(Arc::new(Column::new_unqualified( + "a" + )))), + "expected Min(a), got {got:?}" + ); + assert!( + got.contains(&StatisticsRequest::Min(Arc::new(Column::new_unqualified( + "b" + )))), + "expected Min(b), got {got:?}" + ); + Ok(()) +} + +#[tokio::test] +async fn no_requests_without_a_rule() -> Result<()> { + // Without a rule populating `TableScan::statistics_requests`, the + // provider sees an empty request list — stock DataFusion behavior. + let (provider, last_requests) = make_table(); + let ctx = SessionContext::new(); + ctx.register_table("t", provider)?; + + ctx.sql("SELECT a, b FROM t").await?.collect().await?; + + assert!( + last_requests.lock().unwrap().is_empty(), + "expected no requests without a custom rule" + ); + Ok(()) +} diff --git a/datafusion/core/tests/user_defined/user_defined_aggregates.rs b/datafusion/core/tests/user_defined/user_defined_aggregates.rs index 7d22c5df70dfc..323925bcfaf82 100644 --- a/datafusion/core/tests/user_defined/user_defined_aggregates.rs +++ b/datafusion/core/tests/user_defined/user_defined_aggregates.rs @@ -872,12 +872,22 @@ impl GroupsAccumulator for TestGroupsAccumulator { &mut self, _values: &[ArrayRef], _group_indices: &[usize], - _opt_filter: Option<&arrow::array::BooleanArray>, _total_num_groups: usize, ) -> Result<()> { Ok(()) } + fn convert_to_state( + &self, + values: &[ArrayRef], + _opt_filter: Option<&arrow::array::BooleanArray>, + ) -> Result> { + let len = values.first().map_or(0, |value| value.len()); + Ok(vec![ + Arc::new(PrimitiveArray::::from_value(self.result, len)) + as ArrayRef, + ]) + } fn size(&self) -> usize { size_of::() } diff --git a/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs b/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs index 58a5cb803982b..5b552e5369ef7 100644 --- a/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs +++ b/datafusion/core/tests/user_defined/user_defined_async_scalar_functions.rs @@ -18,14 +18,15 @@ use std::sync::Arc; use arrow::array::{Int32Array, RecordBatch, StringArray}; -use arrow::datatypes::{DataType, Field, Schema}; +use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use async_trait::async_trait; use datafusion::prelude::*; use datafusion_common::test_util::format_batches; use datafusion_common::{Result, assert_batches_eq}; use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; use datafusion_expr::{ - ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, }; fn register_table_and_udf() -> Result { @@ -113,6 +114,110 @@ async fn test_async_udf_metrics() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_async_udf_preserves_result_field_metadata() -> Result<()> { + #[derive(Debug, PartialEq, Eq, Hash, Clone)] + struct AsyncExtensionUDF { + signature: Signature, + } + + impl Default for AsyncExtensionUDF { + fn default() -> Self { + Self { + signature: Signature::exact(vec![DataType::Utf8], Volatility::Volatile), + } + } + } + + impl ScalarUDFImpl for AsyncExtensionUDF { + fn name(&self) -> &str { + "async_extension" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Utf8) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + Ok(args.arg_fields[0] + .as_ref() + .clone() + .with_name(self.name()) + .with_metadata(std::collections::HashMap::from([( + "ARROW:extension:name".to_string(), + "test.async.extension".to_string(), + )])) + .into()) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + panic!("Call invoke_async_with_args instead") + } + } + + #[async_trait] + impl AsyncScalarUDFImpl for AsyncExtensionUDF { + async fn invoke_async_with_args( + &self, + args: ScalarFunctionArgs, + ) -> Result { + Ok(args.args[0].clone()) + } + } + + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("value", DataType::Utf8, false), + ])), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StringArray::from(vec!["one", "two", "three"])), + ], + )?; + + let ctx = SessionContext::new(); + ctx.register_batch("test_table", batch)?; + ctx.register_udf( + AsyncScalarUDF::new(Arc::new(AsyncExtensionUDF::default())).into_scalar_udf(), + ); + + let result = ctx + .sql("SELECT async_extension(value) AS result FROM test_table") + .await? + .collect() + .await?; + + assert_eq!(result[0].schema().field(0).name(), "result"); + assert_eq!( + result[0] + .schema() + .field(0) + .metadata() + .get("ARROW:extension:name"), + Some(&"test.async.extension".to_string()) + ); + + assert_batches_eq!( + &[ + "+--------+", + "| result |", + "+--------+", + "| one |", + "| two |", + "| three |", + "+--------+", + ], + &result + ); + + Ok(()) +} + #[derive(Debug, PartialEq, Eq, Hash, Clone)] struct TestAsyncUDFImpl { batch_size: usize, @@ -162,6 +267,7 @@ impl AsyncScalarUDFImpl for TestAsyncUDFImpl { } /// Simulates calling an async external service +#[expect(clippy::unused_async)] async fn call_external_service(arg1: ColumnarValue) -> Result { Ok(arg1) } diff --git a/datafusion/core/tests/user_defined/user_defined_plan.rs b/datafusion/core/tests/user_defined/user_defined_plan.rs index 505468a19cd37..da7fdd88793e3 100644 --- a/datafusion/core/tests/user_defined/user_defined_plan.rs +++ b/datafusion/core/tests/user_defined/user_defined_plan.rs @@ -67,13 +67,14 @@ use arrow::{ array::Int64Array, datatypes::SchemaRef, record_batch::RecordBatch, util::pretty::pretty_format_batches, }; +use datafusion::catalog::Session; use datafusion::execution::session_state::SessionStateBuilder; use datafusion::{ common::cast::as_int64_array, common::{DFSchemaRef, arrow_datafusion_err}, error::{DataFusionError, Result}, execution::{ - context::{QueryPlanner, SessionState, TaskContext}, + context::{QueryPlanner, TaskContext}, runtime_env::RuntimeEnv, }, logical_expr::{ @@ -98,9 +99,11 @@ use datafusion_expr::{FetchType, InvariantLevel, Projection, SortExpr}; use datafusion_optimizer::AnalyzerRule; use datafusion_optimizer::optimizer::ApplyOrder; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use async_trait::async_trait; use datafusion_common::cast::as_string_view_array; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use futures::{Stream, StreamExt}; /// Execute the specified sql and return the resulting record batches @@ -468,7 +471,7 @@ impl QueryPlanner for TopKQueryPlanner { async fn create_physical_plan( &self, logical_plan: &LogicalPlan, - session_state: &SessionState, + session_state: &dyn Session, ) -> Result> { // Teach the default physical planner how to plan TopK nodes. let physical_planner = @@ -520,7 +523,7 @@ impl OptimizerRule for TopKOptimizerRule { if let LogicalPlan::Sort(Sort { expr, input, .. }) = limit.input.as_ref() && expr.len() == 1 { - // we found a sort with a single sort expr, replace with a a TopK + // we found a sort with a single sort expr, replace with a TopK return Ok(Transformed::yes(LogicalPlan::Extension(Extension { node: Arc::new(TopKPlanNode { k: fetch, @@ -631,7 +634,8 @@ impl ExtensionPlanner for TopKPlanner { node: &dyn UserDefinedLogicalNode, logical_inputs: &[&LogicalPlan], physical_inputs: &[Arc], - _session_state: &SessionState, + _session_state: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, ) -> Result>> { Ok( if let Some(topk_node) = node.as_any().downcast_ref::() { @@ -710,21 +714,36 @@ impl ExecutionPlan for TopKExec { &self.cache } - fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + ]) } fn children(&self) -> Vec<&Arc> { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(TopKExec::new(children[0].clone(), self.k))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Execute one partition and return an iterator over RecordBatch fn execute( &self, @@ -747,18 +766,11 @@ impl ExecutionPlan for TopKExec { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion::physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion/datasource-arrow/Cargo.toml b/datafusion/datasource-arrow/Cargo.toml index 2718e424c6386..6f50135403d69 100644 --- a/datafusion/datasource-arrow/Cargo.toml +++ b/datafusion/datasource-arrow/Cargo.toml @@ -42,6 +42,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } itertools = { workspace = true } @@ -65,3 +66,10 @@ path = "src/mod.rs" # This feature is deprecated, as core functionality in the SpillManager requires all features # it enabled, and will be removed in a future version. compression = [] +# Enables `FileSource::try_to_proto` on `ArrowSource` and the `ArrowScan` decode +# entry point. Mirrors the `proto` feature on `datafusion-datasource`. +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] diff --git a/datafusion/datasource-arrow/src/file_format.rs b/datafusion/datasource-arrow/src/file_format.rs index 9297486ad66e7..c50ad98dfca0b 100644 --- a/datafusion/datasource-arrow/src/file_format.rs +++ b/datafusion/datasource-arrow/src/file_format.rs @@ -37,7 +37,6 @@ use datafusion_common::{ internal_datafusion_err, not_impl_err, }; use datafusion_common_runtime::{JoinSet, SpawnedTask}; -use datafusion_datasource::TableSchema; use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; @@ -45,6 +44,7 @@ use datafusion_datasource::sink::{DataSink, DataSinkExec}; use datafusion_datasource::write::{ ObjectWriterBuilder, SharedBuffer, get_writer_schema, }; +use datafusion_datasource::{TableSchema, TableSchemaBuilder}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::LexRequirement; @@ -197,10 +197,9 @@ impl FileFormat for ArrowFormat { .object_meta .location; - let table_schema = TableSchema::new( - Arc::clone(conf.file_schema()), - conf.table_partition_cols().clone(), - ); + let table_schema = TableSchemaBuilder::from(conf.file_schema()) + .with_table_partition_cols(conf.table_partition_cols().clone()) + .build(); let mut source: Arc = match is_object_in_arrow_ipc_file_format(object_store, object_location).await @@ -357,7 +356,7 @@ impl DisplayAs for ArrowFileSink { } DisplayFormatType::TreeRender => { writeln!(f, "format: arrow")?; - write!(f, "file={}", &self.config.original_url) + write!(f, "file={}", self.config.original_url) } } } @@ -381,7 +380,7 @@ impl DataSink for ArrowFileSink { // Custom implementation of inferring schema. Should eventually be moved upstream to arrow-rs. // See -const ARROW_MAGIC: [u8; 6] = [b'A', b'R', b'R', b'O', b'W', b'1']; +const ARROW_MAGIC: [u8; 6] = *b"ARROW1"; const CONTINUATION_MARKER: [u8; 4] = [0xff; 4]; async fn infer_stream_schema( @@ -549,6 +548,7 @@ mod tests { AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF, }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList}; use object_store::{chunked::ChunkedStore, memory::InMemory}; struct MockSession { @@ -575,6 +575,10 @@ mod tests { &self.config } + fn catalog_list(&self) -> Arc { + Arc::new(EmptyCatalogProviderList) + } + async fn create_physical_plan( &self, _logical_plan: &LogicalPlan, @@ -594,7 +598,7 @@ mod tests { unimplemented!() } - fn higher_order_functions(&self) -> &HashMap> { + fn higher_order_functions(&self) -> &HashMap> { unimplemented!() } diff --git a/datafusion/datasource-arrow/src/source.rs b/datafusion/datasource-arrow/src/source.rs index 061f130f24131..f51e1c100934d 100644 --- a/datafusion/datasource-arrow/src/source.rs +++ b/datafusion/datasource-arrow/src/source.rs @@ -341,7 +341,7 @@ impl FileSource for ArrowSource { // The Arrow IPC stream format doesn't support range-based parallel reading // because it lacks a footer with the information that would be needed to // make range-based parallel reading practical. Without the data in the - // footer you would either need to read the the entire file and record the + // footer you would either need to read the entire file and record the // offsets of the record batches and dictionaries, essentially recreating // the footer's contents, or else each partition would need to read the // entire file up to the correct offset which is a lot of duplicate I/O. @@ -397,15 +397,68 @@ impl FileSource for ArrowSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit projection expressions - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in &self.projection.source { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f) + } + + /// Emit an `ArrowScan` node wrapping the shared base config. + /// + /// Decoding defaults to the IPC file format because protobuf does not + /// distinguish it from the IPC stream format. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ArrowScan( + protobuf::ArrowScanExecNode { + base_conf: Some(base.try_to_proto(ctx)?), + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl ArrowSource { + /// Reconstructs a `DataSourceExec` from a protobuf `ArrowScan`. + /// + /// Defaults to the IPC file format because protobuf does not distinguish it + /// from the IPC stream format. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::file_scan_config::FileScanConfig; + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::ArrowScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not an ArrowScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ArrowScanExecNode is missing required field 'base_conf'" + ) + })?; + + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + let source = Arc::new(ArrowSource::new_file_source(table_schema)); + let scan_conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?; + Ok(DataSourceExec::from_data_source(scan_conf)) } } diff --git a/datafusion/datasource-avro/Cargo.toml b/datafusion/datasource-avro/Cargo.toml index adc2be1cb8f24..70b675d63f427 100644 --- a/datafusion/datasource-avro/Cargo.toml +++ b/datafusion/datasource-avro/Cargo.toml @@ -30,6 +30,15 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +# Enables `FileSource::try_to_proto` on `AvroSource` and the `AvroScan` decode +# entry point. Mirrors the `proto` feature on `datafusion-datasource`. +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] + [dependencies] arrow = { workspace = true } arrow-avro = { workspace = true } @@ -39,6 +48,7 @@ datafusion-common = { workspace = true, features = ["object_store"] } datafusion-datasource = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/datasource-avro/src/source.rs b/datafusion/datasource-avro/src/source.rs index b80d4f462e425..fcc50b559f00b 100644 --- a/datafusion/datasource-avro/src/source.rs +++ b/datafusion/datasource-avro/src/source.rs @@ -173,15 +173,61 @@ impl FileSource for AvroSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit projection expressions - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in &self.projection.source { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f) + } + + /// Emit an `AvroScan` node wrapping the shared base config. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let node = protobuf::AvroScanExecNode { + base_conf: Some(base.try_to_proto(ctx)?), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::AvroScan(node)), + })) + } +} + +#[cfg(feature = "proto")] +impl AvroSource { + /// Reconstructs a `DataSourceExec` from a protobuf `AvroScan`. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::AvroScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not an AvroScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "AvroScanExecNode is missing required field 'base_conf'" + ) + })?; + + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + let source = Arc::new(AvroSource::new(table_schema)); + + let conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?; + Ok(DataSourceExec::from_data_source(conf)) } } diff --git a/datafusion/datasource-csv/Cargo.toml b/datafusion/datasource-csv/Cargo.toml index 295092512742b..7e7195dfda9d5 100644 --- a/datafusion/datasource-csv/Cargo.toml +++ b/datafusion/datasource-csv/Cargo.toml @@ -30,6 +30,14 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +# Enables protobuf serialization hooks for CSV sources and sinks. +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } @@ -41,6 +49,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/datasource-csv/src/file_format.rs b/datafusion/datasource-csv/src/file_format.rs index 9fdd688037682..c0d22b80f08d0 100644 --- a/datafusion/datasource-csv/src/file_format.rs +++ b/datafusion/datasource-csv/src/file_format.rs @@ -158,7 +158,6 @@ impl CsvFormat { .map_err(|e| DataFusionError::ObjectStore(Box::new(e))) .boxed(), ) - .await .map_err(DataFusionError::from) .left_stream(), Err(e) => { @@ -168,9 +167,9 @@ impl CsvFormat { stream.boxed() } - /// Convert a stream of bytes into a stream of of [`Bytes`] containing newline + /// Convert a stream of bytes into a stream of [`Bytes`] containing newline /// delimited CSV records, while accounting for `\` and `"`. - pub async fn read_to_delimited_chunks_from_stream<'a>( + pub fn read_to_delimited_chunks_from_stream<'a>( &self, stream: BoxStream<'a, Result>, ) -> BoxStream<'a, Result> { @@ -393,7 +392,7 @@ impl FileFormat for CsvFormat { .await .map_err(|err| { DataFusionError::Context( - format!("Error when processing CSV file {}", &object.location), + format!("Error when processing CSV file {}", object.location), Box::new(err), ) })?; @@ -759,7 +758,7 @@ impl DisplayAs for CsvSink { } DisplayFormatType::TreeRender => { writeln!(f, "format: csv")?; - write!(f, "file={}", &self.config.original_url) + write!(f, "file={}", self.config.original_url) } } } @@ -826,6 +825,153 @@ impl DataSink for CsvSink { ) -> Result { FileSink::write_all(self, data, context).await } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let input = ctx.encode_child(exec.input())?; + let sort_order = exec.encode_sort_order(ctx)?; + let sink = protobuf::CsvSink::try_from(self)?; + let node = protobuf::CsvSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(exec.schema().as_ref().try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl TryFrom<&CsvSink> for datafusion_proto_models::protobuf::CsvSink { + type Error = DataFusionError; + + fn try_from(value: &CsvSink) -> Result { + Ok(Self { + config: Some(value.config().try_into()?), + writer_options: Some(value.writer_options().try_into()?), + }) + } +} + +#[cfg(feature = "proto")] +impl TryFrom<&datafusion_proto_models::protobuf::CsvSink> for CsvSink { + type Error = DataFusionError; + + fn try_from(value: &datafusion_proto_models::protobuf::CsvSink) -> Result { + let config = + FileSinkConfig::try_from(value.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSink is missing required field 'config'" + ) + })?)?; + let writer_options = value + .writer_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSink is missing required field 'writer_options'" + ) + })? + .try_into()?; + + Ok(Self::new(config, writer_options)) + } +} + +#[cfg(feature = "proto")] +impl CsvSink { + /// Reconstructs a [`DataSinkExec`] containing a `CsvSink` from protobuf. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let sink_node = datafusion_physical_plan::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::CsvSink, + "CsvSink", + ); + let input = ctx.decode_required_child( + sink_node.input.as_deref(), + "CsvSinkExecNode", + "input", + )?; + let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvSinkExecNode is missing required field 'sink'" + ) + })?; + let data_sink = CsvSink::try_from(proto_sink)?; + let sort_order = DataSinkExec::decode_sort_order( + sink_node.sort_order.as_ref(), + ctx, + input.schema().as_ref(), + )?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } +} + +/// Encode a [`CsvFormatFactory`]'s options as their protobuf form. +/// +/// The reverse direction is `From<&protobuf::CsvOptions> for CsvOptions` in +/// `datafusion-proto-models`: `CsvOptions` is a `datafusion-common` type, so +/// that half cannot live here. +#[cfg(feature = "proto")] +impl From<&CsvFormatFactory> for datafusion_proto_models::protobuf::CsvOptions { + fn from(factory: &CsvFormatFactory) -> Self { + if let Some(options) = &factory.options { + datafusion_proto_models::protobuf::CsvOptions { + has_header: options.has_header.map_or(vec![], |v| vec![v as u8]), + delimiter: vec![options.delimiter], + quote: vec![options.quote], + terminator: options.terminator.map_or(vec![], |v| vec![v]), + escape: options.escape.map_or(vec![], |v| vec![v]), + double_quote: options.double_quote.map_or(vec![], |v| vec![v as u8]), + compression: options.compression as i32, + schema_infer_max_rec: options.schema_infer_max_rec.map(|v| v as u64), + date_format: options.date_format.clone().unwrap_or_default(), + datetime_format: options.datetime_format.clone().unwrap_or_default(), + timestamp_format: options.timestamp_format.clone().unwrap_or_default(), + timestamp_tz_format: options + .timestamp_tz_format + .clone() + .unwrap_or_default(), + time_format: options.time_format.clone().unwrap_or_default(), + null_value: options.null_value.clone().unwrap_or_default(), + null_regex: options.null_regex.clone().unwrap_or_default(), + comment: options.comment.map_or(vec![], |v| vec![v]), + newlines_in_values: options + .newlines_in_values + .map_or(vec![], |v| vec![v as u8]), + truncated_rows: options.truncated_rows.map_or(vec![], |v| vec![v as u8]), + compression_level: options.compression_level, + quote_style: options.quote_style as i32, + ignore_leading_whitespace: options + .ignore_leading_whitespace + .map_or(vec![], |v| vec![v as u8]), + ignore_trailing_whitespace: options + .ignore_trailing_whitespace + .map_or(vec![], |v| vec![v as u8]), + } + } else { + datafusion_proto_models::protobuf::CsvOptions::default() + } + } } #[cfg(test)] diff --git a/datafusion/datasource-csv/src/source.rs b/datafusion/datasource-csv/src/source.rs index 611586cee6473..08e4607498e62 100644 --- a/datafusion/datasource-csv/src/source.rs +++ b/datafusion/datasource-csv/src/source.rs @@ -17,25 +17,24 @@ //! Execution plan for reading CSV files +use datafusion_datasource::boundary_stream::AlignedBoundaryStream; use datafusion_datasource::projection::{ProjectionOpener, SplitProjection}; use datafusion_physical_plan::projection::ProjectionExprs; use std::fmt; -use std::io::{Read, Seek, SeekFrom}; +use std::io::Read; use std::sync::Arc; -use std::task::Poll; use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream}; use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener}; use datafusion_datasource::{ - FileRange, ListingTableUrl, PartitionedFile, RangeCalculation, TableSchema, - as_file_source, calculate_range, + FileRange, ListingTableUrl, PartitionedFile, TableSchema, as_file_source, }; use arrow::csv; use datafusion_common::config::CsvOptions; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{DataFusionError, Result}; +use datafusion_common::{DataFusionError, Result, exec_datafusion_err}; use datafusion_common_runtime::JoinSet; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; @@ -314,15 +313,53 @@ impl FileSource for CsvSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit projection expressions - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in &self.projection.source { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f) + } + + /// Emit a `CsvScan` node wrapping the shared base config and CSV options. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let node = protobuf::CsvScanExecNode { + base_conf: Some(base.try_to_proto(ctx)?), + has_header: self.has_header(), + delimiter: proto_byte_to_string(self.delimiter(), "delimiter")?, + quote: proto_byte_to_string(self.quote(), "quote")?, + optional_escape: self + .escape() + .map(|escape| { + Ok::<_, DataFusionError>( + protobuf::csv_scan_exec_node::OptionalEscape::Escape( + proto_byte_to_string(escape, "escape")?, + ), + ) + }) + .transpose()?, + optional_comment: self + .comment() + .map(|comment| { + Ok::<_, DataFusionError>( + protobuf::csv_scan_exec_node::OptionalComment::Comment( + proto_byte_to_string(comment, "comment")?, + ), + ) + }) + .transpose()?, + newlines_in_values: self.newlines_in_values(), + truncate_rows: self.truncate_rows(), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::CsvScan(node)), + })) } } @@ -382,43 +419,53 @@ impl FileOpener for CsvOpener { Ok(Box::pin(async move { // Current partition contains bytes [start_byte, end_byte) (might contain incomplete lines at boundaries) + let file_size = partitioned_file.object_meta.size; + let location = partitioned_file.object_meta.location; + + if let Some(file_range) = partitioned_file.range.as_ref() { + let raw_start: u64 = file_range.start.try_into().map_err(|_| { + exec_datafusion_err!( + "Expected start range to fit in u64, got {}", + file_range.start + ) + })?; + let raw_end: u64 = file_range.end.try_into().map_err(|_| { + exec_datafusion_err!( + "Expected end range to fit in u64, got {}", + file_range.end + ) + })?; + + let aligned_stream = AlignedBoundaryStream::new( + Arc::clone(&store), + location.clone(), + raw_start, + raw_end, + file_size, + terminator.unwrap_or(b'\n'), + ) + .await? + .map_err(DataFusionError::from); + + let decoder = config.builder().build_decoder(); + let input = file_compression_type + .convert_stream(aligned_stream.boxed())? + .fuse(); + let stream = deserialize_stream( + input, + DecoderDeserializer::new(CsvDecoder::new(decoder)), + ); + return Ok(stream.map_err(Into::into).boxed()); + } - let calculated_range = - calculate_range(&partitioned_file, &store, terminator).await?; - - let range = match calculated_range { - RangeCalculation::Range(None) => None, - RangeCalculation::Range(Some(range)) => Some(range.into()), - RangeCalculation::TerminateEarly => { - return Ok( - futures::stream::poll_fn(move |_| Poll::Ready(None)).boxed() - ); - } - }; - - let options = GetOptions { - range, - ..Default::default() - }; - - let result = store - .get_opts(&partitioned_file.object_meta.location, options) - .await?; + // No range specified — read the entire file + let options = GetOptions::default(); + let result = store.get_opts(&location, options).await?; match result.payload { #[cfg(not(target_arch = "wasm32"))] - GetResultPayload::File(mut file, _) => { - let is_whole_file_scanned = partitioned_file.range.is_none(); - let decoder = if is_whole_file_scanned { - // Don't seek if no range as breaks FIFO files - file_compression_type.convert_read(file)? - } else { - file.seek(SeekFrom::Start(result.range.start as _))?; - file_compression_type.convert_read( - file.take((result.range.end - result.range.start) as u64), - )? - }; - + GetResultPayload::File(file, _) => { + let decoder = file_compression_type.convert_read(file)?; let mut reader = config.open(decoder)?; // Use std::iter::from_fn to wrap execution of iterator's next() method. @@ -507,3 +554,97 @@ pub async fn plan_to_csv( Ok(()) } + +#[cfg(feature = "proto")] +fn proto_byte_to_string(b: u8, description: &str) -> Result { + let bytes = &[b]; + let s = std::str::from_utf8(bytes).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Invalid CSV {description}: can not represent {bytes:0x?} as utf8" + ) + })?; + Ok(s.to_owned()) +} + +#[cfg(feature = "proto")] +fn proto_str_to_byte(s: &str, description: &str) -> Result { + datafusion_common::assert_eq_or_internal_err!( + s.len(), + 1, + "Invalid CSV {description}: expected single character, got {s}" + ); + Ok(s.as_bytes()[0]) +} + +#[cfg(feature = "proto")] +impl CsvSource { + /// Reconstructs a `DataSourceExec` from a protobuf `CsvScan`. + /// + /// Custom line terminators are not represented in the wire format. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::config::CsvOptions; + use datafusion_datasource::file_compression_type::FileCompressionType; + use datafusion_datasource::file_scan_config::{ + FileScanConfig, FileScanConfigBuilder, + }; + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::CsvScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a CsvScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "CsvScanExecNode is missing required field 'base_conf'" + ) + })?; + + let escape = match &scan.optional_escape { + Some(protobuf::csv_scan_exec_node::OptionalEscape::Escape(escape)) => { + Some(proto_str_to_byte(escape, "escape")?) + } + None => None, + }; + let comment = match &scan.optional_comment { + Some(protobuf::csv_scan_exec_node::OptionalComment::Comment(comment)) => { + Some(proto_str_to_byte(comment, "comment")?) + } + None => None, + }; + + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + + let csv_options = CsvOptions { + has_header: Some(scan.has_header), + delimiter: proto_str_to_byte(&scan.delimiter, "delimiter")?, + quote: proto_str_to_byte(&scan.quote, "quote")?, + newlines_in_values: Some(scan.newlines_in_values), + truncated_rows: Some(scan.truncate_rows), + ..Default::default() + }; + let source = Arc::new( + CsvSource::new(table_schema) + .with_csv_options(csv_options) + .with_escape(escape) + .with_comment(comment), + ); + + // The compression type is not on the wire; CSV scans always + // deserialize as uncompressed. + let conf = FileScanConfigBuilder::from(FileScanConfig::try_from_proto( + base_conf, ctx, source, + )?) + .with_file_compression_type(FileCompressionType::UNCOMPRESSED) + .build(); + Ok(DataSourceExec::from_data_source(conf)) + } +} diff --git a/datafusion/datasource-json/Cargo.toml b/datafusion/datasource-json/Cargo.toml index b5947ea5c4c67..04192083f583a 100644 --- a/datafusion/datasource-json/Cargo.toml +++ b/datafusion/datasource-json/Cargo.toml @@ -30,6 +30,14 @@ version.workspace = true [package.metadata.docs.rs] all-features = true +[features] +# Enables protobuf serialization hooks for JSON sources and sinks. +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] + [dependencies] arrow = { workspace = true } async-trait = { workspace = true } @@ -41,6 +49,7 @@ datafusion-execution = { workspace = true } datafusion-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } futures = { workspace = true } object_store = { workspace = true } diff --git a/datafusion/datasource-json/src/file_format.rs b/datafusion/datasource-json/src/file_format.rs index 1854fddfb84b3..62d03d67ccd43 100644 --- a/datafusion/datasource-json/src/file_format.rs +++ b/datafusion/datasource-json/src/file_format.rs @@ -429,7 +429,7 @@ impl DisplayAs for JsonSink { } DisplayFormatType::TreeRender => { writeln!(f, "format: json")?; - write!(f, "file={}", &self.config.original_url) + write!(f, "file={}", self.config.original_url) } } } @@ -490,6 +490,105 @@ impl DataSink for JsonSink { ) -> Result { FileSink::write_all(self, data, context).await } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let input = ctx.encode_child(exec.input())?; + let sort_order = exec.encode_sort_order(ctx)?; + let sink = protobuf::JsonSink::try_from(self)?; + let node = protobuf::JsonSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(exec.schema().as_ref().try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::JsonSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl TryFrom<&JsonSink> for datafusion_proto_models::protobuf::JsonSink { + type Error = datafusion_common::DataFusionError; + + fn try_from(value: &JsonSink) -> Result { + Ok(Self { + config: Some(value.config().try_into()?), + writer_options: Some(value.writer_options().try_into()?), + }) + } +} + +#[cfg(feature = "proto")] +impl TryFrom<&datafusion_proto_models::protobuf::JsonSink> for JsonSink { + type Error = datafusion_common::DataFusionError; + + fn try_from(value: &datafusion_proto_models::protobuf::JsonSink) -> Result { + let config = + FileSinkConfig::try_from(value.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSink is missing required field 'config'" + ) + })?)?; + let writer_options = value + .writer_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSink is missing required field 'writer_options'" + ) + })? + .try_into()?; + + Ok(Self::new(config, writer_options)) + } +} + +#[cfg(feature = "proto")] +impl JsonSink { + /// Reconstructs a [`DataSinkExec`] containing a `JsonSink` from protobuf. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let sink_node = datafusion_physical_plan::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::JsonSink, + "JsonSink", + ); + let input = ctx.decode_required_child( + sink_node.input.as_deref(), + "JsonSinkExecNode", + "input", + )?; + let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonSinkExecNode is missing required field 'sink'" + ) + })?; + let data_sink = JsonSink::try_from(proto_sink)?; + let sort_order = DataSinkExec::decode_sort_order( + sink_node.sort_order.as_ref(), + ctx, + input.schema().as_ref(), + )?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } } #[derive(Debug)] @@ -516,3 +615,24 @@ impl Decoder for JsonDecoder { false } } + +/// Encode a [`JsonFormatFactory`]'s options as their protobuf form. +/// +/// The reverse direction is `From<&protobuf::JsonOptions> for JsonOptions` in +/// `datafusion-proto-models`: `JsonOptions` is a `datafusion-common` type, so +/// that half cannot live here. +#[cfg(feature = "proto")] +impl From<&JsonFormatFactory> for datafusion_proto_models::protobuf::JsonOptions { + fn from(factory: &JsonFormatFactory) -> Self { + if let Some(options) = &factory.options { + datafusion_proto_models::protobuf::JsonOptions { + compression: options.compression as i32, + schema_infer_max_rec: options.schema_infer_max_rec.map(|v| v as u64), + compression_level: options.compression_level, + newline_delimited: Some(options.newline_delimited), + } + } else { + datafusion_proto_models::protobuf::JsonOptions::default() + } + } +} diff --git a/datafusion/datasource-json/src/mod.rs b/datafusion/datasource-json/src/mod.rs index f7932c8a21d95..ec93fc9c387e2 100644 --- a/datafusion/datasource-json/src/mod.rs +++ b/datafusion/datasource-json/src/mod.rs @@ -20,7 +20,6 @@ // https://github.com/apache/datafusion/issues/11143 #![cfg_attr(not(test), deny(clippy::clone_on_ref_ptr))] -pub mod boundary_stream; pub mod file_format; pub mod source; pub mod utils; diff --git a/datafusion/datasource-json/src/source.rs b/datafusion/datasource-json/src/source.rs index 2f2f459956f4e..47241c9d99ab5 100644 --- a/datafusion/datasource-json/src/source.rs +++ b/datafusion/datasource-json/src/source.rs @@ -25,12 +25,11 @@ use std::task::{Context, Poll}; use crate::file_format::JsonDecoder; use crate::utils::{ChannelReader, JsonArrayToNdjsonReader}; -use crate::boundary_stream::AlignedBoundaryStream; - use datafusion_common::error::{DataFusionError, Result}; use datafusion_common::exec_datafusion_err; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common_runtime::{JoinSet, SpawnedTask}; +use datafusion_datasource::boundary_stream::AlignedBoundaryStream; use datafusion_datasource::decoder::{DecoderDeserializer, deserialize_stream}; use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_stream::{FileOpenFuture, FileOpener}; @@ -237,15 +236,64 @@ impl FileSource for JsonSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit projection expressions - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in &self.projection.source { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots(self.projection.source.iter(), f) + } + + /// Emit a `JsonScan` node wrapping the shared base config. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let node = protobuf::JsonScanExecNode { + base_conf: Some(base.try_to_proto(ctx)?), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::JsonScan(node)), + })) + } +} + +#[cfg(feature = "proto")] +impl JsonSource { + /// Reconstructs a `DataSourceExec` from a protobuf `JsonScan`. + /// + /// Defaults to newline-delimited JSON because protobuf does not encode the mode. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_datasource::file_scan_config::FileScanConfig; + use datafusion_datasource::source::DataSourceExec; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::JsonScan(scan)) => scan, + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a JsonScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "JsonScanExecNode is missing required field 'base_conf'" + ) + })?; + + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + let source = Arc::new(JsonSource::new(table_schema)); + + let conf = FileScanConfig::try_from_proto(base_conf, ctx, source)?; + Ok(DataSourceExec::from_data_source(conf)) } } diff --git a/datafusion/datasource-parquet/Cargo.toml b/datafusion/datasource-parquet/Cargo.toml index a5855af17a536..a2589af19a6ee 100644 --- a/datafusion/datasource-parquet/Cargo.toml +++ b/datafusion/datasource-parquet/Cargo.toml @@ -32,6 +32,7 @@ all-features = true [dependencies] arrow = { workspace = true } +arrow-schema = { workspace = true } async-trait = { workspace = true } bytes = { workspace = true } datafusion-common = { workspace = true, features = ["object_store", "parquet"] } @@ -45,6 +46,7 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-pruning = { workspace = true } datafusion-session = { workspace = true } futures = { workspace = true } @@ -73,6 +75,11 @@ name = "datafusion_datasource_parquet" path = "src/mod.rs" [features] +proto = [ + "dep:datafusion-proto-models", + "datafusion-datasource/proto", + "datafusion-physical-plan/proto", +] parquet_encryption = [ "parquet/encryption", "datafusion-common/parquet_encryption", @@ -86,3 +93,7 @@ harness = false [[bench]] name = "parquet_struct_filter_pushdown" harness = false + +[[bench]] +name = "parquet_metadata_statistics" +harness = false diff --git a/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs b/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs new file mode 100644 index 0000000000000..46ebd100fde88 --- /dev/null +++ b/datafusion/datasource-parquet/benches/parquet_metadata_statistics.rs @@ -0,0 +1,303 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for deriving DataFusion table statistics from Parquet metadata. +//! +//! This mirrors the structure of Arrow's `arrow_statistics` benchmark: build +//! Parquet metadata once, then repeatedly measure statistics extraction. The +//! benchmark targets the cold planning/statistics path used by listing tables. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BatchSize, BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_datasource_parquet::metadata::DFParquetMetadata; +use parquet::arrow::ArrowSchemaConverter; +use parquet::data_type::ByteArray; +use parquet::file::metadata::{ + ColumnChunkMetaData, FileMetaData, ParquetMetaData, RowGroupMetaData, +}; +use parquet::file::statistics::{Statistics as ParquetStatistics, ValueStatistics}; + +const ROWS_PER_GROUP: usize = 8; + +#[derive(Debug, Copy, Clone)] +struct BenchmarkSpec { + columns: usize, + row_groups: usize, + metadata: MetadataState, +} + +#[derive(Debug, Copy, Clone)] +enum MetadataState { + Full, + Mixed, + None, +} + +impl std::fmt::Display for MetadataState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::Full => write!(f, "full"), + Self::Mixed => write!(f, "mixed"), + Self::None => write!(f, "none"), + } + } +} + +struct BenchmarkCase { + schema: SchemaRef, + metadata: ParquetMetaData, +} + +fn parquet_metadata_statistics(c: &mut Criterion) { + let metadata_states = [ + MetadataState::Full, + MetadataState::Mixed, + MetadataState::None, + ]; + let column_counts = [8, 64, 256]; + let row_group_counts = [1, 32, 128]; + + let mut group = c.benchmark_group("parquet_metadata_statistics"); + + for metadata in metadata_states { + for columns in column_counts { + for row_groups in row_group_counts { + let spec = BenchmarkSpec { + columns, + row_groups, + metadata, + }; + group.bench_function( + BenchmarkId::from_parameter(format!( + "metadata_{}_col_{}_rg_{}", + spec.metadata, spec.columns, spec.row_groups, + )), + |b| { + b.iter_batched( + || BenchmarkCase::new(spec), + |case| { + let statistics = + DFParquetMetadata::statistics_from_parquet_metadata( + black_box(&case.metadata), + black_box(&case.schema), + ) + .expect("statistics extraction failed"); + black_box(statistics); + }, + BatchSize::PerIteration, + ); + }, + ); + } + } + } + + group.finish(); +} + +impl BenchmarkCase { + fn new(spec: BenchmarkSpec) -> Self { + let schema = make_schema(spec.columns); + let metadata = match spec.metadata { + MetadataState::Full => { + make_synthetic_metadata(&schema, spec, full_statistics) + } + MetadataState::Mixed => { + make_synthetic_metadata(&schema, spec, mixed_statistics) + } + MetadataState::None => make_synthetic_metadata(&schema, spec, |_, _, _| None), + }; + + Self { schema, metadata } + } +} + +fn make_synthetic_metadata( + schema: &SchemaRef, + spec: BenchmarkSpec, + statistics: fn(&DataType, usize, usize) -> Option, +) -> ParquetMetaData { + let schema_descr = Arc::new( + ArrowSchemaConverter::new() + .convert(schema.as_ref()) + .expect("failed to convert arrow schema"), + ); + let row_groups = (0..spec.row_groups) + .map(|row_group| { + let columns = schema + .fields() + .iter() + .enumerate() + .map(|(column_idx, field)| { + let mut builder = + ColumnChunkMetaData::builder(schema_descr.column(column_idx)); + if let Some(statistics) = + statistics(field.data_type(), column_idx, row_group) + { + builder = builder.set_statistics(statistics); + } + builder + .set_num_values(ROWS_PER_GROUP as i64) + .build() + .expect("failed to build column metadata") + }) + .collect::>(); + + RowGroupMetaData::builder(Arc::clone(&schema_descr)) + .set_num_rows(ROWS_PER_GROUP as i64) + .set_total_byte_size((spec.columns * ROWS_PER_GROUP * 8) as i64) + .set_column_metadata(columns) + .build() + .expect("failed to build row group metadata") + }) + .collect::>(); + + let file_metadata = FileMetaData::new( + 1, + (spec.row_groups * ROWS_PER_GROUP) as i64, + Some("datafusion parquet metadata benchmark".to_string()), + None, + schema_descr, + None, + ); + + ParquetMetaData::new(file_metadata, row_groups) +} + +fn full_statistics( + data_type: &DataType, + column_idx: usize, + row_group: usize, +) -> Option { + Some(statistics( + data_type, + column_idx, + row_group, + true, + true, + Some(null_count_for_rows()), + )) +} + +fn mixed_statistics( + data_type: &DataType, + column_idx: usize, + row_group: usize, +) -> Option { + if column_idx.is_multiple_of(16) || row_group.is_multiple_of(5) { + return None; + } + + let min_exact = !row_group.is_multiple_of(3); + let max_exact = !row_group.is_multiple_of(4); + let null_count = (!row_group.is_multiple_of(7)).then(null_count_for_rows); + + Some(statistics( + data_type, column_idx, row_group, min_exact, max_exact, null_count, + )) +} + +fn statistics( + data_type: &DataType, + column_idx: usize, + row_group: usize, + min_exact: bool, + max_exact: bool, + null_count: Option, +) -> ParquetStatistics { + let min_row = first_non_null_row(); + let max_row = last_non_null_row(); + + match data_type { + DataType::Int64 => { + let min = min_row.map(|row| value(column_idx, row_group, row)); + let max = max_row.map(|row| value(column_idx, row_group, row)); + ParquetStatistics::Int64( + ValueStatistics::new(min, max, None, null_count, false) + .with_min_is_exact(min_exact) + .with_max_is_exact(max_exact), + ) + } + DataType::Float64 => { + let min = min_row.map(|row| value(column_idx, row_group, row) as f64 * 1.5); + let max = max_row.map(|row| value(column_idx, row_group, row) as f64 * 1.5); + ParquetStatistics::Double( + ValueStatistics::new(min, max, None, null_count, false) + .with_min_is_exact(min_exact) + .with_max_is_exact(max_exact), + ) + } + DataType::Utf8 => { + let min = min_row.map(|row| { + ByteArray::from(string_value(column_idx, row_group, row).into_bytes()) + }); + let max = max_row.map(|row| { + ByteArray::from(string_value(column_idx, row_group, row).into_bytes()) + }); + ParquetStatistics::ByteArray( + ValueStatistics::new(min, max, None, null_count, false) + .with_min_is_exact(min_exact) + .with_max_is_exact(max_exact), + ) + } + other => unreachable!("unsupported benchmark data type: {other:?}"), + } +} + +fn make_schema(columns: usize) -> SchemaRef { + let fields = (0..columns) + .map(|idx| { + let data_type = match idx % 4 { + 0 => DataType::Int64, + 1 => DataType::Float64, + 2 => DataType::Utf8, + _ => DataType::Int64, + }; + Field::new(format!("c{idx:04}"), data_type, true) + }) + .collect::>(); + + Arc::new(Schema::new(fields)) +} + +fn first_non_null_row() -> Option { + (0..ROWS_PER_GROUP).find(|row| !row.is_multiple_of(7)) +} + +fn last_non_null_row() -> Option { + (0..ROWS_PER_GROUP).rev().find(|row| !row.is_multiple_of(7)) +} + +fn null_count_for_rows() -> u64 { + (0..ROWS_PER_GROUP) + .filter(|row| row.is_multiple_of(7)) + .count() as u64 +} + +fn value(column_idx: usize, row_group: usize, row: usize) -> i64 { + (column_idx as i64 * 10_000) + (row_group as i64 * 100) + row as i64 +} + +fn string_value(column_idx: usize, row_group: usize, row: usize) -> String { + format!("s{column_idx:04}_{row_group:04}_{row:04}") +} + +criterion_group!(benches, parquet_metadata_statistics); +criterion_main!(benches); diff --git a/datafusion/datasource-parquet/src/access_plan.rs b/datafusion/datasource-parquet/src/access_plan.rs index edbea39948f09..1e9bae0ff6ba3 100644 --- a/datafusion/datasource-parquet/src/access_plan.rs +++ b/datafusion/datasource-parquet/src/access_plan.rs @@ -17,7 +17,7 @@ use crate::sort::reverse_row_selection; use arrow::datatypes::Schema; -use datafusion_common::{Result, assert_eq_or_internal_err}; +use datafusion_common::{Result, assert_eq_or_internal_err, exec_err}; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr_common::sort_expr::LexOrdering; use log::debug; @@ -104,6 +104,41 @@ pub struct ParquetAccessPlan { fully_matched: Vec, } +/// A file-level row selection for a parquet scan. +/// +/// Attach this type to a [`PartitionedFile`](datafusion_datasource::PartitionedFile) +/// with [`PartitionedFile::with_extension`](datafusion_datasource::PartitionedFile::with_extension) +/// when an external index produces a [`RowSelection`] across the entire parquet +/// file. DataFusion will use parquet metadata to split it into row-group-level +/// access when the file is opened. +#[derive(Debug, Clone, PartialEq)] +pub struct ParquetRowSelection { + selection: RowSelection, +} + +impl ParquetRowSelection { + /// Create a new file-level parquet row selection. + pub fn new(selection: RowSelection) -> Self { + Self { selection } + } + + /// Return a reference to the underlying [`RowSelection`]. + pub fn selection(&self) -> &RowSelection { + &self.selection + } + + /// Convert into the underlying [`RowSelection`]. + pub fn into_inner(self) -> RowSelection { + self.selection + } +} + +impl From for ParquetRowSelection { + fn from(selection: RowSelection) -> Self { + Self::new(selection) + } +} + /// Describes how the parquet reader will access a row group #[derive(Debug, Clone, PartialEq)] pub enum RowGroupAccess { @@ -115,30 +150,108 @@ pub enum RowGroupAccess { Selection(RowSelection), } -/// A consecutive set of row groups that share the same row filter requirement. -#[derive(Debug, Clone, PartialEq)] -pub(crate) struct RowGroupRun { - /// True if this run needs row filter evaluation. - pub(crate) needs_filter: bool, - /// The access plan for this run. - pub(crate) access_plan: ParquetAccessPlan, +impl RowGroupAccess { + /// Return true if this row group should be scanned + pub fn should_scan(&self) -> bool { + match self { + RowGroupAccess::Skip => false, + RowGroupAccess::Scan | RowGroupAccess::Selection(_) => true, + } + } } -impl RowGroupRun { - fn new(needs_filter: bool, access_plan: ParquetAccessPlan) -> Self { +/// Single-pass cursor over a file-level [`RowSelection`]. +/// +/// `take` returns the next selector fragment capped to the requested row count, +/// splitting the current selector when it straddles a row group boundary. +struct OverallRowSelectionCursor { + selector_iter: std::vec::IntoIter, + current: Option, +} + +impl OverallRowSelectionCursor { + fn new(selection: RowSelection) -> Self { + let selectors: Vec = selection.into(); + let mut selector_iter = selectors.into_iter(); + let current = selector_iter.next(); Self { - needs_filter, - access_plan, + selector_iter, + current, } } + + /// Take up to `max_rows` rows from the current selector. + /// + /// If the current selector crosses the requested boundary, this returns the + /// leading fragment and keeps the remaining rows in `self.current` for the + /// next call. + #[inline] + fn take(&mut self, max_rows: usize) -> Option { + let sel = self.current?; + let row_count = sel.row_count.min(max_rows); + self.current = if row_count < sel.row_count { + Some(RowSelector { + row_count: sel.row_count - row_count, + skip: sel.skip, + }) + } else { + self.selector_iter.next() + }; + + Some(RowSelector { + row_count, + skip: sel.skip, + }) + } + + fn remaining_rows(self) -> usize { + self.current.map_or(0, |s| s.row_count) + + self.selector_iter.map(|s| s.row_count).sum::() + } } -impl RowGroupAccess { - /// Return true if this row group should be scanned - pub fn should_scan(&self) -> bool { - match self { - RowGroupAccess::Skip => false, - RowGroupAccess::Scan | RowGroupAccess::Selection(_) => true, +/// Accumulates the selector fragments that belong to one row group. +struct RowGroupAccessBuilder { + /// Selector fragments belonging to this row group. + selectors: Vec, + /// Number of selected rows accumulated for this row group. + selected: usize, + /// Number of skipped rows accumulated for this row group. + skipped: usize, + /// Number of rows still needed to complete this row group. + remaining: usize, +} + +impl RowGroupAccessBuilder { + fn new(row_group_rows: usize) -> Self { + Self { + selectors: Vec::with_capacity(1), + selected: 0, + skipped: 0, + remaining: row_group_rows, + } + } + + #[inline] + fn push(&mut self, selector: RowSelector) { + self.remaining -= selector.row_count; + + if selector.skip { + self.skipped += selector.row_count; + } else { + self.selected += selector.row_count; + } + + self.selectors.push(selector); + } + + fn into_access(self) -> RowGroupAccess { + if self.selected == 0 { + RowGroupAccess::Skip + } else if self.skipped == 0 { + RowGroupAccess::Scan + } else { + RowGroupAccess::Selection(self.selectors.into()) } } } @@ -169,6 +282,60 @@ impl ParquetAccessPlan { } } + /// Create a new `ParquetAccessPlan` from a file-level [`RowSelection`]. + /// + /// The selection is interpreted across all rows in the file, in row group + /// order, and is split into row-group level access using `row_group_meta_data`. + /// Fully skipped row groups become [`RowGroupAccess::Skip`], fully selected + /// row groups become [`RowGroupAccess::Scan`], and partially selected row + /// groups become [`RowGroupAccess::Selection`]. + /// + /// # Errors + /// + /// Returns an error if the selection does not specify exactly the same + /// number of rows as the file metadata. + pub fn try_new_from_overall_row_selection( + selection: RowSelection, + row_group_meta_data: &[RowGroupMetaData], + ) -> Result { + // Keep this as a single pass over the selector stream rather than + // repeatedly calling `RowSelection::split_off` per row group. The + // `split_off` version is simpler, but it clones/retains substantially + // more selector buffer capacity for highly fragmented selections. + let mut cursor = OverallRowSelectionCursor::new(selection); + + let mut selection_rows = 0usize; + let mut file_rows = 0usize; + + let mut row_groups = Vec::with_capacity(row_group_meta_data.len()); + for rg_meta in row_group_meta_data { + let rg_rows = rg_meta.num_rows() as usize; + file_rows += rg_rows; + + let mut builder = RowGroupAccessBuilder::new(rg_rows); + while builder.remaining > 0 { + let Some(selector) = cursor.take(builder.remaining) else { + break; + }; + selection_rows += selector.row_count; + builder.push(selector); + } + + row_groups.push(builder.into_access()); + } + + selection_rows += cursor.remaining_rows(); + + if selection_rows != file_rows { + return exec_err!( + "Invalid Parquet RowSelection. File has {file_rows} rows, \ + but selection specifies {selection_rows} rows." + ); + } + + Ok(Self::new(row_groups)) + } + /// Set the i-th row group to the specified [`RowGroupAccess`] pub fn set(&mut self, idx: usize, access: RowGroupAccess) { let should_scan = access.should_scan(); @@ -213,12 +380,6 @@ impl ParquetAccessPlan { &self.fully_matched } - /// Return true if any scanned row group is fully matched. - fn has_fully_matched(&self) -> bool { - self.row_group_index_iter() - .any(|idx| self.is_fully_matched(idx)) - } - /// Set to scan only the [`RowSelection`] in the specified row group. /// /// Behavior is different depending on the existing access @@ -404,54 +565,6 @@ impl ParquetAccessPlan { self.row_groups } - /// Split this plan into consecutive row group runs that share the same row - /// filter requirement. - pub(crate) fn split_runs(self, needs_filter: bool) -> Vec { - if !needs_filter || !self.has_fully_matched() { - return vec![RowGroupRun::new(needs_filter, self)]; - } - - let num_row_groups = self.row_groups.len(); - let row_groups = self.row_groups; - let fully_matched = self.fully_matched; - let mut runs: Vec = Vec::new(); - - for (idx, (access, fully_matched)) in - row_groups.into_iter().zip(fully_matched).enumerate() - { - if !access.should_scan() { - continue; - } - - let row_group_needs_filter = !fully_matched; - if let Some(run) = runs - .last_mut() - .filter(|run| run.needs_filter == row_group_needs_filter) - { - run.access_plan.set(idx, access); - if fully_matched { - run.access_plan.mark_fully_matched(idx); - } - } else { - let mut run_plan = ParquetAccessPlan::new_none(num_row_groups); - run_plan.set(idx, access); - if fully_matched { - run_plan.mark_fully_matched(idx); - } - runs.push(RowGroupRun::new(row_group_needs_filter, run_plan)); - } - } - - if runs.is_empty() { - vec![RowGroupRun::new( - needs_filter, - ParquetAccessPlan::new_none(num_row_groups), - )] - } else { - runs - } - } - /// Prepare this plan and resolve to the final `PreparedAccessPlan` pub(crate) fn prepare( self, @@ -493,13 +606,24 @@ impl PreparedAccessPlan { /// Reorder row groups by their min statistics for the given sort order. /// /// This helps TopK queries find optimal values first. Row groups are - /// always sorted by min values in ASC order — direction (DESC) is - /// handled separately by `reverse()` which is applied after reorder. + /// lexicographically sorted by per-column min values over the longest + /// prefix of the sort order made of plain columns present in the file + /// schema. The leading column is always sorted ASC by min — direction + /// (DESC) is handled separately by `reverse()` which is applied after + /// reorder. Subsequent columns sort by their direction *relative* to + /// the leading column (and their null placement is flipped when the + /// plan will be reversed), so that the post-`reverse()` order + /// approximates the requested lexicographic order. + /// + /// Secondary sort keys matter when the leading column's min ties + /// across row groups (e.g. `ORDER BY low_cardinality_col, ts LIMIT k`) + /// — without them the reorder is a no-op on such files and the TopK + /// dynamic filter converges only as fast as disk order allows. /// /// Gracefully skips reordering when: /// - There is a row_selection (too complex to remap) /// - 0 or 1 row groups (nothing to reorder) - /// - Sort expression is not a simple column reference + /// - The leading sort expression is not a simple column reference /// - Statistics are unavailable pub(crate) fn reorder_by_statistics( mut self, @@ -518,88 +642,116 @@ impl PreparedAccessPlan { return Ok(self); } - let first_sort_expr = sort_order.first(); - - // Extract column name from sort expression - let column: &Column = match first_sort_expr.expr.downcast_ref::() { - Some(col) => col, - None => { - debug!("Skipping RG reorder: sort expr is not a simple column"); - return Ok(self); - } - }; - - // Expected graceful skip: the sort column lives outside the - // file schema (e.g. a partition column whose ordering came - // through `reversed_satisfies` rather than `column_in_file_schema`). - // Parquet has no per-RG stats for it. Bail out quietly — no - // `debug_assert!` because this is a normal pushdown shape. - if arrow_schema.field_with_name(column.name()).is_err() { - debug!( - "Skipping RG reorder: column `{}` not in file schema", - column.name() - ); - return Ok(self); - } - - // From here, any `StatisticsConverter` / stats read / sort - // failure is unexpected — the column exists in the file - // schema, so building the converter and pulling typed mins - // should succeed on any well-formed parquet file. Trip a - // `debug_assert!` so CI catches regressions, but stay graceful - // in release so a single odd file can't take down a scan. - let converter = match StatisticsConverter::try_new( - column.name(), - arrow_schema, - file_metadata.file_metadata().schema_descr(), - ) { - Ok(c) => c, - Err(e) => { - debug_assert!( - false, - "RG reorder: cannot create stats converter for `{}`: {e}", - column.name(), - ); - return Ok(self); - } - }; - - // Always sort ASC by min values — direction is handled by reverse let rg_metadata: Vec<&RowGroupMetaData> = self .row_group_indexes .iter() .map(|&idx| file_metadata.row_group(idx)) .collect(); - let stat_mins = match converter.row_group_mins(rg_metadata.iter().copied()) { - Ok(vals) => vals, - Err(e) => { - debug_assert!( - false, - "RG reorder: cannot get min values for `{}`: {e}", - column.name(), - ); - return Ok(self); + let leading_descending = sort_order.first().options.descending; + + // Build one `SortColumn` of per-RG mins for each usable prefix + // column of the sort order. The walk stops at the first + // expression that isn't a plain `Column` in the file schema — + // stats for later columns can't refine the order once an + // unresolvable key sits between them and the resolved prefix. + let mut sort_columns: Vec = Vec::new(); + for (i, sort_expr) in sort_order.iter().enumerate() { + let column: &Column = match sort_expr.expr.downcast_ref::() { + Some(col) => col, + None => { + if i == 0 { + debug!("Skipping RG reorder: sort expr is not a simple column"); + return Ok(self); + } + break; + } + }; + + // Expected graceful skip: the sort column lives outside the + // file schema (e.g. a partition column whose ordering came + // through `reversed_satisfies` rather than + // `column_in_file_schema`). Parquet has no per-RG stats for + // it. Bail out quietly — no `debug_assert!` because this is + // a normal pushdown shape. + if arrow_schema.field_with_name(column.name()).is_err() { + if i == 0 { + debug!( + "Skipping RG reorder: column `{}` not in file schema", + column.name() + ); + return Ok(self); + } + break; } - }; - let sort_options = arrow::compute::SortOptions { - descending: false, - nulls_first: first_sort_expr.options.nulls_first, - }; - let sorted_indices = - match arrow::compute::sort_to_indices(&stat_mins, Some(sort_options), None) { - Ok(indices) => indices, + // From here, any `StatisticsConverter` / stats read / sort + // failure is unexpected — the column exists in the file + // schema, so building the converter and pulling typed mins + // should succeed on any well-formed parquet file. Trip a + // `debug_assert!` so CI catches regressions, but stay graceful + // in release so a single odd file can't take down a scan. + let converter = match StatisticsConverter::try_new( + column.name(), + arrow_schema, + file_metadata.file_metadata().schema_descr(), + ) { + Ok(c) => c, Err(e) => { debug_assert!( false, - "RG reorder: arrow sort_to_indices failed for `{}`: {e}", + "RG reorder: cannot create stats converter for `{}`: {e}", column.name(), ); - return Ok(self); + if i == 0 { + return Ok(self); + } + break; } }; + let stat_mins = match converter.row_group_mins(rg_metadata.iter().copied()) { + Ok(vals) => vals, + Err(e) => { + debug_assert!( + false, + "RG reorder: cannot get min values for `{}`: {e}", + column.name(), + ); + if i == 0 { + return Ok(self); + } + break; + } + }; + + // The plan is later `reverse()`d iff the leading column is + // DESC, which flips both value order and null placement of + // every column. Sort each column by its direction relative + // to the leading column (leading itself is therefore always + // ASC), and pre-flip null placement when the reverse is + // coming, so the post-reverse order matches the request. + // Nulls here are row groups with *missing stats*, so their + // placement is a heuristic, not a correctness matter. + let sort_options = arrow::compute::SortOptions { + descending: sort_expr.options.descending != leading_descending, + nulls_first: sort_expr.options.nulls_first != leading_descending, + }; + sort_columns.push(arrow::compute::SortColumn { + values: stat_mins, + options: Some(sort_options), + }); + } + + let sorted_indices = match arrow::compute::lexsort_to_indices(&sort_columns, None) + { + Ok(indices) => indices, + Err(e) => { + debug_assert!(false, "RG reorder: arrow lexsort_to_indices failed: {e}"); + return Ok(self); + } + }; + // Apply the reordering let original_indexes = self.row_group_indexes.clone(); self.row_group_indexes = sorted_indices @@ -758,6 +910,90 @@ mod test { ); } + #[test] + fn test_new_from_overall_row_selection() { + let row_selection = RowSelection::from(vec![ + RowSelector::select(10), + RowSelector::skip(25), + RowSelector::select(10), + RowSelector::skip(15), + RowSelector::select(40), + ]); + + let access_plan = ParquetAccessPlan::try_new_from_overall_row_selection( + row_selection, + &ROW_GROUP_METADATA, + ) + .unwrap(); + + assert_eq!( + access_plan, + ParquetAccessPlan::new(vec![ + RowGroupAccess::Scan, + RowGroupAccess::Skip, + RowGroupAccess::Selection( + vec![ + RowSelector::skip(5), + RowSelector::select(10), + RowSelector::skip(15), + ] + .into() + ), + RowGroupAccess::Scan, + ]) + ); + } + + #[test] + fn test_new_from_overall_row_selection_invalid_row_count() { + let row_selection = RowSelection::from(vec![RowSelector::select(99)]); + + let err = ParquetAccessPlan::try_new_from_overall_row_selection( + row_selection, + &ROW_GROUP_METADATA, + ) + .unwrap_err() + .to_string(); + + assert_contains!( + err, + "Invalid Parquet RowSelection. File has 100 rows, but selection specifies 99 rows" + ); + } + + #[test] + fn test_new_from_overall_row_selection_boundary_splits() { + let row_selection = RowSelection::from(vec![ + RowSelector::skip(5), + RowSelector::select(10), + RowSelector::skip(20), + RowSelector::select(25), + RowSelector::skip(40), + ]); + + let access_plan = ParquetAccessPlan::try_new_from_overall_row_selection( + row_selection, + &ROW_GROUP_METADATA, + ) + .unwrap(); + + assert_eq!( + access_plan, + ParquetAccessPlan::new(vec![ + RowGroupAccess::Selection( + vec![RowSelector::skip(5), RowSelector::select(5)].into() + ), + RowGroupAccess::Selection( + vec![RowSelector::select(5), RowSelector::skip(15)].into() + ), + RowGroupAccess::Selection( + vec![RowSelector::skip(5), RowSelector::select(25)].into() + ), + RowGroupAccess::Skip, + ]) + ); + } + #[test] fn test_invalid_too_few() { let access_plan = ParquetAccessPlan::new(vec![ @@ -1026,4 +1262,172 @@ mod test { assert_eq!(result.row_group_indexes, vec![0, 1]); } + + // ---------------------------------------------------------------- + // multi-column `reorder_by_statistics` tests + // ---------------------------------------------------------------- + + /// Two-column int32 schema named "a", "b". + fn two_col_schema_descr() -> SchemaDescPtr { + use parquet::basic::Type as PhysicalType; + use parquet::schema::types::Type as SchemaType; + let fields = ["a", "b"] + .iter() + .map(|name| { + Arc::new( + SchemaType::primitive_type_builder(name, PhysicalType::INT32) + .build() + .unwrap(), + ) + }) + .collect(); + let schema = SchemaType::group_type_builder("schema") + .with_fields(fields) + .build() + .unwrap(); + Arc::new(SchemaDescriptor::new(Arc::new(schema))) + } + + /// Build a `ParquetMetaData` with one row group per element of + /// `mins`: `(min(a), min(b))` per row group, `min == max`. + fn parquet_metadata_with_two_col_mins(mins: &[(i32, i32)]) -> ParquetMetaData { + let schema_descr = two_col_schema_descr(); + let row_groups: Vec = mins + .iter() + .map(|&(a, b)| { + let columns = [(0, a), (1, b)] + .iter() + .map(|&(col, m)| { + let stats = ParquetStatistics::int32( + Some(m), + Some(m), + None, + Some(0), + false, + ); + ColumnChunkMetaData::builder(schema_descr.column(col)) + .set_statistics(stats) + .set_num_values(100) + .build() + .unwrap() + }) + .collect(); + RowGroupMetaData::builder(schema_descr.clone()) + .set_num_rows(100) + .set_column_metadata(columns) + .build() + .unwrap() + }) + .collect(); + let file_metadata = + FileMetaData::new(0, 0, None, None, schema_descr.clone(), None); + ParquetMetaData::new(file_metadata, row_groups) + } + + fn arrow_schema_ab_int() -> Schema { + Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ]) + } + + fn sort_expr(name: &str, index: usize, descending: bool) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: Arc::new(Column::new(name, index)), + options: SortOptions { + descending, + nulls_first: true, + }, + } + } + + /// `ORDER BY a ASC, b ASC` with the leading key tied everywhere: + /// the secondary key must break the tie, so RGs order by `min(b)`. + #[test] + fn reorder_by_statistics_breaks_leading_ties_with_secondary_column() { + let metadata = + parquet_metadata_with_two_col_mins(&[(1, 300), (1, 100), (1, 200)]); + let plan = PreparedAccessPlan::new(vec![0, 1, 2], None).unwrap(); + let order = + LexOrdering::new(vec![sort_expr("a", 0, false), sort_expr("b", 1, false)]) + .unwrap(); + + let result = plan + .reorder_by_statistics(&order, &metadata, &arrow_schema_ab_int()) + .unwrap(); + + assert_eq!(result.row_group_indexes, vec![1, 2, 0]); + } + + /// `ORDER BY a ASC, b DESC`: the secondary key's direction is + /// honored relative to the leading key, so ties on `min(a)` order + /// by `min(b)` DESC. + #[test] + fn reorder_by_statistics_honors_secondary_direction() { + let metadata = + parquet_metadata_with_two_col_mins(&[(1, 100), (1, 300), (0, 500)]); + let plan = PreparedAccessPlan::new(vec![0, 1, 2], None).unwrap(); + let order = + LexOrdering::new(vec![sort_expr("a", 0, false), sort_expr("b", 1, true)]) + .unwrap(); + + let result = plan + .reorder_by_statistics(&order, &metadata, &arrow_schema_ab_int()) + .unwrap(); + + // a=0 first, then the two a=1 groups by b DESC: 300 before 100. + assert_eq!(result.row_group_indexes, vec![2, 1, 0]); + } + + /// `ORDER BY a DESC, b DESC` is normalized to ASC lexsort here and + /// flipped by the later `reverse()`: both keys sort ASC relative to + /// the leading direction, so reversing yields `(a DESC, b DESC)`. + #[test] + fn reorder_by_statistics_normalizes_desc_desc_for_reverse() { + let metadata = + parquet_metadata_with_two_col_mins(&[(1, 300), (2, 100), (1, 100)]); + let plan = PreparedAccessPlan::new(vec![0, 1, 2], None).unwrap(); + let order = + LexOrdering::new(vec![sort_expr("a", 0, true), sort_expr("b", 1, true)]) + .unwrap(); + + let result = plan + .reorder_by_statistics(&order, &metadata, &arrow_schema_ab_int()) + .unwrap(); + + // ASC lexsort of (a, b): (1,100) < (1,300) < (2,100); the later + // reverse() produces (2,100), (1,300), (1,100) = (a DESC, b DESC). + assert_eq!(result.row_group_indexes, vec![2, 0, 1]); + } + + /// A non-`Column` *secondary* expression stops the stats walk but + /// keeps the leading column's reorder (prefix semantics). + #[test] + fn reorder_by_statistics_keeps_leading_prefix_on_non_column_secondary() { + let metadata = + parquet_metadata_with_two_col_mins(&[(5, 300), (3, 100), (4, 200)]); + let plan = PreparedAccessPlan::new(vec![0, 1, 2], None).unwrap(); + let order = LexOrdering::new(vec![ + sort_expr("a", 0, false), + PhysicalSortExpr { + expr: Arc::new(BinaryExpr::new( + Arc::new(Column::new("b", 1)), + Operator::Plus, + lit(1i32), + )), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + ]) + .unwrap(); + + let result = plan + .reorder_by_statistics(&order, &metadata, &arrow_schema_ab_int()) + .unwrap(); + + // Ordered by min(a) ASC only: 3, 4, 5. + assert_eq!(result.row_group_indexes, vec![1, 2, 0]); + } } diff --git a/datafusion/datasource-parquet/src/bloom_filter.rs b/datafusion/datasource-parquet/src/bloom_filter.rs index 9388aba4385f2..8e55c12ba0896 100644 --- a/datafusion/datasource-parquet/src/bloom_filter.rs +++ b/datafusion/datasource-parquet/src/bloom_filter.rs @@ -33,38 +33,63 @@ use parquet::data_type::Decimal; /// This structure implements [`PruningStatistics`] and is used to prune /// Parquet row groups and data pages based on the query predicate. #[derive(Debug, Clone, Default)] -pub(crate) struct BloomFilterStatistics { - /// Per-column Bloom filters - /// Key: predicate column name - /// Value: - /// * [`Sbbf`] (Bloom filter), - /// * Parquet physical [`Type`] needed to evaluate literals against the filter - column_sbbf: HashMap, +pub struct BloomFilterStatistics { + /// Per-column Bloom filters keyed by predicate column name. + column_sbbf: HashMap, +} + +#[derive(Debug, Clone)] +struct ColumnBloomFilter { + /// [`Sbbf`] (Bloom filter). + sbbf: Sbbf, + /// Parquet physical [`Type`] needed to evaluate literals against the filter. + physical_type: Type, + /// Type length from the Parquet column descriptor. + type_length: i32, } impl BloomFilterStatistics { /// Create an empty [`BloomFilterStatistics`] - pub(crate) fn new() -> Self { + pub fn new() -> Self { Default::default() } /// Create an empty [`BloomFilterStatistics`] with the specified capacity - pub(crate) fn with_capacity(capacity: usize) -> Self { + pub fn with_capacity(capacity: usize) -> Self { Self { column_sbbf: HashMap::with_capacity(capacity), } } - /// Add a Bloom filter and type for the specified column - pub(crate) fn insert(&mut self, column: impl Into, sbbf: Sbbf, ty: Type) { - self.column_sbbf.insert(column.into(), (sbbf, ty)); + /// Add a Bloom filter for the specified column, along with the column's + /// Parquet physical [`Type`] and type length from the column descriptor. + pub fn insert( + &mut self, + column: impl Into, + sbbf: Sbbf, + ty: Type, + type_length: i32, + ) { + self.column_sbbf.insert( + column.into(), + ColumnBloomFilter { + sbbf, + physical_type: ty, + type_length, + }, + ); } /// Helper function for checking if [`Sbbf`] filter contains [`ScalarValue`]. /// /// In case the type of scalar is not supported, returns `true`, assuming that the /// value may be present. - fn check_scalar(sbbf: &Sbbf, value: &ScalarValue, parquet_type: &Type) -> bool { + fn check_scalar( + sbbf: &Sbbf, + value: &ScalarValue, + parquet_type: &Type, + type_length: i32, + ) -> bool { match value { ScalarValue::Utf8(Some(v)) | ScalarValue::Utf8View(Some(v)) @@ -113,8 +138,14 @@ impl BloomFilterStatistics { sbbf.check(&decimal) } Type::FIXED_LEN_BYTE_ARRAY => { - // keep with from_bytes_to_i128 - let b = v.to_be_bytes().to_vec(); + let Ok(type_length) = usize::try_from(type_length) else { + return true; + }; + if type_length == 0 || type_length > 16 { + return true; + } + let b = v.to_be_bytes(); + let b = b[(b.len() - type_length)..].to_vec(); // Use Decimal constructor after https://github.com/apache/arrow-rs/issues/5325 let decimal = Decimal::Bytes { value: b.into(), @@ -125,9 +156,12 @@ impl BloomFilterStatistics { } _ => true, }, - ScalarValue::Dictionary(_, inner) => { - BloomFilterStatistics::check_scalar(sbbf, inner, parquet_type) - } + ScalarValue::Dictionary(_, inner) => BloomFilterStatistics::check_scalar( + sbbf, + inner, + parquet_type, + type_length, + ), _ => true, } } @@ -164,7 +198,7 @@ impl PruningStatistics for BloomFilterStatistics { column: &Column, values: &HashSet, ) -> Option { - let (sbbf, parquet_type) = self.column_sbbf.get(column.name.as_str())?; + let column_bloom_filter = self.column_sbbf.get(column.name.as_str())?; // Bloom filters are probabilistic data structures that can return false // positives (i.e. it might return true even if the value is not @@ -173,7 +207,14 @@ impl PruningStatistics for BloomFilterStatistics { let known_not_present = values .iter() - .map(|value| BloomFilterStatistics::check_scalar(sbbf, value, parquet_type)) + .map(|value| { + BloomFilterStatistics::check_scalar( + &column_bloom_filter.sbbf, + value, + &column_bloom_filter.physical_type, + column_bloom_filter.type_length, + ) + }) // The row group doesn't contain any of the values if // all the checks are false .all(|v| !v); @@ -201,15 +242,28 @@ mod tests { use crate::test_util::ExpectedPruning; use crate::{ParquetAccessPlan, ParquetFileMetrics, RowGroupAccessPlanFilter}; + use arrow::array::Decimal128Array; use arrow::datatypes::{DataType, Field, Schema}; + use bytes::{BufMut, BytesMut}; use datafusion_common::Result; use datafusion_expr::{Expr, col, lit}; use datafusion_physical_expr::planner::logical2physical; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; - use datafusion_pruning::PruningPredicate; - use object_store::ObjectStoreExt; + use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; + use object_store::{ObjectStore, ObjectStoreExt}; + use parquet::arrow::ArrowWriter; use parquet::arrow::ParquetRecordBatchStreamBuilder; - use parquet::arrow::async_reader::ParquetObjectReader; + use parquet::file::properties::{EnabledStatistics, WriterProperties}; + + fn build_test_pruning_predicate( + expr: Arc, + schema: Schema, + ) -> PruningPredicate { + PruningPredicateBuilder::new() + .with_file_schema(Arc::new(schema)) + .try_build(expr) + .unwrap() + } #[tokio::test] async fn test_row_group_bloom_filter_pruning_predicate_simple_expr() { @@ -276,8 +330,7 @@ mod tests { false, ); let expr = logical2physical(&expr, &schema); - let pruning_predicate = - PruningPredicate::try_new(expr, Arc::new(schema)).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, schema); let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate( file_name, @@ -375,6 +428,80 @@ mod tests { .await } + #[tokio::test] + async fn test_row_group_bloom_filter_pruning_predicate_decimal128() { + for precision in [19, 20, 21, 28, 38] { + let scale = 2; + let data = parquet_decimal128_with_bloom_filter( + precision, + scale, + vec![100, 200, 300, 400, 500, 600], + ); + let schema = Schema::new(vec![Field::new( + "decimal_col", + DataType::Decimal128(precision, scale), + true, + )]); + let expr = col("decimal_col").eq(Expr::Literal( + ScalarValue::Decimal128(Some(500), precision, scale), + None, + )); + let expr = logical2physical(&expr, &schema); + let pruning_predicate = build_test_pruning_predicate(expr, schema); + + let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate( + &format!("decimal128-{precision}.parquet"), + data, + &pruning_predicate, + ) + .await + .unwrap(); + + assert_eq!( + pruned_row_groups.access_plan().row_group_indexes(), + vec![2], + "precision {precision}" + ); + } + } + + #[tokio::test] + async fn test_row_group_bloom_filter_pruning_predicate_negative_decimal128() { + for precision in [19, 20, 21, 28, 38] { + let scale = 2; + let data = parquet_decimal128_with_bloom_filter( + precision, + scale, + vec![-100, -200, -300, -400, -500, -600], + ); + let schema = Schema::new(vec![Field::new( + "decimal_col", + DataType::Decimal128(precision, scale), + true, + )]); + let expr = col("decimal_col").eq(Expr::Literal( + ScalarValue::Decimal128(Some(-500), precision, scale), + None, + )); + let expr = logical2physical(&expr, &schema); + let pruning_predicate = build_test_pruning_predicate(expr, schema); + + let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate( + &format!("negative-decimal128-{precision}.parquet"), + data, + &pruning_predicate, + ) + .await + .unwrap(); + + assert_eq!( + pruned_row_groups.access_plan().row_group_indexes(), + vec![2], + "precision {precision}" + ); + } + } + struct BloomFilterTest { file_name: String, schema: Schema, @@ -452,8 +579,7 @@ mod tests { let data = bytes::Bytes::from(std::fs::read(path).unwrap()); let expr = logical2physical(&expr, &schema); - let pruning_predicate = - PruningPredicate::try_new(expr, Arc::new(schema)).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, schema); let pruned_row_groups = test_row_group_bloom_filter_pruning_predicate( &file_name, @@ -467,6 +593,37 @@ mod tests { } } + fn parquet_decimal128_with_bloom_filter( + precision: u8, + scale: i8, + values: Vec, + ) -> bytes::Bytes { + let schema = Arc::new(Schema::new(vec![Field::new( + "decimal_col", + DataType::Decimal128(precision, scale), + true, + )])); + let array = Arc::new( + Decimal128Array::from(values) + .with_precision_and_scale(precision, scale) + .unwrap(), + ) as ArrayRef; + let batch = + arrow::array::RecordBatch::try_new(schema.clone(), vec![array]).unwrap(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(2)) + .set_bloom_filter_enabled(true) + .set_statistics_enabled(EnabledStatistics::None) + .build(); + let mut out = BytesMut::new().writer(); + { + let mut writer = ArrowWriter::try_new(&mut out, schema, Some(props)).unwrap(); + writer.write(&batch).unwrap(); + writer.finish().unwrap(); + } + out.into_inner().freeze() + } + /// Evaluates the pruning predicate on the specified row groups and returns the row groups that are left async fn test_row_group_bloom_filter_pruning_predicate( file_name: &str, @@ -492,17 +649,11 @@ mod tests { let metrics = ExecutionPlanMetricsSet::new(); let file_metrics = ParquetFileMetrics::new(0, object_meta.location.as_ref(), &metrics); - let inner = - ParquetObjectReader::new(Arc::new(in_memory), object_meta.location.clone()) - .with_file_size(object_meta.size); - + let store: Arc = Arc::new(in_memory); let partitioned_file = PartitionedFile::new_from_meta(object_meta); - let reader = ParquetFileReader { - inner, - file_metrics: file_metrics.clone(), - partitioned_file, - }; + let reader = + ParquetFileReader::new(file_metrics.clone(), store, partitioned_file); let mut builder = ParquetRecordBatchStreamBuilder::new(reader).await.unwrap(); let access_plan = ParquetAccessPlan::new_all(builder.metadata().num_row_groups()); @@ -520,6 +671,7 @@ mod tests { column_name.to_string(), column_idx, builder.parquet_schema().column(column_idx).physical_type(), + builder.parquet_schema().column(column_idx).type_length(), )) }) .collect::>(); @@ -532,7 +684,8 @@ mod tests { for idx in pruned_row_groups.row_group_indexes() { let mut bloom_filters = BloomFilterStatistics::with_capacity(parquet_columns.len()); - for (column_name, column_idx, physical_type) in &parquet_columns { + for (column_name, column_idx, physical_type, type_length) in &parquet_columns + { let bf = match builder .get_row_group_column_bloom_filter(idx, *column_idx) .await @@ -545,7 +698,12 @@ mod tests { continue; } }; - bloom_filters.insert(column_name.clone(), bf, *physical_type); + bloom_filters.insert( + column_name.clone(), + bf, + *physical_type, + *type_length, + ); } row_group_bloom_filters[idx] = bloom_filters; } diff --git a/datafusion/datasource-parquet/src/decoder_projection.rs b/datafusion/datasource-parquet/src/decoder_projection.rs new file mode 100644 index 0000000000000..89fdc01af4eda --- /dev/null +++ b/datafusion/datasource-parquet/src/decoder_projection.rs @@ -0,0 +1,161 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Decoder-projection construction for the parquet scan. +//! +//! [`DecoderProjection`] owns the two halves of "project a decoded parquet +//! batch onto the scan's output schema": +//! +//! * the [`ProjectionMask`] installed on the parquet decoder (and on any +//! rebuild performed via `into_builder` at a row-group boundary), and +//! * the per-batch transform ([`DecoderProjection::map`]) that applies the +//! projector and, when needed, rebuilds the batch with the user's +//! `output_schema` to recover metadata / nullability the file schema does +//! not carry. +//! +//! The opener constructs one [`DecoderProjection`] per file via +//! [`DecoderProjection::try_new`] and hands it to the push-decoder stream, +//! which calls [`map`](DecoderProjection::map) on every decoded batch. + +use std::sync::Arc; + +use arrow::array::{RecordBatch, RecordBatchOptions}; +use arrow::datatypes::SchemaRef; + +use datafusion_common::Result; +use datafusion_physical_expr::projection::{ProjectionExprs, Projector}; +use datafusion_physical_expr::utils::reassign_expr_columns; +use datafusion_physical_expr_adapter::replace_columns_with_literals; + +use parquet::arrow::ProjectionMask; +use parquet::schema::types::SchemaDescriptor; + +use crate::opener::{VirtualColumnsState, append_fields}; +use crate::projection_read_plan::build_projection_read_plan; + +/// Per-file decoder projection: the [`ProjectionMask`] installed on the +/// parquet decoder, plus the per-batch transform that maps the decoder's +/// output onto the scan's `output_schema`. +/// +/// Built once per file by the opener via [`Self::try_new`]; the +/// push-decoder stream installs [`Self::projection_mask`] on the decoder +/// (and on any rebuild performed via `into_builder` at a row-group +/// boundary) and calls [`Self::map`] on every decoded batch. +pub(crate) struct DecoderProjection { + projection_mask: ProjectionMask, + projector: Projector, + output_schema: SchemaRef, + /// `true` when the projector's output schema differs from `output_schema` + /// in metadata / nullability and [`map`](Self::map) must rebuild the batch + /// with `output_schema`. + replace_schema: bool, +} + +impl DecoderProjection { + /// Build the decoder projection for a file. + /// + /// `projection` references columns in `physical_file_schema` (i.e. already + /// adapted by the per-file expr adapter); `parquet_schema` is the + /// corresponding parquet [`SchemaDescriptor`]. `output_schema` is what + /// consumers of the scan stream expect. + /// + /// `virtual_state`, when present, describes virtual columns the reader + /// will append to each decoded batch (e.g. parquet `row_number`). Virtual + /// columns are stripped from the projection fed into + /// `build_projection_read_plan` (which only understands file columns) and + /// appended to the stream schema so the projector can resolve them. + pub(crate) fn try_new( + projection: &ProjectionExprs, + physical_file_schema: &SchemaRef, + parquet_schema: &SchemaDescriptor, + output_schema: &SchemaRef, + virtual_state: Option<&VirtualColumnsState>, + ) -> Result { + // Virtual columns are produced by the reader separately from the + // projection mask, so strip them from the expressions we feed into + // `build_projection_read_plan`. We substitute each virtual column + // reference with a null literal; that leaves the remaining Column + // refs (into `physical_file_schema`) intact for + // `ProjectionMask::roots`, which only understands file columns. + let projection_for_read_plan = match virtual_state { + None => projection.clone(), + Some(state) => projection.clone().try_map_exprs(|expr| { + replace_columns_with_literals(expr, state.null_replacements()) + })?, + }; + let read_plan = build_projection_read_plan( + projection_for_read_plan.expr_iter(), + physical_file_schema, + parquet_schema, + ); + + // The reader produces projected file columns followed by any virtual + // columns (`ArrowReaderOptions::with_virtual_columns` appends them to + // each decoded batch). + let stream_schema = match virtual_state { + Some(state) => { + append_fields(&read_plan.projected_schema, state.virtual_columns()) + } + None => Arc::clone(&read_plan.projected_schema), + }; + + // Rebase the projection onto the decoder's stream schema (column + // indices change because the decoder yields only the masked columns). + let rebased_projection = projection + .clone() + .try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?; + let projector = rebased_projection.make_projector(&stream_schema)?; + + // Compare against the projector's *output* schema rather than the + // stream schema, so future widening of the mask (e.g. for post-scan + // filter columns) does not flip this flag. + let replace_schema = projector.output_schema() != output_schema; + + Ok(Self { + projection_mask: read_plan.projection_mask, + projector, + output_schema: Arc::clone(output_schema), + replace_schema, + }) + } + + /// The projection mask to install on every parquet decoder in the scan. + pub(crate) fn projection_mask(&self) -> &ProjectionMask { + &self.projection_mask + } + + /// Map a decoded batch onto the scan's output schema. + /// + /// Applies the [`Projector`] and, when the projector's output schema + /// differs from `output_schema` in metadata or nullability, rebuilds the + /// batch with `output_schema` (some writers emit OPTIONAL fields even when + /// the data has no nulls; some logical schemas carry field-level metadata + /// the file schema does not). + pub(crate) fn map(&self, batch: &RecordBatch) -> Result { + let projected = self.projector.project_batch(batch)?; + if !self.replace_schema { + return Ok(projected); + } + let (_stream_schema, arrays, num_rows) = projected.into_parts(); + let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); + Ok(RecordBatch::try_new_with_options( + Arc::clone(&self.output_schema), + arrays, + &options, + )?) + } +} diff --git a/datafusion/datasource-parquet/src/file_format.rs b/datafusion/datasource-parquet/src/file_format.rs index fe81504e320d7..6358201c06fa5 100644 --- a/datafusion/datasource-parquet/src/file_format.rs +++ b/datafusion/datasource-parquet/src/file_format.rs @@ -24,11 +24,12 @@ use std::sync::Arc; // Re-export so the historical `file_format::*` paths still resolve. #[expect(deprecated)] +pub use crate::schema_coercion::coerce_int96_to_resolution; pub use crate::schema_coercion::{ - Int96Coercer, apply_file_schema_type_coercions, coerce_file_schema_to_string_type, - coerce_file_schema_to_view_type, coerce_int96_to_resolution, - transform_binary_to_string, transform_schema_to_view, + Int96Coercer, apply_file_schema_type_coercions, transform_binary_to_string, + transform_schema_to_view, }; + pub use crate::sink::ParquetSink; use arrow::datatypes::{Fields, Schema, SchemaRef}; @@ -49,6 +50,7 @@ use datafusion_common::{ use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; use datafusion_datasource::sink::DataSinkExec; +use datafusion_datasource::write::get_writer_schema; use datafusion_expr::dml::InsertOp; use datafusion_physical_expr_common::sort_expr::{LexOrdering, LexRequirement}; use datafusion_physical_plan::ExecutionPlan; @@ -296,6 +298,7 @@ async fn get_file_decryption_properties( } #[cfg(not(feature = "parquet_encryption"))] +#[expect(clippy::unused_async)] async fn get_file_decryption_properties( _state: &dyn Session, _options: &TableParquetOptions, @@ -366,7 +369,13 @@ impl FileFormat for ParquetFormat { }) .boxed() // Workaround https://github.com/rust-lang/rust/issues/64552 // fetch schemas concurrently, if requested - .buffer_unordered(state.config_options().execution.meta_fetch_concurrency) + .buffer_unordered( + state + .config_options() + .execution + .meta_fetch_concurrency + .get(), + ) .try_collect() .await?; @@ -526,11 +535,18 @@ impl FileFormat for ParquetFormat { // Convert ordering requirements to Parquet SortingColumns for file metadata let sorting_columns = if let Some(ref requirements) = order_requirements { let ordering: LexOrdering = requirements.clone().into(); + let writer_schema = get_writer_schema(&conf); // In cases like `COPY (... ORDER BY ...) TO ...` the ORDER BY clause // may not be compatible with Parquet sorting columns (e.g. ordering on `random()`). // So if we cannot create a Parquet sorting column from the ordering requirement, // we skip setting sorting columns on the Parquet sink. - lex_ordering_to_sorting_columns(&ordering).ok() + lex_ordering_to_sorting_columns( + &ordering, + conf.output_schema(), + &writer_schema, + ) + .ok() + .filter(|columns| !columns.is_empty()) } else { None }; @@ -626,7 +642,7 @@ pub async fn fetch_parquet_metadata( object_meta: &ObjectMeta, size_hint: Option, decryption_properties: Option<&FileDecryptionProperties>, - file_metadata_cache: Option>, + file_metadata_cache: Option>, ) -> Result> { let decryption_properties = decryption_properties.cloned().map(Arc::new); DFParquetMetadata::new(store, object_meta) @@ -650,7 +666,7 @@ pub async fn fetch_statistics( file: &ObjectMeta, metadata_size_hint: Option, decryption_properties: Option<&FileDecryptionProperties>, - file_metadata_cache: Option>, + file_metadata_cache: Option>, ) -> Result { let decryption_properties = decryption_properties.cloned().map(Arc::new); DFParquetMetadata::new(store, file) @@ -672,3 +688,129 @@ pub fn statistics_from_parquet_meta_calc( ) -> Result { DFParquetMetadata::statistics_from_parquet_metadata(metadata, &table_schema) } + +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf::{self, parquet_column_options, parquet_options}; + +/// Encode a [`ParquetFormatFactory`]'s options as their protobuf form. +/// +/// The reverse direction is `TryFrom<&protobuf::TableParquetOptions> for +/// TableParquetOptions` in `datafusion-proto-models`: `TableParquetOptions` is +/// a `datafusion-common` type, so that half cannot live here. +#[cfg(feature = "proto")] +impl From<&ParquetFormatFactory> for protobuf::TableParquetOptions { + fn from(factory: &ParquetFormatFactory) -> Self { + let global_options = if let Some(ref options) = factory.options { + options.clone() + } else { + return protobuf::TableParquetOptions::default(); + }; + + let column_specific_options = global_options.column_specific_options; + protobuf::TableParquetOptions { + global: Some(protobuf::ParquetOptions { + enable_page_index: global_options.global.enable_page_index, + pruning: global_options.global.pruning, + skip_metadata: global_options.global.skip_metadata, + metadata_size_hint_opt: global_options.global.metadata_size_hint.map(|size| { + parquet_options::MetadataSizeHintOpt::MetadataSizeHint(size as u64) + }), + pushdown_filters: global_options.global.pushdown_filters, + reorder_filters: global_options.global.reorder_filters, + force_filter_selections: global_options.global.force_filter_selections, + data_pagesize_limit: global_options.global.data_pagesize_limit as u64, + write_batch_size: global_options.global.write_batch_size as u64, + writer_version: global_options.global.writer_version.to_string(), + compression_opt: global_options.global.compression.map(|compression| { + parquet_options::CompressionOpt::Compression(compression) + }), + dictionary_enabled_opt: global_options.global.dictionary_enabled.map(|enabled| { + parquet_options::DictionaryEnabledOpt::DictionaryEnabled(enabled) + }), + dictionary_page_size_limit: global_options.global.dictionary_page_size_limit as u64, + statistics_enabled_opt: global_options.global.statistics_enabled.map(|enabled| { + parquet_options::StatisticsEnabledOpt::StatisticsEnabled(enabled) + }), + max_row_group_size: global_options.global.max_row_group_size as u64, + max_in_list_size: global_options.global.max_in_list_size as u64, + created_by: global_options.global.created_by.clone(), + column_index_truncate_length_opt: global_options.global.column_index_truncate_length.map(|length| { + parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length as u64) + }), + statistics_truncate_length_opt: global_options.global.statistics_truncate_length.map(|length| { + parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(length as u64) + }), + data_page_row_count_limit: global_options.global.data_page_row_count_limit as u64, + encoding_opt: global_options.global.encoding.map(|encoding| { + parquet_options::EncodingOpt::Encoding(encoding) + }), + bloom_filter_on_read: global_options.global.bloom_filter_on_read, + bloom_filter_on_write: global_options.global.bloom_filter_on_write, + bloom_filter_fpp_opt: global_options.global.bloom_filter_fpp.map(|fpp| { + parquet_options::BloomFilterFppOpt::BloomFilterFpp(fpp) + }), + bloom_filter_ndv_opt: global_options.global.bloom_filter_ndv.map(|ndv| { + parquet_options::BloomFilterNdvOpt::BloomFilterNdv(ndv) + }), + allow_single_file_parallelism: global_options.global.allow_single_file_parallelism, + maximum_parallel_row_group_writers: global_options.global.maximum_parallel_row_group_writers as u64, + maximum_buffered_record_batches_per_stream: global_options.global.maximum_buffered_record_batches_per_stream as u64, + schema_force_view_types: global_options.global.schema_force_view_types, + binary_as_string: global_options.global.binary_as_string, + skip_arrow_metadata: global_options.global.skip_arrow_metadata, + coerce_int96_opt: global_options.global.coerce_int96.map(|compression| { + parquet_options::CoerceInt96Opt::CoerceInt96(compression) + }), + coerce_int96_tz_opt: global_options.global.coerce_int96_tz.map(|tz| { + parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(tz) + }), + max_predicate_cache_size_opt: global_options.global.max_predicate_cache_size.map(|size| { + parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(size as u64) + }), + max_row_group_bytes_opt: global_options.global.max_row_group_bytes.map(|size| { + parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(size.get() as u64) + }), + content_defined_chunking: Some(protobuf::ParquetCdcOptions { + enabled: global_options.global.content_defined_chunking.enabled, + min_chunk_size: global_options.global.content_defined_chunking.min_chunk_size as u64, + max_chunk_size: global_options.global.content_defined_chunking.max_chunk_size as u64, + norm_level: global_options.global.content_defined_chunking.norm_level, + }), + }), + column_specific_options: column_specific_options.into_iter().map(|(column_name, options)| { + protobuf::ParquetColumnSpecificOptions { + column_name, + options: Some(protobuf::ParquetColumnOptions { + bloom_filter_enabled_opt: options.bloom_filter_enabled.map(|enabled| { + parquet_column_options::BloomFilterEnabledOpt::BloomFilterEnabled(enabled) + }), + encoding_opt: options.encoding.map(|encoding| { + parquet_column_options::EncodingOpt::Encoding(encoding) + }), + dictionary_enabled_opt: options.dictionary_enabled.map(|enabled| { + parquet_column_options::DictionaryEnabledOpt::DictionaryEnabled(enabled) + }), + compression_opt: options.compression.map(|compression| { + parquet_column_options::CompressionOpt::Compression(compression) + }), + statistics_enabled_opt: options.statistics_enabled.map(|enabled| { + parquet_column_options::StatisticsEnabledOpt::StatisticsEnabled(enabled) + }), + bloom_filter_fpp_opt: options.bloom_filter_fpp.map(|fpp| { + parquet_column_options::BloomFilterFppOpt::BloomFilterFpp(fpp) + }), + bloom_filter_ndv_opt: options.bloom_filter_ndv.map(|ndv| { + parquet_column_options::BloomFilterNdvOpt::BloomFilterNdv(ndv) + }), + }) + } + }).collect(), + key_value_metadata: global_options.key_value_metadata + .iter() + .filter_map(|(key, value)| { + value.as_ref().map(|v| (key.clone(), v.clone())) + }) + .collect(), + } + } +} diff --git a/datafusion/datasource-parquet/src/metadata.rs b/datafusion/datasource-parquet/src/metadata.rs index c32e45935636f..38db1e22ee6c0 100644 --- a/datafusion/datasource-parquet/src/metadata.rs +++ b/datafusion/datasource-parquet/src/metadata.rs @@ -18,16 +18,17 @@ //! [`DFParquetMetadata`] for fetching Parquet file metadata, statistics //! and schema information. +use crate::file_format::ObjectStoreFetch; use crate::{Int96Coercer, apply_file_schema_type_coercions}; use arrow::array::{Array, ArrayRef, BooleanArray}; -use arrow::compute::and; use arrow::compute::kernels::cmp::eq; -use arrow::compute::sum; +use arrow::compute::{and, sum}; use arrow::datatypes::{DataType, Schema, SchemaRef, TimeUnit}; use datafusion_common::encryption::FileDecryptionProperties; use datafusion_common::stats::Precision; use datafusion_common::{ - ColumnStatistics, DataFusionError, Result, ScalarValue, Statistics, + ColumnStatistics, DataFusionError, HashMap, Result, ScalarValue, Statistics, + internal_datafusion_err, }; use datafusion_execution::cache::cache_manager::{ CachedFileMetadataEntry, FileMetadata, FileMetadataCache, @@ -43,12 +44,12 @@ use parquet::DecodeResult; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; use parquet::arrow::{parquet_column, parquet_to_arrow_schema}; use parquet::file::metadata::{ - PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder, RowGroupMetaData, - SortingColumn, + PageIndexPolicy, ParquetMetaData, ParquetMetaDataPushDecoder, ParquetMetaDataReader, + RowGroupMetaData, SortingColumn, }; +use parquet::file::statistics::Statistics as ParquetStatistics; use parquet::schema::types::SchemaDescriptor; use std::any::Any; -use std::collections::HashMap; use std::sync::Arc; /// Minimum fraction of row groups that must report NDV statistics for the @@ -65,11 +66,27 @@ const PARTIAL_NDV_THRESHOLD: f64 = 0.75; /// [`ParquetFileReaderFactory`]: crate::ParquetFileReaderFactory #[derive(Debug)] pub struct DFParquetMetadata<'a> { + /// Source of the Parquet file's bytes. store: &'a dyn ObjectStore, + /// Location, size and last-modified time of the target Parquet file. object_meta: &'a ObjectMeta, + /// Hint for the number of trailing bytes to prefetch before parsing the + /// footer, mirroring [`ParquetMetaDataReader::with_prefetch_hint`]. metadata_size_hint: Option, + /// Decryption properties used to read files encrypted with Parquet + /// Modular Encryption, mirroring + /// [`ParquetMetaDataReader::with_decryption_properties`]. decryption_properties: Option>, - file_metadata_cache: Option>, + /// Optional cache of previously fetched [`ParquetMetaData`], keyed by + /// file location. + file_metadata_cache: Option>, + /// Policy controlling whether the Parquet page index (column and offset + /// indexes) is fetched, mirroring + /// [`ParquetMetaDataReader::with_page_index_policy`]. + /// + /// `None` means the effective policy is chosen automatically, see + /// [`DFParquetMetadata::effective_page_index_policy`]. + page_index_policy: Option, /// timeunit to coerce INT96 timestamps to pub coerce_int96: Option, /// Optional timezone applied to INT96-coerced timestamps. @@ -77,6 +94,10 @@ pub struct DFParquetMetadata<'a> { } impl<'a> DFParquetMetadata<'a> { + /// Create a new `DFParquetMetadata` for the given file. + /// + /// Use the `with_*` builder methods to customize behavior + /// before calling [`Self::fetch_metadata`] or [`Self::fetch_schema`]. pub fn new(store: &'a dyn ObjectStore, object_meta: &'a ObjectMeta) -> Self { Self { store, @@ -84,18 +105,29 @@ impl<'a> DFParquetMetadata<'a> { metadata_size_hint: None, decryption_properties: None, file_metadata_cache: None, + page_index_policy: None, coerce_int96: None, coerce_int96_tz: None, } } - /// set metadata size hint + /// Set a hint for the number of trailing bytes to prefetch from the end + /// of the file, equivalent to + /// [`ParquetMetaDataReader::with_prefetch_hint`]. + /// + /// Providing a good estimate of the footer (and, if requested, page index) + /// size can save an extra I/O round trip when fetching metadata from the + /// store. pub fn with_metadata_size_hint(mut self, metadata_size_hint: Option) -> Self { self.metadata_size_hint = metadata_size_hint; self } - /// set decryption properties + /// Set the decryption properties used to read an encrypted Parquet file, + /// equivalent to [`ParquetMetaDataReader::with_decryption_properties`]. + /// + /// Only needed when the target file was written with Parquet Modular + /// Encryption. pub fn with_decryption_properties( mut self, decryption_properties: Option>, @@ -104,32 +136,77 @@ impl<'a> DFParquetMetadata<'a> { self } - /// set file metadata cache + /// Set an optional [`FileMetadataCache`] used to avoid re-fetching + /// [`ParquetMetaData`] for files that have already been read. pub fn with_file_metadata_cache( mut self, - file_metadata_cache: Option>, + file_metadata_cache: Option>, ) -> Self { self.file_metadata_cache = file_metadata_cache; self } - /// Set timeunit to coerce INT96 timestamps to + /// Sets the policy for loading parquet page index structures (column and + /// offset indexes), equivalent to + /// [`ParquetMetaDataReader::with_page_index_policy`]. + /// + /// Passing `None` uses a default automatically, based on whether a metadata + /// cache is configured. + pub fn with_page_index_policy( + mut self, + page_index_policy: Option, + ) -> Self { + self.page_index_policy = page_index_policy; + self + } + + /// Set the [`TimeUnit`] that INT96 timestamp columns should be coerced + /// to when reading the schema. + /// + /// INT96 in Parquet has no defined unit or timezone, so leaving this + /// `None` reads INT96 columns as nanosecond timestamps with no timezone + /// — DataFusion's default behavior. pub fn with_coerce_int96(mut self, time_unit: Option) -> Self { self.coerce_int96 = time_unit; self } /// Set the optional timezone applied to INT96-coerced timestamps. + /// + /// Only used when [`Self::with_coerce_int96`] has also been set, and + /// otherwise has no effect. pub fn with_coerce_int96_tz(mut self, timezone: Option>) -> Self { self.coerce_int96_tz = timezone; self } - /// Fetch parquet metadata from the remote object store + /// Fetch the [`ParquetMetaData`] for this file. + /// + /// Consults the [`FileMetadataCache`] first when one is configured and + /// falls back to reading from the object store via + /// [`ParquetMetaDataPushDecoder`] on a cache miss. pub async fn fetch_metadata(&self) -> Result> { - // implementation to fetch parquet metadata + // fetch_metadata + // │ + // ├─ cache_metadata = encryption check + // ├─ page_index_policy = caller override OR default + // │ + // ├─ CACHE HIT? + // │ │ + // │ ├─ has index OR policy=Skip? → return cache + // │ │ + // │ └─ else (footer only, wants index) + // │ → load_page_index (index bytes only) + // │ → cache_metadata() if allowed + // │ → return + // │ + // └─ CACHE MISS + // → fetch_metadata_from_store(policy) + // → cache_metadata() if allowed + // → return let cache_metadata = !cfg!(feature = "parquet_encryption") || self.decryption_properties.is_none(); + let page_index_policy = self.effective_page_index_policy(cache_metadata); if cache_metadata && let Some(file_metadata_cache) = self.file_metadata_cache.as_ref() @@ -140,9 +217,79 @@ impl<'a> DFParquetMetadata<'a> { .as_any() .downcast_ref::() { - return Ok(Arc::clone(cached_parquet.parquet_metadata())); + let cached_metadata = Arc::clone(cached_parquet.parquet_metadata()); + // Reuse the cache when it already has page index, or when the caller + // asked to skip page index I/O (footer-only metadata is sufficient). + if Self::metadata_has_page_index(cached_metadata.as_ref()) + || page_index_policy == PageIndexPolicy::Skip + { + return Ok(cached_metadata); + } + let metadata = + Self::load_page_index(self.store, self.object_meta, cached_metadata) + .await?; + if cache_metadata { + self.cache_metadata(Arc::clone(&metadata))?; + } + return Ok(metadata); + } + + let metadata = self.fetch_metadata_from_store(page_index_policy).await?; + if cache_metadata { + self.cache_metadata(Arc::clone(&metadata))?; + } + Ok(metadata) + } + + /// Resolve the [`PageIndexPolicy`] to use for a fetch. + fn effective_page_index_policy(&self, cache_metadata: bool) -> PageIndexPolicy { + self.page_index_policy.unwrap_or_else(|| { + // fetching the page index often requires a second IO (after the + // main metadata), so it is not free. + if cache_metadata && self.file_metadata_cache.is_some() { + // When there is a cache available, retrieve the page index + // heuristically on the assumption it will be used multiple + // times + PageIndexPolicy::Optional + } else { + PageIndexPolicy::Skip + } + }) + } + + /// Check whether `metadata` already has both the column index and the + /// offset index populated (see [`ParquetMetaData::column_index`] and + /// [`ParquetMetaData::offset_index`]). + /// + /// Used to decide whether page index I/O can be skipped. + fn metadata_has_page_index(metadata: &ParquetMetaData) -> bool { + metadata.column_index().is_some() && metadata.offset_index().is_some() + } + + /// Store `metadata` in the configured [`FileMetadataCache`], keyed by + /// the file's location. + /// + /// This is a no-op unless a cache has been configured via + /// [`Self::with_file_metadata_cache`]. + fn cache_metadata(&self, metadata: Arc) -> Result<()> { + if let Some(file_metadata_cache) = &self.file_metadata_cache { + file_metadata_cache.put( + &self.object_meta.location, + CachedFileMetadataEntry::new( + self.object_meta.clone(), + Arc::new(CachedParquetMetaData::new(metadata)), + ), + ); } + Ok(()) + } + /// Fetch the full [`ParquetMetaData`] (including footer, and optional + /// page index) from the object store. + async fn fetch_metadata_from_store( + &self, + page_index_policy: PageIndexPolicy, + ) -> Result> { let file_size = self.object_meta.size; let mut decoder = ParquetMetaDataPushDecoder::try_new(file_size) .map_err(DataFusionError::from)?; @@ -153,14 +300,8 @@ impl<'a> DFParquetMetadata<'a> { .with_file_decryption_properties(Some(Arc::clone(decryption_properties))); } - if cache_metadata && self.file_metadata_cache.is_some() { - // Need to retrieve the entire metadata for the caching to be effective. - decoder = decoder.with_page_index_policy(PageIndexPolicy::Optional); - } else { - decoder = decoder.with_page_index_policy(PageIndexPolicy::Skip); - } + decoder = decoder.with_page_index_policy(page_index_policy); - // If we have a size hint, prefetch that many bytes from the end of the file if let Some(hint) = self.metadata_size_hint { let prefetch_start = file_size.saturating_sub(hint as u64); let prefetch_range = prefetch_start..file_size; @@ -199,22 +340,33 @@ impl<'a> DFParquetMetadata<'a> { } }; - let metadata = Arc::new(metadata); + Ok(Arc::new(metadata)) + } - if cache_metadata && let Some(file_metadata_cache) = &self.file_metadata_cache { - file_metadata_cache.put( - &self.object_meta.location, - CachedFileMetadataEntry::new( - self.object_meta.clone(), - Arc::new(CachedParquetMetaData::new(Arc::clone(&metadata))), - ), - ); + /// If `metadata` does not already have a page index, fetch and attach the + /// column and offset indexes. + async fn load_page_index( + store: &dyn ObjectStore, + object_meta: &ObjectMeta, + metadata: Arc, + ) -> Result> { + if metadata.column_index().is_some() && metadata.offset_index().is_some() { + return Ok(metadata); } - - Ok(metadata) + let metadata = + Arc::try_unwrap(metadata).unwrap_or_else(|shared| (*shared).clone()); + let mut reader = ParquetMetaDataReader::new_with_metadata(metadata) + .with_page_index_policy(PageIndexPolicy::Optional); + let fetch = ObjectStoreFetch::new(store, object_meta); + reader + .load_page_index(fetch) + .await + .map_err(DataFusionError::from)?; + Ok(Arc::new(reader.finish().map_err(DataFusionError::from)?)) } - /// Read and parse the schema of the Parquet file + /// Fetch this file's [`ParquetMetaData`] and convert its embedded Thrift + /// schema into an Arrow [`Schema`]. pub async fn fetch_schema(&self) -> Result { let metadata = self.fetch_metadata().await?; @@ -235,7 +387,8 @@ impl<'a> DFParquetMetadata<'a> { Ok(schema) } - /// Return (path, schema) tuple by fetching the schema from Parquet file + /// Convenience wrapper around [`Self::fetch_schema`] that also returns + /// the file's object store [`Path`]. pub(crate) async fn fetch_schema_with_location(&self) -> Result<(Path, Schema)> { let loc_path = self.object_meta.location.clone(); let schema = self.fetch_schema().await?; @@ -353,13 +506,12 @@ impl<'a> DFParquetMetadata<'a> { distinct_counts_array: &mut distinct_counts_array, }; summarize_column_statistics( - file_metadata.schema_descr(), logical_file_schema, - &physical_file_schema, &mut accumulators, idx, &stats_converter, row_groups_metadata, + num_rows, ) .ok(); } @@ -506,119 +658,208 @@ impl StatisticsAccumulators<'_> { } fn summarize_column_statistics( - parquet_schema: &SchemaDescriptor, logical_file_schema: &Schema, - physical_file_schema: &Schema, accumulators: &mut StatisticsAccumulators, logical_schema_index: usize, stats_converter: &StatisticsConverter, row_groups_metadata: &[RowGroupMetaData], + num_rows: usize, ) -> Result<()> { - let max_values = stats_converter.row_group_maxes(row_groups_metadata)?; - let min_values = stats_converter.row_group_mins(row_groups_metadata)?; - let null_counts = stats_converter.row_group_null_counts(row_groups_metadata)?; - let is_max_value_exact_stat = - stats_converter.row_group_is_max_value_exact(row_groups_metadata)?; - let is_min_value_exact_stat = - stats_converter.row_group_is_min_value_exact(row_groups_metadata)?; + let parquet_index = stats_converter.parquet_column_index(); if let Some(max_acc) = &mut accumulators.max_accs[logical_schema_index] { - max_acc.update_batch(&[Arc::clone(&max_values)])?; - - // handle the common special case when all row groups have exact statistics - let exactness = &is_max_value_exact_stat; - if !exactness.is_empty() && exactness.null_count() == 0 && !exactness.has_false() - { - accumulators.is_max_value_exact[logical_schema_index] = Some(true); - } else if !exactness.has_true() { - accumulators.is_max_value_exact[logical_schema_index] = Some(false); - } else { - let val = max_acc.evaluate()?; - accumulators.is_max_value_exact[logical_schema_index] = - has_any_exact_match(&val, &max_values, exactness); - } + accumulators.is_max_value_exact[logical_schema_index] = summarize_bound( + max_acc, + &stats_converter.row_group_maxes(row_groups_metadata)?, + parquet_index, + row_groups_metadata, + ParquetStatistics::max_is_exact, + || Ok(stats_converter.row_group_is_max_value_exact(row_groups_metadata)?), + )?; } if let Some(min_acc) = &mut accumulators.min_accs[logical_schema_index] { - min_acc.update_batch(&[Arc::clone(&min_values)])?; + accumulators.is_min_value_exact[logical_schema_index] = summarize_bound( + min_acc, + &stats_converter.row_group_mins(row_groups_metadata)?, + parquet_index, + row_groups_metadata, + ParquetStatistics::min_is_exact, + || Ok(stats_converter.row_group_is_min_value_exact(row_groups_metadata)?), + )?; + } + + accumulators.null_counts_array[logical_schema_index] = + summarize_null_counts(stats_converter, row_groups_metadata)?; + + accumulators.distinct_counts_array[logical_schema_index] = + summarize_distinct_counts(parquet_index, row_groups_metadata); + + let arrow_field = logical_file_schema.field(logical_schema_index); + accumulators.column_byte_sizes[logical_schema_index] = compute_arrow_column_size( + arrow_field.data_type(), + row_groups_metadata, + parquet_index, + num_rows, + ); + + Ok(()) +} - // handle the common special case when all row groups have exact statistics - let exactness = &is_min_value_exact_stat; - if !exactness.is_empty() && exactness.null_count() == 0 && !exactness.has_false() +/// Feed a column's per-row-group min or max `values` into `acc` and decide +/// whether the resulting bound is exact across all row groups. +/// +/// `is_exact` reads the per-row-group exactness flag straight from the raw +/// parquet statistics. `row_group_exactness` rebuilds the exactness as a Boolean +/// array and is only called for the rare case where row groups disagree. +fn summarize_bound( + acc: &mut A, + values: &ArrayRef, + parquet_index: Option, + row_groups_metadata: &[RowGroupMetaData], + is_exact: impl Fn(&ParquetStatistics) -> bool, + row_group_exactness: impl FnOnce() -> Result, +) -> Result> { + acc.update_batch(&[Arc::clone(values)])?; + + Ok( + match summarize_row_group_exactness(parquet_index, row_groups_metadata, is_exact) { - accumulators.is_min_value_exact[logical_schema_index] = Some(true); - } else if !exactness.has_true() { - accumulators.is_min_value_exact[logical_schema_index] = Some(false); - } else { - let val = min_acc.evaluate()?; - accumulators.is_min_value_exact[logical_schema_index] = - has_any_exact_match(&val, &min_values, exactness); - } + ExactnessSummary::AllExact => Some(true), + ExactnessSummary::NoneExact => Some(false), + ExactnessSummary::Mixed => { + let exactness = row_group_exactness()?; + has_any_exact_match(&acc.evaluate()?, values, &exactness) + } + }, + ) +} + +fn summarize_null_counts( + stats_converter: &StatisticsConverter, + row_groups_metadata: &[RowGroupMetaData], +) -> Result> { + if row_groups_metadata.is_empty() { + return Ok(Precision::Exact(0)); } - accumulators.null_counts_array[logical_schema_index] = match sum(&null_counts) { - Some(null_count) => Precision::Exact(null_count as usize), + let null_counts = stats_converter.row_group_null_counts(row_groups_metadata)?; + + match sum(&null_counts) { + Some(count) => { + // If any row group has an unknown null_count, either because column + // statistics are absent or because the null_count field is omitted, + // report the aggregate as inexact. + if null_counts.null_count() > 0 { + Ok(Precision::Inexact(count as usize)) + } else { + Ok(Precision::Exact(count as usize)) + } + } None => match null_counts.len() { // If sum() returned None we either have no rows or all values are null - 0 => Precision::Exact(0), - _ => Precision::Absent, + 0 => Ok(Precision::Exact(0)), + _ => Ok(Precision::Absent), }, + } +} + +#[derive(Debug, Copy, Clone, PartialEq, Eq)] +enum ExactnessSummary { + AllExact, + NoneExact, + Mixed, +} + +fn summarize_row_group_exactness( + parquet_idx: Option, + row_groups_metadata: &[RowGroupMetaData], + exactness: impl Fn(&ParquetStatistics) -> bool, +) -> ExactnessSummary { + let Some(parquet_idx) = parquet_idx else { + return ExactnessSummary::NoneExact; }; - // This is the same logic as parquet_column but we start from arrow schema index - // instead of looking up by name. - let parquet_index = parquet_column( - parquet_schema, - physical_file_schema, - logical_file_schema.field(logical_schema_index).name(), - ) - .map(|(idx, _)| idx); + summarize_exactness(row_groups_metadata.iter().map(|row_group| { + row_group + .columns() + .get(parquet_idx) + .and_then(|column| column.statistics()) + .map(&exactness) + })) +} - // Extract distinct counts from row group column statistics - accumulators.distinct_counts_array[logical_schema_index] = - if let Some(parquet_idx) = parquet_index { - let num_row_groups = row_groups_metadata.len(); - let distinct_counts: Vec = row_groups_metadata - .iter() - .filter_map(|rg| { - rg.columns() - .get(parquet_idx) - .and_then(|col| col.statistics()) - .and_then(|stats| stats.distinct_count_opt()) - }) - .collect(); +fn summarize_exactness(exactness: I) -> ExactnessSummary +where + I: IntoIterator>, +{ + let mut has_true = false; + let mut has_false_or_null = false; + + for exactness in exactness { + match exactness { + Some(true) => has_true = true, + Some(false) | None => has_false_or_null = true, + } - let coverage = distinct_counts.len() as f64 / num_row_groups.max(1) as f64; + if has_true && has_false_or_null { + return ExactnessSummary::Mixed; + } + } - if coverage < PARTIAL_NDV_THRESHOLD { - Precision::Absent - } else if distinct_counts.len() == 1 && num_row_groups == 1 { - // Single row group with distinct count - use exact value - Precision::Exact(distinct_counts[0] as usize) - } else { - // Multiple row groups - use max as a lower bound estimate - // (can't accurately merge NDV since duplicates may exist across row groups) - match distinct_counts.iter().max() { - Some(&max_ndv) => Precision::Inexact(max_ndv as usize), - None => Precision::Absent, - } - } - } else { - Precision::Absent - }; + if has_true { + ExactnessSummary::AllExact + } else { + ExactnessSummary::NoneExact + } +} - let arrow_field = logical_file_schema.field(logical_schema_index); - accumulators.column_byte_sizes[logical_schema_index] = compute_arrow_column_size( - arrow_field.data_type(), - row_groups_metadata, - parquet_index, - row_groups_metadata - .iter() - .map(|rg| rg.num_rows() as usize) - .sum(), - ); +/// Extract distinct counts from row group column statistics. +fn summarize_distinct_counts( + parquet_idx: Option, + row_groups_metadata: &[RowGroupMetaData], +) -> Precision { + let Some(parquet_idx) = parquet_idx else { + return Precision::Absent; + }; - Ok(()) + let num_row_groups = row_groups_metadata.len(); + if num_row_groups == 0 { + return Precision::Absent; + } + + let required_count = (num_row_groups as f64 * PARTIAL_NDV_THRESHOLD).ceil() as usize; + let mut ndv_count = 0; + let mut max_distinct_count: Option = None; + + for (row_group_idx, row_group) in row_groups_metadata.iter().enumerate() { + if let Some(distinct_count) = row_group + .columns() + .get(parquet_idx) + .and_then(|col| col.statistics()) + .and_then(|stats| stats.distinct_count_opt()) + { + ndv_count += 1; + max_distinct_count = Some(match max_distinct_count { + Some(max) => max.max(distinct_count), + None => distinct_count, + }); + } + + // Return early if there's no chance to reach the required coverage. + let remaining = num_row_groups - row_group_idx - 1; + if ndv_count + remaining < required_count { + return Precision::Absent; + } + } + + match max_distinct_count { + Some(distinct_count) if num_row_groups == 1 => { + Precision::Exact(distinct_count as usize) + } + Some(distinct_count) => Precision::Inexact(distinct_count as usize), + None => Precision::Absent, + } } /// Compute the Arrow in-memory size for a single column @@ -712,10 +953,15 @@ impl FileMetadata for CachedParquetMetaData { /// Convert a [`PhysicalSortExpr`] to a Parquet [`SortingColumn`]. /// -/// Returns `Err` if the expression is not a simple column reference. +/// Returns `Ok(None)` if the referenced column is not in `writer_schema`, such as a +/// hive partition column that is removed before writing the Parquet file. Returns +/// `Err` if the expression is not a simple column reference or references a column +/// outside `input_schema`. pub(crate) fn sort_expr_to_sorting_column( sort_expr: &PhysicalSortExpr, -) -> Result { + input_schema: &Schema, + writer_schema: &Schema, +) -> Result> { let column = sort_expr.expr.downcast_ref::().ok_or_else(|| { DataFusionError::Plan(format!( "Parquet sorting_columns only supports simple column references, \ @@ -724,27 +970,49 @@ pub(crate) fn sort_expr_to_sorting_column( )) })?; - let column_idx: i32 = column.index().try_into().map_err(|_| { + let input_field = input_schema.fields().get(column.index()).ok_or_else(|| { + internal_datafusion_err!( + "Parquet sorting column '{}' references index {} but the input schema has {} columns", + column.name(), + column.index(), + input_schema.fields().len() + ) + })?; + let Some((writer_index, _)) = writer_schema.column_with_name(input_field.name()) + else { + return Ok(None); + }; + + let column_idx: i32 = writer_index.try_into().map_err(|_| { DataFusionError::Plan(format!( - "Column index {} is too large to be represented as i32", - column.index() + "Column index {writer_index} is too large to be represented as i32" )) })?; - Ok(SortingColumn { + Ok(Some(SortingColumn { column_idx, descending: sort_expr.options.descending, nulls_first: sort_expr.options.nulls_first, - }) + })) } /// Convert a [`LexOrdering`] to `Vec` for Parquet. /// -/// Returns `Err` if any expression is not a simple column reference. +/// Columns that are not present in `writer_schema` are omitted from the resulting +/// metadata. Returns `Err` if any expression is not a simple column reference or +/// references a column outside `input_schema`. pub(crate) fn lex_ordering_to_sorting_columns( ordering: &LexOrdering, + input_schema: &Schema, + writer_schema: &Schema, ) -> Result> { - ordering.iter().map(sort_expr_to_sorting_column).collect() + ordering + .iter() + .filter_map(|sort_expr| { + sort_expr_to_sorting_column(sort_expr, input_schema, writer_schema) + .transpose() + }) + .collect() } /// Extracts ordering information from Parquet metadata. @@ -818,6 +1086,57 @@ fn sorting_columns_to_physical_exprs( mod tests { use super::*; use arrow::array::Int32Array; + use arrow::compute::SortOptions; + use arrow::datatypes::Field; + + #[test] + fn test_lex_ordering_to_sorting_columns_uses_writer_schema() -> Result<()> { + let input_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("part", DataType::Utf8, true), + Field::new("b", DataType::Utf8, true), + ]); + let writer_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ]); + let ordering = LexOrdering::new(vec![ + PhysicalSortExpr::new( + Arc::new(Column::new("part", 1)), + SortOptions::default(), + ), + PhysicalSortExpr::new(Arc::new(Column::new("a", 0)), SortOptions::default()), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 2)), + SortOptions { + descending: true, + nulls_first: false, + }, + ), + ]) + .unwrap(); + + let sorting_columns = + lex_ordering_to_sorting_columns(&ordering, &input_schema, &writer_schema)?; + + assert_eq!( + sorting_columns, + vec![ + SortingColumn { + column_idx: 0, + descending: false, + nulls_first: true, + }, + SortingColumn { + column_idx: 1, + descending: true, + nulls_first: false, + }, + ] + ); + + Ok(()) + } #[test] fn test_has_any_exact_match() { @@ -866,6 +1185,30 @@ mod tests { } } + #[test] + fn test_summarize_exactness() { + assert_eq!( + summarize_exactness([Some(true), Some(true)]), + ExactnessSummary::AllExact + ); + assert_eq!( + summarize_exactness([Some(false), None]), + ExactnessSummary::NoneExact + ); + assert_eq!( + summarize_exactness([Some(true), Some(false)]), + ExactnessSummary::Mixed + ); + assert_eq!( + summarize_exactness([Some(true), None]), + ExactnessSummary::Mixed + ); + assert_eq!( + summarize_exactness(std::iter::empty()), + ExactnessSummary::NoneExact + ); + } + mod ndv_tests { use super::*; use arrow::datatypes::Field; @@ -951,6 +1294,92 @@ mod tests { ParquetMetaData::new(file_meta, row_groups) } + #[test] + fn test_summarize_null_counts() { + let schema_descr = create_schema_descr(1); + let arrow_schema = create_arrow_schema(2); + let stats_with_count = + ParquetStatistics::int32(Some(1), Some(10), None, Some(2), false); + let stats_without_count = + ParquetStatistics::int32(Some(1), Some(10), None, None, false); + + let row_groups = vec![ + create_row_group_with_stats( + &schema_descr, + vec![Some(stats_with_count)], + 10, + ), + create_row_group_with_stats( + &schema_descr, + vec![Some(stats_without_count.clone())], + 10, + ), + create_row_group_with_stats(&schema_descr, vec![None], 10), + ]; + let stats_converter = + StatisticsConverter::try_new("col_0", &arrow_schema, &schema_descr) + .unwrap(); + let missing_column_converter = + StatisticsConverter::try_new("col_1", &arrow_schema, &schema_descr) + .unwrap(); + + assert_eq!( + summarize_null_counts(&stats_converter, &row_groups).unwrap(), + Precision::Inexact(2) + ); + assert_eq!( + summarize_null_counts(&missing_column_converter, &row_groups).unwrap(), + Precision::Absent + ); + assert_eq!( + summarize_null_counts(&stats_converter, &[]).unwrap(), + Precision::Exact(0) + ); + assert_eq!( + summarize_null_counts(&missing_column_converter, &[]).unwrap(), + Precision::Exact(0) + ); + + let missing_counts_unknown_converter = + StatisticsConverter::try_new("col_0", &arrow_schema, &schema_descr) + .unwrap() + .with_missing_null_counts_as_zero(false); + assert_eq!( + summarize_null_counts(&missing_counts_unknown_converter, &row_groups) + .unwrap(), + Precision::Inexact(2) + ); + + let row_groups_without_count = vec![ + create_row_group_with_stats( + &schema_descr, + vec![Some(stats_without_count.clone())], + 10, + ), + create_row_group_with_stats( + &schema_descr, + vec![Some(stats_without_count)], + 10, + ), + ]; + assert_eq!( + summarize_null_counts(&stats_converter, &row_groups_without_count) + .unwrap(), + Precision::Exact(0) + ); + + let missing_counts_unknown_converter = + stats_converter.with_missing_null_counts_as_zero(false); + assert_eq!( + summarize_null_counts( + &missing_counts_unknown_converter, + &row_groups_without_count, + ) + .unwrap(), + Precision::Absent + ); + } + #[test] fn test_distinct_count_single_row_group_with_ndv() { // Single row group with distinct count should return Exact diff --git a/datafusion/datasource-parquet/src/metrics.rs b/datafusion/datasource-parquet/src/metrics.rs index 262dde024a527..cbdcb73196b17 100644 --- a/datafusion/datasource-parquet/src/metrics.rs +++ b/datafusion/datasource-parquet/src/metrics.rs @@ -15,9 +15,11 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use datafusion_physical_plan::metrics::{ - Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, MetricCategory, MetricType, - PruningMetrics, RatioMergeStrategy, RatioMetrics, Time, + Count, ExecutionPlanMetricsSet, Gauge, Label, MetricBuilder, MetricCategory, + MetricType, PruningMetrics, RatioMergeStrategy, RatioMetrics, Time, }; /// Stores metrics about the parquet execution for a particular parquet file. @@ -51,6 +53,14 @@ pub struct ParquetFileMetrics { pub limit_pruned_row_groups: PruningMetrics, /// Number of row groups pruned by statistics pub row_groups_pruned_statistics: PruningMetrics, + /// Number of row groups pruned at runtime by a dynamic predicate + /// (e.g. the threshold expression a TopK `SortExec` pushes down). + /// + /// Unlike [`Self::row_groups_pruned_statistics`], which is decided once + /// at access-plan time, this counter reflects row groups that survived + /// the initial pruning but were proved unreachable mid-scan after the + /// dynamic filter tightened. + pub row_groups_pruned_dynamic_filter: Count, /// Total number of bytes scanned pub bytes_scanned: Count, /// Total rows filtered out by predicates pushed into parquet scan @@ -100,37 +110,42 @@ impl ParquetFileMetrics { filename: &str, metrics: &ExecutionPlanMetricsSet, ) -> Self { + // Share the filename label across all per-file metrics to avoid + // allocating the same filename string for each metric. + let filename_label = Label::new("filename", Arc::::from(filename)); + let builder = MetricBuilder::new(metrics).with_label(filename_label); + // ----------------------- // 'summary' level metrics // ----------------------- - let row_groups_pruned_bloom_filter = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let row_groups_pruned_bloom_filter = builder + .clone() .with_type(MetricType::Summary) .pruning_metrics("row_groups_pruned_bloom_filter", partition); - let limit_pruned_row_groups = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let limit_pruned_row_groups = builder + .clone() .with_type(MetricType::Summary) .pruning_metrics("limit_pruned_row_groups", partition); - let row_groups_pruned_statistics = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let row_groups_pruned_statistics = builder + .clone() .with_type(MetricType::Summary) .pruning_metrics("row_groups_pruned_statistics", partition); - let page_index_pages_pruned = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let page_index_pages_pruned = builder + .clone() .with_type(MetricType::Summary) .pruning_metrics("page_index_pages_pruned", partition); - let bytes_scanned = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let bytes_scanned = builder + .clone() .with_type(MetricType::Summary) .with_category(MetricCategory::Bytes) .counter("bytes_scanned", partition); - let metadata_load_time = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let metadata_load_time = builder + .clone() .with_type(MetricType::Summary) .subset_time("metadata_load_time", partition); @@ -138,8 +153,8 @@ impl ParquetFileMetrics { .with_type(MetricType::Summary) .pruning_metrics("files_ranges_pruned_statistics", partition); - let scan_efficiency_ratio = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let scan_efficiency_ratio = builder + .clone() .with_type(MetricType::Summary) .ratio_metrics_with_strategy( "scan_efficiency_ratio", @@ -150,48 +165,52 @@ impl ParquetFileMetrics { // ----------------------- // 'dev' level metrics // ----------------------- - let predicate_evaluation_errors = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let predicate_evaluation_errors = builder + .clone() .with_category(MetricCategory::Rows) .counter("predicate_evaluation_errors", partition); - let pushdown_rows_pruned = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let pushdown_rows_pruned = builder + .clone() .with_category(MetricCategory::Rows) .counter("pushdown_rows_pruned", partition); - let pushdown_rows_matched = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let pushdown_rows_matched = builder + .clone() .with_category(MetricCategory::Rows) .counter("pushdown_rows_matched", partition); - let row_pushdown_eval_time = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let row_pushdown_eval_time = builder + .clone() .subset_time("row_pushdown_eval_time", partition); - let statistics_eval_time = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let statistics_eval_time = builder + .clone() .subset_time("statistics_eval_time", partition); - let bloom_filter_eval_time = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let bloom_filter_eval_time = builder + .clone() .subset_time("bloom_filter_eval_time", partition); - let page_index_eval_time = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let page_index_eval_time = builder + .clone() .subset_time("page_index_eval_time", partition); - let page_index_rows_pruned = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let page_index_rows_pruned = builder + .clone() .pruning_metrics("page_index_rows_pruned", partition); - let predicate_cache_inner_records = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let predicate_cache_inner_records = builder + .clone() .with_category(MetricCategory::Rows) .gauge("predicate_cache_inner_records", partition); - let predicate_cache_records = MetricBuilder::new(metrics) - .with_new_label("filename", filename.to_string()) + let predicate_cache_records = builder .with_category(MetricCategory::Rows) .gauge("predicate_cache_records", partition); + let row_groups_pruned_dynamic_filter = MetricBuilder::new(metrics) + .with_new_label("filename", filename.to_string()) + .with_type(MetricType::Summary) + .counter("row_groups_pruned_dynamic_filter", partition); + Self { files_ranges_pruned_statistics, predicate_evaluation_errors, @@ -211,6 +230,7 @@ impl ParquetFileMetrics { scan_efficiency_ratio, predicate_cache_inner_records, predicate_cache_records, + row_groups_pruned_dynamic_filter, } } @@ -237,4 +257,23 @@ impl ParquetFileMetrics { .counter("page_index_pages_skipped_by_fully_matched", partition); count.add(n); } + + /// Record that page index I/O was skipped because row-group statistics + /// already proved page index could not prune further. + pub(crate) fn add_page_index_load_skipped( + metrics: &ExecutionPlanMetricsSet, + partition: usize, + filename: &str, + n: usize, + ) { + if n == 0 { + return; + } + + let count = MetricBuilder::new(metrics) + .with_new_label("filename", filename.to_string()) + .with_type(MetricType::Summary) + .counter("page_index_load_skipped", partition); + count.add(n); + } } diff --git a/datafusion/datasource-parquet/src/mod.rs b/datafusion/datasource-parquet/src/mod.rs index 6da0cfc4c5371..35f831230b305 100644 --- a/datafusion/datasource-parquet/src/mod.rs +++ b/datafusion/datasource-parquet/src/mod.rs @@ -26,11 +26,14 @@ pub mod access_plan; mod bloom_filter; +mod decoder_projection; pub mod file_format; pub mod metadata; mod metrics; +mod nested_schema_pruning; mod opener; mod page_filter; +mod projection_read_plan; mod push_decoder; mod reader; mod row_filter; @@ -42,9 +45,11 @@ pub mod source; mod supported_predicates; #[cfg(test)] mod test_util; +mod virtual_column; mod writer; -pub use access_plan::{ParquetAccessPlan, RowGroupAccess}; +pub use access_plan::{ParquetAccessPlan, ParquetRowSelection, RowGroupAccess}; +pub use bloom_filter::BloomFilterStatistics; pub use file_format::*; pub use metrics::ParquetFileMetrics; pub use page_filter::PagePruningAccessPlanFilter; @@ -53,10 +58,11 @@ pub use row_filter::build_row_filter; pub use row_filter::can_expr_be_pushed_down_with_schemas; pub use row_group_filter::RowGroupAccessPlanFilter; #[expect(deprecated)] +pub use schema_coercion::coerce_int96_to_resolution; pub use schema_coercion::{ - Int96Coercer, apply_file_schema_type_coercions, coerce_file_schema_to_string_type, - coerce_file_schema_to_view_type, coerce_int96_to_resolution, - transform_binary_to_string, transform_schema_to_view, + Int96Coercer, apply_file_schema_type_coercions, transform_binary_to_string, + transform_schema_to_view, }; pub use sink::ParquetSink; +pub use virtual_column::ParquetVirtualColumn; pub use writer::plan_to_parquet; diff --git a/datafusion/datasource-parquet/src/nested_schema_pruning.rs b/datafusion/datasource-parquet/src/nested_schema_pruning.rs new file mode 100644 index 0000000000000..9768282c3bab0 --- /dev/null +++ b/datafusion/datasource-parquet/src/nested_schema_pruning.rs @@ -0,0 +1,775 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Schema-driven nested projection pruning. +//! +//! When a scan's projection consumes a nested column only through a cast to a +//! *narrower* nested type, for example the file contains +//! `events: List>` but the expression is +//! `CAST(events AS List>)`, the Parquet reader does not need to +//! fetch or decode the leaves the cast target never names. This module +//! computes which Parquet leaves survive such a cast, and the Arrow type the +//! reader will emit for them, by walking the physical and target type trees +//! in parallel and matching struct fields by name (the equivalent of Spark's +//! `ParquetReadSupport.clipParquetSchema`). +//! +//! This situation arises whenever a table's logical schema declares a nested +//! column narrower than the physical Parquet file: the physical expression +//! adapter rewrites the projected column into exactly such a whole-column +//! cast (see `datafusion_physical_expr_adapter`). Engines like Spark +//! communicate nested projection pruning to the scan this way, as a clipped +//! read *schema* rather than as `get_field` expressions. +//! +//! # Safety of clipping +//! +//! The runtime cast for nested types +//! ([`datafusion_common::nested_struct::cast_column`]) consumes source struct +//! children exclusively by looking up the *target* field names, recursively +//! through list wrappers. Physical subtrees not named by the target are +//! provably dead: removing them from the read cannot change the cast's +//! output. That holds for *any* +//! [`CastExpr`](datafusion_physical_expr::expressions::CastExpr) over a +//! nested type, not just the ones the schema adapter inserts: +//! `ColumnarValue::cast_to` routes every +//! cast for which +//! [`requires_nested_struct_cast`](datafusion_common::nested_struct::requires_nested_struct_cast) +//! holds, the same predicate the projection analysis gates on, through +//! `cast_column`. +//! +//! Struct-level nullability is preserved because the Parquet reader +//! reconstructs ancestor validity from the definition levels of any surviving +//! leaf, so every struct level that is clipped must keep at least one leaf. +//! A struct cast with zero field-name overlap at *any* nesting depth would +//! break that: the reader drops a field whose leaves are all masked out, so +//! the emitted type would not match the one predicted here. Such a cast is +//! rejected during physical planning +//! (`datafusion_common::nested_struct::validate_struct_compatibility`, called +//! recursively from `DefaultPhysicalExprAdapter::rewrite`) and by the logical +//! planner's own castability check, so it should never reach this module; if +//! one does anyway (a custom `PhysicalExprAdapter` could build one), +//! [`clip_for_cast`] detects the empty level and declines to clip. +//! +//! The clip is *total*: any type shape it does not understand (maps, +//! dictionaries, wrapper-kind mismatches, ...) keeps all of its leaves, so +//! the worst case is today's behavior of reading the full column. Map values +//! are deliberately not clipped: the runtime cast routes maps through Arrow's +//! positional struct cast, which requires all children to be present. Nor are +//! `ListView`/`LargeListView`/`Dictionary` wrappers clipped here, even though +//! `cast_column` does recurse through them by name. That is a conservative +//! choice (safe, since the worst case is still just a full read) left as a +//! candidate follow-up rather than something this module currently handles. + +use std::collections::HashMap; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, FieldRef, Fields}; + +/// The single child type one level of container nesting wraps, or `None` for +/// a type this module does not descend through (leaves, `Struct`, `Map`, and +/// wrapper kinds this module intentionally does not clip, see the module +/// doc). Shared by [`count_leaves`] and [`contains_struct`], which otherwise +/// need to agree on the exact same set of container variants. +fn nested_child(dt: &DataType) -> Option<&DataType> { + match dt { + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::FixedSizeList(f, _) + | DataType::Map(f, _) => Some(f.data_type()), + DataType::Dictionary(_, value) => Some(value), + DataType::RunEndEncoded(_, value) => Some(value.data_type()), + _ => None, + } +} + +/// Clip `physical` against `cast_target`, returning the Parquet leaves the +/// cast actually consumes (as offsets relative to the root column's first +/// leaf, sorted ascending and non-empty) together with the Arrow type the +/// reader will emit for exactly those leaves. +/// +/// Returns `None` when nothing can be pruned (every leaf is consumed, or the +/// shapes do not allow safe clipping), in which case the caller should read +/// the whole column as before. This function never fails: unknown shapes +/// degrade to keeping all leaves. +pub(crate) fn clip_for_cast( + physical: &DataType, + cast_target: &DataType, +) -> Option<(Vec, DataType)> { + let total = count_leaves(physical); + let mut kept = Vec::new(); + let mut next_leaf = 0; + let mut unclippable = false; + let pruned_type = clip_type( + physical, + cast_target, + &mut next_leaf, + &mut kept, + &mut unclippable, + ); + debug_assert_eq!(next_leaf, total, "leaf accounting must cover the type"); + if unclippable || kept.is_empty() || kept.len() >= total { + return None; + } + Some((kept, pruned_type)) +} + +/// Number of Parquet leaf columns a (Parquet-derived) Arrow type occupies. +pub(crate) fn count_leaves(dt: &DataType) -> usize { + match dt { + DataType::Struct(fields) => { + fields.iter().map(|f| count_leaves(f.data_type())).sum() + } + _ => nested_child(dt).map_or(1, count_leaves), + } +} + +/// Does this type contain a struct at any nesting depth? Used as a fast-path +/// gate: a root with no struct anywhere in its type has no leaves this +/// module could ever clip. +pub(crate) fn contains_struct(dt: &DataType) -> bool { + matches!(dt, DataType::Struct(_)) || nested_child(dt).is_some_and(contains_struct) +} + +/// Above this many target fields, matching physical children against them one +/// by one turns into a quadratic string comparison; build a name lookup +/// instead. Below it the map's allocation costs more than the linear scan it +/// saves (Spark's `ParquetReadSupport.clipParquetGroupFields` builds the map +/// unconditionally; struct widths in practice are small enough that the +/// threshold is worth the branch). +const LINEAR_FIELD_SCAN_MAX: usize = 8; + +/// Find `name` among `fields`, using `by_name` when it was worth building. +/// Duplicate names resolve to the first occurrence either way. +fn lookup_field<'a>( + fields: &'a Fields, + by_name: &Option>, + name: &str, +) -> Option<&'a FieldRef> { + match by_name { + Some(map) => map.get(name).copied(), + None => fields.iter().find(|f| f.name() == name), + } +} + +/// Recursive walker: advances `next_leaf` across every leaf of `physical`, +/// pushing the offsets the cast target consumes into `kept`, and returns the +/// Arrow type the reader emits for those kept leaves. +/// +/// `unclippable` is set when a shape is encountered whose emitted type this +/// module cannot predict; the caller must then read the whole column. The walk +/// still runs to completion so `next_leaf` stays a valid leaf count. +fn clip_type( + physical: &DataType, + target: &DataType, + next_leaf: &mut usize, + kept: &mut Vec, + unclippable: &mut bool, +) -> DataType { + match (physical, target) { + (DataType::Struct(p_children), DataType::Struct(t_children)) => { + let t_by_name = (t_children.len() > LINEAR_FIELD_SCAN_MAX).then(|| { + let mut map = HashMap::with_capacity(t_children.len()); + for tc in t_children.iter() { + map.entry(tc.name().as_str()).or_insert(tc); + } + map + }); + let kept_children: Fields = p_children + .iter() + .filter_map(|pc| { + let Some(tc) = lookup_field(t_children, &t_by_name, pc.name()) else { + skip_leaves(pc.data_type(), next_leaf); + return None; + }; + let before = kept.len(); + let pruned = clip_type( + pc.data_type(), + tc.data_type(), + next_leaf, + kept, + unclippable, + ); + if kept.len() == before { + // This child matched by name but kept no leaves at + // all, which only happens when a nested struct level + // below it shares no field name with its target. The + // reader drops a field whose leaves are all masked + // out, so the emitted type could not be predicted; + // give up on clipping this column entirely rather + // than promise a type the decoder will not produce. + // (`DefaultPhysicalExprAdapter` never builds such a + // cast — `validate_struct_compatibility` rejects a + // zero-overlap struct level at planning time — but a + // custom `PhysicalExprAdapter` could.) + *unclippable = true; + } + Some(field_with_type(pc, pruned)) + }) + .collect(); + DataType::Struct(kept_children) + } + (DataType::List(p_item), DataType::List(t_item)) => { + let pruned = clip_type( + p_item.data_type(), + t_item.data_type(), + next_leaf, + kept, + unclippable, + ); + DataType::List(field_with_type(p_item, pruned)) + } + (DataType::LargeList(p_item), DataType::LargeList(t_item)) => { + let pruned = clip_type( + p_item.data_type(), + t_item.data_type(), + next_leaf, + kept, + unclippable, + ); + DataType::LargeList(field_with_type(p_item, pruned)) + } + // Anything else, leaf pairs, wrapper-kind mismatches, maps, + // dictionaries, fixed-size lists, views, is kept wholesale. + _ => keep_all_leaves(physical, next_leaf, kept), + } +} + +/// Keep every leaf of `dt` (no pruning below this point); returns `dt` +/// unchanged since nothing was clipped. +fn keep_all_leaves( + dt: &DataType, + next_leaf: &mut usize, + kept: &mut Vec, +) -> DataType { + let n = count_leaves(dt); + kept.extend(*next_leaf..*next_leaf + n); + *next_leaf += n; + dt.clone() +} + +fn skip_leaves(dt: &DataType, next_leaf: &mut usize) { + *next_leaf += count_leaves(dt); +} + +/// A projected root column that is consumed through a cast to a narrower +/// nested type (`CAST(col AS target_type)`), recorded during projection +/// analysis. +#[derive(Debug, Clone)] +pub(crate) struct CastColumnAccess { + /// Arrow root column index of the column in the file schema. + pub(crate) root_index: usize, + /// The cast's target type. + pub(crate) target_type: DataType, +} + +/// Rebuild `field` with a new data type, preserving name, nullability and +/// metadata. +pub(crate) fn field_with_type(field: &Field, data_type: DataType) -> FieldRef { + Arc::new(field.clone().with_data_type(data_type)) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn utf8(name: &str) -> Field { + Field::new(name, DataType::Utf8, true) + } + + fn int64(name: &str) -> Field { + Field::new(name, DataType::Int64, true) + } + + fn struct_of(fields: Vec) -> DataType { + DataType::Struct(Fields::from(fields)) + } + + fn list_of(item: DataType) -> DataType { + DataType::List(Arc::new(Field::new("item", item, true))) + } + + #[test] + fn count_leaves_shapes() { + assert_eq!(count_leaves(&DataType::Int32), 1); + assert_eq!(count_leaves(&struct_of(vec![utf8("a"), int64("b")])), 2); + assert_eq!( + count_leaves(&list_of(struct_of(vec![ + utf8("a"), + struct_of(vec![int64("x"), int64("y")]).into_field("s") + ]))), + 3 + ); + let map = DataType::Map( + Arc::new(Field::new( + "entries", + struct_of(vec![utf8("key"), int64("value")]), + false, + )), + false, + ); + assert_eq!(count_leaves(&map), 2); + let dict = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + assert_eq!(count_leaves(&dict), 1); + // Wrapper kinds must be descended through, not counted as one leaf. + // A dictionary or run-end-encoded *value* that is itself a struct has + // as many leaves as the struct: counting it as 1 would misalign every + // later leaf index in the mask. + assert_eq!( + count_leaves(&DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(struct_of(vec![utf8("a"), int64("b")])) + )), + 2 + ); + assert_eq!( + count_leaves(&DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new( + "values", + struct_of(vec![utf8("a"), int64("b")]), + true + )) + )), + 2 + ); + } + + /// [`contains_struct`] gates the projection fast path, so it has to agree + /// with [`count_leaves`] about which wrappers are descended through. + #[test] + fn contains_struct_shapes() { + assert!(!contains_struct(&DataType::Int32)); + assert!(!contains_struct(&list_of(DataType::Int32))); + assert!(contains_struct(&struct_of(vec![int64("a")]))); + assert!(contains_struct(&list_of(struct_of(vec![int64("a")])))); + assert!(contains_struct(&DataType::LargeList(Arc::new(Field::new( + "item", + struct_of(vec![int64("a")]), + true + ))))); + assert!(contains_struct(&DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(struct_of(vec![int64("a")])) + ))); + assert!(!contains_struct(&DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(DataType::Utf8) + ))); + // A map's entries are a struct, so a map always contains one. + assert!(contains_struct(&DataType::Map( + Arc::new(Field::new( + "entries", + struct_of(vec![utf8("key"), int64("value")]), + false + )), + false + ))); + } + + /// `{a, b, c} CAST TO {b}` keeps only b's leaf. + #[test] + fn clip_struct_subset() { + let physical = struct_of(vec![utf8("a"), int64("b"), utf8("c")]); + let target = struct_of(vec![int64("b")]); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![1]); + assert_eq!(emitted, struct_of(vec![int64("b")])); + } + + /// Target field order does not matter: emitted type is in physical order. + #[test] + fn clip_struct_reordered_target() { + let physical = struct_of(vec![utf8("a"), int64("b"), utf8("c")]); + let target = struct_of(vec![utf8("c"), utf8("a")]); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 2]); + assert_eq!(emitted, struct_of(vec![utf8("a"), utf8("c")])); + } + + /// Target fields missing from the physical type are ignored (the runtime + /// cast null-fills them). + #[test] + fn clip_struct_target_field_missing_from_physical() { + let physical = struct_of(vec![utf8("a"), int64("b")]); + let target = struct_of(vec![utf8("a"), int64("z")]); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted, struct_of(vec![utf8("a")])); + } + + /// Leaf-level type mismatch (promotion) still clips: the emitted type + /// keeps the physical leaf type; the cast performs the promotion. + #[test] + fn clip_keeps_physical_leaf_types() { + let physical = + struct_of(vec![Field::new("x", DataType::Int32, true), utf8("pad")]); + let target = struct_of(vec![int64("x")]); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!( + emitted, + struct_of(vec![Field::new("x", DataType::Int32, true)]) + ); + } + + /// Nested struct-in-struct clips at both levels. + #[test] + fn clip_nested_struct() { + let inner_physical = struct_of(vec![int64("x"), utf8("pad_inner")]); + let physical = struct_of(vec![ + inner_physical.clone().into_field("inner"), + utf8("pad_outer"), + ]); + let target = struct_of(vec![struct_of(vec![int64("x")]).into_field("inner")]); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!( + emitted, + struct_of(vec![struct_of(vec![int64("x")]).into_field("inner")]) + ); + } + + /// List, the headline case. + #[test] + fn clip_list_of_struct() { + let physical = list_of(struct_of(vec![int64("x"), utf8("y"), utf8("pad")])); + let target = list_of(struct_of(vec![int64("x"), utf8("y")])); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 1]); + assert_eq!(emitted, list_of(struct_of(vec![int64("x"), utf8("y")]))); + } + + /// Two levels of `list` nesting, the inner one also narrowed, + /// the `events: array>>>` shape + /// reported in `datafusion-comet#4859`, where a sibling struct field at + /// the outer level (`aux`, standing in for that report's + /// `latency_parts`) is dropped entirely rather than clipped. + #[test] + fn clip_two_level_nested_list_of_struct() { + let physical = list_of(struct_of(vec![ + int64("a"), + utf8("pad"), + struct_of(vec![int64("x"), utf8("y")]).into_field("aux"), + list_of(struct_of(vec![int64("g"), utf8("pad2")])).into_field("items"), + ])); + let target = list_of(struct_of(vec![ + int64("a"), + list_of(struct_of(vec![int64("g")])).into_field("items"), + ])); + + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + // a=0, pad=1, aux.x=2, aux.y=3, items.g=4, items.pad2=5: only a and + // items.g survive; pad, all of aux, and items.pad2 are dropped. + assert_eq!(kept, vec![0, 4]); + assert_eq!( + emitted, + list_of(struct_of(vec![ + int64("a"), + list_of(struct_of(vec![int64("g")])).into_field("items"), + ])) + ); + } + + #[test] + fn clip_large_list_of_struct() { + let item = |fields| Arc::new(Field::new("item", struct_of(fields), true)); + let physical = DataType::LargeList(item(vec![int64("x"), utf8("pad")])); + let target = DataType::LargeList(item(vec![int64("x")])); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0]); + assert_eq!(emitted, DataType::LargeList(item(vec![int64("x")]))); + } + + /// Wrapper-kind mismatch cannot be clipped. + #[test] + fn no_clip_on_wrapper_mismatch() { + let physical = list_of(struct_of(vec![int64("x"), utf8("pad")])); + let target = DataType::LargeList(Arc::new(Field::new( + "item", + struct_of(vec![int64("x")]), + true, + ))); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Maps are opaque: never clipped. + #[test] + fn no_clip_on_map() { + let entries = |fields| Arc::new(Field::new("entries", struct_of(fields), false)); + let physical = + DataType::Map(entries(vec![utf8("key"), int64("a"), int64("b")]), false); + let target = DataType::Map(entries(vec![utf8("key"), int64("a")]), false); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Identical types: nothing to prune. + #[test] + fn no_clip_when_identical() { + let t = struct_of(vec![utf8("a"), int64("b")]); + assert!(clip_for_cast(&t, &t).is_none()); + } + + /// Non-nested types: nothing to prune. + #[test] + fn no_clip_on_primitives() { + assert!(clip_for_cast(&DataType::Int32, &DataType::Int64).is_none()); + } + + /// A struct level with zero field-name overlap can't actually reach this + /// code: `validate_struct_compatibility` rejects it during physical + /// planning (see the module doc), so `clip_for_cast` is only ever called + /// with targets that overlap at every nesting level. If it were reached + /// anyway, the generic catch-all keeps every leaf, still safe, just + /// unpruned. + #[test] + fn no_clip_on_zero_overlap() { + let physical = struct_of(vec![utf8("a"), int64("b")]); + let target = struct_of(vec![utf8("z")]); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// A *nested* struct level with zero field-name overlap must not be + /// clipped, even when a sibling keeps leaves. The reader drops a field + /// whose leaves are all masked out (pinned by + /// [`reader_drops_struct_child_with_no_selected_leaves`]), so predicting + /// `{inner: Struct[], c}` here would be a schema the decoder never + /// produces. Read the whole column instead. + #[test] + fn no_clip_when_nested_struct_level_has_no_overlap() { + let physical = struct_of(vec![ + struct_of(vec![int64("a"), int64("b")]).into_field("inner"), + int64("c"), + ]); + let target = struct_of(vec![ + struct_of(vec![int64("z")]).into_field("inner"), + int64("c"), + ]); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Same, one level deeper and behind a list wrapper. + #[test] + fn no_clip_when_nested_list_struct_level_has_no_overlap() { + let physical = struct_of(vec![ + list_of(struct_of(vec![int64("a"), int64("b")])).into_field("items"), + int64("c"), + ]); + let target = struct_of(vec![ + list_of(struct_of(vec![int64("z")])).into_field("items"), + int64("c"), + ]); + assert!(clip_for_cast(&physical, &target).is_none()); + } + + /// Wide structs take the name-map matching path rather than the linear + /// scan; both must produce the same clip. + #[test] + fn clip_wide_struct_matches_by_name() { + let width = LINEAR_FIELD_SCAN_MAX * 4; + let physical = struct_of((0..width).map(|i| int64(&format!("f{i}"))).collect()); + // Even fields only, declared in reverse order: the emitted type is + // still in physical order. + let target = struct_of( + (0..width) + .rev() + .filter(|i| i % 2 == 0) + .map(|i| int64(&format!("f{i}"))) + .collect(), + ); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, (0..width).filter(|i| i % 2 == 0).collect::>()); + assert_eq!( + emitted, + struct_of( + (0..width) + .filter(|i| i % 2 == 0) + .map(|i| int64(&format!("f{i}"))) + .collect() + ) + ); + } + + /// Duplicate physical field names both match the single target field and + /// are both kept, which is what the reader emits for that mask. + #[test] + fn clip_keeps_duplicate_physical_field_names() { + let physical = struct_of(vec![int64("a"), utf8("pad"), int64("a")]); + let target = struct_of(vec![int64("a")]); + let (kept, emitted) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 2]); + assert_eq!(emitted, struct_of(vec![int64("a"), int64("a")])); + } + + /// Pins the arrow-rs behavior the empty-level guard above depends on: a + /// struct child none of whose leaves are selected disappears from the + /// type the reader emits, rather than surviving as an empty struct. + #[test] + fn reader_drops_struct_child_with_no_selected_leaves() { + use arrow::array::{ArrayRef, Int64Array, StructArray}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::arrow::{ArrowWriter, ProjectionMask}; + + let inner_fields = Fields::from(vec![int64("a"), int64("b")]); + let outer_fields = Fields::from(vec![ + Field::new("inner", DataType::Struct(inner_fields.clone()), true), + int64("c"), + ]); + let inner: ArrayRef = Arc::new(StructArray::new( + inner_fields, + vec![ + Arc::new(Int64Array::from(vec![1, 2])) as ArrayRef, + Arc::new(Int64Array::from(vec![3, 4])) as ArrayRef, + ], + None, + )); + let outer = StructArray::new( + outer_fields.clone(), + vec![inner, Arc::new(Int64Array::from(vec![5, 6])) as ArrayRef], + None, + ); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new( + "s", + DataType::Struct(outer_fields), + true, + )])); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(outer)]).unwrap(); + + let file = tempfile::NamedTempFile::new().unwrap(); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + let builder = + ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()).unwrap(); + assert_eq!(builder.parquet_schema().num_columns(), 3); + // Keep only s.c (leaf 2): every leaf of s.inner is masked out. + let mask = ProjectionMask::leaves(builder.parquet_schema(), [2usize]); + let reader = builder.with_projection(mask).build().unwrap(); + let out: Vec = reader.map(|b| b.unwrap()).collect(); + assert_eq!( + out[0].schema().field(0).data_type(), + &struct_of(vec![int64("c")]), + "the fully masked `inner` child is dropped, not emitted as an empty struct" + ); + } + + /// Pins the arrow-rs behavior this module relies on: selecting a subset + /// of leaves under a `List` column with `ProjectionMask::leaves` + /// makes the reader emit exactly the type predicted by [`clip_for_cast`], + /// and null list rows / null struct elements survive (their validity is + /// reconstructed from the surviving leaves' definition levels). + #[test] + fn arrow_reader_emits_clipped_type_for_masked_list_struct() { + use arrow::array::{ + Array, ArrayRef, Int64Array, ListArray, StringArray, StructArray, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer}; + use arrow::record_batch::RecordBatch; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::arrow::{ArrowWriter, ProjectionMask}; + + let item_fields = Fields::from(vec![int64("x"), utf8("y"), utf8("pad")]); + let item_field = Arc::new(Field::new( + "item", + DataType::Struct(item_fields.clone()), + true, + )); + let schema = Arc::new(arrow::datatypes::Schema::new(vec![Field::new( + "events", + DataType::List(Arc::clone(&item_field)), + true, + )])); + + // 3 elements; element 1 is a NULL struct. Rows: [e0, e1], NULL, [e2]. + let columns: Vec = vec![ + Arc::new(Int64Array::from(vec![Some(1), None, Some(3)])), + Arc::new(StringArray::from(vec![Some("a"), None, Some("c")])), + Arc::new(StringArray::from(vec![Some("p0"), None, Some("p2")])), + ]; + let struct_validity = NullBuffer::from(vec![true, false, true]); + let values = StructArray::new(item_fields, columns, Some(struct_validity)); + let list_validity = NullBuffer::from(vec![true, false, true]); + let events = ListArray::new( + item_field, + OffsetBuffer::from_lengths([2, 0, 1]), + Arc::new(values), + Some(list_validity), + ); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(events)]).unwrap(); + + let file = tempfile::NamedTempFile::new().unwrap(); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), schema, None).unwrap(); + writer.write(&batch).unwrap(); + writer.close().unwrap(); + + // Clip to the narrow target {x, y}. + let physical = batch.schema().field(0).data_type().clone(); + let target = list_of(struct_of(vec![int64("x"), utf8("y")])); + let (kept, predicted_type) = clip_for_cast(&physical, &target).unwrap(); + assert_eq!(kept, vec![0, 1]); + + let builder = + ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()).unwrap(); + let mask = ProjectionMask::leaves(builder.parquet_schema(), kept.iter().copied()); + let reader = builder.with_projection(mask).build().unwrap(); + let out: Vec = reader.map(|b| b.unwrap()).collect(); + assert_eq!(out.len(), 1); + let out = &out[0]; + + // Emitted type matches the prediction. + assert_eq!(out.schema().field(0).data_type(), &predicted_type); + + // Null semantics survive the clip. + let events = out.column(0).as_any().downcast_ref::().unwrap(); + assert!(events.is_valid(0)); + assert!(events.is_null(1)); + assert!(events.is_valid(2)); + let structs = events + .values() + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(structs.len(), 3); + assert!(structs.is_valid(0)); + assert!(structs.is_null(1)); + assert!(structs.is_valid(2)); + let x = structs + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(x.value(0), 1); + assert_eq!(x.value(2), 3); + } + + trait IntoField { + fn into_field(self, name: &str) -> Field; + } + + impl IntoField for DataType { + fn into_field(self, name: &str) -> Field { + Field::new(name, self, true) + } + } +} diff --git a/datafusion/datasource-parquet/src/opener/encryption.rs b/datafusion/datasource-parquet/src/opener/encryption.rs index b725198237bbf..498fe8acf7530 100644 --- a/datafusion/datasource-parquet/src/opener/encryption.rs +++ b/datafusion/datasource-parquet/src/opener/encryption.rs @@ -76,6 +76,7 @@ impl EncryptionContext { #[cfg(not(feature = "parquet_encryption"))] #[expect(dead_code)] +#[expect(clippy::unused_async)] impl EncryptionContext { pub(super) async fn get_file_decryption_properties( &self, diff --git a/datafusion/datasource-parquet/src/opener/mod.rs b/datafusion/datasource-parquet/src/opener/mod.rs index 95e0516e8bc27..693e9bd2cbf31 100644 --- a/datafusion/datasource-parquet/src/opener/mod.rs +++ b/datafusion/datasource-parquet/src/opener/mod.rs @@ -24,42 +24,48 @@ use self::early_stop::EarlyStoppingStream; #[cfg(feature = "parquet_encryption")] use self::encryption::EncryptionContext; use crate::access_plan::PreparedAccessPlan; +use crate::decoder_projection::DecoderProjection; use crate::page_filter::PagePruningAccessPlanFilter; -use crate::push_decoder::{DecoderBuilderConfig, PushDecoderStreamState}; -use crate::row_filter::{RowFilterGenerator, build_projection_read_plan}; -use crate::row_group_filter::{BloomFilterStatistics, RowGroupAccessPlanFilter}; +use crate::push_decoder::{ + DecoderBuilderConfig, PushDecoderStreamState, RgPlanEntry, RowGroupPruner, +}; +use crate::row_filter::RowFilterGenerator; +use crate::row_group_filter::RowGroupAccessPlanFilter; use crate::{ - Int96Coercer, ParquetAccessPlan, ParquetFileMetrics, ParquetFileReaderFactory, + BloomFilterStatistics, Int96Coercer, ParquetAccessPlan, ParquetFileMetrics, + ParquetFileReaderFactory, ParquetRowSelection, ParquetVirtualColumn, apply_file_schema_type_coercions, }; use arrow::array::RecordBatch; use arrow::datatypes::DataType; use datafusion_datasource::morsel::{Morsel, MorselPlan, MorselPlanner, Morselizer}; use datafusion_physical_expr::projection::ProjectionExprs; -use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr_adapter::replace_columns_with_literals; +use datafusion_physical_expr_adapter::rewrite::rewrite_input_file_name_in_projection; use std::collections::{HashMap, VecDeque}; use std::fmt; use std::future::Future; use std::mem; use std::sync::Arc; -use arrow::datatypes::{SchemaRef, TimeUnit}; +use arrow::datatypes::{FieldRef, Schema, SchemaRef, TimeUnit}; #[cfg(feature = "parquet_encryption")] use datafusion_common::encryption::FileDecryptionProperties; use datafusion_common::stats::Precision; -use datafusion_common::{ColumnStatistics, Result, ScalarValue, Statistics, exec_err}; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; +use datafusion_common::{ + ColumnStatistics, HashSet, Result, ScalarValue, Statistics, exec_err, internal_err, +}; use datafusion_datasource::{PartitionedFile, TableSchema}; +use datafusion_physical_expr::expressions::{Column, DynamicFilterTracking}; use datafusion_physical_expr::simplifier::PhysicalExprSimplifier; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; -use datafusion_physical_expr_common::physical_expr::{ - PhysicalExpr, is_dynamic_physical_expr, -}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::LexOrdering; use datafusion_physical_plan::metrics::{ BaselineMetrics, Count, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, }; -use datafusion_pruning::{FilePruner, PruningPredicate, build_pruning_predicate}; +use datafusion_pruning::{FilePruner, PruningPredicate, PruningPredicateBuilder}; #[cfg(feature = "parquet_encryption")] use datafusion_common::config::EncryptionFactoryOptions; @@ -74,7 +80,153 @@ use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::parquet_column; use parquet::basic::Type; use parquet::bloom_filter::Sbbf; -use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader}; +use parquet::file::metadata::{PageIndexPolicy, ParquetMetaDataReader, RowGroupMetaData}; + +/// Morselizer-level state for virtual columns, precomputed once per scan +/// partition so each file skips the validator walks, `null_replacements` +/// rebuild, and one of the `append_fields` allocations. +/// +/// Only constructed when the scan actually requests virtual columns; +/// [`ParquetMorselizer`] and [`PreparedParquetOpen`] hold +/// `Option>` so the zero-virtual-column path (the +/// common case) pays nothing. +pub(crate) struct VirtualColumnsState { + /// Shared list of virtual column fields. Cloned as a `Vec` only at the + /// arrow-rs `with_virtual_columns` call site, which takes it by value. + virtual_columns: Arc>, + /// Null-literal substitutions keyed by virtual column name, used to strip + /// virtual-column references from the projection fed into + /// `build_projection_read_plan` (which only understands file columns). + null_replacements: HashMap, + /// `logical_file_schema` with the virtual columns appended. Fed into the + /// per-file expression rewriter so virtual-column references + /// identity-rewrite instead of being replaced with null literals. + logical_schema_with_virtual: SchemaRef, +} + +impl VirtualColumnsState { + /// Validate each field carries a supported arrow virtual extension type + /// and precompute the per-scan derived state. + fn try_new( + virtual_columns: Vec, + logical_file_schema: &SchemaRef, + ) -> Result { + // Gate which extension types we forward to arrow-rs. Adding a new + // supported virtual column means adding a `ParquetVirtualColumn` + // variant — not editing a stringly-typed allowlist here. + for field in &virtual_columns { + ParquetVirtualColumn::try_from(field)?; + } + let null_replacements = virtual_columns + .iter() + .map(|f| ScalarValue::try_from(f.data_type()).map(|v| (f.name().clone(), v))) + .collect::>>()?; + let logical_schema_with_virtual = + append_fields(logical_file_schema, &virtual_columns); + Ok(Self { + virtual_columns: Arc::new(virtual_columns), + null_replacements, + logical_schema_with_virtual, + }) + } + + /// Validated virtual column fields, in declaration order. + pub(crate) fn virtual_columns(&self) -> &[FieldRef] { + &self.virtual_columns + } + + /// Null-literal substitutions keyed by virtual column name. Used to strip + /// virtual-column references from a projection before it is fed into the + /// parquet `ProjectionMask` (which only understands file columns). + pub(crate) fn null_replacements(&self) -> &HashMap { + &self.null_replacements + } +} + +/// Build the per-scan virtual-column state. +/// +/// Two checks run here: +/// - Extension-type allowlist via [`VirtualColumnsState::try_new`]: returns +/// `Err` for unsupported virtual extension types. +/// - Predicate-reference check (when pushdown is enabled): returns `Err` if +/// the predicate references a virtual column. The contract is that callers +/// route filters through +/// [`ParquetSource::try_pushdown_filters`](crate::source::ParquetSource), +/// which classifies virtual-col filters as `PushedDown::No`. Erroring here +/// prevents silent wrong results for callers that bypass that path and set +/// the predicate directly on `ParquetSource`. +/// +/// Returns `None` when the scan has no virtual columns, so callers avoid +/// allocating the shared state on the common path. +pub(crate) fn build_virtual_columns_state( + virtual_columns: &[FieldRef], + logical_file_schema: &SchemaRef, + predicate: Option<&Arc>, + pushdown_filters: bool, +) -> Result>> { + if virtual_columns.is_empty() { + return Ok(None); + } + if pushdown_filters && let Some(predicate) = predicate { + validate_predicate_does_not_reference_virtual_columns( + predicate, + virtual_columns, + )?; + } + let state = + VirtualColumnsState::try_new(virtual_columns.to_vec(), logical_file_schema)?; + Ok(Some(Arc::new(state))) +} + +/// Return `base` unchanged when `extra` is empty; otherwise build a new schema +/// with `extra` appended to `base`'s fields. +pub(crate) fn append_fields(base: &SchemaRef, extra: &[FieldRef]) -> SchemaRef { + if extra.is_empty() { + return Arc::clone(base); + } + let fields = base + .fields() + .iter() + .cloned() + .chain(extra.iter().cloned()) + .collect::>(); + Arc::new(Schema::new(fields)) +} + +/// Reject predicates that reference a virtual column. +/// +/// arrow-rs's `RowFilter` evaluates predicates against a `ProjectionMask` that +/// addresses parquet leaves only; virtual columns (e.g. `row_number`) are +/// synthesized by the reader *after* filter evaluation and cannot be referenced +/// inside a row filter. Silently dropping such a predicate would produce wrong +/// results. +fn validate_predicate_does_not_reference_virtual_columns( + predicate: &Arc, + virtual_columns: &[FieldRef], +) -> Result<()> { + if virtual_columns.is_empty() { + return Ok(()); + } + let virtual_names: HashSet<&str> = + virtual_columns.iter().map(|f| f.name().as_str()).collect(); + let mut offender: Option = None; + predicate.apply(|node: &Arc| { + if let Some(column) = node.downcast_ref::() + && virtual_names.contains(column.name()) + { + offender = Some(column.name().to_string()); + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + })?; + if let Some(name) = offender { + return internal_err!( + "Predicate references virtual column '{name}'; route via \ + ParquetSource::try_pushdown_filters." + ); + } + Ok(()) +} /// Stateless Parquet morselizer implementation. /// @@ -137,10 +289,18 @@ pub(super) struct ParquetMorselizer { /// Maximum size of the predicate cache, in bytes. If none, uses /// the arrow-rs default. pub max_predicate_cache_size: Option, + /// Maximum `IN (...)` list size that the pruning predicate will rewrite + /// into per-value statistics checks. Lists longer than this skip + /// container-level pruning. Sourced from + /// `datafusion.execution.parquet.max_in_list_size`. + pub max_in_list_size: usize, /// Whether to read row groups in reverse order pub reverse_row_groups: bool, /// Optional sort order used to reorder row groups by their min/max statistics. pub sort_order_for_reorder: Option, + /// Per-scan virtual-column state (validation already performed). `None` + /// when no virtual columns are requested — the common path. + pub(crate) virtual_state: Option>, } impl fmt::Debug for ParquetMorselizer { @@ -183,10 +343,10 @@ impl Morselizer for ParquetMorselizer { /// PrepareFilters /// | /// v -/// LoadPageIndex +/// PruneWithStatistics /// | /// v -/// PruneWithStatistics +/// LoadPageIndex? (skipped when all surviving row groups are fully matched) /// | /// v /// LoadBloomFilters @@ -221,10 +381,10 @@ enum ParquetOpenState { /// Specialize any filters for the actual file schema (only known after /// metadata is loaded) PrepareFilters(Box), - /// Loading [Parquet Page Index](https://parquet.apache.org/docs/file-format/pageindex/) - LoadPageIndex(BoxFuture<'static, Result>), /// Pruning Row Groups PruneWithStatistics(Box), + /// Loading [Parquet Page Index](https://parquet.apache.org/docs/file-format/pageindex/) + LoadPageIndex(BoxFuture<'static, Result>), /// Loading bloom filters required for row-group pruning LoadBloomFilters(BoxFuture<'static, Result>), /// Pruning with preloaded Bloom Filters @@ -279,6 +439,11 @@ struct PreparedParquetOpen { output_schema: SchemaRef, projection: ProjectionExprs, predicate: Option>, + /// Per-scan virtual-column state, Arc-cloned from [`ParquetMorselizer`] so + /// each file shares validated fields, precomputed null replacements, and + /// the logical-with-virtual schema. `None` when no virtual columns were + /// requested. + virtual_state: Option>, reorder_predicates: bool, pushdown_filters: bool, force_filter_selections: bool, @@ -291,6 +456,7 @@ struct PreparedParquetOpen { expr_adapter_factory: Arc, predicate_creation_errors: Count, max_predicate_cache_size: Option, + max_in_list_size: usize, reverse_row_groups: bool, sort_order_for_reorder: Option, preserve_order: bool, @@ -384,19 +550,39 @@ impl ParquetOpenState { } ParquetOpenState::PrepareFilters(loaded) => { let prepared_filters = loaded.prepare_filters()?; - Ok(ParquetOpenState::LoadPageIndex( - prepared_filters.load_page_index().boxed(), - )) + Ok(ParquetOpenState::PruneWithStatistics(Box::new( + prepared_filters, + ))) + } + ParquetOpenState::PruneWithStatistics(prepared) => { + let prepared_row_groups = (*prepared).prune_row_groups()?; + if prepared_row_groups.should_load_page_index() { + Ok(ParquetOpenState::LoadPageIndex( + prepared_row_groups.load_page_index().boxed(), + )) + } else { + if prepared_row_groups + .prepared + .page_pruning_predicate + .is_some() + && !prepared_row_groups.row_groups.is_empty() + { + let prepared = &prepared_row_groups.prepared.loaded.prepared; + ParquetFileMetrics::add_page_index_load_skipped( + &prepared.metrics, + prepared.partition_index, + &prepared.file_name, + 1, + ); + } + Ok(ParquetOpenState::LoadBloomFilters( + prepared_row_groups.load_bloom_filters().boxed(), + )) + } } ParquetOpenState::LoadPageIndex(future) => { Ok(ParquetOpenState::LoadPageIndex(future)) } - ParquetOpenState::PruneWithStatistics(prepared) => { - let prepared_row_groups = prepared.prune_row_groups()?; - Ok(ParquetOpenState::LoadBloomFilters( - prepared_row_groups.load_bloom_filters().boxed(), - )) - } ParquetOpenState::LoadBloomFilters(future) => { Ok(ParquetOpenState::LoadBloomFilters(future)) } @@ -510,9 +696,9 @@ impl MorselPlanner for ParquetMorselPlanner { } ParquetOpenState::LoadPageIndex(future) => { Ok(Some(Self::schedule_io(async move { - Ok(ParquetOpenState::PruneWithStatistics(Box::new( - future.await?, - ))) + Ok(ParquetOpenState::LoadBloomFilters( + future.await?.load_bloom_filters().boxed(), + )) }))) } ParquetOpenState::LoadBloomFilters(future) => { @@ -614,22 +800,26 @@ impl ParquetMorselizer { .transpose()?; } + // Replace any `input_file_name()` UDFs in the projection with a literal for this file. + projection = rewrite_input_file_name_in_projection(projection, &file_name)?; + let predicate_creation_errors = MetricBuilder::new(&self.metrics) .with_category(MetricCategory::Rows) .global_counter("num_predicate_creation_errors"); - // Apply literal replacements to projection and predicate - let file_pruner = predicate - .as_ref() - .filter(|p| is_dynamic_physical_expr(p) || partitioned_file.has_statistics()) - .and_then(|p| { - FilePruner::try_new( - Arc::clone(p), - &logical_file_schema, - &partitioned_file, - predicate_creation_errors.clone(), - ) - }); + // `FilePruner::try_new` decides whether a pruner is worthwhile (it needs + // a statistics struct, and either real column statistics or a dynamic + // filter that can prune via partition-value folding) and returns `None` + // otherwise. For a static predicate the pruner's tracker reports no + // changes, so it runs once and adds no ongoing cost. + let file_pruner = predicate.as_ref().and_then(|p| { + FilePruner::try_new( + Arc::clone(p), + &logical_file_schema, + &partitioned_file, + predicate_creation_errors.clone(), + ) + }); Ok(PreparedParquetOpen { partition_index: self.partition_index, @@ -650,6 +840,7 @@ impl ParquetMorselizer { output_schema, projection, predicate, + virtual_state: self.virtual_state.as_ref().map(Arc::clone), reorder_predicates: self.reorder_filters, pushdown_filters: self.pushdown_filters, force_filter_selections: self.force_filter_selections, @@ -662,6 +853,7 @@ impl ParquetMorselizer { expr_adapter_factory: Arc::clone(&self.expr_adapter_factory), predicate_creation_errors, max_predicate_cache_size: self.max_predicate_cache_size, + max_in_list_size: self.max_in_list_size, reverse_row_groups: self.reverse_row_groups, sort_order_for_reorder: self.sort_order_for_reorder.clone(), preserve_order: self.preserve_order, @@ -677,30 +869,21 @@ impl PreparedParquetOpen { /// Returns `None` if the file can be skipped completely. fn prune_file(mut self) -> Result> { // Prune this file using the file level statistics and partition values. - // Since dynamic filters may have been updated since planning it is possible that we are able - // to prune files now that we couldn't prune at planning time. - // It is assumed that there is no point in doing pruning here if the predicate is not dynamic, - // as it would have been done at planning time. - // We'll also check this after every record batch we read, - // and if at some point we are able to prove we can prune the file using just the file level statistics - // we can end the stream early. - // - // Make a FilePruner only if there is either - // 1. a dynamic expr in the predicate - // 2. the file has file-level statistics. - // - // File-level statistics may prune the file without loading - // any row groups or metadata. + // Since dynamic filters may have been updated since planning it is + // possible that we are able to prune files now that we couldn't prune at + // planning time. The `FilePruner` (built when the predicate is dynamic or + // the file carries statistics) also watches any still-active dynamic + // filter, so the + // `EarlyStoppingStream` wrapping the scan can re-check after each batch + // and end the stream early once a tightened filter proves the file can + // be skipped. // - // Dynamic filters may prune the file after initial - // planning, as the dynamic filter is updated during - // execution. - // - // The case where there is a dynamic filter but no - // statistics corresponds to a dynamic filter that - // references partition columns. While rare, this is possible - // e.g. `select * from table order by partition_col limit - // 10` could hit this condition. + // File-level statistics may prune the file without loading any row + // groups or metadata. Partition column predicates are already folded to + // literals (see `replace_columns_with_literals` above), so a dynamic + // filter that references only partition columns can prune here too even + // when the file has no column statistics, e.g. + // `select * from t order by partition_col limit 10`. if let Some(file_pruner) = &mut self.file_pruner && file_pruner.should_prune()? { @@ -723,8 +906,11 @@ impl PreparedParquetOpen { // unnecessary I/O. We decide later if it is needed to evaluate the // pruning predicates. Thus default to not requesting it from the // underlying reader. - let options = + let mut options = ArrowReaderOptions::new().with_page_index_policy(PageIndexPolicy::Skip); + if let Some(schema) = self.partitioned_file.arrow_schema.as_ref() { + options = options.with_schema(Arc::clone(schema)); + } #[cfg(feature = "parquet_encryption")] let mut options = options; #[cfg(feature = "parquet_encryption")] @@ -767,22 +953,22 @@ impl MetadataLoadedParquetOpen { // - The logical file schema: this is the table schema minus any hive partition columns and projections. // This is what the physical file schema is coerced to. // - The physical file schema: this is the schema that the arrow-rs - // parquet reader will actually produce. + // parquet reader will actually produce for the file's columns. Any + // virtual columns (see [`crate::TableSchema::virtual_columns`]) are + // produced separately by the reader and are not part of this schema. let mut physical_file_schema = Arc::clone(reader_metadata.schema()); // The schema loaded from the file may not be the same as the // desired schema (for example if we want to instruct the parquet // reader to read strings using Utf8View instead). Update if necessary + let mut metadata_dirty = false; if let Some(merged) = apply_file_schema_type_coercions( &prepared.logical_file_schema, &physical_file_schema, ) { physical_file_schema = Arc::new(merged); options = options.with_schema(Arc::clone(&physical_file_schema)); - reader_metadata = ArrowReaderMetadata::try_new( - Arc::clone(reader_metadata.metadata()), - options.clone(), - )?; + metadata_dirty = true; } if let Some(ref coerce) = prepared.coerce_int96 @@ -796,6 +982,17 @@ impl MetadataLoadedParquetOpen { { physical_file_schema = Arc::new(merged); options = options.with_schema(Arc::clone(&physical_file_schema)); + metadata_dirty = true; + } + + // Arrow-rs appends virtual columns to the supplied schema internally, + // so any `with_schema` coercion above must stay limited to file columns. + if let Some(state) = prepared.virtual_state.as_ref() { + options = options.with_virtual_columns((*state.virtual_columns).clone())?; + metadata_dirty = true; + } + + if metadata_dirty { reader_metadata = ArrowReaderMetadata::try_new( Arc::clone(reader_metadata.metadata()), options.clone(), @@ -818,11 +1015,32 @@ impl MetadataLoadedParquetOpen { let needs_rewrite = prepared.predicate.is_some() || prepared.logical_file_schema != physical_file_schema; if needs_rewrite { + // When virtual columns are requested, augment the logical and + // physical schemas passed to the rewriter/simplifier with those + // fields. The rewriter identity-rewrites references found in both + // schemas, keeping virtual-column references as `Column` rather + // than replacing them with null literals; the simplifier needs + // them present so it can resolve their data types while walking + // expression trees. We keep `physical_file_schema` itself as the + // pure file schema so downstream predicate pushdown, pruning, and + // row filter construction stay unaffected. + let (logical_for_rewrite, physical_for_rewrite) = + if let Some(state) = prepared.virtual_state.as_ref() { + ( + Arc::clone(&state.logical_schema_with_virtual), + append_fields(&physical_file_schema, &state.virtual_columns), + ) + } else { + ( + Arc::clone(&prepared.logical_file_schema), + Arc::clone(&physical_file_schema), + ) + }; let rewriter = prepared.expr_adapter_factory.create( - Arc::clone(&prepared.logical_file_schema), - Arc::clone(&physical_file_schema), + Arc::clone(&logical_for_rewrite), + Arc::clone(&physical_for_rewrite), )?; - let simplifier = PhysicalExprSimplifier::new(&physical_file_schema); + let simplifier = PhysicalExprSimplifier::new(&physical_for_rewrite); prepared.predicate = prepared .predicate .map(|p| simplifier.simplify(rewriter.rewrite(p)?)) @@ -838,6 +1056,7 @@ impl MetadataLoadedParquetOpen { prepared.predicate.as_ref(), &physical_file_schema, &prepared.predicate_creation_errors, + prepared.max_in_list_size, ); // Only build page pruning predicate if page index is enabled @@ -863,27 +1082,6 @@ impl MetadataLoadedParquetOpen { } impl FiltersPreparedParquetOpen { - /// Load the page index if pruning requires it and metadata did not include it. - async fn load_page_index(mut self) -> Result { - // The page index is not stored inline in the parquet footer so the - // metadata load above may not have read the page index structures yet. - // If we need them for reading and they aren't yet loaded, we need to - // load them now. - if self.page_pruning_predicate.is_some() { - self.loaded.reader_metadata = load_page_index( - self.loaded.reader_metadata, - &mut self.loaded.prepared.async_file_reader, - self.loaded - .options - .clone() - .with_page_index_policy(PageIndexPolicy::Optional), - ) - .await?; - } - - Ok(self) - } - /// Prune row groups using file ranges and parquet metadata. fn prune_row_groups(self) -> Result { let loaded = &self.loaded; @@ -896,7 +1094,7 @@ impl FiltersPreparedParquetOpen { let mut row_groups = RowGroupAccessPlanFilter::new(create_initial_plan( &prepared.file_name, &prepared.extensions, - rg_metadata.len(), + rg_metadata, )?); // If there is a range restricting what parts of the file to read @@ -952,6 +1150,66 @@ impl FiltersPreparedParquetOpen { } impl RowGroupsPrunedParquetOpen { + /// Returns true if the reader would benefit from a page index load, given + /// the current pruning predicate and row group access plan. + /// + /// The page index is used for data page pruning, and it is only useful + /// when: + /// + /// 1. There is at least one row group that may have filtered rows + /// (if it is fully matched we know no rows will be filtered) + /// + /// 2. There is a page index for at least one predicate column (some + /// parquet writers do not write the page index). + fn should_load_page_index(&self) -> bool { + let Some(page_pruning_predicate) = self.prepared.page_pruning_predicate.as_ref() + else { + return false; + }; + let row_groups = &self.row_groups; + let fully_matched = row_groups.is_fully_matched(); + // if all row groups are fully matched, nothing can be pruned + if row_groups.row_group_indexes().all(|idx| fully_matched[idx]) { + return false; + } + + // Check the file's footer metadata to see if a page index was written + // for at least one predicate column in a surviving row group. + // + // Note: offsets are recorded in the footer, so we can determine if a + // page index exists before attempting to read it. + let parquet_metadata = self.prepared.loaded.reader_metadata.metadata(); + let arrow_schema = &self.prepared.loaded.prepared.physical_file_schema; + let parquet_schema = parquet_metadata.file_metadata().schema_descr(); + page_pruning_predicate.predicate_column_names().any(|name| { + let Some((leaf_idx, _)) = parquet_column(parquet_schema, arrow_schema, name) + else { + return false; + }; + row_groups.row_group_indexes().any(|rg_idx| { + let column = parquet_metadata.row_group(rg_idx).column(leaf_idx); + column.column_index_offset().is_some() + && column.offset_index_offset().is_some() + }) + }) + } + + /// Load the page index if pruning requires it and metadata did not include it. + async fn load_page_index(mut self) -> Result { + self.prepared.loaded.reader_metadata = load_page_index( + self.prepared.loaded.reader_metadata.clone(), + &mut self.prepared.loaded.prepared.async_file_reader, + self.prepared + .loaded + .options + .clone() + .with_page_index_policy(PageIndexPolicy::Optional), + ) + .await?; + + Ok(self) + } + /// Load bloom filters needed for pruning when enabled and a pruning predicate exists. async fn load_bloom_filters(mut self) -> Result { let num_row_groups = self @@ -986,7 +1244,7 @@ impl RowGroupsPrunedParquetOpen { mem::replace(&mut prepared.async_file_reader, replacement_reader), reader_metadata, ); - let parquet_columns: Vec<(String, usize, Type)> = predicate + let parquet_columns: Vec<(String, usize, Type, i32)> = predicate .literal_columns() .into_iter() .filter_map(|column_name| { @@ -1000,6 +1258,7 @@ impl RowGroupsPrunedParquetOpen { column_name, column_idx, parquet_schema.column(column_idx).physical_type(), + parquet_schema.column(column_idx).type_length(), )) }) .collect(); @@ -1007,7 +1266,9 @@ impl RowGroupsPrunedParquetOpen { for idx in self.row_groups.row_group_indexes() { let mut row_group_filters = BloomFilterStatistics::with_capacity(parquet_columns.len()); - for (column_name, column_idx, physical_type) in &parquet_columns { + for (column_name, column_idx, physical_type, type_length) in + &parquet_columns + { let bf: Sbbf = match builder .get_row_group_column_bloom_filter(idx, *column_idx) .await @@ -1020,7 +1281,12 @@ impl RowGroupsPrunedParquetOpen { continue; } }; - row_group_filters.insert(column_name, bf, *physical_type); + row_group_filters.insert( + column_name, + bf, + *physical_type, + *type_length, + ); } row_group_bloom_filters[idx] = row_group_filters; } @@ -1156,13 +1422,20 @@ impl RowGroupsPrunedParquetOpen { }; let arrow_reader_metrics = ArrowReaderMetrics::enabled(); - let read_plan = build_projection_read_plan( - prepared.projection.expr_iter(), + + // Build the decoder projection (mask + per-batch transform) in a + // single call. Encapsulating it behind `DecoderProjection` keeps the + // opener's orchestration body focused on filter / decoder / stream + // wiring. + let decoder_projection = DecoderProjection::try_new( + &prepared.projection, &prepared.physical_file_schema, reader_metadata.parquet_schema(), - ); + &prepared.output_schema, + prepared.virtual_state.as_deref(), + )?; - let (decoder, pending_decoders, remaining_limit) = { + let (decoder, rg_plan, has_row_selection) = { let pushdown_predicate = prepared .pushdown_filters .then_some(prepared.predicate.as_ref()) @@ -1175,50 +1448,53 @@ impl RowGroupsPrunedParquetOpen { &prepared.file_metrics, ); - // Split into consecutive runs of row groups that share the same filter - // requirement. Fully matched row groups skip the RowFilter; others need it. - // Reverse the run order for reverse scans so the combined decoder stream - // preserves the requested global row group order. - let mut runs = access_plan.split_runs(row_filter_generator.has_row_filter()); - if prepared.reverse_row_groups { - runs.reverse(); - } - let run_count = runs.len(); - let decoder_limit = prepared.limit.filter(|_| run_count == 1); - let remaining_limit = prepared.limit.filter(|_| run_count > 1); - + // Build the prepared access plan first — `prepare_access_plan` may + // call `reorder_by_statistics` (for `sort_order_for_reorder`) and + // `reverse` (for `reverse_row_groups`), both of which mutate + // `row_group_indexes` to the physical scan order the decoder will + // actually read. We MUST build our `rg_plan` from this reordered + // list, otherwise our per-RG pruner check would consult the + // metadata of a different RG than the decoder is about to yield. let decoder_config = DecoderBuilderConfig { - read_plan: &read_plan, + projection_mask: decoder_projection.projection_mask(), batch_size: prepared.batch_size, arrow_reader_metrics: &arrow_reader_metrics, force_filter_selections: prepared.force_filter_selections, - decoder_limit, + decoder_limit: prepared.limit, }; - // Build a decoder per run. - let mut decoders = VecDeque::with_capacity(runs.len()); - for run in runs { - let prepared_access_plan = prepare_access_plan(run.access_plan)?; - let mut builder = - decoder_config.build(prepared_access_plan, reader_metadata.clone()); - if run.needs_filter { - if let Some(row_filter) = row_filter_generator.next_filter() { - builder = builder.with_row_filter(row_filter); - } - if let Some(max_predicate_cache_size) = - prepared.max_predicate_cache_size - { - builder = builder - .with_max_predicate_cache_size(max_predicate_cache_size); - } + let prepared_access_plan = prepare_access_plan(access_plan)?; + // #24355: a row selection (from page-index pruning, or an externally + // supplied `ParquetRowSelection`) is carried by the decoder as one + // flat selection over the concatenation of the remaining row groups. + // The runtime pruner's `into_builder().with_row_groups(...)` rebuild + // drops row groups without slicing that selection to match, so record + // whether a selection is present and disable runtime pruning below + // when it is (mirroring `reorder_by_statistics`, which also bails when + // a row selection is present). The proper fix that keeps pruning + // under a live selection is tracked in + // https://github.com/apache/arrow-rs/issues/10624 / + // https://github.com/apache/datafusion/issues/24358. + let has_row_selection = prepared_access_plan.row_selection.is_some(); + let rg_plan: VecDeque = prepared_access_plan + .row_group_indexes + .iter() + .copied() + .map(|rg_index| RgPlanEntry { rg_index }) + .collect(); + + let mut builder = + decoder_config.build(prepared_access_plan, reader_metadata.clone()); + if let Some(row_filter) = row_filter_generator.next_filter() { + builder = builder.with_row_filter(row_filter); + if let Some(max_predicate_cache_size) = prepared.max_predicate_cache_size + { + builder = + builder.with_max_predicate_cache_size(max_predicate_cache_size); } - decoders.push_back(builder.build()?); } - let decoder = decoders - .pop_front() - .expect("at least one decoder must be created"); - (decoder, decoders, remaining_limit) + (builder.build()?, rg_plan, has_row_selection) }; let predicate_cache_inner_records = @@ -1226,46 +1502,79 @@ impl RowGroupsPrunedParquetOpen { let predicate_cache_records = prepared.file_metrics.predicate_cache_records.clone(); - // Check if we need to replace the schema to handle things like differing nullability or metadata. - // See note below about file vs. output schema. - let stream_schema = read_plan.projected_schema; - let replace_schema = stream_schema != prepared.output_schema; - - // Rebase column indices to match the narrowed stream schema. - // The projection expressions have indices based on physical_file_schema, - // but the stream only contains the columns selected by the ProjectionMask. - let projection = prepared - .projection - .try_map_exprs(|expr| reassign_expr_columns(expr, &stream_schema))?; - let projector = projection.make_projector(&stream_schema)?; - let output_schema = Arc::clone(&prepared.output_schema); let files_ranges_pruned_statistics = prepared.file_metrics.files_ranges_pruned_statistics.clone(); + + // Build a dynamic row-group pruner only when all three conditions hold: + // 1) the scan has a predicate (so there is something to evaluate), + // 2) the predicate has at least one not-yet-complete dynamic filter + // (`DynamicFilterTracking::Watching`) — static or already-complete + // predicates were fully consumed by `prune_by_statistics` at file + // open, so re-evaluating them per RG boundary would be wasted work, + // 3) there is at least one pending RG that could be skipped. + // The pruner subscribes once to every still-incomplete dynamic filter + // via the `DynamicFilterTracker` watch channel (#22460), so detecting + // a threshold change is a single atomic load — not a tree walk per + // RG check. + // Also disabled when a row selection is live (#24355) — page-index + // pruning is the common source: the pruner rebuilds the decoder via + // `with_row_groups(...)`, which drops row groups without slicing the + // carried selection to match, so pruning under a live selection returns + // wrong results. Decline to prune in that case. + let row_group_pruner = + match (&prepared.predicate, rg_plan.len() > 1, has_row_selection) { + (Some(predicate), true, false) + if matches!( + DynamicFilterTracking::classify(predicate), + DynamicFilterTracking::Watching(_) + ) => + { + Some(RowGroupPruner::new( + Arc::clone(predicate), + Arc::clone(&prepared.physical_file_schema), + Arc::clone(reader_metadata.metadata()), + prepared.predicate_creation_errors.clone(), + prepared.file_metrics.predicate_evaluation_errors.clone(), + prepared.max_in_list_size, + )) + } + _ => None, + }; + let row_groups_pruned_dynamic = prepared + .file_metrics + .row_groups_pruned_dynamic_filter + .clone(); + let stream = PushDecoderStreamState { - decoder, - pending_decoders, - remaining_limit, + decoder: Some(decoder), + active_reader: None, + rg_plan, reader: prepared.async_file_reader, - projector, - output_schema, - replace_schema, + decoder_projection, arrow_reader_metrics, predicate_cache_inner_records, predicate_cache_records, baseline_metrics: prepared.baseline_metrics, + row_group_pruner, + row_groups_pruned_dynamic, } .into_stream(); - // Wrap the stream so a dynamic filter can stop the file scan early. - if let Some(file_pruner) = prepared.file_pruner { - Ok(EarlyStoppingStream::new( - stream, - file_pruner, - files_ranges_pruned_statistics, - ) - .boxed()) - } else { - Ok(stream) + // Wrap the stream so a dynamic filter can stop the file scan early, but + // only when the pruner is still watching a filter that can change + // mid-scan. For a static (or already-complete) predicate the up-front + // `prune_file` check already captured everything that can be pruned, so + // per-batch re-checking would only add overhead. + match prepared.file_pruner { + Some(file_pruner) if file_pruner.is_watching() => { + Ok(EarlyStoppingStream::new( + stream, + file_pruner, + files_ranges_pruned_statistics, + ) + .boxed()) + } + _ => Ok(stream), } } } @@ -1334,29 +1643,44 @@ fn constant_value_from_stats( /// Return the initial [`ParquetAccessPlan`] /// -/// If the user has supplied one as an extension, use that -/// otherwise return a plan that scans all row groups +/// If the user has supplied a parquet access extension, use that; otherwise +/// return a plan that scans all row groups. /// -/// Returns an error if an invalid `ParquetAccessPlan` is provided +/// Returns an error if an invalid parquet access extension is provided. /// /// Note: file_name is only used for error messages fn create_initial_plan( file_name: &str, extensions: &datafusion_datasource::FileExtensions, - row_group_count: usize, + rg_metadata: &[RowGroupMetaData], ) -> Result { - if let Some(access_plan) = extensions.get::() { - let plan_len = access_plan.len(); - if plan_len != row_group_count { - return exec_err!( - "Invalid ParquetAccessPlan for {file_name}. Specified {plan_len} row groups, but file has {row_group_count}" - ); + let row_group_count = rg_metadata.len(); + match ( + extensions.get::(), + extensions.get::(), + ) { + (Some(_), Some(_)) => exec_err!( + "Invalid parquet access extensions for {file_name}. \ + Specify either ParquetAccessPlan or ParquetRowSelection, not both" + ), + (Some(access_plan), None) => { + let plan_len = access_plan.len(); + if plan_len != row_group_count { + return exec_err!( + "Invalid ParquetAccessPlan for {file_name}. Specified {plan_len} row groups, but file has {row_group_count}" + ); + } + Ok(access_plan.clone()) + } + (None, Some(row_selection)) => { + ParquetAccessPlan::try_new_from_overall_row_selection( + row_selection.selection().clone(), + rg_metadata, + ) } - return Ok(access_plan.clone()); + // default to scanning all row groups + (None, None) => Ok(ParquetAccessPlan::new_all(row_group_count)), } - - // default to scanning all row groups - Ok(ParquetAccessPlan::new_all(row_group_count)) } /// Build a page pruning predicate from an optional predicate expression. @@ -1376,13 +1700,14 @@ pub(crate) fn build_pruning_predicates( predicate: Option<&Arc>, file_schema: &SchemaRef, predicate_creation_errors: &Count, + max_in_list_size: usize, ) -> Option> { let predicate = predicate.as_ref()?; - build_pruning_predicate( - Arc::clone(predicate), - file_schema, - predicate_creation_errors, - ) + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(file_schema)) + .with_error_counter(predicate_creation_errors) + .with_max_in_list_size(max_in_list_size) + .build(Arc::clone(predicate)) } /// Returns a `ArrowReaderMetadata` with the page index loaded, loading @@ -1420,17 +1745,24 @@ async fn load_page_index( mod test { use super::*; use super::{ConstantColumns, ParquetMorselizer, constant_columns_from_stats}; - use crate::{DefaultParquetFileReaderFactory, RowGroupAccess}; - use arrow::array::RecordBatch; + use crate::{ + CachedParquetFileReaderFactory, DefaultParquetFileReaderFactory, + ParquetFileReaderFactory, ParquetRowSelection, RowGroupAccess, + }; + use arrow::array::{RecordBatch, record_batch}; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use bytes::{BufMut, BytesMut}; use datafusion_common::{ - ColumnStatistics, ScalarValue, Statistics, internal_err, record_batch, + ColumnStatistics, ScalarValue, Statistics, assert_contains, internal_err, stats::Precision, }; use datafusion_datasource::morsel::{Morsel, Morselizer}; - use datafusion_datasource::{PartitionedFile, TableSchema}; - use datafusion_expr::{col, lit}; + use datafusion_datasource::{PartitionedFile, TableSchema, TableSchemaBuilder}; + use datafusion_execution::cache::cache_manager::{ + CachedFileMetadataEntry, FileMetadataCache, + }; + use datafusion_execution::cache::default_cache::DefaultCache; + use datafusion_expr::{Expr, col, lit}; use datafusion_physical_expr::{ PhysicalExpr, expressions::{Column, DynamicFilterPhysicalExpr, Literal}, @@ -1441,11 +1773,14 @@ mod test { DefaultPhysicalExprAdapterFactory, replace_columns_with_literals, }; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use datafusion_pruning::MAX_IN_LIST_SIZE; use futures::StreamExt; use futures::stream::BoxStream; use object_store::{ObjectStore, ObjectStoreExt, memory::InMemory, path::Path}; - use parquet::arrow::ArrowWriter; + use parquet::arrow::{ArrowSchemaConverter, ArrowWriter}; + use parquet::file::metadata::{ColumnChunkMetaData, FileMetaData, ParquetMetaData}; use parquet::file::properties::WriterProperties; + use parquet::schema::types::SchemaDescPtr; use std::collections::VecDeque; use std::sync::Arc; @@ -1462,6 +1797,7 @@ mod test { predicate: Option>, metadata_size_hint: Option, metrics: ExecutionPlanMetricsSet, + parquet_file_reader_factory: Option>, pushdown_filters: bool, reorder_filters: bool, force_filter_selections: bool, @@ -1470,10 +1806,194 @@ mod test { enable_row_group_stats_pruning: bool, coerce_int96: Option, max_predicate_cache_size: Option, + max_in_list_size: usize, reverse_row_groups: bool, preserve_order: bool, } + #[test] + fn create_initial_plan_from_parquet_row_selection_extension() { + use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; + + let mut extensions = datafusion_datasource::FileExtensions::new(); + extensions.insert(ParquetRowSelection::new(RowSelection::from(vec![ + RowSelector::select(10), + RowSelector::skip(20), + RowSelector::select(30), + ]))); + let rg_metadata = row_group_metadata(&[10, 20, 30]); + + let access_plan = + create_initial_plan("test.parquet", &extensions, &rg_metadata).unwrap(); + + assert_eq!( + access_plan, + ParquetAccessPlan::new(vec![ + RowGroupAccess::Scan, + RowGroupAccess::Skip, + RowGroupAccess::Scan, + ]) + ); + } + + #[test] + fn create_initial_plan_rejects_multiple_access_extensions() { + use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; + + let mut extensions = datafusion_datasource::FileExtensions::new(); + extensions.insert(ParquetAccessPlan::new_all(3)); + extensions.insert(ParquetRowSelection::new(RowSelection::from(vec![ + RowSelector::select(60), + ]))); + let rg_metadata = row_group_metadata(&[10, 20, 30]); + + let err = create_initial_plan("test.parquet", &extensions, &rg_metadata) + .unwrap_err() + .to_string(); + + assert_contains!( + err, + "Specify either ParquetAccessPlan or ParquetRowSelection, not both" + ); + } + + fn row_group_metadata(row_counts: &[i64]) -> Vec { + let schema_descr = test_schema_descr(); + + row_counts + .iter() + .map(|num_rows| { + let column = ColumnChunkMetaData::builder(schema_descr.column(0)) + .set_num_values(*num_rows) + .build() + .unwrap(); + + RowGroupMetaData::builder(Arc::clone(&schema_descr)) + .set_num_rows(*num_rows) + .set_column_metadata(vec![column]) + .build() + .unwrap() + }) + .collect() + } + + #[test] + fn should_load_page_index_checks_predicate_columns() { + // "a" has page index offsets recorded in the footer, "b" does not + let metadata = page_index_metadata(&[("a", true), ("b", false)], 1); + + // predicate on "a": the file has a page index for it, so load it + assert!(should_load_page_index( + metadata.clone(), + Some(col("a").gt(lit(50i32))), + ParquetAccessPlan::new_all(1), + )); + + // predicate on "b": no page index for that column, so skip the load + assert!(!should_load_page_index( + metadata, + Some(col("b").gt(lit(50i32))), + ParquetAccessPlan::new_all(1), + )); + } + + fn test_schema_descr() -> SchemaDescPtr { + let schema = Schema::new(vec![Field::new("a", DataType::Utf8, false)]); + Arc::new(ArrowSchemaConverter::new().convert(&schema).unwrap()) + } + + /// Metadata for a file of Int32 `columns`, where each `(name, + /// has_page_index)` entry controls whether the footer records page index + /// offsets for that column. + fn page_index_metadata( + columns: &[(&str, bool)], + num_row_groups: usize, + ) -> ParquetMetaData { + let arrow_schema = Schema::new( + columns + .iter() + .map(|(name, _)| Field::new(*name, DataType::Int32, false)) + .collect::>(), + ); + let schema_descr = + Arc::new(ArrowSchemaConverter::new().convert(&arrow_schema).unwrap()); + + let row_groups = (0..num_row_groups) + .map(|_| { + let columns = columns + .iter() + .enumerate() + .map(|(idx, (_, has_page_index))| { + let mut builder = + ColumnChunkMetaData::builder(schema_descr.column(idx)) + .set_num_values(10); + if *has_page_index { + builder = builder + .set_column_index_offset(Some(100)) + .set_column_index_length(Some(10)) + .set_offset_index_offset(Some(110)) + .set_offset_index_length(Some(10)); + } + builder.build().unwrap() + }) + .collect(); + RowGroupMetaData::builder(Arc::clone(&schema_descr)) + .set_num_rows(10) + .set_column_metadata(columns) + .build() + .unwrap() + }) + .collect(); + let file_metadata = + FileMetaData::new(1, 10, None, None, Arc::clone(&schema_descr), None); + ParquetMetaData::new(file_metadata, row_groups) + } + + /// Reports [`RowGroupsPrunedParquetOpen::should_load_page_index`] for + /// hand-built parquet `metadata` (no I/O), an optional predicate, and a + /// row group access plan. + fn should_load_page_index( + metadata: ParquetMetaData, + predicate: Option, + plan: ParquetAccessPlan, + ) -> bool { + use crate::RowGroupAccessPlanFilter; + use parquet::arrow::parquet_to_arrow_schema; + + let arrow_schema: SchemaRef = Arc::new( + parquet_to_arrow_schema(metadata.file_metadata().schema_descr(), None) + .unwrap(), + ); + let page_pruning_predicate = predicate.map(|expr| { + let predicate = logical2physical(&expr, &arrow_schema); + build_page_pruning_predicate(&predicate, &arrow_schema) + }); + + let store: Arc = Arc::new(InMemory::new()); + let morselizer = ParquetMorselizerBuilder::new() + .with_store(store) + .with_schema(Arc::clone(&arrow_schema)) + .build(); + let file = PartitionedFile::new("test.parquet".to_string(), 100); + let prepared = morselizer.prepare_open_file(file).unwrap(); + let options = ArrowReaderOptions::new(); + let reader_metadata = + ArrowReaderMetadata::try_new(Arc::new(metadata), options.clone()).unwrap(); + let open = RowGroupsPrunedParquetOpen { + prepared: FiltersPreparedParquetOpen { + loaded: MetadataLoadedParquetOpen { + prepared, + reader_metadata, + options, + }, + pruning_predicate: None, + page_pruning_predicate, + }, + row_groups: RowGroupAccessPlanFilter::new(plan), + }; + open.should_load_page_index() + } + impl ParquetMorselizerBuilder { /// Create a new builder with sensible defaults for tests. fn new() -> Self { @@ -1488,6 +2008,7 @@ mod test { predicate: None, metadata_size_hint: None, metrics: ExecutionPlanMetricsSet::new(), + parquet_file_reader_factory: None, pushdown_filters: false, reorder_filters: false, force_filter_selections: false, @@ -1496,6 +2017,7 @@ mod test { enable_row_group_stats_pruning: false, coerce_int96: None, max_predicate_cache_size: None, + max_in_list_size: MAX_IN_LIST_SIZE, reverse_row_groups: false, preserve_order: false, } @@ -1509,7 +2031,7 @@ mod test { /// Create a simple table schema from a file schema (for files without partition columns). fn with_schema(mut self, file_schema: SchemaRef) -> Self { - self.table_schema = Some(TableSchema::from_file_schema(file_schema)); + self.table_schema = Some(TableSchema::from(file_schema)); self } @@ -1519,12 +2041,28 @@ mod test { self } - /// Set projection by column indices (convenience method for common case). + /// Set projection by column indices. + /// + /// The indices are resolved against the **file schema**, not the full + /// table schema. Callers that need to project partition columns or + /// virtual columns must use [`Self::with_projection`] and construct a + /// [`ProjectionExprs`] against [`TableSchema::table_schema`]. fn with_projection_indices(mut self, indices: &[usize]) -> Self { self.projection_indices = Some(indices.to_vec()); self } + /// Set an explicit projection. + /// + /// Prefer this over [`Self::with_projection_indices`] whenever the + /// projection must reference partition or virtual columns, since + /// `with_projection_indices` resolves its indices against the file + /// schema only. + fn with_projection(mut self, projection: ProjectionExprs) -> Self { + self.projection = Some(projection); + self + } + /// Set the predicate. fn with_predicate(mut self, predicate: Arc) -> Self { self.predicate = Some(predicate); @@ -1555,6 +2093,19 @@ mod test { self } + fn with_metrics(mut self, metrics: ExecutionPlanMetricsSet) -> Self { + self.metrics = metrics; + self + } + + fn with_parquet_file_reader_factory( + mut self, + factory: Arc, + ) -> Self { + self.parquet_file_reader_factory = Some(factory); + self + } + /// Set a row limit. fn with_limit(mut self, limit: usize) -> Self { self.limit = Some(limit); @@ -1567,12 +2118,26 @@ mod test { self } - /// Build the ParquetMorselizer instance. + /// Build the ParquetMorselizer instance, unwrapping validation errors. /// /// # Panics /// - /// Panics if required fields (store, schema/table_schema) are not set. + /// Panics if required fields (store, schema/table_schema) are not set, + /// or if virtual-column validation fails. Use [`Self::try_build`] + /// when the test wants to assert on the validation error. fn build(self) -> ParquetMorselizer { + self.try_build().expect("ParquetMorselizerBuilder::build") + } + + /// Build the ParquetMorselizer instance, returning any morselizer-level + /// validation error (e.g. unsupported virtual extension type, or a + /// predicate that references a virtual column with + /// `pushdown_filters=true`). + /// + /// # Panics + /// + /// Panics if required fields (store, schema/table_schema) are not set. + fn try_build(self) -> Result { let store = self .store .expect("ParquetMorselizerBuilder: store must be set via with_store()"); @@ -1591,7 +2156,14 @@ mod test { ProjectionExprs::from_indices(&all_indices, &file_schema) }; - ParquetMorselizer { + let virtual_state = build_virtual_columns_state( + table_schema.virtual_columns(), + table_schema.file_schema(), + self.predicate.as_ref(), + self.pushdown_filters, + )?; + + Ok(ParquetMorselizer { partition_index: self.partition_index, projection, batch_size: self.batch_size, @@ -1601,9 +2173,11 @@ mod test { table_schema, metadata_size_hint: self.metadata_size_hint, metrics: self.metrics, - parquet_file_reader_factory: Arc::new( - DefaultParquetFileReaderFactory::new(store), - ), + parquet_file_reader_factory: self + .parquet_file_reader_factory + .unwrap_or_else(|| { + Arc::new(DefaultParquetFileReaderFactory::new(store)) as _ + }), pushdown_filters: self.pushdown_filters, reorder_filters: self.reorder_filters, force_filter_selections: self.force_filter_selections, @@ -1621,9 +2195,11 @@ mod test { #[cfg(feature = "parquet_encryption")] encryption_factory: None, max_predicate_cache_size: self.max_predicate_cache_size, + max_in_list_size: self.max_in_list_size, reverse_row_groups: self.reverse_row_groups, sort_order_for_reorder: None, - } + virtual_state, + }) } } @@ -1784,7 +2360,7 @@ mod test { async fn write_parquet( store: Arc, filename: &str, - batch: arrow::record_batch::RecordBatch, + batch: RecordBatch, ) -> usize { write_parquet_batches(store, filename, vec![batch], None).await } @@ -1793,7 +2369,7 @@ mod test { async fn write_parquet_batches( store: Arc, filename: &str, - batches: Vec, + batches: Vec, props: Option, ) -> usize { let mut out = BytesMut::new().writer(); @@ -1811,6 +2387,18 @@ mod test { data_len } + fn counter_metric_value(metrics: &ExecutionPlanMetricsSet, name: &str) -> usize { + use datafusion_physical_plan::metrics::MetricValue; + metrics + .clone_inner() + .sum_by_name(name) + .map(|metric| match metric { + MetricValue::Count { count, .. } => count.value(), + _ => 0, + }) + .unwrap_or(0) + } + fn make_dynamic_expr(expr: Arc) -> Arc { Arc::new(DynamicFilterPhysicalExpr::new( expr.children().into_iter().map(Arc::clone).collect(), @@ -1896,10 +2484,13 @@ mod test { Field::new("a", DataType::Int32, false), ])); - let table_schema_for_opener = TableSchema::new( - file_schema.clone(), - vec![Arc::new(Field::new("part", DataType::Int32, false))], - ); + let table_schema_for_opener = TableSchemaBuilder::from(&file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part", + DataType::Int32, + false, + ))]) + .build(); let make_opener = |predicate| { ParquetMorselizerBuilder::new() .with_store(Arc::clone(&store)) @@ -1965,10 +2556,13 @@ mod test { Field::new("a", DataType::Int32, false), Field::new("b", DataType::Float32, true), ])); - let table_schema_for_opener = TableSchema::new( - file_schema.clone(), - vec![Arc::new(Field::new("part", DataType::Int32, false))], - ); + let table_schema_for_opener = TableSchemaBuilder::from(&file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part", + DataType::Int32, + false, + ))]) + .build(); let make_opener = |predicate| { ParquetMorselizerBuilder::new() .with_store(Arc::clone(&store)) @@ -2037,10 +2631,13 @@ mod test { Field::new("a", DataType::Int32, false), ])); - let table_schema_for_opener = TableSchema::new( - file_schema.clone(), - vec![Arc::new(Field::new("part", DataType::Int32, false))], - ); + let table_schema_for_opener = TableSchemaBuilder::from(&file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part", + DataType::Int32, + false, + ))]) + .build(); let make_opener = |predicate| { ParquetMorselizerBuilder::new() .with_store(Arc::clone(&store)) @@ -2118,10 +2715,13 @@ mod test { Field::new("part", DataType::Int32, false), ])); - let table_schema_for_opener = TableSchema::new( - file_schema.clone(), - vec![Arc::new(Field::new("part", DataType::Int32, false))], - ); + let table_schema_for_opener = TableSchemaBuilder::from(&file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part", + DataType::Int32, + false, + ))]) + .build(); let make_opener = |predicate| { ParquetMorselizerBuilder::new() .with_store(Arc::clone(&store)) @@ -2170,6 +2770,55 @@ mod test { assert_eq!(num_rows, 0); } + #[tokio::test] + async fn test_opener_prioritizes_partitioned_file_schema() { + let store = Arc::new(InMemory::new()) as Arc; + + let batch = record_batch!( + ("a", Int32, vec![Some(1), Some(2), Some(2)]), + ("b", Float32, vec![Some(1.0), Some(2.0), None]) + ) + .unwrap(); + let data_size = + write_parquet(Arc::clone(&store), "test.parquet", batch.clone()).await; + + let schema = batch.schema(); + let query_file = async |schema: SchemaRef| -> Result<(usize, usize)> { + let file = PartitionedFile::new( + "test.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ) + .with_arrow_schema(schema.clone()); + + let predicate = logical2physical(&col("a").eq(lit(1)), &schema); + let opener = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(predicate) + .build(); + + let stream = open_file(&opener, file.clone()).await?; + Ok(count_batches_and_rows(stream).await) + }; + + let (num_batches, num_rows) = + query_file(schema.clone()).await.expect("query_file"); + assert_eq!(num_batches, 1); + assert_eq!(num_rows, 3); + + let mismatching_schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Float64, true), + ]); + assert_eq!( + query_file(SchemaRef::new(mismatching_schema)) + .await + .unwrap_err() + .message(), + "Arrow: Incompatible supplied Arrow schema: data type mismatch for field b: requested Float64 but found Float32" + ); + } + #[tokio::test] async fn test_reverse_scan_row_groups() { use parquet::file::properties::WriterProperties; @@ -2528,6 +3177,197 @@ mod test { ); } + #[test] + fn should_load_page_index_without_predicate() { + assert!(!should_load_page_index( + page_index_metadata(&[("a", true)], 2), + None, + ParquetAccessPlan::new_all(2), + )); + } + + #[test] + fn should_load_page_index_when_surviving_row_groups_not_fully_matched() { + assert!(should_load_page_index( + page_index_metadata(&[("a", true)], 2), + Some(col("a").gt(lit(50i32))), + ParquetAccessPlan::new_all(2), + )); + } + + #[test] + fn should_load_page_index_when_all_surviving_row_groups_fully_matched() { + let mut plan = ParquetAccessPlan::new_all(1); + plan.mark_fully_matched(0); + assert!(!should_load_page_index( + page_index_metadata(&[("a", true)], 1), + Some(col("a").is_not_null()), + plan, + )); + } + + #[tokio::test] + async fn test_page_index_skipped_when_row_groups_fully_matched() { + use parquet::file::properties::WriterProperties; + + let store = Arc::new(InMemory::new()) as Arc; + let values: Vec = (1..=100).collect(); + let batch = record_batch!(( + "a", + Int32, + values.iter().map(|v| Some(*v)).collect::>() + )) + .unwrap(); + let props = WriterProperties::builder() + .set_data_page_row_count_limit(10) + .set_write_batch_size(10) + .build(); + let schema = batch.schema(); + let data_len = write_parquet_batches( + Arc::clone(&store), + "test.parquet", + vec![batch], + Some(props), + ) + .await; + let file = PartitionedFile::new( + "test.parquet".to_string(), + u64::try_from(data_len).unwrap(), + ); + let predicate = logical2physical(&col("a").gt(lit(0i32)), &schema); + let metrics = ExecutionPlanMetricsSet::new(); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&predicate)) + .with_enable_page_index(true) + .with_row_group_stats_pruning(true) + .with_pushdown_filters(false) + .with_metrics(metrics.clone()) + .build(); + + let (_, rows) = + count_batches_and_rows(open_file(&morselizer, file).await.unwrap()).await; + assert_eq!(rows, 100); + assert_eq!(counter_metric_value(&metrics, "page_index_load_skipped"), 1); + } + + #[tokio::test] + async fn test_page_index_skipped_with_cached_reader_factory() { + use parquet::file::properties::WriterProperties; + + let store = Arc::new(InMemory::new()) as Arc; + let metadata_cache: Arc = + Arc::new(DefaultCache::::new( + 64 * 1024 * 1024, + )); + let values: Vec = (1..=100).collect(); + let batch = record_batch!(( + "a", + Int32, + values.iter().map(|v| Some(*v)).collect::>() + )) + .unwrap(); + let props = WriterProperties::builder() + .set_data_page_row_count_limit(10) + .set_write_batch_size(10) + .build(); + let schema = batch.schema(); + let data_len = write_parquet_batches( + Arc::clone(&store), + "test.parquet", + vec![batch], + Some(props), + ) + .await; + let file = PartitionedFile::new( + "test.parquet".to_string(), + u64::try_from(data_len).unwrap(), + ); + let predicate = logical2physical(&col("a").gt(lit(0i32)), &schema); + let metrics = ExecutionPlanMetricsSet::new(); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&predicate)) + .with_enable_page_index(true) + .with_row_group_stats_pruning(true) + .with_pushdown_filters(false) + .with_metrics(metrics.clone()) + .with_parquet_file_reader_factory(Arc::new( + CachedParquetFileReaderFactory::new( + Arc::clone(&store), + Arc::clone(&metadata_cache), + ), + )) + .build(); + + let (_, rows) = + count_batches_and_rows(open_file(&morselizer, file).await.unwrap()).await; + assert_eq!(rows, 100); + assert_eq!(counter_metric_value(&metrics, "page_index_load_skipped"), 1); + + let cached = metadata_cache + .get(&Path::from("test.parquet")) + .expect("metadata cache should contain the file"); + let extra_info = cached.file_metadata.extra_info(); + let page_index_cached = extra_info.get("page_index").map(String::as_str); + assert_eq!( + page_index_cached, + Some("false"), + "cached metadata should not include page index when opener skips it" + ); + } + + #[tokio::test] + async fn test_page_index_loaded_when_not_fully_matched() { + use parquet::file::properties::WriterProperties; + + let store = Arc::new(InMemory::new()) as Arc; + let values: Vec = (1..=100).collect(); + let batch = record_batch!(( + "a", + Int32, + values.iter().map(|v| Some(*v)).collect::>() + )) + .unwrap(); + let props = WriterProperties::builder() + .set_data_page_row_count_limit(10) + .set_write_batch_size(10) + .build(); + let schema = batch.schema(); + let data_len = write_parquet_batches( + Arc::clone(&store), + "test.parquet", + vec![batch], + Some(props), + ) + .await; + let file = PartitionedFile::new( + "test.parquet".to_string(), + u64::try_from(data_len).unwrap(), + ); + let predicate = logical2physical(&col("a").gt(lit(90i32)), &schema); + let metrics = ExecutionPlanMetricsSet::new(); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(Arc::clone(&schema)) + .with_predicate(Arc::clone(&predicate)) + .with_enable_page_index(true) + .with_pushdown_filters(false) + .with_row_group_stats_pruning(false) + .with_metrics(metrics.clone()) + .build(); + + let (_, rows) = + count_batches_and_rows(open_file(&morselizer, file).await.unwrap()).await; + assert_eq!(rows, 10); + assert_eq!(counter_metric_value(&metrics, "page_index_load_skipped"), 0); + } + async fn fully_matched_split_test_file( store: Arc, ) -> (SchemaRef, PartitionedFile) { @@ -2600,89 +3440,484 @@ mod test { assert_eq!(values, vec![7, 4, 5, 6, 3]); } - #[test] - fn test_split_decoder_runs_no_fully_matched() { - // All row groups need filtering: single run. - let plan = ParquetAccessPlan::new(vec![ - RowGroupAccess::Scan, - RowGroupAccess::Scan, - RowGroupAccess::Scan, - ]); - let runs = plan.split_runs(true); - assert_eq!(runs.len(), 1); - assert!(runs[0].needs_filter); - assert_eq!(runs[0].access_plan.row_group_indexes(), vec![0, 1, 2]); - } + /// Helpers for tests that exercise parquet virtual columns + /// (e.g. `row_number`) plumbed through `TableSchema`/`ParquetOpener`. + mod virtual_columns { + use super::*; + use arrow::array::{Array, Int64Array, StringArray}; + use arrow::datatypes::FieldRef; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarUDF; + use datafusion_functions::core::input_file_name::InputFileNameFunc; + use datafusion_physical_expr::{ScalarFunctionExpr, projection::ProjectionExpr}; + use parquet::arrow::RowNumber; + + /// Build a parquet `row_number` virtual column field. Spark's + /// `_tmp_metadata_row_index` is declared nullable, so the default + /// matches that contract; tests that need `nullable=false` can + /// override via `with_nullable`. + fn row_number_field(name: &str, nullable: bool) -> FieldRef { + Arc::new( + Field::new(name, DataType::Int64, nullable) + .with_extension_type(RowNumber), + ) + } - #[test] - fn test_split_decoder_runs_all_fully_matched() { - // All row groups are fully matched: single run, no filter. - let mut plan = ParquetAccessPlan::new(vec![ - RowGroupAccess::Scan, - RowGroupAccess::Scan, - RowGroupAccess::Scan, - ]); - plan.mark_fully_matched(0); - plan.mark_fully_matched(1); - plan.mark_fully_matched(2); + fn input_file_name_expr() -> Arc { + Arc::new(ScalarFunctionExpr::new( + "input_file_name", + Arc::new(ScalarUDF::from(InputFileNameFunc::new())), + vec![], + Arc::new(Field::new("input_file_name", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + )) + } - let runs = plan.split_runs(true); - assert_eq!(runs.len(), 1); - assert!(!runs[0].needs_filter); - assert_eq!(runs[0].access_plan.row_group_indexes(), vec![0, 1, 2]); - } + /// Collect every `Int64` value from the given column in every batch + /// of a stream. Used to verify the `row_number` column end to end. + async fn collect_int64_values( + mut stream: BoxStream<'static, Result>, + column: usize, + ) -> Vec { + let mut out = vec![]; + while let Some(batch) = stream.next().await { + let batch = batch.unwrap(); + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .expect("expected Int64 column"); + for i in 0..array.len() { + assert!( + !array.is_null(i), + "row_number values produced by the reader must not be null" + ); + out.push(array.value(i)); + } + } + out + } - #[test] - fn test_split_decoder_runs_mixed() { - // [F, M, M, F, M] creates 4 runs preserving order. - let mut plan = ParquetAccessPlan::new(vec![ - RowGroupAccess::Scan, // 0: filtered - RowGroupAccess::Scan, // 1: matched - RowGroupAccess::Scan, // 2: matched - RowGroupAccess::Scan, // 3: filtered - RowGroupAccess::Scan, // 4: matched - ]); - plan.mark_fully_matched(1); - plan.mark_fully_matched(2); - plan.mark_fully_matched(4); + /// Write a parquet file containing `num_row_groups` groups of + /// `rows_per_group` rows with a single `value` Int64 column. + /// Values are `0..num_row_groups*rows_per_group`. + async fn write_grouped_file( + store: &Arc, + path: &str, + num_row_groups: usize, + rows_per_group: usize, + ) -> (SchemaRef, usize) { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let mut batches = Vec::with_capacity(num_row_groups); + for g in 0..num_row_groups { + let start = (g * rows_per_group) as i64; + let values: Vec = (start..start + rows_per_group as i64).collect(); + batches.push( + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(values))], + ) + .unwrap(), + ); + } + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(rows_per_group)) + .build(); + let data_size = + write_parquet_batches(Arc::clone(store), path, batches, Some(props)) + .await; + (schema, data_size) + } - let runs = plan.split_runs(true); - assert_eq!(runs.len(), 4); + #[tokio::test] + async fn test_row_index_basic() { + let store = Arc::new(InMemory::new()) as Arc; + let (file_schema, data_size) = + write_grouped_file(&store, "basic.parquet", 1, 5).await; + + let rn_field = row_number_field("row_number", false); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(vec![Arc::clone(&rn_field)]) + .build(); + // Project [value, row_number] — indices in table_schema are + // [0 file:value, 1 virtual:row_number]. + let projection = + ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .with_projection(projection) + .build(); - assert!(runs[0].needs_filter); - assert_eq!(runs[0].access_plan.row_group_indexes(), vec![0]); + let file = PartitionedFile::new( + "basic.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ); + let stream = open_file(&morselizer, file).await.unwrap(); + let row_numbers = collect_int64_values(stream, 1).await; + assert_eq!(row_numbers, vec![0, 1, 2, 3, 4]); + } - assert!(!runs[1].needs_filter); - assert_eq!(runs[1].access_plan.row_group_indexes(), vec![1, 2]); + #[tokio::test] + async fn test_row_index_projection_only() { + let store = Arc::new(InMemory::new()) as Arc; + let (file_schema, data_size) = + write_grouped_file(&store, "proj_only.parquet", 1, 4).await; + + let rn_field = row_number_field("row_number", false); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(vec![Arc::clone(&rn_field)]) + .build(); + // Project only the virtual column (index 1). + let projection = + ProjectionExprs::from_indices(&[1], table_schema.table_schema()); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .with_projection(projection) + .build(); - assert!(runs[2].needs_filter); - assert_eq!(runs[2].access_plan.row_group_indexes(), vec![3]); + let file = PartitionedFile::new( + "proj_only.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ); + let stream = open_file(&morselizer, file).await.unwrap(); + let row_numbers = collect_int64_values(stream, 0).await; + assert_eq!(row_numbers, vec![0, 1, 2, 3]); + } - assert!(!runs[3].needs_filter); - assert_eq!(runs[3].access_plan.row_group_indexes(), vec![4]); - } + #[tokio::test] + async fn test_input_file_name_projection() { + let store = Arc::new(InMemory::new()) as Arc; + let path = "dir/input_file_name.parquet"; + let (file_schema, data_size) = write_grouped_file(&store, path, 1, 3).await; - #[test] - fn test_split_decoder_runs_with_skipped_groups() { - // Skipped row groups are excluded from all runs. - let mut plan = ParquetAccessPlan::new(vec![ - RowGroupAccess::Scan, // 0: filtered - RowGroupAccess::Skip, // 1: pruned - RowGroupAccess::Scan, // 2: matched - RowGroupAccess::Scan, // 3: filtered - ]); - plan.mark_fully_matched(2); + let projection = ProjectionExprs::new([ + ProjectionExpr::new(Arc::new(Column::new("value", 0)), "value"), + ProjectionExpr::new(input_file_name_expr(), "file_name"), + ]); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_schema(file_schema) + .with_projection(projection) + .build(); + + let file = + PartitionedFile::new(path.to_string(), u64::try_from(data_size).unwrap()); + let mut stream = open_file(&morselizer, file).await.unwrap(); + let batch = stream.next().await.unwrap().unwrap(); + assert!(stream.next().await.is_none()); + + assert_eq!(batch.num_columns(), 2); + assert_eq!(batch.schema().field(0).name(), "value"); + assert_eq!(batch.schema().field(1).name(), "file_name"); + + let file_names = batch + .column(1) + .as_any() + .downcast_ref::() + .expect("file_name column should be Utf8"); + assert_eq!(file_names.len(), 3); + for i in 0..file_names.len() { + assert_eq!(file_names.value(i), path); + } + } + + #[tokio::test] + async fn test_row_index_multi_row_group() { + let store = Arc::new(InMemory::new()) as Arc; + let (file_schema, data_size) = + write_grouped_file(&store, "multi_rg.parquet", 3, 100).await; + + let rn_field = row_number_field("row_number", false); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(vec![Arc::clone(&rn_field)]) + .build(); + let projection = + ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .with_projection(projection) + .build(); + + let file = PartitionedFile::new( + "multi_rg.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ); + let stream = open_file(&morselizer, file).await.unwrap(); + let row_numbers = collect_int64_values(stream, 1).await; + let expected: Vec = (0..300).collect(); + assert_eq!(row_numbers, expected); + } - let runs = plan.split_runs(true); - assert_eq!(runs.len(), 3); + #[tokio::test] + async fn test_row_index_with_row_group_skip() { + // 3 row groups of 100 rows. A predicate that excludes the middle + // row group (values 100..200) must leave absolute row numbers + // 0..100 and 200..300 intact — not 0..200. This guards against + // the arrow-rs bug fixed in apache/arrow-rs#8863. + let store = Arc::new(InMemory::new()) as Arc; + let (file_schema, data_size) = + write_grouped_file(&store, "rg_skip.parquet", 3, 100).await; + + let rn_field = row_number_field("row_number", false); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(vec![Arc::clone(&rn_field)]) + .build(); + let projection = + ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + + // `value < 100 OR value >= 200` prunes the middle row group via + // min/max statistics. + let expr = col("value") + .lt(lit(100i64)) + .or(col("value").gt_eq(lit(200i64))); + let predicate = logical2physical(&expr, table_schema.table_schema()); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .with_projection(projection) + .with_predicate(predicate) + .with_row_group_stats_pruning(true) + .build(); + + let file = PartitionedFile::new( + "rg_skip.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ); + let stream = open_file(&morselizer, file).await.unwrap(); + let row_numbers = collect_int64_values(stream, 1).await; + let expected: Vec = (0..100).chain(200..300).collect(); + assert_eq!(row_numbers, expected); + } + + #[tokio::test] + async fn test_row_index_with_partition_cols() { + let store = Arc::new(InMemory::new()) as Arc; + let (file_schema, data_size) = + write_grouped_file(&store, "part=5/data.parquet", 1, 3).await; + + let rn_field = row_number_field("row_number", false); + let partition_col = Arc::new(Field::new("part", DataType::Int32, false)); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::clone(&partition_col)]) + .with_virtual_columns(vec![Arc::clone(&rn_field)]) + .build(); + // table_schema layout: [value(0), part(1), row_number(2)]. + let projection = + ProjectionExprs::from_indices(&[0, 1, 2], table_schema.table_schema()); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .with_projection(projection) + .build(); + + let mut file = PartitionedFile::new( + "part=5/data.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ); + file.partition_values = vec![ScalarValue::Int32(Some(5))]; + + let stream = open_file(&morselizer, file).await.unwrap(); + let mut stream = stream; + let batch = stream.next().await.unwrap().unwrap(); + assert!(stream.next().await.is_none()); - assert!(runs[0].needs_filter); - assert_eq!(runs[0].access_plan.row_group_indexes(), vec![0]); + assert_eq!(batch.num_columns(), 3); + assert_eq!(batch.schema().field(0).name(), "value"); + assert_eq!(batch.schema().field(1).name(), "part"); + assert_eq!(batch.schema().field(2).name(), "row_number"); - assert!(!runs[1].needs_filter); - assert_eq!(runs[1].access_plan.row_group_indexes(), vec![2]); + let part = batch + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert!(part.iter().all(|v| v == Some(5))); - assert!(runs[2].needs_filter); - assert_eq!(runs[2].access_plan.row_group_indexes(), vec![3]); + let rn = batch + .column(2) + .as_any() + .downcast_ref::() + .unwrap(); + let rn_values: Vec = (0..rn.len()).map(|i| rn.value(i)).collect(); + assert_eq!(rn_values, vec![0, 1, 2]); + } + + #[tokio::test] + async fn test_row_index_nullable_int64() { + // Spark declares `_tmp_metadata_row_index` nullable. Verify the + // nullability flag flows through unchanged. + let store = Arc::new(InMemory::new()) as Arc; + let (file_schema, data_size) = + write_grouped_file(&store, "nullable.parquet", 1, 3).await; + + let rn_field = row_number_field("_tmp_metadata_row_index", true); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(vec![Arc::clone(&rn_field)]) + .build(); + let projection = + ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .with_projection(projection) + .build(); + + let file = PartitionedFile::new( + "nullable.parquet".to_string(), + u64::try_from(data_size).unwrap(), + ); + let mut stream = open_file(&morselizer, file).await.unwrap(); + let batch = stream.next().await.unwrap().unwrap(); + + let schema_field = batch.schema().field(1).clone(); + assert_eq!(schema_field.name(), "_tmp_metadata_row_index"); + assert_eq!(schema_field.data_type(), &DataType::Int64); + assert!( + schema_field.is_nullable(), + "nullable flag should be preserved for Spark's row index field" + ); + } + + #[tokio::test] + async fn test_unsupported_virtual_extension_type_rejected() { + // Guard: opener must reject virtual columns carrying extension + // types outside the tested allowlist, rather than silently + // forwarding them to arrow-rs (where they would produce columns + // we have not validated against DataFusion's projection and + // predicate paths). + let store = Arc::new(InMemory::new()) as Arc; + let (file_schema, _data_size) = + write_grouped_file(&store, "unsupported.parquet", 1, 1).await; + + // RowGroupIndex is a real arrow-rs virtual type but is not in + // SUPPORTED_VIRTUAL_EXTENSION_TYPES until a test is added for it. + let rg_field = Arc::new( + Field::new("row_group_index", DataType::Int64, false) + .with_extension_type(parquet::arrow::RowGroupIndex), + ); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(vec![rg_field]) + .build(); + let projection = + ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + + // Validation now happens at morselizer-build time (once per scan + // partition), not once per file inside `prepare_open_file`. + let err = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(&store)) + .with_table_schema(table_schema) + .with_projection(projection) + .try_build() + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("parquet.virtual.row_group_index"), + "error should name the unsupported extension type, got: {msg}" + ); + } + + /// Build a morselizer + file for a 5-row single-row-group parquet at + /// `path`, with a single `row_number` virtual column and the given + /// physical predicate applied to + /// `table_schema = [value(0), row_number(1)]`. + async fn build_pushdown_morselizer( + store: &Arc, + path: &str, + predicate_expr: Expr, + pushdown_filters: bool, + ) -> Result<(ParquetMorselizer, PartitionedFile)> { + let (file_schema, data_size) = write_grouped_file(store, path, 1, 5).await; + let rn_field = row_number_field("row_number", false); + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(vec![Arc::clone(&rn_field)]) + .build(); + let projection = + ProjectionExprs::from_indices(&[0, 1], table_schema.table_schema()); + let predicate = + logical2physical(&predicate_expr, table_schema.table_schema()); + + let morselizer = ParquetMorselizerBuilder::new() + .with_store(Arc::clone(store)) + .with_table_schema(table_schema) + .with_projection(projection) + .with_predicate(predicate) + .with_pushdown_filters(pushdown_filters) + .try_build()?; + + let file = + PartitionedFile::new(path.to_string(), u64::try_from(data_size).unwrap()); + Ok((morselizer, file)) + } + + // The predicate-vs-virtual-column check rejects callers that bypass + // `ParquetSource::try_pushdown_filters` (which keeps virtual-col + // filters above the scan as a `FilterExec`) and set the predicate + // directly on the source with pushdown enabled. Without this guard, + // arrow-rs's `RowFilter` would silently drop the virtual-col conjunct + // and produce wrong results. + #[tokio::test] + async fn test_row_index_predicate_pushdown_mixed_or_errors() { + let store = Arc::new(InMemory::new()) as Arc; + let expr = col("row_number") + .eq(lit(2i64)) + .or(col("value").eq(lit(4i64))); + let err = + build_pushdown_morselizer(&store, "pushdown_mixed.parquet", expr, true) + .await + .unwrap_err(); + assert!( + err.to_string().contains("try_pushdown_filters"), + "error should mention try_pushdown_filters, got: {err}" + ); + } + + #[tokio::test] + async fn test_row_index_predicate_pushdown_virtual_only_errors() { + let store = Arc::new(InMemory::new()) as Arc; + let expr = col("row_number").eq(lit(2i64)); + let err = build_pushdown_morselizer( + &store, + "pushdown_virtual_only.parquet", + expr, + true, + ) + .await + .unwrap_err(); + assert!( + err.to_string().contains("try_pushdown_filters"), + "error should mention try_pushdown_filters, got: {err}" + ); + } + + #[tokio::test] + async fn test_row_index_predicate_allowed_when_pushdown_disabled() { + // Guards the `pushdown_filters=false` path: the predicate is only + // used for stats pruning (a no-op for row_number) and must not + // trip the virtual-column check. + let store = Arc::new(InMemory::new()) as Arc; + let expr = col("row_number").eq(lit(2i64)); + let (morselizer, file) = + build_pushdown_morselizer(&store, "pushdown_off.parquet", expr, false) + .await + .unwrap(); + + let stream = open_file(&morselizer, file).await.unwrap(); + let (_batches, rows) = count_batches_and_rows(stream).await; + assert_eq!(rows, 5); + } } } diff --git a/datafusion/datasource-parquet/src/page_filter.rs b/datafusion/datasource-parquet/src/page_filter.rs index 795a63268b6a9..6bc1aca667981 100644 --- a/datafusion/datasource-parquet/src/page_filter.rs +++ b/datafusion/datasource-parquet/src/page_filter.rs @@ -31,7 +31,7 @@ use arrow::{ use datafusion_common::ScalarValue; use datafusion_common::pruning::PruningStatistics; use datafusion_physical_expr::{PhysicalExpr, split_conjunction}; -use datafusion_pruning::PruningPredicate; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; use log::{debug, trace}; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; @@ -144,10 +144,10 @@ impl PagePruningAccessPlanFilter { let predicates = split_conjunction(expr) .into_iter() .filter_map(|predicate| { - let pp = match PruningPredicate::try_new( - Arc::clone(predicate), - Arc::clone(&schema), - ) { + let pp = match PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(predicate)) + { Ok(pp) => pp, Err(e) => { debug!("Ignoring error creating page pruning predicate: {e}"); @@ -303,7 +303,7 @@ impl PagePruningAccessPlanFilter { debug!( "Use filter and page index to create RowSelection {:?} from predicate: {:?}", - &selection, + selection, predicate.predicate_expr(), ); @@ -376,6 +376,16 @@ impl PagePruningAccessPlanFilter { pub fn filter_number(&self) -> usize { self.predicates.len() } + + /// Returns the names of the columns referenced by the page pruning + /// predicates (each predicate references exactly one column, see + /// [`Self::new`]). + pub(crate) fn predicate_column_names(&self) -> impl Iterator { + self.predicates + .iter() + .filter_map(|p| p.required_columns().single_column()) + .map(|c| c.name()) + } } fn update_selection( diff --git a/datafusion/datasource-parquet/src/projection_read_plan.rs b/datafusion/datasource-parquet/src/projection_read_plan.rs new file mode 100644 index 0000000000000..038180851125a --- /dev/null +++ b/datafusion/datasource-parquet/src/projection_read_plan.rs @@ -0,0 +1,1797 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Resolution of expressions against a Parquet file's schema into a +//! [`ParquetReadPlan`]: the leaf-level [`ProjectionMask`] to install on the +//! decoder plus the Arrow schema the decoder will emit under that mask. +//! +//! This is shared by the opener's projection handling (via +//! [`build_projection_read_plan`]) and row-filter construction (via +//! [`crate::row_filter`]), which both need to translate column and struct +//! field references into Parquet leaf indices. [`PushdownChecker`], the +//! expression traversal that discovers those references, lives here as well +//! so that [`crate::row_filter`] depends on this module and not vice versa. + +use std::collections::{BTreeMap, BTreeSet}; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion_functions::core::input_file_name::InputFileNameFunc; +use parquet::arrow::ProjectionMask; +use parquet::schema::types::SchemaDescriptor; + +use datafusion_common::Result; +use datafusion_common::nested_struct::requires_nested_struct_cast; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; +use datafusion_functions::core::file_row_index::FileRowIndexFunc; +use datafusion_functions::core::getfield::GetFieldFunc; +use datafusion_physical_expr::expressions::{CastExpr, Column, Literal}; +use datafusion_physical_expr::utils::collect_columns; +use datafusion_physical_expr::{PhysicalExpr, ScalarFunctionExpr}; + +use crate::nested_schema_pruning::{ + CastColumnAccess, clip_for_cast, contains_struct, count_leaves, field_with_type, +}; + +/// The result of resolving which Parquet leaf columns and Arrow schema fields +/// are needed to evaluate an expression against a Parquet file +/// +/// This is the shared output of the column resolution pipeline used by both +/// the row filter to build `ArrowPredicate`s and the opener to build `ProjectionMask`s +#[derive(Debug, Clone)] +pub(crate) struct ParquetReadPlan { + /// Projection mask built from leaf column indices in the Parquet schema. + /// Using a `ProjectionMask` directly (rather than raw indices) prevents + /// bugs from accidentally mixing up root vs leaf indices. + pub projection_mask: ProjectionMask, + /// The projected Arrow schema containing only the columns/fields required + /// Struct types are pruned to include only the accessed sub-fields + pub projected_schema: SchemaRef, +} + +/// Records a struct field access via `get_field(struct_col, 'field1', 'field2', ...)`. +/// +/// This allows the row filter to project only the specific Parquet leaf columns +/// needed by the filter, rather than all leaves of the struct. +#[derive(Debug, Clone)] +pub(crate) struct StructFieldAccess { + /// Arrow root column index of the struct in the file schema. + pub(crate) root_index: usize, + /// Field names forming the path into the struct. + /// e.g., `["value"]` for `s['value']`, `["outer", "inner"]` for `s['outer']['inner']`. + pub(crate) field_path: Vec, +} + +/// Trie of nested struct accesses, keyed at the top by the root column index in +/// the file schema and then by field names down each access path. +/// +/// # Example +/// +/// For a filter expression +/// +/// ```sql +/// WHERE s['outer']['a'] > 10 +/// AND s['outer']['b'] < 20 +/// AND s['outer']['inner']['c'] IS NOT NULL +/// ``` +/// +/// where `s` is column index `2` in the file schema, three accesses are +/// recorded — all with `root_index = 2` and paths `["outer","a"]`, +/// `["outer","b"]`, `["outer","inner","c"]`. They produce a trie in which +/// the shared `"outer"` prefix is represented by a single intermediate node: +/// +/// ```text +/// roots: +/// 2 ──► node { selected_here: false } +/// children: +/// "outer" ──► node { selected_here: false } +/// children: +/// "a" ──► { selected_here: true, children: {} } +/// "b" ──► { selected_here: true, children: {} } +/// "inner" ──► { selected_here: false, +/// children: { +/// "c" ──► { selected_here: true, +/// children: {} } +/// } } +/// ``` +#[derive(Debug, Default)] +struct StructAccessTree<'a> { + roots: BTreeMap>, +} + +/// One node in a [`StructAccessTree`]. +/// +/// `selected_here` is `true` when at least one access path terminates at this +/// node. Duplicate paths are idempotent. +#[derive(Debug, Default)] +struct StructAccessNode<'a> { + children: BTreeMap<&'a str, StructAccessNode<'a>>, + selected_here: bool, +} + +impl<'a> StructAccessTree<'a> { + /// Builds a [`StructAccessTree`] from a flat list of accesses. + /// + /// For each [`StructFieldAccess`], walks from the given root index down + /// the field path, creating intermediate nodes as needed, and sets the + /// terminal node's `selected_here` to `true`. Paths sharing a prefix + /// collapse onto common intermediate nodes. + fn from_accesses(accesses: &'a [StructFieldAccess]) -> Self { + let mut tree = Self::default(); + for StructFieldAccess { + root_index, + field_path, + } in accesses + { + let mut node = tree.roots.entry(*root_index).or_default(); + for component in field_path { + node = node.children.entry(component.as_str()).or_default(); + } + node.selected_here = true; + } + tree + } + + /// Returns the node for the given file-schema column index, or `None` if + /// no access path was recorded under that root. + fn root(&self, idx: usize) -> Option<&StructAccessNode<'a>> { + self.roots.get(&idx) + } +} + +/// Traverses a `PhysicalExpr` tree to determine if any column references would +/// prevent the expression from being pushed down to the parquet decoder. +/// +/// An expression cannot be pushed down if it references: +/// - Unsupported nested columns (whole struct references or list fields that are +/// not covered by the supported predicate set) +/// - Columns that don't exist in the file schema +/// +/// Struct field access via `get_field` is supported when the resolved leaf type +/// is primitive (e.g. `get_field(struct_col, 'field') > 5`). +pub(crate) struct PushdownChecker<'schema> { + /// Does the expression require any non-primitive columns (like structs)? + non_primitive_columns: bool, + /// Does the expression reference any columns not present in the file schema? + projected_columns: bool, + /// Does the expression references a ScalarUDF that requires some rewrite + /// and therefore can't be pushed down into the row-filter. + has_unpushable_udfs: bool, + /// Indices into the file schema of columns required to evaluate the expression. + /// Does not include struct columns accessed via `get_field`. + required_columns: Vec, + /// Struct field accesses via `get_field`. + struct_field_accesses: Vec, + /// Whole-column casts to a narrower nested type + /// (`CAST(col AS narrower_struct)`). Only collected when + /// [`Self::with_cast_collection`] enables it (projection analysis); + /// filter pushdown leaves this off. + cast_accesses: Vec, + /// Whether to collect [`Self::cast_accesses`]. + collect_cast_accesses: bool, + /// Whether nested list columns are supported by the predicate semantics. + allow_list_columns: bool, + /// The Arrow schema of the parquet file. + file_schema: &'schema Schema, +} + +impl<'schema> PushdownChecker<'schema> { + pub(crate) fn new(file_schema: &'schema Schema, allow_list_columns: bool) -> Self { + Self { + non_primitive_columns: false, + projected_columns: false, + has_unpushable_udfs: false, + required_columns: Vec::new(), + struct_field_accesses: Vec::new(), + cast_accesses: Vec::new(), + collect_cast_accesses: false, + allow_list_columns, + file_schema, + } + } + + /// Enable collection of whole-column casts to narrower nested types. + pub(crate) fn with_cast_collection(mut self) -> Self { + self.collect_cast_accesses = true; + self + } + + /// Checks whether a struct's root column exists in the file schema and, if so, + /// records its index so the entire struct is decoded for filter evaluation. + /// + /// This is called when we see a `get_field` expression that resolves to a + /// primitive leaf type. We only need the *root* column index because the + /// Parquet reader decodes all leaves of a struct together. + /// + /// # Example + /// + /// Given file schema `{a: Int32, s: Struct(foo: Utf8, bar: Int64)}` and the + /// expression `get_field(s, 'foo') = 'hello'`: + /// + /// - `column_name` = `"s"` (the root struct column) + /// - `file_schema.index_of("s")` returns `1` + /// - We push `1` into `required_columns` + /// - Return `None` (no issue — traversal continues in the caller) + /// + /// If `"s"` is not in the file schema (e.g. a projected-away column), we set + /// `projected_columns = true` and return `Jump` to skip the subtree. + fn check_struct_field_column( + &mut self, + column_name: &str, + field_path: Vec, + ) -> Option { + let Ok(idx) = self.file_schema.index_of(column_name) else { + self.projected_columns = true; + return Some(TreeNodeRecursion::Jump); + }; + + self.struct_field_accesses.push(StructFieldAccess { + root_index: idx, + field_path, + }); + + None + } + + fn check_single_column(&mut self, column_name: &str) -> Option { + let idx = match self.file_schema.index_of(column_name) { + Ok(idx) => idx, + Err(_) => { + // Column does not exist in the file schema, so we can't push this down. + self.projected_columns = true; + return Some(TreeNodeRecursion::Jump); + } + }; + + // Duplicates are handled by dedup() in into_sorted_columns() + self.required_columns.push(idx); + let data_type = self.file_schema.field(idx).data_type(); + + if DataType::is_nested(data_type) { + self.handle_nested_type(data_type) + } else { + None + } + } + + /// Determines whether a nested data type can be pushed down to Parquet decoding. + /// + /// Returns `Some(TreeNodeRecursion::Jump)` if the nested type prevents pushdown, + /// `None` if the type is supported and pushdown can continue. + fn handle_nested_type(&mut self, data_type: &DataType) -> Option { + if self.is_nested_type_supported(data_type) { + None + } else { + // Block pushdown for unsupported nested types: + // - Structs (regardless of predicate support) + // - Lists without supported predicates + self.non_primitive_columns = true; + Some(TreeNodeRecursion::Jump) + } + } + + /// Checks if a nested data type is supported for list column pushdown. + /// + /// List columns are only supported if: + /// 1. The data type is a list variant (List, LargeList, or FixedSizeList) + /// 2. The expression contains supported list predicates (e.g., array_has_all) + fn is_nested_type_supported(&self, data_type: &DataType) -> bool { + let is_list = matches!( + data_type, + DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) + ); + self.allow_list_columns && is_list + } + + #[inline] + pub(crate) fn prevents_pushdown(&self) -> bool { + self.non_primitive_columns || self.projected_columns || self.has_unpushable_udfs + } + + /// Consumes the checker and returns sorted, deduplicated column indices + /// wrapped in a `PushdownColumns` struct. + /// + /// This method sorts the column indices and removes duplicates. The sort + /// is required because downstream code relies on column indices being in + /// ascending order for correct schema projection. + pub(crate) fn into_sorted_columns(mut self) -> PushdownColumns { + self.required_columns.sort_unstable(); + self.required_columns.dedup(); + PushdownColumns { + required_columns: self.required_columns, + struct_field_accesses: self.struct_field_accesses, + cast_accesses: self.cast_accesses, + } + } +} + +impl TreeNodeVisitor<'_> for PushdownChecker<'_> { + type Node = Arc; + + fn f_down(&mut self, node: &Self::Node) -> Result { + // Handle struct field access like `s['foo']['bar'] > 10`. + // + // DataFusion represents nested field access as `get_field(Column("s"), "foo")` + // (or chained: `get_field(get_field(Column("s"), "foo"), "bar")`). + // + // We intercept the outermost `get_field` on the way *down* the tree so + // the visitor never reaches the raw `Column("s")` node. Without this, + // `check_single_column` would see that `s` is a Struct and reject it. + // + // The strategy: + // 1. Match `get_field` whose first arg is a `Column` (the struct root). + // 2. Check that the *resolved* return type is primitive — meaning we've + // drilled all the way to a leaf (e.g. `s['foo']` → Utf8). + // 3. Record the root column index via `check_struct_field_column` and + // return `Jump` to skip visiting the children (the Column and the + // literal field-name args), since we've already handled them. + // + // If the return type is still nested (e.g. `s['nested_struct']` → Struct), + // we fall through and let normal traversal continue, which will + // eventually reject the expression when it hits the struct Column. + if let Some(func) = + ScalarFunctionExpr::try_downcast_func::(node.as_ref()) + { + let args = func.args(); + + if let Some(column) = args.first().and_then(|a| a.downcast_ref::()) { + // for Map columns, get_field performs a runtime key lookup rather than a + // schema-level field access so the entire Map column must be read, + // we skip the struct field optimization and defer to normal Column traversal + let is_map_column = self + .file_schema + .index_of(column.name()) + .ok() + .map(|idx| { + matches!( + self.file_schema.field(idx).data_type(), + DataType::Map(_, _) + ) + }) + .unwrap_or(false); + + let return_type = func.return_type(); + + if !is_map_column + && (!DataType::is_nested(return_type) + || self.is_nested_type_supported(return_type)) + { + // if any field name argument is not a string literal we cannot + // determine the exact leaf path, so we fall back to reading the + // entire struct root column + let field_path = args[1..] + .iter() + .map(|arg| { + arg.downcast_ref::().and_then(|lit| { + lit.value().try_as_str().flatten().map(|s| s.to_string()) + }) + }) + .collect(); + + match field_path { + Some(path) => { + if let Some(recursion) = + self.check_struct_field_column(column.name(), path) + { + return Ok(recursion); + } + } + None => { + // Could not resolve field path — fall back to + // reading the entire struct root column. + if let Some(recursion) = + self.check_single_column(column.name()) + { + return Ok(recursion); + } + } + } + + return Ok(TreeNodeRecursion::Jump); + } + } + } + + // Handle whole-column casts to a narrower nested type, e.g. + // `CAST(events AS List>)` as inserted by the + // physical expression adapter when the logical file schema declares a + // nested column narrower than the physical file. Recording the cast + // target lets the projection read only the leaves the cast consumes + // (see `crate::nested_schema_pruning`). + if self.collect_cast_accesses + && let Some(cast) = node.downcast_ref::() + && let Some(column) = cast.expr().downcast_ref::() + && let Ok(idx) = self.file_schema.index_of(column.name()) + && requires_nested_struct_cast( + self.file_schema.field(idx).data_type(), + cast.cast_type(), + ) + { + self.cast_accesses.push(CastColumnAccess { + root_index: idx, + target_type: cast.cast_type().clone(), + }); + return Ok(TreeNodeRecursion::Jump); + } + + if let Some(column) = node.downcast_ref::() + && let Some(recursion) = self.check_single_column(column.name()) + { + return Ok(recursion); + } + + if ScalarFunctionExpr::try_downcast_func::(node.as_ref()) + .is_some() + || ScalarFunctionExpr::try_downcast_func::(node.as_ref()) + .is_some() + { + self.has_unpushable_udfs = true; + return Ok(TreeNodeRecursion::Jump); + } + + Ok(TreeNodeRecursion::Continue) + } +} + +/// Result of checking which columns are required for filter pushdown. +#[derive(Debug)] +pub(crate) struct PushdownColumns { + /// Sorted, unique column indices into the file schema required to evaluate + /// the filter expression. Must be in ascending order for correct schema + /// projection matching. Does not include struct columns accessed via `get_field`. + pub(crate) required_columns: Vec, + /// Struct field accesses via `get_field`. Each entry records the root struct + /// column index and the field path being accessed. + pub(crate) struct_field_accesses: Vec, + /// Whole-column casts to a narrower nested type. Empty unless cast + /// collection was enabled on the checker. + pub(crate) cast_accesses: Vec, +} + +/// Builds a unified [`ParquetReadPlan`] for a set of projection expressions +/// +/// Unlike [`crate::row_filter::build_parquet_read_plan`] (which is used for +/// filter pushdown and returns `None` when an expression references +/// unsupported nested types or missing columns), this function always +/// succeeds. It collects every column that *can* be resolved in the file and +/// produces a leaf-level projection mask. Columns missing from the file are +/// silently skipped since the projection layer handles those by inserting +/// nulls. +pub(crate) fn build_projection_read_plan( + exprs: impl IntoIterator>, + file_schema: &Schema, + schema_descr: &SchemaDescriptor, +) -> ParquetReadPlan { + // fast path: if every expression is a plain Column reference, skip all + // struct analysis and use root-level projection directly + let exprs = exprs.into_iter().collect::>(); + let all_plain_columns = exprs.iter().all(|e| e.downcast_ref::().is_some()); + + if all_plain_columns { + let mut root_indices: Vec = exprs + .iter() + .map(|e| e.downcast_ref::().unwrap().index()) + .collect(); + root_indices.sort_unstable(); + root_indices.dedup(); + + return root_level_plan(&root_indices, file_schema, schema_descr); + } + + // secondary fast path: if none of the *projected* columns contains a + // struct at any nesting level, there are no leaves to prune and we can + // skip the PushdownChecker traversal and use root-level projection. + // + // Gating on the projected roots rather than on every field of the file + // schema keeps this step O(projected columns): a wide file with a nested + // column the projection never touches should not push the whole + // projection through the slower, name-resolving path. Any column whose + // `index` does not line up with the file schema (a stale `Column` from an + // earlier rewrite) falls through to that path, which resolves by name. + let projected_columns = exprs.iter().flat_map(collect_columns).collect::>(); + let all_resolvable_and_struct_free = projected_columns.iter().all(|col| { + file_schema + .fields() + .get(col.index()) + .is_some_and(|f| f.name() == col.name() && !contains_struct(f.data_type())) + }); + + if all_resolvable_and_struct_free { + let mut root_indices = projected_columns + .iter() + .map(|c| c.index()) + .collect::>(); + root_indices.sort_unstable(); + root_indices.dedup(); + + return root_level_plan(&root_indices, file_schema, schema_descr); + } + + let mut all_root_indices = Vec::new(); + let mut all_struct_accesses = Vec::new(); + let mut all_cast_accesses = Vec::new(); + + for expr in exprs { + let mut checker = PushdownChecker::new(file_schema, true).with_cast_collection(); + let _ = expr.visit(&mut checker); + let columns = checker.into_sorted_columns(); + + all_root_indices.extend_from_slice(&columns.required_columns); + all_struct_accesses.extend(columns.struct_field_accesses); + all_cast_accesses.extend(columns.cast_accesses); + } + + all_root_indices.sort_unstable(); + all_root_indices.dedup(); + + // A whole-column reference reads every leaf of the root, so a cast + // access on the same root would be overridden anyway: drop those up + // front. `all_root_indices` is already sorted, so a binary search + // avoids building a second set just for this filter. + all_cast_accesses.retain(|c| all_root_indices.binary_search(&c.root_index).is_err()); + + if !all_cast_accesses.is_empty() { + return build_read_plan_with_cast_clipping( + file_schema, + schema_descr, + &all_root_indices, + &all_struct_accesses, + &all_cast_accesses, + ); + } + + // when no struct field accesses were found, fall back to root-level projection + // to match the performance of the simple path + if all_struct_accesses.is_empty() { + return root_level_plan(&all_root_indices, file_schema, schema_descr); + } + + let (read_plan, _leaf_indices) = assemble_read_plan( + &all_root_indices, + &all_struct_accesses, + file_schema, + schema_descr, + ); + + read_plan +} + +/// Builds a [`ParquetReadPlan`] when at least one projected root column is +/// consumed through a cast to a narrower nested type. +/// +/// Per root, in ascending root-index order: +/// - roots referenced as whole columns keep every leaf and their full +/// physical field (whole-column reads take precedence; cast accesses on +/// such roots were already dropped by the caller); +/// - roots consumed through a cast, and not also through a `get_field` +/// access on the same root, keep only the leaves the cast target names +/// (see `crate::nested_schema_pruning`); +/// - roots consumed only through `get_field` accesses keep the union of the +/// leaves those accesses reach, as before; +/// - any other referenced root, a cast that can't be safely clipped (see +/// `nested_schema_pruning::clip_for_cast`), a root reached by two casts +/// with *different* targets (a projection can consume the same column +/// through more than one narrowing cast, e.g. +/// `SELECT CAST(s AS STRUCT(a)), CAST(s AS STRUCT(b)) FROM t`; clipping to +/// either target alone would starve the other), or a root reached by both a +/// cast and a `get_field` access (not produced by +/// `DefaultPhysicalExprAdapter`, which always routes a `get_field` over a +/// narrowed column through the same cast rather than a separate access, +/// but a custom `PhysicalExprAdapter` could in principle inject both), +/// falls back to a full read of that root. +fn build_read_plan_with_cast_clipping( + file_schema: &Schema, + schema_descr: &SchemaDescriptor, + whole_root_indices: &[usize], + struct_accesses: &[StructFieldAccess], + cast_accesses: &[CastColumnAccess], +) -> ParquetReadPlan { + let whole_roots: BTreeSet = whole_root_indices.iter().copied().collect(); + let struct_access_roots: BTreeSet = + struct_accesses.iter().map(|a| a.root_index).collect(); + // Every referenced root's Parquet leaves, grouped in one pass over the + // schema descriptor rather than one `leaf_indices_for_roots` scan per + // root (this function may look up several roots). + let leaves_by_root = leaves_grouped_by_root(schema_descr); + + // Root -> (absolute kept leaf indices, cast-clipped Arrow type) for + // roots successfully clipped via a cast. + let mut clipped_by_root: BTreeMap, DataType)> = BTreeMap::new(); + // Roots with a cast access that must fall back to a full read. + let mut fallback_roots: BTreeSet = BTreeSet::new(); + // The cast target already clipped for a root, so a second cast on the + // same root can be recognised as either a repeat (same target: nothing to + // do) or a conflict (different target: neither clip is valid on its own). + let mut clipped_target_by_root: BTreeMap = BTreeMap::new(); + + for access in cast_accesses { + let root = access.root_index; + if whole_roots.contains(&root) || fallback_roots.contains(&root) { + continue; + } + if let Some(previous) = clipped_target_by_root.get(&root) { + if **previous != access.target_type { + // The projection consumes this root through two different + // narrowing casts. Each cast only needs its own leaves, but + // the mask is per column: clipping to the first target would + // silently null-fill whatever the second one needs. Read the + // whole root instead. + clipped_by_root.remove(&root); + clipped_target_by_root.remove(&root); + fallback_roots.insert(root); + } + continue; + } + if struct_access_roots.contains(&root) { + fallback_roots.insert(root); + continue; + } + + let physical_type = file_schema.field(root).data_type(); + let root_leaves = leaves_by_root.get(&root).map_or(&[][..], Vec::as_slice); + + // Defensive: the arrow type's leaf count must agree with the + // Parquet schema (it can diverge if the file embeds a different + // arrow schema). If not, never risk a wrong mask: read the whole + // root. + if root_leaves.len() != count_leaves(physical_type) { + fallback_roots.insert(root); + continue; + } + + match clip_for_cast(physical_type, &access.target_type) { + Some((kept_offsets, pruned_type)) => { + let start = root_leaves[0]; + let absolute = kept_offsets.into_iter().map(|o| start + o).collect(); + clipped_by_root.insert(root, (absolute, pruned_type)); + clipped_target_by_root.insert(root, &access.target_type); + } + // Nothing prunable for this cast: every leaf is consumed. + None => { + fallback_roots.insert(root); + } + } + } + + // `get_field` accesses on roots not already read in full (as a whole + // column, or as a cast that fell back) keep the existing (non-cast) leaf + // resolution. + let get_field_accesses: Vec = struct_accesses + .iter() + .filter(|a| { + // A root carrying a `get_field` access is put into + // `fallback_roots` before any clip is attempted (see the loop + // above), so it can never also be clipped. Assert that rather + // than re-testing it here, so a future reordering trips the + // assert instead of silently changing which leaves are read. + debug_assert!(!clipped_by_root.contains_key(&a.root_index)); + !whole_roots.contains(&a.root_index) + && !fallback_roots.contains(&a.root_index) + }) + .cloned() + .collect(); + + let mut leaf_indices: Vec = Vec::new(); + let mut fields: BTreeMap> = BTreeMap::new(); + + for root in whole_roots.iter().chain(fallback_roots.iter()) { + // A root with no parquet leaves contributes nothing to the mask; + // `ProjectionMask::roots` handles that case the same way, so match it + // rather than indexing and panicking. + if let Some(leaves) = leaves_by_root.get(root) { + leaf_indices.extend(leaves.iter().copied()); + } + fields.insert(*root, Arc::new(file_schema.field(*root).clone())); + } + + for (&root, (kept, pruned_type)) in &clipped_by_root { + leaf_indices.extend(kept.iter().copied()); + fields.insert( + root, + field_with_type(file_schema.field(root), pruned_type.clone()), + ); + } + + if !get_field_accesses.is_empty() { + let get_field_tree = StructAccessTree::from_accesses(&get_field_accesses); + leaf_indices.extend(resolve_struct_field_leaves(&get_field_tree, schema_descr)); + let get_field_schema = build_filter_schema(file_schema, &[], &get_field_tree); + let get_field_roots: BTreeSet = + get_field_accesses.iter().map(|a| a.root_index).collect(); + // `build_filter_schema` emits one field per accessed root in + // ascending root order, which is the order `get_field_roots` iterates + // in, so the two line up positionally. Pairing them beats looking each + // one up by name: no repeated linear scans, and no ambiguity if two + // roots happen to share a name. + debug_assert_eq!(get_field_roots.len(), get_field_schema.fields().len()); + for (root, field) in get_field_roots.iter().zip(get_field_schema.fields()) { + fields.insert(*root, Arc::clone(field)); + } + } + + leaf_indices.sort_unstable(); + leaf_indices.dedup(); + + ParquetReadPlan { + projection_mask: ProjectionMask::leaves( + schema_descr, + leaf_indices.iter().copied(), + ), + projected_schema: Arc::new(Schema::new_with_metadata( + fields.into_values().collect::>(), + file_schema.metadata().clone(), + )), + } +} + +/// Groups every Parquet leaf index by its root (Arrow) column index, in one +/// pass over the schema descriptor. +fn leaves_grouped_by_root( + schema_descr: &SchemaDescriptor, +) -> BTreeMap> { + let mut by_root: BTreeMap> = BTreeMap::new(); + for leaf_idx in 0..schema_descr.num_columns() { + by_root + .entry(schema_descr.get_column_root_idx(leaf_idx)) + .or_default() + .push(leaf_idx); + } + by_root +} + +/// Builds a leaf-level [`ParquetReadPlan`] covering `root_indices` in full plus +/// the individual leaves reached by `struct_field_accesses`. +/// +/// `root_indices` must be sorted, deduplicated indices into `file_schema`. +/// +/// Also returns the resolved Parquet leaf indices, sorted and deduplicated, so +/// callers can size the columns the decoder will read. +pub(crate) fn assemble_read_plan( + root_indices: &[usize], + struct_field_accesses: &[StructFieldAccess], + file_schema: &Schema, + schema_descr: &SchemaDescriptor, +) -> (ParquetReadPlan, Vec) { + let access_tree = StructAccessTree::from_accesses(struct_field_accesses); + + let mut leaf_indices = + leaf_indices_for_roots(root_indices.iter().copied(), schema_descr); + leaf_indices + .extend_from_slice(&resolve_struct_field_leaves(&access_tree, schema_descr)); + leaf_indices.sort_unstable(); + leaf_indices.dedup(); + + let projection_mask = + ProjectionMask::leaves(schema_descr, leaf_indices.iter().copied()); + let projected_schema = build_filter_schema(file_schema, root_indices, &access_tree); + + ( + ParquetReadPlan { + projection_mask, + projected_schema, + }, + leaf_indices, + ) +} + +/// Builds a [`ParquetReadPlan`] that decodes whole root columns. +/// +/// `root_indices` must be sorted, deduplicated indices into `file_schema`. Every +/// leaf below each root is decoded, and the projected schema keeps each root +/// field's full type. Callers that need to decode only some leaves of a struct +/// root must build the plan from leaf indices instead. +fn root_level_plan( + root_indices: &[usize], + file_schema: &Schema, + schema_descr: &SchemaDescriptor, +) -> ParquetReadPlan { + let projection_mask = + ProjectionMask::roots(schema_descr, root_indices.iter().copied()); + let projected_schema = Arc::new( + file_schema + .project(root_indices) + .expect("valid column indices"), + ); + + ParquetReadPlan { + projection_mask, + projected_schema, + } +} + +fn leaf_indices_for_roots( + root_indices: I, + schema_descr: &SchemaDescriptor, +) -> Vec +where + I: IntoIterator, +{ + // Always map root (Arrow) indices to Parquet leaf indices via the schema + // descriptor. Arrow root indices only equal Parquet leaf indices when the + // schema has no group columns (Struct, Map, etc.); when group columns + // exist, their children become separate leaves and shift all subsequent + // leaf indices. + let root_set: BTreeSet<_> = root_indices.into_iter().collect(); + + (0..schema_descr.num_columns()) + .filter(|leaf_idx| { + root_set.contains(&schema_descr.get_column_root_idx(*leaf_idx)) + }) + .collect() +} + +/// Returns the Parquet leaf column indices selected by the access tree. +/// +/// # Matching +/// +/// Iterates Parquet leaves in ascending order (`0..num_columns()`). For each +/// leaf: +/// +/// 1. **Root dispatch.** Look up the leaf's root index — the top-level Arrow +/// column it belongs to — via `SchemaDescriptor::get_column_root_idx`. If +/// that root is absent from the access tree (the filter never touched any +/// field under it), skip the leaf without further work. +/// +/// 2. **Path walk.** Otherwise, take the leaf's dotted column path +/// (`col.path().parts()`), drop the first component (the root field name, +/// already used in step 1), and walk the remaining components against the +/// matching trie subtree via [`leaf_under_tree`]. +/// +/// 3. **Inclusion.** The leaf is added to the result iff the walk reaches a +/// node with `selected_here = true` — either an ancestor along the +/// descent (subsumption: a shallower access subsumes the leaf) or the +/// terminal node reached at the end of the path (exact match). +/// +/// # Returns +/// +/// `Vec` of Parquet leaf column indices. The scan visits each leaf +/// exactly once and pushes in iteration order, so the result is in ascending +/// order and free of duplicates by construction — callers do not need to +/// sort or dedup. +fn resolve_struct_field_leaves( + access_tree: &StructAccessTree<'_>, + schema_descr: &SchemaDescriptor, +) -> Vec { + let mut leaf_indices = Vec::new(); + + for leaf_idx in 0..schema_descr.num_columns() { + let root_idx = schema_descr.get_column_root_idx(leaf_idx); + let Some(root_node) = access_tree.roots.get(&root_idx) else { + continue; + }; + // The first part is the root field name, already used in step 1; walk + // the rest against the tree. + let col = schema_descr.column(leaf_idx); + let Some((_root_name, rest)) = col.path().parts().split_first() else { + continue; + }; + if leaf_under_tree(root_node, rest) { + leaf_indices.push(leaf_idx); + } + } + + leaf_indices +} + +/// True when the leaf path beneath a root is selected by the access tree. +/// +/// A shallower `selected_here` node subsumes deeper accesses: once the walk +/// reaches such a node, every leaf below it is included. +fn leaf_under_tree(mut node: &StructAccessNode<'_>, path: &[String]) -> bool { + for component in path { + if node.selected_here { + return true; + } + let Some(child) = node.children.get(component.as_str()) else { + return false; + }; + node = child; + } + node.selected_here +} + +/// Builds the Arrow schema used to evaluate the filter expression. +/// +/// The returned schema is a **subset** of `file_schema`, restricted to the +/// columns the filter actually touches and (for struct columns accessed +/// only through nested paths) **pruned** to only the accessed fields. +/// +/// # Inputs +/// +/// - `file_schema` — the full file schema; provides the source `Field`s +/// (names, types, nullability, metadata). +/// - `regular_indices` — file-schema column indices the filter references +/// as **whole columns** (non-struct columns, or struct roots referenced +/// in their entirety). Must be sorted, deduplicated. +/// - `access_tree` — the trie of nested struct field accesses recorded by +/// [`PushdownChecker`]. +/// +/// # Behavior +/// +/// The set of columns to include is the union of `regular_indices` and +/// `access_tree.roots.keys()`. For each column index in that union, decide +/// how the field appears in the output: +/// +/// 1. **Whole-column reference** (`idx` is in `regular_indices`). Keep the +/// field's full type unchanged. This is the **whole-root override**: +/// pruning is only valid when a column is accessed *exclusively* through +/// nested field accesses; if any predicate references the whole column, +/// the projected schema must preserve the full type for that column. +/// +/// 2. **Nested-access-only struct root.** Look up the column's node in the +/// access tree and call [`prune_struct_type`] on the field's `DataType` +/// with that node. Wrap the pruned type in a new `Field` carrying the +/// original name and nullability. +/// +/// Column order in the output schema follows ascending file-schema index +/// (via the `BTreeSet` union), matching the order the Parquet reader +/// produces when projecting these columns. +/// +/// # Returns +/// +/// An `Arc` whose fields are a subset of `file_schema`'s, with +/// struct types pruned per the access tree. The schema's metadata is +/// inherited from `file_schema`. +fn build_filter_schema( + file_schema: &Schema, + regular_indices: &[usize], + access_tree: &StructAccessTree<'_>, +) -> SchemaRef { + let regular_set: BTreeSet = regular_indices.iter().copied().collect(); + + let all_indices = regular_indices + .iter() + .copied() + .chain(access_tree.roots.keys().copied()) + .collect::>(); + + let fields = all_indices + .iter() + .map(|&idx| { + let field = file_schema.field(idx); + + // if this column appears as a regular (whole-column) reference, + // keep the full type + // + // Pruning is only valid when the column is accessed exclusively + // through struct field accesses + if regular_set.contains(&idx) { + return Arc::new(field.clone()); + } + + let Some(node) = access_tree.root(idx) else { + return Arc::new(field.clone()); + }; + + let pruned_data_type = prune_struct_type(field.data_type(), node); + Arc::new(Field::new( + field.name(), + pruned_data_type, + field.is_nullable(), + )) + }) + .collect::>(); + + Arc::new(Schema::new_with_metadata( + fields, + file_schema.metadata().clone(), + )) +} + +/// Returns a copy of `dt` with non-accessed struct children removed. +/// +/// # Behavior +/// +/// - If `node.selected_here` is `true`, the input type is returned +/// unchanged. An access path terminates at this node, so the whole +/// subtree (every field of `dt`, recursively) is required. This mirrors +/// the subsumption check in [`leaf_under_tree`] so the projection mask +/// and the projected schema agree even if a producer ever records an +/// access whose `field_path` terminates above a struct. +/// +/// - Otherwise, if `dt` is not a `DataType::Struct`, it is cloned and +/// returned unchanged. The trie only ever guides struct-level pruning; +/// other types pass through. +/// +/// - Otherwise, `dt` is a struct and its fields are iterated in their +/// original order. For each field `f`: +/// 1. Look up `f.name()` in `node.children`. +/// - **Absent.** No access goes through this field. Drop it. +/// - **Present, child node's `selected_here` is `true`.** An access +/// path terminates at this field. Keep the entire subtree by +/// cloning `f` unchanged (`Arc::clone` — no new `Field`). +/// - **Present, child node's `selected_here` is `false`.** Some +/// access goes through this field to a deeper terminal. Recurse +/// into `f.data_type()` with the matching child node, then wrap +/// the pruned type in a fresh `Field` with `f`'s name and +/// nullability. +/// +/// Field ordering is preserved (consumers must match the order the Parquet +/// reader produces when projecting specific leaves). Iterating Arrow's +/// `Fields` directly — rather than iterating `node.children` — is what +/// preserves that order. +/// +/// # Returns +/// +/// A new `DataType::Struct` whose fields are a subset of `dt`'s, restricted +/// to the paths represented by `node`. The original `dt` is not modified. +fn prune_struct_type(dt: &DataType, node: &StructAccessNode<'_>) -> DataType { + if node.selected_here { + // Subsumption: the entire subtree below this node is required. + return dt.clone(); + } + + let DataType::Struct(fields) = dt else { + return dt.clone(); + }; + + let pruned_fields = fields + .iter() + .filter_map(|f| { + let child = node.children.get(f.name().as_str())?; + + let out = if child.selected_here { + // Access path terminates at this field — preserve the whole subtree. + Arc::clone(f) + } else { + // Recurse into nested struct. + let pruned = prune_struct_type(f.data_type(), child); + Arc::new(Field::new(f.name(), pruned, f.is_nullable())) + }; + + Some(out) + }) + .collect::>(); + + DataType::Struct(pruned_fields.into()) +} + +#[cfg(test)] +mod test { + use super::*; + use Column as PhysicalColumn; + use arrow::array::{Int32Array, RecordBatch, StringArray, StructArray}; + use arrow::datatypes::Fields; + use datafusion_common::ScalarValue; + use datafusion_expr::{Expr, col}; + use datafusion_functions::core::get_field; + use datafusion_physical_expr::planner::logical2physical; + use parquet::arrow::ArrowWriter; + use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder; + use parquet::file::metadata::ParquetMetaData; + use tempfile::NamedTempFile; + + #[test] + fn projection_read_plan_preserves_full_struct() { + // Schema: id (Int32), s (Struct{value: Int32, label: Utf8}) + // Parquet leaves: id=0, s.value=1, s.label=2 + let struct_fields: Fields = vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + ] + .into(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(struct_fields.clone()), false), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StructArray::new( + struct_fields, + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])) as _, + Arc::new(StringArray::from(vec!["a", "b", "c"])) as _, + ], + None, + )), + ], + ) + .unwrap(); + + let file = NamedTempFile::new().expect("temp file"); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) + .expect("writer"); + writer.write(&batch).expect("write batch"); + writer.close().expect("close writer"); + + let reader_file = file.reopen().expect("reopen file"); + let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file) + .expect("reader builder"); + let metadata = builder.metadata().clone(); + let file_schema = builder.schema().clone(); + let schema_descr = metadata.file_metadata().schema_descr(); + + // Simulate SELECT * output projection: Column("id") and Column("s") + // Plus a get_field(s, 'value') expression from the pushed-down filter + let exprs: Vec> = vec![ + Arc::new(PhysicalColumn::new("id", 0)), + Arc::new(PhysicalColumn::new("s", 1)), + logical2physical( + &get_field().call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("value".to_string())), None), + ]), + &file_schema, + ), + ]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // The projected schema must have the FULL struct type because Column("s") + // is in the projection. It should NOT be narrowed to Struct{value: Int32}. + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!( + s_field.data_type(), + &DataType::Struct( + vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + ] + .into() + ), + ); + + // all 3 Parquet leaves should be in the projection mask + let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1, 2]); + assert_eq!(read_plan.projection_mask, expected_mask,); + } + + /// Writes the id/struct fixture and returns the schema and metadata a + /// reader sees for it, so callers don't each repeat the reopen + + /// `ParquetRecordBatchReaderBuilder` boilerplate. + /// + /// Schema: id (Int32), s (Struct{value: Int32, label: Utf8, pad: Utf8}). + /// Parquet leaves: id=0, s.value=1, s.label=2, s.pad=3. + fn write_id_struct_file() -> (SchemaRef, Arc) { + let struct_fields: Fields = vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(Field::new("pad", DataType::Utf8, false)), + ] + .into(); + + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(struct_fields.clone()), false), + ])); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(StructArray::new( + struct_fields, + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])) as _, + Arc::new(StringArray::from(vec!["a", "b", "c"])) as _, + Arc::new(StringArray::from(vec!["p0", "p1", "p2"])) as _, + ], + None, + )), + ], + ) + .unwrap(); + + let file = NamedTempFile::new().expect("temp file"); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) + .expect("writer"); + writer.write(&batch).expect("write batch"); + writer.close().expect("close writer"); + + let builder = ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()) + .expect("reader builder"); + (builder.schema().clone(), builder.metadata().clone()) + } + + /// Writes a two-struct-root fixture so tests can combine a cast on one + /// root with an access on another. + /// + /// Schema: a (Struct{p: Int32, q: Utf8}), b (Struct{m: Int32, n: Utf8}). + /// Parquet leaves: a.p=0, a.q=1, b.m=2, b.n=3. + fn write_two_struct_file() -> (SchemaRef, Arc) { + let group = |first: &str, second: &str| -> Fields { + vec![ + Arc::new(Field::new(first, DataType::Int32, false)), + Arc::new(Field::new(second, DataType::Utf8, false)), + ] + .into() + }; + let (a_fields, b_fields) = (group("p", "q"), group("m", "n")); + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Struct(a_fields.clone()), false), + Field::new("b", DataType::Struct(b_fields.clone()), false), + ])); + + let values = |fields: Fields, ints: [i32; 2], strs: [&str; 2]| { + Arc::new(StructArray::new( + fields, + vec![ + Arc::new(Int32Array::from(ints.to_vec())) as _, + Arc::new(StringArray::from(strs.to_vec())) as _, + ], + None, + )) as _ + }; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + values(a_fields, [1, 2], ["a0", "a1"]), + values(b_fields, [3, 4], ["b0", "b1"]), + ], + ) + .unwrap(); + + let file = NamedTempFile::new().expect("temp file"); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) + .expect("writer"); + writer.write(&batch).expect("write batch"); + writer.close().expect("close writer"); + + let builder = ParquetRecordBatchReaderBuilder::try_new(file.reopen().unwrap()) + .expect("reader builder"); + (builder.schema().clone(), builder.metadata().clone()) + } + + /// Builds `CAST(Column(name, index) AS Struct{fields})`. + fn cast_to_struct( + name: &str, + index: usize, + fields: Vec<(&str, DataType)>, + ) -> Arc { + let target = DataType::Struct( + fields + .into_iter() + .map(|(n, dt)| Arc::new(Field::new(n, dt, true))) + .collect::>() + .into(), + ); + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new(name, index)), + target, + None, + )) + } + + /// Builds `get_field(Column(name, index), field)`. + fn get_field_of( + file_schema: &Schema, + name: &str, + field: &str, + ) -> Arc { + logical2physical( + &get_field().call(vec![ + col(name), + Expr::Literal(ScalarValue::Utf8(Some(field.to_string())), None), + ]), + file_schema, + ) + } + + /// Clipping a cast whose only surviving field is *not* the struct's first + /// one: the kept offsets are relative to the root's first leaf and must be + /// rebased onto it. With `s` starting at leaf 1 and `label` at offset 1, + /// getting the arithmetic wrong reads `id` (leaf 0) instead of `s.label`. + #[test] + fn build_projection_read_plan_clips_cast_to_a_non_leading_field() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs = vec![cast_to_struct("s", 1, vec![("label", DataType::Utf8)])]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [2]) + ); + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!( + s_field.data_type(), + &DataType::Struct( + vec![Arc::new(Field::new("label", DataType::Utf8, false))].into() + ), + ); + } + + /// A cast on one root and a `get_field` on a *different* root: each root + /// keeps only what it needs, and both appear in the projected schema in + /// root order. + #[test] + fn build_projection_read_plan_clips_cast_beside_get_field_on_another_root() { + let (file_schema, metadata) = write_two_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs = vec![ + cast_to_struct("a", 0, vec![("p", DataType::Int32)]), + get_field_of(&file_schema, "b", "n"), + ]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // a.p (leaf 0) from the clip, b.n (leaf 3) from the field access. + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [0, 3]) + ); + let field_types = read_plan + .projected_schema + .fields() + .iter() + .map(|f| (f.name().clone(), f.data_type().clone())) + .collect::>(); + assert_eq!( + field_types, + vec![ + ( + "a".to_string(), + DataType::Struct( + vec![Arc::new(Field::new("p", DataType::Int32, false))].into() + ) + ), + ( + "b".to_string(), + DataType::Struct( + vec![Arc::new(Field::new("n", DataType::Utf8, false))].into() + ) + ), + ] + ); + } + + /// Once conflicting cast targets have demoted a root to a full read, a + /// *third* cast on it must not resurrect the clip. + #[test] + fn build_projection_read_plan_keeps_full_read_after_a_third_cast() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs = vec![ + cast_to_struct("s", 1, vec![("value", DataType::Int32)]), + cast_to_struct("s", 1, vec![("label", DataType::Utf8)]), + cast_to_struct("s", 1, vec![("value", DataType::Int32)]), + ]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1, 2, 3]) + ); + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!(s_field.data_type(), file_schema.field(1).data_type()); + } + + /// A whole-column reference wins over a `get_field` access on the same + /// root even when another root is being clipped: `a` keeps every leaf and + /// its full type, `b` keeps only the cast target's. + #[test] + fn build_projection_read_plan_whole_column_beats_get_field_beside_a_clip() { + let (file_schema, metadata) = write_two_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let exprs: Vec> = vec![ + Arc::new(PhysicalColumn::new("a", 0)), + get_field_of(&file_schema, "a", "p"), + cast_to_struct("b", 1, vec![("m", DataType::Int32)]), + ]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // Every leaf of `a` (0, 1) plus b.m (leaf 2). + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [0, 1, 2]) + ); + let a_field = read_plan.projected_schema.field_with_name("a").unwrap(); + assert_eq!( + a_field.data_type(), + file_schema.field(0).data_type(), + "the whole-column reference must keep `a`'s full type" + ); + } + + /// Columns are resolved by *name*: a `Column` whose index points at a + /// different field (a stale index left by an earlier rewrite) must not be + /// taken at face value by the struct fast-path gate. + #[test] + fn build_projection_read_plan_resolves_stale_column_indices_by_name() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + // `s` is at index 1; this claims index 0, which is `id`. + let exprs = vec![cast_to_struct("s", 0, vec![("value", DataType::Int32)])]; + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1]), + "the cast must resolve to `s`, not to whatever sits at index 0" + ); + } + + /// A projection consisting solely of a narrowing cast over a struct root + /// clips the read to the cast target's leaves. + #[test] + fn build_projection_read_plan_clips_cast_over_struct() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let narrow = DataType::Struct( + vec![Arc::new(Field::new("value", DataType::Int32, true))].into(), + ); + let exprs: Vec> = vec![ + Arc::new(PhysicalColumn::new("id", 0)), + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("s", 1)), + narrow.clone(), + None, + )), + ]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // Only id's leaf (0) and s.value's leaf (1) should be read: s.label + // and s.pad are clipped away. + let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1]); + assert_eq!(read_plan.projection_mask, expected_mask); + + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!( + s_field.data_type(), + &DataType::Struct( + vec![Arc::new(Field::new("value", DataType::Int32, false))].into() + ), + ); + } + + /// Two casts on the same root with the *same* target still clip: this is + /// the shape the expression adapter produces when one column is + /// referenced several times (`SELECT s, s FROM narrowed`). + #[test] + fn build_projection_read_plan_clips_repeated_identical_casts() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let narrow = DataType::Struct( + vec![Arc::new(Field::new("value", DataType::Int32, true))].into(), + ); + let cast = || -> Arc { + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("s", 1)), + narrow.clone(), + None, + )) + }; + + let read_plan = + build_projection_read_plan(vec![cast(), cast()], &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1]) + ); + } + + /// Two casts on the same root with *different* targets cannot both be + /// served by one mask: clipping to either target alone would null-fill + /// whatever the other one needs (or fail its runtime struct-compatibility + /// check outright). Read the whole root instead. + #[test] + fn build_projection_read_plan_falls_back_on_conflicting_cast_targets() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let narrow = |name: &str, dt: DataType| -> Arc { + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("s", 1)), + DataType::Struct(vec![Arc::new(Field::new(name, dt, true))].into()), + None, + )) + }; + let exprs = vec![ + narrow("value", DataType::Int32), + narrow("label", DataType::Utf8), + ]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::leaves(schema_descr, [1, 2, 3]), + "every leaf of `s` must be read so both casts see their fields" + ); + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!(s_field.data_type(), file_schema.field(1).data_type()); + } + + /// The struct fast-path gate looks at the *projected* columns, not at + /// every field of the file schema: projecting only `id` produces the same + /// root-level plan it would for a schema with no struct in it at all. + #[test] + fn build_projection_read_plan_ignores_unprojected_struct_columns() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + // Not a bare column, so the all-plain-columns fast path does not apply. + let exprs: Vec> = vec![Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("id", 0)), + DataType::Int64, + None, + ))]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + assert_eq!( + read_plan.projection_mask, + ProjectionMask::roots(schema_descr, [0]) + ); + assert_eq!(read_plan.projected_schema.fields().len(), 1); + } + + /// A root reached by both a narrowing cast and a `get_field` access (not + /// producible by `DefaultPhysicalExprAdapter`, but a custom + /// `PhysicalExprAdapter` could inject both) falls back to a full read of + /// that root rather than attempting to union the two leaf sets. + #[test] + fn build_projection_read_plan_falls_back_when_cast_and_get_field_share_a_root() { + let (file_schema, metadata) = write_id_struct_file(); + let schema_descr = metadata.file_metadata().schema_descr(); + + let narrow = DataType::Struct( + vec![Arc::new(Field::new("value", DataType::Int32, true))].into(), + ); + let exprs: Vec> = vec![ + Arc::new(CastExpr::new( + Arc::new(PhysicalColumn::new("s", 1)), + narrow, + None, + )), + logical2physical( + &get_field().call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("label".to_string())), None), + ]), + &file_schema, + ), + ]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // Every leaf of `s` is read (full fallback), not just value/label. + let expected_mask = ProjectionMask::leaves(schema_descr, [1, 2, 3]); + assert_eq!(read_plan.projection_mask, expected_mask); + + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!( + s_field.data_type(), + &DataType::Struct( + vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(Field::new("pad", DataType::Utf8, false)), + ] + .into() + ), + ); + } + + fn access(root: usize, path: &[&str]) -> StructFieldAccess { + StructFieldAccess { + root_index: root, + field_path: path.iter().map(|&s| s.to_string()).collect(), + } + } + + #[test] + fn struct_access_tree_from_empty_input_has_no_roots() { + let tree = StructAccessTree::from_accesses(&[]); + assert!(tree.roots.is_empty()); + } + + #[test] + fn struct_access_tree_groups_paths_by_root() { + let accesses = [access(0, &["a"]), access(2, &["x"]), access(2, &["y"])]; + let tree = StructAccessTree::from_accesses(&accesses); + + assert_eq!(tree.roots.keys().copied().collect::>(), vec![0, 2]); + let root0 = tree.root(0).unwrap(); + assert!(root0.children.contains_key("a")); + assert!(root0.children["a"].selected_here); + + let root2 = tree.root(2).unwrap(); + assert_eq!( + root2.children.keys().copied().collect::>(), + vec!["x", "y"], + ); + } + + #[test] + fn struct_access_tree_shared_prefix_collapses_into_one_node() { + let accesses = [access(0, &["outer", "a"]), access(0, &["outer", "b"])]; + let tree = StructAccessTree::from_accesses(&accesses); + + let root = tree.root(0).unwrap(); + assert!(!root.selected_here); + + let outer = &root.children["outer"]; + // `outer` itself was never the terminal of an access path. + assert!(!outer.selected_here); + // Both leaves below share the single `outer` node. + assert_eq!( + outer.children.keys().copied().collect::>(), + vec!["a", "b"], + ); + assert!(outer.children["a"].selected_here); + assert!(outer.children["b"].selected_here); + } + + #[test] + fn struct_access_tree_records_both_shallow_and_deep_selection() { + // `s['outer']` (whole subtree) and `s['outer']['a']` (specific leaf) + // both recorded. Consumers honor the shallower selection at walk time; + // the builder simply records both `selected_here` flags. + let accesses = [access(0, &["outer"]), access(0, &["outer", "a"])]; + let tree = StructAccessTree::from_accesses(&accesses); + + let outer = &tree.root(0).unwrap().children["outer"]; + assert!(outer.selected_here); + assert!(outer.children["a"].selected_here); + } + + /// `prune_struct_type` must honor `selected_here` on the input node + /// itself, not only on its children — symmetric with `leaf_under_tree`. + /// Without this guard, a node with `selected_here = true` and no + /// children produces an empty struct (silent drift from the leaf set). + #[test] + fn prune_struct_type_returns_full_type_when_node_is_selected_here() { + let node = StructAccessNode { + selected_here: true, + ..Default::default() + }; + + let s_type = DataType::Struct( + vec![ + Arc::new(Field::new("outer", DataType::Int32, false)), + Arc::new(Field::new("other", DataType::Int32, false)), + ] + .into(), + ); + + let pruned = prune_struct_type(&s_type, &node); + + assert_eq!( + pruned, s_type, + "selected_here on the input node must preserve the full type" + ); + } + + /// Same guard, but for the case where `selected_here` is set on an + /// intermediate node that also has children — e.g. both `s['outer']` + /// and `s['outer']['a']` are recorded. The shallower terminal must + /// keep the entire `outer` subtree, ignoring the deeper child entry. + #[test] + fn prune_struct_type_shallow_selection_subsumes_deeper_children() { + let accesses = [access(0, &["outer"]), access(0, &["outer", "a"])]; + let tree = StructAccessTree::from_accesses(&accesses); + + let outer_type = DataType::Struct( + vec![ + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Field::new("b", DataType::Int32, false)), + ] + .into(), + ); + + let outer_node = &tree.root(0).unwrap().children["outer"]; + let pruned = prune_struct_type(&outer_type, outer_node); + + assert_eq!( + pruned, outer_type, + "shallow selected_here must preserve the whole subtree, \ + not narrow to the deeper child" + ); + } + + /// Mixed whole-root and nested access. + /// Projecting `s` (whole) alongside `get_field(s, 'outer', 'a')` (nested) + /// must preserve the full `s` struct type AND include all `s` leaves in + /// the projection mask. The nested access does not narrow the whole-root + /// reference — `regular_indices` wins over the access tree for that root. + #[test] + fn projection_whole_root_plus_nested_access_keeps_full_struct() { + // Schema: s (Struct{outer: Struct{a, b}}) + // Parquet leaves: s.outer.a=0, s.outer.b=1 + let outer_fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Field::new("b", DataType::Int32, false)), + ] + .into(); + let s_fields: Fields = vec![Arc::new(Field::new( + "outer", + DataType::Struct(outer_fields.clone()), + false, + ))] + .into(); + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(s_fields.clone()), + false, + )])); + + let outer_arr = StructArray::new( + outer_fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![1, 2])) as _, + Arc::new(Int32Array::from(vec![3, 4])) as _, + ], + None, + ); + let s_arr = + StructArray::new(s_fields.clone(), vec![Arc::new(outer_arr) as _], None); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s_arr)]).unwrap(); + + let file = NamedTempFile::new().expect("temp file"); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) + .expect("writer"); + writer.write(&batch).expect("write batch"); + writer.close().expect("close writer"); + + let reader_file = file.reopen().expect("reopen file"); + let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file) + .expect("reader builder"); + let metadata = builder.metadata().clone(); + let file_schema = builder.schema().clone(); + let schema_descr = metadata.file_metadata().schema_descr(); + + // Column("s") (whole struct) + get_field(s, 'outer', 'a') (nested access). + let exprs: Vec> = vec![ + Arc::new(PhysicalColumn::new("s", 0)), + logical2physical( + &get_field().call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("outer".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("a".to_string())), None), + ]), + &file_schema, + ), + ]; + + let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + + // `s` must keep its full nested type — NOT narrowed to Struct{outer: Struct{a}}. + let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + assert_eq!( + s_field.data_type(), + &DataType::Struct(s_fields), + "whole-root reference must preserve the full nested struct type \ + even when a nested access is also recorded" + ); + + // All `s` leaves must be in the projection mask (s.outer.a AND s.outer.b). + let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1]); + assert_eq!( + read_plan.projection_mask, expected_mask, + "whole-root reference must select every leaf under the root" + ); + } +} diff --git a/datafusion/datasource-parquet/src/push_decoder.rs b/datafusion/datasource-parquet/src/push_decoder.rs index 8b71be3e8de96..74d8997198872 100644 --- a/datafusion/datasource-parquet/src/push_decoder.rs +++ b/datafusion/datasource-parquet/src/push_decoder.rs @@ -19,14 +19,18 @@ //! //! This module owns the push-decoder lifecycle: //! -//! - [`DecoderBuilderConfig`] holds the shared options applied to every -//! [`ParquetPushDecoderBuilder`] in a file scan, exposing a single `build` -//! entry point per decoder run. -//! - [`PushDecoderStreamState`] is the per-file stream driver that polls one -//! or more decoders to completion, yielding projected [`RecordBatch`]es. -//! A scan can produce multiple decoders (for example, when fully matched -//! row groups split it into runs with different filter requirements); the -//! state machine drains them in order so the output is contiguous. +//! - [`DecoderBuilderConfig`] holds the shared options applied to the +//! [`ParquetPushDecoderBuilder`] for a file scan, exposing a single `build` +//! entry point. +//! - [`PushDecoderStreamState`] is the per-file stream driver. It owns a +//! **single** [`ParquetPushDecoder`] plus an [`RgPlanEntry`] queue +//! (`rg_plan`) and uses arrow-rs's [`ParquetRecordBatchReader`] iterator +//! to pause at row-group boundaries. At each boundary the optional +//! [`RowGroupPruner`] is consulted; row groups it proves unwinnable are +//! dropped from the head of `rg_plan` and the decoder is rebuilt via +//! [`ParquetPushDecoder::into_builder`] + +//! [`ParquetPushDecoderBuilder::with_row_groups`] so the skipped RGs are +//! bypassed entirely — no decode, no row-filter eval. //! //! The opener constructs both halves and hands the state off to //! [`PushDecoderStreamState::into_stream`] for consumption. @@ -34,31 +38,39 @@ use std::collections::VecDeque; use std::sync::Arc; -use arrow::array::{RecordBatch, RecordBatchOptions}; -use arrow::datatypes::Schema; +use arrow::array::RecordBatch; +use arrow::datatypes::SchemaRef; use futures::StreamExt; use futures::stream::BoxStream; +use log::debug; use parquet::DecodeResult; +use parquet::arrow::ProjectionMask; use parquet::arrow::arrow_reader::metrics::ArrowReaderMetrics; -use parquet::arrow::arrow_reader::{ArrowReaderMetadata, RowSelectionPolicy}; +use parquet::arrow::arrow_reader::{ + ArrowReaderMetadata, ParquetRecordBatchReader, RowSelectionPolicy, +}; use parquet::arrow::async_reader::AsyncFileReader; use parquet::arrow::push_decoder::{ParquetPushDecoder, ParquetPushDecoderBuilder}; +use parquet::file::metadata::ParquetMetaData; -use datafusion_common::{DataFusionError, Result}; -use datafusion_physical_expr::projection::Projector; -use datafusion_physical_plan::metrics::{BaselineMetrics, Gauge}; +use datafusion_common::{DataFusionError, Result, internal_err}; +use datafusion_physical_expr::expressions::DynamicFilterTracking; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_plan::metrics::{BaselineMetrics, Count, Gauge}; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; use crate::access_plan::PreparedAccessPlan; -use crate::row_filter::ParquetReadPlan; +use crate::decoder_projection::DecoderProjection; +use crate::row_group_filter::RowGroupPruningStatistics; -/// Shared options applied to every [`ParquetPushDecoderBuilder`] in a file scan. -/// -/// A single scan may produce multiple decoders (for example, when fully matched -/// row groups split the scan into consecutive runs with different filter -/// requirements). All decoders in that scan share the same projection, batch -/// size, metrics sink, and selection policy. +/// Shared options applied to the [`ParquetPushDecoderBuilder`] for a file +/// scan, and to any later rebuilds performed via +/// [`ParquetPushDecoder::into_builder`] at row-group boundaries (e.g. when +/// the [`RowGroupPruner`] drops subsequent row groups). pub(crate) struct DecoderBuilderConfig<'a> { - pub(crate) read_plan: &'a ParquetReadPlan, + /// Projection mask installed on every decoder in the scan. Sourced from + /// the file's [`DecoderProjection`]. + pub(crate) projection_mask: &'a ProjectionMask, pub(crate) batch_size: usize, pub(crate) arrow_reader_metrics: &'a ArrowReaderMetrics, pub(crate) force_filter_selections: bool, @@ -66,9 +78,9 @@ pub(crate) struct DecoderBuilderConfig<'a> { } impl DecoderBuilderConfig<'_> { - /// Build a [`ParquetPushDecoderBuilder`] for a single decoder run. + /// Build a [`ParquetPushDecoderBuilder`] from a prepared access plan. /// - /// The caller is expected to attach the run-specific + /// The caller is expected to attach the /// [`RowFilter`](parquet::arrow::arrow_reader::RowFilter) and predicate /// cache size on the returned builder. pub(crate) fn build( @@ -77,7 +89,7 @@ impl DecoderBuilderConfig<'_> { metadata: ArrowReaderMetadata, ) -> ParquetPushDecoderBuilder { let mut builder = ParquetPushDecoderBuilder::new_with_metadata(metadata) - .with_projection(self.read_plan.projection_mask.clone()) + .with_projection(self.projection_mask.clone()) .with_batch_size(self.batch_size) .with_metrics(self.arrow_reader_metrics.clone()); if self.force_filter_selections { @@ -94,6 +106,140 @@ impl DecoderBuilderConfig<'_> { } } +#[derive(Debug, Clone)] +pub(crate) struct RgPlanEntry { + pub(crate) rg_index: usize, +} + +/// Runtime row-group pruner driven by a dynamic predicate (e.g. the +/// threshold expression a `TopK` operator pushes down). +/// +/// Mirrors the [`FilePruner`](datafusion_pruning::FilePruner) pattern at +/// the row-group level: subscribes once to every still-incomplete dynamic +/// filter inside the predicate via +/// [`DynamicFilterTracker`](datafusion_physical_expr::expressions::DynamicFilterTracker) +/// and only rebuilds the [`PruningPredicate`] when one of those +/// subscriptions reports an update, then evaluates the cached predicate +/// against the statistics of the requested row groups. +pub(crate) struct RowGroupPruner { + predicate: Arc, + arrow_schema: SchemaRef, + parquet_metadata: Arc, + /// Classifies the predicate's dynamic-filter content. The `Watching` + /// variant carries a tracker that subscribes to every not-yet-complete + /// dynamic filter; for `Static` / `AllComplete` the predicate cannot + /// change so a single up-front `pruning_predicate` build suffices. + tracking: DynamicFilterTracking, + /// First-call sentinel: forces an initial `pruning_predicate` build + /// even when `tracking` is `Static` / `AllComplete`. + needs_initial_build: bool, + /// Cached pruning predicate. `None` means we couldn't build one for the + /// current generation (e.g. the predicate has no analyzable bounds); + /// in that case we conservatively don't prune. + pruning_predicate: Option>, + /// Metric for `build_pruning_predicate` failures (predicate creation). + predicate_creation_errors: Count, + /// Metric for `PruningPredicate::prune` failures (evaluating an + /// already-built predicate against row-group statistics). + predicate_evaluation_errors: Count, + /// Cap on the `IN (...)` list size that the pruning predicate will + /// rewrite into per-value statistics checks. Longer lists skip + /// container-level pruning. Sourced from + /// `datafusion.execution.parquet.max_in_list_size`. + max_in_list_size: usize, +} + +impl RowGroupPruner { + pub(crate) fn new( + predicate: Arc, + arrow_schema: SchemaRef, + parquet_metadata: Arc, + predicate_creation_errors: Count, + predicate_evaluation_errors: Count, + max_in_list_size: usize, + ) -> Self { + let tracking = DynamicFilterTracking::classify(&predicate); + Self { + predicate, + arrow_schema, + parquet_metadata, + tracking, + needs_initial_build: true, + pruning_predicate: None, + predicate_creation_errors, + predicate_evaluation_errors, + max_in_list_size, + } + } + + /// Returns `true` when the statistics for `row_group_indices` prove that + /// every requested row group can be skipped under the current value of + /// the dynamic predicate. + /// + /// On any error (predicate construction, statistics evaluation) the + /// pruner conservatively returns `false` and logs the failure, so a + /// flaky pruning path never silently drops data. + pub(crate) fn should_prune(&mut self, row_group_indices: &[usize]) -> bool { + if row_group_indices.is_empty() { + return false; + } + + // Refresh the cached `PruningPredicate` on the first call and + // whenever a watched dynamic filter has advanced since we last + // looked. `changed()` is a single atomic load per still-incomplete + // filter — no tree walk on every check. + let dynamic_changed = self + .tracking + .watcher() + .is_some_and(|tracker| tracker.changed()); + if self.needs_initial_build || dynamic_changed { + self.pruning_predicate = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&self.arrow_schema)) + .with_error_counter(&self.predicate_creation_errors) + .with_max_in_list_size(self.max_in_list_size) + .build(Arc::clone(&self.predicate)); + self.needs_initial_build = false; + } + + let Some(pp) = self.pruning_predicate.as_ref() else { + return false; + }; + + let row_group_metadatas = row_group_indices + .iter() + .map(|&i| self.parquet_metadata.row_group(i)) + .collect::>(); + let stats = RowGroupPruningStatistics { + parquet_schema: self.parquet_metadata.file_metadata().schema_descr(), + row_group_metadatas, + arrow_schema: self.arrow_schema.as_ref(), + // Match the existing static row-group pruning behavior: when a + // statistic's null count is missing, treat it as zero. This is + // sound for runtime pruning because the predicate only needs to + // prove a row group *cannot* contain matching rows. + missing_null_counts_as_zero: true, + }; + + match pp.prune(&stats) { + // `prune` returns `false` per container that the predicate proves + // cannot contain matching rows. We can skip the run only when + // every requested row group is in that state. + Ok(values) => values.iter().all(|&keep| !keep), + Err(e) => { + // The predicate was already built successfully (we hold `pp`); + // this failure is in *evaluating* it against the row-group + // stats, so it belongs in the evaluation-errors counter, not + // creation-errors. + debug!( + "Ignoring error evaluating runtime row-group pruning predicate: {e}" + ); + self.predicate_evaluation_errors.add(1); + false + } + } + } +} + /// State for a stream that decodes a single Parquet file using a push-based decoder. /// /// The [`transition`](Self::transition) method drives the decoder in a loop: it requests @@ -101,25 +247,31 @@ impl DecoderBuilderConfig<'_> { /// [`ParquetPushDecoder`], and yields projected [`RecordBatch`]es until the file is /// fully consumed. pub(crate) struct PushDecoderStreamState { - pub(crate) decoder: ParquetPushDecoder, - /// Additional decoders to process after the current one finishes. - /// Used when fully matched row groups split the scan into consecutive - /// runs with different filter configurations, maintaining original order. - pub(crate) pending_decoders: VecDeque, - /// Global remaining row limit across all decoder runs. - /// - /// Decoder-local limits are only safe for single-run scans. When the scan - /// is split across multiple decoders, the combined stream limit is enforced - /// here instead. - pub(crate) remaining_limit: Option, + pub(crate) decoder: Option, + pub(crate) active_reader: Option, + pub(crate) rg_plan: VecDeque, pub(crate) reader: Box, - pub(crate) projector: Projector, - pub(crate) output_schema: Arc, - pub(crate) replace_schema: bool, + /// Per-file projection: the mask installed on every decoder and the + /// per-batch transform applied by [`Self::project_batch`]. + pub(crate) decoder_projection: DecoderProjection, pub(crate) arrow_reader_metrics: ArrowReaderMetrics, pub(crate) predicate_cache_inner_records: Gauge, pub(crate) predicate_cache_records: Gauge, pub(crate) baseline_metrics: BaselineMetrics, + /// Dynamic row-group pruner consulted at every row-group boundary. + /// + /// When the file scan was opened with a still-watching dynamic predicate + /// (typically the threshold expression a `TopK` `SortExec` pushed down), + /// we re-evaluate that predicate against the next pending RG's + /// statistics and drop RGs the current threshold proves cannot + /// contribute. The decoder is rebuilt via + /// [`ParquetPushDecoder::into_builder`] + + /// [`ParquetPushDecoderBuilder::with_row_groups`] so the skipped RGs are + /// bypassed entirely. `None` when the scan has no watching dynamic + /// predicate or only one row group remains. + pub(crate) row_group_pruner: Option, + /// Count of row groups skipped at runtime by [`Self::row_group_pruner`]. + pub(crate) row_groups_pruned_dynamic: Count, } impl PushDecoderStreamState { @@ -148,10 +300,99 @@ impl PushDecoderStreamState { /// with `unfold`'s ownership across yield points. async fn transition(mut self) -> Option<(Result, Self)> { loop { - if self.remaining_limit == Some(0) { - return None; + // Step 1: drain a batch from the active reader if any. + if let Some(reader) = self.active_reader.as_mut() { + match reader.next() { + Some(Ok(batch)) => { + let mut timer = self.baseline_metrics.elapsed_compute().timer(); + self.copy_arrow_reader_metrics(); + let result = self.project_batch(&batch); + timer.stop(); + drop(timer); + return Some((result, self)); + } + Some(Err(e)) => { + return Some((Err(DataFusionError::from(e)), self)); + } + None => { + // Reader exhausted: drop and fall through to per-RG + // boundary handling, then try_next_reader. + self.active_reader = None; + } + } + } + + // Step 2: when the decoder is sitting on a row-group boundary, + // scan the entire `rg_plan` and drop every RG the pruner proves + // cannot contribute — head, interior, and tail alike. Evaluating + // per-RG stats against the cached `PruningPredicate` is cheap; + // the expensive part is the `into_builder` rebuild, so we do at + // most one rebuild per boundary regardless of how many RGs were + // dropped. Buffered bytes for already-fetched RGs carry across + // the rebuild. + // + // `into_builder` errors out mid-row-group, so we gate the prune + // pass on `is_at_row_group_boundary()`. When the decoder is + // mid-RG (e.g. byte ranges have been pushed but no reader has + // been handed back yet), step 3 drives it forward and we get + // another chance at the next boundary — the pruner is stateful + // and idempotent, so deferring loses nothing. + let at_boundary = self + .decoder + .as_ref() + .expect("decoder present") + .is_at_row_group_boundary(); + // Only the runtime pruner rebuilds the decoder from `rg_plan`, so + // only it needs `rg_plan` kept in sync with the decoder frontier. + // arrow-rs silently finishes row groups whose post-predicate + // selection is empty without handing back a reader, so without this + // sync `rg_plan` trails the decoder by one and a rebuild re-reads an + // already-delivered row group (#24352). Gating on the pruner also + // avoids the O(remaining row groups) cost of `peek_next_row_group()` + // on ordinary scans that never rebuild. + if at_boundary + && self.row_group_pruner.is_some() + && let Err(e) = self.sync_rg_plan_to_decoder_frontier() + { + return Some((Err(e), self)); + } + if at_boundary && !self.rg_plan.is_empty() { + let mut pruned_count = 0usize; + if let Some(pruner) = self.row_group_pruner.as_mut() { + let mut kept = VecDeque::with_capacity(self.rg_plan.len()); + while let Some(entry) = self.rg_plan.pop_front() { + if pruner.should_prune(&[entry.rg_index]) { + pruned_count += 1; + self.row_groups_pruned_dynamic.add(1); + } else { + kept.push_back(entry); + } + } + self.rg_plan = kept; + } + if pruned_count > 0 { + if self.rg_plan.is_empty() { + return None; + } + let decoder = self.decoder.take().expect("decoder present"); + let new_indices: Vec = + self.rg_plan.iter().map(|e| e.rg_index).collect(); + let rebuilt = match decoder.into_builder() { + Ok(b) => b.with_row_groups(new_indices).build(), + Err(e) => Err(e), + }; + match rebuilt { + Ok(d) => self.decoder = Some(d), + Err(e) => { + return Some((Err(DataFusionError::from(e)), self)); + } + } + } } - match self.decoder.try_decode() { + + // Step 3: drive the decoder. + let decoder = self.decoder.as_mut().expect("decoder present"); + match decoder.try_next_reader() { Ok(DecodeResult::NeedsData(ranges)) => { let data = self .reader @@ -160,43 +401,26 @@ impl PushDecoderStreamState { .map_err(DataFusionError::from); match data { Ok(data) => { - if let Err(e) = self.decoder.push_ranges(ranges, data) { + if let Err(e) = self + .decoder + .as_mut() + .expect("decoder present") + .push_ranges(ranges, data) + { return Some((Err(DataFusionError::from(e)), self)); } } Err(e) => return Some((Err(e), self)), } } - Ok(DecodeResult::Data(batch)) => { - let batch = if let Some(remaining_limit) = self.remaining_limit { - if batch.num_rows() > remaining_limit { - self.remaining_limit = Some(0); - batch.slice(0, remaining_limit) - } else { - self.remaining_limit = - Some(remaining_limit - batch.num_rows()); - batch - } - } else { - batch - }; - let mut timer = self.baseline_metrics.elapsed_compute().timer(); - self.copy_arrow_reader_metrics(); - let result = self.project_batch(&batch); - timer.stop(); - // Release the borrow on baseline_metrics before moving self - drop(timer); - return Some((result, self)); - } - Ok(DecodeResult::Finished) => { - // If there are pending decoders (e.g. for consecutive runs - // with different filter configurations), switch to the next. - if let Some(next) = self.pending_decoders.pop_front() { - self.decoder = next; - continue; - } - return None; + Ok(DecodeResult::Data(reader)) => { + // Pop the RG this reader is for (we already filtered + // pruned ones in step 2, so `rg_plan.front()` is the RG + // the decoder is about to read). + self.rg_plan.pop_front(); + self.active_reader = Some(reader); } + Ok(DecodeResult::Finished) => return None, Err(e) => { return Some((Err(DataFusionError::from(e)), self)); } @@ -204,6 +428,51 @@ impl PushDecoderStreamState { } } + /// Keep `rg_plan.front()` aligned with the row group the decoder will emit + /// next. `try_next_reader` silently finishes row groups whose post-predicate + /// selection is empty (no reader handed back), which would otherwise leave + /// `rg_plan` trailing the decoder by one — a later prune/rebuild would then + /// re-include an already-delivered row group (#24352). + fn sync_rg_plan_to_decoder_frontier(&mut self) -> Result<()> { + match self + .decoder + .as_ref() + .expect("decoder present") + .peek_next_row_group() + .map_err(DataFusionError::from)? + { + Some(actual) => Self::advance_rg_plan_to(&mut self.rg_plan, actual)?, + // Decoder has nothing left to emit — drain our plan so the stream + // finishes cleanly. + None => self.rg_plan.clear(), + } + Ok(()) + } + + /// Pop entries off `rg_plan` until its front is `target`. + /// + /// `target` is the RG the decoder will emit next and must still be in the + /// plan. A missing `target` means the decoder's frontier and `rg_plan` have + /// diverged; we surface that as an internal error rather than silently + /// draining the plan, which would truncate the scan. Kept free-standing on + /// `rg_plan` (rather than `&mut self`) so the pop/guard logic is + /// unit-testable without constructing a full stream state. + fn advance_rg_plan_to( + rg_plan: &mut VecDeque, + target: usize, + ) -> Result<()> { + while let Some(front) = rg_plan.front() { + if front.rg_index == target { + return Ok(()); + } + rg_plan.pop_front(); + } + internal_err!( + "push decoder frontier RG {target} is not in rg_plan; \ + decoder and plan have diverged" + ) + } + /// Copies metrics from ArrowReaderMetrics (the metrics collected by the /// arrow-rs parquet reader) to the parquet file metrics for DataFusion fn copy_arrow_reader_metrics(&self) { @@ -216,24 +485,213 @@ impl PushDecoderStreamState { } fn project_batch(&self, batch: &RecordBatch) -> Result { - let mut batch = self.projector.project_batch(batch)?; - if self.replace_schema { - // Ensure the output batch has the expected schema. - // This handles things like schema level and field level metadata, which may not be present - // in the physical file schema. - // It is also possible for nullability to differ; some writers create files with - // OPTIONAL fields even when there are no nulls in the data. - // In these cases it may make sense for the logical schema to be `NOT NULL`. - // RecordBatch::try_new_with_options checks that if the schema is NOT NULL - // the array cannot contain nulls, amongst other checks. - let (_stream_schema, arrays, num_rows) = batch.into_parts(); - let options = RecordBatchOptions::new().with_row_count(Some(num_rows)); - batch = RecordBatch::try_new_with_options( - Arc::clone(&self.output_schema), - arrays, - &options, - )?; + self.decoder_projection.map(batch) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{Int64Array, RecordBatch}; + use arrow::datatypes::{DataType, Field, Schema}; + use bytes::Bytes; + use datafusion_common::ScalarValue; + use datafusion_expr::Operator; + use datafusion_physical_expr::expressions::{ + BinaryExpr, Column, DynamicFilterPhysicalExpr, lit, + }; + use datafusion_physical_plan::metrics::{ExecutionPlanMetricsSet, MetricBuilder}; + use datafusion_pruning::MAX_IN_LIST_SIZE; + use parquet::arrow::ArrowWriter; + use parquet::file::metadata::ParquetMetaDataPushDecoder; + use parquet::file::properties::WriterProperties; + + /// Build a tiny in-memory Parquet file with three row groups whose `v` + /// column statistics are disjoint: RG0 → 0..1000, RG1 → 1000..2000, + /// RG2 → 2000..3000. Returns (metadata, schema). + fn build_three_rg_file() -> (Arc, SchemaRef) { + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let mut buf = Vec::new(); + let props = WriterProperties::builder() + .set_max_row_group_row_count(Some(1000)) + .build(); + let mut writer = + ArrowWriter::try_new(&mut buf, Arc::clone(&schema), Some(props)).unwrap(); + for rg in 0..3i64 { + let base = rg * 1000; + let vals: Vec = (base..base + 1000).collect(); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int64Array::from(vals))], + ) + .unwrap(); + writer.write(&batch).unwrap(); + writer.flush().unwrap(); } - Ok(batch) + writer.close().unwrap(); + + let file = Bytes::from(buf); + let len = file.len() as u64; + let mut md = ParquetMetaDataPushDecoder::try_new(len).unwrap(); + // One range covering the whole file. Using `expect` rather than + // `allow` per this crate's `clippy::allow-attributes` lint. + #[expect( + clippy::single_range_in_vec_init, + reason = "we want a single range covering the whole file" + )] + let ranges = vec![0..len]; + md.push_ranges(ranges, vec![file]).unwrap(); + let DecodeResult::Data(meta) = md.try_decode().unwrap() else { + panic!("decoding metadata"); + }; + assert_eq!(meta.num_row_groups(), 3, "test fixture must have 3 RGs"); + (Arc::new(meta), schema) + } + + /// Create a fresh `(creation_errors, evaluation_errors)` counter pair + /// for tests. The names mirror the two metrics + /// [`RowGroupPruner::new`] consumes — predicate construction is + /// accounted separately from per-row-group evaluation. + fn pruner_error_counters() -> (Count, Count) { + let metrics = ExecutionPlanMetricsSet::new(); + let creation = + MetricBuilder::new(&metrics).counter("num_predicate_creation_errors", 0); + let evaluation = + MetricBuilder::new(&metrics).counter("predicate_evaluation_errors", 0); + (creation, evaluation) + } + + /// `v > literal` predicate on a single-column schema. + fn gt_predicate(threshold: i64) -> Arc { + Arc::new(BinaryExpr::new( + Arc::new(Column::new("v", 0)), + Operator::Gt, + lit(ScalarValue::Int64(Some(threshold))), + )) + } + + #[test] + fn row_group_pruner_skips_only_disqualified_row_groups() { + let (meta, schema) = build_three_rg_file(); + let (creation, evaluation) = pruner_error_counters(); + let mut pruner = RowGroupPruner::new( + gt_predicate(1500), + Arc::clone(&schema), + Arc::clone(&meta), + creation, + evaluation, + MAX_IN_LIST_SIZE, + ); + + // RG0 (0..1000) is entirely below threshold → fully prunable. + assert!(pruner.should_prune(&[0]), "RG0 should be pruned"); + // RG1 (1000..2000) straddles the threshold → not safe to prune. + assert!(!pruner.should_prune(&[1]), "RG1 must NOT be pruned"); + // RG2 (2000..3000) is entirely above threshold → keep. + assert!(!pruner.should_prune(&[2]), "RG2 must NOT be pruned"); + // Run covering both RG0 and RG1 cannot be skipped — RG1 is alive. + assert!( + !pruner.should_prune(&[0, 1]), + "mixed run with a live RG must NOT be pruned" + ); + // Empty input is a no-op (defensive guard). + assert!(!pruner.should_prune(&[])); + } + + #[test] + fn row_group_pruner_tracks_dynamic_filter_updates() { + let (meta, schema) = build_three_rg_file(); + let dynamic = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("v", 0))], + gt_predicate(500), + )); + let (creation, evaluation) = pruner_error_counters(); + let mut pruner = RowGroupPruner::new( + Arc::clone(&dynamic) as Arc, + Arc::clone(&schema), + Arc::clone(&meta), + creation, + evaluation, + MAX_IN_LIST_SIZE, + ); + + // Initial threshold 500 → only the lower half of RG0 fails, so RG0 + // (0..1000) straddles the threshold and stays alive. + assert!(!pruner.should_prune(&[0])); + assert!(!pruner.should_prune(&[1])); + + // Tighten the threshold via the dynamic filter — TopK fills its + // heap and updates the threshold to 2500. + dynamic + .update(gt_predicate(2500)) + .expect("update threshold"); + + // After the update the pruner must rebuild its `PruningPredicate` + // (driven by the `DynamicFilterTracker`'s change notification) and + // re-evaluate. RG0 and RG1 are both entirely below 2500 now. + assert!( + pruner.should_prune(&[0]), + "RG0 must be pruned after threshold tightens to 2500" + ); + assert!( + pruner.should_prune(&[1]), + "RG1 must be pruned after threshold tightens to 2500" + ); + assert!( + !pruner.should_prune(&[2]), + "RG2 (2000..3000) still straddles 2500" + ); + } + + #[test] + fn row_group_pruner_falls_back_to_conservative_when_predicate_has_no_bounds() { + // A predicate the pruning analyzer can't decompose (e.g. a bare + // column reference of bool type would normally be valid, but a + // non-binary expression on a non-bool column doesn't yield bounds). + // We use `lit(true)` which produces no column references, so + // `build_pruning_predicate` will return None. + let (meta, schema) = build_three_rg_file(); + let (creation, evaluation) = pruner_error_counters(); + let mut pruner = RowGroupPruner::new( + lit(true) as Arc, + Arc::clone(&schema), + Arc::clone(&meta), + creation, + evaluation, + MAX_IN_LIST_SIZE, + ); + // No pruning predicate could be built → conservatively keep RGs. + assert!(!pruner.should_prune(&[0])); + assert!(!pruner.should_prune(&[1])); + assert!(!pruner.should_prune(&[2])); + } + + #[test] + fn advance_rg_plan_to_pops_up_to_target() { + let mut plan: VecDeque = [0usize, 1, 2, 3] + .into_iter() + .map(|rg_index| RgPlanEntry { rg_index }) + .collect(); + PushDecoderStreamState::advance_rg_plan_to(&mut plan, 2).unwrap(); + assert_eq!( + plan.iter().map(|e| e.rg_index).collect::>(), + vec![2, 3], + "must pop the entries before `target` and stop at it", + ); + } + + #[test] + fn advance_rg_plan_to_errors_when_target_absent() { + let mut plan: VecDeque = [0usize, 1, 2] + .into_iter() + .map(|rg_index| RgPlanEntry { rg_index }) + .collect(); + let err = PushDecoderStreamState::advance_rg_plan_to(&mut plan, 5) + .expect_err("a target absent from the plan must be an internal error"); + assert!( + err.to_string().contains("diverged"), + "expected a divergence internal error, got: {err}", + ); } } diff --git a/datafusion/datasource-parquet/src/reader.rs b/datafusion/datasource-parquet/src/reader.rs index 482bf8dced4f8..71b0020f32f64 100644 --- a/datafusion/datasource-parquet/src/reader.rs +++ b/datafusion/datasource-parquet/src/reader.rs @@ -21,18 +21,20 @@ use crate::ParquetFileMetrics; use crate::metadata::DFParquetMetadata; use bytes::Bytes; +use datafusion_common::HashMap; use datafusion_datasource::PartitionedFile; use datafusion_execution::cache::cache_manager::FileMetadata; use datafusion_execution::cache::cache_manager::FileMetadataCache; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use futures::FutureExt; +use futures::TryFutureExt; use futures::future::BoxFuture; -use object_store::ObjectStore; +use object_store::{ObjectStore, ObjectStoreExt}; use parquet::arrow::arrow_reader::ArrowReaderOptions; -use parquet::arrow::async_reader::{AsyncFileReader, ParquetObjectReader}; +use parquet::arrow::async_reader::AsyncFileReader; +use parquet::errors::ParquetError; use parquet::file::metadata::ParquetMetaData; use std::any::Any; -use std::collections::HashMap; use std::fmt::Debug; use std::ops::Range; use std::sync::Arc; @@ -86,62 +88,6 @@ impl DefaultParquetFileReaderFactory { } } -/// Implements [`AsyncFileReader`] for a parquet file in object storage. -/// -/// This implementation uses the [`ParquetObjectReader`] to read data from the -/// object store on demand, as required, tracking the number of bytes read. -/// -/// This implementation does not coalesce I/O operations or cache bytes. Such -/// optimizations can be done either at the object store level or by providing a -/// custom implementation of [`ParquetFileReaderFactory`]. -pub struct ParquetFileReader { - pub file_metrics: ParquetFileMetrics, - pub inner: ParquetObjectReader, - pub partitioned_file: PartitionedFile, -} - -impl AsyncFileReader for ParquetFileReader { - fn get_bytes( - &mut self, - range: Range, - ) -> BoxFuture<'_, parquet::errors::Result> { - let bytes_scanned = range.end - range.start; - self.file_metrics.bytes_scanned.add(bytes_scanned as usize); - self.inner.get_bytes(range) - } - - fn get_byte_ranges( - &mut self, - ranges: Vec>, - ) -> BoxFuture<'_, parquet::errors::Result>> - where - Self: Send, - { - let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); - self.file_metrics.bytes_scanned.add(total as usize); - self.inner.get_byte_ranges(ranges) - } - - fn get_metadata<'a>( - &'a mut self, - options: Option<&'a ArrowReaderOptions>, - ) -> BoxFuture<'a, parquet::errors::Result>> { - self.inner.get_metadata(options) - } -} - -impl Drop for ParquetFileReader { - fn drop(&mut self) { - self.file_metrics - .scan_efficiency_ratio - .add_part(self.file_metrics.bytes_scanned.value()); - // Multiple ParquetFileReaders may run, so we set_total to avoid adding the total multiple times - self.file_metrics - .scan_efficiency_ratio - .set_total(self.partitioned_file.object_meta.size as usize); - } -} - impl ParquetFileReaderFactory for DefaultParquetFileReaderFactory { fn create_reader( &self, @@ -155,40 +101,33 @@ impl ParquetFileReaderFactory for DefaultParquetFileReaderFactory { partitioned_file.object_meta.location.as_ref(), metrics, ); - let store = Arc::clone(&self.store); - let mut inner = ParquetObjectReader::new( - store, - partitioned_file.object_meta.location.clone(), - ) - .with_file_size(partitioned_file.object_meta.size); - - if let Some(hint) = metadata_size_hint { - inner = inner.with_footer_size_hint(hint) - }; - Ok(Box::new(ParquetFileReader { - inner, + let reader = ParquetFileReader::new( file_metrics, + Arc::clone(&self.store), partitioned_file, - })) + ) + .with_metadata_hint(metadata_size_hint); + Ok(Box::new(reader)) } } /// Implementation of [`ParquetFileReaderFactory`] supporting the caching of footer and page /// metadata. Reads and updates the [`FileMetadataCache`] with the [`ParquetMetaData`] data. -/// This reader always loads the entire metadata (including page index, unless the file is -/// encrypted), even if not required by the current query, to ensure it is always available for -/// those that need it. +/// +/// [`ParquetFileReader::get_metadata`] forwards the [`parquet::file::metadata::PageIndexPolicy`] from +/// [`ArrowReaderOptions`] to [`DFParquetMetadata::fetch_metadata`], so callers such as the +/// parquet opener can skip page-index I/O during the initial metadata load. #[derive(Debug)] pub struct CachedParquetFileReaderFactory { store: Arc, - metadata_cache: Arc, + metadata_cache: Arc, } impl CachedParquetFileReaderFactory { pub fn new( store: Arc, - metadata_cache: Arc, + metadata_cache: Arc, ) -> Self { Self { store, @@ -210,69 +149,103 @@ impl ParquetFileReaderFactory for CachedParquetFileReaderFactory { partitioned_file.object_meta.location.as_ref(), metrics, ); - let store = Arc::clone(&self.store); - - let mut inner = ParquetObjectReader::new( - store, - partitioned_file.object_meta.location.clone(), - ) - .with_file_size(partitioned_file.object_meta.size); - if let Some(hint) = metadata_size_hint { - inner = inner.with_footer_size_hint(hint) - }; - - Ok(Box::new(CachedParquetFileReader::new( + let reader = ParquetFileReader::new( file_metrics, Arc::clone(&self.store), - inner, partitioned_file, - Arc::clone(&self.metadata_cache), - metadata_size_hint, - ))) + ) + .with_metadata_hint(metadata_size_hint) + .with_metadata_cache(Some(Arc::clone(&self.metadata_cache))); + + Ok(Box::new(reader)) } } -/// Implements [`AsyncFileReader`] for a Parquet file in object storage. Reads the file metadata -/// from the [`FileMetadataCache`], if available, otherwise reads it directly from the file and then -/// updates the cache. -pub struct CachedParquetFileReader { - pub file_metrics: ParquetFileMetrics, +/// Implements [`AsyncFileReader`] for a parquet file in object storage. +/// +/// This implementation reads data directly from the underlying [`ObjectStore`] +/// on demand, as required, tracking the number of bytes read. +/// +/// When configured via [`Self::with_metadata_cache`], [`Self::get_metadata`] +/// reads footer and page metadata from the cache when available and populates +/// the cache otherwise. Without a cache, metadata is fetched fresh on every call. +/// +/// # Notes +/// +/// This implementation does not coalesce I/O operations or cache bytes. Such +/// optimizations can be done either at the object store level or by providing +/// a custom implementation of [`ParquetFileReaderFactory`]. +pub struct ParquetFileReader { + file_metrics: ParquetFileMetrics, store: Arc, - pub inner: ParquetObjectReader, partitioned_file: PartitionedFile, - metadata_cache: Arc, + metadata_cache: Option>, metadata_size_hint: Option, } -impl CachedParquetFileReader { - pub fn new( +impl ParquetFileReader { + /// Create a new `ParquetFileReader`. + /// + /// By default the reader has no [`FileMetadataCache`] and no metadata + /// size hint, so metadata is fetched fresh on every call (as + /// [`DefaultParquetFileReaderFactory`] does). Use + /// [`Self::with_metadata_cache`] to read and populate a cache (as + /// [`CachedParquetFileReaderFactory`] does), and + /// [`Self::with_metadata_hint`] to set the size hint. + pub(crate) fn new( file_metrics: ParquetFileMetrics, store: Arc, - inner: ParquetObjectReader, partitioned_file: PartitionedFile, - metadata_cache: Arc, - metadata_size_hint: Option, ) -> Self { Self { file_metrics, store, - inner, partitioned_file, - metadata_cache, - metadata_size_hint, + metadata_cache: None, + metadata_size_hint: None, } } + + /// Returns the metrics tracked while reading this file. + pub fn file_metrics(&self) -> &ParquetFileMetrics { + &self.file_metrics + } + + /// Returns the file this reader is reading. + pub fn partitioned_file(&self) -> &PartitionedFile { + &self.partitioned_file + } + + /// Set the [`FileMetadataCache`] for this reader + pub fn with_metadata_cache( + mut self, + metadata_cache: Option>, + ) -> Self { + self.metadata_cache = metadata_cache; + self + } + + /// Set the metadata size hint for this reader. + /// + /// See [`DFParquetMetadata::with_metadata_size_hint`] for more details. + pub fn with_metadata_hint(mut self, metadata_size_hint: Option) -> Self { + self.metadata_size_hint = metadata_size_hint; + self + } } -impl AsyncFileReader for CachedParquetFileReader { +impl AsyncFileReader for ParquetFileReader { fn get_bytes( &mut self, range: Range, ) -> BoxFuture<'_, parquet::errors::Result> { let bytes_scanned = range.end - range.start; self.file_metrics.bytes_scanned.add(bytes_scanned as usize); - self.inner.get_bytes(range) + self.store + .get_range(&self.partitioned_file.object_meta.location, range) + .map_err(|e| ParquetError::External(Box::new(e))) + .boxed() } fn get_byte_ranges( @@ -284,16 +257,21 @@ impl AsyncFileReader for CachedParquetFileReader { { let total: u64 = ranges.iter().map(|r| r.end - r.start).sum(); self.file_metrics.bytes_scanned.add(total as usize); - self.inner.get_byte_ranges(ranges) + async move { + self.store + .get_ranges(&self.partitioned_file.object_meta.location, &ranges) + .await + .map_err(|e| ParquetError::External(Box::new(e))) + } + .boxed() } fn get_metadata<'a>( &'a mut self, - #[cfg_attr(not(feature = "parquet_encryption"), expect(unused_variables))] options: Option<&'a ArrowReaderOptions>, ) -> BoxFuture<'a, parquet::errors::Result>> { let object_meta = self.partitioned_file.object_meta.clone(); - let metadata_cache = Arc::clone(&self.metadata_cache); + let metadata_cache = self.metadata_cache.clone(); async move { #[cfg(feature = "parquet_encryption")] @@ -304,14 +282,17 @@ impl AsyncFileReader for CachedParquetFileReader { #[cfg(not(feature = "parquet_encryption"))] let file_decryption_properties = None; + let page_index_policy = options.map(|o| o.column_index_policy()); + DFParquetMetadata::new(&self.store, &object_meta) .with_decryption_properties(file_decryption_properties) - .with_file_metadata_cache(Some(Arc::clone(&metadata_cache))) + .with_file_metadata_cache(metadata_cache) .with_metadata_size_hint(self.metadata_size_hint) + .with_page_index_policy(page_index_policy) .fetch_metadata() .await .map_err(|e| { - parquet::errors::ParquetError::General(format!( + ParquetError::General(format!( "Failed to fetch metadata for file {}: {e}", object_meta.location, )) @@ -321,7 +302,7 @@ impl AsyncFileReader for CachedParquetFileReader { } } -impl Drop for CachedParquetFileReader { +impl Drop for ParquetFileReader { fn drop(&mut self) { self.file_metrics .scan_efficiency_ratio diff --git a/datafusion/datasource-parquet/src/row_filter.rs b/datafusion/datasource-parquet/src/row_filter.rs index f19dbd6c6fa63..c1a47c896c170 100644 --- a/datafusion/datasource-parquet/src/row_filter.rs +++ b/datafusion/datasource-parquet/src/row_filter.rs @@ -65,31 +65,29 @@ //! - `WHERE s['value'] > 5` — pushed down (accesses a primitive leaf) //! - `WHERE s IS NOT NULL` — not pushed down (references the whole struct) -use std::collections::BTreeSet; use std::sync::Arc; use arrow::array::BooleanArray; -use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::datatypes::{Schema, SchemaRef}; use arrow::error::{ArrowError, Result as ArrowResult}; use arrow::record_batch::RecordBatch; -use datafusion_functions::core::getfield::GetFieldFunc; use parquet::arrow::ProjectionMask; use parquet::arrow::arrow_reader::{ArrowPredicate, RowFilter}; use parquet::file::metadata::ParquetMetaData; -use parquet::schema::types::SchemaDescriptor; use datafusion_common::Result; use datafusion_common::cast::as_boolean_array; -use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}; -use datafusion_physical_expr::ScalarFunctionExpr; -use datafusion_physical_expr::expressions::{Column, Literal}; -use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns}; +use datafusion_common::tree_node::TreeNode; +use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr::{PhysicalExpr, split_conjunction}; use datafusion_physical_plan::metrics; use super::ParquetFileMetrics; use super::supported_predicates::supports_list_predicates; +use crate::projection_read_plan::{ + ParquetReadPlan, PushdownChecker, PushdownColumns, assemble_read_plan, +}; /// A "compiled" predicate passed to `ParquetRecordBatchStream` to perform /// row-level filtering during parquet decoding. @@ -188,22 +186,6 @@ pub(crate) struct FilterCandidate { read_plan: ParquetReadPlan, } -/// The result of resolving which Parquet leaf columns and Arrow schema fields -/// are needed to evaluate an expression against a Parquet file -/// -/// This is the shared output of the column resolution pipeline used by both -/// the row filter to build `ArrowPredicate`s and the opener to build `ProjectionMask`s -#[derive(Debug, Clone)] -pub(crate) struct ParquetReadPlan { - /// Projection mask built from leaf column indices in the Parquet schema. - /// Using a `ProjectionMask` directly (rather than raw indices) prevents - /// bugs from accidentally mixing up root vs leaf indices. - pub projection_mask: ProjectionMask, - /// The projected Arrow schema containing only the columns/fields required - /// Struct types are pruned to include only the accessed sub-fields - pub projected_schema: SchemaRef, -} - /// Helper to build a `FilterCandidate`. /// /// This will do several things: @@ -245,278 +227,6 @@ impl FilterCandidateBuilder { } } -/// Traverses a `PhysicalExpr` tree to determine if any column references would -/// prevent the expression from being pushed down to the parquet decoder. -/// -/// An expression cannot be pushed down if it references: -/// - Unsupported nested columns (whole struct references or list fields that are -/// not covered by the supported predicate set) -/// - Columns that don't exist in the file schema -/// -/// Struct field access via `get_field` is supported when the resolved leaf type -/// is primitive (e.g. `get_field(struct_col, 'field') > 5`). -struct PushdownChecker<'schema> { - /// Does the expression require any non-primitive columns (like structs)? - non_primitive_columns: bool, - /// Does the expression reference any columns not present in the file schema? - projected_columns: bool, - /// Indices into the file schema of columns required to evaluate the expression. - /// Does not include struct columns accessed via `get_field`. - required_columns: Vec, - /// Struct field accesses via `get_field`. - struct_field_accesses: Vec, - /// Whether nested list columns are supported by the predicate semantics. - allow_list_columns: bool, - /// The Arrow schema of the parquet file. - file_schema: &'schema Schema, -} - -impl<'schema> PushdownChecker<'schema> { - fn new(file_schema: &'schema Schema, allow_list_columns: bool) -> Self { - Self { - non_primitive_columns: false, - projected_columns: false, - required_columns: Vec::new(), - struct_field_accesses: Vec::new(), - allow_list_columns, - file_schema, - } - } - - /// Checks whether a struct's root column exists in the file schema and, if so, - /// records its index so the entire struct is decoded for filter evaluation. - /// - /// This is called when we see a `get_field` expression that resolves to a - /// primitive leaf type. We only need the *root* column index because the - /// Parquet reader decodes all leaves of a struct together. - /// - /// # Example - /// - /// Given file schema `{a: Int32, s: Struct(foo: Utf8, bar: Int64)}` and the - /// expression `get_field(s, 'foo') = 'hello'`: - /// - /// - `column_name` = `"s"` (the root struct column) - /// - `file_schema.index_of("s")` returns `1` - /// - We push `1` into `required_columns` - /// - Return `None` (no issue — traversal continues in the caller) - /// - /// If `"s"` is not in the file schema (e.g. a projected-away column), we set - /// `projected_columns = true` and return `Jump` to skip the subtree. - fn check_struct_field_column( - &mut self, - column_name: &str, - field_path: Vec, - ) -> Option { - let Ok(idx) = self.file_schema.index_of(column_name) else { - self.projected_columns = true; - return Some(TreeNodeRecursion::Jump); - }; - - self.struct_field_accesses.push(StructFieldAccess { - root_index: idx, - field_path, - }); - - None - } - - fn check_single_column(&mut self, column_name: &str) -> Option { - let idx = match self.file_schema.index_of(column_name) { - Ok(idx) => idx, - Err(_) => { - // Column does not exist in the file schema, so we can't push this down. - self.projected_columns = true; - return Some(TreeNodeRecursion::Jump); - } - }; - - // Duplicates are handled by dedup() in into_sorted_columns() - self.required_columns.push(idx); - let data_type = self.file_schema.field(idx).data_type(); - - if DataType::is_nested(data_type) { - self.handle_nested_type(data_type) - } else { - None - } - } - - /// Determines whether a nested data type can be pushed down to Parquet decoding. - /// - /// Returns `Some(TreeNodeRecursion::Jump)` if the nested type prevents pushdown, - /// `None` if the type is supported and pushdown can continue. - fn handle_nested_type(&mut self, data_type: &DataType) -> Option { - if self.is_nested_type_supported(data_type) { - None - } else { - // Block pushdown for unsupported nested types: - // - Structs (regardless of predicate support) - // - Lists without supported predicates - self.non_primitive_columns = true; - Some(TreeNodeRecursion::Jump) - } - } - - /// Checks if a nested data type is supported for list column pushdown. - /// - /// List columns are only supported if: - /// 1. The data type is a list variant (List, LargeList, or FixedSizeList) - /// 2. The expression contains supported list predicates (e.g., array_has_all) - fn is_nested_type_supported(&self, data_type: &DataType) -> bool { - let is_list = matches!( - data_type, - DataType::List(_) | DataType::LargeList(_) | DataType::FixedSizeList(_, _) - ); - self.allow_list_columns && is_list - } - - #[inline] - fn prevents_pushdown(&self) -> bool { - self.non_primitive_columns || self.projected_columns - } - - /// Consumes the checker and returns sorted, deduplicated column indices - /// wrapped in a `PushdownColumns` struct. - /// - /// This method sorts the column indices and removes duplicates. The sort - /// is required because downstream code relies on column indices being in - /// ascending order for correct schema projection. - fn into_sorted_columns(mut self) -> PushdownColumns { - self.required_columns.sort_unstable(); - self.required_columns.dedup(); - PushdownColumns { - required_columns: self.required_columns, - struct_field_accesses: self.struct_field_accesses, - } - } -} - -impl TreeNodeVisitor<'_> for PushdownChecker<'_> { - type Node = Arc; - - fn f_down(&mut self, node: &Self::Node) -> Result { - // Handle struct field access like `s['foo']['bar'] > 10`. - // - // DataFusion represents nested field access as `get_field(Column("s"), "foo")` - // (or chained: `get_field(get_field(Column("s"), "foo"), "bar")`). - // - // We intercept the outermost `get_field` on the way *down* the tree so - // the visitor never reaches the raw `Column("s")` node. Without this, - // `check_single_column` would see that `s` is a Struct and reject it. - // - // The strategy: - // 1. Match `get_field` whose first arg is a `Column` (the struct root). - // 2. Check that the *resolved* return type is primitive — meaning we've - // drilled all the way to a leaf (e.g. `s['foo']` → Utf8). - // 3. Record the root column index via `check_struct_field_column` and - // return `Jump` to skip visiting the children (the Column and the - // literal field-name args), since we've already handled them. - // - // If the return type is still nested (e.g. `s['nested_struct']` → Struct), - // we fall through and let normal traversal continue, which will - // eventually reject the expression when it hits the struct Column. - if let Some(func) = - ScalarFunctionExpr::try_downcast_func::(node.as_ref()) - { - let args = func.args(); - - if let Some(column) = args.first().and_then(|a| a.downcast_ref::()) { - // for Map columns, get_field performs a runtime key lookup rather than a - // schema-level field access so the entire Map column must be read, - // we skip the struct field optimization and defer to normal Column traversal - let is_map_column = self - .file_schema - .index_of(column.name()) - .ok() - .map(|idx| { - matches!( - self.file_schema.field(idx).data_type(), - DataType::Map(_, _) - ) - }) - .unwrap_or(false); - - let return_type = func.return_type(); - - if !is_map_column - && (!DataType::is_nested(return_type) - || self.is_nested_type_supported(return_type)) - { - // try to resolve all field name arguments to strinrg literals - // if any argument is not a string literal, we can not determine the exact - // leaf path so we fall back to reading the entire struct root column - let field_path = args[1..] - .iter() - .map(|arg| { - arg.downcast_ref::().and_then(|lit| { - lit.value().try_as_str().flatten().map(|s| s.to_string()) - }) - }) - .collect(); - - match field_path { - Some(path) => { - if let Some(recursion) = - self.check_struct_field_column(column.name(), path) - { - return Ok(recursion); - } - } - None => { - // Could not resolve field path — fall back to - // reading the entire struct root column. - if let Some(recursion) = - self.check_single_column(column.name()) - { - return Ok(recursion); - } - } - } - - return Ok(TreeNodeRecursion::Jump); - } - } - } - - if let Some(column) = node.downcast_ref::() - && let Some(recursion) = self.check_single_column(column.name()) - { - return Ok(recursion); - } - - Ok(TreeNodeRecursion::Continue) - } -} - -/// Describes the nested column behavior for filter pushdown. -/// -/// This enum makes explicit the different states a predicate can be in -/// with respect to nested column handling during Parquet decoding. -/// Result of checking which columns are required for filter pushdown. -#[derive(Debug)] -struct PushdownColumns { - /// Sorted, unique column indices into the file schema required to evaluate - /// the filter expression. Must be in ascending order for correct schema - /// projection matching. Does not include struct columns accessed via `get_field`. - required_columns: Vec, - /// Struct field accesses via `get_field`. Each entry records the root struct - /// column index and the field path being accessed. - struct_field_accesses: Vec, -} - -/// Records a struct field access via `get_field(struct_col, 'field1', 'field2', ...)`. -/// -/// This allows the row filter to project only the specific Parquet leaf columns -/// needed by the filter, rather than all leaves of the struct. -#[derive(Debug, Clone)] -struct StructFieldAccess { - /// Arrow root column index of the struct in the file schema. - root_index: usize, - /// Field names forming the path into the struct. - /// e.g., `["value"]` for `s['value']`, `["outer", "inner"]` for `s['outer']['inner']`. - field_path: Vec, -} - /// Checks if a given expression can be pushed down to the parquet decoder. /// /// Returns `Some(PushdownColumns)` if the expression can be pushed down, @@ -558,344 +268,16 @@ pub(crate) fn build_parquet_read_plan( return Ok(None); }; - let root_indices = &required_columns.required_columns; - - let mut leaf_indices = - leaf_indices_for_roots(root_indices.iter().copied(), schema_descr); - - let struct_leaf_indices = resolve_struct_field_leaves( + let (read_plan, leaf_indices) = assemble_read_plan( + &required_columns.required_columns, &required_columns.struct_field_accesses, file_schema, schema_descr, ); - leaf_indices.extend_from_slice(&struct_leaf_indices); - leaf_indices.sort_unstable(); - leaf_indices.dedup(); let required_bytes = size_of_columns(&leaf_indices, metadata)?; - let projection_mask = - ProjectionMask::leaves(schema_descr, leaf_indices.iter().copied()); - - let projected_schema = build_filter_schema( - file_schema, - root_indices, - &required_columns.struct_field_accesses, - ); - - Ok(Some(( - ParquetReadPlan { - projection_mask, - projected_schema, - }, - required_bytes, - ))) -} - -/// Builds a unified [`ParquetReadPlan`] for a set of projection expressions -/// -/// Unlike [`build_parquet_read_plan`] (which is used for filter pushdown and -/// returns `None` when an expression references unsupported nested types or -/// missing columns), this function always succeeds. It collects every column -/// that *can* be resolved in the file and produces a leaf-level projection -/// mask. Columns missing from the file are silently skipped since the projection -/// layer handles those by inserting nulls. -pub(crate) fn build_projection_read_plan( - exprs: impl IntoIterator>, - file_schema: &Schema, - schema_descr: &SchemaDescriptor, -) -> ParquetReadPlan { - // fast path: if every expression is a plain Column reference, skip all - // struct analysis and use root-level projection directly - let exprs = exprs.into_iter().collect::>(); - let all_plain_columns = exprs.iter().all(|e| e.downcast_ref::().is_some()); - - if all_plain_columns { - let mut root_indices: Vec = exprs - .iter() - .map(|e| e.downcast_ref::().unwrap().index()) - .collect(); - root_indices.sort_unstable(); - root_indices.dedup(); - - let projection_mask = - ProjectionMask::roots(schema_descr, root_indices.iter().copied()); - let projected_schema = Arc::new( - file_schema - .project(&root_indices) - .expect("valid column indices"), - ); - - return ParquetReadPlan { - projection_mask, - projected_schema, - }; - } - - // secondary fast path: if the schema has no struct columns, we can skip - // PushdownChecker traversal and use root-level projection - let has_struct_columns = file_schema - .fields() - .iter() - .any(|f| matches!(f.data_type(), DataType::Struct(_))); - - if !has_struct_columns { - let mut root_indices = exprs - .into_iter() - .flat_map(|e| collect_columns(&e).into_iter().map(|col| col.index())) - .collect::>(); - - root_indices.sort_unstable(); - root_indices.dedup(); - - let projection_mask = - ProjectionMask::roots(schema_descr, root_indices.iter().copied()); - - let projected_schema = Arc::new( - file_schema - .project(&root_indices) - .expect("valid column indices"), - ); - - return ParquetReadPlan { - projection_mask, - projected_schema, - }; - } - - let mut all_root_indices = Vec::new(); - let mut all_struct_accesses = Vec::new(); - - for expr in exprs { - let mut checker = PushdownChecker::new(file_schema, true); - let _ = expr.visit(&mut checker); - let columns = checker.into_sorted_columns(); - - all_root_indices.extend_from_slice(&columns.required_columns); - all_struct_accesses.extend(columns.struct_field_accesses); - } - - all_root_indices.sort_unstable(); - all_root_indices.dedup(); - - // when no struct field accesses were found, fall back to root-level projection - // to match the performance of the simple path - if all_struct_accesses.is_empty() { - let projection_mask = - ProjectionMask::roots(schema_descr, all_root_indices.iter().copied()); - let projected_schema = Arc::new( - file_schema - .project(&all_root_indices) - .expect("valid column indices"), - ); - - return ParquetReadPlan { - projection_mask, - projected_schema, - }; - } - - let leaf_indices = { - let mut out = - leaf_indices_for_roots(all_root_indices.iter().copied(), schema_descr); - let struct_leaf_indices = - resolve_struct_field_leaves(&all_struct_accesses, file_schema, schema_descr); - - out.extend_from_slice(&struct_leaf_indices); - out.sort_unstable(); - out.dedup(); - - out - }; - - let projection_mask = - ProjectionMask::leaves(schema_descr, leaf_indices.iter().copied()); - - let projected_schema = - build_filter_schema(file_schema, &all_root_indices, &all_struct_accesses); - - ParquetReadPlan { - projection_mask, - projected_schema, - } -} - -fn leaf_indices_for_roots( - root_indices: I, - schema_descr: &SchemaDescriptor, -) -> Vec -where - I: IntoIterator, -{ - // Always map root (Arrow) indices to Parquet leaf indices via the schema - // descriptor. Arrow root indices only equal Parquet leaf indices when the - // schema has no group columns (Struct, Map, etc.); when group columns - // exist, their children become separate leaves and shift all subsequent - // leaf indices. - // Struct columns are unsupported. - let root_set: BTreeSet<_> = root_indices.into_iter().collect(); - - (0..schema_descr.num_columns()) - .filter(|leaf_idx| { - root_set.contains(&schema_descr.get_column_root_idx(*leaf_idx)) - }) - .collect() -} - -/// Resolves struct field access to specific Parquet leaf column indices -/// -/// For every `StructFieldAccess`, finds the leaf columns in the Parquet schema -/// whose path matches the struct root name + field path. This avoids reading all -/// leaves of a struct when only specific fields are needed -fn resolve_struct_field_leaves( - accesses: &[StructFieldAccess], - file_schema: &Schema, - schema_descr: &SchemaDescriptor, -) -> Vec { - let mut leaf_indices = Vec::new(); - - for access in accesses { - let root_name = file_schema.field(access.root_index).name(); - let prefix = std::iter::once(root_name.as_str()) - .chain(access.field_path.iter().map(|p| p.as_str())) - .collect::>(); - - for leaf_idx in 0..schema_descr.num_columns() { - let col = schema_descr.column(leaf_idx); - let col_path = col.path().parts(); - - // A leaf matches if its path starts with our prefix. - // e.g., prefix=["s", "value"] matches leaf path ["s", "value"] - // prefix=["s", "outer"] matches ["s", "outer", "inner"] - - // a leaf matches if its path starts with our prefix - // for example: prefix=["s", "value"] matches leaf path ["s", "value"] - // prefix=["s", "outer"] matches ["s", "outer", "inner"] - let leaf_matches_path = col_path.len() >= prefix.len() - && col_path.iter().zip(prefix.iter()).all(|(a, b)| a == b); - - if leaf_matches_path { - leaf_indices.push(leaf_idx); - } - } - } - - leaf_indices -} - -/// Builds a filter schema that includes only the fields actually accessed by the -/// filter expression. -/// -/// For regular (non-struct) columns, the full field type is used. -/// For struct columns accessed via `get_field`, a pruned struct type is created -/// containing only the fields along the access path. Note: it must match the schema -/// that the Parquet reader produces when projecting specific struct leaves -fn build_filter_schema( - file_schema: &Schema, - regular_indices: &[usize], - struct_field_accesses: &[StructFieldAccess], -) -> SchemaRef { - let regular_set: BTreeSet = regular_indices.iter().copied().collect(); - - let all_indices = regular_indices - .iter() - .copied() - .chain( - struct_field_accesses - .iter() - .map(|&StructFieldAccess { root_index, .. }| root_index), - ) - .collect::>(); - - let fields = all_indices - .iter() - .map(|&idx| { - let field = file_schema.field(idx); - - // if this column appears as a regular (whole-column) reference, - // keep the full type - // - // Pruning is only valid when the column is accessed exclusively - // through struct field accesses - if regular_set.contains(&idx) { - return Arc::new(field.clone()); - } - - // collect all field paths that access this root struct column - let field_paths = struct_field_accesses - .iter() - .filter_map( - |&StructFieldAccess { - root_index, - ref field_path, - }| { - (root_index == idx).then_some(field_path.as_slice()) - }, - ) - .collect::>(); - - if field_paths.is_empty() { - return Arc::new(field.clone()); - } - - let pruned_data_type = prune_struct_type(field.data_type(), &field_paths); - Arc::new(Field::new( - field.name(), - pruned_data_type, - field.is_nullable(), - )) - }) - .collect::>(); - - Arc::new(Schema::new_with_metadata( - fields, - file_schema.metadata().clone(), - )) -} - -fn prune_struct_type(dt: &DataType, paths: &[&[String]]) -> DataType { - let DataType::Struct(fields) = dt else { - return dt.clone(); - }; - - let needed = paths - .iter() - .filter_map(|p| p.first().map(|s| s.as_str())) - .collect::>(); - - let pruned_fields = fields - .iter() - .filter_map(|f| { - if !needed.contains(f.name().as_str()) { - return None; - } - - let sub_paths = paths - .iter() - .filter_map(|path| { - if path.first().map(|s| s.as_str()) == Some(f.name()) { - Some(&path[1..]) - } else { - None - } - }) - .filter(|sub| !sub.is_empty()) - .collect::>(); - - let out = if sub_paths.is_empty() { - // Leaf of access path — keep the field as-is. - Arc::clone(f) - } else { - // Recurse into nested struct. - let pruned = prune_struct_type(f.data_type(), &sub_paths); - Arc::new(Field::new(f.name(), pruned, f.is_nullable())) - }; - - Some(out) - }) - .collect::>(); - - DataType::Struct(pruned_fields.into()) + Ok(Some((read_plan, required_bytes))) } /// Checks if a predicate expression can be pushed down to the parquet decoder. @@ -1082,12 +464,11 @@ pub fn build_row_filter( .map(|filters| Some(RowFilter::new(filters))) } -/// Builds row filters for decoder runs. +/// Builds row filters for a parquet decoder. /// -/// A [`RowFilter`] must be owned by a decoder, so scans split across multiple -/// decoder runs need a fresh filter for each run that evaluates row predicates. -/// The first filter is built eagerly during construction so callers can cheaply -/// query [`has_row_filter`](Self::has_row_filter) before splitting the scan. +/// A [`RowFilter`] is owned by a decoder. The first filter is built eagerly +/// during construction so the caller can attach it to the decoder via +/// [`next_filter`](Self::next_filter) without a redundant build call. pub(crate) struct RowFilterGenerator<'a> { predicate: Option<&'a Arc>, physical_file_schema: &'a SchemaRef, @@ -1117,10 +498,6 @@ impl<'a> RowFilterGenerator<'a> { generator } - pub(crate) fn has_row_filter(&self) -> bool { - self.first_row_filter.is_some() - } - pub(crate) fn next_filter(&mut self) -> Option { self.first_row_filter.take().or_else(|| self.build()) } @@ -1149,7 +526,7 @@ impl<'a> RowFilterGenerator<'a> { #[cfg(test)] mod test { use super::*; - use arrow::datatypes::Fields; + use arrow::datatypes::{DataType, Fields}; use datafusion_common::ScalarValue; use arrow::array::{ @@ -1164,6 +541,7 @@ mod test { use datafusion_functions_nested::expr_fn::{ array_has, array_has_all, array_has_any, make_array, }; + use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::planner::logical2physical; use datafusion_physical_expr_adapter::{ DefaultPhysicalExprAdapterFactory, PhysicalExprAdapterFactory, @@ -1176,8 +554,6 @@ mod test { use parquet::file::reader::{FileReader, SerializedFileReader}; use tempfile::NamedTempFile; - use datafusion_physical_expr::expressions::Column as PhysicalColumn; - // List predicates used by the decoder should be accepted for pushdown #[test] fn test_filter_candidate_builder_supports_list_types() { @@ -2034,34 +1410,47 @@ mod test { assert_eq!(file_metrics.pushdown_rows_matched.value(), 2); } + /// Sanity check that the given expression could be evaluated against the given schema without any errors. + /// This will fail if the expression references columns that are not in the schema or if the types of the columns are incompatible, etc. + fn check_expression_can_evaluate_against_schema( + expr: &Arc, + table_schema: &Arc, + ) -> bool { + let batch = RecordBatch::new_empty(Arc::clone(table_schema)); + expr.evaluate(&batch).is_ok() + } + + /// Multiple sibling fields under one struct root: `s['value'] AND s['label']`. + /// The projection mask should include exactly those two leaves (not the third + /// sibling), and the projected schema should be pruned to those siblings. #[test] - fn projection_read_plan_preserves_full_struct() { - // Schema: id (Int32), s (Struct{value: Int32, label: Utf8}) - // Parquet leaves: id=0, s.value=1, s.label=2 + fn get_field_multiple_fields_under_same_root_uses_only_those_leaves() { + // Schema: s (Struct{value: Int32, label: Utf8, extra: Int32}) + // Parquet leaves: s.value=0, s.label=1, s.extra=2 let struct_fields: Fields = vec![ Arc::new(Field::new("value", DataType::Int32, false)), Arc::new(Field::new("label", DataType::Utf8, false)), + Arc::new(Field::new("extra", DataType::Int32, false)), ] .into(); - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Int32, false), - Field::new("s", DataType::Struct(struct_fields.clone()), false), - ])); + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(struct_fields.clone()), + false, + )])); let batch = RecordBatch::try_new( Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), - Arc::new(StructArray::new( - struct_fields, - vec![ - Arc::new(Int32Array::from(vec![10, 20, 30])) as _, - Arc::new(StringArray::from(vec!["a", "b", "c"])) as _, - ], - None, - )), - ], + vec![Arc::new(StructArray::new( + struct_fields.clone(), + vec![ + Arc::new(Int32Array::from(vec![10, 20, 30])) as _, + Arc::new(StringArray::from(vec!["a", "b", "c"])) as _, + Arc::new(Int32Array::from(vec![100, 200, 300])) as _, + ], + None, + ))], ) .unwrap(); @@ -2077,50 +1466,406 @@ mod test { .expect("reader builder"); let metadata = builder.metadata().clone(); let file_schema = builder.schema().clone(); - let schema_descr = metadata.file_metadata().schema_descr(); - - // Simulate SELECT * output projection: Column("id") and Column("s") - // Plus a get_field(s, 'value') expression from the pushed-down filter - let exprs: Vec> = vec![ - Arc::new(PhysicalColumn::new("id", 0)), - Arc::new(PhysicalColumn::new("s", 1)), - logical2physical( - &get_field().call(vec![ - col("s"), - Expr::Literal(ScalarValue::Utf8(Some("value".to_string())), None), - ]), - &file_schema, - ), - ]; - let read_plan = build_projection_read_plan(exprs, &file_schema, schema_descr); + // s['value'] > 5 AND s['label'] = 'b' + let value_expr = get_field() + .call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("value".to_string())), None), + ]) + .gt(Expr::Literal(ScalarValue::Int32(Some(5)), None)); + let label_expr = get_field() + .call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("label".to_string())), None), + ]) + .eq(Expr::Literal( + ScalarValue::Utf8(Some("b".to_string())), + None, + )); + let expr = logical2physical(&value_expr.and(label_expr), &file_schema); + + let candidate = FilterCandidateBuilder::new(expr, Arc::clone(&file_schema)) + .build(&metadata) + .expect("building candidate") + .expect("conjunction of two get_field predicates should be pushable"); - // The projected schema must have the FULL struct type because Column("s") - // is in the projection. It should NOT be narrowed to Struct{value: Int32}. - let s_field = read_plan.projected_schema.field_with_name("s").unwrap(); + // Only s.value (leaf 0) and s.label (leaf 1) should be projected; s.extra (leaf 2) skipped. + let expected_mask = + ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [0, 1]); + assert_eq!( + candidate.read_plan.projection_mask, expected_mask, + "projection_mask should include only the two accessed sibling leaves" + ); + + let s_field = candidate + .read_plan + .projected_schema + .field_with_name("s") + .unwrap(); + let expected_pruned: Fields = vec![ + Arc::new(Field::new("value", DataType::Int32, false)), + Arc::new(Field::new("label", DataType::Utf8, false)), + ] + .into(); assert_eq!( s_field.data_type(), - &DataType::Struct( - vec![ - Arc::new(Field::new("value", DataType::Int32, false)), - Arc::new(Field::new("label", DataType::Utf8, false)), - ] - .into() - ), + &DataType::Struct(expected_pruned), + "projected struct schema should drop the un-accessed `extra` sibling" + ); + } + + /// Two predicates share a nested prefix: `s['outer']['a'] AND s['outer']['b']`. + /// The projection mask should include exactly those two leaves and exclude + /// the cousin under `s['other']` plus `s['outer']['c']`. The projected + /// schema must mirror that shape. + #[test] + fn get_field_nested_shared_prefix_uses_only_prefix_leaves() { + // Schema: s (Struct{outer: Struct{a, b, c}, other: Struct{x}}) + // Parquet leaves: s.outer.a=0, s.outer.b=1, s.outer.c=2, s.other.x=3 + let outer_fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Field::new("b", DataType::Int32, false)), + Arc::new(Field::new("c", DataType::Int32, false)), + ] + .into(); + let other_fields: Fields = + vec![Arc::new(Field::new("x", DataType::Int32, false))].into(); + let s_fields: Fields = vec![ + Arc::new(Field::new( + "outer", + DataType::Struct(outer_fields.clone()), + false, + )), + Arc::new(Field::new( + "other", + DataType::Struct(other_fields.clone()), + false, + )), + ] + .into(); + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(s_fields.clone()), + false, + )])); + + let outer_arr = StructArray::new( + outer_fields, + vec![ + Arc::new(Int32Array::from(vec![1, 2])) as _, + Arc::new(Int32Array::from(vec![10, 20])) as _, + Arc::new(Int32Array::from(vec![100, 200])) as _, + ], + None, + ); + let other_arr = StructArray::new( + other_fields, + vec![Arc::new(Int32Array::from(vec![7, 8])) as _], + None, + ); + let s_arr = StructArray::new( + s_fields, + vec![Arc::new(outer_arr) as _, Arc::new(other_arr) as _], + None, + ); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s_arr)]).unwrap(); + + let file = NamedTempFile::new().expect("temp file"); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) + .expect("writer"); + writer.write(&batch).expect("write batch"); + writer.close().expect("close writer"); + + let reader_file = file.reopen().expect("reopen file"); + let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file) + .expect("reader builder"); + let metadata = builder.metadata().clone(); + let file_schema = builder.schema().clone(); + + // s['outer']['a'] > 0 AND s['outer']['b'] > 0 + let a_expr = get_field() + .call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("outer".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("a".to_string())), None), + ]) + .gt(Expr::Literal(ScalarValue::Int32(Some(0)), None)); + let b_expr = get_field() + .call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("outer".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("b".to_string())), None), + ]) + .gt(Expr::Literal(ScalarValue::Int32(Some(0)), None)); + let expr = logical2physical(&a_expr.and(b_expr), &file_schema); + + let candidate = FilterCandidateBuilder::new(expr, Arc::clone(&file_schema)) + .build(&metadata) + .expect("building candidate") + .expect("shared-prefix nested predicates should be pushable"); + + // Only s.outer.a (0) and s.outer.b (1) — not s.outer.c (2), not s.other.x (3). + let expected_mask = + ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [0, 1]); + assert_eq!( + candidate.read_plan.projection_mask, expected_mask, + "projection_mask should drop cousin and un-accessed sibling leaves" ); - // all3 Parquet leaves should be in the projection mask - let expected_mask = ProjectionMask::leaves(schema_descr, [0, 1, 2]); - assert_eq!(read_plan.projection_mask, expected_mask,); + let s_field = candidate + .read_plan + .projected_schema + .field_with_name("s") + .unwrap(); + let expected_inner: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Field::new("b", DataType::Int32, false)), + ] + .into(); + let expected_outer: Fields = vec![Arc::new(Field::new( + "outer", + DataType::Struct(expected_inner), + false, + ))] + .into(); + assert_eq!( + s_field.data_type(), + &DataType::Struct(expected_outer), + "projected schema should keep only the shared-prefix subtree" + ); } - /// Sanity check that the given expression could be evaluated against the given schema without any errors. - /// This will fail if the expression references columns that are not in the schema or if the types of the columns are incompatible, etc. - fn check_expression_can_evaluate_against_schema( - expr: &Arc, - table_schema: &Arc, - ) -> bool { - let batch = RecordBatch::new_empty(Arc::clone(table_schema)); - expr.evaluate(&batch).is_ok() + /// Two predicates touch disjoint subtrees of the same struct root: + /// `s['outer']['a'] AND s['other']['x']`. Both subtrees must be retained + /// in the projection mask and in the projected schema. + #[test] + fn get_field_disjoint_subtrees_keep_both() { + // Schema: s (Struct{outer: Struct{a, b}, other: Struct{x, y}}) + // Parquet leaves: s.outer.a=0, s.outer.b=1, s.other.x=2, s.other.y=3 + let outer_fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Field::new("b", DataType::Int32, false)), + ] + .into(); + let other_fields: Fields = vec![ + Arc::new(Field::new("x", DataType::Int32, false)), + Arc::new(Field::new("y", DataType::Int32, false)), + ] + .into(); + let s_fields: Fields = vec![ + Arc::new(Field::new( + "outer", + DataType::Struct(outer_fields.clone()), + false, + )), + Arc::new(Field::new( + "other", + DataType::Struct(other_fields.clone()), + false, + )), + ] + .into(); + let schema = Arc::new(Schema::new(vec![Field::new( + "s", + DataType::Struct(s_fields.clone()), + false, + )])); + + let outer_arr = StructArray::new( + outer_fields, + vec![ + Arc::new(Int32Array::from(vec![1, 2])) as _, + Arc::new(Int32Array::from(vec![3, 4])) as _, + ], + None, + ); + let other_arr = StructArray::new( + other_fields, + vec![ + Arc::new(Int32Array::from(vec![5, 6])) as _, + Arc::new(Int32Array::from(vec![7, 8])) as _, + ], + None, + ); + let s_arr = StructArray::new( + s_fields, + vec![Arc::new(outer_arr) as _, Arc::new(other_arr) as _], + None, + ); + let batch = + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(s_arr)]).unwrap(); + + let file = NamedTempFile::new().expect("temp file"); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) + .expect("writer"); + writer.write(&batch).expect("write batch"); + writer.close().expect("close writer"); + + let reader_file = file.reopen().expect("reopen file"); + let builder = ParquetRecordBatchReaderBuilder::try_new(reader_file) + .expect("reader builder"); + let metadata = builder.metadata().clone(); + let file_schema = builder.schema().clone(); + + // s['outer']['a'] > 0 AND s['other']['x'] > 0 + let a_expr = get_field() + .call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("outer".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("a".to_string())), None), + ]) + .gt(Expr::Literal(ScalarValue::Int32(Some(0)), None)); + let x_expr = get_field() + .call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("other".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("x".to_string())), None), + ]) + .gt(Expr::Literal(ScalarValue::Int32(Some(0)), None)); + let expr = logical2physical(&a_expr.and(x_expr), &file_schema); + + let candidate = FilterCandidateBuilder::new(expr, Arc::clone(&file_schema)) + .build(&metadata) + .expect("building candidate") + .expect("disjoint nested predicates should be pushable"); + + // s.outer.a (0) and s.other.x (2); not s.outer.b (1), not s.other.y (3). + let expected_mask = + ProjectionMask::leaves(metadata.file_metadata().schema_descr(), [0, 2]); + assert_eq!( + candidate.read_plan.projection_mask, expected_mask, + "projection_mask should keep one leaf from each disjoint subtree" + ); + + let s_field = candidate + .read_plan + .projected_schema + .field_with_name("s") + .unwrap(); + let expected_outer: Fields = + vec![Arc::new(Field::new("a", DataType::Int32, false))].into(); + let expected_other: Fields = + vec![Arc::new(Field::new("x", DataType::Int32, false))].into(); + let expected_s: Fields = vec![ + Arc::new(Field::new("outer", DataType::Struct(expected_outer), false)), + Arc::new(Field::new("other", DataType::Struct(expected_other), false)), + ] + .into(); + assert_eq!( + s_field.data_type(), + &DataType::Struct(expected_s), + "projected schema should keep one pruned field from each disjoint subtree" + ); + } + + /// End-to-end: shared-prefix nested predicates filter rows correctly during + /// Parquet decoding and report the expected pushdown metrics. + #[test] + fn get_field_end_to_end_shared_prefix_filters_rows() { + // Schema: id (Int32), s (Struct{outer: Struct{a, b}}) + // Parquet leaves: id=0, s.outer.a=1, s.outer.b=2 + let outer_fields: Fields = vec![ + Arc::new(Field::new("a", DataType::Int32, false)), + Arc::new(Field::new("b", DataType::Int32, false)), + ] + .into(); + let s_fields: Fields = vec![Arc::new(Field::new( + "outer", + DataType::Struct(outer_fields.clone()), + false, + ))] + .into(); + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("s", DataType::Struct(s_fields.clone()), false), + ])); + + // +----+--------------------------+ + // | id | s | + // +----+--------------------------+ + // | 1 | {outer: {a: 10, b: 50}} | <- a>5 and b<100 → match + // | 2 | {outer: {a: 0, b: 60}} | <- a>5 fails → drop + // | 3 | {outer: {a: 20, b: 80}} | <- a>5 and b<100 → match + // | 4 | {outer: {a: 30, b: 200}} | <- b<100 fails → drop + // +----+--------------------------+ + let outer_arr = StructArray::new( + outer_fields, + vec![ + Arc::new(Int32Array::from(vec![10, 0, 20, 30])) as _, + Arc::new(Int32Array::from(vec![50, 60, 80, 200])) as _, + ], + None, + ); + let s_arr = StructArray::new(s_fields, vec![Arc::new(outer_arr) as _], None); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4])), + Arc::new(s_arr), + ], + ) + .unwrap(); + + let file = NamedTempFile::new().expect("temp file"); + let mut writer = + ArrowWriter::try_new(file.reopen().unwrap(), Arc::clone(&schema), None) + .expect("writer"); + writer.write(&batch).expect("write batch"); + writer.close().expect("close writer"); + + let reader_file = file.reopen().expect("reopen file"); + let parquet_reader_builder = + ParquetRecordBatchReaderBuilder::try_new(reader_file) + .expect("reader builder"); + let metadata = parquet_reader_builder.metadata().clone(); + let file_schema = parquet_reader_builder.schema().clone(); + + // s['outer']['a'] > 5 AND s['outer']['b'] < 100 + let a_expr = get_field() + .call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("outer".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("a".to_string())), None), + ]) + .gt(Expr::Literal(ScalarValue::Int32(Some(5)), None)); + let b_expr = get_field() + .call(vec![ + col("s"), + Expr::Literal(ScalarValue::Utf8(Some("outer".to_string())), None), + Expr::Literal(ScalarValue::Utf8(Some("b".to_string())), None), + ]) + .lt(Expr::Literal(ScalarValue::Int32(Some(100)), None)); + let expr = logical2physical(&a_expr.and(b_expr), &file_schema); + + let metrics = ExecutionPlanMetricsSet::new(); + let file_metrics = + ParquetFileMetrics::new(0, "shared_prefix_e2e.parquet", &metrics); + + let row_filter = + build_row_filter(&expr, &file_schema, &metadata, false, &file_metrics) + .expect("building row filter") + .expect("row filter should exist"); + + let reader = parquet_reader_builder + .with_row_filter(row_filter) + .build() + .expect("build reader"); + + let mut total_rows = 0; + for batch in reader { + let batch = batch.expect("record batch"); + total_rows += batch.num_rows(); + } + + assert_eq!( + total_rows, 2, + "expected 2 rows matching s.outer.a > 5 AND s.outer.b < 100" + ); + assert_eq!(file_metrics.pushdown_rows_pruned.value(), 2); + assert_eq!(file_metrics.pushdown_rows_matched.value(), 2); } } diff --git a/datafusion/datasource-parquet/src/row_group_filter.rs b/datafusion/datasource-parquet/src/row_group_filter.rs index 07f4fe92cf308..ddf71bb7e6d95 100644 --- a/datafusion/datasource-parquet/src/row_group_filter.rs +++ b/datafusion/datasource-parquet/src/row_group_filter.rs @@ -18,10 +18,8 @@ use std::collections::HashSet; use std::sync::Arc; -use super::{ParquetAccessPlan, ParquetFileMetrics}; -// Re-exported so the existing `crate::row_group_filter::BloomFilterStatistics` -// path keeps resolving for in-crate callers (e.g. `opener`). -pub(crate) use crate::bloom_filter::BloomFilterStatistics; +use super::{ParquetAccessPlan, ParquetFileMetrics, RowGroupAccess}; +use crate::bloom_filter::BloomFilterStatistics; use arrow::array::{ArrayRef, BooleanArray, UInt64Array}; use arrow::datatypes::Schema; use datafusion_common::pruning::PruningStatistics; @@ -31,7 +29,7 @@ use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, IsNullExpr, NotExpr}; use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::{PhysicalExpr, PhysicalExprSimplifier}; -use datafusion_pruning::PruningPredicate; +use datafusion_pruning::{PruningPredicate, PruningPredicateBuilder}; use parquet::arrow::arrow_reader::statistics::StatisticsConverter; use parquet::file::metadata::RowGroupMetaData; use parquet::schema::types::SchemaDescriptor; @@ -190,7 +188,11 @@ impl RowGroupAccessPlanFilter { // find a set of matching row groups that can satisfy the limit for &idx in self.access_plan.row_group_indexes().iter() { if self.access_plan.is_fully_matched(idx) { - let row_group_row_count = rg_metadata[idx].num_rows() as usize; + let row_group_row_count = match &self.access_plan.inner()[idx] { + RowGroupAccess::Skip => continue, + RowGroupAccess::Scan => rg_metadata[idx].num_rows() as usize, + RowGroupAccess::Selection(selection) => selection.row_count(), + }; fully_matched_row_group_indexes.push(idx); fully_matched_rows_count += row_group_row_count; if fully_matched_rows_count >= limit { @@ -211,7 +213,7 @@ impl RowGroupAccessPlanFilter { let mut new_access_plan = ParquetAccessPlan::new_none(rg_metadata.len()); for &idx in &fully_matched_row_group_indexes { - new_access_plan.scan(idx); + new_access_plan.set(idx, self.access_plan.inner()[idx].clone()); new_access_plan.mark_fully_matched(idx); } self.access_plan = new_access_plan; @@ -371,8 +373,9 @@ impl RowGroupAccessPlanFilter { return; }; - let Ok(inverted_predicate) = - PruningPredicate::try_new(inverted_expr, Arc::clone(predicate.schema())) + let Ok(inverted_predicate) = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(predicate.schema())) + .try_build(inverted_expr) else { return; }; @@ -416,7 +419,7 @@ impl RowGroupAccessPlanFilter { /// /// # Panics /// if `row_group_bloom_filters` does not have the same number of row groups as this set - pub(crate) fn prune_by_bloom_filters( + pub fn prune_by_bloom_filters( &mut self, predicate: &PruningPredicate, metrics: &ParquetFileMetrics, @@ -453,12 +456,16 @@ impl RowGroupAccessPlanFilter { } } -/// Wraps a slice of [`RowGroupMetaData`] in a way that implements [`PruningStatistics`] -struct RowGroupPruningStatistics<'a> { - parquet_schema: &'a SchemaDescriptor, - row_group_metadatas: Vec<&'a RowGroupMetaData>, - arrow_schema: &'a Schema, - missing_null_counts_as_zero: bool, +/// Wraps a slice of [`RowGroupMetaData`] in a way that implements [`PruningStatistics`]. +/// +/// Visible to sibling modules so runtime row-group pruners (e.g. the dynamic +/// TopK pruner in `push_decoder.rs`) can reuse this adapter without +/// duplicating the statistics-to-`PruningStatistics` plumbing. +pub(crate) struct RowGroupPruningStatistics<'a> { + pub(crate) parquet_schema: &'a SchemaDescriptor, + pub(crate) row_group_metadatas: Vec<&'a RowGroupMetaData>, + pub(crate) arrow_schema: &'a Schema, + pub(crate) missing_null_counts_as_zero: bool, } impl<'a> RowGroupPruningStatistics<'a> { @@ -535,6 +542,7 @@ mod tests { use datafusion_physical_expr::planner::logical2physical; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; use parquet::arrow::ArrowSchemaConverter; + use parquet::arrow::arrow_reader::{RowSelection, RowSelector}; use parquet::basic::LogicalType; use parquet::data_type::{ByteArray, FixedLenByteArray}; use parquet::file::metadata::ColumnChunkMetaData; @@ -543,6 +551,16 @@ mod tests { schema::types::SchemaDescPtr, }; + fn build_test_pruning_predicate( + expr: Arc, + schema: Arc, + ) -> PruningPredicate { + PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr) + .unwrap() + } + struct PrimitiveTypeField { name: &'static str, physical_ty: PhysicalType, @@ -605,7 +623,7 @@ mod tests { Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); let expr = col("c1").gt(lit(15)); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let field = PrimitiveTypeField::new("c1", PhysicalType::INT32); let schema_descr = get_test_schema_descr(vec![field]); @@ -648,7 +666,7 @@ mod tests { let schema = Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, true)])); let expr = logical2physical(&col("c1").gt(lit(15)), &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let field = PrimitiveTypeField::new("c1", PhysicalType::INT32); let schema_descr = get_test_schema_descr(vec![field]); @@ -701,6 +719,71 @@ mod tests { assert_eq!(row_groups.is_fully_matched(), &vec![false, true, false]); } + #[test] + fn prune_by_limit_preserves_row_selection() { + let field = PrimitiveTypeField::new("c1", PhysicalType::INT32); + let schema_descr = get_test_schema_descr(vec![field]); + let rgm1 = get_row_group_meta_data( + &schema_descr, + vec![ParquetStatistics::int32(None, None, None, Some(0), false)], + ); + let rgm2 = get_row_group_meta_data( + &schema_descr, + vec![ParquetStatistics::int32(None, None, None, Some(0), false)], + ); + let groups = &[rgm1, rgm2]; + + let selection = + RowSelection::from(vec![RowSelector::skip(900), RowSelector::select(100)]); + let mut access_plan = ParquetAccessPlan::new_all(2); + access_plan.scan_selection(0, selection.clone()); + access_plan.mark_fully_matched(0); + access_plan.mark_fully_matched(1); + + let metrics = parquet_file_metrics(); + let mut row_groups = RowGroupAccessPlanFilter::new(access_plan); + row_groups.prune_by_limit(50, groups, &metrics); + + assert_eq!(row_groups.access_plan.row_group_indexes(), vec![0]); + assert_eq!( + row_groups.access_plan.inner(), + &[RowGroupAccess::Selection(selection), RowGroupAccess::Skip] + ); + assert_eq!(row_groups.is_fully_matched(), &vec![true, false]); + } + + #[test] + fn prune_by_limit_counts_only_selected_rows() { + let field = PrimitiveTypeField::new("c1", PhysicalType::INT32); + let schema_descr = get_test_schema_descr(vec![field]); + let rgm1 = get_row_group_meta_data( + &schema_descr, + vec![ParquetStatistics::int32(None, None, None, Some(0), false)], + ); + let rgm2 = get_row_group_meta_data( + &schema_descr, + vec![ParquetStatistics::int32(None, None, None, Some(0), false)], + ); + let groups = &[rgm1, rgm2]; + + let selection = + RowSelection::from(vec![RowSelector::select(10), RowSelector::skip(990)]); + let mut access_plan = ParquetAccessPlan::new_all(2); + access_plan.scan_selection(0, selection.clone()); + access_plan.mark_fully_matched(0); + + let metrics = parquet_file_metrics(); + let mut row_groups = RowGroupAccessPlanFilter::new(access_plan); + row_groups.prune_by_limit(50, groups, &metrics); + + assert_eq!(row_groups.access_plan.row_group_indexes(), vec![0, 1]); + assert_eq!( + row_groups.access_plan.inner(), + &[RowGroupAccess::Selection(selection), RowGroupAccess::Scan] + ); + assert_eq!(row_groups.is_fully_matched(), &vec![true, false]); + } + #[test] fn row_group_pruning_predicate_missing_stats() { use datafusion_expr::{col, lit}; @@ -709,7 +792,7 @@ mod tests { Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); let expr = col("c1").gt(lit(15)); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let field = PrimitiveTypeField::new("c1", PhysicalType::INT32); let schema_descr = get_test_schema_descr(vec![field]); @@ -752,7 +835,7 @@ mod tests { ])); let expr = col("c1").gt(lit(15)).and(col("c2").rem(lit(2)).eq(lit(0))); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let schema_descr = get_test_schema_descr(vec![ PrimitiveTypeField::new("c1", PhysicalType::INT32), @@ -791,7 +874,7 @@ mod tests { // this bypasses the entire predicate expression and no row groups are filtered out let expr = col("c1").gt(lit(15)).or(col("c2").rem(lit(2)).eq(lit(0))); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); // if conditions in predicate are joined with OR and an unsupported expression is used // this bypasses the entire predicate expression and no row groups are filtered out @@ -818,7 +901,7 @@ mod tests { let expr = col("c1").gt(lit(0)); let expr = logical2physical(&expr, &table_schema); let pruning_predicate = - PruningPredicate::try_new(expr, table_schema.clone()).unwrap(); + build_test_pruning_predicate(expr, Arc::clone(&table_schema)); // Model a file schema's column order c2 then c1, which is the opposite // of the table schema @@ -895,7 +978,7 @@ mod tests { let schema_descr = ArrowSchemaConverter::new().convert(&schema).unwrap(); let expr = col("c1").gt(lit(15)).and(col("c2").is_null()); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let groups = gen_row_group_meta_data_for_pruning_predicate(); let metrics = parquet_file_metrics(); @@ -926,7 +1009,7 @@ mod tests { .gt(lit(15)) .and(col("c2").eq(lit(ScalarValue::Boolean(None)))); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let groups = gen_row_group_meta_data_for_pruning_predicate(); let metrics = parquet_file_metrics(); @@ -955,16 +1038,13 @@ mod tests { let schema = Arc::new(Schema::new(vec![Field::new("c1", Decimal128(9, 2), false)])); let field = PrimitiveTypeField::new("c1", PhysicalType::INT32) - .with_logical_type(LogicalType::Decimal { - scale: 2, - precision: 9, - }) + .with_logical_type(LogicalType::decimal(2, 9)) .with_scale(2) .with_precision(9); let schema_descr = get_test_schema_descr(vec![field]); let expr = col("c1").gt(lit(ScalarValue::Decimal128(Some(500), 9, 2))); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let rgm1 = get_row_group_meta_data( &schema_descr, // [1.00, 6.00] @@ -1023,10 +1103,7 @@ mod tests { Arc::new(Schema::new(vec![Field::new("c1", Decimal128(9, 0), false)])); let field = PrimitiveTypeField::new("c1", PhysicalType::INT32) - .with_logical_type(LogicalType::Decimal { - scale: 0, - precision: 9, - }) + .with_logical_type(LogicalType::decimal(0, 9)) .with_scale(0) .with_precision(9); let schema_descr = get_test_schema_descr(vec![field]); @@ -1035,7 +1112,7 @@ mod tests { Decimal128(11, 2), )); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let rgm1 = get_row_group_meta_data( &schema_descr, // [100, 600] @@ -1118,16 +1195,13 @@ mod tests { false, )])); let field = PrimitiveTypeField::new("c1", PhysicalType::INT64) - .with_logical_type(LogicalType::Decimal { - scale: 2, - precision: 18, - }) + .with_logical_type(LogicalType::decimal(2, 18)) .with_scale(2) .with_precision(18); let schema_descr = get_test_schema_descr(vec![field]); let expr = col("c1").lt(lit(ScalarValue::Decimal128(Some(500), 18, 2))); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); let rgm1 = get_row_group_meta_data( &schema_descr, // [6.00, 8.00] @@ -1176,10 +1250,7 @@ mod tests { false, )])); let field = PrimitiveTypeField::new("c1", PhysicalType::FIXED_LEN_BYTE_ARRAY) - .with_logical_type(LogicalType::Decimal { - scale: 2, - precision: 18, - }) + .with_logical_type(LogicalType::decimal(2, 18)) .with_scale(2) .with_precision(18) .with_byte_len(16); @@ -1188,7 +1259,7 @@ mod tests { let left = cast(col("c1"), Decimal128(28, 3)); let expr = left.eq(lit(ScalarValue::Decimal128(Some(100000), 28, 3))); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); // we must use the big-endian when encode the i128 to bytes or vec[u8]. let rgm1 = get_row_group_meta_data( &schema_descr, @@ -1254,10 +1325,7 @@ mod tests { false, )])); let field = PrimitiveTypeField::new("c1", PhysicalType::BYTE_ARRAY) - .with_logical_type(LogicalType::Decimal { - scale: 2, - precision: 18, - }) + .with_logical_type(LogicalType::decimal(2, 18)) .with_scale(2) .with_precision(18) .with_byte_len(16); @@ -1266,7 +1334,7 @@ mod tests { let left = cast(col("c1"), Decimal128(28, 3)); let expr = left.eq(lit(ScalarValue::Decimal128(Some(100000), 28, 3))); let expr = logical2physical(&expr, &schema); - let pruning_predicate = PruningPredicate::try_new(expr, schema.clone()).unwrap(); + let pruning_predicate = build_test_pruning_predicate(expr, Arc::clone(&schema)); // we must use the big-endian when encode the i128 to bytes or vec[u8]. let rgm1 = get_row_group_meta_data( &schema_descr, diff --git a/datafusion/datasource-parquet/src/schema_coercion.rs b/datafusion/datasource-parquet/src/schema_coercion.rs index 4598bb525be32..30cd5d7e65948 100644 --- a/datafusion/datasource-parquet/src/schema_coercion.rs +++ b/datafusion/datasource-parquet/src/schema_coercion.rs @@ -418,118 +418,6 @@ fn coerce_int96_to_resolution_impl( Some(transformed_schema) } -/// Coerces the file schema if the table schema uses a view type. -#[deprecated( - since = "47.0.0", - note = "Use `apply_file_schema_type_coercions` instead" -)] -pub fn coerce_file_schema_to_view_type( - table_schema: &Schema, - file_schema: &Schema, -) -> Option { - let mut transform = false; - let table_fields: HashMap<_, _> = table_schema - .fields - .iter() - .map(|f| { - let dt = f.data_type(); - if dt.equals_datatype(&DataType::Utf8View) - || dt.equals_datatype(&DataType::BinaryView) - { - transform = true; - } - (f.name(), dt) - }) - .collect(); - - if !transform { - return None; - } - - let transformed_fields: Vec> = file_schema - .fields - .iter() - .map( - |field| match (table_fields.get(field.name()), field.data_type()) { - (Some(DataType::Utf8View), DataType::Utf8 | DataType::LargeUtf8) => { - field_with_new_type(field, DataType::Utf8View) - } - ( - Some(DataType::BinaryView), - DataType::Binary | DataType::LargeBinary, - ) => field_with_new_type(field, DataType::BinaryView), - _ => Arc::clone(field), - }, - ) - .collect(); - - Some(Schema::new_with_metadata( - transformed_fields, - file_schema.metadata.clone(), - )) -} - -/// If the table schema uses a string type, coerce the file schema to use a string type. -/// -/// See [`ParquetFormat::binary_as_string`](crate::file_format::ParquetFormat::binary_as_string) for details -#[deprecated( - since = "47.0.0", - note = "Use `apply_file_schema_type_coercions` instead" -)] -pub fn coerce_file_schema_to_string_type( - table_schema: &Schema, - file_schema: &Schema, -) -> Option { - let mut transform = false; - let table_fields: HashMap<_, _> = table_schema - .fields - .iter() - .map(|f| (f.name(), f.data_type())) - .collect(); - let transformed_fields: Vec> = file_schema - .fields - .iter() - .map( - |field| match (table_fields.get(field.name()), field.data_type()) { - // table schema uses string type, coerce the file schema to use string type - ( - Some(DataType::Utf8), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::Utf8) - } - // table schema uses large string type, coerce the file schema to use large string type - ( - Some(DataType::LargeUtf8), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::LargeUtf8) - } - // table schema uses string view type, coerce the file schema to use view type - ( - Some(DataType::Utf8View), - DataType::Binary | DataType::LargeBinary | DataType::BinaryView, - ) => { - transform = true; - field_with_new_type(field, DataType::Utf8View) - } - _ => Arc::clone(field), - }, - ) - .collect(); - - if !transform { - None - } else { - Some(Schema::new_with_metadata( - transformed_fields, - file_schema.metadata.clone(), - )) - } -} - /// Create a new field with the specified data type, copying the other /// properties from the input field fn field_with_new_type(field: &FieldRef, new_type: DataType) -> FieldRef { diff --git a/datafusion/datasource-parquet/src/sink.rs b/datafusion/datasource-parquet/src/sink.rs index a73be8d2e68cf..53f6f1e6b4323 100644 --- a/datafusion/datasource-parquet/src/sink.rs +++ b/datafusion/datasource-parquet/src/sink.rs @@ -33,6 +33,8 @@ use datafusion_datasource::display::FileGroupDisplay; use datafusion_datasource::file_compression_type::FileCompressionType; use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; use datafusion_datasource::sink::DataSink; +#[cfg(feature = "proto")] +use datafusion_datasource::sink::DataSinkExec; use datafusion_datasource::write::demux::DemuxedStreamReceiver; use datafusion_datasource::write::{ ObjectWriterBuilder, SharedBuffer, get_writer_schema, @@ -40,6 +42,8 @@ use datafusion_datasource::write::{ use datafusion_execution::memory_pool::{MemoryConsumer, MemoryPool, MemoryReservation}; use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; +#[cfg(feature = "proto")] +use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::metrics::{ ElapsedComputeFutureExt, ExecutionPlanMetricsSet, MetricBuilder, MetricCategory, MetricsSet, Time, @@ -171,7 +175,7 @@ impl ParquetSink { /// Creates an AsyncArrowWriter which serializes a parquet file to an ObjectStore /// AsyncArrowWriters are used when individual parquet file serialization is not parallelized - async fn create_async_arrow_writer( + fn create_async_arrow_writer( &self, location: &Path, object_store: Arc, @@ -239,6 +243,7 @@ async fn set_writer_encryption_properties( } #[cfg(not(feature = "parquet_encryption"))] +#[expect(clippy::unused_async)] async fn set_writer_encryption_properties( builder: WriterPropertiesBuilder, _runtime: &Arc, @@ -294,16 +299,14 @@ impl FileSink for ParquetSink { // CDC requires the sequential writer: the chunker state lives in ArrowWriter // and persists across row groups. The parallel path bypasses ArrowWriter entirely. if !parquet_opts.global.allow_single_file_parallelism - || parquet_opts.global.use_content_defined_chunking.is_some() + || parquet_opts.global.content_defined_chunking.enabled { - let mut writer = self - .create_async_arrow_writer( - &path, - Arc::clone(&object_store), - context, - parquet_props.clone(), - ) - .await?; + let mut writer = self.create_async_arrow_writer( + &path, + Arc::clone(&object_store), + context, + parquet_props.clone(), + )?; let reservation = MemoryConsumer::new(format!("ParquetSink[{path}]")) .register(context.memory_pool()); file_write_tasks.spawn( @@ -411,6 +414,105 @@ impl DataSink for ParquetSink { ) -> Result { FileSink::write_all(self, data, context).await } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let input = ctx.encode_child(exec.input())?; + let sort_order = exec.encode_sort_order(ctx)?; + let sink = protobuf::ParquetSink::try_from(self)?; + let node = protobuf::ParquetSinkExecNode { + input: Some(Box::new(input)), + sink: Some(sink), + sink_schema: Some(exec.schema().as_ref().try_into()?), + sort_order, + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ParquetSink(Box::new(node))), + })) + } +} + +#[cfg(feature = "proto")] +impl TryFrom<&ParquetSink> for datafusion_proto_models::protobuf::ParquetSink { + type Error = DataFusionError; + + fn try_from(value: &ParquetSink) -> Result { + Ok(Self { + config: Some(value.config().try_into()?), + parquet_options: Some(value.parquet_options().try_into()?), + }) + } +} + +#[cfg(feature = "proto")] +impl TryFrom<&datafusion_proto_models::protobuf::ParquetSink> for ParquetSink { + type Error = DataFusionError; + + fn try_from(value: &datafusion_proto_models::protobuf::ParquetSink) -> Result { + let config = + FileSinkConfig::try_from(value.config.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSink is missing required field 'config'" + ) + })?)?; + let parquet_options = value + .parquet_options + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSink is missing required field 'parquet_options'" + ) + })? + .try_into()?; + + Ok(Self::new(config, parquet_options)) + } +} + +#[cfg(feature = "proto")] +impl ParquetSink { + /// Reconstructs a [`DataSinkExec`] containing a `ParquetSink` from protobuf. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let sink_node = datafusion_physical_plan::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::ParquetSink, + "ParquetSink", + ); + let input = ctx.decode_required_child( + sink_node.input.as_deref(), + "ParquetSinkExecNode", + "input", + )?; + let proto_sink = sink_node.sink.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetSinkExecNode is missing required field 'sink'" + ) + })?; + let data_sink = ParquetSink::try_from(proto_sink)?; + let sort_order = DataSinkExec::decode_sort_order( + sink_node.sort_order.as_ref(), + ctx, + input.schema().as_ref(), + )?; + + Ok(Arc::new(DataSinkExec::new( + input, + Arc::new(data_sink), + sort_order, + ))) + } } /// Consumes a stream of [ArrowLeafColumn] via a channel and serializes them using an [ArrowColumnWriter] diff --git a/datafusion/datasource-parquet/src/sort.rs b/datafusion/datasource-parquet/src/sort.rs index c1cf4e8b7824e..ea33fb0e2ecb2 100644 --- a/datafusion/datasource-parquet/src/sort.rs +++ b/datafusion/datasource-parquet/src/sort.rs @@ -124,9 +124,14 @@ pub fn reverse_row_selection( /// Reorder a file list so the most "promising" files are read first, /// matching `PreparedAccessPlan::reorder_by_statistics` at the -/// row-group level: key off the file's `min(col)`, and let the sort -/// direction follow the request (ASC by `min` for ASC requests, DESC -/// by `min` for DESC requests). +/// row-group level: key lexicographically off the file's per-column +/// `min` for the longest plain-`Column` prefix of the sort order, and +/// let the leading sort direction follow the request (ASC by `min` +/// for ASC requests, DESC by `min` for DESC requests). +/// +/// Secondary sort keys break ties when the leading column's `min` is +/// equal across files (e.g. `ORDER BY low_cardinality_col, ts LIMIT k`), +/// mirroring the row-group level lexicographic reorder. /// /// Keeping both layers consistent matters because they share the same /// convergence story for TopK's dynamic filter: file `i`'s `min` is a @@ -147,51 +152,81 @@ pub(crate) fn reorder_files_by_min_statistics( reverse_row_groups: bool, table_schema: &Schema, ) -> Vec { - let Some((col_name, descending)) = - extract_topk_sort_info(sort_order, reverse_row_groups) - else { + let sort_keys = extract_topk_sort_info(sort_order, reverse_row_groups); + if sort_keys.is_empty() { return files; - }; + } - let Ok(col_idx) = table_schema.index_of(&col_name) else { - return files; - }; + // Resolve names to column indexes; the leading key is required, later + // keys are best-effort (stop at the first unresolvable one). + let mut keys: Vec<(usize, bool)> = Vec::with_capacity(sort_keys.len()); + for (col_name, descending) in &sort_keys { + match table_schema.index_of(col_name) { + Ok(idx) => keys.push((idx, *descending)), + Err(_) if keys.is_empty() => return files, + Err(_) => break, + } + } files.sort_by(|a, b| { - let key_a = file_min_value(a, col_idx); - let key_b = file_min_value(b, col_idx); - match (key_a, key_b) { - (Some(va), Some(vb)) => { - let cmp = va.partial_cmp(&vb).unwrap_or(std::cmp::Ordering::Equal); - if descending { cmp.reverse() } else { cmp } + for &(col_idx, descending) in &keys { + let key_a = file_min_value(a, col_idx); + let key_b = file_min_value(b, col_idx); + let ord = match (key_a, key_b) { + (Some(va), Some(vb)) => { + let cmp = va.partial_cmp(&vb).unwrap_or(std::cmp::Ordering::Equal); + if descending { cmp.reverse() } else { cmp } + } + // Missing stats always sort last, regardless of direction. + (Some(_), None) => std::cmp::Ordering::Less, + (None, Some(_)) => std::cmp::Ordering::Greater, + (None, None) => std::cmp::Ordering::Equal, + }; + if ord != std::cmp::Ordering::Equal { + return ord; } - (Some(_), None) => std::cmp::Ordering::Less, - (None, Some(_)) => std::cmp::Ordering::Greater, - (None, None) => std::cmp::Ordering::Equal, } + std::cmp::Ordering::Equal }); log::debug!( - "Reordered {} files by min({}) {} for TopK optimization", + "Reordered {} files by lexicographic min of {:?} for TopK optimization", files.len(), - col_name, - if descending { "DESC" } else { "ASC" } + sort_keys, ); files } -/// Extract the `(column name, descending)` tuple used by file-level -/// reordering. Returns `None` when the sort order isn't set or the -/// leading sort expression isn't a plain `Column`. +/// Extract the `(column name, descending)` keys used by file-level +/// reordering: the longest prefix of the sort order made of plain +/// `Column` expressions. Returns an empty vec when the sort order isn't +/// set or the leading sort expression isn't a plain `Column`. +/// +/// The leading key's direction is `reverse_row_groups` (the pushdown's +/// authoritative flip decision, which may differ from the raw +/// expression's `descending` in the `reversed_satisfies` case); +/// subsequent keys apply their direction *relative to the leading +/// expression* on top of that flag, so a request like +/// `[a DESC, b ASC]` with `reverse_row_groups=true` sorts by +/// `(min(a) DESC, min(b) ASC)`. fn extract_topk_sort_info( sort_order: Option<&LexOrdering>, reverse_row_groups: bool, -) -> Option<(String, bool)> { - let sort_order = sort_order?; - let first = sort_order.first(); - let col = first.expr.downcast_ref::()?; - Some((col.name().to_string(), reverse_row_groups)) +) -> Vec<(String, bool)> { + let Some(sort_order) = sort_order else { + return vec![]; + }; + let leading_descending = sort_order.first().options.descending; + let mut keys = Vec::new(); + for sort_expr in sort_order.iter() { + let Some(col) = sort_expr.expr.downcast_ref::() else { + break; + }; + let relative_desc = sort_expr.options.descending != leading_descending; + keys.push((col.name().to_string(), reverse_row_groups != relative_desc)); + } + keys } /// File's per-column `min` for the reorder key. diff --git a/datafusion/datasource-parquet/src/source.rs b/datafusion/datasource-parquet/src/source.rs index 2b367cf7600d5..097b4563af5df 100644 --- a/datafusion/datasource-parquet/src/source.rs +++ b/datafusion/datasource-parquet/src/source.rs @@ -24,7 +24,11 @@ use crate::DefaultParquetFileReaderFactory; use crate::ParquetFileReaderFactory; use crate::opener::ParquetMorselizer; use crate::opener::build_pruning_predicates; +use crate::opener::build_virtual_columns_state; use crate::row_filter::can_expr_be_pushed_down_with_schemas; +use arrow_schema::Fields; +use arrow_schema::extension::ExtensionType; +use arrow_schema::{DataType, Field}; use datafusion_common::config::ConfigOptions; #[cfg(feature = "parquet_encryption")] use datafusion_common::config::EncryptionFactoryOptions; @@ -40,9 +44,14 @@ use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; use datafusion_datasource::file_scan_config::FileScanConfig; +use datafusion_functions::core::file_row_index::FileRowIndexFunc; +use datafusion_physical_expr::expressions::{Column, DynamicFilterTracking}; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_expr::{EquivalenceProperties, conjunction}; use datafusion_physical_expr_adapter::DefaultPhysicalExprAdapterFactory; +use datafusion_physical_expr_adapter::rewrite::{ + expr_references_scalar_udf, rewrite_file_row_index_projection, +}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::physical_expr::fmt_sql; use datafusion_physical_plan::DisplayFormatType; @@ -60,6 +69,7 @@ use datafusion_execution::parquet_encryption::EncryptionFactory; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use itertools::Itertools; use object_store::ObjectStore; +use parquet::arrow::RowNumber; #[cfg(feature = "parquet_encryption")] use parquet::encryption::decrypt::FileDecryptionProperties; @@ -208,6 +218,12 @@ use parquet::encryption::decrypt::FileDecryptionProperties; /// used to implement external indexes on top of parquet files and select only /// portions of the files. /// +/// If the external index naturally produces a file-level +/// [`RowSelection`](parquet::arrow::arrow_reader::RowSelection), wrap it in +/// [`ParquetRowSelection`](crate::ParquetRowSelection) and provide it as an +/// extension. DataFusion will use the parquet metadata to split the selection +/// into row-group-level access. +/// /// The `DataSourceExec` will try and reduce any provided `ParquetAccessPlan` /// further based on the contents of `ParquetMetadata` and other settings. /// @@ -347,7 +363,11 @@ impl ParquetSource { self } - /// Set predicate information + /// Set predicate information. + /// + /// Predicates referencing virtual columns must go through + /// [`Self::try_pushdown_filters`]. Passing them here with pushdown + /// enabled trips a debug assert in the opener. #[expect(clippy::needless_pass_by_value)] pub fn with_predicate(&self, predicate: Arc) -> Self { let mut conf = self.clone(); @@ -466,6 +486,14 @@ impl ParquetSource { self.table_parquet_options.global.max_predicate_cache_size } + /// Return the maximum size of an `IN (...)` list that the pruning + /// predicate will rewrite into per-value statistics checks. Lists + /// longer than this skip container-level pruning. Reads from + /// `datafusion.execution.parquet.max_in_list_size`. + pub fn max_in_list_size(&self) -> usize { + self.table_parquet_options.global.max_in_list_size + } + #[cfg(feature = "parquet_encryption")] fn get_encryption_factory_with_config( &self, @@ -585,6 +613,22 @@ impl FileSource for ParquetSource { ); } + // Validate virtual columns (extension-type allowlist) and, when + // pushdown is enabled, reject predicates that reference them. Both + // checks depend only on morselizer-level state, so we pay their cost + // once per scan partition rather than per file. + // + // Gating predicate validation on `pushdown_filters` is deliberate: + // when pushdown is off the predicate stays above the scan as a + // `FilterExec` and resolves virtual columns there; the row-filter + // ban only applies to the pushdown path. + let virtual_state = build_virtual_columns_state( + self.table_schema.virtual_columns(), + self.table_schema.file_schema(), + self.predicate.as_ref(), + self.pushdown_filters(), + )?; + Ok(Box::new(ParquetMorselizer { partition_index: partition, projection: self.projection.clone(), @@ -612,8 +656,10 @@ impl FileSource for ParquetSource { #[cfg(feature = "parquet_encryption")] encryption_factory: self.get_encryption_factory_with_config(), max_predicate_cache_size: self.max_predicate_cache_size(), + max_in_list_size: self.max_in_list_size(), reverse_row_groups: self.reverse_row_groups, sort_order_for_reorder: self.sort_order_for_reorder.clone(), + virtual_state, })) } @@ -648,7 +694,28 @@ impl FileSource for ParquetSource { projection: &ProjectionExprs, ) -> datafusion_common::Result>> { let mut source = self.clone(); - source.projection = self.projection.try_merge(projection)?; + + // If there's no reference to `FileRowIndexFunc` in the projection, we can just merge + // both projections as-is, there's no need to modify the projection first. + if !projection.iter().any(|projection_expr| { + expr_references_scalar_udf::(&projection_expr.expr) + }) { + source.projection = self.projection.try_merge(projection)?; + return Ok(Some(Arc::new(source))); + } + + // If we can find a reference to `FileRowIndexFunc`, we add it as a virtual column + // or re-use an existing one in the table's schema. + let (table_schema, row_index_col) = + table_schema_with_row_index_col(self.table_schema()); + + source.table_schema = table_schema; + source.projection = rewrite_file_row_index_projection( + &self.projection, + projection, + &row_index_col, + )?; + Ok(Some(Arc::new(source))) } @@ -683,6 +750,33 @@ impl FileSource for ParquetSource { write!(f, ", reverse_row_groups=true")?; } + // Plan-time marker for dynamic RG-level pruning: if the + // predicate is dynamic (e.g. a TopK threshold expression), + // the parquet opener will pause the single decoder at row + // group boundaries and consult `RowGroupPruner` to drop + // RGs the current threshold proves unwinnable, rebuilding + // the decoder via `into_builder().with_row_groups(...)` to + // skip them. The actual pruning count appears as + // `row_groups_pruned_dynamic_filter` in EXPLAIN ANALYZE. + // We use `contains_dynamic_filter()` (matches both `Watching` + // and `AllComplete`) rather than the stricter `Watching(_)` + // check the opener uses to construct the pruner. Reason: the + // opener gate is evaluated at file-open time, when a TopK + // threshold has not yet been pushed — at that moment a still- + // useful pruner needs `Watching`. `fmt_extra`, on the other + // hand, is called *also* by `EXPLAIN ANALYZE` after execution + // completes, at which point TopK has marked its dynamic + // filter complete and `classify` returns `AllComplete`. The + // marker is plan-time metadata ("this scan was eligible for + // runtime RG pruning"), so it should still show in that + // post-run rendering. + if let Some(predicate) = self.filter() + && DynamicFilterTracking::classify(&predicate) + .contains_dynamic_filter() + { + write!(f, ", dynamic_rg_pruning=eligible")?; + } + // Try to build the pruning predicates. // These are only generated here because it's useful to have *some* // idea of what pushdown is happening when viewing plans. @@ -697,6 +791,7 @@ impl FileSource for ParquetSource { Some(predicate), self.table_schema.table_schema(), &predicate_creation_errors, + self.max_in_list_size(), ) { let mut guarantees = pruning_predicate .literal_guarantees() @@ -728,7 +823,12 @@ impl FileSource for ParquetSource { filters: Vec>, config: &ConfigOptions, ) -> datafusion_common::Result>> { - let table_schema = self.table_schema.table_schema(); + // Use the schema excluding virtual columns: virtual columns (e.g. + // Parquet `row_number`) are produced by the reader itself and cannot + // be referenced inside a RowFilter, so predicates that reference them + // must not be marked as pushed down — otherwise the scan would + // silently drop them and produce wrong results. + let pushable_schema = self.table_schema.schema_without_virtual_columns(); // Determine if based on configs we should push filters down. // If either the table / scan itself or the config has pushdown enabled, // we will push down the filters. @@ -744,7 +844,7 @@ impl FileSource for ParquetSource { let filters: Vec = filters .into_iter() .map(|filter| { - if can_expr_be_pushed_down_with_schemas(&filter, table_schema) { + if can_expr_be_pushed_down_with_schemas(&filter, pushable_schema) { PushedDownPredicate::supported(filter) } else { PushedDownPredicate::unsupported(filter) @@ -926,15 +1026,12 @@ impl FileSource for ParquetSource { reversed_eq_properties.ordering_satisfy(order.iter().cloned())?; let sort_order = LexOrdering::new(order.iter().cloned()); let column_in_file_schema = sort_order.as_ref().is_some_and(|s| { - s.first() - .expr - .downcast_ref::() - .is_some_and(|col| { - self.table_schema - .file_schema() - .field_with_name(col.name()) - .is_ok() - }) + s.first().expr.downcast_ref::().is_some_and(|col| { + self.table_schema + .file_schema() + .field_with_name(col.name()) + .is_ok() + }) }); if !column_in_file_schema && !reversed_satisfies { @@ -965,22 +1062,208 @@ impl FileSource for ParquetSource { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn PhysicalExpr, + &Arc, ) -> datafusion_common::Result, ) -> datafusion_common::Result { - // Visit predicate (filter) expression if present - let mut tnr = TreeNodeRecursion::Continue; - if let Some(predicate) = &self.predicate { - tnr = tnr.visit_sibling(|| f(predicate.as_ref()))?; + datafusion_physical_plan::apply_expression_roots( + self.predicate + .iter() + .chain(self.projection.iter().map(|proj_expr| &proj_expr.expr)), + f, + ) + } + + /// Emit a `ParquetScan` node wrapping the shared base config plus the + /// Parquet-specific pushdown predicate and `TableParquetOptions`. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + base: &FileScanConfig, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> datafusion_common::Result< + Option, + > { + use datafusion_proto_models::protobuf; + use protobuf::physical_plan_node::PhysicalPlanType; + + let predicate = self + .filter() + .map(|pred| ctx.encode_expr(&pred)) + .transpose()?; + + let node = protobuf::ParquetScanExecNode { + base_conf: Some(base.try_to_proto(ctx)?), + predicate, + parquet_options: Some(self.table_parquet_options().try_into()?), + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::ParquetScan(node)), + })) + } +} + +#[cfg(feature = "proto")] +impl ParquetSource { + /// Reconstructs a `DataSourceExec` from a protobuf `ParquetScan`. + /// + /// Rebuilds the reader factory from the decode context because it is not serialized. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> datafusion_common::Result> { + use crate::CachedParquetFileReaderFactory; + use arrow::datatypes::Schema; + use datafusion_common::config::TableParquetOptions; + use datafusion_datasource::file_scan_config::FileScanConfig; + use datafusion_datasource::source::DataSourceExec; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_proto_models::protobuf; + + let scan = match &node.physical_plan_type { + Some(protobuf::physical_plan_node::PhysicalPlanType::ParquetScan(scan)) => { + scan + } + _ => { + return datafusion_common::internal_err!( + "PhysicalPlanNode is not a ParquetScan" + ); + } + }; + + let base_conf = scan.base_conf.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ParquetScanExecNode is missing required field 'base_conf'" + ) + })?; + + let schema: Arc = Arc::new( + base_conf + .schema + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "FileScanExecConf is missing required field 'schema'" + ) + })? + .try_into()?, + ); + + // The predicate was serialized against the scan's output schema, so it + // must be decoded against the projected schema when a projection is + // present. + let predicate_schema = if !base_conf.projection.is_empty() { + let projected_fields: Vec<_> = base_conf + .projection + .iter() + .map(|&i| schema.field(i as usize).clone()) + .collect(); + Arc::new(Schema::new(projected_fields)) + } else { + schema + }; + + let predicate = scan + .predicate + .as_ref() + .map(|expr| ctx.decode_expr(expr, predicate_schema.as_ref())) + .transpose()?; + + let mut options = TableParquetOptions::default(); + if let Some(table_options) = scan.parquet_options.as_ref() { + options = table_options.try_into()?; } - // Visit projection expressions - for proj_expr in &self.projection { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; + let table_schema = FileScanConfig::parse_table_schema_from_proto(base_conf)?; + let object_store_url = match base_conf.object_store_url.is_empty() { + false => ObjectStoreUrl::parse(&base_conf.object_store_url)?, + true => ObjectStoreUrl::local_filesystem(), + }; + let store = ctx + .task_ctx() + .runtime_env() + .object_store(object_store_url)?; + let metadata_cache = ctx + .task_ctx() + .runtime_env() + .cache_manager + .get_file_metadata_cache(); + let reader_factory = + Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); + + let mut source = ParquetSource::new(table_schema) + .with_parquet_file_reader_factory(reader_factory) + .with_table_parquet_options(options); + + if let Some(predicate) = predicate { + source = source.with_predicate(predicate); } + let base_config = + FileScanConfig::try_from_proto(base_conf, ctx, Arc::new(source))?; + Ok(DataSourceExec::from_data_source(base_config)) + } +} + +/// Returns the a [`TableSchema`] containing a [`RowNumber`] virtual column and a [`Column`] expression referencing its row index column. +/// The expression is then merged into a projection. +/// +/// - If the schema already has a virtual column with the [`RowNumber`] type, it returns the schema unchanged. +/// - If the schema doesn't have the appropriate virtual column, it returns a modified schema with the virtual column appended to it. +fn table_schema_with_row_index_col(table_schema: &TableSchema) -> (TableSchema, Column) { + // If we can find a virtual column with the `RowNumber` type, we just return the schema + // and create the appropriate `column` we're going to use + if let Some((idx, field)) = + table_schema + .virtual_columns() + .iter() + .enumerate() + .find(|(_, field)| { + field + .extension_type_name() + .is_some_and(|name| name == RowNumber::NAME) + }) + { + let virtual_offset = table_schema.file_schema().fields().len() + + table_schema.table_partition_cols().len(); + + return ( + table_schema.clone(), + Column::new(field.name(), virtual_offset + idx), + ); + } - Ok(tnr) + // The hidden field is shared across all files in this scan, but it must + // have a unique table-schema name because later rewrites resolve it by + // column name and index. + let base_row_index_name = "__datafusion_file_row_index"; + let mut row_index_name = base_row_index_name.to_string(); + let mut suffix = 0; + while table_schema + .table_schema() + .field_with_name(&row_index_name) + .is_ok() + { + suffix += 1; + row_index_name = format!("{base_row_index_name}_{suffix}"); } + + let row_index_table_idx = table_schema.table_schema().fields().len(); + let row_index_field = Arc::new( + Field::new(&row_index_name, DataType::Int64, true).with_extension_type(RowNumber), + ); + ( + TableSchema::builder(Arc::clone(table_schema.file_schema())) + .with_table_partition_cols(table_schema.table_partition_cols().clone()) + .with_virtual_columns( + table_schema + .virtual_columns() + .iter() + .cloned() + .chain([row_index_field]) + .collect::(), + ) + .build(), + Column::new(&row_index_name, row_index_table_idx), + ) } #[cfg(test)] @@ -1082,6 +1365,88 @@ mod tests { assert!(source.filter().is_some()); } + /// Render a `ParquetSource`'s `fmt_extra` output as a `String` for + /// inspection in tests. + fn render_fmt_extra(source: &ParquetSource, t: DisplayFormatType) -> String { + use std::fmt::Display; + + struct Wrap<'a> { + source: &'a ParquetSource, + t: DisplayFormatType, + } + impl Display for Wrap<'_> { + fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { + self.source.fmt_extra(self.t, f) + } + } + Wrap { source, t }.to_string() + } + + /// EXPLAIN must surface a `dynamic_rg_pruning=eligible` marker when the + /// predicate carries a `DynamicFilterPhysicalExpr`. This is the + /// plan-time signal that the runtime row-group pruner will fire at + /// every RG boundary. + #[test] + fn fmt_extra_marks_dynamic_predicate_as_pruning_eligible() { + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::{Column, DynamicFilterPhysicalExpr}; + + let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int64, false)])); + let dynamic = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("v", 0))], + lit(true), + )) as Arc; + + let source = + ParquetSource::new(Arc::clone(&schema)).with_predicate(Arc::clone(&dynamic)); + + let rendered = render_fmt_extra(&source, DisplayFormatType::Default); + assert!( + rendered.contains("dynamic_rg_pruning=eligible"), + "expected marker in Default fmt_extra, got: {rendered}" + ); + + let rendered_verbose = render_fmt_extra(&source, DisplayFormatType::Verbose); + assert!( + rendered_verbose.contains("dynamic_rg_pruning=eligible"), + "expected marker in Verbose fmt_extra, got: {rendered_verbose}" + ); + } + + /// EXPLAIN must NOT show the dynamic-RG-pruning marker when the + /// predicate is purely static — the optimization will not fire, so + /// surfacing it would mislead the reader. + #[test] + fn fmt_extra_omits_marker_for_static_predicate() { + use arrow::datatypes::Schema; + + let schema = Arc::new(Schema::empty()); + let predicate = lit(true); + let source = ParquetSource::new(schema).with_predicate(predicate); + + let rendered = render_fmt_extra(&source, DisplayFormatType::Default); + assert!( + !rendered.contains("dynamic_rg_pruning"), + "did not expect marker for static predicate, got: {rendered}" + ); + } + + /// EXPLAIN must NOT show the marker when there is no predicate at all + /// (e.g. unfiltered table scan). + #[test] + fn fmt_extra_omits_marker_when_no_predicate() { + use arrow::datatypes::Schema; + + let schema = Arc::new(Schema::empty()); + let source = ParquetSource::new(schema); + + let rendered = render_fmt_extra(&source, DisplayFormatType::Default); + assert!( + !rendered.contains("dynamic_rg_pruning"), + "did not expect marker for predicate-less scan, got: {rendered}" + ); + } + /// Helpers for the `try_pushdown_sort` regression tests below. mod pushdown_sort_helpers { use super::*; @@ -1308,7 +1673,9 @@ mod tests { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); let partition_b = Arc::new(Field::new("b", DataType::Int32, true)); - let table_schema = TableSchema::new(file_schema, vec![partition_b]); + let table_schema = TableSchema::builder(file_schema) + .with_table_partition_cols(vec![partition_b]) + .build(); let source = ParquetSource::new(table_schema); // EquivalenceProperties is built on the *full* table schema so @@ -1543,6 +1910,69 @@ mod tests { assert_eq!(names(&reordered), vec!["has_min", "no_stats"]); } + /// Multi-column TopK: when the leading column's `min` ties across + /// files, the secondary sort key breaks the tie (lexicographic, + /// mirroring the row-group level reorder). + #[test] + fn reorder_files_breaks_leading_ties_with_secondary_column() { + use datafusion_common::stats::Precision; + use datafusion_common::{ColumnStatistics, ScalarValue, Statistics}; + use datafusion_datasource::PartitionedFile; + use pushdown_sort_helpers::*; + use reorder_files_helpers::*; + + fn file_with_two_mins( + name: &str, + min_a: i32, + min_b: Option, + ) -> PartitionedFile { + let mut pf = PartitionedFile::new(name.to_string(), 0); + let col = |min: Option| ColumnStatistics { + null_count: Precision::Absent, + max_value: Precision::Absent, + min_value: min + .map(|v| Precision::Exact(ScalarValue::Int32(Some(v)))) + .unwrap_or(Precision::Absent), + sum_value: Precision::Absent, + distinct_count: Precision::Absent, + byte_size: Precision::Absent, + }; + pf.statistics = Some(Arc::new(Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: vec![col(Some(min_a)), col(min_b)], + })); + pf + } + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ])); + let mut source = ParquetSource::new(Arc::clone(&schema)); + source.sort_order_for_reorder = Some( + LexOrdering::new(vec![ + sort_expr_on(&schema, "a", false), + sort_expr_on(&schema, "b", false), + ]) + .unwrap(), + ); + + let reordered = source.reorder_files(vec![ + file_with_two_mins("tie_late", 1, Some(300)), + file_with_two_mins("first", 0, Some(999)), + file_with_two_mins("tie_early", 1, Some(100)), + file_with_two_mins("tie_no_b_stats", 1, None), + ]); + + // `first` wins on the leading key; the `a = 1` ties order by + // `min(b)` ASC with missing-`b`-stats last. + assert_eq!( + names(&reordered), + vec!["first", "tie_early", "tie_late", "tie_no_b_stats"] + ); + } + /// When no sort pushdown has fired (`sort_order_for_reorder` is /// `None`), `reorder_files` is a no-op and preserves input order. #[test] @@ -1602,4 +2032,81 @@ mod tests { ); } } + + #[test] + fn test_try_pushdown_filters_rejects_virtual_column_refs() { + // Virtual columns are produced by the reader and cannot be referenced + // inside a RowFilter. `try_pushdown_filters` must report such filters + // as `PushedDown::No` so the FilterExec above the scan stays in + // place — otherwise the scan would silently drop the predicate and + // produce wrong results. + use arrow::datatypes::{DataType, Field, FieldRef, Schema}; + use datafusion_common::config::ConfigOptions; + use datafusion_datasource::TableSchema; + use datafusion_expr::{col, lit as logical_lit}; + use datafusion_functions::core::expr_fn::file_row_index; + use datafusion_physical_expr::planner::logical2physical; + use datafusion_physical_expr_adapter::rewrite::rewrite_file_row_index_expr; + use datafusion_physical_plan::filter_pushdown::PushedDown; + use parquet::arrow::RowNumber; + + let file_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let row_number_field: FieldRef = Arc::new( + Field::new("row_number", DataType::Int64, false) + .with_extension_type(RowNumber), + ); + let table_schema = TableSchema::builder(file_schema) + .with_virtual_columns(vec![row_number_field]) + .build(); + + let source = ParquetSource::new(table_schema).with_pushdown_filters(true); + + let full_schema = source.table_schema.table_schema(); + + let pushable = logical2physical(&col("value").eq(logical_lit(1i64)), full_schema); + let virtual_only = + logical2physical(&col("row_number").eq(logical_lit(2i64)), full_schema); + let mixed = logical2physical( + &col("row_number") + .eq(logical_lit(2i64)) + .or(col("value").eq(logical_lit(4i64))), + full_schema, + ); + let (_, row_index_col) = table_schema_with_row_index_col(source.table_schema()); + let row_index = rewrite_file_row_index_expr( + logical2physical(&file_row_index().gt(logical_lit(2i64)), full_schema), + row_index_col.name(), + row_index_col.index(), + ) + .expect("file_row_index should rewrite to the row_number virtual column"); + + let config = ConfigOptions::default(); + let prop = source + .try_pushdown_filters(vec![pushable, virtual_only, mixed, row_index], &config) + .expect("try_pushdown_filters must not error"); + + assert_eq!(prop.filters.len(), 4); + assert!( + matches!(prop.filters[0], PushedDown::Yes), + "file-column filter should be pushable" + ); + assert!( + matches!(prop.filters[1], PushedDown::No), + "filter referencing only a virtual column must not be pushed down" + ); + assert!( + matches!(prop.filters[2], PushedDown::No), + "filter mixing a virtual column with a file column must not be \ + pushed down (row filter would silently drop it)" + ); + assert!( + matches!(prop.filters[3], PushedDown::No), + "file_row_index() rewrites to a virtual column and must not be \ + pushed down" + ); + } } diff --git a/datafusion/datasource-parquet/src/virtual_column.rs b/datafusion/datasource-parquet/src/virtual_column.rs new file mode 100644 index 0000000000000..2290ad2aeab9d --- /dev/null +++ b/datafusion/datasource-parquet/src/virtual_column.rs @@ -0,0 +1,125 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Typed wrapper for parquet virtual columns. +//! +//! arrow-rs identifies virtual columns via arrow extension types carried on +//! the `FieldRef`. [`ParquetVirtualColumn`] lifts that contract into the type +//! system so callers validate at the boundary (via `TryFrom<&FieldRef>`) +//! rather than string-comparing extension-type names deep inside the reader. + +use arrow::datatypes::FieldRef; +use arrow_schema::extension::ExtensionType; +use datafusion_common::{DataFusionError, Result, not_impl_err}; +use parquet::arrow::RowNumber; +use std::sync::Arc; + +/// A parquet virtual column validated to have a supported arrow extension +/// type. +/// +/// Construct via [`TryFrom<&FieldRef>`]; add a new variant (and update the +/// `TryFrom` impl) when DataFusion gains support for another arrow-rs virtual +/// extension type. +#[derive(Debug, Clone)] +pub enum ParquetVirtualColumn { + /// Absolute row number within the parquet file. Backed by arrow-rs's + /// [`RowNumber`] extension type. + RowNumber(FieldRef), +} + +impl ParquetVirtualColumn { + pub fn field(&self) -> &FieldRef { + match self { + Self::RowNumber(field) => field, + } + } +} + +impl From for FieldRef { + fn from(col: ParquetVirtualColumn) -> Self { + match col { + ParquetVirtualColumn::RowNumber(field) => field, + } + } +} + +impl TryFrom<&FieldRef> for ParquetVirtualColumn { + type Error = DataFusionError; + + fn try_from(field: &FieldRef) -> Result { + let Some(name) = field.extension_type_name() else { + return not_impl_err!( + "Virtual column '{}' is missing an Arrow extension type; \ + supported extension types: [{}]", + field.name(), + RowNumber::NAME + ); + }; + match name { + n if n == RowNumber::NAME => Ok(Self::RowNumber(Arc::clone(field))), + other => not_impl_err!( + "Virtual column '{}' uses unsupported Arrow extension type '{}'; \ + supported types: [{}]. Add a ParquetVirtualColumn variant and \ + a test for this type before wiring it through.", + field.name(), + other, + RowNumber::NAME + ), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::datatypes::{DataType, Field}; + + #[test] + fn row_number_field_converts() { + let field: FieldRef = Arc::new( + Field::new("row_number", DataType::Int64, false) + .with_extension_type(RowNumber), + ); + let col = ParquetVirtualColumn::try_from(&field).expect("valid row_number"); + assert!(matches!(col, ParquetVirtualColumn::RowNumber(_))); + assert_eq!(col.field().name(), "row_number"); + } + + #[test] + fn missing_extension_type_rejected() { + let field: FieldRef = Arc::new(Field::new("plain", DataType::Int64, false)); + let err = ParquetVirtualColumn::try_from(&field).unwrap_err(); + assert!( + err.to_string().contains("missing an Arrow extension type"), + "got: {err}" + ); + } + + #[test] + fn unsupported_extension_type_rejected() { + // RowGroupIndex is a real arrow-rs virtual type not yet in our enum. + let field: FieldRef = Arc::new( + Field::new("row_group_index", DataType::Int64, false) + .with_extension_type(parquet::arrow::RowGroupIndex), + ); + let err = ParquetVirtualColumn::try_from(&field).unwrap_err(); + assert!( + err.to_string().contains("parquet.virtual.row_group_index"), + "error should name the offending extension type, got: {err}" + ); + } +} diff --git a/datafusion/datasource/Cargo.toml b/datafusion/datasource/Cargo.toml index 40e2271f45205..f09447b694f52 100644 --- a/datafusion/datasource/Cargo.toml +++ b/datafusion/datasource/Cargo.toml @@ -34,6 +34,14 @@ all-features = true backtrace = ["datafusion-common/backtrace"] compression = ["async-compression", "liblzma", "bzip2", "flate2", "zstd", "tokio-util"] default = ["compression"] +# Enables protobuf conversions for datasource types, source serialization hooks, +# and the shared `FileScanConfig` <-> proto conversion. Off by default so +# consumers that never serialize plans pay nothing. Mirrors the `proto` feature +# on `datafusion-physical-plan`. +proto = [ + "dep:datafusion-proto-models", + "datafusion-physical-plan/proto", +] [dependencies] arrow = { workspace = true } @@ -56,6 +64,7 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-adapter = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } datafusion-session = { workspace = true } flate2 = { workspace = true, optional = true } futures = { workspace = true } @@ -74,6 +83,7 @@ zstd = { workspace = true, optional = true } [dev-dependencies] criterion = { workspace = true } +datafusion-functions = { workspace = true } insta = { workspace = true } tempfile = { workspace = true } diff --git a/datafusion/datasource-json/src/boundary_stream.rs b/datafusion/datasource/src/boundary_stream.rs similarity index 99% rename from datafusion/datasource-json/src/boundary_stream.rs rename to datafusion/datasource/src/boundary_stream.rs index 847c80279a53e..7b1cfb814df31 100644 --- a/datafusion/datasource-json/src/boundary_stream.rs +++ b/datafusion/datasource/src/boundary_stream.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! Streaming boundary-aligned wrapper for newline-delimited JSON range reads. +//! Streaming boundary-aligned wrapper for newline-delimited JSON and CSV range reads. //! //! [`AlignedBoundaryStream`] wraps a raw byte stream and lazily aligns to //! record (newline) boundaries, avoiding the need for separate `get_opts` @@ -398,7 +398,7 @@ impl Stream for AlignedBoundaryStream { #[cfg(test)] mod tests { use super::*; - use crate::test_utils::{CHUNK_SIZES, make_chunked_store}; + use crate::test_util::{CHUNK_SIZES, make_chunked_store}; use futures::TryStreamExt; async fn collect_stream(stream: AlignedBoundaryStream) -> Vec { diff --git a/datafusion/datasource/src/file.rs b/datafusion/datasource/src/file.rs index 32bee63b54f23..f1a94f2e12363 100644 --- a/datafusion/datasource/src/file.rs +++ b/datafusion/datasource/src/file.rs @@ -359,20 +359,41 @@ pub trait FileSource: Any + Send + Sync { /// - Filter predicates (which may contain dynamic filters) /// - Projection expressions /// - /// The function `f` is called once for each expression. The function should - /// return `TreeNodeRecursion::Continue` to continue visiting other expressions, - /// or `TreeNodeRecursion::Stop` to stop visiting expressions early. + /// The function `f` should be called once per expression unless the function returns + /// [`TreeNodeRecursion::Stop`] to stop iteration. /// - /// Implementations must explicitly visit all expressions. There is no default - /// implementation to ensure that all FileSource implementations handle this correctly. - /// - /// See [`ExecutionPlan::apply_expressions`] for more details and examples. + /// See [`ExecutionPlan::apply_expressions`] for more details and implementation examples. /// /// [`ExecutionPlan::apply_expressions`]: datafusion_physical_plan::ExecutionPlan::apply_expressions fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result; + + /// Serialize this file source into a full [`PhysicalPlanNode`] (a + /// `DataSourceExec` wrapping the `FileScanConfig`), if it knows how. + /// + /// `base` is the shared [`FileScanConfig`] this source is wrapped in; the + /// format-agnostic parts (file groups, schema, statistics, ordering, + /// projection, …) are encoded via + /// [`FileScanConfig::try_to_proto`](crate::file_scan_config::FileScanConfig::try_to_proto), + /// and the concrete source appends its format-specific fields (e.g. CSV + /// delimiter/quote) around it. + /// + /// * `Ok(None)` (the default) — this source has no proto hook yet; the + /// caller falls back to the central downcast chain in `datafusion-proto`. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`FileScanConfig`]: crate::file_scan_config::FileScanConfig + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _base: &FileScanConfig, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } impl dyn FileSource { diff --git a/datafusion/datasource/src/file_scan_config/mod.rs b/datafusion/datasource/src/file_scan_config/mod.rs index 04b74528d5ac1..91dcd5b76fc46 100644 --- a/datafusion/datasource/src/file_scan_config/mod.rs +++ b/datafusion/datasource/src/file_scan_config/mod.rs @@ -20,6 +20,13 @@ pub(crate) mod sort_pushdown; +/// Shared `FileScanConfig` <-> proto conversion, gated on the `proto` feature. +/// Attaches inherent `try_to_proto` / `try_from_proto` / +/// `parse_table_schema_from_proto` helpers to [`FileScanConfig`] used by every +/// file source's `try_to_proto` hook. +#[cfg(feature = "proto")] +mod proto; + use crate::file_groups::FileGroup; use crate::{ PartitionedFile, display::FileGroupsDisplay, file::FileSource, @@ -27,7 +34,7 @@ use crate::{ file_stream::work_source::SharedWorkSource, source::DataSource, statistics::MinMaxStatistics, }; -use arrow::datatypes::FieldRef; +use arrow::datatypes::Fields; use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::TreeNodeRecursion; @@ -40,12 +47,13 @@ use datafusion_execution::{ use datafusion_expr::Operator; use crate::source::OpenArgs; +use datafusion_common::stats::Precision; use datafusion_physical_expr::expressions::{BinaryExpr, Column}; -use datafusion_physical_expr::projection::ProjectionExprs; +use datafusion_physical_expr::projection::{ProjectionExprs, ProjectionMapping}; use datafusion_physical_expr::utils::reassign_expr_columns; use datafusion_physical_expr::{EquivalenceProperties, Partitioning, split_conjunction}; use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory; -use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, is_volatile}; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_plan::SortOrderPushdownResult; use datafusion_physical_plan::coop::cooperative; @@ -114,7 +122,7 @@ use std::{fmt::Debug, fmt::Formatter, fmt::Result as FmtResult, sync::Arc}; /// # fn file_type(&self) -> &str { "parquet" } /// # // Note that this implementation drops the projection on the floor, it is not complete! /// # fn try_pushdown_projection(&self, projection: &ProjectionExprs) -> Result>> { Ok(Some(Arc::new(self.clone()) as Arc)) } -/// # fn apply_expressions(&self, _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } +/// # fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } /// # } /// # impl ParquetSource { /// # fn new(table_schema: impl Into) -> Self { Self {table_schema: table_schema.into()} } @@ -163,6 +171,12 @@ pub struct FileScanConfig { /// DataFusion may attempt to read each partition of files /// concurrently, however files *within* a partition will be read /// sequentially, one after the next. + /// + /// Note that when `datafusion.execution.enable_file_stream_work_stealing` + /// is enabled (the default), files may be reassigned to a different + /// partition at runtime unless `preserve_order` or + /// `partitioned_by_file_group` is set, so a file is not guaranteed to be + /// read by the partition it is grouped under here. pub file_groups: Vec, /// Table constraints pub constraints: Constraints, @@ -202,14 +216,13 @@ pub struct FileScanConfig { /// would be incorrect if there are filters being applied, thus this should be accessed /// via [`FileScanConfig::statistics`]. pub(crate) statistics: Statistics, - /// When true, file_groups are organized by partition column values - /// and output_partitioning will return Hash partitioning on partition columns. - /// This allows the optimizer to skip hash repartitioning for aggregates and joins - /// on partition columns. + /// Declared physical output partitioning for this scan. /// - /// If the number of file partitions > target_partitions, the file partitions will be grouped - /// in a round-robin fashion such that number of file partitions = target_partitions. - pub partitioned_by_file_group: bool, + /// Expressions are against the full table schema, before scan projection or + /// filtering. `ListingTable` validates partition count before building the + /// scan, and direct builders with mismatched counts fall back to + /// `UnknownPartitioning`. + pub output_partitioning: Option, } /// A builder for [`FileScanConfig`]'s. @@ -242,7 +255,9 @@ pub struct FileScanConfig { /// ]; /// /// // Create table schema with file schema and partition columns -/// let table_schema = TableSchema::new(file_schema, partition_cols); +/// let table_schema = TableSchema::builder(file_schema) +/// .with_table_partition_cols(partition_cols) +/// .build(); /// /// // Create a builder for scanning Parquet files from a local filesystem /// let config = FileScanConfigBuilder::new( @@ -276,10 +291,10 @@ pub struct FileScanConfigBuilder { file_groups: Vec, statistics: Option, output_ordering: Vec, + output_partitioning: Option, file_compression_type: Option, batch_size: Option, expr_adapter_factory: Option>, - partitioned_by_file_group: bool, } impl FileScanConfigBuilder { @@ -299,13 +314,13 @@ impl FileScanConfigBuilder { file_groups: vec![], statistics: None, output_ordering: vec![], + output_partitioning: None, file_compression_type: None, limit: None, preserve_order: false, constraints: None, batch_size: None, expr_adapter_factory: None, - partitioned_by_file_group: false, } } @@ -465,6 +480,15 @@ impl FileScanConfigBuilder { self } + /// Set declared physical output partitioning for this scan. + pub fn with_output_partitioning( + mut self, + output_partitioning: Option, + ) -> Self { + self.output_partitioning = output_partitioning; + self + } + /// Set the file compression type pub fn with_file_compression_type( mut self, @@ -494,18 +518,6 @@ impl FileScanConfigBuilder { self } - /// Set whether file groups are organized by partition column values. - /// - /// When set to true, the output partitioning will be declared as Hash partitioning - /// on the partition columns. - pub fn with_partitioned_by_file_group( - mut self, - partitioned_by_file_group: bool, - ) -> Self { - self.partitioned_by_file_group = partitioned_by_file_group; - self - } - /// Build the final [`FileScanConfig`] with all the configured settings. /// /// This method takes ownership of the builder and returns the constructed `FileScanConfig`. @@ -523,10 +535,10 @@ impl FileScanConfigBuilder { file_groups, statistics, output_ordering, + output_partitioning, file_compression_type, batch_size, expr_adapter_factory: expr_adapter, - partitioned_by_file_group, } = self; let constraints = constraints.unwrap_or_default(); @@ -551,7 +563,7 @@ impl FileScanConfigBuilder { batch_size, expr_adapter_factory: expr_adapter, statistics, - partitioned_by_file_group, + output_partitioning, } } } @@ -564,17 +576,118 @@ impl From for FileScanConfigBuilder { file_groups: config.file_groups, statistics: Some(config.statistics), output_ordering: config.output_ordering, + output_partitioning: config.output_partitioning, file_compression_type: Some(config.file_compression_type), limit: config.limit, preserve_order: config.preserve_order, constraints: Some(config.constraints), batch_size: config.batch_size, expr_adapter_factory: config.expr_adapter_factory, - partitioned_by_file_group: config.partitioned_by_file_group, } } } +/// Builds output partitioning over `partition_cols` (resolved to their indices in +/// `schema`) with `partition_count` partitions. Returns `None` when there are no +/// partition columns. Callers use this to declare the output partitioning of a scan +/// whose file groups are organized by partition column values. +pub fn output_partitioning_from_partition_fields( + schema: &Schema, + partition_cols: &Fields, + partition_count: usize, +) -> Option { + if partition_cols.is_empty() { + return None; + } + + let mut exprs: Vec> = Vec::with_capacity(partition_cols.len()); + for partition_col in partition_cols { + let name = partition_col.name(); + let idx = schema + .fields() + .iter() + .position(|field| field.name() == name)?; + exprs.push(Arc::new(Column::new(name, idx))); + } + + Some(Partitioning::Hash(exprs, partition_count)) +} + +fn project_output_partitioning( + partitioning: &Partitioning, + mapping: &ProjectionMapping, + input_schema: &SchemaRef, + partition_count: usize, +) -> Partitioning { + let input_eq_properties = EquivalenceProperties::new(Arc::clone(input_schema)); + match partitioning { + Partitioning::Hash(exprs, _) => { + let projected_exprs = input_eq_properties + .project_expressions(exprs, mapping) + .collect::>>(); + projected_exprs + .map(|exprs| Partitioning::Hash(exprs, partition_count)) + .unwrap_or_else(|| Partitioning::UnknownPartitioning(partition_count)) + } + Partitioning::Range(_) + | Partitioning::RoundRobinBatch(_) + | Partitioning::UnknownPartitioning(_) => { + partitioning.project(mapping, &input_eq_properties) + } + } +} + +/// Returns `true` if merging `outer` into `inner` would duplicate a volatile or +/// non-trivial expression that CSE deduplicated; the caller should then decline +/// the merge. +/// +/// Merging substitutes each `inner` expression into every `outer` reference to +/// it. Since the logical optimizer extracts a repeated expression into a single +/// `inner` entry referenced by column, re-inlining it at more than one +/// reference site undoes that deduplication. An `inner` expression referenced +/// more than once is therefore blocked when it is either: +/// +/// - **volatile** (e.g. `random()`) — evaluating it independently at each site +/// makes references that should share one "locked-in" value diverge (the +/// correctness guard the physical `ProjectionPushdown` and `FilterPushdown` +/// rules also apply via +/// `datafusion_physical_expr_common::physical_expr::is_volatile`); or +/// - **not cheap to recompute** — its placement is not push-to-leaves +/// (`KeepInPlace`: arithmetic, casts, most scalar functions). Leaf-pushable +/// expressions (columns, `get_field`, `input_file_name`) still merge. This +/// matches `try_collapse_projection_chain`. +/// +/// References are counted with multiplicity, so `r + r` counts as two; an +/// expression referenced exactly once has nothing to duplicate. +fn would_duplicate_costly_exprs( + inner: &ProjectionExprs, + outer: &ProjectionExprs, +) -> bool { + use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + + let inner_exprs = inner.as_ref(); + + let mut ref_counts = vec![0usize; inner_exprs.len()]; + for proj_expr in outer.as_ref() { + proj_expr + .expr + .apply(|e| { + if let Some(col) = e.as_ref().downcast_ref::() + && let Some(count) = ref_counts.get_mut(col.index()) + { + *count += 1; + } + Ok(TreeNodeRecursion::Continue) + }) + .expect("infallible closure should not fail"); + } + + ref_counts.iter().enumerate().any(|(idx, &count)| { + let expr = &inner_exprs[idx].expr; + count > 1 && (is_volatile(expr) || !expr.placement().should_push_to_leaves()) + }) +} + impl DataSource for FileScanConfig { fn open( &self, @@ -662,6 +775,10 @@ impl DataSource for FileScanConfig { display_orderings(f, &orderings)?; + if self.output_partitioning.is_some() { + write!(f, ", output_partitioning={}", self.output_partitioning())?; + } + if !self.constraints.is_empty() { write!(f, ", {}", self.constraints)?; } @@ -685,10 +802,9 @@ impl DataSource for FileScanConfig { repartition_file_min_size: usize, output_ordering: Option, ) -> Result>> { - // When files are grouped by partition values, we cannot allow byte-range - // splitting. It would mix rows from different partition values across - // file groups, breaking the Hash partitioning. - if self.partitioned_by_file_group { + // When file groups define output partitioning, repartitioning files + // would invalidate the partition-to-file-group mapping. + if self.output_partitioning.is_some() { return Ok(None); } @@ -704,13 +820,16 @@ impl DataSource for FileScanConfig { /// Returns the output partitioning for this file scan. /// - /// When `partitioned_by_file_group` is true, this returns `Partitioning::Hash` on - /// the Hive partition columns, allowing the optimizer to skip hash repartitioning - /// for aggregates and joins on those columns. + /// When `output_partitioning` is set, this returns the declared partitioning + /// after applying scan projection, allowing the optimizer to skip hash + /// repartitioning for aggregates and joins on the partitioning columns. + /// + /// If projection or partition count validation fails, this returns + /// `UnknownPartitioning`. /// /// Tradeoffs - /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries with - /// `GROUP BY` or `ORDER BY` on partition columns. + /// - Benefit: Eliminates `RepartitionExec` and `SortExec` for queries whose + /// required distribution is satisfied by the scan's output partitioning. /// - Cost: Files are grouped by partition values rather than split by byte /// ranges, which may reduce I/O parallelism when partition sizes are uneven. /// For simple aggregations without `ORDER BY`, this cost may outweigh the benefit. @@ -719,39 +838,37 @@ impl DataSource for FileScanConfig { /// - Idea: Could allow byte-range splitting within partition-aware groups, /// preserving I/O parallelism while maintaining partition semantics. fn output_partitioning(&self) -> Partitioning { - if self.partitioned_by_file_group { - let partition_cols = self.table_partition_cols(); - if !partition_cols.is_empty() { - let projected_schema = match self.projected_schema() { - Ok(schema) => schema, - Err(_) => { - debug!( - "Could not get projected schema, falling back to UnknownPartitioning." - ); - return Partitioning::UnknownPartitioning(self.file_groups.len()); - } - }; - - // Build Column expressions for partition columns based on their - // position in the projected schema - let mut exprs: Vec> = Vec::new(); - for partition_col in partition_cols { - if let Some((idx, _)) = projected_schema - .fields() - .iter() - .enumerate() - .find(|(_, f)| f.name() == partition_col.name()) - { - exprs.push(Arc::new(Column::new(partition_col.name(), idx))); - } - } + let Some(output_partitioning) = self.output_partitioning.clone() else { + return Partitioning::UnknownPartitioning(self.file_groups.len()); + }; + if output_partitioning.partition_count() != self.file_groups.len() { + warn!( + "Declared output partitioning has {} partitions, but file scan has {} file groups. Falling back to UnknownPartitioning.", + output_partitioning.partition_count(), + self.file_groups.len() + ); + return Partitioning::UnknownPartitioning(self.file_groups.len()); + } - if exprs.len() == partition_cols.len() { - return Partitioning::Hash(exprs, self.file_groups.len()); + if let Some(projection) = self.file_source.projection() { + let schema = self.file_source.table_schema().table_schema(); + return match projection.projection_mapping(schema) { + Ok(mapping) => project_output_partitioning( + &output_partitioning, + &mapping, + schema, + self.file_groups.len(), + ), + Err(e) => { + debug!( + "Could not project output partitioning, falling back to UnknownPartitioning: {e}" + ); + Partitioning::UnknownPartitioning(self.file_groups.len()) } - } + }; } - Partitioning::UnknownPartitioning(self.file_groups.len()) + + output_partitioning } /// Computes the effective equivalence properties of this file scan, taking @@ -857,6 +974,17 @@ impl DataSource for FileScanConfig { &self, projection: &ProjectionExprs, ) -> Result>> { + // Don't merge a projection into the scan if it would inline a volatile + // or expensive expression referenced more than once. For a volatile + // expression (e.g. `random()` aliased in a subquery) this would turn a + // single "locked-in" value into multiple independent evaluations (see + // #23220); for an expensive scalar function it would undo CSE and + // re-evaluate the expression at every reference site. + if let Some(inner) = self.file_source.projection() + && would_duplicate_costly_exprs(inner, projection) + { + return Ok(None); + } match self.file_source.try_pushdown_projection(projection)? { Some(new_source) => { let mut new_file_scan_config = self.clone(); @@ -941,14 +1069,19 @@ impl DataSource for FileScanConfig { /// │ → SortExec removed, fetch (LIMIT) pushed to DataSourceExec /// │ /// ├─► FileSource returns Inexact - /// │ (reverse_row_groups=true) - /// │ → SortExec kept, scan optimized + /// │ (e.g. column_in_file_schema: opener will reorder RGs at runtime) + /// │ → rebuild_with_source: sort files by stats; if the post-sort + /// │ file groups are non-overlapping AND the request now validates + /// │ AND no NULLs sit in the sort columns of non-last files, + /// │ upgrade back to Exact (SortExec removed). Otherwise stays + /// │ Inexact and SortExec is kept while the scan is still + /// │ optimised via `sort_order_for_reorder` / `reverse_row_groups`. /// │ /// └─► FileSource returns Unsupported - /// (ordering stripped because files in wrong order) + /// (e.g. expression sort key or partition column) /// → try_sort_file_groups_by_statistics(): /// 1. Sort files within each group by min/max statistics - /// 2. Re-check: non-overlapping + ordering valid? + /// 2. Re-check: non-overlapping + ordering valid + no NULLs? /// YES → Exact → SortExec removed /// NO → Inexact (files reordered, Sort stays) /// ``` @@ -977,8 +1110,42 @@ impl DataSource for FileScanConfig { } } SortOrderPushdownResult::Inexact { inner } => { - Ok(SortOrderPushdownResult::Inexact { - inner: Arc::new(self.rebuild_with_source(inner, false, order)?), + let mut config = self.rebuild_with_source(inner, false, order)?; + // `rebuild_with_source` reorders files by stats; if the + // post-sort files are non-overlapping AND the request now + // validates against the new file groups, `output_ordering` + // is preserved and we can upgrade back to Exact. This + // restores the sort-elimination behaviour that lived in + // the `Unsupported` → `try_sort_file_groups_by_statistics` + // path before #21956 routed `column_in_file_schema` cases + // here. + if config.output_ordering.is_empty() { + return Ok(SortOrderPushdownResult::Inexact { + inner: Arc::new(config), + }); + } + // Upgrading to Exact: the post-sort file groups are + // non-overlapping and each file's declared ordering + // re-validates, so reading the files in their natural + // (declared-sorted) order already yields the requested + // ordering — exactly like the `Unsupported` → Exact path, + // which reads files in natural order too. + // + // Drop the runtime row-group reorder hints the Inexact + // source carried (`sort_order_for_reorder` / + // `reverse_row_groups`) by restoring the original, + // hint-free source. With the `SortExec` removed those + // hints are not just redundant but unsafe: for a DESC + // request the opener sorts row groups ASC-by-min and then + // reverses them, which reorders two row groups within a + // single file that share the same `min` incorrectly + // (e.g. a file `[10,8,8,8]` whose row groups are + // `[10,8]` and `[8,8]` would stream as `8,8,10,8`). + // The `SortExec` used to mask this; once it is gone the + // reordered stream is the final, wrong answer. + config.file_source = Arc::clone(&self.file_source); + Ok(SortOrderPushdownResult::Exact { + inner: Arc::new(config), }) } SortOrderPushdownResult::Unsupported => { @@ -1001,7 +1168,7 @@ impl DataSource for FileScanConfig { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // Delegate to the file source self.file_source.apply_expressions(f) @@ -1011,15 +1178,35 @@ impl DataSource for FileScanConfig { /// during one execution. /// /// This returns `None` when sibling streams must not share work, such as - /// when file order must be preserved or the file groups define the output - /// partitioning needed for the rest of the plan - fn create_sibling_state(&self) -> Option> { - if self.preserve_order || self.partitioned_by_file_group { + /// when file order must be preserved, the file groups define the output + /// partitioning needed for the rest of the plan, or work stealing is + /// disabled via + /// `datafusion.execution.enable_file_stream_work_stealing`. + fn create_sibling_state( + &self, + config: &ConfigOptions, + ) -> Option> { + if self.preserve_order + || self.output_partitioning.is_some() + || !config.execution.enable_file_stream_work_stealing + { return None; } Some(Arc::new(SharedWorkSource::from_config(self)) as Arc) } + + /// Serialize this file scan by delegating to the concrete + /// [`FileSource`]'s + /// [`try_to_proto`](crate::file::FileSource::try_to_proto) hook, passing + /// `self` as the shared spine it needs to emit the base config. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + self.file_source().try_to_proto(self, ctx) + } } impl FileScanConfig { @@ -1068,7 +1255,7 @@ impl FileScanConfig { } /// Get the table partition columns - pub fn table_partition_cols(&self) -> &Vec { + pub fn table_partition_cols(&self) -> &Fields { self.file_source.table_schema().table_partition_cols() } @@ -1078,7 +1265,9 @@ impl FileScanConfig { /// we can't guarantee the statistics are exact because we don't know how many /// rows will be filtered out. pub fn statistics(&self) -> Statistics { - if self.file_source.filter().is_some() { + let filter_may_change_row_count = self.file_source.filter().is_some() + && self.statistics.num_rows != Precision::Exact(0); + if filter_may_change_row_count { self.statistics.clone().to_inexact() } else { self.statistics.clone() @@ -1396,9 +1585,9 @@ mod tests { use std::collections::HashMap; use super::*; - use crate::TableSchema; use crate::source::DataSourceExec; use crate::test_util::col; + use crate::{TableSchema, TableSchemaBuilder}; use crate::{ generate_test_files, test_util::MockSource, tests::aggr_test_schema, verify_sort_integrity, @@ -1413,12 +1602,19 @@ mod tests { use datafusion_execution::TaskContext; use datafusion_expr::SortExpr; use datafusion_physical_expr::PhysicalExpr; + + #[cfg(feature = "proto")] + use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; use datafusion_physical_expr::create_physical_sort_expr; use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::projection::ProjectionExpr; use datafusion_physical_expr::projection::ProjectionExprs; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::execution_plan::collect; + #[cfg(feature = "proto")] + use datafusion_physical_plan::proto::{ExecutionPlanEncode, ExecutionPlanEncodeCtx}; + #[cfg(feature = "proto")] + use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; use futures::FutureExt as _; use futures::StreamExt as _; use futures::stream; @@ -1478,10 +1674,119 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + } + + #[cfg(feature = "proto")] + #[derive(Clone)] + struct ProtoHookSource { + metrics: ExecutionPlanMetricsSet, + table_schema: TableSchema, + } + + #[cfg(feature = "proto")] + impl ProtoHookSource { + fn new(table_schema: TableSchema) -> Self { + Self { + metrics: ExecutionPlanMetricsSet::new(), + table_schema, + } + } + } + + #[cfg(feature = "proto")] + impl FileSource for ProtoHookSource { + fn create_file_opener( + &self, + _object_store: Arc, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result> { + internal_err!("not needed for proto delegation test") + } + + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn with_batch_size(&self, _batch_size: usize) -> Arc { + Arc::new(self.clone()) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "proto-hook-test" + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + + fn try_to_proto( + &self, + _base: &FileScanConfig, + _ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(Some(PhysicalPlanNode::default())) + } + } + + #[cfg(feature = "proto")] + struct UnusedPlanEncoder; + + #[cfg(feature = "proto")] + impl ExecutionPlanEncode for UnusedPlanEncoder { + fn encode_plan( + &self, + _plan: &Arc, + ) -> Result { + internal_err!("not needed for proto delegation test") + } + + fn encode_expr(&self, _expr: &Arc) -> Result { + internal_err!("not needed for proto delegation test") + } + + fn encode_udf(&self, _udf: &ScalarUDF) -> Result>> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udaf(&self, _udaf: &AggregateUDF) -> Result>> { + internal_err!("not needed for proto delegation test") + } + + fn encode_udwf(&self, _udwf: &WindowUDF) -> Result>> { + internal_err!("not needed for proto delegation test") + } + } + + #[cfg(feature = "proto")] + #[test] + fn data_source_exec_delegates_proto_to_file_source() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + false, + )])); + let source = Arc::new(ProtoHookSource::new(TableSchema::from(&schema))); + let config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .build(); + let exec = DataSourceExec::from_data_source(config); + let encoder = UnusedPlanEncoder; + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert_eq!(exec.try_to_proto(&ctx)?, Some(PhysicalPlanNode::default())); + Ok(()) } #[test] @@ -1518,6 +1823,7 @@ mod tests { use chrono::TimeZone; use datafusion_common::DFSchema; use datafusion_expr::execution_props::ExecutionProps; + use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use object_store::{ObjectMeta, path::Path}; struct File { @@ -1731,6 +2037,7 @@ mod tests { &expr, &DFSchema::try_from(Arc::clone(&table_schema))?, &ExecutionProps::default(), + &PhysicalPlanningContext::default(), ) }) .collect::>>()?, @@ -1832,10 +2139,14 @@ mod tests { statistics: Statistics, table_partition_cols: Vec, ) -> FileScanConfig { - let table_schema = TableSchema::new( - file_schema, - table_partition_cols.into_iter().map(Arc::new).collect(), - ); + let table_schema = TableSchema::builder(file_schema) + .with_table_partition_cols( + table_partition_cols + .into_iter() + .map(Arc::new) + .collect::(), + ) + .build(); FileScanConfigBuilder::new( ObjectStoreUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema.clone())), @@ -1851,14 +2162,13 @@ mod tests { let file_schema = aggr_test_schema(); let object_store_url = ObjectStoreUrl::parse("test:///").unwrap(); - let table_schema = TableSchema::new( - Arc::clone(&file_schema), - vec![Arc::new(Field::new( + let table_schema = TableSchemaBuilder::from(&file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( "date", wrap_partition_type_in_dict(DataType::Utf8), false, - ))], - ); + ))]) + .build(); let file_source: Arc = Arc::new(MockSource::new(table_schema.clone())); @@ -1920,7 +2230,7 @@ mod tests { let file_schema = aggr_test_schema(); let object_store_url = ObjectStoreUrl::parse("test:///").unwrap(); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); // Create a file source with a filter let file_source: Arc = Arc::new( @@ -1973,7 +2283,7 @@ mod tests { let file_schema = aggr_test_schema(); let object_store_url = ObjectStoreUrl::parse("test:///").unwrap(); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source: Arc = Arc::new(MockSource::new(table_schema.clone())); @@ -2035,10 +2345,14 @@ mod tests { )]; let file = PartitionedFile::new("test_file.parquet", 100); - let table_schema = TableSchema::new( - Arc::clone(&schema), - partition_cols.iter().map(|f| Arc::new(f.clone())).collect(), - ); + let table_schema = TableSchemaBuilder::from(&schema) + .with_table_partition_cols( + partition_cols + .iter() + .map(|f| Arc::new(f.clone())) + .collect::(), + ) + .build(); let file_source: Arc = Arc::new(MockSource::new(table_schema.clone())); @@ -2074,7 +2388,10 @@ mod tests { Some(vec![0, 2]) ); assert_eq!(new_config.limit, Some(10)); - assert_eq!(*new_config.table_partition_cols(), partition_cols); + assert_eq!( + *new_config.table_partition_cols(), + Fields::from(partition_cols) + ); assert_eq!(new_config.file_groups.len(), 1); assert_eq!(new_config.file_groups[0].len(), 1); assert_eq!( @@ -2087,7 +2404,10 @@ mod tests { #[test] fn test_split_groups_by_statistics_with_target_partitions() -> Result<()> { use datafusion_common::DFSchema; - use datafusion_expr::{col, execution_props::ExecutionProps}; + use datafusion_expr::{ + col, execution_props::ExecutionProps, + physical_planning_context::PhysicalPlanningContext, + }; let schema = Arc::new(Schema::new(vec![Field::new( "value", @@ -2101,7 +2421,13 @@ mod tests { let sort_expr = [col("value").sort(true, false)]; let sort_ordering = sort_expr .map(|expr| { - create_physical_sort_expr(&expr, &df_schema, &exec_props).unwrap() + create_physical_sort_expr( + &expr, + &df_schema, + &exec_props, + &PhysicalPlanningContext::default(), + ) + .unwrap() }) .into(); @@ -2246,9 +2572,8 @@ mod tests { // of just the projected ones. use crate::source::DataSourceExec; - use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; - // Create a schema with 4 columns let schema = Arc::new(Schema::new(vec![ Field::new("col0", DataType::Int32, false), Field::new("col1", DataType::Int32, false), @@ -2284,7 +2609,7 @@ mod tests { let file_group = FileGroup::new(vec![PartitionedFile::new("test.parquet", 1024)]) .with_statistics(Arc::new(file_group_stats)); - let table_schema = TableSchema::new(Arc::clone(&schema), vec![]); + let table_schema = TableSchema::from(&schema); // Create a FileScanConfig with projection: only keep columns 0 and 2 let config = FileScanConfigBuilder::new( @@ -2300,7 +2625,12 @@ mod tests { let exec = DataSourceExec::from_data_source(config); // Get statistics for partition 0 - let partition_stats = exec.partition_statistics(Some(0)).unwrap(); + let partition_stats = StatisticsContext::new() + .compute( + exec.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap(); // Verify that only 2 columns are in the statistics (the projected ones) assert_eq!( @@ -2327,6 +2657,45 @@ mod tests { assert_eq!(partition_stats.total_byte_size, Precision::Exact(800)); } + #[test] + fn test_statistics_with_filter() { + assert_num_rows_with_filter(Precision::Absent, Precision::Absent); + assert_num_rows_with_filter(Precision::Exact(100), Precision::Inexact(100)); + assert_num_rows_with_filter(Precision::Inexact(100), Precision::Inexact(100)); + assert_num_rows_with_filter(Precision::Exact(0), Precision::Exact(0)); + + /// Creates a [`FileScanConfig`] with a filter and calls [`FileScanConfig::statistics`]. + /// Then the function checks the output num_rows stats, given the input num_rows stats. + fn assert_num_rows_with_filter( + input_num_rows: Precision, + expected_num_rows: Precision, + ) { + let schema = Arc::new(Schema::new(vec![Field::new( + "col0", + DataType::Int32, + false, + )])); + + let stats = + Statistics::new_unknown(schema.as_ref()).with_num_rows(input_num_rows); + let file_group = + FileGroup::new(vec![PartitionedFile::new("test.parquet", 1024)]); + + let table_schema = TableSchema::from(&schema); + let config = FileScanConfigBuilder::new( + ObjectStoreUrl::parse("test:///").unwrap(), + Arc::new(MockSource::new(table_schema.clone()).with_filter(Arc::new( + Literal::new(ScalarValue::Boolean(Some(true))), + ))), + ) + .with_file_groups(vec![file_group]) + .with_statistics(stats) + .build(); + + assert_eq!(config.statistics().num_rows, expected_num_rows,); + } + } + /// Regression test for reusing a `DataSourceExec` after its execution-local /// shared work queue has been drained. /// @@ -2422,21 +2791,72 @@ mod tests { vec![partition_col], ); - // partitioned_by_file_group defaults to false + // output_partitioning defaults to None let partitioning = config.output_partitioning(); assert!(matches!(partitioning, Partitioning::UnknownPartitioning(_))); } #[test] - fn test_output_partitioning_no_partition_columns() { + fn test_declared_output_partitioning_projects_with_scan() { let file_schema = aggr_test_schema(); + let output_partitioning = + Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4); + + let mut config = config_for_projection( + Arc::clone(&file_schema), + Some(vec![1, 2]), + Statistics::new_unknown(&file_schema), + vec![], + ); + config.file_groups = vec![ + FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]), + ]; + config.output_partitioning = Some(output_partitioning); + + match config.output_partitioning() { + Partitioning::Hash(exprs, num_partitions) => { + assert_eq!(num_partitions, 4); + assert_eq!(exprs.len(), 1); + let column = exprs[0].downcast_ref::().unwrap(); + assert_eq!(column.name(), "c2"); + assert_eq!(column.index(), 0); + } + _ => panic!("Expected Hash partitioning"), + } + let mut config = config_for_projection( + Arc::clone(&file_schema), + Some(vec![2]), + Statistics::new_unknown(&file_schema), + vec![], + ); + config.file_groups = vec![ + FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]), + FileGroup::new(vec![PartitionedFile::new("f4.parquet".to_string(), 1024)]), + ]; + config.output_partitioning = + Some(Partitioning::Hash(vec![Arc::new(Column::new("c2", 1))], 4)); + + assert!(matches!( + config.output_partitioning(), + Partitioning::UnknownPartitioning(4) + )); + } + + #[test] + fn test_output_partitioning_no_partition_columns() { + let file_schema = aggr_test_schema(); + let config = config_for_projection( Arc::clone(&file_schema), None, Statistics::new_unknown(&file_schema), vec![], // No partition columns ); - config.partitioned_by_file_group = true; let partitioning = config.output_partitioning(); assert!(matches!(partitioning, Partitioning::UnknownPartitioning(_))); @@ -2459,12 +2879,16 @@ mod tests { Statistics::new_unknown(&file_schema), single_partition_col, ); - config.partitioned_by_file_group = true; config.file_groups = vec![ FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]), FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]), FileGroup::new(vec![PartitionedFile::new("f3.parquet".to_string(), 1024)]), ]; + config.output_partitioning = output_partitioning_from_partition_fields( + config.file_source.table_schema().table_schema(), + config.table_partition_cols(), + config.file_groups.len(), + ); let partitioning = config.output_partitioning(); match partitioning { @@ -2488,11 +2912,15 @@ mod tests { Statistics::new_unknown(&file_schema), multiple_partition_cols, ); - config.partitioned_by_file_group = true; config.file_groups = vec![ FileGroup::new(vec![PartitionedFile::new("f1.parquet".to_string(), 1024)]), FileGroup::new(vec![PartitionedFile::new("f2.parquet".to_string(), 1024)]), ]; + config.output_partitioning = output_partitioning_from_partition_fields( + config.file_source.table_schema().table_schema(), + config.table_partition_cols(), + config.file_groups.len(), + ); let partitioning = config.output_partitioning(); match partitioning { @@ -2515,7 +2943,7 @@ mod tests { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(InexactSortPushdownSource::new(table_schema)); let file_groups = vec![FileGroup::new(vec![ @@ -2627,7 +3055,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -2637,7 +3065,7 @@ mod tests { fn sort_pushdown_unsupported_source_files_get_sorted() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(MockSource::new(table_schema)); let file_groups = vec![FileGroup::new(vec![ @@ -2671,7 +3099,7 @@ mod tests { fn sort_pushdown_unsupported_source_already_sorted() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(MockSource::new(table_schema)); let file_groups = vec![FileGroup::new(vec![ @@ -2695,7 +3123,7 @@ mod tests { fn sort_pushdown_unsupported_source_descending_sort() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(MockSource::new(table_schema)); let file_groups = vec![FileGroup::new(vec![ @@ -2734,7 +3162,7 @@ mod tests { fn sort_pushdown_exact_source_non_overlapping_returns_exact() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(ExactSortPushdownSource::new(table_schema)); let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); @@ -2768,7 +3196,7 @@ mod tests { fn sort_pushdown_exact_source_overlapping_downgraded_to_inexact() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(ExactSortPushdownSource::new(table_schema)); let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); @@ -2802,7 +3230,7 @@ mod tests { fn sort_pushdown_exact_source_out_of_order_returns_exact() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(ExactSortPushdownSource::new(table_schema)); let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); @@ -2840,7 +3268,7 @@ mod tests { fn sort_pushdown_unsupported_source_single_file_groups() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(MockSource::new(table_schema)); let file_groups = vec![ @@ -2866,7 +3294,7 @@ mod tests { fn sort_pushdown_unsupported_source_multiple_groups() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(MockSource::new(table_schema)); let file_groups = vec![ @@ -2906,7 +3334,7 @@ mod tests { fn sort_pushdown_unsupported_source_partial_statistics() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(MockSource::new(table_schema)); let file_groups = vec![ @@ -2946,7 +3374,7 @@ mod tests { fn sort_pushdown_inexact_source_with_statistics_sorting() -> Result<()> { let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(InexactSortPushdownSource::new(table_schema)); let file_groups = vec![FileGroup::new(vec![ @@ -2983,7 +3411,7 @@ mod tests { // time (all values in group 0 < group 1), degrading to single-threaded I/O. let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(ExactSortPushdownSource::new(table_schema)); let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); @@ -3041,7 +3469,7 @@ mod tests { // sorting (which would undo the reversal). The result is Inexact. let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, false)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(InexactSortPushdownSource::new(table_schema)); let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); @@ -3108,7 +3536,7 @@ mod tests { // Should NOT upgrade to Exact — NULLs would appear in wrong position. let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(MockSource::new(table_schema)); let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); @@ -3141,7 +3569,7 @@ mod tests { // Files are non-overlapping, no NULLs → should upgrade to Exact let file_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)])); - let table_schema = TableSchema::new(Arc::clone(&file_schema), vec![]); + let table_schema = TableSchema::from(&file_schema); let file_source = Arc::new(MockSource::new(table_schema)); let sort_expr = PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0))); @@ -3166,4 +3594,272 @@ mod tests { ); Ok(()) } + + /// Helper: build a `ProjectionExprs` from `(expr, alias)` pairs. + fn make_projection(pairs: Vec<(Arc, &str)>) -> ProjectionExprs { + ProjectionExprs::new( + pairs + .into_iter() + .map(|(expr, alias)| ProjectionExpr::new(expr, alias)), + ) + } + + /// Helper: create a volatile (non-deterministic) function expression, + /// e.g. `random()`. + fn make_volatile_expr() -> Arc { + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarUDF; + use datafusion_functions::math::random::RandomFunc; + use datafusion_physical_expr::ScalarFunctionExpr; + + Arc::new(ScalarFunctionExpr::new( + "random", + Arc::new(ScalarUDF::from(RandomFunc::new())), + vec![], + Arc::new(Field::new("random", DataType::Float64, false)), + Arc::new(ConfigOptions::default()), + )) + } + + /// Helper: create a deterministic but expensive scalar-function + /// expression, e.g. `abs()`. + fn make_udf_expr(args: Vec>) -> Arc { + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarUDF; + use datafusion_functions::math::abs::AbsFunc; + use datafusion_physical_expr::ScalarFunctionExpr; + + Arc::new(ScalarFunctionExpr::new( + "abs", + Arc::new(ScalarUDF::from(AbsFunc::new())), + args, + Arc::new(Field::new("abs", DataType::Int32, false)), + Arc::new(ConfigOptions::default()), + )) + } + + /// Helper: create a cheap, leaf-pushable scalar function — struct field + /// access `get_field(s, 'x')`, whose placement is `MoveTowardsLeafNodes` + /// when the base is a column and the key is a literal. + fn make_leaf_pushable_expr() -> Arc { + use datafusion_common::config::ConfigOptions; + use datafusion_expr::ScalarUDF; + use datafusion_functions::core::getfield::GetFieldFunc; + use datafusion_physical_expr::ScalarFunctionExpr; + use datafusion_physical_expr::expressions::Literal; + + Arc::new(ScalarFunctionExpr::new( + "get_field", + Arc::new(ScalarUDF::from(GetFieldFunc::new())), + vec![ + Arc::new(Column::new("s", 0)), + Arc::new(Literal::new(ScalarValue::Utf8(Some("x".to_string())))), + ], + Arc::new(Field::new("x", DataType::Int32, true)), + Arc::new(ConfigOptions::default()), + )) + } + + /// Column-only inner projections always merge safely, even when + /// the outer projection references them multiple times. + #[test] + fn test_would_duplicate_allows_column_only_inner() { + let col_a: Arc = Arc::new(Column::new("a", 0)); + let col_b: Arc = Arc::new(Column::new("b", 1)); + + let inner = + make_projection(vec![(Arc::clone(&col_a), "a"), (Arc::clone(&col_b), "b")]); + + // Outer references col 0 twice + let outer = make_projection(vec![ + (Arc::new(Column::new("a", 0)), "x"), + (Arc::new(Column::new("a", 0)), "y"), + ]); + + assert!(!would_duplicate_costly_exprs(&inner, &outer)); + } + + /// A non-trivial computed expression (arithmetic, `KeepInPlace`) referenced + /// multiple times blocks the merge — recomputing it per site is wasteful. + #[test] + fn test_would_duplicate_blocks_computed_multi_ref() { + let col_a: Arc = Arc::new(Column::new("a", 0)); + let col_b: Arc = Arc::new(Column::new("b", 1)); + // Inner: [a + b, b] (index 0 is a non-trivial computed expression) + let inner = make_projection(vec![ + ( + Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + Operator::Plus, + Arc::clone(&col_b), + )), + "sum", + ), + (Arc::clone(&col_b), "b"), + ]); + + // Outer references index 0 twice + let outer = make_projection(vec![ + (Arc::new(Column::new("sum", 0)), "x"), + (Arc::new(Column::new("sum", 0)), "y"), + ]); + + assert!(would_duplicate_costly_exprs(&inner, &outer)); + } + + /// A volatile expression the outer projection does not reference is + /// safe to merge (it is projected away, not duplicated). + #[test] + fn test_would_duplicate_allows_unreferenced_volatile() { + let col_a: Arc = Arc::new(Column::new("a", 0)); + // Inner: [random(), a] + let inner = + make_projection(vec![(make_volatile_expr(), "r"), (Arc::clone(&col_a), "a")]); + + // Outer references only index 1 (the column), not the volatile expr + let outer = make_projection(vec![(Arc::new(Column::new("a", 1)), "a")]); + + assert!(!would_duplicate_costly_exprs(&inner, &outer)); + } + + /// A volatile expression referenced multiple times must block merge: + /// this is the #23220 regression (`random()` aliased then referenced as + /// `x` and `y`). + #[test] + fn test_would_duplicate_blocks_multi_ref_volatile() { + // Inner: [random()] + let inner = make_projection(vec![(make_volatile_expr(), "r")]); + + // Outer references index 0 twice + let outer = make_projection(vec![ + (Arc::new(Column::new("r", 0)), "x"), + (Arc::new(Column::new("r", 0)), "y"), + ]); + + assert!(would_duplicate_costly_exprs(&inner, &outer)); + } + + /// A volatile expression referenced exactly once has nothing to duplicate, + /// so the merge is allowed. + #[test] + fn test_would_duplicate_allows_single_ref_volatile() { + let col_a: Arc = Arc::new(Column::new("a", 0)); + // Inner: [random(), a] + let inner = + make_projection(vec![(make_volatile_expr(), "r"), (Arc::clone(&col_a), "a")]); + + // Outer references the volatile expression exactly once + let outer = make_projection(vec![ + (Arc::new(Column::new("r", 0)), "x"), + (Arc::new(Column::new("a", 1)), "a"), + ]); + + assert!(!would_duplicate_costly_exprs(&inner, &outer)); + } + + /// References are counted with multiplicity, so a single outer expression + /// that duplicates the value (e.g. `r + r`) still blocks the merge. + #[test] + fn test_would_duplicate_blocks_single_expr_self_ref_volatile() { + // Inner: [random()] + let inner = make_projection(vec![(make_volatile_expr(), "r")]); + + // Outer: [r + r] — one expression referencing `random()` twice + let outer = make_projection(vec![( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("r", 0)), + Operator::Plus, + Arc::new(Column::new("r", 0)), + )), + "x", + )]); + + assert!(would_duplicate_costly_exprs(&inner, &outer)); + } + + /// A volatile expression buried inside a larger expression (e.g. + /// `random() + 1`) is still detected and blocks merge. + #[test] + fn test_would_duplicate_blocks_volatile_nested_in_arithmetic() { + // Inner: [random() + 1] + let inner = make_projection(vec![( + Arc::new(BinaryExpr::new( + make_volatile_expr(), + Operator::Plus, + Arc::new(Literal::new(ScalarValue::Float64(Some(1.0)))), + )), + "expr", + )]); + + // Outer references index 0 twice + let outer = make_projection(vec![ + (Arc::new(Column::new("expr", 0)), "x"), + (Arc::new(Column::new("expr", 0)), "y"), + ]); + + assert!(would_duplicate_costly_exprs(&inner, &outer)); + } + + /// Empty projections should not block merging. + #[test] + fn test_would_duplicate_empty_projections() { + let inner = make_projection(vec![]); + let outer = make_projection(vec![]); + assert!(!would_duplicate_costly_exprs(&inner, &outer)); + } + + /// An expensive (scalar-function) expression referenced more than once + /// must block the merge to preserve CSE. + #[test] + fn test_would_duplicate_blocks_multi_ref_expensive() { + let col_a: Arc = Arc::new(Column::new("a", 0)); + // Inner: [abs(a)] + let inner = make_projection(vec![(make_udf_expr(vec![col_a]), "abs_a")]); + + // Outer references index 0 twice + let outer = make_projection(vec![ + (Arc::new(Column::new("abs_a", 0)), "x"), + (Arc::new(Column::new("abs_a", 0)), "y"), + ]); + + assert!(would_duplicate_costly_exprs(&inner, &outer)); + } + + /// An expensive expression referenced only once has nothing to duplicate, + /// so the merge is allowed. + #[test] + fn test_would_duplicate_allows_single_ref_expensive() { + let col_a: Arc = Arc::new(Column::new("a", 0)); + // Inner: [abs(a), a] + let inner = make_projection(vec![ + (make_udf_expr(vec![Arc::clone(&col_a)]), "abs_a"), + (Arc::clone(&col_a), "a"), + ]); + + // Outer references each inner column once + let outer = make_projection(vec![ + (Arc::new(Column::new("abs_a", 0)), "out"), + (Arc::new(Column::new("a", 1)), "a"), + ]); + + assert!(!would_duplicate_costly_exprs(&inner, &outer)); + } + + /// A cheap, leaf-pushable scalar function (placement + /// `MoveTowardsLeafNodes`, e.g. `get_field` / `input_file_name`) still + /// merges even when referenced multiple times — it is meant to be pushed + /// into the scan, so blocking would defeat that optimization. + #[test] + fn test_would_duplicate_allows_leaf_pushable_scalar_function() { + // Inner: [input_file_name()] + let inner = make_projection(vec![(make_leaf_pushable_expr(), "f")]); + + // Outer references index 0 twice + let outer = make_projection(vec![ + (Arc::new(Column::new("f", 0)), "x"), + (Arc::new(Column::new("f", 0)), "y"), + ]); + + assert!(!would_duplicate_costly_exprs(&inner, &outer)); + } } diff --git a/datafusion/datasource/src/file_scan_config/proto.rs b/datafusion/datasource/src/file_scan_config/proto.rs new file mode 100644 index 0000000000000..d7135173c8934 --- /dev/null +++ b/datafusion/datasource/src/file_scan_config/proto.rs @@ -0,0 +1,285 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared serialization of the format-agnostic [`FileScanConfig`] spine. +//! +//! This is the relocated body of `datafusion-proto`'s +//! `serialize_file_scan_config` / `parse_protobuf_file_scan_config`, ported to +//! ride the +//! [`ExecutionPlanEncodeCtx`](datafusion_physical_plan::proto::ExecutionPlanEncodeCtx) / +//! [`ExecutionPlanDecodeCtx`](datafusion_physical_plan::proto::ExecutionPlanDecodeCtx) +//! instead of the raw `PhysicalExtensionCodec` + +//! `PhysicalProtoConverterExtension`. Every +//! `FileSource::try_to_proto` hook (CSV, JSON, Arrow, Parquet, Avro) builds its +//! `*ScanExecNode` around [`FileScanConfig::try_to_proto`] and decodes with +//! [`FileScanConfig::try_from_proto`], keeping a single copy of the shared +//! wire logic. The wire format is byte-for-byte identical to the old central +//! serializer. +//! +//! Child physical expressions (sort orderings, hash/range partitioning, and +//! projection expressions) are (de)serialized through `ctx.encode_expr` / +//! `ctx.decode_expr`; `Schema`, `Statistics`, `Constraints`, and `ScalarValue` +//! go through `datafusion-proto-common`. Nothing here needs the raw codec. + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; +use datafusion_physical_expr::{LexOrdering, Partitioning}; +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; +use datafusion_physical_plan::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; +use datafusion_proto_models::protobuf; + +use crate::file::FileSource; +use crate::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; +use crate::table_schema::TableSchema; + +impl FileScanConfig { + /// Serialize the shared, format-agnostic part of a file scan into a + /// [`protobuf::FileScanExecConf`]. + /// + /// Each concrete [`FileSource::try_to_proto`] + /// wraps the returned value in its own `*ScanExecNode`. Byte-compatible with + /// the former `serialize_file_scan_config` in `datafusion-proto`. + pub fn try_to_proto( + &self, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result { + let file_groups = self + .file_groups + .iter() + .map(TryInto::try_into) + .collect::>>()?; + + let mut output_ordering = vec![]; + for order in &self.output_ordering { + let nodes = sort_exprs_try_to_proto(order.iter(), &ctx.expr_ctx())?; + output_ordering.push(protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes: nodes, + }); + } + + let output_partitioning = self + .output_partitioning + .as_ref() + .map(|partitioning| partitioning.try_to_proto(&ctx.expr_ctx())) + .transpose()?; + + // Fields must be added to the schema so that they can persist in the + // protobuf, and then removed from the schema in `try_from_proto`. + let mut fields = self + .file_schema() + .fields() + .iter() + .cloned() + .collect::>(); + fields.extend(self.table_partition_cols().iter().cloned()); + let schema = + Schema::new(fields).with_metadata(self.file_schema().metadata.clone()); + + let projection_exprs = self + .file_source() + .projection() + .as_ref() + .map(|projection_exprs| { + Ok::<_, DataFusionError>(protobuf::ProjectionExprs { + projections: projection_exprs + .iter() + .map(|expr| { + Ok(protobuf::ProjectionExpr { + alias: expr.alias.to_string(), + expr: Some(ctx.encode_expr(&expr.expr)?), + }) + }) + .collect::>>()?, + }) + }) + .transpose()?; + + Ok(protobuf::FileScanExecConf { + file_groups, + statistics: Some((&self.statistics()).into()), + limit: self.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }), + projection: vec![], + schema: Some((&schema).try_into()?), + table_partition_cols: self + .table_partition_cols() + .iter() + .map(|x| x.name().clone()) + .collect::>(), + object_store_url: self.object_store_url.to_string(), + output_ordering, + constraints: Some(self.constraints.clone().into()), + batch_size: self.batch_size.map(|s| s as u64), + projection_exprs, + output_partitioning, + }) + } + + /// Reconstruct a [`FileScanConfig`] from a [`protobuf::FileScanExecConf`] + /// and a `file_source` the caller has already rebuilt (typically from the + /// table schema via [`FileScanConfig::parse_table_schema_from_proto`]). + /// + /// Byte-compatible with the former `parse_protobuf_file_scan_config`. + pub fn try_from_proto( + conf: &protobuf::FileScanExecConf, + ctx: &ExecutionPlanDecodeCtx<'_>, + file_source: Arc, + ) -> Result { + let schema = parse_file_scan_schema(conf)?; + + let constraints = conf + .constraints + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'constraints'" + ) + })? + .try_into()?; + let statistics = conf + .statistics + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'statistics'" + ) + })? + .try_into()?; + + let file_groups = conf + .file_groups + .iter() + .map(TryInto::try_into) + .collect::>>()?; + + let object_store_url = match conf.object_store_url.is_empty() { + false => ObjectStoreUrl::parse(&conf.object_store_url)?, + true => ObjectStoreUrl::local_filesystem(), + }; + + let mut output_ordering = vec![]; + for node_collection in &conf.output_ordering { + let sort_exprs = sort_exprs_try_from_proto( + &node_collection.physical_sort_expr_nodes, + &ctx.expr_ctx(&schema), + )?; + output_ordering.extend(LexOrdering::new(sort_exprs)); + } + + let output_partitioning = conf + .output_partitioning + .as_ref() + .map(|partitioning| { + Partitioning::try_from_proto(partitioning, &ctx.expr_ctx(&schema)) + }) + .transpose()? + .flatten(); + + // Parse projection expressions if present and apply to the file source. + let file_source = if let Some(proto_projection_exprs) = &conf.projection_exprs { + let projection_exprs: Vec = proto_projection_exprs + .projections + .iter() + .map(|proto_expr| { + let expr = ctx.decode_expr( + proto_expr.expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("ProjectionExpr missing expr field") + })?, + &schema, + )?; + Ok(ProjectionExpr::new(expr, proto_expr.alias.clone())) + }) + .collect::>>()?; + + let projection_exprs = ProjectionExprs::new(projection_exprs); + + file_source + .try_pushdown_projection(&projection_exprs)? + .unwrap_or(file_source) + } else { + file_source + }; + + let config_builder = FileScanConfigBuilder::new(object_store_url, file_source) + .with_file_groups(file_groups) + .with_constraints(constraints) + .with_statistics(statistics) + .with_limit(conf.limit.as_ref().map(|sl| sl.limit as usize)) + .with_output_ordering(output_ordering) + .with_output_partitioning(output_partitioning) + .with_batch_size(conf.batch_size.map(|s| s as usize)); + Ok(config_builder.build()) + } + + /// Parse a [`TableSchema`] (file schema + partition columns) from a + /// [`protobuf::FileScanExecConf`]. File sources use this to rebuild their + /// concrete source before calling [`FileScanConfig::try_from_proto`]. + /// + /// Byte-compatible with the former `parse_table_schema_from_proto`. + pub fn parse_table_schema_from_proto( + conf: &protobuf::FileScanExecConf, + ) -> Result { + let schema = parse_file_scan_schema(conf)?; + + // Reacquire the partition column types from the schema before removing + // them below. + let table_partition_cols = conf + .table_partition_cols + .iter() + .map(|col| Ok(Arc::new(schema.field_with_name(col)?.clone()))) + .collect::>>()?; + + // Remove partition columns from the schema after recreating + // table_partition_cols because the partition columns are not in the + // file. They are present to allow the partition column types to be + // reconstructed after serde. + let file_schema = Arc::new( + Schema::new( + schema + .fields() + .iter() + .filter(|field| !table_partition_cols.contains(field)) + .cloned() + .collect::>(), + ) + .with_metadata(schema.metadata.clone()), + ); + + Ok(TableSchema::builder(file_schema) + .with_table_partition_cols(table_partition_cols) + .build()) + } +} + +/// Parse the full (file + partition columns) schema off the base conf. +fn parse_file_scan_schema(conf: &protobuf::FileScanExecConf) -> Result> { + let schema: Schema = conf + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "FileScanExecConf is missing required field 'schema'" + ) + })? + .try_into()?; + Ok(Arc::new(schema)) +} diff --git a/datafusion/datasource/src/file_scan_config/sort_pushdown.rs b/datafusion/datasource/src/file_scan_config/sort_pushdown.rs index af08ed71b9a6d..3f5beed20fa8d 100644 --- a/datafusion/datasource/src/file_scan_config/sort_pushdown.rs +++ b/datafusion/datasource/src/file_scan_config/sort_pushdown.rs @@ -138,31 +138,76 @@ impl FileScanConfig { false }; - if is_exact && all_non_overlapping { - // Truly exact: within-file ordering guaranteed and files are non-overlapping. - // Keep output_ordering so SortExec can be eliminated for each partition. - // - // We intentionally do NOT redistribute files across groups here. - // The planning-phase bin-packing may interleave file ranges across groups: - // - // Group 0: [f1(1-10), f3(21-30)] ← interleaved with group 1 - // Group 1: [f2(11-20), f4(31-40)] - // - // This interleaving is actually beneficial because SPM pulls from both - // partitions concurrently, keeping parallel I/O active: - // - // SPM: pull P0 [1-10] → pull P1 [11-20] → pull P0 [21-30] → pull P1 [31-40] - // ^^^^^^^^^^^^ ^^^^^^^^^^^^ - // both partitions scanning files simultaneously - // - // If we were to redistribute files consecutively: - // Group 0: [f1(1-10), f2(11-20)] ← all values < group 1 - // Group 1: [f3(21-30), f4(31-40)] + // Decide whether to keep `output_ordering` (i.e. let the outer + // pushdown report `Exact` and drop `SortExec`). + // + // Two paths can produce a keep: + // + // 1. `is_exact && all_non_overlapping`: the source already had + // validated ordering and the post-sort files still don't + // overlap — Exact carries through unchanged. + // + // 2. `!is_exact && all_non_overlapping`: source returned + // `Inexact` because pre-sort `validated_output_ordering()` + // stripped the declaration (files were listed out of order + // on disk). After our stats-based sort the files are now + // non-overlapping — re-validate against the new file + // groups and, if it passes, upgrade back to Exact so the + // outer wrapper drops the `SortExec`. Without this, the + // `Inexact` branch stayed Inexact even when reorder + // restored a perfectly valid ordering, leaving an + // unnecessary `SortExec` above the source (regression + // after #21956's `column_in_file_schema` signal pushed + // this scenario into the Inexact branch instead of the + // `try_sort_file_groups_by_statistics` fallback). + // + // We intentionally do NOT redistribute files across groups here. + // The planning-phase bin-packing may interleave file ranges across groups: + // + // Group 0: [f1(1-10), f3(21-30)] ← interleaved with group 1 + // Group 1: [f2(11-20), f4(31-40)] + // + // This interleaving is actually beneficial because SPM pulls from both + // partitions concurrently, keeping parallel I/O active. + let keep_ordering = match (all_non_overlapping, is_exact) { + // Files still overlap after the stats sort — the combined + // stream isn't ordered, so `output_ordering` must be dropped. + (false, _) => false, + // Source already had validated ordering and the post-sort + // files still don't overlap — Exact carries through. + (true, true) => true, + // Source returned `Inexact`; re-validate against the + // reordered file groups to decide whether to upgrade. // - // SPM would read ALL of group 0 first (values always smaller), then group 1. - // This degrades to single-threaded sequential I/O — the other partition - // sits idle the entire time, losing the parallelism benefit. - } else { + // Same NULL guard as `try_sort_file_groups_by_statistics`: + // we cannot claim Exact if any non-last file contains + // NULLs in the sort columns. With NULLS LAST those + // NULLs sit after all non-null rows in the file, so + // when the next file's non-nulls are smaller than the + // previous file's max, they'd appear *after* the NULLs + // in the concatenated stream — breaking the ordering. + (true, false) => { + let projected_schema = new_config.projected_schema()?; + let projection_indices = new_config + .file_source + .projection() + .as_ref() + .and_then(|p| ordered_column_indices_from_projection(p)); + if any_file_has_nulls_in_sort_columns( + &new_config.file_groups, + order, + &projected_schema, + projection_indices.as_deref(), + ) { + false + } else { + let new_eq_props = new_config.eq_properties(); + new_eq_props.ordering_satisfy(order.iter().cloned())? + } + } + }; + + if !keep_ordering { new_config.output_ordering = vec![]; } @@ -490,7 +535,7 @@ pub(crate) fn validate_orderings( /// file is scanned, the same values for A, B and C can be repeated in /// the same sorted stream /// -///```text +/// ```text /// ┏ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ ━ /// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┃ /// ┃ ┌───────────────┐ ┌──────────────┐ │ diff --git a/datafusion/datasource/src/file_sink_config.rs b/datafusion/datasource/src/file_sink_config.rs index 1abce86a3565f..48dce9a0cdb3e 100644 --- a/datafusion/datasource/src/file_sink_config.rs +++ b/datafusion/datasource/src/file_sink_config.rs @@ -32,6 +32,9 @@ use datafusion_expr::dml::InsertOp; use async_trait::async_trait; use object_store::ObjectStore; +#[cfg(feature = "proto")] +mod proto; + /// Determines how `FileSink` output paths are interpreted. #[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum FileOutputMode { diff --git a/datafusion/datasource/src/file_sink_config/proto.rs b/datafusion/datasource/src/file_sink_config/proto.rs new file mode 100644 index 0000000000000..ed4b5c48bd2af --- /dev/null +++ b/datafusion/datasource/src/file_sink_config/proto.rs @@ -0,0 +1,239 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Protobuf conversion for the format-independent [`FileSinkConfig`]. + +use std::sync::Arc; + +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_expr::dml::InsertOp; +use datafusion_proto_models::protobuf; + +use crate::ListingTableUrl; +use crate::file_groups::FileGroup; +use crate::file_sink_config::{FileOutputMode, FileSinkConfig}; + +impl TryFrom<&FileSinkConfig> for protobuf::FileSinkConfig { + type Error = DataFusionError; + + /// Serialize this shared file-sink configuration without format-specific + /// writer options. + fn try_from(config: &FileSinkConfig) -> Result { + let file_groups = config + .file_group + .iter() + .map(TryInto::try_into) + .collect::>>()?; + let table_paths = config + .table_paths + .iter() + .map(ToString::to_string) + .collect::>(); + let table_partition_cols = config + .table_partition_cols + .iter() + .map(|(name, data_type)| { + Ok(protobuf::PartitionColumn { + name: name.to_owned(), + arrow_type: Some(data_type.try_into()?), + }) + }) + .collect::>>()?; + let insert_op = match config.insert_op { + InsertOp::Append => protobuf::InsertOp::Append, + InsertOp::Overwrite => protobuf::InsertOp::Overwrite, + InsertOp::Replace => protobuf::InsertOp::Replace, + }; + let file_output_mode = match config.file_output_mode { + FileOutputMode::Automatic => protobuf::FileOutputMode::Automatic, + FileOutputMode::SingleFile => protobuf::FileOutputMode::SingleFile, + FileOutputMode::Directory => protobuf::FileOutputMode::Directory, + }; + + Ok(protobuf::FileSinkConfig { + object_store_url: config.object_store_url.to_string(), + file_groups, + table_paths, + output_schema: Some(config.output_schema.as_ref().try_into()?), + table_partition_cols, + keep_partition_by_columns: config.keep_partition_by_columns, + insert_op: insert_op.into(), + file_extension: config.file_extension.clone(), + file_output_mode: file_output_mode.into(), + }) + } +} + +impl TryFrom<&protobuf::FileSinkConfig> for FileSinkConfig { + type Error = DataFusionError; + + /// Reconstruct a shared file-sink configuration from protobuf. + fn try_from(conf: &protobuf::FileSinkConfig) -> Result { + let file_group = FileGroup::new( + conf.file_groups + .iter() + .map(TryInto::try_into) + .collect::>>()?, + ); + let table_paths = conf + .table_paths + .iter() + .map(ListingTableUrl::parse) + .collect::>>()?; + let table_partition_cols = conf + .table_partition_cols + .iter() + .map(|protobuf::PartitionColumn { name, arrow_type }| { + let data_type = arrow_type + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "PartitionColumn is missing required field 'arrow_type'" + ) + })? + .try_into()?; + Ok((name.clone(), data_type)) + }) + .collect::>>()?; + let insert_op = protobuf::InsertOp::try_from(conf.insert_op).map_err(|_| { + internal_datafusion_err!( + "Received a FileSinkConfig message with unknown InsertOp {}", + conf.insert_op + ) + })?; + let insert_op = match insert_op { + protobuf::InsertOp::Append => InsertOp::Append, + protobuf::InsertOp::Overwrite => InsertOp::Overwrite, + protobuf::InsertOp::Replace => InsertOp::Replace, + }; + let file_output_mode = protobuf::FileOutputMode::try_from(conf.file_output_mode) + .map_err(|_| { + internal_datafusion_err!( + "Received a FileSinkConfig message with unknown FileOutputMode {}", + conf.file_output_mode + ) + })?; + let file_output_mode = match file_output_mode { + protobuf::FileOutputMode::Automatic => FileOutputMode::Automatic, + protobuf::FileOutputMode::SingleFile => FileOutputMode::SingleFile, + protobuf::FileOutputMode::Directory => FileOutputMode::Directory, + }; + let output_schema = conf.output_schema.as_ref().ok_or_else(|| { + internal_datafusion_err!( + "FileSinkConfig is missing required field 'output_schema'" + ) + })?; + + Ok(Self { + original_url: String::default(), + object_store_url: ObjectStoreUrl::parse(&conf.object_store_url)?, + file_group, + table_paths, + output_schema: Arc::new(output_schema.try_into()?), + table_partition_cols, + insert_op, + keep_partition_by_columns: conf.keep_partition_by_columns, + file_extension: conf.file_extension.clone(), + file_output_mode, + }) + } +} + +#[cfg(test)] +mod tests { + use arrow::datatypes::Schema; + + use super::*; + + fn valid_file_sink_config() -> protobuf::FileSinkConfig { + protobuf::FileSinkConfig { + object_store_url: ObjectStoreUrl::local_filesystem().to_string(), + output_schema: Some( + (&Schema::empty()) + .try_into() + .expect("empty schema should serialize"), + ), + insert_op: protobuf::InsertOp::Append.into(), + file_output_mode: protobuf::FileOutputMode::Automatic.into(), + ..Default::default() + } + } + + fn assert_decode_error( + mutate: impl FnOnce(&mut protobuf::FileSinkConfig), + expected: impl AsRef, + ) { + let mut conf = valid_file_sink_config(); + mutate(&mut conf); + + let error = + FileSinkConfig::try_from(&conf).expect_err("invalid config should fail"); + match error { + DataFusionError::Internal(message) => { + let message = message + .split_once(DataFusionError::BACK_TRACE_SEP) + .map_or(message.as_str(), |(message, _)| message); + assert_eq!(message, expected.as_ref()); + } + error => panic!("expected internal error, got {error}"), + } + } + + #[test] + fn rejects_unknown_insert_op() { + assert_decode_error( + |conf| conf.insert_op = i32::MAX, + format!( + "Received a FileSinkConfig message with unknown InsertOp {}", + i32::MAX + ), + ); + } + + #[test] + fn rejects_unknown_file_output_mode() { + assert_decode_error( + |conf| conf.file_output_mode = i32::MAX, + format!( + "Received a FileSinkConfig message with unknown FileOutputMode {}", + i32::MAX + ), + ); + } + + #[test] + fn rejects_missing_output_schema() { + assert_decode_error( + |conf| conf.output_schema = None, + "FileSinkConfig is missing required field 'output_schema'", + ); + } + + #[test] + fn rejects_partition_column_without_arrow_type() { + assert_decode_error( + |conf| { + conf.table_partition_cols.push(protobuf::PartitionColumn { + name: "partition".to_string(), + arrow_type: None, + }); + }, + "PartitionColumn is missing required field 'arrow_type'", + ); + } +} diff --git a/datafusion/datasource/src/file_stream/mod.rs b/datafusion/datasource/src/file_stream/mod.rs index e277690cff810..6daed7c338022 100644 --- a/datafusion/datasource/src/file_stream/mod.rs +++ b/datafusion/datasource/src/file_stream/mod.rs @@ -182,6 +182,7 @@ mod tests { use arrow::array::{AsArray, RecordBatch}; use arrow::datatypes::{DataType, Field, Int32Type, Schema}; use datafusion_common::DataFusionError; + use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; @@ -315,7 +316,7 @@ mod tests { let on_error = self.on_error; - let table_schema = TableSchema::new(file_schema, vec![]); + let table_schema = TableSchema::from(file_schema); let config = FileScanConfigBuilder::new( ObjectStoreUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema)), @@ -352,7 +353,7 @@ mod tests { /// Create the smallest valid file scan config for builder validation tests. fn builder_test_config() -> FileScanConfig { - let table_schema = TableSchema::new(Arc::new(Schema::empty()), vec![]); + let table_schema = TableSchema::from(Arc::new(Schema::empty())); FileScanConfigBuilder::new( ObjectStoreUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema)), @@ -1106,13 +1107,13 @@ mod tests { Ok(()) } - /// Verifies that `partitioned_by_file_group` disables shared work stealing. + /// Verifies that declared output partitioning disables shared work stealing. #[tokio::test] - async fn morsel_partitioned_by_file_group_keeps_files_local() -> Result<()> { + async fn morsel_declared_output_partitioning_keeps_files_local() -> Result<()> { // same fixture as `morsel_shared_files_can_be_stolen` but marked as // preserve-partitioned let test = two_partition_morsel_test() - .with_partitioned_by_file_group(true) + .with_declared_output_partitioning(true) .with_file_stream_events(false); insta::assert_snapshot!(test.run().await.unwrap(), @r" @@ -1131,6 +1132,40 @@ mod tests { Ok(()) } + /// Verifies that disabling `enable_file_stream_work_stealing` keeps each + /// stream's files local, so a sibling cannot steal them at runtime. + /// + /// Covers : executors + /// that run each output partition as an isolated task in a separate process + /// (Ballista, datafusion-distributed) poll only their own partition, so the + /// shared work queue would let that one partition drain files belonging to + /// its siblings. Disabling the flag falls back to per-partition file groups. + #[tokio::test] + async fn morsel_disabled_work_stealing_keeps_files_local() -> Result<()> { + // same fixture as `morsel_shared_files_can_be_stolen`, but with work + // stealing disabled via config + let test = two_partition_morsel_test() + .with_enable_file_stream_work_stealing(false) + .with_file_stream_events(false); + + // Even though Partition 1 is polled first, it cannot steal the three + // files assigned to Partition 0; each partition reads only its own. + insta::assert_snapshot!(test.run().await.unwrap(), @r" + ----- Partition 0 ----- + Batch: 101 + Batch: 102 + Batch: 103 + Done + ----- Partition 1 ----- + Batch: 201 + Done + ----- File Stream Events ----- + (omitted due to with_file_stream_events(false)) + "); + + Ok(()) + } + /// Verifies that an empty sibling can immediately steal shared files when /// it is polled before the stream that originally owned them. #[tokio::test] @@ -1216,7 +1251,7 @@ mod tests { let unlimited_config = test.test_config(); let limited_config = test.clone().with_limit(1).test_config(); let shared_work_source = limited_config - .create_sibling_state() + .create_sibling_state(&ConfigOptions::default()) .and_then(|state| state.as_ref().downcast_ref::().cloned()) .expect("shared work source"); let limited_metrics = ExecutionPlanMetricsSet::new(); @@ -1331,7 +1366,8 @@ mod tests { morselizer: MockMorselizer, partition_files: BTreeMap>, preserve_order: bool, - partitioned_by_file_group: bool, + declared_output_partitioning: bool, + enable_file_stream_work_stealing: bool, file_stream_events: bool, build_streams_on_first_read: bool, reads: Vec, @@ -1345,7 +1381,8 @@ mod tests { morselizer: MockMorselizer::new(), partition_files: BTreeMap::new(), preserve_order: false, - partitioned_by_file_group: false, + declared_output_partitioning: false, + enable_file_stream_work_stealing: true, file_stream_events: true, build_streams_on_first_read: false, reads: vec![], @@ -1381,13 +1418,21 @@ mod tests { self } - /// Marks the test scan as pre-partitioned by file group, which should - /// force each stream to keep its own files local. - fn with_partitioned_by_file_group( + /// Declares the test scan's output partitioning, which should force + /// each stream to keep its own files local. + fn with_declared_output_partitioning( mut self, - partitioned_by_file_group: bool, + declared_output_partitioning: bool, ) -> Self { - self.partitioned_by_file_group = partitioned_by_file_group; + self.declared_output_partitioning = declared_output_partitioning; + self + } + + /// Sets `datafusion.execution.enable_file_stream_work_stealing`. When + /// disabled, each stream keeps its own files local instead of sharing a + /// work queue with its siblings. + fn with_enable_file_stream_work_stealing(mut self, enable: bool) -> Self { + self.enable_file_stream_work_stealing = enable; self } @@ -1468,9 +1513,13 @@ mod tests { // `FileStream`s directly, bypassing `DataSourceExec`, so they must // perform the same setup explicitly when exercising sibling-stream // work stealing. - let shared_work_source = config.create_sibling_state().and_then(|state| { - state.as_ref().downcast_ref::().cloned() - }); + let mut options = ConfigOptions::default(); + options.execution.enable_file_stream_work_stealing = + self.enable_file_stream_work_stealing; + let shared_work_source = + config.create_sibling_state(&options).and_then(|state| { + state.as_ref().downcast_ref::().cloned() + }); if !self.build_streams_on_first_read { for partition in build_order { let stream = FileStreamBuilder::new(&config) @@ -1575,10 +1624,22 @@ mod tests { }) .collect::>(); - let table_schema = TableSchema::new( - Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])), - vec![], - ); + let table_schema = + TableSchema::from(Arc::new(Schema::new(vec![Field::new( + "i", + DataType::Int32, + false, + )]))); + // Declaring an output partitioning marks the scan as pre-grouped, which + // keeps each stream's files local (disables shared work stealing). + let output_partitioning = self.declared_output_partitioning.then(|| { + datafusion_physical_expr::Partitioning::Hash( + vec![Arc::new( + datafusion_physical_expr::expressions::Column::new("i", 0), + )], + file_groups.len(), + ) + }); FileScanConfigBuilder::new( ObjectStoreUrl::parse("test:///").unwrap(), Arc::new(MockSource::new(table_schema)), @@ -1586,7 +1647,7 @@ mod tests { .with_file_groups(file_groups) .with_limit(self.limit) .with_preserve_order(self.preserve_order) - .with_partitioned_by_file_group(self.partitioned_by_file_group) + .with_output_partitioning(output_partitioning) .build() } } diff --git a/datafusion/datasource/src/memory.rs b/datafusion/datasource/src/memory.rs index 7c9281dcc2f26..2370ed87a2954 100644 --- a/datafusion/datasource/src/memory.rs +++ b/datafusion/datasource/src/memory.rs @@ -260,16 +260,61 @@ impl DataSource for MemorySourceConfig { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in sort_information - let mut tnr = TreeNodeRecursion::Continue; + Ok(TreeNodeRecursion::Continue) + } + + /// Serialize this `MemorySourceConfig` as a `MemoryScanExecNode` wrapped + /// in a [`PhysicalPlanNode`]. Byte-compatible with the former central + /// `MemoryScan` arm in `datafusion-proto`. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto; + use datafusion_proto_models::protobuf; + + let partitions = self + .partitions + .iter() + .map(|batches| record_batches_to_ipc_bytes(batches)) + .collect::>>()?; + + // Proto3 can't tell `None` from `Some(vec![])`; encode the latter + // as the `[u32::MAX]` sentinel, matching the join/filter nodes. + let projection = match self.projection.as_ref() { + None => Vec::new(), + Some(v) if v.is_empty() => vec![u32::MAX], + Some(v) => v.iter().map(|x| *x as u32).collect(), + }; + + let mut sort_information = Vec::with_capacity(self.sort_information.len()); for ordering in &self.sort_information { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } + let physical_sort_expr_nodes = + sort_exprs_try_to_proto(ordering.iter(), &ctx.expr_ctx())?; + sort_information.push(protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes, + }); } - Ok(tnr) + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::MemoryScan( + protobuf::MemoryScanExecNode { + partitions, + schema: Some(self.schema.as_ref().try_into()?), + projection, + sort_information, + show_sizes: self.show_sizes, + fetch: self.fetch.map(|f| f as u32), + }, + ), + ), + })) } } @@ -622,6 +667,96 @@ impl MemorySourceConfig { } } +#[cfg(feature = "proto")] +impl MemorySourceConfig { + /// Reconstruct a [`DataSourceExec`] wrapping a `MemorySourceConfig` from + /// its protobuf representation. Byte-compatible with the former central + /// `MemoryScan` arm in `datafusion-proto`. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::internal_datafusion_err; + use datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto; + use datafusion_proto_models::protobuf; + + let scan = datafusion_physical_plan::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::MemoryScan, + "MemorySourceConfig", + ); + + let partitions = scan + .partitions + .iter() + .map(|buf| record_batches_from_ipc_bytes(buf)) + .collect::>>()?; + + let proto_schema = scan.schema.as_ref().ok_or_else(|| { + internal_datafusion_err!("schema in MemoryScanExecNode is missing.") + })?; + let schema: SchemaRef = SchemaRef::new(proto_schema.try_into()?); + + // Preserve the empty-projection sentinel written by `try_to_proto`. + let projection = match scan.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), + }; + + let mut sort_information = vec![]; + for ordering in &scan.sort_information { + let sort_exprs = sort_exprs_try_from_proto( + &ordering.physical_sort_expr_nodes, + &ctx.expr_ctx(&schema), + )?; + sort_information.extend(LexOrdering::new(sort_exprs)); + } + + let source = Self::try_new(&partitions, schema, projection)? + .with_limit(scan.fetch.map(|f| f as usize)) + .with_show_sizes(scan.show_sizes) + .try_with_sort_information(sort_information)?; + + Ok(DataSourceExec::from_data_source(source)) + } +} + +/// Encode record batches as Arrow IPC stream bytes; an empty slice encodes to +/// an empty buffer. +#[cfg(feature = "proto")] +fn record_batches_to_ipc_bytes(batches: &[RecordBatch]) -> Result> { + use arrow::ipc::writer::StreamWriter; + + if batches.is_empty() { + return Ok(vec![]); + } + let schema = batches[0].schema(); + let mut buf = Vec::new(); + let mut writer = StreamWriter::try_new(&mut buf, &schema)?; + for batch in batches { + writer.write(batch)?; + } + writer.finish()?; + Ok(buf) +} + +/// Inverse of [`record_batches_to_ipc_bytes`]. +#[cfg(feature = "proto")] +fn record_batches_from_ipc_bytes(buf: &[u8]) -> Result> { + use arrow::ipc::reader::StreamReader; + + if buf.is_empty() { + return Ok(vec![]); + } + let reader = StreamReader::try_new(buf, None)?; + let mut batches = Vec::new(); + for batch in reader { + batches.push(batch?); + } + Ok(batches) +} + /// For use in repartitioning, track the total size and original partition index. /// /// Do not implement clone, in order to avoid unnecessary copying during repartitioning. @@ -868,6 +1003,7 @@ mod tests { use datafusion_common::stats::{ColumnStatistics, Precision}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_plan::expressions::lit; + use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::ExecutionPlan; @@ -1000,7 +1136,7 @@ mod tests { let values = MemorySourceConfig::try_new_as_values(schema, data)?; assert_eq!( - *values.partition_statistics(None)?, + *StatisticsContext::new().compute(values.as_ref(), &StatisticsArgs::new())?, Statistics { num_rows: Precision::Exact(rows), total_byte_size: Precision::Exact(8), // not important diff --git a/datafusion/datasource/src/mod.rs b/datafusion/datasource/src/mod.rs index 84daf608b5182..e415b3e48a02a 100644 --- a/datafusion/datasource/src/mod.rs +++ b/datafusion/datasource/src/mod.rs @@ -28,6 +28,7 @@ //! A table that uses the `ObjectStore` listing capability //! to get the list of files to process. +pub mod boundary_stream; pub mod decoder; pub mod display; pub mod file; @@ -40,6 +41,10 @@ pub mod file_stream; pub mod memory; pub mod morsel; pub mod projection; +/// Protobuf conversions for [`FileRange`], [`PartitionedFile`] and +/// [`FileGroup`](crate::file_groups::FileGroup), gated on the `proto` feature. +#[cfg(feature = "proto")] +mod proto; pub mod schema_adapter; pub mod sink; pub mod source; @@ -54,23 +59,19 @@ pub mod write; pub use self::file::as_file_source; pub use self::url::ListingTableUrl; use crate::file_groups::FileGroup; +use arrow::datatypes::SchemaRef; use chrono::TimeZone; use datafusion_common::stats::Precision; -use datafusion_common::{ColumnStatistics, Result, TableReference, exec_datafusion_err}; +use datafusion_common::{ColumnStatistics, Result, TableReference}; use datafusion_common::{ScalarValue, Statistics}; use datafusion_physical_expr::LexOrdering; -use futures::{Stream, StreamExt}; -use object_store::{GetOptions, GetRange, ObjectStore}; +use futures::Stream; use object_store::{ObjectMeta, path::Path}; -pub use table_schema::TableSchema; -// Remove when add_row_stats is remove -#[expect(deprecated)] -pub use statistics::add_row_stats; pub use statistics::compute_all_files_statistics; use std::any::Any; -use std::ops::Range; use std::pin::Pin; use std::sync::Arc; +pub use table_schema::{TableSchema, TableSchemaBuilder}; /// User-defined per-file extension data, keyed by concrete Rust type. /// @@ -163,12 +164,23 @@ pub struct PartitionedFile { /// The estimated size of the parquet metadata, in bytes pub metadata_size_hint: Option, pub table_reference: Option, + /// A user-provided physical Arrow schema for this file. + /// + /// This schema describes only the columns stored in the file. It must not + /// include partition columns; those are represented separately by + /// [`Self::partition_values`] and the scan's table partition columns. + /// + /// When provided, this field will be used by the Parquet reader to avoid + /// parsing the Arrow schema from the `ARROW:schema` metadata key. Other + /// built-in file sources ignore it for now. + pub arrow_schema: Option, } impl PartitionedFile { /// Create a simple file without metadata or partition pub fn new(path: impl Into, size: u64) -> Self { Self { + arrow_schema: None, object_meta: ObjectMeta { location: Path::from(path.into()), last_modified: chrono::Utc.timestamp_nanos(0), @@ -189,6 +201,7 @@ impl PartitionedFile { /// Create a file from a known ObjectMeta without partition pub fn new_from_meta(object_meta: ObjectMeta) -> Self { Self { + arrow_schema: None, object_meta, partition_values: vec![], range: None, @@ -203,6 +216,7 @@ impl PartitionedFile { /// Create a file range without metadata or partition pub fn new_with_range(path: String, size: u64, start: i64, end: i64) -> Self { Self { + arrow_schema: None, object_meta: ObjectMeta { location: Path::from(path), last_modified: chrono::Utc.timestamp_nanos(0), @@ -221,6 +235,15 @@ impl PartitionedFile { .with_range(start, end) } + /// Provide a physical Arrow schema for this file. + /// + /// The schema must describe only columns stored in the file and must not + /// include partition columns. See [`Self::arrow_schema`] for details. + pub fn with_arrow_schema(mut self, schema: SchemaRef) -> Self { + self.arrow_schema = Some(schema); + self + } + /// Attach partition values to this file. /// This replaces any existing partition values. pub fn with_partition_values(mut self, partition_values: Vec) -> Self { @@ -376,6 +399,7 @@ impl From for PartitionedFile { fn from(object_meta: ObjectMeta) -> Self { PartitionedFile { object_meta, + arrow_schema: None, partition_values: vec![], range: None, statistics: None, @@ -387,119 +411,6 @@ impl From for PartitionedFile { } } -/// Represents the possible outcomes of a range calculation. -/// -/// This enum is used to encapsulate the result of calculating the range of -/// bytes to read from an object (like a file) in an object store. -/// -/// Variants: -/// - `Range(Option>)`: -/// Represents a range of bytes to be read. It contains an `Option` wrapping a -/// `Range`. `None` signifies that the entire object should be read, -/// while `Some(range)` specifies the exact byte range to read. -/// - `TerminateEarly`: -/// Indicates that the range calculation determined no further action is -/// necessary, possibly because the calculated range is empty or invalid. -pub enum RangeCalculation { - Range(Option>), - TerminateEarly, -} - -/// Calculates an appropriate byte range for reading from an object based on the -/// provided metadata. -/// -/// This asynchronous function examines the [`PartitionedFile`] of an object in an object store -/// and determines the range of bytes to be read. The range calculation may adjust -/// the start and end points to align with meaningful data boundaries (like newlines). -/// -/// Returns a `Result` wrapping a [`RangeCalculation`], which is either a calculated byte range or an indication to terminate early. -/// -/// Returns an `Error` if any part of the range calculation fails, such as issues in reading from the object store or invalid range boundaries. -pub async fn calculate_range( - file: &PartitionedFile, - store: &Arc, - terminator: Option, -) -> Result { - let location = &file.object_meta.location; - let file_size = file.object_meta.size; - let newline = terminator.unwrap_or(b'\n'); - - match file.range { - None => Ok(RangeCalculation::Range(None)), - Some(FileRange { start, end }) => { - let start: u64 = start.try_into().map_err(|_| { - exec_datafusion_err!("Expect start range to fit in u64, got {start}") - })?; - let end: u64 = end.try_into().map_err(|_| { - exec_datafusion_err!("Expect end range to fit in u64, got {end}") - })?; - - let start_delta = if start != 0 { - find_first_newline(store, location, start - 1, file_size, newline).await? - } else { - 0 - }; - - if start + start_delta > end { - return Ok(RangeCalculation::TerminateEarly); - } - - let end_delta = if end != file_size { - find_first_newline(store, location, end - 1, file_size, newline).await? - } else { - 0 - }; - - let range = start + start_delta..end + end_delta; - - if range.start >= range.end { - return Ok(RangeCalculation::TerminateEarly); - } - - Ok(RangeCalculation::Range(Some(range))) - } - } -} - -/// Asynchronously finds the position of the first newline character in a specified byte range -/// within an object, such as a file, in an object store. -/// -/// This function scans the contents of the object starting from the specified `start` position -/// up to the `end` position, looking for the first occurrence of a newline character. -/// It returns the position of the first newline relative to the start of the range. -/// -/// Returns a `Result` wrapping a `usize` that represents the position of the first newline character found within the specified range. If no newline is found, it returns the length of the scanned data, effectively indicating the end of the range. -/// -/// The function returns an `Error` if any issues arise while reading from the object store or processing the data stream. -async fn find_first_newline( - object_store: &Arc, - location: &Path, - start: u64, - end: u64, - newline: u8, -) -> Result { - let options = GetOptions { - range: Some(GetRange::Bounded(start..end)), - ..Default::default() - }; - - let result = object_store.get_opts(location, options).await?; - let mut result_stream = result.into_stream(); - - let mut index = 0; - - while let Some(chunk) = result_stream.next().await.transpose()? { - if let Some(position) = chunk.iter().position(|&byte| byte == newline) { - let position = position as u64; - return Ok(index + position); - } - - index += chunk.len() as u64; - } - - Ok(index) -} - /// Generates test files with min-max statistics in different overlap patterns. /// /// Used by tests and benchmarks. @@ -556,6 +467,7 @@ pub fn generate_test_files(num_files: usize, overlap_factor: f64) -> Vec - #[tokio::test] - async fn test_calculate_range_single_line_file() { - use super::{PartitionedFile, RangeCalculation, calculate_range}; - use object_store::ObjectStore; - use object_store::memory::InMemory; - - let content = r#"{"id":1,"data":"aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"}"#; - let file_size = content.len() as u64; - - let store: Arc = Arc::new(InMemory::new()); - let path = Path::from("test.json"); - store.put(&path, content.into()).await.unwrap(); - - let mid = file_size / 2; - let partitioned_file = PartitionedFile::new_with_range( - path.to_string(), - file_size, - mid as i64, - file_size as i64, - ); - - let result = calculate_range(&partitioned_file, &store, None).await; - - assert!(matches!(result, Ok(RangeCalculation::TerminateEarly))); - } } diff --git a/datafusion/datasource/src/projection.rs b/datafusion/datasource/src/projection.rs index ac33a96ca8321..3cf4f29a77a25 100644 --- a/datafusion/datasource/src/projection.rs +++ b/datafusion/datasource/src/projection.rs @@ -26,6 +26,7 @@ use datafusion_physical_expr::{ expressions::{Column, Literal}, projection::{ProjectionExpr, ProjectionExprs}, }; +use datafusion_physical_expr_adapter::rewrite::rewrite_input_file_name_in_projection; use futures::{FutureExt, StreamExt}; use itertools::Itertools; @@ -69,6 +70,7 @@ impl ProjectionOpener { impl FileOpener for ProjectionOpener { fn open(&self, partitioned_file: PartitionedFile) -> Result { let partition_values = partitioned_file.partition_values.clone(); + // Modify any references to partition columns in the projection expressions // and substitute them with literal values from PartitionedFile.partition_values let projection = if self.partition_columns.is_empty() { @@ -80,6 +82,11 @@ impl FileOpener for ProjectionOpener { partition_values, ) }; + // Replace `input_file_name()` with a per-file literal if present. + let projection = rewrite_input_file_name_in_projection( + projection, + partitioned_file.object_meta.location.as_ref(), + )?; let projector = projection.make_projector(&self.input_schema)?; let inner = self.inner.open(partitioned_file)?; @@ -287,22 +294,47 @@ impl SplitProjection { mod test { use std::sync::Arc; - use arrow::array::AsArray; - use arrow::datatypes::{DataType, SchemaRef}; - use datafusion_common::{DFSchema, ScalarValue, record_batch}; - use datafusion_expr::{Expr, col, execution_props::ExecutionProps}; - use datafusion_physical_expr::{create_physical_exprs, projection::ProjectionExpr}; + use arrow::array::{AsArray, RecordBatch, record_batch}; + use arrow::datatypes as arrow_schema; + use arrow::datatypes::{DataType, Field, SchemaRef}; + use datafusion_common::{DFSchema, ScalarValue, config::ConfigOptions}; + use datafusion_expr::{ + Expr, ScalarUDF, col, execution_props::ExecutionProps, + physical_planning_context::PhysicalPlanningContext, + }; + use datafusion_functions::core::input_file_name::InputFileNameFunc; + use datafusion_physical_expr::{ + ScalarFunctionExpr, create_physical_exprs, projection::ProjectionExpr, + }; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use futures::{FutureExt, StreamExt}; use itertools::Itertools; use super::*; + struct StaticBatchOpener { + batch: RecordBatch, + } + + impl FileOpener for StaticBatchOpener { + fn open(&self, _partitioned_file: PartitionedFile) -> Result { + let batch = self.batch.clone(); + Ok(async move { Ok(futures::stream::iter([Ok(batch)]).boxed()) }.boxed()) + } + } + fn create_projection_exprs<'a>( exprs: impl IntoIterator, schema: &SchemaRef, ) -> ProjectionExprs { let df_schema = DFSchema::try_from(Arc::clone(schema)).unwrap(); - let physical_exprs = - create_physical_exprs(exprs, &df_schema, &ExecutionProps::default()).unwrap(); + let physical_exprs = create_physical_exprs( + exprs, + &df_schema, + &ExecutionProps::default(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); let projection_exprs = physical_exprs .into_iter() .enumerate() @@ -311,6 +343,68 @@ mod test { ProjectionExprs::from(projection_exprs) } + fn input_file_name_expr() -> Arc { + Arc::new(ScalarFunctionExpr::new( + "input_file_name", + Arc::new(ScalarUDF::from(InputFileNameFunc::new())), + vec![], + Arc::new(Field::new("input_file_name", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + )) + } + + #[tokio::test] + async fn test_projection_opener_rewrites_input_file_name_with_partitions() { + let file_schema = Schema::new(vec![Field::new("value", DataType::Int32, false)]); + let projection = ProjectionExprs::new([ + ProjectionExpr::new(Arc::new(Column::new("value", 0)), "value"), + ProjectionExpr::new(Arc::new(Column::new("part", 1)), "part"), + ProjectionExpr::new(input_file_name_expr(), "file_name"), + ]); + let split = SplitProjection::new(&file_schema, &projection); + let input_batch = + record_batch!(("value", Int32, vec![10, 20])).expect("input batch"); + + let opener = ProjectionOpener::try_new( + split, + Arc::new(StaticBatchOpener { batch: input_batch }), + &file_schema, + ) + .expect("projection opener"); + + let mut file = PartitionedFile::new("part=west/data.csv", 100); + file.partition_values = vec![ScalarValue::from("west")]; + let mut stream = opener + .open(file) + .expect("open projection") + .await + .expect("inner stream"); + let batch = stream + .next() + .await + .expect("one projected batch") + .expect("projected batch"); + assert!(stream.next().await.is_none()); + + assert_eq!(batch.schema().field(0).name(), "value"); + assert_eq!(batch.schema().field(1).name(), "part"); + assert_eq!(batch.schema().field(2).name(), "file_name"); + + let values = batch + .column(0) + .as_primitive::(); + assert_eq!(values.value(0), 10); + assert_eq!(values.value(1), 20); + + let parts = batch.column(1).as_string::(); + assert_eq!(parts.value(0), "west"); + assert_eq!(parts.value(1), "west"); + + let file_names = batch.column(2).as_string::(); + assert_eq!(file_names.value(0), "part=west/data.csv"); + assert_eq!(file_names.value(1), "part=west/data.csv"); + } + #[test] fn test_split_projection_with_partition_columns() { use arrow::array::AsArray; diff --git a/datafusion/datasource/src/proto.rs b/datafusion/datasource/src/proto.rs new file mode 100644 index 0000000000000..6dc6e2ee45ffd --- /dev/null +++ b/datafusion/datasource/src/proto.rs @@ -0,0 +1,303 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Protobuf conversions for the file-scan leaf types owned by this crate: +//! [`FileRange`], [`PartitionedFile`] and [`FileGroup`]. +//! +//! These are the single copy of that wire logic, used both by the central +//! serializer in `datafusion-proto` and by the per-source `try_to_proto` hooks, +//! so the format cannot drift between them. +//! +//! None of these conversions need a codec or an encode/decode context: every +//! field is plain data or goes through `datafusion-proto-common`. That is why +//! they are plain [`TryFrom`] impls rather than the `try_to_proto(ctx)` / +//! `try_from_proto(node, ctx)` hooks used for plans, expressions and scan +//! configs: the standard trait can express a conversion that takes nothing but +//! the value, and the orphan rule allows it here because one side of each +//! conversion is a type this crate owns. + +use std::sync::Arc; + +use chrono::{TimeZone, Utc}; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err}; +use datafusion_proto_models::protobuf; +use object_store::ObjectMeta; +use object_store::path::Path; + +use crate::file_groups::FileGroup; +use crate::{FileRange, PartitionedFile}; + +impl TryFrom<&FileRange> for protobuf::FileRange { + type Error = DataFusionError; + + fn try_from(range: &FileRange) -> Result { + Ok(protobuf::FileRange { + start: range.start, + end: range.end, + }) + } +} + +impl TryFrom<&protobuf::FileRange> for FileRange { + type Error = DataFusionError; + + fn try_from(range: &protobuf::FileRange) -> Result { + Ok(FileRange { + start: range.start, + end: range.end, + }) + } +} + +impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile { + type Error = DataFusionError; + + fn try_from(file: &PartitionedFile) -> Result { + let last_modified = file.object_meta.last_modified; + let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { + DataFusionError::Plan(format!( + "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" + )) + })? as u64; + Ok(protobuf::PartitionedFile { + arrow_schema: file + .arrow_schema + .as_ref() + .map(|s| s.as_ref().try_into()) + .transpose()?, + path: file.object_meta.location.as_ref().to_owned(), + size: file.object_meta.size, + last_modified_ns, + partition_values: file + .partition_values + .iter() + .map(|v| v.try_into()) + .collect::, _>>()?, + range: file.range.as_ref().map(TryInto::try_into).transpose()?, + statistics: file.statistics.as_ref().map(|s| s.as_ref().into()), + }) + } +} + +impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile { + type Error = DataFusionError; + + fn try_from(file: &protobuf::PartitionedFile) -> Result { + let mut pf = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse(file.path.as_str()).map_err(|e| { + internal_datafusion_err!("Invalid object_store path: {e}") + })?, + last_modified: Utc.timestamp_nanos(file.last_modified_ns as i64), + size: file.size, + e_tag: None, + version: None, + }) + .with_partition_values( + file.partition_values + .iter() + .map(|v| v.try_into()) + .collect::, _>>()?, + ); + if let Some(proto_schema) = file.arrow_schema.as_ref() { + pf = pf.with_arrow_schema(Arc::new( + proto_schema.try_into().map_err(DataFusionError::from)?, + )); + } + if let Some(range) = file.range.as_ref() { + let range = FileRange::try_from(range)?; + pf = pf.with_range(range.start, range.end); + } + if let Some(proto_stats) = file.statistics.as_ref() { + // The wire format carries statistics for the full table schema (file + partition + // columns), so assign directly — `with_statistics` would append the partition + // column stats a second time. + pf.statistics = Some(Arc::new(proto_stats.try_into()?)); + } + Ok(pf) + } +} + +impl TryFrom<&FileGroup> for protobuf::FileGroup { + type Error = DataFusionError; + + fn try_from(group: &FileGroup) -> Result { + Ok(protobuf::FileGroup { + files: group + .files() + .iter() + .map(TryInto::try_into) + .collect::>>()?, + }) + } +} + +impl TryFrom<&protobuf::FileGroup> for FileGroup { + type Error = DataFusionError; + + fn try_from(group: &protobuf::FileGroup) -> Result { + Ok(FileGroup::new( + group + .files + .iter() + .map(TryInto::try_into) + .collect::>>()?, + )) + } +} + +#[cfg(test)] +mod tests { + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{ScalarValue, Statistics}; + + use super::*; + + #[test] + fn partitioned_file_roundtrip_preserves_all_fields() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let pf = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse("foo/bar.parquet")?, + last_modified: Utc.timestamp_nanos(1_000_000_000), + size: 1234, + e_tag: None, + version: None, + }) + .with_partition_values(vec![ScalarValue::from("2024-01-01")]) + .with_range(10, 20) + .with_arrow_schema(Arc::clone(&schema)) + .with_statistics(Arc::new(Statistics::new_unknown(&schema))); + + let encoded = protobuf::PartitionedFile::try_from(&pf)?; + let decoded = PartitionedFile::try_from(&encoded)?; + + assert_eq!(decoded.object_meta.location, pf.object_meta.location); + assert_eq!(decoded.object_meta.size, pf.object_meta.size); + assert_eq!( + decoded.object_meta.last_modified, + pf.object_meta.last_modified + ); + assert_eq!(decoded.partition_values, pf.partition_values); + assert_eq!(decoded.range, pf.range); + assert_eq!(decoded.arrow_schema.as_deref(), Some(schema.as_ref())); + // Statistics span the full table schema (file columns followed by one + // entry per partition column), and survive the round trip intact. + assert_eq!( + pf.statistics.as_ref().unwrap().column_statistics.len(), + schema.fields().len() + pf.partition_values.len() + ); + assert_eq!(decoded.statistics, pf.statistics); + Ok(()) + } + + #[test] + fn partitioned_file_path_roundtrip_percent_encoded() -> Result<()> { + // The wire format carries the *encoded* path, so a location that already + // contains percent escapes must survive without a second round of + // encoding or decoding. + let path_str = "foo/foo%2Fbar/baz%252Fqux"; + let pf = PartitionedFile::new_from_meta(ObjectMeta { + location: Path::parse(path_str)?, + last_modified: Utc.timestamp_nanos(1_000), + size: 42, + e_tag: None, + version: None, + }); + + let encoded = protobuf::PartitionedFile::try_from(&pf)?; + assert_eq!(encoded.path, path_str); + + let decoded = PartitionedFile::try_from(&encoded)?; + assert_eq!(decoded.object_meta.location.as_ref(), path_str); + assert_eq!(decoded.object_meta.location, pf.object_meta.location); + Ok(()) + } + + #[test] + fn partitioned_file_arrow_schema_roundtrip_preserves_metadata() -> Result<()> { + use std::collections::HashMap; + + let arrow_schema = Arc::new(Schema::new_with_metadata( + vec![ + Field::new("id", DataType::Int64, false), + Field::new("value", DataType::Utf8, true).with_metadata(HashMap::from([ + ("field_meta".to_string(), "field_value".to_string()), + ])), + ], + HashMap::from([("schema_meta".to_string(), "schema_value".to_string())]), + )); + let pf = PartitionedFile::new("foo/bar.parquet", 10) + .with_arrow_schema(Arc::clone(&arrow_schema)); + + let encoded = protobuf::PartitionedFile::try_from(&pf)?; + assert!(encoded.arrow_schema.is_some()); + + let decoded = PartitionedFile::try_from(&encoded)?; + assert_eq!(decoded.arrow_schema.as_deref(), Some(arrow_schema.as_ref())); + Ok(()) + } + + #[test] + fn partitioned_file_from_proto_rejects_invalid_path() { + let proto = protobuf::PartitionedFile { + path: "foo//bar.parquet".to_string(), + ..Default::default() + }; + + let err = PartitionedFile::try_from(&proto).unwrap_err(); + assert!( + err.to_string().contains("Invalid object_store path"), + "unexpected error: {err}" + ); + } + + #[test] + fn file_group_from_slice_matches_file_group() -> Result<()> { + // `protobuf::FileGroup: TryFrom<&[T]>` lives in `datafusion-proto-models`, + // generic over the element so that crate never names `PartitionedFile`. + // This is the caller-visible half: the bound resolves via + // `TryFrom<&PartitionedFile> for protobuf::PartitionedFile` above. + let files = vec![ + PartitionedFile::new("a.parquet", 1), + PartitionedFile::new("b.parquet", 2), + ]; + + let from_slice = protobuf::FileGroup::try_from(&files[..])?; + let from_group = protobuf::FileGroup::try_from(&FileGroup::new(files))?; + + assert_eq!(from_slice, from_group); + assert_eq!(from_slice.files.len(), 2); + Ok(()) + } + + #[test] + fn file_group_roundtrip() -> Result<()> { + let group = FileGroup::new(vec![ + PartitionedFile::new("a.parquet", 1), + PartitionedFile::new("b.parquet", 2), + ]); + + let encoded = protobuf::FileGroup::try_from(&group)?; + let decoded = FileGroup::try_from(&encoded)?; + + assert_eq!(decoded.len(), 2); + assert_eq!( + decoded.files()[1].object_meta.location, + group.files()[1].object_meta.location + ); + Ok(()) + } +} diff --git a/datafusion/datasource/src/sink.rs b/datafusion/datasource/src/sink.rs index 2a1f5c4a2fd02..4bf04133b7843 100644 --- a/datafusion/datasource/src/sink.rs +++ b/datafusion/datasource/src/sink.rs @@ -32,8 +32,9 @@ use datafusion_physical_expr_common::sort_expr::{LexRequirement, OrderingRequire use datafusion_physical_plan::metrics::MetricsSet; use datafusion_physical_plan::stream::RecordBatchStreamAdapter; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, Partitioning, - PlanProperties, SendableRecordBatchStream, execute_input_stream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, Partitioning, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, execute_input_stream, }; use async_trait::async_trait; @@ -71,6 +72,21 @@ pub trait DataSink: Any + DisplayAs + Debug + Send + Sync { data: SendableRecordBatchStream, context: &Arc, ) -> Result; + + /// Serialize this sink into a full protobuf plan node, if it knows how. + /// + /// Implementations can use `ctx` to encode the input plan, sink-specific + /// expressions, and [`DataSinkExec::encode_sort_order`]. + /// + /// Returning `Ok(None)` lets the caller try its extension codec instead. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _exec: &DataSinkExec, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } impl dyn DataSink { @@ -145,6 +161,75 @@ impl DataSinkExec { &self.sort_order } + /// Encode the optional sink ordering for a protobuf plan node. + #[cfg(feature = "proto")] + pub fn encode_sort_order( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> + { + use datafusion_physical_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + + self.sort_order + .as_ref() + .map(|requirements| { + requirements + .iter() + .map(|requirement| { + let expr: PhysicalSortExpr = requirement.to_owned().into(); + Ok(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&expr.expr)?)), + asc: !expr.options.descending, + nulls_first: expr.options.nulls_first, + }) + }) + .collect::>>() + .map(|physical_sort_expr_nodes| { + protobuf::PhysicalSortExprNodeCollection { + physical_sort_expr_nodes, + } + }) + }) + .transpose() + } + + /// Decode the optional sink ordering from a protobuf plan node. + #[cfg(feature = "proto")] + pub fn decode_sort_order( + collection: Option< + &datafusion_proto_models::protobuf::PhysicalSortExprNodeCollection, + >, + ctx: &datafusion_physical_plan::proto::ExecutionPlanDecodeCtx<'_>, + schema: &Schema, + ) -> Result> { + use arrow::compute::SortOptions; + use datafusion_physical_expr::PhysicalSortExpr; + + let Some(collection) = collection else { + return Ok(None); + }; + let sort_exprs = collection + .physical_sort_expr_nodes + .iter() + .map(|node| { + let expr = node.expr.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "Unexpected empty physical expression" + ) + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr, schema)?, + options: SortOptions { + descending: !node.asc, + nulls_first: node.nulls_first, + }, + }) + }) + .collect::>>()?; + Ok(LexRequirement::new(sort_exprs.into_iter().map(Into::into))) + } + fn create_schema( input: &Arc, schema: SchemaRef, @@ -190,9 +275,16 @@ impl ExecutionPlan for DataSinkExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { // DataSink is responsible for dynamically partitioning its // own input at execution time, and so requires a single input partition. - vec![Distribution::SinglePartition; self.children().len()] + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition; + self.children().len() + ]) } fn required_input_ordering(&self) -> Vec> { @@ -213,9 +305,10 @@ impl ExecutionPlan for DataSinkExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self::new( Arc::clone(&children[0]), @@ -224,16 +317,20 @@ impl ExecutionPlan for DataSinkExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to sort order requirements if present - if let Some(sort_order) = &self.sort_order { - for req in sort_order.iter() { - f(req.expr.as_ref())?; - } - } Ok(TreeNodeRecursion::Continue) } @@ -274,6 +371,15 @@ impl ExecutionPlan for DataSinkExec { fn metrics(&self) -> Option { self.sink.metrics() } + + /// Delegates protobuf serialization to the underlying sink. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + self.sink().try_to_proto(self, ctx) + } } /// Create a output record batch with a count diff --git a/datafusion/datasource/src/source.rs b/datafusion/datasource/src/source.rs index 420c6b508ce4f..741010c595197 100644 --- a/datafusion/datasource/src/source.rs +++ b/datafusion/datasource/src/source.rs @@ -33,7 +33,8 @@ use datafusion_physical_plan::metrics::{ use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::stream::BatchSplitStream; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, }; use itertools::Itertools; @@ -46,6 +47,7 @@ use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_plan::SortOrderPushdownResult; +use datafusion_physical_plan::StatisticsArgs; use datafusion_physical_plan::filter_pushdown::{ ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown, }; @@ -230,15 +232,15 @@ pub trait DataSource: Any + Send + Sync + Debug { /// This includes filter predicates (which may contain dynamic filters) and any /// other expressions used during data scanning. /// - /// Implementations must override this method. If the data source has no expressions, - /// return `Ok(TreeNodeRecursion::Continue)` immediately. + /// The function `f` should be called once per expression unless the function returns + /// [`TreeNodeRecursion::Stop`] to stop iteration. /// /// See [`ExecutionPlan::apply_expressions`] for more details and implementation examples. /// /// [`ExecutionPlan::apply_expressions`]: datafusion_physical_plan::ExecutionPlan::apply_expressions fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result; /// Injects arbitrary run-time state into this DataSource, returning a new instance @@ -260,9 +262,16 @@ pub trait DataSource: Any + Send + Sync + Debug { /// Create per execution state to share across sibling instances of this /// data source during one execution. /// + /// `config` is the session configuration, so implementations can honor + /// options that disable sibling sharing (returning `None`) for consumers + /// that cannot poll all partitions in one process. + /// /// Returns `None` (the default) if this data source has /// no sibling-shared execution state. - fn create_sibling_state(&self) -> Option> { + fn create_sibling_state( + &self, + _config: &ConfigOptions, + ) -> Option> { None } @@ -273,6 +282,30 @@ pub trait DataSource: Any + Send + Sync + Debug { fn open_with_args(&self, args: OpenArgs) -> Result { self.open(args.partition, args.context) } + + /// Serialize this data source to a full [`PhysicalPlanNode`] (a + /// `DataSourceExec` wrapping this source), if it knows how. + /// + /// This is the `DataSource` analog of + /// [`ExecutionPlan::try_to_proto`]. + /// [`DataSourceExec::try_to_proto`](crate::source::DataSourceExec) delegates + /// to this hook, which for file scans forwards to + /// [`FileSource::try_to_proto`] + /// through the shared [`FileScanConfig`] + /// spine. + /// + /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller falls + /// back to the central downcast chain in `datafusion-proto`. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } } /// Arguments for [`DataSource::open_with_args`] @@ -368,19 +401,30 @@ impl ExecutionPlan for DataSourceExec { Vec::new() } - fn apply_expressions( - &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, - ) -> Result { - // Delegate to the underlying data source - self.data_source.apply_expressions(f) + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) } fn with_new_children( self: Arc, - _: Vec>, + children: Vec>, ) -> Result> { - Ok(self) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + // Delegate to the underlying data source + self.data_source.apply_expressions(f) } /// Implementation of [`ExecutionPlan::repartitioned`] which relies upon the inner [`DataSource::repartitioned`]. @@ -416,7 +460,10 @@ impl ExecutionPlan for DataSourceExec { ) -> Result { let shared_state = self .execution_state - .get_or_init(|| self.data_source.create_sibling_state()) + .get_or_init(|| { + self.data_source + .create_sibling_state(context.session_config().options()) + }) .clone(); let args = OpenArgs::new(partition, Arc::clone(&context)) .with_shared_state(shared_state); @@ -451,8 +498,12 @@ impl ExecutionPlan for DataSourceExec { Some(metrics) } - fn partition_statistics(&self, partition: Option) -> Result> { - self.data_source.partition_statistics(partition) + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + self.data_source.partition_statistics(args.partition()) } fn with_fetch(&self, limit: Option) -> Option> { @@ -563,6 +614,18 @@ impl ExecutionPlan for DataSourceExec { new_exec.execution_state = Arc::new(OnceLock::new()); Ok(Arc::new(new_exec)) } + + /// Delegates serialization to the wrapped [`DataSource`]. For file scans the + /// concrete [`FileSource`] emits the node via its + /// own `try_to_proto` hook, keeping the format-specific wire logic in the + /// format crate. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_plan::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + self.data_source().try_to_proto(ctx) + } } impl DataSourceExec { diff --git a/datafusion/datasource/src/statistics.rs b/datafusion/datasource/src/statistics.rs index 6abfafe9d39d4..491b4e54e8993 100644 --- a/datafusion/datasource/src/statistics.rs +++ b/datafusion/datasource/src/statistics.rs @@ -97,8 +97,16 @@ impl MinMaxStatistics { .zip(s.column_statistics[i].max_value.get_value().cloned()) .ok_or_else(|| plan_datafusion_err!("statistics not found")) } else { - let partition_value = &pv[i - s.column_statistics.len()]; - Ok((partition_value.clone(), partition_value.clone())) + if let Some(partition_value) = + pv.get(i - s.column_statistics.len()) + { + Ok((partition_value.clone(), partition_value.clone())) + } else { + Err(plan_datafusion_err!( + "statistics not found for partition, expected at most {}", + s.column_statistics.len() + )) + } } }) .collect::>>()? @@ -541,14 +549,6 @@ pub fn compute_all_files_statistics( Ok((file_groups_with_stats, statistics)) } -#[deprecated(since = "47.0.0", note = "Use Statistics::add")] -pub fn add_row_stats( - file_num_rows: Precision, - num_rows: Precision, -) -> Precision { - file_num_rows.add(&num_rows) -} - #[cfg(test)] mod tests { use super::*; @@ -890,4 +890,31 @@ mod tests { Ok(()) } + + #[test] + fn min_max_statistics_missing_column_stats_returns_error() { + let schema = test_schema(); + let sort_order = + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + let files = [ + file_with_stats("f1.parquet", Statistics::default()), + file_with_stats("f2.parquet", Statistics::default()), + ]; + + let err = match MinMaxStatistics::new_from_files( + &sort_order, + &schema, + None, + files.iter(), + ) { + Ok(_) => panic!("expected missing statistics error"), + Err(err) => err, + }; + + assert!( + err.to_string() + .contains("statistics not found for partition"), + "unexpected error: {err:?}" + ); + } } diff --git a/datafusion/datasource/src/table_schema.rs b/datafusion/datasource/src/table_schema.rs index 5b7fc4727df05..f1cb86ed7413d 100644 --- a/datafusion/datasource/src/table_schema.rs +++ b/datafusion/datasource/src/table_schema.rs @@ -17,16 +17,23 @@ //! Helper struct to manage table schemas with partition columns -use arrow::datatypes::{FieldRef, SchemaBuilder, SchemaRef}; +use arrow::datatypes::{FieldRef, Fields, SchemaBuilder, SchemaRef}; use std::sync::Arc; /// The overall schema for potentially partitioned data sources. /// /// When reading partitioned data (such as Hive-style partitioning), a [`TableSchema`] -/// consists of two parts: +/// consists of up to three parts: /// 1. **File schema**: The schema of the actual data files on disk /// 2. **Partition columns**: Columns whose values are encoded in the directory structure, /// but not stored in the files themselves +/// 3. **Virtual columns**: Columns produced by the file reader (e.g. Parquet +/// `row_number`) that are not stored in the files +/// +/// The full table schema is composed in that order: file columns, then +/// partition columns, then virtual columns. Consumers that need a different +/// output ordering should use a projection on top of +/// [`TableSchema::table_schema`]. /// /// # Example: Partitioned Table /// @@ -70,30 +77,47 @@ pub struct TableSchema { /// /// These columns are NOT present in the data files but are appended to each /// row during query execution based on the file's location. - table_partition_cols: Arc>, + /// + /// Stored as [`Fields`] (an immutable `Arc<[FieldRef]>`) so that cloning a + /// `TableSchema` is cheap and the partition columns can be shared zero-copy + /// with an existing schema. + table_partition_cols: Fields, + + /// Virtual columns that are generated by the reader rather than read from + /// the data files or the directory structure. + /// + /// For example, a Parquet reader may inject a `row_number` column whose + /// values are produced per file by the reader. Virtual column fields must + /// carry an arrow extension type (e.g. `RowNumber`, `RowGroupIndex`) so the + /// file reader can recognize them. + /// + /// Virtual columns are appended at the end of the table schema, after the + /// file columns and any partition columns (layout: `[file, partition, + /// virtual]`). + virtual_columns: Fields, - /// The complete table schema: file_schema columns followed by partition columns. + /// The complete table schema: file_schema columns, followed by partition + /// columns, followed by virtual columns. /// - /// This is pre-computed during construction by concatenating `file_schema` - /// and `table_partition_cols`, so it can be returned as a cheap reference. + /// This is pre-computed during construction by concatenating the three + /// parts, so it can be returned as a cheap reference. table_schema: SchemaRef, + + /// Schema of file + partition columns, excluding virtual columns. + /// + /// Pre-computed during construction so [`Self::schema_without_virtual_columns`] + /// can return a cheap reference. When there are no virtual columns this + /// shares the same `Arc` as `table_schema`. + schema_without_virtual_columns: SchemaRef, } impl TableSchema { - /// Create a new TableSchema from a file schema and partition columns. - /// - /// The table schema is automatically computed by appending the partition columns - /// to the file schema. + /// Start building a [`TableSchema`] from its (required) file schema. /// - /// You should prefer calling this method over - /// chaining [`TableSchema::from_file_schema`] and [`TableSchema::with_table_partition_cols`] - /// if you have both the file schema and partition columns available at construction time - /// since it avoids re-computing the table schema. - /// - /// # Arguments - /// - /// * `file_schema` - Schema of the data files (without partition columns) - /// * `table_partition_cols` - Partition columns to append to each row + /// Partition columns are optional and added with + /// [`TableSchemaBuilder::with_table_partition_cols`]; the full table schema + /// is computed once by [`TableSchemaBuilder::build`]. This is the preferred + /// way to construct a `TableSchema`. /// /// # Example /// @@ -106,53 +130,54 @@ impl TableSchema { /// Field::new("amount", DataType::Float64, false), /// ])); /// - /// let partition_cols = vec![ - /// Arc::new(Field::new("date", DataType::Utf8, false)), - /// Arc::new(Field::new("region", DataType::Utf8, false)), - /// ]; - /// - /// let table_schema = TableSchema::new(file_schema, partition_cols); + /// let table_schema = TableSchema::builder(file_schema) + /// .with_table_partition_cols(vec![ + /// Arc::new(Field::new("date", DataType::Utf8, false)), + /// Arc::new(Field::new("region", DataType::Utf8, false)), + /// ]) + /// .build(); /// /// // Table schema will have 4 columns: user_id, amount, date, region /// assert_eq!(table_schema.table_schema().fields().len(), 4); /// ``` + pub fn builder(file_schema: SchemaRef) -> TableSchemaBuilder { + TableSchemaBuilder::new(file_schema) + } + + /// Create a new TableSchema from a file schema and partition columns. + /// + /// This is a convenience for + /// `TableSchema::builder(file_schema).with_table_partition_cols(cols).build()`. + #[deprecated( + since = "55.0.0", + note = "use TableSchema::builder(file_schema).with_table_partition_cols(cols).build() (or TableSchema::from(file_schema) for no partition columns)" + )] pub fn new(file_schema: SchemaRef, table_partition_cols: Vec) -> Self { - let mut builder = SchemaBuilder::from(file_schema.as_ref()); - builder.extend(table_partition_cols.iter().cloned()); - Self { - file_schema, - table_partition_cols: Arc::new(table_partition_cols), - table_schema: Arc::new(builder.finish()), - } + TableSchemaBuilder::new(file_schema) + .with_table_partition_cols(table_partition_cols) + .build() } /// Create a new TableSchema with no partition columns. - /// - /// You should prefer calling [`TableSchema::new`] if you have partition columns at - /// construction time since it avoids re-computing the table schema. + #[deprecated( + since = "55.0.0", + note = "use TableSchema::from(file_schema) / file_schema.into()" + )] pub fn from_file_schema(file_schema: SchemaRef) -> Self { - Self::new(file_schema, vec![]) + TableSchemaBuilder::new(file_schema).build() } - /// Add partition columns to an existing TableSchema, returning a new instance. - /// - /// You should prefer calling [`TableSchema::new`] instead of chaining [`TableSchema::from_file_schema`] - /// into [`TableSchema::with_table_partition_cols`] if you have partition columns at construction time - /// since it avoids re-computing the table schema. - pub fn with_table_partition_cols(mut self, partition_cols: Vec) -> Self { - if self.table_partition_cols.is_empty() { - self.table_partition_cols = Arc::new(partition_cols); - } else { - // Append to existing partition columns - let table_partition_cols = Arc::get_mut(&mut self.table_partition_cols).expect( - "Expected to be the sole owner of table_partition_cols since this function accepts mut self", - ); - table_partition_cols.extend(partition_cols); - } - let mut builder = SchemaBuilder::from(self.file_schema.as_ref()); - builder.extend(self.table_partition_cols.iter().cloned()); - self.table_schema = Arc::new(builder.finish()); - self + /// Return a new `TableSchema` with `partition_cols` as its partition columns, + /// replacing any existing ones. Existing virtual columns are preserved. + #[deprecated( + since = "55.0.0", + note = "use TableSchema::builder(file_schema).with_table_partition_cols(cols).build()" + )] + pub fn with_table_partition_cols(self, partition_cols: Vec) -> Self { + TableSchemaBuilder::new(self.file_schema) + .with_table_partition_cols(partition_cols) + .with_virtual_columns(self.virtual_columns) + .build() } /// Get the file schema (without partition columns). @@ -166,28 +191,166 @@ impl TableSchema { /// /// These are the columns derived from the directory structure that /// will be appended to each row during query execution. - pub fn table_partition_cols(&self) -> &Vec { + pub fn table_partition_cols(&self) -> &Fields { &self.table_partition_cols } - /// Get the full table schema (file schema + partition columns). + /// Get the virtual columns. /// - /// This is the complete schema that will be seen by queries, combining - /// both the columns from the files and the partition columns. + /// Virtual columns are produced by the file reader (e.g. Parquet + /// `row_number`) and are not stored in the data files or derived from + /// partition paths. + pub fn virtual_columns(&self) -> &Fields { + &self.virtual_columns + } + + /// Get the full table schema (file schema + partition columns + virtual columns). + /// + /// This is the complete schema that will be seen by queries. Fields appear + /// in the order: file columns, partition columns, virtual columns. pub fn table_schema(&self) -> &SchemaRef { &self.table_schema } + + /// Schema of columns that can be referenced by predicates pushed into the + /// file reader: file columns plus partition columns, excluding virtual + /// columns. + /// + /// Virtual columns are produced by the reader itself (e.g. Parquet + /// `row_number`) and cannot be referenced inside the reader's row filter, + /// so predicates that reference them must stay above the scan. Callers + /// deciding which filters to push down should check against this schema + /// rather than [`Self::table_schema`]. + /// + /// When there are no virtual columns this returns the same schema as + /// [`Self::table_schema`]. + pub fn schema_without_virtual_columns(&self) -> &SchemaRef { + &self.schema_without_virtual_columns + } } impl From for TableSchema { fn from(schema: SchemaRef) -> Self { - Self::from_file_schema(schema) + TableSchemaBuilder::new(schema).build() + } +} + +impl From<&SchemaRef> for TableSchema { + fn from(schema: &SchemaRef) -> Self { + TableSchemaBuilder::new(Arc::clone(schema)).build() + } +} + +/// Builder for [`TableSchema`]. +/// +/// The file schema is the only required input; partition columns and virtual +/// columns are optional. Unlike calling [`TableSchema`]'s setters repeatedly, +/// the builder computes the concatenated table schema exactly once, in +/// [`TableSchemaBuilder::build`]. +/// +/// ``` +/// # use std::sync::Arc; +/// # use arrow::datatypes::{Schema, Field, DataType}; +/// # use datafusion_datasource::TableSchemaBuilder; +/// # let file_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int64, false)])); +/// let table_schema = TableSchemaBuilder::new(file_schema) +/// .with_table_partition_cols(vec![Arc::new(Field::new("date", DataType::Utf8, false))]) +/// .build(); +/// assert_eq!(table_schema.table_partition_cols().len(), 1); +/// ``` +#[derive(Debug, Clone)] +pub struct TableSchemaBuilder { + file_schema: SchemaRef, + table_partition_cols: Fields, + virtual_columns: Fields, +} + +impl TableSchemaBuilder { + /// Create a builder for a `TableSchema` over the given file schema, with no + /// partition or virtual columns yet. + pub fn new(file_schema: SchemaRef) -> Self { + Self { + file_schema, + table_partition_cols: Fields::empty(), + virtual_columns: Fields::empty(), + } + } + + /// Set the partition columns, replacing any previously set. + /// + /// Accepts anything convertible into [`Fields`] (e.g. `Vec` or an + /// existing schema's `Fields`, which is shared zero-copy). + pub fn with_table_partition_cols( + mut self, + table_partition_cols: impl Into, + ) -> Self { + self.table_partition_cols = table_partition_cols.into(); + self + } + + /// Set the virtual columns, replacing any previously set. + /// + /// Virtual columns are produced by the file reader (e.g. Parquet + /// `row_number`) and appended at the end of the table schema. Each field + /// must carry an arrow virtual extension type so the reader can recognize + /// it. + /// + /// Accepts anything convertible into [`Fields`] (e.g. `Vec`). + pub fn with_virtual_columns(mut self, virtual_columns: impl Into) -> Self { + self.virtual_columns = virtual_columns.into(); + self + } + + /// Build the [`TableSchema`], computing the full + /// `file + partition + virtual` schema once. + pub fn build(self) -> TableSchema { + debug_assert!( + self.virtual_columns.iter().enumerate().all(|(i, v)| { + let name = v.name(); + !self.file_schema.fields().iter().any(|f| f.name() == name) + && !self.table_partition_cols.iter().any(|p| p.name() == name) + && !self.virtual_columns[..i].iter().any(|w| w.name() == name) + }), + "virtual column name collides with an existing file, partition, or virtual column" + ); + + let mut builder = SchemaBuilder::from(self.file_schema.as_ref()); + builder.extend(self.table_partition_cols.iter().cloned()); + let (table_schema, schema_without_virtual_columns) = + if self.virtual_columns.is_empty() { + let schema = Arc::new(builder.finish()); + (Arc::clone(&schema), schema) + } else { + let without_virtual = Arc::new(builder.finish()); + let mut builder = SchemaBuilder::from(without_virtual.as_ref()); + builder.extend(self.virtual_columns.iter().cloned()); + (Arc::new(builder.finish()), without_virtual) + }; + TableSchema { + file_schema: self.file_schema, + table_partition_cols: self.table_partition_cols, + virtual_columns: self.virtual_columns, + table_schema, + schema_without_virtual_columns, + } + } +} + +impl From for TableSchemaBuilder { + fn from(schema: SchemaRef) -> Self { + TableSchemaBuilder::new(schema) + } +} + +impl From<&SchemaRef> for TableSchemaBuilder { + fn from(schema: &SchemaRef) -> Self { + TableSchemaBuilder::new(Arc::clone(schema)) } } #[cfg(test)] mod tests { - use super::TableSchema; + use super::{TableSchema, TableSchemaBuilder}; use arrow::datatypes::{DataType, Field, Schema}; use std::sync::Arc; @@ -203,7 +366,9 @@ mod tests { Arc::new(Field::new("region", DataType::Utf8, false)), ]; - let table_schema = TableSchema::new(file_schema.clone(), partition_cols.clone()); + let table_schema = TableSchema::builder(file_schema.clone()) + .with_table_partition_cols(partition_cols.clone()) + .build(); // Verify file schema assert_eq!(table_schema.file_schema().as_ref(), file_schema.as_ref()); @@ -225,55 +390,221 @@ mod tests { } #[test] - fn test_add_multiple_partition_columns() { + fn test_builder_with_partition_cols() { let file_schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); - let initial_partition_cols = - vec![Arc::new(Field::new("country", DataType::Utf8, false))]; + let table_schema = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![ + Arc::new(Field::new("country", DataType::Utf8, false)), + Arc::new(Field::new("year", DataType::Int32, false)), + ]) + .build(); - let table_schema = TableSchema::new(file_schema.clone(), initial_partition_cols); + // File schema is preserved and the partition columns are appended. + assert_eq!(table_schema.file_schema().as_ref(), file_schema.as_ref()); + assert_eq!(table_schema.table_partition_cols().len(), 2); + assert_eq!(table_schema.table_partition_cols()[0].name(), "country"); + assert_eq!(table_schema.table_partition_cols()[1].name(), "year"); - let additional_partition_cols = vec![ - Arc::new(Field::new("city", DataType::Utf8, false)), - Arc::new(Field::new("year", DataType::Int32, false)), - ]; + let expected_schema = Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("country", DataType::Utf8, false), + Field::new("year", DataType::Int32, false), + ]); + assert_eq!(table_schema.table_schema().as_ref(), &expected_schema); + } - let updated_table_schema = - table_schema.with_table_partition_cols(additional_partition_cols); + #[test] + fn test_builder_with_table_partition_cols_replaces() { + // Calling the setter more than once replaces rather than appends. + let file_schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); - // Verify file schema remains unchanged - assert_eq!( - updated_table_schema.file_schema().as_ref(), - file_schema.as_ref() - ); + let table_schema = TableSchemaBuilder::new(file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "country", + DataType::Utf8, + false, + ))]) + .with_table_partition_cols(vec![Arc::new(Field::new( + "city", + DataType::Utf8, + false, + ))]) + .build(); - // Verify partition columns - assert_eq!(updated_table_schema.table_partition_cols().len(), 3); - assert_eq!( - updated_table_schema.table_partition_cols()[0].name(), - "country" - ); - assert_eq!( - updated_table_schema.table_partition_cols()[1].name(), - "city" - ); - assert_eq!( - updated_table_schema.table_partition_cols()[2].name(), - "year" - ); + assert_eq!(table_schema.table_partition_cols().len(), 1); + assert_eq!(table_schema.table_partition_cols()[0].name(), "city"); + } - // Verify full table schema - let expected_fields = vec![ - Field::new("id", DataType::Int32, false), - Field::new("country", DataType::Utf8, false), - Field::new("city", DataType::Utf8, false), - Field::new("year", DataType::Int32, false), - ]; - let expected_schema = Schema::new(expected_fields); - assert_eq!( - updated_table_schema.table_schema().as_ref(), - &expected_schema - ); + #[test] + fn test_builder_accepts_fields_zero_copy() { + // `with_table_partition_cols` accepts an existing schema's `Fields` + // directly (shared via `Arc`, no `Vec` round-trip). + let file_schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let partition_schema = + Schema::new(vec![Field::new("date", DataType::Utf8, false)]); + + let table_schema = TableSchemaBuilder::new(file_schema) + .with_table_partition_cols(partition_schema.fields().clone()) + .build(); + + assert_eq!(table_schema.table_partition_cols().len(), 1); + assert_eq!(table_schema.table_partition_cols()[0].name(), "date"); + } + + #[test] + #[expect(deprecated)] + fn test_deprecated_with_table_partition_cols_replaces() { + // The deprecated setter still works and replaces the partition columns. + // It is safe on a shared clone because partition columns are immutable. + let file_schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)])); + let original = TableSchema::builder(file_schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "country", + DataType::Utf8, + false, + ))]) + .build(); + + let replaced = + original + .clone() + .with_table_partition_cols(vec![Arc::new(Field::new( + "city", + DataType::Utf8, + false, + ))]); + + assert_eq!(replaced.table_partition_cols().len(), 1); + assert_eq!(replaced.table_partition_cols()[0].name(), "city"); + + // The original is untouched. + assert_eq!(original.table_partition_cols().len(), 1); + assert_eq!(original.table_partition_cols()[0].name(), "country"); + } + + #[test] + fn test_builder_with_virtual_columns_layout() { + let file_schema = Arc::new(Schema::new(vec![ + Field::new("user_id", DataType::Int64, false), + Field::new("amount", DataType::Float64, false), + ])); + + let virtual_cols = + vec![Arc::new(Field::new("row_number", DataType::Int64, true))]; + + let partition_cols = vec![Arc::new(Field::new("date", DataType::Utf8, false))]; + + // Apply virtual columns and partition columns in either order on the + // builder; the resulting table schema should always be + // [file, partition, virtual]. + let built_virtual_first = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_virtual_columns(virtual_cols.clone()) + .with_table_partition_cols(partition_cols.clone()) + .build(); + + let built_partition_first = TableSchemaBuilder::new(Arc::clone(&file_schema)) + .with_table_partition_cols(partition_cols.clone()) + .with_virtual_columns(virtual_cols.clone()) + .build(); + + let expected = Schema::new(vec![ + Field::new("user_id", DataType::Int64, false), + Field::new("amount", DataType::Float64, false), + Field::new("date", DataType::Utf8, false), + Field::new("row_number", DataType::Int64, true), + ]); + + for ts in [built_virtual_first, built_partition_first] { + assert_eq!(ts.table_schema().as_ref(), &expected); + assert_eq!(ts.virtual_columns().len(), 1); + assert_eq!(ts.virtual_columns()[0].name(), "row_number"); + assert_eq!(ts.table_partition_cols().len(), 1); + assert_eq!(ts.file_schema().fields().len(), 2); + } + } + + #[test] + #[should_panic(expected = "virtual column name collides")] + #[cfg(debug_assertions)] + fn test_virtual_column_collides_with_file_schema_panics_in_debug() { + let file_schema = Arc::new(Schema::new(vec![Field::new( + "row_number", + DataType::Int64, + false, + )])); + let _ = TableSchemaBuilder::new(file_schema) + .with_virtual_columns(vec![Arc::new(Field::new( + "row_number", + DataType::Int64, + true, + ))]) + .build(); + } + + #[test] + #[should_panic(expected = "virtual column name collides")] + #[cfg(debug_assertions)] + fn test_virtual_column_collides_with_partition_panics_in_debug() { + let file_schema = Arc::new(Schema::new(vec![Field::new( + "user_id", + DataType::Int64, + false, + )])); + let partition_cols = + vec![Arc::new(Field::new("row_number", DataType::Utf8, false))]; + let _ = TableSchemaBuilder::new(file_schema) + .with_table_partition_cols(partition_cols) + .with_virtual_columns(vec![Arc::new(Field::new( + "row_number", + DataType::Int64, + true, + ))]) + .build(); + } + + #[test] + #[should_panic(expected = "virtual column name collides")] + #[cfg(debug_assertions)] + fn test_duplicate_virtual_columns_panic_in_debug() { + let file_schema = Arc::new(Schema::new(vec![Field::new( + "user_id", + DataType::Int64, + false, + )])); + let _ = TableSchemaBuilder::new(file_schema) + .with_virtual_columns(vec![ + Arc::new(Field::new("vc", DataType::Int64, true)), + Arc::new(Field::new("vc", DataType::Int64, true)), + ]) + .build(); + } + + #[test] + #[should_panic(expected = "virtual column name collides")] + #[cfg(debug_assertions)] + fn test_partition_column_added_after_colliding_virtual_panics_in_debug() { + // Builder order is irrelevant: collision check runs in build(). + let file_schema = Arc::new(Schema::new(vec![Field::new( + "user_id", + DataType::Int64, + false, + )])); + let _ = TableSchemaBuilder::new(file_schema) + .with_virtual_columns(vec![Arc::new(Field::new( + "row_number", + DataType::Int64, + true, + ))]) + .with_table_partition_cols(vec![Arc::new(Field::new( + "row_number", + DataType::Utf8, + false, + ))]) + .build(); } } diff --git a/datafusion/datasource/src/test_util.rs b/datafusion/datasource/src/test_util.rs index b59ce58a420a8..20dfae5b3ac79 100644 --- a/datafusion/datasource/src/test_util.rs +++ b/datafusion/datasource/src/test_util.rs @@ -40,7 +40,7 @@ pub(crate) struct MockSource { impl Default for MockSource { fn default() -> Self { let table_schema = - crate::table_schema::TableSchema::new(Arc::new(Schema::empty()), vec![]); + crate::table_schema::TableSchema::from(Arc::new(Schema::empty())); Self { metrics: ExecutionPlanMetricsSet::new(), filter: None, @@ -128,7 +128,7 @@ impl FileSource for MockSource { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -138,3 +138,32 @@ impl FileSource for MockSource { pub(crate) fn col(name: &str, schema: &Schema) -> Result> { Ok(Arc::new(Column::new_with_schema(name, schema)?)) } + +/// Chunk sizes exercised by every parameterised test. +/// +/// `usize::MAX` is intentionally included: `ChunkedStore` treats it as +/// "one chunk containing everything", giving the single-chunk fast path. +pub(crate) const CHUNK_SIZES: &[usize] = &[1, 2, 3, 4, 5, 7, 8, 11, 13, 16, usize::MAX]; + +/// Seed a fresh `InMemory` store with `data` and wrap it in a +/// [`ChunkedStore`] that splits every GET response into `chunk_size`-byte +/// pieces. +pub(crate) async fn make_chunked_store( + data: &[u8], + chunk_size: usize, +) -> (Arc, object_store::path::Path) { + use bytes::Bytes; + use object_store::ObjectStoreExt; + use object_store::PutPayload; + use object_store::chunked::ChunkedStore; + use object_store::memory::InMemory; + use object_store::path::Path; + + let inner = Arc::new(InMemory::new()); + let path = Path::from("test"); + inner + .put(&path, PutPayload::from(Bytes::copy_from_slice(data))) + .await + .unwrap(); + (Arc::new(ChunkedStore::new(inner, chunk_size)), path) +} diff --git a/datafusion/datasource/src/url.rs b/datafusion/datasource/src/url.rs index 14f9b2af0021d..9ac4c5f50d1f7 100644 --- a/datafusion/datasource/src/url.rs +++ b/datafusion/datasource/src/url.rs @@ -18,8 +18,8 @@ use std::sync::Arc; use datafusion_common::{DataFusionError, Result, TableReference}; -use datafusion_execution::cache::TableScopedPath; use datafusion_execution::cache::cache_manager::CachedFileList; +use datafusion_execution::cache::cache_manager::TableScopedPath; use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_session::Session; @@ -523,6 +523,7 @@ mod tests { }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; + use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList}; use object_store::{ CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, PutMultipartOptions, PutPayload, @@ -1191,6 +1192,10 @@ mod tests { &self.config } + fn catalog_list(&self) -> Arc { + Arc::new(EmptyCatalogProviderList) + } + async fn create_physical_plan( &self, _logical_plan: &LogicalPlan, @@ -1210,7 +1215,7 @@ mod tests { unimplemented!() } - fn higher_order_functions(&self) -> &HashMap> { + fn higher_order_functions(&self) -> &HashMap> { unimplemented!() } diff --git a/datafusion/datasource/src/write/demux.rs b/datafusion/datasource/src/write/demux.rs index acc6435acf371..1b3098d309789 100644 --- a/datafusion/datasource/src/write/demux.rs +++ b/datafusion/datasource/src/write/demux.rs @@ -153,9 +153,9 @@ async fn row_count_demuxer( ) -> Result<()> { let exec_options = &context.session_config().options().execution; - let max_rows_per_file = exec_options.soft_max_rows_per_output_file; - let max_buffered_batches = exec_options.max_buffered_batches_per_output_file; - let minimum_parallel_files = exec_options.minimum_parallel_output_files; + let max_rows_per_file = exec_options.soft_max_rows_per_output_file.get(); + let max_buffered_batches = exec_options.max_buffered_batches_per_output_file.get(); + let minimum_parallel_files = exec_options.minimum_parallel_output_files.get(); let mut part_idx = 0; let write_id = rand::distr::Alphanumeric.sample_string(&mut rand::rng(), 16); @@ -305,7 +305,8 @@ async fn hive_style_partitions_demuxer( let write_id = rand::distr::Alphanumeric.sample_string(&mut rand::rng(), 16); let exec_options = &context.session_config().options().execution; - let max_buffered_recordbatches = exec_options.max_buffered_batches_per_output_file; + let max_buffered_recordbatches = + exec_options.max_buffered_batches_per_output_file.get(); // To support non string partition col types, cast the type to &str first let mut value_map: HashMap, Sender> = HashMap::new(); diff --git a/datafusion/datasource/src/write/mod.rs b/datafusion/datasource/src/write/mod.rs index e8d2d17da8ee8..c8c85112f0396 100644 --- a/datafusion/datasource/src/write/mod.rs +++ b/datafusion/datasource/src/write/mod.rs @@ -75,22 +75,6 @@ pub trait BatchSerializer: Sync + Send { fn serialize(&self, batch: RecordBatch, initial: bool) -> Result; } -/// Returns an [`AsyncWrite`] which writes to the given object store location -/// with the specified compression. -/// -/// The writer will have a default buffer size as chosen by [`BufWriter::new`]. -/// -/// We drop the `AbortableWrite` struct and the writer will not try to cleanup on failure. -/// Users can configure automatic cleanup with their cloud provider. -#[deprecated(since = "48.0.0", note = "Use ObjectWriterBuilder::new(...) instead")] -pub async fn create_writer( - file_compression_type: FileCompressionType, - location: &Path, - object_store: Arc, -) -> Result> { - ObjectWriterBuilder::new(file_compression_type, location, object_store).build() -} - /// Converts table schema to writer schema, which may differ in the case /// of hive style partitioning where some columns are removed from the /// underlying files. diff --git a/datafusion/datasource/src/write/orchestration.rs b/datafusion/datasource/src/write/orchestration.rs index 39c91a1c0d676..f75671d950353 100644 --- a/datafusion/datasource/src/write/orchestration.rs +++ b/datafusion/datasource/src/write/orchestration.rs @@ -259,7 +259,7 @@ pub async fn spawn_writer_tasks_and_join( .execution .max_buffered_batches_per_output_file; - let (tx_file_bundle, rx_file_bundle) = mpsc::channel(rb_buffer_size / 2); + let (tx_file_bundle, rx_file_bundle) = mpsc::channel(rb_buffer_size.get() / 2); let (tx_row_cnt, rx_row_cnt) = tokio::sync::oneshot::channel(); let write_coordinator_task = SpawnedTask::spawn(async move { stateless_serialize_and_write_files(rx_file_bundle, tx_row_cnt).await diff --git a/datafusion/doc/src/lib.rs b/datafusion/doc/src/lib.rs index 591a5a62f3b20..11b63ff661f50 100644 --- a/datafusion/doc/src/lib.rs +++ b/datafusion/doc/src/lib.rs @@ -281,7 +281,7 @@ impl DocumentationBuilder { /// /// The argument is rendered like below if None is passed through: /// - /// ```text + /// ```text /// : /// The expression to operate on. Can be a constant, column, or function, and any combination of operators. /// ``` diff --git a/datafusion/doc/src/udf.rs b/datafusion/doc/src/udf.rs index d1f51d919478d..f88db631e60fd 100644 --- a/datafusion/doc/src/udf.rs +++ b/datafusion/doc/src/udf.rs @@ -84,6 +84,14 @@ pub mod scalar_doc_sections { r#"Apache DataFusion uses a [PCRE-like](https://en.wikibooks.org/wiki/Regular_Expressions/Perl-Compatible_Regular_Expressions) regular expression [syntax](https://docs.rs/regex/latest/regex/#syntax) (minus support for several features including look-around and backreferences). + +The following flags are optionally supported in functions: + - **i**: case-insensitive: letters match both upper and lower case + - **m**: multi-line mode: `^` and `$` match begin/end of line + - **s**: allow `.` to match `\n` + - **R**: enables CRLF mode: when multi-line mode is enabled, `\r\n` is used + - **U**: swap the meaning of `x*` and `x*?` + The following regular expression functions are supported:"#, ), }; diff --git a/datafusion/execution/Cargo.toml b/datafusion/execution/Cargo.toml index 06c84d8acb493..c9d4acd3644ba 100644 --- a/datafusion/execution/Cargo.toml +++ b/datafusion/execution/Cargo.toml @@ -55,6 +55,7 @@ sql = [] arrow = { workspace = true } arrow-buffer = { workspace = true } async-trait = { workspace = true } +bytes = { workspace = true } dashmap = { workspace = true } datafusion-common = { workspace = true, default-features = false } datafusion-expr = { workspace = true, default-features = false } @@ -64,9 +65,14 @@ log = { workspace = true } object_store = { workspace = true, features = ["fs"] } parking_lot = { workspace = true } parquet = { workspace = true, optional = true } +pin-project-lite = { workspace = true } rand = { workspace = true } tempfile = { workspace = true } +tokio = { workspace = true } +tokio-util = { workspace = true, features = ["io"] } url = { workspace = true } +[target.'cfg(not(target_arch = "wasm32"))'.dependencies] +tokio = { workspace = true, features = ["fs"] } [dev-dependencies] chrono = { workspace = true } diff --git a/datafusion/execution/src/async_stream.rs b/datafusion/execution/src/async_stream.rs new file mode 100644 index 0000000000000..7ca6ba4850cab --- /dev/null +++ b/datafusion/execution/src/async_stream.rs @@ -0,0 +1,796 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use futures::Stream; +use futures::future::FusedFuture; +use futures::stream::FusedStream; +use parking_lot::Mutex; +use pin_project_lite::pin_project; +use std::ops::DerefMut; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll}; + +/// Creates a [`Stream`] from an async generator function. +/// +/// The `generator` closure receives an [`Emitter`] and runs as an async +/// block. Each `emitter.emit(value).await` call suspends the generator and +/// produces the next item in the stream. The stream ends when the generator +/// future resolves. +/// +/// # Example +/// +/// ``` +/// use datafusion_execution::async_stream; +/// use futures::StreamExt; +/// +/// # #[tokio::main(flavor = "current_thread")] +/// # async fn main() { +/// let stream = async_stream(|mut emitter| async move { +/// for i in 0_i32..3 { +/// emitter.emit(i).await; +/// } +/// }); +/// +/// let values: Vec = stream.collect().await; +/// assert_eq!(values, vec![0, 1, 2]); +/// # } +/// ``` +pub fn async_stream>( + generator: impl FnOnce(Emitter) -> F, +) -> impl FusedStream { + let (emitter, receiver) = tx_rx(); + AsyncStream::new(receiver, generator(emitter)) +} + +/// Creates a fallible [`Stream`] from an async generator function. +/// +/// The `generator` closure receives a [`TryEmitter`] and runs as an +/// async block that returns `Result<(), E>`. Each `emitter.emit(value).await` +/// call suspends the generator and produces `Ok(value)` as the next stream +/// item. The `?` operator can be used inside the generator to short-circuit on +/// errors: the error is emitted as the final `Err(e)` item and the stream +/// ends. The stream also ends when the generator future resolves to `Ok(())`. +/// +/// # Example +/// +/// ``` +/// use datafusion_execution::async_try_stream; +/// use futures::StreamExt; +/// +/// # #[tokio::main(flavor = "current_thread")] +/// # async fn main() { +/// let stream = async_try_stream(|mut emitter| async move { +/// emitter.emit(1_i32).await; +/// emitter.emit(2_i32).await; +/// Err::<(), _>("something went wrong")?; +/// emitter.emit(3_i32).await; // never reached +/// Ok(()) +/// }); +/// +/// let values: Vec> = stream.collect().await; +/// assert_eq!(values, vec![Ok(1), Ok(2), Err("something went wrong")]); +/// # } +/// ``` +pub fn async_try_stream>>( + generator: impl FnOnce(TryEmitter) -> F, +) -> impl FusedStream> { + let (try_emitter, mut emitter, receiver) = try_tx_rx::(); + AsyncStream::new(receiver, async move { + if let Err(e) = generator(try_emitter).await { + // Fill the slot without suspending so this future completes in the same + // poll that yields `Err(e)`: the stream terminates immediately and the + // emitter state is dropped (a consumer may never poll again after an + // error, which would otherwise keep this future suspended inside `emit`) + emitter.set(Err(e)); + } + }) +} + +/// Creates an `Emitter`/`Receiver` pair +fn tx_rx() -> (Emitter, Receiver) { + let slot = Arc::new(Mutex::new(None)); + ( + Emitter { + slot: Arc::clone(&slot), + }, + Receiver { slot }, + ) +} + +/// Creates an `TryEmitter`/`Emitter`/`Receiver` triplet +#[expect( + clippy::type_complexity, + reason = "three-element tuple is clearer than an alias here" +)] +fn try_tx_rx() -> ( + TryEmitter, + Emitter>, + Receiver>, +) { + let slot = Arc::new(Mutex::new(None)); + ( + TryEmitter { + slot: Arc::clone(&slot), + }, + Emitter { + slot: Arc::clone(&slot), + }, + Receiver { slot }, + ) +} + +/// Value slot shared between [`Emitter`] and [`Receiver`]. +/// Use `Arc` to ensure the created `Stream` implementations +/// are both `Send` and `Sync`. +type SlotRef = Arc>>; + +/// A handle for emitting values from an [`async_stream`] generator. +/// +/// The generator closure receives an `Emitter` as its argument. +pub struct Emitter { + slot: SlotRef, +} + +/// A handle for emitting values from an [`async_try_stream`] generator. +/// +/// The generator closure receives a `TryEmitter` as its argument. +pub struct TryEmitter { + slot: SlotRef>, +} + +struct Receiver { + slot: SlotRef, +} + +impl Emitter { + /// Returns a `Future` that emits `value` as the next stream item. + /// + /// The returned future **must be awaited immediately**. On its first poll it + /// yields `Poll::Pending`, handing control back to the stream consumer so it + /// can observe the emitted value. On the next poll (triggered by the + /// consumer calling `poll_next` again) it completes with `Poll::Ready(())`, + /// resuming the generator. + /// + /// # Panics + /// + /// Panics if `emit` is called a second time before the previous future has + /// been awaited, because doing so would silently overwrite the unconsumed + /// value. + pub fn emit(&mut self, value: T) -> impl FusedFuture { + self.set(value); + Emit { done: false } + } + + /// Places `value` in the slot without suspending the generator. Only useful + /// as the very last action before the generator future completes, since + /// nothing yields control back to the consumer in between. + fn set(&mut self, value: T) { + let mut guard = self.slot.lock(); + match guard.deref_mut() { + Some(_) => panic!("Misuse: await was not called after calling emit"), + slot => *slot = Some(value), + } + } +} + +impl TryEmitter { + /// Emits `Ok(value)` as the next stream item and suspends the generator. + /// + /// Behaves identically to [`Emitter::emit`]: the returned future must be + /// awaited immediately and yields `Poll::Pending` on its first poll to + /// transfer control to the stream consumer. + /// + /// # Panics + /// + /// Panics if called before the previous emit future has been awaited. + pub fn emit(&mut self, value: T) -> impl FusedFuture { + let mut guard = self.slot.lock(); + match guard.deref_mut() { + Some(_) => panic!("Misuse: await was not called after calling emit"), + slot => *slot = Some(Ok::(value)), + } + + Emit { done: false } + } +} + +struct Emit { + done: bool, +} + +impl FusedFuture for Emit { + fn is_terminated(&self) -> bool { + self.done + } +} + +impl Future for Emit { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { + if !self.done { + self.done = true; + // Poll::Pending causes the generator to yield, returning control back to the + // calling Stream + Poll::Pending + } else { + Poll::Ready(()) + } + } +} + +pin_project! { + struct AsyncStream { + rx: Receiver, + done: bool, + #[pin] + generator: U, + } +} + +impl AsyncStream { + fn new(rx: Receiver, generator: U) -> AsyncStream { + AsyncStream { + rx, + done: false, + generator, + } + } +} + +impl FusedStream for AsyncStream +where + U: Future, +{ + fn is_terminated(&self) -> bool { + self.done + } +} + +impl Stream for AsyncStream +where + U: Future, +{ + type Item = T; + + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.project(); + + if *this.done { + return Poll::Ready(None); + } + + // The `Option::take` call below ensures the next time poll is called the slot is + // already set to None + debug_assert!(this.rx.slot.lock().is_none()); + let res = this.generator.poll(cx); + *this.done = res.is_ready(); + + match this.rx.slot.lock().take() { + // Generator filled slot -> return next stream item + Some(v) => Poll::Ready(Some(v)), + // Generator did not fill slot and completed -> return None to indicate end of stream + None if *this.done => Poll::Ready(None), + // Generator did not fill slot and not completed -> return Pending since some Future + // other than Emit returned Pending. + None => Poll::Pending, + } + } + + fn size_hint(&self) -> (usize, Option) { + if self.done { (0, Some(0)) } else { (0, None) } + } +} + +#[cfg(test)] +mod test { + use crate::async_stream::Emitter; + use crate::{async_stream, async_try_stream}; + use futures::stream::FusedStream; + use futures::{Stream, StreamExt, pin_mut}; + use std::assert_matches; + use std::pin::Pin; + use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::task::{Context, Poll}; + use tokio::sync::mpsc; + + #[tokio::test] + async fn noop_stream() { + let s = async_stream(|_: Emitter<()>| async {}); + pin_mut!(s); + + assert_eq!(s.next().await, None); + } + + #[tokio::test] + async fn empty_stream() { + let mut ran = false; + + { + let r = &mut ran; + let s = async_stream(|_: Emitter<()>| async { + *r = true; + println!("hello world!"); + }); + pin_mut!(s); + + assert_eq!(s.next().await, None); + } + + assert!(ran); + } + + #[tokio::test] + async fn emit_single_value() { + let s = async_stream(|mut emitter| async move { + emitter.emit("hello").await; + }); + + let values: Vec<_> = s.collect().await; + + assert_eq!(1, values.len()); + assert_eq!("hello", values[0]); + } + + #[tokio::test] + async fn fused() { + let s = async_stream(|mut emitter| async move { + emitter.emit("hello").await; + }); + pin_mut!(s); + + assert!(!s.is_terminated()); + assert_eq!(s.next().await, Some("hello")); + assert_eq!(s.next().await, None); + + assert!(s.is_terminated()); + // This should return None from now on + assert_eq!(s.next().await, None); + } + + #[tokio::test] + async fn emit_multi_value() { + let s = async_stream(|mut emitter| async move { + emitter.emit("hello").await; + emitter.emit("world").await; + emitter.emit("dizzy").await; + }); + + let values: Vec<_> = s.collect().await; + + assert_eq!(3, values.len()); + assert_eq!("hello", values[0]); + assert_eq!("world", values[1]); + assert_eq!("dizzy", values[2]); + } + + #[tokio::test] + #[should_panic = "await was not called after calling emit"] + async fn emit_without_await() { + let s = async_stream(|mut emitter| async move { + #[expect(clippy::let_underscore_future)] + { + let _ = emitter.emit("hello"); + let _ = emitter.emit("world"); + } + }); + + let _: Vec<_> = s.collect().await; + } + + #[tokio::test] + async fn unit_emit_in_select() { + use tokio::select; + + #[expect(clippy::unused_async)] + async fn do_stuff_async() {} + + let s = async_stream(|mut emitter| async move { + select! { + _ = do_stuff_async() => emitter.emit(()).await, + else => emitter.emit(()).await, + } + }); + + let values: Vec<_> = s.collect().await; + assert_eq!(values.len(), 1); + } + + #[tokio::test] + async fn emit_with_select() { + use tokio::select; + + #[expect(clippy::unused_async)] + async fn do_stuff_async() {} + #[expect(clippy::unused_async)] + async fn more_async_work() {} + + let s = async_stream(|mut emitter| async move { + select! { + _ = do_stuff_async() => emitter.emit("hey").await, + _ = more_async_work() => emitter.emit("hey").await, + else => emitter.emit("hey").await, + } + }); + + let values: Vec<_> = s.collect().await; + assert_eq!(values, vec!["hey"]); + } + + #[tokio::test] + async fn return_stream() { + fn build_stream() -> impl Stream { + async_stream(|mut emitter| async move { + emitter.emit(1).await; + emitter.emit(2).await; + emitter.emit(3).await; + }) + } + + let s = build_stream(); + + let values: Vec<_> = s.collect().await; + assert_eq!(3, values.len()); + assert_eq!(1, values[0]); + assert_eq!(2, values[1]); + assert_eq!(3, values[2]); + } + + #[tokio::test] + async fn consume_channel() { + let (tx, mut rx) = mpsc::channel(10); + + let s = async_stream(|mut emitter| async move { + while let Some(v) = rx.recv().await { + emitter.emit(v).await; + } + }); + + pin_mut!(s); + + for i in 0..3 { + assert_matches!(tx.send(i).await, Ok(_)); + assert_eq!(Some(i), s.next().await); + } + + drop(tx); + assert_eq!(None, s.next().await); + } + + #[tokio::test] + async fn borrow_self() { + struct Data(String); + + impl Data { + fn stream(&self) -> impl Stream + '_ { + async_stream(move |mut emitter| async move { + emitter.emit(&self.0[..]).await; + }) + } + } + + let data = Data("hello".to_string()); + let s = data.stream(); + pin_mut!(s); + + assert_eq!(Some("hello"), s.next().await); + } + + #[tokio::test] + async fn stream_in_stream() { + let s = async_stream(|mut emitter| async move { + let s = async_stream(|mut inner_emitter| async move { + for i in 0..3 { + inner_emitter.emit(i).await; + } + }); + + pin_mut!(s); + while let Some(v) = s.next().await { + emitter.emit(v).await; + } + }); + + let values: Vec<_> = s.collect().await; + assert_eq!(3, values.len()); + } + + // Demonstrates that capturing an outer Emitter inside an inner async_stream with a + // different item type is no longer undefined behaviour: the outer emitter writes to its own + // typed slot, so the inner stream never sees any values. The outer stream receives the + // "foo" strings instead because they land in its slot. + #[tokio::test] + async fn stream_in_stream_misuse() { + let s = async_stream(|mut emitter| async move { + let s = async_stream(|_inner_emitter: Emitter| async move { + for _i in 0..3 { + emitter.emit("foo").await; + } + }); + + pin_mut!(s); + while let Some(v) = s.next().await { + println!("{v}"); + } + }); + + let values: Vec<_> = s.collect().await; + assert_eq!(3, values.len()); + } + + #[tokio::test] + async fn emit_non_unpin_value() { + let s: Vec<_> = async_stream(|mut emitter| async move { + for i in 0..3 { + emitter.emit(async move { i }).await; + } + }) + .buffered(1) + .collect() + .await; + + assert_eq!(s, vec![0, 1, 2]); + } + + #[tokio::test] + async fn should_not_call_handler_function_if_not_polled() { + let _ = async_stream(|_: Emitter<()>| async move { + panic!("should not be called"); + }); + } + + #[tokio::test] + async fn should_not_continue_until_next_poll() { + let s = async_stream(|mut emitter| async move { + emitter.emit("hey").await; + panic!("make sure poll based and not push based"); + }); + pin_mut!(s); + let _ = s.next().await; + } + + #[test] + fn inner_try_stream() { + use tokio::select; + + #[expect(clippy::unused_async)] + async fn do_stuff_async() {} + + let _ = async_stream(|mut emitter| async move { + select! { + _ = do_stuff_async() => { + let another_s = async_try_stream(|mut inner_emitter| async move { + inner_emitter.emit(()).await; + Ok(()) + }); + let _: Result<(), ()> = Box::pin(another_s).next().await.unwrap(); + }, + else => {}, + } + emitter.emit(()).await; + }); + } + + #[tokio::test] + async fn single_err() { + let s = async_try_stream(|mut emitter| async move { + if true { + Err("hello")?; + } else { + emitter.emit("world").await; + } + + unreachable!(); + }); + + let values: Vec<_> = s.collect().await; + assert_eq!(1, values.len()); + assert_eq!(Err("hello"), values[0]); + } + + #[tokio::test] + async fn emit_then_err() { + let s = async_try_stream(|mut emitter| async move { + emitter.emit("hello").await; + Err("world")?; + unreachable!(); + }); + + let values: Vec<_> = s.collect().await; + assert_eq!(2, values.len()); + assert_eq!(Ok("hello"), values[0]); + assert_eq!(Err("world"), values[1]); + } + + #[tokio::test] + async fn convert_err() { + struct ErrorA(u8); + #[derive(PartialEq, Debug)] + struct ErrorB(u8); + impl From for ErrorB { + fn from(a: ErrorA) -> ErrorB { + ErrorB(a.0) + } + } + + fn test() -> impl Stream> { + async_try_stream(|mut emitter| async move { + if true { + Err(ErrorA(1))?; + } else { + Err(ErrorB(2))?; + } + emitter.emit("unreachable").await; + Ok(()) + }) + } + + let values: Vec<_> = test().collect().await; + assert_eq!(1, values.len()); + assert_eq!(Err(ErrorB(1)), values[0]); + } + + #[tokio::test] + async fn multi_try() { + fn test() -> impl Stream> { + async_try_stream(|mut emitter| async move { + let a = Ok::<_, String>(Ok::<_, String>(123))??; + for _ in 1..10 { + emitter.emit(a).await; + } + Ok(()) + }) + } + let values: Vec<_> = test().collect().await; + assert_eq!(9, values.len()); + assert_eq!( + std::iter::repeat_n(123, 9).map(Ok).collect::>(), + values + ); + } + + struct DropGuard(Arc); + + impl Drop for DropGuard { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } + } + + #[tokio::test] + async fn generator_freed_on_done() { + let drops = Arc::new(AtomicUsize::new(0)); + let guard = DropGuard(Arc::clone(&drops)); + + let s = async_stream(|mut emitter| async move { + let _guard = guard; + emitter.emit(1).await; + }); + pin_mut!(s); + + assert_eq!(s.next().await, Some(1)); + assert_eq!(s.next().await, None); + + // State captured by the generator is dropped as soon as it completes + // (async blocks drop their locals on return), even though the stream + // itself is still alive + assert_eq!(drops.load(Ordering::SeqCst), 1); + assert_eq!(s.next().await, None); + } + + #[tokio::test] + async fn generator_freed_on_emitted_error() { + let drops = Arc::new(AtomicUsize::new(0)); + let guard = DropGuard(Arc::clone(&drops)); + + let s = async_try_stream(|mut emitter| async move { + let _guard = guard; + emitter.emit(1).await; + Err("boom") + }); + pin_mut!(s); + + assert_eq!(s.next().await, Some(Ok(1))); + assert_eq!(s.next().await, Some(Err("boom"))); + + // The stream terminates in the same poll that yields the error, so the + // generator state is freed even if the consumer never polls again + assert!(s.is_terminated()); + assert_eq!(drops.load(Ordering::SeqCst), 1); + + // Polling again after the error just returns None + assert_eq!(s.next().await, None); + } + + use pin_project_lite::pin_project; + + pin_project! { + struct MyStream { + #[pin] + input: T, + } + } + + impl Stream for MyStream { + type Item = T::Item; + + fn poll_next( + self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + let this = self.project(); + this.input.poll_next(cx) + } + } + + #[tokio::test] + async fn emit_does_not_hold_on_value() { + let waker = futures::task::noop_waker_ref(); + let mut cx = Context::from_waker(waker); + + let run = Arc::::new(AtomicUsize::new(0)); + let moved = Arc::clone(&run); + let s = async_stream(|mut emitter| async move { + for _ in 0..2 { + let before = moved.fetch_add(1, Ordering::SeqCst); + emitter.emit(before).await; + } + }); + + let mut my_stream = Box::pin(MyStream { input: s }); + + #[derive(Debug, PartialEq)] + struct Item { + before: usize, + result: Poll>, + after: usize, + } + + let mut results = vec![]; + + assert_eq!(run.load(Ordering::SeqCst), 0); + + while run.load(Ordering::SeqCst) < 2 { + let before = run.load(Ordering::SeqCst); + let result = my_stream.poll_next_unpin(&mut cx); + let after = run.load(Ordering::SeqCst); + results.push(Item { + before, + result, + after, + }); + } + + assert_eq!( + results, + vec![ + Item { + before: 0, + result: Poll::Ready(Some(0)), + after: 1, + }, + Item { + before: 1, + result: Poll::Ready(Some(1)), + after: 2, + } + ] + ); + } +} diff --git a/datafusion/execution/src/cache/cache_manager.rs b/datafusion/execution/src/cache/cache_manager.rs index 08a8dc9fd9cda..83dcf70975e2b 100644 --- a/datafusion/execution/src/cache/cache_manager.rs +++ b/datafusion/execution/src/cache/cache_manager.rs @@ -15,40 +15,90 @@ // specific language governing permissions and limitations // under the License. -use crate::cache::CacheAccessor; -use crate::cache::DefaultListFilesCache; -use crate::cache::file_statistics_cache::{ - DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, DefaultFileStatisticsCache, - DefaultFilesMetadataCache, -}; -use crate::cache::list_files_cache::ListFilesEntry; -use crate::cache::list_files_cache::TableScopedPath; -use datafusion_common::TableReference; +use crate::cache::default_cache::DefaultCache; +pub use crate::cache::{Cache, CacheValue, SchemaFingerprint, TableScopedPath}; +use datafusion_common::HashMap; use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; -use datafusion_common::stats::Precision; use datafusion_common::{Result, Statistics}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use object_store::ObjectMeta; use object_store::path::Path; use std::any::Any; -use std::collections::HashMap; use std::fmt::{Debug, Formatter}; use std::ops::Deref; use std::sync::Arc; use std::time::Duration; -pub use super::list_files_cache::{ - DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, DEFAULT_LIST_FILES_CACHE_TTL, -}; +pub const DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT: usize = 1024 * 1024; // 1MiB + +pub const DEFAULT_LIST_FILES_CACHE_TTL: Option = None; // Infinite + +pub const DEFAULT_FILE_STATISTICS_MEMORY_LIMIT: usize = 20 * 1024 * 1024; // 20MiB + +pub const DEFAULT_METADATA_CACHE_LIMIT: usize = 50 * 1024 * 1024; // 50M + +/// A cache for file statistics and orderings. +/// +/// This cache stores [`CachedFileMetadata`] which includes: +/// - File metadata for validation (size, last_modified) +/// - Statistics for the file +/// - Ordering information for the file +/// +/// If enabled via [`CacheManagerConfig::with_file_statistics_cache`] this +/// cache avoids inferring the same file statistics repeatedly during the +/// session lifetime. +/// +/// The typical usage pattern is: +/// 1. Call `get(path)` to check for cached value +/// 2. If `Some(cached)`, validate with +/// `cached.is_valid_for(¤t_meta, ¤t_schema_fingerprint)` +/// 3. If invalid or missing, compute new value and call `put(path, new_value)` +/// +/// See [`crate::runtime_env::RuntimeEnv`] for more details +pub type FileStatisticsCache = dyn Cache; + +/// A cache for storing the [`ObjectMeta`]s that result from listing a path. +/// +/// Listing a path means doing an object store "list" operation or `ls` +/// command on the local filesystem. This operation can be expensive, +/// especially when done over remote object stores. +/// +/// The cache key is always the table's base path, ensuring a stable cache key. +/// The cached value is a [`CachedFileList`] containing the files and a timestamp. +/// +/// Partition filtering is done after retrieval using [`CachedFileList::files_matching_prefix`]. +/// +/// See [`crate::runtime_env::RuntimeEnv`] for more details. +pub type ListFilesCache = dyn Cache; + +/// A cache for storing file-embedded metadata. +/// +/// This cache stores per-file metadata in the form of [`CachedFileMetadataEntry`], +/// which includes the [`ObjectMeta`] for validation. +/// +/// For example, the built in [`ListingTable`] uses this cache to avoid parsing +/// Parquet footers multiple times for the same file. +/// +/// The typical usage pattern is: +/// 1. Call `get(path)` to check for cached value +/// 2. If `Some(cached)`, validate with `cached.is_valid_for(¤t_meta)` +/// 3. If invalid or missing, compute new value and call `put(path, new_value)` +/// +/// See [`crate::runtime_env::RuntimeEnv`] for more details. +/// +/// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html +pub type FileMetadataCache = dyn Cache; /// Cached metadata for a file, including statistics and ordering. /// /// This struct embeds the [`ObjectMeta`] used for cache validation, -/// along with the cached statistics and ordering information. +/// the `file_schema` fingerprint, cached statistics, and ordering information. #[derive(Debug, Clone, PartialEq, Eq)] pub struct CachedFileMetadata { /// File metadata used for cache validation (size, last_modified). pub meta: ObjectMeta, + /// Fingerprint of the `file_schema` used to compute `statistics`. + pub schema_fingerprint: Arc, /// Cached statistics for the file, if available. pub statistics: Arc, /// Cached ordering for the file. @@ -59,11 +109,13 @@ impl CachedFileMetadata { /// Create a new cached file metadata entry. pub fn new( meta: ObjectMeta, + schema_fingerprint: Arc, statistics: Arc, ordering: Option, ) -> Self { Self { meta, + schema_fingerprint, statistics, ordering, } @@ -71,43 +123,24 @@ impl CachedFileMetadata { /// Check if this cached entry is still valid for the given metadata. /// - /// Returns true if the file size and last modified time match. - pub fn is_valid_for(&self, current_meta: &ObjectMeta) -> bool { + /// Returns true if the file size, last modified time, and schema match. + pub fn is_valid_for( + &self, + current_meta: &ObjectMeta, + current_schema_fingerprint: &Arc, + ) -> bool { self.meta.size == current_meta.size && self.meta.last_modified == current_meta.last_modified + && (Arc::ptr_eq(&self.schema_fingerprint, current_schema_fingerprint) + || self.schema_fingerprint.as_ref() + == current_schema_fingerprint.as_ref()) } } -/// A cache for file statistics and orderings. -/// -/// This cache stores [`CachedFileMetadata`] which includes: -/// - File metadata for validation (size, last_modified) -/// - Statistics for the file -/// - Ordering information for the file -/// -/// If enabled via [`CacheManagerConfig::with_file_statistics_cache`] this -/// cache avoids inferring the same file statistics repeatedly during the -/// session lifetime. -/// -/// The typical usage pattern is: -/// 1. Call `get(path)` to check for cached value -/// 2. If `Some(cached)`, validate with `cached.is_valid_for(¤t_meta)` -/// 3. If invalid or missing, compute new value and call `put(path, new_value)` -/// -/// See [`crate::runtime_env::RuntimeEnv`] for more details -pub trait FileStatisticsCache: - CacheAccessor -{ - /// Cache memory limit in bytes. - fn cache_limit(&self) -> usize; - - /// Updates the cache with a new memory limit in bytes. - fn update_cache_limit(&self, limit: usize); - - /// Retrieves the information about the entries currently cached. - fn list_entries(&self) -> HashMap; - - fn drop_table_entries(&self, table_ref: &Option) -> Result<()>; +impl CacheValue for CachedFileMetadata { + fn size(&self) -> usize { + DFHeapSize::heap_size(self, &mut DFHeapSizeCtx::default()) + } } impl DFHeapSize for CachedFileMetadata { @@ -118,27 +151,12 @@ impl DFHeapSize for CachedFileMetadata { + self.meta.e_tag.heap_size(ctx) + self.meta.location.as_ref().heap_size(ctx) + self.statistics.heap_size(ctx) + // Do not deep-count `schema_fingerprint`: each ListingTable shares one + // fingerprint across all cached files. //TODO add ordering once LexOrdering/PhysicalExpr implements DFHeapSize } } -/// Represents information about a cached statistics entry. -/// This is used to expose the statistics cache contents to outside modules. -#[derive(Debug, Clone, PartialEq, Eq)] -pub struct FileStatisticsCacheEntry { - pub object_meta: ObjectMeta, - /// Number of table rows. - pub num_rows: Precision, - /// Number of table columns. - pub num_columns: usize, - /// Total table size, in bytes. - pub table_size_bytes: Precision, - /// Size of the statistics entry, in bytes. - pub statistics_size_bytes: usize, - /// Whether ordering information is cached for this file. - pub has_ordering: bool, -} - /// Cached file listing. /// /// TTL expiration is handled internally by the cache implementation. @@ -181,6 +199,32 @@ impl CachedFileList { } } +impl CacheValue for CachedFileList { + fn size(&self) -> usize { + self.files.capacity() * size_of::() + + self + .files + .iter() + .map(meta_heap_bytes) + .reduce(|acc, b| acc + b) + .unwrap_or(0) + } +} + +/// Calculates the number of bytes an [`ObjectMeta`] occupies in the heap. +pub fn meta_heap_bytes(object_meta: &ObjectMeta) -> usize { + let mut size = object_meta.location.as_ref().len(); + + if let Some(e) = &object_meta.e_tag { + size += e.len(); + } + if let Some(v) = &object_meta.version { + size += v.len(); + } + + size +} + impl Deref for CachedFileList { type Target = Arc>; fn deref(&self) -> &Self::Target { @@ -194,38 +238,6 @@ impl From> for CachedFileList { } } -/// Cache for storing the [`ObjectMeta`]s that result from listing a path -/// -/// Listing a path means doing an object store "list" operation or `ls` -/// command on the local filesystem. This operation can be expensive, -/// especially when done over remote object stores. -/// -/// The cache key is always the table's base path, ensuring a stable cache key. -/// The cached value is a [`CachedFileList`] containing the files and a timestamp. -/// -/// Partition filtering is done after retrieval using [`CachedFileList::files_matching_prefix`]. -/// -/// See [`crate::runtime_env::RuntimeEnv`] for more details. -pub trait ListFilesCache: CacheAccessor { - /// Returns the cache's memory limit in bytes. - fn cache_limit(&self) -> usize; - - /// Returns the TTL (time-to-live) for cache entries, if configured. - fn cache_ttl(&self) -> Option; - - /// Updates the cache with a new memory limit in bytes. - fn update_cache_limit(&self, limit: usize); - - /// Updates the cache with a new TTL (time-to-live). - fn update_cache_ttl(&self, ttl: Option); - - /// Retrieves the information about the entries currently cached. - fn list_entries(&self) -> HashMap; - - /// Drop all entries for the given table reference. - fn drop_table_entries(&self, table_ref: &Option) -> Result<()>; -} - /// Generic file-embedded metadata used with [`FileMetadataCache`]. /// /// For example, Parquet footers and page metadata can be represented @@ -240,7 +252,7 @@ pub trait FileMetadata: Any + Send + Sync { /// Returns the size of the metadata in bytes. fn memory_size(&self) -> usize; - /// Returns extra information about this entry (used by [`FileMetadataCache::list_entries`]). + /// Returns extra information about this entry fn extra_info(&self) -> HashMap; } @@ -253,6 +265,12 @@ pub struct CachedFileMetadataEntry { pub file_metadata: Arc, } +impl CacheValue for CachedFileMetadataEntry { + fn size(&self) -> usize { + self.file_metadata.memory_size() + } +} + impl CachedFileMetadataEntry { /// Create a new cached file metadata entry. pub fn new(meta: ObjectMeta, file_metadata: Arc) -> Self { @@ -278,68 +296,6 @@ impl Debug for CachedFileMetadataEntry { } } -/// Cache for file-embedded metadata. -/// -/// This cache stores per-file metadata in the form of [`CachedFileMetadataEntry`], -/// which includes the [`ObjectMeta`] for validation. -/// -/// For example, the built in [`ListingTable`] uses this cache to avoid parsing -/// Parquet footers multiple times for the same file. -/// -/// DataFusion provides a default implementation, [`DefaultFilesMetadataCache`], -/// and users can also provide their own implementations to implement custom -/// caching strategies. -/// -/// The typical usage pattern is: -/// 1. Call `get(path)` to check for cached value -/// 2. If `Some(cached)`, validate with `cached.is_valid_for(¤t_meta)` -/// 3. If invalid or missing, compute new value and call `put(path, new_value)` -/// -/// See [`crate::runtime_env::RuntimeEnv`] for more details. -/// -/// [`ListingTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/listing/struct.ListingTable.html -pub trait FileMetadataCache: CacheAccessor { - /// Returns the cache's memory limit in bytes. - fn cache_limit(&self) -> usize; - - /// Updates the cache with a new memory limit in bytes. - fn update_cache_limit(&self, limit: usize); - - /// Retrieves the information about the entries currently cached. - fn list_entries(&self) -> HashMap; -} - -#[derive(Debug, Clone, PartialEq, Eq)] -/// Represents information about a cached metadata entry. -/// This is used to expose the metadata cache contents to outside modules. -pub struct FileMetadataCacheEntry { - pub object_meta: ObjectMeta, - /// Size of the cached metadata, in bytes. - pub size_bytes: usize, - /// Number of times this entry was retrieved. - pub hits: usize, - /// Additional object-specific information. - pub extra: HashMap, -} - -impl Debug for dyn FileStatisticsCache { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "Cache name: {} with length: {}", self.name(), self.len()) - } -} - -impl Debug for dyn ListFilesCache { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "Cache name: {} with length: {}", self.name(), self.len()) - } -} - -impl Debug for dyn FileMetadataCache { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "Cache name: {} with length: {}", self.name(), self.len()) - } -} - /// Manages various caches used in DataFusion. /// /// Following DataFusion design principles, DataFusion provides default cache @@ -349,28 +305,30 @@ impl Debug for dyn FileMetadataCache { /// See [`CacheManagerConfig`] for configuration options. #[derive(Debug)] pub struct CacheManager { - file_statistic_cache: Option>, - list_files_cache: Option>, - file_metadata_cache: Arc, + file_statistic_cache: Option>, + list_files_cache: Option>, + file_metadata_cache: Arc, } impl CacheManager { pub fn try_new(config: &CacheManagerConfig) -> Result> { - let file_statistic_cache = match &config.file_statistics_cache { - Some(fsc) if config.file_statistics_cache_limit > 0 => { - fsc.update_cache_limit(config.file_statistics_cache_limit); - Some(Arc::clone(fsc)) - } - None if config.file_statistics_cache_limit > 0 => { - let fsc: Arc = Arc::new( - DefaultFileStatisticsCache::new(config.file_statistics_cache_limit), - ); - Some(fsc) - } - _ => None, - }; - - let list_files_cache = match &config.list_files_cache { + let file_statistic_cache: Option> = + match &config.file_statistics_cache { + Some(fsc) if config.file_statistics_cache_limit > 0 => { + fsc.update_cache_limit(config.file_statistics_cache_limit); + Some(Arc::clone(fsc)) + } + None if config.file_statistics_cache_limit > 0 => Some(Arc::new( + DefaultCache::::new( + config.file_statistics_cache_limit, + ) + .with_name("DefaultFileStatisticsCache"), + )), + _ => None, + }; + + let list_files_cache: Option> = match &config.list_files_cache + { Some(lfc) if config.list_files_cache_limit > 0 => { // the cache memory limit or ttl might have changed, ensure they are updated lfc.update_cache_limit(config.list_files_cache_limit); @@ -380,13 +338,13 @@ impl CacheManager { } Some(Arc::clone(lfc)) } - None if config.list_files_cache_limit > 0 => { - let lfc: Arc = Arc::new(DefaultListFilesCache::new( + None if config.list_files_cache_limit > 0 => Some(Arc::new( + DefaultCache::::new_with_ttl( config.list_files_cache_limit, config.list_files_cache_ttl, - )); - Some(lfc) - } + ) + .with_name("DefaultListFilesCache"), + )), _ => None, }; @@ -395,7 +353,10 @@ impl CacheManager { .as_ref() .map(Arc::clone) .unwrap_or_else(|| { - Arc::new(DefaultFilesMetadataCache::new(config.metadata_cache_limit)) + Arc::new( + DefaultCache::new(config.metadata_cache_limit) + .with_name("DefaultFileMetadataCache"), + ) }); // the cache memory limit might have changed, ensure the limit is updated @@ -409,7 +370,7 @@ impl CacheManager { } /// Get the file statistics cache. - pub fn get_file_statistic_cache(&self) -> Option> { + pub fn get_file_statistic_cache(&self) -> Option> { self.file_statistic_cache.clone() } @@ -421,7 +382,7 @@ impl CacheManager { } /// Get the cache for storing the result of listing [`ObjectMeta`]s under the same path. - pub fn get_list_files_cache(&self) -> Option> { + pub fn get_list_files_cache(&self) -> Option> { self.list_files_cache.clone() } @@ -438,7 +399,7 @@ impl CacheManager { } /// Get the file embedded metadata cache. - pub fn get_file_metadata_cache(&self) -> Arc { + pub fn get_file_metadata_cache(&self) -> Arc { Arc::clone(&self.file_metadata_cache) } @@ -448,14 +409,12 @@ impl CacheManager { } } -pub const DEFAULT_METADATA_CACHE_LIMIT: usize = 50 * 1024 * 1024; // 50M - #[derive(Clone)] pub struct CacheManagerConfig { /// Enable caching of file statistics when listing files. /// Enabling the cache avoids repeatedly reading file statistics in a DataFusion session. /// Default is enabled. Currently only Parquet files are supported. - pub file_statistics_cache: Option>, + pub file_statistics_cache: Option>, /// Limit of the file statistics cache, in bytes. Default: 20MiB. pub file_statistics_cache_limit: usize, /// Enable caching of file metadata when listing files. @@ -465,7 +424,7 @@ pub struct CacheManagerConfig { /// Note that if this option is enabled, DataFusion will not see any updates to the underlying /// storage for at least `list_files_cache_ttl` duration. /// Default is enabled. - pub list_files_cache: Option>, + pub list_files_cache: Option>, /// Limit of the `list_files_cache`, in bytes. Default: 1MiB. pub list_files_cache_limit: usize, /// The duration the list files cache will consider an entry valid after insertion. Note that @@ -474,8 +433,8 @@ pub struct CacheManagerConfig { pub list_files_cache_ttl: Option, /// Cache of file-embedded metadata, used to avoid reading it multiple times when processing a /// data file (e.g., Parquet footer and page metadata). - /// If not provided, the [`CacheManager`] will create a [`DefaultFilesMetadataCache`]. - pub file_metadata_cache: Option>, + /// If not provided, the [`CacheManager`] will create it. + pub file_metadata_cache: Option>, /// Limit of the file-embedded metadata cache, in bytes. pub metadata_cache_limit: usize, } @@ -498,7 +457,7 @@ impl CacheManagerConfig { /// Set the cache for file statistics. pub fn with_file_statistics_cache( mut self, - cache: Option>, + cache: Option>, ) -> Self { self.file_statistics_cache = cache; self @@ -513,10 +472,7 @@ impl CacheManagerConfig { /// Set the cache for listing files. /// /// Default is `None` (disabled). - pub fn with_list_files_cache( - mut self, - cache: Option>, - ) -> Self { + pub fn with_list_files_cache(mut self, cache: Option>) -> Self { self.list_files_cache = cache; self } @@ -538,11 +494,9 @@ impl CacheManagerConfig { } /// Sets the cache for file-embedded metadata. - /// - /// Default is a [`DefaultFilesMetadataCache`]. pub fn with_file_metadata_cache( mut self, - cache: Option>, + cache: Option>, ) -> Self { self.file_metadata_cache = cache; self @@ -566,7 +520,7 @@ mod tests { fn test_ttl_preserved_when_not_set_in_config() { // Create a cache with TTL = 1 second let list_file_cache = - DefaultListFilesCache::new(1024, Some(Duration::from_secs(1))); + DefaultCache::new_with_ttl(1024, Some(Duration::from_secs(1))); // Verify the cache has TTL set initially assert_eq!( @@ -603,7 +557,7 @@ mod tests { fn test_ttl_overridden_when_set_in_config() { // Create a cache with TTL = 1 second let list_file_cache = - DefaultListFilesCache::new(1024, Some(Duration::from_secs(1))); + DefaultCache::new_with_ttl(1024, Some(Duration::from_secs(1))); // Put cache in config WITH a different TTL set let config = CacheManagerConfig::default() diff --git a/datafusion/execution/src/cache/default_cache.rs b/datafusion/execution/src/cache/default_cache.rs new file mode 100644 index 0000000000000..bfe326f3a47e1 --- /dev/null +++ b/datafusion/execution/src/cache/default_cache.rs @@ -0,0 +1,2015 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::{Arc, Mutex}; +use std::time::Duration; + +use datafusion_common::TableReference; +use datafusion_common::instant::Instant; +use datafusion_common::{HashMap, Result}; + +use crate::cache::lru_queue::LruQueue; +use crate::cache::{Cache, CacheEntryInfo, CacheKey, CacheValue}; + +/// Source of the current time used by a [`DefaultCache`] when applying TTLs. +pub trait TimeProvider: Send + Sync { + /// Return the current instant. + fn now(&self) -> Instant; +} + +/// [`TimeProvider`] backed by [`Instant::now`]. +/// +/// This is the default time source used by [`DefaultCache`] +#[derive(Debug, Default)] +pub struct SystemTimeProvider; + +impl TimeProvider for SystemTimeProvider { + fn now(&self) -> Instant { + Instant::now() + } +} + +#[derive(Clone)] +struct ValueEntry { + value: V, + expires: Option, +} + +struct DefaultCacheState { + lru_queue: LruQueue>, + hits: HashMap, + memory_limit: usize, + memory_used: usize, + ttl: Option, +} + +impl DefaultCacheState { + fn new(memory_limit: usize, ttl: Option) -> Self { + Self { + lru_queue: LruQueue::new(), + hits: HashMap::new(), + memory_limit, + memory_used: 0, + ttl, + } + } + + fn get(&mut self, key: &K, now: Instant) -> Option { + let entry = self.lru_queue.get(key)?; + if let Some(exp) = entry.expires + && now > exp + { + self.remove(key); + return None; + } + let value = entry.value.clone(); + *self.hits.entry(key.clone()).or_insert(0) += 1; + Some(value) + } + + fn contains_key(&mut self, key: &K, now: Instant) -> bool { + let Some(entry) = self.lru_queue.peek(key) else { + return false; + }; + match entry.expires { + Some(exp) if now > exp => { + self.remove(key); + false + } + _ => true, + } + } + + fn put(&mut self, key: &K, value: V, now: Instant) -> Option { + let value_size = value.size(); + + if value_size == 0 { + return None; + } + + let key_size = key.size(); + let total_size = key_size + value_size; + + if total_size > self.memory_limit { + // Remove potential stale entry + return self.remove(key); + } + + let expires = self.ttl.map(|ttl| now + ttl); + let entry = ValueEntry { value, expires }; + + self.memory_used += total_size; + self.hits.insert(key.clone(), 0); + let old = self.lru_queue.put(key.clone(), entry); + if let Some(old_entry) = &old { + self.memory_used -= key_size; + self.memory_used -= old_entry.value.size(); + } + + self.evict_entries(); + + old.map(|v| v.value) + } + + fn remove(&mut self, key: &K) -> Option { + let entry = self.lru_queue.remove(key)?; + self.memory_used -= key.size(); + self.memory_used -= entry.value.size(); + self.hits.remove(key); + Some(entry.value) + } + + fn evict_entries(&mut self) { + while self.memory_used > self.memory_limit { + let Some((evicted_key, evicted)) = self.lru_queue.pop() else { + // cache is empty while memory_used > memory_limit, cannot happen + log::error!( + "DefaultCache memory accounting bug: memory_used={} but cache is empty", + self.memory_used + ); + debug_assert!(false, "memory_used > limit with empty cache"); + self.memory_used = 0; + return; + }; + self.memory_used -= evicted_key.size(); + self.memory_used -= evicted.value.size(); + self.hits.remove(&evicted_key); + } + } + + fn clear(&mut self) { + self.lru_queue.clear(); + self.hits.clear(); + self.memory_used = 0; + } +} + +/// In-memory [`Cache`] with an LRU eviction policy, byte-based memory limit, +/// and optional per-entry TTL. +/// +/// Entries are evicted in least-recently-used order whenever an insert would +/// push `memory_used` above `memory_limit`. Inserts whose own size exceeds the +/// limit are rejected (and any prior entry under the same key is removed). +/// When a TTL is configured, the expiration is stamped onto each entry at +/// insertion time and checked lazily on access. Entries with size 0 are rejected. +pub struct DefaultCache { + state: Mutex>, + time_provider: Arc, + name: String, +} + +impl DefaultCache { + /// Create a cache with the given memory budget in bytes and no TTL. + pub fn new(memory_limit: usize) -> Self { + Self::new_with_ttl(memory_limit, None) + } + + /// Create a cache with the given memory budget in bytes and an optional + /// TTL applied to every newly inserted entry. + pub fn new_with_ttl(memory_limit: usize, ttl: Option) -> Self { + Self { + state: Mutex::new(DefaultCacheState::new(memory_limit, ttl)), + time_provider: Arc::new(SystemTimeProvider), + name: "DefaultCache".to_string(), + } + } + + /// Override the cache name. + pub fn with_name(mut self, name: impl Into) -> Self { + self.name = name.into(); + self + } + + /// Override the time source used to stamp and check TTLs. + pub fn with_time_provider(mut self, provider: Arc) -> Self { + self.time_provider = provider; + self + } + + /// Number of bytes currently accounted for by live entries. + pub fn memory_used(&self) -> usize { + self.state.lock().unwrap().memory_used + } +} + +impl Cache for DefaultCache { + fn get(&self, key: &K) -> Option { + let now = self.time_provider.now(); + let mut state = self.state.lock().unwrap(); + state.get(key, now) + } + + fn put(&self, key: &K, value: V) -> Option { + let now = self.time_provider.now(); + let mut state = self.state.lock().unwrap(); + state.put(key, value, now) + } + + fn remove(&self, k: &K) -> Option { + let mut state = self.state.lock().unwrap(); + state.remove(k) + } + + fn contains_key(&self, k: &K) -> bool { + let now = self.time_provider.now(); + let mut state = self.state.lock().unwrap(); + state.contains_key(k, now) + } + + fn len(&self) -> usize { + self.state.lock().unwrap().lru_queue.len() + } + + fn clear(&self) { + let mut state = self.state.lock().unwrap(); + state.clear(); + } + + fn name(&self) -> String { + self.name.clone() + } + fn cache_limit(&self) -> usize { + self.state.lock().unwrap().memory_limit + } + + fn update_cache_limit(&self, limit: usize) { + let mut state = self.state.lock().unwrap(); + state.memory_limit = limit; + state.evict_entries(); + } + + fn cache_ttl(&self) -> Option { + self.state.lock().unwrap().ttl + } + + fn update_cache_ttl(&self, ttl: Option) { + let mut state = self.state.lock().unwrap(); + state.ttl = ttl; + } + + fn drop_table_entries(&self, table_ref: &TableReference) -> Result<()> { + let mut state = self.state.lock().unwrap(); + let to_remove: Vec = state + .lru_queue + .keys() + .filter(|k| k.table_ref() == Some(table_ref)) + .cloned() + .collect(); + for k in &to_remove { + state.remove(k); + } + Ok(()) + } + + fn list_entries(&self) -> HashMap> { + let state = self.state.lock().unwrap(); + state + .lru_queue + .list_entries() + .into_iter() + .map(|(k, entry)| { + let hits = state.hits.get(k).copied().unwrap_or(0); + let info = CacheEntryInfo { + value: entry.value.clone(), + size_bytes: entry.value.size(), + hits, + expires: entry.expires, + }; + (k.clone(), info) + }) + .collect() + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::cache::cache_manager::{ + CachedFileList, DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, meta_heap_bytes, + }; + use crate::cache::cache_manager::{ + CachedFileMetadata, DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, + }; + use crate::cache::cache_manager::{CachedFileMetadataEntry, FileMetadata}; + use crate::cache::default_cache::DefaultCache; + use crate::cache::default_cache::TimeProvider; + use crate::cache::{Cache, CacheEntryInfo}; + use crate::cache::{CacheKey, CacheValue}; + use crate::cache::{SchemaFingerprint, TableScopedPath}; + use arrow::array::{Int32Array, ListArray, RecordBatch}; + use arrow::buffer::{OffsetBuffer, ScalarBuffer}; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; + use chrono::DateTime; + use datafusion_common::HashMap; + use datafusion_common::TableReference; + use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; + use datafusion_common::instant::Instant; + use datafusion_common::stats::Precision; + use datafusion_common::{ColumnStatistics, ScalarValue, Statistics}; + use datafusion_expr::ColumnarValue; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; + use object_store::ObjectMeta; + use object_store::path::Path; + use std::sync::Mutex; + use std::thread; + use std::time::Duration; + + pub struct TestFileMetadata { + metadata: String, + } + + impl FileMetadata for TestFileMetadata { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn memory_size(&self) -> usize { + self.metadata.len() + } + + fn extra_info(&self) -> HashMap { + HashMap::from([("extra_info".to_owned(), "abc".to_owned())]) + } + } + + impl PartialEq for CachedFileMetadataEntry { + fn eq(&self, other: &Self) -> bool { + self.meta == other.meta + } + } + + fn create_test_object_meta(path: &str, size: usize) -> ObjectMeta { + ObjectMeta { + location: Path::from(path), + last_modified: DateTime::parse_from_rfc3339("2025-07-29T12:12:12+00:00") + .unwrap() + .into(), + size: size as u64, + e_tag: None, + version: None, + } + } + + #[test] + fn test_default_file_metadata_cache() { + let object_meta = create_test_object_meta("test", 1024); + + let metadata: Arc = Arc::new(TestFileMetadata { + metadata: "retrieved_metadata".to_owned(), + }); + + let cache = DefaultCache::new(1024 * 1024); + + // Cache miss + assert!(cache.get(&object_meta.location).is_none()); + + // Put a value + let cached_entry = + CachedFileMetadataEntry::new(object_meta.clone(), Arc::clone(&metadata)); + cache.put(&object_meta.location, cached_entry); + + // Verify the cached value + assert!(cache.contains_key(&object_meta.location)); + let result = cache.get(&object_meta.location).unwrap(); + let test_file_metadata = Arc::downcast::(result.file_metadata); + assert!(test_file_metadata.is_ok()); + assert_eq!(test_file_metadata.unwrap().metadata, "retrieved_metadata"); + + // Cache hit - check validation + let result2 = cache.get(&object_meta.location).unwrap(); + assert!(result2.is_valid_for(&object_meta)); + + // File size changed - closure should detect invalidity + let object_meta2 = create_test_object_meta("test", 2048); + let result3 = cache.get(&object_meta2.location).unwrap(); + // Cached entry should NOT be valid for new meta + assert!(!result3.is_valid_for(&object_meta2)); + + // Return new entry + let new_entry = + CachedFileMetadataEntry::new(object_meta2.clone(), Arc::clone(&metadata)); + cache.put(&object_meta2.location, new_entry); + + let result4 = cache.get(&object_meta2.location).unwrap(); + assert_eq!(result4.meta.size, 2048); + + // remove + cache.remove(&object_meta.location); + assert!(!cache.contains_key(&object_meta.location)); + + // len and clear + let object_meta3 = create_test_object_meta("test3", 100); + cache.put( + &object_meta.location, + CachedFileMetadataEntry::new(object_meta.clone(), Arc::clone(&metadata)), + ); + cache.put( + &object_meta3.location, + CachedFileMetadataEntry::new(object_meta3.clone(), Arc::clone(&metadata)), + ); + assert_eq!(cache.len(), 2); + cache.clear(); + assert_eq!(cache.len(), 0); + } + + fn generate_test_metadata_with_size( + path: &str, + size: usize, + ) -> (ObjectMeta, Arc) { + let object_meta = ObjectMeta { + location: Path::from(path), + last_modified: chrono::Utc::now(), + size: size as u64, + e_tag: None, + version: None, + }; + let metadata = "a".repeat(size); + let metadata: Arc = Arc::new(TestFileMetadata { metadata }); + + (object_meta, metadata) + } + + #[test] + fn test_default_file_metadata_cache_with_limit() { + // Create a cache with 1000 bytes capacity + 4 keys each key 2 bytes + let cache = DefaultCache::new(1000 + 4 * 2); + + let (object_meta1, metadata1) = generate_test_metadata_with_size("01", 100); + let (object_meta2, metadata2) = generate_test_metadata_with_size("02", 500); + let (object_meta3, metadata3) = generate_test_metadata_with_size("03", 300); + + cache.put( + &object_meta1.location, + CachedFileMetadataEntry::new(object_meta1.clone(), metadata1), + ); + cache.put( + &object_meta2.location, + CachedFileMetadataEntry::new(object_meta2.clone(), metadata2), + ); + cache.put( + &object_meta3.location, + CachedFileMetadataEntry::new(object_meta3.clone(), metadata3), + ); + + // all entries will fit + assert_eq!(cache.len(), 3); + assert_eq!(cache.memory_used(), 906); + assert!(cache.contains_key(&object_meta1.location)); + assert!(cache.contains_key(&object_meta2.location)); + assert!(cache.contains_key(&object_meta3.location)); + + // add a new entry which will remove the least recently used ("1") + let (object_meta4, metadata4) = generate_test_metadata_with_size("04", 200); + cache.put( + &object_meta4.location, + CachedFileMetadataEntry::new(object_meta4.clone(), metadata4), + ); + assert_eq!(cache.len(), 3); + assert_eq!(cache.memory_used(), 1006); + assert!(!cache.contains_key(&object_meta1.location)); + assert!(cache.contains_key(&object_meta4.location)); + + // get entry "2", which will move it to the top of the queue, and add a new one which will + // remove the new least recently used ("3") + let _ = cache.get(&object_meta2.location); + let (object_meta5, metadata5) = generate_test_metadata_with_size("05", 100); + cache.put( + &object_meta5.location, + CachedFileMetadataEntry::new(object_meta5.clone(), metadata5), + ); + assert_eq!(cache.len(), 3); + assert_eq!(cache.memory_used(), 806); + assert!(!cache.contains_key(&object_meta3.location)); + assert!(cache.contains_key(&object_meta5.location)); + + // new entry which will not be able to fit in the 1000 bytes allocated + let (object_meta6, metadata6) = generate_test_metadata_with_size("06", 1200); + cache.put( + &object_meta6.location, + CachedFileMetadataEntry::new(object_meta6.clone(), metadata6), + ); + assert_eq!(cache.len(), 3); + assert_eq!(cache.memory_used(), 806); + assert!(!cache.contains_key(&object_meta6.location)); + + // new entry which is able to fit without removing any entry + let (object_meta7, metadata7) = generate_test_metadata_with_size("07", 200); + cache.put( + &object_meta7.location, + CachedFileMetadataEntry::new(object_meta7.clone(), metadata7), + ); + assert_eq!(cache.len(), 4); + assert_eq!(cache.memory_used(), 1008); + assert!(cache.contains_key(&object_meta7.location)); + + // new entry which will remove all other entries + let (object_meta8, metadata8) = generate_test_metadata_with_size("08", 999); + cache.put( + &object_meta8.location, + CachedFileMetadataEntry::new(object_meta8.clone(), metadata8), + ); + assert_eq!(cache.len(), 1); + assert_eq!(cache.memory_used(), 1001); + assert!(cache.contains_key(&object_meta8.location)); + + // when updating an entry, the previous ones are not unnecessarily removed + let (object_meta9, metadata9) = generate_test_metadata_with_size("09", 300); + let (object_meta10, metadata10) = generate_test_metadata_with_size("10", 200); + let (object_meta11_v1, metadata11_v1) = + generate_test_metadata_with_size("11", 400); + cache.put( + &object_meta9.location, + CachedFileMetadataEntry::new(object_meta9.clone(), metadata9), + ); + cache.put( + &object_meta10.location, + CachedFileMetadataEntry::new(object_meta10.clone(), metadata10), + ); + cache.put( + &object_meta11_v1.location, + CachedFileMetadataEntry::new(object_meta11_v1.clone(), metadata11_v1), + ); + assert_eq!(cache.memory_used(), 906); + assert_eq!(cache.len(), 3); + let (object_meta11_v2, metadata11_v2) = + generate_test_metadata_with_size("11", 500); + cache.put( + &object_meta11_v2.location, + CachedFileMetadataEntry::new(object_meta11_v2.clone(), metadata11_v2), + ); + assert_eq!(cache.memory_used(), 1006); + assert_eq!(cache.len(), 3); + assert!(cache.contains_key(&object_meta9.location)); + assert!(cache.contains_key(&object_meta10.location)); + assert!(cache.contains_key(&object_meta11_v2.location)); + + // when updating an entry that now exceeds the limit, the LRU ("09") needs to be removed + let (object_meta11_v3, metadata11_v3) = + generate_test_metadata_with_size("11", 510); + cache.put( + &object_meta11_v3.location, + CachedFileMetadataEntry::new(object_meta11_v3.clone(), metadata11_v3), + ); + assert_eq!(cache.memory_used(), 714); + assert_eq!(cache.len(), 2); + assert!(cache.contains_key(&object_meta10.location)); + assert!(cache.contains_key(&object_meta11_v3.location)); + + // manually removing an entry that is not the LRU + cache.remove(&object_meta11_v3.location); + assert_eq!(cache.len(), 1); + assert_eq!(cache.memory_used(), 202); + assert!(cache.contains_key(&object_meta10.location)); + assert!(!cache.contains_key(&object_meta11_v3.location)); + + // clear + cache.clear(); + assert_eq!(cache.len(), 0); + assert_eq!(cache.memory_used(), 0); + + // resizing the cache should clear the extra entries + let (object_meta12, metadata12) = generate_test_metadata_with_size("12", 300); + let (object_meta13, metadata13) = generate_test_metadata_with_size("13", 200); + let (object_meta14, metadata14) = generate_test_metadata_with_size("14", 500); + cache.put( + &object_meta12.location, + CachedFileMetadataEntry::new(object_meta12.clone(), metadata12), + ); + cache.put( + &object_meta13.location, + CachedFileMetadataEntry::new(object_meta13.clone(), metadata13), + ); + cache.put( + &object_meta14.location, + CachedFileMetadataEntry::new(object_meta14.clone(), metadata14), + ); + assert_eq!(cache.len(), 3); + assert_eq!(cache.memory_used(), 1006); + cache.update_cache_limit(600); + assert_eq!(cache.len(), 1); + assert_eq!(cache.memory_used(), 502); + assert!(!cache.contains_key(&object_meta12.location)); + assert!(!cache.contains_key(&object_meta13.location)); + assert!(cache.contains_key(&object_meta14.location)); + } + + #[test] + fn test_default_file_metadata_cache_entries_info() { + // Create a cache with 1000 bytes + 4 bytes for 4 keys each key 1 byte + let cache = DefaultCache::new(1000 + 4); + + let (object_meta1, metadata1) = generate_test_metadata_with_size("1", 100); + let (object_meta2, metadata2) = generate_test_metadata_with_size("2", 200); + let (object_meta3, metadata3) = generate_test_metadata_with_size("3", 300); + + // initial entries, all will have hits = 0 + let entry_1 = CachedFileMetadataEntry::new(object_meta1.clone(), metadata1); + let entry_2 = CachedFileMetadataEntry::new(object_meta2.clone(), metadata2); + let entry_3 = CachedFileMetadataEntry::new(object_meta3.clone(), metadata3); + + // Build a cache which fits exactly these 3 entries + + cache.put(&object_meta1.location, entry_1.clone()); + cache.put(&object_meta2.location, entry_2.clone()); + cache.put(&object_meta3.location, entry_3.clone()); + let entries = cache.list_entries(); + + assert_eq!( + entries, + HashMap::from([ + ( + Path::from("1"), + CacheEntryInfo { + value: entry_1.clone(), + size_bytes: 100, + hits: 0, + expires: None, + } + ), + ( + Path::from("2"), + CacheEntryInfo { + value: entry_2.clone(), + size_bytes: 200, + hits: 0, + expires: None, + } + ), + ( + Path::from("3"), + CacheEntryInfo { + value: entry_3.clone(), + size_bytes: 300, + hits: 0, + expires: None, + } + ) + ]) + ); + + // new hit on "1" + let _ = cache.get(&object_meta1.location); + assert_eq!( + cache.list_entries(), + HashMap::from([ + ( + Path::from("1"), + CacheEntryInfo { + value: entry_1.clone(), + size_bytes: 100, + hits: 1, + expires: None, + } + ), + ( + Path::from("2"), + CacheEntryInfo { + value: entry_2.clone(), + size_bytes: 200, + hits: 0, + expires: None, + } + ), + ( + Path::from("3"), + CacheEntryInfo { + value: entry_3.clone(), + size_bytes: 300, + hits: 0, + expires: None, + } + ) + ]) + ); + + // new entry, will evict "2" + let (object_meta4, metadata4) = generate_test_metadata_with_size("4", 600); + let entry_4 = CachedFileMetadataEntry::new(object_meta4.clone(), metadata4); + cache.put(&object_meta4.location, entry_4.clone()); + assert_eq!( + cache.list_entries(), + HashMap::from([ + ( + Path::from("1"), + CacheEntryInfo { + value: entry_1.clone(), + size_bytes: 100, + hits: 1, + expires: None, + } + ), + ( + Path::from("3"), + CacheEntryInfo { + value: entry_3.clone(), + size_bytes: 300, + hits: 0, + expires: None, + } + ), + ( + Path::from("4"), + CacheEntryInfo { + value: entry_4.clone(), + size_bytes: 600, + hits: 0, + expires: None, + } + ) + ]) + ); + + // replace entry "1" + let (object_meta1_new, metadata1_new) = generate_test_metadata_with_size("1", 50); + let entry_1 = + CachedFileMetadataEntry::new(object_meta1_new.clone(), metadata1_new); + cache.put(&object_meta1_new.location, entry_1.clone()); + assert_eq!( + cache.list_entries(), + HashMap::from([ + ( + Path::from("1"), + CacheEntryInfo { + value: entry_1.clone(), + size_bytes: 50, + hits: 0, + expires: None, + } + ), + ( + Path::from("3"), + CacheEntryInfo { + value: entry_3.clone(), + size_bytes: 300, + hits: 0, + expires: None, + } + ), + ( + Path::from("4"), + CacheEntryInfo { + value: entry_4.clone(), + size_bytes: 600, + hits: 0, + expires: None, + } + ) + ]) + ); + + // remove entry "4" + cache.remove(&object_meta4.location); + assert_eq!( + cache.list_entries(), + HashMap::from([ + ( + Path::from("1"), + CacheEntryInfo { + value: entry_1.clone(), + size_bytes: 50, + hits: 0, + expires: None, + } + ), + ( + Path::from("3"), + CacheEntryInfo { + value: entry_3.clone(), + size_bytes: 300, + hits: 0, + expires: None, + } + ) + ]) + ); + + // clear + cache.clear(); + assert_eq!(cache.list_entries(), HashMap::from([])); + } + + fn create_test_meta(path: &str, size: u64) -> ObjectMeta { + ObjectMeta { + location: Path::from(path), + last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") + .unwrap() + .into(), + size, + e_tag: None, + version: None, + } + } + + #[test] + fn test_statistics_cache() { + let meta = create_test_meta("test", 1024); + let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); + + let schema = Schema::new(vec![Field::new( + "test_column", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]); + + let path = TableScopedPath { + path: meta.location.clone(), + table: None, + }; + let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); + + // Cache miss + assert!(cache.get(&path).is_none()); + + // Put a value + let cached_value = CachedFileMetadata::new( + meta.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + cache.put(&path, cached_value); + + // Cache hit + let result = cache.get(&path); + assert!(result.is_some()); + + let cached = result.unwrap(); + assert!(cached.is_valid_for(&meta, &schema_fingerprint)); + + let equivalent_schema_fingerprint = + Arc::new(SchemaFingerprint::from_schema(&schema)); + assert!(!Arc::ptr_eq( + &schema_fingerprint, + &equivalent_schema_fingerprint + )); + assert!(cached.is_valid_for(&meta, &equivalent_schema_fingerprint)); + + let different_schema = Schema::new(vec![Field::new( + "different_column", + DataType::Timestamp(TimeUnit::Second, None), + false, + )]); + let different_schema_fingerprint = + Arc::new(SchemaFingerprint::from_schema(&different_schema)); + assert!(!cached.is_valid_for(&meta, &different_schema_fingerprint)); + + // File size changed - validation should fail + let meta2 = create_test_meta("test", 2048); + + let cached = cache.get(&path).unwrap(); + assert!(!cached.is_valid_for(&meta2, &schema_fingerprint)); + + // Update with new value + let cached_value2 = CachedFileMetadata::new( + meta2.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + cache.put(&path, cached_value2); + + // Test list_entries + let entries = cache.list_entries(); + assert_eq!(entries.len(), 1); + + let path_3 = TableScopedPath { + path: Path::from("test"), + table: None, + }; + + let entry = entries.get(&path_3).unwrap(); + assert_eq!(entry.value.meta.size, 2048); // Should be updated value + } + + #[derive(Clone, Debug, PartialEq, Eq, Hash)] + struct MockExpr {} + + impl std::fmt::Display for MockExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "MockExpr") + } + } + + impl PhysicalExpr for MockExpr { + fn data_type( + &self, + _input_schema: &Schema, + ) -> datafusion_common::Result { + Ok(DataType::Int32) + } + + fn nullable(&self, _input_schema: &Schema) -> datafusion_common::Result { + Ok(false) + } + + fn evaluate( + &self, + _batch: &RecordBatch, + ) -> datafusion_common::Result { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> datafusion_common::Result> { + assert!(children.is_empty()); + Ok(self) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "MockExpr") + } + } + + fn ordering() -> LexOrdering { + let expr = Arc::new(MockExpr {}) as Arc; + LexOrdering::new(vec![PhysicalSortExpr::new_default(expr)]).unwrap() + } + + #[test] + fn test_ordering_cache() { + let meta = create_test_meta("test.parquet", 100); + let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); + + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); + + // Cache statistics with no ordering + let cached_value = CachedFileMetadata::new( + meta.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + None, // No ordering yet + ); + + let path = TableScopedPath { + path: meta.location.clone(), + table: None, + }; + + cache.put(&path, cached_value); + + let result = cache.get(&path).unwrap(); + assert!(result.ordering.is_none()); + + // Update to add ordering + let mut cached = cache.get(&path).unwrap(); + if cached.is_valid_for(&meta, &schema_fingerprint) && cached.ordering.is_none() { + cached.ordering = Some(ordering()); + } + cache.put(&path, cached); + + let result2 = cache.get(&path).unwrap(); + assert!(result2.ordering.is_some()); + + // Verify list_entries shows has_ordering = true + let entries = cache.list_entries(); + assert_eq!(entries.len(), 1); + assert!(entries.get(&path).unwrap().value.ordering.is_some()); + } + + #[test] + fn test_cache_invalidation_on_file_modification() { + let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); + let path = TableScopedPath { + path: Path::from("test.parquet"), + table: None, + }; + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); + + let meta_v1 = create_test_meta("test.parquet", 100); + + // Cache initial value + let cached_value = CachedFileMetadata::new( + meta_v1.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + cache.put(&path, cached_value); + + // File modified (size changed) + let meta_v2 = create_test_meta("test.parquet", 200); + + let cached = cache.get(&path).unwrap(); + // Should not be valid for new meta + assert!(!cached.is_valid_for(&meta_v2, &schema_fingerprint)); + + // Compute new value and update + let new_cached = CachedFileMetadata::new( + meta_v2.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + cache.put(&path, new_cached); + + // Should have new metadata + let result = cache.get(&path).unwrap(); + assert_eq!(result.meta.size, 200); + } + + #[test] + fn test_ordering_cache_invalidation_on_file_modification() { + let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); + let path = TableScopedPath { + path: Path::from("test.parquet"), + table: None, + }; + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); + + // Cache with original metadata and ordering + let meta_v1 = ObjectMeta { + location: path.path.clone(), + last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") + .unwrap() + .into(), + size: 100, + e_tag: None, + version: None, + }; + let ordering_v1 = ordering(); + let cached_v1 = CachedFileMetadata::new( + meta_v1.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + Some(ordering_v1), + ); + cache.put(&path, cached_v1); + + // Verify cached ordering is valid + let cached = cache.get(&path).unwrap(); + assert!(cached.is_valid_for(&meta_v1, &schema_fingerprint)); + assert!(cached.ordering.is_some()); + + // File modified (size changed) + let meta_v2 = ObjectMeta { + location: path.path.clone(), + last_modified: DateTime::parse_from_rfc3339("2022-09-28T10:00:00+02:00") + .unwrap() + .into(), + size: 200, // Changed + e_tag: None, + version: None, + }; + + // Cache entry exists but should be invalid for new metadata + let cached = cache.get(&path).unwrap(); + assert!(!cached.is_valid_for(&meta_v2, &schema_fingerprint)); + + // Cache new version with different ordering + let ordering_v2 = ordering(); // New ordering instance + let cached_v2 = CachedFileMetadata::new( + meta_v2.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + Some(ordering_v2), + ); + cache.put(&path, cached_v2); + + // Old metadata should be invalid + let cached = cache.get(&path).unwrap(); + assert!(!cached.is_valid_for(&meta_v1, &schema_fingerprint)); + + // New metadata should be valid + assert!(cached.is_valid_for(&meta_v2, &schema_fingerprint)); + assert!(cached.ordering.is_some()); + } + + #[test] + fn test_list_entries() { + let cache = DefaultCache::new(DEFAULT_FILE_STATISTICS_MEMORY_LIMIT); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); + + let meta1 = create_test_meta("test1.parquet", 100); + + let cached_value_1 = CachedFileMetadata::new( + meta1.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + None, + ); + + let path_1 = TableScopedPath { + path: meta1.location.clone(), + table: None, + }; + + cache.put(&path_1, cached_value_1.clone()); + let meta2 = create_test_meta("test2.parquet", 200); + let cached_value_2 = CachedFileMetadata::new( + meta2.clone(), + Arc::clone(&schema_fingerprint), + Arc::new(Statistics::new_unknown(&schema)), + Some(ordering()), + ); + + let path_2 = TableScopedPath { + path: meta2.location.clone(), + table: None, + }; + + cache.put(&path_2, cached_value_2.clone()); + + let entries = cache.list_entries(); + assert_eq!( + entries, + HashMap::from([ + ( + path_1, + CacheEntryInfo { + value: cached_value_1, + hits: 0, + size_bytes: 373, + expires: None, + } + ), + ( + path_2, + CacheEntryInfo { + value: cached_value_2, + hits: 0, + size_bytes: 373, + expires: None, + } + ), + ]) + ); + } + + #[test] + fn test_cache_entry_added_when_entries_are_within_cache_limit() { + let (meta_1, value_1) = + create_cached_file_metadata_with_stats("test1.parquet", 10); + let (meta_2, value_2) = + create_cached_file_metadata_with_stats("test2.parquet", 10); + let (meta_3, value_3) = + create_cached_file_metadata_with_stats("test3.parquet", 10); + + let mut ctx = DFHeapSizeCtx::default(); + + let limit_for_2_entries = meta_1.location.as_ref().heap_size(&mut ctx) + + value_1.heap_size(&mut ctx) + + meta_2.location.as_ref().heap_size(&mut ctx) + + value_2.heap_size(&mut ctx); + + // create a cache with a limit which fits exactly 2 entries + let cache = DefaultCache::new(limit_for_2_entries); + let path_1 = TableScopedPath { + path: meta_1.location.clone(), + table: None, + }; + + let path_2 = TableScopedPath { + path: meta_2.location.clone(), + table: None, + }; + + cache.put(&path_1, value_1.clone()); + cache.put(&path_2, value_2.clone()); + + assert_eq!(cache.len(), 2); + assert_eq!(cache.memory_used(), limit_for_2_entries); + + let result_1 = cache.get(&path_1); + let result_2 = cache.get(&path_2); + assert_eq!(result_1.unwrap(), value_1); + assert_eq!(result_2.unwrap(), value_2); + + let path_3 = TableScopedPath { + path: meta_3.location.clone(), + table: None, + }; + + // adding the third entry evicts the first entry + cache.put(&path_3, value_3.clone()); + assert_eq!(cache.len(), 2); + assert_eq!(cache.memory_used(), limit_for_2_entries); + + let result_1 = cache.get(&path_1); + assert!(result_1.is_none()); + + let result_2 = cache.get(&path_2); + let result_3 = cache.get(&path_3); + + assert_eq!(result_2.unwrap(), value_2); + assert_eq!(result_3.unwrap(), value_3); + + // add the third entry again, making sure memory usage remains the same + cache.put(&path_3, value_3.clone()); + assert_eq!(cache.memory_used(), limit_for_2_entries); + cache.put(&path_3, value_3.clone()); + assert_eq!(cache.memory_used(), limit_for_2_entries); + + let mut ctx = DFHeapSizeCtx::default(); + cache.remove(&path_2); + assert_eq!(cache.len(), 1); + assert_eq!( + cache.memory_used(), + meta_3.location.as_ref().heap_size(&mut ctx) + value_3.heap_size(&mut ctx) + ); + + cache.clear(); + assert_eq!(cache.len(), 0); + assert_eq!(cache.memory_used(), 0); + } + + #[test] + fn test_cache_rejects_entry_which_is_too_large() { + let (meta, value_too_large) = + create_cached_file_metadata_with_stats("test1.parquet", 10); + let mut ctx = DFHeapSizeCtx::default(); + let limit_less_than_the_entry = value_too_large.clone().heap_size(&mut ctx) - 1; + + // create a cache with a size less than the entry + let cache = DefaultCache::new(limit_less_than_the_entry); + + let path_1 = TableScopedPath { + path: meta.location.clone(), + table: None, + }; + + cache.put(&path_1, value_too_large.clone()); + + assert_eq!(cache.len(), 0); + assert_eq!(cache.memory_used(), 0); + + // Test stale entry is removed when oversized entry is added + let (_, value_fits) = create_cached_file_metadata_with_stats("test1.parquet", 7); + cache.put(&path_1, value_fits.clone()); + + assert_eq!(cache.len(), 1); + assert_eq!(cache.memory_used(), 1514); + + // now add an entry which is over the limit and make sure the old stale entry is removed + let stale_entry = cache.put(&path_1, value_too_large.clone()); + assert_eq!(stale_entry, Some(value_fits)); + + assert_eq!(cache.len(), 0); + assert_eq!(cache.memory_used(), 0); + } + + fn create_cached_file_metadata_with_stats( + file_name: &str, + series_size: i32, + ) -> (ObjectMeta, CachedFileMetadata) { + let series: Vec = (0..=series_size).collect(); + let values = Int32Array::from(series); + let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, series_size + 1])); + let field = Arc::new(Field::new_list_field(DataType::Int32, false)); + let list_array = ListArray::new(field, offsets, Arc::new(values), None); + + let column_statistics = ColumnStatistics { + null_count: Precision::Exact(1), + max_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), + min_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), + sum_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), + distinct_count: Precision::Exact(10), + byte_size: Precision::Absent, + }; + + let stats = Statistics { + num_rows: Precision::Exact(100), + total_byte_size: Precision::Exact(100), + column_statistics: vec![column_statistics.clone()], + }; + let mut ctx = DFHeapSizeCtx::default(); + let object_meta = create_test_meta(file_name, stats.heap_size(&mut ctx) as u64); + let schema = Schema::new(vec![Field::new("list", DataType::Int32, true)]); + let schema_fingerprint = Arc::new(SchemaFingerprint::from_schema(&schema)); + let value = CachedFileMetadata::new( + object_meta.clone(), + schema_fingerprint, + Arc::new(stats.clone()), + None, + ); + (object_meta, value) + } + + struct MockTimeProvider { + base: Instant, + offset: Mutex, + } + + impl MockTimeProvider { + fn new() -> Self { + Self { + base: Instant::now(), + offset: Mutex::new(Duration::ZERO), + } + } + + fn inc(&self, duration: Duration) { + let mut offset = self.offset.lock().unwrap(); + *offset += duration; + } + } + + impl TimeProvider for MockTimeProvider { + fn now(&self) -> Instant { + self.base + *self.offset.lock().unwrap() + } + } + + /// Helper function to create a test ObjectMeta with a specific path and location string size + fn create_object_meta(path: &str, location_size: usize) -> ObjectMeta { + // Create a location string of the desired size by padding with zeros + let location_str = if location_size > path.len() { + format!("{}{}", path, "0".repeat(location_size - path.len())) + } else { + path.to_string() + }; + + ObjectMeta { + location: Path::from(location_str), + last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") + .unwrap() + .into(), + size: 1024, + e_tag: None, + version: None, + } + } + + /// Helper function to create a TableScopedPath and a CachedFileList with at least meta_size bytes + fn create_test_list_files_entry( + path: &str, + count: usize, + meta_size: usize, + table: Option, + ) -> (TableScopedPath, CachedFileList) { + let key = TableScopedPath { + table, + path: Path::from(path), + }; + let metas: Vec = (0..count) + .map(|i| create_object_meta(&format!("file{i}"), meta_size)) + .collect(); + let value = CachedFileList::new(metas); + (key, value) + } + + #[test] + fn test_basic_operations() { + let cache = DefaultCache::new(DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT); + let table_ref = Some(TableReference::from("table")); + let path = Path::from("test_path"); + let key = TableScopedPath { + table: table_ref.clone(), + path, + }; + + // Initially cache is empty + assert!(!cache.contains_key(&key)); + assert_eq!(cache.len(), 0); + + // Cache miss - get returns None + assert!(cache.get(&key).is_none()); + + // Put a value + let meta = create_test_object_meta("file1", 50); + cache.put(&key, CachedFileList::new(vec![meta])); + + // Entry should be cached + assert!(cache.contains_key(&key)); + assert_eq!(cache.len(), 1); + let result = cache.get(&key).unwrap(); + assert_eq!(result.files.len(), 1); + + // Remove the entry + let removed = cache.remove(&key).unwrap(); + assert_eq!(removed.files.len(), 1); + assert!(!cache.contains_key(&key)); + assert_eq!(cache.len(), 0); + + // Put multiple entries + let (key1, value1) = + create_test_list_files_entry("path1", 2, 50, table_ref.clone()); + let (key2, value2) = create_test_list_files_entry("path2", 3, 50, table_ref); + cache.put(&key1, value1.clone()); + cache.put(&key2, value2.clone()); + assert_eq!(cache.len(), 2); + + // List cache entries + assert_eq!( + cache.list_entries(), + HashMap::from([ + ( + key1.clone(), + CacheEntryInfo { + value: value1.clone(), + size_bytes: value1.size(), + hits: 0, + expires: None, + } + ), + ( + key2.clone(), + CacheEntryInfo { + value: value2.clone(), + size_bytes: value2.size(), + hits: 0, + expires: None, + } + ) + ]) + ); + + // Clear all entries + cache.clear(); + assert_eq!(cache.len(), 0); + assert!(!cache.contains_key(&key1)); + assert!(!cache.contains_key(&key2)); + } + + #[test] + fn test_lru_eviction_basic() { + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 100, table_ref.clone()); + let (key2, value2) = + create_test_list_files_entry("path2", 1, 100, table_ref.clone()); + let (key3, value3) = + create_test_list_files_entry("path3", 1, 100, table_ref.clone()); + + let entry_size = key1.size() + value1.size(); + + // Set cache limit to exactly fit all 3 entries + let cache = DefaultCache::new(entry_size * 3); + + // All three entries should fit + cache.put(&key1, value1); + cache.put(&key2, value2); + cache.put(&key3, value3); + assert_eq!(cache.len(), 3); + assert!(cache.contains_key(&key1)); + assert!(cache.contains_key(&key2)); + assert!(cache.contains_key(&key3)); + + // Adding a new entry should evict path1 (LRU) + let (key4, value4) = create_test_list_files_entry("path4", 1, 100, table_ref); + cache.put(&key4, value4); + + assert_eq!(cache.len(), 3); + assert!(!cache.contains_key(&key1)); // Evicted + assert!(cache.contains_key(&key2)); + assert!(cache.contains_key(&key3)); + assert!(cache.contains_key(&key4)); + } + + #[test] + fn test_lru_ordering_after_access() { + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 100, table_ref.clone()); + let (key2, value2) = + create_test_list_files_entry("path2", 1, 100, table_ref.clone()); + let (key3, value3) = + create_test_list_files_entry("path3", 1, 100, table_ref.clone()); + + // Set cache limit to fit exactly three entries + let cache = DefaultCache::new((key1.size() + value1.size()) * 3); + + cache.put(&key1, value1); + cache.put(&key2, value2); + cache.put(&key3, value3); + assert_eq!(cache.len(), 3); + + // Access path1 to move it to front (MRU) + // Order is now: path2 (LRU), path3, path1 (MRU) + let _ = cache.get(&key1); + + // Adding a new entry should evict path2 (the LRU) + let (key4, value4) = create_test_list_files_entry("path4", 1, 100, table_ref); + cache.put(&key4, value4); + + assert_eq!(cache.len(), 3); + assert!(cache.contains_key(&key1)); // Still present (recently accessed) + assert!(!cache.contains_key(&key2)); // Evicted (was LRU) + assert!(cache.contains_key(&key3)); + assert!(cache.contains_key(&key4)); + } + + #[test] + fn test_reject_too_large() { + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 100, table_ref.clone()); + let (key2, value2) = + create_test_list_files_entry("path2", 1, 100, table_ref.clone()); + + // Set cache limit to fit both entries + let cache = DefaultCache::new((key1.size() + value1.size()) * 2); + + cache.put(&key1, value1); + cache.put(&key2, value2); + assert_eq!(cache.len(), 2); + + // Try to add an entry that's too large to fit in the cache + // The entry is not stored (too large) + let (key_large, value_large) = + create_test_list_files_entry("large", 1, 1000, table_ref); + cache.put(&key_large, value_large); + + // Large entry should not be added + assert!(!cache.contains_key(&key_large)); + assert_eq!(cache.len(), 2); + assert!(cache.contains_key(&key1)); + assert!(cache.contains_key(&key2)); + } + + #[test] + fn test_multiple_evictions() { + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 100, table_ref.clone()); + let (key2, value2) = + create_test_list_files_entry("path2", 1, 100, table_ref.clone()); + let (key3, value3) = + create_test_list_files_entry("path3", 1, 100, table_ref.clone()); + + let entry_size = key1.size() + value1.size(); + + // Set cache limit for exactly 3 entries + let cache = DefaultCache::new(entry_size * 3); + + cache.put(&key1, value1); + cache.put(&key2, value2); + cache.put(&key3, value3); + assert_eq!(cache.len(), 3); + + // Add a large entry that requires evicting 2 entries + let (key_large, value_large) = + create_test_list_files_entry("large", 1, 200, table_ref); + cache.put(&key_large, value_large); + + // path1 and path2 should be evicted (both LRU), path3 and path_large remain + assert_eq!(cache.len(), 2); + assert!(!cache.contains_key(&key1)); // Evicted + assert!(!cache.contains_key(&key2)); // Evicted + assert!(cache.contains_key(&key3)); + assert!(cache.contains_key(&key_large)); + } + + #[test] + fn test_cache_limit_resize() { + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 100, table_ref.clone()); + let (key2, value2) = + create_test_list_files_entry("path2", 1, 100, table_ref.clone()); + let (key3, value3) = create_test_list_files_entry("path3", 1, 100, table_ref); + + let entry_size = key1.size() + value1.size(); + + let cache = DefaultCache::new(entry_size * 3); + + // Add three entries + cache.put(&key1, value1); + cache.put(&key2, value2); + cache.put(&key3, value3); + assert_eq!(cache.len(), 3); + + // Resize cache to only fit one entry + cache.update_cache_limit(entry_size); + + // Should keep only the most recent entry (path3, the MRU) + assert_eq!(cache.len(), 1); + assert!(cache.contains_key(&key3)); + // Earlier entries (LRU) should be evicted + assert!(!cache.contains_key(&key1)); + assert!(!cache.contains_key(&key2)); + } + + #[test] + fn test_entry_update_with_size_change() { + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 100, table_ref.clone()); + let (key2, value2) = + create_test_list_files_entry("path2", 1, 100, table_ref.clone()); + let (key3, value3_v1) = + create_test_list_files_entry("path3", 1, 100, table_ref.clone()); + + let entry_size = key1.size() + value1.size(); + + let cache = DefaultCache::new(entry_size * 3); + + // Add three entries + cache.put(&key1, value1); + cache.put(&key2, value2.clone()); + cache.put(&key3, value3_v1); + assert_eq!(cache.len(), 3); + + // Update path3 with same size - should not cause eviction + let (_, value3_v2) = + create_test_list_files_entry("path3", 1, 100, table_ref.clone()); + cache.put(&key3, value3_v2); + + assert_eq!(cache.len(), 3); + assert!(cache.contains_key(&key1)); + assert!(cache.contains_key(&key2)); + assert!(cache.contains_key(&key3)); + + // Update path3 with larger size that requires evicting path1 (LRU) + let (_, value3_v3) = create_test_list_files_entry("path3", 1, 200, table_ref); + cache.put(&key3, value3_v3.clone()); + + assert_eq!(cache.len(), 2); + assert!(!cache.contains_key(&key1)); // Evicted (was LRU) + assert!(cache.contains_key(&key2)); + assert!(cache.contains_key(&key3)); + + // List cache entries + assert_eq!( + cache.list_entries(), + HashMap::from([ + ( + key2, + CacheEntryInfo { + value: value2.clone(), + size_bytes: value2.size(), + hits: 0, + expires: None, + } + ), + ( + key3, + CacheEntryInfo { + value: value3_v3.clone(), + size_bytes: value3_v3.size(), + hits: 0, + expires: None, + } + ) + ]) + ); + } + + #[test] + fn test_cache_with_ttl() { + let ttl = Duration::from_millis(100); + + let mock_time = Arc::new(MockTimeProvider::new()); + let cache = DefaultCache::new_with_ttl(10000, Some(ttl)) + .with_time_provider(Arc::clone(&mock_time) as Arc); + + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 2, 50, table_ref.clone()); + let (key2, value2) = create_test_list_files_entry("path2", 2, 50, table_ref); + cache.put(&key1, value1.clone()); + cache.put(&key2, value2.clone()); + + // Entries should be accessible immediately + assert!(cache.get(&key1).is_some()); + assert!(cache.get(&key2).is_some()); + // List cache entries + assert_eq!( + cache.list_entries(), + HashMap::from([ + ( + key1.clone(), + CacheEntryInfo { + value: value1.clone(), + size_bytes: value1.size(), + hits: 1, + expires: mock_time.now().checked_add(ttl), + } + ), + ( + key2.clone(), + CacheEntryInfo { + value: value2.clone(), + size_bytes: value2.size(), + hits: 1, + expires: mock_time.now().checked_add(ttl), + } + ) + ]) + ); + // Wait for TTL to expire + mock_time.inc(Duration::from_millis(150)); + + // Entries should now return None when observed through contains_key + assert!(!cache.contains_key(&key1)); + assert_eq!(cache.len(), 1); // key1 was removed by contains_key() + assert!(!cache.contains_key(&key2)); + assert_eq!(cache.len(), 0); // key2 was removed by contains_key() + } + + #[test] + fn test_cache_with_ttl_and_lru() { + let ttl = Duration::from_millis(200); + + let mock_time = Arc::new(MockTimeProvider::new()); + let cache = DefaultCache::new_with_ttl(1100, Some(ttl)) + .with_time_provider(Arc::clone(&mock_time) as Arc); + + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 400, table_ref.clone()); + let (key2, value2) = + create_test_list_files_entry("path2", 1, 400, table_ref.clone()); + + let (key3, value3) = create_test_list_files_entry("path3", 1, 400, table_ref); + cache.put(&key1, value1); + mock_time.inc(Duration::from_millis(50)); + cache.put(&key2, value2); + mock_time.inc(Duration::from_millis(50)); + + // path3 should evict path1 due to size limit + cache.put(&key3, value3); + assert!(!cache.contains_key(&key1)); // Evicted by LRU + assert!(cache.contains_key(&key2)); + assert!(cache.contains_key(&key3)); + + mock_time.inc(Duration::from_millis(151)); + + assert!(!cache.contains_key(&key2)); // Expired + assert!(cache.contains_key(&key3)); // Still valid + } + + #[test] + fn test_ttl_expiration_in_get() { + let ttl = Duration::from_millis(100); + let cache = DefaultCache::new_with_ttl(10000, Some(ttl)); + + let table_ref = Some(TableReference::from("table")); + let (key, value) = create_test_list_files_entry("path", 2, 50, table_ref); + + // Cache the entry + cache.put(&key, value.clone()); + + // Entry should be accessible immediately + let result = cache.get(&key); + assert!(result.is_some()); + assert_eq!(result.unwrap().files.len(), 2); + + // Wait for TTL to expire + thread::sleep(Duration::from_millis(150)); + + // Get should return None because entry expired + let result2 = cache.get(&key); + assert!(result2.is_none()); + } + + #[test] + fn test_meta_heap_bytes_calculation() { + // Test with minimal ObjectMeta (no e_tag, no version) + let meta1 = ObjectMeta { + location: Path::from("test"), + last_modified: chrono::Utc::now(), + size: 100, + e_tag: None, + version: None, + }; + assert_eq!(meta_heap_bytes(&meta1), 4); // Just the location string "test" + + // Test with e_tag + let meta2 = ObjectMeta { + location: Path::from("test"), + last_modified: chrono::Utc::now(), + size: 100, + e_tag: Some("etag123".to_string()), + version: None, + }; + assert_eq!(meta_heap_bytes(&meta2), 4 + 7); // location (4) + e_tag (7) + + // Test with version + let meta3 = ObjectMeta { + location: Path::from("test"), + last_modified: chrono::Utc::now(), + size: 100, + e_tag: None, + version: Some("v1.0".to_string()), + }; + assert_eq!(meta_heap_bytes(&meta3), 4 + 4); // location (4) + version (4) + + // Test with both e_tag and version + let meta4 = ObjectMeta { + location: Path::from("test"), + last_modified: chrono::Utc::now(), + size: 100, + e_tag: Some("tag".to_string()), + version: Some("ver".to_string()), + }; + assert_eq!(meta_heap_bytes(&meta4), 4 + 3 + 3); // location (4) + e_tag (3) + version (3) + } + + #[test] + fn test_memory_tracking() { + let cache = DefaultCache::new(1000); + + // Verify cache starts with 0 memory used + { + assert_eq!(cache.memory_used(), 0); + } + + // Add entry and verify memory tracking + let table_ref = Some(TableReference::from("table")); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 100, table_ref.clone()); + cache.put(&key1, value1.clone()); + let entry_size_1 = key1.size() + value1.size(); + { + assert_eq!(cache.memory_used(), entry_size_1); + } + + // Add another entry + let (key2, value2) = + create_test_list_files_entry("path2", 1, 200, table_ref.clone()); + cache.put(&key2, value2.clone()); + let entry_size_2 = key2.size() + value2.size(); + + { + assert_eq!(cache.memory_used(), entry_size_1 + entry_size_2); + } + + // Remove first entry and verify memory decreases + cache.remove(&key1); + { + assert_eq!(cache.memory_used(), entry_size_2); + } + + // Clear and verify memory is 0 + cache.clear(); + { + assert_eq!(cache.memory_used(), 0); + } + } + + // Prefix filtering tests using CachedFileList::filter_by_prefix + + /// Helper function to create ObjectMeta with a specific location path + fn create_object_meta_with_path(location: &str) -> ObjectMeta { + ObjectMeta { + location: Path::from(location), + last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") + .unwrap() + .into(), + size: 1024, + e_tag: None, + version: None, + } + } + + #[test] + fn test_prefix_filtering() { + let cache = DefaultCache::new(100000); + + // Create files for a partitioned table + let table_base = Path::from("my_table"); + let files = vec![ + create_object_meta_with_path("my_table/a=1/file1.parquet"), + create_object_meta_with_path("my_table/a=1/file2.parquet"), + create_object_meta_with_path("my_table/a=2/file3.parquet"), + create_object_meta_with_path("my_table/a=2/file4.parquet"), + ]; + + // Cache the full table listing + let table_ref = Some(TableReference::from("table")); + let key = TableScopedPath { + table: table_ref, + path: table_base, + }; + cache.put(&key, CachedFileList::new(files)); + + let result = cache.get(&key).unwrap(); + + // Filter for partition a=1 + let prefix_a1 = Some(Path::from("my_table/a=1")); + let filtered = result.files_matching_prefix(&prefix_a1); + assert_eq!(filtered.len(), 2); + assert!( + filtered + .iter() + .all(|m| m.location.as_ref().starts_with("my_table/a=1")) + ); + + // Filter for partition a=2 + let prefix_a2 = Some(Path::from("my_table/a=2")); + let filtered_2 = result.files_matching_prefix(&prefix_a2); + assert_eq!(filtered_2.len(), 2); + assert!( + filtered_2 + .iter() + .all(|m| m.location.as_ref().starts_with("my_table/a=2")) + ); + + // No filter returns all + let all = result.files_matching_prefix(&None); + assert_eq!(all.len(), 4); + } + + #[test] + fn test_prefix_no_matching_files() { + let cache = DefaultCache::new(100000); + + let table_base = Path::from("my_table"); + let files = vec![ + create_object_meta_with_path("my_table/a=1/file1.parquet"), + create_object_meta_with_path("my_table/a=2/file2.parquet"), + ]; + + let table_ref = Some(TableReference::from("table")); + let key = TableScopedPath { + table: table_ref, + path: table_base, + }; + cache.put(&key, CachedFileList::new(files)); + let result = cache.get(&key).unwrap(); + + // Query for partition a=3 which doesn't exist + let prefix_a3 = Some(Path::from("my_table/a=3")); + let filtered = result.files_matching_prefix(&prefix_a3); + assert!(filtered.is_empty()); + } + + #[test] + fn test_nested_partitions() { + let cache = DefaultCache::new(100000); + + let table_base = Path::from("events"); + let files = vec![ + create_object_meta_with_path( + "events/year=2024/month=01/day=01/file1.parquet", + ), + create_object_meta_with_path( + "events/year=2024/month=01/day=02/file2.parquet", + ), + create_object_meta_with_path( + "events/year=2024/month=02/day=01/file3.parquet", + ), + create_object_meta_with_path( + "events/year=2025/month=01/day=01/file4.parquet", + ), + ]; + + let table_ref = Some(TableReference::from("table")); + let key = TableScopedPath { + table: table_ref, + path: table_base, + }; + cache.put(&key, CachedFileList::new(files)); + let result = cache.get(&key).unwrap(); + + // Filter for year=2024/month=01 + let prefix_month = Some(Path::from("events/year=2024/month=01")); + let filtered = result.files_matching_prefix(&prefix_month); + assert_eq!(filtered.len(), 2); + + // Filter for year=2024 + let prefix_year = Some(Path::from("events/year=2024")); + let filtered_year = result.files_matching_prefix(&prefix_year); + assert_eq!(filtered_year.len(), 3); + } + + #[test] + fn test_drop_table_entries() { + let cache = DefaultCache::new(DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT); + + let table_ref1 = TableReference::from("table1"); + let table_ref2 = TableReference::from("table2"); + let (key1, value1) = + create_test_list_files_entry("path1", 1, 100, Some(table_ref1.clone())); + let (key2, value2) = + create_test_list_files_entry("path2", 1, 100, Some(table_ref1.clone())); + let (key3, value3) = + create_test_list_files_entry("path3", 1, 100, Some(table_ref2.clone())); + + cache.put(&key1, value1); + cache.put(&key2, value2); + cache.put(&key3, value3); + + cache.drop_table_entries(&table_ref1).unwrap(); + + assert!(!cache.contains_key(&key1)); + assert!(!cache.contains_key(&key2)); + assert!(cache.contains_key(&key3)); + } +} diff --git a/datafusion/execution/src/cache/file_metadata_cache.rs b/datafusion/execution/src/cache/file_metadata_cache.rs deleted file mode 100644 index 5e899d7dd9f8b..0000000000000 --- a/datafusion/execution/src/cache/file_metadata_cache.rs +++ /dev/null @@ -1,764 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::{collections::HashMap, sync::Mutex}; - -use object_store::path::Path; - -use crate::cache::{ - CacheAccessor, - cache_manager::{CachedFileMetadataEntry, FileMetadataCache, FileMetadataCacheEntry}, - lru_queue::LruQueue, -}; - -/// Handles the inner state of the [`DefaultFilesMetadataCache`] struct. -struct DefaultFilesMetadataCacheState { - lru_queue: LruQueue, - memory_limit: usize, - memory_used: usize, - cache_hits: HashMap, -} - -impl DefaultFilesMetadataCacheState { - fn new(memory_limit: usize) -> Self { - Self { - lru_queue: LruQueue::new(), - memory_limit, - memory_used: 0, - cache_hits: HashMap::new(), - } - } - - /// Returns the respective entry from the cache, if it exists. - /// If the entry exists, it becomes the most recently used. - fn get(&mut self, k: &Path) -> Option { - self.lru_queue.get(k).cloned().inspect(|_| { - *self.cache_hits.entry(k.clone()).or_insert(0) += 1; - }) - } - - /// Checks if the metadata is currently cached. - /// The LRU queue is not updated. - fn contains_key(&self, k: &Path) -> bool { - self.lru_queue.peek(k).is_some() - } - - /// Adds a new key-value pair to cache, meaning LRU entries might be evicted if required. - /// If the key is already in the cache, the previous metadata is returned. - /// If the size of the metadata is greater than the `memory_limit`, the value is not inserted. - fn put( - &mut self, - key: Path, - value: CachedFileMetadataEntry, - ) -> Option { - let value_size = value.file_metadata.memory_size(); - - // no point in trying to add this value to the cache if it cannot fit entirely - if value_size > self.memory_limit { - return None; - } - - self.cache_hits.insert(key.clone(), 0); - // if the key is already in the cache, the old value is removed - let old_value = self.lru_queue.put(key, value); - self.memory_used += value_size; - if let Some(ref old_entry) = old_value { - self.memory_used -= old_entry.file_metadata.memory_size(); - } - - self.evict_entries(); - - old_value - } - - /// Evicts entries from the LRU cache until `memory_used` is lower than `memory_limit`. - fn evict_entries(&mut self) { - while self.memory_used > self.memory_limit { - if let Some(removed) = self.lru_queue.pop() { - self.memory_used -= removed.1.file_metadata.memory_size(); - } else { - // cache is empty while memory_used > memory_limit, cannot happen - debug_assert!( - false, - "cache is empty while memory_used > memory_limit, cannot happen" - ); - return; - } - } - } - - /// Removes an entry from the cache and returns it, if it exists. - fn remove(&mut self, k: &Path) -> Option { - if let Some(old_entry) = self.lru_queue.remove(k) { - self.memory_used -= old_entry.file_metadata.memory_size(); - self.cache_hits.remove(k); - Some(old_entry) - } else { - None - } - } - - /// Returns the number of entries currently cached. - fn len(&self) -> usize { - self.lru_queue.len() - } - - /// Removes all entries from the cache. - fn clear(&mut self) { - self.lru_queue.clear(); - self.memory_used = 0; - self.cache_hits.clear(); - } -} - -/// Default implementation of [`FileMetadataCache`] -/// -/// Collected file embedded metadata cache. -/// -/// The metadata for each file is validated by comparing the cached [`ObjectMeta`] -/// (size and last_modified) against the current file state using `cached.is_valid_for(¤t_meta)`. -/// -/// # Internal details -/// -/// The `memory_limit` controls the maximum size of the cache, which uses a -/// Least Recently Used eviction algorithm. When adding a new entry, if the total -/// size of the cached entries exceeds `memory_limit`, the least recently used entries -/// are evicted until the total size is lower than `memory_limit`. -/// -/// [`ObjectMeta`]: object_store::ObjectMeta -pub struct DefaultFilesMetadataCache { - // the state is wrapped in a Mutex to ensure the operations are atomic - state: Mutex, -} - -impl DefaultFilesMetadataCache { - /// Create a new instance of [`DefaultFilesMetadataCache`]. - /// - /// # Arguments - /// `memory_limit`: the maximum size of the cache, in bytes - // - pub fn new(memory_limit: usize) -> Self { - Self { - state: Mutex::new(DefaultFilesMetadataCacheState::new(memory_limit)), - } - } - - /// Returns the size of the cached memory, in bytes. - pub fn memory_used(&self) -> usize { - let state = self.state.lock().unwrap(); - state.memory_used - } -} - -impl CacheAccessor for DefaultFilesMetadataCache { - fn get(&self, key: &Path) -> Option { - let mut state = self.state.lock().unwrap(); - state.get(key) - } - - fn put( - &self, - key: &Path, - value: CachedFileMetadataEntry, - ) -> Option { - let mut state = self.state.lock().unwrap(); - state.put(key.clone(), value) - } - - fn remove(&self, k: &Path) -> Option { - let mut state = self.state.lock().unwrap(); - state.remove(k) - } - - fn contains_key(&self, k: &Path) -> bool { - let state = self.state.lock().unwrap(); - state.contains_key(k) - } - - fn len(&self) -> usize { - let state = self.state.lock().unwrap(); - state.len() - } - - fn clear(&self) { - let mut state = self.state.lock().unwrap(); - state.clear(); - } - - fn name(&self) -> String { - "DefaultFilesMetadataCache".to_string() - } -} - -impl FileMetadataCache for DefaultFilesMetadataCache { - fn cache_limit(&self) -> usize { - let state = self.state.lock().unwrap(); - state.memory_limit - } - - fn update_cache_limit(&self, limit: usize) { - let mut state = self.state.lock().unwrap(); - state.memory_limit = limit; - state.evict_entries(); - } - - fn list_entries(&self) -> HashMap { - let state = self.state.lock().unwrap(); - let mut entries = HashMap::::new(); - - for (path, entry) in state.lru_queue.list_entries() { - entries.insert( - path.clone(), - FileMetadataCacheEntry { - object_meta: entry.meta.clone(), - size_bytes: entry.file_metadata.memory_size(), - hits: *state.cache_hits.get(path).expect("entry must exist"), - extra: entry.file_metadata.extra_info(), - }, - ); - } - - entries - } -} - -#[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::Arc; - - use crate::cache::CacheAccessor; - use crate::cache::cache_manager::{ - CachedFileMetadataEntry, FileMetadata, FileMetadataCache, FileMetadataCacheEntry, - }; - use crate::cache::file_metadata_cache::DefaultFilesMetadataCache; - use object_store::ObjectMeta; - use object_store::path::Path; - - pub struct TestFileMetadata { - metadata: String, - } - - impl FileMetadata for TestFileMetadata { - fn as_any(&self) -> &dyn std::any::Any { - self - } - - fn memory_size(&self) -> usize { - self.metadata.len() - } - - fn extra_info(&self) -> HashMap { - HashMap::from([("extra_info".to_owned(), "abc".to_owned())]) - } - } - - fn create_test_object_meta(path: &str, size: usize) -> ObjectMeta { - ObjectMeta { - location: Path::from(path), - last_modified: chrono::DateTime::parse_from_rfc3339( - "2025-07-29T12:12:12+00:00", - ) - .unwrap() - .into(), - size: size as u64, - e_tag: None, - version: None, - } - } - - #[test] - fn test_default_file_metadata_cache() { - let object_meta = create_test_object_meta("test", 1024); - - let metadata: Arc = Arc::new(TestFileMetadata { - metadata: "retrieved_metadata".to_owned(), - }); - - let cache = DefaultFilesMetadataCache::new(1024 * 1024); - - // Cache miss - assert!(cache.get(&object_meta.location).is_none()); - - // Put a value - let cached_entry = - CachedFileMetadataEntry::new(object_meta.clone(), Arc::clone(&metadata)); - cache.put(&object_meta.location, cached_entry); - - // Verify the cached value - assert!(cache.contains_key(&object_meta.location)); - let result = cache.get(&object_meta.location).unwrap(); - let test_file_metadata = Arc::downcast::(result.file_metadata); - assert!(test_file_metadata.is_ok()); - assert_eq!(test_file_metadata.unwrap().metadata, "retrieved_metadata"); - - // Cache hit - check validation - let result2 = cache.get(&object_meta.location).unwrap(); - assert!(result2.is_valid_for(&object_meta)); - - // File size changed - closure should detect invalidity - let object_meta2 = create_test_object_meta("test", 2048); - let result3 = cache.get(&object_meta2.location).unwrap(); - // Cached entry should NOT be valid for new meta - assert!(!result3.is_valid_for(&object_meta2)); - - // Return new entry - let new_entry = - CachedFileMetadataEntry::new(object_meta2.clone(), Arc::clone(&metadata)); - cache.put(&object_meta2.location, new_entry); - - let result4 = cache.get(&object_meta2.location).unwrap(); - assert_eq!(result4.meta.size, 2048); - - // remove - cache.remove(&object_meta.location); - assert!(!cache.contains_key(&object_meta.location)); - - // len and clear - let object_meta3 = create_test_object_meta("test3", 100); - cache.put( - &object_meta.location, - CachedFileMetadataEntry::new(object_meta.clone(), Arc::clone(&metadata)), - ); - cache.put( - &object_meta3.location, - CachedFileMetadataEntry::new(object_meta3.clone(), Arc::clone(&metadata)), - ); - assert_eq!(cache.len(), 2); - cache.clear(); - assert_eq!(cache.len(), 0); - } - - fn generate_test_metadata_with_size( - path: &str, - size: usize, - ) -> (ObjectMeta, Arc) { - let object_meta = ObjectMeta { - location: Path::from(path), - last_modified: chrono::Utc::now(), - size: size as u64, - e_tag: None, - version: None, - }; - let metadata: Arc = Arc::new(TestFileMetadata { - metadata: "a".repeat(size), - }); - - (object_meta, metadata) - } - - #[test] - fn test_default_file_metadata_cache_with_limit() { - let cache = DefaultFilesMetadataCache::new(1000); - let (object_meta1, metadata1) = generate_test_metadata_with_size("1", 100); - let (object_meta2, metadata2) = generate_test_metadata_with_size("2", 500); - let (object_meta3, metadata3) = generate_test_metadata_with_size("3", 300); - - cache.put( - &object_meta1.location, - CachedFileMetadataEntry::new(object_meta1.clone(), metadata1), - ); - cache.put( - &object_meta2.location, - CachedFileMetadataEntry::new(object_meta2.clone(), metadata2), - ); - cache.put( - &object_meta3.location, - CachedFileMetadataEntry::new(object_meta3.clone(), metadata3), - ); - - // all entries will fit - assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 900); - assert!(cache.contains_key(&object_meta1.location)); - assert!(cache.contains_key(&object_meta2.location)); - assert!(cache.contains_key(&object_meta3.location)); - - // add a new entry which will remove the least recently used ("1") - let (object_meta4, metadata4) = generate_test_metadata_with_size("4", 200); - cache.put( - &object_meta4.location, - CachedFileMetadataEntry::new(object_meta4.clone(), metadata4), - ); - assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 1000); - assert!(!cache.contains_key(&object_meta1.location)); - assert!(cache.contains_key(&object_meta4.location)); - - // get entry "2", which will move it to the top of the queue, and add a new one which will - // remove the new least recently used ("3") - let _ = cache.get(&object_meta2.location); - let (object_meta5, metadata5) = generate_test_metadata_with_size("5", 100); - cache.put( - &object_meta5.location, - CachedFileMetadataEntry::new(object_meta5.clone(), metadata5), - ); - assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 800); - assert!(!cache.contains_key(&object_meta3.location)); - assert!(cache.contains_key(&object_meta5.location)); - - // new entry which will not be able to fit in the 1000 bytes allocated - let (object_meta6, metadata6) = generate_test_metadata_with_size("6", 1200); - cache.put( - &object_meta6.location, - CachedFileMetadataEntry::new(object_meta6.clone(), metadata6), - ); - assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 800); - assert!(!cache.contains_key(&object_meta6.location)); - - // new entry which is able to fit without removing any entry - let (object_meta7, metadata7) = generate_test_metadata_with_size("7", 200); - cache.put( - &object_meta7.location, - CachedFileMetadataEntry::new(object_meta7.clone(), metadata7), - ); - assert_eq!(cache.len(), 4); - assert_eq!(cache.memory_used(), 1000); - assert!(cache.contains_key(&object_meta7.location)); - - // new entry which will remove all other entries - let (object_meta8, metadata8) = generate_test_metadata_with_size("8", 999); - cache.put( - &object_meta8.location, - CachedFileMetadataEntry::new(object_meta8.clone(), metadata8), - ); - assert_eq!(cache.len(), 1); - assert_eq!(cache.memory_used(), 999); - assert!(cache.contains_key(&object_meta8.location)); - - // when updating an entry, the previous ones are not unnecessarily removed - let (object_meta9, metadata9) = generate_test_metadata_with_size("9", 300); - let (object_meta10, metadata10) = generate_test_metadata_with_size("10", 200); - let (object_meta11_v1, metadata11_v1) = - generate_test_metadata_with_size("11", 400); - cache.put( - &object_meta9.location, - CachedFileMetadataEntry::new(object_meta9.clone(), metadata9), - ); - cache.put( - &object_meta10.location, - CachedFileMetadataEntry::new(object_meta10.clone(), metadata10), - ); - cache.put( - &object_meta11_v1.location, - CachedFileMetadataEntry::new(object_meta11_v1.clone(), metadata11_v1), - ); - assert_eq!(cache.memory_used(), 900); - assert_eq!(cache.len(), 3); - let (object_meta11_v2, metadata11_v2) = - generate_test_metadata_with_size("11", 500); - cache.put( - &object_meta11_v2.location, - CachedFileMetadataEntry::new(object_meta11_v2.clone(), metadata11_v2), - ); - assert_eq!(cache.memory_used(), 1000); - assert_eq!(cache.len(), 3); - assert!(cache.contains_key(&object_meta9.location)); - assert!(cache.contains_key(&object_meta10.location)); - assert!(cache.contains_key(&object_meta11_v2.location)); - - // when updating an entry that now exceeds the limit, the LRU ("9") needs to be removed - let (object_meta11_v3, metadata11_v3) = - generate_test_metadata_with_size("11", 501); - cache.put( - &object_meta11_v3.location, - CachedFileMetadataEntry::new(object_meta11_v3.clone(), metadata11_v3), - ); - assert_eq!(cache.memory_used(), 701); - assert_eq!(cache.len(), 2); - assert!(cache.contains_key(&object_meta10.location)); - assert!(cache.contains_key(&object_meta11_v3.location)); - - // manually removing an entry that is not the LRU - cache.remove(&object_meta11_v3.location); - assert_eq!(cache.len(), 1); - assert_eq!(cache.memory_used(), 200); - assert!(cache.contains_key(&object_meta10.location)); - assert!(!cache.contains_key(&object_meta11_v3.location)); - - // clear - cache.clear(); - assert_eq!(cache.len(), 0); - assert_eq!(cache.memory_used(), 0); - - // resizing the cache should clear the extra entries - let (object_meta12, metadata12) = generate_test_metadata_with_size("12", 300); - let (object_meta13, metadata13) = generate_test_metadata_with_size("13", 200); - let (object_meta14, metadata14) = generate_test_metadata_with_size("14", 500); - cache.put( - &object_meta12.location, - CachedFileMetadataEntry::new(object_meta12.clone(), metadata12), - ); - cache.put( - &object_meta13.location, - CachedFileMetadataEntry::new(object_meta13.clone(), metadata13), - ); - cache.put( - &object_meta14.location, - CachedFileMetadataEntry::new(object_meta14.clone(), metadata14), - ); - assert_eq!(cache.len(), 3); - assert_eq!(cache.memory_used(), 1000); - cache.update_cache_limit(600); - assert_eq!(cache.len(), 1); - assert_eq!(cache.memory_used(), 500); - assert!(!cache.contains_key(&object_meta12.location)); - assert!(!cache.contains_key(&object_meta13.location)); - assert!(cache.contains_key(&object_meta14.location)); - } - - #[test] - fn test_default_file_metadata_cache_entries_info() { - let cache = DefaultFilesMetadataCache::new(1000); - let (object_meta1, metadata1) = generate_test_metadata_with_size("1", 100); - let (object_meta2, metadata2) = generate_test_metadata_with_size("2", 200); - let (object_meta3, metadata3) = generate_test_metadata_with_size("3", 300); - - // initial entries, all will have hits = 0 - cache.put( - &object_meta1.location, - CachedFileMetadataEntry::new(object_meta1.clone(), metadata1), - ); - cache.put( - &object_meta2.location, - CachedFileMetadataEntry::new(object_meta2.clone(), metadata2), - ); - cache.put( - &object_meta3.location, - CachedFileMetadataEntry::new(object_meta3.clone(), metadata3), - ); - assert_eq!( - cache.list_entries(), - HashMap::from([ - ( - Path::from("1"), - FileMetadataCacheEntry { - object_meta: object_meta1.clone(), - size_bytes: 100, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("2"), - FileMetadataCacheEntry { - object_meta: object_meta2.clone(), - size_bytes: 200, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("3"), - FileMetadataCacheEntry { - object_meta: object_meta3.clone(), - size_bytes: 300, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ) - ]) - ); - - // new hit on "1" - let _ = cache.get(&object_meta1.location); - assert_eq!( - cache.list_entries(), - HashMap::from([ - ( - Path::from("1"), - FileMetadataCacheEntry { - object_meta: object_meta1.clone(), - size_bytes: 100, - hits: 1, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("2"), - FileMetadataCacheEntry { - object_meta: object_meta2.clone(), - size_bytes: 200, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("3"), - FileMetadataCacheEntry { - object_meta: object_meta3.clone(), - size_bytes: 300, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ) - ]) - ); - - // new entry, will evict "2" - let (object_meta4, metadata4) = generate_test_metadata_with_size("4", 600); - cache.put( - &object_meta4.location, - CachedFileMetadataEntry::new(object_meta4.clone(), metadata4), - ); - assert_eq!( - cache.list_entries(), - HashMap::from([ - ( - Path::from("1"), - FileMetadataCacheEntry { - object_meta: object_meta1.clone(), - size_bytes: 100, - hits: 1, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("3"), - FileMetadataCacheEntry { - object_meta: object_meta3.clone(), - size_bytes: 300, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("4"), - FileMetadataCacheEntry { - object_meta: object_meta4.clone(), - size_bytes: 600, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ) - ]) - ); - - // replace entry "1" - let (object_meta1_new, metadata1_new) = generate_test_metadata_with_size("1", 50); - cache.put( - &object_meta1_new.location, - CachedFileMetadataEntry::new(object_meta1_new.clone(), metadata1_new), - ); - assert_eq!( - cache.list_entries(), - HashMap::from([ - ( - Path::from("1"), - FileMetadataCacheEntry { - object_meta: object_meta1_new.clone(), - size_bytes: 50, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("3"), - FileMetadataCacheEntry { - object_meta: object_meta3.clone(), - size_bytes: 300, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("4"), - FileMetadataCacheEntry { - object_meta: object_meta4.clone(), - size_bytes: 600, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ) - ]) - ); - - // remove entry "4" - cache.remove(&object_meta4.location); - assert_eq!( - cache.list_entries(), - HashMap::from([ - ( - Path::from("1"), - FileMetadataCacheEntry { - object_meta: object_meta1_new.clone(), - size_bytes: 50, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ), - ( - Path::from("3"), - FileMetadataCacheEntry { - object_meta: object_meta3.clone(), - size_bytes: 300, - hits: 0, - extra: HashMap::from([( - "extra_info".to_owned(), - "abc".to_owned() - )]), - } - ) - ]) - ); - - // clear - cache.clear(); - assert_eq!(cache.list_entries(), HashMap::from([])); - } -} diff --git a/datafusion/execution/src/cache/file_statistics_cache.rs b/datafusion/execution/src/cache/file_statistics_cache.rs deleted file mode 100644 index 12f0bb1b8af88..0000000000000 --- a/datafusion/execution/src/cache/file_statistics_cache.rs +++ /dev/null @@ -1,744 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use crate::cache::cache_manager::{ - CachedFileMetadata, FileStatisticsCache, FileStatisticsCacheEntry, -}; -use crate::cache::{CacheAccessor, TableScopedPath}; -use std::collections::HashMap; -use std::sync::Mutex; - -pub use crate::cache::DefaultFilesMetadataCache; -use crate::cache::lru_queue::LruQueue; -use datafusion_common::TableReference; -use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; - -/// Default implementation of [`FileStatisticsCache`] -/// -/// Stores cached file metadata (statistics and orderings) for files. -/// -/// The typical usage pattern is: -/// 1. Call `get(path)` to check for cached value -/// 2. If `Some(cached)`, validate with `cached.is_valid_for(¤t_meta)` -/// 3. If invalid or missing, compute new value and call `put(path, new_value)` -/// -/// # Internal details -/// -/// The `memory_limit` controls the maximum size of the cache, which uses a -/// Least Recently Used eviction algorithm. When adding a new entry, if the total -/// size of the cached entries exceeds `memory_limit`, the least recently used entries -/// are evicted until the total size is lower than `memory_limit`. -/// -/// -/// [`FileStatisticsCache`]: crate::cache::cache_manager::FileStatisticsCache -#[derive(Default)] -pub struct DefaultFileStatisticsCache { - state: Mutex, -} - -impl DefaultFileStatisticsCache { - pub fn new(memory_limit: usize) -> Self { - Self { - state: Mutex::new(DefaultFileStatisticsCacheState::new(memory_limit)), - } - } - - /// Returns the size of the cached memory, in bytes. - pub fn memory_used(&self) -> usize { - let state = self.state.lock().unwrap(); - state.memory_used - } -} - -struct DefaultFileStatisticsCacheState { - lru_queue: LruQueue, - memory_limit: usize, - memory_used: usize, -} - -pub const DEFAULT_FILE_STATISTICS_MEMORY_LIMIT: usize = 20 * 1024 * 1024; // 20MiB - -impl Default for DefaultFileStatisticsCacheState { - fn default() -> Self { - Self { - lru_queue: LruQueue::new(), - memory_limit: DEFAULT_FILE_STATISTICS_MEMORY_LIMIT, - memory_used: 0, - } - } -} - -impl DefaultFileStatisticsCacheState { - fn new(memory_limit: usize) -> Self { - Self { - lru_queue: LruQueue::new(), - memory_limit, - memory_used: 0, - } - } - fn get(&mut self, key: &TableScopedPath) -> Option { - self.lru_queue.get(key).cloned() - } - - fn put( - &mut self, - key: &TableScopedPath, - value: CachedFileMetadata, - ) -> Option { - let mut ctx = DFHeapSizeCtx::default(); - let key_size = key.heap_size(&mut ctx); - let entry_size = value.heap_size(&mut ctx); - - if entry_size + key_size > self.memory_limit { - // Remove potential stale entry - return self.remove(key); - } - - self.memory_used += entry_size; - self.memory_used += key_size; - - let old_value = self.lru_queue.put(key.clone(), value); - if let Some(old_entry) = &old_value { - let mut ctx = DFHeapSizeCtx::default(); - self.memory_used -= old_entry.heap_size(&mut ctx); - self.memory_used -= key_size; - } - - self.evict_entries(); - - old_value - } - - fn remove(&mut self, k: &TableScopedPath) -> Option { - if let Some(old_entry) = self.lru_queue.remove(k) { - let mut ctx = DFHeapSizeCtx::default(); - self.memory_used -= k.heap_size(&mut ctx); - self.memory_used -= old_entry.heap_size(&mut ctx); - Some(old_entry) - } else { - None - } - } - - fn contains_key(&self, k: &TableScopedPath) -> bool { - self.lru_queue.contains_key(k) - } - - fn len(&self) -> usize { - self.lru_queue.len() - } - - fn clear(&mut self) { - self.lru_queue.clear(); - self.memory_used = 0; - } - - fn evict_entries(&mut self) { - while self.memory_used > self.memory_limit { - if let Some(removed) = self.lru_queue.pop() { - let mut ctx = DFHeapSizeCtx::default(); - self.memory_used -= removed.0.heap_size(&mut ctx); - self.memory_used -= removed.1.heap_size(&mut ctx); - } else { - // cache is empty while memory_used > memory_limit, cannot happen - log::error!( - "File statistics cache memory accounting bug: memory_used={} but cache is empty. \ - Please report this to the Apache DataFusion developers.", - self.memory_used - ); - debug_assert!( - false, - "memory_used={} but cache is empty", - self.memory_used - ); - self.memory_used = 0; - return; - } - } - } -} -impl CacheAccessor for DefaultFileStatisticsCache { - fn get(&self, key: &TableScopedPath) -> Option { - let mut state = self.state.lock().unwrap(); - state.get(key) - } - - fn put( - &self, - key: &TableScopedPath, - value: CachedFileMetadata, - ) -> Option { - let mut state = self.state.lock().unwrap(); - state.put(key, value) - } - - fn remove(&self, key: &TableScopedPath) -> Option { - let mut state = self.state.lock().unwrap(); - state.remove(key) - } - - fn contains_key(&self, k: &TableScopedPath) -> bool { - let state = self.state.lock().unwrap(); - state.contains_key(k) - } - - fn len(&self) -> usize { - let state = self.state.lock().unwrap(); - state.len() - } - - fn clear(&self) { - let mut state = self.state.lock().unwrap(); - state.clear(); - } - - fn name(&self) -> String { - "DefaultFileStatisticsCache".to_string() - } -} - -impl FileStatisticsCache for DefaultFileStatisticsCache { - fn cache_limit(&self) -> usize { - let state = self.state.lock().unwrap(); - state.memory_limit - } - - fn update_cache_limit(&self, limit: usize) { - let mut state = self.state.lock().unwrap(); - state.memory_limit = limit; - state.evict_entries(); - } - - fn list_entries(&self) -> HashMap { - let mut entries = HashMap::::new(); - let mut ctx = DFHeapSizeCtx::default(); - for entry in self.state.lock().unwrap().lru_queue.list_entries() { - let path = entry.0.clone(); - let cached = entry.1; - entries.insert( - path, - FileStatisticsCacheEntry { - object_meta: cached.meta.clone(), - num_rows: cached.statistics.num_rows, - num_columns: cached.statistics.column_statistics.len(), - table_size_bytes: cached.statistics.total_byte_size, - statistics_size_bytes: cached.statistics.heap_size(&mut ctx), - has_ordering: cached.ordering.is_some(), - }, - ); - } - - entries - } - - fn drop_table_entries( - &self, - table_ref: &Option, - ) -> datafusion_common::Result<()> { - let mut state = self.state.lock().unwrap(); - let mut table_paths = vec![]; - for (path, _) in state.lru_queue.list_entries() { - if path.table == *table_ref { - table_paths.push(path.clone()); - } - } - for path in table_paths { - state.remove(&path); - } - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use crate::cache::cache_manager::{ - CachedFileMetadata, FileStatisticsCache, FileStatisticsCacheEntry, - }; - use arrow::array::{Int32Array, ListArray, RecordBatch}; - use arrow::buffer::{OffsetBuffer, ScalarBuffer}; - use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; - use chrono::DateTime; - use datafusion_common::heap_size::DFHeapSizeCtx; - use datafusion_common::stats::Precision; - use datafusion_common::{ColumnStatistics, ScalarValue, Statistics}; - use datafusion_expr::ColumnarValue; - use datafusion_physical_expr_common::physical_expr::PhysicalExpr; - use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; - use object_store::ObjectMeta; - use object_store::path::Path; - use std::sync::Arc; - - fn create_test_meta(path: &str, size: u64) -> ObjectMeta { - ObjectMeta { - location: Path::from(path), - last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") - .unwrap() - .into(), - size, - e_tag: None, - version: None, - } - } - - #[test] - fn test_statistics_cache() { - let meta = create_test_meta("test", 1024); - let cache = DefaultFileStatisticsCache::default(); - - let schema = Schema::new(vec![Field::new( - "test_column", - DataType::Timestamp(TimeUnit::Second, None), - false, - )]); - - let path = TableScopedPath { - path: meta.location.clone(), - table: None, - }; - - // Cache miss - assert!(cache.get(&path).is_none()); - - // Put a value - let cached_value = CachedFileMetadata::new( - meta.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&path, cached_value); - - // Cache hit - let result = cache.get(&path); - assert!(result.is_some()); - - let cached = result.unwrap(); - assert!(cached.is_valid_for(&meta)); - - // File size changed - validation should fail - let meta2 = create_test_meta("test", 2048); - - let path_2 = TableScopedPath { - path: meta2.location.clone(), - table: None, - }; - - let cached = cache.get(&path_2).unwrap(); - assert!(!cached.is_valid_for(&meta2)); - - // Update with new value - let cached_value2 = CachedFileMetadata::new( - meta2.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&path_2, cached_value2); - - // Test list_entries - let entries = cache.list_entries(); - assert_eq!(entries.len(), 1); - - let path_3 = TableScopedPath { - path: Path::from("test"), - table: None, - }; - - let entry = entries.get(&path_3).unwrap(); - assert_eq!(entry.object_meta.size, 2048); // Should be updated value - } - - #[derive(Clone, Debug, PartialEq, Eq, Hash)] - struct MockExpr {} - - impl std::fmt::Display for MockExpr { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "MockExpr") - } - } - - impl PhysicalExpr for MockExpr { - fn data_type( - &self, - _input_schema: &Schema, - ) -> datafusion_common::Result { - Ok(DataType::Int32) - } - - fn nullable(&self, _input_schema: &Schema) -> datafusion_common::Result { - Ok(false) - } - - fn evaluate( - &self, - _batch: &RecordBatch, - ) -> datafusion_common::Result { - unimplemented!() - } - - fn children(&self) -> Vec<&Arc> { - vec![] - } - - fn with_new_children( - self: Arc, - children: Vec>, - ) -> datafusion_common::Result> { - assert!(children.is_empty()); - Ok(self) - } - - fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, "MockExpr") - } - } - - fn ordering() -> LexOrdering { - let expr = Arc::new(MockExpr {}) as Arc; - LexOrdering::new(vec![PhysicalSortExpr::new_default(expr)]).unwrap() - } - - #[test] - fn test_ordering_cache() { - let meta = create_test_meta("test.parquet", 100); - let cache = DefaultFileStatisticsCache::default(); - - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - - // Cache statistics with no ordering - let cached_value = CachedFileMetadata::new( - meta.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, // No ordering yet - ); - - let path = TableScopedPath { - path: meta.location.clone(), - table: None, - }; - - cache.put(&path, cached_value); - - let result = cache.get(&path).unwrap(); - assert!(result.ordering.is_none()); - - // Update to add ordering - let mut cached = cache.get(&path).unwrap(); - if cached.is_valid_for(&meta) && cached.ordering.is_none() { - cached.ordering = Some(ordering()); - } - cache.put(&path, cached); - - let result2 = cache.get(&path).unwrap(); - assert!(result2.ordering.is_some()); - - // Verify list_entries shows has_ordering = true - let entries = cache.list_entries(); - assert_eq!(entries.len(), 1); - assert!(entries.get(&path).unwrap().has_ordering); - } - - #[test] - fn test_cache_invalidation_on_file_modification() { - let cache = DefaultFileStatisticsCache::default(); - let path = TableScopedPath { - path: Path::from("test.parquet"), - table: None, - }; - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - - let meta_v1 = create_test_meta("test.parquet", 100); - - // Cache initial value - let cached_value = CachedFileMetadata::new( - meta_v1.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&path, cached_value); - - // File modified (size changed) - let meta_v2 = create_test_meta("test.parquet", 200); - - let cached = cache.get(&path).unwrap(); - // Should not be valid for new meta - assert!(!cached.is_valid_for(&meta_v2)); - - // Compute new value and update - let new_cached = CachedFileMetadata::new( - meta_v2.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - cache.put(&path, new_cached); - - // Should have new metadata - let result = cache.get(&path).unwrap(); - assert_eq!(result.meta.size, 200); - } - - #[test] - fn test_ordering_cache_invalidation_on_file_modification() { - let cache = DefaultFileStatisticsCache::default(); - let path = TableScopedPath { - path: Path::from("test.parquet"), - table: None, - }; - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - - // Cache with original metadata and ordering - let meta_v1 = ObjectMeta { - location: path.path.clone(), - last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") - .unwrap() - .into(), - size: 100, - e_tag: None, - version: None, - }; - let ordering_v1 = ordering(); - let cached_v1 = CachedFileMetadata::new( - meta_v1.clone(), - Arc::new(Statistics::new_unknown(&schema)), - Some(ordering_v1), - ); - cache.put(&path, cached_v1); - - // Verify cached ordering is valid - let cached = cache.get(&path).unwrap(); - assert!(cached.is_valid_for(&meta_v1)); - assert!(cached.ordering.is_some()); - - // File modified (size changed) - let meta_v2 = ObjectMeta { - location: path.path.clone(), - last_modified: DateTime::parse_from_rfc3339("2022-09-28T10:00:00+02:00") - .unwrap() - .into(), - size: 200, // Changed - e_tag: None, - version: None, - }; - - // Cache entry exists but should be invalid for new metadata - let cached = cache.get(&path).unwrap(); - assert!(!cached.is_valid_for(&meta_v2)); - - // Cache new version with different ordering - let ordering_v2 = ordering(); // New ordering instance - let cached_v2 = CachedFileMetadata::new( - meta_v2.clone(), - Arc::new(Statistics::new_unknown(&schema)), - Some(ordering_v2), - ); - cache.put(&path, cached_v2); - - // Old metadata should be invalid - let cached = cache.get(&path).unwrap(); - assert!(!cached.is_valid_for(&meta_v1)); - - // New metadata should be valid - assert!(cached.is_valid_for(&meta_v2)); - assert!(cached.ordering.is_some()); - } - - #[test] - fn test_list_entries() { - let cache = DefaultFileStatisticsCache::default(); - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); - - let meta1 = create_test_meta("test1.parquet", 100); - - let cached_value = CachedFileMetadata::new( - meta1.clone(), - Arc::new(Statistics::new_unknown(&schema)), - None, - ); - - let path_1 = TableScopedPath { - path: meta1.location.clone(), - table: None, - }; - - cache.put(&path_1, cached_value); - let meta2 = create_test_meta("test2.parquet", 200); - let cached_value = CachedFileMetadata::new( - meta2.clone(), - Arc::new(Statistics::new_unknown(&schema)), - Some(ordering()), - ); - - let path_2 = TableScopedPath { - path: meta2.location.clone(), - table: None, - }; - - cache.put(&path_2, cached_value); - - let entries = cache.list_entries(); - assert_eq!( - entries, - HashMap::from([ - ( - path_1, - FileStatisticsCacheEntry { - object_meta: meta1, - num_rows: Precision::Absent, - num_columns: 1, - table_size_bytes: Precision::Absent, - statistics_size_bytes: 360, - has_ordering: false, - } - ), - ( - path_2, - FileStatisticsCacheEntry { - object_meta: meta2, - num_rows: Precision::Absent, - num_columns: 1, - table_size_bytes: Precision::Absent, - statistics_size_bytes: 360, - has_ordering: true, - } - ), - ]) - ); - } - - #[test] - fn test_cache_entry_added_when_entries_are_within_cache_limit() { - let (meta_1, value_1) = create_cached_file_metadata_with_stats("test1.parquet"); - let (meta_2, value_2) = create_cached_file_metadata_with_stats("test2.parquet"); - let (meta_3, value_3) = create_cached_file_metadata_with_stats("test3.parquet"); - - let mut ctx = DFHeapSizeCtx::default(); - - let limit_for_2_entries = meta_1.location.as_ref().heap_size(&mut ctx) - + value_1.heap_size(&mut ctx) - + meta_2.location.as_ref().heap_size(&mut ctx) - + value_2.heap_size(&mut ctx); - - // create a cache with a limit which fits exactly 2 entries - let cache = DefaultFileStatisticsCache::new(limit_for_2_entries); - let path_1 = TableScopedPath { - path: meta_1.location.clone(), - table: None, - }; - - let path_2 = TableScopedPath { - path: meta_2.location.clone(), - table: None, - }; - - cache.put(&path_1, value_1.clone()); - cache.put(&path_2, value_2.clone()); - - assert_eq!(cache.len(), 2); - assert_eq!(cache.memory_used(), limit_for_2_entries); - - let result_1 = cache.get(&path_1); - let result_2 = cache.get(&path_2); - assert_eq!(result_1.unwrap(), value_1); - assert_eq!(result_2.unwrap(), value_2); - - let path_3 = TableScopedPath { - path: meta_3.location.clone(), - table: None, - }; - - // adding the third entry evicts the first entry - cache.put(&path_3, value_3.clone()); - assert_eq!(cache.len(), 2); - assert_eq!(cache.memory_used(), limit_for_2_entries); - - let result_1 = cache.get(&path_1); - assert!(result_1.is_none()); - - let result_2 = cache.get(&path_2); - let result_3 = cache.get(&path_3); - - assert_eq!(result_2.unwrap(), value_2); - assert_eq!(result_3.unwrap(), value_3); - - // add the third entry again, making sure memory usage remains the same - cache.put(&path_3, value_3.clone()); - assert_eq!(cache.memory_used(), limit_for_2_entries); - cache.put(&path_3, value_3.clone()); - assert_eq!(cache.memory_used(), limit_for_2_entries); - - let mut ctx = DFHeapSizeCtx::default(); - cache.remove(&path_2); - assert_eq!(cache.len(), 1); - assert_eq!( - cache.memory_used(), - meta_3.location.as_ref().heap_size(&mut ctx) + value_3.heap_size(&mut ctx) - ); - - cache.clear(); - assert_eq!(cache.len(), 0); - assert_eq!(cache.memory_used(), 0); - } - - #[test] - fn test_cache_rejects_entry_which_is_too_large() { - let (meta, value) = create_cached_file_metadata_with_stats("test1.parquet"); - let mut ctx = DFHeapSizeCtx::default(); - let limit_less_than_the_entry = value.heap_size(&mut ctx) - 1; - - // create a cache with a size less than the entry - let cache = DefaultFileStatisticsCache::new(limit_less_than_the_entry); - - let path_1 = TableScopedPath { - path: meta.location.clone(), - table: None, - }; - - cache.put(&path_1, value); - - assert_eq!(cache.len(), 0); - assert_eq!(cache.memory_used(), 0); - } - - fn create_cached_file_metadata_with_stats( - file_name: &str, - ) -> (ObjectMeta, CachedFileMetadata) { - let series: Vec = (0..=10).collect(); - let values = Int32Array::from(series); - let offsets = OffsetBuffer::new(ScalarBuffer::from(vec![0, 11])); - let field = Arc::new(Field::new_list_field(DataType::Int32, false)); - let list_array = ListArray::new(field, offsets, Arc::new(values), None); - - let column_statistics = ColumnStatistics { - null_count: Precision::Exact(1), - max_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), - min_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), - sum_value: Precision::Exact(ScalarValue::List(Arc::new(list_array.clone()))), - distinct_count: Precision::Exact(10), - byte_size: Precision::Absent, - }; - - let stats = Statistics { - num_rows: Precision::Exact(100), - total_byte_size: Precision::Exact(100), - column_statistics: vec![column_statistics.clone()], - }; - let mut ctx = DFHeapSizeCtx::default(); - let object_meta = create_test_meta(file_name, stats.heap_size(&mut ctx) as u64); - let value = - CachedFileMetadata::new(object_meta.clone(), Arc::new(stats.clone()), None); - (object_meta, value) - } -} diff --git a/datafusion/execution/src/cache/list_files_cache.rs b/datafusion/execution/src/cache/list_files_cache.rs deleted file mode 100644 index a3cdf7c5e9110..0000000000000 --- a/datafusion/execution/src/cache/list_files_cache.rs +++ /dev/null @@ -1,1236 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use crate::cache::{ - CacheAccessor, - cache_manager::{CachedFileList, ListFilesCache}, - lru_queue::LruQueue, -}; - -use std::fmt::{Debug, Display, Formatter}; -use std::mem::size_of; -use std::{ - collections::HashMap, - sync::{Arc, Mutex}, - time::Duration, -}; - -use datafusion_common::TableReference; -use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; -use datafusion_common::instant::Instant; -use object_store::{ObjectMeta, path::Path}; - -pub trait TimeProvider: Send + Sync + 'static { - fn now(&self) -> Instant; -} - -#[derive(Debug, Default)] -pub struct SystemTimeProvider; - -impl TimeProvider for SystemTimeProvider { - fn now(&self) -> Instant { - Instant::now() - } -} - -/// Default implementation of [`ListFilesCache`] -/// -/// Caches file metadata for file listing operations. -/// -/// # Internal details -/// -/// The `memory_limit` parameter controls the maximum size of the cache, which uses a Least -/// Recently Used eviction algorithm. When adding a new entry, if the total number of entries in -/// the cache exceeds `memory_limit`, the least recently used entries are evicted until the total -/// size is lower than the `memory_limit`. -/// -/// # Cache API -/// -/// Uses `get` and `put` methods for cache operations. TTL validation is handled internally - -/// expired entries return `None` from `get`. -pub struct DefaultListFilesCache { - state: Mutex, - time_provider: Arc, -} - -impl Default for DefaultListFilesCache { - fn default() -> Self { - Self::new(DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, None) - } -} - -impl DefaultListFilesCache { - /// Creates a new instance of [`DefaultListFilesCache`]. - /// - /// # Arguments - /// * `memory_limit` - The maximum size of the cache, in bytes. - /// * `ttl` - The TTL (time-to-live) of entries in the cache. - pub fn new(memory_limit: usize, ttl: Option) -> Self { - Self { - state: Mutex::new(DefaultListFilesCacheState::new(memory_limit, ttl)), - time_provider: Arc::new(SystemTimeProvider), - } - } - - #[cfg(test)] - pub(crate) fn with_time_provider(mut self, provider: Arc) -> Self { - self.time_provider = provider; - self - } -} - -#[derive(Clone, PartialEq, Debug)] -pub struct ListFilesEntry { - pub metas: CachedFileList, - pub size_bytes: usize, - pub expires: Option, -} - -impl ListFilesEntry { - fn try_new( - cached_file_list: CachedFileList, - ttl: Option, - now: Instant, - ) -> Option { - let size_bytes = (cached_file_list.files.capacity() * size_of::()) - + cached_file_list - .files - .iter() - .map(meta_heap_bytes) - .reduce(|acc, b| acc + b)?; - - Some(Self { - metas: cached_file_list, - size_bytes, - expires: ttl.map(|t| now + t), - }) - } -} - -/// Calculates the number of bytes an [`ObjectMeta`] occupies in the heap. -fn meta_heap_bytes(object_meta: &ObjectMeta) -> usize { - let mut size = object_meta.location.as_ref().len(); - - if let Some(e) = &object_meta.e_tag { - size += e.len(); - } - if let Some(v) = &object_meta.version { - size += v.len(); - } - - size -} - -/// The default memory limit for the [`DefaultListFilesCache`] -pub const DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT: usize = 1024 * 1024; // 1MiB - -/// The default cache TTL for the [`DefaultListFilesCache`] -pub const DEFAULT_LIST_FILES_CACHE_TTL: Option = None; // Infinite - -/// Key for [`DefaultListFilesCache`] -/// -/// Each entry is scoped to its use within a specific table so that the cache -/// can differentiate between identical paths in different tables, and -/// table-level cache invalidation. -#[derive(PartialEq, Eq, Hash, Clone, Debug)] -pub struct TableScopedPath { - pub table: Option, - pub path: Path, -} - -/// Handles the inner state of the [`DefaultListFilesCache`] struct. -pub struct DefaultListFilesCacheState { - lru_queue: LruQueue, - memory_limit: usize, - memory_used: usize, - ttl: Option, -} - -impl Default for DefaultListFilesCacheState { - fn default() -> Self { - Self { - lru_queue: LruQueue::new(), - memory_limit: DEFAULT_LIST_FILES_CACHE_MEMORY_LIMIT, - memory_used: 0, - ttl: DEFAULT_LIST_FILES_CACHE_TTL, - } - } -} - -impl DFHeapSize for TableScopedPath { - fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { - self.path.as_ref().heap_size(ctx) + self.table.heap_size(ctx) - } -} - -impl Display for TableScopedPath { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - if let Some(table) = &self.table { - write!(f, "{}, {}", self.path, table) - } else { - write!(f, "{}", self.path) - } - } -} - -impl DefaultListFilesCacheState { - fn new(memory_limit: usize, ttl: Option) -> Self { - Self { - lru_queue: LruQueue::new(), - memory_limit, - memory_used: 0, - ttl, - } - } - - /// Gets an entry from the cache, checking for expiration. - /// - /// Returns the cached file list if it exists and hasn't expired. - /// If the entry has expired, it is removed from the cache. - fn get(&mut self, key: &TableScopedPath, now: Instant) -> Option { - let entry = self.lru_queue.get(key)?; - - // Check expiration - if let Some(exp) = entry.expires - && now > exp - { - self.remove(key); - return None; - } - - Some(entry.metas.clone()) - } - - /// Checks if the respective entry is currently cached. - /// - /// If the entry has expired by `now` it is removed from the cache. - /// - /// The LRU queue is not updated. - fn contains_key(&mut self, k: &TableScopedPath, now: Instant) -> bool { - let Some(entry) = self.lru_queue.peek(k) else { - return false; - }; - - match entry.expires { - Some(exp) if now > exp => { - self.remove(k); - false - } - _ => true, - } - } - - /// Adds a new key-value pair to cache expiring at `now` + the TTL. - /// - /// This means that LRU entries might be evicted if required. - /// If the key is already in the cache, the previous entry is returned. - /// If the size of the entry is greater than the `memory_limit`, the value is not inserted. - fn put( - &mut self, - key: &TableScopedPath, - value: CachedFileList, - now: Instant, - ) -> Option { - let entry = ListFilesEntry::try_new(value, self.ttl, now)?; - let entry_size = entry.size_bytes; - - // no point in trying to add this value to the cache if it cannot fit entirely - if entry_size > self.memory_limit { - return None; - } - - // if the key is already in the cache, the old value is removed - let old_value = self.lru_queue.put(key.clone(), entry); - self.memory_used += entry_size; - - if let Some(entry) = &old_value { - self.memory_used -= entry.size_bytes; - } - - self.evict_entries(); - - old_value.map(|v| v.metas) - } - - /// Evicts entries from the LRU cache until `memory_used` is lower than `memory_limit`. - fn evict_entries(&mut self) { - while self.memory_used > self.memory_limit { - if let Some(removed) = self.lru_queue.pop() { - self.memory_used -= removed.1.size_bytes; - } else { - // cache is empty while memory_used > memory_limit, cannot happen - debug_assert!( - false, - "cache is empty while memory_used > memory_limit, cannot happen" - ); - return; - } - } - } - - /// Removes an entry from the cache and returns it, if it exists. - fn remove(&mut self, k: &TableScopedPath) -> Option { - if let Some(entry) = self.lru_queue.remove(k) { - self.memory_used -= entry.size_bytes; - Some(entry.metas) - } else { - None - } - } - - /// Returns the number of entries currently cached. - fn len(&self) -> usize { - self.lru_queue.len() - } - - /// Removes all entries from the cache. - fn clear(&mut self) { - self.lru_queue.clear(); - self.memory_used = 0; - } -} - -impl CacheAccessor for DefaultListFilesCache { - fn get(&self, key: &TableScopedPath) -> Option { - let mut state = self.state.lock().unwrap(); - let now = self.time_provider.now(); - state.get(key, now) - } - - fn put( - &self, - key: &TableScopedPath, - value: CachedFileList, - ) -> Option { - let mut state = self.state.lock().unwrap(); - let now = self.time_provider.now(); - state.put(key, value, now) - } - - fn remove(&self, k: &TableScopedPath) -> Option { - let mut state = self.state.lock().unwrap(); - state.remove(k) - } - - fn contains_key(&self, k: &TableScopedPath) -> bool { - let mut state = self.state.lock().unwrap(); - let now = self.time_provider.now(); - state.contains_key(k, now) - } - - fn len(&self) -> usize { - let state = self.state.lock().unwrap(); - state.len() - } - - fn clear(&self) { - let mut state = self.state.lock().unwrap(); - state.clear(); - } - - fn name(&self) -> String { - String::from("DefaultListFilesCache") - } -} - -impl ListFilesCache for DefaultListFilesCache { - fn cache_limit(&self) -> usize { - let state = self.state.lock().unwrap(); - state.memory_limit - } - - fn cache_ttl(&self) -> Option { - let state = self.state.lock().unwrap(); - state.ttl - } - - fn update_cache_limit(&self, limit: usize) { - let mut state = self.state.lock().unwrap(); - state.memory_limit = limit; - state.evict_entries(); - } - - fn update_cache_ttl(&self, ttl: Option) { - let mut state = self.state.lock().unwrap(); - state.ttl = ttl; - state.evict_entries(); - } - - fn list_entries(&self) -> HashMap { - let state = self.state.lock().unwrap(); - let mut entries = HashMap::::new(); - for (path, entry) in state.lru_queue.list_entries() { - entries.insert(path.clone(), entry.clone()); - } - entries - } - - fn drop_table_entries( - &self, - table_ref: &Option, - ) -> datafusion_common::Result<()> { - let mut state = self.state.lock().unwrap(); - let mut table_paths = vec![]; - for (path, _) in state.lru_queue.list_entries() { - if path.table == *table_ref { - table_paths.push(path.clone()); - } - } - for path in table_paths { - state.remove(&path); - } - Ok(()) - } -} - -#[cfg(test)] -mod tests { - use super::*; - use chrono::DateTime; - use std::thread; - - struct MockTimeProvider { - base: Instant, - offset: Mutex, - } - - impl MockTimeProvider { - fn new() -> Self { - Self { - base: Instant::now(), - offset: Mutex::new(Duration::ZERO), - } - } - - fn inc(&self, duration: Duration) { - let mut offset = self.offset.lock().unwrap(); - *offset += duration; - } - } - - impl TimeProvider for MockTimeProvider { - fn now(&self) -> Instant { - self.base + *self.offset.lock().unwrap() - } - } - - /// Helper function to create a test ObjectMeta with a specific path and location string size - fn create_test_object_meta(path: &str, location_size: usize) -> ObjectMeta { - // Create a location string of the desired size by padding with zeros - let location_str = if location_size > path.len() { - format!("{}{}", path, "0".repeat(location_size - path.len())) - } else { - path.to_string() - }; - - ObjectMeta { - location: Path::from(location_str), - last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") - .unwrap() - .into(), - size: 1024, - e_tag: None, - version: None, - } - } - - /// Helper function to create a CachedFileList with at least meta_size bytes - fn create_test_list_files_entry( - path: &str, - count: usize, - meta_size: usize, - ) -> (Path, CachedFileList, usize) { - let metas: Vec = (0..count) - .map(|i| create_test_object_meta(&format!("file{i}"), meta_size)) - .collect(); - - // Calculate actual size using the same logic as ListFilesEntry::try_new - let size = (metas.capacity() * size_of::()) - + metas.iter().map(meta_heap_bytes).sum::(); - - (Path::from(path), CachedFileList::new(metas), size) - } - - #[test] - fn test_basic_operations() { - let cache = DefaultListFilesCache::default(); - let table_ref = Some(TableReference::from("table")); - let path = Path::from("test_path"); - let key = TableScopedPath { - table: table_ref.clone(), - path, - }; - - // Initially cache is empty - assert!(!cache.contains_key(&key)); - assert_eq!(cache.len(), 0); - - // Cache miss - get returns None - assert!(cache.get(&key).is_none()); - - // Put a value - let meta = create_test_object_meta("file1", 50); - cache.put(&key, CachedFileList::new(vec![meta])); - - // Entry should be cached - assert!(cache.contains_key(&key)); - assert_eq!(cache.len(), 1); - let result = cache.get(&key).unwrap(); - assert_eq!(result.files.len(), 1); - - // Remove the entry - let removed = cache.remove(&key).unwrap(); - assert_eq!(removed.files.len(), 1); - assert!(!cache.contains_key(&key)); - assert_eq!(cache.len(), 0); - - // Put multiple entries - let (path1, value1, size1) = create_test_list_files_entry("path1", 2, 50); - let (path2, value2, size2) = create_test_list_files_entry("path2", 3, 50); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref, - path: path2, - }; - cache.put(&key1, value1.clone()); - cache.put(&key2, value2.clone()); - assert_eq!(cache.len(), 2); - - // List cache entries - assert_eq!( - cache.list_entries(), - HashMap::from([ - ( - key1.clone(), - ListFilesEntry { - metas: value1, - size_bytes: size1, - expires: None, - } - ), - ( - key2.clone(), - ListFilesEntry { - metas: value2, - size_bytes: size2, - expires: None, - } - ) - ]) - ); - - // Clear all entries - cache.clear(); - assert_eq!(cache.len(), 0); - assert!(!cache.contains_key(&key1)); - assert!(!cache.contains_key(&key2)); - } - - #[test] - fn test_lru_eviction_basic() { - let (path1, value1, size) = create_test_list_files_entry("path1", 1, 100); - let (path2, value2, _) = create_test_list_files_entry("path2", 1, 100); - let (path3, value3, _) = create_test_list_files_entry("path3", 1, 100); - - // Set cache limit to exactly fit all three entries - let cache = DefaultListFilesCache::new(size * 3, None); - - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref.clone(), - path: path2, - }; - let key3 = TableScopedPath { - table: table_ref.clone(), - path: path3, - }; - - // All three entries should fit - cache.put(&key1, value1); - cache.put(&key2, value2); - cache.put(&key3, value3); - assert_eq!(cache.len(), 3); - assert!(cache.contains_key(&key1)); - assert!(cache.contains_key(&key2)); - assert!(cache.contains_key(&key3)); - - // Adding a new entry should evict path1 (LRU) - let (path4, value4, _) = create_test_list_files_entry("path4", 1, 100); - let key4 = TableScopedPath { - table: table_ref, - path: path4, - }; - cache.put(&key4, value4); - - assert_eq!(cache.len(), 3); - assert!(!cache.contains_key(&key1)); // Evicted - assert!(cache.contains_key(&key2)); - assert!(cache.contains_key(&key3)); - assert!(cache.contains_key(&key4)); - } - - #[test] - fn test_lru_ordering_after_access() { - let (path1, value1, size) = create_test_list_files_entry("path1", 1, 100); - let (path2, value2, _) = create_test_list_files_entry("path2", 1, 100); - let (path3, value3, _) = create_test_list_files_entry("path3", 1, 100); - - // Set cache limit to fit exactly three entries - let cache = DefaultListFilesCache::new(size * 3, None); - - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref.clone(), - path: path2, - }; - let key3 = TableScopedPath { - table: table_ref.clone(), - path: path3, - }; - - cache.put(&key1, value1); - cache.put(&key2, value2); - cache.put(&key3, value3); - assert_eq!(cache.len(), 3); - - // Access path1 to move it to front (MRU) - // Order is now: path2 (LRU), path3, path1 (MRU) - let _ = cache.get(&key1); - - // Adding a new entry should evict path2 (the LRU) - let (path4, value4, _) = create_test_list_files_entry("path4", 1, 100); - let key4 = TableScopedPath { - table: table_ref, - path: path4, - }; - cache.put(&key4, value4); - - assert_eq!(cache.len(), 3); - assert!(cache.contains_key(&key1)); // Still present (recently accessed) - assert!(!cache.contains_key(&key2)); // Evicted (was LRU) - assert!(cache.contains_key(&key3)); - assert!(cache.contains_key(&key4)); - } - - #[test] - fn test_reject_too_large() { - let (path1, value1, size) = create_test_list_files_entry("path1", 1, 100); - let (path2, value2, _) = create_test_list_files_entry("path2", 1, 100); - - // Set cache limit to fit both entries - let cache = DefaultListFilesCache::new(size * 2, None); - - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref.clone(), - path: path2, - }; - cache.put(&key1, value1); - cache.put(&key2, value2); - assert_eq!(cache.len(), 2); - - // Try to add an entry that's too large to fit in the cache - // The entry is not stored (too large) - let (path_large, value_large, _) = create_test_list_files_entry("large", 1, 1000); - let key_large = TableScopedPath { - table: table_ref, - path: path_large, - }; - cache.put(&key_large, value_large); - - // Large entry should not be added - assert!(!cache.contains_key(&key_large)); - assert_eq!(cache.len(), 2); - assert!(cache.contains_key(&key1)); - assert!(cache.contains_key(&key2)); - } - - #[test] - fn test_multiple_evictions() { - let (path1, value1, size) = create_test_list_files_entry("path1", 1, 100); - let (path2, value2, _) = create_test_list_files_entry("path2", 1, 100); - let (path3, value3, _) = create_test_list_files_entry("path3", 1, 100); - - // Set cache limit for exactly 3 entries - let cache = DefaultListFilesCache::new(size * 3, None); - - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref.clone(), - path: path2, - }; - let key3 = TableScopedPath { - table: table_ref.clone(), - path: path3, - }; - cache.put(&key1, value1); - cache.put(&key2, value2); - cache.put(&key3, value3); - assert_eq!(cache.len(), 3); - - // Add a large entry that requires evicting 2 entries - let (path_large, value_large, _) = create_test_list_files_entry("large", 1, 200); - let key_large = TableScopedPath { - table: table_ref, - path: path_large, - }; - cache.put(&key_large, value_large); - - // path1 and path2 should be evicted (both LRU), path3 and path_large remain - assert_eq!(cache.len(), 2); - assert!(!cache.contains_key(&key1)); // Evicted - assert!(!cache.contains_key(&key2)); // Evicted - assert!(cache.contains_key(&key3)); - assert!(cache.contains_key(&key_large)); - } - - #[test] - fn test_cache_limit_resize() { - let (path1, value1, size) = create_test_list_files_entry("path1", 1, 100); - let (path2, value2, _) = create_test_list_files_entry("path2", 1, 100); - let (path3, value3, _) = create_test_list_files_entry("path3", 1, 100); - - let cache = DefaultListFilesCache::new(size * 3, None); - - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref.clone(), - path: path2, - }; - let key3 = TableScopedPath { - table: table_ref, - path: path3, - }; - // Add three entries - cache.put(&key1, value1); - cache.put(&key2, value2); - cache.put(&key3, value3); - assert_eq!(cache.len(), 3); - - // Resize cache to only fit one entry - cache.update_cache_limit(size); - - // Should keep only the most recent entry (path3, the MRU) - assert_eq!(cache.len(), 1); - assert!(cache.contains_key(&key3)); - // Earlier entries (LRU) should be evicted - assert!(!cache.contains_key(&key1)); - assert!(!cache.contains_key(&key2)); - } - - #[test] - fn test_entry_update_with_size_change() { - let (path1, value1, size) = create_test_list_files_entry("path1", 1, 100); - let (path2, value2, size2) = create_test_list_files_entry("path2", 1, 100); - let (path3, value3_v1, _) = create_test_list_files_entry("path3", 1, 100); - - let cache = DefaultListFilesCache::new(size * 3, None); - - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref.clone(), - path: path2, - }; - let key3 = TableScopedPath { - table: table_ref, - path: path3, - }; - // Add three entries - cache.put(&key1, value1); - cache.put(&key2, value2.clone()); - cache.put(&key3, value3_v1); - assert_eq!(cache.len(), 3); - - // Update path3 with same size - should not cause eviction - let (_, value3_v2, _) = create_test_list_files_entry("path3", 1, 100); - cache.put(&key3, value3_v2); - - assert_eq!(cache.len(), 3); - assert!(cache.contains_key(&key1)); - assert!(cache.contains_key(&key2)); - assert!(cache.contains_key(&key3)); - - // Update path3 with larger size that requires evicting path1 (LRU) - let (_, value3_v3, size3_v3) = create_test_list_files_entry("path3", 1, 200); - cache.put(&key3, value3_v3.clone()); - - assert_eq!(cache.len(), 2); - assert!(!cache.contains_key(&key1)); // Evicted (was LRU) - assert!(cache.contains_key(&key2)); - assert!(cache.contains_key(&key3)); - - // List cache entries - assert_eq!( - cache.list_entries(), - HashMap::from([ - ( - key2, - ListFilesEntry { - metas: value2, - size_bytes: size2, - expires: None, - } - ), - ( - key3, - ListFilesEntry { - metas: value3_v3, - size_bytes: size3_v3, - expires: None, - } - ) - ]) - ); - } - - #[test] - fn test_cache_with_ttl() { - let ttl = Duration::from_millis(100); - - let mock_time = Arc::new(MockTimeProvider::new()); - let cache = DefaultListFilesCache::new(10000, Some(ttl)) - .with_time_provider(Arc::clone(&mock_time) as Arc); - - let (path1, value1, size1) = create_test_list_files_entry("path1", 2, 50); - let (path2, value2, size2) = create_test_list_files_entry("path2", 2, 50); - - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref, - path: path2, - }; - cache.put(&key1, value1.clone()); - cache.put(&key2, value2.clone()); - - // Entries should be accessible immediately - assert!(cache.get(&key1).is_some()); - assert!(cache.get(&key2).is_some()); - // List cache entries - assert_eq!( - cache.list_entries(), - HashMap::from([ - ( - key1.clone(), - ListFilesEntry { - metas: value1, - size_bytes: size1, - expires: mock_time.now().checked_add(ttl), - } - ), - ( - key2.clone(), - ListFilesEntry { - metas: value2, - size_bytes: size2, - expires: mock_time.now().checked_add(ttl), - } - ) - ]) - ); - // Wait for TTL to expire - mock_time.inc(Duration::from_millis(150)); - - // Entries should now return None when observed through contains_key - assert!(!cache.contains_key(&key1)); - assert_eq!(cache.len(), 1); // key1 was removed by contains_key() - assert!(!cache.contains_key(&key2)); - assert_eq!(cache.len(), 0); // key2 was removed by contains_key() - } - - #[test] - fn test_cache_with_ttl_and_lru() { - let ttl = Duration::from_millis(200); - - let mock_time = Arc::new(MockTimeProvider::new()); - let cache = DefaultListFilesCache::new(1000, Some(ttl)) - .with_time_provider(Arc::clone(&mock_time) as Arc); - - let (path1, value1, _) = create_test_list_files_entry("path1", 1, 400); - let (path2, value2, _) = create_test_list_files_entry("path2", 1, 400); - let (path3, value3, _) = create_test_list_files_entry("path3", 1, 400); - - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref.clone(), - path: path2, - }; - let key3 = TableScopedPath { - table: table_ref, - path: path3, - }; - cache.put(&key1, value1); - mock_time.inc(Duration::from_millis(50)); - cache.put(&key2, value2); - mock_time.inc(Duration::from_millis(50)); - - // path3 should evict path1 due to size limit - cache.put(&key3, value3); - assert!(!cache.contains_key(&key1)); // Evicted by LRU - assert!(cache.contains_key(&key2)); - assert!(cache.contains_key(&key3)); - - mock_time.inc(Duration::from_millis(151)); - - assert!(!cache.contains_key(&key2)); // Expired - assert!(cache.contains_key(&key3)); // Still valid - } - - #[test] - fn test_ttl_expiration_in_get() { - let ttl = Duration::from_millis(100); - let cache = DefaultListFilesCache::new(10000, Some(ttl)); - - let (path, value, _) = create_test_list_files_entry("path", 2, 50); - let table_ref = Some(TableReference::from("table")); - let key = TableScopedPath { - table: table_ref, - path, - }; - - // Cache the entry - cache.put(&key, value.clone()); - - // Entry should be accessible immediately - let result = cache.get(&key); - assert!(result.is_some()); - assert_eq!(result.unwrap().files.len(), 2); - - // Wait for TTL to expire - thread::sleep(Duration::from_millis(150)); - - // Get should return None because entry expired - let result2 = cache.get(&key); - assert!(result2.is_none()); - } - - #[test] - fn test_meta_heap_bytes_calculation() { - // Test with minimal ObjectMeta (no e_tag, no version) - let meta1 = ObjectMeta { - location: Path::from("test"), - last_modified: chrono::Utc::now(), - size: 100, - e_tag: None, - version: None, - }; - assert_eq!(meta_heap_bytes(&meta1), 4); // Just the location string "test" - - // Test with e_tag - let meta2 = ObjectMeta { - location: Path::from("test"), - last_modified: chrono::Utc::now(), - size: 100, - e_tag: Some("etag123".to_string()), - version: None, - }; - assert_eq!(meta_heap_bytes(&meta2), 4 + 7); // location (4) + e_tag (7) - - // Test with version - let meta3 = ObjectMeta { - location: Path::from("test"), - last_modified: chrono::Utc::now(), - size: 100, - e_tag: None, - version: Some("v1.0".to_string()), - }; - assert_eq!(meta_heap_bytes(&meta3), 4 + 4); // location (4) + version (4) - - // Test with both e_tag and version - let meta4 = ObjectMeta { - location: Path::from("test"), - last_modified: chrono::Utc::now(), - size: 100, - e_tag: Some("tag".to_string()), - version: Some("ver".to_string()), - }; - assert_eq!(meta_heap_bytes(&meta4), 4 + 3 + 3); // location (4) + e_tag (3) + version (3) - } - - #[test] - fn test_entry_creation() { - // Test with empty vector - let empty_list = CachedFileList::new(vec![]); - let now = Instant::now(); - let entry = ListFilesEntry::try_new(empty_list, None, now); - assert!(entry.is_none()); - - // Validate entry size - let metas: Vec = (0..5) - .map(|i| create_test_object_meta(&format!("file{i}"), 30)) - .collect(); - let cached_list = CachedFileList::new(metas); - let entry = ListFilesEntry::try_new(cached_list, None, now).unwrap(); - assert_eq!(entry.metas.files.len(), 5); - // Size should be: capacity * sizeof(ObjectMeta) + (5 * 30) for heap bytes - let expected_size = (entry.metas.files.capacity() * size_of::()) - + (entry.metas.files.len() * 30); - assert_eq!(entry.size_bytes, expected_size); - - // Test with TTL - let meta = create_test_object_meta("file", 50); - let ttl = Duration::from_secs(10); - let cached_list = CachedFileList::new(vec![meta]); - let entry = ListFilesEntry::try_new(cached_list, Some(ttl), now).unwrap(); - assert!(entry.expires.unwrap() > now); - } - - #[test] - fn test_memory_tracking() { - let cache = DefaultListFilesCache::new(1000, None); - - // Verify cache starts with 0 memory used - { - let state = cache.state.lock().unwrap(); - assert_eq!(state.memory_used, 0); - } - - // Add entry and verify memory tracking - let (path1, value1, size1) = create_test_list_files_entry("path1", 1, 100); - let table_ref = Some(TableReference::from("table")); - let key1 = TableScopedPath { - table: table_ref.clone(), - path: path1, - }; - cache.put(&key1, value1); - { - let state = cache.state.lock().unwrap(); - assert_eq!(state.memory_used, size1); - } - - // Add another entry - let (path2, value2, size2) = create_test_list_files_entry("path2", 1, 200); - let key2 = TableScopedPath { - table: table_ref.clone(), - path: path2, - }; - cache.put(&key2, value2); - { - let state = cache.state.lock().unwrap(); - assert_eq!(state.memory_used, size1 + size2); - } - - // Remove first entry and verify memory decreases - cache.remove(&key1); - { - let state = cache.state.lock().unwrap(); - assert_eq!(state.memory_used, size2); - } - - // Clear and verify memory is 0 - cache.clear(); - { - let state = cache.state.lock().unwrap(); - assert_eq!(state.memory_used, 0); - } - } - - // Prefix filtering tests using CachedFileList::filter_by_prefix - - /// Helper function to create ObjectMeta with a specific location path - fn create_object_meta_with_path(location: &str) -> ObjectMeta { - ObjectMeta { - location: Path::from(location), - last_modified: DateTime::parse_from_rfc3339("2022-09-27T22:36:00+02:00") - .unwrap() - .into(), - size: 1024, - e_tag: None, - version: None, - } - } - - #[test] - fn test_prefix_filtering() { - let cache = DefaultListFilesCache::new(100000, None); - - // Create files for a partitioned table - let table_base = Path::from("my_table"); - let files = vec![ - create_object_meta_with_path("my_table/a=1/file1.parquet"), - create_object_meta_with_path("my_table/a=1/file2.parquet"), - create_object_meta_with_path("my_table/a=2/file3.parquet"), - create_object_meta_with_path("my_table/a=2/file4.parquet"), - ]; - - // Cache the full table listing - let table_ref = Some(TableReference::from("table")); - let key = TableScopedPath { - table: table_ref, - path: table_base, - }; - cache.put(&key, CachedFileList::new(files)); - - let result = cache.get(&key).unwrap(); - - // Filter for partition a=1 - let prefix_a1 = Some(Path::from("my_table/a=1")); - let filtered = result.files_matching_prefix(&prefix_a1); - assert_eq!(filtered.len(), 2); - assert!( - filtered - .iter() - .all(|m| m.location.as_ref().starts_with("my_table/a=1")) - ); - - // Filter for partition a=2 - let prefix_a2 = Some(Path::from("my_table/a=2")); - let filtered_2 = result.files_matching_prefix(&prefix_a2); - assert_eq!(filtered_2.len(), 2); - assert!( - filtered_2 - .iter() - .all(|m| m.location.as_ref().starts_with("my_table/a=2")) - ); - - // No filter returns all - let all = result.files_matching_prefix(&None); - assert_eq!(all.len(), 4); - } - - #[test] - fn test_prefix_no_matching_files() { - let cache = DefaultListFilesCache::new(100000, None); - - let table_base = Path::from("my_table"); - let files = vec![ - create_object_meta_with_path("my_table/a=1/file1.parquet"), - create_object_meta_with_path("my_table/a=2/file2.parquet"), - ]; - - let table_ref = Some(TableReference::from("table")); - let key = TableScopedPath { - table: table_ref, - path: table_base, - }; - cache.put(&key, CachedFileList::new(files)); - let result = cache.get(&key).unwrap(); - - // Query for partition a=3 which doesn't exist - let prefix_a3 = Some(Path::from("my_table/a=3")); - let filtered = result.files_matching_prefix(&prefix_a3); - assert!(filtered.is_empty()); - } - - #[test] - fn test_nested_partitions() { - let cache = DefaultListFilesCache::new(100000, None); - - let table_base = Path::from("events"); - let files = vec![ - create_object_meta_with_path( - "events/year=2024/month=01/day=01/file1.parquet", - ), - create_object_meta_with_path( - "events/year=2024/month=01/day=02/file2.parquet", - ), - create_object_meta_with_path( - "events/year=2024/month=02/day=01/file3.parquet", - ), - create_object_meta_with_path( - "events/year=2025/month=01/day=01/file4.parquet", - ), - ]; - - let table_ref = Some(TableReference::from("table")); - let key = TableScopedPath { - table: table_ref, - path: table_base, - }; - cache.put(&key, CachedFileList::new(files)); - let result = cache.get(&key).unwrap(); - - // Filter for year=2024/month=01 - let prefix_month = Some(Path::from("events/year=2024/month=01")); - let filtered = result.files_matching_prefix(&prefix_month); - assert_eq!(filtered.len(), 2); - - // Filter for year=2024 - let prefix_year = Some(Path::from("events/year=2024")); - let filtered_year = result.files_matching_prefix(&prefix_year); - assert_eq!(filtered_year.len(), 3); - } - - #[test] - fn test_drop_table_entries() { - let cache = DefaultListFilesCache::default(); - - let (path1, value1, _) = create_test_list_files_entry("path1", 1, 100); - let (path2, value2, _) = create_test_list_files_entry("path2", 1, 100); - let (path3, value3, _) = create_test_list_files_entry("path3", 1, 100); - - let table_ref1 = Some(TableReference::from("table1")); - let key1 = TableScopedPath { - table: table_ref1.clone(), - path: path1, - }; - let key2 = TableScopedPath { - table: table_ref1.clone(), - path: path2, - }; - - let table_ref2 = Some(TableReference::from("table2")); - let key3 = TableScopedPath { - table: table_ref2.clone(), - path: path3, - }; - - cache.put(&key1, value1); - cache.put(&key2, value2); - cache.put(&key3, value3); - - cache.drop_table_entries(&table_ref1).unwrap(); - - assert!(!cache.contains_key(&key1)); - assert!(!cache.contains_key(&key2)); - assert!(cache.contains_key(&key3)); - } -} diff --git a/datafusion/execution/src/cache/lru_queue.rs b/datafusion/execution/src/cache/lru_queue.rs index fb3d158ced425..a19f13865fd3d 100644 --- a/datafusion/execution/src/cache/lru_queue.rs +++ b/datafusion/execution/src/cache/lru_queue.rs @@ -212,6 +212,12 @@ impl LruQueue { pub fn list_entries(&self) -> HashMap<&K, &V> { self.data.iter().map(|(k, (_, v))| (k, v)).collect() } + + /// Returns an iterator over references to the keys currently in the queue. + /// The order is unspecified and does not reflect the LRU order. + pub fn keys(&self) -> impl Iterator { + self.data.keys() + } } #[cfg(test)] diff --git a/datafusion/execution/src/cache/mod.rs b/datafusion/execution/src/cache/mod.rs index 76bd660e6c7d5..f47a3f3ca49f3 100644 --- a/datafusion/execution/src/cache/mod.rs +++ b/datafusion/execution/src/cache/mod.rs @@ -16,41 +16,32 @@ // under the License. pub mod cache_manager; -pub mod file_statistics_cache; pub mod lru_queue; -mod file_metadata_cache; -mod list_files_cache; +pub mod default_cache; -pub use file_metadata_cache::DefaultFilesMetadataCache; -pub use list_files_cache::DefaultListFilesCache; -pub use list_files_cache::ListFilesEntry; -pub use list_files_cache::TableScopedPath; +use datafusion_common::arrow::datatypes::{DataType, Schema}; +use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; +use datafusion_common::instant::Instant; +use datafusion_common::{HashMap, TableReference}; +use object_store::path::Path; +use std::collections::hash_map::DefaultHasher; +use std::fmt::{Debug, Display, Formatter}; +use std::hash::{Hash, Hasher}; +use std::time::Duration; /// Base trait for cache implementations with common operations. /// /// This trait provides the fundamental cache operations (`get`, `put`, `remove`, etc.) -/// that all cache types share. Specific cache traits like [`cache_manager::FileStatisticsCache`], -/// [`cache_manager::ListFilesCache`], and [`cache_manager::FileMetadataCache`] extend this -/// trait with their specialized methods. +/// that all cache types share. /// /// ## Thread Safety /// /// Implementations must handle their own locking via internal mutability, as methods do not /// take mutable references and may be accessed by multiple concurrent queries. /// -/// ## Validation Pattern -/// -/// Validation metadata (e.g., file size, last modified time) should be embedded in the -/// value type `V`. The typical usage pattern is: -/// 1. Call `get(key)` to check for cached value -/// 2. If `Some(cached)`, validate with `cached.is_valid_for(¤t_meta)` -/// 3. If invalid or missing, compute new value and call `put(key, new_value)` -pub trait CacheAccessor: Send + Sync { +pub trait Cache: Send + Sync { /// Get a cached entry if it exists. - /// - /// Returns the cached value without any validation. The caller should - /// validate the returned value if freshness matters. fn get(&self, key: &K) -> Option; /// Store a value in the cache. @@ -77,4 +68,209 @@ pub trait CacheAccessor: Send + Sync { /// Return the cache name. fn name(&self) -> String; + + /// Current memory budget, in bytes. + fn cache_limit(&self) -> usize; + + /// Change the memory budget in bytes. + fn update_cache_limit(&self, limit: usize); + + /// Time-to-live applied to newly inserted entries, or `None` if entries + /// never expire on their own. + fn cache_ttl(&self) -> Option; + + /// Change the TTL applied to subsequent inserts. + fn update_cache_ttl(&self, _ttl: Option); + + /// Invalidate every entry associated with `table_ref`. + fn drop_table_entries( + &self, + table_ref: &TableReference, + ) -> datafusion_common::Result<()>; + + /// Snapshot of all current entries with per-entry metadata (size, hits, + /// expiration) for diagnostics and observability. + fn list_entries(&self) -> HashMap>; +} + +/// Key type for entries stored in a [`Cache`]. +pub trait CacheKey: Clone + Eq + Hash + Send + Sync + Debug { + /// Size of the key in bytes, used for cache memory accounting. + fn size(&self) -> usize; + + /// Table this key is associated with, or `None` if the key is not + /// table-scoped. + fn table_ref(&self) -> Option<&TableReference>; +} + +/// Value type for entries stored in a [`Cache`]. +pub trait CacheValue: Clone + Send + Sync { + /// Size of the value in bytes used for cache memory accounting. + fn size(&self) -> usize; +} + +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct CacheEntryInfo { + pub value: V, + pub size_bytes: usize, + pub hits: usize, + pub expires: Option, +} + +impl Debug for dyn Cache { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "Cache name: {} with length: {}", self.name(), self.len()) + } +} + +impl CacheKey for Path { + fn size(&self) -> usize { + self.as_ref().heap_size(&mut DFHeapSizeCtx::default()) + } + + fn table_ref(&self) -> Option<&TableReference> { + None + } +} + +impl CacheKey for TableScopedPath { + fn size(&self) -> usize { + DFHeapSize::heap_size(self, &mut DFHeapSizeCtx::default()) + } + + fn table_ref(&self) -> Option<&TableReference> { + self.table.as_ref() + } +} + +/// Each entry is scoped to its use within a specific table so that the cache +/// can differentiate between identical paths in different tables, and +/// table-level cache invalidation. +#[derive(PartialEq, Eq, Hash, Clone, Debug)] +pub struct TableScopedPath { + pub table: Option, + pub path: Path, +} + +impl DFHeapSize for TableScopedPath { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.path.as_ref().heap_size(ctx) + self.table.heap_size(ctx) + } +} + +impl Display for TableScopedPath { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + if let Some(table) = &self.table { + write!(f, "{}, {}", self.path, table) + } else { + write!(f, "{}", self.path) + } + } +} + +/// A fingerprint of the `file_schema` used to compute a file's statistics. +/// +/// Captures exactly the attributes that determine the layout and meaning of +/// `Statistics::column_statistics`: each column's name, data type and +/// nullability, in order. It deliberately excludes field/schema metadata, which +/// cannot affect statistics — including it would needlessly fragment the cache. +#[derive(Clone, Debug)] +pub struct SchemaFingerprint { + columns: Vec<(String, DataType, bool)>, + /// Precomputed hash of `columns`, so hashing a key on every cache lookup is + /// O(1) rather than O(schema width). Computed once in `from_schema` with a + /// fixed-seed hasher so it is stable across keys; `PartialEq` still compares + /// `columns` exactly, so a hash collision can never make two different + /// schemas share a cache entry. + hash: u64, +} + +impl SchemaFingerprint { + /// Builds a fingerprint from the `file_schema` used to compute statistics + /// (the schema of the columns physically read, not the full table schema — + /// partition columns and their statistics are handled separately). + pub fn from_schema(file_schema: &Schema) -> Self { + let columns: Vec<(String, DataType, bool)> = file_schema + .fields() + .iter() + .map(|f| (f.name().clone(), f.data_type().clone(), f.is_nullable())) + .collect(); + let mut hasher = DefaultHasher::new(); + columns.hash(&mut hasher); + Self { + columns, + hash: hasher.finish(), + } + } +} + +impl PartialEq for SchemaFingerprint { + fn eq(&self, other: &Self) -> bool { + // Cheap hash gate first, then an exact comparison so collisions are safe. + self.hash == other.hash && self.columns == other.columns + } +} + +impl Eq for SchemaFingerprint {} + +impl Hash for SchemaFingerprint { + fn hash(&self, state: &mut H) { + state.write_u64(self.hash); + } +} + +impl DFHeapSize for SchemaFingerprint { + fn heap_size(&self, ctx: &mut DFHeapSizeCtx) -> usize { + self.columns.heap_size(ctx) + } +} + +#[cfg(test)] +mod schema_fingerprint_tests { + use super::*; + use datafusion_common::arrow::datatypes::Field; + + fn fp(fields: Vec) -> SchemaFingerprint { + SchemaFingerprint::from_schema(&Schema::new(fields)) + } + + /// `from_schema` must capture nullability and field order — the two + /// attributes most easily dropped by a wrong implementation. + #[test] + fn fingerprint_captures_nullability_and_order() { + assert_ne!( + fp(vec![Field::new("id", DataType::Int64, false)]), + fp(vec![Field::new("id", DataType::Int64, true)]), + "nullability must affect the fingerprint", + ); + + let ab = fp(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Utf8, true), + ]); + let ba = fp(vec![ + Field::new("b", DataType::Utf8, true), + Field::new("a", DataType::Int64, false), + ]); + assert_ne!(ab, ba, "field order must affect the fingerprint"); + } + + /// Metadata must NOT affect the fingerprint: it cannot change column + /// statistics, so including it would needlessly fragment the cache. + #[test] + fn fingerprint_ignores_metadata() { + let plain = fp(vec![Field::new("id", DataType::Int64, false)]); + + let field_md = SchemaFingerprint::from_schema(&Schema::new(vec![ + Field::new("id", DataType::Int64, false) + .with_metadata([("note".to_string(), "x".to_string())].into()), + ])); + assert_eq!(plain, field_md, "field metadata must be ignored"); + + let schema_md = SchemaFingerprint::from_schema( + &Schema::new(vec![Field::new("id", DataType::Int64, false)]) + .with_metadata([("k".to_string(), "v".to_string())].into()), + ); + assert_eq!(plain, schema_md, "schema metadata must be ignored"); + } } diff --git a/datafusion/execution/src/config.rs b/datafusion/execution/src/config.rs index b2917a4583628..efaedebdadb33 100644 --- a/datafusion/execution/src/config.rs +++ b/datafusion/execution/src/config.rs @@ -19,7 +19,7 @@ use std::{collections::HashMap, sync::Arc}; use datafusion_common::{ Result, ScalarValue, - config::{ConfigExtension, ConfigOptions, SpillCompression}, + config::{ConfigExtension, ConfigNonZeroUsize, ConfigOptions, SpillCompression}, extensions::Extensions, }; @@ -51,7 +51,7 @@ use datafusion_common::{ /// .set_bool("datafusion.execution.parquet.pushdown_filters", true); /// /// assert_eq!(config.batch_size(), 1234); -/// assert_eq!(config.options().execution.batch_size, 1234); +/// assert_eq!(config.options().execution.batch_size.get(), 1234); /// assert_eq!(config.options().execution.parquet.pushdown_filters, true); /// ``` /// @@ -60,15 +60,16 @@ use datafusion_common::{ /// /// ``` /// # use datafusion_execution::config::SessionConfig; -/// # use datafusion_common::ScalarValue; +/// # use datafusion_common::config::ConfigNonZeroUsize; /// # /// let mut config = SessionConfig::new(); -/// config.options_mut().execution.batch_size = 1234; +/// config.options_mut().execution.batch_size = ConfigNonZeroUsize::try_new(1234)?; /// config.options_mut().execution.parquet.pushdown_filters = true; /// # /// # assert_eq!(config.batch_size(), 1234); -/// # assert_eq!(config.options().execution.batch_size, 1234); +/// # assert_eq!(config.options().execution.batch_size.get(), 1234); /// # assert_eq!(config.options().execution.parquet.pushdown_filters, true); +/// # datafusion_common::Result::<()>::Ok(()) /// ``` /// /// ## Built-in options @@ -137,7 +138,7 @@ impl SessionConfig { /// use datafusion_execution::config::SessionConfig; /// /// let config = SessionConfig::new(); - /// assert!(config.options().execution.batch_size > 0); + /// assert!(config.options().execution.batch_size.get() > 0); /// ``` pub fn options(&self) -> &Arc { &self.options @@ -148,11 +149,13 @@ impl SessionConfig { /// Can be used to set configuration options. /// /// ``` + /// use datafusion_common::config::ConfigNonZeroUsize; /// use datafusion_execution::config::SessionConfig; /// /// let mut config = SessionConfig::new(); - /// config.options_mut().execution.batch_size = 1024; - /// assert_eq!(config.options().execution.batch_size, 1024); + /// config.options_mut().execution.batch_size = ConfigNonZeroUsize::try_new(1024)?; + /// assert_eq!(config.options().execution.batch_size.get(), 1024); + /// # datafusion_common::Result::<()>::Ok(()) /// ``` pub fn options_mut(&mut self) -> &mut ConfigOptions { Arc::make_mut(&mut self.options) @@ -186,9 +189,8 @@ impl SessionConfig { /// Customize batch size pub fn with_batch_size(mut self, n: usize) -> Self { - // batch size must be greater than zero - assert!(n > 0); - self.options_mut().execution.batch_size = n; + self.options_mut().execution.batch_size = + ConfigNonZeroUsize::try_new(n).expect("batch size must be greater than zero"); self } @@ -391,7 +393,7 @@ impl SessionConfig { /// Get the currently configured batch size pub fn batch_size(&self) -> usize { - self.options.execution.batch_size + self.options.execution.batch_size.get() } /// Enables or disables the coalescence of small batches into larger batches @@ -512,7 +514,7 @@ impl SessionConfig { /// Extensions are opaque and the types are unknown to DataFusion itself, which makes them extremely flexible. [^1] /// /// Extensions are stored within an [`Arc`] so they do NOT require [`Clone`]. The are immutable. If you need to - /// modify their state over their lifetime -- e.g. for caches -- you need to establish some for of interior mutability. + /// modify their state over their lifetime -- e.g. for caches -- you need to establish some form of interior mutability. /// /// Extensions are indexed by their type `T`. If multiple values of the same type are provided, only the last one /// will be kept. diff --git a/datafusion/execution/src/disk_manager.rs b/datafusion/execution/src/disk_manager.rs index 1a14bd239a61a..313379f01291f 100644 --- a/datafusion/execution/src/disk_manager.rs +++ b/datafusion/execution/src/disk_manager.rs @@ -17,36 +17,49 @@ //! [`DiskManager`]: Manages files generated during query execution -use datafusion_common::{ - DataFusionError, Result, config_err, resources_datafusion_err, resources_err, -}; +use crate::spill_file::{SpillFile, SpillWriter, TempFileFactory}; +use bytes::Bytes; +use datafusion_common::human_readable_size; +use datafusion_common::{DataFusionError, Result, config_err, resources_datafusion_err}; +#[cfg(not(target_arch = "wasm32"))] +use futures::StreamExt; use log::debug; use parking_lot::Mutex; use rand::{Rng, rng}; +use std::fmt::Debug; use std::path::{Path, PathBuf}; use std::sync::Arc; use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; use tempfile::{Builder, NamedTempFile, TempDir}; - -use datafusion_common::human_readable_size; - pub const DEFAULT_MAX_TEMP_DIRECTORY_SIZE: u64 = 100 * 1024 * 1024 * 1024; // 100GB +pub const DEFAULT_MAX_SPILL_MERGE_FAN_IN: usize = 0; /// Builder pattern for the [DiskManager] structure -#[derive(Clone, Debug)] +#[derive(Clone)] pub struct DiskManagerBuilder { /// The storage mode of the disk manager mode: DiskManagerMode, /// The maximum amount of data (in bytes) stored inside the temporary directories. /// Default to 100GB max_temp_directory_size: u64, + /// Maximum number of spill files opened by one external merge pass. + /// A value of 0 means unlimited. + max_spill_merge_fan_in: usize, +} +impl Debug for DiskManagerBuilder { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DiskManagerBuilder") + .field("mode", &self.mode) + .field("max_temp_directory_size", &self.max_temp_directory_size) + .finish() + } } - impl Default for DiskManagerBuilder { fn default() -> Self { Self { mode: DiskManagerMode::OsTmpDirectory, max_temp_directory_size: DEFAULT_MAX_TEMP_DIRECTORY_SIZE, + max_spill_merge_fan_in: DEFAULT_MAX_SPILL_MERGE_FAN_IN, } } } @@ -61,6 +74,26 @@ impl DiskManagerBuilder { self } + /// Configure a custom factory for creating temporary spill files. + /// + /// This sets the disk manager mode to [`DiskManagerMode::Custom`], so + /// operators that spill during query execution create files through the + /// provided [`TempFileFactory`] instead of using local temporary files. + pub fn set_temp_file_factory(&mut self, temp_file_factory: Arc) { + self.mode = DiskManagerMode::Custom(temp_file_factory); + } + + /// Configure a custom factory for creating temporary spill files. + /// + /// See details on [`Self::set_temp_file_factory`]. + pub fn with_temp_file_factory( + mut self, + temp_file_factory: Arc, + ) -> Self { + self.set_temp_file_factory(temp_file_factory); + self + } + pub fn set_max_temp_directory_size(&mut self, value: u64) { self.max_temp_directory_size = value; } @@ -70,14 +103,25 @@ impl DiskManagerBuilder { self } + pub fn set_max_spill_merge_fan_in(&mut self, value: usize) { + self.max_spill_merge_fan_in = value; + } + + pub fn with_max_spill_merge_fan_in(mut self, value: usize) -> Self { + self.set_max_spill_merge_fan_in(value); + self + } + /// Create a DiskManager given the builder pub fn build(self) -> Result { match self.mode { DiskManagerMode::OsTmpDirectory => Ok(DiskManager { local_dirs: Mutex::new(Some(vec![])), - max_temp_directory_size: self.max_temp_directory_size, + max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size), + max_spill_merge_fan_in: AtomicUsize::new(self.max_spill_merge_fan_in), used_disk_space: Arc::new(AtomicU64::new(0)), active_files_count: Arc::new(AtomicUsize::new(0)), + factory: None, }), DiskManagerMode::Directories(conf_dirs) => { let local_dirs = create_local_dirs(&conf_dirs)?; @@ -86,22 +130,34 @@ impl DiskManagerBuilder { ); Ok(DiskManager { local_dirs: Mutex::new(Some(local_dirs)), - max_temp_directory_size: self.max_temp_directory_size, + max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size), + max_spill_merge_fan_in: AtomicUsize::new(self.max_spill_merge_fan_in), used_disk_space: Arc::new(AtomicU64::new(0)), active_files_count: Arc::new(AtomicUsize::new(0)), + factory: None, }) } DiskManagerMode::Disabled => Ok(DiskManager { local_dirs: Mutex::new(None), - max_temp_directory_size: self.max_temp_directory_size, + max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size), + max_spill_merge_fan_in: AtomicUsize::new(self.max_spill_merge_fan_in), used_disk_space: Arc::new(AtomicU64::new(0)), active_files_count: Arc::new(AtomicUsize::new(0)), + factory: None, + }), + DiskManagerMode::Custom(factory) => Ok(DiskManager { + local_dirs: Mutex::new(None), + max_temp_directory_size: AtomicU64::new(self.max_temp_directory_size), + max_spill_merge_fan_in: AtomicUsize::new(self.max_spill_merge_fan_in), + used_disk_space: Arc::new(AtomicU64::new(0)), + active_files_count: Arc::new(AtomicUsize::new(0)), + factory: Some(factory), }), } } } -#[derive(Clone, Debug, Default)] +#[derive(Clone, Default)] pub enum DiskManagerMode { /// Create a new [DiskManager] that creates temporary files within /// a temporary directory chosen by the OS @@ -113,53 +169,26 @@ pub enum DiskManagerMode { /// at random for each temporary file created. Directories(Vec), - /// Disable disk manager, attempts to create temporary files will error - Disabled, -} - -/// Configuration for temporary disk access -#[deprecated(since = "48.0.0", note = "Use DiskManagerBuilder instead")] -#[derive(Debug, Clone, Default)] -#[allow(clippy::allow_attributes)] -#[allow(deprecated)] -pub enum DiskManagerConfig { - /// Use the provided [DiskManager] instance - Existing(Arc), - - /// Create a new [DiskManager] that creates temporary files within - /// a temporary directory chosen by the OS - #[default] - NewOs, - - /// Create a new [DiskManager] that creates temporary files within - /// the specified directories - NewSpecified(Vec), + /// Create a new [DiskManager] with a cutstom backend + Custom(Arc), /// Disable disk manager, attempts to create temporary files will error Disabled, } -#[expect(deprecated)] -impl DiskManagerConfig { - /// Create temporary files in a temporary directory chosen by the OS - pub fn new() -> Self { - Self::default() - } - - /// Create temporary files using the provided disk manager - pub fn new_existing(existing: Arc) -> Self { - Self::Existing(existing) - } - - /// Create temporary files in the specified directories - pub fn new_specified(paths: Vec) -> Self { - Self::NewSpecified(paths) +impl Debug for DiskManagerMode { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::OsTmpDirectory => write!(f, "OsTmpDirectory"), + Self::Directories(dirs) => f.debug_tuple("Directories").field(dirs).finish(), + Self::Disabled => write!(f, "Disabled"), + Self::Custom(_) => write!(f, "Custom(Arc)"), + } } } /// Manages files generated during query execution, e.g. spill files generated /// while processing dataset larger than available memory. -#[derive(Debug)] pub struct DiskManager { /// TempDirs to put temporary files in. /// @@ -167,15 +196,31 @@ pub struct DiskManager { /// If `None` an error will be returned (configured not to spill) local_dirs: Mutex>>>, /// The maximum amount of data (in bytes) stored inside the temporary directories. - /// Default to 100GB - max_temp_directory_size: u64, + /// Default to 100GB. Stored as `AtomicU64` so it can be adjusted at runtime + /// without requiring exclusive (`&mut`) access to the `DiskManager`. + max_temp_directory_size: AtomicU64, + /// Maximum number of spill files opened by one external merge pass. + /// A value of 0 preserves the memory-driven, unbounded behavior. + max_spill_merge_fan_in: AtomicUsize, /// Used disk space in the temporary directories. Now only spilled data for /// external executors are counted. used_disk_space: Arc, /// Number of active temporary files created by this disk manager active_files_count: Arc, + /// Factory + factory: Option>, +} +impl Debug for DiskManager { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DiskManager") + .field("local_dirs", &self.local_dirs) + .field("max_temp_directory_size", &self.max_temp_directory_size) + .field("used_disk_space", &self.used_disk_space) + .field("active_files_count", &self.active_files_count) + .field("factory", &self.factory.is_some()) + .finish() + } } - /// Information about the current disk usage for spilling #[derive(Debug, Clone, Copy)] pub struct SpillingProgress { @@ -191,69 +236,46 @@ impl DiskManager { DiskManagerBuilder::default() } - /// Create a DiskManager given the configuration - #[expect(deprecated)] - #[deprecated(since = "48.0.0", note = "Use DiskManager::builder() instead")] - pub fn try_new(config: DiskManagerConfig) -> Result> { - match config { - DiskManagerConfig::Existing(manager) => Ok(manager), - DiskManagerConfig::NewOs => Ok(Arc::new(Self { - local_dirs: Mutex::new(Some(vec![])), - max_temp_directory_size: DEFAULT_MAX_TEMP_DIRECTORY_SIZE, - used_disk_space: Arc::new(AtomicU64::new(0)), - active_files_count: Arc::new(AtomicUsize::new(0)), - })), - DiskManagerConfig::NewSpecified(conf_dirs) => { - let local_dirs = create_local_dirs(&conf_dirs)?; - debug!( - "Created local dirs {local_dirs:?} as DataFusion working directory" - ); - Ok(Arc::new(Self { - local_dirs: Mutex::new(Some(local_dirs)), - max_temp_directory_size: DEFAULT_MAX_TEMP_DIRECTORY_SIZE, - used_disk_space: Arc::new(AtomicU64::new(0)), - active_files_count: Arc::new(AtomicUsize::new(0)), - })) - } - DiskManagerConfig::Disabled => Ok(Arc::new(Self { - local_dirs: Mutex::new(None), - max_temp_directory_size: DEFAULT_MAX_TEMP_DIRECTORY_SIZE, - used_disk_space: Arc::new(AtomicU64::new(0)), - active_files_count: Arc::new(AtomicUsize::new(0)), - })), - } - } - + /// Atomically set the max temp directory size at runtime. + /// + /// Takes `&self`, so it works through `Arc` without requiring + /// exclusive access. Takes effect immediately for subsequent spill writes. + /// + /// Use this when you need to adjust the limit dynamically while queries + /// are running (e.g., adapting to available disk space). pub fn set_max_temp_directory_size( - &mut self, + &self, max_temp_directory_size: u64, ) -> Result<()> { // If the disk manager is disabled and `max_temp_directory_size` is not 0, // this operation is not meaningful, fail early. - if self.local_dirs.lock().is_none() && max_temp_directory_size != 0 { + if self.local_dirs.lock().is_none() + && max_temp_directory_size != 0 + && self.factory.is_none() + { return config_err!( "Cannot set max temp directory size for a disk manager that spilling is disabled" ); } - self.max_temp_directory_size = max_temp_directory_size; + self.max_temp_directory_size + .store(max_temp_directory_size, Ordering::Relaxed); Ok(()) } + #[deprecated( + since = "54.0.0", + note = "Use `set_max_temp_directory_size` directly, it now takes &self" + )] pub fn set_arc_max_temp_directory_size( - this: &mut Arc, + this: &Arc, max_temp_directory_size: u64, ) -> Result<()> { - if let Some(inner) = Arc::get_mut(this) { - inner.set_max_temp_directory_size(max_temp_directory_size)?; - Ok(()) - } else { - config_err!("DiskManager should be a single instance") - } + this.set_max_temp_directory_size(max_temp_directory_size) } pub fn with_max_temp_directory_size( - mut self, + self, max_temp_directory_size: u64, ) -> Result { self.set_max_temp_directory_size(max_temp_directory_size)?; @@ -266,7 +288,23 @@ impl DiskManager { /// Returns the maximum temporary directory size in bytes pub fn max_temp_directory_size(&self) -> u64 { - self.max_temp_directory_size + self.max_temp_directory_size.load(Ordering::Relaxed) + } + + /// Atomically set the maximum spill merge fan-in. + /// + /// A value of 0 disables the cap. Values of 1 are accepted but external + /// merge code will still merge at least two spill streams to make progress. + pub fn set_max_spill_merge_fan_in(&self, max_spill_merge_fan_in: usize) { + self.max_spill_merge_fan_in + .store(max_spill_merge_fan_in, Ordering::Relaxed); + } + + /// Returns the maximum number of spill files opened by one merge pass. + /// + /// A value of 0 means unlimited. + pub fn max_spill_merge_fan_in(&self) -> usize { + self.max_spill_merge_fan_in.load(Ordering::Relaxed) } /// Returns the current spilling progress @@ -294,7 +332,7 @@ impl DiskManager { /// files. If this returns false, any call to `create_tmp_file` /// will error. pub fn tmp_files_enabled(&self) -> bool { - self.local_dirs.lock().is_some() + self.factory.is_some() || self.local_dirs.lock().is_some() } /// Return a temporary file from a randomized choice in the configured locations @@ -304,7 +342,11 @@ impl DiskManager { pub fn create_tmp_file( self: &Arc, request_description: &str, - ) -> Result { + ) -> Result> { + // Delegate to custom backend if configured + if let Some(factory) = &self.factory { + return factory.create_temp_file(request_description); + } let mut guard = self.local_dirs.lock(); let local_dirs = guard.as_mut().ok_or_else(|| { resources_datafusion_err!( @@ -327,7 +369,7 @@ impl DiskManager { let dir_index = rng().random_range(0..local_dirs.len()); self.active_files_count.fetch_add(1, Ordering::Relaxed); - Ok(RefCountedTempFile { + Ok(Arc::new(RefCountedTempFile { parent_temp_dir: Arc::clone(&local_dirs[dir_index]), tempfile: Arc::new( Builder::new() @@ -336,19 +378,13 @@ impl DiskManager { ), current_file_disk_usage: Arc::new(AtomicU64::new(0)), disk_manager: Arc::clone(self), - }) + })) } } /// A wrapper around a [`NamedTempFile`] that also contains /// a reference to its parent temporary directory. /// -/// # Note -/// After any modification to the underlying file (e.g., writing data to it), the caller -/// must invoke [`Self::update_disk_usage`] to update the global disk usage counter. -/// This ensures the disk manager can properly enforce usage limits configured by -/// [`DiskManager::with_max_temp_directory_size`]. -/// /// This type is Clone-able, allowing multiple references to the same underlying file. /// The file is deleted only when the last reference is dropped. /// @@ -364,8 +400,7 @@ pub struct RefCountedTempFile { parent_temp_dir: Arc, /// The underlying temporary file, wrapped in Arc to allow cloning tempfile: Arc, - /// Tracks the current disk usage of this temporary file. See - /// [`Self::update_disk_usage`] for more details. + /// Tracks the current disk usage of this temporary file. /// /// This is wrapped in `Arc` so that all clones share the same /// disk usage tracking, preventing incorrect accounting when clones are dropped. @@ -394,46 +429,7 @@ impl RefCountedTempFile { self.tempfile.as_ref() } - /// Updates the global disk usage counter after modifications to the underlying file. - /// - /// # Errors - /// - Returns an error if the global disk usage exceeds the configured limit. - pub fn update_disk_usage(&mut self) -> Result<()> { - // Get new file size from OS - let metadata = self.tempfile.as_file().metadata()?; - let new_disk_usage = metadata.len(); - - // Get the old disk usage - let old_disk_usage = self.current_file_disk_usage.load(Ordering::Relaxed); - - // Update the global disk usage by: - // 1. Subtracting the old file size from the global counter - self.disk_manager - .used_disk_space - .fetch_sub(old_disk_usage, Ordering::Relaxed); - // 2. Adding the new file size to the global counter - self.disk_manager - .used_disk_space - .fetch_add(new_disk_usage, Ordering::Relaxed); - - // 3. Check if the updated global disk usage exceeds the configured limit - let global_disk_usage = self.disk_manager.used_disk_space.load(Ordering::Relaxed); - if global_disk_usage > self.disk_manager.max_temp_directory_size { - return resources_err!( - "The used disk space during the spilling process has exceeded the allowable limit of {}. \ - Please try increasing the config: `datafusion.runtime.max_temp_directory_size`.", - human_readable_size(self.disk_manager.max_temp_directory_size as usize) - ); - } - - // 4. Update the local file size tracking - self.current_file_disk_usage - .store(new_disk_usage, Ordering::Relaxed); - - Ok(()) - } - - pub fn current_disk_usage(&self) -> u64 { + fn current_disk_usage(&self) -> u64 { self.current_file_disk_usage.load(Ordering::Relaxed) } } @@ -473,6 +469,122 @@ fn create_local_dirs(local_dirs: &[PathBuf]) -> Result>> { .collect() } +pub struct FileSpillWriter { + file: std::fs::File, + disk_manager: Arc, + current_file_disk_usage: Arc, +} + +impl std::io::Write for FileSpillWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let len = buf.len() as u64; + if len == 0 { + return Ok(0); + } + + let new_global = self + .disk_manager + .used_disk_space + .fetch_add(len, Ordering::Relaxed) + + len; + + let limit = self.disk_manager.max_temp_directory_size(); + + if new_global > limit { + self.disk_manager + .used_disk_space + .fetch_sub(len, Ordering::Relaxed); + + return Err(std::io::Error::other(format!( + "The used disk space during the spilling process has exceeded the allowable limit of {}. \ + Please try increasing the config: `datafusion.runtime.max_temp_directory_size`.", + human_readable_size(limit as usize) + ))); + } + + self.file.write_all(buf).map_err(DataFusionError::IoError)?; + + self.current_file_disk_usage + .fetch_add(len, Ordering::Relaxed); + + Ok(buf.len()) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.file.flush() + } +} + +impl SpillWriter for FileSpillWriter { + fn finish(&mut self) -> Result<()> { + // flush() is handled by Arrow, nothing left to do here + Ok(()) + } +} + +impl SpillFile for RefCountedTempFile { + fn path(&self) -> Option<&Path> { + Some(self.tempfile.path()) + } + + fn size(&self) -> Option { + Some(self.current_disk_usage()) + } + #[cfg(not(target_arch = "wasm32"))] + fn read_stream( + &self, + ) -> Result> + Send>>> + { + let path = self.path().to_owned(); + + let stream = + futures::stream::once(async move { + tokio::fs::File::open(&path) + .await + .map_err(DataFusionError::IoError) + }) + .flat_map( + |open_result| -> std::pin::Pin< + Box> + Send>, + > { + match open_result { + Ok(file) => Box::pin( + // Use a 128KB read buffer. The default 8KB causes excessive async + // poll overhead when reading multi-MB spill files back into memory. + tokio_util::io::ReaderStream::with_capacity(file, 128 * 1024) + .map(|r| r.map_err(DataFusionError::IoError)), + ), + Err(e) => Box::pin(futures::stream::once(async move { Err(e) })), + } + }, + ); + + Ok(Box::pin(stream)) + } + + #[cfg(target_arch = "wasm32")] + fn read_stream( + &self, + ) -> Result> + Send>>> + { + datafusion_common::exec_err!( + "Default OS file spilling is not supported on WASM. Configure DiskManager with a Custom TempFileFactory." + ) + } + + fn open_writer(&self) -> Result> { + let file = self + .tempfile + .as_file() + .try_clone() + .map_err(DataFusionError::IoError)?; + Ok(Box::new(FileSpillWriter { + file, + disk_manager: Arc::clone(&self.disk_manager), + current_file_disk_usage: Arc::clone(&self.current_file_disk_usage), + })) + } +} #[cfg(test)] mod tests { use super::*; @@ -492,7 +604,10 @@ mod tests { // the returned tempfile file should be in the temp directory let local_dirs = local_dir_snapshot(&dm); - assert_path_in_dirs(actual.path(), local_dirs.iter().map(|p| p.as_path())); + assert_path_in_dirs( + actual.path().unwrap(), + local_dirs.iter().map(|p| p.as_path()), + ); Ok(()) } @@ -524,7 +639,7 @@ mod tests { let actual = dm.create_tmp_file("Testing")?; // the file should be in one of the specified local directories - assert_path_in_dirs(actual.path(), local_dirs.into_iter()); + assert_path_in_dirs(actual.path().unwrap(), local_dirs.into_iter()); Ok(()) } @@ -538,13 +653,17 @@ mod tests { .unwrap(), ); assert!(!manager.tmp_files_enabled()); - assert_eq!( - manager - .create_tmp_file("Testing") - .unwrap_err() - .strip_backtrace(), - "Resources exhausted: Memory Exhausted while Testing (DiskManager is disabled)", - ) + match manager.create_tmp_file("Testing") { + Err(e) => { + assert_eq!( + e.strip_backtrace(), + "Resources exhausted: Memory Exhausted while Testing (DiskManager is disabled)" + ); + } + Ok(_) => { + panic!("Expected DiskManager to fail creating a file when disabled!") + } + } } #[test] @@ -577,7 +696,7 @@ mod tests { // Test for the case using OS arranged temporary directory let dm = Arc::new(DiskManagerBuilder::default().build()?); let temp_file = dm.create_tmp_file("Testing")?; - let temp_file_path = temp_file.path().to_owned(); + let temp_file_path = temp_file.path().unwrap().to_owned(); assert!(temp_file_path.exists()); drop(dm); @@ -599,7 +718,7 @@ mod tests { .build()?, ); let temp_file = dm.create_tmp_file("Testing")?; - let temp_file_path = temp_file.path().to_owned(); + let temp_file_path = temp_file.path().unwrap().to_owned(); assert!(temp_file_path.exists()); drop(dm); @@ -613,30 +732,26 @@ mod tests { #[test] fn test_disk_usage_basic() -> Result<()> { - use std::io::Write; - let dm = Arc::new(DiskManagerBuilder::default().build()?); - let mut temp_file = dm.create_tmp_file("Testing")?; - + let temp_file = dm.create_tmp_file("Testing")?; + let mut writer = temp_file.open_writer()?; // Initially, disk usage should be 0 assert_eq!(dm.used_disk_space(), 0); - assert_eq!(temp_file.current_disk_usage(), 0); + assert_eq!(temp_file.size().unwrap(), 0); // Write some data to the file - temp_file.inner().as_file().write_all(b"hello world")?; - temp_file.update_disk_usage()?; + writer.write_all(b"hello world")?; // Disk usage should now reflect the written data - let expected_usage = temp_file.current_disk_usage(); + let expected_usage = temp_file.size().unwrap(); assert!(expected_usage > 0); assert_eq!(dm.used_disk_space(), expected_usage); // Write more data - temp_file.inner().as_file().write_all(b" more data")?; - temp_file.update_disk_usage()?; + writer.write_all(b"more_data")?; // Disk usage should increase - let new_usage = temp_file.current_disk_usage(); + let new_usage = temp_file.size().unwrap(); assert!(new_usage > expected_usage); assert_eq!(dm.used_disk_space(), new_usage); @@ -651,64 +766,60 @@ mod tests { #[test] fn test_disk_usage_with_clones() -> Result<()> { - use std::io::Write; - let dm = Arc::new(DiskManagerBuilder::default().build()?); - let mut temp_file = dm.create_tmp_file("Testing")?; + let temp_file = dm.create_tmp_file("Testing")?; // Write some data - temp_file.inner().as_file().write_all(b"test data")?; - temp_file.update_disk_usage()?; + let mut writer = temp_file.open_writer()?; + writer.write_all(b"test data")?; - let usage_after_write = temp_file.current_disk_usage(); + let usage_after_write = temp_file.size().unwrap(); assert!(usage_after_write > 0); assert_eq!(dm.used_disk_space(), usage_after_write); // Clone the file - let clone1 = temp_file.clone(); - let clone2 = temp_file.clone(); + let clone1 = Arc::clone(&temp_file); + let clone2 = Arc::clone(&temp_file); // All clones should see the same disk usage - assert_eq!(clone1.current_disk_usage(), usage_after_write); - assert_eq!(clone2.current_disk_usage(), usage_after_write); - + assert_eq!(clone1.size().unwrap(), usage_after_write); + assert_eq!(clone2.size().unwrap(), usage_after_write); // Global disk usage should still be the same (not multiplied by number of clones) assert_eq!(dm.used_disk_space(), usage_after_write); // Write more data through one clone - clone1.inner().as_file().write_all(b" more data")?; - let mut mutable_clone1 = clone1; - mutable_clone1.update_disk_usage()?; + let mut clone_writer = clone1.open_writer()?; + clone_writer.write_all(b" more data")?; - let new_usage = mutable_clone1.current_disk_usage(); + let new_usage = clone1.size().unwrap(); assert!(new_usage > usage_after_write); - // All clones should see the updated disk usage - assert_eq!(temp_file.current_disk_usage(), new_usage); - assert_eq!(clone2.current_disk_usage(), new_usage); - assert_eq!(mutable_clone1.current_disk_usage(), new_usage); + assert_eq!(temp_file.size().unwrap(), new_usage); + assert_eq!(clone2.size().unwrap(), new_usage); + assert_eq!(clone1.size().unwrap(), new_usage); // Global disk usage should reflect the new size (not multiplied) assert_eq!(dm.used_disk_space(), new_usage); // Drop one clone - drop(mutable_clone1); + drop(clone_writer); + drop(clone1); // Disk usage should NOT change (other clones still exist) assert_eq!(dm.used_disk_space(), new_usage); - assert_eq!(temp_file.current_disk_usage(), new_usage); - assert_eq!(clone2.current_disk_usage(), new_usage); + assert_eq!(temp_file.size().unwrap(), new_usage); + assert_eq!(clone2.size().unwrap(), new_usage); // Drop another clone drop(clone2); // Disk usage should still NOT change (original still exists) assert_eq!(dm.used_disk_space(), new_usage); - assert_eq!(temp_file.current_disk_usage(), new_usage); + assert_eq!(temp_file.size().unwrap(), new_usage); // Drop the original + drop(writer); drop(temp_file); - // Now disk usage should return to 0 (last reference dropped) assert_eq!(dm.used_disk_space(), 0); @@ -717,29 +828,27 @@ mod tests { #[test] fn test_disk_usage_clones_dropped_out_of_order() -> Result<()> { - use std::io::Write; - let dm = Arc::new(DiskManagerBuilder::default().build()?); - let mut temp_file = dm.create_tmp_file("Testing")?; + let temp_file = dm.create_tmp_file("Testing")?; + let mut writer = temp_file.open_writer()?; // Write data - temp_file.inner().as_file().write_all(b"test")?; - temp_file.update_disk_usage()?; + writer.write_all(b"test")?; - let usage = temp_file.current_disk_usage(); + let usage = temp_file.size().unwrap(); assert_eq!(dm.used_disk_space(), usage); // Create multiple clones - let clone1 = temp_file.clone(); - let clone2 = temp_file.clone(); - let clone3 = temp_file.clone(); + let clone1 = Arc::clone(&temp_file); + let clone2 = Arc::clone(&temp_file); + let clone3 = Arc::clone(&temp_file); // Drop the original first (out of order) drop(temp_file); // Disk usage should still be tracked (clones exist) assert_eq!(dm.used_disk_space(), usage); - assert_eq!(clone1.current_disk_usage(), usage); + assert_eq!(clone1.size().unwrap(), usage); // Drop clones in different order drop(clone2); @@ -759,25 +868,24 @@ mod tests { #[test] fn test_disk_usage_multiple_files() -> Result<()> { - use std::io::Write; - let dm = Arc::new(DiskManagerBuilder::default().build()?); // Create multiple temp files - let mut file1 = dm.create_tmp_file("Testing1")?; - let mut file2 = dm.create_tmp_file("Testing2")?; + let file1 = dm.create_tmp_file("Testing1")?; + let file2 = dm.create_tmp_file("Testing2")?; + + let mut writer1 = file1.open_writer()?; + let mut writer2 = file2.open_writer()?; // Write to first file - file1.inner().as_file().write_all(b"file1")?; - file1.update_disk_usage()?; - let usage1 = file1.current_disk_usage(); + writer1.write_all(b"file1")?; + let usage1 = file1.size().unwrap(); assert_eq!(dm.used_disk_space(), usage1); // Write to second file - file2.inner().as_file().write_all(b"file2 data")?; - file2.update_disk_usage()?; - let usage2 = file2.current_disk_usage(); + writer2.write_all(b"file2 data")?; + let usage2 = file2.size().unwrap(); // Global usage should be sum of both files assert_eq!(dm.used_disk_space(), usage1 + usage2); @@ -796,4 +904,275 @@ mod tests { Ok(()) } + + #[test] + fn test_dynamic_limit_adjustment_through_shared_ref() -> Result<()> { + // Verify that set_max_temp_directory_size works through &self (not &mut self). + // This is the key behavioral change: the limit can be adjusted at runtime + // without exclusive access, enabling dynamic resize while queries are running. + let dm = DiskManager::builder() + .with_max_temp_directory_size(1024) + .build()?; + let dm = Arc::new(dm); + + assert_eq!(dm.max_temp_directory_size(), 1024); + + // Adjust through shared reference (simulates concurrent access via Arc) + dm.set_max_temp_directory_size(2048)?; + assert_eq!(dm.max_temp_directory_size(), 2048); + + // Can also decrease + dm.set_max_temp_directory_size(512)?; + assert_eq!(dm.max_temp_directory_size(), 512); + + Ok(()) + } + + #[test] + fn test_dynamic_limit_concurrent_access() -> Result<()> { + // Verify that multiple threads can read and write the limit concurrently + let dm = Arc::new( + DiskManager::builder() + .with_max_temp_directory_size(1000) + .build()?, + ); + + let handles: Vec<_> = (0..8) + .map(|i| { + let dm = Arc::clone(&dm); + std::thread::spawn(move || { + // Each thread sets a different limit and reads it back + let new_limit = (i + 1) * 1000; + dm.set_max_temp_directory_size(new_limit).unwrap(); + // Read should return SOME value set by one of the threads + let current = dm.max_temp_directory_size(); + assert!((1000..=8000).contains(¤t)); + }) + }) + .collect(); + + for h in handles { + h.join().unwrap(); + } + + // Final value should be one of the values set by threads + let final_val = dm.max_temp_directory_size(); + assert!((1000..=8000).contains(&final_val)); + + Ok(()) + } + + #[test] + fn test_max_spill_merge_fan_in_builder_and_dynamic_update() -> Result<()> { + let dm = Arc::new( + DiskManager::builder() + .with_max_spill_merge_fan_in(8) + .build()?, + ); + + assert_eq!(dm.max_spill_merge_fan_in(), 8); + + dm.set_max_spill_merge_fan_in(4); + assert_eq!(dm.max_spill_merge_fan_in(), 4); + + dm.set_max_spill_merge_fan_in(0); + assert_eq!(dm.max_spill_merge_fan_in(), 0); + + Ok(()) + } + + #[test] + fn test_disabled_disk_manager_rejects_nonzero_limit() -> Result<()> { + let dm = DiskManager::builder() + .with_mode(DiskManagerMode::Disabled) + .build()?; + let dm = Arc::new(dm); + + // Setting non-zero limit on disabled DiskManager should error + let result = dm.set_max_temp_directory_size(1024); + assert!(result.is_err()); + + // Setting zero is OK + assert!(dm.set_max_temp_directory_size(0).is_ok()); + + Ok(()) + } + + #[test] + fn test_limit_decrease_below_current_usage() -> Result<()> { + // Scenario: DiskManager has 100GB limit, currently using 80GB. + // Admin lowers limit to 60GB. What happens? + // + // Expected behavior: + // - Existing spill files remain on disk (not deleted) + // - used_disk_space still reports 80GB + // - New spill writes FAIL immediately (80GB > 60GB new limit) + // - Once old queries complete and release their files (used drops below 60GB), + // new spill writes succeed again + // + // This demonstrates graceful degradation: lowering the limit doesn't + // reclaim existing files (would break running queries), but prevents + // additional spilling until usage drops naturally. + let dm = DiskManager::builder() + .with_max_temp_directory_size(100 * 1024 * 1024 * 1024) // 100GB + .build()?; + let dm = Arc::new(dm); + + // Simulate 80GB of existing spill usage + dm.used_disk_space + .store(80 * 1024 * 1024 * 1024, Ordering::Relaxed); + + assert_eq!(dm.max_temp_directory_size(), 100 * 1024 * 1024 * 1024); + assert_eq!(dm.used_disk_space(), 80 * 1024 * 1024 * 1024); + + // Lower the limit to 60GB (below current usage) + dm.set_max_temp_directory_size(60 * 1024 * 1024 * 1024)?; + assert_eq!(dm.max_temp_directory_size(), 60 * 1024 * 1024 * 1024); + + // Current usage (80GB) now exceeds the new limit (60GB). + // The used_disk_space is NOT reclaimed — existing files stay. + assert_eq!(dm.used_disk_space(), 80 * 1024 * 1024 * 1024); + + // Any attempt to write MORE would be rejected at the SpillWriter level + // because used_disk_space(80GB) > max_temp_directory_size(60GB). + // (SpillWriter check: `global_disk_usage > limit` returns ResourcesExhausted) + + // Simulate old queries completing: usage drops to 50GB + dm.used_disk_space + .store(50 * 1024 * 1024 * 1024, Ordering::Relaxed); + + // Now usage (50GB) < limit (60GB) — new spill writes would succeed again + assert!(dm.used_disk_space() < dm.max_temp_directory_size()); + + Ok(()) + } + + #[test] + fn test_limit_decrease_with_concurrent_queries() -> Result<()> { + // Scenario: Multiple threads spilling while limit is lowered concurrently. + // Demonstrates that: + // 1. In-flight spills that started before the limit change complete normally + // (they already incremented used_disk_space) + // 2. New spills after the limit change respect the new lower limit + // 3. No data corruption or panics from concurrent access + let dm = Arc::new( + DiskManager::builder() + .with_max_temp_directory_size(100 * 1024 * 1024) // 100MB + .build()?, + ); + + let barrier = Arc::new(std::sync::Barrier::new(5)); + + // 4 threads simulate concurrent spilling + let spill_handles: Vec<_> = (0..4) + .map(|_| { + let dm = Arc::clone(&dm); + let barrier = Arc::clone(&barrier); + std::thread::spawn(move || { + barrier.wait(); + // Simulate spill: increment used_disk_space + dm.used_disk_space + .fetch_add(10 * 1024 * 1024, Ordering::Relaxed); + std::thread::sleep(std::time::Duration::from_millis(10)); + // Simulate cleanup + dm.used_disk_space + .fetch_sub(10 * 1024 * 1024, Ordering::Relaxed); + }) + }) + .collect(); + + // 1 thread lowers the limit mid-flight + let dm_resize = Arc::clone(&dm); + let resize_barrier = Arc::clone(&barrier); + let resize_handle = std::thread::spawn(move || { + resize_barrier.wait(); + // Lower limit while spills are in progress + dm_resize + .set_max_temp_directory_size(30 * 1024 * 1024) // 30MB + .unwrap(); + }); + + for h in spill_handles { + h.join().unwrap(); + } + resize_handle.join().unwrap(); + + // After all threads complete: + // - Limit is 30MB (last set by resize thread) + // - used_disk_space is 0 (all spills cleaned up) + // - No panics, no corruption + assert_eq!(dm.max_temp_directory_size(), 30 * 1024 * 1024); + assert_eq!(dm.used_disk_space(), 0); + + Ok(()) + } + + #[test] + fn test_rollback_on_limit_exceeded_then_drop_returns_to_zero() -> Result<()> { + // This test verifies that lowering the limit, failing a spill write, + // and then dropping the file leaves used_disk_space at zero. + // + // Without the rollback fix, the global counter would be permanently + // inflated by the delta between the new and old file sizes. + + let dm = Arc::new( + DiskManager::builder() + .with_max_temp_directory_size(10 * 1024 * 1024) // 10MB + .build()?, + ); + + let file = dm.create_tmp_file("test_rollback")?; + + let mut writer = file.open_writer()?; + + // Create a temp file and write some data + { + let data = vec![0u8; 1024]; // 1KB + writer.write_all(&data)?; + } + + let usage_after_first_write = dm.used_disk_space(); + assert!(usage_after_first_write > 0); + + // Write more data to grow the file + { + let data = vec![0u8; 4 * 1024]; // 4KB more + writer.write_all(&data)?; + } + + let usage_after_second_write = dm.used_disk_space(); + assert!(usage_after_second_write > usage_after_first_write); + + // Now lower the limit to 1 byte — below current usage + dm.set_max_temp_directory_size(1)?; + + // Write even more data + { + let data = vec![0u8; 2 * 1024]; // 2KB more + + // This write should FAIL (exceeds new 1-byte limit) + let result = writer.write_all(&data); + + assert!(result.is_err()); + assert!( + result + .unwrap_err() + .to_string() + .contains("exceeded the allowable limit") + ); + } + + // Critical check: used_disk_space should still equal the LAST + // successful update (before the failed one), not be inflated + assert_eq!(dm.used_disk_space(), usage_after_second_write); + + // Drop the writer and file — should subtract the last successful file size + drop(writer); + drop(file); + + // After drop: used_disk_space must be zero (no leak) + assert_eq!(dm.used_disk_space(), 0); + + Ok(()) + } } diff --git a/datafusion/execution/src/lib.rs b/datafusion/execution/src/lib.rs index 1a8da9459ae10..5af7064f1cb8b 100644 --- a/datafusion/execution/src/lib.rs +++ b/datafusion/execution/src/lib.rs @@ -27,6 +27,7 @@ //! DataFusion execution configuration and runtime structures +mod async_stream; pub mod cache; pub mod config; pub mod disk_manager; @@ -35,6 +36,7 @@ pub mod object_store; #[cfg(feature = "parquet_encryption")] pub mod parquet_encryption; pub mod runtime_env; +pub mod spill_file; mod stream; mod task; @@ -44,7 +46,9 @@ pub mod registry { }; } +pub use async_stream::{Emitter, TryEmitter, async_stream, async_try_stream}; pub use disk_manager::DiskManager; pub use registry::FunctionRegistry; +pub use spill_file::{SpillFile, SpillWriter, TempFileFactory}; pub use stream::{RecordBatchStream, SendableRecordBatchStream}; pub use task::{TaskContext, TaskContextProvider}; diff --git a/datafusion/execution/src/memory_pool/mod.rs b/datafusion/execution/src/memory_pool/mod.rs index 2b36ee7f40add..40a79d136b84e 100644 --- a/datafusion/execution/src/memory_pool/mod.rs +++ b/datafusion/execution/src/memory_pool/mod.rs @@ -24,6 +24,7 @@ use std::fmt::Display; use std::hash::{Hash, Hasher}; use std::{cmp::Ordering, sync::Arc, sync::atomic}; +mod peak_recording; mod pool; #[cfg(feature = "arrow_buffer_pool")] @@ -36,6 +37,7 @@ pub mod proxy { pub use datafusion_common::{ human_readable_count, human_readable_duration, human_readable_size, units, }; +pub use peak_recording::*; pub use pool::*; /// Tracks and potentially limits memory use across operators during execution. diff --git a/datafusion/execution/src/memory_pool/peak_recording.rs b/datafusion/execution/src/memory_pool/peak_recording.rs new file mode 100644 index 0000000000000..b407cc0eaf36b --- /dev/null +++ b/datafusion/execution/src/memory_pool/peak_recording.rs @@ -0,0 +1,377 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Records the peak [`MemoryPool`] reservation reached during a benchmark. +//! +//! DataFusion's [`MemoryPool`] deliberately accounts for only the "large" +//! allocations that scale with input size; intermediate batches flowing between +//! operators are assumed to be small and are left untracked. The [`MemoryPool`] +//! documentation therefore advises reserving "some overhead (e.g. 10%)" on top +//! of the configured limit. +//! +//! Nothing reports what that overhead actually is, because the peak reservation +//! itself is never recorded — [`MemoryPool::reserved`] is a live value that has +//! usually fallen back to zero by the time a query finishes. This module records +//! the high-water mark so benchmarks can emit it alongside the peak RSS that +//! `print_memory_stats` already prints, making the gap between the two +//! measurable. +//! +//! This is measurement only: nothing here enforces a relationship between the +//! two numbers. +//! +//! What lands in the peak is whatever the pool accounts for, so this follows +//! the accounting rather than fixing it in place. Arrow-side reservations made +//! through `ArrowMemoryPool` are included, because that adapter grows a +//! DataFusion reservation against the pool it wraps; nothing claims buffers +//! today, but the peak picks it up when something does. + +use std::{ + fmt::{Debug, Display, Formatter}, + sync::{ + Arc, + atomic::{AtomicUsize, Ordering}, + }, +}; + +use super::{MemoryConsumer, MemoryLimit, MemoryPool, MemoryReservation}; +use datafusion_common::Result; + +/// Wraps a [`MemoryPool`], recording the high-water mark of +/// [`MemoryPool::reserved`] as reservations come and go. +/// +/// Every method delegates to the wrapped pool, so wrapping does not change how +/// memory is granted, limited, or reported. The one thing it does change is +/// downcasting: `rt.memory_pool.downcast_ref::()` now finds this +/// wrapper instead of the pool it wraps. Nothing in the benchmarks relies on +/// that, and [`Self::from_pool`] uses the same mechanism to find the recorder. +/// +/// Both high-water marks are held per instance, so a benchmark that builds a +/// fresh runtime per query gets a reading scoped to that query without any +/// coordination. +/// +/// # Example +/// +/// ``` +/// # use std::sync::Arc; +/// # use datafusion_execution::memory_pool::{GreedyMemoryPool, MemoryConsumer, MemoryPool, PeakRecordingPool}; +/// let recording = Arc::new(PeakRecordingPool::new(Arc::new(GreedyMemoryPool::new(1024)))); +/// let pool: Arc = Arc::clone(&recording) as _; +/// +/// let reservation = MemoryConsumer::new("example").register(&pool); +/// reservation.try_grow(512)?; +/// reservation.shrink(512); +/// +/// // The pool is back to empty, but the high-water mark is retained. +/// assert_eq!(pool.reserved(), 0); +/// assert_eq!(recording.peak_reserved(), 512); +/// +/// // The recorder can also be recovered from the pool it was installed as. +/// assert_eq!(PeakRecordingPool::from_pool(&*pool).unwrap().peak_reserved(), 512); +/// # Ok::<(), datafusion_common::DataFusionError>(()) +/// ``` +pub struct PeakRecordingPool { + inner: Arc, + /// Running total of everything granted through this wrapper, kept so the + /// peak can be maintained without asking `inner` for its total. + reserved: AtomicUsize, + /// High-water mark since the last [`PeakRecordingPool::reset_peak`]. + peak: AtomicUsize, + /// High-water mark since this pool was created. Never reset. + max: AtomicUsize, +} + +impl PeakRecordingPool { + /// Wrap `inner`, recording its peak reservation from here on. + /// + /// `inner` is expected to be empty: the running total starts at zero, so + /// anything reserved before wrapping is not counted. + pub fn new(inner: Arc) -> Self { + Self { + inner, + reserved: AtomicUsize::new(0), + peak: AtomicUsize::new(0), + max: AtomicUsize::new(0), + } + } + + /// The recorder installed as `pool`, if there is one. + /// + /// Returns `None` whenever a benchmark runs without a memory limit, since + /// `CommonOpt::runtime_env_builder` only installs the wrapper alongside a + /// pool it has a limit for. + pub fn from_pool(pool: &dyn MemoryPool) -> Option<&Self> { + pool.downcast_ref::() + } + + /// Peak reservation, in bytes, since the last [`Self::reset_peak`]. + pub fn peak_reserved(&self) -> usize { + self.peak.load(Ordering::Relaxed) + } + + /// Peak reservation, in bytes, since this pool was created. + /// + /// Unlike [`Self::peak_reserved`] this is never reset, so it reports the + /// peak across every query that shared this pool. + pub fn max_reserved(&self) -> usize { + self.max.load(Ordering::Relaxed) + } + + /// Reset the value returned by [`Self::peak_reserved`] to what is reserved + /// right now, so the next reading covers only what follows. + /// + /// `BenchmarkRun::start_new_case` calls this, giving each benchmark query + /// its own reading. Anything still held when a query starts — data the + /// benchmark loaded up front, say — stays in the reading, since the query + /// runs with those bytes reserved. + pub fn reset_peak(&self) { + self.peak + .store(self.reserved.load(Ordering::Relaxed), Ordering::Relaxed); + } + + /// Add `additional` granted bytes to the running total and publish it to + /// both high-water marks. + /// + /// Accumulating deltas rather than reading [`MemoryPool::reserved`] keeps + /// the wrapped pool's own bookkeeping off this path: `FairSpillPool` takes + /// its state lock to answer `reserved()`, which would double the lock + /// traffic of every accounted allocation in the benchmark being measured. + /// The total stays exact because the trait grants exactly what is asked + /// for — `grow` is infallible and `try_grow` either grants `additional` or + /// returns an error, leaving the reservation untouched. + fn record(&self, additional: usize) { + let reserved = + self.reserved.fetch_add(additional, Ordering::Relaxed) + additional; + self.peak.fetch_max(reserved, Ordering::Relaxed); + self.max.fetch_max(reserved, Ordering::Relaxed); + } +} + +impl Debug for PeakRecordingPool { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PeakRecordingPool") + .field("inner", &self.inner) + .field("peak", &self.peak_reserved()) + .field("max", &self.max_reserved()) + .finish() + } +} + +impl Display for PeakRecordingPool { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + // Deferring to the wrapped pool keeps `SHOW ALL`-style output and error + // messages identical to running without the wrapper. + Display::fmt(&self.inner, f) + } +} + +impl MemoryPool for PeakRecordingPool { + fn name(&self) -> &str { + self.inner.name() + } + + fn register(&self, consumer: &MemoryConsumer) { + self.inner.register(consumer); + } + + fn unregister(&self, consumer: &MemoryConsumer) { + self.inner.unregister(consumer); + } + + fn grow(&self, reservation: &MemoryReservation, additional: usize) { + self.inner.grow(reservation, additional); + self.record(additional); + } + + fn shrink(&self, reservation: &MemoryReservation, shrink: usize) { + self.inner.shrink(reservation, shrink); + self.reserved.fetch_sub(shrink, Ordering::Relaxed); + } + + fn try_grow(&self, reservation: &MemoryReservation, additional: usize) -> Result<()> { + self.inner.try_grow(reservation, additional)?; + self.record(additional); + Ok(()) + } + + fn reserved(&self) -> usize { + self.inner.reserved() + } + + fn memory_limit(&self) -> MemoryLimit { + self.inner.memory_limit() + } +} + +#[cfg(test)] +mod tests { + use crate::memory_pool::GreedyMemoryPool; + + use super::*; + + /// A recording pool over a `GreedyMemoryPool`, returned both as the + /// recorder (to read the marks) and as the pool reservations register with. + fn pool(limit: usize) -> (Arc, Arc) { + let recording = Arc::new(PeakRecordingPool::new(Arc::new( + GreedyMemoryPool::new(limit), + ))); + let pool = Arc::clone(&recording) as Arc; + (recording, pool) + } + + #[test] + fn records_high_water_mark_across_reservations() { + let (recording, pool) = pool(1024); + + let a = MemoryConsumer::new("a").register(&pool); + let b = MemoryConsumer::new("b").register(&pool); + + a.try_grow(300).unwrap(); + b.try_grow(400).unwrap(); + // Peak of the sum, not the largest single reservation. + assert_eq!(recording.peak_reserved(), 700); + + a.shrink(300); + b.try_grow(100).unwrap(); + + // Falling back below the peak leaves it untouched, and the later growth + // does not reach it. + assert_eq!(pool.reserved(), 500); + assert_eq!(recording.peak_reserved(), 700); + } + + #[test] + fn failed_growth_does_not_move_the_peak() { + let (recording, pool) = pool(1024); + + let reservation = MemoryConsumer::new("a").register(&pool); + reservation.try_grow(600).unwrap(); + reservation + .try_grow(600) + .expect_err("should exceed the 1024 byte pool"); + + assert_eq!(recording.peak_reserved(), 600); + } + + #[test] + fn reset_clears_the_window_but_not_the_run_maximum() { + let (recording, pool) = pool(1024); + + let reservation = MemoryConsumer::new("a").register(&pool); + reservation.try_grow(800).unwrap(); + reservation.shrink(800); + + recording.reset_peak(); + assert_eq!(recording.peak_reserved(), 0); + assert_eq!(recording.max_reserved(), 800); + + reservation.try_grow(100).unwrap(); + assert_eq!(recording.peak_reserved(), 100); + assert_eq!(recording.max_reserved(), 800); + } + + #[test] + fn reset_keeps_what_is_still_reserved() { + let (recording, pool) = pool(1024); + + // Something a benchmark loaded up front and holds across queries. + let held = MemoryConsumer::new("held").register(&pool); + held.try_grow(300).unwrap(); + + recording.reset_peak(); + assert_eq!(recording.peak_reserved(), 300); + + let query = MemoryConsumer::new("query").register(&pool); + query.try_grow(200).unwrap(); + assert_eq!(recording.peak_reserved(), 500); + } + + #[test] + fn marks_are_per_instance() { + let (one, one_pool) = pool(1024); + let (two, _two_pool) = pool(1024); + + MemoryConsumer::new("a") + .register(&one_pool) + .try_grow(512) + .unwrap(); + + assert_eq!(one.peak_reserved(), 512); + assert_eq!(two.peak_reserved(), 0); + } + + #[test] + fn is_recoverable_from_the_pool_it_is_installed_as() { + let (recording, pool) = pool(1024); + + MemoryConsumer::new("a") + .register(&pool) + .try_grow(512) + .unwrap(); + + let found = PeakRecordingPool::from_pool(&*pool).expect("recorder installed"); + assert_eq!(found.peak_reserved(), recording.peak_reserved()); + + // A pool with no recorder in front of it reports nothing. + let plain: Arc = Arc::new(GreedyMemoryPool::new(1024)); + assert!(PeakRecordingPool::from_pool(&*plain).is_none()); + } + + #[test] + fn delegates_limit_and_name_to_the_wrapped_pool() { + let inner: Arc = Arc::new(GreedyMemoryPool::new(4096)); + let wrapped = PeakRecordingPool::new(Arc::clone(&inner)); + + assert_eq!(wrapped.name(), inner.name()); + assert_eq!(wrapped.to_string(), inner.to_string()); + assert!(matches!(wrapped.memory_limit(), MemoryLimit::Finite(4096))); + } + + /// Arrow-side reservations reach the recorder too. + /// + /// [`ArrowMemoryPool`] implements Arrow's `MemoryPool` by growing a + /// DataFusion [`MemoryReservation`] against the pool it wraps, so a buffer + /// claimed through it lands in `grow` here. Nothing in DataFusion claims + /// buffers yet (see apache/datafusion#22898), but when something does, the + /// bytes show up in this peak without further changes — as long as the + /// adapter is built from the `RuntimeEnv`'s pool, which is the wrapped one. + /// This test pins that. + /// + /// Only compiled with `--features arrow_buffer_pool`, since that's what + /// gates `crate::memory_pool::arrow` and `arrow_buffer::MemoryPool` in the + /// first place; not part of this crate's default feature set. + #[cfg(feature = "arrow_buffer_pool")] + #[test] + fn records_reservations_arriving_through_the_arrow_adapter() { + use crate::memory_pool::arrow::ArrowMemoryPool; + use arrow_buffer::MemoryPool as ArrowMemoryPoolTrait; + + let (recording, pool) = pool(4096); + + let arrow_pool = + ArrowMemoryPool::new(Arc::clone(&pool), MemoryConsumer::new("arrow")); + let reservation = arrow_pool.reserve(1024); + + // The Arrow-side reservation is visible as DataFusion pool usage... + assert_eq!(pool.reserved(), 1024); + assert_eq!(recording.peak_reserved(), 1024); + + // ...and dropping it releases the bytes while the peak is retained. + drop(reservation); + assert_eq!(pool.reserved(), 0); + assert_eq!(recording.peak_reserved(), 1024); + } +} diff --git a/datafusion/execution/src/memory_pool/pool.rs b/datafusion/execution/src/memory_pool/pool.rs index 52b601d5cd78b..d854cbd627cec 100644 --- a/datafusion/execution/src/memory_pool/pool.rs +++ b/datafusion/execution/src/memory_pool/pool.rs @@ -64,7 +64,7 @@ impl MemoryPool for UnboundedMemoryPool { impl Display for UnboundedMemoryPool { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { let used = self.used.load(Ordering::Relaxed); - write!(f, "{}(used: {})", &self.name(), human_readable_size(used)) + write!(f, "{}(used: {})", self.name(), human_readable_size(used)) } } @@ -135,7 +135,7 @@ impl Display for GreedyMemoryPool { write!( f, "{}(used: {}, pool_size: {})", - &self.name(), + self.name(), human_readable_size(used), human_readable_size(self.pool_size) ) @@ -290,7 +290,7 @@ impl Display for FairSpillPool { write!( f, "{}(pool_size: {})", - &self.name(), + self.name(), human_readable_size(self.pool_size), ) } @@ -416,9 +416,9 @@ impl Display for TrackConsumersPool { write!( f, "{}(inner_pool: {}, num_of_top_consumers: {})", - &self.name(), - &self.inner, - &self.top, + self.name(), + self.inner, + self.top, ) } } diff --git a/datafusion/execution/src/runtime_env.rs b/datafusion/execution/src/runtime_env.rs index 5b90f28a141ef..fcfe51267e65f 100644 --- a/datafusion/execution/src/runtime_env.rs +++ b/datafusion/execution/src/runtime_env.rs @@ -18,8 +18,7 @@ //! Execution [`RuntimeEnv`] environment that manages access to object //! store, memory manager, disk manager. -#[expect(deprecated)] -use crate::disk_manager::{DiskManagerConfig, SpillingProgress}; +use crate::disk_manager::SpillingProgress; use crate::{ disk_manager::{DiskManager, DiskManagerBuilder, DiskManagerMode}, memory_pool::{ @@ -91,57 +90,77 @@ impl Debug for RuntimeEnv { } } -/// Creates runtime configuration entries with the provided values -/// -/// This helper function defines the structure and metadata for all runtime configuration -/// entries to avoid duplication between `RuntimeEnv::config_entries()` and -/// `RuntimeEnvBuilder::entries()`. -fn create_runtime_config_entries( +struct RuntimeConfigValues { memory_limit: Option, max_temp_directory_size: Option, + max_spill_merge_fan_in: Option, temp_directory: Option, metadata_cache_limit: Option, list_files_cache_limit: Option, list_files_cache_ttl: Option, file_statistics_cache_limit: Option, -) -> Vec { - vec![ - ConfigEntry { - key: "datafusion.runtime.memory_limit".to_string(), - value: memory_limit, - description: "Maximum memory limit for query execution. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", - }, - ConfigEntry { - key: "datafusion.runtime.max_temp_directory_size".to_string(), - value: max_temp_directory_size, - description: "Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", - }, - ConfigEntry { - key: "datafusion.runtime.temp_directory".to_string(), - value: temp_directory, - description: "The path to the temporary file directory.", - }, - ConfigEntry { - key: "datafusion.runtime.metadata_cache_limit".to_string(), - value: metadata_cache_limit, - description: "Maximum memory to use for file metadata cache such as Parquet metadata. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", - }, - ConfigEntry { - key: "datafusion.runtime.list_files_cache_limit".to_string(), - value: list_files_cache_limit, - description: "Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", - }, - ConfigEntry { - key: "datafusion.runtime.list_files_cache_ttl".to_string(), - value: list_files_cache_ttl, - description: "TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes.", - }, - ConfigEntry { - key: "datafusion.runtime.file_statistics_cache_limit".to_string(), - value: file_statistics_cache_limit, - description: "Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", - }, - ] +} + +impl RuntimeConfigValues { + /// Creates runtime configuration entries with the provided values. + /// + /// This defines the structure and metadata for all runtime configuration + /// entries to avoid duplication between `RuntimeEnv::config_entries()` and + /// `RuntimeEnvBuilder::entries()`. + fn into_config_entries(self) -> Vec { + let Self { + memory_limit, + max_temp_directory_size, + max_spill_merge_fan_in, + temp_directory, + metadata_cache_limit, + list_files_cache_limit, + list_files_cache_ttl, + file_statistics_cache_limit, + } = self; + vec![ + ConfigEntry { + key: "datafusion.runtime.memory_limit".to_string(), + value: memory_limit, + description: "Maximum memory limit for query execution. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", + }, + ConfigEntry { + key: "datafusion.runtime.max_temp_directory_size".to_string(), + value: max_temp_directory_size, + description: "Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", + }, + ConfigEntry { + key: "datafusion.runtime.max_spill_merge_fan_in".to_string(), + value: max_spill_merge_fan_in, + description: "Maximum number of spill files opened by one external merge pass. Use 0 for unlimited. Values below 2 still use 2 so a merge can make progress.", + }, + ConfigEntry { + key: "datafusion.runtime.temp_directory".to_string(), + value: temp_directory, + description: "The path to the temporary file directory.", + }, + ConfigEntry { + key: "datafusion.runtime.metadata_cache_limit".to_string(), + value: metadata_cache_limit, + description: "Maximum memory to use for file metadata cache such as Parquet metadata. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", + }, + ConfigEntry { + key: "datafusion.runtime.list_files_cache_limit".to_string(), + value: list_files_cache_limit, + description: "Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", + }, + ConfigEntry { + key: "datafusion.runtime.list_files_cache_ttl".to_string(), + value: list_files_cache_ttl, + description: "TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes.", + }, + ConfigEntry { + key: "datafusion.runtime.file_statistics_cache_limit".to_string(), + value: file_statistics_cache_limit, + description: "Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes.", + }, + ] + } } impl RuntimeEnv { @@ -269,6 +288,8 @@ impl RuntimeEnv { let max_temp_dir_size = self.disk_manager.max_temp_directory_size(); let max_temp_dir_value = format_byte_size(max_temp_dir_size); + let max_spill_merge_fan_in = + self.disk_manager.max_spill_merge_fan_in().to_string(); let temp_paths = self.disk_manager.temp_dir_paths(); let temp_dir_value = if temp_paths.is_empty() { @@ -310,15 +331,17 @@ impl RuntimeEnv { .expect("File statistics cache size conversion failed"), ); - create_runtime_config_entries( - memory_limit_value, - Some(max_temp_dir_value), - temp_dir_value, - Some(metadata_cache_value), - Some(list_files_cache_value), + RuntimeConfigValues { + memory_limit: memory_limit_value, + max_temp_directory_size: Some(max_temp_dir_value), + max_spill_merge_fan_in: Some(max_spill_merge_fan_in), + temp_directory: temp_dir_value, + metadata_cache_limit: Some(metadata_cache_value), + list_files_cache_limit: Some(list_files_cache_value), list_files_cache_ttl, - Some(file_statistics_cache_value), - ) + file_statistics_cache_limit: Some(file_statistics_cache_value), + } + .into_config_entries() } } @@ -333,9 +356,8 @@ impl Default for RuntimeEnv { /// See example on [`RuntimeEnv`] #[derive(Clone)] pub struct RuntimeEnvBuilder { - #[expect(deprecated)] /// DiskManager to manage temporary disk file usage - pub disk_manager: DiskManagerConfig, + pub disk_manager: Option>, /// DiskManager builder to manager temporary disk file usage pub disk_manager_builder: Option, /// [`MemoryPool`] from which to allocate memory @@ -371,14 +393,6 @@ impl RuntimeEnvBuilder { } } - #[expect(deprecated)] - #[deprecated(since = "48.0.0", note = "Use with_disk_manager_builder instead")] - /// Customize disk manager - pub fn with_disk_manager(mut self, disk_manager: DiskManagerConfig) -> Self { - self.disk_manager = disk_manager; - self - } - /// Customize the disk manager builder pub fn with_disk_manager_builder(mut self, disk_manager: DiskManagerBuilder) -> Self { self.disk_manager_builder = Some(disk_manager); @@ -435,6 +449,14 @@ impl RuntimeEnvBuilder { self.with_disk_manager_builder(builder.with_max_temp_directory_size(size)) } + /// Limit the number of spill files opened by one external merge pass. + /// + /// A value of 0 means unlimited. + pub fn with_max_spill_merge_fan_in(mut self, fan_in: usize) -> Self { + let builder = self.disk_manager_builder.take().unwrap_or_default(); + self.with_disk_manager_builder(builder.with_max_spill_merge_fan_in(fan_in)) + } + /// Specify the limit of the file-embedded metadata cache, in bytes. pub fn with_metadata_cache_limit(mut self, limit: usize) -> Self { self.cache_manager = self.cache_manager.with_metadata_cache_limit(limit); @@ -472,14 +494,15 @@ impl RuntimeEnvBuilder { let memory_pool = memory_pool.unwrap_or_else(|| Arc::new(UnboundedMemoryPool::default())); + let disk_manager: Arc = match (disk_manager, disk_manager_builder) { + (_, Some(builder)) => Arc::new(builder.build()?), + (Some(manager), None) => manager, + (None, None) => Arc::new(DiskManagerBuilder::default().build()?), + }; + Ok(RuntimeEnv { memory_pool, - disk_manager: if let Some(builder) = disk_manager_builder { - Arc::new(builder.build()?) - } else { - #[expect(deprecated)] - DiskManager::try_new(disk_manager)? - }, + disk_manager, cache_manager: CacheManager::try_new(&cache_manager)?, object_store_registry, #[cfg(feature = "parquet_encryption")] @@ -511,10 +534,7 @@ impl RuntimeEnvBuilder { }; Self { - #[expect(deprecated)] - disk_manager: DiskManagerConfig::Existing(Arc::clone( - &runtime_env.disk_manager, - )), + disk_manager: Some(Arc::clone(&runtime_env.disk_manager)), disk_manager_builder: None, memory_pool: Some(Arc::clone(&runtime_env.memory_pool)), cache_manager: cache_config, @@ -528,15 +548,17 @@ impl RuntimeEnvBuilder { /// Returns a list of all available runtime configurations with their current values and descriptions pub fn entries(&self) -> Vec { - create_runtime_config_entries( - None, - Some("100G".to_string()), - None, - Some("50M".to_owned()), - Some("1M".to_owned()), - None, - Some("20M".to_owned()), - ) + RuntimeConfigValues { + memory_limit: None, + max_temp_directory_size: Some("100G".to_string()), + max_spill_merge_fan_in: Some("0".to_string()), + temp_directory: None, + metadata_cache_limit: Some("50M".to_owned()), + list_files_cache_limit: Some("1M".to_owned()), + list_files_cache_ttl: None, + file_statistics_cache_limit: Some("20M".to_owned()), + } + .into_config_entries() } /// Generate documentation that can be included in the user guide diff --git a/datafusion/execution/src/spill_file.rs b/datafusion/execution/src/spill_file.rs new file mode 100644 index 0000000000000..dca5da23f53e1 --- /dev/null +++ b/datafusion/execution/src/spill_file.rs @@ -0,0 +1,52 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use bytes::Bytes; +use datafusion_common::Result; +use futures::Stream; +use std::path::Path; +use std::pin::Pin; +use std::sync::Arc; + +/// Abstraction over a spill file backend. +/// Implementations handle their own quota enforcement and blocking concerns. +pub trait SpillFile: Send + Sync { + /// Returns the OS path if this is a local file, None otherwise. + fn path(&self) -> Option<&Path> { + None + } + + /// Returns current size in bytes if cheaply available. + fn size(&self) -> Option; + + /// Returns file contents as an async stream of byte chunks. + fn read_stream(&self) -> Result> + Send>>>; + + /// Opens a writer for appending data to this file. + fn open_writer(&self) -> Result>; +} + +/// Writer for spill file backends. +pub trait SpillWriter: std::io::Write + Send { + /// Intended for close/sync/commit operations. + fn finish(&mut self) -> Result<()>; +} + +/// Factory for creating spill files. +pub trait TempFileFactory: Send + Sync { + fn create_temp_file(&self, description: &str) -> Result>; +} diff --git a/datafusion/execution/src/task.rs b/datafusion/execution/src/task.rs index 0de0c937f2211..1c1a717d19c79 100644 --- a/datafusion/execution/src/task.rs +++ b/datafusion/execution/src/task.rs @@ -52,14 +52,14 @@ use std::{collections::HashMap, sync::Arc}; pub struct TaskContext { /// Session Id session_id: String, - /// Optional Task Identify + /// Optional task identity task_id: Option, /// Session configuration session_config: SessionConfig, /// Scalar functions associated with this task context scalar_functions: HashMap>, /// Higher order functions associated with this task context - higher_order_functions: HashMap>, + higher_order_functions: HashMap>, /// Aggregate functions associated with this task context aggregate_functions: HashMap>, /// Window functions associated with this task context @@ -98,7 +98,7 @@ impl TaskContext { session_id: String, session_config: SessionConfig, scalar_functions: HashMap>, - higher_order_functions: HashMap>, + higher_order_functions: HashMap>, aggregate_functions: HashMap>, window_functions: HashMap>, runtime: Arc, @@ -144,7 +144,7 @@ impl TaskContext { &self.scalar_functions } - pub fn higher_order_functions(&self) -> &HashMap> { + pub fn higher_order_functions(&self) -> &HashMap> { &self.higher_order_functions } @@ -167,6 +167,12 @@ impl TaskContext { self.runtime = runtime; self } + + /// Update the `task_id` + pub fn with_task_id(mut self, task_id: String) -> Self { + self.task_id = Some(task_id); + self + } } impl FunctionRegistry for TaskContext { @@ -182,7 +188,7 @@ impl FunctionRegistry for TaskContext { }) } - fn higher_order_function(&self, name: &str) -> Result> { + fn higher_order_function(&self, name: &str) -> Result> { let result = self.higher_order_functions.get(name); result.cloned().ok_or_else(|| { @@ -236,8 +242,8 @@ impl FunctionRegistry for TaskContext { fn register_higher_order_function( &mut self, - function: Arc, - ) -> Result>> { + function: Arc, + ) -> Result>> { function.aliases().iter().for_each(|alias| { self.higher_order_functions .insert(alias.clone(), Arc::clone(&function)); diff --git a/datafusion/expr-common/src/accumulator.rs b/datafusion/expr-common/src/accumulator.rs index 59fb6a595206a..7e9a4ae525ea3 100644 --- a/datafusion/expr-common/src/accumulator.rs +++ b/datafusion/expr-common/src/accumulator.rs @@ -92,6 +92,8 @@ pub trait Accumulator: Send + Sync + Debug + std::any::Any { /// /// "Allocated" means that for internal containers such as `Vec`, /// the `capacity` should be used not the `len`. + /// + /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; /// Returns the intermediate state of the accumulator, consuming the diff --git a/datafusion/expr-common/src/casts.rs b/datafusion/expr-common/src/casts.rs index dad589e4bfe9f..3518c02772672 100644 --- a/datafusion/expr-common/src/casts.rs +++ b/datafusion/expr-common/src/casts.rs @@ -28,10 +28,36 @@ use arrow::datatypes::{ MAX_DECIMAL128_FOR_EACH_PRECISION, MIN_DECIMAL32_FOR_EACH_PRECISION, MIN_DECIMAL64_FOR_EACH_PRECISION, MIN_DECIMAL128_FOR_EACH_PRECISION, TimeUnit, }; -use arrow::temporal_conversions::{MICROSECONDS, MILLISECONDS, NANOSECONDS}; +use arrow::temporal_conversions::{ + MICROSECONDS, MILLISECONDS, MILLISECONDS_IN_DAY, NANOSECONDS, +}; use datafusion_common::ScalarValue; -/// Convert a literal value from one data type to another +/// Convert a literal [`ScalarValue`] to `target_type`, preserving the exact value. +/// +/// Returns `None` if the value cannot be represented in `target_type` +/// *exactly*. +/// +/// This is a restricted, value-preserving cast used to rewrite comparison +/// predicates of the form `CAST(col AS target_type) literal` into +/// `col try_cast_literal_to_type(literal, col_type)`. That rewrite is +/// only valid when the cast cannot change the comparison result. +/// +/// # Supported Casts +/// * numeric → numeric, including integers, decimals, `Date32`/`Date64` and +/// `Timestamp`s, rejecting values outside the target's range or that would +/// lose decimal digits +/// * string → string between `Utf8`, `LargeUtf8` and `Utf8View` +/// * wrapping a value into, or unwrapping it out of, a `Dictionary` whose value +/// type matches the literal's type +/// * `Binary` → `FixedSizeBinary` of the matching length +/// * `Timestamp` → `Timestamp` cast between different time units is allowed even +/// though it can truncate (for example nanoseconds → seconds), and a unit +/// conversion that overflows yields a `NULL` literal rather than `None`. +/// +/// # See Also +/// - [`ScalarValue::cast_to`]: a general-purpose cast that can lose information +/// or change a value's meaning. pub fn try_cast_literal_to_type( lit_value: &ScalarValue, target_type: &DataType, @@ -74,11 +100,72 @@ fn is_date_type(data_type: &DataType) -> bool { /// For example, `CAST(ts AS DATE) = DATE '2024-01-01'` means "any timestamp /// during that day", but unwrapping it to `ts = TIMESTAMP '2024-01-01 /// 00:00:00'` matches only midnight. +/// +/// An identity cast (`from_type == to_type`, e.g. `Date32 -> Date32`) never +/// changes comparison semantics and is therefore not lossy. +/// +/// A cast between the two date types (`Date32` <-> `Date64`) is not pre-filtered +/// as lossy here, because whether it loses information is a per-value question +/// rather than a per-type one. `Date32` -> `Date64` is always exact (a day scaled +/// to midnight in milliseconds). `Date64` -> `Date32` is exact only when the value +/// lands on a day boundary: Arrow nominally defines `Date64` as whole days encoded +/// in milliseconds, but arrow-rs does not enforce that (see arrow-rs#5288), so a +/// `Date64` carrying sub-day milliseconds would lose them. This is not a licence to +/// drop them - [`try_cast_numeric_literal`] returns `None` for a `Date64` value not +/// divisible by 86_400_000, so an inexact `Date64` -> `Date32` fold never happens. fn is_lossy_temporal_cast(from_type: &DataType, to_type: &DataType) -> bool { + if from_type == to_type { + return false; + } + if is_date_type(from_type) && is_date_type(to_type) { + return false; + } (is_date_type(from_type) && to_type.is_temporal()) || (is_date_type(to_type) && from_type.is_temporal()) } +/// Returns true when casting a timestamp from `from_type` to `to_type` loses +/// timestamp precision. +/// +/// This is used by comparison cast unwrapping to avoid rewrites such as +/// `CAST(ts_ns AS timestamp(ms)) = lit_ms` -> `ts_ns = lit_ns`. The original +/// predicate can match any nanosecond value in the same millisecond, while the +/// rewritten predicate only matches the exact millisecond boundary. +pub fn is_timestamp_precision_narrowing_cast( + from_type: &DataType, + to_type: &DataType, +) -> bool { + let (DataType::Timestamp(from_unit, _), DataType::Timestamp(to_unit, _)) = + (from_type, to_type) + else { + return false; + }; + + timestamp_unit_scale(from_unit) > timestamp_unit_scale(to_unit) +} + +/// Returns true when casting a date column from `from_type` to `to_type` narrows +/// `Date64` (milliseconds) to `Date32` (days). +/// +/// Like [`is_timestamp_precision_narrowing_cast`], this guards comparison cast +/// unwrapping against a many-to-one column cast. `CAST(date64 AS Date32) = lit_day` +/// matches any millisecond within that day, but the rewritten `date64 = lit_ms` +/// matches only midnight. Arrow does not require `Date64` values to be whole days +/// (see arrow-rs#5288), so the column may carry sub-day values the planner cannot +/// see; the widening direction (`Date32 -> Date64`) is injective and stays allowed. +pub fn is_date_narrowing_cast(from_type: &DataType, to_type: &DataType) -> bool { + matches!((from_type, to_type), (DataType::Date64, DataType::Date32)) +} + +fn timestamp_unit_scale(unit: &TimeUnit) -> i128 { + match unit { + TimeUnit::Second => 1, + TimeUnit::Millisecond => MILLISECONDS as i128, + TimeUnit::Microsecond => MICROSECONDS as i128, + TimeUnit::Nanosecond => NANOSECONDS as i128, + } +} + /// Returns true if unwrap_cast_in_comparison supports this numeric type fn is_supported_numeric_type(data_type: &DataType) -> bool { matches!( @@ -118,6 +205,36 @@ fn is_supported_binary_type(data_type: &DataType) -> bool { matches!(data_type, DataType::Binary | DataType::FixedSizeBinary(_)) } +/// Scale a `Date32`/`Date64` literal value into the units of `target_type`, +/// returning `None` when the conversion is not exact. +/// +/// `Date32` counts **days** since the Unix epoch while `Date64` counts +/// **milliseconds** since the Unix epoch, so a cross conversion scales by +/// [`MILLISECONDS_IN_DAY`]: +/// * `Date32` -> `Date64` is always exact: `days * MILLISECONDS_IN_DAY` +/// (guarded against `i64`/`i128` overflow). +/// * `Date64` -> `Date32` is exact only when the millisecond value lands on a +/// whole-day boundary; otherwise it returns `None` so the cast unwrap is +/// skipped (correct for every operator, including `=`). +/// +/// For a same-type date cast or a date/integer cast the generic `mul` +/// multiplier already applies, so this returns `value * mul`. +fn scale_date_literal( + value: i128, + from_type: &DataType, + target_type: &DataType, + mul: i128, +) -> Option { + const MILLIS_PER_DAY: i128 = MILLISECONDS_IN_DAY as i128; + match (from_type, target_type) { + (DataType::Date32, DataType::Date64) => value.checked_mul(MILLIS_PER_DAY), + (DataType::Date64, DataType::Date32) => { + (value % MILLIS_PER_DAY == 0).then_some(value / MILLIS_PER_DAY) + } + _ => value.checked_mul(mul), + } +} + /// Convert a numeric value from one numeric data type to another fn try_cast_numeric_literal( lit_value: &ScalarValue, @@ -193,8 +310,12 @@ fn try_cast_numeric_literal( ScalarValue::UInt16(Some(v)) => (*v as i128).checked_mul(mul), ScalarValue::UInt32(Some(v)) => (*v as i128).checked_mul(mul), ScalarValue::UInt64(Some(v)) => (*v as i128).checked_mul(mul), - ScalarValue::Date32(Some(v)) => (*v as i128).checked_mul(mul), - ScalarValue::Date64(Some(v)) => (*v as i128).checked_mul(mul), + ScalarValue::Date32(Some(v)) => { + scale_date_literal(*v as i128, &lit_data_type, target_type, mul) + } + ScalarValue::Date64(Some(v)) => { + scale_date_literal(*v as i128, &lit_data_type, target_type, mul) + } ScalarValue::TimestampSecond(Some(v), _) => (*v as i128).checked_mul(mul), ScalarValue::TimestampMillisecond(Some(v), _) => (*v as i128).checked_mul(mul), ScalarValue::TimestampMicrosecond(Some(v), _) => (*v as i128).checked_mul(mul), @@ -760,6 +881,224 @@ mod tests { ); } + #[test] + fn test_try_cast_identity_date_allowed() { + // An identity Date cast (e.g. `CAST(date_col AS DATE)` where the column + // is already Date32) must fold: it never changes comparison semantics, + // so `try_cast_literal_to_type` should return the same value rather than + // treating it as a lossy temporal cast. + expect_cast( + ScalarValue::Date32(Some(19_723)), + DataType::Date32, + ExpectedCast::Value(ScalarValue::Date32(Some(19_723))), + ); + + expect_cast( + ScalarValue::Date64(Some(1_704_067_200_000)), + DataType::Date64, + ExpectedCast::Value(ScalarValue::Date64(Some(1_704_067_200_000))), + ); + + // is_lossy_temporal_cast must classify an identity cast as non-lossy. + assert!(!is_lossy_temporal_cast( + &DataType::Date32, + &DataType::Date32 + )); + assert!(!is_lossy_temporal_cast( + &DataType::Date64, + &DataType::Date64 + )); + } + + #[test] + fn test_try_cast_between_date32_and_date64() { + // 2025-01-01 is day 20089 since the Unix epoch, which is + // 20089 * 86_400_000 = 1_735_689_600_000 milliseconds. + const DAY_2025_01_01: i32 = 20089; + const MS_2025_01_01: i64 = 1_735_689_600_000; + assert_eq!(DAY_2025_01_01 as i64 * MILLISECONDS_IN_DAY, MS_2025_01_01); + + // Date32 -> Date64 is always exact (days scaled up to milliseconds). + expect_cast( + ScalarValue::Date32(Some(DAY_2025_01_01)), + DataType::Date64, + ExpectedCast::Value(ScalarValue::Date64(Some(MS_2025_01_01))), + ); + + // Date64 -> Date32 is exact only on a whole-day boundary. + expect_cast( + ScalarValue::Date64(Some(MS_2025_01_01)), + DataType::Date32, + ExpectedCast::Value(ScalarValue::Date32(Some(DAY_2025_01_01))), + ); + + // A Date64 value that is not on a day boundary cannot be represented as + // a Date32 exactly, so no rewrite is produced. + expect_cast( + ScalarValue::Date64(Some(MS_2025_01_01 + 1)), + DataType::Date32, + ExpectedCast::NoValue, + ); + expect_cast( + ScalarValue::Date64(Some(MS_2025_01_01 - 1)), + DataType::Date32, + ExpectedCast::NoValue, + ); + + // The epoch and negative (pre-epoch) days round-trip exactly. + expect_cast( + ScalarValue::Date32(Some(0)), + DataType::Date64, + ExpectedCast::Value(ScalarValue::Date64(Some(0))), + ); + expect_cast( + ScalarValue::Date32(Some(-1)), + DataType::Date64, + ExpectedCast::Value(ScalarValue::Date64(Some(-MILLISECONDS_IN_DAY))), + ); + expect_cast( + ScalarValue::Date64(Some(-MILLISECONDS_IN_DAY)), + DataType::Date32, + ExpectedCast::Value(ScalarValue::Date32(Some(-1))), + ); + + // Same-type date casts remain identity conversions. + expect_cast( + ScalarValue::Date32(Some(DAY_2025_01_01)), + DataType::Date32, + ExpectedCast::Value(ScalarValue::Date32(Some(DAY_2025_01_01))), + ); + expect_cast( + ScalarValue::Date64(Some(MS_2025_01_01)), + DataType::Date64, + ExpectedCast::Value(ScalarValue::Date64(Some(MS_2025_01_01))), + ); + } + + #[test] + fn test_is_lossy_temporal_cast_date_pairs() { + // Date <-> Date is let through the pre-filter (per-value exactness is + // enforced downstream in try_cast_numeric_literal, not here). + assert!(!is_lossy_temporal_cast( + &DataType::Date32, + &DataType::Date64 + )); + assert!(!is_lossy_temporal_cast( + &DataType::Date64, + &DataType::Date32 + )); + // Identity is not lossy. + assert!(!is_lossy_temporal_cast( + &DataType::Date32, + &DataType::Date32 + )); + // Date <-> Timestamp remains lossy. + let ts = DataType::Timestamp(TimeUnit::Millisecond, None); + assert!(is_lossy_temporal_cast(&DataType::Date32, &ts)); + assert!(is_lossy_temporal_cast(&ts, &DataType::Date32)); + } + + #[test] + fn test_timestamp_precision_narrowing_cast() { + let ts_ns = DataType::Timestamp(TimeUnit::Nanosecond, None); + let ts_us = DataType::Timestamp(TimeUnit::Microsecond, None); + let ts_ms = DataType::Timestamp(TimeUnit::Millisecond, None); + let ts_s = DataType::Timestamp(TimeUnit::Second, None); + + assert!(is_timestamp_precision_narrowing_cast(&ts_ns, &ts_ms)); + assert!(is_timestamp_precision_narrowing_cast(&ts_us, &ts_s)); + assert!(!is_timestamp_precision_narrowing_cast(&ts_ms, &ts_ns)); + assert!(!is_timestamp_precision_narrowing_cast(&ts_ms, &ts_ms)); + assert!(!is_timestamp_precision_narrowing_cast( + &DataType::Int64, + &ts_ms + )); + } + + #[test] + fn test_is_date_narrowing_cast() { + // Only Date64 -> Date32 narrows (ms -> days, many-to-one). + assert!(is_date_narrowing_cast(&DataType::Date64, &DataType::Date32)); + // The widening direction is injective and must not be flagged. + assert!(!is_date_narrowing_cast( + &DataType::Date32, + &DataType::Date64 + )); + // Identity and non-date pairs are not date-narrowing casts. + assert!(!is_date_narrowing_cast( + &DataType::Date32, + &DataType::Date32 + )); + assert!(!is_date_narrowing_cast( + &DataType::Date64, + &DataType::Date64 + )); + assert!(!is_date_narrowing_cast(&DataType::Int64, &DataType::Date32)); + } + + #[test] + fn test_scale_date_literal_exactness_and_overflow() { + const MS_PER_DAY: i128 = MILLISECONDS_IN_DAY as i128; + + // Date32 -> Date64 is always exact: days scaled to midnight milliseconds. + // 2025-01-01 is day 20089 = 1_735_689_600_000 ms. + assert_eq!( + scale_date_literal(20089, &DataType::Date32, &DataType::Date64, 1), + Some(1_735_689_600_000) + ); + assert_eq!( + scale_date_literal(0, &DataType::Date32, &DataType::Date64, 1), + Some(0) + ); + // Negative (pre-epoch) whole day: 1969-12-31 is day -1 = -86_400_000 ms. + assert_eq!( + scale_date_literal(-1, &DataType::Date32, &DataType::Date64, 1), + Some(-86_400_000) + ); + + // Date64 -> Date32 is exact only on a whole-day boundary. + assert_eq!( + scale_date_literal( + 1_735_689_600_000, + &DataType::Date64, + &DataType::Date32, + 1 + ), + Some(20089) + ); + assert_eq!( + scale_date_literal(-86_400_000, &DataType::Date64, &DataType::Date32, 1), + Some(-1) + ); + // Sub-day values are not exactly representable as a Date32, in both the + // positive and the pre-epoch negative direction -> None (no fold). + assert_eq!( + scale_date_literal( + 1_735_732_800_000, + &DataType::Date64, + &DataType::Date32, + 1 + ), + None + ); + assert_eq!( + scale_date_literal(-43_200_000, &DataType::Date64, &DataType::Date32, 1), + None + ); + + // Extremes: a Date32 at i32::MIN / i32::MAX widens with checked i128 + // arithmetic, producing the exact millisecond value without overflow or + // panic. + assert_eq!( + scale_date_literal(i32::MAX as i128, &DataType::Date32, &DataType::Date64, 1), + Some(i32::MAX as i128 * MS_PER_DAY) + ); + assert_eq!( + scale_date_literal(i32::MIN as i128, &DataType::Date32, &DataType::Date64, 1), + Some(i32::MIN as i128 * MS_PER_DAY) + ); + } + #[test] fn test_try_cast_to_type_unsupported() { // int64 to list diff --git a/datafusion/expr-common/src/columnar_value.rs b/datafusion/expr-common/src/columnar_value.rs index bc6b8177ab3cf..ef9192c3569d9 100644 --- a/datafusion/expr-common/src/columnar_value.rs +++ b/datafusion/expr-common/src/columnar_value.rs @@ -18,9 +18,12 @@ //! [`ColumnarValue`] represents the result of evaluating an expression. use arrow::{ - array::{Array, ArrayRef, Date32Array, Date64Array, NullArray}, + array::{ + Array, ArrayRef, Date32Array, Date64Array, NullArray, TimestampMicrosecondArray, + TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, + }, compute::{CastOptions, kernels, max, min}, - datatypes::DataType, + datatypes::{DataType, TimeUnit}, util::pretty::pretty_format_columns, }; use datafusion_common::internal_datafusion_err; @@ -28,7 +31,10 @@ use datafusion_common::{ Result, ScalarValue, format::DEFAULT_CAST_OPTIONS, internal_err, - scalar::{date_to_timestamp_multiplier, ensure_timestamp_in_bounds}, + scalar::{ + date_to_timestamp_multiplier, ensure_timestamp_in_bounds, + timestamp_to_timestamp_multiplier, + }, }; use std::fmt; use std::sync::Arc; @@ -319,7 +325,9 @@ fn cast_array_by_name( ) { datafusion_common::nested_struct::cast_column(array, cast_type, cast_options) } else { - ensure_date_array_timestamp_bounds(array, cast_type)?; + if !cast_options.safe { + ensure_temporal_array_timestamp_bounds(array, cast_type)?; + } Ok(kernels::cast::cast_with_options( array, cast_type, @@ -328,12 +336,14 @@ fn cast_array_by_name( } } -fn ensure_date_array_timestamp_bounds( +fn ensure_temporal_array_timestamp_bounds( array: &ArrayRef, cast_type: &DataType, ) -> Result<()> { let source_type = array.data_type().clone(); - let Some(multiplier) = date_to_timestamp_multiplier(&source_type, cast_type) else { + let Some(multiplier) = date_to_timestamp_multiplier(&source_type, cast_type) + .or_else(|| timestamp_to_timestamp_multiplier(&source_type, cast_type)) + else { return Ok(()); }; @@ -367,7 +377,55 @@ fn ensure_date_array_timestamp_bounds( })?; (min(arr), max(arr)) } - _ => return Ok(()), // Not a date type, nothing to do + DataType::Timestamp(TimeUnit::Second, _) => { + let arr = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!( + "Expected TimestampSecondArray but found {}", + array.data_type() + ) + })?; + (min(arr), max(arr)) + } + DataType::Timestamp(TimeUnit::Millisecond, _) => { + let arr = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!( + "Expected TimestampMillisecondArray but found {}", + array.data_type() + ) + })?; + (min(arr), max(arr)) + } + DataType::Timestamp(TimeUnit::Microsecond, _) => { + let arr = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!( + "Expected TimestampMicrosecondArray but found {}", + array.data_type() + ) + })?; + (min(arr), max(arr)) + } + DataType::Timestamp(TimeUnit::Nanosecond, _) => { + let arr = array + .as_any() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!( + "Expected TimestampNanosecondArray but found {}", + array.data_type() + ) + })?; + (min(arr), max(arr)) + } + _ => return Ok(()), // Not a temporal type that needs checking. }; // Only validate the min and max values instead of all elements @@ -694,4 +752,48 @@ mod tests { "unexpected error: {err}" ); } + + #[test] + fn cast_timestamp_array_to_timestamp_overflow() { + let overflow_value = i64::MAX / 1_000_000_000 + 1; + let array: ArrayRef = + Arc::new(TimestampSecondArray::from(vec![Some(overflow_value)])); + let value = ColumnarValue::Array(array); + let result = + value.cast_to(&DataType::Timestamp(TimeUnit::Nanosecond, None), None); + let err = result.expect_err("expected overflow to be detected"); + assert!( + err.to_string() + .contains("converted value exceeds the representable i64 range"), + "unexpected error: {err}" + ); + } + + #[test] + fn safe_cast_timestamp_array_to_timestamp_overflow_returns_null() { + let overflow_value = i64::MAX / 1_000_000_000 + 1; + let array: ArrayRef = + Arc::new(TimestampSecondArray::from(vec![Some(overflow_value)])); + let value = ColumnarValue::Array(array); + let safe_options = CastOptions { + safe: true, + ..DEFAULT_CAST_OPTIONS + }; + + let casted = value + .cast_to( + &DataType::Timestamp(TimeUnit::Nanosecond, None), + Some(&safe_options), + ) + .expect("expected safe cast to return null"); + + let ColumnarValue::Array(array) = casted else { + panic!("expected array after cast"); + }; + let array = array + .as_any() + .downcast_ref::() + .expect("expected TimestampNanosecondArray"); + assert!(array.is_null(0)); + } } diff --git a/datafusion/expr-common/src/groups_accumulator.rs b/datafusion/expr-common/src/groups_accumulator.rs index 9053f7a8eab9f..5c01418e04ce7 100644 --- a/datafusion/expr-common/src/groups_accumulator.rs +++ b/datafusion/expr-common/src/groups_accumulator.rs @@ -18,7 +18,7 @@ //! Vectorized [`GroupsAccumulator`] use arrow::array::{ArrayRef, BooleanArray}; -use datafusion_common::{Result, not_impl_err}; +use datafusion_common::{Result, utils::split_vec_min_alloc}; /// Describes how many rows should be emitted during grouping. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -45,13 +45,7 @@ impl EmitTo { // Take the entire vector, leave new (empty) vector std::mem::take(v) } - Self::First(n) => { - // get end n+1,.. values into t - let mut t = v.split_off(*n); - // leave n+1,.. in v - std::mem::swap(v, &mut t); - t - } + Self::First(n) => split_vec_min_alloc(v, *n), } } } @@ -189,12 +183,14 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// /// * `values`: arrays produced from previously calling `state` on other accumulators. /// - /// Other arguments are the same as for [`Self::update_batch`]. + /// Other arguments are the same as for [`Self::update_batch`], except that + /// there is no `opt_filter` — aggregate filters are applied during the + /// partial (update) phase, so by the time intermediate states are merged + /// no per-row filtering is needed. fn merge_batch( &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()>; @@ -235,22 +231,66 @@ pub trait GroupsAccumulator: Send + std::any::Any { /// [`Accumulator::state`]: crate::accumulator::Accumulator::state fn convert_to_state( &self, - _values: &[ArrayRef], - _opt_filter: Option<&BooleanArray>, - ) -> Result> { - not_impl_err!("Input batch conversion to state not implemented") - } - - /// Returns `true` if [`Self::convert_to_state`] is implemented to support - /// intermediate aggregate state conversion. - fn supports_convert_to_state(&self) -> bool { - false - } + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result>; /// Amount of memory used to store the state of this accumulator, /// in bytes. /// /// This function is called once per batch, so it should be `O(n)` to /// compute, not `O(num_groups)` + /// + /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; } + +#[cfg(test)] +mod tests { + use super::EmitTo; + + /// When `n` is small relative to `len`, the old `split_off(n) + swap` pattern had + /// two allocation problems: + /// + /// 1. The returned Vec kept the original large backing allocation even though it + /// only contains `n` elements (wasted capacity on a short-lived value). + /// 2. `split_off` allocated a fresh Vec for the `len - n` remaining elements, + /// even though that side is much larger than `n` — the expensive side to + /// allocate. + /// + /// `split_vec_min_alloc` fixes both: when `n * 2 <= len` it uses + /// `drain(0..n).collect()`, allocating only `n` elements for the emitted prefix + /// and keeping the original large backing in the remaining accumulator. + #[test] + fn take_needed_first_small_n_allocates_minimally() { + let mut v: Vec = Vec::with_capacity(128); + v.extend(0..20i32); + let original_capacity = v.capacity(); // 128 + + // n=4, n*2=8 <= len=20 -> drain branch in split_vec_min_alloc + let emitted = EmitTo::First(4).take_needed(&mut v); + + assert_eq!(emitted, vec![0, 1, 2, 3]); + assert_eq!(v, (4..20i32).collect::>()); + + // The emitted prefix must NOT carry the original large allocation. + // Old split_off+swap returned a Vec with capacity=128 for only 4 elements. + assert!( + emitted.capacity() <= 4, + "emitted prefix capacity {} should be ~n=4, not the original {}", + emitted.capacity(), + original_capacity, + ); + + // The remaining accumulator must retain the original large allocation so + // that incoming groups don't immediately force a realloc. + // Old split_off+swap left the remaining vec with a small fresh allocation. + assert_eq!( + v.capacity(), + original_capacity, + "remaining vec capacity {} should equal original {}", + v.capacity(), + original_capacity, + ); + } +} diff --git a/datafusion/expr-common/src/interval_arithmetic.rs b/datafusion/expr-common/src/interval_arithmetic.rs index e2f8198c92845..68541e1e6b32c 100644 --- a/datafusion/expr-common/src/interval_arithmetic.rs +++ b/datafusion/expr-common/src/interval_arithmetic.rs @@ -910,10 +910,16 @@ impl Interval { if data_type.is_integer() || matches!( data_type, - DataType::Date32 | DataType::Date64 | DataType::Timestamp(_, _) + DataType::Date32 + | DataType::Date64 + | DataType::Timestamp(_, _) + | DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) ) { - self.upper.distance(&self.lower).map(|diff| diff as u64) + self.upper.distance_u64(&self.lower) } else if data_type.is_floating() { // Negative numbers are sorted in the reverse order. To // always have a positive difference after the subtraction, @@ -944,7 +950,7 @@ impl Interval { // Cardinality calculations are not implemented for this data type yet: None } - .map(|result| result + 1) + .and_then(|result| result.checked_add(1)) } /// Reflects an [`Interval`] around the point zero. @@ -4158,6 +4164,28 @@ mod tests { )?; assert_eq!(interval.cardinality().unwrap(), 1_000_000_001); + // Decimal types + let interval = Interval::try_new( + ScalarValue::Decimal128(Some(100), 10, 2), + ScalarValue::Decimal128(Some(110), 10, 2), + )?; + assert_eq!(interval.cardinality().unwrap(), 11); + Ok(()) + } + + #[test] + fn test_cardinality_full_integer_range_does_not_overflow() -> Result<()> { + let interval = Interval::try_new( + ScalarValue::Int64(Some(i64::MIN)), + ScalarValue::Int64(Some(i64::MAX)), + )?; + assert_eq!(interval.cardinality(), None); + + let interval = Interval::try_new( + ScalarValue::UInt64(Some(0)), + ScalarValue::UInt64(Some(u64::MAX)), + )?; + assert_eq!(interval.cardinality(), None); Ok(()) } diff --git a/datafusion/expr-common/src/operator.rs b/datafusion/expr-common/src/operator.rs index a078a27f2a302..b15e770802799 100644 --- a/datafusion/expr-common/src/operator.rs +++ b/datafusion/expr-common/src/operator.rs @@ -36,15 +36,15 @@ pub enum Operator { Plus, /// Subtraction Minus, - /// Multiplication operator, like `*` + /// Multiplication Multiply, - /// Division operator, like `/` + /// Division Divide, - /// Remainder operator, like `%` + /// Remainder Modulo, - /// Logical AND, like `&&` + /// Logical AND And, - /// Logical OR, like `||` + /// Logical OR Or, /// `IS DISTINCT FROM` (see [`distinct`]) /// @@ -80,20 +80,20 @@ pub enum Operator { BitwiseShiftRight, /// Bitwise left, like `<<` BitwiseShiftLeft, - /// String concat + /// String concatenation, like `||` StringConcat, /// At arrow, like `@>`. /// /// Currently only supported to be used with lists: /// ```sql - /// select [1,3] <@ [1,2,3] + /// select [1,2,3] @> [1,3] /// ``` AtArrow, /// Arrow at, like `<@`. /// /// Currently only supported to be used with lists: /// ```sql - /// select [1,2,3] @> [1,3] + /// select [1,3] <@ [1,2,3] /// ``` ArrowAt, /// Arrow, like `->`. @@ -120,7 +120,7 @@ pub enum Operator { /// /// Not implemented in DataFusion yet. IntegerDivide, - /// Hash Minis, like `#-` + /// Hash Minus, like `#-` /// /// Not implemented in DataFusion yet. HashMinus, @@ -163,6 +163,10 @@ impl Operator { Operator::ILikeMatch => Some(Operator::NotILikeMatch), Operator::NotLikeMatch => Some(Operator::LikeMatch), Operator::NotILikeMatch => Some(Operator::ILikeMatch), + Operator::RegexMatch => Some(Operator::RegexNotMatch), + Operator::RegexIMatch => Some(Operator::RegexNotIMatch), + Operator::RegexNotMatch => Some(Operator::RegexMatch), + Operator::RegexNotIMatch => Some(Operator::RegexIMatch), Operator::Plus | Operator::Minus | Operator::Multiply @@ -170,10 +174,6 @@ impl Operator { | Operator::Modulo | Operator::And | Operator::Or - | Operator::RegexMatch - | Operator::RegexIMatch - | Operator::RegexNotMatch - | Operator::RegexNotIMatch | Operator::BitwiseAnd | Operator::BitwiseOr | Operator::BitwiseXor @@ -377,7 +377,8 @@ impl Operator { | Operator::Question | Operator::QuestionAnd | Operator::QuestionPipe - | Operator::Colon => true, + | Operator::Colon + | Operator::StringConcat => true, // E.g. `TRUE OR NULL` is `TRUE` Operator::Or @@ -385,11 +386,53 @@ impl Operator { | Operator::And // IS DISTINCT FROM and IS NOT DISTINCT FROM always return a TRUE/FALSE value, never NULL | Operator::IsDistinctFrom - | Operator::IsNotDistinctFrom - // DataFusion string concatenation operator treats NULL as an empty string - | Operator::StringConcat => false, + | Operator::IsNotDistinctFrom => false, } } + + /// Parse an `Operator` from the string name `datafusion-proto` uses on the + /// wire (the `Debug` name of the variant, e.g. `"Eq"`). + /// + /// Returns `None` for names with no binary-operator counterpart. This is + /// the canonical proto-string mapping, shared by `datafusion-proto` + /// (logical plans) and `PhysicalExpr` decoders such as `BinaryExpr`, so the + /// mapping is not duplicated across crates. + pub fn from_proto_name(name: &str) -> Option { + Some(match name { + "And" => Operator::And, + "Or" => Operator::Or, + "Eq" => Operator::Eq, + "NotEq" => Operator::NotEq, + "LtEq" => Operator::LtEq, + "Lt" => Operator::Lt, + "Gt" => Operator::Gt, + "GtEq" => Operator::GtEq, + "Plus" => Operator::Plus, + "Minus" => Operator::Minus, + "Multiply" => Operator::Multiply, + "Divide" => Operator::Divide, + "Modulo" => Operator::Modulo, + "IsDistinctFrom" => Operator::IsDistinctFrom, + "IsNotDistinctFrom" => Operator::IsNotDistinctFrom, + "BitwiseAnd" => Operator::BitwiseAnd, + "BitwiseOr" => Operator::BitwiseOr, + "BitwiseXor" => Operator::BitwiseXor, + "BitwiseShiftLeft" => Operator::BitwiseShiftLeft, + "BitwiseShiftRight" => Operator::BitwiseShiftRight, + "RegexIMatch" => Operator::RegexIMatch, + "RegexMatch" => Operator::RegexMatch, + "RegexNotIMatch" => Operator::RegexNotIMatch, + "RegexNotMatch" => Operator::RegexNotMatch, + "LikeMatch" => Operator::LikeMatch, + "ILikeMatch" => Operator::ILikeMatch, + "NotLikeMatch" => Operator::NotLikeMatch, + "NotILikeMatch" => Operator::NotILikeMatch, + "StringConcat" => Operator::StringConcat, + "AtArrow" => Operator::AtArrow, + "ArrowAt" => Operator::ArrowAt, + _ => return None, + }) + } } impl fmt::Display for Operator { diff --git a/datafusion/expr-common/src/signature.rs b/datafusion/expr-common/src/signature.rs index 3e941f00c2ee3..f0010f0a05014 100644 --- a/datafusion/expr-common/src/signature.rs +++ b/datafusion/expr-common/src/signature.rs @@ -880,11 +880,6 @@ impl TypeSignature { } } - #[deprecated(since = "46.0.0", note = "See get_example_types instead")] - pub fn get_possible_types(&self) -> Vec> { - self.get_example_types() - } - /// Return example acceptable types for this `TypeSignature`' /// /// Returns a `Vec` for each argument to the function @@ -1054,6 +1049,8 @@ pub enum Coercion { Exact { /// The required type for the argument desired_type: TypeSignatureClass, + /// Physical encoding preservation requested by the function. + encoding_preservation: EncodingPreservation, }, /// Coercion that accepts the desired type and can implicitly coerce from other types. @@ -1062,12 +1059,44 @@ pub enum Coercion { desired_type: TypeSignatureClass, /// Rules for implicit coercion from other types implicit_coercion: ImplicitCoercion, + /// Physical encoding preservation requested by the function. + encoding_preservation: EncodingPreservation, }, } +/// Controls whether a [`Coercion`] preserves an argument's physical encoding +/// (e.g. dictionary) instead of materializing it to the coerced value type. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, PartialOrd, Hash)] +pub struct EncodingPreservation { + preserve_dictionary: bool, +} + +impl EncodingPreservation { + /// Preserve dictionary encoding and coerce only the dictionary values. + pub const fn dictionary() -> Self { + Self { + preserve_dictionary: true, + } + } + + /// Preserve dictionary encoding and coerce only the dictionary values. + pub const fn with_dictionary(mut self) -> Self { + self.preserve_dictionary = true; + self + } + + /// Returns whether dictionary encoding should be preserved. + pub const fn preserve_dictionary(self) -> bool { + self.preserve_dictionary + } +} + impl Coercion { pub fn new_exact(desired_type: TypeSignatureClass) -> Self { - Self::Exact { desired_type } + Self::Exact { + desired_type, + encoding_preservation: EncodingPreservation::default(), + } } /// Create a new coercion with implicit coercion rules. @@ -1085,6 +1114,37 @@ impl Coercion { allowed_source_types, default_casted_type, }, + encoding_preservation: EncodingPreservation::default(), + } + } + + pub fn with_encoding_preservation( + mut self, + encoding_preservation: EncodingPreservation, + ) -> Self { + match &mut self { + Coercion::Exact { + encoding_preservation: current, + .. + } + | Coercion::Implicit { + encoding_preservation: current, + .. + } => *current = encoding_preservation, + } + self + } + + pub fn encoding_preservation(&self) -> EncodingPreservation { + match self { + Coercion::Exact { + encoding_preservation, + .. + } + | Coercion::Implicit { + encoding_preservation, + .. + } => *encoding_preservation, } } @@ -1108,7 +1168,7 @@ impl Coercion { pub fn desired_type(&self) -> &TypeSignatureClass { match self { - Coercion::Exact { desired_type } => desired_type, + Coercion::Exact { desired_type, .. } => desired_type, Coercion::Implicit { desired_type, .. } => desired_type, } } @@ -1133,6 +1193,7 @@ impl PartialEq for Coercion { fn eq(&self, other: &Self) -> bool { self.desired_type() == other.desired_type() && self.implicit_coercion() == other.implicit_coercion() + && self.encoding_preservation() == other.encoding_preservation() } } @@ -1140,6 +1201,7 @@ impl Hash for Coercion { fn hash(&self, state: &mut H) { self.desired_type().hash(state); self.implicit_coercion().hash(state); + self.encoding_preservation().hash(state); } } @@ -2183,6 +2245,20 @@ mod tests { assert_snapshot!(implicit_with_multiple_sources, @"Int64"); } + #[test] + fn test_coercion_encoding_preservation_affects_equality() { + assert!(!EncodingPreservation::default().preserve_dictionary()); + let preserve_dictionary = EncodingPreservation::dictionary(); + assert!(preserve_dictionary.preserve_dictionary()); + + let default = Coercion::new_exact(TypeSignatureClass::Native(logical_string())); + let preserving = default + .clone() + .with_encoding_preservation(preserve_dictionary); + + assert_ne!(default, preserving); + } + #[test] fn test_to_string_repr_coercible() { use insta::assert_snapshot; diff --git a/datafusion/expr-common/src/sort_properties.rs b/datafusion/expr-common/src/sort_properties.rs index 5d17a34a96fbc..74d644f79faef 100644 --- a/datafusion/expr-common/src/sort_properties.rs +++ b/datafusion/expr-common/src/sort_properties.rs @@ -140,9 +140,62 @@ pub struct ExprProperties { /// the expression. Used to compute reliable bounds. pub range: Interval, /// Indicates whether the expression preserves lexicographical ordering - /// of its inputs. For example, string concatenation preserves ordering, - /// while addition does not. + /// of its inputs. + /// + /// This is a *non-strict* (monotone) property: inputs advancing in + /// lexicographical order never make the output decrease, but distinct + /// inputs may map to equal outputs (ties). See + /// [`Self::strictly_order_preserving`] for the strict variant and an + /// explanation of the difference. pub preserves_lex_ordering: bool, + /// Indicates whether the expression is strictly order-preserving with + /// respect to its inputs that are `Ordered`: the output is ordered in the + /// same direction, equal outputs can only result from equal values of + /// those inputs (i.e. the mapping is one-to-one), and nulls map to nulls. + /// + /// i.e. setting this to true means that `a.cmp(b) == f(a).cmp(f(b))` + /// + /// # Difference from [`Self::preserves_lex_ordering`] + /// + /// The two properties differ in both their premise and their strictness: + /// + /// - `preserves_lex_ordering` assumes the inputs advance in + /// *lexicographical* order (a later input may decrease whenever an + /// earlier one increases), and only promises a non-decreasing output, + /// allowing distinct inputs to collapse into equal outputs; `floor`, + /// `date_trunc` and narrowing casts do exactly that. + /// - `strictly_order_preserving` assumes every `Ordered` input advances + /// *simultaneously* (component-wise, which is what actually holds when + /// all of them are sorted in the data), and promises a strict output: + /// equal outputs only from equal inputs. + /// + /// For an expression with a single ordered input the premises coincide, + /// and this field is simply the stronger claim: it implies + /// `preserves_lex_ordering`. With multiple ordered inputs, neither + /// implies the other: a lexicographical-ordering-preserving expression + /// need not be strict (distinct inputs may still produce equal outputs), + /// while `a + b` over two ordered, overflow-free inputs is strict but not + /// lexicographical (under the lexicographical premise `b` may decrease + /// while `a` increases, making the sum decrease). + /// + /// The distinction matters for suffix sort keys. Optimizers use this + /// field to substitute a sort key with an expression computed from it: + /// if data is sorted by `[x, y]`, it is also sorted by `[expr(x), y]`. + /// That claim requires `y` to be sorted within each run of equal + /// `expr(x)` values, which only holds if equal outputs imply equal `x` + /// values. With a merely monotone expression such as `floor`, one output + /// run can span several `x` groups, and `y` restarts at each group: + /// + /// ```text + /// sorted by [x, y]: (1.2, 5), (1.8, 1), (2.5, 3) + /// [floor(x), y]: (1, 5), (1, 1), (2, 3) <-- y not sorted within + /// the "1" run + /// ``` + /// + /// Hence a monotone expression only justifies the length-1 ordering + /// `[expr(x)]`, while a strictly order-preserving one keeps the entire + /// suffix valid. When in doubt, set to `false`. + pub strictly_order_preserving: bool, } impl ExprProperties { @@ -153,6 +206,7 @@ impl ExprProperties { sort_properties: SortProperties::default(), range: Interval::make_unbounded(&DataType::Null).unwrap(), preserves_lex_ordering: false, + strictly_order_preserving: false, } } @@ -173,4 +227,14 @@ impl ExprProperties { self.preserves_lex_ordering = preserves_lex_ordering; self } + + /// Sets whether the expression is strictly order-preserving and returns + /// the modified instance. + pub fn with_strictly_order_preserving( + mut self, + strictly_order_preserving: bool, + ) -> Self { + self.strictly_order_preserving = strictly_order_preserving; + self + } } diff --git a/datafusion/expr-common/src/statistics.rs b/datafusion/expr-common/src/statistics.rs index c94c181615aed..034358b043135 100644 --- a/datafusion/expr-common/src/statistics.rs +++ b/datafusion/expr-common/src/statistics.rs @@ -1694,3 +1694,47 @@ mod tests { all_ops.into_iter().collect() } } + +use std::sync::Arc; + +use datafusion_common::Column; + +/// A statistic a caller would like a provider to supply, if it can do so +/// cheaply. +/// +/// A small, query-aware extension to the existing `Statistics` model: instead +/// of "give me everything you have for every column", a caller can ask for a +/// specific list of stats by name. `StatisticsRequest` is just that vocabulary +/// — DataFusion itself does not populate or consume it. It exists so a request +/// can be threaded from a `TableScan` (see `TableScan::statistics_requests`) +/// through `ScanArgs::statistics_requests` to a `TableProvider`, which is enough +/// for a query-aware statistics feature to be implemented outside of DataFusion. +/// +/// Each variant maps onto a field of [`datafusion_common::Statistics`] / +/// [`datafusion_common::ColumnStatistics`], so a provider that already +/// populates one can answer the request trivially. +/// +/// The per-column variants hold an `Arc` rather than an owned +/// [`Column`] (which carries owned strings) so cloning a request — and the +/// `BTreeSet` stored on `TableScan`, which is cloned with +/// the plan during optimization — stays cheap. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)] +pub enum StatisticsRequest { + /// Smallest non-null value of `column`. + Min(Arc), + /// Largest non-null value of `column`. + Max(Arc), + /// Number of NULLs in `column`. + NullCount(Arc), + /// Number of distinct values in `column` (exact or estimated). + DistinctCount(Arc), + /// Sum of values in `column` (numerics, widened per + /// `ColumnStatistics::sum_value`). + Sum(Arc), + /// Encoded/output byte size of `column`. + ByteSize(Arc), + /// Number of rows in the container (table / file). + RowCount, + /// Total byte size of the container's output. + TotalByteSize, +} diff --git a/datafusion/expr-common/src/type_coercion/binary.rs b/datafusion/expr-common/src/type_coercion/binary.rs index aec87ec5ff853..77ef1f59f7bb8 100644 --- a/datafusion/expr-common/src/type_coercion/binary.rs +++ b/datafusion/expr-common/src/type_coercion/binary.rs @@ -237,7 +237,7 @@ impl<'a> BinaryTypeCoercer<'a> { }) } StringConcat => { - string_concat_coercion(lhs, rhs).map(Signature::uniform).ok_or_else(|| { + string_concat_coercion(lhs, rhs).ok_or_else(|| { plan_datafusion_err!( "Cannot infer common string type for string concat operation {} {} {}", self.lhs, self.op, self.rhs ) @@ -267,6 +267,23 @@ impl<'a> BinaryTypeCoercer<'a> { ret: Int64, }); } + Plus | Minus if is_time_interval_arithmetic(lhs, rhs, self.op) => { + // `time ± interval` yields a `time` wrapped within the 24-hour clock, + // matching PostgreSQL and DuckDB (e.g. `time '23:30' + interval '2 hours'` + // is `01:30:00`). The interval is normalized to `MonthDayNano`; the time + // operand keeps its own unit and is also the result type -- mirroring + // `timestamp/date + interval`, which preserve their unit and apply the + // interval at that resolution. So, like `timestamp(s) + interval + // '1 nanosecond'`, `time(s) + interval '1 nanosecond'` is a no-op rather + // than widening the type. + let (lhs, rhs, ret) = match (lhs, rhs) { + (Interval(_), time) => { + (Interval(MonthDayNano), time.clone(), time.clone()) + } + (time, _) => (time.clone(), Interval(MonthDayNano), time.clone()), + }; + return Ok(Signature { lhs, rhs, ret }); + } Plus | Minus | Multiply | Divide | Modulo => { if let Ok(ret) = self.get_result(lhs, rhs) { @@ -362,6 +379,23 @@ fn is_date_minus_date(lhs: &DataType, rhs: &DataType) -> bool { ) } +/// Returns true for `time + interval`, `interval + time`, or `time - interval`. +/// +/// These follow PostgreSQL/DuckDB semantics where the result is a `time` value +/// wrapped within the 24-hour clock, rather than being widened to an interval. +fn is_time_interval_arithmetic(lhs: &DataType, rhs: &DataType, op: &Operator) -> bool { + use DataType::{Interval, Time32, Time64}; + match op { + Operator::Plus => matches!( + (lhs, rhs), + (Time32(_) | Time64(_), Interval(_)) | (Interval(_), Time32(_) | Time64(_)) + ), + // `interval - time` is not meaningful, so only `time - interval` is accepted. + Operator::Minus => matches!((lhs, rhs), (Time32(_) | Time64(_), Interval(_))), + _ => false, + } +} + /// Coercion rules for mathematics operators between decimal and non-decimal types. fn math_decimal_coercion( lhs_type: &DataType, @@ -654,11 +688,8 @@ pub fn type_union_resolution(data_types: &[DataType]) -> Option { // For example, // i64 and decimal(7, 2) are expect to get coerced type decimal(22, 2) // numeric string ('1') and numeric (2) are expect to get coerced type numeric (1, 2) - if let Some(t) = type_union_resolution_coercion(data_type, candidate_t) { - candidate_type = Some(t); - } else { - return None; - } + let t = type_union_resolution_coercion(data_type, candidate_t)?; + candidate_type = Some(t); } else { candidate_type = Some(data_type.clone()); } @@ -743,14 +774,11 @@ fn type_union_resolution_coercion( ) -> Option { for rhs_field in rhs.iter() { if lhs_field.name() == rhs_field.name() { - if let Some(t) = type_union_resolution_coercion( + let t = type_union_resolution_coercion( lhs_field.data_type(), rhs_field.data_type(), - ) { - return Some(t); - } else { - return None; - } + )?; + return Some(t); } } @@ -774,6 +802,7 @@ fn type_union_resolution_coercion( } _ => binary_numeric_coercion(lhs_type, rhs_type) .or_else(|| list_coercion(lhs_type, rhs_type, type_union_resolution_coercion)) + .or_else(|| map_coercion(lhs_type, rhs_type, type_union_resolution_coercion)) .or_else(|| temporal_coercion_nonstrict_timezone(lhs_type, rhs_type)) .or_else(|| string_coercion(lhs_type, rhs_type)) .or_else(|| null_coercion(lhs_type, rhs_type)) @@ -938,6 +967,7 @@ pub fn comparison_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option T)` extracts +/// values from the matching variant; rows whose active variant cannot be +/// cast to `T` become NULL. +/// +/// Identical union types are already handled by the `equals_datatype` fast path +/// in [`comparison_coercion`]; coercing between two different union types is not +/// supported. +fn union_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option { + use arrow::datatypes::DataType::*; + + match (lhs_type, rhs_type) { + (Union(fields, _), opaque) | (opaque, Union(fields, _)) => fields + .iter() + .any(|(_, f)| can_cast_types(f.data_type(), opaque)) + .then(|| opaque.clone()), + _ => None, + } +} + /// Returns the output type of applying mathematics operations such as /// `+` to arguments of `lhs_type` and `rhs_type`. fn mathematics_numerical_coercion( @@ -1612,50 +1664,55 @@ fn ree_coercion( /// 1. At least one side of lhs and rhs should be string type (Utf8 / LargeUtf8) /// 2. Data type of the other side should be able to cast to string type /// 3. Binary and string types cannot be mixed -fn string_concat_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option { +fn string_concat_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option { use arrow::datatypes::DataType::*; - string_coercion(lhs_type, rhs_type).or_else(|| match (lhs_type, rhs_type) { - // Allow pure binary + binary - ( - Binary | LargeBinary | BinaryView | FixedSizeBinary(_), - Binary | LargeBinary | BinaryView | FixedSizeBinary(_), - ) => { - // Coerce fixed-sized binary to variable-sized `Binary` to make uniform signature - // with the `Binary` result - let lhs_type = match lhs_type { - FixedSizeBinary(_) => &Binary, - val => val, - }; - let rhs_type = match rhs_type { - FixedSizeBinary(_) => &Binary, - val => val, - }; - binary_coercion(lhs_type, rhs_type) - } - // Deny other mixed binary + string combinations - ( - Binary | LargeBinary | BinaryView | FixedSizeBinary(_), - Utf8 | LargeUtf8 | Utf8View, - ) => None, - ( - Utf8 | LargeUtf8 | Utf8View, - Binary | LargeBinary | BinaryView | FixedSizeBinary(_), - ) => None, - // Predicate-based coercion rules are following - (Utf8View, from_type) | (from_type, Utf8View) => { - string_concat_internal_coercion(from_type, &Utf8View) - } - (Utf8, from_type) | (from_type, Utf8) => { - string_concat_internal_coercion(from_type, &Utf8) - } - (LargeUtf8, from_type) | (from_type, LargeUtf8) => { - string_concat_internal_coercion(from_type, &LargeUtf8) - } - (Dictionary(_, lhs_value_type), Dictionary(_, rhs_value_type)) => { - string_coercion(lhs_value_type, rhs_value_type).or(None) - } - _ => None, - }) + + string_coercion(lhs_type, rhs_type) + .map(Signature::uniform) + .or_else(|| match (lhs_type, rhs_type) { + // Allow concatenation of mixed fixed size binary + (FixedSizeBinary(l), FixedSizeBinary(r)) => Some(Signature { + lhs: lhs_type.clone(), + rhs: rhs_type.clone(), + ret: FixedSizeBinary(l + r), + }), + // Allow pure binary + binary + ( + Binary | LargeBinary | BinaryView | FixedSizeBinary(_), + Binary | LargeBinary | BinaryView | FixedSizeBinary(_), + ) => { + // Coerce fixed-sized binary to variable-sized `Binary` to make uniform signature + // with the `Binary` result + let lhs_type = match lhs_type { + FixedSizeBinary(_) => &Binary, + val => val, + }; + let rhs_type = match rhs_type { + FixedSizeBinary(_) => &Binary, + val => val, + }; + binary_coercion(lhs_type, rhs_type).map(Signature::uniform) + } + // Predicate-based coercion rules are following, + // including mixed binary + string combinations + (Utf8View, from_type) | (from_type, Utf8View) => { + string_concat_internal_coercion(from_type, &Utf8View) + .map(Signature::uniform) + } + (Utf8, from_type) | (from_type, Utf8) => { + string_concat_internal_coercion(from_type, &Utf8).map(Signature::uniform) + } + (LargeUtf8, from_type) | (from_type, LargeUtf8) => { + string_concat_internal_coercion(from_type, &LargeUtf8) + .map(Signature::uniform) + } + (Dictionary(_, lhs_value_type), Dictionary(_, rhs_value_type)) => { + string_coercion(lhs_value_type, rhs_value_type) + .or(None) + .map(Signature::uniform) + } + _ => None, + }) } fn array_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option { @@ -2048,22 +2105,10 @@ fn temporal_coercion(lhs_type: &DataType, rhs_type: &DataType) -> Option TimeUnit { use arrow::datatypes::TimeUnit::*; match (lhs_unit, rhs_unit) { - (Second, Millisecond) => Second, - (Second, Microsecond) => Second, - (Second, Nanosecond) => Second, - (Millisecond, Second) => Second, - (Millisecond, Microsecond) => Millisecond, - (Millisecond, Nanosecond) => Millisecond, - (Microsecond, Second) => Second, - (Microsecond, Millisecond) => Millisecond, - (Microsecond, Nanosecond) => Microsecond, - (Nanosecond, Second) => Second, - (Nanosecond, Millisecond) => Millisecond, - (Nanosecond, Microsecond) => Microsecond, - (l, r) => { - assert_eq!(l, r); - *l - } + (Second, Second) => Second, + (Nanosecond, _) | (_, Nanosecond) => Nanosecond, + (Microsecond, _) | (_, Microsecond) => Microsecond, + (Millisecond, _) | (_, Millisecond) => Millisecond, } } diff --git a/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs b/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs index eb5622fedb8aa..70a8fc0e35a15 100644 --- a/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs +++ b/datafusion/expr-common/src/type_coercion/binary/tests/arithmetic.rs @@ -40,8 +40,8 @@ fn test_date_timestamp_arithmetic_error() -> Result<()> { &DataType::Timestamp(Millisecond, None), ) .get_input_types()?; - assert_eq!(lhs, DataType::Timestamp(Millisecond, None)); - assert_eq!(rhs, DataType::Timestamp(Millisecond, None)); + assert_eq!(lhs, DataType::Timestamp(Nanosecond, None)); + assert_eq!(rhs, DataType::Timestamp(Nanosecond, None)); let err = BinaryTypeCoercer::new(&DataType::Date32, &Operator::Plus, &DataType::Date64) diff --git a/datafusion/expr-common/src/type_coercion/binary/tests/comparison.rs b/datafusion/expr-common/src/type_coercion/binary/tests/comparison.rs index f8bff3ca90ecf..cfa3bbe189929 100644 --- a/datafusion/expr-common/src/type_coercion/binary/tests/comparison.rs +++ b/datafusion/expr-common/src/type_coercion/binary/tests/comparison.rs @@ -575,6 +575,24 @@ fn test_type_coercion_compare() -> Result<()> { Operator::Eq, DataType::Timestamp(Second, Some("Europe/Brussels".into())) ); + test_coercion_binary_rule!( + DataType::Timestamp(Second, None), + DataType::Timestamp(Millisecond, None), + Operator::Eq, + DataType::Timestamp(Millisecond, None) + ); + test_coercion_binary_rule!( + DataType::Timestamp(Second, Some("America/New_York".into())), + DataType::Timestamp(Nanosecond, Some("Europe/Brussels".into())), + Operator::Lt, + DataType::Timestamp(Nanosecond, Some("America/New_York".into())) + ); + test_coercion_binary_rule!( + DataType::Timestamp(Microsecond, None), + DataType::Timestamp(Nanosecond, None), + Operator::GtEq, + DataType::Timestamp(Nanosecond, None) + ); // list let inner_field = Arc::new(Field::new_list_field(DataType::Int64, true)); @@ -872,6 +890,78 @@ fn test_type_union_coercion_prefers_string() { ); } +#[test] +fn test_type_union_coercion_prefers_finer_timestamp_unit() { + assert_eq!( + type_union_coercion( + &DataType::Timestamp(Second, None), + &DataType::Timestamp(Millisecond, None), + ), + Some(DataType::Timestamp(Millisecond, None)) + ); + assert_eq!( + type_union_resolution(&[ + DataType::Timestamp(Second, None), + DataType::Timestamp(Nanosecond, None), + ]), + Some(DataType::Timestamp(Nanosecond, None)) + ); +} + +/// Tests that `type_union_resolution` unifies Map types by recursing into the +/// key/value types, so a Map whose value type is Null (e.g. `MAP {'k': NULL}`) +/// unifies with a concretely-typed Map in a VALUES list. +/// See . +#[test] +fn test_type_union_resolution_map() { + fn map_type(value_type: DataType) -> DataType { + DataType::Map( + Arc::new(Field::new( + "entries", + DataType::Struct(Fields::from(vec![ + Field::new("key", DataType::Utf8, false), + Field::new("value", value_type, true), + ])), + false, + )), + false, + ) + } + + // Null value type unifies with a concrete value type, in both orders + assert_eq!( + type_union_resolution(&[map_type(DataType::Int64), map_type(DataType::Null)]), + Some(map_type(DataType::Int64)) + ); + assert_eq!( + type_union_resolution(&[map_type(DataType::Null), map_type(DataType::Int64)]), + Some(map_type(DataType::Int64)) + ); + + // Numeric value types widen following the scalar rules + assert_eq!( + type_union_resolution(&[ + map_type(DataType::Int64), + map_type(DataType::Null), + map_type(DataType::Float64), + ]), + Some(map_type(DataType::Float64)) + ); + + // Map cannot unify with a non-Map composite type + assert_eq!( + type_union_resolution(&[ + map_type(DataType::Int64), + DataType::Struct(Fields::from(vec![Field::new( + "key", + DataType::Utf8, + false + )])), + ]), + None + ); +} + /// Tests that comparison operators coerce to numeric when comparing /// numeric and string types. #[test] @@ -935,7 +1025,8 @@ fn test_string_concat_coercion() -> Result<()> { DataType::FixedSizeBinary(4), DataType::FixedSizeBinary(16), Operator::StringConcat, - DataType::Binary + DataType::FixedSizeBinary(4), + DataType::FixedSizeBinary(16) ); test_coercion_binary_rule!( DataType::FixedSizeBinary(4), @@ -976,19 +1067,18 @@ fn test_string_concat_coercion() -> Result<()> { DataType::Binary, DataType::LargeBinary, DataType::BinaryView, - DataType::FixedSizeBinary(8), ] { - assert!( - BinaryTypeCoercer::new(&binary_dt, &Operator::StringConcat, &string_dt,) - .get_input_types() - .is_err(), - "{binary_dt} || {string_dt}" + test_coercion_binary_rule!( + &binary_dt, + &string_dt, + Operator::StringConcat, + string_dt ); - assert!( - BinaryTypeCoercer::new(&string_dt, &Operator::StringConcat, &binary_dt,) - .get_input_types() - .is_err(), - "{string_dt} || {binary_dt}" + test_coercion_binary_rule!( + &string_dt, + &binary_dt, + Operator::StringConcat, + string_dt ); } } diff --git a/datafusion/expr-common/src/type_coercion/binary/tests/mod.rs b/datafusion/expr-common/src/type_coercion/binary/tests/mod.rs index e4653d4955eb0..f771b10d9313f 100644 --- a/datafusion/expr-common/src/type_coercion/binary/tests/mod.rs +++ b/datafusion/expr-common/src/type_coercion/binary/tests/mod.rs @@ -27,12 +27,18 @@ use super::*; /// - op: The binary operator (e.g., "+", "-", etc.) /// - expected_type: The type both sides should be coerced to macro_rules! test_coercion_binary_rule { - ($LHS_TYPE:expr, $RHS_TYPE:expr, $OP:expr, $RESULT_TYPE:expr) => {{ + ($LHS_TYPE:expr, $RHS_TYPE:expr, $OP:expr, $RESULT_TYPE:expr) => { let (lhs, rhs) = BinaryTypeCoercer::new(&$LHS_TYPE, &$OP, &$RHS_TYPE).get_input_types()?; assert_eq!(lhs, $RESULT_TYPE); assert_eq!(rhs, $RESULT_TYPE); - }}; + }; + ($LHS_TYPE:expr, $RHS_TYPE:expr, $OP:expr, $L_RESULT_TYPE:expr, $R_RESULT_TYPE:expr) => { + let (lhs, rhs) = + BinaryTypeCoercer::new(&$LHS_TYPE, &$OP, &$RHS_TYPE).get_input_types()?; + assert_eq!(lhs, $L_RESULT_TYPE); + assert_eq!(rhs, $R_RESULT_TYPE); + }; } /// Tests that coercion for a binary operator between one type and multiple right-hand side types diff --git a/datafusion/expr/Cargo.toml b/datafusion/expr/Cargo.toml index 8cec01feb30b5..4fe7b65f6d05f 100644 --- a/datafusion/expr/Cargo.toml +++ b/datafusion/expr/Cargo.toml @@ -42,6 +42,10 @@ name = "datafusion_expr" [features] default = ["sql"] +# Enables protobuf conversions for the expression types owned by this crate. +# Off by default so consumers that never serialize plans pay nothing. Mirrors +# the `proto` feature on `datafusion-datasource` and friends. +proto = ["dep:datafusion-proto-common", "dep:datafusion-proto-models"] recursive_protection = ["dep:recursive"] sql = ["sqlparser"] @@ -56,6 +60,8 @@ datafusion-expr-common = { workspace = true } datafusion-functions-aggregate-common = { workspace = true } datafusion-functions-window-common = { workspace = true } datafusion-physical-expr-common = { workspace = true } +datafusion-proto-common = { workspace = true, optional = true } +datafusion-proto-models = { workspace = true, optional = true } indexmap = { workspace = true } itertools = { workspace = true } recursive = { workspace = true, optional = true } diff --git a/datafusion/expr/src/execution_props.rs b/datafusion/expr/src/execution_props.rs index 649f74ed3997c..7c5369d1144dd 100644 --- a/datafusion/expr/src/execution_props.rs +++ b/datafusion/expr/src/execution_props.rs @@ -18,14 +18,9 @@ use crate::var_provider::{VarProvider, VarType}; use chrono::{DateTime, Utc}; use datafusion_common::HashMap; -use datafusion_common::ScalarValue; -use datafusion_common::TableReference; use datafusion_common::alias::AliasGenerator; use datafusion_common::config::ConfigOptions; -use datafusion_common::{Result, internal_err}; -use std::fmt; -use std::hash::{Hash, Hasher}; -use std::sync::{Arc, Mutex}; +use std::sync::Arc; /// Holds properties and scratch state used while optimizing a [`LogicalPlan`] /// and translating it into an executable physical plan, such as the statement @@ -64,16 +59,6 @@ pub struct ExecutionProps { pub config_options: Option>, /// Providers for scalar variables pub var_providers: Option>>, - /// Maps each logical `Subquery` to its index in `subquery_results`. - /// Populated by the physical planner before calling `create_physical_expr`. - pub subquery_indexes: HashMap, - /// Shared results container for uncorrelated scalar subquery values. - /// Populated at execution time by `ScalarSubqueryExec`. - pub subquery_results: ScalarSubqueryResults, - /// Maps each lambda variable name to its lambda qualifier generated - /// during physical planning. Populated by the physical planner for - /// each lambda before calling `create_physical_expr`. - pub lambda_variable_qualifier: HashMap, } impl Default for ExecutionProps { @@ -90,9 +75,6 @@ impl ExecutionProps { alias_generator: Arc::new(AliasGenerator::new()), config_options: None, var_providers: None, - subquery_indexes: HashMap::new(), - subquery_results: ScalarSubqueryResults::default(), - lambda_variable_qualifier: HashMap::new(), } } @@ -151,119 +133,6 @@ impl ExecutionProps { pub fn config_options(&self) -> Option<&Arc> { self.config_options.as_ref() } - - /// Adds a mapping for each variable to the given qualifier. Existing - /// variables with conflicting names get's shadowed - pub fn with_qualified_lambda_variables( - mut self, - qualifier: &TableReference, - variables: &[String], - ) -> Self { - for var in variables { - self.lambda_variable_qualifier - .entry_ref(var) - .insert(qualifier.clone()); - } - - self - } -} - -/// Index of a scalar subquery within a [`ScalarSubqueryResults`] container. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] -pub struct SubqueryIndex(usize); - -impl SubqueryIndex { - /// Creates a new subquery index. - pub const fn new(index: usize) -> Self { - Self(index) - } - - /// Returns the underlying slot index. - pub const fn as_usize(self) -> usize { - self.0 - } -} - -/// Shared results container for uncorrelated scalar subqueries. -/// -/// Each entry corresponds to one scalar subquery, identified by its index. -/// Each slot is populated at execution time by `ScalarSubqueryExec`, read by -/// `ScalarSubqueryExpr` instances that share this container, and cleared when -/// the plan is reset for re-execution. -#[derive(Clone, Default)] -pub struct ScalarSubqueryResults { - slots: Arc>>>, -} - -impl ScalarSubqueryResults { - /// Creates a new shared results container with `n` empty slots. - pub fn new(n: usize) -> Self { - Self { - slots: Arc::new((0..n).map(|_| Mutex::new(None)).collect()), - } - } - - /// Returns the scalar value stored at `index`, if it has been populated. - pub fn get(&self, index: SubqueryIndex) -> Option { - let slot = self.slots.get(index.as_usize())?; - slot.lock().unwrap().clone() - } - - /// Stores `value` in the slot at `index`. - pub fn set(&self, index: SubqueryIndex, value: ScalarValue) -> Result<()> { - let Some(slot) = self.slots.get(index.as_usize()) else { - return internal_err!( - "ScalarSubqueryResults: result index {} is out of bounds", - index.as_usize() - ); - }; - - let mut slot = slot.lock().unwrap(); - if slot.is_some() { - return internal_err!( - "ScalarSubqueryResults: result for index {} was already populated", - index.as_usize() - ); - } - *slot = Some(value); - - Ok(()) - } - - /// Clears all populated results so the container can be reused. - pub fn clear(&self) { - for slot in self.slots.iter() { - *slot.lock().unwrap() = None; - } - } - - /// Returns true if `this` and `other` point to the same shared container. - pub fn ptr_eq(this: &Self, other: &Self) -> bool { - Arc::ptr_eq(&this.slots, &other.slots) - } -} - -impl fmt::Debug for ScalarSubqueryResults { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_list() - .entries(self.slots.iter().map(|slot| slot.lock().unwrap().clone())) - .finish() - } -} - -impl PartialEq for ScalarSubqueryResults { - fn eq(&self, other: &Self) -> bool { - Self::ptr_eq(self, other) - } -} - -impl Eq for ScalarSubqueryResults {} - -impl Hash for ScalarSubqueryResults { - fn hash(&self, state: &mut H) { - Arc::as_ptr(&self.slots).hash(state); - } } #[cfg(test)] @@ -274,44 +143,8 @@ mod test { fn debug() { let props = ExecutionProps::new(); assert_eq!( - "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None, subquery_indexes: {}, subquery_results: [], lambda_variable_qualifier: {} }", + "ExecutionProps { query_execution_start_time: None, alias_generator: AliasGenerator { next_id: 1 }, config_options: None, var_providers: None }", format!("{props:?}") ); } - - #[test] - fn scalar_subquery_results_set_and_get() -> Result<()> { - let results = ScalarSubqueryResults::new(1); - assert_eq!(results.get(SubqueryIndex::new(0)), None); - - results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?; - assert_eq!( - results.get(SubqueryIndex::new(0)), - Some(ScalarValue::Int32(Some(42))) - ); - assert!( - results - .set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7))) - .is_err() - ); - - Ok(()) - } - - #[test] - fn scalar_subquery_results_clear() -> Result<()> { - let results = ScalarSubqueryResults::new(1); - results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?; - - results.clear(); - - assert_eq!(results.get(SubqueryIndex::new(0)), None); - results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7)))?; - assert_eq!( - results.get(SubqueryIndex::new(0)), - Some(ScalarValue::Int32(Some(7))) - ); - - Ok(()) - } } diff --git a/datafusion/expr/src/expr.rs b/datafusion/expr/src/expr.rs index d6276b944c334..f9c0662e682e8 100644 --- a/datafusion/expr/src/expr.rs +++ b/datafusion/expr/src/expr.rs @@ -92,7 +92,7 @@ impl From for NullTreatment { /// /// For example the expression `A + 1` will be represented as /// -///```text +/// ```text /// BinaryExpr { /// left: Expr::Column("A"), /// op: Operator::Plus, @@ -265,7 +265,7 @@ impl From for NullTreatment { /// /// [`ExplainFormat::Tree`]: crate::logical_plan::ExplainFormat::Tree /// -///``` +/// ``` /// # use datafusion_expr::{lit, col}; /// let expr = col("c1") + lit(42); /// assert_eq!(format!("{}", expr.human_display()), "c1 + 42"); @@ -301,7 +301,7 @@ impl From for NullTreatment { /// Rewrite an expression, replacing references to column "a" in an /// to the literal `42`: /// -/// ``` +/// ``` /// # use datafusion_common::tree_node::{Transformed, TreeNode}; /// # use datafusion_expr::{col, Expr, lit}; /// // expression a = 5 AND b = 6 @@ -437,14 +437,14 @@ pub enum Expr { #[derive(Clone, Eq, PartialOrd, Debug)] pub struct HigherOrderFunction { /// The function - pub func: Arc, + pub func: Arc, /// List of expressions to feed to the functions as arguments pub args: Vec, } impl HigherOrderFunction { /// Create a new `HigherOrderFunction` from a [`HigherOrderUDF`] - pub fn new(func: Arc, args: Vec) -> Self { + pub fn new(func: Arc, args: Vec) -> Self { Self { func, args } } @@ -452,7 +452,7 @@ impl HigherOrderFunction { self.func.name() } - /// Invokes the inner function [`HigherOrderUDF::lambda_parameters`] + /// Invokes the inner function [`crate::HigherOrderUDFImpl::lambda_parameters`] /// using the arguments of this invocation. This expression lambda /// variables must be already resolved either by coming from the /// default sql planner or by calling [Expr::resolve_lambda_variables] @@ -662,7 +662,7 @@ pub fn intersect_metadata_for_union<'a>( } Some(current) => { // Only keep keys that exist in both with the same value - current.retain(|k, v| metadata.get(k) == Some(v)); + current.retain(|k, v| metadata.get(k) == Some(&*v)); } } } @@ -671,22 +671,43 @@ pub fn intersect_metadata_for_union<'a>( } /// UNNEST expression. +/// +/// When `outer` is `true`, the unnest should preserve `NULL` and empty input +/// lists by emitting a single `NULL` output row for each. When `false` (the +/// historical default), the behavior is identical to the plain `UNNEST(col)` +/// SQL form: `NULL` and empty input lists are dropped from the output. #[derive(Clone, PartialEq, Eq, PartialOrd, Hash, Debug)] pub struct Unnest { pub expr: Box, + /// Outer-unnest behavior: also expand empty input lists into a single + /// `NULL` output row (in addition to preserving `NULL` input rows). + pub outer: bool, } impl Unnest { - /// Create a new Unnest expression. + /// Create a new Unnest expression with default (non-outer) semantics. pub fn new(expr: Expr) -> Self { Self { expr: Box::new(expr), + outer: false, } } - /// Create a new Unnest expression. + /// Create a new Unnest expression with default (non-outer) semantics. pub fn new_boxed(boxed: Box) -> Self { - Self { expr: boxed } + Self { + expr: boxed, + outer: false, + } + } + + /// Create a new Unnest expression with outer-unnest semantics: `NULL` + /// and empty input lists each produce a single `NULL` output row. + pub fn new_outer(expr: Expr) -> Self { + Self { + expr: Box::new(expr), + outer: true, + } } } @@ -2186,6 +2207,29 @@ impl Expr { rewrite_placeholder(item, expr.as_ref(), schema)?; } } + Expr::InSubquery(InSubquery { + expr, + subquery, + negated: _, + }) => { + rewrite_placeholder_from_subquery( + "InSubquery", + expr.as_mut(), + subquery, + )?; + } + Expr::SetComparison(SetComparison { + expr, + subquery, + op: _, + quantifier: _, + }) => { + rewrite_placeholder_from_subquery( + "SetComparison", + expr.as_mut(), + subquery, + )?; + } Expr::Like(Like { expr, pattern, .. }) | Expr::SimilarTo(Like { expr, pattern, .. }) => { rewrite_placeholder(pattern.as_mut(), expr.as_ref(), schema)?; @@ -2260,6 +2304,7 @@ impl Expr { pub fn spans(&self) -> Option<&Spans> { match self { Expr::Column(col) => Some(&col.spans), + Expr::Not(inner) | Expr::Negative(inner) => inner.spans(), _ => None, } } @@ -2407,11 +2452,19 @@ impl NormalizeEq for Expr { | (Expr::IsNotTrue(self_expr), Expr::IsNotTrue(other_expr)) | (Expr::IsNotFalse(self_expr), Expr::IsNotFalse(other_expr)) | (Expr::IsNotUnknown(self_expr), Expr::IsNotUnknown(other_expr)) - | (Expr::Negative(self_expr), Expr::Negative(other_expr)) - | ( - Expr::Unnest(Unnest { expr: self_expr }), - Expr::Unnest(Unnest { expr: other_expr }), - ) => self_expr.normalize_eq(other_expr), + | (Expr::Negative(self_expr), Expr::Negative(other_expr)) => { + self_expr.normalize_eq(other_expr) + } + ( + Expr::Unnest(Unnest { + expr: self_expr, + outer: self_outer, + }), + Expr::Unnest(Unnest { + expr: other_expr, + outer: other_outer, + }), + ) => self_outer == other_outer && self_expr.normalize_eq(other_expr), ( Expr::Between(Between { expr: self_expr, @@ -2859,7 +2912,9 @@ impl HashNode for Expr { field.hash(state); column.hash(state); } - Expr::Unnest(Unnest { expr: _expr }) => {} + Expr::Unnest(Unnest { expr: _expr, outer }) => { + outer.hash(state); + } Expr::HigherOrderFunction(HigherOrderFunction { func, args: _args }) => { func.hash(state); } @@ -2911,6 +2966,26 @@ macro_rules! expr_vec_fmt { .join(", ") }}; } +/// Infer an untyped placeholder on the left of a single-column subquery predicate from the subquery projection +fn rewrite_placeholder_from_subquery( + kind: &str, + expr: &mut Expr, + subquery: &Subquery, +) -> Result<()> { + let subquery_schema = subquery.subquery.schema(); + match &subquery_schema.fields()[..] { + [subquery_field] => { + let column = + Expr::Column(Column::new_unqualified(subquery_field.name().clone())); + rewrite_placeholder(expr, &column, subquery_schema) + } + _ => plan_err!( + "{kind} should only return one column, but found {}: {}", + subquery_schema.fields().len(), + subquery_schema.field_names().join(", ") + ), + } +} struct SchemaDisplay<'a>(&'a Expr); impl Display for SchemaDisplay<'_> { @@ -3079,8 +3154,9 @@ impl Display for SchemaDisplay<'_> { } Expr::Negative(expr) => write!(f, "(- {})", SchemaDisplay(expr)), Expr::Not(expr) => write!(f, "NOT {}", SchemaDisplay(expr)), - Expr::Unnest(Unnest { expr }) => { - write!(f, "UNNEST({})", SchemaDisplay(expr)) + Expr::Unnest(Unnest { expr, outer }) => { + let name = if *outer { "UNNEST_OUTER" } else { "UNNEST" }; + write!(f, "{name}({})", SchemaDisplay(expr)) } Expr::ScalarFunction(ScalarFunction { func, args }) => { match func.schema_name(args) { @@ -3354,8 +3430,9 @@ impl Display for SqlDisplay<'_> { } Expr::Negative(expr) => write!(f, "(- {})", SqlDisplay(expr)), Expr::Not(expr) => write!(f, "NOT {}", SqlDisplay(expr)), - Expr::Unnest(Unnest { expr }) => { - write!(f, "UNNEST({})", SqlDisplay(expr)) + Expr::Unnest(Unnest { expr, outer }) => { + let name = if *outer { "UNNEST_OUTER" } else { "UNNEST" }; + write!(f, "{name}({})", SqlDisplay(expr)) } Expr::SimilarTo(Like { negated, @@ -3708,7 +3785,7 @@ impl Display for Expr { } }, Expr::Placeholder(Placeholder { id, .. }) => write!(f, "{id}"), - Expr::Unnest(Unnest { expr }) => { + Expr::Unnest(Unnest { expr, .. }) => { write!(f, "{UNNEST_COLUMN_PREFIX}({expr})") } Expr::HigherOrderFunction(fun) => { @@ -3817,6 +3894,214 @@ mod test { } } + #[test] + fn infer_placeholder_in_subquery() { + // WHERE $1 IN (SELECT a FROM t) + let subquery_field = Field::new("a", DataType::Int32, false); + let subquery_schema = Arc::new( + DFSchema::from_unqualified_fields( + vec![subquery_field].into(), + Default::default(), + ) + .unwrap(), + ); + let subquery = Subquery { + subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: subquery_schema, + })), + outer_ref_columns: vec![], + spans: Spans::new(), + }; + + let in_subquery = Expr::InSubquery(InSubquery { + expr: Box::new(Expr::Placeholder(Placeholder { + id: "$1".to_string(), + field: None, + })), + subquery, + negated: false, + }); + + let outer_schema = DFSchema::empty(); + let (inferred_expr, contains_placeholder) = + in_subquery.infer_placeholder_types(&outer_schema).unwrap(); + + assert!(contains_placeholder); + + match inferred_expr { + Expr::InSubquery(in_subquery) => match *in_subquery.expr { + Expr::Placeholder(placeholder) => { + let inferred = placeholder.field.expect("placeholder field"); + assert_eq!(inferred.data_type(), &DataType::Int32); + assert!(inferred.is_nullable()); + } + _ => panic!("Expected Placeholder expression in InSubquery"), + }, + _ => panic!("Expected InSubquery expression"), + } + } + + #[test] + fn infer_placeholder_not_in_subquery() { + // WHERE $1 NOT IN (SELECT a FROM t) + let subquery_field = Field::new("a", DataType::Int32, false); + let subquery_schema = Arc::new( + DFSchema::from_unqualified_fields( + vec![subquery_field].into(), + Default::default(), + ) + .unwrap(), + ); + let subquery = Subquery { + subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: subquery_schema, + })), + outer_ref_columns: vec![], + spans: Spans::new(), + }; + + let not_in_subquery = Expr::InSubquery(InSubquery { + expr: Box::new(Expr::Placeholder(Placeholder { + id: "$1".to_string(), + field: None, + })), + subquery, + negated: true, + }); + + let outer_schema = DFSchema::empty(); + let (inferred_expr, contains_placeholder) = not_in_subquery + .infer_placeholder_types(&outer_schema) + .unwrap(); + + assert!(contains_placeholder); + + match inferred_expr { + Expr::InSubquery(in_subquery) => { + assert!(in_subquery.negated, "negated flag must be preserved"); + match *in_subquery.expr { + Expr::Placeholder(placeholder) => { + let inferred = placeholder.field.expect("placeholder field"); + assert_eq!(inferred.data_type(), &DataType::Int32); + assert!(inferred.is_nullable()); + } + _ => { + panic!("Expected Placeholder expression in InSubquery") + } + } + } + _ => panic!("Expected InSubquery expression"), + } + } + + #[test] + fn infer_placeholder_set_comparison_any() { + // WHERE $1 = ANY (SELECT a FROM t) -- parallel to infer_placeholder_in_subquery + let subquery_field = Field::new("a", DataType::Int32, false); + let subquery_schema = Arc::new( + DFSchema::from_unqualified_fields( + vec![subquery_field].into(), + Default::default(), + ) + .unwrap(), + ); + let subquery = Subquery { + subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: subquery_schema, + })), + outer_ref_columns: vec![], + spans: Spans::new(), + }; + + let set_cmp = Expr::SetComparison(SetComparison { + expr: Box::new(Expr::Placeholder(Placeholder { + id: "$1".to_string(), + field: None, + })), + subquery, + op: Operator::Eq, + quantifier: SetQuantifier::Any, + }); + + let outer_schema = DFSchema::empty(); + let (inferred_expr, contains_placeholder) = + set_cmp.infer_placeholder_types(&outer_schema).unwrap(); + + assert!(contains_placeholder); + + match inferred_expr { + Expr::SetComparison(sc) => { + assert_eq!(sc.quantifier, SetQuantifier::Any); + match *sc.expr { + Expr::Placeholder(p) => { + let inferred = + p.field.expect("placeholder field should be Int32"); + assert_eq!(inferred.data_type(), &DataType::Int32); + assert!(inferred.is_nullable()); + } + _ => panic!("Expected Placeholder expression in SetComparison"), + } + } + _ => panic!("Expected SetComparison expression"), + } + } + + #[test] + fn infer_placeholder_set_comparison_all() { + // WHERE $1 <> ALL (SELECT a FROM t) + let subquery_field = Field::new("a", DataType::Int32, false); + let subquery_schema = Arc::new( + DFSchema::from_unqualified_fields( + vec![subquery_field].into(), + Default::default(), + ) + .unwrap(), + ); + let subquery = Subquery { + subquery: Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: subquery_schema, + })), + outer_ref_columns: vec![], + spans: Spans::new(), + }; + + let set_cmp = Expr::SetComparison(SetComparison { + expr: Box::new(Expr::Placeholder(Placeholder { + id: "$1".to_string(), + field: None, + })), + subquery, + op: Operator::NotEq, + quantifier: SetQuantifier::All, + }); + + let outer_schema = DFSchema::empty(); + let (inferred_expr, contains_placeholder) = + set_cmp.infer_placeholder_types(&outer_schema).unwrap(); + + assert!(contains_placeholder); + + match inferred_expr { + Expr::SetComparison(sc) => { + assert_eq!(sc.quantifier, SetQuantifier::All); + match *sc.expr { + Expr::Placeholder(p) => { + let inferred = + p.field.expect("placeholder field should be Int32"); + assert_eq!(inferred.data_type(), &DataType::Int32); + assert!(inferred.is_nullable()); + } + _ => panic!("Expected Placeholder expression in SetComparison"), + } + } + _ => panic!("Expected SetComparison expression"), + } + } + #[test] fn infer_placeholder_like_and_similar_to() { // name LIKE $1 @@ -3926,6 +4211,24 @@ mod test { Ok(()) } + #[test] + fn format_decimal_literal() { + let expr = lit(ScalarValue::Decimal128(Some(1), 1, 1)); + assert_eq!("Decimal128(0.1,1,1)", format!("{expr}")); + assert_eq!("Decimal128(0.1,1,1)", expr.schema_name().to_string()); + assert_eq!("0.1", expr.human_display().to_string()); + + let expr = lit(ScalarValue::Decimal128(Some(120), 3, 2)); + assert_eq!("Decimal128(1.20,3,2)", format!("{expr}")); + assert_eq!("Decimal128(1.20,3,2)", expr.schema_name().to_string()); + assert_eq!("1.20", expr.human_display().to_string()); + + let null_expr = lit(ScalarValue::Decimal128(None, 10, 2)); + assert_eq!("Decimal128(NULL,10,2)", format!("{null_expr}")); + assert_eq!("Decimal128(NULL,10,2)", null_expr.schema_name().to_string()); + assert_eq!("NULL", null_expr.human_display().to_string()); + } + #[test] fn test_partial_ord() { // Test validates that partial ord is defined for Expr, not diff --git a/datafusion/expr/src/expr_fn.rs b/datafusion/expr/src/expr_fn.rs index 9d711113e4f74..b1a5a12d155ce 100644 --- a/datafusion/expr/src/expr_fn.rs +++ b/datafusion/expr/src/expr_fn.rs @@ -386,10 +386,11 @@ pub fn when(when: Expr, then: Expr) -> CaseBuilder { CaseBuilder::new(None, vec![when], vec![then], None) } -/// Create a Unnest expression +/// Create a Unnest expression with default (non-outer) semantics. pub fn unnest(expr: Expr) -> Expr { Expr::Unnest(Unnest { expr: Box::new(expr), + outer: false, }) } diff --git a/datafusion/expr/src/expr_rewriter/mod.rs b/datafusion/expr/src/expr_rewriter/mod.rs index eab8114d6910b..7a6ac3fc8b062 100644 --- a/datafusion/expr/src/expr_rewriter/mod.rs +++ b/datafusion/expr/src/expr_rewriter/mod.rs @@ -87,13 +87,16 @@ pub fn normalize_col_with_schemas_and_ambiguity_check( using_columns: &[HashSet], ) -> Result { // Normalize column inside Unnest - if let Expr::Unnest(Unnest { expr }) = expr { + if let Expr::Unnest(Unnest { expr, outer }) = expr { let e = normalize_col_with_schemas_and_ambiguity_check( expr.as_ref().clone(), schemas, using_columns, )?; - return Ok(Expr::Unnest(Unnest { expr: Box::new(e) })); + return Ok(Expr::Unnest(Unnest { + expr: Box::new(e), + outer, + })); } expr.transform(|expr| { @@ -483,7 +486,7 @@ mod test { normalize_col_with_schemas_and_ambiguity_check(expr, &[&schemas], &[]) .unwrap_err() .strip_backtrace(); - let expected = "Schema error: No field named b. \ + let expected = "Schema error: No field named b.\n\ Valid fields are \"tableA\".a."; assert_eq!(error, expected); } diff --git a/datafusion/expr/src/expr_schema.rs b/datafusion/expr/src/expr_schema.rs index 039bbad65a660..8927fcf4d0bbe 100644 --- a/datafusion/expr/src/expr_schema.rs +++ b/datafusion/expr/src/expr_schema.rs @@ -157,7 +157,7 @@ impl ExprSchemable for Expr { Expr::Cast(Cast { field, .. }) | Expr::TryCast(TryCast { field, .. }) => { Ok(field.data_type().clone()) } - Expr::Unnest(Unnest { expr }) => { + Expr::Unnest(Unnest { expr, .. }) => { let arg_data_type = expr.get_type(schema)?; // Unnest's output type is the inner type of the list match arg_data_type { @@ -366,7 +366,14 @@ impl ExprSchemable for Expr { | Expr::IsNotUnknown(_) | Expr::Exists { .. } => Ok(false), Expr::SetComparison(_) => Ok(true), - Expr::InSubquery(InSubquery { expr, .. }) => expr.nullable(input_schema), + Expr::InSubquery(InSubquery { expr, subquery, .. }) => { + let expr_nullable = expr.nullable(input_schema)?; + let subquery_nullable = subquery.subquery.schema().fields().first().ok_or_else(|| { + plan_datafusion_err!("subquery must return exactly one column of data to compare against") + })?.is_nullable(); + + Ok(expr_nullable | subquery_nullable) + } Expr::ScalarSubquery(subquery) => { Ok(subquery.subquery.schema().field(0).is_nullable()) } @@ -796,8 +803,13 @@ mod tests { use std::collections::HashMap; use super::*; - use crate::{and, col, lit, not, or, out_ref_col_with_metadata, when}; + use crate::logical_plan::builder::LogicalTableSource; + use crate::{ + LogicalPlanBuilder, and, col, in_subquery, lit, not, or, + out_ref_col_with_metadata, when, + }; + use arrow::datatypes::Schema; use datafusion_common::{DFSchema, assert_or_internal_err}; macro_rules! test_is_expr_nullable { @@ -1192,6 +1204,76 @@ mod tests { } } + /// A scan of `t`, whose single column `a` has the given nullability. + fn scan_t(a_nullable: bool) -> LogicalPlanBuilder { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, a_nullable)]); + let source = Arc::new(LogicalTableSource::new(Arc::new(schema))); + LogicalPlanBuilder::scan("t", source, None).unwrap() + } + + #[test] + fn in_subquery_nullability() { + // `x IN (SELECT a FROM t)` evaluates to NULL when `x` is NULL, and when `x` + // matches no row while `a` contains a NULL. So it is nullable exactly when + // either the compared expression or the subquery's output column is. + let cases = [ + (false, false, false), + (false, true, true), + (true, false, true), + (true, true, true), + ]; + + for (x_nullable, a_nullable, expected) in cases { + let subquery = scan_t(a_nullable) + .project(vec![col("a")]) + .unwrap() + .build() + .unwrap(); + let expr = in_subquery(col("x"), Arc::new(subquery)); + let schema = MockExprSchema::new().with_nullable(x_nullable); + + assert_eq!(expr.nullable(&schema).unwrap(), expected); + } + } + + #[test] + fn in_subquery_nullability_uses_subquery_output_schema() { + // `DISTINCT` carries no expressions of its own, but its output column is still + // nullable, so the `IN` expression must be nullable too. + let subquery = scan_t(true) + .project(vec![col("a")]) + .unwrap() + .distinct() + .unwrap() + .build() + .unwrap(); + let expr = in_subquery(col("x"), Arc::new(subquery)); + assert!(expr.nullable(&MockExprSchema::new()).unwrap()); + + // A computed projection's expressions reference `t.a`, which does not appear in + // the subquery's output schema, so nullability must be read off that schema's + // single column rather than by resolving the projection's expressions against it. + let subquery = scan_t(false) + .project(vec![col("a") + lit(1)]) + .unwrap() + .build() + .unwrap(); + let expr = in_subquery(col("x"), Arc::new(subquery)); + assert!(!expr.nullable(&MockExprSchema::new()).unwrap()); + } + + #[test] + fn in_subquery_nullability_errors_for_no_subquery_columns() { + let subquery = LogicalPlanBuilder::empty(false).build().unwrap(); + let expr = in_subquery(col("x"), Arc::new(subquery)); + + let err = expr.nullable(&MockExprSchema::new()).unwrap_err(); + assert_eq!( + err.strip_backtrace(), + "Error during planning: subquery must return exactly one column of data to compare against" + ); + } + #[test] fn test_scalar_variable() { let mut meta = HashMap::new(); diff --git a/datafusion/expr/src/higher_order_function.rs b/datafusion/expr/src/higher_order_function.rs index 00522ad97b9e2..c300be8f6fcfe 100644 --- a/datafusion/expr/src/higher_order_function.rs +++ b/datafusion/expr/src/higher_order_function.rs @@ -22,8 +22,9 @@ use crate::expr::{ schema_name_from_exprs_comma_separated_without_space, }; use crate::type_coercion::functions::value_fields_with_higher_order_udf; +use crate::udf_eq::UdfEq; use crate::{ColumnarValue, Documentation, Expr, ExprSchemable}; -use arrow::array::{ArrayRef, RecordBatch}; +use arrow::array::{ArrayRef, RecordBatch, RecordBatchOptions}; use arrow::datatypes::{DataType, FieldRef, Schema}; use arrow_schema::SchemaRef; use datafusion_common::config::ConfigOptions; @@ -67,14 +68,14 @@ pub enum HigherOrderTypeSignature { /// function. /// /// If this signature is specified, - /// DataFusion will call [`HigherOrderUDF::coerce_value_types`] to prepare argument types. + /// DataFusion will call [`HigherOrderUDFImpl::coerce_value_types`] to prepare argument types. UserDefined, /// One or more lambdas or arguments with arbitrary types VariadicAny, /// The specified number of lambdas or arguments with arbitrary types. Any(usize), /// Exactly the specified arguments in the given order, with arbitrary types. - /// DataFusion will call [`HigherOrderUDF::coerce_value_types`] to prepare the value + /// DataFusion will call [`HigherOrderUDFImpl::coerce_value_types`] to prepare the value /// argument types. Exact(Vec>), } @@ -91,7 +92,7 @@ pub struct HigherOrderSignature { pub type_signature: HigherOrderTypeSignature, /// The volatility of the function. See [Volatility] for more information. pub volatility: Volatility, - /// The max number of times to call [HigherOrderUDF::lambda_parameters] before raising an error. + /// The max number of times to call [HigherOrderUDFImpl::lambda_parameters] before raising an error. /// Used to guard against implementations that causes an infinite loop by endlessly returning /// [LambdaParametersProgress::Partial]. Defaults to 256 pub lambda_parameters_max_iterations: usize, @@ -137,7 +138,7 @@ impl HigherOrderSignature { } /// Exactly the specified arguments in the given order, with arbitrary types. - /// DataFusion will call [`HigherOrderUDF::coerce_value_types`] to prepare the value + /// DataFusion will call [`HigherOrderUDFImpl::coerce_value_types`] to prepare the value /// argument types. /// /// # Example @@ -158,13 +159,13 @@ impl HigherOrderSignature { } } -impl PartialEq for dyn HigherOrderUDF { +impl PartialEq for dyn HigherOrderUDFImpl { fn eq(&self, other: &Self) -> bool { self.dyn_eq(other as _) } } -impl PartialOrd for dyn HigherOrderUDF { +impl PartialOrd for dyn HigherOrderUDFImpl { fn partial_cmp(&self, other: &Self) -> Option { let mut cmp = self.name().cmp(other.name()); if cmp == Ordering::Equal { @@ -193,15 +194,15 @@ impl PartialOrd for dyn HigherOrderUDF { } } -impl Eq for dyn HigherOrderUDF {} +impl Eq for dyn HigherOrderUDFImpl {} -impl Hash for dyn HigherOrderUDF { +impl Hash for dyn HigherOrderUDFImpl { fn hash(&self, state: &mut H) { self.dyn_hash(state) } } -/// Arguments passed to [`HigherOrderUDF::invoke_with_args`] when invoking a +/// Arguments passed to [`HigherOrderUDFImpl::invoke_with_args`] when invoking a /// higher order function. #[derive(Debug, Clone)] pub struct HigherOrderFunctionArgs { @@ -210,7 +211,7 @@ pub struct HigherOrderFunctionArgs { /// Field associated with each arg, if it exists /// For lambdas, it will be the field of the result of /// the lambda if evaluated with the parameters - /// returned from [`HigherOrderUDF::lambda_parameters`] + /// returned from [`HigherOrderUDFImpl::lambda_parameters`] pub arg_fields: Vec>, /// The number of rows in record batch being evaluated pub number_rows: usize, @@ -238,6 +239,26 @@ pub struct LambdaArgument { /// For example, for `array_transform([2], v -> -v)`, /// this will be `vec![Field::new("v", DataType::Int32, true)]` params: Vec, + /// Indices into [`Self::params`] of the parameters that are actually + /// referenced by [`Self::body`] (taking nested-lambda shadowing into + /// account), in the original declaration order of `params`. + /// + /// [`Self::evaluate`] only evaluates and pushes the closures whose + /// corresponding parameter index appears here, so unused declared + /// parameters leave no slot in the merged batch and the body's compressed + /// column indices line up directly with what the evaluator built. + /// + /// Callers who already have a `LambdaExpr` should pass + /// `LambdaExpr::used_param_indices()` directly to [`Self::new`] — both + /// are indices into the same positionally-aligned `params` list. + /// + /// Every index here must be `< params.len()`; see the precondition on + /// [`Self::new`]. + /// + /// Relies on captures sorting before this lambda's own params in the + /// planner's (un-projected) index space, which is what makes + /// `captures ++ used_params` below line up with the projected body. + used_param_indices: Vec, /// The body of the lambda /// /// For example, for `array_transform([2], v -> -v)`, @@ -256,26 +277,45 @@ pub struct LambdaArgument { } impl LambdaArgument { + /// # Preconditions + /// + /// Every index in `used_param_indices` must be `< params.len()`; + /// violating this panics on out-of-bounds indexing below. Callers should + /// pass `LambdaExpr::used_param_indices()`, which always indexes into the + /// same `params` list, rather than constructing indices by hand. pub fn new( params: Vec, body: Arc, captures: Option, + used_param_indices: &[usize], ) -> Self { - let fields = match &captures { + debug_assert!( + used_param_indices.iter().all(|i| *i < params.len()), + "used_param_indices contains an index out of bounds for params \ + (len {}): {:?}", + params.len(), + used_param_indices + ); + + let used_param_indices = used_param_indices.to_vec(); + let effective_params = used_param_indices.iter().map(|i| Arc::clone(¶ms[*i])); + + let fields: Vec = match &captures { Some(batch) => batch .schema_ref() .fields() .iter() .cloned() - .chain(params.clone()) + .chain(effective_params) .collect(), - None => params.clone(), + None => effective_params.collect(), }; let schema = Arc::new(Schema::new(fields)); Self { params, + used_param_indices, body, schema, captures, @@ -284,7 +324,12 @@ impl LambdaArgument { /// Evaluate this lambda /// `args` should evaluate to the value of each parameter - /// of the correspondent lambda returned in [HigherOrderUDF::lambda_parameters]. + /// of the correspondent lambda returned in [HigherOrderUDFImpl::lambda_parameters]. + /// + /// Only the closures in `args` for parameters the lambda body actually + /// references are called; closures for declared-but-unused parameters + /// are skipped entirely. Callers should not rely on every closure in + /// `args` being invoked. /// /// `spread_captures` is responsible for transforming the captured column arrays /// so they align with the evaluation batch. Captures are snapshotted from the @@ -343,6 +388,7 @@ impl LambdaArgument { spread_captures.as_ref(), Arc::clone(&self.schema), &self.params, + &self.used_param_indices, args, )?; @@ -354,6 +400,7 @@ fn merge_captures_with_variables( captures: Option<&RecordBatch>, schema: SchemaRef, params: &[FieldRef], + used_param_indices: &[usize], variables: &[&dyn Fn() -> Result], ) -> Result { if variables.len() < params.len() { @@ -364,23 +411,42 @@ fn merge_captures_with_variables( ); } + let push_param_arrays = |columns: &mut Vec| -> Result<()> { + for &i in used_param_indices { + columns.push(variables[i]()?); + } + Ok(()) + }; + let columns = match captures { Some(captures) => { let mut columns = captures.columns().to_vec(); - - for arg in &variables[..params.len()] { - columns.push(arg()?); - } - + push_param_arrays(&mut columns)?; + columns + } + None => { + let mut columns = Vec::with_capacity(used_param_indices.len()); + push_param_arrays(&mut columns)?; columns } - None => variables - .iter() - .take(params.len()) - .map(|arg| arg()) - .collect::>()?, }; + if columns.is_empty() { + // No columns to derive a row count from, so borrow one variable's + // array length instead (all variables have the same length). + let row_count = variables.first().ok_or_else(|| { + internal_datafusion_err!( + "merge_captures_with_variables: no variables to derive a row count from" + ) + })?()? + .len(); + return Ok(RecordBatch::try_new_with_options( + schema, + vec![], + &RecordBatchOptions::new().with_row_count(Some(row_count)), + )?); + } + Ok(RecordBatch::try_new(schema, columns)?) } @@ -390,13 +456,13 @@ fn merge_captures_with_variables( /// such as the type of the arguments, any scalar arguments and if the /// arguments can (ever) be null /// -/// See [`HigherOrderUDF::return_field_from_args`] for more information +/// See [`HigherOrderUDFImpl::return_field_from_args`] for more information #[derive(Clone, Debug)] pub struct HigherOrderReturnFieldArgs<'a> { /// The data types of the arguments to the function /// /// If argument `i` to the function is a lambda, it will be the field of the result of the - /// lambda if evaluated with the parameters returned from [`HigherOrderUDF::lambda_parameters`] + /// lambda if evaluated with the parameters returned from [`HigherOrderUDFImpl::lambda_parameters`] /// /// For example, with `array_transform([1], v -> v == 5)` /// this field will be @@ -426,19 +492,19 @@ pub enum ValueOrLambda { } /// Represents a step during the resolution of the parameters of all lambdas of a given -/// higher-order function via [HigherOrderUDF::lambda_parameters]. It's valid that the +/// higher-order function via [HigherOrderUDFImpl::lambda_parameters]. It's valid that the /// fields of a given lambda changes between steps, and is up to the implementation to /// provide during the function evaluation the parameters that matches the fields returned -/// at the [LambdaParametersProgress::Complete] step. See [HigherOrderUDF::lambda_parameters] +/// at the [LambdaParametersProgress::Complete] step. See [HigherOrderUDFImpl::lambda_parameters] /// docs for more details pub enum LambdaParametersProgress { /// The parameters of some lambdas are unknown due to a dependency on another lambda output field /// or are placeholders due to a dependency on it's own output field. It's perfectly valid to /// contain only `Some`'s and not a single `None`, representing lambdas that depends only on itself - /// and not on others. [HigherOrderUDF::lambda_parameters] will be called again with the output + /// and not on others. [HigherOrderUDFImpl::lambda_parameters] will be called again with the output /// field of all lambdas with known parameters. Partial(Vec>>), - /// There are no unmet dependencies and all parameters are known, [HigherOrderUDF::lambda_parameters] + /// There are no unmet dependencies and all parameters are known, [HigherOrderUDFImpl::lambda_parameters] /// will not be called again Complete(Vec>), } @@ -448,10 +514,13 @@ pub enum LambdaParametersProgress { /// This trait exposes the full API for implementing user defined functions and /// can be used to implement any function. /// +/// New higher order functions typically implement this trait and are then +/// wrapped in a [`HigherOrderUDF`] for registration with DataFusion. +/// /// See [`array_transform.rs`] for a commented complete implementation /// /// [`array_transform.rs`]: https://github.com/apache/datafusion/blob/main/datafusion/functions-nested/src/array_transform.rs -pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any { +pub trait HigherOrderUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { /// Returns this function's name fn name(&self) -> &str; @@ -546,11 +615,11 @@ pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any { /// /// For functions which lambda parameters depends on the output of other lambdas, or on their own lambda, /// this can return [LambdaParametersProgress::Partial] until all dependencies are met. Note that for - /// lambda with cyclic dependencies, you likely want to use [HigherOrderUDF::coerce_values_for_lambdas] too. + /// lambda with cyclic dependencies, you likely want to use [HigherOrderUDFImpl::coerce_values_for_lambdas] too. /// Take as an example a flexible array_reduce with the signature `(arr: [V], initial_value: I, (ACC, V) -> ACC, (ACC) -> O) -> O`. /// It has a cyclic dependency in the merge lambda, and a dependency of the finish lambda in the merge lambda, /// and only requires the initial value to be *coercible* to the output of the merge lambda, which is defined by - /// it's [HigherOrderUDF::coerce_values_for_lambdas] implementation. The expression + /// it's [HigherOrderUDFImpl::coerce_values_for_lambdas] implementation. The expression /// /// `array_reduce([1.2, 2.1], 0, (acc, v) -> acc + v + 1.5, v -> v > 5.1)` /// @@ -658,7 +727,7 @@ pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any { ) -> Result; /// Coerce value arguments of a function call to types that the function can evaluate also taking into - /// account the *output type of it's lambdas*. This differs from [HigherOrderUDF::coerce_value_types] + /// account the *output type of it's lambdas*. This differs from [HigherOrderUDFImpl::coerce_value_types] /// that only has access to the type of it's value arguments because it's called before the output type /// of lambdas are known. /// @@ -744,7 +813,7 @@ pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any { /// Setting this to true prevents certain optimizations such as common /// subexpression elimination /// - /// When overriding this function to return `true`, [HigherOrderUDF::conditional_arguments] can also be + /// When overriding this function to return `true`, [HigherOrderUDFImpl::conditional_arguments] can also be /// overridden to report more accurately which arguments are eagerly evaluated and which ones /// lazily. fn short_circuits(&self) -> bool { @@ -768,7 +837,7 @@ pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any { /// Implementations must ensure that the two returned `Vec`s are disjunct, /// and that each argument from `args` is present in one the two `Vec`s. /// - /// When overriding this function, [HigherOrderUDF::short_circuits] must + /// When overriding this function, [HigherOrderUDFImpl::short_circuits] must /// be overridden to return `true`. fn conditional_arguments<'a>( &self, @@ -783,7 +852,7 @@ pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any { /// Coerce value arguments of a function call to types that the function can evaluate. /// Note that if you need to coerce values based on the output type of lambdas, you - /// must use [HigherOrderUDF::coerce_values_for_lambdas], as this function is used before + /// must use [HigherOrderUDFImpl::coerce_values_for_lambdas], as this function is used before /// the output type of lambdas are known /// /// See the [type coercion module](crate::type_coercion) @@ -806,7 +875,7 @@ pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any { ) } - /// Returns the documentation for this HigherOrderUDF. + /// Returns the documentation for this function. /// /// Documentation can be accessed programmatically as well as generating /// publicly facing documentation. @@ -815,6 +884,296 @@ pub trait HigherOrderUDF: Debug + DynEq + DynHash + Send + Sync + Any { } } +/// Logical representation of a Higher Order User Defined Function. +/// +/// A higher order function takes one or more lambda arguments in addition to +/// regular value arguments. This struct contains the information DataFusion +/// needs to plan and invoke functions you supply such as name, type signature, +/// return type, and actual implementation. +#[derive(Debug, Clone)] +pub struct HigherOrderUDF { + inner: Arc, +} + +impl PartialEq for HigherOrderUDF { + fn eq(&self, other: &Self) -> bool { + self.inner.as_ref().dyn_eq(other.inner.as_ref()) + } +} + +impl PartialOrd for HigherOrderUDF { + fn partial_cmp(&self, other: &Self) -> Option { + let mut cmp = self.name().cmp(other.name()); + if cmp == Ordering::Equal { + cmp = self.signature().partial_cmp(other.signature())?; + } + if cmp == Ordering::Equal { + cmp = self.aliases().partial_cmp(other.aliases())?; + } + // Contract for PartialOrd and PartialEq consistency requires that + // a == b if and only if partial_cmp(a, b) == Some(Equal). + if cmp == Ordering::Equal && self != other { + // Functions may have other properties besides name and signature + // that differentiate two instances (e.g. type, or arbitrary parameters). + // We cannot return Some(Equal) in such case. + return None; + } + debug_assert!( + cmp == Ordering::Equal || self != other, + "Detected incorrect implementation of PartialEq when comparing functions: '{}' and '{}'. \ + The functions compare as equal, but they are not equal based on general properties that \ + the PartialOrd implementation observes,", + self.name(), + other.name() + ); + Some(cmp) + } +} + +impl Eq for HigherOrderUDF {} + +impl Hash for HigherOrderUDF { + fn hash(&self, state: &mut H) { + self.inner.dyn_hash(state) + } +} + +impl HigherOrderUDF { + /// Create a new `HigherOrderUDF` from a [`HigherOrderUDFImpl`] trait object. + /// + /// Note this is the same as using the `From` impl (`HigherOrderUDF::from`). + pub fn new_from_impl(fun: F) -> HigherOrderUDF + where + F: HigherOrderUDFImpl + 'static, + { + Self::new_from_shared_impl(Arc::new(fun)) + } + + /// Create a new `HigherOrderUDF` from a shared [`HigherOrderUDFImpl`] trait object. + pub fn new_from_shared_impl(fun: Arc) -> HigherOrderUDF { + Self { inner: fun } + } + + /// Return the underlying [`HigherOrderUDFImpl`] trait object for this function. + pub fn inner(&self) -> &Arc { + &self.inner + } + + /// Adds additional names that can be used to invoke this function, in + /// addition to `name`. + /// + /// If you implement [`HigherOrderUDFImpl`] directly you should return aliases + /// directly. + pub fn with_aliases(self, aliases: impl IntoIterator) -> Self { + Self::new_from_impl(AliasedHigherOrderUDFImpl::new( + Arc::clone(&self.inner), + aliases, + )) + } + + /// Returns this function's name. + /// + /// See [`HigherOrderUDFImpl::name`] for more details. + pub fn name(&self) -> &str { + self.inner.name() + } + + /// Returns the aliases for this function. + /// + /// See [`HigherOrderUDF::with_aliases`] for more details. + pub fn aliases(&self) -> &[String] { + self.inner.aliases() + } + + /// Returns this function's schema_name. + /// + /// See [`HigherOrderUDFImpl::schema_name`] for more details. + pub fn schema_name(&self, args: &[Expr]) -> Result { + self.inner.schema_name(args) + } + + /// Returns this function's [`HigherOrderSignature`]. + pub fn signature(&self) -> &HigherOrderSignature { + self.inner.signature() + } + + /// Returns the parameters of all lambdas of this function for the current step. + /// + /// See [`HigherOrderUDFImpl::lambda_parameters`] for more details. + pub fn lambda_parameters( + &self, + step: usize, + fields: &[ValueOrLambda>], + ) -> Result { + self.inner.lambda_parameters(step, fields) + } + + /// Coerce value arguments based on lambda output types. + /// + /// See [`HigherOrderUDFImpl::coerce_values_for_lambdas`] for more details. + pub fn coerce_values_for_lambdas( + &self, + fields: &[ValueOrLambda], + ) -> Result>> { + self.inner.coerce_values_for_lambdas(fields) + } + + /// Returns the return field of the function given its arguments. + /// + /// See [`HigherOrderUDFImpl::return_field_from_args`] for more details. + pub fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result { + self.inner.return_field_from_args(args) + } + + /// Whether List or LargeList arguments should have non-empty null sublists + /// cleaned before invoking this function. + pub fn clear_null_values(&self) -> bool { + self.inner.clear_null_values() + } + + /// Invoke the function returning the appropriate result. + /// + /// See [`HigherOrderUDFImpl::invoke_with_args`] for more details. + pub fn invoke_with_args( + &self, + args: HigherOrderFunctionArgs, + ) -> Result { + self.inner.invoke_with_args(args) + } + + /// Returns true if some of this function's subexpressions may not be evaluated. + /// + /// See [`HigherOrderUDFImpl::short_circuits`] for more details. + pub fn short_circuits(&self) -> bool { + self.inner.short_circuits() + } + + /// Returns which arguments are evaluated eagerly vs lazily. + /// + /// See [`HigherOrderUDFImpl::conditional_arguments`] for more details. + pub fn conditional_arguments<'a>( + &self, + args: &'a [Expr], + ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> { + self.inner.conditional_arguments(args) + } + + /// Coerce value arguments of a function call to types that the function can evaluate. + /// + /// See [`HigherOrderUDFImpl::coerce_value_types`] for more details. + pub fn coerce_value_types(&self, arg_types: &[DataType]) -> Result> { + self.inner.coerce_value_types(arg_types) + } + + /// Returns the documentation for this function, if any. + pub fn documentation(&self) -> Option<&Documentation> { + self.inner.documentation() + } +} + +impl From for HigherOrderUDF +where + F: HigherOrderUDFImpl + 'static, +{ + fn from(fun: F) -> Self { + Self::new_from_impl(fun) + } +} + +/// `HigherOrderUDFImpl` that adds aliases to the underlying function. It is +/// better to implement [`HigherOrderUDFImpl`], which supports aliases, directly +/// if possible. +#[derive(Debug, PartialEq, Eq, Hash)] +struct AliasedHigherOrderUDFImpl { + inner: UdfEq>, + aliases: Vec, +} + +impl AliasedHigherOrderUDFImpl { + fn new( + inner: Arc, + new_aliases: impl IntoIterator, + ) -> Self { + let mut aliases = inner.aliases().to_vec(); + aliases.extend(new_aliases.into_iter().map(|s| s.to_string())); + Self { + inner: inner.into(), + aliases, + } + } +} + +#[warn(clippy::missing_trait_methods)] // Delegates, so it should implement every single trait method +impl HigherOrderUDFImpl for AliasedHigherOrderUDFImpl { + fn name(&self) -> &str { + self.inner.name() + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn schema_name(&self, args: &[Expr]) -> Result { + self.inner.schema_name(args) + } + + fn signature(&self) -> &HigherOrderSignature { + self.inner.signature() + } + + fn lambda_parameters( + &self, + step: usize, + fields: &[ValueOrLambda>], + ) -> Result { + self.inner.lambda_parameters(step, fields) + } + + fn coerce_values_for_lambdas( + &self, + fields: &[ValueOrLambda], + ) -> Result>> { + self.inner.coerce_values_for_lambdas(fields) + } + + fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result { + self.inner.return_field_from_args(args) + } + + fn clear_null_values(&self) -> bool { + self.inner.clear_null_values() + } + + fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { + self.inner.invoke_with_args(args) + } + + fn short_circuits(&self) -> bool { + self.inner.short_circuits() + } + + fn conditional_arguments<'a>( + &self, + args: &'a [Expr], + ) -> Option<(Vec<&'a Expr>, Vec<&'a Expr>)> { + self.inner.conditional_arguments(args) + } + + fn coerce_value_types(&self, arg_types: &[DataType]) -> Result> { + self.inner.coerce_value_types(arg_types) + } + + fn documentation(&self) -> Option<&Documentation> { + self.inner.documentation() + } +} + pub(crate) fn resolve_lambda_variables( expr: Expr, schema: &DFSchema, @@ -854,7 +1213,7 @@ pub(crate) fn resolve_lambda_variables( } fn resolve_higher_order_function( - func: Arc, + func: Arc, args: Vec, schema: &DFSchema, // a map of lambda variable name => a never empty stack of fields [ [..shadowed], in_scope ] @@ -1083,8 +1442,8 @@ mod tests { use datafusion_expr_common::signature::Volatility; use crate::{ - Expr, HigherOrderSignature, HigherOrderUDF, LambdaParametersProgress, - ValueOrLambda, col, + Expr, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, + LambdaParametersProgress, ValueOrLambda, col, expr::{HigherOrderFunction, LambdaVariable}, lambda, lambda_var, lit, }; @@ -1095,7 +1454,7 @@ mod tests { field: &'static str, signature: HigherOrderSignature, } - impl HigherOrderUDF for TestHigherOrderUDF { + impl HigherOrderUDFImpl for TestHigherOrderUDF { fn name(&self) -> &str { self.name } @@ -1158,12 +1517,12 @@ mod tests { assert_eq!(b.partial_cmp(&o), Some(Ordering::Less)); } - fn test_func(name: &'static str, parameter: &'static str) -> Arc { - Arc::new(TestHigherOrderUDF { + fn test_func(name: &'static str, parameter: &'static str) -> Arc { + Arc::new(HigherOrderUDF::new_from_impl(TestHigherOrderUDF { name, field: parameter, signature: HigherOrderSignature::variadic_any(Volatility::Immutable), - }) + })) } fn hash(value: &T) -> u64 { @@ -1177,7 +1536,7 @@ mod tests { signature: HigherOrderSignature, } - impl HigherOrderUDF for MockArrayReduce { + impl HigherOrderUDFImpl for MockArrayReduce { fn name(&self) -> &str { "array_reduce" } @@ -1274,9 +1633,9 @@ mod tests { )])) .unwrap(); - let func = Arc::new(MockArrayReduce { + let func = Arc::new(HigherOrderUDF::new_from_impl(MockArrayReduce { signature: HigherOrderSignature::variadic_any(Volatility::Immutable), - }) as _; + })); /* array_reduce( @@ -1387,4 +1746,109 @@ mod tests { Some(Arc::new(Field::new(name, dt, nullable))), )) } + + /// A physical expression that reads the column at a fixed index of the + /// batch it is evaluated against, for exercising [`LambdaArgument`] + /// directly without depending on `datafusion-physical-expr`. + #[derive(Debug, Eq, PartialEq, Hash)] + struct ColumnAt(usize); + + impl std::fmt::Display for ColumnAt { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "column_at({})", self.0) + } + } + + impl PhysicalExpr for ColumnAt { + fn evaluate(&self, batch: &RecordBatch) -> Result { + Ok(ColumnarValue::Array(Arc::clone(batch.column(self.0)))) + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{self}") + } + } + + /// `(k, v) -> v` with only `v` used must push `v`'s array, not `k`'s. + #[test] + fn test_lambda_argument_evaluate_pushes_only_used_param() { + use arrow::array::Int32Array; + + let k_field = Arc::new(Field::new("k", DataType::Int32, true)); + let v_field = Arc::new(Field::new("v", DataType::Int32, true)); + + let body = Arc::new(ColumnAt(0)) as Arc; + let lambda_arg = LambdaArgument::new(vec![k_field, v_field], body, None, &[1]); + + let k_values: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300])); + let v_values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let k_closure = || -> Result { Ok(Arc::clone(&k_values)) }; + let v_closure = || -> Result { Ok(Arc::clone(&v_values)) }; + let args: Vec<&dyn Fn() -> Result> = vec![&k_closure, &v_closure]; + + let result = lambda_arg + .evaluate(&args, |arrays| Ok(arrays.to_vec())) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &Int32Array::from(vec![1, 2, 3]), + "body should read v's values, not k's" + ); + } + + /// Same as above, but with a capture occupying the leading slot. + #[test] + fn test_lambda_argument_evaluate_pushes_only_used_param_with_captures() { + use arrow::array::Int32Array; + + let cap_field = Arc::new(Field::new("cap", DataType::Int32, true)); + let k_field = Arc::new(Field::new("k", DataType::Int32, true)); + let v_field = Arc::new(Field::new("v", DataType::Int32, true)); + + let body = Arc::new(ColumnAt(1)) as Arc; + + let cap_values: ArrayRef = Arc::new(Int32Array::from(vec![9, 9, 9])); + let captures = RecordBatch::try_new( + Arc::new(Schema::new(vec![cap_field])), + vec![cap_values], + ) + .unwrap(); + + let lambda_arg = + LambdaArgument::new(vec![k_field, v_field], body, Some(captures), &[1]); + + let k_values: ArrayRef = Arc::new(Int32Array::from(vec![100, 200, 300])); + let v_values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let k_closure = || -> Result { Ok(Arc::clone(&k_values)) }; + let v_closure = || -> Result { Ok(Arc::clone(&v_values)) }; + let args: Vec<&dyn Fn() -> Result> = vec![&k_closure, &v_closure]; + + let result = lambda_arg + .evaluate(&args, |arrays| Ok(arrays.to_vec())) + .unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &Int32Array::from(vec![1, 2, 3]), + "body should read v's values, not k's or the capture's" + ); + } } diff --git a/datafusion/expr/src/lib.rs b/datafusion/expr/src/lib.rs index da7f20783bd06..75041c701454a 100644 --- a/datafusion/expr/src/lib.rs +++ b/datafusion/expr/src/lib.rs @@ -55,6 +55,7 @@ pub mod expr_rewriter; pub mod expr_schema; pub mod extension_types; pub mod function; +pub mod physical_planning_context; pub mod select_expr; pub mod groups_accumulator { pub use datafusion_expr_common::groups_accumulator::*; @@ -68,6 +69,11 @@ pub mod dml { pub use crate::logical_plan::dml::*; } pub mod planner; +/// Protobuf conversions for [`WindowFrame`], [`WindowFrameBound`], +/// [`WindowFrameUnits`], [`MergeIntoClauseKind`](dml::MergeIntoClauseKind) and +/// [`NullTreatment`](expr::NullTreatment), gated on the `proto` feature. +#[cfg(feature = "proto")] +mod proto; pub mod registry; pub mod simplify; pub mod sort_properties { @@ -101,8 +107,8 @@ pub use datafusion_expr_common::groups_accumulator::{EmitTo, GroupsAccumulator}; pub use datafusion_expr_common::operator::Operator; pub use datafusion_expr_common::placement::ExpressionPlacement; pub use datafusion_expr_common::signature::{ - ArrayFunctionArgument, ArrayFunctionSignature, Coercion, Signature, - TIMEZONE_WILDCARD, TypeSignature, TypeSignatureClass, Volatility, + ArrayFunctionArgument, ArrayFunctionSignature, Coercion, EncodingPreservation, + Signature, TIMEZONE_WILDCARD, TypeSignature, TypeSignatureClass, Volatility, }; pub use datafusion_expr_common::type_coercion::binary; pub use expr::{ @@ -117,8 +123,8 @@ pub use function::{ }; pub use higher_order_function::{ HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature, - HigherOrderTypeSignature, HigherOrderUDF, LambdaArgument, LambdaParametersProgress, - ValueOrLambda, + HigherOrderTypeSignature, HigherOrderUDF, HigherOrderUDFImpl, LambdaArgument, + LambdaParametersProgress, ValueOrLambda, }; pub use literal::{ Literal, TimestampLiteral, lit, lit_timestamp_nano, lit_with_metadata, diff --git a/datafusion/expr/src/logical_plan/builder.rs b/datafusion/expr/src/logical_plan/builder.rs index 017a123eb035b..a3d2c6e17adf9 100644 --- a/datafusion/expr/src/logical_plan/builder.rs +++ b/datafusion/expr/src/logical_plan/builder.rs @@ -33,8 +33,8 @@ use crate::expr_rewriter::{ use crate::logical_plan::{ Aggregate, Analyze, Distinct, DistinctOn, EmptyRelation, Explain, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Prepare, - Projection, Repartition, Sort, SubqueryAlias, TableScan, Union, Unnest, Values, - Window, + Projection, Repartition, Sort, SubqueryAlias, TableScanBuilder, Union, Unnest, + Values, Window, }; use crate::select_expr::SelectExpr; use crate::utils::{ @@ -192,12 +192,13 @@ impl LogicalPlanBuilder { // Ensure that the recursive term has the same field types as the static term let coerced_recursive_term = coerce_plan_expr_for_schema(recursive_term, self.plan.schema())?; - Ok(Self::from(LogicalPlan::RecursiveQuery(RecursiveQuery { + let recursive_query = RecursiveQuery::try_new( name, - static_term: self.plan, - recursive_term: Arc::new(coerced_recursive_term), + self.plan, + Arc::new(coerced_recursive_term), is_distinct, - }))) + )?; + Ok(Self::from(LogicalPlan::RecursiveQuery(recursive_query))) } /// Create a values list based relation, and the schema is inferred from data, consuming @@ -282,7 +283,7 @@ impl LogicalPlanBuilder { && !can_cast_types(&data_type, field_type) { return exec_err!( - "type mismatch and can't cast to got {} and {}", + "Types don't match and no valid cast exists, received data of type {} for field of type {}", data_type, field_type ); @@ -302,6 +303,7 @@ impl LogicalPlanBuilder { for j in 0..n_cols { let mut common_type: Option = None; let mut common_metadata: Option = None; + let mut nullable = false; for (i, row) in values.iter().enumerate() { let value = &row[j]; let metadata = value.metadata(&schema)?; @@ -316,13 +318,17 @@ impl LogicalPlanBuilder { } else { common_metadata = Some(metadata.clone()); } + if !nullable && value.nullable(&schema)? { + nullable = true; + } let data_type = value.get_type(&schema)?; if data_type == DataType::Null { continue; } if let Some(prev_type) = common_type { - // get common type of each column values. + // Widen the running type so that it can hold both the + // previously seen rows and this row's value. let data_types = vec![prev_type.clone(), data_type.clone()]; let Some(new_type) = type_union_resolution(&data_types) else { return plan_err!( @@ -334,13 +340,13 @@ impl LogicalPlanBuilder { common_type = Some(data_type); } } - // assuming common_type was not set, and no error, therefore the type should be NULL - // since the code loop skips NULL - fields.push_with_metadata( - common_type.unwrap_or(DataType::Null), - true, - common_metadata, - ); + // If common_type is not set, every value in this column had type + // NULL. A DataType::Null field is always nullable. + let (data_type, nullable) = match common_type { + Some(t) => (t, nullable), + None => (DataType::Null, true), + }; + fields.push_with_metadata(data_type, nullable, common_metadata); } Self::infer_inner(values, fields, &schema) @@ -510,8 +516,11 @@ impl LogicalPlanBuilder { filters: Vec, fetch: Option, ) -> Result { - let table_scan = - TableScan::try_new(table_name, table_source, projection, filters, fetch)?; + let table_scan = TableScanBuilder::new(table_name, table_source) + .with_projection(projection) + .with_filters(filters) + .with_fetch(fetch) + .build()?; // Inline TableScan if table_scan.filters.is_empty() @@ -1327,8 +1336,11 @@ impl LogicalPlanBuilder { if explain_option.analyze { Ok(Self::new(LogicalPlan::Analyze(Analyze { verbose: explain_option.verbose, + format: explain_option.format, input: self.plan, schema, + analyze_level: explain_option.analyze_level, + analyze_categories: explain_option.analyze_categories, }))) } else { let stringified_plans = @@ -1341,6 +1353,7 @@ impl LogicalPlanBuilder { stringified_plans, schema, logical_optimization_succeeded: false, + show_statistics: explain_option.show_statistics, }))) } } @@ -2887,6 +2900,48 @@ mod tests { Ok(()) } + #[test] + fn plan_builder_aggregate_rejects_nested_aggregates() -> Result<()> { + // https://github.com/apache/datafusion/issues/23812 + let err = table_scan( + Some("employee_csv"), + &employee_schema(), + Some(vec![0, 3, 4]), + )? + .aggregate(vec![col("id")], vec![sum(sum(col("salary")))]) + .expect_err("nested aggregates should be rejected"); + + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(employee_csv.salary)' is nested inside 'sum(sum(employee_csv.salary))'" + ); + + Ok(()) + } + + #[test] + fn plan_builder_window_rejects_nested_window_functions() -> Result<()> { + // https://github.com/apache/datafusion/issues/23812 + let sum_over = |arg| { + Expr::from(expr::WindowFunction::new( + crate::WindowFunctionDefinition::AggregateUDF( + crate::test::function_stub::sum_udaf(), + ), + vec![arg], + )) + }; + let err = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![4]))? + .window(vec![sum_over(sum_over(col("salary")))]) + .expect_err("nested window functions should be rejected"); + + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Window function calls cannot be nested: 'sum(employee_csv.salary) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(employee_csv.salary) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'" + ); + + Ok(()) + } + #[test] fn test_join_metadata() -> Result<()> { let left_schema = DFSchema::new_with_metadata( @@ -2979,4 +3034,25 @@ mod tests { ] ); } + + #[test] + fn test_values_with_schema_type_mismatch_error_message() { + // Date32 field, but the value is a Boolean, which cannot be cast to Date32. + let schema = Arc::new( + DFSchema::from_unqualified_fields( + vec![Field::new("a", DataType::Date32, false)].into(), + HashMap::new(), + ) + .unwrap(), + ); + + let err = LogicalPlanBuilder::values_with_schema(vec![vec![lit(true)]], &schema) + .unwrap_err(); + + assert_eq!( + err.strip_backtrace(), + "Execution error: Types don't match and no valid cast exists, \ + received data of type Boolean for field of type Date32" + ); + } } diff --git a/datafusion/expr/src/logical_plan/ddl.rs b/datafusion/expr/src/logical_plan/ddl.rs index 5779fb0c4ea5b..51d88e43c1576 100644 --- a/datafusion/expr/src/logical_plan/ddl.rs +++ b/datafusion/expr/src/logical_plan/ddl.rs @@ -38,8 +38,10 @@ use sqlparser::ast::Ident; /// Various types of DDL (CREATE / DROP) catalog manipulation #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub enum DdlStatement { - /// Creates an external table. - CreateExternalTable(CreateExternalTable), + /// Creates an external table. Boxed to keep `LogicalPlan` enum size down + /// — `CreateExternalTable` is ~312 bytes, dwarfing every other variant + /// in the plan tree and forcing the whole enum to that width. + CreateExternalTable(Box), /// Creates an in memory table. CreateMemoryTable(CreateMemoryTable), /// Creates a new view. @@ -56,8 +58,9 @@ pub enum DdlStatement { DropView(DropView), /// Drops a catalog schema DropCatalogSchema(DropCatalogSchema), - /// Create function statement - CreateFunction(CreateFunction), + /// Create function statement. Boxed for the same reason as + /// [`Self::CreateExternalTable`] (~288 bytes). + CreateFunction(Box), /// Drop function statement DropFunction(DropFunction), } @@ -66,9 +69,7 @@ impl DdlStatement { /// Get a reference to the logical plan's schema pub fn schema(&self) -> &DFSchemaRef { match self { - DdlStatement::CreateExternalTable(CreateExternalTable { schema, .. }) => { - schema - } + DdlStatement::CreateExternalTable(ce) => &ce.schema, DdlStatement::CreateMemoryTable(CreateMemoryTable { input, .. }) | DdlStatement::CreateView(CreateView { input, .. }) => input.schema(), DdlStatement::CreateCatalogSchema(CreateCatalogSchema { schema, .. }) => { @@ -79,7 +80,7 @@ impl DdlStatement { DdlStatement::DropTable(DropTable { schema, .. }) => schema, DdlStatement::DropView(DropView { schema, .. }) => schema, DdlStatement::DropCatalogSchema(DropCatalogSchema { schema, .. }) => schema, - DdlStatement::CreateFunction(CreateFunction { schema, .. }) => schema, + DdlStatement::CreateFunction(cf) => &cf.schema, DdlStatement::DropFunction(DropFunction { schema, .. }) => schema, } } @@ -131,11 +132,9 @@ impl DdlStatement { impl Display for Wrapper<'_> { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self.0 { - DdlStatement::CreateExternalTable(CreateExternalTable { - name, - constraints, - .. - }) => { + DdlStatement::CreateExternalTable(ce) => { + let name = &ce.name; + let constraints = &ce.constraints; if constraints.is_empty() { write!(f, "CreateExternalTable: {name:?}") } else { @@ -191,7 +190,8 @@ impl DdlStatement { "DropCatalogSchema: {name:?} if not exist:={if_exists} cascade:={cascade}" ) } - DdlStatement::CreateFunction(CreateFunction { name, .. }) => { + DdlStatement::CreateFunction(cf) => { + let name = &cf.name; write!(f, "CreateFunction: name {name:?}") } DdlStatement::DropFunction(DropFunction { name, .. }) => { @@ -211,8 +211,12 @@ pub struct CreateExternalTable { pub schema: DFSchemaRef, /// The table name pub name: TableReference, - /// The physical location - pub location: String, + /// The physical locations of the table files. + /// + /// More than one location may be supplied (for example + /// `CREATE EXTERNAL TABLE ... LOCATION ('a.parquet', 'b.parquet')`), in which + /// case the files are read together as a single table. + pub locations: Vec, /// The file type of physical file pub file_type: String, /// Partition Columns @@ -266,7 +270,7 @@ impl CreateExternalTable { ) -> CreateExternalTableBuilder { CreateExternalTableBuilder { name: name.into(), - location: location.into(), + locations: vec![location.into()], file_type: file_type.into(), schema, table_partition_cols: vec![], @@ -289,7 +293,7 @@ impl CreateExternalTable { #[derive(Debug, Clone)] pub struct CreateExternalTableBuilder { name: TableReference, - location: String, + locations: Vec, file_type: String, schema: DFSchemaRef, table_partition_cols: Vec, @@ -311,6 +315,16 @@ impl CreateExternalTableBuilder { self } + /// Set the physical locations of the table files, replacing the single + /// location supplied to [`CreateExternalTable::builder`]. + /// + /// When more than one location is provided the files are read together as + /// a single table. + pub fn with_locations(mut self, locations: Vec) -> Self { + self.locations = locations; + self + } + /// Set the if_not_exists flag pub fn with_if_not_exists(mut self, if_not_exists: bool) -> Self { self.if_not_exists = if_not_exists; @@ -373,7 +387,7 @@ impl CreateExternalTableBuilder { CreateExternalTable { schema: self.schema, name: self.name, - location: self.location, + locations: self.locations, file_type: self.file_type, table_partition_cols: self.table_partition_cols, if_not_exists: self.if_not_exists, @@ -394,7 +408,7 @@ impl Hash for CreateExternalTable { fn hash(&self, state: &mut H) { self.schema.hash(state); self.name.hash(state); - self.location.hash(state); + self.locations.hash(state); self.file_type.hash(state); self.table_partition_cols.hash(state); self.if_not_exists.hash(state); @@ -413,8 +427,8 @@ impl PartialOrd for CreateExternalTable { struct ComparableCreateExternalTable<'a> { /// The table name pub name: &'a TableReference, - /// The physical location - pub location: &'a String, + /// The physical locations + pub locations: &'a Vec, /// The file type of physical file pub file_type: &'a String, /// Partition Columns @@ -432,7 +446,7 @@ impl PartialOrd for CreateExternalTable { } let comparable_self = ComparableCreateExternalTable { name: &self.name, - location: &self.location, + locations: &self.locations, file_type: &self.file_type, table_partition_cols: &self.table_partition_cols, if_not_exists: &self.if_not_exists, @@ -443,7 +457,7 @@ impl PartialOrd for CreateExternalTable { }; let comparable_other = ComparableCreateExternalTable { name: &other.name, - location: &other.location, + locations: &other.locations, file_type: &other.file_type, table_partition_cols: &other.table_partition_cols, if_not_exists: &other.if_not_exists, diff --git a/datafusion/expr/src/logical_plan/display.rs b/datafusion/expr/src/logical_plan/display.rs index 58c7feb616179..09f41c94f64fa 100644 --- a/datafusion/expr/src/logical_plan/display.rs +++ b/datafusion/expr/src/logical_plan/display.rs @@ -515,6 +515,23 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { "Partitioning Key": hash_expr }) } + Partitioning::Range(range) => { + let range_expr: Vec = + range.ordering().iter().map(|e| format!("{e}")).collect(); + let split_points: Vec = range + .split_points() + .iter() + .map(|e| format!("{e}")) + .collect(); + + json!({ + "Node Type": "Repartition", + "Partitioning Scheme": "Range", + "Partition Count": range.partition_count(), + "Partitioning Key": range_expr, + "Split Points": split_points + }) + } Partitioning::DistributeBy(expr) => { let dist_by_expr: Vec = expr.iter().map(|e| format!("{e}")).collect(); @@ -617,11 +634,7 @@ impl<'a, 'b> PgJsonVisitor<'a, 'b> { let list_type_columns = list_col_indices .iter() .map(|(i, unnest_info)| { - format!( - "{}|depth={:?}", - &input_columns[*i].to_string(), - unnest_info.depth - ) + format!("{}|depth={:?}", input_columns[*i], unnest_info.depth) }) .collect::>(); let struct_type_columns = struct_col_indices diff --git a/datafusion/expr/src/logical_plan/dml.rs b/datafusion/expr/src/logical_plan/dml.rs index b668cbfe2cc35..7717dfaff7a33 100644 --- a/datafusion/expr/src/logical_plan/dml.rs +++ b/datafusion/expr/src/logical_plan/dml.rs @@ -23,9 +23,9 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::file_options::file_type::FileType; -use datafusion_common::{DFSchemaRef, TableReference}; +use datafusion_common::{DFSchemaRef, Result, TableReference, internal_err}; -use crate::{LogicalPlan, TableSource}; +use crate::{Expr, LogicalPlan, TableSource}; /// Operator that copies the contents of a database to file(s) #[derive(Clone)] @@ -227,7 +227,11 @@ impl PartialOrd for DmlStatement { /// The type of DML operation to perform. /// /// See [`DmlStatement`] for more details. +/// +/// Marked `#[non_exhaustive]` so adding new variants in future releases is +/// not a SemVer break for downstream matchers. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +#[non_exhaustive] pub enum WriteOp { /// `INSERT INTO` operation Insert(InsertOp), @@ -239,6 +243,8 @@ pub enum WriteOp { Ctas, /// `TRUNCATE` operation Truncate, + /// `MERGE INTO` operation + MergeInto(Box), } impl WriteOp { @@ -250,6 +256,7 @@ impl WriteOp { WriteOp::Update => "Update", WriteOp::Ctas => "Ctas", WriteOp::Truncate => "Truncate", + WriteOp::MergeInto(_) => "MergeInto", } } } @@ -291,6 +298,196 @@ impl Display for InsertOp { } } +/// Describes a MERGE INTO operation's parameters. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct MergeIntoOp { + /// The join condition from `ON `. + pub on: Expr, + /// The WHEN clauses, in the order they appeared in the SQL. + pub clauses: Vec, +} + +impl MergeIntoOp { + /// Count of top-level [`Expr`]s owned by this operation (no allocation). + /// + /// Matches the length of [`Self::exprs`] and the `exprs` vec consumed by + /// [`Self::with_new_exprs`]. + fn expr_count(&self) -> usize { + 1 + self + .clauses + .iter() + .map(|c| { + c.predicate.is_some() as usize + + match &c.action { + MergeIntoAction::Update(a) => a.len(), + MergeIntoAction::Insert { values, .. } => values.len(), + MergeIntoAction::Delete => 0, + } + }) + .sum::() + } + + /// Top-level [`Expr`]s in stable order: `on`, then per-clause predicate + /// (if any) and action value expressions. + pub fn exprs(&self) -> Vec<&Expr> { + let mut out = Vec::with_capacity(self.expr_count()); + out.push(&self.on); + for clause in &self.clauses { + if let Some(predicate) = &clause.predicate { + out.push(predicate); + } + match &clause.action { + MergeIntoAction::Update(assignments) => { + out.extend(assignments.iter().map(|(_, value)| value)); + } + MergeIntoAction::Insert { values, .. } => { + out.extend(values.iter()); + } + MergeIntoAction::Delete => {} + } + } + out + } + + /// Rebuild this `MergeIntoOp` from a flat vector of new expressions, in + /// the same order produced by [`Self::exprs`]. The clause kinds, action + /// kinds, column lists, and presence/absence of each predicate are + /// preserved from `self`. + pub fn with_new_exprs(&self, exprs: Vec) -> Result { + let expected = self.expr_count(); + if exprs.len() != expected { + return internal_err!( + "MergeIntoOp::with_new_exprs expected {expected} expressions, got {}", + exprs.len() + ); + } + let mut iter = exprs.into_iter(); + let on = iter.next().expect("non-empty by length check"); + let clauses = self + .clauses + .iter() + .map(|clause| { + let predicate = clause + .predicate + .is_some() + .then(|| iter.next().expect("non-empty by length check")); + let action = match &clause.action { + MergeIntoAction::Update(assignments) => { + let assignments = assignments + .iter() + .map(|(name, _)| { + ( + name.clone(), + iter.next().expect("non-empty by length check"), + ) + }) + .collect(); + MergeIntoAction::Update(assignments) + } + MergeIntoAction::Insert { columns, values } => { + let values = values + .iter() + .map(|_| iter.next().expect("non-empty by length check")) + .collect(); + MergeIntoAction::Insert { + columns: columns.clone(), + values, + } + } + MergeIntoAction::Delete => MergeIntoAction::Delete, + }; + MergeIntoClause { + kind: clause.kind, + predicate, + action, + } + }) + .collect(); + Ok(Self { on, clauses }) + } +} + +/// A single WHEN clause within a MERGE INTO statement. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct MergeIntoClause { + /// Whether this fires on matched or unmatched rows. + pub kind: MergeIntoClauseKind, + /// Optional additional predicate (`AND `). + pub predicate: Option, + /// The action to take. + pub action: MergeIntoAction, +} + +/// Which rows a MERGE WHEN clause applies to. +/// +/// Mirrors `sqlparser::ast::MergeClauseKind` so that the SQL spelling is +/// preserved through the logical plan. +/// +/// **Note on `NotMatched` vs `NotMatchedByTarget`:** these two variants are +/// semantically identical — both describe a source row that has no matching +/// target row. `NotMatched` is the SQL standard short form (used by +/// Snowflake, Postgres, SQL Server); `NotMatchedByTarget` is BigQuery's +/// explicit form added for symmetry with `NotMatchedBySource`. Downstream +/// consumers (planners, table providers, optimizers) MUST treat the two +/// variants identically. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Hash)] +pub enum MergeIntoClauseKind { + /// `WHEN MATCHED` + Matched, + /// `WHEN NOT MATCHED` — see type-level note for the equivalence with + /// [`NotMatchedByTarget`](Self::NotMatchedByTarget). + NotMatched, + /// `WHEN NOT MATCHED BY TARGET` — see type-level note for the + /// equivalence with [`NotMatched`](Self::NotMatched). + NotMatchedByTarget, + /// `WHEN NOT MATCHED BY SOURCE` + NotMatchedBySource, +} + +impl MergeIntoClauseKind { + /// True if this clause fires on a source row that has no matching target + /// row. Returns `true` for both [`NotMatched`](Self::NotMatched) and + /// [`NotMatchedByTarget`](Self::NotMatchedByTarget) (see the type-level + /// note explaining why those two variants are semantically identical). + /// + /// Prefer this predicate over hand-written `matches!` arms so the + /// `NotMatched`/`NotMatchedByTarget` equivalence is enforced in one place. + pub fn is_not_matched_by_target(&self) -> bool { + matches!(self, Self::NotMatched | Self::NotMatchedByTarget) + } + + /// Collapse the SQL-spelling variants into the canonical three semantic + /// categories: [`Matched`](Self::Matched), + /// [`NotMatchedByTarget`](Self::NotMatchedByTarget) (covering both + /// "NOT MATCHED" spellings), and + /// [`NotMatchedBySource`](Self::NotMatchedBySource). + /// + /// Use this in downstream `match` expressions when the SQL spelling + /// distinction does not matter — e.g. in planners, optimizers, or + /// table-provider dispatch. + pub fn canonical(self) -> Self { + match self { + Self::NotMatched => Self::NotMatchedByTarget, + other => other, + } + } +} + +/// The action for a single WHEN clause. +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub enum MergeIntoAction { + /// `UPDATE SET col1 = expr1, col2 = expr2, ...`, stored as + /// `(column_name, value_expr)` pairs. + Update(Vec<(String, Expr)>), + /// `INSERT (col1, col2, ...) VALUES (expr1, expr2, ...)`. `columns` may + /// be empty, meaning all columns. + Insert { + columns: Vec, + values: Vec, + }, + Delete, +} + fn make_count_schema() -> DFSchemaRef { Arc::new( Schema::new(vec![Field::new("count", DataType::UInt64, false)]) @@ -298,3 +495,103 @@ fn make_count_schema() -> DFSchemaRef { .unwrap(), ) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::{col, lit}; + + #[test] + fn write_op_merge_into_name_and_display() { + let op = WriteOp::MergeInto(Box::new(MergeIntoOp { + on: col("id").eq(col("source_id")), + clauses: vec![MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: Some(col("qty").gt(lit(0_i64))), + action: MergeIntoAction::Update(vec![( + "qty".to_string(), + col("source_qty"), + )]), + }], + })); + assert_eq!(op.name(), "MergeInto"); + assert_eq!(format!("{op}"), "MergeInto"); + } + + #[test] + fn merge_into_clause_kind_is_not_matched_by_target() { + assert!(!MergeIntoClauseKind::Matched.is_not_matched_by_target()); + assert!(MergeIntoClauseKind::NotMatched.is_not_matched_by_target()); + assert!(MergeIntoClauseKind::NotMatchedByTarget.is_not_matched_by_target()); + assert!(!MergeIntoClauseKind::NotMatchedBySource.is_not_matched_by_target()); + } + + #[test] + fn merge_into_clause_kind_canonical_collapses_not_matched() { + assert_eq!( + MergeIntoClauseKind::NotMatched.canonical(), + MergeIntoClauseKind::NotMatchedByTarget + ); + assert_eq!( + MergeIntoClauseKind::NotMatchedByTarget.canonical(), + MergeIntoClauseKind::NotMatchedByTarget + ); + assert_eq!( + MergeIntoClauseKind::Matched.canonical(), + MergeIntoClauseKind::Matched + ); + assert_eq!( + MergeIntoClauseKind::NotMatchedBySource.canonical(), + MergeIntoClauseKind::NotMatchedBySource + ); + } + + #[test] + fn merge_into_op_exprs_round_trip() { + let op = MergeIntoOp { + on: col("id").eq(col("source_id")), + clauses: vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: Some(col("qty").gt(lit(0_i64))), + action: MergeIntoAction::Update(vec![ + ("qty".to_string(), col("source_qty")), + ("price".to_string(), col("source_price")), + ]), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec!["id".to_string(), "qty".to_string()], + values: vec![col("source_id"), col("source_qty")], + }, + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatchedBySource, + predicate: Some(col("active").eq(lit(true))), + action: MergeIntoAction::Delete, + }, + ], + }; + let exprs = op.exprs(); + assert_eq!(exprs.len(), 7); + + let owned: Vec = exprs.into_iter().cloned().collect(); + let rebuilt = op.with_new_exprs(owned).unwrap(); + assert_eq!(op, rebuilt); + } + + #[test] + fn merge_into_op_with_new_exprs_length_mismatch() { + let op = MergeIntoOp { + on: col("id").eq(col("source_id")), + clauses: vec![], + }; + let err = op.with_new_exprs(vec![]).unwrap_err(); + assert!( + err.to_string().contains("expected 1 expressions, got 0"), + "unexpected error: {err}" + ); + } +} diff --git a/datafusion/expr/src/logical_plan/extension.rs b/datafusion/expr/src/logical_plan/extension.rs index fe324d40fd952..e1ee273968676 100644 --- a/datafusion/expr/src/logical_plan/extension.rs +++ b/datafusion/expr/src/logical_plan/extension.rs @@ -314,7 +314,7 @@ pub trait UserDefinedLogicalNodeCore: } } -/// Automatically derive UserDefinedLogicalNode to `UserDefinedLogicalNode` +/// Automatically derive `UserDefinedLogicalNode` from `UserDefinedLogicalNodeCore` /// to avoid boiler plate for implementing `as_any`, `Hash`, `PartialEq` and `PartialOrd`. impl UserDefinedLogicalNode for T { fn as_any(&self) -> &dyn Any { diff --git a/datafusion/expr/src/logical_plan/invariants.rs b/datafusion/expr/src/logical_plan/invariants.rs index 0889afd08fee4..d6867d1ceb112 100644 --- a/datafusion/expr/src/logical_plan/invariants.rs +++ b/datafusion/expr/src/logical_plan/invariants.rs @@ -21,7 +21,7 @@ use datafusion_common::{ }; use crate::{ - Aggregate, Expr, Filter, Join, JoinType, LogicalPlan, Window, + Aggregate, DmlStatement, Expr, Filter, Join, JoinType, LogicalPlan, Window, WriteOp, expr::{Exists, InSubquery, SetComparison}, expr_rewriter::strip_outer_reference, utils::{collect_subquery_cols, split_conjunction}, @@ -253,7 +253,11 @@ pub fn check_subquery_expr( | LogicalPlan::TableScan(_) | LogicalPlan::Window(_) | LogicalPlan::Aggregate(_) - | LogicalPlan::Join(_) => Ok(()), + | LogicalPlan::Join(_) + | LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + .. + }) => Ok(()), _ => plan_err!( "In/Exist/SetComparison subquery can only be used in \ Projection, Filter, TableScan, Window functions, Aggregate and Join plan nodes, \ diff --git a/datafusion/expr/src/logical_plan/mod.rs b/datafusion/expr/src/logical_plan/mod.rs index c2b01868c97f3..4766c3f33379f 100644 --- a/datafusion/expr/src/logical_plan/mod.rs +++ b/datafusion/expr/src/logical_plan/mod.rs @@ -36,14 +36,17 @@ pub use ddl::{ CreateFunctionBody, CreateIndex, CreateMemoryTable, CreateView, DdlStatement, DropCatalogSchema, DropFunction, DropTable, DropView, OperateFunctionArg, }; -pub use dml::{DmlStatement, WriteOp}; +pub use dml::{ + DmlStatement, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, + WriteOp, +}; pub use plan::{ Aggregate, Analyze, ColumnUnnestList, DescribeTable, Distinct, DistinctOn, EmptyRelation, Explain, ExplainOption, Extension, FetchType, Filter, Join, JoinConstraint, JoinType, Limit, LogicalPlan, Partitioning, PlanType, Projection, - RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, Subquery, - SubqueryAlias, TableScan, ToStringifiedPlan, Union, Unnest, Values, Window, - projection_schema, + RangePartitioning, RecursiveQuery, Repartition, SkipType, Sort, StringifiedPlan, + Subquery, SubqueryAlias, TableScan, TableScanBuilder, ToStringifiedPlan, Union, + Unnest, Values, Window, projection_schema, }; pub use statement::{ Deallocate, Execute, Prepare, ResetVariable, SetVariable, Statement, diff --git a/datafusion/expr/src/logical_plan/plan.rs b/datafusion/expr/src/logical_plan/plan.rs index c572b202f03ce..1a141ea52a13a 100644 --- a/datafusion/expr/src/logical_plan/plan.rs +++ b/datafusion/expr/src/logical_plan/plan.rs @@ -18,7 +18,7 @@ //! Logical plan types use std::cmp::Ordering; -use std::collections::{HashMap, HashSet}; +use std::collections::{BTreeSet, HashMap, HashSet}; use std::fmt::{self, Debug, Display, Formatter}; use std::hash::{Hash, Hasher}; use std::sync::{Arc, LazyLock}; @@ -39,10 +39,11 @@ use crate::expr_rewriter::{ }; use crate::logical_plan::display::{GraphvizVisitor, IndentVisitor}; use crate::logical_plan::extension::UserDefinedLogicalNode; -use crate::logical_plan::{DmlStatement, Statement}; +use crate::logical_plan::{DmlStatement, Statement, WriteOp}; use crate::utils::{ - enumerate_grouping_sets, exprlist_to_fields, find_out_reference_exprs, - grouping_set_expr_count, grouping_set_to_exprlist, merge_schema, split_conjunction, + check_aggregate_and_window_nesting, enumerate_grouping_sets, exprlist_to_fields, + find_out_reference_exprs, grouping_set_expr_count, grouping_set_to_exprlist, + merge_schema, split_conjunction, }; use crate::{ BinaryExpr, CreateMemoryTable, CreateView, Execute, Expr, ExprSchemable, GroupingSet, @@ -50,9 +51,11 @@ use crate::{ WindowFunctionDefinition, build_join_schema, expr_vec_fmt, requalify_sides_if_needed, }; +use crate::statistics::StatisticsRequest; +use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, FieldRef, Schema, SchemaRef}; use datafusion_common::cse::{NormalizeEq, Normalizeable}; -use datafusion_common::format::ExplainFormat; +use datafusion_common::format::{ExplainAnalyzeCategories, ExplainFormat, MetricType}; use datafusion_common::metadata::check_metadata_with_storage_equal; use datafusion_common::tree_node::{ Transformed, TreeNode, TreeNodeContainer, TreeNodeRecursion, @@ -60,10 +63,12 @@ use datafusion_common::tree_node::{ use datafusion_common::{ Column, Constraints, DFSchema, DFSchemaRef, DataFusionError, Dependency, FunctionalDependence, FunctionalDependencies, NullEquality, ParamValues, Result, - ScalarValue, Spans, TableReference, UnnestOptions, aggregate_functional_dependencies, - assert_eq_or_internal_err, assert_or_internal_err, internal_err, plan_err, + ScalarValue, Spans, SplitPoint, TableReference, UnnestOptions, + aggregate_functional_dependencies, assert_eq_or_internal_err, assert_or_internal_err, + internal_err, plan_err, validate_range_split_points, }; use indexmap::IndexSet; +use itertools::Itertools as _; // backwards compatibility use crate::display::PgJsonVisitor; @@ -353,10 +358,7 @@ impl LogicalPlan { LogicalPlan::Copy(CopyTo { output_schema, .. }) => output_schema, LogicalPlan::Ddl(ddl) => ddl.schema(), LogicalPlan::Unnest(Unnest { schema, .. }) => schema, - LogicalPlan::RecursiveQuery(RecursiveQuery { static_term, .. }) => { - // we take the schema of the static term as the schema of the entire recursive query - static_term.schema() - } + LogicalPlan::RecursiveQuery(RecursiveQuery { schema, .. }) => schema, } } @@ -740,7 +742,14 @@ impl LogicalPlan { }; Ok(LogicalPlan::Distinct(distinct)) } - LogicalPlan::RecursiveQuery(_) => Ok(self), + LogicalPlan::RecursiveQuery(RecursiveQuery { + name, + static_term, + recursive_term, + is_distinct, + schema: _, + }) => RecursiveQuery::try_new(name, static_term, recursive_term, is_distinct) + .map(LogicalPlan::RecursiveQuery), LogicalPlan::Analyze(_) => Ok(self), LogicalPlan::Explain(_) => Ok(self), LogicalPlan::TableScan(_) => Ok(self), @@ -802,12 +811,20 @@ impl LogicalPlan { op, .. }) => { - self.assert_no_expressions(expr)?; let input = self.only_input(inputs)?; + let op = match op { + WriteOp::MergeInto(merge_op) => { + WriteOp::MergeInto(Box::new(merge_op.with_new_exprs(expr)?)) + } + other => { + self.assert_no_expressions(expr)?; + other.clone() + } + }; Ok(LogicalPlan::Dml(DmlStatement::new( table_name.clone(), Arc::clone(target), - op.clone(), + op, Arc::new(input), ))) } @@ -864,6 +881,32 @@ impl LogicalPlan { input: Arc::new(input), })) } + Partitioning::Range(range) => { + if expr.len() != range.ordering().len() { + return internal_err!( + "Incorrect number of expressions for Range partitioning" + ); + } + let input = self.only_input(inputs)?; + let ordering = range + .ordering() + .iter() + .zip(expr) + .map(|(sort_expr, expr)| SortExpr { + expr, + asc: sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }) + .collect(); + let range = RangePartitioning::try_new( + ordering, + range.split_points().to_vec(), + )?; + Ok(LogicalPlan::Repartition(Repartition { + partitioning_scheme: Partitioning::Range(range), + input: Arc::new(input), + })) + } Partitioning::DistributeBy(_) => { let input = self.only_input(inputs)?; Ok(LogicalPlan::Repartition(Repartition { @@ -1080,20 +1123,24 @@ impl LogicalPlan { }) => { self.assert_no_expressions(expr)?; let (static_term, recursive_term) = self.only_two_inputs(inputs)?; - Ok(LogicalPlan::RecursiveQuery(RecursiveQuery { - name: name.clone(), - static_term: Arc::new(static_term), - recursive_term: Arc::new(recursive_term), - is_distinct: *is_distinct, - })) + RecursiveQuery::try_new( + name.clone(), + Arc::new(static_term), + Arc::new(recursive_term), + *is_distinct, + ) + .map(LogicalPlan::RecursiveQuery) } LogicalPlan::Analyze(a) => { self.assert_no_expressions(expr)?; let input = self.only_input(inputs)?; Ok(LogicalPlan::Analyze(Analyze { verbose: a.verbose, + format: a.format.clone(), schema: Arc::clone(&a.schema), input: Arc::new(input), + analyze_level: a.analyze_level, + analyze_categories: a.analyze_categories.clone(), })) } LogicalPlan::Explain(e) => { @@ -1106,6 +1153,7 @@ impl LogicalPlan { stringified_plans: e.stringified_plans.clone(), schema: Arc::clone(&e.schema), logical_optimization_succeeded: e.logical_optimization_succeeded, + show_statistics: e.show_statistics, })) } LogicalPlan::Statement(Statement::Prepare(Prepare { @@ -1147,12 +1195,20 @@ impl LogicalPlan { options, .. }) => { - self.assert_no_expressions(expr)?; + let exec_columns = if expr.is_empty() { + columns.clone() + } else { + expr.into_iter() + .map(|e| match e { + Expr::Column(c) => Ok(c), + other => internal_err!( + "Expected Expr::Column for Unnest exec_columns, got {other:?}" + ), + }) + .collect::>>()? + }; let input = self.only_input(inputs)?; - // Update schema with unnested column type. - let new_plan = - unnest_with_options(input, columns.clone(), options.clone())?; - Ok(new_plan) + Ok(unnest_with_options(input, exec_columns, options.clone())?) } } } @@ -1686,7 +1742,7 @@ impl LogicalPlan { /// ``` pub fn display_indent(&self) -> impl Display + '_ { // Boilerplate structure to wrap LogicalPlan with something - // that that can be formatted + // that can be formatted struct Wrapper<'a>(&'a LogicalPlan); impl Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -1732,7 +1788,7 @@ impl LogicalPlan { /// ``` pub fn display_indent_schema(&self) -> impl Display + '_ { // Boilerplate structure to wrap LogicalPlan with something - // that that can be formatted + // that can be formatted struct Wrapper<'a>(&'a LogicalPlan); impl Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -1752,7 +1808,7 @@ impl LogicalPlan { /// Users can use this format to visualize the plan in existing plan visualization tools, for example [dalibo](https://explain.dalibo.com/) pub fn display_pg_json(&self) -> impl Display + '_ { // Boilerplate structure to wrap LogicalPlan with something - // that that can be formatted + // that can be formatted struct Wrapper<'a>(&'a LogicalPlan); impl Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -1798,7 +1854,7 @@ impl LogicalPlan { /// ``` pub fn display_graphviz(&self) -> impl Display + '_ { // Boilerplate structure to wrap LogicalPlan with something - // that that can be formatted + // that can be formatted struct Wrapper<'a>(&'a LogicalPlan); impl Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -1849,7 +1905,7 @@ impl LogicalPlan { /// ``` pub fn display(&self) -> impl Display + '_ { // Boilerplate structure to wrap LogicalPlan with something - // that that can be formatted + // that can be formatted struct Wrapper<'a>(&'a LogicalPlan); impl Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -2034,6 +2090,7 @@ impl LogicalPlan { filter, join_constraint, join_type, + null_aware, .. }) => { let join_expr: Vec = @@ -2042,6 +2099,8 @@ impl LogicalPlan { .as_ref() .map(|expr| format!(" Filter: {expr}")) .unwrap_or_else(|| "".to_string()); + let null_aware_expr = + if *null_aware { " null_aware" } else { "" }; let join_type = if filter.is_none() && keys.is_empty() && *join_type == JoinType::Inner @@ -2061,15 +2120,17 @@ impl LogicalPlan { filter_expr )?; } + write!(f, "{null_aware_expr}")?; Ok(()) } JoinConstraint::Using => { write!( f, - "{} Join: Using {}{}", + "{} Join: Using {}{}{}", join_type, join_expr.join(", "), filter_expr, + null_aware_expr, ) } } @@ -2091,6 +2152,9 @@ impl LogicalPlan { n ) } + Partitioning::Range(range) => { + write!(f, "Repartition: {range}") + } Partitioning::DistributeBy(expr) => { let dist_by_expr: Vec = expr.iter().map(|e| format!("{e}")).collect(); @@ -2167,8 +2231,7 @@ impl LogicalPlan { .map(|(i, unnest_info)| { format!( "{}|depth={}", - &input_columns[*i].to_string(), - unnest_info.depth + input_columns[*i], unnest_info.depth ) }) .collect::>(); @@ -2257,7 +2320,7 @@ impl PartialOrd for EmptyRelation { /// intermediate table, then empty the intermediate table. /// /// [Postgres Docs]: https://www.postgresql.org/docs/current/queries-with.html#QUERIES-WITH-RECURSIVE -#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +#[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct RecursiveQuery { /// Name of the query pub name: String, @@ -2269,6 +2332,90 @@ pub struct RecursiveQuery { /// Should the output of the recursive term be deduplicated (`UNION`) or /// not (`UNION ALL`). pub is_distinct: bool, + /// Schema exposed to parent plans after reconciling the static and recursive terms. + pub schema: DFSchemaRef, +} + +impl PartialOrd for RecursiveQuery { + fn partial_cmp(&self, other: &Self) -> Option { + match self.name.partial_cmp(&other.name) { + Some(Ordering::Equal) => { + match self.static_term.partial_cmp(&other.static_term) { + Some(Ordering::Equal) => { + match self.recursive_term.partial_cmp(&other.recursive_term) { + Some(Ordering::Equal) => { + self.is_distinct.partial_cmp(&other.is_distinct) + } + cmp => cmp, + } + } + cmp => cmp, + } + } + cmp => cmp, + } + // If the query definition compares equal but the derived schema differs, + // return `None` instead of contradicting `PartialEq` with `Some(Equal)`. + // TODO (https://github.com/apache/datafusion/issues/17477) avoid recomparing all fields + .filter(|cmp| *cmp != Ordering::Equal || self == other) + } +} + +impl RecursiveQuery { + pub fn try_new( + name: String, + static_term: Arc, + recursive_term: Arc, + is_distinct: bool, + ) -> Result { + let schema = + recursive_query_output_schema(static_term.schema(), recursive_term.schema())?; + Ok(Self { + name, + static_term, + recursive_term, + is_distinct, + schema, + }) + } +} + +/// Compute a recursive query's output schema by considering both its static and +/// recursive terms. +/// +/// Field names, types, and metadata come from the static term. A field is +/// nullable if either the static or the recursive term produces a nullable +/// value in that position, matching how `UNION` reconciles branch nullability. +/// +/// Functional dependencies are intentionally dropped: the recursive term +/// appends rows that can duplicate values the static term guarantees unique, so +/// any FDs carried by the static term may not hold over the combined output. +fn recursive_query_output_schema( + static_schema: &DFSchemaRef, + recursive_schema: &DFSchemaRef, +) -> Result { + if static_schema.fields().len() != recursive_schema.fields().len() { + return Err(DataFusionError::Plan(format!( + "Non-recursive term and recursive term must have the same number of columns ({} != {})", + static_schema.fields().len(), + recursive_schema.fields().len() + ))); + } + + let fields = static_schema + .iter() + .zip(recursive_schema.fields()) + .map(|((qualifier, static_field), recursive_field)| { + let nullable = static_field.is_nullable() || recursive_field.is_nullable(); + ( + qualifier.cloned(), + static_field.as_ref().clone().with_nullable(nullable).into(), + ) + }) + .collect::>(); + + DFSchema::new_with_metadata(fields, static_schema.metadata().clone()) + .map(DFSchemaRef::new) } /// Values expression. See @@ -2495,6 +2642,19 @@ pub struct Filter { } impl Filter { + /// Create a new filter operator. + /// + /// Skips the type-checking and dealiasing done in [Self::try_new]. + /// For internal use in DataFusion only. + /// + /// **Preconditions:** + /// - the `predicate` expression returns a boolean value + /// - the `predicate` expression is not aliased + #[doc(hidden)] + pub fn new(predicate: Expr, input: Arc) -> Self { + Self { predicate, input } + } + /// Create a new filter operator. /// /// Notes: as Aliases have no effect on the output of a filter operator, @@ -2503,13 +2663,6 @@ impl Filter { Self::try_new_internal(predicate, input) } - /// Create a new filter operator for a having clause. - /// This is similar to a filter, but its having flag is set to true. - #[deprecated(since = "48.0.0", note = "Use `try_new` instead")] - pub fn try_new_with_having(predicate: Expr, input: Arc) -> Result { - Self::try_new_internal(predicate, input) - } - fn is_allowed_filter_type(data_type: &DataType) -> bool { match data_type { // Interpret NULL as a missing boolean value. @@ -2635,6 +2788,11 @@ pub struct Window { impl Window { /// Create a new window operator. pub fn try_new(window_expr: Vec, input: Arc) -> Result { + // Reject e.g. `sum(sum(x) OVER ()) OVER ()` here rather than letting it + // reach physical planning, which has no equivalent for a nested window + // function. + check_aggregate_and_window_nesting(window_expr.iter())?; + let fields: Vec<(Option, Arc)> = input .schema() .iter() @@ -2780,6 +2938,12 @@ pub struct TableScan { pub filters: Vec, /// Optional number of rows to read pub fetch: Option, + /// Statistics the planner would like the provider to answer for this + /// scan, typically attached by a custom optimizer rule from the + /// surrounding plan (e.g. Min/Max for sort keys). + /// + /// A [`BTreeSet`], not a `Vec` to keep the resulting plan deterministic. + pub statistics_requests: BTreeSet, } impl Debug for TableScan { @@ -2854,6 +3018,7 @@ impl Hash for TableScan { impl TableScan { /// Initialize TableScan with appropriate schema from the given /// arguments. + #[deprecated(since = "54.0.0", note = "use `TableScanBuilder` instead")] pub fn try_new( table_name: impl Into, table_source: Arc, @@ -2861,14 +3026,92 @@ impl TableScan { filters: Vec, fetch: Option, ) -> Result { - let table_name = table_name.into(); + TableScanBuilder::new(table_name, table_source) + .with_projection(projection) + .with_filters(filters) + .with_fetch(fetch) + .build() + } +} + +/// Builder for [`TableScan`]. +/// +/// Prefer this over constructing a [`TableScan`] directly: it derives the +/// `projected_schema` from the source schema and projection, and is resilient +/// to new fields being added to [`TableScan`]. An existing scan can be turned +/// back into a builder with `TableScanBuilder::from(scan)`, tweaked, and +/// rebuilt with [`TableScanBuilder::build`]. +pub struct TableScanBuilder { + table_name: TableReference, + source: Arc, + projection: Option>, + filters: Vec, + fetch: Option, + statistics_requests: BTreeSet, +} + +impl TableScanBuilder { + /// Create a new builder for a scan of `source` named `table_name`. + pub fn new( + table_name: impl Into, + source: Arc, + ) -> Self { + Self { + table_name: table_name.into(), + source, + projection: None, + filters: vec![], + fetch: None, + statistics_requests: BTreeSet::new(), + } + } + + /// Set the column projection (indices into the source schema). + pub fn with_projection(mut self, projection: Option>) -> Self { + self.projection = projection; + self + } + + /// Set the filter expressions offered to the table provider. + pub fn with_filters(mut self, filters: Vec) -> Self { + self.filters = filters; + self + } + + /// Set the maximum number of rows to read. + pub fn with_fetch(mut self, fetch: Option) -> Self { + self.fetch = fetch; + self + } + + /// Set the statistics requests for the scan. See + /// [`TableScan::statistics_requests`]. + pub fn with_statistics_requests( + mut self, + statistics_requests: BTreeSet, + ) -> Self { + self.statistics_requests = statistics_requests; + self + } + + /// Build the [`TableScan`], deriving its `projected_schema` from the + /// source schema and projection. + pub fn build(self) -> Result { + let TableScanBuilder { + table_name, + source, + projection, + filters, + fetch, + statistics_requests, + } = self; if table_name.table().is_empty() { return plan_err!("table_name cannot be empty"); } - let schema = table_source.schema(); + let schema = source.schema(); let func_dependencies = FunctionalDependencies::new_from_constraints( - table_source.constraints(), + source.constraints(), schema.fields.len(), ); let projected_schema = projection @@ -2894,17 +3137,31 @@ impl TableScan { })?; let projected_schema = Arc::new(projected_schema); - Ok(Self { + Ok(TableScan { table_name, - source: table_source, + source, projection, projected_schema, filters, fetch, + statistics_requests, }) } } +impl From for TableScanBuilder { + fn from(scan: TableScan) -> Self { + Self { + table_name: scan.table_name, + source: scan.source, + projection: scan.projection, + filters: scan.filters, + fetch: scan.fetch, + statistics_requests: scan.statistics_requests, + } + } +} + // Repartition the plan based on a partitioning scheme. #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub struct Repartition { @@ -3216,6 +3473,15 @@ pub struct ExplainOption { pub analyze: bool, /// Output syntax/format pub format: ExplainFormat, + /// Statement-level override for `datafusion.explain.show_statistics`. + /// `None` means "fall back to session config". + pub show_statistics: Option, + /// Statement-level override for `datafusion.explain.analyze_level`. + /// `None` means "fall back to session config". + pub analyze_level: Option, + /// Statement-level override for `datafusion.explain.analyze_categories`. + /// `None` means "fall back to session config". + pub analyze_categories: Option, } impl Default for ExplainOption { @@ -3224,6 +3490,9 @@ impl Default for ExplainOption { verbose: false, analyze: false, format: ExplainFormat::Indent, + show_statistics: None, + analyze_level: None, + analyze_categories: None, } } } @@ -3246,6 +3515,30 @@ impl ExplainOption { self.format = format; self } + + /// Builder-style setter for a statement-level override of + /// `datafusion.explain.show_statistics`. + pub fn with_show_statistics(mut self, show_statistics: Option) -> Self { + self.show_statistics = show_statistics; + self + } + + /// Builder-style setter for a statement-level override of + /// `datafusion.explain.analyze_level`. + pub fn with_analyze_level(mut self, analyze_level: Option) -> Self { + self.analyze_level = analyze_level; + self + } + + /// Builder-style setter for a statement-level override of + /// `datafusion.explain.analyze_categories`. + pub fn with_analyze_categories( + mut self, + analyze_categories: Option, + ) -> Self { + self.analyze_categories = analyze_categories; + self + } } /// Produces a relation with string representations of @@ -3269,6 +3562,9 @@ pub struct Explain { pub schema: DFSchemaRef, /// Used by physical planner to check if should proceed with planning pub logical_optimization_succeeded: bool, + /// Statement-level override for `datafusion.explain.show_statistics`. + /// When `None`, the session-config value is used. + pub show_statistics: Option, } // Manual implementation needed because of `schema` field. Comparison excludes this field. @@ -3284,18 +3580,22 @@ impl PartialOrd for Explain { pub stringified_plans: &'a Vec, /// Used by physical planner to check if should proceed with planning pub logical_optimization_succeeded: &'a bool, + /// Statement-level override for show_statistics + pub show_statistics: &'a Option, } let comparable_self = ComparableExplain { verbose: &self.verbose, plan: &self.plan, stringified_plans: &self.stringified_plans, logical_optimization_succeeded: &self.logical_optimization_succeeded, + show_statistics: &self.show_statistics, }; let comparable_other = ComparableExplain { verbose: &other.verbose, plan: &other.plan, stringified_plans: &other.stringified_plans, logical_optimization_succeeded: &other.logical_optimization_succeeded, + show_statistics: &other.show_statistics, }; comparable_self .partial_cmp(&comparable_other) @@ -3310,13 +3610,24 @@ impl PartialOrd for Explain { pub struct Analyze { /// Should extra detail be included? pub verbose: bool, + /// Output syntax/format for the rendered physical plan + metrics. + pub format: ExplainFormat, /// The logical plan that is being EXPLAIN ANALYZE'd pub input: Arc, /// The output schema of the explain (2 columns of text) pub schema: DFSchemaRef, + /// Statement-level override for `datafusion.explain.analyze_level`. + /// When `None`, the session-config value is used. + pub analyze_level: Option, + /// Statement-level override for `datafusion.explain.analyze_categories`. + /// When `None`, the session-config value is used. + pub analyze_categories: Option, } -// Manual implementation needed because of `schema` field. Comparison excludes this field. +// Manual implementation needed because of `schema` field and the lack of +// `PartialOrd` on `MetricType` / `ExplainAnalyzeCategories`. Ordering is +// defined over `(verbose, input)` and then falls back to `==` for the +// remaining statement-level override fields. impl PartialOrd for Analyze { fn partial_cmp(&self, other: &Self) -> Option { match self.verbose.partial_cmp(&other.verbose) { @@ -3595,6 +3906,10 @@ impl Aggregate { group_expr: Vec, aggr_expr: Vec, ) -> Result { + // Reject e.g. `sum(sum(x))` here rather than letting it reach physical + // planning, which has no equivalent for a nested aggregate. + check_aggregate_and_window_nesting(group_expr.iter().chain(aggr_expr.iter()))?; + let group_expr = enumerate_grouping_sets(group_expr)?; let is_grouping_set = matches!(group_expr.as_slice(), [Expr::GroupingSet(_)]); @@ -3839,8 +4154,12 @@ fn calc_func_dependencies_for_project( exprs: &[Expr], input: &LogicalPlan, ) -> Result { + // Sentinel for projection outputs that do not map back to any input field. + const COMPUTED_EXPR_INDEX: usize = usize::MAX; + let input_fields = input.schema().field_names(); - // Calculate expression indices (if present) in the input schema. + // Map each projection output position to its input column index. + // A projection expression can produce multiple output columns, such as `*`. let proj_indices = exprs .iter() .map(|expr| match expr { @@ -3856,30 +4175,33 @@ fn calc_func_dependencies_for_project( Ok::<_, DataFusionError>( wildcard_fields .into_iter() - .filter_map(|(qualifier, f)| { + .map(|(qualifier, f)| { let flat_name = qualifier .map(|t| format!("{}.{}", t, f.name())) .unwrap_or_else(|| f.name().clone()); - input_fields.iter().position(|item| *item == flat_name) + input_fields + .iter() + .position(|item| *item == flat_name) + .unwrap_or(COMPUTED_EXPR_INDEX) }) .collect::>(), ) } Expr::Alias(alias) => { let name = format!("{}", alias.expr); - Ok(input_fields + let input_index = input_fields .iter() .position(|item| *item == name) - .map(|i| vec![i]) - .unwrap_or(vec![])) + .unwrap_or(COMPUTED_EXPR_INDEX); + Ok(vec![input_index]) } _ => { let name = format!("{expr}"); - Ok(input_fields + let input_index = input_fields .iter() .position(|item| *item == name) - .map(|i| vec![i]) - .unwrap_or(vec![])) + .unwrap_or(COMPUTED_EXPR_INDEX); + Ok(vec![input_index]) } }) .collect::>>()? @@ -4137,11 +4459,16 @@ impl Debug for Subquery { } } -/// Logical partitioning schemes supported by [`LogicalPlan::Repartition`] +/// Logical partitioning schemes. /// -/// See [`Partitioning`] for more details on partitioning +/// A scheme can describe either requested repartitioning in +/// [`LogicalPlan::Repartition`] or a partitioning property declared by a source. +/// Some schemes are only valid as metadata until planner support is added. /// -/// [`Partitioning`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/enum.Partitioning.html# +/// For physical execution partitioning, see +/// [`datafusion_physical_expr::Partitioning`]. +/// +/// [`datafusion_physical_expr::Partitioning`]: https://docs.rs/datafusion/latest/datafusion/physical_expr/enum.Partitioning.html# #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] pub enum Partitioning { /// Allocate batches using a round-robin algorithm and the specified number of partitions @@ -4149,10 +4476,118 @@ pub enum Partitioning { /// Allocate rows based on a hash of one of more expressions and the specified number /// of partitions. Hash(Vec, usize), + /// Partition rows by ranges. + /// See [`RangePartitioning`] for the logical contract. + Range(RangePartitioning), /// The DISTRIBUTE BY clause is used to repartition the data based on the input expressions DistributeBy(Vec), } +impl Partitioning { + /// Return the number of partitions, if known. + pub fn partition_count(&self) -> Option { + match self { + Self::RoundRobinBatch(partition_count) | Self::Hash(_, partition_count) => { + Some(*partition_count) + } + Self::Range(range) => Some(range.partition_count()), + Self::DistributeBy(_) => None, + } + } +} + +/// Logical range partitioning. +/// +/// [`RangePartitioning`] describes an ordered logical key space with split points. +/// +/// - `ordering` defines the partitioning key and ordering using logical +/// [`SortExpr`]s. +/// - `split_points` define the boundaries between adjacent partitions. +/// +/// Comparisons use the lexicographic order defined by `ordering`, +/// including `ASC`/`DESC` and null ordering. Split points must be ordered +/// according to that ordering, and each split point must have one value per +/// ordering expression. See [`SplitPoint`] for the shared boundary contract. +/// +/// The expressions are resolved against the declaring plan's schema. This +/// constructor does not validate split point value types against the resolved +/// expression types. Like other user-specified data properties such as +/// sortedness, if a source declares range partitioning, it is responsible for +/// placing each row in the partition described by the split points. DataFusion +/// will not validate this is upheld. +/// +/// NOTE: Range-aware optimizer and execution behavior will be introduced +/// incrementally. See +/// . +#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash)] +pub struct RangePartitioning { + /// Ordered logical partitioning key. + ordering: Vec, + /// Boundaries between adjacent partitions. + split_points: Vec, +} + +impl RangePartitioning { + /// Creates logical range partitioning metadata and validates split point + /// shape and ordering. + pub fn try_new( + ordering: Vec, + split_points: Vec, + ) -> Result { + if ordering.is_empty() { + return plan_err!("Range partitioning requires non-empty ordering"); + } + + validate_range_split_points(&split_points, &logical_sort_options(&ordering))?; + + Ok(Self { + ordering, + split_points, + }) + } + + /// Return the number of partitions. + pub fn partition_count(&self) -> usize { + self.split_points.len() + 1 + } + + /// Returns the ordering that defines the range key. + pub fn ordering(&self) -> &[SortExpr] { + &self.ordering + } + + /// Returns the ordered split points between partitions. + pub fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } +} + +fn logical_sort_options(ordering: &[SortExpr]) -> Vec { + ordering + .iter() + .map(|sort_expr| SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }) + .collect() +} + +impl Display for RangePartitioning { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + let ordering = self.ordering().iter().map(ToString::to_string).join(", "); + let split_points = self + .split_points() + .iter() + .map(ToString::to_string) + .join(", "); + write!( + f, + "Range([{ordering}], [{split_points}], {})", + self.partition_count() + ) + } +} + /// Represent the unnesting operation on a list column, such as the recursion depth and /// the output column name after unnesting /// @@ -4480,6 +4915,45 @@ mod tests { use insta::{assert_debug_snapshot, assert_snapshot}; use std::hash::DefaultHasher; + /// `LogicalPlan` is moved/swapped on every step of the planning hot path + /// (every `mem::take` in an in-place rewriter, every `Arc` + /// write, every owned `map_*` traversal). Its size is set by the largest + /// variant, so an oversized variant balloons cost for every other variant. + /// + /// Today the size-setter should be `Join` (~176 bytes); `DdlStatement` is + /// boxed precisely so it does not dominate. If you grow a variant, please + /// box the new large fields rather than letting this number creep up — + /// see the analogous `test_size_of_expr` in `expr.rs`. + #[test] + fn test_size_of_logical_plan() { + // `LogicalPlan` enum on aarch64 / x86_64. Today this matches + // `Join`'s 176 bytes (the enum discriminant fits in `Join`'s + // alignment padding); if `Join` grows or another variant overtakes + // it, this number will move with the new size-setter. + assert_eq!(size_of::(), 176); + // `DdlStatement` is `Ddl(DdlStatement)`'s payload; keep it below the + // `Join` ceiling so it never re-becomes the size-setter. + assert!( + size_of::() < size_of::(), + "DdlStatement ({} bytes) should stay smaller than Join ({} bytes); \ + box the new large variant rather than letting it dominate `LogicalPlan`.", + size_of::(), + size_of::(), + ); + // Sanity check the two boxed variants stay boxed (so the payload + // sits on the heap, not in the enum). + assert_eq!( + size_of::>(), + 8, + "CreateExternalTable should be Box'd inside DdlStatement" + ); + assert_eq!( + size_of::>(), + 8, + "CreateFunction should be Box'd inside DdlStatement" + ); + } + fn employee_schema() -> Schema { Schema::new(vec![ Field::new("id", DataType::Int32, false), @@ -4490,6 +4964,210 @@ mod tests { ]) } + #[test] + fn projection_with_leading_computed_column_preserves_pk() -> Result<()> { + let constraints = + Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]); + let source = Arc::new( + LogicalTableSource::new(Arc::new(employee_schema())) + .with_constraints(constraints), + ); + let plan = LogicalPlanBuilder::scan("employee_csv", source, None)? + .project(vec![ + lit(1i32).alias("__common_expr_1"), + col("id"), + col("first_name"), + col("salary"), + ])? + .build()?; + + let deps = plan.schema().functional_dependencies(); + assert_eq!(deps.len(), 1); + assert_eq!(deps[0].source_indices, vec![1]); + + Ok(()) + } + + #[test] + fn projection_with_leading_computed_column_and_wildcard_preserves_pk() -> Result<()> { + let constraints = + Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]); + let source = Arc::new( + LogicalTableSource::new(Arc::new(employee_schema())) + .with_constraints(constraints), + ); + let plan = LogicalPlanBuilder::scan("employee_csv", source, None)? + .project(vec![ + SelectExpr::Expression(lit(1i32).alias("__common_expr_1")), + SelectExpr::Wildcard(Default::default()), + ])? + .build()?; + + let deps = plan.schema().functional_dependencies(); + assert_eq!(plan.schema().fields().len(), 6); + assert_eq!(deps.len(), 1); + assert_eq!(deps[0].source_indices, vec![1]); + assert_eq!(deps[0].target_indices, vec![0, 1, 2, 3, 4, 5]); + + Ok(()) + } + + #[test] + fn projection_with_wildcard_expr_before_pk_preserves_pk() -> Result<()> { + let constraints = + Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]); + let source = Arc::new( + LogicalTableSource::new(Arc::new(employee_schema())) + .with_constraints(constraints), + ); + let input = LogicalPlanBuilder::scan("employee_csv", source, None)?.build()?; + #[expect(deprecated)] + let projection = Projection::try_new( + vec![ + Expr::Wildcard { + qualifier: None, + options: Box::new(crate::expr::WildcardOptions::default()), + }, + col("employee_csv.id"), + ], + Arc::new(input), + )?; + + let deps = projection.schema.functional_dependencies(); + assert_eq!(deps.len(), 1); + assert_eq!(deps[0].source_indices, vec![1]); + + Ok(()) + } + + fn i32_split_point(value: i32) -> SplitPoint { + SplitPoint::new(vec![ScalarValue::Int32(Some(value))]) + } + + fn null_i32_split_point() -> SplitPoint { + SplitPoint::new(vec![ScalarValue::Int32(None)]) + } + + #[test] + fn logical_range_partitioning_validates_shape() { + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10), i32_split_point(20)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + + let range = RangePartitioning::try_new( + vec![col("id").sort(false, true)], + vec![i32_split_point(20), i32_split_point(10)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + + let err = RangePartitioning::try_new(vec![], vec![]).unwrap_err(); + assert!(err.to_string().contains("non-empty ordering")); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true), col("salary").sort(true, true)], + vec![i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split point 0 has width 1, but ordering has width 2") + ); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(20), i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split points must be strictly ordered") + ); + + let err = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10), i32_split_point(10)], + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("split points must be strictly ordered") + ); + + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![null_i32_split_point(), i32_split_point(10)], + ) + .unwrap(); + assert_eq!(range.partition_count(), 3); + } + + #[test] + fn logical_partitioning_reports_known_partition_count() -> Result<()> { + let range = RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10)], + )?; + + assert_eq!(Partitioning::RoundRobinBatch(4).partition_count(), Some(4)); + assert_eq!( + Partitioning::Hash(vec![col("id")], 8).partition_count(), + Some(8) + ); + assert_eq!(Partitioning::Range(range).partition_count(), Some(2)); + assert_eq!( + Partitioning::DistributeBy(vec![col("id")]).partition_count(), + None + ); + + Ok(()) + } + + #[test] + fn logical_range_partitioning_participates_in_expression_rewrite() -> Result<()> { + let input = + table_scan(Some("employee_csv"), &employee_schema(), None)?.build()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(input), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![i32_split_point(10)], + )?), + }); + + let mut visited_exprs = vec![]; + plan.apply_expressions(|expr| { + visited_exprs.push(expr.to_string()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(visited_exprs, vec!["id"]); + + let plan = plan + .map_expressions(|expr| { + if expr == col("id") { + Ok(Transformed::yes(col("salary"))) + } else { + Ok(Transformed::no(expr)) + } + })? + .data; + + let LogicalPlan::Repartition(Repartition { + partitioning_scheme: Partitioning::Range(range), + .. + }) = plan + else { + unreachable!("expected range repartition"); + }; + assert_eq!(range.ordering()[0].expr, col("salary")); + assert_eq!(range.partition_count(), 2); + + Ok(()) + } + fn display_plan() -> Result { let plan1 = table_scan(Some("employee_csv"), &employee_schema(), Some(vec![3]))? .build()?; @@ -4500,6 +5178,74 @@ mod tests { .build() } + fn recursive_term_scan(name: &str, fields: Vec) -> Result> { + Ok(Arc::new( + table_scan(Some(name), &Schema::new(fields), None)?.build()?, + )) + } + + #[test] + fn recursive_query_widens_nullability_per_column() -> Result<()> { + // Column `a` is non-nullable in both terms and must stay non-nullable; + // column `b` is non-nullable in the static term but nullable in the + // recursive term, so the output must widen it to nullable. + let static_term = recursive_term_scan( + "static", + vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ], + )?; + let recursive_term = recursive_term_scan( + "rec", + vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, true), + ], + )?; + + let query = + RecursiveQuery::try_new("t".to_string(), static_term, recursive_term, false)?; + + // Names and types are taken from the static term. + assert_eq!(query.schema.field(0).name(), "a"); + assert_eq!(query.schema.field(1).name(), "b"); + assert_eq!(query.schema.field(0).data_type(), &DataType::Int32); + assert_eq!(query.schema.field(1).data_type(), &DataType::Int32); + // Nullability is widened independently per column. + assert!(!query.schema.field(0).is_nullable()); + assert!(query.schema.field(1).is_nullable()); + // `schema()` returns the widened recursive-query schema. + assert_eq!( + LogicalPlan::RecursiveQuery(query.clone()).schema(), + &query.schema + ); + Ok(()) + } + + #[test] + fn recursive_query_rejects_column_count_mismatch() -> Result<()> { + let static_term = + recursive_term_scan("static", vec![Field::new("a", DataType::Int32, false)])?; + let recursive_term = recursive_term_scan( + "rec", + vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ], + )?; + + let err = + RecursiveQuery::try_new("t".to_string(), static_term, recursive_term, false) + .unwrap_err(); + assert!( + err.strip_backtrace() + .contains("must have the same number of columns"), + "unexpected error: {err}" + ); + Ok(()) + } + #[test] fn test_display_indent() -> Result<()> { let plan = display_plan()?; @@ -5139,6 +5885,7 @@ mod tests { projected_schema: Arc::clone(&schema), filters: vec![], fetch: None, + statistics_requests: BTreeSet::new(), })); let col = schema.field_names()[0].clone(); @@ -5169,6 +5916,7 @@ mod tests { projected_schema: Arc::clone(&unique_schema), filters: vec![], fetch: None, + statistics_requests: BTreeSet::new(), })); let col = schema.field_names()[0].clone(); @@ -5955,4 +6703,53 @@ mod tests { Ok(()) } + + #[test] + fn test_unnest_with_new_exprs_accepts_expressions() -> Result<()> { + use crate::LogicalPlanBuilder; + use arrow::datatypes::{DataType, Field, Schema}; + + let schema = Schema::new(vec![ + Field::new("list_col", DataType::new_list(DataType::Int32, true), true), + Field::new("other_col", DataType::Int32, true), + ]); + let plan = table_scan(Some("t"), &schema, None)?.build()?; + let unnest_plan = LogicalPlanBuilder::from(plan) + .unnest_column("list_col")? + .build()?; + + let exprs = unnest_plan.expressions(); + assert!(!exprs.is_empty(), "Unnest should expose exec_columns"); + assert_eq!(exprs.len(), 1); + assert!(matches!(&exprs[0], Expr::Column(c) if c.name == "list_col")); + + let inputs: Vec = + unnest_plan.inputs().into_iter().cloned().collect(); + let rebuilt = unnest_plan.with_new_exprs(exprs, inputs)?; + assert_eq!(rebuilt.schema(), unnest_plan.schema()); + + Ok(()) + } + + #[test] + fn test_unnest_with_new_exprs_empty_preserves_columns() -> Result<()> { + use crate::LogicalPlanBuilder; + use arrow::datatypes::{DataType, Field, Schema}; + + let schema = Schema::new(vec![ + Field::new("list_col", DataType::new_list(DataType::Int32, true), true), + Field::new("other_col", DataType::Int32, true), + ]); + let plan = table_scan(Some("t"), &schema, None)?.build()?; + let unnest_plan = LogicalPlanBuilder::from(plan) + .unnest_column("list_col")? + .build()?; + + let inputs: Vec = + unnest_plan.inputs().into_iter().cloned().collect(); + let rebuilt = unnest_plan.with_new_exprs(vec![], inputs)?; + assert_eq!(rebuilt.schema(), unnest_plan.schema()); + + Ok(()) + } } diff --git a/datafusion/expr/src/logical_plan/tree_node.rs b/datafusion/expr/src/logical_plan/tree_node.rs index ef9382a57209a..c4c1d743b58b6 100644 --- a/datafusion/expr/src/logical_plan/tree_node.rs +++ b/datafusion/expr/src/logical_plan/tree_node.rs @@ -37,12 +37,15 @@ //! * [`LogicalPlan::with_new_exprs`]: Create a new plan with different expressions //! * [`LogicalPlan::expressions`]: Return a copy of the plan's expressions +use std::sync::Arc; + +use crate::logical_plan::plan::RangePartitioning; use crate::{ Aggregate, Analyze, CreateMemoryTable, CreateView, DdlStatement, Distinct, DistinctOn, DmlStatement, Execute, Explain, Expr, Extension, Filter, Join, Limit, LogicalPlan, Partitioning, Prepare, Projection, RecursiveQuery, Repartition, Sort, Statement, Subquery, SubqueryAlias, TableScan, Union, Unnest, UserDefinedLogicalNode, - Values, Window, dml::CopyTo, + Values, Window, WriteOp, builder::unnest_with_options, dml::CopyTo, }; use datafusion_common::tree_node::TreeNodeRefContainer; @@ -203,6 +206,7 @@ impl TreeNode for LogicalPlan { stringified_plans, schema, logical_optimization_succeeded, + show_statistics, }) => plan.map_elements(f)?.update_data(|plan| { LogicalPlan::Explain(Explain { verbose, @@ -211,17 +215,24 @@ impl TreeNode for LogicalPlan { stringified_plans, schema, logical_optimization_succeeded, + show_statistics, }) }), LogicalPlan::Analyze(Analyze { verbose, + format, input, schema, + analyze_level, + analyze_categories, }) => input.map_elements(f)?.update_data(|input| { LogicalPlan::Analyze(Analyze { verbose, + format, input, schema, + analyze_level, + analyze_categories, }) }), LogicalPlan::Dml(DmlStatement { @@ -329,13 +340,18 @@ impl TreeNode for LogicalPlan { static_term, recursive_term, is_distinct, + schema, }) => (static_term, recursive_term).map_elements(f)?.update_data( |(static_term, recursive_term)| { + // Ordinary child rewrites preserve derived schemas. Call + // `LogicalPlan::recompute_schema` when child schemas should + // be reconciled again. LogicalPlan::RecursiveQuery(RecursiveQuery { name, static_term, recursive_term, is_distinct, + schema, }) }, ), @@ -414,6 +430,7 @@ impl LogicalPlan { Partitioning::Hash(expr, _) | Partitioning::DistributeBy(expr) => { expr.apply_elements(f) } + Partitioning::Range(range) => range.ordering().to_vec().apply_elements(f), Partitioning::RoundRobinBatch(_) => Ok(TreeNodeRecursion::Continue), }, LogicalPlan::Window(Window { window_expr, .. }) => { @@ -463,6 +480,10 @@ impl LogicalPlan { } _ => Ok(TreeNodeRecursion::Continue), }, + LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(merge_op), + .. + }) => merge_op.exprs().apply_ref_elements(f), // plans without expressions LogicalPlan::EmptyRelation(_) | LogicalPlan::RecursiveQuery(_) @@ -519,6 +540,19 @@ impl LogicalPlan { Partitioning::DistributeBy(expr) => expr .map_elements(f)? .update_data(Partitioning::DistributeBy), + Partitioning::Range(range) => { + let split_points = range.split_points().to_vec(); + range + .ordering() + .to_vec() + .map_elements(f)? + .map_data(|ordering| { + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) + })? + } Partitioning::RoundRobinBatch(_) => Transformed::no(partitioning_scheme), } .update_data(|partitioning_scheme| { @@ -615,6 +649,7 @@ impl LogicalPlan { projected_schema, filters, fetch, + statistics_requests, }) => filters.map_elements(f)?.update_data(|filters| { LogicalPlan::TableScan(TableScan { table_name, @@ -623,6 +658,7 @@ impl LogicalPlan { projected_schema, filters, fetch, + statistics_requests, }) }), LogicalPlan::Distinct(Distinct::On(DistinctOn { @@ -656,9 +692,60 @@ impl LogicalPlan { _ => Transformed::no(stmt), } .update_data(LogicalPlan::Statement), + LogicalPlan::Unnest(Unnest { + input, + exec_columns, + options, + .. + }) => { + let exprs: Vec = + exec_columns.into_iter().map(Expr::Column).collect(); + exprs.map_elements(f)?.map_data(|mapped_exprs| { + let new_columns = mapped_exprs + .into_iter() + .map(|e| match e { + Expr::Column(c) => Ok(c), + other => internal_err!( + "Expected Expr::Column for Unnest exec_columns, got {other:?}" + ), + }) + .collect::>>()?; + // Rebuild through `unnest_with_options` so the derived + // `list_type_columns`, `struct_type_columns`, + // `dependency_indices`, and `schema` are recomputed from + // the (possibly rewritten) columns rather than carried over + // stale. This keeps `map_expressions` consistent with + // `with_new_exprs`. + unnest_with_options( + Arc::unwrap_or_clone(input), + new_columns, + options, + ) + })? + } + LogicalPlan::Dml(DmlStatement { + table_name, + target, + op: WriteOp::MergeInto(merge_op), + input, + output_schema, + }) => { + let owned_exprs: Vec = + merge_op.exprs().into_iter().cloned().collect(); + owned_exprs.map_elements(f)?.transform_data(|new_exprs| { + Ok(Transformed::no(LogicalPlan::Dml(DmlStatement { + table_name, + target, + op: WriteOp::MergeInto(Box::new( + merge_op.with_new_exprs(new_exprs)?, + )), + input, + output_schema, + }))) + })? + } // plans without expressions LogicalPlan::EmptyRelation(_) - | LogicalPlan::Unnest(_) | LogicalPlan::RecursiveQuery(_) | LogicalPlan::Subquery(_) | LogicalPlan::SubqueryAlias(_) @@ -841,6 +928,32 @@ impl LogicalPlan { }) } + /// Returns true if any expression in this node contains a subquery + /// (Exists, InSubquery, SetComparison, or ScalarSubquery). + fn has_subquery_expressions(&self) -> bool { + let mut found = false; + let _ = self.apply_expressions(|expr| { + if found { + return Ok(TreeNodeRecursion::Stop); + } + expr.apply(|e| { + if matches!( + e, + Expr::Exists(_) + | Expr::InSubquery(_) + | Expr::SetComparison(_) + | Expr::ScalarSubquery(_) + ) { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + }); + found + } + /// Similarly to [`Self::map_children`], rewrites all subqueries that may /// appear in expressions such as `IN (SELECT ...)` using `f`. /// @@ -849,6 +962,14 @@ impl LogicalPlan { self, mut f: F, ) -> Result> { + // Fast path: skip the expensive ownership-based expression traversal + // when this node has no subquery expressions. This avoids + // map_expressions → transform_down walking every expression node + // via consume+recreate just to find no subqueries. + if !self.has_subquery_expressions() { + return Ok(Transformed::no(self)); + } + self.map_expressions(|expr| { expr.transform_down(|expr| match expr { Expr::Exists(Exists { subquery, negated }) => { diff --git a/datafusion/expr/src/physical_planning_context.rs b/datafusion/expr/src/physical_planning_context.rs new file mode 100644 index 0000000000000..b2e579ea3c7cb --- /dev/null +++ b/datafusion/expr/src/physical_planning_context.rs @@ -0,0 +1,260 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt; +use std::hash::{Hash, Hasher}; +use std::sync::{Arc, Mutex}; + +use datafusion_common::{HashMap, Result, ScalarValue, TableReference, internal_err}; + +/// Context used while converting a logical plan subtree into a physical plan. +/// +/// Unlike [`ExecutionProps`](crate::execution_props::ExecutionProps), which +/// applies to the overall planning and execution of a query, this context can +/// differ between recursively planned subtrees. It currently carries: +/// +/// * the state needed to create physical expressions for +/// [`Expr::ScalarSubquery`] nodes that read from a shared +/// [`ScalarSubqueryResults`] container, and +/// * the qualifiers assigned to the [`Expr::LambdaVariable`]s that are in scope. +/// +/// The physical planner builds this context from the set of uncorrelated scalar +/// subqueries it has scheduled for a subtree. It is then passed explicitly +/// through `create_physical_expr` so that function can find the slot index for +/// each [`Subquery`]. While planning the body of a lambda, +/// `create_physical_expr` extends the context with the lambda's parameters via +/// [`Self::with_qualified_lambda_variables`]. +/// +/// An empty [`PhysicalPlanningContext`] (the [`Default`]) is what every +/// non-physical-planner caller passes; if such a caller encounters a scalar +/// subquery, `create_physical_expr` returns a `not_impl_err`. +/// +/// [`Expr::ScalarSubquery`]: crate::Expr::ScalarSubquery +/// [`Expr::LambdaVariable`]: crate::Expr::LambdaVariable +/// [`Subquery`]: crate::logical_plan::Subquery +#[derive(Clone, Debug, Default)] +pub struct PhysicalPlanningContext { + /// Behind an `Arc` because the context is cloned for each lambda body that + /// is planned, and the indexes are the same for the whole subtree. + indexes: Arc>, + results: ScalarSubqueryResults, + /// Maps each lambda variable name in scope to the qualifier generated for + /// its lambda during physical planning. + lambda_variable_qualifier: HashMap, +} + +impl PhysicalPlanningContext { + /// Create a [`PhysicalPlanningContext`] from an index map and a shared + /// results container. The index map must use the same indices as slots in + /// `results`. + pub fn new( + indexes: HashMap, + results: ScalarSubqueryResults, + ) -> Self { + Self { + indexes: Arc::new(indexes), + results, + lambda_variable_qualifier: HashMap::new(), + } + } + + /// Returns the slot index assigned to `subquery`, if any. + pub fn index_of( + &self, + subquery: &crate::logical_plan::Subquery, + ) -> Option { + self.indexes.get(subquery).copied() + } + + /// Returns the shared results container. + pub fn results(&self) -> &ScalarSubqueryResults { + &self.results + } + + /// Adds a mapping for each variable to the given qualifier. Existing + /// variables with conflicting names are shadowed. + pub fn with_qualified_lambda_variables( + mut self, + qualifier: &TableReference, + variables: &[String], + ) -> Self { + for var in variables { + self.lambda_variable_qualifier + .entry_ref(var) + .insert(qualifier.clone()); + } + + self + } + + /// Returns the qualifier of the lambda variable `name`, if it is in scope. + pub fn lambda_variable_qualifier(&self, name: &str) -> Option<&TableReference> { + self.lambda_variable_qualifier.get(name) + } +} + +/// Index of a scalar subquery within a [`ScalarSubqueryResults`] container. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +pub struct SubqueryIndex(usize); + +impl SubqueryIndex { + /// Creates a new subquery index. + pub const fn new(index: usize) -> Self { + Self(index) + } + + /// Returns the underlying slot index. + pub const fn as_usize(self) -> usize { + self.0 + } +} + +/// Shared results container for uncorrelated scalar subqueries. +/// +/// Each entry corresponds to one scalar subquery, identified by its index. +/// Each slot is populated at execution time by `ScalarSubqueryExec`, read by +/// `ScalarSubqueryExpr` instances that share this container, and cleared when +/// the plan is reset for re-execution. +#[derive(Clone, Default)] +pub struct ScalarSubqueryResults { + slots: Arc>>>, +} + +impl ScalarSubqueryResults { + /// Creates a new shared results container with `n` empty slots. + pub fn new(n: usize) -> Self { + Self { + slots: Arc::new((0..n).map(|_| Mutex::new(None)).collect()), + } + } + + /// Returns the scalar value stored at `index`, if it has been populated. + pub fn get(&self, index: SubqueryIndex) -> Option { + let slot = self.slots.get(index.as_usize())?; + slot.lock().unwrap().clone() + } + + /// Stores `value` in the slot at `index`. + pub fn set(&self, index: SubqueryIndex, value: ScalarValue) -> Result<()> { + let Some(slot) = self.slots.get(index.as_usize()) else { + return internal_err!( + "ScalarSubqueryResults: result index {} is out of bounds", + index.as_usize() + ); + }; + + let mut slot = slot.lock().unwrap(); + if slot.is_some() { + return internal_err!( + "ScalarSubqueryResults: result for index {} was already populated", + index.as_usize() + ); + } + *slot = Some(value); + + Ok(()) + } + + /// Clears all populated results so the container can be reused. + pub fn clear(&self) { + for slot in self.slots.iter() { + *slot.lock().unwrap() = None; + } + } + + /// Returns true if `this` and `other` point to the same shared container. + pub fn ptr_eq(this: &Self, other: &Self) -> bool { + Arc::ptr_eq(&this.slots, &other.slots) + } +} + +impl fmt::Debug for ScalarSubqueryResults { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_list() + .entries(self.slots.iter().map(|slot| slot.lock().unwrap().clone())) + .finish() + } +} + +impl PartialEq for ScalarSubqueryResults { + fn eq(&self, other: &Self) -> bool { + Self::ptr_eq(self, other) + } +} + +impl Eq for ScalarSubqueryResults {} + +impl Hash for ScalarSubqueryResults { + fn hash(&self, state: &mut H) { + Arc::as_ptr(&self.slots).hash(state); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn scalar_subquery_results_set_and_get() -> Result<()> { + let results = ScalarSubqueryResults::new(1); + assert_eq!(results.get(SubqueryIndex::new(0)), None); + + results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?; + assert_eq!( + results.get(SubqueryIndex::new(0)), + Some(ScalarValue::Int32(Some(42))) + ); + assert!( + results + .set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7))) + .is_err() + ); + + Ok(()) + } + + #[test] + fn lambda_variables_shadow_outer_scope() { + let outer = TableReference::bare("lambda_1"); + let inner = TableReference::bare("lambda_2"); + + let ctx = PhysicalPlanningContext::default() + .with_qualified_lambda_variables(&outer, &["x".to_string(), "y".to_string()]) + .with_qualified_lambda_variables(&inner, &["y".to_string()]); + + assert_eq!(ctx.lambda_variable_qualifier("x"), Some(&outer)); + assert_eq!(ctx.lambda_variable_qualifier("y"), Some(&inner)); + assert_eq!(ctx.lambda_variable_qualifier("z"), None); + } + + #[test] + fn scalar_subquery_results_clear() -> Result<()> { + let results = ScalarSubqueryResults::new(1); + results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(42)))?; + + results.clear(); + + assert_eq!(results.get(SubqueryIndex::new(0)), None); + results.set(SubqueryIndex::new(0), ScalarValue::Int32(Some(7)))?; + assert_eq!( + results.get(SubqueryIndex::new(0)), + Some(ScalarValue::Int32(Some(7))) + ); + + Ok(()) + } +} diff --git a/datafusion/expr/src/planner.rs b/datafusion/expr/src/planner.rs index d69f4ac5fe23f..7aaf3a98cbe5d 100644 --- a/datafusion/expr/src/planner.rs +++ b/datafusion/expr/src/planner.rs @@ -61,7 +61,8 @@ pub trait ContextProvider { not_impl_err!("Table Functions are not supported") } - /// Provides an intermediate table that is used to store the results of a CTE during execution + /// Provides an intermediate table that is used to expose a recursive CTE + /// self-reference during planning and execution. /// /// CTE stands for "Common Table Expression" /// @@ -72,6 +73,9 @@ pub trait ContextProvider { /// of the sql crate (for example [`CteWorkTable`]). /// /// The [`ContextProvider`] provides a way to "hide" this dependency. + /// The schema argument is the schema to expose for scans of the recursive + /// self-reference, which may be more conservative than the final recursive + /// query output schema. /// /// [`SqlToRel`]: https://docs.rs/datafusion/latest/datafusion/sql/planner/struct.SqlToRel.html /// [`CteWorkTable`]: https://docs.rs/datafusion/latest/datafusion/datasource/cte_worktable/struct.CteWorkTable.html @@ -104,7 +108,7 @@ pub trait ContextProvider { fn get_function_meta(&self, name: &str) -> Option>; /// Return the higher order function with a given name, if any - fn get_higher_order_meta(&self, name: &str) -> Option>; + fn get_higher_order_meta(&self, name: &str) -> Option>; /// Return the aggregate function with a given name, if any fn get_aggregate_meta(&self, name: &str) -> Option>; diff --git a/datafusion/expr/src/predicate_bounds.rs b/datafusion/expr/src/predicate_bounds.rs index 992d9f88bb14a..6b672221a7d06 100644 --- a/datafusion/expr/src/predicate_bounds.rs +++ b/datafusion/expr/src/predicate_bounds.rs @@ -183,6 +183,10 @@ impl PredicateBoundsEvaluator<'_> { Expr::BinaryExpr(BinaryExpr { op, .. }) if op.returns_null_on_null() => { self.is_null_if_any_child_null(expr) } + // Strict scalar functions return NULL when any argument is NULL. + Expr::ScalarFunction(func) if func.func.is_strict() => { + self.is_null_if_any_child_null(expr) + } Expr::Alias(_) | Expr::Cast(_) | Expr::Like(_) @@ -235,8 +239,9 @@ mod tests { use crate::expr::ScalarFunction; use crate::predicate_bounds::evaluate_bounds; use crate::{ - Expr, binary_expr, col, create_udf, is_false, is_not_false, is_not_null, - is_not_true, is_not_unknown, is_null, is_true, is_unknown, lit, not, + Expr, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, binary_expr, col, + create_udf, is_false, is_not_false, is_not_null, is_not_true, is_not_unknown, + is_null, is_true, is_unknown, lit, not, }; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::{DFSchema, Result, ScalarValue}; @@ -666,6 +671,29 @@ mod tests { } } + #[test] + fn evaluate_bounds_strict_udf_is_null_when_child_null() { + let col = col("col"); + let strict_func = make_test_udf_expr("strict_test", true, vec![col.clone()]); + let non_strict_func = + make_test_udf_expr("non_strict_test", false, vec![col.clone()]); + let schema = DFSchema::try_from(Schema::new(vec![Field::new( + "col", + DataType::UInt8, + true, + )])) + .unwrap(); + + assert_eq!( + evaluate_bounds(&is_not_null(strict_func), Some(&col), &schema).unwrap(), + NullableInterval::FALSE, + ); + assert_eq!( + evaluate_bounds(&is_not_null(non_strict_func), Some(&col), &schema).unwrap(), + NullableInterval::TRUE_OR_FALSE, + ); + } + fn make_scalar_func_expr() -> Expr { let scalar_func_impl = |_: &[ColumnarValue]| Ok(ColumnarValue::Scalar(ScalarValue::Null)); @@ -678,4 +706,51 @@ mod tests { ); Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(udf), vec![])) } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdf { + name: &'static str, + signature: Signature, + strict: bool, + } + + impl TestUdf { + fn new(name: &'static str, strict: bool) -> Self { + Self { + name, + signature: Signature::uniform( + 1, + vec![DataType::UInt8], + Volatility::Immutable, + ), + strict, + } + } + } + + impl ScalarUDFImpl for TestUdf { + fn name(&self) -> &str { + self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::UInt8) + } + + fn is_strict(&self) -> bool { + self.strict + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + unimplemented!() + } + } + + fn make_test_udf_expr(name: &'static str, strict: bool, args: Vec) -> Expr { + ScalarUDF::from(TestUdf::new(name, strict)).call(args) + } } diff --git a/datafusion/expr/src/proto.rs b/datafusion/expr/src/proto.rs new file mode 100644 index 0000000000000..00b340210807d --- /dev/null +++ b/datafusion/expr/src/proto.rs @@ -0,0 +1,246 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Protobuf conversions for the expression types owned by this crate: +//! [`WindowFrame`], [`WindowFrameBound`], [`WindowFrameUnits`], +//! [`MergeIntoClauseKind`](crate::dml::MergeIntoClauseKind) and +//! [`NullTreatment`](crate::expr::NullTreatment). +//! +//! These are plain [`From`] / [`TryFrom`] impls rather than something taking a +//! codec: every field is either an enum tag or a [`ScalarValue`], so the +//! conversion needs nothing but the value itself. The orphan rule allows them +//! here because one side of each conversion is a type this crate owns. +//! +//! [`ScalarValue`]: datafusion_common::ScalarValue + +use datafusion_common::ScalarValue; +use datafusion_proto_common::{FromProtoError, ToProtoError}; +use datafusion_proto_models::protobuf; + +use crate::dml::MergeIntoClauseKind; +use crate::expr::NullTreatment; +use crate::{WindowFrame, WindowFrameBound, WindowFrameUnits}; + +impl From for WindowFrameUnits { + fn from(units: protobuf::WindowFrameUnits) -> Self { + match units { + protobuf::WindowFrameUnits::Rows => Self::Rows, + protobuf::WindowFrameUnits::Range => Self::Range, + protobuf::WindowFrameUnits::Groups => Self::Groups, + } + } +} + +impl From for protobuf::WindowFrameUnits { + fn from(units: WindowFrameUnits) -> Self { + match units { + WindowFrameUnits::Rows => Self::Rows, + WindowFrameUnits::Range => Self::Range, + WindowFrameUnits::Groups => Self::Groups, + } + } +} + +impl TryFrom for WindowFrameBound { + type Error = FromProtoError; + + fn try_from(bound: protobuf::WindowFrameBound) -> Result { + let bound_type = + protobuf::WindowFrameBoundType::try_from(bound.window_frame_bound_type) + .map_err(|_| { + FromProtoError::unknown( + "WindowFrameBoundType", + bound.window_frame_bound_type, + ) + })?; + match bound_type { + protobuf::WindowFrameBoundType::CurrentRow => Ok(Self::CurrentRow), + protobuf::WindowFrameBoundType::Preceding => match bound.bound_value { + Some(x) => Ok(Self::Preceding(ScalarValue::try_from(&x)?)), + None => Ok(Self::Preceding(ScalarValue::UInt64(None))), + }, + protobuf::WindowFrameBoundType::Following => match bound.bound_value { + Some(x) => Ok(Self::Following(ScalarValue::try_from(&x)?)), + None => Ok(Self::Following(ScalarValue::UInt64(None))), + }, + } + } +} + +impl TryFrom<&WindowFrameBound> for protobuf::WindowFrameBound { + type Error = ToProtoError; + + fn try_from(bound: &WindowFrameBound) -> Result { + Ok(match bound { + WindowFrameBound::CurrentRow => Self { + window_frame_bound_type: protobuf::WindowFrameBoundType::CurrentRow + .into(), + bound_value: None, + }, + WindowFrameBound::Preceding(v) => Self { + window_frame_bound_type: protobuf::WindowFrameBoundType::Preceding.into(), + bound_value: Some(v.try_into()?), + }, + WindowFrameBound::Following(v) => Self { + window_frame_bound_type: protobuf::WindowFrameBoundType::Following.into(), + bound_value: Some(v.try_into()?), + }, + }) + } +} + +impl TryFrom for WindowFrame { + type Error = FromProtoError; + + fn try_from(window: protobuf::WindowFrame) -> Result { + let units = WindowFrameUnits::from( + protobuf::WindowFrameUnits::try_from(window.window_frame_units).map_err( + |_| { + FromProtoError::unknown("WindowFrameUnits", window.window_frame_units) + }, + )?, + ); + let start_bound = WindowFrameBound::try_from( + window + .start_bound + .ok_or_else(|| FromProtoError::required("start_bound"))?, + )?; + let end_bound = window + .end_bound + .map(|end_bound| match end_bound { + protobuf::window_frame::EndBound::Bound(end_bound) => { + WindowFrameBound::try_from(end_bound) + } + }) + .transpose()? + .unwrap_or(WindowFrameBound::CurrentRow); + Ok(WindowFrame::new_bounds(units, start_bound, end_bound)) + } +} + +impl TryFrom<&WindowFrame> for protobuf::WindowFrame { + type Error = ToProtoError; + + fn try_from(window: &WindowFrame) -> Result { + Ok(Self { + window_frame_units: protobuf::WindowFrameUnits::from(window.units).into(), + start_bound: Some((&window.start_bound).try_into()?), + end_bound: Some(protobuf::window_frame::EndBound::Bound( + (&window.end_bound).try_into()?, + )), + }) + } +} + +impl From for MergeIntoClauseKind { + fn from(kind: protobuf::merge_into_clause_node::Kind) -> Self { + match kind { + protobuf::merge_into_clause_node::Kind::Matched => Self::Matched, + protobuf::merge_into_clause_node::Kind::NotMatched => Self::NotMatched, + protobuf::merge_into_clause_node::Kind::NotMatchedByTarget => { + Self::NotMatchedByTarget + } + protobuf::merge_into_clause_node::Kind::NotMatchedBySource => { + Self::NotMatchedBySource + } + } + } +} + +impl From for protobuf::merge_into_clause_node::Kind { + fn from(kind: MergeIntoClauseKind) -> Self { + match kind { + MergeIntoClauseKind::Matched => Self::Matched, + MergeIntoClauseKind::NotMatched => Self::NotMatched, + MergeIntoClauseKind::NotMatchedByTarget => Self::NotMatchedByTarget, + MergeIntoClauseKind::NotMatchedBySource => Self::NotMatchedBySource, + } + } +} + +impl From for NullTreatment { + fn from(t: protobuf::NullTreatment) -> Self { + match t { + protobuf::NullTreatment::RespectNulls => Self::RespectNulls, + protobuf::NullTreatment::IgnoreNulls => Self::IgnoreNulls, + } + } +} + +impl From for protobuf::NullTreatment { + fn from(t: NullTreatment) -> Self { + match t { + NullTreatment::RespectNulls => Self::RespectNulls, + NullTreatment::IgnoreNulls => Self::IgnoreNulls, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn window_frame_roundtrip() -> Result<(), Box> { + let frame = WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::Preceding(ScalarValue::UInt64(Some(2))), + WindowFrameBound::Following(ScalarValue::UInt64(Some(3))), + ); + + let encoded = protobuf::WindowFrame::try_from(&frame)?; + let decoded = WindowFrame::try_from(encoded)?; + + assert_eq!(decoded.units, frame.units); + assert_eq!(decoded.start_bound, frame.start_bound); + assert_eq!(decoded.end_bound, frame.end_bound); + Ok(()) + } + + #[test] + fn window_frame_from_proto_rejects_missing_start_bound() { + let proto = protobuf::WindowFrame { + window_frame_units: protobuf::WindowFrameUnits::Rows.into(), + start_bound: None, + end_bound: None, + }; + + let err = WindowFrame::try_from(proto).unwrap_err(); + assert!( + err.to_string().contains("start_bound"), + "unexpected error: {err}" + ); + } + + #[test] + fn missing_end_bound_decodes_as_current_row() -> Result<(), Box> + { + let proto = protobuf::WindowFrame { + window_frame_units: protobuf::WindowFrameUnits::Rows.into(), + start_bound: Some(protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::CurrentRow + .into(), + bound_value: None, + }), + end_bound: None, + }; + + let decoded = WindowFrame::try_from(proto)?; + assert_eq!(decoded.end_bound, WindowFrameBound::CurrentRow); + Ok(()) + } +} diff --git a/datafusion/expr/src/registry.rs b/datafusion/expr/src/registry.rs index f03cc5936c6ed..4b9744d9573b6 100644 --- a/datafusion/expr/src/registry.rs +++ b/datafusion/expr/src/registry.rs @@ -56,7 +56,7 @@ pub trait FunctionRegistry { /// Returns a reference to the user defined higher order function named /// `name`. - fn higher_order_function(&self, name: &str) -> Result>; + fn higher_order_function(&self, name: &str) -> Result>; /// Returns a reference to the user defined aggregate function (udaf) named /// `name`. @@ -81,8 +81,8 @@ pub trait FunctionRegistry { /// for example if the registry is read only. fn register_higher_order_function( &mut self, - _function: Arc, - ) -> Result>> { + _function: Arc, + ) -> Result>> { not_impl_err!("Registering HigherOrderUDF") } /// Registers a new [`AggregateUDF`], returning any previously registered @@ -122,7 +122,7 @@ pub trait FunctionRegistry { fn deregister_higher_order_function( &mut self, _name: &str, - ) -> Result>> { + ) -> Result>> { not_impl_err!("Deregistering HigherOrderUDF") } @@ -198,7 +198,7 @@ pub struct MemoryFunctionRegistry { /// Window Functions udwfs: HashMap>, /// Higher Order Functions - higher_order_functions: HashMap>, + higher_order_functions: HashMap>, } impl MemoryFunctionRegistry { @@ -219,7 +219,7 @@ impl FunctionRegistry for MemoryFunctionRegistry { .ok_or_else(|| plan_datafusion_err!("Function {name} not found")) } - fn higher_order_function(&self, name: &str) -> Result> { + fn higher_order_function(&self, name: &str) -> Result> { self.higher_order_functions .get(name) .cloned() @@ -245,8 +245,8 @@ impl FunctionRegistry for MemoryFunctionRegistry { } fn register_higher_order_function( &mut self, - function: Arc, - ) -> Result>> { + function: Arc, + ) -> Result>> { Ok(self .higher_order_functions .insert(function.name().into(), function)) diff --git a/datafusion/expr/src/sql.rs b/datafusion/expr/src/sql.rs index d582a0f6b95d1..23e8d2f63d941 100644 --- a/datafusion/expr/src/sql.rs +++ b/datafusion/expr/src/sql.rs @@ -38,7 +38,7 @@ pub struct IlikeSelectItem { impl Display for IlikeSelectItem { fn fmt(&self, f: &mut Formatter) -> fmt::Result { - write!(f, "ILIKE '{}'", &self.pattern)?; + write!(f, "ILIKE '{}'", self.pattern)?; Ok(()) } } diff --git a/datafusion/expr/src/tree_node.rs b/datafusion/expr/src/tree_node.rs index 010441b5a25d1..941fd22ea179f 100644 --- a/datafusion/expr/src/tree_node.rs +++ b/datafusion/expr/src/tree_node.rs @@ -49,7 +49,7 @@ impl TreeNode for Expr { ) -> Result { match self { Expr::Alias(Alias { expr, .. }) - | Expr::Unnest(Unnest { expr }) + | Expr::Unnest(Unnest { expr, .. }) | Expr::Not(expr) | Expr::IsNotNull(expr) | Expr::IsTrue(expr) @@ -150,9 +150,9 @@ impl TreeNode for Expr { quantifier, }) }), - Expr::Unnest(Unnest { expr, .. }) => expr + Expr::Unnest(Unnest { expr, outer }) => expr .map_elements(f)? - .update_data(|expr| Expr::Unnest(Unnest { expr })), + .update_data(|expr| Expr::Unnest(Unnest { expr, outer })), Expr::Alias(Alias { expr, relation, diff --git a/datafusion/expr/src/type_coercion/functions.rs b/datafusion/expr/src/type_coercion/functions.rs index 1f625e33d31ef..ec3ab6f441827 100644 --- a/datafusion/expr/src/type_coercion/functions.rs +++ b/datafusion/expr/src/type_coercion/functions.rs @@ -33,7 +33,9 @@ use datafusion_common::utils::{ use datafusion_common::{ Result, exec_err, internal_err, plan_err, types::NativeType, utils::list_ndims, }; -use datafusion_expr_common::signature::ArrayFunctionArgument; +use datafusion_expr_common::signature::{ + ArrayFunctionArgument, EncodingPreservation, TypeSignatureClass, +}; use datafusion_expr_common::type_coercion::binary::type_union_resolution; use datafusion_expr_common::{ signature::{ArrayFunctionSignature, FIXED_SIZE_LIST_WILDCARD, TIMEZONE_WILDCARD}, @@ -158,7 +160,7 @@ pub fn fields_with_udf( /// argument must be coerced to match `signature`. /// For lambda arguments, returns a clone of the associated data /// -/// Note this does not invokes [HigherOrderUDF::coerce_values_for_lambdas]. +/// Note this does not invokes [crate::HigherOrderUDFImpl::coerce_values_for_lambdas]. /// If that's required, use [value_fields_with_higher_order_udf_and_lambdas] /// instead /// @@ -166,7 +168,7 @@ pub fn fields_with_udf( /// [`type_coercion`](crate::type_coercion) module. pub fn value_fields_with_higher_order_udf( current_fields: &[ValueOrLambda], - func: &dyn HigherOrderUDF, + func: &HigherOrderUDF, ) -> Result>> { match func.signature().type_signature { HigherOrderTypeSignature::UserDefined => { @@ -306,7 +308,7 @@ pub fn value_fields_with_higher_order_udf( } /// Performs type coercion for higher order function arguments, -/// including those defined by [HigherOrderUDF::coerce_values_for_lambdas], +/// including those defined by [crate::HigherOrderUDFImpl::coerce_values_for_lambdas], /// if it returns `Some(...)` instead of the default `None`. Note that /// compared to [value_fields_with_higher_order_udf], this function requires /// the [ValueOrLambda::Lambda] variant to contain the output field of the lambda. @@ -319,7 +321,7 @@ pub fn value_fields_with_higher_order_udf( /// [`type_coercion`](crate::type_coercion) module. pub fn value_fields_with_higher_order_udf_and_lambdas( current_fields: &[ValueOrLambda], - func: &dyn HigherOrderUDF, + func: &HigherOrderUDF, ) -> Result>> { let mut new_fields = value_fields_with_higher_order_udf(current_fields, func)?; @@ -586,6 +588,52 @@ fn get_valid_types( arguments: &[ArrayFunctionArgument], array_coercion: Option<&ListCoercion>, ) -> Result>> { + fn rebuild_array_type( + current_type: &DataType, + element_type: &DataType, + nullable: bool, + large_list: bool, + fixed_size: Option, + ) -> DataType { + // Preserve the original list field when possible so field name or + // metadata differences do not introduce otherwise unnecessary casts. + let field = match current_type { + DataType::List(field) + | DataType::LargeList(field) + | DataType::FixedSizeList(field, _) => Some(Arc::new( + field + .as_ref() + .clone() + .with_data_type(element_type.clone()) + .with_nullable(nullable), + )), + _ => None, + }; + + if large_list { + field.map_or_else( + || DataType::new_large_list(element_type.clone(), nullable), + DataType::LargeList, + ) + } else if let Some(size) = fixed_size { + field.map_or_else( + || { + DataType::new_fixed_size_list( + element_type.clone(), + size, + nullable, + ) + }, + |field| DataType::FixedSizeList(field, size), + ) + } else { + field.map_or_else( + || DataType::new_list(element_type.clone(), nullable), + DataType::List, + ) + } + } + if current_types.len() != arguments.len() { return Ok(vec![vec![]]); } @@ -657,21 +705,13 @@ fn get_valid_types( ArrayFunctionArgument::Array => { if current_type.is_null() { DataType::Null - } else if large_list { - DataType::new_large_list( - element_type.clone(), - is_nested_item_nullable.unwrap_or(true), - ) - } else if let Some(size) = list_sizes.next() { - DataType::new_fixed_size_list( - element_type.clone(), - size, - is_nested_item_nullable.unwrap_or(true), - ) } else { - DataType::new_list( - element_type.clone(), + rebuild_array_type( + current_type, + &element_type, is_nested_item_nullable.unwrap_or(true), + large_list, + list_sizes.next(), ) } } @@ -835,9 +875,56 @@ fn get_valid_types( TypeSignature::Coercible(param_types) => { function_length_check(function_name, current_types.len(), param_types.len())?; + fn coercion_value_type<'a>( + current_type: &'a DataType, + desired_type: &TypeSignatureClass, + ) -> &'a DataType { + if matches!(desired_type, TypeSignatureClass::Any) { + return current_type; + } + + match current_type { + DataType::Dictionary(_, value_type) => { + coercion_value_type(value_type, desired_type) + } + _ => current_type, + } + } + + fn preserve_encoding( + current_type: &DataType, + casted_type: DataType, + desired_type: &TypeSignatureClass, + encoding_preservation: EncodingPreservation, + ) -> DataType { + if matches!(desired_type, TypeSignatureClass::Any) { + return casted_type; + } + + match current_type { + DataType::Dictionary(key_type, value_type) => { + let casted_type = preserve_encoding( + value_type, + casted_type, + desired_type, + encoding_preservation, + ); + if encoding_preservation.preserve_dictionary() { + DataType::Dictionary(key_type.clone(), Box::new(casted_type)) + } else { + casted_type + } + } + _ => casted_type, + } + } + let mut new_types = Vec::with_capacity(current_types.len()); for (current_type, param) in current_types.iter().zip(param_types.iter()) { let current_native_type: NativeType = current_type.into(); + let encoding_preservation = param.encoding_preservation(); + let coercion_value_type = + coercion_value_type(current_type, param.desired_type()); if param .desired_type() @@ -845,9 +932,14 @@ fn get_valid_types( { let casted_type = param .desired_type() - .default_casted_type(¤t_native_type, current_type)?; + .default_casted_type(¤t_native_type, coercion_value_type)?; - new_types.push(casted_type); + new_types.push(preserve_encoding( + current_type, + casted_type, + param.desired_type(), + encoding_preservation, + )); } else if param .allowed_source_types() .iter() @@ -856,8 +948,13 @@ fn get_valid_types( // If the condition is met which means `implicit coercion`` is provided so we can safely unwrap let default_casted_type = param.default_casted_type().unwrap(); let casted_type = - default_casted_type.default_cast_for(current_type)?; - new_types.push(casted_type); + default_casted_type.default_cast_for(coercion_value_type)?; + new_types.push(preserve_encoding( + current_type, + casted_type, + param.desired_type(), + encoding_preservation, + )); } else { let hint = if matches!(current_native_type, NativeType::Binary) { "\n\nHint: Binary types are not automatically coerced to String. Use CAST(column AS VARCHAR) to convert Binary data to String." @@ -987,12 +1084,8 @@ fn maybe_data_types( // attempt to coerce. // TODO: Replace with `can_cast_types` after failing cases are resolved // (they need new signature that returns exactly valid types instead of list of possible valid types). - if let Some(coerced_type) = coerced_from(valid_type, current_type) { - new_type.push(coerced_type) - } else { - // not possible - return None; - } + let coerced_type = coerced_from(valid_type, current_type)?; + new_type.push(coerced_type) } } Some(new_type) @@ -1105,13 +1198,11 @@ fn coerced_from<'a>( ) => Some(type_into.clone()), ( Timestamp(TimeUnit::Nanosecond, None), - Null | Timestamp(_, None) | Date32 | Utf8 | LargeUtf8, + Null | Timestamp(_, None) | Date32 | Date64 | Utf8 | LargeUtf8 | Utf8View, ) => Some(type_into.clone()), - (Interval(_), Null | Utf8 | LargeUtf8) => Some(type_into.clone()), - // We can go into a Utf8View from a Utf8 or LargeUtf8 - (Utf8View, Utf8 | LargeUtf8 | Null) => Some(type_into.clone()), + (Interval(_), Null | Utf8 | LargeUtf8 | Utf8View) => Some(type_into.clone()), // Any type can be coerced into strings - (Utf8 | LargeUtf8, _) => Some(type_into.clone()), + (Utf8 | LargeUtf8 | Utf8View, _) => Some(type_into.clone()), // We can go into a BinaryView from a Binary or LargeBinary (BinaryView, Binary | LargeBinary | Null) => Some(type_into.clone()), (Null, _) if can_cast_types(type_from, type_into) => Some(type_into.clone()), @@ -1169,18 +1260,18 @@ fn coerced_from<'a>( mod tests { use crate::{ HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, HigherOrderSignature, - Volatility, + HigherOrderUDFImpl, Volatility, }; use super::*; use arrow::datatypes::IntervalUnit; use datafusion_common::{ assert_contains, - types::{logical_binary, logical_int64}, + types::{logical_binary, logical_int64, logical_string}, }; use datafusion_expr_common::{ columnar_value::ColumnarValue, - signature::{Coercion, TypeSignatureClass}, + signature::{Coercion, EncodingPreservation, TypeSignatureClass}, }; #[test] @@ -1664,6 +1755,31 @@ mod tests { Ok(()) } + #[test] + fn test_get_valid_types_array_and_index_preserves_list_field_name() -> Result<()> { + let struct_fields = vec![ + Field::new("id", DataType::Utf8, true), + Field::new("prim", DataType::Boolean, true), + ]; + let current_type = DataType::List(Arc::new(Field::new( + "element", + DataType::Struct(struct_fields.into()), + true, + ))); + let signature = Signature::array_and_index(Volatility::Immutable); + + assert_eq!( + get_valid_types( + "array_element", + &signature.type_signature, + &[current_type.clone(), DataType::Int64], + )?, + vec![vec![current_type, DataType::Int64]] + ); + + Ok(()) + } + #[test] fn test_get_valid_types_element_and_array() -> Result<()> { let function = "element_and_array"; @@ -1754,16 +1870,196 @@ mod tests { ))?; assert_eq!(vec![DataType::Int64], output); - // Dictionary gets passed through if we use TypeSignatureClass apart from Native - let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?; + // Any always preserves the original physical type + let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Any))?; assert_eq!(vec![dictionary.clone()], output); + let output = dictionary_input( + Coercion::new_exact(TypeSignatureClass::Any) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?; + assert_eq!(vec![dictionary.clone()], output); + + // Typed non-Native classes materialize dictionaries by default + let output = dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?; + assert_eq!(vec![DataType::Int64], output); + let output = dictionary_input(Coercion::new_implicit( TypeSignatureClass::Integer, vec![], NativeType::Int64, ))?; - assert_eq!(vec![dictionary.clone()], output); + assert_eq!(vec![DataType::Int64], output); + + // Typed non-Native classes preserve dictionaries only when requested + let output = dictionary_input( + Coercion::new_exact(TypeSignatureClass::Integer) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?; + assert_eq!(vec![dictionary], output); + + Ok(()) + } + + #[test] + fn test_coercible_dictionary_preserves_encoding() -> Result<()> { + fn dictionary_input( + value_type: DataType, + coercion: Coercion, + ) -> Result> { + fields_with_udf( + &[Field::new( + "field", + DataType::Dictionary(Box::new(DataType::Int8), Box::new(value_type)), + true, + ) + .into()], + &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)), + ) + .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect()) + } + + let coercion = Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()); + + assert_eq!( + dictionary_input(DataType::LargeUtf8, coercion.clone())?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::LargeUtf8), + )] + ); + assert_eq!( + dictionary_input( + DataType::BinaryView, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_string()), + vec![TypeSignatureClass::Native(logical_binary())], + NativeType::String, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Utf8View), + )] + ); + // Contrast: without encoding_preservation, Native strips dictionary entirely + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ), + )?, + vec![DataType::Int64] + ); + // With encoding_preservation, dictionary wrapper is preserved, value coerced + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Int64), + )] + ); + // Without encoding_preservation, non-Native classes materialize dictionaries + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Integer, + vec![], + NativeType::Int64, + ), + )?, + vec![DataType::Int32] + ); + // With encoding_preservation, non-Native classes preserve dictionaries + assert_eq!( + dictionary_input( + DataType::Int32, + Coercion::new_implicit( + TypeSignatureClass::Integer, + vec![], + NativeType::Int64, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?, + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Int32), + )] + ); + + Ok(()) + } + + #[test] + fn test_coercible_nested_dictionary() -> Result<()> { + let nested_dictionary = DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Dictionary( + Box::new(DataType::Int16), + Box::new(DataType::Int32), + )), + ); + let nested_dictionary_input = |coercion| -> Result> { + fields_with_udf( + &[Field::new("field", nested_dictionary.clone(), true).into()], + &MockUdf(Signature::coercible(vec![coercion], Volatility::Immutable)), + ) + .map(|v| v.into_iter().map(|f| f.data_type().clone()).collect()) + }; + + // Without preservation, recursively unwrap dictionaries to the unchanged leaf. + let output = + nested_dictionary_input(Coercion::new_exact(TypeSignatureClass::Integer))?; + assert_eq!(vec![DataType::Int32], output); + + // With preservation, restore the complete dictionary stack around the leaf. + let output = nested_dictionary_input( + Coercion::new_exact(TypeSignatureClass::Integer) + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?; + assert_eq!(vec![nested_dictionary.clone()], output); + + let int64_coercion = || { + Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ) + }; + + // Without preservation, materialize the coerced leaf type. + let output = nested_dictionary_input(int64_coercion())?; + assert_eq!(vec![DataType::Int64], output); + + // With preservation, restore the complete dictionary stack around the coerced leaf. + let output = nested_dictionary_input( + int64_coercion() + .with_encoding_preservation(EncodingPreservation::dictionary()), + )?; + assert_eq!( + vec![DataType::Dictionary( + Box::new(DataType::Int8), + Box::new(DataType::Dictionary( + Box::new(DataType::Int16), + Box::new(DataType::Int64), + )), + )], + output + ); Ok(()) } @@ -1901,7 +2197,7 @@ mod tests { coerced_value_types: Vec, } - impl HigherOrderUDF for MockHigherOrderUDF { + impl HigherOrderUDFImpl for MockHigherOrderUDF { fn name(&self) -> &str { "mock_higher_order_function" } @@ -1962,10 +2258,10 @@ mod tests { #[test] fn test_higher_order_function_user_defined_type_coercion() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::user_defined(Volatility::Immutable), coerced_value_types: vec![DataType::new_large_list(DataType::Int32, false)], - }; + }); let new_fields = value_fields_with_higher_order_udf( &[ @@ -1996,10 +2292,10 @@ mod tests { #[test] fn test_higher_order_function_coerce_values_for_lambdas() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::variadic_any(Volatility::Immutable), coerced_value_types: vec![], - }; + }); let new_fields = value_fields_with_higher_order_udf_and_lambdas( &[ @@ -2032,10 +2328,10 @@ mod tests { #[test] fn test_higher_order_function_user_defined_type_coercion_bad_args() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::user_defined(Volatility::Immutable), coerced_value_types: vec![DataType::Int32], - }; + }); let err = value_fields_with_higher_order_udf::<()>(&[], &fun).unwrap_err(); @@ -2047,10 +2343,10 @@ mod tests { #[test] fn test_higher_order_function_faulty_user_defined_type_coercion() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::user_defined(Volatility::Immutable), coerced_value_types: vec![DataType::Int32, DataType::Int32], - }; + }); let err = value_fields_with_higher_order_udf::<()>( &[ValueOrLambda::Value(Arc::new(Field::new( @@ -2070,10 +2366,10 @@ mod tests { #[test] fn test_higher_order_function_any_signature() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::any(1, Volatility::Immutable), coerced_value_types: vec![], - }; + }); let new_fields = value_fields_with_higher_order_udf(&[ValueOrLambda::Lambda(())], &fun) @@ -2085,10 +2381,10 @@ mod tests { #[test] fn test_higher_order_function_any_signature_bad_args() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::any(1, Volatility::Immutable), coerced_value_types: vec![], - }; + }); let err = value_fields_with_higher_order_udf::<()>(&[], &fun).unwrap_err(); @@ -2100,13 +2396,13 @@ mod tests { #[test] fn test_higher_order_function_exact_signature() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::exact( vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())], Volatility::Immutable, ), coerced_value_types: vec![DataType::new_large_list(DataType::Int32, false)], - }; + }); let new_fields = value_fields_with_higher_order_udf( &[ @@ -2137,13 +2433,13 @@ mod tests { #[test] fn test_higher_order_function_exact_signature_wrong_value_count() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::exact( vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())], Volatility::Immutable, ), coerced_value_types: vec![], - }; + }); let err = value_fields_with_higher_order_udf::<()>( &[ValueOrLambda::Lambda(()), ValueOrLambda::Lambda(())], @@ -2159,13 +2455,13 @@ mod tests { #[test] fn test_higher_order_function_exact_signature_wrong_lambda_count() { - let fun = MockHigherOrderUDF { + let fun = HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::exact( vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())], Volatility::Immutable, ), coerced_value_types: vec![], - }; + }); let err = value_fields_with_higher_order_udf::<()>( &[ diff --git a/datafusion/expr/src/udf.rs b/datafusion/expr/src/udf.rs index 6a3aa31a8609f..2de3be4c10fa4 100644 --- a/datafusion/expr/src/udf.rs +++ b/datafusion/expr/src/udf.rs @@ -209,6 +209,14 @@ impl ScalarUDF { self.inner.aliases() } + /// Returns true if this function always returns NULL when any argument is + /// NULL. + /// + /// See [`ScalarUDFImpl::is_strict`] for more details. + pub fn is_strict(&self) -> bool { + self.inner.is_strict() + } + /// Returns this function's [`Signature`] (what input types are accepted). /// /// See [`ScalarUDFImpl::signature`] for more details. @@ -372,6 +380,11 @@ impl ScalarUDF { self.inner.preserves_lex_ordering(inputs) } + /// See [`ScalarUDFImpl::strictly_order_preserving`] for more details. + pub fn strictly_order_preserving(&self, inputs: &[ExprProperties]) -> Result { + self.inner.strictly_order_preserving(inputs) + } + /// See [`ScalarUDFImpl::coerce_types`] for more details. pub fn coerce_types(&self, arg_types: &[DataType]) -> Result> { self.inner.coerce_types(arg_types) @@ -693,6 +706,20 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { true } + /// Returns true if this function always returns NULL when any argument is + /// NULL. + /// + /// Strict functions are NULL-propagating: if any argument evaluates to + /// NULL, the function result is guaranteed to be NULL. Optimizer rules can + /// use this property when reasoning about expression nullability and + /// null-rejecting filters. + /// + /// Defaults to `false` because user-defined functions may choose to accept + /// NULL inputs and produce non-NULL results. + fn is_strict(&self) -> bool { + false + } + /// Invoke the function returning the appropriate result. /// /// # Performance @@ -958,11 +985,19 @@ pub trait ScalarUDFImpl: Debug + DynEq + DynHash + Send + Sync + Any { /// Returns true if the function preserves lexicographical ordering based on /// the input ordering. /// - /// For example, `concat(a || b)` preserves lexicographical ordering, but `abs(a)` does not. + /// See [`ExprProperties::preserves_lex_ordering`] for more details fn preserves_lex_ordering(&self, _inputs: &[ExprProperties]) -> Result { Ok(false) } + /// Returns true if the function is strictly order-preserving with respect + /// to its `Ordered` inputs, i.e. `a.cmp(b) == f(a).cmp(f(b))`. + /// + /// See [`ExprProperties::strictly_order_preserving`] for more details + fn strictly_order_preserving(&self, _inputs: &[ExprProperties]) -> Result { + Ok(false) + } + /// Coerce arguments of a function call to types that the function can evaluate. /// /// This function is only called if [`ScalarUDFImpl::signature`] returns @@ -1103,6 +1138,10 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.is_nullable(args, schema) } + fn is_strict(&self) -> bool { + self.inner.is_strict() + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { self.inner.invoke_with_args(args) } @@ -1170,6 +1209,10 @@ impl ScalarUDFImpl for AliasedScalarUDFImpl { self.inner.preserves_lex_ordering(inputs) } + fn strictly_order_preserving(&self, inputs: &[ExprProperties]) -> Result { + self.inner.strictly_order_preserving(inputs) + } + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { self.inner.coerce_types(arg_types) } diff --git a/datafusion/expr/src/udf_eq.rs b/datafusion/expr/src/udf_eq.rs index 5fb0266aef5dd..8766b483137f4 100644 --- a/datafusion/expr/src/udf_eq.rs +++ b/datafusion/expr/src/udf_eq.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::{AggregateUDFImpl, HigherOrderUDF, ScalarUDFImpl, WindowUDFImpl}; +use crate::{AggregateUDFImpl, HigherOrderUDFImpl, ScalarUDFImpl, WindowUDFImpl}; use std::any::Any; use std::fmt::Debug; use std::hash::{DefaultHasher, Hash, Hasher}; @@ -94,7 +94,7 @@ impl UdfPointer for Arc { } } -impl UdfPointer for Arc { +impl UdfPointer for Arc { fn equals(&self, other: &Self::Target) -> bool { self.as_ref().dyn_eq(other) } diff --git a/datafusion/expr/src/utils.rs b/datafusion/expr/src/utils.rs index 22abb454d4e6b..7f79c5cf18c4a 100644 --- a/datafusion/expr/src/utils.rs +++ b/datafusion/expr/src/utils.rs @@ -34,8 +34,8 @@ use datafusion_common::tree_node::{ }; use datafusion_common::utils::get_at_indices; use datafusion_common::{ - Column, DFSchema, DFSchemaRef, HashMap, Result, TableReference, internal_err, - plan_err, + Column, DFSchema, DFSchemaRef, DataFusionError, Diagnostic, HashMap, Result, Span, + TableReference, internal_err, plan_datafusion_err, plan_err, }; #[cfg(not(feature = "sql"))] @@ -652,6 +652,106 @@ pub fn find_aggregate_exprs<'a>(exprs: impl IntoIterator) -> Ve }) } +/// Returns an error if any of `exprs` nests aggregate or window function calls +/// in a way that has no physical equivalent: an aggregate call may not contain +/// another aggregate call (`sum(sum(x))`) or a window call +/// (`sum(sum(x) OVER ())`), and a window call may not contain another window +/// call (`sum(sum(x) OVER ()) OVER ()`). The reverse nesting, an aggregate used +/// as the argument of a window call (`sum(sum(x)) OVER ()`), is legal: there the +/// aggregate is evaluated by the `Aggregate` node and the window function is +/// evaluated on top of its result. +/// +/// Such expressions are not valid SQL either, so they are rejected while the +/// logical plan is built rather than failing later with an error that does not +/// point back at the original SQL. +/// +/// [`Aggregate::try_new`] and [`Window::try_new`] call this, so the SQL planner +/// and the `DataFrame`/`LogicalPlanBuilder` paths are checked without callers +/// invoking it directly. The lower-level `try_new_with_schema` constructors and +/// building a `Window` from its public fields bypass the check, so a caller +/// that constructs those nodes by hand should call this itself. +/// +/// [`Aggregate::try_new`]: crate::logical_plan::Aggregate::try_new +/// [`Window::try_new`]: crate::logical_plan::Window::try_new +pub(crate) fn check_aggregate_and_window_nesting<'a>( + exprs: impl IntoIterator, +) -> Result<()> { + for expr in exprs { + expr.apply(|outer| { + if !matches!(outer, Expr::AggregateFunction(_) | Expr::WindowFunction(_)) { + return Ok(TreeNodeRecursion::Continue); + } + + // Look for an illegally nested call in the arguments, `FILTER`, + // `ORDER BY` and `PARTITION BY` of this call + let mut err = None; + outer.apply_children(|child| { + child.apply(|inner| { + err = illegal_nesting_err(outer, inner); + if err.is_some() { + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + + match err { + Some(err) => Err(err), + None => Ok(TreeNodeRecursion::Continue), + } + })?; + } + Ok(()) +} + +/// The planning error for a call to `inner` nested inside a call to `outer`, or +/// `None` if that nesting is legal. +fn illegal_nesting_err(outer: &Expr, inner: &Expr) -> Option { + // Messages follow PostgreSQL, which rejects the same three cases + let (message, help) = match (outer, inner) { + (Expr::AggregateFunction(_), Expr::AggregateFunction(_)) => ( + "Aggregate function calls cannot be nested", + format!("Compute '{inner}' in an inner query and aggregate its result"), + ), + (Expr::AggregateFunction(_), Expr::WindowFunction(_)) => ( + "Aggregate function calls cannot contain window function calls", + format!("Compute '{inner}' in an inner query and aggregate its result"), + ), + (Expr::WindowFunction(_), Expr::WindowFunction(_)) => ( + "Window function calls cannot be nested", + format!("Compute '{inner}' in an inner query and use its result here"), + ), + // Anything else, including an aggregate inside a window call + _ => return None, + }; + + Some( + plan_datafusion_err!("{message}: '{inner}' is nested inside '{outer}'") + .with_diagnostic( + Diagnostic::new_error(message, first_span(inner)).with_help(help, None), + ), + ) +} + +/// Best effort source location for `expr`: the first [`Span`] found in its +/// subtree. Only some expressions (currently columns) carry spans, so pointing +/// at e.g. the column of `sum(x)` is the closest we can get to the location of +/// the whole expression. +fn first_span(expr: &Expr) -> Option { + let mut span = None; + expr.apply(|e| { + span = e.spans().and_then(|spans| spans.first()); + if span.is_some() { + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + .ok()?; + span +} + /// Collect all deeply nested `Expr::WindowFunction`. They are returned in order of occurrence /// (depth first), with duplicates omitted. pub fn find_window_exprs<'a>(exprs: impl IntoIterator) -> Vec { @@ -1917,4 +2017,91 @@ mod tests { substr(string: String, start_pos: Int64, length: Int64) "); } + + /// `sum() OVER ()` + fn sum_over(args: Vec) -> Expr { + Expr::from(WindowFunction::new( + WindowFunctionDefinition::AggregateUDF(sum_udaf()), + args, + )) + } + + #[test] + fn test_check_aggregate_and_window_nesting_ok() -> Result<()> { + use crate::test::function_stub::{count, sum}; + + let exprs = [ + // a plain aggregate, and one wrapped in a scalar expression + sum(col("a")), + count(col("a")) + lit(1), + // a window function over a column, and over an aggregate + sum_over(vec![col("a")]), + sum_over(vec![sum(col("a"))]), + ]; + + check_aggregate_and_window_nesting(exprs.iter())?; + Ok(()) + } + + #[test] + fn test_check_aggregate_and_window_nesting_err() { + use crate::test::function_stub::{count, sum}; + use insta::assert_snapshot; + + // an aggregate directly inside an aggregate + let err = check_aggregate_and_window_nesting([&sum(sum(col("a")))]).unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(a)' is nested inside 'sum(sum(a))'" + ); + + // nested below another expression in the arguments + let err = check_aggregate_and_window_nesting([&sum(col("a") + count(col("b")))]) + .unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'COUNT(b)' is nested inside 'sum(a + COUNT(b))'" + ); + + // nested in the FILTER of an aggregate + let filtered = sum(col("a")) + .filter(sum(col("b")).gt(lit(0))) + .build() + .unwrap(); + let err = check_aggregate_and_window_nesting([&filtered]).unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(b)' is nested inside 'sum(a) FILTER (WHERE sum(b) > Int32(0))'" + ); + + // nested in the ORDER BY of an aggregate + let ordered = sum(col("a")) + .order_by(vec![Sort::new(sum(col("b")), true, false)]) + .build() + .unwrap(); + let err = check_aggregate_and_window_nesting([&ordered]).unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(b)' is nested inside 'sum(a) ORDER BY [sum(b) ASC NULLS LAST]'" + ); + + // a window function inside an aggregate + let err = check_aggregate_and_window_nesting([&sum(sum_over(vec![col("a")]))]) + .unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot contain window function calls: 'sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)'" + ); + + // a window function inside a window function + let err = + check_aggregate_and_window_nesting([&sum_over(vec![sum_over(vec![col( + "a", + )])])]) + .unwrap_err(); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Window function calls cannot be nested: 'sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(a) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'" + ); + } } diff --git a/datafusion/expr/src/window_state.rs b/datafusion/expr/src/window_state.rs index f8d4609d3690c..b4d3d09069b14 100644 --- a/datafusion/expr/src/window_state.rs +++ b/datafusion/expr/src/window_state.rs @@ -101,6 +101,31 @@ impl WindowAggState { Ok(()) } + /// Returns true when this state is fully up to date with the partition's + /// buffered batch, meaning another evaluation pass over the partition could + /// not produce any new results or change any state: + /// + /// - `last_calculated_index` has reached the end of the partition's + /// buffered batch, so every row of this partition that has arrived so + /// far already has a result. + /// - When a partition ends, a final evaluation pass is needed to bring + /// derived state up to date. + #[inline] + pub fn is_up_to_date_with( + &self, + partition_batch_state: &PartitionBatchState, + ) -> bool { + let all_rows_have_results = + self.last_calculated_index == partition_batch_state.record_batch.num_rows(); + if all_rows_have_results { + debug_assert_eq!(self.n_row_result_missing, 0); + } + + // `self.is_end` holds the flag as of the previous evaluation pass. + let partition_just_ended = !self.is_end && partition_batch_state.is_end; + all_rows_have_results && !partition_just_ended + } + pub fn new(out_type: &DataType) -> Result { let empty_out_col = ScalarValue::try_from(out_type)?.to_array_of_size(0)?; Ok(Self { @@ -248,14 +273,9 @@ impl WindowFrameContext { pub struct PartitionBatchState { /// The record batch belonging to current partition pub record_batch: RecordBatch, - /// The record batch that contains the most recent row at the input. - /// Please note that this batch doesn't necessarily have the same partitioning - /// with `record_batch`. Keeping track of this batch enables us to prune - /// `record_batch` when cardinality of the partition is sparse. - pub most_recent_row: Option, /// Flag indicating whether we have received all data for this partition pub is_end: bool, - /// Number of rows emitted for each partition + /// Number of rows emitted for this partition since the last pruning pass pub n_out_row: usize, } @@ -263,7 +283,6 @@ impl PartitionBatchState { pub fn new(schema: SchemaRef) -> Self { Self { record_batch: RecordBatch::new_empty(schema), - most_recent_row: None, is_end: false, n_out_row: 0, } @@ -272,7 +291,6 @@ impl PartitionBatchState { pub fn new_with_batch(batch: RecordBatch) -> Self { Self { record_batch: batch, - most_recent_row: None, is_end: false, n_out_row: 0, } @@ -283,12 +301,6 @@ impl PartitionBatchState { concat_batches(&self.record_batch.schema(), [&self.record_batch, batch])?; Ok(()) } - - pub fn set_most_recent_row(&mut self, batch: RecordBatch) { - // It is enough for the batch to contain only a single row (the rest - // are not necessary). - self.most_recent_row = Some(batch); - } } /// This structure encapsulates all the state information we require as we scan diff --git a/datafusion/ffi/Cargo.toml b/datafusion/ffi/Cargo.toml index 7eed11c0c69e8..affcff3dbdcd9 100644 --- a/datafusion/ffi/Cargo.toml +++ b/datafusion/ffi/Cargo.toml @@ -63,7 +63,7 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-optimizer = { workspace = true } datafusion-physical-plan = { workspace = true } -datafusion-proto = { workspace = true } +datafusion-proto = { workspace = true, default-features = false } datafusion-proto-common = { workspace = true } datafusion-session = { workspace = true } futures = { workspace = true } @@ -71,7 +71,7 @@ libloading = "0.9" log = { workspace = true } prost = { workspace = true } semver = "1.0.28" -stabby = "72.1.1" +stabby = "72.1.2" tokio = { workspace = true } [dev-dependencies] @@ -83,10 +83,11 @@ datafusion-functions-window = { workspace = true } doc-comment = { workspace = true } [features] +default = ["parquet"] integration-tests = [ "datafusion-functions", "datafusion-functions-aggregate", "datafusion-functions-table", "datafusion-functions-window", ] -tarpaulin_include = [] # Exists only to prevent warnings on stable and still have accurate coverage +parquet = ["datafusion-proto/parquet"] diff --git a/datafusion/ffi/src/arrow_wrappers.rs b/datafusion/ffi/src/arrow_wrappers.rs index 1c921b0f83b1e..62fb36f836785 100644 --- a/datafusion/ffi/src/arrow_wrappers.rs +++ b/datafusion/ffi/src/arrow_wrappers.rs @@ -49,7 +49,6 @@ impl From for WrappedSchema { /// Since going through the FFI always has the potential to fail, we need to catch these errors, /// give the user a warning, and return some kind of result. In this case we default to an /// empty schema. -#[cfg(not(tarpaulin_include))] fn catch_df_schema_error(e: &ArrowError) -> Schema { error!( "Unable to convert from FFI_ArrowSchema to DataFusion Schema in FFI_PlanProperties. {e}" diff --git a/datafusion/ffi/src/execution_plan.rs b/datafusion/ffi/src/execution_plan.rs index ddad605081745..d7ee5dace30cc 100644 --- a/datafusion/ffi/src/execution_plan.rs +++ b/datafusion/ffi/src/execution_plan.rs @@ -25,7 +25,8 @@ use datafusion_common::{DataFusionError, Result, Statistics}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr_common::metrics::MetricsSet; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, StatisticsArgs, StatisticsContext, }; use stabby::string::String as SString; use stabby::vec::Vec as SVec; @@ -33,6 +34,7 @@ use tokio::runtime::Handle; use crate::config::FFI_ConfigOptions; use crate::execution::FFI_TaskContext; +use crate::physical_expr::FFI_PhysicalExpr; use crate::physical_expr::metrics::FFI_MetricsSet; use crate::plan_properties::FFI_PlanProperties; use crate::record_batch_stream::FFI_RecordBatchStream; @@ -50,6 +52,14 @@ pub struct FFI_ExecutionPlan { /// Return a vector of children plans pub children: unsafe extern "C" fn(plan: &Self) -> SVec, + /// Return the physical expression roots owned by this plan node. + pub apply_expressions: + unsafe extern "C" fn(plan: &Self) -> FFI_Result>, + + /// Return the dynamic expressions produced by this plan node. + pub dynamic_expressions_produced: + unsafe extern "C" fn(plan: &Self) -> SVec, + pub with_new_children: unsafe extern "C" fn(plan: &Self, children: SVec) -> FFI_Result, @@ -91,6 +101,9 @@ pub struct FFI_ExecutionPlan { /// Release the memory of the private data when it is no longer being used. pub release: unsafe extern "C" fn(arg: &mut Self), + /// Return the major DataFusion version number of this provider. + pub version: unsafe extern "C" fn() -> u64, + /// Internal data. This is only to be accessed by the provider of the plan. /// A [`ForeignExecutionPlan`] should never attempt to access this data. pub private_data: *mut c_void, @@ -138,6 +151,27 @@ unsafe extern "C" fn children_fn_wrapper( .collect() } +unsafe extern "C" fn apply_expressions_fn_wrapper( + plan: &FFI_ExecutionPlan, +) -> FFI_Result> { + let mut expressions = SVec::new(); + let result = plan.inner().apply_expressions(&mut |expr| { + expressions.push(FFI_PhysicalExpr::from(Arc::clone(expr))); + Ok(TreeNodeRecursion::Continue) + }); + sresult!(result.map(|_| expressions)) +} + +unsafe extern "C" fn dynamic_expressions_produced_fn_wrapper( + plan: &FFI_ExecutionPlan, +) -> SVec { + plan.inner() + .dynamic_expressions_produced() + .into_iter() + .map(FFI_PhysicalExpr::from) + .collect() +} + unsafe extern "C" fn with_new_children_fn_wrapper( plan: &FFI_ExecutionPlan, children: SVec, @@ -151,7 +185,10 @@ unsafe extern "C" fn with_new_children_fn_wrapper( .collect(); let children = sresult_return!(children); - let new_plan = sresult_return!(inner_plan.with_new_children(children)); + let new_plan = sresult_return!(inner_plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute) + )); FFI_Result::Ok(FFI_ExecutionPlan::new(new_plan, runtime)) } @@ -210,8 +247,11 @@ unsafe extern "C" fn partition_statistics_fn_wrapper( partition: FFI_Option, ) -> FFI_Result> { let partition: Option = partition.into(); - plan.inner() - .partition_statistics(partition) + StatisticsContext::new() + .compute( + plan.inner().as_ref(), + &StatisticsArgs::new().with_partition(partition), + ) .map(|stats| SVec::from(serialize_statistics(stats.as_ref()).as_slice())) .into() } @@ -265,7 +305,7 @@ fn pass_runtime_to_children( // If the parent is foreign and the child is local to this library, then when // we called `children()` above we will get something other than a // `ForeignExecutionPlan`. In this case wrap the plan in a `ForeignExecutionPlan` - // because when we call `with_new_children` below it will extract the + // because when we call `replace_children` below it will extract the // FFI plan that does contain the runtime. if plan_is_foreign && !child.is::() { updated_children = true; @@ -278,7 +318,12 @@ fn pass_runtime_to_children( }) .collect::>>()?; if updated_children { - Arc::clone(plan).with_new_children(children).map(Some) + Arc::clone(plan) + .replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + .map(Some) } else { Ok(None) } @@ -303,6 +348,8 @@ impl FFI_ExecutionPlan { Self { properties: properties_fn_wrapper, children: children_fn_wrapper, + apply_expressions: apply_expressions_fn_wrapper, + dynamic_expressions_produced: dynamic_expressions_produced_fn_wrapper, with_new_children: with_new_children_fn_wrapper, name: name_fn_wrapper, execute: execute_fn_wrapper, @@ -311,6 +358,7 @@ impl FFI_ExecutionPlan { partition_statistics: partition_statistics_fn_wrapper, clone: clone_fn_wrapper, release: release_fn_wrapper, + version: crate::version, private_data: Box::into_raw(private_data) as *mut c_void, library_marker_id: crate::get_library_marker_id, } @@ -413,9 +461,10 @@ impl ExecutionPlan for ForeignExecutionPlan { self.children.iter().collect() } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { let children = children .into_iter() @@ -427,6 +476,16 @@ impl ExecutionPlan for ForeignExecutionPlan { (&new_plan).try_into() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -442,17 +501,28 @@ impl ExecutionPlan for ForeignExecutionPlan { fn apply_expressions( &self, f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.properties.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + let expressions = + df_result!(unsafe { (self.plan.apply_expressions)(&self.plan) })?; + datafusion_physical_plan::apply_expression_roots( + expressions.iter().map(|expression| { + let expression: Arc = + expression.into(); + expression + }), + f, + ) + } + + fn dynamic_expressions_produced( + &self, + ) -> Vec> { + unsafe { (self.plan.dynamic_expressions_produced)(&self.plan) } + .iter() + .map(>::from) + .collect() } fn repartitioned( @@ -487,8 +557,9 @@ impl ExecutionPlan for ForeignExecutionPlan { #[cfg(any(test, feature = "integration-tests"))] pub mod tests { - use datafusion_physical_plan::Partitioning; + use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; + use datafusion_physical_plan::{Partitioning, PhysicalExpr}; use super::*; @@ -496,6 +567,8 @@ pub mod tests { pub struct EmptyExec { props: Arc, children: Vec>, + expressions: Vec>, + dynamic_expressions: Vec>, metrics: Option, statistics: Option, } @@ -510,6 +583,8 @@ pub mod tests { Boundedness::Bounded, )), children: Vec::default(), + expressions: Vec::default(), + dynamic_expressions: Vec::default(), metrics: None, statistics: None, } @@ -524,6 +599,22 @@ pub mod tests { self.statistics = Some(statistics); self } + + pub fn with_expressions( + mut self, + expressions: Vec>, + ) -> Self { + self.expressions = expressions; + self + } + + pub fn with_dynamic_expressions( + mut self, + dynamic_expressions: Vec>, + ) -> Self { + self.dynamic_expressions = dynamic_expressions; + self + } } impl DisplayAs for EmptyExec { @@ -549,18 +640,31 @@ pub mod tests { self.children.iter().collect() } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(EmptyExec { props: Arc::clone(&self.props), children, + expressions: self.expressions.clone(), + dynamic_expressions: self.dynamic_expressions.clone(), metrics: self.metrics.clone(), statistics: self.statistics.clone(), })) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -569,13 +673,18 @@ pub mod tests { unimplemented!() } + fn dynamic_expressions_produced(&self) -> Vec> { + self.dynamic_expressions.iter().map(Arc::clone).collect() + } + fn metrics(&self) -> Option { self.metrics.clone() } - fn partition_statistics( + fn statistics_from_inputs( &self, - _partition: Option, + _input_stats: &[Arc], + _args: &StatisticsArgs, ) -> Result> { Ok(Arc::new(self.statistics.clone().unwrap_or_else(|| { Statistics::new_unknown(self.props.eq_properties.schema()) @@ -584,21 +693,16 @@ pub mod tests { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, - ) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.props.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + datafusion_physical_plan::apply_expression_roots(&self.expressions, f) } } + pub(crate) fn create_dynamic_filter() -> Arc { + Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))) + } + #[test] fn test_round_trip_ffi_execution_plan() -> Result<()> { let schema = Arc::new(arrow::datatypes::Schema::new(vec![ @@ -628,6 +732,61 @@ pub mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_apply_expressions() -> Result<()> { + let schema = Arc::new(arrow::datatypes::Schema::empty()); + let dynamic_filter = create_dynamic_filter(); + let expected_id = dynamic_filter + .expression_id() + .expect("dynamic filters always have an expression ID"); + let expression: Arc = Arc::clone(&dynamic_filter) as _; + let original_plan = + Arc::new(EmptyExec::new(schema).with_expressions(vec![expression])); + + let mut ffi_plan = FFI_ExecutionPlan::new(original_plan, None); + ffi_plan.library_marker_id = crate::mock_foreign_marker_id; + let foreign_plan: Arc = (&ffi_plan).try_into()?; + + let mut retained = None; + foreign_plan.apply_expressions(&mut |expr| { + retained = Some(Arc::clone(expr)); + Ok(TreeNodeRecursion::Continue) + })?; + drop(foreign_plan); + + assert_eq!( + retained.and_then(|expr| expr.expression_id()), + Some(expected_id) + ); + Ok(()) + } + + #[test] + fn test_ffi_execution_plan_dynamic_expressions_produced() -> Result<()> { + let schema = Arc::new(arrow::datatypes::Schema::empty()); + let dynamic_filter = create_dynamic_filter(); + let expected_id = dynamic_filter + .expression_id() + .expect("dynamic filters always have an expression ID"); + let expression: Arc = Arc::clone(&dynamic_filter) as _; + let original_plan = + Arc::new(EmptyExec::new(schema).with_dynamic_expressions(vec![expression])); + + let mut ffi_plan = FFI_ExecutionPlan::new(original_plan, None); + ffi_plan.library_marker_id = crate::mock_foreign_marker_id; + let foreign_plan: Arc = (&ffi_plan).try_into()?; + foreign_plan.check_invariants( + datafusion_physical_plan::execution_plan::InvariantLevel::Always, + )?; + + let produced = foreign_plan.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + assert_eq!(produced[0].expression_id(), Some(expected_id)); + drop(foreign_plan); + assert_eq!(produced[0].expression_id(), Some(expected_id)); + Ok(()) + } + #[test] fn test_ffi_execution_plan_children() -> Result<()> { let schema = Arc::new(arrow::datatypes::Schema::new(vec![ @@ -648,7 +807,10 @@ pub mod tests { assert_eq!(parent_foreign.children().len(), 0); assert_eq!(child_foreign.children().len(), 0); - let parent_foreign = parent_foreign.with_new_children(vec![child_foreign])?; + let parent_foreign = parent_foreign.replace_children( + vec![child_foreign], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; assert_eq!(parent_foreign.children().len(), 1); // Version 2: Adding child to the local plan @@ -658,7 +820,10 @@ pub mod tests { let child_foreign = >::try_from(&child_local)?; let parent_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let parent_plan = parent_plan.with_new_children(vec![child_foreign])?; + let parent_plan = parent_plan.replace_children( + vec![child_foreign], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; let mut parent_local = FFI_ExecutionPlan::new(parent_plan, None); parent_local.library_marker_id = crate::mock_foreign_marker_id; let parent_foreign = >::try_from(&parent_local)?; @@ -705,27 +870,34 @@ pub mod tests { Ok(()) } - #[test] - fn test_ffi_execution_plan_partition_statistics_round_trip() -> Result<()> { + /// Build an `EmptyExec` carrying `statistics`, then export it across the + /// (mock) FFI boundary and return the resulting foreign plan. + #[cfg(test)] + fn export_empty_exec_over_ffi( + schema: &arrow::datatypes::SchemaRef, + statistics: Option, + ) -> Result> { + let mut plan = EmptyExec::new(Arc::clone(schema)); + if let Some(statistics) = statistics { + plan = plan.with_statistics(statistics); + } + let mut local = FFI_ExecutionPlan::new(Arc::new(plan), None); + local.library_marker_id = crate::mock_foreign_marker_id; + let foreign: Arc = (&local).try_into()?; + Ok(foreign) + } + + /// Schema and a fully-populated `Statistics` (including `ScalarValue`-typed + /// min/max) shared by the FFI statistics round-trip tests. + #[cfg(test)] + fn stats_round_trip_fixture() -> (arrow::datatypes::SchemaRef, Statistics) { use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; let schema = Arc::new(arrow::datatypes::Schema::new(vec![ arrow::datatypes::Field::new("a", arrow::datatypes::DataType::Int32, true), ])); - - // Plans without explicit statistics return Statistics::new_unknown across - // the boundary. - let bare_plan = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let mut bare_local = FFI_ExecutionPlan::new(bare_plan, None); - bare_local.library_marker_id = crate::mock_foreign_marker_id; - let bare_foreign: Arc = (&bare_local).try_into()?; - let bare_stats = bare_foreign.partition_statistics(None)?; - assert_eq!(bare_stats.as_ref(), &Statistics::new_unknown(&schema)); - - // Plans with statistics round-trip them faithfully, including - // ScalarValue-typed min/max. - let original_stats = Statistics { + let statistics = Statistics { num_rows: Precision::Exact(7), total_byte_size: Precision::Inexact(128), column_statistics: vec![ColumnStatistics { @@ -737,18 +909,72 @@ pub mod tests { byte_size: Precision::Exact(28), }], }; - let stats_plan = Arc::new( - EmptyExec::new(Arc::clone(&schema)).with_statistics(original_stats.clone()), + (schema, statistics) + } + + /// Statistics survive an FFI round trip when queried through the + /// **deprecated** `partition_statistics` entry point on the foreign plan. + #[test] + #[expect(deprecated)] + fn test_ffi_execution_plan_partition_statistics_round_trip() -> Result<()> { + let (schema, original_stats) = stats_round_trip_fixture(); + + // A plan without explicit statistics reports new_unknown. + let bare = export_empty_exec_over_ffi(&schema, None)?; + assert_eq!( + bare.partition_statistics(None)?.as_ref(), + &Statistics::new_unknown(&schema) + ); + + // A plan with statistics round-trips them for overall and per-partition queries. + let with_stats = + export_empty_exec_over_ffi(&schema, Some(original_stats.clone()))?; + assert_eq!( + with_stats.partition_statistics(None)?.as_ref(), + &original_stats + ); + assert_eq!( + with_stats.partition_statistics(Some(1))?.as_ref(), + &original_stats ); - let mut stats_local = FFI_ExecutionPlan::new(stats_plan, None); - stats_local.library_marker_id = crate::mock_foreign_marker_id; - let stats_foreign: Arc = (&stats_local).try_into()?; - let observed = stats_foreign.partition_statistics(None)?; - assert_eq!(observed.as_ref(), &original_stats); + Ok(()) + } + + /// Same round trip as + /// [`test_ffi_execution_plan_partition_statistics_round_trip`], but queried + /// through the **new** `StatisticsContext::compute` entry point. + #[test] + fn test_ffi_execution_plan_statistics_context_round_trip() -> Result<()> { + let (schema, original_stats) = stats_round_trip_fixture(); + + // A plan without explicit statistics reports new_unknown. + let bare = export_empty_exec_over_ffi(&schema, None)?; + assert_eq!( + StatisticsContext::new() + .compute(bare.as_ref(), &StatisticsArgs::new())? + .as_ref(), + &Statistics::new_unknown(&schema) + ); - let observed_partition = stats_foreign.partition_statistics(Some(1))?; - assert_eq!(observed_partition.as_ref(), &original_stats); + // A plan with statistics round-trips them for overall and per-partition queries. + let with_stats = + export_empty_exec_over_ffi(&schema, Some(original_stats.clone()))?; + assert_eq!( + StatisticsContext::new() + .compute(with_stats.as_ref(), &StatisticsArgs::new())? + .as_ref(), + &original_stats + ); + assert_eq!( + StatisticsContext::new() + .compute( + with_stats.as_ref(), + &StatisticsArgs::new().with_partition(Some(1)), + )? + .as_ref(), + &original_stats + ); Ok(()) } diff --git a/datafusion/ffi/src/expr/expr_properties.rs b/datafusion/ffi/src/expr/expr_properties.rs index 5b37cc6a28535..584f774c7b26e 100644 --- a/datafusion/ffi/src/expr/expr_properties.rs +++ b/datafusion/ffi/src/expr/expr_properties.rs @@ -29,6 +29,7 @@ pub struct FFI_ExprProperties { sort_properties: FFI_SortProperties, range: FFI_Interval, preserves_lex_ordering: bool, + strictly_order_preserving: bool, } impl TryFrom<&ExprProperties> for FFI_ExprProperties { @@ -41,6 +42,7 @@ impl TryFrom<&ExprProperties> for FFI_ExprProperties { sort_properties, range, preserves_lex_ordering: value.preserves_lex_ordering, + strictly_order_preserving: value.strictly_order_preserving, }) } } @@ -54,6 +56,7 @@ impl TryFrom for ExprProperties { sort_properties, range, preserves_lex_ordering: value.preserves_lex_ordering, + strictly_order_preserving: value.strictly_order_preserving, }) } } diff --git a/datafusion/ffi/src/lib.rs b/datafusion/ffi/src/lib.rs index 4df6c4b570f34..de8f8cba9ca9b 100644 --- a/datafusion/ffi/src/lib.rs +++ b/datafusion/ffi/src/lib.rs @@ -36,8 +36,10 @@ pub mod ffi_option; pub mod insert_op; pub mod physical_expr; pub mod physical_optimizer; +pub mod placement; pub mod plan_properties; pub mod proto; +pub mod query_planner; pub mod record_batch_stream; pub mod schema_provider; pub mod session; diff --git a/datafusion/ffi/src/physical_expr/metrics.rs b/datafusion/ffi/src/physical_expr/metrics.rs index ebef728e0520d..763cc3f079a01 100644 --- a/datafusion/ffi/src/physical_expr/metrics.rs +++ b/datafusion/ffi/src/physical_expr/metrics.rs @@ -128,6 +128,9 @@ pub struct FFI_RatioMetrics { } /// FFI-stable mirror of [`MetricValue`]. +/// +/// This is part of the stable ABI and must not be reordered. New variants must be +/// appended at the end. #[repr(C, u8)] #[derive(Debug, Clone)] pub enum FFI_MetricValue { @@ -170,6 +173,10 @@ pub enum FFI_MetricValue { display: SString, as_usize_value: u64, }, + PeakMemoryUsage { + name: SString, + gauge: u64, + }, } // ----------------------------------------------------------------------------- @@ -425,6 +432,10 @@ impl From<&MetricValue> for FFI_MetricValue { name: SString::from(name.as_ref()), gauge: gauge.value() as u64, }, + MetricValue::PeakMemoryUsage { name, gauge } => Self::PeakMemoryUsage { + name: SString::from(name.as_ref()), + gauge: gauge.value() as u64, + }, MetricValue::Time { name, time } => Self::Time { name: SString::from(name.as_ref()), time_ns: time.value() as u64, @@ -481,6 +492,10 @@ impl From for MetricValue { name: Cow::Owned(name.into()), gauge: gauge_from_value(gauge), }, + FFI_MetricValue::PeakMemoryUsage { name, gauge } => Self::PeakMemoryUsage { + name: Cow::Owned(name.into()), + gauge: gauge_from_value(gauge), + }, FFI_MetricValue::Time { name, time_ns } => Self::Time { name: Cow::Owned(name.into()), time: time_from_nanos(time_ns), @@ -624,6 +639,13 @@ mod tests { gauge, }); + let peak_memory = Gauge::new(); + peak_memory.add(44); + assert_value_roundtrip(MetricValue::PeakMemoryUsage { + name: Cow::Borrowed("peak_mem_used"), + gauge: peak_memory, + }); + let time = Time::new(); time.add_duration(std::time::Duration::from_nanos(33)); assert_value_roundtrip(MetricValue::Time { diff --git a/datafusion/ffi/src/physical_expr/mod.rs b/datafusion/ffi/src/physical_expr/mod.rs index 9a3ee273936c3..8e6676b6e41bf 100644 --- a/datafusion/ffi/src/physical_expr/mod.rs +++ b/datafusion/ffi/src/physical_expr/mod.rs @@ -120,6 +120,8 @@ pub struct FFI_PhysicalExpr { pub is_volatile_node: unsafe extern "C" fn(&Self) -> bool, + pub expression_id: unsafe extern "C" fn(&Self) -> FFI_Option, + // Display trait pub display: unsafe extern "C" fn(&Self) -> SString, @@ -387,6 +389,13 @@ unsafe extern "C" fn is_volatile_node_fn_wrapper(expr: &FFI_PhysicalExpr) -> boo let expr = expr.inner(); expr.is_volatile_node() } + +unsafe extern "C" fn expression_id_fn_wrapper( + expr: &FFI_PhysicalExpr, +) -> FFI_Option { + expr.inner().expression_id().into() +} + unsafe extern "C" fn display_fn_wrapper(expr: &FFI_PhysicalExpr) -> SString { let expr = expr.inner(); format!("{expr}").into() @@ -434,6 +443,7 @@ unsafe extern "C" fn clone_fn_wrapper(expr: &FFI_PhysicalExpr) -> FFI_PhysicalEx snapshot: snapshot_fn_wrapper, snapshot_generation: snapshot_generation_fn_wrapper, is_volatile_node: is_volatile_node_fn_wrapper, + expression_id: expression_id_fn_wrapper, display: display_fn_wrapper, hash: hash_fn_wrapper, clone: clone_fn_wrapper, @@ -477,6 +487,7 @@ impl From> for FFI_PhysicalExpr { snapshot: snapshot_fn_wrapper, snapshot_generation: snapshot_generation_fn_wrapper, is_volatile_node: is_volatile_node_fn_wrapper, + expression_id: expression_id_fn_wrapper, display: display_fn_wrapper, hash: hash_fn_wrapper, clone: clone_fn_wrapper, @@ -713,6 +724,10 @@ impl PhysicalExpr for ForeignPhysicalExpr { fn is_volatile_node(&self) -> bool { unsafe { (self.expr.is_volatile_node)(&self.expr) } } + + fn expression_id(&self) -> Option { + unsafe { (self.expr.expression_id)(&self.expr) }.into() + } } impl Eq for ForeignPhysicalExpr {} @@ -747,7 +762,9 @@ mod tests { use datafusion_expr::interval_arithmetic::Interval; #[expect(deprecated)] use datafusion_expr::statistics::Distribution; - use datafusion_physical_expr::expressions::{Column, NegativeExpr, NotExpr}; + use datafusion_physical_expr::expressions::{ + Column, DynamicFilterPhysicalExpr, NegativeExpr, NotExpr, lit, + }; use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, fmt_sql}; use crate::physical_expr::FFI_PhysicalExpr; @@ -762,6 +779,21 @@ mod tests { (original, foreign_expr) } + #[test] + fn ffi_physical_expr_expression_id() { + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let expected_id = dynamic_filter + .expression_id() + .expect("dynamic filters always have an expression ID"); + let expression: Arc = + Arc::::clone(&dynamic_filter); + let mut ffi_expr = FFI_PhysicalExpr::from(expression); + ffi_expr.library_marker_id = crate::mock_foreign_marker_id; + + let foreign_expr: Arc = (&ffi_expr).into(); + assert_eq!(foreign_expr.expression_id(), Some(expected_id)); + } + fn test_record_batch() -> RecordBatch { record_batch!(("a", Int32, [1, 2, 3])).unwrap() } diff --git a/datafusion/ffi/src/physical_expr/partitioning.rs b/datafusion/ffi/src/physical_expr/partitioning.rs index 434b6a097e645..2a9a8528c6c3e 100644 --- a/datafusion/ffi/src/physical_expr/partitioning.rs +++ b/datafusion/ffi/src/physical_expr/partitioning.rs @@ -17,20 +17,35 @@ use std::sync::Arc; -use datafusion_physical_expr::Partitioning; +use datafusion_common::{DataFusionError, ScalarValue, SplitPoint}; +use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, +}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use stabby::vec::Vec as SVec; +use crate::arrow_wrappers::WrappedArray; use crate::physical_expr::FFI_PhysicalExpr; +use crate::physical_expr::sort::FFI_PhysicalSortExpr; + +/// A stable struct for sharing [`RangePartitioning`] across FFI boundaries. +/// See [`RangePartitioning`] for the descriptions of each field. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_RangePartitioning { + split_points: SVec>, + ordering: SVec, +} /// A stable struct for sharing [`Partitioning`] across FFI boundaries. -/// See ['Partitioning'] for the meaning of each variant. +/// See [`Partitioning`] for the meaning of each variant. #[repr(C)] #[derive(Debug)] pub enum FFI_Partitioning { RoundRobinBatch(usize), Hash(SVec, usize), UnknownPartitioning(usize), + Range(FFI_RangePartitioning), } impl From<&Partitioning> for FFI_Partitioning { @@ -45,44 +60,130 @@ impl From<&Partitioning> for FFI_Partitioning { .collect(); Self::Hash(exprs, *size) } + Partitioning::Range(range) => { + // Producer-side conversion should be infallible at ABI boundary + let split_points = range + .split_points() + .iter() + .map(|split_point| { + split_point + .values() + .iter() + .map(|value| { + WrappedArray::try_from(value).expect( + "ScalarValue in RangePartitioning should convert to WrappedArray", + ) + }) + .collect() + }) + .collect(); + let ordering = range + .ordering() + .iter() + .map(FFI_PhysicalSortExpr::from) + .collect(); + Self::Range(FFI_RangePartitioning { + split_points, + ordering, + }) + } Partitioning::UnknownPartitioning(size) => Self::UnknownPartitioning(*size), } } } -impl From<&FFI_Partitioning> for Partitioning { - fn from(value: &FFI_Partitioning) -> Self { - match value { +impl TryFrom for Partitioning { + type Error = DataFusionError; + + fn try_from(value: FFI_Partitioning) -> Result { + Ok(match value { FFI_Partitioning::RoundRobinBatch(size) => { - Partitioning::RoundRobinBatch(*size) + Partitioning::RoundRobinBatch(size) } FFI_Partitioning::Hash(exprs, size) => { let exprs = exprs.iter().map(>::from).collect(); - Self::Hash(exprs, *size) + Self::Hash(exprs, size) + } + FFI_Partitioning::Range(range) => { + let split_points = range + .split_points + .into_iter() + .map(|split_point| { + split_point + .into_iter() + .map(ScalarValue::try_from) + .collect::, _>>() + .map(SplitPoint::new) + }) + .collect::, _>>()?; + + let ordering = + LexOrdering::new(range.ordering.iter().map(PhysicalSortExpr::from)) + .ok_or_else(|| { + DataFusionError::Internal( + "FFI Range partitioning ordering must be non-empty" + .to_string(), + ) + })?; + + Self::Range(RangePartitioning::try_new(ordering, split_points)?) } FFI_Partitioning::UnknownPartitioning(size) => { - Self::UnknownPartitioning(*size) + Self::UnknownPartitioning(size) } - } + }) } } #[cfg(test)] mod tests { - use datafusion_physical_expr::Partitioning; - use datafusion_physical_expr::expressions::lit; + use std::sync::Arc; + + use arrow_schema::SortOptions; + use datafusion_common::{Result, ScalarValue, SplitPoint}; + use datafusion_physical_expr::expressions::{Column, lit}; + use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, + }; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use stabby::vec::Vec as SVec; - use crate::physical_expr::partitioning::FFI_Partitioning; + use crate::physical_expr::partitioning::{FFI_Partitioning, FFI_RangePartitioning}; + + fn range_partitioning() -> Result { + let a = Arc::new(Column::new("a", 0)) as Arc; + let b = Arc::new(Column::new("b", 1)) as Arc; + let ordering = LexOrdering::new([ + PhysicalSortExpr::new(a, SortOptions::default()), + PhysicalSortExpr::new(b, SortOptions::new(true, false)), + ]) + .expect("non-empty ordering"); + let split_points = vec![ + SplitPoint::new(vec![ + ScalarValue::Int64(Some(10)), + ScalarValue::Utf8(Some("a".to_string())), + ]), + SplitPoint::new(vec![ + ScalarValue::Int64(Some(20)), + ScalarValue::Utf8(Some("b".to_string())), + ]), + ]; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points, + )?)) + } #[test] - fn round_trip_ffi_partitioning() { + fn round_trip_ffi_partitioning() -> Result<()> { for partitioning in [ Partitioning::RoundRobinBatch(10), Partitioning::Hash(vec![lit(1)], 10), Partitioning::UnknownPartitioning(10), + range_partitioning()?, ] { let ffi_partitioning: FFI_Partitioning = (&partitioning).into(); - let returned: Partitioning = (&ffi_partitioning).into(); + let returned: Partitioning = ffi_partitioning.try_into()?; if let Partitioning::UnknownPartitioning(return_size) = returned { let Partitioning::UnknownPartitioning(original_size) = partitioning @@ -94,5 +195,32 @@ mod tests { assert_eq!(partitioning, returned); } } + + Ok(()) + } + + #[test] + fn round_trip_ffi_range_partitioning_compound_key() -> Result<()> { + let partitioning = range_partitioning()?; + + let ffi_partitioning: FFI_Partitioning = (&partitioning).into(); + let returned: Partitioning = ffi_partitioning.try_into()?; + assert_eq!(partitioning, returned); + + Ok(()) + } + + #[test] + fn ffi_range_partitioning_rejects_empty_ordering() { + let ffi_partitioning = FFI_Partitioning::Range(FFI_RangePartitioning { + split_points: SVec::new(), + ordering: SVec::new(), + }); + + let err = Partitioning::try_from(ffi_partitioning).unwrap_err(); + assert!( + err.to_string().contains("ordering must be non-empty"), + "{err}" + ); } } diff --git a/datafusion/ffi/src/physical_optimizer.rs b/datafusion/ffi/src/physical_optimizer.rs index 84dc40ce8f46c..3fb213208327b 100644 --- a/datafusion/ffi/src/physical_optimizer.rs +++ b/datafusion/ffi/src/physical_optimizer.rs @@ -21,7 +21,7 @@ use std::sync::Arc; use async_trait::async_trait; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; -use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_optimizer::{PhysicalOptimizerContext, PhysicalOptimizerRule}; use datafusion_physical_plan::ExecutionPlan; use stabby::string::String as SString; use tokio::runtime::Handle; @@ -31,6 +31,84 @@ use crate::execution_plan::FFI_ExecutionPlan; use crate::util::FFI_Result; use crate::{df_result, sresult_return}; +/// A stable struct for sharing [`PhysicalOptimizerContext`] across FFI boundaries. +/// +/// This provides access to configuration options for optimizer rules that need +/// extended context beyond the plan itself. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_PhysicalOptimizerContext { + pub config_options: + unsafe extern "C" fn(&FFI_PhysicalOptimizerContext) -> FFI_ConfigOptions, + + /// Release the memory of the private data. + pub release: unsafe extern "C" fn(&mut FFI_PhysicalOptimizerContext), + + /// Internal data. Only accessed by the provider. + pub private_data: *const c_void, +} + +unsafe impl Send for FFI_PhysicalOptimizerContext {} +unsafe impl Sync for FFI_PhysicalOptimizerContext {} + +struct OptimizerContextPrivateData { + config: ConfigOptions, +} + +impl FFI_PhysicalOptimizerContext { + pub fn new(context: &dyn PhysicalOptimizerContext) -> Self { + let private_data = Box::new(OptimizerContextPrivateData { + config: context.config_options().clone(), + }); + let private_data = Box::into_raw(private_data) as *const c_void; + + Self { + config_options: context_config_options_fn, + release: context_release_fn, + private_data, + } + } + + fn inner(&self) -> &OptimizerContextPrivateData { + unsafe { &*(self.private_data as *const OptimizerContextPrivateData) } + } +} + +impl Drop for FFI_PhysicalOptimizerContext { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +unsafe extern "C" fn context_config_options_fn( + ctx: &FFI_PhysicalOptimizerContext, +) -> FFI_ConfigOptions { + FFI_ConfigOptions::from(&ctx.inner().config) +} + +unsafe extern "C" fn context_release_fn(ctx: &mut FFI_PhysicalOptimizerContext) { + if !ctx.private_data.is_null() { + unsafe { + let _ = Box::from_raw(ctx.private_data as *mut OptimizerContextPrivateData); + } + ctx.private_data = std::ptr::null(); + } +} + +/// Reconstructed [`PhysicalOptimizerContext`] on the consumer side of FFI. +/// +/// `StatisticsRegistry` is not plumbed because it contains trait object vtables +/// that are only valid within the originating library. +struct ForeignOptimizerContext { + config: ConfigOptions, +} + +impl PhysicalOptimizerContext for ForeignOptimizerContext { + fn config_options(&self) -> &ConfigOptions { + &self.config + } +} + /// A stable struct for sharing [`PhysicalOptimizerRule`] across FFI boundaries. #[repr(C)] #[derive(Debug)] @@ -55,6 +133,12 @@ pub struct FFI_PhysicalOptimizerRule { /// Return the major DataFusion version number of this rule. pub version: unsafe extern "C" fn() -> u64, + pub optimize_with_context: unsafe extern "C" fn( + &Self, + plan: &FFI_ExecutionPlan, + context: &FFI_PhysicalOptimizerContext, + ) -> FFI_Result, + /// Internal data. This is only to be accessed by the provider of the rule. /// A [`ForeignPhysicalOptimizerRule`] should never attempt to access this data. pub private_data: *mut c_void, @@ -98,6 +182,23 @@ unsafe extern "C" fn optimize_fn_wrapper( FFI_Result::Ok(FFI_ExecutionPlan::new(optimized_plan, runtime)) } +unsafe extern "C" fn optimize_with_context_fn_wrapper( + rule: &FFI_PhysicalOptimizerRule, + plan: &FFI_ExecutionPlan, + context: &FFI_PhysicalOptimizerContext, +) -> FFI_Result { + let runtime = rule.runtime(); + let inner = rule.inner(); + let plan: Arc = sresult_return!(plan.try_into()); + let config = sresult_return!(ConfigOptions::try_from(unsafe { + (context.config_options)(context) + })); + let foreign_ctx = ForeignOptimizerContext { config }; + let optimized_plan = sresult_return!(inner.optimize_with_context(plan, &foreign_ctx)); + + FFI_Result::Ok(FFI_ExecutionPlan::new(optimized_plan, runtime)) +} + unsafe extern "C" fn name_fn_wrapper(rule: &FFI_PhysicalOptimizerRule) -> SString { let rule = rule.inner(); rule.name().into() @@ -127,6 +228,7 @@ unsafe extern "C" fn clone_fn_wrapper( FFI_PhysicalOptimizerRule { optimize: optimize_fn_wrapper, + optimize_with_context: optimize_with_context_fn_wrapper, name: name_fn_wrapper, schema_check: schema_check_fn_wrapper, clone: clone_fn_wrapper, @@ -160,6 +262,7 @@ impl FFI_PhysicalOptimizerRule { Self { optimize: optimize_fn_wrapper, + optimize_with_context: optimize_with_context_fn_wrapper, name: name_fn_wrapper, schema_check: schema_check_fn_wrapper, clone: clone_fn_wrapper, @@ -220,6 +323,24 @@ impl PhysicalOptimizerRule for ForeignPhysicalOptimizerRule { (&optimized_plan).try_into() } + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, + ) -> Result> { + let ffi_context = FFI_PhysicalOptimizerContext::new(context); + let plan = FFI_ExecutionPlan::new(plan, None); + + let optimized_plan = unsafe { + df_result!((self.rule.optimize_with_context)( + &self.rule, + &plan, + &ffi_context + ))? + }; + (&optimized_plan).try_into() + } + fn name(&self) -> &str { &self.name } @@ -236,8 +357,11 @@ mod tests { use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; - use datafusion_physical_optimizer::PhysicalOptimizerRule; + use datafusion_physical_optimizer::{ + ConfigOnlyContext, PhysicalOptimizerContext, PhysicalOptimizerRule, + }; use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::operator_statistics::StatisticsRegistry; use super::*; use crate::execution_plan::tests::EmptyExec; @@ -265,6 +389,39 @@ mod tests { } } + /// A rule that returns an error from `optimize` but succeeds when + /// called via `optimize_with_context`, proving the context path is taken. + #[derive(Debug)] + struct ContextAwareRule; + + impl PhysicalOptimizerRule for ContextAwareRule { + fn optimize( + &self, + _plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + Err(datafusion_common::DataFusionError::Plan( + "optimize should not be called directly".to_string(), + )) + } + + fn optimize_with_context( + &self, + plan: Arc, + _context: &dyn PhysicalOptimizerContext, + ) -> Result> { + Ok(plan) + } + + fn name(&self) -> &str { + "context_aware_rule" + } + + fn schema_check(&self) -> bool { + true + } + } + fn create_test_plan() -> Arc { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); @@ -374,4 +531,70 @@ mod tests { Ok(()) } + + #[test] + fn test_optimize_with_context_round_trip() -> Result<()> { + let rule: Arc = + Arc::new(ContextAwareRule); + + let mut ffi_rule = FFI_PhysicalOptimizerRule::new(rule, None); + ffi_rule.library_marker_id = crate::mock_foreign_marker_id; + + let foreign_rule: Arc = + (&ffi_rule).into(); + + let plan = create_test_plan(); + let config = ConfigOptions::new(); + let context = ConfigOnlyContext::new(&config); + + let optimized = foreign_rule.optimize_with_context(plan, &context)?; + assert_eq!(optimized.name(), "empty-exec"); + + Ok(()) + } + + /// Tests that `optimize_with_context` works even when the caller supplies a + /// statistics registry. The registry cannot survive the FFI round-trip (it + /// contains trait object vtables that are library-local), so the provider + /// side will always see `None`. This test verifies the context-aware path + /// still succeeds in that scenario. + #[test] + fn test_optimize_with_context_with_registry() -> Result<()> { + let rule: Arc = + Arc::new(ContextAwareRule); + + let mut ffi_rule = FFI_PhysicalOptimizerRule::new(rule, None); + ffi_rule.library_marker_id = crate::mock_foreign_marker_id; + + let foreign_rule: Arc = + (&ffi_rule).into(); + + struct ContextWithRegistry { + config: ConfigOptions, + registry: StatisticsRegistry, + } + + impl PhysicalOptimizerContext for ContextWithRegistry { + fn config_options(&self) -> &ConfigOptions { + &self.config + } + + fn statistics_registry(&self) -> Option<&StatisticsRegistry> { + Some(&self.registry) + } + } + + let ctx = ContextWithRegistry { + config: ConfigOptions::new(), + registry: StatisticsRegistry::default_with_builtin_providers(), + }; + + let plan = create_test_plan(); + // The optimize_with_context path works, but the registry is not + // available on the provider side (it will be None). + let optimized = foreign_rule.optimize_with_context(plan, &ctx)?; + assert_eq!(optimized.name(), "empty-exec"); + + Ok(()) + } } diff --git a/datafusion/ffi/src/placement.rs b/datafusion/ffi/src/placement.rs new file mode 100644 index 0000000000000..837f0e3aad647 --- /dev/null +++ b/datafusion/ffi/src/placement.rs @@ -0,0 +1,72 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion_expr::ExpressionPlacement; + +#[expect(non_camel_case_types)] +#[repr(u8)] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum FFI_ExpressionPlacement { + Literal, + Column, + MoveTowardsLeafNodes, + KeepInPlace, +} + +impl From for FFI_ExpressionPlacement { + fn from(value: ExpressionPlacement) -> Self { + match value { + ExpressionPlacement::Literal => Self::Literal, + ExpressionPlacement::Column => Self::Column, + ExpressionPlacement::MoveTowardsLeafNodes => Self::MoveTowardsLeafNodes, + ExpressionPlacement::KeepInPlace => Self::KeepInPlace, + } + } +} + +impl From for ExpressionPlacement { + fn from(value: FFI_ExpressionPlacement) -> Self { + match value { + FFI_ExpressionPlacement::Literal => Self::Literal, + FFI_ExpressionPlacement::Column => Self::Column, + FFI_ExpressionPlacement::MoveTowardsLeafNodes => Self::MoveTowardsLeafNodes, + FFI_ExpressionPlacement::KeepInPlace => Self::KeepInPlace, + } + } +} + +#[cfg(test)] +mod tests { + use datafusion::logical_expr::ExpressionPlacement; + + use super::FFI_ExpressionPlacement; + + fn test_round_trip_placement(placement: ExpressionPlacement) { + let ffi_placement: FFI_ExpressionPlacement = placement.into(); + let round_trip: ExpressionPlacement = ffi_placement.into(); + + assert_eq!(placement, round_trip); + } + + #[test] + fn test_all_round_trip_placement() { + test_round_trip_placement(ExpressionPlacement::Literal); + test_round_trip_placement(ExpressionPlacement::Column); + test_round_trip_placement(ExpressionPlacement::MoveTowardsLeafNodes); + test_round_trip_placement(ExpressionPlacement::KeepInPlace); + } +} diff --git a/datafusion/ffi/src/plan_properties.rs b/datafusion/ffi/src/plan_properties.rs index b286ee2d7d30c..09ef26af32349 100644 --- a/datafusion/ffi/src/plan_properties.rs +++ b/datafusion/ffi/src/plan_properties.rs @@ -20,7 +20,7 @@ use std::sync::Arc; use arrow::datatypes::SchemaRef; use datafusion_common::error::{DataFusionError, Result}; -use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_physical_plan::PlanProperties; use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; @@ -172,6 +172,7 @@ impl TryFrom for PlanProperties { .unwrap_or_default(); let partitioning = unsafe { (ffi_props.output_partitioning)(&ffi_props) }; + let partitioning = Partitioning::try_from(partitioning)?; let eq_properties = if sort_exprs.is_empty() { EquivalenceProperties::new(Arc::new(schema)) @@ -187,7 +188,7 @@ impl TryFrom for PlanProperties { Ok(PlanProperties::new( eq_properties, - (&partitioning).into(), + partitioning, emission_type, boundedness, )) @@ -260,13 +261,15 @@ impl From for EmissionType { #[cfg(test)] mod tests { + use arrow::datatypes::{DataType, Field, Schema}; use datafusion::physical_expr::PhysicalSortExpr; use datafusion::physical_plan::Partitioning; + use datafusion_common::{ScalarValue, SplitPoint}; + use datafusion_physical_expr::{LexOrdering, RangePartitioning}; use super::*; fn create_test_props() -> Result { - use arrow::datatypes::{DataType, Field, Schema}; let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Float32, false)])); @@ -282,6 +285,25 @@ mod tests { )) } + fn create_range_test_props() -> Result { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let col = datafusion::physical_plan::expressions::col("a", &schema)?; + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(col)]) + .expect("non-empty ordering"); + let split_points = vec![ + SplitPoint::new(vec![ScalarValue::Int64(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int64(Some(20))]), + ]; + let range = RangePartitioning::try_new(ordering, split_points)?; + + Ok(PlanProperties::new( + EquivalenceProperties::new(schema), + Partitioning::Range(range), + EmissionType::Incremental, + Boundedness::Bounded, + )) + } + #[test] fn test_round_trip_ffi_plan_properties() -> Result<()> { let original_props = create_test_props()?; @@ -314,4 +336,22 @@ mod tests { Ok(()) } + + #[test] + fn test_round_trip_ffi_plan_properties_range_partitioning() -> Result<()> { + let original_props = create_range_test_props()?; + + let mut local_props_ptr = FFI_PlanProperties::from(&original_props); + local_props_ptr.library_marker_id = crate::mock_foreign_marker_id; + + let foreign_props: PlanProperties = local_props_ptr.try_into()?; + + assert_eq!( + format!("{:?}", foreign_props.output_partitioning()), + format!("{:?}", original_props.output_partitioning()) + ); + assert_eq!(format!("{foreign_props:?}"), format!("{original_props:?}")); + + Ok(()) + } } diff --git a/datafusion/ffi/src/proto/logical_extension_codec.rs b/datafusion/ffi/src/proto/logical_extension_codec.rs index 97aa5c901a636..ed2c594f1bc02 100644 --- a/datafusion/ffi/src/proto/logical_extension_codec.rs +++ b/datafusion/ffi/src/proto/logical_extension_codec.rs @@ -99,7 +99,7 @@ pub struct FFI_LogicalExtensionCodec { try_encode_udwf: unsafe extern "C" fn(&Self, node: FFI_WindowUDF) -> FFI_Result>, - pub task_ctx_provider: FFI_TaskContextProvider, + pub(crate) task_ctx_provider: FFI_TaskContextProvider, /// Used to create a clone on the provider of the execution plan. This should /// only need to be called by the receiver of the plan. @@ -295,7 +295,7 @@ impl Drop for FFI_LogicalExtensionCodec { impl FFI_LogicalExtensionCodec { /// Creates a new [`FFI_LogicalExtensionCodec`]. pub fn new( - codec: Arc, + codec: Arc, runtime: Option, task_ctx_provider: impl Into, ) -> Self { @@ -712,14 +712,12 @@ mod tests { #[test] fn ffi_logical_extension_codec_local_bypass() { - let codec = - Arc::new(TestExtensionCodec {}) as Arc; + let codec = Arc::new(TestExtensionCodec {}) as Arc; let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); let mut ffi_codec = FFI_LogicalExtensionCodec::new(Arc::clone(&codec), None, task_ctx_provider); - let codec = codec as Arc; // Verify local libraries can be downcast to their original let foreign_codec: Arc = (&ffi_codec).into(); assert!(arc_ptr_eq(&foreign_codec, &codec)); diff --git a/datafusion/ffi/src/proto/physical_extension_codec.rs b/datafusion/ffi/src/proto/physical_extension_codec.rs index 60d9d03dbd6dd..95d2ed68a6ea3 100644 --- a/datafusion/ffi/src/proto/physical_extension_codec.rs +++ b/datafusion/ffi/src/proto/physical_extension_codec.rs @@ -25,7 +25,10 @@ use datafusion_expr::{ AggregateUDF, AggregateUDFImpl, ScalarUDF, ScalarUDFImpl, WindowUDF, WindowUDFImpl, }; use datafusion_physical_plan::ExecutionPlan; -use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use datafusion_proto::physical_plan::{ + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalProtoConverterExtension, +}; use stabby::slice::Slice as SSlice; use stabby::str::Str as SStr; @@ -89,7 +92,7 @@ pub struct FFI_PhysicalExtensionCodec { unsafe extern "C" fn(&Self, node: FFI_WindowUDF) -> FFI_Result>, /// Access the current [`TaskContext`]. - task_ctx_provider: FFI_TaskContextProvider, + pub(crate) task_ctx_provider: FFI_TaskContextProvider, /// Used to create a clone on the provider of the execution plan. This should /// only need to be called by the receiver of the plan. @@ -145,8 +148,12 @@ unsafe extern "C" fn try_decode_fn_wrapper( .collect::>>(); let inputs = sresult_return!(inputs); - let plan = - sresult_return!(codec.try_decode(buf.as_ref(), &inputs, task_ctx.as_ref())); + let plan = sresult_return!(codec.try_decode( + buf.as_ref(), + &inputs, + task_ctx.as_ref(), + &DefaultPhysicalProtoConverter {}, + )); FFI_Result::Ok(FFI_ExecutionPlan::new(plan, runtime)) } @@ -160,7 +167,11 @@ unsafe extern "C" fn try_encode_fn_wrapper( let plan: Arc = sresult_return!((&node).try_into()); let mut bytes = Vec::new(); - sresult_return!(codec.try_encode(plan, &mut bytes)); + sresult_return!(codec.try_encode( + plan, + &mut bytes, + &DefaultPhysicalProtoConverter {} + )); FFI_Result::Ok(bytes.into_iter().collect()) } @@ -270,7 +281,7 @@ impl Drop for FFI_PhysicalExtensionCodec { impl FFI_PhysicalExtensionCodec { /// Creates a new [`FFI_PhysicalExtensionCodec`]. pub fn new( - codec: Arc, + codec: Arc, runtime: Option, task_ctx_provider: impl Into, ) -> Self { @@ -335,6 +346,7 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec { buf: &[u8], inputs: &[Arc], _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { let inputs = inputs .iter() @@ -348,7 +360,12 @@ impl PhysicalExtensionCodec for ForeignPhysicalExtensionCodec { Ok(plan) } - fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { let plan = FFI_ExecutionPlan::new(node, None); let bytes = df_result!(unsafe { (self.0.try_encode)(&self.0, plan) })?; @@ -426,7 +443,10 @@ pub(crate) mod tests { use datafusion_functions_aggregate::sum::Sum; use datafusion_functions_window::rank::{Rank, RankType}; use datafusion_physical_plan::ExecutionPlan; - use datafusion_proto::physical_plan::PhysicalExtensionCodec; + use datafusion_proto::physical_plan::{ + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalProtoConverterExtension, + }; use crate::execution_plan::tests::EmptyExec; use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; @@ -449,6 +469,7 @@ pub(crate) mod tests { buf: &[u8], _inputs: &[Arc], _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { if buf[0] != Self::MAGIC_NUMBER { return exec_err!( @@ -467,6 +488,7 @@ pub(crate) mod tests { &self, node: Arc, buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { buf.push(Self::MAGIC_NUMBER); @@ -587,10 +609,18 @@ pub(crate) mod tests { let exec = create_test_exec(); let input_execs = [create_test_exec()]; let mut bytes = Vec::new(); - foreign_codec.try_encode(Arc::clone(&exec), &mut bytes)?; - - let returned_exec = - foreign_codec.try_decode(&bytes, &input_execs, ctx.task_ctx().as_ref())?; + foreign_codec.try_encode( + Arc::clone(&exec), + &mut bytes, + &DefaultPhysicalProtoConverter {}, + )?; + + let returned_exec = foreign_codec.try_decode( + &bytes, + &input_execs, + ctx.task_ctx().as_ref(), + &DefaultPhysicalProtoConverter {}, + )?; assert!(returned_exec.is::()); @@ -665,14 +695,12 @@ pub(crate) mod tests { #[test] fn ffi_physical_extension_codec_local_bypass() { - let codec = - Arc::new(TestExtensionCodec {}) as Arc; + let codec = Arc::new(TestExtensionCodec {}) as Arc; let (_ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); let mut ffi_codec = FFI_PhysicalExtensionCodec::new(Arc::clone(&codec), None, task_ctx_provider); - let codec = codec as Arc; // Verify local libraries can be downcast to their original let foreign_codec: Arc = (&ffi_codec).into(); assert!(arc_ptr_eq(&foreign_codec, &codec)); diff --git a/datafusion/ffi/src/query_planner.rs b/datafusion/ffi/src/query_planner.rs new file mode 100644 index 0000000000000..6d895d65c5fc1 --- /dev/null +++ b/datafusion/ffi/src/query_planner.rs @@ -0,0 +1,445 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! FFI support for [`QueryPlanner`]. +//! +//! A typical deployment has three libraries. Library A (for example, +//! `datafusion-python`) owns the [`Session`] and codec registry. Library B owns +//! a custom table provider and its extension nodes. Library C (for example, +//! Ballista or `datafusion-distributed`) owns the query planner. A serializes a +//! logical plan and invokes C, while `FFI_SessionRef` lets C call session +//! services in A. C deserializes the logical plan, creates a physical plan, +//! serializes that result, and returns it for A to deserialize. The logical and +//! physical extension codecs preserve nodes supplied by B. +//! +//! The physical result is serialized instead of returned as an +//! [`crate::execution_plan::FFI_ExecutionPlan`]. An FFI execution-plan handle is +//! a foreign trait-object proxy, so even a built-in plan created in C cannot be +//! downcast to its concrete +//! type in A. Serialization reconstructs known plan nodes with A's local Rust +//! type identities, allowing A's optimizers and other consumers to downcast +//! them. Extension codecs control how custom nodes are reconstructed. +//! +//! A node returned by B while C is planning is still foreign to C unless a +//! codec boundary reconstructs it in C. The query-planner boundary guarantees +//! that C-local serializable nodes, and extension nodes understood by the +//! configured codecs, are reconstructed for A when the completed plan returns. +//! +//! # Delegating back to library A +//! +//! C commonly wants A's built-in planning as a starting point, then rewrites the +//! result. A must export its planner *before* installing C's planner on the +//! session, and C must retain that handle: after the swap, +//! [`Session::query_planner`] reports C's own planner, and +//! [`Session::create_physical_plan`] dispatches to it, so either one is a +//! self-call. Delegating to the retained handle is safe, because DataFusion's +//! built-in physical planner never re-dispatches through [`Session`]. +//! +//! Retain the planner rather than the session. [`FFI_QueryPlanner`] owns a +//! reference-counted planner, so it outlives A's original session, whereas +//! `FFI_SessionRef` borrows its session with the lifetime erased. + +use std::ffi::c_void; +use std::sync::Arc; + +use async_ffi::{FfiFuture, FutureExt}; +use async_trait::async_trait; +use datafusion_common::error::{DataFusionError, Result}; +use datafusion_expr::LogicalPlan; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_proto::bytes::{ + logical_plan_from_bytes_with_extension_codec, + logical_plan_to_bytes_with_extension_codec, + physical_plan_from_bytes_with_extension_codec, + physical_plan_to_bytes_with_extension_codec, +}; +use datafusion_proto::logical_plan::LogicalExtensionCodec; +use datafusion_proto::physical_plan::PhysicalExtensionCodec; +use datafusion_session::{QueryPlanner, Session}; +use stabby::vec::Vec as SVec; +use tokio::runtime::Handle; + +use crate::execution::FFI_TaskContextProvider; +use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::session::{FFI_SessionRef, ForeignSession}; +use crate::util::FFI_Result; +use crate::{df_result, sresult_return}; + +/// An ABI-stable handle to a [`QueryPlanner`] owned by another library. +/// +/// The Rust-facing adapters serialize the input [`LogicalPlan`] and resulting +/// [`ExecutionPlan`]; callers do not invoke the byte-oriented function pointer +/// directly. +#[repr(C)] +#[derive(Debug)] +pub struct FFI_QueryPlanner { + create_physical_plan: unsafe extern "C" fn( + &Self, + logical_plan_serialized: SVec, + session: FFI_SessionRef, + ) -> FfiFuture>>, + + /// Codec used to encode and decode logical plans and extension nodes. + logical_codec: FFI_LogicalExtensionCodec, + + /// Codec used to encode and decode physical plans and extension nodes. + physical_codec: FFI_PhysicalExtensionCodec, + + /// Used to create a clone of the query planner. + clone: unsafe extern "C" fn(planner: &Self) -> Self, + + /// Release the memory of the private data when it is no longer being used. + release: unsafe extern "C" fn(arg: &mut Self), + + /// Return the major DataFusion version number of this planner. + pub version: unsafe extern "C" fn() -> u64, + + /// Internal data. This is only to be accessed by the provider of the planner. + /// A [`ForeignQueryPlanner`] should never attempt to access this data. + private_data: *mut c_void, + + /// Utility to identify when FFI objects are accessed locally through + /// the foreign interface. See [`crate::get_library_marker_id`]. + pub library_marker_id: extern "C" fn() -> usize, +} + +unsafe impl Send for FFI_QueryPlanner {} +unsafe impl Sync for FFI_QueryPlanner {} + +struct QueryPlannerPrivateData { + planner: Arc, +} + +impl FFI_QueryPlanner { + fn inner(&self) -> &Arc { + let private_data = self.private_data as *const QueryPlannerPrivateData; + unsafe { &(*private_data).planner } + } +} + +unsafe extern "C" fn create_physical_plan_fn_wrapper( + planner: &FFI_QueryPlanner, + logical_plan_serialized: SVec, + session: FFI_SessionRef, +) -> FfiFuture>> { + let internal_planner = Arc::clone(planner.inner()); + let logical_codec: Arc = (&planner.logical_codec).into(); + let physical_codec: Arc = + (&planner.physical_codec).into(); + + async move { + let mut foreign_session = None; + let session = sresult_return!( + session + .as_local() + .map(Ok::<&dyn Session, DataFusionError>) + .unwrap_or_else(|| { + foreign_session = Some(ForeignSession::try_from(&session)?); + Ok(foreign_session.as_ref().unwrap()) + }) + ); + + let logical_plan = sresult_return!(logical_plan_from_bytes_with_extension_codec( + logical_plan_serialized.as_slice(), + session.task_ctx().as_ref(), + logical_codec.as_ref(), + )); + + let physical_plan = sresult_return!( + internal_planner + .create_physical_plan(&logical_plan, session) + .await + ); + let physical_plan = sresult_return!(physical_plan_to_bytes_with_extension_codec( + physical_plan, + physical_codec.as_ref(), + )); + + FFI_Result::Ok(SVec::from(physical_plan.as_ref())) + } + .into_ffi() +} + +unsafe extern "C" fn release_fn_wrapper(planner: &mut FFI_QueryPlanner) { + unsafe { + debug_assert!(!planner.private_data.is_null()); + let private_data = + Box::from_raw(planner.private_data as *mut QueryPlannerPrivateData); + drop(private_data); + planner.private_data = std::ptr::null_mut(); + } +} + +unsafe extern "C" fn clone_fn_wrapper(planner: &FFI_QueryPlanner) -> FFI_QueryPlanner { + let old_planner = Arc::clone(planner.inner()); + + let private_data = Box::into_raw(Box::new(QueryPlannerPrivateData { + planner: old_planner, + })) as *mut c_void; + + FFI_QueryPlanner { + create_physical_plan: create_physical_plan_fn_wrapper, + logical_codec: planner.logical_codec.clone(), + physical_codec: planner.physical_codec.clone(), + clone: clone_fn_wrapper, + release: release_fn_wrapper, + version: super::version, + private_data, + library_marker_id: crate::get_library_marker_id, + } +} + +impl Drop for FFI_QueryPlanner { + fn drop(&mut self) { + unsafe { (self.release)(self) } + } +} + +impl Clone for FFI_QueryPlanner { + fn clone(&self) -> Self { + unsafe { (self.clone)(self) } + } +} + +impl FFI_QueryPlanner { + /// Creates an [`FFI_QueryPlanner`] with native extension codecs. + /// + /// Both codecs are required so that the caller states which extension nodes + /// survive the boundary. Pass + /// [`DefaultLogicalExtensionCodec`](datafusion_proto::logical_plan::DefaultLogicalExtensionCodec) + /// and + /// [`DefaultPhysicalExtensionCodec`](datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec) + /// when no custom nodes are involved. `runtime` and `task_ctx_provider` + /// support codec callbacks across the FFI boundary. + pub fn new( + planner: Arc, + runtime: Option, + task_ctx_provider: impl Into, + logical_codec: Arc, + physical_codec: Arc, + ) -> Self { + let task_ctx_provider = task_ctx_provider.into(); + let logical_codec = FFI_LogicalExtensionCodec::new( + logical_codec, + runtime.clone(), + task_ctx_provider.clone(), + ); + let physical_codec = + FFI_PhysicalExtensionCodec::new(physical_codec, runtime, task_ctx_provider); + Self::new_with_ffi_codecs(planner, logical_codec, physical_codec) + } + + /// Creates an [`FFI_QueryPlanner`] using prebuilt FFI extension codecs. + /// + /// If `planner` is already foreign, this re-exports its original FFI handle + /// rather than adding another wrapper layer. The handle still adopts the + /// codecs supplied here, so they are never silently discarded. + pub fn new_with_ffi_codecs( + planner: Arc, + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + ) -> Self { + let any_ref: &dyn std::any::Any = planner.as_ref(); + if let Some(planner) = any_ref.downcast_ref::() { + let mut planner = planner.0.clone(); + planner.logical_codec = logical_codec; + planner.physical_codec = physical_codec; + return planner; + } + + let private_data = Box::new(QueryPlannerPrivateData { planner }); + + Self { + create_physical_plan: create_physical_plan_fn_wrapper, + logical_codec, + physical_codec, + clone: clone_fn_wrapper, + release: release_fn_wrapper, + version: super::version, + private_data: Box::into_raw(private_data) as *mut c_void, + library_marker_id: crate::get_library_marker_id, + } + } + + /// Creates a physical plan through this planner's FFI interface. + /// + /// This serializes `logical_plan`, exports `session` as an + /// `FFI_SessionRef`, invokes the planner's owning library, and + /// deserializes its physical-plan response. `session_runtime` is attached + /// to the exported session for callbacks that need its Tokio runtime. + /// + /// The [`QueryPlanner`] implementation for [`ForeignQueryPlanner`] cannot + /// obtain the session owner's runtime from the trait API, so it calls this + /// method with `None`. Embedders that own the runtime and need session + /// callbacks to enter it must call this method directly with `Some(handle)`. + pub async fn create_physical_plan_with_session_runtime( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + session_runtime: Option, + ) -> Result> { + let codec: Arc = (&self.logical_codec).into(); + let logical_plan = + logical_plan_to_bytes_with_extension_codec(logical_plan, codec.as_ref())?; + let logical_plan = SVec::from(logical_plan.as_ref()); + let task_ctx = session.task_ctx(); + let session = FFI_SessionRef::new_with_ffi_codecs( + session, + session_runtime, + self.logical_codec.clone(), + self.physical_codec.clone(), + ); + + let physical_plan = unsafe { + df_result!((self.create_physical_plan)(self, logical_plan, session).await)? + }; + let physical_codec: Arc = + (&self.physical_codec).into(); + + physical_plan_from_bytes_with_extension_codec( + physical_plan.as_slice(), + task_ctx.as_ref(), + physical_codec.as_ref(), + ) + } +} + +/// Consumer-side [`QueryPlanner`] adapter for an [`FFI_QueryPlanner`]. +/// +/// Calls serialize the logical plan, invoke the producing library, and +/// deserialize its physical-plan response. +#[derive(Debug)] +pub struct ForeignQueryPlanner(pub FFI_QueryPlanner); + +unsafe impl Send for ForeignQueryPlanner {} +unsafe impl Sync for ForeignQueryPlanner {} + +impl From<&FFI_QueryPlanner> for Arc { + fn from(planner: &FFI_QueryPlanner) -> Self { + if (planner.library_marker_id)() == crate::get_library_marker_id() { + Arc::clone(planner.inner()) + } else { + Arc::new(ForeignQueryPlanner(planner.clone())) + } + } +} + +#[async_trait] +impl QueryPlanner for ForeignQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + self.0 + .create_physical_plan_with_session_runtime(logical_plan, session, None) + .await + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion::prelude::SessionContext; + use datafusion_common::Result; + use datafusion_execution::TaskContextProvider; + use datafusion_expr::LogicalPlanBuilder; + use datafusion_physical_plan::empty::EmptyExec; + use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; + use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; + + use super::*; + + #[derive(Debug)] + struct EmptyQueryPlanner; + + #[async_trait] + impl QueryPlanner for EmptyQueryPlanner { + async fn create_physical_plan( + &self, + _logical_plan: &LogicalPlan, + _session: &dyn Session, + ) -> Result> { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + Ok(Arc::new(EmptyExec::new(schema))) + } + } + + fn create_ffi_query_planner(ctx: Arc) -> FFI_QueryPlanner { + let task_ctx_provider = Arc::clone(&ctx) as Arc; + FFI_QueryPlanner::new( + Arc::new(EmptyQueryPlanner), + None, + &task_ctx_provider, + Arc::new(DefaultLogicalExtensionCodec {}), + Arc::new(DefaultPhysicalExtensionCodec {}), + ) + } + + #[test] + fn test_ffi_query_planner_local_bypass() { + let ctx = Arc::new(SessionContext::new()); + let ffi_planner = create_ffi_query_planner(ctx); + let planner: Arc = (&ffi_planner).into(); + let any_ref: &dyn std::any::Any = planner.as_ref(); + assert!(any_ref.downcast_ref::().is_some()); + } + + #[tokio::test] + async fn test_round_trip_ffi_query_planner_create_physical_plan() -> Result<()> { + let ctx = Arc::new(SessionContext::new()); + let mut ffi_planner = create_ffi_query_planner(Arc::clone(&ctx)); + ffi_planner.library_marker_id = crate::mock_foreign_marker_id; + + let planner: Arc = (&ffi_planner).into(); + let any_ref: &dyn std::any::Any = planner.as_ref(); + assert!(any_ref.downcast_ref::().is_some()); + + let logical_plan = LogicalPlanBuilder::empty(false).build()?; + let state = ctx.state(); + let physical_plan = planner.create_physical_plan(&logical_plan, &state).await?; + assert_eq!(physical_plan.name(), "EmptyExec"); + assert!(physical_plan.is::()); + + Ok(()) + } + + #[tokio::test] + async fn test_create_physical_plan_with_session_runtime() -> Result<()> { + let ctx = Arc::new(SessionContext::new()); + let ffi_planner = create_ffi_query_planner(Arc::clone(&ctx)); + let logical_plan = LogicalPlanBuilder::empty(false).build()?; + let state = ctx.state(); + + let physical_plan = ffi_planner + .create_physical_plan_with_session_runtime( + &logical_plan, + &state, + Some(Handle::current()), + ) + .await?; + + assert_eq!(physical_plan.name(), "EmptyExec"); + assert!(physical_plan.is::()); + + Ok(()) + } +} diff --git a/datafusion/ffi/src/record_batch_stream.rs b/datafusion/ffi/src/record_batch_stream.rs index 74709848cbb7f..5a92cbfe5fe78 100644 --- a/datafusion/ffi/src/record_batch_stream.rs +++ b/datafusion/ffi/src/record_batch_stream.rs @@ -218,8 +218,8 @@ impl Drop for FFI_RecordBatchStream { mod tests { use std::sync::Arc; + use arrow::array::record_batch; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion::common::record_batch; use datafusion::error::Result; use datafusion::execution::SendableRecordBatchStream; use datafusion::test_util::bounded_stream; diff --git a/datafusion/ffi/src/session/mod.rs b/datafusion/ffi/src/session/mod.rs index dfc9d1c7dfebd..83f842508ab2c 100644 --- a/datafusion/ffi/src/session/mod.rs +++ b/datafusion/ffi/src/session/mod.rs @@ -15,10 +15,25 @@ // specific language governing permissions and limitations // under the License. +//! FFI support for [`Session`]. +//! +//! # Delegating physical planning +//! +//! Consider a session owned by library A that uses a query planner owned by +//! library C. After A installs C's planner, [`ForeignSession::query_planner`] +//! returns C's planner and [`ForeignSession::create_physical_plan`] dispatches +//! to C's planner. C must not call `create_physical_plan`, or invoke the planner +//! returned by `query_planner`, to delegate planning back to A. Repeating either +//! self-call recurses until the stack is exhausted. +//! +//! To delegate safely, A must export its original planner before installing C's +//! planner, and C must retain and invoke that planner directly. See the +//! [`crate::query_planner`] module for details. + use std::any::Any; use std::collections::HashMap; use std::ffi::c_void; -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use arrow_schema::SchemaRef; use arrow_schema::ffi::FFI_ArrowSchema; @@ -37,12 +52,18 @@ use datafusion_expr::{ }; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; -use datafusion_proto::bytes::{logical_plan_from_bytes, logical_plan_to_bytes}; +use datafusion_proto::bytes::{ + logical_plan_from_bytes, logical_plan_from_bytes_with_extension_codec, + logical_plan_to_bytes, logical_plan_to_bytes_with_extension_codec, +}; use datafusion_proto::logical_plan::LogicalExtensionCodec; use datafusion_proto::logical_plan::from_proto::parse_expr; use datafusion_proto::logical_plan::to_proto::serialize_expr; +use datafusion_proto::physical_plan::DefaultPhysicalExtensionCodec; use datafusion_proto::protobuf::LogicalExprNode; -use datafusion_session::Session; +use datafusion_session::{ + CatalogProviderList, PhysicalOptimizerRule, QueryPlanner, Session, +}; use prost::Message; use stabby::str::Str as SStr; @@ -51,10 +72,14 @@ use stabby::vec::Vec as SVec; use tokio::runtime::Handle; use crate::arrow_wrappers::WrappedSchema; +use crate::catalog_provider_list::FFI_CatalogProviderList; use crate::execution::FFI_TaskContext; use crate::execution_plan::FFI_ExecutionPlan; use crate::physical_expr::FFI_PhysicalExpr; +use crate::physical_optimizer::FFI_PhysicalOptimizerRule; use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::query_planner::FFI_QueryPlanner; use crate::session::config::FFI_SessionConfig; use crate::udaf::FFI_AggregateUDF; use crate::udf::FFI_ScalarUDF; @@ -83,6 +108,15 @@ pub(crate) struct FFI_SessionRef { config: unsafe extern "C" fn(&Self) -> FFI_SessionConfig, + catalog_list: unsafe extern "C" fn(&Self) -> FFI_CatalogProviderList, + + query_planner: unsafe extern "C" fn(&Self) -> FFI_QueryPlanner, + + optimize: unsafe extern "C" fn( + &Self, + logical_plan_serialized: SVec, + ) -> FFI_Result>, + create_physical_plan: unsafe extern "C" fn( &Self, @@ -107,8 +141,12 @@ pub(crate) struct FFI_SessionRef { task_ctx: unsafe extern "C" fn(&Self) -> FFI_TaskContext, + physical_optimizers: unsafe extern "C" fn(&Self) -> SVec, + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + /// Used to create a clone on the provider of the registry. This should /// only need to be called by the receiver of the plan. clone: unsafe extern "C" fn(plan: &Self) -> Self, @@ -132,12 +170,12 @@ unsafe impl Send for FFI_SessionRef {} unsafe impl Sync for FFI_SessionRef {} struct SessionPrivateData<'a> { - session: &'a (dyn Session + Send + Sync), + session: &'a dyn Session, runtime: Option, } impl FFI_SessionRef { - fn inner(&self) -> &(dyn Session + Send + Sync) { + fn inner(&self) -> &dyn Session { let private_data = self.private_data as *const SessionPrivateData; unsafe { (*private_data).session } } @@ -160,6 +198,46 @@ unsafe extern "C" fn config_fn_wrapper(session: &FFI_SessionRef) -> FFI_SessionC session.config().into() } +unsafe extern "C" fn catalog_list_fn_wrapper( + session: &FFI_SessionRef, +) -> FFI_CatalogProviderList { + FFI_CatalogProviderList::new_with_ffi_codec( + session.inner().catalog_list(), + unsafe { session.runtime() }.clone(), + session.logical_codec.clone(), + ) +} + +unsafe extern "C" fn query_planner_fn_wrapper( + session: &FFI_SessionRef, +) -> FFI_QueryPlanner { + FFI_QueryPlanner::new_with_ffi_codecs( + session.inner().query_planner(), + session.logical_codec.clone(), + session.physical_codec.clone(), + ) +} + +unsafe extern "C" fn optimize_fn_wrapper( + session: &FFI_SessionRef, + logical_plan_serialized: SVec, +) -> FFI_Result> { + let logical_codec: Arc = (&session.logical_codec).into(); + let inner = session.inner(); + let logical_plan = sresult_return!(logical_plan_from_bytes_with_extension_codec( + logical_plan_serialized.as_slice(), + inner.task_ctx().as_ref(), + logical_codec.as_ref(), + )); + let optimized_plan = sresult_return!(inner.optimize(&logical_plan)); + let optimized_plan = sresult_return!(logical_plan_to_bytes_with_extension_codec( + &optimized_plan, + logical_codec.as_ref(), + )); + + FFI_Result::Ok(SVec::from(optimized_plan.as_ref())) +} + unsafe extern "C" fn create_physical_plan_fn_wrapper( session: &FFI_SessionRef, logical_plan_serialized: SVec, @@ -258,6 +336,7 @@ fn table_options_to_rhash(mut options: TableOptions) -> SVec<(SString, SString)> "datafusion_ffi.table_current_format".into(), match current_format { ConfigFileType::JSON => "json", + #[cfg(feature = "parquet")] ConfigFileType::PARQUET => "parquet", ConfigFileType::CSV => "csv", } @@ -289,6 +368,18 @@ unsafe extern "C" fn task_ctx_fn_wrapper(session: &FFI_SessionRef) -> FFI_TaskCo session.inner().task_ctx().into() } +unsafe extern "C" fn physical_optimizers_fn_wrapper( + session: &FFI_SessionRef, +) -> SVec { + let runtime = unsafe { session.runtime().clone() }; + session + .inner() + .physical_optimizers() + .iter() + .map(|rule| FFI_PhysicalOptimizerRule::new(Arc::clone(rule), runtime.clone())) + .collect() +} + unsafe extern "C" fn release_fn_wrapper(provider: &mut FFI_SessionRef) { unsafe { let private_data = @@ -309,6 +400,9 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_SessionRef) -> FFI_SessionR FFI_SessionRef { session_id: session_id_fn_wrapper, config: config_fn_wrapper, + catalog_list: catalog_list_fn_wrapper, + query_planner: query_planner_fn_wrapper, + optimize: optimize_fn_wrapper, create_physical_plan: create_physical_plan_fn_wrapper, create_physical_expr: create_physical_expr_fn_wrapper, scalar_functions: scalar_functions_fn_wrapper, @@ -317,7 +411,9 @@ unsafe extern "C" fn clone_fn_wrapper(provider: &FFI_SessionRef) -> FFI_SessionR table_options: table_options_fn_wrapper, default_table_options: default_table_options_fn_wrapper, task_ctx: task_ctx_fn_wrapper, + physical_optimizers: physical_optimizers_fn_wrapper, logical_codec: provider.logical_codec.clone(), + physical_codec: provider.physical_codec.clone(), clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -335,14 +431,60 @@ impl Drop for FFI_SessionRef { } impl FFI_SessionRef { - /// Creates a new [`FFI_SessionRef`]. + /// Creates a new [`FFI_SessionRef`] with a default physical extension codec. + /// + /// The synthesized [`DefaultPhysicalExtensionCodec`] supports built-in physical + /// nodes only. A query planner obtained through this session reference therefore + /// cannot encode or decode custom physical extension nodes. Use + /// [`Self::new_with_ffi_codecs`] with matching logical and physical codecs when + /// custom physical nodes must cross the FFI boundary. + /// + /// The physical codec wrapper requires a + /// [`FFI_TaskContextProvider`](crate::execution::FFI_TaskContextProvider), but this + /// constructor has only a session reference and a logical codec. It therefore + /// reuses the logical codec's provider. The provider may be owned by another + /// library; this is safe, but it must remain live and return the task context + /// intended for codec callbacks. The default physical codec does not successfully + /// decode extension nodes, so callers that need such callbacks must instead use + /// [`Self::new_with_ffi_codecs`] with an explicitly configured physical codec and + /// task context provider. pub fn new( - session: &(dyn Session + Send + Sync), + session: &dyn Session, + runtime: Option, + logical_codec: FFI_LogicalExtensionCodec, + ) -> Self { + // `Session` provides a TaskContext but not the reference-counted + // TaskContextProvider needed by the FFI codec. Reuse the provider associated + // with the logical codec under the assumptions documented above. + let physical_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(DefaultPhysicalExtensionCodec {}), + runtime.clone(), + logical_codec.task_ctx_provider.clone(), + ); + Self::new_with_ffi_codecs(session, runtime, logical_codec, physical_codec) + } + + /// Creates a new [`FFI_SessionRef`] using existing FFI codecs. + /// + /// The codecs must form a matching pair that can round-trip every logical and + /// physical extension node exposed through the session. Their task context + /// providers must remain live and return contexts appropriate for their decode + /// callbacks. + /// + /// If `session` is already foreign, this re-exports its original FFI handle + /// rather than adding another wrapper layer. The handle adopts the codecs + /// supplied here while retaining its original private data and runtime. + pub fn new_with_ffi_codecs( + session: &dyn Session, runtime: Option, logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, ) -> Self { if let Some(session) = session.as_any().downcast_ref::() { - return session.session.clone(); + let mut session = session.session.clone(); + session.logical_codec = logical_codec; + session.physical_codec = physical_codec; + return session; } let private_data = Box::new(SessionPrivateData { session, runtime }); @@ -350,6 +492,9 @@ impl FFI_SessionRef { Self { session_id: session_id_fn_wrapper, config: config_fn_wrapper, + catalog_list: catalog_list_fn_wrapper, + query_planner: query_planner_fn_wrapper, + optimize: optimize_fn_wrapper, create_physical_plan: create_physical_plan_fn_wrapper, create_physical_expr: create_physical_expr_fn_wrapper, scalar_functions: scalar_functions_fn_wrapper, @@ -358,7 +503,9 @@ impl FFI_SessionRef { table_options: table_options_fn_wrapper, default_table_options: default_table_options_fn_wrapper, task_ctx: task_ctx_fn_wrapper, + physical_optimizers: physical_optimizers_fn_wrapper, logical_codec, + physical_codec, clone: clone_fn_wrapper, release: release_fn_wrapper, @@ -371,27 +518,39 @@ impl FFI_SessionRef { /// This wrapper struct exists on the receiver side of the FFI interface, so it has /// no guarantees about being able to access the data in `private_data`. Any functions -/// defined on this struct must only use the stable functions provided in -/// FFI_Session to interact with the foreign table provider. +/// defined on this struct must use only the stable function pointers in +/// `FFI_SessionRef` to interact with the foreign session. +/// +/// # Query planner delegation +/// +/// If the session owner installed the current foreign query planner, +/// [`Session::create_physical_plan`] dispatches back to that planner and +/// [`Session::query_planner`] returns that planner. The planner must retain and +/// invoke the session owner's previous planner instead of using either method to +/// delegate back to the session. Otherwise, repeated delegation exhausts the +/// stack. See [`crate::query_planner`] for details. #[derive(Debug)] pub struct ForeignSession { session: FFI_SessionRef, config: SessionConfig, + catalog_list: Arc, scalar_functions: HashMap>, - higher_order_functions: HashMap>, + higher_order_functions: HashMap>, aggregate_functions: HashMap>, window_functions: HashMap>, extension_types: ExtensionTypeRegistryRef, table_options: TableOptions, runtime_env: Arc, props: ExecutionProps, + query_planner: OnceLock>, + physical_optimizers: OnceLock>>, } unsafe impl Send for ForeignSession {} unsafe impl Sync for ForeignSession {} impl FFI_SessionRef { - pub fn as_local(&self) -> Option<&(dyn Session + Send + Sync)> { + pub fn as_local(&self) -> Option<&dyn Session> { if (self.library_marker_id)() == crate::get_library_marker_id() { return Some(self.inner()); } @@ -409,6 +568,9 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { let config = (session.config)(session); let config = SessionConfig::try_from(&config)?; + let ffi_catalog_list = (session.catalog_list)(session); + let catalog_list = (&ffi_catalog_list).into(); + let scalar_functions = (session.scalar_functions)(session) .into_iter() .map(|kv_pair| { @@ -442,10 +604,10 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { ) }) .collect(); - Ok(Self { session: session.clone(), config, + catalog_list, table_options, scalar_functions, higher_order_functions: HashMap::new(), @@ -454,6 +616,8 @@ impl TryFrom<&FFI_SessionRef> for ForeignSession { extension_types: Arc::new(MemoryExtensionTypeRegistry::default()), runtime_env: Default::default(), props: Default::default(), + query_planner: OnceLock::new(), + physical_optimizers: OnceLock::new(), }) } } @@ -476,6 +640,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption let formats = [ ConfigFileType::CSV, ConfigFileType::JSON, + #[cfg(feature = "parquet")] ConfigFileType::PARQUET, ]; for format in formats { @@ -483,6 +648,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption // included in the formats list above and in the extension check below. let format_name = match &format { ConfigFileType::CSV => "csv", + #[cfg(feature = "parquet")] ConfigFileType::PARQUET => "parquet", ConfigFileType::JSON => "json", }; @@ -504,7 +670,6 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption .unwrap_or_else(|err| log::warn!("Error parsing table options: {err}")); } } - let extension_options: HashMap = options .iter() .filter_map(|(k, v)| { @@ -525,6 +690,7 @@ fn table_options_from_rhashmap(options: SVec<(SString, SString)>) -> TableOption table_options.current_format = current_format.and_then(|format| match format.as_str() { "csv" => Some(ConfigFileType::CSV), + #[cfg(feature = "parquet")] "parquet" => Some(ConfigFileType::PARQUET), "json" => Some(ConfigFileType::JSON), _ => None, @@ -546,6 +712,35 @@ impl Session for ForeignSession { self.config.options() } + fn catalog_list(&self) -> Arc { + Arc::clone(&self.catalog_list) + } + + fn query_planner(&self) -> Arc { + Arc::clone(self.query_planner.get_or_init(|| unsafe { + let planner = (self.session.query_planner)(&self.session); + (&planner).into() + })) + } + + fn optimize(&self, plan: &LogicalPlan) -> datafusion_common::Result { + unsafe { + let codec: Arc = + (&self.session.logical_codec).into(); + let logical_plan = + logical_plan_to_bytes_with_extension_codec(plan, codec.as_ref())?; + let optimized_plan = df_result!((self.session.optimize)( + &self.session, + SVec::from(logical_plan.as_ref()), + ))?; + logical_plan_from_bytes_with_extension_codec( + optimized_plan.as_slice(), + self.task_ctx().as_ref(), + codec.as_ref(), + ) + } + } + async fn create_physical_plan( &self, logical_plan: &LogicalPlan, @@ -586,11 +781,20 @@ impl Session for ForeignSession { } } + fn physical_optimizers(&self) -> &[Arc] { + self.physical_optimizers.get_or_init(|| unsafe { + (self.session.physical_optimizers)(&self.session) + .into_iter() + .map(|rule| (&rule).into()) + .collect() + }) + } + fn scalar_functions(&self) -> &HashMap> { &self.scalar_functions } - fn higher_order_functions(&self) -> &HashMap> { + fn higher_order_functions(&self) -> &HashMap> { &self.higher_order_functions } @@ -645,8 +849,10 @@ impl Session for ForeignSession { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::{AtomicUsize, Ordering}; use arrow_schema::{DataType, Field, Schema}; + use datafusion::catalog::MemoryCatalogProvider; use datafusion::execution::SessionStateBuilder; use datafusion_common::DataFusionError; use datafusion_expr::col; @@ -655,13 +861,69 @@ mod tests { use super::*; + static QUERY_PLANNER_CALLS: AtomicUsize = AtomicUsize::new(0); + static PHYSICAL_OPTIMIZER_CALLS: AtomicUsize = AtomicUsize::new(0); + + unsafe extern "C" fn counting_query_planner( + session: &FFI_SessionRef, + ) -> FFI_QueryPlanner { + QUERY_PLANNER_CALLS.fetch_add(1, Ordering::Relaxed); + unsafe { query_planner_fn_wrapper(session) } + } + + unsafe extern "C" fn counting_physical_optimizers( + session: &FFI_SessionRef, + ) -> SVec { + PHYSICAL_OPTIMIZER_CALLS.fetch_add(1, Ordering::Relaxed); + unsafe { physical_optimizers_fn_wrapper(session) } + } + + #[test] + fn test_foreign_session_lazily_loads_planning_state() -> Result<(), DataFusionError> { + QUERY_PLANNER_CALLS.store(0, Ordering::Relaxed); + PHYSICAL_OPTIMIZER_CALLS.store(0, Ordering::Relaxed); + + let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); + let logical_codec = FFI_LogicalExtensionCodec::new( + Arc::new(DefaultLogicalExtensionCodec {}), + None, + task_ctx_provider, + ); + let state = ctx.state(); + let mut local_session = FFI_SessionRef::new(&state, None, logical_codec); + local_session.query_planner = counting_query_planner; + local_session.physical_optimizers = counting_physical_optimizers; + + let mut foreign_session = ForeignSession::try_from(&local_session)?; + assert_eq!(QUERY_PLANNER_CALLS.load(Ordering::Relaxed), 0); + assert_eq!(PHYSICAL_OPTIMIZER_CALLS.load(Ordering::Relaxed), 0); + + // `FFI_SessionRef::clone` restores the standard function pointers, so + // instrument the clone retained by `ForeignSession` as well. + foreign_session.session.query_planner = counting_query_planner; + foreign_session.session.physical_optimizers = counting_physical_optimizers; + + foreign_session.query_planner(); + foreign_session.query_planner(); + assert_eq!(QUERY_PLANNER_CALLS.load(Ordering::Relaxed), 1); + + foreign_session.physical_optimizers(); + foreign_session.physical_optimizers(); + assert_eq!(PHYSICAL_OPTIMIZER_CALLS.load(Ordering::Relaxed), 1); + + Ok(()) + } + #[tokio::test] async fn test_ffi_session() -> Result<(), DataFusionError> { let (ctx, task_ctx_provider) = crate::util::tests::test_session_and_ctx(); let mut table_options = TableOptions::default(); table_options.csv.has_header = Some(true); table_options.json.schema_infer_max_rec = Some(10); - table_options.parquet.global.coerce_int96 = Some("123456789".into()); + #[cfg(feature = "parquet")] + { + table_options.parquet.global.coerce_int96 = Some("123456789".into()); + } table_options.current_format = Some(ConfigFileType::JSON); let state = SessionStateBuilder::new_from_existing(ctx.state()) @@ -687,7 +949,30 @@ mod tests { assert_eq!(foreign_session.session_id(), state.session_id()); + let foreign_catalog_list = foreign_session.catalog_list(); + assert_eq!( + foreign_catalog_list.catalog_names(), + state.catalog_list().catalog_names() + ); + foreign_catalog_list.register_catalog( + "foreign_registered".to_owned(), + Arc::new(MemoryCatalogProvider::new()), + ); + assert!(state.catalog_list().catalog("foreign_registered").is_some()); + let logical_plan = LogicalPlan::default(); + assert_eq!(foreign_session.optimize(&logical_plan)?, logical_plan); + assert_eq!( + foreign_session.physical_optimizers().len(), + state.physical_optimizers().len() + ); + assert!(foreign_session.statistics_registry().is_none()); + let planned = foreign_session + .query_planner() + .create_physical_plan(&logical_plan, &foreign_session) + .await?; + assert_eq!(planned.name(), "EmptyExec"); + let physical_plan = foreign_session.create_physical_plan(&logical_plan).await?; assert_eq!( format!("{physical_plan:?}"), diff --git a/datafusion/ffi/src/table_provider.rs b/datafusion/ffi/src/table_provider.rs index 5a4b2fa27256f..ee9377bff064e 100644 --- a/datafusion/ffi/src/table_provider.rs +++ b/datafusion/ffi/src/table_provider.rs @@ -263,7 +263,7 @@ unsafe extern "C" fn scan_fn_wrapper( let session = sresult_return!( session .as_local() - .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>) + .map(Ok::<&dyn Session, DataFusionError>) .unwrap_or_else(|| { foreign_session = Some(ForeignSession::try_from(&session)?); Ok(foreign_session.as_ref().unwrap()) @@ -314,7 +314,7 @@ unsafe extern "C" fn insert_into_fn_wrapper( let session = sresult_return!( session .as_local() - .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>) + .map(Ok::<&dyn Session, DataFusionError>) .unwrap_or_else(|| { foreign_session = Some(ForeignSession::try_from(&session)?); Ok(foreign_session.as_ref().unwrap()) diff --git a/datafusion/ffi/src/table_provider_factory.rs b/datafusion/ffi/src/table_provider_factory.rs index 3ce8841614bc0..63ebb51bb1db8 100644 --- a/datafusion/ffi/src/table_provider_factory.rs +++ b/datafusion/ffi/src/table_provider_factory.rs @@ -152,7 +152,7 @@ impl FFI_TableProviderFactory { let plan = LogicalPlanNode::decode(cmd_serialized.as_ref()) .map_err(|e| DataFusionError::Internal(format!("{e:?}")))?; match plan.try_into_logical_plan(&task_ctx, logical_codec.as_ref())? { - LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) => Ok(cmd), + LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) => Ok(*cmd), _ => Err(DataFusionError::Internal( "Invalid logical plan in FFI_TableProviderFactory.".to_owned(), )), @@ -211,7 +211,7 @@ async fn create_fn_wrapper_impl( let mut foreign_session = None; let session = session .as_local() - .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>) + .map(Ok::<&dyn Session, DataFusionError>) .unwrap_or_else(|| { foreign_session = Some(ForeignSession::try_from(&session)?); Ok(foreign_session.as_ref().unwrap()) @@ -272,7 +272,7 @@ impl ForeignTableProviderFactory { let logical_codec: Arc = (&self.0.logical_codec).into(); - let plan = LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)); + let plan = LogicalPlan::Ddl(DdlStatement::CreateExternalTable(Box::new(cmd))); let plan: LogicalPlanNode = AsLogicalPlan::try_from_logical_plan(&plan, logical_codec.as_ref())?; @@ -368,7 +368,7 @@ mod tests { let cmd = CreateExternalTable { schema: Schema::empty().to_dfschema_ref()?, name: TableReference::bare("test_table"), - location: "test".to_string(), + locations: vec!["test".to_string()], file_type: "test".to_string(), table_partition_cols: vec![], if_not_exists: false, @@ -406,7 +406,7 @@ mod tests { let cmd = CreateExternalTable { schema: Schema::empty().to_dfschema_ref()?, name: TableReference::bare("cloned_test"), - location: "test".to_string(), + locations: vec!["test".to_string()], file_type: "test".to_string(), table_partition_cols: vec![], if_not_exists: false, diff --git a/datafusion/ffi/src/tests/async_provider.rs b/datafusion/ffi/src/tests/async_provider.rs index 011d3f0a0a343..83057d8c45db3 100644 --- a/datafusion/ffi/src/tests/async_provider.rs +++ b/datafusion/ffi/src/tests/async_provider.rs @@ -31,13 +31,15 @@ use std::sync::Arc; use arrow::array::RecordBatch; use arrow::datatypes::Schema; use async_trait::async_trait; -use datafusion_catalog::TableProvider; +use datafusion_catalog::{MemoryCatalogProvider, TableProvider}; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, exec_err}; use datafusion_execution::RecordBatchStream; use datafusion_expr::Expr; use datafusion_physical_expr::{EquivalenceProperties, Partitioning}; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, +}; use datafusion_session::Session; use futures::Stream; use tokio::runtime::Handle; @@ -136,11 +138,28 @@ impl TableProvider for AsyncTableProvider { async fn scan( &self, - _state: &dyn Session, + state: &dyn Session, _projection: Option<&Vec>, _filters: &[Expr], _limit: Option, ) -> Result> { + let catalog = state.catalog_list().catalog("datafusion").ok_or_else(|| { + datafusion_common::exec_datafusion_err!("missing datafusion catalog") + })?; + let schema = catalog.schema("public").ok_or_else(|| { + datafusion_common::exec_datafusion_err!("missing public schema") + })?; + if schema.table("external_table").await?.is_none() { + return exec_err!("missing external_table"); + } + + // Register a catalog from the dynamically loaded library so the host + // can verify that catalog mutations cross the FFI boundary as well. + state.catalog_list().register_catalog( + "ffi_registered".to_owned(), + Arc::new(MemoryCatalogProvider::new()), + ); + Ok(Arc::new(AsyncTestExecutionPlan::new( self.batch_request.clone(), self.batch_receiver.resubscribe(), @@ -194,13 +213,24 @@ impl ExecutionPlan for AsyncTestExecutionPlan { Vec::default() } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -214,18 +244,11 @@ impl ExecutionPlan for AsyncTestExecutionPlan { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion_physical_plan::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in the output ordering from equivalence properties - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.properties.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } diff --git a/datafusion/ffi/src/tests/catalog.rs b/datafusion/ffi/src/tests/catalog.rs index 0c02de5d049ae..b0b0858a8a3d7 100644 --- a/datafusion/ffi/src/tests/catalog.rs +++ b/datafusion/ffi/src/tests/catalog.rs @@ -48,8 +48,8 @@ pub struct FixedSchemaProvider { } pub fn fruit_table() -> Arc { + use arrow::array::record_batch; use arrow::datatypes::{DataType, Field}; - use datafusion_common::record_batch; let schema = Arc::new(Schema::new(vec![ Field::new("units", DataType::Int32, true), diff --git a/datafusion/ffi/src/tests/mod.rs b/datafusion/ffi/src/tests/mod.rs index 62e62d82359b5..74310d9c28e2f 100644 --- a/datafusion/ffi/src/tests/mod.rs +++ b/datafusion/ffi/src/tests/mod.rs @@ -17,32 +17,35 @@ use std::sync::Arc; -use arrow::array::RecordBatch; +use arrow::array::{RecordBatch, record_batch}; use arrow_schema::{DataType, Field, Schema}; use async_provider::create_async_table_provider; use async_trait::async_trait; use catalog::create_catalog_provider; use datafusion_catalog::MemTable; use datafusion_catalog::{Session, TableProvider}; -use datafusion_common::record_batch; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, Statistics}; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{Expr, TableType}; +use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_plan::ExecutionPlan; use sync_provider::create_sync_table_provider; use udf_udaf_udwf::{ - create_ffi_abs_func, create_ffi_random_func, create_ffi_rank_func, - create_ffi_stddev_func, create_ffi_sum_func, create_ffi_table_func, + create_ffi_abs_func, create_ffi_first_value_func, create_ffi_random_func, + create_ffi_rank_func, create_ffi_stddev_func, create_ffi_sum_func, + create_ffi_table_func, }; use crate::catalog_provider::FFI_CatalogProvider; use crate::catalog_provider_list::FFI_CatalogProviderList; use crate::config::extension_options::FFI_ExtensionOptions; use crate::execution_plan::FFI_ExecutionPlan; -use crate::execution_plan::tests::EmptyExec; +use crate::execution_plan::tests::{EmptyExec, create_dynamic_filter}; use crate::physical_optimizer::FFI_PhysicalOptimizerRule; use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::query_planner::FFI_QueryPlanner; use crate::table_provider::FFI_TableProvider; use crate::table_provider_factory::FFI_TableProviderFactory; use crate::tests::catalog::create_catalog_provider_list; @@ -50,11 +53,13 @@ use crate::udaf::FFI_AggregateUDF; use crate::udf::FFI_ScalarUDF; use crate::udtf::FFI_TableFunction; use crate::udwf::FFI_WindowUDF; +use crate::util::FFI_Option; mod async_provider; pub mod catalog; pub mod config; mod physical_optimizer; +mod query_planner; mod sync_provider; mod table_provider_factory; mod udf_udaf_udwf; @@ -90,6 +95,8 @@ pub struct ForeignLibraryModule { pub create_timezone_udf: extern "C" fn() -> FFI_ScalarUDF, + pub create_placement_udf: extern "C" fn() -> FFI_ScalarUDF, + pub create_table_function: extern "C" fn(FFI_LogicalExtensionCodec) -> FFI_TableFunction, @@ -106,6 +113,10 @@ pub struct ForeignLibraryModule { pub create_empty_exec: extern "C" fn() -> FFI_ExecutionPlan, + pub create_exec_with_expressions: extern "C" fn() -> FFI_ExecutionPlan, + + pub create_exec_with_dynamic_expressions: extern "C" fn() -> FFI_ExecutionPlan, + pub create_exec_with_statistics: extern "C" fn() -> FFI_ExecutionPlan, pub create_table_with_statistics: @@ -113,7 +124,20 @@ pub struct ForeignLibraryModule { pub create_physical_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, + pub create_context_aware_optimizer_rule: extern "C" fn() -> FFI_PhysicalOptimizerRule, + + /// Construct a query planner. When `library_a_planner` is provided the + /// planner delegates to it, as library C does after library A swaps planners. + pub create_query_planner: extern "C" fn( + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + library_a_planner: FFI_Option, + ) -> FFI_QueryPlanner, + pub version: extern "C" fn() -> u64, + + /// Create an aggregate UDAF using first_value + pub create_first_value_udaf: extern "C" fn() -> FFI_AggregateUDF, } pub fn create_test_schema() -> Arc { @@ -158,6 +182,21 @@ pub(crate) extern "C" fn create_empty_exec() -> FFI_ExecutionPlan { FFI_ExecutionPlan::new(plan, None) } +pub(crate) extern "C" fn create_exec_with_expressions() -> FFI_ExecutionPlan { + let schema = Arc::new(Schema::empty()); + let expression: Arc = create_dynamic_filter(); + let plan = Arc::new(EmptyExec::new(schema).with_expressions(vec![expression])); + FFI_ExecutionPlan::new(plan, None) +} + +pub(crate) extern "C" fn create_exec_with_dynamic_expressions() -> FFI_ExecutionPlan { + let schema = Arc::new(Schema::empty()); + let expression: Arc = create_dynamic_filter(); + let plan = + Arc::new(EmptyExec::new(schema).with_dynamic_expressions(vec![expression])); + FFI_ExecutionPlan::new(plan, None) +} + /// Returns canonical statistics used by both the producer and consumer sides of /// the integration tests so round-trips can be asserted without hard-coding /// the values in two places. @@ -249,16 +288,23 @@ pub extern "C" fn datafusion_ffi_get_module() -> ForeignLibraryModule { create_scalar_udf: create_ffi_abs_func, create_nullary_udf: create_ffi_random_func, create_timezone_udf: udf_udaf_udwf::create_timezone_func, + create_placement_udf: udf_udaf_udwf::create_placement_func, create_table_function: create_ffi_table_func, create_sum_udaf: create_ffi_sum_func, create_stddev_udaf: create_ffi_stddev_func, create_rank_udwf: create_ffi_rank_func, create_extension_options: config::create_extension_options, create_empty_exec, + create_exec_with_expressions, + create_exec_with_dynamic_expressions, create_exec_with_statistics, create_table_with_statistics, create_physical_optimizer_rule: physical_optimizer::create_physical_optimizer_rule, + create_context_aware_optimizer_rule: + physical_optimizer::create_context_aware_optimizer_rule, + create_query_planner: query_planner::create_query_planner, version: super::version, + create_first_value_udaf: create_ffi_first_value_func, } } diff --git a/datafusion/ffi/src/tests/physical_optimizer.rs b/datafusion/ffi/src/tests/physical_optimizer.rs index 2476526125b06..581f454e5259e 100644 --- a/datafusion/ffi/src/tests/physical_optimizer.rs +++ b/datafusion/ffi/src/tests/physical_optimizer.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; -use datafusion_physical_optimizer::PhysicalOptimizerRule; +use datafusion_physical_optimizer::{PhysicalOptimizerContext, PhysicalOptimizerRule}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::limit::GlobalLimitExec; @@ -52,3 +52,45 @@ pub(crate) extern "C" fn create_physical_optimizer_rule() -> FFI_PhysicalOptimiz let rule: Arc = Arc::new(AddLimitRule); FFI_PhysicalOptimizerRule::new(rule, None) } + +/// A rule that returns an error from `optimize()` (proving the context path must +/// be taken) but succeeds in `optimize_with_context()` by wrapping the plan in a +/// `GlobalLimitExec`. +#[derive(Debug)] +struct ContextAwareAddLimitRule; + +impl PhysicalOptimizerRule for ContextAwareAddLimitRule { + fn optimize( + &self, + _plan: Arc, + _config: &ConfigOptions, + ) -> Result> { + Err(datafusion_common::DataFusionError::Plan( + "optimize should not be called directly; use optimize_with_context" + .to_string(), + )) + } + + fn optimize_with_context( + &self, + plan: Arc, + _context: &dyn PhysicalOptimizerContext, + ) -> Result> { + Ok(Arc::new(GlobalLimitExec::new(plan, 0, Some(10)))) + } + + fn name(&self) -> &str { + "context_aware_add_limit_rule" + } + + fn schema_check(&self) -> bool { + true + } +} + +pub(crate) extern "C" fn create_context_aware_optimizer_rule() -> FFI_PhysicalOptimizerRule +{ + let rule: Arc = + Arc::new(ContextAwareAddLimitRule); + FFI_PhysicalOptimizerRule::new(rule, None) +} diff --git a/datafusion/ffi/src/tests/query_planner.rs b/datafusion/ffi/src/tests/query_planner.rs new file mode 100644 index 0000000000000..90d713f8a6336 --- /dev/null +++ b/datafusion/ffi/src/tests/query_planner.rs @@ -0,0 +1,172 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::any::Any; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema}; +use async_trait::async_trait; +use datafusion_catalog::default_table_source::source_as_provider; +use datafusion_common::{Result, exec_err}; +use datafusion_expr::LogicalPlan; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::union::UnionExec; +use datafusion_session::{QueryPlanner, Session}; + +use crate::execution_plan::ForeignExecutionPlan; +use crate::proto::logical_extension_codec::FFI_LogicalExtensionCodec; +use crate::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; +use crate::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; +use crate::session::ForeignSession; +use crate::table_provider::ForeignTableProvider; +use crate::util::FFI_Option; + +#[derive(Debug)] +struct TestQueryPlanner; + +#[async_trait] +impl QueryPlanner for TestQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + if let LogicalPlan::TableScan(scan) = logical_plan { + if session.as_any().downcast_ref::().is_none() { + return exec_err!("library A's session was not foreign to library C"); + } + + let provider = source_as_provider(&scan.source)?; + if provider.downcast_ref::().is_none() { + return exec_err!("library B's provider was not foreign to library C"); + } + let library_b_plan = provider + .scan(session, scan.projection.as_ref(), &scan.filters, scan.fetch) + .await?; + + if !library_b_plan.is::() { + return exec_err!("library B's plan unexpectedly downcast as C-local"); + } + + let plan = UnionExec::try_new(vec![ + Arc::clone(&library_b_plan), + Arc::clone(&library_b_plan), + ])?; + if !plan.is::() { + return exec_err!("library C could not downcast its local UnionExec"); + } + return Ok(plan); + } + + let query_planner = session.query_planner(); + let planner_any: &dyn Any = query_planner.as_ref(); + if planner_any.downcast_ref::().is_none() { + return exec_err!("query planner did not cross the FFI boundary"); + } + session.optimize(logical_plan)?; + if session.physical_optimizers().is_empty() { + return exec_err!("physical optimizers did not cross the FFI boundary"); + } + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + Ok(Arc::new(EmptyExec::new(schema))) + } +} + +/// Library C's planner for the planner-swap deployment. +/// +/// It holds the query planner library A exported *before* A swapped this planner +/// into its session, so delegating to it cannot re-enter library C. +#[derive(Debug)] +struct SwappedQueryPlanner { + library_a_planner: Arc, +} + +#[async_trait] +impl QueryPlanner for SwappedQueryPlanner { + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result> { + if session.as_any().downcast_ref::().is_none() { + return exec_err!("library A's session was not foreign to library C"); + } + + // After the swap, the planner installed on library A's session is this + // planner, so `session.query_planner()` and `session.create_physical_plan()` + // are both self-references. Assert the hazard instead of triggering it: + // calling either would recurse until the stack is exhausted. + let installed = session.query_planner(); + let installed: &dyn Any = installed.as_ref(); + if installed.downcast_ref::().is_none() { + return exec_err!( + "expected the swapped session to report library C's own planner" + ); + } + + // Delegate to library A. The result crosses the FFI boundary as + // serialized bytes, so library C receives nodes carrying its own local + // Rust type identities. + let plan = self + .library_a_planner + .create_physical_plan(logical_plan, session) + .await?; + + if plan.is::() { + return exec_err!("library A's plan was opaque to library C"); + } + let Some(sort) = plan.downcast_ref::() else { + return exec_err!( + "library C could not downcast library A's SortExec; got {}", + plan.name() + ); + }; + // Library B's scan is still foreign to library C. Only a codec boundary + // reconstructs it, and library A's codec hands back an A-local node. + if !sort.input().is::() { + return exec_err!("library B's scan unexpectedly downcast as C-local"); + } + + Ok(UnionExec::try_new(vec![ + Arc::clone(&plan), + Arc::clone(&plan), + ])?) + } +} + +/// Creates library C's query planner. +/// +/// `library_a_planner` is the planner library A exported before swapping this one +/// onto its session. When it is absent the planner does its own planning instead +/// of delegating. +pub extern "C" fn create_query_planner( + logical_codec: FFI_LogicalExtensionCodec, + physical_codec: FFI_PhysicalExtensionCodec, + library_a_planner: FFI_Option, +) -> FFI_QueryPlanner { + let planner: Arc = match library_a_planner.as_ref() { + Some(library_a_planner) => Arc::new(SwappedQueryPlanner { + library_a_planner: library_a_planner.into(), + }), + None => Arc::new(TestQueryPlanner), + }; + + FFI_QueryPlanner::new_with_ffi_codecs(planner, logical_codec, physical_codec) +} diff --git a/datafusion/ffi/src/tests/udf_udaf_udwf.rs b/datafusion/ffi/src/tests/udf_udaf_udwf.rs index 399a2cc6be5cd..830c639c743d6 100644 --- a/datafusion/ffi/src/tests/udf_udaf_udwf.rs +++ b/datafusion/ffi/src/tests/udf_udaf_udwf.rs @@ -20,12 +20,15 @@ use std::sync::Arc; use arrow_schema::DataType; use datafusion_catalog::TableFunctionImpl; use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::sort_properties::ExprProperties; use datafusion_expr::{ - AggregateUDF, ColumnarValue, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, - Volatility, WindowUDF, + AggregateUDF, ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDF, + ScalarUDFImpl, Signature, Volatility, WindowUDF, }; use datafusion_functions::math::abs::AbsFunc; use datafusion_functions::math::random::RandomFunc; +use datafusion_functions_aggregate::first_last::FirstValue; use datafusion_functions_aggregate::stddev::Stddev; use datafusion_functions_aggregate::sum::Sum; use datafusion_functions_table::generate_series::RangeFunc; @@ -102,6 +105,13 @@ impl ScalarUDFImpl for TimeZoneUDF { let tz = args.config_options.execution.time_zone.clone(); Ok(ColumnarValue::Scalar(ScalarValue::from(tz))) } + + fn with_updated_config(&self, config: &ConfigOptions) -> Option { + config.execution.time_zone.as_ref()?; + Some(ScalarUDF::from(Self { + signature: self.signature.clone(), + })) + } } pub(crate) extern "C" fn create_timezone_func() -> FFI_ScalarUDF { @@ -112,6 +122,63 @@ pub(crate) extern "C" fn create_timezone_func() -> FFI_ScalarUDF { udf.into() } +#[derive(Debug, PartialEq, Eq, Hash)] +struct PlacementUDF { + signature: Signature, +} + +impl ScalarUDFImpl for PlacementUDF { + fn name(&self) -> &str { + "placement_udf" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type( + &self, + _arg_types: &[DataType], + ) -> datafusion_common::Result { + Ok(DataType::Int64) + } + + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + datafusion_common::internal_err!("placement_udf is not meant to be invoked") + } + + fn placement(&self, args: &[ExpressionPlacement]) -> ExpressionPlacement { + // Push to the leaves only for a (Column, Literal) pairing, so the + // test catches dropped, reordered, or truncated arguments. + if matches!( + args, + [ExpressionPlacement::Column, ExpressionPlacement::Literal] + ) { + ExpressionPlacement::MoveTowardsLeafNodes + } else { + ExpressionPlacement::KeepInPlace + } + } + + fn preserves_lex_ordering( + &self, + inputs: &[ExprProperties], + ) -> datafusion_common::Result { + Ok(inputs.iter().all(|input| input.preserves_lex_ordering)) + } +} + +pub(crate) extern "C" fn create_placement_func() -> FFI_ScalarUDF { + let udf: Arc = Arc::new(ScalarUDF::from(PlacementUDF { + signature: Signature::uniform(1, vec![DataType::Int64], Volatility::Immutable), + })); + + udf.into() +} + pub(crate) extern "C" fn create_ffi_table_func( codec: FFI_LogicalExtensionCodec, ) -> FFI_TableFunction { @@ -126,6 +193,12 @@ pub(crate) extern "C" fn create_ffi_sum_func() -> FFI_AggregateUDF { udaf.into() } +pub(crate) extern "C" fn create_ffi_first_value_func() -> FFI_AggregateUDF { + let udaf: Arc = Arc::new(FirstValue::new().into()); + + udaf.into() +} + pub(crate) extern "C" fn create_ffi_stddev_func() -> FFI_AggregateUDF { let udaf: Arc = Arc::new(Stddev::new().into()); diff --git a/datafusion/ffi/src/tests/utils.rs b/datafusion/ffi/src/tests/utils.rs index e1374c786266b..3119ab96d4032 100644 --- a/datafusion/ffi/src/tests/utils.rs +++ b/datafusion/ffi/src/tests/utils.rs @@ -21,29 +21,6 @@ use datafusion_common::{DataFusionError, Result}; use crate::tests::ForeignLibraryModule; -/// Compute the path to the built cdylib. Checks debug, release, and ci profile dirs. -fn compute_library_dir(target_path: &Path) -> PathBuf { - let debug_dir = target_path.join("debug"); - let release_dir = target_path.join("release"); - let ci_dir = target_path.join("ci"); - - let all_dirs = vec![debug_dir.clone(), release_dir, ci_dir]; - - all_dirs - .into_iter() - .filter(|dir| dir.join("deps").exists()) - .filter_map(|dir| { - dir.join("deps") - .metadata() - .and_then(|m| m.modified()) - .ok() - .map(|date| (dir, date)) - }) - .max_by_key(|(_, date)| *date) - .map(|(dir, _)| dir) - .unwrap_or(debug_dir) -} - /// Find the cdylib file for datafusion_ffi in the given directory. fn find_cdylib(deps_dir: &Path) -> Result { let lib_prefix = if cfg!(target_os = "windows") { @@ -71,23 +48,26 @@ fn find_cdylib(deps_dir: &Path) -> Result { )) } -pub fn get_module() -> Result { - let expected_version = crate::version(); - - let crate_root = Path::new(env!("CARGO_MANIFEST_DIR")); - let target_dir = crate_root - .parent() - .expect("Failed to find crate parent") - .parent() - .expect("Failed to find workspace root") - .join("target"); +/// Locate the built `datafusion_ffi` cdylib. +/// +/// The cdylib sits next to the running test binary, so this follows Cargo's +/// actual output directory and is robust to the active profile and a custom +/// `--target-dir` (e.g. `cargo llvm-cov`). +fn find_library() -> Result { + let exe = + std::env::current_exe().map_err(|e| DataFusionError::External(Box::new(e)))?; + let deps_dir = exe.parent().ok_or_else(|| { + DataFusionError::External("Failed to find test binary directory".into()) + })?; + find_cdylib(deps_dir) +} - let library_dir = compute_library_dir(target_dir.as_path()); - let lib_path = find_cdylib(&library_dir.join("deps"))?; +fn load_module(lib_path: &Path) -> Result { + let expected_version = crate::version(); // Load the library using libloading let lib = unsafe { - libloading::Library::new(&lib_path) + libloading::Library::new(lib_path) .map_err(|e| DataFusionError::External(Box::new(e)))? }; @@ -105,3 +85,40 @@ pub fn get_module() -> Result { Ok(module) } + +pub fn get_module() -> Result { + load_module(&find_library()?) +} + +/// Load an independent copy of the integration-test cdylib. +/// +/// Copying to a unique path makes the dynamic loader create a separate image +/// with its own library marker and Rust object graph. +pub fn get_module_copy(name: &str) -> Result { + let source = find_library()?; + let file_name = source + .file_name() + .and_then(|name| name.to_str()) + .ok_or_else(|| DataFusionError::External("Invalid cdylib filename".into()))?; + // Windows cannot remove a loaded DLL, so use a stable name that bounds the + // retained test artifacts to one file per library role. + #[cfg(target_os = "windows")] + let destination = source.with_file_name(format!("{name}_{file_name}")); + #[cfg(not(target_os = "windows"))] + let destination = + source.with_file_name(format!("{}_{}_{}", std::process::id(), name, file_name)); + + std::fs::copy(&source, &destination) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + match load_module(&destination) { + Ok(module) => { + #[cfg(not(target_os = "windows"))] + let _ = std::fs::remove_file(destination); + Ok(module) + } + Err(error) => { + let _ = std::fs::remove_file(destination); + Err(error) + } + } +} diff --git a/datafusion/ffi/src/udaf/groups_accumulator.rs b/datafusion/ffi/src/udaf/groups_accumulator.rs index 1600bef39da45..4d1b0b4be0a2b 100644 --- a/datafusion/ffi/src/udaf/groups_accumulator.rs +++ b/datafusion/ffi/src/udaf/groups_accumulator.rs @@ -64,7 +64,6 @@ pub struct FFI_GroupsAccumulator { accumulator: &mut Self, values: SVec, group_indices: SVec, - opt_filter: FFI_Option, total_num_groups: usize, ) -> FFI_Result<()>, @@ -74,8 +73,6 @@ pub struct FFI_GroupsAccumulator { opt_filter: FFI_Option, ) -> FFI_Result>, - pub supports_convert_to_state: bool, - /// Release the memory of the private data when it is no longer being used. pub release: unsafe extern "C" fn(accumulator: &mut Self), @@ -195,21 +192,14 @@ unsafe extern "C" fn merge_batch_fn_wrapper( accumulator: &mut FFI_GroupsAccumulator, values: SVec, group_indices: SVec, - opt_filter: FFI_Option, total_num_groups: usize, ) -> FFI_Result<()> { unsafe { let accumulator = accumulator.inner_mut(); let values = sresult_return!(process_values(values)); let group_indices: Vec = group_indices.into_iter().collect(); - let opt_filter = sresult_return!(process_opt_filter(opt_filter)); - sresult!(accumulator.merge_batch( - &values, - &group_indices, - opt_filter.as_ref(), - total_num_groups - )) + sresult!(accumulator.merge_batch(&values, &group_indices, total_num_groups)) } } @@ -255,7 +245,6 @@ impl From> for FFI_GroupsAccumulator { return accumulator.accumulator; } - let supports_convert_to_state = accumulator.supports_convert_to_state(); let private_data = GroupsAccumulatorPrivateData { accumulator }; Self { @@ -265,7 +254,6 @@ impl From> for FFI_GroupsAccumulator { state: state_fn_wrapper, merge_batch: merge_batch_fn_wrapper, convert_to_state: convert_to_state_fn_wrapper, - supports_convert_to_state, release: release_fn_wrapper, private_data: Box::into_raw(Box::new(private_data)) as *mut c_void, @@ -379,7 +367,6 @@ impl GroupsAccumulator for ForeignGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { unsafe { @@ -388,20 +375,11 @@ impl GroupsAccumulator for ForeignGroupsAccumulator { .map(WrappedArray::try_from) .collect::, ArrowError>>()?; let group_indices = group_indices.iter().cloned().collect(); - let opt_filter = opt_filter - .map(|bool_array| to_ffi(&bool_array.to_data())) - .transpose()? - .map(|(array, schema)| WrappedArray { - array, - schema: WrappedSchema(schema), - }) - .into(); df_result!((self.accumulator.merge_batch)( &mut self.accumulator, values.into_iter().collect(), group_indices, - opt_filter, total_num_groups )) } @@ -439,10 +417,6 @@ impl GroupsAccumulator for ForeignGroupsAccumulator { .collect() } } - - fn supports_convert_to_state(&self) -> bool { - self.accumulator.supports_convert_to_state - } } #[repr(C)] @@ -517,8 +491,7 @@ mod tests { let second_states = vec![make_array(create_array!(Boolean, vec![false]).to_data())]; - let opt_filter = create_array!(Boolean, vec![true]); - foreign_accum.merge_batch(&second_states, &[0], Some(opt_filter.as_ref()), 1)?; + foreign_accum.merge_batch(&second_states, &[0], 1)?; let groups_bool = foreign_accum.evaluate(EmitTo::All)?; assert_eq!(groups_bool.len(), 1); assert_eq!( diff --git a/datafusion/ffi/src/udaf/mod.rs b/datafusion/ffi/src/udaf/mod.rs index c4f8fb1254e84..b3a087e5d0022 100644 --- a/datafusion/ffi/src/udaf/mod.rs +++ b/datafusion/ffi/src/udaf/mod.rs @@ -145,6 +145,10 @@ pub struct FFI_AggregateUDF { /// the foreign interface. See [`crate::get_library_marker_id`] and /// the crate's `README.md` for more information. pub library_marker_id: extern "C" fn() -> usize, + + /// FFI equivalent to [`AggregateUDF::supports_null_handling_clause`] + pub supports_null_handling_clause: + unsafe extern "C" fn(udaf: &FFI_AggregateUDF) -> bool, } unsafe impl Send for FFI_AggregateUDF {} @@ -327,6 +331,12 @@ unsafe extern "C" fn order_sensitivity_fn_wrapper( unsafe { udaf.inner().order_sensitivity().into() } } +unsafe extern "C" fn supports_null_handling_clause_fn_wrapper( + udaf: &FFI_AggregateUDF, +) -> bool { + unsafe { udaf.inner().supports_null_handling_clause() } +} + unsafe extern "C" fn coerce_types_fn_wrapper( udaf: &FFI_AggregateUDF, arg_types: SVec, @@ -401,6 +411,7 @@ impl From> for FFI_AggregateUDF { release: release_fn_wrapper, private_data: Box::into_raw(private_data) as *mut c_void, library_marker_id: crate::get_library_marker_id, + supports_null_handling_clause: supports_null_handling_clause_fn_wrapper, } } } @@ -595,6 +606,10 @@ impl AggregateUDFImpl for ForeignAggregateUDF { unsafe { (self.udaf.order_sensitivity)(&self.udaf).into() } } + fn supports_null_handling_clause(&self) -> bool { + unsafe { (self.udaf.supports_null_handling_clause)(&self.udaf) } + } + fn simplify(&self) -> Option { None } @@ -774,6 +789,19 @@ mod tests { Ok(()) } + #[test] + fn test_supports_null_handling_clause() -> Result<()> { + let first_value = create_test_foreign_udaf( + datafusion::functions_aggregate::first_last::FirstValue::new(), + )?; + assert!(first_value.supports_null_handling_clause()); + + let sum = create_test_foreign_udaf(Sum::new())?; + assert!(!sum.supports_null_handling_clause()); + + Ok(()) + } + #[test] fn test_beneficial_ordering() -> Result<()> { let foreign_udaf = create_test_foreign_udaf( diff --git a/datafusion/ffi/src/udf/mod.rs b/datafusion/ffi/src/udf/mod.rs index ff18a30e4ba19..fa08cbd042330 100644 --- a/datafusion/ffi/src/udf/mod.rs +++ b/datafusion/ffi/src/udf/mod.rs @@ -26,10 +26,11 @@ use arrow::ffi::{FFI_ArrowSchema, from_ffi, to_ffi}; use arrow_schema::FieldRef; use datafusion_common::config::ConfigOptions; use datafusion_common::{DataFusionError, Result, internal_err}; +use datafusion_expr::sort_properties::ExprProperties; use datafusion_expr::type_coercion::functions::fields_with_udf; use datafusion_expr::{ - ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, - Signature, + ColumnarValue, ExpressionPlacement, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, + ScalarUDFImpl, Signature, }; use return_type_args::{ FFI_ReturnFieldArgs, ForeignReturnFieldArgs, ForeignReturnFieldArgsOwned, @@ -41,8 +42,10 @@ use stabby::vec::Vec as SVec; use crate::arrow_wrappers::{WrappedArray, WrappedSchema}; use crate::config::FFI_ConfigOptions; use crate::expr::columnar_value::FFI_ColumnarValue; +use crate::expr::expr_properties::FFI_ExprProperties; +use crate::placement::FFI_ExpressionPlacement; use crate::util::{ - FFI_Result, rvec_wrapped_to_vec_datatype, vec_datatype_to_rvec_wrapped, + FFI_Option, FFI_Result, rvec_wrapped_to_vec_datatype, vec_datatype_to_rvec_wrapped, }; use crate::volatility::FFI_Volatility; use crate::{df_result, sresult, sresult_return}; @@ -91,6 +94,14 @@ pub struct FFI_ScalarUDF { arg_types: SVec, ) -> FFI_Result>, + /// FFI equivalent to the `placement` of a [`ScalarUDFImpl`]. Returns the + /// placement hint for the underlying [`ScalarUDF`] given each argument's + /// placement. Infallible, so it returns the value directly, not an `FFI_Result`. + pub placement: unsafe extern "C" fn( + udf: &Self, + args: SVec, + ) -> FFI_ExpressionPlacement, + /// Used to create a clone on the provider of the udf. This should /// only need to be called by the receiver of the udf. pub clone: unsafe extern "C" fn(udf: &Self) -> Self, @@ -106,6 +117,19 @@ pub struct FFI_ScalarUDF { /// the foreign interface. See [`crate::get_library_marker_id`] and /// the crate's `README.md` for more information. pub library_marker_id: extern "C" fn() -> usize, + + /// FFI equivalent to [`ScalarUDFImpl::preserves_lex_ordering`]. + pub preserves_lex_ordering: unsafe extern "C" fn( + udf: &Self, + inputs: SVec, + ) -> FFI_Result, + + /// FFI equivalent to [`ScalarUDFImpl::with_updated_config`]. + pub with_updated_config: + unsafe extern "C" fn( + udf: &Self, + config: FFI_ConfigOptions, + ) -> FFI_Result>, } unsafe impl Send for FFI_ScalarUDF {} @@ -157,6 +181,46 @@ unsafe extern "C" fn coerce_types_fn_wrapper( sresult!(vec_datatype_to_rvec_wrapped(&return_types)) } +unsafe extern "C" fn placement_fn_wrapper( + udf: &FFI_ScalarUDF, + args: SVec, +) -> FFI_ExpressionPlacement { + let args = args + .into_iter() + .map(ExpressionPlacement::from) + .collect::>(); + + udf.inner().placement(&args).into() +} + +unsafe extern "C" fn preserves_lex_ordering_fn_wrapper( + udf: &FFI_ScalarUDF, + inputs: SVec, +) -> FFI_Result { + let result = inputs + .into_iter() + .map(ExprProperties::try_from) + .collect::>>() + .and_then(|inputs| udf.inner().preserves_lex_ordering(&inputs)); + + sresult!(result) +} + +unsafe extern "C" fn with_updated_config_fn_wrapper( + udf: &FFI_ScalarUDF, + config: FFI_ConfigOptions, +) -> FFI_Result> { + let config = sresult_return!(ConfigOptions::try_from(config)); + + let updated: Option = udf + .inner() + .inner() + .with_updated_config(&config) + .map(|updated| Arc::new(updated).into()); + + FFI_Result::Ok(updated.into()) +} + unsafe extern "C" fn invoke_with_args_fn_wrapper( udf: &FFI_ScalarUDF, args: SVec, @@ -250,10 +314,13 @@ impl From> for FFI_ScalarUDF { invoke_with_args: invoke_with_args_fn_wrapper, return_field_from_args: return_field_from_args_fn_wrapper, coerce_types: coerce_types_fn_wrapper, + placement: placement_fn_wrapper, clone: clone_fn_wrapper, release: release_fn_wrapper, private_data: Box::into_raw(private_data) as *mut c_void, library_marker_id: crate::get_library_marker_id, + preserves_lex_ordering: preserves_lex_ordering_fn_wrapper, + with_updated_config: with_updated_config_fn_wrapper, } } } @@ -281,6 +348,21 @@ pub struct ForeignScalarUDF { unsafe impl Send for ForeignScalarUDF {} unsafe impl Sync for ForeignScalarUDF {} +impl ForeignScalarUDF { + fn new(udf: FFI_ScalarUDF) -> Self { + let name = udf.name.to_string(); + let signature = Signature::user_defined((&udf.volatility).into()); + let aliases = udf.aliases.iter().map(|s| s.to_string()).collect(); + + Self { + name, + aliases, + udf, + signature, + } + } +} + impl PartialEq for ForeignScalarUDF { fn eq(&self, other: &Self) -> bool { let Self { @@ -312,22 +394,22 @@ impl Hash for ForeignScalarUDF { } } +impl From for Arc { + fn from(udf: FFI_ScalarUDF) -> Self { + if (udf.library_marker_id)() == crate::get_library_marker_id() { + Arc::clone(udf.inner().inner()) + } else { + Arc::new(ForeignScalarUDF::new(udf)) + } + } +} + impl From<&FFI_ScalarUDF> for Arc { fn from(udf: &FFI_ScalarUDF) -> Self { if (udf.library_marker_id)() == crate::get_library_marker_id() { Arc::clone(udf.inner().inner()) } else { - let name = udf.name.to_string(); - let signature = Signature::user_defined((&udf.volatility).into()); - - let aliases = udf.aliases.iter().map(|s| s.to_string()).collect(); - - Arc::new(ForeignScalarUDF { - name, - udf: udf.clone(), - aliases, - signature, - }) + Arc::new(ForeignScalarUDF::new(udf.clone())) } } } @@ -427,12 +509,101 @@ impl ScalarUDFImpl for ForeignScalarUDF { Ok(rvec_wrapped_to_vec_datatype(&result_types)?) } } + + fn placement(&self, args: &[ExpressionPlacement]) -> ExpressionPlacement { + let args = args + .iter() + .map(|p| FFI_ExpressionPlacement::from(*p)) + .collect::>(); + + let result = unsafe { (self.udf.placement)(&self.udf, args) }; + + result.into() + } + + fn preserves_lex_ordering(&self, inputs: &[ExprProperties]) -> Result { + inputs + .iter() + .map(FFI_ExprProperties::try_from) + .collect::>>() + .and_then(|inputs| { + let result = + unsafe { (self.udf.preserves_lex_ordering)(&self.udf, inputs) }; + df_result!(result) + }) + } + + fn with_updated_config(&self, config: &ConfigOptions) -> Option { + let config: FFI_ConfigOptions = config.into(); + + let result = unsafe { (self.udf.with_updated_config)(&self.udf, config) }; + + let updated = match df_result!(result) { + Ok(updated) => updated.into_option()?, + Err(error) => { + log::warn!("Unable to update scalar UDF configuration over FFI: {error}"); + return None; + } + }; + + Some(ScalarUDF::new_from_shared_impl(updated.into())) + } } #[cfg(test)] mod tests { use super::*; + #[derive(Debug, PartialEq, Eq, Hash)] + struct PlacementUDF { + signature: Signature, + } + + impl ScalarUDFImpl for PlacementUDF { + fn name(&self) -> &str { + "placement_udf" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + internal_err!("placement_udf is not meant to be invoked") + } + + fn placement(&self, args: &[ExpressionPlacement]) -> ExpressionPlacement { + // Push to the leaves only for a (Column, Literal) pairing, so the + // test catches dropped, reordered, or truncated arguments. + if matches!( + args, + [ExpressionPlacement::Column, ExpressionPlacement::Literal] + ) { + ExpressionPlacement::MoveTowardsLeafNodes + } else { + ExpressionPlacement::KeepInPlace + } + } + + fn preserves_lex_ordering(&self, inputs: &[ExprProperties]) -> Result { + if inputs.is_empty() { + return internal_err!("preserves_lex_ordering requires an input"); + } + + Ok(inputs.iter().all(|input| input.preserves_lex_ordering)) + } + + fn with_updated_config(&self, _config: &ConfigOptions) -> Option { + Some(ScalarUDF::from(Self { + signature: self.signature.clone(), + })) + } + } + #[test] fn test_round_trip_scalar_udf() -> Result<()> { let original_udf = datafusion::functions::math::abs::AbsFunc::new(); @@ -444,6 +615,11 @@ mod tests { let foreign_udf: Arc = (&local_udf).into(); assert_eq!(original_udf.name(), foreign_udf.name()); + assert!( + foreign_udf + .with_updated_config(&ConfigOptions::default()) + .is_none() + ); Ok(()) } @@ -467,4 +643,66 @@ mod tests { Ok(()) } + + #[test] + fn test_ffi_udf_placement_round_trip() -> Result<()> { + use datafusion_expr::Volatility; + + let original_udf = Arc::new(ScalarUDF::from(PlacementUDF { + signature: Signature::uniform( + 1, + vec![DataType::Int64], + Volatility::Immutable, + ), + })); + + let mut ffi_udf = FFI_ScalarUDF::from(original_udf); + + // Force the foreign path so the call travels through the FFI vtable + // rather than downcasting back to the original local type. + ffi_udf.library_marker_id = crate::mock_foreign_marker_id; + let foreign_udf: Arc = (&ffi_udf).into(); + assert!(foreign_udf.is::()); + + // Without the plumbing the override is dropped and every call is + // KeepInPlace. The three cases also check the arguments survive the + // round trip in order. + assert_eq!( + foreign_udf + .placement(&[ExpressionPlacement::Column, ExpressionPlacement::Literal]), + ExpressionPlacement::MoveTowardsLeafNodes + ); + assert_eq!( + foreign_udf + .placement(&[ExpressionPlacement::Literal, ExpressionPlacement::Column]), + ExpressionPlacement::KeepInPlace + ); + assert_eq!(foreign_udf.placement(&[]), ExpressionPlacement::KeepInPlace); + + let preserves = ExprProperties::new_unknown().with_preserves_lex_ordering(true); + let does_not_preserve = ExprProperties::new_unknown(); + + assert!( + foreign_udf + .preserves_lex_ordering(std::slice::from_ref(&preserves)) + .unwrap() + ); + assert!( + !foreign_udf + .preserves_lex_ordering(&[preserves, does_not_preserve]) + .unwrap() + ); + assert!(foreign_udf.preserves_lex_ordering(&[]).is_err()); + + let updated = foreign_udf + .with_updated_config(&ConfigOptions::default()) + .expect("provider should return an updated UDF"); + assert_eq!( + updated + .placement(&[ExpressionPlacement::Column, ExpressionPlacement::Literal]), + ExpressionPlacement::MoveTowardsLeafNodes + ); + + Ok(()) + } } diff --git a/datafusion/ffi/src/udtf.rs b/datafusion/ffi/src/udtf.rs index 0a111028798d1..fa28519d58de5 100644 --- a/datafusion/ffi/src/udtf.rs +++ b/datafusion/ffi/src/udtf.rs @@ -155,7 +155,7 @@ unsafe extern "C" fn call_with_args_wrapper( let session = sresult_return!( session .as_local() - .map(Ok::<&(dyn Session + Send + Sync), DataFusionError>) + .map(Ok::<&dyn Session, DataFusionError>) .unwrap_or_else(|| { foreign_session = Some(ForeignSession::try_from(&session)?); Ok(foreign_session.as_ref().unwrap()) diff --git a/datafusion/ffi/tests/ffi_execution_plan.rs b/datafusion/ffi/tests/ffi_execution_plan.rs index bd84e064de4c2..4067d7eb49b2a 100644 --- a/datafusion/ffi/tests/ffi_execution_plan.rs +++ b/datafusion/ffi/tests/ffi_execution_plan.rs @@ -21,14 +21,20 @@ mod tests { use arrow::datatypes::Schema; use arrow_schema::DataType; use datafusion_common::DataFusionError; - use datafusion_ffi::execution_plan::FFI_ExecutionPlan; - use datafusion_ffi::execution_plan::ForeignExecutionPlan; - use datafusion_ffi::execution_plan::{ExecutionPlanPrivateData, tests::EmptyExec}; + use datafusion_common::tree_node::TreeNodeRecursion; + use datafusion_ffi::execution_plan::{ + ExecutionPlanPrivateData, FFI_ExecutionPlan, ForeignExecutionPlan, + tests::EmptyExec, + }; use datafusion_ffi::tests::utils::get_module; - use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::execution_plan::InvariantLevel; + use datafusion_physical_plan::{ + ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions, + }; use std::sync::Arc; #[test] + #[expect(deprecated)] fn test_ffi_execution_plan_partition_statistics_cross_library() -> Result<(), DataFusionError> { let module = get_module()?; @@ -62,6 +68,47 @@ mod tests { Ok(()) } + #[test] + fn test_ffi_execution_plan_expressions_cross_library() -> Result<(), DataFusionError> + { + let module = get_module()?; + let plan = (module.create_exec_with_expressions)(); + let plan: Arc = (&plan).try_into()?; + assert!(plan.is::()); + + let mut retained = None; + plan.apply_expressions(&mut |expr| { + retained = Some(Arc::clone(expr)); + Ok(TreeNodeRecursion::Continue) + })?; + drop(plan); + + assert!( + retained + .as_ref() + .and_then(|expr| expr.expression_id()) + .is_some() + ); + Ok(()) + } + + #[test] + fn test_ffi_execution_plan_dynamic_expressions_cross_library() + -> Result<(), DataFusionError> { + let module = get_module()?; + let plan = (module.create_exec_with_dynamic_expressions)(); + let plan: Arc = (&plan).try_into()?; + assert!(plan.is::()); + plan.check_invariants(InvariantLevel::Always)?; + + let produced = plan.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + assert!(produced[0].expression_id().is_some()); + drop(plan); + assert!(produced[0].expression_id().is_some()); + Ok(()) + } + #[test] fn test_ffi_execution_plan_new_sets_runtimes_on_children() -> Result<(), DataFusionError> { @@ -91,7 +138,10 @@ mod tests { let grandchild_plan = generate_local_plan(); - let child_plan = child_plan.with_new_children(vec![grandchild_plan])?; + let child_plan = child_plan.replace_children( + vec![grandchild_plan], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; unsafe { // Originally the runtime is not set. We go through the unsafe casting @@ -106,7 +156,10 @@ mod tests { assert!((*grandchild_private_data).runtime.is_none()); } - let parent_plan = generate_local_plan().with_new_children(vec![child_plan])?; + let parent_plan = generate_local_plan().replace_children( + vec![child_plan], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; // Adding the grandchild beneath this FFI plan should get the runtime passed down. let runtime = tokio::runtime::Builder::new_current_thread() diff --git a/datafusion/ffi/tests/ffi_integration.rs b/datafusion/ffi/tests/ffi_integration.rs index 6a6b6b3100cdb..86f953e262ead 100644 --- a/datafusion/ffi/tests/ffi_integration.rs +++ b/datafusion/ffi/tests/ffi_integration.rs @@ -58,6 +58,15 @@ mod tests { assert!(results.contains(&create_record_batch(6, 1))); assert!(results.contains(&create_record_batch(7, 5))); + if !synchronous { + assert!( + ctx.state() + .catalog_list() + .catalog("ffi_registered") + .is_some() + ); + } + Ok(()) } @@ -100,7 +109,7 @@ mod tests { let cmd = CreateExternalTable { schema: Schema::empty().to_dfschema_ref()?, name: TableReference::bare("cloned_test"), - location: "test".to_string(), + locations: vec!["test".to_string()], file_type: "test".to_string(), table_partition_cols: vec![], if_not_exists: false, diff --git a/datafusion/ffi/tests/ffi_physical_optimizer.rs b/datafusion/ffi/tests/ffi_physical_optimizer.rs index d860fda340ae6..d8baf522889e8 100644 --- a/datafusion/ffi/tests/ffi_physical_optimizer.rs +++ b/datafusion/ffi/tests/ffi_physical_optimizer.rs @@ -25,7 +25,7 @@ mod tests { use datafusion_ffi::execution_plan::tests::EmptyExec; use datafusion_ffi::physical_optimizer::ForeignPhysicalOptimizerRule; use datafusion_ffi::tests::utils::get_module; - use datafusion_physical_optimizer::PhysicalOptimizerRule; + use datafusion_physical_optimizer::{ConfigOnlyContext, PhysicalOptimizerRule}; use datafusion_physical_plan::ExecutionPlan; fn create_test_plan() -> Arc { @@ -66,4 +66,30 @@ mod tests { Ok(()) } + + #[test] + fn test_ffi_physical_optimizer_rule_with_context() -> Result<(), DataFusionError> { + let module = get_module()?; + + let ffi_rule = (module.create_context_aware_optimizer_rule)(); + + let foreign_rule: Arc = + (&ffi_rule).into(); + + // Verify that plain optimize fails (proving we need context path) + let plan = create_test_plan(); + let config = ConfigOptions::new(); + assert!(foreign_rule.optimize(plan, &config).is_err()); + + // Verify context-aware path works + let plan = create_test_plan(); + let context = ConfigOnlyContext::new(&config); + let optimized = foreign_rule.optimize_with_context(plan, &context)?; + + assert_eq!(optimized.name(), "GlobalLimitExec"); + assert_eq!(optimized.children().len(), 1); + assert_eq!(optimized.children()[0].name(), "empty-exec"); + + Ok(()) + } } diff --git a/datafusion/ffi/tests/ffi_query_planner.rs b/datafusion/ffi/tests/ffi_query_planner.rs new file mode 100644 index 0000000000000..c72b8c3e889ae --- /dev/null +++ b/datafusion/ffi/tests/ffi_query_planner.rs @@ -0,0 +1,348 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod utils; + +#[cfg(feature = "integration-tests")] +mod tests { + use std::sync::{Arc, OnceLock, Weak}; + + use arrow::datatypes::SchemaRef; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::SessionContext; + use datafusion_catalog::TableProvider; + use datafusion_common::{ + DataFusionError, Result, TableReference, exec_err, not_impl_err, + }; + use datafusion_execution::{TaskContext, TaskContextProvider}; + use datafusion_expr::logical_plan::Extension; + use datafusion_expr::{LogicalPlan, col}; + use datafusion_ffi::execution::FFI_TaskContextProvider; + use datafusion_ffi::execution_plan::ForeignExecutionPlan; + use datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec; + use datafusion_ffi::proto::physical_extension_codec::FFI_PhysicalExtensionCodec; + use datafusion_ffi::query_planner::{FFI_QueryPlanner, ForeignQueryPlanner}; + use datafusion_ffi::table_provider::ForeignTableProvider; + use datafusion_ffi::tests::{ + create_test_schema, + utils::{get_module, get_module_copy}, + }; + use datafusion_ffi::util::FFI_Option; + use datafusion_physical_plan::ExecutionPlan; + use datafusion_physical_plan::empty::EmptyExec; + use datafusion_physical_plan::sorts::sort::SortExec; + use datafusion_physical_plan::union::UnionExec; + use datafusion_proto::logical_plan::LogicalExtensionCodec; + use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, + PhysicalProtoConverterExtension, + }; + use datafusion_session::QueryPlanner; + + #[tokio::test] + async fn test_ffi_query_planner() -> Result<(), DataFusionError> { + let module = get_module()?; + let (ctx, logical_codec) = crate::utils::ctx_and_codec(); + let task_ctx_provider = Arc::clone(&ctx) as Arc; + let task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let physical_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(DefaultPhysicalExtensionCodec {}), + None, + task_ctx_provider, + ); + + let ffi_planner = (module.create_query_planner)( + logical_codec, + physical_codec, + FFI_Option::None, + ); + let planner: Arc = (&ffi_planner).into(); + + let any_ref: &dyn std::any::Any = planner.as_ref(); + assert!(any_ref.downcast_ref::().is_some()); + + let logical_plan = datafusion_expr::LogicalPlanBuilder::empty(false).build()?; + let state = ctx.state(); + let physical_plan = planner.create_physical_plan(&logical_plan, &state).await?; + + assert_eq!(physical_plan.name(), "EmptyExec"); + assert!(physical_plan.is::()); + + Ok(()) + } + + /// Test-only codec that preserves library B's table provider while the logical + /// plan crosses between library A and library C. + /// + /// Encoding writes a fixed identifier and stores a weak reference to the + /// provider. Decoding validates the identifier and upgrades that reference. + /// This works because all three test libraries run in one process and library + /// A's session continues to own the provider. + /// + /// This is not a general serialization format for table providers. A + /// cross-process deployment must provide its own codec that either resolves a + /// stable identifier through shared state or reconstructs the provider from a + /// portable, provider-specific description. DataFusion passes the table + /// reference, schema, and task context separately to the decoder. + #[derive(Debug, Default)] + struct LibraryALogicalCodec { + library_b_provider: OnceLock>, + } + + impl LogicalExtensionCodec for LibraryALogicalCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[LogicalPlan], + _ctx: &TaskContext, + ) -> Result { + not_impl_err!("logical extension nodes are not used in this test") + } + + fn try_encode(&self, _node: &Extension, _buf: &mut Vec) -> Result<()> { + not_impl_err!("logical extension nodes are not used in this test") + } + + fn try_decode_table_provider( + &self, + buf: &[u8], + _table_ref: &TableReference, + _schema: SchemaRef, + _ctx: &TaskContext, + ) -> Result> { + if buf != b"library-b-provider" { + return exec_err!("unexpected library B provider payload"); + } + self.library_b_provider + .get() + .and_then(Weak::upgrade) + .ok_or_else(|| DataFusionError::Plan("missing library B provider".into())) + } + + fn try_encode_table_provider( + &self, + _table_ref: &TableReference, + node: Arc, + buf: &mut Vec, + ) -> Result<()> { + self.library_b_provider + .get_or_init(|| Arc::downgrade(&node)); + buf.extend_from_slice(b"library-b-provider"); + Ok(()) + } + } + + /// Library A's physical codec reconstructs B's opaque foreign plan as an + /// A-local test plan when the result returns from library C. + /// + /// Encoding sees B's node in one of two shapes. When A serializes a plan it + /// built itself, B's scan is a [`ForeignExecutionPlan`]. When library C + /// serializes a plan containing a node A previously handed it, the FFI handle + /// unwraps back to its home library, so A is asked to encode the very + /// [`EmptyExec`] its own `try_decode` produced. + #[derive(Debug)] + struct LibraryAPhysicalCodec; + + impl PhysicalExtensionCodec for LibraryAPhysicalCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + if buf != b"library-b-empty-exec" || !inputs.is_empty() { + return exec_err!("unexpected library B execution plan payload"); + } + Ok(Arc::new(EmptyExec::new(create_test_schema()))) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + if !node.is::() && !node.is::() { + return exec_err!( + "expected library B's plan to be foreign or A-local; got {}", + node.name() + ); + } + buf.extend_from_slice(b"library-b-empty-exec"); + Ok(()) + } + } + + #[tokio::test] + async fn test_three_library_query_planner_restores_type_identity() -> Result<()> { + // Library A: datafusion-python owns the session and codec registry. + let state = SessionStateBuilder::new_with_default_features() + .with_physical_optimizer_rules(vec![]) + .build(); + let ctx = Arc::new(SessionContext::new_with_state(state)); + let task_ctx_provider = Arc::clone(&ctx) as Arc; + let ffi_task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let logical_codec = FFI_LogicalExtensionCodec::new( + Arc::new(LibraryALogicalCodec::default()), + None, + ffi_task_ctx_provider.clone(), + ); + let physical_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(LibraryAPhysicalCodec), + None, + ffi_task_ctx_provider, + ); + + let library_b = get_module_copy("query_planner_library_b")?; + let library_c = get_module_copy("query_planner_library_c")?; + + // Library B: reuse the synchronous table provider from the existing + // FFI integration-test module. + let ffi_provider = (library_b.create_table)(true, logical_codec.clone()); + let provider: Arc = (&ffi_provider).into(); + assert!(provider.downcast_ref::().is_some()); + ctx.register_table("library_b", provider)?; + let logical_plan = ctx.table("library_b").await?.into_optimized_plan()?; + + // Library C: a foreign query planner sees B's scan result as opaque, + // but can downcast its own UnionExec. Its result is serialized rather + // than returned as FFI_ExecutionPlan. + let ffi_planner = (library_c.create_query_planner)( + logical_codec, + physical_codec, + FFI_Option::None, + ); + let planner: Arc = (&ffi_planner).into(); + let planner_any: &dyn std::any::Any = planner.as_ref(); + assert!(planner_any.downcast_ref::().is_some()); + + let state = ctx.state(); + let physical_plan = planner.create_physical_plan(&logical_plan, &state).await?; + + // Deserialization in A reconstructs the full result as A-local + // concrete nodes, including the plans that originated in B. + assert!(physical_plan.is::()); + assert!(!physical_plan.is::()); + let children = physical_plan.children(); + assert_eq!(children.len(), 2); + assert!(children.iter().all(|child| child.is::())); + assert!( + children + .iter() + .all(|child| !child.is::()) + ); + + Ok(()) + } + + /// Exercises the deployment library C actually uses: library A hands its own + /// query planner to C, then installs C's planner on the session it already + /// owns. C plans by delegating back to A's captured planner. + /// + /// This is the case that requires serialized plans in both directions. C must + /// downcast the nodes A produced in order to rewrite them, and A must downcast + /// the nodes C produced in order to run its own passes over the result. + #[tokio::test] + async fn test_query_planner_swap_round_trips_type_identity() -> Result<()> { + // Library A: datafusion-python owns the session and codec registry. The + // physical optimizer rules are cleared so the assertions below observe + // planning alone. + let state = SessionStateBuilder::new_with_default_features() + .with_physical_optimizer_rules(vec![]) + .build(); + let ctx = Arc::new(SessionContext::new_with_state(state)); + let task_ctx_provider = Arc::clone(&ctx) as Arc; + let ffi_task_ctx_provider = FFI_TaskContextProvider::from(&task_ctx_provider); + let logical_codec = FFI_LogicalExtensionCodec::new( + Arc::new(LibraryALogicalCodec::default()), + None, + ffi_task_ctx_provider.clone(), + ); + let physical_codec = FFI_PhysicalExtensionCodec::new( + Arc::new(LibraryAPhysicalCodec), + None, + ffi_task_ctx_provider, + ); + + let library_b = get_module_copy("planner_swap_library_b")?; + let library_c = get_module_copy("planner_swap_library_c")?; + + // Library B: a table provider that is foreign to both A and C. + let ffi_provider = (library_b.create_table)(true, logical_codec.clone()); + let provider: Arc = (&ffi_provider).into(); + ctx.register_table("library_b", provider)?; + + // Library A exports its default planner *before* the swap. Fetching it + // afterwards through `FFI_SessionRef::query_planner` would hand library C + // its own planner back. + let library_a_planner = Arc::clone(ctx.state().query_planner()); + let ffi_library_a_planner = FFI_QueryPlanner::new_with_ffi_codecs( + library_a_planner, + logical_codec.clone(), + physical_codec.clone(), + ); + + // Library C: builds its planner around A's planner. + let ffi_planner = (library_c.create_query_planner)( + logical_codec, + physical_codec, + FFI_Option::Some(ffi_library_a_planner), + ); + let library_c_planner: Arc = + (&ffi_planner).into(); + let planner_any: &dyn std::any::Any = library_c_planner.as_ref(); + assert!(planner_any.downcast_ref::().is_some()); + + // Library A swaps C's planner into the session it already owns. Mutating + // the existing state keeps the `Arc` identity stable, so + // the task context provider captured by the codecs above stays current. + let state_ref = ctx.state_ref(); + let swapped = SessionStateBuilder::new_from_existing(state_ref.read().clone()) + .with_query_planner(library_c_planner) + .build(); + *state_ref.write() = swapped; + + // A sort keeps a well-known, non-extension node at the root of A's + // physical plan. A projection or limit would be pushed into the scan, + // leaving only library B's opaque node for C to inspect. + let logical_plan = ctx + .table("library_b") + .await? + .sort(vec![col("a").sort(true, true)])? + .into_optimized_plan()?; + + // Planning now runs A -> C -> A -> C -> A across three library images. + let physical_plan = ctx.state().create_physical_plan(&logical_plan).await?; + + // Library A reconstructs C's result as A-local concrete nodes, including + // the plan that originated in B. + assert!(physical_plan.is::()); + assert!(!physical_plan.is::()); + let children = physical_plan.children(); + assert_eq!(children.len(), 2); + for child in &children { + let sort = child + .downcast_ref::() + .expect("library A could not downcast the SortExec it planned"); + assert!(sort.input().is::()); + assert!(!sort.input().is::()); + } + + Ok(()) + } +} diff --git a/datafusion/ffi/tests/ffi_udaf.rs b/datafusion/ffi/tests/ffi_udaf.rs index 7df3404d7421b..090151416e4e9 100644 --- a/datafusion/ffi/tests/ffi_udaf.rs +++ b/datafusion/ffi/tests/ffi_udaf.rs @@ -21,8 +21,7 @@ mod tests { use std::sync::Arc; - use arrow::array::Float64Array; - use datafusion::common::record_batch; + use arrow::array::{Float64Array, record_batch}; use datafusion::error::Result; use datafusion::logical_expr::{AggregateUDF, AggregateUDFImpl}; use datafusion::prelude::{SessionContext, col}; @@ -67,6 +66,22 @@ mod tests { Ok(()) } + #[test] + fn test_supports_null_handling_clause() -> Result<()> { + let module = get_module()?; + + let ffi_first_value_func = (module.create_first_value_udaf)(); + let foreign_first_value_func: Arc = + (&ffi_first_value_func).into(); + assert!(foreign_first_value_func.supports_null_handling_clause()); + + let ffi_sum_func = (module.create_sum_udaf)(); + let foreign_sum_func: Arc = (&ffi_sum_func).into(); + assert!(!foreign_sum_func.supports_null_handling_clause()); + + Ok(()) + } + #[tokio::test] async fn test_ffi_grouping_udaf() -> Result<()> { let module = get_module()?; diff --git a/datafusion/ffi/tests/ffi_udf.rs b/datafusion/ffi/tests/ffi_udf.rs index 6e6cb31f53133..10e0bb5cc1c80 100644 --- a/datafusion/ffi/tests/ffi_udf.rs +++ b/datafusion/ffi/tests/ffi_udf.rs @@ -19,14 +19,15 @@ /// when the feature integration-tests is built #[cfg(feature = "integration-tests")] mod tests { - use arrow::array::{Array, AsArray}; + use arrow::array::{Array, AsArray, record_batch}; use arrow::datatypes::DataType; - use datafusion::common::record_batch; + use datafusion::common::config::ConfigOptions; use datafusion::error::Result; - use datafusion::logical_expr::{ScalarUDF, ScalarUDFImpl}; + use datafusion::logical_expr::{ExpressionPlacement, ScalarUDF, ScalarUDFImpl}; use datafusion::prelude::{SessionContext, col}; use datafusion_execution::config::SessionConfig; use datafusion_expr::lit; + use datafusion_expr::sort_properties::ExprProperties; use datafusion_ffi::tests::create_record_batch; use datafusion_ffi::tests::utils::get_module; use std::sync::Arc; @@ -91,6 +92,36 @@ mod tests { Ok(()) } + /// Checks planning-property overrides across the FFI boundary. + #[tokio::test] + async fn test_scalar_udf_placement() -> Result<()> { + let module = get_module()?; + + let ffi_placement_func = (module.create_placement_udf)(); + let foreign_func: Arc = (&ffi_placement_func).into(); + + // The override pushes to the leaves only for (Column, Literal), so these + // also check the arguments cross the boundary in order. + assert_eq!( + foreign_func + .placement(&[ExpressionPlacement::Column, ExpressionPlacement::Literal]), + ExpressionPlacement::MoveTowardsLeafNodes + ); + assert_eq!( + foreign_func + .placement(&[ExpressionPlacement::Literal, ExpressionPlacement::Column]), + ExpressionPlacement::KeepInPlace + ); + + let preserves = ExprProperties::new_unknown().with_preserves_lex_ordering(true); + let does_not_preserve = ExprProperties::new_unknown(); + + assert!(foreign_func.preserves_lex_ordering(std::slice::from_ref(&preserves))?); + assert!(!foreign_func.preserves_lex_ordering(&[preserves, does_not_preserve])?); + + Ok(()) + } + #[tokio::test] async fn test_config_on_scalar_udf() -> Result<()> { let module = get_module()?; @@ -127,4 +158,30 @@ mod tests { Ok(()) } + + /// Validates that a provider's `with_updated_config` override survives the + /// FFI boundary (the trait default returns `None`). + #[test] + fn test_with_updated_config_on_scalar_udf() -> Result<()> { + let module = get_module()?; + + let ffi_udf = (module.create_timezone_udf)(); + let foreign_udf: Arc = (&ffi_udf).into(); + + assert!( + foreign_udf + .with_updated_config(&ConfigOptions::default()) + .is_none() + ); + + let mut options = ConfigOptions::default(); + options.execution.time_zone = Some("AEST".into()); + + let updated = foreign_udf + .with_updated_config(&options) + .expect("provider should return an updated UDF"); + assert_eq!(updated.name(), "TimeZoneUDF"); + + Ok(()) + } } diff --git a/datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs b/datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs index 0a4c1692baa84..0394a8391ad70 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/avg_distinct/decimal.rs @@ -16,14 +16,14 @@ // under the License. use arrow::{ - array::{ArrayRef, ArrowNumericType}, - datatypes::{ - Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType, i256, - }, + array::{ArrayRef, ArrowNativeTypeOp, ArrowNumericType}, + compute::DecimalCast, + datatypes::{ArrowNativeType, DecimalType}, }; -use datafusion_common::{Result, ScalarValue}; +use datafusion_common::{Result, ScalarValue, exec_datafusion_err, exec_err}; use datafusion_expr_common::accumulator::Accumulator; use std::fmt::Debug; +use std::marker::PhantomData; use std::mem::size_of_val; use crate::aggregate::sum_distinct::DistinctSumAccumulator; @@ -31,33 +31,46 @@ use crate::utils::DecimalAverager; /// Generic implementation of `AVG DISTINCT` for Decimal types. /// Handles both all Arrow decimal types (32, 64, 128 and 256 bits). +/// +/// The distinct values are stored in the input type `I`; only the intermediate +/// sum is computed in the (never narrower) sum type `S` so it cannot overflow +/// `I`'s native type. #[derive(Debug)] -pub struct DecimalDistinctAvgAccumulator { - sum_accumulator: DistinctSumAccumulator, +pub struct DecimalDistinctAvgAccumulator< + I: DecimalType + Debug, + S: DecimalType + Debug = I, +> { + sum_accumulator: DistinctSumAccumulator, sum_scale: i8, target_precision: u8, target_scale: i8, + _sum_type: PhantomData, } -impl DecimalDistinctAvgAccumulator { +impl DecimalDistinctAvgAccumulator { pub fn with_decimal_params( sum_scale: i8, target_precision: u8, target_scale: i8, ) -> Self { - let data_type = T::TYPE_CONSTRUCTOR(T::MAX_PRECISION, sum_scale); + let data_type = I::TYPE_CONSTRUCTOR(I::MAX_PRECISION, sum_scale); Self { sum_accumulator: DistinctSumAccumulator::new(&data_type), sum_scale, target_precision, target_scale, + _sum_type: PhantomData, } } } -impl Accumulator - for DecimalDistinctAvgAccumulator +impl Accumulator for DecimalDistinctAvgAccumulator +where + I: DecimalType + ArrowNumericType + Debug, + S: DecimalType + ArrowNumericType + Debug, + I::Native: Into + DecimalCast, + S::Native: DecimalCast, { fn state(&mut self) -> Result> { self.sum_accumulator.state() @@ -72,78 +85,43 @@ impl Accumulator } fn evaluate(&mut self) -> Result { - if self.sum_accumulator.distinct_count() == 0 { - return ScalarValue::new_primitive::( - None, - &T::TYPE_CONSTRUCTOR(self.target_precision, self.target_scale), - ); + let out_type = I::TYPE_CONSTRUCTOR(self.target_precision, self.target_scale); + let count = self.sum_accumulator.distinct_count(); + if count == 0 { + return ScalarValue::new_primitive::(None, &out_type); } - let sum_scalar = self.sum_accumulator.evaluate()?; - - match sum_scalar { - ScalarValue::Decimal32(Some(sum), _, _) => { - let decimal_averager = DecimalAverager::::try_new( - self.sum_scale, - self.target_precision, - self.target_scale, - )?; - let avg = decimal_averager - .avg(sum, self.sum_accumulator.distinct_count() as i32)?; - Ok(ScalarValue::Decimal32( - Some(avg), - self.target_precision, - self.target_scale, - )) - } - ScalarValue::Decimal64(Some(sum), _, _) => { - let decimal_averager = DecimalAverager::::try_new( - self.sum_scale, - self.target_precision, - self.target_scale, - )?; - let avg = decimal_averager - .avg(sum, self.sum_accumulator.distinct_count() as i64)?; - Ok(ScalarValue::Decimal64( - Some(avg), - self.target_precision, - self.target_scale, - )) - } - ScalarValue::Decimal128(Some(sum), _, _) => { - let decimal_averager = DecimalAverager::::try_new( - self.sum_scale, - self.target_precision, - self.target_scale, - )?; - let avg = decimal_averager - .avg(sum, self.sum_accumulator.distinct_count() as i128)?; - Ok(ScalarValue::Decimal128( - Some(avg), - self.target_precision, - self.target_scale, - )) - } - ScalarValue::Decimal256(Some(sum), _, _) => { - let decimal_averager = DecimalAverager::::try_new( - self.sum_scale, - self.target_precision, - self.target_scale, - )?; - // `distinct_count` returns `u64`, but `avg` expects `i256` - // first convert `u64` to `i128`, then convert `i128` to `i256` to avoid overflow - let distinct_cnt: i128 = self.sum_accumulator.distinct_count() as i128; - let count: i256 = i256::from_i128(distinct_cnt); - let avg = decimal_averager.avg(sum, count)?; - Ok(ScalarValue::Decimal256( - Some(avg), - self.target_precision, - self.target_scale, - )) - } - - _ => unreachable!("Unsupported decimal type: {:?}", sum_scalar), + // Sum the distinct input values in the wider `S` so the total cannot + // overflow the input's native width (mirrors the non-distinct path). + let mut sum = S::Native::usize_as(0); + for value in self.sum_accumulator.distinct_values() { + sum = sum.add_wrapping(value.into()); } + + let Some(count) = S::Native::from_usize(count) else { + return exec_err!( + "Arithmetic overflow in avg: the distinct count {count} cannot \ + be represented in the sum type" + ); + }; + + let averager = DecimalAverager::::try_new( + self.sum_scale, + self.target_precision, + self.target_scale, + )?; + // Narrowing the average back to the (never wider) output type cannot + // fail in practice: `DecimalAverager::avg` validates the average + // against the output precision, whose bound fits the output's native + // type by construction + let avg = + I::Native::from_decimal(averager.avg(sum, count)?).ok_or_else(|| { + exec_datafusion_err!( + "Arithmetic overflow in avg: the computed average does not fit \ + the output type" + ) + })?; + ScalarValue::new_primitive::(Some(avg), &out_type) } fn size(&self) -> usize { @@ -160,6 +138,9 @@ mod tests { use arrow::array::{ Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, }; + use arrow::datatypes::{ + Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, i256, + }; use std::sync::Arc; #[test] @@ -279,4 +260,94 @@ mod tests { Ok(()) } + + // The overflow regression tests below use odd-count ranges symmetric + // around a center value, so the exact sum is `count * center` and the + // average is exactly `center`. + + #[test] + fn test_decimal32_distinct_avg_widens_to_decimal64() -> Result<()> { + // 42951 distinct values centered on 50000: + // sum = 42951 * 50000 = 2,147,550,000 > i32::MAX + let array = Decimal32Array::from_iter_values(28525..=71475) + .with_precision_and_scale(5, 0)?; + + let mut accumulator = DecimalDistinctAvgAccumulator::< + Decimal32Type, + Decimal64Type, + >::with_decimal_params(0, 9, 4); + accumulator.update_batch(&[Arc::new(array)])?; + + assert_eq!( + accumulator.evaluate()?, + ScalarValue::Decimal32(Some(500_000_000), 9, 4) + ); + + Ok(()) + } + + #[test] + fn test_decimal32_distinct_avg_widens_to_decimal128() -> Result<()> { + // 21477 distinct values centered on 99999: + // sum = 21477 * 99999 = 2,147,678,523 > i32::MAX + let array = Decimal32Array::from_iter_values(89261..=110737) + .with_precision_and_scale(9, 0)?; + + let mut accumulator = DecimalDistinctAvgAccumulator::< + Decimal32Type, + Decimal128Type, + >::with_decimal_params(0, 9, 4); + accumulator.update_batch(&[Arc::new(array)])?; + + assert_eq!( + accumulator.evaluate()?, + ScalarValue::Decimal32(Some(999_990_000), 9, 4) + ); + + Ok(()) + } + + #[test] + fn test_decimal64_distinct_avg_widens_to_decimal128() -> Result<()> { + // 92235 distinct values centered on 10^14 - 1: + // sum = 92235 * (10^14 - 1) ~= 9.22e18 > i64::MAX + let center: i64 = 100_000_000_000_000 - 1; + let array = Decimal64Array::from_iter_values(center - 46117..=center + 46117) + .with_precision_and_scale(18, 0)?; + + let mut accumulator = DecimalDistinctAvgAccumulator::< + Decimal64Type, + Decimal128Type, + >::with_decimal_params(0, 18, 4); + accumulator.update_batch(&[Arc::new(array)])?; + + assert_eq!( + accumulator.evaluate()?, + ScalarValue::Decimal64(Some(999_999_999_999_990_000), 18, 4) + ); + + Ok(()) + } + + #[test] + fn test_decimal128_distinct_avg_widens_to_decimal256() -> Result<()> { + // 21477 distinct values ending at 10^34 - 1, centered on 10^34 - 10739: + // sum = 21477 * (10^34 - 10739) ~= 2.15e38 > i128::MAX + let center: i128 = 10_i128.pow(34) - 10739; + let array = Decimal128Array::from_iter_values(center - 10738..=center + 10738) + .with_precision_and_scale(34, 0)?; + + let mut accumulator = DecimalDistinctAvgAccumulator::< + Decimal128Type, + Decimal256Type, + >::with_decimal_params(0, 38, 4); + accumulator.update_batch(&[Arc::new(array)])?; + + assert_eq!( + accumulator.evaluate()?, + ScalarValue::Decimal128(Some(center * 10_000), 38, 4) + ); + + Ok(()) + } } diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs index 83cc5cded8361..bb706aa614dbc 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct.rs @@ -28,5 +28,6 @@ pub use native::Bitmap65536DistinctCountAccumulator; pub use native::Bitmap65536DistinctCountAccumulatorI16; pub use native::BoolArray256DistinctCountAccumulator; pub use native::BoolArray256DistinctCountAccumulatorI8; +pub use native::BooleanDistinctCountAccumulator; pub use native::FloatDistinctCountAccumulator; pub use native::PrimitiveDistinctCountAccumulator; diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs index d370d59c90012..10aa21c3acad2 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/groups.rs @@ -16,7 +16,8 @@ // under the License. use arrow::array::{ - ArrayRef, AsArray, BooleanArray, Int64Array, ListArray, PrimitiveArray, + Array, ArrayRef, AsArray, BooleanArray, Int64Array, ListArray, ListBuilder, + PrimitiveArray, PrimitiveBuilder, }; use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ArrowPrimitiveType, Field}; @@ -160,7 +161,6 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> datafusion_common::Result<()> { debug_assert_eq!(values.len(), 1); @@ -183,9 +183,122 @@ where Ok(()) } + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> datafusion_common::Result> { + debug_assert_eq!(values.len(), 1); + let arr = values[0].as_primitive::(); + + let values_builder = PrimitiveBuilder::::with_capacity(arr.len()); + let mut builder = ListBuilder::new(values_builder) + .with_field(Arc::new(Field::new_list_field(T::DATA_TYPE, true))); + + for row in 0..arr.len() { + let included = arr.is_valid(row) + && opt_filter + .is_none_or(|filter| filter.is_valid(row) && filter.value(row)); + if included { + builder.values().append_value(arr.value(row)); + } + builder.append(true); + } + + Ok(vec![Arc::new(builder.finish())]) + } fn size(&self) -> usize { size_of::() + self.seen.capacity() * (size_of::<(usize, T::Native)>() + size_of::()) + self.counts.capacity() * size_of::() } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::datatypes::Int32Type; + use datafusion_common::Result; + + #[test] + fn convert_to_state_roundtrips_through_merge() -> Result<()> { + let values = Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(2), + None, + Some(3), + Some(4), + Some(5), + Some(5), + ])) as ArrayRef; + let filter = BooleanArray::from(vec![ + Some(true), + Some(true), + Some(true), + Some(true), + None, + Some(true), + Some(true), + Some(true), + ]); + let group_indices = vec![0usize, 1, 0, 1, 0, 0, 0, 0]; + + let mut direct = PrimitiveDistinctCountGroupsAccumulator::::new(); + direct.update_batch( + std::slice::from_ref(&values), + &group_indices, + Some(&filter), + 2, + )?; + let direct = direct.evaluate(EmitTo::All)?; + + let converter = PrimitiveDistinctCountGroupsAccumulator::::new(); + let state = + converter.convert_to_state(std::slice::from_ref(&values), Some(&filter))?; + assert_eq!(state[0].null_count(), 0); + let mut merged = PrimitiveDistinctCountGroupsAccumulator::::new(); + merged.merge_batch(&state, &group_indices, 2)?; + let merged = merged.evaluate(EmitTo::All)?; + + assert_eq!( + direct.as_any().downcast_ref::().unwrap(), + merged.as_any().downcast_ref::().unwrap() + ); + Ok(()) + } + + #[test] + fn convert_to_state_preserves_empty_and_filtered_rows() -> Result<()> { + let converter = PrimitiveDistinctCountGroupsAccumulator::::new(); + let empty_values = + Arc::new(Int32Array::from(Vec::>::new())) as ArrayRef; + let state = + converter.convert_to_state(std::slice::from_ref(&empty_values), None)?; + assert_eq!(state[0].len(), 0); + assert_eq!(state[0].null_count(), 0); + + let values = Arc::new(Int32Array::from(vec![Some(1), Some(2), None])) as ArrayRef; + let filter = BooleanArray::from(vec![Some(false), None, Some(false)]); + let group_indices = vec![0usize, 1, 0]; + + let state = + converter.convert_to_state(std::slice::from_ref(&values), Some(&filter))?; + assert_eq!(state[0].len(), values.len()); + assert_eq!(state[0].null_count(), 0); + let list_state = state[0].as_list::(); + for row in 0..list_state.len() { + assert_eq!(list_state.value_length(row), 0); + } + + let mut merged = PrimitiveDistinctCountGroupsAccumulator::::new(); + merged.merge_batch(&state, &group_indices, 2)?; + let result = merged.evaluate(EmitTo::All)?; + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &Int64Array::from(vec![0, 0]) + ); + Ok(()) + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs index fb9cfb379a26e..00c1a47b9eafb 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/count_distinct/native.rs @@ -26,14 +26,16 @@ use std::hash::Hash; use std::mem::size_of_val; use std::sync::Arc; +use arrow::array::Array; use arrow::array::ArrayRef; +use arrow::array::BooleanArray; use arrow::array::PrimitiveArray; use arrow::array::types::ArrowPrimitiveType; use arrow::datatypes::DataType; use datafusion_common::hash_utils::RandomState; use datafusion_common::ScalarValue; -use datafusion_common::cast::{as_list_array, as_primitive_array}; +use datafusion_common::cast::{as_boolean_array, as_list_array, as_primitive_array}; use datafusion_common::utils::SingleRowListArrayBuilder; use datafusion_common::utils::memory::estimate_memory_size; use datafusion_expr_common::accumulator::Accumulator; @@ -85,11 +87,15 @@ where } let arr = as_primitive_array::(&values[0])?; - arr.iter().for_each(|value| { - if let Some(value) = value { + if arr.null_count() == 0 { + // Fast path: no nulls, so skip the per-element validity check and + // insert directly from the values buffer (mirrors `merge_batch`). + self.values.extend(arr.values().iter().copied()); + } else { + arr.iter().flatten().for_each(|value| { self.values.insert(value); - } - }); + }); + } Ok(()) } @@ -518,3 +524,140 @@ impl Accumulator for Bitmap65536DistinctCountAccumulatorI16 { size_of_val(self) + 8192 } } + +/// Optimized COUNT DISTINCT accumulator for `Boolean` using two flags. +/// +/// Tracks whether `false` and `true` have been observed; nulls are skipped. +/// Result is always 0, 1, or 2. +#[derive(Debug)] +pub struct BooleanDistinctCountAccumulator { + has_seen_false: bool, + has_seen_true: bool, +} + +impl BooleanDistinctCountAccumulator { + pub fn new() -> Self { + Self { + has_seen_false: false, + has_seen_true: false, + } + } + + #[inline] + fn seen_both(&self) -> bool { + self.has_seen_false && self.has_seen_true + } + + #[inline] + fn count(&self) -> i64 { + (self.has_seen_false as u8 + self.has_seen_true as u8) as i64 + } + + /// Update flags from a `BooleanArray`, short-circuiting per-flag once set. + #[inline] + fn observe(&mut self, arr: &BooleanArray) { + if !self.has_seen_false && arr.has_false() { + self.has_seen_false = true; + } + if !self.has_seen_true && arr.has_true() { + self.has_seen_true = true; + } + } +} + +impl Default for BooleanDistinctCountAccumulator { + fn default() -> Self { + Self::new() + } +} + +impl Accumulator for BooleanDistinctCountAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> datafusion_common::Result<()> { + if values.is_empty() || self.seen_both() { + return Ok(()); + } + + let arr = as_boolean_array(&values[0])?; + self.observe(arr); + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> datafusion_common::Result<()> { + if states.is_empty() || self.seen_both() { + return Ok(()); + } + + let arr = as_list_array(&states[0])?; + arr.iter().try_for_each(|maybe_list| { + if self.seen_both() { + return Ok(()); + } + if let Some(list) = maybe_list { + self.observe(as_boolean_array(&list)?); + }; + Ok(()) + }) + } + + fn state(&mut self) -> datafusion_common::Result> { + let mut values: Vec = Vec::with_capacity(2); + if self.has_seen_false { + values.push(false); + } + if self.has_seen_true { + values.push(true); + } + + let arr = Arc::new(BooleanArray::from(values)); + Ok(vec![ + SingleRowListArrayBuilder::new(arr).build_list_scalar(), + ]) + } + + fn evaluate(&mut self) -> datafusion_common::Result { + Ok(ScalarValue::Int64(Some(self.count()))) + } + + fn size(&self) -> usize { + size_of_val(self) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::datatypes::Int64Type; + + #[test] + fn update_batch_null_free_fast_path_agrees_with_general_path() { + // The null-free fast path must produce the same distinct set as the + // general (validity-checking) path. + let dense: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3, 2, 1])); + let sparse: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + None, + Some(3), + Some(2), + Some(1), + ])); + + let mut dense_acc = + PrimitiveDistinctCountAccumulator::::new(&DataType::Int64); + dense_acc + .update_batch(std::slice::from_ref(&dense)) + .unwrap(); + + let mut sparse_acc = + PrimitiveDistinctCountAccumulator::::new(&DataType::Int64); + sparse_acc + .update_batch(std::slice::from_ref(&sparse)) + .unwrap(); + + // Both should count the 3 distinct non-null values {1, 2, 3}. + assert_eq!(dense_acc.evaluate().unwrap(), ScalarValue::Int64(Some(3))); + assert_eq!(sparse_acc.evaluate().unwrap(), ScalarValue::Int64(Some(3))); + } +} diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs index ad2a21bb4733c..b5610419166df 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator.rs @@ -375,13 +375,12 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { self.invoke_per_accumulator( values, group_indices, - opt_filter, + None, total_num_groups, |accumulator, values_to_accumulate| { accumulator.merge_batch(values_to_accumulate)?; @@ -442,10 +441,6 @@ impl GroupsAccumulator for GroupsAccumulatorAdapter { Ok(arrays) } - - fn supports_convert_to_state(&self) -> bool { - true - } } /// Extension trait for [`Vec`] to account for allocations. diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs index d1d8924a2c3e8..77bb7598e2747 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/bool_op.rs @@ -132,11 +132,10 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // update / merge are the same - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } fn size(&self) -> usize { @@ -157,8 +156,4 @@ where Ok(vec![Arc::new(values_filtered)]) } - - fn supports_convert_to_state(&self) -> bool { - true - } } diff --git a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs index a81b89e1e46f1..c5d74978664c9 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/groups_accumulator/prim_op.rs @@ -131,11 +131,10 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // update / merge are the same - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } /// Converts an input batch directly to a state batch @@ -190,11 +189,6 @@ where Ok(vec![Arc::new(state_values)]) } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn size(&self) -> usize { self.values.capacity() * size_of::() + self.null_state.size() } diff --git a/datafusion/functions-aggregate-common/src/aggregate/sum_distinct/numeric.rs b/datafusion/functions-aggregate-common/src/aggregate/sum_distinct/numeric.rs index e5a23597c44ad..2119c06b48aaf 100644 --- a/datafusion/functions-aggregate-common/src/aggregate/sum_distinct/numeric.rs +++ b/datafusion/functions-aggregate-common/src/aggregate/sum_distinct/numeric.rs @@ -50,6 +50,13 @@ impl DistinctSumAccumulator { pub fn distinct_count(&self) -> usize { self.values.values.len() } + + /// Iterates the distinct values collected so far. `AVG(DISTINCT)` re-sums + /// them in a wider type instead of using [`Self::evaluate`]'s input-typed + /// sum. + pub(crate) fn distinct_values(&self) -> impl Iterator + '_ { + self.values.values.iter().map(|v| v.0) + } } impl Accumulator for DistinctSumAccumulator { diff --git a/datafusion/functions-aggregate-common/src/tdigest.rs b/datafusion/functions-aggregate-common/src/tdigest.rs index a7450f0eb52e9..8db7d0bc8a541 100644 --- a/datafusion/functions-aggregate-common/src/tdigest.rs +++ b/datafusion/functions-aggregate-common/src/tdigest.rs @@ -31,8 +31,8 @@ use arrow::datatypes::DataType; use arrow::datatypes::Float64Type; -use datafusion_common::ScalarValue; use datafusion_common::cast::as_primitive_array; +use datafusion_common::{DataFusionError, ScalarValue, exec_err}; use std::cmp::Ordering; use std::mem::{size_of, size_of_val}; @@ -148,6 +148,23 @@ impl TDigest { self.max_size } + /// The sum of all values ingested into this digest. + #[inline] + pub fn sum(&self) -> f64 { + self.sum + } + + /// The centroids that make up this digest, ordered by mean. + /// + /// Together with the [`Self::sum()`], [`Self::max_size()`], + /// [`Self::count()`], [`Self::max()`], and [`Self::min()`] accessors this + /// exposes the full serialized state of the digest without packing it into + /// a [`ScalarValue`] list. See [`Self::try_from_parts()`] for the inverse. + #[inline] + pub fn centroids(&self) -> &[Centroid] { + &self.centroids + } + /// Size in bytes including `Self`. pub fn size(&self) -> usize { size_of_val(self) + (size_of::() * self.centroids.capacity()) @@ -611,6 +628,74 @@ impl TDigest { centroids, } } + + /// Construct a [`TDigest`] directly from its constituent parts, validating + /// the inputs. + /// + /// Together with the [`Self::centroids()`], [`Self::sum()`], + /// [`Self::max_size()`], [`Self::count()`], [`Self::max()`], and + /// [`Self::min()`] accessors, this allows a digest to be serialized into and + /// restored from a caller's own format without round-tripping through a + /// [`ScalarValue`] list (the non-Arrow counterpart to + /// [`Self::from_scalar_state()`]). + /// + /// Unlike [`Self::from_scalar_state()`], this validates its inputs, returning + /// an error rather than a silently wrong digest when handed corrupt state. + /// Callers who trust their data can `unwrap()`. + /// + /// # Errors + /// + /// Returns an error if: + /// - `min` and `max` are both finite but `max < min`; + /// - the `centroids` are not sorted in non-decreasing order by mean (the + /// order produced by [`Self::centroids()`]); or + /// - any centroid weight is not finite and strictly positive + /// ([`Self::estimate_quantile()`] divides by a centroid's weight, so a + /// zero, negative, or non-finite weight yields silently wrong results). + pub fn try_from_parts( + max_size: usize, + sum: f64, + count: f64, + max: f64, + min: f64, + centroids: Vec, + ) -> Result { + if min.is_finite() && max.is_finite() && max.total_cmp(&min).is_lt() { + return exec_err!( + "invalid TDigest state: max ({max}) is less than min ({min})" + ); + } + + for pair in centroids.windows(2) { + if pair[0].cmp_mean(&pair[1]).is_gt() { + return exec_err!( + "invalid TDigest state: centroids must be sorted by mean, \ + but {} precedes {}", + pair[0].mean(), + pair[1].mean() + ); + } + } + + for centroid in ¢roids { + if !(centroid.weight().is_finite() && centroid.weight() > 0.0) { + return exec_err!( + "invalid TDigest state: centroid weight must be finite and \ + positive, got {}", + centroid.weight() + ); + } + } + + Ok(Self { + max_size, + sum, + count, + max, + min, + centroids, + }) + } } #[cfg(debug_assertions)] @@ -760,4 +845,147 @@ mod tests { // The result should be approximately equal to the input value assert!((result - 15.699999988079073).abs() < 1e-10); } + + // A representative set of digests covering the empty, single-value and + // heavily-compressed cases, used to exercise the `try_from_parts`/accessor + // external-state contract. + fn sample_digests() -> Vec { + vec![ + // Empty: no values ingested, so max/min are NaN and centroids empty. + TDigest::new(100), + // A single value. + TDigest::new(100).merge_unsorted_f64(vec![42.0]), + // Many values, forcing compression down to `max_size` centroids. + TDigest::new(100).merge_unsorted_f64((1..=10_000).map(f64::from).collect()), + // A different shape and `max_size`. + TDigest::new(50) + .merge_unsorted_f64((1..=5_000).map(|v| f64::from(v).sqrt()).collect()), + ] + } + + const QUANTILE_GRID: [f64; 9] = [0.0, 0.01, 0.1, 0.25, 0.5, 0.75, 0.9, 0.99, 1.0]; + + // Rebuild a digest purely from its public accessors via `try_from_parts`. + fn rebuild_via_parts(t: &TDigest) -> TDigest { + TDigest::try_from_parts( + t.max_size(), + t.sum(), + t.count(), + t.max(), + t.min(), + t.centroids().to_vec(), + ) + .expect("digest built from real accessors is valid") + } + + #[test] + fn test_from_parts_roundtrip() { + for t in sample_digests() { + let rebuilt = rebuild_via_parts(&t); + + // The serialized state must be identical. `to_scalar_state()` + // compares `Float64` by bit pattern, so this also holds for the + // empty digest whose max/min are NaN. + assert_eq!(rebuilt.to_scalar_state(), t.to_scalar_state()); + + // Quantile estimates must be bitwise-equal across the grid. + for q in QUANTILE_GRID { + assert_eq!( + rebuilt.estimate_quantile(q).to_bits(), + t.estimate_quantile(q).to_bits(), + "quantile {q} diverged after try_from_parts roundtrip" + ); + } + } + } + + #[test] + fn test_from_parts_equals_original() { + // For digests without NaN fields, use the strongest available equality: + // the derived `PartialEq` on `TDigest`. (The empty digest is excluded + // because NaN != NaN under the derived comparison; it is covered by + // `test_from_parts_roundtrip` via `to_scalar_state`.) + for t in sample_digests().into_iter().filter(|t| t.count() > 0.0) { + let rebuilt = rebuild_via_parts(&t); + assert_eq!(rebuilt, t); + } + } + + #[test] + fn test_accessors_agree_with_scalar_state() { + for t in sample_digests() { + let state = t.to_scalar_state(); + + // `sum()` matches the sum field packed into the scalar state. + assert_eq!(ScalarValue::Float64(Some(t.sum())), state[1]); + + // `centroids()` matches the flat mean/weight pairs in the list. + let flattened: Vec = t + .centroids() + .iter() + .flat_map(|c| [c.mean(), c.weight()]) + .map(|v| ScalarValue::Float64(Some(v))) + .collect(); + let expected = ScalarValue::new_list_nullable(&flattened, &DataType::Float64); + assert_eq!(ScalarValue::List(expected), state[5]); + } + } + + #[test] + fn test_from_parts_rejects_max_less_than_min() { + let err = TDigest::try_from_parts( + 100, + 3.0, + 2.0, + 1.0, // max + 5.0, // min > max + vec![Centroid::new(1.0, 1.0), Centroid::new(5.0, 1.0)], + ) + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("max") && msg.contains("less than min"), + "unexpected error message: {msg}" + ); + } + + #[test] + fn test_from_parts_rejects_unsorted_centroids() { + let err = TDigest::try_from_parts( + 100, + 6.0, + 3.0, + 3.0, + 1.0, + // Means out of order: 3.0 precedes 1.0. + vec![Centroid::new(3.0, 1.0), Centroid::new(1.0, 1.0)], + ) + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("sorted by mean"), + "unexpected error message: {msg}" + ); + } + + #[test] + fn test_from_parts_rejects_non_positive_weight() { + // A zero weight would divide-by-zero inside `estimate_quantile`. + for bad_weight in [0.0, -1.0, f64::NAN, f64::INFINITY] { + let err = TDigest::try_from_parts( + 100, + 1.0, + bad_weight, + 1.0, + 1.0, + vec![Centroid::new(1.0, bad_weight)], + ) + .unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("weight must be finite and"), + "weight {bad_weight}: unexpected error message: {msg}" + ); + } + } } diff --git a/datafusion/functions-aggregate/Cargo.toml b/datafusion/functions-aggregate/Cargo.toml index 778e6a24bf00e..5abea16e2cc81 100644 --- a/datafusion/functions-aggregate/Cargo.toml +++ b/datafusion/functions-aggregate/Cargo.toml @@ -50,8 +50,8 @@ datafusion-functions-aggregate-common = { workspace = true } datafusion-macros = { workspace = true } datafusion-physical-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } -foldhash = "0.2" half = { workspace = true } +hashbrown = { workspace = true } log = { workspace = true } num-traits = { workspace = true } @@ -95,3 +95,14 @@ harness = false [[bench]] name = "percentile_cont" harness = false + +[[bench]] +name = "sliding_max" +harness = false + +[[bench]] +name = "variance" +harness = false + +[features] +force_hash_collisions = ["datafusion-common/force_hash_collisions"] diff --git a/datafusion/functions-aggregate/benches/approx_distinct.rs b/datafusion/functions-aggregate/benches/approx_distinct.rs index cc85c2163c180..2ab783d9acf05 100644 --- a/datafusion/functions-aggregate/benches/approx_distinct.rs +++ b/datafusion/functions-aggregate/benches/approx_distinct.rs @@ -15,17 +15,25 @@ // specific language governing permissions and limitations // under the License. +use std::hint::black_box; use std::sync::Arc; use arrow::array::{ - ArrayRef, Int8Array, Int16Array, Int64Array, StringArray, StringViewArray, - UInt8Array, UInt16Array, + ArrayRef, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, + Int8Array, Int16Array, Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, + IntervalYearMonthArray, StringArray, StringViewArray, UInt8Array, UInt16Array, +}; +use arrow::datatypes::{ + DataType, Field, IntervalDayTime, IntervalMonthDayNano, IntervalUnit, Schema, i256, }; -use arrow::datatypes::{DataType, Field, Schema}; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_expr::function::AccumulatorArgs; -use datafusion_expr::{Accumulator, AggregateUDFImpl}; +use datafusion_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, +}; use datafusion_functions_aggregate::approx_distinct::ApproxDistinct; +use datafusion_physical_expr::GroupsAccumulatorAdapter; +use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::col; use rand::rngs::StdRng; use rand::{Rng, SeedableRng}; @@ -34,6 +42,17 @@ const BATCH_SIZE: usize = 8192; const SHORT_STRING_LENGTH: usize = 8; const LONG_STRING_LENGTH: usize = 20; +// Grouped (high-cardinality `GROUP BY`) benchmark parameters. +const N_GROUPS: usize = 50_000; +const AVG_ROWS_PER_GROUP: usize = 8; +const STRING_POOL_SIZE: usize = 100_000; + +const DECIMAL32_PRECISION: u8 = 9; +const DECIMAL64_PRECISION: u8 = 18; +const DECIMAL128_PRECISION: u8 = 10; +const DECIMAL256_PRECISION: u8 = 40; +const DECIMAL_SCALE: i8 = 2; + fn prepare_accumulator(data_type: DataType) -> Box { let schema = Arc::new(Schema::new(vec![Field::new("f", data_type, true)])); let expr = col("f", &schema).unwrap(); @@ -51,6 +70,52 @@ fn prepare_accumulator(data_type: DataType) -> Box { ApproxDistinct::new().accumulator(accumulator_args).unwrap() } +/// Creates a `Decimal32Array` from a pool of `n_distinct` values. +fn create_decimal32_array(n_distinct: usize) -> Decimal32Array { + let mut rng = StdRng::seed_from_u64(42); + let pool: Vec = (0..n_distinct).map(|i| i as i32 * 50).collect(); + (0..BATCH_SIZE) + .map(|_| Some(pool[rng.random_range(0..pool.len())])) + .collect::() + .with_precision_and_scale(DECIMAL32_PRECISION, DECIMAL_SCALE) + .unwrap() +} + +/// Creates a `Decimal64Array` from a pool of `n_distinct` values. +fn create_decimal64_array(n_distinct: usize) -> Decimal64Array { + let mut rng = StdRng::seed_from_u64(42); + let pool: Vec = (0..n_distinct).map(|i| i as i64 * 50).collect(); + (0..BATCH_SIZE) + .map(|_| Some(pool[rng.random_range(0..pool.len())])) + .collect::() + .with_precision_and_scale(DECIMAL64_PRECISION, DECIMAL_SCALE) + .unwrap() +} + +/// Creates a `Decimal128Array` from a pool of `n_distinct` values. +fn create_decimal128_array(n_distinct: usize) -> Decimal128Array { + let mut rng = StdRng::seed_from_u64(42); + let pool: Vec = (0..n_distinct).map(|i| i as i128 * 50).collect(); + (0..BATCH_SIZE) + .map(|_| Some(pool[rng.random_range(0..pool.len())])) + .collect::() + .with_precision_and_scale(DECIMAL128_PRECISION, DECIMAL_SCALE) + .unwrap() +} + +/// Creates a `Decimal256Array` from a pool of `n_distinct` values. +fn create_decimal256_array(n_distinct: usize) -> Decimal256Array { + let mut rng = StdRng::seed_from_u64(42); + let pool: Vec = (0..n_distinct) + .map(|i| i256::from_i128(i as i128 * 50)) + .collect(); + (0..BATCH_SIZE) + .map(|_| Some(pool[rng.random_range(0..pool.len())])) + .collect::() + .with_precision_and_scale(DECIMAL256_PRECISION, DECIMAL_SCALE) + .unwrap() +} + /// Creates an Int64Array where values are drawn from `0..n_distinct`. fn create_i64_array(n_distinct: usize) -> Int64Array { let mut rng = StdRng::seed_from_u64(42); @@ -91,6 +156,38 @@ fn create_i16_array(n_distinct: usize) -> Int16Array { .collect() } +/// Creates an `IntervalYearMonthArray` where values are drawn from `0..n_distinct`. +fn create_interval_year_month_array(n_distinct: usize) -> IntervalYearMonthArray { + let mut rng = StdRng::seed_from_u64(42); + (0..BATCH_SIZE) + .map(|_| Some(rng.random_range(0..n_distinct as i32))) + .collect() +} + +/// Creates an `IntervalDayTimeArray` where values are drawn from a pool of +/// `n_distinct` values. +fn create_interval_day_time_array(n_distinct: usize) -> IntervalDayTimeArray { + let mut rng = StdRng::seed_from_u64(42); + let pool: Vec = (0..n_distinct) + .map(|i| IntervalDayTime::new(i as i32, i as i32 * 100)) + .collect(); + (0..BATCH_SIZE) + .map(|_| Some(pool[rng.random_range(0..pool.len())])) + .collect() +} + +/// Creates an `IntervalMonthDayNanoArray` where values are drawn from a pool of +/// `n_distinct` values. +fn create_interval_month_day_nano_array(n_distinct: usize) -> IntervalMonthDayNanoArray { + let mut rng = StdRng::seed_from_u64(42); + let pool: Vec = (0..n_distinct) + .map(|i| IntervalMonthDayNano::new(i as i32, i as i32, i as i64 * 1_000)) + .collect(); + (0..BATCH_SIZE) + .map(|_| Some(pool[rng.random_range(0..pool.len())])) + .collect() +} + /// Creates a pool of `n_distinct` random strings of the given length. fn create_string_pool(n_distinct: usize, string_length: usize) -> Vec { let mut rng = StdRng::seed_from_u64(42); @@ -214,7 +311,281 @@ fn approx_distinct_benchmark(c: &mut Criterion) { .unwrap() }) }); + + // Decimal32 + let values = Arc::new(create_decimal32_array(200)) as ArrayRef; + c.bench_function("approx_distinct decimal32", |b| { + b.iter(|| { + let mut accumulator = prepare_accumulator(DataType::Decimal32( + DECIMAL32_PRECISION, + DECIMAL_SCALE, + )); + accumulator + .update_batch(std::slice::from_ref(&values)) + .unwrap() + }) + }); + + // Decimal64 + let values = Arc::new(create_decimal64_array(200)) as ArrayRef; + c.bench_function("approx_distinct decimal64", |b| { + b.iter(|| { + let mut accumulator = prepare_accumulator(DataType::Decimal64( + DECIMAL64_PRECISION, + DECIMAL_SCALE, + )); + accumulator + .update_batch(std::slice::from_ref(&values)) + .unwrap() + }) + }); + + // Decimal128 + let values = Arc::new(create_decimal128_array(200)) as ArrayRef; + c.bench_function("approx_distinct decimal128", |b| { + b.iter(|| { + let mut accumulator = prepare_accumulator(DataType::Decimal128( + DECIMAL128_PRECISION, + DECIMAL_SCALE, + )); + accumulator + .update_batch(std::slice::from_ref(&values)) + .unwrap() + }) + }); + + // Decimal256 + let values = Arc::new(create_decimal256_array(200)) as ArrayRef; + c.bench_function("approx_distinct decimal256", |b| { + b.iter(|| { + let mut accumulator = prepare_accumulator(DataType::Decimal256( + DECIMAL256_PRECISION, + DECIMAL_SCALE, + )); + accumulator + .update_batch(std::slice::from_ref(&values)) + .unwrap() + }) + }); + + // Interval benchmarks + for pct in [80, 99] { + let n_distinct = BATCH_SIZE * pct / 100; + + // IntervalYearMonth + let values = Arc::new(create_interval_year_month_array(n_distinct)) as ArrayRef; + c.bench_function( + &format!("approx_distinct interval year_month {pct}% distinct"), + |b| { + b.iter(|| { + let mut accumulator = + prepare_accumulator(DataType::Interval(IntervalUnit::YearMonth)); + accumulator + .update_batch(std::slice::from_ref(&values)) + .unwrap() + }) + }, + ); + + // IntervalDayTime + let values = Arc::new(create_interval_day_time_array(n_distinct)) as ArrayRef; + c.bench_function( + &format!("approx_distinct interval day_time {pct}% distinct"), + |b| { + b.iter(|| { + let mut accumulator = + prepare_accumulator(DataType::Interval(IntervalUnit::DayTime)); + accumulator + .update_batch(std::slice::from_ref(&values)) + .unwrap() + }) + }, + ); + + // IntervalMonthDayNano + let values = + Arc::new(create_interval_month_day_nano_array(n_distinct)) as ArrayRef; + c.bench_function( + &format!("approx_distinct interval month_day_nano {pct}% distinct"), + |b| { + b.iter(|| { + let mut accumulator = prepare_accumulator(DataType::Interval( + IntervalUnit::MonthDayNano, + )); + accumulator + .update_batch(std::slice::from_ref(&values)) + .unwrap() + }) + }, + ); + } +} + +/// Build a `GroupsAccumulator` the same way the aggregate operator does: use the +/// specialized one if the function supports it, otherwise fall back to wrapping +/// the per-group `Accumulator` in a `GroupsAccumulatorAdapter`. +fn prepare_groups_accumulator(data_type: DataType) -> Box { + let schema = Arc::new(Schema::new(vec![Field::new("f", data_type, true)])); + let expr = col("f", &schema).unwrap(); + let udf = Arc::new(AggregateUDF::from(ApproxDistinct::new())); + let agg = Arc::new( + AggregateExprBuilder::new(udf, vec![expr]) + .schema(schema) + .alias("approx_distinct(f)") + .build() + .unwrap(), + ); + + if agg.groups_accumulator_supported() { + agg.create_groups_accumulator().unwrap() + } else { + let agg = Arc::clone(&agg); + let factory = move || agg.create_accumulator(); + Box::new(GroupsAccumulatorAdapter::new(factory)) + } +} + +fn grouped_total_rows() -> usize { + N_GROUPS * AVG_ROWS_PER_GROUP +} + +/// A random group index in `0..N_GROUPS` for each row of a batch. +fn make_group_indices(rng: &mut StdRng) -> Vec { + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..N_GROUPS)) + .collect() +} + +/// Pre-build all input batches `(values, group_indices)` for the grouped run, so +/// the measured loop only times the accumulator, not data generation. +fn build_grouped_batches(data_type: &DataType) -> Vec<(ArrayRef, Vec)> { + let n_batches = grouped_total_rows().div_ceil(BATCH_SIZE); + let mut rng = StdRng::seed_from_u64(7); + let pool = create_string_pool(STRING_POOL_SIZE, SHORT_STRING_LENGTH); + + (0..n_batches) + .map(|_| { + let group_indices = make_group_indices(&mut rng); + let values: ArrayRef = match data_type { + DataType::Int64 => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(rng.random::())) + .collect::(), + ), + DataType::Utf8 => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(pool[rng.random_range(0..pool.len())].as_str())) + .collect::(), + ), + DataType::Utf8View => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(pool[rng.random_range(0..pool.len())].as_str())) + .collect::(), + ), + DataType::Decimal32(p, s) => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(rng.random::())) + .collect::() + .with_precision_and_scale(*p, *s) + .unwrap(), + ), + DataType::Decimal64(p, s) => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(rng.random::())) + .collect::() + .with_precision_and_scale(*p, *s) + .unwrap(), + ), + DataType::Decimal128(p, s) => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(rng.random::() as i128)) + .collect::() + .with_precision_and_scale(*p, *s) + .unwrap(), + ), + DataType::Decimal256(p, s) => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(i256::from_i128(rng.random::() as i128))) + .collect::() + .with_precision_and_scale(*p, *s) + .unwrap(), + ), + DataType::Interval(IntervalUnit::YearMonth) => Arc::new( + (0..BATCH_SIZE) + .map(|_| Some(rng.random::())) + .collect::(), + ), + DataType::Interval(IntervalUnit::DayTime) => Arc::new( + (0..BATCH_SIZE) + .map(|_| { + Some(IntervalDayTime::new( + rng.random::(), + rng.random::(), + )) + }) + .collect::(), + ), + DataType::Interval(IntervalUnit::MonthDayNano) => Arc::new( + (0..BATCH_SIZE) + .map(|_| { + Some(IntervalMonthDayNano::new( + rng.random::(), + rng.random::(), + rng.random::(), + )) + }) + .collect::(), + ), + other => panic!("unsupported grouped bench type: {other}"), + }; + (values, group_indices) + }) + .collect() +} + +/// Benchmark grouped `approx_distinct` over many groups. Each iteration feeds all batches into a +/// fresh accumulator and emits the result for every group. +fn approx_distinct_grouped_benchmark(c: &mut Criterion) { + let mut group = c.benchmark_group("approx_distinct_grouped"); + group.sample_size(10); + + for data_type in [ + DataType::Int64, + DataType::Utf8, + DataType::Utf8View, + DataType::Decimal32(DECIMAL32_PRECISION, DECIMAL_SCALE), + DataType::Decimal64(DECIMAL64_PRECISION, DECIMAL_SCALE), + DataType::Decimal128(DECIMAL128_PRECISION, DECIMAL_SCALE), + DataType::Decimal256(DECIMAL256_PRECISION, DECIMAL_SCALE), + DataType::Interval(IntervalUnit::YearMonth), + DataType::Interval(IntervalUnit::DayTime), + DataType::Interval(IntervalUnit::MonthDayNano), + ] { + let batches = build_grouped_batches(&data_type); + let label = format!("{data_type:?} {N_GROUPS} groups"); + group.bench_function(&label, |b| { + b.iter(|| { + let mut acc = prepare_groups_accumulator(data_type.clone()); + for (values, group_indices) in &batches { + acc.update_batch( + std::slice::from_ref(values), + group_indices, + None, + N_GROUPS, + ) + .unwrap(); + } + black_box(acc.evaluate(EmitTo::All).unwrap()); + }) + }); + } + + group.finish(); } -criterion_group!(benches, approx_distinct_benchmark); +criterion_group!( + benches, + approx_distinct_benchmark, + approx_distinct_grouped_benchmark +); criterion_main!(benches); diff --git a/datafusion/functions-aggregate/benches/array_agg.rs b/datafusion/functions-aggregate/benches/array_agg.rs index b0d8148c3ea65..d7e5a511078a5 100644 --- a/datafusion/functions-aggregate/benches/array_agg.rs +++ b/datafusion/functions-aggregate/benches/array_agg.rs @@ -20,11 +20,14 @@ use std::sync::Arc; use arrow::array::{ Array, ArrayRef, ArrowPrimitiveType, AsArray, ListArray, NullBufferBuilder, + StringArray, }; -use arrow::datatypes::{Field, Int64Type}; +use arrow::datatypes::{DataType, Field, Int64Type}; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_expr::Accumulator; -use datafusion_functions_aggregate::array_agg::ArrayAggAccumulator; +use datafusion_functions_aggregate::array_agg::{ + ArrayAggAccumulator, DistinctArrayAggAccumulator, +}; use arrow::buffer::OffsetBuffer; use arrow::util::bench_util::create_primitive_array; @@ -191,5 +194,101 @@ fn array_agg_benchmark(c: &mut Criterion) { ); } -criterion_group!(benches, array_agg_benchmark); +/// A realistic pool of database names with variable lengths. +const DB_NAMES: &[&str] = &[ + "postgres", + "mysql", + "oracle", + "mssql", + "mongodb", + "redis", + "elasticsearch", + "cassandra", + "dynamodb", + "bigquery", + "snowflake", + "redshift", + "databricks", + "clickhouse", + "duckdb", + "cockroachdb", + "tidb", + "mariadb", + "sqlite", + "neo4j", + "influxdb", + "timescaledb", + "yugabytedb", + "planetscale", + "singlestore", +]; + +/// Low-cardinality: every row is drawn uniformly from `DB_NAMES` (~25 distinct +/// values across 8 192 rows). Exercises the hot duplicate path. +fn create_string_array_low_cardinality(size: usize) -> StringArray { + let mut rng = StdRng::seed_from_u64(42); + StringArray::from_iter_values( + (0..size).map(|_| DB_NAMES[rng.random_range(0..DB_NAMES.len())]), + ) +} + +/// High-cardinality: `db_name_pct` fraction of rows are drawn from `DB_NAMES`; +/// the rest are near-unique random hex strings ("id_XXXXXXXX"). +/// With 8 192 rows and a 32-bit space the collision probability among the +/// random strings is < 1 %, giving ~7 800 distinct values in total. +fn create_string_array_high_cardinality(size: usize, db_name_pct: f32) -> StringArray { + let mut rng = StdRng::seed_from_u64(42); + let strings: Vec = (0..size) + .map(|_| { + if rng.random::() < db_name_pct { + DB_NAMES[rng.random_range(0..DB_NAMES.len())].to_string() + } else { + format!("id_{:08x}", rng.random::()) + } + }) + .collect(); + StringArray::from_iter_values(strings.iter().map(String::as_str)) +} + +fn distinct_update_batch_bench( + c: &mut Criterion, + name: &str, + values: &ArrayRef, + ignore_nulls: bool, +) { + c.bench_function(name, |b| { + b.iter(|| { + DistinctArrayAggAccumulator::try_new(&DataType::Utf8, None, ignore_nulls) + .unwrap() + .update_batch(std::slice::from_ref(values)) + .unwrap() + }) + }); +} + +fn distinct_array_agg_benchmark(c: &mut Criterion) { + // --- Low cardinality: ~25 distinct DB names in 8 192 rows --------------- + // Realistic production scenario: most rows are duplicates, the HashSet + // saturates quickly and the rest of the batch is pure dedup overhead. + let values = Arc::new(create_string_array_low_cardinality(8192)) as ArrayRef; + distinct_update_batch_bench( + c, + "distinct_array_agg utf8 low cardinality (~25 distinct)", + &values, + false, + ); + + // --- High cardinality: ~5 % DB names, ~95 % near-unique random strings -- + // Worst-case scenario: almost every row is a new distinct value, so the + // accumulator pays the full insertion cost for nearly every row. + let values = Arc::new(create_string_array_high_cardinality(8192, 0.05)) as ArrayRef; + distinct_update_batch_bench( + c, + "distinct_array_agg utf8 high cardinality (~7800 distinct, 5% db names)", + &values, + false, + ); +} + +criterion_group!(benches, array_agg_benchmark, distinct_array_agg_benchmark); criterion_main!(benches); diff --git a/datafusion/functions-aggregate/benches/first_last.rs b/datafusion/functions-aggregate/benches/first_last.rs index 1d18e1c7dcd44..235f11ff30f63 100644 --- a/datafusion/functions-aggregate/benches/first_last.rs +++ b/datafusion/functions-aggregate/benches/first_last.rs @@ -15,10 +15,16 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, BooleanArray, Int64Array}; +use arrow::array::{ + Array, ArrayRef, BooleanArray, Int64Array, ListArray, MapArray, StringArray, + StructArray, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::compute::SortOptions; -use arrow::datatypes::{DataType, Field, Int64Type, Schema}; -use arrow::util::bench_util::{create_boolean_array, create_primitive_array}; +use arrow::datatypes::{DataType, Field, Fields, Float64Type, Int64Type, Schema}; +use arrow::util::bench_util::{ + create_boolean_array, create_primitive_array, create_string_array_with_len, +}; use datafusion_common::instant::Instant; use std::hint::black_box; use std::sync::Arc; @@ -29,14 +35,21 @@ use datafusion_expr::{ use datafusion_functions_aggregate::first_last::{ FirstValue, LastValue, TrivialFirstValueAccumulator, TrivialLastValueAccumulator, }; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::GroupsAccumulatorAdapter; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::col; use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; -fn prepare_groups_accumulator(is_first: bool) -> Box { +/// Build a `GroupsAccumulator` for an arbitrary value type, so the nested-type +/// (`Struct` / `List`) fast paths added for `first_value` / `last_value` can be +/// exercised with the same harness as the primitive ones. +fn prepare_typed_groups_accumulator( + is_first: bool, + value_type: DataType, +) -> Box { let schema = Arc::new(Schema::new(vec![ - Field::new("value", DataType::Int64, true), + Field::new("value", value_type.clone(), true), Field::new("ord", DataType::Int64, true), ])); @@ -46,11 +59,12 @@ fn prepare_groups_accumulator(is_first: bool) -> Box { options: SortOptions::default(), }; - let value_field: Arc = Field::new("value", DataType::Int64, true).into(); - let accumulator_args = AccumulatorArgs { + let value_field: Arc = Field::new("value", value_type.clone(), true).into(); + let value_expr = col("value", &schema).unwrap(); + let make_args = || AccumulatorArgs { return_field: Arc::clone(&value_field), schema: &schema, - expr_fields: &[value_field], + expr_fields: std::slice::from_ref(&value_field), ignore_nulls: false, order_bys: std::slice::from_ref(&sort_expr), is_reversed: false, @@ -60,20 +74,81 @@ fn prepare_groups_accumulator(is_first: bool) -> Box { "LAST_VALUE(value ORDER BY ord)" }, is_distinct: false, - exprs: &[col("value", &schema).unwrap()], + exprs: std::slice::from_ref(&value_expr), }; + // Mirror the planner: use the native GroupsAccumulator when this value type + // is supported and otherwise fall back to a GroupsAccumulatorAdapter around + // one per-group Accumulator. Deciding with `groups_accumulator_supported` + // (rather than catching `create_groups_accumulator` errors) keeps genuine + // construction failures loud. The same case then runs the fallback on a + // build without native nested support and the native path on one with it, + // so a before/after benchmark run surfaces the win directly. + let supported = if is_first { + FirstValue::new().groups_accumulator_supported(make_args()) + } else { + LastValue::new().groups_accumulator_supported(make_args()) + }; + if !supported { + return build_fallback_adapter(is_first, value_type); + } if is_first { FirstValue::new() - .create_groups_accumulator(accumulator_args) + .create_groups_accumulator(make_args()) .unwrap() } else { LastValue::new() - .create_groups_accumulator(accumulator_args) + .create_groups_accumulator(make_args()) .unwrap() } } +/// Build the *fallback* grouped accumulator for a value type: a +/// `GroupsAccumulatorAdapter` wrapping one per-group `Accumulator`. This is +/// exactly what nested value types (`List` / `Struct` / `Map`) used before +/// they gained a native `GroupsAccumulator`, and it is what the planner still +/// selects when `groups_accumulator_supported` returns `false`. Benching this +/// side by side with `prepare_typed_groups_accumulator` (the native path) +/// shows the win from the native `GroupsAccumulator`. +fn build_fallback_adapter( + is_first: bool, + value_type: DataType, +) -> Box { + Box::new(GroupsAccumulatorAdapter::new(move || { + let schema = Arc::new(Schema::new(vec![ + Field::new("value", value_type.clone(), true), + Field::new("ord", DataType::Int64, true), + ])); + let sort_expr = PhysicalSortExpr { + expr: col("ord", &schema)?, + options: SortOptions::default(), + }; + let value_field: Arc = + Field::new("value", value_type.clone(), true).into(); + let value_expr = col("value", &schema)?; + let accumulator_args = AccumulatorArgs { + return_field: Arc::clone(&value_field), + schema: &schema, + expr_fields: std::slice::from_ref(&value_field), + ignore_nulls: false, + order_bys: std::slice::from_ref(&sort_expr), + is_reversed: false, + name: if is_first { + "FIRST_VALUE(value ORDER BY ord)" + } else { + "LAST_VALUE(value ORDER BY ord)" + }, + is_distinct: false, + exprs: std::slice::from_ref(&value_expr), + }; + if is_first { + FirstValue::new().accumulator(accumulator_args) + } else { + LastValue::new().accumulator(accumulator_args) + } + })) +} + fn create_trivial_accumulator( is_first: bool, ignore_nulls: bool, @@ -104,11 +179,13 @@ fn evaluate_bench( ) { let n = values.len(); let group_indices: Vec = (0..n).map(|i| i % num_groups).collect(); + let value_type = values.data_type().clone(); c.bench_function(name, |b| { b.iter_batched( || { - let mut accumulator = prepare_groups_accumulator(is_first); + let mut accumulator = + prepare_typed_groups_accumulator(is_first, value_type.clone()); accumulator .update_batch( &[Arc::clone(&values), Arc::clone(&ord)], @@ -139,6 +216,7 @@ fn update_bench( ) { let n = values.len(); let group_indices: Vec = (0..n).map(|i| i % num_groups).collect(); + let value_type = values.data_type().clone(); // Initialize with worst-case ordering so update_batch forces rows comparison for all groups. let worst_ord: ArrayRef = Arc::new(Int64Array::from(vec![ @@ -153,7 +231,8 @@ fn update_bench( c.bench_function(name, |b| { b.iter_batched( || { - let mut accumulator = prepare_groups_accumulator(is_first); + let mut accumulator = + prepare_typed_groups_accumulator(is_first, value_type.clone()); accumulator .update_batch( &[Arc::clone(&values), Arc::clone(&worst_ord)], @@ -197,6 +276,7 @@ fn merge_bench( let n = values.len(); let group_indices: Vec = (0..n).map(|i| i % num_groups).collect(); let is_set: ArrayRef = Arc::new(BooleanArray::from(vec![true; n])); + let value_type = values.data_type().clone(); // Initialize with worst-case ordering so update_batch forces rows comparison for all groups. let worst_ord: ArrayRef = Arc::new(Int64Array::from(vec![ @@ -212,7 +292,8 @@ fn merge_bench( b.iter_batched( || { // Prebuild accumulator - let mut accumulator = prepare_groups_accumulator(is_first); + let mut accumulator = + prepare_typed_groups_accumulator(is_first, value_type.clone()); accumulator .update_batch( &[Arc::clone(&values), Arc::clone(&worst_ord)], @@ -235,7 +316,6 @@ fn merge_bench( Arc::clone(&is_set), ], &group_indices, - opt_filter, num_groups, ) .unwrap(), @@ -271,6 +351,167 @@ fn trivial_update_bench( }); } +/// A top-level validity buffer with roughly `null_density` nulls, so the +/// generated nested arrays have null *values* (not just null inner +/// fields/elements) — matching the `nulls={pct}%` semantics of the primitive +/// benchmarks, where the value itself is null. Returns `None` at 0% so the +/// arrays stay fully valid. Derived from arrow's own null generator for a +/// deterministic, density-accurate pattern. +fn top_level_nulls(n: usize, null_density: f32) -> Option { + create_primitive_array::(n, null_density) + .nulls() + .cloned() +} + +/// A 3-field struct value column `Struct`. `null_density` +/// controls both the struct-level null values and the inner field nulls. +fn create_struct_array(n: usize, null_density: f32) -> ArrayRef { + let a = Arc::new(create_primitive_array::(n, null_density)) as ArrayRef; + let b = + Arc::new(create_string_array_with_len::(n, null_density, 16)) as ArrayRef; + let d = Arc::new(create_primitive_array::(n, null_density)) as ArrayRef; + let fields = Fields::from(vec![ + Field::new("c0", DataType::Int64, true), + Field::new("c1", DataType::Utf8, true), + Field::new("c2", DataType::Float64, true), + ]); + Arc::new(StructArray::new( + fields, + vec![a, b, d], + top_level_nulls(n, null_density), + )) +} + +/// A `List` value column with fixed-size lists of `list_len` elements. +fn create_list_array(n: usize, list_len: usize, null_density: f32) -> ArrayRef { + let child = Arc::new(create_primitive_array::( + n * list_len, + null_density, + )) as ArrayRef; + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(list_len, n)); + let field = Arc::new(Field::new_list_field(DataType::Int64, true)); + Arc::new(ListArray::new( + field, + offsets, + child, + top_level_nulls(n, null_density), + )) +} + +/// A `Map` value column with `entries_per_row` entries per row. +/// Values carry `null_density` nulls (keys are never null), matching the null +/// treatment of the struct / list generators. +fn create_map_array(n: usize, entries_per_row: usize, null_density: f32) -> ArrayRef { + let total = n * entries_per_row; + let values = + Arc::new(create_primitive_array::(total, null_density)) as ArrayRef; + let keys = Arc::new(StringArray::from_iter_values( + (0..total).map(|idx| format!("k{}", idx % entries_per_row)), + )) as ArrayRef; + let entry_fields = Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", DataType::Int64, true), + ]); + let entries = StructArray::new(entry_fields.clone(), vec![keys, values], None); + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(entries_per_row, n)); + let map_field = + Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)); + Arc::new(MapArray::new( + map_field, + offsets, + entries, + top_level_nulls(n, null_density), + false, + )) +} + +/// A composite `List>` column — a list whose +/// elements are structs (the "array of records" shape). Exercises the +/// nested-within-nested case, which the generic value-state path must also +/// handle. +fn create_list_of_struct_array(n: usize, list_len: usize, null_density: f32) -> ArrayRef { + let total = n * list_len; + let a = + Arc::new(create_primitive_array::(total, null_density)) as ArrayRef; + let b = + Arc::new(create_string_array_with_len::(total, null_density, 8)) as ArrayRef; + let struct_fields = Fields::from(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Utf8, true), + ]); + let child = + Arc::new(StructArray::new(struct_fields.clone(), vec![a, b], None)) as ArrayRef; + let offsets = OffsetBuffer::from_lengths(std::iter::repeat_n(list_len, n)); + let list_field = + Arc::new(Field::new_list_field(DataType::Struct(struct_fields), true)); + Arc::new(ListArray::new( + list_field, + offsets, + child, + top_level_nulls(n, null_density), + )) +} + +fn first_last_nested_benchmark(c: &mut Criterion) { + const N: usize = 65536; + const NUM_GROUPS: usize = 1024; + + let ord = Arc::new(create_primitive_array::(N, 0.0)) as ArrayRef; + + for pct in [0, 90] { + let null_density = (pct as f32) / 100.0; + + // One column per nested value type. Each type gets the same treatment + // as the primitive first_value / last_value benchmarks: update and + // merge (both first and last) plus evaluate, at 0% and 90% nulls. On a + // build without native nested support these run the fallback adapter; + // with this PR they run the native GroupsAccumulator, so the benchmark + // bot's before/after diff shows the win per type. + let columns: [(&str, ArrayRef); 4] = [ + ("struct(i64,utf8,f64)", create_struct_array(N, null_density)), + ("list[4]", create_list_array(N, 4, null_density)), + ("map", create_map_array(N, 4, null_density)), + ( + "list[4]", + create_list_of_struct_array(N, 4, null_density), + ), + ]; + + for (type_label, values) in columns { + for (fn_label, is_first) in [("first_value", true), ("last_value", false)] { + update_bench( + c, + is_first, + &format!("{fn_label} update_bench {type_label} nulls={pct}%"), + values.clone(), + ord.clone(), + None, + NUM_GROUPS, + ); + merge_bench( + c, + is_first, + &format!("{fn_label} merge_bench {type_label} nulls={pct}%"), + values.clone(), + ord.clone(), + None, + NUM_GROUPS, + ); + } + evaluate_bench( + c, + true, + EmitTo::All, + &format!("first_value evaluate_bench {type_label} nulls={pct}%, all"), + values.clone(), + ord.clone(), + None, + NUM_GROUPS, + ); + } + } +} + fn first_last_benchmark(c: &mut Criterion) { const N: usize = 65536; const NUM_GROUPS: usize = 1024; @@ -355,5 +596,5 @@ fn first_last_benchmark(c: &mut Criterion) { } } -criterion_group!(benches, first_last_benchmark); +criterion_group!(benches, first_last_benchmark, first_last_nested_benchmark); criterion_main!(benches); diff --git a/datafusion/functions-aggregate/benches/sliding_max.rs b/datafusion/functions-aggregate/benches/sliding_max.rs new file mode 100644 index 0000000000000..d5de001a1a79d --- /dev/null +++ b/datafusion/functions-aggregate/benches/sliding_max.rs @@ -0,0 +1,113 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, Int64Array, StringArray}; +use arrow::datatypes::DataType; +use criterion::{BenchmarkId, Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_expr::Accumulator; +use datafusion_functions_aggregate::min_max::SlidingMaxAccumulator; +use rand::Rng; +use rand::SeedableRng; +use rand::rngs::StdRng; +use std::sync::Arc; + +fn generate_random_i64(size: usize) -> Vec { + let mut rng = StdRng::seed_from_u64(42); + (0..size).map(|_| rng.random_range(0..1_000_000)).collect() +} + +fn generate_random_strings(size: usize) -> Vec { + let mut rng = StdRng::seed_from_u64(42); + (0..size) + .map(|_| { + let len = rng.random_range(10..40); + (0..len) + .map(|_| rng.random_range(b'a'..=b'z') as char) + .collect() + }) + .collect() +} + +/// Simulates a sliding window by calling update_batch and retract_batch +/// on SlidingMaxAccumulator, mirroring how the query engine uses it. +fn bench_sliding_max_for( + c: &mut Criterion, + label: &str, + data_type: &DataType, + array: &ArrayRef, + data_size: usize, + window_size: usize, +) { + let mut group = c.benchmark_group(format!("sliding_window_max_{label}")); + group.throughput(Throughput::Elements(data_size as u64)); + + group.bench_with_input( + BenchmarkId::new("sliding_max", window_size), + &window_size, + |b, &w| { + b.iter(|| { + let mut acc = SlidingMaxAccumulator::try_new(data_type).unwrap(); + // Warm up the window + let init_batch = array.slice(0, w); + acc.update_batch(&[init_batch]).unwrap(); + + // Slide: for each subsequent element, add it and retract one + for i in w..data_size { + let new_val = array.slice(i, 1); + let old_val = array.slice(i - w, 1); + acc.update_batch(&[new_val]).unwrap(); + acc.retract_batch(&[old_val]).unwrap(); + std::hint::black_box(acc.evaluate().unwrap()); + } + }); + }, + ); + + group.finish(); +} + +fn bench_sliding_max(c: &mut Criterion) { + let data_size = 50_000; + + let i64_data: Vec = generate_random_i64(data_size); + let str_data: Vec = generate_random_strings(data_size); + + let i64_array: ArrayRef = Arc::new(Int64Array::from(i64_data)); + let str_array: ArrayRef = Arc::new(StringArray::from(str_data)); + + for window_size in [100, 1000, 5000] { + bench_sliding_max_for( + c, + "int64", + &DataType::Int64, + &i64_array, + data_size, + window_size, + ); + bench_sliding_max_for( + c, + "utf8", + &DataType::Utf8, + &str_array, + data_size, + window_size, + ); + } +} + +criterion_group!(benches, bench_sliding_max); +criterion_main!(benches); diff --git a/datafusion/functions-aggregate/benches/variance.rs b/datafusion/functions-aggregate/benches/variance.rs new file mode 100644 index 0000000000000..ef55bf32b8843 --- /dev/null +++ b/datafusion/functions-aggregate/benches/variance.rs @@ -0,0 +1,83 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Float64Array}; +use criterion::{BatchSize, Criterion, criterion_group, criterion_main}; +use datafusion_expr::Accumulator; +use datafusion_functions_aggregate::variance::VarianceAccumulator; +use datafusion_functions_aggregate_common::stats::StatsType; + +const BATCH_SIZE: usize = 8192; + +fn batch_array(null_stride: Option) -> ArrayRef { + let values = (0..BATCH_SIZE) + .map(|idx| { + if null_stride.is_some_and(|stride| idx % stride == 0) { + None + } else { + Some(idx as f64) + } + }) + .collect::>(); + Arc::new(Float64Array::from(values)) as ArrayRef +} + +fn update_bench(c: &mut Criterion, name: &str, batch: &ArrayRef) { + c.bench_function(name, |b| { + b.iter(|| { + let mut acc = VarianceAccumulator::try_new(StatsType::Sample).unwrap(); + acc.update_batch(std::slice::from_ref(batch)).unwrap(); + black_box(acc.evaluate().unwrap()) + }) + }); +} + +fn retract_bench(c: &mut Criterion, name: &str, batch: &ArrayRef) { + c.bench_function(name, |b| { + b.iter_batched( + || { + let mut acc = VarianceAccumulator::try_new(StatsType::Sample).unwrap(); + // Accumulate two batches so that retracting one leaves the + // accumulator with rows remaining, as in a sliding window. + acc.update_batch(std::slice::from_ref(batch)).unwrap(); + acc.update_batch(std::slice::from_ref(batch)).unwrap(); + acc + }, + |mut acc| { + acc.retract_batch(std::slice::from_ref(batch)).unwrap(); + black_box(acc.evaluate().unwrap()) + }, + BatchSize::SmallInput, + ) + }); +} + +fn variance_benchmark(c: &mut Criterion) { + let no_nulls = batch_array(None); + let with_nulls = batch_array(Some(10)); + + update_bench(c, "variance update_batch f64 no_nulls", &no_nulls); + update_bench(c, "variance update_batch f64 with_nulls", &with_nulls); + retract_bench(c, "variance retract_batch f64 no_nulls", &no_nulls); + retract_bench(c, "variance retract_batch f64 with_nulls", &with_nulls); +} + +criterion_group!(benches, variance_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions-aggregate/src/any_value.rs b/datafusion/functions-aggregate/src/any_value.rs new file mode 100644 index 0000000000000..dc3bd23d806fc --- /dev/null +++ b/datafusion/functions-aggregate/src/any_value.rs @@ -0,0 +1,125 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Defines the ANY_VALUE aggregation. + +use std::fmt::Debug; +use std::hash::Hash; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::{Result, not_impl_err}; +use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; +use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name}; +use datafusion_expr::{ + Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility, +}; +use datafusion_macros::user_doc; + +use crate::first_last::TrivialFirstValueAccumulator; + +make_udaf_expr_and_func!( + AnyValue, + any_value, + expression, + "Returns an arbitrary non-null value", + any_value_udaf +); + +#[user_doc( + doc_section(label = "General Functions"), + description = "Returns an arbitrary non-null value from a group, or NULL if the group contains only NULL values.", + syntax_example = "any_value(expression)", + sql_example = r#"```sql +> SELECT any_value(column_name) FROM table_name; ++------------------------+ +| any_value(column_name) | ++------------------------+ +| arbitrary_value | ++------------------------+ +```"#, + standard_argument(name = "expression",) +)] +#[derive(PartialEq, Eq, Hash, Debug)] +pub struct AnyValue { + signature: Signature, +} + +impl Default for AnyValue { + fn default() -> Self { + Self::new() + } +} + +impl AnyValue { + pub fn new() -> Self { + Self { + signature: Signature::any(1, Volatility::Immutable), + } + } +} + +impl AggregateUDFImpl for AnyValue { + fn name(&self) -> &str { + "any_value" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + not_impl_err!("Not called because return_field is implemented") + } + + fn return_field(&self, arg_fields: &[FieldRef]) -> Result { + Ok(Arc::new( + Field::new(self.name(), arg_fields[0].data_type().clone(), true) + .with_metadata(arg_fields[0].metadata().clone()), + )) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { + TrivialFirstValueAccumulator::try_new(acc_args.return_field.data_type(), true) + .map(|acc| Box::new(acc) as _) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result> { + Ok(vec![ + Field::new( + format_state_name(args.name, "any_value"), + args.return_type().clone(), + true, + ) + .into(), + Field::new( + format_state_name(args.name, "any_value_is_set"), + DataType::Boolean, + true, + ) + .into(), + ]) + } + + fn order_sensitivity(&self) -> AggregateOrderSensitivity { + AggregateOrderSensitivity::Insensitive + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} diff --git a/datafusion/functions-aggregate/src/approx_distinct.rs b/datafusion/functions-aggregate/src/approx_distinct.rs index cc42b6c22bdbe..1746edd8239f2 100644 --- a/datafusion/functions-aggregate/src/approx_distinct.rs +++ b/datafusion/functions-aggregate/src/approx_distinct.rs @@ -17,19 +17,23 @@ //! Defines physical expressions that can evaluated at runtime during query execution -use crate::hyperloglog::{HLL_HASH_STATE, HyperLogLog}; -use arrow::array::{Array, BinaryArray, StringViewArray}; +use crate::hyperloglog::{HLL_HASH_STATE, HyperLogLog, NUM_REGISTERS, count_from_hashes}; use arrow::array::{ - GenericBinaryArray, GenericStringArray, OffsetSizeTrait, PrimitiveArray, + Array, ArrayRef, BinaryArray, BinaryBuilder, BooleanArray, PrimitiveArray, + UInt64Array, }; +use arrow::buffer::NullBuffer; use arrow::datatypes::{ - ArrowPrimitiveType, Date32Type, Date64Type, FieldRef, Int32Type, Int64Type, + ArrowPrimitiveType, DataType, Date32Type, Date64Type, Decimal32Type, Decimal64Type, + Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType, + DurationNanosecondType, DurationSecondType, Field, FieldRef, Int32Type, Int64Type, + IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, IntervalYearMonthType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt32Type, UInt64Type, }; -use arrow::{array::ArrayRef, datatypes::DataType, datatypes::Field}; use datafusion_common::ScalarValue; +use datafusion_common::hash_utils::create_hashes; use datafusion_common::{ DataFusionError, Result, downcast_value, internal_datafusion_err, internal_err, not_impl_err, @@ -37,17 +41,21 @@ use datafusion_common::{ use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, Documentation, Signature, Volatility, + Accumulator, AggregateUDFImpl, Documentation, EmitTo, GroupsAccumulator, Signature, + Volatility, }; use datafusion_functions_aggregate_common::aggregate::count_distinct::{ Bitmap65536DistinctCountAccumulator, Bitmap65536DistinctCountAccumulatorI16, BoolArray256DistinctCountAccumulator, BoolArray256DistinctCountAccumulatorI8, + BooleanDistinctCountAccumulator, }; +use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filter_to_nulls; use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator; use datafusion_macros::user_doc; use std::fmt::{Debug, Formatter}; -use std::hash::{BuildHasher, Hash}; -use std::marker::PhantomData; +use std::hash::Hash; +use std::mem::{size_of, size_of_val}; +use std::sync::Arc; make_udaf_expr_and_func!( ApproxDistinct, @@ -117,6 +125,73 @@ impl Accumulator for ApproxDistinctBitmapWrapper { } } +#[derive(Debug)] +struct HLLAccumulator { + hll: HyperLogLog, + hashes: Vec, +} + +impl HLLAccumulator { + pub fn new() -> Self { + Self { + hll: HyperLogLog::new(), + hashes: Vec::new(), + } + } +} + +impl Accumulator for HLLAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let array = values[0].as_ref(); + self.hashes.clear(); + self.hashes.resize(array.len(), 0); + create_hashes([array], &HLL_HASH_STATE, &mut self.hashes)?; + + match array.logical_nulls() { + None => { + for &hash in &self.hashes { + self.hll.add_hashed(hash); + } + } + Some(nulls) => { + for row in 0..array.len() { + if nulls.is_valid(row) { + self.hll.add_hashed(self.hashes[row]); + } + } + } + } + Ok(()) + } + + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + assert_eq!(1, states.len(), "expect only 1 element in the states"); + let binary_array = downcast_value!(states[0], BinaryArray); + for v in binary_array.iter() { + let v = v.ok_or_else(|| { + internal_datafusion_err!("Impossibly got empty binary array from states") + })?; + let other = v.try_into()?; + self.hll.merge(&other); + } + Ok(()) + } + + fn state(&mut self) -> Result> { + let value = ScalarValue::from(&self.hll); + Ok(vec![value]) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::UInt64(Some(self.hll.count() as u64))) + } + + fn size(&self) -> usize { + size_of_val(self) + self.hashes.capacity() * size_of::() + } +} + +/// Specialize the numeric case for extra performance. #[derive(Debug)] struct NumericHLLAccumulator where @@ -138,159 +213,413 @@ where } } -#[derive(Debug)] -struct StringHLLAccumulator +impl Accumulator for NumericHLLAccumulator where - T: OffsetSizeTrait, + T: ArrowPrimitiveType + Debug, + T::Native: Hash, { - hll: HyperLogLog, - phantom_data: PhantomData, -} + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let array: &PrimitiveArray = downcast_value!(values[0], PrimitiveArray, T); + self.hll.extend(array.into_iter().flatten()); + Ok(()) + } -impl StringHLLAccumulator -where - T: OffsetSizeTrait, -{ - pub fn new() -> Self { - Self { - hll: HyperLogLog::new(), - phantom_data: PhantomData, + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + assert_eq!(1, states.len(), "expect only 1 element in the states"); + let binary_array = downcast_value!(states[0], BinaryArray); + for v in binary_array.iter() { + let v = v.ok_or_else(|| { + internal_datafusion_err!("Impossibly got empty binary array from states") + })?; + let other = v.try_into()?; + self.hll.merge(&other); } + Ok(()) + } + + fn state(&mut self) -> Result> { + let value = ScalarValue::from(&self.hll); + Ok(vec![value]) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::UInt64(Some(self.hll.count() as u64))) + } + + fn size(&self) -> usize { + size_of_val(self) } } -#[derive(Debug)] -struct StringViewHLLAccumulator { - hll: HyperLogLog, +/// Maximum number of distinct hashes kept in the sparse representation of a +/// per-group sketch before it is promoted to a dense [`HyperLogLog`]. +/// +/// A dense sketch always occupies [`NUM_REGISTERS`] (16 KiB) regardless of how +/// many values it has seen. The vast majority of groups in a high-cardinality +/// `GROUP BY` only observe a handful of distinct values, so keeping their state +/// as a small list of hashes saves a huge amount of memory (both while +/// aggregating and when serializing the partial state for the final phase). +const SPARSE_LIMIT: usize = 256; + +/// Per-group HyperLogLog state used by [`HllGroupsAccumulator`]. +/// +/// Starts out as a compact list of the (deduplicated) hashes observed for the +/// group and only switches to a full dense [`HyperLogLog`] once it has seen more +/// than [`SPARSE_LIMIT`] distinct values. Folding the stored hashes into a dense +/// sketch produces exactly the same registers as adding the original values one +/// by one, so the cardinality estimate is identical to the per-group +/// [`Accumulator`] path. +#[derive(Clone, Debug)] +enum GroupHll { + /// Distinct hashes seen so far. May contain duplicates between compactions. + Sparse(Vec), + Dense(Box>), } -impl StringViewHLLAccumulator { - pub fn new() -> Self { - Self { - hll: HyperLogLog::new(), - } +impl Default for GroupHll { + fn default() -> Self { + GroupHll::Sparse(Vec::new()) } } -#[derive(Debug)] -struct BinaryHLLAccumulator -where - T: OffsetSizeTrait, -{ - hll: HyperLogLog<[u8]>, - phantom_data: PhantomData, +/// Fold a slice of pre-computed hashes into a fresh [`HyperLogLog`] sketch. +fn fold_sparse_to_hll(hashes: &[u64]) -> HyperLogLog { + let mut hll = HyperLogLog::::new(); + for &h in hashes { + hll.add_hashed(h); + } + hll } -impl BinaryHLLAccumulator -where - T: OffsetSizeTrait, -{ - pub fn new() -> Self { - Self { - hll: HyperLogLog::new(), - phantom_data: PhantomData, +impl GroupHll { + /// Add a pre-computed hash, returning the change in heap-allocated bytes so + /// the accumulator can track its memory usage incrementally. + #[inline] + fn add_hash(&mut self, hash: u64) -> isize { + match self { + GroupHll::Dense(hll) => { + hll.add_hashed(hash); + 0 + } + GroupHll::Sparse(v) => { + let cap_before = v.capacity(); + v.push(hash); + if v.len() >= 2 * SPARSE_LIMIT { + return self.compact_or_promote(cap_before); + } + ((v.capacity() - cap_before) * size_of::()) as isize + } } } -} -macro_rules! default_accumulator_impl { - () => { - fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - assert_eq!(1, states.len(), "expect only 1 element in the states"); - let binary_array = downcast_value!(states[0], BinaryArray); - for v in binary_array.iter() { - let v = v.ok_or_else(|| { - internal_datafusion_err!( - "Impossibly got empty binary array from states" - ) - })?; - let other = v.try_into()?; - self.hll.merge(&other); + /// Deduplicate the sparse hash list and, if it still exceeds + /// [`SPARSE_LIMIT`] distinct values, promote it to a dense sketch. + #[cold] + fn compact_or_promote(&mut self, cap_before: usize) -> isize { + let GroupHll::Sparse(v) = self else { + return 0; + }; + v.sort_unstable(); + v.dedup(); + if v.len() > SPARSE_LIMIT { + // cap_before is the capacity already reflected in allocated_bytes. + // Any reallocation caused by the triggering push was never counted and + // is also freed here, so the two cancel out. + *self = GroupHll::Dense(Box::new(fold_sparse_to_hll(v))); + (NUM_REGISTERS as isize) - ((cap_before * size_of::()) as isize) + } else { + // Account for any Vec growth caused by the triggering push. + // sort/dedup do not reallocate, so v.capacity() is the post-push capacity. + ((v.capacity() - cap_before) * size_of::()) as isize + } + } + + /// Merge a serialized state (produced by [`Self::serialize`] or by the + /// per-group [`Accumulator`]) into this sketch. + fn merge_serialized(&mut self, bytes: &[u8]) -> Result { + if bytes.is_empty() { + return Ok(0); + } + if bytes.len() == NUM_REGISTERS { + let other: HyperLogLog = bytes.try_into()?; + Ok(self.merge_dense(&other)) + } else { + if !bytes.len().is_multiple_of(size_of::()) { + return internal_err!( + "approx_distinct: malformed sparse state: length {} is not a multiple of {}", + bytes.len(), + size_of::() + ); + } + if bytes.len() > SPARSE_LIMIT * size_of::() { + return internal_err!( + "approx_distinct: malformed sparse state: length {} exceeds sparse limit {}", + bytes.len(), + SPARSE_LIMIT * size_of::() + ); } - Ok(()) + let mut delta = 0; + for chunk in bytes.chunks_exact(size_of::()) { + let h = u64::from_le_bytes(chunk.try_into().unwrap()); + delta += self.add_hash(h); + } + Ok(delta) } + } - fn state(&mut self) -> Result> { - let value = ScalarValue::from(&self.hll); - Ok(vec![value]) + /// Merge a dense sketch into this one, promoting to dense if necessary. + fn merge_dense(&mut self, other: &HyperLogLog) -> isize { + match self { + GroupHll::Dense(hll) => { + hll.merge(other); + 0 + } + GroupHll::Sparse(v) => { + let cap_before = v.capacity(); + let mut hll = other.clone(); + for &h in v.iter() { + hll.add_hashed(h); + } + *self = GroupHll::Dense(Box::new(hll)); + (NUM_REGISTERS as isize) - ((cap_before * size_of::()) as isize) + } } + } - fn evaluate(&mut self) -> Result { - Ok(ScalarValue::UInt64(Some(self.hll.count() as u64))) + /// The approximate number of distinct values seen by this group. + fn count(&self) -> u64 { + match self { + GroupHll::Dense(hll) => hll.count() as u64, + // Estimate directly from the stored hashes; this produces exactly the + // same value as folding them into a dense sketch but avoids + // allocating and scanning a 16 KiB register array for every group. + GroupHll::Sparse(v) => count_from_hashes(v) as u64, } + } - fn size(&self) -> usize { - // HLL has static size - std::mem::size_of_val(self) + /// Heap bytes held by this sketch. Mirrors the deltas accrued in + /// [`Self::add_hash`] / [`Self::merge_dense`] so emitting a group can + /// precisely reverse them. + fn heap_bytes(&self) -> usize { + match self { + GroupHll::Sparse(v) => v.capacity() * size_of::(), + GroupHll::Dense(_) => NUM_REGISTERS, } - }; + } + + /// Serialize the sketch into `scratch` (which is cleared first). A dense + /// sketch is written as its raw [`NUM_REGISTERS`] registers (wire-compatible + /// with the per-group [`Accumulator`]); a sparse sketch is written as its + /// distinct hashes in little-endian order unless it has crossed + /// [`SPARSE_LIMIT`], in which case it is emitted as dense state so the final + /// merge path accepts it. + fn serialize(&mut self, scratch: &mut Vec) { + scratch.clear(); + match self { + GroupHll::Dense(hll) => { + let registers: &[u8] = (**hll).as_ref(); + scratch.extend_from_slice(registers); + } + GroupHll::Sparse(v) => { + v.sort_unstable(); + v.dedup(); + if v.len() > SPARSE_LIMIT { + scratch.extend_from_slice(fold_sparse_to_hll(v).as_ref()); + } else { + for &h in v.iter() { + scratch.extend_from_slice(&h.to_le_bytes()); + } + } + } + } + } } -impl Accumulator for BinaryHLLAccumulator -where - T: OffsetSizeTrait, -{ - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let array: &GenericBinaryArray = - downcast_value!(values[0], GenericBinaryArray, T); - // flatten because we would skip nulls - self.hll.extend(array.into_iter().flatten()); - Ok(()) +/// A [`GroupsAccumulator`] for `approx_distinct` that keeps one adaptive +/// (sparse → dense) HyperLogLog sketch per group. +/// +/// This is dramatically faster than the generic `GroupsAccumulatorAdapter` +/// fallback for high-cardinality `GROUP BY`s: it processes the whole input in a +/// single vectorized pass (no per-group `take`/slice and no dynamic dispatch), +/// and the sparse representation avoids allocating a 16 KiB sketch for every +/// group when most groups only see a few distinct values. +/// +/// +/// # Example +/// +/// For `SELECT k, approx_distinct(v) FROM t GROUP BY k`, each group owns one +/// independent sketch: +/// +/// ```text +/// group state +/// a Sparse([h1, h2, h3, h2]) +/// b Dense(HLL registers) +/// ... +/// ``` +/// +/// Group `a` has fewer than [`SPARSE_LIMIT`] distinct hashes, so it stays in +/// the sparse representation. Before emitting state or estimating the count, the +/// hash list is sorted and deduplicated to `[h1, h2, h3]`, then those hashes are +/// interpreted exactly as if they had been added to a dense [`HyperLogLog`]. +/// +/// Group `b` has crossed the sparse limit, so its hashes have already been +/// replayed into a dense sketch. New values for `b` update the dense registers +/// directly, and serialized state is the raw [`NUM_REGISTERS`]-byte register +/// array. +struct HllGroupsAccumulator { + /// Per-group sketches, indexed by `group_index`. + groups: Vec, + /// Incrementally maintained estimate of heap bytes used by `groups`. + allocated_bytes: usize, + /// Reused workspace for vectorized value hashing. + hashes: Vec, +} + +impl HllGroupsAccumulator { + fn new() -> Self { + Self { + groups: Vec::new(), + allocated_bytes: 0, + hashes: Vec::new(), + } + } + + #[inline] + fn ensure_groups(&mut self, total_num_groups: usize) { + if total_num_groups > self.groups.len() { + self.groups.resize_with(total_num_groups, GroupHll::default); + } } - default_accumulator_impl!(); + #[inline] + fn apply_delta(&mut self, delta: isize) { + self.allocated_bytes = + (self.allocated_bytes as isize).saturating_add(delta).max(0) as usize; + } } -impl Accumulator for StringViewHLLAccumulator { - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let array: &StringViewArray = downcast_value!(values[0], StringViewArray); - - // When all strings are stored inline in the StringView (≤ 12 bytes), - // hash the raw u128 view directly instead of materializing a &str. - if array.data_buffers().is_empty() { - for (i, &view) in array.views().iter().enumerate() { - if !array.is_null(i) { - self.hll.add_hashed(HLL_HASH_STATE.hash_one(view)); +impl GroupsAccumulator for HllGroupsAccumulator { + fn update_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + self.ensure_groups(total_num_groups); + let array = values[0].as_ref(); + self.hashes.clear(); + self.hashes.resize(array.len(), 0); + create_hashes([array], &HLL_HASH_STATE, &mut self.hashes)?; + + let mut delta: isize = 0; + // Pre-combine value-nulls and filter into one mask so the update loop + // only visits rows that should affect the sketch. + let filter_nulls = opt_filter.map(filter_to_nulls); + let value_nulls = array.logical_nulls(); + let combined_nulls = + NullBuffer::union(filter_nulls.as_ref(), value_nulls.as_ref()); + match combined_nulls { + None => { + for (row, &hash) in self.hashes.iter().enumerate() { + delta += self.groups[group_indices[row]].add_hash(hash); + } + } + Some(nulls) => { + for row in nulls.valid_indices() { + delta += self.groups[group_indices[row]].add_hash(self.hashes[row]); } } - } else { - self.hll.extend(array.iter().flatten()); } - + self.apply_delta(delta); Ok(()) } - default_accumulator_impl!(); -} - -impl Accumulator for StringHLLAccumulator -where - T: OffsetSizeTrait, -{ - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let array: &GenericStringArray = - downcast_value!(values[0], GenericStringArray, T); - // flatten because we would skip nulls - self.hll.extend(array.into_iter().flatten()); + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + self.ensure_groups(total_num_groups); + let states = downcast_value!(values[0], BinaryArray); + let mut delta: isize = 0; + for (row, &group_index) in group_indices.iter().enumerate() { + if states.is_valid(row) { + delta += self.groups[group_index].merge_serialized(states.value(row))?; + } + } + self.apply_delta(delta); Ok(()) } - default_accumulator_impl!(); -} + fn evaluate(&mut self, emit_to: EmitTo) -> Result { + let groups = emit_to.take_needed(&mut self.groups); + let mut freed = 0; + let counts: UInt64Array = groups + .iter() + .map(|g| { + freed += g.heap_bytes(); + Some(g.count()) + }) + .collect(); + // The emitted groups have been removed; reclaim their tracked bytes. + self.allocated_bytes = self.allocated_bytes.saturating_sub(freed); + Ok(Arc::new(counts)) + } -impl Accumulator for NumericHLLAccumulator -where - T: ArrowPrimitiveType + Debug, - T::Native: Hash, -{ - fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let array: &PrimitiveArray = downcast_value!(values[0], PrimitiveArray, T); - // flatten because we would skip nulls - self.hll.extend(array.into_iter().flatten()); - Ok(()) + fn state(&mut self, emit_to: EmitTo) -> Result> { + let mut groups = emit_to.take_needed(&mut self.groups); + let mut builder = BinaryBuilder::new(); + let mut scratch: Vec = Vec::new(); + let mut freed = 0; + for g in groups.iter_mut() { + freed += g.heap_bytes(); + g.serialize(&mut scratch); + builder.append_value(&scratch); + } + // The emitted groups have been removed; reclaim their tracked bytes. + self.allocated_bytes = self.allocated_bytes.saturating_sub(freed); + Ok(vec![Arc::new(builder.finish())]) } - default_accumulator_impl!(); + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + assert_eq!(values.len(), 1, "single argument to convert_to_state"); + let array = values[0].as_ref(); + let mut hashes = vec![0; array.len()]; + create_hashes([array], &HLL_HASH_STATE, &mut hashes)?; + + let filter_nulls = opt_filter.map(filter_to_nulls); + let value_nulls = array.logical_nulls(); + let combined_nulls = + NullBuffer::union(filter_nulls.as_ref(), value_nulls.as_ref()); + + let mut builder = BinaryBuilder::new(); + let mut scratch = Vec::new(); + for (row, hash) in hashes.into_iter().enumerate() { + if combined_nulls + .as_ref() + .is_none_or(|nulls| nulls.is_valid(row)) + { + scratch.clear(); + scratch.extend_from_slice(&hash.to_le_bytes()); + builder.append_value(&scratch); + } else { + builder.append_value([]); + } + } + + Ok(vec![Arc::new(builder.finish())]) + } + fn size(&self) -> usize { + self.groups.capacity() * size_of::() + + self.allocated_bytes + + self.hashes.capacity() * size_of::() + } } impl Debug for ApproxDistinct { @@ -336,10 +665,13 @@ impl ApproxDistinct { } #[cold] -fn get_small_int_approx_accumulator( +fn get_fixed_domain_approx_accumulator( data_type: &DataType, ) -> Result> { match data_type { + DataType::Boolean => Ok(Box::new(ApproxDistinctBitmapWrapper { + inner: BooleanDistinctCountAccumulator::new(), + })), DataType::UInt8 => Ok(Box::new(ApproxDistinctBitmapWrapper { inner: BoolArray256DistinctCountAccumulator::new(), })), @@ -357,7 +689,10 @@ fn get_small_int_approx_accumulator( } #[cold] -fn get_small_int_state_field(name: &str, data_type: &DataType) -> Result> { +fn get_fixed_domain_state_field( + name: &str, + data_type: &DataType, +) -> Result> { Ok(vec![ Field::new_list( format_state_name(name, "approx_distinct"), @@ -381,6 +716,14 @@ impl AggregateUDFImpl for ApproxDistinct { Ok(DataType::UInt64) } + fn default_value(&self, _data_type: &DataType) -> Result { + Ok(ScalarValue::UInt64(Some(0))) + } + + fn is_nullable(&self) -> bool { + false + } + fn state_fields(&self, args: StateFieldsArgs) -> Result> { let data_type = args.input_fields[0].data_type(); match data_type { @@ -392,9 +735,11 @@ impl AggregateUDFImpl for ApproxDistinct { ) .into(), ]), - DataType::UInt8 | DataType::Int8 | DataType::UInt16 | DataType::Int16 => { - get_small_int_state_field(args.name, data_type) - } + DataType::Boolean + | DataType::UInt8 + | DataType::Int8 + | DataType::UInt16 + | DataType::Int16 => get_fixed_domain_state_field(args.name, data_type), _ => Ok(vec![ Field::new( format_state_name(args.name, "hll_registers"), @@ -409,9 +754,14 @@ impl AggregateUDFImpl for ApproxDistinct { fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { let data_type = acc_args.expr_fields[0].data_type(); + // For primitive types, use specialized accumulators for better performance. let accumulator: Box = match data_type { - DataType::UInt8 | DataType::Int8 | DataType::UInt16 | DataType::Int16 => { - return get_small_int_approx_accumulator(data_type); + DataType::Boolean + | DataType::UInt8 + | DataType::Int8 + | DataType::UInt16 + | DataType::Int16 => { + return get_fixed_domain_approx_accumulator(data_type); } DataType::UInt32 => Box::new(NumericHLLAccumulator::::new()), DataType::UInt64 => Box::new(NumericHLLAccumulator::::new()), @@ -443,11 +793,54 @@ impl AggregateUDFImpl for ApproxDistinct { DataType::Timestamp(TimeUnit::Nanosecond, _) => { Box::new(NumericHLLAccumulator::::new()) } - DataType::Utf8 => Box::new(StringHLLAccumulator::::new()), - DataType::LargeUtf8 => Box::new(StringHLLAccumulator::::new()), - DataType::Utf8View => Box::new(StringViewHLLAccumulator::new()), - DataType::Binary => Box::new(BinaryHLLAccumulator::::new()), - DataType::LargeBinary => Box::new(BinaryHLLAccumulator::::new()), + DataType::Interval(IntervalUnit::YearMonth) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Interval(IntervalUnit::DayTime) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Interval(IntervalUnit::MonthDayNano) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Decimal32(_, _) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Decimal64(_, _) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Decimal128(_, _) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Decimal256(_, _) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Duration(TimeUnit::Second) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Duration(TimeUnit::Millisecond) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Duration(TimeUnit::Microsecond) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Duration(TimeUnit::Nanosecond) => { + Box::new(NumericHLLAccumulator::::new()) + } + DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::BinaryView + | DataType::FixedSizeBinary(_) + | DataType::List(_) + | DataType::LargeList(_) + | DataType::FixedSizeList(_, _) + | DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::Map(_, _) + | DataType::Struct(_) + | DataType::Union(_, _) + | DataType::LargeBinary => Box::new(HLLAccumulator::new()), DataType::Null => { Box::new(NoopAccumulator::new(ScalarValue::UInt64(Some(0)))) } @@ -460,7 +853,567 @@ impl AggregateUDFImpl for ApproxDistinct { Ok(accumulator) } + fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { + is_hll_groups_type(args.expr_fields[0].data_type()) + } + + fn create_groups_accumulator( + &self, + args: AccumulatorArgs, + ) -> Result> { + let data_type = args.expr_fields[0].data_type(); + if is_hll_groups_type(data_type) { + Ok(Box::new(HllGroupsAccumulator::new())) + } else { + not_impl_err!( + "GroupsAccumulator for 'approx_distinct' is not implemented for data type {data_type}" + ) + } + } + fn documentation(&self) -> Option<&Documentation> { self.doc() } } + +/// Returns true for the data types backed by the HyperLogLog +/// [`HllGroupsAccumulator`]. The fixed-domain types (booleans / small ints) and +/// `Null` fall back to the per-group [`Accumulator`] path. +fn is_hll_groups_type(data_type: &DataType) -> bool { + matches!( + data_type, + DataType::UInt32 + | DataType::UInt64 + | DataType::Int32 + | DataType::Int64 + | DataType::Date32 + | DataType::Date64 + | DataType::Time32(TimeUnit::Second) + | DataType::Time32(TimeUnit::Millisecond) + | DataType::Time64(TimeUnit::Microsecond) + | DataType::Time64(TimeUnit::Nanosecond) + | DataType::Timestamp(TimeUnit::Second, _) + | DataType::Timestamp(TimeUnit::Millisecond, _) + | DataType::Timestamp(TimeUnit::Microsecond, _) + | DataType::Timestamp(TimeUnit::Nanosecond, _) + | DataType::Interval(IntervalUnit::YearMonth) + | DataType::Interval(IntervalUnit::DayTime) + | DataType::Interval(IntervalUnit::MonthDayNano) + | DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Duration(_) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Utf8View + | DataType::Binary + | DataType::BinaryView + | DataType::FixedSizeBinary(_) + | DataType::LargeBinary + | DataType::List(_) + | DataType::LargeList(_) + | DataType::FixedSizeList(_, _) + | DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::Map(_, _) + | DataType::Struct(_) + | DataType::Union(_, _) + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::hash::BuildHasher; + + #[cfg(not(feature = "force_hash_collisions"))] + mod real_hash_test { + use super::*; + use arrow::array::{ + AsArray, Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, + Int64Array, IntervalDayTimeArray, IntervalMonthDayNanoArray, + IntervalYearMonthArray, StringViewArray, + }; + use arrow::datatypes::{IntervalDayTime, IntervalMonthDayNano, i256}; + use std::sync::Arc; + // A string longer than the 12-byte inline limit + const LONG: &str = "this string is definitely longer than twelve bytes"; + + fn distinct_count(acc: &mut HLLAccumulator) -> u64 { + match acc.evaluate().unwrap() { + ScalarValue::UInt64(Some(v)) => v, + other => panic!("unexpected evaluate result: {other:?}"), + } + } + + fn assert_count_numerical_acc_and_group_acc(array: ArrayRef, expected: u64) + where + T: ArrowPrimitiveType + Debug, + T::Native: Hash, + { + assert!( + is_hll_groups_type(array.data_type()), + "{} should be groups-capable", + array.data_type() + ); + + let mut acc = NumericHLLAccumulator::::new(); + acc.update_batch(&[Arc::clone(&array)]).unwrap(); + let per_group_count = match acc.evaluate().unwrap() { + ScalarValue::UInt64(Some(v)) => v, + other => panic!("unexpected evaluate result: {other:?}"), + }; + + let group_indices = vec![0usize; array.len()]; + let mut acc = HllGroupsAccumulator::new(); + acc.update_batch(std::slice::from_ref(&array), &group_indices, None, 1) + .unwrap(); + let groups_count = acc + .evaluate(EmitTo::All) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + + assert_eq!( + per_group_count, + groups_count, + "paths disagree for {}", + array.data_type() + ); + assert_eq!( + per_group_count, + expected, + "wrong count for {}", + array.data_type() + ); + } + + #[test] + fn decimal_support_numerical_acc_and_group_acc() { + let decimal_32: ArrayRef = Arc::new( + Decimal32Array::from(vec![ + 1i32, + 2, + 2, + 3, + 3, + 3, + 0, + 0, + 123_456_789, + 999_999_999, + 999_999_999, + ]) + .with_precision_and_scale(9, 2) + .unwrap(), + ); + assert_count_numerical_acc_and_group_acc::(decimal_32, 6); + + let decimal_64: ArrayRef = Arc::new( + Decimal64Array::from(vec![ + 1i64, + 2, + 2, + 3, + 3, + 3, + 0, + 0, + 1_234_567_890_123, + 9_999_999_999_999, + 9_999_999_999_999, + ]) + .with_precision_and_scale(18, 2) + .unwrap(), + ); + assert_count_numerical_acc_and_group_acc::(decimal_64, 6); + + let decimal_128: ArrayRef = Arc::new( + Decimal128Array::from(vec![ + 1i128, + 2, + 2, + 3, + 3, + 3, + 0, + 0, + 1_234_567_890, + 9_999_999_999, + 9_999_999_999, + ]) + .with_precision_and_scale(38, 2) + .unwrap(), + ); + assert_count_numerical_acc_and_group_acc::(decimal_128, 6); + + let big_256_a = + i256::from_string("123456789012345678901234567890123456").unwrap(); + let big_256_b = + i256::from_string("987654321098765432109876543210987654").unwrap(); + + let decimal_256: ArrayRef = Arc::new( + Decimal256Array::from(vec![ + i256::from_i128(1), + i256::from_i128(2), + i256::from_i128(2), + i256::from_i128(3), + i256::from_i128(3), + i256::from_i128(3), + i256::from_i128(0), + i256::from_i128(0), + big_256_a, + big_256_b, + big_256_b, + ]) + .with_precision_and_scale(40, 2) + .unwrap(), + ); + assert_count_numerical_acc_and_group_acc::(decimal_256, 6); + } + + #[test] + fn interval_support_numerical_acc_and_group_acc() { + let year_month: ArrayRef = + Arc::new(IntervalYearMonthArray::from(vec![1, 2, 2, 3, 3, 3, 0, 0])); + assert_count_numerical_acc_and_group_acc::( + year_month, 4, + ); + + let day_time: ArrayRef = Arc::new(IntervalDayTimeArray::from(vec![ + IntervalDayTime::new(1, 0), + IntervalDayTime::new(1, 0), + IntervalDayTime::new(1, 5), + IntervalDayTime::new(2, 0), + ])); + assert_count_numerical_acc_and_group_acc::(day_time, 3); + + let month_day_nano: ArrayRef = + Arc::new(IntervalMonthDayNanoArray::from(vec![ + IntervalMonthDayNano::new(1, 0, 0), + IntervalMonthDayNano::new(1, 0, 0), + IntervalMonthDayNano::new(1, 0, 5), + IntervalMonthDayNano::new(0, 2, 0), + IntervalMonthDayNano::new(0, 0, 0), + ])); + assert_count_numerical_acc_and_group_acc::( + month_day_nano, + 4, + ); + } + + /// `approx_distinct(v) FILTER (WHERE nullable_bool)` — a NULL filter row + /// must not be counted (null filter is treated the same as false). + #[test] + fn update_batch_nullable_filter_excludes_null_filter_rows() { + let values: ArrayRef = Arc::new(Int64Array::from(vec![1i64, 2, 3, 4, 5])); + // row 0: filter=true, row 1: filter=NULL, row 2: filter=false, + // row 3: filter=NULL, row 4: filter=true + let filter = + BooleanArray::from(vec![Some(true), None, Some(false), None, Some(true)]); + + let mut acc = HllGroupsAccumulator::new(); + // put all rows in group 0 + let group_indices = vec![0usize; 5]; + acc.update_batch(&[values], &group_indices, Some(&filter), 1) + .unwrap(); + + // Only rows 0 and 4 (values 1 and 5) should be counted. + let result = acc.evaluate(EmitTo::All).unwrap(); + let counts = result.as_any().downcast_ref::().unwrap(); + // reference: hash 1 and 5 into a dense sketch + let expected = reference_count(&[h(1), h(5)]); + assert_eq!(counts.value(0), expected); + } + + #[test] + fn groups_convert_to_state_roundtrips_through_merge() { + let values: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + Some(2), + Some(2), + None, + Some(3), + ])); + let filter = BooleanArray::from(vec![ + Some(true), + Some(true), + Some(true), + Some(true), + None, + ]); + let group_indices = vec![0usize, 1, 0, 1, 0]; + + let mut direct = HllGroupsAccumulator::new(); + direct + .update_batch( + std::slice::from_ref(&values), + &group_indices, + Some(&filter), + 2, + ) + .unwrap(); + let direct = direct + .evaluate(EmitTo::All) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + + let converter = HllGroupsAccumulator::new(); + let state = converter + .convert_to_state(std::slice::from_ref(&values), Some(&filter)) + .unwrap(); + assert_eq!(state[0].null_count(), 0); + let mut merged = HllGroupsAccumulator::new(); + merged.merge_batch(&state, &group_indices, 2).unwrap(); + let merged = merged + .evaluate(EmitTo::All) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + + assert_eq!(direct, merged); + } + + #[test] + fn groups_convert_to_state_preserves_empty_and_filtered_rows() { + let converter = HllGroupsAccumulator::new(); + let empty_values: ArrayRef = + Arc::new(Int64Array::from(Vec::>::new())); + let state = converter + .convert_to_state(std::slice::from_ref(&empty_values), None) + .unwrap(); + assert_eq!(state[0].len(), 0); + assert_eq!(state[0].null_count(), 0); + + let values: ArrayRef = + Arc::new(Int64Array::from(vec![Some(1), Some(2), None])); + let filter = BooleanArray::from(vec![Some(false), None, Some(false)]); + let group_indices = vec![0usize, 1, 0]; + let state = converter + .convert_to_state(std::slice::from_ref(&values), Some(&filter)) + .unwrap(); + assert_eq!(state[0].len(), values.len()); + assert_eq!(state[0].null_count(), 0); + let state = state[0].as_any().downcast_ref::().unwrap(); + for row in 0..state.len() { + assert_eq!(state.value(row), b""); + } + + let mut merged = HllGroupsAccumulator::new(); + merged + .merge_batch(&[Arc::new(state.clone())], &group_indices, 2) + .unwrap(); + let result = merged + .evaluate(EmitTo::All) + .unwrap() + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_eq!(result, UInt64Array::from(vec![0, 0])); + } + + /// Regression: a short (≤ 12-byte) Utf8View string must hash identically + /// in an all-inline batch and in a mixed batch that also contains a long + /// string (which forces a data buffer). + #[test] + fn utf8view_groups_short_string_hashed_consistently_across_batches() { + // Batch 1: all-inline (no data buffers) — "aaa" is hashed as u128 view. + let batch1: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", "bbb"])); + assert!(batch1.as_string_view().data_buffers().is_empty()); + + // Batch 2: mixed — LONG forces a data buffer; "aaa" must still be + // hashed as u128 view so it matches its appearance in batch 1. + let batch2: ArrayRef = Arc::new(StringViewArray::from(vec!["aaa", LONG])); + assert!(!batch2.as_string_view().data_buffers().is_empty()); + + let group_indices = vec![0usize, 0]; + let mut acc = HllGroupsAccumulator::new(); + acc.update_batch(&[batch1], &group_indices, None, 1) + .unwrap(); + acc.update_batch(&[batch2], &group_indices, None, 1) + .unwrap(); + + // True distinct values: {"aaa", "bbb", LONG} == 3. + let result = acc.evaluate(EmitTo::All).unwrap(); + let counts = result.as_any().downcast_ref::().unwrap(); + assert_eq!(counts.value(0), 3); + } + + /// Regression: a short (≤ 12-byte) Utf8View string must hash identically + /// regardless of which batch it appears in — all-inline or mixed. + #[test] + fn utf8view_acc_split_batches_match_single_mixed_batch() { + // Multiset: {"aaa" x2, "bbb", LONG}, so 3 distinct values. + let mixed: ArrayRef = + Arc::new(StringViewArray::from(vec!["aaa", "bbb", LONG, "aaa"])); + let mut acc_single = HLLAccumulator::new(); + acc_single.update_batch(&[mixed]).unwrap(); + + // Same multiset, but split so "aaa" lands in both an all-inline batch + // and a batch with a data buffer (forced by LONG). + let inline_only: ArrayRef = + Arc::new(StringViewArray::from(vec!["aaa", "bbb"])); + let with_buffer: ArrayRef = + Arc::new(StringViewArray::from(vec!["aaa", LONG])); + assert!(inline_only.as_string_view().data_buffers().is_empty()); + assert!(!with_buffer.as_string_view().data_buffers().is_empty()); + + let mut acc_split = HLLAccumulator::new(); + acc_split.update_batch(&[inline_only]).unwrap(); + acc_split.update_batch(&[with_buffer]).unwrap(); + + assert_eq!( + distinct_count(&mut acc_single), + distinct_count(&mut acc_split) + ); + assert_eq!(distinct_count(&mut acc_single), 3); + } + } + + fn h(v: u64) -> u64 { + HLL_HASH_STATE.hash_one(v) + } + + /// Reference count: fold the given distinct hashes straight into a dense + /// HyperLogLog. The grouped sketch must agree with this exactly. + fn reference_count(hashes: &[u64]) -> u64 { + let mut hll = HyperLogLog::::new(); + for &hash in hashes { + hll.add_hashed(hash); + } + hll.count() as u64 + } + + fn serialize(g: &mut GroupHll) -> Vec { + let mut buf = Vec::new(); + g.serialize(&mut buf); + buf + } + + #[test] + fn sparse_stays_sparse_for_small_groups() { + let mut g = GroupHll::default(); + let hashes: Vec = (0..50).map(h).collect(); + for &hash in &hashes { + g.add_hash(hash); + } + // duplicates must not change the estimate or trigger promotion + for &hash in &hashes { + g.add_hash(hash); + } + assert!( + matches!(g, GroupHll::Sparse(_)), + "small group must be sparse" + ); + assert_eq!(g.count(), reference_count(&hashes)); + // sparse serialized state is far smaller than a dense 16 KiB sketch + // and must not exceed the sparse limit contract enforced by merge_serialized + let serialized = serialize(&mut g); + assert!(serialized.len() < NUM_REGISTERS); + assert!(serialized.len() <= SPARSE_LIMIT * size_of::()); + } + + #[test] + fn promotes_to_dense_for_large_groups() { + let mut g = GroupHll::default(); + let hashes: Vec = (0..(SPARSE_LIMIT as u64 * 4)).map(h).collect(); + for &hash in &hashes { + g.add_hash(hash); + } + assert!(matches!(g, GroupHll::Dense(_)), "large group must be dense"); + assert_eq!(g.count(), reference_count(&hashes)); + } + + #[test] + fn serialize_then_merge_roundtrips() { + for n in [0u64, 10, SPARSE_LIMIT as u64 * 4] { + let hashes: Vec = (0..n).map(h).collect(); + let mut src = GroupHll::default(); + for &hash in &hashes { + src.add_hash(hash); + } + let bytes = serialize(&mut src); + let mut dst = GroupHll::default(); + dst.merge_serialized(&bytes).unwrap(); + assert_eq!(dst.count(), reference_count(&hashes), "n = {n}"); + } + } + + #[test] + fn sparse_limit_group_serializes_as_mergeable_sparse_state() { + let hashes: Vec = (0..SPARSE_LIMIT as u64).map(h).collect(); + let mut src = GroupHll::default(); + for &hash in &hashes { + src.add_hash(hash); + } + assert!(matches!(src, GroupHll::Sparse(_))); + + let bytes = serialize(&mut src); + assert_eq!(bytes.len(), SPARSE_LIMIT * size_of::()); + + let mut dst = GroupHll::default(); + dst.merge_serialized(&bytes).unwrap(); + assert_eq!(dst.count(), reference_count(&hashes)); + } + + #[test] + fn medium_sparse_group_serializes_as_mergeable_dense_state() { + let n = SPARSE_LIMIT as u64 + 44; + let hashes: Vec = (0..n).map(h).collect(); + let mut src = GroupHll::default(); + for &hash in &hashes { + src.add_hash(hash); + } + assert!( + matches!(src, GroupHll::Sparse(_)), + "group should not promote during update before the compaction threshold" + ); + + let bytes = serialize(&mut src); + assert_eq!(bytes.len(), NUM_REGISTERS); + + let mut dst = GroupHll::default(); + dst.merge_serialized(&bytes).unwrap(); + assert_eq!(dst.count(), reference_count(&hashes)); + } + + #[test] + fn merge_combines_disjoint_groups() { + // sparse + sparse, sparse + dense, dense + dense + let left: Vec = (0..100).map(h).collect(); + let right: Vec = (100..(SPARSE_LIMIT as u64 * 4)).map(h).collect(); + let all: Vec = left.iter().chain(right.iter()).copied().collect(); + + let mut a = GroupHll::default(); + for &hash in &left { + a.add_hash(hash); + } + let mut b = GroupHll::default(); + for &hash in &right { + b.add_hash(hash); + } + let b_bytes = serialize(&mut b); + a.merge_serialized(&b_bytes).unwrap(); + assert_eq!(a.count(), reference_count(&all)); + } + + #[test] + fn empty_group_counts_zero() { + let mut g = GroupHll::default(); + assert_eq!(g.count(), 0); + let bytes = serialize(&mut g); + assert!(bytes.is_empty()); + let mut dst = GroupHll::default(); + dst.merge_serialized(&bytes).unwrap(); + assert_eq!(dst.count(), 0); + } +} diff --git a/datafusion/functions-aggregate/src/approx_percentile_cont.rs b/datafusion/functions-aggregate/src/approx_percentile_cont.rs index 3f1adcca12362..ea8fea1b1bc29 100644 --- a/datafusion/functions-aggregate/src/approx_percentile_cont.rs +++ b/datafusion/functions-aggregate/src/approx_percentile_cont.rs @@ -301,6 +301,13 @@ impl AggregateUDFImpl for ApproxPercentileCont { } fn return_type(&self, arg_types: &[DataType]) -> Result { + // Defensive: the public signature already restricts callers to 2 or 3 + // arguments. This guards against aggregate planning accidentally + // feeding state-field types (e.g. from `PartialReduce`) back into + // `return_type`, which would otherwise silently choose the wrong type. + if arg_types.len() > 3 { + return plan_err!("approx_percentile_cont requires at most 3 arguments"); + } if !arg_types[0].is_numeric() { return plan_err!("approx_percentile_cont requires numeric input types"); } diff --git a/datafusion/functions-aggregate/src/array_agg.rs b/datafusion/functions-aggregate/src/array_agg.rs index 24edaaff1f09d..0e02ff118678f 100644 --- a/datafusion/functions-aggregate/src/array_agg.rs +++ b/datafusion/functions-aggregate/src/array_agg.rs @@ -18,7 +18,7 @@ //! `ARRAY_AGG` aggregate implementation: [`ArrayAgg`] use std::cmp::Ordering; -use std::collections::{HashSet, VecDeque}; +use std::collections::VecDeque; use std::mem::{size_of, size_of_val, take}; use std::sync::Arc; @@ -27,14 +27,19 @@ use arrow::array::{ UInt32Array, new_empty_array, }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; -use arrow::compute::{SortOptions, filter}; +use arrow::compute::{SortOptions, cast, filter}; use arrow::datatypes::{DataType, Field, FieldRef, Fields}; +use arrow::row::{OwnedRow, Row, RowConverter, Rows, SortField}; use datafusion_common::cast::as_list_array; +use datafusion_common::hash_utils::{RandomState, create_hashes}; +use datafusion_common::utils::proxy::HashTableAllocExt; use datafusion_common::utils::{ SingleRowListArrayBuilder, compare_rows, get_row_at_idx, take_function_args, }; -use datafusion_common::{Result, ScalarValue, assert_eq_or_internal_err, exec_err}; +use datafusion_common::{ + Result, ScalarValue, assert_eq_or_internal_err, exec_err, internal_err, +}; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ @@ -47,6 +52,7 @@ use datafusion_functions_aggregate_common::order::AggregateOrderSensitivity; use datafusion_functions_aggregate_common::utils::ordering_fields; use datafusion_macros::user_doc; use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; +use hashbrown::hash_table::HashTable; make_udaf_expr_and_func!( ArrayAgg, @@ -730,7 +736,6 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "one argument to merge_batch"); @@ -792,11 +797,6 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator { Ok(vec![Arc::new(list_array)]) } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn size(&self) -> usize { self.batches .iter() @@ -812,14 +812,67 @@ impl GroupsAccumulator for ArrayAggGroupsAccumulator { } } +/// Resources that are allocated lazily on the first `update_batch` call, +/// once the concrete runtime Arrow type is known. +/// +/// Grouping all three fields together makes the "either all present or all +/// absent" invariant explicit in the type system, replacing the scattered +/// `.expect()` calls that would otherwise be needed. +#[derive(Debug)] +struct DistinctState { + /// Converts Arrow arrays to/from the comparable row format. + converter: RowConverter, + /// One owned encoded row per live distinct value, indexed by group index. + /// Compacted via swap-remove on eviction so there are never dead slots. + group_rows: Vec, + /// Live refcount per group index. `counts[i]` is how many times the value + /// at `group_rows[i]` is currently present in the window frame. + counts: Vec, + /// Hash of the encoded row at group index `i`, kept in sync with + /// `group_rows` and `counts`. Needed to patch the map on swap-remove + /// eviction without re-encoding the moved row. + row_hashes: Vec, + /// Temporary buffer for encoding an incoming batch; reused across calls. + rows_buffer: Rows, +} + #[derive(Debug)] pub struct DistinctArrayAggAccumulator { - values: HashSet, + /// Lazily allocated on the first `update_batch`; `None` until then. + state: Option, + /// Hash table storing `(hash, group_index)`. Only contains live entries + /// (those whose count is > 0). Evicted on `retract_batch` when count + /// drops to zero. + map: HashTable<(u64, usize)>, + /// Heap size of `map` in bytes, tracked for `size()` reporting. + map_size: usize, + /// Reused buffer for batch hashes. + hashes_buffer: Vec, + /// Random state used by `create_hashes`. + random_state: RandomState, datatype: DataType, sort_options: Option, ignore_nulls: bool, } +/// Returns `true` if `dt` is, or recursively contains, a `Dictionary` type. +/// +/// `RowConverter` always decodes to the physical (non-dictionary) type, so a +/// cast back to the declared logical type is required when this is true. +fn datatype_contains_dictionary(dt: &DataType) -> bool { + match dt { + DataType::Dictionary(_, _) => true, + DataType::List(f) + | DataType::LargeList(f) + | DataType::FixedSizeList(f, _) + | DataType::Map(f, _) => datatype_contains_dictionary(f.data_type()), + DataType::Struct(fields) => fields + .iter() + .any(|f| datatype_contains_dictionary(f.data_type())), + _ => false, + } +} + impl DistinctArrayAggAccumulator { pub fn try_new( datatype: &DataType, @@ -827,12 +880,37 @@ impl DistinctArrayAggAccumulator { ignore_nulls: bool, ) -> Result { Ok(Self { - values: HashSet::new(), + state: None, + map: HashTable::new(), + map_size: 0, + hashes_buffer: Vec::new(), + random_state: RandomState::default(), datatype: datatype.clone(), sort_options, ignore_nulls, }) } + + /// Lazily initialises the `DistinctState` on the first call, using the + /// actual runtime column type. + fn ensure_state(&mut self, data_type: &DataType) -> Result<()> { + if self.state.is_none() { + let sort_field = match self.sort_options { + Some(opts) => SortField::new_with_options(data_type.clone(), opts), + None => SortField::new(data_type.clone()), + }; + let converter = RowConverter::new(vec![sort_field])?; + let rows_buffer = converter.empty_rows(0, 0); + self.state = Some(DistinctState { + converter, + group_rows: Vec::new(), + counts: Vec::new(), + row_hashes: Vec::new(), + rows_buffer, + }); + } + Ok(()) + } } impl Accumulator for DistinctArrayAggAccumulator { @@ -846,22 +924,76 @@ impl Accumulator for DistinctArrayAggAccumulator { } let val = &values[0]; - let nulls = if self.ignore_nulls { - val.logical_nulls() + + // Filter nulls out upfront when ignore_nulls is set so they are + // never inserted into the dedup state. + let filtered; + let col: &ArrayRef = if self.ignore_nulls { + if let Some(nulls) = val.logical_nulls() { + if nulls.null_count() > 0 { + let mask: BooleanArray = nulls.iter().map(Some).collect(); + filtered = filter(val.as_ref(), &mask)?; + &filtered + } else { + val + } + } else { + val + } } else { - None + val }; - let nulls = nulls.as_ref(); - if nulls.is_none_or(|nulls| nulls.null_count() < val.len()) { - for i in 0..val.len() { - if nulls.is_none_or(|nulls| nulls.is_valid(i)) { - self.values - .insert(ScalarValue::try_from_array(val, i)?.compacted()); + if col.is_empty() { + return Ok(()); + } + + self.ensure_state(col.data_type())?; + + // Encode the entire incoming batch into rows_buffer in one pass. + let DistinctState { + converter, + group_rows, + counts, + row_hashes, + rows_buffer, + } = self.state.as_mut().unwrap(); + rows_buffer.clear(); + converter.append(rows_buffer, std::slice::from_ref(col))?; + + // Pre-compute all hashes for the batch in one SIMD-friendly pass. + self.hashes_buffer.clear(); + self.hashes_buffer.resize(col.len(), 0); + create_hashes( + std::slice::from_ref(col), + &self.random_state, + &mut self.hashes_buffer, + )?; + + for (row_idx, &hash) in self.hashes_buffer.iter().enumerate() { + let row = rows_buffer.row(row_idx); + let entry = self.map.find_mut(hash, |&(h, group_idx)| { + h == hash && group_rows[group_idx].row() == row + }); + match entry { + Some((_, group_idx)) => { + // Already known: just increment the live refcount. + counts[*group_idx] += 1; + } + None => { + // New distinct value: own the encoded row, record it. + let new_group_idx = group_rows.len(); + group_rows.push(row.owned()); + counts.push(1); + row_hashes.push(hash); + self.map.insert_accounted( + (hash, new_group_idx), + |&(h, _)| h, + &mut self.map_size, + ); } } } - Ok(()) } @@ -872,6 +1004,7 @@ impl Accumulator for DistinctArrayAggAccumulator { assert_eq_or_internal_err!(states.len(), 1, "expects single state"); + // The DISTINCT state is `List`. states[0] .as_list::() .iter() @@ -880,49 +1013,180 @@ impl Accumulator for DistinctArrayAggAccumulator { } fn evaluate(&mut self) -> Result { - let mut values: Vec = self.values.iter().cloned().collect(); - if values.is_empty() { + if self.map.is_empty() { return Ok(ScalarValue::new_null_list(self.datatype.clone(), true, 1)); } - if let Some(opts) = self.sort_options { - let mut delayed_cmp_err = Ok(()); - values.sort_by(|a, b| { - if a.is_null() { - return match opts.nulls_first { - true => Ordering::Less, - false => Ordering::Greater, - }; - } - if b.is_null() { - return match opts.nulls_first { - true => Ordering::Greater, - false => Ordering::Less, - }; - } - match opts.descending { - true => b.try_cmp(a), - false => a.try_cmp(b), - } - .unwrap_or_else(|err| { - delayed_cmp_err = Err(err); - Ordering::Equal - }) - }); - delayed_cmp_err?; + let DistinctState { + converter, + group_rows, + .. + } = self + .state + .as_ref() + .expect("state must be set when map is non-empty"); + + // Collect the group indices of all live entries. + let mut live_indices: Vec = + self.map.iter().map(|&(_, group_idx)| group_idx).collect(); + + // If ORDER BY was specified, the RowConverter bakes the sort direction + // into the row bytes, so lexicographic sort gives the correct order. + if self.sort_options.is_some() { + live_indices + .sort_unstable_by(|&a, &b| group_rows[a].row().cmp(&group_rows[b].row())); + } + + // Decode the selected rows back into an Arrow array. + let rows: Vec> = + live_indices.iter().map(|&i| group_rows[i].row()).collect(); + let arrays = converter.convert_rows(rows)?; + + // `convert_rows` always returns the physical (non-dictionary) type. + // Cast back to the declared logical type when they differ AND the + // declared type contains a Dictionary somewhere (directly or nested + // inside a Struct, List, etc.) — that is the only case where + // RowConverter strips the logical type. + let decoded = if arrays[0].data_type() != &self.datatype + && datatype_contains_dictionary(&self.datatype) + { + cast(arrays[0].as_ref(), &self.datatype)? + } else { + Arc::clone(&arrays[0]) }; + let values: Vec = (0..decoded.len()) + .map(|i| ScalarValue::try_from_array(decoded.as_ref(), i)) + .collect::>()?; + let arr = ScalarValue::new_list(&values, &self.datatype, true); Ok(ScalarValue::List(arr)) } + fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + if values.is_empty() { + return Ok(()); + } + + assert_eq_or_internal_err!(values.len(), 1, "expects single batch"); + + let val = &values[0]; + + // Mirror the null-filtering logic from update_batch so we only + // retract values that were actually inserted. + let filtered; + let col: &ArrayRef = if self.ignore_nulls { + if let Some(nulls) = val.logical_nulls() { + if nulls.null_count() > 0 { + let mask: BooleanArray = nulls.iter().map(Some).collect(); + filtered = filter(val.as_ref(), &mask)?; + &filtered + } else { + val + } + } else { + val + } + } else { + val + }; + + if col.is_empty() { + return Ok(()); + } + + let DistinctState { + converter, + group_rows, + counts, + row_hashes, + rows_buffer, + } = self + .state + .as_mut() + .expect("retract_batch called before update_batch"); + + rows_buffer.clear(); + converter.append(rows_buffer, std::slice::from_ref(col))?; + + self.hashes_buffer.clear(); + self.hashes_buffer.resize(col.len(), 0); + create_hashes( + std::slice::from_ref(col), + &self.random_state, + &mut self.hashes_buffer, + )?; + + for (row_idx, &hash) in self.hashes_buffer.iter().enumerate() { + let row = rows_buffer.row(row_idx); + match self.map.find_entry(hash, |&(h, group_idx)| { + h == hash && group_rows[group_idx].row() == row + }) { + Err(_) => { + return internal_err!( + "DistinctArrayAggAccumulator::retract_batch: \ + value not present in state" + ); + } + Ok(occupied) => { + let (_, dead_idx) = *occupied.get(); + counts[dead_idx] -= 1; + if counts[dead_idx] == 0 { + occupied.remove(); + // Compact via swap-remove: move the last slot into the + // dead slot so group_rows / counts / row_hashes stay + // dense with no dead entries. + let last_idx = group_rows.len() - 1; + if dead_idx != last_idx { + // Patch the map entry that points to last_idx so + // it points to dead_idx instead. + let last_hash = row_hashes[last_idx]; + self.map + .find_mut(last_hash, |&(_, idx)| idx == last_idx) + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "DistinctArrayAggAccumulator: map is missing \ + group index {last_idx} during swap-remove \ + compaction" + ) + })? + .1 = dead_idx; + } + group_rows.swap_remove(dead_idx); + counts.swap_remove(dead_idx); + row_hashes.swap_remove(dead_idx); + } + } + } + } + Ok(()) + } + + fn supports_retract_batch(&self) -> bool { + true + } + fn size(&self) -> usize { - size_of_val(self) + ScalarValue::size_of_hashset(&self.values) - - size_of_val(&self.values) + size_of_val(self) + + self + .state + .as_ref() + .map(|s| { + s.group_rows + .iter() + .map(|r| r.row().data().len()) + .sum::() + + s.group_rows.capacity() * size_of::() + + s.counts.capacity() * size_of::() + + s.row_hashes.capacity() * size_of::() + + s.rows_buffer.size() + + s.converter.size() + }) + .unwrap_or(0) + + self.map_size + + self.hashes_buffer.capacity() * size_of::() + self.datatype.size() - size_of_val(&self.datatype) - - size_of_val(&self.sort_options) - + size_of::>() } } @@ -1008,7 +1272,13 @@ impl OrderSensitiveArrayAggAccumulator { } else { (0..fields.len()) .map(|i| { - let column_values = self.ordering_values.iter().map(|x| x[i].clone()); + let column_values: Box> = if self + .reverse + { + Box::new(self.ordering_values.iter().rev().map(|x| x[i].clone())) + } else { + Box::new(self.ordering_values.iter().map(|x| x[i].clone())) + }; ScalarValue::iter_to_array(column_values) }) .collect::>()? @@ -1471,15 +1741,17 @@ mod tests { acc2.update_batch(&[data(["b", "c", "a"])])?; acc1 = merge(acc1, acc2)?; - assert_eq!(acc1.size(), 282); + assert_eq!(acc1.size(), 174); Ok(()) } #[test] fn does_not_over_account_memory_distinct() -> Result<()> { - let (mut acc1, mut acc2) = ArrayAggAccumulatorBuilder::string() - .distinct() - .build_two()?; + let (mut acc1, mut acc2) = ArrayAggAccumulatorBuilder::new(DataType::List( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + )) + .distinct() + .build_two()?; acc1.update_batch(&[string_list_data([ vec!["a", "b", "c"], @@ -1488,17 +1760,18 @@ mod tests { acc2.update_batch(&[string_list_data([vec!["e", "f", "g"]])])?; acc1 = merge(acc1, acc2)?; - // without compaction, the size is 16660 - assert_eq!(acc1.size(), 1660); + assert_eq!(acc1.size(), 2274); Ok(()) } #[test] fn does_not_over_account_memory_ordered() -> Result<()> { - let mut acc = ArrayAggAccumulatorBuilder::string() - .order_by_col("col", SortOptions::new(false, false)) - .build()?; + let mut acc = ArrayAggAccumulatorBuilder::new(DataType::List(Arc::new( + Field::new_list_field(DataType::Utf8, true), + ))) + .order_by_col("col", SortOptions::new(false, false)) + .build()?; acc.update_batch(&[string_list_data([ vec!["a", "b", "c"], @@ -1512,6 +1785,231 @@ mod tests { Ok(()) } + #[test] + fn ordered_aggregate_nested_nullability_mismatch_issue_24022() -> Result<()> { + use arrow::array::{Int32Array, Int64Array, StructArray}; + use datafusion_physical_expr::expressions::Column; + + let requested_element_type = + DataType::Struct(Fields::from(vec![Field::new("n", DataType::Int32, true)])); + let inferred_field = Field::new("n", DataType::Int32, false); + + let ordering_dtype = DataType::Int64; + let schema = Schema::new(vec![ + Field::new("val", requested_element_type.clone(), true), + Field::new("ord", DataType::Int64, true), + ]); + let ord_expr = Arc::new( + Column::new_with_schema("ord", &schema).expect("column not in schema"), + ) as Arc; + + let asc_opts = SortOptions { + descending: false, + nulls_first: false, + }; + let asc_ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::clone(&ord_expr), + asc_opts, + )]) + .unwrap(); + + let mut acc = OrderSensitiveArrayAggAccumulator::try_new( + &requested_element_type, + std::slice::from_ref(&ordering_dtype), + asc_ordering, + /*is_input_pre_ordered=*/ true, + /*reverse=*/ false, + /*ignore_nulls=*/ false, + )?; + + let value_arr = Arc::new(StructArray::from(vec![( + Arc::new(inferred_field), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )])) as ArrayRef; + + let ord_arr = Arc::new(Int64Array::from(vec![0i64])) as ArrayRef; + + acc.update_batch(&[value_arr, ord_arr])?; + + let evaluated = acc.evaluate()?; + + if let ScalarValue::List(arr) = evaluated { + assert_eq!( + arr.data_type(), + &DataType::List(Arc::new(Field::new_list_field( + requested_element_type.clone(), + true + ))) + ); + + let expected_struct_array = StructArray::from(vec![( + Arc::new(Field::new("n", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]); + let expected_array = Arc::new(expected_struct_array) as ArrayRef; + assert_eq!(&arr.value(0), &expected_array); + } else { + panic!("Expected ScalarValue::List"); + } + + Ok(()) + } + + #[test] + fn distinct_aggregate_nested_nullability_mismatch_issue_24022() -> Result<()> { + use arrow::array::{Int32Array, StructArray}; + use datafusion_common::ScalarValue; + + let requested_element_type = + DataType::Struct(Fields::from(vec![Field::new("n", DataType::Int32, true)])); + let inferred_field = Field::new("n", DataType::Int32, false); + + let mut acc = DistinctArrayAggAccumulator::try_new( + &requested_element_type, + None, + /*ignore_nulls=*/ false, + )?; + + let value_arr = Arc::new(StructArray::from(vec![( + Arc::new(inferred_field), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )])) as ArrayRef; + + acc.update_batch(&[value_arr])?; + + let evaluated = acc.evaluate()?; + + if let ScalarValue::List(arr) = evaluated { + assert_eq!( + arr.data_type(), + &DataType::List(Arc::new(Field::new_list_field( + requested_element_type.clone(), + true + ))) + ); + + let expected_struct_array = StructArray::from(vec![( + Arc::new(Field::new("n", DataType::Int32, true)), + Arc::new(Int32Array::from(vec![1])) as ArrayRef, + )]); + let expected_array = Arc::new(expected_struct_array) as ArrayRef; + assert_eq!(&arr.value(0), &expected_array); + } else { + panic!("Expected ScalarValue::List"); + } + + Ok(()) + } + + // Reproduces the bug where `state()` emits reversed values but non-reversed + // orderings when the optimizer sets is_input_pre_ordered=true + reverse=true + // (DESC aggregate with ASC pre-sorted input). The partial states are fed into + // a final accumulator via merge_batch; without the fix the ordering keys and + // values are mismatched so the final sort produces wrong order. + #[test] + fn desc_order_partial_final_merge_correct() -> Result<()> { + use arrow::array::Int64Array; + use datafusion_physical_expr::expressions::Column; + + let schema = Schema::new(vec![ + Field::new("val", DataType::Int64, true), + Field::new("ord", DataType::Int64, true), + ]); + let ord_expr = Arc::new( + Column::new_with_schema("ord", &schema).expect("column not in schema"), + ) as Arc; + + // ordering_req for partial = [ord ASC] (reversed, because input is pre-sorted ASC + // and the user wants DESC — the optimizer reverses the requirement) + let asc_opts = SortOptions { + descending: false, + nulls_first: false, + }; + let desc_opts = SortOptions { + descending: true, + nulls_first: false, + }; + + let asc_ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::clone(&ord_expr), + asc_opts, + )]) + .unwrap(); + let desc_ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + Arc::clone(&ord_expr), + desc_opts, + )]) + .unwrap(); + + let ordering_dtype = DataType::Int64; + + // Partial acc A: sees rows [0,1,2] arriving in ASC order (pre-ordered). + // is_input_pre_ordered=true, reverse=true, ordering_req=[ASC]. + let mut partial_a = OrderSensitiveArrayAggAccumulator::try_new( + &DataType::Int64, + std::slice::from_ref(&ordering_dtype), + asc_ordering.clone(), + /*is_input_pre_ordered=*/ true, + /*reverse=*/ true, + /*ignore_nulls=*/ false, + )?; + let vals_a = Arc::new(Int64Array::from(vec![0i64, 1, 2])) as ArrayRef; + let ords_a = Arc::new(Int64Array::from(vec![0i64, 1, 2])) as ArrayRef; + partial_a.update_batch(&[vals_a, ords_a])?; + let state_a = partial_a + .state()? + .iter() + .map(|v| v.to_array()) + .collect::>>()?; + + // Partial acc B: sees rows [3,4,5] arriving in ASC order. + let mut partial_b = OrderSensitiveArrayAggAccumulator::try_new( + &DataType::Int64, + std::slice::from_ref(&ordering_dtype), + asc_ordering, + /*is_input_pre_ordered=*/ true, + /*reverse=*/ true, + /*ignore_nulls=*/ false, + )?; + let vals_b = Arc::new(Int64Array::from(vec![3i64, 4, 5])) as ArrayRef; + let ords_b = Arc::new(Int64Array::from(vec![3i64, 4, 5])) as ArrayRef; + partial_b.update_batch(&[vals_b, ords_b])?; + let state_b = partial_b + .state()? + .iter() + .map(|v| v.to_array()) + .collect::>>()?; + + // Final acc: not optimized — ordering_req=[DESC], reverse=false. + let mut final_acc = OrderSensitiveArrayAggAccumulator::try_new( + &DataType::Int64, + std::slice::from_ref(&ordering_dtype), + desc_ordering, + /*is_input_pre_ordered=*/ false, + /*reverse=*/ false, + /*ignore_nulls=*/ false, + )?; + final_acc.merge_batch(&state_a)?; + final_acc.merge_batch(&state_b)?; + let result = final_acc.evaluate()?; + + let ScalarValue::List(list) = result else { + return datafusion_common::internal_err!("expected List"); + }; + let result_vals: Vec = list + .values() + .as_any() + .downcast_ref::() + .unwrap() + .iter() + .map(|v| v.unwrap()) + .collect(); + + // Expected DESC: [5, 4, 3, 2, 1, 0] + assert_eq!(result_vals, vec![5i64, 4, 3, 2, 1, 0]); + Ok(()) + } + struct ArrayAggAccumulatorBuilder { return_field: FieldRef, distinct: bool, @@ -1526,15 +2024,19 @@ mod tests { fn new(data_type: DataType) -> Self { Self { - return_field: Field::new("f", data_type.clone(), true).into(), + return_field: Field::new( + "f", + DataType::List(Arc::new(Field::new_list_field( + data_type.clone(), + true, + ))), + true, + ) + .into(), distinct: false, order_bys: vec![], schema: Schema { - fields: Fields::from(vec![Field::new( - "col", - DataType::new_list(data_type, true), - true, - )]), + fields: Fields::from(vec![Field::new("col", data_type, true)]), metadata: Default::default(), }, } @@ -1851,7 +2353,7 @@ mod tests { // Merge acc2's state into acc1 let state = acc2.state(EmitTo::All)?; - acc1.merge_batch(&state, &[0, 1], None, 2)?; + acc1.merge_batch(&state, &[0, 1], 2)?; // Another update_batch on acc1 after the merge let values: ArrayRef = Arc::new(Int32Array::from(vec![5, 6])); @@ -1920,7 +2422,7 @@ mod tests { // Feed state into a new accumulator via merge_batch let mut acc2 = ArrayAggGroupsAccumulator::new(DataType::Int32, false); - acc2.merge_batch(&state, &[0, 0, 1], None, 2)?; + acc2.merge_batch(&state, &[0, 0, 1], 2)?; // Group 0 received rows 0 ([1]) and 1 ([NULL]) → [1, NULL] let vals = eval_i32_lists(&mut acc2, EmitTo::All)?; @@ -1950,7 +2452,7 @@ mod tests { // Feed state into a new accumulator via merge_batch let mut acc2 = ArrayAggGroupsAccumulator::new(DataType::Int32, true); - acc2.merge_batch(&state, &[0, 0, 1, 1], None, 2)?; + acc2.merge_batch(&state, &[0, 0, 1, 1], 2)?; // Group 0: received [1] and null (skipped) → [1] let vals = eval_i32_lists(&mut acc2, EmitTo::All)?; @@ -2300,4 +2802,331 @@ mod tests { Ok(()) } + + // ---- DistinctArrayAggAccumulator retract_batch tests ---- + + // Build a DISTINCT accumulator with ascending sort so evaluate output is + // deterministic regardless of HashMap iteration order. + fn distinct_acc(ignore_nulls: bool) -> Result { + DistinctArrayAggAccumulator::try_new( + &DataType::Utf8, + Some(SortOptions::default()), + ignore_nulls, + ) + } + + #[test] + fn distinct_retract_duplicate_remains() -> Result<()> { + // Canonical regression for the HashSet-can't-retract bug: a value + // that appears multiple times in-frame must survive retraction of + // a single occurrence. + let mut acc = distinct_acc(false)?; + + // Feed [A, A, B] across two batches to exercise multi-batch state. + acc.update_batch(&[data(["A", "A"])])?; + acc.update_batch(&[data(["B"])])?; + assert_eq!(print_nulls(str_arr(acc.evaluate()?)?), vec!["A", "B"]); + + // Retract a single A — the other A is still in the frame. + acc.retract_batch(&[data(["A"])])?; + assert_eq!(print_nulls(str_arr(acc.evaluate()?)?), vec!["A", "B"]); + + // Retract the remaining A — only B left. + acc.retract_batch(&[data(["A"])])?; + assert_eq!(print_nulls(str_arr(acc.evaluate()?)?), vec!["B"]); + + Ok(()) + } + + #[test] + fn distinct_retract_full_removal() -> Result<()> { + let mut acc = distinct_acc(false)?; + + acc.update_batch(&[data(["A", "B"])])?; + acc.retract_batch(&[data(["A", "B"])])?; + + let result = acc.evaluate()?; + assert!( + matches!(&result, ScalarValue::List(arr) if arr.is_null(0)), + "expected null list after full retract, got {result:?}" + ); + + Ok(()) + } + + #[test] + fn distinct_retract_ignore_nulls_skips() -> Result<()> { + // ignore_nulls=true: NULL never enters state on update, so retract + // must also skip NULL — otherwise we'd error on the missing key. + let mut acc = distinct_acc(true)?; + + acc.update_batch(&[data([Some("A"), None, Some("B")])])?; + assert_eq!(print_nulls(str_arr(acc.evaluate()?)?), vec!["A", "B"]); + + // Retract [A, NULL] — the NULL is skipped, only A is removed. + acc.retract_batch(&[data([Some("A"), None])])?; + assert_eq!(print_nulls(str_arr(acc.evaluate()?)?), vec!["B"]); + + Ok(()) + } + + #[test] + fn distinct_retract_null_tracked() -> Result<()> { + // ignore_nulls=false: NULL enters state with a refcount and must + // retract symmetrically; the NULL key must be removed at zero + // (else evaluate still emits a NULL element). + let mut acc = distinct_acc(false)?; + + acc.update_batch(&[data([Some("A"), None, None])])?; + // With nulls_first=true (SortOptions default), NULL sorts before A. + assert_eq!(print_nulls(str_arr(acc.evaluate()?)?), vec!["NULL", "A"]); + + // Retract one NULL — count drops to 1, key still present. + acc.retract_batch(&[data::, 1>([None])])?; + assert_eq!(print_nulls(str_arr(acc.evaluate()?)?), vec!["NULL", "A"]); + + // Retract the remaining NULL — key is removed. + acc.retract_batch(&[data::, 1>([None])])?; + assert_eq!(print_nulls(str_arr(acc.evaluate()?)?), vec!["A"]); + + Ok(()) + } + + #[test] + fn distinct_supports_retract_batch() -> Result<()> { + let acc = distinct_acc(false)?; + assert!(acc.supports_retract_batch()); + + let acc_ignore = distinct_acc(true)?; + assert!(acc_ignore.supports_retract_batch()); + + Ok(()) + } + + #[test] + fn distinct_merge_then_evaluate_regression() -> Result<()> { + // Non-window path: state -> merge_batch -> evaluate must still + // produce the union of distinct values across partitions. + let mut acc1 = distinct_acc(false)?; + let mut acc2 = distinct_acc(false)?; + + acc1.update_batch(&[data(["A", "A", "B"])])?; + acc2.update_batch(&[data(["A", "C"])])?; + + let state = acc2.state()?; + let state_arrs: Vec = state + .into_iter() + .map(|sv| sv.to_array_of_size(1)) + .collect::>>()?; + acc1.merge_batch(&state_arrs)?; + + assert_eq!(print_nulls(str_arr(acc1.evaluate()?)?), vec!["A", "B", "C"]); + + Ok(()) + } + + #[test] + fn distinct_array_agg_utf8_deduplicates() -> Result<()> { + use arrow::array::StringArray; + + // 7 rows with 4 distinct values, each duplicate appearing twice. + let input: ArrayRef = Arc::new(StringArray::from(vec![ + "postgres", "mysql", "postgres", "redis", "mysql", "duckdb", "redis", + ])); + + let mut acc = DistinctArrayAggAccumulator::try_new(&DataType::Utf8, None, false)?; + acc.update_batch(&[input])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + let inner = arr.value(0); + let strings = inner + .as_any() + .downcast_ref::() + .expect("inner array should be StringArray"); + + // HashSet ordering is nondeterministic — sort before asserting. + let mut values: Vec<&str> = + (0..strings.len()).map(|i| strings.value(i)).collect(); + values.sort_unstable(); + + assert_eq!(values, vec!["duckdb", "mysql", "postgres", "redis"]); + Ok(()) + } + + #[test] + fn distinct_array_agg_int64_deduplicates() -> Result<()> { + use arrow::array::Int64Array; + + // 7 rows with 4 distinct values, each duplicate appearing twice. + let input: ArrayRef = Arc::new(Int64Array::from(vec![1i64, 2, 1, 3, 2, 4, 3])); + + let mut acc = + DistinctArrayAggAccumulator::try_new(&DataType::Int64, None, false)?; + acc.update_batch(&[input])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + let inner = arr.value(0); + let ints = inner + .as_any() + .downcast_ref::() + .expect("inner array should be Int64Array"); + + let mut values: Vec = (0..ints.len()).map(|i| ints.value(i)).collect(); + values.sort_unstable(); + + assert_eq!(values, vec![1i64, 2, 3, 4]); + Ok(()) + } + + #[test] + fn distinct_array_agg_float64_deduplicates() -> Result<()> { + use arrow::array::Float64Array; + + // 7 rows with 4 distinct values, each duplicate appearing twice. + let input: ArrayRef = Arc::new(Float64Array::from(vec![ + 1.0f64, 2.5, 1.0, 3.75, 2.5, 4.0, 3.75, + ])); + + let mut acc = + DistinctArrayAggAccumulator::try_new(&DataType::Float64, None, false)?; + acc.update_batch(&[input])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + let inner = arr.value(0); + let floats = inner + .as_any() + .downcast_ref::() + .expect("inner array should be Float64Array"); + + // f64 has no Ord — use total_cmp for a stable sort. + let mut values: Vec = (0..floats.len()).map(|i| floats.value(i)).collect(); + values.sort_unstable_by(|a, b| a.total_cmp(b)); + + assert_eq!(values, vec![1.0f64, 2.5, 3.75, 4.0]); + Ok(()) + } + + #[test] + fn distinct_array_agg_dictionary_preserves_type() -> Result<()> { + use arrow::array::{DictionaryArray, Int32Array, StringArray}; + + // Dictionary(Int32, Utf8) input with duplicates. + let keys = Int32Array::from(vec![0, 1, 0, 2, 1]); // "a", "b", "a", "c", "b" + let values = StringArray::from(vec!["a", "b", "c"]); + let dict: ArrayRef = Arc::new(DictionaryArray::new(keys, Arc::new(values))); + + let datatype = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let mut acc = DistinctArrayAggAccumulator::try_new(&datatype, None, false)?; + acc.update_batch(&[dict])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + // The element type of the returned list must stay Dictionary(Int32, Utf8), + // not be silently widened to Utf8. + assert_eq!( + arr.values().data_type(), + &datatype, + "element type must be Dictionary(Int32, Utf8), got {}", + arr.values().data_type() + ); + + // There should be exactly 3 distinct values. + assert_eq!(arr.value(0).len(), 3); + Ok(()) + } + + #[test] + fn distinct_array_agg_date32_deduplicates() -> Result<()> { + use arrow::array::Date32Array; + + // 7 rows with 4 distinct dates (days since epoch), each duplicate appearing twice. + let input: ArrayRef = Arc::new(Date32Array::from(vec![ + 100i32, 200, 100, 300, 200, 400, 300, + ])); + + let mut acc = + DistinctArrayAggAccumulator::try_new(&DataType::Date32, None, false)?; + acc.update_batch(&[input])?; + + let result = acc.evaluate()?; + let ScalarValue::List(arr) = &result else { + panic!("expected ScalarValue::List, got {result:?}"); + }; + + let inner = arr.value(0); + let dates = inner + .as_any() + .downcast_ref::() + .expect("inner array should be Date32Array"); + + let mut values: Vec = (0..dates.len()).map(|i| dates.value(i)).collect(); + values.sort_unstable(); + + assert_eq!(values, vec![100i32, 200, 300, 400]); + Ok(()) + } + + #[test] + fn distinct_retract_memory_is_bounded() -> Result<()> { + use arrow::array::Int64Array; + + // Emulates a sliding window where each value enters and immediately + // leaves. Only CARDINALITY distinct values are ever live at once; + // memory must not grow with the number of rows processed. + const CARDINALITY: i64 = 10; + const WARMUP_ROWS: i64 = 1_000; + const EXTRA_ROWS: i64 = 20_000; + + let mut acc = + DistinctArrayAggAccumulator::try_new(&DataType::Int64, None, false)?; + + let slide = |acc: &mut DistinctArrayAggAccumulator, rows: i64| -> Result<()> { + for i in 0..rows { + let value: ArrayRef = Arc::new(Int64Array::from(vec![i % CARDINALITY])); + acc.update_batch(std::slice::from_ref(&value))?; + acc.retract_batch(std::slice::from_ref(&value))?; + } + Ok(()) + }; + + // Let every buffer reach its steady state before taking a baseline. + slide(&mut acc, WARMUP_ROWS)?; + let baseline = acc.size(); + + slide(&mut acc, EXTRA_ROWS)?; + let grown = acc.size(); + + assert!( + grown <= 2 * baseline, + "size() must not grow with the number of retracted rows: \ + {baseline} bytes after {WARMUP_ROWS} rows, \ + {grown} bytes after {} rows", + WARMUP_ROWS + EXTRA_ROWS + ); + + // Everything was retracted so evaluate must return null. + let result = acc.evaluate()?; + assert!( + matches!(&result, ScalarValue::List(arr) if arr.is_null(0)), + "expected null list after retracting every row, got {result:?}" + ); + + Ok(()) + } } diff --git a/datafusion/functions-aggregate/src/average.rs b/datafusion/functions-aggregate/src/average.rs index ddeb9b0870a16..e5030bf39e409 100644 --- a/datafusion/functions-aggregate/src/average.rs +++ b/datafusion/functions-aggregate/src/average.rs @@ -22,17 +22,19 @@ use arrow::array::{ BooleanArray, PrimitiveArray, PrimitiveBuilder, UInt64Array, }; -use arrow::compute::sum; +use arrow::compute::{DecimalCast, sum}; use arrow::datatypes::{ ArrowNativeType, DECIMAL32_MAX_PRECISION, DECIMAL32_MAX_SCALE, DECIMAL64_MAX_PRECISION, DECIMAL64_MAX_SCALE, DECIMAL128_MAX_PRECISION, DECIMAL128_MAX_SCALE, DECIMAL256_MAX_PRECISION, DECIMAL256_MAX_SCALE, DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DecimalType, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, - DurationSecondType, Field, FieldRef, Float64Type, TimeUnit, UInt64Type, i256, + DurationSecondType, Field, FieldRef, Float64Type, TimeUnit, UInt64Type, }; use datafusion_common::types::{NativeType, logical_float64}; -use datafusion_common::{Result, ScalarValue, exec_err, not_impl_err}; +use datafusion_common::{ + Result, ScalarValue, exec_datafusion_err, exec_err, internal_err, not_impl_err, +}; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ @@ -51,6 +53,7 @@ use datafusion_functions_aggregate_common::utils::DecimalAverager; use datafusion_macros::user_doc; use log::debug; use std::fmt::Debug; +use std::marker::PhantomData; use std::mem::{size_of, size_of_val}; use std::sync::Arc; @@ -125,6 +128,85 @@ impl Default for Avg { } } +/// Digits reserved above the input precision for `avg`'s intermediate sum: 4 for +/// the scale-up [`DecimalAverager`] applies before dividing (`Avg::return_type` +/// adds 4 to the scale), 9 for the row count. +/// +/// The 9 is a row budget. A sum of `n` rows of `Decimal(p, _)` is bounded by +/// `n * 10^p`, so a sum type with `p + 4 + 9` digits holds `10^9` rows. The sum +/// wraps on overflow, like the `sum` aggregate; the budget is what puts that out +/// of reach. `Decimal256` input near max precision is the exception: no wider +/// type exists, so its sum keeps only whatever headroom `Decimal256(76, _)` has +/// left, as before this budget was introduced. +const AVG_SUM_HEADROOM_DIGITS: u8 = 13; + +/// The narrowest decimal that can accumulate `avg`'s sum over `data_type`, never +/// narrower than `data_type` itself. Other types accumulate as themselves. +fn avg_sum_data_type(data_type: &DataType) -> DataType { + let (precision, scale, input_max_precision) = match data_type { + DataType::Decimal32(precision, scale) => { + (*precision, *scale, DECIMAL32_MAX_PRECISION) + } + DataType::Decimal64(precision, scale) => { + (*precision, *scale, DECIMAL64_MAX_PRECISION) + } + DataType::Decimal128(precision, scale) => { + (*precision, *scale, DECIMAL128_MAX_PRECISION) + } + DataType::Decimal256(precision, scale) => { + (*precision, *scale, DECIMAL256_MAX_PRECISION) + } + data_type => return data_type.clone(), + }; + + let required = precision + .saturating_add(AVG_SUM_HEADROOM_DIGITS) + .max(input_max_precision); + + // `required` always exceeds `DECIMAL32_MAX_PRECISION`, so a `Decimal32` sum is + // never wide enough, not even for `Decimal32` input + if required <= DECIMAL64_MAX_PRECISION { + DataType::Decimal64(DECIMAL64_MAX_PRECISION, scale) + } else if required <= DECIMAL128_MAX_PRECISION { + DataType::Decimal128(DECIMAL128_MAX_PRECISION, scale) + } else { + DataType::Decimal256(DECIMAL256_MAX_PRECISION, scale) + } +} + +/// Instantiates `$builder::` for every decimal pair that +/// [`avg_sum_data_type`] can produce. +macro_rules! decimal_avg_dispatch { + ($input:expr, $sum:expr, $builder:ident, $($arg:expr),*) => { + match ($input, $sum) { + (DataType::Decimal32(..), DataType::Decimal64(..)) => { + $builder::($($arg),*) + } + (DataType::Decimal32(..), DataType::Decimal128(..)) => { + $builder::($($arg),*) + } + (DataType::Decimal64(..), DataType::Decimal64(..)) => { + $builder::($($arg),*) + } + (DataType::Decimal64(..), DataType::Decimal128(..)) => { + $builder::($($arg),*) + } + (DataType::Decimal128(..), DataType::Decimal128(..)) => { + $builder::($($arg),*) + } + (DataType::Decimal128(..), DataType::Decimal256(..)) => { + $builder::($($arg),*) + } + (DataType::Decimal256(..), DataType::Decimal256(..)) => { + $builder::($($arg),*) + } + (input, sum) => { + internal_err!("avg cannot accumulate {input} as {sum}") + } + } + }; +} + impl AggregateUDFImpl for Avg { fn name(&self) -> &str { "avg" @@ -179,39 +261,19 @@ impl AggregateUDFImpl for Avg { // Numeric types are converted to Float64 via `coerce_avg_type` during logical plan creation (Float64, _) => Ok(Box::new(Float64DistinctAvgAccumulator::default())), - ( - Decimal32(_, scale), - Decimal32(target_precision, target_scale), - ) => Ok(Box::new(DecimalDistinctAvgAccumulator::::with_decimal_params( - *scale, - *target_precision, - *target_scale, - ))), - ( - Decimal64(_, scale), - Decimal64(target_precision, target_scale), - ) => Ok(Box::new(DecimalDistinctAvgAccumulator::::with_decimal_params( - *scale, - *target_precision, - *target_scale, - ))), - ( - Decimal128(_, scale), - Decimal128(target_precision, target_scale), - ) => Ok(Box::new(DecimalDistinctAvgAccumulator::::with_decimal_params( - *scale, - *target_precision, - *target_scale, - ))), - - ( - Decimal256(_, scale), - Decimal256(target_precision, target_scale), - ) => Ok(Box::new(DecimalDistinctAvgAccumulator::::with_decimal_params( - *scale, - *target_precision, - *target_scale, - ))), + (Decimal32(..), Decimal32(..)) + | (Decimal64(..), Decimal64(..)) + | (Decimal128(..), Decimal128(..)) + | (Decimal256(..), Decimal256(..)) => { + let sum_data_type = avg_sum_data_type(data_type); + decimal_avg_dispatch!( + data_type, + &sum_data_type, + decimal_distinct_avg_accumulator, + &sum_data_type, + acc_args.return_type() + ) + } (dt, return_type) => exec_err!( "AVG(DISTINCT) for ({} --> {}) not supported", @@ -222,51 +284,19 @@ impl AggregateUDFImpl for Avg { } else { match (&data_type, acc_args.return_type()) { (Float64, Float64) => Ok(Box::::default()), - ( - Decimal32(sum_precision, sum_scale), - Decimal32(target_precision, target_scale), - ) => Ok(Box::new(DecimalAvgAccumulator:: { - sum: None, - count: 0, - sum_scale: *sum_scale, - sum_precision: *sum_precision, - target_precision: *target_precision, - target_scale: *target_scale, - })), - ( - Decimal64(sum_precision, sum_scale), - Decimal64(target_precision, target_scale), - ) => Ok(Box::new(DecimalAvgAccumulator:: { - sum: None, - count: 0, - sum_scale: *sum_scale, - sum_precision: *sum_precision, - target_precision: *target_precision, - target_scale: *target_scale, - })), - ( - Decimal128(sum_precision, sum_scale), - Decimal128(target_precision, target_scale), - ) => Ok(Box::new(DecimalAvgAccumulator:: { - sum: None, - count: 0, - sum_scale: *sum_scale, - sum_precision: *sum_precision, - target_precision: *target_precision, - target_scale: *target_scale, - })), - - ( - Decimal256(sum_precision, sum_scale), - Decimal256(target_precision, target_scale), - ) => Ok(Box::new(DecimalAvgAccumulator:: { - sum: None, - count: 0, - sum_scale: *sum_scale, - sum_precision: *sum_precision, - target_precision: *target_precision, - target_scale: *target_scale, - })), + (Decimal32(..), Decimal32(..)) + | (Decimal64(..), Decimal64(..)) + | (Decimal128(..), Decimal128(..)) + | (Decimal256(..), Decimal256(..)) => { + let sum_data_type = avg_sum_data_type(data_type); + decimal_avg_dispatch!( + data_type, + &sum_data_type, + decimal_avg_accumulator, + sum_data_type.clone(), + acc_args.return_type().clone() + ) + } (Duration(time_unit), Duration(result_unit)) => { Ok(Box::new(DurationAvgAccumulator { @@ -314,17 +344,14 @@ impl AggregateUDFImpl for Avg { .into(), ]) } else { + let sum_data_type = avg_sum_data_type(args.input_fields[0].data_type()); Ok(vec![ Field::new( format_state_name(args.name, "count"), DataType::UInt64, true, ), - Field::new( - format_state_name(args.name, "sum"), - args.input_fields[0].data_type().clone(), - true, - ), + Field::new(format_state_name(args.name, "sum"), sum_data_type, true), ] .into_iter() .map(Arc::new) @@ -361,83 +388,18 @@ impl AggregateUDFImpl for Avg { |sum: f64, count: u64| Ok(sum / count as f64), ))) } - ( - Decimal32(_sum_precision, sum_scale), - Decimal32(target_precision, target_scale), - ) => { - let decimal_averager = DecimalAverager::::try_new( - *sum_scale, - *target_precision, - *target_scale, - )?; - - let avg_fn = - move |sum: i32, count: u64| decimal_averager.avg(sum, count as i32); - - Ok(Box::new(AvgGroupsAccumulator::::new( + (Decimal32(..), Decimal32(..)) + | (Decimal64(..), Decimal64(..)) + | (Decimal128(..), Decimal128(..)) + | (Decimal256(..), Decimal256(..)) => { + let sum_data_type = avg_sum_data_type(data_type); + decimal_avg_dispatch!( data_type, - args.return_field.data_type(), - avg_fn, - ))) - } - ( - Decimal64(_sum_precision, sum_scale), - Decimal64(target_precision, target_scale), - ) => { - let decimal_averager = DecimalAverager::::try_new( - *sum_scale, - *target_precision, - *target_scale, - )?; - - let avg_fn = - move |sum: i64, count: u64| decimal_averager.avg(sum, count as i64); - - Ok(Box::new(AvgGroupsAccumulator::::new( - data_type, - args.return_field.data_type(), - avg_fn, - ))) - } - ( - Decimal128(_sum_precision, sum_scale), - Decimal128(target_precision, target_scale), - ) => { - let decimal_averager = DecimalAverager::::try_new( - *sum_scale, - *target_precision, - *target_scale, - )?; - - let avg_fn = - move |sum: i128, count: u64| decimal_averager.avg(sum, count as i128); - - Ok(Box::new(AvgGroupsAccumulator::::new( - data_type, - args.return_field.data_type(), - avg_fn, - ))) - } - - ( - Decimal256(_sum_precision, sum_scale), - Decimal256(target_precision, target_scale), - ) => { - let decimal_averager = DecimalAverager::::try_new( - *sum_scale, - *target_precision, - *target_scale, - )?; - - let avg_fn = move |sum: i256, count: u64| { - decimal_averager.avg(sum, i256::from_usize(count as usize).unwrap()) - }; - - Ok(Box::new(AvgGroupsAccumulator::::new( - data_type, - args.return_field.data_type(), - avg_fn, - ))) + &sum_data_type, + decimal_avg_groups_accumulator, + &sum_data_type, + args.return_field.data_type() + ) } (Duration(time_unit), Duration(_result_unit)) => { @@ -500,6 +462,117 @@ impl AggregateUDFImpl for Avg { } } +/// The precision and scale of a decimal `DataType` +fn decimal_parts(data_type: &DataType) -> Result<(u8, i8)> { + match data_type { + DataType::Decimal32(precision, scale) + | DataType::Decimal64(precision, scale) + | DataType::Decimal128(precision, scale) + | DataType::Decimal256(precision, scale) => Ok((*precision, *scale)), + data_type => internal_err!("expected a decimal type, got {data_type}"), + } +} + +fn decimal_avg_fn( + sum_scale: i8, + target_precision: u8, + target_scale: i8, +) -> Result Result + Send + Sync + 'static> +where + I: DecimalType, + S: DecimalType, + I::Native: DecimalCast, + S::Native: DecimalCast, +{ + let decimal_averager = + DecimalAverager::::try_new(sum_scale, target_precision, target_scale)?; + + Ok(move |sum, count: u64| { + let Some(count) = usize::try_from(count).ok().and_then(S::Native::from_usize) + else { + return exec_err!( + "Arithmetic overflow in avg: the row count {count} cannot be \ + represented in the sum type" + ); + }; + + // Narrowing the average back to the (never wider) output type cannot + // fail in practice: `DecimalAverager::avg` validates the average + // against the output precision, whose bound fits the output's native + // type by construction + I::Native::from_decimal(decimal_averager.avg(sum, count)?).ok_or_else(|| { + exec_datafusion_err!( + "Arithmetic overflow in avg: the computed average does not fit \ + the output type" + ) + }) + }) +} + +fn decimal_avg_accumulator( + sum_data_type: DataType, + return_data_type: DataType, +) -> Result> +where + I: DecimalType + ArrowNumericType + Debug + Send + Sync, + S: DecimalType + ArrowNumericType + Debug + Send + Sync, + I::Native: Into + DecimalCast, + S::Native: DecimalCast, +{ + let (_, sum_scale) = decimal_parts(&sum_data_type)?; + let (target_precision, target_scale) = decimal_parts(&return_data_type)?; + let avg_fn = decimal_avg_fn::(sum_scale, target_precision, target_scale)?; + + Ok(Box::new(DecimalAvgAccumulator::::new( + sum_data_type, + return_data_type, + avg_fn, + ))) +} + +fn decimal_distinct_avg_accumulator( + sum_data_type: &DataType, + return_data_type: &DataType, +) -> Result> +where + I: DecimalType + ArrowNumericType + Debug + Send + Sync, + S: DecimalType + ArrowNumericType + Debug + Send + Sync, + I::Native: Into + DecimalCast, + S::Native: DecimalCast, +{ + let (_, sum_scale) = decimal_parts(sum_data_type)?; + let (target_precision, target_scale) = decimal_parts(return_data_type)?; + + Ok(Box::new( + DecimalDistinctAvgAccumulator::::with_decimal_params( + sum_scale, + target_precision, + target_scale, + ), + )) +} + +fn decimal_avg_groups_accumulator( + sum_data_type: &DataType, + return_data_type: &DataType, +) -> Result> +where + I: DecimalType + ArrowNumericType + Debug + Send + Sync, + S: DecimalType + ArrowNumericType + Debug + Send + Sync, + I::Native: Into + DecimalCast, + S::Native: DecimalCast, +{ + let (_, sum_scale) = decimal_parts(sum_data_type)?; + let (target_precision, target_scale) = decimal_parts(return_data_type)?; + let avg_fn = decimal_avg_fn::(sum_scale, target_precision, target_scale)?; + + Ok(Box::new(AvgGroupsAccumulator::::new( + sum_data_type, + return_data_type, + avg_fn, + ))) +} + /// An accumulator to compute the average #[derive(Debug, Default)] pub struct AvgAccumulator { @@ -567,24 +640,104 @@ impl Accumulator for AvgAccumulator { } } -/// An accumulator to compute the average for decimals -#[derive(Debug)] -struct DecimalAvgAccumulator { - sum: Option, +/// An accumulator to compute the average for decimals. +/// +/// `I` is the input (and output) decimal type. `S` is the type used to accumulate +/// the sum, chosen by [`avg_sum_data_type`] so the running total does not overflow. +struct DecimalAvgAccumulator +where + I: DecimalType + ArrowNumericType + Debug, + S: DecimalType + ArrowNumericType + Debug, + I::Native: Into, + F: Fn(S::Native, u64) -> Result, +{ + sum: Option, count: u64, - sum_scale: i8, - sum_precision: u8, - target_precision: u8, - target_scale: i8, + sum_data_type: DataType, + return_data_type: DataType, + avg_fn: F, + _phantom: PhantomData, +} + +impl Debug for DecimalAvgAccumulator +where + I: DecimalType + ArrowNumericType + Debug, + S: DecimalType + ArrowNumericType + Debug, + I::Native: Into, + F: Fn(S::Native, u64) -> Result, +{ + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("DecimalAvgAccumulator") + .field("sum", &self.sum) + .field("count", &self.count) + .field("sum_data_type", &self.sum_data_type) + .field("return_data_type", &self.return_data_type) + .finish_non_exhaustive() + } } -impl Accumulator for DecimalAvgAccumulator { +impl DecimalAvgAccumulator +where + I: DecimalType + ArrowNumericType + Debug, + S: DecimalType + ArrowNumericType + Debug, + I::Native: Into, + F: Fn(S::Native, u64) -> Result, +{ + fn new(sum_data_type: DataType, return_data_type: DataType, avg_fn: F) -> Self { + Self { + sum: None, + count: 0, + sum_data_type, + return_data_type, + avg_fn, + _phantom: PhantomData, + } + } +} + +/// Sums `values` into the wider `S`. +/// +/// Wraps on overflow, matching the `sum` aggregate and [`arrow::compute::sum`]. +/// [`avg_sum_data_type`] gives `S` enough headroom that this is unreachable for +/// any realistic row count. +fn decimal_sum_as(values: &PrimitiveArray) -> Option +where + I: DecimalType + ArrowNumericType, + S: DecimalType + ArrowNumericType, + I::Native: Into, +{ + // Matches `arrow::compute::sum`: an empty or all-null input has no sum + if values.null_count() == values.len() { + return None; + } + + let mut sum = S::Native::default(); + if values.null_count() == 0 { + for value in values.values() { + sum = sum.add_wrapping((*value).into()); + } + } else { + for value in values.iter().flatten() { + sum = sum.add_wrapping(value.into()); + } + } + + Some(sum) +} + +impl Accumulator for DecimalAvgAccumulator +where + I: DecimalType + ArrowNumericType + Debug, + S: DecimalType + ArrowNumericType + Debug, + I::Native: Into, + F: Fn(S::Native, u64) -> Result + Send + Sync + 'static, +{ fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let values = values[0].as_primitive::(); + let values = values[0].as_primitive::(); self.count += (values.len() - values.null_count()) as u64; - if let Some(x) = sum(values) { - let v = self.sum.get_or_insert_with(T::Native::default); + if let Some(x) = decimal_sum_as::(values) { + let v = self.sum.unwrap_or_default(); self.sum = Some(v.add_wrapping(x)); } Ok(()) @@ -597,22 +750,10 @@ impl Accumulator for DecimalAvgAccumu let v = if self.count == 0 { None } else { - self.sum - .map(|v| { - DecimalAverager::::try_new( - self.sum_scale, - self.target_precision, - self.target_scale, - )? - .avg(v, T::Native::from_usize(self.count as usize).unwrap()) - }) - .transpose()? + self.sum.map(|v| (self.avg_fn)(v, self.count)).transpose()? }; - ScalarValue::new_primitive::( - v, - &T::TYPE_CONSTRUCTOR(self.target_precision, self.target_scale), - ) + ScalarValue::new_primitive::(v, &self.return_data_type) } fn size(&self) -> usize { @@ -622,10 +763,7 @@ impl Accumulator for DecimalAvgAccumu fn state(&mut self) -> Result> { Ok(vec![ ScalarValue::from(self.count), - ScalarValue::new_primitive::( - self.sum, - &T::TYPE_CONSTRUCTOR(self.sum_precision, self.sum_scale), - )?, + ScalarValue::new_primitive::(self.sum, &self.sum_data_type)?, ]) } @@ -634,17 +772,18 @@ impl Accumulator for DecimalAvgAccumu self.count += sum(states[0].as_primitive::()).unwrap_or_default(); // sums are summed - if let Some(x) = sum(states[1].as_primitive::()) { - let v = self.sum.get_or_insert_with(T::Native::default); + if let Some(x) = sum(states[1].as_primitive::()) { + let v = self.sum.unwrap_or_default(); self.sum = Some(v.add_wrapping(x)); } Ok(()) } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let values = values[0].as_primitive::(); + let values = values[0].as_primitive::(); self.count -= (values.len() - values.null_count()) as u64; - if let Some(x) = sum(values) { - self.sum = Some(self.sum.unwrap().sub_wrapping(x)); + if let Some(x) = decimal_sum_as::(values) { + let v = self.sum.unwrap_or_default(); + self.sum = Some(v.sub_wrapping(x)); } Ok(()) } @@ -760,16 +899,21 @@ impl Accumulator for DurationAvgAccumulator { } } -/// An accumulator to compute the average of `[PrimitiveArray]`. +/// An accumulator to compute the average of `[PrimitiveArray]`. /// Stores values as native types, and does overflow checking /// /// F: Function that calculates the average value from a sum of -/// T::Native and a total count +/// S::Native and a total count +/// +/// `I` is the input (and output) type. `S` is a possibly wider type used to +/// accumulate the sum so it does not overflow. #[derive(Debug)] -struct AvgGroupsAccumulator +struct AvgGroupsAccumulator where - T: ArrowNumericType + Send, - F: Fn(T::Native, u64) -> Result + Send + 'static, + I: ArrowNumericType + Send, + S: ArrowNumericType + Send, + I::Native: Into, + F: Fn(S::Native, u64) -> Result + Send + 'static, { /// The type of the internal sum sum_data_type: DataType, @@ -781,24 +925,28 @@ where counts: Vec, /// Sums per group, stored as the native type - sums: Vec, + sums: Vec, /// Track nulls in the input / filters null_state: NullState, /// Function that computes the final average (value / count) avg_fn: F, + + _phantom: PhantomData, } -impl AvgGroupsAccumulator +impl AvgGroupsAccumulator where - T: ArrowNumericType + Send, - F: Fn(T::Native, u64) -> Result + Send + 'static, + I: ArrowNumericType + Send, + S: ArrowNumericType + Send, + I::Native: Into, + F: Fn(S::Native, u64) -> Result + Send + 'static, { pub fn new(sum_data_type: &DataType, return_data_type: &DataType, avg_fn: F) -> Self { debug!( "AvgGroupsAccumulator ({}, sum type: {sum_data_type}) --> {return_data_type}", - std::any::type_name::() + std::any::type_name::() ); Self { @@ -808,14 +956,17 @@ where sums: vec![], null_state: NullState::new(), avg_fn, + _phantom: PhantomData, } } } -impl GroupsAccumulator for AvgGroupsAccumulator +impl GroupsAccumulator for AvgGroupsAccumulator where - T: ArrowNumericType + Send, - F: Fn(T::Native, u64) -> Result + Send + 'static, + I: ArrowNumericType + Send, + S: ArrowNumericType + Send, + I::Native: Into, + F: Fn(S::Native, u64) -> Result + Send + 'static, { fn update_batch( &mut self, @@ -825,11 +976,12 @@ where total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "single argument to update_batch"); - let values = values[0].as_primitive::(); + let values = values[0].as_primitive::(); // increment counts, update sums self.counts.resize(total_num_groups, 0); - self.sums.resize(total_num_groups, T::default_value()); + self.sums.resize(total_num_groups, S::default_value()); + self.null_state.accumulate( group_indices, values, @@ -838,7 +990,7 @@ where |group_index, new_value| { // SAFETY: group_index is guaranteed to be in bounds let sum = unsafe { self.sums.get_unchecked_mut(group_index) }; - *sum = sum.add_wrapping(new_value); + *sum = sum.add_wrapping(new_value.into()); self.counts[group_index] += 1; }, @@ -859,10 +1011,10 @@ where // don't evaluate averages with null inputs to avoid errors on null values - let array: PrimitiveArray = if let Some(nulls) = &nulls + let array: PrimitiveArray = if let Some(nulls) = &nulls && nulls.null_count() > 0 { - let mut builder = PrimitiveBuilder::::with_capacity(nulls.len()) + let mut builder = PrimitiveBuilder::::with_capacity(nulls.len()) .with_data_type(self.return_data_type.clone()); let iter = sums.into_iter().zip(counts).zip(nulls.iter()); @@ -875,7 +1027,7 @@ where } builder.finish() } else { - let averages: Vec = sums + let averages: Vec = sums .into_iter() .zip(counts) .map(|(sum, count)| (self.avg_fn)(sum, count)) @@ -895,7 +1047,7 @@ where let counts = UInt64Array::new(counts.into(), nulls.clone()); // zero copy let sums = emit_to.take_needed(&mut self.sums); - let sums = PrimitiveArray::::new(sums.into(), nulls) // zero copy + let sums = PrimitiveArray::::new(sums.into(), nulls) // zero copy .with_data_type(self.sum_data_type.clone()); Ok(vec![ @@ -908,19 +1060,18 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 2, "two arguments to merge_batch"); // first batch is counts, second is partial sums let partial_counts = values[0].as_primitive::(); - let partial_sums = values[1].as_primitive::(); + let partial_sums = values[1].as_primitive::(); // update counts with partial counts self.counts.resize(total_num_groups, 0); self.null_state.accumulate( group_indices, partial_counts, - opt_filter, + None, total_num_groups, |group_index, partial_count| { // SAFETY: group_index is guaranteed to be in bounds @@ -930,13 +1081,13 @@ where ); // update sums - self.sums.resize(total_num_groups, T::default_value()); + self.sums.resize(total_num_groups, S::default_value()); self.null_state.accumulate( group_indices, partial_sums, - opt_filter, + None, total_num_groups, - |group_index, new_value: ::Native| { + |group_index, new_value: ::Native| { // SAFETY: group_index is guaranteed to be in bounds let sum = unsafe { self.sums.get_unchecked_mut(group_index) }; *sum = sum.add_wrapping(new_value); @@ -951,10 +1102,27 @@ where values: &[ArrayRef], opt_filter: Option<&BooleanArray>, ) -> Result> { - let sums = values[0] - .as_primitive::() - .clone() - .with_data_type(self.sum_data_type.clone()); + // When the sum type equals the input type (`I == S`: `Float64`, + // `Duration`, `Decimal256`, and any decimal whose precision already + // leaves [`avg_sum_data_type`] enough headroom) the input is already a + // valid sum array and is reused as is; the downcast is by Rust type, so + // it succeeds even when precision differs. Otherwise every value is + // widened. + let sums = match values[0].as_any().downcast_ref::>() { + Some(sums) => sums.clone().with_data_type(self.sum_data_type.clone()), + None => { + let values = values[0].as_primitive::(); + // Values under null slots are widened too rather than branching per + // element; `set_nulls` below masks them out again. + let sums: Vec = values + .values() + .iter() + .map(|value| (*value).into()) + .collect(); + PrimitiveArray::::new(sums.into(), values.nulls().cloned()) + .with_data_type(self.sum_data_type.clone()) + } + }; let counts = UInt64Array::from_value(1, sums.len()); let nulls = filtered_null_mask(opt_filter, &sums); @@ -965,12 +1133,265 @@ where Ok(vec![Arc::new(counts) as ArrayRef, Arc::new(sums)]) } + fn size(&self) -> usize { + // Heap buffers + self.counts.capacity() * size_of::() + + self.sums.capacity() * size_of::() + // Vec struct overhead (ptr, len, cap) for each field + + size_of::>() + + size_of::>() + // Null tracking buffers + + self.null_state.size() + } +} - fn supports_convert_to_state(&self) -> bool { - true +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{ + Decimal32Array, Decimal64Array, Decimal128Array, Decimal256Array, + DurationSecondArray, Float64Array, + }; + use arrow::datatypes::{Schema, i256}; + + struct AvgCase { + name: &'static str, + values: ArrayRef, + return_type: DataType, + sum_type: DataType, + expected: ScalarValue, } - fn size(&self) -> usize { - self.counts.capacity() * size_of::() + self.sums.capacity() * size_of::() + fn with_avg_args( + input_type: &DataType, + return_type: &DataType, + f: impl FnOnce(AccumulatorArgs) -> R, + ) -> R { + let schema = Schema::empty(); + let expr_field = Arc::new(Field::new("a", input_type.clone(), true)); + let return_field = Arc::new(Field::new("avg", return_type.clone(), true)); + + f(AccumulatorArgs { + return_field, + schema: &schema, + expr_fields: &[expr_field], + ignore_nulls: false, + order_bys: &[], + is_distinct: false, + name: "avg", + is_reversed: false, + exprs: &[], + }) + } + + fn avg_groups_accumulator( + input_type: &DataType, + return_type: &DataType, + ) -> Result> { + with_avg_args(input_type, return_type, |args| { + Avg::new().create_groups_accumulator(args) + }) + } + + fn avg_accumulator( + input_type: &DataType, + return_type: &DataType, + ) -> Result> { + with_avg_args(input_type, return_type, |args| Avg::new().accumulator(args)) + } + + fn avg_state_fields( + input_type: &DataType, + return_type: &DataType, + ) -> Result> { + let input_field = Arc::new(Field::new("a", input_type.clone(), true)); + let return_field = Arc::new(Field::new("avg", return_type.clone(), true)); + + Avg::new().state_fields(StateFieldsArgs { + name: "avg", + input_fields: &[input_field], + return_field, + ordering_fields: &[], + is_distinct: false, + }) + } + + fn avg_cases() -> Result> { + const ROWS: usize = 21_476; + const DECIMAL32_VALUE: i32 = 99_999; + const DECIMAL64_ROWS: usize = 92_235; + const DECIMAL64_VALUE: i64 = 99_999_999_999_999; + const DECIMAL128_ROWS: usize = 21_476; + const DECIMAL128_VALUE: i128 = 9_999_999_999_999_999_999_999_999_999_999_999; + + Ok(vec![ + AvgCase { + name: "float64", + values: Arc::new(Float64Array::from(vec![10.0, 20.0])), + return_type: DataType::Float64, + sum_type: DataType::Float64, + expected: ScalarValue::Float64(Some(15.0)), + }, + AvgCase { + name: "decimal32", + values: Arc::new( + Decimal32Array::from(vec![Some(DECIMAL32_VALUE); ROWS]) + .with_precision_and_scale(5, 0)?, + ), + return_type: DataType::Decimal32(9, 4), + sum_type: DataType::Decimal64(18, 0), + expected: ScalarValue::Decimal32(Some(DECIMAL32_VALUE * 10_000), 9, 4), + }, + AvgCase { + name: "decimal64", + values: Arc::new( + Decimal64Array::from(vec![Some(DECIMAL64_VALUE); DECIMAL64_ROWS]) + .with_precision_and_scale(14, 0)?, + ), + return_type: DataType::Decimal64(18, 4), + sum_type: DataType::Decimal128(38, 0), + expected: ScalarValue::Decimal64(Some(DECIMAL64_VALUE * 10_000), 18, 4), + }, + AvgCase { + name: "decimal128", + values: Arc::new( + Decimal128Array::from(vec![Some(DECIMAL128_VALUE); DECIMAL128_ROWS]) + .with_precision_and_scale(34, 0)?, + ), + return_type: DataType::Decimal128(38, 4), + sum_type: DataType::Decimal256(76, 0), + expected: ScalarValue::Decimal128(Some(DECIMAL128_VALUE * 10_000), 38, 4), + }, + AvgCase { + name: "decimal256", + values: Arc::new( + Decimal256Array::from(vec![i256::from_i128(10), i256::from_i128(20)]) + .with_precision_and_scale(50, 0)?, + ), + return_type: DataType::Decimal256(54, 4), + sum_type: DataType::Decimal256(76, 0), + expected: ScalarValue::Decimal256(Some(i256::from_i128(150_000)), 54, 4), + }, + // A `Decimal128` whose precision leaves room for the sum stays on + // `i128` rather than widening to the emulated `i256` arithmetic + AvgCase { + name: "decimal128_with_headroom", + values: Arc::new( + Decimal128Array::from(vec![100_000, 200_000]) + .with_precision_and_scale(20, 4)?, + ), + return_type: DataType::Decimal128(24, 8), + sum_type: DataType::Decimal128(38, 4), + expected: ScalarValue::Decimal128(Some(1_500_000_000), 24, 8), + }, + // A `Decimal32` at max precision needs more than `Decimal64` can hold + // once `DecimalAverager` scales the sum up, so it accumulates as `i128` + AvgCase { + name: "decimal32_max_precision", + values: Arc::new( + Decimal32Array::from(vec![10, 20]).with_precision_and_scale(9, 0)?, + ), + return_type: DataType::Decimal32(9, 4), + sum_type: DataType::Decimal128(38, 0), + expected: ScalarValue::Decimal32(Some(150_000), 9, 4), + }, + // One duration unit suffices: all four units instantiate the same + // `S = I` generic code + AvgCase { + name: "duration_second", + values: Arc::new(DurationSecondArray::from(vec![10, 20])), + return_type: DataType::Duration(TimeUnit::Second), + sum_type: DataType::Duration(TimeUnit::Second), + expected: ScalarValue::DurationSecond(Some(15)), + }, + ]) + } + + #[test] + fn avg_accumulator_evaluate_and_state_types() -> Result<()> { + for case in avg_cases()? { + let input_type = case.values.data_type(); + let state_fields = avg_state_fields(input_type, &case.return_type)?; + let mut acc = avg_accumulator(input_type, &case.return_type)?; + acc.update_batch(std::slice::from_ref(&case.values))?; + + let state = acc.state()?; + assert_eq!( + &state[0].data_type(), + state_fields[0].data_type(), + "{}", + case.name + ); + assert_eq!( + &state[1].data_type(), + state_fields[1].data_type(), + "{}", + case.name + ); + assert_eq!(acc.evaluate()?, case.expected, "{}", case.name); + } + + Ok(()) + } + + #[test] + fn avg_groups_state_types_match_state_fields() -> Result<()> { + for case in avg_cases()? { + let input_type = case.values.data_type(); + let state_fields = avg_state_fields(input_type, &case.return_type)?; + let acc = avg_groups_accumulator(input_type, &case.return_type)?; + let state = acc.convert_to_state(std::slice::from_ref(&case.values), None)?; + + assert_eq!( + state_fields[0].data_type(), + &DataType::UInt64, + "{}", + case.name + ); + assert_eq!(state_fields[1].data_type(), &case.sum_type, "{}", case.name); + assert_eq!(state[0].data_type(), &DataType::UInt64, "{}", case.name); + assert_eq!(state[1].data_type(), &case.sum_type, "{}", case.name); + } + + Ok(()) + } + + #[test] + fn avg_groups_convert_to_state_roundtrip() -> Result<()> { + for case in avg_cases()? { + let input_type = case.values.data_type(); + let partial = avg_groups_accumulator(input_type, &case.return_type)?; + let mut final_acc = avg_groups_accumulator(input_type, &case.return_type)?; + let state = + partial.convert_to_state(std::slice::from_ref(&case.values), None)?; + final_acc.merge_batch(&state, &vec![0; case.values.len()], 1)?; + + let result = final_acc.evaluate(EmitTo::All)?; + assert_eq!(result.data_type(), &case.return_type, "{}", case.name); + assert_eq!( + ScalarValue::try_from_array(result.as_ref(), 0)?, + case.expected, + "{}", + case.name + ); + } + + Ok(()) + } + + /// The widened sum fits, but the average does not fit the output type once + /// `DecimalAverager` rescales it: avg must error rather than silently wrap + #[test] + fn avg_errors_when_average_exceeds_output_precision() -> Result<()> { + let values: ArrayRef = Arc::new( + Decimal32Array::from(vec![999_999_999]).with_precision_and_scale(9, 0)?, + ); + let return_type = DataType::Decimal32(9, 4); + let mut acc = avg_accumulator(values.data_type(), &return_type)?; + + acc.update_batch(&[values])?; + assert!(acc.evaluate().is_err()); + + Ok(()) } } diff --git a/datafusion/functions-aggregate/src/correlation.rs b/datafusion/functions-aggregate/src/correlation.rs index 2621fcf0bf3c7..b9bc57dfa989c 100644 --- a/datafusion/functions-aggregate/src/correlation.rs +++ b/datafusion/functions-aggregate/src/correlation.rs @@ -281,6 +281,10 @@ impl Accumulator for CorrelationAccumulator { self.stddev2.retract_batch(&values[1..2])?; Ok(()) } + + fn supports_retract_batch(&self) -> bool { + true + } } #[derive(Default)] @@ -485,11 +489,60 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { ]) } + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + assert_eq!(values.len(), 2, "two arguments to convert_to_state"); + let array_x = downcast_array::(&values[0]); + let array_y = downcast_array::(&values[1]); + + let len = array_x.len(); + let mut counts = Vec::with_capacity(len); + let mut sum_x = Vec::with_capacity(len); + let mut sum_y = Vec::with_capacity(len); + let mut sum_xy = Vec::with_capacity(len); + let mut sum_xx = Vec::with_capacity(len); + let mut sum_yy = Vec::with_capacity(len); + + for row in 0..len { + let included = array_x.is_valid(row) + && array_y.is_valid(row) + && opt_filter + .is_none_or(|filter| filter.is_valid(row) && filter.value(row)); + if included { + let x = array_x.value(row); + let y = array_y.value(row); + counts.push(1); + sum_x.push(x); + sum_y.push(y); + sum_xy.push(x * y); + sum_xx.push(x * x); + sum_yy.push(y * y); + } else { + counts.push(0); + sum_x.push(0.0); + sum_y.push(0.0); + sum_xy.push(0.0); + sum_xx.push(0.0); + sum_yy.push(0.0); + } + } + + Ok(vec![ + Arc::new(UInt64Array::from(counts)), + Arc::new(Float64Array::from(sum_x)), + Arc::new(Float64Array::from(sum_y)), + Arc::new(Float64Array::from(sum_xy)), + Arc::new(Float64Array::from(sum_xx)), + Arc::new(Float64Array::from(sum_yy)), + ]) + } fn merge_batch( &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // Resize vectors to accommodate total number of groups @@ -508,11 +561,6 @@ impl GroupsAccumulator for CorrelationGroupsAccumulator { let partial_sum_xx = values[4].as_primitive::(); let partial_sum_yy = values[5].as_primitive::(); - assert!( - opt_filter.is_none(), - "aggregate filter should be applied in partial stage, there should be no filter in final stage" - ); - accumulate_correlation_states( group_indices, ( @@ -595,4 +643,90 @@ mod tests { }); assert!(result.is_err()); } + + #[test] + fn convert_to_state_roundtrips_through_merge() -> Result<()> { + let x = Arc::new(Float64Array::from(vec![ + Some(1.0), + Some(2.0), + None, + Some(4.0), + Some(8.0), + Some(16.0), + Some(32.0), + ])) as ArrayRef; + let y = Arc::new(Float64Array::from(vec![ + Some(2.0), + Some(4.0), + Some(6.0), + None, + Some(16.0), + Some(32.0), + Some(64.0), + ])) as ArrayRef; + let filter = BooleanArray::from(vec![ + Some(true), + Some(false), + Some(true), + Some(true), + None, + Some(true), + Some(true), + ]); + let values = vec![x, y]; + let group_indices = vec![0, 1, 0, 1, 0, 0, 0]; + + let mut direct = CorrelationGroupsAccumulator::new(); + direct.update_batch(&values, &group_indices, Some(&filter), 2)?; + let direct = direct.evaluate(EmitTo::All)?; + + let converter = CorrelationGroupsAccumulator::new(); + let state = converter.convert_to_state(&values, Some(&filter))?; + let mut merged = CorrelationGroupsAccumulator::new(); + merged.merge_batch(&state, &group_indices, 2)?; + let merged = merged.evaluate(EmitTo::All)?; + + assert_eq!( + direct.as_any().downcast_ref::().unwrap(), + merged.as_any().downcast_ref::().unwrap() + ); + Ok(()) + } + + #[test] + fn convert_to_state_preserves_empty_and_filtered_rows() -> Result<()> { + let converter = CorrelationGroupsAccumulator::new(); + let empty_values = vec![ + Arc::new(Float64Array::from(Vec::>::new())) as ArrayRef, + Arc::new(Float64Array::from(Vec::>::new())) as ArrayRef, + ]; + let state = converter.convert_to_state(&empty_values, None)?; + for state_array in &state { + assert_eq!(state_array.len(), 0); + assert_eq!(state_array.null_count(), 0); + } + + let values = vec![ + Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), None])) as ArrayRef, + Arc::new(Float64Array::from(vec![Some(2.0), None, Some(4.0)])) as ArrayRef, + ]; + let filter = BooleanArray::from(vec![Some(false), None, Some(false)]); + let group_indices = vec![0, 1, 0]; + let state = converter.convert_to_state(&values, Some(&filter))?; + for state_array in &state { + assert_eq!(state_array.len(), values[0].len()); + assert_eq!(state_array.null_count(), 0); + } + + let counts = state[0].as_any().downcast_ref::().unwrap(); + assert_eq!(counts, &UInt64Array::from(vec![0, 0, 0])); + + let mut merged = CorrelationGroupsAccumulator::new(); + merged.merge_batch(&state, &group_indices, 2)?; + let result = merged.evaluate(EmitTo::All)?; + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result.null_count(), 2); + Ok(()) + } } diff --git a/datafusion/functions-aggregate/src/count.rs b/datafusion/functions-aggregate/src/count.rs index eab36d4951a9c..1e72d8ac3d5b1 100644 --- a/datafusion/functions-aggregate/src/count.rs +++ b/datafusion/functions-aggregate/src/count.rs @@ -29,6 +29,7 @@ use arrow::{ }, }; use datafusion_common::hash_utils::RandomState; +use datafusion_common::heap_size::{DFHeapSize, DFHeapSizeCtx}; use datafusion_common::{ HashMap, Result, ScalarValue, downcast_value, exec_err, internal_err, not_impl_err, stats::Precision, utils::expr::COUNT_STAR_EXPANSION, @@ -555,7 +556,17 @@ impl Accumulator for SlidingDistinctCountAccumulator { } fn size(&self) -> usize { + // Mirrors `DistinctCountAccumulator::full_size`: self + HashMap + // bucket array + per-key inner heap + DataType inner heap. size_of_val(self) + + (size_of::() + size_of::()) * self.counts.capacity() + + self + .counts + .keys() + .map(|k| k.size() - size_of_val(k)) + .sum::() + + self.data_type.size() + - size_of_val(&self.data_type) } } @@ -665,8 +676,6 @@ impl GroupsAccumulator for CountGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - // Since aggregate filter should be applied in partial stage, in final stage there should be no filter - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "one argument to merge_batch"); @@ -765,13 +774,8 @@ impl GroupsAccumulator for CountGroupsAccumulator { Ok(vec![state_array]) } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn size(&self) -> usize { - self.counts.capacity() * size_of::() + self.counts.heap_size(&mut DFHeapSizeCtx::default()) } } @@ -929,6 +933,23 @@ mod tests { )?) } + #[test] + fn count_groups_size_includes_vec_capacity() -> Result<()> { + let mut acc = CountGroupsAccumulator::new(); + let empty_size = acc.size(); + assert_eq!(empty_size, 0); + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); + acc.update_batch(&[values], &[0, 1, 2], None, 3)?; + + assert!(acc.counts.capacity() > 0); + let allocated_size = acc.counts.heap_size(&mut DFHeapSizeCtx::default()); + assert_eq!(allocated_size, acc.counts.capacity() * size_of::()); + assert_eq!(acc.size(), allocated_size); + assert!(acc.size() > empty_size); + + Ok(()) + } + #[test] fn count_accumulator_nulls() -> Result<()> { let mut accumulator = CountAccumulator::new(); diff --git a/datafusion/functions-aggregate/src/covariance.rs b/datafusion/functions-aggregate/src/covariance.rs index 18d602ab33940..bd7c8a039076a 100644 --- a/datafusion/functions-aggregate/src/covariance.rs +++ b/datafusion/functions-aggregate/src/covariance.rs @@ -305,6 +305,14 @@ impl Accumulator for CovarianceAccumulator { _ => continue, }; + if self.count <= 1 { + self.count = 0; + self.mean1 = 0.0; + self.mean2 = 0.0; + self.algo_const = 0.0; + continue; + } + let new_count = self.count - 1; let delta1 = self.mean1 - value1; let new_mean1 = delta1 / new_count as f64 + self.mean1; @@ -373,4 +381,8 @@ impl Accumulator for CovarianceAccumulator { fn size(&self) -> usize { size_of_val(self) } + + fn supports_retract_batch(&self) -> bool { + true + } } diff --git a/datafusion/functions-aggregate/src/first_last.rs b/datafusion/functions-aggregate/src/first_last.rs index 1935f29c4cfe8..c56cd73dbeabe 100644 --- a/datafusion/functions-aggregate/src/first_last.rs +++ b/datafusion/functions-aggregate/src/first_last.rs @@ -51,7 +51,7 @@ use datafusion_physical_expr_common::sort_expr::LexOrdering; mod state; -use state::{BytesValueState, PrimitiveValueState, ValueState}; +use state::{BytesValueState, GenericValueState, PrimitiveValueState, ValueState}; create_func!(FirstValue, first_value_udaf); create_func!(LastValue, last_value_udaf); @@ -171,6 +171,23 @@ fn create_groups_accumulator( BytesValueState::try_new(data_type.clone())?, ), + // Nested / composite types fall through to a generic ScalarValue-backed + // state. Slower per-batch than the primitive/bytes fast paths but still + // avoids the per-row ScalarValue churn of the per-group `Accumulator` + // path: winner extraction happens once per group per batch, not once + // per candidate row. + DataType::List(_) + | DataType::LargeList(_) + | DataType::ListView(_) + | DataType::LargeListView(_) + | DataType::FixedSizeList(_, _) + | DataType::Struct(_) + | DataType::Map(_, _) => create_groups_accumulator_helper( + args, + is_first, + GenericValueState::new(data_type.clone()), + ), + _ => internal_err!( "GroupsAccumulator not supported for {}({})", function_name, @@ -209,6 +226,13 @@ fn groups_accumulator_supported(args: &AccumulatorArgs) -> bool { | Binary | LargeBinary | BinaryView + | List(_) + | LargeList(_) + | ListView(_) + | LargeListView(_) + | FixedSizeList(_, _) + | Struct(_) + | Map(_, _) ) } @@ -555,8 +579,15 @@ impl FirstLastGroupsAccumulator { for (idx_in_val, group_idx) in group_indices.iter().enumerate() { let group_idx = *group_idx; - let passed_filter = opt_filter.is_none_or(|x| x.value(idx_in_val)); - let is_set = is_set_arr.is_none_or(|x| x.value(idx_in_val)); + // A row passes the FILTER clause only when the predicate is + // `true`; rows whose predicate evaluates to `null` are excluded. + let passed_filter = + opt_filter.is_none_or(|x| x.is_valid(idx_in_val) && x.value(idx_in_val)); + // `is_set_arr` carries the user FILTER clause (including its + // nulls) when the state was produced by `convert_to_state`, so + // the validity check is required here as well (#22666). + let is_set = + is_set_arr.is_none_or(|x| x.is_valid(idx_in_val) && x.value(idx_in_val)); if !passed_filter || !is_set { continue; @@ -671,7 +702,6 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator, total_num_groups: usize, ) -> Result<()> { self.resize_states(total_num_groups); @@ -690,7 +720,7 @@ impl GroupsAccumulator for FirstLastGroupsAccumulator GroupsAccumulator for FirstLastGroupsAccumulator() + self.extreme_of_each_group_buf.1.capacity() / 8 } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn convert_to_state( &self, values: &[ArrayRef], @@ -1192,7 +1217,7 @@ impl Accumulator for TrivialLastValueAccumulator { if let Some(last) = filtered_states.last() && !last.is_empty() { - self.last = ScalarValue::try_from_array(last, 0)?; + self.last = ScalarValue::try_from_array(last, last.len() - 1)?; self.is_set = true; } Ok(()) @@ -1416,6 +1441,7 @@ mod tests { use arrow::{ array::{BooleanArray, Int64Array, ListArray, PrimitiveArray, StringArray}, + buffer::NullBuffer, compute::SortOptions, datatypes::Schema, }; @@ -1523,7 +1549,21 @@ mod tests { let merged_state = last_accumulator.state()?; assert_eq!(merged_state.len(), state1.len()); + assert_eq!(last_accumulator.evaluate()?, ScalarValue::Int64(Some(10))); + + Ok(()) + } + + #[test] + fn test_trivial_last_value_merge_all_flags_false() -> Result<()> { + let mut acc = TrivialLastValueAccumulator::try_new(&DataType::Int64, false)?; + let states: Vec = vec![ + Arc::new(Int64Array::from(vec![None, None])), + Arc::new(BooleanArray::from(vec![false, false])), + ]; + acc.merge_batch(&states)?; + assert_eq!(acc.evaluate()?, ScalarValue::Int64(None)); Ok(()) } @@ -1587,12 +1627,7 @@ mod tests { group_acc.compute_size_of_orderings() ); - group_acc.merge_batch( - &state, - &[0, 1, 2], - Some(&BooleanArray::from(vec![true, false, false])), - 3, - )?; + group_acc.merge_batch(&state, &[0, 1, 2], 3)?; assert_eq!( group_acc.size_of_orderings, @@ -1608,8 +1643,11 @@ mod tests { let binding = group_acc.evaluate(EmitTo::All)?; let eval_result = binding.as_any().downcast_ref::().unwrap(); + // group 0 keeps merged value=1 (ordering=1). + // group 1 keeps merged value=-6 (ordering=-6 < 6, so -6 is "first"). + // group 2 had no merged value (is_set=false), so update_batch value=6 wins. let expect: PrimitiveArray = - Int64Array::from(vec![Some(1), Some(6), Some(6), None]); + Int64Array::from(vec![Some(1), Some(-6), Some(6), None]); assert_eq!(eval_result, &expect); @@ -1680,7 +1718,7 @@ mod tests { group_acc.compute_size_of_orderings() ); - group_acc.merge_batch(&s, &Vec::from_iter(0..s[0].len()), None, 100)?; + group_acc.merge_batch(&s, &Vec::from_iter(0..s[0].len()), 100)?; assert_eq!( group_acc.size_of_orderings, group_acc.compute_size_of_orderings() @@ -1753,12 +1791,7 @@ mod tests { ]; assert_eq!(state, expected_state); - group_acc.merge_batch( - &state, - &[0, 1, 2], - Some(&BooleanArray::from(vec![true, false, false])), - 3, - )?; + group_acc.merge_batch(&state, &[0, 1, 2], 3)?; val_with_orderings.clear(); val_with_orderings.push(Arc::new(Int64Array::from(vec![66, 6]))); @@ -1769,6 +1802,10 @@ mod tests { let binding = group_acc.evaluate(EmitTo::All)?; let eval_result = binding.as_any().downcast_ref::().unwrap(); + // group 0: merged value=1 (ordering=1, is_set=true), update not called. + // group 1: merged value=-6 (ordering=-6, is_set=true); update ordering=66 > -6 + // → LAST_VALUE keeps the higher ordering, so group 1 becomes 66. + // group 2: is_set=false after merge; update_batch sets it to 6. let expect: PrimitiveArray = Int64Array::from(vec![Some(1), Some(66), Some(6), None]); @@ -1777,6 +1814,118 @@ mod tests { Ok(()) } + /// Rows whose FILTER predicate evaluates to `null` must not pass the + /// filter, even when the underlying value bit at the null slot is `true` + /// (#22666). + #[test] + fn test_group_acc_filter_null_predicate() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("c", DataType::Int64, true), + ])); + + let sort_keys = [PhysicalSortExpr { + expr: col("c", &schema).unwrap(), + options: SortOptions::default(), + }]; + + let mut group_acc = FirstLastGroupsAccumulator::try_new( + PrimitiveValueState::::new(DataType::Int64), + sort_keys.into(), + true, + &[DataType::Int64], + true, + )?; + + let val_with_orderings: Vec = vec![ + Arc::new(Int64Array::from(vec![10, 20, 30])), + Arc::new(Int64Array::from(vec![10, 20, 30])), + ]; + + // Row 0: predicate is null (but its value bit is true, as produced by + // kernels such as `b < 1` when the null slot's underlying value is 0) + // Row 1: predicate is false + // Row 2: predicate is true + let filter = BooleanArray::new( + BooleanBuffer::from(vec![false, true, false, true]), + Some(NullBuffer::from(BooleanBuffer::from(vec![ + true, false, true, true, + ]))), + ) + .slice(1, 3); + assert_eq!(filter.offset(), 1); + + group_acc.update_batch(&val_with_orderings, &[0, 0, 1], Some(&filter), 2)?; + + let binding = group_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_any().downcast_ref::().unwrap(); + + // Group 0 has no row with a `true` predicate, so it must stay unset. + // Group 1 takes the only row with a `true` predicate. + let expect: PrimitiveArray = Int64Array::from(vec![None, Some(30)]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + + /// `convert_to_state` stores the user FILTER clause (including its nulls) + /// in the `is_set` state column, so `merge_batch` must not treat a null + /// `is_set` entry with a set value bit as "is set" (#22666). + #[test] + fn test_group_acc_merge_null_is_set() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("c", DataType::Int64, true), + ])); + + let sort_keys = [PhysicalSortExpr { + expr: col("c", &schema).unwrap(), + options: SortOptions::default(), + }]; + + let group_acc = FirstLastGroupsAccumulator::try_new( + PrimitiveValueState::::new(DataType::Int64), + sort_keys.clone().into(), + true, + &[DataType::Int64], + true, + )?; + + let val_with_orderings: Vec = vec![ + Arc::new(Int64Array::from(vec![10, 20])), + Arc::new(Int64Array::from(vec![10, 20])), + ]; + + // Same null-with-set-value-bit filter as above, carried into the state + let filter = BooleanArray::new( + BooleanBuffer::from(vec![true, true]), + Some(NullBuffer::from(BooleanBuffer::from(vec![false, true]))), + ); + + let state = group_acc.convert_to_state(&val_with_orderings, Some(&filter))?; + assert_eq!(state.len(), 3); + + let mut merging_acc = FirstLastGroupsAccumulator::try_new( + PrimitiveValueState::::new(DataType::Int64), + sort_keys.into(), + true, + &[DataType::Int64], + true, + )?; + + merging_acc.merge_batch(&state, &[0, 0], 1)?; + + let binding = merging_acc.evaluate(EmitTo::All)?; + let eval_result = binding.as_any().downcast_ref::().unwrap(); + + // Only the second row is valid and passes; the null-predicate row must + // be skipped even though its value bit is true. + let expect: PrimitiveArray = Int64Array::from(vec![Some(20)]); + assert_eq!(eval_result, &expect); + + Ok(()) + } + #[test] fn test_first_list_acc_size() -> Result<()> { fn size_after_batch(values: &[ArrayRef]) -> Result { @@ -1922,4 +2071,319 @@ mod tests { Ok(()) } + + /// End-to-end integration test for the nested-type support added to + /// [`FirstLastGroupsAccumulator`]: build the accumulator directly with a + /// [`GenericValueState`] for `List` and verify that winners are + /// selected correctly across multiple batches. + /// + /// Mirrors the shape produced by SQL like: + /// ```sql + /// SELECT first_value(list_col ORDER BY o DESC) FROM t GROUP BY p + /// ``` + /// which previously fell back to the per-group `Accumulator` path and + /// blew up on wide payloads. + #[test] + fn test_first_group_acc_list_int32() -> Result<()> { + let value_type = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let schema = Arc::new(Schema::new(vec![ + Field::new("val", value_type.clone(), true), + Field::new("ord", DataType::Int64, true), + ])); + let sort_keys = [PhysicalSortExpr { + expr: col("ord", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }]; + + let mut group_acc = FirstLastGroupsAccumulator::try_new( + GenericValueState::new(value_type.clone()), + sort_keys.into(), + false, + &[DataType::Int64], + /* pick_first = */ true, + )?; + + // Batch 1: four rows across two groups. + // Winners (largest ord per group with pick_first=true + DESC): + // group 0 -> ord=30 -> [3, 3, 3] + // group 1 -> ord=40 -> [4, 4, 4, 4] + let values_1 = ListArray::from_iter_primitive::([ + Some(vec![Some(1)]), + Some(vec![Some(2), Some(2)]), + Some(vec![Some(3), Some(3), Some(3)]), + Some(vec![Some(4), Some(4), Some(4), Some(4)]), + ]); + let orderings_1 = Int64Array::from(vec![10, 20, 30, 40]); + group_acc.update_batch( + &[ + Arc::new(values_1) as ArrayRef, + Arc::new(orderings_1) as ArrayRef, + ], + &[0, 1, 0, 1], + None, + 2, + )?; + + // Batch 2: group 0 gets a new winner ord=50 -> [9, 9]; group 1 + // keeps its previous winner (5 < 40). + let values_2 = ListArray::from_iter_primitive::([ + Some(vec![Some(9), Some(9)]), + Some(vec![Some(8)]), + ]); + let orderings_2 = Int64Array::from(vec![50, 5]); + group_acc.update_batch( + &[ + Arc::new(values_2) as ArrayRef, + Arc::new(orderings_2) as ArrayRef, + ], + &[0, 1], + None, + 2, + )?; + + let result = group_acc.evaluate(EmitTo::All)?; + let result = result.as_list::(); + assert_eq!(result.len(), 2); + let g0 = result.value(0); + let g0 = g0.as_primitive::(); + assert_eq!(g0.len(), 2); + assert_eq!(g0.value(0), 9); + assert_eq!(g0.value(1), 9); + let g1 = result.value(1); + let g1 = g1.as_primitive::(); + assert_eq!(g1.len(), 4); + for i in 0..4 { + assert_eq!(g1.value(i), 4); + } + Ok(()) + } + + /// Regression test for the wide-payload memory blow-up: run the full + /// aggregate loop over a batch large enough that the per-group + /// `Accumulator` path would have generated N * batch-worth of state + /// (via `ScalarValue::List` clones) and verify that the reported + /// accumulator size stays proportional to `#groups`, not `#rows`. + #[test] + fn test_first_group_acc_list_size_bounded_by_groups() -> Result<()> { + let value_type = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let schema = Arc::new(Schema::new(vec![ + Field::new("val", value_type.clone(), true), + Field::new("ord", DataType::Int64, true), + ])); + let sort_keys = [PhysicalSortExpr { + expr: col("ord", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }]; + let mut group_acc = FirstLastGroupsAccumulator::try_new( + GenericValueState::new(value_type), + sort_keys.into(), + false, + &[DataType::Int64], + true, + )?; + + // 10 groups × 10_000 candidate rows per group (100_000 total). Each + // list value has ~10 elements. Under the old per-group `Accumulator` + // + Arc-slice code path this would pin every batch in memory. + const GROUPS: usize = 10; + const ROWS_PER_GROUP: usize = 10_000; + const N: usize = GROUPS * ROWS_PER_GROUP; + let values = ListArray::from_iter_primitive::( + repeat_with(|| Some(vec![Some(1_i32); 10])).take(N), + ); + let orderings = Int64Array::from((0..N as i64).collect::>()); + let group_indices: Vec = (0..N).map(|i| i % GROUPS).collect(); + + group_acc.update_batch( + &[ + Arc::new(values) as ArrayRef, + Arc::new(orderings) as ArrayRef, + ], + &group_indices, + None, + GROUPS, + )?; + + // Sanity: the retained size must be small — well under what a single + // input batch worth of list buffers would occupy. The exact number is + // implementation-dependent, but should be O(GROUPS * per-list), not + // O(N * per-list). + let size = group_acc.size(); + assert!( + size < 100_000, + "accumulator size {size} bytes is not bounded by #groups (10 groups × ~10 int32 list elements)" + ); + + // Winner per group is the row with the largest ord — with our layout + // that's the last row assigned to each group. + let result = group_acc.evaluate(EmitTo::All)?; + let result = result.as_list::(); + assert_eq!(result.len(), GROUPS); + for g in 0..GROUPS { + let winner = result.value(g); + let winner = winner.as_primitive::(); + assert_eq!(winner.len(), 10); + for i in 0..10 { + assert_eq!(winner.value(i), 1); + } + } + Ok(()) + } + + /// End-to-end memory-savings regression test. + /// + /// Streams many independent batches of wide `List` payload through + /// the accumulator, dropping each source batch immediately after feeding + /// it in. The test then verifies three things: + /// + /// 1. The accumulator still emits the correct winners after every + /// source batch has been dropped (proves that stored values are + /// owned copies, not `Arc` slices into batches that no longer + /// exist). + /// 2. No buffer of any past source batch is shared by the emitted + /// output — the raw data-buffer pointer of every source batch is + /// recorded, and the final output's buffers must not alias any of + /// them (proves `compact()` copied the winners into owned memory). + /// 3. The accumulator's reported `size()` stays bounded by + /// `#groups * per-group-cost`, independent of `#batches * #rows`. + /// + /// This is the regression test for the wide-payload pinning behaviour + /// that motivated this PR. + #[test] + fn test_first_group_acc_list_no_source_batch_pinning() -> Result<()> { + let value_type = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let schema = Arc::new(Schema::new(vec![ + Field::new("val", value_type.clone(), true), + Field::new("ord", DataType::Int64, true), + ])); + let sort_keys = [PhysicalSortExpr { + expr: col("ord", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }]; + let mut group_acc = FirstLastGroupsAccumulator::try_new( + GenericValueState::new(value_type), + sort_keys.into(), + false, + &[DataType::Int64], + true, + )?; + + const GROUPS: usize = 4; + const BATCHES: usize = 50; + const ROWS_PER_BATCH: usize = 256; + + // Record the raw pointer of each source batch's Int32 value-data + // buffer. If `compact()` did its job, the accumulator's final + // output must not share any of these pointers — every winner + // value should have been copied into an owned buffer. + let mut source_value_ptrs: Vec<*const u8> = Vec::with_capacity(BATCHES); + + // Track the running-max ord we have fed to each group so the test's + // "expected winner" oracle matches the accumulator's choice. + let mut expected_ord = [i64::MIN; GROUPS]; + let mut expected_val_repeat = [0_i32; GROUPS]; + + for batch in 0..BATCHES { + // Each batch's list values are `[batch as i32; group_idx + 1]` + // — a distinct payload per (batch, row) so we can verify the + // winner by content. + let values = ListArray::from_iter_primitive::( + (0..ROWS_PER_BATCH).map(|i| { + let g = i % GROUPS; + Some(vec![Some(batch as i32); g + 1]) + }), + ); + let orderings = Int64Array::from( + (0..ROWS_PER_BATCH as i64) + .map(|i| batch as i64 * ROWS_PER_BATCH as i64 + i) + .collect::>(), + ); + let group_indices: Vec = + (0..ROWS_PER_BATCH).map(|i| i % GROUPS).collect(); + + // Update the oracle: the last row in this batch that hits each + // group has the largest ord for that group in this batch. + for i in (0..ROWS_PER_BATCH).rev() { + let g = i % GROUPS; + let ord = batch as i64 * ROWS_PER_BATCH as i64 + i as i64; + if ord > expected_ord[g] { + expected_ord[g] = ord; + expected_val_repeat[g] = batch as i32; + } + } + + // Capture the raw pointer of this batch's Int32 value-data + // buffer *before* handing ownership to the accumulator. Int32 + // arrays have a single value buffer at index 0. + source_value_ptrs.push(values.values().to_data().buffers()[0].as_ptr()); + + let values_arc: Arc = Arc::new(values); + let orderings_arc: Arc = Arc::new(orderings); + + group_acc.update_batch( + &[values_arc, orderings_arc], + &group_indices, + None, + GROUPS, + )?; + + // Drop happens implicitly at end of scope. + } + + // (2) Size is bounded by #groups. The exact number is + // implementation-dependent but should be orders of magnitude below + // `BATCHES * ROWS_PER_BATCH * per-list-cost` (the amount that would + // be retained under the old Arc-slice pinning bug). + let size = group_acc.size(); + assert!( + size < 10_000, + "accumulator size {size} bytes is not bounded by #groups \ + (expected O({GROUPS}) not O({BATCHES} * {ROWS_PER_BATCH}))" + ); + + // (1) Winners are still readable and match the oracle. + let result = group_acc.evaluate(EmitTo::All)?; + let result_list = result.as_list::(); + assert_eq!(result_list.len(), GROUPS); + for (g, expected_repeat) in expected_val_repeat.iter().enumerate().take(GROUPS) { + let winner = result_list.value(g); + let winner = winner.as_primitive::(); + assert_eq!(winner.len(), g + 1, "winner list length for group {g}"); + for i in 0..winner.len() { + assert_eq!( + winner.value(i), + *expected_repeat, + "winner payload mismatch for group {g}" + ); + } + } + + // (3) The critical byte-level check: the emitted output's Int32 + // value-data buffer must NOT share a raw pointer with any of the + // source batches. If `compact()` were omitted, `list_array.value(i)` + // would yield a slice whose backing buffer points into the source + // batch — the accumulator would then either pin the batch or emit + // an output that shares its buffer. + let result_values_ptr = result_list.values().to_data().buffers()[0].as_ptr(); + for (i, src_ptr) in source_value_ptrs.iter().enumerate() { + assert_ne!( + *src_ptr, result_values_ptr, + "emitted result's Int32 value buffer aliases source batch \ + {i}'s buffer; compact() is not making an owned copy" + ); + } + Ok(()) + } } diff --git a/datafusion/functions-aggregate/src/first_last/state.rs b/datafusion/functions-aggregate/src/first_last/state.rs index cd7114bf04f9c..d99b4f6ecc6da 100644 --- a/datafusion/functions-aggregate/src/first_last/state.rs +++ b/datafusion/functions-aggregate/src/first_last/state.rs @@ -25,7 +25,7 @@ use arrow::array::{ }; use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::DataType; -use datafusion_common::{Result, internal_err}; +use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::EmitTo; pub(crate) trait ValueState: Send + Sync { @@ -290,6 +290,78 @@ impl BytesValueState { } } +/// Fallback state for arbitrary Arrow types (List, LargeList, Struct, Map, ...) +/// that are not covered by [`PrimitiveValueState`] or [`BytesValueState`]. +/// +/// Stores one [`ScalarValue`] per group. Winners are identified by the +/// vectorized comparator in the enclosing accumulator, so the per-row +/// allocation cost of the fallback `Accumulator` path is avoided: +/// `ScalarValue::try_from_array` is called once per group per batch (at +/// winner-update time), not once per candidate row. +pub(crate) struct GenericValueState { + vals: Vec>, + data_type: DataType, + /// Cached total heap size of `vals`, updated on each mutation to avoid + /// walking the vector on every `size()` call. + total_size: usize, +} + +impl GenericValueState { + pub(crate) fn new(data_type: DataType) -> Self { + Self { + vals: vec![], + data_type, + total_size: 0, + } + } +} + +impl ValueState for GenericValueState { + fn resize(&mut self, new_size: usize) { + if new_size < self.vals.len() { + for v in self.vals[new_size..].iter().flatten() { + self.total_size -= v.size(); + } + } + self.vals.resize(new_size, None); + } + + fn update(&mut self, group_idx: usize, array: &ArrayRef, idx: usize) -> Result<()> { + if let Some(v) = &self.vals[group_idx] { + self.total_size -= v.size(); + } + let mut scalar = ScalarValue::try_from_array(array, idx)?; + // `try_from_array` for nested types returns Arc slices into the source + // batch buffers, so a single stored winner would pin the entire batch + // in memory. Compact copies the referenced bytes into an owned buffer + // so old batches can be dropped as new ones arrive. + scalar.compact(); + self.total_size += scalar.size(); + self.vals[group_idx] = Some(scalar); + Ok(()) + } + + fn take(&mut self, emit_to: EmitTo) -> Result { + let taken = emit_to.take_needed(&mut self.vals); + let taken_size: usize = taken.iter().flatten().map(|v| v.size()).sum(); + self.total_size -= taken_size; + + let default = ScalarValue::try_from(&self.data_type)?; + let scalars = taken + .into_iter() + .map(|opt| opt.unwrap_or_else(|| default.clone())) + .collect::>(); + if scalars.is_empty() { + return Ok(arrow::array::new_empty_array(&self.data_type)); + } + ScalarValue::iter_to_array(scalars) + } + + fn size(&self) -> usize { + self.vals.capacity() * size_of::>() + self.total_size + } +} + pub(crate) fn take_need( bool_buf_builder: &mut BooleanBufferBuilder, emit_to: EmitTo, @@ -312,9 +384,12 @@ pub(crate) fn take_need( mod tests { use super::*; use arrow::array::{ - BinaryArray, BinaryViewArray, LargeBinaryArray, LargeStringArray, StringArray, - StringViewArray, + Array, BinaryArray, BinaryViewArray, FixedSizeListArray, Int32Array, + Int32Builder, LargeBinaryArray, LargeListArray, LargeStringArray, ListBuilder, + MapArray, StringArray, StringBuilder, StringViewArray, StructArray, }; + use arrow::buffer::{OffsetBuffer, ScalarBuffer}; + use arrow::datatypes::{DataType, Field, Fields}; #[test] fn test_bytes_value_state_utf8() -> Result<()> { @@ -459,4 +534,382 @@ mod tests { Ok(()) } + + // ---------- GenericValueState (nested types) ---------- + + /// Build a `List` array with three rows: `["a"]`, `["b", "c"]`, + /// `["d", "e", "f"]`. Used by several tests. + fn make_list_utf8_array() -> ArrayRef { + let mut builder = ListBuilder::new(StringBuilder::new()); + builder.values().append_value("a"); + builder.append(true); + builder.values().append_value("b"); + builder.values().append_value("c"); + builder.append(true); + builder.values().append_value("d"); + builder.values().append_value("e"); + builder.values().append_value("f"); + builder.append(true); + Arc::new(builder.finish()) + } + + #[test] + fn test_generic_value_state_list_utf8() -> Result<()> { + let list_utf8 = + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))); + let mut state = GenericValueState::new(list_utf8.clone()); + state.resize(2); + + let array = make_list_utf8_array(); + + // group 0 <- ["a"] ; group 1 <- ["b", "c"] + state.update(0, &array, 0)?; + state.update(1, &array, 1)?; + + // Overwrite group 0 with the wider ["d", "e", "f"] (size-accounting + // must decrement the old value before adding the new one). + let size_after_first = state.total_size; + state.update(0, &array, 2)?; + assert!( + state.total_size > 0, + "total_size must remain positive after overwrite" + ); + // The overwrite replaced group 0's payload; the delta relative to the + // previous state should equal `new.size() - old.size()`. If the caller + // forgot to subtract the old size, `total_size` would drift upward. + let expected_delta = { + let new_scalar = { + let mut s = ScalarValue::try_from_array(&array, 2)?; + s.compact(); + s + }; + let old_scalar = { + let mut s = ScalarValue::try_from_array(&array, 0)?; + s.compact(); + s + }; + new_scalar.size() as isize - old_scalar.size() as isize + }; + assert_eq!( + state.total_size as isize - size_after_first as isize, + expected_delta, + "size accounting drifted after overwrite" + ); + + let result = state.take(EmitTo::All)?; + let result = result.as_list::(); + assert_eq!(result.len(), 2); + let g0 = result.value(0); + let g0 = g0.as_any().downcast_ref::().unwrap(); + assert_eq!(g0.value(0), "d"); + assert_eq!(g0.value(1), "e"); + assert_eq!(g0.value(2), "f"); + let g1 = result.value(1); + let g1 = g1.as_any().downcast_ref::().unwrap(); + assert_eq!(g1.value(0), "b"); + assert_eq!(g1.value(1), "c"); + + assert_eq!(state.total_size, 0, "state must be fully drained"); + Ok(()) + } + + #[test] + fn test_generic_value_state_struct() -> Result<()> { + let fields = Fields::from(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, false), + ]); + let struct_type = DataType::Struct(fields.clone()); + + let id = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let name = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef; + let struct_array = + Arc::new(StructArray::new(fields, vec![id, name], None)) as ArrayRef; + + let mut state = GenericValueState::new(struct_type); + state.resize(2); + state.update(0, &struct_array, 0)?; + state.update(1, &struct_array, 2)?; + + let out = state.take(EmitTo::All)?; + let out = out.as_any().downcast_ref::().unwrap(); + assert_eq!(out.len(), 2); + + let out_id = out.column(0).as_any().downcast_ref::().unwrap(); + let out_name = out + .column(1) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(out_id.value(0), 1); + assert_eq!(out_id.value(1), 3); + assert_eq!(out_name.value(0), "a"); + assert_eq!(out_name.value(1), "c"); + Ok(()) + } + + #[test] + fn test_generic_value_state_large_list() -> Result<()> { + let field = Arc::new(Field::new("item", DataType::Int32, true)); + let large_list_type = DataType::LargeList(Arc::clone(&field)); + + let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let offsets: OffsetBuffer = + OffsetBuffer::new(ScalarBuffer::from(vec![0_i64, 2, 5, 6])); + let array = Arc::new(LargeListArray::new(field, offsets, Arc::new(values), None)) + as ArrayRef; + + let mut state = GenericValueState::new(large_list_type); + state.resize(2); + state.update(0, &array, 0)?; // [1, 2] + state.update(1, &array, 2)?; // [6] + + let out = state.take(EmitTo::All)?; + let out = out.as_list::(); + assert_eq!(out.len(), 2); + let g0 = out.value(0); + let g0 = g0.as_any().downcast_ref::().unwrap(); + assert_eq!(g0.len(), 2); + assert_eq!(g0.value(0), 1); + assert_eq!(g0.value(1), 2); + let g1 = out.value(1); + let g1 = g1.as_any().downcast_ref::().unwrap(); + assert_eq!(g1.len(), 1); + assert_eq!(g1.value(0), 6); + Ok(()) + } + + #[test] + fn test_generic_value_state_fixed_size_list() -> Result<()> { + let field = Arc::new(Field::new("item", DataType::Int32, true)); + let fsl_type = DataType::FixedSizeList(Arc::clone(&field), 2); + + let values = Int32Array::from(vec![1, 2, 3, 4, 5, 6]); + let array = Arc::new(FixedSizeListArray::new(field, 2, Arc::new(values), None)) + as ArrayRef; + + let mut state = GenericValueState::new(fsl_type); + state.resize(2); + state.update(0, &array, 0)?; // [1, 2] + state.update(1, &array, 2)?; // [5, 6] + + let out = state.take(EmitTo::All)?; + let out = out + .as_any() + .downcast_ref::() + .expect("emitted FixedSizeListArray"); + assert_eq!(out.len(), 2); + let g0 = out.value(0); + let g0 = g0.as_any().downcast_ref::().unwrap(); + assert_eq!(g0.value(0), 1); + assert_eq!(g0.value(1), 2); + let g1 = out.value(1); + let g1 = g1.as_any().downcast_ref::().unwrap(); + assert_eq!(g1.value(0), 5); + assert_eq!(g1.value(1), 6); + Ok(()) + } + + #[test] + fn test_generic_value_state_map() -> Result<()> { + // Map with two entries: {"a": 1, "b": 2}, {"c": 3} + let keys = Arc::new(StringArray::from(vec!["a", "b", "c"])) as ArrayRef; + let values = Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef; + let entry_fields = Fields::from(vec![ + Field::new("keys", DataType::Utf8, false), + Field::new("values", DataType::Int32, true), + ]); + let entries = StructArray::new(entry_fields.clone(), vec![keys, values], None); + let offsets: OffsetBuffer = + OffsetBuffer::new(ScalarBuffer::from(vec![0_i32, 2, 3])); + let map_field = + Arc::new(Field::new("entries", DataType::Struct(entry_fields), false)); + let map_array = Arc::new(MapArray::new( + Arc::clone(&map_field), + offsets, + entries, + None, + false, + )) as ArrayRef; + let map_type = DataType::Map(map_field, false); + + let mut state = GenericValueState::new(map_type); + state.resize(2); + state.update(0, &map_array, 0)?; // {"a": 1, "b": 2} + state.update(1, &map_array, 1)?; // {"c": 3} + + let out = state.take(EmitTo::All)?; + let out = out.as_any().downcast_ref::().unwrap(); + assert_eq!(out.len(), 2); + assert_eq!(out.value_length(0), 2); + assert_eq!(out.value_length(1), 1); + Ok(()) + } + + #[test] + fn test_generic_value_state_emit_first() -> Result<()> { + let list_utf8 = + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))); + let mut state = GenericValueState::new(list_utf8); + state.resize(3); + + let array = make_list_utf8_array(); + state.update(0, &array, 0)?; + state.update(1, &array, 1)?; + state.update(2, &array, 2)?; + + let after_all_updates = state.total_size; + assert!(after_all_updates > 0); + + // Emit the first 2 groups; remaining group 2 stays. + let head = state.take(EmitTo::First(2))?; + let head = head.as_list::(); + assert_eq!(head.len(), 2); + let h0 = head.value(0); + let h0 = h0.as_any().downcast_ref::().unwrap(); + assert_eq!(h0.value(0), "a"); + let h1 = head.value(1); + let h1 = h1.as_any().downcast_ref::().unwrap(); + assert_eq!(h1.value(0), "b"); + assert_eq!(h1.value(1), "c"); + + // After partial emit, `total_size` shrank but is still positive. + assert!(state.total_size > 0); + assert!(state.total_size < after_all_updates); + + let tail = state.take(EmitTo::All)?; + let tail = tail.as_list::(); + assert_eq!(tail.len(), 1); + let t0 = tail.value(0); + let t0 = t0.as_any().downcast_ref::().unwrap(); + assert_eq!(t0.value(0), "d"); + assert_eq!(t0.value(1), "e"); + assert_eq!(t0.value(2), "f"); + + assert_eq!(state.total_size, 0); + Ok(()) + } + + #[test] + fn test_generic_value_state_update_null() -> Result<()> { + // List with rows: [1, 2], NULL + let mut builder = ListBuilder::new(Int32Builder::new()); + builder.values().append_value(1); + builder.values().append_value(2); + builder.append(true); + builder.append(false); // null entry + let array: ArrayRef = Arc::new(builder.finish()); + + let list_type = array.data_type().clone(); + let mut state = GenericValueState::new(list_type); + state.resize(1); + + // group 0 = [1, 2] + state.update(0, &array, 0)?; + let size_after_value = state.total_size; + assert!(size_after_value > 0); + + // Overwrite group 0 with NULL. The size accounting must subtract the + // previous value's size and then add the null-scalar's size; the point + // of this test is that `total_size` stays consistent (no drift) and + // the null is emitted correctly. + state.update(0, &array, 1)?; + // Recomputing from scratch must match the cached total_size. + let recomputed: usize = state.vals.iter().flatten().map(|v| v.size()).sum(); + assert_eq!( + state.total_size, recomputed, + "total_size drifted after null update" + ); + + let out = state.take(EmitTo::All)?; + let out = out.as_list::(); + assert_eq!(out.len(), 1); + assert!(out.is_null(0)); + assert_eq!(state.total_size, 0); + Ok(()) + } + + #[test] + fn test_generic_value_state_compact_releases_parent_batch() -> Result<()> { + // Regression test for the memory-pinning bug: without compact(), + // `ScalarValue::try_from_array` on a List column produces a + // ScalarValue whose child values array is an Arrow slice pointing + // into the *source* batch's underlying byte buffer. That means the + // source batch's memory stays alive until every extracted winner + // is dropped, even if the outer ListArray is released. `compact()` + // must copy the referenced bytes into a fresh owned buffer. + // + // Correctly detecting this requires comparing the raw buffer + // pointer of the source `Utf8` value-data buffer against the raw + // buffer pointer of the stored winner's value-data buffer. Checking + // `Arc::strong_count` on the outer `ArrayRef` is not sufficient, + // because `list_array.value(idx)` returns a sliced child that keeps + // its own Arc chain independent of the outer ListArray. + let list_utf8 = + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))); + let mut state = GenericValueState::new(list_utf8); + state.resize(1); + + let array: ArrayRef = make_list_utf8_array(); + + // Capture the raw pointer of the *source* Utf8 value-data buffer. + // Utf8Array has two buffers: offsets (buffer 0) and value bytes + // (buffer 1). Comparing buffer 1 is the direct check for byte + // pinning. + let source_values_ptr = + array.as_list::().values().to_data().buffers()[1].as_ptr(); + + state.update(0, &array, 0)?; + drop(array); + + // Directly probe the stored ScalarValue's underlying values buffer. + let stored_values_ptr = match state + .vals + .first() + .and_then(|opt| opt.as_ref()) + .expect("group 0 should have a stored value") + { + ScalarValue::List(list_arr) => { + list_arr.values().to_data().buffers()[1].as_ptr() + } + other => panic!("expected ScalarValue::List, got {other:?}"), + }; + + assert_ne!( + source_values_ptr, stored_values_ptr, + "compact() failed: stored ScalarValue still shares the source \ + batch's Utf8 value-data buffer, meaning the batch is pinned in \ + memory even after the outer ArrayRef is dropped" + ); + + // Data must still be readable from the stored copy. + let out = state.take(EmitTo::All)?; + let out = out.as_list::(); + assert_eq!(out.len(), 1); + let g0 = out.value(0); + let g0 = g0.as_any().downcast_ref::().unwrap(); + assert_eq!(g0.value(0), "a"); + Ok(()) + } + + #[test] + fn test_generic_value_state_resize_shrink_recovers_size() -> Result<()> { + let list_utf8 = + DataType::List(Arc::new(Field::new("item", DataType::Utf8, true))); + let mut state = GenericValueState::new(list_utf8); + state.resize(3); + + let array = make_list_utf8_array(); + state.update(0, &array, 0)?; + state.update(1, &array, 1)?; + state.update(2, &array, 2)?; + let full_size = state.total_size; + assert!(full_size > 0); + + // Shrinking must subtract the dropped groups' sizes from total_size. + state.resize(1); + assert!(state.total_size > 0); + assert!(state.total_size < full_size); + Ok(()) + } } diff --git a/datafusion/functions-aggregate/src/hyperloglog.rs b/datafusion/functions-aggregate/src/hyperloglog.rs index 3861800847edb..9968e5a98194f 100644 --- a/datafusion/functions-aggregate/src/hyperloglog.rs +++ b/datafusion/functions-aggregate/src/hyperloglog.rs @@ -42,7 +42,7 @@ use std::marker::PhantomData; const HLL_P: usize = 14_usize; /// The number of bits of the hash value used determining the number of leading zeros const HLL_Q: usize = 64_usize - HLL_P; -const NUM_REGISTERS: usize = 1_usize << HLL_P; +pub(crate) const NUM_REGISTERS: usize = 1_usize << HLL_P; /// Mask to obtain index into the registers const HLL_P_MASK: u64 = (NUM_REGISTERS as u64) - 1; @@ -55,14 +55,7 @@ where phantom: PhantomData, } -/// Fixed seed for the hashing so that values are consistent across runs -/// -/// Note that when we later move on to have serialized HLL register binaries -/// shared across cluster, this HLL_HASH_STATE will have to be consistent across all -/// parties otherwise we might have corruption. So ideally for later this seed -/// shall be part of the serialized form (or stay unchanged across versions). -pub(crate) const HLL_HASH_STATE: foldhash::quality::FixedState = - foldhash::quality::FixedState::with_seed(0); +pub(crate) use datafusion_common::hash_utils::HLL_RANDOM_STATE as HLL_HASH_STATE; impl Default for HyperLogLog where @@ -93,9 +86,8 @@ where } } - /// choice of hash function: foldhash is already an dependency - /// and it fits the requirements of being a 64bit hash with - /// reasonable performance. + /// The HLL hash state is shared through `datafusion_common::hash_utils` + /// so sketches remain compatible across accumulators. #[inline] fn hash_value(&self, obj: &T) -> u64 { HLL_HASH_STATE.hash_one(obj) @@ -145,16 +137,69 @@ where /// Guess the number of unique elements seen by the HyperLogLog. pub fn count(&self) -> usize { - let histogram = self.get_histogram(); - let m = NUM_REGISTERS as f64; - let mut z = m * hll_tau((m - histogram[HLL_Q + 1] as f64) / m); - for i in histogram[1..=HLL_Q].iter().rev() { - z += *i as f64; - z *= 0.5; + count_from_histogram(&self.get_histogram()) + } +} + +/// Compute `index` and `rho` (register value) for a precomputed hash, exactly as +/// [`HyperLogLog::add_hashed`] does. +#[inline] +pub(crate) fn register_for_hash(hash: u64) -> (usize, u8) { + let index = (hash & HLL_P_MASK) as usize; + let rho = (((hash >> HLL_P) | (1_u64 << HLL_Q)).trailing_zeros() + 1) as u8; + (index, rho) +} + +/// Estimate the cardinality of a set of precomputed hashes without +/// materializing a full [`NUM_REGISTERS`]-byte register array. +/// +/// This is equivalent to adding every hash to a fresh [`HyperLogLog`] via +/// [`HyperLogLog::add_hashed`] and calling [`HyperLogLog::count`], but only does +/// work proportional to the number of hashes. It is used to cheaply estimate the +/// many small groups produced by a high-cardinality `GROUP BY`, where allocating +/// and scanning a 16 KiB sketch per group would dominate the runtime. +/// +/// `hashes` may contain duplicates (duplicate hashes are idempotent). +pub(crate) fn count_from_hashes(hashes: &[u64]) -> usize { + if hashes.is_empty() { + return 0; + } + // For each touched register index keep the maximum rho. Sorting by + // (index, rho) groups equal indices together with the max rho last. + let mut idx_rho: Vec<(usize, u8)> = + hashes.iter().map(|&hash| register_for_hash(hash)).collect(); + idx_rho.sort_unstable(); + + let mut histogram = [0u32; HLL_Q + 2]; + let mut touched = 0u32; + let mut i = 0; + while i < idx_rho.len() { + let index = idx_rho[i].0; + let mut max_rho = idx_rho[i].1; + i += 1; + while i < idx_rho.len() && idx_rho[i].0 == index { + max_rho = idx_rho[i].1; // ascending rho => last is the max + i += 1; } - z += m * hll_sigma(histogram[0] as f64 / m); - (0.5 / 2_f64.ln() * m * m / z).round() as usize + histogram[max_rho as usize] += 1; + touched += 1; + } + // All remaining registers are still zero. + histogram[0] = NUM_REGISTERS as u32 - touched; + count_from_histogram(&histogram) +} + +/// Apply the HyperLogLog cardinality estimator to a register histogram. +#[inline] +fn count_from_histogram(histogram: &[u32; HLL_Q + 2]) -> usize { + let m = NUM_REGISTERS as f64; + let mut z = m * hll_tau((m - histogram[HLL_Q + 1] as f64) / m); + for i in histogram[1..=HLL_Q].iter().rev() { + z += *i as f64; + z *= 0.5; } + z += m * hll_sigma(histogram[0] as f64 / m); + (0.5 / 2_f64.ln() * m * m / z).round() as usize } /// Helper function sigma as defined in diff --git a/datafusion/functions-aggregate/src/lib.rs b/datafusion/functions-aggregate/src/lib.rs index 1b9996220d882..e3f2714abbf25 100644 --- a/datafusion/functions-aggregate/src/lib.rs +++ b/datafusion/functions-aggregate/src/lib.rs @@ -65,6 +65,7 @@ #[macro_use] pub mod macros; +pub mod any_value; pub mod approx_distinct; pub mod approx_median; pub mod approx_percentile_cont; @@ -102,6 +103,7 @@ use std::sync::Arc; /// Fluent-style API for creating `Expr`s pub mod expr_fn { + pub use super::any_value::any_value; pub use super::approx_distinct::approx_distinct; pub use super::approx_median::approx_median; pub use super::approx_percentile_cont::approx_percentile_cont; @@ -147,6 +149,7 @@ pub mod expr_fn { /// Returns all default aggregate functions pub fn all_default_aggregate_functions() -> Vec> { vec![ + any_value::any_value_udaf(), array_agg::array_agg_udaf(), first_last::first_value_udaf(), first_last::last_value_udaf(), diff --git a/datafusion/functions-aggregate/src/median.rs b/datafusion/functions-aggregate/src/median.rs index 02a49ab6dcca0..fb74da87c7fc8 100644 --- a/datafusion/functions-aggregate/src/median.rs +++ b/datafusion/functions-aggregate/src/median.rs @@ -39,10 +39,11 @@ use arrow::datatypes::{ ArrowNativeType, ArrowPrimitiveType, Decimal32Type, Decimal64Type, FieldRef, }; +use datafusion_common::hash_utils::RandomState; use datafusion_common::types::{NativeType, logical_float64}; use datafusion_common::{ - DataFusionError, Result, ScalarValue, assert_eq_or_internal_err, - internal_datafusion_err, + DataFusionError, Result, ScalarValue, assert_eq_or_internal_err, exec_datafusion_err, + internal_datafusion_err, internal_err, }; use datafusion_expr::function::StateFieldsArgs; use datafusion_expr::{ @@ -52,6 +53,7 @@ use datafusion_expr::{ use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filtered_null_mask; +use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator; use datafusion_functions_aggregate_common::utils::{GenericDistinctBuffer, Hashable}; use datafusion_macros::user_doc; use std::collections::HashMap; @@ -137,6 +139,17 @@ impl AggregateUDFImpl for Median { } fn state_fields(&self, args: StateFieldsArgs) -> Result> { + if args.input_fields[0].data_type().is_null() { + return Ok(vec![ + Field::new( + format_state_name(args.name, self.name()), + DataType::Null, + true, + ) + .into(), + ]); + } + //Intermediate state is a list of the elements we have collected so far let field = Field::new_list_field(args.input_fields[0].data_type().clone(), true); let state_name = if args.is_distinct { @@ -173,6 +186,10 @@ impl AggregateUDFImpl for Median { } let dt = acc_args.expr_fields[0].data_type().clone(); + if dt.is_null() { + return Ok(Box::new(NoopAccumulator::default())); + } + downcast_integer! { dt => (helper, dt), DataType::Float16 => helper!(Float16Type, dt), @@ -191,7 +208,7 @@ impl AggregateUDFImpl for Median { } fn groups_accumulator_supported(&self, args: AccumulatorArgs) -> bool { - !args.is_distinct + !args.is_distinct && !args.expr_fields[0].data_type().is_null() } fn create_groups_accumulator( @@ -282,8 +299,18 @@ impl Accumulator for MedianAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let values = values[0].as_primitive::(); - self.all_values.reserve(values.len() - values.null_count()); - self.all_values.extend(values.iter().flatten()); + let additional = values.len() - values.null_count(); + self.all_values.try_reserve(additional).map_err(|e| { + exec_datafusion_err!( + "failed to reserve {additional} values for median accumulator: {e}" + ) + })?; + if values.null_count() > 0 { + self.all_values.extend(values.iter().flatten()); + } else { + // Fast path: no nulls, so the values buffer can be appended wholesale. + self.all_values.extend_from_slice(values.values()); + } Ok(()) } @@ -305,11 +332,19 @@ impl Accumulator for MedianAccumulator { } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let mut to_remove: HashMap, usize> = HashMap::new(); + let mut to_remove: HashMap, usize, RandomState> = + HashMap::default(); let arr = values[0].as_primitive::(); - for value in arr.iter().flatten() { - *to_remove.entry(Hashable(value)).or_default() += 1; + if arr.null_count() > 0 { + for value in arr.iter().flatten() { + *to_remove.entry(Hashable(value)).or_default() += 1; + } + } else { + // Fast path: no nulls, so skip the per-element validity check. + for value in arr.values().iter() { + *to_remove.entry(Hashable(*value)).or_default() += 1; + } } let mut i = 0; @@ -330,6 +365,15 @@ impl Accumulator for MedianAccumulator { i += 1; } } + + // Retracting values that are not tracked means the accumulator state + // has diverged from the window frame; continuing would silently + // produce wrong results, so surface it as an error. + if !to_remove.is_empty() { + return internal_err!( + "median retract_batch: retracted value(s) not present in the window" + ); + } Ok(()) } @@ -388,8 +432,6 @@ impl GroupsAccumulator for MedianGroupsAccumulator, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "one argument to merge_batch"); @@ -532,18 +574,14 @@ impl GroupsAccumulator for MedianGroupsAccumulator bool { - true - } - fn size(&self) -> usize { self.group_values .iter() - .map(|values| values.capacity() * size_of::()) + .map(|values| values.capacity() * size_of::()) .sum::() - // account for size of self.grou_values too - + self.group_values.capacity() * size_of::>() + // account for size of self.group_values too + + self.group_values.capacity() * size_of::>() + + size_of::>>() } } @@ -628,3 +666,58 @@ fn calculate_median(values: &mut [T::Native]) -> Option MedianAccumulator { + MedianAccumulator { + data_type: DataType::Float64, + all_values: vec![], + } + } + + #[test] + fn retract_batch_errors_on_untracked_value() { + let mut acc = median_accumulator(); + let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); + acc.update_batch(std::slice::from_ref(&values)).unwrap(); + + let retract: ArrayRef = Arc::new(Float64Array::from(vec![3.0])); + let err = acc + .retract_batch(std::slice::from_ref(&retract)) + .unwrap_err() + .to_string(); + assert!( + err.contains("not present in the window"), + "unexpected error: {err}" + ); + } + + #[test] + fn update_batch_with_and_without_nulls_agree() { + // The null-free fast path must accumulate the same values as the + // general path. + let dense: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])); + let sparse: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.0), + None, + Some(2.0), + None, + Some(3.0), + ])); + + let mut dense_acc = median_accumulator(); + dense_acc + .update_batch(std::slice::from_ref(&dense)) + .unwrap(); + let mut sparse_acc = median_accumulator(); + sparse_acc + .update_batch(std::slice::from_ref(&sparse)) + .unwrap(); + + assert_eq!(dense_acc.all_values, sparse_acc.all_values); + } +} diff --git a/datafusion/functions-aggregate/src/min_max.rs b/datafusion/functions-aggregate/src/min_max.rs index f4eaaab853464..41643747e8a42 100644 --- a/datafusion/functions-aggregate/src/min_max.rs +++ b/datafusion/functions-aggregate/src/min_max.rs @@ -52,7 +52,8 @@ use datafusion_expr::{ use datafusion_expr::{GroupsAccumulator, StatisticsArgs}; use datafusion_macros::user_doc; use half::f16; -use std::mem::size_of_val; +use std::collections::VecDeque; +use std::mem::{size_of, size_of_val}; use std::ops::Deref; fn get_min_max_result_type(input_types: &[DataType]) -> Result> { @@ -380,7 +381,8 @@ impl AggregateUDFImpl for Max { #[derive(Debug)] pub struct SlidingMaxAccumulator { - max: ScalarValue, + /// Typed NULL returned when the window contains no non-null values + empty_value: ScalarValue, moving_max: MovingMax, } @@ -388,30 +390,38 @@ impl SlidingMaxAccumulator { /// new max accumulator pub fn try_new(datatype: &DataType) -> Result { Ok(Self { - max: ScalarValue::try_from(datatype)?, + empty_value: ScalarValue::try_from(datatype)?, moving_max: MovingMax::::new(), }) } + + fn current_max(&self) -> ScalarValue { + match self.moving_max.max() { + Some(res) => res.clone(), + None => self.empty_value.clone(), + } + } } impl Accumulator for SlidingMaxAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { for idx in 0..values[0].len() { let val = ScalarValue::try_from_array(&values[0], idx)?; - self.moving_max.push(val); - } - if let Some(res) = self.moving_max.max() { - self.max = res.clone(); + if !val.is_null() { + self.moving_max.push(val); + } } Ok(()) } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - for _idx in 0..values[0].len() { - (self.moving_max).pop(); - } - if let Some(res) = self.moving_max.max() { - self.max = res.clone(); + // We assume that values are retracted in the order they were added, so + // the retracted values must be the oldest elements of `moving_max`. + // NULLs are never pushed, so be sure to only pop once per non-NULL + // value. + let valid_count = values[0].len() - values[0].logical_null_count(); + for _ in 0..valid_count { + self.moving_max.pop(); } Ok(()) } @@ -421,11 +431,11 @@ impl Accumulator for SlidingMaxAccumulator { } fn state(&mut self) -> Result> { - Ok(vec![self.max.clone()]) + Ok(vec![self.current_max()]) } fn evaluate(&mut self) -> Result { - Ok(self.max.clone()) + Ok(self.current_max()) } fn supports_retract_batch(&self) -> bool { @@ -433,7 +443,9 @@ impl Accumulator for SlidingMaxAccumulator { } fn size(&self) -> usize { - size_of_val(self) - size_of_val(&self.max) + self.max.size() + size_of_val(self) - size_of_val(&self.empty_value) + + self.empty_value.size() + + self.moving_max.heap_size(|sv| sv.size() - size_of_val(sv)) } } @@ -664,22 +676,30 @@ impl AggregateUDFImpl for Min { #[derive(Debug)] pub struct SlidingMinAccumulator { - min: ScalarValue, + /// Typed NULL returned when the window contains no non-null values + empty_value: ScalarValue, moving_min: MovingMin, } impl SlidingMinAccumulator { pub fn try_new(datatype: &DataType) -> Result { Ok(Self { - min: ScalarValue::try_from(datatype)?, + empty_value: ScalarValue::try_from(datatype)?, moving_min: MovingMin::::new(), }) } + + fn current_min(&self) -> ScalarValue { + match self.moving_min.min() { + Some(res) => res.clone(), + None => self.empty_value.clone(), + } + } } impl Accumulator for SlidingMinAccumulator { fn state(&mut self) -> Result> { - Ok(vec![self.min.clone()]) + Ok(vec![self.current_min()]) } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { @@ -689,21 +709,17 @@ impl Accumulator for SlidingMinAccumulator { self.moving_min.push(val); } } - if let Some(res) = self.moving_min.min() { - self.min = res.clone(); - } Ok(()) } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - for idx in 0..values[0].len() { - let val = ScalarValue::try_from_array(&values[0], idx)?; - if !val.is_null() { - (self.moving_min).pop(); - } - } - if let Some(res) = self.moving_min.min() { - self.min = res.clone(); + // We assume that values are retracted in the order they were added, so + // the retracted values must be the oldest elements of `moving_min`. + // NULLs are never pushed, so be sure to only pop once per non-NULL + // value. + let valid_count = values[0].len() - values[0].logical_null_count(); + for _ in 0..valid_count { + self.moving_min.pop(); } Ok(()) } @@ -713,7 +729,7 @@ impl Accumulator for SlidingMinAccumulator { } fn evaluate(&mut self) -> Result { - Ok(self.min.clone()) + Ok(self.current_min()) } fn supports_retract_batch(&self) -> bool { @@ -721,77 +737,55 @@ impl Accumulator for SlidingMinAccumulator { } fn size(&self) -> usize { - size_of_val(self) - size_of_val(&self.min) + self.min.size() + size_of_val(self) - size_of_val(&self.empty_value) + + self.empty_value.size() + + self.moving_min.heap_size(|sv| sv.size() - size_of_val(sv)) } } /// Keep track of the minimum value in a sliding window. /// -/// The implementation is taken from -/// -/// `moving min max` provides one data structure for keeping track of the -/// minimum value and one for keeping track of the maximum value in a sliding -/// window. -/// -/// Each element is stored with the current min/max. One stack to push and another one for pop. If pop stack is empty, -/// push to this stack all elements popped from first stack while updating their current min/max. Now pop from -/// the second stack (MovingMin/Max struct works as a queue). To find the minimum element of the queue, -/// look at the smallest/largest two elements of the individual stacks, then take the minimum of those two values. -/// -/// The complexity of the operations are -/// - O(1) for getting the minimum/maximum -/// - O(1) for push -/// - amortized O(1) for pop -/// -/// ``` -/// # use datafusion_functions_aggregate::min_max::MovingMin; -/// let mut moving_min = MovingMin::::new(); -/// moving_min.push(2); -/// moving_min.push(1); -/// moving_min.push(3); -/// -/// assert_eq!(moving_min.min(), Some(&1)); -/// assert_eq!(moving_min.pop(), Some(2)); +/// `MovingMin` keeps track of the minimum value in a sliding window using a +/// monotonic deque. Each element is stored with its sequence number, and the +/// deque maintains candidate elements in ascending value order. /// -/// assert_eq!(moving_min.min(), Some(&1)); -/// assert_eq!(moving_min.pop(), Some(1)); -/// -/// assert_eq!(moving_min.min(), Some(&3)); -/// assert_eq!(moving_min.pop(), Some(3)); -/// -/// assert_eq!(moving_min.min(), None); -/// assert_eq!(moving_min.pop(), None); -/// ``` +/// Complexity: +/// - O(1) for getting the minimum +/// - amortized O(1) for push +/// - O(1) for pop #[derive(Debug)] -pub struct MovingMin { - push_stack: Vec<(T, T)>, - pop_stack: Vec<(T, T)>, +pub(crate) struct MovingMin { + deque: VecDeque<(u64, T)>, + push_seq: u64, + pop_seq: u64, } -impl Default for MovingMin { +impl Default for MovingMin { fn default() -> Self { Self { - push_stack: Vec::new(), - pop_stack: Vec::new(), + deque: VecDeque::new(), + push_seq: 0, + pop_seq: 0, } } } -impl MovingMin { - /// Creates a new `MovingMin` to keep track of the minimum in a sliding - /// window. +impl MovingMin { + /// Creates a new `MovingMin` to keep track of the minimum in a sliding window. #[inline] pub fn new() -> Self { Self::default() } - /// Creates a new `MovingMin` to keep track of the minimum in a sliding - /// window with `capacity` allocated slots. + /// Creates a new `MovingMin` to keep track of the minimum in a sliding window with + /// `capacity` allocated slots. + #[cfg(test)] #[inline] pub fn with_capacity(capacity: usize) -> Self { Self { - push_stack: Vec::with_capacity(capacity), - pop_stack: Vec::with_capacity(capacity), + deque: VecDeque::with_capacity(capacity), + push_seq: 0, + pop_seq: 0, } } @@ -799,105 +793,113 @@ impl MovingMin { /// empty. #[inline] pub fn min(&self) -> Option<&T> { - match (self.push_stack.last(), self.pop_stack.last()) { - (None, None) => None, - (Some((_, min)), None) => Some(min), - (None, Some((_, min))) => Some(min), - (Some((_, a)), Some((_, b))) => Some(if a < b { a } else { b }), - } + self.deque.front().map(|(_, val)| val) + } + + #[inline] + fn check_invariants(&self) { + debug_assert!(self.pop_seq <= self.push_seq); + debug_assert!( + self.deque + .front() + .is_none_or(|&(front_seq, _)| front_seq >= self.pop_seq) + ); } /// Pushes a new element into the sliding window. #[inline] pub fn push(&mut self, val: T) { - self.push_stack.push(match self.push_stack.last() { - Some((_, min)) => { - if val > *min { - (val, min.clone()) - } else { - (val.clone(), val) - } - } - None => (val.clone(), val), - }); + let seq = self.push_seq; + self.push_seq += 1; + while self.deque.back().is_some_and(|back_val| back_val.1 >= val) { + self.deque.pop_back(); + } + self.deque.push_back((seq, val)); + + self.check_invariants(); } - /// Removes and returns the last value of the sliding window. + /// Removes the oldest value from the sliding window. + /// + /// If the window is empty, this is a no-op. #[inline] - pub fn pop(&mut self) -> Option { - if self.pop_stack.is_empty() { - match self.push_stack.pop() { - Some((val, _)) => { - let mut last = (val.clone(), val); - self.pop_stack.push(last.clone()); - while let Some((val, _)) = self.push_stack.pop() { - let min = if last.1 < val { - last.1.clone() - } else { - val.clone() - }; - last = (val.clone(), min); - self.pop_stack.push(last.clone()); - } - } - None => return None, - } + pub fn pop(&mut self) { + if self.is_empty() { + return; + } + let seq = self.pop_seq; + self.pop_seq += 1; + if self + .deque + .front() + .is_some_and(|front_val| front_val.0 == seq) + { + self.deque.pop_front(); } - self.pop_stack.pop().map(|(val, _)| val) + + self.check_invariants(); } /// Returns the number of elements stored in the sliding window. - #[inline] + #[cfg(test)] pub fn len(&self) -> usize { - self.push_stack.len() + self.pop_stack.len() + (self.push_seq - self.pop_seq) as usize } /// Returns `true` if the moving window contains no elements. #[inline] pub fn is_empty(&self) -> bool { - self.len() == 0 + self.push_seq == self.pop_seq + } + + /// Heap bytes owned by the deque plus each stored `T`'s + /// heap payload as reported by `elem_heap`. Excludes `size_of::()`. + #[inline] + fn heap_size(&self, elem_heap: impl Fn(&T) -> usize) -> usize { + moving_deque_heap_size(&self.deque, elem_heap) } } +/// Shared implementation for [`MovingMin::heap_size`] and +/// [`MovingMax::heap_size`]. Both share the same deque layout. +#[inline] +fn moving_deque_heap_size( + deque: &VecDeque<(u64, T)>, + elem_heap: impl Fn(&T) -> usize, +) -> usize { + let buffers = deque.capacity() * size_of::<(u64, T)>(); + let elems: usize = deque.iter().map(|(_, val)| elem_heap(val)).sum(); + buffers + elems +} + /// Keep track of the maximum value in a sliding window. /// -/// See [`MovingMin`] for more details. -/// -/// ``` -/// # use datafusion_functions_aggregate::min_max::MovingMax; -/// let mut moving_max = MovingMax::::new(); -/// moving_max.push(2); -/// moving_max.push(3); -/// moving_max.push(1); -/// -/// assert_eq!(moving_max.max(), Some(&3)); -/// assert_eq!(moving_max.pop(), Some(2)); -/// -/// assert_eq!(moving_max.max(), Some(&3)); -/// assert_eq!(moving_max.pop(), Some(3)); -/// -/// assert_eq!(moving_max.max(), Some(&1)); -/// assert_eq!(moving_max.pop(), Some(1)); +/// `MovingMax` keeps track of the maximum value in a sliding window using a +/// monotonic deque. Each element is stored with its sequence number, and the +/// deque maintains candidate elements in descending value order. /// -/// assert_eq!(moving_max.max(), None); -/// assert_eq!(moving_max.pop(), None); -/// ``` +/// Complexity: +/// - O(1) for getting the maximum +/// - amortized O(1) for push +/// - O(1) for pop #[derive(Debug)] -pub struct MovingMax { - push_stack: Vec<(T, T)>, - pop_stack: Vec<(T, T)>, +pub(crate) struct MovingMax { + deque: VecDeque<(u64, T)>, + push_seq: u64, + pop_seq: u64, } -impl Default for MovingMax { +impl Default for MovingMax { fn default() -> Self { Self { - push_stack: Vec::new(), - pop_stack: Vec::new(), + deque: VecDeque::new(), + push_seq: 0, + pop_seq: 0, } } } -impl MovingMax { +impl MovingMax { /// Creates a new `MovingMax` to keep track of the maximum in a sliding window. #[inline] pub fn new() -> Self { @@ -906,74 +908,83 @@ impl MovingMax { /// Creates a new `MovingMax` to keep track of the maximum in a sliding window with /// `capacity` allocated slots. + #[cfg(test)] #[inline] pub fn with_capacity(capacity: usize) -> Self { Self { - push_stack: Vec::with_capacity(capacity), - pop_stack: Vec::with_capacity(capacity), + deque: VecDeque::with_capacity(capacity), + push_seq: 0, + pop_seq: 0, } } /// Returns the maximum of the sliding window or `None` if the window is empty. #[inline] pub fn max(&self) -> Option<&T> { - match (self.push_stack.last(), self.pop_stack.last()) { - (None, None) => None, - (Some((_, max)), None) => Some(max), - (None, Some((_, max))) => Some(max), - (Some((_, a)), Some((_, b))) => Some(if a > b { a } else { b }), - } + self.deque.front().map(|(_, val)| val) + } + + #[inline] + fn check_invariants(&self) { + debug_assert!(self.pop_seq <= self.push_seq); + debug_assert!( + self.deque + .front() + .is_none_or(|&(front_seq, _)| front_seq >= self.pop_seq) + ); } /// Pushes a new element into the sliding window. #[inline] pub fn push(&mut self, val: T) { - self.push_stack.push(match self.push_stack.last() { - Some((_, max)) => { - if val < *max { - (val, max.clone()) - } else { - (val.clone(), val) - } - } - None => (val.clone(), val), - }); + let seq = self.push_seq; + self.push_seq += 1; + while self.deque.back().is_some_and(|back_val| back_val.1 <= val) { + self.deque.pop_back(); + } + self.deque.push_back((seq, val)); + + self.check_invariants(); } - /// Removes and returns the last value of the sliding window. + /// Removes the oldest value from the sliding window. + /// + /// If the window is empty, this is a no-op. #[inline] - pub fn pop(&mut self) -> Option { - if self.pop_stack.is_empty() { - match self.push_stack.pop() { - Some((val, _)) => { - let mut last = (val.clone(), val); - self.pop_stack.push(last.clone()); - while let Some((val, _)) = self.push_stack.pop() { - let max = if last.1 > val { - last.1.clone() - } else { - val.clone() - }; - last = (val.clone(), max); - self.pop_stack.push(last.clone()); - } - } - None => return None, - } + pub fn pop(&mut self) { + if self.is_empty() { + return; + } + let seq = self.pop_seq; + self.pop_seq += 1; + if self + .deque + .front() + .is_some_and(|front_val| front_val.0 == seq) + { + self.deque.pop_front(); } - self.pop_stack.pop().map(|(val, _)| val) + + self.check_invariants(); } /// Returns the number of elements stored in the sliding window. - #[inline] + #[cfg(test)] pub fn len(&self) -> usize { - self.push_stack.len() + self.pop_stack.len() + (self.push_seq - self.pop_seq) as usize } /// Returns `true` if the moving window contains no elements. #[inline] pub fn is_empty(&self) -> bool { - self.len() == 0 + self.push_seq == self.pop_seq + } + + /// Heap bytes owned by the deque plus each stored `T`'s + /// heap payload as reported by `elem_heap`. Excludes `size_of::()`. + #[inline] + fn heap_size(&self, elem_heap: impl Fn(&T) -> usize) -> usize { + moving_deque_heap_size(&self.deque, elem_heap) } } @@ -1195,6 +1206,58 @@ mod tests { Ok(()) } + #[test] + fn sliding_min_all_null_window() -> Result<()> { + let mut min_acc = SlidingMinAccumulator::try_new(&DataType::Int32)?; + + let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), None])); + min_acc.update_batch(&[Arc::clone(&values)])?; + assert_eq!(min_acc.evaluate()?, ScalarValue::Int32(Some(3))); + + // Retract `3`; the window now contains only the NULL + let retracted: ArrayRef = Arc::new(Int32Array::from(vec![Some(3)])); + min_acc.retract_batch(&[Arc::clone(&retracted)])?; + assert_eq!(min_acc.evaluate()?, ScalarValue::Int32(None)); + + // A subsequent non-null value must be picked up again + let update: ArrayRef = Arc::new(Int32Array::from(vec![Some(7)])); + min_acc.update_batch(&[Arc::clone(&update)])?; + assert_eq!(min_acc.evaluate()?, ScalarValue::Int32(Some(7))); + + // Retracting the NULL row must not pop the remaining value + let null_row: ArrayRef = Arc::new(Int32Array::from(vec![None::])); + min_acc.retract_batch(&[Arc::clone(&null_row)])?; + assert_eq!(min_acc.evaluate()?, ScalarValue::Int32(Some(7))); + + Ok(()) + } + + #[test] + fn sliding_max_all_null_window() -> Result<()> { + let mut max_acc = SlidingMaxAccumulator::try_new(&DataType::Int32)?; + + let values: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), None])); + max_acc.update_batch(&[Arc::clone(&values)])?; + assert_eq!(max_acc.evaluate()?, ScalarValue::Int32(Some(3))); + + // Retract `3`; the window now contains only the NULL + let retracted: ArrayRef = Arc::new(Int32Array::from(vec![Some(3)])); + max_acc.retract_batch(&[Arc::clone(&retracted)])?; + assert_eq!(max_acc.evaluate()?, ScalarValue::Int32(None)); + + // A subsequent non-null value must be picked up again + let update: ArrayRef = Arc::new(Int32Array::from(vec![Some(7)])); + max_acc.update_batch(&[Arc::clone(&update)])?; + assert_eq!(max_acc.evaluate()?, ScalarValue::Int32(Some(7))); + + // Retracting the NULL row must not disturb the remaining value + let null_row: ArrayRef = Arc::new(Int32Array::from(vec![None::])); + max_acc.retract_batch(&[Arc::clone(&null_row)])?; + assert_eq!(max_acc.evaluate()?, ScalarValue::Int32(Some(7))); + + Ok(()) + } + #[test] fn moving_min_tests() -> Result<()> { moving_min_i32(100, 10)?; @@ -1213,6 +1276,95 @@ mod tests { Ok(()) } + #[test] + fn moving_min_max_heap_size_i32() { + // Fixed-width `T` has no per-element heap payload, so `heap_size` + // reports exactly the buffer's capacity in bytes. + let mut moving_min = MovingMin::::with_capacity(4); + let mut moving_max = MovingMax::::with_capacity(4); + let elem = |_: &i32| 0; + + let buffer_only = moving_min.deque.capacity() * size_of::<(u64, i32)>(); + assert_eq!(moving_min.heap_size(elem), buffer_only); + assert_eq!(moving_max.heap_size(elem), buffer_only); + + for i in 0..3 { + moving_min.push(i); + moving_max.push(i); + } + // Elements sit inside the pre-allocated buffers, so still buffer-only. + assert_eq!(moving_min.heap_size(elem), buffer_only); + assert_eq!(moving_max.heap_size(elem), buffer_only); + } + + #[test] + fn moving_min_max_heap_size_counts_elems() { + let mut moving_min = MovingMin::::with_capacity(2); + let mut moving_max = MovingMax::::with_capacity(2); + let elem = |s: &String| s.capacity(); + + moving_min.push("abcdef".to_string()); + moving_max.push("abcdef".to_string()); + + let buffers = moving_min.deque.capacity() * size_of::<(u64, String)>(); + let elems = 6; + assert_eq!(moving_min.heap_size(elem), buffers + elems); + assert_eq!(moving_max.heap_size(elem), buffers + elems); + } + + #[test] + fn test_moving_min_max_empty_pop() { + let mut moving_min = MovingMin::::new(); + moving_min.pop(); // empty pop is a no-op + assert_eq!(moving_min.len(), 0); + assert!(moving_min.is_empty()); + // Verify it still works correctly after empty pop + moving_min.push(10); + moving_min.push(20); + assert_eq!(moving_min.min(), Some(&10)); + moving_min.pop(); + assert_eq!(moving_min.min(), Some(&20)); + + let mut moving_max = MovingMax::::new(); + moving_max.pop(); // empty pop is a no-op + assert_eq!(moving_max.len(), 0); + assert!(moving_max.is_empty()); + // Verify it still works correctly after empty pop + moving_max.push(20); + moving_max.push(10); + assert_eq!(moving_max.max(), Some(&20)); + moving_max.pop(); + assert_eq!(moving_max.max(), Some(&10)); + } + + #[test] + fn test_moving_min_max_duplicate_heavy() { + let mut moving_min = MovingMin::::new(); + let mut moving_max = MovingMax::::new(); + + // Push duplicates + for _ in 0..5 { + moving_min.push(5); + moving_max.push(5); + } + + assert_eq!(moving_min.len(), 5); + assert_eq!(moving_max.len(), 5); + + // Ensure min/max query works and we can pop all duplicates correctly + for i in (1..=5).rev() { + assert_eq!(moving_min.len(), i); + assert_eq!(moving_max.len(), i); + assert_eq!(moving_min.min(), Some(&5)); + assert_eq!(moving_max.max(), Some(&5)); + moving_min.pop(); + moving_max.pop(); + } + + assert!(moving_min.is_empty()); + assert!(moving_max.is_empty()); + } + #[test] fn test_min_max_coerce_types() { // the coerced types is same with input types diff --git a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs index e4ac7eccf5692..efeaea314c4f5 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_bytes.rs @@ -27,6 +27,8 @@ use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls: use std::mem::size_of; use std::sync::Arc; +use datafusion_common::utils::split_vec_min_alloc; + /// Implements fast Min/Max [`GroupsAccumulator`] for "bytes" types ([`StringArray`], /// [`BinaryArray`], [`StringViewArray`], etc) /// @@ -307,11 +309,10 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // min/max are their own states (no transition needed) - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } fn convert_to_state( @@ -324,11 +325,6 @@ impl GroupsAccumulator for MinMaxBytesAccumulator { let output = apply_filter_as_nulls(&values[0], opt_filter)?; Ok(vec![output]) } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn size(&self) -> usize { self.inner.size() } @@ -493,7 +489,7 @@ impl MinMaxBytesState { ) } EmitTo::First(n) => { - let first_min_maxes: Vec<_> = self.min_max.drain(..n).collect(); + let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n); let first_data_capacity: usize = first_min_maxes .iter() .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) diff --git a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs index 796fd586ca5c8..d1bac4e2f90db 100644 --- a/datafusion/functions-aggregate/src/min_max/min_max_struct.rs +++ b/datafusion/functions-aggregate/src/min_max/min_max_struct.rs @@ -30,6 +30,8 @@ use datafusion_common::{ use datafusion_expr::{EmitTo, GroupsAccumulator}; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::apply_filter_as_nulls; +use datafusion_common::utils::split_vec_min_alloc; + /// Accumulator for MIN/MAX operations on Struct data types. /// /// This accumulator tracks the minimum or maximum struct value encountered @@ -116,7 +118,7 @@ impl GroupsAccumulator for MinMaxStructAccumulator { let mut copy = MutableArrayData::new(min_maxes_refs, true, min_maxes_data.len()); for (i, item) in min_maxes_data.iter().enumerate() { - copy.extend(i, 0, item.len()); + copy.try_extend(i, 0, item.len())?; } let result = copy.freeze(); assert_eq!(&self.inner.data_type, result.data_type()); @@ -132,11 +134,10 @@ impl GroupsAccumulator for MinMaxStructAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // min/max are their own states (no transition needed) - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } fn convert_to_state( @@ -149,11 +150,6 @@ impl GroupsAccumulator for MinMaxStructAccumulator { let output = apply_filter_as_nulls(&values[0], opt_filter)?; Ok(vec![output]) } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn size(&self) -> usize { self.inner.size() } @@ -282,7 +278,7 @@ impl MinMaxStructState { ) } EmitTo::First(n) => { - let first_min_maxes: Vec<_> = self.min_max.drain(..n).collect(); + let first_min_maxes = split_vec_min_alloc(&mut self.min_max, n); let first_data_capacity: usize = first_min_maxes .iter() .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) diff --git a/datafusion/functions-aggregate/src/percentile_cont.rs b/datafusion/functions-aggregate/src/percentile_cont.rs index 256388c216f00..3a98900bbb446 100644 --- a/datafusion/functions-aggregate/src/percentile_cont.rs +++ b/datafusion/functions-aggregate/src/percentile_cont.rs @@ -32,13 +32,16 @@ use arrow::{ use num_traits::AsPrimitive; use arrow::array::ArrowNativeTypeOp; +use datafusion_common::hash_utils::RandomState; use datafusion_common::internal_err; use datafusion_common::types::{NativeType, logical_float64}; +use datafusion_common::utils::memory::estimate_memory_size; use datafusion_functions_aggregate_common::noop_accumulator::NoopAccumulator; use crate::min_max::{max_udaf, min_udaf}; use datafusion_common::{ - Result, ScalarValue, internal_datafusion_err, utils::take_function_args, + Result, ScalarValue, exec_datafusion_err, internal_datafusion_err, + utils::{SingleRowListArrayBuilder, take_function_args}, }; use datafusion_expr::utils::format_state_name; use datafusion_expr::{ @@ -53,7 +56,7 @@ use datafusion_expr::{ }; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::accumulate::accumulate; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::nulls::filtered_null_mask; -use datafusion_functions_aggregate_common::utils::{GenericDistinctBuffer, Hashable}; +use datafusion_functions_aggregate_common::utils::Hashable; use datafusion_macros::user_doc; use crate::utils::validate_percentile_expr; @@ -420,8 +423,18 @@ where fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let values = values[0].as_primitive::(); - self.all_values.reserve(values.len() - values.null_count()); - self.all_values.extend(values.iter().flatten()); + let additional = values.len() - values.null_count(); + self.all_values.try_reserve(additional).map_err(|e| { + exec_datafusion_err!( + "failed to reserve {additional} values for percentile_cont accumulator: {e}" + ) + })?; + if values.null_count() > 0 { + self.all_values.extend(values.iter().flatten()); + } else { + // Fast path: no nulls, so the values buffer can be appended wholesale. + self.all_values.extend_from_slice(values.values()); + } Ok(()) } @@ -441,11 +454,19 @@ where } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - let mut to_remove: HashMap, usize> = HashMap::new(); + let mut to_remove: HashMap, usize, RandomState> = + HashMap::default(); let arr = values[0].as_primitive::(); - for value in arr.iter().flatten() { - *to_remove.entry(Hashable(value)).or_default() += 1; + if arr.null_count() > 0 { + for value in arr.iter().flatten() { + *to_remove.entry(Hashable(value)).or_default() += 1; + } + } else { + // Fast path: no nulls, so skip the per-element validity check. + for value in arr.values().iter() { + *to_remove.entry(Hashable(*value)).or_default() += 1; + } } let mut i = 0; @@ -466,6 +487,15 @@ where i += 1; } } + + // Retracting values that are not tracked means the accumulator state + // has diverged from the window frame; continuing would silently + // produce wrong results, so surface it as an error. + if !to_remove.is_empty() { + return internal_err!( + "percentile_cont retract_batch: retracted value(s) not present in the window" + ); + } Ok(()) } @@ -531,8 +561,6 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - // Since aggregate filter should be applied in partial stage, in final stage there should be no filter - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 1, "one argument to merge_batch"); @@ -648,11 +676,6 @@ where Ok(vec![Arc::new(converted_list_array)]) } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn size(&self) -> usize { self.group_values .iter() @@ -663,16 +686,28 @@ where } } +/// Sliding-window–capable accumulator for `percentile_cont(DISTINCT ...)`. +/// +/// Distinct values are tracked with a per-value multiplicity count (how many +/// rows currently in the window carry that value) rather than a plain set, so +/// that `retract_batch` only drops a value once *all* of its occurrences have +/// left the window frame. The percentile is then computed over the set of keys +/// with a positive count. #[derive(Debug)] struct DistinctPercentileContAccumulator { - distinct_values: GenericDistinctBuffer, + /// Distinct value -> number of in-window rows carrying it. + /// + /// Uses the same fast (foldhash) `RandomState` as the shared + /// `GenericDistinctBuffer` rather than the standard library's default + /// SipHash, which is considerably slower for this hot path. + counts: HashMap, usize, RandomState>, percentile: f64, } impl DistinctPercentileContAccumulator { fn new(percentile: f64) -> Self { Self { - distinct_values: GenericDistinctBuffer::new(T::DATA_TYPE), + counts: HashMap::default(), percentile, } } @@ -685,26 +720,59 @@ where f64: AsPrimitive, { fn state(&mut self) -> Result> { - self.distinct_values.state() + // Emit the distinct keys as a single List scalar, matching the state + // shape declared in `state_fields` (a List of the input type). Counts + // are window-local bookkeeping and are intentionally not serialized: + // cross-partition merges only need the distinct key set. + let arr = Arc::new( + PrimitiveArray::::from_iter_values(self.counts.keys().map(|v| v.0)) + .with_data_type(T::DATA_TYPE), + ); + Ok(vec![ + SingleRowListArrayBuilder::new(arr).build_list_scalar(), + ]) } fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { - self.distinct_values.update_batch(values) + // `values` may carry extra argument columns (e.g. the percentile + // literal); only the first column holds the aggregated values. + let arr = values[0].as_primitive::(); + if arr.null_count() > 0 { + for value in arr.iter().flatten() { + *self.counts.entry(Hashable(value)).or_default() += 1; + } + } else { + // Fast path: no nulls, so skip the per-element validity check. + for value in arr.values().iter() { + *self.counts.entry(Hashable(*value)).or_default() += 1; + } + } + Ok(()) } fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { - self.distinct_values.merge_batch(states) + let list = states[0].as_list::(); + for values in list.iter().flatten() { + let arr = values.as_primitive::(); + for value in arr.iter().flatten() { + *self.counts.entry(Hashable(value)).or_default() += 1; + } + } + Ok(()) } fn evaluate(&mut self) -> Result { - let mut values: Vec = - self.distinct_values.values.iter().map(|v| v.0).collect(); + let mut values: Vec = self.counts.keys().map(|v| v.0).collect(); let value = calculate_percentile::(&mut values, self.percentile); ScalarValue::new_primitive::(value, &T::DATA_TYPE) } fn size(&self) -> usize { - size_of_val(self) + self.distinct_values.size() + estimate_memory_size::<(Hashable, usize)>( + self.counts.capacity(), + size_of_val(self), + ) + .unwrap() } fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { @@ -713,8 +781,32 @@ where } let arr = values[0].as_primitive::(); - for value in arr.iter().flatten() { - self.distinct_values.values.remove(&Hashable(value)); + let mut decrement = |value: T::Native| { + match self.counts.get_mut(&Hashable(value)) { + Some(count) => { + *count -= 1; + if *count == 0 { + self.counts.remove(&Hashable(value)); + } + Ok(()) + } + // Retracting a value that isn't tracked means the accumulator + // state has diverged from the window frame; continuing would + // silently produce wrong results, so surface it as an error. + None => internal_err!( + "percentile_cont(DISTINCT) retract_batch: retracted a value not present in the window" + ), + } + }; + if arr.null_count() > 0 { + for value in arr.iter().flatten() { + decrement(value)?; + } + } else { + // Fast path: no nulls, so skip the per-element validity check. + for value in arr.values().iter() { + decrement(*value)?; + } } Ok(()) } @@ -806,18 +898,60 @@ where #[cfg(test)] mod tests { - use super::calculate_percentile; + use super::*; + use arrow::array::Float64Array; use half::f16; + #[test] + fn retract_batch_errors_on_untracked_value() { + let mut acc = PercentileContAccumulator::::new(0.5); + let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); + acc.update_batch(std::slice::from_ref(&values)).unwrap(); + + let retract: ArrayRef = Arc::new(Float64Array::from(vec![3.0])); + let err = acc + .retract_batch(std::slice::from_ref(&retract)) + .unwrap_err() + .to_string(); + assert!( + err.contains("not present in the window"), + "unexpected error: {err}" + ); + } + + #[test] + fn update_batch_with_and_without_nulls_agree() { + // The null-free fast path must accumulate the same values as the + // general path. + let dense: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])); + let sparse: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.0), + None, + Some(2.0), + None, + Some(3.0), + ])); + + let mut dense_acc = PercentileContAccumulator::::new(0.5); + dense_acc + .update_batch(std::slice::from_ref(&dense)) + .unwrap(); + let mut sparse_acc = PercentileContAccumulator::::new(0.5); + sparse_acc + .update_batch(std::slice::from_ref(&sparse)) + .unwrap(); + + assert_eq!(dense_acc.all_values, sparse_acc.all_values); + } + #[test] fn f16_interpolation_does_not_overflow_to_nan() { // Regression test for https://github.com/apache/datafusion/issues/18945 // Interpolating between 0 and the max finite f16 value previously overflowed // intermediate f16 computations and produced NaN. let mut values = vec![f16::from_f32(0.0), f16::from_f32(65504.0)]; - let result = - calculate_percentile::(&mut values, 0.5) - .expect("non-empty input"); + let result = calculate_percentile::(&mut values, 0.5) + .expect("non-empty input"); let result_f = result.to_f32(); assert!( !result_f.is_nan(), diff --git a/datafusion/functions-aggregate/src/regr.rs b/datafusion/functions-aggregate/src/regr.rs index 3a68672abb949..3d5bbf1eda24e 100644 --- a/datafusion/functions-aggregate/src/regr.rs +++ b/datafusion/functions-aggregate/src/regr.rs @@ -457,6 +457,18 @@ impl AggregateUDFImpl for Regr { } } + fn default_value(&self, _data_type: &DataType) -> Result { + if self.regr_type == RegrType::Count { + Ok(ScalarValue::UInt64(Some(0))) + } else { + Ok(ScalarValue::Float64(None)) + } + } + + fn is_nullable(&self) -> bool { + self.regr_type != RegrType::Count + } + fn accumulator(&self, _acc_args: AccumulatorArgs) -> Result> { Ok(Box::new(RegrAccumulator::try_new(&self.regr_type)?)) } diff --git a/datafusion/functions-aggregate/src/stddev.rs b/datafusion/functions-aggregate/src/stddev.rs index 68e38a3b8db07..15511bf4a565f 100644 --- a/datafusion/functions-aggregate/src/stddev.rs +++ b/datafusion/functions-aggregate/src/stddev.rs @@ -22,7 +22,7 @@ use std::hash::Hash; use std::mem::align_of_val; use std::sync::Arc; -use arrow::array::Float64Array; +use arrow::array::{BooleanArray, Float64Array}; use arrow::datatypes::FieldRef; use arrow::{array::ArrayRef, datatypes::DataType, datatypes::Field}; use datafusion_common::ScalarValue; @@ -318,7 +318,7 @@ impl GroupsAccumulator for StddevGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&arrow::array::BooleanArray>, + opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { self.variance @@ -329,11 +329,10 @@ impl GroupsAccumulator for StddevGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&arrow::array::BooleanArray>, total_num_groups: usize, ) -> Result<()> { self.variance - .merge_batch(values, group_indices, opt_filter, total_num_groups) + .merge_batch(values, group_indices, total_num_groups) } fn evaluate(&mut self, emit_to: datafusion_expr::EmitTo) -> Result { @@ -346,6 +345,13 @@ impl GroupsAccumulator for StddevGroupsAccumulator { self.variance.state(emit_to) } + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + self.variance.convert_to_state(values, opt_filter) + } fn size(&self) -> usize { self.variance.size() } diff --git a/datafusion/functions-aggregate/src/string_agg.rs b/datafusion/functions-aggregate/src/string_agg.rs index f0757818afb93..3fe2b0a186ae3 100644 --- a/datafusion/functions-aggregate/src/string_agg.rs +++ b/datafusion/functions-aggregate/src/string_agg.rs @@ -413,11 +413,10 @@ impl GroupsAccumulator for StringAggGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { // State is always LargeUtf8, which update_batch already handles. - self.update_batch(values, group_indices, opt_filter, total_num_groups) + self.update_batch(values, group_indices, None, total_num_groups) } fn convert_to_state( @@ -433,11 +432,6 @@ impl GroupsAccumulator for StringAggGroupsAccumulator { }; Ok(vec![result]) } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn size(&self) -> usize { self.total_data_bytes + self.values.capacity() * size_of::>() @@ -898,7 +892,7 @@ mod tests { // Simulate a second accumulator's state (LargeUtf8 partial strings) let partial_state: ArrayRef = Arc::new(LargeStringArray::from(vec!["c,d", "e"])); - acc.merge_batch(&[partial_state], &[0, 1], None, 2)?; + acc.merge_batch(&[partial_state], &[0, 1], 2)?; let result = evaluate_groups(&mut acc, EmitTo::All); assert_eq!( diff --git a/datafusion/functions-aggregate/src/sum.rs b/datafusion/functions-aggregate/src/sum.rs index 81efea1df22b1..71932c5f0b3f7 100644 --- a/datafusion/functions-aggregate/src/sum.rs +++ b/datafusion/functions-aggregate/src/sum.rs @@ -24,10 +24,12 @@ use arrow::datatypes::{ DECIMAL128_MAX_PRECISION, DECIMAL256_MAX_PRECISION, DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, DurationSecondType, FieldRef, - Float64Type, Int64Type, TimeUnit, UInt64Type, + Float64Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, IntervalUnit, + IntervalYearMonthType, TimeUnit, UInt64Type, }; use datafusion_common::hash_utils::RandomState; use datafusion_common::internal_err; +use datafusion_common::stats::Precision; use datafusion_common::types::{ NativeType, logical_float64, logical_int8, logical_int16, logical_int32, logical_int64, logical_uint8, logical_uint16, logical_uint32, logical_uint64, @@ -39,13 +41,14 @@ use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::{AggregateOrderSensitivity, format_state_name}; use datafusion_expr::{ Accumulator, AggregateUDFImpl, Coercion, Documentation, Expr, GroupsAccumulator, - Operator, ReversedUDAF, SetMonotonicity, Signature, TypeSignature, + Operator, ReversedUDAF, SetMonotonicity, Signature, StatisticsArgs, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_functions_aggregate_common::aggregate::groups_accumulator::prim_op::PrimitiveGroupsAccumulator; use datafusion_functions_aggregate_common::aggregate::sum_distinct::DistinctSumAccumulator; use datafusion_macros::user_doc; -use std::mem::size_of_val; +use datafusion_physical_expr::expressions::{CastExpr, Column}; +use std::mem::{size_of, size_of_val}; make_udaf_expr_and_func!( Sum, @@ -117,6 +120,21 @@ macro_rules! downcast_sum { $args.return_field.data_type().clone() ) } + DataType::Interval(IntervalUnit::YearMonth) => { + $helper!( + IntervalYearMonthType, + $args.return_field.data_type().clone() + ) + } + DataType::Interval(IntervalUnit::DayTime) => { + $helper!(IntervalDayTimeType, $args.return_field.data_type().clone()) + } + DataType::Interval(IntervalUnit::MonthDayNano) => { + $helper!( + IntervalMonthDayNanoType, + $args.return_field.data_type().clone() + ) + } _ => { not_impl_err!( "Sum not supported for {}: {}", @@ -186,6 +204,9 @@ impl Sum { TypeSignature::Coercible(vec![Coercion::new_exact( TypeSignatureClass::Duration, )]), + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Interval, + )]), ], Volatility::Immutable, ), @@ -232,6 +253,7 @@ impl AggregateUDFImpl for Sum { Ok(DataType::Decimal256(new_precision, *scale)) } DataType::Duration(time_unit) => Ok(DataType::Duration(*time_unit)), + DataType::Interval(interval_unit) => Ok(DataType::Interval(*interval_unit)), other => { exec_err!("[return_type] SUM not supported for {}", other) } @@ -303,13 +325,18 @@ impl AggregateUDFImpl for Sum { args: AccumulatorArgs, ) -> Result> { if args.is_distinct { - // distinct path: use our sliding‐window distinct‐sum - macro_rules! helper_distinct { - ($t:ty, $dt:expr) => { - Ok(Box::new(SlidingDistinctSumAccumulator::try_new(&$dt)?)) - }; + // distinct path: [`SlidingDistinctSumAccumulator`] only implements + // Int64, so gate the supported type here rather than dispatching + // through `downcast_sum!`, which accepts every SUM type + match args.return_field.data_type() { + DataType::Int64 => Ok(Box::new(SlidingDistinctSumAccumulator::try_new( + &DataType::Int64, + )?)), + _ => not_impl_err!( + "SUM(DISTINCT) over sliding window frames is only supported for Int64, got {}", + args.expr_fields[0].data_type() + ), } - downcast_sum!(args, helper_distinct) } else { // non‐distinct path: existing sliding sum macro_rules! helper { @@ -385,6 +412,58 @@ impl AggregateUDFImpl for Sum { // SUM(arg) + lit * COUNT(arg) Ok(Some(sum_agg + (lit.clone() * count_agg))) } + + fn value_from_stats(&self, statistics_args: &StatisticsArgs) -> Option { + if statistics_args.is_distinct { + return None; + } + + let [expr] = statistics_args.exprs else { + return None; + }; + + let (col_expr, cast_type) = match expr.downcast_ref::() { + Some(col_expr) => (col_expr, None), + None => { + let cast_expr = expr.downcast_ref::()?; + let col_expr = cast_expr.expr().downcast_ref::()?; + (col_expr, Some(cast_expr.cast_type())) + } + }; + + let col_stats = statistics_args + .statistics + .column_statistics + .get(col_expr.index())?; + + // Replacing SUM with a literal is only valid for exact statistics. + // `cast_to_sum_type` also widens small integer stats to the SQL SUM + // return type, e.g. Int32 statistics become an Int64 SUM value. + let Precision::Exact(val) = col_stats.sum_value.cast_to_sum_type() else { + return None; + }; + if val.is_null() { + return None; + } + + // SUM coercion can introduce a physical CAST around the input column + // (`SUM(Int32)` becomes `SUM(CAST(Int32 AS Int64))`). Only use the + // column's raw sum stats when the widened stats value matches that + // cast target and the aggregate return type. + if let Some(cast_type) = cast_type { + let value_type = val.data_type(); + if cast_type != statistics_args.return_type || &value_type != cast_type { + return None; + } + return Some(val); + } + + if &val.data_type() == statistics_args.return_type { + Some(val) + } else { + val.cast_to(statistics_args.return_type).ok() + } + } } /// This accumulator computes SUM incrementally @@ -525,7 +604,9 @@ impl SlidingDistinctSumAccumulator { pub fn try_new(data_type: &DataType) -> Result { // TODO support other numeric types if *data_type != DataType::Int64 { - return exec_err!("SlidingDistinctSumAccumulator only supports Int64"); + return exec_err!( + "SlidingDistinctSumAccumulator only supports Int64, got {data_type}" + ); } Ok(Self { counts: HashMap::default(), @@ -533,29 +614,65 @@ impl SlidingDistinctSumAccumulator { data_type: data_type.clone(), }) } + + fn update_value(&mut self, value: i64) { + let cnt = self.counts.entry(value).or_insert(0); + if *cnt == 0 { + // first occurrence in window + self.sum = self.sum.wrapping_add(value); + } + *cnt += 1; + } + + fn retract_value(&mut self, value: i64) { + if let Some(cnt) = self.counts.get_mut(&value) { + *cnt -= 1; + if *cnt == 0 { + // last copy leaving window + self.sum = self.sum.wrapping_sub(value); + self.counts.remove(&value); + } + } + } + + fn apply_valid_values( + &mut self, + arr: &arrow::array::PrimitiveArray, + mut op: F, + ) where + F: FnMut(&mut Self, i64), + { + if arr.null_count() == 0 { + for &value in arr.values() { + op(self, value); + } + } else { + for (idx, &value) in arr.values().iter().enumerate() { + if arr.is_valid(idx) { + op(self, value); + } + } + } + } } impl Accumulator for SlidingDistinctSumAccumulator { fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let arr = values[0].as_primitive::(); - for &v in arr.values() { - let cnt = self.counts.entry(v).or_insert(0); - if *cnt == 0 { - // first occurrence in window - self.sum = self.sum.wrapping_add(v); - } - *cnt += 1; - } + self.apply_valid_values(arr, Self::update_value); Ok(()) } fn evaluate(&mut self) -> Result { // O(1) wrap of running sum - Ok(ScalarValue::Int64(Some(self.sum))) + Ok(ScalarValue::Int64( + (!self.counts.is_empty()).then_some(self.sum), + )) } fn size(&self) -> usize { - size_of_val(self) + // Estimate the owned map buckets; implementation-specific control bytes are excluded. + size_of_val(self) + self.counts.capacity() * size_of::<(i64, usize)>() } fn state(&mut self) -> Result> { @@ -581,11 +698,7 @@ impl Accumulator for SlidingDistinctSumAccumulator { if let ScalarValue::Int64(Some(v)) = ScalarValue::try_from_array(&*maybe_inner, idx)? { - let cnt = self.counts.entry(v).or_insert(0); - if *cnt == 0 { - self.sum = self.sum.wrapping_add(v); - } - *cnt += 1; + self.update_value(v); } } } @@ -594,16 +707,7 @@ impl Accumulator for SlidingDistinctSumAccumulator { fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let arr = values[0].as_primitive::(); - for &v in arr.values() { - if let Some(cnt) = self.counts.get_mut(&v) { - *cnt -= 1; - if *cnt == 0 { - // last copy leaving window - self.sum = self.sum.wrapping_sub(v); - self.counts.remove(&v); - } - } - } + self.apply_valid_values(arr, Self::retract_value); Ok(()) } @@ -611,3 +715,161 @@ impl Accumulator for SlidingDistinctSumAccumulator { true } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::{ + array::{Decimal128Array, Int64Array}, + buffer::{NullBuffer, ScalarBuffer}, + }; + use std::{ + mem::{size_of, size_of_val}, + sync::Arc, + }; + + #[test] + fn sliding_distinct_sum_ignores_null_slots() -> Result<()> { + let mut acc = SlidingDistinctSumAccumulator::try_new(&DataType::Int64)?; + + let values: ArrayRef = Arc::new(Int64Array::new( + ScalarBuffer::from(vec![42, 5, 5]), + Some(NullBuffer::from(vec![false, true, true])), + )); + acc.update_batch(&[values])?; + assert_eq!(acc.evaluate()?, ScalarValue::Int64(Some(5))); + + let retract: ArrayRef = Arc::new(Int64Array::new( + ScalarBuffer::from(vec![42, 5]), + Some(NullBuffer::from(vec![false, true])), + )); + acc.retract_batch(&[retract])?; + assert_eq!(acc.evaluate()?, ScalarValue::Int64(Some(5))); + + let retract_last: ArrayRef = + Arc::new(Int64Array::new(ScalarBuffer::from(vec![5]), None)); + acc.retract_batch(&[retract_last])?; + assert_eq!(acc.evaluate()?, ScalarValue::Int64(None)); + + Ok(()) + } + + fn expected_sliding_distinct_sum_size(acc: &SlidingDistinctSumAccumulator) -> usize { + size_of_val(acc) + acc.counts.capacity() * size_of::<(i64, usize)>() + } + + #[test] + fn sliding_distinct_sum_size_includes_hash_map_capacity() -> Result<()> { + let mut acc = SlidingDistinctSumAccumulator::try_new(&DataType::Int64)?; + let empty_size = acc.size(); + let values: ArrayRef = Arc::new(Int64Array::from(vec![1, 2, 3])); + acc.update_batch(&[Arc::clone(&values)])?; + + let expected = expected_sliding_distinct_sum_size(&acc); + assert!(acc.counts.capacity() > 0); + assert_eq!(acc.size(), expected); + assert!(acc.size() > empty_size); + + let initial_capacity = acc.counts.capacity(); + let additional_values: ArrayRef = + Arc::new(Int64Array::from_iter(4..4 + initial_capacity as i64 + 1)); + acc.update_batch(&[Arc::clone(&additional_values)])?; + + let grown_size = expected_sliding_distinct_sum_size(&acc); + assert!(acc.counts.capacity() > initial_capacity); + assert_eq!(acc.size(), grown_size); + assert!(acc.size() > expected); + + acc.retract_batch(&[values])?; + acc.retract_batch(&[additional_values])?; + assert!(acc.counts.is_empty()); + assert_eq!(acc.size(), grown_size); + + Ok(()) + } + + #[test] + fn sliding_distinct_sum_returns_null_for_all_null_frame() -> Result<()> { + let mut acc = SlidingDistinctSumAccumulator::try_new(&DataType::Int64)?; + + let values: ArrayRef = Arc::new(Int64Array::new( + ScalarBuffer::from(vec![99]), + Some(NullBuffer::from(vec![false])), + )); + acc.update_batch(&[values])?; + assert_eq!(acc.evaluate()?, ScalarValue::Int64(None)); + + Ok(()) + } + + #[test] + fn decimal_sum_accumulator_uses_widened_return_type() -> Result<()> { + let values: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(99_999), Some(99_999)]) + .with_precision_and_scale(5, 2)?, + ); + let mut acc = SumAccumulator::::new(DataType::Decimal128(15, 2)); + + acc.update_batch(&[values])?; + + assert_eq!( + acc.evaluate()?, + ScalarValue::Decimal128(Some(199_998), 15, 2) + ); + Ok(()) + } + + #[test] + fn sum_value_from_stats_widens_small_integer_sum() { + let statistics = datafusion_common::Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: vec![datafusion_common::ColumnStatistics { + sum_value: Precision::Exact(ScalarValue::Int32(Some(10))), + ..Default::default() + }], + }; + let return_type = DataType::Int64; + let expr: Arc = + Arc::new(Column::new("a", 0)); + let exprs = vec![expr]; + let statistics_args = StatisticsArgs { + statistics: &statistics, + return_type: &return_type, + is_distinct: false, + exprs: &exprs, + }; + + assert_eq!( + Sum::new().value_from_stats(&statistics_args), + Some(ScalarValue::Int64(Some(10))) + ); + } + + #[test] + fn sum_value_from_stats_casts_decimal_sum_to_return_type() { + let statistics = datafusion_common::Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Absent, + column_statistics: vec![datafusion_common::ColumnStatistics { + sum_value: Precision::Exact(ScalarValue::Decimal128(Some(12345), 5, 2)), + ..Default::default() + }], + }; + let return_type = DataType::Decimal128(15, 2); + let expr: Arc = + Arc::new(Column::new("a", 0)); + let exprs = vec![expr]; + let statistics_args = StatisticsArgs { + statistics: &statistics, + return_type: &return_type, + is_distinct: false, + exprs: &exprs, + }; + + assert_eq!( + Sum::new().value_from_stats(&statistics_args), + Some(ScalarValue::Decimal128(Some(12345), 15, 2)) + ); + } +} diff --git a/datafusion/functions-aggregate/src/variance.rs b/datafusion/functions-aggregate/src/variance.rs index ce3e00b9ffd91..b8e52f849a7cc 100644 --- a/datafusion/functions-aggregate/src/variance.rs +++ b/datafusion/functions-aggregate/src/variance.rs @@ -326,6 +326,23 @@ fn update(count: u64, mean: f64, m2: f64, value: f64) -> (u64, f64, f64) { (new_count, new_mean, new_m2) } +/// Inverse of [`update`]: removes a previously accumulated value. Retracting +/// from a state with one or zero values resets the state to empty. +#[inline] +fn retract(count: u64, mean: f64, m2: f64, value: f64) -> (u64, f64, f64) { + if count <= 1 { + return (0, 0.0, 0.0); + } + + let new_count = count - 1; + let delta1 = mean - value; + let new_mean = delta1 / new_count as f64 + mean; + let delta2 = new_mean - value; + let new_m2 = m2 - delta1 * delta2; + + (new_count, new_mean, new_m2) +} + impl Accumulator for VarianceAccumulator { fn state(&mut self) -> Result> { Ok(vec![ @@ -348,15 +365,8 @@ impl Accumulator for VarianceAccumulator { fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { let arr = as_float64_array(&values[0])?; for value in arr.iter().flatten() { - let new_count = self.count - 1; - let delta1 = self.mean - value; - let new_mean = delta1 / new_count as f64 + self.mean; - let delta2 = new_mean - value; - let new_m2 = self.m2 - delta1 * delta2; - - self.count -= 1; - self.mean = new_mean; - self.m2 = new_m2; + (self.count, self.mean, self.m2) = + retract(self.count, self.mean, self.m2, value) } Ok(()) @@ -521,8 +531,6 @@ impl GroupsAccumulator for VarianceGroupsAccumulator { &mut self, values: &[ArrayRef], group_indices: &[usize], - // Since aggregate filter should be applied in partial stage, in final stage there should be no filter - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 3, "two arguments to merge_batch"); @@ -575,6 +583,39 @@ impl GroupsAccumulator for VarianceGroupsAccumulator { ]) } + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + assert_eq!(values.len(), 1, "single argument to convert_to_state"); + let values = as_float64_array(&values[0])?; + + let len = values.len(); + let mut counts = Vec::with_capacity(len); + let mut means = Vec::with_capacity(len); + let mut m2s = Vec::with_capacity(len); + + for row in 0..len { + if values.is_valid(row) + && opt_filter + .is_none_or(|filter| filter.is_valid(row) && filter.value(row)) + { + counts.push(1); + means.push(values.value(row)); + } else { + counts.push(0); + means.push(0.0); + } + m2s.push(0.0); + } + + Ok(vec![ + Arc::new(UInt64Array::new(counts.into(), None)), + Arc::new(Float64Array::new(means.into(), None)), + Arc::new(Float64Array::new(m2s.into(), None)), + ]) + } fn size(&self) -> usize { self.m2s.capacity() * size_of::() + self.means.capacity() * size_of::() @@ -653,6 +694,75 @@ mod tests { use super::*; + #[test] + fn update_batch_ignores_nulls() -> Result<()> { + // An array with nulls must accumulate the same values as a dense + // array of its non-null values. + let dense: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])); + let sparse: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.0), + None, + Some(2.0), + Some(3.0), + None, + Some(4.0), + ])); + + let mut dense_acc = VarianceAccumulator::try_new(StatsType::Sample)?; + dense_acc.update_batch(std::slice::from_ref(&dense))?; + let mut sparse_acc = VarianceAccumulator::try_new(StatsType::Sample)?; + sparse_acc.update_batch(std::slice::from_ref(&sparse))?; + + // Sample variance of {1, 2, 3, 4} is 5/3 (all steps are exact in f64). + assert_eq!(dense_acc.evaluate()?, ScalarValue::Float64(Some(5.0 / 3.0))); + assert_eq!(dense_acc.evaluate()?, sparse_acc.evaluate()?); + Ok(()) + } + + #[test] + fn retract_batch_ignores_nulls() -> Result<()> { + let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0, 4.0])); + let dense_retract: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); + let sparse_retract: ArrayRef = + Arc::new(Float64Array::from(vec![Some(1.0), None, Some(2.0)])); + + let mut dense_acc = VarianceAccumulator::try_new(StatsType::Sample)?; + dense_acc.update_batch(std::slice::from_ref(&values))?; + dense_acc.retract_batch(std::slice::from_ref(&dense_retract))?; + let mut sparse_acc = VarianceAccumulator::try_new(StatsType::Sample)?; + sparse_acc.update_batch(std::slice::from_ref(&values))?; + sparse_acc.retract_batch(std::slice::from_ref(&sparse_retract))?; + + // Sample variance of the remaining {3, 4} is 0.5 (all steps are exact + // in f64). + assert_eq!(dense_acc.evaluate()?, ScalarValue::Float64(Some(0.5))); + assert_eq!(dense_acc.evaluate()?, sparse_acc.evaluate()?); + Ok(()) + } + + #[test] + fn retract_batch_resets_when_underflowing() -> Result<()> { + // Retracting more values than were accumulated resets to the empty + // state, with or without nulls in the retracted batch. + let values: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); + let dense_retract: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0, 3.0])); + let sparse_retract: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(1.0), + None, + Some(2.0), + Some(3.0), + ])); + + for retract in [&dense_retract, &sparse_retract] { + let mut acc = VarianceAccumulator::try_new(StatsType::Sample)?; + acc.update_batch(std::slice::from_ref(&values))?; + acc.retract_batch(std::slice::from_ref(retract))?; + assert_eq!(acc.get_count(), 0); + assert_eq!(acc.evaluate()?, ScalarValue::Float64(None)); + } + Ok(()) + } + #[test] fn test_groups_accumulator_merge_empty_states() -> Result<()> { let state_1 = vec![ @@ -666,12 +776,97 @@ mod tests { Arc::new(Float64Array::from(vec![1.0])), ]; let mut acc = VarianceGroupsAccumulator::new(StatsType::Sample); - acc.merge_batch(&state_1, &[0], None, 1)?; - acc.merge_batch(&state_2, &[0], None, 1)?; + acc.merge_batch(&state_1, &[0], 1)?; + acc.merge_batch(&state_2, &[0], 1)?; let result = acc.evaluate(EmitTo::All)?; let result = result.as_any().downcast_ref::().unwrap(); assert_eq!(result.len(), 1); assert_eq!(result.value(0), 1.0); Ok(()) } + + #[test] + fn convert_to_state_roundtrips_through_merge() -> Result<()> { + let values = Arc::new(Float64Array::from(vec![ + Some(1.0), + Some(2.0), + None, + Some(4.0), + Some(8.0), + Some(16.0), + Some(32.0), + ])) as ArrayRef; + let filter = BooleanArray::from(vec![ + Some(true), + Some(false), + Some(true), + None, + Some(true), + Some(true), + Some(true), + ]); + let group_indices = vec![0, 1, 0, 1, 0, 0, 0]; + + let mut direct = VarianceGroupsAccumulator::new(StatsType::Sample); + direct.update_batch( + std::slice::from_ref(&values), + &group_indices, + Some(&filter), + 2, + )?; + let direct = direct.evaluate(EmitTo::All)?; + + let converter = VarianceGroupsAccumulator::new(StatsType::Sample); + let state = + converter.convert_to_state(std::slice::from_ref(&values), Some(&filter))?; + let mut merged = VarianceGroupsAccumulator::new(StatsType::Sample); + merged.merge_batch(&state, &group_indices, 2)?; + let merged = merged.evaluate(EmitTo::All)?; + + let direct = direct.as_any().downcast_ref::().unwrap(); + let merged = merged.as_any().downcast_ref::().unwrap(); + assert_eq!(direct.len(), merged.len()); + for row in 0..direct.len() { + assert_eq!(direct.is_null(row), merged.is_null(row)); + if direct.is_valid(row) { + assert!((direct.value(row) - merged.value(row)).abs() < 1e-12); + } + } + Ok(()) + } + + #[test] + fn convert_to_state_preserves_empty_and_filtered_rows() -> Result<()> { + let converter = VarianceGroupsAccumulator::new(StatsType::Sample); + let empty_values = + Arc::new(Float64Array::from(Vec::>::new())) as ArrayRef; + let state = + converter.convert_to_state(std::slice::from_ref(&empty_values), None)?; + for state_array in &state { + assert_eq!(state_array.len(), 0); + assert_eq!(state_array.null_count(), 0); + } + + let values = + Arc::new(Float64Array::from(vec![Some(1.0), Some(2.0), None])) as ArrayRef; + let filter = BooleanArray::from(vec![Some(false), None, Some(false)]); + let group_indices = vec![0, 1, 0]; + let state = + converter.convert_to_state(std::slice::from_ref(&values), Some(&filter))?; + for state_array in &state { + assert_eq!(state_array.len(), values.len()); + assert_eq!(state_array.null_count(), 0); + } + + let counts = state[0].as_any().downcast_ref::().unwrap(); + assert_eq!(counts, &UInt64Array::from(vec![0, 0, 0])); + + let mut merged = VarianceGroupsAccumulator::new(StatsType::Sample); + merged.merge_batch(&state, &group_indices, 2)?; + let result = merged.evaluate(EmitTo::All)?; + let result = result.as_any().downcast_ref::().unwrap(); + assert_eq!(result.len(), 2); + assert_eq!(result.null_count(), 2); + Ok(()) + } } diff --git a/datafusion/functions-nested/benches/array_has.rs b/datafusion/functions-nested/benches/array_has.rs index f5e66d56c0efe..1a64f1cc4b160 100644 --- a/datafusion/functions-nested/benches/array_has.rs +++ b/datafusion/functions-nested/benches/array_has.rs @@ -15,7 +15,9 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, Int64Array, ListArray, StringArray}; +use arrow::array::{ + ArrayRef, Int64Array, LargeStringArray, ListArray, StringArray, StringViewArray, +}; use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, Field}; use criterion::{ @@ -43,17 +45,234 @@ fn criterion_benchmark(c: &mut Criterion) { for &size in &array_sizes { bench_array_has(c, size); + bench_array_has_array(c, size); bench_array_has_all(c, size); bench_array_has_any(c, size); } // Specific benchmarks for string arrays (common use case) bench_array_has_strings(c); + bench_array_has_array_strings(c); bench_array_has_all_strings(c); bench_array_has_any_strings(c); // Benchmark for array_has_any with one scalar arg bench_array_has_any_scalar(c); + + // Array-needle fast-path profile: null patterns, list length, row height. + bench_array_has_array_null_patterns(c); + bench_array_has_array_by_size(c); + bench_array_has_array_by_rows(c); +} + +/// Invoke `array_has` once with an array (column) needle -- exercises the +/// `array_has_dispatch_for_array` fast path. +fn run_array_needle_case( + group: &mut criterion::BenchmarkGroup<'_, criterion::measurement::WallTime>, + id: String, + haystack: ArrayRef, + needle: ArrayRef, + rows: usize, +) { + let config_options = Arc::new(ConfigOptions::default()); + let return_field: Arc = Field::new("result", DataType::Boolean, true).into(); + let arg_fields: Vec> = vec![ + Field::new("arr", haystack.data_type().clone(), false).into(), + Field::new("el", needle.data_type().clone(), false).into(), + ]; + let args = vec![ColumnarValue::Array(haystack), ColumnarValue::Array(needle)]; + group.bench_function(id, |b| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: rows, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }); + }); +} + +/// Build a `List` of `array_size` string elements per row (`{prefix}{i}`) with +/// the given element type (`Utf8` / `LargeUtf8` / `Utf8View`) and null density. +/// The prefix controls element length: a short one stays inline in a `Utf8View` +/// (<= 12 bytes), a long one spills to the data buffer. +fn string_list_array( + num_rows: usize, + array_size: usize, + null_density: f64, + prefix: &str, + element_type: &DataType, +) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(SEED); + let data = (0..num_rows * array_size).map(|_| { + if rng.random::() < null_density { + None + } else { + Some(format!("{prefix}{}", rng.random_range(0..array_size))) + } + }); + let values: ArrayRef = match element_type { + DataType::Utf8 => Arc::new(data.collect::()), + DataType::LargeUtf8 => Arc::new(data.collect::()), + DataType::Utf8View => Arc::new(data.collect::()), + other => panic!("unsupported string element type: {other}"), + }; + let offsets = (0..=num_rows) + .map(|i| (i * array_size) as i32) + .collect::>(); + Arc::new( + ListArray::try_new( + Arc::new(Field::new("item", element_type.clone(), true)), + OffsetBuffer::new(offsets.into()), + values, + None, + ) + .unwrap(), + ) +} + +/// Build a string needle column (one value per row) of the given element type. +fn string_value_array( + num_rows: usize, + range: usize, + prefix: &str, + element_type: &DataType, +) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(SEED + 2); + let data = + (0..num_rows).map(|_| Some(format!("{prefix}{}", rng.random_range(0..range)))); + match element_type { + DataType::Utf8 => Arc::new(data.collect::()), + DataType::LargeUtf8 => Arc::new(data.collect::()), + DataType::Utf8View => Arc::new(data.collect::()), + other => panic!("unsupported string element type: {other}"), + } +} + +/// Array needle, fixed list length (64), across null patterns. i64 covers +/// no-nulls found/not-found, 30% nulls found/not-found, all-null, and a +/// null-fill collision. Each string +/// element type (`Utf8`, `LargeUtf8`, `Utf8View`) covers no nulls / 30% nulls at +/// both short (inline, <= 12 byte) and long (> 12 byte, shared-prefix) element +/// lengths, plus all-null. `not_found` shifts the needle out of the value range. +fn bench_array_has_array_null_patterns(c: &mut Criterion) { + let (rows, size) = (10_000usize, 64usize); + let s = size as i64; + let mut group = c.benchmark_group("array_has_array_null_patterns"); + run_array_needle_case( + &mut group, + "i64/no_nulls".to_string(), + create_int64_list_array(rows, size, 0.0), + create_int64_value_array(rows, s, 0), + rows, + ); + // Worst case for the all-valid fold: non-null, no match -> the whole row is + // scanned (the branchless OR-reduction never short-circuits). + run_array_needle_case( + &mut group, + "i64/no_nulls_not_found".to_string(), + create_int64_list_array(rows, size, 0.0), + create_int64_value_array(rows, s, s), + rows, + ); + run_array_needle_case( + &mut group, + "i64/nulls30_found".to_string(), + create_int64_list_array(rows, size, 0.3), + create_int64_value_array(rows, s, 0), + rows, + ); + run_array_needle_case( + &mut group, + "i64/nulls30_not_found".to_string(), + create_int64_list_array(rows, size, 0.3), + create_int64_value_array(rows, s, s), + rows, + ); + run_array_needle_case( + &mut group, + "i64/all_null".to_string(), + create_int64_list_array(rows, size, 1.0), + create_int64_value_array(rows, s, 0), + rows, + ); + run_array_needle_case( + &mut group, + "i64/collision".to_string(), + create_int64_list_array(rows, size, 1.0), + create_int64_value_array(rows, 1, 0), + rows, + ); + // Short elements stay inline in a `Utf8View` (<= 12 bytes); long elements + // share a 4-byte prefix (the realistic case where the view prefix can't + // reject, forcing a buffer compare). `_short` / `_long` labels distinguish + // them; all-null has no content so it is length-independent. + let short = "value_"; // "value_0".."value_63": <= 8 bytes, inline + let long = "long_element_string_value_"; // ~28 bytes, spills to the buffer + for (type_label, element_type) in [ + ("utf8", DataType::Utf8), + ("largeutf8", DataType::LargeUtf8), + ("utf8view", DataType::Utf8View), + ] { + for (len_label, prefix) in [("short", short), ("long", long)] { + for (pat_label, density) in [("no_nulls", 0.0), ("nulls30", 0.3)] { + run_array_needle_case( + &mut group, + format!("{type_label}_{len_label}/{pat_label}"), + string_list_array(rows, size, density, prefix, &element_type), + string_value_array(rows, size, prefix, &element_type), + rows, + ); + } + } + run_array_needle_case( + &mut group, + format!("{type_label}/all_null"), + string_list_array(rows, size, 1.0, short, &element_type), + string_value_array(rows, size, short, &element_type), + rows, + ); + } + group.finish(); +} + +/// Array needle, i64, 30% element nulls, not found, across list lengths. +fn bench_array_has_array_by_size(c: &mut Criterion) { + let rows = 10_000usize; + let mut group = c.benchmark_group("array_has_array_by_size"); + for size in [8usize, 32, 128, 256, 512, 1024] { + let s = size as i64; + run_array_needle_case( + &mut group, + size.to_string(), + create_int64_list_array(rows, size, 0.3), + create_int64_value_array(rows, s, s), + rows, + ); + } + group.finish(); +} + +/// Array needle, i64, 8 elems/row, 30% nulls, not found, across row counts. +fn bench_array_has_array_by_rows(c: &mut Criterion) { + let (size, s) = (8usize, 8i64); + let mut group = c.benchmark_group("array_has_array_by_rows"); + for rows in [10_000usize, 100_000, 1_000_000] { + run_array_needle_case( + &mut group, + rows.to_string(), + create_int64_list_array(rows, size, 0.3), + create_int64_value_array(rows, s, s), + rows, + ); + } + group.finish(); } fn bench_array_has(c: &mut Criterion, array_size: usize) { @@ -119,6 +338,136 @@ fn bench_array_has(c: &mut Criterion, array_size: usize) { group.finish(); } +/// Benchmarks array_has where the needle is an array (a column with one value +/// per row) rather than a scalar. +fn bench_array_has_array(c: &mut Criterion, array_size: usize) { + let mut group = c.benchmark_group("array_has_array_i64"); + let haystack = create_int64_list_array(NUM_ROWS, array_size, NULL_DENSITY); + let config_options = Arc::new(ConfigOptions::default()); + let return_field: Arc = Field::new("result", DataType::Boolean, true).into(); + let arg_fields: Vec> = vec![ + Field::new("arr", haystack.data_type().clone(), false).into(), + Field::new("el", DataType::Int64, false).into(), + ]; + + // Needle values drawn from the same range as the haystack values, so many + // rows find a match (and the inner loop can short-circuit). + let needle_found = create_int64_value_array(NUM_ROWS, array_size as i64, 0); + let args_found = vec![ + ColumnarValue::Array(haystack.clone()), + ColumnarValue::Array(needle_found), + ]; + group.bench_with_input( + BenchmarkId::new("found", array_size), + &array_size, + |b, _| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args_found.clone(), + arg_fields: arg_fields.clone(), + number_rows: NUM_ROWS, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }) + }, + ); + + // Needle values outside the haystack range: never matches, so every row + // scans its full element list (worst case for the inner loop). + let needle_not_found = + create_int64_value_array(NUM_ROWS, array_size as i64, array_size as i64); + let args_not_found = vec![ + ColumnarValue::Array(haystack.clone()), + ColumnarValue::Array(needle_not_found), + ]; + group.bench_with_input( + BenchmarkId::new("not_found", array_size), + &array_size, + |b, _| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args_not_found.clone(), + arg_fields: arg_fields.clone(), + number_rows: NUM_ROWS, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }) + }, + ); + + group.finish(); +} + +fn bench_array_has_array_strings(c: &mut Criterion) { + let mut group = c.benchmark_group("array_has_array_strings"); + let config_options = Arc::new(ConfigOptions::default()); + let return_field: Arc = Field::new("result", DataType::Boolean, true).into(); + + let sizes = vec![10, 100, 500]; + + for &size in &sizes { + let haystack = create_string_list_array(NUM_ROWS, size, NULL_DENSITY); + let arg_fields: Vec> = vec![ + Field::new("arr", haystack.data_type().clone(), false).into(), + Field::new("el", DataType::Utf8, false).into(), + ]; + + let needle_found = create_string_value_array(NUM_ROWS, size, "value_"); + let args_found = vec![ + ColumnarValue::Array(haystack.clone()), + ColumnarValue::Array(needle_found), + ]; + group.bench_with_input(BenchmarkId::new("found", size), &size, |b, _| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args_found.clone(), + arg_fields: arg_fields.clone(), + number_rows: NUM_ROWS, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }) + }); + + let needle_not_found = create_string_value_array(NUM_ROWS, size, "missing_"); + let args_not_found = vec![ + ColumnarValue::Array(haystack.clone()), + ColumnarValue::Array(needle_not_found), + ]; + group.bench_with_input(BenchmarkId::new("not_found", size), &size, |b, _| { + let udf = ArrayHas::new(); + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args_not_found.clone(), + arg_fields: arg_fields.clone(), + number_rows: NUM_ROWS, + return_field: return_field.clone(), + config_options: config_options.clone(), + }) + .unwrap(), + ) + }) + }); + } + + group.finish(); +} + fn bench_array_has_all(c: &mut Criterion, array_size: usize) { let mut group = c.benchmark_group("array_has_all"); let haystack = create_int64_list_array(NUM_ROWS, array_size, NULL_DENSITY); @@ -659,6 +1008,30 @@ fn create_int64_list_array( ) } +/// Create an `Int64Array` of `num_rows` non-null values in `[offset, offset + +/// range)`, used as an array needle for `array_has`. +fn create_int64_value_array(num_rows: usize, range: i64, offset: i64) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(SEED + 2); + let values = (0..num_rows) + .map(|_| Some(rng.random_range(0..range) + offset)) + .collect::(); + Arc::new(values) +} + +/// Create a `StringArray` of `num_rows` non-null values like "{prefix}{idx}" +/// where `idx` is drawn from `[0, range)`, used as an array needle for +/// `array_has`. +fn create_string_value_array(num_rows: usize, range: usize, prefix: &str) -> ArrayRef { + let mut rng = StdRng::seed_from_u64(SEED + 2); + let values = (0..num_rows) + .map(|_| { + let idx = rng.random_range(0..range); + Some(format!("{prefix}{idx}")) + }) + .collect::(); + Arc::new(values) +} + /// Like `create_int64_list_array` but values are offset so they won't /// appear in a standard list array (useful for "not found" benchmarks). fn create_int64_list_array_with_offset( diff --git a/datafusion/functions-nested/benches/arrays_zip.rs b/datafusion/functions-nested/benches/arrays_zip.rs index bc82b2978cc42..812e5e3dbec8a 100644 --- a/datafusion/functions-nested/benches/arrays_zip.rs +++ b/datafusion/functions-nested/benches/arrays_zip.rs @@ -109,7 +109,7 @@ fn bench_arrays_zip(c: &mut Criterion, name: &str, null_density: f64) { } fn criterion_benchmark(c: &mut Criterion) { - bench_arrays_zip(c, "arrays_zip_no_nulls_8192", 0.0); + bench_arrays_zip(c, "arrays_zip_perfect_zip_8192", 0.0); bench_arrays_zip(c, "arrays_zip_10pct_nulls_8192", 0.1); } diff --git a/datafusion/functions-nested/benches/map.rs b/datafusion/functions-nested/benches/map.rs index 67e7f314d2515..9cc4289ca1f1c 100644 --- a/datafusion/functions-nested/benches/map.rs +++ b/datafusion/functions-nested/benches/map.rs @@ -28,8 +28,7 @@ use datafusion_expr::planner::ExprPlanner; use datafusion_expr::{ColumnarValue, Expr, ScalarFunctionArgs}; use datafusion_functions_nested::map::map_udf; use datafusion_functions_nested::planner::NestedFunctionPlanner; -use rand::Rng; -use rand::prelude::ThreadRng; +use rand::prelude::*; use std::collections::HashSet; use std::hash::Hash; use std::hint::black_box; @@ -38,10 +37,7 @@ use std::sync::Arc; const MAP_ROWS: usize = 1000; const MAP_KEYS_PER_ROW: usize = 1000; -fn gen_unique_values( - rng: &mut ThreadRng, - mut make_value: impl FnMut(i32) -> T, -) -> Vec +fn gen_unique_values(rng: &mut StdRng, mut make_value: impl FnMut(i32) -> T) -> Vec where T: Eq + Hash, { @@ -64,15 +60,15 @@ fn gen_repeat_values(values: &[T], repeats: usize) -> Vec { repeated } -fn gen_utf8_values(rng: &mut ThreadRng) -> Vec { +fn gen_utf8_values(rng: &mut StdRng) -> Vec { gen_unique_values(rng, |value| value.to_string()) } -fn gen_binary_values(rng: &mut ThreadRng) -> Vec> { +fn gen_binary_values(rng: &mut StdRng) -> Vec> { gen_unique_values(rng, |value| value.to_le_bytes().to_vec()) } -fn gen_primitive_values(rng: &mut ThreadRng) -> Vec { +fn gen_primitive_values(rng: &mut StdRng) -> Vec { gen_unique_values(rng, |value| value) } @@ -122,7 +118,7 @@ fn bench_map_case(c: &mut Criterion, name: &str, keys: ArrayRef, values: ArrayRe fn criterion_benchmark(c: &mut Criterion) { c.bench_function("make_map_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let keys = gen_utf8_values(&mut rng); let values = gen_primitive_values(&mut rng); let mut buffer = Vec::new(); @@ -143,7 +139,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); }); - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let values = Arc::new(Int32Array::from(gen_repeat_values( &gen_primitive_values(&mut rng), MAP_ROWS, diff --git a/datafusion/functions-nested/src/array_add.rs b/datafusion/functions-nested/src/array_add.rs new file mode 100644 index 0000000000000..dd170911fb8e0 --- /dev/null +++ b/datafusion/functions-nested/src/array_add.rs @@ -0,0 +1,203 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ScalarUDFImpl`] definitions for array_add function. + +use crate::utils::{coerce_array_math_arg_types, make_scalar_function}; +use arrow::array::{ + Array, ArrayRef, Float64Array, GenericListArray, NullBufferBuilder, OffsetSizeTrait, +}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{ + DataType, + DataType::{LargeList, List}, + Field, +}; +use datafusion_common::cast::{as_float64_array, as_generic_list_array}; +use datafusion_common::{Result, exec_err, utils::take_function_args}; +use datafusion_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_macros::user_doc; +use std::sync::Arc; + +make_udf_expr_and_func!( + ArrayAdd, + array_add, + array1 array2, + "returns the element-wise sum of two numeric arrays.", + array_add_udf +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns the element-wise sum of two numeric arrays of equal length, computed as `array1[i] + array2[i]` per position. NULL is propagated per element: if either input element at position `i` is NULL, the corresponding output element is NULL (positions are preserved). Returns NULL if either entire input array is NULL. Errors if the per-row lengths differ. Returns an empty array if both inputs are empty.", + syntax_example = "array_add(array1, array2)", + sql_example = r#"```sql +> select array_add([1.0, 2.0, 3.0], [10.0, 20.0, 30.0]); ++---------------------------------------------------------+ +| array_add(List([1.0,2.0,3.0]),List([10.0,20.0,30.0])) | ++---------------------------------------------------------+ +| [11.0, 22.0, 33.0] | ++---------------------------------------------------------+ +```"#, + argument( + name = "array1", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ), + argument( + name = "array2", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArrayAdd { + signature: Signature, + aliases: Vec, +} + +impl Default for ArrayAdd { + fn default() -> Self { + Self::new() + } +} + +impl ArrayAdd { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + aliases: vec!["list_add".to_string()], + } + } +} + +impl ScalarUDFImpl for ArrayAdd { + fn name(&self) -> &str { + "array_add" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + // After `coerce_types`, both args share the same List/LargeList shape. + Ok(arg_types[0].clone()) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [_, _] = take_function_args(self.name(), arg_types)?; + coerce_array_math_arg_types(self.name(), arg_types) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(array_add_inner)(&args.args) + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} + +fn array_add_inner(args: &[ArrayRef]) -> Result { + let [array1, array2] = take_function_args("array_add", args)?; + match (array1.data_type(), array2.data_type()) { + (List(_), List(_)) => general_array_add::(array1, array2), + (LargeList(_), LargeList(_)) => general_array_add::(array1, array2), + (arg_type1, arg_type2) => exec_err!( + "array_add received unexpected types after coercion: {arg_type1} and {arg_type2}" + ), + } +} + +fn general_array_add( + lhs: &ArrayRef, + rhs: &ArrayRef, +) -> Result { + let lhs = as_generic_list_array::(lhs)?; + let rhs = as_generic_list_array::(rhs)?; + + let lhs_values = as_float64_array(lhs.values())?; + let rhs_values = as_float64_array(rhs.values())?; + let lhs_offsets = lhs.value_offsets(); + let rhs_offsets = rhs.value_offsets(); + + // Row-level validity: a row is valid iff both sides are valid at that row. + let row_nulls = NullBuffer::union(lhs.nulls(), rhs.nulls()); + + let mut out_values: Vec = Vec::with_capacity(lhs_values.len()); + let mut out_inner_nulls = NullBufferBuilder::new(lhs_values.len()); + let mut out_offsets = Vec::::with_capacity(lhs.len() + 1); + out_offsets.push(O::zero()); + + for row in 0..lhs.len() { + // Whole-row NULL on either side -> NULL output row, no elements. + if row_nulls.as_ref().is_some_and(|nb| nb.is_null(row)) { + out_offsets.push(out_offsets[row]); + continue; + } + + let start1 = lhs_offsets[row].as_usize(); + let len1 = lhs.value_length(row).as_usize(); + let start2 = rhs_offsets[row].as_usize(); + let len2 = rhs.value_length(row).as_usize(); + + if len1 != len2 { + return exec_err!( + "array_add requires both list inputs to have the same length per row, got {len1} and {len2} at row {row}" + ); + } + + let l_slice = lhs_values.slice(start1, len1); + let r_slice = rhs_values.slice(start2, len2); + + let l_vals = l_slice.values(); + let r_vals = r_slice.values(); + + for i in 0..len1 { + out_values.push(l_vals[i] + r_vals[i]); + } + + // Per-element validity: position `i` is valid iff both lhs[i] and rhs[i] + // are valid. `NullBuffer::union` returns `None` when both sides are + // entirely valid. + match NullBuffer::union(l_slice.nulls(), r_slice.nulls()) { + Some(nb) => out_inner_nulls.append_buffer(&nb), + None => out_inner_nulls.append_n_non_nulls(len1), + } + + out_offsets.push(out_offsets[row] + O::usize_as(len1)); + } + + let values_array = Arc::new(Float64Array::new( + out_values.into(), + out_inner_nulls.finish(), + )); + let field = Arc::new(Field::new_list_field(DataType::Float64, true)); + + Ok(Arc::new(GenericListArray::::try_new( + field, + OffsetBuffer::new(out_offsets.into()), + values_array, + row_nulls, + )?)) +} diff --git a/datafusion/functions-nested/src/array_any_match.rs b/datafusion/functions-nested/src/array_any_match.rs index 3ce43a23c2124..0f620f18bd8f2 100644 --- a/datafusion/functions-nested/src/array_any_match.rs +++ b/datafusion/functions-nested/src/array_any_match.rs @@ -15,28 +15,26 @@ // specific language governing permissions and limitations // under the License. -//! [`HigherOrderUDF`] definitions for array_any_match function. +//! [`datafusion_expr::HigherOrderUDF`] definitions for array_any_match function. use arrow::{ - array::{Array, AsArray, BooleanArray, BooleanBuilder, new_null_array}, + array::{Array, BooleanArray, BooleanBuilder}, buffer::NullBuffer, - compute::take_arrays, - datatypes::{ArrowNativeType, DataType, Field, FieldRef}, -}; -use datafusion_common::{ - Result, exec_datafusion_err, exec_err, plan_err, - utils::{ - adjust_offsets_for_slice, list_values, list_values_row_number, take_function_args, - }, + datatypes::{DataType, Field, FieldRef}, }; +use datafusion_common::{Result, plan_err, utils::take_function_args}; use datafusion_expr::{ ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, - HigherOrderSignature, HigherOrderUDF, LambdaParametersProgress, ValueOrLambda, + HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, Volatility, }; use datafusion_macros::user_doc; use std::{fmt::Debug, sync::Arc}; +use crate::lambda_utils::{ + SingleListLambdaResult, coerce_single_list_arg, evaluate_single_list_predicate, +}; + make_higher_order_function_expr_and_func!( ArrayAnyMatch, array_any_match, @@ -106,7 +104,7 @@ fn any_match_for_range( if any_null { None } else { Some(false) } } -impl HigherOrderUDF for ArrayAnyMatch { +impl HigherOrderUDFImpl for ArrayAnyMatch { fn name(&self) -> &str { "array_any_match" } @@ -120,30 +118,7 @@ impl HigherOrderUDF for ArrayAnyMatch { } fn coerce_value_types(&self, arg_types: &[DataType]) -> Result> { - let [list] = arg_types else { - return plan_err!( - "{} function requires 1 value argument, got {}", - self.name(), - arg_types.len() - ); - }; - - let coerced = match list { - DataType::List(_) | DataType::LargeList(_) => list.clone(), - DataType::ListView(field) | DataType::FixedSizeList(field, _) => { - DataType::List(Arc::clone(field)) - } - DataType::LargeListView(field) => DataType::LargeList(Arc::clone(field)), - _ => { - return plan_err!( - "{} expected a list as first argument, got {}", - self.name(), - list - ); - } - }; - - Ok(vec![coerced]) + coerce_single_list_arg(self.name(), arg_types) } fn lambda_parameters( @@ -171,85 +146,35 @@ impl HigherOrderUDF for ArrayAnyMatch { &self, args: HigherOrderReturnFieldArgs, ) -> Result> { - let [ValueOrLambda::Value(list), _] = + let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(lambda)] = take_function_args(self.name(), args.arg_fields)? else { return plan_err!("{} expects a value as first argument", self.name()); }; - let nullable = list.is_nullable(); + let nullable = list.is_nullable() || lambda.is_nullable(); Ok(Arc::new(Field::new("", DataType::Boolean, nullable))) } fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { - let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(lambda)] = - take_function_args(self.name(), &args.args)? - else { - return exec_err!("{} expects a value followed by a lambda", self.name()); + let evaluated = match evaluate_single_list_predicate(self.name(), &args)? { + SingleListLambdaResult::EarlyReturn(v) => return Ok(v), + SingleListLambdaResult::Ready(v) => v, }; - let list_array = list.to_array(args.number_rows)?; - - // fast path: fully null input — also required for FixedSizeList which can't be - // handled by clear_null_values when fully null - if list_array.null_count() == list_array.len() { - return Ok(ColumnarValue::Array(new_null_array( - args.return_type(), - list_array.len(), - ))); - } - - let list_values = list_values(&list_array)?; - - let values_param = || Ok(Arc::clone(&list_values)); - - let predicate_results = lambda - .evaluate(&[&values_param], |arrays| { - let indices = list_values_row_number(&list_array)?; - Ok(take_arrays(arrays, &indices, None)?) - })? - .into_array(list_values.len())?; - - let predicate_bool = predicate_results - .as_any() - .downcast_ref::() - .ok_or_else(|| { - exec_datafusion_err!( - "{} predicate must return boolean array", - self.name() - ) - })?; - - let mut values = BooleanBuilder::with_capacity(list_array.len()); - - // Maps predicate results (flat over all elements) back to one Boolean per row. - // Uses adjusted offsets so sliced lists index correctly into the predicate array. - macro_rules! process_list { - ($list_typed:expr) => {{ - let offsets = adjust_offsets_for_slice($list_typed); - for i in 0..$list_typed.len() { - let start = offsets[i].as_usize(); - let end = offsets[i + 1].as_usize(); - // any_match_for_range returns None when nulls poison the result; - // null rows produce an empty range and return Some(false), but their - // null bit is preserved by attaching the original null bitmap below. - values.append_option(any_match_for_range(predicate_bool, start, end)); - } - }}; - } + let predicate = evaluated.boolean_predicate(self.name())?; - match list_array.data_type() { - DataType::List(_) => { - process_list!(list_array.as_list::()); - } - DataType::LargeList(_) => { - process_list!(list_array.as_list::()); - } - other => return exec_err!("expected list, got {other}"), + let mut values = BooleanBuilder::with_capacity(evaluated.len()); + for i in 0..evaluated.len() { + let (start, end) = evaluated.row_range(i); + // any_match_for_range returns None when nulls poison the result; + // null rows produce an empty range and return Some(false), but their + // null bit is preserved by attaching the original null bitmap below. + values.append_option(any_match_for_range(&predicate, start, end)); } let (boolean_buffer, predicate_nulls) = values.finish().into_parts(); // Merge: a row is null if the input list row was null or the predicate returned null. - let nulls = NullBuffer::union(list_array.nulls(), predicate_nulls.as_ref()); + let nulls = NullBuffer::union(evaluated.nulls(), predicate_nulls.as_ref()); Ok(ColumnarValue::Array(Arc::new(BooleanArray::new( boolean_buffer, nulls, @@ -272,14 +197,19 @@ mod tests { }; use datafusion_common::{DFSchema, Result}; use datafusion_expr::{ - Expr, col, + Expr, HigherOrderReturnFieldArgs, HigherOrderUDFImpl, ValueOrLambda, col, execution_props::ExecutionProps, expr::{HigherOrderFunction, LambdaVariable}, lambda, lit, + physical_planning_context::PhysicalPlanningContext, }; use datafusion_physical_expr::create_physical_expr; - use crate::array_any_match::array_any_match_higher_order_function; + use crate::array_any_match::{ArrayAnyMatch, array_any_match_higher_order_function}; + use crate::lambda_utils::test_utils::{ + create_i32_large_list, create_i32_list, eval_hof_on_i32_list, + eval_hof_on_i32_list_with_outer, v, + }; fn run_any_match( list: impl arrow::array::Array + Clone + 'static, @@ -311,6 +241,7 @@ mod tests { )), &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), )? .evaluate(&RecordBatch::try_new( Arc::clone(schema.inner()), @@ -344,6 +275,7 @@ mod tests { )), &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), )? .evaluate(&RecordBatch::try_new( Arc::clone(schema.inner()), @@ -413,6 +345,44 @@ mod tests { Ok(()) } + #[test] + fn test_any_match_return_field_nullability() -> Result<()> { + for list_nullable in [true, false] { + for lambda_nullable in [true, false] { + let list = Arc::new(Field::new( + "list", + DataType::new_list(DataType::Int32, true), + list_nullable, + )); + let lambda = + Arc::new(Field::new("predicate", DataType::Boolean, lambda_nullable)); + let arg_fields = [ + ValueOrLambda::Value(Arc::clone(&list)), + ValueOrLambda::Lambda(Arc::clone(&lambda)), + ]; + let scalar_arguments = [None, None]; + + let result = ArrayAnyMatch::new().return_field_from_args( + HigherOrderReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalar_arguments, + }, + )?; + + assert_eq!( + result, + Arc::new(Field::new( + "", + DataType::Boolean, + list_nullable || lambda_nullable, + )) + ); + } + } + + Ok(()) + } + // Predicate must not be evaluated on elements belonging to null rows. // The 10 in the null row would satisfy x > 5, but the row result must be None. #[test] @@ -480,4 +450,44 @@ mod tests { ); Ok(()) } + + #[test] + fn test_any_match_large_list_parity() -> Result<()> { + let list = create_i32_large_list( + vec![1, 2, 3], + OffsetBuffer::::from_lengths(vec![3]), + None, + ); + let result = eval_hof_on_i32_list( + array_any_match_higher_order_function(), + list, + v().gt(lit(2i32)), + )?; + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &BooleanArray::from(vec![Some(true)]) + ); + Ok(()) + } + + #[test] + fn test_any_match_captured_outer_column() -> Result<()> { + let list = create_i32_list( + vec![1, 50, 4, 50, 7, 50], + OffsetBuffer::::from_lengths(vec![2, 2, 2]), + None, + ); + let number = Int32Array::from(vec![10, 40, 60]); + let result = eval_hof_on_i32_list_with_outer( + array_any_match_higher_order_function(), + list, + number, + v().gt(col("number")), + )?; + assert_eq!( + result.as_any().downcast_ref::().unwrap(), + &BooleanArray::from(vec![Some(true), Some(true), Some(false)]) + ); + Ok(()) + } } diff --git a/datafusion/functions-nested/src/array_avg.rs b/datafusion/functions-nested/src/array_avg.rs new file mode 100644 index 0000000000000..8133d3cac86f0 --- /dev/null +++ b/datafusion/functions-nested/src/array_avg.rs @@ -0,0 +1,174 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ScalarUDFImpl`] definitions for array_avg function. + +use crate::utils::make_scalar_function; +use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait}; +use arrow::datatypes::{ + DataType, + DataType::{FixedSizeList, LargeList, List, Null}, + Field, +}; +use datafusion_common::cast::{as_float64_array, as_generic_list_array}; +use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only}; +use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args}; +use datafusion_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_macros::user_doc; +use std::sync::Arc; + +make_udf_expr_and_func!( + ArrayAvg, + array_avg, + array, + "returns the arithmetic mean of elements in a numeric array.", + array_avg_udf +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns the arithmetic mean (sum divided by count) of the elements of the input array. NULL elements are skipped (per SQL aggregate convention) and excluded from the count. Returns NULL if the input row is NULL, every element is NULL, or the array is empty.", + syntax_example = "array_avg(array)", + sql_example = r#"```sql +> select array_avg([1.0, 2.0, 3.0]); ++----------------------------+ +| array_avg(List([1.0,2.0,3.0])) | ++----------------------------+ +| 2.0 | ++----------------------------+ +```"#, + argument( + name = "array", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArrayAvg { + signature: Signature, + aliases: Vec, +} + +impl Default for ArrayAvg { + fn default() -> Self { + Self::new() + } +} + +impl ArrayAvg { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + aliases: vec!["list_avg".to_string()], + } + } +} + +impl ScalarUDFImpl for ArrayAvg { + fn name(&self) -> &str { + "array_avg" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [arg_type] = take_function_args(self.name(), arg_types)?; + let coercion = Some(&ListCoercion::FixedSizedListToList); + + if !matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) { + return plan_err!("{} does not support type {arg_type}", self.name()); + } + + let coerced = if matches!(arg_type, Null) { + List(Arc::new(Field::new_list_field(DataType::Float64, true))) + } else { + coerced_type_with_base_type_only(arg_type, &DataType::Float64, coercion) + }; + + Ok(vec![coerced]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(array_avg_inner)(&args.args) + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} + +fn array_avg_inner(args: &[ArrayRef]) -> Result { + let [array] = take_function_args("array_avg", args)?; + match array.data_type() { + List(_) => general_array_avg::(array), + LargeList(_) => general_array_avg::(array), + arg_type => { + internal_err!("array_avg received unexpected type after coercion: {arg_type}") + } + } +} + +fn general_array_avg(array: &ArrayRef) -> Result { + let list_array = as_generic_list_array::(array)?; + let values = as_float64_array(list_array.values())?; + let offsets = list_array.value_offsets(); + + let mut builder = Float64Array::builder(list_array.len()); + + for row in 0..list_array.len() { + if list_array.is_null(row) { + builder.append_null(); + continue; + } + + let start = offsets[row].as_usize(); + let end = offsets[row + 1].as_usize(); + + // Skip NULL elements per SQL aggregate convention (matches PostgreSQL + // AVG, DuckDB list_avg, Spark aggregate). Empty arrays and all-NULL + // arrays both yield NULL — same behavior as SQL AVG over an empty + // set or all-NULL column. + let mut sum = 0.0_f64; + let mut count: u64 = 0; + for i in start..end { + if values.is_valid(i) { + sum += values.value(i); + count += 1; + } + } + + if count > 0 { + builder.append_value(sum / count as f64); + } else { + builder.append_null(); + } + } + + Ok(Arc::new(builder.finish())) +} diff --git a/datafusion/functions-nested/src/array_compact.rs b/datafusion/functions-nested/src/array_compact.rs index 11be494b5b20f..4222d6264bebe 100644 --- a/datafusion/functions-nested/src/array_compact.rs +++ b/datafusion/functions-nested/src/array_compact.rs @@ -130,14 +130,22 @@ fn compact_list( field: &Arc, ) -> Result { let values = list_array.values(); - - // Fast path: no nulls in values, return input unchanged - if values.null_count() == 0 { + // Use logical nulls so element types without a validity buffer + // (e.g. NullArray) are still treated as null. + let Some(values_nulls) = values.logical_nulls() else { + // Fast path: no validity buffer, no nulls to remove + return Ok(Arc::new(list_array.clone())); + }; + let values_null_count = values_nulls.null_count(); + if values_null_count == 0 { + // Fast path: validity buffer present but no nulls set return Ok(Arc::new(list_array.clone())); } + let list_nulls = list_array.nulls(); + let list_offsets = list_array.offsets(); let original_data = values.to_data(); - let capacity = original_data.len() - values.null_count(); + let capacity = original_data.len() - values_null_count; let mut offsets = Vec::::with_capacity(list_array.len() + 1); offsets.push(O::zero()); let mut mutable = MutableArrayData::with_capacities( @@ -147,25 +155,25 @@ fn compact_list( ); for row_index in 0..list_array.len() { - if list_array.nulls().is_some_and(|n| n.is_null(row_index)) { + if list_nulls.is_some_and(|n| n.is_null(row_index)) { offsets.push(offsets[row_index]); continue; } - let start = list_array.offsets()[row_index].as_usize(); - let end = list_array.offsets()[row_index + 1].as_usize(); - let mut copied = 0usize; + let start = list_offsets[row_index].as_usize(); + let end = list_offsets[row_index + 1].as_usize(); + let row_null_count = values_nulls.slice(start, end - start).null_count(); + let kept = (end - start) - row_null_count; // Batch consecutive non-null elements into single extend() calls // to reduce per-element overhead. For [1, 2, NULL, 3, 4] this // produces 2 extend calls (0..2, 3..5) instead of 4 individual ones. let mut batch_start: Option = None; for i in start..end { - if values.is_null(i) { + if values_nulls.is_null(i) { // Null breaks the current batch — flush it if let Some(bs) = batch_start { - mutable.extend(0, bs, i); - copied += i - bs; + mutable.try_extend(0, bs, i)?; batch_start = None; } } else if batch_start.is_none() { @@ -174,11 +182,10 @@ fn compact_list( } // Flush any remaining batch after the loop if let Some(bs) = batch_start { - mutable.extend(0, bs, end); - copied += end - bs; + mutable.try_extend(0, bs, end)?; } - offsets.push(offsets[row_index] + O::usize_as(copied)); + offsets.push(offsets[row_index] + O::usize_as(kept)); } let new_values = make_array(mutable.freeze()); @@ -186,6 +193,6 @@ fn compact_list( Arc::clone(field), OffsetBuffer::new(offsets.into()), new_values, - list_array.nulls().cloned(), + list_nulls.cloned(), )?)) } diff --git a/datafusion/functions-nested/src/array_filter.rs b/datafusion/functions-nested/src/array_filter.rs index f8b7fc35404a8..3439699433272 100644 --- a/datafusion/functions-nested/src/array_filter.rs +++ b/datafusion/functions-nested/src/array_filter.rs @@ -15,31 +15,28 @@ // specific language governing permissions and limitations // under the License. -//! [`HigherOrderUDF`] definitions for array_filter function. +//! [`datafusion_expr::HigherOrderUDF`] definitions for array_filter function. use arrow::{ array::{ Array, ArrayRef, AsArray, BooleanArray, LargeListArray, ListArray, - OffsetBufferBuilder, OffsetSizeTrait, new_empty_array, + OffsetSizeTrait, new_empty_array, }, buffer::{OffsetBuffer, ScalarBuffer}, - compute::{filter as arrow_filter, take_arrays}, + compute::filter as arrow_filter, datatypes::{DataType, Field, FieldRef}, }; -use datafusion_common::{ - Result, ScalarValue, exec_err, - utils::{adjust_offsets_for_slice, list_values_row_number}, -}; +use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, - HigherOrderSignature, HigherOrderUDF, LambdaParametersProgress, ValueOrLambda, + HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, Volatility, }; use datafusion_macros::user_doc; use std::sync::Arc; use crate::lambda_utils::{ - ListValuesResult, coerce_single_list_arg, extract_list_values, + SingleListLambdaResult, coerce_single_list_arg, evaluate_single_list_predicate, single_list_lambda_parameters, value_lambda_pair, }; @@ -96,7 +93,7 @@ impl ArrayFilter { } } -impl HigherOrderUDF for ArrayFilter { +impl HigherOrderUDFImpl for ArrayFilter { fn name(&self) -> &str { "array_filter" } @@ -130,12 +127,9 @@ impl HigherOrderUDF for ArrayFilter { } fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { - let (list, lambda) = value_lambda_pair(self.name(), &args.args)?; - let list_array = list.to_array(args.number_rows)?; - - let list_values = match extract_list_values(&list_array, args.return_type())? { - ListValuesResult::EarlyReturn(v) => return Ok(v), - ListValuesResult::Values(v) => v, + let evaluated = match evaluate_single_list_predicate(self.name(), &args)? { + SingleListLambdaResult::EarlyReturn(v) => return Ok(v), + SingleListLambdaResult::Ready(v) => v, }; let field = match args.return_field.data_type() { @@ -149,56 +143,47 @@ impl HigherOrderUDF for ArrayFilter { } }; - let values_param = || Ok(Arc::clone(&list_values)); - let predicate_output = lambda.evaluate(&[&values_param], |arrays| { - let indices = list_values_row_number(&list_array)?; - Ok(take_arrays(arrays, &indices, None)?) - })?; - // Scalar predicate short-circuit: x -> true or x -> false/null - if let ColumnarValue::Scalar(ScalarValue::Boolean(b)) = &predicate_output { + if let ColumnarValue::Scalar(ScalarValue::Boolean(b)) = + &evaluated.evaluated_result + { return match b { - Some(true) => Ok(ColumnarValue::Array(list_array)), + Some(true) => Ok(ColumnarValue::Array(evaluated.original_list)), _ => Ok(ColumnarValue::Array(empty_filtered_list( - &list_array, + &evaluated.original_list, field, )?)), }; } - let predicate = predicate_output.into_array(list_values.len())?; - let Some(predicate) = predicate.as_any().downcast_ref::() else { - return exec_err!( - "{} lambda must return boolean, got {}", - self.name(), - predicate.data_type() - ); - }; + let predicate = evaluated.boolean_predicate(self.name())?; // ListView and LargeListView are coerced to List/LargeList by coerce_value_types. - let filtered_list = match list_array.data_type() { + let filtered_list = match evaluated.original_list.data_type() { DataType::List(_) => { - let list = list_array.as_list::(); - let adjusted_offsets = adjust_offsets_for_slice(list); - let (filtered_values, new_offsets) = - filter_list_values(&list_values, predicate, &adjusted_offsets)?; + let (filtered_values, new_offsets) = filter_list_values( + &evaluated.flattened_values, + &predicate, + &evaluated.adjusted_offsets::(), + )?; Arc::new(ListArray::new( field, new_offsets, filtered_values, - list.nulls().cloned(), + evaluated.nulls().cloned(), )) as ArrayRef } DataType::LargeList(_) => { - let large_list = list_array.as_list::(); - let adjusted_offsets = adjust_offsets_for_slice(large_list); - let (filtered_values, new_offsets) = - filter_list_values(&list_values, predicate, &adjusted_offsets)?; + let (filtered_values, new_offsets) = filter_list_values( + &evaluated.flattened_values, + &predicate, + &evaluated.adjusted_offsets::(), + )?; Arc::new(LargeListArray::new( field, new_offsets, filtered_values, - large_list.nulls().cloned(), + evaluated.nulls().cloned(), )) } other => exec_err!("expected list, got {other}")?, @@ -252,13 +237,11 @@ fn filter_list_values( offsets: &OffsetBuffer, ) -> Result<(ArrayRef, OffsetBuffer)> { let num_sublists = offsets.len().saturating_sub(1); - let mut builder = OffsetBufferBuilder::::new(num_sublists); - let has_nulls = predicate.null_count() > 0; - for i in 0..num_sublists { + let new_offsets = OffsetBuffer::::from_lengths((0..num_sublists).map(|i| { let start = offsets[i].as_usize(); let end = offsets[i + 1].as_usize(); - let count = if has_nulls { + if has_nulls { (start..end) .filter(|&j| predicate.is_valid(j) && predicate.value(j)) .count() @@ -267,11 +250,8 @@ fn filter_list_values( .values() .slice(start, end - start) .count_set_bits() - }; - builder.push_length(count); - } - - let new_offsets = builder.finish(); + } + })); if new_offsets.last() == offsets.last() { return Ok((Arc::clone(values), offsets.clone())); @@ -289,9 +269,14 @@ mod tests { buffer::{NullBuffer, OffsetBuffer}, }; + use arrow::array::Int32Array; + use crate::array_filter::array_filter_higher_order_function; - use crate::lambda_utils::test_utils::{create_i32_list, eval_hof_on_i32_list, v}; - use datafusion_expr::lit; + use crate::lambda_utils::test_utils::{ + create_i32_large_list, create_i32_list, eval_hof_on_i32_list, + eval_hof_on_i32_list_with_outer, v, + }; + use datafusion_expr::{col, lit}; fn keep_greater_than_two( list: impl Array + Clone + 'static, @@ -461,4 +446,45 @@ mod tests { ); assert_eq!(actual, &expected); } + + #[test] + fn filter_large_list_parity() { + let list = create_i32_large_list( + vec![1, 2, 3, 4, 5], + OffsetBuffer::::from_lengths(vec![5]), + None, + ); + let res = keep_greater_than_two(list).unwrap(); + let actual = res.as_list::(); + let expected = create_i32_large_list( + vec![3, 4, 5], + OffsetBuffer::::from_lengths(vec![3]), + None, + ); + assert_eq!(actual, &expected); + } + + #[test] + fn filter_captured_outer_column() { + let list = create_i32_list( + vec![1, 50, 4, 50, 7, 50], + OffsetBuffer::::from_lengths(vec![2, 2, 2]), + None, + ); + let number = Int32Array::from(vec![10, 40, 60]); + let res = eval_hof_on_i32_list_with_outer( + array_filter_higher_order_function(), + list, + number, + v().gt(col("number")), + ) + .unwrap(); + let actual = res.as_list::(); + let expected = create_i32_list( + vec![50, 50], + OffsetBuffer::::from_lengths(vec![1, 1, 0]), + None, + ); + assert_eq!(actual, &expected); + } } diff --git a/datafusion/functions-nested/src/array_first.rs b/datafusion/functions-nested/src/array_first.rs new file mode 100644 index 0000000000000..615dc47394379 --- /dev/null +++ b/datafusion/functions-nested/src/array_first.rs @@ -0,0 +1,433 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`datafusion_expr::HigherOrderUDF`] definitions for array_first function. + +use arrow::{ + array::{Array, BooleanArray, UInt64Array, UInt64Builder}, + compute::take, + datatypes::{DataType, FieldRef}, +}; +use datafusion_common::{Result, exec_err, plan_err}; +use datafusion_expr::{ + ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, + Volatility, +}; +use datafusion_macros::user_doc; +use std::sync::Arc; + +use crate::lambda_utils::{ + EvaluatedListLambda, SingleListLambdaResult, coerce_single_list_arg, + evaluate_single_list_predicate, single_list_lambda_parameters, value_lambda_pair, +}; + +make_higher_order_function_expr_and_func!( + ArrayFirst, + array_first, + array lambda, + "returns the first element of an array that satisfies the predicate", + array_first_higher_order_function +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns the first element of an array that satisfies the given predicate. Returns null if the array is empty or no element matches. A predicate that returns null for an element is treated as not matching.", + syntax_example = "array_first(array, predicate)", + sql_example = r#"```sql +> select array_first([1, 2, 3, 4], x -> x > 2); ++----------------------------------------+ +| array_first([1,2,3,4],x -> x > 2) | ++----------------------------------------+ +| 3 | ++----------------------------------------+ +```"#, + argument( + name = "array", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ), + argument( + name = "predicate", + description = "Lambda predicate that returns a boolean. The first element for which it returns true is returned." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArrayFirst { + signature: HigherOrderSignature, + aliases: Vec, +} + +impl Default for ArrayFirst { + fn default() -> Self { + Self::new() + } +} + +impl ArrayFirst { + pub fn new() -> Self { + Self { + signature: HigherOrderSignature::exact( + vec![ValueOrLambda::Value(()), ValueOrLambda::Lambda(())], + Volatility::Immutable, + ), + aliases: vec![String::from("list_first")], + } + } +} + +impl HigherOrderUDFImpl for ArrayFirst { + fn name(&self) -> &str { + "array_first" + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn coerce_value_types(&self, arg_types: &[DataType]) -> Result> { + coerce_single_list_arg(self.name(), arg_types) + } + + fn lambda_parameters( + &self, + _step: usize, + fields: &[ValueOrLambda>], + ) -> Result { + single_list_lambda_parameters(self.name(), fields) + } + + fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result { + let (list, _lambda) = value_lambda_pair(self.name(), args.arg_fields)?; + + let element_field = match list.data_type() { + DataType::List(field) | DataType::LargeList(field) => field, + other => { + return plan_err!( + "{} expected a list as first argument, got {other}", + self.name() + ); + } + }; + + // The result is a single element of the array. It is always nullable + // because an empty array (or no matching element) yields null. + Ok(Arc::new( + element_field + .as_ref() + .clone() + .with_name("") + .with_nullable(true), + )) + } + + fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { + let evaluated = match evaluate_single_list_predicate(self.name(), &args)? { + SingleListLambdaResult::EarlyReturn(v) => return Ok(v), + SingleListLambdaResult::Ready(v) => v, + }; + + let predicate = evaluated.boolean_predicate(self.name())?; + let indices = match evaluated.original_list.data_type() { + DataType::List(_) | DataType::LargeList(_) => { + first_match_indices(&evaluated, &predicate) + } + other => return exec_err!("expected list, got {other}"), + }; + + let result = take(evaluated.flattened_values.as_ref(), &indices, None)?; + Ok(ColumnarValue::Array(result)) + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} + +/// Builds a `UInt64` index array (one entry per sublist) pointing at the first +/// element whose predicate is true, or null when no element matches. Indices are +/// absolute into the (sliced) flat values array, so `take` gathers the matches. +/// +/// A null predicate value is treated as not matching. The matched element itself +/// may be null and is still returned. +fn first_match_indices( + evaluated: &EvaluatedListLambda, + predicate: &BooleanArray, +) -> UInt64Array { + let mut builder = UInt64Builder::with_capacity(evaluated.len()); + + for i in 0..evaluated.len() { + let (start, end) = evaluated.row_range(i); + + match (start..end).find(|&j| predicate.is_valid(j) && predicate.value(j)) { + Some(j) => builder.append_value(j as u64), + None => builder.append_null(), + } + } + + builder.finish() +} + +#[cfg(test)] +mod tests { + use arrow::{ + array::{Array, AsArray, Int32Array, StringArray}, + buffer::{NullBuffer, OffsetBuffer}, + datatypes::Int32Type, + }; + + use crate::array_first::array_first_higher_order_function; + use crate::lambda_utils::test_utils::{ + create_i32_large_list, create_i32_list, eval_hof_on_i32_list, + eval_hof_on_i32_list_with_outer, v, + }; + use datafusion_common::Result; + use datafusion_expr::{col, lit}; + + fn first_greater_than_two( + list: impl Array + Clone + 'static, + ) -> Result { + eval_hof_on_i32_list(array_first_higher_order_function(), list, v().gt(lit(2i32))) + } + + // predicate: (100 / v) > 5; panics on divide by zero if v == 0 is evaluated + fn first_where_hundred_div_gt_five( + list: impl Array + Clone + 'static, + ) -> Result { + eval_hof_on_i32_list( + array_first_higher_order_function(), + list, + (lit(100i32) / v()).gt(lit(5i32)), + ) + } + + #[test] + fn test_first_basic() -> Result<()> { + let list = create_i32_list( + vec![1, 2, 3, 4, 5], + OffsetBuffer::::from_lengths(vec![5]), + None, + ); + let res = first_greater_than_two(list)?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![Some(3)]) + ); + Ok(()) + } + + #[test] + fn test_first_no_match_is_null() -> Result<()> { + let list = + create_i32_list(vec![1, 2], OffsetBuffer::::from_lengths(vec![2]), None); + let res = first_greater_than_two(list)?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![None]) + ); + Ok(()) + } + + #[test] + fn test_first_empty_array_is_null() -> Result<()> { + let list = create_i32_list( + Vec::::new(), + OffsetBuffer::::from_lengths(vec![0]), + None, + ); + let res = first_greater_than_two(list)?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![None]) + ); + Ok(()) + } + + #[test] + fn test_first_multiple_sublists() -> Result<()> { + // [1,5] -> 5, [2,4,3] -> 4, [1,2] -> null + let list = create_i32_list( + vec![1, 5, 2, 4, 3, 1, 2], + OffsetBuffer::::from_lengths(vec![2, 3, 2]), + None, + ); + let res = first_greater_than_two(list)?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![Some(5), Some(4), None]) + ); + Ok(()) + } + + #[test] + fn test_first_null_predicate_element_is_skipped() -> Result<()> { + // [1, NULL, 4] with v > 2: the NULL element's predicate is null and is + // skipped, so the first match is 4. + let list = create_i32_list( + Int32Array::from(vec![Some(1), None, Some(4)]), + OffsetBuffer::::from_lengths(vec![3]), + None, + ); + let res = first_greater_than_two(list)?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![Some(4)]) + ); + Ok(()) + } + + #[test] + fn test_first_matched_null_element_is_returned() -> Result<()> { + // [1, NULL, 3] with `v IS NULL`: the first match is the null element, + // which is returned as null. + let list = create_i32_list( + Int32Array::from(vec![Some(1), None, Some(3)]), + OffsetBuffer::::from_lengths(vec![3]), + None, + ); + let res = eval_hof_on_i32_list( + array_first_higher_order_function(), + list, + v().is_null(), + )?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![None]) + ); + Ok(()) + } + + // The 0 in the null row would divide by zero if the predicate were evaluated + // on it. The result for the null row must be null. + #[test] + fn test_first_does_not_evaluate_predicate_on_null_row_values() -> Result<()> { + let list = create_i32_list( + vec![1, 2, 0, 4, 5], + OffsetBuffer::::from_lengths(vec![3, 2]), + Some(NullBuffer::from(vec![false, true])), + ); + let res = first_where_hundred_div_gt_five(list)?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![None, Some(4)]) + ); + Ok(()) + } + + // The 0 before the slice offset would divide by zero if evaluated. + #[test] + fn test_first_does_not_evaluate_predicate_on_unreachable_values() -> Result<()> { + // sublists: [0], [4,5], [50,100]; slice away the first + let list = create_i32_list( + vec![0, 4, 5, 50, 100], + OffsetBuffer::::from_lengths(vec![1, 2, 2]), + None, + ) + .slice(1, 2); + let res = first_where_hundred_div_gt_five(list)?; + // [4,5]: 100/4=25>5 -> 4. [50,100]: 2>5 false, 1>5 false -> null + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![Some(4), None]) + ); + Ok(()) + } + + #[test] + fn test_first_eagerly_evaluates_predicate_after_match() { + // Although 4 is the first match, the predicate is evaluated for the + // later 0 in the same sublist and produces a division-by-zero error. + let list = + create_i32_list(vec![4, 0], OffsetBuffer::::from_lengths(vec![2]), None); + + let err = first_where_hundred_div_gt_five(list).unwrap_err(); + assert!( + err.to_string().contains("Divide by zero"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_first_large_list_parity() -> Result<()> { + let list = create_i32_large_list( + vec![1, 2, 3, 4, 5], + OffsetBuffer::::from_lengths(vec![5]), + None, + ); + let res = first_greater_than_two(list)?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![Some(3)]) + ); + Ok(()) + } + + #[test] + fn test_first_captured_outer_column() -> Result<()> { + let list = create_i32_list( + vec![1, 50, 4, 50, 7, 50], + OffsetBuffer::::from_lengths(vec![2, 2, 2]), + None, + ); + let number = Int32Array::from(vec![10, 40, 60]); + let res = eval_hof_on_i32_list_with_outer( + array_first_higher_order_function(), + list, + number, + v().gt(col("number")), + )?; + assert_eq!( + res.as_primitive::(), + &Int32Array::from(vec![Some(50), Some(50), None]) + ); + Ok(()) + } + + #[test] + fn test_first_string_elements() -> Result<()> { + use arrow::array::ListArray; + use arrow::datatypes::{DataType, Field}; + use datafusion_expr::Expr; + use datafusion_expr::expr::LambdaVariable; + use std::sync::Arc; + + // ['a', 'bb', 'ccc'] with v > 'a' -> 'bb' (exercises take on a non-primitive type) + let values = StringArray::from(vec!["a", "bb", "ccc"]); + let list = ListArray::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + OffsetBuffer::::from_lengths(vec![3]), + Arc::new(values), + None, + ); + + let x = Expr::LambdaVariable(LambdaVariable::new( + "v".to_string(), + Some(Arc::new(Field::new("v", DataType::Utf8, true))), + )); + let body = x.gt(lit("a")); + + let res = eval_hof_on_i32_list(array_first_higher_order_function(), list, body)?; + assert_eq!(res.as_string::(), &StringArray::from(vec![Some("bb")])); + Ok(()) + } +} diff --git a/datafusion/functions-nested/src/array_has.rs b/datafusion/functions-nested/src/array_has.rs index 04818258f040b..0f680469f6023 100644 --- a/datafusion/functions-nested/src/array_has.rs +++ b/datafusion/functions-nested/src/array_has.rs @@ -18,11 +18,13 @@ //! [`ScalarUDFImpl`] definitions for array_has, array_has_all and array_has_any functions. use arrow::array::{ - Array, ArrayRef, AsArray, BooleanArray, BooleanBufferBuilder, Datum, Scalar, - StringArrayType, + Array, ArrayRef, ArrowNativeTypeOp, ArrowPrimitiveType, AsArray, BooleanArray, + BooleanBufferBuilder, Datum, MAX_INLINE_VIEW_LEN, PrimitiveArray, Scalar, + StringArrayType, StringViewArray, }; -use arrow::buffer::{BooleanBuffer, NullBuffer}; +use arrow::buffer::{BooleanBuffer, NullBuffer, OffsetBuffer}; use arrow::datatypes::DataType; +use arrow::downcast_primitive_array; use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::cast::{as_fixed_size_list_array, as_generic_list_array}; use datafusion_common::utils::string_utils::string_array_to_vec; @@ -323,11 +325,85 @@ impl<'a> ArrayWrapper<'a> { } } +/// Evaluate `array_has` with an array (per-row) needle. +/// +/// Primitive and string element types take a per-type fast path; nested (and any +/// other) element types fall back to the per-row `eq` kernel, which allocates a +/// `BooleanArray` per row. fn array_has_dispatch_for_array<'a>( haystack: ArrayWrapper<'a>, needle: &ArrayRef, ) -> Result { let combined_nulls = NullBuffer::union(haystack.nulls(), needle.nulls()); + let needle = needle.as_ref(); + + // Rebase offsets to 0 with `OffsetBuffer::subtract` so `offsets[i]` indexes + // `visible_values` directly (the haystack may be a sliced list). + let raw = OffsetBuffer::new( + haystack + .offsets() + .map(|o| o as i64) + .collect::>() + .into(), + ); + let first = raw[0]; + let visible_values = haystack + .values() + .slice(first as usize, (raw[raw.len() - 1] - first) as usize); + let visible_values = visible_values.as_ref(); + let offsets: Vec = raw.subtract(first).iter().map(|&o| o as usize).collect(); + + // Fast path for primitive/string elements whose (coerced) type matches the + // needle; a type mismatch or a nested type falls through to the per-row kernel. + let fast_path = if visible_values.data_type() != needle.data_type() { + None + } else { + downcast_primitive_array! { + visible_values => { + // The element-null path makes several passes over the values, so + // past a large average list length the per-row `eq` kernel is + // faster -- bail to it. The single-pass all-valid path has no such + // crossover, so only bail when elements are null. + let num_rows = offsets.len() - 1; + if num_rows > 0 + && offsets[num_rows] / num_rows > NULL_FAST_PATH_MAX_LEN + && visible_values.null_count() > 0 + { + None + } else { + Some(array_has_array_primitive( + visible_values, needle, &offsets, + combined_nulls.as_ref(), + )) + } + }, + DataType::Utf8 => Some(array_has_array_string( + visible_values.as_string::(), + needle.as_string::(), + &offsets, + combined_nulls.as_ref(), + )), + DataType::LargeUtf8 => Some(array_has_array_string( + visible_values.as_string::(), + needle.as_string::(), + &offsets, + combined_nulls.as_ref(), + )), + DataType::Utf8View => Some(array_has_array_string_view( + visible_values.as_string_view(), + needle.as_string_view(), + &offsets, + combined_nulls.as_ref(), + )), + _ => None, + } + }; + + if let Some(values) = fast_path { + return Ok(Arc::new(BooleanArray::new(values, combined_nulls))); + } + + // Fallback: per-row `eq` kernel (nested element types, or a type mismatch). let mut result = BooleanBufferBuilder::new(haystack.len()); for (i, arr) in haystack.iter().enumerate() { if combined_nulls.as_ref().is_some_and(|n| n.is_null(i)) { @@ -344,6 +420,146 @@ fn array_has_dispatch_for_array<'a>( Ok(Arc::new(BooleanArray::new(result.finish(), combined_nulls))) } +/// Average list length past which the element-null path loses to the per-row +/// `eq` kernel and bails to it (empirically measured). +const NULL_FAST_PATH_MAX_LEN: usize = 512; + +/// Primitive fast path, two branches on element validity: +/// +/// 1. No nulls: branchless OR-reduction over the raw slice (auto-vectorizes). +/// 2. Nulls: AND the equality bitmap with validity (a null slot's value is +/// arbitrary), then reduce each row to "any bit set". Chunked to bound the +/// expanded needle. +fn array_has_array_primitive( + values: &PrimitiveArray, + needle: &dyn Array, + offsets: &[usize], + combined_nulls: Option<&NullBuffer>, +) -> BooleanBuffer +where + T::Native: ArrowNativeTypeOp, +{ + let needle = needle.as_primitive::(); + let num_rows = offsets.len() - 1; + let value_slice = values.values(); + let needle_slice = needle.values(); + + let Some(element_nulls) = values.nulls() else { + return BooleanBuffer::collect_bool(num_rows, |i| { + if combined_nulls.is_some_and(|n| n.is_null(i)) { + return false; + } + // `needle[i]` is non-null here: combined_nulls covers the needle nulls. + let needle_val = needle_slice[i]; + let start = offsets[i]; + let end = offsets[i + 1]; + value_slice[start..end] + .iter() + .fold(false, |acc, &v| acc | v.is_eq(needle_val)) + }); + }; + + // Case 2 (see fn doc), chunked like the all/any kernels. + let mut result = BooleanBufferBuilder::new(num_rows); + let mut needle_expanded: Vec = Vec::new(); + for chunk_start in (0..num_rows).step_by(ROW_CONVERSION_CHUNK_SIZE) { + let chunk_end = (chunk_start + ROW_CONVERSION_CHUNK_SIZE).min(num_rows); + let elem_start = offsets[chunk_start]; + let elem_end = offsets[chunk_end]; + + // Expand the per-row needle across this chunk's elements (reused scratch), + // then compare in one vectorizable pass and mask out null elements. + needle_expanded.clear(); + for i in chunk_start..chunk_end { + needle_expanded.extend(std::iter::repeat_n( + needle_slice[i], + offsets[i + 1] - offsets[i], + )); + } + let chunk_values = &value_slice[elem_start..elem_end]; + let eq_bits = BooleanBuffer::collect_bool(chunk_values.len(), |k| { + chunk_values[k].is_eq(needle_expanded[k]) + }); + let matched = &eq_bits + & &element_nulls + .inner() + .slice(elem_start, elem_end - elem_start); + + for i in chunk_start..chunk_end { + if combined_nulls.is_some_and(|n| n.is_null(i)) { + result.append(false); + continue; + } + let start = offsets[i] - elem_start; + let end = offsets[i + 1] - elem_start; + result.append(matched.slice(start, end - start).has_true()); + } + } + result.finish() +} + +/// String fast path, generic over the offset width (`Utf8` / `LargeUtf8`). +fn array_has_array_string<'a, S: StringArrayType<'a> + Copy>( + values: S, + needle: S, + offsets: &[usize], + combined_nulls: Option<&NullBuffer>, +) -> BooleanBuffer { + let num_rows = offsets.len() - 1; + BooleanBuffer::collect_bool(num_rows, |i| { + if combined_nulls.is_some_and(|n| n.is_null(i)) { + return false; + } + // `needle[i]` is non-null here: combined_nulls covers the needle nulls. + let needle_val = needle.value(i); + let start = offsets[i]; + let end = offsets[i + 1]; + // Compare the value first and only consult validity on a match (see the + // primitive path for why this is correct and faster on no-match scans). + (start..end).any(|k| values.value(k) == needle_val && !values.is_null(k)) + }) +} + +/// `Utf8View` variant of [`array_has_array_string`]: compare the packed 128-bit +/// views directly so the length + 4-byte prefix reject non-matches without +/// touching the data buffer, and an inline value matches on the view alone. A +/// longer view is only materialized to confirm a candidate; validity is +/// consulted only on a view match. +fn array_has_array_string_view( + values: &StringViewArray, + needle: &StringViewArray, + offsets: &[usize], + combined_nulls: Option<&NullBuffer>, +) -> BooleanBuffer { + let num_rows = offsets.len() - 1; + let value_views = values.views(); + let needle_views = needle.views(); + BooleanBuffer::collect_bool(num_rows, |i| { + if combined_nulls.is_some_and(|n| n.is_null(i)) { + return false; + } + // `needle[i]` is non-null here: combined_nulls covers the needle nulls. + let needle_view = needle_views[i]; + // Low 32 bits are the byte length; the next 32 are the inline prefix. + let needle_inline = (needle_view as u32) <= MAX_INLINE_VIEW_LEN; + let needle_lo = needle_view as u64; + let needle_val = needle.value(i); + let start = offsets[i]; + let end = offsets[i + 1]; + (start..end).any(|k| { + let v = value_views[k]; + let matched = if needle_inline { + // Inline: the whole view is the canonical value (zero padded). + v == needle_view + } else { + // Longer: reject on length + prefix, then confirm the bytes. + (v as u64) == needle_lo && values.value(k) == needle_val + }; + matched && !values.is_null(k) + }) + }) +} + fn array_has_dispatch_for_scalar( haystack: ArrayWrapper<'_>, needle: &dyn Datum, @@ -778,22 +994,22 @@ fn array_has_any_with_scalar_general( #[user_doc( doc_section(label = "Array Functions"), - description = "Returns true if all elements of sub-array exist in array.", - syntax_example = "array_has_all(array, sub-array)", + description = "Returns true if all elements of sub_array exist in array.", + syntax_example = "array_has_all(array, sub_array)", sql_example = r#"```sql > select array_has_all([1, 2, 3, 4], [2, 3]); -+--------------------------------------------+ ++---------------------------------------------+ | array_has_all(List([1,2,3,4]), List([2,3])) | -+--------------------------------------------+ -| true | -+--------------------------------------------+ ++---------------------------------------------+ +| true | ++---------------------------------------------+ ```"#, argument( name = "array", description = "Array expression. Can be a constant, column, or function, and any combination of array operators." ), argument( - name = "sub-array", + name = "sub_array", description = "Array expression. Can be a constant, column, or function, and any combination of array operators." ) )] @@ -1311,4 +1527,63 @@ mod tests { &[Some(true), Some(true)], ); } + + /// Invoke `array_has` with the needle as an array (a column with one value + /// per row). This exercises `array_has_dispatch_for_array` and its fast path. + fn invoke_array_has_array(haystack: ArrayRef, needle: ArrayRef) -> ArrayRef { + let num_rows = haystack.len(); + let haystack_type = haystack.data_type().clone(); + let needle_type = needle.data_type().clone(); + ArrayHas::new() + .invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(haystack), ColumnarValue::Array(needle)], + arg_fields: vec![ + Arc::new(Field::new("haystack", haystack_type, false)), + Arc::new(Field::new("needle", needle_type, false)), + ], + number_rows: num_rows, + return_field: Arc::new(Field::new("return", DataType::Boolean, true)), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() + .into_array(num_rows) + .unwrap() + } + + #[test] + fn test_array_has_array_needle_sliced() { + // Offset normalization for sliced haystacks must keep the element ranges + // and the needle column aligned, for both `List` (offsets from the + // buffer) and `FixedSizeList` (offsets computed as `i * value_length`). + // Slicing is an execution artifact SQL/SLT can't force, so this stays a + // unit test; value-level behavior is covered by `array/array_has.slt`. + let full = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(10), Some(20), Some(30)]), // needle 20 -> true + Some(vec![Some(40)]), // needle 41 -> false + Some(vec![Some(50), Some(60)]), // needle 60 -> true + Some(vec![Some(70)]), + ]); + let sliced_haystack: ArrayRef = Arc::new(full.slice(1, 3)); + let sliced_needle: ArrayRef = + Arc::new(Int32Array::from(vec![999, 20, 41, 60, 999]).slice(1, 3)); + let result = invoke_array_has_array(sliced_haystack, sliced_needle); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false), Some(true)] + ); + + // Sliced FixedSizeList (width 2; rows 1..=2 of + // [[1,2],[11,12],[21,22],[31,32]] visible) with an aligned needle column. + let field = Arc::new(Field::new("item", DataType::Int32, true)); + let fsl_values = Arc::new(Int32Array::from(vec![1, 2, 11, 12, 21, 22, 31, 32])); + let fsl: ArrayRef = + Arc::new(FixedSizeListArray::new(field, 2, fsl_values, None).slice(1, 2)); + let needle: ArrayRef = Arc::new(Int32Array::from(vec![11, 99])); + let result = invoke_array_has_array(fsl, needle); + assert_eq!( + result.as_boolean().iter().collect::>(), + vec![Some(true), Some(false)] + ); + } } diff --git a/datafusion/functions-nested/src/array_normalize.rs b/datafusion/functions-nested/src/array_normalize.rs index 0ff7674032d7f..f7da07e5f6e69 100644 --- a/datafusion/functions-nested/src/array_normalize.rs +++ b/datafusion/functions-nested/src/array_normalize.rs @@ -19,9 +19,9 @@ use crate::utils::make_scalar_function; use arrow::array::{ - Array, ArrayRef, Float64Array, GenericListArray, NullBufferBuilder, - OffsetBufferBuilder, OffsetSizeTrait, + Array, ArrayRef, Float64Array, GenericListArray, NullBufferBuilder, OffsetSizeTrait, }; +use arrow::buffer::OffsetBuffer; use arrow::datatypes::{ DataType, DataType::{FixedSizeList, LargeList, List, Null}, @@ -144,13 +144,14 @@ fn general_array_normalize(arrays: &[ArrayRef]) -> Result = Vec::with_capacity(values.len()); - let mut new_offsets = OffsetBufferBuilder::::new(list_array.len()); + let mut new_offsets = Vec::::with_capacity(list_array.len() + 1); + new_offsets.push(O::zero()); let mut nulls = NullBufferBuilder::new(list_array.len()); for row in 0..list_array.len() { if list_array.is_null(row) { nulls.append_null(); - new_offsets.push_length(0); + new_offsets.push(new_offsets[row]); continue; } @@ -161,7 +162,7 @@ fn general_array_normalize(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result(arrays: &[ArrayRef]) -> Result::try_new( field, - new_offsets.finish(), + OffsetBuffer::new(new_offsets.into()), values_array, nulls.finish(), )?)) diff --git a/datafusion/functions-nested/src/array_product.rs b/datafusion/functions-nested/src/array_product.rs new file mode 100644 index 0000000000000..a5cef43142fa0 --- /dev/null +++ b/datafusion/functions-nested/src/array_product.rs @@ -0,0 +1,174 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ScalarUDFImpl`] definitions for array_product function. + +use crate::utils::make_scalar_function; +use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait}; +use arrow::datatypes::{ + DataType, + DataType::{FixedSizeList, LargeList, List, Null}, + Field, +}; +use datafusion_common::cast::{as_float64_array, as_generic_list_array}; +use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only}; +use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args}; +use datafusion_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_macros::user_doc; +use std::sync::Arc; + +make_udf_expr_and_func!( + ArrayProduct, + array_product, + array, + "returns the product of the elements of a numeric array.", + array_product_udf +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns the product of the elements in the input numeric array. \ + NULL elements inside the array are skipped (matching SQL aggregate \ + convention). Returns NULL if the input is NULL, every element is \ + NULL, or the array is empty. The result is always returned as \ + `Float64`.", + syntax_example = "array_product(array)", + sql_example = r#"```sql +> select array_product([1.0, 2.0, 3.0]); ++------------------------------------+ +| array_product(List([1.0,2.0,3.0])) | ++------------------------------------+ +| 6.0 | ++------------------------------------+ +```"#, + argument( + name = "array", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArrayProduct { + signature: Signature, + aliases: Vec, +} + +impl Default for ArrayProduct { + fn default() -> Self { + Self::new() + } +} + +impl ArrayProduct { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + aliases: vec!["list_product".to_string()], + } + } +} + +impl ScalarUDFImpl for ArrayProduct { + fn name(&self) -> &str { + "array_product" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [arg_type] = take_function_args(self.name(), arg_types)?; + let coercion = Some(&ListCoercion::FixedSizedListToList); + + if !matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) { + return plan_err!("{} does not support type {arg_type}", self.name()); + } + + let coerced = if matches!(arg_type, Null) { + List(Arc::new(Field::new_list_field(DataType::Float64, true))) + } else { + coerced_type_with_base_type_only(arg_type, &DataType::Float64, coercion) + }; + + Ok(vec![coerced]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(array_product_inner)(&args.args) + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} + +fn array_product_inner(args: &[ArrayRef]) -> Result { + let [array] = take_function_args("array_product", args)?; + match array.data_type() { + List(_) => general_array_product::(args), + LargeList(_) => general_array_product::(args), + arg_type => internal_err!( + "array_product received unexpected type after coercion: {arg_type}" + ), + } +} + +fn general_array_product(arrays: &[ArrayRef]) -> Result { + let list_array = as_generic_list_array::(&arrays[0])?; + let values = as_float64_array(list_array.values())?; + let offsets = list_array.value_offsets(); + + let mut builder = Float64Array::builder(list_array.len()); + + for row in 0..list_array.len() { + if list_array.is_null(row) { + builder.append_null(); + continue; + } + + let start = offsets[row].as_usize(); + let end = offsets[row + 1].as_usize(); + + let mut prod = 1.0_f64; + let mut any_valid = false; + for i in start..end { + if values.is_valid(i) { + prod *= values.value(i); + any_valid = true; + } + } + + if any_valid { + builder.append_value(prod); + } else { + builder.append_null(); + } + } + + Ok(Arc::new(builder.finish())) +} diff --git a/datafusion/functions-nested/src/array_scale.rs b/datafusion/functions-nested/src/array_scale.rs new file mode 100644 index 0000000000000..a6be910d20b12 --- /dev/null +++ b/datafusion/functions-nested/src/array_scale.rs @@ -0,0 +1,219 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ScalarUDFImpl`] definitions for array_scale function. + +use crate::utils::make_scalar_function; +use arrow::array::{Array, ArrayRef, Float64Array, GenericListArray, OffsetSizeTrait}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::datatypes::{ + DataType, + DataType::{FixedSizeList, LargeList, List, Null}, + Field, +}; +use datafusion_common::cast::{as_float64_array, as_generic_list_array}; +use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only}; +use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args}; +use datafusion_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_macros::user_doc; +use std::sync::Arc; + +make_udf_expr_and_func!( + ArrayScale, + array_scale, + array scalar, + "scales each element of a numeric array by a scalar.", + array_scale_udf +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns a new array with each element of the input array multiplied by a scalar value, computed as `array[i] * scalar`. Returns NULL if the input row is NULL or the scalar is NULL. If a NULL element appears in the input array at position `i`, the result element at position `i` is NULL. Returns an empty array for an empty input array.", + syntax_example = "array_scale(array, scalar)", + sql_example = r#"```sql +> select array_scale([1.0, 2.0, 3.0], 2.0); ++----------------------------------+ +| array_scale(List([1.0,2.0,3.0]),Float64(2.0)) | ++----------------------------------+ +| [2.0, 4.0, 6.0] | ++----------------------------------+ +```"#, + argument( + name = "array", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ), + argument( + name = "scalar", + description = "Numeric scalar to multiply each element by. Can be a constant or column expression." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArrayScale { + signature: Signature, + aliases: Vec, +} + +impl Default for ArrayScale { + fn default() -> Self { + Self::new() + } +} + +impl ArrayScale { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + aliases: vec!["list_scale".to_string()], + } + } +} + +impl ScalarUDFImpl for ArrayScale { + fn name(&self) -> &str { + "array_scale" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + // After `coerce_types`, `arg_types[0]` is one of List(Float64) or LargeList(Float64). + Ok(arg_types[0].clone()) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [array_type, scalar_type] = take_function_args(self.name(), arg_types)?; + let coercion = Some(&ListCoercion::FixedSizedListToList); + + if !matches!( + array_type, + Null | List(_) | LargeList(_) | FixedSizeList(..) + ) { + return plan_err!( + "{} first argument must be a list type, got {array_type}", + self.name() + ); + } + + if !scalar_type.is_numeric() && !matches!(scalar_type, Null) { + return plan_err!( + "{} second argument must be numeric, got {scalar_type}", + self.name() + ); + } + + let coerced_array = if matches!(array_type, Null) { + List(Arc::new(Field::new_list_field(DataType::Float64, true))) + } else { + coerced_type_with_base_type_only(array_type, &DataType::Float64, coercion) + }; + + Ok(vec![coerced_array, DataType::Float64]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(array_scale_inner)(&args.args) + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} + +fn array_scale_inner(args: &[ArrayRef]) -> Result { + let [array, scalar] = take_function_args("array_scale", args)?; + match array.data_type() { + List(_) => general_array_scale::(array, scalar), + LargeList(_) => general_array_scale::(array, scalar), + arg_type => internal_err!( + "array_scale received unexpected type after coercion: {arg_type}" + ), + } +} + +fn general_array_scale( + array: &ArrayRef, + scalar: &ArrayRef, +) -> Result { + let list_array = as_generic_list_array::(array)?; + let scalar_array = as_float64_array(scalar)?; + + let values = as_float64_array(list_array.values())?; + let offsets = list_array.value_offsets(); + + // A row is null whenever either input row is null. The scalar applies + // uniformly across the array, so a null scalar makes the whole row + // undefined; union the two row-level null buffers in a single pass + // rather than tracking row nulls inside the value loop. + let row_nulls = NullBuffer::union(list_array.nulls(), scalar_array.nulls()); + + let mut value_builder = Float64Array::builder(values.len()); + let mut new_offsets = Vec::::with_capacity(list_array.len() + 1); + new_offsets.push(O::zero()); + + for row in 0..list_array.len() { + if row_nulls.as_ref().is_some_and(|nb| nb.is_null(row)) { + new_offsets.push(new_offsets[row]); + continue; + } + + let start = offsets[row].as_usize(); + let end = offsets[row + 1].as_usize(); + let len = end - start; + let scalar_val = scalar_array.value(row); + + let slice = values.slice(start, len); + + // Per-element NULL propagation for NULL elements inside the array. + for i in 0..len { + if slice.is_null(i) { + value_builder.append_null(); + } else { + value_builder.append_value(slice.value(i) * scalar_val); + } + } + + new_offsets.push(new_offsets[row] + O::usize_as(len)); + } + + let values_array = Arc::new(value_builder.finish()); + + // Preserve the inner field from the input array (including any user + // metadata). After `coerce_types` the inner type is Float64, but the + // input may still carry field-level annotations worth keeping. + let field = match list_array.data_type() { + List(f) | LargeList(f) => Arc::clone(f), + other => { + return internal_err!("array_scale unexpected list type: {other}"); + } + }; + + Ok(Arc::new(GenericListArray::::try_new( + field, + OffsetBuffer::new(new_offsets.into()), + values_array, + row_nulls, + )?)) +} diff --git a/datafusion/functions-nested/src/array_subtract.rs b/datafusion/functions-nested/src/array_subtract.rs new file mode 100644 index 0000000000000..24600da04f74e --- /dev/null +++ b/datafusion/functions-nested/src/array_subtract.rs @@ -0,0 +1,130 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ScalarUDFImpl`] definitions for array_subtract function. + +use crate::utils::{ + array_math_binary_op, coerce_array_math_arg_types, make_scalar_function, +}; +use arrow::array::ArrayRef; +use arrow::datatypes::{ + DataType, + DataType::{LargeList, List}, +}; +use datafusion_common::{Result, exec_err, utils::take_function_args}; +use datafusion_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_macros::user_doc; + +make_udf_expr_and_func!( + ArraySubtract, + array_subtract, + array1 array2, + "returns the element-wise difference of two numeric arrays.", + array_subtract_udf +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns the element-wise difference of two numeric arrays of equal length, computed as `array1[i] - array2[i]` per position. NULL is propagated per element: if either input element at position `i` is NULL, the corresponding output element is NULL (positions are preserved). Returns NULL if either entire input array is NULL. Errors if the per-row lengths differ. Returns an empty array if both inputs are empty.", + syntax_example = "array_subtract(array1, array2)", + sql_example = r#"```sql +> select array_subtract([10.0, 20.0, 30.0], [1.0, 2.0, 3.0]); ++--------------------------------------------------------------+ +| array_subtract(List([10.0,20.0,30.0]),List([1.0,2.0,3.0])) | ++--------------------------------------------------------------+ +| [9.0, 18.0, 27.0] | ++--------------------------------------------------------------+ +```"#, + argument( + name = "array1", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ), + argument( + name = "array2", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArraySubtract { + signature: Signature, + aliases: Vec, +} + +impl Default for ArraySubtract { + fn default() -> Self { + Self::new() + } +} + +impl ArraySubtract { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + aliases: vec!["list_subtract".to_string()], + } + } +} + +impl ScalarUDFImpl for ArraySubtract { + fn name(&self) -> &str { + "array_subtract" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [_, _] = take_function_args(self.name(), arg_types)?; + coerce_array_math_arg_types(self.name(), arg_types) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(array_subtract_inner)(&args.args) + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} + +fn array_subtract_inner(args: &[ArrayRef]) -> Result { + let [array1, array2] = take_function_args("array_subtract", args)?; + let sub = |a: f64, b: f64| a - b; + match (array1.data_type(), array2.data_type()) { + (List(_), List(_)) => { + array_math_binary_op::("array_subtract", array1, array2, sub) + } + (LargeList(_), LargeList(_)) => { + array_math_binary_op::("array_subtract", array1, array2, sub) + } + (arg_type1, arg_type2) => exec_err!( + "array_subtract received unexpected types after coercion: {arg_type1} and {arg_type2}" + ), + } +} diff --git a/datafusion/functions-nested/src/array_sum.rs b/datafusion/functions-nested/src/array_sum.rs new file mode 100644 index 0000000000000..d115355f5cbb9 --- /dev/null +++ b/datafusion/functions-nested/src/array_sum.rs @@ -0,0 +1,174 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`ScalarUDFImpl`] definitions for array_sum function. + +use crate::utils::make_scalar_function; +use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait}; +use arrow::datatypes::{ + DataType, + DataType::{FixedSizeList, LargeList, List, Null}, + Field, +}; +use datafusion_common::cast::{as_float64_array, as_generic_list_array}; +use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only}; +use datafusion_common::{Result, internal_err, plan_err, utils::take_function_args}; +use datafusion_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_macros::user_doc; +use std::sync::Arc; + +make_udf_expr_and_func!( + ArraySum, + array_sum, + array, + "returns the sum of elements in a numeric array.", + array_sum_udf +); + +#[user_doc( + doc_section(label = "Array Functions"), + description = "Returns the sum of the elements of the input array, computed as `array[0] + array[1] + ...`. NULL elements are skipped (per SQL aggregate convention). Returns NULL if the input row is NULL, every element is NULL, or the array is empty.", + syntax_example = "array_sum(array)", + sql_example = r#"```sql +> select array_sum([1.0, 2.0, 3.0]); ++----------------------------+ +| array_sum(List([1.0,2.0,3.0])) | ++----------------------------+ +| 6.0 | ++----------------------------+ +```"#, + argument( + name = "array", + description = "Array expression. Can be a constant, column, or function, and any combination of array operators." + ) +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct ArraySum { + signature: Signature, + aliases: Vec, +} + +impl Default for ArraySum { + fn default() -> Self { + Self::new() + } +} + +impl ArraySum { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + aliases: vec!["list_sum".to_string()], + } + } +} + +impl ScalarUDFImpl for ArraySum { + fn name(&self) -> &str { + "array_sum" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [arg_type] = take_function_args(self.name(), arg_types)?; + let coercion = Some(&ListCoercion::FixedSizedListToList); + + if !matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) { + return plan_err!("{} does not support type {arg_type}", self.name()); + } + + let coerced = if matches!(arg_type, Null) { + List(Arc::new(Field::new_list_field(DataType::Float64, true))) + } else { + coerced_type_with_base_type_only(arg_type, &DataType::Float64, coercion) + }; + + Ok(vec![coerced]) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(array_sum_inner)(&args.args) + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} + +fn array_sum_inner(args: &[ArrayRef]) -> Result { + let [array] = take_function_args("array_sum", args)?; + match array.data_type() { + List(_) => general_array_sum::(array), + LargeList(_) => general_array_sum::(array), + arg_type => { + internal_err!("array_sum received unexpected type after coercion: {arg_type}") + } + } +} + +fn general_array_sum(array: &ArrayRef) -> Result { + let list_array = as_generic_list_array::(array)?; + let values = as_float64_array(list_array.values())?; + let offsets = list_array.value_offsets(); + + let mut builder = Float64Array::builder(list_array.len()); + + for row in 0..list_array.len() { + if list_array.is_null(row) { + builder.append_null(); + continue; + } + + let start = offsets[row].as_usize(); + let end = offsets[row + 1].as_usize(); + + // Skip NULL elements per SQL aggregate convention (matches PostgreSQL + // array_sum, DuckDB list_sum, Spark aggregate). Empty arrays and + // all-NULL arrays both yield NULL — same behavior as SQL SUM over + // an empty set or all-NULL column. + let mut sum = 0.0_f64; + let mut any_valid = false; + for i in start..end { + if values.is_valid(i) { + sum += values.value(i); + any_valid = true; + } + } + + if any_valid { + builder.append_value(sum); + } else { + builder.append_null(); + } + } + + Ok(Arc::new(builder.finish())) +} diff --git a/datafusion/functions-nested/src/array_transform.rs b/datafusion/functions-nested/src/array_transform.rs index a0415749f45e2..e07952722ec0d 100644 --- a/datafusion/functions-nested/src/array_transform.rs +++ b/datafusion/functions-nested/src/array_transform.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! [`HigherOrderUDF`] definitions for array_transform function. +//! [`datafusion_expr::HigherOrderUDF`] definitions for array_transform function. use arrow::{ array::{Array, ArrayRef, AsArray, LargeListArray, ListArray}, @@ -28,7 +28,7 @@ use datafusion_common::{ }; use datafusion_expr::{ ColumnarValue, Documentation, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, - HigherOrderSignature, HigherOrderUDF, LambdaParametersProgress, ValueOrLambda, + HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, ValueOrLambda, Volatility, }; use datafusion_macros::user_doc; @@ -50,7 +50,7 @@ make_higher_order_function_expr_and_func!( #[user_doc( doc_section(label = "Array Functions"), description = "transforms the values of an array", - syntax_example = "array_transform(array, x -> x*2)", + syntax_example = "array_transform(array, lambda)", sql_example = r#"```sql > select array_transform([1, 2, 3, 4, 5], x -> x*2); +-------------------------------------------+ @@ -63,7 +63,10 @@ make_higher_order_function_expr_and_func!( name = "array", description = "Array expression. Can be a constant, column, or function, and any combination of array operators." ), - argument(name = "lambda", description = "Lambda") + argument( + name = "lambda", + description = "The lambda function used to transform each value of the array." + ) )] #[derive(Debug, PartialEq, Eq, Hash)] pub struct ArrayTransform { @@ -89,7 +92,7 @@ impl ArrayTransform { } } -impl HigherOrderUDF for ArrayTransform { +impl HigherOrderUDFImpl for ArrayTransform { fn name(&self) -> &str { "array_transform" } diff --git a/datafusion/functions-nested/src/arrays_zip.rs b/datafusion/functions-nested/src/arrays_zip.rs index 5f1cb9dedf408..c574821707cc0 100644 --- a/datafusion/functions-nested/src/arrays_zip.rs +++ b/datafusion/functions-nested/src/arrays_zip.rs @@ -22,7 +22,7 @@ use arrow::array::{ Array, ArrayRef, Capacities, ListArray, MutableArrayData, NullBufferBuilder, StructArray, new_null_array, }; -use arrow::buffer::OffsetBuffer; +use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::DataType::{FixedSizeList, LargeList, List, Null}; use arrow::datatypes::{DataType, Field, Fields}; use datafusion_common::cast::{ @@ -44,7 +44,7 @@ struct ListColumnView { /// Pre-computed per-row start offsets (length = num_rows + 1). offsets: Vec, /// Null bitmap from the input array (None means no nulls). - nulls: Option, + nulls: Option, } impl ListColumnView { @@ -130,7 +130,7 @@ impl ScalarUDFImpl for ArraysZip { return exec_err!("arrays_zip expects array arguments, got {dt}"); } }; - fields.push(Field::new(format!("{}", i + 1), element_type, true)); + fields.push(Field::new(arrays_zip_field_name(i), element_type, true)); } Ok(List(Arc::new(Field::new_list_field( @@ -163,8 +163,13 @@ fn arrays_zip_inner(args: &[ArrayRef]) -> Result { return exec_err!("arrays_zip requires at least one argument"); } + let field_names = arrays_zip_field_names(args.len()); let num_rows = args[0].len(); + if let Some(result) = try_perfect_list_zip(args, &field_names)? { + return Ok(result); + } + // Build a type-erased ListColumnView for each argument. // None means the argument is Null-typed (all nulls, no backing data). let mut views: Vec> = Vec::with_capacity(args.len()); @@ -225,8 +230,8 @@ fn arrays_zip_inner(args: &[ArrayRef]) -> Result { let struct_fields: Fields = element_types .iter() - .enumerate() - .map(|(i, dt)| Field::new(format!("{}", i + 1), dt.clone(), true)) + .zip(field_names.iter()) + .map(|(dt, name)| Field::new(name.clone(), dt.clone(), true)) .collect::>() .into(); @@ -275,15 +280,15 @@ fn arrays_zip_inner(args: &[ArrayRef]) -> Result { let end = v.offsets[row_idx + 1]; let len = end - start; let builder = builders[col_idx].as_mut().unwrap(); - builder.extend(0, start, end); + builder.try_extend(0, start, end)?; if len < max_len { - builder.extend_nulls(max_len - len); + builder.try_extend_nulls(max_len - len)?; } } _ => { // Null list entry or None (Null-typed) arg — all nulls. if let Some(builder) = builders[col_idx].as_mut() { - builder.extend_nulls(max_len); + builder.try_extend_nulls(max_len)?; } } } @@ -327,3 +332,282 @@ fn arrays_zip_inner(args: &[ArrayRef]) -> Result { Ok(Arc::new(result)) } + +fn arrays_zip_field_name(index: usize) -> String { + (index + 1).to_string() +} + +fn arrays_zip_field_names(len: usize) -> Vec { + (0..len).map(arrays_zip_field_name).collect() +} + +/// Fast path for regular List inputs whose existing buffers already match the +/// zipped output: all offsets and values lengths match, and null rows cover no +/// values. This lets us reuse offsets and child values instead of rebuilding. +fn try_perfect_list_zip( + args: &[ArrayRef], + field_names: &[String], +) -> Result> { + debug_assert_eq!(args.len(), field_names.len()); + + let mut list_arrays = Vec::with_capacity(args.len()); + let mut struct_fields = Vec::with_capacity(args.len()); + + for (arg, field_name) in args.iter().zip(field_names) { + let arr = match arg.data_type() { + List(field) => { + struct_fields.push(Field::new( + field_name.clone(), + field.data_type().clone(), + true, + )); + as_list_array(arg)? + } + _ => return Ok(None), + }; + + list_arrays.push(arr); + } + + let first = list_arrays[0]; + let num_rows = first.len(); + let offsets = first.offsets().clone(); + let values_len = first.values().len(); + + // Reusing the child arrays is only valid when every list uses the exact + // same row boundaries and exposes the same total number of child values. + for arr in &list_arrays { + if arr.values().len() != values_len || arr.offsets() != &offsets { + return Ok(None); + } + } + + let nulls = if list_arrays.iter().any(|arr| arr.null_count() != 0) { + let first_nulls = first.nulls(); + if list_arrays.iter().all(|arr| arr.nulls() == first_nulls) { + first_nulls.cloned() + } else { + // Match the general path: arrays_zip only marks an output row null + // when every concrete input list is null. Mixed null and non-null + // empty lists still produce a non-null empty list, but mixed null + // rows with values must fall back to preserve field-level nulls. + let mut null_builder = NullBufferBuilder::new(num_rows); + for row_idx in 0..num_rows { + let mut all_null = true; + + for arr in &list_arrays { + if arr.is_null(row_idx) { + if arr.offsets()[row_idx + 1] != arr.offsets()[row_idx] { + return Ok(None); + } + } else { + all_null = false; + } + } + + if all_null { + null_builder.append_null(); + } else { + null_builder.append_non_null(); + } + } + + null_builder.finish() + } + } else { + None + }; + + let struct_columns = list_arrays + .iter() + .map(|arr| Arc::clone(arr.values())) + .collect::>(); + let struct_array = + StructArray::try_new(Fields::from(struct_fields), struct_columns, None)?; + let result = ListArray::try_new( + Arc::new(Field::new_list_field( + struct_array.data_type().clone(), + true, + )), + offsets, + Arc::new(struct_array), + nulls, + )?; + + Ok(Some(Arc::new(result))) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int64Array; + use arrow::buffer::NullBuffer; + + fn list(values: Vec, offsets: Vec) -> Arc { + list_with_validity(values, offsets, None) + } + + fn list_with_validity( + values: Vec, + offsets: Vec, + valid: Option>, + ) -> Arc { + Arc::new( + ListArray::try_new( + Arc::new(Field::new_list_field(DataType::Int64, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(Int64Array::from(values)), + valid.map(NullBuffer::from), + ) + .unwrap(), + ) + } + + #[test] + fn perfect_zip_reuses_input_values_and_offsets() { + let left = list(vec![1, 2, 3, 4, 5, 6], vec![0, 2, 3, 6]); + let right = list(vec![10, 20, 30, 40, 50, 60], vec![0, 2, 3, 6]); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + let values = result + .values() + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(result.offsets().ptr_eq(left.offsets())); + assert!(Arc::ptr_eq(values.column(0), left.values())); + assert!(Arc::ptr_eq(values.column(1), right.values())); + } + + #[test] + fn perfect_zip_uses_supplied_field_names() { + let left = list(vec![1, 2, 3], vec![0, 1, 3]); + let right = list(vec![10, 20, 30], vec![0, 1, 3]); + let field_names = vec!["left".to_string(), "right".to_string()]; + + let result = try_perfect_list_zip( + &[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ], + &field_names, + ) + .unwrap() + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + let values = result + .values() + .as_any() + .downcast_ref::() + .unwrap(); + let names = values + .fields() + .iter() + .map(|field| field.name().as_str()) + .collect::>(); + + assert_eq!(names, vec!["left", "right"]); + } + + #[test] + fn perfect_zip_reuses_zero_length_null_rows() { + let left = list_with_validity( + vec![1, 2, 3, 4], + vec![0, 2, 2, 4], + Some(vec![true, false, true]), + ); + let right = list_with_validity( + vec![10, 20, 30, 40], + vec![0, 2, 2, 4], + Some(vec![true, false, true]), + ); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + + assert!(result.offsets().ptr_eq(left.offsets())); + assert!(result.is_null(1)); + } + + #[test] + fn perfect_zip_preserves_mixed_null_empty_rows() { + let left = + list_with_validity(vec![], vec![0, 0, 0, 0], Some(vec![false, true, false])); + let right = + list_with_validity(vec![], vec![0, 0, 0, 0], Some(vec![true, false, false])); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + + assert!(result.offsets().ptr_eq(left.offsets())); + assert!(!result.is_null(0)); + assert!(!result.is_null(1)); + assert!(result.is_null(2)); + } + + #[test] + fn perfect_zip_reuses_null_rows_with_hidden_values() { + let left = + list_with_validity(vec![1, 2, 3, 4], vec![0, 2, 4], Some(vec![true, false])); + let right = list_with_validity( + vec![10, 20, 30, 40], + vec![0, 2, 4], + Some(vec![true, false]), + ); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + + assert!(result.offsets().ptr_eq(left.offsets())); + assert_eq!(result.value_offsets(), &[0, 2, 4]); + assert!(result.is_null(1)); + } + + #[test] + fn mixed_null_row_with_hidden_values_uses_general_path() { + let left = + list_with_validity(vec![1, 2, 3, 4], vec![0, 2, 4], Some(vec![true, false])); + let right = list_with_validity( + vec![10, 20, 30, 40], + vec![0, 2, 4], + Some(vec![true, true]), + ); + + let result = arrays_zip_inner(&[ + Arc::clone(&left) as ArrayRef, + Arc::clone(&right) as ArrayRef, + ]) + .unwrap(); + let result = result.as_any().downcast_ref::().unwrap(); + let values = result + .values() + .as_any() + .downcast_ref::() + .unwrap(); + + assert!(!result.offsets().ptr_eq(left.offsets())); + assert_eq!(result.value_offsets(), &[0, 2, 4]); + assert!(values.column(0).is_null(2)); + assert!(values.column(0).is_null(3)); + assert!(!values.column(1).is_null(2)); + assert!(!values.column(1).is_null(3)); + } +} diff --git a/datafusion/functions-nested/src/cardinality.rs b/datafusion/functions-nested/src/cardinality.rs index d21bb72a457a8..38def2d4e4afe 100644 --- a/datafusion/functions-nested/src/cardinality.rs +++ b/datafusion/functions-nested/src/cardinality.rs @@ -23,10 +23,15 @@ use arrow::array::{ }; use arrow::datatypes::{ DataType, - DataType::{LargeList, List, Map, Null, UInt64}, + DataType::{ + FixedSizeList, LargeList, LargeListView, List, ListView, Map, Null, UInt64, + }, }; use datafusion_common::Result; -use datafusion_common::cast::{as_large_list_array, as_list_array, as_map_array}; +use datafusion_common::cast::{ + as_fixed_size_list_array, as_large_list_array, as_large_list_view_array, + as_list_array, as_list_view_array, as_map_array, +}; use datafusion_common::exec_err; use datafusion_common::utils::{ListCoercion, take_function_args}; use datafusion_expr::{ @@ -146,14 +151,50 @@ fn generic_list_cardinality( let result = array .iter() .map(|arr| match arr { - Some(arr) if arr.is_empty() => Ok(Some(0u64)), - arr => match crate::utils::compute_array_dims(arr)? { - Some(vector) => { - Ok(Some(vector.iter().map(|x| x.unwrap()).product::())) - } - None => Ok(None), - }, + Some(arr) => value_cardinality(&arr).map(Some), + None => Ok(None), }) .collect::>()?; Ok(Arc::new(result) as ArrayRef) } + +fn value_cardinality(array: &ArrayRef) -> Result { + match array.data_type() { + List(_) => { + let list = as_list_array(&array)?; + sum_list_cardinality(list.iter()) + } + LargeList(_) => { + let list = as_large_list_array(&array)?; + sum_list_cardinality(list.iter()) + } + ListView(_) => { + let list = as_list_view_array(&array)?; + sum_list_cardinality(list.iter()) + } + LargeListView(_) => { + let list = as_large_list_view_array(&array)?; + sum_list_cardinality(list.iter()) + } + FixedSizeList(..) => { + let list = as_fixed_size_list_array(&array)?; + sum_list_cardinality(list.iter()) + } + _ => Ok(array.len() as u64), + } +} + +fn sum_list_cardinality(mut iter: I) -> Result +where + I: Iterator>, +{ + iter.try_fold(0u64, |total, arr| { + let value_count = match arr { + Some(arr) => value_cardinality(&arr)?, + None => 0, + }; + total.checked_add(value_count).ok_or_else(|| { + datafusion_common::exec_datafusion_err!("cardinality overflowed u64") + }) + }) +} diff --git a/datafusion/functions-nested/src/concat.rs b/datafusion/functions-nested/src/concat.rs index 8d06140889a55..1f03a0b17014e 100644 --- a/datafusion/functions-nested/src/concat.rs +++ b/datafusion/functions-nested/src/concat.rs @@ -20,26 +20,29 @@ use std::sync::Arc; use crate::make_array::make_array_inner; -use crate::utils::{align_array_dimensions, check_datatypes, make_scalar_function}; +use crate::utils::{ + align_array_dimensions, check_datatypes, list_inner_field, list_type_with_element, + make_scalar_function, +}; use arrow::array::{ Array, ArrayData, ArrayRef, Capacities, GenericListArray, MutableArrayData, OffsetSizeTrait, }; use arrow::buffer::{NullBuffer, OffsetBuffer}; -use arrow::datatypes::{DataType, Field}; +use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::Result; use datafusion_common::utils::{ ListCoercion, base_type, coerced_type_with_base_type_only, }; use datafusion_common::{ cast::as_generic_list_array, - exec_err, plan_err, + exec_err, internal_err, plan_err, utils::{list_ndims, take_function_args}, }; use datafusion_expr::binary::type_union_resolution; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - Volatility, + ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, + Signature, Volatility, }; use datafusion_macros::user_doc; use itertools::Itertools; @@ -104,17 +107,26 @@ impl ScalarUDFImpl for ArrayAppend { &self.signature } - fn return_type(&self, arg_types: &[DataType]) -> Result { - let [array_type, element_type] = take_function_args(self.name(), arg_types)?; - if array_type.is_null() { - Ok(DataType::new_list(element_type.clone(), true)) - } else { - Ok(array_type.clone()) - } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let [array_field, element_field] = + take_function_args(self.name(), args.arg_fields)?; + let data_type = append_prepend_return_type( + array_field.data_type(), + element_field.data_type(), + element_field.is_nullable(), + ); + Ok(Arc::new(Field::new(self.name(), data_type, true))) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_append_inner)(&args.args) + let return_type = args.return_field.data_type().clone(); + make_scalar_function(|args: &[ArrayRef]| array_append_inner(args, &return_type))( + &args.args, + ) } fn aliases(&self) -> &[String] { @@ -186,17 +198,26 @@ impl ScalarUDFImpl for ArrayPrepend { &self.signature } - fn return_type(&self, arg_types: &[DataType]) -> Result { - let [element_type, array_type] = take_function_args(self.name(), arg_types)?; - if array_type.is_null() { - Ok(DataType::new_list(element_type.clone(), true)) - } else { - Ok(array_type.clone()) - } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let [element_field, array_field] = + take_function_args(self.name(), args.arg_fields)?; + let data_type = append_prepend_return_type( + array_field.data_type(), + element_field.data_type(), + element_field.is_nullable(), + ); + Ok(Arc::new(Field::new(self.name(), data_type, true))) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_prepend_inner)(&args.args) + let return_type = args.return_field.data_type().clone(); + make_scalar_function(|args: &[ArrayRef]| array_prepend_inner(args, &return_type))( + &args.args, + ) } fn aliases(&self) -> &[String] { @@ -375,13 +396,38 @@ pub fn array_concat_inner(args: &[ArrayRef]) -> Result { args[0].len(), ))) } else if large_list { - concat_internal::(args) + concat_internal::(args, None) + } else { + concat_internal::(args, None) + } +} + +/// Return type shared by `array_append` and `array_prepend`: the input list +/// type, except that its inner field is nullable whenever the appended or +/// prepended element may be null. +fn append_prepend_return_type( + array_type: &DataType, + element_type: &DataType, + element_nullable: bool, +) -> DataType { + if array_type.is_null() { + DataType::new_list(element_type.clone(), true) } else { - concat_internal::(args) + list_type_with_element(array_type, element_nullable) } } -fn concat_internal(args: &[ArrayRef]) -> Result { +/// Concatenates the list arrays in `args` row-wise. +/// +/// `field` is the list field the output must carry. `array_concat` passes `None` +/// because its `return_type` derives a fresh field from the unified element +/// types, which is what deriving the field from the aligned inputs reproduces. +/// `array_append` / `array_prepend` promise their input's field verbatim and so +/// must pass it in explicitly. +fn concat_internal( + args: &[ArrayRef], + field: Option<&FieldRef>, +) -> Result { let args = align_array_dimensions::(args.to_vec())?; let list_arrays = args @@ -432,17 +478,20 @@ fn concat_internal(args: &[ArrayRef]) -> Result { let start = list_array.offsets()[row_idx].to_usize().unwrap(); let end = list_array.offsets()[row_idx + 1].to_usize().unwrap(); if start < end { - mutable.extend(arr_idx, start, end); + mutable.try_extend(arr_idx, start, end)?; } } offsets.push(O::usize_as(mutable.len())); } - let data_type = list_arrays[0].value_type(); + let field = match field { + Some(field) => Arc::clone(field), + None => Arc::new(Field::new_list_field(list_arrays[0].value_type(), true)), + }; let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(data_type, true)), + field, OffsetBuffer::new(offsets.into()), arrow::array::make_array(data), valid, @@ -451,22 +500,26 @@ fn concat_internal(args: &[ArrayRef]) -> Result { // Kernel functions -fn array_append_inner(args: &[ArrayRef]) -> Result { +fn array_append_inner(args: &[ArrayRef], return_type: &DataType) -> Result { let [array, values] = take_function_args("array_append", args)?; match array.data_type() { DataType::Null => make_array_inner(&[Arc::clone(values)]), - DataType::List(_) => general_append_and_prepend::(args, true), - DataType::LargeList(_) => general_append_and_prepend::(args, true), + DataType::List(_) => general_append_and_prepend::(args, true, return_type), + DataType::LargeList(_) => { + general_append_and_prepend::(args, true, return_type) + } arg_type => exec_err!("array_append does not support type {arg_type}"), } } -fn array_prepend_inner(args: &[ArrayRef]) -> Result { +fn array_prepend_inner(args: &[ArrayRef], return_type: &DataType) -> Result { let [values, array] = take_function_args("array_prepend", args)?; match array.data_type() { DataType::Null => make_array_inner(&[Arc::clone(values)]), - DataType::List(_) => general_append_and_prepend::(args, false), - DataType::LargeList(_) => general_append_and_prepend::(args, false), + DataType::List(_) => general_append_and_prepend::(args, false, return_type), + DataType::LargeList(_) => { + general_append_and_prepend::(args, false, return_type) + } arg_type => exec_err!("array_prepend does not support type {arg_type}"), } } @@ -474,6 +527,7 @@ fn array_prepend_inner(args: &[ArrayRef]) -> Result { fn general_append_and_prepend( args: &[ArrayRef], is_append: bool, + return_type: &DataType, ) -> Result where i64: TryInto, @@ -490,14 +544,22 @@ where (list_array, element_array) }; + let name = if is_append { + "array_append" + } else { + "array_prepend" + }; + let field = list_inner_field(name, return_type)?; + let res = match list_array.value_type() { - DataType::List(_) => concat_internal::(args)?, - DataType::LargeList(_) => concat_internal::(args)?, - data_type => { + DataType::List(_) | DataType::LargeList(_) => { + concat_internal::(args, Some(&field))? + } + _ => { return generic_append_and_prepend::( list_array, element_array, - &data_type, + field, is_append, ); } @@ -516,7 +578,7 @@ where /// /// * `list_array` - A reference to the ListArray to which elements will be appended/prepended. /// * `element_array` - A reference to the Array containing elements to be appended/prepended. -/// * `field` - A reference to the Field describing the data type of the arrays. +/// * `field` - The list field the output must carry, taken from the promised return type. /// * `is_append` - A boolean flag indicating whether to append (`true`) or prepend (`false`) elements. /// /// # Examples @@ -528,7 +590,7 @@ where fn generic_append_and_prepend( list_array: &GenericListArray, element_array: &ArrayRef, - data_type: &DataType, + field: FieldRef, is_append: bool, ) -> Result where @@ -553,11 +615,11 @@ where let start = offset_window[0].to_usize().unwrap(); let end = offset_window[1].to_usize().unwrap(); if is_append { - mutable.extend(values_index, start, end); - mutable.extend(element_index, row_index, row_index + 1); + mutable.try_extend(values_index, start, end)?; + mutable.try_extend(element_index, row_index, row_index + 1)?; } else { - mutable.extend(element_index, row_index, row_index + 1); - mutable.extend(values_index, start, end); + mutable.try_extend(element_index, row_index, row_index + 1)?; + mutable.try_extend(values_index, start, end)?; } offsets.push(offsets[row_index] + O::usize_as(end - start + 1)); } @@ -565,7 +627,7 @@ where let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(data_type.to_owned(), true)), + field, OffsetBuffer::new(offsets.into()), arrow::array::make_array(data), None, diff --git a/datafusion/functions-nested/src/dimension.rs b/datafusion/functions-nested/src/dimension.rs index 01fb81d878e0b..7e9a10f362562 100644 --- a/datafusion/functions-nested/src/dimension.rs +++ b/datafusion/functions-nested/src/dimension.rs @@ -122,7 +122,7 @@ make_udf_expr_and_func!( #[user_doc( doc_section(label = "Array Functions"), description = "Returns the number of dimensions of the array.", - syntax_example = "array_ndims(array, element)", + syntax_example = "array_ndims(array)", sql_example = r#"```sql > select array_ndims([[1, 2, 3], [4, 5, 6]]); +----------------------------------+ @@ -134,8 +134,7 @@ make_udf_expr_and_func!( argument( name = "array", description = "Array expression. Can be a constant, column, or function, and any combination of array operators." - ), - argument(name = "element", description = "Array element.") + ) )] #[derive(Debug, PartialEq, Eq, Hash)] pub(super) struct ArrayNdims { diff --git a/datafusion/functions-nested/src/distance.rs b/datafusion/functions-nested/src/distance.rs index edf1806b66c2d..c9aec816676a7 100644 --- a/datafusion/functions-nested/src/distance.rs +++ b/datafusion/functions-nested/src/distance.rs @@ -18,9 +18,7 @@ //! [ScalarUDFImpl] definitions for array_distance function. use crate::utils::make_scalar_function; -use arrow::array::{ - Array, ArrayRef, Float64Array, LargeListArray, ListArray, OffsetSizeTrait, -}; +use arrow::array::{Array, ArrayRef, Float64Array, OffsetSizeTrait}; use arrow::datatypes::{ DataType, DataType::{FixedSizeList, LargeList, List, Null}, @@ -35,7 +33,6 @@ use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; -use datafusion_functions::downcast_arg; use datafusion_macros::user_doc; use itertools::Itertools; use std::sync::Arc; @@ -44,13 +41,13 @@ make_udf_expr_and_func!( ArrayDistance, array_distance, array, - "returns the Euclidean distance between two numeric arrays.", + "returns the Euclidean distance between two one-dimensional numeric arrays.", array_distance_udf ); #[user_doc( doc_section(label = "Array Functions"), - description = "Returns the Euclidean distance between two input arrays of equal length.", + description = "Returns the Euclidean distance between two one-dimensional input arrays of equal length.", syntax_example = "array_distance(array1, array2)", sql_example = r#"```sql > select array_distance([1, 2], [1, 4]); @@ -106,16 +103,30 @@ impl ScalarUDFImpl for ArrayDistance { fn coerce_types(&self, arg_types: &[DataType]) -> Result> { let [_, _] = take_function_args(self.name(), arg_types)?; let coercion = Some(&ListCoercion::FixedSizedListToList); - let arg_types = arg_types.iter().map(|arg_type| { - if matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) { + let arg_types = arg_types.iter().map(|arg_type| match arg_type { + Null => Ok(coerced_type_with_base_type_only( + arg_type, + &DataType::Float64, + coercion, + )), + List(field) | LargeList(field) | FixedSizeList(field, _) => { + // Distance between nested lists is not supported + if matches!( + field.data_type(), + List(_) | LargeList(_) | FixedSizeList(..) + ) { + return plan_err!( + "{} only supports one-dimensional arrays, got {arg_type}", + self.name() + ); + } Ok(coerced_type_with_base_type_only( arg_type, &DataType::Float64, coercion, )) - } else { - plan_err!("{} does not support type {arg_type}", self.name()) } + _ => plan_err!("{} does not support type {arg_type}", self.name()), }); arg_types.try_collect() @@ -172,43 +183,6 @@ fn compute_array_distance( None => return Ok(None), }; - let mut value1 = value1; - let mut value2 = value2; - - loop { - match value1.data_type() { - List(_) => { - if downcast_arg!(value1, ListArray).null_count() > 0 { - return Ok(None); - } - value1 = downcast_arg!(value1, ListArray).value(0); - } - LargeList(_) => { - if downcast_arg!(value1, LargeListArray).null_count() > 0 { - return Ok(None); - } - value1 = downcast_arg!(value1, LargeListArray).value(0); - } - _ => break, - } - - match value2.data_type() { - List(_) => { - if downcast_arg!(value2, ListArray).null_count() > 0 { - return Ok(None); - } - value2 = downcast_arg!(value2, ListArray).value(0); - } - LargeList(_) => { - if downcast_arg!(value2, LargeListArray).null_count() > 0 { - return Ok(None); - } - value2 = downcast_arg!(value2, LargeListArray).value(0); - } - _ => break, - } - } - // Check for NULL values inside the arrays if value1.null_count() != 0 || value2.null_count() != 0 { return Ok(None); diff --git a/datafusion/functions-nested/src/empty.rs b/datafusion/functions-nested/src/empty.rs index 262eb4935c968..6db412d29b0d8 100644 --- a/datafusion/functions-nested/src/empty.rs +++ b/datafusion/functions-nested/src/empty.rs @@ -122,9 +122,19 @@ fn array_empty_inner(args: &[ArrayRef]) -> Result { } fn general_array_empty(array: &ArrayRef) -> Result { - let result = as_generic_list_array::(array)? - .iter() - .map(|arr| arr.map(|arr| arr.is_empty())) - .collect::(); + let result = as_generic_list_array::(array)?; + let is_empty_iter = result.offsets().lengths().map(|n| n == 0); + // SAFETY: this is safe since the iterator lengths is exact size and + // trusted - it maps over fixed known number of elements + let output_buffer = unsafe { BooleanArray::from_trusted_len_iter(is_empty_iter) }; + + let (values, _) = output_buffer.into_parts(); + + // Add the nulls + let result = BooleanArray::new( + values, + result.nulls().filter(|n| n.null_count() > 0).cloned(), + ); + Ok(Arc::new(result)) } diff --git a/datafusion/functions-nested/src/except.rs b/datafusion/functions-nested/src/except.rs index 12ed6c2e186f4..dbf815c0ec539 100644 --- a/datafusion/functions-nested/src/except.rs +++ b/datafusion/functions-nested/src/except.rs @@ -27,7 +27,7 @@ use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::compute::take; use arrow::datatypes::{DataType, FieldRef}; use arrow::row::{RowConverter, SortField}; -use datafusion_common::utils::{ListCoercion, take_function_args}; +use datafusion_common::utils::{ListCoercion, normalize_float_zero, take_function_args}; use datafusion_common::{HashSet, Result, internal_err}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -169,16 +169,21 @@ fn general_except( ) -> Result> { let converter = RowConverter::new(vec![SortField::new(l.value_type())])?; + // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) groups + // ±0 together for both the rhs lookup set and the lhs probe. + let l_values_norm = normalize_float_zero(l.values()); + let r_values_norm = normalize_float_zero(r.values()); + // Only convert the visible portion of the values array. For sliced // ListArrays, values() returns the full underlying array but only // elements between the first and last offset are referenced. let l_first = l.offsets()[0].as_usize(); let l_len = l.offsets()[l.len()].as_usize() - l_first; - let l_values = converter.convert_columns(&[l.values().slice(l_first, l_len)])?; + let l_values = converter.convert_columns(&[l_values_norm.slice(l_first, l_len)])?; let r_first = r.offsets()[0].as_usize(); let r_len = r.offsets()[r.len()].as_usize() - r_first; - let r_values = converter.convert_columns(&[r.values().slice(r_first, r_len)])?; + let r_values = converter.convert_columns(&[r_values_norm.slice(r_first, r_len)])?; let mut offsets = Vec::::with_capacity(l.len() + 1); offsets.push(OffsetSize::usize_as(0)); @@ -223,11 +228,11 @@ fn general_except( } else if OffsetSize::IS_LARGE { let indices = UInt64Array::from(indices.into_iter().map(|i| i as u64).collect::>()); - take(l.values().as_ref(), &indices, None)? + take(l_values_norm.as_ref(), &indices, None)? } else { let indices = UInt32Array::from(indices.into_iter().map(|i| i as u32).collect::>()); - take(l.values().as_ref(), &indices, None)? + take(l_values_norm.as_ref(), &indices, None)? }; Ok(GenericListArray::::new( diff --git a/datafusion/functions-nested/src/extract.rs b/datafusion/functions-nested/src/extract.rs index fe1e31e8d5efb..cb7a316b289a9 100644 --- a/datafusion/functions-nested/src/extract.rs +++ b/datafusion/functions-nested/src/extract.rs @@ -23,9 +23,8 @@ use arrow::array::{ }; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; -use arrow::datatypes::{ - DataType::{FixedSizeList, LargeList, LargeListView, List, ListView, Null}, - Field, +use arrow::datatypes::DataType::{ + FixedSizeList, LargeList, LargeListView, List, ListView, Null, }; use datafusion_common::cast::as_large_list_array; use datafusion_common::cast::as_list_array; @@ -48,7 +47,7 @@ use datafusion_expr::{ use datafusion_macros::user_doc; use std::sync::Arc; -use crate::utils::make_scalar_function; +use crate::utils::{list_inner_field, make_scalar_function}; // Create static instances of ScalarUDFs for each function make_udf_expr_and_func!( @@ -256,9 +255,9 @@ where let end = offset_window[1]; let len = end - start; - // array is null - if array.is_null(row_index) { - mutable.extend_nulls(1); + // array or index is null + if array.is_null(row_index) || indexes.is_null(row_index) { + mutable.try_extend_nulls(1)?; continue; } @@ -266,10 +265,10 @@ where if let Some(index) = index { let start = start.as_usize() + index.as_usize(); - mutable.extend(0, start, start + 1_usize); + mutable.try_extend(0, start, start + 1_usize)?; } else { // Index out of bounds - mutable.extend_nulls(1); + mutable.try_extend_nulls(1)?; } } @@ -289,7 +288,7 @@ pub fn array_slice(array: Expr, begin: Expr, end: Expr, stride: Option) -> #[user_doc( doc_section(label = "Array Functions"), description = "Returns a slice of the array based on 1-indexed start and end positions.", - syntax_example = "array_slice(array, begin, end)", + syntax_example = "array_slice(array, begin, end[, stride])", sql_example = r#"```sql > select array_slice([1, 2, 3, 4, 5, 6, 7, 8], 3, 6); +--------------------------------------------------------+ @@ -622,9 +621,16 @@ where let values = array.values(); let original_data = values.to_data(); let capacity = Capacities::Array(original_data.len()); - + // Carry the input's list field through to the output so that the returned + // type matches the one promised by `return_type` / `return_field_from_args`, + // including the field name, nullability and metadata. + let field = list_inner_field("general_array_slice", array.data_type())?; + + // `use_nulls` is false because we never call `try_extend_nulls`: null rows are + // emitted as empty slices. Arrow still allocates a validity buffer on its own + // if the child array has nulls. let mut mutable = - MutableArrayData::with_capacities(vec![&original_data], true, capacity); + MutableArrayData::with_capacities(vec![&original_data], false, capacity); // We have the slice syntax compatible with DuckDB v0.8.1. // The rule `adjusted_from_index` and `adjusted_to_index` follows the rule of array_slice in duckdb. @@ -638,9 +644,11 @@ where let end = offset_window[1]; let len = end - start; + // The row is null, so its contents are never observed. Emit an empty + // slice rather than a null child element: the input's list field may be + // non-nullable, in which case a null child would be invalid. if nulls.as_ref().is_some_and(|n| n.is_null(row_index)) { - mutable.extend_nulls(1); - offsets.push(offsets[row_index] + O::usize_as(1)); + offsets.push(offsets[row_index]); continue; } @@ -665,14 +673,14 @@ where } => { let start_index = (start + rel_start).to_usize().unwrap(); let end_index = (start + rel_start + slice_len).to_usize().unwrap(); - mutable.extend(0, start_index, end_index); + mutable.try_extend(0, start_index, end_index)?; offsets.push(offsets[row_index] + slice_len); } SlicePlan::Indices(indices) => { let count = indices.len(); for rel_index in indices { let absolute_index = (start + rel_index).to_usize().unwrap(); - mutable.extend(0, absolute_index, absolute_index + 1); + mutable.try_extend(0, absolute_index, absolute_index + 1)?; } offsets.push(offsets[row_index] + O::usize_as(count)); } @@ -682,7 +690,7 @@ where let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(array.value_type(), true)), + field, OffsetBuffer::::new(offsets.into()), arrow::array::make_array(data), nulls, @@ -704,12 +712,15 @@ where let field = match array.data_type() { ListView(field) | LargeListView(field) => Arc::clone(field), other => { - return internal_err!("array_slice got unexpected data type: {}", other); + return internal_err!( + "general_list_view_array_slice got unexpected data type: {other}" + ); } }; + // See the note on `use_nulls` in `general_array_slice`. let mut mutable = - MutableArrayData::with_capacities(vec![&original_data], true, capacity); + MutableArrayData::with_capacities(vec![&original_data], false, capacity); // We must build `offsets` and `sizes` buffers manually as ListView does not enforce // monotonically increasing offsets. @@ -754,7 +765,7 @@ where } => { let start_index = (start + rel_start).to_usize().unwrap(); let end_index = (start + rel_start + slice_len).to_usize().unwrap(); - mutable.extend(0, start_index, end_index); + mutable.try_extend(0, start_index, end_index)?; offsets.push(current_offset); sizes.push(slice_len); current_offset += slice_len; @@ -763,7 +774,7 @@ where let count = indices.len(); for rel_index in indices { let absolute_index = (start + rel_index).to_usize().unwrap(); - mutable.extend(0, absolute_index, absolute_index + 1); + mutable.try_extend(0, absolute_index, absolute_index + 1)?; } let length = O::usize_as(count); offsets.push(current_offset); @@ -970,7 +981,7 @@ where #[user_doc( doc_section(label = "Array Functions"), - description = "Returns the first non-null element in the array.", + description = "Returns the first non-null element in the array. Returns NULL if the array is empty or NULL.", syntax_example = "array_any_value(array)", sql_example = r#"```sql > select array_any_value([NULL, 1, 2, 3]); @@ -1062,10 +1073,18 @@ where for (row_index, offset_window) in array.offsets().windows(2).enumerate() { let start = offset_window[0]; + let end = offset_window[1]; - // array is null + // the list element is null if array.is_null(row_index) { - mutable.extend_nulls(1); + mutable.try_extend_nulls(1)?; + continue; + } + + // the list element is empty; there is no value to take, so the result + // is NULL. + if start == end { + mutable.try_extend_nulls(1)?; continue; } @@ -1077,16 +1096,16 @@ where row_nulls_buffer.valid_indices().next() { let index = start.as_usize() + first_non_null_index; - mutable.extend(0, index, index + 1) + mutable.try_extend(0, index, index + 1)?; } else { // all the elements in the array are null - mutable.extend_nulls(1); + mutable.try_extend_nulls(1)?; } } None => { // no nulls are present in the array so take the first element let index = start.as_usize(); - mutable.extend(0, index, index + 1); + mutable.try_extend(0, index, index + 1)?; } } } @@ -1107,7 +1126,7 @@ mod tests { }; use arrow::array::{ListArray, RecordBatch}; use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; - use arrow::datatypes::{DataType, Field}; + use arrow::datatypes::{DataType, Field, Int32Type}; use datafusion_common::{Column, DFSchema, Result, assert_batches_eq}; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::{Expr, ExprSchemable}; @@ -1198,6 +1217,26 @@ mod tests { Ok(()) } + #[test] + fn test_array_element_null_index_with_non_zero_buffer_returns_null() -> Result<()> { + let list_array = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![Some(4)]), + Some(vec![Some(5)]), + ]); + let indexes = Int64Array::new( + ScalarBuffer::from(vec![1, 1, 1]), + Some(NullBuffer::from(vec![true, false, true])), + ); + + let result = general_array_element(&list_array, &indexes)?; + let expected = Int32Array::from(vec![Some(1), None, Some(5)]); + + assert_eq!(result.as_primitive::(), &expected); + + Ok(()) + } + #[test] fn test_array_any_null_handling() -> Result<()> { let values: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])); diff --git a/datafusion/functions-nested/src/lambda_utils.rs b/datafusion/functions-nested/src/lambda_utils.rs index cb8682d4bd18b..4b01ae314e4c7 100644 --- a/datafusion/functions-nested/src/lambda_utils.rs +++ b/datafusion/functions-nested/src/lambda_utils.rs @@ -17,13 +17,19 @@ //! Shared utilities for `(array, lambda)` style higher-order functions. -use arrow::array::ArrayRef; -use arrow::datatypes::{DataType, FieldRef}; +use arrow::array::{ArrayRef, AsArray, BooleanArray, OffsetSizeTrait, new_null_array}; +use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::compute::take_arrays; +use arrow::datatypes::{ArrowNativeType, DataType, FieldRef}; +use datafusion_common::utils::{adjust_offsets_for_slice, list_values_row_number}; use datafusion_common::{ Result, ScalarValue, plan_err, utils::{list_values, take_function_args}, }; -use datafusion_expr::{ColumnarValue, LambdaParametersProgress, ValueOrLambda}; +use datafusion_common::{exec_datafusion_err, exec_err}; +use datafusion_expr::{ + ColumnarValue, HigherOrderFunctionArgs, LambdaParametersProgress, ValueOrLambda, +}; use std::sync::Arc; /// Extracts a `(value, lambda)` pair from a [`ValueOrLambda`] slice. @@ -65,6 +71,7 @@ pub(crate) fn coerce_single_list_arg( DataType::List(Arc::clone(field)) } DataType::LargeListView(field) => DataType::LargeList(Arc::clone(field)), + DataType::Null => DataType::new_list(DataType::Null, true), _ => return plan_err!("{name} expected a list as first argument, got {list}"), }; @@ -125,12 +132,225 @@ pub(crate) fn extract_list_values( Ok(ListValuesResult::Values(values)) } +pub(crate) enum SingleListLambdaResult { + EarlyReturn(ColumnarValue), + Ready(EvaluatedListLambda), +} + +pub(crate) struct EvaluatedListLambda { + pub original_list: ArrayRef, + pub flattened_values: ArrayRef, + pub evaluated_result: ColumnarValue, + row_offsets: Vec, +} + +impl EvaluatedListLambda { + pub(crate) fn len(&self) -> usize { + self.original_list.len() + } + + pub(crate) fn nulls(&self) -> Option<&NullBuffer> { + self.original_list.nulls() + } + + pub(crate) fn row_range(&self, i: usize) -> (usize, usize) { + (self.row_offsets[i], self.row_offsets[i + 1]) + } + + pub(crate) fn adjusted_offsets(&self) -> OffsetBuffer { + OffsetBuffer::from_lengths(self.row_offsets.windows(2).map(|w| w[1] - w[0])) + } + + pub(crate) fn boolean_predicate(&self, name: &str) -> Result { + let arr = self + .evaluated_result + .clone() + .into_array(self.flattened_values.len())?; + + let predicate = arr.as_any().downcast_ref::().ok_or_else(|| { + exec_datafusion_err!("{} predicate must return boolean array", name) + })?; + + Ok(predicate.clone()) + } +} + +fn adjusted_row_offsets(list: &ArrayRef) -> Result> { + Ok(match list.data_type() { + DataType::List(_) => adjust_offsets_for_slice(list.as_list::()) + .iter() + .map(|o| o.as_usize()) + .collect(), + DataType::LargeList(_) => adjust_offsets_for_slice(list.as_list::()) + .iter() + .map(|o| o.as_usize()) + .collect(), + other => return exec_err!("expected list, got {other}"), + }) +} + +fn evaluate_single_list_lambda( + name: &str, + args: &HigherOrderFunctionArgs, +) -> Result { + let (original_list, lambda) = value_lambda_pair(name, &args.args)?; + let original_list = original_list.to_array(args.number_rows)?; + + if original_list.null_count() == original_list.len() { + return Ok(SingleListLambdaResult::EarlyReturn(ColumnarValue::Array( + new_null_array(args.return_type(), original_list.len()), + ))); + } + + let flattened_values = list_values(&original_list)?; + let values_param = || Ok(Arc::clone(&flattened_values)); + + let evaluated_result = lambda.evaluate(&[&values_param], |arrays| { + let indices = list_values_row_number(&original_list)?; + Ok(take_arrays(arrays, &indices, None)?) + })?; + + let row_offsets = adjusted_row_offsets(&original_list)?; + + Ok(SingleListLambdaResult::Ready(EvaluatedListLambda { + original_list, + flattened_values, + evaluated_result, + row_offsets, + })) +} + +pub(crate) fn evaluate_single_list_predicate( + name: &str, + args: &HigherOrderFunctionArgs, +) -> Result { + let result = evaluate_single_list_lambda(name, args)?; + let SingleListLambdaResult::Ready(evaluated_list_lambda) = &result else { + return Ok(result); + }; + + match &evaluated_list_lambda.evaluated_result { + ColumnarValue::Scalar(ScalarValue::Boolean(_)) => Ok(result), + ColumnarValue::Scalar(scalar) => exec_err!( + "{name} lambda must return boolean, got {}", + scalar.data_type() + ), + ColumnarValue::Array(array) if array.as_any().is::() => Ok(result), + ColumnarValue::Array(array) => exec_err!( + "{name} lambda must return boolean, got {}", + array.data_type() + ), + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::{ + array::ArrayRef, + buffer::{NullBuffer, OffsetBuffer}, + datatypes::{DataType, Field}, + }; + use datafusion_common::Result; + + use super::{adjusted_row_offsets, coerce_single_list_arg}; + use crate::lambda_utils::test_utils::{create_i32_large_list, create_i32_list}; + + #[test] + fn adjusted_row_offsets_matches_list_lengths() -> Result<()> { + let list = create_i32_list( + vec![1, 2, 3, 4, 5], + OffsetBuffer::::from_lengths(vec![2, 0, 3]), + None, + ); + let list = Arc::new(list) as ArrayRef; + assert_eq!(adjusted_row_offsets(&list)?, vec![0, 2, 2, 5]); + Ok(()) + } + + #[test] + fn adjusted_row_offsets_on_sliced_list() -> Result<()> { + let list = create_i32_list( + vec![10, 1, 2, 3, 4], + OffsetBuffer::::from_lengths(vec![1, 2, 2]), + None, + ) + .slice(1, 2); + let list = Arc::new(list) as ArrayRef; + assert_eq!(adjusted_row_offsets(&list)?, vec![0, 2, 4]); + Ok(()) + } + + #[test] + fn adjusted_row_offsets_null_rows_keep_backing_lengths() -> Result<()> { + let list = create_i32_list( + vec![1, 99, 100, 2], + OffsetBuffer::::from_lengths(vec![1, 2, 1]), + Some(NullBuffer::from(vec![true, false, true])), + ); + let list = Arc::new(list) as ArrayRef; + assert_eq!(adjusted_row_offsets(&list)?, vec![0, 1, 3, 4]); + Ok(()) + } + + #[test] + fn adjusted_row_offsets_large_list_parity() -> Result<()> { + let list = create_i32_large_list( + vec![1, 2, 3, 4], + OffsetBuffer::::from_lengths(vec![1, 3]), + None, + ); + let list = Arc::new(list) as ArrayRef; + assert_eq!(adjusted_row_offsets(&list)?, vec![0, 1, 4]); + Ok(()) + } + + #[test] + fn coerce_single_list_arg_supports_advertised_list_likes() -> Result<()> { + let field = Arc::new(Field::new_list_field(DataType::Int32, true)); + assert_eq!( + coerce_single_list_arg("test", &[DataType::List(Arc::clone(&field))])?, + vec![DataType::List(Arc::clone(&field))] + ); + assert_eq!( + coerce_single_list_arg("test", &[DataType::LargeList(Arc::clone(&field))])?, + vec![DataType::LargeList(Arc::clone(&field))] + ); + assert_eq!( + coerce_single_list_arg( + "test", + &[DataType::FixedSizeList(Arc::clone(&field), 3)] + )?, + vec![DataType::List(Arc::clone(&field))] + ); + assert_eq!( + coerce_single_list_arg("test", &[DataType::ListView(Arc::clone(&field))])?, + vec![DataType::List(Arc::clone(&field))] + ); + assert_eq!( + coerce_single_list_arg( + "test", + &[DataType::LargeListView(Arc::clone(&field))] + )?, + vec![DataType::LargeList(field)] + ); + Ok(()) + } + + #[test] + fn coerce_single_list_arg_rejects_non_list() { + let err = coerce_single_list_arg("test", &[DataType::Int32]).unwrap_err(); + assert!(err.to_string().contains("expected a list")); + } +} + #[cfg(test)] pub(crate) mod test_utils { use std::{collections::HashMap, sync::Arc}; use arrow::{ - array::{Array, ArrayRef, Int32Array, ListArray, RecordBatch}, + array::{Array, ArrayRef, Int32Array, LargeListArray, ListArray, RecordBatch}, buffer::{NullBuffer, OffsetBuffer}, datatypes::{DataType, Field}, }; @@ -140,6 +360,7 @@ pub(crate) mod test_utils { execution_props::ExecutionProps, expr::{HigherOrderFunction, LambdaVariable}, lambda, + physical_planning_context::PhysicalPlanningContext, }; use datafusion_physical_expr::create_physical_expr; @@ -152,8 +373,17 @@ pub(crate) mod test_utils { ListArray::new(list_field, offsets, Arc::new(values.into()), nulls) } + pub(crate) fn create_i32_large_list( + values: impl Into, + offsets: OffsetBuffer, + nulls: Option, + ) -> LargeListArray { + let list_field = Arc::new(Field::new_list_field(DataType::Int32, true)); + LargeListArray::new(list_field, offsets, Arc::new(values.into()), nulls) + } + pub(crate) fn eval_hof_on_i32_list( - func: Arc, + func: Arc, list: impl Array + Clone + 'static, lambda_body: Expr, ) -> Result { @@ -174,6 +404,7 @@ pub(crate) mod test_utils { )), &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), )? .evaluate(&RecordBatch::try_new( Arc::clone(schema.inner()), @@ -182,6 +413,39 @@ pub(crate) mod test_utils { .into_array(list.len()) } + /// Evaluates a HOF whose lambda body may capture an outer `number` column. + pub(crate) fn eval_hof_on_i32_list_with_outer( + func: Arc, + list: impl Array + Clone + 'static, + number: Int32Array, + lambda_body: Expr, + ) -> Result { + assert_eq!(list.len(), number.len()); + let schema = DFSchema::from_unqualified_fields( + vec![ + Field::new("list", list.data_type().clone(), list.is_nullable()), + Field::new("number", DataType::Int32, true), + ] + .into(), + HashMap::new(), + )?; + + create_physical_expr( + &Expr::HigherOrderFunction(HigherOrderFunction::new( + func, + vec![col("list"), lambda(["v"], lambda_body)], + )), + &schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )? + .evaluate(&RecordBatch::try_new( + Arc::clone(schema.inner()), + vec![Arc::new(list.clone()), Arc::new(number)], + )?)? + .into_array(list.len()) + } + pub(crate) fn v() -> Expr { Expr::LambdaVariable(LambdaVariable::new( "v".to_string(), diff --git a/datafusion/functions-nested/src/length.rs b/datafusion/functions-nested/src/length.rs index 9579c3c9cd658..24c79c40d7d5d 100644 --- a/datafusion/functions-nested/src/length.rs +++ b/datafusion/functions-nested/src/length.rs @@ -49,7 +49,7 @@ make_udf_expr_and_func!( #[user_doc( doc_section(label = "Array Functions"), description = "Returns the length of the array dimension.", - syntax_example = "array_length(array, dimension)", + syntax_example = "array_length(array[, dimension])", sql_example = r#"```sql > select array_length([1, 2, 3, 4, 5], 1); +-------------------------------------------+ @@ -62,7 +62,7 @@ make_udf_expr_and_func!( name = "array", description = "Array expression. Can be a constant, column, or function, and any combination of array operators." ), - argument(name = "dimension", description = "Array dimension.") + argument(name = "dimension", description = "Array dimension. Default is 1") )] #[derive(Debug, PartialEq, Eq, Hash)] pub struct ArrayLength { diff --git a/datafusion/functions-nested/src/lib.rs b/datafusion/functions-nested/src/lib.rs index 1e6dc68cb23ae..2c7bd25d7dbcd 100644 --- a/datafusion/functions-nested/src/lib.rs +++ b/datafusion/functions-nested/src/lib.rs @@ -40,13 +40,18 @@ pub mod macros; #[macro_use] pub mod macros_lambda; +pub mod array_add; pub mod array_any_match; -pub(crate) mod lambda_utils; - +pub mod array_avg; pub mod array_compact; pub mod array_filter; +pub mod array_first; pub mod array_has; pub mod array_normalize; +pub mod array_product; +pub mod array_scale; +pub mod array_subtract; +pub mod array_sum; pub mod array_transform; pub mod arrays_zip; pub mod cardinality; @@ -60,6 +65,7 @@ pub mod expr_ext; pub mod extract; pub mod flatten; pub mod inner_product; +pub(crate) mod lambda_utils; pub mod length; pub mod make_array; pub mod map; @@ -89,13 +95,20 @@ use std::sync::Arc; /// Fluent-style API for creating `Expr`s pub mod expr_fn { + pub use super::array_add::array_add; pub use super::array_any_match::array_any_match; + pub use super::array_avg::array_avg; pub use super::array_compact::array_compact; pub use super::array_filter::array_filter; + pub use super::array_first::array_first; pub use super::array_has::array_has; pub use super::array_has::array_has_all; pub use super::array_has::array_has_any; pub use super::array_normalize::array_normalize; + pub use super::array_product::array_product; + pub use super::array_scale::array_scale; + pub use super::array_subtract::array_subtract; + pub use super::array_sum::array_sum; pub use super::array_transform::array_transform; pub use super::arrays_zip::arrays_zip; pub use super::cardinality::cardinality; @@ -171,6 +184,12 @@ pub fn all_default_nested_functions() -> Vec> { empty::array_empty_udf(), length::array_length_udf(), array_normalize::array_normalize_udf(), + array_add::array_add_udf(), + array_avg::array_avg_udf(), + array_product::array_product_udf(), + array_scale::array_scale_udf(), + array_subtract::array_subtract_udf(), + array_sum::array_sum_udf(), cosine_distance::cosine_distance_udf(), inner_product::inner_product_udf(), distance::array_distance_udf(), @@ -201,10 +220,11 @@ pub fn all_default_nested_functions() -> Vec> { ] } -pub fn all_default_higher_order_functions() -> Vec> { +pub fn all_default_higher_order_functions() -> Vec> { vec![ array_any_match::array_any_match_higher_order_function(), array_filter::array_filter_higher_order_function(), + array_first::array_first_higher_order_function(), array_transform::array_transform_higher_order_function(), ] } @@ -220,7 +240,7 @@ pub fn register_all(registry: &mut dyn FunctionRegistry) -> Result<()> { Ok(()) as Result<()> })?; - let functions: Vec> = all_default_higher_order_functions(); + let functions: Vec> = all_default_higher_order_functions(); functions.into_iter().try_for_each(|function| { let existing_function = registry.register_higher_order_function(function)?; if let Some(existing_function) = existing_function { diff --git a/datafusion/functions-nested/src/macros_lambda.rs b/datafusion/functions-nested/src/macros_lambda.rs index 8c15d8aed13b6..c8fe670844b2d 100644 --- a/datafusion/functions-nested/src/macros_lambda.rs +++ b/datafusion/functions-nested/src/macros_lambda.rs @@ -95,11 +95,11 @@ macro_rules! create_higher_order { ($UDF:ident, $HIGHER_ORDER_UDF_FN:ident, $CTOR:path) => { #[doc = concat!("HigherOrderFunction that returns a [`HigherOrderUDF`](datafusion_expr::HigherOrderUDF) for ")] #[doc = stringify!($UDF)] - pub fn $HIGHER_ORDER_UDF_FN() -> std::sync::Arc { + pub fn $HIGHER_ORDER_UDF_FN() -> std::sync::Arc { // Singleton instance of [`$UDF`], ensures the UDF is only created once - static INSTANCE: std::sync::LazyLock> = + static INSTANCE: std::sync::LazyLock> = std::sync::LazyLock::new(|| { - std::sync::Arc::new($CTOR()) + std::sync::Arc::new(datafusion_expr::HigherOrderUDF::new_from_impl($CTOR())) }); std::sync::Arc::clone(&INSTANCE) } diff --git a/datafusion/functions-nested/src/make_array.rs b/datafusion/functions-nested/src/make_array.rs index 32af5df2c6019..ba746cd9cf686 100644 --- a/datafusion/functions-nested/src/make_array.rs +++ b/datafusion/functions-nested/src/make_array.rs @@ -50,7 +50,7 @@ make_udf_expr_and_func!( #[user_doc( doc_section(label = "Array Functions"), description = "Returns an array using the specified input expressions.", - syntax_example = "make_array(expression1[, ..., expression_n])", + syntax_example = "make_array([expression1, ..., expression_n])", sql_example = r#"```sql > select make_array(1, 2, 3, 4, 5); +----------------------------------------------------------+ @@ -224,9 +224,9 @@ pub fn array_array( && !arg.is_null(row_idx) && arg.is_valid(row_idx) { - mutable.extend(arr_idx, row_idx, row_idx + 1); + mutable.try_extend(arr_idx, row_idx, row_idx + 1)?; } else { - mutable.extend_nulls(1); + mutable.try_extend_nulls(1)?; } } offsets.push(O::usize_as(mutable.len())); diff --git a/datafusion/functions-nested/src/map.rs b/datafusion/functions-nested/src/map.rs index c7418e9021494..660442f8a3bfd 100644 --- a/datafusion/functions-nested/src/map.rs +++ b/datafusion/functions-nested/src/map.rs @@ -63,47 +63,33 @@ fn can_evaluate_to_const(args: &[ColumnarValue]) -> bool { .all(|arg| matches!(arg, ColumnarValue::Scalar(_))) } -fn make_map_batch(args: &[ColumnarValue]) -> Result { +fn into_array_and_type( + arg: ColumnarValue, + rows: usize, + expand_scalar: bool, +) -> Result<(ArrayRef, DataType)> { + let data_type = arg.data_type(); + let array = if expand_scalar { + arg.into_array(rows)? + } else { + get_first_array_ref(&arg)? + }; + + Ok((array, data_type)) +} + +fn make_map_batch(args: Vec, number_rows: usize) -> Result { + let can_evaluate_to_const = can_evaluate_to_const(&args); let [keys_arg, values_arg] = take_function_args("make_map", args)?; + let expand_scalar = !can_evaluate_to_const; - let can_evaluate_to_const = can_evaluate_to_const(args); - - let keys = get_first_array_ref(keys_arg)?; - let key_array = keys.as_ref(); - - match keys_arg { - ColumnarValue::Array(_) => match key_array.data_type() { - DataType::List(_) => keys - .as_list::() - .iter() - .flatten() - .try_for_each(|row| validate_map_keys(row.as_ref()))?, - DataType::LargeList(_) => keys - .as_list::() - .iter() - .flatten() - .try_for_each(|row| validate_map_keys(row.as_ref()))?, - DataType::FixedSizeList(_, _) => { - keys.as_fixed_size_list() - .iter() - .flatten() - .try_for_each(|row| validate_map_keys(row.as_ref()))? - } - data_type => { - return exec_err!( - "Expected list, large_list or fixed_size_list, got {:?}", - data_type - ); - } - }, - ColumnarValue::Scalar(_) => { - validate_map_keys(key_array)?; - } - } + let (keys, keys_data_type) = + into_array_and_type(keys_arg, number_rows, expand_scalar)?; + let (values, _) = into_array_and_type(values_arg, number_rows, expand_scalar)?; - let values = get_first_array_ref(values_arg)?; + validate_map_keys_for_data_type(&keys, &keys_data_type, can_evaluate_to_const)?; - make_map_batch_internal(&keys, &values, can_evaluate_to_const, &keys_arg.data_type()) + make_map_batch_internal(&keys, &values, can_evaluate_to_const, &keys_data_type) } fn validate_unique_primitive_keys(array: &dyn Array) -> Result<()> @@ -227,6 +213,38 @@ fn validate_map_keys(array: &dyn Array) -> Result<()> { } } +fn validate_map_keys_for_data_type( + keys: &ArrayRef, + keys_data_type: &DataType, + can_evaluate_to_const: bool, +) -> Result<()> { + if can_evaluate_to_const { + return validate_map_keys(keys.as_ref()); + } + + match keys_data_type { + DataType::List(_) => keys + .as_list::() + .iter() + .flatten() + .try_for_each(|row| validate_map_keys(row.as_ref())), + DataType::LargeList(_) => keys + .as_list::() + .iter() + .flatten() + .try_for_each(|row| validate_map_keys(row.as_ref())), + DataType::FixedSizeList(_, _) => keys + .as_fixed_size_list() + .iter() + .flatten() + .try_for_each(|row| validate_map_keys(row.as_ref())), + data_type => exec_err!( + "Expected list, large_list or fixed_size_list, got {:?}", + data_type + ), + } +} + fn get_first_array_ref(columnar_value: &ColumnarValue) -> Result { match columnar_value { ColumnarValue::Scalar(value) => match value { @@ -310,7 +328,7 @@ fn make_map_batch_internal( doc_section(label = "Map Functions"), description = "Returns an Arrow map with the specified key-value pairs.\n\n\ The `make_map` function creates a map from two lists: one for keys and one for values. Each key must be unique and non-null.", - syntax_example = "map(key, value)\nmap(key: value)\nmake_map(['key1', 'key2'], ['value1', 'value2'])", + syntax_example = "map(key, value)\nmap {key: value}\nmake_map(['key1', 'key2'], ['value1', 'value2'])", sql_example = r#" ```sql -- Using map function @@ -399,7 +417,7 @@ impl ScalarUDFImpl for MapFunc { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_map_batch(&args.args) + make_map_batch(args.args, args.number_rows) } fn documentation(&self) -> Option<&Documentation> { @@ -716,10 +734,13 @@ mod tests { let values_array = Arc::new(value_builder.finish()); // Call make_map_batch - should succeed - let result = make_map_batch(&[ - ColumnarValue::Array(keys_array), - ColumnarValue::Array(values_array), - ]); + let result = make_map_batch( + vec![ + ColumnarValue::Array(keys_array), + ColumnarValue::Array(values_array), + ], + 3, + ); assert!(result.is_ok(), "Should handle NULL maps correctly"); @@ -764,10 +785,13 @@ mod tests { let values_array = Arc::new(value_builder.finish()); // Call make_map_batch - should fail - let result = make_map_batch(&[ - ColumnarValue::Array(keys_array), - ColumnarValue::Array(values_array), - ]); + let result = make_map_batch( + vec![ + ColumnarValue::Array(keys_array), + ColumnarValue::Array(values_array), + ], + 1, + ); assert!(result.is_err(), "Should reject null keys within maps"); @@ -812,10 +836,13 @@ mod tests { let values_array = Arc::new(value_builder.finish()); // Call make_map_batch - should succeed - let result = make_map_batch(&[ - ColumnarValue::Array(keys_array), - ColumnarValue::Array(values_array), - ]); + let result = make_map_batch( + vec![ + ColumnarValue::Array(keys_array), + ColumnarValue::Array(values_array), + ], + 2, + ); assert!( result.is_ok(), @@ -882,10 +909,13 @@ mod tests { let values_array = Arc::new(value_builder.finish()); // Call make_map_batch - should succeed - let result = make_map_batch(&[ - ColumnarValue::Array(keys_array), - ColumnarValue::Array(values_array), - ]); + let result = make_map_batch( + vec![ + ColumnarValue::Array(keys_array), + ColumnarValue::Array(values_array), + ], + 3, + ); assert!( result.is_ok(), diff --git a/datafusion/functions-nested/src/map_extract.rs b/datafusion/functions-nested/src/map_extract.rs index aab0d013a4152..40340ec2cf635 100644 --- a/datafusion/functions-nested/src/map_extract.rs +++ b/datafusion/functions-nested/src/map_extract.rs @@ -105,6 +105,11 @@ impl ScalarUDFImpl for MapExtract { fn return_type(&self, arg_types: &[DataType]) -> Result { let [map_type, _] = take_function_args(self.name(), arg_types)?; + + if map_type.is_null() { + return Ok(DataType::Null); + } + let map_fields = get_map_entry_field(map_type)?; Ok(DataType::List(Arc::new(Field::new_list_field( map_fields.last().unwrap().data_type().clone(), @@ -123,6 +128,10 @@ impl ScalarUDFImpl for MapExtract { fn coerce_types(&self, arg_types: &[DataType]) -> Result> { let [map_type, _] = take_function_args(self.name(), arg_types)?; + if map_type.is_null() { + return Ok(arg_types.to_vec()); + } + let field = get_map_entry_field(map_type)?; Ok(vec![ map_type.clone(), @@ -161,10 +170,10 @@ fn general_map_extract_inner( match value_index { Some(index) => { - mutable.extend(0, start + index, start + index + 1); + mutable.try_extend(0, start + index, start + index + 1)?; } None => { - mutable.extend_nulls(1); + mutable.try_extend_nulls(1)?; } } offsets.push(offsets[row_index] + 1); @@ -185,6 +194,7 @@ fn map_extract_inner(args: &[ArrayRef]) -> Result { let map_array = match map_arg.data_type() { DataType::Map(_, _) => as_map_array(&map_arg)?, + DataType::Null => return Ok(Arc::clone(map_arg)), _ => return exec_err!("The first argument in map_extract must be a map"), }; diff --git a/datafusion/functions-nested/src/position.rs b/datafusion/functions-nested/src/position.rs index d65620ede38e6..2a0134b4b96ff 100644 --- a/datafusion/functions-nested/src/position.rs +++ b/datafusion/functions-nested/src/position.rs @@ -56,7 +56,7 @@ make_udf_expr_and_func!( #[user_doc( doc_section(label = "Array Functions"), description = "Returns the position of the first occurrence of the specified element in the array, or NULL if not found. Comparisons are done using `IS DISTINCT FROM` semantics, so NULL is considered to match NULL.", - syntax_example = "array_position(array, element)\narray_position(array, element, index)", + syntax_example = "array_position(array, element[, index])", sql_example = r#"```sql > select array_position([1, 2, 2, 3, 1, 4], 2); +----------------------------------------------+ @@ -78,7 +78,7 @@ make_udf_expr_and_func!( argument(name = "element", description = "Element to search for in the array."), argument( name = "index", - description = "Index at which to start searching (1-indexed)." + description = "Index at which to start searching (1-indexed). Defaults to searching from the start" ) )] #[derive(Debug, PartialEq, Eq, Hash)] diff --git a/datafusion/functions-nested/src/remove.rs b/datafusion/functions-nested/src/remove.rs index d0f838ddad12a..491e823c9a21d 100644 --- a/datafusion/functions-nested/src/remove.rs +++ b/datafusion/functions-nested/src/remove.rs @@ -18,16 +18,17 @@ //! [`ScalarUDFImpl`] definitions for array_remove, array_remove_n, array_remove_all functions. use crate::utils; -use crate::utils::make_scalar_function; use arrow::array::{ - Array, ArrayRef, Capacities, GenericListArray, MutableArrayData, OffsetSizeTrait, - cast::AsArray, make_array, + Array, ArrayRef, Capacities, GenericListArray, MutableArrayData, NullBufferBuilder, + OffsetSizeTrait, Scalar, cast::AsArray, make_array, new_null_array, }; -use arrow::buffer::{NullBuffer, OffsetBuffer}; +use arrow::buffer::OffsetBuffer; use arrow::datatypes::{DataType, FieldRef}; use datafusion_common::cast::as_int64_array; use datafusion_common::utils::ListCoercion; -use datafusion_common::{Result, exec_err, internal_err, utils::take_function_args}; +use datafusion_common::{ + Result, ScalarValue, exec_err, internal_err, utils::take_function_args, +}; use datafusion_expr::{ ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, @@ -109,11 +110,31 @@ impl ScalarUDFImpl for ArrayRemove { &self, args: datafusion_expr::ReturnFieldArgs, ) -> Result { - Ok(Arc::clone(&args.arg_fields[0])) + let array_field = args.arg_fields[0].as_ref().clone(); + let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); + Ok(Arc::new(array_field.with_nullable(nullable))) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_remove_inner)(&args.args) + let [list_arg, element_arg] = take_function_args(self.name(), &args.args)?; + let num_rows = args.number_rows; + let list_array = list_arg.to_array(num_rows)?; + match element_arg { + ColumnarValue::Scalar(scalar_element) + if !scalar_element.is_null() + && !scalar_element.data_type().is_nested() => + { + let result = + array_remove_with_scalar_args(&list_array, scalar_element, 1i64)?; + Ok(ColumnarValue::Array(result)) + } + element_arg => { + let element_array = element_arg.to_array(num_rows)?; + let result = + array_remove_internal(&list_array, &element_array, &[Some(1)])?; + Ok(ColumnarValue::Array(result)) + } + } } fn aliases(&self) -> &[String] { @@ -210,11 +231,39 @@ impl ScalarUDFImpl for ArrayRemoveN { &self, args: datafusion_expr::ReturnFieldArgs, ) -> Result { - Ok(Arc::clone(&args.arg_fields[0])) + let array_field = args.arg_fields[0].as_ref().clone(); + let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); + Ok(Arc::new(array_field.with_nullable(nullable))) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_remove_n_inner)(&args.args) + let [list_arg, element_arg, max_arg] = + take_function_args(self.name(), &args.args)?; + let num_rows = args.number_rows; + let list_array = list_arg.to_array(num_rows)?; + match (element_arg, max_arg) { + ( + ColumnarValue::Scalar(scalar_element), + ColumnarValue::Scalar(scalar_max), + ) if !scalar_element.is_null() && !scalar_element.data_type().is_nested() => { + let ScalarValue::Int64(Some(n)) = scalar_max else { + return Ok(ColumnarValue::Array(new_null_array( + list_array.data_type(), + num_rows, + ))); + }; + let result = + array_remove_with_scalar_args(&list_array, scalar_element, *n)?; + Ok(ColumnarValue::Array(result)) + } + (element_arg, max_arg) => { + let element_array = element_arg.to_array(num_rows)?; + let max_array = max_arg.to_array(num_rows)?; + let arr_n = as_int64_array(&max_array)?.iter().collect::>(); + let result = array_remove_internal(&list_array, &element_array, &arr_n)?; + Ok(ColumnarValue::Array(result)) + } + } } fn aliases(&self) -> &[String] { @@ -300,11 +349,34 @@ impl ScalarUDFImpl for ArrayRemoveAll { &self, args: datafusion_expr::ReturnFieldArgs, ) -> Result { - Ok(Arc::clone(&args.arg_fields[0])) + let array_field = args.arg_fields[0].as_ref().clone(); + let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); + Ok(Arc::new(array_field.with_nullable(nullable))) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_remove_all_inner)(&args.args) + let [list_arg, element_arg] = take_function_args(self.name(), &args.args)?; + let num_rows = args.number_rows; + let list_array = list_arg.to_array(num_rows)?; + match element_arg { + ColumnarValue::Scalar(scalar_element) + if !scalar_element.is_null() + && !scalar_element.data_type().is_nested() => + { + let result = + array_remove_with_scalar_args(&list_array, scalar_element, i64::MAX)?; + Ok(ColumnarValue::Array(result)) + } + element_arg => { + let element_array = element_arg.to_array(num_rows)?; + let result = array_remove_internal( + &list_array, + &element_array, + &[Some(i64::MAX)], + )?; + Ok(ColumnarValue::Array(result)) + } + } } fn aliases(&self) -> &[String] { @@ -316,31 +388,10 @@ impl ScalarUDFImpl for ArrayRemoveAll { } } -fn array_remove_inner(args: &[ArrayRef]) -> Result { - let [array, element] = take_function_args("array_remove", args)?; - - let arr_n = vec![1; array.len()]; - array_remove_internal(array, element, &arr_n) -} - -fn array_remove_n_inner(args: &[ArrayRef]) -> Result { - let [array, element, max] = take_function_args("array_remove_n", args)?; - - let arr_n = as_int64_array(max)?.values().to_vec(); - array_remove_internal(array, element, &arr_n) -} - -fn array_remove_all_inner(args: &[ArrayRef]) -> Result { - let [array, element] = take_function_args("array_remove_all", args)?; - - let arr_n = vec![i64::MAX; array.len()]; - array_remove_internal(array, element, &arr_n) -} - fn array_remove_internal( array: &ArrayRef, element_array: &ArrayRef, - arr_n: &[i64], + arr_n: &[Option], ) -> Result { match array.data_type() { DataType::List(_) => { @@ -351,12 +402,36 @@ fn array_remove_internal( let list_array = array.as_list::(); general_remove::(list_array, element_array, arr_n) } + DataType::Null => Ok(new_null_array(array.data_type(), array.len())), array_type => { exec_err!("array_remove_all does not support type '{array_type}'.") } } } +/// Fast path for `array_remove` when the needle is a non-null, non-nested scalar. +/// Dispatches to the bulk `not_distinct` comparison kernel. +fn array_remove_with_scalar_args( + array: &ArrayRef, + scalar_needle: &ScalarValue, + max_removals: i64, +) -> Result { + match array.data_type() { + DataType::List(_) => { + let list_array = array.as_list::(); + general_remove_with_scalar::(list_array, scalar_needle, max_removals) + } + DataType::LargeList(_) => { + let list_array = array.as_list::(); + general_remove_with_scalar::(list_array, scalar_needle, max_removals) + } + DataType::Null => Ok(new_null_array(array.data_type(), array.len())), + array_type => exec_err!( + "array_remove/array_remove_n/array_remove_all does not support type '{array_type}'." + ), + } +} + /// For each element of `list_array[i]`, removed up to `arr_n[i]` occurrences /// of `element_array[i]`. /// @@ -377,7 +452,7 @@ fn array_remove_internal( fn general_remove( list_array: &GenericListArray, element_array: &ArrayRef, - arr_n: &[i64], + arr_n: &[Option], ) -> Result { let list_field = match list_array.data_type() { DataType::List(field) | DataType::LargeList(field) => field, @@ -390,7 +465,7 @@ fn general_remove( }; let original_data = list_array.values().to_data(); // Build up the offsets for the final output array - let mut offsets = Vec::::with_capacity(arr_n.len() + 1); + let mut offsets = Vec::::with_capacity(list_array.len() + 1); offsets.push(OffsetSize::zero()); let mut mutable = MutableArrayData::with_capacities( @@ -398,20 +473,28 @@ fn general_remove( false, Capacities::Array(original_data.len()), ); - - // Pre-compute combined null bitmap - let nulls = NullBuffer::union(list_array.nulls(), element_array.nulls()); + let mut valid = NullBufferBuilder::new(list_array.len()); for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() { - if nulls.as_ref().is_some_and(|nulls| nulls.is_null(row_index)) { + if list_array.is_null(row_index) || element_array.is_null(row_index) { offsets.push(offsets[row_index]); + valid.append_null(); continue; } + let n = if arr_n.len() == 1 { + arr_n[0] + } else { + arr_n[row_index] + }; + let Some(n) = n else { + offsets.push(offsets[row_index]); + valid.append_null(); + continue; + }; + let start = offset_window[0].to_usize().unwrap(); let end = offset_window[1].to_usize().unwrap(); - // n is the number of elements to remove in this row - let n = arr_n[row_index]; // compare each element in the list, `false` means the element matches and should be removed let eq_array = utils::compare_element_to_list( @@ -425,8 +508,9 @@ fn general_remove( // Fast path: no elements to remove, copy entire row if num_to_remove == 0 { - mutable.extend(0, start, end); + mutable.try_extend(0, start, end)?; offsets.push(offsets[row_index] + OffsetSize::usize_as(end - start)); + valid.append_non_null(); continue; } @@ -440,7 +524,7 @@ fn general_remove( if keep == Some(false) && removed < max_removals { // Flush pending batch before skipping this element if let Some(bs) = pending_batch_to_retain { - mutable.extend(0, start + bs, start + i); + mutable.try_extend(0, start + bs, start + i)?; copied += i - bs; pending_batch_to_retain = None; } @@ -452,10 +536,111 @@ fn general_remove( // Flush remaining batch if let Some(bs) = pending_batch_to_retain { - mutable.extend(0, start + bs, start + eq_array.len()); + mutable.try_extend(0, start + bs, start + eq_array.len())?; copied += eq_array.len() - bs; } + offsets.push(offsets[row_index] + OffsetSize::usize_as(copied)); + valid.append_non_null(); + } + + let new_values = make_array(mutable.freeze()); + Ok(Arc::new(GenericListArray::::try_new( + Arc::clone(list_field), + OffsetBuffer::new(offsets.into()), + new_values, + valid.finish(), + )?)) +} + +/// For each element of `list_array[i]`, removes up to `max_removals` occurrences +/// of the scalar needle. +/// +/// This is a specialized version of `general_remove` for scalar elements that +/// uses bulk comparison for better performance. +fn general_remove_with_scalar( + list_array: &GenericListArray, + scalar_needle: &ScalarValue, + max_removals: i64, +) -> Result { + if max_removals <= 0 { + return Ok(Arc::new(list_array.clone())); + } + + let list_field = match list_array.data_type() { + DataType::List(field) | DataType::LargeList(field) => field, + _ => { + return exec_err!( + "Expected List or LargeList data type, got {:?}", + list_array.data_type() + ); + } + }; + + let list_offsets = list_array.offsets(); + let first_offset = list_offsets[0].to_usize().unwrap(); + let last_offset = list_offsets[list_offsets.len() - 1].to_usize().unwrap(); + let values_range_len = last_offset - first_offset; + let values_slice = list_array.values().slice(first_offset, values_range_len); + let original_data = values_slice.to_data(); + let mut offsets = Vec::::with_capacity(list_array.len() + 1); + offsets.push(OffsetSize::zero()); + + let mut mutable = MutableArrayData::with_capacities( + vec![&original_data], + false, + Capacities::Array(original_data.len()), + ); + let nulls = list_array.nulls().cloned(); + let needle = scalar_needle.to_array_of_size(1)?; + let remove_mask = arrow_ord::cmp::not_distinct(&values_slice, &Scalar::new(needle))?; + let remove_bits = remove_mask.values(); + + for (row_index, offset_window) in list_offsets.windows(2).enumerate() { + if nulls.as_ref().is_some_and(|nulls| nulls.is_null(row_index)) { + offsets.push(offsets[row_index]); + continue; + } + + let start = offset_window[0].to_usize().unwrap() - first_offset; + let end = offset_window[1].to_usize().unwrap() - first_offset; + let row_len = end - start; + + let row_remove_bits = remove_bits.slice(start, row_len); + let num_to_remove = row_remove_bits.count_set_bits(); + + if num_to_remove == 0 { + mutable.try_extend(0, start, end)?; + offsets.push(offsets[row_index] + OffsetSize::usize_as(row_len)); + continue; + } + + let removals_to_apply = max_removals.min(num_to_remove as i64) as usize; + + // Iterate only over the removal positions via set_indices. This is + // efficient when the number of removals is small relative to the row + // length (common case), since it skips over retained elements. + let mut removed = 0usize; + let mut copied = 0usize; + let mut prev_end = start; + for remove_pos in row_remove_bits.set_indices() { + let abs_pos = start + remove_pos; + if abs_pos > prev_end { + mutable.try_extend(0, prev_end, abs_pos)?; + copied += abs_pos - prev_end; + } + prev_end = abs_pos + 1; + removed += 1; + if removed == removals_to_apply { + break; + } + } + // Copy the remaining tail after the last removal + if prev_end < end { + mutable.try_extend(0, prev_end, end)?; + copied += end - prev_end; + } + offsets.push(offsets[row_index] + OffsetSize::usize_as(copied)); } @@ -472,8 +657,10 @@ fn general_remove( mod tests { use crate::remove::{ArrayRemove, ArrayRemoveAll, ArrayRemoveN}; use arrow::array::{ - Array, ArrayRef, AsArray, GenericListArray, ListArray, OffsetSizeTrait, + Array, ArrayRef, AsArray, GenericListArray, Int32Array, Int64Array, ListArray, + OffsetSizeTrait, }; + use arrow::buffer::{NullBuffer, ScalarBuffer}; use arrow::datatypes::{DataType, Field, Int32Type}; use datafusion_common::ScalarValue; use datafusion_expr::{ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl}; @@ -485,25 +672,34 @@ mod tests { fn test_array_remove_nullability() { for nullability in [true, false] { for item_nullability in [true, false] { - let input_field = Arc::new(Field::new( - "num", - DataType::new_list(DataType::Int32, item_nullability), - nullability, - )); - let args_fields = vec![ - Arc::clone(&input_field), - Arc::new(Field::new("a", DataType::Int32, false)), - ]; - let scalar_args = vec![None, Some(&ScalarValue::Int32(Some(1)))]; - - let result = ArrayRemove::new() - .return_field_from_args(ReturnFieldArgs { - arg_fields: &args_fields, - scalar_arguments: &scalar_args, - }) - .unwrap(); - - assert_eq!(result, input_field); + for element_nullability in [true, false] { + let input_field = Arc::new(Field::new( + "num", + DataType::new_list(DataType::Int32, item_nullability), + nullability, + )); + let args_fields = vec![ + Arc::clone(&input_field), + Arc::new(Field::new("a", DataType::Int32, element_nullability)), + ]; + let scalar_args = vec![None, Some(&ScalarValue::Int32(Some(1)))]; + + let result = ArrayRemove::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &args_fields, + scalar_arguments: &scalar_args, + }) + .unwrap(); + + let expected = Arc::new( + input_field + .as_ref() + .clone() + .with_nullable(nullability || element_nullability), + ); + + assert_eq!(result, expected); + } } } } @@ -512,30 +708,47 @@ mod tests { fn test_array_remove_n_nullability() { for nullability in [true, false] { for item_nullability in [true, false] { - let input_field = Arc::new(Field::new( - "num", - DataType::new_list(DataType::Int32, item_nullability), - nullability, - )); - let args_fields = vec![ - Arc::clone(&input_field), - Arc::new(Field::new("a", DataType::Int32, false)), - Arc::new(Field::new("b", DataType::Int64, false)), - ]; - let scalar_args = vec![ - None, - Some(&ScalarValue::Int32(Some(1))), - Some(&ScalarValue::Int64(Some(1))), - ]; - - let result = ArrayRemoveN::new() - .return_field_from_args(ReturnFieldArgs { - arg_fields: &args_fields, - scalar_arguments: &scalar_args, - }) - .unwrap(); - - assert_eq!(result, input_field); + for element_nullability in [true, false] { + for count_nullability in [true, false] { + let input_field = Arc::new(Field::new( + "num", + DataType::new_list(DataType::Int32, item_nullability), + nullability, + )); + let args_fields = vec![ + Arc::clone(&input_field), + Arc::new(Field::new( + "a", + DataType::Int32, + element_nullability, + )), + Arc::new(Field::new("b", DataType::Int64, count_nullability)), + ]; + let scalar_args = vec![ + None, + Some(&ScalarValue::Int32(Some(1))), + Some(&ScalarValue::Int64(Some(1))), + ]; + + let result = ArrayRemoveN::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &args_fields, + scalar_arguments: &scalar_args, + }) + .unwrap(); + + let expected_nullable = + nullability || element_nullability || count_nullability; + let expected = Arc::new( + input_field + .as_ref() + .clone() + .with_nullable(expected_nullable), + ); + + assert_eq!(result, expected); + } + } } } } @@ -544,19 +757,33 @@ mod tests { fn test_array_remove_all_nullability() { for nullability in [true, false] { for item_nullability in [true, false] { - let input_field = Arc::new(Field::new( - "num", - DataType::new_list(DataType::Int32, item_nullability), - nullability, - )); - let result = ArrayRemoveAll::new() - .return_field_from_args(ReturnFieldArgs { - arg_fields: &[Arc::clone(&input_field)], - scalar_arguments: &[None], - }) - .unwrap(); - - assert_eq!(result, input_field); + for element_nullability in [true, false] { + let input_field = Arc::new(Field::new( + "num", + DataType::new_list(DataType::Int32, item_nullability), + nullability, + )); + let args_fields = vec![ + Arc::clone(&input_field), + Arc::new(Field::new("a", DataType::Int32, element_nullability)), + ]; + let scalar_args = vec![None, Some(&ScalarValue::Int32(Some(1)))]; + let result = ArrayRemoveAll::new() + .return_field_from_args(ReturnFieldArgs { + arg_fields: &args_fields, + scalar_arguments: &scalar_args, + }) + .unwrap(); + + let expected = Arc::new( + input_field + .as_ref() + .clone() + .with_nullable(nullability || element_nullability), + ); + + assert_eq!(result, expected); + } } } } @@ -734,6 +961,58 @@ mod tests { assert_array_remove_n(input_list, expected_list, element_to_remove, 2); } + #[test] + fn test_array_remove_n_null_count_returns_null() { + let array: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(2)]), + Some(vec![Some(4), Some(2)]), + ])); + let element: ArrayRef = Arc::new(Int32Array::from(vec![2, 2])); + let max: ArrayRef = Arc::new(Int64Array::new( + ScalarBuffer::from(vec![1, 1]), + Some(NullBuffer::from(vec![true, false])), + )); + + let udf = ArrayRemoveN::new(); + let args_fields = vec![ + Arc::new(Field::new("num", array.data_type().clone(), false)), + Arc::new(Field::new("el", DataType::Int32, false)), + Arc::new(Field::new("count", DataType::Int64, true)), + ]; + let scalar_args = vec![None, None, None]; + let return_field = udf + .return_field_from_args(ReturnFieldArgs { + arg_fields: &args_fields, + scalar_arguments: &scalar_args, + }) + .unwrap(); + let result = udf + .invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(array), + ColumnarValue::Array(element), + ColumnarValue::Array(max), + ], + arg_fields: args_fields, + number_rows: 2, + return_field, + config_options: Arc::new(Default::default()), + }) + .unwrap(); + let expected = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + None, + ]); + + match result { + ColumnarValue::Array(array) => { + assert_eq!(array.as_list::(), &expected); + } + _ => panic!("Expected ColumnarValue::Array"), + } + } + fn assert_array_remove_n( input_list: ArrayRef, expected_list: GenericListArray, diff --git a/datafusion/functions-nested/src/repeat.rs b/datafusion/functions-nested/src/repeat.rs index ceec748a6e776..d7dff21141429 100644 --- a/datafusion/functions-nested/src/repeat.rs +++ b/datafusion/functions-nested/src/repeat.rs @@ -31,15 +31,19 @@ use arrow::datatypes::{ }; use datafusion_common::cast::{as_int64_array, as_large_list_array, as_list_array}; use datafusion_common::types::{NativeType, logical_int64}; -use datafusion_common::{DataFusionError, Result}; +use datafusion_common::{Result, exec_datafusion_err}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; use datafusion_expr_common::signature::{Coercion, TypeSignatureClass}; use datafusion_macros::user_doc; +use std::mem::size_of; use std::sync::Arc; +const ARRAY_REPEAT_LENGTH_EXCEEDED: &str = + "array_repeat: requested length exceeds maximum array size"; + make_udf_expr_and_func!( ArrayRepeat, array_repeat, @@ -175,28 +179,25 @@ fn general_repeat( array: &ArrayRef, count_array: &Int64Array, ) -> Result { - let total_repeated_values: usize = (0..count_array.len()) - .map(|i| get_count_with_validity(count_array, i)) - .sum(); + let total_repeated_values = + (0..count_array.len()).try_fold(0usize, |total, idx| { + total + .checked_add(repeat_count(count_array, idx).unwrap_or_default()) + .ok_or_else(|| { + exec_datafusion_err!( + "array_repeat: total repeated values overflowed usize" + ) + }) + })?; + ensure_repeated_values_fit::(total_repeated_values)?; + let (offsets, _) = build_repeat_offsets::(count_array)?; let mut take_indices = Vec::with_capacity(total_repeated_values); - let mut offsets = Vec::with_capacity(count_array.len() + 1); - offsets.push(O::zero()); - let mut running_offset = 0usize; for idx in 0..count_array.len() { - let count = get_count_with_validity(count_array, idx); - running_offset = running_offset.checked_add(count).ok_or_else(|| { - DataFusionError::Execution( - "array_repeat: running_offset overflowed usize".to_string(), - ) - })?; - let offset = O::from_usize(running_offset).ok_or_else(|| { - DataFusionError::Execution(format!( - "array_repeat: offset {running_offset} exceeds the maximum value for offset type" - )) - })?; - offsets.push(offset); + let Some(count) = repeat_count(count_array, idx) else { + continue; + }; take_indices.extend(std::iter::repeat_n(idx as u64, count)); } @@ -231,47 +232,44 @@ fn general_list_repeat( count_array: &Int64Array, ) -> Result { let list_offsets = list_array.value_offsets(); + let (outer_offsets, outer_total) = build_repeat_offsets::(count_array)?; // calculate capacities for pre-allocation - let mut outer_total = 0usize; let mut inner_total = 0usize; for i in 0..count_array.len() { - let count = get_count_with_validity(count_array, i); - if count > 0 { - outer_total += count; - if list_array.is_valid(i) { - let len = list_offsets[i + 1].to_usize().unwrap() - - list_offsets[i].to_usize().unwrap(); - inner_total += len * count; - } + let Some(count) = repeat_count(count_array, i) else { + continue; + }; + if count > 0 && list_array.is_valid(i) { + let len = list_offsets[i + 1].to_usize().unwrap() + - list_offsets[i].to_usize().unwrap(); + inner_total = + checked_repeat_len_add(inner_total, checked_repeat_len_mul(len, count)?)?; + ensure_repeated_values_fit::(inner_total)?; } } // Build inner structures - let mut inner_offsets = Vec::with_capacity(outer_total + 1); + let inner_offsets_capacity = checked_offset_slots_capacity::(outer_total)?; + let mut inner_offsets = Vec::with_capacity(inner_offsets_capacity); let mut take_indices = Vec::with_capacity(inner_total); let mut inner_nulls = BooleanBufferBuilder::new(outer_total); let mut inner_running = 0usize; inner_offsets.push(O::zero()); for row_idx in 0..count_array.len() { - let count = get_count_with_validity(count_array, row_idx); + let Some(count) = repeat_count(count_array, row_idx) else { + continue; + }; let list_is_valid = list_array.is_valid(row_idx); let start = list_offsets[row_idx].to_usize().unwrap(); let end = list_offsets[row_idx + 1].to_usize().unwrap(); let row_len = end - start; for _ in 0..count { - inner_running = inner_running.checked_add(row_len).ok_or_else(|| { - DataFusionError::Execution( - "array_repeat: inner offset overflowed usize".to_string(), - ) - })?; - let offset = O::from_usize(inner_running).ok_or_else(|| { - DataFusionError::Execution(format!( - "array_repeat: offset {inner_running} exceeds the maximum value for offset type" - )) - })?; + inner_running = checked_repeat_len_add(inner_running, row_len)?; + ensure_repeated_values_fit::(inner_running)?; + let offset = checked_repeat_offset::(inner_running)?; inner_offsets.push(offset); inner_nulls.append(list_is_valid); if list_is_valid { @@ -293,30 +291,185 @@ fn general_list_repeat( Some(NullBuffer::new(inner_nulls.finish())), )?; - // Build outer ListArray Ok(Arc::new(GenericListArray::::try_new( Arc::new(Field::new_list_field( list_array.data_type().to_owned(), true, )), - OffsetBuffer::::from_lengths( - count_array - .iter() - .map(|c| c.map(|v| if v > 0 { v as usize } else { 0 }).unwrap_or(0)), - ), + OffsetBuffer::new(outer_offsets.into()), Arc::new(inner_list), count_array.nulls().cloned(), )?)) } -/// Helper function to get count from count_array at given index -/// Return 0 for null values or non-positive count. +fn build_repeat_offsets( + count_array: &Int64Array, +) -> Result<(Vec, usize)> { + let offsets_capacity = checked_offset_slots_capacity::(count_array.len())?; + let mut offsets = Vec::with_capacity(offsets_capacity); + offsets.push(O::zero()); + let mut running_offset = 0usize; + + for idx in 0..count_array.len() { + let Some(count) = repeat_count(count_array, idx) else { + offsets.push(*offsets.last().unwrap()); + continue; + }; + running_offset = checked_repeat_len_add(running_offset, count)?; + ensure_repeated_values_fit::(running_offset)?; + let offset = checked_repeat_offset::(running_offset)?; + offsets.push(offset); + } + + Ok((offsets, running_offset)) +} + +fn checked_repeat_len_add(lhs: usize, rhs: usize) -> Result { + lhs.checked_add(rhs) + .ok_or_else(|| exec_datafusion_err!("{}", ARRAY_REPEAT_LENGTH_EXCEEDED)) +} + +fn checked_repeat_len_mul(lhs: usize, rhs: usize) -> Result { + lhs.checked_mul(rhs) + .ok_or_else(|| exec_datafusion_err!("{}", ARRAY_REPEAT_LENGTH_EXCEEDED)) +} + +fn ensure_repeated_values_fit(len: usize) -> Result<()> { + ensure_vec_capacity::(len)?; + checked_repeat_offset::(len)?; + + Ok(()) +} + +fn ensure_vec_capacity(len: usize) -> Result<()> { + if len > max_vec_elements::() { + return Err(exec_datafusion_err!("{}", ARRAY_REPEAT_LENGTH_EXCEEDED)); + } + + Ok(()) +} + +fn checked_offset_slots_capacity(len: usize) -> Result { + let capacity = checked_repeat_len_add(len, 1)?; + ensure_vec_capacity::(capacity)?; + + Ok(capacity) +} + +fn checked_repeat_offset(offset: usize) -> Result { + O::from_usize(offset).ok_or_else(|| { + exec_datafusion_err!( + "array_repeat: offset {offset} exceeds the maximum value for offset type" + ) + }) +} + +fn max_vec_elements() -> usize { + let element_size = size_of::(); + (isize::MAX as usize) + .checked_div(element_size) + .unwrap_or(usize::MAX) +} + +/// Helper function to get count from count_array at given index. +/// Returns `None` for NULL values and `Some(0)` for non-positive counts. #[inline] -fn get_count_with_validity(count_array: &Int64Array, idx: usize) -> usize { +fn repeat_count(count_array: &Int64Array, idx: usize) -> Option { if count_array.is_null(idx) { - 0 + None } else { let c = count_array.value(idx); - if c > 0 { c as usize } else { 0 } + Some(if c > 0 { c as usize } else { 0 }) + } +} + +#[cfg(test)] +mod tests { + use super::{array_repeat_inner, general_list_repeat, general_repeat}; + use arrow::array::{Array, ArrayRef, AsArray, Int32Array, Int64Array, ListArray}; + use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow::datatypes::{Field, Int32Type}; + use datafusion_common::Result; + use std::sync::Arc; + + #[test] + fn test_array_repeat_null_count_stays_null() -> Result<()> { + let array: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3])); + let counts = Int64Array::new( + ScalarBuffer::from(vec![2, 1, 1]), + Some(NullBuffer::from(vec![true, false, true])), + ); + + let result = general_repeat::(&array, &counts)?; + let expected = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(1)]), + None, + Some(vec![Some(3)]), + ]); + + assert_eq!(result.as_list::(), &expected); + + Ok(()) + } + + #[test] + fn test_array_repeat_nested_null_count_stays_null() -> Result<()> { + let list_array = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(3), Some(4)]), + Some(vec![Some(5)]), + ]); + let counts = Int64Array::new( + ScalarBuffer::from(vec![2, 1, 1]), + Some(NullBuffer::from(vec![true, false, true])), + ); + + let result = general_list_repeat::(&list_array, &counts)?; + let repeated_values = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![Some(1), Some(2)]), + Some(vec![Some(5)]), + ]); + let expected = ListArray::new( + Arc::new(Field::new_list_field( + repeated_values.data_type().clone(), + true, + )), + OffsetBuffer::new(ScalarBuffer::from(vec![0, 2, 2, 3])), + Arc::new(repeated_values), + Some(NullBuffer::from(vec![true, false, true])), + ); + + assert_eq!(result.as_list::(), &expected); + + Ok(()) + } + + #[test] + fn scalar_count_exceeding_max_array_size_returns_error() { + let element: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let count: ArrayRef = Arc::new(Int64Array::from(vec![i64::MAX])); + + let err = array_repeat_inner(&[element, count]).unwrap_err(); + assert!( + err.to_string().starts_with( + "Execution error: array_repeat: requested length exceeds maximum array size" + ), + "unexpected error: {err}" + ); + } + + #[test] + fn scalar_count_exceeding_list_offset_limit_returns_error() { + let element: ArrayRef = Arc::new(Int64Array::from(vec![1])); + let count: ArrayRef = Arc::new(Int64Array::from(vec![i32::MAX as i64 + 1])); + + let err = array_repeat_inner(&[element, count]).unwrap_err(); + assert!( + err.to_string().starts_with( + "Execution error: array_repeat: offset 2147483648 exceeds the maximum value for offset type" + ), + "unexpected error: {err}" + ); } } diff --git a/datafusion/functions-nested/src/replace.rs b/datafusion/functions-nested/src/replace.rs index a9a53a3cb989f..4bfd0c0dbecfe 100644 --- a/datafusion/functions-nested/src/replace.rs +++ b/datafusion/functions-nested/src/replace.rs @@ -19,22 +19,23 @@ use arrow::array::{ Array, ArrayRef, AsArray, Capacities, GenericListArray, MutableArrayData, - NullBufferBuilder, OffsetSizeTrait, new_null_array, + NullBufferBuilder, OffsetSizeTrait, Scalar, new_null_array, }; -use arrow::datatypes::{DataType, Field}; - use arrow::buffer::OffsetBuffer; +use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::cast::as_int64_array; use datafusion_common::utils::ListCoercion; -use datafusion_common::{Result, exec_err, utils::take_function_args}; +use datafusion_common::{ + Result, ScalarValue, exec_err, internal_err, utils::take_function_args, +}; use datafusion_expr::{ ArrayFunctionArgument, ArrayFunctionSignature, ColumnarValue, Documentation, - ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, Volatility, + ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, + Volatility, }; use datafusion_macros::user_doc; -use crate::utils::compare_element_to_list; -use crate::utils::make_scalar_function; +use crate::utils::{compare_element_to_list, list_inner_field, list_type_with_element}; use std::sync::Arc; @@ -120,12 +121,45 @@ impl ScalarUDFImpl for ArrayReplace { &self.signature } - fn return_type(&self, args: &[DataType]) -> Result { - Ok(args[0].clone()) + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + replace_return_field(self.name(), args.arg_fields) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_replace_inner)(&args.args) + let return_type = args.return_field.data_type().clone(); + let [list_arg, from_arg, to_arg] = take_function_args(self.name(), &args.args)?; + let num_rows = args.number_rows; + let list_array = list_arg.to_array(num_rows)?; + match (from_arg, to_arg) { + (ColumnarValue::Scalar(scalar_from), ColumnarValue::Scalar(scalar_to)) => { + let result = array_replace_with_scalar_args( + self.name(), + &list_array, + scalar_from, + scalar_to, + 1i64, + &return_type, + )?; + Ok(ColumnarValue::Array(result)) + } + (from_arg, to_arg) => { + let from_array = from_arg.to_array(num_rows)?; + let to_array = to_arg.to_array(num_rows)?; + let result = array_replace_internal( + self.name(), + &list_array, + &from_array, + &to_array, + &[Some(1)], + &return_type, + )?; + Ok(ColumnarValue::Array(result)) + } + } } fn aliases(&self) -> &[String] { @@ -195,12 +229,57 @@ impl ScalarUDFImpl for ArrayReplaceN { &self.signature } - fn return_type(&self, args: &[DataType]) -> Result { - Ok(args[0].clone()) + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + replace_return_field(self.name(), args.arg_fields) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_replace_n_inner)(&args.args) + let return_type = args.return_field.data_type().clone(); + let [list_arg, from_arg, to_arg, max_arg] = + take_function_args(self.name(), &args.args)?; + let num_rows = args.number_rows; + let list_array = list_arg.to_array(num_rows)?; + match (from_arg, to_arg, max_arg) { + ( + ColumnarValue::Scalar(scalar_from), + ColumnarValue::Scalar(scalar_to), + ColumnarValue::Scalar(scalar_max), + ) => { + let ScalarValue::Int64(Some(n)) = scalar_max else { + return Ok(ColumnarValue::Array(new_null_array( + &return_type, + num_rows, + ))); + }; + let result = array_replace_with_scalar_args( + self.name(), + &list_array, + scalar_from, + scalar_to, + *n, + &return_type, + )?; + Ok(ColumnarValue::Array(result)) + } + (from_arg, to_arg, max_arg) => { + let from_array = from_arg.to_array(num_rows)?; + let to_array = to_arg.to_array(num_rows)?; + let max_array = max_arg.to_array(num_rows)?; + let result = array_replace_n_inner( + self.name(), + &list_array, + &from_array, + &to_array, + &max_array, + &return_type, + )?; + Ok(ColumnarValue::Array(result)) + } + } } fn aliases(&self) -> &[String] { @@ -268,12 +347,45 @@ impl ScalarUDFImpl for ArrayReplaceAll { &self.signature } - fn return_type(&self, args: &[DataType]) -> Result { - Ok(args[0].clone()) + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + replace_return_field(self.name(), args.arg_fields) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(array_replace_all_inner)(&args.args) + let return_type = args.return_field.data_type().clone(); + let [list_arg, from_arg, to_arg] = take_function_args(self.name(), &args.args)?; + let num_rows = args.number_rows; + let list_array = list_arg.to_array(num_rows)?; + match (from_arg, to_arg) { + (ColumnarValue::Scalar(scalar_from), ColumnarValue::Scalar(scalar_to)) => { + let result = array_replace_with_scalar_args( + self.name(), + &list_array, + scalar_from, + scalar_to, + i64::MAX, + &return_type, + )?; + Ok(ColumnarValue::Array(result)) + } + (from_arg, to_arg) => { + let from_array = from_arg.to_array(num_rows)?; + let to_array = to_arg.to_array(num_rows)?; + let result = array_replace_internal( + self.name(), + &list_array, + &from_array, + &to_array, + &[Some(i64::MAX)], + &return_type, + )?; + Ok(ColumnarValue::Array(result)) + } + } } fn aliases(&self) -> &[String] { @@ -285,6 +397,24 @@ impl ScalarUDFImpl for ArrayReplaceAll { } } +/// Return field shared by `array_replace`, `array_replace_n` and +/// `array_replace_all`: the input list type, except that its inner field is +/// nullable whenever the replacement element may be null. +fn replace_return_field(name: &str, arg_fields: &[FieldRef]) -> Result { + // `array` is at index 0 and `to` at index 2 for all three functions. + // `from` never contributes values to the output, so `to` is the only + // argument besides `array` that can affect the output's type. + let [array_field, _from_field, to_field, ..] = arg_fields else { + return exec_err!( + "{name} expects at least 3 arguments, got {}", + arg_fields.len() + ); + }; + let data_type = + list_type_with_element(array_field.data_type(), to_field.is_nullable()); + Ok(Arc::new(Field::new(name, data_type, true))) +} + /// For each element of `list_array[i]`, replaces up to `arr_n[i]` occurrences /// of `from_array[i]`, `to_array[i]`. /// @@ -306,10 +436,12 @@ fn general_replace( list_array: &GenericListArray, from_array: &ArrayRef, to_array: &ArrayRef, - arr_n: &[i64], + arr_n: &[Option], + field: FieldRef, ) -> Result { // Build up the offsets for the final output array - let mut offsets: Vec = vec![O::usize_as(0)]; + let mut offsets: Vec = Vec::with_capacity(list_array.len() + 1); + offsets.push(O::usize_as(0)); let values = list_array.values(); let original_data = values.to_data(); let to_data = to_array.to_data(); @@ -331,6 +463,17 @@ fn general_replace( continue; } + let n = if arr_n.len() == 1 { + arr_n[0] + } else { + arr_n[row_index] + }; + let Some(n) = n else { + offsets.push(offsets[row_index]); + valid.append_null(); + continue; + }; + let start = offset_window[0]; let end = offset_window[1]; @@ -343,16 +486,15 @@ fn general_replace( let original_idx = O::usize_as(0); let replace_idx = O::usize_as(1); - let n = arr_n[row_index]; let mut counter = 0; // All elements are false, no need to replace, just copy original data if n <= 0 || !eq_array.has_true() { - mutable.extend( + mutable.try_extend( original_idx.to_usize().unwrap(), start.to_usize().unwrap(), end.to_usize().unwrap(), - ); + )?; offsets.push(offsets[row_index] + (end - start)); valid.append_non_null(); continue; @@ -364,21 +506,25 @@ fn general_replace( if to_replace == Some(true) && counter < n { // Flush any pending retain run before emitting the replacement. if let Some(rs) = pending_retain.take() { - mutable.extend( + mutable.try_extend( original_idx.to_usize().unwrap(), (start + rs).to_usize().unwrap(), (start + i).to_usize().unwrap(), - ); + )?; } - mutable.extend(replace_idx.to_usize().unwrap(), row_index, row_index + 1); + mutable.try_extend( + replace_idx.to_usize().unwrap(), + row_index, + row_index + 1, + )?; counter += 1; if counter == n { // copy original data for any matches past n - mutable.extend( + mutable.try_extend( original_idx.to_usize().unwrap(), (start + i).to_usize().unwrap() + 1, end.to_usize().unwrap(), - ); + )?; break; } } else if pending_retain.is_none() { @@ -391,11 +537,11 @@ fn general_replace( if counter < n && let Some(rs) = pending_retain { - mutable.extend( + mutable.try_extend( original_idx.to_usize().unwrap(), (start + rs).to_usize().unwrap(), end.to_usize().unwrap(), - ); + )?; } offsets.push(offsets[row_index] + (end - start)); @@ -405,70 +551,280 @@ fn general_replace( let data = mutable.freeze(); Ok(Arc::new(GenericListArray::::try_new( - Arc::new(Field::new_list_field(list_array.value_type(), true)), + field, OffsetBuffer::::new(offsets.into()), arrow::array::make_array(data), valid.finish(), )?)) } -fn array_replace_inner(args: &[ArrayRef]) -> Result { - let [array, from, to] = take_function_args("array_replace", args)?; +/// Replaces up to `max_replacements` occurrences of `needle` with the single +/// element in `to_array` for each row in `list_array`. +/// +/// This is a specialized fast path for the all-scalar case that uses a single +/// bulk `not_distinct` comparison over only the visible values range, then +/// iterates match positions via `set_indices` instead of scanning every bit. +fn general_replace_with_scalar( + list_array: &GenericListArray, + needle: &Scalar, + scalar_to: &ScalarValue, + max_replacements: i64, + field: FieldRef, +) -> Result { + // No replacement needed, but the output still has to carry the promised + // field, which may be more nullable than the input's. + if max_replacements <= 0 { + return Ok(Arc::new(GenericListArray::::try_new( + field, + list_array.offsets().clone(), + Arc::clone(list_array.values()), + list_array.nulls().cloned(), + )?)); + } + + let first_offset = list_array.offsets()[0].to_usize().unwrap(); + let last_offset = list_array.offsets()[list_array.len()].to_usize().unwrap(); + let visible_values = list_array + .values() + .slice(first_offset, last_offset - first_offset); - // replace at most one occurrence for each element - let arr_n = vec![1; array.len()]; - match array.data_type() { - DataType::List(_) => { - let list_array = array.as_list::(); - general_replace::(list_array, from, to, &arr_n) + let to_array = scalar_to.to_array_of_size(1)?; + let original_data = visible_values.to_data(); + let to_data = to_array.to_data(); + let capacity = Capacities::Array(original_data.len()); + + let mut mutable = MutableArrayData::with_capacities( + vec![&original_data, &to_data], + false, + capacity, + ); + + let mut offsets = Vec::::with_capacity(list_array.len() + 1); + offsets.push(O::zero()); + + // Single bulk comparison over the visible values only. + let match_bitmap = arrow_ord::cmp::not_distinct(&visible_values, needle)?; + let match_bits = match_bitmap.values(); + + for (row_index, offset_window) in list_array.offsets().windows(2).enumerate() { + // Offsets relative to visible_values (subtract first_offset). + let start = offset_window[0].to_usize().unwrap() - first_offset; + let end = offset_window[1].to_usize().unwrap() - first_offset; + let row_len = end - start; + + if list_array.is_null(row_index) { + offsets.push(offsets[row_index]); + continue; } - DataType::LargeList(_) => { - let list_array = array.as_list::(); - general_replace::(list_array, from, to, &arr_n) + + // Slice the match bits to this row and iterate only over true positions. + let row_bits = match_bits.slice(start, row_len); + let mut match_positions = row_bits + .set_indices() + .take(max_replacements as usize) + .peekable(); + if match_positions.peek().is_none() { + mutable.try_extend(0, start, end)?; + offsets.push(offsets[row_index] + O::usize_as(row_len)); + continue; } - DataType::Null => Ok(new_null_array(array.data_type(), 1)), - array_type => exec_err!("array_replace does not support type '{array_type}'."), + + // Iterate only over the positions that match using set_indices, + // which is more efficient than scanning every bit because the number + // of matches is typically much smaller than the total array size. + let mut prev_end = 0usize; + for match_pos in match_positions { + // Retain elements before this match. + if match_pos > prev_end { + mutable.try_extend(0, start + prev_end, start + match_pos)?; + } + // Emit the replacement element. + mutable.try_extend(1, 0, 1)?; + prev_end = match_pos + 1; + } + + // Copy remaining elements after the last replacement. + if prev_end < row_len { + mutable.try_extend(0, start + prev_end, end)?; + } + + offsets.push(offsets[row_index] + O::usize_as(row_len)); } + + let data = mutable.freeze(); + + Ok(Arc::new(GenericListArray::::try_new( + field, + OffsetBuffer::new(offsets.into()), + arrow::array::make_array(data), + list_array.nulls().cloned(), + )?)) } -fn array_replace_n_inner(args: &[ArrayRef]) -> Result { - let [array, from, to, max] = take_function_args("array_replace_n", args)?; +/// Fast path for `array_replace` when all arguments are scalars. +/// +/// Uses a single bulk `not_distinct` comparison instead of per-row comparisons. +fn array_replace_with_scalar_args( + name: &str, + list_array: &ArrayRef, + scalar_from: &ScalarValue, + scalar_to: &ScalarValue, + max_replacements: i64, + return_type: &DataType, +) -> Result { + // `not_distinct` doesn't support nested types, fall back to the generic array path. + if scalar_from.data_type().is_nested() { + let num_rows = list_array.len(); + let from_array = scalar_from.to_array_of_size(num_rows)?; + let to_array = scalar_to.to_array_of_size(num_rows)?; + return array_replace_internal( + name, + list_array, + &from_array, + &to_array, + &vec![Some(max_replacements); num_rows], + return_type, + ); + } + + let needle = Scalar::new(scalar_from.to_array_of_size(1)?); + match list_array.data_type() { + DataType::List(_) => general_replace_with_scalar::( + list_array.as_list::(), + &needle, + scalar_to, + max_replacements, + list_inner_field(name, return_type)?, + ), + DataType::LargeList(_) => general_replace_with_scalar::( + list_array.as_list::(), + &needle, + scalar_to, + max_replacements, + list_inner_field(name, return_type)?, + ), + DataType::Null => Ok(new_null_array(return_type, list_array.len())), + array_type => exec_err!("{name} does not support type '{array_type}'."), + } +} - // replace the specified number of occurrences - let arr_n = as_int64_array(max)?.values().to_vec(); +fn array_replace_internal( + name: &str, + array: &ArrayRef, + from: &ArrayRef, + to: &ArrayRef, + arr_n: &[Option], + return_type: &DataType, +) -> Result { match array.data_type() { - DataType::List(_) => { - let list_array = array.as_list::(); - general_replace::(list_array, from, to, &arr_n) - } - DataType::LargeList(_) => { - let list_array = array.as_list::(); - general_replace::(list_array, from, to, &arr_n) - } - DataType::Null => Ok(new_null_array(array.data_type(), 1)), - array_type => { - exec_err!("array_replace_n does not support type '{array_type}'.") - } + DataType::List(_) => general_replace::( + array.as_list::(), + from, + to, + arr_n, + list_inner_field(name, return_type)?, + ), + DataType::LargeList(_) => general_replace::( + array.as_list::(), + from, + to, + arr_n, + list_inner_field(name, return_type)?, + ), + DataType::Null => Ok(new_null_array(return_type, array.len())), + array_type => exec_err!("{name} does not support type '{array_type}'."), } } -fn array_replace_all_inner(args: &[ArrayRef]) -> Result { - let [array, from, to] = take_function_args("array_replace_all", args)?; +fn array_replace_n_inner( + name: &str, + array: &ArrayRef, + from: &ArrayRef, + to: &ArrayRef, + max: &ArrayRef, + return_type: &DataType, +) -> Result { + let arr_n = as_int64_array(max)?.iter().collect::>(); + array_replace_internal(name, array, from, to, &arr_n, return_type) +} - // replace all occurrences (up to "i64::MAX") - let arr_n = vec![i64::MAX; array.len()]; - match array.data_type() { - DataType::List(_) => { - let list_array = array.as_list::(); - general_replace::(list_array, from, to, &arr_n) - } - DataType::LargeList(_) => { - let list_array = array.as_list::(); - general_replace::(list_array, from, to, &arr_n) - } - DataType::Null => Ok(new_null_array(array.data_type(), 1)), - array_type => { - exec_err!("array_replace_all does not support type '{array_type}'.") - } +#[cfg(test)] +mod tests { + use super::{ArrayReplaceN, array_replace_n_inner}; + use arrow::array::{ArrayRef, AsArray, Int32Array, Int64Array, ListArray}; + use arrow::buffer::{NullBuffer, ScalarBuffer}; + use arrow::datatypes::{DataType, Field, Int32Type}; + use datafusion_common::{Result, ScalarValue, config::ConfigOptions}; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; + use std::sync::Arc; + + #[test] + fn test_array_replace_n_null_max_returns_null() -> Result<()> { + let array: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![Some(4), Some(2)]), + ])); + let from: ArrayRef = Arc::new(Int32Array::from(vec![2, 2])); + let to: ArrayRef = Arc::new(Int32Array::from(vec![9, 9])); + let max: ArrayRef = Arc::new(Int64Array::new( + ScalarBuffer::from(vec![1, 1]), + Some(NullBuffer::from(vec![true, false])), + )); + + let result = array_replace_n_inner( + "array_replace_n", + &array, + &from, + &to, + &max, + array.data_type(), + )?; + let expected = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(9), Some(3)]), + None, + ]); + + assert_eq!(result.as_list::(), &expected); + + Ok(()) + } + + #[test] + fn test_array_replace_n_scalar_null_max_returns_null() -> Result<()> { + let array: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![Some(4), Some(2)]), + ])); + let array_field = Arc::new(Field::new("array", array.data_type().clone(), true)); + + let result = ArrayReplaceN::new().invoke_with_args(ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::clone(&array)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ColumnarValue::Scalar(ScalarValue::Int32(Some(9))), + ColumnarValue::Scalar(ScalarValue::Int64(None)), + ], + arg_fields: vec![ + Arc::clone(&array_field), + Arc::new(Field::new("from", DataType::Int32, false)), + Arc::new(Field::new("to", DataType::Int32, false)), + Arc::new(Field::new("max", DataType::Int64, true)), + ], + number_rows: array.len(), + return_field: Arc::clone(&array_field), + config_options: Arc::new(ConfigOptions::default()), + })?; + + let result = result.into_array(array.len())?; + let expected = ListArray::from_iter_primitive::(vec![ + Option::>>::None, + Option::>>::None, + ]); + + assert_eq!(result.as_list::(), &expected); + + Ok(()) } } diff --git a/datafusion/functions-nested/src/resize.rs b/datafusion/functions-nested/src/resize.rs index 243f3531f9150..e08149ec0f938 100644 --- a/datafusion/functions-nested/src/resize.rs +++ b/datafusion/functions-nested/src/resize.rs @@ -49,8 +49,8 @@ make_udf_expr_and_func!( #[user_doc( doc_section(label = "Array Functions"), - description = "Resizes the list to contain size elements. Initializes new elements with value or empty if value is not set.", - syntax_example = "array_resize(array, size, value)", + description = "Resizes the list to contain size elements.", + syntax_example = "array_resize(array, size[, value])", sql_example = r#"```sql > select array_resize([1, 2, 3], 5, 0); +-------------------------------------+ @@ -66,7 +66,7 @@ make_udf_expr_and_func!( argument(name = "size", description = "New size of given array."), argument( name = "value", - description = "Defines new elements' value or empty if value is not set." + description = "If expanding the array, defines the values to fill in. Defaults to null." ) )] #[derive(Debug, PartialEq, Eq, Hash)] @@ -203,7 +203,7 @@ fn general_list_resize>( let mut max_extra: usize = 0; let mut output_values_len: usize = 0; for (row_index, offset_window) in array.offsets().windows(2).enumerate() { - if array.is_null(row_index) { + if array.is_null(row_index) || count_array.is_null(row_index) { continue; } let target_count = count_array.value(row_index).to_usize().ok_or_else(|| { @@ -219,6 +219,14 @@ fn general_list_resize>( } } + if output_values_len > max_resize_values(&data_type) + || O::from_usize(output_values_len).is_none() + { + return exec_err!( + "array_resize: resulting array of {output_values_len} elements exceeds the maximum array size" + ); + } + // The fast path is valid when at least one row grows and every row would // use the same fill value. let use_bulk_fill = max_extra > 0 @@ -256,7 +264,7 @@ fn general_list_resize>( &original_data, &default_value_data, output_values_len, - |mutable, _, extra_count| mutable.extend(1, 0, extra_count), + |mutable, _, extra_count| Ok(mutable.try_extend(1, 0, extra_count)?), ) } else { // Slow path: rows may need different fill values, so append from the @@ -278,8 +286,9 @@ fn general_list_resize>( output_values_len, |mutable, row_index, extra_count| { for _ in 0..extra_count { - mutable.extend(1, row_index, row_index + 1); + mutable.try_extend(1, row_index, row_index + 1)?; } + Ok(()) }, ) } @@ -296,7 +305,7 @@ fn build_resized_list( ) -> Result where O: OffsetSizeTrait + TryInto, - F: FnMut(&mut MutableArrayData, usize, usize), + F: FnMut(&mut MutableArrayData, usize, usize) -> Result<()>, { let capacity = Capacities::Array(output_values_len); let mut offsets = vec![O::usize_as(0)]; @@ -308,7 +317,7 @@ where let mut null_builder = NullBufferBuilder::new(array.len()); for (row_index, offset_window) in array.offsets().windows(2).enumerate() { - if array.is_null(row_index) { + if array.is_null(row_index) || count_array.is_null(row_index) { null_builder.append_null(); offsets.push(offsets[row_index]); continue; @@ -323,11 +332,11 @@ where if start + count > offset_window[1] { let extra_count = (start + count - offset_window[1]).to_usize().unwrap(); let end = offset_window[1]; - mutable.extend(0, start.to_usize().unwrap(), end.to_usize().unwrap()); - append_fill_values(&mut mutable, row_index, extra_count); + mutable.try_extend(0, start.to_usize().unwrap(), end.to_usize().unwrap())?; + append_fill_values(&mut mutable, row_index, extra_count)?; } else { let end = start + count; - mutable.extend(0, start.to_usize().unwrap(), end.to_usize().unwrap()); + mutable.try_extend(0, start.to_usize().unwrap(), end.to_usize().unwrap())?; }; offsets.push(offsets[row_index] + count); } @@ -341,3 +350,119 @@ where null_builder.finish(), )?)) } + +/// Largest element count whose eager value buffer stays within `isize::MAX` +/// bytes, so `array_resize` rejects oversized results instead of panicking. +/// Only primitive and `FixedSizeBinary` leaves are byte-exact. +fn max_resize_values(value_type: &DataType) -> usize { + let element_width = match value_type { + DataType::FixedSizeBinary(size) if *size > 0 => *size as usize, + _ => value_type.primitive_width().unwrap_or(size_of::()), + }; + + (isize::MAX as usize) / element_width.max(1) +} + +#[cfg(test)] +mod tests { + use super::array_resize_inner; + use arrow::array::{ + ArrayRef, AsArray, FixedSizeBinaryArray, Int64Array, LargeListArray, ListArray, + }; + use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; + use arrow::datatypes::{DataType, Field, Int32Type, Int64Type}; + use datafusion_common::Result; + use std::sync::Arc; + + #[test] + fn test_array_resize_null_size_returns_null() -> Result<()> { + let array: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + Some(vec![Some(4), Some(5)]), + ])); + let size: ArrayRef = Arc::new(Int64Array::new( + ScalarBuffer::from(vec![2, 1]), + Some(NullBuffer::from(vec![true, false])), + )); + + let result = array_resize_inner(&[array, size])?; + let expected = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + None, + ]); + + assert_eq!(result.as_list::(), &expected); + + Ok(()) + } + + #[test] + fn test_array_resize_large_size_errors_without_panicking() { + let array: ArrayRef = + Arc::new(ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1)]), + ])); + let size: ArrayRef = Arc::new(Int64Array::from(vec![i64::MAX])); + let fill: ArrayRef = Arc::new(Int64Array::from(vec![0])); + + let err = array_resize_inner(&[array, size, fill]).unwrap_err(); + assert!( + err.to_string().contains("exceeds the maximum array size"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_array_resize_fixed_size_binary_large_size_errors_without_panicking() { + let values = + FixedSizeBinaryArray::try_from_iter(vec![vec![0u8; 32]].into_iter()).unwrap(); + let elem_field = + Arc::new(Field::new_list_field(DataType::FixedSizeBinary(32), true)); + let offsets = OffsetBuffer::::new(vec![0i64, 1].into()); + let array: ArrayRef = Arc::new(LargeListArray::new( + elem_field, + offsets, + Arc::new(values) as ArrayRef, + None, + )); + // Passes the width-16 bound (isize::MAX / 16) but overflows at width 32. + let size: ArrayRef = Arc::new(Int64Array::from(vec![400_000_000_000_000_000i64])); + + let err = array_resize_inner(&[array, size]).unwrap_err(); + assert!( + err.to_string().contains("exceeds the maximum array size"), + "unexpected error: {err}" + ); + } + + #[test] + fn test_array_resize_accumulates_values_across_rows() { + // Each row's target (6e17) is individually under the width-8 cap + // (isize::MAX / 8), but their sum (1.2e18) exceeds it, so the guard + // must reject based on the accumulated total rather than per row. + let values = Int64Array::from(vec![1, 2]); + let offsets = OffsetBuffer::::new(vec![0i64, 1, 2].into()); + let elem_field = Arc::new(Field::new_list_field(DataType::Int64, true)); + let array: ArrayRef = Arc::new(LargeListArray::new( + elem_field, + offsets, + Arc::new(values) as ArrayRef, + None, + )); + let size: ArrayRef = Arc::new(Int64Array::from(vec![ + 600_000_000_000_000_000i64, + 600_000_000_000_000_000i64, + ])); + + let err = array_resize_inner(&[array, size]).unwrap_err(); + assert!( + err.to_string().contains("1200000000000000000"), + "expected accumulated total in error: {err}" + ); + assert!( + err.to_string().contains("exceeds the maximum array size"), + "unexpected error: {err}" + ); + } +} diff --git a/datafusion/functions-nested/src/set_ops.rs b/datafusion/functions-nested/src/set_ops.rs index 2ad08e2d43c02..2214d3d35bb7b 100644 --- a/datafusion/functions-nested/src/set_ops.rs +++ b/datafusion/functions-nested/src/set_ops.rs @@ -28,7 +28,7 @@ use arrow::datatypes::DataType::{LargeList, List, Null}; use arrow::datatypes::{DataType, Field, FieldRef}; use arrow::row::{RowConverter, SortField}; use datafusion_common::cast::{as_large_list_array, as_list_array}; -use datafusion_common::utils::ListCoercion; +use datafusion_common::utils::{ListCoercion, normalize_float_zero}; use datafusion_common::{ Result, assert_eq_or_internal_err, exec_err, internal_err, utils::take_function_args, }; @@ -351,21 +351,28 @@ fn generic_set_lists( let converter = RowConverter::new(vec![SortField::new(l.value_type())])?; + // Normalize -0.0 → +0.0 so RowConverter (which uses IEEE 754 totalOrder + // and treats ±0 as distinct) groups them together. Use the normalized + // arrays for both row conversion and the final output values. + let l_values_norm = normalize_float_zero(l.values()); + let r_values_norm = normalize_float_zero(r.values()); + // Only convert the visible portion of the values array. For sliced // ListArrays, values() returns the full underlying array but only // elements between the first and last offset are referenced. let l_first = l.offsets()[0].as_usize(); let l_len = l.offsets()[l.len()].as_usize() - l_first; - let rows_l = converter.convert_columns(&[l.values().slice(l_first, l_len)])?; + let l_values = l_values_norm.slice(l_first, l_len); + let rows_l = converter.convert_columns(&[Arc::clone(&l_values)])?; let r_first = r.offsets()[0].as_usize(); let r_len = r.offsets()[r.len()].as_usize() - r_first; - let rows_r = converter.convert_columns(&[r.values().slice(r_first, r_len)])?; + let r_values = r_values_norm.slice(r_first, r_len); + let rows_r = converter.convert_columns(&[Arc::clone(&r_values)])?; - // Combine the *sliced* value arrays so 0-based indices from the row - // converter map directly into the concatenated array. - let l_values = l.values().slice(l_first, l_len); - let r_values = r.values().slice(r_first, r_len); + // Indices from the row converter are 0-based in the per-side slice; + // concatenating those same slices lets indices map directly into the + // combined values array. let combined_values = concat(&[l_values.as_ref(), r_values.as_ref()])?; let r_offset = l_len; @@ -558,13 +565,18 @@ fn general_array_distinct( let converter = RowConverter::new(vec![SortField::new(dt.clone())])?; + // Normalize -0.0 → +0.0 so RowConverter (which uses IEEE 754 totalOrder + // and treats ±0 as distinct) groups them together, and so the output + // carries the canonical sign. + let values_norm = normalize_float_zero(array.values()); + // Only convert the visible portion of the values array. For sliced // ListArrays, values() returns the full underlying array but only // elements between the first and last offset are referenced. let first_offset = value_offsets[0].as_usize(); let visible_len = value_offsets[array.len()].as_usize() - first_offset; let rows = - converter.convert_columns(&[array.values().slice(first_offset, visible_len)])?; + converter.convert_columns(&[values_norm.slice(first_offset, visible_len)])?; let mut indices: Vec = Vec::with_capacity(rows.num_rows()); let mut seen = HashSet::new(); @@ -593,19 +605,19 @@ fn general_array_distinct( } // Gather distinct values in a single pass, using the computed `indices`. - // Indices are absolute positions in array.values() (first_offset was added - // back when collecting them), so we can take directly from the full values. + // Indices are absolute positions in the (normalized) values array, so we + // can take directly from the full values. // Use UInt64Array for LargeList to support values arrays exceeding u32::MAX. let final_values = if indices.is_empty() { new_empty_array(&dt) } else if OffsetSize::IS_LARGE { let indices = UInt64Array::from(indices.into_iter().map(|i| i as u64).collect::>()); - take(array.values().as_ref(), &indices, None)? + take(values_norm.as_ref(), &indices, None)? } else { let indices = UInt32Array::from(indices.into_iter().map(|i| i as u32).collect::>()); - take(array.values().as_ref(), &indices, None)? + take(values_norm.as_ref(), &indices, None)? }; Ok(Arc::new(GenericListArray::::try_new( diff --git a/datafusion/functions-nested/src/sort.rs b/datafusion/functions-nested/src/sort.rs index 0a34cce6b965f..f4f9148f760bf 100644 --- a/datafusion/functions-nested/src/sort.rs +++ b/datafusion/functions-nested/src/sort.rs @@ -55,7 +55,7 @@ make_udf_expr_and_func!( #[user_doc( doc_section(label = "Array Functions"), description = "Sort array.", - syntax_example = "array_sort(array, desc, nulls_first)", + syntax_example = "array_sort(array[, order[, nulls_order]])", sql_example = r#"```sql > select array_sort([3, 1, 2]); +-----------------------------+ @@ -63,17 +63,23 @@ make_udf_expr_and_func!( +-----------------------------+ | [1, 2, 3] | +-----------------------------+ +> select array_sort([3, 1, NULL, 2], 'desc', 'nulls last'); ++--------------------------------------------------+ +| array_sort(List(3,1,NULL,2),'desc','nulls last') | ++--------------------------------------------------+ +| [3, 2, 1, NULL] | ++--------------------------------------------------+ ```"#, argument( name = "array", description = "Array expression. Can be a constant, column, or function, and any combination of array operators." ), argument( - name = "desc", + name = "order", description = "Whether to sort in ascending (`ASC`) or descending (`DESC`) order. The default is `ASC`." ), argument( - name = "nulls_first", + name = "nulls_order", description = "Whether to sort nulls first (`NULLS FIRST`) or last (`NULLS LAST`). The default is `NULLS FIRST`." ) )] @@ -471,12 +477,7 @@ fn take_by_indices( fn rebase_offsets( offsets: &OffsetBuffer, ) -> OffsetBuffer { - if offsets[0].as_usize() == 0 { - offsets.clone() - } else { - let rebased: Vec = offsets.iter().map(|o| *o - offsets[0]).collect(); - OffsetBuffer::new(rebased.into()) - } + offsets.clone().subtract(offsets[0]) } fn order_desc(modifier: &str) -> Result { diff --git a/datafusion/functions-nested/src/utils.rs b/datafusion/functions-nested/src/utils.rs index eeff003e8e766..9822b6121e695 100644 --- a/datafusion/functions-nested/src/utils.rs +++ b/datafusion/functions-nested/src/utils.rs @@ -19,21 +19,73 @@ use std::sync::Arc; -use arrow::datatypes::{DataType, Field, Fields}; +use arrow::datatypes::{DataType, Field, FieldRef, Fields}; use arrow::array::{ - Array, ArrayRef, BooleanArray, GenericListArray, OffsetSizeTrait, Scalar, + Array, ArrayRef, BooleanArray, Float64Array, GenericListArray, NullBufferBuilder, + OffsetSizeTrait, Scalar, }; -use arrow::buffer::OffsetBuffer; +use arrow::buffer::{NullBuffer, OffsetBuffer}; use datafusion_common::cast::{ - as_fixed_size_list_array, as_large_list_array, as_large_list_view_array, - as_list_array, as_list_view_array, + as_fixed_size_list_array, as_float64_array, as_generic_list_array, + as_large_list_array, as_large_list_view_array, as_list_array, as_list_view_array, }; use datafusion_common::{Result, ScalarValue, exec_err, internal_err, plan_err}; use datafusion_expr::ColumnarValue; use itertools::Itertools as _; +/// Computes the return type of a function that produces a list with the same +/// inner field as `array_type`, plus an element that may be null when +/// `element_nullable` is set. +/// +/// The inner field is carried over from `array_type` verbatim — name, metadata +/// and all — so that the type promised at planning time is the one the kernel +/// can actually build. Its nullability is widened when `element_nullable` is +/// set, because a nullable new element may introduce nulls into a list whose +/// elements were previously declared non-nullable. +/// +/// Types other than `List`/`LargeList` are returned unchanged; callers handle +/// `Null` themselves and the kernels reject anything else at execution time. +pub(crate) fn list_type_with_element( + array_type: &DataType, + element_nullable: bool, +) -> DataType { + match array_type { + DataType::List(field) => { + DataType::List(widen_nullability(field, element_nullable)) + } + DataType::LargeList(field) => { + DataType::LargeList(widen_nullability(field, element_nullable)) + } + other => other.clone(), + } +} + +fn widen_nullability(field: &FieldRef, nullable: bool) -> FieldRef { + if nullable && !field.is_nullable() { + Arc::new(field.as_ref().clone().with_nullable(true)) + } else { + Arc::clone(field) + } +} + +/// Extracts the inner field of a `List`/`LargeList` type, so that a kernel can +/// build a list array carrying exactly that field. +/// +/// Used both on an input's type and on the type promised by +/// [`ScalarUDFImpl::return_field_from_args`]. Anything else is a bug in the +/// caller's dispatch, hence the internal error; `context` names the kernel so +/// that error identifies where the bad dispatch happened. +/// +/// [`ScalarUDFImpl::return_field_from_args`]: datafusion_expr::ScalarUDFImpl::return_field_from_args +pub(crate) fn list_inner_field(context: &str, data_type: &DataType) -> Result { + match data_type { + DataType::List(field) | DataType::LargeList(field) => Ok(Arc::clone(field)), + other => internal_err!("{context} got unexpected data type: {other}"), + } +} + pub(crate) fn check_datatypes(name: &str, args: &[&ArrayRef]) -> Result<()> { let data_type = args[0].data_type(); if !args.iter().all(|arg| { @@ -276,6 +328,142 @@ pub(crate) fn get_map_entry_field(data_type: &DataType) -> Result<&Fields> { } } +/// Shared `coerce_types` impl for array-math UDFs whose kernels expect +/// `List` / `LargeList` (e.g. `array_add`, `cosine_distance`, +/// `inner_product`, `array_normalize`). +/// +/// Each input must be `Null`, `List`, `LargeList`, or `FixedSizeList`; otherwise +/// returns a plan error naming `name`. `FixedSizeList` is widened to `List`, +/// `Null` is coerced to a list of `Float64`, and if any input is `LargeList` +/// the rest are widened to `LargeList` so the runtime sees a homogeneous pair. +pub(crate) fn coerce_array_math_arg_types( + name: &str, + arg_types: &[DataType], +) -> Result> { + use DataType::{FixedSizeList, LargeList, List, Null}; + use datafusion_common::utils::{ListCoercion, coerced_type_with_base_type_only}; + + let coercion = Some(&ListCoercion::FixedSizedListToList); + + for arg_type in arg_types { + if !matches!(arg_type, Null | List(_) | LargeList(_) | FixedSizeList(..)) { + return plan_err!("{name} does not support type {arg_type}"); + } + } + + // If any input is `LargeList`, both sides must be widened to `LargeList` + // so the runtime dispatch in `inner_product_inner` sees a homogeneous + // pair. Follows the pattern in `ArrayConcat::coerce_types`. + let any_large_list = arg_types.iter().any(|t| matches!(t, LargeList(_))); + + let coerced = arg_types + .iter() + .map(|arg_type| { + if matches!(arg_type, Null) { + let field = Arc::new(Field::new_list_field(DataType::Float64, true)); + return if any_large_list { + LargeList(field) + } else { + List(field) + }; + } + let coerced = + coerced_type_with_base_type_only(arg_type, &DataType::Float64, coercion); + match coerced { + List(field) if any_large_list => LargeList(field), + other => other, + } + }) + .collect(); + + Ok(coerced) +} + +/// Element-wise binary operation kernel for two `Float64` lists of equal per-row +/// length. The caller is responsible for type-dispatching on `O` (`i32` for +/// `List`, `i64` for `LargeList`). +/// +/// Semantics: +/// - whole-row NULL on either side → NULL output row, length 0 +/// - per-element NULL on either side → NULL at that output position +/// - per-row length mismatch → exec error tagged with `op_name` +/// +/// `op_name` flows into the error message; `op` is the per-element scalar op +/// (e.g. `|a, b| a + b` for `array_add`, `|a, b| a - b` for `array_subtract`). +pub(crate) fn array_math_binary_op( + op_name: &str, + lhs: &ArrayRef, + rhs: &ArrayRef, + op: F, +) -> Result +where + O: OffsetSizeTrait, + F: Fn(f64, f64) -> f64, +{ + let lhs = as_generic_list_array::(lhs)?; + let rhs = as_generic_list_array::(rhs)?; + + let lhs_values = as_float64_array(lhs.values())?; + let rhs_values = as_float64_array(rhs.values())?; + let lhs_offsets = lhs.value_offsets(); + let rhs_offsets = rhs.value_offsets(); + + let row_nulls = NullBuffer::union(lhs.nulls(), rhs.nulls()); + + let mut out_values: Vec = Vec::with_capacity(lhs_values.len()); + let mut out_inner_nulls = NullBufferBuilder::new(lhs_values.len()); + let mut out_offsets = Vec::::with_capacity(lhs.len() + 1); + out_offsets.push(O::zero()); + + for row in 0..lhs.len() { + if row_nulls.as_ref().is_some_and(|nb| nb.is_null(row)) { + out_offsets.push(out_offsets[row]); + continue; + } + + let start1 = lhs_offsets[row].as_usize(); + let len1 = lhs.value_length(row).as_usize(); + let start2 = rhs_offsets[row].as_usize(); + let len2 = rhs.value_length(row).as_usize(); + + if len1 != len2 { + return exec_err!( + "{op_name} requires both list inputs to have the same length per row, got {len1} and {len2} at row {row}" + ); + } + + let l_slice = lhs_values.slice(start1, len1); + let r_slice = rhs_values.slice(start2, len2); + + let l_vals = l_slice.values(); + let r_vals = r_slice.values(); + + for i in 0..len1 { + out_values.push(op(l_vals[i], r_vals[i])); + } + + match NullBuffer::union(l_slice.nulls(), r_slice.nulls()) { + Some(nb) => out_inner_nulls.append_buffer(&nb), + None => out_inner_nulls.append_n_non_nulls(len1), + } + + out_offsets.push(out_offsets[row] + O::usize_as(len1)); + } + + let values_array = Arc::new(Float64Array::new( + out_values.into(), + out_inner_nulls.finish(), + )); + let field = Arc::new(Field::new_list_field(DataType::Float64, true)); + + Ok(Arc::new(GenericListArray::::try_new( + field, + OffsetBuffer::new(out_offsets.into()), + values_array, + row_nulls, + )?)) +} + #[cfg(test)] mod tests { use super::*; diff --git a/datafusion/functions-table/src/generate_series.rs b/datafusion/functions-table/src/generate_series.rs index 52baa7e6cf8ef..f5e4df13899df 100644 --- a/datafusion/functions-table/src/generate_series.rs +++ b/datafusion/functions-table/src/generate_series.rs @@ -27,7 +27,7 @@ use async_trait::async_trait; use datafusion_catalog::TableFunctionImpl; use datafusion_catalog::TableProvider; use datafusion_catalog::{Session, TableFunctionArgs}; -use datafusion_common::{Result, ScalarValue, plan_err}; +use datafusion_common::{Result, ScalarValue, plan_datafusion_err, plan_err}; use datafusion_expr::{Expr, TableType}; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::expressions::Column; @@ -79,9 +79,18 @@ pub trait SeriesValue: fmt::Debug + Clone + Send + Sync + 'static { /// Check if we've reached the end of the series fn should_stop(&self, end: Self, step: &Self::StepType, include_end: bool) -> bool; - /// Advance to the next value in the series + /// Advance to the next value in the series. fn advance(&mut self, step: &Self::StepType) -> Result<()>; + /// Advance to the next value, adjusting the end of the series if needed. + /// + /// The default implementation preserves the behavior of [`Self::advance`]. + /// Implementations can override this method when they need to handle an + /// overflow by terminating the series after the current value. + fn advance_with_end(&mut self, _end: &mut Self, step: &Self::StepType) -> Result<()> { + self.advance(step) + } + /// Create an Arrow array from a vector of values fn create_array(&self, values: Vec) -> Result; @@ -105,6 +114,22 @@ impl SeriesValue for i64 { Ok(()) } + fn advance_with_end(&mut self, end: &mut Self, step: &Self::StepType) -> Result<()> { + if let Some(next) = self.checked_add(*step) { + *self = next; + } else { + // Advancing would overflow: clamp `end` so the series stops after + // the current (last reachable) value instead of panicking or + // wrapping around. + *end = if *step > 0 { + self.saturating_sub(1) + } else { + self.saturating_add(1) + }; + } + Ok(()) + } + fn create_array(&self, values: Vec) -> Result { Ok(Arc::new(Int64Array::from(values))) } @@ -172,6 +197,27 @@ impl SeriesValue for TimestampValue { Ok(()) } + fn advance_with_end(&mut self, end: &mut Self, step: &Self::StepType) -> Result<()> { + let tz = self + .parsed_tz + .unwrap_or_else(|| Tz::from_str("+00:00").unwrap()); + if let Some(next_ts) = + TimestampNanosecondType::add_month_day_nano(self.value, *step, tz) + { + self.value = next_ts; + } else { + // Advancing would exceed the timestamp range. Clamp `end` so the + // series terminates after the current (last reachable) value. + let step_negative = step.months < 0 || step.days < 0 || step.nanoseconds < 0; + end.value = if step_negative { + self.value.saturating_add(1) + } else { + self.value.saturating_sub(1) + }; + } + Ok(()) + } + fn create_array(&self, values: Vec) -> Result { let array = TimestampNanosecondArray::from(values); @@ -259,6 +305,7 @@ impl GenerateSeriesTable { end: *end, step: *step, current: *start, + finished: false, batch_size, include_end: *include_end, name, @@ -299,6 +346,7 @@ impl GenerateSeriesTable { parsed_tz: Some(parsed_tz), tz_str: tz.clone(), }, + finished: false, batch_size, include_end: *include_end, name, @@ -328,6 +376,7 @@ impl GenerateSeriesTable { parsed_tz: None, tz_str: None, }, + finished: false, batch_size, include_end: *include_end, name, @@ -369,6 +418,7 @@ pub struct GenericSeriesState { step: T::StepType, batch_size: usize, current: T, + finished: bool, include_end: bool, name: &'static str, } @@ -409,6 +459,10 @@ impl LazyBatchGenerator for GenericSeriesState { } fn generate_next_batch(&mut self) -> Result> { + if self.finished { + return Ok(None); + } + let mut buf = Vec::with_capacity(self.batch_size); while buf.len() < self.batch_size @@ -417,7 +471,24 @@ impl LazyBatchGenerator for GenericSeriesState { .should_stop(self.end.clone(), &self.step, self.include_end) { buf.push(self.current.to_value_type()); - self.current.advance(&self.step)?; + if self + .current + .should_stop(self.end.clone(), &self.step, false) + { + self.finished = true; + break; + } + + let original_end = self.end.clone(); + self.current.advance_with_end(&mut self.end, &self.step)?; + if self + .current + .should_stop(self.end.clone(), &self.step, self.include_end) + { + self.end = original_end; + self.finished = true; + break; + } } if buf.is_empty() { @@ -432,6 +503,7 @@ impl LazyBatchGenerator for GenericSeriesState { fn reset_state(&self) -> Arc> { let mut new = self.clone(); new.current = new.start.clone(); + new.finished = false; Arc::new(RwLock::new(new)) } } @@ -484,7 +556,7 @@ impl TableProvider for GenerateSeriesTable { _filters: &[Expr], _limit: Option, ) -> Result> { - let batch_size = state.config_options().execution.batch_size; + let batch_size = state.config_options().execution.batch_size.get(); let generator = self.as_generator(batch_size)?; let mut exec = LazyMemoryExec::try_new(self.schema(), vec![generator])? .with_projection(projection.cloned()); @@ -740,8 +812,20 @@ impl GenerateSeriesFuncImpl { // Date32 is days since 1970-01-01, so multiply by nanoseconds per day const NANOS_PER_DAY: i64 = 24 * 60 * 60 * 1_000_000_000; - let start_ts = start_date as i64 * NANOS_PER_DAY; - let end_ts = end_date as i64 * NANOS_PER_DAY; + // Dates outside the nanosecond timestamp range (1677-09-21 to + // 2262-04-11) cannot be represented; return an error instead of + // panicking (debug) or silently wrapping (release). + let date_to_ts_nanos = |date: i32, arg: &str| { + (date as i64).checked_mul(NANOS_PER_DAY).ok_or_else(|| { + plan_datafusion_err!( + "{arg} for {} is out of range of nanosecond timestamps", + self.name + ) + }) + }; + + let start_ts = date_to_ts_nanos(start_date, "First argument")?; + let end_ts = date_to_ts_nanos(end_date, "Second argument")?; // Validate step interval validate_interval_step(step_interval)?; @@ -804,11 +888,40 @@ mod generate_series_tests { end: 5, step: 1, current: 1, + finished: false, + batch_size: 8192, + include_end: true, + name: "test", + }; + let batch = state.generate_next_batch()?.expect("missing batch"); + + let state_reset = state.reset_state(); + let reset_batch = state_reset + .write() + .generate_next_batch()? + .expect("missing reset batch"); + + assert_eq!(batch, reset_batch); + + Ok(()) + } + + #[test] + fn test_generic_series_state_reset_after_overflow() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let mut state = GenericSeriesState:: { + schema, + start: i64::MAX - 1, + end: i64::MAX, + step: 2, + current: i64::MAX - 1, + finished: false, batch_size: 8192, include_end: true, name: "test", }; let batch = state.generate_next_batch()?.expect("missing batch"); + assert!(state.generate_next_batch()?.is_none()); let state_reset = state.reset_state(); let reset_batch = state_reset diff --git a/datafusion/functions-window/src/lead_lag.rs b/datafusion/functions-window/src/lead_lag.rs index de4071c0ceda7..fea4a1a4aadda 100644 --- a/datafusion/functions-window/src/lead_lag.rs +++ b/datafusion/functions-window/src/lead_lag.rs @@ -18,6 +18,8 @@ //! `lead` and `lag` window function implementations use crate::utils::{get_scalar_value_from_args, get_signed_integer}; +use arrow::array::UInt64Builder; +use arrow::compute::{interleave, take}; use arrow::datatypes::FieldRef; use datafusion_common::arrow::array::ArrayRef; use datafusion_common::arrow::datatypes::DataType; @@ -419,6 +421,52 @@ fn offset_magnitude(offset: i64) -> usize { } } +enum ShiftIndexBuilder { + Take(UInt64Builder), + Interleave(Vec<(usize, usize)>), +} + +impl ShiftIndexBuilder { + fn new(capacity: usize, default_is_null: bool) -> Self { + if default_is_null { + Self::Take(UInt64Builder::with_capacity(capacity)) + } else { + Self::Interleave(Vec::with_capacity(capacity)) + } + } + + fn append_option(&mut self, index: Option) { + match self { + Self::Take(indices) => { + indices.append_option(index.map(|index| index as u64)); + } + Self::Interleave(indices) => { + // `interleave` receives `[array, default]`. + indices.push(index.map_or((1, 0), |index| (0, index))); + } + } + } + + fn finish( + self, + array: &ArrayRef, + default_value: &ScalarValue, + ) -> Result { + match self { + Self::Take(mut indices) => { + let indices = indices.finish(); + take(array.as_ref(), &indices, None) + .map_err(|error| arrow_datafusion_err!(error)) + } + Self::Interleave(indices) => { + let default = default_value.to_array_of_size(1)?; + interleave(&[array.as_ref(), default.as_ref()], &indices) + .map_err(|error| arrow_datafusion_err!(error)) + } + } + } +} + impl WindowShiftEvaluator { fn is_lag(&self) -> bool { // Mode is LAG, when shift_offset is positive @@ -433,49 +481,57 @@ fn evaluate_all_with_ignore_null( default_value: &ScalarValue, is_lag: bool, ) -> Result { - let valid_indices: Vec = - array.nulls().unwrap().valid_indices().collect::>(); - let direction = !is_lag; - let new_array_results: Result, DataFusionError> = (0..array.len()) - .map(|id| { - let result_index = match valid_indices.binary_search(&id) { - Ok(pos) => if direction { - pos.checked_add(offset as usize) - } else { - pos.checked_sub(offset.unsigned_abs() as usize) - } - .and_then(|new_pos| { - if new_pos < valid_indices.len() { - Some(valid_indices[new_pos]) - } else { - None - } - }), - Err(pos) => if direction { - pos.checked_add(offset as usize) - } else if pos > 0 { - pos.checked_sub(offset.unsigned_abs() as usize) - } else { - None - } - .and_then(|new_pos| { - if new_pos < valid_indices.len() { - Some(valid_indices[new_pos]) - } else { - None - } - }), + if offset == 0 { + return Ok(Arc::clone(array)); + } + + // Arrays without NULLs do not necessarily have a null bitmap. + let Some(nulls) = array.nulls() else { + return shift_with_default_value(array, offset, default_value); + }; + + let shift = offset_magnitude(offset); + if shift >= array.len() { + return default_value.to_array_of_size(array.len()); + } + + let mut indices = ShiftIndexBuilder::new(array.len(), default_value.is_null()); + if is_lag { + let mut preceding = VecDeque::new(); + for index in 0..array.len() { + let result_index = if preceding.len() == shift { + preceding.front().copied() + } else { + None }; + indices.append_option(result_index); - match result_index { - Some(index) => ScalarValue::try_from_array(array, index), - None => Ok(default_value.clone()), + if nulls.is_valid(index) { + if preceding.len() == shift { + preceding.pop_front(); + } + preceding.push_back(index); + } + } + } else { + let mut following = VecDeque::new(); + let mut next_index = 0; + for index in 0..array.len() { + while following.front().is_some_and(|next| *next <= index) { + following.pop_front(); + } + next_index = next_index.max(index.saturating_add(1)); + while following.len() < shift && next_index < array.len() { + if nulls.is_valid(next_index) { + following.push_back(next_index); + } + next_index += 1; } - }) - .collect(); + indices.append_option(following.get(shift - 1).copied()); + } + } - let new_array = new_array_results?; - ScalarValue::iter_to_array(new_array) + indices.finish(array, default_value) } // TODO: change the original arrow::compute::kernels::window::shift impl to support an optional default value fn shift_with_default_value( @@ -688,7 +744,8 @@ impl PartitionEvaluator for WindowShiftEvaluator { mod tests { use super::*; use arrow::array::*; - use datafusion_common::cast::as_int32_array; + use arrow::datatypes::Int8Type; + use datafusion_common::cast::{as_dictionary_array, as_int32_array, as_string_array}; use datafusion_physical_expr::expressions::{Column, Literal}; fn test_i32_result( @@ -838,4 +895,136 @@ mod tests { .collect::(), ) } + + #[test] + fn test_evaluate_all_with_ignore_null() -> Result<()> { + let input: ArrayRef = Arc::new(Int32Array::from(vec![ + None, + Some(10), + None, + Some(20), + Some(30), + None, + ])); + + let cases = [ + ( + 1, + ScalarValue::Int32(None), + Int32Array::from(vec![ + None, + None, + Some(10), + Some(10), + Some(20), + Some(30), + ]), + ), + ( + -1, + ScalarValue::Int32(None), + Int32Array::from(vec![ + Some(10), + Some(20), + Some(20), + Some(30), + None, + None, + ]), + ), + ( + 2, + ScalarValue::Int32(Some(-1)), + Int32Array::from(vec![ + Some(-1), + Some(-1), + Some(-1), + Some(-1), + Some(10), + Some(20), + ]), + ), + ( + -2, + ScalarValue::Int32(Some(-1)), + Int32Array::from(vec![ + Some(20), + Some(30), + Some(30), + Some(-1), + Some(-1), + Some(-1), + ]), + ), + ( + 0, + ScalarValue::Int32(Some(-1)), + Int32Array::from(vec![None, Some(10), None, Some(20), Some(30), None]), + ), + ]; + + for (offset, default_value, expected) in cases { + let actual = evaluate_all_with_ignore_null( + &input, + offset, + &default_value, + offset > 0, + )?; + assert_eq!(expected, *as_int32_array(&actual)?); + } + Ok(()) + } + + #[test] + fn test_ignore_nulls_dictionary_with_bounded_keys() -> Result<()> { + let keys = + Int8Array::from_iter(std::iter::once(None).chain((0_i8..=127).map(Some))); + let values = + StringArray::from_iter_values((0..128).map(|index| format!("value-{index}"))); + let input: ArrayRef = Arc::new(DictionaryArray::::try_new( + keys, + Arc::new(values), + )?); + let default_value = ScalarValue::Dictionary( + Box::new(DataType::Int8), + Box::new(ScalarValue::Utf8(Some("default".to_string()))), + ); + + let actual = evaluate_all_with_ignore_null(&input, 1, &default_value, true)?; + let actual = as_dictionary_array::(actual.as_ref())?; + let values = as_string_array(actual.values().as_ref())?; + + assert_eq!(actual.len(), 129); + assert_eq!(values.len(), 128); + for index in 0..2 { + let key = actual.key(index).expect("non-null default"); + assert_eq!(values.value(key), "default"); + } + for index in 2..actual.len() { + let key = actual.key(index).expect("selected value"); + assert_eq!(values.value(key), format!("value-{}", index - 2)); + } + Ok(()) + } + + #[test] + fn test_ignore_nulls_without_null_bitmap() -> Result<()> { + let input = Int32Array::from(vec![1, 2, 3]); + assert!(input.nulls().is_none()); + let input: ArrayRef = Arc::new(input); + + for (offset, expected) in [ + (1, Int32Array::from(vec![None, Some(1), Some(2)])), + (-1, Int32Array::from(vec![Some(2), Some(3), None])), + ] { + let actual = evaluate_all_with_ignore_null( + &input, + offset, + &ScalarValue::Int32(None), + offset > 0, + )?; + assert_eq!(expected, *as_int32_array(&actual)?); + } + Ok(()) + } } diff --git a/datafusion/functions-window/src/nth_value.rs b/datafusion/functions-window/src/nth_value.rs index 437b4ecdb370a..df723772166a6 100644 --- a/datafusion/functions-window/src/nth_value.rs +++ b/datafusion/functions-window/src/nth_value.rs @@ -125,6 +125,14 @@ impl NthValue { } } +fn validate_nth_value_n(n: i64) -> Result { + if n == i64::MIN { + return exec_err!("The second argument of nth_value must not be i64::MIN"); + } + + Ok(n) +} + static FIRST_VALUE_DOCUMENTATION: LazyLock = LazyLock::new(|| { Documentation::builder( DOC_SECTION_ANALYTICAL, @@ -287,6 +295,7 @@ impl WindowUDFImpl for NthValue { .map(|v| get_signed_integer(&v)) { Some(Ok(n)) => { + let n = validate_nth_value_n(n)?; if partition_evaluator_args.is_reversed() { -n } else { @@ -660,4 +669,24 @@ mod tests { )?; Ok(()) } + + #[test] + fn nth_value_i64_min_returns_error() { + let expr = Arc::new(Column::new("c3", 0)) as Arc; + let n_value = Arc::new(Literal::new(ScalarValue::Int64(Some(i64::MIN)))) + as Arc; + + let err = NthValue::nth() + .partition_evaluator(PartitionEvaluatorArgs::new( + &[expr, n_value], + &[Field::new("f", DataType::Int32, true).into()], + false, + false, + )) + .unwrap_err(); + + assert!(err.to_string().starts_with( + "Execution error: The second argument of nth_value must not be i64::MIN" + )); + } } diff --git a/datafusion/functions/Cargo.toml b/datafusion/functions/Cargo.toml index d6a6693d862cc..a170e9f07c39f 100644 --- a/datafusion/functions/Cargo.toml +++ b/datafusion/functions/Cargo.toml @@ -67,7 +67,7 @@ name = "datafusion_functions" [dependencies] arrow = { workspace = true } arrow-buffer = { workspace = true } -base64 = { version = "0.22", optional = true } +base64 = { version = "0.23", optional = true } blake2 = { version = "^0.10.2", optional = true } blake3 = { version = "1.8", optional = true } chrono = { workspace = true } @@ -98,6 +98,16 @@ env_logger = { workspace = true } rand = { workspace = true } tokio = { workspace = true, features = ["macros", "rt", "sync"] } +[[bench]] +harness = false +name = "replace_scalar" +required-features = ["string_expressions"] + +[[bench]] +harness = false +name = "round_dense" +required-features = ["math_expressions"] + [[bench]] harness = false name = "ascii" @@ -153,6 +163,11 @@ harness = false name = "to_hex" required-features = ["string_expressions"] +[[bench]] +harness = false +name = "regexp_match" +required-features = ["regex_expressions"] + [[bench]] harness = false name = "regx" @@ -182,6 +197,11 @@ harness = false name = "date_trunc" required-features = ["datetime_expressions"] +[[bench]] +harness = false +name = "date_part" +required-features = ["datetime_expressions"] + [[bench]] harness = false name = "to_char" @@ -212,6 +232,11 @@ harness = false name = "atan2" required-features = ["math_expressions"] +[[bench]] +harness = false +name = "power" +required-features = ["math_expressions"] + [[bench]] harness = false name = "substr_index" @@ -230,6 +255,10 @@ required-features = ["string_expressions"] [[bench]] harness = false name = "upper" + +[[bench]] +harness = false +name = "upper_unicode" required-features = ["string_expressions"] [[bench]] @@ -286,6 +315,11 @@ harness = false name = "trunc" required-features = ["math_expressions"] +[[bench]] +harness = false +name = "trunc_precision" +required-features = ["math_expressions"] + [[bench]] harness = false name = "initcap" @@ -296,6 +330,11 @@ harness = false name = "find_in_set" required-features = ["unicode_expressions"] +[[bench]] +harness = false +name = "find_in_set_literal" +required-features = ["unicode_expressions"] + [[bench]] harness = false name = "contains" @@ -316,6 +355,15 @@ harness = false name = "regexp_count" required-features = ["regex_expressions"] +[[bench]] +harness = false +name = "regexp_instr" +required-features = ["regex_expressions"] + +[[bench]] +harness = false +name = "get_field" + [[bench]] harness = false name = "crypto" @@ -355,3 +403,8 @@ required-features = ["math_expressions"] harness = false name = "round" required-features = ["math_expressions"] + +[[bench]] +harness = false +name = "dictionary_encoding" +required-features = ["string_expressions", "unicode_expressions"] diff --git a/datafusion/functions/benches/concat.rs b/datafusion/functions/benches/concat.rs index 0fb910800e3bc..6736625be0365 100644 --- a/datafusion/functions/benches/concat.rs +++ b/datafusion/functions/benches/concat.rs @@ -23,8 +23,8 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::string::concat; -use rand::Rng; use rand::distr::Alphanumeric; +use rand::prelude::*; use std::hint::black_box; use std::sync::Arc; @@ -48,17 +48,20 @@ fn create_array_args_view(size: usize) -> Vec { ] } -fn generate_random_string(str_len: usize) -> String { - rand::rng() - .sample_iter(&Alphanumeric) +fn generate_random_string(rng: &mut StdRng, str_len: usize) -> String { + rng.sample_iter(&Alphanumeric) .take(str_len) .map(char::from) .collect() } -fn create_scalar_args(count: usize, str_len: usize) -> Vec { +fn create_scalar_args( + rng: &mut StdRng, + count: usize, + str_len: usize, +) -> Vec { std::iter::repeat_with(|| { - let s = generate_random_string(str_len); + let s = generate_random_string(rng, str_len); ColumnarValue::Scalar(ScalarValue::Utf8(Some(s))) }) .take(count) @@ -67,6 +70,7 @@ fn create_scalar_args(count: usize, str_len: usize) -> Vec { fn criterion_benchmark(c: &mut Criterion) { // Benchmark for array concat + let mut rng = StdRng::seed_from_u64(0); for size in [1024, 4096, 8192] { let args = create_array_args(size, 32); let arg_fields = args @@ -138,7 +142,7 @@ fn criterion_benchmark(c: &mut Criterion) { } // Benchmark for scalar concat - let scalar_args = create_scalar_args(10, 100); + let scalar_args = create_scalar_args(&mut rng, 10, 100); let scalar_arg_fields = scalar_args .iter() .enumerate() diff --git a/datafusion/functions/benches/concat_ws.rs b/datafusion/functions/benches/concat_ws.rs index 97d6d96411d73..d437f38773f78 100644 --- a/datafusion/functions/benches/concat_ws.rs +++ b/datafusion/functions/benches/concat_ws.rs @@ -23,8 +23,8 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::string::concat_ws; -use rand::Rng; use rand::distr::Alphanumeric; +use rand::prelude::*; use std::hint::black_box; use std::sync::Arc; @@ -38,9 +38,8 @@ fn create_array_args(size: usize, str_len: usize) -> Vec { ] } -fn generate_random_string(str_len: usize) -> String { - rand::rng() - .sample_iter(&Alphanumeric) +fn generate_random_string(rng: &mut StdRng, str_len: usize) -> String { + rng.sample_iter(&Alphanumeric) .take(str_len) .map(char::from) .collect() @@ -53,8 +52,9 @@ fn create_scalar_args(count: usize, str_len: usize) -> Vec { ",".to_string(), )))); + let mut rng = StdRng::seed_from_u64(0); for _ in 0..count { - let s = generate_random_string(str_len); + let s = generate_random_string(&mut rng, str_len); args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some(s)))); } args diff --git a/datafusion/functions/benches/date_bin.rs b/datafusion/functions/benches/date_bin.rs index 28dee96987261..bae1438fa4d5b 100644 --- a/datafusion/functions/benches/date_bin.rs +++ b/datafusion/functions/benches/date_bin.rs @@ -25,10 +25,9 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::date_bin; -use rand::Rng; -use rand::rngs::ThreadRng; +use rand::prelude::*; -fn timestamps(rng: &mut ThreadRng) -> TimestampSecondArray { +fn timestamps(rng: &mut StdRng) -> TimestampSecondArray { let mut seconds = vec![]; for _ in 0..1000 { seconds.push(rng.random_range(0..1_000_000)); @@ -39,7 +38,7 @@ fn timestamps(rng: &mut ThreadRng) -> TimestampSecondArray { fn criterion_benchmark(c: &mut Criterion) { c.bench_function("date_bin_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let timestamps_array = Arc::new(timestamps(&mut rng)) as ArrayRef; let batch_len = timestamps_array.len(); let interval = ColumnarValue::Scalar(ScalarValue::new_interval_dt(0, 1_000_000)); diff --git a/datafusion/functions/benches/date_part.rs b/datafusion/functions/benches/date_part.rs new file mode 100644 index 0000000000000..fb93ebd03b4e2 --- /dev/null +++ b/datafusion/functions/benches/date_part.rs @@ -0,0 +1,349 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::types::{IntervalDayTime, IntervalMonthDayNano}; +use arrow::array::{ + Array, ArrayRef, Date32Array, Date64Array, DurationNanosecondArray, + IntervalDayTimeArray, IntervalMonthDayNanoArray, IntervalYearMonthArray, + Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, + Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampNanosecondArray, TimestampSecondArray, +}; +use arrow::datatypes::{DataType, Field}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; +use datafusion_functions::datetime::date_part; +use rand::prelude::StdRng; +use rand::{Rng, SeedableRng}; + +const BATCH_SIZE: usize = 1000; +const TS_BOUND: i64 = 2_006_463_600; +const SEC_DAY: i64 = 86_400; +const DAYS_SINCE_EPOCH: i64 = TS_BOUND / SEC_DAY; + +fn generate_timestamp_ns_array(rng: &mut StdRng) -> TimestampNanosecondArray { + TimestampNanosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND * 1_000_000_000)) + .collect::>(), + ) +} + +fn generate_timestamp_us_array(rng: &mut StdRng) -> TimestampMicrosecondArray { + TimestampMicrosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND * 1_000_000)) + .collect::>(), + ) +} + +fn generate_timestamp_ms_array(rng: &mut StdRng) -> TimestampMillisecondArray { + TimestampMillisecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND * 1_000)) + .collect::>(), + ) +} + +fn generate_timestamp_s_array(rng: &mut StdRng) -> TimestampSecondArray { + TimestampSecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND)) + .collect::>(), + ) +} + +fn generate_date32_array(rng: &mut StdRng) -> Date32Array { + // Provide days since epoch + Date32Array::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..DAYS_SINCE_EPOCH as i32)) + .collect::>(), + ) +} + +fn generate_date64_array(rng: &mut StdRng) -> Date64Array { + // Provide milliseconds since epoch aligned to day boundaries + Date64Array::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..DAYS_SINCE_EPOCH) * SEC_DAY * 1_000) + .collect::>(), + ) +} + +fn generate_time32_second_array(rng: &mut StdRng) -> Time32SecondArray { + Time32SecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..SEC_DAY as i32)) + .collect::>(), + ) +} + +fn generate_time32_millisecond_array(rng: &mut StdRng) -> Time32MillisecondArray { + Time32MillisecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..(SEC_DAY * 1_000) as i32)) + .collect::>(), + ) +} + +fn generate_time64_microsecond_array(rng: &mut StdRng) -> Time64MicrosecondArray { + Time64MicrosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..SEC_DAY * 1_000_000)) + .collect::>(), + ) +} + +fn generate_time64_nanosecond_array(rng: &mut StdRng) -> Time64NanosecondArray { + Time64NanosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..SEC_DAY * 1_000_000_000)) + .collect::>(), + ) +} + +fn generate_interval_year_month_array(rng: &mut StdRng) -> IntervalYearMonthArray { + let years = 10; + IntervalYearMonthArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..12 * years)) + .collect::>(), + ) +} + +fn generate_interval_day_time_array(rng: &mut StdRng) -> IntervalDayTimeArray { + IntervalDayTimeArray::from( + (0..BATCH_SIZE) + .map(|_| IntervalDayTime { + days: rng.random_range(0..365), + milliseconds: rng.random_range(0..(SEC_DAY * 1_000) as i32), + }) + .collect::>(), + ) +} + +fn generate_interval_mdn_array(rng: &mut StdRng) -> IntervalMonthDayNanoArray { + IntervalMonthDayNanoArray::from( + (0..BATCH_SIZE) + .map(|_| IntervalMonthDayNano { + months: rng.random_range(0..12), + days: rng.random_range(0..365), + nanoseconds: rng.random_range(0..SEC_DAY * 1_000_000_000), + }) + .collect::>(), + ) +} + +fn generate_duration_nanosecond_array(rng: &mut StdRng) -> DurationNanosecondArray { + DurationNanosecondArray::from( + (0..BATCH_SIZE) + .map(|_| rng.random_range(0..TS_BOUND * 1_000_000_000)) + .collect::>(), + ) +} + +fn bench_date_part( + c: &mut Criterion, + udf: &Arc, + bench_name: &str, + part: &str, + array: ArrayRef, + return_type: DataType, +) { + let batch_len = array.len(); + let part_cv = ColumnarValue::Scalar(ScalarValue::Utf8(Some(part.to_string()))); + let array_cv = ColumnarValue::Array(array); + let return_field = Arc::new(Field::new("date_part", return_type, true)); + let arg_fields = vec![ + Field::new("a", part_cv.data_type(), true).into(), + Field::new("b", array_cv.data_type(), true).into(), + ]; + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_function(bench_name, |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![part_cv.clone(), array_cv.clone()], + arg_fields: arg_fields.clone(), + number_rows: batch_len, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .expect("date_part should work on valid values"), + ) + }) + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = StdRng::seed_from_u64(42); + + let ts_s = Arc::new(generate_timestamp_s_array(&mut rng)) as ArrayRef; + let ts_ms = Arc::new(generate_timestamp_ms_array(&mut rng)) as ArrayRef; + let ts_us = Arc::new(generate_timestamp_us_array(&mut rng)) as ArrayRef; + let ts_ns = Arc::new(generate_timestamp_ns_array(&mut rng)) as ArrayRef; + let time32_s = Arc::new(generate_time32_second_array(&mut rng)) as ArrayRef; + let time32_ms = Arc::new(generate_time32_millisecond_array(&mut rng)) as ArrayRef; + let time64_us = Arc::new(generate_time64_microsecond_array(&mut rng)) as ArrayRef; + let time64_ns = Arc::new(generate_time64_nanosecond_array(&mut rng)) as ArrayRef; + let interval_ym = Arc::new(generate_interval_year_month_array(&mut rng)) as ArrayRef; + let interval_dt = Arc::new(generate_interval_day_time_array(&mut rng)) as ArrayRef; + let interval_mdn = Arc::new(generate_interval_mdn_array(&mut rng)) as ArrayRef; + let duration_ns = Arc::new(generate_duration_nanosecond_array(&mut rng)) as ArrayRef; + let date32 = Arc::new(generate_date32_array(&mut rng)) as ArrayRef; + let date64 = Arc::new(generate_date64_array(&mut rng)) as ArrayRef; + + let udf = date_part(); + + for part in ["year", "month", "week", "day", "hour", "minute"] { + for (name, array) in + [("s", &ts_s), ("ms", &ts_ms), ("us", &ts_us), ("ns", &ts_ns)] + { + bench_date_part( + c, + &udf, + &format!("date_part_{part}_{name}_1000"), + part, + Arc::clone(array), + DataType::Int32, + ); + } + } + for part in ["year", "month", "week", "day"] { + bench_date_part( + c, + &udf, + &format!("date_part_{part}_date32_1000"), + part, + Arc::clone(&date32), + DataType::Int32, + ); + bench_date_part( + c, + &udf, + &format!("date_part_{part}_date64_1000"), + part, + Arc::clone(&date64), + DataType::Int32, + ); + } + + for part in ["second", "millisecond", "microsecond"] { + for (name, array) in + [("s", &ts_s), ("ms", &ts_ms), ("us", &ts_us), ("ns", &ts_ns)] + { + bench_date_part( + c, + &udf, + &format!("date_part_{part}_{name}_1000"), + part, + Arc::clone(array), + DataType::Int32, + ); + } + bench_date_part( + c, + &udf, + &format!("date_part_{part}_date32_1000"), + part, + Arc::clone(&date32), + DataType::Int32, + ); + bench_date_part( + c, + &udf, + &format!("date_part_{part}_date64_1000"), + part, + Arc::clone(&date64), + DataType::Int32, + ); + } + + for (name, array) in [("s", &ts_s), ("ms", &ts_ms), ("us", &ts_us), ("ns", &ts_ns)] { + bench_date_part( + c, + &udf, + &format!("date_part_nanosecond_{name}_1000"), + "nanosecond", + Arc::clone(array), + DataType::Int64, + ); + } + bench_date_part( + c, + &udf, + "date_part_nanosecond_date32_1000", + "nanosecond", + Arc::clone(&date32), + DataType::Int64, + ); + bench_date_part( + c, + &udf, + "date_part_nanosecond_date64_1000", + "nanosecond", + Arc::clone(&date64), + DataType::Int64, + ); + + for (name, array) in [ + ("s", &ts_s), + ("ms", &ts_ms), + ("us", &ts_us), + ("ns", &ts_ns), + ("date32", &date32), + ("date64", &date64), + ("time32_s", &time32_s), + ("time32_ms", &time32_ms), + ("time64_us", &time64_us), + ("time64_ns", &time64_ns), + ("interval_ym", &interval_ym), + ("interval_dt", &interval_dt), + ("interval_mdn", &interval_mdn), + ("duration_ns", &duration_ns), + ] { + bench_date_part( + c, + &udf, + &format!("date_part_epoch_{name}_1000"), + "epoch", + Arc::clone(array), + DataType::Float64, + ); + } + + for part in ["quarter", "isoyear", "doy", "dow", "isodow"] { + bench_date_part( + c, + &udf, + &format!("date_part_{part}_timestamp_ns_1000"), + part, + Arc::clone(&ts_ns), + DataType::Int32, + ); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/date_trunc.rs b/datafusion/functions/benches/date_trunc.rs index 0668a1cc5085c..e2372fff2a02e 100644 --- a/datafusion/functions/benches/date_trunc.rs +++ b/datafusion/functions/benches/date_trunc.rs @@ -18,52 +18,64 @@ use std::hint::black_box; use std::sync::Arc; -use arrow::array::{Array, ArrayRef, TimestampSecondArray}; +use arrow::array::{Array, ArrayRef, TimestampNanosecondArray, TimestampSecondArray}; use arrow::datatypes::Field; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs}; use datafusion_functions::datetime::date_trunc; -use rand::Rng; -use rand::rngs::ThreadRng; +use rand::rngs::StdRng; +use rand::{Rng, SeedableRng}; -fn timestamps(rng: &mut ThreadRng) -> TimestampSecondArray { - let mut seconds = vec![]; - for _ in 0..1000 { - seconds.push(rng.random_range(0..1_000_000)); - } +const NUM_ROWS: usize = 1000; +const NANOS_PER_SECOND: i64 = 1_000_000_000; +/// Roughly 30 years, so that values span many months, quarters and years. +const RANGE_SECONDS: i64 = 30 * 365 * 24 * 60 * 60; - TimestampSecondArray::from(seconds) +fn seedable_rng() -> StdRng { + StdRng::seed_from_u64(42) } -fn criterion_benchmark(c: &mut Criterion) { - c.bench_function("date_trunc_minute_1000", |b| { - let mut rng = rand::rng(); - let timestamps_array = Arc::new(timestamps(&mut rng)) as ArrayRef; - let batch_len = timestamps_array.len(); - let precision = - ColumnarValue::Scalar(ScalarValue::Utf8(Some("minute".to_string()))); - let timestamps = ColumnarValue::Array(timestamps_array); - let udf = date_trunc(); - let args = vec![precision, timestamps]; - let arg_fields = args - .iter() - .enumerate() - .map(|(idx, arg)| { - Field::new(format!("arg_{idx}"), arg.data_type(), true).into() - }) - .collect::>(); +fn second_timestamps() -> TimestampSecondArray { + let mut rng = seedable_rng(); + (0..NUM_ROWS) + .map(|_| Some(rng.random_range(0..1_000_000i64))) + .collect() +} - let scalar_arguments = vec![None; arg_fields.len()]; - let return_field = udf - .return_field_from_args(ReturnFieldArgs { - arg_fields: &arg_fields, - scalar_arguments: &scalar_arguments, - }) - .unwrap(); - let config_options = Arc::new(ConfigOptions::default()); +fn nanosecond_timestamps() -> TimestampNanosecondArray { + let mut rng = seedable_rng(); + (0..NUM_ROWS) + .map(|_| { + let seconds = rng.random_range(-RANGE_SECONDS..RANGE_SECONDS); + Some(seconds * NANOS_PER_SECOND + rng.random_range(0..NANOS_PER_SECOND)) + }) + .collect() +} + +fn run_benchmark(c: &mut Criterion, name: &str, granularity: &str, array: ArrayRef) { + let batch_len = array.len(); + let precision = + ColumnarValue::Scalar(ScalarValue::Utf8(Some(granularity.to_string()))); + let udf = date_trunc(); + let args = vec![precision, ColumnarValue::Array(array)]; + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into()) + .collect::>(); + let scalar_arguments = vec![None; arg_fields.len()]; + let return_field = udf + .return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalar_arguments, + }) + .unwrap(); + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_function(name, |b| { b.iter(|| { black_box( udf.invoke_with_args(ScalarFunctionArgs { @@ -79,5 +91,23 @@ fn criterion_benchmark(c: &mut Criterion) { }); } +fn criterion_benchmark(c: &mut Criterion) { + let seconds: ArrayRef = Arc::new(second_timestamps()); + run_benchmark(c, "date_trunc_minute_1000", "minute", Arc::clone(&seconds)); + run_benchmark(c, "date_trunc_month_second_1000", "month", seconds); + + // Coarse granularities on an untimezoned array: these need calendar + // arithmetic rather than a plain division. + let nanos: ArrayRef = Arc::new(nanosecond_timestamps()); + for granularity in ["week", "month", "quarter", "year"] { + run_benchmark( + c, + &format!("date_trunc_{granularity}_nanos_1000"), + granularity, + Arc::clone(&nanos), + ); + } +} + criterion_group!(benches, criterion_benchmark); criterion_main!(benches); diff --git a/datafusion/functions/benches/dictionary_encoding.rs b/datafusion/functions/benches/dictionary_encoding.rs new file mode 100644 index 0000000000000..4ba04a4940e61 --- /dev/null +++ b/datafusion/functions/benches/dictionary_encoding.rs @@ -0,0 +1,104 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, DictionaryArray}; +use arrow::compute::cast; +use arrow::datatypes::{Field, Int32Type}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::type_coercion::functions::fields_with_udf; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; + +const NUM_ROWS: usize = 8_192; +const DICTIONARY_CARDINALITIES: [usize; 4] = [10, 100, 1_000, 8_192]; + +fn create_string_dictionary(cardinality: usize) -> ArrayRef { + let values = (0..NUM_ROWS) + .map(|index| Some(format!("value_{:04}", index % cardinality))) + .collect::>(); + Arc::new( + values + .iter() + .map(|value| value.as_deref()) + .collect::>(), + ) +} + +fn benchmark_dictionary_string_udfs(c: &mut Criterion) { + let udfs = [ + ("ascii", datafusion_functions::string::ascii()), + ("bit_length", datafusion_functions::string::bit_length()), + ("btrim", datafusion_functions::string::btrim()), + ( + "character_length", + datafusion_functions::unicode::character_length(), + ), + ("initcap", datafusion_functions::unicode::initcap()), + ("ltrim", datafusion_functions::string::ltrim()), + ("octet_length", datafusion_functions::string::octet_length()), + ("reverse", datafusion_functions::unicode::reverse()), + ("rtrim", datafusion_functions::string::rtrim()), + ]; + let config_options = Arc::new(ConfigOptions::default()); + + for cardinality in DICTIONARY_CARDINALITIES { + let dictionary = create_string_dictionary(cardinality); + let mut group = c.benchmark_group(format!( + "dictionary_encoding/string/cardinality_{cardinality}" + )); + for (name, udf) in &udfs { + let input_field = + Field::new("a", dictionary.data_type().clone(), false).into(); + let coerced_field = fields_with_udf(&[input_field], udf.as_ref()) + .unwrap() + .into_iter() + .next() + .unwrap(); + let coerced_type = coerced_field.data_type(); + let return_type = + udf.return_type(std::slice::from_ref(coerced_type)).unwrap(); + let return_field = Field::new("f", return_type, false).into(); + let input = if dictionary.data_type() == coerced_type { + Arc::clone(&dictionary) + } else { + cast(dictionary.as_ref(), coerced_type).unwrap() + }; + + group.bench_function(*name, |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::clone(&input))], + arg_fields: vec![Arc::clone(&coerced_field)], + number_rows: NUM_ROWS, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + } + group.finish(); + } +} + +criterion_group!(benches, benchmark_dictionary_string_udfs); +criterion_main!(benches); diff --git a/datafusion/functions/benches/encoding.rs b/datafusion/functions/benches/encoding.rs index 0b8f0c5c51a58..451baff5183fd 100644 --- a/datafusion/functions/benches/encoding.rs +++ b/datafusion/functions/benches/encoding.rs @@ -27,10 +27,35 @@ use std::sync::Arc; fn criterion_benchmark(c: &mut Criterion) { let decode = encoding::decode(); + let encode = encoding::encode(); let config_options = Arc::new(ConfigOptions::default()); for size in [1024, 4096, 8192] { let bin_array = Arc::new(create_binary_array::(size, 0.2)); + + c.bench_function(&format!("hex_encode/{size}"), |b| { + let method = ColumnarValue::Scalar("hex".into()); + let arg_fields = vec![ + Field::new("a", bin_array.data_type().to_owned(), true).into(), + Field::new("b", method.data_type().to_owned(), true).into(), + ]; + let args = vec![ColumnarValue::Array(bin_array.clone()), method]; + let return_field = Field::new("f", DataType::Utf8, true).into(); + + b.iter(|| { + black_box( + encode + .invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); c.bench_function(&format!("base64_decode/{size}"), |b| { let method = ColumnarValue::Scalar("base64".into()); let encoded = encoding::encode() diff --git a/datafusion/functions/benches/find_in_set_literal.rs b/datafusion/functions/benches/find_in_set_literal.rs new file mode 100644 index 0000000000000..013c7c2081668 --- /dev/null +++ b/datafusion/functions/benches/find_in_set_literal.rs @@ -0,0 +1,98 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks the `find_in_set(column, constant_list)` path where the set is a +//! scalar literal. A long list exercises the pre-built lookup; a short list +//! stays on the per-row linear scan. + +use arrow::array::StringArray; +use arrow::datatypes::{DataType, Field}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use rand::prelude::StdRng; +use rand::{Rng, SeedableRng}; +use std::hint::black_box; +use std::sync::Arc; + +const N_ROWS: usize = 8192; + +/// Builds a string column whose values are drawn from `entries` plus a small +/// fraction of misses, so both hits and misses are exercised. +fn build_column(entries: &[String]) -> StringArray { + let mut rng = StdRng::seed_from_u64(42); + let values: Vec> = (0..N_ROWS) + .map(|_| { + let r = rng.random::(); + if r < 0.1 { + None + } else if r < 0.4 { + Some("__miss__".to_string()) + } else { + let idx = rng.random_range(0..entries.len()); + Some(entries[idx].clone()) + } + }) + .collect(); + StringArray::from(values) +} + +fn bench_case(c: &mut Criterion, label: &str, num_entries: usize) { + let find_in_set = datafusion_functions::unicode::find_in_set(); + let entries: Vec = (0..num_entries).map(|i| format!("item{i}")).collect(); + let list = entries.join(","); + + let column = build_column(&entries); + let args = vec![ + ColumnarValue::Array(Arc::new(column)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(list))), + ]; + let arg_fields = args + .iter() + .map(|arg| Field::new("a", arg.data_type().clone(), true).into()) + .collect::>(); + let return_field = Arc::new(Field::new("f", DataType::Int32, true)); + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_with_input( + BenchmarkId::new("find_in_set_literal", label), + &num_entries, + |b, _| { + b.iter(|| { + black_box(find_in_set.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: N_ROWS, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + })) + }) + }, + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + // Short list stays on the linear scan (below the lookup threshold). + bench_case(c, "short_list_4", 4); + // Long lists exercise the pre-built lookup. + bench_case(c, "long_list_64", 64); + bench_case(c, "long_list_256", 256); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/gcd.rs b/datafusion/functions/benches/gcd.rs index 3c72a46e6643d..ca49415b0f679 100644 --- a/datafusion/functions/benches/gcd.rs +++ b/datafusion/functions/benches/gcd.rs @@ -25,12 +25,12 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::math::gcd; -use rand::Rng; +use rand::prelude::*; use std::hint::black_box; use std::sync::Arc; fn generate_i64_array(n_rows: usize) -> ArrayRef { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let values = (0..n_rows) .map(|_| rng.random_range(0..1000)) .collect::>(); diff --git a/datafusion/functions/benches/get_field.rs b/datafusion/functions/benches/get_field.rs new file mode 100644 index 0000000000000..8a5fd0a1e2fa9 --- /dev/null +++ b/datafusion/functions/benches/get_field.rs @@ -0,0 +1,95 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +extern crate criterion; + +use arrow::array::{ArrayRef, Int32Builder, MapBuilder, StringBuilder}; +use arrow::datatypes::{DataType, Field}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_functions::core::get_field; +use std::hint::black_box; +use std::sync::Arc; + +/// A map array with `size` rows, each holding `entries` key/value pairs. +/// Every tenth row is null. +fn map_array(size: usize, entries: usize) -> ArrayRef { + let mut builder = MapBuilder::new(None, StringBuilder::new(), Int32Builder::new()); + for row in 0..size { + if row % 10 == 0 { + builder.append(false).unwrap(); + continue; + } + for entry in 0..entries { + builder.keys().append_value(format!("key_{entry}")); + builder.values().append_value((row * entry) as i32); + } + builder.append(true).unwrap(); + } + Arc::new(builder.finish()) +} + +fn bench_get_field( + c: &mut Criterion, + name: &str, + size: usize, + entries: usize, + key: &str, +) { + let udf = get_field(); + let args = vec![ + ColumnarValue::Array(map_array(size, entries)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(key.to_string()))), + ]; + let arg_fields = vec![ + Field::new("map", args[0].data_type(), true).into(), + Field::new("key", DataType::Utf8, false).into(), + ]; + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_function(name, |b| { + b.iter(|| { + black_box( + udf.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Int32, true).into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + // First key: the match is found immediately, so the per-row overhead + // dominates. + bench_get_field(c, "get_field_map_1024_entries_4_first", 1024, 4, "key_0"); + // Last key: every entry of the row is compared before the match. + bench_get_field(c, "get_field_map_1024_entries_4_last", 1024, 4, "key_3"); + bench_get_field(c, "get_field_map_1024_entries_16_last", 1024, 16, "key_15"); + // Key that is not present in any row. + bench_get_field(c, "get_field_map_1024_entries_4_missing", 1024, 4, "key_9"); + bench_get_field(c, "get_field_map_8192_entries_4_last", 8192, 4, "key_3"); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/lcm.rs b/datafusion/functions/benches/lcm.rs index 247c0ec749d15..5a4e5d2bced7d 100644 --- a/datafusion/functions/benches/lcm.rs +++ b/datafusion/functions/benches/lcm.rs @@ -24,12 +24,12 @@ use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::math::lcm; -use rand::Rng; +use rand::prelude::*; use std::hint::black_box; use std::sync::Arc; fn generate_i64_array(n_rows: usize) -> ArrayRef { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let values = (0..n_rows) .map(|_| rng.random_range(0..1000)) .collect::>(); diff --git a/datafusion/functions/benches/make_date.rs b/datafusion/functions/benches/make_date.rs index 1c7b61ec60497..2e82a871eb0d6 100644 --- a/datafusion/functions/benches/make_date.rs +++ b/datafusion/functions/benches/make_date.rs @@ -25,10 +25,9 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::make_date; -use rand::Rng; -use rand::rngs::ThreadRng; +use rand::prelude::*; -fn years(rng: &mut ThreadRng) -> Int32Array { +fn years(rng: &mut StdRng) -> Int32Array { let mut years = vec![]; for _ in 0..8192 { years.push(rng.random_range(1900..2050)); @@ -37,7 +36,7 @@ fn years(rng: &mut ThreadRng) -> Int32Array { Int32Array::from(years) } -fn months(rng: &mut ThreadRng) -> Int32Array { +fn months(rng: &mut StdRng) -> Int32Array { let mut months = vec![]; for _ in 0..8192 { months.push(rng.random_range(1..13)); @@ -46,7 +45,7 @@ fn months(rng: &mut ThreadRng) -> Int32Array { Int32Array::from(months) } -fn days(rng: &mut ThreadRng) -> Int32Array { +fn days(rng: &mut StdRng) -> Int32Array { let mut days = vec![]; for _ in 0..8192 { days.push(rng.random_range(1..29)); @@ -56,7 +55,7 @@ fn days(rng: &mut ThreadRng) -> Int32Array { } fn criterion_benchmark(c: &mut Criterion) { c.bench_function("make_date_col_col_col_8192", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let years_array = Arc::new(years(&mut rng)) as ArrayRef; let batch_len = years_array.len(); let years = ColumnarValue::Array(years_array); @@ -86,7 +85,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("make_date_scalar_col_col_8192", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let year = ColumnarValue::Scalar(ScalarValue::Int32(Some(2025))); let months_arr = Arc::new(months(&mut rng)) as ArrayRef; let batch_len = months_arr.len(); @@ -116,7 +115,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("make_date_scalar_scalar_col_8192", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let year = ColumnarValue::Scalar(ScalarValue::Int32(Some(2025))); let month = ColumnarValue::Scalar(ScalarValue::Int32(Some(11))); let day_arr = Arc::new(days(&mut rng)); diff --git a/datafusion/functions/benches/nanvl.rs b/datafusion/functions/benches/nanvl.rs index 206eebd81eb81..d3d2c7ebff998 100644 --- a/datafusion/functions/benches/nanvl.rs +++ b/datafusion/functions/benches/nanvl.rs @@ -108,6 +108,80 @@ fn criterion_benchmark(c: &mut Criterion) { bench.iter(|| black_box(nanvl_fn.invoke_with_args(args.clone()).unwrap())) }); } + + // Partially-null array benchmarks exercise the null-aware match arms that + // the fully-populated benchmarks above never reach: only-x-null, + // only-y-null, and both-null. + let bench_pair = + |c: &mut Criterion, name: &str, x: ArrayRef, y: ArrayRef, size: usize| { + c.bench_function(name, |bench| { + let args = ScalarFunctionArgs { + args: vec![ + ColumnarValue::Array(Arc::clone(&x)), + ColumnarValue::Array(Arc::clone(&y)), + ], + arg_fields: vec![ + Field::new("a", DataType::Float64, true).into(), + Field::new("b", DataType::Float64, true).into(), + ], + number_rows: size, + return_field: Field::new("f", DataType::Float64, true).into(), + config_options: Arc::clone(&config_options), + }; + bench.iter(|| black_box(nanvl_fn.invoke_with_args(args.clone()).unwrap())) + }); + }; + + for size in [1024, 4096, 8192] { + // `x` mixes non-NaN, NaN, and null so every code path is taken. + let x_nulls: ArrayRef = Arc::new(Float64Array::from( + (0..size) + .map(|i| match i % 3 { + 0 => Some(1.0), + 1 => Some(f64::NAN), + _ => None, + }) + .collect::>(), + )); + // `x` without nulls, alternating non-NaN and NaN. + let x_full: ArrayRef = Arc::new(Float64Array::from( + (0..size) + .map(|i| if i % 2 == 0 { 1.0 } else { f64::NAN }) + .collect::>(), + )); + // `y` with roughly a quarter nulls. + let y_nulls: ArrayRef = Arc::new(Float64Array::from( + (0..size) + .map(|i| if i % 4 == 3 { None } else { Some(2.0) }) + .collect::>(), + )); + let y_full: ArrayRef = Arc::new(Float64Array::from(vec![2.0; size])); + + // (Some, None): only `x` has nulls. + bench_pair( + c, + &format!("nanvl/array_f64_x_nulls/{size}"), + Arc::clone(&x_nulls), + Arc::clone(&y_full), + size, + ); + // (None, Some): only `y` has nulls. + bench_pair( + c, + &format!("nanvl/array_f64_y_nulls/{size}"), + Arc::clone(&x_full), + Arc::clone(&y_nulls), + size, + ); + // (Some, Some): both inputs have nulls. + bench_pair( + c, + &format!("nanvl/array_f64_both_nulls/{size}"), + Arc::clone(&x_nulls), + Arc::clone(&y_nulls), + size, + ); + } } criterion_group!(benches, criterion_benchmark); diff --git a/datafusion/functions/benches/overlay.rs b/datafusion/functions/benches/overlay.rs index 4554cc435e738..0b7fff5989d1f 100644 --- a/datafusion/functions/benches/overlay.rs +++ b/datafusion/functions/benches/overlay.rs @@ -21,24 +21,39 @@ use arrow::datatypes::{DataType, Field}; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; -use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; use helper::gen_string_array; use std::hint::black_box; use std::sync::Arc; -fn criterion_benchmark(c: &mut Criterion) { - const N_ROWS: usize = 8192; +#[expect(clippy::too_many_arguments)] +fn bench_overlay( + c: &mut Criterion, + name: &str, + overlay: &ScalarUDF, + n_rows: usize, + null_density: f32, + utf8_density: f32, + is_string_view: bool, + with_for: bool, +) { const STR_LEN: usize = 128; - let overlay = datafusion_functions::core::overlay(); - let config_options = Arc::new(ConfigOptions::default()); - - let mut args = gen_string_array(N_ROWS, STR_LEN, 0.1, 0.5, false); - args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some( - "DataFusion".to_string(), - )))); + let mut args = + gen_string_array(n_rows, STR_LEN, null_density, utf8_density, is_string_view); + // The substring scalar's type must match the string column's type (the + // function dispatches per-type without coercion). + let substr = "DataFusion".to_string(); + let substr_scalar = if is_string_view { + ScalarValue::Utf8View(Some(substr)) + } else { + ScalarValue::Utf8(Some(substr)) + }; + args.push(ColumnarValue::Scalar(substr_scalar)); args.push(ColumnarValue::Scalar(ScalarValue::Int64(Some(32)))); - args.push(ColumnarValue::Scalar(ScalarValue::Int64(Some(8)))); + if with_for { + args.push(ColumnarValue::Scalar(ScalarValue::Int64(Some(8)))); + } let arg_fields = args .iter() @@ -46,15 +61,16 @@ fn criterion_benchmark(c: &mut Criterion) { .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into()) .collect::>(); let return_field = Arc::new(Field::new("f", DataType::Utf8, true)); + let config_options = Arc::new(ConfigOptions::default()); - c.bench_function("overlay_StringArray_utf8_scalar_args", |b| { + c.bench_function(name, |b| { b.iter(|| { black_box( overlay .invoke_with_args(ScalarFunctionArgs { args: args.clone(), arg_fields: arg_fields.clone(), - number_rows: N_ROWS, + number_rows: n_rows, return_field: Arc::clone(&return_field), config_options: Arc::clone(&config_options), }) @@ -64,5 +80,121 @@ fn criterion_benchmark(c: &mut Criterion) { }); } +fn criterion_benchmark(c: &mut Criterion) { + const N_ROWS: usize = 8192; + const MIXED_UTF8: f32 = 0.5; + let overlay = datafusion_functions::core::overlay(); + + // Null-density variants on StringArray (mixed ASCII/UTF-8, 4-arg form). + bench_overlay( + c, + "overlay_StringArray_low_nulls", + &overlay, + N_ROWS, + 0.1, + MIXED_UTF8, + false, + true, + ); + bench_overlay( + c, + "overlay_StringArray_high_nulls", + &overlay, + N_ROWS, + 0.9, + MIXED_UTF8, + false, + true, + ); + bench_overlay( + c, + "overlay_StringArray_no_nulls", + &overlay, + N_ROWS, + 0.0, + MIXED_UTF8, + false, + true, + ); + + // Content variants on StringArray (no nulls, 4-arg form). Pair against + // `overlay_StringArray_no_nulls` to isolate the impact of UTF-8 density. + bench_overlay( + c, + "overlay_StringArray_ascii", + &overlay, + N_ROWS, + 0.0, + 0.0, + false, + true, + ); + bench_overlay( + c, + "overlay_StringArray_all_utf8", + &overlay, + N_ROWS, + 0.0, + 1.0, + false, + true, + ); + + // 3-arg form (no FOR clause), where the replace length is derived from + // the substring per row. + bench_overlay( + c, + "overlay_StringArray_no_for", + &overlay, + N_ROWS, + 0.0, + MIXED_UTF8, + false, + false, + ); + + // StringViewArray counterparts. + bench_overlay( + c, + "overlay_StringViewArray_low_nulls", + &overlay, + N_ROWS, + 0.1, + MIXED_UTF8, + true, + true, + ); + bench_overlay( + c, + "overlay_StringViewArray_ascii", + &overlay, + N_ROWS, + 0.0, + 0.0, + true, + true, + ); + bench_overlay( + c, + "overlay_StringViewArray_all_utf8", + &overlay, + N_ROWS, + 0.0, + 1.0, + true, + true, + ); + bench_overlay( + c, + "overlay_StringViewArray_no_for", + &overlay, + N_ROWS, + 0.0, + MIXED_UTF8, + true, + false, + ); +} + criterion_group!(benches, criterion_benchmark); criterion_main!(benches); diff --git a/datafusion/functions/benches/pad.rs b/datafusion/functions/benches/pad.rs index c71d5a7161a66..78ebf12236a70 100644 --- a/datafusion/functions/benches/pad.rs +++ b/datafusion/functions/benches/pad.rs @@ -28,8 +28,8 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::unicode; -use rand::Rng; -use rand::distr::{Distribution, Uniform}; +use rand::distr::Uniform; +use rand::prelude::*; use std::hint::black_box; use std::sync::Arc; use std::time::Duration; @@ -51,7 +51,7 @@ fn create_unicode_string_array( size: usize, null_density: f32, ) -> arrow::array::GenericStringArray { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let mut builder = GenericStringBuilder::::new(); for i in 0..size { if rng.random::() < null_density { @@ -67,7 +67,7 @@ fn create_unicode_string_view_array( size: usize, null_density: f32, ) -> arrow::array::StringViewArray { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let mut builder = StringViewBuilder::with_capacity(size); for i in 0..size { if rng.random::() < null_density { @@ -104,7 +104,7 @@ where dist: Uniform::new_inclusive::(0, len as i64), }; - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); (0..size) .map(|_| { if rng.random::() < null_density { diff --git a/datafusion/functions/benches/power.rs b/datafusion/functions/benches/power.rs new file mode 100644 index 0000000000000..5336e42ebe59b --- /dev/null +++ b/datafusion/functions/benches/power.rs @@ -0,0 +1,140 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Microbenchmark for `power(decimal_array, int_*)`. +//! +//! Covers both array- and scalar-shaped integer exponents on a Decimal +//! base. Both shapes are dispatched to the native per-row decimal kernel; +//! the bench guards against any future change that routes either shape +//! through a Float64 round-trip, which is measurably slower than the +//! decimal kernel for the cases the kernel can handle. + +extern crate criterion; + +use arrow::array::{Decimal128Array, Int64Array}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; +use datafusion_functions::math::power; +use std::hint::black_box; +use std::sync::Arc; + +fn make_decimal_array(size: usize, precision: u8, scale: i8) -> Decimal128Array { + // Use a fixed unscaled value (250) so the bench is independent of `scale`. + // The four-arm dispatch in `power` only cares about the Decimal variant + // and the exponent's shape, not the numeric value. + let arr = Decimal128Array::from(vec![250i128; size]); + arr.with_precision_and_scale(precision, scale).unwrap() +} + +fn make_int_array(size: usize, value: i64) -> Int64Array { + Int64Array::from(vec![value; size]) +} + +fn run_power( + power_fn: &ScalarUDF, + args: &[ColumnarValue], + arg_fields: &[FieldRef], + return_field: &FieldRef, + config_options: &Arc, + num_rows: usize, +) { + black_box( + power_fn + .invoke_with_args(ScalarFunctionArgs { + args: args.to_vec(), + arg_fields: arg_fields.to_vec(), + number_rows: num_rows, + return_field: Arc::clone(return_field), + config_options: Arc::clone(config_options), + }) + .unwrap(), + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + let power_fn = power(); + let config_options = Arc::new(ConfigOptions::default()); + let precision: u8 = 20; + let scale: i8 = 2; + let decimal_ty = DataType::Decimal128(precision, scale); + + // Exponents are bounded by what the native decimal kernel can handle + // without overflowing the i128 intermediate; see + // + let exponents = [2i64, 4, 8]; + + for size in [1024usize, 8192] { + let base_arr = Arc::new(make_decimal_array(size, precision, scale)); + let base_field: FieldRef = Field::new("base", decimal_ty.clone(), true).into(); + let exp_field: FieldRef = Field::new("exp", DataType::Int64, true).into(); + let return_field: FieldRef = Field::new("r", decimal_ty.clone(), true).into(); + let arg_fields = vec![base_field, exp_field]; + + for &exp in &exponents { + let exp_arr = Arc::new(make_int_array(size, exp)); + let array_args = vec![ + ColumnarValue::Array(base_arr.clone()), + ColumnarValue::Array(exp_arr), + ]; + c.bench_function( + &format!( + "power decimal({precision},{scale}) array x int array, exp={exp}, n={size}" + ), + |b| { + b.iter(|| { + run_power( + &power_fn, + &array_args, + &arg_fields, + &return_field, + &config_options, + size, + ) + }) + }, + ); + + let scalar_args = vec![ + ColumnarValue::Array(base_arr.clone()), + ColumnarValue::Scalar(ScalarValue::Int64(Some(exp))), + ]; + c.bench_function( + &format!( + "power decimal({precision},{scale}) array x int scalar, exp={exp}, n={size}" + ), + |b| { + b.iter(|| { + run_power( + &power_fn, + &scalar_args, + &arg_fields, + &return_field, + &config_options, + size, + ) + }) + }, + ); + } + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/regexp_instr.rs b/datafusion/functions/benches/regexp_instr.rs new file mode 100644 index 0000000000000..9ac630d8c4b6e --- /dev/null +++ b/datafusion/functions/benches/regexp_instr.rs @@ -0,0 +1,99 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::Int64Array; +use arrow::array::OffsetSizeTrait; +use arrow::datatypes::{DataType, Field}; +use arrow::util::bench_util::create_string_array_with_len; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{DataFusionError, ScalarValue}; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_functions::regex; +use std::hint::black_box; +use std::sync::Arc; + +fn create_args( + size: usize, + str_len: usize, + with_start: bool, +) -> Vec { + let string_array = Arc::new(create_string_array_with_len::(size, 0.1, str_len)); + let pattern = ColumnarValue::Scalar(ScalarValue::Utf8(Some("a".to_string()))); + + if with_start { + let start_array = Arc::new(Int64Array::from( + (0..size).map(|i| (i % 10 + 1) as i64).collect::>(), + )); + vec![ + ColumnarValue::Array(string_array), + pattern, + ColumnarValue::Array(start_array), + ] + } else { + vec![ColumnarValue::Array(string_array), pattern] + } +} + +fn invoke_regexp_instr_with_args( + args: Vec, + number_rows: usize, +) -> Result { + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into()) + .collect::>(); + let config_options = Arc::new(ConfigOptions::default()); + + regex::regexp_instr().invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Field::new("f", DataType::Int64, true).into(), + config_options: Arc::clone(&config_options), + }) +} + +fn criterion_benchmark(c: &mut Criterion) { + let size = 1024; + + for str_len in [32, 128] { + let args = create_args::(size, str_len, false); + c.bench_function( + &format!("regexp_instr_no_start [size={size}, str_len={str_len}]"), + |b| { + b.iter(|| { + black_box(invoke_regexp_instr_with_args(args.clone(), size).unwrap()) + }) + }, + ); + + let args = create_args::(size, str_len, true); + c.bench_function( + &format!("regexp_instr_with_start [size={size}, str_len={str_len}]"), + |b| { + b.iter(|| { + black_box(invoke_regexp_instr_with_args(args.clone(), size).unwrap()) + }) + }, + ); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/regexp_match.rs b/datafusion/functions/benches/regexp_match.rs new file mode 100644 index 0000000000000..d5929df07c81f --- /dev/null +++ b/datafusion/functions/benches/regexp_match.rs @@ -0,0 +1,137 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks `regexp_match` through `invoke_with_args`, which is how a query +//! plan calls it. The pattern (and flags) are literals, as in +//! `regexp_match(col, '[a-z]+')`. + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, StringArray}; +use arrow::compute::cast; +use arrow::datatypes::{DataType, Field}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_functions::regex::regexpmatch::RegexpMatchFunc; +use rand::Rng; +use rand::distr::Alphanumeric; +use rand::rngs::ThreadRng; + +const SIZE: usize = 1000; +const PATTERN: &str = ".*([A-Z]{1}).*"; + +fn data(rng: &mut ThreadRng) -> StringArray { + (0..SIZE) + .map(|_| { + rng.sample_iter(&Alphanumeric) + .take(7) + .map(char::from) + .collect::() + }) + .collect::>() + .into() +} + +fn run(c: &mut Criterion, name: &str, values: &ArrayRef, args: &[ColumnarValue]) { + let func = RegexpMatchFunc::new(); + let arg_fields: Vec<_> = args + .iter() + .enumerate() + .map(|(idx, arg)| Field::new(format!("arg_{idx}"), arg.data_type(), true).into()) + .collect(); + let return_field = Arc::new(Field::new_list( + "f", + Field::new_list_field(values.data_type().clone(), true), + true, + )); + let config_options = Arc::new(ConfigOptions::default()); + + c.bench_function(name, |b| { + b.iter(|| { + black_box( + func.invoke_with_args(ScalarFunctionArgs { + args: args.to_vec(), + arg_fields: arg_fields.clone(), + number_rows: SIZE, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .expect("regexp_match should work on valid values"), + ) + }) + }); +} + +fn criterion_benchmark(c: &mut Criterion) { + let mut rng = rand::rng(); + let utf8 = Arc::new(data(&mut rng)) as ArrayRef; + let utf8view = cast(&utf8, &DataType::Utf8View).unwrap(); + + run( + c, + "regexp_match_1000 literal pattern", + &utf8, + &[ + ColumnarValue::Array(Arc::clone(&utf8)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(PATTERN.to_string()))), + ], + ); + + run( + c, + "regexp_match_1000 literal pattern and flags", + &utf8, + &[ + ColumnarValue::Array(Arc::clone(&utf8)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(PATTERN.to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some("i".to_string()))), + ], + ); + + run( + c, + "regexp_match_1000 literal pattern utf8view", + &utf8view, + &[ + ColumnarValue::Array(Arc::clone(&utf8view)), + ColumnarValue::Scalar(ScalarValue::Utf8View(Some(PATTERN.to_string()))), + ], + ); + + // Covers the path where the pattern varies per row and so cannot be + // compiled once for the whole array. + let patterns = Arc::new(StringArray::from( + (0..SIZE) + .map(|i| if i % 2 == 0 { PATTERN } else { "^(A).*" }) + .collect::>(), + )) as ArrayRef; + run( + c, + "regexp_match_1000 pattern array", + &utf8, + &[ + ColumnarValue::Array(Arc::clone(&utf8)), + ColumnarValue::Array(patterns), + ], + ); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/regx.rs b/datafusion/functions/benches/regx.rs index a46b548236d08..dd263e41f6fc5 100644 --- a/datafusion/functions/benches/regx.rs +++ b/datafusion/functions/benches/regx.rs @@ -32,11 +32,9 @@ use datafusion_functions::regex::regexpinstr::regexp_instr_func; use datafusion_functions::regex::regexplike::{RegexpLikeFunc, regexp_like}; use datafusion_functions::regex::regexpmatch::regexp_match; use datafusion_functions::regex::regexpreplace::regexp_replace; -use rand::Rng; use rand::distr::Alphanumeric; -use rand::prelude::IndexedRandom; -use rand::rngs::ThreadRng; -fn data(rng: &mut ThreadRng) -> StringArray { +use rand::prelude::*; +fn data(rng: &mut StdRng) -> StringArray { let mut data: Vec = vec![]; for _ in 0..1000 { data.push( @@ -50,7 +48,7 @@ fn data(rng: &mut ThreadRng) -> StringArray { StringArray::from(data) } -fn regex(rng: &mut ThreadRng) -> StringArray { +fn regex(rng: &mut StdRng) -> StringArray { let samples = [ ".*([A-Z]{1}).*".to_string(), "^(A).*".to_string(), @@ -66,7 +64,7 @@ fn regex(rng: &mut ThreadRng) -> StringArray { StringArray::from(data) } -fn start(rng: &mut ThreadRng) -> Int64Array { +fn start(rng: &mut StdRng) -> Int64Array { let mut data: Vec = vec![]; for _ in 0..1000 { data.push(rng.random_range(1..5)); @@ -75,7 +73,7 @@ fn start(rng: &mut ThreadRng) -> Int64Array { Int64Array::from(data) } -fn n(rng: &mut ThreadRng) -> Int64Array { +fn n(rng: &mut StdRng) -> Int64Array { let mut data: Vec = vec![]; for _ in 0..1000 { data.push(rng.random_range(1..5)); @@ -84,7 +82,7 @@ fn n(rng: &mut ThreadRng) -> Int64Array { Int64Array::from(data) } -fn flags(rng: &mut ThreadRng) -> StringArray { +fn flags(rng: &mut StdRng) -> StringArray { let samples = [Some("i".to_string()), Some("im".to_string()), None]; let mut sb = StringBuilder::new(); for _ in 0..1000 { @@ -99,7 +97,7 @@ fn flags(rng: &mut ThreadRng) -> StringArray { sb.finish() } -fn subexp(rng: &mut ThreadRng) -> Int64Array { +fn subexp(rng: &mut StdRng) -> Int64Array { let mut data: Vec = vec![]; for _ in 0..1000 { data.push(rng.random_range(1..5)); @@ -112,7 +110,7 @@ fn criterion_benchmark(c: &mut Criterion) { let regexp_like_func = RegexpLikeFunc::new(); let config_options = Arc::new(ConfigOptions::default()); c.bench_function("regexp_count_1000 string", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = Arc::new(data(&mut rng)) as ArrayRef; let regex = Arc::new(regex(&mut rng)) as ArrayRef; let start = Arc::new(start(&mut rng)) as ArrayRef; @@ -132,7 +130,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_count_1000 utf8view", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = cast(&data(&mut rng), &DataType::Utf8View).unwrap(); let regex = cast(®ex(&mut rng), &DataType::Utf8View).unwrap(); let start = Arc::new(start(&mut rng)) as ArrayRef; @@ -152,7 +150,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_instr_1000 string", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = Arc::new(data(&mut rng)) as ArrayRef; let regex = Arc::new(regex(&mut rng)) as ArrayRef; let start = Arc::new(start(&mut rng)) as ArrayRef; @@ -176,7 +174,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_instr_1000 utf8view", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = cast(&data(&mut rng), &DataType::Utf8View).unwrap(); let regex = cast(®ex(&mut rng), &DataType::Utf8View).unwrap(); let start = Arc::new(start(&mut rng)) as ArrayRef; @@ -198,7 +196,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_like_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = Arc::new(data(&mut rng)) as ArrayRef; let regex = Arc::new(regex(&mut rng)) as ArrayRef; let flags = Arc::new(flags(&mut rng)) as ArrayRef; @@ -212,7 +210,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_like_1000 utf8view", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = cast(&data(&mut rng), &DataType::Utf8View).unwrap(); let regex = cast(®ex(&mut rng), &DataType::Utf8View).unwrap(); let flags = cast(&flags(&mut rng), &DataType::Utf8View).unwrap(); @@ -252,7 +250,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_match_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = Arc::new(data(&mut rng)) as ArrayRef; let regex = Arc::new(regex(&mut rng)) as ArrayRef; let flags = Arc::new(flags(&mut rng)) as ArrayRef; @@ -270,7 +268,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_match_1000 utf8view", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = cast(&data(&mut rng), &DataType::Utf8View).unwrap(); let regex = cast(®ex(&mut rng), &DataType::Utf8View).unwrap(); let flags = cast(&flags(&mut rng), &DataType::Utf8View).unwrap(); @@ -288,7 +286,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_replace_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = Arc::new(data(&mut rng)) as ArrayRef; let regex = Arc::new(regex(&mut rng)) as ArrayRef; let flags = Arc::new(flags(&mut rng)) as ArrayRef; @@ -310,7 +308,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("regexp_replace_1000 utf8view", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data = cast(&data(&mut rng), &DataType::Utf8View).unwrap(); let regex = cast(®ex(&mut rng), &DataType::Utf8View).unwrap(); let flags = cast(&flags(&mut rng), &DataType::Utf8View).unwrap(); diff --git a/datafusion/functions/benches/replace.rs b/datafusion/functions/benches/replace.rs index b117968bad039..7ad198995a028 100644 --- a/datafusion/functions/benches/replace.rs +++ b/datafusion/functions/benches/replace.rs @@ -162,36 +162,6 @@ fn criterion_benchmark(c: &mut Criterion) { } } - // Empty-`from` path: insert `to` between every char of the input and at - // both ends. - if size == 1024 { - for &str_len in &[32_usize, 128] { - let args = create_args::(size, str_len, false, 0, 3, 0.0); - group.bench_function( - format!("replace_string_empty_from [size={size}, str_len={str_len}]"), - |b| { - b.iter(|| { - let args_cloned = args.clone(); - black_box(invoke_replace_with_args(args_cloned, size)) - }) - }, - ); - - let args = create_args::(size, str_len, true, 0, 3, 0.0); - group.bench_function( - format!( - "replace_string_view_empty_from [size={size}, str_len={str_len}]" - ), - |b| { - b.iter(|| { - let args_cloned = args.clone(); - black_box(invoke_replace_with_args(args_cloned, size)) - }) - }, - ); - } - } - group.finish(); } } diff --git a/datafusion/functions/benches/replace_scalar.rs b/datafusion/functions/benches/replace_scalar.rs new file mode 100644 index 0000000000000..e64c12e8ebf40 --- /dev/null +++ b/datafusion/functions/benches/replace_scalar.rs @@ -0,0 +1,78 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks the common `replace(column, 'lit', 'lit')` shape where the +//! `from`/`to` arguments are scalars, exercising the scalar-argument fast path. + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field}; +use arrow::util::bench_util::create_string_array_with_len; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_functions::string; +use std::hint::black_box; +use std::sync::Arc; + +fn run(c: &mut Criterion, size: usize, str_len: usize, from: &str, to: &str) { + let haystack: ArrayRef = + Arc::new(create_string_array_with_len::(size, 0.1, str_len)); + let args = vec![ + ColumnarValue::Array(Arc::clone(&haystack)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(from.to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(to.to_string()))), + ]; + let arg_fields = args + .iter() + .enumerate() + .map(|(i, a)| Field::new(format!("arg_{i}"), a.data_type(), true).into()) + .collect::>(); + let config_options = Arc::new(ConfigOptions::default()); + let func = string::replace(); + + c.bench_function( + &format!("replace_scalar from={from:?} [size={size}, str_len={str_len}]"), + |b| { + b.iter(|| { + black_box( + func.invoke_with_args(ScalarFunctionArgs { + args: args.clone(), + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Field::new("f", DataType::Utf8, true).into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }, + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + let size = 8192; + for str_len in [16_usize, 32, 64] { + // Multi-character patterns exercise the substring-finder path, where + // hoisting the finder out of the per-row loop matters most. + run(c, size, str_len, "ab", "XYZ"); + run(c, size, str_len, "the", "a"); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/round_dense.rs b/datafusion/functions/benches/round_dense.rs new file mode 100644 index 0000000000000..2c37849bde489 --- /dev/null +++ b/datafusion/functions/benches/round_dense.rs @@ -0,0 +1,94 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Microbenchmark for `round(float_array, scalar_decimal_places)` over a +//! Float column with no NULLs — the dense elementwise-rounding path. + +use arrow::array::ArrayRef; +use arrow::datatypes::{DataType, Field, Float32Type, Float64Type}; +use arrow::util::bench_util::create_primitive_array; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; +use datafusion_functions::math::round::RoundFunc; +use std::hint::black_box; +use std::sync::Arc; + +fn criterion_benchmark(c: &mut Criterion) { + let round_fn = RoundFunc::new(); + let config_options = Arc::new(ConfigOptions::default()); + + for size in [1024usize, 4096, 8192] { + // Float64, no nulls. + let f64_array: ArrayRef = + Arc::new(create_primitive_array::(size, 0.0)); + let f64_args = vec![ + ColumnarValue::Array(Arc::clone(&f64_array)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ]; + c.bench_with_input(BenchmarkId::new("round_dense_f64", size), &size, |b, _| { + b.iter(|| { + black_box( + round_fn + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + arg_fields: vec![ + Field::new("a", DataType::Float64, false).into(), + Field::new("b", DataType::Int32, false).into(), + ], + number_rows: size, + return_field: Field::new("f", DataType::Float64, false) + .into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + + // Float32, no nulls. + let f32_array: ArrayRef = + Arc::new(create_primitive_array::(size, 0.0)); + let f32_args = vec![ + ColumnarValue::Array(Arc::clone(&f32_array)), + ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + ]; + c.bench_with_input(BenchmarkId::new("round_dense_f32", size), &size, |b, _| { + b.iter(|| { + black_box( + round_fn + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + arg_fields: vec![ + Field::new("a", DataType::Float32, false).into(), + Field::new("b", DataType::Int32, false).into(), + ], + number_rows: size, + return_field: Field::new("f", DataType::Float32, false) + .into(), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/to_char.rs b/datafusion/functions/benches/to_char.rs index 350a55a37135c..8a9497bb33aa7 100644 --- a/datafusion/functions/benches/to_char.rs +++ b/datafusion/functions/benches/to_char.rs @@ -27,12 +27,10 @@ use datafusion_common::ScalarValue; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::to_char; -use rand::Rng; -use rand::prelude::IndexedRandom; -use rand::rngs::ThreadRng; +use rand::prelude::*; fn pick_date_in_range( - rng: &mut ThreadRng, + rng: &mut StdRng, start_date: NaiveDate, end_date: NaiveDate, ) -> NaiveDate { @@ -41,7 +39,7 @@ fn pick_date_in_range( start_date + TimeDelta::try_days(random_days).unwrap() } -fn generate_date32_array(rng: &mut ThreadRng) -> Date32Array { +fn generate_date32_array(rng: &mut StdRng) -> Date32Array { let mut data: Vec = vec![]; let unix_days_from_ce = NaiveDate::from_ymd_opt(1970, 1, 1) .unwrap() @@ -62,7 +60,7 @@ fn generate_date32_array(rng: &mut ThreadRng) -> Date32Array { Date32Array::from(data) } -fn generate_date64_array(rng: &mut ThreadRng) -> Date64Array { +fn generate_date64_array(rng: &mut StdRng) -> Date64Array { let start_date = "1970-01-01" .parse::() .expect("Date should parse"); @@ -96,21 +94,21 @@ const DATETIME_PATTERNS: [&str; 8] = [ "%c", ]; -fn pick_date_pattern(rng: &mut ThreadRng) -> String { +fn pick_date_pattern(rng: &mut StdRng) -> String { (*DATE_PATTERNS .choose(rng) .expect("Empty list of date patterns")) .to_string() } -fn pick_date_time_pattern(rng: &mut ThreadRng) -> String { +fn pick_date_time_pattern(rng: &mut StdRng) -> String { (*DATETIME_PATTERNS .choose(rng) .expect("Empty list of date time patterns")) .to_string() } -fn pick_date_and_date_time_mixed_pattern(rng: &mut ThreadRng) -> String { +fn pick_date_and_date_time_mixed_pattern(rng: &mut StdRng) -> String { match rng.random_bool(0.5) { true => pick_date_pattern(rng), false => pick_date_time_pattern(rng), @@ -118,8 +116,8 @@ fn pick_date_and_date_time_mixed_pattern(rng: &mut ThreadRng) -> String { } fn generate_pattern_array( - rng: &mut ThreadRng, - pick_fn: impl Fn(&mut ThreadRng) -> String, + rng: &mut StdRng, + pick_fn: impl Fn(&mut StdRng) -> String, ) -> StringArray { let mut data = Vec::with_capacity(1000); @@ -130,15 +128,15 @@ fn generate_pattern_array( StringArray::from(data) } -fn generate_date_pattern_array(rng: &mut ThreadRng) -> StringArray { +fn generate_date_pattern_array(rng: &mut StdRng) -> StringArray { generate_pattern_array(rng, pick_date_pattern) } -fn generate_datetime_pattern_array(rng: &mut ThreadRng) -> StringArray { +fn generate_datetime_pattern_array(rng: &mut StdRng) -> StringArray { generate_pattern_array(rng, pick_date_time_pattern) } -fn generate_mixed_pattern_array(rng: &mut ThreadRng) -> StringArray { +fn generate_mixed_pattern_array(rng: &mut StdRng) -> StringArray { generate_pattern_array(rng, pick_date_and_date_time_mixed_pattern) } @@ -146,7 +144,7 @@ fn criterion_benchmark(c: &mut Criterion) { let config_options = Arc::new(ConfigOptions::default()); c.bench_function("to_char_array_date_only_patterns_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data_arr = generate_date32_array(&mut rng); let batch_len = data_arr.len(); let data = ColumnarValue::Array(Arc::new(data_arr) as ArrayRef); @@ -173,7 +171,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("to_char_array_datetime_patterns_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data_arr = generate_date64_array(&mut rng); let batch_len = data_arr.len(); let data = ColumnarValue::Array(Arc::new(data_arr) as ArrayRef); @@ -200,7 +198,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("to_char_array_mixed_patterns_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data_arr = generate_date64_array(&mut rng); let batch_len = data_arr.len(); let data = ColumnarValue::Array(Arc::new(data_arr) as ArrayRef); @@ -227,7 +225,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("to_char_scalar_date_only_pattern_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data_arr = generate_date32_array(&mut rng); let batch_len = data_arr.len(); let data = ColumnarValue::Array(Arc::new(data_arr) as ArrayRef); @@ -253,7 +251,7 @@ fn criterion_benchmark(c: &mut Criterion) { }); c.bench_function("to_char_scalar_datetime_pattern_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data_arr = generate_date64_array(&mut rng); let batch_len = data_arr.len(); let data = ColumnarValue::Array(Arc::new(data_arr) as ArrayRef); @@ -285,7 +283,7 @@ fn criterion_benchmark(c: &mut Criterion) { // Covers full fallback (every row triggers the cast) c.bench_function("to_char_array_date32_datetime_patterns_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data_arr = generate_date32_array(&mut rng); let batch_len = data_arr.len(); let data = ColumnarValue::Array(Arc::new(data_arr) as ArrayRef); @@ -313,7 +311,7 @@ fn criterion_benchmark(c: &mut Criterion) { // Covers partial fallback (roughly half the rows trigger it) c.bench_function("to_char_array_date32_mixed_patterns_1000", |b| { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); let data_arr = generate_date32_array(&mut rng); let batch_len = data_arr.len(); let data = ColumnarValue::Array(Arc::new(data_arr) as ArrayRef); diff --git a/datafusion/functions/benches/to_local_time.rs b/datafusion/functions/benches/to_local_time.rs index 42d1e271980e8..04440bf0ac28a 100644 --- a/datafusion/functions/benches/to_local_time.rs +++ b/datafusion/functions/benches/to_local_time.rs @@ -24,17 +24,16 @@ use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::to_local_time; -use rand::Rng; -use rand::rngs::ThreadRng; +use rand::prelude::*; -fn timestamps(rng: &mut ThreadRng) -> TimestampNanosecondArray { +fn timestamps(rng: &mut StdRng) -> TimestampNanosecondArray { let nanos: Vec = (0..100_000) .map(|_| rng.random_range(0..1_000_000_000_000_000_000i64)) .collect(); TimestampNanosecondArray::from(nanos).with_timezone("America/New_York") } -fn timestamps_with_nulls(rng: &mut ThreadRng) -> TimestampNanosecondArray { +fn timestamps_with_nulls(rng: &mut StdRng) -> TimestampNanosecondArray { let values: Vec> = (0..100_000) .map(|_| { if rng.random_range(0..10u32) == 0 { @@ -73,7 +72,7 @@ fn bench_to_local_time(c: &mut Criterion, name: &str, array: ArrayRef) { } fn criterion_benchmark(c: &mut Criterion) { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); bench_to_local_time( c, "to_local_time_no_nulls_100k", diff --git a/datafusion/functions/benches/to_time.rs b/datafusion/functions/benches/to_time.rs index 6b3aa192415a3..f4499e2a7d0ba 100644 --- a/datafusion/functions/benches/to_time.rs +++ b/datafusion/functions/benches/to_time.rs @@ -24,10 +24,9 @@ use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_common::config::ConfigOptions; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::datetime::to_time; -use rand::Rng; -use rand::rngs::ThreadRng; +use rand::prelude::*; -fn random_time_string(rng: &mut ThreadRng) -> String { +fn random_time_string(rng: &mut StdRng) -> String { format!( "{:02}:{:02}:{:02}.{:06}", rng.random_range(0..24u32), @@ -37,12 +36,12 @@ fn random_time_string(rng: &mut ThreadRng) -> String { ) } -fn time_strings(rng: &mut ThreadRng) -> StringArray { +fn time_strings(rng: &mut StdRng) -> StringArray { let strings: Vec = (0..100_000).map(|_| random_time_string(rng)).collect(); StringArray::from(strings) } -fn time_strings_with_nulls(rng: &mut ThreadRng) -> StringArray { +fn time_strings_with_nulls(rng: &mut StdRng) -> StringArray { let values: Vec> = (0..100_000) .map(|_| { if rng.random_range(0..10u32) == 0 { @@ -81,7 +80,7 @@ fn bench_to_time(c: &mut Criterion, name: &str, array: ArrayRef) { } fn criterion_benchmark(c: &mut Criterion) { - let mut rng = rand::rng(); + let mut rng = StdRng::seed_from_u64(0); bench_to_time(c, "to_time_no_nulls_100k", Arc::new(time_strings(&mut rng))); bench_to_time( c, diff --git a/datafusion/functions/benches/translate.rs b/datafusion/functions/benches/translate.rs index d0568ba0f5355..adde7b4bd763d 100644 --- a/datafusion/functions/benches/translate.rs +++ b/datafusion/functions/benches/translate.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::OffsetSizeTrait; +use arrow::array::{GenericStringArray, OffsetSizeTrait}; use arrow::datatypes::{DataType, Field}; use arrow::util::bench_util::create_string_array_with_len; use criterion::{Criterion, SamplingMode, criterion_group, criterion_main}; @@ -23,10 +23,37 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::{DataFusionError, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; use datafusion_functions::unicode; +use rand::SeedableRng; +use rand::prelude::IndexedRandom; +use rand::rngs::StdRng; use std::hint::black_box; use std::sync::Arc; use std::time::Duration; +// Mix of 2-byte (Greek) and 3-byte (CJK/Hangul) UTF-8 to exercise +// variable-width char paths in translate. +const NON_ASCII_ALPHABET: &[char] = &[ + 'α', 'β', 'γ', 'δ', 'ε', 'ζ', 'η', 'θ', 'ι', 'κ', 'λ', 'μ', 'ν', 'ξ', 'ο', 'π', 'ρ', + 'σ', 'τ', 'υ', 'φ', 'χ', 'ψ', 'ω', '日', '本', '語', '中', '文', '한', '국', '어', +]; + +fn create_non_ascii_string_array( + size: usize, + char_count: usize, + seed: u64, +) -> GenericStringArray { + let mut rng = StdRng::seed_from_u64(seed); + (0..size) + .map(|_| { + Some( + (0..char_count) + .map(|_| *NON_ASCII_ALPHABET.choose(&mut rng).unwrap()) + .collect::(), + ) + }) + .collect() +} + fn create_args_array_from_to( size: usize, str_len: usize, @@ -42,6 +69,25 @@ fn create_args_array_from_to( ] } +fn create_args_array_from_to_non_ascii( + size: usize, + str_len: usize, +) -> Vec { + let string_array = Arc::new(create_non_ascii_string_array::( + size, + str_len, + 0xA110_AAAA, + )); + let from_array = Arc::new(create_non_ascii_string_array::(size, 3, 0xA110_BBBB)); + let to_array = Arc::new(create_non_ascii_string_array::(size, 2, 0xA110_CCCC)); + + vec![ + ColumnarValue::Array(string_array), + ColumnarValue::Array(from_array), + ColumnarValue::Array(to_array), + ] +} + fn create_args_scalar_from_to( size: usize, str_len: usize, @@ -91,6 +137,17 @@ fn criterion_benchmark(c: &mut Criterion) { }) }); + let args = create_args_array_from_to_non_ascii::(size, str_len); + group.bench_function( + format!("array_from_to_non_ascii [str_len={str_len}]"), + |b| { + b.iter(|| { + let args_cloned = args.clone(); + black_box(invoke_translate_with_args(args_cloned, size)) + }) + }, + ); + let args = create_args_scalar_from_to::(size, str_len); group.bench_function(format!("scalar_from_to [str_len={str_len}]"), |b| { b.iter(|| { diff --git a/datafusion/functions/benches/trunc_precision.rs b/datafusion/functions/benches/trunc_precision.rs new file mode 100644 index 0000000000000..5d75694d6ed2f --- /dev/null +++ b/datafusion/functions/benches/trunc_precision.rs @@ -0,0 +1,91 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks the `trunc(value, precision)` array path where `precision` is a +//! constant (scalar) argument. + +use arrow::datatypes::{DataType, Field, Float32Type, Float64Type}; +use arrow::util::bench_util::create_primitive_array; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs}; +use datafusion_functions::math::trunc; +use std::hint::black_box; +use std::sync::Arc; + +fn criterion_benchmark(c: &mut Criterion) { + let trunc = trunc(); + let config_options = Arc::new(ConfigOptions::default()); + + for size in [1024, 4096, 8192] { + let f64_array = Arc::new(create_primitive_array::(size, 0.2)); + let f64_args = vec![ + ColumnarValue::Array(f64_array), + ColumnarValue::Scalar(ScalarValue::Int64(Some(3))), + ]; + let arg_fields = vec![ + Field::new("a", DataType::Float64, true).into(), + Field::new("p", DataType::Int64, false).into(), + ]; + let return_field = Field::new("f", DataType::Float64, true).into(); + c.bench_function(&format!("trunc f64 precision array: {size}"), |b| { + b.iter(|| { + black_box( + trunc + .invoke_with_args(ScalarFunctionArgs { + args: f64_args.clone(), + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + + let f32_array = Arc::new(create_primitive_array::(size, 0.2)); + let f32_args = vec![ + ColumnarValue::Array(f32_array), + ColumnarValue::Scalar(ScalarValue::Int64(Some(3))), + ]; + let arg_fields = vec![ + Field::new("a", DataType::Float32, true).into(), + Field::new("p", DataType::Int64, false).into(), + ]; + let return_field = Field::new("f", DataType::Float32, true).into(); + c.bench_function(&format!("trunc f32 precision array: {size}"), |b| { + b.iter(|| { + black_box( + trunc + .invoke_with_args(ScalarFunctionArgs { + args: f32_args.clone(), + arg_fields: arg_fields.clone(), + number_rows: size, + return_field: Arc::clone(&return_field), + config_options: Arc::clone(&config_options), + }) + .unwrap(), + ) + }) + }); + } +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/benches/upper_unicode.rs b/datafusion/functions/benches/upper_unicode.rs new file mode 100644 index 0000000000000..2748c85e74854 --- /dev/null +++ b/datafusion/functions/benches/upper_unicode.rs @@ -0,0 +1,90 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks `upper` on non-ASCII input, which exercises the +//! character-streaming case-conversion path (not the ASCII fast path). + +use std::hint::black_box; +use std::sync::Arc; + +use arrow::array::{ArrayRef, LargeStringArray, StringArray}; +use arrow::datatypes::{DataType, Field}; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::config::ConfigOptions; +use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDF}; +use datafusion_functions::string; + +// A pool of non-ASCII words so `is_ascii()` is false and the Unicode path runs. +const WORDS: [&str; 8] = [ + "café", + "straße", + "αλφα", + "こんにちは", + "münchen", + "naïve", + " órdenes", + "tschüß", +]; + +fn build_values(size: usize) -> Vec> { + (0..size) + .map(|i| { + if i % 10 == 0 { + None + } else { + // Concatenate a few words for a longer, mixed value. + let a = WORDS[i % WORDS.len()]; + let b = WORDS[(i * 7 + 3) % WORDS.len()]; + Some(format!("{a} {b} {a}")) + } + }) + .collect() +} + +fn invoke(func: &ScalarUDF, array: ArrayRef, dt: DataType) { + let len = array.len(); + let config_options = Arc::new(ConfigOptions::default()); + black_box( + func.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(array)], + arg_fields: vec![Field::new("a", dt.clone(), true).into()], + number_rows: len, + return_field: Field::new("f", dt, true).into(), + config_options, + }) + .unwrap(), + ); +} + +fn criterion_benchmark(c: &mut Criterion) { + let upper = string::upper(); + let size = 4096; + let values = build_values(size); + + let utf8: ArrayRef = Arc::new(StringArray::from(values.clone())); + let large: ArrayRef = Arc::new(LargeStringArray::from(values)); + + c.bench_function("upper_unicode_utf8", |b| { + b.iter(|| invoke(&upper, Arc::clone(&utf8), DataType::Utf8)) + }); + c.bench_function("upper_unicode_large_utf8", |b| { + b.iter(|| invoke(&upper, Arc::clone(&large), DataType::LargeUtf8)) + }); +} + +criterion_group!(benches, criterion_benchmark); +criterion_main!(benches); diff --git a/datafusion/functions/src/binaries.rs b/datafusion/functions/src/binaries.rs new file mode 100644 index 0000000000000..861b7574cea19 --- /dev/null +++ b/datafusion/functions/src/binaries.rs @@ -0,0 +1,257 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::strings::{ColumnarValueRef, ConcatBuilder}; +use arrow::array::{ + Array, ArrayDataBuilder, ArrayRef, BinaryViewArray, GenericBinaryArray, + OffsetSizeTrait, make_view, +}; +use arrow_buffer::{ArrowNativeType, Buffer, MutableBuffer, NullBuffer, ScalarBuffer}; +use datafusion_common::{Result, exec_datafusion_err, exec_err, internal_err}; +use std::marker::PhantomData; +use std::sync::Arc; + +pub(crate) struct ConcatGenericBinaryBuilder { + offsets_buffer: MutableBuffer, + value_buffer: MutableBuffer, + _phantom: PhantomData, +} +pub(crate) type ConcatBinaryBuilder = ConcatGenericBinaryBuilder; +pub(crate) type ConcatLargeBinaryBuilder = ConcatGenericBinaryBuilder; + +impl ConcatGenericBinaryBuilder { + pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self { + let capacity = item_capacity + .checked_add(1) + .map(|i| i.saturating_mul(size_of::())) + .expect("capacity integer overflow"); + + let mut offsets_buffer = MutableBuffer::with_capacity(capacity); + // SAFETY: the first offset value is definitely not going to exceed the bounds. + unsafe { offsets_buffer.push_unchecked(O::usize_as(0)) }; + Self { + offsets_buffer, + value_buffer: MutableBuffer::with_capacity(data_capacity), + _phantom: PhantomData, + } + } +} + +impl ConcatBuilder + for ConcatGenericBinaryBuilder +{ + fn write( + &mut self, + column: &ColumnarValueRef, + i: usize, + ) -> Result<()> { + match column { + ColumnarValueRef::Scalar(s) => { + self.value_buffer.extend_from_slice(s); + } + ColumnarValueRef::NullableBinaryArray(array) => { + if !CHECK_VALID || array.is_valid(i) { + self.value_buffer.extend_from_slice(array.value(i)); + } + } + ColumnarValueRef::NullableLargeBinaryArray(array) => { + if !CHECK_VALID || array.is_valid(i) { + self.value_buffer.extend_from_slice(array.value(i)); + } + } + ColumnarValueRef::NullableBinaryViewArray(array) => { + if !CHECK_VALID || array.is_valid(i) { + self.value_buffer.extend_from_slice(array.value(i)); + } + } + ColumnarValueRef::NonNullableBinaryArray(array) => { + self.value_buffer.extend_from_slice(array.value(i)); + } + ColumnarValueRef::NonNullableLargeBinaryArray(array) => { + self.value_buffer.extend_from_slice(array.value(i)); + } + ColumnarValueRef::NonNullableBinaryViewArray(array) => { + self.value_buffer.extend_from_slice(array.value(i)); + } + _ => { + return exec_err!( + "concat: unexpected column type for binary builder: {column:?}" + ); + } + } + Ok(()) + } + + fn append_offset(&mut self) -> Result<()> { + let next_offset: O = O::from_usize(self.value_buffer.len()) + .ok_or_else(|| exec_datafusion_err!("byte array offset overflow"))?; + self.offsets_buffer.push(next_offset); + Ok(()) + } + + /// Finalize the builder into a concrete [`GenericBinaryArray`]. + /// + /// # Errors + /// + /// Returns an error when: + /// + /// - the provided `null_buffer` is not the same length as the `offsets_buffer`. + fn finish(self, null_buffer: Option) -> Result { + let row_count = self.offsets_buffer.len() / size_of::() - 1; + if let Some(ref null_buffer) = null_buffer + && null_buffer.len() != row_count + { + return internal_err!( + "Null buffer and offsets buffer must be the same length" + ); + } + let array_builder = ArrayDataBuilder::new(GenericBinaryArray::::DATA_TYPE) + .len(row_count) + .add_buffer(self.offsets_buffer.into()) + .add_buffer(self.value_buffer.into()) + .nulls(null_buffer); + // SAFETY: all data that was appended was valid and the values + // and offsets were created correctly + let array_data = unsafe { array_builder.build_unchecked() }; + let array = GenericBinaryArray::::from(array_data); + Ok(Arc::new(array)) + } +} + +/// Builder used by `concat`/`concat_ws` to assemble a [`BinaryViewArray`] one +/// row at a time from multiple input columns. +/// +/// Each row is written via repeated `write` calls (one per input +/// fragment) followed by a single `append_offset` to commit the row +/// as a single binary view. The output null buffer is supplied by the caller +/// at `finish` time, avoiding per-row NULL handling work. +/// +pub(crate) struct ConcatBinaryViewBuilder { + views: Vec, + data: Vec, + block: Vec, +} + +impl ConcatBinaryViewBuilder { + pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self { + Self { + views: Vec::with_capacity(item_capacity), + data: Vec::with_capacity(data_capacity), + block: vec![], + } + } +} + +impl ConcatBuilder for ConcatBinaryViewBuilder { + fn write( + &mut self, + column: &ColumnarValueRef, + i: usize, + ) -> Result<()> { + match column { + ColumnarValueRef::Scalar(s) => { + self.block.extend_from_slice(s); + } + ColumnarValueRef::NullableBinaryArray(array) => { + if !CHECK_VALID || array.is_valid(i) { + self.block.extend_from_slice(array.value(i)); + } + } + ColumnarValueRef::NullableLargeBinaryArray(array) => { + if !CHECK_VALID || array.is_valid(i) { + self.block.extend_from_slice(array.value(i)); + } + } + ColumnarValueRef::NullableBinaryViewArray(array) => { + if !CHECK_VALID || array.is_valid(i) { + self.block.extend_from_slice(array.value(i)); + } + } + ColumnarValueRef::NonNullableBinaryArray(array) => { + self.block.extend_from_slice(array.value(i)); + } + ColumnarValueRef::NonNullableLargeBinaryArray(array) => { + self.block.extend_from_slice(array.value(i)); + } + ColumnarValueRef::NonNullableBinaryViewArray(array) => { + self.block.extend_from_slice(array.value(i)); + } + _ => { + return exec_err!( + "concat: unexpected column type for binary view builder: {column:?}" + ); + } + } + Ok(()) + } + + /// Finalizes the current row by converting the accumulated data into a + /// StringView and appending it to the views buffer. + fn append_offset(&mut self) -> Result<()> { + let v = &self.block; + if v.len() > 12 { + let offset: u32 = self + .data + .len() + .try_into() + .map_err(|_| exec_datafusion_err!("byte array offset overflow"))?; + self.data.extend_from_slice(v); + self.views.push(make_view(v, 0, offset)); + } else { + self.views.push(make_view(v, 0, 0)); + } + + self.block.clear(); + Ok(()) + } + + /// Finalize the builder into a concrete [`BinaryViewArray`]. + /// + /// # Errors + /// + /// Returns an error when: + /// + /// - the provided `null_buffer` length does not match the row count. + fn finish(self, null_buffer: Option) -> Result { + if let Some(ref nulls) = null_buffer + && nulls.len() != self.views.len() + { + return internal_err!( + "Null buffer length ({}) must match row count ({})", + nulls.len(), + self.views.len() + ); + } + + let buffers: Vec = if self.data.is_empty() { + vec![] + } else { + vec![Buffer::from(self.data)] + }; + + // SAFETY: views were constructed with correct lengths, offsets, and + // prefixes. + let array = unsafe { + BinaryViewArray::new_unchecked( + ScalarBuffer::from(self.views), + buffers, + null_buffer, + ) + }; + Ok(Arc::new(array)) + } +} diff --git a/datafusion/functions/src/core/file_row_index.rs b/datafusion/functions/src/core/file_row_index.rs new file mode 100644 index 0000000000000..7b2667a8b8768 --- /dev/null +++ b/datafusion/functions/src/core/file_row_index.rs @@ -0,0 +1,96 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Implementation of the `file_row_index` scalar function. + +use arrow::datatypes::DataType; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, exec_err}; +use datafusion_doc::Documentation; +use datafusion_expr::{ + ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_macros::user_doc; + +/// Scalar UDF implementation for `file_row_index()`. +/// +/// File sources that can expose per-file row indexes rewrite this placeholder +/// function into a source-provided physical expression. Direct evaluation +/// returns an error because there is no file context outside a scan. +#[user_doc( + doc_section(label = "Other Functions"), + description = r#"Returns the zero-based row offset within the source file +that produced the current row. + +The value is scoped to one file, so rows from different files in the same scan +can have the same row index. This function is intended to be rewritten at +file-scan time. If the input file is not known (for example, if this function +is evaluated outside a file scan, or was not pushed down into one), direct +evaluation returns an error. +"#, + syntax_example = "file_row_index()", + sql_example = r#"```sql +SELECT file_row_index() FROM t; +```"# +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct FileRowIndexFunc { + signature: Signature, +} + +impl Default for FileRowIndexFunc { + fn default() -> Self { + Self::new() + } +} + +impl FileRowIndexFunc { + pub fn new() -> Self { + Self { + signature: Signature::nullary(Volatility::Volatile), + } + } +} + +impl ScalarUDFImpl for FileRowIndexFunc { + fn name(&self) -> &str { + "file_row_index" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, args: &[DataType]) -> Result { + let [] = take_function_args(self.name(), args)?; + Ok(DataType::Int64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [] = take_function_args(self.name(), args.args)?; + exec_err!("file_row_index() is source dependent and cannot be evaluated directly") + } + + fn placement(&self, _args: &[ExpressionPlacement]) -> ExpressionPlacement { + ExpressionPlacement::MoveTowardsLeafNodes + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} diff --git a/datafusion/functions/src/core/getfield.rs b/datafusion/functions/src/core/getfield.rs index b7092afcee492..6ec874fb672d1 100644 --- a/datafusion/functions/src/core/getfield.rs +++ b/datafusion/functions/src/core/getfield.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; +use std::sync::{Arc, OnceLock}; use arrow::array::{ Array, BooleanArray, Capacities, MutableArrayData, Scalar, cast::AsArray, make_array, @@ -37,6 +37,9 @@ use datafusion_expr::{ }; use datafusion_macros::user_doc; +use super::named_struct::NamedStructFunc; +use super::r#struct::StructFunc; + #[user_doc( doc_section(label = "Other Functions"), description = r#"Returns a field within a map or a struct with the given key. @@ -126,22 +129,22 @@ fn process_map_array( let mut mutable = MutableArrayData::with_capacities(vec![&original_data], true, capacity); + let offsets = map_array.value_offsets(); + // Scan the comparison result in place: slicing it per entry would allocate + // a new array for every row of the map. Map keys are non-null by + // definition, so the comparison result carries no nulls to check here. + let matches = keys.values(); + for entry in 0..map_array.len() { - let start = map_array.value_offsets()[entry] as usize; - let end = map_array.value_offsets()[entry + 1] as usize; + let start = offsets[entry] as usize; + let end = offsets[entry + 1] as usize; - let maybe_matched = keys - .slice(start, end - start) - .iter() - .enumerate() - .find(|(_, t)| t.unwrap()); + let matched = (start..end).find(|&i| matches.value(i)); - if maybe_matched.is_none() { - mutable.extend_nulls(1); - continue; + match matched { + Some(i) => mutable.try_extend(0, i, i + 1)?, + None => mutable.try_extend_nulls(1)?, } - let (match_offset, _) = maybe_matched.unwrap(); - mutable.extend(0, start + match_offset, start + match_offset + 1); } let data = mutable.freeze(); @@ -174,14 +177,14 @@ fn process_map_with_nested_key( let mut found_match = false; for i in start..end { if comparator(i, 0).is_eq() { - mutable.extend(0, i, i + 1); + mutable.try_extend(0, i, i + 1)?; found_match = true; break; } } if !found_match { - mutable.extend_nulls(1); + mutable.try_extend_nulls(1)?; } } @@ -249,6 +252,120 @@ fn extract_single_field(base: ColumnarValue, name: ScalarValue) -> Result Arc { + static GET_FIELD_UDF: OnceLock> = OnceLock::new(); + Arc::clone( + GET_FIELD_UDF + .get_or_init(|| Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::new()))), + ) +} + +/// Try to simplify a `get_field` call whose base is an inline struct +/// constructor by resolving the field access at plan time. +/// +/// Handles both struct constructors: +/// * `named_struct('a', x, 'b', y)` — fields are looked up by name. +/// * `struct(x, y)` — fields are positional and named `c0`, `c1`, ... +/// +/// For example: +/// * `get_field(named_struct('min', a, 'max', b), 'max')` => `b` +/// * `get_field(struct(a, b), 'c1')` => `b` +/// +/// `args` is the (already flattened) argument list of the `get_field` call: +/// `[base, field_name, rest_of_path...]`. When extra path elements remain +/// after resolving the first one (`get_field(named_struct('s', inner), 's', 'k')`), +/// the resolved value is re-wrapped in a `get_field` call for the remaining +/// path so the simplifier can recurse into it on the next pass. +/// +/// Returns `None` — leaving the expression untouched — whenever the rewrite +/// cannot be proven safe, e.g. a non-literal field name, a `named_struct` +/// with a non-literal field name (which might shadow the requested field at +/// runtime), or a field the constructor does not produce. +/// +/// Replacing the access with the selected field expression drops the +/// expressions for the other (unaccessed) fields, so they are no longer +/// evaluated — e.g. `get_field(named_struct('a', 1/0, 'b', x), 'b')` becomes +/// `x` and the `1/0` is never evaluated. This is intentional and matches the +/// optimizer's contract for immutable expressions: a simplification may drop +/// sub-expressions whose value is not observed. +fn simplify_get_field_over_struct_constructor(args: &[Expr]) -> Option { + let [base, field_name, rest @ ..] = args else { + return None; + }; + + // The accessed field name must be a non-empty string literal. + let Expr::Literal(field_name, _) = field_name else { + return None; + }; + let field_name = field_name + .try_as_str() + .flatten() + .filter(|s| !s.is_empty())?; + + let Expr::ScalarFunction(ScalarFunction { + func, + args: ctor_args, + }) = base + else { + return None; + }; + + let value = if func.inner().is::() { + // named_struct(name1, value1, name2, value2, ...) + if !ctor_args.len().is_multiple_of(2) { + return None; + } + let mut matched = None; + for pair in ctor_args.chunks_exact(2) { + // Every name must be a literal string: a non-literal name appearing + // *before* the first match could evaluate to `field_name` at runtime + // and become the real first match (Arrow's `column_by_name` returns + // the first match), so we cannot resolve the access. + // + // We conservatively bail on *any* non-literal name. Once a literal + // match has been found, a later non-literal name is in fact harmless + // — it can never precede the first match — so bailing there is a + // deliberate approximation we accept to keep this check simple, not a + // correctness requirement. + let Expr::Literal(name, _) = &pair[0] else { + return None; + }; + let name = name.try_as_str().flatten()?; + // `column_by_name` resolves to the first match, so do the same. + if matched.is_none() && name == field_name { + matched = Some(&pair[1]); + } + } + matched?.clone() + } else if func.inner().is::() { + // struct(value0, value1, ...) produces fields named c0, c1, ... + let index: usize = field_name.strip_prefix('c')?.parse().ok()?; + // Reject non-canonical spellings (e.g. "c01") that name no real field. + if format!("c{index}") != field_name { + return None; + } + ctor_args.get(index)?.clone() + } else { + return None; + }; + + if rest.is_empty() { + return Some(value); + } + + // Remaining path elements: re-wrap as get_field(value, rest...) and let + // the simplifier resolve the rest on a subsequent pass. + let mut new_args = Vec::with_capacity(rest.len() + 1); + new_args.push(value); + new_args.extend_from_slice(rest); + Some(Expr::ScalarFunction(ScalarFunction::new_udf( + get_field_udf(), + new_args, + ))) +} + impl GetFieldFunc { pub fn new() -> Self { Self { @@ -479,14 +596,12 @@ impl ScalarUDFImpl for GetFieldFunc { // Flatten all nested get_field calls in a single pass // Pattern: get_field(get_field(get_field(base, a), b), c) => get_field(base, a, b, c) - - // Collect path arguments from all nested levels - let mut path_args_stack = Vec::new(); + // + // `path_args_stack` collects each level's field-name arguments, + // outermost first; it is reversed below to restore access order. + let mut path_args_stack = vec![&args[1..]]; let mut current_expr = &args[0]; - // Push the outermost path arguments first - path_args_stack.push(&args[1..]); - // Walk down the chain of nested get_field calls let base_expr = loop { if let Expr::ScalarFunction(ScalarFunction { @@ -506,28 +621,30 @@ impl ScalarUDFImpl for GetFieldFunc { break current_expr; }; - // If no nested get_field calls were found, return original - if path_args_stack.len() == args.len() - 1 { - return Ok(ExprSimplifyResult::Original(args)); - } + // Whether any nested get_field calls were collapsed above. + let did_flatten = path_args_stack.len() > 1; - // If we found any nested get_field calls, flatten them - // Build merged args: [base, ...all_path_args_in_correct_order] + // Build merged args: [base, ...all path args in access order]. + // The stack holds path slices outermost-first, so iterate in reverse. let mut merged_args = vec![base_expr.clone()]; - - // Add path args in reverse order (innermost to outermost) - // Stack is: [outermost_paths, ..., innermost_paths] - // We want: [base, innermost_paths, ..., outermost_paths] for path_slice in path_args_stack.iter().rev() { merged_args.extend_from_slice(path_slice); } - Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction( - ScalarFunction::new_udf( - Arc::new(ScalarUDF::new_from_impl(GetFieldFunc::new())), - merged_args, - ), - ))) + // Resolve field accesses against an inline struct constructor: + // get_field(named_struct('min', a, 'max', b), 'max') => b + if let Some(simplified) = simplify_get_field_over_struct_constructor(&merged_args) + { + return Ok(ExprSimplifyResult::Simplified(simplified)); + } + + if did_flatten { + return Ok(ExprSimplifyResult::Simplified(Expr::ScalarFunction( + ScalarFunction::new_udf(get_field_udf(), merged_args), + ))); + } + + Ok(ExprSimplifyResult::Original(args)) } fn coerce_types(&self, arg_types: &[DataType]) -> Result> { @@ -828,4 +945,187 @@ mod tests { let args = vec![ExpressionPlacement::Literal, ExpressionPlacement::Literal]; assert_eq!(func.placement(&args), ExpressionPlacement::KeepInPlace); } + + // --- get_field over struct constructor simplification -------------------- + + use datafusion_common::Column; + use datafusion_expr::simplify::SimplifyContext; + + /// A non-empty string literal expression. + fn lit_str(s: &str) -> Expr { + Expr::Literal(ScalarValue::Utf8(Some(s.to_string())), None) + } + + /// A column reference expression. + fn col(name: &str) -> Expr { + Expr::Column(Column::from_name(name)) + } + + fn scalar_fn(udf: ScalarUDF, args: Vec) -> Expr { + Expr::ScalarFunction(ScalarFunction::new_udf(Arc::new(udf), args)) + } + + /// `named_struct(name1, value1, name2, value2, ...)`. + fn named_struct(pairs: Vec<(&str, Expr)>) -> Expr { + let args = pairs + .into_iter() + .flat_map(|(name, value)| [lit_str(name), value]) + .collect(); + scalar_fn(ScalarUDF::new_from_impl(NamedStructFunc::new()), args) + } + + /// `struct(value0, value1, ...)`. + fn struct_fn(values: Vec) -> Expr { + scalar_fn(ScalarUDF::new_from_impl(StructFunc::new()), values) + } + + /// `get_field(args...)`. + fn get_field(args: Vec) -> Expr { + scalar_fn(ScalarUDF::new_from_impl(GetFieldFunc::new()), args) + } + + /// Run `GetFieldFunc::simplify` once and return the rewritten expression, + /// panicking if the input was left unchanged. + fn simplified(args: Vec) -> Expr { + match GetFieldFunc::new() + .simplify(args, &SimplifyContext::default()) + .unwrap() + { + ExprSimplifyResult::Simplified(expr) => expr, + ExprSimplifyResult::Original(args) => { + panic!("expected the expression to be simplified, got {args:?}") + } + } + } + + /// Assert that `GetFieldFunc::simplify` leaves the arguments unchanged. + fn assert_not_simplified(args: Vec) { + match GetFieldFunc::new() + .simplify(args.clone(), &SimplifyContext::default()) + .unwrap() + { + ExprSimplifyResult::Original(unchanged) => assert_eq!(unchanged, args), + ExprSimplifyResult::Simplified(expr) => { + panic!("expected no simplification, got {expr:?}") + } + } + } + + #[test] + fn simplify_get_field_named_struct_returns_matching_value() { + // get_field(named_struct('min', a, 'max', b), 'max') => b + let args = vec![ + named_struct(vec![("min", col("a")), ("max", col("b"))]), + lit_str("max"), + ]; + assert_eq!(simplified(args), col("b")); + } + + #[test] + fn simplify_get_field_named_struct_first_field() { + // get_field(named_struct('min', a, 'max', b), 'min') => a + let args = vec![ + named_struct(vec![("min", col("a")), ("max", col("b"))]), + lit_str("min"), + ]; + assert_eq!(simplified(args), col("a")); + } + + #[test] + fn simplify_get_field_named_struct_duplicate_names_picks_first() { + // Arrow's `column_by_name` resolves to the first match; mirror that. + let args = vec![ + named_struct(vec![("k", col("a")), ("k", col("b"))]), + lit_str("k"), + ]; + assert_eq!(simplified(args), col("a")); + } + + #[test] + fn simplify_get_field_struct_positional() { + // get_field(struct(a, b), 'c1') => b + let args = vec![struct_fn(vec![col("a"), col("b")]), lit_str("c1")]; + assert_eq!(simplified(args), col("b")); + } + + #[test] + fn simplify_get_field_nested_named_struct() { + // get_field(named_struct('s', named_struct('k', x)), 's', 'k') + // => get_field(named_struct('k', x), 'k') (first pass) + // => x (second pass) + let args = vec![ + named_struct(vec![("s", named_struct(vec![("k", col("x"))]))]), + lit_str("s"), + lit_str("k"), + ]; + let first_pass = simplified(args); + let Expr::ScalarFunction(ScalarFunction { args, .. }) = first_pass else { + panic!("expected a get_field call after the first pass") + }; + assert_eq!(simplified(args), col("x")); + } + + #[test] + fn simplify_get_field_flattens_then_resolves_named_struct() { + // get_field(get_field(named_struct('s', named_struct('k', x)), 's'), 'k') + // flattens to get_field(named_struct(...), 's', 'k') and resolves 's'. + let args = vec![ + get_field(vec![ + named_struct(vec![("s", named_struct(vec![("k", col("x"))]))]), + lit_str("s"), + ]), + lit_str("k"), + ]; + let expected = get_field(vec![named_struct(vec![("k", col("x"))]), lit_str("k")]); + assert_eq!(simplified(args), expected); + } + + #[test] + fn simplify_get_field_dynamic_field_name_left_alone() { + // A non-literal field name cannot be resolved at plan time. + let args = vec![named_struct(vec![("a", col("x"))]), col("field_name")]; + assert_not_simplified(args); + } + + #[test] + fn simplify_get_field_null_field_name_left_alone() { + // A NULL string literal field name resolves to no field, so the + // `try_as_str().flatten()` guard must leave the expression untouched. + let null_field_name = Expr::Literal(ScalarValue::Utf8(None), None); + let args = vec![named_struct(vec![("a", col("x"))]), null_field_name]; + assert_not_simplified(args); + } + + #[test] + fn simplify_get_field_dynamic_struct_name_left_alone() { + // A non-literal name inside named_struct could shadow the requested + // field at runtime, so the rewrite must bail out entirely. + let named_struct_with_dynamic_name = scalar_fn( + ScalarUDF::new_from_impl(NamedStructFunc::new()), + vec![col("dynamic_name"), col("x")], + ); + let args = vec![named_struct_with_dynamic_name, lit_str("a")]; + assert_not_simplified(args); + } + + #[test] + fn simplify_get_field_missing_field_left_alone() { + // The named_struct does not produce field 'missing'. + let args = vec![named_struct(vec![("a", col("x"))]), lit_str("missing")]; + assert_not_simplified(args); + } + + #[test] + fn simplify_get_field_non_canonical_struct_field_left_alone() { + // 'c01' is not a real field name produced by `struct(...)`. + let args = vec![struct_fn(vec![col("a"), col("b")]), lit_str("c01")]; + assert_not_simplified(args); + } + + #[test] + fn simplify_get_field_column_base_left_alone() { + // A plain column base is not a struct constructor. + let args = vec![col("s"), lit_str("a")]; + assert_not_simplified(args); + } } diff --git a/datafusion/functions/src/core/input_file_name.rs b/datafusion/functions/src/core/input_file_name.rs new file mode 100644 index 0000000000000..a47e9daaf8d3c --- /dev/null +++ b/datafusion/functions/src/core/input_file_name.rs @@ -0,0 +1,95 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`InputFileNameFunc`]: Implementation of the `input_file_name` function. + +use arrow::datatypes::DataType; +use datafusion_common::{exec_err, utils::take_function_args}; +use datafusion_doc::Documentation; +use datafusion_expr::{ + ColumnarValue, ExpressionPlacement, ScalarFunctionArgs, ScalarUDFImpl, Signature, + Volatility, +}; +use datafusion_macros::user_doc; + +#[user_doc( + doc_section(label = "Other Functions"), + description = r#"Returns the path of the input file that produced the current row. + +Note: file paths/URIs may be sensitive metadata depending on your environment. + +This function is intended to be rewritten at file-scan time (when the file is +known). If the input file is not known (for example, if this function is +evaluated outside a file scan, or was not pushed down into one), direct evaluation returns an error. +"#, + syntax_example = "input_file_name()", + sql_example = r#"```sql +SELECT input_file_name() FROM t; +```"# +)] +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct InputFileNameFunc { + signature: Signature, +} + +impl Default for InputFileNameFunc { + fn default() -> Self { + Self::new() + } +} + +impl InputFileNameFunc { + pub fn new() -> Self { + Self { + signature: Signature::nullary(Volatility::Volatile), + } + } +} + +impl ScalarUDFImpl for InputFileNameFunc { + fn name(&self) -> &str { + "input_file_name" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> datafusion_common::Result { + let [] = take_function_args(self.name(), arg_types)?; + Ok(DataType::Utf8) + } + + fn invoke_with_args( + &self, + args: ScalarFunctionArgs, + ) -> datafusion_common::Result { + let [] = take_function_args(self.name(), args.args)?; + + exec_err!( + "input_file_name() is source dependent and cannot be evaluated directly" + ) + } + + fn placement(&self, _args: &[ExpressionPlacement]) -> ExpressionPlacement { + ExpressionPlacement::MoveTowardsLeafNodes + } + + fn documentation(&self) -> Option<&Documentation> { + self.doc() + } +} diff --git a/datafusion/functions/src/core/mod.rs b/datafusion/functions/src/core/mod.rs index 5657f9d88810c..3f7a562a27ec7 100644 --- a/datafusion/functions/src/core/mod.rs +++ b/datafusion/functions/src/core/mod.rs @@ -28,9 +28,11 @@ pub mod arrowtypeof; pub mod cast_to_type; pub mod coalesce; pub mod expr_ext; +pub mod file_row_index; pub mod getfield; pub mod greatest; mod greatest_least_utils; +pub mod input_file_name; pub mod least; pub mod named_struct; pub mod nullif; @@ -67,6 +69,8 @@ make_udf_function!(version::VersionFunc, version); make_udf_function!(arrow_metadata::ArrowMetadataFunc, arrow_metadata); make_udf_function!(with_metadata::WithMetadataFunc, with_metadata); make_udf_function!(arrow_field::ArrowFieldFunc, arrow_field); +make_udf_function!(file_row_index::FileRowIndexFunc, file_row_index); +make_udf_function!(input_file_name::InputFileNameFunc, input_file_name); pub mod expr_fn { use datafusion_expr::{Expr, Literal}; @@ -115,7 +119,12 @@ pub mod expr_fn { arrow_metadata, "Returns the metadata of the input expression", args, - ),( + ), + ( + input_file_name, + "Returns the path of the input file that produced the current row", + ), + ( with_metadata, "Attaches Arrow field metadata (key/value pairs) to the input expression", args, @@ -143,6 +152,9 @@ pub mod expr_fn { union_tag, "Returns the name of the currently selected field in the union", arg1 + ),( + file_row_index, + "Returns the offset of the row within its source file", )); #[doc = "Returns the value of the field with the given name from the struct"] @@ -195,6 +207,8 @@ pub fn functions() -> Vec> { union_extract(), union_tag(), version(), + input_file_name(), r#struct(), + file_row_index(), ] } diff --git a/datafusion/functions/src/core/overlay.rs b/datafusion/functions/src/core/overlay.rs index 2d99af9e783bb..c1f3353a8f413 100644 --- a/datafusion/functions/src/core/overlay.rs +++ b/datafusion/functions/src/core/overlay.rs @@ -15,11 +15,16 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; - -use arrow::array::{ArrayRef, GenericStringArray, OffsetSizeTrait}; +use arrow::array::{ + Array, ArrayRef, GenericStringArray, Int64Array, OffsetSizeTrait, StringArrayType, + StringViewArray, +}; +use arrow::buffer::NullBuffer; use arrow::datatypes::DataType; +use crate::strings::{ + BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringWriter, +}; use crate::utils::{make_scalar_function, utf8_to_str_type}; use datafusion_common::cast::{ as_generic_string_array, as_int64_array, as_string_view_array, @@ -112,106 +117,79 @@ impl ScalarUDFImpl for OverlayFunc { } } -/// Converts a 0-based character index into a byte index suitable for UTF-8 -/// slicing. -fn byte_index_for_char(string: &str, char_idx: usize, is_ascii: bool) -> usize { - if is_ascii { - char_idx.min(string.len()) - } else { - string - .char_indices() - .nth(char_idx) - .map_or(string.len(), |(byte_idx, _)| byte_idx) +/// Computes the byte ranges of `string` to keep around the replaced span: the +/// prefix is `string[..prefix_end]` and the suffix is `string[suffix_start..]`. +/// +/// `start_pos` is a 1-based character position; the caller must ensure it is +/// `>= 1`. `replace_len` is the number of characters of `string` to replace, +/// and may be negative (in which case `suffix_start <= prefix_end` and the +/// result re-emits part of the original string). +/// +/// Matches PostgreSQL semantics for codepoint indices past the end of +/// `string`: `prefix_end` and `suffix_start` clamp to `string.len()`. +fn overlay_bounds(string: &str, start_pos: i64, replace_len: i64) -> (usize, usize) { + let start_char_idx = start_pos - 1; + let end_char_idx = start_char_idx.saturating_add(replace_len); + + if string.is_ascii() { + // ASCII fast path: byte index == codepoint index. + let len = string.len() as i64; + let prefix_end = start_char_idx.clamp(0, len) as usize; + let suffix_start = end_char_idx.clamp(0, len) as usize; + return (prefix_end, suffix_start); + } + + let prefix_target = usize::try_from(start_char_idx).unwrap_or(usize::MAX); + let suffix_target = usize::try_from(end_char_idx.max(0)).unwrap_or(usize::MAX); + let target_max = prefix_target.max(suffix_target); + + // Single forward pass over codepoint boundaries records both targets. + // Either target falls through to `string.len()` if past the codepoint + // count. + let mut prefix_byte = string.len(); + let mut suffix_byte = string.len(); + for (count, (byte_idx, _)) in string.char_indices().enumerate() { + if count == prefix_target { + prefix_byte = byte_idx; + } + if count == suffix_target { + suffix_byte = byte_idx; + } + if count == target_max { + break; + } } + (prefix_byte, suffix_byte) } -/// Builds the OVERLAY result for a single (non-null) row. -/// -/// `start_pos` is a 1-based character position; `replace_len` is the number -/// of characters of `string` to replace with `characters`. -fn overlay_one( +/// Appends the overlay result for one non-null row into `builder`. +#[inline] +fn apply_overlay( string: &str, characters: &str, start_pos: i64, replace_len: i64, -) -> Result { + builder: &mut B, +) -> Result<()> { if start_pos < 1 { - return exec_err!("negative substring length not allowed"); + return exec_err!("overlay start position must be at least 1: {start_pos}"); } + let (prefix_end, suffix_start) = overlay_bounds(string, start_pos, replace_len); + builder.append_with(|w| { + w.write_str(&string[..prefix_end]); + w.write_str(characters); + w.write_str(&string[suffix_start..]); + }); + Ok(()) +} - let is_ascii = string.is_ascii(); - let string_char_len = if is_ascii { - string.len() as i64 +#[inline] +fn char_count(characters: &str) -> i64 { + if characters.is_ascii() { + characters.len() as i64 } else { - string.chars().count() as i64 - }; - - // Convert SQL's 1-based character position into 0-based character indexes. - // `start_char_idx` is the first replaced character; `end_char_idx` is the - // first character after the replaced span. - // - // No upper-bound check on `start_char_idx`: when it exceeds `string_char_len` - // we want the whole string as the prefix (PostgreSQL-compatible "insert past - // end" semantics). - let start_char_idx = start_pos - 1; - let end_char_idx = start_char_idx.saturating_add(replace_len); - - let prefix_char_idx = usize::try_from(start_char_idx).unwrap_or(usize::MAX); - let prefix_end_byte = byte_index_for_char(string, prefix_char_idx, is_ascii); - - let mut res = String::with_capacity(string.len() + characters.len()); - res.push_str(&string[..prefix_end_byte]); - res.push_str(characters); - - if end_char_idx < string_char_len { - let suffix_char_idx = usize::try_from(end_char_idx.max(0)).unwrap_or(usize::MAX); - let suffix_start_byte = byte_index_for_char(string, suffix_char_idx, is_ascii); - res.push_str(&string[suffix_start_byte..]); + characters.chars().count() as i64 } - Ok(res) -} - -macro_rules! process_overlay { - // Three argument case - ($string_array:expr, $characters_array:expr, $pos_array:expr) => {{ - $string_array - .iter() - .zip($characters_array.iter()) - .zip($pos_array.iter()) - .map(|((string, characters), start_pos)| { - match (string, characters, start_pos) { - (Some(string), Some(characters), Some(start_pos)) => { - let replace_len = characters.chars().count() as i64; - overlay_one(string, characters, start_pos, replace_len).map(Some) - } - _ => Ok(None), - } - }) - .collect::>>() - }}; - - // Four argument case - ($string_array:expr, $characters_array:expr, $pos_array:expr, $len_array:expr) => {{ - $string_array - .iter() - .zip($characters_array.iter()) - .zip($pos_array.iter()) - .zip($len_array.iter()) - .map(|(((string, characters), start_pos), replace_len)| { - match (string, characters, start_pos, replace_len) { - ( - Some(string), - Some(characters), - Some(start_pos), - Some(replace_len), - ) => { - overlay_one(string, characters, start_pos, replace_len).map(Some) - } - _ => Ok(None), - } - }) - .collect::>>() - }}; } /// `OVERLAY(string PLACING substring FROM start [FOR count])` @@ -232,44 +210,122 @@ fn overlay(args: &[ArrayRef]) -> Result { args.len() ); } + let pos_array = as_int64_array(&args[2])?; + let len_array = if args.len() == 4 { + Some(as_int64_array(&args[3])?) + } else { + None + }; + if args[0].data_type() == &DataType::Utf8View { - string_view_overlay::(args) + let string_array = as_string_view_array(&args[0])?; + let characters_array = as_string_view_array(&args[1])?; + let data_capacity = visible_view_bytes(string_array) + .saturating_add(visible_view_bytes(characters_array)); + let builder = GenericStringArrayBuilder::::with_capacity( + string_array.len(), + data_capacity, + ); + overlay_inner( + string_array, + characters_array, + pos_array, + len_array, + builder, + ) } else { - string_overlay::(args) + let string_array = as_generic_string_array::(&args[0])?; + let characters_array = as_generic_string_array::(&args[1])?; + let data_capacity = visible_offset_bytes(string_array) + .saturating_add(visible_offset_bytes(characters_array)); + let builder = GenericStringArrayBuilder::::with_capacity( + string_array.len(), + data_capacity, + ); + overlay_inner( + string_array, + characters_array, + pos_array, + len_array, + builder, + ) } } -fn string_overlay(args: &[ArrayRef]) -> Result { - let string_array = as_generic_string_array::(&args[0])?; - let characters_array = as_generic_string_array::(&args[1])?; - let pos_array = as_int64_array(&args[2])?; +/// Drives the per-row OVERLAY computation. A null in any input array +/// produces a null output. +fn overlay_inner<'a, V, B>( + string_array: V, + characters_array: V, + pos_array: &Int64Array, + len_array: Option<&Int64Array>, + mut builder: B, +) -> Result +where + V: StringArrayType<'a, Item = &'a str> + Copy, + B: BulkNullStringArrayBuilder, +{ + let len = string_array.len(); + let nulls = NullBuffer::union_many([ + string_array.nulls(), + characters_array.nulls(), + pos_array.nulls(), + len_array.and_then(|a| a.nulls()), + ]); - let result = if args.len() == 4 { - let len_array = as_int64_array(&args[3])?; - process_overlay!(string_array, characters_array, pos_array, len_array)? + if let Some(nulls_ref) = nulls.as_ref() { + for i in 0..len { + if nulls_ref.is_null(i) { + builder.append_placeholder(); + continue; + } + // SAFETY: `i < len`, and null bitmap check implies not-null + let string = unsafe { string_array.value_unchecked(i) }; + let characters = unsafe { characters_array.value_unchecked(i) }; + let start_pos = unsafe { pos_array.value_unchecked(i) }; + let replace_len = match len_array { + Some(arr) => unsafe { arr.value_unchecked(i) }, + None => char_count(characters), + }; + apply_overlay(string, characters, start_pos, replace_len, &mut builder)?; + } } else { - process_overlay!(string_array, characters_array, pos_array)? - }; - Ok(Arc::new(result) as ArrayRef) + for i in 0..len { + // SAFETY: `i < len`, and no null bitmap means no nulls + let string = unsafe { string_array.value_unchecked(i) }; + let characters = unsafe { characters_array.value_unchecked(i) }; + let start_pos = unsafe { pos_array.value_unchecked(i) }; + let replace_len = match len_array { + Some(arr) => unsafe { arr.value_unchecked(i) }, + None => char_count(characters), + }; + apply_overlay(string, characters, start_pos, replace_len, &mut builder)?; + } + } + builder.finish(nulls) } -fn string_view_overlay(args: &[ArrayRef]) -> Result { - let string_array = as_string_view_array(&args[0])?; - let characters_array = as_string_view_array(&args[1])?; - let pos_array = as_int64_array(&args[2])?; +/// Bytes referenced by the visible window of `array`, computed from the +/// per-view lengths. +fn visible_view_bytes(array: &StringViewArray) -> usize { + array.lengths().map(|l| l as usize).sum() +} - let result = if args.len() == 4 { - let len_array = as_int64_array(&args[3])?; - process_overlay!(string_array, characters_array, pos_array, len_array)? - } else { - process_overlay!(string_array, characters_array, pos_array)? - }; - Ok(Arc::new(result) as ArrayRef) +/// Bytes referenced by the visible window of `array`, derived from the offset +/// buffer. +fn visible_offset_bytes(array: &GenericStringArray) -> usize { + let offsets = array.value_offsets(); + // `value_offsets()` always has `array.len() + 1` entries (≥1). + let first = offsets.first().copied().unwrap_or_default(); + let last = offsets.last().copied().unwrap_or_default(); + last.as_usize() - first.as_usize() } #[cfg(test)] mod tests { - use arrow::array::{Int64Array, StringArray}; + use std::sync::Arc; + + use arrow::array::StringArray; use super::*; diff --git a/datafusion/functions/src/crypto/md5.rs b/datafusion/functions/src/crypto/md5.rs index 178aebf0fbd41..b1206d2e423cc 100644 --- a/datafusion/functions/src/crypto/md5.rs +++ b/datafusion/functions/src/crypto/md5.rs @@ -21,6 +21,7 @@ use datafusion_common::{ cast::as_binary_array, internal_err, types::{logical_binary, logical_string}, + utils::hex::{HexCase, encode_bytes}, utils::take_function_args, }; use datafusion_expr::{ @@ -98,22 +99,6 @@ impl ScalarUDFImpl for Md5Func { } } -/// Hex encoding lookup table for fast byte-to-hex conversion -const HEX_CHARS_LOWER: &[u8; 16] = b"0123456789abcdef"; - -/// Fast hex encoding using a lookup table instead of format strings. -/// This is significantly faster than using `write!("{:02x}")` for each byte. -#[inline] -fn hex_encode(data: impl AsRef<[u8]>) -> String { - let bytes = data.as_ref(); - let mut s = String::with_capacity(bytes.len() * 2); - for &b in bytes { - s.push(HEX_CHARS_LOWER[(b >> 4) as usize] as char); - s.push(HEX_CHARS_LOWER[(b & 0x0f) as usize] as char); - } - s -} - fn md5(args: &[ColumnarValue]) -> Result { let [data] = take_function_args("md5", args)?; let value = digest_process(data, DigestAlgorithm::Md5)?; @@ -122,13 +107,15 @@ fn md5(args: &[ColumnarValue]) -> Result { Ok(match value { ColumnarValue::Array(array) => { let binary_array = as_binary_array(&array)?; - let string_array: StringViewArray = - binary_array.iter().map(|opt| opt.map(hex_encode)).collect(); + let string_array: StringViewArray = binary_array + .iter() + .map(|opt| opt.map(|b| encode_bytes(b, HexCase::Lower))) + .collect(); ColumnarValue::Array(Arc::new(string_array)) } - ColumnarValue::Scalar(ScalarValue::Binary(opt)) => { - ColumnarValue::Scalar(ScalarValue::Utf8View(opt.map(hex_encode))) - } + ColumnarValue::Scalar(ScalarValue::Binary(opt)) => ColumnarValue::Scalar( + ScalarValue::Utf8View(opt.map(|b| encode_bytes(&b, HexCase::Lower))), + ), _ => return internal_err!("Impossibly got invalid results from digest"), }) } diff --git a/datafusion/functions/src/datetime/common.rs b/datafusion/functions/src/datetime/common.rs index 2db64beafa9b7..9a7f94bd5973f 100644 --- a/datafusion/functions/src/datetime/common.rs +++ b/datafusion/functions/src/datetime/common.rs @@ -32,7 +32,7 @@ use chrono::{DateTime, TimeZone, Utc}; use datafusion_common::cast::as_generic_string_array; use datafusion_common::{ DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err, - internal_datafusion_err, unwrap_or_internal_err, + internal_datafusion_err, }; use datafusion_expr::ColumnarValue; @@ -353,9 +353,9 @@ where // if the first argument is a scalar utf8 all arguments are expected to be scalar utf8 ColumnarValue::Scalar(scalar) => match scalar.try_as_str() { Some(a) => { - let a = a.as_ref(); - // ASK: Why do we trust `a` to be non-null at this point? - let a = unwrap_or_internal_err!(a); + let Some(a) = a.as_ref() else { + return Ok(ColumnarValue::Scalar(scalar_value(dt, None)?)); + }; let mut ret = None; @@ -384,7 +384,10 @@ where } } - unwrap_or_internal_err!(ret) + match ret { + Some(ret) => ret, + None => Ok(ColumnarValue::Scalar(scalar_value(dt, None)?)), + } } other => { exec_err!("Unsupported data type {other:?} for function {name}") @@ -483,12 +486,21 @@ where if let Some(x) = x { for arg in args { let v = match arg { - ColumnarValue::Array(a) => match a.data_type() { - DataType::Utf8View => Ok(a.as_string_view().value(pos)), - DataType::LargeUtf8 => Ok(a.as_string::().value(pos)), - DataType::Utf8 => Ok(a.as_string::().value(pos)), - other => exec_err!("Unexpected type encountered '{other}'"), - }, + ColumnarValue::Array(a) => { + if a.is_null(pos) { + continue; + } + match a.data_type() { + DataType::Utf8View => Ok(a.as_string_view().value(pos)), + DataType::LargeUtf8 => { + Ok(a.as_string::().value(pos)) + } + DataType::Utf8 => Ok(a.as_string::().value(pos)), + other => { + exec_err!("Unexpected type encountered '{other}'") + } + } + } ColumnarValue::Scalar(s) => match s.try_as_str() { Some(Some(v)) => Ok(v), Some(None) => continue, // null string diff --git a/datafusion/functions/src/datetime/current_date.rs b/datafusion/functions/src/datetime/current_date.rs index d07a3b1caf13b..e93a64e8cc090 100644 --- a/datafusion/functions/src/datetime/current_date.rs +++ b/datafusion/functions/src/datetime/current_date.rs @@ -35,9 +35,7 @@ Returns the current date in the session time zone. The `current_date()` return value is determined at query time and will return the same date, no matter when in the query plan the function executes. "#, - syntax_example = r#"current_date() - (optional) SET datafusion.execution.time_zone = '+00:00'; - SELECT current_date();"#, + syntax_example = "current_date()", sql_example = r#"```sql > SELECT current_date(); +----------------+ diff --git a/datafusion/functions/src/datetime/current_time.rs b/datafusion/functions/src/datetime/current_time.rs index 92f4ae5e66f02..b93fb07d2b6f2 100644 --- a/datafusion/functions/src/datetime/current_time.rs +++ b/datafusion/functions/src/datetime/current_time.rs @@ -38,9 +38,7 @@ The `current_time()` return value is determined at query time and will return th The session time zone can be set using the statement 'SET datafusion.execution.time_zone = desired time zone'. The time zone can be a value like +00:00, 'Europe/London' etc. "#, - syntax_example = r#"current_time() - (optional) SET datafusion.execution.time_zone = '+00:00'; - SELECT current_time();"#, + syntax_example = "current_time()", sql_example = r#"```sql > SELECT current_time(); +--------------------+ diff --git a/datafusion/functions/src/datetime/date_bin.rs b/datafusion/functions/src/datetime/date_bin.rs index c26623c46b0c1..1338df3aa916f 100644 --- a/datafusion/functions/src/datetime/date_bin.rs +++ b/datafusion/functions/src/datetime/date_bin.rs @@ -31,9 +31,12 @@ use arrow::datatypes::{ DataType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType, TimeUnit, }; +use arrow::error::ArrowError; use arrow::temporal_conversions::NANOSECONDS_IN_DAY; use datafusion_common::cast::as_primitive_array; -use datafusion_common::{Result, ScalarValue, exec_err, not_impl_err, plan_err}; +use datafusion_common::{ + Result, ScalarValue, exec_datafusion_err, exec_err, not_impl_err, plan_err, +}; use datafusion_expr::TypeSignature::Exact; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ @@ -51,7 +54,7 @@ Calculates time intervals and returns the start of the interval nearest to the s For example, if you "bin" or "window" data into 15 minute intervals, an input timestamp of `2023-01-01T18:18:18Z` will be updated to the start time of the 15 minute bin it is in: `2023-01-01T18:15:00Z`. "#, - syntax_example = "date_bin(interval, expression, origin-timestamp)", + syntax_example = "date_bin(interval, expression[, origin_timestamp])", sql_example = r#"```sql -- Bin the timestamp into 1 day intervals > SELECT date_bin(interval '1 day', time) as bin @@ -92,7 +95,7 @@ FROM VALUES (TIME '02:18:18'), (TIME '19:00:03') t(time); description = "Time expression to operate on. Can be a constant, column, or function." ), argument( - name = "origin-timestamp", + name = "origin_timestamp", description = r#"Optional. Starting point used to determine bin boundaries. If not specified defaults 1970-01-01T00:00:00Z (the UNIX epoch in UTC). The following intervals are supported: - nanoseconds @@ -322,26 +325,61 @@ impl Interval { // return time in nanoseconds that the source timestamp falls into based on the stride and origin fn date_bin_nanos_interval(stride_nanos: i64, source: i64, origin: i64) -> Result { let time_diff = source.checked_sub(origin).ok_or_else(|| { - arrow::error::ArrowError::InvalidArgumentError(format!( + ArrowError::InvalidArgumentError(format!( "date_bin source timestamp {source} - origin {origin} overflows i64" )) })?; // distance from origin to bin - let time_delta = compute_distance(time_diff, stride_nanos); + let time_delta = compute_distance(time_diff, stride_nanos)?; - Ok(origin + time_delta) + origin.checked_add(time_delta).ok_or_else(|| { + ArrowError::InvalidArgumentError(format!( + "date_bin origin {origin} + delta {time_delta} overflows i64" + )) + .into() + }) } // distance from origin to bin -fn compute_distance(time_diff: i64, stride: i64) -> i64 { - let time_delta = time_diff - (time_diff % stride); +fn compute_distance(time_diff: i64, stride: i64) -> Result { + let remainder = time_diff.checked_rem(stride).ok_or_else(|| { + ArrowError::InvalidArgumentError(format!( + "date_bin compute_distance time_diff {time_diff} % stride {stride} overflows i64" + )) + })?; + let time_delta = time_diff.checked_sub(remainder).ok_or_else(|| { + ArrowError::InvalidArgumentError(format!( + "date_bin compute_distance time_diff {time_diff} - remainder {remainder} overflows i64" + )) + })?; if time_diff < 0 && stride > 1 && time_delta != time_diff { // The origin is later than the source timestamp, round down to the previous bin - time_delta - stride + time_delta.checked_sub(stride).ok_or_else(|| { + ArrowError::InvalidArgumentError(format!( + "date_bin compute_distance time_delta {time_delta} - stride {stride} overflows i64" + )) + .into() + }) + } else { + Ok(time_delta) + } +} + +// Shift `origin_date` by `month_delta` months, mapping an out-of-range result to +// the same error the binning paths reported when this was written inline. +fn shift_months(origin_date: DateTime, month_delta: i64) -> Result> { + if month_delta < 0 { + origin_date + .checked_sub_months(Months::new(month_delta.unsigned_abs() as u32)) + .ok_or_else(|| { + exec_datafusion_err!("DATE_BIN month subtraction out of range") + }) } else { - time_delta + origin_date + .checked_add_months(Months::new(month_delta as u32)) + .ok_or_else(|| exec_datafusion_err!("DATE_BIN month addition out of range")) } } @@ -357,39 +395,15 @@ fn date_bin_months_interval(stride_months: i64, source: i64, origin: i64) -> Res - origin_date.month() as i32; // distance from origin to bin - let month_delta = compute_distance(month_diff as i64, stride_months); + let month_delta = compute_distance(month_diff as i64, stride_months)?; - let mut bin_time = if month_delta < 0 { - match origin_date - .checked_sub_months(Months::new(month_delta.unsigned_abs() as u32)) - { - Some(dt) => dt, - None => return exec_err!("DATE_BIN month subtraction out of range"), - } - } else { - match origin_date.checked_add_months(Months::new(month_delta as u32)) { - Some(dt) => dt, - None => return exec_err!("DATE_BIN month addition out of range"), - } - }; + let mut bin_time = shift_months(origin_date, month_delta)?; // If origin is not midnight of first date of the month, the bin_time may be larger than the source // In this case, we need to move back to previous bin if bin_time > source_date { let month_delta = month_delta - stride_months; - bin_time = if month_delta < 0 { - match origin_date - .checked_sub_months(Months::new(month_delta.unsigned_abs() as u32)) - { - Some(dt) => dt, - None => return exec_err!("DATE_BIN month subtraction out of range"), - } - } else { - match origin_date.checked_add_months(Months::new(month_delta as u32)) { - Some(dt) => dt, - None => return exec_err!("DATE_BIN month addition out of range"), - } - }; + bin_time = shift_months(origin_date, month_delta)?; } match bin_time.timestamp_nanos_opt() { Some(nanos) => Ok(nanos), @@ -398,14 +412,86 @@ fn date_bin_months_interval(stride_months: i64, source: i64, origin: i64) -> Res } fn to_utc_date_time(nanos: i64) -> Result> { - let secs = nanos / NANOS_PER_SEC; - let nsec = (nanos % NANOS_PER_SEC) as u32; + // Keep negative sub-second values normalized as seconds + non-negative nanos. + let secs = nanos.div_euclid(NANOS_PER_SEC); + let nsec = nanos.rem_euclid(NANOS_PER_SEC) as u32; match DateTime::from_timestamp(secs, nsec) { Some(dt) => Ok(dt), None => exec_err!("Invalid timestamp value"), } } +fn timestamp_scale() -> i64 { + match T::UNIT { + Nanosecond => 1, + Microsecond => NANOS_PER_MICRO, + Millisecond => NANOS_PER_MILLI, + Second => NANOSECONDS, + } +} + +// Scale to nanoseconds and report overflow as a normal error. +fn checked_scale_to_nanos(x: i64, scale: i64) -> Result { + match x.checked_mul(scale) { + Some(scaled) => Ok(scaled), + None => exec_err!("date_bin timestamp value {x} * scale {scale} overflows i64"), + } +} + +// Per-row failures map to NULL, so use Option in the hot path. +#[inline] +fn scale_and_bin_to_nanos( + value: i64, + scale: i64, + origin: i64, + stride: i64, + stride_fn: BinFunction, +) -> Option { + value + .checked_mul(scale) + .and_then(|scaled| stride_fn(stride, scaled, origin).ok()) +} + +// Per-row timestamp binning shared by scalar and array paths. +// Source-value failures become None, which callers map to NULL. +#[inline] +fn date_bin_timestamp_value( + value: i64, + origin: i64, + stride: i64, + stride_fn: BinFunction, +) -> Option { + let scale = timestamp_scale::(); + scale_and_bin_to_nanos(value, scale, origin, stride, stride_fn) + .map(|binned| binned / scale) +} + +// Per-row TIME binning shared by scalar and array paths. +// The modulo keeps the result within a single day before unscaling. +#[inline] +fn date_bin_time_value( + value: i64, + scale: i64, + origin: i64, + stride: i64, + stride_fn: BinFunction, +) -> Option { + scale_and_bin_to_nanos(value, scale, origin, stride, stride_fn) + .map(|binned| (binned % NANOSECONDS_IN_DAY) / scale) +} + +fn validate_time_stride(stride: &Interval) -> Result<()> { + match stride { + Interval::Months(m) if *m > 0 => { + exec_err!("DATE_BIN stride for TIME input must be less than 1 day") + } + Interval::Nanoseconds(ns) if *ns >= NANOSECONDS_IN_DAY => { + exec_err!("DATE_BIN stride for TIME input must be less than 1 day") + } + _ => Ok(()), + } +} + // Supported intervals: // 1. IntervalDayTime: this means that the stride is in days, hours, minutes, seconds and milliseconds // We will assume month interval won't be converted into this type @@ -476,83 +562,20 @@ fn date_bin_impl( (*v, false) } ColumnarValue::Scalar(ScalarValue::Time32Millisecond(Some(v))) => { - match stride { - Interval::Months(m) => { - if m > 0 { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - Interval::Nanoseconds(ns) => { - if ns >= NANOSECONDS_IN_DAY { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - } - - (*v as i64 * NANOS_PER_MILLI, true) + validate_time_stride(&stride)?; + // TIME origins can come from reinterpret casts, so scale defensively. + (checked_scale_to_nanos(*v as i64, NANOS_PER_MILLI)?, true) } ColumnarValue::Scalar(ScalarValue::Time32Second(Some(v))) => { - match stride { - Interval::Months(m) => { - if m > 0 { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - Interval::Nanoseconds(ns) => { - if ns >= NANOSECONDS_IN_DAY { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - } - - (*v as i64 * NANOS_PER_SEC, true) + validate_time_stride(&stride)?; + (checked_scale_to_nanos(*v as i64, NANOS_PER_SEC)?, true) } ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(v))) => { - match stride { - Interval::Months(m) => { - if m > 0 { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - Interval::Nanoseconds(ns) => { - if ns >= NANOSECONDS_IN_DAY { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - } - - (*v * NANOS_PER_MICRO, true) + validate_time_stride(&stride)?; + (checked_scale_to_nanos(*v, NANOS_PER_MICRO)?, true) } ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(Some(v))) => { - match stride { - Interval::Months(m) => { - if m > 0 { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - Interval::Nanoseconds(ns) => { - if ns >= NANOSECONDS_IN_DAY { - return exec_err!( - "DATE_BIN stride for TIME input must be less than 1 day" - ); - } - } - } - + validate_time_stride(&stride)?; (*v, true) } ColumnarValue::Scalar(v) => { @@ -575,108 +598,85 @@ fn date_bin_impl( return exec_err!("DATE_BIN stride must be non-zero"); } - fn stride_map_fn( - origin: i64, - stride: i64, - stride_fn: BinFunction, - ) -> impl Fn(i64) -> Result { - let scale = match T::UNIT { - Nanosecond => 1, - Microsecond => NANOS_PER_MICRO, - Millisecond => NANOS_PER_MILLI, - Second => NANOSECONDS, - }; - move |x: i64| match stride_fn(stride, x * scale, origin) { - Ok(result) => Ok(result / scale), - Err(e) => Err(e), + // A TIME source requires a TIME origin. This shared-input check is ordered + // after stride/origin parsing and the zero-stride check so error ordering is + // unchanged, and replaces the per-arm guards in the TIME branches below. + if !is_time { + match array.data_type() { + Time32(_) => { + return exec_err!("DATE_BIN with Time32 source requires Time32 origin"); + } + Time64(_) => { + return exec_err!("DATE_BIN with Time64 source requires Time64 origin"); + } + _ => {} } } Ok(match array { ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(v, tz_opt)) => { - let apply_stride_fn = - stride_map_fn::(origin, stride, stride_fn); ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( - v.and_then(|val| apply_stride_fn(val).ok()), + v.and_then(|x| { + date_bin_timestamp_value::( + x, origin, stride, stride_fn, + ) + }), tz_opt.clone(), )) } ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond(v, tz_opt)) => { - let apply_stride_fn = - stride_map_fn::(origin, stride, stride_fn); ColumnarValue::Scalar(ScalarValue::TimestampMicrosecond( - v.and_then(|val| apply_stride_fn(val).ok()), + v.and_then(|x| { + date_bin_timestamp_value::( + x, origin, stride, stride_fn, + ) + }), tz_opt.clone(), )) } ColumnarValue::Scalar(ScalarValue::TimestampMillisecond(v, tz_opt)) => { - let apply_stride_fn = - stride_map_fn::(origin, stride, stride_fn); ColumnarValue::Scalar(ScalarValue::TimestampMillisecond( - v.and_then(|val| apply_stride_fn(val).ok()), + v.and_then(|x| { + date_bin_timestamp_value::( + x, origin, stride, stride_fn, + ) + }), tz_opt.clone(), )) } ColumnarValue::Scalar(ScalarValue::TimestampSecond(v, tz_opt)) => { - let apply_stride_fn = - stride_map_fn::(origin, stride, stride_fn); ColumnarValue::Scalar(ScalarValue::TimestampSecond( - v.and_then(|val| apply_stride_fn(val).ok()), + v.and_then(|x| { + date_bin_timestamp_value::( + x, origin, stride, stride_fn, + ) + }), tz_opt.clone(), )) } ColumnarValue::Scalar(ScalarValue::Time32Millisecond(v)) => { - if !is_time { - return exec_err!("DATE_BIN with Time32 source requires Time32 origin"); - } let result = v.and_then(|x| { - match stride_fn(stride, x as i64 * NANOS_PER_MILLI, origin) { - Ok(binned_nanos) => { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - Some((nanos / NANOS_PER_MILLI) as i32) - } - Err(_) => None, - } + date_bin_time_value(x as i64, NANOS_PER_MILLI, origin, stride, stride_fn) + .map(|binned| binned as i32) }); ColumnarValue::Scalar(ScalarValue::Time32Millisecond(result)) } ColumnarValue::Scalar(ScalarValue::Time32Second(v)) => { - if !is_time { - return exec_err!("DATE_BIN with Time32 source requires Time32 origin"); - } let result = v.and_then(|x| { - match stride_fn(stride, x as i64 * NANOS_PER_SEC, origin) { - Ok(binned_nanos) => { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - Some((nanos / NANOS_PER_SEC) as i32) - } - Err(_) => None, - } + date_bin_time_value(x as i64, NANOS_PER_SEC, origin, stride, stride_fn) + .map(|binned| binned as i32) }); ColumnarValue::Scalar(ScalarValue::Time32Second(result)) } ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(v)) => { - if !is_time { - return exec_err!("DATE_BIN with Time64 source requires Time64 origin"); - } - let result = v.and_then(|x| match stride_fn(stride, x, origin) { - Ok(binned_nanos) => Some(binned_nanos % (NANOSECONDS_IN_DAY)), - Err(_) => None, - }); + let result = + v.and_then(|x| date_bin_time_value(x, 1, origin, stride, stride_fn)); ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(result)) } ColumnarValue::Scalar(ScalarValue::Time64Microsecond(v)) => { - if !is_time { - return exec_err!("DATE_BIN with Time64 source requires Time64 origin"); - } - let result = - v.and_then(|x| match stride_fn(stride, x * NANOS_PER_MICRO, origin) { - Ok(binned_nanos) => { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - Some(nanos / NANOS_PER_MICRO) - } - Err(_) => None, - }); + let result = v.and_then(|x| { + date_bin_time_value(x, NANOS_PER_MICRO, origin, stride, stride_fn) + }); ColumnarValue::Scalar(ScalarValue::Time64Microsecond(result)) } ColumnarValue::Array(array) => { @@ -691,20 +691,11 @@ fn date_bin_impl( T: ArrowTimestampType, { let array = as_primitive_array::(array)?; - let scale = match T::UNIT { - Nanosecond => 1, - Microsecond => NANOS_PER_MICRO, - Millisecond => NANOS_PER_MILLI, - Second => NANOSECONDS, - }; - - let result: PrimitiveArray = array.try_unary(|val| { - stride_fn(stride, val * scale, origin) - .map(|binned| binned / scale) - .map_err(|e| { - arrow::error::ArrowError::ComputeError(e.to_string()) - }) - })?; + + // Per-row errors become NULL, matching scalar behavior. + let result: PrimitiveArray = array.unary_opt(|val| { + date_bin_timestamp_value::(val, origin, stride, stride_fn) + }); let array = result.with_timezone_opt(tz_opt.clone()); Ok(ColumnarValue::Array(Arc::new(array))) @@ -732,80 +723,54 @@ fn date_bin_impl( )? } Time32(Millisecond) => { - if !is_time { - return exec_err!( - "DATE_BIN with Time32 source requires Time32 origin" - ); - } let array = array.as_primitive::(); let result: PrimitiveArray = - array.try_unary(|x| { - stride_fn(stride, x as i64 * NANOS_PER_MILLI, origin) - .map(|binned_nanos| { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - (nanos / NANOS_PER_MILLI) as i32 - }) - .map_err(|e| { - arrow::error::ArrowError::ComputeError(e.to_string()) - }) - })?; + array.unary_opt(|x| { + date_bin_time_value( + x as i64, + NANOS_PER_MILLI, + origin, + stride, + stride_fn, + ) + .map(|binned| binned as i32) + }); ColumnarValue::Array(Arc::new(result)) } Time32(Second) => { - if !is_time { - return exec_err!( - "DATE_BIN with Time32 source requires Time32 origin" - ); - } let array = array.as_primitive::(); - let result: PrimitiveArray = - array.try_unary(|x| { - stride_fn(stride, x as i64 * NANOS_PER_SEC, origin) - .map(|binned_nanos| { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - (nanos / NANOS_PER_SEC) as i32 - }) - .map_err(|e| { - arrow::error::ArrowError::ComputeError(e.to_string()) - }) - })?; + let result: PrimitiveArray = array.unary_opt(|x| { + date_bin_time_value( + x as i64, + NANOS_PER_SEC, + origin, + stride, + stride_fn, + ) + .map(|binned| binned as i32) + }); ColumnarValue::Array(Arc::new(result)) } Time64(Microsecond) => { - if !is_time { - return exec_err!( - "DATE_BIN with Time64 source requires Time64 origin" - ); - } let array = array.as_primitive::(); let result: PrimitiveArray = - array.try_unary(|x| { - stride_fn(stride, x * NANOS_PER_MICRO, origin) - .map(|binned_nanos| { - let nanos = binned_nanos % (NANOSECONDS_IN_DAY); - nanos / NANOS_PER_MICRO - }) - .map_err(|e| { - arrow::error::ArrowError::ComputeError(e.to_string()) - }) - })?; + array.unary_opt(|x| { + date_bin_time_value( + x, + NANOS_PER_MICRO, + origin, + stride, + stride_fn, + ) + }); ColumnarValue::Array(Arc::new(result)) } Time64(Nanosecond) => { - if !is_time { - return exec_err!( - "DATE_BIN with Time64 source requires Time64 origin" - ); - } let array = array.as_primitive::(); let result: PrimitiveArray = - array.try_unary(|x| { - stride_fn(stride, x, origin) - .map(|binned_nanos| binned_nanos % (NANOSECONDS_IN_DAY)) - .map_err(|e| { - arrow::error::ArrowError::ComputeError(e.to_string()) - }) - })?; + array.unary_opt(|x| { + date_bin_time_value(x, 1, origin, stride, stride_fn) + }); ColumnarValue::Array(Arc::new(result)) } _ => { @@ -861,6 +826,31 @@ mod tests { DateBinFunc::new().invoke_with_args(args) } + fn assert_null_scalar(value: ColumnarValue, expected_type: DataType) { + let ColumnarValue::Scalar(value) = value else { + panic!("expected scalar, got {value:?}"); + }; + assert_eq!(value.data_type(), expected_type); + assert!(value.is_null(), "expected NULL, got {value:?}"); + } + + fn assert_array_null_then_valid(value: ColumnarValue, expected_type: DataType) { + let ColumnarValue::Array(array) = value else { + panic!("expected array, got {value:?}"); + }; + assert_eq!(array.data_type(), &expected_type); + assert!(array.is_null(0), "expected NULL at row 0"); + assert!(array.is_valid(1), "expected valid value at row 1"); + } + + fn assert_overflow_error(result: Result) { + let err = result.expect_err("expected overflow error"); + assert!( + err.strip_backtrace().contains("overflows i64"), + "unexpected error: {err}" + ); + } + #[test] fn test_date_bin() { let return_field = &Arc::new(Field::new( @@ -1341,4 +1331,281 @@ mod tests { assert!(val.is_none(), "Expected None for out of range operation"); } } + + #[test] + fn test_date_bin_compute_distance_i64_min() { + // Regression for #22215: date_bin_nanos_interval on a source near i64::MIN + // previously panicked inside compute_distance with "attempt to subtract with overflow". + // Now it must return a normal Err that the scalar pipeline maps to NULL. + let result = date_bin_nanos_interval(3, i64::MIN, 0); + assert!( + result.is_err(), + "expected Err for source=i64::MIN, got {result:?}" + ); + + let return_field = &Arc::new(Field::new( + "f", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::new_interval_mdn(0, 0, 3)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(i64::MIN), None)), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)), + ]; + let result = invoke_date_bin_with_args(args, 1, return_field); + assert!(result.is_ok(), "expected Ok with NULL, got {result:?}"); + if let ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(val, _)) = + result.unwrap() + { + assert!( + val.is_none(), + "Expected None for compute_distance overflow, got {val:?}" + ); + } else { + panic!("Expected TimestampNanosecond scalar"); + } + } + + #[test] + fn test_date_bin_scale_overflow_returns_null() { + // Scaling non-nanosecond timestamps to nanoseconds can overflow. + use arrow::array::{ + ArrayRef, TimestampMicrosecondArray, TimestampMillisecondArray, + TimestampSecondArray, + }; + + let scalar_cases = [ + ScalarValue::TimestampSecond(Some(i64::MAX), None), + ScalarValue::TimestampMillisecond(Some(i64::MAX), None), + ScalarValue::TimestampMicrosecond(Some(i64::MAX), None), + ]; + for source in scalar_cases { + let expected_type = source.data_type(); + let return_field = Arc::new(Field::new("f", expected_type.clone(), true)); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), + ColumnarValue::Scalar(source), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)), + ]; + let result = invoke_date_bin_with_args(args, 1, &return_field) + .unwrap_or_else(|e| panic!("expected Ok for {expected_type}, got {e:?}")); + assert_null_scalar(result, expected_type); + } + + let array_cases: Vec = vec![ + Arc::new(TimestampSecondArray::from(vec![Some(i64::MAX), Some(0)])), + Arc::new(TimestampMillisecondArray::from(vec![ + Some(i64::MAX), + Some(0), + ])), + Arc::new(TimestampMicrosecondArray::from(vec![ + Some(i64::MAX), + Some(0), + ])), + ]; + for array in array_cases { + let dt = array.data_type().clone(); + let return_field = Arc::new(Field::new("f", dt.clone(), true)); + let args = vec![ + ColumnarValue::Scalar(ScalarValue::new_interval_dt(1, 0)), + ColumnarValue::Array(array), + ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)), + ]; + let result = invoke_date_bin_with_args(args, 2, &return_field) + .unwrap_or_else(|e| panic!("expected Ok for {dt:?}, got {e:?}")); + assert_array_null_then_valid(result, dt); + } + } + + #[test] + fn test_date_bin_time64_micro_overflow_handling() { + // Time64(Microsecond) can hold out-of-range values after reinterpret casts. + use arrow::array::Time64MicrosecondArray; + + let data_type = DataType::Time64(TimeUnit::Microsecond); + let return_field = &Arc::new(Field::new("f", data_type.clone(), true)); + let stride = || ColumnarValue::Scalar(ScalarValue::new_interval_dt(0, 1000)); + let origin = || ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(0))); + + // Out-of-range source values are per-row data, so they become NULL. + let args = vec![ + stride(), + ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(i64::MAX))), + origin(), + ]; + let result = invoke_date_bin_with_args(args, 1, return_field).unwrap(); + assert_null_scalar(result, data_type.clone()); + + let array = Arc::new(Time64MicrosecondArray::from(vec![Some(i64::MAX), Some(0)])); + let args = vec![stride(), ColumnarValue::Array(array), origin()]; + let result = invoke_date_bin_with_args(args, 2, return_field).unwrap(); + assert_array_null_then_valid(result, data_type); + + let bad_origin = + || ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(i64::MAX))); + + // Out-of-range origins are shared inputs, so they return an error. + let args = vec![ + stride(), + ColumnarValue::Scalar(ScalarValue::Time64Microsecond(Some(0))), + bad_origin(), + ]; + assert_overflow_error(invoke_date_bin_with_args(args, 1, return_field)); + + let array = Arc::new(Time64MicrosecondArray::from(vec![Some(0), Some(1)])); + let args = vec![stride(), ColumnarValue::Array(array), bad_origin()]; + assert_overflow_error(invoke_date_bin_with_args(args, 2, return_field)); + } + + // Compare scalar execution with a one-row array for the same input. + fn assert_scalar_array_parity( + stride: ScalarValue, + source: ScalarValue, + origin: ScalarValue, + ) { + let return_field = Arc::new(Field::new("f", source.data_type().clone(), true)); + + let scalar_args = vec![ + ColumnarValue::Scalar(stride.clone()), + ColumnarValue::Scalar(source.clone()), + ColumnarValue::Scalar(origin.clone()), + ]; + let scalar_result = invoke_date_bin_with_args(scalar_args, 1, &return_field) + .expect("scalar path should not error"); + let ColumnarValue::Scalar(scalar_value) = scalar_result else { + panic!("expected scalar result, got {scalar_result:?}"); + }; + + let source_array = source.to_array().expect("source value to array"); + let array_args = vec![ + ColumnarValue::Scalar(stride), + ColumnarValue::Array(source_array), + ColumnarValue::Scalar(origin), + ]; + let array_result = invoke_date_bin_with_args(array_args, 1, &return_field) + .expect("array path should not error"); + let ColumnarValue::Array(array) = array_result else { + panic!("expected array result, got {array_result:?}"); + }; + let array_value = + ScalarValue::try_from_array(&array, 0).expect("array row to scalar"); + + assert_eq!( + scalar_value, array_value, + "scalar and array results diverged for source {source:?}" + ); + } + + #[test] + fn test_date_bin_scalar_array_parity() { + // Negative sub-second timestamp with a month interval. This is the case + // that previously diverged (scalar value vs array execution error) + // before #22610; both paths must now agree on the same non-NULL value. + assert_scalar_array_parity( + ScalarValue::new_interval_mdn(1, 0, 0), + ScalarValue::TimestampNanosecond(Some(-1), None), + ScalarValue::TimestampNanosecond(Some(0), None), + ); + + // Source scaling overflow -> NULL in both paths. + assert_scalar_array_parity( + ScalarValue::new_interval_dt(1, 0), + ScalarValue::TimestampSecond(Some(i64::MAX), None), + ScalarValue::TimestampNanosecond(Some(0), None), + ); + + // Month interval out-of-range binning -> NULL in both paths. + assert_scalar_array_parity( + ScalarValue::new_interval_mdn(1637426858, 0, 0), + ScalarValue::TimestampMillisecond(Some(1040292460), None), + ScalarValue::TimestampNanosecond( + Some(string_to_timestamp_nanos("1984-01-07 00:00:00").unwrap()), + None, + ), + ); + } + + #[test] + fn test_date_bin_time_source_requires_time_origin() { + // A TIME source combined with a non-TIME (timestamp) origin is rejected + // with a unit-specific message. This is the shared-input guard that was + // hoisted out of the per-type match arms; cover scalar and array for + // both Time32 and Time64 so the error text stays put. + use arrow::array::{Time32MillisecondArray, Time64NanosecondArray}; + + let stride = || ColumnarValue::Scalar(ScalarValue::new_interval_dt(0, 1000)); + let ts_origin = + || ColumnarValue::Scalar(ScalarValue::TimestampNanosecond(Some(0), None)); + + let assert_msg = |args: Vec, dt: DataType, msg: &str| { + let return_field = Arc::new(Field::new("f", dt, true)); + assert_eq!( + invoke_date_bin_with_args(args, 1, &return_field) + .err() + .unwrap() + .strip_backtrace(), + msg + ); + }; + + let time32_msg = + "Execution error: DATE_BIN with Time32 source requires Time32 origin"; + assert_msg( + vec![ + stride(), + ColumnarValue::Scalar(ScalarValue::Time32Millisecond(Some(0))), + ts_origin(), + ], + DataType::Time32(TimeUnit::Millisecond), + time32_msg, + ); + assert_msg( + vec![ + stride(), + ColumnarValue::Array(Arc::new(Time32MillisecondArray::from(vec![Some( + 0, + )]))), + ts_origin(), + ], + DataType::Time32(TimeUnit::Millisecond), + time32_msg, + ); + + let time64_msg = + "Execution error: DATE_BIN with Time64 source requires Time64 origin"; + assert_msg( + vec![ + stride(), + ColumnarValue::Scalar(ScalarValue::Time64Nanosecond(Some(0))), + ts_origin(), + ], + DataType::Time64(TimeUnit::Nanosecond), + time64_msg, + ); + assert_msg( + vec![ + stride(), + ColumnarValue::Array(Arc::new(Time64NanosecondArray::from(vec![Some( + 0, + )]))), + ts_origin(), + ], + DataType::Time64(TimeUnit::Nanosecond), + time64_msg, + ); + } + + #[test] + fn test_date_bin_compute_distance_rem_overflow() { + // Regression for #22215: `time_diff % stride` panics with "attempt to + // calculate the remainder with overflow" when `time_diff == i64::MIN` + // and `stride == -1`. Now it must return a normal Err that the scalar + // pipeline maps to NULL. + let result = date_bin_nanos_interval(-1, i64::MIN, 0); + assert!( + result.is_err(), + "expected Err for time_diff=i64::MIN, stride=-1, got {result:?}" + ); + } } diff --git a/datafusion/functions/src/datetime/date_part.rs b/datafusion/functions/src/datetime/date_part.rs index 3c405d388bcab..e3f67db905615 100644 --- a/datafusion/functions/src/datetime/date_part.rs +++ b/datafusion/functions/src/datetime/date_part.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::iter::repeat_n; use std::str::FromStr; use std::sync::Arc; @@ -240,16 +241,8 @@ impl ScalarUDFImpl for DatePartFunc { "doy" => date_part(array.as_ref(), DatePart::DayOfYear)?, "dow" => date_part(array.as_ref(), DatePart::DayOfWeekSunday0)?, "isodow" => { - // Postgres `isodow` is 1..=7 with Mon=1. Arrow's - // `DayOfWeekMonday0` returns 0..=6 with Mon=0; shift by - // +1 to match Postgres. TODO: switch to a future - // `DatePart::DayOfWeekMonday1` upstream variant once it - // exists, so this kernel-then-add becomes a single call. - let zero_based = - date_part(array.as_ref(), DatePart::DayOfWeekMonday0)?; - let int_arr = as_int32_array(&zero_based)?; - let one_based: Int32Array = int_arr.unary(|v| v + 1); - Arc::new(one_based) as ArrayRef + // Postgres `isodow` is 1..=7 with Mon=1 + date_part(array.as_ref(), DatePart::DayOfWeekMonday1)? } "epoch" => epoch(array.as_ref())?, _ => return exec_err!("Date part '{part}' not supported"), @@ -398,7 +391,7 @@ fn part_normalization(part: &str) -> &str { /// Invoke [`date_part`] on an `array` (e.g. Timestamp) and convert the /// result to a total number of seconds, milliseconds, microseconds or -/// nanoseconds +/// nanoseconds as an `Int32Array` fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result { // Nanosecond is neither supported in Postgres nor DuckDB, to avoid dealing // with overflow and precision issue we don't support nanosecond @@ -406,6 +399,19 @@ fn seconds_as_i32(array: &dyn Array, unit: TimeUnit) -> Result { return not_impl_err!("Date part {unit:?} not supported"); } + // Fast path with seconds - no need to compute nanoseconds + if unit == Second { + return Ok(date_part(array, DatePart::Second)?); + } + + // Fast path for Date32 and Date64 - no seconds + if array.data_type() == &Date32 || array.data_type() == &Date64 { + return Ok(Arc::new(Int32Array::from_iter_values_with_nulls( + repeat_n(0, array.len()), + array.nulls().cloned(), + ))); + } + let conversion_factor = match unit { Second => 1_000_000_000, Millisecond => 1_000_000, @@ -547,6 +553,14 @@ fn epoch(array: &dyn Array) -> Result { /// `nanosecond`s in each second, so representing up to 60 seconds as /// nanoseconds can be values up to 60 billion, which does not fit in Int32. fn seconds_ns(array: &dyn Array) -> Result { + // Fast path for Date32 and Date64 - no nanoseconds + if array.data_type() == &Date32 || array.data_type() == &Date64 { + return Ok(Arc::new(Int64Array::from_iter_values_with_nulls( + repeat_n(0, array.len()), + array.nulls().cloned(), + ))); + } + let secs = date_part(array, DatePart::Second)?; // This assumes array is primitive and not a dictionary let secs = as_int32_array(secs.as_ref())?; diff --git a/datafusion/functions/src/datetime/date_trunc.rs b/datafusion/functions/src/datetime/date_trunc.rs index 784f593c2529d..6dcd7a666d0a6 100644 --- a/datafusion/functions/src/datetime/date_trunc.rs +++ b/datafusion/functions/src/datetime/date_trunc.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::fmt; use std::num::NonZeroI64; use std::ops::{Add, Sub}; use std::str::FromStr; @@ -22,7 +23,6 @@ use std::sync::Arc; use arrow::array::temporal_conversions::{ MICROSECONDS, MILLISECONDS, NANOSECONDS, as_datetime_with_timezone, - timestamp_ns_to_datetime, }; use arrow::array::timezone::Tz; use arrow::array::types::{ @@ -135,6 +135,24 @@ impl DateTruncGranularity { } } +impl fmt::Display for DateTruncGranularity { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let value = match self { + Self::Microsecond => "microsecond", + Self::Millisecond => "millisecond", + Self::Second => "second", + Self::Minute => "minute", + Self::Hour => "hour", + Self::Day => "day", + Self::Week => "week", + Self::Month => "month", + Self::Quarter => "quarter", + Self::Year => "year", + }; + f.write_str(value) + } +} + #[user_doc( doc_section(label = "Time and Date Functions"), description = "Truncates a timestamp or time value to a specified precision.", @@ -443,6 +461,7 @@ const NANOS_PER_MILLISECOND: i64 = NANOSECONDS / MILLISECONDS; const NANOS_PER_SECOND: i64 = NANOSECONDS; const NANOS_PER_MINUTE: i64 = 60 * NANOS_PER_SECOND; const NANOS_PER_HOUR: i64 = 60 * NANOS_PER_MINUTE; +const NANOS_PER_DAY: i64 = 24 * NANOS_PER_HOUR; const MICROS_PER_MILLISECOND: i64 = MICROSECONDS / MILLISECONDS; const MICROS_PER_SECOND: i64 = MICROSECONDS; @@ -572,52 +591,143 @@ where fn _date_trunc_coarse_with_tz( granularity: DateTruncGranularity, - value: Option>, + value: DateTime, ) -> Result> { - if let Some(value) = value { - let local = value.naive_local(); - let truncated = _date_trunc_coarse::(granularity, Some(local))?; - let truncated = truncated.and_then(|truncated| { - match truncated.and_local_timezone(value.timezone()) { - LocalResult::None => { - // This can happen if the date_trunc operation moves the time into - // an hour that doesn't exist due to daylight savings. On known example where - // this can happen is with historic dates in the America/Sao_Paulo time zone. - // To account for this adjust the time by a few hours, convert to local time, - // and then adjust the time back. - truncated - .sub(TimeDelta::try_hours(3).unwrap()) - .and_local_timezone(value.timezone()) - .single() - .map(|v| v.add(TimeDelta::try_hours(3).unwrap())) - } - LocalResult::Single(datetime) => Some(datetime), - LocalResult::Ambiguous(datetime1, datetime2) => { - // Because we are truncating from an equally or more specific time - // the original time must have been within the ambiguous local time - // period. Therefore the offset of one of these times should match the - // offset of the original time. - if datetime1.offset().fix() == value.offset().fix() { - Some(datetime1) - } else { - Some(datetime2) - } + let local = value.naive_local(); + let truncated = _date_trunc_coarse::(granularity, Some(local))?; + let truncated = truncated.and_then(|truncated| { + match truncated.and_local_timezone(value.timezone()) { + LocalResult::None => { + // This can happen if the date_trunc operation moves the time into + // an hour that doesn't exist due to daylight savings. On known example where + // this can happen is with historic dates in the America/Sao_Paulo time zone. + // To account for this adjust the time by a few hours, convert to local time, + // and then adjust the time back. + truncated + .sub(TimeDelta::try_hours(3).unwrap()) + .and_local_timezone(value.timezone()) + .single() + .map(|v| v.add(TimeDelta::try_hours(3).unwrap())) + } + LocalResult::Single(datetime) => Some(datetime), + LocalResult::Ambiguous(datetime1, datetime2) => { + // Because we are truncating from an equally or more specific time + // the original time must have been within the ambiguous local time + // period. Therefore the offset of one of these times should match the + // offset of the original time. + if datetime1.offset().fix() == value.offset().fix() { + Some(datetime1) + } else { + Some(datetime2) } } - }); - Ok(truncated.and_then(|value| value.timestamp_nanos_opt())) + } + }); + Ok(truncated.and_then(|value| value.timestamp_nanos_opt())) +} + +// The two helpers below duplicate `chrono::NaiveDate::{from_epoch_days, +// to_epoch_days}`. They are kept separate because chrono's versions round trip +// through a validated `NaiveDate`: `from_epoch_days` computes year flags and +// returns an `Option`, and reading the year/month/day back out decodes them from +// its packed representation. These helpers stay in plain integers, which is all +// the truncation below needs. + +/// Days from the Unix epoch to 0000-03-01, the epoch used by the civil calendar +/// conversions below. +const DAYS_EPOCH_SHIFT: i64 = 719_468; + +/// Days in a 400 year era of the proleptic Gregorian calendar. +const DAYS_PER_ERA: i64 = 146_097; + +/// Splits a day count relative to the Unix epoch into a proleptic Gregorian +/// year, month (1-12) and day of month (1-31). +/// +/// This is a port of Howard Hinnant's `civil_from_days`, which documents the +/// derivation of the constants and the March-based year used below: +/// +fn civil_from_days(days: i64) -> (i64, i64, i64) { + let z = days + DAYS_EPOCH_SHIFT; + let era = z.div_euclid(DAYS_PER_ERA); + let day_of_era = z.rem_euclid(DAYS_PER_ERA); + let year_of_era = (day_of_era - day_of_era / 1460 + day_of_era / 36524 + - day_of_era / 146_096) + / 365; + let day_of_year = + day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100); + // Month index with March as 0, so that the leap day falls at the end of the year. + let month_index = (5 * day_of_year + 2) / 153; + let day = day_of_year - (153 * month_index + 2) / 5 + 1; + let month = if month_index < 10 { + month_index + 3 } else { - _date_trunc_coarse::(granularity, None)?; - Ok(None) - } + month_index - 9 + }; + let year = year_of_era + era * 400 + i64::from(month <= 2); + (year, month, day) +} + +/// Inverse of [`civil_from_days`]: the day count relative to the Unix epoch for +/// the given proleptic Gregorian date. +/// +/// This is a port of Howard Hinnant's `days_from_civil`, which documents the +/// derivation of the constants and the March-based year used below: +/// +fn days_from_civil(year: i64, month: i64, day: i64) -> i64 { + let year = year - i64::from(month <= 2); + let era = year.div_euclid(400); + let year_of_era = year.rem_euclid(400); + let month_index = if month > 2 { month - 3 } else { month + 9 }; + let day_of_year = (153 * month_index + 2) / 5 + day - 1; + let day_of_era = + year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year; + era * DAYS_PER_ERA + day_of_era - DAYS_EPOCH_SHIFT } +/// Truncates a UTC nanosecond timestamp with integer arithmetic. Truncating on +/// the calendar directly avoids converting every value to a `NaiveDateTime` and +/// rebuilding it field by field. +/// +/// Returns `None` when the truncated timestamp is no longer representable as +/// nanoseconds since the epoch, which the caller reports as an out of range +/// error. fn _date_trunc_coarse_without_tz( granularity: DateTruncGranularity, - value: Option, -) -> Result> { - let value = _date_trunc_coarse::(granularity, value)?; - Ok(value.and_then(|value| value.and_utc().timestamp_nanos_opt())) + value: i64, +) -> Option { + let truncate_to = |unit: i64| value.checked_sub(value.rem_euclid(unit)); + let days = || value.div_euclid(NANOS_PER_DAY); + let nanos_from_days = |days: i64| days.checked_mul(NANOS_PER_DAY); + + match granularity { + // Sub-second granularities are applied by the caller, which rescales + // the nanoseconds to the time unit of the array. + DateTruncGranularity::Millisecond | DateTruncGranularity::Microsecond => { + Some(value) + } + DateTruncGranularity::Second => truncate_to(NANOS_PER_SECOND), + DateTruncGranularity::Minute => truncate_to(NANOS_PER_MINUTE), + DateTruncGranularity::Hour => truncate_to(NANOS_PER_HOUR), + DateTruncGranularity::Day => nanos_from_days(days()), + DateTruncGranularity::Week => { + let days = days(); + // `Weekday::num_days_from_monday` for the epoch (a Thursday) is 3. + nanos_from_days(days - (days + 3).rem_euclid(7)) + } + DateTruncGranularity::Month => { + let days = days(); + let (_, _, day_of_month) = civil_from_days(days); + nanos_from_days(days - (day_of_month - 1)) + } + DateTruncGranularity::Quarter => { + let (year, month, _) = civil_from_days(days()); + nanos_from_days(days_from_civil(year, 1 + 3 * ((month - 1) / 3), 1)) + } + DateTruncGranularity::Year => { + let (year, _, _) = civil_from_days(days()); + nanos_from_days(days_from_civil(year, 1, 1)) + } + } } /// Truncates the single `value`, expressed in nanoseconds since the @@ -629,24 +739,23 @@ fn date_trunc_coarse( value: i64, tz: Option, ) -> Result { + let input = value; let value = match tz { Some(tz) => { // Use chrono DateTime to clear the various fields because need to clear per timezone, // and NaiveDateTime (ISO 8601) has no concept of timezones let value = as_datetime_with_timezone::(value, tz) .ok_or(exec_datafusion_err!("Timestamp {value} out of range"))?; - _date_trunc_coarse_with_tz(granularity, Some(value)) + _date_trunc_coarse_with_tz(granularity, value)? } - None => { - // Use chrono NaiveDateTime to clear the various fields, if we don't have a timezone. - let value = timestamp_ns_to_datetime(value) - .ok_or_else(|| exec_datafusion_err!("Timestamp {value} out of range"))?; - _date_trunc_coarse_without_tz(granularity, Some(value)) - } - }?; + None => _date_trunc_coarse_without_tz(granularity, value), + }; - // `with_x(0)` are infallible because `0` are always a valid - Ok(value.unwrap()) + value.ok_or_else(|| { + exec_datafusion_err!( + "Timestamp {input} out of range after truncating to {granularity}" + ) + }) } /// Fast path for fine granularities (hour and smaller) that can be handled @@ -879,6 +988,19 @@ mod tests { }); } + #[test] + fn date_trunc_out_of_range_lower_bound_returns_error() { + let timestamp = string_to_timestamp_nanos("1677-09-22T00:00:00Z").unwrap(); + let err = date_trunc_coarse(DateTruncGranularity::Year, timestamp, None) + .unwrap_err() + .to_string(); + + assert!( + err.contains("out of range after truncating to year"), + "{err}" + ); + } + #[test] fn test_date_trunc_timezones() { let cases = [ diff --git a/datafusion/functions/src/datetime/from_unixtime.rs b/datafusion/functions/src/datetime/from_unixtime.rs index 4787c75b610b6..85494f3abff73 100644 --- a/datafusion/functions/src/datetime/from_unixtime.rs +++ b/datafusion/functions/src/datetime/from_unixtime.rs @@ -22,6 +22,7 @@ use arrow::datatypes::TimeUnit::Second; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::{Result, ScalarValue, exec_err, internal_err}; use datafusion_expr::TypeSignature::Exact; +use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, @@ -147,6 +148,24 @@ impl ScalarUDFImpl for FromUnixtimeFunc { } } + fn output_ordering(&self, inputs: &[ExprProperties]) -> Result { + // The optional timezone argument must be a constant string and only + // affects the display metadata, not the stored epoch value, so the + // output ordering follows the first argument. + Ok(inputs[0].sort_properties) + } + + fn preserves_lex_ordering(&self, _inputs: &[ExprProperties]) -> Result { + Ok(true) + } + + fn strictly_order_preserving(&self, _inputs: &[ExprProperties]) -> Result { + // `from_unixtime` stores the input's exact `Int64` value as a + // `Timestamp(Second)`: the mapping is one-to-one, order-preserving, + // and maps nulls to nulls. + Ok(true) + } + fn documentation(&self) -> Option<&Documentation> { self.doc() } diff --git a/datafusion/functions/src/datetime/make_date.rs b/datafusion/functions/src/datetime/make_date.rs index dc1328742f24e..3a6b76ed86eb0 100644 --- a/datafusion/functions/src/datetime/make_date.rs +++ b/datafusion/functions/src/datetime/make_date.rs @@ -17,10 +17,10 @@ use std::sync::Arc; -use arrow::array::builder::PrimitiveBuilder; use arrow::array::cast::AsArray; use arrow::array::types::{Date32Type, Int32Type}; use arrow::array::{Array, PrimitiveArray}; +use arrow::buffer::NullBuffer; use arrow::datatypes::DataType; use arrow::datatypes::DataType::Date32; use chrono::prelude::*; @@ -139,24 +139,27 @@ impl ScalarUDFImpl for MakeDateFunc { let months = months.as_primitive::(); let days = days.as_primitive::(); - let mut builder: PrimitiveBuilder = - PrimitiveArray::builder(len); + let nulls = + NullBuffer::union_many([years.nulls(), months.nulls(), days.nulls()]); + let mut values = Vec::with_capacity(len); for i in 0..len { // match postgresql behaviour which returns null for any null input - if years.is_null(i) || months.is_null(i) || days.is_null(i) { - builder.append_null(); + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { + values.push(0); } else { make_date_inner( years.value(i), months.value(i), days.value(i), - |days: i32| builder.append_value(days), + |days: i32| values.push(days), )?; } } - Ok(ColumnarValue::Array(Arc::new(builder.finish()))) + Ok(ColumnarValue::Array(Arc::new( + PrimitiveArray::::new(values.into(), nulls), + ))) } } } @@ -197,3 +200,88 @@ fn make_date_inner( exec_err!("Unable to parse date from {year}, {month}, {day}") } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::Int32Array; + use arrow::datatypes::Field; + use datafusion_common::config::ConfigOptions; + + fn invoke(args: Vec, number_rows: usize) -> Result { + let arg_fields = args + .iter() + .map(|a| Field::new("a", a.data_type(), true).into()) + .collect::>(); + MakeDateFunc::new().invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows, + return_field: Field::new("f", Date32, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }) + } + + #[test] + fn test_make_date_array() { + let years = ColumnarValue::Array(Arc::new(Int32Array::from(vec![ + Some(1970), + Some(1970), + ]))); + let months = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(1)]))); + let days = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(2)]))); + + let ColumnarValue::Array(arr) = invoke(vec![years, months, days], 2).unwrap() + else { + panic!("expected array result"); + }; + let arr = arr.as_primitive::(); + // Days since the unix epoch. + assert_eq!(arr.value(0), 0); + assert_eq!(arr.value(1), 1); + } + + #[test] + fn test_make_date_null_propagation() { + // A NULL in any component column yields a NULL row. + let years = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(2000), None]))); + let months = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(6), Some(6)]))); + let days = ColumnarValue::Array(Arc::new(Int32Array::from(vec![None, Some(15)]))); + + let ColumnarValue::Array(arr) = invoke(vec![years, months, days], 2).unwrap() + else { + panic!("expected array result"); + }; + let arr = arr.as_primitive::(); + assert!(arr.is_null(0)); + assert!(arr.is_null(1)); + } + + #[test] + fn test_make_date_scalar_array_mix() { + let year = ColumnarValue::Scalar(ScalarValue::Int32(Some(1970))); + let month = ColumnarValue::Scalar(ScalarValue::Int32(Some(1))); + let days = + ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1), Some(3)]))); + + let ColumnarValue::Array(arr) = invoke(vec![year, month, days], 2).unwrap() + else { + panic!("expected array result"); + }; + let arr = arr.as_primitive::(); + assert_eq!(arr.value(0), 0); + assert_eq!(arr.value(1), 2); + } + + #[test] + fn test_make_date_out_of_range_errors() { + let years = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(2000)]))); + let months = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(13)]))); + let days = ColumnarValue::Array(Arc::new(Int32Array::from(vec![Some(1)]))); + assert!(invoke(vec![years, months, days], 1).is_err()); + } +} diff --git a/datafusion/functions/src/datetime/to_char.rs b/datafusion/functions/src/datetime/to_char.rs index 5accddd07f2b4..1d3847117420a 100644 --- a/datafusion/functions/src/datetime/to_char.rs +++ b/datafusion/functions/src/datetime/to_char.rs @@ -57,10 +57,6 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo argument( name = "format", description = "A [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) string to use to convert the expression." - ), - argument( - name = "day", - description = "Day to use when making the date. Can be a constant, column or function, and any combination of arithmetic operators." ) )] #[derive(Debug, PartialEq, Eq, Hash)] diff --git a/datafusion/functions/src/datetime/to_date.rs b/datafusion/functions/src/datetime/to_date.rs index cd75ac6bed3ac..668c6ce029751 100644 --- a/datafusion/functions/src/datetime/to_date.rs +++ b/datafusion/functions/src/datetime/to_date.rs @@ -38,7 +38,7 @@ Integers and doubles are interpreted as days since the unix epoch (`1970-01-01T0 Returns the corresponding date. Note: `to_date` returns Date32, which represents its values as the number of days since unix epoch(`1970-01-01`) stored as signed 32 bit value. The largest supported date value is `9999-12-31`.", - syntax_example = "to_date('2017-05-31', '%Y-%m-%d')", + syntax_example = "to_date(expression[, format1, ..., format_n])", sql_example = r#"```sql > select to_date('2023-01-31'); +-------------------------------+ @@ -61,7 +61,7 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo name = "format_n", description = r"Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully parse the expression - an error will be returned." + an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL." ) )] #[derive(Debug, PartialEq, Eq, Hash)] @@ -100,7 +100,7 @@ impl ToDateFunc { args, |s, format| { string_to_timestamp_millis_formatted(s, format) - .map(|n| n / (24 * 60 * 60 * 1_000)) + .map(|n| n.div_euclid(24 * 60 * 60 * 1_000)) .and_then(|v| { v.try_into().map_err(|_| { internal_datafusion_err!("Unable to cast to Date32 for converting from i64 to i32 failed") @@ -519,4 +519,44 @@ mod tests { panic!("Conversion of {date_str} succeeded, but should have failed. "); } } + + /// A NULL format must be skipped even when its slot still holds parseable + /// bytes, otherwise it can silently win over a later valid format. + #[test] + fn test_to_date_null_format_slot_retaining_bytes() { + use arrow::buffer::NullBuffer; + + // The first format physically holds "%d/%m/%Y", but is marked NULL. + let (offsets, values, _) = + GenericStringArray::::from(vec!["%d/%m/%Y"]).into_parts(); + let formats = + GenericStringArray::new(offsets, values, Some(NullBuffer::new_null(1))); + assert!(formats.is_null(0)); + assert_eq!(formats.value(0), "%d/%m/%Y"); + + // Without the validity check, the first format parses this as 2023-02-01 + // and incorrectly wins over the valid second format. + let values = GenericStringArray::::from(vec!["01/02/2023"]); + let fallback_formats = GenericStringArray::::from(vec!["%m/%d/%Y"]); + let res = invoke_to_date_with_args( + vec![ + ColumnarValue::Array(Arc::new(values)), + ColumnarValue::Array(Arc::new(formats)), + ColumnarValue::Array(Arc::new(fallback_formats)), + ], + 1, + ) + .unwrap(); + + let ColumnarValue::Array(res) = res else { + panic!("expected an array result"); + }; + let res = res.as_any().downcast_ref::().unwrap(); + + assert!(!res.is_null(0)); + assert_eq!( + res.value(0), + Date32Type::parse_formatted("01/02/2023", "%m/%d/%Y").unwrap() + ); + } } diff --git a/datafusion/functions/src/datetime/to_time.rs b/datafusion/functions/src/datetime/to_time.rs index 94aa49fbbad2f..f5fe59cbb87b0 100644 --- a/datafusion/functions/src/datetime/to_time.rs +++ b/datafusion/functions/src/datetime/to_time.rs @@ -22,7 +22,7 @@ use arrow::array::types::Time64NanosecondType; use arrow::array::{Array, PrimitiveArray, StringArrayType}; use arrow::datatypes::DataType; use arrow::datatypes::DataType::*; -use chrono::NaiveTime; +use chrono::format::{Item, Parsed, StrftimeItems, parse}; use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -47,7 +47,7 @@ Timestamps will have the time portion extracted. Returns the corresponding time. Note: `to_time` returns Time64(Nanosecond), which represents the time of day in nanoseconds since midnight.", - syntax_example = "to_time('12:30:45', '%H:%M:%S')", + syntax_example = "to_time(expression[, format1, ..., format_n])", sql_example = r#"```sql > select to_time('12:30:45'); +---------------------------+ @@ -141,6 +141,7 @@ impl ScalarUDFImpl for ToTimeFunc { /// Convert string arguments to time (standalone function, not a method on ToTimeFunc) fn string_to_time(args: &[ColumnarValue]) -> Result { let formats = collect_formats(args)?; + let formats = compile_formats(&formats); match &args[0] { ColumnarValue::Scalar(ScalarValue::Utf8(s)) @@ -207,10 +208,25 @@ fn timestamp_to_time(arg: &ColumnarValue) -> Result { arg.cast_to(&Time64(arrow::datatypes::TimeUnit::Nanosecond), None) } +struct CompiledTimeFormat<'a> { + source: &'a str, + items: Vec>, +} + +fn compile_formats<'a>(formats: &[&'a str]) -> Vec> { + formats + .iter() + .map(|source| CompiledTimeFormat { + source, + items: StrftimeItems::new(source).collect(), + }) + .collect() +} + /// Parse time array using the provided formats fn parse_time_array<'a, A: StringArrayType<'a>>( array: &A, - formats: &[&str], + formats: &[CompiledTimeFormat<'_>], ) -> Result> { let mut values = Vec::with_capacity(array.len()); for i in 0..array.len() { @@ -224,10 +240,12 @@ fn parse_time_array<'a, A: StringArrayType<'a>>( } /// Parse time string using provided formats -fn parse_time_with_formats(s: &str, formats: &[&str]) -> Result { +fn parse_time_with_formats(s: &str, formats: &[CompiledTimeFormat<'_>]) -> Result { for format in formats { - if let Ok(time) = NaiveTime::parse_from_str(s, format) { - // Use Arrow's time_to_time64ns function instead of custom implementation + let mut parsed = Parsed::new(); + if parse(&mut parsed, s, format.items.iter()).is_ok() + && let Ok(time) = parsed.to_naive_time() + { return Ok(time_to_time64ns(time)); } } @@ -235,5 +253,8 @@ fn parse_time_with_formats(s: &str, formats: &[&str]) -> Result { "Error parsing '{}' as time. Tried formats: {:?}", s, formats + .iter() + .map(|format| format.source) + .collect::>() ) } diff --git a/datafusion/functions/src/datetime/to_timestamp.rs b/datafusion/functions/src/datetime/to_timestamp.rs index 405f6ff3c7b13..1b45910f7261c 100644 --- a/datafusion/functions/src/datetime/to_timestamp.rs +++ b/datafusion/functions/src/datetime/to_timestamp.rs @@ -26,11 +26,11 @@ use arrow::array::{ use arrow::datatypes::DataType::*; use arrow::datatypes::TimeUnit::{Microsecond, Millisecond, Nanosecond, Second}; use arrow::datatypes::{ - ArrowTimestampType, DataType, TimestampMicrosecondType, TimestampMillisecondType, - TimestampNanosecondType, TimestampSecondType, + ArrowTimestampType, DECIMAL128_MAX_PRECISION, DataType, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, }; use datafusion_common::config::ConfigOptions; -use datafusion_common::{Result, ScalarType, ScalarValue, exec_err}; +use datafusion_common::{Result, ScalarType, ScalarValue, exec_datafusion_err, exec_err}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, @@ -81,7 +81,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo description = r#" Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully -parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is +parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. +Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. "# ) @@ -131,7 +132,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo description = r#" Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully -parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is +parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. +Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. "# ) @@ -181,7 +183,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo description = r#" Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully -parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is +parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. +Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. "# ) @@ -231,7 +234,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo description = r#" Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully -parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is +parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. +Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. "# ) @@ -280,7 +284,8 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo description = r#" Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully -parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is +parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. +Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. "# ) @@ -332,14 +337,31 @@ impl_to_timestamp_constructors!(ToTimestampMillisFunc); impl_to_timestamp_constructors!(ToTimestampMicrosFunc); impl_to_timestamp_constructors!(ToTimestampNanosFunc); -fn decimal_to_nanoseconds(value: i128, scale: i8) -> i64 { +fn decimal_to_nanoseconds(value: i128, scale: i8) -> Result { let nanos_exponent = 9_i16 - scale as i16; + let power = 10_i128 + .checked_pow(nanos_exponent.unsigned_abs() as u32) + .ok_or_else(|| { + exec_datafusion_err!( + "Decimal value {value} with scale {scale} overflows timestamp nanoseconds" + ) + })?; + let timestamp_nanos = if nanos_exponent >= 0 { - value * 10_i128.pow(nanos_exponent as u32) + value.checked_mul(power).ok_or_else(|| { + exec_datafusion_err!( + "Decimal value {value} with scale {scale} overflows timestamp nanoseconds" + ) + })? } else { - value / 10_i128.pow(nanos_exponent.unsigned_abs() as u32) + value / power }; - timestamp_nanos as i64 + + i64::try_from(timestamp_nanos).map_err(|_| { + exec_datafusion_err!( + "Decimal value {value} with scale {scale} overflows timestamp nanoseconds" + ) + }) } fn decimal128_to_timestamp_nanos( @@ -348,7 +370,7 @@ fn decimal128_to_timestamp_nanos( ) -> Result { match arg { ColumnarValue::Scalar(ScalarValue::Decimal128(Some(value), _, scale)) => { - let timestamp_nanos = decimal_to_nanoseconds(*value, *scale); + let timestamp_nanos = decimal_to_nanoseconds(*value, *scale)?; Ok(ColumnarValue::Scalar(ScalarValue::TimestampNanosecond( Some(timestamp_nanos), tz, @@ -362,8 +384,8 @@ fn decimal128_to_timestamp_nanos( let scale = decimal_arr.scale(); let result: TimestampNanosecondArray = decimal_arr .iter() - .map(|v| v.map(|val| decimal_to_nanoseconds(val, scale))) - .collect(); + .map(|v| v.map(|val| decimal_to_nanoseconds(val, scale)).transpose()) + .collect::>()?; let result = result.with_timezone_opt(tz); Ok(ColumnarValue::Array(Arc::new(result))) } @@ -474,7 +496,8 @@ impl ScalarUDFImpl for ToTimestampFunc { _ => exec_err!("Invalid Float64 value for to_timestamp"), }, Decimal32(_, _) | Decimal64(_, _) | Decimal256(_, _) => { - let arg = args[0].cast_to(&Decimal128(38, 9), None)?; + let arg = + args[0].cast_to(&Decimal128(DECIMAL128_MAX_PRECISION, 9), None)?; decimal128_to_timestamp_nanos(&arg, tz) } Decimal128(_, _) => decimal128_to_timestamp_nanos(&args[0], tz), @@ -947,6 +970,37 @@ mod tests { Ok(()) } + #[test] + fn to_timestamp_decimal128_overflow_returns_error() { + let value = "99999999999999999999999999999999999999" + .parse::() + .unwrap(); + let err = decimal128_to_timestamp_nanos( + &ColumnarValue::Scalar(ScalarValue::Decimal128(Some(value), 38, 0)), + None, + ) + .unwrap_err() + .to_string(); + + assert_contains!(err, "overflows timestamp nanoseconds"); + } + + #[test] + fn to_timestamp_decimal128_array_overflow_returns_error() { + let value = "99999999999999999999999999999999999999" + .parse::() + .unwrap(); + let array = Decimal128Array::from(vec![Some(value)]) + .with_precision_and_scale(38, 0) + .unwrap(); + let err = + decimal128_to_timestamp_nanos(&ColumnarValue::Array(Arc::new(array)), None) + .unwrap_err() + .to_string(); + + assert_contains!(err, "overflows timestamp nanoseconds"); + } + #[test] fn to_timestamp_with_formats_arrays_and_nulls() -> Result<()> { // ensure that arrow array implementation is wired up and handles nulls correctly @@ -1830,19 +1884,19 @@ mod tests { #[test] fn test_decimal_to_nanoseconds_negative_scale() { // scale -2: internal value 5 represents 5 * 10^2 = 500 seconds - let nanos = decimal_to_nanoseconds(5, -2); + let nanos = decimal_to_nanoseconds(5, -2).unwrap(); assert_eq!(nanos, 500_000_000_000); // 500 seconds in nanoseconds // scale -1: internal value 10 represents 10 * 10^1 = 100 seconds - let nanos = decimal_to_nanoseconds(10, -1); + let nanos = decimal_to_nanoseconds(10, -1).unwrap(); assert_eq!(nanos, 100_000_000_000); // scale 0: internal value 5 represents 5 seconds - let nanos = decimal_to_nanoseconds(5, 0); + let nanos = decimal_to_nanoseconds(5, 0).unwrap(); assert_eq!(nanos, 5_000_000_000); // scale 3: internal value 1500 represents 1.5 seconds - let nanos = decimal_to_nanoseconds(1500, 3); + let nanos = decimal_to_nanoseconds(1500, 3).unwrap(); assert_eq!(nanos, 1_500_000_000); } } diff --git a/datafusion/functions/src/datetime/to_unixtime.rs b/datafusion/functions/src/datetime/to_unixtime.rs index 9fcfd254ca74d..5b9734c05d7be 100644 --- a/datafusion/functions/src/datetime/to_unixtime.rs +++ b/datafusion/functions/src/datetime/to_unixtime.rs @@ -56,7 +56,7 @@ Integers, unsigned integers, and floats are interpreted as seconds since the uni ), argument( name = "format_n", - description = "Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully parse the expression an error will be returned." + description = "Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL." ) )] #[derive(Debug, PartialEq, Eq, Hash)] diff --git a/datafusion/functions/src/encoding/inner.rs b/datafusion/functions/src/encoding/inner.rs index ad156f735b33b..850e312abdb40 100644 --- a/datafusion/functions/src/encoding/inner.rs +++ b/datafusion/functions/src/encoding/inner.rs @@ -24,7 +24,7 @@ use arrow::{ }, datatypes::DataType, }; -use arrow_buffer::{Buffer, OffsetBufferBuilder}; +use arrow_buffer::{Buffer, OffsetBuffer}; use base64::{ Engine as _, engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig}, @@ -33,7 +33,10 @@ use datafusion_common::{ DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err, internal_err, not_impl_err, plan_err, types::{NativeType, logical_string}, - utils::take_function_args, + utils::{ + hex::{HexCase, encode_bytes as encode_hex, encode_bytes_to_slice}, + take_function_args, + }, }; use datafusion_expr::{ Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -292,9 +295,10 @@ fn decode_array(array: &ArrayRef, encoding: Encoding) -> Result { } DataType::BinaryView => { let array = array.as_binary_view(); - // Don't know if there is a more strict upper bound we can infer - // for view arrays byte data size. - encoding.decode_array::<_, i32>(&array, array.get_buffer_memory_size()) + encoding.decode_array::<_, i32>( + &array, + array.lengths().map(|l| l as usize).sum::(), + ) } DataType::LargeBinary => { let array = array.as_binary::(); @@ -369,7 +373,7 @@ impl Encoding { match self { Self::Base64 => BASE64_ENGINE.encode(value), Self::Base64Padded => BASE64_ENGINE_PADDED.encode(value), - Self::Hex => hex::encode(value), + Self::Hex => encode_hex(value, HexCase::Lower), } } @@ -410,11 +414,7 @@ impl Encoding { .collect(); Ok(Arc::new(array)) } - Self::Hex => { - let array: GenericStringArray = - array.iter().map(|x| x.map(hex::encode)).collect(); - Ok(Arc::new(array)) - } + Self::Hex => hex_encode_array::<_, OutputOffset>(array), } } @@ -459,6 +459,41 @@ impl Encoding { } } +/// Hex-encode a binary array into a string array, writing the lowercase hex +/// digits directly into a single pre-sized value buffer. Each input byte maps +/// to exactly two hex characters, so the output size is known up front and no +/// per-element `String` is allocated. +fn hex_encode_array<'a, InputBinaryArray, OutputOffset>( + array: &InputBinaryArray, +) -> Result +where + InputBinaryArray: BinaryArrayType<'a>, + OutputOffset: OffsetSizeTrait, +{ + let total_input_bytes: usize = array.iter().flatten().map(|v| v.len()).sum(); + + let mut values = vec![0u8; total_input_bytes * 2]; + let mut offsets = Vec::::with_capacity(array.len() + 1); + offsets.push(OutputOffset::zero()); + + let mut pos = 0usize; + for v in array.iter() { + if let Some(v) = v { + let out_len = v.len() * 2; + encode_bytes_to_slice(v, HexCase::Lower, &mut values[pos..pos + out_len])?; + pos += out_len; + } + offsets.push(OutputOffset::usize_as(pos)); + } + + let array = GenericStringArray::::try_new( + OffsetBuffer::new(offsets.into()), + Buffer::from_vec(values), + array.nulls().cloned(), + )?; + Ok(Arc::new(array)) +} + fn delegated_decode<'a, DecodeFunction, InputBinaryArray, OutputOffset>( decode: DecodeFunction, input: &InputBinaryArray, @@ -470,22 +505,21 @@ where OutputOffset: OffsetSizeTrait, { let mut values = vec![0; conservative_upper_bound_size]; - let mut offsets = OffsetBufferBuilder::new(input.len()); + let mut offsets = Vec::::with_capacity(input.len() + 1); + offsets.push(OutputOffset::zero()); let mut total_bytes_decoded = 0; for v in input.iter() { if let Some(v) = v { let cursor = &mut values[total_bytes_decoded..]; let decoded = decode(v, cursor)?; total_bytes_decoded += decoded; - offsets.push_length(decoded); - } else { - offsets.push_length(0); } + offsets.push(OutputOffset::usize_as(total_bytes_decoded)); } // We reserved an upper bound size for the values buffer, but we only use the actual size values.truncate(total_bytes_decoded); let binary_array = GenericBinaryArray::::try_new( - offsets.finish(), + OffsetBuffer::new(offsets.into()), Buffer::from_vec(values), input.nulls().cloned(), )?; @@ -494,7 +528,7 @@ where #[cfg(test)] mod tests { - use arrow::array::BinaryArray; + use arrow::array::{ArrayBuilder, BinaryArray, BinaryViewBuilder}; use arrow_buffer::OffsetBuffer; use super::*; @@ -519,4 +553,14 @@ mod tests { let size = estimate_byte_data_size(&array); assert_eq!(size, 31); } + + #[test] + fn test_estimate_view_size() { + let mut builder = BinaryViewBuilder::new().with_deduplicate_strings(); + for _ in 0..1000 { + builder.append_value([65u8; 64]); + } + let arr = ArrayBuilder::finish(&mut builder); + decode_array(&arr, Encoding::Base64).unwrap(); + } } diff --git a/datafusion/functions/src/lib.rs b/datafusion/functions/src/lib.rs index 7e753d7f35eb3..14d1743770883 100644 --- a/datafusion/functions/src/lib.rs +++ b/datafusion/functions/src/lib.rs @@ -141,6 +141,7 @@ make_stub_package!(unicode, "unicode_expressions"); #[cfg(any(feature = "datetime_expressions", feature = "unicode_expressions"))] pub mod planner; +pub mod binaries; pub mod strings; pub mod utils; diff --git a/datafusion/functions/src/macros.rs b/datafusion/functions/src/macros.rs index 71528b4d16bf0..8a6607c46b45e 100644 --- a/datafusion/functions/src/macros.rs +++ b/datafusion/functions/src/macros.rs @@ -207,9 +207,22 @@ macro_rules! downcast_arg { /// $NAME: the name of the function /// $UNARY_FUNC: the unary function to apply to the argument /// $OUTPUT_ORDERING: the output ordering calculation method of the function +/// $STRICT: whether the function returns NULL when any argument is NULL /// $GET_DOC: the function to get the documentation of the UDF macro_rules! make_math_unary_udf { - ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $GET_DOC:expr) => { + ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $STRICT:expr, $GET_DOC:expr) => { + make_math_unary_udf!( + $UDF, + $NAME, + $UNARY_FUNC, + $OUTPUT_ORDERING, + $EVALUATE_BOUNDS, + $STRICT, + $GET_DOC, + None:: Result<()>> + ); + }; + ($UDF:ident, $NAME:ident, $UNARY_FUNC:ident, $OUTPUT_ORDERING:expr, $EVALUATE_BOUNDS:expr, $STRICT:expr, $GET_DOC:expr, $VALIDATOR:expr) => { $crate::make_udf_function!($NAME::$UDF, $NAME); mod $NAME { @@ -218,6 +231,7 @@ macro_rules! make_math_unary_udf { use arrow::array::{ArrayRef, AsArray}; use arrow::datatypes::{DataType, Float32Type, Float64Type}; + use arrow::error::ArrowError; use datafusion_common::{Result, exec_err}; use datafusion_expr::interval_arithmetic::Interval; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; @@ -233,11 +247,10 @@ macro_rules! make_math_unary_udf { impl $UDF { pub fn new() -> Self { - use DataType::*; Self { signature: Signature::uniform( 1, - vec![Float64, Float32], + vec![DataType::Float64, DataType::Float32], Volatility::Immutable, ), } @@ -258,11 +271,14 @@ macro_rules! make_math_unary_udf { match arg_type { DataType::Float32 => Ok(DataType::Float32), - // For other types (possible values float64/null/int), use Float64 _ => Ok(DataType::Float64), } } + fn is_strict(&self) -> bool { + $STRICT + } + fn output_ordering( &self, input: &[ExprProperties], @@ -280,16 +296,38 @@ macro_rules! make_math_unary_udf { ) -> Result { let args = ColumnarValue::values_to_arrays(&args.args)?; let arr: ArrayRef = match args[0].data_type() { - DataType::Float64 => Arc::new( - args[0] + DataType::Float64 => { + let values = args[0] .as_primitive::() - .unary::<_, Float64Type>(|x: f64| f64::$UNARY_FUNC(x)), - ) as ArrayRef, - DataType::Float32 => Arc::new( - args[0] + .try_unary::<_, Float64Type, _>( + |x: f64| -> std::result::Result { + if let Some(validate) = $VALIDATOR { + validate(x).map_err(|error| { + ArrowError::ComputeError(error.to_string()) + })?; + } + + Ok(f64::$UNARY_FUNC(x)) + }, + )?; + Arc::new(values) as ArrayRef + } + DataType::Float32 => { + let values = args[0] .as_primitive::() - .unary::<_, Float32Type>(|x: f32| f32::$UNARY_FUNC(x)), - ) as ArrayRef, + .try_unary::<_, Float32Type, _>( + |x: f32| -> std::result::Result { + if let Some(validate) = $VALIDATOR { + validate(x as f64).map_err(|error| { + ArrowError::ComputeError(error.to_string()) + })?; + } + + Ok(f32::$UNARY_FUNC(x)) + }, + )?; + Arc::new(values) as ArrayRef + } other => { return exec_err!( "Unsupported data type {other:?} for function {}", @@ -311,16 +349,21 @@ macro_rules! make_math_unary_udf { /// Macro to create a binary math UDF. /// -/// A binary math function takes two arguments of types Float32 or Float64, -/// applies a binary floating function to the argument, and returns a value of the same type. +/// A binary math function takes two numeric arguments. When both arguments are +/// Float32 the function is evaluated in single precision and returns Float32. +/// Any other combination of numeric (or null) argument types is coerced to +/// Float64 and returns Float64; in particular integers are widened to Float64 +/// rather than Float32 so that values needing more than 24 bits of mantissa are +/// not silently rounded. /// /// $UDF: the name of the UDF struct that implements `ScalarUDFImpl` /// $NAME: the name of the function /// $BINARY_FUNC: the binary function to apply to the argument /// $OUTPUT_ORDERING: the output ordering calculation method of the function +/// $STRICT: whether the function returns NULL when any argument is NULL /// $GET_DOC: the function to get the documentation of the UDF macro_rules! make_math_binary_udf { - ($UDF:ident, $NAME:ident, $BINARY_FUNC:ident, $OUTPUT_ORDERING:expr, $GET_DOC:expr) => { + ($UDF:ident, $NAME:ident, $BINARY_FUNC:ident, $OUTPUT_ORDERING:expr, $STRICT:expr, $GET_DOC:expr) => { $crate::make_udf_function!($NAME::$UDF, $NAME); mod $NAME { @@ -331,7 +374,6 @@ macro_rules! make_math_binary_udf { use arrow::datatypes::{DataType, Float32Type, Float64Type}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue, internal_err}; - use datafusion_expr::TypeSignature; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, @@ -345,13 +387,18 @@ macro_rules! make_math_binary_udf { impl $UDF { pub fn new() -> Self { - use DataType::*; Self { - signature: Signature::one_of( - vec![ - TypeSignature::Exact(vec![Float32, Float32]), - TypeSignature::Exact(vec![Float64, Float64]), - ], + // Float64 is listed first so that integer (and other + // non-float) arguments coerce to Float64 rather than + // Float32; genuine Float32 arguments still match + // exactly and stay in single precision. Coercing + // integers to Float64 matters for correctness: Float32 + // has only a 24-bit mantissa, so widening a large + // integer to Float32 would round it before the function + // is ever applied. + signature: Signature::uniform( + 2, + vec![DataType::Float64, DataType::Float32], Volatility::Immutable, ), } @@ -368,15 +415,16 @@ macro_rules! make_math_binary_udf { } fn return_type(&self, arg_types: &[DataType]) -> Result { - let arg_type = &arg_types[0]; - - match arg_type { - DataType::Float32 => Ok(DataType::Float32), - // For other types (possible values float64/null/int), use Float64 + match (&arg_types[0], &arg_types[1]) { + (DataType::Float32, DataType::Float32) => Ok(DataType::Float32), _ => Ok(DataType::Float64), } } + fn is_strict(&self) -> bool { + $STRICT + } + fn output_ordering( &self, input: &[ExprProperties], diff --git a/datafusion/functions/src/math/abs.rs b/datafusion/functions/src/math/abs.rs index 02ac89756d919..5c5c24a1a65f5 100644 --- a/datafusion/functions/src/math/abs.rs +++ b/datafusion/functions/src/math/abs.rs @@ -158,6 +158,10 @@ impl ScalarUDFImpl for AbsFunc { Ok(arg_types[0].clone()) } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let args = ColumnarValue::values_to_arrays(&args.args)?; let [input] = take_function_args(self.name(), args)?; diff --git a/datafusion/functions/src/math/ceil.rs b/datafusion/functions/src/math/ceil.rs index 395cb4eae03f5..7b2c0c35e4cad 100644 --- a/datafusion/functions/src/math/ceil.rs +++ b/datafusion/functions/src/math/ceil.rs @@ -89,6 +89,10 @@ impl ScalarUDFImpl for CeilFunc { } } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let arg = &args.args[0]; diff --git a/datafusion/functions/src/math/common.rs b/datafusion/functions/src/math/common.rs new file mode 100644 index 0000000000000..9bb6f6fe1e35c --- /dev/null +++ b/datafusion/functions/src/math/common.rs @@ -0,0 +1,320 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::ArrowNativeTypeOp; +use arrow::error::ArrowError; +use num_traits::{CheckedMul, CheckedNeg, Signed}; +use std::fmt::Display; +use std::mem::swap; +use std::ops::RemAssign; + +/// A gcd helper to compute GCD using Euclidean GCD algorithm +/// on non-negative numbers (scalars and decimals) +fn gcd_helper(a: T, b: T) -> Result +where + T: ArrowNativeTypeOp + RemAssign + CheckedNeg, +{ + debug_assert!(a >= T::ZERO); + debug_assert!(b >= T::ZERO); + let (mut a, mut b) = if a > b { (a, b) } else { (b, a) }; + + while b != T::ZERO { + swap(&mut a, &mut b); + b %= a; + } + + Ok(a) +} + +/// Computes gcd of two unsigned integers using Binary GCD algorithm +/// Faster, works with integers only +pub(crate) fn unsigned_gcd(mut a: u64, mut b: u64) -> u64 { + if a == 0 { + return b; + } + if b == 0 { + return a; + } + + let shift = (a | b).trailing_zeros(); + a >>= a.trailing_zeros(); + loop { + b >>= b.trailing_zeros(); + if a > b { + swap(&mut a, &mut b); + } + b -= a; + if b == 0 { + return a << shift; + } + } +} + +/// Computes gcd of two signed numbers (integers or decimals), +/// checking for output integer overflow +pub(crate) fn gcd_signed(x: T, y: T) -> Result +where + T: ArrowNativeTypeOp + RemAssign + Signed + CheckedNeg, +{ + // Make absolute values, keeping type + let a = if x.is_positive() { + x + } else { + x.checked_neg() + .ok_or_else(|| ArrowError::ComputeError("Signed integer overflow".into()))? + }; + let b = if y.is_positive() { + y + } else { + y.checked_neg() + .ok_or_else(|| ArrowError::ComputeError("Signed integer overflow".into()))? + }; + // Call with signed numbers + gcd_helper(a, b) +} + +/// Computes gcd of two signed integers +pub(crate) fn gcd_signed_int(x: i64, y: i64) -> Result { + let a = x.unsigned_abs(); + let b = y.unsigned_abs(); + + // Call with unsigned numbers + let r = unsigned_gcd(a, b); + // gcd(i64::MIN, i64::MIN) = u64::MIN.unsigned_abs() cannot fit into i64 + r.try_into().map_err(|_| { + ArrowError::ComputeError(format!("Signed integer overflow in GCD({x}, {y})")) + }) +} + +/// Computes lcm of two signed numbers (integers or decimals) +pub(crate) fn lcm_signed(x: T, y: T) -> Result +where + T: ArrowNativeTypeOp + RemAssign + Signed + CheckedNeg + CheckedMul + Display, +{ + if x == T::ZERO || y == T::ZERO { + return Ok(T::ZERO); + } + + // Make absolute values, keeping type + let a = if x.is_positive() { + x + } else { + x.checked_neg() + .ok_or_else(|| ArrowError::ComputeError("Signed integer overflow".into()))? + }; + let b = if y.is_positive() { + y + } else { + y.checked_neg() + .ok_or_else(|| ArrowError::ComputeError("Signed integer overflow".into()))? + }; + // Call with signed numbers + let gcd = gcd_helper(a, b)?; + // gcd is not zero since both a and b are not zero, so the division is safe. + (a / gcd).checked_mul(&b).ok_or_else(|| { + ArrowError::ComputeError(format!("Signed integer overflow in LCM({x}, {y})")) + }) +} + +/// Computes lcm of two signed integers, +/// checking for output integer overflow +pub(crate) fn lcm_signed_int(x: i64, y: i64) -> Result { + if x == 0 || y == 0 { + return Ok(0); + } + + let a = x.unsigned_abs(); + let b = y.unsigned_abs(); + + let gcd = gcd_helper::(a, b)?; + // gcd is not zero since both a and b are not zero, so the division is safe. + (a / gcd) + .checked_mul(b) + .and_then(|v| i64::try_from(v).ok()) + .ok_or_else(|| { + ArrowError::ComputeError(format!("Signed integer overflow in LCM({x}, {y})")) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow_buffer::i256; + + const GCD_COMMON_TEST_CASES: [(i64, i64, i64); 18] = [ + // Basic cases + (48, 18, 6), + (54, 24, 6), + (100, 50, 50), + (17, 19, 1), + (21, 14, 7), + // Edge cases with 0 + (0, 0, 0), + (0, 5, 5), + (10, 0, 10), + // Same numbers + (7, 7, 7), + (100, 100, 100), + // One is 1 + (1, 1, 1), + (1, 100, 1), + (999, 1, 1), + // Large numbers + (1000000, 500000, 500000), + (123456, 789012, 12), + (999999, 111111, 111111), + // Powers of 2 + (64, 128, 64), + (1024, 2048, 1024), + ]; + + const LCM_COMMON_TEST_CASES: [(i64, i64, i64); 18] = [ + // Basic cases + (48, 18, 144), + (54, 24, 216), + (100, 50, 100), + (17, 19, 323), + (21, 14, 42), + // Edge cases with 0 + (0, 0, 0), + (0, 5, 0), + (10, 0, 0), + // Same numbers + (7, 7, 7), + (100, 100, 100), + // One is 1 + (1, 1, 1), + (1, 100, 100), + (999, 1, 999), + // Large numbers + (1_000_000, 500_000, 1_000_000), + (123_456, 789_012, 8_117_355_456), + (999_999, 111_111, 999_999), + // Powers of 2 + (64, 128, 128), + (1024, 2048, 2048), + ]; + + #[test] + fn test_gcd_i64() { + let test_cases: Vec<(i64, i64, i64)> = [ + GCD_COMMON_TEST_CASES.into(), + vec![ + // Max value cases + (1, i64::MAX, 1), + (i64::MAX, 1, 1), + (i64::MAX, i64::MAX, i64::MAX), + ], + ] + .concat(); + + // Success cases + for (a, b, expected) in test_cases { + let actual_euclidean = gcd_signed(a, b).expect("should succeed"); + assert_eq!( + actual_euclidean, expected, + "gcd_signed({a}, {b}) expected {expected}, actual {actual_euclidean}" + ); + let actual_binary: i64 = + unsigned_gcd(a.try_into().unwrap(), b.try_into().unwrap()) + .try_into() + .expect("overflow"); + assert_eq!( + actual_binary, expected, + "unsigned_gcd({a}, {b}) expected {expected}, actual {actual_binary}" + ); + } + } + + #[test] + fn test_gcd_decimal() { + let test_cases: Vec<(i256, i256, i256)> = [ + GCD_COMMON_TEST_CASES + .iter() + .map(|&(a, b, c)| (i256::from(a), i256::from(b), i256::from(c))) + .collect(), + vec![ + (i256::from(1), i256::MAX, i256::from(1)), + (i256::MAX, i256::from(1), i256::from(1)), + (i256::MAX, i256::MAX, i256::MAX), + ], + ] + .concat(); + + // Success cases + for (a, b, expected) in test_cases { + let actual = gcd_signed(a, b).expect("should succeed"); + assert_eq!( + actual, expected, + "euclid_gcd({a}, {b}) expected {expected}, actual {actual}" + ); + } + } + + #[test] + fn test_lcm_i64() { + let test_cases: Vec<(i64, i64, i64)> = [ + LCM_COMMON_TEST_CASES.into(), + vec![ + // Negative inputs - LCM is always non-negative + (-6, 4, 12), + (-4, -6, 12), + // Max value cases + (1, i64::MAX, i64::MAX), + (i64::MAX, 1, i64::MAX), + (i64::MAX, i64::MAX, i64::MAX), + ], + ] + .concat(); + + for (a, b, expected) in test_cases { + let actual = lcm_signed_int(a, b).expect("should succeed"); + assert_eq!( + actual, expected, + "lcm_signed_int({a}, {b}) expected {expected}, actual {actual}" + ); + } + } + + #[test] + fn test_lcm_decimal() { + let test_cases: Vec<(i256, i256, i256)> = [ + LCM_COMMON_TEST_CASES + .iter() + .map(|&(a, b, c)| (i256::from(a), i256::from(b), i256::from(c))) + .collect(), + vec![ + // Negative inputs - LCM is always non-negative + (i256::from(-6_i64), i256::from(4_i64), i256::from(12_i64)), + (i256::from(-4_i64), i256::from(-6_i64), i256::from(12_i64)), + // Max value cases + (i256::from(1_i64), i256::MAX, i256::MAX), + (i256::MAX, i256::from(1_i64), i256::MAX), + (i256::MAX, i256::MAX, i256::MAX), + ], + ] + .concat(); + + for (a, b, expected) in test_cases { + let actual = lcm_signed(a, b).expect("should succeed"); + assert_eq!( + actual, expected, + "lcm_signed({a}, {b}) expected {expected}, actual {actual}" + ); + } + } +} diff --git a/datafusion/functions/src/math/cot.rs b/datafusion/functions/src/math/cot.rs index 24f0a412e3a8a..ca207778f7eb7 100644 --- a/datafusion/functions/src/math/cot.rs +++ b/datafusion/functions/src/math/cot.rs @@ -86,6 +86,10 @@ impl ScalarUDFImpl for CotFunc { } } + fn is_strict(&self) -> bool { + true + } + fn documentation(&self) -> Option<&Documentation> { self.doc() } diff --git a/datafusion/functions/src/math/factorial.rs b/datafusion/functions/src/math/factorial.rs index 3b4f973f19d62..f4e9b60dd3799 100644 --- a/datafusion/functions/src/math/factorial.rs +++ b/datafusion/functions/src/math/factorial.rs @@ -76,6 +76,10 @@ impl ScalarUDFImpl for FactorialFunc { Ok(Int64) } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [arg] = take_function_args(self.name(), args.args)?; diff --git a/datafusion/functions/src/math/floor.rs b/datafusion/functions/src/math/floor.rs index e02aa141c5b71..4ab6e0eb5effd 100644 --- a/datafusion/functions/src/math/floor.rs +++ b/datafusion/functions/src/math/floor.rs @@ -129,6 +129,10 @@ impl ScalarUDFImpl for FloorFunc { } } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let arg = &args.args[0]; diff --git a/datafusion/functions/src/math/gcd.rs b/datafusion/functions/src/math/gcd.rs index 8b92c454d9b4c..6a4e69620e060 100644 --- a/datafusion/functions/src/math/gcd.rs +++ b/datafusion/functions/src/math/gcd.rs @@ -17,16 +17,22 @@ use arrow::array::{ArrayRef, AsArray, PrimitiveArray}; use arrow::compute::try_binary; -use arrow::datatypes::{DataType, Int64Type}; -use arrow::error::ArrowError; -use std::mem::swap; +use arrow::datatypes::{ + DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Int64Type, +}; use std::sync::Arc; -use datafusion_common::{Result, ScalarValue, exec_err, internal_datafusion_err}; +use crate::math::common::{gcd_signed, gcd_signed_int, unsigned_gcd}; +use crate::utils::calculate_binary_decimal_math_cast; +use datafusion_common::utils::take_function_args; +use datafusion_common::{ + Result, ScalarValue, exec_err, internal_datafusion_err, plan_err, +}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; +use datafusion_expr_common::type_coercion::binary::decimal_coercion; use datafusion_macros::user_doc; #[user_doc( @@ -58,11 +64,7 @@ impl Default for GcdFunc { impl GcdFunc { pub fn new() -> Self { Self { - signature: Signature::uniform( - 2, - vec![DataType::Int64], - Volatility::Immutable, - ), + signature: Signature::user_defined(Volatility::Immutable), } } } @@ -76,37 +78,127 @@ impl ScalarUDFImpl for GcdFunc { &self.signature } - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Int64) + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) + } + + fn is_strict(&self) -> bool { + true + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [arg1, arg2] = take_function_args(self.name(), arg_types)?; + + let coerced_type = match (arg1, arg2) { + (DataType::Null, _) | (_, DataType::Null) => Ok(DataType::Int64), + (lhs, rhs) if lhs.is_integer() && rhs.is_integer() => Ok(DataType::Int64), + (lhs, rhs) if lhs.is_decimal() || rhs.is_decimal() => { + decimal_coercion(lhs, rhs).map(Ok).unwrap_or_else(|| { + plan_err!( + "Unsupported argument types {lhs:?} and {rhs:?} for function {}", + self.name() + ) + }) + } + (lhs, rhs) => { + plan_err!( + "Unsupported argument types {lhs:?} and {rhs:?} for function {}", + self.name() + ) + } + }?; + Ok(vec![coerced_type.clone(), coerced_type]) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let number_rows = args.number_rows; let args: [ColumnarValue; 2] = args.args.try_into().map_err(|_| { internal_datafusion_err!("Expected 2 arguments for function gcd") })?; - match args { - [ColumnarValue::Array(a), ColumnarValue::Array(b)] => { - compute_gcd_for_arrays(&a, &b) + if args[0].data_type() == DataType::Int64 { + // Optimized path for both integers + match args { + [ColumnarValue::Array(a), ColumnarValue::Array(b)] => { + compute_gcd_for_arrays(&a, &b) + } + [ + ColumnarValue::Scalar(ScalarValue::Int64(a)), + ColumnarValue::Scalar(ScalarValue::Int64(b)), + ] => match (a, b) { + (Some(a), Some(b)) => Ok(ColumnarValue::Scalar(ScalarValue::Int64( + Some(gcd_signed_int(a, b)?), + ))), + _ => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))), + }, + [ + ColumnarValue::Array(a), + ColumnarValue::Scalar(ScalarValue::Int64(b)), + ] => compute_gcd_with_scalar(&a, b), + [ + ColumnarValue::Scalar(ScalarValue::Int64(a)), + ColumnarValue::Array(b), + ] => compute_gcd_with_scalar(&b, a), + _ => exec_err!("Unsupported argument types for function gcd"), } - [ - ColumnarValue::Scalar(ScalarValue::Int64(a)), - ColumnarValue::Scalar(ScalarValue::Int64(b)), - ] => match (a, b) { - (Some(a), Some(b)) => Ok(ColumnarValue::Scalar(ScalarValue::Int64( - Some(compute_gcd(a, b)?), - ))), - _ => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))), - }, - [ - ColumnarValue::Array(a), - ColumnarValue::Scalar(ScalarValue::Int64(b)), - ] => compute_gcd_with_scalar(&a, b), - [ - ColumnarValue::Scalar(ScalarValue::Int64(a)), - ColumnarValue::Array(b), - ] => compute_gcd_with_scalar(&b, a), - _ => exec_err!("Unsupported argument types for function gcd"), + } else { + // Decimal path: convert left to array and use generic helper + let left = args[0].to_array(number_rows)?; + let right = &args[1]; + + let arr: ArrayRef = match (left.data_type(), right.data_type()) { + ( + lhs @ DataType::Decimal32(precision, scale), + rhs @ DataType::Decimal32(_, _), + ) if *lhs == rhs => calculate_binary_decimal_math_cast::< + Decimal32Type, + Decimal32Type, + Decimal32Type, + _, + >( + &left, right, gcd_signed, *precision, *scale, lhs + )?, + ( + lhs @ DataType::Decimal64(precision, scale), + rhs @ DataType::Decimal64(_, _), + ) if *lhs == rhs => calculate_binary_decimal_math_cast::< + Decimal64Type, + Decimal64Type, + Decimal64Type, + _, + >( + &left, right, gcd_signed, *precision, *scale, lhs + )?, + ( + lhs @ DataType::Decimal128(precision, scale), + rhs @ DataType::Decimal128(_, _), + ) if *lhs == rhs => calculate_binary_decimal_math_cast::< + Decimal128Type, + Decimal128Type, + Decimal128Type, + _, + >( + &left, right, gcd_signed, *precision, *scale, lhs + )?, + ( + lhs @ DataType::Decimal256(precision, scale), + rhs @ DataType::Decimal256(_, _), + ) if *lhs == rhs => calculate_binary_decimal_math_cast::< + Decimal256Type, + Decimal256Type, + Decimal256Type, + _, + >( + &left, right, gcd_signed, *precision, *scale, lhs + )?, + (lhs, rhs) => { + exec_err!( + "Unsupported data types {lhs:?} and {rhs:?} for function {}", + self.name() + ) + }?, + }; + Ok(ColumnarValue::Array(arr)) } } @@ -118,7 +210,7 @@ impl ScalarUDFImpl for GcdFunc { fn compute_gcd_for_arrays(a: &ArrayRef, b: &ArrayRef) -> Result { let a = a.as_primitive::(); let b = b.as_primitive::(); - try_binary(a, b, compute_gcd) + try_binary(a, b, gcd_signed_int) .map(|arr: PrimitiveArray| { ColumnarValue::Array(Arc::new(arr) as ArrayRef) }) @@ -141,44 +233,37 @@ fn compute_gcd_with_scalar(arr: &ArrayRef, scalar: Option) -> Result { let result: PrimitiveArray = - prim.try_unary(|val| compute_gcd(val, scalar_value))?; + prim.try_unary(|val| gcd_signed_int(val, scalar_value))?; Ok(ColumnarValue::Array(Arc::new(result) as ArrayRef)) } None => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))), } } -/// Computes gcd of two unsigned integers using Binary GCD algorithm. -pub(super) fn unsigned_gcd(mut a: u64, mut b: u64) -> u64 { - if a == 0 { - return b; - } - if b == 0 { - return a; - } +#[cfg(test)] +mod tests { + use super::*; - let shift = (a | b).trailing_zeros(); - a >>= a.trailing_zeros(); - loop { - b >>= b.trailing_zeros(); - if a > b { - swap(&mut a, &mut b); - } - b -= a; - if b == 0 { - return a << shift; - } - } -} + #[test] + fn test_coercion() { + let mut coerced = GcdFunc::new() + .coerce_types(&[DataType::Int64, DataType::Int32]) + .expect("coercion should succeed"); + assert_eq!(coerced, vec![DataType::Int64, DataType::Int64]); + + coerced = GcdFunc::new() + .coerce_types(&[DataType::Decimal128(10, 2), DataType::Int32]) + .expect("coercion should succeed"); + + assert_eq!( + coerced, + vec![DataType::Decimal128(12, 2), DataType::Decimal128(12, 2)] + ); -/// Computes greatest common divisor using Binary GCD algorithm. -pub fn compute_gcd(x: i64, y: i64) -> Result { - let a = x.unsigned_abs(); - let b = y.unsigned_abs(); - let r = unsigned_gcd(a, b); - // The result can be up to 2^63 (e.g. gcd(i64::MIN, 0) or - // gcd(i64::MIN, i64::MIN)), which does not fit into i64. - r.try_into().map_err(|_| { - ArrowError::ComputeError(format!("Signed integer overflow in GCD({x}, {y})")) - }) + coerced = GcdFunc::new() + .coerce_types(&[DataType::Decimal128(10, 2), DataType::Null]) + .expect("coercion should succeed"); + + assert_eq!(coerced, vec![DataType::Int64, DataType::Int64]); + } } diff --git a/datafusion/functions/src/math/iszero.rs b/datafusion/functions/src/math/iszero.rs index de6fc669692ee..62cfdd4c839ec 100644 --- a/datafusion/functions/src/math/iszero.rs +++ b/datafusion/functions/src/math/iszero.rs @@ -85,6 +85,10 @@ impl ScalarUDFImpl for IsZeroFunc { Ok(Boolean) } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [arg] = take_function_args(self.name(), args.args)?; diff --git a/datafusion/functions/src/math/lcm.rs b/datafusion/functions/src/math/lcm.rs index 9398e9f8d6e00..248e4b93ffd8b 100644 --- a/datafusion/functions/src/math/lcm.rs +++ b/datafusion/functions/src/math/lcm.rs @@ -15,25 +15,22 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; - -use arrow::array::{ArrayRef, AsArray, PrimitiveArray}; -use arrow::compute::try_binary; -use arrow::datatypes::DataType; -use arrow::datatypes::DataType::Int64; -use arrow::datatypes::Int64Type; +use arrow::array::ArrayRef; +use arrow::datatypes::{ + DataType, Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Int64Type, +}; -use arrow::error::ArrowError; -use datafusion_common::{Result, exec_err}; +use crate::math::common::{lcm_signed, lcm_signed_int}; +use crate::utils::{calculate_binary_decimal_math_cast, calculate_binary_math}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, exec_err, plan_err}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; +use datafusion_expr_common::type_coercion::binary::decimal_coercion; use datafusion_macros::user_doc; -use super::gcd::unsigned_gcd; -use crate::utils::make_scalar_function; - #[user_doc( doc_section(label = "Math Functions"), description = "Returns the least common multiple of `expression_x` and `expression_y`. Returns 0 if either input is zero.", @@ -62,9 +59,8 @@ impl Default for LcmFunc { impl LcmFunc { pub fn new() -> Self { - use DataType::*; Self { - signature: Signature::uniform(2, vec![Int64], Volatility::Immutable), + signature: Signature::user_defined(Volatility::Immutable), } } } @@ -78,49 +74,104 @@ impl ScalarUDFImpl for LcmFunc { &self.signature } - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(Int64) + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) } - fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(lcm, vec![])(&args.args) + fn is_strict(&self) -> bool { + true } - fn documentation(&self) -> Option<&Documentation> { - self.doc() + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + let [arg1, arg2] = take_function_args(self.name(), arg_types)?; + + let coerced_type = match (arg1, arg2) { + (DataType::Null, _) | (_, DataType::Null) => Ok(DataType::Int64), + (lhs, rhs) if lhs.is_integer() && rhs.is_integer() => Ok(DataType::Int64), + (lhs, rhs) if lhs.is_decimal() || rhs.is_decimal() => { + decimal_coercion(lhs, rhs).map(Ok).unwrap_or_else(|| { + plan_err!( + "Unsupported argument types {lhs:?} and {rhs:?} for function {}", + self.name() + ) + }) + } + (lhs, rhs) => { + plan_err!( + "Unsupported argument types {lhs:?} and {rhs:?} for function {}", + self.name() + ) + } + }?; + Ok(vec![coerced_type.clone(), coerced_type]) } -} -/// Lcm SQL function -fn lcm(args: &[ArrayRef]) -> Result { - let compute_lcm = |x: i64, y: i64| -> Result { - if x == 0 || y == 0 { - return Ok(0); - } - - // lcm(x, y) = |x| * |y| / gcd(|x|, |y|) - let a = x.unsigned_abs(); - let b = y.unsigned_abs(); - let gcd = unsigned_gcd(a, b); - // gcd is not zero since both a and b are not zero, so the division is safe. - (a / gcd) - .checked_mul(b) - .and_then(|v| i64::try_from(v).ok()) - .ok_or_else(|| { - ArrowError::ComputeError(format!( - "Signed integer overflow in LCM({x}, {y})" - )) - }) - }; - - match args[0].data_type() { - Int64 => { - let arg1 = args[0].as_primitive::(); - let arg2 = args[1].as_primitive::(); + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let left = &args.args[0].to_array(args.number_rows)?; + let right = &args.args[1]; + + let arr: ArrayRef = match (left.data_type(), right.data_type()) { + (DataType::Int64, _) => calculate_binary_math::< + Int64Type, + Int64Type, + Int64Type, + _, + >(&left, right, lcm_signed_int)?, + ( + lhs @ DataType::Decimal32(precision, scale), + rhs @ DataType::Decimal32(_, _), + ) if *lhs == rhs => { + calculate_binary_decimal_math_cast::< + Decimal32Type, + Decimal32Type, + Decimal32Type, + _, + >(&left, right, lcm_signed, *precision, *scale, lhs)? + } + ( + lhs @ DataType::Decimal64(precision, scale), + rhs @ DataType::Decimal64(_, _), + ) if *lhs == rhs => { + calculate_binary_decimal_math_cast::< + Decimal64Type, + Decimal64Type, + Decimal64Type, + _, + >(&left, right, lcm_signed, *precision, *scale, lhs)? + } + ( + lhs @ DataType::Decimal128(precision, scale), + rhs @ DataType::Decimal128(_, _), + ) if *lhs == rhs => { + calculate_binary_decimal_math_cast::< + Decimal128Type, + Decimal128Type, + Decimal128Type, + _, + >(&left, right, lcm_signed, *precision, *scale, lhs)? + } + ( + lhs @ DataType::Decimal256(precision, scale), + rhs @ DataType::Decimal256(_, _), + ) if *lhs == rhs => { + calculate_binary_decimal_math_cast::< + Decimal256Type, + Decimal256Type, + Decimal256Type, + _, + >(&left, right, lcm_signed, *precision, *scale, lhs)? + } + (lhs, rhs) => { + return exec_err!( + "Unsupported data types {lhs:?} and {rhs:?} for function {}", + self.name() + ); + } + }; + Ok(ColumnarValue::Array(arr)) + } - let result: PrimitiveArray = try_binary(arg1, arg2, compute_lcm)?; - Ok(Arc::new(result) as ArrayRef) - } - other => exec_err!("Unsupported data type {other:?} for function lcm"), + fn documentation(&self) -> Option<&Documentation> { + self.doc() } } diff --git a/datafusion/functions/src/math/log.rs b/datafusion/functions/src/math/log.rs index ac94f78e0c723..732cfff6cf053 100644 --- a/datafusion/functions/src/math/log.rs +++ b/datafusion/functions/src/math/log.rs @@ -203,6 +203,10 @@ impl ScalarUDFImpl for LogFunc { } } + fn is_strict(&self) -> bool { + true + } + fn output_ordering(&self, input: &[ExprProperties]) -> Result { let (base_sort_properties, num_sort_properties) = if input.len() == 1 { // log(x) defaults to log(10, x) @@ -358,23 +362,27 @@ impl ScalarUDFImpl for LogFunc { } else { lit(ScalarValue::new_ten(&number_datatype)?) }; + let base_nullable = info.nullable(&base)?; match number { Expr::Literal(value, _) - if value == ScalarValue::new_one(&number_datatype)? => + if value == ScalarValue::new_one(&number_datatype)? && !base_nullable => { Ok(ExprSimplifyResult::Simplified(lit(ScalarValue::new_zero( &info.get_data_type(&base)?, )?))) } Expr::ScalarFunction(ScalarFunction { func, mut args }) - if is_pow(&func) && args.len() == 2 && base == args[0] => + if is_pow(&func) + && args.len() == 2 + && base == args[0] + && !base_nullable => { let b = args.pop().unwrap(); // length checked above Ok(ExprSimplifyResult::Simplified(b)) } number => { - if number == base { + if number == base && !base_nullable { Ok(ExprSimplifyResult::Simplified(lit(ScalarValue::new_one( &number_datatype, )?))) @@ -1165,7 +1173,12 @@ mod tests { #[test] fn test_log_decimal256_large() { // Large Decimal256 values that don't fit in i128 now use f64 fallback - let arg_field = Field::new("a", DataType::Decimal256(38, 0), false).into(); + let arg_field = Field::new( + "a", + DataType::Decimal256(DECIMAL256_MAX_PRECISION, 0), + false, + ) + .into(); let args = ScalarFunctionArgs { args: vec![ ColumnarValue::Array(Arc::new(Decimal256Array::from(vec![ diff --git a/datafusion/functions/src/math/mod.rs b/datafusion/functions/src/math/mod.rs index 610e773d68fd0..4b79866895d84 100644 --- a/datafusion/functions/src/math/mod.rs +++ b/datafusion/functions/src/math/mod.rs @@ -18,12 +18,14 @@ //! "math" DataFusion functions use crate::math::monotonicity::*; +use datafusion_common::{Result, exec_err}; use datafusion_expr::ScalarUDF; use std::sync::Arc; pub mod abs; pub mod bounds; pub mod ceil; +mod common; pub mod cot; mod decimal; pub mod factorial; @@ -42,6 +44,14 @@ pub mod round; pub mod signum; pub mod trunc; +fn validate_sqrt_input(value: f64) -> Result<()> { + if value < 0.0 { + exec_err!("cannot take square root of a negative number") + } else { + Ok(()) + } +} + // Create UDFs make_udf_function!(abs::AbsFunc, abs); make_math_unary_udf!( @@ -50,6 +60,7 @@ make_math_unary_udf!( acos, super::acos_order, super::bounds::acos_bounds, + true, super::get_acos_doc ); make_math_unary_udf!( @@ -58,6 +69,7 @@ make_math_unary_udf!( acosh, super::acosh_order, super::bounds::acosh_bounds, + true, super::get_acosh_doc ); make_math_unary_udf!( @@ -66,6 +78,7 @@ make_math_unary_udf!( asin, super::asin_order, super::bounds::asin_bounds, + true, super::get_asin_doc ); make_math_unary_udf!( @@ -74,6 +87,7 @@ make_math_unary_udf!( asinh, super::asinh_order, super::bounds::unbounded_bounds, + true, super::get_asinh_doc ); make_math_unary_udf!( @@ -82,6 +96,7 @@ make_math_unary_udf!( atan, super::atan_order, super::bounds::atan_bounds, + true, super::get_atan_doc ); make_math_unary_udf!( @@ -90,6 +105,7 @@ make_math_unary_udf!( atanh, super::atanh_order, super::bounds::unbounded_bounds, + true, super::get_atanh_doc ); make_math_binary_udf!( @@ -97,6 +113,7 @@ make_math_binary_udf!( atan2, atan2, super::atan2_order, + true, super::get_atan2_doc ); make_math_unary_udf!( @@ -105,6 +122,7 @@ make_math_unary_udf!( cbrt, super::cbrt_order, super::bounds::unbounded_bounds, + true, super::get_cbrt_doc ); make_udf_function!(ceil::CeilFunc, ceil); @@ -114,6 +132,7 @@ make_math_unary_udf!( cos, super::cos_order, super::bounds::cos_bounds, + true, super::get_cos_doc ); make_math_unary_udf!( @@ -122,6 +141,7 @@ make_math_unary_udf!( cosh, super::cosh_order, super::bounds::cosh_bounds, + true, super::get_cosh_doc ); make_udf_function!(cot::CotFunc, cot); @@ -131,6 +151,7 @@ make_math_unary_udf!( to_degrees, super::degrees_order, super::bounds::unbounded_bounds, + true, super::get_degrees_doc ); make_math_unary_udf!( @@ -139,6 +160,7 @@ make_math_unary_udf!( exp, super::exp_order, super::bounds::exp_bounds, + true, super::get_exp_doc ); make_udf_function!(factorial::FactorialFunc, factorial); @@ -154,6 +176,7 @@ make_math_unary_udf!( ln, super::ln_order, super::bounds::unbounded_bounds, + true, super::get_ln_doc ); make_math_unary_udf!( @@ -162,6 +185,7 @@ make_math_unary_udf!( log2, super::log2_order, super::bounds::unbounded_bounds, + true, super::get_log2_doc ); make_math_unary_udf!( @@ -170,6 +194,7 @@ make_math_unary_udf!( log10, super::log10_order, super::bounds::unbounded_bounds, + true, super::get_log10_doc ); make_udf_function!(nanvl::NanvlFunc, nanvl); @@ -181,6 +206,7 @@ make_math_unary_udf!( to_radians, super::radians_order, super::bounds::radians_bounds, + true, super::get_radians_doc ); make_udf_function!(random::RandomFunc, random); @@ -192,6 +218,7 @@ make_math_unary_udf!( sin, super::sin_order, super::bounds::sin_bounds, + true, super::get_sin_doc ); make_math_unary_udf!( @@ -200,6 +227,7 @@ make_math_unary_udf!( sinh, super::sinh_order, super::bounds::unbounded_bounds, + true, super::get_sinh_doc ); make_math_unary_udf!( @@ -208,7 +236,9 @@ make_math_unary_udf!( sqrt, super::sqrt_order, super::bounds::sqrt_bounds, - super::get_sqrt_doc + true, + super::get_sqrt_doc, + Some(super::validate_sqrt_input) ); make_math_unary_udf!( TanFunc, @@ -216,6 +246,7 @@ make_math_unary_udf!( tan, super::tan_order, super::bounds::unbounded_bounds, + true, super::get_tan_doc ); make_math_unary_udf!( @@ -224,10 +255,146 @@ make_math_unary_udf!( tanh, super::tanh_order, super::bounds::tanh_bounds, + true, super::get_tanh_doc ); make_udf_function!(trunc::TruncFunc, trunc); +#[cfg(test)] +mod strict_tests { + use super::*; + use arrow::datatypes::Field; + use datafusion_common::ScalarValue; + use datafusion_expr::{ + ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDF, + }; + use std::sync::Arc; + + #[test] + fn strict_math_functions_propagate_nulls() { + let cases = vec![ + (abs(), vec![ScalarValue::from(1.0)]), + (acos(), vec![ScalarValue::from(0.5)]), + (acosh(), vec![ScalarValue::from(1.5)]), + (asin(), vec![ScalarValue::from(0.5)]), + (asinh(), vec![ScalarValue::from(0.5)]), + (atan(), vec![ScalarValue::from(0.5)]), + ( + atan2(), + vec![ScalarValue::from(0.5), ScalarValue::from(1.0)], + ), + (atanh(), vec![ScalarValue::from(0.5)]), + (cbrt(), vec![ScalarValue::from(8.0)]), + (ceil(), vec![ScalarValue::from(1.5)]), + (cos(), vec![ScalarValue::from(0.5)]), + (cosh(), vec![ScalarValue::from(0.5)]), + (cot(), vec![ScalarValue::from(0.5)]), + (degrees(), vec![ScalarValue::from(0.5)]), + (exp(), vec![ScalarValue::from(0.5)]), + (factorial(), vec![ScalarValue::from(5_i64)]), + (floor(), vec![ScalarValue::from(1.5)]), + ( + gcd(), + vec![ScalarValue::from(48_i64), ScalarValue::from(18_i64)], + ), + (isnan(), vec![ScalarValue::from(1.0)]), + (iszero(), vec![ScalarValue::from(1.0)]), + ( + lcm(), + vec![ScalarValue::from(4_i64), ScalarValue::from(5_i64)], + ), + (ln(), vec![ScalarValue::from(2.0)]), + (log(), vec![ScalarValue::from(10.0)]), + ( + log(), + vec![ScalarValue::from(10.0), ScalarValue::from(100.0)], + ), + (log2(), vec![ScalarValue::from(2.0)]), + (log10(), vec![ScalarValue::from(10.0)]), + ( + power(), + vec![ScalarValue::from(2.0), ScalarValue::from(3.0)], + ), + (radians(), vec![ScalarValue::from(90.0)]), + (round(), vec![ScalarValue::from(1.5)]), + ( + round(), + vec![ScalarValue::from(1.5), ScalarValue::from(1_i32)], + ), + (signum(), vec![ScalarValue::from(-1.0)]), + (sin(), vec![ScalarValue::from(0.5)]), + (sinh(), vec![ScalarValue::from(0.5)]), + (sqrt(), vec![ScalarValue::from(4.0)]), + (tan(), vec![ScalarValue::from(0.5)]), + (tanh(), vec![ScalarValue::from(0.5)]), + (trunc(), vec![ScalarValue::from(1.5)]), + ( + trunc(), + vec![ScalarValue::from(1.5), ScalarValue::from(1_i64)], + ), + ]; + + for (func, valid_args) in cases { + assert!(func.is_strict(), "{} should be marked strict", func.name()); + + for null_mask in 0..(1 << valid_args.len()) { + let mut args = valid_args.clone(); + for (arg_idx, arg) in args.iter_mut().enumerate() { + if null_mask & (1 << arg_idx) != 0 { + *arg = ScalarValue::try_new_null(&arg.data_type()).unwrap(); + } + } + + let result = + invoke_with_scalar_args(&func, args).unwrap_or_else(|error| { + panic!( + "{} failed for NULL mask {null_mask:b}: {error}", + func.name() + ) + }); + let expected_null = null_mask != 0; + let result = result.into_array(1).unwrap(); + assert_eq!( + result.null_count() == result.len(), + expected_null, + "{} returned {result:?} for NULL mask {null_mask:0width$b}", + func.name(), + width = valid_args.len(), + ); + } + } + } + + fn invoke_with_scalar_args( + func: &ScalarUDF, + args: Vec, + ) -> Result { + let arg_fields = args + .iter() + .enumerate() + .map(|(idx, arg)| { + Arc::new(Field::new( + format!("arg_{idx}"), + arg.data_type(), + arg.is_null(), + )) + }) + .collect::>(); + let scalar_arguments = args.iter().map(Some).collect::>(); + let return_field = func.return_field_from_args(ReturnFieldArgs { + arg_fields: &arg_fields, + scalar_arguments: &scalar_arguments, + })?; + func.invoke_with_args(ScalarFunctionArgs { + args: args.into_iter().map(ColumnarValue::Scalar).collect(), + arg_fields, + number_rows: 1, + return_field, + config_options: Arc::new(Default::default()), + }) + } +} + pub mod expr_fn { export_functions!( (abs, "returns the absolute value of a given number", num), diff --git a/datafusion/functions/src/math/monotonicity.rs b/datafusion/functions/src/math/monotonicity.rs index 4a0db9ef0cf7a..d223446d96e90 100644 --- a/datafusion/functions/src/math/monotonicity.rs +++ b/datafusion/functions/src/math/monotonicity.rs @@ -154,7 +154,7 @@ static DOCUMENTATION_ASINH: LazyLock = LazyLock::new(|| { ) .with_standard_argument("numeric_expression", Some("Numeric")) .with_sql_example( - r#" ```sql + r#" ```sql > SELECT asinh(1); +------------+ | asinh(1) | @@ -184,7 +184,7 @@ static DOCUMENTATION_ATAN: LazyLock = LazyLock::new(|| { .with_standard_argument("numeric_expression", Some("Numeric")) .with_sql_example( r#"```sql - > SELECT atan(1); +> SELECT atan(1); +-----------+ | atan(1) | +-----------+ @@ -223,7 +223,7 @@ static DOCUMENTATION_ATANH: LazyLock = ) .with_standard_argument("numeric_expression", Some("Numeric")) .with_sql_example(r#"```sql - > SELECT atanh(0.5); +> SELECT atanh(0.5); +-------------+ | atanh(0.5) | +-------------+ @@ -262,11 +262,11 @@ Can be a constant, column, or function, and any combination of arithmetic operat ) .with_sql_example(r#"```sql > SELECT atan2(1, 1); -+------------+ -| atan2(1,1) | -+------------+ -| 0.7853982 | -+------------+ ++--------------------+ +| atan2(1,1) | ++--------------------+ +| 0.7853981633974483 | ++--------------------+ ```"#) .build() }); @@ -394,7 +394,7 @@ static DOCUMENTATION_DEGREES: LazyLock = LazyLock::new(|| { .with_standard_argument("numeric_expression", Some("Numeric")) .with_sql_example( r#"```sql - > SELECT degrees(pi()); +> SELECT degrees(pi()); +------------+ | degrees(0) | +------------+ @@ -719,12 +719,12 @@ static DOCUMENTATION_TANH: LazyLock = LazyLock::new(|| { .with_standard_argument("numeric_expression", Some("Numeric")) .with_sql_example( r#"```sql - > SELECT tanh(20); - +----------+ - | tanh(20) | - +----------+ - | 1.0 | - +----------+ +> SELECT tanh(20); ++----------+ +| tanh(20) | ++----------+ +| 1.0 | ++----------+ ```"#, ) .build() @@ -761,6 +761,7 @@ mod tests { .unwrap(), sort_properties: sp, preserves_lex_ordering: false, + strictly_order_preserving: false, } } diff --git a/datafusion/functions/src/math/nans.rs b/datafusion/functions/src/math/nans.rs index c5ea2fa079a45..c313db30378bf 100644 --- a/datafusion/functions/src/math/nans.rs +++ b/datafusion/functions/src/math/nans.rs @@ -83,6 +83,10 @@ impl ScalarUDFImpl for IsNanFunc { Ok(DataType::Boolean) } + fn is_strict(&self) -> bool { + true + } + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [arg] = take_function_args(self.name(), args.args)?; diff --git a/datafusion/functions/src/math/nanvl.rs b/datafusion/functions/src/math/nanvl.rs index 251e98bb72c03..cc146983067be 100644 --- a/datafusion/functions/src/math/nanvl.rs +++ b/datafusion/functions/src/math/nanvl.rs @@ -17,16 +17,20 @@ use std::sync::Arc; -use arrow::array::{ArrayRef, AsArray, Float16Array, Float32Array, Float64Array}; +use arrow::array::builder::NullBufferBuilder; +use arrow::array::{Array, ArrayRef, AsArray, PrimitiveArray}; use arrow::datatypes::DataType::{Float16, Float32, Float64}; -use arrow::datatypes::{DataType, Float16Type, Float32Type, Float64Type}; +use arrow::datatypes::{ + ArrowPrimitiveType, DataType, Float16Type, Float32Type, Float64Type, +}; +use datafusion_common::types::{NativeType, logical_float64}; use datafusion_common::{Result, ScalarValue, exec_err, utils::take_function_args}; -use datafusion_expr::TypeSignature::Exact; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - Volatility, + Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, + TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; +use num_traits::Float; #[user_doc( doc_section(label = "Math Functions"), @@ -63,12 +67,32 @@ impl Default for NanvlFunc { impl NanvlFunc { pub fn new() -> Self { + // Non-float numerics (integers, decimals) and NULL coerce to Float64, + // which represents as many inputs as possible before rounding. + let non_float = Coercion::new_implicit( + TypeSignatureClass::Native(logical_float64()), + vec![TypeSignatureClass::Integer, TypeSignatureClass::Decimal], + NativeType::Float64, + ); + // Any numeric (including floats) coerces to Float64. + let to_float64 = Coercion::new_implicit( + TypeSignatureClass::Native(logical_float64()), + vec![TypeSignatureClass::Numeric], + NativeType::Float64, + ); Self { signature: Signature::one_of( vec![ - Exact(vec![Float16, Float16]), - Exact(vec![Float32, Float32]), - Exact(vec![Float64, Float64]), + // If either argument is a non-float numeric (or NULL), both + // are computed in Float64. Two arms cover either argument + // order. + TypeSignature::Coercible(vec![non_float.clone(), to_float64.clone()]), + TypeSignature::Coercible(vec![to_float64, non_float]), + // Otherwise both arguments are floats; preserve their + // (widest common) precision rather than widening to Float64. + TypeSignature::Exact(vec![Float16, Float16]), + TypeSignature::Exact(vec![Float32, Float32]), + TypeSignature::Exact(vec![Float64, Float64]), ], Volatility::Immutable, ), @@ -86,9 +110,9 @@ impl ScalarUDFImpl for NanvlFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - match &arg_types[0] { - Float16 => Ok(Float16), - Float32 => Ok(Float32), + match (&arg_types[0], &arg_types[1]) { + (Float16, Float16) => Ok(Float16), + (Float32, Float32) => Ok(Float32), _ => Ok(Float64), } } @@ -97,16 +121,10 @@ impl ScalarUDFImpl for NanvlFunc { let [x, y] = take_function_args(self.name(), args.args)?; match (x, y) { - (ColumnarValue::Scalar(ScalarValue::Float16(Some(v))), y) if v.is_nan() => { - Ok(y) - } - (ColumnarValue::Scalar(ScalarValue::Float32(Some(v))), y) if v.is_nan() => { - Ok(y) - } - (ColumnarValue::Scalar(ScalarValue::Float64(Some(v))), y) if v.is_nan() => { - Ok(y) - } + // Scalar x: return y if x is NaN, otherwise x (which may be NULL). + (ColumnarValue::Scalar(ref x), y) if scalar_is_nan(x) => Ok(y), (x @ ColumnarValue::Scalar(_), _) => Ok(x), + // At least one argument is an array: evaluate element-wise. (x, y) => { let args = ColumnarValue::values_to_arrays(&[x, y])?; Ok(ColumnarValue::Array(nanvl(&args)?)) @@ -119,52 +137,95 @@ impl ScalarUDFImpl for NanvlFunc { } } +fn scalar_is_nan(scalar: &ScalarValue) -> bool { + match scalar { + ScalarValue::Float16(Some(v)) => v.is_nan(), + ScalarValue::Float32(Some(v)) => v.is_nan(), + ScalarValue::Float64(Some(v)) => v.is_nan(), + _ => false, + } +} + /// Nanvl SQL function /// /// - x is NaN -> output is y (which may itself be NULL) /// - otherwise -> output is x (which may itself be NULL) fn nanvl(args: &[ArrayRef]) -> Result { match args[0].data_type() { - Float64 => { - let x = args[0].as_primitive::(); - let y = args[1].as_primitive::(); - let result: Float64Array = x - .iter() - .zip(y.iter()) - .map(|(x_value, y_value)| match x_value { - Some(x_value) if x_value.is_nan() => y_value, - _ => x_value, - }) - .collect(); - Ok(Arc::new(result) as ArrayRef) - } - Float32 => { - let x = args[0].as_primitive::(); - let y = args[1].as_primitive::(); - let result: Float32Array = x + Float64 => Ok(Arc::new(nanvl_impl::( + args[0].as_primitive(), + args[1].as_primitive(), + ))), + Float32 => Ok(Arc::new(nanvl_impl::( + args[0].as_primitive(), + args[1].as_primitive(), + ))), + Float16 => Ok(Arc::new(nanvl_impl::( + args[0].as_primitive(), + args[1].as_primitive(), + ))), + other => exec_err!("Unsupported data type {other:?} for function nanvl"), + } +} + +/// Element-wise `nanvl`: selects `y[i]` where `x[i]` is `NaN`, otherwise `x[i]` +/// (a null `x` selects `x`, i.e. propagates null). +/// +/// This produces output identical to collecting an iterator of `Option`s but +/// splits out a null-free fast path that iterates the raw value slices, +/// skipping per-element validity checks and `Option` handling. The null-aware +/// path builds its null buffer lazily via [`NullBufferBuilder`]. +fn nanvl_impl(x: &PrimitiveArray, y: &PrimitiveArray) -> PrimitiveArray +where + T: ArrowPrimitiveType, + T::Native: Float, +{ + let xv = x.values(); + let yv = y.values(); + + match (x.nulls(), y.nulls()) { + // No nulls in either input means no nulls in the output, so we can + // iterate values directly and avoid the null bookkeeping entirely. + (None, None) => { + let values: Vec = xv .iter() - .zip(y.iter()) - .map(|(x_value, y_value)| match x_value { - Some(x_value) if x_value.is_nan() => y_value, - _ => x_value, - }) + .zip(yv.iter()) + .map( + |(&x_value, &y_value)| { + if x_value.is_nan() { y_value } else { x_value } + }, + ) .collect(); - Ok(Arc::new(result) as ArrayRef) + PrimitiveArray::::new(values.into(), None) } - Float16 => { - let x = args[0].as_primitive::(); - let y = args[1].as_primitive::(); - let result: Float16Array = x - .iter() - .zip(y.iter()) - .map(|(x_value, y_value)| match x_value { - Some(x_value) if x_value.is_nan() => y_value, - _ => x_value, - }) - .collect(); - Ok(Arc::new(result) as ArrayRef) + _ => { + let len = x.len(); + let mut nulls = NullBufferBuilder::new(len); + let mut values = Vec::with_capacity(len); + for i in 0..len { + // `y` is only consulted when `x` is a (non-null) NaN, matching + // the original short-circuiting match. + if x.is_valid(i) { + let x_value = xv[i]; + if x_value.is_nan() { + if y.is_valid(i) { + values.push(yv[i]); + nulls.append_non_null(); + } else { + values.push(T::Native::default()); + nulls.append_null(); + } + } else { + values.push(x_value); + nulls.append_non_null(); + } + } else { + values.push(T::Native::default()); + nulls.append_null(); + } + } + PrimitiveArray::::new(values.into(), nulls.finish()) } - other => exec_err!("Unsupported data type {other:?} for function nanvl"), } } @@ -174,7 +235,7 @@ mod test { use crate::math::nanvl::nanvl; - use arrow::array::{ArrayRef, Float32Array, Float64Array}; + use arrow::array::{Array, ArrayRef, Float32Array, Float64Array}; use datafusion_common::cast::{as_float32_array, as_float64_array}; #[test] @@ -212,4 +273,91 @@ mod test { assert_eq!(floats.value(2), 3.0); assert!(floats.value(3).is_nan()); } + + #[test] + fn test_nanvl_f64_with_nulls() { + // Covers the null-aware path and null propagation: + // - x null -> null (regardless of y) + // - x NaN, y non-null -> y + // - x NaN, y null -> null + // - x non-NaN -> x + let args: Vec = vec![ + Arc::new(Float64Array::from(vec![ + None, + Some(f64::NAN), + Some(f64::NAN), + Some(2.5), + ])), // x + Arc::new(Float64Array::from(vec![ + Some(9.0), + Some(6.0), + None, + Some(7.0), + ])), // y + ]; + + let result = nanvl(&args).expect("failed to initialize function nanvl"); + let floats = + as_float64_array(&result).expect("failed to initialize function nanvl"); + + assert_eq!(floats.len(), 4); + assert!(floats.is_null(0)); + assert_eq!(floats.value(1), 6.0); + assert!(floats.is_null(2)); + assert_eq!(floats.value(3), 2.5); + } + + #[test] + fn test_nanvl_f64_only_y_nulls() { + // `x` has no nulls, `y` does: + // - x non-NaN -> x + // - x NaN, y non-null -> y + // - x NaN, y null -> null (propagated from y) + let args: Vec = vec![ + Arc::new(Float64Array::from(vec![1.0, f64::NAN, f64::NAN, 4.0])), // x + Arc::new(Float64Array::from(vec![ + Some(5.0), + Some(6.0), + None, + Some(8.0), + ])), // y + ]; + + let result = nanvl(&args).expect("failed to initialize function nanvl"); + let floats = + as_float64_array(&result).expect("failed to initialize function nanvl"); + + assert_eq!(floats.len(), 4); + assert_eq!(floats.value(0), 1.0); + assert_eq!(floats.value(1), 6.0); + assert!(floats.is_null(2)); + assert_eq!(floats.value(3), 4.0); + } + + #[test] + fn test_nanvl_f64_only_x_nulls() { + // `x` has nulls, `y` does not: + // - x null -> null (propagated from x) + // - x NaN -> y + // - x non-NaN -> x + let args: Vec = vec![ + Arc::new(Float64Array::from(vec![ + None, + Some(f64::NAN), + Some(3.0), + None, + ])), // x + Arc::new(Float64Array::from(vec![5.0, 6.0, 7.0, 8.0])), // y + ]; + + let result = nanvl(&args).expect("failed to initialize function nanvl"); + let floats = + as_float64_array(&result).expect("failed to initialize function nanvl"); + + assert_eq!(floats.len(), 4); + assert!(floats.is_null(0)); + assert_eq!(floats.value(1), 6.0); + assert_eq!(floats.value(2), 3.0); + assert!(floats.is_null(3)); + } } diff --git a/datafusion/functions/src/math/power.rs b/datafusion/functions/src/math/power.rs index 3fe30a1ffa86a..30ac401b5ff7d 100644 --- a/datafusion/functions/src/math/power.rs +++ b/datafusion/functions/src/math/power.rs @@ -18,25 +18,20 @@ //! Math function: `power()`. use super::log::LogFunc; -use crate::utils::{calculate_binary_decimal_math, calculate_binary_math}; +use crate::utils::calculate_binary_math; use arrow::array::{Array, ArrayRef}; -use arrow::datatypes::i256; -use arrow::datatypes::{ - ArrowNativeType, ArrowNativeTypeOp, DataType, Decimal32Type, Decimal64Type, - Decimal128Type, Decimal256Type, Float64Type, Int64Type, -}; +use arrow::datatypes::{DataType, Float64Type}; use arrow::error::ArrowError; -use datafusion_common::types::{NativeType, logical_float64, logical_int64}; +use datafusion_common::types::{NativeType, logical_float64}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, Expr, ScalarFunctionArgs, ScalarUDF, - ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, lit, + Cast, Coercion, ColumnarValue, Documentation, Expr, ScalarFunctionArgs, ScalarUDF, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, lit, }; use datafusion_macros::user_doc; -use num_traits::{NumCast, ToPrimitive}; /// Matches PostgreSQL: `power(0::float8, negative)` is undefined (IEEE 754 would yield infinity). #[inline] @@ -78,328 +73,18 @@ impl Default for PowerFunc { impl PowerFunc { pub fn new() -> Self { - let integer = Coercion::new_implicit( - TypeSignatureClass::Native(logical_int64()), - vec![TypeSignatureClass::Integer], - NativeType::Int64, - ); - let decimal = Coercion::new_exact(TypeSignatureClass::Decimal); let float = Coercion::new_implicit( TypeSignatureClass::Native(logical_float64()), vec![TypeSignatureClass::Numeric], NativeType::Float64, ); Self { - signature: Signature::one_of( - vec![ - TypeSignature::Coercible(vec![decimal.clone(), integer]), - TypeSignature::Coercible(vec![decimal.clone(), float.clone()]), - TypeSignature::Coercible(vec![float; 2]), - ], - Volatility::Immutable, - ), + signature: Signature::coercible(vec![float; 2], Volatility::Immutable), aliases: vec![String::from("pow")], } } } -/// Binary function to calculate a math power to integer exponent -/// for scaled integer types. -/// -/// Formula -/// The power for a scaled integer `b` is -/// -/// ```text -/// (b * 10^(-s)) ^ e -/// ``` -/// However, the result should be scaled back from scale 0 to scale `s`, -/// which is done by multiplying by `10^s`. -/// At the end, the formula is: -/// -/// ```text -/// b^e * 10^(-s * e) * 10^s = b^e / 10^(s * (e-1)) -/// ``` -/// Example of 2.5 ^ 4 = 39: -/// 2.5 is represented as 25 with scale 1 -/// The unscaled result is 25^4 = 390625 -/// Scale it back to 1: 390625 / 10^4 = 39 -fn pow_decimal_int(base: T, scale: i8, exp: i64) -> Result -where - T: ArrowNativeType + ArrowNativeTypeOp + ToPrimitive + NumCast + Copy, -{ - // Negative exponent: fall back to float computation - if exp < 0 { - return pow_decimal_float(base, scale, exp as f64); - } - - let exp: u32 = exp.try_into().map_err(|_| { - ArrowError::ArithmeticOverflow(format!("Unsupported exp value: {exp}")) - })?; - // Handle edge case for exp == 0 - // If scale < 0, 10^scale (e.g., 10^-2 = 0.01) becomes 0 in integer arithmetic. - if exp == 0 { - return if scale >= 0 { - T::usize_as(10).pow_checked(scale as u32).map_err(|_| { - ArrowError::ArithmeticOverflow(format!( - "Cannot make unscale factor for {scale} and {exp}" - )) - }) - } else { - Ok(T::ZERO) - }; - } - let powered: T = base.pow_checked(exp).map_err(|_| { - ArrowError::ArithmeticOverflow(format!("Cannot raise base {base:?} to exp {exp}")) - })?; - - // Calculate the scale adjustment: s * (e - 1) - // We use i64 to prevent overflow during the intermediate multiplication - let mul_exp = (scale as i64).wrapping_mul(exp as i64 - 1); - - if mul_exp == 0 { - return Ok(powered); - } - - // If mul_exp is positive, we divide (standard case). - // If mul_exp is negative, we multiply (negative scale case). - if mul_exp > 0 { - let div_factor: T = - T::usize_as(10).pow_checked(mul_exp as u32).map_err(|_| { - ArrowError::ArithmeticOverflow(format!( - "Cannot make div factor for {scale} and {exp}" - )) - })?; - powered.div_checked(div_factor) - } else { - // mul_exp is negative, so we multiply by 10^(-mul_exp) - let abs_exp = mul_exp.checked_neg().ok_or_else(|| { - ArrowError::ArithmeticOverflow( - "Overflow while negating scale exponent".to_string(), - ) - })?; - let mul_factor: T = - T::usize_as(10).pow_checked(abs_exp as u32).map_err(|_| { - ArrowError::ArithmeticOverflow(format!( - "Cannot make mul factor for {scale} and {exp}" - )) - })?; - powered.mul_checked(mul_factor) - } -} - -/// Binary function to calculate a math power to float exponent -/// for scaled integer types. -fn pow_decimal_float(base: T, scale: i8, exp: f64) -> Result -where - T: ArrowNativeType + ArrowNativeTypeOp + ToPrimitive + NumCast + Copy, -{ - if exp.is_finite() && exp.trunc() == exp && exp >= 0f64 && exp < u32::MAX as f64 { - return pow_decimal_int(base, scale, exp as i64); - } - - if !exp.is_finite() { - return Err(ArrowError::ComputeError(format!( - "Cannot use non-finite exp: {exp}" - ))); - } - - pow_decimal_float_fallback(base, scale, exp) -} - -/// Compute the f64 power result and scale it back. -/// Returns the rounded i128 result for conversion to target type. -#[inline] -fn compute_pow_f64_result( - base_f64: f64, - scale: i8, - exp: f64, -) -> Result { - let result_f64 = float64_power_checked(base_f64, exp)?; - - if !result_f64.is_finite() { - return Err(ArrowError::ArithmeticOverflow(format!( - "Result of {base_f64}^{exp} is not finite" - ))); - } - - let scale_factor = 10f64.powi(scale as i32); - let result_scaled = result_f64 * scale_factor; - let result_rounded = result_scaled.round(); - - if result_rounded.abs() > i128::MAX as f64 { - return Err(ArrowError::ArithmeticOverflow(format!( - "Result {result_rounded} is too large for the target decimal type" - ))); - } - - Ok(result_rounded as i128) -} - -/// Convert i128 result to target decimal native type using NumCast. -/// Returns error if value overflows the target type. -#[inline] -fn decimal_from_i128(value: i128) -> Result -where - T: NumCast, -{ - NumCast::from(value).ok_or_else(|| { - ArrowError::ArithmeticOverflow(format!( - "Value {value} is too large for the target decimal type" - )) - }) -} - -/// Fallback implementation using f64 for negative or non-integer exponents. -/// This handles cases that cannot be computed using integer arithmetic. -fn pow_decimal_float_fallback(base: T, scale: i8, exp: f64) -> Result -where - T: ToPrimitive + NumCast + Copy, -{ - if scale < 0 { - return Err(ArrowError::NotYetImplemented(format!( - "Negative scale is not yet supported: {scale}" - ))); - } - - let scale_factor = 10f64.powi(scale as i32); - let base_f64 = base.to_f64().ok_or_else(|| { - ArrowError::ComputeError("Cannot convert base to f64".to_string()) - })? / scale_factor; - - let result_i128 = compute_pow_f64_result(base_f64, scale, exp)?; - - decimal_from_i128(result_i128) -} - -/// Decimal256 specialized float exponent version. -fn pow_decimal256_float(base: i256, scale: i8, exp: f64) -> Result { - if exp.is_finite() && exp.trunc() == exp && exp >= 0f64 && exp < u32::MAX as f64 { - return pow_decimal256_int(base, scale, exp as i64); - } - - if !exp.is_finite() { - return Err(ArrowError::ComputeError(format!( - "Cannot use non-finite exp: {exp}" - ))); - } - - pow_decimal256_float_fallback(base, scale, exp) -} - -/// Decimal256 specialized integer exponent version. -fn pow_decimal256_int(base: i256, scale: i8, exp: i64) -> Result { - if exp < 0 { - return pow_decimal256_float(base, scale, exp as f64); - } - - let exp: u32 = exp.try_into().map_err(|_| { - ArrowError::ArithmeticOverflow(format!("Unsupported exp value: {exp}")) - })?; - - if exp == 0 { - return if scale >= 0 { - i256::from_i128(10).pow_checked(scale as u32).map_err(|_| { - ArrowError::ArithmeticOverflow(format!( - "Cannot make unscale factor for {scale} and {exp}" - )) - }) - } else { - Ok(i256::from_i128(0)) - }; - } - - let powered: i256 = base.pow_checked(exp).map_err(|_| { - ArrowError::ArithmeticOverflow(format!("Cannot raise base {base:?} to exp {exp}")) - })?; - - let mul_exp = (scale as i64).wrapping_mul(exp as i64 - 1); - - if mul_exp == 0 { - return Ok(powered); - } - - if mul_exp > 0 { - let div_factor: i256 = - i256::from_i128(10) - .pow_checked(mul_exp as u32) - .map_err(|_| { - ArrowError::ArithmeticOverflow(format!( - "Cannot make div factor for {scale} and {exp}" - )) - })?; - powered.div_checked(div_factor) - } else { - let abs_exp = mul_exp.checked_neg().ok_or_else(|| { - ArrowError::ArithmeticOverflow( - "Overflow while negating scale exponent".to_string(), - ) - })?; - let mul_factor: i256 = - i256::from_i128(10) - .pow_checked(abs_exp as u32) - .map_err(|_| { - ArrowError::ArithmeticOverflow(format!( - "Cannot make mul factor for {scale} and {exp}" - )) - })?; - powered.mul_checked(mul_factor) - } -} - -/// Fallback implementation for Decimal256. -fn pow_decimal256_float_fallback( - base: i256, - scale: i8, - exp: f64, -) -> Result { - if scale < 0 { - return Err(ArrowError::NotYetImplemented(format!( - "Negative scale is not yet supported: {scale}" - ))); - } - - let scale_factor = 10f64.powi(scale as i32); - let base_f64 = base.to_f64().ok_or_else(|| { - ArrowError::ComputeError("Cannot convert base to f64".to_string()) - })? / scale_factor; - - let result_i128 = compute_pow_f64_result(base_f64, scale, exp)?; - - // i256 can be constructed from i128 directly - Ok(i256::from_i128(result_i128)) -} - -/// Fallback implementation for decimal power when exponent is an array. -/// Casts decimal to float64, computes power, and casts back to original decimal type. -/// This is used for performance when exponent varies per-row. -fn pow_decimal_with_float_fallback( - base: &ArrayRef, - exponent: &ColumnarValue, - num_rows: usize, -) -> Result { - use arrow::compute::cast; - - let original_type = base.data_type().clone(); - let base_f64 = cast(base.as_ref(), &DataType::Float64)?; - - let exp_f64 = match exponent { - ColumnarValue::Array(arr) => cast(arr.as_ref(), &DataType::Float64)?, - ColumnarValue::Scalar(scalar) => { - let scalar_f64 = scalar.cast_to(&DataType::Float64)?; - scalar_f64.to_array_of_size(num_rows)? - } - }; - - let result_f64 = calculate_binary_math::( - &base_f64, - &ColumnarValue::Array(exp_f64), - float64_power_checked, - )?; - - let result = cast(result_f64.as_ref(), &original_type)?; - Ok(ColumnarValue::Array(result)) -} - impl ScalarUDFImpl for PowerFunc { fn name(&self) -> &str { "power" @@ -410,11 +95,12 @@ impl ScalarUDFImpl for PowerFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - if arg_types[0].is_null() { - Ok(DataType::Float64) - } else { - Ok(arg_types[0].clone()) - } + let [_base, _exponent] = take_function_args(self.name(), arg_types)?; + Ok(DataType::Float64) + } + + fn is_strict(&self) -> bool { + true } fn aliases(&self) -> &[String] { @@ -423,25 +109,8 @@ impl ScalarUDFImpl for PowerFunc { fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let [base, exponent] = take_function_args(self.name(), &args.args)?; - - // For decimal types, only use native decimal - // operations when we have a scalar exponent. When the exponent is an array, - // fall back to float computation for better performance. - let use_float_fallback = matches!( - base.data_type(), - DataType::Decimal32(_, _) - | DataType::Decimal64(_, _) - | DataType::Decimal128(_, _) - | DataType::Decimal256(_, _) - ) && matches!(exponent, ColumnarValue::Array(_)); - let base = base.to_array(args.number_rows)?; - // If decimal with array exponent, cast to float and compute - if use_float_fallback { - return pow_decimal_with_float_fallback(&base, exponent, args.number_rows); - } - let arr: ArrayRef = match (base.data_type(), exponent.data_type()) { (DataType::Float64, DataType::Float64) => { calculate_binary_math::( @@ -450,108 +119,6 @@ impl ScalarUDFImpl for PowerFunc { float64_power_checked, )? } - (DataType::Decimal32(precision, scale), DataType::Int64) => { - calculate_binary_decimal_math::( - &base, - exponent, - |b, e| pow_decimal_int(b, *scale, e), - *precision, - *scale, - )? - } - (DataType::Decimal32(precision, scale), DataType::Float64) => { - calculate_binary_decimal_math::< - Decimal32Type, - Float64Type, - Decimal32Type, - _, - >( - &base, - exponent, - |b, e| pow_decimal_float(b, *scale, e), - *precision, - *scale, - )? - } - (DataType::Decimal64(precision, scale), DataType::Int64) => { - calculate_binary_decimal_math::( - &base, - exponent, - |b, e| pow_decimal_int(b, *scale, e), - *precision, - *scale, - )? - } - (DataType::Decimal64(precision, scale), DataType::Float64) => { - calculate_binary_decimal_math::< - Decimal64Type, - Float64Type, - Decimal64Type, - _, - >( - &base, - exponent, - |b, e| pow_decimal_float(b, *scale, e), - *precision, - *scale, - )? - } - (DataType::Decimal128(precision, scale), DataType::Int64) => { - calculate_binary_decimal_math::< - Decimal128Type, - Int64Type, - Decimal128Type, - _, - >( - &base, - exponent, - |b, e| pow_decimal_int(b, *scale, e), - *precision, - *scale, - )? - } - (DataType::Decimal128(precision, scale), DataType::Float64) => { - calculate_binary_decimal_math::< - Decimal128Type, - Float64Type, - Decimal128Type, - _, - >( - &base, - exponent, - |b, e| pow_decimal_float(b, *scale, e), - *precision, - *scale, - )? - } - (DataType::Decimal256(precision, scale), DataType::Int64) => { - calculate_binary_decimal_math::< - Decimal256Type, - Int64Type, - Decimal256Type, - _, - >( - &base, - exponent, - |b, e| pow_decimal256_int(b, *scale, e), - *precision, - *scale, - )? - } - (DataType::Decimal256(precision, scale), DataType::Float64) => { - calculate_binary_decimal_math::< - Decimal256Type, - Float64Type, - Decimal256Type, - _, - >( - &base, - exponent, - |b, e| pow_decimal256_float(b, *scale, e), - *precision, - *scale, - )? - } (base_type, exp_type) => { return internal_err!( "Unsupported data types for base {base_type:?} and exponent {exp_type:?} for power" @@ -573,31 +140,56 @@ impl ScalarUDFImpl for PowerFunc { let [base, exponent] = take_function_args("power", args)?; let base_type = info.get_data_type(&base)?; let exponent_type = info.get_data_type(&exponent)?; + let base_nullable = info.nullable(&base)?; + let return_type = + self.return_type(&[base_type.clone(), exponent_type.clone()])?; // Null propagation if base_type.is_null() || exponent_type.is_null() { - let return_type = self.return_type(&[base_type, exponent_type])?; return Ok(ExprSimplifyResult::Simplified(lit( ScalarValue::Null.cast_to(&return_type)? ))); } + // `simplify` runs on the logical expression *before* type coercion, + // so a simplified sub-expression may still carry its original type + // rather than the Float64 that `power` is declared to return. Cast it + // back when needed to preserve the schema the optimizer already + // committed to — e.g. `power(int_col, 1)` simplifies to `int_col`, + // and the `b` in `power(b, log(b, uint_col))` simplifies to `uint_col`, + // both of which must become Float64. + let cast_to_return_type = |expr: Expr, expr_type: &DataType| { + if expr_type == &return_type { + expr + } else { + Expr::Cast(Cast::new(Box::new(expr), return_type.clone())) + } + }; + match exponent { Expr::Literal(value, _) - if value == ScalarValue::new_zero(&exponent_type)? => + if value == ScalarValue::new_zero(&exponent_type)? && !base_nullable => { Ok(ExprSimplifyResult::Simplified(lit(ScalarValue::new_one( - &base_type, + &return_type, )?))) } Expr::Literal(value, _) if value == ScalarValue::new_one(&exponent_type)? => { - Ok(ExprSimplifyResult::Simplified(base)) + Ok(ExprSimplifyResult::Simplified(cast_to_return_type( + base, &base_type, + ))) } Expr::ScalarFunction(ScalarFunction { func, mut args }) - if is_log(&func) && args.len() == 2 && base == args[0] => + if is_log(&func) + && args.len() == 2 + && base == args[0] + && !base_nullable => { let b = args.pop().unwrap(); // length checked above - Ok(ExprSimplifyResult::Simplified(b)) + let b_type = info.get_data_type(&b)?; + Ok(ExprSimplifyResult::Simplified(cast_to_return_type( + b, &b_type, + ))) } _ => Ok(ExprSimplifyResult::Original(vec![base, exponent])), } @@ -617,69 +209,6 @@ fn is_log(func: &ScalarUDF) -> bool { mod tests { use super::*; - #[test] - fn test_pow_decimal128_helper() { - // Expression: 2.5 ^ 4 = 39.0625 - assert_eq!(pow_decimal_int(25i128, 1, 4).unwrap(), 390i128); - assert_eq!(pow_decimal_int(2500i128, 3, 4).unwrap(), 39062i128); - assert_eq!(pow_decimal_int(25000i128, 4, 4).unwrap(), 390625i128); - - // Expression: 25 ^ 4 = 390625 - assert_eq!(pow_decimal_int(25i128, 0, 4).unwrap(), 390625i128); - - // Expressions for edge cases - assert_eq!(pow_decimal_int(25i128, 1, 1).unwrap(), 25i128); - assert_eq!(pow_decimal_int(25i128, 0, 1).unwrap(), 25i128); - assert_eq!(pow_decimal_int(25i128, 0, 0).unwrap(), 1i128); - assert_eq!(pow_decimal_int(25i128, 1, 0).unwrap(), 10i128); - - assert_eq!(pow_decimal_int(25i128, -1, 4).unwrap(), 390625000i128); - } - - #[test] - fn test_pow_decimal_float_fallback() { - // Test negative exponent: 4^(-1) = 0.25 - // 4 with scale 2 = 400, result should be 25 (0.25 with scale 2) - let result: i128 = pow_decimal_float(400i128, 2, -1.0).unwrap(); - assert_eq!(result, 25); - - // Test non-integer exponent: 4^0.5 = 2 - // 4 with scale 2 = 400, result should be 200 (2.0 with scale 2) - let result: i128 = pow_decimal_float(400i128, 2, 0.5).unwrap(); - assert_eq!(result, 200); - - // Test 8^(1/3) = 2 (cube root) - // 8 with scale 1 = 80, result should be 20 (2.0 with scale 1) - let result: i128 = pow_decimal_float(80i128, 1, 1.0 / 3.0).unwrap(); - assert_eq!(result, 20); - - // Test negative base with integer exponent still works - // (-2)^3 = -8 - // -2 with scale 1 = -20, result should be -80 (-8.0 with scale 1) - let result: i128 = pow_decimal_float(-20i128, 1, 3.0).unwrap(); - assert_eq!(result, -80); - - // Test positive integer exponent goes through fast path - // 2.5^4 = 39.0625 - // 25 with scale 1, result should be 390 (39.0 with scale 1) - truncated - let result: i128 = pow_decimal_float(25i128, 1, 4.0).unwrap(); - assert_eq!(result, 390); // Uses integer path - - // Test non-finite exponent returns error - assert!(pow_decimal_float(100i128, 2, f64::NAN).is_err()); - assert!(pow_decimal_float(100i128, 2, f64::INFINITY).is_err()); - - // PostgreSQL: zero to a negative power is undefined - assert!(pow_decimal_float(0i128, 2, -1.0).is_err()); - } - - #[test] - fn test_pow_decimal256_zero_to_negative_exp_errors() { - assert!(pow_decimal256_float(i256::ZERO, 2, -1.0).is_err()); - // Negative integer exponent uses pow_decimal256_float via pow_decimal256_int - assert!(pow_decimal256_int(i256::ZERO, 2, -1).is_err()); - } - #[test] fn test_float64_power_checked_zero_negative_exp() { assert_eq!(float64_power_checked(0.0, 1.0).unwrap(), 0.0); diff --git a/datafusion/functions/src/math/round.rs b/datafusion/functions/src/math/round.rs index 78016c0f52f71..10500810a56b4 100644 --- a/datafusion/functions/src/math/round.rs +++ b/datafusion/functions/src/math/round.rs @@ -15,15 +15,17 @@ // specific language governing permissions and limitations // under the License. -use crate::utils::{calculate_binary_decimal_math, calculate_binary_math}; +use crate::utils::{calculate_binary_decimal_math_cast, calculate_binary_math}; -use arrow::array::ArrayRef; +use arrow::array::{Array, ArrayRef, AsArray}; use arrow::datatypes::DataType::{ - Decimal32, Decimal64, Decimal128, Decimal256, Float32, Float64, + Decimal32, Decimal64, Decimal128, Decimal256, Float32, Float64, Int8, Int16, Int32, + Int64, UInt8, UInt16, UInt32, UInt64, }; use arrow::datatypes::{ - ArrowNativeTypeOp, DataType, Decimal32Type, Decimal64Type, Decimal128Type, - Decimal256Type, DecimalType, Float32Type, Float64Type, Int32Type, + ArrowNativeTypeOp, ArrowPrimitiveType, DataType, Decimal32Type, Decimal64Type, + Decimal128Type, Decimal256Type, DecimalType, Float32Type, Float64Type, Int8Type, + Int16Type, Int32Type, Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; use arrow::datatypes::{Field, FieldRef}; use arrow::error::ArrowError; @@ -37,6 +39,7 @@ use datafusion_expr::{ ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; +use num_traits::{PrimInt, Signed, cast, checked_pow}; use std::sync::Arc; fn output_scale_for_decimal(precision: u8, input_scale: i8, decimal_places: i32) -> i8 { @@ -185,6 +188,7 @@ impl RoundFunc { vec![TypeSignatureClass::Integer], NativeType::Int32, ); + let integer = Coercion::new_exact(TypeSignatureClass::Integer); let float32 = Coercion::new_exact(TypeSignatureClass::Native(logical_float32())); let float64 = Coercion::new_implicit( TypeSignatureClass::Native(logical_float64()), @@ -199,6 +203,11 @@ impl RoundFunc { decimal_places.clone(), ]), TypeSignature::Coercible(vec![decimal]), + TypeSignature::Coercible(vec![ + integer.clone(), + decimal_places.clone(), + ]), + TypeSignature::Coercible(vec![integer]), TypeSignature::Coercible(vec![ float32.clone(), decimal_places.clone(), @@ -218,6 +227,10 @@ impl ScalarUDFImpl for RoundFunc { "round" } + fn is_strict(&self) -> bool { + true + } + fn signature(&self) -> &Signature { &self.signature } @@ -245,6 +258,7 @@ impl ScalarUDFImpl for RoundFunc { // extra precision to accommodate potential carry-over. let return_type = match input_type { + input_type if input_type.is_integer() => input_type.clone(), Float32 => Float32, Decimal32(precision, scale) => calculate_new_precision_scale::< Decimal32Type, @@ -308,6 +322,9 @@ impl ScalarUDFImpl for RoundFunc { }; match (value_scalar, args.return_type()) { + (value_scalar, return_type) if return_type.is_integer() => { + round_integer_scalar(value_scalar, return_type, dp) + } (ScalarValue::Float32(Some(v)), _) => { let rounded = round_float(*v, dp)?; Ok(ColumnarValue::Scalar(ScalarValue::from(rounded))) @@ -468,25 +485,25 @@ fn round_columnar( let decimal_places_is_array = matches!(decimal_places, ColumnarValue::Array(_)); let arr: ArrayRef = match (value_array.data_type(), return_type) { - (Float64, _) => { - let result = calculate_binary_math::( - value_array.as_ref(), - decimal_places, - round_float::, - )?; - result as _ - } - (Float32, _) => { - let result = calculate_binary_math::( - value_array.as_ref(), - decimal_places, - round_float::, - )?; - result as _ + (input_type, return_type) + if input_type == return_type && return_type.is_integer() => + { + match decimal_places { + ColumnarValue::Scalar(ScalarValue::Int32(Some(dp))) if *dp >= 0 => { + value_array + } + _ => round_integer_array( + value_array.as_ref(), + decimal_places, + return_type, + )?, + } } + (Float64, _) => round_float_column::(&value_array, decimal_places)?, + (Float32, _) => round_float_column::(&value_array, decimal_places)?, (Decimal32(input_precision, scale), Decimal32(precision, new_scale)) => { // reduce scale to reclaim integer precision - let result = calculate_binary_decimal_math::< + let result = calculate_binary_decimal_math_cast::< Decimal32Type, Int32Type, Decimal32Type, @@ -518,11 +535,12 @@ fn round_columnar( }, *precision, *new_scale, + &Int32, )?; result as _ } (Decimal64(input_precision, scale), Decimal64(precision, new_scale)) => { - let result = calculate_binary_decimal_math::< + let result = calculate_binary_decimal_math_cast::< Decimal64Type, Int32Type, Decimal64Type, @@ -551,11 +569,12 @@ fn round_columnar( }, *precision, *new_scale, + &Int32, )?; result as _ } (Decimal128(input_precision, scale), Decimal128(precision, new_scale)) => { - let result = calculate_binary_decimal_math::< + let result = calculate_binary_decimal_math_cast::< Decimal128Type, Int32Type, Decimal128Type, @@ -584,11 +603,12 @@ fn round_columnar( }, *precision, *new_scale, + &Int32, )?; result as _ } (Decimal256(input_precision, scale), Decimal256(precision, new_scale)) => { - let result = calculate_binary_decimal_math::< + let result = calculate_binary_decimal_math_cast::< Decimal256Type, Int32Type, Decimal256Type, @@ -617,6 +637,7 @@ fn round_columnar( }, *precision, *new_scale, + &Int32, )?; result as _ } @@ -630,15 +651,257 @@ fn round_columnar( } } -fn round_float(value: T, decimal_places: i32) -> Result +fn round_signed_integer( + value: T, + decimal_places: i32, + type_name: &str, +) -> Result where - T: num_traits::Float, + T: PrimInt + Signed, +{ + if decimal_places >= 0 || value == T::zero() { + return Ok(value); + } + + let ten = cast::<_, T>(10).expect("10 fits in all integer types"); + let Some(factor) = checked_pow(ten, decimal_places.unsigned_abs() as usize) else { + return Ok(T::zero()); + }; + + let two = cast::<_, T>(2).expect("2 fits in all integer types"); + let one = T::one(); + let threshold = factor / two; + let mut quotient = value / factor; + let remainder = value % factor; + + if remainder >= threshold { + quotient = quotient.checked_add(&one).ok_or_else(|| { + ArrowError::ComputeError(format!("Overflow while rounding {type_name}")) + })?; + } else if remainder <= -threshold { + quotient = quotient.checked_sub(&one).ok_or_else(|| { + ArrowError::ComputeError(format!("Overflow while rounding {type_name}")) + })?; + } + + quotient.checked_mul(&factor).ok_or_else(|| { + ArrowError::ComputeError(format!("Overflow while rounding {type_name}")) + }) +} + +fn round_unsigned_integer( + value: T, + decimal_places: i32, + type_name: &str, +) -> Result +where + T: PrimInt, { - let factor = T::from(10_f64.powi(decimal_places)).ok_or_else(|| { + if decimal_places >= 0 || value == T::zero() { + return Ok(value); + } + + let ten = cast::<_, T>(10).expect("10 fits in all integer types"); + let Some(factor) = checked_pow(ten, decimal_places.unsigned_abs() as usize) else { + return Ok(T::zero()); + }; + + let two = cast::<_, T>(2).expect("2 fits in all integer types"); + let one = T::one(); + let threshold = factor / two; + let mut quotient = value / factor; + let remainder = value % factor; + + if remainder >= threshold { + quotient = quotient.checked_add(&one).ok_or_else(|| { + ArrowError::ComputeError(format!("Overflow while rounding {type_name}")) + })?; + } + + quotient.checked_mul(&factor).ok_or_else(|| { + ArrowError::ComputeError(format!("Overflow while rounding {type_name}")) + }) +} + +fn round_integer_scalar( + value: &ScalarValue, + return_type: &DataType, + decimal_places: i32, +) -> Result { + match (value, return_type) { + (ScalarValue::Int8(Some(v)), Int8) => Ok(ColumnarValue::Scalar( + ScalarValue::Int8(Some(round_signed_integer(*v, decimal_places, "Int8")?)), + )), + (ScalarValue::Int16(Some(v)), Int16) => Ok(ColumnarValue::Scalar( + ScalarValue::Int16(Some(round_signed_integer(*v, decimal_places, "Int16")?)), + )), + (ScalarValue::Int32(Some(v)), Int32) => Ok(ColumnarValue::Scalar( + ScalarValue::Int32(Some(round_signed_integer(*v, decimal_places, "Int32")?)), + )), + (ScalarValue::Int64(Some(v)), Int64) => Ok(ColumnarValue::Scalar( + ScalarValue::Int64(Some(round_signed_integer(*v, decimal_places, "Int64")?)), + )), + (ScalarValue::UInt8(Some(v)), UInt8) => { + Ok(ColumnarValue::Scalar(ScalarValue::UInt8(Some( + round_unsigned_integer(*v, decimal_places, "UInt8")?, + )))) + } + (ScalarValue::UInt16(Some(v)), UInt16) => { + Ok(ColumnarValue::Scalar(ScalarValue::UInt16(Some( + round_unsigned_integer(*v, decimal_places, "UInt16")?, + )))) + } + (ScalarValue::UInt32(Some(v)), UInt32) => { + Ok(ColumnarValue::Scalar(ScalarValue::UInt32(Some( + round_unsigned_integer(*v, decimal_places, "UInt32")?, + )))) + } + (ScalarValue::UInt64(Some(v)), UInt64) => { + Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some( + round_unsigned_integer(*v, decimal_places, "UInt64")?, + )))) + } + _ => internal_err!( + "Unexpected integer round input/output types: {} -> {}", + value.data_type(), + return_type + ), + } +} + +macro_rules! round_integer_array { + ($ARRAY:expr, $DP:expr, $ARRAY_TYPE:ty, $ROUND_FN:ident, $TYPE_NAME:expr) => {{ + let array = $ARRAY.as_primitive::<$ARRAY_TYPE>(); + + let result = calculate_binary_math::<$ARRAY_TYPE, Int32Type, $ARRAY_TYPE, _>( + array, + $DP, + |v, dp| $ROUND_FN(v, dp, $TYPE_NAME), + )?; + + Ok(result as ArrayRef) + }}; +} + +fn round_integer_array( + value_array: &dyn Array, + decimal_places: &ColumnarValue, + return_type: &DataType, +) -> Result { + match return_type { + Int8 => round_integer_array!( + value_array, + decimal_places, + Int8Type, + round_signed_integer, + "Int8" + ), + Int16 => round_integer_array!( + value_array, + decimal_places, + Int16Type, + round_signed_integer, + "Int16" + ), + Int32 => round_integer_array!( + value_array, + decimal_places, + Int32Type, + round_signed_integer, + "Int32" + ), + Int64 => round_integer_array!( + value_array, + decimal_places, + Int64Type, + round_signed_integer, + "Int64" + ), + UInt8 => round_integer_array!( + value_array, + decimal_places, + UInt8Type, + round_unsigned_integer, + "UInt8" + ), + UInt16 => round_integer_array!( + value_array, + decimal_places, + UInt16Type, + round_unsigned_integer, + "UInt16" + ), + UInt32 => round_integer_array!( + value_array, + decimal_places, + UInt32Type, + round_unsigned_integer, + "UInt32" + ), + UInt64 => round_integer_array!( + value_array, + decimal_places, + UInt64Type, + round_unsigned_integer, + "UInt64" + ), + _ => internal_err!("Unexpected return type for integer round: {return_type}"), + } +} + +/// Rounds a float array to `decimal_places`. +/// +/// The shared `calculate_binary_math` kernel routes through `try_unary` and +/// re-evaluates `round_float` (including `10f64.powi(decimal_places)` and a +/// `Result` check) for every element. When `decimal_places` is a non-null +/// scalar, the scaling factor can instead be hoisted out of the loop and the +/// infallible `unary` kernel used, which the compiler can autovectorize. +/// `unary` also computes over null slots, but it carries the input null buffer +/// through to the output, so those values stay masked. +fn round_float_column( + value_array: &ArrayRef, + decimal_places: &ColumnarValue, +) -> Result +where + PT: ArrowPrimitiveType, + PT::Native: num_traits::Float, +{ + // Bring `Float` into scope so `.round()` resolves on the `PT::Native` + // projection below. + use num_traits::Float; + + if let ColumnarValue::Scalar(ScalarValue::Int32(Some(decimal_places))) = + decimal_places + { + let factor = round_factor::(*decimal_places)?; + let result = value_array + .as_primitive::() + .unary::<_, PT>(|value| (value * factor).round() / factor); + return Ok(Arc::new(result) as ArrayRef); + } + + let result = calculate_binary_math::( + value_array.as_ref(), + decimal_places, + round_float::, + )?; + Ok(result as _) +} + +/// Computes the power-of-ten scaling factor used to round to `decimal_places`. +fn round_factor(decimal_places: i32) -> Result { + T::from(10_f64.powi(decimal_places)).ok_or_else(|| { ArrowError::ComputeError(format!( "Invalid value for decimal places: {decimal_places}" )) - })?; + }) +} + +fn round_float(value: T, decimal_places: i32) -> Result +where + T: num_traits::Float, +{ + let factor = round_factor::(decimal_places)?; Ok((value * factor).round() / factor) } @@ -728,6 +991,7 @@ mod test { use std::sync::Arc; use arrow::array::{ArrayRef, Float32Array, Float64Array, Int64Array}; + use arrow::datatypes::DataType; use datafusion_common::DataFusionError; use datafusion_common::ScalarValue; use datafusion_common::cast::{as_float32_array, as_float64_array}; @@ -793,6 +1057,35 @@ mod test { assert_eq!(floats, &expected); } + /// A scalar `decimal_places` takes the hoisted-factor `unary` path, which + /// computes over null slots as well. The nulls must survive into the output. + #[test] + fn test_round_f64_scalar_decimal_places_preserves_nulls() { + let value: ArrayRef = Arc::new(Float64Array::from(vec![ + Some(125.2345), + None, + Some(-1.555), + None, + ])); + + let result = super::round_columnar( + &ColumnarValue::Array(value), + &ColumnarValue::Scalar(ScalarValue::Int32(Some(2))), + 4, + &DataType::Float64, + ) + .expect("failed to initialize function round"); + let ColumnarValue::Array(result) = result else { + panic!("expected an array result"); + }; + let floats = + as_float64_array(&result).expect("failed to initialize function round"); + + let expected = Float64Array::from(vec![Some(125.23), None, Some(-1.56), None]); + + assert_eq!(floats, &expected); + } + #[test] fn test_round_f32_one_input() { let args: Vec = vec![ diff --git a/datafusion/functions/src/math/signum.rs b/datafusion/functions/src/math/signum.rs index 8c8eeacf12394..05b78fcffe2a7 100644 --- a/datafusion/functions/src/math/signum.rs +++ b/datafusion/functions/src/math/signum.rs @@ -86,6 +86,10 @@ impl ScalarUDFImpl for SignumFunc { } } + fn is_strict(&self) -> bool { + true + } + fn output_ordering(&self, input: &[ExprProperties]) -> Result { // Non-decreasing for all real numbers x. Ok(input[0].sort_properties) diff --git a/datafusion/functions/src/math/trunc.rs b/datafusion/functions/src/math/trunc.rs index 991ad0e9c470d..bb8bea8ae75de 100644 --- a/datafusion/functions/src/math/trunc.rs +++ b/datafusion/functions/src/math/trunc.rs @@ -15,22 +15,32 @@ // specific language governing permissions and limitations // under the License. +use std::ops::{Div, Mul}; use std::sync::Arc; -use crate::utils::make_scalar_function; +use crate::utils::{calculate_binary_decimal_math_cast, make_scalar_function}; use arrow::array::{ArrayRef, AsArray, PrimitiveArray}; -use arrow::datatypes::DataType::{Float32, Float64}; -use arrow::datatypes::{DataType, Float32Type, Float64Type, Int64Type}; +use arrow::datatypes::DataType::{ + Decimal32, Decimal64, Decimal128, Decimal256, Float32, Float64, +}; +use arrow::datatypes::{ + ArrowPrimitiveType, DataType, Decimal32Type, Decimal64Type, Decimal128Type, + Decimal256Type, DecimalType, Float32Type, Float64Type, Int64Type, +}; use datafusion_common::ScalarValue::Int64; -use datafusion_common::{Result, ScalarValue, exec_err}; -use datafusion_expr::TypeSignature::Exact; +use datafusion_common::types::{ + NativeType, logical_float32, logical_float64, logical_int64, +}; +use datafusion_common::{Result, ScalarValue, exec_err, plan_err}; use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; use datafusion_expr::{ ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; +use datafusion_expr_common::signature::{Coercion, TypeSignature, TypeSignatureClass}; use datafusion_macros::user_doc; +use num_traits::{Float, NumCast, One, Zero, pow}; #[user_doc( doc_section(label = "Math Functions"), @@ -68,19 +78,38 @@ impl Default for TruncFunc { impl TruncFunc { pub fn new() -> Self { - use DataType::*; + let decimal = Coercion::new_exact(TypeSignatureClass::Decimal); + let decimal_places = Coercion::new_implicit( + TypeSignatureClass::Native(logical_int64()), + vec![TypeSignatureClass::Integer], + NativeType::Int64, + ); + let float32 = Coercion::new_exact(TypeSignatureClass::Native(logical_float32())); + let float64 = Coercion::new_implicit( + TypeSignatureClass::Native(logical_float64()), + vec![TypeSignatureClass::Numeric], + NativeType::Float64, + ); Self { // math expressions expect 1 argument of type f64 or f32 // priority is given to f64 because e.g. `sqrt(1i32)` is in IR (real numbers) and thus we // return the best approximation for it (in f64). // We accept f32 because in this case it is clear that the best approximation - // will be as good as the number of digits in the number + // Decimal arguments are accepted to handle large values properly signature: Signature::one_of( vec![ - Exact(vec![Float32, Int64]), - Exact(vec![Float64, Int64]), - Exact(vec![Float64]), - Exact(vec![Float32]), + TypeSignature::Coercible(vec![ + decimal.clone(), + decimal_places.clone(), + ]), + TypeSignature::Coercible(vec![decimal]), + TypeSignature::Coercible(vec![ + float32.clone(), + decimal_places.clone(), + ]), + TypeSignature::Coercible(vec![float32]), + TypeSignature::Coercible(vec![float64.clone(), decimal_places]), + TypeSignature::Coercible(vec![float64]), ], Volatility::Immutable, ), @@ -93,14 +122,25 @@ impl ScalarUDFImpl for TruncFunc { "trunc" } + fn is_strict(&self) -> bool { + true + } + fn signature(&self) -> &Signature { &self.signature } fn return_type(&self, arg_types: &[DataType]) -> Result { - match arg_types[0] { + match &arg_types[0] { Float32 => Ok(Float32), - _ => Ok(Float64), + Float64 => Ok(Float64), + dt if dt.is_decimal() => Ok(dt.clone()), + DataType::Null => Ok(Float64), + _ => plan_err!( + "Unsupported data type {:?} for function {}", + arg_types[0], + self.name() + ), } } @@ -122,6 +162,12 @@ impl ScalarUDFImpl for TruncFunc { } }; + // Whether an explicit precision argument was supplied. The array fast + // paths below must only apply to the two-argument form: single-argument + // `trunc(x)` uses a different zero handling (mapping `-0.0` to `0.0`) + // that must be preserved. + let has_precision_arg = args.args.len() == 2; + // Scalar fast path using tuple matching for (value, precision) match (&args.args[0], precision) { // Null cases @@ -146,6 +192,74 @@ impl ScalarUDFImpl for TruncFunc { compute_truncate32(*v, p) }))), ), + ( + ColumnarValue::Scalar(ScalarValue::Decimal32( + Some(v), + lprecision, + lscale, + )), + Some(p), + ) => Ok(ColumnarValue::Scalar(ScalarValue::Decimal32( + Some(compute_truncate_decimal::(*v, *lscale, p)), + *lprecision, + *lscale, + ))), + ( + ColumnarValue::Scalar(ScalarValue::Decimal64( + Some(v), + lprecision, + lscale, + )), + Some(p), + ) => Ok(ColumnarValue::Scalar(ScalarValue::Decimal64( + Some(compute_truncate_decimal::(*v, *lscale, p)), + *lprecision, + *lscale, + ))), + ( + ColumnarValue::Scalar(ScalarValue::Decimal128( + Some(v), + lprecision, + lscale, + )), + Some(p), + ) => Ok(ColumnarValue::Scalar(ScalarValue::Decimal128( + Some(compute_truncate_decimal::(*v, *lscale, p)), + *lprecision, + *lscale, + ))), + ( + ColumnarValue::Scalar(ScalarValue::Decimal256( + Some(v), + lprecision, + lscale, + )), + Some(p), + ) => Ok(ColumnarValue::Scalar(ScalarValue::Decimal256( + Some(compute_truncate_decimal::(*v, *lscale, p)), + *lprecision, + *lscale, + ))), + + // Array value with a constant (scalar) precision: hoist the power + // of ten out of the per-element loop instead of broadcasting the + // scalar into a full precision array and recomputing `10^p` for + // every element (see `truncate_float_array`). + (ColumnarValue::Array(arr), Some(p)) + if has_precision_arg && arr.data_type() == &Float64 => + { + Ok(ColumnarValue::Array(truncate_float_array::( + arr, p, + ))) + } + (ColumnarValue::Array(arr), Some(p)) + if has_precision_arg && arr.data_type() == &Float32 => + { + Ok(ColumnarValue::Array(truncate_float_array::( + arr, p, + ))) + } + // Array path for everything else _ => make_scalar_function(trunc, vec![])(&args.args), } @@ -234,27 +348,139 @@ fn trunc(args: &[ArrayRef]) -> Result { } _ => exec_err!("trunc function requires a scalar or array for precision"), }, + Decimal32(lprecision, lscale) => Ok(calculate_binary_decimal_math_cast::< + Decimal32Type, + Int64Type, + Decimal32Type, + _, + >( + num.as_ref(), + &precision, + |v, y| Ok(compute_truncate_decimal::(v, *lscale, y)), + *lprecision, + *lscale, + &DataType::Int64, + )? as ArrayRef), + Decimal64(lprecision, lscale) => Ok(calculate_binary_decimal_math_cast::< + Decimal64Type, + Int64Type, + Decimal64Type, + _, + >( + num.as_ref(), + &precision, + |v, y| Ok(compute_truncate_decimal::(v, *lscale, y)), + *lprecision, + *lscale, + &DataType::Int64, + )? as ArrayRef), + Decimal128(lprecision, lscale) => Ok(calculate_binary_decimal_math_cast::< + Decimal128Type, + Int64Type, + Decimal128Type, + _, + >( + num.as_ref(), + &precision, + |v, y| Ok(compute_truncate_decimal::(v, *lscale, y)), + *lprecision, + *lscale, + &DataType::Int64, + )? as ArrayRef), + Decimal256(lprecision, lscale) => Ok(calculate_binary_decimal_math_cast::< + Decimal256Type, + Int64Type, + Decimal256Type, + _, + >( + num.as_ref(), + &precision, + |v, y| Ok(compute_truncate_decimal::(v, *lscale, y)), + *lprecision, + *lscale, + &DataType::Int64, + )? as ArrayRef), other => exec_err!("Unsupported data type {other:?} for function trunc"), } } -fn compute_truncate32(x: f32, y: i64) -> f32 { - let factor = 10.0_f32.powi(y as i32); +/// Truncates `x` using a pre-computed `factor` of `10^precision`. Taking the +/// factor as an argument lets callers hoist `10^precision` out of a per-element +/// loop when the precision is constant. +fn truncate_with_factor(x: F, factor: F) -> F { (x * factor).trunc() / factor } +/// Truncates every element of a float array to `precision` decimal places, +/// computing the `10^precision` factor once and reusing it for every element. +fn truncate_float_array(arr: &ArrayRef, precision: i64) -> ArrayRef +where + T: ArrowPrimitiveType, + T::Native: Float, +{ + let factor = ::from(10.0_f64) + .unwrap() + .powi(precision as i32); + Arc::new( + arr.as_primitive::() + .unary::<_, T>(|x| truncate_with_factor(x, factor)), + ) +} + +fn compute_truncate32(x: f32, y: i64) -> f32 { + truncate_with_factor(x, 10.0_f32.powi(y as i32)) +} + fn compute_truncate64(x: f64, y: i64) -> f64 { - let factor = 10.0_f64.powi(y as i32); - (x * factor).trunc() / factor + truncate_with_factor(x, 10.0_f64.powi(y as i32)) +} + +/// Truncates a decimal value to `truncate_precision` fractional digits. +/// If `truncate_precision` is positive, clear that amount of trailing low-order digits +/// If `truncate_precision` is negative, it also clears digits before a decimal point +/// +/// Example: +/// Truncating number 12.3456 (123456 as i128 with scale=4) to 1 digit produces 12.3. +/// It makes exp = 4-1 = 3; factor = 10^3 = 1000; result = (123456 / 1000) * 1000 = 123000 +/// It is a decimal 12.3 with scale=4 +/// +/// Truncating number 12.3456 to -1 digit produces 10.0. +/// It makes exp = 4-(-1) = 5; factor = 10^5 = 100000; result = (123456 / 100000) * 100000 = 100000 +/// It is a decimal 10.0 with scale=4 +fn compute_truncate_decimal( + x: T::Native, + scale: i8, + truncate_precision: i64, +) -> T::Native +where + T: DecimalType, + T::Native: Copy + From + One + Zero + Div + Mul, +{ + // How many trailing digits of decimal to clear + let exp = (scale as i64).saturating_sub(truncate_precision); + if exp <= 0 { + // Keep more digits than actually stored, so nothing to truncate + x + } else if exp >= T::MAX_PRECISION as i64 { + // Drop more digits that can be stored, return 0 without overflowing `pow` + T::Native::zero() + } else { + let base = T::Native::from(10_i32); + let exp = exp as usize; + let factor = pow::(base, exp); + // Result is (x / factor) * factor, so (x/factor) drops extra digits + (x / factor) * factor + } } #[cfg(test)] mod test { use std::sync::Arc; - use crate::math::trunc::trunc; + use crate::math::trunc::{compute_truncate_decimal, trunc}; use arrow::array::{ArrayRef, Float32Array, Float64Array, Int64Array}; + use arrow::datatypes::Decimal128Type; use datafusion_common::cast::{as_float32_array, as_float64_array}; #[test] @@ -328,4 +554,60 @@ mod test { assert_eq!(floats.value(3), 123.0); assert_eq!(floats.value(4), -321.0); } + + #[test] + fn test_compute_truncate_decimal128() { + // number 12.3456 (scale 4) truncated to 3 places = 12.345 + assert_eq!( + compute_truncate_decimal::(123_456, 4, 3), + 123_450 + ); + // number 12.3456 (scale 4) truncated to 1 place = 12.3 + assert_eq!( + compute_truncate_decimal::(123_456, 4, 1), + 123_000 + ); + + // requesting more places = no change + assert_eq!( + compute_truncate_decimal::(123_456, 4, 10), + 123_456 + ); + + // truncating to 0 places = whole number 12 + assert_eq!( + compute_truncate_decimal::(123_456, 4, 0), + 120_000 + ); + + // number 12.3456 (scale 2) truncated to -1 places = 10 + assert_eq!( + compute_truncate_decimal::(123_456, 4, -1), + 100_000 + ); + + // number 12.3456 (scale 2) truncated to -3 places = 0 + assert_eq!( + compute_truncate_decimal::(123_456, 4, -3), + 0 + ); + + // number 1234.56 (scale 2) truncated to -3 places = 1000 + assert_eq!( + compute_truncate_decimal::(123_456, 2, -3), + 100_000 + ); + + // out of scale + assert_eq!( + compute_truncate_decimal::(123_456, 4, -900), + 0 + ); + + // truncation rounds towards zero: -12.3456 = -12.345 + assert_eq!( + compute_truncate_decimal::(-123_456, 4, 3), + -123_450 + ); + } } diff --git a/datafusion/functions/src/regex/mod.rs b/datafusion/functions/src/regex/mod.rs index 75cc5d9514cbd..67241712038b9 100644 --- a/datafusion/functions/src/regex/mod.rs +++ b/datafusion/functions/src/regex/mod.rs @@ -146,6 +146,22 @@ where Ok(result) } +/// Maps `start`, a 1-based character position, to a byte offset in `value`. +/// Positions `1..=n` (for an `n`-character string) map to the corresponding +/// character's first byte; position `n + 1`, the end of the string, maps to +/// `value.len()`. Returns `None` for larger positions. Callers must validate +/// `start >= 1`. +pub(crate) fn start_to_byte_offset(value: &str, start: i64) -> Option { + // If `start - 1` does not fit in `usize`, it is necessarily past the end + // of the string. + let start_index = usize::try_from(start - 1).ok()?; + value + .char_indices() + .map(|(offset, _)| offset) + .chain(std::iter::once(value.len())) + .nth(start_index) +} + pub fn compile_regex(regex: &str, flags: Option<&str>) -> Result { let pattern = match flags { None | Some("") => regex.to_string(), @@ -164,3 +180,32 @@ pub fn compile_regex(regex: &str, flags: Option<&str>) -> Result select regexp_count('abcAbAbc', 'abc', 2, 'i'); +---------------------------------------------------------------+ @@ -49,16 +49,11 @@ use std::sync::Arc; standard_argument(name = "regexp", prefix = "Regular"), argument( name = "start", - description = "- **start**: Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function." + description = "Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function." ), argument( name = "flags", - description = r#"Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: - - **i**: case-insensitive: letters match both upper and lower case - - **m**: multi-line mode: ^ and $ match begin/end of line - - **s**: allow . to match \n - - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used - - **U**: swap the meaning of x* and x*?"# + description = r#"Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags."# ) )] #[derive(Debug, PartialEq, Eq, Hash)] @@ -267,48 +262,37 @@ fn regexp_count_inner<'a, S>( where S: StringArrayType<'a>, { - let (regex_scalar, is_regex_scalar) = if is_regex_scalar || regex_array.len() == 1 { - ( - (!regex_array.is_null(0)).then(|| regex_array.value(0)), - true, - ) + // Treat single-element arrays as scalars, broadcast to every row. An + // absent optional argument behaves like a scalar set to its default. + let is_regex_scalar = is_regex_scalar || regex_array.len() == 1; + let is_start_scalar = + start_array.is_none_or(|array| is_start_scalar || array.len() == 1); + let is_flags_scalar = + flags_array.is_none_or(|array| is_flags_scalar || array.len() == 1); + + // A NULL in any scalar argument produces a NULL result for every row + if (is_regex_scalar && regex_array.is_null(0)) + || (is_start_scalar && start_array.is_some_and(|array| array.is_null(0))) + || (is_flags_scalar && flags_array.is_some_and(|array| array.is_null(0))) + { + return Ok(Arc::new(Int64Array::new_null(values.len()))); + } + + let regex_scalar = is_regex_scalar.then(|| regex_array.value(0)); + // An absent `start` defaults to 1 + let start_scalar = + is_start_scalar.then(|| start_array.map_or(1, |array| array.value(0))); + // A `flags_scalar` of None means no flags were supplied + let flags_scalar = if is_flags_scalar { + flags_array.map(|array| array.value(0)) } else { - (None, false) + None }; - let (start_array, start_scalar, is_start_scalar) = - if let Some(start_array) = start_array { - if is_start_scalar || start_array.len() == 1 { - (None, Some(start_array.value(0)), true) - } else { - (Some(start_array), None, false) - } - } else { - (None, Some(1), true) - }; - - let (flags_array, flags_scalar, is_flags_scalar) = - if let Some(flags_array) = flags_array { - if is_flags_scalar || flags_array.len() == 1 { - (None, Some(flags_array.value(0)), true) - } else { - (Some(flags_array), None, false) - } - } else { - (None, None, true) - }; - let mut regex_cache = HashMap::new(); - match (is_regex_scalar, is_start_scalar, is_flags_scalar) { - (true, true, true) => { - let regex = match regex_scalar { - None => { - return Ok(Arc::new(Int64Array::from(vec![0; values.len()]))); - } - Some(regex) => regex, - }; - + match (regex_scalar, is_start_scalar, is_flags_scalar) { + (Some(regex), true, true) => { let pattern = compile_regex(regex, flags_scalar)?; Ok(Arc::new( @@ -318,14 +302,7 @@ where .collect::>()?, )) } - (true, true, false) => { - let regex = match regex_scalar { - None => { - return Ok(Arc::new(Int64Array::from(vec![0; values.len()]))); - } - Some(regex) => regex, - }; - + (Some(regex), true, false) => { let flags_array = flags_array.unwrap(); if values.len() != flags_array.len() { return Err(ArrowError::ComputeError(format!( @@ -340,21 +317,21 @@ where .iter() .zip(flags_array.iter()) .map(|(value, flags)| { - let pattern = - compile_and_cache_regex(regex, flags, &mut regex_cache)?; + let Some(flags) = flags else { + return Ok(None); + }; + + let pattern = compile_and_cache_regex( + regex, + Some(flags), + &mut regex_cache, + )?; count_matches(value, pattern, start_scalar) }) .collect::>()?, )) } - (true, false, true) => { - let regex = match regex_scalar { - None => { - return Ok(Arc::new(Int64Array::from(vec![0; values.len()]))); - } - Some(regex) => regex, - }; - + (Some(regex), false, true) => { let pattern = compile_regex(regex, flags_scalar)?; let start_array = start_array.unwrap(); @@ -367,14 +344,7 @@ where .collect::>()?, )) } - (true, false, false) => { - let regex = match regex_scalar { - None => { - return Ok(Arc::new(Int64Array::from(vec![0; values.len()]))); - } - Some(regex) => regex, - }; - + (Some(regex), false, false) => { let flags_array = flags_array.unwrap(); if values.len() != flags_array.len() { return Err(ArrowError::ComputeError(format!( @@ -391,15 +361,19 @@ where flags_array.iter() ) .map(|(value, start, flags)| { + let Some(flags) = flags else { + return Ok(None); + }; + let pattern = - compile_and_cache_regex(regex, flags, &mut regex_cache)?; + compile_and_cache_regex(regex, Some(flags), &mut regex_cache)?; count_matches(value, pattern, start) }) .collect::>()?, )) } - (false, true, true) => { + (None, true, true) => { if values.len() != regex_array.len() { return Err(ArrowError::ComputeError(format!( "regex_array must be the same length as values array; got {} and {}", @@ -413,9 +387,8 @@ where .iter() .zip(regex_array.iter()) .map(|(value, regex)| { - let regex = match regex { - None => return Ok(0), - Some(regex) => regex, + let Some(regex) = regex else { + return Ok(None); }; let pattern = compile_and_cache_regex( @@ -428,7 +401,7 @@ where .collect::>()?, )) } - (false, true, false) => { + (None, true, false) => { if values.len() != regex_array.len() { return Err(ArrowError::ComputeError(format!( "regex_array must be the same length as values array; got {} and {}", @@ -449,20 +422,22 @@ where Ok(Arc::new( izip!(values.iter(), regex_array.iter(), flags_array.iter()) .map(|(value, regex, flags)| { - let regex = match regex { - None => return Ok(0), - Some(regex) => regex, + let (Some(regex), Some(flags)) = (regex, flags) else { + return Ok(None); }; - let pattern = - compile_and_cache_regex(regex, flags, &mut regex_cache)?; + let pattern = compile_and_cache_regex( + regex, + Some(flags), + &mut regex_cache, + )?; count_matches(value, pattern, start_scalar) }) .collect::>()?, )) } - (false, false, true) => { + (None, false, true) => { if values.len() != regex_array.len() { return Err(ArrowError::ComputeError(format!( "regex_array must be the same length as values array; got {} and {}", @@ -483,9 +458,8 @@ where Ok(Arc::new( izip!(values.iter(), regex_array.iter(), start_array.iter()) .map(|(value, regex, start)| { - let regex = match regex { - None => return Ok(0), - Some(regex) => regex, + let Some(regex) = regex else { + return Ok(None); }; let pattern = compile_and_cache_regex( @@ -498,7 +472,7 @@ where .collect::>()?, )) } - (false, false, false) => { + (None, false, false) => { if values.len() != regex_array.len() { return Err(ArrowError::ComputeError(format!( "regex_array must be the same length as values array; got {} and {}", @@ -533,13 +507,12 @@ where flags_array.iter() ) .map(|(value, regex, start, flags)| { - let regex = match regex { - None => return Ok(0), - Some(regex) => regex, + let (Some(regex), Some(flags)) = (regex, flags) else { + return Ok(None); }; let pattern = - compile_and_cache_regex(regex, flags, &mut regex_cache)?; + compile_and_cache_regex(regex, Some(flags), &mut regex_cache)?; count_matches(value, pattern, start) }) .collect::>()?, @@ -552,45 +525,23 @@ fn count_matches( value: Option<&str>, pattern: &Regex, start: Option, -) -> Result { - let value = match value { - None => return Ok(0), - Some(value) => value, +) -> Result, ArrowError> { + // A NULL value or start position produces a NULL result. + let (Some(value), Some(start)) = (value, start) else { + return Ok(None); }; - if let Some(start) = start { - if start < 1 { - return Err(ArrowError::ComputeError( - "regexp_count() requires start to be 1 based".to_string(), - )); - } - - let char_len = value.chars().count(); - let start_index = (start as usize).saturating_sub(1); - - if start_index > char_len { - return Ok(0); - } - - // Find the byte offset for the start position (1-based character index) - let byte_offset = if start_index == char_len { - value.len() - } else { - value - .char_indices() - .nth(start_index) - .map(|(idx, _)| idx) - .unwrap_or(value.len()) - }; - - // Use string slicing instead of collecting chars into a new String - let find_slice = &value[byte_offset..]; - let count = pattern.find_iter(find_slice).count(); - Ok(count as i64) - } else { - let count = pattern.find_iter(value).count(); - Ok(count as i64) + if start < 1 { + return Err(ArrowError::ComputeError( + "regexp_count() requires start to be 1 based".to_string(), + )); } + + let Some(byte_offset) = start_to_byte_offset(value, start) else { + return Ok(Some(0)); + }; + let count = pattern.find_iter(&value[byte_offset..]).count(); + Ok(Some(count as i64)) } #[cfg(test)] @@ -625,6 +576,24 @@ mod tests { test_case_sensitive_regexp_count_array_complex::(); test_case_regexp_count_cache_check::>(); + + test_regexp_count_null_scalars(); + + test_regexp_count_null_array_rows::>(); + test_regexp_count_null_array_rows::>(); + test_regexp_count_null_array_rows::(); + + test_regexp_count_null_start_array::>(); + test_regexp_count_null_start_array::>(); + test_regexp_count_null_start_array::(); + + test_regexp_count_null_flags_array::>(); + test_regexp_count_null_flags_array::>(); + test_regexp_count_null_flags_array::(); + + test_regexp_count_null_scalar_regex_array_values::>(); + test_regexp_count_null_scalar_regex_array_values::>(); + test_regexp_count_null_scalar_regex_array_values::(); } fn regexp_count_with_scalar_values(args: &[ScalarValue]) -> Result { @@ -988,6 +957,129 @@ mod tests { assert_eq!(re.as_ref(), &expected); } + fn test_regexp_count_null_scalars() { + // A NULL in any scalar argument produces a NULL result. + let cases: Vec> = vec![ + vec![ScalarValue::Utf8(None), ScalarValue::Utf8(None)], + vec![ + ScalarValue::Utf8(None), + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Int64(Some(1)), + ScalarValue::Utf8(Some("i".to_string())), + ], + vec![ + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Utf8(None), + ScalarValue::Int64(Some(1)), + ScalarValue::Utf8(Some("i".to_string())), + ], + vec![ + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Int64(None), + ScalarValue::Utf8(Some("i".to_string())), + ], + vec![ + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Int64(Some(1)), + ScalarValue::Utf8(None), + ], + ]; + + for args in cases { + let re = regexp_count_with_scalar_values(&args); + match re { + Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => { + assert_eq!(v, None, "regexp_count null scalar test failed"); + } + _ => panic!("Unexpected result"), + } + } + } + + fn test_regexp_count_null_array_rows() + where + A: From>> + Array + 'static, + { + let values = A::from(vec![ + None, + Some("abc"), + Some("abc"), + Some("abc"), + Some("abc"), + ]); + let regex = A::from(vec![ + Some("abc"), + None, + Some("abc"), + Some("abc"), + Some("abc"), + ]); + let start = Int64Array::from(vec![Some(1), Some(1), None, Some(1), Some(1)]); + let flags = A::from(vec![Some("i"), Some("i"), Some("i"), None, Some("i")]); + + let expected = Int64Array::from(vec![None, None, None, None, Some(1)]); + + let re = regexp_count_func(&[ + Arc::new(values), + Arc::new(regex), + Arc::new(start), + Arc::new(flags), + ]) + .unwrap(); + assert_eq!(re.as_ref(), &expected); + } + + fn test_regexp_count_null_start_array() + where + A: From> + Array + 'static, + { + let values = A::from(vec!["abc", "abcb"]); + let regex = A::from(vec!["b"]); + let start = Int64Array::from(vec![Some(1), None]); + + let expected = Int64Array::from(vec![Some(1), None]); + + let re = regexp_count_func(&[Arc::new(values), Arc::new(regex), Arc::new(start)]) + .unwrap(); + assert_eq!(re.as_ref(), &expected); + } + + fn test_regexp_count_null_flags_array() + where + A: From> + From>> + Array + 'static, + { + let values: A = vec!["aB", "aB"].into(); + let regex: A = vec!["b"].into(); + let start = Int64Array::from(vec![1]); + let flags: A = vec![None, Some("i")].into(); + + let expected = Int64Array::from(vec![None, Some(1)]); + + let re = regexp_count_func(&[ + Arc::new(values), + Arc::new(regex), + Arc::new(start), + Arc::new(flags), + ]) + .unwrap(); + assert_eq!(re.as_ref(), &expected); + } + + fn test_regexp_count_null_scalar_regex_array_values() + where + A: From> + From>> + Array + 'static, + { + let values: A = vec!["abc", "abcabc"].into(); + let regex: A = vec![Option::<&str>::None].into(); + + let expected = Int64Array::from(vec![None::, None]); + + let re = regexp_count_func(&[Arc::new(values), Arc::new(regex)]).unwrap(); + assert_eq!(re.as_ref(), &expected); + } + fn test_case_regexp_count_cache_check() where A: From> + Array + 'static, diff --git a/datafusion/functions/src/regex/regexpinstr.rs b/datafusion/functions/src/regex/regexpinstr.rs index d46e4452dbab1..96152297fbc87 100644 --- a/datafusion/functions/src/regex/regexpinstr.rs +++ b/datafusion/functions/src/regex/regexpinstr.rs @@ -16,8 +16,9 @@ // under the License. use arrow::array::{ - Array, ArrayRef, AsArray, Datum, Int64Array, PrimitiveArray, StringArrayType, + Array, ArrayRef, AsArray, Datum, Int64Array, Int64Builder, StringArrayType, }; +use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, Int64Type}; use arrow::datatypes::{ DataType::Int64, DataType::LargeUtf8, DataType::Utf8, DataType::Utf8View, @@ -29,12 +30,12 @@ use datafusion_expr::{ TypeSignature::Exact, TypeSignature::Uniform, Volatility, }; use datafusion_macros::user_doc; -use itertools::izip; use regex::Regex; use std::collections::HashMap; +use std::collections::hash_map::Entry; use std::sync::Arc; -use crate::regex::compile_and_cache_regex; +use crate::regex::{compile_regex, start_to_byte_offset}; #[user_doc( doc_section(label = "Regular Expression Functions"), @@ -52,20 +53,15 @@ use crate::regex::compile_and_cache_regex; standard_argument(name = "regexp", prefix = "Regular"), argument( name = "start", - description = "- **start**: Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function. Defaults to 1" + description = "Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function. Defaults to 1" ), argument( name = "N", - description = "- **N**: Optional The N-th occurrence of pattern to find. Defaults to 1 (first match). Can be a constant, column, or function." + description = "Optional The N-th occurrence of pattern to find. Defaults to 1 (first match). Can be a constant, column, or function." ), argument( name = "flags", - description = r#"Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: - - **i**: case-insensitive: letters match both upper and lower case - - **m**: multi-line mode: ^ and $ match begin/end of line - - **s**: allow . to match \n - - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used - - **U**: swap the meaning of x* and x*?"# + description = r#"Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags."# ), argument( name = "subexpr", @@ -240,7 +236,7 @@ fn regexp_instr( ®ex_array.as_string::(), start_array.map(|start| start.as_primitive::()), nth_array.map(|nth| nth.as_primitive::()), - Some(flags_array.as_string::()), + Some(&flags_array.as_string::()), subexpr_array.map(|subexpr| subexpr.as_primitive::()), ), (LargeUtf8, LargeUtf8, None) => regexp_instr_inner( @@ -256,7 +252,7 @@ fn regexp_instr( ®ex_array.as_string::(), start_array.map(|start| start.as_primitive::()), nth_array.map(|nth| nth.as_primitive::()), - Some(flags_array.as_string::()), + Some(&flags_array.as_string::()), subexpr_array.map(|subexpr| subexpr.as_primitive::()), ), (Utf8View, Utf8View, None) => regexp_instr_inner( @@ -272,7 +268,7 @@ fn regexp_instr( ®ex_array.as_string_view(), start_array.map(|start| start.as_primitive::()), nth_array.map(|nth| nth.as_primitive::()), - Some(flags_array.as_string_view()), + Some(&flags_array.as_string_view()), subexpr_array.map(|subexpr| subexpr.as_primitive::()), ), _ => Err(ArrowError::ComputeError( @@ -286,120 +282,97 @@ fn regexp_instr_inner<'a, S>( regex_array: &S, start_array: Option<&Int64Array>, nth_array: Option<&Int64Array>, - flags_array: Option, + flags_array: Option<&S>, subexp_array: Option<&Int64Array>, ) -> Result where S: StringArrayType<'a>, { let len = values.len(); + let mut regex_cache = RegexCache::default(); + let mut result = Int64Builder::with_capacity(len); + + // A NULL in any argument produces a NULL result + let nulls = NullBuffer::union_many([ + values.nulls(), + regex_array.nulls(), + start_array.and_then(|array| array.nulls()), + nth_array.and_then(|array| array.nulls()), + flags_array.and_then(|array| array.nulls()), + subexp_array.and_then(|array| array.nulls()), + ]); + + for i in 0..len { + if nulls.as_ref().is_some_and(|nulls| nulls.is_null(i)) { + result.append_null(); + continue; + } - let default_start_array = PrimitiveArray::::from(vec![1; len]); - let start_array = start_array.unwrap_or(&default_start_array); - let start_input: Vec = (0..start_array.len()) - .map(|i| start_array.value(i)) // handle nulls as 0 - .collect(); - - let default_nth_array = PrimitiveArray::::from(vec![1; len]); - let nth_array = nth_array.unwrap_or(&default_nth_array); - let nth_input: Vec = (0..nth_array.len()) - .map(|i| nth_array.value(i)) // handle nulls as 0 - .collect(); - - let flags_input = match flags_array { - Some(flags) => flags.iter().collect(), - None => vec![None; len], - }; + let value = values.value(i); + let regex = regex_array.value(i); + let flags = flags_array.map(|array| array.value(i)); + let pattern = regex_cache.get_or_compile(regex, flags)?; - let default_subexp_array = PrimitiveArray::::from(vec![0; len]); - let subexp_array = subexp_array.unwrap_or(&default_subexp_array); - let subexp_input: Vec = (0..subexp_array.len()) - .map(|i| subexp_array.value(i)) // handle nulls as 0 - .collect(); - - let mut regex_cache = HashMap::new(); - - let result: Result>, ArrowError> = izip!( - values.iter(), - regex_array.iter(), - start_input.iter(), - nth_input.iter(), - flags_input.iter(), - subexp_input.iter() - ) - .map(|(value, regex, start, nth, flags, subexp)| match regex { - None => Ok(None), - Some("") => Ok(Some(0)), - Some(regex) => get_index( - value, - regex, - *start, - *nth, - *subexp, - *flags, - &mut regex_cache, - ), - }) - .collect(); - Ok(Arc::new(Int64Array::from(result?))) -} + // The defaults apply when the optional argument was not supplied. + let start = start_array.map_or(1, |array| array.value(i)); + let nth = nth_array.map_or(1, |array| array.value(i)); + let subexp = subexp_array.map_or(0, |array| array.value(i)); -fn handle_subexp( - pattern: &Regex, - search_slice: &str, - subexpr: i64, - value: &str, - byte_start_offset: usize, -) -> Result, ArrowError> { - if let Some(captures) = pattern.captures(search_slice) - && let Some(matched) = captures.get(subexpr as usize) - { - // Convert byte offset relative to search_slice back to 1-based character offset - // relative to the original `value` string. - let start_char_offset = - value[..byte_start_offset + matched.start()].chars().count() as i64 + 1; - return Ok(Some(start_char_offset)); + result.append_value(get_index(value, pattern, start, nth, subexp)?); } - Ok(Some(0)) // Return 0 if the subexpression was not found + + Ok(Arc::new(result.finish())) } -fn get_nth_match( - pattern: &Regex, - search_slice: &str, - n: i64, - byte_start_offset: usize, - value: &str, -) -> Result, ArrowError> { - if let Some(mat) = pattern.find_iter(search_slice).nth((n - 1) as usize) { - // Convert byte offset relative to search_slice back to 1-based character offset - // relative to the original `value` string. - let match_start_byte_offset = byte_start_offset + mat.start(); - let match_start_char_offset = - value[..match_start_byte_offset].chars().count() as i64 + 1; - Ok(Some(match_start_char_offset)) - } else { - Ok(Some(0)) // Return 0 if the N-th match was not found +/// Compiles the patterns seen so far, keyed by `(pattern, flags)`. +/// +/// Patterns are addressed by index rather than by reference so that `last` can +/// memoize the previous row's pattern without holding a borrow of `indices` +/// across rows. A literal pattern yields the same string on every row, so that +/// memo means the common case never hashes a key. +#[derive(Default)] +struct RegexCache<'a> { + compiled: Vec, + indices: HashMap<(&'a str, Option<&'a str>), usize>, + last: Option<((&'a str, Option<&'a str>), usize)>, +} + +impl<'a> RegexCache<'a> { + fn get_or_compile( + &mut self, + regex: &'a str, + flags: Option<&'a str>, + ) -> Result<&Regex, ArrowError> { + let key = (regex, flags); + let index = match self.last { + Some((last_key, index)) if last_key == key => index, + _ => { + let index = match self.indices.entry(key) { + Entry::Occupied(entry) => *entry.get(), + Entry::Vacant(entry) => { + self.compiled.push(compile_regex(regex, flags)?); + *entry.insert(self.compiled.len() - 1) + } + }; + self.last = Some((key, index)); + index + } + }; + Ok(&self.compiled[index]) } } -fn get_index<'strings, 'cache>( - value: Option<&str>, - pattern: &'strings str, + +/// Returns the 1-based character position of the `n`-th match of `pattern` in +/// `value`, or 0 if there is no such match. The search begins at the 1-based +/// character position `start`. A positive `subexpr` selects that capture group +/// of the first match instead of the `n`-th match. +fn get_index( + value: &str, + pattern: &Regex, start: i64, n: i64, subexpr: i64, - flags: Option<&'strings str>, - regex_cache: &'cache mut HashMap<(&'strings str, Option<&'strings str>), Regex>, -) -> Result, ArrowError> -where - 'strings: 'cache, -{ - let value = match value { - None => return Ok(None), - Some("") => return Ok(Some(0)), - Some(value) => value, - }; - let pattern: &Regex = compile_and_cache_regex(pattern, flags, regex_cache)?; - // println!("get_index: value = {}, pattern = {}, start = {}, n = {}, subexpr = {}, flags = {:?}", value, pattern, start, n, subexpr, flags); +) -> Result { if start < 1 { return Err(ArrowError::ComputeError( "regexp_instr() requires start to be 1-based".to_string(), @@ -412,31 +385,30 @@ where )); } - // --- Simplified byte_start_offset calculation --- - let total_chars = value.chars().count() as i64; - let byte_start_offset: usize = if start > total_chars { - // If start is beyond the total characters, it means we start searching - // after the string effectively. No matches possible. - return Ok(Some(0)); - } else { - // Get the byte offset for the (start - 1)-th character (0-based) - value - .char_indices() - .nth((start - 1) as usize) - .map(|(idx, _)| idx) - .unwrap_or(0) // Should not happen if start is valid and <= total_chars + let Some(byte_start_offset) = start_to_byte_offset(value, start) else { + return Ok(0); }; - // --- End simplified calculation --- - let search_slice = &value[byte_start_offset..]; - // Handle subexpression capturing first, as it takes precedence - if subexpr > 0 { - return handle_subexp(pattern, search_slice, subexpr, value, byte_start_offset); - } + // A subexpression, when requested, takes precedence over the N-th match. + let match_start = if subexpr > 0 { + pattern + .captures(search_slice) + .and_then(|captures| captures.get(subexpr as usize)) + .map(|matched| matched.start()) + } else { + // `n` is 1-based, `nth` is 0-based. + pattern + .find_iter(search_slice) + .nth((n - 1) as usize) + .map(|matched| matched.start()) + }; - // Use nth to get the N-th match (n is 1-based, nth is 0-based) - get_nth_match(pattern, search_slice, n, byte_start_offset, value) + // Convert the byte offset within `search_slice` back to a 1-based character + // offset within `value`. + Ok(match_start.map_or(0, |offset| { + value[..byte_start_offset + offset].chars().count() as i64 + 1 + })) } #[cfg(test)] @@ -445,6 +417,7 @@ mod tests { use arrow::array::{GenericStringArray, StringViewArray}; use arrow::datatypes::Field; use datafusion_common::config::ConfigOptions; + use itertools::izip; #[test] fn test_regexp_instr() { test_case_sensitive_regexp_instr_nulls(); @@ -464,6 +437,20 @@ mod tests { test_case_sensitive_regexp_instr_array_nth::>(); test_case_sensitive_regexp_instr_array_nth::>(); test_case_sensitive_regexp_instr_array_nth::(); + + test_case_sensitive_regexp_instr_empty_pattern::>(); + test_case_sensitive_regexp_instr_empty_pattern::>(); + test_case_sensitive_regexp_instr_empty_pattern::(); + + test_case_sensitive_regexp_instr_zero_width_pattern::>(); + test_case_sensitive_regexp_instr_zero_width_pattern::>(); + test_case_sensitive_regexp_instr_zero_width_pattern::(); + + test_regexp_instr_null_scalar_args(); + + test_regexp_instr_null_array_rows::>(); + test_regexp_instr_null_array_rows::>(); + test_regexp_instr_null_array_rows::(); } fn regexp_instr_with_scalar_values(args: &[ScalarValue]) -> Result { @@ -492,7 +479,7 @@ mod tests { fn test_case_sensitive_regexp_instr_nulls() { let v = ""; let r = ""; - let expected = 0; + let expected = 1; let regex_sv = ScalarValue::Utf8(Some(r.to_string())); let re = regexp_instr_with_scalar_values(&[v.to_string().into(), regex_sv]); // let res_exp = re.unwrap(); @@ -502,6 +489,29 @@ mod tests { } _ => panic!("Unexpected result"), } + + for (value, regex) in [ + ( + ScalarValue::Utf8(None), + ScalarValue::Utf8(Some(String::new())), + ), + ( + ScalarValue::LargeUtf8(None), + ScalarValue::LargeUtf8(Some(String::new())), + ), + ( + ScalarValue::Utf8View(None), + ScalarValue::Utf8View(Some(String::new())), + ), + ] { + let re = regexp_instr_with_scalar_values(&[value, regex]); + match re { + Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => { + assert_eq!(v, None, "regexp_instr NULL scalar test failed"); + } + _ => panic!("Unexpected result"), + } + } } fn test_case_sensitive_regexp_instr_scalar() { let values = [ @@ -762,6 +772,126 @@ mod tests { }); } + fn test_regexp_instr_null_scalar_args() { + // A NULL in any argument produces a NULL result + let cases: Vec> = vec![ + // NULL start + vec![ + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Utf8(Some("b".to_string())), + ScalarValue::Int64(None), + ], + // NULL N + vec![ + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Utf8(Some("b".to_string())), + ScalarValue::Int64(Some(1)), + ScalarValue::Int64(None), + ], + // NULL flags + vec![ + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Utf8(Some("b".to_string())), + ScalarValue::Int64(Some(1)), + ScalarValue::Int64(Some(1)), + ScalarValue::Utf8(None), + ], + // NULL subexpr + vec![ + ScalarValue::Utf8(Some("abc".to_string())), + ScalarValue::Utf8(Some("(b)".to_string())), + ScalarValue::Int64(Some(1)), + ScalarValue::Int64(Some(1)), + ScalarValue::Utf8(Some("i".to_string())), + ScalarValue::Int64(None), + ], + ]; + + for args in cases { + let re = regexp_instr_with_scalar_values(&args); + match re { + Ok(ColumnarValue::Scalar(ScalarValue::Int64(v))) => { + assert_eq!(v, None, "regexp_instr null scalar test failed"); + } + _ => panic!("Unexpected result"), + } + } + } + + fn test_regexp_instr_null_array_rows() + where + A: From>> + Array + 'static, + { + let values = A::from(vec![ + None, + Some("abc"), + Some("abc"), + Some("abc"), + Some("abc"), + Some("abc"), + Some("abc"), + ]); + let regex = A::from(vec![ + Some("b"), + None, + Some("b"), + Some("b"), + Some("b"), + Some("(b)"), + Some("b"), + ]); + let start = Int64Array::from(vec![ + Some(1), + Some(1), + None, + Some(1), + Some(1), + Some(1), + Some(1), + ]); + let nth = Int64Array::from(vec![ + Some(1), + Some(1), + Some(1), + None, + Some(1), + Some(1), + Some(1), + ]); + let flags = A::from(vec![ + Some(""), + Some(""), + Some(""), + Some(""), + None, + Some("i"), + Some(""), + ]); + let subexp = Int64Array::from(vec![ + Some(0), + Some(0), + Some(0), + Some(0), + Some(0), + None, + Some(0), + ]); + + let expected = + Int64Array::from(vec![None, None, None, None, None, None, Some(2)]); + + let re = regexp_instr_func(&[ + Arc::new(values), + Arc::new(regex), + Arc::new(start), + Arc::new(nth), + Arc::new(flags), + Arc::new(subexp), + ]) + .unwrap(); + assert_eq!(re.as_ref(), &expected); + } + fn test_case_sensitive_regexp_instr_array() where A: From> + Array + 'static, @@ -813,4 +943,38 @@ mod tests { .unwrap(); assert_eq!(re.as_ref(), &expected); } + + fn test_case_sensitive_regexp_instr_empty_pattern() + where + A: From> + Array + 'static, + { + let values = A::from(vec!["abc", "", "abc", "abc", "😀"]); + let regex = A::from(vec!["", "", "", "", ""]); + let start = Int64Array::from(vec![1, 1, 4, 5, 1]); + let nth = Int64Array::from(vec![1, 1, 1, 1, 2]); + let expected = Int64Array::from(vec![1, 1, 4, 0, 2]); + + let re = regexp_instr_func(&[ + Arc::new(values), + Arc::new(regex), + Arc::new(start), + Arc::new(nth), + ]) + .unwrap(); + assert_eq!(re.as_ref(), &expected); + } + + fn test_case_sensitive_regexp_instr_zero_width_pattern() + where + A: From> + Array + 'static, + { + let values = A::from(vec!["abc"]); + let regex = A::from(vec!["x*"]); + let start = Int64Array::from(vec![4]); + let expected = Int64Array::from(vec![4]); + + let re = regexp_instr_func(&[Arc::new(values), Arc::new(regex), Arc::new(start)]) + .unwrap(); + assert_eq!(re.as_ref(), &expected); + } } diff --git a/datafusion/functions/src/regex/regexplike.rs b/datafusion/functions/src/regex/regexplike.rs index 56754b13db227..e7b31b767a4b0 100644 --- a/datafusion/functions/src/regex/regexplike.rs +++ b/datafusion/functions/src/regex/regexplike.rs @@ -61,12 +61,7 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo standard_argument(name = "regexp", prefix = "Regular"), argument( name = "flags", - description = r#"Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: - - **i**: case-insensitive: letters match both upper and lower case - - **m**: multi-line mode: ^ and $ match begin/end of line - - **s**: allow . to match \n - - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used - - **U**: swap the meaning of x* and x*?"# + description = r#"Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags."# ) )] #[derive(Debug, PartialEq, Eq, Hash)] diff --git a/datafusion/functions/src/regex/regexpmatch.rs b/datafusion/functions/src/regex/regexpmatch.rs index 34153d9c8ab96..ce7c437a54520 100644 --- a/datafusion/functions/src/regex/regexpmatch.rs +++ b/datafusion/functions/src/regex/regexpmatch.rs @@ -16,7 +16,7 @@ // under the License. //! Regex expressions -use arrow::array::{Array, ArrayRef, AsArray}; +use arrow::array::{Array, ArrayRef, AsArray, Datum}; use arrow::compute::kernels::regexp; use arrow::datatypes::DataType; use arrow::datatypes::Field; @@ -57,12 +57,7 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo ), argument( name = "flags", - description = r#"Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: - - **i**: case-insensitive: letters match both upper and lower case - - **m**: multi-line mode: ^ and $ match begin/end of line - - **s**: allow . to match \n - - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used - - **U**: swap the meaning of x* and x*?"# + description = r#"Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags."# ) )] #[derive(Debug, PartialEq, Eq, Hash)] @@ -116,6 +111,14 @@ impl ScalarUDFImpl for RegexpMatchFunc { fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { let args = &args.args; + + // A literal pattern is the common case, and handing it to the kernel as + // a scalar lets the regex be compiled once for the whole array. Any + // other argument shape falls through to the general path below. + if let Some(result) = regexp_match_scalar_pattern(args)? { + return Ok(ColumnarValue::Array(result)); + } + let len = args .iter() .fold(Option::::None, |acc, arg| match arg { @@ -145,6 +148,61 @@ impl ScalarUDFImpl for RegexpMatchFunc { } } +/// Runs `regexp_match` with the pattern (and flags, if given) passed to the +/// kernel as scalar [`Datum`]s, so the regex is compiled once for the whole +/// array. +/// +/// Applies when the values are an array, the pattern is a non-null scalar of +/// the same string type as the values, and the flags, if given, are a scalar of +/// that same type and are not the unsupported "global" flag. +/// +/// Returns `Ok(None)` for every other argument shape, leaving the caller's +/// general path to materialize each argument as an array, zip the rows, and +/// raise whatever error the shape warrants. +fn regexp_match_scalar_pattern(args: &[ColumnarValue]) -> Result> { + let (values, pattern, flags) = match args { + [values, pattern] => (values, pattern, None), + [values, pattern, flags] => (values, pattern, Some(flags)), + _ => return Ok(None), + }; + + let (ColumnarValue::Array(values), ColumnarValue::Scalar(pattern)) = + (values, pattern) + else { + return Ok(None); + }; + let flags = match flags { + // An array of flags has to be zipped with the values row by row. + Some(ColumnarValue::Array(_)) => return Ok(None), + Some(ColumnarValue::Scalar(flags)) => Some(flags), + None => None, + }; + + // The kernel requires the values, the pattern and the flags to share one + // string type. + let value_type = values.data_type(); + + if !matches!(pattern.try_as_str(), Some(Some(_))) + || &pattern.data_type() != value_type + || flags.is_some_and(|flags| { + flags.try_as_str() == Some(Some("g")) || &flags.data_type() != value_type + }) + { + return Ok(None); + } + + let pattern = pattern.to_scalar()?; + let flags = flags.map(ScalarValue::to_scalar).transpose()?; + + regexp::regexp_match( + values, + &pattern, + flags.as_ref().map(|flags| flags as &dyn Datum), + ) + .map(Some) + .map_err(|e| arrow_datafusion_err!(e)) +} + pub fn regexp_match(args: &[ArrayRef]) -> Result { match args.len() { 2 => regexp::regexp_match(&args[0], &args[1], None) @@ -257,4 +315,71 @@ mod tests { "Error during planning: regexp_match() does not support the \"global\" option" ); } + + /// The literal-pattern fast path must agree with the general path that + /// zips a pattern array with the values, for every argument shape. + #[test] + fn test_scalar_pattern_matches_array_pattern() { + use super::{RegexpMatchFunc, ScalarValue}; + use arrow::array::{Array, ArrayRef}; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::{ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl}; + + let values = Arc::new(StringArray::from(vec![ + Some("abc"), + Some("ABC"), + None, + Some(""), + Some("a-b-c"), + ])) as ArrayRef; + + for pattern in ["([a-z])(b)?", "^(A)", "no-match", "", "[a-z]+"] { + for flags in [None, Some("i")] { + let mut scalar_args = vec![ + ColumnarValue::Array(Arc::clone(&values)), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(pattern.to_string()))), + ]; + let mut array_args = vec![ + Arc::clone(&values), + Arc::new(StringArray::from(vec![pattern; values.len()])) as ArrayRef, + ]; + if let Some(flags) = flags { + scalar_args.push(ColumnarValue::Scalar(ScalarValue::Utf8(Some( + flags.to_string(), + )))); + array_args + .push(Arc::new(StringArray::from(vec![flags; values.len()])) + as ArrayRef); + } + + let arg_fields = scalar_args + .iter() + .enumerate() + .map(|(idx, arg)| { + Field::new(format!("arg_{idx}"), arg.data_type(), true).into() + }) + .collect(); + let actual = RegexpMatchFunc::new() + .invoke_with_args(ScalarFunctionArgs { + args: scalar_args, + arg_fields, + number_rows: values.len(), + return_field: Field::new_list( + "f", + Field::new_list_field(DataType::Utf8, true), + true, + ) + .into(), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() + .to_array(values.len()) + .unwrap(); + + let expected = regexp_match(&array_args).unwrap(); + assert_eq!(&actual, &expected, "pattern={pattern:?} flags={flags:?}"); + } + } + } } diff --git a/datafusion/functions/src/regex/regexpreplace.rs b/datafusion/functions/src/regex/regexpreplace.rs index 215dd33324375..ec4afbad47d04 100644 --- a/datafusion/functions/src/regex/regexpreplace.rs +++ b/datafusion/functions/src/regex/regexpreplace.rs @@ -79,13 +79,7 @@ Additional examples can be found [here](https://github.com/apache/datafusion/blo ), argument( name = "flags", - description = r#"Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: -- **g**: (global) Search globally and don't return after the first match -- **i**: case-insensitive: letters match both upper and lower case -- **m**: multi-line mode: ^ and $ match begin/end of line -- **s**: allow . to match \n -- **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used -- **U**: swap the meaning of x* and x*?"# + description = r#"Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags."# ) )] #[derive(Debug, PartialEq, Eq, Hash)] diff --git a/datafusion/functions/src/string/ascii.rs b/datafusion/functions/src/string/ascii.rs index bb5a8d0125a70..db539a4d11719 100644 --- a/datafusion/functions/src/string/ascii.rs +++ b/datafusion/functions/src/string/ascii.rs @@ -15,13 +15,16 @@ // specific language governing permissions and limitations // under the License. +use crate::utils::transform_leaf_type_preserving_encoding; use arrow::array::{ArrayRef, AsArray, Int32Array, StringArrayType}; use arrow::datatypes::DataType; use arrow::error::ArrowError; use datafusion_common::types::logical_string; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue, internal_err}; -use datafusion_expr::{ColumnarValue, Documentation, TypeSignatureClass}; +use datafusion_expr::{ + ColumnarValue, Documentation, EncodingPreservation, TypeSignatureClass, +}; use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility}; use datafusion_expr_common::signature::Coercion; use datafusion_macros::user_doc; @@ -63,9 +66,10 @@ impl AsciiFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -81,8 +85,8 @@ impl ScalarUDFImpl for AsciiFunc { &self.signature } - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Int32) + fn return_type(&self, arg_types: &[DataType]) -> Result { + transform_leaf_type_preserving_encoding(&arg_types[0], &|_| Ok(DataType::Int32)) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -90,24 +94,7 @@ impl ScalarUDFImpl for AsciiFunc { match arg { ColumnarValue::Scalar(scalar) => { - if scalar.is_null() { - return Ok(ColumnarValue::Scalar(ScalarValue::Int32(None))); - } - - match scalar { - ScalarValue::Utf8(Some(s)) - | ScalarValue::LargeUtf8(Some(s)) - | ScalarValue::Utf8View(Some(s)) => { - let result = s.chars().next().map_or(0, |c| c as i32); - Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(result)))) - } - _ => { - internal_err!( - "Unexpected data type {:?} for function ascii", - scalar.data_type() - ) - } - } + Ok(ColumnarValue::Scalar(ascii_scalar(&scalar)?)) } ColumnarValue::Array(array) => Ok(ColumnarValue::Array(ascii(&[array])?)), } @@ -118,22 +105,71 @@ impl ScalarUDFImpl for AsciiFunc { } } +fn ascii_scalar(scalar: &ScalarValue) -> Result { + match scalar { + ScalarValue::Utf8(value) + | ScalarValue::LargeUtf8(value) + | ScalarValue::Utf8View(value) => { + Ok(ScalarValue::Int32(value.as_deref().map(first_char_code))) + } + ScalarValue::Dictionary(key_type, value) => Ok(ScalarValue::Dictionary( + key_type.clone(), + Box::new(ascii_scalar(value)?), + )), + _ => internal_err!( + "Unexpected data type {:?} for function ascii", + scalar.data_type() + ), + } +} + +/// Returns the Unicode scalar value of the first character of `s`, or 0 when +/// `s` is empty. Reads the leading byte first so the common all-ASCII case +/// avoids constructing a `char` iterator and decoding a multi-byte sequence. +#[inline] +fn first_char_code(s: &str) -> i32 { + match s.as_bytes().first() { + None => 0, + // ASCII byte: the codepoint equals the byte value. + Some(&b) if b < 0x80 => b as i32, + // Leading byte of a multi-byte sequence: decode the first char. + Some(_) => s.chars().next().map_or(0, |c| c as i32), + } +} + fn calculate_ascii<'a, V>(array: &V) -> Result where V: StringArrayType<'a, Item = &'a str>, { - let values: Vec<_> = (0..array.len()) - .map(|i| { - if array.is_null(i) { - 0 - } else { - let s = array.value(i); - s.chars().next().map_or(0, |c| c as i32) - } - }) - .collect(); - - let array = Int32Array::new(values.into(), array.nulls().cloned()); + let len = array.len(); + let nulls = array.nulls().cloned(); + + // Split the null-handling out of the hot loop: when there is no null + // buffer every index is valid, so we can skip the per-element null check + // and use unchecked accessors. + let values: Vec = match nulls { + Some(ref n) => (0..len) + .map(|i| { + if n.is_null(i) { + 0 + } else { + // SAFETY: `n.is_null(i)` was false, so `i` is a valid, + // non-null index. + let s = unsafe { array.value_unchecked(i) }; + first_char_code(s) + } + }) + .collect(), + None => (0..len) + .map(|i| { + // SAFETY: no null buffer means every index in `0..len` is valid. + let s = unsafe { array.value_unchecked(i) }; + first_char_code(s) + }) + .collect(), + }; + + let array = Int32Array::new(values.into(), nulls); Ok(Arc::new(array)) } @@ -153,6 +189,11 @@ pub fn ascii(args: &[ArrayRef]) -> Result { let string_array = args[0].as_string_view(); Ok(calculate_ascii(&string_array)?) } + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let converted = ascii(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(converted)) + } _ => internal_err!("Unsupported data type"), } } diff --git a/datafusion/functions/src/string/bit_length.rs b/datafusion/functions/src/string/bit_length.rs index 76d8bb73bba87..4af22f5db5b5f 100644 --- a/datafusion/functions/src/string/bit_length.rs +++ b/datafusion/functions/src/string/bit_length.rs @@ -18,13 +18,13 @@ use arrow::compute::kernels::length::bit_length; use arrow::datatypes::DataType; -use crate::utils::utf8_to_int_type; +use crate::utils::{transform_leaf_type_preserving_encoding, utf8_to_int_type}; use datafusion_common::types::logical_string; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -59,9 +59,10 @@ impl BitLengthFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -78,7 +79,9 @@ impl ScalarUDFImpl for BitLengthFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - utf8_to_int_type(&arg_types[0], "bit_length") + transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| { + utf8_to_int_type(data_type, "bit_length") + }) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -86,18 +89,7 @@ impl ScalarUDFImpl for BitLengthFunc { match array { ColumnarValue::Array(v) => Ok(ColumnarValue::Array(bit_length(v.as_ref())?)), - ColumnarValue::Scalar(v) => match v { - ScalarValue::Utf8(v) => Ok(ColumnarValue::Scalar(ScalarValue::Int32( - v.as_ref().map(|x| (x.len() * 8) as i32), - ))), - ScalarValue::LargeUtf8(v) => Ok(ColumnarValue::Scalar( - ScalarValue::Int64(v.as_ref().map(|x| (x.len() * 8) as i64)), - )), - ScalarValue::Utf8View(v) => Ok(ColumnarValue::Scalar( - ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32)), - )), - _ => unreachable!("bit length"), - }, + ColumnarValue::Scalar(v) => Ok(ColumnarValue::Scalar(bit_length_scalar(v))), } } @@ -105,3 +97,21 @@ impl ScalarUDFImpl for BitLengthFunc { self.doc() } } + +fn bit_length_scalar(value: &ScalarValue) -> ScalarValue { + match value { + ScalarValue::Utf8(v) => { + ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32)) + } + ScalarValue::LargeUtf8(v) => { + ScalarValue::Int64(v.as_ref().map(|x| (x.len() * 8) as i64)) + } + ScalarValue::Utf8View(v) => { + ScalarValue::Int32(v.as_ref().map(|x| (x.len() * 8) as i32)) + } + ScalarValue::Dictionary(key_type, value) => { + ScalarValue::Dictionary(key_type.clone(), Box::new(bit_length_scalar(value))) + } + _ => unreachable!("bit length"), + } +} diff --git a/datafusion/functions/src/string/btrim.rs b/datafusion/functions/src/string/btrim.rs index 279f444d9ffe7..82e1f7d2c778d 100644 --- a/datafusion/functions/src/string/btrim.rs +++ b/datafusion/functions/src/string/btrim.rs @@ -16,30 +16,41 @@ // under the License. use crate::string::common::*; -use crate::utils::{make_scalar_function, utf8_to_str_type}; -use arrow::array::{ArrayRef, OffsetSizeTrait}; +use crate::utils::make_scalar_function; +use arrow::array::{ArrayRef, AsArray}; use arrow::datatypes::DataType; use datafusion_common::types::logical_string; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; use std::sync::Arc; /// Returns the longest string with leading and trailing characters removed. If the characters are not specified, spaces are removed. /// btrim('xyxtrimyyx', 'xyz') = 'trim' -fn btrim(args: &[ArrayRef]) -> Result { - let use_string_view = args[0].data_type() == &DataType::Utf8View; +fn btrim(args: &[ArrayRef]) -> Result { let args = if args.len() > 1 { let arg1 = arrow::compute::kernels::cast::cast(&args[1], args[0].data_type())?; vec![Arc::clone(&args[0]), arg1] } else { args.to_owned() }; - general_trim::(&args, use_string_view) + match args[0].data_type() { + DataType::Utf8 => general_trim::(&args, false), + DataType::LargeUtf8 => general_trim::(&args, false), + DataType::Utf8View => general_trim::(&args, true), + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let trimmed = btrim(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(trimmed)) + } + other => exec_err!( + "Unsupported data type {other:?} for function btrim, expected Utf8, LargeUtf8 or Utf8View." + ), + } } #[user_doc( @@ -85,9 +96,12 @@ impl BTrimFunc { Coercion::new_exact(TypeSignatureClass::Native(logical_string())), Coercion::new_exact(TypeSignatureClass::Native(logical_string())), ]), - TypeSignature::Coercible(vec![Coercion::new_exact( - TypeSignatureClass::Native(logical_string()), - )]), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::dictionary(), + ), + ]), ], Volatility::Immutable, ), @@ -106,28 +120,11 @@ impl ScalarUDFImpl for BTrimFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - if arg_types[0] == DataType::Utf8View { - Ok(DataType::Utf8View) - } else { - utf8_to_str_type(&arg_types[0], "btrim") - } + Ok(arg_types[0].clone()) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - match args.args[0].data_type() { - DataType::Utf8 | DataType::Utf8View => make_scalar_function( - btrim::, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - DataType::LargeUtf8 => make_scalar_function( - btrim::, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - other => exec_err!( - "Unsupported data type {other:?} for function btrim,\ - expected Utf8, LargeUtf8 or Utf8View." - ), - } + make_scalar_function(btrim, vec![Hint::Pad, Hint::AcceptsSingular])(&args.args) } fn aliases(&self) -> &[String] { diff --git a/datafusion/functions/src/string/common.rs b/datafusion/functions/src/string/common.rs index 2732ba4f86ef2..11ebf7d3d62dd 100644 --- a/datafusion/functions/src/string/common.rs +++ b/datafusion/functions/src/string/common.rs @@ -21,13 +21,13 @@ use std::sync::Arc; use crate::strings::{ GenericStringArrayBuilder, STRING_VIEW_INIT_BLOCK_SIZE, STRING_VIEW_MAX_BLOCK_SIZE, - StringViewArrayBuilder, append_view, + StringViewArrayBuilder, StringWriter, append_view, }; use arrow::array::{ - Array, ArrayRef, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, + Array, ArrayRef, AsArray, GenericStringArray, NullBufferBuilder, OffsetSizeTrait, StringViewArray, new_null_array, }; -use arrow::buffer::{Buffer, OffsetBuffer, ScalarBuffer}; +use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; @@ -262,6 +262,65 @@ fn trim_and_append_view( } } +/// Builds the trimmed output array by writing the trimmed slices straight into +/// the value buffer, rather than collecting through a string builder. +/// +/// Every trimmed value is a substring of its input, so the byte range the input +/// spans bounds the output's. Reserving that much up front means one allocation +/// and no growth during the copy, and it also guarantees the running offset stays +/// within `T` (the input array's own offsets already fit). +/// +/// `nulls` becomes the output null buffer; null rows contribute no bytes. +/// `trim_row` is called only for non-null rows, with the row index and its value, +/// and must return a subslice of the value it is given. +fn build_trimmed( + string_array: &GenericStringArray, + nulls: Option, + mut trim_row: F, +) -> ArrayRef +where + F: for<'a> FnMut(usize, &'a str) -> &'a str, +{ + let len = string_array.len(); + let input_offsets = string_array.value_offsets(); + let start = input_offsets.first().unwrap().as_usize(); + let end = input_offsets.last().unwrap().as_usize(); + + let mut values: Vec = Vec::with_capacity(end - start); + let mut offsets: Vec = Vec::with_capacity(len + 1); + offsets.push(T::usize_as(0)); + + match &nulls { + // Keeping the null check out of the all-valid path leaves it branch-free. + None => { + for i in 0..len { + // SAFETY: `i` is in bounds. + let s = unsafe { string_array.value_unchecked(i) }; + values.extend_from_slice(trim_row(i, s).as_bytes()); + offsets.push(T::usize_as(values.len())); + } + } + Some(validity) => { + for i in 0..len { + if validity.is_valid(i) { + // SAFETY: `i` is in bounds. + let s = unsafe { string_array.value_unchecked(i) }; + values.extend_from_slice(trim_row(i, s).as_bytes()); + } + offsets.push(T::usize_as(values.len())); + } + } + } + + let offsets = OffsetBuffer::new(ScalarBuffer::from(offsets)); + // SAFETY: trimming splits `s` on char boundaries, so the value buffer is a + // concatenation of valid UTF-8; the offsets are monotonic and end at its length. + let array = unsafe { + GenericStringArray::::new_unchecked(offsets, Buffer::from_vec(values), nulls) + }; + Arc::new(array) +} + /// Applies the trim function to the given string array(s) /// and returns a new string array with the trimmed values. /// @@ -273,12 +332,11 @@ fn string_trim(args: &[ArrayRef]) -> Result { // Trim spaces by default - let result = string_array - .iter() - .map(|string| string.map(|s| Tr::trim_ascii_char(s, b' ').0)) - .collect::>(); - - Ok(Arc::new(result) as ArrayRef) + Ok(build_trimmed( + string_array, + string_array.nulls().cloned(), + |_, s| Tr::trim_ascii_char(s, b' ').0, + )) } 2 => { let characters_array = as_generic_string_array::(&args[1])?; @@ -293,29 +351,31 @@ fn string_trim(args: &[ArrayRef]) -> Result = characters_array.value(0).chars().collect(); - let result = string_array - .iter() - .map(|item| item.map(|s| Tr::trim(s, &pattern).0)) - .collect::>(); - return Ok(Arc::new(result) as ArrayRef); + return Ok(build_trimmed( + string_array, + string_array.nulls().cloned(), + |_, s| Tr::trim(s, &pattern).0, + )); } + // Indexing `characters_array` per row below requires the two arguments + // to line up. + if characters_array.len() != string_array.len() { + return exec_err!( + "Function TRIM was called with mismatched argument lengths" + ); + } + + // A row is null if either argument is null. + let nulls = NullBuffer::union(string_array.nulls(), characters_array.nulls()); + // Per-row pattern - must compute pattern chars for each row let mut pattern: Vec = Vec::new(); - let result = string_array - .iter() - .zip(characters_array.iter()) - .map(|(string, characters)| match (string, characters) { - (Some(s), Some(c)) => { - pattern.clear(); - pattern.extend(c.chars()); - Some(Tr::trim(s, &pattern).0) - } - _ => None, - }) - .collect::>(); - - Ok(Arc::new(result) as ArrayRef) + Ok(build_trimmed(string_array, nulls, |i, s| { + pattern.clear(); + pattern.extend(characters_array.value(i).chars()); + Tr::trim(s, &pattern).0 + })) } other => { exec_err!( @@ -342,69 +402,129 @@ fn unicode_case(s: &str, lower: bool) -> String { } } +/// Writes the case-converted form of `s` directly into `w`. +/// +/// Uppercasing is a context-free character mapping, so each character is +/// mapped and streamed straight into the output buffer, avoiding the +/// intermediate `String` that `str::to_uppercase` allocates per row. +/// +/// Lowercasing is *not* context-free — `str::to_lowercase` applies the +/// special Greek final-sigma rule (Σ becomes ς at the end of a word but σ +/// elsewhere), which a per-character mapping cannot reproduce — so it keeps +/// using `str::to_lowercase`. +#[inline] +fn write_unicode_case(w: &mut impl StringWriter, s: &str, lower: bool) { + if lower { + w.write_str(&s.to_lowercase()); + } else { + for c in s.chars() { + for upper in c.to_uppercase() { + w.write_char(upper); + } + } + } +} + fn case_conversion( args: &[ColumnarValue], lower: bool, name: &str, ) -> Result { match &args[0] { - ColumnarValue::Array(array) => match array.data_type() { - DataType::Utf8 => Ok(ColumnarValue::Array(case_conversion_array::( - array, lower, - )?)), - DataType::LargeUtf8 => Ok(ColumnarValue::Array( - case_conversion_array::(array, lower)?, - )), - DataType::Utf8View => { - let string_array = as_string_view_array(array)?; - if string_array.is_ascii() { - return Ok(ColumnarValue::Array(Arc::new( - case_conversion_utf8view_ascii(string_array, lower), - ))); - } - let item_len = string_array.len(); - // Null-preserving: reuse the input null buffer as the output null buffer. - let nulls = string_array.nulls().cloned(); - let mut builder = StringViewArrayBuilder::with_capacity(item_len); - - if let Some(ref n) = nulls { - for i in 0..item_len { - if n.is_null(i) { - builder.append_placeholder(); - } else { - // SAFETY: `n.is_null(i)` was false in the branch above. - let s = unsafe { string_array.value_unchecked(i) }; - builder.append_value(&unicode_case(s, lower)); - } - } - } else { - for i in 0..item_len { - // SAFETY: no null buffer means every index is valid. - let s = unsafe { string_array.value_unchecked(i) }; - builder.append_value(&unicode_case(s, lower)); - } - } + ColumnarValue::Array(array) => Ok(ColumnarValue::Array( + case_conversion_columnar_array(array, lower, name)?, + )), + ColumnarValue::Scalar(scalar) => Ok(ColumnarValue::Scalar( + case_conversion_scalar(scalar, lower, name)?, + )), + } +} - Ok(ColumnarValue::Array(Arc::new(builder.finish(nulls)?))) - } - other => exec_err!("Unsupported data type {other:?} for function {name}"), - }, - ColumnarValue::Scalar(scalar) => match scalar { - ScalarValue::Utf8(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8(result))) - } - ScalarValue::LargeUtf8(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(result))) - } - ScalarValue::Utf8View(a) => { - let result = a.as_ref().map(|x| unicode_case(x, lower)); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(result))) +fn case_conversion_scalar( + scalar: &ScalarValue, + lower: bool, + name: &str, +) -> Result { + match scalar { + ScalarValue::Utf8(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::Utf8(result)) + } + ScalarValue::LargeUtf8(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::LargeUtf8(result)) + } + ScalarValue::Utf8View(a) => { + let result = a.as_ref().map(|x| unicode_case(x, lower)); + Ok(ScalarValue::Utf8View(result)) + } + ScalarValue::Dictionary(key_type, value) => { + let converted = case_conversion_scalar(value.as_ref(), lower, name)?; + Ok(ScalarValue::Dictionary( + key_type.clone(), + Box::new(converted), + )) + } + other => exec_err!("Unsupported data type {other:?} for function {name}"), + } +} + +fn case_conversion_columnar_array( + array: &ArrayRef, + lower: bool, + name: &str, +) -> Result { + match array.data_type() { + DataType::Utf8 => case_conversion_array::(array, lower), + DataType::LargeUtf8 => case_conversion_array::(array, lower), + DataType::Utf8View => case_conversion_utf8view(array, lower), + DataType::Dictionary(_, _) => case_conversion_dictionary(array, lower, name), + other => exec_err!("Unsupported data type {other:?} for function {name}"), + } +} + +fn case_conversion_utf8view(array: &ArrayRef, lower: bool) -> Result { + let string_array = as_string_view_array(array)?; + if string_array.is_ascii() { + return Ok(Arc::new(case_conversion_utf8view_ascii( + string_array, + lower, + ))); + } + let item_len = string_array.len(); + // Null-preserving: reuse the input null buffer as the output null buffer. + let nulls = string_array.nulls().cloned(); + let mut builder = StringViewArrayBuilder::with_capacity(item_len); + + if let Some(ref n) = nulls { + for i in 0..item_len { + if n.is_null(i) { + builder.try_append_placeholder()?; + } else { + // SAFETY: `n.is_null(i)` was false in the branch above. + let s = unsafe { string_array.value_unchecked(i) }; + builder.try_append_value(&unicode_case(s, lower))?; } - other => exec_err!("Unsupported data type {other:?} for function {name}"), - }, + } + } else { + for i in 0..item_len { + // SAFETY: no null buffer means every index is valid. + let s = unsafe { string_array.value_unchecked(i) }; + builder.try_append_value(&unicode_case(s, lower))?; + } } + + Ok(Arc::new(builder.finish(nulls)?)) +} + +fn case_conversion_dictionary( + array: &ArrayRef, + lower: bool, + name: &str, +) -> Result { + let dictionary = array.as_any_dictionary(); + let converted = case_conversion_columnar_array(dictionary.values(), lower, name)?; + Ok(dictionary.with_values(converted)) } fn case_conversion_array( @@ -431,18 +551,18 @@ fn case_conversion_array( if let Some(ref n) = nulls { for i in 0..item_len { if n.is_null(i) { - builder.append_placeholder(); + builder.try_append_placeholder()?; } else { // SAFETY: `n.is_null(i)` was false in the branch above. let s = unsafe { string_array.value_unchecked(i) }; - builder.append_value(&unicode_case(s, lower)); + builder.try_append_with(|w| write_unicode_case(w, s, lower))?; } } } else { for i in 0..item_len { // SAFETY: no null buffer means every index is valid. let s = unsafe { string_array.value_unchecked(i) }; - builder.append_value(&unicode_case(s, lower)); + builder.try_append_with(|w| write_unicode_case(w, s, lower))?; } } Ok(Arc::new(builder.finish(nulls)?)) @@ -520,6 +640,10 @@ fn case_conversion_utf8view_ascii_inner u8>( block_size = block_size.saturating_mul(2); } let to_reserve = len.max(block_size as usize); + #[expect( + clippy::disallowed_methods, + reason = "StringView's block size bounds growth, so reserve cannot overflow capacity arithmetically. This hot loop intentionally avoids the extra `try_reserve` checks. It remains subject to allocator failure/OOM, which must be managed externally." + )] in_progress.reserve(to_reserve); } @@ -595,15 +719,10 @@ fn case_conversion_ascii_array( let values = Buffer::from_vec(converted); // Shift offsets from `start`-based to 0-based so they index into `values`. - let offsets = if start == 0 { - string_array.offsets().clone() - } else { - let s = O::usize_as(start); - let rebased: Vec = value_offsets.iter().map(|&o| o - s).collect(); - // SAFETY: subtracting a constant from monotonic offsets preserves - // monotonicity, and `start` is the minimum offset, so no underflow. - unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(rebased)) } - }; + let offsets = string_array + .offsets() + .clone() + .subtract(string_array.offsets()[0]); let nulls = string_array.nulls().cloned(); // SAFETY: offsets are monotonic and in-bounds for `values`; nulls diff --git a/datafusion/functions/src/string/concat.rs b/datafusion/functions/src/string/concat.rs index b10db23472c99..1c1f6d640798a 100644 --- a/datafusion/functions/src/string/concat.rs +++ b/datafusion/functions/src/string/concat.rs @@ -15,17 +15,16 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{Array, as_largestring_array}; -use arrow::datatypes::DataType; -use datafusion_expr::sort_properties::ExprProperties; -use std::sync::Arc; - +use crate::binaries::{ + ConcatBinaryBuilder, ConcatBinaryViewBuilder, ConcatLargeBinaryBuilder, +}; use crate::string::concat; use crate::strings::{ - ColumnarValueRef, ConcatLargeStringBuilder, ConcatStringBuilder, - ConcatStringViewBuilder, + ColumnarValueRef, ConcatBuilder, ConcatLargeStringBuilder, ConcatStringBuilder, + ConcatStringViewBuilder, widest_binary_type, widest_string_type, }; -use datafusion_common::cast::{as_binary_array, as_string_array, as_string_view_array}; +use arrow::array::Array; +use arrow::datatypes::DataType; use datafusion_common::{ Result, ScalarValue, exec_datafusion_err, internal_err, plan_err, }; @@ -67,27 +66,18 @@ impl Default for ConcatFunc { impl ConcatFunc { pub fn new() -> Self { - use DataType::*; Self { - signature: Signature::variadic( - vec![Utf8View, Utf8, LargeUtf8, Binary], - Volatility::Immutable, - ), + // Use `Signature::UserDefined` to allow different argument types. + // `Variadic` requires every argument to be coerced to the same string type, + // so the UDF cannot distinguish between binary and string inputs. + signature: Signature::user_defined(Volatility::Immutable), } } } -fn deduce_return_type(arg_types: &[DataType]) -> DataType { - use DataType::*; - if arg_types.contains(&Utf8View) { - Utf8View - } else if arg_types.contains(&LargeUtf8) { - LargeUtf8 - } else { - Utf8 - } -} - +// Supports string + string concatenation, binary + binary concatenation, +// and mixed string + binary concatenation (binary is coerced to the widest +// string type). impl ScalarUDFImpl for ConcatFunc { fn name(&self) -> &str { "concat" @@ -97,9 +87,18 @@ impl ScalarUDFImpl for ConcatFunc { &self.signature } - /// Match the return type to the input types to avoid unnecessary casts. On + /// Coerce all arguments to the widest type within the binary / string family + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if arg_types.is_empty() { + plan_err!("concat does not support zero arguments") + } else { + coerce_arg_types(arg_types) + } + } + /// mixed inputs, prefer Utf8View; prefer LargeUtf8 over Utf8 to avoid /// potential overflow on LargeUtf8 input. + /// For binaries, use the similar hierarchy fn return_type(&self, arg_types: &[DataType]) -> Result { Ok(deduce_return_type(arg_types)) } @@ -107,11 +106,9 @@ impl ScalarUDFImpl for ConcatFunc { /// Concatenates the text representations of all the arguments. NULL arguments are ignored. /// concat('abcde', 2, NULL, 22) = 'abcde222' fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let return_datatype = args.return_type().clone(); let ScalarFunctionArgs { args, .. } = args; - let arg_types: Vec = args.iter().map(|c| c.data_type()).collect(); - let return_datatype = deduce_return_type(&arg_types); - let array_len = args.iter().find_map(|x| match x { ColumnarValue::Array(array) => Some(array.len()), _ => None, @@ -126,7 +123,14 @@ impl ScalarUDFImpl for ConcatFunc { }; if let ScalarValue::Binary(Some(value)) = scalar { values.push(value); + } else if let ScalarValue::LargeBinary(Some(value)) = scalar { + values.push(value); + } else if let ScalarValue::BinaryView(Some(value)) = scalar { + values.push(value); + } else if scalar.is_null() { + // null binary scalar: skip (consistent with null string behaviour) } else { + // String case match scalar.try_as_str() { Some(Some(v)) => values.push(v.as_bytes()), Some(None) => {} // null literal @@ -138,20 +142,42 @@ impl ScalarUDFImpl for ConcatFunc { } } let concat_bytes = values.concat(); - let result = std::str::from_utf8(&concat_bytes) - .map_err(|_| exec_datafusion_err!("invalid UTF-8 in binary literal"))? - .to_string(); return match return_datatype { DataType::Utf8View => { + let result = std::str::from_utf8(&concat_bytes) + .map_err(|_| { + exec_datafusion_err!("invalid UTF-8 in binary literal") + })? + .to_string(); Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(result)))) } DataType::Utf8 => { + let result = std::str::from_utf8(&concat_bytes) + .map_err(|_| { + exec_datafusion_err!("invalid UTF-8 in binary literal") + })? + .to_string(); Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result)))) } DataType::LargeUtf8 => { + let result = std::str::from_utf8(&concat_bytes) + .map_err(|_| { + exec_datafusion_err!("invalid UTF-8 in binary literal") + })? + .to_string(); Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(result)))) } + DataType::Binary => Ok(ColumnarValue::Scalar(ScalarValue::Binary(Some( + concat_bytes, + )))), + // Serves LargeBinary and FixedSizeBinary inputs + DataType::LargeBinary => Ok(ColumnarValue::Scalar( + ScalarValue::LargeBinary(Some(concat_bytes)), + )), + DataType::BinaryView => Ok(ColumnarValue::Scalar( + ScalarValue::BinaryView(Some(concat_bytes)), + )), other => { plan_err!("Concat function does not support datatype of {other}") } @@ -164,121 +190,46 @@ impl ScalarUDFImpl for ConcatFunc { let mut columns = Vec::with_capacity(args.len()); for arg in &args { - match arg { - ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value)) - | ColumnarValue::Scalar(ScalarValue::LargeUtf8(maybe_value)) - | ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => { - if let Some(s) = maybe_value { - data_size += s.len() * len; - columns.push(ColumnarValueRef::Scalar(s.as_bytes())); - } - } - ColumnarValue::Scalar(ScalarValue::Binary(maybe_value)) => { - if let Some(b) = maybe_value { - // data_size is a capacity hint, so doesn't matter if it is chars or bytes - data_size += b.len() * len; - columns.push(ColumnarValueRef::Scalar(b.as_slice())); - } - } - ColumnarValue::Array(array) => { - match array.data_type() { - DataType::Utf8 => { - let string_array = as_string_array(array)?; - - data_size += string_array.values().len(); - let column = if array.is_nullable() { - ColumnarValueRef::NullableArray(string_array) - } else { - ColumnarValueRef::NonNullableArray(string_array) - }; - columns.push(column); - } - DataType::LargeUtf8 => { - let string_array = as_largestring_array(array); - - data_size += string_array.values().len(); - let column = if array.is_nullable() { - ColumnarValueRef::NullableLargeStringArray(string_array) - } else { - ColumnarValueRef::NonNullableLargeStringArray( - string_array, - ) - }; - columns.push(column); - } - DataType::Utf8View => { - let string_array = as_string_view_array(array)?; - - // This is an estimate; in particular, it will - // undercount arrays of short strings (<= 12 bytes). - data_size += string_array.total_buffer_bytes_used(); - let column = if array.is_nullable() { - ColumnarValueRef::NullableStringViewArray(string_array) - } else { - ColumnarValueRef::NonNullableStringViewArray(string_array) - }; - columns.push(column); - } - DataType::Binary => { - let string_array = as_binary_array(array)?; - - data_size += string_array.values().len(); - let column = if array.is_nullable() { - ColumnarValueRef::NullableBinaryArray(string_array) - } else { - ColumnarValueRef::NonNullableBinaryArray(string_array) - }; - columns.push(column); - } - other => { - return plan_err!( - "Input was {other} which is not a supported datatype for concat function" - ); - } - }; - } - _ => unreachable!("concat"), + if let Some(column) = + ColumnarValueRef::from_columnar_value(arg, &mut data_size, len, 1, false)? + { + columns.push(column); } } match return_datatype { - DataType::Utf8 => { - let mut builder = ConcatStringBuilder::with_capacity(len, data_size); - for i in 0..len { - columns - .iter() - .for_each(|column| builder.write::(column, i)); - builder.append_offset()?; - } - - let string_array = builder.finish(None)?; - Ok(ColumnarValue::Array(Arc::new(string_array))) - } - DataType::Utf8View => { - let mut builder = ConcatStringViewBuilder::with_capacity(len, data_size); - for i in 0..len { - columns - .iter() - .for_each(|column| builder.write::(column, i)); - builder.append_offset()?; - } - - let string_array = builder.finish(None)?; - Ok(ColumnarValue::Array(Arc::new(string_array))) - } - DataType::LargeUtf8 => { - let mut builder = ConcatLargeStringBuilder::with_capacity(len, data_size); - for i in 0..len { - columns - .iter() - .for_each(|column| builder.write::(column, i)); - builder.append_offset()?; - } - - let string_array = builder.finish(None)?; - Ok(ColumnarValue::Array(Arc::new(string_array))) - } - _ => unreachable!(), + DataType::Utf8 => build_concat( + ConcatStringBuilder::with_capacity(len, data_size), + &columns, + len, + ), + DataType::Utf8View => build_concat( + ConcatStringViewBuilder::with_capacity(len, data_size), + &columns, + len, + ), + DataType::LargeUtf8 => build_concat( + ConcatLargeStringBuilder::with_capacity(len, data_size), + &columns, + len, + ), + DataType::Binary => build_concat( + ConcatBinaryBuilder::with_capacity(len, data_size), + &columns, + len, + ), + // Serves LargeBinary and FixedSizeBinary inputs + DataType::LargeBinary => build_concat( + ConcatLargeBinaryBuilder::with_capacity(len, data_size), + &columns, + len, + ), + DataType::BinaryView => build_concat( + ConcatBinaryViewBuilder::with_capacity(len, data_size), + &columns, + len, + ), + _ => unreachable!("concat"), } } @@ -301,13 +252,72 @@ impl ScalarUDFImpl for ConcatFunc { fn documentation(&self) -> Option<&Documentation> { self.doc() } +} + +pub(crate) fn deduce_return_type(arg_types: &[DataType]) -> DataType { + use DataType::*; + if arg_types.contains(&BinaryView) { + BinaryView + } else if arg_types.contains(&LargeBinary) { + // Serves LargeBinary and FixedSizeBinary inputs + LargeBinary + } else if arg_types.contains(&Binary) { + Binary + } else if arg_types.contains(&Utf8View) { + Utf8View + } else if arg_types.contains(&LargeUtf8) { + LargeUtf8 + } else { + Utf8 + } +} - fn preserves_lex_ordering(&self, _inputs: &[ExprProperties]) -> Result { - Ok(true) +/// Coerce all arguments to the widest type within the binary / string family +pub(crate) fn coerce_arg_types(arg_types: &[DataType]) -> Result> { + let has_binary = arg_types.iter().any(|dt| dt.is_binary()); + let has_string = arg_types.iter().any(|dt| dt.is_string()); + if has_binary && has_string { + // Mixed string+binary: coerce everything to the widest string type + // This behaviour is seen for Spark, DuckDB + Ok(vec![widest_string_type(arg_types); arg_types.len()]) + } else if has_binary { + // Pure binary+binary concatenation: coerce to the widest binary type + Ok(vec![widest_binary_type(arg_types); arg_types.len()]) + } else { + // Pure string+string concatenation: coerce to the widest string type + Ok(vec![widest_string_type(arg_types); arg_types.len()]) } } +/// Build a `concats` output array using a generic [`ConcatBuilder`]. +fn build_concat( + mut builder: B, + columns: &[ColumnarValueRef], + len: usize, +) -> Result { + for i in 0..len { + for column in columns { + builder.write::(column, i)?; + } + builder.append_offset()?; + } + + let array = builder.finish(None)?; + Ok(ColumnarValue::Array(array)) +} + pub(crate) fn simplify_concat(args: Vec) -> Result { + // Skip simplification when binary literals are present, because it + // handles only strings + for arg in &args { + match arg { + Expr::Literal(dt, _) if dt.data_type().is_binary() => { + return Ok(ExprSimplifyResult::Original(args)); + } + _ => {} + } + } + let mut new_args = Vec::with_capacity(args.len()); let mut contiguous_scalar = "".to_string(); @@ -396,10 +406,13 @@ mod tests { use super::*; use crate::utils::test::test_function; use DataType::*; - use arrow::array::{ArrayRef, StringArray}; + use arrow::array::{ + ArrayRef, BinaryArray, BinaryViewArray, LargeBinaryArray, StringArray, + }; use arrow::array::{LargeStringArray, StringViewArray}; use arrow::datatypes::Field; use datafusion_common::config::ConfigOptions; + use std::sync::Arc; #[test] fn test_functions() -> Result<()> { @@ -471,38 +484,95 @@ mod tests { Utf8View, StringViewArray ); + Ok(()) + } + + #[test] + fn test_scalar_binary() -> Result<()> { + test_function!( + ConcatFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Binary(Some( + "Café".as_bytes().into() + ))), + ColumnarValue::Scalar(ScalarValue::Binary(Some("cc".as_bytes().into()))), + ], + Ok(Some("Cafécc".as_bytes())), + &[u8], + Binary, + BinaryArray + ); test_function!( ConcatFunc::new(), vec![ ColumnarValue::Scalar(ScalarValue::Binary(Some( "Café".as_bytes().into() ))), - ColumnarValue::Scalar(ScalarValue::Utf8(None)), - ColumnarValue::Scalar(ScalarValue::Utf8(Some("cc".to_string()))), + ColumnarValue::Scalar(ScalarValue::LargeBinary(Some( + "cc".as_bytes().into() + ))), ], - Ok(Some("Cafécc")), - &str, - Utf8, - StringArray + Ok(Some("Cafécc".as_bytes())), + &[u8], + LargeBinary, + LargeBinaryArray ); test_function!( ConcatFunc::new(), vec![ - ColumnarValue::Scalar(ScalarValue::Binary(Some(Vec::from( - "Café".as_bytes() - )))), - ColumnarValue::Scalar(ScalarValue::Binary(Some("cc".as_bytes().into()))), + ColumnarValue::Scalar(ScalarValue::Binary(Some( + "Café".as_bytes().into() + ))), + ColumnarValue::Scalar(ScalarValue::BinaryView(Some( + "cc".as_bytes().into() + ))), ], - Ok(Some("Cafécc")), - &str, - Utf8, - StringArray + Ok(Some("Cafécc".as_bytes())), + &[u8], + BinaryView, + BinaryViewArray + ); + test_function!( + ConcatFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::BinaryView(Some( + "Café".as_bytes().into() + ))), + ColumnarValue::Scalar(ScalarValue::BinaryView(Some( + "cc".as_bytes().into() + ))), + ], + Ok(Some("Cafécc".as_bytes())), + &[u8], + BinaryView, + BinaryViewArray + ); + // Skip one Binary(None) + test_function!( + ConcatFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Binary(None)), + ColumnarValue::Scalar(ScalarValue::Binary(Some(b"hello".to_vec()))), + ], + Ok(Some(b"hello".as_ref())), + &[u8], + Binary, + BinaryArray + ); + // Skip all Binary(None), producing an empty array + test_function!( + ConcatFunc::new(), + vec![ColumnarValue::Scalar(ScalarValue::Binary(None))], + Ok(Some(b"".as_ref())), + &[u8], + Binary, + BinaryArray ); Ok(()) } #[test] - fn concat() -> Result<()> { + fn test_array_string() -> Result<()> { let c0 = ColumnarValue::Array(Arc::new(StringArray::from(vec!["foo", "bar", "baz"]))); let c1 = ColumnarValue::Scalar(ScalarValue::Utf8(Some(",".to_string()))); @@ -532,7 +602,7 @@ mod tests { args: vec![c0, c1, c2, c3, c4], arg_fields, number_rows: 3, - return_field: Field::new("f", Utf8, true).into(), + return_field: Field::new("f", Utf8View, true).into(), config_options: Arc::new(ConfigOptions::default()), }; @@ -548,4 +618,55 @@ mod tests { } Ok(()) } + + #[test] + fn test_array_binary() -> Result<()> { + let c0 = ColumnarValue::Array(Arc::new(BinaryArray::from_vec(vec![ + b"foo", b"bar", b"baz", + ]))); + let c1 = ColumnarValue::Scalar(ScalarValue::LargeBinary(Some(b",".to_vec()))); + let c2 = ColumnarValue::Array(Arc::new(BinaryArray::from_opt_vec(vec![ + Some(b"x"), + None, + Some(b"z"), + ]))); + let c3 = ColumnarValue::Scalar(ScalarValue::BinaryView(Some(b",".to_vec()))); + let c4 = ColumnarValue::Array(Arc::new(BinaryViewArray::from_iter(vec![ + Some(b"a"), + None, + Some(b"b"), + ]))); + let arg_fields = vec![ + Field::new("a", Binary, true), + Field::new("a", LargeBinary, true), + Field::new("a", Binary, true), + Field::new("a", BinaryView, true), + Field::new("a", BinaryView, true), + ] + .into_iter() + .map(Arc::new) + .collect::>(); + + let args = ScalarFunctionArgs { + args: vec![c0, c1, c2, c3, c4], + arg_fields, + number_rows: 3, + return_field: Field::new("f", BinaryView, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }; + + let result = ConcatFunc::new().invoke_with_args(args)?; + let expected = Arc::new(BinaryViewArray::from_iter(vec![ + Some(b"foo,x,a".to_vec()), + Some(b"bar,,".to_vec()), + Some(b"baz,z,b".to_vec()), + ])) as ArrayRef; + match &result { + ColumnarValue::Array(array) => { + assert_eq!(&expected, array); + } + _ => panic!(), + } + Ok(()) + } } diff --git a/datafusion/functions/src/string/concat_ws.rs b/datafusion/functions/src/string/concat_ws.rs index 2c2d4bd42165b..8cb6869974813 100644 --- a/datafusion/functions/src/string/concat_ws.rs +++ b/datafusion/functions/src/string/concat_ws.rs @@ -16,20 +16,18 @@ // under the License. use arrow::array::Array; -use std::sync::Arc; - use arrow::datatypes::DataType; +use crate::binaries::{ + ConcatBinaryBuilder, ConcatBinaryViewBuilder, ConcatLargeBinaryBuilder, +}; use crate::string::concat; -use crate::string::concat::simplify_concat; +use crate::string::concat::{coerce_arg_types, deduce_return_type, simplify_concat}; use crate::string::concat_ws; use crate::strings::{ - ColumnarValueRef, ConcatLargeStringBuilder, ConcatStringBuilder, + ColumnarValueRef, ConcatBuilder, ConcatLargeStringBuilder, ConcatStringBuilder, ConcatStringViewBuilder, }; -use datafusion_common::cast::{ - as_large_string_array, as_string_array, as_string_view_array, -}; use datafusion_common::{Result, ScalarValue, exec_err, internal_err, plan_err}; use datafusion_expr::expr::ScalarFunction; use datafusion_expr::simplify::{ExprSimplifyResult, SimplifyContext}; @@ -76,12 +74,11 @@ impl Default for ConcatWsFunc { impl ConcatWsFunc { pub fn new() -> Self { - use DataType::*; Self { - signature: Signature::variadic( - vec![Utf8View, Utf8, LargeUtf8], - Volatility::Immutable, - ), + // Use `Signature::UserDefined` to allow different argument types. + // `Variadic` requires every argument to be coerced to the same string type, + // so the UDF cannot distinguish between binary and string inputs. + signature: Signature::user_defined(Volatility::Immutable), } } } @@ -95,25 +92,29 @@ impl ScalarUDFImpl for ConcatWsFunc { &self.signature } - /// Match the return type to the input types to avoid unnecessary casts. On - /// mixed inputs, prefer Utf8View; prefer LargeUtf8 over Utf8 to avoid - /// potential overflow on LargeUtf8 input. - fn return_type(&self, arg_types: &[DataType]) -> Result { - use DataType::*; - if arg_types.contains(&Utf8View) { - Ok(Utf8View) - } else if arg_types.contains(&LargeUtf8) { - Ok(LargeUtf8) + /// Coerce all arguments to the widest type within the binary / string family + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if arg_types.len() < 2 { + plan_err!( + "concat_ws expects at least 2 arguments, got {}", + arg_types.len() + ) } else { - Ok(Utf8) + coerce_arg_types(arg_types) } } + /// Match the return type to the input types. Delegates to `concat` implementation. + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(deduce_return_type(arg_types)) + } + /// Concatenates all but the first argument, with separators. The first /// argument is used as the separator string, and should not be NULL. Other /// NULL arguments are ignored. /// concat_ws(',', 'abcde', 2, NULL, 22) = 'abcde,2,22' fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let return_datatype = args.return_type().clone(); let ScalarFunctionArgs { args, .. } = args; if args.len() < 2 { @@ -123,14 +124,9 @@ impl ScalarUDFImpl for ConcatWsFunc { ); } - let return_datatype = if args.iter().any(|c| c.data_type() == DataType::Utf8View) - { - DataType::Utf8View - } else if args.iter().any(|c| c.data_type() == DataType::LargeUtf8) { - DataType::LargeUtf8 - } else { - DataType::Utf8 - }; + let arg_types: Vec = args.iter().map(|c| c.data_type()).collect(); + + let with_binary = arg_types.iter().any(|dt| dt.is_binary()); let array_len = args.iter().find_map(|x| match x { ColumnarValue::Array(array) => Some(array.len()), @@ -142,47 +138,101 @@ impl ScalarUDFImpl for ConcatWsFunc { let ColumnarValue::Scalar(scalar) = &args[0] else { unreachable!() }; - let sep = match scalar.try_as_str() { - Some(Some(s)) => s, - Some(None) => { - // null literal string - return match return_datatype { - DataType::Utf8View => { - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(None))) + + return if with_binary { + // Binary scalar path + let sep_bytes: &[u8] = match scalar { + ScalarValue::Binary(Some(v)) + | ScalarValue::LargeBinary(Some(v)) + | ScalarValue::BinaryView(Some(v)) => v.as_slice(), + ScalarValue::FixedSizeBinary(_, Some(v)) => v.as_slice(), + scalar if scalar.is_null() => { + return Ok(null_scalar(&return_datatype)); + } + other => { + return internal_err!("Expected binary separator, got {other:?}"); + } + }; + + let mut values: Vec<&[u8]> = Vec::with_capacity(args.len() - 1); + for arg in &args[1..] { + let ColumnarValue::Scalar(s) = arg else { + unreachable!() + }; + match s { + ScalarValue::Binary(Some(v)) + | ScalarValue::LargeBinary(Some(v)) + | ScalarValue::BinaryView(Some(v)) => values.push(v.as_slice()), + ScalarValue::FixedSizeBinary(_, Some(v)) => { + values.push(v.as_slice()) } - DataType::LargeUtf8 => { - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(None))) + // skip null + scalar if scalar.is_null() => {} + other => { + return internal_err!("Expected binary value, got {other:?}"); } - _ => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))), - }; + } } - None => return internal_err!("Expected string literal, got {scalar:?}"), - }; + let result = values.join(sep_bytes); - let mut values = Vec::with_capacity(args.len() - 1); - for arg in &args[1..] { - let ColumnarValue::Scalar(scalar) = arg else { - unreachable!() - }; - - match scalar.try_as_str() { - Some(Some(v)) => values.push(v), - Some(None) => {} // null literal string + match return_datatype { + DataType::Binary => { + Ok(ColumnarValue::Scalar(ScalarValue::Binary(Some(result)))) + } + DataType::LargeBinary => Ok(ColumnarValue::Scalar( + ScalarValue::LargeBinary(Some(result)), + )), + DataType::BinaryView => { + Ok(ColumnarValue::Scalar(ScalarValue::BinaryView(Some(result)))) + } + other => { + plan_err!("concat_ws does not support return type {other}") + } + } + } else { + // String scalar path + let sep = match scalar.try_as_str() { + Some(Some(s)) => s, + Some(None) => { + return Ok(null_scalar(&return_datatype)); + } None => { return internal_err!("Expected string literal, got {scalar:?}"); } - } - } - let result = values.join(sep); + }; + + let mut values = Vec::with_capacity(args.len() - 1); + for arg in &args[1..] { + let ColumnarValue::Scalar(scalar) = arg else { + unreachable!() + }; - return match return_datatype { - DataType::Utf8View => { - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(result)))) + match scalar.try_as_str() { + Some(Some(v)) => values.push(v), + Some(None) => {} // null literal string + None => { + return internal_err!( + "Expected string literal, got {scalar:?}" + ); + } + } } - DataType::LargeUtf8 => { - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(result)))) + let result = values.join(sep); + + match return_datatype { + DataType::Utf8View => { + Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(result)))) + } + DataType::LargeUtf8 => { + Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(result)))) + } + DataType::Utf8 => { + Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result)))) + } + other => { + plan_err!("concat_ws does not support return type {other}") + } } - _ => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result)))), }; } @@ -190,190 +240,66 @@ impl ScalarUDFImpl for ConcatWsFunc { let len = array_len.unwrap(); let mut data_size = 0; - // parse sep - let sep = match &args[0] { - ColumnarValue::Scalar(scalar) => match scalar.try_as_str() { - Some(Some(s)) => { - data_size += s.len() * len * (args.len() - 2); // estimate - ColumnarValueRef::Scalar(s.as_bytes()) - } - Some(None) => { - return match return_datatype { - DataType::Utf8View => { - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(None))) - } - DataType::LargeUtf8 => { - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(None))) - } - _ => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))), - }; - } - None => { - return internal_err!("Expected string separator, got {scalar:?}"); - } - }, - ColumnarValue::Array(array) => match array.data_type() { - DataType::Utf8 => { - let string_array = as_string_array(array)?; - data_size += string_array.values().len() * (args.len() - 2); - if array.is_nullable() { - ColumnarValueRef::NullableArray(string_array) - } else { - ColumnarValueRef::NonNullableArray(string_array) - } - } - DataType::LargeUtf8 => { - let string_array = as_large_string_array(array)?; - data_size += string_array.values().len() * (args.len() - 2); - if array.is_nullable() { - ColumnarValueRef::NullableLargeStringArray(string_array) - } else { - ColumnarValueRef::NonNullableLargeStringArray(string_array) - } - } - DataType::Utf8View => { - let string_array = as_string_view_array(array)?; - data_size += - string_array.total_buffer_bytes_used() * (args.len() - 2); - if array.is_nullable() { - ColumnarValueRef::NullableStringViewArray(string_array) - } else { - ColumnarValueRef::NonNullableStringViewArray(string_array) - } - } - other => { - return plan_err!( - "Input was {other} which is not a supported datatype for concat_ws separator" - ); - } - }, - }; + let sep_column = &args[0]; + + // A null scalar separator makes the entire result null for all rows. + if matches!(sep_column, ColumnarValue::Scalar(s) if s.is_null()) { + return Ok(null_scalar(&return_datatype)); + } + + let sep: ColumnarValueRef = ColumnarValueRef::from_columnar_value(sep_column, &mut data_size, len, args.len() - 2, true)? + .map(Ok) + .unwrap_or_else(|| plan_err!( + "Input {sep_column} which is not a supported datatype for concat_ws separator" + ))?; let mut columns = Vec::with_capacity(args.len() - 1); for arg in &args[1..] { - match arg { - ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value)) - | ColumnarValue::Scalar(ScalarValue::LargeUtf8(maybe_value)) - | ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => { - if let Some(s) = maybe_value { - data_size += s.len() * len; - columns.push(ColumnarValueRef::Scalar(s.as_bytes())); - } - } - ColumnarValue::Array(array) => { - match array.data_type() { - DataType::Utf8 => { - let string_array = as_string_array(array)?; - - data_size += string_array.values().len(); - let column = if array.is_nullable() { - ColumnarValueRef::NullableArray(string_array) - } else { - ColumnarValueRef::NonNullableArray(string_array) - }; - columns.push(column); - } - DataType::LargeUtf8 => { - let string_array = as_large_string_array(array)?; - - data_size += string_array.values().len(); - let column = if array.is_nullable() { - ColumnarValueRef::NullableLargeStringArray(string_array) - } else { - ColumnarValueRef::NonNullableLargeStringArray( - string_array, - ) - }; - columns.push(column); - } - DataType::Utf8View => { - let string_array = as_string_view_array(array)?; - - // This is an estimate; in particular, it will - // undercount arrays of short strings (<= 12 bytes). - data_size += string_array.total_buffer_bytes_used(); - let column = if array.is_nullable() { - ColumnarValueRef::NullableStringViewArray(string_array) - } else { - ColumnarValueRef::NonNullableStringViewArray(string_array) - }; - columns.push(column); - } - other => { - return plan_err!( - "Input was {other} which is not a supported datatype for concat_ws function." - ); - } - }; - } - _ => unreachable!(), + if let Some(column) = + ColumnarValueRef::from_columnar_value(arg, &mut data_size, len, 1, false)? + { + columns.push(column); } } match return_datatype { - DataType::Utf8View => { - let mut builder = ConcatStringViewBuilder::with_capacity(len, data_size); - for i in 0..len { - if !sep.is_valid(i) { - builder.append_offset()?; - continue; - } - let mut first = true; - for column in &columns { - if column.is_valid(i) { - if !first { - builder.write::(&sep, i); - } - builder.write::(column, i); - first = false; - } - } - builder.append_offset()?; - } - Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls())?))) - } - DataType::LargeUtf8 => { - let mut builder = ConcatLargeStringBuilder::with_capacity(len, data_size); - for i in 0..len { - if !sep.is_valid(i) { - builder.append_offset()?; - continue; - } - let mut first = true; - for column in &columns { - if column.is_valid(i) { - if !first { - builder.write::(&sep, i); - } - builder.write::(column, i); - first = false; - } - } - builder.append_offset()?; - } - Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls())?))) - } - _ => { - let mut builder = ConcatStringBuilder::with_capacity(len, data_size); - for i in 0..len { - if !sep.is_valid(i) { - builder.append_offset()?; - continue; - } - let mut first = true; - for column in &columns { - if column.is_valid(i) { - if !first { - builder.write::(&sep, i); - } - builder.write::(column, i); - first = false; - } - } - builder.append_offset()?; - } - Ok(ColumnarValue::Array(Arc::new(builder.finish(sep.nulls())?))) - } + DataType::Utf8 => build_concat_ws( + ConcatStringBuilder::with_capacity(len, data_size), + &sep, + &columns, + len, + ), + DataType::LargeUtf8 => build_concat_ws( + ConcatLargeStringBuilder::with_capacity(len, data_size), + &sep, + &columns, + len, + ), + DataType::Utf8View => build_concat_ws( + ConcatStringViewBuilder::with_capacity(len, data_size), + &sep, + &columns, + len, + ), + DataType::Binary => build_concat_ws( + ConcatBinaryBuilder::with_capacity(len, data_size), + &sep, + &columns, + len, + ), + DataType::LargeBinary => build_concat_ws( + ConcatLargeBinaryBuilder::with_capacity(len, data_size), + &sep, + &columns, + len, + ), + DataType::BinaryView => build_concat_ws( + ConcatBinaryViewBuilder::with_capacity(len, data_size), + &sep, + &columns, + len, + ), + other => plan_err!("concat_ws does not support return type {other}"), } } @@ -398,6 +324,41 @@ impl ScalarUDFImpl for ConcatWsFunc { } } +/// Build a `concat_ws` output array using a generic [`ConcatBuilder`]. +/// Write non-null column values per row, inserting the separator between them +fn build_concat_ws( + mut builder: B, + sep: &ColumnarValueRef, + columns: &[ColumnarValueRef], + len: usize, +) -> Result { + for i in 0..len { + if !sep.is_valid(i) { + builder.append_offset()?; + continue; + } + let mut first = true; + for column in columns { + if column.is_valid(i) { + if !first { + builder.write::(sep, i)?; + } + builder.write::(column, i)?; + first = false; + } + } + builder.append_offset()?; + } + let array = builder.finish(sep.nulls())?; + Ok(ColumnarValue::Array(array)) +} + +fn null_scalar(dt: &DataType) -> ColumnarValue { + ColumnarValue::Scalar( + ScalarValue::try_new_null(dt).unwrap_or(ScalarValue::Utf8(None)), + ) +} + fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result { // Preserve the delimiter's string type for any new literals produced // during simplification. @@ -406,6 +367,17 @@ fn simplify_concat_ws(delimiter: &Expr, args: &[Expr]) -> Result DataType::Utf8, }; + // Shortcut for binary delimiters + if delimiter_type.is_binary() { + let mut args = args + .iter() + .filter(|x| !is_null(x)) + .cloned() + .collect::>(); + args.insert(0, delimiter.clone()); + return Ok(ExprSimplifyResult::Original(args)); + } + let typed_lit = |s: String| -> Expr { match delimiter_type { DataType::LargeUtf8 => lit(ScalarValue::LargeUtf8(Some(s))), @@ -532,8 +504,11 @@ mod tests { use std::sync::Arc; use crate::string::concat_ws::ConcatWsFunc; - use arrow::array::{Array, ArrayRef, LargeStringArray, StringArray, StringViewArray}; - use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View}; + use arrow::array::{ + Array, ArrayRef, BinaryArray, LargeBinaryArray, LargeStringArray, StringArray, + StringViewArray, + }; + use arrow::datatypes::DataType::{Binary, LargeBinary, LargeUtf8, Utf8, Utf8View}; use arrow::datatypes::Field; use datafusion_common::Result; use datafusion_common::ScalarValue; @@ -934,4 +909,87 @@ mod tests { Ok(()) } + + #[test] + fn concat_ws_binary_scalars() -> Result<()> { + let c0 = ColumnarValue::Scalar(ScalarValue::Binary(Some(b"|".to_vec()))); + let c1 = ColumnarValue::Scalar(ScalarValue::Binary(Some(b"aa".to_vec()))); + let c2 = ColumnarValue::Scalar(ScalarValue::Binary(None)); + let c3 = ColumnarValue::Scalar(ScalarValue::Binary(Some(b"cc".to_vec()))); + + let arg_fields = vec![ + Field::new("a", Binary, true).into(), + Field::new("a", Binary, true).into(), + Field::new("a", Binary, true).into(), + Field::new("a", Binary, true).into(), + ]; + let args = ScalarFunctionArgs { + args: vec![c0, c1, c2, c3], + arg_fields, + number_rows: 1, + return_field: Field::new("f", Binary, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }; + let result = ConcatWsFunc::new().invoke_with_args(args)?; + match result { + ColumnarValue::Scalar(ScalarValue::Binary(Some(v))) => { + assert_eq!(v, b"aa|cc"); + } + other => panic!("Expected Binary scalar, got {other:?}"), + } + + Ok(()) + } + + #[test] + fn concat_ws_binary_arrays() -> Result<()> { + for c1_large_binary in [false, true] { + let c0 = ColumnarValue::Scalar(ScalarValue::Binary(Some(b",".to_vec()))); + let c1 = if c1_large_binary { + ColumnarValue::Array(Arc::new(LargeBinaryArray::from_vec(vec![ + b"foo".as_ref(), + b"bar", + b"baz", + ]))) + } else { + ColumnarValue::Array(Arc::new(BinaryArray::from_vec(vec![ + b"foo".as_ref(), + b"bar", + b"baz", + ]))) + }; + let c2 = + ColumnarValue::Array(Arc::new(LargeBinaryArray::from_opt_vec(vec![ + Some(b"x".as_ref()), + None, + Some(b"z"), + ]))); + + let arg_fields = vec![ + Field::new("a", Binary, true).into(), + Field::new("a", Binary, true).into(), + Field::new("a", LargeBinary, true).into(), + ]; + let args = ScalarFunctionArgs { + args: vec![c0, c1, c2], + arg_fields, + number_rows: 3, + return_field: Field::new("f", LargeBinary, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }; + + let result = ConcatWsFunc::new().invoke_with_args(args)?; + let expected = Arc::new(LargeBinaryArray::from_opt_vec(vec![ + Some(b"foo,x".as_ref()), + Some(b"bar"), + Some(b"baz,z"), + ])) as ArrayRef; + match &result { + ColumnarValue::Array(array) => assert_eq!(&expected, array), + _ => panic!("Expected array result"), + } + } + + Ok(()) + } } diff --git a/datafusion/functions/src/string/lower.rs b/datafusion/functions/src/string/lower.rs index 57cbe1d8779f0..88f2c800e9e0c 100644 --- a/datafusion/functions/src/string/lower.rs +++ b/datafusion/functions/src/string/lower.rs @@ -21,8 +21,8 @@ use crate::string::common::to_lower; use datafusion_common::Result; use datafusion_common::types::logical_string; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -57,9 +57,10 @@ impl LowerFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } diff --git a/datafusion/functions/src/string/ltrim.rs b/datafusion/functions/src/string/ltrim.rs index e49ffeb0541ff..04e33253ed6df 100644 --- a/datafusion/functions/src/string/ltrim.rs +++ b/datafusion/functions/src/string/ltrim.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, OffsetSizeTrait}; +use arrow::array::{ArrayRef, AsArray}; use arrow::datatypes::DataType; use std::sync::Arc; @@ -25,22 +25,33 @@ use datafusion_common::types::logical_string; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; /// Returns the longest string with leading characters removed. If the characters are not specified, spaces are removed. /// ltrim('zzzytest', 'xyz') = 'test' -fn ltrim(args: &[ArrayRef]) -> Result { - let use_string_view = args[0].data_type() == &DataType::Utf8View; +fn ltrim(args: &[ArrayRef]) -> Result { let args = if args.len() > 1 { let arg1 = arrow::compute::kernels::cast::cast(&args[1], args[0].data_type())?; vec![Arc::clone(&args[0]), arg1] } else { args.to_owned() }; - general_trim::(&args, use_string_view) + match args[0].data_type() { + DataType::Utf8 => general_trim::(&args, false), + DataType::LargeUtf8 => general_trim::(&args, false), + DataType::Utf8View => general_trim::(&args, true), + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let trimmed = ltrim(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(trimmed)) + } + other => exec_err!( + "Unsupported data type {other:?} for function ltrim, expected Utf8, LargeUtf8 or Utf8View." + ), + } } #[user_doc( @@ -90,9 +101,12 @@ impl LtrimFunc { Coercion::new_exact(TypeSignatureClass::Native(logical_string())), Coercion::new_exact(TypeSignatureClass::Native(logical_string())), ]), - TypeSignature::Coercible(vec![Coercion::new_exact( - TypeSignatureClass::Native(logical_string()), - )]), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::dictionary(), + ), + ]), ], Volatility::Immutable, ), @@ -114,20 +128,7 @@ impl ScalarUDFImpl for LtrimFunc { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - match args.args[0].data_type() { - DataType::Utf8 | DataType::Utf8View => make_scalar_function( - ltrim::, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - DataType::LargeUtf8 => make_scalar_function( - ltrim::, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - other => exec_err!( - "Unsupported data type {other:?} for function ltrim,\ - expected Utf8, LargeUtf8 or Utf8View." - ), - } + make_scalar_function(ltrim, vec![Hint::Pad, Hint::AcceptsSingular])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/string/octet_length.rs b/datafusion/functions/src/string/octet_length.rs index ecffb2a6de7af..02df262ee27aa 100644 --- a/datafusion/functions/src/string/octet_length.rs +++ b/datafusion/functions/src/string/octet_length.rs @@ -18,13 +18,13 @@ use arrow::compute::kernels::length::length; use arrow::datatypes::DataType; -use crate::utils::utf8_to_int_type; +use crate::utils::{transform_leaf_type_preserving_encoding, utf8_to_int_type}; use datafusion_common::types::logical_string; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -59,9 +59,10 @@ impl OctetLengthFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -78,7 +79,9 @@ impl ScalarUDFImpl for OctetLengthFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - utf8_to_int_type(&arg_types[0], "octet_length") + transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| { + utf8_to_int_type(data_type, "octet_length") + }) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -86,18 +89,7 @@ impl ScalarUDFImpl for OctetLengthFunc { match array { ColumnarValue::Array(v) => Ok(ColumnarValue::Array(length(v.as_ref())?)), - ColumnarValue::Scalar(v) => match v { - ScalarValue::Utf8(v) => Ok(ColumnarValue::Scalar(ScalarValue::Int32( - v.as_ref().map(|x| x.len() as i32), - ))), - ScalarValue::LargeUtf8(v) => Ok(ColumnarValue::Scalar( - ScalarValue::Int64(v.as_ref().map(|x| x.len() as i64)), - )), - ScalarValue::Utf8View(v) => Ok(ColumnarValue::Scalar( - ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)), - )), - _ => unreachable!("OctetLengthFunc"), - }, + ColumnarValue::Scalar(v) => Ok(ColumnarValue::Scalar(octet_length_scalar(v))), } } @@ -106,6 +98,23 @@ impl ScalarUDFImpl for OctetLengthFunc { } } +fn octet_length_scalar(value: &ScalarValue) -> ScalarValue { + match value { + ScalarValue::Utf8(v) => ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)), + ScalarValue::LargeUtf8(v) => { + ScalarValue::Int64(v.as_ref().map(|x| x.len() as i64)) + } + ScalarValue::Utf8View(v) => { + ScalarValue::Int32(v.as_ref().map(|x| x.len() as i32)) + } + ScalarValue::Dictionary(key_type, value) => ScalarValue::Dictionary( + key_type.clone(), + Box::new(octet_length_scalar(value)), + ), + _ => unreachable!("OctetLengthFunc"), + } +} + #[cfg(test)] mod tests { use std::sync::Arc; diff --git a/datafusion/functions/src/string/repeat.rs b/datafusion/functions/src/string/repeat.rs index b551d2ac707a9..a53f1e2e4fc42 100644 --- a/datafusion/functions/src/string/repeat.rs +++ b/datafusion/functions/src/string/repeat.rs @@ -26,7 +26,9 @@ use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View}; use datafusion_common::cast::as_int64_array; use datafusion_common::types::{NativeType, logical_int64, logical_string}; use datafusion_common::utils::take_function_args; -use datafusion_common::{DataFusionError, Result, ScalarValue, exec_err, internal_err}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err, internal_err, +}; use datafusion_expr::{ColumnarValue, Documentation, Volatility}; use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Signature}; use datafusion_expr_common::signature::{Coercion, TypeSignatureClass}; @@ -166,7 +168,21 @@ fn compute_repeat(s: &str, count: i64, max_size: usize) -> Result { if count <= 0 { return Ok(String::new()); } - let result_len = s.len().saturating_mul(count as usize); + let result_len = repeat_len(s.len(), count, max_size)?; + debug_assert!(result_len <= max_size); + let count = repeat_count(count, max_size)?; + Ok(s.repeat(count)) +} + +fn repeat_len(string_len: usize, count: i64, max_size: usize) -> Result { + let count = repeat_count(count, max_size)?; + let result_len = string_len.checked_mul(count).ok_or_else(|| { + exec_datafusion_err!( + "string size overflow on repeat, max size is {}, but got {}", + max_size, + usize::MAX + ) + })?; if result_len > max_size { return exec_err!( "string size overflow on repeat, max size is {}, but got {}", @@ -174,7 +190,18 @@ fn compute_repeat(s: &str, count: i64, max_size: usize) -> Result { result_len ); } - Ok(s.repeat(count as usize)) + Ok(result_len) +} + +fn repeat_count(count: i64, max_size: usize) -> Result { + match usize::try_from(count) { + Ok(count) => Ok(count), + Err(_) => exec_err!( + "string size overflow on repeat, max size is {}, but got {}", + max_size, + usize::MAX + ), + } } /// Repeats string the specified number of times. @@ -227,22 +254,22 @@ fn calculate_capacities<'a, S>( where S: StringArrayType<'a>, { - let mut total_capacity = 0; - let mut max_item_capacity = 0; + let mut total_capacity = 0usize; + let mut max_item_capacity = 0usize; string_array.iter().zip(number_array.iter()).try_for_each( |(string, number)| -> Result<(), DataFusionError> { match (string, number) { (Some(string), Some(number)) if number >= 0 => { - let item_capacity = string.len() * number as usize; - if item_capacity > max_str_len { - return exec_err!( - "string size overflow on repeat, max size is {}, but got {}", - max_str_len, - number as usize * string.len() - ); - } - total_capacity += item_capacity; + let item_capacity = repeat_len(string.len(), number, max_str_len)?; + total_capacity = + total_capacity.checked_add(item_capacity).ok_or_else(|| { + exec_datafusion_err!( + "string size overflow on repeat, max size is {}, but got {}", + max_str_len, + usize::MAX + ) + })?; max_item_capacity = max_item_capacity.max(item_capacity); } _ => (), @@ -487,6 +514,18 @@ mod tests { assert_sliced_offset_output::(result); } + #[test] + fn test_repeat_string_array_overflow() { + let strings: ArrayRef = Arc::new(StringArray::from(vec![Some("abc")])); + let counts: ArrayRef = Arc::new(Int64Array::from(vec![Some(i64::MAX)])); + + let err = super::repeat(&strings, &counts).unwrap_err().to_string(); + assert!( + err.contains("string size overflow on repeat"), + "unexpected error: {err}" + ); + } + #[test] fn test_repeat_sliced_large_string_with_null_offset() { let (strings, counts) = diff --git a/datafusion/functions/src/string/replace.rs b/datafusion/functions/src/string/replace.rs index 769727999ea05..549b8e1a3b0f9 100644 --- a/datafusion/functions/src/string/replace.rs +++ b/datafusion/functions/src/string/replace.rs @@ -17,13 +17,12 @@ use std::sync::Arc; -use arrow::array::{Array, ArrayRef, OffsetSizeTrait}; +use arrow::array::{ArrayRef, OffsetSizeTrait, StringArrayType}; use arrow::buffer::NullBuffer; use arrow::datatypes::DataType; +use memchr::memmem; -use crate::strings::{ - BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringWriter, -}; +use crate::strings::{GenericStringArrayBuilder, StringWriter}; use crate::utils::{make_scalar_function, utf8_to_str_type}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; use datafusion_common::types::logical_string; @@ -130,6 +129,43 @@ impl ScalarUDFImpl for ReplaceFunc { } } + // Fast path: when `from` and `to` are non-null scalars we can + // pre-build a substring finder once and reuse it for every haystack + // row, mirroring the scalar-argument fast paths in + // `strpos`/`translate`/`split_part`. + if let ( + ColumnarValue::Array(haystack), + ColumnarValue::Scalar(from), + ColumnarValue::Scalar(to), + ) = (&converted_args[0], &converted_args[1], &converted_args[2]) + && let (Some(Some(from)), Some(Some(to))) = + (from.try_as_str(), to.try_as_str()) + { + let result = match coercion_type { + DataType::Utf8 => replace_scalar::<_, i32>( + as_generic_string_array::(haystack)?, + from, + to, + ), + DataType::LargeUtf8 => replace_scalar::<_, i64>( + as_generic_string_array::(haystack)?, + from, + to, + ), + DataType::Utf8View => replace_scalar::<_, i32>( + as_string_view_array(haystack)?, + from, + to, + ), + other => { + return exec_err!( + "Unsupported coercion data type {other:?} for function replace" + ); + } + }; + return result.map(ColumnarValue::Array); + } + match coercion_type { DataType::Utf8 => { make_scalar_function(replace::, vec![])(&converted_args) @@ -164,40 +200,7 @@ fn replace_view(args: &[ArrayRef]) -> Result { let from_array = as_string_view_array(&args[1])?; let to_array = as_string_view_array(&args[2])?; - let len = string_array.len(); - let mut builder = GenericStringArrayBuilder::::with_capacity(len, 0); - let nulls = NullBuffer::union_many([ - string_array.nulls(), - from_array.nulls(), - to_array.nulls(), - ]); - - // Hoist the nulls.is_some() check out of the loop. LLVM does not always - // unswitch this loop on its own (the Utf8View body is large enough to - // exceed its cost-benefit threshold). - if let Some(nulls_ref) = nulls.as_ref() { - for i in 0..len { - if nulls_ref.is_null(i) { - builder.append_placeholder(); - continue; - } - // SAFETY: union of input nulls is non-null at i, so each input is too. - let string = unsafe { string_array.value_unchecked(i) }; - let from = unsafe { from_array.value_unchecked(i) }; - let to = unsafe { to_array.value_unchecked(i) }; - apply_replace(&mut builder, string, from, to); - } - } else { - for i in 0..len { - // SAFETY: i < len, and no input has a null buffer. - let string = unsafe { string_array.value_unchecked(i) }; - let from = unsafe { from_array.value_unchecked(i) }; - let to = unsafe { to_array.value_unchecked(i) }; - apply_replace(&mut builder, string, from, to); - } - } - - Ok(Arc::new(builder.finish(nulls)?) as ArrayRef) + replace_arrays::<_, i32>(string_array, from_array, to_array) } /// Replaces all occurrences in string of substring from with substring to. @@ -207,49 +210,68 @@ fn replace(args: &[ArrayRef]) -> Result { let from_array = as_generic_string_array::(&args[1])?; let to_array = as_generic_string_array::(&args[2])?; + replace_arrays::<_, T>(string_array, from_array, to_array) +} + +fn replace_arrays<'a, S, O>( + string_array: S, + from_array: S, + to_array: S, +) -> Result +where + S: StringArrayType<'a> + Copy, + O: OffsetSizeTrait, +{ let len = string_array.len(); - let mut builder = GenericStringArrayBuilder::::with_capacity(len, 0); let nulls = NullBuffer::union_many([ string_array.nulls(), from_array.nulls(), to_array.nulls(), ]); + build_replaced::(len, nulls, |builder, i| { + // SAFETY: build_replaced only calls this for rows that are non-null in + // the union buffer, so every input array is non-null at i. + let string = unsafe { string_array.value_unchecked(i) }; + let from = unsafe { from_array.value_unchecked(i) }; + let to = unsafe { to_array.value_unchecked(i) }; + apply_replace(builder, string, from, to, None) + }) +} - // Hoist the nulls.is_some() check out of the loop. LLVM unswitches this - // automatically today, but kept explicit so the no-nulls fast path is not - // contingent on the optimizer's cost heuristic. +/// Appends `len` rows to a fresh string builder: a null placeholder for each +/// null row and `append_row` for each non-null row. The `nulls.is_some()` check +/// is hoisted out of the loop so the all-non-null case does not depend on LLVM +/// loop-unswitching heuristics. +fn build_replaced( + len: usize, + nulls: Option, + mut append_row: impl FnMut(&mut GenericStringArrayBuilder, usize) -> Result<()>, +) -> Result { + let mut builder = GenericStringArrayBuilder::::with_capacity(len, 0); if let Some(nulls_ref) = nulls.as_ref() { for i in 0..len { if nulls_ref.is_null(i) { - builder.append_placeholder(); - continue; + builder.try_append_placeholder()?; + } else { + append_row(&mut builder, i)?; } - // SAFETY: union of input nulls is non-null at i, so each input is too. - let string = unsafe { string_array.value_unchecked(i) }; - let from = unsafe { from_array.value_unchecked(i) }; - let to = unsafe { to_array.value_unchecked(i) }; - apply_replace(&mut builder, string, from, to); } } else { for i in 0..len { - // SAFETY: i < len, and no input has a null buffer. - let string = unsafe { string_array.value_unchecked(i) }; - let from = unsafe { from_array.value_unchecked(i) }; - let to = unsafe { to_array.value_unchecked(i) }; - apply_replace(&mut builder, string, from, to); + append_row(&mut builder, i)?; } } - Ok(Arc::new(builder.finish(nulls)?) as ArrayRef) } #[inline] -fn apply_replace( - builder: &mut B, +fn apply_replace( + builder: &mut GenericStringArrayBuilder, string: &str, from: &str, to: &str, -) { + finder: Option<&memmem::Finder>, +) -> Result<()> { // Hot path: single ASCII byte → single ASCII byte. An ASCII byte (< 0x80) // cannot appear inside a multi-byte UTF-8 sequence, so any multi-byte // sequences in `string` pass through unchanged and output stays valid @@ -259,44 +281,101 @@ fn apply_replace( && to_byte.is_ascii() { // SAFETY: see the contract above. - unsafe { - builder.append_byte_map(string.as_bytes(), |b| { + return unsafe { + builder.try_append_byte_map(string.as_bytes(), |b| { if b == from_byte { to_byte } else { b } - }); - } - return; + }) + }; } if from.is_empty() { - // Empty `from`: insert `to` before each character and at both ends. - builder.append_with(|w| { - w.write_str(to); - for ch in string.chars() { - w.write_char(ch); - w.write_str(to); - } - }); - return; + // PostgreSQL returns the input unchanged when `from` is empty (#22253). + return builder.try_append_value(string); } - builder.append_with(|w| replace_into_writer(w, string, from, to)); + builder.try_append_with(|w| replace_into_writer(w, string, from, to, finder)) +} + +/// Writes `string` into `w` with every non-overlapping occurrence of `from` +/// replaced by `to`. When `finder` is `Some`, matches are located with the +/// pre-built finder (the scalar fast path, where `from` is constant across all +/// rows); otherwise `str::match_indices` builds a searcher per call. +/// +/// Both `string` and `from` are valid UTF-8, and UTF-8 is self-synchronizing, +/// so a byte match of `from` can only start on a char boundary of `string`; the +/// slices below are therefore always valid. +#[inline] +fn replace_into_writer( + w: &mut W, + string: &str, + from: &str, + to: &str, + finder: Option<&memmem::Finder>, +) { + match finder { + Some(finder) => write_replaced( + w, + string, + to, + from.len(), + finder.find_iter(string.as_bytes()), + ), + None => write_replaced( + w, + string, + to, + from.len(), + string.match_indices(from).map(|(start, _)| start), + ), + } } +/// Copies `string` into `w`, replacing the `from_len`-byte substring at each +/// byte offset yielded by `starts` with `to`. `starts` must be ascending and +/// non-overlapping, as produced by both `memmem::Finder::find_iter` and +/// `str::match_indices`. #[inline] -fn replace_into_writer(w: &mut W, string: &str, from: &str, to: &str) { +fn write_replaced( + w: &mut W, + string: &str, + to: &str, + from_len: usize, + starts: impl Iterator, +) { let mut last_end = 0; - for (start, _part) in string.match_indices(from) { + for start in starts { w.write_str(&string[last_end..start]); w.write_str(to); - last_end = start + from.len(); + last_end = start + from_len; } w.write_str(&string[last_end..]); } +/// Fast path for a `from`/`to` pair that is constant across all rows. The +/// substring finder is built once and reused for every haystack value, which +/// avoids the per-row searcher construction incurred by `str::match_indices`. +fn replace_scalar<'a, S, O>(haystack: S, from: &str, to: &str) -> Result +where + S: StringArrayType<'a> + Copy, + O: OffsetSizeTrait, +{ + // `from` and `to` are non-null scalars, so the output nulls are exactly the + // haystack's nulls (matching the null union computed by the general path). + let nulls = haystack.nulls().cloned(); + // Built once and reused for every row. + let finder = memmem::Finder::new(from.as_bytes()); + build_replaced::(haystack.len(), nulls, |builder, i| { + // SAFETY: build_replaced only calls this for non-null rows. + let string = unsafe { haystack.value_unchecked(i) }; + apply_replace(builder, string, from, to, Some(&finder)) + }) +} + #[cfg(test)] mod tests { use super::*; use crate::utils::test::test_function; + use arrow::array::Array; use arrow::array::LargeStringArray; use arrow::array::StringArray; use arrow::datatypes::DataType::{LargeUtf8, Utf8}; @@ -346,6 +425,105 @@ mod tests { StringArray ); + test_function!( + ReplaceFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(String::from("abc")))), + ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(String::from("")))), + ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(String::from("x")))), + ], + Ok(Some("abc")), + &str, + LargeUtf8, + LargeStringArray + ); + Ok(()) } + + /// The scalar-argument fast path must produce output that is bit-identical + /// to the general (array-argument) path for every kind of pattern. + #[test] + fn scalar_fast_path_matches_general() { + use arrow::array::{ArrayRef, StringViewArray}; + use arrow::datatypes::Field; + use datafusion_common::config::ConfigOptions; + use std::sync::Arc; + + let rows = vec![ + Some("hello world"), + None, + Some("aaaa"), + Some(""), + Some("a.b.c.d"), + Some("úñîçödé abcúñ"), + Some("mississippi"), + Some(" double spaces "), + ]; + // Covers byte-map (single ASCII → single ASCII), deletion (empty `to`), + // empty `from`, multi-byte `to`, and multi-byte non-ASCII `from`. + let cases = [ + (" ", "_"), + ("a", "X"), + ("ss", "Z"), + ("", "Q"), + ("a", "yy"), + ("úñ", "A"), + (".", ""), + ("i", "II"), + ]; + + let invoke = |haystack: &ArrayRef, + from: ColumnarValue, + to: ColumnarValue| + -> ArrayRef { + let args = vec![ColumnarValue::Array(Arc::clone(haystack)), from, to]; + let arg_fields = args + .iter() + .enumerate() + .map(|(i, a)| Field::new(format!("a{i}"), a.data_type(), true).into()) + .collect(); + match ReplaceFunc::new() + .invoke_with_args(ScalarFunctionArgs { + args, + arg_fields, + number_rows: haystack.len(), + return_field: Field::new("f", Utf8, true).into(), + config_options: Arc::new(ConfigOptions::default()), + }) + .unwrap() + { + ColumnarValue::Array(a) => a, + ColumnarValue::Scalar(s) => s.to_array_of_size(haystack.len()).unwrap(), + } + }; + + for (from, to) in cases { + let n = rows.len(); + for haystack in [ + Arc::new(StringArray::from(rows.clone())) as ArrayRef, + Arc::new(LargeStringArray::from(rows.clone())) as ArrayRef, + Arc::new(StringViewArray::from(rows.clone())) as ArrayRef, + ] { + // scalar `from`/`to` -> new fast path + let fast = invoke( + &haystack, + ColumnarValue::Scalar(ScalarValue::Utf8(Some(from.to_string()))), + ColumnarValue::Scalar(ScalarValue::Utf8(Some(to.to_string()))), + ); + // array `from`/`to` -> general path + let general = invoke( + &haystack, + ColumnarValue::Array(Arc::new(StringArray::from(vec![from; n]))), + ColumnarValue::Array(Arc::new(StringArray::from(vec![to; n]))), + ); + assert_eq!( + &fast, + &general, + "mismatch for from={from:?} to={to:?} on {:?}", + haystack.data_type() + ); + } + } + } } diff --git a/datafusion/functions/src/string/rtrim.rs b/datafusion/functions/src/string/rtrim.rs index 05ad9e855976d..d1126cb2418ae 100644 --- a/datafusion/functions/src/string/rtrim.rs +++ b/datafusion/functions/src/string/rtrim.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, OffsetSizeTrait}; +use arrow::array::{ArrayRef, AsArray}; use arrow::datatypes::DataType; use std::sync::Arc; @@ -25,22 +25,33 @@ use datafusion_common::types::logical_string; use datafusion_common::{Result, exec_err}; use datafusion_expr::function::Hint; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignature, TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; /// Returns the longest string with trailing characters removed. If the characters are not specified, spaces are removed. /// rtrim('testxxzx', 'xyz') = 'test' -fn rtrim(args: &[ArrayRef]) -> Result { - let use_string_view = args[0].data_type() == &DataType::Utf8View; +fn rtrim(args: &[ArrayRef]) -> Result { let args = if args.len() > 1 { let arg1 = arrow::compute::kernels::cast::cast(&args[1], args[0].data_type())?; vec![Arc::clone(&args[0]), arg1] } else { args.to_owned() }; - general_trim::(&args, use_string_view) + match args[0].data_type() { + DataType::Utf8 => general_trim::(&args, false), + DataType::LargeUtf8 => general_trim::(&args, false), + DataType::Utf8View => general_trim::(&args, true), + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let trimmed = rtrim(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(trimmed)) + } + other => exec_err!( + "Unsupported data type {other:?} for function rtrim, expected Utf8, LargeUtf8 or Utf8View." + ), + } } #[user_doc( @@ -90,9 +101,12 @@ impl RtrimFunc { Coercion::new_exact(TypeSignatureClass::Native(logical_string())), Coercion::new_exact(TypeSignatureClass::Native(logical_string())), ]), - TypeSignature::Coercible(vec![Coercion::new_exact( - TypeSignatureClass::Native(logical_string()), - )]), + TypeSignature::Coercible(vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation( + EncodingPreservation::dictionary(), + ), + ]), ], Volatility::Immutable, ), @@ -114,20 +128,7 @@ impl ScalarUDFImpl for RtrimFunc { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - match args.args[0].data_type() { - DataType::Utf8 | DataType::Utf8View => make_scalar_function( - rtrim::, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - DataType::LargeUtf8 => make_scalar_function( - rtrim::, - vec![Hint::Pad, Hint::AcceptsSingular], - )(&args.args), - other => exec_err!( - "Unsupported data type {other:?} for function rtrim,\ - expected Utf8, LargeUtf8 or Utf8View." - ), - } + make_scalar_function(rtrim, vec![Hint::Pad, Hint::AcceptsSingular])(&args.args) } fn documentation(&self) -> Option<&Documentation> { diff --git a/datafusion/functions/src/string/split_part.rs b/datafusion/functions/src/string/split_part.rs index 1994c65bcf326..9b73a1af88501 100644 --- a/datafusion/functions/src/string/split_part.rs +++ b/datafusion/functions/src/string/split_part.rs @@ -15,13 +15,15 @@ // specific language governing permissions and limitations // under the License. +use crate::strings::{ + BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringViewArrayBuilder, +}; use crate::utils::utf8_to_str_type; use arrow::array::{ - Array, ArrayRef, AsArray, ByteView, GenericStringBuilder, Int64Array, - StringArrayType, StringLikeArrayBuilder, StringViewArray, StringViewBuilder, + Array, ArrayRef, AsArray, ByteView, Int64Array, StringArrayType, StringViewArray, make_view, new_null_array, }; -use arrow::buffer::ScalarBuffer; +use arrow::buffer::{NullBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::ScalarValue; use datafusion_common::cast::as_int64_array; @@ -167,7 +169,7 @@ impl ScalarUDFImpl for SplitPartFunc { let result = match args[0].data_type() { DataType::Utf8View => split_part_for_delimiter_type!( &args[0].as_string_view(), - StringViewBuilder::with_capacity(inferred_length) + StringViewArrayBuilder::with_capacity(inferred_length) ), DataType::Utf8 => { let str_arr = &args[0].as_string::(); @@ -176,7 +178,7 @@ impl ScalarUDFImpl for SplitPartFunc { // pre-allocating the full input data size. split_part_for_delimiter_type!( str_arr, - GenericStringBuilder::::with_capacity( + GenericStringArrayBuilder::::with_capacity( inferred_length, inferred_length, ) @@ -187,7 +189,7 @@ impl ScalarUDFImpl for SplitPartFunc { // Conservative under-estimate; see Utf8 comment above. split_part_for_delimiter_type!( str_arr, - GenericStringBuilder::::with_capacity( + GenericStringArrayBuilder::::with_capacity( inferred_length, inferred_length, ) @@ -293,7 +295,7 @@ fn split_part_scalar( arr, delimiter, position, - GenericStringBuilder::::with_capacity(arr.len(), arr.len()), + GenericStringArrayBuilder::::with_capacity(arr.len(), arr.len()), ) } DataType::LargeUtf8 => { @@ -303,7 +305,7 @@ fn split_part_scalar( arr, delimiter, position, - GenericStringBuilder::::with_capacity(arr.len(), arr.len()), + GenericStringArrayBuilder::::with_capacity(arr.len(), arr.len()), ) } other => exec_err!("Unsupported string type {other:?} for split_part"), @@ -323,7 +325,7 @@ fn split_part_scalar_impl<'a, S, B>( ) -> Result where S: StringArrayType<'a> + Copy, - B: StringLikeArrayBuilder, + B: BulkNullStringArrayBuilder, { if delimiter.is_empty() { // PostgreSQL: empty delimiter treats input as a single field, @@ -367,16 +369,31 @@ where fn map_strings<'a, S, B, F>(string_array: S, mut builder: B, f: F) -> Result where S: StringArrayType<'a> + Copy, - B: StringLikeArrayBuilder, + B: BulkNullStringArrayBuilder, F: Fn(&'a str) -> Option<&'a str>, { - for string in string_array.iter() { - match string { - Some(s) => builder.append_value(f(s).unwrap_or("")), - None => builder.append_null(), + let item_len = string_array.len(); + let nulls = string_array.nulls().cloned(); + + if let Some(ref n) = nulls { + for i in 0..item_len { + if n.is_null(i) { + builder.append_placeholder(); + } else { + // SAFETY: `n.is_null(i)` was false in the branch above. + let s = unsafe { string_array.value_unchecked(i) }; + builder.append_value(f(s).unwrap_or("")); + } + } + } else { + for i in 0..item_len { + // SAFETY: no null buffer means every index is valid. + let s = unsafe { string_array.value_unchecked(i) }; + builder.append_value(f(s).unwrap_or("")); } } - Ok(Arc::new(builder.finish()) as ArrayRef) + + builder.finish(nulls) } /// Finds the `n`th (0-based) split part using a pre-built `memmem::Finder`. @@ -390,10 +407,8 @@ fn split_nth_finder<'a>( let bytes = string.as_bytes(); let mut start = 0; for _ in 0..n { - match finder.find(&bytes[start..]) { - Some(pos) => start += pos + delim_len, - None => return None, - } + let pos = finder.find(&bytes[start..])?; + start += pos + delim_len } match finder.find(&bytes[start..]) { Some(pos) => Some(&string[start..start + pos]), @@ -413,10 +428,8 @@ fn rsplit_nth_finder<'a>( let bytes = string.as_bytes(); let mut end = bytes.len(); for _ in 0..n { - match finder.rfind(&bytes[..end]) { - Some(pos) => end = pos, - None => return None, - } + let pos = finder.rfind(&bytes[..end])?; + end = pos } match finder.rfind(&bytes[..end]) { Some(pos) => Some(&string[pos + delim_len..end]), @@ -543,58 +556,82 @@ fn split_part_impl<'a, StringArrType, DelimiterArrType, B>( where StringArrType: StringArrayType<'a>, DelimiterArrType: StringArrayType<'a>, - B: StringLikeArrayBuilder, + B: BulkNullStringArrayBuilder, { - for ((string, delimiter), n) in string_array - .iter() - .zip(delimiter_array.iter()) - .zip(n_array.iter()) - { - match (string, delimiter, n) { - (Some(string), Some(delimiter), Some(n)) => { - let result = match n.cmp(&0) { - std::cmp::Ordering::Greater => { - let idx: usize = (n - 1).try_into().map_err(|_| { - exec_datafusion_err!( - "split_part index {n} exceeds maximum supported value" - ) - })?; - if delimiter.is_empty() { - // Match PostgreSQL's behavior: empty delimiter - // treats input as a single field, so only position - // 1 returns data. - (n == 1).then_some(string) - } else { - split_nth(string, delimiter, idx) - } - } - std::cmp::Ordering::Less => { - let idx: usize = - (n.unsigned_abs() - 1).try_into().map_err(|_| { - exec_datafusion_err!( - "split_part index {n} exceeds minimum supported value" - ) - })?; - if delimiter.is_empty() { - // Match PostgreSQL's behavior: empty delimiter - // treats input as a single field, so only position - // -1 returns data. - (n == -1).then_some(string) - } else { - rsplit_nth(string, delimiter, idx) - } - } - std::cmp::Ordering::Equal => { - return exec_err!("field position must not be zero"); - } - }; - builder.append_value(result.unwrap_or("")); + let nulls = NullBuffer::union_many([ + string_array.nulls(), + delimiter_array.nulls(), + n_array.nulls(), + ]); + + if let Some(ref n) = nulls { + for i in 0..string_array.len() { + if n.is_null(i) { + builder.append_placeholder(); + continue; } - _ => builder.append_null(), + + // SAFETY: the union null buffer is valid at `i`, so each input is valid. + let string = unsafe { string_array.value_unchecked(i) }; + let delimiter = unsafe { delimiter_array.value_unchecked(i) }; + let position = unsafe { n_array.value_unchecked(i) }; + append_split_part(string, delimiter, position, &mut builder)?; + } + } else { + for i in 0..string_array.len() { + // SAFETY: no input has a null buffer, so every index is valid. + let string = unsafe { string_array.value_unchecked(i) }; + let delimiter = unsafe { delimiter_array.value_unchecked(i) }; + let position = unsafe { n_array.value_unchecked(i) }; + append_split_part(string, delimiter, position, &mut builder)?; } } - Ok(Arc::new(builder.finish()) as ArrayRef) + builder.finish(nulls) +} + +#[inline] +fn append_split_part( + string: &str, + delimiter: &str, + n: i64, + builder: &mut B, +) -> Result<()> { + let result = match n.cmp(&0) { + std::cmp::Ordering::Greater => { + let idx: usize = (n - 1).try_into().map_err(|_| { + exec_datafusion_err!( + "split_part index {n} exceeds maximum supported value" + ) + })?; + if delimiter.is_empty() { + // Match PostgreSQL's behavior: empty delimiter treats input + // as a single field, so only position 1 returns data. + (n == 1).then_some(string) + } else { + split_nth(string, delimiter, idx) + } + } + std::cmp::Ordering::Less => { + let idx: usize = (n.unsigned_abs() - 1).try_into().map_err(|_| { + exec_datafusion_err!( + "split_part index {n} exceeds minimum supported value" + ) + })?; + if delimiter.is_empty() { + // Match PostgreSQL's behavior: empty delimiter treats input + // as a single field, so only position -1 returns data. + (n == -1).then_some(string) + } else { + rsplit_nth(string, delimiter, idx) + } + } + std::cmp::Ordering::Equal => { + return exec_err!("field position must not be zero"); + } + }; + builder.append_value(result.unwrap_or("")); + Ok(()) } #[cfg(test)] diff --git a/datafusion/functions/src/string/to_hex.rs b/datafusion/functions/src/string/to_hex.rs index 497a0a1206922..9f239c2aed93e 100644 --- a/datafusion/functions/src/string/to_hex.rs +++ b/datafusion/functions/src/string/to_hex.rs @@ -24,6 +24,7 @@ use arrow::datatypes::{ Int64Type, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; use datafusion_common::cast::as_primitive_array; +use datafusion_common::utils::hex::{HexCase, ToHex}; use datafusion_common::{Result, ScalarValue, exec_err, internal_err}; use datafusion_expr::{ Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -31,9 +32,6 @@ use datafusion_expr::{ }; use datafusion_macros::user_doc; -/// Hex lookup table for fast conversion -const HEX_CHARS: &[u8; 16] = b"0123456789abcdef"; - /// Converts the number to its equivalent hexadecimal representation. /// to_hex(2147483647) = '7fffffff' fn to_hex_array(array: &ArrayRef) -> Result @@ -59,8 +57,7 @@ where // Process all values directly (including null slots - we write empty strings for nulls) // The null bitmap will mark which entries are actually null for value in integer_array.values() { - let hex_len = value.write_hex_to_buffer(&mut hex_buffer); - values.extend_from_slice(&hex_buffer[16 - hex_len..]); + values.extend_from_slice(value.write_hex(HexCase::Lower, &mut hex_buffer)); offsets.push(values.len() as i32); } @@ -79,100 +76,9 @@ where #[inline] fn to_hex_scalar(value: T) -> String { let mut hex_buffer = [0u8; 16]; - let hex_len = value.write_hex_to_buffer(&mut hex_buffer); - // SAFETY: hex_buffer is ASCII hex digits - unsafe { std::str::from_utf8_unchecked(&hex_buffer[16 - hex_len..]).to_string() } -} - -/// Trait for converting integer types to hexadecimal in a buffer -trait ToHex: ArrowNativeType { - /// Write hex representation to buffer and return the number of hex digits written. - /// The hex digits are written right-aligned in the buffer (starting from position 16 - len). - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize; -} - -/// Write unsigned value to hex buffer and return the number of digits written. -/// Digits are written right-aligned in the buffer. -#[inline] -fn write_unsigned_hex_to_buffer(value: u64, buffer: &mut [u8; 16]) -> usize { - if value == 0 { - buffer[15] = b'0'; - return 1; - } - - // Write hex digits from right to left - let mut pos = 16; - let mut v = value; - while v > 0 { - pos -= 1; - buffer[pos] = HEX_CHARS[(v & 0xf) as usize]; - v >>= 4; - } - - 16 - pos -} - -/// Write signed value to hex buffer (two's complement for negative) and return digit count -#[inline] -fn write_signed_hex_to_buffer(value: i64, buffer: &mut [u8; 16]) -> usize { - // For negative values, use two's complement representation (same as casting to u64) - write_unsigned_hex_to_buffer(value as u64, buffer) -} - -impl ToHex for i8 { - #[inline] - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize { - write_signed_hex_to_buffer(self as i64, buffer) - } -} - -impl ToHex for i16 { - #[inline] - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize { - write_signed_hex_to_buffer(self as i64, buffer) - } -} - -impl ToHex for i32 { - #[inline] - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize { - write_signed_hex_to_buffer(self as i64, buffer) - } -} - -impl ToHex for i64 { - #[inline] - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize { - write_signed_hex_to_buffer(self, buffer) - } -} - -impl ToHex for u8 { - #[inline] - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize { - write_unsigned_hex_to_buffer(self as u64, buffer) - } -} - -impl ToHex for u16 { - #[inline] - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize { - write_unsigned_hex_to_buffer(self as u64, buffer) - } -} - -impl ToHex for u32 { - #[inline] - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize { - write_unsigned_hex_to_buffer(self as u64, buffer) - } -} - -impl ToHex for u64 { - #[inline] - fn write_hex_to_buffer(self, buffer: &mut [u8; 16]) -> usize { - write_unsigned_hex_to_buffer(self, buffer) - } + let hex = value.write_hex(HexCase::Lower, &mut hex_buffer); + // SAFETY: hex holds only ASCII hex digits. + unsafe { std::str::from_utf8_unchecked(hex).to_string() } } #[user_doc( diff --git a/datafusion/functions/src/string/upper.rs b/datafusion/functions/src/string/upper.rs index c0ac90b1bc598..789ab2c046203 100644 --- a/datafusion/functions/src/string/upper.rs +++ b/datafusion/functions/src/string/upper.rs @@ -20,8 +20,8 @@ use arrow::datatypes::DataType; use datafusion_common::Result; use datafusion_common::types::logical_string; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -56,9 +56,10 @@ impl UpperFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } diff --git a/datafusion/functions/src/strings.rs b/datafusion/functions/src/strings.rs index 1d02def4765cc..c788c6fb1f33f 100644 --- a/datafusion/functions/src/strings.rs +++ b/datafusion/functions/src/strings.rs @@ -19,18 +19,40 @@ use std::marker::PhantomData; use std::mem::size_of; use std::sync::Arc; -use datafusion_common::{Result, exec_datafusion_err, internal_err}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err, internal_err, + plan_err, +}; use arrow::array::{ - Array, ArrayAccessor, ArrayDataBuilder, ArrayRef, BinaryArray, ByteView, - GenericStringArray, LargeStringArray, OffsetSizeTrait, StringArray, StringViewArray, - make_view, + Array, ArrayAccessor, ArrayDataBuilder, ArrayRef, BinaryArray, BinaryViewArray, + ByteView, GenericStringArray, LargeBinaryArray, LargeStringArray, OffsetSizeTrait, + StringArray, StringViewArray, as_largestring_array, make_view, }; use arrow::buffer::{Buffer, MutableBuffer, NullBuffer, ScalarBuffer}; use arrow::datatypes::DataType; +use arrow_buffer::ArrowNativeType; +use datafusion_common::cast::{ + as_binary_array, as_binary_view_array, as_large_binary_array, as_string_array, + as_string_view_array, +}; +use datafusion_expr_common::columnar_value::ColumnarValue; + +/// Trait abstracting concatenating string and binary collections. +pub(crate) trait ConcatBuilder { + fn write( + &mut self, + column: &ColumnarValueRef, + i: usize, + ) -> Result<()>; -/// Builder used by `concat`/`concat_ws` to assemble a [`StringArray`] one row -/// at a time from multiple input columns. + fn append_offset(&mut self) -> Result<()>; + + fn finish(self, null_buffer: Option) -> Result; +} + +/// Builder used by `concat`/`concat_ws` to assemble a [`GenericStringArray`] +/// (`StringArray` or `LargeStringArray`) one row at a time from multiple input columns. /// /// Each row is written via repeated `write` calls (one per input fragment) /// followed by a single `append_offset` to commit the row. The output null @@ -39,39 +61,46 @@ use arrow::datatypes::DataType; /// /// For the common "produce one `&str` per row" pattern, prefer /// `GenericStringArrayBuilder` instead. -pub(crate) struct ConcatStringBuilder { +pub(crate) struct ConcatGenericStringBuilder { offsets_buffer: MutableBuffer, value_buffer: MutableBuffer, - /// If true, a safety check is required during the `finish` call - tainted: bool, + _phantom: PhantomData, } +pub(crate) type ConcatStringBuilder = ConcatGenericStringBuilder; +pub(crate) type ConcatLargeStringBuilder = ConcatGenericStringBuilder; -impl ConcatStringBuilder { +impl ConcatGenericStringBuilder { pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self { let capacity = item_capacity .checked_add(1) - .map(|i| i.saturating_mul(size_of::())) + .map(|i| i.saturating_mul(size_of::())) .expect("capacity integer overflow"); let mut offsets_buffer = MutableBuffer::with_capacity(capacity); // SAFETY: the first offset value is definitely not going to exceed the bounds. - unsafe { offsets_buffer.push_unchecked(0_i32) }; + unsafe { offsets_buffer.push_unchecked(O::usize_as(0)) }; Self { offsets_buffer, value_buffer: MutableBuffer::with_capacity(data_capacity), - tainted: false, + _phantom: PhantomData, } } +} - pub fn write( +impl ConcatBuilder + for ConcatGenericStringBuilder +{ + fn write( &mut self, column: &ColumnarValueRef, i: usize, - ) { + ) -> Result<()> { match column { ColumnarValueRef::Scalar(s) => { + std::str::from_utf8(s).map_err(|_| { + exec_datafusion_err!("concat: scalar bytes are not valid UTF-8") + })?; self.value_buffer.extend_from_slice(s); - self.tainted = true; } ColumnarValueRef::NullableArray(array) => { if !CHECK_VALID || array.is_valid(i) { @@ -91,12 +120,6 @@ impl ConcatStringBuilder { .extend_from_slice(array.value(i).as_bytes()); } } - ColumnarValueRef::NullableBinaryArray(array) => { - if !CHECK_VALID || array.is_valid(i) { - self.value_buffer.extend_from_slice(array.value(i)); - } - self.tainted = true; - } ColumnarValueRef::NonNullableArray(array) => { self.value_buffer .extend_from_slice(array.value(i).as_bytes()); @@ -109,32 +132,31 @@ impl ConcatStringBuilder { self.value_buffer .extend_from_slice(array.value(i).as_bytes()); } - ColumnarValueRef::NonNullableBinaryArray(array) => { - self.value_buffer.extend_from_slice(array.value(i)); - self.tainted = true; + _ => { + return exec_err!( + "concat: unexpected column type for string builder: {column:?}" + ); } } + Ok(()) } - pub fn append_offset(&mut self) -> Result<()> { - let next_offset: i32 = self - .value_buffer - .len() - .try_into() - .map_err(|_| exec_datafusion_err!("byte array offset overflow"))?; + fn append_offset(&mut self) -> Result<()> { + let next_offset: O = O::from_usize(self.value_buffer.len()) + .ok_or_else(|| exec_datafusion_err!("byte array offset overflow"))?; self.offsets_buffer.push(next_offset); Ok(()) } - /// Finalize the builder into a concrete [`StringArray`]. + /// Finalize the builder into a concrete [`GenericStringArray`]. /// /// # Errors /// /// Returns an error when: /// /// - the provided `null_buffer` is not the same length as the `offsets_buffer`. - pub fn finish(self, null_buffer: Option) -> Result { - let row_count = self.offsets_buffer.len() / size_of::() - 1; + fn finish(self, null_buffer: Option) -> Result { + let row_count = self.offsets_buffer.len() / size_of::() - 1; if let Some(ref null_buffer) = null_buffer && null_buffer.len() != row_count { @@ -142,22 +164,16 @@ impl ConcatStringBuilder { "Null buffer and offsets buffer must be the same length" ); } - let array_builder = ArrayDataBuilder::new(DataType::Utf8) + let array_builder = ArrayDataBuilder::new(GenericStringArray::::DATA_TYPE) .len(row_count) .add_buffer(self.offsets_buffer.into()) .add_buffer(self.value_buffer.into()) .nulls(null_buffer); - if self.tainted { - // Raw binary arrays with possible invalid utf-8 were used, - // so let ArrayDataBuilder perform validation - let array_data = array_builder.build()?; - Ok(StringArray::from(array_data)) - } else { - // SAFETY: all data that was appended was valid UTF8 and the values - // and offsets were created correctly - let array_data = unsafe { array_builder.build_unchecked() }; - Ok(StringArray::from(array_data)) - } + // SAFETY: all data that was appended was valid UTF8 and the values + // and offsets were created correctly + let array_data = unsafe { array_builder.build_unchecked() }; + let array = GenericStringArray::::from(array_data); + Ok(Arc::new(array)) } } @@ -175,8 +191,6 @@ pub(crate) struct ConcatStringViewBuilder { views: Vec, data: Vec, block: Vec, - /// If true, a safety check is required during the `append_offset` call - tainted: bool, } impl ConcatStringViewBuilder { @@ -185,19 +199,22 @@ impl ConcatStringViewBuilder { views: Vec::with_capacity(item_capacity), data: Vec::with_capacity(data_capacity), block: vec![], - tainted: false, } } +} - pub fn write( +impl ConcatBuilder for ConcatStringViewBuilder { + fn write( &mut self, column: &ColumnarValueRef, i: usize, - ) { + ) -> Result<()> { match column { ColumnarValueRef::Scalar(s) => { + std::str::from_utf8(s).map_err(|_| { + exec_datafusion_err!("concat: scalar bytes are not valid UTF-8") + })?; self.block.extend_from_slice(s); - self.tainted = true; } ColumnarValueRef::NullableArray(array) => { if !CHECK_VALID || array.is_valid(i) { @@ -214,12 +231,6 @@ impl ConcatStringViewBuilder { self.block.extend_from_slice(array.value(i).as_bytes()); } } - ColumnarValueRef::NullableBinaryArray(array) => { - if !CHECK_VALID || array.is_valid(i) { - self.block.extend_from_slice(array.value(i)); - } - self.tainted = true; - } ColumnarValueRef::NonNullableArray(array) => { self.block.extend_from_slice(array.value(i).as_bytes()); } @@ -229,21 +240,18 @@ impl ConcatStringViewBuilder { ColumnarValueRef::NonNullableStringViewArray(array) => { self.block.extend_from_slice(array.value(i).as_bytes()); } - ColumnarValueRef::NonNullableBinaryArray(array) => { - self.block.extend_from_slice(array.value(i)); - self.tainted = true; + _ => { + return exec_err!( + "concat: unexpected column type for string view builder: {column:?}" + ); } } + Ok(()) } /// Finalizes the current row by converting the accumulated data into a /// StringView and appending it to the views buffer. - pub fn append_offset(&mut self) -> Result<()> { - if self.tainted { - std::str::from_utf8(&self.block) - .map_err(|_| exec_datafusion_err!("invalid UTF-8 in binary literal"))?; - } - + fn append_offset(&mut self) -> Result<()> { let v = &self.block; if v.len() > 12 { let offset: u32 = self @@ -258,7 +266,6 @@ impl ConcatStringViewBuilder { } self.block.clear(); - self.tainted = false; Ok(()) } @@ -269,7 +276,7 @@ impl ConcatStringViewBuilder { /// Returns an error when: /// /// - the provided `null_buffer` length does not match the row count. - pub fn finish(self, null_buffer: Option) -> Result { + fn finish(self, null_buffer: Option) -> Result { if let Some(ref nulls) = null_buffer && nulls.len() != self.views.len() { @@ -287,8 +294,8 @@ impl ConcatStringViewBuilder { }; // SAFETY: views were constructed with correct lengths, offsets, and - // prefixes. UTF-8 validity was checked in append_offset() for any row - // where tainted data (e.g., binary literals) was appended. + // prefixes. All input fragments came from string arrays or string + // scalars, all of which are valid UTF-8. let array = unsafe { StringViewArray::new_unchecked( ScalarBuffer::from(self.views), @@ -296,135 +303,7 @@ impl ConcatStringViewBuilder { null_buffer, ) }; - Ok(array) - } -} - -/// Builder used by `concat`/`concat_ws` to assemble a [`LargeStringArray`] one -/// row at a time from multiple input columns. See [`ConcatStringBuilder`] for -/// details on the row-composition contract. -/// -/// For the common "produce one `&str` per row" pattern, prefer -/// `GenericStringArrayBuilder` instead. -pub(crate) struct ConcatLargeStringBuilder { - offsets_buffer: MutableBuffer, - value_buffer: MutableBuffer, - /// If true, a safety check is required during the `finish` call - tainted: bool, -} - -impl ConcatLargeStringBuilder { - pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self { - let capacity = item_capacity - .checked_add(1) - .map(|i| i.saturating_mul(size_of::())) - .expect("capacity integer overflow"); - - let mut offsets_buffer = MutableBuffer::with_capacity(capacity); - // SAFETY: the first offset value is definitely not going to exceed the bounds. - unsafe { offsets_buffer.push_unchecked(0_i64) }; - Self { - offsets_buffer, - value_buffer: MutableBuffer::with_capacity(data_capacity), - tainted: false, - } - } - - pub fn write( - &mut self, - column: &ColumnarValueRef, - i: usize, - ) { - match column { - ColumnarValueRef::Scalar(s) => { - self.value_buffer.extend_from_slice(s); - self.tainted = true; - } - ColumnarValueRef::NullableArray(array) => { - if !CHECK_VALID || array.is_valid(i) { - self.value_buffer - .extend_from_slice(array.value(i).as_bytes()); - } - } - ColumnarValueRef::NullableLargeStringArray(array) => { - if !CHECK_VALID || array.is_valid(i) { - self.value_buffer - .extend_from_slice(array.value(i).as_bytes()); - } - } - ColumnarValueRef::NullableStringViewArray(array) => { - if !CHECK_VALID || array.is_valid(i) { - self.value_buffer - .extend_from_slice(array.value(i).as_bytes()); - } - } - ColumnarValueRef::NullableBinaryArray(array) => { - if !CHECK_VALID || array.is_valid(i) { - self.value_buffer.extend_from_slice(array.value(i)); - } - self.tainted = true; - } - ColumnarValueRef::NonNullableArray(array) => { - self.value_buffer - .extend_from_slice(array.value(i).as_bytes()); - } - ColumnarValueRef::NonNullableLargeStringArray(array) => { - self.value_buffer - .extend_from_slice(array.value(i).as_bytes()); - } - ColumnarValueRef::NonNullableStringViewArray(array) => { - self.value_buffer - .extend_from_slice(array.value(i).as_bytes()); - } - ColumnarValueRef::NonNullableBinaryArray(array) => { - self.value_buffer.extend_from_slice(array.value(i)); - self.tainted = true; - } - } - } - - pub fn append_offset(&mut self) -> Result<()> { - let next_offset: i64 = self - .value_buffer - .len() - .try_into() - .map_err(|_| exec_datafusion_err!("byte array offset overflow"))?; - self.offsets_buffer.push(next_offset); - Ok(()) - } - - /// Finalize the builder into a concrete [`LargeStringArray`]. - /// - /// # Errors - /// - /// Returns an error when: - /// - /// - the provided `null_buffer` is not the same length as the `offsets_buffer`. - pub fn finish(self, null_buffer: Option) -> Result { - let row_count = self.offsets_buffer.len() / size_of::() - 1; - if let Some(ref null_buffer) = null_buffer - && null_buffer.len() != row_count - { - return internal_err!( - "Null buffer and offsets buffer must be the same length" - ); - } - let array_builder = ArrayDataBuilder::new(DataType::LargeUtf8) - .len(row_count) - .add_buffer(self.offsets_buffer.into()) - .add_buffer(self.value_buffer.into()) - .nulls(null_buffer); - if self.tainted { - // Raw binary arrays with possible invalid utf-8 were used, - // so let ArrayDataBuilder perform validation - let array_data = array_builder.build()?; - Ok(LargeStringArray::from(array_data)) - } else { - // SAFETY: all data that was appended was valid Large UTF8 and the values - // and offsets were created correctly - let array_data = unsafe { array_builder.build_unchecked() }; - Ok(LargeStringArray::from(array_data)) - } + Ok(Arc::new(array)) } } @@ -453,6 +332,24 @@ pub(crate) struct GenericStringArrayBuilder { _phantom: PhantomData, } +fn offset_overflow_error() -> DataFusionError { + exec_datafusion_err!( + "byte array offset overflow: output size exceeds {} bytes", + O::MAX_OFFSET + ) +} + +fn string_view_overflow_error(field: &str) -> DataFusionError { + exec_datafusion_err!("byte array offset overflow: {field} exceeds i32::MAX") +} + +fn try_offset(len: usize) -> Result { + if len > O::MAX_OFFSET { + return Err(offset_overflow_error::()); + } + Ok(O::usize_as(len)) +} + impl GenericStringArrayBuilder { pub fn with_capacity(item_capacity: usize, data_capacity: usize) -> Self { let capacity = item_capacity @@ -470,30 +367,107 @@ impl GenericStringArrayBuilder { } } + #[inline] + fn try_push_offset_for_len(&mut self, len: usize) -> Result<()> { + let next_offset = try_offset::(len)?; + self.offsets_buffer.push(next_offset); + Ok(()) + } + + #[inline] + fn try_append_bytes(&mut self, additional_len: usize, append: F) -> Result<()> + where + F: FnOnce(&mut MutableBuffer), + { + let next_len = self + .value_buffer + .len() + .checked_add(additional_len) + .ok_or_else(offset_overflow_error::)?; + let next_offset = try_offset::(next_len)?; + append(&mut self.value_buffer); + debug_assert_eq!(self.value_buffer.len(), next_len); + self.offsets_buffer.push(next_offset); + Ok(()) + } + + /// Fallible variant of [`Self::append_value`]. + /// + /// # Errors + /// + /// Returns an error if the cumulative byte length exceeds this builder's + /// offset type limit. + #[inline] + pub fn try_append_value(&mut self, value: &str) -> Result<()> { + self.try_append_bytes(value.len(), |value_buffer| { + value_buffer.extend_from_slice(value.as_bytes()); + }) + } + + /// Fallible variant of [`Self::append_placeholder`]. + /// + /// # Errors + /// + /// Returns an error if the current cumulative byte length exceeds this + /// builder's offset type limit. + #[inline] + pub fn try_append_placeholder(&mut self) -> Result<()> { + self.try_push_offset_for_len(self.value_buffer.len())?; + self.placeholder_count += 1; + Ok(()) + } + /// See [`BulkNullStringArrayBuilder::append_value`]. /// + /// Note: new call sites that need recoverable overflow handling should + /// prefer [`Self::try_append_value`]. + /// /// # Panics /// /// Panics if the cumulative byte length exceeds `O::MAX`. #[inline] pub fn append_value(&mut self, value: &str) { - self.value_buffer.extend_from_slice(value.as_bytes()); - let next_offset = - O::from_usize(self.value_buffer.len()).expect("byte array offset overflow"); - self.offsets_buffer.push(next_offset); + self.try_append_value(value) + .expect("byte array offset overflow"); } /// See [`BulkNullStringArrayBuilder::append_placeholder`]. + /// + /// Note: new call sites that need recoverable overflow handling should + /// prefer [`Self::try_append_placeholder`]. #[inline] pub fn append_placeholder(&mut self) { - let next_offset = - O::from_usize(self.value_buffer.len()).expect("byte array offset overflow"); - self.offsets_buffer.push(next_offset); - self.placeholder_count += 1; + self.try_append_placeholder() + .expect("byte array offset overflow"); + } + + /// Fallible variant of [`Self::append_byte_map`]. + /// + /// # Safety + /// + /// The bytes produced by applying `map` to each byte of `src`, in order, + /// must form valid UTF-8. + /// + /// # Errors + /// + /// Returns an error if the cumulative byte length exceeds this builder's + /// offset type limit. + #[inline] + pub unsafe fn try_append_byte_map u8>( + &mut self, + src: &[u8], + mut map: F, + ) -> Result<()> { + self.try_append_bytes(src.len(), |value_buffer| { + value_buffer.extend(src.iter().map(|&b| map(b))); + }) } /// See [`BulkNullStringArrayBuilder::append_byte_map`]. /// + /// Note: new call sites that need recoverable overflow handling should + /// prefer [`Self::try_append_byte_map`]. + /// /// # Safety /// /// The bytes produced by applying `map` to each byte of `src`, in order, @@ -503,15 +477,46 @@ impl GenericStringArrayBuilder { /// /// Panics if the cumulative byte length exceeds `O::MAX`. #[inline] - pub unsafe fn append_byte_map u8>(&mut self, src: &[u8], mut map: F) { - self.value_buffer.extend(src.iter().map(|&b| map(b))); - let next_offset = - O::from_usize(self.value_buffer.len()).expect("byte array offset overflow"); + pub unsafe fn append_byte_map u8>(&mut self, src: &[u8], map: F) { + // SAFETY: caller upholds this method's UTF-8 contract. + unsafe { self.try_append_byte_map(src, map) } + .expect("byte array offset overflow"); + } + + /// Fallible variant of [`Self::append_with`]. + /// + /// # Errors + /// + /// Returns an error if the cumulative byte length exceeds this builder's + /// offset type limit. + #[inline] + pub fn try_append_with(&mut self, f: F) -> Result<()> + where + F: FnOnce(&mut GenericStringWriter<'_>), + { + let old_len = self.value_buffer.len(); + let mut writer = GenericStringWriter { + value_buffer: &mut self.value_buffer, + }; + f(&mut writer); + let next_offset = match try_offset::(self.value_buffer.len()) { + Ok(offset) => offset, + Err(e) => { + // SAFETY: `old_len` was the initialized length before `f` wrote to + // this owned buffer, so shrinking back preserves initialized data. + unsafe { self.value_buffer.set_len(old_len) }; + return Err(e); + } + }; self.offsets_buffer.push(next_offset); + Ok(()) } /// See [`BulkNullStringArrayBuilder::append_with`]. /// + /// Note: new call sites that need recoverable overflow handling should + /// prefer [`Self::try_append_with`]. + /// /// # Panics /// /// Panics if the cumulative byte length exceeds `O::MAX`. @@ -520,6 +525,9 @@ impl GenericStringArrayBuilder { where F: FnOnce(&mut GenericStringWriter<'_>), { + // Do not delegate to `try_append_with`: it rolls back value_buffer on + // overflow before returning Err, which would change this infallible + // method's state if its panic is caught. let mut writer = GenericStringWriter { value_buffer: &mut self.value_buffer, }; @@ -680,39 +688,75 @@ impl StringViewArrayBuilder { self.block_size } - /// See [`BulkNullStringArrayBuilder::append_value`]. + /// Fallible variant of [`Self::append_value`]. /// - /// # Panics + /// # Errors /// - /// Panics if the value length, the in-progress buffer offset, or the - /// number of completed buffers exceeds `i32::MAX`. The ByteView spec - /// uses signed 32-bit integers for these fields; exceeding `i32::MAX` - /// would produce an array that does not round-trip through Arrow IPC - /// (see ). + /// Returns an error if the value length, in-progress buffer offset, or + /// number of completed buffers exceeds `i32::MAX`. The ByteView spec uses + /// signed 32-bit integers for these fields; exceeding `i32::MAX` would + /// produce an array that does not round-trip through Arrow IPC (see + /// ). #[inline] - pub fn append_value(&mut self, value: &str) { + pub fn try_append_value(&mut self, value: &str) -> Result<()> { let v = value.as_bytes(); - let length: u32 = - i32::try_from(v.len()).expect("value length exceeds i32::MAX") as u32; + let length: u32 = i32::try_from(v.len()) + .map_err(|_| string_view_overflow_error("value length"))? + as u32; if length <= 12 { self.views.push(make_view(v, 0, 0)); - return; + return Ok(()); } - let required_cap = self.in_progress.len() + length as usize; - if self.in_progress.capacity() < required_cap { - self.flush_in_progress(); - let to_reserve = (length as usize).max(self.next_block_size() as usize); - self.in_progress.reserve(to_reserve); - } + self.try_ensure_long_capacity(length)?; let offset: u32 = i32::try_from(self.in_progress.len()) - .expect("offset exceeds i32::MAX") as u32; + .map_err(|_| string_view_overflow_error("offset"))? + as u32; + let buffer_index: u32 = i32::try_from(self.completed.len()) + .map_err(|_| string_view_overflow_error("buffer count"))? + as u32; self.in_progress.extend_from_slice(v); - self.views.push(self.make_long_view(length, offset, v)); + self.views.push(Self::make_long_view_checked( + length, + buffer_index, + offset, + v, + )); + Ok(()) + } + + /// See [`BulkNullStringArrayBuilder::append_value`]. + /// + /// Note: new call sites that need recoverable overflow handling should + /// prefer [`Self::try_append_value`]. + /// + /// # Panics + /// + /// Panics under the same conditions that [`Self::try_append_value`] returns + /// an error. + #[inline] + pub fn append_value(&mut self, value: &str) { + self.try_append_value(value) + .expect("byte array offset overflow"); + } + + /// Fallible variant of [`Self::append_placeholder`]. + /// + /// # Errors + /// + /// This currently cannot fail; it returns `Result` for API symmetry with + /// other fallible append methods. + #[inline] + pub fn try_append_placeholder(&mut self) -> Result<()> { + self.append_placeholder(); + Ok(()) } /// See [`BulkNullStringArrayBuilder::append_placeholder`]. + /// + /// Note: new call sites that need recoverable overflow handling should + /// prefer [`Self::try_append_placeholder`]. #[inline] pub fn append_placeholder(&mut self) { // Zero-length inline view — `length` field is 0, no buffer ref. @@ -720,18 +764,34 @@ impl StringViewArrayBuilder { self.placeholder_count += 1; } - /// Ensure the in-progress block has room for `length` more bytes, - /// flushing the current block and starting a new (doubled) one if not. - /// Caller must invoke this only when no bytes of the current row are - /// yet in `in_progress` — flushing mid-row would orphan partial data. + /// Fallible variant of [`Self::ensure_long_capacity`]. #[inline] - fn ensure_long_capacity(&mut self, length: u32) { - let required_cap = self.in_progress.len() + length as usize; + fn try_ensure_long_capacity(&mut self, length: u32) -> Result<()> { + let required_cap = self + .in_progress + .len() + .checked_add(length as usize) + .ok_or_else(|| string_view_overflow_error("string view block size"))?; if self.in_progress.capacity() < required_cap { self.flush_in_progress(); let to_reserve = (length as usize).max(self.next_block_size() as usize); + #[expect( + clippy::disallowed_methods, + reason = "StringView's block size bounds growth, so reserve cannot overflow capacity arithmetically. This hot loop intentionally avoids the extra `try_reserve` checks. It remains subject to allocator failure/OOM, which must be managed externally." + )] self.in_progress.reserve(to_reserve); } + Ok(()) + } + + /// Ensure the in-progress block has room for `length` more bytes, + /// flushing the current block and starting a new (doubled) one if not. + /// Caller must invoke this only when no bytes of the current row are + /// yet in `in_progress` — flushing mid-row would orphan partial data. + #[inline] + fn ensure_long_capacity(&mut self, length: u32) { + self.try_ensure_long_capacity(length) + .expect("byte array offset overflow"); } /// Encode a long-form view referencing `length` bytes already written @@ -742,10 +802,12 @@ impl StringViewArrayBuilder { /// function is `[inline(never)]` and has to handle short strings, so /// building the view here ourselves is faster. #[inline] - fn make_long_view(&self, length: u32, offset: u32, prefix_bytes: &[u8]) -> u128 { - let buffer_index: u32 = i32::try_from(self.completed.len()) - .expect("buffer count exceeds i32::MAX") - as u32; + fn make_long_view_checked( + length: u32, + buffer_index: u32, + offset: u32, + prefix_bytes: &[u8], + ) -> u128 { ByteView { length, // length > 12, so prefix_bytes has at least 4 bytes. @@ -756,6 +818,14 @@ impl StringViewArrayBuilder { .into() } + #[inline] + fn make_long_view(&self, length: u32, offset: u32, prefix_bytes: &[u8]) -> u128 { + let buffer_index: u32 = i32::try_from(self.completed.len()) + .expect("buffer count exceeds i32::MAX") + as u32; + Self::make_long_view_checked(length, buffer_index, offset, prefix_bytes) + } + /// See [`BulkNullStringArrayBuilder::append_byte_map`]. /// /// # Safety @@ -1157,6 +1227,10 @@ pub(crate) enum ColumnarValueRef<'a> { NonNullableStringViewArray(&'a StringViewArray), NullableBinaryArray(&'a BinaryArray), NonNullableBinaryArray(&'a BinaryArray), + NullableLargeBinaryArray(&'a LargeBinaryArray), + NonNullableLargeBinaryArray(&'a LargeBinaryArray), + NullableBinaryViewArray(&'a BinaryViewArray), + NonNullableBinaryViewArray(&'a BinaryViewArray), } impl ColumnarValueRef<'_> { @@ -1167,11 +1241,15 @@ impl ColumnarValueRef<'_> { | Self::NonNullableArray(_) | Self::NonNullableLargeStringArray(_) | Self::NonNullableStringViewArray(_) - | Self::NonNullableBinaryArray(_) => true, + | Self::NonNullableBinaryArray(_) + | Self::NonNullableLargeBinaryArray(_) + | Self::NonNullableBinaryViewArray(_) => true, Self::NullableArray(array) => array.is_valid(i), Self::NullableStringViewArray(array) => array.is_valid(i), Self::NullableLargeStringArray(array) => array.is_valid(i), Self::NullableBinaryArray(array) => array.is_valid(i), + Self::NullableLargeBinaryArray(array) => array.is_valid(i), + Self::NullableBinaryViewArray(array) => array.is_valid(i), } } @@ -1182,13 +1260,168 @@ impl ColumnarValueRef<'_> { | Self::NonNullableArray(_) | Self::NonNullableStringViewArray(_) | Self::NonNullableLargeStringArray(_) - | Self::NonNullableBinaryArray(_) => None, + | Self::NonNullableBinaryArray(_) + | Self::NonNullableLargeBinaryArray(_) + | Self::NonNullableBinaryViewArray(_) => None, Self::NullableArray(array) => array.nulls().cloned(), Self::NullableStringViewArray(array) => array.nulls().cloned(), Self::NullableLargeStringArray(array) => array.nulls().cloned(), Self::NullableBinaryArray(array) => array.nulls().cloned(), + Self::NullableLargeBinaryArray(array) => array.nulls().cloned(), + Self::NullableBinaryViewArray(array) => array.nulls().cloned(), } } + + /// Parse a [`ColumnarValue`] argument into `ColumnarValueRef`. + /// Returns `None` when the argument is null or null scalar + /// Returns an error when a columnar value type is not supported. + /// Shared by `concat` and `concat_ws`. + pub(crate) fn from_columnar_value<'a>( + col: &'a ColumnarValue, + data_size: &mut usize, + len: usize, + size_factor: usize, + convert_to_str: bool, + ) -> Result>> { + match col { + ColumnarValue::Scalar(ScalarValue::Utf8(maybe_value)) + | ColumnarValue::Scalar(ScalarValue::LargeUtf8(maybe_value)) + | ColumnarValue::Scalar(ScalarValue::Utf8View(maybe_value)) => { + if let Some(s) = maybe_value { + *data_size += s.len() * len * size_factor; + Ok(Some(ColumnarValueRef::Scalar(s.as_bytes()))) + } else { + Ok(None) + } + } + ColumnarValue::Scalar(ScalarValue::Binary(maybe_value)) + | ColumnarValue::Scalar(ScalarValue::LargeBinary(maybe_value)) + | ColumnarValue::Scalar(ScalarValue::BinaryView(maybe_value)) + | ColumnarValue::Scalar(ScalarValue::FixedSizeBinary(_, maybe_value)) => { + if let Some(b) = maybe_value { + *data_size += b.len() * len * size_factor; + Ok(Some(ColumnarValueRef::Scalar(b.as_slice()))) + } else { + Ok(None) + } + } + ColumnarValue::Scalar(scalar) if scalar.is_null() => { + // null scalar is skipped + Ok(None) + } + ColumnarValue::Scalar(scalar) if convert_to_str => { + match scalar.try_as_str() { + Some(Some(s)) => { + *data_size += s.len() * len * size_factor; + Ok(Some(ColumnarValueRef::Scalar(s.as_bytes()))) + } + Some(None) => unreachable!("null handled above"), + None => { + internal_err!("Expected string or binary, got {scalar:?}") + } + } + } + ColumnarValue::Array(array) => match array.data_type() { + DataType::Utf8 => { + let string_array = as_string_array(array)?; + *data_size += string_array.values().len() * size_factor; + let column = if array.is_nullable() { + ColumnarValueRef::NullableArray(string_array) + } else { + ColumnarValueRef::NonNullableArray(string_array) + }; + Ok(Some(column)) + } + DataType::LargeUtf8 => { + let string_array = as_largestring_array(array); + *data_size += string_array.values().len() * size_factor; + let column = if array.is_nullable() { + ColumnarValueRef::NullableLargeStringArray(string_array) + } else { + ColumnarValueRef::NonNullableLargeStringArray(string_array) + }; + Ok(Some(column)) + } + DataType::Utf8View => { + let string_array = as_string_view_array(array)?; + *data_size += string_array.total_buffer_bytes_used() * size_factor; + let column = if array.is_nullable() { + ColumnarValueRef::NullableStringViewArray(string_array) + } else { + ColumnarValueRef::NonNullableStringViewArray(string_array) + }; + Ok(Some(column)) + } + DataType::Binary => { + let binary_array = as_binary_array(array)?; + *data_size += binary_array.values().len() * size_factor; + let column = if array.is_nullable() { + ColumnarValueRef::NullableBinaryArray(binary_array) + } else { + ColumnarValueRef::NonNullableBinaryArray(binary_array) + }; + Ok(Some(column)) + } + DataType::LargeBinary => { + let binary_array = as_large_binary_array(array)?; + *data_size += binary_array.values().len() * size_factor; + let column = if array.is_nullable() { + ColumnarValueRef::NullableLargeBinaryArray(binary_array) + } else { + ColumnarValueRef::NonNullableLargeBinaryArray(binary_array) + }; + Ok(Some(column)) + } + DataType::BinaryView => { + let binary_array = as_binary_view_array(array)?; + *data_size += binary_array.total_buffer_bytes_used() * size_factor; + let column = if array.is_nullable() { + ColumnarValueRef::NullableBinaryViewArray(binary_array) + } else { + ColumnarValueRef::NonNullableBinaryViewArray(binary_array) + }; + Ok(Some(column)) + } + other => { + plan_err!( + "Input was {other} which is not a supported datatype for concat function" + ) + } + }, + _ => { + plan_err!( + "Input was {col} which is not a supported datatype for concat function" + ) + } + } + } +} + +/// Return the widest binary type found in `types`. +/// Order: `BinaryView` > `LargeBinary` / `FixedSizeBinary` > `Binary`. +pub(crate) fn widest_binary_type(types: &[DataType]) -> DataType { + if types.iter().any(|t| matches!(t, DataType::BinaryView)) { + DataType::BinaryView + } else if types + .iter() + .any(|t| matches!(t, DataType::LargeBinary | DataType::FixedSizeBinary(_))) + { + DataType::LargeBinary + } else { + DataType::Binary + } +} + +/// Return the widest string type found in `types`. +/// Order: `Utf8View` > `LargeUtf8` > `Utf8`. +pub(crate) fn widest_string_type(types: &[DataType]) -> DataType { + if types.iter().any(|t| matches!(t, DataType::Utf8View)) { + DataType::Utf8View + } else if types.iter().any(|t| matches!(t, DataType::LargeUtf8)) { + DataType::LargeUtf8 + } else { + DataType::Utf8 + } } #[cfg(test)] @@ -1380,6 +1613,90 @@ mod tests { assert_finish_errs_on_length_mismatch(StringViewArrayBuilder::with_capacity(2)); } + #[test] + fn generic_string_builder_try_append_success_path() { + let mut builder = GenericStringArrayBuilder::::with_capacity(4, 16); + builder.try_append_value("abc").unwrap(); + builder.try_append_placeholder().unwrap(); + // SAFETY: ASCII input and output. + unsafe { + builder + .try_append_byte_map(b"de", |b| b.to_ascii_uppercase()) + .unwrap(); + } + builder + .try_append_with(|w| { + w.write_str("f"); + w.write_char('é'); + }) + .unwrap(); + + let nulls = Some(NullBuffer::from(vec![true, false, true, true])); + let array = builder.finish(nulls).unwrap(); + assert_eq!( + &array, + &StringArray::from(vec![Some("abc"), None, Some("DE"), Some("fé")]) + ); + } + + #[test] + fn generic_string_builder_mixed_append_success_path() { + let mut builder = GenericStringArrayBuilder::::with_capacity(4, 16); + builder.append_value("ab"); + builder.try_append_value("cd").unwrap(); + // SAFETY: ASCII input and output. + unsafe { + builder.append_byte_map(b"ef", |b| b.to_ascii_uppercase()); + builder + .try_append_byte_map(b"gh", |b| b.to_ascii_uppercase()) + .unwrap(); + } + + let array = builder.finish(None).unwrap(); + assert_eq!( + &array, + &StringArray::from(vec![Some("ab"), Some("cd"), Some("EF"), Some("GH")]) + ); + } + + #[test] + fn string_view_builder_try_append_success_path() { + let mut builder = StringViewArrayBuilder::with_capacity(3); + builder.try_append_value("abc").unwrap(); + builder.try_append_placeholder().unwrap(); + builder.try_append_value("a long string value").unwrap(); + + let nulls = Some(NullBuffer::from(vec![true, false, true])); + let array = builder.finish(nulls).unwrap(); + assert_eq!(array.value(0), "abc"); + assert!(array.is_null(1)); + assert_eq!(array.value(2), "a long string value"); + } + + #[test] + fn generic_string_builder_try_offset_overflow() { + let err = try_offset::(i32::MAX as usize + 1) + .unwrap_err() + .to_string(); + assert!( + err.contains("byte array offset overflow"), + "unexpected error: {err}" + ); + } + + #[test] + fn generic_string_builder_try_append_bytes_overflow() { + let mut builder = GenericStringArrayBuilder::::with_capacity(0, 0); + let err = builder + .try_append_bytes(i32::MAX as usize + 1, |_| unreachable!()) + .unwrap_err() + .to_string(); + assert!( + err.contains("byte array offset overflow"), + "unexpected error: {err}" + ); + } + #[test] #[cfg(debug_assertions)] #[should_panic(expected = "placeholder rows")] diff --git a/datafusion/functions/src/unicode/character_length.rs b/datafusion/functions/src/unicode/character_length.rs index 465b15ace1d10..9f0d952a02636 100644 --- a/datafusion/functions/src/unicode/character_length.rs +++ b/datafusion/functions/src/unicode/character_length.rs @@ -15,16 +15,19 @@ // specific language governing permissions and limitations // under the License. -use crate::utils::{make_scalar_function, utf8_to_int_type}; +use crate::utils::{ + make_scalar_function, transform_leaf_type_preserving_encoding, utf8_to_int_type, +}; use arrow::array::{ Array, ArrayRef, ArrowPrimitiveType, AsArray, OffsetSizeTrait, PrimitiveArray, StringArrayType, }; use arrow::datatypes::{ArrowNativeType, DataType, Int32Type, Int64Type}; use datafusion_common::Result; +use datafusion_common::types::{NativeType, logical_string}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; use std::sync::Arc; @@ -59,11 +62,16 @@ impl Default for CharacterLengthFunc { impl CharacterLengthFunc { pub fn new() -> Self { - use DataType::*; Self { - signature: Signature::uniform( - 1, - vec![Utf8, LargeUtf8, Utf8View], + signature: Signature::coercible( + vec![ + Coercion::new_implicit( + TypeSignatureClass::Native(logical_string()), + vec![TypeSignatureClass::Any], + NativeType::String, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), aliases: vec![String::from("length"), String::from("char_length")], @@ -81,7 +89,9 @@ impl ScalarUDFImpl for CharacterLengthFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - utf8_to_int_type(&arg_types[0], "character_length") + transform_leaf_type_preserving_encoding(&arg_types[0], &|data_type| { + utf8_to_int_type(data_type, "character_length") + }) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -114,6 +124,11 @@ fn character_length(args: &[ArrayRef]) -> Result { let string_array = args[0].as_string_view(); character_length_general::(&string_array) } + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let converted = character_length(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(converted)) + } _ => unreachable!("CharacterLengthFunc"), } } diff --git a/datafusion/functions/src/unicode/common.rs b/datafusion/functions/src/unicode/common.rs index 092f2b8003b1b..5dc8f334da8a3 100644 --- a/datafusion/functions/src/unicode/common.rs +++ b/datafusion/functions/src/unicode/common.rs @@ -50,6 +50,16 @@ pub(crate) fn try_as_scalar_i64(cv: &ColumnarValue) -> Option { } } +/// Estimates data capacity for `pad` based on `length_array` with row length. +/// For ASCII, one row is at most `target_len` bytes. +/// For UTF8, it could be larger +pub(crate) fn pad_data_capacity(length_array: &Int64Array) -> usize { + length_array + .iter() + .flatten() + .fold(0, |acc, len| acc.saturating_add(len as usize)) +} + /// A trait for `left` and `right` byte slicing operations pub(crate) trait LeftRightSlicer { fn slice(string: &str, n: i64) -> Range; @@ -115,16 +125,22 @@ pub(crate) enum StringCharLen { /// Calculate the byte length of the substring of `n` chars from string `string` #[inline] fn left_right_byte_length(string: &str, n: i64) -> usize { + let abs = n.unsigned_abs().min(usize::MAX as u64) as usize; + // For ASCII input every character is exactly one byte, so the byte offset of + // the n-th codepoint is just the (clamped) character count. This avoids the + // per-character `char_indices()` scan of the general path. match n.cmp(&0) { + Ordering::Equal => 0, + // `abs` chars trimmed from the end: keep the leading `len - abs`. + Ordering::Less if string.is_ascii() => string.len().saturating_sub(abs), Ordering::Less => string .char_indices() - .nth_back((n.unsigned_abs().min(usize::MAX as u64) - 1) as usize) + .nth_back(abs - 1) .map(|(index, _)| index) .unwrap_or(0), - Ordering::Equal => 0, - Ordering::Greater => { - byte_offset_of_char(string, n.unsigned_abs().min(usize::MAX as u64) as usize) - } + // First `abs` chars, but never past the end of the string. + Ordering::Greater if string.is_ascii() => abs.min(string.len()), + Ordering::Greater => byte_offset_of_char(string, abs), } } @@ -151,79 +167,20 @@ pub(crate) fn general_left_right( } } -/// Returns true if all offsets in the array fit in i32, meaning the values -/// buffer can be referenced by StringView's offset field. -fn values_fit_in_i32(string_array: &GenericStringArray) -> bool { - string_array - .offsets() - .last() - .map(|offset| offset.as_usize() <= i32::MAX as usize) - .unwrap_or(true) -} - /// `left`/`right` for Utf8/LargeUtf8 input. -/// -/// When offsets fit in i32, produces a zero-copy `StringViewArray` with views -/// pointing into the input values buffer. Otherwise falls back to building a -/// `StringViewArray` by copying. fn general_left_right_array( string_array: &GenericStringArray, n_array: &Int64Array, ) -> Result { - if !values_fit_in_i32(string_array) { - let result = string_array - .iter() - .zip(n_array.iter()) - .map(|(string, n)| match (string, n) { - (Some(string), Some(n)) => Some(&string[F::slice(string, n)]), - _ => None, - }) - .collect::(); - return Ok(Arc::new(result) as ArrayRef); - } - - let len = string_array.len(); - let offsets = string_array.value_offsets(); - let nulls = NullBuffer::union(string_array.nulls(), n_array.nulls()); - - let mut views_buf = Vec::with_capacity(len); - let mut has_out_of_line = false; - - for (i, offset) in offsets.iter().enumerate().take(len) { - if nulls.as_ref().is_some_and(|n| n.is_null(i)) { - views_buf.push(0); - continue; - } - - // SAFETY: we just checked validity above - let string = unsafe { string_array.value_unchecked(i) }; - let n = n_array.value(i); - let range = F::slice(string, n); - let result_bytes = &string.as_bytes()[range.clone()]; - if result_bytes.len() > 12 { - has_out_of_line = true; - } - - let buf_offset = offset.as_usize() as u32 + range.start as u32; - views_buf.push(make_view(result_bytes, 0, buf_offset)); - } - - let views = ScalarBuffer::from(views_buf); - let data_buffers = if has_out_of_line { - vec![string_array.values().clone()] - } else { - vec![] - }; - - // SAFETY: - // - Each view is produced by `make_view` with correct bytes and offset - // - Out-of-line views reference buffer index 0, which is the original - // values buffer included in data_buffers when has_out_of_line is true - // - values_fit_in_i32 guarantees all offsets fit in i32 - unsafe { - let array = StringViewArray::new_unchecked(views, data_buffers, nulls); - Ok(Arc::new(array) as ArrayRef) - } + let result = string_array + .iter() + .zip(n_array.iter()) + .map(|(string, n)| match (string, n) { + (Some(string), Some(n)) => Some(&string[F::slice(string, n)]), + _ => None, + }) + .collect::>(); + Ok(Arc::new(result) as ArrayRef) } /// `general_left_right` for StringViewArray input. diff --git a/datafusion/functions/src/unicode/find_in_set.rs b/datafusion/functions/src/unicode/find_in_set.rs index 0a83eb3ed61ef..fa23532406ce1 100644 --- a/datafusion/functions/src/unicode/find_in_set.rs +++ b/datafusion/functions/src/unicode/find_in_set.rs @@ -25,7 +25,7 @@ use arrow_buffer::NullBuffer; use crate::utils::utf8_to_int_type; use datafusion_common::{ - Result, ScalarValue, exec_err, internal_err, utils::take_function_args, + HashMap, Result, ScalarValue, exec_err, internal_err, utils::take_function_args, }; use datafusion_expr::TypeSignature::Exact; use datafusion_expr::{ @@ -316,6 +316,11 @@ where Ok(Arc::new(PrimitiveArray::::new(values.into(), nulls)) as ArrayRef) } +/// Minimum set length at which a pre-built lookup beats a per-row linear scan. +/// Below this, the linear scan's small constant factor wins, so short sets are +/// left untouched to avoid regressing them. +const FIND_IN_SET_LOOKUP_THRESHOLD: usize = 16; + fn find_in_set_right_literal<'a, T, V>( string_array: V, str_list: &[&str], @@ -329,16 +334,34 @@ where let nulls = string_array.nulls().cloned(); let zero = T::Native::from_usize(0).unwrap(); + // The set (`str_list`) is constant across all rows. For a large set, the + // per-row `position` linear scan is O(set_len). Building a lookup from each + // distinct entry to its 1-based position once turns each row into an O(1) + // probe (first occurrence wins, exactly matching `position`). Below the + // threshold the linear scan's small constant factor is faster, so the map is + // built at most once here rather than per row. + let map: Option> = + (str_list.len() >= FIND_IN_SET_LOOKUP_THRESHOLD).then(|| { + let mut map = HashMap::with_capacity(str_list.len()); + for (idx, entry) in str_list.iter().enumerate() { + map.entry(*entry).or_insert(idx + 1); + } + map + }); + let values: Vec = (0..len) .map(|i| { if nulls.as_ref().is_some_and(|n| n.is_null(i)) { return zero; } let string = string_array.value(i); - let position = str_list - .iter() - .position(|s| *s == string) - .map_or(0, |idx| idx + 1); + let position = match &map { + Some(map) => map.get(string).copied().unwrap_or(0), + None => str_list + .iter() + .position(|s| *s == string) + .map_or(0, |idx| idx + 1), + }; T::Native::from_usize(position).unwrap() }) .collect(); @@ -545,4 +568,46 @@ mod tests { ], Int32Array::from(vec![None::; 3]) ); + + // Exercises both the lookup-map path (list length >= threshold) and the + // linear-scan path (short list), including a duplicate entry to confirm the + // first occurrence wins in both. + #[test] + fn test_right_literal_lookup_matches_linear() { + use super::find_in_set_right_literal; + use arrow::datatypes::Int32Type; + + // 40 unique entries plus a duplicate of "item5" appended at index 40, so + // the length is well over FIND_IN_SET_LOOKUP_THRESHOLD. + let mut long_list: Vec = (0..40).map(|i| format!("item{i}")).collect(); + long_list.push("item5".to_string()); + let long_refs: Vec<&str> = long_list.iter().map(|s| s.as_str()).collect(); + let short_refs = ["a", "b", "c"]; + + let strings = StringArray::from(vec![ + Some("item0"), + Some("item39"), + Some("item5"), + Some("missing"), + None, + Some("b"), + ]); + + let long = + find_in_set_right_literal::(&strings, &long_refs).unwrap(); + let long = long.as_any().downcast_ref::().unwrap(); + assert_eq!(long.value(0), 1); + assert_eq!(long.value(1), 40); + assert_eq!(long.value(2), 6); // first occurrence of "item5" + assert_eq!(long.value(3), 0); + assert!(long.is_null(4)); + assert_eq!(long.value(5), 0); + + let short = + find_in_set_right_literal::(&strings, &short_refs).unwrap(); + let short = short.as_any().downcast_ref::().unwrap(); + assert_eq!(short.value(0), 0); + assert!(short.is_null(4)); + assert_eq!(short.value(5), 2); // "b" at position 2 + } } diff --git a/datafusion/functions/src/unicode/initcap.rs b/datafusion/functions/src/unicode/initcap.rs index 711b2c49b09f6..0332ab5d4427f 100644 --- a/datafusion/functions/src/unicode/initcap.rs +++ b/datafusion/functions/src/unicode/initcap.rs @@ -17,18 +17,17 @@ use std::sync::Arc; -use arrow::array::{Array, ArrayRef, GenericStringArray, OffsetSizeTrait}; -use arrow::buffer::{Buffer, OffsetBuffer}; +use arrow::array::{Array, ArrayRef, AsArray, GenericStringArray, OffsetSizeTrait}; +use arrow::buffer::Buffer; use arrow::datatypes::DataType; use crate::strings::{GenericStringArrayBuilder, StringViewArrayBuilder}; -use crate::utils::{make_scalar_function, utf8_to_str_type}; use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; use datafusion_common::types::logical_string; use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ - Coercion, ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -64,9 +63,10 @@ impl InitcapFunc { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Native( - logical_string(), - ))], + vec![ + Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -83,54 +83,16 @@ impl ScalarUDFImpl for InitcapFunc { } fn return_type(&self, arg_types: &[DataType]) -> Result { - if let DataType::Utf8View = arg_types[0] { - Ok(DataType::Utf8View) - } else { - utf8_to_str_type(&arg_types[0], "initcap") - } + Ok(arg_types[0].clone()) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let arg = &args.args[0]; - - // Scalar fast path - handle directly without array conversion - if let ColumnarValue::Scalar(scalar) = arg { - return match scalar { - ScalarValue::Utf8(None) - | ScalarValue::LargeUtf8(None) - | ScalarValue::Utf8View(None) => Ok(arg.clone()), - ScalarValue::Utf8(Some(s)) => { - let mut result = String::new(); - initcap_string(s, &mut result); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some(result)))) - } - ScalarValue::LargeUtf8(Some(s)) => { - let mut result = String::new(); - initcap_string(s, &mut result); - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some(result)))) - } - ScalarValue::Utf8View(Some(s)) => { - let mut result = String::new(); - initcap_string(s, &mut result); - Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some(result)))) - } - other => { - exec_err!( - "Unsupported data type {:?} for function `initcap`", - other.data_type() - ) - } - }; - } - - // Array path - let args = &args.args; - match args[0].data_type() { - DataType::Utf8 => make_scalar_function(initcap::, vec![])(args), - DataType::LargeUtf8 => make_scalar_function(initcap::, vec![])(args), - DataType::Utf8View => make_scalar_function(initcap_utf8view, vec![])(args), - other => { - exec_err!("Unsupported data type {other:?} for function `initcap`") + match &args.args[0] { + ColumnarValue::Scalar(scalar) => { + Ok(ColumnarValue::Scalar(initcap_scalar(scalar)?)) + } + ColumnarValue::Array(array) => { + Ok(ColumnarValue::Array(initcap_array(array)?)) } } } @@ -140,6 +102,55 @@ impl ScalarUDFImpl for InitcapFunc { } } +fn initcap_scalar(scalar: &ScalarValue) -> Result { + match scalar { + ScalarValue::Utf8(None) + | ScalarValue::LargeUtf8(None) + | ScalarValue::Utf8View(None) => Ok(scalar.clone()), + ScalarValue::Utf8(Some(s)) => { + let mut result = String::new(); + initcap_string(s, &mut result); + Ok(ScalarValue::Utf8(Some(result))) + } + ScalarValue::LargeUtf8(Some(s)) => { + let mut result = String::new(); + initcap_string(s, &mut result); + Ok(ScalarValue::LargeUtf8(Some(result))) + } + ScalarValue::Utf8View(Some(s)) => { + let mut result = String::new(); + initcap_string(s, &mut result); + Ok(ScalarValue::Utf8View(Some(result))) + } + ScalarValue::Dictionary(key_type, value) => Ok(ScalarValue::Dictionary( + key_type.clone(), + Box::new(initcap_scalar(value)?), + )), + other => { + exec_err!( + "Unsupported data type {:?} for function `initcap`", + other.data_type() + ) + } + } +} + +fn initcap_array(array: &ArrayRef) -> Result { + match array.data_type() { + DataType::Utf8 => initcap::(&[Arc::clone(array)]), + DataType::LargeUtf8 => initcap::(&[Arc::clone(array)]), + DataType::Utf8View => initcap_utf8view(&[Arc::clone(array)]), + DataType::Dictionary(_, _) => { + let dictionary = array.as_any_dictionary(); + let converted = initcap_array(dictionary.values())?; + Ok(dictionary.with_values(converted)) + } + other => { + exec_err!("Unsupported data type {other:?} for function `initcap`") + } + } +} + /// Converts the first letter of each word to uppercase and the rest to /// lowercase. Words are sequences of alphanumeric characters separated by /// non-alphanumeric characters. @@ -166,12 +177,12 @@ fn initcap(args: &[ArrayRef]) -> Result { if let Some(ref n) = nulls { for i in 0..len { if n.is_null(i) { - builder.append_placeholder(); + builder.try_append_placeholder()?; } else { // SAFETY: not null per check above. let s = unsafe { string_array.value_unchecked(i) }; initcap_string(s, &mut container); - builder.append_value(&container); + builder.try_append_value(&container)?; } } } else { @@ -179,7 +190,7 @@ fn initcap(args: &[ArrayRef]) -> Result { // SAFETY: no null buffer means every index is valid. let s = unsafe { string_array.value_unchecked(i) }; initcap_string(s, &mut container); - builder.append_value(&container); + builder.try_append_value(&container)?; } } @@ -217,17 +228,10 @@ fn initcap_ascii_array( } let values = Buffer::from_vec(out); - let out_offsets = if first_offset == 0 { - offsets.clone() - } else { - // For sliced arrays, we need to rebase the offsets to reflect that the - // output only contains the bytes in the visible slice. - let rebased_offsets = offsets - .iter() - .map(|offset| T::usize_as(offset.as_usize() - first_offset)) - .collect::>(); - OffsetBuffer::::new(rebased_offsets.into()) - }; + + // Rebase offsets for sliced arrays to reflect that the + // output only contains the bytes in the visible slice. + let out_offsets = offsets.clone().subtract(offsets[0]); // SAFETY: ASCII case conversion preserves byte length, so the original // string boundaries are preserved. `out_offsets` is either identical to diff --git a/datafusion/functions/src/unicode/left.rs b/datafusion/functions/src/unicode/left.rs index 423ab4d5dc54b..0788e69d92528 100644 --- a/datafusion/functions/src/unicode/left.rs +++ b/datafusion/functions/src/unicode/left.rs @@ -79,8 +79,8 @@ impl ScalarUDFImpl for LeftFunc { &self.signature } - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Utf8View) + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) } /// Returns first n characters in the string, or when n is negative, returns all but last |n| characters. @@ -108,8 +108,8 @@ impl ScalarUDFImpl for LeftFunc { #[cfg(test)] mod tests { - use arrow::array::{Array, StringViewArray}; - use arrow::datatypes::DataType::Utf8View; + use arrow::array::{Array, LargeStringArray, StringArray, StringViewArray}; + use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View}; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; @@ -127,8 +127,19 @@ mod tests { ], Ok(Some("ab")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray + ); + test_function!( + LeftFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some("abcde".to_string()))), + ColumnarValue::Scalar(ScalarValue::from(2i64)), + ], + Ok(Some("ab")), + &str, + LargeUtf8, + LargeStringArray ); test_function!( LeftFunc::new(), @@ -138,8 +149,8 @@ mod tests { ], Ok(Some("abcde")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( LeftFunc::new(), @@ -149,8 +160,8 @@ mod tests { ], Ok(Some("abc")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( LeftFunc::new(), @@ -160,8 +171,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( LeftFunc::new(), @@ -171,8 +182,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( LeftFunc::new(), @@ -182,8 +193,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( LeftFunc::new(), @@ -193,8 +204,8 @@ mod tests { ], Ok(None), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( LeftFunc::new(), @@ -204,8 +215,8 @@ mod tests { ], Ok(None), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( LeftFunc::new(), @@ -215,8 +226,8 @@ mod tests { ], Ok(Some("joséé")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( LeftFunc::new(), @@ -226,8 +237,8 @@ mod tests { ], Ok(Some("joséé")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); #[cfg(not(feature = "unicode_expressions"))] test_function!( @@ -240,8 +251,8 @@ mod tests { "function left requires compilation with feature flag: unicode_expressions." ), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); // StringView cases @@ -307,8 +318,8 @@ mod tests { ], Ok(Some(expected.as_str())), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); } diff --git a/datafusion/functions/src/unicode/lpad.rs b/datafusion/functions/src/unicode/lpad.rs index d27bc8633e730..40bffeecf422a 100644 --- a/datafusion/functions/src/unicode/lpad.rs +++ b/datafusion/functions/src/unicode/lpad.rs @@ -178,7 +178,8 @@ impl ScalarUDFImpl for LPadFunc { } use super::common::{ - StringCharLen, char_count_or_boundary, try_as_scalar_i64, try_as_scalar_str, + StringCharLen, char_count_or_boundary, pad_data_capacity, try_as_scalar_i64, + try_as_scalar_str, }; /// Optimized lpad for constant target_len and fill arguments. @@ -373,7 +374,10 @@ where T: OffsetSizeTrait, { let array = if let Some(fill_array) = fill_array { - let mut builder: GenericStringBuilder = GenericStringBuilder::new(); + let mut builder: GenericStringBuilder = GenericStringBuilder::with_capacity( + string_array.len(), + pad_data_capacity(length_array), + ); let mut fill_chars_buf = Vec::new(); for ((string, target_len), fill) in string_array @@ -449,7 +453,10 @@ where builder.finish() } else { - let mut builder: GenericStringBuilder = GenericStringBuilder::new(); + let mut builder: GenericStringBuilder = GenericStringBuilder::with_capacity( + string_array.len(), + pad_data_capacity(length_array), + ); for (string, target_len) in string_array.iter().zip(length_array.iter()) { if let (Some(string), Some(target_len)) = (string, target_len) { diff --git a/datafusion/functions/src/unicode/reverse.rs b/datafusion/functions/src/unicode/reverse.rs index 813dcb5f504dd..9dfc25fbdfe07 100644 --- a/datafusion/functions/src/unicode/reverse.rs +++ b/datafusion/functions/src/unicode/reverse.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::sync::Arc; + use crate::strings::{ BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringViewArrayBuilder, }; @@ -22,10 +24,11 @@ use crate::utils::make_scalar_function; use DataType::{LargeUtf8, Utf8, Utf8View}; use arrow::array::{Array, ArrayRef, AsArray, StringArrayType}; use arrow::datatypes::DataType; -use datafusion_common::{Result, exec_err}; +use datafusion_common::Result; +use datafusion_common::types::{NativeType, logical_string}; use datafusion_expr::{ - ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, - Volatility, + Coercion, ColumnarValue, Documentation, EncodingPreservation, ScalarFunctionArgs, + ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, }; use datafusion_macros::user_doc; @@ -56,11 +59,16 @@ impl Default for ReverseFunc { impl ReverseFunc { pub fn new() -> Self { - use DataType::*; Self { - signature: Signature::uniform( - 1, - vec![Utf8View, Utf8, LargeUtf8], + signature: Signature::coercible( + vec![ + Coercion::new_implicit( + TypeSignatureClass::Native(logical_string()), + vec![TypeSignatureClass::Any], + NativeType::String, + ) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } @@ -81,13 +89,7 @@ impl ScalarUDFImpl for ReverseFunc { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - let args = &args.args; - match args[0].data_type() { - Utf8 | Utf8View | LargeUtf8 => make_scalar_function(reverse, vec![])(args), - other => { - exec_err!("Unsupported data type {other:?} for function reverse") - } - } + make_scalar_function(reverse, vec![])(&args.args) } fn documentation(&self) -> Option<&Documentation> { @@ -113,6 +115,11 @@ fn reverse(args: &[ArrayRef]) -> Result { &args[0].as_string_view(), StringViewArrayBuilder::with_capacity(len), ), + DataType::Dictionary(_, _) => { + let dictionary = args[0].as_any_dictionary(); + let converted = reverse(&[Arc::clone(dictionary.values())])?; + Ok(dictionary.with_values(converted)) + } _ => unreachable!( "Reverse can only be applied to Utf8View, Utf8 and LargeUtf8 types" ), diff --git a/datafusion/functions/src/unicode/right.rs b/datafusion/functions/src/unicode/right.rs index 0ed170fef72d7..21fb0690a11a2 100644 --- a/datafusion/functions/src/unicode/right.rs +++ b/datafusion/functions/src/unicode/right.rs @@ -79,8 +79,8 @@ impl ScalarUDFImpl for RightFunc { &self.signature } - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Utf8View) + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) } /// Returns right n characters in the string, or when n is negative, returns all but first |n| characters. @@ -108,8 +108,8 @@ impl ScalarUDFImpl for RightFunc { #[cfg(test)] mod tests { - use arrow::array::{Array, StringViewArray}; - use arrow::datatypes::DataType::Utf8View; + use arrow::array::{Array, LargeStringArray, StringArray, StringViewArray}; + use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View}; use datafusion_common::{Result, ScalarValue}; use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; @@ -127,8 +127,19 @@ mod tests { ], Ok(Some("de")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray + ); + test_function!( + RightFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some("abcde".to_string()))), + ColumnarValue::Scalar(ScalarValue::from(2i64)), + ], + Ok(Some("de")), + &str, + LargeUtf8, + LargeStringArray ); test_function!( RightFunc::new(), @@ -138,8 +149,8 @@ mod tests { ], Ok(Some("abcde")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( RightFunc::new(), @@ -149,8 +160,8 @@ mod tests { ], Ok(Some("cde")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( RightFunc::new(), @@ -160,8 +171,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( RightFunc::new(), @@ -171,8 +182,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( RightFunc::new(), @@ -182,8 +193,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( RightFunc::new(), @@ -193,8 +204,8 @@ mod tests { ], Ok(None), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( RightFunc::new(), @@ -204,8 +215,8 @@ mod tests { ], Ok(None), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( RightFunc::new(), @@ -215,8 +226,8 @@ mod tests { ], Ok(Some("érend")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( RightFunc::new(), @@ -226,8 +237,8 @@ mod tests { ], Ok(Some("éérend")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); #[cfg(not(feature = "unicode_expressions"))] test_function!( @@ -240,8 +251,8 @@ mod tests { "function right requires compilation with feature flag: unicode_expressions." ), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); // StringView cases @@ -304,8 +315,8 @@ mod tests { ], Ok(Some(expected.as_str())), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); } diff --git a/datafusion/functions/src/unicode/rpad.rs b/datafusion/functions/src/unicode/rpad.rs index b3e14f93526ab..784a2037cfbe1 100644 --- a/datafusion/functions/src/unicode/rpad.rs +++ b/datafusion/functions/src/unicode/rpad.rs @@ -178,7 +178,8 @@ impl ScalarUDFImpl for RPadFunc { } use super::common::{ - StringCharLen, char_count_or_boundary, try_as_scalar_i64, try_as_scalar_str, + StringCharLen, char_count_or_boundary, pad_data_capacity, try_as_scalar_i64, + try_as_scalar_str, }; /// Optimized rpad for constant target_len and fill arguments. @@ -372,7 +373,10 @@ where T: OffsetSizeTrait, { let array = if let Some(fill_array) = fill_array { - let mut builder: GenericStringBuilder = GenericStringBuilder::new(); + let mut builder: GenericStringBuilder = GenericStringBuilder::with_capacity( + string_array.len(), + pad_data_capacity(length_array), + ); let mut fill_chars_buf = Vec::new(); for ((string, target_len), fill) in string_array @@ -450,7 +454,10 @@ where builder.finish() } else { - let mut builder: GenericStringBuilder = GenericStringBuilder::new(); + let mut builder: GenericStringBuilder = GenericStringBuilder::with_capacity( + string_array.len(), + pad_data_capacity(length_array), + ); for (string, target_len) in string_array.iter().zip(length_array.iter()) { if let (Some(string), Some(target_len)) = (string, target_len) { diff --git a/datafusion/functions/src/unicode/substr.rs b/datafusion/functions/src/unicode/substr.rs index 903c03857e370..0cae2152248e0 100644 --- a/datafusion/functions/src/unicode/substr.rs +++ b/datafusion/functions/src/unicode/substr.rs @@ -17,11 +17,11 @@ use std::sync::Arc; -use crate::strings::{StringViewArrayBuilder, append_view}; +use crate::strings::append_view; use crate::utils::make_scalar_function; use arrow::array::{ Array, ArrayRef, AsArray, GenericStringArray, Int64Array, OffsetSizeTrait, - StringArrayType, StringViewArray, make_view, + StringArrayType, StringViewArray, }; use arrow::buffer::{NullBuffer, ScalarBuffer}; use arrow::datatypes::DataType; @@ -111,9 +111,8 @@ impl ScalarUDFImpl for SubstrFunc { &self.signature } - // `SubstrFunc` always generates `Utf8View` output for its efficiency. - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Utf8View) + fn return_type(&self, arg_types: &[DataType]) -> Result { + Ok(arg_types[0].clone()) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -319,131 +318,37 @@ fn string_view_substr( } } -fn values_fit_in_i32(string_array: &GenericStringArray) -> bool { - // The Arrow spec defines StringView offset fields as signed 32-bit - // integers, so the maximum representable offset is i32::MAX. - string_array - .offsets() - .last() - .map(|offset| offset.as_usize() <= i32::MAX as usize) - .unwrap_or(true) -} - -#[inline] -fn append_view_from_buffer( - views_buf: &mut Vec, - substr: &str, - byte_offset: usize, -) -> bool { - let byte_offset = - u32::try_from(byte_offset).expect("validated string buffer offset fits in i32"); - let view = make_view(substr.as_bytes(), 0, byte_offset); - views_buf.push(view); - substr.len() > 12 -} - -#[expect(clippy::needless_range_loop)] fn generic_string_substr( string_array: &GenericStringArray, args: &[ArrayRef], ) -> Result { - // We'd like to return a StringViewArray that points into the input string - // array's values buffer. Since the Arrow spec defines StringView offsets - // as i32, we can't use this approach when the values buffer is >2GB, so - // fallback to copying. - if !values_fit_in_i32(string_array) { - return generic_string_substr_copy(string_array, args); - } - let start_array = as_int64_array(&args[0])?; let count_array_opt = args.get(1).map(|a| as_int64_array(a)).transpose()?; let is_ascii = enable_ascii_fast_path(&string_array, start_array, count_array_opt); - let offsets = string_array.value_offsets(); - let mut views_buf = Vec::with_capacity(string_array.len()); - let mut has_out_of_line = false; - - // Combine null bitmaps from all inputs in bulk. let nulls = NullBuffer::union_many([ string_array.nulls(), start_array.nulls(), count_array_opt.and_then(|a| a.nulls()), ]); - for i in 0..string_array.len() { - if nulls.as_ref().is_some_and(|n| n.is_null(i)) { - views_buf.push(0); - continue; - } - - let string = string_array.value(i); - let source_offset = offsets[i].as_usize(); - let start = start_array.value(i); - let count = count_array_opt.map(|a| a.value(i)); - - let (byte_start, byte_end) = get_true_start_end(string, start, count, is_ascii)?; - has_out_of_line |= append_view_from_buffer( - &mut views_buf, - &string[byte_start..byte_end], - source_offset + byte_start, - ); - } + let result = (0..string_array.len()) + .map(|i| { + if nulls.as_ref().is_some_and(|n| n.is_null(i)) { + return Ok(None); + } - let views_buf = ScalarBuffer::from(views_buf); + let string = string_array.value(i); + let start = start_array.value(i); + let count = count_array_opt.map(|a| a.value(i)); - // If all result strings are stored inline, we don't need to retain the - // input string array. - let data_buffers = if has_out_of_line { - vec![string_array.values().clone()] - } else { - vec![] - }; + let (byte_start, byte_end) = + get_true_start_end(string, start, count, is_ascii)?; + Ok(Some(&string[byte_start..byte_end])) + }) + .collect::>>()?; - // Safety: - // (1) The blocks of the given views are all provided - // (2) Each referenced range in the source values buffer is within bounds - unsafe { - let array = StringViewArray::new_unchecked(views_buf, data_buffers, nulls); - Ok(Arc::new(array) as ArrayRef) - } -} - -// Fallback for `generic_string_substr` if we can't use zerocopy because the -// input string array is too large. -fn generic_string_substr_copy( - string_array: &GenericStringArray, - args: &[ArrayRef], -) -> Result { - let start_array = as_int64_array(&args[0])?; - let count_array_opt = args.get(1).map(|a| as_int64_array(a)).transpose()?; - - let is_ascii = enable_ascii_fast_path(&string_array, start_array, count_array_opt); - - // Combine null bitmaps from all inputs in bulk. - let nulls = NullBuffer::union_many([ - string_array.nulls(), - start_array.nulls(), - count_array_opt.and_then(|a| a.nulls()), - ]); - - let len = string_array.len(); - let mut result_builder = StringViewArrayBuilder::with_capacity(len); - - for i in 0..len { - if nulls.as_ref().is_some_and(|n| n.is_null(i)) { - result_builder.append_placeholder(); - continue; - } - - let string = string_array.value(i); - let start = start_array.value(i); - let count = count_array_opt.map(|a| a.value(i)); - - let (byte_start, byte_end) = get_true_start_end(string, start, count, is_ascii)?; - result_builder.append_value(&string[byte_start..byte_end]); - } - - Ok(Arc::new(result_builder.finish(nulls)?) as ArrayRef) + Ok(Arc::new(result) as ArrayRef) } #[cfg(test)] @@ -451,9 +356,10 @@ mod tests { use std::sync::Arc; use arrow::array::{ - Array, ArrayRef, AsArray, Int64Array, StringArray, StringViewArray, + Array, ArrayRef, AsArray, Int64Array, LargeStringArray, StringArray, + StringViewArray, }; - use arrow::datatypes::DataType::Utf8View; + use arrow::datatypes::DataType::{LargeUtf8, Utf8, Utf8View}; use datafusion_common::{Result, ScalarValue, exec_err}; use datafusion_expr::{ColumnarValue, ScalarUDFImpl}; @@ -563,8 +469,21 @@ mod tests { ], Ok(Some("alphabet")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray + ); + test_function!( + SubstrFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::LargeUtf8(Some( + "alphabet".to_string() + ))), + ColumnarValue::Scalar(ScalarValue::from(0i64)), + ], + Ok(Some("alphabet")), + &str, + LargeUtf8, + LargeStringArray ); test_function!( SubstrFunc::new(), @@ -574,8 +493,8 @@ mod tests { ], Ok(Some("ésoj")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -585,8 +504,8 @@ mod tests { ], Ok(Some("joséésoj")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -596,8 +515,8 @@ mod tests { ], Ok(Some("alphabet")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -607,8 +526,8 @@ mod tests { ], Ok(Some("lphabet")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -618,8 +537,8 @@ mod tests { ], Ok(Some("phabet")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -629,8 +548,8 @@ mod tests { ], Ok(Some("alphabet")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -640,8 +559,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -651,8 +570,8 @@ mod tests { ], Ok(None), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -663,8 +582,8 @@ mod tests { ], Ok(Some("ph")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -675,8 +594,8 @@ mod tests { ], Ok(Some("phabet")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -687,8 +606,8 @@ mod tests { ], Ok(Some("alph")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); // starting from 5 (10 + -5) test_function!( @@ -700,8 +619,8 @@ mod tests { ], Ok(Some("alph")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); // starting from -1 (4 + -5) test_function!( @@ -713,8 +632,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); // starting from 0 (5 + -5) test_function!( @@ -726,8 +645,8 @@ mod tests { ], Ok(Some("")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -738,8 +657,8 @@ mod tests { ], Ok(None), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -750,8 +669,8 @@ mod tests { ], Ok(None), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -762,8 +681,8 @@ mod tests { ], exec_err!("negative count not allowed: -1"), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -774,8 +693,8 @@ mod tests { ], Ok(Some("és")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); #[cfg(not(feature = "unicode_expressions"))] test_function!( @@ -788,8 +707,8 @@ mod tests { "function substr requires compilation with feature flag: unicode_expressions." ), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -799,8 +718,8 @@ mod tests { ], exec_err!("start position overflow: -9223372036854775808"), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -811,8 +730,8 @@ mod tests { ], exec_err!("start position overflow: -9223372036854775808"), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); test_function!( SubstrFunc::new(), @@ -823,8 +742,8 @@ mod tests { ], Ok(Some("arge count")), &str, - Utf8View, - StringViewArray + Utf8, + StringArray ); Ok(()) @@ -832,7 +751,6 @@ mod tests { #[test] fn test_sliced_string_array_array_args() -> Result<()> { - // Use strings longer than 12 bytes so the result views are out-of-line. let string_array = Arc::new(StringArray::from(vec![ "skipped_prefix_value", "alphabet_long_string", @@ -843,7 +761,7 @@ mod tests { let count_array = Arc::new(Int64Array::from(vec![15, 14])) as ArrayRef; let result = super::substr(&[string_array, start_array, count_array])?; - let result = result.as_string_view(); + let result = result.as_string::(); assert_eq!(result.value(0), "phabet_long_str"); assert_eq!(result.value(1), "ésojanother_lo"); diff --git a/datafusion/functions/src/unicode/substrindex.rs b/datafusion/functions/src/unicode/substrindex.rs index d122a34a9fc38..f9f0bafa04309 100644 --- a/datafusion/functions/src/unicode/substrindex.rs +++ b/datafusion/functions/src/unicode/substrindex.rs @@ -281,7 +281,7 @@ where for i in 0..num_rows { if nulls.as_ref().is_some_and(|n| n.is_null(i)) { - builder.append_placeholder(); + builder.try_append_placeholder()?; continue; } // SAFETY: `i < num_rows` and the union of input nulls is valid at i, @@ -289,7 +289,7 @@ where let string = unsafe { string_array.value_unchecked(i) }; let delimiter = unsafe { delimiter_array.value_unchecked(i) }; let n = unsafe { count_array.value_unchecked(i) }; - builder.append_value(substr_index_slice(string, delimiter, n)); + builder.try_append_value(substr_index_slice(string, delimiter, n))?; } Ok(Arc::new(builder.finish(nulls)?) as ArrayRef) @@ -487,13 +487,13 @@ where let nulls = string_array.nulls().cloned(); for i in 0..string_array.len() { if nulls.as_ref().is_some_and(|n| n.is_null(i)) { - builder.append_placeholder(); + builder.try_append_placeholder()?; continue; } // SAFETY: `i < string_array.len()` and `nulls` is valid at i, so the // input is also valid at i. let s = unsafe { string_array.value_unchecked(i) }; - builder.append_value(f(s)); + builder.try_append_value(f(s))?; } Ok(Arc::new(builder.finish(nulls)?) as ArrayRef) } @@ -797,6 +797,27 @@ mod tests { Ok(()) } + #[test] + fn test_substr_index_all_nulls() -> Result<()> { + use super::substr_index_general; + use crate::strings::GenericStringArrayBuilder; + + let strings = StringArray::from(vec![None::<&str>, None]); + let delimiters = StringArray::from(vec![None::<&str>, Some(".")]); + let counts = Int64Array::from(vec![None, None]); + + let result = substr_index_general( + &strings, + &delimiters, + &counts, + GenericStringArrayBuilder::::with_capacity(strings.len(), 0), + )?; + let result = result.as_string::(); + assert_eq!(result, &StringArray::from(vec![None::<&str>, None])); + + Ok(()) + } + #[test] fn test_substr_index_utf8view_array_sliced() -> Result<()> { use super::substr_index_view; diff --git a/datafusion/functions/src/unicode/translate.rs b/datafusion/functions/src/unicode/translate.rs index 29dc660b86f62..85e83897f41da 100644 --- a/datafusion/functions/src/unicode/translate.rs +++ b/datafusion/functions/src/unicode/translate.rs @@ -15,13 +15,16 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ - ArrayAccessor, ArrayIter, ArrayRef, AsArray, LargeStringBuilder, StringBuilder, - StringLikeArrayBuilder, StringViewBuilder, -}; +use arrow::array::{Array, ArrayRef, AsArray, GenericStringArray, StringArrayType}; +use arrow::buffer::NullBuffer; use arrow::datatypes::DataType; use datafusion_common::HashMap; +use super::common::try_as_scalar_str; +use crate::strings::{ + BulkNullStringArrayBuilder, GenericStringArrayBuilder, StringViewArrayBuilder, + StringWriter, +}; use crate::utils::make_scalar_function; use datafusion_common::{Result, exec_err}; use datafusion_expr::TypeSignature::Exact; @@ -96,14 +99,7 @@ impl ScalarUDFImpl for TranslateFunc { try_as_scalar_str(&args.args[1]), try_as_scalar_str(&args.args[2]), ) { - let to_chars: Vec = to_str.chars().collect(); - - let mut from_map: HashMap = HashMap::new(); - for (index, c) in from_str.chars().enumerate() { - from_map.entry(c).or_insert(index); - } - - let ascii_table = build_ascii_translate_table(from_str, to_str); + let table = build_translate_table(from_str, to_str); let string_array = args.args[0].to_array_of_size(args.number_rows)?; let len = string_array.len(); @@ -111,38 +107,24 @@ impl ScalarUDFImpl for TranslateFunc { let result = match string_array.data_type() { DataType::Utf8View => { let arr = string_array.as_string_view(); - let builder = StringViewBuilder::with_capacity(len); - translate_with_map( - arr, - &from_map, - &to_chars, - ascii_table.as_ref(), - builder, - ) + let builder = StringViewArrayBuilder::with_capacity(len); + translate_with_table(&arr, &table, builder) } DataType::Utf8 => { let arr = string_array.as_string::(); - let builder = - StringBuilder::with_capacity(len, arr.value_data().len()); - translate_with_map( - arr, - &from_map, - &to_chars, - ascii_table.as_ref(), - builder, - ) + let builder = GenericStringArrayBuilder::::with_capacity( + len, + arr.value_data().len(), + ); + translate_with_table(&arr, &table, builder) } DataType::LargeUtf8 => { let arr = string_array.as_string::(); - let builder = - LargeStringBuilder::with_capacity(len, arr.value_data().len()); - translate_with_map( - arr, - &from_map, - &to_chars, - ascii_table.as_ref(), - builder, - ) + let builder = GenericStringArrayBuilder::::with_capacity( + len, + arr.value_data().len(), + ); + translate_with_table(&arr, &table, builder) } other => { return exec_err!( @@ -162,8 +144,6 @@ impl ScalarUDFImpl for TranslateFunc { } } -use super::common::try_as_scalar_str; - fn invoke_translate(args: &[ArrayRef]) -> Result { let len = args[0].len(); match args[0].data_type() { @@ -171,24 +151,28 @@ fn invoke_translate(args: &[ArrayRef]) -> Result { let string_array = args[0].as_string_view(); let from_array = args[1].as_string::(); let to_array = args[2].as_string::(); - let builder = StringViewBuilder::with_capacity(len); - translate(string_array, from_array, to_array, builder) + let builder = StringViewArrayBuilder::with_capacity(len); + translate(&string_array, from_array, to_array, builder) } DataType::Utf8 => { let string_array = args[0].as_string::(); let from_array = args[1].as_string::(); let to_array = args[2].as_string::(); - let builder = - StringBuilder::with_capacity(len, string_array.value_data().len()); - translate(string_array, from_array, to_array, builder) + let builder = GenericStringArrayBuilder::::with_capacity( + len, + string_array.value_data().len(), + ); + translate(&string_array, from_array, to_array, builder) } DataType::LargeUtf8 => { let string_array = args[0].as_string::(); let from_array = args[1].as_string::(); let to_array = args[2].as_string::(); - let builder = - LargeStringBuilder::with_capacity(len, string_array.value_data().len()); - translate(string_array, from_array, to_array, builder) + let builder = GenericStringArrayBuilder::::with_capacity( + len, + string_array.value_data().len(), + ); + translate(&string_array, from_array, to_array, builder) } other => { exec_err!("Unsupported data type {other:?} for function translate") @@ -196,69 +180,89 @@ fn invoke_translate(args: &[ArrayRef]) -> Result { } } -/// Replaces each character in string that matches a character in the from set with the corresponding character in the to set. If from is longer than to, occurrences of the extra characters in from are deleted. +/// Replaces each character in string that matches a character in the from set +/// with the corresponding character in the to set. If from is longer than to, +/// occurrences of the extra characters in from are deleted. +/// /// translate('12345', '143', 'ax') = 'a2x5' -fn translate<'a, V, B, O>( - string_array: V, - from_array: B, - to_array: B, +fn translate<'a, S, O>( + string_array: &S, + from_array: &GenericStringArray, + to_array: &GenericStringArray, mut builder: O, ) -> Result where - V: ArrayAccessor, - B: ArrayAccessor, - O: StringLikeArrayBuilder, + S: StringArrayType<'a>, + O: BulkNullStringArrayBuilder, { - let string_array_iter = ArrayIter::new(string_array); - let from_array_iter = ArrayIter::new(from_array); - let to_array_iter = ArrayIter::new(to_array); - - let mut from_map: HashMap = HashMap::new(); - let mut to_chars: Vec = Vec::new(); - let mut result_buf = String::new(); - - for ((string, from), to) in string_array_iter.zip(from_array_iter).zip(to_array_iter) - { - match (string, from, to) { - (Some(string), Some(from), Some(to)) => { - from_map.clear(); - to_chars.clear(); - result_buf.clear(); - - for (index, c) in from.chars().enumerate() { - from_map.entry(c).or_insert(index); - } + let mut from_map: HashMap> = HashMap::new(); + let len = string_array.len(); + let nulls = NullBuffer::union_many([ + string_array.nulls(), + from_array.nulls(), + to_array.nulls(), + ]); + + if let Some(nulls_ref) = nulls.as_ref() { + for i in 0..len { + if nulls_ref.is_null(i) { + builder.append_placeholder(); + continue; + } - to_chars.extend(to.chars()); + // SAFETY: union of input nulls is non-null at i, so each input is too. + let string = unsafe { string_array.value_unchecked(i) }; + let from = unsafe { from_array.value_unchecked(i) }; + let to = unsafe { to_array.value_unchecked(i) }; + append_translated_row(&mut builder, string, from, to, &mut from_map); + } + } else { + for i in 0..len { + // SAFETY: i < len, and no input has a null buffer. + let string = unsafe { string_array.value_unchecked(i) }; + let from = unsafe { from_array.value_unchecked(i) }; + let to = unsafe { to_array.value_unchecked(i) }; + append_translated_row(&mut builder, string, from, to, &mut from_map); + } + } - translate_char_by_char(string, &from_map, &to_chars, &mut result_buf); + builder.finish(nulls) +} - builder.append_value(&result_buf); - } - _ => builder.append_null(), - } +#[inline] +fn append_translated_row( + builder: &mut B, + string: &str, + from: &str, + to: &str, + from_map: &mut HashMap>, +) { + if let Some(ascii_table) = build_ascii_translate_table(from, to) { + append_translated_ascii(builder, string, &ascii_table); + return; + } + + from_map.clear(); + let mut to_iter = to.chars(); + for c in from.chars() { + let replacement = to_iter.next(); + from_map.entry(c).or_insert(replacement); } - Ok(builder.finish()) + builder.append_with(|w| write_translated_chars(w, string, from_map)); } -/// Translate `input` character-by-character using `from_map` and `to_chars`, -/// appending the result to `buf`. #[inline] -fn translate_char_by_char( +fn write_translated_chars( + w: &mut W, input: &str, - from_map: &HashMap, - to_chars: &[char], - buf: &mut String, + from_map: &HashMap>, ) { for c in input.chars() { match from_map.get(&c) { - Some(n) => { - if let Some(&replacement) = to_chars.get(*n) { - buf.push(replacement); - } - } - None => buf.push(c), + Some(Some(r)) => w.write_char(*r), + Some(None) => {} // delete: `from` had no corresponding `to` char + None => w.write_char(c), } } } @@ -268,86 +272,170 @@ fn translate_char_by_char( /// value > 127 works since valid ASCII is 0–127. const ASCII_DELETE: u8 = 0xFF; -/// If `from` and `to` are both ASCII, build a fixed-size lookup table for -/// translation. Each entry maps an input byte to its replacement byte, or to -/// [`ASCII_DELETE`] if the character should be removed. Returns `None` if -/// either string contains non-ASCII characters. -fn build_ascii_translate_table(from: &str, to: &str) -> Option<[u8; 128]> { +/// Lookup table for ASCII-only translation. Entries 0..128 map input bytes to +/// replacement bytes, or `ASCII_DELETE` if the character should be deleted. +/// Entries 128..256 map to themselves so non-ASCII bytes pass through +/// unchanged. +#[derive(Debug)] +struct AsciiTranslateTable { + map: [u8; 256], + has_delete: bool, +} + +/// We use a byte-indexed table when both `from` and `to` strings are ASCII, +/// otherwise a char-indexed map where `None` means delete. +#[expect( + clippy::large_enum_variant, + reason = "one instance per call, passed by reference" +)] +enum TranslateTable { + Byte(AsciiTranslateTable), + Char(HashMap>), +} + +#[inline] +fn build_translate_table(from: &str, to: &str) -> TranslateTable { + if let Some(ascii) = build_ascii_translate_table(from, to) { + return TranslateTable::Byte(ascii); + } + let mut from_map: HashMap> = HashMap::with_capacity(from.len()); + let mut to_iter = to.chars(); + for c in from.chars() { + let replacement = to_iter.next(); + from_map.entry(c).or_insert(replacement); + } + TranslateTable::Char(from_map) +} + +/// Returns `None` if either string contains non-ASCII characters. +fn build_ascii_translate_table(from: &str, to: &str) -> Option { if !from.is_ascii() || !to.is_ascii() { return None; } - let mut table = [0u8; 128]; - for i in 0..128u8 { - table[i as usize] = i; - } + let to_bytes = to.as_bytes(); + let mut map = std::array::from_fn::(|i| i as u8); let mut seen = [false; 128]; + let mut has_delete = false; + for (i, from_byte) in from.bytes().enumerate() { let idx = from_byte as usize; if !seen[idx] { seen[idx] = true; if i < to_bytes.len() { - table[idx] = to_bytes[i]; + map[idx] = to_bytes[i]; } else { - table[idx] = ASCII_DELETE; + map[idx] = ASCII_DELETE; + has_delete = true; } } } - Some(table) + + Some(AsciiTranslateTable { map, has_delete }) +} + +#[inline] +fn append_translated_ascii( + builder: &mut B, + input: &str, + table: &AsciiTranslateTable, +) { + // Fast path: equal-length byte-to-byte map when no deletions. + if !table.has_delete { + // SAFETY: ASCII source bytes map to ASCII replacements; non-ASCII + // bytes 128..256 map to themselves, so multi-byte UTF-8 sequences + // pass through unchanged. Output length equals input length and + // remains valid UTF-8. + unsafe { + builder.append_byte_map(input.as_bytes(), |b| table.map[b as usize]); + } + } else { + builder.append_with(|w| write_translated_ascii(w, input, table)); + } } -/// Optimized translate for constant `from` and `to` arguments: uses a pre-built -/// translation map instead of rebuilding it for every row. When an ASCII byte -/// lookup table is provided, ASCII input rows use the lookup table; non-ASCII -/// inputs fall back to the char-based map. -fn translate_with_map<'a, V, O>( - string_array: V, - from_map: &HashMap, - to_chars: &[char], - ascii_table: Option<&[u8; 128]>, +#[inline] +fn write_translated_ascii( + w: &mut W, + input: &str, + table: &AsciiTranslateTable, +) { + let bytes = input.as_bytes(); + let mut copy_start = 0; + + for (i, &b) in bytes.iter().enumerate() { + let mapped = table.map[b as usize]; + if mapped == b { + continue; + } + + if copy_start < i { + w.write_str(&input[copy_start..i]); + } + if mapped != ASCII_DELETE { + w.write_char(mapped as char); + } + copy_start = i + 1; + } + + if copy_start < input.len() { + w.write_str(&input[copy_start..]); + } +} + +fn translate_with_table<'a, S, O>( + string_array: &S, + table: &TranslateTable, mut builder: O, ) -> Result where - V: ArrayAccessor, - O: StringLikeArrayBuilder, + S: StringArrayType<'a>, + O: BulkNullStringArrayBuilder, { - let mut result_buf = String::new(); - let mut ascii_buf: Vec = Vec::new(); - - for string in ArrayIter::new(string_array) { - match string { - Some(s) => { - // Fast path: byte-level table lookup for ASCII strings - if let Some(table) = ascii_table - && s.is_ascii() - { - ascii_buf.clear(); - for &b in s.as_bytes() { - let mapped = table[b as usize]; - if mapped != ASCII_DELETE { - ascii_buf.push(mapped); - } - } - // SAFETY: all bytes are ASCII, hence valid UTF-8. - builder.append_value(unsafe { - std::str::from_utf8_unchecked(&ascii_buf) - }); - } else { - result_buf.clear(); - translate_char_by_char(s, from_map, to_chars, &mut result_buf); - builder.append_value(&result_buf); - } + let len = string_array.len(); + let nulls = string_array.nulls().cloned(); + + if let Some(nulls_ref) = nulls.as_ref() { + for i in 0..len { + if nulls_ref.is_null(i) { + builder.append_placeholder(); + continue; } - None => builder.append_null(), + + // SAFETY: input null buffer is non-null at i. + let s = unsafe { string_array.value_unchecked(i) }; + apply_translate_table(&mut builder, s, table); + } + } else { + for i in 0..len { + // SAFETY: no null buffer means every index is valid. + let s = unsafe { string_array.value_unchecked(i) }; + apply_translate_table(&mut builder, s, table); } } - Ok(builder.finish()) + builder.finish(nulls) +} + +#[inline] +fn apply_translate_table( + builder: &mut B, + input: &str, + table: &TranslateTable, +) { + match table { + TranslateTable::Byte(t) => append_translated_ascii(builder, input, t), + TranslateTable::Char(m) => { + builder.append_with(|w| write_translated_chars(w, input, m)) + } + } } #[cfg(test)] mod tests { - use arrow::array::{Array, StringArray, StringViewArray}; + use std::sync::Arc; + + use arrow::array::{Array, ArrayRef, StringArray, StringViewArray}; use arrow::datatypes::DataType::{Utf8, Utf8View}; use datafusion_common::{Result, ScalarValue}; @@ -430,8 +518,7 @@ mod tests { Utf8, StringArray ); - // Non-ASCII input with ASCII scalar from/to: exercises the - // char-based fallback within translate_with_map. + // Non-ASCII input with ASCII scalar from/to. test_function!( TranslateFunc::new(), vec![ @@ -502,4 +589,27 @@ mod tests { Ok(()) } + + #[test] + fn test_array_args_with_nulls() -> Result<()> { + let string_array = Arc::new(StringArray::from(vec![ + Some("café!"), + Some("abc"), + Some("abc"), + ])) as ArrayRef; + let from_array = + Arc::new(StringArray::from(vec![Some("!"), Some("a"), None])) as ArrayRef; + let to_array = + Arc::new(StringArray::from(vec![Some(""), Some("x"), Some("y")])) as ArrayRef; + + let result = super::invoke_translate(&[string_array, from_array, to_array])?; + let result = result.as_any().downcast_ref::().unwrap(); + + assert_eq!(result.len(), 3); + assert_eq!(result.value(0), "café"); + assert_eq!(result.value(1), "xbc"); + assert!(result.is_null(2)); + + Ok(()) + } } diff --git a/datafusion/functions/src/utils.rs b/datafusion/functions/src/utils.rs index b9bde1454994c..b93bdb0b0d3bb 100644 --- a/datafusion/functions/src/utils.rs +++ b/datafusion/functions/src/utils.rs @@ -74,6 +74,28 @@ get_optimal_return_type!(utf8_to_str_type, DataType::LargeUtf8, DataType::Utf8); // `utf8_to_int_type`: returns either a Int32 or Int64 based on the input type size. get_optimal_return_type!(utf8_to_int_type, DataType::Int64, DataType::Int32); +/// Transforms the leaf type while preserving supported encoding containers. +/// +/// Keep encoded type handling centralized here so additional encodings can be +/// supported without changing each function's return type implementation. +pub(crate) fn transform_leaf_type_preserving_encoding( + arg_type: &DataType, + transform: &F, +) -> Result +where + F: Fn(&DataType) -> Result, +{ + match arg_type { + DataType::Dictionary(key_type, value_type) => Ok(DataType::Dictionary( + key_type.clone(), + Box::new(transform_leaf_type_preserving_encoding( + value_type, transform, + )?), + )), + _ => transform(arg_type), + } +} + /// Creates a scalar function implementation for the given function. /// * `inner` - the function to be executed /// * `hints` - hints to be used when expanding scalars to arrays @@ -133,6 +155,72 @@ pub fn calculate_binary_math( right: &ColumnarValue, fun: F, ) -> Result>> +where + L: ArrowPrimitiveType, + R: ArrowPrimitiveType, + O: ArrowPrimitiveType, + F: Fn(L::Native, R::Native) -> Result, + R::Native: TryFrom, +{ + calculate_binary_math_cast::(left, right, fun, &R::DATA_TYPE) +} + +/// Computes a binary math function for input arrays using a specified function +/// and applies rescaling to given precision and scale. +/// Generic types: +/// - `L`: Left array decimal type +/// - `R`: Right array primitive type +/// - `O`: Output array decimal type +/// - `F`: Functor computing `fun(l: L, r: R) -> Result` +#[deprecated( + since = "55.0.0", + note = "Use `calculate_binary_decimal_math_cast` instead" +)] +pub fn calculate_binary_decimal_math( + left: &dyn Array, + right: &ColumnarValue, + fun: F, + precision: u8, + scale: i8, +) -> Result>> +where + L: DecimalType, + R: ArrowPrimitiveType, + O: DecimalType, + F: Fn(L::Native, R::Native) -> Result, + R::Native: TryFrom, +{ + calculate_binary_decimal_math_cast::( + left, + right, + fun, + precision, + scale, + &R::DATA_TYPE, + ) +} + +/// Computes a binary math function for input arrays using a specified function. +/// +/// It casts the right operand to `cast_target` instead of the default `R::DATA_TYPE` to preserve +/// the right operand scale. +/// +/// # Type Parameters +/// - `L`: Left array primitive type +/// - `R`: Right array primitive type +/// - `O`: Output array primitive type +/// - `F`: Functor computing `fun(l: L, r: R) -> Result` +/// # Arguments +/// - `left`: Left input array +/// - `right`: Right input array or scalar value +/// - `fun`: Function of type `F` +/// - `cast_target`: Data type to cast right operand to before applying function +fn calculate_binary_math_cast( + left: &dyn Array, + right: &ColumnarValue, + fun: F, + cast_target: &DataType, +) -> Result>> where L: ArrowPrimitiveType, R: ArrowPrimitiveType, @@ -141,7 +229,7 @@ where R::Native: TryFrom, { let left = left.as_primitive::(); - let right = right.cast_to(&R::DATA_TYPE, None)?; + let right = right.cast_to(cast_target, None)?; let result = match right { ColumnarValue::Scalar(scalar) => { if scalar.is_null() { @@ -151,9 +239,7 @@ where } else { let right = R::Native::try_from(scalar.clone()).map_err(|_| { DataFusionError::NotImplemented(format!( - "Cannot convert scalar value {} to {}", - &scalar, - R::DATA_TYPE + "Cannot convert scalar value {scalar} to {cast_target}" )) })?; left.try_unary::<_, O, _>(|lvalue| fun(lvalue, right))? @@ -168,18 +254,30 @@ where } /// Computes a binary math function for input arrays using a specified function -/// and apply rescaling to given precision and scale. -/// Generic types: +/// and applies rescaling to given precision and scale. +/// +/// It casts the right operand to `cast_target` instead of the default `R::DATA_TYPE` to preserve +/// the right operand scale. +/// +/// # Type Parameters /// - `L`: Left array decimal type /// - `R`: Right array primitive type /// - `O`: Output array decimal type /// - `F`: Functor computing `fun(l: L, r: R) -> Result` -pub fn calculate_binary_decimal_math( +/// # Arguments +/// - `left`: Left input array +/// - `right`: Right input array or scalar value +/// - `fun`: Function of type `F` +/// - `precision`: Precision to apply to output decimal array +/// - `scale`: Scale to apply to output decimal array +/// - `cast_target`: Data type to cast right operand to before applying function +pub fn calculate_binary_decimal_math_cast( left: &dyn Array, right: &ColumnarValue, fun: F, precision: u8, scale: i8, + cast_target: &DataType, ) -> Result>> where L: DecimalType, @@ -188,7 +286,8 @@ where F: Fn(L::Native, R::Native) -> Result, R::Native: TryFrom, { - let result_array = calculate_binary_math::(left, right, fun)?; + let result_array = + calculate_binary_math_cast::(left, right, fun, cast_target)?; Ok(Arc::new( result_array .as_ref() diff --git a/datafusion/macros/Cargo.toml b/datafusion/macros/Cargo.toml index 91f1dde62aaac..d5ab6a8fff624 100644 --- a/datafusion/macros/Cargo.toml +++ b/datafusion/macros/Cargo.toml @@ -46,4 +46,4 @@ proc-macro = true [dependencies] datafusion-doc = { workspace = true } quote = "1.0.44" -syn = { version = "2.0.117", features = ["full"] } +syn = { version = "3.0.2", features = ["full"] } diff --git a/datafusion/macros/src/user_doc.rs b/datafusion/macros/src/user_doc.rs index ce9e7d55ef103..2dde56a4ff694 100644 --- a/datafusion/macros/src/user_doc.rs +++ b/datafusion/macros/src/user_doc.rs @@ -38,7 +38,7 @@ use syn::{DeriveInput, LitStr, parse_macro_input}; /// #[user_doc( /// doc_section(label = "Time and Date Functions"), /// description = r"Converts a value to a date (`YYYY-MM-DD`).", -/// syntax_example = "to_date('2017-05-31', '%Y-%m-%d')", +/// syntax_example = "to_date(expression[, ..., format_n])", /// sql_example = r#"```sql /// > select to_date('2023-01-31'); /// +-----------------------------+ @@ -77,7 +77,7 @@ use syn::{DeriveInput, LitStr, parse_macro_input}; /// description: None, /// }, /// r"Converts a value to a date (`YYYY-MM-DD`).".to_string(), -/// "to_date('2017-05-31', '%Y-%m-%d')".to_string(), +/// "to_date(expression[, ..., format_n])".to_string(), /// ) /// .with_sql_example( /// r#"```sql diff --git a/datafusion/optimizer/src/analyzer/function_rewrite.rs b/datafusion/optimizer/src/analyzer/function_rewrite.rs index 9faa60d939fe3..a66e3ccc0cf8a 100644 --- a/datafusion/optimizer/src/analyzer/function_rewrite.rs +++ b/datafusion/optimizer/src/analyzer/function_rewrite.rs @@ -23,9 +23,9 @@ use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{DFSchema, Result}; use crate::utils::NamePreserver; -use datafusion_expr::LogicalPlan; use datafusion_expr::expr_rewriter::FunctionRewrite; use datafusion_expr::utils::merge_schema; +use datafusion_expr::{DmlStatement, LogicalPlan, WriteOp}; use std::sync::Arc; /// Analyzer rule that invokes [`FunctionRewrite`]s on expressions @@ -58,6 +58,23 @@ impl ApplyFunctionRewrites { schema.merge(&source_schema); } + // MERGE expressions reference the target table, which is not one of + // `plan.inputs()`. Rebuild the target schema from the DML's + // `table_name` and `target` so those columns resolve. + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let target_schema = DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?; + schema.merge(&target_schema); + } + let name_preserver = NamePreserver::new(&plan); plan.map_expressions(|expr| { diff --git a/datafusion/optimizer/src/analyzer/type_coercion.rs b/datafusion/optimizer/src/analyzer/type_coercion.rs index 7b81feab47a99..d11c3e7435fde 100644 --- a/datafusion/optimizer/src/analyzer/type_coercion.rs +++ b/datafusion/optimizer/src/analyzer/type_coercion.rs @@ -43,7 +43,7 @@ use datafusion_expr::expr_rewriter::coerce_plan_expr_for_schema; use datafusion_expr::expr_schema::cast_subquery; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::type_coercion::binary::{ - comparison_coercion, like_coercion, type_union_coercion, + comparison_coercion, like_coercion, regex_coercion, type_union_coercion, }; use datafusion_expr::type_coercion::functions::{ UDFCoercionExt, fields_with_udf, value_fields_with_higher_order_udf_and_lambdas, @@ -57,9 +57,10 @@ use datafusion_expr::type_coercion::{ }; use datafusion_expr::utils::merge_schema; use datafusion_expr::{ - Cast, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, Projection, Union, - ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, is_false, - is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, lit, not, + Cast, DmlStatement, Expr, ExprSchemable, Join, Limit, LogicalPlan, Operator, + Projection, Union, ValueOrLambda, WindowFrame, WindowFrameBound, WindowFrameUnits, + WriteOp, is_false, is_not_false, is_not_true, is_not_unknown, is_true, is_unknown, + lit, not, }; /// Performs type coercion by determining the schema @@ -128,6 +129,21 @@ fn analyze_internal( schema.merge(&source_schema); } + // MERGE expressions (ON / WHEN clauses) reference the target table, which + // is not one of `plan.inputs()`. Rebuild the target schema from the DML's + // `table_name` and `target` so those columns resolve during coercion. + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let target_schema = + DFSchema::try_from_qualified_schema(table_name.clone(), &target.schema())?; + schema.merge(&target_schema); + } + // merge the outer schema for correlated subqueries // like case: // select t2.c2 from t1 where t1.c1 in (select t2.c1 from t2 where t2.c2=t1.c3) @@ -177,10 +193,60 @@ impl<'a> TypeCoercionRewriter<'a> { LogicalPlan::Join(join) => self.coerce_join(join), LogicalPlan::Union(union) => Self::coerce_union(union), LogicalPlan::Limit(limit) => Self::coerce_limit(limit), + LogicalPlan::Dml(dml) => self.coerce_dml(dml), _ => Ok(plan), } } + fn coerce_dml(&self, mut dml: DmlStatement) -> Result { + let WriteOp::MergeInto(merge_op) = &dml.op else { + return Ok(LogicalPlan::Dml(dml)); + }; + + let target_schema = DFSchema::try_from_qualified_schema( + dml.table_name.clone(), + &dml.target.schema(), + )?; + let mut merge_op = (**merge_op).clone(); + merge_op.on = self.coerce_predicate(merge_op.on, "MERGE ON condition")?; + for clause in &mut merge_op.clauses { + clause.predicate = clause + .predicate + .take() + .map(|expr| self.coerce_predicate(expr, "MERGE WHEN condition")) + .transpose()?; + + match &mut clause.action { + datafusion_expr::dml::MergeIntoAction::Update(assignments) => { + for (column, value) in assignments { + let field = target_schema.field_with_unqualified_name(column)?; + *value = value.clone().cast_to(field.data_type(), self.schema)?; + } + } + datafusion_expr::dml::MergeIntoAction::Insert { columns, values } => { + if columns.is_empty() { + for (value, field) in + values.iter_mut().zip(target_schema.fields()) + { + *value = + value.clone().cast_to(field.data_type(), self.schema)?; + } + } else { + for (column, value) in columns.iter().zip(values) { + let field = + target_schema.field_with_unqualified_name(column)?; + *value = + value.clone().cast_to(field.data_type(), self.schema)?; + } + } + } + datafusion_expr::dml::MergeIntoAction::Delete => {} + } + } + dml.op = WriteOp::MergeInto(Box::new(merge_op)); + Ok(LogicalPlan::Dml(dml)) + } + /// Coerce join equality expressions and join filter /// /// Joins must be treated specially as their equality expressions are stored @@ -212,7 +278,7 @@ impl<'a> TypeCoercionRewriter<'a> { // Join filter must be boolean join.filter = join .filter - .map(|expr| self.coerce_join_filter(expr)) + .map(|expr| self.coerce_predicate(expr, "Join condition")) .transpose()?; Ok(LogicalPlan::Join(join)) @@ -280,12 +346,14 @@ impl<'a> TypeCoercionRewriter<'a> { })) } - fn coerce_join_filter(&self, expr: Expr) -> Result { + fn coerce_predicate(&self, expr: Expr, description: &str) -> Result { let expr_type = expr.get_type(self.schema)?; match expr_type { DataType::Boolean => Ok(expr), DataType::Null => expr.cast_to(&DataType::Boolean, self.schema), - other => plan_err!("Join condition must be boolean type, but got {other:?}"), + other => { + plan_err!("{description} must be boolean type, but got {other:?}") + } } } @@ -442,6 +510,38 @@ impl<'a> TypeCoercionRewriter<'a> { Ok(e) } + + /// Coerce the value and pattern expressions of a string pattern matching + /// expression (`LIKE`, `ILIKE` or `SIMILAR TO`) to a common type using + /// the provided coercion rules. `LIKE` can preserve a dictionary-encoded + /// value expression, while regex array kernels require both operands to + /// have the same physical string type. + fn coerce_like_operands( + &self, + expr: Expr, + pattern: Expr, + coercion: fn(&DataType, &DataType) -> Option, + op_name: &str, + preserve_utf8_dictionary: bool, + ) -> Result<(Box, Box)> { + let left_type = expr.get_type(self.schema)?; + let right_type = pattern.get_type(self.schema)?; + let coerced_type = coercion(&left_type, &right_type).ok_or_else(|| { + plan_datafusion_err!( + "There isn't a common type to coerce {left_type} and {right_type} in {op_name} expression" + ) + })?; + let expr = match left_type { + DataType::Dictionary(_, inner) + if preserve_utf8_dictionary && *inner == DataType::Utf8 => + { + Box::new(expr) + } + _ => Box::new(expr.cast_to(&coerced_type, self.schema)?), + }; + let pattern = Box::new(pattern.cast_to(&coerced_type, self.schema)?); + Ok((expr, pattern)) + } } impl TreeNodeRewriter for TypeCoercionRewriter<'_> { @@ -588,23 +688,14 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { escape_char, case_insensitive, }) => { - let left_type = expr.get_type(self.schema)?; - let right_type = pattern.get_type(self.schema)?; - let coerced_type = like_coercion(&left_type, &right_type).ok_or_else(|| { - let op_name = if case_insensitive { - "ILIKE" - } else { - "LIKE" - }; - plan_datafusion_err!( - "There isn't a common type to coerce {left_type} and {right_type} in {op_name} expression" - ) - })?; - let expr = match left_type { - DataType::Dictionary(_, inner) if *inner == DataType::Utf8 => expr, - _ => Box::new(expr.cast_to(&coerced_type, self.schema)?), - }; - let pattern = Box::new(pattern.cast_to(&coerced_type, self.schema)?); + let op_name = if case_insensitive { "ILIKE" } else { "LIKE" }; + let (expr, pattern) = self.coerce_like_operands( + *expr, + *pattern, + like_coercion, + op_name, + true, + )?; Ok(Transformed::yes(Expr::Like(Like::new( negated, expr, @@ -613,6 +704,32 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { case_insensitive, )))) } + Expr::SimilarTo(Like { + negated, + expr, + pattern, + escape_char, + case_insensitive, + }) => { + // `SIMILAR TO` is planned as a regex operator, so its operands + // must be coerced to a common string type using the same + // coercion rules as the physical regex operators. Otherwise + // mismatched operand types panic during execution. + let (expr, pattern) = self.coerce_like_operands( + *expr, + *pattern, + regex_coercion, + "SIMILAR TO", + false, + )?; + Ok(Transformed::yes(Expr::SimilarTo(Like::new( + negated, + expr, + pattern, + escape_char, + case_insensitive, + )))) + } Expr::BinaryExpr(BinaryExpr { left, op, right }) => { let (left, right) = self.coerce_binary_op(*left, self.schema, op, *right, self.schema)?; @@ -715,6 +832,12 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { }) => { let new_expr = coerce_arguments_for_signature(args, self.schema, func.as_ref())?; + + let filter = filter + .map(|filter| filter.cast_to(&DataType::Boolean, self.schema)) + .transpose()? + .map(Box::new); + Ok(Transformed::yes(Expr::AggregateFunction( expr::AggregateFunction::new_udf( func, @@ -752,6 +875,11 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { } }; + let filter = filter + .map(|filter| filter.cast_to(&DataType::Boolean, self.schema)) + .transpose()? + .map(Box::new); + let new_expr = Expr::from(WindowFunction { fun, params: expr::WindowFunctionParams { @@ -801,7 +929,6 @@ impl TreeNodeRewriter for TypeCoercionRewriter<'_> { | Expr::Column(_) | Expr::ScalarVariable(_, _) | Expr::Literal(_, _) - | Expr::SimilarTo(_) | Expr::IsNotNull(_) | Expr::IsNull(_) | Expr::Cast(_) @@ -1311,7 +1438,9 @@ mod test { use crate::assert_analyzed_plan_with_config_eq_snapshot; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{TransformedResult, TreeNode}; - use datafusion_common::{DFSchema, DFSchemaRef, Result, ScalarValue, Spans}; + use datafusion_common::{ + DFSchema, DFSchemaRef, Result, ScalarValue, Spans, TableReference, + }; use datafusion_expr::expr::{self, InSubquery, Like, ScalarFunction}; use datafusion_expr::logical_plan::{EmptyRelation, Projection, Sort}; use datafusion_expr::test::function_stub::avg_udaf; @@ -1322,7 +1451,6 @@ mod test { col, create_udaf, is_true, lit, }; use datafusion_functions_aggregate::average::AvgAccumulator; - use datafusion_sql::TableReference; fn empty() -> Arc { Arc::new(LogicalPlan::EmptyRelation(EmptyRelation { @@ -1475,6 +1603,88 @@ mod test { ) } + #[test] + fn merge_into_resolves_and_coerces_target_and_source_columns() -> Result<()> { + use datafusion_expr::dml::{ + MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, + }; + use datafusion_expr::logical_plan::table_scan; + use datafusion_expr::{DmlStatement, WriteOp}; + + // Target table `target(id: UInt32)`. + let target_table_name = TableReference::bare("target"); + let target_arrow_schema = + Schema::new(vec![Field::new("id", DataType::UInt32, false)]); + let target_plan = + table_scan(Some(target_table_name.clone()), &target_arrow_schema, None)? + .build()?; + let target_source = match &target_plan { + LogicalPlan::TableScan(ts) => Arc::clone(&ts.source), + _ => unreachable!("table_scan() always builds a TableScan"), + }; + + // Source plan `source(id: Int64)` — deliberately a different numeric + // type than `target.id` so the `ON` comparison needs a CAST. + let source_arrow_schema = + Schema::new(vec![Field::new("id", DataType::Int64, false)]); + let source_plan = + table_scan(Some("source"), &source_arrow_schema, None)?.build()?; + + // `ON target.id = source.id`. Resolving `target.id` requires the + // target schema to be visible to the analyzer, which only sees + // `plan.inputs()` (the source plan) by default. + let on = col("target.id").eq(col("source.id")); + let merge_op = MergeIntoOp { + on, + clauses: vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: None, + action: MergeIntoAction::Update(vec![( + "id".to_string(), + col("source.id"), + )]), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec!["id".to_string()], + values: vec![col("source.id")], + }, + }, + ], + }; + let plan = LogicalPlan::Dml(DmlStatement::new( + target_table_name, + target_source, + WriteOp::MergeInto(Box::new(merge_op)), + Arc::new(source_plan), + )); + + let analyzed = Analyzer::with_rules(vec![Arc::new(TypeCoercion::new())]) + .execute_and_check(plan, &ConfigOptions::default(), |_, _| {})?; + let LogicalPlan::Dml(dml) = analyzed else { + panic!("expected Dml"); + }; + let WriteOp::MergeInto(merge_op) = dml.op else { + panic!("expected MergeInto"); + }; + assert_eq!( + merge_op.on.to_string(), + "CAST(target.id AS Int64) = source.id" + ); + let MergeIntoAction::Update(assignments) = &merge_op.clauses[0].action else { + panic!("expected UPDATE"); + }; + assert_eq!(assignments[0].1.to_string(), "CAST(source.id AS UInt32)"); + let MergeIntoAction::Insert { values, .. } = &merge_op.clauses[1].action else { + panic!("expected INSERT"); + }; + assert_eq!(values[0].to_string(), "CAST(source.id AS UInt32)"); + Ok(()) + } + #[test] fn coerce_utf8view_output() -> Result<()> { // Plan A @@ -1772,6 +1982,31 @@ mod test { } } + #[derive(Debug, Hash, PartialEq, Eq)] + struct TestArrayElementUDF; + + impl ScalarUDFImpl for TestArrayElementUDF { + fn name(&self) -> &str { + "TestArrayElementUDF" + } + + fn signature(&self) -> &Signature { + static SIGNATURE: std::sync::LazyLock = + std::sync::LazyLock::new(|| { + Signature::array_and_index(Volatility::Immutable) + }); + &SIGNATURE + } + + fn return_type(&self, _args: &[DataType]) -> Result { + Ok(Utf8) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + Ok(ColumnarValue::Scalar(ScalarValue::from("a"))) + } + } + #[test] fn scalar_udf() -> Result<()> { let empty = empty(); @@ -2203,6 +2438,113 @@ mod test { Ok(()) } + #[test] + fn similar_to_for_type_coercion() -> Result<()> { + // similar to : utf8 similar to "abc" + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::new_utf8("abc"))); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(Utf8); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r#" + Projection: a SIMILAR TO Utf8("abc") + EmptyRelation: rows=0 + "# + )?; + + // NULL pattern is coerced to a typed NULL instead of panicking + // (https://github.com/apache/datafusion/issues/22886) + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::Null)); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(Utf8); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r" + Projection: a SIMILAR TO CAST(NULL AS Utf8) + EmptyRelation: rows=0 + " + )?; + + // Utf8View value and Utf8 pattern are coerced to Utf8View + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::new_utf8("abc"))); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(DataType::Utf8View); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r#" + Projection: a SIMILAR TO CAST(Utf8("abc") AS Utf8View) + EmptyRelation: rows=0 + "# + )?; + + // Utf8 value and Utf8View pattern are coerced to Utf8View + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::Utf8View(Some("abc".to_string())))); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(Utf8); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r#" + Projection: CAST(a AS Utf8View) SIMILAR TO Utf8View("abc") + EmptyRelation: rows=0 + "# + )?; + + // Dictionary values are coerced to the common regex operand type + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::new_utf8("abc"))); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(DataType::Dictionary( + Box::new(DataType::Int32), + Box::new(Utf8), + )); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r#" + Projection: CAST(a AS Utf8) SIMILAR TO Utf8("abc") + EmptyRelation: rows=0 + "# + )?; + + // incompatible types are a planning error, not a panic + let expr = Box::new(col("a")); + let pattern = Box::new(lit(ScalarValue::new_utf8("abc"))); + let similar_to_expr = + Expr::SimilarTo(Like::new(false, expr, pattern, None, false)); + let empty = empty_with_type(DataType::Int64); + let plan = + LogicalPlan::Projection(Projection::try_new(vec![similar_to_expr], empty)?); + assert_type_coercion_error( + plan, + "There isn't a common type to coerce Int64 and Utf8 in SIMILAR TO expression", + )?; + + Ok(()) + } + #[test] fn unknown_for_type_coercion() -> Result<()> { // unknown @@ -2669,6 +3011,33 @@ mod test { ) } + #[test] + fn array_element_preserves_parquet_list_field_name() -> Result<()> { + let list_type = DataType::List(Arc::new(Field::new( + "element", + DataType::Struct( + vec![ + Field::new("id", Utf8, true), + Field::new("prim", DataType::Boolean, true), + ] + .into(), + ), + true, + ))); + + let expr = ScalarUDF::from(TestArrayElementUDF).call(vec![col("a"), lit(1_i64)]); + let empty = empty_with_type(list_type); + let plan = LogicalPlan::Projection(Projection::try_new(vec![expr], empty)?); + + assert_analyzed_plan_eq!( + plan, + @r#" + Projection: TestArrayElementUDF(a, Int64(1)) + EmptyRelation: rows=0 + "# + ) + } + #[test] fn interval_plus_timestamp() -> Result<()> { // SELECT INTERVAL '1' YEAR + '2000-01-01T00:00:00'::timestamp; diff --git a/datafusion/optimizer/src/common_subexpr_eliminate.rs b/datafusion/optimizer/src/common_subexpr_eliminate.rs index 2775d62144c56..41d09db7c2bbe 100644 --- a/datafusion/optimizer/src/common_subexpr_eliminate.rs +++ b/datafusion/optimizer/src/common_subexpr_eliminate.rs @@ -826,6 +826,9 @@ fn extract_expressions(expr: &Expr, result: &mut Vec) { let col = Column::new(qualifier, field_name); result.push(Expr::Column(col)) } + result.push(Expr::Column(Column::from_name( + Aggregate::INTERNAL_GROUPING_ID, + ))); } else { let (qualifier, field_name) = expr.qualified_name(); let col = Column::new(qualifier, field_name); @@ -1106,6 +1109,27 @@ mod test { ) } + #[test] + fn common_aggregate_grouping_set_preserves_internal_id() -> Result<()> { + let plan = LogicalPlanBuilder::from(test_table_scan()?) + .aggregate( + vec![grouping_set(vec![vec![col("a")]])], + vec![avg(col("b")).alias("first"), avg(col("b")).alias("second")], + )? + .filter(col(Aggregate::INTERNAL_GROUPING_ID).eq(lit(0_u8)))? + .build()?; + + assert_optimized_plan_equal!( + plan, + @ r" + Filter: __grouping_id = UInt8(0) + Projection: test.a, __grouping_id, __common_expr_1 AS first, __common_expr_1 AS second + Aggregate: groupBy=[[GROUPING SETS ((test.a))]], aggr=[[avg(test.b) AS __common_expr_1]] + TableScan: test + " + ) + } + #[test] fn subexpr_in_same_order() -> Result<()> { let table_scan = test_table_scan()?; @@ -1288,20 +1312,31 @@ mod test { #[test] fn test_extract_expressions_from_grouping_set() -> Result<()> { - let mut result = Vec::with_capacity(3); + let mut result = Vec::with_capacity(4); let grouping = grouping_set(vec![vec![col("a"), col("b")], vec![col("c")]]); extract_expressions(&grouping, &mut result); - assert!(result.len() == 3); + assert_eq!( + result, + vec![ + col("a"), + col("b"), + col("c"), + col(Aggregate::INTERNAL_GROUPING_ID), + ] + ); Ok(()) } #[test] fn test_extract_expressions_from_grouping_set_with_identical_expr() -> Result<()> { - let mut result = Vec::with_capacity(2); + let mut result = Vec::with_capacity(3); let grouping = grouping_set(vec![vec![col("a"), col("b")], vec![col("a")]]); extract_expressions(&grouping, &mut result); - assert!(result.len() == 2); + assert_eq!( + result, + vec![col("a"), col("b"), col(Aggregate::INTERNAL_GROUPING_ID),] + ); Ok(()) } diff --git a/datafusion/optimizer/src/decorrelate.rs b/datafusion/optimizer/src/decorrelate.rs index 2a71205c64c8b..9490af0e59749 100644 --- a/datafusion/optimizer/src/decorrelate.rs +++ b/datafusion/optimizer/src/decorrelate.rs @@ -35,8 +35,8 @@ use datafusion_expr::utils::{ collect_subquery_cols, conjunction, find_join_exprs, split_conjunction, }; use datafusion_expr::{ - BinaryExpr, Cast, EmptyRelation, Expr, FetchType, LogicalPlan, LogicalPlanBuilder, - Operator, expr, lit, + BinaryExpr, Cast, EmptyRelation, Expr, ExprSchemable, FetchType, LogicalPlan, + LogicalPlanBuilder, Operator, expr, lit, }; /// This struct rewrite the sub query plan by pull up the correlated @@ -512,18 +512,12 @@ fn agg_exprs_evaluation_result_on_empty_batch( let result_expr = e .clone() .transform_up(|expr| { - let new_expr = match expr { - Expr::AggregateFunction(expr::AggregateFunction { func, .. }) => { - if func.name() == "count" { - Transformed::yes(Expr::Literal( - ScalarValue::Int64(Some(0)), - None, - )) - } else { - Transformed::yes(Expr::Literal(ScalarValue::Null, None)) - } - } - _ => Transformed::no(expr), + let new_expr = if let Expr::AggregateFunction(agg) = &expr { + let return_type = expr.get_type(schema.as_ref())?; + let default_value = agg.func.default_value(&return_type)?; + Transformed::yes(Expr::Literal(default_value, None)) + } else { + Transformed::no(expr) }; Ok(new_expr) }) diff --git a/datafusion/optimizer/src/eliminate_cross_join.rs b/datafusion/optimizer/src/eliminate_cross_join.rs index 8306d4b54c256..95b70da443d88 100644 --- a/datafusion/optimizer/src/eliminate_cross_join.rs +++ b/datafusion/optimizer/src/eliminate_cross_join.rs @@ -20,7 +20,7 @@ use crate::{OptimizerConfig, OptimizerRule}; use std::sync::Arc; use crate::join_key_set::JoinKeySet; -use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{NullEquality, Result}; use datafusion_expr::expr::{BinaryExpr, Expr}; use datafusion_expr::logical_plan::{ @@ -85,6 +85,17 @@ impl OptimizerRule for EliminateCrossJoin { plan: LogicalPlan, config: &dyn OptimizerConfig, ) -> Result> { + // Fast path: nothing to do if the plan contains no `Join` nodes. + // Without this guard the rule still falls through to + // `rewrite_children`, which walks the entire plan, processes + // uncorrelated subqueries, and rewrites every direct child via + // `map_children` (clone-on-write) — paid by every query in the + // logical optimizer pipeline. Same shape as the + // `plan_has_subqueries` fast-path landed in #22298. + if !plan_has_joins(&plan) { + return Ok(Transformed::no(plan)); + } + let plan_schema = Arc::clone(plan.schema()); let mut possible_join_keys = JoinKeySet::new(); let mut all_inputs: Vec = vec![]; @@ -207,6 +218,34 @@ impl OptimizerRule for EliminateCrossJoin { } } +/// Returns `true` if `plan` contains at least one [`LogicalPlan::Join`] +/// node, either directly in its tree *or* inside an embedded subquery +/// plan reachable through `Expr::ScalarSubquery` / `Expr::InSubquery` +/// / `Expr::Exists` / `Expr::SetComparison`. +/// +/// Used as a fast-path gate at the top of [`EliminateCrossJoin::rewrite`] +/// so that join-free plans skip the full recursive rewrite. Subquery +/// traversal matters because `rewrite_children` also dives into +/// uncorrelated subqueries via `map_uncorrelated_subqueries`; ignoring +/// them here would skip optimizing a `CROSS JOIN` that sits only inside +/// an `IN (SELECT ... FROM a, b)`-style predicate. +/// +/// `LogicalPlan::apply_with_subqueries` already implements the +/// "walk this node + every child + every subquery plan" traversal we +/// need, so the helper is a thin wrapper around it. +fn plan_has_joins(plan: &LogicalPlan) -> bool { + let mut found = false; + let _ = plan.apply_with_subqueries(|node| { + if matches!(node, LogicalPlan::Join(_)) { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }); + found +} + fn rewrite_children( optimizer: &impl OptimizerRule, plan: LogicalPlan, @@ -1418,4 +1457,102 @@ mod tests { Ok(()) } + + // ---------------- fast-path tests ---------------- + + /// `plan_has_joins` detects a `Join` at the root of the plan. + #[test] + fn plan_has_joins_detects_root_join() -> Result<()> { + let plan = LogicalPlanBuilder::from(test_table_scan_with_name("t1")?) + .cross_join(test_table_scan_with_name("t2")?)? + .build()?; + assert!(plan_has_joins(&plan)); + Ok(()) + } + + /// `plan_has_joins` detects a `Join` nested under other operators. + #[test] + fn plan_has_joins_detects_nested_join() -> Result<()> { + let plan = LogicalPlanBuilder::from(test_table_scan_with_name("t1")?) + .cross_join(test_table_scan_with_name("t2")?)? + .filter(col("t1.a").eq(col("t2.a")))? + .project(vec![col("t1.a")])? + .build()?; + assert!(plan_has_joins(&plan)); + Ok(()) + } + + /// Join-free plans return `false` so the fast-path in `rewrite` can + /// bail out before doing any recursion. + #[test] + fn plan_has_joins_returns_false_for_join_free_plan() -> Result<()> { + let plan = LogicalPlanBuilder::from(test_table_scan_with_name("t1")?) + .filter(col("a").gt(lit(0_i32)))? + .project(vec![col("a"), col("b")])? + .build()?; + assert!(!plan_has_joins(&plan)); + Ok(()) + } + + /// `plan_has_joins` walks into embedded subquery plans — e.g. an + /// outer `Filter` whose predicate is `IN (SELECT ... FROM a, b)` + /// where the inner plan contains a `CROSS JOIN`. Without this the + /// fast-path would silently skip optimizing joins-in-subqueries + /// because `LogicalPlan::apply` doesn't descend into subquery + /// plan trees. + #[test] + fn plan_has_joins_detects_join_inside_subquery() -> Result<()> { + use datafusion_expr::in_subquery; + + // Subquery plan that itself contains a join. + let subquery_plan = + LogicalPlanBuilder::from(test_table_scan_with_name("sub_t1")?) + .cross_join(test_table_scan_with_name("sub_t2")?)? + .project(vec![col("sub_t1.a")])? + .build()?; + + // Outer plan with NO direct Join — only the IN subquery reaches one. + let outer = LogicalPlanBuilder::from(test_table_scan_with_name("t1")?) + .filter(in_subquery(col("a"), Arc::new(subquery_plan)))? + .project(vec![col("a")])? + .build()?; + + assert!( + plan_has_joins(&outer), + "plan_has_joins must descend into subquery plans" + ); + Ok(()) + } + + /// `EliminateCrossJoin::rewrite` short-circuits on join-free plans: + /// no recursion into `rewrite_children`, no `Transformed::yes`, + /// the plan comes back identical. + #[test] + fn rewrite_short_circuits_when_plan_has_no_joins() -> Result<()> { + let plan = LogicalPlanBuilder::from(test_table_scan_with_name("t1")?) + .filter(col("a").gt(lit(0_i32)))? + .project(vec![col("a"), col("b")])? + .build()?; + + let starting_display = plan.display_indent_schema().to_string(); + let starting_schema = Arc::clone(plan.schema()); + + let rule = EliminateCrossJoin::new(); + let Transformed { + transformed, + data: optimized_plan, + .. + } = rule.rewrite(plan, &OptimizerContext::new())?; + + assert!( + !transformed, + "join-free plan should not be marked as transformed" + ); + assert_eq!(&starting_schema, optimized_plan.schema()); + assert_eq!( + starting_display, + optimized_plan.display_indent_schema().to_string() + ); + Ok(()) + } } diff --git a/datafusion/optimizer/src/eliminate_group_by_constant.rs b/datafusion/optimizer/src/eliminate_group_by_constant.rs index e21241ba7d993..f0efe96668dba 100644 --- a/datafusion/optimizer/src/eliminate_group_by_constant.rs +++ b/datafusion/optimizer/src/eliminate_group_by_constant.rs @@ -64,10 +64,14 @@ impl OptimizerRule for EliminateGroupByConstant { .group_expr .iter() .partition(|expr| is_redundant_group_expr(expr, &group_by_columns)); - - if redundant.is_empty() - || (required.is_empty() && aggregate.aggr_expr.is_empty()) - { + // Return now if no simplification can be done. We also bail out + // if applying the optimization would eliminate all of the + // grouping expressions (e.g., GROUP BY on only constant + // expressions): this would turn a grouped aggregate into an + // ungrouped aggregate, which changes query semantics (grouped + // aggregates produce an empty result set on an empty input, + // whereas ungrouped aggregates return a single row). + if redundant.is_empty() || required.is_empty() { return Ok(Transformed::no(LogicalPlan::Aggregate(aggregate))); } @@ -221,16 +225,15 @@ mod tests { } #[test] - fn test_eliminate_constant() -> Result<()> { + fn test_no_op_only_constant_with_aggregate() -> Result<()> { let scan = test_table_scan()?; let plan = LogicalPlanBuilder::from(scan) .aggregate(vec![lit("test"), lit(123u32)], vec![count(col("c"))])? .build()?; assert_optimized_plan_equal!(plan, @r#" - Projection: Utf8("test"), UInt32(123), count(test.c) - Aggregate: groupBy=[[]], aggr=[[count(test.c)]] - TableScan: test + Aggregate: groupBy=[[Utf8("test"), UInt32(123)]], aggr=[[count(test.c)]] + TableScan: test "#) } diff --git a/datafusion/optimizer/src/eliminate_join.rs b/datafusion/optimizer/src/eliminate_join.rs index 885910c1e4182..56aa8887065be 100644 --- a/datafusion/optimizer/src/eliminate_join.rs +++ b/datafusion/optimizer/src/eliminate_join.rs @@ -15,19 +15,150 @@ // specific language governing permissions and limitations // under the License. -//! [`EliminateJoin`] rewrites `INNER JOIN` with `true`/`null` -use crate::optimizer::ApplyOrder; +//! [`EliminateJoin`] rewrites joins to simpler forms to make them cheaper +//! to evaluate. We implement three distinct rewrites: +//! +//! * An inner join can be rewritten to an empty relation if the join condition +//! is trivially false. +//! +//! * An inner join `L ⋈ R` can be rewritten to a left semi join `L ⋉ R` +//! (`LeftSemi`), which keeps the rows of L that have a match in R and outputs +//! only L's columns. The rewrite to `L ⋉ R` is valid when both of the +//! following are true: +//! +//! 1. None of R's columns are referenced above the join. +//! 2. R does not observably multiply L's rows. This holds when either the +//! join's ancestors are duplicate-insensitive (e.g., DISTINCT) or we can use +//! functional dependencies to prove that each L row matches at most one R +//! row (R is provably unique on the join keys). +//! +//! * A left outer join `L ⟕ R` can be removed entirely, i.e. replaced by `L`, +//! under the same two conditions. Unlike an inner join, a left join +//! preserves every row of L whether or not it has a match in R, so when R's +//! columns are unused and R cannot multiply L's rows the join has no +//! observable effect at all. Such joins commonly appear in generated SQL +//! and in queries over views that join in lookup tables the query does not +//! read. A join filter does not prevent this rewrite: for a left join it +//! only decides whether a left row is matched or null-padded, and either +//! way the row is emitted. Symmetrically, a right outer join `L ⟖ R` can be +//! replaced by `R` when L's columns are unused and L cannot multiply R's +//! rows. +//! +//! # Overview +//! +//! `rewrite_subtree` walks the plan top-down, threading two pieces of context +//! down to each join: +//! +//! * `live` — which of the join's output columns are referenced above it. It is +//! propagated top-down: each node asks its children only for the columns it +//! needs from them, so a projection or aggregate asks for just the columns its +//! expressions reference, dropping the rest (the narrowing); a join splits the +//! set across its two inputs. +//! * `duplicate_insensitive` — whether emitting each row once instead of many +//! times will not change the output. A duplicate-collapsing node (e.g., +//! DISTINCT, GROUP BY with no aggregate functions, or the existence side of a +//! semi/anti/mark join) sets it `true` for its subtree, and it propagates +//! downward until a node that makes the row count observable again (a `LIMIT`, +//! a top-N sort, ...) clears it. It is therefore fixed by the nearest such +//! node, not by the whole ancestor chain: a collapsing node shields its subtree, +//! so a duplicate-sensitive node further above does not matter. +//! +//! At each join, `rewritten_join_type` combines this context with the side's +//! functional dependencies to choose `Inner`, `LeftSemi`, or `RightSemi`, or +//! to eliminate the join entirely in favor of its preserved input. Most +//! node types just forward the context to their single child via +//! `rewrite_single_input`; nodes that alter column requirements or +//! duplicate-sensitivity (projection, aggregate, sort, ...) adjust it first. +use crate::utils::for_each_referenced_index; use crate::{OptimizerConfig, OptimizerRule}; -use datafusion_common::tree_node::Transformed; -use datafusion_common::{Result, ScalarValue}; -use datafusion_expr::JoinType::Inner; +use datafusion_common::tree_node::{Transformed, TreeNode}; +use datafusion_common::{ + DFSchema, Dependency, HashSet, NullEquality, Result, ScalarValue, +}; use datafusion_expr::{ - Expr, - logical_plan::{EmptyRelation, LogicalPlan}, + Expr, JoinType, + logical_plan::{ + Aggregate, Distinct, DistinctOn, EmptyRelation, Filter, Join, Limit, LogicalPlan, + Partitioning, Projection, Repartition, Sort, SubqueryAlias, + }, }; +use std::sync::Arc; + +/// The columns that are "live" at a plan node, i.e., which of its output +/// columns are referenced by an ancestor node. Represented as a set of column +/// indices, relative to the node's schema. +/// +/// See the module-level docs for how this set is threaded down the plan and +/// narrowed or split at each node. +#[derive(Debug, Default, Clone)] +struct LiveColumns(HashSet); + +impl LiveColumns { + fn new() -> Self { + Self(HashSet::new()) + } + + /// Every column of `schema` is live. + fn all(schema: &DFSchema) -> Self { + Self((0..schema.fields().len()).collect()) + } + + /// The columns of `schema` referenced by any of `exprs`. + fn try_new<'a>( + exprs: impl IntoIterator, + schema: &DFSchema, + ) -> Result { + let mut live = Self::new(); + live.extend_from(exprs, schema)?; + Ok(live) + } -/// Eliminates joins when join condition is false. -/// Replaces joins when inner join condition is true with a cross join. + /// Inserts the index, within `schema`, of every column referenced by any of + /// `exprs`, including columns reached through correlated subquery outer + /// references. + fn extend_from<'a>( + &mut self, + exprs: impl IntoIterator, + schema: &DFSchema, + ) -> Result<()> { + for expr in exprs { + for_each_referenced_index(expr, schema, |idx| { + self.0.insert(idx); + })?; + } + Ok(()) + } + + fn insert(&mut self, idx: usize) { + self.0.insert(idx); + } + + fn is_empty(&self) -> bool { + self.0.is_empty() + } + + /// Splits live columns spanning a join's combined output (the left input's + /// columns first, then the right input's) into the per-side sets, rebasing + /// the right side's indices to start at zero. `left_len` is the number of + /// columns contributed by the left input. + fn split_at(&self, left_len: usize) -> (Self, Self) { + let mut left = Self::new(); + let mut right = Self::new(); + for &idx in &self.0 { + if idx < left_len { + left.insert(idx); + } else { + right.insert(idx - left_len); + } + } + (left, right) + } +} + +/// Rewrites an inner join to a semi join when one input only filters the +/// other, removes an outer join whose non-preserved side is unused and cannot +/// multiply the preserved side's rows, and replaces an always-false inner join +/// with an empty relation. #[derive(Default, Debug)] pub struct EliminateJoin; @@ -42,44 +173,489 @@ impl OptimizerRule for EliminateJoin { "eliminate_join" } - fn apply_order(&self) -> Option { - Some(ApplyOrder::TopDown) - } - fn rewrite( &self, plan: LogicalPlan, _config: &dyn OptimizerConfig, ) -> Result> { - match plan { - LogicalPlan::Join(join) if join.join_type == Inner && join.on.is_empty() => { - match join.filter { - Some(Expr::Literal(ScalarValue::Boolean(Some(false)), _)) => Ok( - Transformed::yes(LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: false, - schema: join.schema, - })), - ), - _ => Ok(Transformed::no(LogicalPlan::Join(join))), + let live = LiveColumns::all(plan.schema()); + rewrite_subtree(plan, live, false) + } +} + +/// Rewrites `plan` and everything below it, including joins nested inside +/// subquery expressions. +/// +/// [`rewrite_node`] handles the node itself and recurses into its plan +/// children; this wrapper additionally descends into the node's own subquery +/// expressions. Each subquery is seeded as a fresh root, since its columns are +/// independent of the enclosing plan's `live` set. +fn rewrite_subtree( + plan: LogicalPlan, + live: LiveColumns, + duplicate_insensitive: bool, +) -> Result> { + rewrite_node(plan, live, duplicate_insensitive)?.transform_data(|plan| { + plan.map_subqueries(|subquery| { + let live = LiveColumns::all(subquery.schema()); + rewrite_subtree(subquery, live, false) + }) + }) +} + +fn rewrite_node( + plan: LogicalPlan, + live: LiveColumns, + duplicate_insensitive: bool, +) -> Result> { + match plan { + // The only arm that rewrites a join; the rest just thread context down to one. + LogicalPlan::Join(join) => rewrite_join(join, &live, duplicate_insensitive), + LogicalPlan::Projection(Projection { + expr, + input, + schema, + .. + }) => { + // Narrows `live` to the columns the projection's expressions reference. + let child_live = LiveColumns::try_new(&expr, input.schema())?; + rewrite_single_input(input, child_live, duplicate_insensitive, |input| { + Ok(LogicalPlan::Projection(Projection::try_new_with_schema( + expr, input, schema, + )?)) + }) + } + LogicalPlan::Filter(Filter { + predicate, input, .. + }) => { + // Adds the predicate's columns to `live` (a side used only by the filter stays live). + let mut child_live = live; + child_live.extend_from([&predicate], input.schema())?; + rewrite_single_input(input, child_live, duplicate_insensitive, |input| { + Ok(LogicalPlan::Filter(Filter::new(predicate, input))) + }) + } + LogicalPlan::Aggregate(Aggregate { + input, + group_expr, + aggr_expr, + schema, + .. + }) => { + // Narrows `live` to the grouping and aggregate expressions' columns. + let child_live = LiveColumns::try_new( + group_expr.iter().chain(&aggr_expr), + input.schema(), + )?; + + // A grouping aggregate with no aggregate functions (`GROUP BY` with + // an empty `aggr_expr`) only observes which group-key values exist, + // not how many rows produced them, so its input is duplicate- + // insensitive. + let child_duplicate_insensitive = + !group_expr.is_empty() && aggr_expr.is_empty(); + + rewrite_single_input( + input, + child_live, + child_duplicate_insensitive, + |input| { + Ok(LogicalPlan::Aggregate(Aggregate::try_new_with_schema( + input, group_expr, aggr_expr, schema, + )?)) + }, + ) + } + LogicalPlan::Distinct(Distinct::All(input)) => { + // `SELECT DISTINCT *` is equivalent to a no-aggregate `GROUP BY` + // over every input column, so the input is duplicate-insensitive, + // but every column is part of the dedup key. + let child_live = LiveColumns::all(input.schema()); + rewrite_single_input(input, child_live, true, |input| { + Ok(LogicalPlan::Distinct(Distinct::All(input))) + }) + } + LogicalPlan::Distinct(Distinct::On(DistinctOn { + on_expr, + select_expr, + sort_expr, + input, + schema, + })) => { + // `DISTINCT ON (on) select [ORDER BY sort]` is a no-aggregate + // `GROUP BY` on the columns it reads, so its input is duplicate- + // insensitive; the live columns are exactly those of the + // ON/SELECT/ORDER BY expressions. + let mut child_live = + LiveColumns::try_new(on_expr.iter().chain(&select_expr), input.schema())?; + if let Some(sort_expr) = &sort_expr { + child_live + .extend_from(sort_expr.iter().map(|s| &s.expr), input.schema())?; + } + + rewrite_single_input(input, child_live, true, |input| { + Ok(LogicalPlan::Distinct(Distinct::On(DistinctOn { + on_expr, + select_expr, + sort_expr, + input, + schema, + }))) + }) + } + LogicalPlan::Sort(Sort { expr, input, fetch }) => { + // Adds the sort-key columns to `live`. + let mut child_live = live; + child_live.extend_from(expr.iter().map(|s| &s.expr), input.schema())?; + + // A `fetch` (top-N) makes the row count observable, so duplicate- + // insensitivity does not survive past it. + let child_duplicate_insensitive = duplicate_insensitive && fetch.is_none(); + rewrite_single_input( + input, + child_live, + child_duplicate_insensitive, + |input| Ok(LogicalPlan::Sort(Sort { expr, input, fetch })), + ) + } + LogicalPlan::Limit(Limit { skip, fetch, input }) => { + // LIMIT makes the row count observable, so it clears duplicate-insensitivity. + rewrite_single_input(input, live, false, |input| { + Ok(LogicalPlan::Limit(Limit { skip, fetch, input })) + }) + } + LogicalPlan::SubqueryAlias(SubqueryAlias { input, alias, .. }) => { + // Re-aliases columns 1:1, so `live` and duplicate-sensitivity pass through unchanged. + rewrite_single_input(input, live, duplicate_insensitive, |input| { + Ok(LogicalPlan::SubqueryAlias(SubqueryAlias::try_new( + input, alias, + )?)) + }) + } + LogicalPlan::Repartition(Repartition { + input, + partitioning_scheme, + }) => { + // Adds any partitioning-key columns to `live`; duplicate-sensitivity is unchanged. + let mut child_live = live; + match &partitioning_scheme { + Partitioning::Hash(exprs, _) | Partitioning::DistributeBy(exprs) => { + child_live.extend_from(exprs, input.schema())?; + } + Partitioning::Range(range) => { + child_live.extend_from( + range.ordering().iter().map(|sort_expr| &sort_expr.expr), + input.schema(), + )?; } + Partitioning::RoundRobinBatch(_) => {} } - _ => Ok(Transformed::no(plan)), + rewrite_single_input(input, child_live, duplicate_insensitive, |input| { + Ok(LogicalPlan::Repartition(Repartition { + input, + partitioning_scheme, + })) + }) } + // Conservatively treat any other plan node as a fresh root, since we are + // not sure of its semantics with respect to duplicates or live columns. + _ => plan.map_children(|child| { + let live = LiveColumns::all(child.schema()); + rewrite_subtree(child, live, false) + }), + } +} + +/// Recurses into a single-input node's child, threading `child_live` and +/// `duplicate_insensitive` down, then rebuilds the node from the (possibly +/// rewritten) child via `rebuild`. The child's `Transformed` flag is preserved, +/// so the node is reported as changed exactly when its child changed. +fn rewrite_single_input( + input: Arc, + child_live: LiveColumns, + duplicate_insensitive: bool, + rebuild: F, +) -> Result> +where + F: FnOnce(Arc) -> Result, +{ + rewrite_subtree( + Arc::unwrap_or_clone(input), + child_live, + duplicate_insensitive, + )? + .map_data(|input| rebuild(Arc::new(input))) +} + +fn rewrite_join( + join: Join, + live: &LiveColumns, + duplicate_insensitive: bool, +) -> Result> { + if join.join_type == JoinType::Inner + && join.on.is_empty() + && matches!( + join.filter.as_ref(), + Some(Expr::Literal(ScalarValue::Boolean(Some(false)), _)) + ) + { + return Ok(Transformed::yes(LogicalPlan::EmptyRelation( + EmptyRelation { + produce_one_row: false, + schema: join.schema, + }, + ))); } - fn supports_rewrite(&self) -> bool { - true + let (visible_left, visible_right) = split_join_output_columns(&join, live); + + let rewritten_join_type = match rewritten_join_type( + &join, + &visible_left, + &visible_right, + duplicate_insensitive, + ) { + JoinRewrite::ReplaceWithLeft => { + let left = rewrite_subtree( + Arc::unwrap_or_clone(join.left), + visible_left, + duplicate_insensitive, + )?; + return Ok(Transformed::yes(left.data)); + } + JoinRewrite::ReplaceWithRight => { + let right = rewrite_subtree( + Arc::unwrap_or_clone(join.right), + visible_right, + duplicate_insensitive, + )?; + return Ok(Transformed::yes(right.data)); + } + JoinRewrite::Join(join_type) => join_type, + }; + + let (mut left_live, mut right_live) = match rewritten_join_type { + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { + (visible_left, LiveColumns::new()) + } + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { + (LiveColumns::new(), visible_right) + } + _ => (visible_left, visible_right), + }; + + add_join_condition_columns(&join, &mut left_live, &mut right_live)?; + + let (left_dup_insensitive, right_dup_insensitive) = + child_duplicate_insensitivity(rewritten_join_type, duplicate_insensitive); + + let left = rewrite_subtree( + Arc::unwrap_or_clone(join.left), + left_live, + left_dup_insensitive, + )?; + let right = rewrite_subtree( + Arc::unwrap_or_clone(join.right), + right_live, + right_dup_insensitive, + )?; + + let changed = + left.transformed || right.transformed || rewritten_join_type != join.join_type; + let left = Arc::new(left.data); + let right = Arc::new(right.data); + + if changed { + // The join type or an input changed, so the output schema may have + // narrowed; recompute it via `try_new`. + Ok(Transformed::yes(LogicalPlan::Join(Join::try_new( + left, + right, + join.on, + join.filter, + rewritten_join_type, + join.join_constraint, + join.null_equality, + join.null_aware, + )?))) + } else { + // Nothing changed; reassemble the join reusing its existing schema rather + // than recomputing it. + Ok(Transformed::no(LogicalPlan::Join(Join { + left, + right, + on: join.on, + filter: join.filter, + join_type: join.join_type, + join_constraint: join.join_constraint, + schema: join.schema, + null_equality: join.null_equality, + null_aware: join.null_aware, + }))) } } +/// Returns which join inputs can safely ignore duplicate rows from their own +/// descendants. For semi/anti/mark joins, duplicates from the existence side do +/// not change the result even when the parent itself is duplicate-sensitive. +fn child_duplicate_insensitivity( + join_type: JoinType, + duplicate_insensitive: bool, +) -> (bool, bool) { + match join_type { + JoinType::Inner => (duplicate_insensitive, duplicate_insensitive), + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { + (duplicate_insensitive, true) + } + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { + (true, duplicate_insensitive) + } + JoinType::Left | JoinType::Right | JoinType::Full => (false, false), + } +} + +/// The rewrite chosen for a join by [`rewritten_join_type`]. +enum JoinRewrite { + /// Keep the join, with this (possibly rewritten) join type. + Join(JoinType), + /// The join has no observable effect; replace it with its left input. + ReplaceWithLeft, + /// The join has no observable effect; replace it with its right input. + ReplaceWithRight, +} + +/// Chooses a cheaper form for a join: removes an outer join whose non-preserved +/// side is redundant, or rewrites an inner join to a semi join when the +/// removed side has no parent-visible columns and either the parent ignores +/// duplicate output rows or the removed side is unique on the join keys. +fn rewritten_join_type( + join: &Join, + visible_left: &LiveColumns, + visible_right: &LiveColumns, + duplicate_insensitive: bool, +) -> JoinRewrite { + // A side is redundant when nothing above the join references its columns + // and it cannot multiply the other side's rows (the ancestors are + // duplicate-insensitive, or the side is unique on the join keys). + let can_remove_right = visible_right.is_empty() + && (duplicate_insensitive + || side_unique_on_join( + join.right.schema(), + join.on.iter().map(|(_, right)| right), + join.null_equality, + )); + + // A LEFT JOIN preserves every left row, so with a redundant right side the + // join has no observable effect and can be replaced by its left input. A + // join filter cannot prevent this: it only decides whether a left row is + // matched or null-padded, and either way the row is emitted. + if join.join_type == JoinType::Left && can_remove_right { + return JoinRewrite::ReplaceWithLeft; + } + let can_remove_left = visible_left.is_empty() + && (duplicate_insensitive + || side_unique_on_join( + join.left.schema(), + join.on.iter().map(|(left, _)| left), + join.null_equality, + )); + + // Symmetrical rule for RIGHT JOIN removal (same explanation as above for the left-join case) + if join.join_type == JoinType::Right && can_remove_left { + return JoinRewrite::ReplaceWithRight; + } + + if join.join_type != JoinType::Inner || join.on.is_empty() { + return JoinRewrite::Join(join.join_type); + } + + if can_remove_right { + return JoinRewrite::Join(JoinType::LeftSemi); + } + if can_remove_left { + return JoinRewrite::Join(JoinType::RightSemi); + } + + JoinRewrite::Join(JoinType::Inner) +} + +fn add_join_condition_columns( + join: &Join, + left_live: &mut LiveColumns, + right_live: &mut LiveColumns, +) -> Result<()> { + left_live.extend_from(join.on.iter().map(|(l, _)| l), join.left.schema())?; + right_live.extend_from(join.on.iter().map(|(_, r)| r), join.right.schema())?; + + if let Some(filter) = &join.filter { + left_live.extend_from([filter], join.left.schema())?; + right_live.extend_from([filter], join.right.schema())?; + } + + Ok(()) +} + +fn split_join_output_columns( + join: &Join, + live: &LiveColumns, +) -> (LiveColumns, LiveColumns) { + let left_len = join.left.schema().fields().len(); + match join.join_type { + JoinType::Inner | JoinType::Left | JoinType::Right | JoinType::Full => { + live.split_at(left_len) + } + // A semi/anti/mark join outputs only the surviving side's columns, with + // the same index space, so `live` passes straight through to that side. + JoinType::LeftSemi | JoinType::LeftAnti | JoinType::LeftMark => { + (live.clone(), LiveColumns::new()) + } + JoinType::RightSemi | JoinType::RightAnti | JoinType::RightMark => { + (LiveColumns::new(), live.clone()) + } + } +} + +fn side_unique_on_join<'a>( + schema: &DFSchema, + join_exprs: impl Iterator, + null_equality: NullEquality, +) -> bool { + let join_key_indices = join_exprs + .filter_map(|expr| match expr { + Expr::Alias(alias) => alias.expr.as_ref().try_as_col(), + _ => expr.try_as_col(), + }) + .filter_map(|column| schema.maybe_index_of_column(column)) + .collect::>(); + + schema.functional_dependencies().iter().any(|dependency| { + dependency.mode == Dependency::Single + && (!dependency.nullable || null_equality == NullEquality::NullEqualsNothing) + && dependency + .source_indices + .iter() + .all(|idx| join_key_indices.contains(idx)) + }) +} + #[cfg(test)] mod tests { use crate::OptimizerContext; use crate::assert_optimized_plan_eq_snapshot; use crate::eliminate_join::EliminateJoin; - use datafusion_common::Result; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{ + Constraint, Constraints, NullEquality, Result, ScalarValue, SplitPoint, + }; use datafusion_expr::JoinType::Inner; - use datafusion_expr::{lit, logical_plan::builder::LogicalPlanBuilder}; + use datafusion_expr::{ + Expr, JoinType, Partitioning, RangePartitioning, col, exists, lit, + logical_plan::builder::{ + LogicalPlanBuilder, table_scan, table_source_with_constraints, + }, + out_ref_col, + }; + use datafusion_functions_aggregate::expr_fn::count; use std::sync::Arc; macro_rules! assert_optimized_plan_equal { @@ -110,4 +686,591 @@ mod tests { assert_optimized_plan_equal!(plan, @"EmptyRelation: rows=0") } + + #[test] + fn inner_to_left_semi_when_removed_side_is_unique() -> Result<()> { + let plan = left_join_right_with_constraints(primary_key_on_id())? + .project(vec![col("l.x")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.x + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn inner_to_left_semi_when_removed_side_is_unique_with_join_filter() -> Result<()> { + let right = scan("r", &test_schema(), primary_key_on_id())?; + let plan = + LogicalPlanBuilder::from(scan("l", &test_schema(), Constraints::default())?) + .join( + right, + Inner, + (vec!["l.id"], vec!["r.id"]), + Some(col("r.y").gt(col("l.x"))), + )? + .project(vec![col("l.x")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.x + LeftSemi Join: l.id = r.id Filter: r.y > l.x + TableScan: l + TableScan: r + ") + } + + #[test] + fn inner_to_right_semi_when_removed_side_is_unique() -> Result<()> { + let plan = left_with_constraints_join_right(primary_key_on_id())? + .project(vec![col("r.y")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: r.y + RightSemi Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn inner_to_left_semi_for_duplicate_insensitive_parent() -> Result<()> { + let plan = left_join_right()? + .aggregate(vec![col("l.x")], Vec::::new())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[l.x]], aggr=[[]] + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn aggregate_with_aggregates_is_not_duplicate_insensitive() -> Result<()> { + // A `GROUP BY` *with* aggregate functions observes how many rows fall in + // each group, so its input is not duplicate-insensitive. With a non-unique + // right side the join must stay an inner join: collapsing it to a semi + // join would drop matching duplicates and undercount `count(l.id)`. + let plan = left_join_right()? + .aggregate(vec![col("l.x")], vec![count(col("l.id"))])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[l.x]], aggr=[[count(l.id)]] + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn duplicate_insensitive_context_propagates_through_join_tree() -> Result<()> { + let left = scan("l", &test_schema(), Constraints::default())?; + let middle = scan("m", &test_schema(), Constraints::default())?; + let right = scan("r", &test_schema(), Constraints::default())?; + + let left_join_middle = LogicalPlanBuilder::from(left) + .join(middle, Inner, (vec!["l.id"], vec!["m.id"]), None)? + .build()?; + + let plan = LogicalPlanBuilder::from(left_join_middle) + .join(right, Inner, (vec!["l.id"], vec!["r.id"]), None)? + .aggregate(vec![col("l.x")], Vec::::new())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[l.x]], aggr=[[]] + LeftSemi Join: l.id = r.id + LeftSemi Join: l.id = m.id + TableScan: l + TableScan: m + TableScan: r + ") + } + + #[test] + fn projection_does_not_rewrite_without_uniqueness() -> Result<()> { + let plan = left_join_right()?.project(vec![col("l.x")])?.build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.x + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn required_filter_column_prevents_duplicate_insensitive_rewrite() -> Result<()> { + let plan = left_join_right()? + .filter(col("r.y").gt(lit(10_i32)))? + .aggregate(vec![col("l.x")], Vec::::new())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[l.x]], aggr=[[]] + Filter: r.y > Int32(10) + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn distinct_star_keeps_unreferenced_side() -> Result<()> { + // `SELECT DISTINCT *` deduplicates on every join-output column, including + // the right side's. With a non-unique right side the inner join can + // multiply left rows into distinct `(l, r)` combinations, so the join + // must not be rewritten to a semi join (which would drop the right + // columns from the DISTINCT key and undercount the result). This holds + // even when the right side is unique on the join keys: its columns are + // part of the DISTINCT key regardless. + let plan = left_join_right()? + .distinct()? + .project(vec![col("l.x")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.x + Distinct: + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn distinct_drops_unreferenced_side_when_projected() -> Result<()> { + // `SELECT DISTINCT l.x` projects the right side away below the DISTINCT, + // leaving it outside the dedup key. Like a no-aggregate `GROUP BY l.x`, + // the DISTINCT makes the input duplicate-insensitive, so the inner join + // collapses to a semi join even though the right side is not unique. + let plan = left_join_right()? + .project(vec![col("l.x")])? + .distinct()? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Distinct: + Projection: l.x + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn correlated_subquery_outer_ref_prevents_rewrite() -> Result<()> { + // The aggregate makes the parent duplicate-insensitive, so absent any + // other use of the right side the join would collapse to a semi join. + // But the `EXISTS` subquery correlates on `r.y`, so the right side is + // still needed and the join must stay an inner join. Otherwise the + // semi join would drop `r`, orphaning the correlated `r.y` reference. + let subquery = + LogicalPlanBuilder::from(scan("s", &test_schema(), Constraints::default())?) + .filter(col("s.id").eq(out_ref_col(DataType::Int32, "r.y")))? + .project(vec![lit(1)])? + .build()?; + + let plan = left_join_right()? + .filter(exists(Arc::new(subquery)))? + .aggregate(vec![col("l.x")], Vec::::new())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[l.x]], aggr=[[]] + Filter: EXISTS () + Subquery: + Projection: Int32(1) + Filter: s.id = outer_ref(r.y) + TableScan: s + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn inner_to_semi_inside_uncorrelated_subquery() -> Result<()> { + // A join nested inside a (not-yet-decorrelated) subquery is still + // rewritten, because `rewrite_subtree` descends into subquery plans + // itself via `map_subqueries`. Here the subquery's projection keeps + // only `l.x` and the removed side `r` is unique (PK), so the inner join + // collapses to a semi join. + let subquery = left_join_right_with_constraints(primary_key_on_id())? + .project(vec![col("l.x")])? + .build()?; + + let plan = LogicalPlanBuilder::from(scan( + "outer", + &test_schema(), + Constraints::default(), + )?) + .filter(exists(Arc::new(subquery)))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: EXISTS () + Subquery: + Projection: l.x + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + TableScan: outer + ") + } + + #[test] + fn inner_to_semi_inside_correlated_subquery() -> Result<()> { + // `map_subqueries` descends into correlated subqueries too, not just + // uncorrelated ones, so a join inside one is still rewritten. The + // subquery correlates on `outer.id` (via the filter), but that reference + // and the projection touch only `l`; `r` is unique (PK) and unreferenced, + // so the inner join inside the subquery collapses to a semi join. + let subquery = left_join_right_with_constraints(primary_key_on_id())? + .filter(col("l.x").eq(out_ref_col(DataType::Int32, "outer.id")))? + .project(vec![col("l.x")])? + .build()?; + + let plan = LogicalPlanBuilder::from(scan( + "outer", + &test_schema(), + Constraints::default(), + )?) + .filter(exists(Arc::new(subquery)))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: EXISTS () + Subquery: + Projection: l.x + Filter: l.x = outer_ref(outer.id) + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + TableScan: outer + ") + } + + #[test] + fn nullable_unique_rewrites_under_null_equals_nothing() -> Result<()> { + // A `UNIQUE` (rather than `PRIMARY KEY`) constraint marks the key as + // nullable. Under the default `NullEqualsNothing` join semantics a null + // key matches nothing, so a unique side still yields at most one match + // per left row and the inner join can become a semi join. + let left = scan("l", &test_schema(), Constraints::default())?; + let right = scan("r", &test_schema(), unique_on_x())?; + let plan = LogicalPlanBuilder::from(left) + .join(right, Inner, (vec!["l.x"], vec!["r.x"]), None)? + .project(vec![col("l.id")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.id + LeftSemi Join: l.x = r.x + TableScan: l + TableScan: r + ") + } + + #[test] + fn nullable_unique_does_not_rewrite_under_null_equals_null() -> Result<()> { + // With `NullEqualsNull` semantics two null keys compare equal, so a + // nullable `UNIQUE` key no longer guarantees at most one match per left + // row: several null-keyed right rows could match a null-keyed left row. + // Uniqueness on the join keys is therefore not established and the inner + // join must be preserved. + let left = scan("l", &test_schema(), Constraints::default())?; + let right = scan("r", &test_schema(), unique_on_x())?; + let plan = LogicalPlanBuilder::from(left) + .join_detailed( + right, + Inner, + (vec!["l.x"], vec!["r.x"]), + None, + NullEquality::NullEqualsNull, + )? + .project(vec![col("l.id")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.id + Inner Join: l.x = r.x + TableScan: l + TableScan: r + ") + } + + #[test] + fn composite_unique_rewrites_when_join_covers_all_key_columns() -> Result<()> { + // The removed side is unique on the composite key `(id, x)`. The join + // equates both key columns, so each left row matches at most one right + // row and the inner join can become a semi join. + let left = scan("l", &test_schema(), Constraints::default())?; + let right = scan("r", &test_schema(), composite_primary_key_on_id_x())?; + let plan = LogicalPlanBuilder::from(left) + .join( + right, + Inner, + (vec!["l.id", "l.x"], vec!["r.id", "r.x"]), + None, + )? + .project(vec![col("l.y")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.y + LeftSemi Join: l.id = r.id, l.x = r.x + TableScan: l + TableScan: r + ") + } + + #[test] + fn composite_unique_does_not_rewrite_when_join_misses_a_key_column() -> Result<()> { + // The removed side is unique only on the *composite* key `(id, x)`. The + // join equates `id` but not `x`, so a left row may match many right rows + // (those sharing its `id` but differing in `x`). Uniqueness on the join + // keys is not established, so the inner join must be preserved. This + // guards the requirement that the join cover *every* column of the + // unique key, not just some. + let left = scan("l", &test_schema(), Constraints::default())?; + let right = scan("r", &test_schema(), composite_primary_key_on_id_x())?; + let plan = LogicalPlanBuilder::from(left) + .join(right, Inner, (vec!["l.id"], vec!["r.id"]), None)? + .project(vec![col("l.y")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.y + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn top_n_sort_blocks_duplicate_insensitive_rewrite() -> Result<()> { + // A top-N `Sort` (one with a `fetch`) makes the row count observable, so + // the duplicate-insensitivity established by the `GROUP BY` does not survive + // past it. With a non-unique right side the join must stay an inner join: a + // semi join could drop matching duplicates and change which rows fall within + // the top N. + let plan = left_join_right()? + .sort_with_limit(vec![col("l.x").sort(true, false)], Some(5))? + .aggregate(vec![col("l.x")], Vec::::new())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[l.x]], aggr=[[]] + Sort: l.x ASC NULLS LAST, fetch=5 + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn sort_without_fetch_preserves_duplicate_insensitive_rewrite() -> Result<()> { + // A `Sort` without a `fetch` does not make the row count observable, so it + // forwards the parent's duplicate-insensitivity to the join unchanged + // (sorting before or after duplicate removal is equivalent). The non-unique + // right side is unreferenced, so the inner join collapses to a semi join. + let plan = left_join_right()? + .sort(vec![col("l.x").sort(true, false)])? + .aggregate(vec![col("l.x")], Vec::::new())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[l.x]], aggr=[[]] + Sort: l.x ASC NULLS LAST + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn limit_blocks_duplicate_insensitive_rewrite() -> Result<()> { + // `LIMIT` makes the row count observable, clearing the duplicate- + // insensitivity established by the `GROUP BY`. With a non-unique right side + // the join must stay an inner join, since a semi join could drop matching + // duplicates and change which rows the limit returns. + let plan = left_join_right()? + .limit(0, Some(5))? + .aggregate(vec![col("l.x")], Vec::::new())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Aggregate: groupBy=[[l.x]], aggr=[[]] + Limit: skip=0, fetch=5 + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn repartition_hash_key_keeps_removed_side_live() -> Result<()> { + // The projection keeps only `l.x`, and the right side is unique (PK), so + // absent any other use of `r` the inner join would collapse to a semi join. + // But the `Repartition` hashes on `r.y`, which keeps the right side live, so + // the join must stay an inner join to preserve `r.y` for the partitioning. + let plan = left_join_right_with_constraints(primary_key_on_id())? + .repartition(Partitioning::Hash(vec![col("r.y")], 4))? + .project(vec![col("l.x")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.x + Repartition: Hash(r.y) partition_count=4 + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn repartition_range_key_keeps_removed_side_live() -> Result<()> { + // The projection keeps only `l.x`, and the right side is unique (PK), so + // absent any other use of `r` the inner join would collapse to a semi join. + // But the `Repartition` ranges on `r.y`, which keeps the right side live, so + // the join must stay an inner join to preserve `r.y` for the partitioning. + let plan = left_join_right_with_constraints(primary_key_on_id())? + .repartition(Partitioning::Range(RangePartitioning::try_new( + vec![col("r.y").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?))? + .project(vec![col("l.x")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.x + Repartition: Range([r.y ASC NULLS FIRST], [(10)], 2) + Inner Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn distinct_on_enables_semi_join_rewrite() -> Result<()> { + // `DISTINCT ON (l.x)` is a no-aggregate `GROUP BY` on the columns it reads, + // so it makes its input duplicate-insensitive. The non-unique right side is + // unreferenced, so the inner join collapses to a semi join. + let plan = left_join_right()? + .distinct_on(vec![col("l.x")], vec![col("l.x")], None)? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + DistinctOn: on_expr=[[l.x]], select_expr=[[l.x]], sort_expr=[[]] + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + #[test] + fn existing_semi_join_passes_through_unchanged() -> Result<()> { + // A join that is already a semi join is threaded through unchanged: the rule + // only rewrites inner joins. This exercises the context-propagation paths for + // a non-inner join type, whose existence side contributes no live columns. + let left = scan("l", &test_schema(), Constraints::default())?; + let right = scan("r", &test_schema(), Constraints::default())?; + let plan = LogicalPlanBuilder::from(left) + .join( + right, + JoinType::LeftSemi, + (vec!["l.id"], vec!["r.id"]), + None, + )? + .project(vec![col("l.x")])? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Projection: l.x + LeftSemi Join: l.id = r.id + TableScan: l + TableScan: r + ") + } + + fn left_join_right() -> Result { + left_join_right_with_constraints(Constraints::default()) + } + + fn left_join_right_with_constraints( + right_constraints: Constraints, + ) -> Result { + let left = scan("l", &test_schema(), Constraints::default())?; + let right = scan("r", &test_schema(), right_constraints)?; + + LogicalPlanBuilder::from(left).join( + right, + Inner, + (vec!["l.id"], vec!["r.id"]), + None, + ) + } + + fn left_with_constraints_join_right( + left_constraints: Constraints, + ) -> Result { + let left = scan("l", &test_schema(), left_constraints)?; + let right = scan("r", &test_schema(), Constraints::default())?; + + LogicalPlanBuilder::from(left).join( + right, + Inner, + (vec!["l.id"], vec!["r.id"]), + None, + ) + } + + fn scan( + name: &str, + schema: &Schema, + constraints: Constraints, + ) -> Result { + if constraints.is_empty() { + table_scan(Some(name), schema, None)?.build() + } else { + LogicalPlanBuilder::scan( + name, + table_source_with_constraints(schema, constraints), + None, + )? + .build() + } + } + + fn test_schema() -> Schema { + Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("x", DataType::Int32, true), + Field::new("y", DataType::Int32, true), + ]) + } + + fn primary_key_on_id() -> Constraints { + Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0])]) + } + + /// A nullable unique key on column `x` (index 1). `Unique` (unlike + /// `PrimaryKey`) marks the dependency as nullable, which is what gates the + /// rewrite on the join's `null_equality`. + fn unique_on_x() -> Constraints { + Constraints::new_unverified(vec![Constraint::Unique(vec![1])]) + } + + /// A composite primary key spanning columns `id` and `x` (indices 0 and 1). + fn composite_primary_key_on_id_x() -> Constraints { + Constraints::new_unverified(vec![Constraint::PrimaryKey(vec![0, 1])]) + } } diff --git a/datafusion/optimizer/src/eliminate_outer_join.rs b/datafusion/optimizer/src/eliminate_outer_join.rs index cd060469b2990..a327da66f8ae2 100644 --- a/datafusion/optimizer/src/eliminate_outer_join.rs +++ b/datafusion/optimizer/src/eliminate_outer_join.rs @@ -15,39 +15,66 @@ // specific language governing permissions and limitations // under the License. -//! [`EliminateOuterJoin`] converts `LEFT/RIGHT/FULL` joins to `INNER` joins +//! [`EliminateOuterJoin`] rewrites outer joins to simpler join types when +//! filters make the outer rows unnecessary (e.g. `LEFT`/`RIGHT` to `INNER`, +//! and `FULL` to `LEFT`/`RIGHT`/`INNER`). +use crate::push_down_filter::replace_cols_by_name; use crate::{OptimizerConfig, OptimizerRule}; -use datafusion_common::{Column, DFSchema, Result}; -use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan}; +use datafusion_common::{Column, DFSchema, Result, qualified_name}; +use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan, Projection}; use datafusion_expr::{Expr, Filter, Operator}; use crate::optimizer::ApplyOrder; use datafusion_common::tree_node::Transformed; use datafusion_expr::expr::{BinaryExpr, Cast, InList, Like, TryCast}; +use std::collections::HashMap; use std::sync::Arc; +/// Attempt to simplify outer joins when filters make their null-padded +/// rows impossible to observe. /// -/// Attempt to replace outer joins with inner joins. +/// Outer joins are generally more expensive than inner joins and can block +/// predicate pushdown and other optimizations. When a filter above an outer +/// join removes every row the join would add for unmatched input rows, the +/// join can be changed to a cheaper join type. /// -/// Outer joins are typically more expensive to compute at runtime -/// than inner joins and prevent various forms of predicate pushdown -/// and other optimizations, so removing them if possible is beneficial. +/// For example: /// -/// Inner joins filter out rows that do match. Outer joins pass rows -/// that do not match padded with nulls. If there is a filter in the -/// query that would filter any such null rows after the join the rows -/// introduced by the outer join are filtered. +/// ```sql +/// SELECT ... +/// FROM a LEFT JOIN b ON ... +/// WHERE b.xx = 100 +/// ``` /// -/// For example, in the `select ... from a left join b on ... where b.xx = 100;` +/// For unmatched rows from `a`, the LEFT JOIN would produce a row with +/// `b.xx` set to NULL. The predicate `b.xx = 100` does not pass for those +/// rows, so the query does not need the LEFT JOIN's null-padded output and +/// the join can be rewritten as an inner join. /// -/// For rows when `b.xx` is null (as it would be after an outer join), -/// the `b.xx = 100` predicate filters them out and there is no -/// need to produce null rows for output. +/// The same reasoning can also simplify FULL joins to LEFT, RIGHT, or INNER +/// joins when filters remove the rows padded on one or both sides. /// -/// Generally, an outer join can be rewritten to inner join if the -/// filters from the WHERE clause return false while any inputs are -/// null and columns of those quals are come from nullable side of -/// outer join. +/// This rule looks for a filter above an outer join: +/// +/// ```text +/// Filter(predicate) +/// Join(LEFT/RIGHT/FULL) +/// ``` +/// +/// It also handles plan shapes where projection pruning has inserted one or +/// more Projection nodes between the filter and join: +/// +/// ```text +/// Filter(predicate over projection output) +/// Projection(...) +/// ... +/// Join(LEFT/RIGHT/FULL) +/// ``` +/// +/// In the projection case, the rule rewrites a copy of the predicate through +/// each Projection so it can analyze the predicate against the Join inputs. +/// The original filter predicate and Projection nodes are preserved when the +/// plan is rebuilt. #[derive(Default, Debug)] pub struct EliminateOuterJoin; @@ -77,60 +104,127 @@ impl OptimizerRule for EliminateOuterJoin { plan: LogicalPlan, _config: &dyn OptimizerConfig, ) -> Result> { - match plan { - LogicalPlan::Filter(mut filter) => match Arc::unwrap_or_clone(filter.input) { + let LogicalPlan::Filter(filter) = plan else { + return Ok(Transformed::no(plan)); + }; + + // Descend through one or more Projection nodes until we find a Join. + // For each Projection we encounter, rewrite a working copy of the + // predicate by replacing references to projection output columns with + // the expressions that define them. Keep the filter's original + // predicate intact for eventual use in the rebuilt plan; the rewritten + // predicate is used only for the null-rejection analysis. + let mut rewritten_predicate = filter.predicate.clone(); + let mut projections: Vec = Vec::new(); + let mut cur = Arc::clone(&filter.input); + + let new_join = loop { + match cur.as_ref() { + LogicalPlan::Projection(p) => { + rewritten_predicate = + inline_through_projection(rewritten_predicate, p)?; + let next = Arc::clone(&p.input); + projections.push(p.clone()); + cur = next; + } LogicalPlan::Join(join) => { - let mut non_nullable_cols: Vec = vec![]; - - extract_non_nullable_columns( - &filter.predicate, - &mut non_nullable_cols, - join.left.schema(), - join.right.schema(), - true, - ); - - let new_join_type = if join.join_type.is_outer() { - let mut left_non_nullable = false; - let mut right_non_nullable = false; - for col in non_nullable_cols.iter() { - if join.left.schema().has_column(col) { - left_non_nullable = true; - } - if join.right.schema().has_column(col) { - right_non_nullable = true; - } - } - eliminate_outer( - join.join_type, - left_non_nullable, - right_non_nullable, - ) - } else { - join.join_type + let Some(new_join) = try_simplify_join(join, &rewritten_predicate) + else { + return Ok(Transformed::no(LogicalPlan::Filter(filter))); }; - - let new_join = Arc::new(LogicalPlan::Join(Join { - left: join.left, - right: join.right, - join_type: new_join_type, - join_constraint: join.join_constraint, - on: join.on.clone(), - filter: join.filter.clone(), - schema: Arc::clone(&join.schema), - null_equality: join.null_equality, - null_aware: join.null_aware, - })); - Filter::try_new(filter.predicate, new_join) - .map(|f| Transformed::yes(LogicalPlan::Filter(f))) + break new_join; } - filter_input => { - filter.input = Arc::new(filter_input); - Ok(Transformed::no(LogicalPlan::Filter(filter))) + _ => { + return Ok(Transformed::no(LogicalPlan::Filter(filter))); } - }, - _ => Ok(Transformed::no(plan)), - } + } + }; + + let rebuilt_inner = rewrap_projections(new_join, projections); + Filter::try_new(filter.predicate, Arc::new(rebuilt_inner)) + .map(|f| Transformed::yes(LogicalPlan::Filter(f))) + } +} + +/// Attempt to simplify an outer join by analyzing `predicate` for +/// null-rejection. If the predicate filters out rows padded with NULLs on one +/// or both sides, return a copy of `join` rewritten to an equivalent join type +/// that omits those rows in the first place; otherwise return `None`. +fn try_simplify_join(join: &Join, predicate: &Expr) -> Option { + if !join.join_type.is_outer() { + return None; + } + + let null_rejecting_sides = extract_null_rejecting_sides( + predicate, + join.left.schema(), + join.right.schema(), + true, + ); + + let new_join_type = eliminate_outer( + join.join_type, + null_rejecting_sides.left, + null_rejecting_sides.right, + ); + if new_join_type == join.join_type { + return None; + } + + Some(LogicalPlan::Join(Join { + left: Arc::clone(&join.left), + right: Arc::clone(&join.right), + join_type: new_join_type, + join_constraint: join.join_constraint, + on: join.on.clone(), + filter: join.filter.clone(), + schema: Arc::clone(&join.schema), + null_equality: join.null_equality, + null_aware: join.null_aware, + })) +} + +/// Substitute the projection's output column references in `predicate` with +/// the projection's defining expressions (stripped of any `Alias` wrapper). +/// The result expresses `predicate` over the projection's *input* schema. +/// +/// Unlike `PushDownFilter`, this rule does not change expression evaluation +/// behavior (in fact, the rewritten expressions are only used for analysis +/// purposes). Therefore, function volatility and `MoveTowardsLeafNodes` +/// placement can be ignored here. +fn inline_through_projection(predicate: Expr, p: &Projection) -> Result { + let mut map: HashMap = HashMap::new(); + for ((qualifier, field), expr) in p.schema.iter().zip(p.expr.iter()) { + map.insert( + qualified_name(qualifier, field.name()), + unalias(expr).clone(), + ); + } + replace_cols_by_name(predicate, &map) +} + +/// Re-attach a stack of projections above `new_inner`, restoring the original +/// plan shape with the new (possibly retyped) join at the bottom. Projection +/// schemas are reused as-is; only nullability of columns sourced from the +/// formerly-outer side may have changed, and the existing rule already takes +/// this looser-schema approach at the join itself. +fn rewrap_projections( + new_inner: LogicalPlan, + projections: Vec, +) -> LogicalPlan { + let mut current = new_inner; + for mut p in projections.into_iter().rev() { + p.input = Arc::new(current); + current = LogicalPlan::Projection(p); + } + current +} + +fn unalias(expr: &Expr) -> &Expr { + if let Expr::Alias(a) = expr { + unalias(&a.expr) + } else { + expr } } @@ -139,212 +233,159 @@ pub fn eliminate_outer( left_non_nullable: bool, right_non_nullable: bool, ) -> JoinType { - let mut new_join_type = join_type; - match join_type { - JoinType::Left if right_non_nullable => { - new_join_type = JoinType::Inner; + match (join_type, left_non_nullable, right_non_nullable) { + (JoinType::Left, _, true) => JoinType::Inner, + (JoinType::Right, true, _) => JoinType::Inner, + (JoinType::Full, true, true) => JoinType::Inner, + (JoinType::Full, true, false) => JoinType::Left, + (JoinType::Full, false, true) => JoinType::Right, + _ => join_type, + } +} + +#[derive(Default, Clone, Copy, Debug, PartialEq, Eq)] +struct NullRejectingSides { + left: bool, + right: bool, +} + +impl NullRejectingSides { + /// The join side(s) a column belongs to. + /// + /// A bare column reference is null-rejecting on its own side: if the column + /// is NULL, every null-propagating operator above it yields NULL and the row + /// is filtered. + fn for_column(col: &Column, left_schema: &DFSchema, right_schema: &DFSchema) -> Self { + Self { + left: left_schema.has_column(col), + right: right_schema.has_column(col), } - JoinType::Left => {} - JoinType::Right if left_non_nullable => { - new_join_type = JoinType::Inner; + } + + fn union(self, other: Self) -> Self { + Self { + left: self.left || other.left, + right: self.right || other.right, } - JoinType::Right => {} - JoinType::Full => { - if left_non_nullable && right_non_nullable { - new_join_type = JoinType::Inner; - } else if left_non_nullable { - new_join_type = JoinType::Left; - } else if right_non_nullable { - new_join_type = JoinType::Right; - } + } + + fn intersection(self, other: Self) -> Self { + Self { + left: self.left && other.left, + right: self.right && other.right, } - _ => {} } - new_join_type } -/// Recursively traverses expr, if expr returns false when -/// any inputs are null, treats columns of both sides as non_nullable columns. +/// Compute which join sides are null-rejected by `expr` in a WHERE clause. +/// For each marked side, rows padded with NULLs on that side are guaranteed to +/// evaluate to NULL or false and be filtered out. /// -/// For and/or expr, extracts from all sub exprs and merges the columns. -/// For or expr, if one of sub exprs returns true, discards all columns from or expr. -/// For IS NOT NULL/NOT expr, always returns false for NULL input. -/// extracts columns from these exprs. -/// For all other exprs, fall through -fn extract_non_nullable_columns( +/// `left_schema` and `right_schema` map column references to join sides. +/// `top_level` is true only while walking the root WHERE context; nested +/// contexts are more conservative because their boolean result may be combined +/// by an enclosing expression. +fn extract_null_rejecting_sides( expr: &Expr, - non_nullable_cols: &mut Vec, left_schema: &Arc, right_schema: &Arc, top_level: bool, -) { +) -> NullRejectingSides { match expr { Expr::Column(col) => { - non_nullable_cols.push(col.clone()); + NullRejectingSides::for_column(col, left_schema, right_schema) } Expr::BinaryExpr(BinaryExpr { left, op, right }) => match op { - // If one of the inputs are null for these operators, the results should be false. - Operator::Eq - | Operator::NotEq - | Operator::Lt - | Operator::LtEq - | Operator::Gt - | Operator::GtEq => { - extract_non_nullable_columns( - left, - non_nullable_cols, - left_schema, - right_schema, - false, - ); - extract_non_nullable_columns( - right, - non_nullable_cols, - left_schema, - right_schema, - false, - ) - } Operator::And | Operator::Or => { - // treat And as Or if does not from top level, such as - // not (c1 < 10 and c2 > 100) - if top_level && *op == Operator::And { - extract_non_nullable_columns( - left, - non_nullable_cols, - left_schema, - right_schema, - top_level, - ); - extract_non_nullable_columns( - right, - non_nullable_cols, - left_schema, - right_schema, - top_level, - ); - return; - } - - let mut left_non_nullable_cols: Vec = vec![]; - let mut right_non_nullable_cols: Vec = vec![]; - - extract_non_nullable_columns( + let left_sides = extract_null_rejecting_sides( left, - &mut left_non_nullable_cols, left_schema, right_schema, top_level, ); - extract_non_nullable_columns( + let right_sides = extract_null_rejecting_sides( right, - &mut right_non_nullable_cols, left_schema, right_schema, top_level, ); - // for query: select *** from a left join b where b.c1 ... or b.c2 ... - // this can be eliminated to inner join. - // for query: select *** from a left join b where a.c1 ... or b.c2 ... - // this can not be eliminated. - // If columns of relation exist in both sub exprs, any columns of this relation - // can be added to non nullable columns. - if !left_non_nullable_cols.is_empty() - && !right_non_nullable_cols.is_empty() - { - for left_col in &left_non_nullable_cols { - for right_col in &right_non_nullable_cols { - if (left_schema.has_column(left_col) - && left_schema.has_column(right_col)) - || (right_schema.has_column(left_col) - && right_schema.has_column(right_col)) - { - non_nullable_cols.push(left_col.clone()); - break; - } - } - } + // Top-level AND: each conjunct is an independent WHERE filter, + // so side evidence from either branch is sufficient. + // Nested AND is handled like OR because the enclosing context + // may still let a NULL-padded row pass. + if top_level && *op == Operator::And { + left_sides.union(right_sides) + } else { + // OR (and nested AND): a NULL-padded row is rejected only + // if both branches reject NULLs for the same side. + left_sides.intersection(right_sides) } } - _ => {} + // Other NULL-on-NULL operators preserve null rejection from either + // operand. + op if op.returns_null_on_null() => { + let left_sides = + extract_null_rejecting_sides(left, left_schema, right_schema, false); + let right_sides = + extract_null_rejecting_sides(right, left_schema, right_schema, false); + left_sides.union(right_sides) + } + // Other operators, notably IS [ NOT ] DISTINCT FROM, are not + // NULL-propagating and provide no side-level rejection evidence. + _ => NullRejectingSides::default(), }, - Expr::Not(arg) => extract_non_nullable_columns( - arg, - non_nullable_cols, - left_schema, - right_schema, - false, - ), - Expr::IsNotNull(arg) => { - if !top_level { - return; + Expr::Not(arg) | Expr::Negative(arg) => { + extract_null_rejecting_sides(arg, left_schema, right_schema, false) + } + // These wrappers return FALSE on NULL input, so they reject NULLs only + // when they are themselves in the root WHERE context. Under another + // expression, that FALSE can be transformed into a NULL-accepting result + // (for example by NOT), so recurse only at the top level. + Expr::IsNotNull(arg) + | Expr::IsTrue(arg) + | Expr::IsFalse(arg) + | Expr::IsNotUnknown(arg) => { + if top_level { + extract_null_rejecting_sides(arg, left_schema, right_schema, false) + } else { + NullRejectingSides::default() } - extract_non_nullable_columns( - arg, - non_nullable_cols, - left_schema, - right_schema, - false, - ) } Expr::Cast(Cast { expr, field: _ }) - | Expr::TryCast(TryCast { expr, field: _ }) => extract_non_nullable_columns( - expr, - non_nullable_cols, - left_schema, - right_schema, - false, - ), - // IN list and BETWEEN are null-rejecting on the input expression: - // if the input column is NULL, the result is NULL (filtered out), - // regardless of whether the list/range contains NULLs. - Expr::InList(InList { expr, .. }) => extract_non_nullable_columns( - expr, - non_nullable_cols, - left_schema, - right_schema, - false, - ), - Expr::Between(between) => extract_non_nullable_columns( - &between.expr, - non_nullable_cols, - left_schema, - right_schema, - false, - ), - // LIKE is null-rejecting: if either the input column or the pattern - // is NULL, the result is NULL (filtered out by WHERE). - Expr::Like(Like { expr, pattern, .. }) => { - extract_non_nullable_columns( - expr, - non_nullable_cols, - left_schema, - right_schema, - false, - ); - extract_non_nullable_columns( - pattern, - non_nullable_cols, - left_schema, - right_schema, - false, - ); + | Expr::TryCast(TryCast { expr, field: _ }) => { + extract_null_rejecting_sides(expr, left_schema, right_schema, false) } - // IS TRUE, IS FALSE, and IS NOT UNKNOWN are null-rejecting: - // if the input is NULL, they return false (filtered out by WHERE). - // Note: IS NOT TRUE, IS NOT FALSE, and IS UNKNOWN are NOT null-rejecting - // because they return true for NULL input. - Expr::IsTrue(arg) | Expr::IsFalse(arg) | Expr::IsNotUnknown(arg) => { - extract_non_nullable_columns( - arg, - non_nullable_cols, - left_schema, - right_schema, - false, - ) + // IN list and BETWEEN reject NULLs from their input expression; list + // values and range bounds do not affect which join side is padded. + Expr::InList(InList { expr, .. }) => { + extract_null_rejecting_sides(expr, left_schema, right_schema, false) } - _ => {} + Expr::Between(between) => { + extract_null_rejecting_sides(&between.expr, left_schema, right_schema, false) + } + Expr::Like(Like { expr, pattern, .. }) => { + let expr_sides = + extract_null_rejecting_sides(expr, left_schema, right_schema, false); + let pattern_sides = + extract_null_rejecting_sides(pattern, left_schema, right_schema, false); + expr_sides.union(pattern_sides) + } + // Strict scalar functions are NULL-propagating: if any argument from a + // padded join side is NULL, the function result is NULL, and an + // enclosing NULL-rejecting predicate filters the row out. + Expr::ScalarFunction(func) if func.func.is_strict() => func + .args + .iter() + .map(|arg| { + extract_null_rejecting_sides(arg, left_schema, right_schema, false) + }) + .fold(NullRejectingSides::default(), NullRejectingSides::union), + // Everything else is conservative: NULL-accepting predicates such as + // IS NULL / IS NOT TRUE / IS NOT FALSE / IS UNKNOWN must not eliminate + // an outer join, and non-strict functions/subqueries/accessors/literals + // have no uniform NULL-propagation contract here. + _ => NullRejectingSides::default(), } } @@ -357,12 +398,54 @@ mod tests { use arrow::datatypes::DataType; use datafusion_common::ScalarValue; use datafusion_expr::{ + ColumnarValue, Operator::{And, Or}, - binary_expr, cast, col, lit, + ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, Volatility, binary_expr, + cast, col, lit, logical_plan::builder::LogicalPlanBuilder, - try_cast, + not, try_cast, }; + #[test] + fn null_rejecting_sides_union() { + let left_side = NullRejectingSides { + left: true, + right: false, + }; + let right_side = NullRejectingSides { + left: false, + right: true, + }; + + assert_eq!( + left_side.union(right_side), + NullRejectingSides { + left: true, + right: true, + } + ); + } + + #[test] + fn null_rejecting_sides_intersection() { + let both_sides = NullRejectingSides { + left: true, + right: true, + }; + let right_side = NullRejectingSides { + left: false, + right: true, + }; + + assert_eq!( + both_sides.intersection(right_side), + NullRejectingSides { + left: false, + right: true, + } + ); + } + macro_rules! assert_optimized_plan_equal { ( $plan:expr, @@ -379,6 +462,57 @@ mod tests { }}; } + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdf { + name: &'static str, + signature: Signature, + strict: bool, + } + + impl TestUdf { + fn new(name: &'static str, strict: bool) -> Self { + Self { + name, + signature: Signature::uniform( + 1, + vec![DataType::UInt32], + Volatility::Immutable, + ), + strict, + } + } + } + + impl ScalarUDFImpl for TestUdf { + fn name(&self) -> &str { + self.name + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::UInt32) + } + + fn is_strict(&self) -> bool { + self.strict + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + unimplemented!() + } + } + + fn strict_udf(arg: Expr) -> Expr { + ScalarUDF::from(TestUdf::new("strict_test", true)).call(vec![arg]) + } + + fn non_strict_udf(arg: Expr) -> Expr { + ScalarUDF::from(TestUdf::new("non_strict_test", false)).call(vec![arg]) + } + #[test] fn eliminate_left_with_null() -> Result<()> { let t1 = test_table_scan_with_name("t1")?; @@ -427,6 +561,98 @@ mod tests { ") } + #[test] + fn eliminate_left_with_strict_function() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(strict_udf(col("t2.b")).gt(lit(5u32)))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: strict_test(t2.b) > UInt32(5) + Inner Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_left_with_non_strict_function() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(non_strict_udf(col("t2.b")).gt(lit(5u32)))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: non_strict_test(t2.b) > UInt32(5) + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn eliminate_left_with_nested_strict_is_not_null() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(strict_udf(strict_udf(col("t2.b"))).is_not_null())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: strict_test(strict_test(t2.b)) IS NOT NULL + Inner Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_left_with_strict_function_is_null() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(strict_udf(col("t2.b")).is_null())? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: strict_test(t2.b) IS NULL + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + #[test] fn eliminate_right_with_or() -> Result<()> { let t1 = test_table_scan_with_name("t1")?; @@ -896,6 +1122,83 @@ mod tests { ") } + #[test] + fn no_eliminate_left_with_not_is_true() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // NOT( IS TRUE) is equivalent to ( IS NOT TRUE): TRUE when + // is FALSE OR NULL. So `WHERE NOT((t2.b > 5) IS TRUE)` accepts + // rows where t2.b is NULL (because t2.b > 5 is NULL → IS TRUE is + // FALSE → NOT FALSE = TRUE). The LEFT JOIN must NOT be converted. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(not(col("t2.b").gt(lit(5u32)).is_true()))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: NOT t2.b > UInt32(5) IS TRUE + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_left_with_not_is_false() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // Same shape, IS FALSE: NOT( IS FALSE) accepts NULL on the + // inner column. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(not(col("t2.b").gt(lit(5u32)).is_false()))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: NOT t2.b > UInt32(5) IS FALSE + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_left_with_not_is_not_unknown() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // Same shape, IS NOT UNKNOWN: NOT( IS NOT UNKNOWN) is + // equivalent to ( IS UNKNOWN), which is TRUE when is NULL. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(not(col("t2.b").gt(lit(5u32)).is_not_unknown()))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: NOT t2.b > UInt32(5) IS NOT UNKNOWN + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + #[test] fn eliminate_full_with_type_cast() -> Result<()> { let t1 = test_table_scan_with_name("t1")?; @@ -1171,4 +1474,241 @@ mod tests { TableScan: t2 ") } + + // ----- Filter pierces a Projection to reach the Join ----- + + #[test] + fn eliminate_left_through_projection() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // Filter → Projection → LeftJoin is the shape produced by projection + // pruning in queries such as TPC-DS q49, where the post-join + // Projection sits between the filter and the join. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .project(vec![col("t1.a"), col("t2.b").alias("bb")])? + .filter(col("bb").gt(lit(10u32)))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: bb > UInt32(10) + Projection: t1.a, t2.b AS bb + Inner Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_left_through_projection_with_or_cross_side() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // After inlining the filter is still t1.b > 10 OR t2.b < 20, which + // is null-tolerant when t2 is NULL (the t1.b clause can still hold). + // The LEFT JOIN must be preserved. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .project(vec![col("t1.b").alias("x"), col("t2.b").alias("y")])? + .filter(binary_expr( + col("x").gt(lit(10u32)), + Or, + col("y").lt(lit(20u32)), + ))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: x > UInt32(10) OR y < UInt32(20) + Projection: t1.b AS x, t2.b AS y + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_left_through_projection_with_only_left_filter() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // A filter that constrains only the preserved (left) side of a + // LEFT JOIN does not justify converting it to INNER — the LEFT + // would still pass nullable right-side rows that the filter + // accepts. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .project(vec![col("t1.b").alias("x"), col("t2.b")])? + .filter(col("x").gt(lit(10u32)))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: x > UInt32(10) + Projection: t1.b AS x, t2.b + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn eliminate_left_with_arithmetic_predicate() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // t2.b * 2 + 1 > 10 is null-rejecting on t2.b: arithmetic + // operators propagate NULL, so the whole expression is NULL when + // t2.b is NULL, and NULL > 10 is filtered out by WHERE. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter( + binary_expr( + binary_expr(col("t2.b"), Operator::Multiply, lit(2u32)), + Operator::Plus, + lit(1u32), + ) + .gt(lit(10u32)), + )? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: t2.b * UInt32(2) + UInt32(1) > UInt32(10) + Inner Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + #[test] + fn eliminate_left_with_negative_predicate() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // Unary minus propagates NULL: -NULL is NULL, so `WHERE -t2.b > 0` + // is null-rejecting on t2.b. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(Expr::Negative(Box::new(col("t2.b"))).gt(lit(0u32)))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: (- t2.b) > UInt32(0) + Inner Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_left_with_is_distinct_from() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // IS DISTINCT FROM is NOT null-rejecting: t2.b IS DISTINCT FROM 5 is + // true when t2.b is NULL (NULL is distinct from 5). Padding rows from + // a LEFT JOIN would survive the filter, so the LEFT JOIN must stay. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(binary_expr( + col("t2.b"), + Operator::IsDistinctFrom, + lit(5u32), + ))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: t2.b IS DISTINCT FROM UInt32(5) + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_left_with_is_not_distinct_from() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // IS NOT DISTINCT FROM is also not null-rejecting: t2.b IS NOT + // DISTINCT FROM NULL is true when t2.b is NULL. The LEFT JOIN must + // stay. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .filter(binary_expr( + col("t2.b"), + Operator::IsNotDistinctFrom, + lit(ScalarValue::UInt32(None)), + ))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: t2.b IS NOT DISTINCT FROM UInt32(NULL) + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn no_eliminate_through_non_transparent() -> Result<()> { + let t1 = test_table_scan_with_name("t1")?; + let t2 = test_table_scan_with_name("t2")?; + + // Limit is intentionally not treated as transparent: a Limit below + // the Filter changes which rows survive, so swapping LEFT→INNER + // beneath it could yield a different surviving-row set even when + // the filter is null-rejecting on the right side. + let plan = LogicalPlanBuilder::from(t1) + .join( + t2, + JoinType::Left, + (vec![Column::from_name("a")], vec![Column::from_name("a")]), + None, + )? + .limit(0, Some(5))? + .filter(col("t2.b").gt(lit(10u32)))? + .build()?; + + assert_optimized_plan_equal!(plan, @r" + Filter: t2.b > UInt32(10) + Limit: skip=0, fetch=5 + Left Join: t1.a = t2.a + TableScan: t1 + TableScan: t2 + ") + } } diff --git a/datafusion/optimizer/src/extract_leaf_expressions.rs b/datafusion/optimizer/src/extract_leaf_expressions.rs index c5c5610aeaed9..b855f224c420b 100644 --- a/datafusion/optimizer/src/extract_leaf_expressions.rs +++ b/datafusion/optimizer/src/extract_leaf_expressions.rs @@ -21,7 +21,7 @@ //! [`ExtractLeafExpressions`] (pass 1) and [`PushDownLeafProjections`] (pass 2). use indexmap::{IndexMap, IndexSet}; -use std::collections::HashMap; +use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; use datafusion_common::alias::AliasGenerator; @@ -827,12 +827,8 @@ fn split_and_push_projection( let original_schema = proj.schema.as_ref(); let mut recovery_exprs: Vec = Vec::with_capacity(proj.expr.len()); - let mut needs_recovery = false; let mut has_new_extractions = false; let mut proj_exprs_captured: usize = 0; - // Track standalone column expressions (Case B) to detect column refs - // from extracted aliases (Case A) that aren't also standalone expressions. - let mut standalone_columns: IndexSet = IndexSet::new(); for (expr, (qualifier, field)) in proj.expr.iter().zip(original_schema.iter()) { if let Expr::Alias(alias) = expr @@ -854,7 +850,6 @@ fn split_and_push_projection( } else if let Expr::Column(col) = expr { // Plain column pass-through — track it in the extractor extractors[0].columns_needed.insert(col.clone()); - standalone_columns.insert(col.clone()); recovery_exprs.push(expr.clone()); proj_exprs_captured += 1; } else { @@ -875,7 +870,6 @@ fn split_and_push_projection( original_name != &expr_name }; let recovery_expr = if needs_alias { - needs_recovery = true; transformed_expr .clone() .alias_qualified(qualifier.cloned(), original_name) @@ -883,14 +877,6 @@ fn split_and_push_projection( transformed_expr.clone() }; - // If the expression was transformed (i.e., has extracted sub-parts), - // it differs from what the pushed projection outputs → needs recovery. - // Also, any non-column, non-__datafusion_extracted expression needs recovery - // because the pushed extraction projection won't output it directly. - if transformed.transformed || !matches!(expr, Expr::Column(_)) { - needs_recovery = true; - } - recovery_exprs.push(recovery_expr); } } @@ -913,17 +899,6 @@ fn split_and_push_projection( return Ok(None); } - // If columns_needed has entries that aren't standalone projection columns - // (i.e., they came from column refs inside extracted aliases), a merge - // into an inner projection will widen the schema with those extra columns, - // requiring a recovery projection to restore the original schema. - if columns_needed - .iter() - .any(|c| !standalone_columns.contains(c)) - { - needs_recovery = true; - } - // ── Phase 2: Push down ────────────────────────────────────────────── let proj_input = Arc::clone(&proj.input); let pushed = push_extraction_pairs( @@ -959,6 +934,37 @@ fn split_and_push_projection( } }; + // The recovery projection restores the original projection's output. We need + // it whenever `base_plan` no longer exposes the same set of output column + // names, which happens two ways: + // * a column is *renamed* — a transformed expression now surfaces as its + // internal `__datafusion_extracted_*` alias instead of the original name; + // * a column is *leaked* — pushing the projection down widens `base_plan` + // with an inner extraction projection's *other* extracted aliases bubbling + // up through a Filter. A schema-caching parent like SubqueryAlias then + // keeps a stale schema (see `map_children` in `logical_plan/tree_node.rs`) + // and the later `optimize_projections` pass fails to resolve columns. + // + // Both are captured by comparing the *set of unqualified field names*. We + // compare by unqualified name rather than the full qualified schema on + // purpose: extracted aliases are globally unique, so name-only comparison is + // unambiguous for them, while it ignores the benign column reordering and the + // `SubqueryAlias` re-qualification (`sub.__datafusion_extracted_1` vs + // `__datafusion_extracted_1`) that a qualified/ordered comparison would + // spuriously treat as drift, stacking redundant recovery projections. + let base_names: BTreeSet<&str> = base_plan + .schema() + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect(); + let original_names: BTreeSet<&str> = original_schema + .fields() + .iter() + .map(|f| f.name().as_str()) + .collect(); + let needs_recovery = base_names != original_names; + // Wrap with recovery projection if the output schema changed if needs_recovery { let recovery = LogicalPlan::Projection(Projection::try_new( @@ -1143,6 +1149,12 @@ fn try_push_into_inputs( return Ok(None); } + // Unnest may output a column with the same name but different value/type + // than its input column. Name-based routing cannot distinguish those. + if matches!(node, LogicalPlan::Unnest(_)) { + return Ok(None); + } + // SubqueryAlias remaps qualifiers between input and output. // Rewrite pairs/columns from alias-space to input-space before routing. let remapped = if let LogicalPlan::SubqueryAlias(sa) = node { @@ -1611,9 +1623,10 @@ mod tests { ## After Pushdown Projection: __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")) - Filter: __datafusion_extracted_1 = Utf8("active") - Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2 - TableScan: test projection=[user] + Projection: test.user, __datafusion_extracted_2 + Filter: __datafusion_extracted_1 = Utf8("active") + Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2 + TableScan: test projection=[user] ## Optimized Projection: __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")) @@ -1674,13 +1687,17 @@ mod tests { TableScan: test projection=[user] ## After Pushdown + Projection: test.user, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("label")) + Projection: test.user, __datafusion_extracted_2 + Filter: __datafusion_extracted_1 > Int32(150) + Projection: leaf_udf(test.user, Utf8("value")) AS __datafusion_extracted_1, test.user, leaf_udf(test.user, Utf8("label")) AS __datafusion_extracted_2 + TableScan: test projection=[user] + + ## Optimized Projection: test.user, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("label")) Filter: __datafusion_extracted_1 > Int32(150) Projection: leaf_udf(test.user, Utf8("value")) AS __datafusion_extracted_1, test.user, leaf_udf(test.user, Utf8("label")) AS __datafusion_extracted_2 TableScan: test projection=[user] - - ## Optimized - (same as after pushdown) "#) } @@ -1885,19 +1902,15 @@ mod tests { TableScan: test projection=[id, user] ## After Pushdown - Projection: test.id, test.user - Filter: __datafusion_extracted_1 IS NOT NULL - Filter: __datafusion_extracted_2 = Utf8("active") - Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1 - TableScan: test projection=[id, user] - - ## Optimized Projection: test.id, test.user Filter: __datafusion_extracted_1 IS NOT NULL Projection: test.id, test.user, __datafusion_extracted_1 Filter: __datafusion_extracted_2 = Utf8("active") Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1 TableScan: test projection=[id, user] + + ## Optimized + (same as after pushdown) "#) } @@ -2006,9 +2019,10 @@ mod tests { ## After Pushdown Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")), COUNT(Int32(1)) Aggregate: groupBy=[[__datafusion_extracted_1]], aggr=[[COUNT(Int32(1))]] - Filter: __datafusion_extracted_2 = Utf8("active") - Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1 - TableScan: test projection=[user] + Projection: test.user, __datafusion_extracted_1 + Filter: __datafusion_extracted_2 = Utf8("active") + Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_1 + TableScan: test projection=[user] ## Optimized Projection: __datafusion_extracted_1 AS leaf_udf(test.user,Utf8("name")), COUNT(Int32(1)) @@ -2085,19 +2099,15 @@ mod tests { TableScan: test projection=[a, b, c] ## After Pushdown - Projection: test.a, test.b, test.c - Filter: __datafusion_extracted_1 = Int32(2) - Filter: __datafusion_extracted_2 = Int32(1) - Projection: leaf_udf(test.a, Utf8("x")) AS __datafusion_extracted_2, test.a, test.b, test.c, leaf_udf(test.b, Utf8("y")) AS __datafusion_extracted_1 - TableScan: test projection=[a, b, c] - - ## Optimized Projection: test.a, test.b, test.c Filter: __datafusion_extracted_1 = Int32(2) Projection: test.a, test.b, test.c, __datafusion_extracted_1 Filter: __datafusion_extracted_2 = Int32(1) Projection: leaf_udf(test.a, Utf8("x")) AS __datafusion_extracted_2, test.a, test.b, test.c, leaf_udf(test.b, Utf8("y")) AS __datafusion_extracted_1 TableScan: test projection=[a, b, c] + + ## Optimized + (same as after pushdown) "#) } @@ -2307,21 +2317,15 @@ mod tests { ## After Pushdown Projection: test.id, test.user, right.id, right.user Filter: __datafusion_extracted_1 = Utf8("active") - Inner Join: __datafusion_extracted_2 = __datafusion_extracted_3 - Projection: leaf_udf(test.user, Utf8("id")) AS __datafusion_extracted_2, test.id, test.user, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1 - TableScan: test projection=[id, user] - Projection: leaf_udf(right.user, Utf8("id")) AS __datafusion_extracted_3, right.id, right.user - TableScan: right projection=[id, user] - - ## Optimized - Projection: test.id, test.user, right.id, right.user - Filter: __datafusion_extracted_1 = Utf8("active") - Projection: test.id, test.user, __datafusion_extracted_1, right.id, right.user + Projection: test.id, test.user, right.id, right.user, __datafusion_extracted_1 Inner Join: __datafusion_extracted_2 = __datafusion_extracted_3 Projection: leaf_udf(test.user, Utf8("id")) AS __datafusion_extracted_2, test.id, test.user, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1 TableScan: test projection=[id, user] Projection: leaf_udf(right.user, Utf8("id")) AS __datafusion_extracted_3, right.id, right.user TableScan: right projection=[id, user] + + ## Optimized + (same as after pushdown) "#) } @@ -2674,10 +2678,11 @@ mod tests { ## After Pushdown Projection: __datafusion_extracted_2 AS leaf_udf(sub.user,Utf8("name")) - Filter: __datafusion_extracted_1 = Utf8("active") - SubqueryAlias: sub - Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, test.user - TableScan: test projection=[user] + Projection: sub.user, __datafusion_extracted_2 + Filter: __datafusion_extracted_1 = Utf8("active") + SubqueryAlias: sub + Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, test.user + TableScan: test projection=[user] ## Optimized Projection: __datafusion_extracted_2 AS leaf_udf(sub.user,Utf8("name")) @@ -2849,9 +2854,10 @@ mod tests { ## After Pushdown Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")) - Filter: __datafusion_extracted_1 = Utf8("active") - Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2 - TableScan: test projection=[id, user] + Projection: test.id, test.user, __datafusion_extracted_2 + Filter: __datafusion_extracted_1 = Utf8("active") + Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2 + TableScan: test projection=[id, user] ## Optimized Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")) @@ -2886,9 +2892,10 @@ mod tests { ## After Pushdown Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("status")) - Filter: __datafusion_extracted_1 > Int32(5) - Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2 - TableScan: test projection=[id, user] + Projection: test.id, test.user, __datafusion_extracted_2 + Filter: __datafusion_extracted_1 > Int32(5) + Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2 + TableScan: test projection=[id, user] ## Optimized Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("status")) @@ -2940,11 +2947,12 @@ mod tests { ## After Pushdown Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_3 AS leaf_udf(right.user,Utf8("status")) - Left Join: Filter: test.id = right.id AND __datafusion_extracted_1 > Int32(5) - Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, test.id, test.user - TableScan: test projection=[id, user] - Projection: leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_1, right.id, right.user, leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_3 - TableScan: right projection=[id, user] + Projection: test.id, test.user, right.id, right.user, __datafusion_extracted_2, __datafusion_extracted_3 + Left Join: Filter: test.id = right.id AND __datafusion_extracted_1 > Int32(5) + Projection: leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, test.id, test.user + TableScan: test projection=[id, user] + Projection: leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_1, right.id, right.user, leaf_udf(right.user, Utf8("status")) AS __datafusion_extracted_3 + TableScan: right projection=[id, user] ## Optimized Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_3 AS leaf_udf(right.user,Utf8("status")) @@ -2985,9 +2993,10 @@ mod tests { ## After Pushdown Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_3 AS leaf_udf(test.user,Utf8("status")) - Filter: __datafusion_extracted_1 > Int32(5) - Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_3 - TableScan: test projection=[id, user] + Projection: test.id, test.user, __datafusion_extracted_2, __datafusion_extracted_3 + Filter: __datafusion_extracted_1 > Int32(5) + Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("name")) AS __datafusion_extracted_2, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_3 + TableScan: test projection=[id, user] ## Optimized Projection: test.id, __datafusion_extracted_2 AS leaf_udf(test.user,Utf8("name")), __datafusion_extracted_3 AS leaf_udf(test.user,Utf8("status")) @@ -3035,4 +3044,123 @@ mod tests { Ok(()) } + + /// Regression test: the optimizer must not push extractions through + /// `Unnest`. + /// + /// `try_push_into_inputs` routes extracted pairs to inputs by column name. + /// `Unnest` can emit an output column with the same name as its input + /// column but a different value/type (the unnested element), so name-based + /// routing cannot tell the two apart. `try_push_into_inputs` therefore + /// treats `Unnest` as a barrier and bails instead of pushing through it + /// (see the `matches!(node, LogicalPlan::Unnest(_))` guard there). + #[test] + fn test_no_push_through_unnest() -> Result<()> { + use arrow::datatypes::{DataType, Field, Schema}; + + let schema = Schema::new(vec![ + Field::new("list_col", DataType::new_list(DataType::Int32, true), true), + Field::new("other_col", DataType::Int32, true), + ]); + let table_scan = + datafusion_expr::logical_plan::table_scan(Some("t"), &schema, None)? + .build()?; + let plan = LogicalPlanBuilder::from(table_scan) + .unnest_column("list_col")? + .filter(leaf_udf(col("list_col"), "x").eq(lit(1i32)))? + .build()?; + + let ctx = OptimizerContext::new().with_max_passes(1); + let optimizer = Optimizer::with_rules(vec![ + Arc::new(ExtractLeafExpressions::new()), + Arc::new(PushDownLeafProjections::new()), + ]); + let optimized = optimizer.optimize(plan, &ctx, |_, _| {})?; + + insta::assert_snapshot!(format!("{optimized}"), @r#" + Projection: list_col, t.other_col + Filter: __datafusion_extracted_1 = Int32(1) + Projection: leaf_udf(list_col, Utf8("x")) AS __datafusion_extracted_1, list_col, t.other_col + Unnest: lists[t.list_col|depth=1] structs[] + TableScan: t + "#); + + Ok(()) + } + + /// Regression test: a leaf expression used in **both** the filter and the + /// projection, with the **bare base column** also projected, over a + /// `SubqueryAlias` whose projection emits an **extra column the outer query + /// never consumes** (`synth`). + /// + /// This reproduces a production failure where the leaf-pushdown passes drop + /// the bare passthrough column from an intermediate schema, causing the + /// subsequent `optimize_projections` run to fail with: + /// `Schema error: No field named __datafusion_extracted_N`. + /// + /// Equivalent SQL: + /// ```sql + /// CREATE VIEW v AS SELECT user, id, id + 1 AS synth FROM test; + /// SELECT user['status'], user, id FROM v WHERE user['status'] IS NOT NULL; + /// ``` + #[test] + fn test_subquery_alias_with_unconsumed_column() -> Result<()> { + let table_scan = test_table_scan_with_struct()?; + + // This is the plan shape *after* `push_down_filter` has run: it pushes + // the `leaf_udf(...)` filter down through the `SubqueryAlias` and below + // the view's inner projection. The filter and the outer projection now + // each contain the same leaf expression but are separated by the + // `SubqueryAlias`, so they extract into two *independent* aliases + // (`__datafusion_extracted_1` from the filter, `__datafusion_extracted_2` + // from the projection) instead of deduplicating into one. + // + // The view projects an extra `synth` column the outer query never + // consumes — without it the bug does not manifest. + let inner = LogicalPlanBuilder::from(table_scan) + .filter(leaf_udf(col("user"), "status").is_not_null())? + .project(vec![ + col("user"), + col("id"), + (col("id") + lit(1u32)).alias("synth"), + ])? + .alias("v")? + .build()?; + + // Outer projection: leaf expr + the bare base column + id. + let plan = LogicalPlanBuilder::from(inner) + .project(vec![ + leaf_udf(col("v.user"), "status"), + col("v.user"), + col("v.id"), + ])? + .build()?; + + // Run the leaf-pushdown passes followed by `optimize_projections`, + // exactly as the default optimizer schedules them. `optimize_projections` + // is what prunes the unused `synth` column and validates the plan; if the + // leaf passes drop the bare `v.user` passthrough column it fails with + // `Schema error: No field named __datafusion_extracted_N`. + let ctx = OptimizerContext::new(); + let optimizer = Optimizer::with_rules(vec![ + Arc::new(ExtractLeafExpressions::new()), + Arc::new(PushDownLeafProjections::new()), + Arc::new(OptimizeProjections::new()), + ]); + let optimized = optimizer.optimize(plan, &ctx, |_, _| {})?; + + // The bare `test.user` passthrough column is preserved and the view's + // output schema (`user`, `id`, `__datafusion_extracted_2`) is restored + // by a recovery projection, so `optimize_projections` succeeds. + insta::assert_snapshot!(format!("{optimized}"), @r#" + Projection: __datafusion_extracted_2 AS leaf_udf(v.user,Utf8("status")), v.user, v.id + SubqueryAlias: v + Projection: test.user, test.id, __datafusion_extracted_2 + Filter: __datafusion_extracted_1 IS NOT NULL + Projection: leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_1, test.id, test.user, leaf_udf(test.user, Utf8("status")) AS __datafusion_extracted_2 + TableScan: test projection=[id, user] + "#); + + Ok(()) + } } diff --git a/datafusion/optimizer/src/filter_null_join_keys.rs b/datafusion/optimizer/src/filter_null_join_keys.rs index c8f419d3e543e..e3de8048a879d 100644 --- a/datafusion/optimizer/src/filter_null_join_keys.rs +++ b/datafusion/optimizer/src/filter_null_join_keys.rs @@ -52,6 +52,7 @@ impl OptimizerRule for FilterNullJoinKeys { match plan { LogicalPlan::Join(mut join) if !join.on.is_empty() + && !join.null_aware && join.null_equality == NullEquality::NullEqualsNothing => { let (left_preserved, right_preserved) = @@ -359,4 +360,50 @@ mod tests { let t2 = table_scan(Some("t2"), &schema, None)?.build()?; Ok((t1, t2)) } + + #[test] + fn null_aware_left_mark_join_keys_not_filtered() -> Result<()> { + let (t1, t2) = test_tables()?; + let plan = build_null_aware_plan(t1, t2, JoinType::LeftMark)?; + + assert_optimized_plan_equal!(plan, @r" + LeftMark Join: t1.id = t2.optional_id null_aware + TableScan: t1 + TableScan: t2 + ") + } + + #[test] + fn null_aware_left_anti_join_keys_not_filtered() -> Result<()> { + let (t1, t2) = test_tables()?; + let plan = build_null_aware_plan(t1, t2, JoinType::LeftAnti)?; + + assert_optimized_plan_equal!(plan, @r" + LeftAnti Join: t1.id = t2.optional_id null_aware + TableScan: t1 + TableScan: t2 + ") + } + + /// A join whose nullable right key would get an `IS NOT NULL` filter if it + /// were not null-aware. + fn build_null_aware_plan( + left_table: LogicalPlan, + right_table: LogicalPlan, + join_type: JoinType, + ) -> Result { + LogicalPlanBuilder::from(left_table) + .join_detailed_with_options( + right_table, + join_type, + ( + vec![Column::from_qualified_name("t1.id")], + vec![Column::from_qualified_name("t2.optional_id")], + ), + None, + NullEquality::NullEqualsNothing, + true, + )? + .build() + } } diff --git a/datafusion/optimizer/src/optimize_projections/mod.rs b/datafusion/optimizer/src/optimize_projections/mod.rs index bc923706a44b0..80aceb8cad44c 100644 --- a/datafusion/optimizer/src/optimize_projections/mod.rs +++ b/datafusion/optimizer/src/optimize_projections/mod.rs @@ -29,15 +29,13 @@ use datafusion_common::{ }; use datafusion_expr::expr::Alias; use datafusion_expr::{ - Aggregate, Distinct, EmptyRelation, Expr, Projection, TableScan, Unnest, Window, - logical_plan::LogicalPlan, + Aggregate, Distinct, EmptyRelation, Expr, Projection, TableScanBuilder, Unnest, + Window, logical_plan::LogicalPlan, }; use crate::optimize_projections::required_indices::RequiredIndices; use crate::utils::NamePreserver; -use datafusion_common::tree_node::{ - Transformed, TreeNode, TreeNodeContainer, TreeNodeRecursion, -}; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeContainer}; /// Optimizer rule to prune unnecessary columns from intermediate schemas /// inside the [`LogicalPlan`]. This rule: @@ -269,23 +267,15 @@ fn optimize_projections( .transform_data(|plan| optimize_subqueries(plan, config)); } LogicalPlan::TableScan(table_scan) => { - let TableScan { - table_name, - source, - projection, - filters, - fetch, - projected_schema: _, - } = table_scan; - // Get indices referred to in the original (schema with all fields) // given projected indices. - let projection = match &projection { + let projection = match &table_scan.projection { Some(projection) => indices.into_mapped_indices(|idx| projection[idx]), None => indices.into_inner(), }; - let new_scan = - TableScan::try_new(table_name, source, Some(projection), filters, fetch)?; + let new_scan = TableScanBuilder::from(table_scan) + .with_projection(Some(projection)) + .build()?; return Transformed::yes(LogicalPlan::TableScan(new_scan)) .transform_data(|plan| optimize_subqueries(plan, config)); @@ -381,39 +371,35 @@ fn optimize_projections( // These operators have no inputs, so stop the optimization process. return Ok(Transformed::no(plan)); } - LogicalPlan::RecursiveQuery(recursive) => { - // Only allow subqueries that reference the current CTE; nested subqueries are not yet - // supported for projection pushdown for simplicity. - // TODO: be able to do projection pushdown on recursive CTEs with subqueries - if plan_contains_other_subqueries( - recursive.static_term.as_ref(), - &recursive.name, - ) || plan_contains_other_subqueries( - recursive.recursive_term.as_ref(), - &recursive.name, - ) { - return Ok(Transformed::no(plan)); - } - - plan.inputs() - .into_iter() - .map(|input| { - indices - .clone() - .with_projection_beneficial() - .with_plan_exprs(&plan, input.schema()) - }) - .collect::>>()? + LogicalPlan::RecursiveQuery(_) => { + // optimize the static and recursive terms: treat each recursive CTE term like a + // standalone subquery: optimize its internals, but do not push parent required indices + // through the RecursiveQuery boundary, as this can otherwise lead to bugs + // (see: https://github.com/apache/datafusion/issues/22249) + return plan.map_children(|c| { + let indices = RequiredIndices::new_for_all_exprs(&c); + optimize_projections(c, config, indices) + }); } LogicalPlan::Join(join) => { let left_len = join.left.schema().fields().len(); let right_len = join.right.schema().fields().len(); let (left_req_indices, right_req_indices) = split_join_requirements(left_len, right_len, indices, &join.join_type); - let left_indices = + let mut left_indices = left_req_indices.with_plan_exprs(&plan, join.left.schema())?; - let right_indices = + let mut right_indices = right_req_indices.with_plan_exprs(&plan, join.right.schema())?; + // Ensure an empty mark join still has a column to qualify mark + match join.join_type { + JoinType::LeftMark if right_indices.indices().is_empty() => { + right_indices = right_indices.append(&[0]); + } + JoinType::RightMark if left_indices.indices().is_empty() => { + left_indices = left_indices.append(&[0]); + } + _ => {} + } // Joins benefit from "small" input tables (lower memory usage). // Therefore, each child benefits from projection: vec![ @@ -536,6 +522,30 @@ fn optimize_subqueries( /// - `Ok(None)`: Signals that merge is not beneficial (and has not taken place). /// - `Err(error)`: An error occurred during the function call. fn merge_consecutive_projections(proj: Projection) -> Result> { + // Collapse the whole chain in one pass; otherwise an N-deep chain needs + // N outer optimizer passes to fully fold. + let mut current = proj; + let mut transformed_any = false; + loop { + let Transformed { + data, transformed, .. + } = merge_consecutive_projections_one_level(current)?; + current = data; + if !transformed { + break; + } + transformed_any = true; + } + Ok(if transformed_any { + Transformed::yes(current) + } else { + Transformed::no(current) + }) +} + +fn merge_consecutive_projections_one_level( + proj: Projection, +) -> Result> { let Projection { expr, input, @@ -876,64 +886,6 @@ pub fn is_projection_unnecessary( )) } -/// Returns true if the plan subtree contains any subqueries that are not the -/// CTE reference itself. This treats any non-CTE [`LogicalPlan::SubqueryAlias`] -/// node (including aliased relations) as a blocker, along with expression-level -/// subqueries like scalar, EXISTS, or IN. These cases prevent projection -/// pushdown for now because we cannot safely reason about their column usage. -fn plan_contains_other_subqueries(plan: &LogicalPlan, cte_name: &str) -> bool { - if let LogicalPlan::SubqueryAlias(alias) = plan - && alias.alias.table() != cte_name - && !subquery_alias_targets_recursive_cte(alias.input.as_ref(), cte_name) - { - return true; - } - - let mut found = false; - plan.apply_expressions(|expr| { - if expr_contains_subquery(expr) { - found = true; - Ok(TreeNodeRecursion::Stop) - } else { - Ok(TreeNodeRecursion::Continue) - } - }) - .expect("expression traversal never fails"); - if found { - return true; - } - - plan.inputs() - .into_iter() - .any(|child| plan_contains_other_subqueries(child, cte_name)) -} - -fn expr_contains_subquery(expr: &Expr) -> bool { - expr.exists(|e| match e { - Expr::ScalarSubquery(_) | Expr::Exists(_) | Expr::InSubquery(_) => Ok(true), - _ => Ok(false), - }) - // Safe unwrap since we are doing a simple boolean check - .unwrap() -} - -fn subquery_alias_targets_recursive_cte(plan: &LogicalPlan, cte_name: &str) -> bool { - match plan { - LogicalPlan::TableScan(scan) => scan.table_name.table() == cte_name, - LogicalPlan::SubqueryAlias(alias) => { - subquery_alias_targets_recursive_cte(alias.input.as_ref(), cte_name) - } - _ => { - let inputs = plan.inputs(); - if inputs.len() == 1 { - subquery_alias_targets_recursive_cte(inputs[0], cte_name) - } else { - false - } - } - } -} - #[cfg(test)] mod tests { use std::cmp::Ordering; @@ -2446,6 +2398,62 @@ mod tests { ) } + // Stacked filter-less LeftMark joins (from `= ANY` / `<> ALL`) must keep + // each `mark` qualified so they don't collide. + #[test] + fn optimize_projections_stacked_mark_joins_keep_qualified_mark() -> Result<()> { + let person = test_table_scan_with_name("person")?; + + let aliased_scan = |table: &str, alias: &str| -> Result { + LogicalPlanBuilder::from(test_table_scan_with_name(table)?) + .project(vec![col(format!("{table}.a"))])? + .alias(alias)? + .build() + }; + + let plan = LogicalPlanBuilder::from(person) + .join_on( + aliased_scan("s1", "__correlated_sq_1")?, + JoinType::LeftMark, + vec![lit(true)], + )? + .join_on( + aliased_scan("s2", "__correlated_sq_2")?, + JoinType::LeftMark, + vec![lit(true)], + )? + .join_on( + aliased_scan("s3", "__correlated_sq_3")?, + JoinType::LeftMark, + vec![lit(true)], + )? + .filter( + col("__correlated_sq_1.mark") + .or(col("__correlated_sq_2.mark")) + .and(not(col("__correlated_sq_3.mark"))), + )? + .project(vec![col("person.a")])? + .build()?; + + assert_optimized_plan_equal!( + plan, + @r" + Projection: person.a + Filter: (__correlated_sq_1.mark OR __correlated_sq_2.mark) AND NOT __correlated_sq_3.mark + LeftMark Join: Filter: Boolean(true) + LeftMark Join: Filter: Boolean(true) + LeftMark Join: Filter: Boolean(true) + TableScan: person projection=[a] + SubqueryAlias: __correlated_sq_1 + TableScan: s1 projection=[a] + SubqueryAlias: __correlated_sq_2 + TableScan: s2 projection=[a] + SubqueryAlias: __correlated_sq_3 + TableScan: s3 projection=[a] + " + ) + } + fn observe(_plan: &LogicalPlan, _rule: &dyn OptimizerRule) {} fn optimize(plan: LogicalPlan) -> Result { diff --git a/datafusion/optimizer/src/optimize_projections/required_indices.rs b/datafusion/optimizer/src/optimize_projections/required_indices.rs index 5e73a9fbeceda..33f0d48721a8b 100644 --- a/datafusion/optimizer/src/optimize_projections/required_indices.rs +++ b/datafusion/optimizer/src/optimize_projections/required_indices.rs @@ -17,7 +17,8 @@ //! [`RequiredIndices`] helper for OptimizeProjection -use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; +use crate::utils::for_each_referenced_index; +use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Column, DFSchemaRef, Result}; use datafusion_expr::{Expr, LogicalPlan}; @@ -112,29 +113,8 @@ impl RequiredIndices { /// * `input_schema`: The input schema to analyze for index requirements. /// * `expr`: An expression for which we want to find necessary field indices. fn add_expr(&mut self, input_schema: &DFSchemaRef, expr: &Expr) { - // `apply` does not descend into subqueries, so recurse manually to - // handle those cases. - expr.apply(|e| { - match e { - Expr::Column(c) | Expr::OuterReferenceColumn(_, c) => { - if let Some(idx) = input_schema.maybe_index_of_column(c) { - self.indices.push(idx); - } - } - Expr::ScalarSubquery(sub) => { - self.add_exprs(input_schema, &sub.outer_ref_columns); - } - Expr::Exists(ex) => { - self.add_exprs(input_schema, &ex.subquery.outer_ref_columns); - } - Expr::InSubquery(isq) => { - self.add_exprs(input_schema, &isq.subquery.outer_ref_columns); - } - _ => {} - } - Ok(TreeNodeRecursion::Continue) - }) - .expect("traversal is infallible"); + for_each_referenced_index(expr, input_schema, |idx| self.indices.push(idx)) + .expect("traversal is infallible"); } /// Like [`Self::add_expr`], but for multiple expressions. diff --git a/datafusion/optimizer/src/optimizer.rs b/datafusion/optimizer/src/optimizer.rs index 31f8088f79c98..db7ad8475273a 100644 --- a/datafusion/optimizer/src/optimizer.rs +++ b/datafusion/optimizer/src/optimizer.rs @@ -28,9 +28,18 @@ use log::{debug, warn}; use datafusion_common::alias::AliasGenerator; use datafusion_common::config::ConfigOptions; use datafusion_common::instant::Instant; -use datafusion_common::tree_node::{Transformed, TreeNodeRewriter}; +use datafusion_common::tree_node::{ + Transformed, TreeNode, TreeNodeRecursion, TreeNodeRewriter, +}; use datafusion_common::{DFSchema, DataFusionError, HashSet, Result, internal_err}; +use datafusion_expr::dml::CopyTo; use datafusion_expr::logical_plan::LogicalPlan; +use datafusion_expr::{ + Aggregate, Analyze, CreateMemoryTable, CreateView, DdlStatement, Distinct, + DistinctOn, DmlStatement, Explain, Expr, Extension, Filter, Join, Limit, Projection, + RecursiveQuery, Repartition, Sort, Statement, Subquery, SubqueryAlias, Union, Unnest, + Window, +}; use crate::common_subexpr_eliminate::CommonSubexprEliminate; use crate::decorrelate_lateral_join::DecorrelateLateralJoin; @@ -359,6 +368,226 @@ impl TreeNodeRewriter for Rewriter<'_> { } } +/// Applies `f` to each child (input) of `plan` in place, using +/// [`Arc::make_mut`] for copy-on-write semantics on `Arc` +/// children. When the `Arc` refcount is 1 (the common case here) +/// `Arc::make_mut` hands out a `&mut` without cloning; when it is >1 the +/// inner value is cloned first. +/// +/// Returns `Ok(true)` if any child was modified by `f`. +/// +/// This is deliberately private to the optimizer rather than a method on +/// [`LogicalPlan`]: it is an implementation detail of in-place rewriting, and +/// the `Arc::make_mut` approach does not generalize to the other tree types +/// (`Expr` children are `Box`ed; `PhysicalExpr`/`ExecutionPlan` children are +/// `Arc`, which `Arc::make_mut` cannot handle). If `TreeNode` ever +/// grows an in-place traversal this logic can move there. +/// +/// # Error semantics +/// +/// If `f` returns `Err` for a child, that error is returned immediately; +/// children visited earlier keep whatever modifications `f` already applied +/// to them — they are **not** rolled back. +fn map_children_mut Result>( + plan: &mut LogicalPlan, + mut f: F, +) -> Result { + Ok(match plan { + LogicalPlan::Projection(Projection { input, .. }) + | LogicalPlan::Filter(Filter { input, .. }) + | LogicalPlan::Repartition(Repartition { input, .. }) + | LogicalPlan::Window(Window { input, .. }) + | LogicalPlan::Aggregate(Aggregate { input, .. }) + | LogicalPlan::Sort(Sort { input, .. }) + | LogicalPlan::Limit(Limit { input, .. }) + | LogicalPlan::SubqueryAlias(SubqueryAlias { input, .. }) + | LogicalPlan::Analyze(Analyze { input, .. }) + | LogicalPlan::Dml(DmlStatement { input, .. }) + | LogicalPlan::Copy(CopyTo { input, .. }) + | LogicalPlan::Unnest(Unnest { input, .. }) => f(Arc::make_mut(input))?, + LogicalPlan::Subquery(Subquery { subquery, .. }) => f(Arc::make_mut(subquery))?, + LogicalPlan::Join(Join { left, right, .. }) => { + let l = f(Arc::make_mut(left))?; + let r = f(Arc::make_mut(right))?; + l || r + } + LogicalPlan::Union(Union { inputs, .. }) => { + let mut changed = false; + for input in inputs { + changed |= f(Arc::make_mut(input))?; + } + changed + } + LogicalPlan::Distinct(Distinct::All(input)) => f(Arc::make_mut(input))?, + LogicalPlan::Distinct(Distinct::On(DistinctOn { input, .. })) => { + f(Arc::make_mut(input))? + } + LogicalPlan::Explain(Explain { plan, .. }) => f(Arc::make_mut(plan))?, + LogicalPlan::Ddl(DdlStatement::CreateMemoryTable(CreateMemoryTable { + input, + .. + })) + | LogicalPlan::Ddl(DdlStatement::CreateView(CreateView { input, .. })) => { + f(Arc::make_mut(input))? + } + LogicalPlan::RecursiveQuery(RecursiveQuery { + static_term, + recursive_term, + .. + }) => { + let s = f(Arc::make_mut(static_term))?; + let r = f(Arc::make_mut(recursive_term))?; + s || r + } + LogicalPlan::Statement(Statement::Prepare(p)) => f(Arc::make_mut(&mut p.input))?, + LogicalPlan::Extension(Extension { node }) => { + let inputs = node.inputs(); + if inputs.is_empty() { + false + } else { + // Extension nodes don't expose mutable children, + // fall back to the ownership-based API + let mut changed = false; + let exprs = node.expressions(); + let new_inputs: Vec = inputs + .into_iter() + .map(|input| { + let mut plan = input.clone(); + if f(&mut plan)? { + changed = true; + } + Ok(plan) + }) + .collect::>>()?; + if changed { + *node = node.with_exprs_and_inputs(exprs, new_inputs)?; + } + changed + } + } + // plans without inputs + LogicalPlan::TableScan { .. } + | LogicalPlan::EmptyRelation { .. } + | LogicalPlan::Values { .. } + | LogicalPlan::DescribeTable(_) + | LogicalPlan::Ddl(DdlStatement::CreateExternalTable(_)) + | LogicalPlan::Ddl(DdlStatement::CreateCatalogSchema(_)) + | LogicalPlan::Ddl(DdlStatement::CreateCatalog(_)) + | LogicalPlan::Ddl(DdlStatement::CreateIndex(_)) + | LogicalPlan::Ddl(DdlStatement::DropTable(_)) + | LogicalPlan::Ddl(DdlStatement::DropView(_)) + | LogicalPlan::Ddl(DdlStatement::DropCatalogSchema(_)) + | LogicalPlan::Ddl(DdlStatement::CreateFunction(_)) + | LogicalPlan::Ddl(DdlStatement::DropFunction(_)) + | LogicalPlan::Statement(_) => false, + }) +} + +/// Rewrites a plan tree in place using `Arc::make_mut` for +/// copy-on-write semantics on `Arc` children. +/// +/// This avoids the `Arc::unwrap_or_clone` + `Arc::new` cycle that the +/// ownership-based `TreeNode::rewrite` performs at every child node. +/// +/// # Error semantics +/// +/// On `Err`, `*plan` is left in an **unspecified** state and must not be used. +/// Note this is different than consuming APIs such as [`TreeNode::rewrite`] +/// where the original plan is freed and no longer available on error +#[cfg_attr(feature = "recursive_protection", recursive::recursive)] +fn rewrite_plan_in_place( + plan: &mut LogicalPlan, + apply_order: ApplyOrder, + rule: &dyn OptimizerRule, + config: &dyn OptimizerConfig, +) -> Result { + // f_down phase + let mut changed = false; + if apply_order == ApplyOrder::TopDown { + // `rule.rewrite()` takes the plan by value, so bridge the `&mut` to an + // owned value with `std::mem::take`. `LogicalPlan::default()` is a cheap + // empty placeholder (shared empty schema, no allocation) and is + // overwritten with the rule's output on the next line. + let owned = std::mem::take(plan); + let result = rule.rewrite(owned, config)?; + *plan = result.data; + changed |= result.transformed; + // Respect TreeNodeRecursion::Stop/Jump from the rule + if result.tnr == TreeNodeRecursion::Stop { + return Ok(changed); + } + } + + let mut child_schema_changed = false; + let children_changed = map_children_mut(plan, |child| { + let old_schema = Arc::clone(child.schema()); + let child_changed = rewrite_plan_in_place(child, apply_order, rule, config)?; + if child_changed && old_schema.as_ref() != child.schema().as_ref() { + child_schema_changed = true; + } + Ok(child_changed) + })?; + changed |= children_changed; + + if child_schema_changed { + // Child rewrites can change their output schemas. Recompute the current + // node before later rules use positional requirements from that schema. + let owned = std::mem::take(plan); + *plan = owned.recompute_schema()?; + } + + // f_up phase + if apply_order == ApplyOrder::BottomUp { + let owned = std::mem::take(plan); + let result = rule.rewrite(owned, config)?; + *plan = result.data; + changed |= result.transformed; + } + + Ok(changed) +} + +/// Returns true if the plan contains any subquery expressions +/// (EXISTS, IN subquery, scalar subquery, set comparison). +/// +/// Used to determine whether the more expensive `rewrite_with_subqueries` +/// traversal is needed. When the plan has no subqueries, the cheaper +/// `rewrite` traversal is sufficient since all plan nodes are reachable +/// via direct children. +fn plan_has_subqueries(plan: &LogicalPlan) -> bool { + let mut found = false; + let _ = plan.apply(|node| { + if found { + return Ok(TreeNodeRecursion::Stop); + } + node.apply_expressions(|expr| { + if found { + return Ok(TreeNodeRecursion::Stop); + } + expr.apply(|e| { + if matches!( + e, + Expr::Exists(_) + | Expr::InSubquery(_) + | Expr::SetComparison(_) + | Expr::ScalarSubquery(_) + ) { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + }); + found +} + impl Optimizer { /// Optimizes the logical plan by applying optimizer rules, and /// invoking observer function after each call @@ -388,6 +617,12 @@ impl Optimizer { while i < options.optimizer.max_passes { log_plan(&format!("Optimizer input (pass {i})"), &new_plan); + // Track subquery presence across the pass. Refresh after changed + // rules so decorrelation can move later rules onto the in-place + // path; that path refreshes parent schemas after child schemas + // change. + let mut has_subqueries = plan_has_subqueries(&new_plan); + for rule in &self.rules { // If skipping failed rules, copy plan before attempting to rewrite // as rewriting is destructive @@ -400,9 +635,42 @@ impl Optimizer { let result = match rule.apply_order() { // optimizer handles recursion - Some(apply_order) => new_plan.rewrite_with_subqueries( - &mut Rewriter::new(apply_order, rule.as_ref(), config), - ), + Some(apply_order) => { + if has_subqueries { + // Plans with subqueries need the full + // rewrite_with_subqueries traversal to + // recurse into subquery plans. + new_plan.rewrite_with_subqueries( + &mut Rewriter::new( + apply_order, + rule.as_ref(), + config, + ), + ) + } else { + // No subqueries: use in-place rewriting + // with Arc::make_mut for zero-cost CoW on + // children, avoiding Arc unwrap/rewrap. + // + // On error `new_plan` is left in an unspecified + // state (see `rewrite_plan_in_place`); the result + // handling below discards it, restoring `prev_plan` + // when `skip_failed_rules` is set or propagating + // the error otherwise. + rewrite_plan_in_place( + &mut new_plan, + apply_order, + rule.as_ref(), + config, + ) + .map(|transformed| { + Transformed::new_transformed( + std::mem::take(&mut new_plan), + transformed, + ) + }) + } + } // rule handles recursion itself None => { rule.rewrite(new_plan, config) @@ -433,6 +701,7 @@ impl Optimizer { new_plan = data; observer(&new_plan, rule.as_ref()); if transformed { + has_subqueries = plan_has_subqueries(&new_plan); log_plan(rule.name(), &new_plan); } else { debug!( @@ -516,13 +785,15 @@ mod tests { use datafusion_common::tree_node::Transformed; use datafusion_common::{ - DFSchema, DFSchemaRef, DataFusionError, Result, assert_contains, plan_err, + Column, DFSchema, DFSchemaRef, DataFusionError, Result, assert_contains, plan_err, }; use datafusion_expr::logical_plan::EmptyRelation; - use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, Projection, col, lit}; + use datafusion_expr::{ + Expr, JoinType, LogicalPlan, LogicalPlanBuilder, Projection, col, lit, + }; use crate::optimizer::Optimizer; - use crate::test::test_table_scan; + use crate::test::{test_table_scan, test_table_scan_with_name}; use crate::{OptimizerConfig, OptimizerContext, OptimizerRule}; use super::ApplyOrder; @@ -606,6 +877,34 @@ mod tests { Ok(()) } + #[test] + fn in_place_rewrite_recomputes_parent_schema_when_child_schema_changes() -> Result<()> + { + let left = LogicalPlanBuilder::from(test_table_scan_with_name("left")?) + .project(vec![col("left.a"), col("left.b"), col("left.c")])? + .build()?; + let right = LogicalPlanBuilder::from(test_table_scan_with_name("right")?) + .project(vec![col("right.a"), col("right.b"), col("right.c")])? + .build()?; + let mut plan = LogicalPlanBuilder::from(left) + .join_on(right, JoinType::Inner, [col("left.a").eq(col("right.a"))])? + .build()?; + + assert_eq!(plan.schema().fields().len(), 6); + + let changed = super::rewrite_plan_in_place( + &mut plan, + ApplyOrder::TopDown, + &KeepOnlyAProjectionRule {}, + &OptimizerContext::new(), + )?; + + assert!(changed); + assert_eq!(plan.schema().fields().len(), 2); + assert!(plan.schema().has_column_with_unqualified_name("a")); + Ok(()) + } + #[test] fn optimizer_detects_plan_equal_to_the_initial() -> Result<()> { // Run a goofy optimizer, which rotates projection columns @@ -723,6 +1022,40 @@ mod tests { } } + #[derive(Default, Debug)] + struct KeepOnlyAProjectionRule {} + + impl OptimizerRule for KeepOnlyAProjectionRule { + fn name(&self) -> &str { + "keep_only_a_projection" + } + + fn apply_order(&self) -> Option { + Some(ApplyOrder::TopDown) + } + + fn supports_rewrite(&self) -> bool { + true + } + + fn rewrite( + &self, + plan: LogicalPlan, + _config: &dyn OptimizerConfig, + ) -> Result> { + let projection = match plan { + LogicalPlan::Projection(p) => p, + _ => return Ok(Transformed::no(plan)), + }; + + let expr = Expr::from(Column::from(projection.schema.qualified_field(0))); + + Ok(Transformed::yes(LogicalPlan::Projection( + Projection::try_new(vec![expr], Arc::clone(&projection.input))?, + ))) + } + } + /// A goofy rule doing rotation of columns in all projections. /// /// Useful to test cycle detection. diff --git a/datafusion/optimizer/src/push_down_filter.rs b/datafusion/optimizer/src/push_down_filter.rs index 46e129ad4bdd3..cf54ae254746d 100644 --- a/datafusion/optimizer/src/push_down_filter.rs +++ b/datafusion/optimizer/src/push_down_filter.rs @@ -30,21 +30,22 @@ use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, }; use datafusion_common::{ - Column, DFSchema, Result, assert_eq_or_internal_err, assert_or_internal_err, - internal_err, plan_err, qualified_name, + Column, DFSchema, Result, assert_eq_or_internal_err, internal_err, plan_err, + qualified_name, }; use datafusion_expr::expr::WindowFunction; use datafusion_expr::expr_rewriter::replace_col; -use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan, TableScan, Union}; +use datafusion_expr::logical_plan::{Join, JoinType, LogicalPlan}; use datafusion_expr::utils::{ conjunction, expr_to_columns, split_conjunction, split_conjunction_owned, }; use datafusion_expr::{ - BinaryExpr, Expr, Filter, Operator, Projection, TableProviderFilterPushDown, and, or, + BinaryExpr, Distinct, Expr, Filter, Operator, Projection, + TableProviderFilterPushDown, and, or, }; use crate::optimizer::ApplyOrder; -use crate::simplify_expressions::simplify_predicates; +use crate::simplify_expressions::{reorder_predicates, simplify_predicates}; use crate::utils::{ ColumnReference, has_all_column_refs, is_restrict_null_predicate, schema_columns, }; @@ -447,16 +448,13 @@ fn push_down_all_join( let mut on_filter_join_conditions = vec![]; let (on_left_preserved, on_right_preserved) = on_lr_is_preserved(join.join_type); - - if !on_filter.is_empty() { - for on in on_filter { - if on_left_preserved && checker.is_left_only(&on) { - left_push.push(on) - } else if on_right_preserved && checker.is_right_only(&on) { - right_push.push(on) - } else { - on_filter_join_conditions.push(on) - } + for on in on_filter { + if on_left_preserved && checker.is_left_only(&on) { + left_push.push(on) + } else if on_right_preserved && checker.is_right_only(&on) { + right_push.push(on) + } else { + on_filter_join_conditions.push(on) } } @@ -498,41 +496,46 @@ fn push_down_all_join( )); } + // Add any new join conditions as the non join predicates + let join_conditions_empty = join_conditions.is_empty(); + join_conditions.extend(on_filter_join_conditions); + join.filter = conjunction(join_conditions); + + if join_conditions_empty && left_push.is_empty() && right_push.is_empty() { + // wrap the join on the filter whose predicates must be kept, if any + return Ok(Transformed::no(with_filters( + keep_predicates, + LogicalPlan::Join(join), + ))); + } + if let Some(predicate) = conjunction(left_push) { - join.left = Arc::new(LogicalPlan::Filter(Filter::try_new(predicate, join.left)?)); + join.left = Arc::new(LogicalPlan::Filter(Filter::new(predicate, join.left))); } + if let Some(predicate) = conjunction(right_push) { - join.right = - Arc::new(LogicalPlan::Filter(Filter::try_new(predicate, join.right)?)); + join.right = Arc::new(LogicalPlan::Filter(Filter::new(predicate, join.right))); } - // Add any new join conditions as the non join predicates - join_conditions.extend(on_filter_join_conditions); - join.filter = conjunction(join_conditions); - // wrap the join on the filter whose predicates must be kept, if any - let plan = LogicalPlan::Join(join); - let plan = if let Some(predicate) = conjunction(keep_predicates) { - LogicalPlan::Filter(Filter::try_new(predicate, Arc::new(plan))?) - } else { - plan - }; - Ok(Transformed::yes(plan)) + Ok(Transformed::yes(with_filters( + keep_predicates, + LogicalPlan::Join(join), + ))) } fn push_down_join( - join: Join, - parent_predicate: Option<&Expr>, + mut join: Join, + parent_predicate: Option, ) -> Result> { // Split the parent predicate into individual conjunctive parts. - let predicates = parent_predicate - .map_or_else(Vec::new, |pred| split_conjunction_owned(pred.clone())); + let predicates = parent_predicate.map_or_else(Vec::new, split_conjunction_owned); // Extract conjunctions from the JOIN's ON filter, if present. let on_filters = join .filter - .as_ref() - .map_or_else(Vec::new, |filter| split_conjunction_owned(filter.clone())); + .take() + .map_or_else(Vec::new, split_conjunction_owned); // Are there any new join predicates that can be inferred from the filter expressions? let inferred_join_predicates = with_debug_timing("infer_join_predicates", || { @@ -573,6 +576,17 @@ fn infer_join_predicates( predicates: &[Expr], on_filters: &[Expr], ) -> Result> { + // Null-aware joins (e.g. `NOT IN` with a nullable subquery) rely on SQL + // three-valued logic: a NULL join key on the right/subquery side makes the + // predicate UNKNOWN and empties the result, so those NULLs must reach the + // join. Inferring an equi-key predicate here would rewrite a left-side + // predicate onto the right side and, because the inferred predicate must be + // null-rejecting, drop the subquery's NULL rows and produce wrong results. + // Skip inference entirely for null-aware joins. + if join.null_aware { + return Ok(vec![]); + } + // Only allow both side key is column. let join_col_keys = join .on @@ -773,13 +787,11 @@ impl OptimizerRule for PushDownFilter { plan: LogicalPlan, config: &dyn OptimizerConfig, ) -> Result> { - let _ = config.options(); + let _ = config; if let LogicalPlan::Join(join) = plan { return push_down_join(join, None); }; - let plan_schema = Arc::clone(plan.schema()); - let LogicalPlan::Filter(mut filter) = plan else { return Ok(Transformed::no(plan)); }; @@ -788,6 +800,7 @@ impl OptimizerRule for PushDownFilter { let old_predicate_len = predicate.len(); let new_predicates = with_debug_timing("simplify_predicates", || simplify_predicates(predicate))?; + if log_enabled!(Level::Debug) { debug!( "push_down_filter: simplify_predicates old_count={}, new_count={}", @@ -795,7 +808,14 @@ impl OptimizerRule for PushDownFilter { new_predicates.len() ); } - if old_predicate_len != new_predicates.len() { + + // Place cheap predicates before expensive ones, so the `AND` + // evaluator's right-side short-circuit can skip evaluating expensive + // predicates on rows that have already been filtered out. + let (new_predicates, reorder_changed) = reorder_predicates(new_predicates); + + let count_changed = old_predicate_len != new_predicates.len(); + if count_changed || reorder_changed { let Some(new_predicate) = conjunction(new_predicates) else { // new_predicates is empty - remove the filter entirely // Return the child plan without the filter @@ -812,44 +832,48 @@ impl OptimizerRule for PushDownFilter { } match Arc::unwrap_or_clone(filter.input) { - LogicalPlan::Filter(child_filter) => { - // child filters first to preserve execution order - let new_predicates = split_conjunction_owned(child_filter.predicate) - .into_iter() - .chain(split_conjunction_owned(filter.predicate)) - // use IndexSet to remove duplicates while preserving predicate order - .collect::>(); + LogicalPlan::Filter(mut child_filter) => { + // Child filters first to preserve execution order. + // Use IndexSet to remove duplicates while preserving predicate order. + let new_predicates: IndexSet = + split_conjunction_owned(child_filter.predicate) + .into_iter() + .chain(split_conjunction_owned(filter.predicate)) + .collect(); let Some(new_predicate) = conjunction(new_predicates) else { return plan_err!("at least one expression exists"); }; - let new_filter = LogicalPlan::Filter(Filter::try_new( - new_predicate, - child_filter.input, - )?); - - self.rewrite(new_filter, config) + child_filter.predicate = new_predicate; + self.rewrite(LogicalPlan::Filter(child_filter), config) } - LogicalPlan::Repartition(repartition) => { - let new_filter = - Filter::try_new(filter.predicate, Arc::clone(&repartition.input)) - .map(LogicalPlan::Filter)?; - insert_below(LogicalPlan::Repartition(repartition), new_filter) + LogicalPlan::Repartition(mut repartition) => { + filter.input = repartition.input; + repartition.input = Arc::new(LogicalPlan::Filter(filter)); + Ok(Transformed::yes(LogicalPlan::Repartition(repartition))) } LogicalPlan::Distinct(distinct) => { - let new_filter = - Filter::try_new(filter.predicate, Arc::clone(distinct.input())) - .map(LogicalPlan::Filter)?; - insert_below(LogicalPlan::Distinct(distinct), new_filter) + let distinct = match distinct { + Distinct::All(input) => { + filter.input = input; + Distinct::All(Arc::new(LogicalPlan::Filter(filter))) + } + Distinct::On(mut distinct) => { + filter.input = distinct.input; + distinct.input = Arc::new(LogicalPlan::Filter(filter)); + Distinct::On(distinct) + } + }; + + Ok(Transformed::yes(LogicalPlan::Distinct(distinct))) } - LogicalPlan::Sort(sort) => { - let new_filter = - Filter::try_new(filter.predicate, Arc::clone(&sort.input)) - .map(LogicalPlan::Filter)?; - insert_below(LogicalPlan::Sort(sort), new_filter) + LogicalPlan::Sort(mut sort) => { + filter.input = sort.input; + sort.input = Arc::new(LogicalPlan::Filter(filter)); + Ok(Transformed::yes(LogicalPlan::Sort(sort))) } - LogicalPlan::SubqueryAlias(subquery_alias) => { + LogicalPlan::SubqueryAlias(mut subquery_alias) => { let mut replace_map = HashMap::new(); for (i, (qualifier, field)) in subquery_alias.input.schema().iter().enumerate() @@ -861,30 +885,24 @@ impl OptimizerRule for PushDownFilter { Expr::Column(Column::new(qualifier.cloned(), field.name())), ); } - let new_predicate = replace_cols_by_name(filter.predicate, &replace_map)?; - let new_filter = LogicalPlan::Filter(Filter::try_new( - new_predicate, - Arc::clone(&subquery_alias.input), - )?); - insert_below(LogicalPlan::SubqueryAlias(subquery_alias), new_filter) + filter.predicate = replace_cols_by_name(filter.predicate, &replace_map)?; + filter.input = subquery_alias.input; + subquery_alias.input = Arc::new(LogicalPlan::Filter(filter)); + Ok(Transformed::yes(LogicalPlan::SubqueryAlias(subquery_alias))) } LogicalPlan::Projection(projection) => { let predicates = split_conjunction_owned(filter.predicate.clone()); - let (new_projection, keep_predicate) = + let (mut result, keep_predicates) = rewrite_projection(predicates, projection)?; - if new_projection.transformed { - match keep_predicate { - None => Ok(new_projection), - Some(keep_predicate) => new_projection.map_data(|child_plan| { - Filter::try_new(keep_predicate, Arc::new(child_plan)) - .map(LogicalPlan::Filter) - }), - } + if result.transformed { + result.data = with_filters(keep_predicates, result.data) } else { - filter.input = Arc::new(new_projection.data); - Ok(Transformed::no(LogicalPlan::Filter(filter))) + filter.input = Arc::new(result.data); + result.data = LogicalPlan::Filter(filter) } + + Ok(result) } LogicalPlan::Unnest(mut unnest) => { let predicates = split_conjunction_owned(filter.predicate.clone()); @@ -895,11 +913,10 @@ impl OptimizerRule for PushDownFilter { for idx in &unnest.struct_type_columns { let (sub_qualifier, field) = unnest.input.schema().qualified_field(*idx); - let field_name = field.name().clone(); - if let DataType::Struct(children) = field.data_type() { + let field_name = field.name(); for child in children { - let child_name = child.name().clone(); + let child_name = child.name(); unnest_struct_columns.push(Column::new( sub_qualifier.cloned(), format!("{field_name}.{child_name}"), @@ -942,29 +959,21 @@ impl OptimizerRule for PushDownFilter { // Filter // Unnest Input (Projection) - let unnest_input = std::mem::take(&mut unnest.input); - - let filter_with_unnest_input = LogicalPlan::Filter(Filter::try_new( - conjunction(non_unnest_predicates).unwrap(), // Safe to unwrap since non_unnest_predicates is not empty. - unnest_input, - )?); - + // Safe to unwrap since non_unnest_predicates is not empty. + filter.predicate = conjunction(non_unnest_predicates).unwrap(); + filter.input = unnest.input; // Directly assign new filter plan as the new unnest's input. // The new filter plan will go through another rewrite pass since the rule itself // is applied recursively to all the child from top to down - let unnest_plan = - insert_below(LogicalPlan::Unnest(unnest), filter_with_unnest_input)?; - - match conjunction(unnest_predicates) { - None => Ok(unnest_plan), - Some(predicate) => Ok(Transformed::yes(LogicalPlan::Filter( - Filter::try_new(predicate, Arc::new(unnest_plan.data))?, - ))), - } + unnest.input = Arc::new(LogicalPlan::Filter(filter)); + Ok(Transformed::yes(with_filters( + unnest_predicates, + LogicalPlan::Unnest(unnest), + ))) } - LogicalPlan::Union(ref union) => { + LogicalPlan::Union(mut union) => { let mut inputs = Vec::with_capacity(union.inputs.len()); - for input in &union.inputs { + for input in union.inputs { let mut replace_map = HashMap::new(); for (i, (qualifier, field)) in input.schema().iter().enumerate() { let (union_qualifier, union_field) = @@ -977,72 +986,51 @@ impl OptimizerRule for PushDownFilter { let push_predicate = replace_cols_by_name(filter.predicate.clone(), &replace_map)?; - inputs.push(Arc::new(LogicalPlan::Filter(Filter::try_new( + inputs.push(Arc::new(LogicalPlan::Filter(Filter::new( push_predicate, - Arc::clone(input), - )?))) + input, + )))) } - Ok(Transformed::yes(LogicalPlan::Union(Union { - inputs, - schema: Arc::clone(&plan_schema), - }))) + + union.inputs = inputs; + Ok(Transformed::yes(LogicalPlan::Union(union))) } - LogicalPlan::Aggregate(agg) => { + LogicalPlan::Aggregate(mut agg) => { // We can push down Predicate which in groupby_expr. - let group_expr_columns = agg - .group_expr - .iter() - .map(|e| { - let (relation, name) = e.qualified_name(); - Column::new(relation, name) - }) - .collect::>(); + let group_expr_columns = expr_columns(&agg.group_expr); - let predicates = split_conjunction_owned(filter.predicate); + // As for plan Filter: Column(a+b) > 0 -- Agg: groupby:[Column(a)+Column(b)] + // After push, we need to replace `a+b` with Column(a)+Column(b) + // So we need create a replace_map, add {`a+b` --> Expr(Column(a)+Column(b))} + let mut replace_map = HashMap::new(); + for expr in &agg.group_expr { + replace_map.insert(expr.schema_name().to_string(), unalias(expr)); + } + let predicates = split_conjunction_owned(filter.predicate); let mut keep_predicates = vec![]; let mut push_predicates = vec![]; for expr in predicates { let cols = expr.column_refs(); if cols.iter().all(|c| group_expr_columns.contains(c)) { - push_predicates.push(expr); + push_predicates.push(replace_cols_by_name(expr, &replace_map)?); } else { keep_predicates.push(expr); } } - // As for plan Filter: Column(a+b) > 0 -- Agg: groupby:[Column(a)+Column(b)] - // After push, we need to replace `a+b` with Column(a)+Column(b) - // So we need create a replace_map, add {`a+b` --> Expr(Column(a)+Column(b))} - let mut replace_map = HashMap::new(); - for expr in &agg.group_expr { - replace_map.insert(expr.schema_name().to_string(), expr.clone()); - } - let replaced_push_predicates = push_predicates - .into_iter() - .map(|expr| replace_cols_by_name(expr, &replace_map)) - .collect::>>()?; - - let agg_input = Arc::clone(&agg.input); - Transformed::yes(LogicalPlan::Aggregate(agg)) - .transform_data(|new_plan| { - // If we have a filter to push, we push it down to the input of the aggregate - if let Some(predicate) = conjunction(replaced_push_predicates) { - let new_filter = make_filter(predicate, agg_input)?; - insert_below(new_plan, new_filter) - } else { - Ok(Transformed::no(new_plan)) - } - })? - .map_data(|child_plan| { - // if there are any remaining predicates we can't push, add them - // back as a filter - if let Some(predicate) = conjunction(keep_predicates) { - make_filter(predicate, Arc::new(child_plan)) - } else { - Ok(child_plan) - } - }) + // If we have a filter to push, we push it down to the input of the aggregate + let result = if let Some(predicate) = conjunction(push_predicates) { + filter.predicate = predicate; + filter.input = agg.input; + agg.input = Arc::new(LogicalPlan::Filter(filter)); + Transformed::yes(LogicalPlan::Aggregate(agg)) + } else { + Transformed::no(LogicalPlan::Aggregate(agg)) + }; + + // If there are any remaining predicates we can't push, add them back as a filter + result.map_data(|plan| Ok(with_filters(keep_predicates, plan))) } // Tries to push filters based on the partition key(s) of the window function(s) used. // Example: @@ -1054,22 +1042,16 @@ impl OptimizerRule for PushDownFilter { // Filter: (b > 1) and (c > 1) // Window: func() PARTITION BY [a] ... // Filter: (a > 1) - LogicalPlan::Window(window) => { + LogicalPlan::Window(mut window) => { // Retrieve the set of potential partition keys where we can push filters by. // Unlike aggregations, where there is only one statement per SELECT, there can be // multiple window functions, each with potentially different partition keys. // Therefore, we need to ensure that any potential partition key returned is used in // ALL window functions. Otherwise, filters cannot be pushed by through that column. - let extract_partition_keys = |func: &WindowFunction| { - func.params - .partition_by - .iter() - .map(|c| { - let (relation, name) = c.qualified_name(); - Column::new(relation, name) - }) - .collect::>() - }; + fn extract_partition_keys(func: &WindowFunction) -> HashSet { + expr_columns(&func.params.partition_by) + } + let potential_partition_keys = window .window_expr .iter() @@ -1119,31 +1101,22 @@ impl OptimizerRule for PushDownFilter { // place, so we can use `push_predicates` directly. This is consistent with other // optimizers, such as the one used by Postgres. - let window_input = Arc::clone(&window.input); - Transformed::yes(LogicalPlan::Window(window)) - .transform_data(|new_plan| { - // If we have a filter to push, we push it down to the input of the window - if let Some(predicate) = conjunction(push_predicates) { - let new_filter = make_filter(predicate, window_input)?; - insert_below(new_plan, new_filter) - } else { - Ok(Transformed::no(new_plan)) - } - })? - .map_data(|child_plan| { - // if there are any remaining predicates we can't push, add them - // back as a filter - if let Some(predicate) = conjunction(keep_predicates) { - make_filter(predicate, Arc::new(child_plan)) - } else { - Ok(child_plan) - } - }) + // If we have a filter to push, we push it down to the input of the aggregate + let result = if let Some(predicate) = conjunction(push_predicates) { + filter.predicate = predicate; + filter.input = window.input; + window.input = Arc::new(LogicalPlan::Filter(filter)); + Transformed::yes(LogicalPlan::Window(window)) + } else { + Transformed::no(LogicalPlan::Window(window)) + }; + + // If there are any remaining predicates we can't push, add them back as a filter + result.map_data(|plan| Ok(with_filters(keep_predicates, plan))) } - LogicalPlan::Join(join) => push_down_join(join, Some(&filter.predicate)), - LogicalPlan::TableScan(scan) => { + LogicalPlan::Join(join) => push_down_join(join, Some(filter.predicate)), + LogicalPlan::TableScan(mut scan) => { let filter_predicates = split_conjunction(&filter.predicate); - // Filters containing scalar subqueries cannot be pushed to // providers because the subquery result is not available // until execution time. @@ -1169,13 +1142,21 @@ impl OptimizerRule for PushDownFilter { non_volatile_filters.len() ); + if supported_filters + .iter() + .all(|res| res == &TableProviderFilterPushDown::Unsupported) + { + filter.input = Arc::new(LogicalPlan::TableScan(scan)); + return Ok(Transformed::no(LogicalPlan::Filter(filter))); + } + // Compose scan filters from non-volatile filters of `Exact` or `Inexact` pushdown type - let zip = non_volatile_filters.into_iter().zip(supported_filters); + let zip = non_volatile_filters.iter().zip(supported_filters.iter()); let new_scan_filters = zip .clone() - .filter(|(_, res)| res != &TableProviderFilterPushDown::Unsupported) - .map(|(pred, _)| pred); + .filter(|(_, res)| *res != &TableProviderFilterPushDown::Unsupported) + .map(|(&pred, _)| pred); // Add new scan filters let new_scan_filters: Vec = scan @@ -1186,28 +1167,31 @@ impl OptimizerRule for PushDownFilter { .cloned() .collect(); + if supported_filters + .iter() + .all(|res| res == &TableProviderFilterPushDown::Inexact) + && scan.filters == new_scan_filters + { + filter.input = Arc::new(LogicalPlan::TableScan(scan)); + return Ok(Transformed::no(LogicalPlan::Filter(filter))); + } else { + scan.filters = new_scan_filters; + } + // Compose predicates to be of `Unsupported` or `Inexact` pushdown type, // and also include volatile and subquery-containing filters let new_predicate: Vec = zip - .filter(|(_, res)| res != &TableProviderFilterPushDown::Exact) - .map(|(pred, _)| pred) + .filter(|(_, res)| *res != &TableProviderFilterPushDown::Exact) + .map(|(&pred, _)| pred) .chain(volatile_filters) .chain(subquery_filters) .cloned() .collect(); - let new_scan = LogicalPlan::TableScan(TableScan { - filters: new_scan_filters, - ..scan - }); - - Transformed::yes(new_scan).transform_data(|new_scan| { - if let Some(predicate) = conjunction(new_predicate) { - make_filter(predicate, Arc::new(new_scan)).map(Transformed::yes) - } else { - Ok(Transformed::no(new_scan)) - } - }) + Ok(Transformed::yes(with_filters( + new_predicate, + LogicalPlan::TableScan(scan), + ))) } LogicalPlan::Extension(extension_plan) => { // This check prevents the Filter from being removed when the extension node has no children, @@ -1222,17 +1206,16 @@ impl OptimizerRule for PushDownFilter { // determine if we can push any predicates down past the extension node // each element is true for push, false to keep - let predicate_push_or_keep = split_conjunction(&filter.predicate) - .iter() - .map(|expr| { - let cols = expr.column_refs(); - if cols.iter().any(|c| prevent_cols.contains(&c.name)) { - Ok(false) // No push (keep) - } else { - Ok(true) // push - } - }) - .collect::>>()?; + let predicate_push_or_keep: Vec = + split_conjunction(&filter.predicate) + .iter() + .map(|expr| { + !expr + .column_refs() + .iter() + .any(|c| prevent_cols.contains(&c.name)) + }) + .collect(); // all predicates are kept, no changes needed if predicate_push_or_keep.iter().all(|&x| !x) { @@ -1254,33 +1237,25 @@ impl OptimizerRule for PushDownFilter { } } - let new_children = match conjunction(push_predicates) { - Some(predicate) => extension_plan - .node - .inputs() - .into_iter() - .map(|child| { - Ok(LogicalPlan::Filter(Filter::try_new( - predicate.clone(), - Arc::new(child.clone()), - )?)) - }) - .collect::>>()?, - None => extension_plan.node.inputs().into_iter().cloned().collect(), - }; + // Unwrap - push_predicates is not empty, predicate_push_or_keep checked. + let predicate = conjunction(push_predicates).unwrap(); + let new_children = extension_plan + .node + .inputs() + .into_iter() + .map(|child| { + LogicalPlan::Filter(Filter::new( + predicate.clone(), + Arc::new(child.clone()), + )) + }) + .collect(); + // extension with new inputs. - let child_plan = LogicalPlan::Extension(extension_plan); - let new_extension = - child_plan.with_new_exprs(child_plan.expressions(), new_children)?; - - let new_plan = match conjunction(keep_predicates) { - Some(predicate) => LogicalPlan::Filter(Filter::try_new( - predicate, - Arc::new(new_extension), - )?), - None => new_extension, - }; - Ok(Transformed::yes(new_plan)) + let extension = LogicalPlan::Extension(extension_plan); + let new_plan = + extension.with_new_exprs(extension.expressions(), new_children)?; + Ok(Transformed::yes(with_filters(keep_predicates, new_plan))) } child => { filter.input = Arc::new(child); @@ -1320,22 +1295,19 @@ impl OptimizerRule for PushDownFilter { fn rewrite_projection( predicates: Vec, mut projection: Projection, -) -> Result<(Transformed, Option)> { +) -> Result<(Transformed, Vec)> { // Partition projection expressions into non-pushable vs pushable. // Non-pushable expressions are volatile (must not be duplicated) or // MoveTowardsLeafNodes (cheap expressions like get_field where re-inlining // into a filter causes optimizer instability — ExtractLeafExpressions will // undo the push-down, creating an infinite loop that runs until the // iteration limit is hit). - let (non_pushable_map, pushable_map): (HashMap<_, _>, HashMap<_, _>) = projection + let (non_pushable_map, pushable_map) = projection .schema .iter() .zip(projection.expr.iter()) .map(|((qualifier, field), expr)| { - // strip alias, as they should not be part of filters - let expr = expr.clone().unalias(); - - (qualified_name(qualifier, field.name()), expr) + (qualified_name(qualifier, field.name()), unalias(expr)) }) .partition(|(_, value)| { value.is_volatile() @@ -1352,67 +1324,30 @@ fn rewrite_projection( } } - match conjunction(push_predicates) { - Some(expr) => { - // re-write all filters based on this projection - // E.g. in `Filter: b\n Projection: a > 1 as b`, we can swap them, but the filter must be "a > 1" - let new_filter = LogicalPlan::Filter(Filter::try_new( - replace_cols_by_name(expr, &pushable_map)?, - std::mem::take(&mut projection.input), - )?); - - projection.input = Arc::new(new_filter); - - Ok(( - Transformed::yes(LogicalPlan::Projection(projection)), - conjunction(keep_predicates), - )) - } - None => Ok(( - Transformed::no(LogicalPlan::Projection(projection)), - conjunction(keep_predicates), - )), - } + let projection = if let Some(expr) = conjunction(push_predicates) { + // re-write all filters based on this projection + // E.g. in `Filter: b\n Projection: a > 1 as b`, we can swap them, but the filter must be "a > 1" + projection.input = Arc::new(LogicalPlan::Filter(Filter::new( + replace_cols_by_name(expr, &pushable_map)?, + projection.input, + ))); + + Transformed::yes(LogicalPlan::Projection(projection)) + } else { + Transformed::no(LogicalPlan::Projection(projection)) + }; + + Ok((projection, keep_predicates)) } /// Creates a new LogicalPlan::Filter node. +/// +/// Deprecated: use [`Filter::try_new`] directly. +#[deprecated] pub fn make_filter(predicate: Expr, input: Arc) -> Result { Filter::try_new(predicate, input).map(LogicalPlan::Filter) } -/// Replace the existing child of the single input node with `new_child`. -/// -/// Starting: -/// ```text -/// plan -/// child -/// ``` -/// -/// Ending: -/// ```text -/// plan -/// new_child -/// ``` -fn insert_below( - plan: LogicalPlan, - new_child: LogicalPlan, -) -> Result> { - let mut new_child = Some(new_child); - let transformed_plan = plan.map_children(|_child| { - if let Some(new_child) = new_child.take() { - Ok(Transformed::yes(new_child)) - } else { - // already took the new child - internal_err!("node had more than one input") - } - })?; - - // make sure we did the actual replacement - assert_or_internal_err!(new_child.is_none(), "node had no inputs"); - - Ok(transformed_plan) -} - impl PushDownFilter { #[expect(missing_docs)] pub fn new() -> Self { @@ -1439,41 +1374,64 @@ where /// replaces columns by its name on the projection. pub fn replace_cols_by_name( e: Expr, - replace_map: &HashMap, + replace_map: &HashMap>, ) -> Result { e.transform_up(|expr| { - Ok(if let Expr::Column(c) = &expr { - match replace_map.get(&c.flat_name()) { - Some(new_c) => Transformed::yes(new_c.clone()), - None => Transformed::no(expr), - } + if let Expr::Column(c) = &expr + && let Some(new_expr) = replace_map.get(&c.flat_name()) + { + Ok(Transformed::yes(new_expr.as_ref().clone())) } else { - Transformed::no(expr) - }) + Ok(Transformed::no(expr)) + } }) .data() } +/// Unalias expression reference. +fn unalias(expr: &Expr) -> &Expr { + if let Expr::Alias(alias) = expr { + unalias(&alias.expr) + } else { + expr + } +} + /// check whether the expression uses the columns in `check_map`. -fn contain(e: &Expr, check_map: &HashMap) -> bool { +fn contain(e: &Expr, check_map: &HashMap) -> bool { let mut is_contain = false; e.apply(|expr| { - Ok(if let Expr::Column(c) = &expr { - match check_map.get(&c.flat_name()) { - Some(_) => { - is_contain = true; - TreeNodeRecursion::Stop - } - None => TreeNodeRecursion::Continue, - } + if let Expr::Column(c) = &expr + && check_map.contains_key(&c.flat_name()) + { + is_contain = true; + Ok(TreeNodeRecursion::Stop) } else { - TreeNodeRecursion::Continue - }) + Ok(TreeNodeRecursion::Continue) + } }) .unwrap(); is_contain } +fn with_filters(predicates: Vec, plan: LogicalPlan) -> LogicalPlan { + if let Some(predicate) = conjunction(predicates) { + LogicalPlan::Filter(Filter::new(predicate, Arc::new(plan))) + } else { + plan + } +} + +fn expr_columns(exprs: &[Expr]) -> HashSet { + exprs + .iter() + .map(|expr| { + let (relation, name) = expr.qualified_name(); + Column::new(relation, name) + }) + .collect() +} + #[cfg(test)] mod tests { use std::cmp::Ordering; @@ -1487,9 +1445,9 @@ mod tests { use datafusion_expr::logical_plan::table_scan; use datafusion_expr::{ ColumnarValue, ExprFunctionExt, Extension, LogicalPlanBuilder, - ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TableSource, TableType, - UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, in_list, - in_subquery, lit, + ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, TableScan, TableSource, + TableType, UserDefinedLogicalNodeCore, Volatility, WindowFunctionDefinition, col, + in_list, in_subquery, lit, }; use crate::OptimizerContext; @@ -1536,6 +1494,17 @@ mod tests { }}; } + /// For testing that we don't return [Transformed::yes] when not necessary, + /// as it triggers rebuilding parent plan nodes. + macro_rules! assert_plan_not_transformed { + ($plan:expr) => {{ + let transformed = PushDownFilter::new() + .rewrite($plan, &OptimizerContext::new()) + .expect("failed to optimize plan"); + assert!(!transformed.transformed); + }}; + } + #[test] fn filter_before_projection() -> Result<()> { let table_scan = test_table_scan()?; @@ -1684,6 +1653,8 @@ mod tests { .aggregate(vec![col("a")], vec![sum(col("b")).alias("b")])? .filter(col("b").gt(lit(10i64)))? .build()?; + assert_plan_not_transformed!(plan.clone()); + // filter of aggregate is after aggregation since they are non-commutative assert_optimized_plan_equal!( plan, @@ -1876,6 +1847,7 @@ mod tests { .window(vec![window])? .filter(col("c").gt(lit(10i64)))? .build()?; + assert_plan_not_transformed!(plan.clone()); assert_optimized_plan_equal!( plan, @@ -3101,6 +3073,7 @@ mod tests { Some(filter), )? .build()?; + assert_plan_not_transformed!(plan.clone()); // not part of the test, just good to know: assert_snapshot!(plan, @@ -3165,6 +3138,7 @@ mod tests { projection, source: Arc::new(test_provider), fetch: None, + statistics_requests: std::collections::BTreeSet::new(), }); Ok(LogicalPlanBuilder::from(table_scan)) @@ -3207,15 +3181,16 @@ mod tests { let plan = table_scan_with_pushdown_provider(TableProviderFilterPushDown::Inexact)?; - let optimized_plan = PushDownFilter::new() + let optimized = PushDownFilter::new() .rewrite(plan, &OptimizerContext::new()) - .expect("failed to optimize plan") - .data; + .expect("failed to optimize plan"); + assert!(optimized.transformed); + assert_plan_not_transformed!(optimized.data.clone()); // Optimizing the same plan multiple times should produce the same plan // each time. assert_optimized_plan_equal!( - optimized_plan, + optimized.data, @r" Filter: a = Int64(1) TableScan: test, partial_filters=[a = Int64(1)] @@ -3227,6 +3202,7 @@ mod tests { fn filter_with_table_provider_unsupported() -> Result<()> { let plan = table_scan_with_pushdown_provider(TableProviderFilterPushDown::Unsupported)?; + assert_plan_not_transformed!(plan.clone()); assert_optimized_plan_equal!( plan, @@ -3861,6 +3837,51 @@ mod tests { ) } + /// Regression test: for a null-aware LeftAnti join (the shape produced by + /// `NOT IN` with a nullable subquery), a right-side predicate must NOT be + /// inferred onto the join. Inference would push a null-rejecting predicate + /// to the subquery side, dropping its NULL rows and breaking the + /// three-valued `NOT IN` semantics. + #[test] + fn null_aware_left_anti_join_no_inferred_pushdown() -> Result<()> { + let table_scan = test_table_scan_with_name("test1")?; + let left = LogicalPlanBuilder::from(table_scan) + .project(vec![col("a"), col("b")])? + .build()?; + let right_table_scan = test_table_scan_with_name("test2")?; + let right = LogicalPlanBuilder::from(right_table_scan) + .project(vec![col("a"), col("b")])? + .build()?; + let plan = LogicalPlanBuilder::from(left) + .join_detailed_with_options( + right, + JoinType::LeftAnti, + ( + vec![Column::from_qualified_name("test1.a")], + vec![Column::from_qualified_name("test2.a")], + ), + None, + datafusion_common::NullEquality::NullEqualsNothing, + true, + )? + .filter(col("test1.a").gt(lit(2u32)))? + .build()?; + + // The left-side filter is pushed to the left input, but — unlike the + // non-null-aware `left_anti_join` test — no `test2.a > 2` predicate is + // inferred onto the right/subquery side. + assert_optimized_plan_equal!( + plan, + @r" + LeftAnti Join: test1.a = test2.a null_aware + Projection: test1.a, test1.b + TableScan: test1, full_filters=[test1.a > UInt32(2)] + Projection: test2.a, test2.b + TableScan: test2 + " + ) + } + #[test] fn left_anti_join_with_filters() -> Result<()> { let table_scan = test_table_scan_with_name("test1")?; @@ -4217,7 +4238,7 @@ mod tests { plan, @r" Projection: a, b - Filter: t.a > Int32(5) AND t.b > Int32(10) AND TestScalarUDF() > Float64(0.1) + Filter: TestScalarUDF() > Float64(0.1) AND t.a > Int32(5) AND t.b > Int32(10) TableScan: test " ) diff --git a/datafusion/optimizer/src/replace_distinct_aggregate.rs b/datafusion/optimizer/src/replace_distinct_aggregate.rs index 06df61e766615..cc2616379057a 100644 --- a/datafusion/optimizer/src/replace_distinct_aggregate.rs +++ b/datafusion/optimizer/src/replace_distinct_aggregate.rs @@ -22,7 +22,7 @@ use crate::{OptimizerConfig, OptimizerRule}; use std::sync::Arc; use datafusion_common::tree_node::Transformed; -use datafusion_common::{Column, Result}; +use datafusion_common::{Column, Dependency, Result}; use datafusion_expr::expr_rewriter::normalize_cols; use datafusion_expr::utils::expand_wildcard; use datafusion_expr::{Aggregate, Distinct, DistinctOn, Expr, LogicalPlan}; @@ -101,9 +101,14 @@ impl OptimizerRule for ReplaceDistinctWithAggregate { let field_count = input.schema().fields().len(); for dep in input.schema().functional_dependencies().iter() { - // If distinct is exactly the same with a previous GROUP BY, we can - // simply remove it: - if dep.source_indices.len() >= field_count + // If the input is already unique on all of its columns (e.g. + // it is a GROUP BY over exactly these columns), the DISTINCT + // is a no-op and we can simply remove it. The dependency mode + // must be `Single`: a `Multi` dependence (e.g. a former key + // downgraded by a join) means equal rows may occur multiple + // times, so the DISTINCT still has work to do. + if dep.mode == Dependency::Single + && dep.source_indices.len() >= field_count && dep.source_indices[..field_count] .iter() .enumerate() diff --git a/datafusion/optimizer/src/rewrite_set_comparison.rs b/datafusion/optimizer/src/rewrite_set_comparison.rs index c8c35b518743a..18712c5335205 100644 --- a/datafusion/optimizer/src/rewrite_set_comparison.rs +++ b/datafusion/optimizer/src/rewrite_set_comparison.rs @@ -25,7 +25,7 @@ use datafusion_common::{Column, DFSchema, ExprSchema, Result, ScalarValue, plan_ use datafusion_expr::expr::{self, Exists, SetComparison, SetQuantifier}; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::logical_plan::builder::LogicalPlanBuilder; -use datafusion_expr::{Expr, LogicalPlan, lit}; +use datafusion_expr::{DmlStatement, Expr, LogicalPlan, WriteOp, lit}; use std::sync::Arc; use datafusion_expr::utils::merge_schema; @@ -44,7 +44,19 @@ impl RewriteSetComparison { } fn rewrite_plan(&self, plan: LogicalPlan) -> Result> { - let schema = merge_schema(&plan.inputs()); + let mut schema = merge_schema(&plan.inputs()); + if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + schema.merge(&DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?); + } plan.map_expressions(|expr| { expr.transform_up(|expr| rewrite_set_comparison(expr, &schema)) }) diff --git a/datafusion/optimizer/src/scalar_subquery_to_join.rs b/datafusion/optimizer/src/scalar_subquery_to_join.rs index 76d22c7fb374b..44011a125ba96 100644 --- a/datafusion/optimizer/src/scalar_subquery_to_join.rs +++ b/datafusion/optimizer/src/scalar_subquery_to_join.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -//! [`ScalarSubqueryToJoin`] rewriting correlated scalar subquery filters to `JOIN`s +//! [`ScalarSubqueryToJoin`] rewriting scalar subquery filters to `JOIN`s use std::collections::{BTreeSet, HashMap}; use std::sync::Arc; @@ -34,11 +34,16 @@ use datafusion_common::{Column, Result, ScalarValue, assert_or_internal_err, pla use datafusion_expr::expr_rewriter::create_col_from_scalar_expr; use datafusion_expr::logical_plan::{JoinType, Subquery}; use datafusion_expr::utils::conjunction; -use datafusion_expr::{EmptyRelation, Expr, LogicalPlan, LogicalPlanBuilder, expr}; +use datafusion_expr::{Expr, LogicalPlan, LogicalPlanBuilder, lit, not, when}; -/// Optimizer rule that rewrites correlated scalar subquery filters to joins and -/// places an additional projection on top of the filter, to preserve the -/// original schema. +/// Optimizer rule that rewrites scalar subquery filters to joins and places an +/// additional projection on top of the filter to preserve the original schema. +/// +/// When [`datafusion_common::config::OptimizerOptions::enable_physical_uncorrelated_scalar_subquery`] is +/// true (the default), only *correlated* scalar subqueries are rewritten here; +/// uncorrelated ones are left for physical execution via `ScalarSubqueryExec`. +/// When the option is false, all scalar subqueries — correlated and +/// uncorrelated — are rewritten to left joins by this rule. #[derive(Default, Debug)] pub struct ScalarSubqueryToJoin {} @@ -63,10 +68,12 @@ impl ScalarSubqueryToJoin { &self, predicate: &Expr, alias_gen: &Arc, + physical_uncorrelated: bool, ) -> Result<(Vec<(Subquery, String)>, Expr)> { let mut extract = ExtractScalarSubQuery { sub_query_info: vec![], alias_gen, + physical_uncorrelated, }; predicate .clone() @@ -88,15 +95,23 @@ impl OptimizerRule for ScalarSubqueryToJoin { ) -> Result> { match plan { LogicalPlan::Filter(filter) => { + let physical_uncorrelated = config + .options() + .optimizer + .enable_physical_uncorrelated_scalar_subquery; // Optimization: skip the rest of the rule and its copies if - // there are no scalar subqueries - if !contains_correlated_scalar_subquery(&filter.predicate) { + // there are no scalar subqueries this rule should rewrite + if !contains_scalar_subquery_to_rewrite( + &filter.predicate, + physical_uncorrelated, + ) { return Ok(Transformed::no(LogicalPlan::Filter(filter))); } let (subqueries, mut rewrite_expr) = self.extract_subquery_exprs( &filter.predicate, config.alias_generator(), + physical_uncorrelated, )?; assert_or_internal_err!( @@ -107,18 +122,17 @@ impl OptimizerRule for ScalarSubqueryToJoin { // iterate through all subqueries in predicate, turning each into a left join let mut cur_input = filter.input.as_ref().clone(); for (subquery, alias) in subqueries { - if let Some((optimized_subquery, expr_check_map)) = + if let Some((optimized_subquery, compensation_exprs)) = build_join(&subquery, &cur_input, &alias)? { - if !expr_check_map.is_empty() { + if !compensation_exprs.is_empty() { rewrite_expr = rewrite_expr .transform_up(|expr| { - // replace column references with entry in map, if it exists - if let Some(map_expr) = expr + if let Some(compensation_expr) = expr .try_as_col() - .and_then(|col| expr_check_map.get(col)) + .and_then(|col| compensation_exprs.get(col)) { - Ok(Transformed::yes(map_expr.clone())) + Ok(Transformed::yes(compensation_expr.clone())) } else { Ok(Transformed::no(expr)) } @@ -142,13 +156,15 @@ impl OptimizerRule for ScalarSubqueryToJoin { Ok(Transformed::yes(new_plan)) } LogicalPlan::Projection(projection) => { + let physical_uncorrelated = config + .options() + .optimizer + .enable_physical_uncorrelated_scalar_subquery; // Optimization: skip the rest of the rule and its copies if there - // are no correlated scalar subqueries - if !projection - .expr - .iter() - .any(contains_correlated_scalar_subquery) - { + // are no scalar subqueries this rule should rewrite + if !projection.expr.iter().any(|expr| { + contains_scalar_subquery_to_rewrite(expr, physical_uncorrelated) + }) { return Ok(Transformed::no(LogicalPlan::Projection(projection))); } @@ -157,8 +173,11 @@ impl OptimizerRule for ScalarSubqueryToJoin { let mut rewrite_exprs: Vec = Vec::with_capacity(projection.expr.len()); for (idx, expr) in projection.expr.iter().enumerate() { - let (subqueries, rewrite_expr) = - self.extract_subquery_exprs(expr, config.alias_generator())?; + let (subqueries, rewrite_expr) = self.extract_subquery_exprs( + expr, + config.alias_generator(), + physical_uncorrelated, + )?; for (_, alias) in &subqueries { alias_to_index.insert(alias.clone(), idx); } @@ -172,22 +191,21 @@ impl OptimizerRule for ScalarSubqueryToJoin { // iterate through all subqueries in predicate, turning each into a left join let mut cur_input = projection.input.as_ref().clone(); for (subquery, alias) in all_subqueries { - if let Some((optimized_subquery, expr_check_map)) = + if let Some((optimized_subquery, compensation_exprs)) = build_join(&subquery, &cur_input, &alias)? { cur_input = optimized_subquery; - if !expr_check_map.is_empty() + if !compensation_exprs.is_empty() && let Some(&idx) = alias_to_index.get(&alias) { let new_expr = rewrite_exprs[idx] .clone() .transform_up(|expr| { - // replace column references with entry in map, if it exists - if let Some(map_expr) = expr + if let Some(compensation_expr) = expr .try_as_col() - .and_then(|col| expr_check_map.get(col)) + .and_then(|col| compensation_exprs.get(col)) { - Ok(Transformed::yes(map_expr.clone())) + Ok(Transformed::yes(compensation_expr.clone())) } else { Ok(Transformed::no(expr)) } @@ -230,12 +248,20 @@ impl OptimizerRule for ScalarSubqueryToJoin { } } -/// Returns true if the expression contains a correlated scalar subquery, false -/// otherwise. Uncorrelated scalar subqueries are handled by the physical -/// planner via `ScalarSubqueryExec` and do not need to be converted to joins. -fn contains_correlated_scalar_subquery(expr: &Expr) -> bool { +/// Returns true if the expression contains a scalar subquery that this rule +/// should rewrite to a join. +/// +/// When `enable_physical_uncorrelated_scalar_subquery` is true (the default) only +/// correlated scalar subqueries are rewritten — uncorrelated ones are handled +/// by the physical planner via `ScalarSubqueryExec`. When it is false, all +/// scalar subqueries (correlated and uncorrelated) are rewritten. +fn contains_scalar_subquery_to_rewrite(expr: &Expr, physical_uncorrelated: bool) -> bool { expr.exists(|expr| { - Ok(matches!(expr, Expr::ScalarSubquery(sq) if !sq.outer_ref_columns.is_empty())) + Ok(matches!( + expr, + Expr::ScalarSubquery(sq) + if !physical_uncorrelated || !sq.outer_ref_columns.is_empty() + )) }) .expect("Inner is always Ok") } @@ -243,6 +269,7 @@ fn contains_correlated_scalar_subquery(expr: &Expr) -> bool { struct ExtractScalarSubQuery<'a> { sub_query_info: Vec<(Subquery, String)>, alias_gen: &'a Arc, + physical_uncorrelated: bool, } impl TreeNodeRewriter for ExtractScalarSubQuery<'_> { @@ -250,9 +277,13 @@ impl TreeNodeRewriter for ExtractScalarSubQuery<'_> { fn f_down(&mut self, expr: Expr) -> Result> { match expr { - // Skip uncorrelated scalar subqueries + // Match scalar subqueries this rule should rewrite to a join. When + // `physical_uncorrelated` is true, only correlated subqueries are + // rewritten — uncorrelated ones are handled later by the physical + // planner. When false, both are rewritten. Expr::ScalarSubquery(ref subquery) - if !subquery.outer_ref_columns.is_empty() => + if !self.physical_uncorrelated + || !subquery.outer_ref_columns.is_empty() => { let subquery = subquery.clone(); let scalar_expr = subquery @@ -285,90 +316,100 @@ impl TreeNodeRewriter for ExtractScalarSubQuery<'_> { /// /// ```text /// select c.id from customers c -/// left join (select c_id, avg(total) as val from orders group by c_id) o on o.c_id = c.c_id -/// where c.balance > o.val +/// left join (select c_id, avg(total) from orders group by c_id) o +/// on o.c_id = c.id +/// where c.balance > o."avg(total)" /// ``` /// -/// Or a query like: -/// -/// ```text -/// select id from customers where balance > -/// (select avg(total) from orders) -/// ``` -/// -/// and optimizes it into: -/// -/// ```text -/// select c.id from customers c -/// left join (select avg(total) as val from orders) a -/// where c.balance > a.val -/// ``` +/// When [`datafusion_common::config::OptimizerOptions::enable_physical_uncorrelated_scalar_subquery`] is +/// false, this function also handles uncorrelated scalar subqueries, rewriting +/// them as a `Left Join: Filter: Boolean(true)` instead of leaving them for +/// `ScalarSubqueryExec`. /// /// # Arguments /// -/// * `query_info` - The subquery portion of the `where` (select avg(total) from orders) -/// * `filter_input` - The non-subquery portion (from customers) -/// * `outer_others` - Any additional parts to the `where` expression (and c.x = y) -/// * `subquery_alias` - Subquery aliases +/// * `subquery` - The scalar subquery to rewrite (correlated, or uncorrelated +/// when `enable_physical_uncorrelated_scalar_subquery` is false). +/// * `outer_input` - The outer plan that the decorrelated subquery is +/// left-joined onto — the input of the `Filter` or `Projection` node +/// that contained the subquery. +/// * `subquery_alias` - The unique alias assigned to the decorrelated +/// subquery; used both to qualify the join condition and to produce +/// column references for the caller to substitute. +/// +/// Returns `Ok(None)` if the subquery cannot be decorrelated. On success, +/// returns the rewritten outer plan and a map from each count-bug-affected +/// column to its `CASE WHEN __always_true IS NULL THEN ... END` compensation +/// expression, which the caller must substitute into any expression that +/// references those columns. fn build_join( subquery: &Subquery, - filter_input: &LogicalPlan, + outer_input: &LogicalPlan, subquery_alias: &str, ) -> Result)>> { + // `build_join` also handles uncorrelated scalar subqueries (as a left + // join with `Boolean(true)`) when the + // `enable_physical_uncorrelated_scalar_subquery` option is disabled. let subquery_plan = subquery.subquery.as_ref(); let mut pull_up = PullUpCorrelatedExpr::new().with_need_handle_count_bug(true); - let new_plan = subquery_plan.clone().rewrite(&mut pull_up).data()?; + let decorrelated_subquery = subquery_plan.clone().rewrite(&mut pull_up).data()?; if !pull_up.can_pull_up { return Ok(None); } - let collected_count_expr_map = - pull_up.collected_count_expr_map.get(&new_plan).cloned(); - let sub_query_alias = LogicalPlanBuilder::from(new_plan) + let collected_count_expr_map = pull_up + .collected_count_expr_map + .get(&decorrelated_subquery) + .cloned(); + let aliased_subquery = LogicalPlanBuilder::from(decorrelated_subquery) .alias(subquery_alias.to_string())? .build()?; - let mut all_correlated_cols = BTreeSet::new(); - pull_up + let all_correlated_cols: BTreeSet = pull_up .correlated_subquery_cols_map .values() - .for_each(|cols| all_correlated_cols.extend(cols.clone())); + .flatten() + .cloned() + .collect(); - // alias the join filter + // Correlated columns now live in the decorrelated subquery's output, + // so re-qualify them with the subquery alias. let join_filter_opt = conjunction(pull_up.join_filters).map_or(Ok(None), |filter| { replace_qualified_name(filter, &all_correlated_cols, subquery_alias).map(Some) })?; - // join our sub query into the main plan - let new_plan = if join_filter_opt.is_none() { - match filter_input { - LogicalPlan::EmptyRelation(EmptyRelation { - produce_one_row: true, - schema: _, - }) => sub_query_alias, - _ => { - // if not correlated, group down to 1 row and left join on that (preserving row count) - LogicalPlanBuilder::from(filter_input.clone()) - .join_on( - sub_query_alias, - JoinType::Left, - vec![Expr::Literal(ScalarValue::Boolean(Some(true)), None)], - )? - .build()? - } - } - } else { - // left join if correlated, grouping by the join keys so we don't change row count - LogicalPlanBuilder::from(filter_input.clone()) - .join_on(sub_query_alias, JoinType::Left, join_filter_opt)? - .build()? - }; - let mut computation_project_expr = HashMap::new(); + // When pull-up did not extract any usable join keys (a correlated subquery + // whose predicate references only outer columns), fall back to `ON true`: + // the decorrelated subquery still yields at most one row per outer row + // because its aggregate is grouped by the (empty) set of correlated inner + // columns. + let join_filter = join_filter_opt.or_else(|| Some(lit(true))); + + let new_plan = LogicalPlanBuilder::from(outer_input.clone()) + .join_on(aliased_subquery, JoinType::Left, join_filter)? + .build()?; + + // Add count-bug compensation for each of the subquery's projected + // expressions that yield non-NULL values on empty input. We wrap each + // such expression in a CASE that substitutes the empty-input value + // when the LEFT JOIN produced synthetic right-side NULLs (no inner + // row matched), and uses the actual right-side value (which may + // itself be NULL) otherwise. + let mut compensation_exprs = HashMap::new(); if let Some(expr_map) = collected_count_expr_map { + let mut expr_rewrite = TypeCoercionRewriter { + schema: new_plan.schema(), + }; + let having_arm = pull_up + .pull_up_having_expr + .as_ref() + .map(|f| (not(f.clone()), lit(ScalarValue::Null))); for (name, result) in expr_map { if evaluates_to_null(result.clone(), result.column_refs())? { - // If expr always returns null when column is null, skip processing + // Aggregates whose empty-input value is NULL (max/min/sum/…) + // need no compensation: the LEFT JOIN already produces NULL + // for unmatched outer rows. continue; } @@ -376,42 +417,21 @@ fn build_join( Column::new(Some(subquery_alias), UN_MATCHED_ROW_INDICATOR); // Qualify with the subquery alias to avoid ambiguity when the // outer table has a column with the same name as the aggregate. - let value_col = Column::new(Some(subquery_alias), name.clone()); - - let computer_expr = if let Some(filter) = &pull_up.pull_up_having_expr { - Expr::Case(expr::Case { - expr: None, - when_then_expr: vec![ - ( - Box::new(Expr::IsNull(Box::new(Expr::Column(indicator_col)))), - Box::new(result), - ), - ( - Box::new(Expr::Not(Box::new(filter.clone()))), - Box::new(Expr::Literal(ScalarValue::Null, None)), - ), - ], - else_expr: Some(Box::new(Expr::Column(value_col.clone()))), - }) - } else { - Expr::Case(expr::Case { - expr: None, - when_then_expr: vec![( - Box::new(Expr::IsNull(Box::new(Expr::Column(indicator_col)))), - Box::new(result), - )], - else_expr: Some(Box::new(Expr::Column(value_col.clone()))), - }) - }; - let mut expr_rewrite = TypeCoercionRewriter { - schema: new_plan.schema(), - }; - computation_project_expr - .insert(value_col, computer_expr.rewrite(&mut expr_rewrite).data()?); + let value_col = Column::new(Some(subquery_alias), name); + + let mut builder = when(Expr::Column(indicator_col).is_null(), result); + if let Some((when_expr, then_expr)) = &having_arm { + builder = builder.when(when_expr.clone(), then_expr.clone()); + } + let compensation_expr = builder.otherwise(Expr::Column(value_col.clone()))?; + compensation_exprs.insert( + value_col, + compensation_expr.rewrite(&mut expr_rewrite).data()?, + ); } } - Ok(Some((new_plan, computation_project_expr))) + Ok(Some((new_plan, compensation_exprs))) } #[cfg(test)] @@ -425,7 +445,7 @@ mod tests { use datafusion_expr::test::function_stub::sum; use crate::assert_optimized_plan_eq_display_indent_snapshot; - use datafusion_expr::{Between, col, lit, out_ref_col, scalar_subquery}; + use datafusion_expr::{Between, col, expr, out_ref_col, scalar_subquery}; use datafusion_functions_aggregate::min_max::{max, min}; macro_rules! assert_optimized_plan_equal { @@ -837,7 +857,7 @@ mod tests { assert_optimized_plan_equal!( plan, @r#" - Projection: customer.c_custkey, CASE WHEN __scalar_sq_1.__always_true IS NULL THEN CASE WHEN CAST(NULL AS Boolean) THEN Utf8("a") ELSE Utf8("b") END ELSE __scalar_sq_1.CASE WHEN max(orders.o_totalprice) THEN Utf8("a") ELSE Utf8("b") END END AS CASE WHEN max(orders.o_totalprice) THEN Utf8("a") ELSE Utf8("b") END [c_custkey:Int64, CASE WHEN max(orders.o_totalprice) THEN Utf8("a") ELSE Utf8("b") END:Utf8;N] + Projection: customer.c_custkey, CASE WHEN __scalar_sq_1.__always_true IS NULL THEN CASE WHEN CAST(Float64(NULL) AS Boolean) THEN Utf8("a") ELSE Utf8("b") END ELSE __scalar_sq_1.CASE WHEN max(orders.o_totalprice) THEN Utf8("a") ELSE Utf8("b") END END AS CASE WHEN max(orders.o_totalprice) THEN Utf8("a") ELSE Utf8("b") END [c_custkey:Int64, CASE WHEN max(orders.o_totalprice) THEN Utf8("a") ELSE Utf8("b") END:Utf8;N] Left Join: Filter: customer.c_custkey = __scalar_sq_1.o_custkey [c_custkey:Int64, c_name:Utf8, CASE WHEN max(orders.o_totalprice) THEN Utf8("a") ELSE Utf8("b") END:Utf8;N, o_custkey:Int64;N, __always_true:Boolean;N] TableScan: customer [c_custkey:Int64, c_name:Utf8] SubqueryAlias: __scalar_sq_1 [CASE WHEN max(orders.o_totalprice) THEN Utf8("a") ELSE Utf8("b") END:Utf8, o_custkey:Int64, __always_true:Boolean] @@ -1177,4 +1197,52 @@ mod tests { " ) } + + #[test] + fn uncorrelated_scalar_subquery_rewritten_when_flag_off() -> Result<()> { + use datafusion_common::config::ConfigOptions; + + let sq = Arc::new( + LogicalPlanBuilder::from(scan_tpch_table("orders")) + .aggregate(Vec::::new(), vec![max(col("orders.o_custkey"))])? + .project(vec![max(col("orders.o_custkey"))])? + .build()?, + ); + + let plan = LogicalPlanBuilder::from(scan_tpch_table("customer")) + .filter(col("customer.c_custkey").eq(scalar_subquery(sq)))? + .project(vec![col("customer.c_custkey")])? + .build()?; + + let mut options = ConfigOptions::default(); + options + .optimizer + .enable_physical_uncorrelated_scalar_subquery = false; + let context = crate::OptimizerContext::new_with_config_options(Arc::new(options)); + + let rule: Arc = + Arc::new(ScalarSubqueryToJoin::new()); + let optimizer = crate::Optimizer::with_rules(vec![rule]); + let optimized_plan = optimizer + .optimize(plan, &context, |_, _| {}) + .expect("failed to optimize plan"); + let formatted_plan = optimized_plan.display_indent_schema(); + + insta::assert_snapshot!( + formatted_plan, + @r" + Projection: customer.c_custkey [c_custkey:Int64] + Projection: customer.c_custkey, customer.c_name [c_custkey:Int64, c_name:Utf8] + Filter: customer.c_custkey = __scalar_sq_1.max(orders.o_custkey) [c_custkey:Int64, c_name:Utf8, max(orders.o_custkey):Int64;N] + Left Join: Filter: Boolean(true) [c_custkey:Int64, c_name:Utf8, max(orders.o_custkey):Int64;N] + TableScan: customer [c_custkey:Int64, c_name:Utf8] + SubqueryAlias: __scalar_sq_1 [max(orders.o_custkey):Int64;N] + Projection: max(orders.o_custkey) [max(orders.o_custkey):Int64;N] + Aggregate: groupBy=[[]], aggr=[[max(orders.o_custkey)]] [max(orders.o_custkey):Int64;N] + TableScan: orders [o_orderkey:Int64, o_custkey:Int64, o_orderstatus:Utf8, o_totalprice:Float64;N] + " + ); + + Ok(()) + } } diff --git a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs index 143d8eae695af..2b606687d47a3 100644 --- a/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs +++ b/datafusion/optimizer/src/simplify_expressions/expr_simplifier.rs @@ -40,6 +40,7 @@ use datafusion_common::{ tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter}, }; use datafusion_expr::expr::HigherOrderFunction; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ BinaryExpr, Case, ColumnarValue, Expr, ExprSchemable, Like, Operator, Volatility, and, binary::BinaryTypeCoercer, lit, or, preimage::PreimageResult, @@ -707,11 +708,15 @@ impl ConstEvaluator { return ConstSimplifyResult::NotSimplified(s, m); } - let phys_expr = - match create_physical_expr(&expr, &DUMMY_DF_SCHEMA, &self.execution_props) { - Ok(e) => e, - Err(err) => return ConstSimplifyResult::SimplifyRuntimeError(err, expr), - }; + let phys_expr = match create_physical_expr( + &expr, + &DUMMY_DF_SCHEMA, + &self.execution_props, + &PhysicalPlanningContext::default(), + ) { + Ok(e) => e, + Err(err) => return ConstSimplifyResult::SimplifyRuntimeError(err, expr), + }; let metadata = phys_expr .return_field(DUMMY_BATCH.schema_ref()) .ok() @@ -2541,6 +2546,36 @@ mod tests { assert_eq!(simplify(expr_b), expected_b); } + /// `c3_non_null IN (SELECT a FROM t)`, where `a` has the given nullability. + fn in_subquery_expr(a_nullable: bool) -> Expr { + let schema = Schema::new(vec![Field::new("a", DataType::Int64, a_nullable)]); + let source = Arc::new(LogicalTableSource::new(Arc::new(schema))); + let subquery = LogicalPlanBuilder::scan("t", source, None) + .unwrap() + .project(vec![col("a")]) + .unwrap() + .build() + .unwrap(); + + in_subquery(col("c3_non_null"), Arc::new(subquery)) + } + + #[test] + fn test_simplify_eq_not_self_in_subquery() { + // `expr_a`: even though `c3_non_null` is non-nullable, the `IN` evaluates to NULL + // when `c3_non_null` matches no row and the subquery's `a` contains a NULL. So the + // expression is nullable and `A = A` must not fold to `true`. + let expr_a = in_subquery_expr(true); + let expected_a = expr_a.clone().is_not_null().or(lit_bool_null()); + + // `expr_b`: neither side can be NULL, so the `IN` is non-nullable and `A = A` is true. + let expr_b = in_subquery_expr(false); + let expected_b = lit(true); + + assert_eq!(simplify(expr_a.clone().eq(expr_a)), expected_a); + assert_eq!(simplify(expr_b.clone().eq(expr_b)), expected_b); + } + #[test] fn test_simplify_or_true() { let expr_a = col("c2").or(lit(true)); @@ -2919,6 +2954,21 @@ mod tests { } } + #[test] + fn test_simplify_concat_by_null() { + let null = Expr::Literal(ScalarValue::Utf8(None), None); + // A || null --> null + { + let expr = binary_expr(col("c1"), Operator::StringConcat, null.clone()); + assert_eq!(simplify(expr), null); + } + // null || A --> null + { + let expr = binary_expr(null.clone(), Operator::StringConcat, col("c1")); + assert_eq!(simplify(expr), null); + } + } + #[test] fn test_simplify_composed_bitwise_and() { // ((c2 > 5) & (c1 < 6)) & (c2 > 5) --> (c2 > 5) & (c1 < 6) @@ -3047,17 +3097,6 @@ mod tests { #[test] fn test_simplify_negated_bitwise_and() { - // !c4 & c4 --> 0 - let expr = (-col("c4_non_null")) & col("c4_non_null"); - let expected = lit(0u32); - - assert_eq!(simplify(expr), expected); - // c4 & !c4 --> 0 - let expr = col("c4_non_null") & (-col("c4_non_null")); - let expected = lit(0u32); - - assert_eq!(simplify(expr), expected); - // !c3 & c3 --> 0 let expr = (-col("c3_non_null")) & col("c3_non_null"); let expected = lit(0i64); @@ -3072,18 +3111,6 @@ mod tests { #[test] fn test_simplify_negated_bitwise_or() { - // !c4 | c4 --> -1 - let expr = (-col("c4_non_null")) | col("c4_non_null"); - let expected = lit(-1i32); - - assert_eq!(simplify(expr), expected); - - // c4 | !c4 --> -1 - let expr = col("c4_non_null") | (-col("c4_non_null")); - let expected = lit(-1i32); - - assert_eq!(simplify(expr), expected); - // !c3 | c3 --> -1 let expr = (-col("c3_non_null")) | col("c3_non_null"); let expected = lit(-1i64); @@ -3099,18 +3126,6 @@ mod tests { #[test] fn test_simplify_negated_bitwise_xor() { - // !c4 ^ c4 --> -1 - let expr = (-col("c4_non_null")) ^ col("c4_non_null"); - let expected = lit(-1i32); - - assert_eq!(simplify(expr), expected); - - // c4 ^ !c4 --> -1 - let expr = col("c4_non_null") ^ (-col("c4_non_null")); - let expected = lit(-1i32); - - assert_eq!(simplify(expr), expected); - // !c3 ^ c3 --> -1 let expr = (-col("c3_non_null")) ^ col("c3_non_null"); let expected = lit(-1i64); @@ -3538,6 +3553,32 @@ mod tests { assert_no_change(regex_match(col("c1"), lit("foo|bar|baz|blarg|bozo|etc"))); } + #[test] + fn test_simplify_not_regex_match() { + let pattern = || lit("foo.*"); + + // NOT (c1 ~ pattern) --> c1 !~ pattern + assert_eq!( + simplify(regex_match(col("c1"), pattern()).not()), + regex_not_match(col("c1"), pattern()), + ); + // NOT (c1 !~ pattern) --> c1 ~ pattern + assert_eq!( + simplify(regex_not_match(col("c1"), pattern()).not()), + regex_match(col("c1"), pattern()), + ); + // NOT (c1 ~* pattern) --> c1 !~* pattern + assert_eq!( + simplify(regex_imatch(col("c1"), pattern()).not()), + regex_not_imatch(col("c1"), pattern()), + ); + // NOT (c1 !~* pattern) --> c1 ~* pattern + assert_eq!( + simplify(regex_not_imatch(col("c1"), pattern()).not()), + regex_imatch(col("c1"), pattern()), + ); + } + #[track_caller] fn assert_no_change(expr: Expr) { let optimized = simplify(expr.clone()); diff --git a/datafusion/optimizer/src/simplify_expressions/mod.rs b/datafusion/optimizer/src/simplify_expressions/mod.rs index 89c79d3fb4203..e0b53b79d468c 100644 --- a/datafusion/optimizer/src/simplify_expressions/mod.rs +++ b/datafusion/optimizer/src/simplify_expressions/mod.rs @@ -22,6 +22,7 @@ pub mod expr_simplifier; mod inlist_simplifier; mod linear_aggregates; mod regex; +mod reorder_predicates; pub mod simplify_exprs; pub mod simplify_literal; mod simplify_predicates; @@ -33,6 +34,7 @@ mod utils; pub use datafusion_expr::simplify::SimplifyContext; pub use expr_simplifier::*; +pub(crate) use reorder_predicates::reorder_predicates; pub use simplify_exprs::*; pub use simplify_predicates::simplify_predicates; diff --git a/datafusion/optimizer/src/simplify_expressions/regex.rs b/datafusion/optimizer/src/simplify_expressions/regex.rs index b341c328e992a..f04d9476c42fe 100644 --- a/datafusion/optimizer/src/simplify_expressions/regex.rs +++ b/datafusion/optimizer/src/simplify_expressions/regex.rs @@ -283,20 +283,23 @@ fn partial_anchored_literal_to_like(v: &[Hir]) -> Option { /// Extracts a string literal expression assuming that [`is_anchored_literal`] /// returned true. -fn anchored_literal_to_expr(v: &[Hir]) -> Option { +fn anchored_literal_to_expr(v: &[Hir], string_scalar: &StringScalar) -> Option { match v.len() { - 2 => Some(lit("")), + 2 => Some(string_scalar.to_expr("")), 3 => { let HirKind::Literal(l) = v[1].kind() else { return None; }; - like_str_from_literal(l).map(lit) + like_str_from_literal(l).map(|s| string_scalar.to_expr(s)) } _ => None, } } -fn anchored_alternation_to_exprs(v: &[Hir]) -> Option> { +fn anchored_alternation_to_exprs( + v: &[Hir], + string_scalar: &StringScalar, +) -> Option> { if 3 != v.len() { return None; } @@ -308,7 +311,8 @@ fn anchored_alternation_to_exprs(v: &[Hir]) -> Option> { for hir in alters { let mut is_safe = false; if let HirKind::Literal(l) = hir.kind() - && let Some(safe_literal) = str_from_literal(l).map(lit) + && let Some(safe_literal) = + str_from_literal(l).map(|s| string_scalar.to_expr(s)) { literals.push(safe_literal); is_safe = true; @@ -321,7 +325,9 @@ fn anchored_alternation_to_exprs(v: &[Hir]) -> Option> { return Some(literals); } else if let HirKind::Literal(l) = sub.kind() { - if let Some(safe_literal) = str_from_literal(l).map(lit) { + if let Some(safe_literal) = + str_from_literal(l).map(|s| string_scalar.to_expr(s)) + { return Some(vec![safe_literal]); } return None; @@ -351,12 +357,18 @@ fn lower_simple( )); } HirKind::Concat(inner) if is_anchored_literal(inner) => { - return anchored_literal_to_expr(inner).map(|right| { - mode.expr_matches_literal(Box::new(left.clone()), Box::new(right)) + return anchored_literal_to_expr(inner, string_scalar).map(|right| { + if mode.i { + // Case-insensitive: use ILIKE for exact match (no wildcards) + mode.expr(Box::new(left.clone()), Box::new(right)) + } else { + // Case-sensitive: use Eq / NotEq + mode.expr_matches_literal(Box::new(left.clone()), Box::new(right)) + } }); } - HirKind::Concat(inner) if is_anchored_capture(inner) => { - return anchored_alternation_to_exprs(inner) + HirKind::Concat(inner) if !mode.i && is_anchored_capture(inner) => { + return anchored_alternation_to_exprs(inner, string_scalar) .map(|right| left.clone().in_list(right, mode.not)); } HirKind::Concat(inner) => { @@ -386,20 +398,17 @@ fn lower_alt( let mut accu: Option = None; for part in alts { - if let Some(expr) = lower_simple(mode, left, part, string_scalar) { - accu = match accu { - Some(accu) => { - if mode.not { - Some(accu.and(expr)) - } else { - Some(accu.or(expr)) - } + let expr = lower_simple(mode, left, part, string_scalar)?; + accu = match accu { + Some(accu) => { + if mode.not { + Some(accu.and(expr)) + } else { + Some(accu.or(expr)) } - None => Some(expr), - }; - } else { - return None; - } + } + None => Some(expr), + }; } Some(accu.expect("at least two alts")) diff --git a/datafusion/optimizer/src/simplify_expressions/reorder_predicates.rs b/datafusion/optimizer/src/simplify_expressions/reorder_predicates.rs new file mode 100644 index 0000000000000..221fa5d20c58c --- /dev/null +++ b/datafusion/optimizer/src/simplify_expressions/reorder_predicates.rs @@ -0,0 +1,193 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Reorder conjunctive (`AND`) predicates so that cheap predicates run before +//! expensive ones. +//! +//! DataFusion's `AND` evaluator short-circuits the right-hand side when the +//! left-hand side keeps few rows, so leading with a cheap predicate shrinks +//! the batch that expensive ones see. +//! +//! The cost of evaluating a predicate is assessed with a simple, conservative +//! heuristic: we define an allow-list of cheap operations, and consider an +//! expression to be cheap if it consists ONLY of cheap operations; everything +//! else is considered expensive. The sort of stable, so order within each +//! class is preserved. +//! +//! This reordering scheme is intentionally simple; many enhancements are +//! possible (e.g., consider both cost and selectivity, build a more complex +//! cost model, add estimated evaluation cost for individual UDFs). + +use datafusion_common::tree_node::TreeNode; +use datafusion_expr::{BinaryExpr, Expr, Operator}; + +/// Stable partition of `predicates`: cheap first, then expensive. +/// +/// Returns `(predicates, changed)`. When `changed` is `false` the input was +/// already cheap-first and the caller can skip rebuilding the conjunction. +pub(crate) fn reorder_predicates(predicates: Vec) -> (Vec, bool) { + if predicates.len() <= 1 { + return (predicates, false); + } + + // Volatile predicates may have observable side-effects and reordering + // conjuncts can change how many times they evaluate. Preserve user order + // if any predicate contains a volatile expression. + if predicates.iter().any(Expr::is_volatile) { + return (predicates, false); + } + + let classes: Vec = predicates.iter().map(is_cheap_predicate).collect(); + + // A reorder is needed iff an expensive predicate precedes a cheap one + let needs_reorder = classes.windows(2).any(|w| !w[0] && w[1]); + if !needs_reorder { + return (predicates, false); + } + + let mut cheap = Vec::with_capacity(predicates.len()); + let mut expensive = Vec::new(); + for (p, is_cheap) in predicates.into_iter().zip(classes) { + if is_cheap { + cheap.push(p); + } else { + expensive.push(p); + } + } + cheap.extend(expensive); + (cheap, true) +} + +/// Returns true if every node in `expr`'s tree is cheap. +fn is_cheap_predicate(expr: &Expr) -> bool { + !expr + .exists(|node| Ok(!is_cheap_node(node))) + .expect("is_cheap_node is infallible") +} + +/// Returns true if `expr` is itself cheap. +/// +/// We use a simple, conservative heuristic to determine if an expression is +/// cheap to evaluate: we enumerate known-cheap operations (e.g., equality +/// comparisons, negations, casts), and consider anything outside this list to +/// be expensive. New/unrecognized expressions therefore default to being +/// expensive. +fn is_cheap_node(expr: &Expr) -> bool { + match expr { + // Direct reads and literals. + Expr::Column(_) + | Expr::Literal(_, _) + | Expr::ScalarVariable(_, _) + | Expr::Placeholder(_) + | Expr::OuterReferenceColumn(_, _) + | Expr::LambdaVariable(_) + // Wrappers; children are walked separately by `is_cheap_predicate`. + | Expr::Alias(_) + // Single-row unary predicates and arithmetic negation. + | Expr::Not(_) + | Expr::Negative(_) + | Expr::IsNull(_) + | Expr::IsNotNull(_) + | Expr::IsTrue(_) + | Expr::IsFalse(_) + | Expr::IsUnknown(_) + | Expr::IsNotTrue(_) + | Expr::IsNotFalse(_) + | Expr::IsNotUnknown(_) + // Composite cheap forms; child expressions are walked separately. + | Expr::Between(_) + | Expr::Case(_) + | Expr::Cast(_) + | Expr::TryCast(_) + | Expr::InList(_) => true, + // BinaryExpr is cheap unless the operator is LIKE or regexp matching. + Expr::BinaryExpr(BinaryExpr { op, .. }) => !matches!( + op, + Operator::LikeMatch + | Operator::ILikeMatch + | Operator::NotLikeMatch + | Operator::NotILikeMatch + | Operator::RegexMatch + | Operator::RegexIMatch + | Operator::RegexNotMatch + | Operator::RegexNotIMatch + ), + _ => false, + } +} + +#[cfg(test)] +mod tests { + use super::*; + use datafusion_expr::{col, lit}; + + #[test] + fn like_predicate_moves_after_equality() { + let cheap = col("a").eq(lit(1)); + let expensive = col("b").like(lit("%foo%")); + let (out, changed) = reorder_predicates(vec![expensive.clone(), cheap.clone()]); + assert_eq!(out, vec![cheap, expensive]); + assert!(changed); + } + + #[test] + fn order_among_cheap_predicates_is_preserved() { + let p1 = col("a").eq(lit(1)); + let p2 = col("b").eq(lit(2)); + let p3 = col("c").eq(lit(3)); + let input = vec![p1.clone(), p2.clone(), p3.clone()]; + let (out, changed) = reorder_predicates(input.clone()); + assert_eq!(out, input); + assert!(!changed); + } + + #[test] + fn order_among_expensive_predicates_is_preserved() { + let p1 = col("a").like(lit("%a%")); + let p2 = Expr::BinaryExpr(BinaryExpr::new( + Box::new(col("b")), + Operator::RegexMatch, + Box::new(lit("foo")), + )); + let p3 = col("c").like(lit("%c%")); + let input = vec![p1.clone(), p2.clone(), p3.clone()]; + let (out, changed) = reorder_predicates(input.clone()); + assert_eq!(out, input); + assert!(!changed); + } + + #[test] + fn already_cheap_first_reports_no_change() { + let cheap = col("a").eq(lit(1)); + let expensive = col("b").like(lit("%a%")); + let input = vec![cheap.clone(), expensive.clone()]; + let (out, changed) = reorder_predicates(input.clone()); + assert_eq!(out, input); + assert!(!changed); + } + + #[test] + fn nested_expensive_under_not_is_expensive() { + // The top node is `Not`, which is on the cheap allow-list. The walk + // must descend into the `Like` to flag this predicate as expensive. + let cheap = col("a").eq(lit(1)); + let nested = Expr::Not(Box::new(col("b").like(lit("%foo%")))); + let (out, changed) = reorder_predicates(vec![nested.clone(), cheap.clone()]); + assert_eq!(out, vec![cheap, nested]); + assert!(changed); + } +} diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs index 3e495f5355103..0e72a17abc9f7 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_exprs.rs @@ -21,12 +21,12 @@ use std::sync::Arc; use datafusion_common::tree_node::{Transformed, TreeNode}; use datafusion_common::{Column, DFSchema, DFSchemaRef, DataFusionError, Result}; -use datafusion_expr::Expr; use datafusion_expr::logical_plan::{Aggregate, LogicalPlan, Projection}; use datafusion_expr::simplify::SimplifyContext; use datafusion_expr::utils::{ columnize_expr, find_aggregate_exprs, grouping_set_to_exprlist, merge_schema, }; +use datafusion_expr::{DmlStatement, Expr, WriteOp}; use super::ExprSimplifier; use crate::optimizer::ApplyOrder; @@ -77,7 +77,20 @@ impl SimplifyExpressions { plan: LogicalPlan, config: &dyn OptimizerConfig, ) -> Result> { - let schema = if !plan.inputs().is_empty() { + let schema = if let LogicalPlan::Dml(DmlStatement { + op: WriteOp::MergeInto(_), + table_name, + target, + .. + }) = &plan + { + let mut schema = merge_schema(&plan.inputs()); + schema.merge(&DFSchema::try_from_qualified_schema( + table_name.clone(), + &target.schema(), + )?); + DFSchemaRef::new(schema) + } else if !plan.inputs().is_empty() { DFSchemaRef::new(merge_schema(&plan.inputs())) } else if let LogicalPlan::TableScan(scan) = &plan { // When predicates are pushed into a table scan, there is no input diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_literal.rs b/datafusion/optimizer/src/simplify_expressions/simplify_literal.rs index 72e9dbc99dfae..2236a7e55bc52 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_literal.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_literal.rs @@ -62,7 +62,7 @@ where .simplify(expr.clone()) .map_err(|err| plan_datafusion_err!("Cannot simplify {expr:?}: {err}"))?; let coerced_expr: Expr = simplifier.coerce(simplified_expr, schema.as_ref())?; - log::debug!("Coerced expression: {:?}", &coerced_expr); + log::debug!("Coerced expression: {coerced_expr:?}"); match coerced_expr { Expr::Literal(scalar_value, _) => { diff --git a/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs b/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs index e811ce7313102..356f2711b708e 100644 --- a/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs +++ b/datafusion/optimizer/src/simplify_expressions/simplify_predicates.rs @@ -63,12 +63,14 @@ pub fn simplify_predicates(predicates: Vec) -> Result> { | Operator::Eq, right, }) => { - let left_col = extract_column_from_expr(left); - let right_col = extract_column_from_expr(right); - if let (Some(col), Some(_)) = (&left_col, right.as_literal()) { - column_predicates.entry(col.clone()).or_default().push(pred); - } else if let (Some(_), Some(col)) = (left.as_literal(), &right_col) { - column_predicates.entry(col.clone()).or_default().push(pred); + if let (Some(col), Some(_)) = + (extract_column_from_expr(left), right.as_literal()) + { + column_predicates.entry(col).or_default().push(pred); + } else if let (Some(_), Some(col)) = + (left.as_literal(), extract_column_from_expr(right)) + { + column_predicates.entry(col).or_default().push(pred); } else { other_predicates.push(pred); } diff --git a/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs b/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs index a5b65d0d8e7a4..ef0bfa516fe41 100644 --- a/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs +++ b/datafusion/optimizer/src/simplify_expressions/unwrap_cast.rs @@ -59,7 +59,10 @@ use datafusion_common::{Result, ScalarValue}; use datafusion_common::{internal_err, tree_node::Transformed}; use datafusion_expr::{BinaryExpr, lit}; use datafusion_expr::{Cast, Expr, Operator, TryCast, simplify::SimplifyContext}; -use datafusion_expr_common::casts::{is_supported_type, try_cast_literal_to_type}; +use datafusion_expr_common::casts::{ + is_date_narrowing_cast, is_supported_type, is_timestamp_precision_narrowing_cast, + try_cast_literal_to_type, +}; pub(super) fn unwrap_cast_in_comparison_for_binary( info: &SimplifyContext, @@ -113,10 +116,14 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary( match (expr, literal) { ( Expr::TryCast(TryCast { - expr: left_expr, .. + expr: left_expr, + field, + .. }) | Expr::Cast(Cast { - expr: left_expr, .. + expr: left_expr, + field, + .. }), Expr::Literal(lit_val, _), ) => { @@ -128,6 +135,12 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_binary( return false; }; + if is_timestamp_precision_narrowing_cast(&expr_type, field.data_type()) + || is_date_narrowing_cast(&expr_type, field.data_type()) + { + return false; + } + if cast_literal_to_type_with_op(lit_val, &expr_type, op).is_some() { return true; } @@ -146,10 +159,14 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist( list: &[Expr], ) -> bool { let (Expr::TryCast(TryCast { - expr: left_expr, .. + expr: left_expr, + field, + .. }) | Expr::Cast(Cast { - expr: left_expr, .. + expr: left_expr, + field, + .. })) = expr else { return false; @@ -163,6 +180,12 @@ pub(super) fn is_cast_expr_and_support_unwrap_cast_in_comparison_for_inlist( return false; } + if is_timestamp_precision_narrowing_cast(&expr_type, field.data_type()) + || is_date_narrowing_cast(&expr_type, field.data_type()) + { + return false; + } + for right in list { let Ok(right_type) = info.get_data_type(right) else { return false; @@ -586,6 +609,25 @@ mod tests { assert_eq!(optimize_test(expr_lt, &schema), expected); } + #[test] + fn test_not_unwrap_cast_timestamp_precision_narrowing() { + let schema = expr_test_schema(); + let expr_input = cast(col("ts_nano_none"), timestamp_millis_none_type()) + .eq(lit_timestamp_millis_none(1)); + + assert_eq!(optimize_test(expr_input.clone(), &schema), expr_input); + } + + #[test] + fn test_unwrap_cast_timestamp_precision_widening() { + let schema = expr_test_schema(); + let expr_input = cast(col("ts_millis_none"), timestamp_nano_none_type()) + .eq(lit_timestamp_nano_none(1_000_000)); + let expected = col("ts_millis_none").eq(lit_timestamp_millis_none(1)); + + assert_eq!(optimize_test(expr_input, &schema), expected); + } + fn optimize_test(expr: Expr, schema: &DFSchemaRef) -> Expr { let simplifier = ExprSimplifier::new( SimplifyContext::builder() @@ -607,6 +649,7 @@ mod tests { Field::new("c5", DataType::Float32, false), Field::new("c6", DataType::UInt32, false), Field::new("ts_nano_none", timestamp_nano_none_type(), false), + Field::new("ts_millis_none", timestamp_millis_none_type(), false), Field::new("ts_nano_utf", timestamp_nano_utc_type(), false), Field::new("str1", DataType::Utf8, false), Field::new("largestr", DataType::LargeUtf8, false), @@ -643,6 +686,10 @@ mod tests { lit(ScalarValue::TimestampNanosecond(Some(ts), None)) } + fn lit_timestamp_millis_none(ts: i64) -> Expr { + lit(ScalarValue::TimestampMillisecond(Some(ts), None)) + } + fn lit_timestamp_nano_utc(ts: i64) -> Expr { let utc = Some("+0:00".into()); lit(ScalarValue::TimestampNanosecond(Some(ts), utc)) @@ -652,6 +699,10 @@ mod tests { DataType::Timestamp(TimeUnit::Nanosecond, None) } + fn timestamp_millis_none_type() -> DataType { + DataType::Timestamp(TimeUnit::Millisecond, None) + } + // this is the type that now() returns fn timestamp_nano_utc_type() -> DataType { let utc = Some("+0:00".into()); diff --git a/datafusion/optimizer/src/simplify_expressions/utils.rs b/datafusion/optimizer/src/simplify_expressions/utils.rs index b0908b47602f7..89bb762d59ce2 100644 --- a/datafusion/optimizer/src/simplify_expressions/utils.rs +++ b/datafusion/optimizer/src/simplify_expressions/utils.rs @@ -17,7 +17,6 @@ //! Utility functions for expression simplification -use arrow::datatypes::i256; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::{ Case, Expr, Like, Operator, @@ -25,47 +24,6 @@ use datafusion_expr::{ expr_fn::{and, bitwise_and, bitwise_or, or}, }; -pub static POWS_OF_TEN: [i128; 38] = [ - 1, - 10, - 100, - 1000, - 10000, - 100000, - 1000000, - 10000000, - 100000000, - 1000000000, - 10000000000, - 100000000000, - 1000000000000, - 10000000000000, - 100000000000000, - 1000000000000000, - 10000000000000000, - 100000000000000000, - 1000000000000000000, - 10000000000000000000, - 100000000000000000000, - 1000000000000000000000, - 10000000000000000000000, - 100000000000000000000000, - 1000000000000000000000000, - 10000000000000000000000000, - 100000000000000000000000000, - 1000000000000000000000000000, - 10000000000000000000000000000, - 100000000000000000000000000000, - 1000000000000000000000000000000, - 10000000000000000000000000000000, - 100000000000000000000000000000000, - 1000000000000000000000000000000000, - 10000000000000000000000000000000000, - 100000000000000000000000000000000000, - 1000000000000000000000000000000000000, - 10000000000000000000000000000000000000, -]; - /// returns true if `needle` is found in a chain of search_op /// expressions. Such as: (A AND B) AND C fn expr_contains_inner(expr: &Expr, needle: &Expr, search_op: Operator) -> bool { @@ -139,54 +97,26 @@ pub fn delete_xor_in_complex_expr(expr: &Expr, needle: &Expr, is_left: bool) -> } pub fn is_zero(s: &Expr) -> bool { - match s { - Expr::Literal(ScalarValue::Int8(Some(0)), _) - | Expr::Literal(ScalarValue::Int16(Some(0)), _) - | Expr::Literal(ScalarValue::Int32(Some(0)), _) - | Expr::Literal(ScalarValue::Int64(Some(0)), _) - | Expr::Literal(ScalarValue::UInt8(Some(0)), _) - | Expr::Literal(ScalarValue::UInt16(Some(0)), _) - | Expr::Literal(ScalarValue::UInt32(Some(0)), _) - | Expr::Literal(ScalarValue::UInt64(Some(0)), _) => true, - Expr::Literal(ScalarValue::Float32(Some(v)), _) if *v == 0. => true, - Expr::Literal(ScalarValue::Float64(Some(v)), _) if *v == 0. => true, - Expr::Literal(ScalarValue::Decimal128(Some(v), _p, _s), _) if *v == 0 => true, - Expr::Literal(ScalarValue::Decimal256(Some(v), _p, _s), _) - if *v == i256::ZERO => - { - true - } - _ => false, + if let Expr::Literal(sv, _) = s + && sv.data_type().is_numeric() + { + // unwrap safe since numeric types always have a 0 value + sv == &ScalarValue::new_zero(&sv.data_type()).unwrap() + } else { + false } } pub fn is_one(s: &Expr) -> bool { - match s { - Expr::Literal(ScalarValue::Int8(Some(1)), _) - | Expr::Literal(ScalarValue::Int16(Some(1)), _) - | Expr::Literal(ScalarValue::Int32(Some(1)), _) - | Expr::Literal(ScalarValue::Int64(Some(1)), _) - | Expr::Literal(ScalarValue::UInt8(Some(1)), _) - | Expr::Literal(ScalarValue::UInt16(Some(1)), _) - | Expr::Literal(ScalarValue::UInt32(Some(1)), _) - | Expr::Literal(ScalarValue::UInt64(Some(1)), _) => true, - Expr::Literal(ScalarValue::Float32(Some(v)), _) if *v == 1. => true, - Expr::Literal(ScalarValue::Float64(Some(v)), _) if *v == 1. => true, - Expr::Literal(ScalarValue::Decimal128(Some(v), _p, s), _) => { - *s >= 0 - && POWS_OF_TEN - .get(*s as usize) - .map(|x| x == v) - .unwrap_or_default() - } - Expr::Literal(ScalarValue::Decimal256(Some(v), _p, s), _) => { - *s >= 0 - && match i256::from(10).checked_pow(*s as u32) { - Some(res) => res == *v, - None => false, - } - } - _ => false, + if let Expr::Literal(sv, _) = s + && sv.data_type().is_numeric() + // there are edge cases like negative scale decimals not being able to + // create a one value so this can fail + && let Ok(one) = ScalarValue::new_one(&sv.data_type()) + { + sv == &one + } else { + false } } diff --git a/datafusion/optimizer/src/utils.rs b/datafusion/optimizer/src/utils.rs index ad151d1ddb8e0..4ea1589cfa7df 100644 --- a/datafusion/optimizer/src/utils.rs +++ b/datafusion/optimizer/src/utils.rs @@ -24,10 +24,12 @@ use arrow::array::{Array, RecordBatch, new_null_array}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::TableReference; use datafusion_common::cast::as_boolean_array; -use datafusion_common::tree_node::{TransformedResult, TreeNode}; +use datafusion_common::tree_node::{TransformedResult, TreeNode, TreeNodeRecursion}; use datafusion_common::{Column, DFSchema, Result, ScalarValue}; use datafusion_expr::execution_props::ExecutionProps; +use datafusion_expr::expr::{Exists, InSubquery, SetComparison}; use datafusion_expr::expr_rewriter::replace_col; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ColumnarValue, Expr, logical_plan::LogicalPlan}; use datafusion_physical_expr::create_physical_expr; use log::{debug, trace}; @@ -37,6 +39,56 @@ use std::sync::Arc; /// as it was initially placed here and then moved elsewhere. pub use datafusion_expr::expr_rewriter::NamePreserver; +/// Invokes `f` with the index, within `schema`, of every column referenced by +/// `expr` — including columns reached through a correlated subquery's outer +/// references. Columns absent from `schema` are skipped. +/// +/// A subquery's own plan is intentionally not traversed: its internal columns +/// index into its own schema, not `schema`; only the outer (correlated) columns +/// it references from `schema` are relevant. The comparison expression of an +/// `IN`/set-comparison subquery is reached by the normal expression walk. +/// +/// This is the shared primitive behind the top-down "which of a node's output +/// columns does an ancestor still need" analyses, namely +/// [`OptimizeProjections`](crate::optimize_projections::OptimizeProjections) +/// and [`EliminateJoin`](crate::eliminate_join::EliminateJoin). The two keep +/// their own required-index containers (an ordered set vs. a hash set), so this +/// reports indices through a callback rather than populating a shared type. +pub(crate) fn for_each_referenced_index( + expr: &Expr, + schema: &DFSchema, + mut f: impl FnMut(usize), +) -> Result<()> { + visit_referenced_indices(expr, schema, &mut f) +} + +fn visit_referenced_indices( + expr: &Expr, + schema: &DFSchema, + f: &mut dyn FnMut(usize), +) -> Result<()> { + expr.apply(|expr| { + match expr { + Expr::Column(column) | Expr::OuterReferenceColumn(_, column) => { + if let Some(idx) = schema.maybe_index_of_column(column) { + f(idx); + } + } + Expr::Exists(Exists { subquery, .. }) + | Expr::InSubquery(InSubquery { subquery, .. }) + | Expr::SetComparison(SetComparison { subquery, .. }) + | Expr::ScalarSubquery(subquery) => { + for outer in &subquery.outer_ref_columns { + visit_referenced_indices(outer, schema, f)?; + } + } + _ => {} + } + Ok(TreeNodeRecursion::Continue) + })?; + Ok(()) +} + /// Returns true if `expr` contains all columns in `schema_cols` pub(crate) fn has_all_column_refs( expr: &Expr, @@ -182,8 +234,13 @@ fn evaluate_expr_with_null_column<'a>( let replaced_predicate = replace_col(predicate, &join_cols_to_replace)?; let coerced_predicate = coerce(replaced_predicate, &input_schema)?; - create_physical_expr(&coerced_predicate, &input_schema, &execution_props)? - .evaluate(&input_batch) + create_physical_expr( + &coerced_predicate, + &input_schema, + &execution_props, + &PhysicalPlanningContext::default(), + )? + .evaluate(&input_batch) } fn coerce(expr: Expr, schema: &DFSchema) -> Result { diff --git a/datafusion/optimizer/tests/optimizer_integration.rs b/datafusion/optimizer/tests/optimizer_integration.rs index e61e6467930e6..26b48c5e1f352 100644 --- a/datafusion/optimizer/tests/optimizer_integration.rs +++ b/datafusion/optimizer/tests/optimizer_integration.rs @@ -56,8 +56,7 @@ fn init() { #[test] fn recursive_cte_with_nested_subquery() -> Result<()> { - // Covers bailout path in `plan_contains_other_subqueries`, ensuring nested subqueries - // within recursive CTE branches prevent projection pushdown. + // projection optimization is applied to recursive CTEs even with nested subqueries let sql = r#" WITH RECURSIVE numbers(id, level) AS ( SELECT sub.id, sub.level FROM ( @@ -75,21 +74,20 @@ fn recursive_cte_with_nested_subquery() -> Result<()> { assert_snapshot!( format!("{plan}"), - @r" + @" SubqueryAlias: numbers - Projection: sub.id AS id, sub.level AS level - RecursiveQuery: is_distinct=false - Projection: sub.id, sub.level - SubqueryAlias: sub - Projection: test.col_int32 AS id, Int64(1) AS level - TableScan: test - Projection: t.col_int32, numbers.level + Int64(1) - Inner Join: CAST(t.col_int32 AS Int64) = CAST(numbers.id AS Int64) + Int64(1) - SubqueryAlias: t - Filter: CAST(test.col_int32 AS Int64) IS NOT NULL - TableScan: test - Filter: CAST(numbers.id AS Int64) + Int64(1) IS NOT NULL - TableScan: numbers + RecursiveQuery: is_distinct=false + Projection: sub.id AS id, sub.level AS level + SubqueryAlias: sub + Projection: test.col_int32 AS id, Int64(1) AS level + TableScan: test projection=[col_int32] + Projection: t.col_int32, numbers.level + Int64(1) + Inner Join: CAST(t.col_int32 AS Int64) = CAST(numbers.id AS Int64) + Int64(1) + SubqueryAlias: t + Filter: CAST(test.col_int32 AS Int64) IS NOT NULL + TableScan: test projection=[col_int32] + Filter: CAST(numbers.id AS Int64) + Int64(1) IS NOT NULL + TableScan: numbers projection=[id, level] " ); @@ -277,13 +275,12 @@ fn intersect() -> Result<()> { format!("{plan}"), @r" LeftSemi Join: left.col_int32 = test.col_int32, left.col_utf8 = test.col_utf8 - Aggregate: groupBy=[[left.col_int32, left.col_utf8]], aggr=[[]] - LeftSemi Join: left.col_int32 = right.col_int32, left.col_utf8 = right.col_utf8 - Aggregate: groupBy=[[left.col_int32, left.col_utf8]], aggr=[[]] - SubqueryAlias: left - TableScan: test projection=[col_int32, col_utf8] - SubqueryAlias: right + LeftSemi Join: left.col_int32 = right.col_int32, left.col_utf8 = right.col_utf8 + Aggregate: groupBy=[[left.col_int32, left.col_utf8]], aggr=[[]] + SubqueryAlias: left TableScan: test projection=[col_int32, col_utf8] + SubqueryAlias: right + TableScan: test projection=[col_int32, col_utf8] TableScan: test projection=[col_int32, col_utf8] " ); @@ -527,12 +524,10 @@ fn select_correlated_predicate_subquery_with_uppercase_ident() { " ); } - #[test] -fn recursive_cte_projection_pushdown() -> Result<()> { - // Test that projection pushdown works with recursive CTEs by ensuring - // only the required columns are projected from the base table, even when - // the CTE definition includes unused columns +fn recursive_cte_outer_projection_pushdown() -> Result<()> { + // projection optimization of a recursive CTE based on the outer query's projected columns is + // not done as this can lead to bugs (see: https://github.com/apache/datafusion/issues/22249). let sql = "WITH RECURSIVE nodes AS (\ SELECT col_int32 AS id, col_utf8 AS name, col_uint32 AS extra FROM test \ UNION ALL \ @@ -540,18 +535,19 @@ fn recursive_cte_projection_pushdown() -> Result<()> { ) SELECT id FROM nodes"; let plan = test_sql(sql)?; - // The optimizer successfully performs projection pushdown by only selecting the needed - // columns from the base table and recursive table, eliminating unused columns + // col_int32, col_utf8, and col_uint32 and projected from test since they are used in the + // recursive CTE, even though the outer query only requires col_int32 assert_snapshot!( format!("{plan}"), @r" SubqueryAlias: nodes - RecursiveQuery: is_distinct=false - Projection: test.col_int32 AS id - TableScan: test projection=[col_int32] - Projection: CAST(CAST(nodes.id AS Int64) + Int64(1) AS Int32) - Filter: nodes.id < Int32(3) - TableScan: nodes projection=[id] + Projection: id + RecursiveQuery: is_distinct=false + Projection: test.col_int32 AS id, test.col_utf8 AS name, test.col_uint32 AS extra + TableScan: test projection=[col_int32, col_uint32, col_utf8] + Projection: CAST(CAST(nodes.id AS Int64) + Int64(1) AS Int32), nodes.name, nodes.extra + Filter: nodes.id < Int32(3) + TableScan: nodes projection=[id, name, extra] " ); Ok(()) @@ -570,47 +566,19 @@ fn recursive_cte_with_aliased_self_reference() -> Result<()> { format!("{plan}"), @r" SubqueryAlias: nodes - RecursiveQuery: is_distinct=false - Projection: test.col_int32 AS id - TableScan: test projection=[col_int32] - Projection: CAST(CAST(child.id AS Int64) + Int64(1) AS Int32) - SubqueryAlias: child - Filter: nodes.id < Int32(3) - TableScan: nodes projection=[id] + Projection: id + RecursiveQuery: is_distinct=false + Projection: test.col_int32 AS id, test.col_utf8 AS name + TableScan: test projection=[col_int32, col_utf8] + Projection: CAST(CAST(child.id AS Int64) + Int64(1) AS Int32), child.name + SubqueryAlias: child + Filter: nodes.id < Int32(3) + TableScan: nodes projection=[id, name] ", ); Ok(()) } -#[test] -fn recursive_cte_with_unused_columns() -> Result<()> { - // Test projection pushdown with a recursive CTE where the base case - // includes columns that are never used in the recursive part or final result - let sql = "WITH RECURSIVE series AS (\ - SELECT 1 AS n, col_utf8, col_uint32, col_date32 FROM test WHERE col_int32 = 1 \ - UNION ALL \ - SELECT n + 1, col_utf8, col_uint32, col_date32 FROM series WHERE n < 3\ - ) SELECT n FROM series"; - let plan = test_sql(sql)?; - - // The optimizer successfully performs projection pushdown by eliminating unused columns - // even when they're defined in the CTE but not actually needed - assert_snapshot!( - format!("{plan}"), - @r" - SubqueryAlias: series - RecursiveQuery: is_distinct=false - Projection: Int64(1) AS n - Filter: test.col_int32 = Int32(1) - TableScan: test projection=[col_int32] - Projection: series.n + Int64(1) - Filter: series.n < Int64(3) - TableScan: series projection=[n] - " - ); - Ok(()) -} - #[test] /// Asserts the minimal plan shape once projection pushdown succeeds for a recursive CTE. /// Unlike the previous two tests that retain extra columns in either the base or recursive @@ -824,10 +792,9 @@ fn extension_node_does_not_block_projection_pruning() -> Result<()> { OpaqueRequirementsExtension Sort: t.a ASC NULLS FIRST, t.ts ASC NULLS FIRST Projection: t.a, CAST(t.ts AS Timestamp(ms, "UTC")) AS ts - Projection: t.a, t.ts - Filter: __common_expr_3 > TimestampMillisecond(1000, Some("UTC")) AND __common_expr_3 < TimestampMillisecond(2000, Some("UTC")) - Projection: CAST(t.ts AS Timestamp(ms, "UTC")) AS __common_expr_3, t.a, t.ts - TableScan: t projection=[a, ts], partial_filters=[t.ts > TimestampNanosecond(1000000000, None), t.ts < TimestampNanosecond(2000000000, None), CAST(t.ts AS Timestamp(ms, "UTC")) > TimestampMillisecond(1000, Some("UTC")), CAST(t.ts AS Timestamp(ms, "UTC")) < TimestampMillisecond(2000, Some("UTC"))] + Filter: __common_expr_3 > TimestampMillisecond(1000, Some("UTC")) AND __common_expr_3 < TimestampMillisecond(2000, Some("UTC")) + Projection: CAST(t.ts AS Timestamp(ms, "UTC")) AS __common_expr_3, t.a, t.ts + TableScan: t projection=[a, ts], partial_filters=[CAST(t.ts AS Timestamp(ms, "UTC")) > TimestampMillisecond(1000, Some("UTC")), CAST(t.ts AS Timestamp(ms, "UTC")) < TimestampMillisecond(2000, Some("UTC"))] "#, ); @@ -882,7 +849,7 @@ impl ContextProvider for MyContextProvider { fn get_higher_order_meta( &self, _name: &str, - ) -> Option> { + ) -> Option> { None } diff --git a/datafusion/physical-expr-adapter/src/lib.rs b/datafusion/physical-expr-adapter/src/lib.rs index ea4db19ee110e..b224d8f4b8fe9 100644 --- a/datafusion/physical-expr-adapter/src/lib.rs +++ b/datafusion/physical-expr-adapter/src/lib.rs @@ -24,6 +24,7 @@ //! Physical expression schema adaptation utilities for DataFusion +pub mod rewrite; pub mod schema_rewriter; pub use schema_rewriter::{ diff --git a/datafusion/physical-expr-adapter/src/rewrite.rs b/datafusion/physical-expr-adapter/src/rewrite.rs new file mode 100644 index 0000000000000..7345a587ee6a4 --- /dev/null +++ b/datafusion/physical-expr-adapter/src/rewrite.rs @@ -0,0 +1,337 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Rewrite expressions in preparation for files being scanned, such as scan-metadata scalar UDFs. +//! +//! Functions like [`file_row_index()`] and [`input_file_name()`] are placeholders +//! whose value is only known during a file scan. The helpers here replace those +//! UDFs with ordinary physical expressions bound to the current file: a column +//! reference into a source-provided row-index column, or a per-file literal, etc. +//! +//! [`file_row_index()`]: datafusion_functions::core::file_row_index::FileRowIndexFunc +//! [`input_file_name()`]: datafusion_functions::core::input_file_name::InputFileNameFunc + +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field}; +use datafusion_common::{ + Result, ScalarValue, + tree_node::{Transformed, TreeNode, TreeNodeRecursion}, +}; +use datafusion_expr::ScalarUDFImpl; +use datafusion_functions::core::file_row_index::FileRowIndexFunc; +use datafusion_functions::core::input_file_name::InputFileNameFunc; +use datafusion_physical_expr::ScalarFunctionExpr; +use datafusion_physical_expr::expressions::{CastExpr, Column, Literal}; +use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + +/// Return true if a [`PhysicalExpr`] references scalar UDF `T`. +/// +/// This matches the concrete [`ScalarUDFImpl`] type rather than the function +/// name, so unrelated UDFs with the same name are not treated as matches. +pub fn expr_references_scalar_udf( + expr: &Arc, +) -> bool { + let mut found = false; + + expr.apply(|node| { + if ScalarFunctionExpr::try_downcast_func::(node.as_ref()).is_some() { + found = true; + return Ok(TreeNodeRecursion::Stop); + } + Ok(TreeNodeRecursion::Continue) + }) + .expect("Infallible traversal of PhysicalExpr tree failed"); + + found +} + +/// Rewrite occurrences of scalar UDF `T` in a [`PhysicalExpr`] using +/// `replacement`. +/// +/// The rewrite matches the concrete [`ScalarUDFImpl`] type rather than the +/// function name. `replacement` is called with each matching +/// [`ScalarFunctionExpr`] after its children have been rewritten. +fn rewrite_scalar_udf( + expr: Arc, + mut replacement: F, +) -> Result> +where + T: ScalarUDFImpl, + F: FnMut(&ScalarFunctionExpr) -> Result>, +{ + expr.transform_up(|node| { + if let Some(scalar_fn) = ScalarFunctionExpr::try_downcast_func::(node.as_ref()) + { + Ok(Transformed::yes(replacement(scalar_fn)?)) + } else { + Ok(Transformed::no(node)) + } + }) + .map(|transformed| transformed.data) +} + +/// Rewrite [`file_row_index()`][FileRowIndexFunc] in a [`PhysicalExpr`] to +/// read from a source-provided row-index column. +/// +/// `row_index_idx` is the index of `row_index_name` in the schema that the +/// rewritten expression will be evaluated against. The rewrite uses ordinary +/// physical expressions: a [`Column`] that reads the source row-index values +/// wrapped in a [`CastExpr`] that exposes the public `file_row_index: Int64` +/// return field without source-specific extension metadata. +pub fn rewrite_file_row_index_expr( + expr: Arc, + row_index_name: &str, + row_index_idx: usize, +) -> Result> { + rewrite_scalar_udf::(expr, |_| { + let source = Arc::new(Column::new(row_index_name, row_index_idx)); + let target_field = Arc::new(Field::new("file_row_index", DataType::Int64, true)); + Ok(Arc::new(CastExpr::new_with_target_field( + source, + target_field, + None, + ))) + }) +} + +/// Rewrite [`file_row_index()`][FileRowIndexFunc] in pushed [`ProjectionExprs`] +/// to read from a source-provided row-index column. +/// +/// +/// For example if `row_index_column` is `__datafusion_row_idx` this function rewrites all +/// instances of [`file_row_index()`][FileRowIndexFunc] to +/// `__datafusion_row_index` [`Column`] references. +/// +/// `base_projection` is the current projection already pushed into a source. +/// The row-index source column is appended to that base projection if it is not +/// already present. `projection` is rewritten to read from the projected +/// row-index column and then merged on top of the extended base projection. +pub fn rewrite_file_row_index_projection( + base_projection: &ProjectionExprs, + projection: &ProjectionExprs, + row_index_col: &Column, +) -> Result { + let mut base_exprs = base_projection.as_ref().to_vec(); + let row_index_projection_idx = + base_projection.projected_column_position(row_index_col); + + // If the column doesn't exist in the projection yet + if row_index_projection_idx.is_none() { + base_exprs.push(ProjectionExpr { + expr: Arc::new(row_index_col.clone()), + alias: row_index_col.name().to_owned(), + }); + } + + let rewritten_projection = projection.clone().try_map_exprs(|expr| { + rewrite_file_row_index_expr( + expr, + row_index_col.name(), + row_index_projection_idx.unwrap_or(base_exprs.len() - 1), + ) + })?; + + ProjectionExprs::new(base_exprs).try_merge(&rewritten_projection) +} + +/// Rewrite [`input_file_name()`][InputFileNameFunc] in pushed +/// [`ProjectionExprs`] to a per-file [`Literal`] holding `file_name`. +/// +/// If the projection contains no [`input_file_name()`][InputFileNameFunc] UDF it +/// is returned unchanged, without allocating the literal or rebuilding the +/// projection tree (the common case for queries that don't use the function). +pub fn rewrite_input_file_name_in_projection( + projection: ProjectionExprs, + file_name: &str, +) -> Result { + if !projection + .iter() + .any(|p| expr_references_scalar_udf::(&p.expr)) + { + return Ok(projection); + } + + let file_name_lit = + Arc::new(Literal::new(ScalarValue::Utf8(Some(file_name.to_string())))) + as Arc; + + projection.try_map_exprs(|expr| { + rewrite_scalar_udf::(expr, |_| { + Ok(Arc::clone(&file_name_lit)) + }) + }) +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::datatypes::Schema; + use datafusion_common::config::ConfigOptions; + use datafusion_expr::{Operator, ScalarUDF}; + use datafusion_physical_expr::expressions; + use std::collections::HashMap; + + fn file_row_index_expr() -> Arc { + Arc::new(ScalarFunctionExpr::new( + "file_row_index", + Arc::new(ScalarUDF::from(FileRowIndexFunc::new())), + vec![], + Arc::new(Field::new("file_row_index", DataType::Int64, true)), + Arc::new(ConfigOptions::default()), + )) + } + + fn input_file_name_expr() -> Arc { + Arc::new(ScalarFunctionExpr::new( + "input_file_name", + Arc::new(ScalarUDF::from(InputFileNameFunc::new())), + vec![], + Arc::new(Field::new("input_file_name", DataType::Utf8, true)), + Arc::new(ConfigOptions::default()), + )) + } + + #[test] + fn test_rewrite_scalar_udf_replaces_nested_typed_udf() -> Result<()> { + let expr = Arc::new(expressions::BinaryExpr::new( + file_row_index_expr(), + Operator::Plus, + expressions::lit(ScalarValue::Int64(Some(1))), + )) as Arc; + + let rewritten = rewrite_scalar_udf::(expr, |_| { + Ok(expressions::lit(ScalarValue::Int64(Some(7)))) + })?; + + let binary = rewritten + .downcast_ref::() + .expect("rewritten expression should remain binary"); + assert_eq!(binary.op(), &Operator::Plus); + + let left = binary + .left() + .downcast_ref::() + .expect("left side should be rewritten to a literal"); + assert_eq!(left.value(), &ScalarValue::Int64(Some(7))); + + let right = binary + .right() + .downcast_ref::() + .expect("right side should remain the original literal"); + assert_eq!(right.value(), &ScalarValue::Int64(Some(1))); + Ok(()) + } + + #[test] + fn test_rewrite_input_file_name_in_projection() -> Result<()> { + let file_name = "part=west/data.parquet"; + let projection = ProjectionExprs::new([ + ProjectionExpr::new(input_file_name_expr(), "file_name"), + ProjectionExpr::new( + Arc::new(expressions::BinaryExpr::new( + input_file_name_expr(), + Operator::Eq, + expressions::lit(ScalarValue::Utf8(Some(file_name.to_string()))), + )), + "matches_file", + ), + ]); + + let rewritten = rewrite_input_file_name_in_projection(projection, file_name)?; + let rewritten = rewritten.as_ref(); + assert_eq!(rewritten[0].alias, "file_name"); + assert_eq!(rewritten[1].alias, "matches_file"); + + let file_name_lit = rewritten[0] + .expr + .downcast_ref::() + .expect("input_file_name should rewrite to a literal"); + assert_eq!( + file_name_lit.value(), + &ScalarValue::Utf8(Some(file_name.to_string())) + ); + + let binary = rewritten[1] + .expr + .downcast_ref::() + .expect("nested expression should remain binary"); + assert_eq!(binary.op(), &Operator::Eq); + + let left = binary + .left() + .downcast_ref::() + .expect("nested input_file_name should rewrite to a literal"); + assert_eq!( + left.value(), + &ScalarValue::Utf8(Some(file_name.to_string())) + ); + + let right = binary + .right() + .downcast_ref::() + .expect("comparison literal should remain unchanged"); + assert_eq!( + right.value(), + &ScalarValue::Utf8(Some(file_name.to_string())) + ); + Ok(()) + } + + #[test] + fn test_rewrite_file_row_index_expr_to_source_column() -> Result<()> { + let expr = rewrite_file_row_index_expr( + file_row_index_expr(), + "__datafusion_file_row_index", + 2, + )?; + + let cast_expr = expr + .downcast_ref::() + .expect("file row index expression should be a cast"); + assert_eq!(cast_expr.cast_type(), &DataType::Int64); + let target_field = cast_expr.target_field(); + assert_eq!(target_field.name(), "file_row_index"); + assert_eq!(target_field.data_type(), &DataType::Int64); + assert!(target_field.is_nullable()); + assert!(target_field.metadata().is_empty()); + + let source = cast_expr + .expr() + .downcast_ref::() + .expect("source column"); + assert_eq!(source.name(), "__datafusion_file_row_index"); + assert_eq!(source.index(), 2); + + let input_schema = Schema::new(vec![ + Field::new("value", DataType::Int64, true), + Field::new("__datafusion_file_row_index", DataType::Int64, false) + .with_metadata(HashMap::from([( + "source".to_string(), + "virtual".to_string(), + )])), + ]); + let return_field = expr.return_field(&input_schema)?; + assert_eq!(return_field.name(), "file_row_index"); + assert_eq!(return_field.data_type(), &DataType::Int64); + assert!(return_field.is_nullable()); + assert!(return_field.metadata().is_empty()); + Ok(()) + } +} diff --git a/datafusion/physical-expr-adapter/src/schema_rewriter.rs b/datafusion/physical-expr-adapter/src/schema_rewriter.rs index 9fb4950317ff8..ef25af7d920fb 100644 --- a/datafusion/physical-expr-adapter/src/schema_rewriter.rs +++ b/datafusion/physical-expr-adapter/src/schema_rewriter.rs @@ -34,6 +34,7 @@ use datafusion_common::{ }; use datafusion_functions::core::getfield::GetFieldFunc; use datafusion_physical_expr::PhysicalExprSimplifier; +use datafusion_physical_expr::expressions::Literal; use datafusion_physical_expr::projection::{ProjectionExprs, Projector}; use datafusion_physical_expr::{ ScalarFunctionExpr, @@ -311,7 +312,7 @@ impl DefaultPhysicalExprAdapterRewriter { None => return Ok(None), }; - let lit = match field_name_expr.downcast_ref::() { + let lit = match field_name_expr.downcast_ref::() { Some(lit) => lit, None => return Ok(None), }; @@ -364,7 +365,7 @@ impl DefaultPhysicalExprAdapterRewriter { }; let null_value = ScalarValue::Null.cast_to(logical_struct_field.data_type())?; - Ok(Some(Arc::new(expressions::Literal::new_with_metadata( + Ok(Some(Arc::new(Literal::new_with_metadata( null_value, Some(FieldMetadata::from(logical_struct_field.as_ref())), )))) @@ -411,12 +412,10 @@ impl DefaultPhysicalExprAdapterRewriter { // If the column is missing from the physical schema fill it in with nulls. // For a different behavior, provide a custom `PhysicalExprAdapter` implementation. let null_value = ScalarValue::Null.cast_to(logical_field.data_type())?; - return Ok(Transformed::yes(Arc::new( - expressions::Literal::new_with_metadata( - null_value, - Some(FieldMetadata::from(logical_field)), - ), - ))); + return Ok(Transformed::yes(Arc::new(Literal::new_with_metadata( + null_value, + Some(FieldMetadata::from(logical_field)), + )))); }; let fields_match = logical_field == physical_field.as_ref(); @@ -432,7 +431,7 @@ impl DefaultPhysicalExprAdapterRewriter { // We need a cast expression whenever the logical and physical fields differ, // whether that difference is only metadata/nullability or also data type. // TODO: add optimization to move the cast from the column to literal expressions in the case of `col = 123` - // since that's much cheaper to evalaute. + // since that's much cheaper to evaluate. // See https://github.com/apache/datafusion/issues/15780#issuecomment-2824716928 validate_data_type_compatibility( resolved_column.name(), @@ -628,10 +627,11 @@ mod tests { use super::*; use arrow::array::{ Array, BooleanArray, GenericListArray, Int32Array, Int64Array, RecordBatch, - RecordBatchOptions, StringArray, StringViewArray, StructArray, + RecordBatchOptions, StringArray, StringViewArray, StructArray, record_batch, }; + use arrow::datatypes as arrow_schema; use arrow::datatypes::{Field, Fields, Schema}; - use datafusion_common::{assert_contains, record_batch}; + use datafusion_common::assert_contains; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{Column, Literal, col}; diff --git a/datafusion/physical-expr-common/Cargo.toml b/datafusion/physical-expr-common/Cargo.toml index 0e4748b81d3ff..903f5a6a901ac 100644 --- a/datafusion/physical-expr-common/Cargo.toml +++ b/datafusion/physical-expr-common/Cargo.toml @@ -40,11 +40,18 @@ workspace = true [lib] name = "datafusion_physical_expr_common" +[features] +default = [] +# Enables the `PhysicalExpr::to_proto` hook used by `datafusion-proto`. +# Off by default so crates that never serialize plans pay nothing. +proto = ["dep:datafusion-proto-models"] + [dependencies] arrow = { workspace = true } chrono = { workspace = true } datafusion-common = { workspace = true } datafusion-expr-common = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } hashbrown = { workspace = true } indexmap = { workspace = true } itertools = { workspace = true } @@ -58,3 +65,7 @@ rand = { workspace = true } [[bench]] harness = false name = "compare_nested" + +[[bench]] +harness = false +name = "arrow_bytes_map" diff --git a/datafusion/physical-expr-common/benches/arrow_bytes_map.rs b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs new file mode 100644 index 0000000000000..7c8cdc3b4c50e --- /dev/null +++ b/datafusion/physical-expr-common/benches/arrow_bytes_map.rs @@ -0,0 +1,82 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, StringArray}; +use criterion::{Criterion, Throughput, criterion_group, criterion_main}; +use datafusion_physical_expr_common::binary_map::{ArrowBytesMap, OutputType}; +use std::hint::black_box; +use std::sync::Arc; + +const NUM_ROWS: usize = 8192; + +fn make_short_strings(cardinality: usize) -> ArrayRef { + let values = (0..NUM_ROWS).map(|index| format!("{:04x}", index % cardinality)); + Arc::new(StringArray::from_iter_values(values)) +} + +fn make_long_strings(cardinality: usize) -> ArrayRef { + let values = (0..NUM_ROWS).map(|index| { + let value = (index % cardinality) as u32; + format!( + "{value:08x}{:08x}{:08x}{:08x}", + value.wrapping_mul(17), + value.wrapping_mul(31), + value.wrapping_mul(127) + ) + }); + Arc::new(StringArray::from_iter_values(values)) +} + +fn bench_arrow_bytes_map(c: &mut Criterion) { + let cases = [ + // Exercises inline entry storage while still growing the output buffer. + ("short_unique", make_short_strings(NUM_ROWS)), + // Exercises repeated buffer growth and out-of-line entry storage. + ("long_unique", make_long_strings(NUM_ROWS)), + // Fits the distinct values in the initial buffer and repeats comparisons. + ("long_low_cardinality", make_long_strings(128)), + ]; + + let mut group = c.benchmark_group("arrow_bytes_map"); + group.throughput(Throughput::Elements(NUM_ROWS as u64)); + + for (name, values) in cases { + group.bench_function(name, |b| { + b.iter(|| { + let mut map = ArrowBytesMap::::new(OutputType::Utf8); + let mut next_payload = 0; + map.insert_if_new( + &values, + |_| { + let payload = next_payload; + next_payload += 1; + payload + }, + |payload| { + black_box(payload); + }, + ); + black_box(map.into_state()) + }) + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_arrow_bytes_map); +criterion_main!(benches); diff --git a/datafusion/physical-expr-common/src/binary_map.rs b/datafusion/physical-expr-common/src/binary_map.rs index ad184d6500d56..44ca35c7f8708 100644 --- a/datafusion/physical-expr-common/src/binary_map.rs +++ b/datafusion/physical-expr-common/src/binary_map.rs @@ -19,12 +19,12 @@ //! StringArray / LargeStringArray / BinaryArray / LargeBinaryArray. use arrow::array::{ - Array, ArrayRef, BufferBuilder, GenericBinaryArray, GenericStringArray, - NullBufferBuilder, OffsetSizeTrait, + Array, ArrayRef, GenericBinaryArray, GenericStringArray, NullBufferBuilder, + OffsetSizeTrait, cast::AsArray, types::{ByteArrayType, GenericBinaryType, GenericStringType}, }; -use arrow::buffer::{NullBuffer, OffsetBuffer, ScalarBuffer}; +use arrow::buffer::{Buffer, NullBuffer, OffsetBuffer, ScalarBuffer}; use arrow::datatypes::DataType; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; @@ -218,8 +218,8 @@ where map: hashbrown::hash_table::HashTable>, /// Total size of the map in bytes map_size: usize, - /// In progress arrow `Buffer` containing all values - buffer: BufferBuilder, + /// In progress buffer containing all values + buffer: Vec, /// Offsets into `buffer` for each distinct value. These offsets as used /// directly to create the final `GenericBinaryArray`. The `i`th string is /// stored in the range `offsets[i]..offsets[i+1]` in `buffer`. Null values @@ -248,7 +248,7 @@ where output_type, map: hashbrown::hash_table::HashTable::with_capacity(INITIAL_MAP_CAPACITY), map_size: 0, - buffer: BufferBuilder::new(INITIAL_BUFFER_CAPACITY), + buffer: Vec::with_capacity(INITIAL_BUFFER_CAPACITY), offsets: vec![O::default()], // first offset is always 0 random_state: RandomState::default(), hashes_buffer: vec![], @@ -405,7 +405,7 @@ where // Put the small values into buffer and offsets so it appears // the output array, but store the actual bytes inline for // comparison - self.buffer.append_slice(value); + self.buffer.extend_from_slice(value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); let new_header = Entry { @@ -433,7 +433,7 @@ where // Need to compare the bytes in the buffer // SAFETY: buffer is only appended to, and we correctly inserted values and offsets let existing_value = - unsafe { self.buffer.as_slice().get_unchecked(header.range()) }; + unsafe { self.buffer.get_unchecked(header.range()) }; value == existing_value }); @@ -446,7 +446,7 @@ where // appears the output array, and store that offset // so the bytes can be compared if needed let offset = self.buffer.len(); // offset of start for data - self.buffer.append_slice(value); + self.buffer.extend_from_slice(value); self.offsets.push(O::usize_as(self.buffer.len())); let payload = make_payload_fn(Some(value)); @@ -488,7 +488,7 @@ where map: _, map_size: _, offsets, - mut buffer, + buffer, random_state: _, hashes_buffer: _, null, @@ -502,7 +502,7 @@ where // SAFETY: the offsets were constructed correctly in `insert_if_new` -- // monotonically increasing, overflows were checked. let offsets = unsafe { OffsetBuffer::new_unchecked(ScalarBuffer::from(offsets)) }; - let values = buffer.finish(); + let values = Buffer::from_vec(buffer); match output_type { OutputType::Binary => { diff --git a/datafusion/physical-expr-common/src/datum.rs b/datafusion/physical-expr-common/src/datum.rs index bd5790507f662..d23fb30db6c4a 100644 --- a/datafusion/physical-expr-common/src/datum.rs +++ b/datafusion/physical-expr-common/src/datum.rs @@ -23,6 +23,7 @@ use arrow::compute::kernels::cmp::{ }; use arrow::compute::{SortOptions, ilike, like, nilike, nlike}; use arrow::error::ArrowError; +use datafusion_common::utils::{normalize_float_zero, normalize_float_zero_scalar}; use datafusion_common::{Result, ScalarValue}; use datafusion_common::{arrow_datafusion_err, assert_or_internal_err, internal_err}; use datafusion_expr_common::columnar_value::ColumnarValue; @@ -84,7 +85,22 @@ pub fn apply_cmp( } }; - apply(lhs, rhs, |l, r| Ok(Arc::new(f(l, r)?))) + // Arrow's comparison kernels use IEEE 754 totalOrder semantics for + // floats, which treats `-0.0` and `+0.0` as distinct. Normalize float + // operands so SQL semantics (`+0.0 == -0.0`) hold. No-op for + // non-float types. + let lhs = normalize_cmp_input(lhs); + let rhs = normalize_cmp_input(rhs); + apply(&lhs, &rhs, |l, r| Ok(Arc::new(f(l, r)?))) + } +} + +fn normalize_cmp_input(cv: &ColumnarValue) -> ColumnarValue { + match cv { + ColumnarValue::Array(a) => ColumnarValue::Array(normalize_float_zero(a)), + ColumnarValue::Scalar(s) => { + ColumnarValue::Scalar(normalize_float_zero_scalar(s.clone())) + } } } diff --git a/datafusion/physical-expr-common/src/metrics/builder.rs b/datafusion/physical-expr-common/src/metrics/builder.rs index e9c0b76af2582..7d5a18f535369 100644 --- a/datafusion/physical-expr-common/src/metrics/builder.rs +++ b/datafusion/physical-expr-common/src/metrics/builder.rs @@ -25,13 +25,15 @@ use crate::metrics::{ }; use super::{ - Count, ExecutionPlanMetricsSet, Gauge, Label, Metric, MetricValue, Time, Timestamp, + Count, ExecutionPlanMetricsSet, Gauge, Label, LabelValue, Metric, MetricValue, Time, + Timestamp, }; /// Structure for constructing metrics, counters, timers, etc. /// /// Note the use of `Cow<..>` is to avoid allocations in the common -/// case of constant strings +/// case of constant strings. Dynamically created label strings are shared when +/// [`Label`] values are cloned. /// /// ```rust /// use datafusion_physical_expr_common::metrics::*; @@ -47,6 +49,7 @@ use super::{ /// .with_new_label("filename", "my_awesome_file.parquet") /// .counter("num_bytes", partition); /// ``` +#[derive(Clone)] pub struct MetricBuilder<'a> { /// Location that the metric created by this builder will be added do metrics: &'a ExecutionPlanMetricsSet, @@ -108,7 +111,10 @@ impl<'a> MetricBuilder<'a> { name: impl Into>, value: impl Into>, ) -> Self { - self.with_label(Label::new(name.into(), value.into())) + self.with_label(Label::new( + LabelValue::from(name.into()), + LabelValue::from(value.into()), + )) } /// Set the partition of the metric being constructed @@ -243,6 +249,23 @@ impl<'a> MetricBuilder<'a> { gauge } + /// Consumes self and creates a new [`Gauge`] for recording peak memory + /// usage in bytes. + pub fn peak_memory_usage( + self, + gauge_name: impl Into>, + partition: usize, + ) -> Gauge { + let gauge = Gauge::new(); + self.with_category(MetricCategory::Bytes) + .with_partition(partition) + .build(MetricValue::PeakMemoryUsage { + name: gauge_name.into(), + gauge: gauge.clone(), + }); + gauge + } + /// Consume self and create a new Timer for recording the elapsed /// CPU time spent by an operator pub fn elapsed_compute(self, partition: usize) -> Time { diff --git a/datafusion/physical-expr-common/src/metrics/mod.rs b/datafusion/physical-expr-common/src/metrics/mod.rs index eecd8cfabd5eb..146c039c75f6a 100644 --- a/datafusion/physical-expr-common/src/metrics/mod.rs +++ b/datafusion/physical-expr-common/src/metrics/mod.rs @@ -30,6 +30,7 @@ use parking_lot::Mutex; use std::{ borrow::Cow, fmt::{self, Debug, Display}, + hash::{Hash, Hasher}, sync::Arc, vec::IntoIter, }; @@ -308,6 +309,7 @@ impl MetricsSet { MetricValue::SpilledRows(_) => false, MetricValue::CurrentMemoryUsage(_) => false, MetricValue::Gauge { name, .. } => name == metric_name, + MetricValue::PeakMemoryUsage { name, .. } => name == metric_name, MetricValue::StartTimestamp(_) => false, MetricValue::EndTimestamp(_) => false, MetricValue::PruningMetrics { name, .. } => name == metric_name, @@ -416,6 +418,21 @@ impl MetricsSet { .collect::>(); Self { metrics } } + + /// Returns a new `MetricsSet` filtered by metric name. + /// Only metrics with the names appearing the list will be kept. + pub fn filter_by_names(self, names: &[String]) -> Self { + if names.is_empty() { + return Self { metrics: vec![] }; + } + + let metrics = self + .metrics + .into_iter() + .filter(|metric| names.iter().any(|name| name == metric.value().name())) + .collect::>(); + Self { metrics } + } } impl Display for MetricsSet { @@ -519,20 +536,19 @@ impl From for ExecutionPlanMetricsSet { /// telemetry], /// etc. /// -/// As the name and value are expected to mostly be constant strings, -/// use a [`Cow`] to avoid copying / allocations in this common case. +/// As the name and value are expected to often be constant strings, borrowed +/// static strings avoid allocations in that common case. Dynamic strings are +/// stored behind [`Arc`] so cloning labels does not copy the underlying +/// string data. #[derive(Debug, Clone, PartialEq, Eq, Hash)] pub struct Label { - name: Cow<'static, str>, - value: Cow<'static, str>, + name: LabelValue, + value: LabelValue, } impl Label { /// Create a new [`Label`] - pub fn new( - name: impl Into>, - value: impl Into>, - ) -> Self { + pub fn new(name: impl Into, value: impl Into) -> Self { let name = name.into(); let value = value.into(); Self { name, value } @@ -540,12 +556,12 @@ impl Label { /// Returns the name of this label pub fn name(&self) -> &str { - self.name.as_ref() + self.name.as_str() } /// Returns the value of this label pub fn value(&self) -> &str { - self.value.as_ref() + self.value.as_str() } } @@ -555,6 +571,89 @@ impl Display for Label { } } +/// A label name or value. +/// +/// String literals preserve the existing allocation-free path. Dynamic strings +/// can be stored behind [`Arc`], so cloning a [`Label`] only increments an +/// atomic reference count and does not allocate or copy the underlying string +/// data. +#[derive(Clone)] +pub struct LabelValue(LabelValueInner); + +/// Internal representation for label names and values. +/// +/// `LabelValue` is public because `Label::new` accepts it, but these storage +/// variants are implementation details. Keeping them private prevents external +/// code from constructing or matching on `Static` and `Shared` directly. +#[derive(Clone)] +enum LabelValueInner { + Static(&'static str), + Shared(Arc), +} + +impl LabelValue { + /// Return this label value as a string slice. + pub fn as_str(&self) -> &str { + match &self.0 { + LabelValueInner::Static(value) => value, + LabelValueInner::Shared(value) => value.as_ref(), + } + } +} + +impl From<&'static str> for LabelValue { + fn from(value: &'static str) -> Self { + Self(LabelValueInner::Static(value)) + } +} + +impl From for LabelValue { + fn from(value: String) -> Self { + Self(LabelValueInner::Shared(Arc::from(value))) + } +} + +impl From> for LabelValue { + fn from(value: Arc) -> Self { + Self(LabelValueInner::Shared(value)) + } +} + +impl From> for LabelValue { + fn from(value: Cow<'static, str>) -> Self { + match value { + Cow::Borrowed(value) => value.into(), + Cow::Owned(value) => value.into(), + } + } +} + +impl PartialEq for LabelValue { + fn eq(&self, other: &Self) -> bool { + self.as_str() == other.as_str() + } +} + +impl Eq for LabelValue {} + +impl Hash for LabelValue { + fn hash(&self, state: &mut H) { + self.as_str().hash(state); + } +} + +impl Debug for LabelValue { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Debug::fmt(self.as_str(), f) + } +} + +impl Display for LabelValue { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + Display::fmt(self.as_str(), f) + } +} + #[cfg(test)] mod tests { use std::time::Duration; @@ -609,6 +708,18 @@ mod tests { assert_eq!("output_rows{partition=2, foo=bar}=66", metric.to_string()) } + #[test] + fn test_label_owned_and_borrowed_values_are_equal() { + let borrowed = Label::new("foo", "bar"); + let owned = Label::new("foo".to_string(), "bar".to_string()); + let shared = Label::new("foo", Arc::::from("bar")); + + assert_eq!(borrowed, owned); + assert_eq!(borrowed, shared); + assert_eq!(borrowed.to_string(), owned.to_string()); + assert_eq!(borrowed.to_string(), shared.to_string()); + } + #[test] fn test_output_rows() { let metrics = ExecutionPlanMetricsSet::new(); @@ -870,4 +981,29 @@ mod tests { metric_names(&metrics) ); } + + #[test] + fn test_filter_by_names() { + let metrics = ExecutionPlanMetricsSet::new(); + MetricBuilder::new(&metrics).output_rows(0); + MetricBuilder::new(&metrics).counter("custom_counter", 0); + + assert!( + metrics + .clone_inner() + .filter_by_names(&[]) + .iter() + .next() + .is_none() + ); + + let names = vec!["output_rows".to_string()]; + let filtered = metrics.clone_inner().filter_by_names(&names); + + assert_eq!(filtered.iter().count(), 1); + assert_eq!( + filtered.iter().next().unwrap().value().name(), + "output_rows" + ); + } } diff --git a/datafusion/physical-expr-common/src/metrics/value.rs b/datafusion/physical-expr-common/src/metrics/value.rs index ef0087c20d91f..232fefcc5f47e 100644 --- a/datafusion/physical-expr-common/src/metrics/value.rs +++ b/datafusion/physical-expr-common/src/metrics/value.rs @@ -672,6 +672,13 @@ pub enum MetricValue { /// The value of the metric gauge: Gauge, }, + /// Operator defined peak memory usage in bytes. + PeakMemoryUsage { + /// The provided name of this metric + name: Cow<'static, str>, + /// The value of the metric + gauge: Gauge, + }, /// Operator defined time Time { /// The provided name of this metric @@ -744,6 +751,13 @@ impl PartialEq for MetricValue { name: other_name, gauge: other_gauge, }, + ) + | ( + MetricValue::PeakMemoryUsage { name, gauge }, + MetricValue::PeakMemoryUsage { + name: other_name, + gauge: other_gauge, + }, ) => name == other_name && gauge == other_gauge, ( MetricValue::Time { name, time }, @@ -810,7 +824,9 @@ impl MetricValue { Self::CurrentMemoryUsage(_) => "mem_used", Self::ElapsedCompute(_) => "elapsed_compute", Self::Count { name, .. } => name.borrow(), - Self::Gauge { name, .. } => name.borrow(), + Self::Gauge { name, .. } | Self::PeakMemoryUsage { name, .. } => { + name.borrow() + } Self::Time { name, .. } => name.borrow(), Self::StartTimestamp(_) => "start_timestamp", Self::EndTimestamp(_) => "end_timestamp", @@ -833,7 +849,9 @@ impl MetricValue { Self::CurrentMemoryUsage(used) => used.value(), Self::ElapsedCompute(time) => time.value(), Self::Count { count, .. } => count.value(), - Self::Gauge { gauge, .. } => gauge.value(), + Self::Gauge { gauge, .. } | Self::PeakMemoryUsage { gauge, .. } => { + gauge.value() + } Self::Time { time, .. } => time.value(), Self::StartTimestamp(timestamp) => timestamp .value() @@ -875,6 +893,10 @@ impl MetricValue { name: name.clone(), gauge: Gauge::new(), }, + Self::PeakMemoryUsage { name, .. } => Self::PeakMemoryUsage { + name: name.clone(), + gauge: Gauge::new(), + }, Self::Time { name, .. } => Self::Time { name: name.clone(), time: Time::new(), @@ -933,6 +955,12 @@ impl MetricValue { Self::Gauge { gauge: other_gauge, .. }, + ) + | ( + Self::PeakMemoryUsage { gauge, .. }, + Self::PeakMemoryUsage { + gauge: other_gauge, .. + }, ) => gauge.add(other_gauge.value()), (Self::ElapsedCompute(time), Self::ElapsedCompute(other_time)) | ( @@ -1029,6 +1057,7 @@ impl MetricValue { "page_index_pages_skipped_by_fully_matched" => 8, _ => 14, }, + Self::PeakMemoryUsage { .. } => 13, Self::Gauge { .. } => 15, Self::Time { .. } => 16, Self::Ratio { .. } => 17, @@ -1064,6 +1093,10 @@ impl Display for MetricValue { let readable_size = human_readable_size(gauge.value()); write!(f, "{readable_size}") } + Self::PeakMemoryUsage { gauge, .. } => { + let readable_size = human_readable_size(gauge.value()); + write!(f, "{readable_size}") + } Self::Gauge { gauge, .. } => { // Generic gauge metrics - format with human-readable count write!(f, "{}", human_readable_count(gauge.value())) @@ -1525,6 +1558,18 @@ mod tests { "100.0 MB" ); + // Test PeakMemoryUsage formatting (should use size, not count) + let peak_mem_gauge = Gauge::new(); + peak_mem_gauge.add(100 * MB as usize); + assert_eq!( + MetricValue::PeakMemoryUsage { + name: "peak_mem_used".into(), + gauge: peak_mem_gauge.clone() + } + .to_string(), + "100.0 MB" + ); + // Test custom Gauge formatting (should use count) let custom_gauge = Gauge::new(); custom_gauge.add(50_000); diff --git a/datafusion/physical-expr-common/src/physical_expr.rs b/datafusion/physical-expr-common/src/physical_expr.rs index 7b3f7dcc76c87..59393e75786bd 100644 --- a/datafusion/physical-expr-common/src/physical_expr.rs +++ b/datafusion/physical-expr-common/src/physical_expr.rs @@ -477,6 +477,304 @@ pub trait PhysicalExpr: Any + Send + Sync + Display + Debug + DynEq + DynHash { fn expression_id(&self) -> Option { None } + + /// Serialize this expression to a [`PhysicalExprNode`] proto message. + /// + /// Returning `Ok(None)` means "this expression does not know how to + /// serialize itself"; the caller (typically `datafusion-proto`) will fall + /// back to its existing codec / extension paths. This matches today's + /// behavior for expressions that aren't built into `datafusion-proto`. + /// + /// Returning `Ok(Some(node))` means the expression has serialized itself + /// fully; the caller should not try any further fallback path. + /// + /// Returning `Err(_)` means a real serialization failure (e.g. the + /// expression knows it should serialize but a child failed). + /// + /// The motivating use case is letting expressions with private state + /// (e.g. `DynamicFilterPhysicalExpr`'s `RwLock`-protected inner fields) + /// reach into their own internals for `try_to_proto`/`try_from_proto` + /// without having to expose `pub` accessors to `datafusion-proto`. See + /// . + /// + /// The `try_` prefix matches the fallible `try_from_proto` decode + /// constructors; both sides of the round-trip are fallible and named + /// consistently. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } +} + +/// Encode-side context for [`PhysicalExpr::try_to_proto`]. +/// +/// Expression authors only ever see [`proto_encode::PhysicalExprEncodeCtx`]: +/// a concrete struct with stable methods. Internally it dispatches to a +/// [`proto_encode::PhysicalExprEncode`] implementor that lives in +/// `datafusion-proto`, which is what lets `physical-expr-common` stay free +/// of `datafusion-proto` as a dep. +/// +/// More specialized helpers (e.g. encoding UDFs/UDAFs/UDWFs through the +/// extension codec) can be added to the context as expressions migrate; +/// today they're not required because the encoder forwards to the existing +/// codec via the proto converter. +#[cfg(feature = "proto")] +pub mod proto_encode { + use std::sync::Arc; + + use datafusion_common::Result; + use datafusion_proto_models::protobuf::PhysicalExprNode; + + use super::PhysicalExpr; + + /// Encoder context handed to [`super::PhysicalExpr::try_to_proto`]. + /// + /// Wraps an internal [`PhysicalExprEncode`] trait object so callers see a + /// stable concrete type while implementations can evolve in + /// `datafusion-proto`. + pub struct PhysicalExprEncodeCtx<'a> { + encoder: &'a dyn PhysicalExprEncode, + } + + impl<'a> PhysicalExprEncodeCtx<'a> { + /// Construct a new encode context. Typically called by + /// `datafusion-proto`; expression authors receive `&PhysicalExprEncodeCtx`. + pub fn new(encoder: &'a dyn PhysicalExprEncode) -> Self { + Self { encoder } + } + + /// Encode a child expression. Routes through the configured encoder + /// so dedup-aware encoding is preserved. + pub fn encode_child( + &self, + expr: &Arc, + ) -> Result { + self.encoder.encode(expr) + } + + /// Encode a sequence of child expressions, preserving order. + /// + /// Convenience wrapper over [`Self::encode_child`] for expressions + /// holding a `repeated` proto field (e.g. the `list` of an `InList`). + /// The first encode error short-circuits. + pub fn encode_children_expressions<'b, I>( + &self, + exprs: I, + ) -> Result> + where + I: IntoIterator>, + { + exprs + .into_iter() + .map(|expr| self.encode_child(expr)) + .collect() + } + } + + /// Internal dispatch trait. Implementors live in `datafusion-proto` and + /// wrap the existing `PhysicalExtensionCodec` + + /// `PhysicalProtoConverterExtension` plumbing. Expression authors should + /// use [`PhysicalExprEncodeCtx`] instead of calling this directly. + pub trait PhysicalExprEncode { + /// Encode an expression to a protobuf node. + fn encode(&self, expr: &Arc) -> Result; + } +} + +/// Decode-side counterpart to [`proto_encode`]. +/// +/// Expression authors implement an associated `try_from_proto` on their +/// concrete type, with the signature +/// +/// ```ignore +/// fn try_from_proto( +/// node: &PhysicalExprNode, +/// ctx: &PhysicalExprDecodeCtx<'_>, +/// ) -> Result> +/// ``` +/// +/// It takes the whole [`PhysicalExprNode`] — the exact inverse of what +/// [`PhysicalExpr::try_to_proto`] returns — so the constructor can also see +/// outer-node fields such as `expr_id`. The central match in +/// `datafusion-proto` dispatches `ExprType` variants to these constructors. +/// +/// As with the encode side, the public surface is a struct (not a `&dyn` +/// trait) so future fields/helpers (registries for third-party expressions, +/// schema-resolution caches, etc.) can be added without changing the +/// signature every expression depends on. +/// +/// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode +#[cfg(feature = "proto")] +pub mod proto_decode { + use std::sync::Arc; + + use arrow::datatypes::Schema; + use datafusion_common::Result; + use datafusion_proto_models::protobuf::PhysicalExprNode; + + use super::PhysicalExpr; + + /// Open the outer [`PhysicalExprNode`] and assert it carries the expected + /// `ExprType` variant, returning the inner payload (auto-derefs through + /// `Box`) or bailing with an `Internal` error. + /// + /// Every `try_from_proto` starts with the same six-line `match`: + /// + /// ```ignore + /// let try_cast = match &node.expr_type { + /// Some(protobuf::physical_expr_node::ExprType::TryCast(x)) => x.as_ref(), + /// _ => return internal_err!("PhysicalExprNode is not a TryCastExpr"), + /// }; + /// ``` + /// + /// With this macro that collapses to: + /// + /// ```ignore + /// let try_cast = expect_expr_variant!( + /// node, + /// protobuf::physical_expr_node::ExprType::TryCast, + /// "TryCastExpr", + /// ); + /// ``` + /// + /// Pass the variant as a `::` path so the macro stays agnostic to how + /// the caller imports the proto types. + #[macro_export] + macro_rules! expect_expr_variant { + ($node:expr, $variant:path, $expr_name:literal $(,)?) => {{ + match &$node.expr_type { + ::core::option::Option::Some($variant(inner)) => inner, + _ => { + return ::datafusion_common::internal_err!(concat!( + "PhysicalExprNode is not a ", + $expr_name + )); + } + } + }}; + } + #[doc(inline)] + pub use expect_expr_variant; + + /// Decoder context handed to per-expression `try_from_proto` constructors. + /// + /// Wraps an internal [`PhysicalExprDecode`] trait object plus a borrowed + /// schema. The trait stays an implementation detail of `datafusion-proto`; + /// expression authors only see this struct. + pub struct PhysicalExprDecodeCtx<'a> { + schema: &'a Schema, + decoder: &'a dyn PhysicalExprDecode, + } + + impl<'a> PhysicalExprDecodeCtx<'a> { + /// Construct a new decode context. Typically called by + /// `datafusion-proto`; expression authors receive + /// `&PhysicalExprDecodeCtx`. + pub fn new(schema: &'a Schema, decoder: &'a dyn PhysicalExprDecode) -> Self { + Self { schema, decoder } + } + + /// The schema bound to this decode context. Use it for column lookups, + /// data-type resolution, etc. + pub fn schema(&self) -> &Schema { + self.schema + } + + /// Decode an expression node, recursing into child sub-expressions. + /// + /// Routes built-in `ExprType` variants through `datafusion-proto`'s + /// central match and forwards extension nodes to the registered codec + /// (today via [`PhysicalExtensionCodec::try_decode_expr`]; later via + /// a per-type registry — see #21835). + /// + /// [`PhysicalExtensionCodec::try_decode_expr`]: https://docs.rs/datafusion-proto/latest/datafusion_proto/physical_plan/trait.PhysicalExtensionCodec.html#method.try_decode_expr + pub fn decode(&self, node: &PhysicalExprNode) -> Result> { + self.decoder.decode(node, self.schema) + } + + /// Decode a required child node, erroring if it is absent. + /// + /// Proto child expressions are encoded as `Option>`; + /// pass the field directly (e.g. `node.expr.as_deref()`). `expr_name` + /// is the expression being decoded (e.g. `"InListExpr"`) and `field` + /// the proto field (e.g. `"expr"`); both are woven into the error so + /// it names *where* the missing field is, without each author + /// hand-rolling the string. + pub fn decode_required_expression( + &self, + node: Option<&PhysicalExprNode>, + expr_name: &str, + field: &str, + ) -> Result> { + let node = node.ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "{expr_name} is missing required field '{field}'" + ) + })?; + self.decode(node) + } + + /// Decode a sequence of child nodes, preserving order. + /// + /// Convenience wrapper over [`Self::decode`] for expressions holding a + /// `repeated` proto field (e.g. the `list` of an `InList`). The first + /// decode error short-circuits. + pub fn decode_children_expressions<'b, I>( + &self, + nodes: I, + ) -> Result>> + where + I: IntoIterator, + { + nodes.into_iter().map(|node| self.decode(node)).collect() + } + } + + /// Unwrap a required non-expression proto field. + /// + /// Mirrors [`PhysicalExprDecodeCtx::decode_required_expression`] for proto + /// fields that aren't [`PhysicalExprNode`]s — e.g. the `arrow_type` of a + /// `PhysicalCastNode` or the `scalar` of a `PhysicalLiteralNode`. Keeps + /// the "missing required field" message format identical across + /// expressions: + /// + /// ```ignore + /// let arrow_type = require_proto_field( + /// cast_expr.arrow_type.as_ref(), + /// "CastExpr", + /// "arrow_type", + /// )?; + /// ``` + pub fn require_proto_field( + opt: Option, + expr_name: &str, + field: &str, + ) -> Result { + opt.ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "{expr_name} is missing required field '{field}'" + ) + }) + } + + /// Internal dispatch trait. Implementors live in `datafusion-proto`. + /// Expression authors should use [`PhysicalExprDecodeCtx`] instead of + /// calling this directly. + pub trait PhysicalExprDecode { + /// Decode a proto node into a concrete `PhysicalExpr`. The schema is + /// passed alongside so implementations can support recursive children + /// and rebind the context per call (e.g. for nested plans). + fn decode( + &self, + node: &PhysicalExprNode, + schema: &Schema, + ) -> Result>; + } } #[deprecated( @@ -696,6 +994,12 @@ pub fn snapshot_generation(expr: &Arc) -> u64 { /// Check if the given `PhysicalExpr` is dynamic. /// Internally this calls [`snapshot_generation`] to check if the generation is non-zero, /// any dynamic `PhysicalExpr` should have a non-zero generation. +#[deprecated( + since = "55.0.0", + note = "Downcast to `DynamicFilterPhysicalExpr`, or use \ + `DynamicFilterTracking::classify(expr).contains_dynamic_filter()` from \ + `datafusion_physical_expr`" +)] pub fn is_dynamic_physical_expr(expr: &Arc) -> bool { // If the generation is non-zero, then this `PhysicalExpr` is dynamic. snapshot_generation(expr) != 0 @@ -908,3 +1212,83 @@ mod test { ); } } + +#[cfg(all(test, feature = "proto"))] +mod proto_helper_tests { + use datafusion_common::DataFusionError; + use datafusion_proto_models::protobuf::{ + self, PhysicalColumn, PhysicalExprNode, physical_expr_node, + }; + + use crate::expect_expr_variant; + use crate::physical_expr::proto_decode::require_proto_field; + + fn column_node() -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::Column(PhysicalColumn { + name: "a".to_string(), + index: 0, + })), + } + } + + #[test] + fn require_proto_field_returns_inner() { + let v = require_proto_field(Some(7_u32), "FooExpr", "answer").unwrap(); + assert_eq!(v, 7); + } + + #[test] + fn require_proto_field_reports_missing() { + let err = require_proto_field::(None, "FooExpr", "answer").unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) + if msg.contains("FooExpr is missing required field 'answer'") + )); + } + + fn expect_column( + node: &PhysicalExprNode, + ) -> Result<&PhysicalColumn, DataFusionError> { + let inner = + expect_expr_variant!(node, physical_expr_node::ExprType::Column, "Column",); + Ok(inner) + } + + #[test] + fn expect_expr_variant_returns_inner_payload() { + let node = column_node(); + let col = expect_column(&node).unwrap(); + assert_eq!(col.name, "a"); + } + + #[test] + fn expect_expr_variant_rejects_wrong_variant() { + let node = PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::Negative(Box::new( + protobuf::PhysicalNegativeNode { expr: None }, + ))), + }; + let err = expect_column(&node).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a Column") + )); + } + + #[test] + fn expect_expr_variant_rejects_missing_expr_type() { + let node = PhysicalExprNode { + expr_id: None, + expr_type: None, + }; + let err = expect_column(&node).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a Column") + )); + } +} diff --git a/datafusion/physical-expr-common/src/sort_expr.rs b/datafusion/physical-expr-common/src/sort_expr.rs index 84ffb92eaa600..6e8dbccdb7c0e 100644 --- a/datafusion/physical-expr-common/src/sort_expr.rs +++ b/datafusion/physical-expr-common/src/sort_expr.rs @@ -183,6 +183,122 @@ impl PhysicalSortExpr { } } +/// Protobuf conversions for [`PhysicalSortExpr`]. +/// +/// This is the flat [`PhysicalSortExprNode`] representation used wherever the +/// wire format stores an ordering (scan output orderings, range partitioning, +/// window frames, …). It is *not* the `PhysicalExprNode::Sort` wrapping that +/// `SortExec` uses for its own `expr` field. +/// +/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode +#[cfg(feature = "proto")] +impl PhysicalSortExpr { + /// Serialize this sort expression, encoding its child expression through + /// `ctx`. + pub fn try_to_proto( + &self, + ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result { + Ok(datafusion_proto_models::protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_child(&self.expr)?)), + asc: !self.options.descending, + nulls_first: self.options.nulls_first, + }) + } + + /// Reconstruct a [`PhysicalSortExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalSortExprNode, + ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result { + let expr = ctx.decode_required_expression( + node.expr.as_deref(), + "PhysicalSortExpr", + "expr", + )?; + Ok(PhysicalSortExpr { + expr, + options: SortOptions { + descending: !node.asc, + nulls_first: node.nulls_first, + }, + }) + } +} + +/// Serialize a sequence of sort expressions into the flat +/// [`PhysicalSortExprNode`] list the wire format uses for an ordering. +/// +/// Accepts anything that yields [`PhysicalSortExpr`]s by value or by reference, +/// so a [`LexOrdering`], a `&[PhysicalSortExpr]`, or a [`LexRequirement`] +/// mapped through [`PhysicalSortExpr::from`] all work: +/// +/// ```ignore +/// let nodes = sort_exprs_try_to_proto(ordering.iter(), ctx)?; +/// let nodes = sort_exprs_try_to_proto( +/// requirement.iter().map(|req| PhysicalSortExpr::from(req.clone())), +/// ctx, +/// )?; +/// ``` +/// +/// The `PhysicalSortExprNodeCollection` message some plans use is just this +/// list in a wrapper, so those callers wrap the result themselves rather than +/// this function guessing which shape they mean. +/// +/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode +#[cfg(feature = "proto")] +pub fn sort_exprs_try_to_proto>( + exprs: impl IntoIterator, + ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, +) -> Result> { + exprs + .into_iter() + .map(|expr| expr.borrow().try_to_proto(ctx)) + .collect() +} + +/// Reconstruct a sequence of sort expressions from the flat +/// [`PhysicalSortExprNode`] list, the counterpart of +/// [`sort_exprs_try_to_proto`]. +/// +/// Returns the expressions rather than a [`LexOrdering`] or a +/// [`LexRequirement`], because callers differ in what an empty list means: +/// `LexOrdering::new` / `LexRequirement::new` return `None` for it, which is +/// "no ordering declared" for a scan and an error for an operator that requires +/// one. Callers with the former convention can use +/// [`optional_ordering_try_from_proto`] instead. +/// +/// [`PhysicalSortExprNode`]: datafusion_proto_models::protobuf::PhysicalSortExprNode +#[cfg(feature = "proto")] +pub fn sort_exprs_try_from_proto( + nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode], + ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, +) -> Result> { + nodes + .iter() + .map(|node| PhysicalSortExpr::try_from_proto(node, ctx)) + .collect() +} + +/// Serialize an optional [`LexOrdering`], encoding `None` as an empty list. +#[cfg(feature = "proto")] +pub fn optional_ordering_try_to_proto( + ordering: Option<&LexOrdering>, + ctx: &crate::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, +) -> Result> { + sort_exprs_try_to_proto(ordering.into_iter().flatten(), ctx) +} + +/// Counterpart of [`optional_ordering_try_to_proto`]: an empty list decodes +/// as `None`. +#[cfg(feature = "proto")] +pub fn optional_ordering_try_from_proto( + nodes: &[datafusion_proto_models::protobuf::PhysicalSortExprNode], + ctx: &crate::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, +) -> Result> { + Ok(LexOrdering::new(sort_exprs_try_from_proto(nodes, ctx)?)) +} + impl PartialEq for PhysicalSortExpr { fn eq(&self, other: &Self) -> bool { self.options == other.options && self.expr.eq(&other.expr) diff --git a/datafusion/physical-expr-common/src/utils.rs b/datafusion/physical-expr-common/src/utils.rs index e469885f83316..5dadcdcabb180 100644 --- a/datafusion/physical-expr-common/src/utils.rs +++ b/datafusion/physical-expr-common/src/utils.rs @@ -370,21 +370,21 @@ fn scatter_fallback( let mut true_pos = 0; let mask_array = BooleanArray::new(mask.clone(), None); - SlicesIterator::new(&mask_array).for_each(|(start, end)| { + for (start, end) in SlicesIterator::new(&mask_array) { // the gap needs to be filled with nulls if start > filled { - mutable.extend_nulls(start - filled); + mutable.try_extend_nulls(start - filled)?; } // fill with truthy values let len = end - start; - mutable.extend(0, true_pos, true_pos + len); + mutable.try_extend(0, true_pos, true_pos + len)?; true_pos += len; filled = end; - }); + } // the remaining part is falsy if filled < output_len { - mutable.extend_nulls(output_len - filled); + mutable.try_extend_nulls(output_len - filled)?; } let data = mutable.freeze(); @@ -614,11 +614,9 @@ mod tests { #[test] fn scatter_fixed_size_binary_test() -> Result<()> { - let truthy = Arc::new(FixedSizeBinaryArray::from(vec![ - &[1u8, 2][..], - &[3, 4][..], - &[5, 6][..], - ])); + let truthy = Arc::new(FixedSizeBinaryArray::try_from_iter( + vec![&[1u8, 2][..], &[3, 4][..], &[5, 6][..]].into_iter(), + )?); let mask = BooleanArray::from(vec![true, false, true, false, true]); let result = scatter(&mask, truthy.as_ref())?; diff --git a/datafusion/physical-expr/Cargo.toml b/datafusion/physical-expr/Cargo.toml index b755353d75658..65ef2a3ceb216 100644 --- a/datafusion/physical-expr/Cargo.toml +++ b/datafusion/physical-expr/Cargo.toml @@ -42,6 +42,12 @@ name = "datafusion_physical_expr" [features] recursive_protection = ["dep:recursive"] +# Forwards the `proto` feature to `datafusion-physical-expr-common`, exposing +# `PhysicalExpr::to_proto` and letting expressions in this crate implement it. +proto = [ + "dep:datafusion-proto-models", + "datafusion-physical-expr-common/proto", +] [dependencies] arrow = { workspace = true } @@ -50,6 +56,7 @@ datafusion-expr = { workspace = true } datafusion-expr-common = { workspace = true } datafusion-functions-aggregate-common = { workspace = true } datafusion-physical-expr-common = { workspace = true } +datafusion-proto-models = { workspace = true, optional = true } hashbrown = { workspace = true } indexmap = { workspace = true } itertools = { workspace = true, features = ["use_std"] } diff --git a/datafusion/physical-expr/benches/binary_op.rs b/datafusion/physical-expr/benches/binary_op.rs index 99fc40fa1c91b..f170561652070 100644 --- a/datafusion/physical-expr/benches/binary_op.rs +++ b/datafusion/physical-expr/benches/binary_op.rs @@ -15,11 +15,11 @@ // specific language governing permissions and limitations // under the License. +use arrow::{array::StringArray, record_batch::RecordBatch}; use arrow::{ - array::BooleanArray, + array::{BooleanArray, Date32Array, Date64Array}, datatypes::{DataType, Field, Schema}, }; -use arrow::{array::StringArray, record_batch::RecordBatch}; use criterion::{Criterion, criterion_group, criterion_main}; use datafusion_expr::{Operator, and, binary_expr, col, lit, or}; use datafusion_physical_expr::{ @@ -30,6 +30,9 @@ use datafusion_physical_expr::{ use std::hint::black_box; use std::sync::Arc; +const DATE_ARRAY_LEN: usize = 8192; +const MILLIS_PER_DAY: i64 = 86_400_000; + /// Generates BooleanArrays with different true/false distributions for benchmarking. /// /// Returns a vector of tuples containing scenario name and corresponding BooleanArray. @@ -309,6 +312,81 @@ fn create_record_batch( Ok(rbs) } -criterion_group!(benches, benchmark_binary_op_in_short_circuit); +fn make_date32_batch(null_percent: f64) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Date32, true), + Field::new("b", DataType::Date32, true), + ])); + + let left = Date32Array::from_iter((0..DATE_ARRAY_LEN).map(|i| { + (null_percent == 0.0 || i % (1.0 / null_percent) as usize != 0) + .then_some(18_000 + i as i32) + })); + let right = Date32Array::from_iter((0..DATE_ARRAY_LEN).map(|i| { + (null_percent == 0.0 || i % (1.0 / null_percent) as usize != 0) + .then_some(17_000 + (i % 365) as i32) + })); + + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(left), Arc::new(right)]) + .unwrap() +} + +fn make_date64_batch(null_percent: f64) -> RecordBatch { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Date64, true), + Field::new("b", DataType::Date64, true), + ])); + + let left = Date64Array::from_iter((0..DATE_ARRAY_LEN).map(|i| { + (null_percent == 0.0 || i % (1.0 / null_percent) as usize != 0) + .then_some((18_000 + i as i64) * MILLIS_PER_DAY) + })); + let right = Date64Array::from_iter((0..DATE_ARRAY_LEN).map(|i| { + (null_percent == 0.0 || i % (1.0 / null_percent) as usize != 0) + .then_some((17_000 + (i % 365) as i64) * MILLIS_PER_DAY) + })); + + RecordBatch::try_new(Arc::clone(&schema), vec![Arc::new(left), Arc::new(right)]) + .unwrap() +} + +/// Benchmark Date32 column subtraction. +fn benchmark_date32_subtract(c: &mut Criterion) { + for (name, null_percent) in [("no_nulls", 0.0), ("20_percent_nulls", 0.2)] { + let batch = make_date32_batch(null_percent); + let expr = BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Minus, + Arc::new(Column::new("b", 1)), + ); + + c.bench_function(&format!("date32_subtract/{name}"), |b| { + b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())) + }); + } +} + +/// Benchmark Date64 column subtraction. +fn benchmark_date64_subtract(c: &mut Criterion) { + for (name, null_percent) in [("no_nulls", 0.0), ("20_percent_nulls", 0.2)] { + let batch = make_date64_batch(null_percent); + let expr = BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Minus, + Arc::new(Column::new("b", 1)), + ); + + c.bench_function(&format!("date64_subtract/{name}"), |b| { + b.iter(|| black_box(expr.evaluate(black_box(&batch)).unwrap())) + }); + } +} + +criterion_group!( + benches, + benchmark_binary_op_in_short_circuit, + benchmark_date32_subtract, + benchmark_date64_subtract +); criterion_main!(benches); diff --git a/datafusion/physical-expr/benches/in_list_strategy.rs b/datafusion/physical-expr/benches/in_list_strategy.rs index 5c4922fdcf8a9..c69af192b9cdd 100644 --- a/datafusion/physical-expr/benches/in_list_strategy.rs +++ b/datafusion/physical-expr/benches/in_list_strategy.rs @@ -34,9 +34,10 @@ //! | Case | Types | Characteristics | List Sizes Tested | //! |------|-------|-----------------|-------------------| //! | Narrow integer cases | UInt8 | small value domain | 4, 16 | -//! | Narrow integer cases | Int16 | larger value domain | 4, 64, 256 | +//! | Narrow integer cases | Int16, Float16 | larger value domain | 4, 64, 256 | //! | 32-bit primitive cases | Int32, Float32 | small and large lists | 4, 32, 64, 256 | //! | 64-bit primitive cases | Int64, TimestampNs | small and large lists | 4, 16, 32, 128 | +//! | 128-bit interval cases | IntervalMonthDayNano | small lists | 4 | //! | Utf8 short-string cases | Utf8 | 8-byte strings | 4, 64, 256 | //! | Utf8 long-string cases | Utf8 | 24-byte strings | 4, 64, 256 | //! | Utf8View short-string cases | Utf8View | 8-byte strings | 4, 16, 64, 256 | @@ -45,12 +46,14 @@ //! | Shared-prefix string cases | Utf8, Utf8View | same prefix, different suffix | 16, 32, 64 | //! | Fixed-size binary cases | FixedSizeBinary(16) | fixed-width binary values | 4, 64, 256, 10000 | +use arrow::array::types::IntervalMonthDayNano; use arrow::array::*; -use arrow::datatypes::{Field, Int32Type, Schema}; +use arrow::datatypes::{Field, Int32Type, IntervalMonthDayNanoType, Schema}; use arrow::record_batch::RecordBatch; use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; use datafusion_common::ScalarValue; use datafusion_physical_expr::expressions::{col, in_list, lit}; +use half::f16; use rand::distr::Alphanumeric; use rand::prelude::*; use std::sync::Arc; @@ -392,6 +395,23 @@ fn bench_narrow_integer(c: &mut Criterion) { ); } } + + // Float16: same 65,536-value bit-pattern domain as Int16/UInt16. + for list_size in [4, 64, 256] { + for match_pct in MATCH_RATES { + bench_numeric::( + c, + "narrow_integer", + &format!("f16/list={list_size}/match={match_pct}%"), + &NumericBenchConfig::new( + list_size, + match_pct as f64 / 100.0, + |rng| f16::from_f32(rng.random::() * 1000.0), + |v| ScalarValue::Float16(Some(v)), + ), + ); + } + } } // ============================================================================= @@ -510,6 +530,28 @@ fn bench_timestamp_ns(c: &mut Criterion) { } } +fn bench_interval_month_day_nano(c: &mut Criterion) { + for match_pct in MATCH_RATES { + bench_numeric::( + c, + "interval_month_day_nano", + &format!("small_list/list=4/match={match_pct}%"), + &NumericBenchConfig::new( + 4, + match_pct as f64 / 100.0, + |rng| { + IntervalMonthDayNanoType::make_value( + rng.random_range(-120..=120), + rng.random_range(-31..=31), + rng.random_range(-1_000_000_000..=1_000_000_000), + ) + }, + |v| ScalarValue::IntervalMonthDayNano(Some(v)), + ), + ); + } +} + // ============================================================================= // UTF8 STRING CASE BENCHMARKS // ============================================================================= @@ -993,7 +1035,7 @@ fn bench_fixed_size_binary_inner( .collect(); let refs: Vec<&[u8]> = values.iter().map(|v| v.as_slice()).collect(); - let array = FixedSizeBinaryArray::from(refs); + let array = FixedSizeBinaryArray::try_from_iter(refs.into_iter()).unwrap(); let schema = Schema::new(vec![Field::new("a", array.data_type().clone(), true)]); let exprs: Vec<_> = haystack @@ -1031,7 +1073,7 @@ fn bench_fixed_size_binary(c: &mut Criterion) { criterion_group! { name = benches; config = Criterion::default(); - targets = bench_narrow_integer, bench_primitive, bench_f32, bench_timestamp_ns, bench_utf8, bench_utf8view, bench_dictionary, bench_nulls, bench_fixed_size_binary + targets = bench_narrow_integer, bench_primitive, bench_f32, bench_timestamp_ns, bench_interval_month_day_nano, bench_utf8, bench_utf8view, bench_dictionary, bench_nulls, bench_fixed_size_binary } criterion_main!(benches); diff --git a/datafusion/physical-expr/src/aggregate.rs b/datafusion/physical-expr/src/aggregate.rs index e5d55aba4f51c..013779cf8c102 100644 --- a/datafusion/physical-expr/src/aggregate.rs +++ b/datafusion/physical-expr/src/aggregate.rs @@ -51,6 +51,7 @@ use datafusion_expr::execution_props::ExecutionProps; use datafusion_expr::expr::{ AggregateFunction, AggregateFunctionParams, NullTreatment, physical_name, }; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{AggregateUDF, Expr, ReversedUDAF, SetMonotonicity}; use datafusion_expr_common::accumulator::Accumulator; use datafusion_expr_common::groups_accumulator::GroupsAccumulator; @@ -423,6 +424,7 @@ pub struct LoweredAggregateBuilder<'a> { logical_input_schema: &'a DFSchema, physical_input_schema: &'a Schema, execution_props: &'a ExecutionProps, + planning_ctx: &'a PhysicalPlanningContext, } impl<'a> LoweredAggregateBuilder<'a> { @@ -430,12 +432,17 @@ impl<'a> LoweredAggregateBuilder<'a> { /// /// `logical_input_schema` is used to resolve logical expressions such as /// columns, while `physical_input_schema` is the input schema used by the - /// physical aggregate expression. + /// physical aggregate expression. `planning_ctx` is used when creating + /// physical expressions that reference uncorrelated scalar subqueries. + /// Callers creating physical aggregates outside of physical planning should + /// pass `&PhysicalPlanningContext::default()`, in which case converting a + /// scalar-subquery expression returns a planning error. pub fn new( expr: &'a Expr, logical_input_schema: &'a DFSchema, physical_input_schema: &'a Schema, execution_props: &'a ExecutionProps, + planning_ctx: &'a PhysicalPlanningContext, ) -> Self { Self { expr, @@ -446,6 +453,7 @@ impl<'a> LoweredAggregateBuilder<'a> { logical_input_schema, physical_input_schema, execution_props, + planning_ctx, } } @@ -484,6 +492,7 @@ impl<'a> LoweredAggregateBuilder<'a> { logical_input_schema, physical_input_schema, execution_props, + planning_ctx, } = self; let (name, human_display, output_metadata, expr) = lower_aggregate_display( @@ -515,16 +524,29 @@ impl<'a> LoweredAggregateBuilder<'a> { physical_name(&expr)? }; - let physical_args = - create_physical_exprs(args, logical_input_schema, execution_props)?; + let physical_args = create_physical_exprs( + args, + logical_input_schema, + execution_props, + planning_ctx, + )?; let filter = filter .as_ref() .map(|filter| { - create_physical_expr(filter, logical_input_schema, execution_props) + create_physical_expr( + filter, + logical_input_schema, + execution_props, + planning_ctx, + ) }) .transpose()?; - let order_bys = - create_physical_sort_exprs(order_by, logical_input_schema, execution_props)?; + let order_bys = create_physical_sort_exprs( + order_by, + logical_input_schema, + execution_props, + planning_ctx, + )?; let ignore_nulls = null_treatment.unwrap_or(NullTreatment::RespectNulls) == NullTreatment::IgnoreNulls; @@ -858,7 +880,7 @@ impl AggregateFunctionExpr { // `retract_batch` method will not be called. In this case // having retract_batch is not a requirement. // - // This approach is a a bit different than window function + // This approach is a bit different than window function // approach. In window function (when they use a window frame) // they get all the desired range during evaluation. if !accumulator.supports_retract_batch() { @@ -1162,6 +1184,7 @@ mod tests { &logical_schema, &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), ) .build()?; @@ -1185,6 +1208,7 @@ mod tests { &logical_schema, &schema, &ExecutionProps::new(), + &PhysicalPlanningContext::default(), ) .with_human_display(expr.human_display().to_string()) .build()?; diff --git a/datafusion/physical-expr/src/analysis.rs b/datafusion/physical-expr/src/analysis.rs index 1dca36b75f9f5..a00fc19ae9c02 100644 --- a/datafusion/physical-expr/src/analysis.rs +++ b/datafusion/physical-expr/src/analysis.rs @@ -350,6 +350,7 @@ mod tests { use datafusion_common::{DFSchema, ScalarValue, assert_contains, stats::Precision}; use datafusion_expr::{ Expr, col, execution_props::ExecutionProps, interval_arithmetic::Interval, lit, + physical_planning_context::PhysicalPlanningContext, }; use crate::{AnalysisContext, create_physical_expr, expressions::Column}; @@ -412,8 +413,13 @@ mod tests { for (expr, lower, upper) in test_cases { let boundaries = ExprBoundaries::try_new_unbounded(&schema).unwrap(); let df_schema = DFSchema::try_from(Arc::clone(&schema)).unwrap(); - let physical_expr = - create_physical_expr(&expr, &df_schema, &ExecutionProps::new()).unwrap(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); let analysis_result = analyze( &physical_expr, AnalysisContext::new(boundaries), @@ -453,8 +459,13 @@ mod tests { for expr in test_cases { let boundaries = ExprBoundaries::try_new_unbounded(&schema).unwrap(); let df_schema = DFSchema::try_from(Arc::clone(&schema)).unwrap(); - let physical_expr = - create_physical_expr(&expr, &df_schema, &ExecutionProps::new()).unwrap(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); let analysis_result = analyze( &physical_expr, AnalysisContext::new(boundaries), @@ -475,8 +486,13 @@ mod tests { let expected_error = "OR operator cannot yet propagate true intervals"; let boundaries = ExprBoundaries::try_new_unbounded(&schema).unwrap(); let df_schema = DFSchema::try_from(Arc::clone(&schema)).unwrap(); - let physical_expr = - create_physical_expr(&expr, &df_schema, &ExecutionProps::new()).unwrap(); + let physical_expr = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); let analysis_error = analyze( &physical_expr, AnalysisContext::new(boundaries), diff --git a/datafusion/physical-expr/src/async_scalar_function.rs b/datafusion/physical-expr/src/async_scalar_function.rs index 5612e63b530e7..e8ee9a69481df 100644 --- a/datafusion/physical-expr/src/async_scalar_function.rs +++ b/datafusion/physical-expr/src/async_scalar_function.rs @@ -88,12 +88,9 @@ impl AsyncFuncExpr { } /// Return the output field generated by evaluating this function - pub fn field(&self, input_schema: &Schema) -> Result { - Ok(Field::new( - &self.name, - self.func.data_type(input_schema)?, - self.func.nullable(input_schema)?, - )) + #[deprecated(since = "55.0.0", note = "Use return_field instead")] + pub fn field(&self, _input_schema: &Schema) -> Result { + Ok(self.return_field.as_ref().clone().with_name(&self.name)) } /// Return the ideal batch size for this function @@ -211,6 +208,12 @@ impl PhysicalExpr for AsyncFuncExpr { self.func.data_type(input_schema) } + fn return_field(&self, _input_schema: &Schema) -> Result { + Ok(Arc::new( + self.return_field.as_ref().clone().with_name(&self.name), + )) + } + fn nullable(&self, input_schema: &Schema) -> Result { self.func.nullable(input_schema) } diff --git a/datafusion/physical-expr/src/equivalence/class.rs b/datafusion/physical-expr/src/equivalence/class.rs index d00a4a32278f0..1f9a6a583cc44 100644 --- a/datafusion/physical-expr/src/equivalence/class.rs +++ b/datafusion/physical-expr/src/equivalence/class.rs @@ -551,7 +551,19 @@ impl EquivalenceGroup { sort_exprs .into_iter() .map(|sort_expr| self.normalize_sort_expr(sort_expr)) - .filter(|sort_expr| self.is_expr_constant(&sort_expr.expr).is_none()) + .filter(|sort_expr| !self.is_uniform_constant(&sort_expr.expr)) + } + + /// Returns `true` when `expr` is a *globally* constant column, safe to drop + /// from a required ordering. Only [`AcrossPartitions::Uniform`] qualifies; a + /// [`AcrossPartitions::Heterogeneous`] value is constant within a partition + /// but varies across partitions, so it still discriminates the order once + /// partitions are merged and must be kept. + fn is_uniform_constant(&self, expr: &Arc) -> bool { + matches!( + self.is_expr_constant(expr), + Some(AcrossPartitions::Uniform(_)) + ) } /// Normalizes the given sort requirement according to this group. The @@ -582,7 +594,7 @@ impl EquivalenceGroup { sort_reqs .into_iter() .map(|req| self.normalize_sort_requirement(req)) - .filter(|req| self.is_expr_constant(&req.expr).is_none()) + .filter(|req| !self.is_uniform_constant(&req.expr)) } /// Perform an indirect projection of `expr` by consulting the equivalence diff --git a/datafusion/physical-expr/src/equivalence/ordering.rs b/datafusion/physical-expr/src/equivalence/ordering.rs index 2ce8a8d246fe7..499187a603979 100644 --- a/datafusion/physical-expr/src/equivalence/ordering.rs +++ b/datafusion/physical-expr/src/equivalence/ordering.rs @@ -329,7 +329,7 @@ mod tests { EquivalenceClass, EquivalenceGroup, EquivalenceProperties, OrderingEquivalenceClass, convert_to_orderings, convert_to_sort_exprs, }; - use crate::expressions::{BinaryExpr, Column, col}; + use crate::expressions::{BinaryExpr, CastExpr, Column, col}; use crate::utils::tests::TestScalarUDF; use crate::{ AcrossPartitions, ConstExpr, PhysicalExpr, PhysicalExprRef, PhysicalSortExpr, @@ -376,6 +376,45 @@ mod tests { Ok(()) } + #[test] + fn test_ordering_satisfy_strictly_order_preserving() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int64, true), + ])); + let col_a = col("a", &schema)?; + let col_b = col("b", &schema)?; + let asc = SortOptions::default(); + let sort_a = PhysicalSortExpr::new(Arc::clone(&col_a), asc); + let sort_b = PhysicalSortExpr::new(Arc::clone(&col_b), asc); + let eq_properties = EquivalenceProperties::new_with_orderings( + Arc::clone(&schema), + [vec![sort_a.clone(), sort_b.clone()]], + ); + + assert!(eq_properties.ordering_satisfy(vec![sort_a.clone(), sort_b.clone()])?); + assert!(eq_properties.ordering_satisfy(vec![sort_a.clone()])?); + + // A widening cast is strictly order-preserving: `a` is constant + // within each group of equal `CAST(a AS BIGINT)` values, so `b` + // remains sorted within those groups. + let widening = Arc::new(CastExpr::new(Arc::clone(&col_a), DataType::Int64, None)) + as PhysicalExprRef; + let sort_widening = PhysicalSortExpr::new(widening, asc); + assert!(eq_properties.ordering_satisfy(vec![sort_widening, sort_b.clone()])?); + + // A narrowing cast is only monotonic: it satisfies as a leading key, + // but it may collapse distinct `a` values, so `b` is not guaranteed + // to be sorted within its tie groups. + let narrowing = Arc::new(CastExpr::new(Arc::clone(&col_a), DataType::Int16, None)) + as PhysicalExprRef; + let sort_narrowing = PhysicalSortExpr::new(narrowing, asc); + assert!(eq_properties.ordering_satisfy(vec![sort_narrowing.clone()])?); + assert!(!eq_properties.ordering_satisfy(vec![sort_narrowing, sort_b.clone()])?); + + Ok(()) + } + #[test] fn test_ordering_satisfy_with_equivalence2() -> Result<()> { let test_schema = create_test_schema()?; @@ -486,8 +525,8 @@ mod tests { vec![col_e], // requirement [a ASC, c ASC, a+b ASC], vec![(col_a, options), (col_c, options), (&a_plus_b, options)], - // expected: requirement is satisfied. - true, + // expected: requirement is not satisfied because addition can wrap. + false, ), // ------------ TEST CASE 4 ------------ ( @@ -633,8 +672,8 @@ mod tests { vec![col_e], // requirement [c ASC, d ASC, a + b ASC], vec![(col_c, options), (col_d, options), (&a_plus_b, options)], - // expected: requirement is satisfied. - true, + // expected: requirement is not satisfied because addition can wrap. + false, ), ]; diff --git a/datafusion/physical-expr/src/equivalence/properties/dependency.rs b/datafusion/physical-expr/src/equivalence/properties/dependency.rs index 2ebc71559fcf4..bd8bef84de2d8 100644 --- a/datafusion/physical-expr/src/equivalence/properties/dependency.rs +++ b/datafusion/physical-expr/src/equivalence/properties/dependency.rs @@ -632,10 +632,10 @@ mod tests { ]); let test_cases = vec![ - // d + b + // d + b can wrap ( Arc::new(BinaryExpr::new(col_d, Operator::Plus, Arc::clone(&col_b))) as _, - SortProperties::Ordered(option_asc), + SortProperties::Unordered, ), // b (col_b, SortProperties::Ordered(option_asc)), @@ -717,8 +717,8 @@ mod tests { (vec![col_b], vec![]), // TEST CASE 5 (vec![col_d], vec![(col_d, option_asc)]), - // TEST CASE 5 - (vec![&a_plus_d], vec![(&a_plus_d, option_asc)]), + // TEST CASE 5: a + d is not ordered because addition can wrap. + (vec![&a_plus_d], vec![]), // TEST CASE 6 ( vec![col_b, col_d], @@ -1011,7 +1011,7 @@ mod tests { } #[test] - fn test_ordering_equivalence_with_lex_monotonic_concat() -> Result<()> { + fn test_ordering_equivalence_with_non_lex_monotonic_concat() -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("a", DataType::Utf8, false), Field::new("b", DataType::Utf8, false), @@ -1033,28 +1033,23 @@ mod tests { // Assume existing ordering is [c ASC, a ASC, b ASC] let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - eq_properties.add_ordering([ + let initial_ordering: LexOrdering = [ PhysicalSortExpr::new_default(Arc::clone(&col_c)).asc(), PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), - ]); + ] + .into(); + + eq_properties.add_ordering(initial_ordering.clone()); // Add equality condition c = concat(a, b) eq_properties.add_equal_conditions(Arc::clone(&col_c), a_concat_b)?; let orderings = eq_properties.oeq_class(); - let expected_ordering1 = [PhysicalSortExpr::new_default(col_c).asc()].into(); - let expected_ordering2 = [ - PhysicalSortExpr::new_default(col_a).asc(), - PhysicalSortExpr::new_default(col_b).asc(), - ] - .into(); - - // The ordering should be [c ASC] and [a ASC, b ASC] - assert_eq!(orderings.len(), 2); - assert!(orderings.contains(&expected_ordering1)); - assert!(orderings.contains(&expected_ordering2)); + // The ordering should remain unchanged since concat is not lex-monotonic + assert_eq!(orderings.len(), 1); + assert!(orderings.contains(&initial_ordering)); Ok(()) } @@ -1101,55 +1096,6 @@ mod tests { Ok(()) } - #[test] - fn test_ordering_equivalence_with_concat_equality() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - Field::new("c", DataType::Utf8, false), - ])); - - let col_a = col("a", &schema)?; - let col_b = col("b", &schema)?; - let col_c = col("c", &schema)?; - - let a_concat_b = Arc::new(ScalarFunctionExpr::new( - "concat", - concat(), - vec![Arc::clone(&col_a), Arc::clone(&col_b)], - Field::new("f", DataType::Utf8, true).into(), - Arc::new(ConfigOptions::default()), - )) as _; - - // Assume existing ordering is [concat(a, b) ASC, a ASC, b ASC] - let mut eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - eq_properties.add_ordering([ - PhysicalSortExpr::new_default(Arc::clone(&a_concat_b)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_a)).asc(), - PhysicalSortExpr::new_default(Arc::clone(&col_b)).asc(), - ]); - - // Add equality condition c = concat(a, b) - eq_properties.add_equal_conditions(col_c, Arc::clone(&a_concat_b))?; - - let orderings = eq_properties.oeq_class(); - - let expected_ordering1 = [PhysicalSortExpr::new_default(a_concat_b).asc()].into(); - let expected_ordering2 = [ - PhysicalSortExpr::new_default(col_a).asc(), - PhysicalSortExpr::new_default(col_b).asc(), - ] - .into(); - - // The ordering should be [c ASC] and [a ASC, b ASC] - assert_eq!(orderings.len(), 2); - assert!(orderings.contains(&expected_ordering1)); - assert!(orderings.contains(&expected_ordering2)); - - Ok(()) - } - #[test] fn test_requirements_compatible() -> Result<()> { let schema = Arc::new(Schema::new(vec![ diff --git a/datafusion/physical-expr/src/equivalence/properties/joins.rs b/datafusion/physical-expr/src/equivalence/properties/joins.rs index 536badba435d3..d41293615f6c0 100644 --- a/datafusion/physical-expr/src/equivalence/properties/joins.rs +++ b/datafusion/physical-expr/src/equivalence/properties/joins.rs @@ -210,7 +210,7 @@ mod tests { &[], )?; let err_msg = - format!("expected: {:?}, actual:{:?}", expected, &join_eq.oeq_class); + format!("expected: {:?}, actual:{:?}", expected, join_eq.oeq_class); assert_eq!(join_eq.oeq_class.len(), expected.len(), "{err_msg}"); for ordering in join_eq.oeq_class { assert!( diff --git a/datafusion/physical-expr/src/equivalence/properties/mod.rs b/datafusion/physical-expr/src/equivalence/properties/mod.rs index bb74cd1d9c7b3..22b3382f50638 100644 --- a/datafusion/physical-expr/src/equivalence/properties/mod.rs +++ b/datafusion/physical-expr/src/equivalence/properties/mod.rs @@ -33,13 +33,13 @@ use self::dependency::{ use crate::equivalence::{ AcrossPartitions, EquivalenceGroup, OrderingEquivalenceClass, ProjectionMapping, }; -use crate::expressions::{CastExpr, Column, Literal, with_new_schema}; +use crate::expressions::{Column, Literal, with_new_schema}; use crate::{ ConstExpr, LexOrdering, LexRequirement, PhysicalExpr, PhysicalSortExpr, PhysicalSortRequirement, }; -use arrow::datatypes::{DataType, SchemaRef}; +use arrow::datatypes::SchemaRef; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_common::{Constraint, Constraints, HashMap, Result, plan_err}; use datafusion_expr::interval_arithmetic::Interval; @@ -195,24 +195,30 @@ impl OrderingEquivalenceCache { } impl EquivalenceProperties { - /// Helper used by the ordering equivalence rule when considering whether a - /// cast-bearing expression can replace an existing sort key without - /// invalidating the ordering. + /// Helper used by the ordering equivalence rule when considering whether + /// an expression can replace an existing sort key without invalidating + /// the ordering. /// - /// The substitution is only allowed when the cast wraps the very same child - /// expression that the original sort used and the casted type is a - /// widening/order-preserving conversion. Without those restrictions, a - /// narrowing cast could collapse distinct values and violate the existing + /// The substitution is only allowed when, treating the sort key as the + /// only ordered input, the expression reports the same ordering *and* + /// that it is a one-to-one, order-preserving function of it (see + /// [`ExprProperties::strictly_order_preserving`]). For example, a + /// widening `CAST` of the sort key qualifies, while a narrowing one does + /// not, as it could collapse distinct values and violate the existing /// sort order. - fn substitute_cast_ordering( + fn substitute_order_preserving_ordering( r_expr: Arc, sort_expr: &PhysicalSortExpr, - expr_type: &DataType, + schema: &SchemaRef, ) -> Option { - let cast_expr = r_expr.downcast_ref::()?; - - (cast_expr.expr().eq(&sort_expr.expr) - && CastExpr::check_bigger_cast(cast_expr.cast_type(), expr_type)) + if r_expr.eq(&sort_expr.expr) { + // No point in substituting an expression with itself. + return None; + } + let dependencies = Dependencies::new(std::iter::once(sort_expr.clone())); + let properties = get_expr_properties(&r_expr, &dependencies, schema).ok()?; + (properties.strictly_order_preserving + && properties.sort_properties == SortProperties::Ordered(sort_expr.options)) .then(|| PhysicalSortExpr::new(r_expr, sort_expr.options)) } @@ -482,6 +488,7 @@ impl EquivalenceProperties { sort_properties: SortProperties::Ordered(next.options), range: Interval::make_unbounded(&data_type)?, preserves_lex_ordering: true, + strictly_order_preserving: true, }); } // Check if the expression is monotonic in all arguments: @@ -626,24 +633,55 @@ impl EquivalenceProperties { if !satisfy { return Ok(false); } - // Treat satisfied keys as constants in subsequent iterations. We - // can do this because the "next" key only matters in a lexicographical - // ordering when the keys to its left have the same values. - // - // Note that these expressions are not properly "constants". This is just - // an implementation strategy confined to this function. - // - // For example, assume that the requirement is `[a ASC, (b + c) ASC]`, - // and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. - // From the analysis above, we know that `[a ASC]` is satisfied. Then, - // we add column `a` as constant to the algorithm state. This enables us - // to deduce that `(b + c) ASC` is satisfied, given `a` is constant. - let const_expr = ConstExpr::from(element.expr); - eq_properties.add_constants(std::iter::once(const_expr))?; + // Treat satisfied keys (and the sub-expressions they pin down) as + // constants in subsequent iterations. See + // [`Self::add_satisfied_key_constants`] for the rationale. + eq_properties.add_satisfied_key_constants(element.expr)?; } Ok(true) } + /// Registers a satisfied sort key as a constant for subsequent iterations + /// of the ordering satisfaction checks. We can do this because the "next" + /// key only matters in a lexicographical ordering when the keys to its + /// left have the same values (i.e. within a single tie group). Note that + /// these expressions are not properly "constants"; this is just an + /// implementation strategy confined to the satisfaction checks. + /// + /// For example, assume that the requirement is `[a ASC, (b + c) ASC]`, + /// and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. + /// Once we deduce that `[a ASC]` is satisfied, we add column `a` as a + /// constant to the algorithm state. This enables us to deduce that + /// `(b + c) ASC` is satisfied, given `a` is constant. + /// + /// In addition to the key itself, this also registers any sub-expressions + /// whose values the key pins down: if an expression is strictly + /// order-preserving, equal outputs imply equal values of its ordered + /// children, so within a tie group of the key those children are constant + /// as well. For example, if data is sorted by `[a, b]`, the requirement + /// `[CAST(a AS BIGINT) ASC, b ASC]` is satisfied: `a` is constant within + /// each group of equal `CAST(a AS BIGINT)` values, and hence `b` is + /// sorted within each such group. + fn add_satisfied_key_constants(&mut self, expr: Arc) -> Result<()> { + let mut stack = vec![expr]; + while let Some(expr) = stack.pop() { + let properties = self.get_expr_properties(Arc::clone(&expr)); + if properties.strictly_order_preserving { + for child in expr.children() { + let child_properties = self.get_expr_properties(Arc::clone(child)); + if matches!( + child_properties.sort_properties, + SortProperties::Ordered(_) + ) { + stack.push(Arc::clone(child)); + } + } + } + self.add_constants(std::iter::once(ConstExpr::from(expr)))?; + } + Ok(()) + } + /// Returns the number of consecutive sort expressions (starting from the /// left) that are satisfied by the existing ordering. fn common_sort_prefix_length(&self, normal_ordering: &LexOrdering) -> Result { @@ -676,20 +714,10 @@ impl EquivalenceProperties { // many we've satisfied so far: return Ok(idx); } - // Treat satisfied keys as constants in subsequent iterations. We - // can do this because the "next" key only matters in a lexicographical - // ordering when the keys to its left have the same values. - // - // Note that these expressions are not properly "constants". This is just - // an implementation strategy confined to this function. - // - // For example, assume that the requirement is `[a ASC, (b + c) ASC]`, - // and existing equivalent orderings are `[a ASC, b ASC]` and `[c ASC]`. - // From the analysis above, we know that `[a ASC]` is satisfied. Then, - // we add column `a` as constant to the algorithm state. This enables us - // to deduce that `(b + c) ASC` is satisfied, given `a` is constant. - let const_expr = ConstExpr::from(Arc::clone(&element.expr)); - eq_properties.add_constants(std::iter::once(const_expr))? + // Treat satisfied keys (and the sub-expressions they pin down) as + // constants in subsequent iterations. See + // [`Self::add_satisfied_key_constants`] for the rationale. + eq_properties.add_satisfied_key_constants(Arc::clone(&element.expr))?; } // All sort expressions are satisfied, return full length: Ok(full_length) @@ -840,7 +868,9 @@ impl EquivalenceProperties { /// /// TODO: Handle all scenarios that allow substitution; e.g. when `x` is /// sorted, `atan(x + 1000)` should also be substituted. For now, we - /// only consider single-column `CAST` expressions. + /// consider widening `CAST` expressions and single-child expressions + /// that declare themselves one-to-one order-preserving via + /// [`ExprProperties::strictly_order_preserving`]. fn substitute_oeq_class( schema: &SchemaRef, mapping: &ProjectionMapping, @@ -852,21 +882,17 @@ impl EquivalenceProperties { order .into_iter() .map(|sort_expr| { - // The sort expression comes from this schema, so the - // following call to `unwrap` is safe. - let expr_type = sort_expr.expr.data_type(schema).unwrap(); let original_sort_expr = sort_expr.clone(); - // TODO: Add one-to-one analysis for ScalarFunctions. mapping .iter() .map(|(source, _target)| source) .filter(|source| expr_refers(source, &original_sort_expr.expr)) .cloned() .filter_map(|r_expr| { - Self::substitute_cast_ordering( + Self::substitute_order_preserving_ordering( r_expr, &original_sort_expr, - &expr_type, + schema, ) }) .chain(std::iter::once(sort_expr)) @@ -1314,7 +1340,18 @@ impl EquivalenceProperties { if let (Some(data_type), Some(AcrossPartitions::Uniform(Some(value)))) = (data_type, &mut eq_class.constant) { - *value = value.cast_to(&data_type)?; + match value.cast_to(&data_type) { + Ok(cast_value) => *value = cast_value, + Err(_) => { + // This is optimizer metadata. If a stale constant + // value cannot be represented after schema rewrite, + // drop the constant instead of failing planning. + eq_class.constant = None; + } + } + } + if eq_class.is_trivial() { + continue; } eq_classes.push(eq_class); } @@ -1396,7 +1433,10 @@ fn update_properties( } else if node.expr.is::() { // We have a Column, which is the other possible leaf node type: node.data.range = - Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)? + Interval::make_unbounded(&node.expr.data_type(eq_properties.schema())?)?; + // A column is the identity mapping of itself, which is trivially + // strict: + node.data.strictly_order_preserving = true; } // Now, check what we know about orderings: let normal_expr = eq_properties @@ -1458,23 +1498,36 @@ fn get_expr_properties( schema: &SchemaRef, ) -> Result { if let Some(column_order) = dependencies.iter().find(|&order| expr.eq(&order.expr)) { - // If exact match is found, return its ordering. + // If exact match is found, return its ordering. This is a base case + // of the recursion: the expression is treated as an atomic ordered + // input from here on, so `strictly_order_preserving` states only that + // it is a one-to-one mapping *of itself* (the identity), which holds + // for any expression. It makes no claim about the expression being + // one-to-one in its own inputs (e.g. `floor(x)` as a sort key), and + // it does not need to: parent expressions are substituted for this + // sort key, so their strictness only has to be relative to it. Ok(ExprProperties { sort_properties: SortProperties::Ordered(column_order.options), range: Interval::make_unbounded(&expr.data_type(schema)?)?, preserves_lex_ordering: false, + strictly_order_preserving: true, }) } else if expr.downcast_ref::().is_some() { Ok(ExprProperties { sort_properties: SortProperties::Unordered, range: Interval::make_unbounded(&expr.data_type(schema)?)?, preserves_lex_ordering: false, + // A base case of the recursion: a column is the identity mapping + // of itself, which is trivially one-to-one. + strictly_order_preserving: true, }) } else if let Some(literal) = expr.downcast_ref::() { Ok(ExprProperties { sort_properties: SortProperties::Singleton, range: literal.value().into(), preserves_lex_ordering: true, + // Vacuously true: a literal has no ordered inputs. + strictly_order_preserving: true, }) } else { // Find orderings of its children diff --git a/datafusion/physical-expr/src/equivalence/properties/union.rs b/datafusion/physical-expr/src/equivalence/properties/union.rs index d77129472a8ba..ea4094e75159a 100644 --- a/datafusion/physical-expr/src/equivalence/properties/union.rs +++ b/datafusion/physical-expr/src/equivalence/properties/union.rs @@ -311,7 +311,7 @@ mod tests { use crate::equivalence::tests::{create_test_schema, parse_sort_expr}; use crate::expressions::col; - use arrow::datatypes::{DataType, Field, Schema}; + use arrow::datatypes::{DataType, Field, Schema, TimeUnit}; use datafusion_common::ScalarValue; use itertools::Itertools; @@ -899,6 +899,36 @@ mod tests { Ok(()) } + #[test] + fn test_union_drops_unrepresentable_constant_value_after_schema_rewrite() -> Result<()> + { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "ticker", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )])); + let output_schema = Arc::new(Schema::new(vec![Field::new( + "timestamp", + DataType::Timestamp(TimeUnit::Nanosecond, None), + true, + )])); + + let ticker = col("ticker", &input_schema)?; + let stale_value = ScalarValue::Utf8(Some("ESU6".to_owned())); + let const_expr = ConstExpr::new( + Arc::clone(&ticker), + AcrossPartitions::Uniform(Some(stale_value)), + ); + + let mut input = EquivalenceProperties::new(input_schema); + input.add_constants(vec![const_expr])?; + + let union_props = calculate_union(vec![input], output_schema)?; + assert!(union_props.constants().is_empty()); + + Ok(()) + } + /// Return a new schema with the same types, but new field names /// /// The new field names are the old field names with `text` appended. diff --git a/datafusion/physical-expr/src/expressions/binary.rs b/datafusion/physical-expr/src/expressions/binary.rs index b92668fe9bd0d..1bd49696bbdca 100644 --- a/datafusion/physical-expr/src/expressions/binary.rs +++ b/datafusion/physical-expr/src/expressions/binary.rs @@ -19,14 +19,13 @@ mod kernels; use crate::PhysicalExpr; use crate::intervals::cp_solver::{propagate_arithmetic, propagate_comparison}; +use std::cmp::Ordering; use std::hash::Hash; use std::sync::Arc; use arrow::array::*; use arrow::compute::kernels::boolean::{and_kleene, or_kleene}; -use arrow::compute::kernels::concat_elements::{ - concat_element_binary, concat_elements_utf8, -}; +use arrow::compute::kernels::concat_elements::concat_elements_dyn; use arrow::compute::{SlicesIterator, cast, filter_record_batch}; use arrow::datatypes::*; use arrow::error::ArrowError; @@ -35,7 +34,7 @@ use datafusion_common::{Result, ScalarValue, internal_err, not_impl_err}; use datafusion_expr::binary::BinaryTypeCoercer; use datafusion_expr::interval_arithmetic::{Interval, apply_operator}; -use datafusion_expr::sort_properties::ExprProperties; +use datafusion_expr::sort_properties::{ExprProperties, SortProperties}; #[expect(deprecated)] use datafusion_expr::statistics::Distribution::{Bernoulli, Gaussian}; #[expect(deprecated)] @@ -50,8 +49,7 @@ use kernels::{ bitwise_and_dyn, bitwise_and_dyn_scalar, bitwise_or_dyn, bitwise_or_dyn_scalar, bitwise_shift_left_dyn, bitwise_shift_left_dyn_scalar, bitwise_shift_right_dyn, bitwise_shift_right_dyn_scalar, bitwise_xor_dyn, bitwise_xor_dyn_scalar, - concat_elements_binary_view_array, concat_elements_utf8view, regex_match_dyn, - regex_match_dyn_scalar, + regex_match_dyn, regex_match_dyn_scalar, }; /// Binary expression @@ -121,6 +119,68 @@ impl BinaryExpr { pub fn op(&self) -> &Operator { &self.op } + + /// Wrapping on overflow breaks monotonicity (e.g. the sum of two + /// ascending `UInt8` columns can wrap back to small values), so the + /// derived ordering is kept only when overflow is impossible. `time ± + /// interval` wraps around the 24-hour clock even in checked mode, so it + /// never preserves ordering. + fn arithmetic_sort_properties( + &self, + sort_properties: SortProperties, + l_range: &Interval, + r_range: &Interval, + range: &Interval, + ) -> SortProperties { + if sort_properties == SortProperties::Singleton { + return sort_properties; + } + let wraps_in_domain = match self.op { + Operator::Plus => { + is_time_plus_interval(&l_range.data_type(), &r_range.data_type()) + } + Operator::Minus => { + is_time_minus_interval(&l_range.data_type(), &r_range.data_type()) + } + _ => false, + }; + let cannot_overflow = !range.is_unbounded() + && !unsigned_subtraction_may_underflow(self.op, l_range, r_range, range); + if !wraps_in_domain && (self.fail_on_overflow || cannot_overflow) { + sort_properties + } else { + SortProperties::Unordered + } + } +} + +/// Returns `true` unless `l_range - r_range` provably stays within an unsigned +/// domain. +/// +/// [`Interval`] standardizes an underflowed (i.e. `null`) lower bound of an +/// unsigned type back to zero, so an apparently bounded result range is not +/// enough to rule out wrapping here -- e.g. `[0, 10] - [0, 10]` over `UInt32` +/// yields `[0, 10]` even though `0 - 10` wraps to `u32::MAX`. Compare the +/// endpoints that produce the smallest difference instead. +fn unsigned_subtraction_may_underflow( + op: Operator, + l_range: &Interval, + r_range: &Interval, + range: &Interval, +) -> bool { + if op != Operator::Minus || !range.data_type().is_unsigned_integer() { + return false; + } + let (smallest_lhs, largest_rhs) = (l_range.lower(), r_range.upper()); + if smallest_lhs.is_null() || largest_rhs.is_null() { + return true; + } + // Operands of differing types compare as incomparable, in which case we + // conservatively assume an underflow is possible. + !matches!( + smallest_lhs.partial_cmp(largest_rhs), + Some(Ordering::Greater | Ordering::Equal) + ) } impl std::fmt::Display for BinaryExpr { @@ -177,82 +237,283 @@ fn is_date_minus_date(lhs: &DataType, rhs: &DataType) -> bool { ) } -/// Computes the difference between two dates and returns the result as Int64 (days) -/// This aligns with PostgreSQL, DuckDB, and MySQL behavior where date - date returns an integer +/// Milliseconds per day, used for Date64 subtraction. +const MILLIS_PER_DAY: i64 = 86_400_000; + +/// Evaluates `Date32 - Date32` or `Date64 - Date64`, returning the difference in +/// whole days as `Int64`. /// -/// Implementation: Uses Arrow's sub_wrapping to get Duration, then converts to Int64 days +/// This matches the behavior of PostgreSQL, DuckDB, and MySQL, where +/// `date - date` yields an integer day count rather than an interval. fn apply_date_subtraction( lhs: &ColumnarValue, rhs: &ColumnarValue, ) -> Result { - use arrow::compute::kernels::numeric::sub_wrapping; + match (lhs.data_type(), rhs.data_type()) { + (DataType::Date32, DataType::Date32) => { + subtract_date_to_days::(lhs, rhs, |l, r| l - r) + } + (DataType::Date64, DataType::Date64) => { + subtract_date_to_days::(lhs, rhs, |l, r| { + l.wrapping_sub(r) / MILLIS_PER_DAY + }) + } + (_, _) => unreachable!("apply_date_subtraction called with non-date types"), + } +} - // Use Arrow's sub_wrapping to compute the Duration result - let duration_result = apply(lhs, rhs, sub_wrapping)?; +/// Generic date subtraction: operates directly on the native primitive values +/// of `T` (i32 for Date32, i64 for Date64), applying `day_diff_fn` to produce +/// an Int64 day count. +fn subtract_date_to_days( + lhs: &ColumnarValue, + rhs: &ColumnarValue, + day_diff_fn: impl Fn(i64, i64) -> i64, +) -> Result +where + T::Native: Copy + Into, +{ + /// Extract the date value as `i64`. Returns `None` for null scalars. + fn date_scalar_to_i64( + scalar: &ScalarValue, + ) -> Result> { + match scalar { + ScalarValue::Date32(value) if P::DATA_TYPE == DataType::Date32 => { + Ok(value.map(i64::from)) + } + ScalarValue::Date64(value) if P::DATA_TYPE == DataType::Date64 => Ok(*value), + other => { + internal_err!( + "{} date scalar expected, got: {}", + P::DATA_TYPE, + other.data_type() + ) + } + } + } - // Convert Duration to Int64 (days) - match duration_result { - ColumnarValue::Array(array) => { - let int64_array = duration_to_days(&array)?; - Ok(ColumnarValue::Array(int64_array)) + match (lhs, rhs) { + (ColumnarValue::Array(left), ColumnarValue::Array(right)) => { + let left = left.as_primitive::(); + let right = right.as_primitive::(); + let result: Int64Array = + arrow::compute::binary::<_, _, _, Int64Type>(left, right, |l, r| { + day_diff_fn(l.into(), r.into()) + })?; + Ok(ColumnarValue::Array(Arc::new(result))) } - ColumnarValue::Scalar(scalar) => { - // Convert scalar Duration to Int64 days - let array = scalar.to_array_of_size(1)?; - let int64_array = duration_to_days(&array)?; - let int64_scalar = ScalarValue::try_from_array(int64_array.as_ref(), 0)?; - Ok(ColumnarValue::Scalar(int64_scalar)) + (ColumnarValue::Array(left), ColumnarValue::Scalar(right)) => { + let left = left.as_primitive::(); + match date_scalar_to_i64::(right)? { + Some(right_val) => { + let result: Int64Array = + left.unary(|l| day_diff_fn(l.into(), right_val)); + Ok(ColumnarValue::Array(Arc::new(result))) + } + None => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))), + } + } + (ColumnarValue::Scalar(left), ColumnarValue::Array(right)) => { + let right = right.as_primitive::(); + match date_scalar_to_i64::(left)? { + Some(left_val) => { + let result: Int64Array = + right.unary(|r| day_diff_fn(left_val, r.into())); + Ok(ColumnarValue::Array(Arc::new(result))) + } + None => Ok(ColumnarValue::Scalar(ScalarValue::Int64(None))), + } + } + (ColumnarValue::Scalar(left), ColumnarValue::Scalar(right)) => { + let left_val = date_scalar_to_i64::(left)?; + let right_val = date_scalar_to_i64::(right)?; + Ok(ColumnarValue::Scalar(ScalarValue::Int64( + left_val.zip(right_val).map(|(l, r)| day_diff_fn(l, r)), + ))) } } } -/// Converts a Duration array to Int64 days -/// Handles different Duration time units (Second, Millisecond, Microsecond, Nanosecond) -fn duration_to_days(array: &ArrayRef) -> Result { - use datafusion_common::cast::{ - as_duration_microsecond_array, as_duration_millisecond_array, - as_duration_nanosecond_array, as_duration_second_array, +/// Returns true for `time + interval` or `interval + time`. +fn is_time_plus_interval(lhs: &DataType, rhs: &DataType) -> bool { + matches!( + (lhs, rhs), + ( + DataType::Time32(_) | DataType::Time64(_), + DataType::Interval(_) + ) | ( + DataType::Interval(_), + DataType::Time32(_) | DataType::Time64(_) + ) + ) +} + +/// Returns true for `time - interval`. +fn is_time_minus_interval(lhs: &DataType, rhs: &DataType) -> bool { + matches!( + (lhs, rhs), + ( + DataType::Time32(_) | DataType::Time64(_), + DataType::Interval(_) + ) + ) +} + +/// Evaluates `time + interval`, `interval + time`, or `time - interval`, returning a +/// `time` wrapped within the 24-hour clock to match PostgreSQL and DuckDB (e.g. +/// `time '23:30' + interval '2 hours'` is `01:30:00`). arrow's arithmetic kernels do +/// not implement time-of-day arithmetic, so it is handled here. +/// +/// The result keeps the input time's unit; the interval (normalized to `MonthDayNano` +/// by the coercion layer) is applied at nanosecond precision and floored to that unit, +/// mirroring `timestamp(unit) + interval`. Only the sub-day portion of the interval +/// affects a time-of-day -- whole months and days are ignored, matching PostgreSQL. The +/// floor is applied after the sign, so `time(s) + interval '1 nanosecond'` is a no-op +/// while `time(s) - interval '1 nanosecond'` rolls back a second, exactly as the +/// timestamp case does. +fn apply_time_interval( + lhs: &ColumnarValue, + rhs: &ColumnarValue, + subtract: bool, +) -> Result { + // The `time` operand determines the result type; the other is the interval. + let (time, interval) = if matches!(lhs.data_type(), DataType::Interval(_)) { + (rhs, lhs) + } else { + (lhs, rhs) }; - const SECONDS_PER_DAY: i64 = 86_400; - const MILLIS_PER_DAY: i64 = 86_400_000; - const MICROS_PER_DAY: i64 = 86_400_000_000; - const NANOS_PER_DAY: i64 = 86_400_000_000_000; - - match array.data_type() { - DataType::Duration(TimeUnit::Second) => { - let duration_array = as_duration_second_array(array)?; - let result: Int64Array = duration_array - .iter() - .map(|v| v.map(|val| val / SECONDS_PER_DAY)) - .collect(); - Ok(Arc::new(result)) + // Dispatch on the time unit; `ns_per_unit` converts the interval's nanoseconds to + // that unit, and the arithmetic is done (and wrapped) at that resolution. + match time.data_type() { + DataType::Time32(TimeUnit::Second) => wrap_time_interval::( + time, + interval, + subtract, + 1_000_000_000, + ), + DataType::Time32(TimeUnit::Millisecond) => { + wrap_time_interval::( + time, interval, subtract, 1_000_000, + ) } - DataType::Duration(TimeUnit::Millisecond) => { - let duration_array = as_duration_millisecond_array(array)?; - let result: Int64Array = duration_array - .iter() - .map(|v| v.map(|val| val / MILLIS_PER_DAY)) - .collect(); - Ok(Arc::new(result)) + DataType::Time64(TimeUnit::Microsecond) => { + wrap_time_interval::(time, interval, subtract, 1_000) } - DataType::Duration(TimeUnit::Microsecond) => { - let duration_array = as_duration_microsecond_array(array)?; - let result: Int64Array = duration_array - .iter() - .map(|v| v.map(|val| val / MICROS_PER_DAY)) - .collect(); - Ok(Arc::new(result)) + DataType::Time64(TimeUnit::Nanosecond) => { + wrap_time_interval::(time, interval, subtract, 1) } - DataType::Duration(TimeUnit::Nanosecond) => { - let duration_array = as_duration_nanosecond_array(array)?; - let result: Int64Array = duration_array - .iter() - .map(|v| v.map(|val| val / NANOS_PER_DAY)) - .collect(); - Ok(Arc::new(result)) + other => internal_err!("time operand expected, got: {other}"), + } +} + +/// Adds or subtracts an interval to/from a `time` of arrow primitive type `T`, wrapping +/// the result within the 24-hour clock and keeping the type `T`. `ns_per_unit` is the +/// number of nanoseconds in one unit of `T` (e.g. `1_000` for microseconds). +fn wrap_time_interval( + time: &ColumnarValue, + interval: &ColumnarValue, + subtract: bool, + ns_per_unit: i64, +) -> Result +where + T::Native: Copy + Into + TryFrom, +{ + /// Nanoseconds in a 24-hour day. + const DAY_NANOS: i64 = 86_400_000_000_000; + // Units in a 24-hour day, at `T`'s resolution. + let day_units = DAY_NANOS / ns_per_unit; + + // Wraps `time ± interval` into `[0, day_units)`. The interval is reduced modulo a day + // (so the sum stays within `i64`), applied at nanosecond precision, then floored to + // `T`'s unit -- matching `timestamp(unit) ± interval`. Because the floor is applied + // after the sign, `time(s) - interval '1 nanosecond'` rolls back a full second, just + // as the timestamp case does, while `time(s) + interval '1 nanosecond'` is a no-op. + // `div_euclid`/`rem_euclid` floor toward negative infinity, so the wrapped value stays + // in `[0, day_units)`, which always fits `T::Native`. + let wrap = |time_unit: i64, iv: IntervalMonthDayNano| -> T::Native { + let iv_ns = iv.nanoseconds % DAY_NANOS; + let signed_ns = if subtract { -iv_ns } else { iv_ns }; + let delta = signed_ns.div_euclid(ns_per_unit); + let wrapped = (time_unit + delta).rem_euclid(day_units); + T::Native::try_from(wrapped).unwrap_or_default() + }; + + /// Extracts an `Interval(MonthDayNano)` scalar. + fn interval_scalar(scalar: &ScalarValue) -> Result> { + match scalar { + ScalarValue::IntervalMonthDayNano(value) => Ok(*value), + other => internal_err!( + "Interval(MonthDayNano) scalar expected, got: {}", + other.data_type() + ), + } + } + + /// Extracts a time scalar as its unit count since midnight. + fn time_scalar_units(scalar: &ScalarValue) -> Result> { + match scalar { + ScalarValue::Time32Second(value) | ScalarValue::Time32Millisecond(value) => { + Ok(value.map(i64::from)) + } + ScalarValue::Time64Microsecond(value) + | ScalarValue::Time64Nanosecond(value) => Ok(*value), + other => { + internal_err!("time scalar expected, got: {}", other.data_type()) + } + } + } + + /// Builds a time scalar of type `P` from a unit count. + fn time_scalar(value: Option) -> ScalarValue { + match P::DATA_TYPE { + DataType::Time32(TimeUnit::Second) => { + ScalarValue::Time32Second(value.map(|v| v as i32)) + } + DataType::Time32(TimeUnit::Millisecond) => { + ScalarValue::Time32Millisecond(value.map(|v| v as i32)) + } + DataType::Time64(TimeUnit::Microsecond) => { + ScalarValue::Time64Microsecond(value) + } + _ => ScalarValue::Time64Nanosecond(value), + } + } + + match (time, interval) { + (ColumnarValue::Array(time), ColumnarValue::Array(interval)) => { + let time = time.as_primitive::(); + let interval = interval.as_primitive::(); + let result: PrimitiveArray = + arrow::compute::binary(time, interval, |t, iv| wrap(t.into(), iv))?; + Ok(ColumnarValue::Array(Arc::new(result))) + } + (ColumnarValue::Array(time), ColumnarValue::Scalar(interval)) => { + let time = time.as_primitive::(); + match interval_scalar(interval)? { + Some(iv) => { + let result: PrimitiveArray = time.unary(|t| wrap(t.into(), iv)); + Ok(ColumnarValue::Array(Arc::new(result))) + } + None => Ok(ColumnarValue::Scalar(time_scalar::(None))), + } + } + (ColumnarValue::Scalar(time), ColumnarValue::Array(interval)) => { + let interval = interval.as_primitive::(); + match time_scalar_units(time)? { + Some(t) => { + let result: PrimitiveArray = interval.unary(|iv| wrap(t, iv)); + Ok(ColumnarValue::Array(Arc::new(result))) + } + None => Ok(ColumnarValue::Scalar(time_scalar::(None))), + } + } + (ColumnarValue::Scalar(time), ColumnarValue::Scalar(interval)) => { + let result = time_scalar_units(time)? + .zip(interval_scalar(interval)?) + .map(|(t, iv)| wrap(t, iv).into()); + Ok(ColumnarValue::Scalar(time_scalar::(result))) } - other => internal_err!("duration_to_days expected Duration type, got: {}", other), } } @@ -284,41 +545,50 @@ impl PhysicalExpr for BinaryExpr { let rhs = self.right.evaluate(batch)?; return Ok(rhs); } - ShortCircuitStrategy::PreSelection(selection) => { - // The function `evaluate_selection` was not called for filtering and calculation, - // as it takes into account cases where the selection contains null values. - let batch = filter_record_batch(batch, selection)?; - let right_ret = self.right.evaluate(&batch)?; + ShortCircuitStrategy::PreSelection { mask, fill_value } => { + // `mask` selects the rows whose result depends on the RHS; the + // unselected rows are all `fill_value` (see `ShortCircuitStrategy`). + // + // Use `filter_record_batch` directly because `evaluate_selection` + // scatters the RHS back to the original batch length. + let selection_batch = filter_record_batch(batch, &mask)?; + let right_ret = self.right.evaluate(&selection_batch)?; match &right_ret { ColumnarValue::Array(array) => { - // When the array on the right is all true or all false, skip the scatter process let boolean_array = array.as_boolean(); - if boolean_array.null_count() == 0 && !boolean_array.has_false() { - return Ok(lhs); - } else if boolean_array.null_count() == 0 - && !boolean_array.has_true() - { - // If the right-hand array is returned at this point,the lengths will be inconsistent; - // returning a scalar can avoid this issue - return Ok(ColumnarValue::Scalar(ScalarValue::Boolean( - Some(false), - ))); + // If the RHS is uniform on the selected rows, the whole + // expression collapses and no scatter is needed. + if boolean_array.null_count() == 0 { + let rhs_value = if !boolean_array.has_false() { + Some(true) + } else if !boolean_array.has_true() { + Some(false) + } else { + None + }; + if let Some(rhs_value) = rhs_value { + return Ok(uniform_pre_selection_result( + rhs_value, fill_value, lhs, + )); + } } - return pre_selection_scatter(selection, Some(boolean_array)); + return pre_selection_scatter( + &mask, + Some(boolean_array), + fill_value, + ); } ColumnarValue::Scalar(scalar) => { if let ScalarValue::Boolean(v) = scalar { - // When the scalar is true or false, skip the scatter process + // A scalar RHS applies uniformly to all selected rows. if let Some(v) = v { - if *v { - return Ok(lhs); - } else { - return Ok(right_ret); - } + return Ok(uniform_pre_selection_result( + *v, fill_value, lhs, + )); } else { - return pre_selection_scatter(selection, None); + return pre_selection_scatter(&mask, None, fill_value); } } else { return internal_err!( @@ -338,6 +608,18 @@ impl PhysicalExpr for BinaryExpr { let input_schema = schema.as_ref(); match self.op { + // `time ± interval` returns a wrapped `time` (PostgreSQL/DuckDB + // semantics); arrow's arithmetic kernels don't implement it. + Operator::Plus + if is_time_plus_interval(&left_data_type, &right_data_type) => + { + return apply_time_interval(&lhs, &rhs, false); + } + Operator::Minus + if is_time_minus_interval(&left_data_type, &right_data_type) => + { + return apply_time_interval(&lhs, &rhs, true); + } Operator::Plus if self.fail_on_overflow => return apply(&lhs, &rhs, add), Operator::Plus => return apply(&lhs, &rhs, add_wrapping), // Special case: Date - Date returns Int64 (days difference) @@ -541,45 +823,69 @@ impl PhysicalExpr for BinaryExpr { let (l_order, l_range) = (children[0].sort_properties, &children[0].range); let (r_order, r_range) = (children[1].sort_properties, &children[1].range); match self.op() { - Operator::Plus => Ok(ExprProperties { - sort_properties: l_order.add(&r_order), - range: l_range.add(r_range)?, - preserves_lex_ordering: false, - }), - Operator::Minus => Ok(ExprProperties { - sort_properties: l_order.sub(&r_order), - range: l_range.sub(r_range)?, - preserves_lex_ordering: false, - }), + Operator::Plus => { + let range = l_range.add(r_range)?; + Ok(ExprProperties { + sort_properties: self.arithmetic_sort_properties( + l_order.add(&r_order), + l_range, + r_range, + &range, + ), + range, + preserves_lex_ordering: false, + strictly_order_preserving: false, + }) + } + Operator::Minus => { + let range = l_range.sub(r_range)?; + Ok(ExprProperties { + sort_properties: self.arithmetic_sort_properties( + l_order.sub(&r_order), + l_range, + r_range, + &range, + ), + range, + preserves_lex_ordering: false, + strictly_order_preserving: false, + }) + } Operator::Gt => Ok(ExprProperties { sort_properties: l_order.gt_or_gteq(&r_order), range: l_range.gt(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::GtEq => Ok(ExprProperties { sort_properties: l_order.gt_or_gteq(&r_order), range: l_range.gt_eq(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::Lt => Ok(ExprProperties { sort_properties: r_order.gt_or_gteq(&l_order), range: l_range.lt(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::LtEq => Ok(ExprProperties { sort_properties: r_order.gt_or_gteq(&l_order), range: l_range.lt_eq(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::And => Ok(ExprProperties { sort_properties: r_order.and_or(&l_order), range: l_range.and(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), Operator::Or => Ok(ExprProperties { sort_properties: r_order.and_or(&l_order), range: l_range.or(r_range)?, preserves_lex_ordering: false, + strictly_order_preserving: false, }), _ => Ok(ExprProperties::new_unknown()), } @@ -610,6 +916,108 @@ impl PhysicalExpr for BinaryExpr { write!(f, " {} ", self.op)?; write_child(f, self.right.as_ref(), precedence) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + // Linearize a nested binary expression tree of the same operator + // into a flat vector of operands to avoid deep recursion in proto. + let op = self.op; + let mut operand_refs: Vec<&Arc> = vec![&self.right]; + let mut current_expr: &BinaryExpr = self; + loop { + match current_expr.left.downcast_ref::() { + Some(bin) if bin.op == op => { + operand_refs.push(&bin.right); + current_expr = bin; + } + _ => { + operand_refs.push(¤t_expr.left); + break; + } + } + } + // Reverse so operands are ordered from left innermost to right outermost. + operand_refs.reverse(); + + let operands = ctx.encode_children_expressions(operand_refs)?; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::BinaryExpr( + Box::new(protobuf::PhysicalBinaryExprNode { + l: None, + r: None, + op: format!("{op:?}"), + operands, + }), + )), + })) + } +} + +#[cfg(feature = "proto")] +impl BinaryExpr { + /// Reconstruct a [`BinaryExpr`] (or a left-deep tree of them when the proto + /// uses the linearized `operands` form) from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`] — the exact inverse of what + /// [`PhysicalExpr::try_to_proto`] produces — so every expression's + /// `try_from_proto` shares one signature. The operator string is parsed + /// via the canonical [`Operator::from_proto_name`] mapping, so no `op` + /// argument needs to be threaded in by the caller. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + /// [`PhysicalExpr::try_to_proto`]: datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto + /// [`PhysicalExprDecodeCtx::decode`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::decode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + let node = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::BinaryExpr, + "BinaryExpr", + ); + let op = Operator::from_proto_name(&node.op).ok_or_else(|| { + datafusion_common::DataFusionError::Internal(format!( + "Unsupported binary operator '{}'", + node.op + )) + })?; + + if !node.operands.is_empty() { + // New linearized format: reduce the flat operands list back into + // a nested binary expression tree. + let operands = ctx.decode_children_expressions(&node.operands)?; + + if operands.len() < 2 { + return internal_err!( + "A binary expression must always have at least 2 operands" + ); + } + + Ok(operands + .into_iter() + .reduce(|left, right| { + Arc::new(BinaryExpr::new(left, op, right)) as Arc + }) + .expect("Binary expression could not be reduced to a single expression.")) + } else { + // Legacy format with l/r fields. + let left = + ctx.decode_required_expression(node.l.as_deref(), "BinaryExpr", "left")?; + let right = + ctx.decode_required_expression(node.r.as_deref(), "BinaryExpr", "right")?; + Ok(Arc::new(BinaryExpr::new(left, op, right))) + } + } } /// Casts dictionary array to result type for binary numerical operators. Such operators @@ -713,7 +1121,7 @@ impl BinaryExpr { BitwiseXor => bitwise_xor_dyn(left, right), BitwiseShiftRight => bitwise_shift_right_dyn(left, right), BitwiseShiftLeft => bitwise_shift_left_dyn(left, right), - StringConcat => concat_elements(&left, &right), + StringConcat => concat_elements_dyn(&left, &right).map_err(|e| e.into()), AtArrow | ArrowAt | Arrow | LongArrow | HashArrow | HashLongArrow | AtAt | HashMinus | AtQuestion | Question | QuestionAnd | QuestionPipe | IntegerDivide | Colon => { @@ -726,16 +1134,28 @@ impl BinaryExpr { } } -enum ShortCircuitStrategy<'a> { +enum ShortCircuitStrategy { None, ReturnLeft, ReturnRight, - PreSelection(&'a BooleanArray), + /// Evaluate the right-hand side only on the rows selected by `mask`, then + /// scatter the results back, filling the unselected rows with `fill_value`. + /// + /// - For `AND`, `mask` selects the rows where the LHS is `true` and + /// `fill_value` is `false` (rows where the LHS is `false` are `false`). + /// - For `OR`, `mask` selects the rows where the LHS is `false` and + /// `fill_value` is `true` (rows where the LHS is `true` are `true`). + PreSelection { + mask: BooleanArray, + fill_value: bool, + }, } /// Based on the results calculated from the left side of the short-circuit operation, -/// if the proportion of `true` is less than 0.2 and the current operation is an `and`, -/// the `RecordBatch` will be filtered in advance. +/// pre-selection filters the `RecordBatch` before evaluating the right-hand side when +/// the side that cannot short-circuit the operator is rare: +/// - for `AND`, when the proportion of `true` is less than or equal to 0.2 +/// - for `OR`, when the proportion of `false` is less than or equal to 0.2 const PRE_SELECTION_THRESHOLD: f32 = 0.2; /// Checks if a logical operator (`AND`/`OR`) can short-circuit evaluation based on the left-hand side (lhs) result. @@ -744,24 +1164,21 @@ const PRE_SELECTION_THRESHOLD: f32 = 0.2; /// - For `AND`: /// - if LHS is all false => short-circuit → return LHS /// - if LHS is all true => short-circuit → return RHS -/// - if LHS is mixed and true_count/sum_count <= [`PRE_SELECTION_THRESHOLD`] -> pre-selection +/// - if LHS is mixed and true_count / len <= [`PRE_SELECTION_THRESHOLD`] -> pre-selection /// - For `OR`: /// - if LHS is all true => short-circuit → return LHS /// - if LHS is all false => short-circuit → return RHS +/// - if LHS is mixed and false_count / len <= [`PRE_SELECTION_THRESHOLD`] -> pre-selection /// # Arguments /// * `lhs` - The left-hand side (lhs) columnar value (array or scalar) -/// * `lhs` - The left-hand side (lhs) columnar value (array or scalar) /// * `op` - The logical operator (`AND` or `OR`) /// /// # Implementation Notes /// 1. Only works with Boolean-typed arguments (other types automatically return `false`) /// 2. Handles both scalar values and array values /// 3. For arrays, uses optimized bit counting techniques for boolean arrays -fn check_short_circuit<'a>( - lhs: &'a ColumnarValue, - op: &Operator, -) -> ShortCircuitStrategy<'a> { - // Quick reject for non-logical operators,and quick judgment when op is and +fn check_short_circuit(lhs: &ColumnarValue, op: &Operator) -> ShortCircuitStrategy { + // Only logical operators can use this path. let is_and = match op { Operator::And => true, Operator::Or => false, @@ -789,36 +1206,42 @@ fn check_short_circuit<'a>( let true_count = bool_array.values().count_set_bits(); if is_and { - // For AND, prioritize checking for all-false (short circuit case) - // Uses optimized false_count() method provided by Arrow - - // Short circuit if all values are false if true_count == 0 { return ShortCircuitStrategy::ReturnLeft; } - // If no false values, then all must be true if true_count == len { return ShortCircuitStrategy::ReturnRight; } - // determine if we can pre-selection if true_count as f32 / len as f32 <= PRE_SELECTION_THRESHOLD { - return ShortCircuitStrategy::PreSelection(bool_array); + // Select rows where the LHS is true; rows where the LHS + // is false are false regardless of the RHS. + return ShortCircuitStrategy::PreSelection { + mask: bool_array.clone(), + fill_value: false, + }; } } else { - // For OR, prioritize checking for all-true (short circuit case) - // Uses optimized true_count() method provided by Arrow - - // Short circuit if all values are true if true_count == len { return ShortCircuitStrategy::ReturnLeft; } - // If no true values, then all must be false if true_count == 0 { return ShortCircuitStrategy::ReturnRight; } + + let false_count = len - true_count; + if false_count as f32 / len as f32 <= PRE_SELECTION_THRESHOLD { + // Select rows where the LHS is false; rows where the LHS + // is true are true regardless of the RHS. The LHS has no + // nulls here, so negating its bits is infallible. + let mask = BooleanArray::new(!bool_array.values(), None); + return ShortCircuitStrategy::PreSelection { + mask, + fill_value: true, + }; + } } } } @@ -841,62 +1264,54 @@ fn check_short_circuit<'a>( ShortCircuitStrategy::None } -/// Creates a new boolean array based on the evaluation of the right expression, -/// but only for positions where the left_result is true. +/// Collapses a pre-selected expression whose RHS is uniformly `rhs_value` across +/// every selected row, avoiding a scatter: +/// - when it equals `fill_value`, every row is `fill_value` (a scalar); +/// - otherwise the selected rows already equal the RHS, which matches the LHS +/// there, and the unselected rows are the LHS value too, so the result is `lhs`. +fn uniform_pre_selection_result( + rhs_value: bool, + fill_value: bool, + lhs: ColumnarValue, +) -> ColumnarValue { + if rhs_value == fill_value { + ColumnarValue::Scalar(ScalarValue::Boolean(Some(fill_value))) + } else { + lhs + } +} + +/// Creates a boolean array by scattering compact RHS results into the positions +/// selected by `mask`. /// -/// This function is used for short-circuit evaluation optimization of logical AND operations: -/// - When left_result has few true values, we only evaluate the right expression for those positions -/// - Values are copied from right_array where left_result is true -/// - All other positions are filled with false values +/// This function is used for short-circuit evaluation optimization of logical AND/OR operations: +/// - Only selected rows are evaluated on the RHS +/// - Values are copied from `right_result` where `mask` is true +/// - All other positions are filled with `fill_value` (`false` for AND, `true` for OR) /// /// # Parameters -/// - `left_result` Boolean array with selection mask (typically from left side of AND) +/// - `mask` Boolean array with the rows whose result depends on the RHS /// - `right_result` Result of evaluating right side of expression (only for selected positions) +/// - `fill_value` The value for the unselected positions (`false` for AND, `true` for OR) /// /// # Returns -/// A combined ColumnarValue with values from right_result where left_result is true -/// -/// # Example -/// Initial Data: { 1, 2, 3, 4, 5 } -/// Left Evaluation -/// (Condition: Equal to 2 or 3) -/// ↓ -/// Filtered Data: {2, 3} -/// Left Bitmap: { 0, 1, 1, 0, 0 } -/// ↓ -/// Right Evaluation -/// (Condition: Even numbers) -/// ↓ -/// Right Data: { 2 } -/// Right Bitmap: { 1, 0 } -/// ↓ -/// Combine Results -/// Final Bitmap: { 0, 1, 0, 0, 0 } -/// -/// # Note -/// Perhaps it would be better to modify `left_result` directly without creating a copy? -/// In practice, `left_result` should have only one owner, so making changes should be safe. -/// However, this is difficult to achieve under the immutable constraints of [`Arc`] and [`BooleanArray`]. +/// A combined `ColumnarValue` with the same length as `mask`. fn pre_selection_scatter( - left_result: &BooleanArray, + mask: &BooleanArray, right_result: Option<&BooleanArray>, + fill_value: bool, ) -> Result { - let result_len = left_result.len(); + let result_len = mask.len(); let mut result_array_builder = BooleanArray::builder(result_len); - // keep track of current position we have in right boolean array let mut right_array_pos = 0; - - // keep track of how much is filled let mut last_end = 0; - // reduce if condition in for_each match right_result { Some(right_result) => { - SlicesIterator::new(left_result).for_each(|(start, end)| { - // the gap needs to be filled with false + SlicesIterator::new(mask).for_each(|(start, end)| { if start > last_end { - result_array_builder.append_n(start - last_end, false); + result_array_builder.append_n(start - last_end, fill_value); } // copy values from right array for this slice @@ -910,13 +1325,11 @@ fn pre_selection_scatter( last_end = end; }); } - None => SlicesIterator::new(left_result).for_each(|(start, end)| { - // the gap needs to be filled with false + None => SlicesIterator::new(mask).for_each(|(start, end)| { if start > last_end { - result_array_builder.append_n(start - last_end, false); + result_array_builder.append_n(start - last_end, fill_value); } - // append nulls for this slice derictly let len = end - start; result_array_builder.append_nulls(len); @@ -924,49 +1337,15 @@ fn pre_selection_scatter( }), } - // Fill any remaining positions with false + // Fill any remaining positions with `fill_value` if last_end < result_len { - result_array_builder.append_n(result_len - last_end, false); + result_array_builder.append_n(result_len - last_end, fill_value); } let boolean_result = result_array_builder.finish(); Ok(ColumnarValue::Array(Arc::new(boolean_result))) } -fn concat_elements(left: &ArrayRef, right: &ArrayRef) -> Result { - Ok(match left.data_type() { - DataType::Utf8 => Arc::new(concat_elements_utf8( - left.as_string::(), - right.as_string::(), - )?), - DataType::LargeUtf8 => Arc::new(concat_elements_utf8( - left.as_string::(), - right.as_string::(), - )?), - DataType::Utf8View => Arc::new(concat_elements_utf8view( - left.as_string_view(), - right.as_string_view(), - )?), - DataType::Binary => Arc::new(concat_element_binary::( - left.as_binary(), - right.as_binary(), - )?), - DataType::LargeBinary => Arc::new(concat_element_binary::( - left.as_binary(), - right.as_binary(), - )?), - DataType::BinaryView => Arc::new(concat_elements_binary_view_array( - left.as_binary_view(), - right.as_binary_view(), - )?), - other => { - return internal_err!( - "Data type {other:?} not supported for binary operation 'concat_elements' on string arrays" - ); - } - }) -} - /// Create a binary expression whose arguments are correctly coerced. /// This function errors if it is not possible to coerce the arguments /// to computational types supported by the operator. @@ -1001,12 +1380,181 @@ mod tests { use crate::expressions::{Column, Literal, col, lit, try_cast}; use datafusion_expr::lit as expr_lit; - use datafusion_common::plan_datafusion_err; + use datafusion_common::{assert_contains, plan_datafusion_err}; use datafusion_physical_expr_common::physical_expr::fmt_sql; use crate::planner::logical2physical; use arrow::array::BooleanArray; + use arrow::compute::SortOptions; use datafusion_expr::col as logical_col; + + #[test] + fn test_arithmetic_ordering_overflow() -> Result<()> { + let asc = SortProperties::Ordered(Default::default()); + let ordered = |range: Interval| ExprProperties { + sort_properties: asc, + range, + preserves_lex_ordering: false, + strictly_order_preserving: false, + }; + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let a_plus_b = + BinaryExpr::new(col("a", &schema)?, Operator::Plus, col("b", &schema)?); + let unbounded = [ + ordered(Interval::make_unbounded(&DataType::Int32)?), + ordered(Interval::make_unbounded(&DataType::Int32)?), + ]; + let bounded = [ + ordered(Interval::make(Some(0), Some(10))?), + ordered(Interval::make(Some(0), Some(10))?), + ]; + + // Unknown ranges: the sum may overflow and wrap, so it is unordered. + assert_eq!( + a_plus_b.get_properties(&unbounded)?.sort_properties, + SortProperties::Unordered + ); + // Bounded ranges that cannot overflow keep the ordering, as does + // checked arithmetic, which errors instead of wrapping. + assert_eq!(a_plus_b.get_properties(&bounded)?.sort_properties, asc); + let checked = a_plus_b.with_fail_on_overflow(true); + assert_eq!(checked.get_properties(&unbounded)?.sort_properties, asc); + + // `time + interval` wraps around the 24-hour clock even in checked + // mode, so it never preserves ordering. + let time = DataType::Time64(TimeUnit::Nanosecond); + let interval = DataType::Interval(IntervalUnit::MonthDayNano); + let schema = Schema::new(vec![ + Field::new("t", time.clone(), false), + Field::new("i", interval.clone(), false), + ]); + let time_plus_interval = + BinaryExpr::new(col("t", &schema)?, Operator::Plus, col("i", &schema)?) + .with_fail_on_overflow(true); + let time_props = [ + ordered(Interval::make_unbounded(&time)?), + ordered(Interval::make_unbounded(&interval)?), + ]; + assert_eq!( + time_plus_interval + .get_properties(&time_props)? + .sort_properties, + SortProperties::Unordered + ); + + Ok(()) + } + + /// `a - b` only derives an ordering when `a` and `b` are ordered in + /// opposite directions, so every case below pairs an ascending left-hand + /// side with a descending right-hand side. + #[test] + fn test_subtraction_ordering_overflow() -> Result<()> { + let asc = SortProperties::Ordered(SortOptions { + descending: false, + nulls_first: true, + }); + let desc = SortProperties::Ordered(SortOptions { + descending: true, + nulls_first: true, + }); + let props = |sort_properties, range| ExprProperties { + sort_properties, + range, + preserves_lex_ordering: false, + strictly_order_preserving: false, + }; + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let a_minus_b = + BinaryExpr::new(col("a", &schema)?, Operator::Minus, col("b", &schema)?); + + // Signed minimum: the difference can underflow past `i32::MIN` and + // wrap around to large positive values. + let signed_underflow = [ + props(asc, Interval::make(Some(i32::MIN), Some(0))?), + props(desc, Interval::make(Some(0), Some(i32::MAX))?), + ]; + assert_eq!( + a_minus_b.get_properties(&signed_underflow)?.sort_properties, + SortProperties::Unordered + ); + // The very same ranges keep the ordering under checked arithmetic, + // which errors instead of wrapping. + let checked = a_minus_b.clone().with_fail_on_overflow(true); + assert_eq!( + checked.get_properties(&signed_underflow)?.sort_properties, + asc + ); + // Ranges whose difference stays inside `Int32` are safe. + let signed_safe = [ + props(asc, Interval::make(Some(0), Some(10))?), + props(desc, Interval::make(Some(0), Some(10))?), + ]; + assert_eq!(a_minus_b.get_properties(&signed_safe)?.sort_properties, asc); + + let schema = Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ]); + let a_minus_b = + BinaryExpr::new(col("a", &schema)?, Operator::Minus, col("b", &schema)?); + + // Unsigned underflow: the ranges overlap, so `0 - 1` wraps to + // `u32::MAX` even though both operands are bounded. + let unsigned_underflow = [ + props(asc, Interval::make(Some(0_u32), Some(10_u32))?), + props(desc, Interval::make(Some(0_u32), Some(10_u32))?), + ]; + assert_eq!( + a_minus_b + .get_properties(&unsigned_underflow)? + .sort_properties, + SortProperties::Unordered + ); + // A left-hand range that always dominates the right-hand one cannot + // underflow. + let unsigned_safe = [ + props(asc, Interval::make(Some(10_u32), Some(20_u32))?), + props(desc, Interval::make(Some(0_u32), Some(5_u32))?), + ]; + assert_eq!( + a_minus_b.get_properties(&unsigned_safe)?.sort_properties, + asc + ); + + // `time - interval` wraps around the 24-hour clock even in checked + // mode, so it never preserves ordering. + let time = DataType::Time64(TimeUnit::Nanosecond); + let interval = DataType::Interval(IntervalUnit::MonthDayNano); + let schema = Schema::new(vec![ + Field::new("t", time.clone(), false), + Field::new("i", interval.clone(), false), + ]); + let time_minus_interval = + BinaryExpr::new(col("t", &schema)?, Operator::Minus, col("i", &schema)?) + .with_fail_on_overflow(true); + let time_props = [ + props(asc, Interval::make_unbounded(&time)?), + props(desc, Interval::make_unbounded(&interval)?), + ]; + assert_eq!( + time_minus_interval + .get_properties(&time_props)? + .sort_properties, + SortProperties::Unordered + ); + + Ok(()) + } + /// Performs a binary operation, applying any type coercion necessary fn binary_op( left: Arc, @@ -1910,6 +2458,82 @@ mod tests { Ok(()) } + #[test] + fn date32_minus_date32_returns_int64_days() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Date32, true), + Field::new("b", DataType::Date32, true), + ])); + let a = Arc::new(Date32Array::from(vec![ + Some(18_901), + Some(18_901), + None, + Some(18_900), + ])); + let b = Arc::new(Date32Array::from(vec![ + Some(18_898), + Some(18_904), + Some(18_900), + None, + ])); + + apply_arithmetic::( + schema, + vec![a, b], + Operator::Minus, + Int64Array::from(vec![Some(3), Some(-3), None, None]), + )?; + + Ok(()) + } + + #[test] + fn date64_minus_date64_returns_int64_days() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Date64, true), + Field::new("b", DataType::Date64, true), + ])); + let a = Arc::new(Date64Array::from(vec![ + Some(18_901 * MILLIS_PER_DAY), + Some(18_901 * MILLIS_PER_DAY), + None, + Some(18_900 * MILLIS_PER_DAY), + ])); + let b = Arc::new(Date64Array::from(vec![ + Some(18_898 * MILLIS_PER_DAY), + Some(18_904 * MILLIS_PER_DAY), + Some(18_900 * MILLIS_PER_DAY), + None, + ])); + + apply_arithmetic::( + schema, + vec![a, b], + Operator::Minus, + Int64Array::from(vec![Some(3), Some(-3), None, None]), + )?; + + Ok(()) + } + + #[test] + fn date32_minus_null_scalar_returns_int64_null_scalar() -> Result<()> { + let result = apply_date_subtraction( + &ColumnarValue::Array(Arc::new(Date32Array::from(vec![ + Some(18_901), + Some(18_900), + ]))), + &ColumnarValue::Scalar(ScalarValue::Date32(None)), + )?; + + assert!(matches!( + result, + ColumnarValue::Scalar(ScalarValue::Int64(None)) + )); + + Ok(()) + } + #[test] fn minus_op_dict() -> Result<()> { let schema = Schema::new(vec![ @@ -2977,6 +3601,105 @@ mod tests { Ok(()) } + #[test] + fn regex_scalar_with_dictionary_nulls() -> Result<()> { + let dictionary_values = Arc::new(StringArray::from(vec![ + Some("abc"), + None, + Some("ABC"), + Some("def"), + ])); + let keys = UInt32Array::from(vec![Some(0), None, Some(1), Some(2), Some(3)]); + let dictionary = + Arc::new(DictionaryArray::try_new(keys, dictionary_values)?) as ArrayRef; + let utf8 = cast(&dictionary, &DataType::Utf8)?; + let pattern = ScalarValue::Utf8(Some("^abc$".to_string())); + let dictionary_schema = Arc::new(Schema::new(vec![Field::new( + "a", + dictionary.data_type().clone(), + true, + )])); + let utf8_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, true)])); + + let evaluate = + |schema: &SchemaRef, array: &ArrayRef, op: Operator| -> Result { + let expr = binary(col("a", schema)?, op, lit(pattern.clone()), schema)?; + let batch = + RecordBatch::try_new(Arc::clone(schema), vec![Arc::clone(array)])?; + Ok(expr + .evaluate(&batch)? + .into_array(batch.num_rows()) + .expect("Failed to convert to array")) + }; + + for (op, expected) in [ + ( + Operator::RegexMatch, + BooleanArray::from(vec![ + Some(true), + None, + None, + Some(false), + Some(false), + ]), + ), + ( + Operator::RegexIMatch, + BooleanArray::from(vec![Some(true), None, None, Some(true), Some(false)]), + ), + ( + Operator::RegexNotMatch, + BooleanArray::from(vec![Some(false), None, None, Some(true), Some(true)]), + ), + ( + Operator::RegexNotIMatch, + BooleanArray::from(vec![ + Some(false), + None, + None, + Some(false), + Some(true), + ]), + ), + ] { + let dictionary_result = evaluate(&dictionary_schema, &dictionary, op)?; + let utf8_result = evaluate(&utf8_schema, &utf8, op)?; + + assert_eq!(dictionary_result.as_ref(), &expected); + assert_eq!(&dictionary_result, &utf8_result); + } + + Ok(()) + } + + #[test] + fn regex_mismatched_array_types_error() -> Result<()> { + // The analyzer coerces both operands of a regex operator to a common + // string type, but an expression that bypasses it (e.g. constructed + // directly) must return an error instead of panicking + // (https://github.com/apache/datafusion/issues/22886) + let schema = Schema::new(vec![ + Field::new("a", DataType::Utf8View, true), + Field::new("b", DataType::Utf8, true), + ]); + let a = Arc::new(StringViewArray::from(vec!["user auth failed"])) as ArrayRef; + let b = Arc::new(StringArray::from(vec!["(auth|login)"])) as ArrayRef; + + // construct the expression directly, without coercion + let expr = binary( + col("a", &schema)?, + Operator::RegexMatch, + col("b", &schema)?, + &schema, + )?; + let batch = RecordBatch::try_new(Arc::new(schema), vec![a, b])?; + let err = expr.evaluate(&batch).unwrap_err(); + assert_contains!(err.to_string(), "failed to downcast array"); + + Ok(()) + } + #[test] fn or_with_nulls_op() -> Result<()> { let schema = Schema::new(vec![ @@ -5019,14 +5742,17 @@ mod tests { let ColumnarValue::Array(array) = &left_value else { panic!("Expected ColumnarValue::Array"); }; - let ShortCircuitStrategy::PreSelection(value) = + let ShortCircuitStrategy::PreSelection { mask, fill_value } = check_short_circuit(&left_value, &Operator::And) else { panic!("Expected ShortCircuitStrategy::PreSelection"); }; + // For AND, the mask selects the rows where the LHS is true and the + // unselected rows are filled with `false`. + assert!(!fill_value); let expected_boolean_arr: Vec<_> = as_boolean_array(array).unwrap().iter().collect(); - let boolean_arr: Vec<_> = value.iter().collect(); + let boolean_arr: Vec<_> = mask.iter().collect(); assert_eq!(expected_boolean_arr, boolean_arr); // op: OR left: all true @@ -5037,10 +5763,33 @@ mod tests { ShortCircuitStrategy::ReturnLeft )); - // op: OR left: not all true + // 20% false: OR can pre-select the false rows. let left_expr: Arc = logical2physical(&logical_col("a").gt(expr_lit(2)), &schema); let left_value = left_expr.evaluate(&batch).unwrap(); + let ColumnarValue::Array(array) = &left_value else { + panic!("Expected ColumnarValue::Array"); + }; + let ShortCircuitStrategy::PreSelection { mask, fill_value } = + check_short_circuit(&left_value, &Operator::Or) + else { + panic!("Expected ShortCircuitStrategy::PreSelection"); + }; + // For OR, the mask selects the rows where the LHS is false (the negation + // of the LHS) and the unselected rows are filled with `true`. + assert!(fill_value); + let negated_lhs: Vec<_> = as_boolean_array(array) + .unwrap() + .iter() + .map(|v| v.map(|b| !b)) + .collect(); + let boolean_arr: Vec<_> = mask.iter().collect(); + assert_eq!(negated_lhs, boolean_arr); + + // 60% false: OR falls back to normal evaluation. + let left_expr: Arc = + logical2physical(&logical_col("a").gt(expr_lit(4)), &schema); + let left_value = left_expr.evaluate(&batch).unwrap(); assert!(matches!( check_short_circuit(&left_value, &Operator::Or), ShortCircuitStrategy::None @@ -5144,15 +5893,10 @@ mod tests { )); } - /// Test for [pre_selection_scatter] - /// Since [check_short_circuit] ensures that the left side does not contain null and is neither all_true nor all_false, as well as not being empty, - /// the following tests have been designed: - /// 1. Test sparse left with interleaved true/false - /// 2. Test multiple consecutive true blocks - /// 3. Test multiple consecutive true blocks - /// 4. Test single true at first position - /// 5. Test single true at last position - /// 6. Test nulls in right array + /// Test for [pre_selection_scatter]. + /// + /// `check_short_circuit` only calls this helper with a non-empty, + /// non-null mask that is neither all true nor all false. #[test] fn test_pre_selection_scatter() { fn create_bool_array(bools: Vec) -> BooleanArray { @@ -5165,7 +5909,7 @@ mod tests { let left = create_bool_array(vec![true, false, true, false, true]); let right = create_bool_array(vec![false, true, false]); - let result = pre_selection_scatter(&left, Some(&right)).unwrap(); + let result = pre_selection_scatter(&left, Some(&right), false).unwrap(); let result_arr = result.into_array(left.len()).unwrap(); let expected = create_bool_array(vec![false, false, true, false, false]); @@ -5179,7 +5923,7 @@ mod tests { create_bool_array(vec![false, true, true, false, true, true, true]); let right = create_bool_array(vec![true, false, false, true, false]); - let result = pre_selection_scatter(&left, Some(&right)).unwrap(); + let result = pre_selection_scatter(&left, Some(&right), false).unwrap(); let result_arr = result.into_array(left.len()).unwrap(); let expected = @@ -5193,7 +5937,7 @@ mod tests { let left = create_bool_array(vec![true, false, false]); let right = create_bool_array(vec![false]); - let result = pre_selection_scatter(&left, Some(&right)).unwrap(); + let result = pre_selection_scatter(&left, Some(&right), false).unwrap(); let result_arr = result.into_array(left.len()).unwrap(); let expected = create_bool_array(vec![false, false, false]); @@ -5206,7 +5950,7 @@ mod tests { let left = create_bool_array(vec![false, false, true]); let right = create_bool_array(vec![false]); - let result = pre_selection_scatter(&left, Some(&right)).unwrap(); + let result = pre_selection_scatter(&left, Some(&right), false).unwrap(); let result_arr = result.into_array(left.len()).unwrap(); let expected = create_bool_array(vec![false, false, false]); @@ -5219,7 +5963,7 @@ mod tests { let left = create_bool_array(vec![false, true, false, true]); let right = BooleanArray::from(vec![None, Some(false)]); - let result = pre_selection_scatter(&left, Some(&right)).unwrap(); + let result = pre_selection_scatter(&left, Some(&right), false).unwrap(); let result_arr = result.into_array(left.len()).unwrap(); let expected = BooleanArray::from(vec![ @@ -5230,6 +5974,38 @@ mod tests { ]); assert_eq!(&expected, result_arr.as_boolean()); } + // OR semantics: selected rows take the RHS, unselected rows become true. + { + // Selection (LHS false rows): [T, F, T, F, T] + // Right (RHS on those rows): [F, T, F] + let left = create_bool_array(vec![true, false, true, false, true]); + let right = create_bool_array(vec![false, true, false]); + + let result = pre_selection_scatter(&left, Some(&right), true).unwrap(); + let result_arr = result.into_array(left.len()).unwrap(); + + // selected rows take the RHS value; unselected rows are `true` + let expected = create_bool_array(vec![false, true, true, true, false]); + assert_eq!(&expected, result_arr.as_boolean()); + } + // OR semantics with nulls in the right array. + { + // Selection (LHS false rows): [F, T, F, T] + // Right: [None, Some(false)] + let left = create_bool_array(vec![false, true, false, true]); + let right = BooleanArray::from(vec![None, Some(false)]); + + let result = pre_selection_scatter(&left, Some(&right), true).unwrap(); + let result_arr = result.into_array(left.len()).unwrap(); + + let expected = BooleanArray::from(vec![ + Some(true), // unselected => true + None, // null from right + Some(true), // unselected => true + Some(false), + ]); + assert_eq!(&expected, result_arr.as_boolean()); + } } #[test] @@ -5256,6 +6032,89 @@ mod tests { ); } + #[test] + fn test_or_false_preselection_returns_lhs() { + // `c OR false` over a mostly-true `c` triggers OR pre-selection; the + // result must equal `c`. + let schema = + Arc::new(Schema::new(vec![Field::new("c", DataType::Boolean, false)])); + let c_array = + Arc::new(BooleanArray::from(vec![true, false, true, true, true])) as ArrayRef; + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![Arc::clone(&c_array)]) + .unwrap(); + + let expr = logical2physical(&logical_col("c").or(expr_lit(false)), &schema); + + let result = expr.evaluate(&batch).unwrap(); + let ColumnarValue::Array(result_arr) = result else { + panic!("Expected ColumnarValue::Array"); + }; + + let expected: Vec<_> = c_array.as_boolean().iter().collect(); + let actual: Vec<_> = result_arr.as_boolean().iter().collect(); + assert_eq!( + expected, actual, + "OR with FALSE must equal LHS even with PreSelection" + ); + } + + #[test] + fn test_or_preselection_matches_kleene() { + // The OR pre-selection path must match full-batch Kleene OR. + use arrow::compute::kernels::boolean::or_kleene; + + let schema = Arc::new(Schema::new(vec![ + Field::new("c", DataType::Boolean, true), + Field::new("d", DataType::Boolean, true), + ])); + + // `c` is mostly true (2/10 false => 20% <= threshold) so OR pre-selects. + let c = BooleanArray::from(vec![ + true, true, false, true, true, true, true, false, true, true, + ]); + + let d_cases = vec![ + // Mixed RHS with nulls exercises scatter and null copy. + BooleanArray::from(vec![ + Some(false), + Some(true), + Some(true), + Some(false), + Some(false), + Some(true), + Some(false), + None, + Some(true), + None, + ]), + // RHS true on selected rows exercises the uniform-fill path. + BooleanArray::from(vec![Some(true); 10]), + // RHS false on selected rows exercises the return-LHS path. + BooleanArray::from(vec![Some(false); 10]), + ]; + + for d in d_cases { + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(c.clone()) as ArrayRef, + Arc::new(d.clone()) as ArrayRef, + ], + ) + .unwrap(); + + let expr = logical2physical(&logical_col("c").or(logical_col("d")), &schema); + let result = expr.evaluate(&batch).unwrap().into_array(c.len()).unwrap(); + + let expected = or_kleene(&c, &d).unwrap(); + assert_eq!( + expected, + *result.as_boolean(), + "OR pre-selection must match Kleene OR for d = {d:?}" + ); + } + } + #[test] fn test_evaluate_bounds_int32() { let schema = Schema::new(vec![ diff --git a/datafusion/physical-expr/src/expressions/binary/kernels.rs b/datafusion/physical-expr/src/expressions/binary/kernels.rs index e573d7ece2afa..a123fba1f9da2 100644 --- a/datafusion/physical-expr/src/expressions/binary/kernels.rs +++ b/datafusion/physical-expr/src/expressions/binary/kernels.rs @@ -18,7 +18,6 @@ //! This module contains computation kernels that are specific to //! datafusion and not (yet) targeted to port upstream to arrow use arrow::array::*; -use arrow::buffer::{MutableBuffer, NullBuffer}; use arrow::compute::kernels::bitwise::{ bitwise_and, bitwise_and_scalar, bitwise_or, bitwise_or_scalar, bitwise_shift_left, bitwise_shift_left_scalar, bitwise_shift_right, bitwise_shift_right_scalar, @@ -27,9 +26,8 @@ use arrow::compute::kernels::bitwise::{ use arrow::compute::kernels::boolean::not; use arrow::compute::kernels::comparison::{regexp_is_match, regexp_is_match_scalar}; use arrow::datatypes::DataType; -use arrow::error::ArrowError; use datafusion_common::{Result, ScalarValue}; -use datafusion_common::{internal_err, plan_err}; +use datafusion_common::{exec_err, internal_err, plan_err}; use std::sync::Arc; @@ -161,104 +159,30 @@ create_left_integral_dyn_scalar_kernel!( bitwise_shift_left_scalar ); -/// Concatenates two `StringViewArray`s element-wise. -/// If either element is `Null`, the result element is also `Null`. -/// -/// # Errors -/// - Returns an error if the input arrays have different lengths. -/// - Returns an error if any concatenated string exceeds `u32::MAX` (≈4 GB) in length. -pub fn concat_elements_utf8view( - left: &StringViewArray, - right: &StringViewArray, -) -> std::result::Result { - if left.len() != right.len() { - return Err(ArrowError::ComputeError(format!( - "Arrays must have the same length: {} != {}", - left.len(), - right.len() - ))); - } - let mut result = StringViewBuilder::with_capacity(left.len()); - - // Avoid reallocations by writing to a reused buffer (note we could be even - // more efficient by creating the view directly here and avoid the buffer - // but that would be more complex) - let mut buffer = String::new(); - - // Pre-compute combined null bitmap, so the per-row NULL check is more - // efficient - let nulls = NullBuffer::union(left.nulls(), right.nulls()); - - for i in 0..left.len() { - if nulls.as_ref().is_some_and(|n| n.is_null(i)) { - result.append_null(); - } else { - let l = left.value(i); - let r = right.value(i); - buffer.clear(); - buffer.push_str(l); - buffer.push_str(r); - result.try_append_value(&buffer)?; - } - } - Ok(result.finish()) -} - -/// Concatenates two `BinaryViewArray`s element-wise. -/// If either element is `Null`, the result element is also `Null`. -/// -/// # Errors -/// - Returns an error if the input arrays have different lengths. -/// - Returns an error if any concatenated string exceeds `u32::MAX` in length. -pub fn concat_elements_binary_view_array( - left: &BinaryViewArray, - right: &BinaryViewArray, -) -> std::result::Result { - if left.len() != right.len() { - return Err(ArrowError::ComputeError(format!( - "Arrays must have the same length: {} != {}", - left.len(), - right.len() - ))); - } - let mut result = BinaryViewBuilder::with_capacity(left.len()); - - // Avoid reallocations by writing to a reused buffer (note we could be even - // more efficient by creating the view directly here and avoid the buffer - // but that would be more complex) - let mut buffer = MutableBuffer::new(0); - - // Pre-compute combined null bitmap, so the per-row NULL check is more - // efficient - let nulls = NullBuffer::union(left.nulls(), right.nulls()); - - for i in 0..left.len() { - if nulls.as_ref().is_some_and(|n| n.is_null(i)) { - result.append_null(); - } else { - let l = left.value(i); - let r = right.value(i); - buffer.clear(); - buffer.extend_from_slice(l); - buffer.extend_from_slice(r); - // No try-version of append_value - result.try_append_value(&buffer)?; - } - } - Ok(result.finish()) -} - /// Invoke a compute kernel on a pair of binary data arrays with flags macro_rules! regexp_is_match_flag { ($LEFT:expr, $RIGHT:expr, $ARRAYTYPE:ident, $NOT:expr, $FLAG:expr) => {{ - let ll = $LEFT - .as_any() - .downcast_ref::<$ARRAYTYPE>() - .expect("failed to downcast array"); - let rr = $RIGHT - .as_any() - .downcast_ref::<$ARRAYTYPE>() - .expect("failed to downcast array"); + // The analyzer coerces both operands to a common string type, but + // expressions that bypass it may still reach here with mismatched + // types, which must surface as an error rather than a panic. + let ll = match $LEFT.as_any().downcast_ref::<$ARRAYTYPE>() { + Some(ll) => ll, + None => { + return exec_err!( + "failed to downcast array to {} for operation 'regex_match_dyn'", + stringify!($ARRAYTYPE) + ); + } + }; + let rr = match $RIGHT.as_any().downcast_ref::<$ARRAYTYPE>() { + Some(rr) => rr, + None => { + return exec_err!( + "failed to downcast array to {} for operation 'regex_match_dyn'", + stringify!($ARRAYTYPE) + ); + } + }; let flag = if $FLAG { Some($ARRAYTYPE::from(vec!["i"; ll.len()])) @@ -299,10 +223,15 @@ pub(crate) fn regex_match_dyn( /// Invoke a compute kernel on a data array and a scalar value with flag macro_rules! regexp_is_match_flag_scalar { ($LEFT:expr, $RIGHT:expr, $ARRAYTYPE:ident, $NOT:expr, $FLAG:expr) => {{ - let ll = $LEFT - .as_any() - .downcast_ref::<$ARRAYTYPE>() - .expect("failed to downcast array"); + let ll = match $LEFT.as_any().downcast_ref::<$ARRAYTYPE>() { + Some(ll) => ll, + None => { + return Some(exec_err!( + "failed to downcast array to {} for operation 'regex_match_dyn_scalar'", + stringify!($ARRAYTYPE) + )); + } + }; if let Some(Some(string_value)) = $RIGHT.try_as_str() { let flag = $FLAG.then_some("i"); @@ -341,7 +270,8 @@ pub(crate) fn regex_match_dyn_scalar( regexp_is_match_flag_scalar!(left, right, LargeStringArray, not_match, flag) } DataType::Dictionary(_, _) => { - let values = left.as_any_dictionary().values(); + let dictionary = left.as_any_dictionary(); + let values = dictionary.values(); match values.data_type() { DataType::Utf8 => regexp_is_match_flag_scalar!(values, right, StringArray, not_match, flag), @@ -351,16 +281,15 @@ pub(crate) fn regex_match_dyn_scalar( "Data type {} not supported as a dictionary value type for operation 'regex_match_dyn_scalar' on string array", other ), - }.map( - // downcast_dictionary_array duplicates code per possible key type, so we aim to do all prep work before - |evaluated_values| downcast_dictionary_array! { - left => { - let unpacked_dict = evaluated_values.take_iter(left.keys().iter().map(|opt| opt.map(|v| v as _))).collect::(); - Arc::new(unpacked_dict) as ArrayRef - }, - _ => unreachable!(), - } - ) + } + .and_then(|evaluated_values| { + // Expand back to rows while preserving nulls from both keys and values. + Ok(arrow::compute::take( + evaluated_values.as_ref(), + dictionary.keys(), + None, + )?) + }) } other => internal_err!( "Data type {} not supported for operation 'regex_match_dyn_scalar' on string array", diff --git a/datafusion/physical-expr/src/expressions/case.rs b/datafusion/physical-expr/src/expressions/case.rs index 568ecb9cf336b..17288a9737699 100644 --- a/datafusion/physical-expr/src/expressions/case.rs +++ b/datafusion/physical-expr/src/expressions/case.rs @@ -19,7 +19,9 @@ mod literal_lookup_table; use super::{Column, Literal}; use crate::PhysicalExpr; -use crate::expressions::{LambdaVariable, lit, try_cast}; +use crate::expressions::{ + CastExpr, LambdaVariable, NegativeExpr, NotExpr, lit, try_cast, +}; use arrow::array::*; use arrow::compute::kernels::zip::zip; use arrow::compute::{ @@ -1278,7 +1280,11 @@ impl PhysicalExpr for CaseExpr { // it would evaluate to null. // Replace the `then` expression with `NULL` in the `when` expression - let with_null = match replace_with_null(w, t.as_ref(), input_schema) { + let with_null = match replace_with_null( + w, + unwrap_certainly_null_expr(t.as_ref()), + input_schema, + ) { Err(e) => return Some(Err(e)), Ok(e) => e, }; @@ -1410,6 +1416,86 @@ impl PhysicalExpr for CaseExpr { } write!(f, "END") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Case(Box::new( + protobuf::PhysicalCaseNode { + expr: self + .expr() + .map(|expr| ctx.encode_child(expr).map(Box::new)) + .transpose()?, + when_then_expr: self + .when_then_expr() + .iter() + .map(|(when_expr, then_expr)| { + Ok(protobuf::PhysicalWhenThen { + when_expr: Some(ctx.encode_child(when_expr)?), + then_expr: Some(ctx.encode_child(then_expr)?), + }) + }) + .collect::>>()?, + else_expr: self + .else_expr() + .map(|expr| ctx.encode_child(expr).map(Box::new)) + .transpose()?, + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl CaseExpr { + /// Reconstruct a [`CaseExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let case = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Case, + "CaseExpr", + ); + + Ok(Arc::new(CaseExpr::try_new( + case.expr + .as_deref() + .map(|expr| ctx.decode(expr)) + .transpose()?, + case.when_then_expr + .iter() + .map(|when_then| { + Ok(( + ctx.decode_required_expression( + when_then.when_expr.as_ref(), + "CaseExpr", + "when_expr", + )?, + ctx.decode_required_expression( + when_then.then_expr.as_ref(), + "CaseExpr", + "then_expr", + )?, + )) + }) + .collect::>>()?, + case.else_expr + .as_deref() + .map(|expr| ctx.decode(expr)) + .transpose()?, + )?)) + } } /// Attempts to const evaluate the given `predicate`. @@ -1457,6 +1543,25 @@ fn replace_with_null( Ok(with_null) } +/// Returns the innermost [`PhysicalExpr`] that is provably null if `expr` is null. +/// +/// Keep this in sync with the logical-plan equivalent, `unwrap_certainly_null_expr` +/// in `datafusion/expr/src/expr_schema.rs`. If the two disagree on which wrappers +/// are null-preserving, `CASE` nullability computed by the logical and physical +/// planners can diverge and cause a schema mismatch during planning. +/// See for rationale. +fn unwrap_certainly_null_expr(expr: &dyn PhysicalExpr) -> &dyn PhysicalExpr { + if let Some(expr) = expr.downcast_ref::() { + unwrap_certainly_null_expr(expr.arg().as_ref()) + } else if let Some(expr) = expr.downcast_ref::() { + unwrap_certainly_null_expr(expr.arg().as_ref()) + } else if let Some(expr) = expr.downcast_ref::() { + unwrap_certainly_null_expr(expr.expr.as_ref()) + } else { + expr + } +} + /// Create a CASE expression pub fn case( expr: Option>, @@ -2497,10 +2602,45 @@ mod tests { let zero = lit(0); let foo_eq_zero = binary(Arc::clone(&foo), Operator::Eq, Arc::clone(&zero), &schema)?; + let cast_foo = cast(Arc::clone(&foo), &schema, DataType::Int64)?; + let negative_foo = expressions::negative(Arc::clone(&foo), &schema)?; assert_not_nullable(when_then_else(&foo_is_not_null, &foo, &zero)?, &schema); assert_not_nullable(when_then_else(¬_foo_is_null, &foo, &zero)?, &schema); assert_not_nullable(when_then_else(&foo_eq_zero, &foo, &zero)?, &schema); + assert_not_nullable( + when_then_else(&foo_is_not_null, &cast_foo, &lit(0i64))?, + &schema, + ); + assert_not_nullable( + when_then_else(&foo_is_not_null, &negative_foo, &zero)?, + &schema, + ); + + // Nested null-preserving wrappers must be unwrapped recursively. `CAST(-foo)` + // still collapses `foo IS NOT NULL` to `false`, so the branch is + // unreachable-as-null and the `CASE` is not nullable. + let cast_negative_foo = cast( + expressions::negative(Arc::clone(&foo), &schema)?, + &schema, + DataType::Int64, + )?; + assert_not_nullable( + when_then_else(&foo_is_not_null, &cast_negative_foo, &lit(0i64))?, + &schema, + ); + + // `TRY_CAST` is intentionally NOT treated as null-preserving: it yields + // NULL on a failed cast even for a non-null input, so a guarded `TRY_CAST` + // branch is still reachable-as-null and the `CASE` stays nullable. This must + // stay consistent with the logical planner (`unwrap_certainly_null_expr` in + // `datafusion/expr/src/expr_schema.rs`); unwrapping it on only one side would + // reintroduce a logical/physical schema mismatch. + let try_cast_foo = try_cast(Arc::clone(&foo), &schema, DataType::Int64)?; + assert_nullable( + when_then_else(&foo_is_not_null, &try_cast_foo, &lit(0i64))?, + &schema, + ); assert_not_nullable( when_then_else( @@ -2622,6 +2762,23 @@ mod tests { &schema, ); + let boolean_schema = + Schema::new(vec![Field::new("predicate", DataType::Boolean, true)]); + let predicate = col("predicate", &boolean_schema)?; + let predicate_is_not_null = is_not_null(Arc::clone(&predicate))?; + let not_predicate = expressions::not(Arc::clone(&predicate))?; + assert_not_nullable( + when_then_else(&predicate_is_not_null, ¬_predicate, &lit(false))?, + &boolean_schema, + ); + + // Nested `NOT` is likewise unwrapped recursively. + let not_not_predicate = expressions::not(Arc::clone(¬_predicate))?; + assert_not_nullable( + when_then_else(&predicate_is_not_null, ¬_not_predicate, &lit(false))?, + &boolean_schema, + ); + Ok(()) } @@ -3193,3 +3350,177 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::col; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf; + use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalWhenThen}; + + fn proto_case_fixture() -> CaseExpr { + let schema = Schema::new(vec![Field::new("a", DataType::Boolean, true)]); + CaseExpr::try_new( + Some(col("a", &schema).unwrap()), + vec![(lit(true), lit(1_i32))], + Some(lit(0_i32)), + ) + .unwrap() + } + + fn proto_when_then( + when_expr: Option, + then_expr: Option, + ) -> PhysicalWhenThen { + PhysicalWhenThen { + when_expr, + then_expr, + } + } + + fn proto_case_node( + expr: Option>, + when_then_expr: Vec, + else_expr: Option>, + ) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Case(Box::new( + protobuf::PhysicalCaseNode { + expr, + when_then_expr, + else_expr, + }, + ))), + } + } + + #[test] + fn try_to_proto_encodes_case_expr() { + let case = proto_case_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = case + .try_to_proto(&ctx) + .unwrap() + .expect("CaseExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let case_node = match node.expr_type { + Some(protobuf::physical_expr_node::ExprType::Case(boxed)) => *boxed, + other => panic!("expected a CaseExpr node, got {other:?}"), + }; + assert!(case_node.expr.is_some()); + assert_eq!(case_node.when_then_expr.len(), 1); + assert!(case_node.when_then_expr[0].when_expr.is_some()); + assert!(case_node.when_then_expr[0].then_expr.is_some()); + assert!(case_node.else_expr.is_some()); + } + + #[test] + fn try_to_proto_propagates_child_encode_error() { + let case = proto_case_fixture(); + // Call 1 is the optional CASE expr, call 2 is the WHEN expr. + let encoder = StubEncoder::failing_on(2); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let err = case.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2"))); + } + + #[test] + fn try_from_proto_decodes_case_expr() { + let node = proto_case_node( + Some(Box::new(column_node("case"))), + vec![proto_when_then( + Some(column_node("when")), + Some(column_node("then")), + )], + Some(Box::new(column_node("else"))), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = CaseExpr::try_from_proto(&node, &ctx).unwrap(); + let case = decoded + .downcast_ref::() + .expect("decoded expr should be a CaseExpr"); + + assert!(case.expr().is_some()); + assert_eq!(case.when_then_expr().len(), 1); + assert!(case.else_expr().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_case_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a CaseExpr")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_when_expr() { + let node = proto_case_node( + None, + vec![proto_when_then(None, Some(column_node("then")))], + None, + ); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("CaseExpr is missing required field 'when_expr'")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_then_expr() { + let node = proto_case_node( + None, + vec![proto_when_then(Some(column_node("when")), None)], + None, + ); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("CaseExpr is missing required field 'then_expr'")) + ); + } + + #[test] + fn try_from_proto_propagates_child_decode_error() { + let node = proto_case_node( + Some(Box::new(column_node("case"))), + vec![proto_when_then( + Some(column_node("when")), + Some(column_node("then")), + )], + Some(Box::new(column_node("else"))), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(2); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CaseExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2"))); + } +} diff --git a/datafusion/physical-expr/src/expressions/cast.rs b/datafusion/physical-expr/src/expressions/cast.rs index ad214a89ceb71..dbb91e365af90 100644 --- a/datafusion/physical-expr/src/expressions/cast.rs +++ b/datafusion/physical-expr/src/expressions/cast.rs @@ -179,11 +179,8 @@ impl CastExpr { | (UInt8, UInt16 | UInt32 | UInt64) | (UInt16, UInt32 | UInt64) | (UInt32, UInt64) - | ( - Int8 | Int16 | Int32 | UInt8 | UInt16 | UInt32, - Float32 | Float64 - ) - | (Int64 | UInt64, Float64) + | (Int8 | Int16 | UInt8 | UInt16, Float32) + | (Int8 | Int16 | Int32 | UInt8 | UInt16 | UInt32, Float64) | (Utf8, LargeUtf8) ) } @@ -214,8 +211,18 @@ pub(crate) fn cast_expr_properties( target_type: &DataType, ) -> Result { let unbounded = Interval::make_unbounded(target_type)?; - if is_order_preserving_cast_family(&child.range.data_type(), target_type) { - Ok(child.clone().with_range(unbounded)) + let source_type = child.range.data_type(); + // A widening cast is additionally one-to-one, so it is strictly + // order-preserving; a narrowing cast may collapse distinct values, + // breaking the ordering of subsequent sort keys. + let bigger_cast = CastExpr::check_bigger_cast(target_type, &source_type); + if is_order_preserving_cast_family(&source_type, target_type) || bigger_cast { + Ok(child + .clone() + .with_range(unbounded) + .with_strictly_order_preserving( + child.strictly_order_preserving && bigger_cast, + )) } else { Ok(ExprProperties::new_unknown().with_range(unbounded)) } @@ -298,6 +305,61 @@ impl PhysicalExpr for CastExpr { write!(f, ")") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new( + protobuf::PhysicalCastNode { + expr: Some(Box::new(ctx.encode_child(self.expr())?)), + arrow_type: Some(self.cast_type().try_into()?), + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl CastExpr { + /// Reconstruct a [`CastExpr`] from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`] so the decode signature matches + /// other migrated expressions and can inspect outer-node metadata if + /// needed in the future. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::internal_datafusion_err; + use datafusion_common::internal_err; + use datafusion_proto_models::protobuf; + + let cast_expr = match &node.expr_type { + Some(protobuf::physical_expr_node::ExprType::Cast(cast_expr)) => { + cast_expr.as_ref() + } + _ => return internal_err!("PhysicalExprNode is not a CastExpr"), + }; + + let expr = ctx.decode_required_expression( + cast_expr.expr.as_deref(), + "CastExpr", + "expr", + )?; + let arrow_type = cast_expr.arrow_type.as_ref().ok_or_else(|| { + internal_datafusion_err!("CastExpr is missing required field 'arrow_type'") + })?; + + Ok(Arc::new(CastExpr::new(expr, arrow_type.try_into()?, None))) + } } /// Return a PhysicalExpression representing `expr` casted to @@ -1153,4 +1215,190 @@ mod tests { Ok(()) } + + #[test] + fn test_check_bigger_cast_precision_loss() { + use DataType::*; + + // Exact conversions without precision loss + assert!(CastExpr::check_bigger_cast(&Int16, &Int8)); + assert!(CastExpr::check_bigger_cast(&Int64, &Int32)); + assert!(CastExpr::check_bigger_cast(&Float32, &Int16)); + assert!(CastExpr::check_bigger_cast(&Float32, &UInt16)); + assert!(CastExpr::check_bigger_cast(&Float64, &Int32)); + assert!(CastExpr::check_bigger_cast(&Float64, &UInt32)); + assert!(CastExpr::check_bigger_cast(&LargeUtf8, &Utf8)); + + // Precision-losing int-to-float conversions should return false + assert!(!CastExpr::check_bigger_cast(&Float32, &Int32)); + assert!(!CastExpr::check_bigger_cast(&Float32, &UInt32)); + assert!(!CastExpr::check_bigger_cast(&Float64, &Int64)); + assert!(!CastExpr::check_bigger_cast(&Float64, &UInt64)); + + // Signed <-> Unsigned conversions should return false (not order-preserving due to negative values) + assert!(!CastExpr::check_bigger_cast(&UInt16, &Int8)); + assert!(!CastExpr::check_bigger_cast(&UInt32, &Int16)); + assert!(!CastExpr::check_bigger_cast(&Int16, &UInt8)); + } +} + +/// Tests for the `try_to_proto` / `try_from_proto` hooks. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::datafusion_common::ArrowType; + use datafusion_proto_models::protobuf::{ + PhysicalCastNode, PhysicalExprNode, physical_expr_node, + }; + + /// A `CastExpr` over an `Int32` column, casting to `Int64`. + fn proto_cast_fixture() -> CastExpr { + let schema = Schema::new(vec![Field::new("a", Int32, false)]); + CastExpr::new(col("a", &schema).unwrap(), Int64, None) + } + + fn proto_int64_arrow_type() -> ArrowType { + (&Int64).try_into().unwrap() + } + + /// Build a `CastExpr` proto node with the given child and target type. + fn proto_cast_node( + expr: Option>, + arrow_type: Option, + ) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::Cast(Box::new( + PhysicalCastNode { expr, arrow_type }, + ))), + } + } + + #[test] + fn try_to_proto_encodes_cast_expr() { + let cast = proto_cast_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = cast + .try_to_proto(&ctx) + .unwrap() + .expect("CastExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let cast_node = match node.expr_type { + Some(physical_expr_node::ExprType::Cast(cast_node)) => *cast_node, + other => panic!("expected a Cast node, got {other:?}"), + }; + assert!(cast_node.expr.is_some()); + + let arrow_type = cast_node + .arrow_type + .as_ref() + .expect("cast type should be encoded"); + let data_type: DataType = arrow_type.try_into().unwrap(); + assert_eq!(data_type, Int64); + } + + #[test] + fn try_to_proto_propagates_child_encode_error() { + let cast = proto_cast_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let err = cast.try_to_proto(&ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("call 1") + )); + } + + #[test] + fn try_from_proto_decodes_cast_expr() { + let node = proto_cast_node( + Some(Box::new(column_node("a"))), + Some(proto_int64_arrow_type()), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = CastExpr::try_from_proto(&node, &ctx).unwrap(); + let cast = decoded + .downcast_ref::() + .expect("decoded expr should be a CastExpr"); + + assert_eq!(cast.cast_type(), &Int64); + assert!(cast.expr().downcast_ref::().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_cast_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) + if msg.contains("PhysicalExprNode is not a CastExpr") + )); + } + + #[test] + fn try_from_proto_rejects_missing_expr() { + let node = proto_cast_node(None, Some(proto_int64_arrow_type())); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) + if msg.contains("CastExpr is missing required field 'expr'") + )); + } + + #[test] + fn try_from_proto_rejects_missing_arrow_type() { + let node = proto_cast_node(Some(Box::new(column_node("a"))), None); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) + if msg.contains("CastExpr is missing required field 'arrow_type'") + )); + } + + #[test] + fn try_from_proto_propagates_child_decode_error() { + let node = proto_cast_node( + Some(Box::new(column_node("a"))), + Some(proto_int64_arrow_type()), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = CastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("call 1") + )); + } } diff --git a/datafusion/physical-expr/src/expressions/column.rs b/datafusion/physical-expr/src/expressions/column.rs index 7d4b0e7e2f396..482ab6ef1e787 100644 --- a/datafusion/physical-expr/src/expressions/column.rs +++ b/datafusion/physical-expr/src/expressions/column.rs @@ -146,6 +146,63 @@ impl PhysicalExpr for Column { fn placement(&self) -> ExpressionPlacement { ExpressionPlacement::Column } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Column(self.into())), + })) + } +} + +#[cfg(feature = "proto")] +impl From<&datafusion_proto_models::protobuf::PhysicalColumn> for Column { + fn from(c: &datafusion_proto_models::protobuf::PhysicalColumn) -> Self { + Column::new(&c.name, c.index as usize) + } +} + +#[cfg(feature = "proto")] +impl From<&Column> for datafusion_proto_models::protobuf::PhysicalColumn { + fn from(c: &Column) -> Self { + Self { + name: c.name.clone(), + index: c.index as u32, + } + } +} + +#[cfg(feature = "proto")] +impl Column { + /// Reconstruct a [`Column`] from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`] — the exact inverse of what + /// [`PhysicalExpr::try_to_proto`] produces — so every expression's + /// `try_from_proto` shares one signature. The decode context is currently + /// unused, but is threaded through so that future expressions with child + /// sub-expressions can recurse via [`PhysicalExprDecodeCtx::decode`]. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + /// [`PhysicalExpr::try_to_proto`]: datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto + /// [`PhysicalExprDecodeCtx::decode`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::decode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + let column = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Column, + "Column", + ); + Ok(Arc::new(Column::from(column))) + } } impl Column { diff --git a/datafusion/physical-expr/src/expressions/dynamic_filters.rs b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs similarity index 63% rename from datafusion/physical-expr/src/expressions/dynamic_filters.rs rename to datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs index 5b9de882160aa..eb3d457de82ad 100644 --- a/datafusion/physical-expr/src/expressions/dynamic_filters.rs +++ b/datafusion/physical-expr/src/expressions/dynamic_filters/mod.rs @@ -22,13 +22,23 @@ use tokio::sync::watch; use crate::PhysicalExpr; use arrow::datatypes::{DataType, Schema}; +#[cfg(feature = "proto")] +use datafusion_common::internal_datafusion_err; use datafusion_common::{ Result, tree_node::{Transformed, TransformedResult, TreeNode}, }; + use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::physical_expr::DynHash; +mod tracker; +pub use tracker::{DynamicFilterTracker, DynamicFilterTracking}; + +/// Per-generation cache of the remapped current expression for +/// [`DynamicFilterPhysicalExpr::current`]. See the field docs there. +type CurrentExprCache = Arc)>>>; + /// State of a dynamic filter, tracking both updates and completion. #[derive(Debug, Clone, Copy, PartialEq, Eq)] enum FilterState { @@ -56,7 +66,6 @@ impl FilterState { /// For more background, please also see the [Dynamic Filters: Passing Information Between Operators During Execution for 25x Faster Queries blog] /// /// [Dynamic Filters: Passing Information Between Operators During Execution for 25x Faster Queries blog]: https://datafusion.apache.org/blog/2025/09/10/dynamic-filters -#[derive(Debug)] pub struct DynamicFilterPhysicalExpr { /// The original children of this PhysicalExpr, if any. /// This is necessary because the dynamic filter may be initialized with a placeholder (e.g. `lit(true)`) @@ -66,6 +75,16 @@ pub struct DynamicFilterPhysicalExpr { /// If any of the children were remapped / modified (e.g. to adjust for projections) we need to keep track of the new children /// so that when we update `current()` in subsequent iterations we can re-apply the replacements. remapped_children: Option>>, + /// Cache of the last (generation, remapped-expression) pair returned by + /// [`Self::current`]. `current()` is hot on the per-batch RowFilter path; + /// when the inner generation hasn't changed (common — updates fire once + /// per HashJoin build or once per TopK threshold refresh, but `evaluate` + /// is called per batch), the cache serves the remapped expression + /// without re-running the `transform_up` tree walk in + /// [`Self::remap_children`]. Reset on `update()` (by generation bump) + /// and populated with `None` on `with_new_children` (each derived + /// filter owns its own cache). + current_cache: CurrentExprCache, /// The source of dynamic filters. inner: Arc>, /// Broadcasts filter state (updates and completion) to all waiters. @@ -77,27 +96,40 @@ pub struct DynamicFilterPhysicalExpr { nullable: Arc>>, } +impl std::fmt::Debug for DynamicFilterPhysicalExpr { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + // Manual impl deliberately omits `current_cache`: it is a pure + // optimization artifact whose contents depend on whether + // `current()` has been called, and roundtrip tests (e.g. in + // `datafusion-proto`) compare `format!("{:?}", ..)` output. + f.debug_struct("DynamicFilterPhysicalExpr") + .field("children", &self.children) + .field("remapped_children", &self.remapped_children) + .field("inner", &self.inner) + .field("state_watch", &self.state_watch) + .field("data_type", &self.data_type) + .field("nullable", &self.nullable) + .finish() + } +} + /// Atomic internal state of a [`DynamicFilterPhysicalExpr`]. /// /// `expression_id` lives here because it identifies the actual filter expression `expr`. /// Derived `DynamicFilterPhysicalExpr`s (e.g. via [`PhysicalExpr::with_new_children`]) are /// the same logical filter and must report the same `expression_id`. -/// -/// **Warning:** exposed publicly solely so that proto (de)serialization in -/// `datafusion-proto` can read and rebuild this state. Do not treat this type -/// or its layout as a stable API. #[derive(Clone, Debug)] -pub struct Inner { +struct Inner { /// A unique identifier for the expression. - pub expression_id: u64, + expression_id: u64, /// A counter that gets incremented every time the expression is updated so that we can track changes cheaply. /// This is used for [`PhysicalExpr::snapshot_generation`] to have a cheap check for changes. - pub generation: u64, - pub expr: Arc, + generation: u64, + expr: Arc, /// Flag for quick synchronous check if filter is complete. /// This is redundant with the watch channel state, but allows us to return immediately /// from `wait_complete()` without subscribing if already complete. - pub is_complete: bool, + is_complete: bool, } impl Inner { @@ -187,6 +219,7 @@ impl DynamicFilterPhysicalExpr { children, remapped_children: None, // Initially no remapped children inner: Arc::new(RwLock::new(Inner::new(inner))), + current_cache: Arc::new(RwLock::new(None)), state_watch, data_type: Arc::new(RwLock::new(None)), nullable: Arc::new(RwLock::new(None)), @@ -230,9 +263,48 @@ impl DynamicFilterPhysicalExpr { /// Get the current expression. /// This will return the current expression with any children /// remapped to match calls to [`PhysicalExpr::with_new_children`]. + /// + /// Called per batch on the RowFilter path (via + /// [`PhysicalExpr::evaluate`]). The remap walk is O(tree size) and, for + /// dynamic filters that carry a large `InListExpr` (join key IN list), + /// dominated by `InListExpr::with_new_children` cloning the whole list. + /// The inner generation only changes when [`Self::update`] fires, so we + /// cache the remapped expression per generation and return it directly + /// on subsequent per-batch calls. pub fn current(&self) -> Result> { - let expr = Arc::clone(self.inner.read().expr()); - Self::remap_children(&self.children, self.remapped_children.as_ref(), expr) + // Fast path: cache hit for the current generation. + let (expr, generation) = { + let inner = self.inner.read(); + (Arc::clone(inner.expr()), inner.generation) + }; + if let Some((cached_gen, cached_expr)) = self.current_cache.read().as_ref() + && *cached_gen == generation + { + return Ok(Arc::clone(cached_expr)); + } + // Slow path: (re)compute the remap and store it under a write lock. + let remapped = + Self::remap_children(&self.children, self.remapped_children.as_ref(), expr)?; + // Only publish our result if it is strictly newer than whatever is + // currently cached. Without this guard a slow computation that + // observed an older `inner` could clobber a newer entry that a + // concurrent caller has already published (see #23532 review), which + // would force subsequent readers to redo the remap for the newer + // generation. Same-generation writes are also skipped: the cached + // and about-to-write remaps are semantically identical (same input + // expression, same remapped_children), so overwriting is redundant + // and only wastes a write-lock take. + { + let mut cache = self.current_cache.write(); + let should_write = match cache.as_ref() { + Some((cached_gen, _)) => generation > *cached_gen, + None => true, + }; + if should_write { + *cache = Some((generation, Arc::clone(&remapped))); + } + } + Ok(remapped) } /// Update the current expression and notify all waiters. @@ -326,6 +398,31 @@ impl DynamicFilterPhysicalExpr { .await; } + /// Returns `true` if this filter has been marked complete via + /// [`Self::mark_complete`] and will therefore never change again. + pub(crate) fn is_complete(&self) -> bool { + self.inner.read().is_complete + } + + /// Subscribe to this filter's updates for cheap, synchronous change + /// detection. + /// + /// The returned [`DynamicFilterSubscription`] lets a consumer poll whether + /// the filter's expression has advanced since it last looked, without + /// re-walking a predicate tree or re-deriving a generation on every check. + /// This is the building block used by [`DynamicFilterTracker`] to watch + /// every dynamic filter inside a (possibly composite) predicate. + pub(crate) fn subscribe(&self) -> DynamicFilterSubscription { + let mut receiver = self.state_watch.subscribe(); + // Mark the current state as already-seen so the first `observe()` only + // reports updates that happen *after* subscription. + let last_generation = receiver.borrow_and_update().generation(); + DynamicFilterSubscription { + receiver, + last_generation, + } + } + /// Check if this dynamic filter is being actively used by any consumers. /// /// Returns `true` if there are references beyond the producer (e.g., the HashJoinExec @@ -337,6 +434,10 @@ impl DynamicFilterPhysicalExpr { /// We check both Arc counts to handle two cases: /// - Transformed filters (via `with_new_children`) share the inner Arc (inner count > 1) /// - Direct clones (via `Arc::clone`) increment the outer count (outer count > 1) + #[deprecated( + since = "55.0.0", + note = "Traverse ExecutionPlan::apply_expressions and compare PhysicalExpr::expression_id instead" + )] pub fn is_used(self: &Arc) -> bool { // Strong count > 1 means at least one consumer is holding a reference beyond the producer. Arc::strong_count(self) > 1 || Arc::strong_count(&self.inner) > 1 @@ -362,29 +463,10 @@ impl DynamicFilterPhysicalExpr { write!(f, " ]") } - /// Return the filter's original children (before any remapping). - /// - /// **Warning:** intended only for `datafusion-proto` (de)serialization. - /// Not a stable API. - pub fn original_children(&self) -> &[Arc] { - &self.children - } - - /// Return the filter's remapped children, if any have been set via - /// [`PhysicalExpr::with_new_children`]. - /// - /// **Warning:** intended only for `datafusion-proto` (de)serialization. - /// Not a stable API. - pub fn remapped_children(&self) -> Option<&[Arc]> { - self.remapped_children.as_deref() - } - /// Rebuild a `DynamicFilterPhysicalExpr` from its stored parts. Used by /// proto deserialization. - /// - /// **Warning:** intended only for `datafusion-proto` (de)serialization. - /// Not a stable API. - pub fn from_parts( + #[cfg(any(test, feature = "proto"))] + fn from_parts( children: Vec>, remapped_children: Option>>, inner: Inner, @@ -404,19 +486,12 @@ impl DynamicFilterPhysicalExpr { children, remapped_children, inner: Arc::new(RwLock::new(inner)), + current_cache: Arc::new(RwLock::new(None)), state_watch, data_type: Arc::new(RwLock::new(None)), nullable: Arc::new(RwLock::new(None)), } } - - /// Return a clone of the atomically-captured `Inner` state. - /// - /// **Warning:** intended only for `datafusion-proto` (de)serialization. - /// Not a stable API. - pub fn inner(&self) -> Inner { - self.inner.read().clone() - } } impl PhysicalExpr for DynamicFilterPhysicalExpr { @@ -437,6 +512,9 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr { remapped_children: Some(children), // Note: expression_id is preserved inner: Arc::clone(&self.inner), + // Fresh cache per derived filter — remap depends on this + // instance's `remapped_children`, which just changed. + current_cache: Arc::new(RwLock::new(None)), state_watch: self.state_watch.clone(), data_type: Arc::clone(&self.data_type), nullable: Arc::clone(&self.nullable), @@ -520,6 +598,182 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr { fn expression_id(&self) -> Option { Some(self.inner.read().expression_id) } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use datafusion_proto_models::protobuf::physical_expr_node::ExprType; + + let children = self + .children + .iter() + .map(|c| ctx.encode_child(c)) + .collect::>>()?; + + let remapped_children = match &self.remapped_children { + Some(remapped) => remapped + .iter() + .map(|c| ctx.encode_child(c)) + .collect::>>()?, + None => vec![], + }; + + let inner = self.inner.read().clone(); + let inner_expr = Box::new(ctx.encode_child(&inner.expr)?); + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: Some(inner.expression_id), + expr_type: Some(ExprType::DynamicFilter(Box::new( + protobuf::PhysicalDynamicFilterNode { + children, + remapped_children, + generation: inner.generation, + inner_expr: Some(inner_expr), + is_complete: inner.is_complete, + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl DynamicFilterPhysicalExpr { + /// Reconstruct a [`DynamicFilterPhysicalExpr`] from a proto node. + /// + /// Called by the `ExprType::DynamicFilter` arm in `datafusion-proto`'s + /// `parse_physical_expr_with_converter`. Follows the same + /// `PhysicalExprDecodeCtx`-based pattern used by `Column`, `BinaryExpr`, etc. + pub fn try_from_proto( + proto: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf::physical_expr_node::ExprType; + + let ExprType::DynamicFilter(df) = proto.expr_type.as_ref().ok_or_else(|| { + internal_datafusion_err!("Missing expr_type in PhysicalExprNode") + })? + else { + return Err(internal_datafusion_err!("Expected DynamicFilter expr_type")); + }; + + // Decode original children + let children = df + .children + .iter() + .map(|c| ctx.decode(c)) + .collect::>>()?; + + // Decode remapped children (empty vec means None) + let remapped_children = if df.remapped_children.is_empty() { + None + } else { + Some( + df.remapped_children + .iter() + .map(|c| ctx.decode(c)) + .collect::>>()?, + ) + }; + + // Decode the inner expression + let inner_expr_proto = df.inner_expr.as_ref().ok_or_else(|| { + internal_datafusion_err!("Missing inner_expr in PhysicalDynamicFilterNode") + })?; + let inner_expr = ctx.decode(inner_expr_proto)?; + + // Restore the expression_id from the outer PhysicalExprNode + let expression_id = proto.expr_id.ok_or_else(|| { + internal_datafusion_err!( + "Missing expr_id in PhysicalExprNode for DynamicFilter" + ) + })?; + + let inner = Inner { + expression_id, + generation: df.generation, + expr: inner_expr, + is_complete: df.is_complete, + }; + + Ok(Arc::new(Self::from_parts( + children, + remapped_children, + inner, + ))) + } +} + +/// The result of polling a [`DynamicFilterSubscription`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) struct DynamicFilterChange { + /// The filter's expression advanced since the previous observation. + pub(crate) changed: bool, + /// The filter has been marked complete; it will never change again and the + /// subscription can be dropped. + pub(crate) complete: bool, +} + +/// A cheap, synchronous handle for observing updates to a single +/// [`DynamicFilterPhysicalExpr`]. +/// +/// Obtained via [`DynamicFilterPhysicalExpr::subscribe`]. Steady-state polling +/// via [`Self::observe`] is a single atomic load (the underlying +/// [`tokio::sync::watch`] version counter); the lock is only taken when the +/// filter has actually been updated. +#[derive(Debug)] +pub(crate) struct DynamicFilterSubscription { + receiver: watch::Receiver, + /// Last generation we reported as "seen". Used to distinguish a real + /// expression update from a bare [`DynamicFilterPhysicalExpr::mark_complete`] + /// (which re-broadcasts the current generation without changing the + /// expression). + last_generation: u64, +} + +impl DynamicFilterSubscription { + /// Observe the latest state of the filter. + /// + /// Reports whether the filter's expression advanced since the previous call + /// and whether it has since been marked complete. Cheap when nothing has + /// changed: a single atomic comparison with no lock acquisition. + pub(crate) fn observe(&mut self) -> DynamicFilterChange { + match self.receiver.has_changed() { + Ok(true) => { + let state = *self.receiver.borrow_and_update(); + let changed = state.generation() > self.last_generation; + if changed { + self.last_generation = state.generation(); + } + DynamicFilterChange { + changed, + complete: matches!(state, FilterState::Complete { .. }), + } + } + Ok(false) => DynamicFilterChange { + changed: false, + complete: false, + }, + // The watch sender lives inside the predicate's + // `DynamicFilterPhysicalExpr`, which the owner of this subscription + // keeps alive, so observing a dropped sender signals a bug rather + // than normal completion. Flag it loudly in debug builds; in release + // degrade to "complete" (no further updates are possible) instead of + // silently masking it. + Err(_) => { + debug_assert!( + false, + "DynamicFilterSubscription observed a dropped watch sender; \ + the owning predicate should keep it alive" + ); + DynamicFilterChange { + changed: false, + complete: true, + } + } + } + } } /// An atomic counter used to generate monotonic u64 ids. @@ -546,6 +800,25 @@ impl ExpressionIdAtomicCounter { /// file and be made public for other expressions to use. static EXPR_ID_SOURCE: ExpressionIdAtomicCounter = ExpressionIdAtomicCounter::new(); +#[cfg(test)] +impl DynamicFilterPhysicalExpr { + /// Test-only clone that produces a fresh outer instance sharing the + /// same `inner`. Used by the concurrent stress test to obtain a + /// standalone `Arc` without going through `with_new_children` + /// (which would clear `remapped_children`). + fn clone_with_remapped_children_for_test(&self) -> Self { + Self { + children: self.children.clone(), + remapped_children: self.remapped_children.clone(), + inner: Arc::clone(&self.inner), + current_cache: Arc::new(RwLock::new(None)), + state_watch: self.state_watch.clone(), + data_type: Arc::clone(&self.data_type), + nullable: Arc::clone(&self.nullable), + } + } +} + #[cfg(test)] mod test { use crate::{ @@ -822,6 +1095,10 @@ mod test { } #[test] + #[expect( + deprecated, + reason = "covers the deprecated API during its retention period" + )] fn test_is_used() { let filter = Arc::new(DynamicFilterPhysicalExpr::new( vec![], @@ -1003,22 +1280,19 @@ mod test { // Capture the parts and reconstruct. `expression_id` rides in `inner`. let reconstructed = DynamicFilterPhysicalExpr::from_parts( - reassigned.original_children().to_vec(), - reassigned.remapped_children().map(|r| r.to_vec()), - reassigned.inner(), + reassigned.children.to_vec(), + reassigned.remapped_children.as_ref().map(|r| r.to_vec()), + reassigned.inner.read().clone(), ); + assert_eq!(reassigned.children, reconstructed.children); assert_eq!( - reassigned.original_children(), - reconstructed.original_children(), - ); - assert_eq!( - reassigned.remapped_children(), - reconstructed.remapped_children(), + reassigned.remapped_children, + reconstructed.remapped_children, ); assert_eq!(reassigned.expression_id(), reconstructed.expression_id()); - let r = reassigned.inner(); - let c = reconstructed.inner(); + let r = reassigned.inner.read().clone(); + let c = reconstructed.inner.read().clone(); assert_eq!(r.generation, c.generation); assert_eq!(r.is_complete, c.is_complete); assert_eq!(format!("{:?}", r.expr), format!("{:?}", c.expr)); @@ -1087,4 +1361,238 @@ mod test { "mark_complete() must not change expression_id", ); } + + /// Repeated `current()` at the same generation must return the exact same + /// `Arc` — the cache serves without re-running `remap_children`. + #[test] + fn test_current_cache_hits_within_generation() { + let table_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let col_a = col("a", &table_schema).unwrap(); + let expr = Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + datafusion_expr::Operator::Gt, + lit(10) as Arc, + )); + // Force the remap path to actually run: give the filter a distinct + // `remapped_children`. Without this, `remap_children` returns the + // input Arc unchanged and every call is trivially pointer-equal. + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + expr as Arc, + )); + let remapped_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let derived = reassign_expr_columns( + Arc::clone(&filter) as Arc, + &remapped_schema, + ) + .unwrap(); + let derived = derived + .downcast_ref::() + .expect("derived filter must be a DynamicFilterPhysicalExpr"); + + // First call populates the cache. Second and third must return the + // *same* Arc — proving `remap_children` did not run again. + let first = derived.current().unwrap(); + let second = derived.current().unwrap(); + let third = derived.current().unwrap(); + assert!( + Arc::ptr_eq(&first, &second), + "current() should return the cached Arc within a generation", + ); + assert!(Arc::ptr_eq(&second, &third)); + } + + /// `update()` bumps the generation; the next `current()` must return a + /// fresh remapped expression, not the stale cached one. + #[test] + fn test_current_cache_invalidates_on_update() { + let table_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let col_a = col("a", &table_schema).unwrap(); + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + lit(10) as Arc, + )); + // Remap to force the cache path. + let derived = reassign_expr_columns( + Arc::clone(&filter) as Arc, + &table_schema, + ) + .unwrap(); + let derived = derived + .downcast_ref::() + .expect("derived filter must be a DynamicFilterPhysicalExpr"); + + let before = derived.current().unwrap(); + // Bump the generation with a distinct expression. + filter + .update(Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + datafusion_expr::Operator::Gt, + lit(42) as Arc, + )) as Arc) + .unwrap(); + let after = derived.current().unwrap(); + assert!( + !Arc::ptr_eq(&before, &after), + "current() must return a fresh Arc after update() bumps the generation", + ); + assert_ne!(format!("{before:?}"), format!("{after:?}")); + } + + /// `with_new_children` produces a derived filter with its own cache slot; + /// populating one filter's cache must not leak into the other. + #[test] + fn test_current_cache_is_per_derived_filter() { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ])); + // Original expression references `a`. Each derived filter remaps `a` + // to a *different* column so remap_children returns distinct exprs + // per filter (and thus distinct cached Arcs). + let col_a = col("a", &schema).unwrap(); + let expr = Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + datafusion_expr::Operator::Gt, + lit(10) as Arc, + )); + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + expr as Arc, + )); + + let d1 = Arc::clone(&filter) + .with_new_children(vec![col("b", &schema).unwrap()]) + .unwrap(); + let d2 = Arc::clone(&filter) + .with_new_children(vec![col("c", &schema).unwrap()]) + .unwrap(); + let d1 = d1.downcast_ref::().unwrap(); + let d2 = d2.downcast_ref::().unwrap(); + + let d1_first = d1.current().unwrap(); + let d2_first = d2.current().unwrap(); + // Distinct remap_children paths produce distinct cached Arcs. + assert!(!Arc::ptr_eq(&d1_first, &d2_first)); + assert_ne!(format!("{d1_first:?}"), format!("{d2_first:?}")); + // Subsequent calls each hit their own cache. + assert!(Arc::ptr_eq(&d1_first, &d1.current().unwrap())); + assert!(Arc::ptr_eq(&d2_first, &d2.current().unwrap())); + } + + /// Stress-test the cache under concurrent readers and periodic writes. + /// + /// Motivation: prod scans run with tens/hundreds of partitions, each + /// calling `current()` on the same `Arc` per + /// batch, while the producer (HashJoin build / TopK) fires `update()` + /// on a separate task. A caching bug that only shows up under + /// contention (torn read, ABA-style Arc lifetime issue, cache monotonicity + /// violation) would be invisible in single-threaded tests. This test + /// hot-loops many readers against a writer and asserts the invariants + /// that matter: + /// 1. `current()` never panics and always returns a valid `Arc`. + /// 2. Cache generation never regresses (monotonic non-decreasing). + /// 3. After the writer stops, the cache eventually converges to the + /// final `inner.generation`. + #[test] + fn test_current_cache_concurrent_readers_and_writer() { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::thread; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let col_a = col("a", &schema).unwrap(); + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + lit(true) as Arc, + )); + // Force the remap path: give the derived filter distinct + // remapped_children so `current()` actually runs `remap_children` + // instead of the short-circuit `Arc::clone(&expr)`. + let derived = + reassign_expr_columns(Arc::clone(&filter) as Arc, &schema) + .unwrap(); + // Re-wrap in Arc for cross-thread sharing. + let derived: Arc = Arc::new( + derived + .downcast_ref::() + .expect("derived is DynamicFilterPhysicalExpr") + .clone_with_remapped_children_for_test(), + ); + + let stop = Arc::new(AtomicBool::new(false)); + const READERS: usize = 8; + const READER_ITERS: usize = 5_000; + const WRITER_ITERS: i32 = 200; + + let mut readers = Vec::with_capacity(READERS); + for _ in 0..READERS { + let d = Arc::clone(&derived); + let stop = Arc::clone(&stop); + readers.push(thread::spawn(move || { + let mut last_seen_gen: u64 = 0; + for _ in 0..READER_ITERS { + if stop.load(Ordering::Relaxed) { + break; + } + let expr = d.current().expect("current must not fail"); + // Cheap sanity: the returned Arc's Debug must be + // formattable — proves it's a valid PhysicalExpr. + let _ = format!("{expr:?}"); + // Cache generation observed by this reader must never + // decrease across successive calls on the same filter. + let cached = d + .current_cache + .read() + .as_ref() + .map(|(g, _)| *g) + .unwrap_or(0); + assert!( + cached >= last_seen_gen, + "cache generation regressed: {cached} < {last_seen_gen}", + ); + last_seen_gen = cached; + } + })); + } + + let f_writer = Arc::clone(&filter); + let writer = thread::spawn(move || { + for i in 0..WRITER_ITERS { + f_writer + .update(lit(i) as Arc) + .expect("update must succeed"); + // Yield to give readers a chance to see intermediate states. + thread::yield_now(); + } + }); + + writer.join().expect("writer thread panicked"); + stop.store(true, Ordering::Relaxed); + for h in readers { + h.join().expect("reader thread panicked"); + } + + // After the writer is done, one final `current()` should sync the + // cache to the latest generation. + let _ = derived.current().unwrap(); + let (cache_gen, _) = derived + .current_cache + .read() + .as_ref() + .expect("cache populated after final current()") + .clone(); + let inner_gen = derived.inner.read().generation; + assert_eq!( + cache_gen, inner_gen, + "final cache generation must match inner.generation", + ); + // Writer bumps generation once per update, so final generation is + // starting-generation + WRITER_ITERS. Starting is 1, so final is + // WRITER_ITERS + 1. + assert_eq!(inner_gen, WRITER_ITERS as u64 + 1); + } } diff --git a/datafusion/physical-expr/src/expressions/dynamic_filters/tracker.rs b/datafusion/physical-expr/src/expressions/dynamic_filters/tracker.rs new file mode 100644 index 0000000000000..fd4c18b07e2cd --- /dev/null +++ b/datafusion/physical-expr/src/expressions/dynamic_filters/tracker.rs @@ -0,0 +1,331 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Tracking changes to the dynamic filters inside a predicate. +//! +//! Several operators (Parquet file/row-group pruning, remote execution, ...) +//! hold a predicate that *may* contain one or more +//! [`DynamicFilterPhysicalExpr`] nodes which are updated during execution +//! (e.g. a `TopK` tightening its threshold, or a `HashJoinExec` publishing the +//! build-side bounds). These consumers repeatedly ask two questions: +//! +//! 1. *"Does this predicate contain anything that can still change?"* — to +//! decide whether it is worth setting up runtime re-pruning at all. +//! 2. *"Has it changed since I last looked?"* — to decide whether to rebuild an +//! expensive derived artifact (e.g. a `PruningPredicate`). +//! +//! Historically each call site answered these by recursively folding +//! [`PhysicalExpr::snapshot_generation`] over the whole tree on *every* check +//! and diffing the resulting `u64`. [`DynamicFilterTracker`] replaces that with +//! a single up-front walk that subscribes to each still-incomplete dynamic +//! filter; subsequent checks only poll the (shrinking) set of subscriptions, +//! each of which is a cheap atomic load in the common "nothing changed" case. +//! +//! [`PhysicalExpr::snapshot_generation`]: crate::PhysicalExpr::snapshot_generation + +use std::sync::Arc; + +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + +use crate::PhysicalExpr; + +use super::{DynamicFilterPhysicalExpr, DynamicFilterSubscription}; + +/// Classification of a predicate according to the dynamic filters it contains. +/// +/// Produced by [`DynamicFilterTracking::classify`] with a single tree walk so +/// callers can answer both "is it worth pruning at all?" and "do I need to keep +/// watching?" without traversing the predicate twice. +#[derive(Debug)] +pub enum DynamicFilterTracking { + /// The predicate contains no [`DynamicFilterPhysicalExpr`] at all. It is + /// fully static and will never change. + Static, + /// The predicate contains one or more dynamic filters, but all of them have + /// already been marked complete. Their *current* values may differ from + /// what was known at planning time (so a one-shot prune is still + /// worthwhile), but they will not change again — there is nothing to watch. + AllComplete, + /// The predicate contains at least one dynamic filter that can still change. + /// The embedded [`DynamicFilterTracker`] should be polled to detect updates. + Watching(DynamicFilterTracker), +} + +impl DynamicFilterTracking { + /// Walk `predicate` once and classify its dynamic-filter content, + /// subscribing to every filter that is not yet complete. + pub fn classify(predicate: &Arc) -> Self { + let mut subscriptions = Vec::new(); + let mut found_any = false; + predicate + .apply(|expr| { + if let Some(filter) = expr.downcast_ref::() { + found_any = true; + // Already-complete filters can never change again, so there + // is no point subscribing to them. + if !filter.is_complete() { + subscriptions.push(filter.subscribe()); + } + } + Ok(TreeNodeRecursion::Continue) + }) + .expect("traversal closure is infallible"); + + if !found_any { + DynamicFilterTracking::Static + } else if subscriptions.is_empty() { + DynamicFilterTracking::AllComplete + } else { + DynamicFilterTracking::Watching(DynamicFilterTracker { subscriptions }) + } + } + + /// `true` if the predicate contains any dynamic filter (complete or not), + /// i.e. its value may differ from what was known at planning time and is + /// therefore worth re-evaluating at least once. + pub fn contains_dynamic_filter(&self) -> bool { + !matches!(self, DynamicFilterTracking::Static) + } + + /// Mutable access to the underlying tracker when there is still something to + /// watch. + pub fn watcher(&mut self) -> Option<&mut DynamicFilterTracker> { + match self { + DynamicFilterTracking::Watching(tracker) => Some(tracker), + _ => None, + } + } +} + +/// Watches every still-incomplete [`DynamicFilterPhysicalExpr`] reachable from a +/// predicate and reports, cheaply, whether any of them has been updated since +/// the last check. +/// +/// Obtain one from [`DynamicFilterTracking::classify`] via +/// [`DynamicFilterTracking::watcher`]; the `Watching` variant carries it only +/// when there is at least one dynamic filter that can still change. +#[derive(Debug)] +pub struct DynamicFilterTracker { + /// Subscriptions to the not-yet-complete dynamic filters. Entries are + /// dropped as their filters complete, so the set only shrinks. + subscriptions: Vec, +} + +impl DynamicFilterTracker { + /// Returns `true` if any watched filter's expression has advanced since the + /// previous call. + /// + /// Filters that have completed are dropped from the watch set as they are + /// observed; once every filter has completed this is a no-op that always + /// returns `false`. + pub fn changed(&mut self) -> bool { + let mut changed = false; + self.subscriptions.retain_mut(|subscription| { + let change = subscription.observe(); + changed |= change.changed; + // Keep the subscription only while the filter can still change. + !change.complete + }); + changed + } +} + +#[cfg(test)] +impl DynamicFilterTracker { + /// Build a tracker directly, or `None` if `predicate` has no dynamic filter + /// that can still change. Test-only; production builds a tracker via + /// [`DynamicFilterTracking::classify`] + [`DynamicFilterTracking::watcher`]. + fn try_new(predicate: &Arc) -> Option { + match DynamicFilterTracking::classify(predicate) { + DynamicFilterTracking::Watching(tracker) => Some(tracker), + DynamicFilterTracking::Static | DynamicFilterTracking::AllComplete => None, + } + } + + /// `true` once every watched filter has completed and been dropped. + fn is_exhausted(&self) -> bool { + self.subscriptions.is_empty() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use crate::expressions::{BinaryExpr, col, lit}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_expr::Operator; + + /// `col > ` where the dynamic filter starts as `lit(true)`. + fn dynamic_predicate() -> (Arc, Arc) { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let column = col("a", &schema).unwrap(); + let filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&column)], + lit(true), + )); + let predicate = Arc::new(BinaryExpr::new( + column, + Operator::Gt, + Arc::clone(&filter) as Arc, + )) as Arc; + (predicate, filter) + } + + #[test] + fn static_predicate_is_not_watched() { + let predicate = lit(true); + assert!(matches!( + DynamicFilterTracking::classify(&predicate), + DynamicFilterTracking::Static + )); + assert!(DynamicFilterTracker::try_new(&predicate).is_none()); + } + + #[test] + fn already_complete_filter_is_not_watched() { + let (predicate, filter) = dynamic_predicate(); + filter.mark_complete(); + + match DynamicFilterTracking::classify(&predicate) { + DynamicFilterTracking::AllComplete => {} + other => panic!("expected AllComplete, got {other:?}"), + } + // Still reported as dynamic (worth a one-shot prune)... + assert!(DynamicFilterTracking::classify(&predicate).contains_dynamic_filter()); + // ...but there is nothing to watch. + assert!(DynamicFilterTracker::try_new(&predicate).is_none()); + } + + #[test] + fn detects_update_exactly_once() { + let (predicate, filter) = dynamic_predicate(); + let mut tracker = DynamicFilterTracker::try_new(&predicate) + .expect("predicate has an incomplete dynamic filter"); + + // No update yet. + assert!(!tracker.changed()); + + filter.update(lit(false)).unwrap(); + // The update is reported once... + assert!(tracker.changed()); + // ...and not repeatedly. + assert!(!tracker.changed()); + } + + #[test] + fn update_before_subscribe_is_not_reported() { + let (predicate, filter) = dynamic_predicate(); + + // An update that happens *before* the tracker subscribes must not be + // reported on the first poll: `subscribe()` snapshots the current + // generation via `borrow_and_update()`, so only post-subscription + // updates count. + filter.update(lit(false)).unwrap(); + + let mut tracker = DynamicFilterTracker::try_new(&predicate) + .expect("predicate has an incomplete dynamic filter"); + assert!(!tracker.changed()); + + // A subsequent update is still reported. + filter.update(lit(true)).unwrap(); + assert!(tracker.changed()); + } + + #[test] + fn mark_complete_does_not_count_as_a_change() { + let (predicate, filter) = dynamic_predicate(); + let mut tracker = DynamicFilterTracker::try_new(&predicate).unwrap(); + + filter.update(lit(false)).unwrap(); + assert!(tracker.changed()); + + // `mark_complete()` re-broadcasts the current generation without + // changing the expression: it must not trigger a spurious rebuild. + filter.mark_complete(); + assert!(!tracker.changed()); + // The filter has completed, so the tracker drains itself. + assert!(tracker.is_exhausted()); + } + + #[test] + fn coalesced_update_then_complete_is_one_change() { + let (predicate, filter) = dynamic_predicate(); + let mut tracker = DynamicFilterTracker::try_new(&predicate).unwrap(); + + // Update and complete before the tracker gets a chance to observe. + // The watch channel only retains the latest value, so the tracker sees + // `Complete` directly; it must still report the (final) change once. + filter.update(lit(false)).unwrap(); + filter.mark_complete(); + + assert!(tracker.changed()); + assert!(tracker.is_exhausted()); + assert!(!tracker.changed()); + } + + #[test] + fn watches_multiple_filters_independently() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ]); + let col_a = col("a", &schema).unwrap(); + let col_b = col("b", &schema).unwrap(); + let filter_a = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_a)], + lit(true), + )); + let filter_b = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::clone(&col_b)], + lit(true), + )); + let predicate = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + col_a, + Operator::Gt, + Arc::clone(&filter_a) as Arc, + )), + Operator::And, + Arc::new(BinaryExpr::new( + col_b, + Operator::Lt, + Arc::clone(&filter_b) as Arc, + )), + )) as Arc; + + let mut tracker = DynamicFilterTracker::try_new(&predicate).unwrap(); + assert!(!tracker.changed()); + + filter_a.update(lit(false)).unwrap(); + assert!(tracker.changed()); + assert!(!tracker.changed()); + + filter_b.update(lit(false)).unwrap(); + assert!(tracker.changed()); + assert!(!tracker.changed()); + + // Completing one filter leaves the other still watched. + filter_a.mark_complete(); + assert!(!tracker.changed()); + assert!(!tracker.is_exhausted()); + + filter_b.mark_complete(); + assert!(!tracker.changed()); + assert!(tracker.is_exhausted()); + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list.rs b/datafusion/physical-expr/src/expressions/in_list.rs index e2251d8e63fa7..874e149b58328 100644 --- a/datafusion/physical-expr/src/expressions/in_list.rs +++ b/datafusion/physical-expr/src/expressions/in_list.rs @@ -37,7 +37,9 @@ use datafusion_common::{ use datafusion_expr::{ColumnarValue, expr_vec_fmt}; mod array_static_filter; +mod branchless_filter; mod primitive_filter; +mod result; mod static_filter; mod strategy; @@ -246,6 +248,32 @@ impl InListExpr { Ok(Self::new(expr, list, negated, static_filter)) } + + #[cfg(feature = "proto")] + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let node = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::InList, + "InList", + ); + + let expr = + ctx.decode_required_expression(node.expr.as_deref(), "InListExpr", "expr")?; + let list = ctx.decode_children_expressions(&node.list)?; + + Ok(Arc::new(InListExpr::try_new( + expr, + list, + node.negated, + ctx.schema(), + )?)) + } } impl std::fmt::Display for InListExpr { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { @@ -442,6 +470,25 @@ impl PhysicalExpr for InListExpr { } write!(f, ")") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::InList(Box::new( + protobuf::PhysicalInListNode { + expr: Some(Box::new(ctx.encode_child(&self.expr)?)), + list: ctx.encode_children_expressions(&self.list)?, + negated: self.negated, + }, + ))), + })) + } } impl PartialEq for InListExpr { @@ -2854,10 +2901,9 @@ mod tests { #[test] fn test_in_list_esoteric_types() -> Result<()> { - // Test esoteric/less common types to validate the transform and mapping flow. - // These types are reinterpreted to base primitive types (e.g., Timestamp -> UInt64, - // Interval -> Decimal128, Float16 -> UInt16). We just need to verify basic - // functionality works - no need for comprehensive null handling tests. + // Test less common types covered by IN-list evaluation. Some of these + // use specialized filters, and others fall back to the generic path; + // this keeps the end-to-end behavior covered either way. // Helper: simple IN test that expects [Some(true), Some(false)] let test_type = |data_type: DataType, @@ -2880,7 +2926,7 @@ mod tests { Ok(()) }; - // Timestamp types (all units map to Int64 -> UInt64) + // Timestamp types test_type( DataType::Timestamp(TimeUnit::Second, None), Arc::new(TimestampSecondArray::from(vec![Some(1000), Some(2000)])), @@ -2914,7 +2960,7 @@ mod tests { ], )?; - // Time32 and Time64 (map to Int32 -> UInt32 and Int64 -> UInt64 respectively) + // Time32 and Time64 test_type( DataType::Time32(TimeUnit::Second), Arc::new(Time32SecondArray::from(vec![Some(3600), Some(7200)])), @@ -2960,7 +3006,7 @@ mod tests { ], )?; - // Duration types (map to Int64 -> UInt64) + // Duration types test_type( DataType::Duration(TimeUnit::Second), Arc::new(DurationSecondArray::from(vec![Some(86400), Some(172800)])), @@ -3006,7 +3052,7 @@ mod tests { ], )?; - // Interval types (map to 16-byte Decimal128Type) + // Interval types test_type( DataType::Interval(IntervalUnit::YearMonth), Arc::new(IntervalYearMonthArray::from(vec![Some(12), Some(24)])), @@ -3068,8 +3114,7 @@ mod tests { ], )?; - // Decimal256 (maps to Decimal128Type for 16-byte width) - // Need to use with_precision_and_scale() to set the metadata + // Decimal256. Need to use with_precision_and_scale() to set the metadata. let precision = 38; let scale = 10; test_type( @@ -3479,6 +3524,7 @@ mod tests { DataType::UInt16, DataType::UInt32, DataType::UInt64, + DataType::Float16, DataType::Float32, DataType::Float64, ]; @@ -3821,3 +3867,163 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col, lit}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalInListNode, physical_expr_node, + }; + + /// Build an `InListExpr` proto node with the given children. + fn in_list_node( + expr: Option>, + list: Vec, + negated: bool, + ) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::InList(Box::new( + PhysicalInListNode { + expr, + list, + negated, + }, + ))), + } + } + + /// An `InListExpr` over a column with one literal value. + fn in_list_fixture() -> InListExpr { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + InListExpr::try_new(col("a", &schema).unwrap(), vec![lit(1)], false, &schema) + .unwrap() + } + + #[test] + fn try_to_proto_encodes_in_list() { + let in_list = in_list_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = in_list + .try_to_proto(&ctx) + .unwrap() + .expect("InListExpr should encode to Some(node)"); + + // Built-in exprs never set expr_id; only dynamic filters do. + assert!(node.expr_id.is_none()); + let in_list_node = match node.expr_type { + Some(physical_expr_node::ExprType::InList(boxed)) => *boxed, + other => panic!("expected an InList node, got {other:?}"), + }; + assert!(!in_list_node.negated); + assert!(in_list_node.expr.is_some()); + assert_eq!(in_list_node.list.len(), 1); + } + + #[test] + fn try_to_proto_propagates_expr_encode_error() { + let in_list = in_list_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = in_list.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_to_proto_propagates_list_encode_error() { + let in_list = in_list_fixture(); + // Call 1 is for `expr`, Call 2 is for the first element of `list` + let encoder = StubEncoder::failing_on(2); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = in_list.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2"))); + } + + #[test] + fn try_from_proto_decodes_in_list() { + let node = in_list_node( + Some(Box::new(column_node("a"))), + vec![column_node("b")], + true, + ); + let schema = Schema::new(vec![Field::new("decoded", DataType::Int32, true)]); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = InListExpr::try_from_proto(&node, &ctx).unwrap(); + let in_list = decoded + .downcast_ref::() + .expect("decoded expr should be an InListExpr"); + + assert!(in_list.negated()); + assert!(in_list.expr().downcast_ref::().is_some()); + assert_eq!(in_list.list().len(), 1); + } + + #[test] + fn try_from_proto_rejects_non_in_list_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = InListExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a InList") + )); + } + + #[test] + fn try_from_proto_rejects_missing_expr() { + let node = in_list_node(None, vec![column_node("b")], false); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = InListExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("InListExpr is missing required field 'expr'") + )); + } + + #[test] + fn try_from_proto_propagates_expr_decode_error() { + let node = in_list_node( + Some(Box::new(column_node("a"))), + vec![column_node("b")], + false, + ); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = InListExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_from_proto_propagates_list_decode_error() { + let node = in_list_node( + Some(Box::new(column_node("a"))), + vec![column_node("b")], + false, + ); + let schema = Schema::empty(); + // Call 1 is `expr`, Call 2 is the first element of `list` + let decoder = StubDecoder::failing_on(2); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = InListExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2"))); + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list/array_static_filter.rs b/datafusion/physical-expr/src/expressions/in_list/array_static_filter.rs index 93bfcd49600d0..75e92dbcc59b4 100644 --- a/datafusion/physical-expr/src/expressions/in_list/array_static_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/array_static_filter.rs @@ -23,11 +23,11 @@ use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::compute::{SortOptions, take}; use arrow::datatypes::DataType; use arrow::util::bit_iterator::BitIndexIterator; -use datafusion_common::HashMap; use datafusion_common::Result; use datafusion_common::hash_utils::{RandomState, with_hashes}; -use hashbrown::hash_map::RawEntryMut; +use hashbrown::HashTable; +use super::result::build_in_list_result; use super::static_filter::StaticFilter; /// Static filter for InList that stores the array and hash set for O(1) lookups @@ -35,11 +35,92 @@ use super::static_filter::StaticFilter; pub(super) struct ArrayStaticFilter { in_array: ArrayRef, state: RandomState, - /// Used to provide a lookup from value to in list index + /// Stores indices into `in_array` for O(1) lookups. + table: HashTable, +} + +impl ArrayStaticFilter { + /// Computes a [`StaticFilter`] for the provided [`Array`] if there + /// are nulls present or there are more than the configured number of + /// elements. /// - /// Note: usize::hash is not used, instead the raw entry - /// API is used to store entries w.r.t their value - map: HashMap, + /// Note: This is split into a separate function as higher-rank trait bounds currently + /// cause type inference to misbehave + pub(super) fn try_new(in_array: ArrayRef) -> Result { + // Null type has no natural order - return empty hash set + if in_array.data_type() == &DataType::Null { + return Ok(ArrayStaticFilter { + in_array, + state: RandomState::default(), + table: HashTable::new(), + }); + } + + let state = RandomState::default(); + let table = Self::build_haystack_table(&in_array, &state)?; + + Ok(Self { + in_array, + state, + table, + }) + } + + fn build_haystack_table( + haystack: &ArrayRef, + state: &RandomState, + ) -> Result> { + let mut table = HashTable::new(); + + with_hashes([haystack.as_ref()], state, |hashes| -> Result<()> { + let cmp = make_comparator(haystack, haystack, SortOptions::default())?; + + let insert_value = |idx| { + let hash = hashes[idx]; + // Only insert if not already present (deduplication) + if table.find(hash, |&x| cmp(x, idx).is_eq()).is_none() { + table.insert_unique(hash, idx, |&x| hashes[x]); + } + }; + + match haystack.nulls() { + Some(nulls) => { + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + .for_each(insert_value) + } + None => (0..haystack.len()).for_each(insert_value), + } + + Ok(()) + })?; + + Ok(table) + } + + fn find_needles_in_haystack( + &self, + needles: &dyn Array, + negated: bool, + ) -> Result { + let needle_nulls = needles.logical_nulls(); + let haystack_has_nulls = self.in_array.null_count() != 0; + + with_hashes([needles], &self.state, |needle_hashes| { + let cmp = make_comparator(needles, &self.in_array, SortOptions::default())?; + + Ok(build_in_list_result( + needles.len(), + needle_nulls.as_ref(), + haystack_has_nulls, + negated, + #[inline(always)] + |i| { + let hash = needle_hashes[i]; + self.table.find(hash, |&idx| cmp(i, idx).is_eq()).is_some() + }, + )) + }) + } } impl StaticFilter for ArrayStaticFilter { @@ -76,85 +157,6 @@ impl StaticFilter for ArrayStaticFilter { _ => {} } - let needle_nulls = v.logical_nulls(); - let needle_nulls = needle_nulls.as_ref(); - let haystack_has_nulls = self.in_array.null_count() != 0; - - with_hashes([v], &self.state, |hashes| { - let cmp = make_comparator(v, &self.in_array, SortOptions::default())?; - Ok((0..v.len()) - .map(|i| { - // SQL three-valued logic: null IN (...) is always null - if needle_nulls.is_some_and(|nulls| nulls.is_null(i)) { - return None; - } - - let hash = hashes[i]; - let contains = self - .map - .raw_entry() - .from_hash(hash, |idx| cmp(i, *idx).is_eq()) - .is_some(); - - match contains { - true => Some(!negated), - false if haystack_has_nulls => None, - false => Some(negated), - } - }) - .collect()) - }) - } -} - -impl ArrayStaticFilter { - /// Computes a [`StaticFilter`] for the provided [`Array`] if there - /// are nulls present or there are more than the configured number of - /// elements. - /// - /// Note: This is split into a separate function as higher-rank trait bounds currently - /// cause type inference to misbehave - pub(super) fn try_new(in_array: ArrayRef) -> Result { - // Null type has no natural order - return empty hash set - if in_array.data_type() == &DataType::Null { - return Ok(ArrayStaticFilter { - in_array, - state: RandomState::default(), - map: HashMap::with_hasher(()), - }); - } - - let state = RandomState::default(); - let mut map: HashMap = HashMap::with_hasher(()); - - with_hashes([&in_array], &state, |hashes| -> Result<()> { - let cmp = make_comparator(&in_array, &in_array, SortOptions::default())?; - - let insert_value = |idx| { - let hash = hashes[idx]; - if let RawEntryMut::Vacant(v) = map - .raw_entry_mut() - .from_hash(hash, |x| cmp(*x, idx).is_eq()) - { - v.insert_with_hasher(hash, idx, (), |x| hashes[*x]); - } - }; - - match in_array.nulls() { - Some(nulls) => { - BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) - .for_each(insert_value) - } - None => (0..in_array.len()).for_each(insert_value), - } - - Ok(()) - })?; - - Ok(Self { - in_array, - state, - map, - }) + self.find_needles_in_haystack(v, negated) } } diff --git a/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs new file mode 100644 index 0000000000000..cd0cbd0de59a8 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/branchless_filter.rs @@ -0,0 +1,578 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Fast membership tests for small, fixed-width primitive `IN` lists. +//! +//! # Why use a branchless filter? +//! +//! For a short list such as `x IN (10, 20, 30)`, it can be faster to compare +//! `x` with all three values than to build and search a hash table. +//! +//! "Branchless" means that the filter always checks every list value. It +//! combines the answers with `|`, while `||` would stop at the first match. +//! This regular sequence of comparisons is easier for the compiler and CPU to +//! optimize. +//! +//! # How does it work? +//! +//! When the filter is built, it stores the non-null list values and chooses a +//! comparison function for that list length. Only this small function is +//! specialized for each length. The rest of [`BranchlessFilter`] is shared, +//! which keeps the generated code small. +//! +//! Some Arrow types share the same in-memory representation. For example, a +//! `Float32` and a `UInt32` both use four bytes per value. The filter compares +//! those stored bits through an unsigned type of the same size, without copying +//! the value buffer. A bit pattern is simply the bytes Arrow uses to store a +//! value. Comparing it preserves details such as `0.0` versus `-0.0` and +//! different NaN values. [`BranchlessFilterType`] defines these safe, +//! same-sized mappings and checks their sizes at compile time. +//! +//! The fast path is intentionally limited to short lists: +//! +//! - 16 values for 1-byte types +//! - 8 values for 2-byte types +//! - 32 values for 4-byte types +//! - 16 values for 8-byte types +//! - 4 values for 16-byte types +//! +//! These numbers do not follow one size-based pattern. One- and two-byte +//! values have an especially efficient next step: every possible bit pattern +//! fits in a compact bitmap. For a longer list, the bitmap filter turns on one +//! bit for each listed value, then checks membership with a direct bit lookup. +//! This becomes a better fit before a 64- or 128-comparison branchless chain +//! would be useful. Wider types have too many possible values for such a +//! bitmap, so their limits are tuned separately. +//! +//! Larger lists use the standard filter strategy, including bitmap filters for +//! one- and two-byte types. +//! +//! # What about nulls? +//! +//! Null list entries are omitted from the comparison chain but counted by the +//! filter. Evaluation first records which values matched, then +//! [`build_result_from_contains`] combines it with input nulls, list nulls, and +//! `NOT IN` to produce the usual SQL null behavior. + +use std::mem::size_of; + +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray, PrimitiveArray}; +use arrow::buffer::{BooleanBuffer, ScalarBuffer}; +use arrow::datatypes::*; +use arrow::util::bit_iterator::BitIndexIterator; +use datafusion_common::{Result, exec_datafusion_err, internal_datafusion_err}; + +use super::result::build_result_from_contains; +use super::static_filter::{StaticFilter, handle_dictionary}; + +pub(super) type BranchlessNative = + <::CompareType as ArrowPrimitiveType>::Native; + +/// Maximum list size for branchless lookup on 1-byte primitives. +/// +/// Sixteen 1-byte values fit in one 128-bit SIMD vector, so this keeps the +/// branchless list small enough for a single vectorized membership check. +const BRANCHLESS_MAX_1B: usize = 16; + +/// Maximum list size for branchless lookup on 2-byte primitives. +/// +/// Eight 2-byte values fit in one 128-bit SIMD vector, so this keeps the +/// branchless list small enough for a single vectorized membership check. +const BRANCHLESS_MAX_2B: usize = 8; + +/// Maximum list size for branchless lookup on 4-byte primitives. +/// +/// Thirty-two 4-byte values keep the inline list at 128 bytes. Beyond that, +/// the comparison chain and filter footprint grow enough that the hash/generic +/// fallback is a better fit. +const BRANCHLESS_MAX_4B: usize = 32; + +/// Maximum list size for branchless lookup on 8-byte primitives. +/// +/// Sixteen 8-byte values use the same 128-byte inline-list budget as 4-byte +/// primitives. Larger lists are left to the hash/generic fallback. +const BRANCHLESS_MAX_8B: usize = 16; + +/// Maximum list size for branchless lookup on 16-byte primitives. +/// +/// These comparisons are wider, so this path is limited to four values. +/// Larger lists are left to the generic fallback. +const BRANCHLESS_MAX_16B: usize = 4; + +/// Arrow primitive types supported by [`BranchlessFilter`]. +/// +/// `T` is the logical Arrow type accepted by the filter. `CompareType` is the +/// same-width type used for the fixed comparison chain. Signed integers, +/// floats, and temporal values use an unsigned comparison type so they compare +/// by their raw bit pattern. +pub(super) trait BranchlessFilterType: + ArrowPrimitiveType + Send + Sync + 'static +{ + type CompareType: ArrowPrimitiveType + Send + Sync + 'static; + + /// Maximum number of non-null IN-list values to handle with + /// [`BranchlessFilter`] for this primitive type. + const MAX_LIST_LEN: usize; +} + +macro_rules! branchless_filter_type { + ($logical:ty, $compare:ty, $max_len:expr) => { + // The branchless filter reads the same Arrow value buffer as the + // comparison type. That is only valid when both native types have the + // same width, so catch any bad mapping here at compile time. + const _: () = assert!( + size_of::<<$logical as ArrowPrimitiveType>::Native>() + == size_of::<<$compare as ArrowPrimitiveType>::Native>(), + "BranchlessFilterType::CompareType must use the same native width" + ); + + impl BranchlessFilterType for $logical { + type CompareType = $compare; + const MAX_LIST_LEN: usize = $max_len; + } + }; +} + +branchless_filter_type!(Int8Type, UInt8Type, BRANCHLESS_MAX_1B); +branchless_filter_type!(UInt8Type, UInt8Type, BRANCHLESS_MAX_1B); +branchless_filter_type!(Int16Type, UInt16Type, BRANCHLESS_MAX_2B); +branchless_filter_type!(UInt16Type, UInt16Type, BRANCHLESS_MAX_2B); +branchless_filter_type!(Float16Type, UInt16Type, BRANCHLESS_MAX_2B); + +branchless_filter_type!(Int32Type, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(UInt32Type, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(Float32Type, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(Date32Type, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(Time32SecondType, UInt32Type, BRANCHLESS_MAX_4B); +branchless_filter_type!(Time32MillisecondType, UInt32Type, BRANCHLESS_MAX_4B); + +branchless_filter_type!(Int64Type, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(UInt64Type, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(Float64Type, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(Date64Type, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(Time64MicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(Time64NanosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(TimestampSecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(TimestampMillisecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(TimestampMicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(TimestampNanosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(DurationSecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(DurationMillisecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(DurationMicrosecondType, UInt64Type, BRANCHLESS_MAX_8B); +branchless_filter_type!(DurationNanosecondType, UInt64Type, BRANCHLESS_MAX_8B); + +branchless_filter_type!(Decimal128Type, Decimal128Type, BRANCHLESS_MAX_16B); +branchless_filter_type!( + IntervalMonthDayNanoType, + IntervalMonthDayNanoType, + BRANCHLESS_MAX_16B +); + +/// Checks each input value against the `IN`-list values. +type MembershipCheck = fn(in_list_values: &[C], input_values: &[C]) -> BooleanBuffer; + +/// A branchless filter for fixed-width primitive `IN` lists up to +/// `T::MAX_LIST_LEN` values. +/// +/// The filter stores the non-null `IN`-list values in a slice and chooses a +/// comparison function for that length. Keeping the length out of +/// `BranchlessFilter` avoids generating a full copy of the filter for every +/// supported length. +pub(super) struct BranchlessFilter { + expected_data_type: DataType, + null_count: usize, + in_list_values: Box<[BranchlessNative]>, + check_values: MembershipCheck>, +} + +impl BranchlessFilter +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq, +{ + pub(super) fn try_new(in_array: &ArrayRef) -> Result { + let in_array = in_array.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("BranchlessFilter: expected {} array", T::DATA_TYPE) + })?; + let non_null_count = in_array.len() - in_array.null_count(); + // `try_new` can be called on its own, so check the limit here too. + if non_null_count > T::MAX_LIST_LEN { + return Err(internal_datafusion_err!( + "BranchlessFilter: supports at most {} non-null values, got {non_null_count}", + T::MAX_LIST_LEN + )); + } + + let all_values = branchless_values::(in_array); + let mut in_list_values = Vec::with_capacity(non_null_count); + + match in_array.nulls() { + None => { + in_list_values.extend(all_values.iter().copied()); + } + Some(nulls) => { + for row in + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + { + in_list_values.push(all_values[row]); + } + } + } + + debug_assert_eq!(in_list_values.len(), non_null_count); + let in_list_values = in_list_values.into_boxed_slice(); + let check_values = membership_check_for_len::(in_list_values.len()); + + Ok(Self { + expected_data_type: in_array.data_type().clone(), + null_count: in_array.null_count(), + in_list_values, + check_values, + }) + } +} + +impl StaticFilter for BranchlessFilter +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq + Send + Sync, +{ + fn null_count(&self) -> usize { + self.null_count + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + + // Arrow compatibility ignores timestamp timezone and decimal precision/scale + // while still requiring the same primitive representation. + if !PrimitiveArray::::is_compatible(v.data_type()) { + return Err(exec_datafusion_err!( + "BranchlessFilter: expected {} array, got {}", + self.expected_data_type, + v.data_type() + )); + } + + let v = v.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("BranchlessFilter: expected {} array", T::DATA_TYPE) + })?; + let input_values = branchless_values::(v); + let matches = + (self.check_values)(self.in_list_values.as_ref(), input_values.as_ref()); + Ok(build_result_from_contains( + v.nulls(), + self.null_count > 0, + negated, + matches, + )) + } +} + +/// Picks the comparison function for `len` non-null `IN`-list values. +/// +/// A length of zero is used when the list contains only nulls. The comparisons +/// return false, and the caller then applies the usual SQL null behavior. +fn membership_check_for_len(len: usize) -> MembershipCheck> +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq, +{ + macro_rules! choose { + ($($n:literal),* $(,)?) => { + match len { + $($n => check_values::, $n>,)* + _ => unreachable!("list length exceeds the configured limit"), + } + }; + } + + // Avoid creating checks for lengths a type does not support. + match T::MAX_LIST_LEN { + 4 => choose!(0, 1, 2, 3, 4), + 8 => choose!(0, 1, 2, 3, 4, 5, 6, 7, 8), + 16 => choose!(0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16), + 32 => choose!( + 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, + 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, + ), + _ => unreachable!("list-size limits must be 4, 8, 16, or 32"), + } +} + +#[inline] +fn check_values( + in_list_values: &[C], + input_values: &[C], +) -> BooleanBuffer +where + C: Copy + PartialEq, +{ + let in_list_values: &[C; N] = in_list_values + .try_into() + .expect("comparison length matches IN-list values"); + + BooleanBuffer::collect_bool(input_values.len(), |i| { + // SAFETY: `collect_bool` invokes this closure for indices in + // `0..input_values.len()`. + let input_value = unsafe { *input_values.get_unchecked(i) }; + // `|` checks every list value; `||` would stop after the first match. + in_list_values + .iter() + .fold(false, |acc, &value| acc | (value == input_value)) + }) +} + +fn branchless_values(array: &PrimitiveArray) -> ScalarBuffer> +where + T: BranchlessFilterType, +{ + let data = array.to_data(); + ScalarBuffer::>::new( + data.buffers()[0].clone(), + data.offset(), + data.len(), + ) +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{ + Decimal128Array, Float16Array, Float32Array, Float64Array, Int8Array, + IntervalMonthDayNanoArray, TimestampMillisecondArray, TimestampNanosecondArray, + UInt8Array, UInt16Array, + }; + use half::f16; + + use super::*; + + fn assert_contains( + filter: &dyn StaticFilter, + needles: &dyn Array, + expected: Vec>, + ) -> Result<()> { + assert_eq!( + filter.contains(needles, false)?, + BooleanArray::from(expected) + ); + Ok(()) + } + + #[test] + fn branchless_filter_u8_handles_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = UInt8Array::from(vec![Some(1), Some(2), None, Some(3)]); + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, Some(false)]) + ); + + Ok(()) + } + + #[test] + fn branchless_filter_all_null_list_preserves_sql_null_semantics() -> Result<()> { + let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![None, None])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = UInt8Array::from(vec![Some(1), None]); + let expected = BooleanArray::from(vec![None, None]); + + assert_eq!(filter.contains(&needles, false)?, expected); + assert_eq!(filter.contains(&needles, true)?, expected); + + Ok(()) + } + + #[test] + fn branchless_filter_i8_handles_signed_boundaries_and_slices() -> Result<()> { + let haystack: ArrayRef = Arc::new( + Int8Array::from(vec![Some(99), Some(i8::MIN), None, Some(-1), Some(42)]) + .slice(1, 3), + ); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = + Int8Array::from(vec![Some(7), Some(i8::MIN), Some(-1), None]).slice(1, 3); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(true), None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), Some(false), None]) + ); + + let wrong_type = UInt8Array::from(vec![Some(128), Some(u8::MAX)]); + let err = filter.contains(&wrong_type, false).unwrap_err().to_string(); + assert!(err.contains("expected Int8 array, got UInt8"), "{err}"); + + Ok(()) + } + + #[test] + fn branchless_filter_f16_handles_bit_patterns_and_slices() -> Result<()> { + let nan_a = f16::from_bits(0x7e01); + let nan_b = f16::from_bits(0x7e02); + let haystack: ArrayRef = Arc::new( + Float16Array::from(vec![ + Some(f16::from_f32(9.0)), + Some(f16::from_f32(-0.0)), + Some(nan_a), + None, + ]) + .slice(1, 3), + ); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = Float16Array::from(vec![ + Some(f16::from_f32(0.0)), + Some(f16::from_f32(-0.0)), + Some(nan_a), + Some(nan_b), + None, + ]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![None, Some(true), Some(true), None, None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![None, Some(false), Some(false), None, None]) + ); + + let wrong_type = UInt16Array::from(vec![Some(0x8000), Some(0x7e01)]); + let err = filter.contains(&wrong_type, false).unwrap_err().to_string(); + assert!(err.contains("expected Float16 array, got UInt16"), "{err}"); + + Ok(()) + } + + #[test] + fn branchless_filter_floats_use_bit_equality() -> Result<()> { + let nan_a = f32::from_bits(0x7fc0_0001); + let nan_b = f32::from_bits(0x7fc0_0002); + let haystack: ArrayRef = + Arc::new(Float32Array::from(vec![Some(-0.0), Some(nan_a)])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = + Float32Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + ); + + let nan_a = f64::from_bits(0x7ff8_0000_0000_0001); + let nan_b = f64::from_bits(0x7ff8_0000_0000_0002); + let haystack: ArrayRef = + Arc::new(Float64Array::from(vec![Some(-0.0), Some(nan_a)])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = + Float64Array::from(vec![Some(0.0), Some(-0.0), Some(nan_a), Some(nan_b)]); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(false), Some(true), Some(true), Some(false)]) + ); + + Ok(()) + } + + #[test] + fn branchless_filter_timestamp_uses_physical_compatibility() -> Result<()> { + let haystack: ArrayRef = Arc::new( + TimestampNanosecondArray::from(vec![Some(1), Some(3)]).with_timezone("UTC"), + ); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = TimestampNanosecondArray::from(vec![Some(1), Some(2), None]) + .with_timezone("UTC"); + + assert_contains(&filter, &needles, vec![Some(true), Some(false), None])?; + + let different_timezone = TimestampNanosecondArray::from(vec![Some(1), Some(2)]) + .with_timezone("Europe/Paris"); + assert_contains(&filter, &different_timezone, vec![Some(true), Some(false)])?; + + let different_unit = TimestampMillisecondArray::from(vec![Some(1)]); + let err = filter + .contains(&different_unit, false) + .unwrap_err() + .to_string(); + assert!(err.contains("Timestamp(ns"), "{err}"); + assert!(err.contains("Timestamp(ms"), "{err}"); + + Ok(()) + } + + #[test] + fn branchless_filter_decimal128_handles_precision_scale_and_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new( + Decimal128Array::from(vec![Some(12345), None, Some(-700), Some(42)]) + .with_precision_and_scale(10, 2)?, + ); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = + Decimal128Array::from(vec![Some(12345), Some(999), None, Some(-700)]) + .with_precision_and_scale(10, 2)?; + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, Some(false)]) + ); + + let compatible_metadata = + Decimal128Array::from(vec![Some(12345)]).with_precision_and_scale(11, 3)?; + assert_contains(&filter, &compatible_metadata, vec![Some(true)])?; + + Ok(()) + } + + #[test] + fn branchless_filter_interval_month_day_nano_handles_nulls() -> Result<()> { + let one_month = IntervalMonthDayNanoType::make_value(1, 0, 0); + let two_days = IntervalMonthDayNanoType::make_value(0, 2, 0); + let three_nanos = IntervalMonthDayNanoType::make_value(0, 0, 3); + let absent = IntervalMonthDayNanoType::make_value(4, 5, 6); + let haystack: ArrayRef = Arc::new(IntervalMonthDayNanoArray::from(vec![ + Some(one_month), + None, + Some(two_days), + Some(three_nanos), + ])); + let filter = BranchlessFilter::::try_new(&haystack)?; + let needles = IntervalMonthDayNanoArray::from(vec![ + Some(one_month), + Some(absent), + None, + Some(three_nanos), + ]); + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, Some(false)]) + ); + + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs index 2c084a1cb247b..8f8d9bad04afa 100644 --- a/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/primitive_filter.rs @@ -15,16 +15,212 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ - Array, ArrayRef, AsArray, BooleanArray, downcast_array, downcast_dictionary_array, -}; +//! Optimized primitive type filters for InList expressions. +//! +//! This module provides membership tests for Arrow primitive types. + +use arrow::array::{Array, ArrayRef, AsArray, BooleanArray}; use arrow::buffer::{BooleanBuffer, NullBuffer}; -use arrow::compute::take; use arrow::datatypes::*; +use arrow::util::bit_iterator::BitIndexIterator; use datafusion_common::{HashSet, Result, exec_datafusion_err}; use std::hash::{Hash, Hasher}; -use super::static_filter::StaticFilter; +use super::result::build_in_list_result; +use super::static_filter::{StaticFilter, handle_dictionary}; + +/// Storage for the bits used by [`BitmapFilter`]. +/// +/// `BitmapFilter` represents an `IN` list with one bit for each possible +/// value, so membership checks become direct bit tests. This trait lets the +/// same filter code use different storage sizes for different integer widths. +pub(super) trait BitmapStorage: Send + Sync { + fn new_zeroed() -> Self; + fn set_bit(&mut self, index: usize); + fn get_bit(&self, index: usize) -> bool; +} + +// `UInt8` has 256 possible values, 0 through 255. One bit per value takes +// 256 bits, which fits in four `u64` words. +impl BitmapStorage for [u64; 4] { + #[inline] + fn new_zeroed() -> Self { + [0u64; 4] + } + #[inline] + fn set_bit(&mut self, index: usize) { + self[index / 64] |= 1u64 << (index % 64); + } + #[inline(always)] + fn get_bit(&self, index: usize) -> bool { + (self[index / 64] >> (index % 64)) & 1 != 0 + } +} + +// `UInt16` has 65,536 possible values. One bit per value takes 65,536 bits, +// which is 1,024 `u64` words, or 8 KiB. Box the array so the filter stores a +// pointer instead of carrying an 8 KiB array inline. +impl BitmapStorage for Box<[u64; 1024]> { + #[inline] + fn new_zeroed() -> Self { + Box::new([0u64; 1024]) + } + #[inline] + fn set_bit(&mut self, index: usize) { + self[index / 64] |= 1u64 << (index % 64); + } + #[inline(always)] + fn get_bit(&self, index: usize) -> bool { + (self[index / 64] >> (index % 64)) & 1 != 0 + } +} + +/// Arrow primitive types supported by [`BitmapFilter`]. +/// +/// Arrow already defines the Rust value type as `T::Native`. This trait only +/// supplies the bitmap storage size and maps values to their bit-pattern index +/// for the primitive domains that are small enough to represent with one bit +/// per possible value. +pub(super) trait BitmapFilterType: + ArrowPrimitiveType + Send + Sync + 'static +{ + type Storage: BitmapStorage; + + /// Returns the index in the bitmap to check for this value. + fn index(value: Self::Native) -> usize; +} + +/// `Int8` has 256 possible bit patterns, so four `u64` words cover the full domain. +impl BitmapFilterType for Int8Type { + type Storage = [u64; 4]; + + #[inline(always)] + fn index(value: Self::Native) -> usize { + // Reinterpret the signed value's bit pattern into a bitmap index. + value as u8 as usize + } +} + +/// `UInt8` has 256 possible values, so four `u64` words cover the full domain. +impl BitmapFilterType for UInt8Type { + type Storage = [u64; 4]; + + #[inline(always)] + fn index(value: Self::Native) -> usize { + value as usize + } +} + +/// `Int16` has 65,536 possible bit patterns, so 1,024 `u64` words cover the full +/// domain. +impl BitmapFilterType for Int16Type { + type Storage = Box<[u64; 1024]>; + + #[inline(always)] + fn index(value: Self::Native) -> usize { + // Reinterpret the signed value's bit pattern into a bitmap index. + value as u16 as usize + } +} + +/// `UInt16` has 65,536 possible values, so 1,024 `u64` words cover the full +/// domain. +impl BitmapFilterType for UInt16Type { + type Storage = Box<[u64; 1024]>; + + #[inline(always)] + fn index(value: Self::Native) -> usize { + value as usize + } +} + +/// `Float16` has 65,536 possible bit patterns, so 1,024 `u64` words cover the +/// full domain. +impl BitmapFilterType for Float16Type { + type Storage = Box<[u64; 1024]>; + + #[inline(always)] + fn index(value: Self::Native) -> usize { + value.to_bits() as usize + } +} + +/// `IN` filter backed by one bit per possible value. +/// +/// Building the filter scans the non-null values in the IN-list and turns on +/// the bit selected by each value. Evaluating input values checks the same bit +/// position. Null handling and `NOT IN` inversion are handled by +/// `build_in_list_result`. +pub(super) struct BitmapFilter { + null_count: usize, + bits: T::Storage, +} + +impl BitmapFilter +where + T: BitmapFilterType, +{ + pub(super) fn try_new(in_array: &ArrayRef) -> Result { + let prim_array = in_array.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("BitmapFilter: expected {} array", T::DATA_TYPE) + })?; + let mut bits = T::Storage::new_zeroed(); + let values = prim_array.values(); + match prim_array.nulls() { + None => { + for &v in values { + bits.set_bit(T::index(v)); + } + } + Some(nulls) => { + for i in + BitIndexIterator::new(nulls.validity(), nulls.offset(), nulls.len()) + { + bits.set_bit(T::index(values[i])); + } + } + } + Ok(Self { + null_count: prim_array.null_count(), + bits, + }) + } + + #[inline(always)] + fn check(&self, needle: T::Native) -> bool { + self.bits.get_bit(T::index(needle)) + } +} + +impl StaticFilter for BitmapFilter +where + T: BitmapFilterType, +{ + fn null_count(&self) -> usize { + self.null_count + } + + fn contains(&self, v: &dyn Array, negated: bool) -> Result { + handle_dictionary!(self, v, negated); + let v = v.as_primitive_opt::().ok_or_else(|| { + exec_datafusion_err!("BitmapFilter: expected {} array", T::DATA_TYPE) + })?; + let input_values = v.values(); + Ok(build_in_list_result( + v.len(), + v.nulls(), + self.null_count > 0, + negated, + #[inline(always)] + |i| { + // SAFETY: `build_in_list_result` invokes this closure for + // indices in `0..v.len()`, which matches `input_values.len()`. + let needle = unsafe { *input_values.get_unchecked(i) }; + self.check(needle) + }, + )) + } +} /// Wrapper for f32 that implements Hash and Eq using bit comparison. /// This treats NaN values as equal to each other when they have the same bit pattern. @@ -94,9 +290,13 @@ macro_rules! primitive_static_filter { impl $Name { pub(super) fn try_new(in_array: &ArrayRef) -> Result { - let in_array = in_array - .as_primitive_opt::<$ArrowType>() - .ok_or_else(|| exec_datafusion_err!("Failed to downcast an array to a '{}' array", stringify!($ArrowType)))?; + let in_array = + in_array.as_primitive_opt::<$ArrowType>().ok_or_else(|| { + exec_datafusion_err!( + "Failed to downcast an array to a '{}' array", + stringify!($ArrowType) + ) + })?; let mut values = HashSet::with_capacity(in_array.len()); let null_count = in_array.null_count(); @@ -115,19 +315,14 @@ macro_rules! primitive_static_filter { } fn contains(&self, v: &dyn Array, negated: bool) -> Result { - // Handle dictionary arrays by recursing on the values - downcast_dictionary_array! { - v => { - let values_contains = self.contains(v.values().as_ref(), negated)?; - let result = take(&values_contains, v.keys(), None)?; - return Ok(downcast_array(result.as_ref())) - } - _ => {} - } + handle_dictionary!(self, v, negated); - let v = v - .as_primitive_opt::<$ArrowType>() - .ok_or_else(|| exec_datafusion_err!("Failed to downcast an array to a '{}' array", stringify!($ArrowType)))?; + let v = v.as_primitive_opt::<$ArrowType>().ok_or_else(|| { + exec_datafusion_err!( + "Failed to downcast an array to a '{}' array", + stringify!($ArrowType) + ) + })?; let haystack_has_nulls = self.null_count > 0; let needle_values = v.values(); @@ -188,8 +383,10 @@ macro_rules! primitive_static_filter { } (true, true) => { // Both have nulls - combine needle nulls with haystack-induced nulls - let needle_validity = needle_nulls.map(|n| n.inner().clone()) - .unwrap_or_else(|| BooleanBuffer::new_set(needle_values.len())); + let needle_validity = + needle_nulls.map(|n| n.inner().clone()).unwrap_or_else( + || BooleanBuffer::new_set(needle_values.len()), + ); // Valid when original "in set" is true (see above) let haystack_validity = if negated { @@ -210,13 +407,8 @@ macro_rules! primitive_static_filter { }; } -// Generate specialized filters for all integer primitive types -primitive_static_filter!(Int8StaticFilter, Int8Type); -primitive_static_filter!(Int16StaticFilter, Int16Type); primitive_static_filter!(Int32StaticFilter, Int32Type); primitive_static_filter!(Int64StaticFilter, Int64Type); -primitive_static_filter!(UInt8StaticFilter, UInt8Type); -primitive_static_filter!(UInt16StaticFilter, UInt16Type); primitive_static_filter!(UInt32StaticFilter, UInt32Type); primitive_static_filter!(UInt64StaticFilter, UInt64Type); @@ -231,3 +423,165 @@ macro_rules! float_static_filter { // Generate specialized filters for float types using ordered wrappers float_static_filter!(Float32StaticFilter, Float32Type, OrderedFloat32); float_static_filter!(Float64StaticFilter, Float64Type, OrderedFloat64); + +#[cfg(test)] +mod tests { + use super::*; + use std::sync::Arc; + + use arrow::array::{ + DictionaryArray, Float16Array, Int8Array, Int16Array, UInt8Array, UInt16Array, + }; + use half::f16; + + fn assert_contains( + filter: &dyn StaticFilter, + needles: &dyn Array, + expected: Vec>, + ) -> Result<()> { + assert_eq!( + filter.contains(needles, false)?, + BooleanArray::from(expected) + ); + Ok(()) + } + + #[test] + fn bitmap_filter_u8_handles_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)])); + let filter = BitmapFilter::::try_new(&haystack)?; + let needles = UInt8Array::from(vec![Some(1), Some(2), None, Some(3)]); + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)])?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, None, Some(false)]) + ); + + Ok(()) + } + + #[test] + fn bitmap_filter_u8_handles_dictionary_needles() -> Result<()> { + let haystack: ArrayRef = Arc::new(UInt8Array::from(vec![Some(1), None, Some(3)])); + let filter = BitmapFilter::::try_new(&haystack)?; + + let keys = Int8Array::from(vec![Some(0), Some(1), None, Some(2)]); + let values = Arc::new(UInt8Array::from(vec![Some(1), Some(2), Some(3)])); + let needles = DictionaryArray::try_new(keys, values)?; + + assert_contains(&filter, &needles, vec![Some(true), None, None, Some(true)]) + } + + #[test] + fn bitmap_filter_i8_handles_signed_boundaries_and_slices() -> Result<()> { + let haystack: ArrayRef = Arc::new( + Int8Array::from(vec![Some(99), Some(i8::MIN), None, Some(-1), Some(42)]) + .slice(1, 3), + ); + let filter = BitmapFilter::::try_new(&haystack)?; + let needles = + Int8Array::from(vec![Some(7), Some(i8::MIN), Some(-1), None]).slice(1, 3); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(true), None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), Some(false), None]) + ); + + Ok(()) + } + + #[test] + fn bitmap_filter_u16_handles_boundaries_and_nulls() -> Result<()> { + let haystack: ArrayRef = Arc::new(UInt16Array::from(vec![ + Some(0), + None, + Some(1024), + Some(u16::MAX), + ])); + let filter = BitmapFilter::::try_new(&haystack)?; + let needles = + UInt16Array::from(vec![Some(0), Some(1), Some(1024), Some(u16::MAX), None]); + + assert_contains( + &filter, + &needles, + vec![Some(true), None, Some(true), Some(true), None], + )?; + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, Some(false), Some(false), None]) + ); + + Ok(()) + } + + #[test] + fn bitmap_filter_i16_handles_signed_boundaries_and_slices() -> Result<()> { + let haystack: ArrayRef = Arc::new( + Int16Array::from(vec![ + Some(123), + Some(i16::MIN), + None, + Some(-1), + Some(i16::MAX), + ]) + .slice(1, 4), + ); + let filter = BitmapFilter::::try_new(&haystack)?; + let needles = + Int16Array::from(vec![Some(0), Some(i16::MIN), Some(7), Some(i16::MAX)]) + .slice(1, 3); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), None, Some(true)]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), None, Some(false)]) + ); + + Ok(()) + } + + #[test] + fn bitmap_filter_f16_handles_bit_patterns_and_slices() -> Result<()> { + let nan_a = f16::from_bits(0x7e01); + let nan_b = f16::from_bits(0x7e02); + let haystack: ArrayRef = Arc::new( + Float16Array::from(vec![ + Some(f16::from_f32(9.0)), + Some(f16::from_f32(1.5)), + None, + Some(f16::from_f32(-0.0)), + Some(nan_a), + ]) + .slice(1, 4), + ); + let filter = BitmapFilter::::try_new(&haystack)?; + let needles = Float16Array::from(vec![ + Some(f16::from_f32(0.0)), + Some(f16::from_f32(-0.0)), + Some(nan_a), + Some(nan_b), + None, + ]) + .slice(1, 4); + + assert_eq!( + filter.contains(&needles, false)?, + BooleanArray::from(vec![Some(true), Some(true), None, None]) + ); + assert_eq!( + filter.contains(&needles, true)?, + BooleanArray::from(vec![Some(false), Some(false), None, None]) + ); + + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list/result.rs b/datafusion/physical-expr/src/expressions/in_list/result.rs new file mode 100644 index 0000000000000..3ebdbfe19f743 --- /dev/null +++ b/datafusion/physical-expr/src/expressions/in_list/result.rs @@ -0,0 +1,105 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Result building helpers for InList operations. +//! +//! This module provides unified logic for building BooleanArray results +//! from IN list membership tests, handling null propagation correctly +//! according to SQL three-valued logic. + +use arrow::array::BooleanArray; +use arrow::buffer::{BooleanBuffer, NullBuffer}; + +// Truth table for (needle_nulls, haystack_has_nulls, negated): +// (Some, true, false) => values: valid & contains, nulls: valid & contains +// (None, true, false) => values: contains, nulls: contains +// (Some, true, true) => values: valid & !contains, nulls: valid & contains +// (None, true, true) => values: !contains, nulls: contains +// (Some, false, false) => values: valid & contains, nulls: valid +// (Some, false, true) => values: valid & !contains, nulls: valid +// (None, false, false) => values: contains, nulls: none +// (None, false, true) => values: !contains, nulls: none + +/// Builds a BooleanArray result for IN list operations. +/// +/// This function handles the null propagation logic for SQL IN lists: +/// - If the needle value is null, the result is null +/// - If the needle is not in the set and the haystack has nulls, the result is null +/// - Otherwise, the result is true/false based on membership and negation +/// +/// This version computes contains for all positions, including nulls, then applies +/// null masking via bitmap operations. +#[inline] +pub(crate) fn build_in_list_result( + len: usize, + needle_nulls: Option<&NullBuffer>, + haystack_has_nulls: bool, + negated: bool, + contains: C, +) -> BooleanArray +where + C: FnMut(usize) -> bool, +{ + let contains_buf = BooleanBuffer::collect_bool(len, contains); + build_result_from_contains(needle_nulls, haystack_has_nulls, negated, contains_buf) +} + +/// Builds a BooleanArray result from a pre-computed contains buffer. +/// +/// This version does not assume contains_buf is pre-masked at null positions. +/// It handles nulls using bitmap operations. +#[inline] +pub(crate) fn build_result_from_contains( + needle_nulls: Option<&NullBuffer>, + haystack_has_nulls: bool, + negated: bool, + contains_buf: BooleanBuffer, +) -> BooleanArray { + match (needle_nulls, haystack_has_nulls, negated) { + // Haystack has nulls: result is null unless value is found. + (Some(v), true, false) => { + // values: valid & contains, nulls: valid & contains + let values = v.inner() & &contains_buf; + BooleanArray::new(values.clone(), Some(NullBuffer::new(values))) + } + (None, true, false) => { + BooleanArray::new(contains_buf.clone(), Some(NullBuffer::new(contains_buf))) + } + (Some(v), true, true) => { + // NOT IN with nulls: false if found, null if not found or needle null. + // values: valid & !contains, nulls: valid & contains + let valid = v.inner(); + let values = valid & &(!&contains_buf); + let nulls = valid & &contains_buf; + BooleanArray::new(values, Some(NullBuffer::new(nulls))) + } + (None, true, true) => { + BooleanArray::new(!&contains_buf, Some(NullBuffer::new(contains_buf))) + } + // Haystack has no nulls: result validity follows needle validity. + (Some(v), false, false) => { + // values: valid & contains, nulls: valid + BooleanArray::new(v.inner() & &contains_buf, Some(v.clone())) + } + (Some(v), false, true) => { + // values: valid & !contains, nulls: valid + BooleanArray::new(v.inner() & &(!&contains_buf), Some(v.clone())) + } + (None, false, false) => BooleanArray::new(contains_buf, None), + (None, false, true) => BooleanArray::new(!&contains_buf, None), + } +} diff --git a/datafusion/physical-expr/src/expressions/in_list/static_filter.rs b/datafusion/physical-expr/src/expressions/in_list/static_filter.rs index 218bd27950266..3c964d4183474 100644 --- a/datafusion/physical-expr/src/expressions/in_list/static_filter.rs +++ b/datafusion/physical-expr/src/expressions/in_list/static_filter.rs @@ -35,3 +35,20 @@ pub(super) trait StaticFilter { /// implementation unwraps the dictionary and operates on its values. fn contains(&self, v: &dyn Array, negated: bool) -> Result; } + +/// Evaluate dictionary-encoded needles by applying a filter to dictionary +/// values and remapping the result through the keys. +macro_rules! handle_dictionary { + ($self:ident, $v:ident, $negated:ident) => { + arrow::array::downcast_dictionary_array! { + $v => { + let values_contains = $self.contains($v.values().as_ref(), $negated)?; + let result = arrow::compute::take(&values_contains, $v.keys(), None)?; + return Ok(arrow::array::downcast_array(result.as_ref())) + } + _ => {} + } + }; +} + +pub(super) use handle_dictionary; diff --git a/datafusion/physical-expr/src/expressions/in_list/strategy.rs b/datafusion/physical-expr/src/expressions/in_list/strategy.rs index b7ee3dd1a3b9d..d5ca8154a92f6 100644 --- a/datafusion/physical-expr/src/expressions/in_list/strategy.rs +++ b/datafusion/physical-expr/src/expressions/in_list/strategy.rs @@ -19,39 +19,179 @@ use std::sync::Arc; use arrow::array::ArrayRef; use arrow::compute::cast; -use arrow::datatypes::DataType; +use arrow::datatypes::{ + DataType, Date32Type, Date64Type, Decimal128Type, DurationMicrosecondType, + DurationMillisecondType, DurationNanosecondType, DurationSecondType, Float16Type, + Float32Type, Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, + IntervalMonthDayNanoType, IntervalUnit, Time32MillisecondType, Time32SecondType, + Time64MicrosecondType, Time64NanosecondType, TimeUnit, TimestampMicrosecondType, + TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, + UInt16Type, UInt32Type, UInt64Type, +}; use datafusion_common::Result; use super::array_static_filter::ArrayStaticFilter; +use super::branchless_filter::{ + BranchlessFilter, BranchlessFilterType, BranchlessNative, +}; use super::primitive_filter::*; use super::static_filter::StaticFilter; -pub(super) fn instantiate_static_filter( - in_array: ArrayRef, -) -> Result> { +type StaticFilterRef = Arc; + +pub(super) fn instantiate_static_filter(in_array: ArrayRef) -> Result { + let in_array = flatten_dictionary_haystack(in_array)?; + + if let Some(filter) = instantiate_branchless_filter(&in_array)? { + return Ok(filter); + } + + instantiate_standard_filter(in_array) +} + +fn flatten_dictionary_haystack(in_array: ArrayRef) -> Result { // Flatten dictionary-encoded haystacks to their value type so that // specialized filters (e.g. Int32StaticFilter) are used instead of // falling through to the generic ArrayStaticFilter. - let in_array = match in_array.data_type() { - DataType::Dictionary(_, value_type) => cast(&in_array, value_type.as_ref())?, - _ => in_array, - }; match in_array.data_type() { - // Integer primitive types - DataType::Int8 => Ok(Arc::new(Int8StaticFilter::try_new(&in_array)?)), - DataType::Int16 => Ok(Arc::new(Int16StaticFilter::try_new(&in_array)?)), + DataType::Dictionary(_, value_type) => Ok(cast(&in_array, value_type.as_ref())?), + _ => Ok(in_array), + } +} + +fn instantiate_branchless_filter(in_array: &ArrayRef) -> Result> { + let non_null_count = in_array.len() - in_array.null_count(); + + macro_rules! filter { + ($arrow_type:ty) => { + branchless_filter::<$arrow_type>(in_array, non_null_count) + }; + } + + match in_array.data_type() { + DataType::Int8 => filter!(Int8Type), + DataType::UInt8 => filter!(UInt8Type), + DataType::Int16 => filter!(Int16Type), + DataType::UInt16 => filter!(UInt16Type), + DataType::Float16 => filter!(Float16Type), + DataType::Int32 => filter!(Int32Type), + DataType::UInt32 => filter!(UInt32Type), + DataType::Float32 => filter!(Float32Type), + DataType::Date32 => filter!(Date32Type), + DataType::Time32(unit) => match unit { + TimeUnit::Second => filter!(Time32SecondType), + TimeUnit::Millisecond => filter!(Time32MillisecondType), + _ => Ok(None), + }, + DataType::Int64 => filter!(Int64Type), + DataType::UInt64 => filter!(UInt64Type), + DataType::Float64 => filter!(Float64Type), + DataType::Date64 => filter!(Date64Type), + DataType::Time64(unit) => match unit { + TimeUnit::Microsecond => filter!(Time64MicrosecondType), + TimeUnit::Nanosecond => filter!(Time64NanosecondType), + _ => Ok(None), + }, + DataType::Timestamp(unit, _) => match unit { + TimeUnit::Second => filter!(TimestampSecondType), + TimeUnit::Millisecond => filter!(TimestampMillisecondType), + TimeUnit::Microsecond => filter!(TimestampMicrosecondType), + TimeUnit::Nanosecond => filter!(TimestampNanosecondType), + }, + DataType::Duration(unit) => match unit { + TimeUnit::Second => filter!(DurationSecondType), + TimeUnit::Millisecond => filter!(DurationMillisecondType), + TimeUnit::Microsecond => filter!(DurationMicrosecondType), + TimeUnit::Nanosecond => filter!(DurationNanosecondType), + }, + DataType::Decimal128(_, _) => filter!(Decimal128Type), + DataType::Interval(IntervalUnit::MonthDayNano) => { + filter!(IntervalMonthDayNanoType) + } + _ => Ok(None), + } +} + +fn instantiate_standard_filter(in_array: ArrayRef) -> Result { + match in_array.data_type() { + DataType::Int8 => bitmap_filter::(&in_array), + DataType::UInt8 => bitmap_filter::(&in_array), + DataType::Int16 => bitmap_filter::(&in_array), + DataType::UInt16 => bitmap_filter::(&in_array), + DataType::Float16 => bitmap_filter::(&in_array), DataType::Int32 => Ok(Arc::new(Int32StaticFilter::try_new(&in_array)?)), DataType::Int64 => Ok(Arc::new(Int64StaticFilter::try_new(&in_array)?)), - DataType::UInt8 => Ok(Arc::new(UInt8StaticFilter::try_new(&in_array)?)), - DataType::UInt16 => Ok(Arc::new(UInt16StaticFilter::try_new(&in_array)?)), DataType::UInt32 => Ok(Arc::new(UInt32StaticFilter::try_new(&in_array)?)), DataType::UInt64 => Ok(Arc::new(UInt64StaticFilter::try_new(&in_array)?)), // Float primitive types (use ordered wrappers for Hash/Eq) DataType::Float32 => Ok(Arc::new(Float32StaticFilter::try_new(&in_array)?)), DataType::Float64 => Ok(Arc::new(Float64StaticFilter::try_new(&in_array)?)), _ => { - /* fall through to generic implementation for unsupported types (Struct, etc.) */ + // Fall through to generic implementation for unsupported types + // (Struct, etc.). Ok(Arc::new(ArrayStaticFilter::try_new(in_array)?)) } } } + +fn bitmap_filter(in_array: &ArrayRef) -> Result +where + T: BitmapFilterType, +{ + Ok(Arc::new(BitmapFilter::::try_new(in_array)?)) +} + +fn branchless_filter( + in_array: &ArrayRef, + non_null_count: usize, +) -> Result> +where + T: BranchlessFilterType, + BranchlessNative: Copy + PartialEq + Send + Sync, +{ + // Larger lists use the standard filter. `try_new` checks the limit again. + if non_null_count > T::MAX_LIST_LEN { + return Ok(None); + } + + Ok(Some(Arc::new(BranchlessFilter::::try_new(in_array)?))) +} + +#[cfg(test)] +mod tests { + use arrow::array::UInt32Array; + use arrow::datatypes::UInt32Type; + + use super::super::branchless_filter::BranchlessFilterType; + use super::*; + + fn uint32_array(values: Vec>) -> ArrayRef { + Arc::new(UInt32Array::from(values)) + } + + #[test] + fn branchless_routing_respects_max_list_len() -> Result<()> { + let max_len = ::MAX_LIST_LEN; + + let values = (0..max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!(instantiate_branchless_filter(&uint32_array(values))?.is_some()); + + let values = (0..=max_len) + .map(|value| Some(value as u32)) + .collect::>(); + assert!(instantiate_branchless_filter(&uint32_array(values))?.is_none()); + + Ok(()) + } + + #[test] + fn branchless_routing_handles_zero_non_null_values() -> Result<()> { + let array = uint32_array(vec![None; 3]); + + assert!(instantiate_branchless_filter(&array)?.is_some()); + + Ok(()) + } +} diff --git a/datafusion/physical-expr/src/expressions/is_not_null.rs b/datafusion/physical-expr/src/expressions/is_not_null.rs index 86acf0a4ea116..3f3b7d16e543a 100644 --- a/datafusion/physical-expr/src/expressions/is_not_null.rs +++ b/datafusion/physical-expr/src/expressions/is_not_null.rs @@ -22,8 +22,7 @@ use arrow::{ datatypes::{DataType, Schema}, record_batch::RecordBatch, }; -use datafusion_common::Result; -use datafusion_common::ScalarValue; +use datafusion_common::{Result, ScalarValue}; use datafusion_expr::ColumnarValue; use std::hash::Hash; use std::sync::Arc; @@ -103,6 +102,48 @@ impl PhysicalExpr for IsNotNullExpr { self.arg.fmt_sql(f)?; write!(f, " IS NOT NULL") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr( + Box::new(protobuf::PhysicalIsNotNull { + expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + }), + )), + })) + } +} + +#[cfg(feature = "proto")] +impl IsNotNullExpr { + /// Reconstruct an [`IsNotNullExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let node = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::IsNotNullExpr, + "IsNotNullExpr", + ); + let expr = ctx.decode_required_expression( + node.expr.as_deref(), + "IsNotNullExpr", + "expr", + )?; + + Ok(Arc::new(IsNotNullExpr::new(expr))) + } } /// Create an IS NOT NULL expression @@ -213,3 +254,109 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalIsNotNull, physical_expr_node, + }; + + fn is_not_null_node(expr: Option>) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::IsNotNullExpr(Box::new( + PhysicalIsNotNull { expr }, + ))), + } + } + + fn is_not_null_fixture() -> IsNotNullExpr { + let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]); + IsNotNullExpr::new(col("a", &schema).unwrap()) + } + + #[test] + fn try_to_proto_encodes_is_not_null_expr() { + let is_not_null = is_not_null_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = is_not_null + .try_to_proto(&ctx) + .unwrap() + .expect("IsNotNullExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let is_not_null_node = match node.expr_type { + Some(physical_expr_node::ExprType::IsNotNullExpr(boxed)) => *boxed, + other => panic!("expected an IsNotNullExpr node, got {other:?}"), + }; + assert!(is_not_null_node.expr.is_some()); + } + + #[test] + fn try_to_proto_propagates_expr_encode_error() { + let is_not_null = is_not_null_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = is_not_null.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_from_proto_decodes_is_not_null_expr() { + let node = is_not_null_node(Some(Box::new(column_node("a")))); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap(); + let is_not_null = decoded + .downcast_ref::() + .expect("decoded expr should be an IsNotNullExpr"); + assert!(is_not_null.arg().downcast_ref::().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_is_not_null_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a IsNotNullExpr")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_expr() { + let node = is_not_null_node(None); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("IsNotNullExpr is missing required field 'expr'")) + ); + } + + #[test] + fn try_from_proto_propagates_expr_decode_error() { + let node = is_not_null_node(Some(Box::new(column_node("a")))); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = IsNotNullExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } +} diff --git a/datafusion/physical-expr/src/expressions/is_null.rs b/datafusion/physical-expr/src/expressions/is_null.rs index 8534ddb8d104f..da008a1cfb821 100644 --- a/datafusion/physical-expr/src/expressions/is_null.rs +++ b/datafusion/physical-expr/src/expressions/is_null.rs @@ -22,8 +22,7 @@ use arrow::{ datatypes::{DataType, Schema}, record_batch::RecordBatch, }; -use datafusion_common::Result; -use datafusion_common::ScalarValue; +use datafusion_common::{Result, ScalarValue}; use datafusion_expr::ColumnarValue; use std::hash::Hash; use std::sync::Arc; @@ -102,6 +101,45 @@ impl PhysicalExpr for IsNullExpr { self.arg.fmt_sql(f)?; write!(f, " IS NULL") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::IsNullExpr( + Box::new(protobuf::PhysicalIsNull { + expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + }), + )), + })) + } +} + +#[cfg(feature = "proto")] +impl IsNullExpr { + /// Reconstruct an [`IsNullExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let node = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::IsNullExpr, + "IsNullExpr", + ); + let expr = + ctx.decode_required_expression(node.expr.as_deref(), "IsNullExpr", "expr")?; + + Ok(Arc::new(IsNullExpr::new(expr))) + } } /// Create an IS NULL expression @@ -224,3 +262,109 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalIsNull, physical_expr_node, + }; + + fn is_null_node(expr: Option>) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::IsNullExpr(Box::new( + PhysicalIsNull { expr }, + ))), + } + } + + fn is_null_fixture() -> IsNullExpr { + let schema = Schema::new(vec![Field::new("a", DataType::Int64, true)]); + IsNullExpr::new(col("a", &schema).unwrap()) + } + + #[test] + fn try_to_proto_encodes_is_null_expr() { + let is_null = is_null_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = is_null + .try_to_proto(&ctx) + .unwrap() + .expect("IsNullExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let is_null_node = match node.expr_type { + Some(physical_expr_node::ExprType::IsNullExpr(boxed)) => *boxed, + other => panic!("expected an IsNullExpr node, got {other:?}"), + }; + assert!(is_null_node.expr.is_some()); + } + + #[test] + fn try_to_proto_propagates_expr_encode_error() { + let is_null = is_null_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = is_null.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_from_proto_decodes_is_null_expr() { + let node = is_null_node(Some(Box::new(column_node("a")))); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = IsNullExpr::try_from_proto(&node, &ctx).unwrap(); + let is_null = decoded + .downcast_ref::() + .expect("decoded expr should be an IsNullExpr"); + assert!(is_null.arg().downcast_ref::().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_is_null_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = IsNullExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a IsNullExpr")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_expr() { + let node = is_null_node(None); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = IsNullExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("IsNullExpr is missing required field 'expr'")) + ); + } + + #[test] + fn try_from_proto_propagates_expr_decode_error() { + let node = is_null_node(Some(Box::new(column_node("a")))); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = IsNullExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } +} diff --git a/datafusion/physical-expr/src/expressions/lambda.rs b/datafusion/physical-expr/src/expressions/lambda.rs index 9275821ae9150..95bb5db0b328e 100644 --- a/datafusion/physical-expr/src/expressions/lambda.rs +++ b/datafusion/physical-expr/src/expressions/lambda.rs @@ -31,7 +31,7 @@ use arrow::{ }; use datafusion_common::{ HashMap, plan_err, - tree_node::{Transformed, TreeNode, TreeNodeRecursion}, + tree_node::{Transformed, TreeNode, TreeNodeRecursion, TreeNodeVisitor}, }; use datafusion_common::{HashSet, Result, internal_err}; use datafusion_expr::ColumnarValue; @@ -43,6 +43,7 @@ pub struct LambdaExpr { body: Arc, projected_body: Arc, projection: Vec, + used_param_indices: Vec, } // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 [https://github.com/apache/datafusion/issues/13196] @@ -60,7 +61,7 @@ impl Hash for LambdaExpr { } impl LambdaExpr { - /// Create a new lambda expression with the given parameters and body + /// Create a new lambda expression with the given parameters and body. pub fn try_new(params: Vec, body: Arc) -> Result { if !all_unique(¶ms) { return plan_err!( @@ -75,27 +76,30 @@ impl LambdaExpr { } fn new(params: Vec, body: Arc) -> Self { - let mut used_column_indices = HashSet::new(); + let own_params: HashSet = params.iter().cloned().collect(); - body.apply(|node| { - if let Some(col) = node.downcast_ref::() { - used_column_indices.insert(col.index()); - } else if let Some(var) = node.downcast_ref::() { - used_column_indices.insert(var.index()); - } - - Ok(TreeNodeRecursion::Continue) - }) - .expect("closure should be infallible"); + let mut visitor = CollectUsedVisitor { + own_params: &own_params, + used_indices: HashSet::new(), + used_param_names: HashSet::new(), + shadow_stack: Vec::new(), + }; + body.visit(&mut visitor).expect("visitor is infallible"); + let CollectUsedVisitor { + used_indices, + used_param_names, + .. + } = visitor; - let mut projection = used_column_indices.into_iter().collect::>(); + let mut projection = used_indices.into_iter().collect::>(); projection.sort(); let column_index_map = projection .iter() + .copied() .enumerate() - .map(|(projected, original)| (*original, projected)) + .map(|(new_idx, original)| (original, new_idx)) .collect::>(); let projected_body = Arc::clone(&body) @@ -124,11 +128,19 @@ impl LambdaExpr { .expect("closure should be infallible") .data; + let used_param_indices = params + .iter() + .enumerate() + .filter(|(_, name)| used_param_names.contains(*name)) + .map(|(i, _)| i) + .collect(); + Self { params, body, projected_body, projection, + used_param_indices, } } @@ -142,6 +154,27 @@ impl LambdaExpr { &self.body } + #[cfg(feature = "proto")] + /// Reconstruct a [`LambdaExpr`] from a proto node. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let lambda = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Lambda, + "LambdaExpr", + ); + + Ok(Arc::new(LambdaExpr::try_new( + lambda.params.clone(), + ctx.decode_required_expression(lambda.body.as_deref(), "LambdaExpr", "body")?, + )?)) + } + pub(crate) fn projection(&self) -> &[usize] { &self.projection } @@ -149,6 +182,75 @@ impl LambdaExpr { pub(crate) fn projected_body(&self) -> &Arc { &self.projected_body } + + /// Indices into [`params`](Self::params) of the parameters the body + /// actually references, in declaration order. See `CollectUsedVisitor` + /// in this module. + /// + /// Relies on the planner appending each lambda's own params after + /// captures, matching the `captures ++ used_params` layout + /// `LambdaArgument::new` builds. + pub fn used_param_indices(&self) -> &[usize] { + &self.used_param_indices + } +} + +/// Walks the body of a [`LambdaExpr`] and collects, on a single pass: +/// +/// * `used_indices` — every `Column` / `LambdaVariable` index referenced +/// anywhere in the tree (including inside nested lambdas). This drives +/// the `projection` used to slice the outer batch. +/// * `used_param_names` — the subset of *this* lambda's `own_params` that +/// the body actually references. +/// +/// A nested lambda can declare its own parameter with the same name as +/// one of `own_params` — a distinct variable that happens to reuse the +/// name (variable shadowing). E.g. in +/// `(k, v) -> func(col, (k, v2) -> k + v2 + v)`, the inner `k` is not +/// `own_params`' `k`; only `v` should flow up as used, not `k`. +/// +/// `shadow_stack` holds one frame per nested `LambdaExpr` currently being +/// visited, each frame being that lambda's own parameter names. A +/// `LambdaVariable` only counts toward `used_param_names` if its name +/// isn't in any active frame (i.e. not shadowed). +/// +/// The stack is maintained via `TreeNodeVisitor`'s `f_down` / `f_up`: +/// push a frame when entering a nested [`LambdaExpr`], pop it when leaving. +struct CollectUsedVisitor<'a> { + own_params: &'a HashSet, + used_indices: HashSet, + used_param_names: HashSet, + shadow_stack: Vec>, +} + +impl TreeNodeVisitor<'_> for CollectUsedVisitor<'_> { + type Node = Arc; + + fn f_down(&mut self, node: &Self::Node) -> Result { + if let Some(col) = node.downcast_ref::() { + self.used_indices.insert(col.index()); + } else if let Some(var) = node.downcast_ref::() { + self.used_indices.insert(var.index()); + + let name = var.name(); + let shadowed = self.shadow_stack.iter().any(|frame| frame.contains(name)); + if !shadowed && self.own_params.contains(name) { + self.used_param_names.insert(name.to_string()); + } + } else if let Some(nested) = node.downcast_ref::() { + self.shadow_stack + .push(nested.params.iter().cloned().collect()); + } + + Ok(TreeNodeRecursion::Continue) + } + + fn f_up(&mut self, node: &Self::Node) -> Result { + if node.downcast_ref::().is_some() { + self.shadow_stack.pop(); + } + Ok(TreeNodeRecursion::Continue) + } } impl std::fmt::Display for LambdaExpr { @@ -193,9 +295,27 @@ impl PhysicalExpr for LambdaExpr { fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "({}) -> {}", self.params.join(", "), self.body) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Lambda(Box::new( + protobuf::PhysicalLambdaExprNode { + params: self.params().to_vec(), + body: Some(Box::new(ctx.encode_child(self.body())?)), + }, + ))), + })) + } } -/// Create a lambda expression +/// Create a lambda expression. pub fn lambda( params: impl IntoIterator>, body: Arc, @@ -234,10 +354,15 @@ fn check_async_udf(body: &Arc) -> Result<()> { #[cfg(test)] mod tests { - use crate::expressions::{NoOp, lambda::lambda}; - use arrow::{array::RecordBatch, datatypes::Schema}; + use crate::expressions::{Column, LambdaVariable, NoOp, lambda::lambda}; + use arrow::{ + array::RecordBatch, + datatypes::{DataType, Field, Schema}, + }; use std::sync::Arc; + use super::LambdaExpr; + #[test] fn test_lambda_evaluate() { let lambda = lambda(["a"], Arc::new(NoOp::new())).unwrap(); @@ -249,4 +374,125 @@ mod tests { fn test_lambda_duplicate_name() { assert!(lambda(["a", "a"], Arc::new(NoOp::new())).is_err()); } + + /// A two-parameter lambda whose body only references the second + /// parameter (`v`) must report only `v` as used. The higher-order + /// function uses this set to push only `v` into the merged batch, so + /// the body's compressed `LambdaVariable` index for `v` lines up with + /// the batch layout. + #[test] + fn test_used_params_collects_only_referenced_param() { + let v_field = Arc::new(Field::new("v", DataType::Int32, true)); + let body = Arc::new(LambdaVariable::new(1, Arc::clone(&v_field))); + + let lambda = + LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); + + assert_eq!(lambda.projection(), &[1]); + assert_eq!(lambda.used_param_indices(), &[1]); + } + + /// A body that references neither declared parameter reports no used params. + #[test] + fn test_used_params_all_unused() { + let body = Arc::new(NoOp::new()); + + let lambda = + LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); + + assert!(lambda.projection().is_empty()); + assert!(lambda.used_param_indices().is_empty()); + } + + /// A three-parameter lambda that skips the middle parameter reports only the ends as used. + #[test] + fn test_used_params_three_params_middle_unused() { + let a_field = Arc::new(Field::new("a", DataType::Int32, true)); + let c_field = Arc::new(Field::new("c", DataType::Int32, true)); + let body = Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(LambdaVariable::new(0, Arc::clone(&a_field))), + datafusion_expr::Operator::Plus, + Arc::new(LambdaVariable::new(2, Arc::clone(&c_field))), + )); + + let lambda = LambdaExpr::try_new( + vec!["a".to_string(), "b".to_string(), "c".to_string()], + body, + ) + .unwrap(); + + assert_eq!(lambda.used_param_indices(), &[0, 2]); + } + + /// Referencing params out of declaration order still reports both as used. + #[test] + fn test_used_params_both_used_in_reverse_reference_order() { + let k_field = Arc::new(Field::new("k", DataType::Int32, true)); + let v_field = Arc::new(Field::new("v", DataType::Int32, true)); + let body = Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(LambdaVariable::new(1, Arc::clone(&v_field))), + datafusion_expr::Operator::Plus, + Arc::new(LambdaVariable::new(0, Arc::clone(&k_field))), + )); + + let lambda = + LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], body).unwrap(); + + assert_eq!(lambda.projection(), &[0, 1]); + assert_eq!(lambda.used_param_indices(), &[0, 1]); + } + + /// Inside a nested lambda that re-declares one of the outer parameter + /// names, only the non-shadowed outer references should be reported as + /// used by the outer lambda. In + /// `(k, v) -> func(col, (k, v2) -> k + v2 + v)` the inner `k` shadows + /// the outer `k`, so the outer lambda must only see `v` as used. + #[test] + fn test_used_params_handles_shadowing_inside_nested_lambda() { + let outer_k_field = Arc::new(Field::new("k", DataType::Int32, true)); + let outer_v_field = Arc::new(Field::new("v", DataType::Int32, true)); + let inner_v2_field = Arc::new(Field::new("v2", DataType::Int32, true)); + + // Inner lambda body references "k" (inner's), "v2" (inner's), and + // "v" (outer's). Build it directly with the dense compressed + // indices the inner LambdaExpr::new would produce: sorted referenced + // indices, so the names alone matter here — what matters for + // shadow tracking is the names, not the indices. + let inner_body: Arc = + Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(LambdaVariable::new(1, Arc::clone(&outer_k_field))), + datafusion_expr::Operator::Plus, + Arc::new(LambdaVariable::new(2, Arc::clone(&inner_v2_field))), + )), + datafusion_expr::Operator::Plus, + Arc::new(LambdaVariable::new(0, Arc::clone(&outer_v_field))), + )); + let inner_lambda = Arc::new( + LambdaExpr::try_new(vec!["k".to_string(), "v2".to_string()], inner_body) + .unwrap(), + ); + + // Outer body wraps the inner lambda in a binary op next to a + // regular column reference so the walk has something non-trivial + // to descend through. The outer body references the inner lambda + // via `inner_lambda`. + let outer_body: Arc = + Arc::new(crate::expressions::BinaryExpr::new( + Arc::new(Column::new("col", 0)), + datafusion_expr::Operator::Plus, + inner_lambda, + )); + + let outer_lambda = + LambdaExpr::try_new(vec!["k".to_string(), "v".to_string()], outer_body) + .unwrap(); + + assert_eq!( + outer_lambda.used_param_indices(), + &[1], + "only outer's `v` (index 1) should be reported as used; `k` (index 0) is \ + shadowed inside the nested lambda" + ); + } } diff --git a/datafusion/physical-expr/src/expressions/lambda_variable.rs b/datafusion/physical-expr/src/expressions/lambda_variable.rs index 1c130ab12e9bb..f7e69100208a3 100644 --- a/datafusion/physical-expr/src/expressions/lambda_variable.rs +++ b/datafusion/physical-expr/src/expressions/lambda_variable.rs @@ -72,6 +72,32 @@ impl LambdaVariable { pub fn field(&self) -> &FieldRef { &self.field } + + #[cfg(feature = "proto")] + /// Reconstruct a [`LambdaVariable`] from a proto node. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::{ + expect_expr_variant, physical_expr::proto_decode::require_proto_field, + }; + use datafusion_proto_models::protobuf; + + let var = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::LambdaVariable, + "LambdaVariable", + ); + + Ok(Arc::new(LambdaVariable::new( + var.index as usize, + Arc::new( + require_proto_field(var.field.as_ref(), "LambdaVariable", "field")? + .try_into()?, + ), + ))) + } } impl std::fmt::Display for LambdaVariable { @@ -135,6 +161,24 @@ impl PhysicalExpr for LambdaVariable { fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}@{}", self.name(), self.index) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::LambdaVariable( + protobuf::PhysicalLambdaVariableExprNode { + index: self.index() as u32, + field: Some(self.field().as_ref().try_into()?), + }, + )), + })) + } } /// Create a lambda variable expression diff --git a/datafusion/physical-expr/src/expressions/like.rs b/datafusion/physical-expr/src/expressions/like.rs index 07ceb4e7d7d49..7535f109a0a92 100644 --- a/datafusion/physical-expr/src/expressions/like.rs +++ b/datafusion/physical-expr/src/expressions/like.rs @@ -145,6 +145,65 @@ impl PhysicalExpr for LikeExpr { write!(f, " {} ", self.op_name())?; self.pattern.fmt_sql(f) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::LikeExpr(Box::new( + protobuf::PhysicalLikeExprNode { + negated: self.negated, + case_insensitive: self.case_insensitive, + expr: Some(Box::new(ctx.encode_child(&self.expr)?)), + pattern: Some(Box::new(ctx.encode_child(&self.pattern)?)), + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl LikeExpr { + /// Reconstruct a [`LikeExpr`] from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`] so the decode signature matches + /// other migrated expressions and can inspect outer-node metadata if + /// needed in the future. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let like_expr = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::LikeExpr, + "LikeExpr", + ); + + Ok(Arc::new(LikeExpr::new( + like_expr.negated, + like_expr.case_insensitive, + ctx.decode_required_expression( + like_expr.expr.as_deref(), + "LikeExpr", + "expr", + )?, + ctx.decode_required_expression( + like_expr.pattern.as_deref(), + "LikeExpr", + "pattern", + )?, + ))) + } } /// used for optimize Dictionary like @@ -283,3 +342,189 @@ mod test { Ok(()) } } + +/// Tests for the `try_to_proto` / `try_from_proto` hooks. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalLikeExprNode, physical_expr_node, + }; + + /// Build a `LikeExpr` proto node with the given children. + fn like_node( + negated: bool, + case_insensitive: bool, + expr: Option>, + pattern: Option>, + ) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::LikeExpr(Box::new( + PhysicalLikeExprNode { + negated, + case_insensitive, + expr, + pattern, + }, + ))), + } + } + + /// A `LikeExpr` over two `Utf8` columns with both flags set, so the + /// `negated` / `case_insensitive` wiring is actually exercised. + fn like_fixture() -> LikeExpr { + let schema = Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Utf8, false), + ]); + LikeExpr::new( + true, + true, + col("a", &schema).unwrap(), + col("b", &schema).unwrap(), + ) + } + + #[test] + fn try_to_proto_encodes_like_expr() { + let like = like_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = like + .try_to_proto(&ctx) + .unwrap() + .expect("LikeExpr should encode to Some(node)"); + + // Built-in exprs never set expr_id; only dynamic filters do. + assert!(node.expr_id.is_none()); + let like_node = match node.expr_type { + Some(physical_expr_node::ExprType::LikeExpr(boxed)) => *boxed, + other => panic!("expected a LikeExpr node, got {other:?}"), + }; + assert!(like_node.negated); + assert!(like_node.case_insensitive); + assert!(like_node.expr.is_some()); + assert!(like_node.pattern.is_some()); + } + + #[test] + fn try_to_proto_propagates_expr_encode_error() { + let like = like_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = like.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_to_proto_propagates_pattern_encode_error() { + let like = like_fixture(); + let encoder = StubEncoder::failing_on(2); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = like.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2"))); + } + + #[test] + fn try_from_proto_decodes_like_expr() { + let node = like_node( + true, + true, + Some(Box::new(column_node("a"))), + Some(Box::new(column_node("b"))), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = LikeExpr::try_from_proto(&node, &ctx).unwrap(); + let like = decoded + .downcast_ref::() + .expect("decoded expr should be a LikeExpr"); + assert!(like.negated()); + assert!(like.case_insensitive()); + assert!(like.expr().downcast_ref::().is_some()); + assert!(like.pattern().downcast_ref::().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_like_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = LikeExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a LikeExpr") + )); + } + + #[test] + fn try_from_proto_rejects_missing_expr() { + let node = like_node(false, false, None, Some(Box::new(column_node("b")))); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = LikeExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("LikeExpr is missing required field 'expr'") + )); + } + + #[test] + fn try_from_proto_rejects_missing_pattern() { + let node = like_node(false, false, Some(Box::new(column_node("a"))), None); + let schema = Schema::empty(); + // `expr` is present, so it is decoded before the missing-`pattern` + // check fires; use a decoder that succeeds for that first child. + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = LikeExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) if msg.contains("LikeExpr is missing required field 'pattern'") + )); + } + + #[test] + fn try_from_proto_propagates_expr_decode_error() { + let node = like_node( + false, + false, + Some(Box::new(column_node("a"))), + Some(Box::new(column_node("b"))), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = LikeExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_from_proto_propagates_pattern_decode_error() { + let node = like_node( + false, + false, + Some(Box::new(column_node("a"))), + Some(Box::new(column_node("b"))), + ); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(2); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = LikeExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 2"))); + } +} diff --git a/datafusion/physical-expr/src/expressions/literal.rs b/datafusion/physical-expr/src/expressions/literal.rs index 7351158c54e31..a7af824230780 100644 --- a/datafusion/physical-expr/src/expressions/literal.rs +++ b/datafusion/physical-expr/src/expressions/literal.rs @@ -123,6 +123,8 @@ impl PhysicalExpr for Literal { sort_properties: SortProperties::Singleton, range: Interval::try_new(self.value().clone(), self.value().clone())?, preserves_lex_ordering: true, + // Vacuously true: a literal has no ordered inputs. + strictly_order_preserving: true, }) } @@ -133,6 +135,41 @@ impl PhysicalExpr for Literal { fn placement(&self) -> ExpressionPlacement { ExpressionPlacement::Literal } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Literal( + (&self.value).try_into()?, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl Literal { + /// Reconstruct a [`Literal`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let scalar_proto = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Literal, + "Literal", + ); + let value = ScalarValue::try_from(scalar_proto)?; + Ok(Arc::new(Literal::new(value))) + } } /// Create a literal expression @@ -190,3 +227,103 @@ mod tests { Ok(()) } } + +/// Tests for the `try_to_proto` / `try_from_proto` hooks. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::proto_test_util::{StubEncoder, UnreachableDecoder, column_node}; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::physical_expr_node; + + fn i32_literal() -> Literal { + Literal::new(ScalarValue::Int32(Some(42))) + } + + // ── try_to_proto ───────────────────────────────────────────────────────── + + #[test] + fn try_to_proto_encodes_literal() { + let literal = i32_literal(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = literal + .try_to_proto(&ctx) + .unwrap() + .expect("Literal should encode to Some(node)"); + + // Literal nodes never set expr_id. + assert!(node.expr_id.is_none()); + // Variant must be Literal, not any other expr type. + assert!(matches!( + node.expr_type, + Some(physical_expr_node::ExprType::Literal(_)) + )); + } + + #[test] + fn try_to_proto_null_literal() { + let literal = Literal::new(ScalarValue::Int32(None)); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = literal + .try_to_proto(&ctx) + .unwrap() + .expect("null Literal should encode to Some(node)"); + + assert!(matches!( + node.expr_type, + Some(physical_expr_node::ExprType::Literal(_)) + )); + + // Decode and verify the null payload round-trips correctly. + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let decoded = Literal::try_from_proto(&node, &dec_ctx).unwrap(); + let lit = decoded + .downcast_ref::() + .expect("decoded expr should be a Literal"); + assert_eq!(lit.value(), &ScalarValue::Int32(None)); + } + + // ── try_from_proto ─────────────────────────────────────────────────────── + + #[test] + fn try_from_proto_roundtrip() { + let original = i32_literal(); + let encoder = StubEncoder::ok(); + let enc_ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = original + .try_to_proto(&enc_ctx) + .unwrap() + .expect("should encode"); + + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = Literal::try_from_proto(&node, &dec_ctx).unwrap(); + let lit = decoded + .downcast_ref::() + .expect("decoded expr should be a Literal"); + assert_eq!(lit.value(), &ScalarValue::Int32(Some(42))); + } + + #[test] + fn try_from_proto_rejects_non_literal_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = Literal::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(ref msg) if msg.contains("PhysicalExprNode is not a Literal")) + ); + } +} diff --git a/datafusion/physical-expr/src/expressions/mod.rs b/datafusion/physical-expr/src/expressions/mod.rs index 7cf874c448ea0..035dd5d5072b0 100644 --- a/datafusion/physical-expr/src/expressions/mod.rs +++ b/datafusion/physical-expr/src/expressions/mod.rs @@ -45,7 +45,9 @@ pub use case::{CaseExpr, case}; pub use cast::{CastExpr, cast}; pub use column::{Column, col, with_new_schema}; pub use datafusion_expr::utils::format_state_name; -pub use dynamic_filters::{DynamicFilterPhysicalExpr, Inner as DynamicFilterInner}; +pub use dynamic_filters::{ + DynamicFilterPhysicalExpr, DynamicFilterTracker, DynamicFilterTracking, +}; pub use in_list::{InListExpr, in_list}; pub use is_not_null::{IsNotNullExpr, is_not_null}; pub use is_null::{IsNullExpr, is_null}; diff --git a/datafusion/physical-expr/src/expressions/negative.rs b/datafusion/physical-expr/src/expressions/negative.rs index e2bda4c8aaf49..c894c12784dc5 100644 --- a/datafusion/physical-expr/src/expressions/negative.rs +++ b/datafusion/physical-expr/src/expressions/negative.rs @@ -166,6 +166,8 @@ impl PhysicalExpr for NegativeExpr { sort_properties: -children[0].sort_properties, range: children[0].range.clone().arithmetic_negate()?, preserves_lex_ordering: false, + // Negation is one-to-one but reverses the ordering direction. + strictly_order_preserving: false, }) } @@ -174,6 +176,45 @@ impl PhysicalExpr for NegativeExpr { self.arg.fmt_sql(f)?; write!(f, ")") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Negative(Box::new( + protobuf::PhysicalNegativeNode { + expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl NegativeExpr { + /// Reconstruct a [`NegativeExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let n = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::Negative, + "Negative", + ); + let expr = + ctx.decode_required_expression(n.expr.as_deref(), "NegativeExpr", "expr")?; + + Ok(Arc::new(NegativeExpr::new(expr))) + } } /// Creates a unary expression NEGATIVE @@ -402,3 +443,111 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalNegativeNode, physical_expr_node, + }; + + /// Build a `NegativeExpr` proto node with the given children. + fn negative_node(expr: Option>) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::Negative(Box::new( + PhysicalNegativeNode { expr }, + ))), + } + } + + /// A `NegativeExpr` over a column of type Int32. + fn negative_fixture() -> NegativeExpr { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + NegativeExpr::new(col("a", &schema).unwrap()) + } + + #[test] + fn try_to_proto_encodes_negative_expr() { + let negative = negative_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = negative + .try_to_proto(&ctx) + .unwrap() + .expect("NegativeExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let negative_node = match node.expr_type { + Some(physical_expr_node::ExprType::Negative(boxed)) => *boxed, + other => panic!("expected a NegativeExpr node, got {other:?}"), + }; + assert!(negative_node.expr.is_some()); + } + + #[test] + fn try_to_proto_propagates_expr_encode_error() { + let negative = negative_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = negative.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_from_proto_decodes_negative_expr() { + let node = negative_node(Some(Box::new(column_node("a")))); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = NegativeExpr::try_from_proto(&node, &ctx).unwrap(); + let negative = decoded + .downcast_ref::() + .expect("decoded expr should be a NegativeExpr"); + assert!(negative.arg().downcast_ref::().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_negative_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = NegativeExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a Negative")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_expr() { + let node = negative_node(None); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = NegativeExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("NegativeExpr is missing required field 'expr'")) + ); + } + + #[test] + fn try_from_proto_propagates_expr_decode_error() { + let node = negative_node(Some(Box::new(column_node("a")))); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = NegativeExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } +} diff --git a/datafusion/physical-expr/src/expressions/not.rs b/datafusion/physical-expr/src/expressions/not.rs index b63effdbb9c88..f856dd568a8da 100644 --- a/datafusion/physical-expr/src/expressions/not.rs +++ b/datafusion/physical-expr/src/expressions/not.rs @@ -181,6 +181,45 @@ impl PhysicalExpr for NotExpr { write!(f, "NOT ")?; self.arg.fmt_sql(f) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::NotExpr(Box::new( + protobuf::PhysicalNot { + expr: Some(Box::new(ctx.encode_child(&self.arg)?)), + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl NotExpr { + /// Reconstruct a [`NotExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let not_expr = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::NotExpr, + "NotExpr", + ); + let expr = + ctx.decode_required_expression(not_expr.expr.as_deref(), "NotExpr", "expr")?; + + Ok(Arc::new(NotExpr::new(expr))) + } } /// Creates a unary expression NOT @@ -357,3 +396,112 @@ mod tests { Arc::clone(&SCHEMA) } } + +/// Tests for the `try_to_proto` / `try_from_proto` hooks. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalNot, physical_expr_node, + }; + + /// Build a `NotExpr` proto node with the given child. + fn not_node(expr: Option>) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::NotExpr(Box::new( + PhysicalNot { expr }, + ))), + } + } + + /// A `NotExpr` over a boolean column. + fn not_fixture() -> NotExpr { + let schema = Schema::new(vec![Field::new("a", DataType::Boolean, true)]); + NotExpr::new(col("a", &schema).unwrap()) + } + + #[test] + fn try_to_proto_encodes_not_expr() { + let not = not_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = not + .try_to_proto(&ctx) + .unwrap() + .expect("NotExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let not_node = match node.expr_type { + Some(physical_expr_node::ExprType::NotExpr(boxed)) => *boxed, + other => panic!("expected a NotExpr node, got {other:?}"), + }; + assert!(not_node.expr.is_some()); + } + + #[test] + fn try_to_proto_propagates_expr_encode_error() { + let not = not_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = not.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_from_proto_decodes_not_expr() { + let node = not_node(Some(Box::new(column_node("a")))); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = NotExpr::try_from_proto(&node, &ctx).unwrap(); + let not = decoded + .downcast_ref::() + .expect("decoded expr should be a NotExpr"); + assert!(not.arg().downcast_ref::().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_not_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = NotExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a NotExpr")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_expr() { + let node = not_node(None); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = NotExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("NotExpr is missing required field 'expr'")) + ); + } + + #[test] + fn try_from_proto_propagates_expr_decode_error() { + let node = not_node(Some(Box::new(column_node("a")))); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = NotExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } +} diff --git a/datafusion/physical-expr/src/expressions/try_cast.rs b/datafusion/physical-expr/src/expressions/try_cast.rs index ba59d113acaab..65b953fd181b7 100644 --- a/datafusion/physical-expr/src/expressions/try_cast.rs +++ b/datafusion/physical-expr/src/expressions/try_cast.rs @@ -119,6 +119,56 @@ impl PhysicalExpr for TryCastExpr { self.expr.fmt_sql(f)?; write!(f, " AS {:?})", self.cast_type) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::TryCast(Box::new( + protobuf::PhysicalTryCastNode { + expr: Some(Box::new(ctx.encode_child(&self.expr)?)), + arrow_type: Some(self.cast_type().try_into()?), + }, + ))), + })) + } +} + +#[cfg(feature = "proto")] +impl TryCastExpr { + /// Reconstruct a [`TryCastExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_physical_expr_common::physical_expr::proto_decode::require_proto_field; + use datafusion_proto_models::protobuf; + + let try_cast = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::TryCast, + "TryCastExpr", + ); + let expr = ctx.decode_required_expression( + try_cast.expr.as_deref(), + "TryCastExpr", + "expr", + )?; + let arrow_type = require_proto_field( + try_cast.arrow_type.as_ref(), + "TryCastExpr", + "arrow_type", + )?; + let cast_type: DataType = arrow_type.try_into()?; + + Ok(Arc::new(TryCastExpr::new(expr, cast_type))) + } } /// Return a PhysicalExpression representing `expr` casted to @@ -593,3 +643,143 @@ mod tests { Ok(()) } } + +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::expressions::{Column, col}; + use crate::proto_test_util::{ + StubDecoder, StubEncoder, UnreachableDecoder, column_node, + }; + use arrow::datatypes::Field; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::datafusion_common::ArrowType; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalTryCastNode, physical_expr_node, + }; + + fn try_cast_fixture() -> TryCastExpr { + let schema = Schema::new(vec![Field::new("a", DataType::Utf8, true)]); + TryCastExpr::new(col("a", &schema).unwrap(), DataType::Int32) + } + + fn int32_arrow_type() -> ArrowType { + (&DataType::Int32).try_into().unwrap() + } + + fn try_cast_node( + expr: Option>, + arrow_type: Option, + ) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::TryCast(Box::new( + PhysicalTryCastNode { expr, arrow_type }, + ))), + } + } + + #[test] + fn try_to_proto_encodes_try_cast_expr() { + let try_cast = try_cast_fixture(); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = try_cast + .try_to_proto(&ctx) + .unwrap() + .expect("TryCastExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let try_cast_node = match node.expr_type { + Some(physical_expr_node::ExprType::TryCast(boxed)) => *boxed, + other => panic!("expected a TryCastExpr node, got {other:?}"), + }; + assert!(try_cast_node.expr.is_some()); + + let arrow_type = try_cast_node + .arrow_type + .as_ref() + .expect("try cast type should be encoded"); + let data_type: DataType = arrow_type.try_into().unwrap(); + assert_eq!(data_type, DataType::Int32); + } + + #[test] + fn try_to_proto_propagates_child_encode_error() { + let try_cast = try_cast_fixture(); + let encoder = StubEncoder::failing_on(1); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + let err = try_cast.try_to_proto(&ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } + + #[test] + fn try_from_proto_decodes_try_cast_expr() { + let node = + try_cast_node(Some(Box::new(column_node("a"))), Some(int32_arrow_type())); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = TryCastExpr::try_from_proto(&node, &ctx).unwrap(); + let try_cast = decoded + .downcast_ref::() + .expect("decoded expr should be a TryCastExpr"); + + assert_eq!(try_cast.cast_type(), &DataType::Int32); + assert!(try_cast.expr().downcast_ref::().is_some()); + } + + #[test] + fn try_from_proto_rejects_non_try_cast_node() { + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("PhysicalExprNode is not a TryCastExpr")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_expr() { + let node = try_cast_node(None, Some(int32_arrow_type())); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("TryCastExpr is missing required field 'expr'")) + ); + } + + #[test] + fn try_from_proto_rejects_missing_arrow_type() { + let node = try_cast_node(Some(Box::new(column_node("a"))), None); + let schema = Schema::empty(); + let decoder = StubDecoder::ok(); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + matches!(err, DataFusionError::Internal(msg) if msg.contains("TryCastExpr is missing required field 'arrow_type'")) + ); + } + + #[test] + fn try_from_proto_propagates_child_decode_error() { + let node = + try_cast_node(Some(Box::new(column_node("a"))), Some(int32_arrow_type())); + let schema = Schema::empty(); + let decoder = StubDecoder::failing_on(1); + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = TryCastExpr::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!(err, DataFusionError::Internal(msg) if msg.contains("call 1"))); + } +} diff --git a/datafusion/physical-expr/src/expressions/unknown_column.rs b/datafusion/physical-expr/src/expressions/unknown_column.rs index 4969fc33743c7..ed85f20dd274b 100644 --- a/datafusion/physical-expr/src/expressions/unknown_column.rs +++ b/datafusion/physical-expr/src/expressions/unknown_column.rs @@ -27,6 +27,7 @@ use arrow::{ record_batch::RecordBatch, }; use datafusion_common::{Result, internal_err}; + use datafusion_expr::ColumnarValue; #[derive(Debug, Clone, Eq)] @@ -84,6 +85,42 @@ impl PhysicalExpr for UnKnownColumn { fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { std::fmt::Display::fmt(self, f) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::UnknownColumn( + protobuf::UnknownColumn { + name: self.name.clone(), + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl UnKnownColumn { + /// Reconstruct an [`UnKnownColumn`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_proto_models::protobuf; + + let unknown_col = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::UnknownColumn, + "UnKnownColumn", + ); + Ok(Arc::new(UnKnownColumn::new(&unknown_col.name))) + } } impl Hash for UnKnownColumn { @@ -99,3 +136,103 @@ impl PartialEq for UnKnownColumn { false } } + +/// Tests for the `try_to_proto` / `try_from_proto` hooks. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::proto_test_util::{StubEncoder, UnreachableDecoder, column_node}; + use arrow::datatypes::Schema; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{self, physical_expr_node}; + + // ── try_to_proto ───────────────────────────────────────────────────────── + + #[test] + fn try_to_proto_encodes_unknown_column() { + let expr = UnKnownColumn::new("my_col"); + let encoder = StubEncoder::ok(); + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = expr + .try_to_proto(&ctx) + .unwrap() + .expect("UnKnownColumn should encode to Some(node)"); + + // Built-in exprs never set expr_id; only dynamic filters do. + assert!(node.expr_id.is_none()); + + // Verify the encoded name matches the original. + let protobuf::UnknownColumn { name } = match node.expr_type { + Some(physical_expr_node::ExprType::UnknownColumn(c)) => c, + other => panic!("expected UnknownColumn proto node, got {other:?}"), + }; + assert_eq!(name, "my_col"); + } + + // ── try_from_proto ─────────────────────────────────────────────────────── + + #[test] + fn try_from_proto_decodes_name() { + let node = protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::UnknownColumn( + protobuf::UnknownColumn { + name: "my_col".to_string(), + }, + )), + }; + let schema = Schema::empty(); + // UnKnownColumn has no child exprs so the decoder is never called. + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = UnKnownColumn::try_from_proto(&node, &ctx).unwrap(); + let col = decoded + .downcast_ref::() + .expect("decoded expr should be an UnKnownColumn"); + assert_eq!(col.name(), "my_col"); + } + + #[test] + fn try_from_proto_rejects_non_unknown_column_node() { + // column_node produces an ExprType::Column node, not UnknownColumn. + let node = column_node("a"); + let schema = Schema::empty(); + let decoder = UnreachableDecoder; + let ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = UnKnownColumn::try_from_proto(&node, &ctx).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(ref msg) + if msg.contains("PhysicalExprNode is not a UnKnownColumn") + )); + } + + // ── roundtrip ──────────────────────────────────────────────────────────── + + #[test] + fn unknown_column_proto_roundtrip() { + let expr = UnKnownColumn::new("col_b"); + let encoder = StubEncoder::ok(); + let enc_ctx = PhysicalExprEncodeCtx::new(&encoder); + + let node = expr + .try_to_proto(&enc_ctx) + .unwrap() + .expect("UnKnownColumn should encode to Some(node)"); + + let schema = Schema::empty(); + // UnKnownColumn has no child exprs so the decoder is never called. + let decoder = UnreachableDecoder; + let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let decoded = UnKnownColumn::try_from_proto(&node, &dec_ctx).unwrap(); + let col = decoded + .downcast_ref::() + .expect("decoded expr should be an UnKnownColumn"); + assert_eq!(col.name(), "col_b"); + } +} diff --git a/datafusion/physical-expr/src/higher_order_function.rs b/datafusion/physical-expr/src/higher_order_function.rs index 801e69ea8fb69..e28b38bd7c8c1 100644 --- a/datafusion/physical-expr/src/higher_order_function.rs +++ b/datafusion/physical-expr/src/higher_order_function.rs @@ -69,7 +69,7 @@ enum ArgSlot { /// Physical expression of a higher order function pub struct HigherOrderFunctionExpr { /// A shared instance of the higher-order function - fun: Arc, + fun: Arc, /// The name of the higher-order function name: String, /// List of expressions to feed to the function as arguments @@ -125,7 +125,7 @@ impl HigherOrderFunctionExpr { /// Note that lambda arguments must be present directly in args as [LambdaExpr], /// and not as a wrapped child of any arg pub fn try_new_with_schema( - fun: Arc, + fun: Arc, args: Vec>, schema: &Schema, config_options: Arc, @@ -172,7 +172,7 @@ impl HigherOrderFunctionExpr { } /// Get the higher order function implementation - pub fn fun(&self) -> &dyn HigherOrderUDF { + pub fn fun(&self) -> &HigherOrderUDF { self.fun.as_ref() } @@ -200,7 +200,7 @@ impl HigherOrderFunctionExpr { } /// Resolve every lambda's parameter list. Returns an empty `Vec` when - /// there are no lambdas, avoiding the [`HigherOrderUDF::lambda_parameters`] + /// there are no lambdas, avoiding the [`datafusion_expr::HigherOrderUDFImpl::lambda_parameters`] /// virtual call entirely. fn resolve_lambda_parameters( &self, @@ -353,6 +353,7 @@ impl PhysicalExpr for HigherOrderFunctionExpr { } else { Some(batch.project(&projection)?) }, + lambda.used_param_indices(), ))) } ArgSlot::Value => { @@ -509,17 +510,20 @@ mod tests { use super::*; use crate::HigherOrderFunctionExpr; + use crate::create_physical_expr; use crate::expressions::Column; use crate::expressions::NoOp; use crate::expressions::lambda; use crate::expressions::not; - use arrow::array::NullArray; use arrow::array::RecordBatchOptions; + use arrow::array::{ArrayRef, Int32Array}; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::Result; use datafusion_common::assert_contains; + use datafusion_expr::execution_props::ExecutionProps; + use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::{ - HigherOrderFunctionArgs, HigherOrderSignature, HigherOrderUDF, + HigherOrderFunctionArgs, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, }; use datafusion_expr_common::columnar_value::ColumnarValue; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; @@ -531,7 +535,7 @@ mod tests { signature: HigherOrderSignature, } - impl HigherOrderUDF for MockHigherOrderUDF { + impl HigherOrderUDFImpl for MockHigherOrderUDF { fn name(&self) -> &str { "mock_function" } @@ -545,9 +549,11 @@ mod tests { _step: usize, _fields: &[ValueOrLambda>], ) -> Result { - Ok(LambdaParametersProgress::Complete(vec![vec![Arc::new( - Field::new("", DataType::Null, true), - )]])) + // Offer two params; single-param lambdas just ignore the second. + Ok(LambdaParametersProgress::Complete(vec![vec![ + Arc::new(Field::new("", DataType::Int32, true)), + Arc::new(Field::new("", DataType::Int32, true)), + ]])) } fn return_field_from_args( @@ -567,7 +573,18 @@ mod tests { ) -> Result { match &args.args[0] { ValueOrLambda::Lambda(lambda) => lambda.evaluate( - &[&|| Ok(Arc::new(NullArray::new(args.number_rows)))], + &[ + // Sentinel for the first param, distinct from the second's value. + &|| { + Ok(Arc::new(Int32Array::from(vec![-1000; args.number_rows])) + as ArrayRef) + }, + &|| { + Ok(Arc::new(Int32Array::from_iter_values( + (0..args.number_rows as i32).map(|i| 10 * (i + 1)), + )) as ArrayRef) + }, + ], |arrays| Ok(arrays.to_vec()), ), ValueOrLambda::Value(value) => Ok(value.clone()), @@ -578,14 +595,14 @@ mod tests { #[test] fn test_higher_order_function_volatile_node() { // Create a volatile UDF - let volatile_udf = Arc::new(MockHigherOrderUDF { + let volatile_udf = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::variadic_any(Volatility::Volatile), - }); + })); // Create a non-volatile UDF - let stable_udf = Arc::new(MockHigherOrderUDF { + let stable_udf = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::variadic_any(Volatility::Stable), - }); + })); let schema = Schema::new(vec![Field::new("a", DataType::Float32, false)]); let args = vec![Arc::new(Column::new("a", 0)) as Arc]; @@ -620,9 +637,9 @@ mod tests { #[test] fn test_higher_order_function_wrapped_lambda() { - let fun = Arc::new(MockHigherOrderUDF { + let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::variadic_any(Volatility::Stable), - }); + })); let expected = ScalarValue::Int32(Some(42)); @@ -657,9 +674,9 @@ mod tests { #[test] fn test_higher_order_function_badly_wrapped_lambda() { - let fun = Arc::new(MockHigherOrderUDF { + let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::variadic_any(Volatility::Stable), - }); + })); let hof = HigherOrderFunctionExpr::try_new_with_schema( fun, @@ -694,9 +711,9 @@ mod tests { #[test] fn test_higher_order_function_unexpected_lambda() { - let fun = Arc::new(MockHigherOrderUDF { + let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF { signature: HigherOrderSignature::variadic_any(Volatility::Stable), - }); + })); let hof = HigherOrderFunctionExpr::try_new_with_schema( fun, @@ -715,4 +732,54 @@ mod tests { "mock_function received a lambda via with_new_children at position 0 that wasn't a lambda before" ); } + + /// Exercises the real planner end to end (not hand-picked indices) to + /// check the "captures before own-params" layout invariant. + #[test] + fn test_higher_order_function_two_lambda_params_capture_and_unused_param() { + use datafusion_common::DFSchema; + use datafusion_expr::expr::{HigherOrderFunction, LambdaVariable}; + use datafusion_expr::{Expr, col, lambda as logical_lambda}; + + let fun = Arc::new(HigherOrderUDF::new_from_impl(MockHigherOrderUDF { + signature: HigherOrderSignature::variadic_any(Volatility::Stable), + })); + + // Body uses capture "a" and param "v"; param "k" is left unused. + let v = Expr::LambdaVariable(LambdaVariable::new( + "v".to_string(), + Some(Arc::new(Field::new("v", DataType::Int32, true))), + )); + let body = col("a") + v; + let lambda_expr = logical_lambda(["k", "v"], body); + + let schema = DFSchema::from_unqualified_fields( + vec![Field::new("a", DataType::Int32, false)].into(), + std::collections::HashMap::new(), + ) + .unwrap(); + + let physical_expr = create_physical_expr( + &Expr::HigherOrderFunction(HigherOrderFunction::new(fun, vec![lambda_expr])), + &schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) + .unwrap(); + + let batch = RecordBatch::try_new( + Arc::clone(schema.inner()), + vec![Arc::new(Int32Array::from(vec![1, 2, 3])) as ArrayRef], + ) + .unwrap(); + + let result = physical_expr.evaluate(&batch).unwrap(); + let ColumnarValue::Array(result) = result else { + unreachable!() + }; + + // a + v; k's sentinel (-1000) must not leak into the result. + let expected = Int32Array::from(vec![11, 22, 33]); + assert_eq!(result.as_ref(), &expected); + } } diff --git a/datafusion/physical-expr/src/lib.rs b/datafusion/physical-expr/src/lib.rs index 848bf81d15979..80e9f88b510ed 100644 --- a/datafusion/physical-expr/src/lib.rs +++ b/datafusion/physical-expr/src/lib.rs @@ -40,6 +40,9 @@ mod partitioning; mod physical_expr; pub mod planner; pub mod projection; +/// Shared test helpers for the `try_to_proto` / `try_from_proto` unit tests +#[cfg(all(test, feature = "proto"))] +pub(crate) mod proto_test_util; mod scalar_function; pub mod scalar_subquery; pub mod simplifier; @@ -55,14 +58,19 @@ pub mod execution_props { pub use aggregate::groups_accumulator::{GroupsAccumulatorAdapter, NullState}; pub use analysis::{AnalysisContext, ExprBoundaries, analyze}; +pub use datafusion_common::SplitPoint; pub use equivalence::{ AcrossPartitions, ConstExpr, EquivalenceProperties, calculate_union, }; -pub use partitioning::{Distribution, Partitioning}; +pub use expressions::{DynamicFilterTracker, DynamicFilterTracking}; +pub use partitioning::{ + Distribution, Partitioning, PartitioningSatisfaction, RangePartitioning, +}; pub use physical_expr::{ add_offset_to_expr, add_offset_to_physical_sort_exprs, create_lex_ordering, - create_ordering, create_physical_sort_expr, create_physical_sort_exprs, - physical_exprs_bag_equal, physical_exprs_contains, physical_exprs_equal, + create_ordering, create_physical_partitioning, create_physical_sort_expr, + create_physical_sort_exprs, physical_exprs_bag_equal, physical_exprs_contains, + physical_exprs_equal, }; pub use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, PhysicalExprRef}; diff --git a/datafusion/physical-expr/src/partitioning.rs b/datafusion/physical-expr/src/partitioning.rs index d24c60b63e6bd..98f082f7256db 100644 --- a/datafusion/physical-expr/src/partitioning.rs +++ b/datafusion/physical-expr/src/partitioning.rs @@ -19,9 +19,16 @@ use crate::{ EquivalenceProperties, PhysicalExpr, equivalence::ProjectionMapping, - expressions::UnKnownColumn, physical_exprs_equal, + expressions::UnKnownColumn, physical_exprs_contains, physical_exprs_equal, }; +pub use datafusion_common::SplitPoint; +use datafusion_common::{Result, validate_range_split_points}; use datafusion_physical_expr_common::physical_expr::format_physical_expr_list; +use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; use std::fmt; use std::fmt::Display; use std::sync::Arc; @@ -117,6 +124,8 @@ pub enum Partitioning { /// Allocate rows based on a hash of one of more expressions and the specified number of /// partitions Hash(Vec>, usize), + /// Partition rows by source-declared ranges + Range(RangePartitioning), /// Unknown partitioning scheme with a known number of partitions UnknownPartitioning(usize), } @@ -133,6 +142,7 @@ impl Display for Partitioning { .join(", "); write!(f, "Hash([{phy_exprs_str}], {size})") } + Partitioning::Range(range) => write!(f, "{range}"), Partitioning::UnknownPartitioning(size) => { write!(f, "UnknownPartitioning({size})") } @@ -140,6 +150,191 @@ impl Display for Partitioning { } } +/// Physical range partitioning. +/// +/// [`RangePartitioning`] describes an ordered key space with split points. +/// +/// - `ordering` defines the partitioning key and ordering. +/// - `split_points` define the boundaries between adjacent partitions. +/// +/// Comparisons use the lexicographic order defined by `ordering`, including +/// `ASC`/`DESC` and null ordering. Split points must be strictly ordered +/// according to that ordering, and each split point must have one value per +/// ordering expression. See [`SplitPoint`] for the shared boundary convention. +/// +/// Like other user-specified data properties such as sortedness, if a source +/// declares range partitioning, it is responsible for placing each row in the +/// partition described by the split points. DataFusion will not validate this is +/// upheld. +/// +/// For a single range key: +/// +/// ```text +/// ordering = [date ASC NULLS LAST] +/// split_points = [ +/// (2022-01-01), +/// (2023-01-01), +/// ] +/// +/// partition 0: date before 2022-01-01 +/// partition 1: date between 2022-01-01 (inclusive) and 2023-01-01 (exclusive) +/// partition 2: date at/after 2023-01-01 +/// ``` +/// +/// The same model extends to compound keys. +/// For `ordering = [time ASC, city ASC]`, split points are ordered +/// lexicographically by `(time, city)`: +/// +/// ```text +/// ordering = [time ASC NULLS LAST, city ASC NULLS LAST] +/// split_points = [ +/// (2022, Allston), +/// (2023, Allston), +/// ] +/// +/// partition 0: keys before (2022, Allston) +/// partition 1: keys between (2022, Allston) and (2023, Allston) +/// partition 2: keys at/after (2023, Allston) +/// ``` +/// +/// NOTE: Optimizer and execution behavior for this partitioning is intentionally +/// not implemented and will be introduced incrementally. See +/// . +#[derive(Debug, Clone, PartialEq)] +pub struct RangePartitioning { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Boundaries between adjacent partitions. + split_points: Vec, +} + +impl RangePartitioning { + /// Creates range partitioning metadata without validating split points. + /// + /// Use [`Self::try_new`] to validate the contract documented on + /// [`RangePartitioning`]. + pub fn new(ordering: LexOrdering, split_points: Vec) -> Self { + Self { + ordering, + split_points, + } + } + + /// Creates range partitioning metadata and validates split point shape and + /// ordering. + pub fn try_new(ordering: LexOrdering, split_points: Vec) -> Result { + validate_range_split_points( + &split_points, + &ordering + .iter() + .map(|sort_expr| sort_expr.options) + .collect::>(), + )?; + Ok(Self::new(ordering, split_points)) + } + + /// Returns the ordering that defines the range key. + pub fn ordering(&self) -> &LexOrdering { + &self.ordering + } + + /// Returns the ordered split points between partitions. + pub fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Returns the number of partitions. + pub fn partition_count(&self) -> usize { + self.split_points.len() + 1 + } + + /// Calculates the range partitioning after applying the given projection. + /// + /// Returns `None` if any range key cannot be projected or if projection + /// collapses distinct range keys into duplicate output expressions. + fn project( + &self, + mapping: &ProjectionMapping, + input_eq_properties: &EquivalenceProperties, + ) -> Option { + let exprs = self + .ordering + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + let projected_exprs = input_eq_properties + .project_expressions(&exprs, mapping) + .collect::>>()?; + let sort_exprs = self + .ordering + .iter() + .zip(projected_exprs) + .map(|(sort_expr, expr)| PhysicalSortExpr::new(expr, sort_expr.options)) + .collect::>(); + let ordering = LexOrdering::new(sort_exprs)?; + if ordering.len() != self.ordering.len() { + return None; + } + + Some(Self { + ordering, + split_points: self.split_points.clone(), + }) + } +} + +impl Display for RangePartitioning { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + let split_points = format_range_split_points(&self.split_points); + write!( + f, + "Range([{}], [{}], {})", + self.ordering, + split_points, + self.partition_count() + ) + } +} + +fn format_range_split_points(split_points: &[SplitPoint]) -> String { + split_points + .iter() + .map(ToString::to_string) + .collect::>() + .join(", ") +} + +fn equivalent_exprs( + left: &[Arc], + right: &[Arc], + eq_properties: &EquivalenceProperties, +) -> bool { + if physical_exprs_equal(left, right) { + return true; + } + + let eq_groups = eq_properties.eq_group(); + if eq_groups.is_empty() { + return false; + } + + let normalized_left = normalize_exprs(left, eq_properties); + let normalized_right = normalize_exprs(right, eq_properties); + + physical_exprs_equal(&normalized_left, &normalized_right) +} + +fn normalize_exprs( + exprs: &[Arc], + eq_properties: &EquivalenceProperties, +) -> Vec> { + let eq_groups = eq_properties.eq_group(); + exprs + .iter() + .map(|expr| eq_groups.normalize_expr(Arc::clone(expr))) + .collect() +} + /// Represents how a [`Partitioning`] satisfies a [`Distribution`] requirement. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum PartitioningSatisfaction { @@ -167,6 +362,7 @@ impl Partitioning { use Partitioning::*; match self { RoundRobinBatch(n) | Hash(_, n) | UnknownPartitioning(n) => *n, + Range(range) => range.partition_count(), } } @@ -182,11 +378,9 @@ impl Partitioning { return false; } - subset_exprs.iter().all(|subset_expr| { - superset_exprs - .iter() - .any(|superset_expr| subset_expr.eq(superset_expr)) - }) + subset_exprs + .iter() + .all(|subset_expr| physical_exprs_contains(superset_exprs, subset_expr)) } #[deprecated(since = "52.0.0", note = "Use satisfaction instead")] @@ -201,6 +395,10 @@ impl Partitioning { /// Returns how this [`Partitioning`] satisfies the partitioning scheme mandated /// by the `required` [`Distribution`]. + #[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" + )] pub fn satisfaction( &self, required: &Distribution, @@ -212,88 +410,265 @@ impl Partitioning { Distribution::SinglePartition if self.partition_count() == 1 => { PartitioningSatisfaction::Exact } - // When partition count is 1, hash requirement is satisfied. - Distribution::HashPartitioned(_) if self.partition_count() == 1 => { + // When partition count is 1, key partitioning is satisfied. + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) + if self.partition_count() == 1 => + { PartitioningSatisfaction::Exact } - Distribution::HashPartitioned(required_exprs) => match self { + Distribution::HashPartitioned(required_exprs) + | Distribution::KeyPartitioned(required_exprs) => match self { // Here we do not check the partition count for hash partitioning and assumes the partition count // and hash functions in the system are the same. In future if we plan to support storage partition-wise joins, // then we need to have the partition count and hash functions validation. - Partitioning::Hash(partition_exprs, _) => { - // Empty hash partitioning is invalid - if partition_exprs.is_empty() || required_exprs.is_empty() { - return PartitioningSatisfaction::NotSatisfied; - } - - // Fast path: exact match - if physical_exprs_equal(required_exprs, partition_exprs) { - return PartitioningSatisfaction::Exact; - } - - // Normalization path using equivalence groups - let eq_groups = eq_properties.eq_group(); - if !eq_groups.is_empty() { - let normalized_required_exprs = required_exprs - .iter() - .map(|e| eq_groups.normalize_expr(Arc::clone(e))) - .collect::>(); - let normalized_partition_exprs = partition_exprs - .iter() - .map(|e| eq_groups.normalize_expr(Arc::clone(e))) - .collect::>(); - if physical_exprs_equal( - &normalized_required_exprs, - &normalized_partition_exprs, - ) { - return PartitioningSatisfaction::Exact; - } - - if allow_subset - && Self::is_subset_partitioning( - &normalized_partition_exprs, - &normalized_required_exprs, - ) - { - return PartitioningSatisfaction::Subset; - } - } else if allow_subset - && Self::is_subset_partitioning(partition_exprs, required_exprs) - { - return PartitioningSatisfaction::Subset; - } - + Partitioning::Hash(partition_exprs, _) => Self::key_satisfaction( + partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ), + Partitioning::Range(range) => { + let partition_exprs = range + .ordering() + .iter() + .map(|sort_expr| Arc::clone(&sort_expr.expr)) + .collect::>(); + Self::key_satisfaction( + &partition_exprs, + required_exprs, + eq_properties, + allow_subset, + ) + } + Partitioning::RoundRobinBatch(_) + | Partitioning::UnknownPartitioning(_) => { PartitioningSatisfaction::NotSatisfied } - _ => PartitioningSatisfaction::NotSatisfied, }, - _ => PartitioningSatisfaction::NotSatisfied, + Distribution::SinglePartition => PartitioningSatisfaction::NotSatisfied, } } + fn key_satisfaction( + partition_exprs: &[Arc], + required_exprs: &[Arc], + eq_properties: &EquivalenceProperties, + allow_subset: bool, + ) -> PartitioningSatisfaction { + if partition_exprs.is_empty() || required_exprs.is_empty() { + return PartitioningSatisfaction::NotSatisfied; + } + + if equivalent_exprs(required_exprs, partition_exprs, eq_properties) { + return PartitioningSatisfaction::Exact; + } + + let eq_groups = eq_properties.eq_group(); + if !eq_groups.is_empty() { + if allow_subset { + let normalized_partition_exprs = + normalize_exprs(partition_exprs, eq_properties); + let normalized_required_exprs = + normalize_exprs(required_exprs, eq_properties); + if Self::is_subset_partitioning( + &normalized_partition_exprs, + &normalized_required_exprs, + ) { + return PartitioningSatisfaction::Subset; + } + } + } else if allow_subset + && Self::is_subset_partitioning(partition_exprs, required_exprs) + { + return PartitioningSatisfaction::Subset; + } + + PartitioningSatisfaction::NotSatisfied + } + /// Calculate the output partitioning after applying the given projection. pub fn project( &self, mapping: &ProjectionMapping, input_eq_properties: &EquivalenceProperties, ) -> Self { - if let Partitioning::Hash(exprs, part) = self { - let normalized_exprs = input_eq_properties - .project_expressions(exprs, mapping) - .zip(exprs) - .map(|(proj_expr, expr)| { - proj_expr.unwrap_or_else(|| { - Arc::new(UnKnownColumn::new(&expr.to_string())) + match self { + Partitioning::Hash(exprs, part) => { + let normalized_exprs = input_eq_properties + .project_expressions(exprs, mapping) + .zip(exprs) + .map(|(proj_expr, expr)| { + proj_expr.unwrap_or_else(|| { + Arc::new(UnKnownColumn::new(&expr.to_string())) + }) }) - }) - .collect(); - Partitioning::Hash(normalized_exprs, *part) - } else { - self.clone() + .collect(); + Partitioning::Hash(normalized_exprs, *part) + } + Partitioning::Range(range) => { + if let Some(projected) = range.project(mapping, input_eq_properties) { + Partitioning::Range(projected) + } else { + Partitioning::UnknownPartitioning(range.partition_count()) + } + } + Partitioning::RoundRobinBatch(_) | Partitioning::UnknownPartitioning(_) => { + self.clone() + } } } } +/// Protobuf conversions for [`Partitioning`]. +/// +/// Child expressions (hash keys, range orderings) and `ScalarValue` split +/// points are (de)serialized through the expression-level context, so this is +/// the single copy of the partitioning wire format: `RepartitionExec` and +/// `datafusion-proto`'s central serializer route through it, and the remaining +/// per-plan migrations (`FileScanConfig` and friends) are meant to do the same +/// rather than grow another copy. +/// +/// [`protobuf::Partitioning`]: datafusion_proto_models::protobuf::Partitioning +#[cfg(feature = "proto")] +impl Partitioning { + /// Serialize this partitioning into its protobuf representation. + pub fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result { + use datafusion_proto_models::protobuf; + + let partition_method = match self { + Partitioning::RoundRobinBatch(n) => { + protobuf::partitioning::PartitionMethod::RoundRobin(wire_partition_count( + *n, + )?) + } + Partitioning::Hash(exprs, n) => { + protobuf::partitioning::PartitionMethod::Hash( + protobuf::PhysicalHashRepartition { + hash_expr: ctx.encode_children_expressions(exprs)?, + partition_count: wire_partition_count(*n)?, + }, + ) + } + Partitioning::Range(range) => { + let sort_expr = sort_exprs_try_to_proto(range.ordering().iter(), ctx)?; + let split_point = range + .split_points() + .iter() + .map(|split_point| { + let value = split_point + .values() + .iter() + .map(|value| value.try_into().map_err(Into::into)) + .collect::>>()?; + Ok(protobuf::PhysicalRangeSplitPoint { value }) + }) + .collect::>>()?; + protobuf::partitioning::PartitionMethod::Range( + protobuf::PhysicalRangePartitioning { + sort_expr, + split_point, + }, + ) + } + Partitioning::UnknownPartitioning(n) => { + protobuf::partitioning::PartitionMethod::Unknown(wire_partition_count( + *n, + )?) + } + }; + Ok(protobuf::Partitioning { + partition_method: Some(partition_method), + }) + } + + /// Reconstruct a [`Partitioning`] from its protobuf representation. + /// + /// Returns `Ok(None)` when the message carries no `partition_method`, which + /// the wire format uses to mean "no output partitioning declared"; callers + /// for which it is required should turn that into their own error. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::Partitioning, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::{ScalarValue, internal_datafusion_err, internal_err}; + use datafusion_proto_models::protobuf; + + let Some(partition_method) = node.partition_method.as_ref() else { + return Ok(None); + }; + let partitioning = match partition_method { + protobuf::partitioning::PartitionMethod::RoundRobin(n) => { + Partitioning::RoundRobinBatch(partition_count(*n)?) + } + protobuf::partitioning::PartitionMethod::Hash(hash) => { + let exprs = hash + .hash_expr + .iter() + .map(|expr| ctx.decode(expr)) + .collect::>>()?; + Partitioning::Hash(exprs, partition_count(hash.partition_count)?) + } + protobuf::partitioning::PartitionMethod::Unknown(n) => { + Partitioning::UnknownPartitioning(partition_count(*n)?) + } + protobuf::partitioning::PartitionMethod::Range(range) => { + let sort_exprs = sort_exprs_try_from_proto(&range.sort_expr, ctx)?; + let sort_expr_count = sort_exprs.len(); + let ordering = LexOrdering::new(sort_exprs).ok_or_else(|| { + internal_datafusion_err!( + "Range partitioning requires non-empty ordering" + ) + })?; + if ordering.len() != sort_expr_count { + return internal_err!( + "Range partitioning ordering must not contain duplicate expressions" + ); + } + let split_points = range + .split_point + .iter() + .map(|split_point| { + let values = split_point + .value + .iter() + .map(|value| ScalarValue::try_from(value).map_err(Into::into)) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect::>>()?; + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?) + } + }; + Ok(Some(partitioning)) + } +} + +/// Narrow a wire partition count to `usize`. +#[cfg(feature = "proto")] +fn partition_count(count: u64) -> Result { + usize::try_from(count).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Partition count {count} exceeds usize::MAX" + ) + }) +} + +/// Widen a partition count to its `u64` wire representation. +/// +/// The mirror of [`partition_count`]: an out-of-range count is an error on both +/// sides rather than a silent truncation on the way out. +#[cfg(feature = "proto")] +fn wire_partition_count(count: usize) -> Result { + u64::try_from(count).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Partition count {count} exceeds u64::MAX" + ) + }) +} + impl PartialEq for Partitioning { fn eq(&self, other: &Partitioning) -> bool { match (self, other) { @@ -306,6 +681,7 @@ impl PartialEq for Partitioning { { true } + (Partitioning::Range(left), Partitioning::Range(right)) => left == right, _ => false, } } @@ -319,11 +695,19 @@ pub enum Distribution { UnspecifiedDistribution, /// A single partition is required SinglePartition, + /// Deprecated historical name for [`Distribution::KeyPartitioned`]. + /// See for details. + #[deprecated(since = "55.0.0", note = "Use Distribution::KeyPartitioned")] + HashPartitioned(Vec>), /// Requires children to be distributed in such a way that the same /// values of the keys end up in the same partition - HashPartitioned(Vec>), + KeyPartitioned(Vec>), } +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] impl Distribution { /// Creates a `Partitioning` that satisfies this `Distribution` pub fn create_partitioning(self, partition_count: usize) -> Partitioning { @@ -332,13 +716,17 @@ impl Distribution { Partitioning::UnknownPartitioning(partition_count) } Distribution::SinglePartition => Partitioning::UnknownPartitioning(1), - Distribution::HashPartitioned(expr) => { + Distribution::HashPartitioned(expr) | Distribution::KeyPartitioned(expr) => { Partitioning::Hash(expr, partition_count) } } } } +#[expect( + deprecated, + reason = "HashPartitioned display is preserved during the KeyPartitioned migration" +)] impl Display for Distribution { fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { match self { @@ -347,6 +735,9 @@ impl Display for Distribution { Distribution::HashPartitioned(exprs) => { write!(f, "HashPartitioned[{}])", format_physical_expr_list(exprs)) } + Distribution::KeyPartitioned(exprs) => { + write!(f, "KeyPartitioned[{}])", format_physical_expr_list(exprs)) + } } } } @@ -356,56 +747,181 @@ mod tests { use super::*; use crate::expressions::Column; + use crate::projection::ProjectionTargets; - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::Result; + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_common::{Result, ScalarValue}; + + struct PartitioningTestFixture { + schema: SchemaRef, + cols: Vec>, + eq_properties: EquivalenceProperties, + } + + impl PartitioningTestFixture { + fn new(fields: Vec<(&str, DataType)>) -> Result { + let schema = Arc::new(Schema::new( + fields + .iter() + .map(|(name, data_type)| Field::new(*name, data_type.clone(), false)) + .collect::>(), + )); + let cols = fields + .iter() + .map(|(name, _)| { + Ok(Arc::new(Column::new_with_schema(name, &schema)?) + as Arc) + }) + .collect::>()?; + let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + + Ok(Self { + schema, + cols, + eq_properties, + }) + } + + fn int64(names: &[&str]) -> Result { + Self::new(names.iter().map(|name| (*name, DataType::Int64)).collect()) + } + + fn col(&self, index: usize) -> Arc { + Arc::clone(&self.cols[index]) + } + + fn cols( + &self, + indices: impl IntoIterator, + ) -> Vec> { + indices.into_iter().map(|index| self.col(index)).collect() + } + + fn hash_partitioning( + &self, + indices: impl IntoIterator, + partition_count: usize, + ) -> Partitioning { + Partitioning::Hash(self.cols(indices), partition_count) + } + + fn key_distribution( + &self, + indices: impl IntoIterator, + ) -> Distribution { + Distribution::KeyPartitioned(self.cols(indices)) + } + + fn range_sort_expr( + &self, + index: usize, + options: SortOptions, + ) -> PhysicalSortExpr { + PhysicalSortExpr::new(self.col(index), options) + } + + fn range_ordering( + &self, + indices: impl IntoIterator, + ) -> LexOrdering { + LexOrdering::new( + indices + .into_iter() + .map(|index| PhysicalSortExpr::new_default(self.col(index))), + ) + .expect("ordering must not be empty") + } + + fn range( + &self, + indices: impl IntoIterator, + split_points: Vec, + ) -> RangePartitioning { + RangePartitioning::try_new(self.range_ordering(indices), split_points) + .expect("test range partitioning should be valid") + } + + fn range_partitioning( + &self, + indices: impl IntoIterator, + split_points: Vec, + ) -> Partitioning { + Partitioning::Range(self.range(indices, split_points)) + } + + fn range_partitioning_with_ordering( + &self, + ordering: LexOrdering, + split_points: Vec, + ) -> Partitioning { + Partitioning::Range( + RangePartitioning::try_new(ordering, split_points) + .expect("test range partitioning should be valid"), + ) + } + } + + fn assert_satisfaction( + desc: &str, + partitioning: &Partitioning, + required: &Distribution, + eq_properties: &EquivalenceProperties, + expected_with_subset: PartitioningSatisfaction, + expected_without_subset: PartitioningSatisfaction, + ) { + assert_eq!( + partitioning.satisfaction(required, eq_properties, true), + expected_with_subset, + "Failed for {desc} with subset enabled" + ); + assert_eq!( + partitioning.satisfaction(required, eq_properties, false), + expected_without_subset, + "Failed for {desc} with subset disabled" + ); + } #[test] + #[expect( + deprecated, + reason = "test intentionally covers deprecated HashPartitioned compatibility" + )] fn partitioning_satisfy_distribution() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("column_1", DataType::Int64, false), - Field::new("column_2", DataType::Utf8, false), - ])); - - let partition_exprs1: Vec> = vec![ - Arc::new(Column::new_with_schema("column_1", &schema).unwrap()), - Arc::new(Column::new_with_schema("column_2", &schema).unwrap()), - ]; - - let partition_exprs2: Vec> = vec![ - Arc::new(Column::new_with_schema("column_2", &schema).unwrap()), - Arc::new(Column::new_with_schema("column_1", &schema).unwrap()), - ]; + let fixture = PartitioningTestFixture::new(vec![ + ("column_1", DataType::Int64), + ("column_2", DataType::Utf8), + ])?; let distribution_types = vec![ Distribution::UnspecifiedDistribution, Distribution::SinglePartition, - Distribution::HashPartitioned(partition_exprs1.clone()), + Distribution::HashPartitioned(fixture.cols([0, 1])), + fixture.key_distribution([0, 1]), ]; let single_partition = Partitioning::UnknownPartitioning(1); let unspecified_partition = Partitioning::UnknownPartitioning(10); let round_robin_partition = Partitioning::RoundRobinBatch(10); - let hash_partition1 = Partitioning::Hash(partition_exprs1, 10); - let hash_partition2 = Partitioning::Hash(partition_exprs2, 10); - let eq_properties = EquivalenceProperties::new(schema); + let hash_partition1 = fixture.hash_partitioning([0, 1], 10); + let hash_partition2 = fixture.hash_partitioning([1, 0], 10); for distribution in distribution_types { let result = ( single_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), unspecified_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), round_robin_partition - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), hash_partition1 - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), hash_partition2 - .satisfaction(&distribution, &eq_properties, true) + .satisfaction(&distribution, &fixture.eq_properties, true) .is_satisfied(), ); @@ -416,7 +932,7 @@ mod tests { Distribution::SinglePartition => { assert_eq!(result, (true, false, false, false, false)) } - Distribution::HashPartitioned(_) => { + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) => { assert_eq!(result, (true, false, false, true, false)) } } @@ -426,141 +942,129 @@ mod tests { } #[test] - fn test_partitioning_satisfy_by_subset() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + #[expect( + deprecated, + reason = "test intentionally covers deprecated HashPartitioned compatibility" + )] + fn deprecated_hash_partitioned_matches_key_partitioned() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let partitioning = fixture.hash_partitioning([0, 1], 4); + let hash_distribution = Distribution::HashPartitioned(fixture.cols([0, 1])); + let key_distribution = fixture.key_distribution([0, 1]); + + assert_eq!( + partitioning.satisfaction(&hash_distribution, &fixture.eq_properties, false), + partitioning.satisfaction(&key_distribution, &fixture.eq_properties, false) + ); + assert_eq!( + hash_distribution.create_partitioning(4), + key_distribution.create_partitioning(4) + ); + + Ok(()) + } + + #[test] + fn hash_partitioning_key_distribution_satisfaction() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; + let unknown: Arc = Arc::new(UnKnownColumn::new("dropped")); let test_cases = vec![ ( - "Hash([a]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), + "exact: KeyPartitioned([a, b]) satisfied by Hash([a, b])", + fixture.hash_partitioning([0, 1], 4), + fixture.key_distribution([0, 1]), + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ), + ( + "subset: KeyPartitioned([a, b]) satisfied by Hash([a])", + fixture.hash_partitioning([0], 4), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + "subset: KeyPartitioned([a, b, c]) satisfied by Hash([b])", + fixture.hash_partitioning([1], 4), + fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), + "subset reordered: KeyPartitioned([a, b, c]) satisfied by Hash([b, a])", + fixture.hash_partitioning([1, 0], 4), + fixture.key_distribution([0, 1, 2]), PartitioningSatisfaction::Subset, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([b]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), - PartitioningSatisfaction::Subset, + "superset: KeyPartitioned([a]) not satisfied by Hash([a, b])", + fixture.hash_partitioning([0, 1], 4), + fixture.key_distribution([0]), + PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([b, a]) vs Hash([a, b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), - PartitioningSatisfaction::Subset, + "superset: KeyPartitioned([a, b]) not satisfied by Hash([a, b, c])", + fixture.hash_partitioning([0, 1, 2], 4), + fixture.key_distribution([0, 1]), + PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), - ]; - - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } - - Ok(()) - } - - #[test] - fn test_partitioning_current_superset() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let test_cases = vec![ ( - "Hash([a, b]) vs Hash([a])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + "partial overlap: KeyPartitioned([a, b]) not satisfied by Hash([a, c])", + fixture.hash_partitioning([0, 2], 4), + fixture.key_distribution([0, 1]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b, c]) vs Hash([a])", - Partitioning::Hash( - vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_c)], - 4, - ), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), + "no overlap: KeyPartitioned([b, c]) not satisfied by Hash([a])", + fixture.hash_partitioning([0], 4), + fixture.key_distribution([1, 2]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), ( - "Hash([a, b, c]) vs Hash([a, b])", - Partitioning::Hash( - vec![Arc::clone(&col_a), Arc::clone(&col_b), Arc::clone(&col_c)], - 4, - ), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), + "unknown partition expr", + Partitioning::Hash(vec![Arc::clone(&unknown)], 4), + fixture.key_distribution([0, 1]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ( + "unknown required expr", + fixture.hash_partitioning([0, 1], 4), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ( + "same unknown expr", + Partitioning::Hash(vec![Arc::clone(&unknown)], 4), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown)]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ( + "unknown partition expr is not a valid subset", + Partitioning::Hash(vec![Arc::clone(&unknown)], 4), + Distribution::KeyPartitioned(vec![Arc::clone(&unknown), fixture.col(0)]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ( + "empty hash partitioning", + Partitioning::Hash(vec![], 4), + fixture.key_distribution([0]), + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ), + ( + "empty key distribution", + fixture.hash_partitioning([0], 4), + Distribution::KeyPartitioned(vec![]), PartitioningSatisfaction::NotSatisfied, PartitioningSatisfaction::NotSatisfied, ), @@ -569,280 +1073,458 @@ mod tests { for (desc, partition, required, expected_with_subset, expected_without_subset) in test_cases { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" + assert_satisfaction( + desc, + &partition, + &required, + &fixture.eq_properties, + expected_with_subset, + expected_without_subset, ); } Ok(()) } + fn int_split_point(values: impl IntoIterator) -> SplitPoint { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::Int64(Some(value))) + .collect(), + ) + } + + fn assert_range_try_new_error( + ordering: LexOrdering, + split_points: Vec, + expected: &str, + ) { + let error = RangePartitioning::try_new(ordering, split_points) + .unwrap_err() + .to_string(); + assert!(error.contains(expected), "{error}"); + } + #[test] - fn test_partitioning_partial_overlap() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); - - let test_cases = vec![( - "Partial overlap: Hash([a, c]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_c)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a), Arc::clone(&col_b)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - )]; + fn test_range_partitioning_metadata() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + + let range_partitioning = + fixture.range([0], vec![int_split_point([10]), int_split_point([20])]); + assert_eq!(range_partitioning.ordering()[0].to_string(), "a@0 ASC"); + assert_eq!( + range_partitioning.split_points(), + &[int_split_point([10]), int_split_point([20])] + ); + let partitioning = Partitioning::Range(range_partitioning); + + assert_eq!(partitioning.partition_count(), 3); + assert_eq!( + partitioning.to_string(), + "Range([a@0 ASC], [(10), (20)], 3)" + ); - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); + Ok(()) + } - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } + #[test] + fn test_range_partitioning_try_new_validates_split_points() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let asc_a = fixture.range_ordering([0]); + let ordering_ab = fixture.range_ordering([0, 1]); + + assert_range_try_new_error( + ordering_ab.clone(), + vec![int_split_point([10])], + "split point 0 has width 1, but ordering has width 2", + ); + + RangePartitioning::try_new( + [fixture.range_sort_expr(0, SortOptions::new(true, false))].into(), + vec![int_split_point([20]), int_split_point([10])], + )?; + + assert_range_try_new_error( + asc_a, + vec![int_split_point([20]), int_split_point([10])], + "split points must be strictly ordered", + ); + + assert_range_try_new_error( + [fixture.range_sort_expr(0, SortOptions::new(false, false))].into(), + vec![ + SplitPoint::new(vec![ScalarValue::Int64(None)]), + int_split_point([10]), + ], + "split points must be strictly ordered", + ); + + RangePartitioning::try_new( + ordering_ab.clone(), + vec![int_split_point([10, 20]), int_split_point([10, 30])], + )?; + + assert_range_try_new_error( + ordering_ab, + vec![int_split_point([10, 30]), int_split_point([10, 20])], + "split points must be strictly ordered", + ); Ok(()) } #[test] - fn test_partitioning_no_overlap() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - Field::new("c", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let col_c: Arc = - Arc::new(Column::new_with_schema("c", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + fn test_range_partitioning_project_preserves_or_degrades() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let range_partitioning = fixture.range_partitioning_with_ordering( + [fixture.range_sort_expr(1, SortOptions::new(true, false))].into(), + vec![int_split_point([10])], + ); + + let keep_b_mapping = ProjectionMapping::from_indices(&[1], &fixture.schema)?; + let projected = + range_partitioning.project(&keep_b_mapping, &fixture.eq_properties); + assert_eq!( + projected.to_string(), + "Range([b@0 DESC NULLS LAST], [(10)], 2)" + ); + + let drop_b_mapping = ProjectionMapping::from_indices(&[0], &fixture.schema)?; + let projected = + range_partitioning.project(&drop_b_mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); - let test_cases = vec![ + Ok(()) + } + + #[test] + fn test_range_partitioning_project_degrades_if_ordering_collapses() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b"])?; + let target: Arc = Arc::new(Column::new("x", 0)); + let range_partitioning = + fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); + let mapping = ProjectionMapping::from_iter([ ( - "Hash([a]) vs Hash([b, c])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_b), - Arc::clone(&col_c), - ]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, + fixture.col(0), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), ), ( - "Hash([a, b]) vs Hash([c])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_c)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, + fixture.col(1), + ProjectionTargets::from(vec![(Arc::clone(&target), 0)]), ), - ]; + ]); - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); - - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } + let projected = range_partitioning.project(&mapping, &fixture.eq_properties); + let Partitioning::UnknownPartitioning(partition_count) = projected else { + panic!("expected UnknownPartitioning, got {projected:?}"); + }; + assert_eq!(partition_count, 2); Ok(()) } #[test] - fn test_partitioning_exact_match() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + fn range_partitioning_key_distribution_satisfaction() -> Result<()> { + let fixture = PartitioningTestFixture::int64(&["a", "b", "c"])?; + let range_a = fixture.range_partitioning([0], vec![int_split_point([10])]); + let range_ab = + fixture.range_partitioning([0, 1], vec![int_split_point([10, 100])]); + + assert_satisfaction( + "exact single key", + &range_a, + &fixture.key_distribution([0]), + &fixture.eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); + assert_satisfaction( + "exact compound key", + &range_ab, + &fixture.key_distribution([0, 1]), + &fixture.eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); + assert_satisfaction( + "subset key", + &range_a, + &fixture.key_distribution([0, 1]), + &fixture.eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + assert_satisfaction( + "incompatible key", + &range_a, + &fixture.key_distribution([1]), + &fixture.eq_properties, + PartitioningSatisfaction::NotSatisfied, + PartitioningSatisfaction::NotSatisfied, + ); + + let mut eq_properties = fixture.eq_properties.clone(); + eq_properties.add_equal_conditions(fixture.col(0), fixture.col(2))?; + assert_satisfaction( + "equivalent subset key", + &range_a, + &fixture.key_distribution([1, 2]), + &eq_properties, + PartitioningSatisfaction::Subset, + PartitioningSatisfaction::NotSatisfied, + ); + + let mut eq_properties = fixture.eq_properties.clone(); + eq_properties.add_equal_conditions(fixture.col(0), fixture.col(1))?; + assert_satisfaction( + "equivalent exact key", + &range_a, + &fixture.key_distribution([1]), + &eq_properties, + PartitioningSatisfaction::Exact, + PartitioningSatisfaction::Exact, + ); - let test_cases = vec![ - ( - "Hash([a, b]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), - PartitioningSatisfaction::Exact, - PartitioningSatisfaction::Exact, - ), - ( - "Hash([a]) vs Hash([a])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), - PartitioningSatisfaction::Exact, - PartitioningSatisfaction::Exact, - ), - ]; + Ok(()) + } +} - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); +#[cfg(all(test, feature = "proto"))] +mod ordering_proto_tests { + use std::sync::Arc; - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_physical_expr_common::sort_expr::{ + LexRequirement, PhysicalSortExpr, PhysicalSortRequirement, + sort_exprs_try_from_proto, sort_exprs_try_to_proto, + }; - Ok(()) + use crate::expressions::Column; + use crate::proto_test_util::{StubDecoder, StubEncoder}; + + fn schema() -> Schema { + Schema::new(vec![Field::new("a", DataType::Int32, false)]) } - #[test] - fn test_partitioning_unknown() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Int64, false), - ])); - - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let col_b: Arc = - Arc::new(Column::new_with_schema("b", &schema)?); - let unknown: Arc = Arc::new(UnKnownColumn::new("dropped")); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); + fn sort_expr(descending: bool, nulls_first: bool) -> PhysicalSortExpr { + PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions { + descending, + nulls_first, + }, + ) + } - let test_cases = vec![ - ( - "Hash([unknown]) vs Hash([a, b])", - Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::HashPartitioned(vec![ - Arc::clone(&col_a), - Arc::clone(&col_b), - ]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ( - "Hash([a, b]) vs Hash([unknown])", - Partitioning::Hash(vec![Arc::clone(&col_a), Arc::clone(&col_b)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&unknown)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ( - "Hash([unknown]) vs Hash([unknown])", - Partitioning::Hash(vec![Arc::clone(&unknown)], 4), - Distribution::HashPartitioned(vec![Arc::clone(&unknown)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ]; + #[test] + fn sort_exprs_round_trip_preserves_options_and_order() { + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let exprs = vec![sort_expr(true, false), sort_expr(false, true)]; + + let nodes = sort_exprs_try_to_proto(&exprs, &encode_ctx).unwrap(); + // `asc` is the inverse of `descending` on the wire. + assert_eq!( + nodes + .iter() + .map(|node| (node.asc, node.nulls_first)) + .collect::>(), + vec![(false, false), (true, true)] + ); + + let schema = schema(); + let decoder = StubDecoder::ok(); + let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let decoded = sort_exprs_try_from_proto(&nodes, &decode_ctx).unwrap(); + assert_eq!( + decoded.iter().map(|expr| expr.options).collect::>(), + exprs.iter().map(|expr| expr.options).collect::>() + ); + } - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); + #[test] + fn sort_exprs_accepts_owned_requirements() { + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let requirement = LexRequirement::from([PhysicalSortRequirement::new( + Arc::new(Column::new("a", 0)), + Some(SortOptions { + descending: true, + nulls_first: true, + }), + )]); + + let nodes = sort_exprs_try_to_proto( + requirement + .iter() + .map(|req| PhysicalSortExpr::from(req.clone())), + &encode_ctx, + ) + .unwrap(); + + assert_eq!(nodes.len(), 1); + assert!(!nodes[0].asc); + assert!(nodes[0].nulls_first); + } - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" - ); - } + #[test] + fn sort_exprs_propagate_encode_errors() { + let encoder = StubEncoder::failing_on(2); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let exprs = vec![sort_expr(false, false), sort_expr(true, true)]; - Ok(()) + let err = sort_exprs_try_to_proto(&exprs, &encode_ctx).unwrap_err(); + assert!(err.to_string().contains("stub encode failure on call 2")); } #[test] - fn test_partitioning_empty_hash() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + fn sort_exprs_reject_missing_inner_expr() { + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let mut nodes = + sort_exprs_try_to_proto(&[sort_expr(false, false)], &encode_ctx).unwrap(); + nodes[0].expr = None; + + let schema = schema(); + let decoder = StubDecoder::ok(); + let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let err = sort_exprs_try_from_proto(&nodes, &decode_ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalSortExpr is missing required field 'expr'") + ); + } +} - let col_a: Arc = - Arc::new(Column::new_with_schema("a", &schema)?); - let eq_properties = EquivalenceProperties::new(Arc::clone(&schema)); +/// Partition counts are `usize` in memory and `u64` on the wire, so every +/// counted [`Partitioning`] variant crosses a width boundary in both +/// directions. These pin that neither crossing wraps or panics. +#[cfg(all(test, feature = "proto"))] +mod partition_count_proto_tests { + use std::sync::Arc; - let test_cases = vec![ - ( - "Hash([]) vs Hash([a])", - Partitioning::Hash(vec![], 4), - Distribution::HashPartitioned(vec![Arc::clone(&col_a)]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ( - "Hash([a]) vs Hash([])", - Partitioning::Hash(vec![Arc::clone(&col_a)], 4), - Distribution::HashPartitioned(vec![]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ( - "Hash([]) vs Hash([])", - Partitioning::Hash(vec![], 4), - Distribution::HashPartitioned(vec![]), - PartitioningSatisfaction::NotSatisfied, - PartitioningSatisfaction::NotSatisfied, - ), - ]; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr_common::physical_expr::PhysicalExpr; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf; - for (desc, partition, required, expected_with_subset, expected_without_subset) in - test_cases - { - let result = partition.satisfaction(&required, &eq_properties, true); - assert_eq!( - result, expected_with_subset, - "Failed for {desc} with subset enabled" - ); + use super::{Partitioning, partition_count, wire_partition_count}; + use crate::expressions::Column; + use crate::proto_test_util::{StubDecoder, StubEncoder, column_node}; + + fn partitioning_node( + method: protobuf::partitioning::PartitionMethod, + ) -> protobuf::Partitioning { + protobuf::Partitioning { + partition_method: Some(method), + } + } - let result = partition.satisfaction(&required, &eq_properties, false); - assert_eq!( - result, expected_without_subset, - "Failed for {desc} with subset disabled" + /// The counted variants, each carrying `count`. `Range` is excluded: it + /// derives its partition count from its split points rather than reading + /// one off the wire. + fn counted_methods(count: u64) -> Vec { + use protobuf::partitioning::PartitionMethod; + + vec![ + PartitionMethod::RoundRobin(count), + PartitionMethod::Unknown(count), + PartitionMethod::Hash(protobuf::PhysicalHashRepartition { + hash_expr: vec![column_node("a")], + partition_count: count, + }), + ] + } + + #[test] + fn partition_count_round_trips_at_the_usize_ceiling() { + // `usize::MAX` is the largest count that can exist in memory, so it has + // to widen onto the wire and narrow back unchanged. + let wire = wire_partition_count(usize::MAX).unwrap(); + assert_eq!(wire, u64::try_from(usize::MAX).unwrap()); + assert_eq!(partition_count(wire).unwrap(), usize::MAX); + } + + #[test] + fn out_of_range_partition_count_is_reported_not_wrapped() { + // A count wider than the target's `usize` can only be reached by + // decoding on a narrower host than the one that encoded. That used to + // wrap (`as usize`) or panic (`unwrap`); it is an error now. On a + // 64-bit target every `u64` fits, so the same input has to decode + // losslessly instead of being rejected. + let narrowed = partition_count(u64::MAX); + + #[cfg(target_pointer_width = "64")] + assert_eq!(narrowed.unwrap(), usize::MAX); + + #[cfg(not(target_pointer_width = "64"))] + assert!( + narrowed + .unwrap_err() + .to_string() + .contains("Partition count 18446744073709551615 exceeds usize::MAX") + ); + } + + #[test] + fn try_from_proto_narrows_every_counted_variant() { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let decoder = StubDecoder::ok(); + let decode_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + for method in counted_methods(u64::MAX) { + let decoded = + Partitioning::try_from_proto(&partitioning_node(method), &decode_ctx); + + #[cfg(target_pointer_width = "64")] + assert_eq!(decoded.unwrap().unwrap().partition_count(), usize::MAX); + + #[cfg(not(target_pointer_width = "64"))] + assert!( + decoded + .unwrap_err() + .to_string() + .contains("exceeds usize::MAX") ); } + } - Ok(()) + #[test] + fn try_to_proto_widens_every_counted_variant() { + use protobuf::partitioning::PartitionMethod; + + let encoder = StubEncoder::ok(); + let encode_ctx = PhysicalExprEncodeCtx::new(&encoder); + let hash_key: Arc = Arc::new(Column::new("a", 0)); + + let encoded = [ + Partitioning::RoundRobinBatch(usize::MAX), + Partitioning::UnknownPartitioning(usize::MAX), + Partitioning::Hash(vec![hash_key], usize::MAX), + ] + .iter() + .map(|partitioning| { + match partitioning + .try_to_proto(&encode_ctx) + .unwrap() + .partition_method + { + Some(PartitionMethod::RoundRobin(n) | PartitionMethod::Unknown(n)) => n, + Some(PartitionMethod::Hash(hash)) => hash.partition_count, + other => panic!("expected a counted partition method, got {other:?}"), + } + }) + .collect::>(); + + // Every variant widens to the same wire value, with no truncation. + assert_eq!(encoded, vec![u64::try_from(usize::MAX).unwrap(); 3]); } } diff --git a/datafusion/physical-expr/src/physical_expr.rs b/datafusion/physical-expr/src/physical_expr.rs index 77ede76e1daa8..d45d0fe14902e 100644 --- a/datafusion/physical-expr/src/physical_expr.rs +++ b/datafusion/physical-expr/src/physical_expr.rs @@ -21,15 +21,18 @@ use crate::expressions::{self, Column}; use crate::{LexOrdering, PhysicalSortExpr, create_physical_expr}; use arrow::compute::SortOptions; -use arrow::datatypes::{Schema, SchemaRef}; +use arrow::datatypes::{DataType, Schema, SchemaRef}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{DFSchema, HashMap}; +use datafusion_common::{DFSchema, HashMap, ScalarValue, SplitPoint}; use datafusion_common::{Result, plan_err}; use datafusion_expr::execution_props::ExecutionProps; -use datafusion_expr::{Expr, SortExpr}; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion_expr::{Expr, Partitioning as LogicalPartitioning, SortExpr}; +use datafusion_expr_common::casts::try_cast_literal_to_type; use itertools::izip; // Exports: +use crate::{Partitioning, RangePartitioning}; pub(crate) use datafusion_physical_expr_common::physical_expr::PhysicalExpr; /// Adds the `offset` value to `Column` indices inside `expr`. This function is @@ -58,7 +61,7 @@ pub fn physical_exprs_contains( ) -> bool { physical_exprs .iter() - .any(|physical_expr| physical_expr.eq(expr)) + .any(|physical_expr| physical_expr.as_ref().eq(expr.as_ref())) } /// Checks whether the given physical expression slices are equal. @@ -66,7 +69,8 @@ pub fn physical_exprs_equal( lhs: &[Arc], rhs: &[Arc], ) -> bool { - lhs.len() == rhs.len() && izip!(lhs, rhs).all(|(lhs, rhs)| lhs.eq(rhs)) + lhs.len() == rhs.len() + && izip!(lhs, rhs).all(|(lhs, rhs)| lhs.as_ref().eq(rhs.as_ref())) } /// Checks whether the given physical expression slices are equal in the sense @@ -187,35 +191,150 @@ pub fn create_lex_ordering( exprs, &df_schema, execution_props, + &PhysicalPlanningContext::default(), )?)); } Ok(all_sort_orders) } /// Create a physical sort expression from a logical expression +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_physical_sort_expr( e: &SortExpr, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result { - create_physical_expr(&e.expr, input_dfschema, execution_props).map(|expr| { - let options = SortOptions::new(!e.asc, e.nulls_first); - PhysicalSortExpr::new(expr, options) - }) + create_physical_expr(&e.expr, input_dfschema, execution_props, planning_ctx).map( + |expr| { + let options = SortOptions::new(!e.asc, e.nulls_first); + PhysicalSortExpr::new(expr, options) + }, + ) } /// Create vector of physical sort expression from a vector of logical expression +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_physical_sort_exprs( exprs: &[SortExpr], input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result> { exprs .iter() - .map(|e| create_physical_sort_expr(e, input_dfschema, execution_props)) + .map(|e| { + create_physical_sort_expr(e, input_dfschema, execution_props, planning_ctx) + }) .collect() } +/// Create physical partitioning from logical partitioning. +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. +pub fn create_physical_partitioning( + partitioning: &LogicalPartitioning, + input_dfschema: &DFSchema, + execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, +) -> Result { + match partitioning { + LogicalPartitioning::RoundRobinBatch(n) => Ok(Partitioning::RoundRobinBatch(*n)), + LogicalPartitioning::Hash(exprs, partition_count) => { + let exprs = exprs + .iter() + .map(|expr| { + create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + ) + }) + .collect::>>()?; + Ok(Partitioning::Hash(exprs, *partition_count)) + } + LogicalPartitioning::Range(range) => { + let ordering = create_physical_sort_exprs( + range.ordering(), + input_dfschema, + execution_props, + planning_ctx, + )?; + let Some(ordering) = LexOrdering::new(ordering) else { + return plan_err!("Range partitioning requires non-empty ordering"); + }; + let split_points = normalize_range_split_points( + &ordering, + range.split_points(), + input_dfschema.as_arrow(), + )?; + let range = RangePartitioning::try_new(ordering, split_points)?; + Ok(Partitioning::Range(range)) + } + LogicalPartitioning::DistributeBy(_) => { + datafusion_common::not_impl_err!( + "Physical plan does not support DistributeBy partitioning" + ) + } + } +} + +fn normalize_range_split_points( + ordering: &LexOrdering, + split_points: &[SplitPoint], + schema: &Schema, +) -> Result> { + split_points + .iter() + .enumerate() + .map(|(split_idx, split_point)| { + let values = split_point + .values() + .iter() + .zip(ordering.iter()) + .enumerate() + .map(|(value_idx, (value, sort_expr))| { + let target_type = sort_expr.expr.data_type(schema)?; + normalize_range_split_point_value( + value, + &target_type, + split_idx, + value_idx, + ) + }) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect() +} + +fn normalize_range_split_point_value( + value: &ScalarValue, + target_type: &DataType, + split_idx: usize, + value_idx: usize, +) -> Result { + let value_type = value.data_type(); + if &value_type == target_type { + return Ok(value.clone()); + } + + if let Some(casted) = try_cast_literal_to_type(value, target_type) { + // Split points define physical partition boundaries, so normalization + // must reject casts that would change the advertised boundary. + if try_cast_literal_to_type(&casted, &value_type).as_ref() == Some(value) { + return Ok(casted); + } + } + + plan_err!( + "Range output partitioning split point {split_idx} value {value_idx} with type {value_type} cannot be represented exactly as ordering expression type {target_type}" + ) +} + pub fn add_offset_to_physical_sort_exprs( sort_exprs: impl IntoIterator, offset: isize, @@ -233,7 +352,7 @@ pub fn add_offset_to_physical_sort_exprs( mod tests { use super::*; - use crate::expressions::{BinaryExpr, Literal}; + use crate::expressions::{BinaryExpr, Literal, UnKnownColumn}; use crate::physical_expr::{ physical_exprs_bag_equal, physical_exprs_contains, physical_exprs_equal, }; @@ -279,6 +398,12 @@ mod tests { // below expressions are not inside physical_exprs assert!(!physical_exprs_contains(&physical_exprs, &col_c_expr)); assert!(!physical_exprs_contains(&physical_exprs, &lit1)); + + let unknown = Arc::new(UnKnownColumn::new("unknown")) as Arc; + assert!(!physical_exprs_contains( + std::slice::from_ref(&unknown), + &unknown + )); } #[test] @@ -309,6 +434,12 @@ mod tests { assert!(!physical_exprs_equal(&vec1, &vec3)); assert!(!physical_exprs_bag_equal(&vec1, &vec2)); assert!(!physical_exprs_bag_equal(&vec1, &vec3)); + + let unknown = Arc::new(UnKnownColumn::new("unknown")) as Arc; + assert!(!physical_exprs_equal( + std::slice::from_ref(&unknown), + std::slice::from_ref(&unknown) + )); } #[test] diff --git a/datafusion/physical-expr/src/planner.rs b/datafusion/physical-expr/src/planner.rs index d0d0508a106a5..f80d1b15bdc59 100644 --- a/datafusion/physical-expr/src/planner.rs +++ b/datafusion/physical-expr/src/planner.rs @@ -37,6 +37,7 @@ use datafusion_expr::expr::{ Alias, Cast, HigherOrderFunction, InList, Lambda, LambdaVariable, Placeholder, ScalarFunction, }; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; use datafusion_expr::var_provider::VarType; use datafusion_expr::var_provider::is_system_variables; use datafusion_expr::{ @@ -63,6 +64,7 @@ use datafusion_expr::{ /// # use datafusion_expr::{Expr, col, lit}; /// # use datafusion_physical_expr::create_physical_expr; /// # use datafusion_expr::execution_props::ExecutionProps; +/// # use datafusion_expr::physical_planning_context::PhysicalPlanningContext; /// // For a logical expression `a = 1`, we can create a physical expression /// let expr = col("a").eq(lit(1)); /// // To create a PhysicalExpr we need 1. a schema @@ -70,8 +72,11 @@ use datafusion_expr::{ /// let df_schema = DFSchema::try_from(schema).unwrap(); /// // 2. ExecutionProps /// let props = ExecutionProps::new(); -/// // We can now create a PhysicalExpr: -/// let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap(); +/// // We can now create a PhysicalExpr. Expressions with no scalar +/// // subqueries use an empty `PhysicalPlanningContext`: +/// let physical_expr = +/// create_physical_expr(&expr, &df_schema, &props, &PhysicalPlanningContext::default()) +/// .unwrap(); /// ``` /// /// # Example: Executing a PhysicalExpr to obtain [ColumnarValue] @@ -83,12 +88,15 @@ use datafusion_expr::{ /// # use datafusion_expr::{Expr, col, lit, ColumnarValue}; /// # use datafusion_physical_expr::create_physical_expr; /// # use datafusion_expr::execution_props::ExecutionProps; +/// # use datafusion_expr::physical_planning_context::PhysicalPlanningContext; /// # let expr = col("a").eq(lit(1)); /// # let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); /// # let df_schema = DFSchema::try_from(schema.clone()).unwrap(); /// # let props = ExecutionProps::new(); /// // Given a PhysicalExpr, for `a = 1` we can evaluate it against a RecordBatch like this: -/// let physical_expr = create_physical_expr(&expr, &df_schema, &props).unwrap(); +/// let physical_expr = +/// create_physical_expr(&expr, &df_schema, &props, &PhysicalPlanningContext::default()) +/// .unwrap(); /// // Input of [1,2,3] /// let input_batch = RecordBatch::try_from_iter(vec![ /// ("a", Arc::new(Int32Array::from(vec![1, 2, 3])) as _) @@ -111,11 +119,22 @@ use datafusion_expr::{ /// * `e` - The logical expression /// * `input_dfschema` - The DataFusion schema for the input, used to resolve `Column` references /// to qualified or unqualified fields by name. +/// * `execution_props` - Per-execution properties such as the query start time. +/// * `planning_ctx` - The [`PhysicalPlanningContext`] used to resolve +/// `Expr::ScalarSubquery` and `Expr::LambdaVariable` nodes. The physical +/// planner threads the subquery index map and shared results container from +/// its `ScalarSubqueryExec` construction into calls to +/// `create_physical_expr`; the lambda variable qualifiers are added by this +/// function itself as it descends into lambda bodies. Callers creating +/// physical expressions outside of physical planning should pass +/// `&PhysicalPlanningContext::default()`; converting a scalar subquery then returns a +/// planning error. #[cfg_attr(feature = "recursive_protection", recursive::recursive)] pub fn create_physical_expr( e: &Expr, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result> { let input_schema = input_dfschema.as_arrow(); @@ -131,7 +150,12 @@ pub fn create_physical_expr( new_metadata, ))) } else { - Ok(create_physical_expr(expr, input_dfschema, execution_props)?) + Ok(create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?) } } Expr::Column(c) => { @@ -167,12 +191,22 @@ pub fn create_physical_expr( Operator::IsNotDistinctFrom, lit(true), ); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsNotTrue(expr) => { let binary_op = binary_expr(expr.as_ref().clone(), Operator::IsDistinctFrom, lit(true)); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsFalse(expr) => { let binary_op = binary_expr( @@ -180,12 +214,22 @@ pub fn create_physical_expr( Operator::IsNotDistinctFrom, lit(false), ); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsNotFalse(expr) => { let binary_op = binary_expr(expr.as_ref().clone(), Operator::IsDistinctFrom, lit(false)); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsUnknown(expr) => { let binary_op = binary_expr( @@ -193,7 +237,12 @@ pub fn create_physical_expr( Operator::IsNotDistinctFrom, Expr::Literal(ScalarValue::Boolean(None), None), ); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::IsNotUnknown(expr) => { let binary_op = binary_expr( @@ -201,12 +250,27 @@ pub fn create_physical_expr( Operator::IsDistinctFrom, Expr::Literal(ScalarValue::Boolean(None), None), ); - create_physical_expr(&binary_op, input_dfschema, execution_props) + create_physical_expr( + &binary_op, + input_dfschema, + execution_props, + planning_ctx, + ) } Expr::BinaryExpr(BinaryExpr { left, op, right }) => { // Create physical expressions for left and right operands - let lhs = create_physical_expr(left, input_dfschema, execution_props)?; - let rhs = create_physical_expr(right, input_dfschema, execution_props)?; + let lhs = create_physical_expr( + left, + input_dfschema, + execution_props, + planning_ctx, + )?; + let rhs = create_physical_expr( + right, + input_dfschema, + execution_props, + planning_ctx, + )?; // Note that the logical planner is responsible // for type coercion on the arguments (e.g. if one // argument was originally Int32 and one was @@ -229,10 +293,18 @@ pub fn create_physical_expr( "LIKE does not support escape_char other than the backslash (\\)" ); } - let physical_expr = - create_physical_expr(expr, input_dfschema, execution_props)?; - let physical_pattern = - create_physical_expr(pattern, input_dfschema, execution_props)?; + let physical_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + let physical_pattern = create_physical_expr( + pattern, + input_dfschema, + execution_props, + planning_ctx, + )?; like( *negated, *case_insensitive, @@ -251,10 +323,18 @@ pub fn create_physical_expr( if escape_char.is_some() { return exec_err!("SIMILAR TO does not support escape_char yet"); } - let physical_expr = - create_physical_expr(expr, input_dfschema, execution_props)?; - let physical_pattern = - create_physical_expr(pattern, input_dfschema, execution_props)?; + let physical_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + let physical_pattern = create_physical_expr( + pattern, + input_dfschema, + execution_props, + planning_ctx, + )?; similar_to(*negated, *case_insensitive, physical_expr, physical_pattern) } Expr::Case(case) => { @@ -263,6 +343,7 @@ pub fn create_physical_expr( e.as_ref(), input_dfschema, execution_props, + planning_ctx, )?) } else { None @@ -272,10 +353,18 @@ pub fn create_physical_expr( .iter() .map(|(w, t)| (w.as_ref(), t.as_ref())) .unzip(); - let when_expr = - create_physical_exprs(when_expr, input_dfschema, execution_props)?; - let then_expr = - create_physical_exprs(then_expr, input_dfschema, execution_props)?; + let when_expr = create_physical_exprs( + when_expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + let then_expr = create_physical_exprs( + then_expr, + input_dfschema, + execution_props, + planning_ctx, + )?; let when_then_expr: Vec<(Arc, Arc)> = when_expr .iter() @@ -288,6 +377,7 @@ pub fn create_physical_expr( e.as_ref(), input_dfschema, execution_props, + planning_ctx, )?) } else { None @@ -295,7 +385,7 @@ pub fn create_physical_expr( Ok(expressions::case(expr, when_then_expr, else_expr)?) } Expr::Cast(Cast { expr, field }) => expressions::cast_with_target_field( - create_physical_expr(expr, input_dfschema, execution_props)?, + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, input_schema, Arc::clone(field), None, @@ -314,31 +404,45 @@ pub fn create_physical_expr( } expressions::try_cast( - create_physical_expr(expr, input_dfschema, execution_props)?, + create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?, input_schema, field.data_type().clone(), ) } - Expr::Not(expr) => { - expressions::not(create_physical_expr(expr, input_dfschema, execution_props)?) - } + Expr::Not(expr) => expressions::not(create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?), Expr::Negative(expr) => expressions::negative( - create_physical_expr(expr, input_dfschema, execution_props)?, + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx)?, input_schema, ), Expr::IsNull(expr) => expressions::is_null(create_physical_expr( expr, input_dfschema, execution_props, + planning_ctx, )?), Expr::IsNotNull(expr) => expressions::is_not_null(create_physical_expr( expr, input_dfschema, execution_props, + planning_ctx, )?), Expr::ScalarFunction(ScalarFunction { func, args }) => { - let physical_args = - create_physical_exprs(args, input_dfschema, execution_props)?; + let physical_args = create_physical_exprs( + args, + input_dfschema, + execution_props, + planning_ctx, + )?; let config_options = match execution_props.config_options.as_ref() { Some(config_options) => Arc::clone(config_options), None => Arc::new(ConfigOptions::default()), @@ -357,9 +461,20 @@ pub fn create_physical_expr( low, high, }) => { - let value_expr = create_physical_expr(expr, input_dfschema, execution_props)?; - let low_expr = create_physical_expr(low, input_dfschema, execution_props)?; - let high_expr = create_physical_expr(high, input_dfschema, execution_props)?; + let value_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; + let low_expr = + create_physical_expr(low, input_dfschema, execution_props, planning_ctx)?; + let high_expr = create_physical_expr( + high, + input_dfschema, + execution_props, + planning_ctx, + )?; // rewrite the between into the two binary operators let binary_expr = binary( @@ -394,17 +509,25 @@ pub fn create_physical_expr( Ok(expressions::lit(ScalarValue::Boolean(None))) } _ => { - let value_expr = - create_physical_expr(expr, input_dfschema, execution_props)?; + let value_expr = create_physical_expr( + expr, + input_dfschema, + execution_props, + planning_ctx, + )?; - let list_exprs = - create_physical_exprs(list, input_dfschema, execution_props)?; + let list_exprs = create_physical_exprs( + list, + input_dfschema, + execution_props, + planning_ctx, + )?; expressions::in_list(value_expr, list_exprs, negated, input_schema) } }, Expr::ScalarSubquery(sq) => { - match execution_props.subquery_indexes.get(sq) { - Some(&index) => { + match planning_ctx.index_of(sq) { + Some(index) => { let schema = sq.subquery.schema(); if schema.fields().len() != 1 { return plan_err!( @@ -418,7 +541,7 @@ pub fn create_physical_expr( dt, nullable, index, - execution_props.subquery_results.clone(), + planning_ctx.results().clone(), ))) } None => { @@ -491,13 +614,23 @@ pub fn create_physical_expr( input_dfschema.metadata().clone(), )?; - let execution_props = execution_props + let planning_ctx = planning_ctx .clone() .with_qualified_lambda_variables(&qualifier, &lambda.params); - create_physical_expr(arg, &lambda_schema, &execution_props) + create_physical_expr( + arg, + &lambda_schema, + execution_props, + &planning_ctx, + ) } - _ => create_physical_expr(arg, input_dfschema, execution_props), + _ => create_physical_expr( + arg, + input_dfschema, + execution_props, + planning_ctx, + ), }) .collect::>()?; @@ -515,7 +648,7 @@ pub fn create_physical_expr( } Expr::Lambda(Lambda { params, body }) => expressions::lambda( params, - create_physical_expr(body, input_dfschema, execution_props)?, + create_physical_expr(body, input_dfschema, execution_props, planning_ctx)?, ), Expr::LambdaVariable(LambdaVariable { name, @@ -526,12 +659,14 @@ pub fn create_physical_expr( plan_datafusion_err!("unresolved LambdaVariable {name}") })?; - let qualifier = execution_props - .lambda_variable_qualifier - .get(name) - .ok_or_else(|| { - plan_datafusion_err!("qualifier for lambda variable {name} not found") - })?; + let qualifier = + planning_ctx + .lambda_variable_qualifier(name) + .ok_or_else(|| { + plan_datafusion_err!( + "qualifier for lambda variable {name} not found" + ) + })?; let index = input_dfschema .index_of_column_by_name(Some(qualifier), name) @@ -572,17 +707,22 @@ pub fn create_physical_expr( } /// Create vector of Physical Expression from a vector of logical expression +/// +/// See [`create_physical_expr`] for details on the `planning_ctx` argument. pub fn create_physical_exprs<'a, I>( exprs: I, input_dfschema: &DFSchema, execution_props: &ExecutionProps, + planning_ctx: &PhysicalPlanningContext, ) -> Result>> where I: IntoIterator, { exprs .into_iter() - .map(|expr| create_physical_expr(expr, input_dfschema, execution_props)) + .map(|expr| { + create_physical_expr(expr, input_dfschema, execution_props, planning_ctx) + }) .collect() } @@ -591,7 +731,13 @@ pub fn logical2physical(expr: &Expr, schema: &Schema) -> Arc { // TODO this makes a deep copy of the Schema. Should take SchemaRef instead and avoid deep copy let df_schema = schema.clone().to_dfschema().unwrap(); let execution_props = ExecutionProps::new(); - create_physical_expr(expr, &df_schema, &execution_props).unwrap() + create_physical_expr( + expr, + &df_schema, + &execution_props, + &PhysicalPlanningContext::default(), + ) + .unwrap() } #[cfg(test)] @@ -608,7 +754,12 @@ mod tests { fn lower_cast_expr(expr: &Expr, schema: &Schema) -> Result> { let df_schema = DFSchema::try_from(schema.clone())?; - create_physical_expr(expr, &df_schema, &ExecutionProps::new()) + create_physical_expr( + expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + ) } fn as_planner_cast(physical: &Arc) -> &expressions::CastExpr { @@ -623,7 +774,12 @@ mod tests { let schema = Schema::new(vec![Field::new("letter", DataType::Utf8, false)]); let df_schema = DFSchema::try_from_qualified_schema("data", &schema)?; - let p = create_physical_expr(&expr, &df_schema, &ExecutionProps::new())?; + let p = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; let batch = RecordBatch::try_new( Arc::new(schema), @@ -728,8 +884,12 @@ mod tests { let df_schema = DFSchema::try_from(schema)?; // This should not stack overflow - let _physical_expr = - create_physical_expr(&expr, &df_schema, &ExecutionProps::new())?; + let _physical_expr = create_physical_expr( + &expr, + &df_schema, + &ExecutionProps::new(), + &PhysicalPlanningContext::default(), + )?; Ok(()) } diff --git a/datafusion/physical-expr/src/projection.rs b/datafusion/physical-expr/src/projection.rs index ca999479fa916..0e8876f017379 100644 --- a/datafusion/physical-expr/src/projection.rs +++ b/datafusion/physical-expr/src/projection.rs @@ -69,6 +69,14 @@ impl PartialEq for ProjectionExpr { impl Eq for ProjectionExpr {} +/// Enables [`ProjectionExpr`] to be treated as a reference to its wrapped +/// [`Arc`] using [`AsRef::as_ref`]. +impl AsRef> for ProjectionExpr { + fn as_ref(&self) -> &Arc { + &self.expr + } +} + impl std::fmt::Display for ProjectionExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { if self.expr.to_string() == self.alias { @@ -539,6 +547,56 @@ impl ProjectionExprs { }) } + /// Create a new [`Projector`] using field and schema metadata from + /// `projected_schema`. + /// + /// Field names, data types, and nullability are still derived from the physical + /// projection expressions and `input_schema`; only field and schema metadata are + /// taken from `projected_schema`. + /// + /// # Errors + /// + /// Returns an error if the projection cannot be applied to `input_schema`, or if + /// `projected_schema` has a different number of fields than the projection. + pub fn make_projector_with_schema_metadata( + &self, + input_schema: &Schema, + projected_schema: &Schema, + ) -> Result { + let output_schema = self.project_schema(input_schema)?; + if output_schema.fields().len() != projected_schema.fields().len() { + return Err(internal_datafusion_err!( + "Projection has {} output fields but metadata schema has {} fields", + output_schema.fields().len(), + projected_schema.fields().len() + )); + } + + let fields = output_schema + .fields() + .iter() + .zip(projected_schema.fields()) + .map(|(field, projected_field)| { + Arc::new( + field + .as_ref() + .clone() + .with_metadata(projected_field.metadata().clone()), + ) + }) + .collect::>(); + let output_schema = Arc::new(Schema::new_with_metadata( + fields, + projected_schema.metadata().clone(), + )); + + Ok(Projector { + projection: self.clone(), + output_schema, + expression_metrics: None, + }) + } + pub fn create_expression_metrics( &self, metrics: &ExecutionPlanMetricsSet, @@ -661,7 +719,7 @@ impl ProjectionExprs { for proj_expr in self.exprs.iter() { let expr = &proj_expr.expr; let col_stats = if let Some(col) = expr.downcast_ref::() { - std::mem::take(&mut stats.column_statistics[col.index()]) + column_statistics_at(&stats.column_statistics, col.index()) } else if let Some(literal) = expr.downcast_ref::() { // Handle literal expressions (constants) by calculating proper statistics let data_type = expr.data_type(output_schema)?; @@ -725,6 +783,60 @@ impl ProjectionExprs { stats.column_statistics = column_statistics; Ok(stats) } + + /// Returns the output position of `column` if this projection contains it. + /// + /// This only matches projection expressions that are exactly [`Column`] expressions. + /// Computed expressions, even if they reference `column`, do not match. The + /// comparison uses [`Column`] equality, so both the name and index must match. + /// If the same column appears more than once, this returns the first matching + /// position. + /// + /// # Example + /// + /// ```rust + /// use datafusion_common::ScalarValue; + /// use datafusion_physical_expr::expressions::{Column, Literal}; + /// use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; + /// use std::sync::Arc; + /// + /// let projection = ProjectionExprs::new([ + /// ProjectionExpr::new(Arc::new(Column::new("b", 1)), "b"), + /// ProjectionExpr::new( + /// Arc::new(Literal::new(ScalarValue::Int32(Some(42)))), + /// "answer", + /// ), + /// ProjectionExpr::new(Arc::new(Column::new("a", 0)), "a"), + /// ]); + /// + /// assert_eq!( + /// projection.projected_column_position(&Column::new("b", 1)), + /// Some(0) + /// ); + /// assert_eq!( + /// projection.projected_column_position(&Column::new("a", 0)), + /// Some(2) + /// ); + /// + /// // The literal projection is not a Column expression. + /// assert_eq!( + /// projection.projected_column_position(&Column::new("answer", 1)), + /// None + /// ); + /// + /// // Columns not present in the projection also return None. + /// assert_eq!( + /// projection.projected_column_position(&Column::new("c", 2)), + /// None + /// ); + /// ``` + pub fn projected_column_position(&self, column: &Column) -> Option { + self.iter().position(|expr| { + expr.expr + .downcast_ref::() + .is_some_and(|projected| projected == column) + }) + } } /// Propagate column statistics through CAST projections. Other expressions @@ -736,7 +848,7 @@ fn project_column_statistics_through_expr( column_stats: &[ColumnStatistics], ) -> ColumnStatistics { if let Some(col) = expr.downcast_ref::() { - return column_stats[col.index()].clone(); + return column_statistics_at(column_stats, col.index()); } let Some(cast_expr) = expr.downcast_ref::() else { return ColumnStatistics::new_unknown(); @@ -744,6 +856,22 @@ fn project_column_statistics_through_expr( let inner_stats = project_column_statistics_through_expr(cast_expr.expr.as_ref(), column_stats); let target_type = cast_expr.cast_type(); + + // A cast whose source values are already of the target `DataType` never + // changes any value -- see `cast_array_by_name`'s same-type fast path in + // `ColumnarValue::cast_to`. In that case every statistic, not just + // min/max, carries over unchanged (this is what a cast that only + // re-stamps a column's nullability, as `UnionExec`/`InterleaveExec` + // insert, looks like here). + let already_target_type = matches!( + (inner_stats.min_value.get_value(), inner_stats.max_value.get_value()), + (Some(min), Some(max)) + if min.data_type() == *target_type && max.data_type() == *target_type + ); + if already_target_type { + return inner_stats; + } + ColumnStatistics { min_value: inner_stats .min_value @@ -760,6 +888,16 @@ fn project_column_statistics_through_expr( } } +fn column_statistics_at( + column_stats: &[ColumnStatistics], + index: usize, +) -> ColumnStatistics { + column_stats + .get(index) + .cloned() + .unwrap_or_else(ColumnStatistics::new_unknown) +} + impl<'a> IntoIterator for &'a ProjectionExprs { type Item = &'a ProjectionExpr; type IntoIter = std::slice::Iter<'a, ProjectionExpr>; @@ -963,8 +1101,6 @@ pub fn update_expr( return Ok(Transformed::no(expr)); }; if unproject { - state = RewriteState::RewrittenValid; - // Update the index of `column`: let projected_expr = projected_exprs.get(column.index()).ok_or_else(|| { internal_datafusion_err!( "Column index {} out of bounds for projected expressions of length {}", @@ -972,6 +1108,17 @@ pub fn update_expr( projected_exprs.len() ) })?; + // Skip rebuilding the parent if substituting with an equal + // Column (e.g. pass-through `c0@0` -> `c0@0` during chained + // projection collapse). Without this, every CASE/BinaryExpr + // containing such a Column is reconstructed unnecessarily. + if let Some(projected_col) = + projected_expr.expr.downcast_ref::() + && projected_col == column + { + return Ok(Transformed::no(expr)); + } + state = RewriteState::RewrittenValid; Ok(Transformed::yes(Arc::clone(&projected_expr.expr))) } else { // default to invalid, in case we can't find the relevant column @@ -1504,8 +1651,6 @@ pub(crate) mod tests { vec![("a_new", option_asc), ("b_new", option_asc)], // [a_new ASC, d_new ASC] vec![("a_new", option_asc), ("d_new", option_asc)], - // [a_new ASC, b+d ASC] - vec![("a_new", option_asc), ("b+d", option_asc)], ], ), // ------- TEST CASE 8 ---------- @@ -1587,12 +1732,6 @@ pub(crate) mod tests { ("b_new", option_asc), ("c_new", option_asc), ], - // [a_new ASC, b_new ASC, c+d ASC] - vec![ - ("a_new", option_asc), - ("b_new", option_asc), - ("c+d", option_asc), - ], ], ), // ------- TEST CASE 11 ---------- @@ -1614,8 +1753,6 @@ pub(crate) mod tests { vec![ // [a_new ASC, b_new ASC] vec![("a_new", option_asc), ("b_new", option_asc)], - // [a_new ASC, b + d ASC] - vec![("a_new", option_asc), ("b+d", option_asc)], ], ), // ------- TEST CASE 12 ---------- @@ -1697,30 +1834,12 @@ pub(crate) mod tests { ], // expected vec![ - // [a_new ASC, d_new ASC, b+e ASC] - vec![ - ("a_new", option_asc), - ("d_new", option_asc), - ("b+e", option_asc), - ], - // [d_new ASC, a_new ASC, b+e ASC] - vec![ - ("d_new", option_asc), - ("a_new", option_asc), - ("b+e", option_asc), - ], - // [c_new ASC, d_new ASC, b+e ASC] - vec![ - ("c_new", option_asc), - ("d_new", option_asc), - ("b+e", option_asc), - ], - // [d_new ASC, c_new ASC, b+e ASC] - vec![ - ("d_new", option_asc), - ("c_new", option_asc), - ("b+e", option_asc), - ], + // [a_new ASC] + vec![("a_new", option_asc)], + // [c_new ASC] + vec![("c_new", option_asc)], + // [d_new ASC] + vec![("d_new", option_asc)], ], ), // ------- TEST CASE 15 ---------- @@ -1742,12 +1861,8 @@ pub(crate) mod tests { ], // expected vec![ - // [a_new ASC, d_new ASC, b+e ASC] - vec![ - ("a_new", option_asc), - ("c_new", option_asc), - ("a+b", option_asc), - ], + // [a_new ASC, c_new ASC] + vec![("a_new", option_asc), ("c_new", option_asc)], ], ), // ------- TEST CASE 16 ---------- @@ -1772,8 +1887,6 @@ pub(crate) mod tests { vec![ // [a_new ASC, b_new ASC] vec![("a_new", option_asc), ("b_new", option_asc)], - // [a_new ASC, b_new ASC] - vec![("a_new", option_asc), ("b+e", option_asc)], // [c_new ASC, b_new DESC] vec![("c_new", option_asc), ("b_new", option_desc)], ], @@ -2046,7 +2159,6 @@ pub(crate) mod tests { let projection_mapping = ProjectionMapping::try_new(proj_exprs, &schema)?; let output_schema = output_schema(&projection_mapping, &schema)?; - let col_a_plus_b_new = &col("a+b", &output_schema)?; let col_c_new = &col("c_new", &output_schema)?; let col_d_new = &col("d_new", &output_schema)?; @@ -2064,18 +2176,10 @@ pub(crate) mod tests { vec![], // expected vec![ - // [d_new ASC, c_new ASC, a+b ASC] - vec![ - (col_d_new, option_asc), - (col_c_new, option_asc), - (col_a_plus_b_new, option_asc), - ], - // [c_new ASC, d_new ASC, a+b ASC] - vec![ - (col_c_new, option_asc), - (col_d_new, option_asc), - (col_a_plus_b_new, option_asc), - ], + // [c_new ASC] + vec![(col_c_new, option_asc)], + // [d_new ASC] + vec![(col_d_new, option_asc)], ], ), // ---------- TEST CASE 2 ------------ @@ -2091,18 +2195,10 @@ pub(crate) mod tests { vec![(col_e, col_a)], // expected vec![ - // [d_new ASC, c_new ASC, a+b ASC] - vec![ - (col_d_new, option_asc), - (col_c_new, option_asc), - (col_a_plus_b_new, option_asc), - ], - // [c_new ASC, d_new ASC, a+b ASC] - vec![ - (col_c_new, option_asc), - (col_d_new, option_asc), - (col_a_plus_b_new, option_asc), - ], + // [c_new ASC] + vec![(col_c_new, option_asc)], + // [d_new ASC] + vec![(col_d_new, option_asc)], ], ), // ---------- TEST CASE 3 ------------ @@ -2193,6 +2289,43 @@ pub(crate) mod tests { Schema::new(vec![field_0, field_1, field_2]) } + #[test] + fn test_projected_column_position_returns_output_position() { + let projection = ProjectionExprs::new([ + ProjectionExpr::new(Arc::new(Column::new("col2", 2)), "col2"), + ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "col0"), + ]); + + assert_eq!( + projection.projected_column_position(&Column::new("col2", 2)), + Some(0) + ); + assert_eq!( + projection.projected_column_position(&Column::new("col0", 0)), + Some(1) + ); + } + + #[test] + fn test_projected_column_position_returns_none_for_non_column_or_missing() { + let projection = ProjectionExprs::new([ + ProjectionExpr::new( + Arc::new(Literal::new(ScalarValue::Int64(Some(42)))), + "col1", + ), + ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "col0"), + ]); + + assert_eq!( + projection.projected_column_position(&Column::new("col1", 1)), + None + ); + assert_eq!( + projection.projected_column_position(&Column::new("col2", 2)), + None + ); + } + #[test] fn test_stats_projection_columns_only() { let source = get_stats(); @@ -2825,6 +2958,35 @@ pub(crate) mod tests { Ok(()) } + #[test] + fn test_project_statistics_with_same_type_cast_is_exact_passthrough() -> Result<()> { + // A cast to the column's own `DataType` (e.g. one that only re-stamps + // nullability via `CastExpr::new_with_target_field`, as `UnionExec`/ + // `InterleaveExec` insert) never changes any value, so every + // statistic -- not just min/max -- should carry over unchanged. + let input_stats = get_stats(); + let col0_stats = input_stats.column_statistics[0].clone(); + let input_schema = get_schema(); + + let projection = ProjectionExprs::new(vec![ProjectionExpr { + expr: Arc::new(CastExpr::new( + Arc::new(Column::new("col0", 0)), + DataType::Int64, + None, + )), + alias: "casted".to_string(), + }]); + + let output_stats = projection.project_statistics( + input_stats, + &projection.project_schema(&input_schema)?, + )?; + + assert_eq!(output_stats.column_statistics[0], col0_stats); + + Ok(()) + } + #[test] fn test_project_statistics_with_cast() -> Result<()> { let input_stats = get_stats(); @@ -2857,6 +3019,107 @@ pub(crate) mod tests { Ok(()) } + #[test] + fn test_project_statistics_duplicate_column() -> Result<()> { + let input_stats = get_stats(); + let col0 = input_stats.column_statistics[0].clone(); + let projection = ProjectionExprs::new([ + ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "a"), + ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "b"), + ]); + + let output_schema = projection.project_schema(&get_schema())?; + let output_stats = projection.project_statistics(input_stats, &output_schema)?; + + assert_eq!(output_stats.column_statistics, vec![col0.clone(), col0]); + Ok(()) + } + + #[test] + fn test_project_statistics_column_and_cast() -> Result<()> { + let input_stats = get_stats(); + let col0 = input_stats.column_statistics[0].clone(); + let projection = ProjectionExprs::new([ + ProjectionExpr::new(Arc::new(Column::new("col0", 0)), "num"), + ProjectionExpr::new( + Arc::new(CastExpr::new( + Arc::new(Column::new("col0", 0)), + DataType::Int32, + None, + )), + "casted", + ), + ]); + + let output_schema = projection.project_schema(&get_schema())?; + let output_stats = projection.project_statistics(input_stats, &output_schema)?; + + assert_eq!(output_stats.column_statistics[0], col0); + assert_eq!( + output_stats.column_statistics[1], + ColumnStatistics { + min_value: Precision::Exact(ScalarValue::Int32(Some(-4))), + max_value: Precision::Exact(ScalarValue::Int32(Some(21))), + distinct_count: Precision::Exact(5), + null_count: Precision::Exact(0), + sum_value: Precision::Absent, + byte_size: Precision::Absent, + } + ); + + Ok(()) + } + + #[test] + fn test_project_statistics_missing_column_stats_are_unknown() -> Result<()> { + let mut input_stats = get_stats(); + let input_schema = get_schema(); + input_stats.column_statistics.truncate(2); + + // The schema has col2, but the statistics do not. This can happen for + // source-provided virtual columns that are available at execution time + // but not represented in file-level statistics. + let projection = ProjectionExprs::new(vec![ + ProjectionExpr { + expr: Arc::new(Column::new("col2", 2)), + alias: "virtual_col".to_string(), + }, + ProjectionExpr { + expr: Arc::new(CastExpr::new( + Arc::new(Column::new("col2", 2)), + DataType::Float64, + None, + )), + alias: "casted_virtual_col".to_string(), + }, + ProjectionExpr { + expr: Arc::new(Column::new("col0", 0)), + alias: "physical_col".to_string(), + }, + ]); + + let output_stats = projection.project_statistics( + input_stats, + &projection.project_schema(&input_schema)?, + )?; + + assert_eq!(output_stats.column_statistics.len(), 3); + assert_eq!( + output_stats.column_statistics[0], + ColumnStatistics::new_unknown() + ); + assert_eq!( + output_stats.column_statistics[1], + ColumnStatistics::new_unknown() + ); + assert_eq!( + output_stats.column_statistics[2].max_value, + Precision::Exact(ScalarValue::Int64(Some(21))) + ); + + Ok(()) + } + #[test] fn test_project_statistics_primitive_width_only() -> Result<()> { let input_stats = get_stats(); diff --git a/datafusion/physical-expr/src/proto_test_util.rs b/datafusion/physical-expr/src/proto_test_util.rs new file mode 100644 index 0000000000000..ab280335800b9 --- /dev/null +++ b/datafusion/physical-expr/src/proto_test_util.rs @@ -0,0 +1,141 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Shared test helpers for proto serialization / deserialization in expression unit tests +//! without depending on `datafusion-proto` (which would create circular deps). + +use std::cell::Cell; +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{DataFusionError, Result}; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecode; +use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncode; +use datafusion_proto_models::protobuf::{self, PhysicalExprNode, physical_expr_node}; + +use crate::expressions::Column; + +/// A proto node for a `Column`, useful as a stand-in child node when building +/// an expression's proto representation in tests. +pub(crate) fn column_node(name: &str) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: name.to_string(), + index: 0, + }, + )), + } +} + +/// Decoder stub for driving `try_from_proto`: returns a fixed `Column` for each +/// child node, optionally failing on the Nth `decode` call so the +/// `ctx.decode(..)?` error arms can be exercised. +pub(crate) struct StubDecoder { + fail_on_call: Option, + calls: Cell, +} + +impl StubDecoder { + /// Always succeeds, returning a placeholder `Column` per child. + pub(crate) fn ok() -> Self { + Self { + fail_on_call: None, + calls: Cell::new(0), + } + } + + /// Fails on the `call`-th invocation (1-based), succeeding otherwise. + pub(crate) fn failing_on(call: usize) -> Self { + Self { + fail_on_call: Some(call), + calls: Cell::new(0), + } + } +} + +impl PhysicalExprDecode for StubDecoder { + fn decode( + &self, + _node: &PhysicalExprNode, + _schema: &Schema, + ) -> Result> { + let call = self.calls.get() + 1; + self.calls.set(call); + if Some(call) == self.fail_on_call { + return Err(DataFusionError::Internal(format!( + "stub decode failure on call {call}" + ))); + } + Ok(Arc::new(Column::new("decoded", 0))) + } +} + +/// Decoder that must never run: used to assert that the reject paths of a +/// `try_from_proto` (wrong node, missing child) bail out before decoding. +pub(crate) struct UnreachableDecoder; + +impl PhysicalExprDecode for UnreachableDecoder { + fn decode( + &self, + _node: &PhysicalExprNode, + _schema: &Schema, + ) -> Result> { + unreachable!("decode must not be reached when the node is rejected") + } +} + +/// Encoder stub for driving `try_to_proto`: emits a placeholder `Column` node +/// for each child, optionally failing on the Nth `encode` call so the +/// `ctx.encode_child(..)?` error arms can be exercised. +pub(crate) struct StubEncoder { + fail_on_call: Option, + calls: Cell, +} + +impl StubEncoder { + /// Always succeeds, emitting a placeholder `Column` node per child. + pub(crate) fn ok() -> Self { + Self { + fail_on_call: None, + calls: Cell::new(0), + } + } + + /// Fails on the `call`-th invocation (1-based), succeeding otherwise. + pub(crate) fn failing_on(call: usize) -> Self { + Self { + fail_on_call: Some(call), + calls: Cell::new(0), + } + } +} + +impl PhysicalExprEncode for StubEncoder { + fn encode(&self, _expr: &Arc) -> Result { + let call = self.calls.get() + 1; + self.calls.set(call); + if Some(call) == self.fail_on_call { + return Err(DataFusionError::Internal(format!( + "stub encode failure on call {call}" + ))); + } + Ok(column_node("child")) + } +} diff --git a/datafusion/physical-expr/src/scalar_function.rs b/datafusion/physical-expr/src/scalar_function.rs index 418d005c971ea..6a5ab219aa8dd 100644 --- a/datafusion/physical-expr/src/scalar_function.rs +++ b/datafusion/physical-expr/src/scalar_function.rs @@ -316,6 +316,7 @@ impl PhysicalExpr for ScalarFunctionExpr { fn get_properties(&self, children: &[ExprProperties]) -> Result { let sort_properties = self.fun.output_ordering(children)?; let preserves_lex_ordering = self.fun.preserves_lex_ordering(children)?; + let strictly_order_preserving = self.fun.strictly_order_preserving(children)?; let children_range = children .iter() .map(|props| &props.range) @@ -326,6 +327,7 @@ impl PhysicalExpr for ScalarFunctionExpr { sort_properties, range, preserves_lex_ordering, + strictly_order_preserving, }) } diff --git a/datafusion/physical-expr/src/scalar_subquery.rs b/datafusion/physical-expr/src/scalar_subquery.rs index ea00847151e66..473b52a5cb45c 100644 --- a/datafusion/physical-expr/src/scalar_subquery.rs +++ b/datafusion/physical-expr/src/scalar_subquery.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field, FieldRef, Schema}; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_datafusion_err}; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_expr_common::columnar_value::ColumnarValue; use datafusion_expr_common::sort_properties::{ExprProperties, SortProperties}; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; @@ -59,22 +59,34 @@ impl ScalarSubqueryExpr { } } + pub fn results(&self) -> &ScalarSubqueryResults { + &self.results + } + + #[deprecated( + since = "55.0.0", + note = "was only used for proto serialization, which no longer needs it; use `return_field` for type/nullability. It will be removed in 61.0.0 or 6 months after 55.0.0 is released, whichever is longer." + )] pub fn data_type(&self) -> &DataType { &self.data_type } + #[deprecated( + since = "55.0.0", + note = "was only used for proto serialization, which no longer needs it; use `return_field` for type/nullability. It will be removed in 61.0.0 or 6 months after 55.0.0 is released, whichever is longer." + )] pub fn nullable(&self) -> bool { self.nullable } /// Returns the index of this subquery in the shared results container. + #[deprecated( + since = "55.0.0", + note = "was only used for proto serialization, which no longer needs it. It will be removed in 61.0.0 or 6 months after 55.0.0 is released, whichever is longer." + )] pub fn index(&self) -> SubqueryIndex { self.index } - - pub fn results(&self) -> &ScalarSubqueryResults { - &self.results - } } impl fmt::Display for ScalarSubqueryExpr { @@ -139,6 +151,69 @@ impl PhysicalExpr for ScalarSubqueryExpr { fn fmt_sql(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "(scalar subquery)") } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarSubquery( + protobuf::PhysicalScalarSubqueryExprNode { + data_type: Some((&self.data_type).try_into()?), + nullable: self.nullable, + index: u32::try_from(self.index.as_usize()).map_err(|_| { + internal_datafusion_err!( + "scalar subquery index {} does not fit in u32", + self.index.as_usize() + ) + })?, + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl ScalarSubqueryExpr { + /// Reconstruct a [`ScalarSubqueryExpr`] from its protobuf representation. + /// + /// Unlike other expressions, this takes a third argument: the shared + /// [`ScalarSubqueryResults`] container. That container is a runtime-only + /// `Arc` shared with the surrounding `ScalarSubqueryExec` and is not part of + /// the wire format, so it cannot be reconstructed here or carried on the + /// decode context (which lives in a crate that cannot depend on + /// `datafusion-expr`). The match arm in `from_proto.rs` fetches it from the + /// plan-level decode context and passes it in. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + results: &ScalarSubqueryResults, + ) -> Result> { + use datafusion_physical_expr_common::expect_expr_variant; + use datafusion_physical_expr_common::physical_expr::proto_decode::require_proto_field; + use datafusion_proto_models::protobuf; + + let sq = expect_expr_variant!( + node, + protobuf::physical_expr_node::ExprType::ScalarSubquery, + "ScalarSubqueryExpr", + ); + let data_type = require_proto_field( + sq.data_type.as_ref(), + "ScalarSubqueryExpr", + "data_type", + )? + .try_into()?; + Ok(Arc::new(ScalarSubqueryExpr::new( + data_type, + sq.nullable, + SubqueryIndex::new(sq.index as usize), + results.clone(), + ))) + } } #[cfg(test)] @@ -238,3 +313,123 @@ mod tests { assert_ne!(e1a, e3); } } + +/// Tests for the `try_to_proto` / `try_from_proto` hooks. +#[cfg(all(test, feature = "proto"))] +mod proto_tests { + use super::*; + use crate::proto_test_util::{StubEncoder, UnreachableDecoder, column_node}; + use datafusion_common::DataFusionError; + use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; + use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; + use datafusion_proto_models::protobuf::{ + PhysicalExprNode, PhysicalScalarSubqueryExprNode, physical_expr_node, + }; + + /// Build a `ScalarSubquery` proto node directly, with control over each + /// field, so the decode error paths can be exercised independently. + fn proto_scalar_subquery_node( + data_type: Option, + nullable: bool, + index: u32, + ) -> PhysicalExprNode { + PhysicalExprNode { + expr_id: None, + expr_type: Some(physical_expr_node::ExprType::ScalarSubquery( + PhysicalScalarSubqueryExprNode { + data_type, + nullable, + index, + }, + )), + } + } + + #[test] + fn round_trips_through_proto() { + // A three-slot results container so index 2 is meaningful. + let results = ScalarSubqueryResults::new(3); + let expr = ScalarSubqueryExpr::new( + DataType::Int32, + true, + SubqueryIndex::new(2), + results.clone(), + ); + + // Encode: the expression serializes itself via try_to_proto. + let encoder = StubEncoder::ok(); + let enc_ctx = PhysicalExprEncodeCtx::new(&encoder); + let node = expr + .try_to_proto(&enc_ctx) + .unwrap() + .expect("ScalarSubqueryExpr should encode to Some(node)"); + + assert!(node.expr_id.is_none()); + let sq = match &node.expr_type { + Some(physical_expr_node::ExprType::ScalarSubquery(sq)) => sq, + other => panic!("expected a ScalarSubquery node, got {other:?}"), + }; + assert!(sq.nullable); + assert_eq!(sq.index, 2); + let encoded_type: DataType = sq + .data_type + .as_ref() + .expect("data_type encoded") + .try_into() + .unwrap(); + assert_eq!(encoded_type, DataType::Int32); + + // Decode: reconstruct from the proto node, threading in the shared + // results container the surrounding exec would provide. + let decoder = UnreachableDecoder; + let schema = Schema::empty(); + let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + let decoded = + ScalarSubqueryExpr::try_from_proto(&node, &dec_ctx, &results).unwrap(); + let decoded = decoded + .downcast_ref::() + .expect("decoded expr should be a ScalarSubqueryExpr"); + + // data_type + nullable survive the round-trip (observed via return_field). + let field = decoded.return_field(&Schema::empty()).unwrap(); + assert_eq!(field.data_type(), &DataType::Int32); + assert!(field.is_nullable()); + + // Same shared container + same index → equal to the original. + assert_eq!(decoded, &expr); + } + + #[test] + fn rejects_non_scalar_subquery_node() { + let node = column_node("a"); + let results = ScalarSubqueryResults::new(1); + let decoder = UnreachableDecoder; + let schema = Schema::empty(); + let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = + ScalarSubqueryExpr::try_from_proto(&node, &dec_ctx, &results).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) + if msg.contains("PhysicalExprNode is not a ScalarSubqueryExpr") + )); + } + + #[test] + fn rejects_missing_data_type() { + let node = proto_scalar_subquery_node(None, false, 0); + let results = ScalarSubqueryResults::new(1); + let decoder = UnreachableDecoder; + let schema = Schema::empty(); + let dec_ctx = PhysicalExprDecodeCtx::new(&schema, &decoder); + + let err = + ScalarSubqueryExpr::try_from_proto(&node, &dec_ctx, &results).unwrap_err(); + assert!(matches!( + err, + DataFusionError::Internal(msg) + if msg.contains("ScalarSubqueryExpr is missing required field 'data_type'") + )); + } +} diff --git a/datafusion/physical-expr/src/simplifier/unwrap_cast.rs b/datafusion/physical-expr/src/simplifier/unwrap_cast.rs index 4f4dfb2c20a81..3e67fc8291a4e 100644 --- a/datafusion/physical-expr/src/simplifier/unwrap_cast.rs +++ b/datafusion/physical-expr/src/simplifier/unwrap_cast.rs @@ -36,7 +36,10 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Schema}; use datafusion_common::{Result, ScalarValue, tree_node::Transformed}; use datafusion_expr::Operator; -use datafusion_expr_common::casts::try_cast_literal_to_type; +use datafusion_expr_common::casts::{ + is_date_narrowing_cast, is_timestamp_precision_narrowing_cast, + try_cast_literal_to_type, +}; use crate::PhysicalExpr; use crate::expressions::{BinaryExpr, CastExpr, Literal, TryCastExpr, lit}; @@ -60,13 +63,14 @@ fn try_unwrap_cast_binary( schema: &Schema, ) -> Result>> { // Case 1: cast(left_expr) op literal - if let (Some((inner_expr, _cast_type)), Some(literal)) = ( + if let (Some((inner_expr, cast_type)), Some(literal)) = ( extract_cast_info(binary.left()), binary.right().downcast_ref::(), ) && binary.op().supports_propagation() && let Some(unwrapped) = try_unwrap_cast_comparison( Arc::clone(inner_expr), literal.value(), + cast_type, *binary.op(), schema, )? @@ -75,7 +79,7 @@ fn try_unwrap_cast_binary( } // Case 2: literal op cast(right_expr) - if let (Some(literal), Some((inner_expr, _cast_type))) = ( + if let (Some(literal), Some((inner_expr, cast_type))) = ( binary.left().downcast_ref::(), extract_cast_info(binary.right()), ) { @@ -85,6 +89,7 @@ fn try_unwrap_cast_binary( && let Some(unwrapped) = try_unwrap_cast_comparison( Arc::clone(inner_expr), literal.value(), + cast_type, swapped_op, schema, )? @@ -118,12 +123,19 @@ fn extract_cast_info( fn try_unwrap_cast_comparison( inner_expr: Arc, literal_value: &ScalarValue, + cast_type: &DataType, op: Operator, schema: &Schema, ) -> Result>> { // Get the data type of the inner expression let inner_type = inner_expr.data_type(schema)?; + if is_timestamp_precision_narrowing_cast(&inner_type, cast_type) + || is_date_narrowing_cast(&inner_type, cast_type) + { + return Ok(None); + } + // Try to cast the literal to the inner expression's type if let Some(casted_literal) = try_cast_literal_to_type(literal_value, &inner_type) { let literal_expr = lit(casted_literal); @@ -138,7 +150,7 @@ fn try_unwrap_cast_comparison( mod tests { use super::*; use crate::expressions::col; - use arrow::datatypes::Field; + use arrow::datatypes::{Field, TimeUnit}; use datafusion_common::tree_node::TreeNode; /// Check if an expression is a cast expression @@ -222,6 +234,23 @@ mod tests { assert_eq!(*optimized_binary.op(), Operator::Gt); } + #[test] + fn test_no_unwrap_date64_to_date32_narrowing() { + let schema = Schema::new(vec![Field::new("d64", DataType::Date64, false)]); + + // cast(d64 AS Date32) = Date32(20089) must NOT unwrap: narrowing a Date64 + // column to Date32 truncates milliseconds to the day (many-to-one), so the + // rewritten `d64 = ` would drop sub-day rows. + let column_expr = col("d64", &schema).unwrap(); + let cast_expr = Arc::new(CastExpr::new(column_expr, DataType::Date32, None)); + let literal_expr = lit(ScalarValue::Date32(Some(20089))); + let binary_expr = + Arc::new(BinaryExpr::new(cast_expr, Operator::Eq, literal_expr)); + + let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap(); + assert!(!result.transformed); + } + #[test] fn test_no_unwrap_when_types_unsupported() { let schema = Schema::new(vec![Field::new("f1", DataType::Float32, false)]); @@ -548,6 +577,59 @@ mod tests { assert!(!result.transformed); } + #[test] + fn test_not_unwrap_timestamp_precision_narrowing() { + let schema = Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Nanosecond, None), + false, + )]); + + let column_expr = col("ts", &schema).unwrap(); + let cast_expr = Arc::new(CastExpr::new( + column_expr, + DataType::Timestamp(TimeUnit::Millisecond, None), + None, + )); + let literal_expr = lit(ScalarValue::TimestampMillisecond(Some(1), None)); + let binary_expr = + Arc::new(BinaryExpr::new(cast_expr, Operator::Eq, literal_expr)); + + let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap(); + + assert!(!result.transformed); + } + + #[test] + fn test_unwrap_timestamp_precision_widening() { + let schema = Schema::new(vec![Field::new( + "ts", + DataType::Timestamp(TimeUnit::Millisecond, None), + false, + )]); + + let column_expr = col("ts", &schema).unwrap(); + let cast_expr = Arc::new(CastExpr::new( + column_expr, + DataType::Timestamp(TimeUnit::Nanosecond, None), + None, + )); + let literal_expr = lit(ScalarValue::TimestampNanosecond(Some(1_000_000), None)); + let binary_expr = + Arc::new(BinaryExpr::new(cast_expr, Operator::Eq, literal_expr)); + + let result = unwrap_cast_in_comparison(binary_expr, &schema).unwrap(); + + assert!(result.transformed); + let optimized_binary = result.data.downcast_ref::().unwrap(); + assert!(!is_cast_expr(optimized_binary.left())); + let right_literal = optimized_binary.right().downcast_ref::().unwrap(); + assert_eq!( + right_literal.value(), + &ScalarValue::TimestampMillisecond(Some(1), None) + ); + } + #[test] fn test_complex_nested_expression() { let schema = test_schema(); diff --git a/datafusion/physical-expr/src/window/aggregate.rs b/datafusion/physical-expr/src/window/aggregate.rs index 1ff13d107c036..7cfdcb167f80a 100644 --- a/datafusion/physical-expr/src/window/aggregate.rs +++ b/datafusion/physical-expr/src/window/aggregate.rs @@ -23,7 +23,9 @@ use std::sync::Arc; use crate::aggregate::AggregateFunctionExpr; use crate::window::standard::add_new_ordering_expr_with_partition_by; -use crate::window::window_expr::{AggregateWindowExpr, WindowFn, filter_array}; +use crate::window::window_expr::{ + AggregateWindowExpr, WindowEvalContext, WindowFn, filter_array, +}; use crate::window::{ PartitionBatches, PartitionWindowAggStates, SlidingAggregateWindowExpr, WindowExpr, }; @@ -148,8 +150,9 @@ impl WindowExpr for PlainAggregateWindowExpr { &self, partition_batches: &PartitionBatches, window_agg_state: &mut PartitionWindowAggStates, + eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { - self.aggregate_evaluate_stateful(partition_batches, window_agg_state)?; + self.aggregate_evaluate_stateful(partition_batches, window_agg_state, eval_ctx)?; // Update window frame range for each partition. As we know that // non-sliding aggregations will never call `retract_batch`, this value diff --git a/datafusion/physical-expr/src/window/mod.rs b/datafusion/physical-expr/src/window/mod.rs index b45e35440ac20..79b9a9580af89 100644 --- a/datafusion/physical-expr/src/window/mod.rs +++ b/datafusion/physical-expr/src/window/mod.rs @@ -28,5 +28,6 @@ pub use standard_window_function_expr::StandardWindowFunctionExpr; pub use window_expr::PartitionBatches; pub use window_expr::PartitionKey; pub use window_expr::PartitionWindowAggStates; +pub use window_expr::WindowEvalContext; pub use window_expr::WindowExpr; pub use window_expr::WindowState; diff --git a/datafusion/physical-expr/src/window/sliding_aggregate.rs b/datafusion/physical-expr/src/window/sliding_aggregate.rs index a71df3ec88472..a39334f057dcb 100644 --- a/datafusion/physical-expr/src/window/sliding_aggregate.rs +++ b/datafusion/physical-expr/src/window/sliding_aggregate.rs @@ -22,7 +22,9 @@ use std::ops::Range; use std::sync::Arc; use crate::aggregate::AggregateFunctionExpr; -use crate::window::window_expr::{AggregateWindowExpr, WindowFn, filter_array}; +use crate::window::window_expr::{ + AggregateWindowExpr, WindowEvalContext, WindowFn, filter_array, +}; use crate::window::{ PartitionBatches, PartitionWindowAggStates, PlainAggregateWindowExpr, WindowExpr, }; @@ -102,8 +104,9 @@ impl WindowExpr for SlidingAggregateWindowExpr { &self, partition_batches: &PartitionBatches, window_agg_state: &mut PartitionWindowAggStates, + eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { - self.aggregate_evaluate_stateful(partition_batches, window_agg_state) + self.aggregate_evaluate_stateful(partition_batches, window_agg_state, eval_ctx) } fn partition_by(&self) -> &[Arc] { @@ -207,6 +210,23 @@ impl AggregateWindowExpr for SlidingAggregateWindowExpr { filter_mask: Option<&BooleanArray>, ) -> Result { if cur_range.start == cur_range.end { + // Keep the accumulator synchronized with `last_range`. RANGE frames + // can become empty between two non-empty frames when the ORDER BY + // values contain gaps. + let retract_bound = last_range.end - last_range.start; + if retract_bound > 0 { + let slice_mask = + filter_mask.map(|m| m.slice(last_range.start, retract_bound)); + let retract: Vec = value_slice + .iter() + .map(|v| v.slice(last_range.start, retract_bound)) + .map(|arr| match &slice_mask { + Some(m) => filter_array(&arr, m), + None => Ok(arr), + }) + .collect::>>()?; + accumulator.retract_batch(&retract)? + } self.aggregate .default_value(self.aggregate.field().data_type()) } else { diff --git a/datafusion/physical-expr/src/window/standard.rs b/datafusion/physical-expr/src/window/standard.rs index 46f3cabbadd48..278b66c373f31 100644 --- a/datafusion/physical-expr/src/window/standard.rs +++ b/datafusion/physical-expr/src/window/standard.rs @@ -22,7 +22,7 @@ use std::ops::Range; use std::sync::Arc; use super::{StandardWindowFunctionExpr, WindowExpr}; -use crate::window::window_expr::{WindowFn, get_orderby_values}; +use crate::window::window_expr::{WindowEvalContext, WindowFn, get_orderby_values}; use crate::window::{PartitionBatches, PartitionWindowAggStates, WindowState}; use crate::{EquivalenceProperties, PhysicalExpr}; @@ -128,7 +128,7 @@ impl WindowExpr for StandardWindowExpr { let mut window_frame_ctx = WindowFrameContext::new(Arc::clone(&self.window_frame), sort_options); let mut last_range = Range { start: 0, end: 0 }; - // We iterate on each row to calculate window frame range and and window function result + // We iterate on each row to calculate window frame range and window function result for idx in 0..num_rows { let range = window_frame_ctx.calculate_range( order_bys_ref, @@ -157,6 +157,7 @@ impl WindowExpr for StandardWindowExpr { &self, partition_batches: &PartitionBatches, window_agg_state: &mut PartitionWindowAggStates, + _eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { let field = self.expr.field()?; let out_type = field.data_type(); @@ -175,6 +176,7 @@ impl WindowExpr for StandardWindowExpr { .or_insert(WindowState { state: new_state.clone(), window_fn: WindowFn::Builtin(evaluator), + published: false, }) }; let evaluator = match &mut window_state.window_fn { diff --git a/datafusion/physical-expr/src/window/window_expr.rs b/datafusion/physical-expr/src/window/window_expr.rs index 0f0ec647a50ae..1c52ea4ea6d0e 100644 --- a/datafusion/physical-expr/src/window/window_expr.rs +++ b/datafusion/physical-expr/src/window/window_expr.rs @@ -30,9 +30,11 @@ use arrow::compute::kernels::sort::SortColumn; use arrow::datatypes::FieldRef; use arrow::record_batch::RecordBatch; use datafusion_common::cast::as_boolean_array; +use datafusion_common::hash_utils::RandomState; use datafusion_common::utils::compare_rows; use datafusion_common::{ - Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, internal_err, + Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err, + internal_err, }; use datafusion_expr::window_state::{ PartitionBatchState, WindowAggState, WindowFrameContext, WindowFrameStateGroups, @@ -98,10 +100,14 @@ pub trait WindowExpr: Send + Sync + Debug { /// Evaluate the window function against the batch. This function facilitates /// stateful, bounded-memory implementations. + /// + /// `eval_ctx` carries stream-level (cross-partition) information; see + /// [`WindowEvalContext`]. fn evaluate_stateful( &self, _partition_batches: &PartitionBatches, _window_agg_state: &mut PartitionWindowAggStates, + _eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { internal_err!("evaluate_stateful is not implemented for {}", self.name()) } @@ -225,9 +231,18 @@ pub trait AggregateWindowExpr: WindowExpr { &self, partition_batches: &PartitionBatches, window_agg_state: &mut PartitionWindowAggStates, + eval_ctx: &WindowEvalContext<'_>, ) -> Result<()> { let field = self.field()?; let out_type = field.data_type(); + // Every partition consults the same most recent input row, so its + // ORDER BY values can be evaluated once, outside the per-partition + // loop. + let most_recent_row_order_bys = eval_ctx + .most_recent_row + .map(|batch| self.order_by_columns(batch)) + .transpose()? + .map(get_orderby_values); for (partition_row, partition_batch_state) in partition_batches.iter() { if !window_agg_state.contains_key(partition_row) { let accumulator = self.get_accumulator()?; @@ -236,6 +251,7 @@ pub trait AggregateWindowExpr: WindowExpr { WindowState { state: WindowAggState::new(out_type)?, window_fn: WindowFn::Aggregate(accumulator), + published: false, }, ); }; @@ -248,7 +264,12 @@ pub trait AggregateWindowExpr: WindowExpr { }; let state = &mut window_state.state; let record_batch = &partition_batch_state.record_batch; - let most_recent_row = partition_batch_state.most_recent_row.as_ref(); + + // Skip partitions that cannot produce anything new until they + // either receive rows or reach their end. + if state.is_up_to_date_with(partition_batch_state) { + continue; + } // If there is no window state context, initialize it. let window_frame_ctx = state.window_frame_ctx.get_or_insert_with(|| { @@ -258,7 +279,7 @@ pub trait AggregateWindowExpr: WindowExpr { let out_col = self.get_result_column( accumulator, record_batch, - most_recent_row, + most_recent_row_order_bys.as_deref(), // Start search from the last range &mut state.window_frame_range, window_frame_ctx, @@ -276,7 +297,8 @@ pub trait AggregateWindowExpr: WindowExpr { /// # Arguments /// * `accumulator`: The accumulator to use for the calculation. /// * `record_batch`: batch belonging to the current partition (see [`PartitionBatchState`]). - /// * `most_recent_row`: the batch that contains the most recent row, if available (see [`PartitionBatchState`]). + /// * `most_recent_row_order_bys`: ORDER BY values of the most recent input + /// row, if available (see [`WindowExpr::evaluate_stateful`]). /// * `last_range`: The last range of rows that were processed (see [`WindowAggState`]). /// * `window_frame_ctx`: Details about the window frame (see [`WindowFrameContext`]). /// * `idx`: The index of the current row in the record batch. @@ -286,7 +308,7 @@ pub trait AggregateWindowExpr: WindowExpr { &self, accumulator: &mut Box, record_batch: &RecordBatch, - most_recent_row: Option<&RecordBatch>, + most_recent_row_order_bys: Option<&[ArrayRef]>, last_range: &mut Range, window_frame_ctx: &mut WindowFrameContext, mut idx: usize, @@ -326,10 +348,6 @@ pub trait AggregateWindowExpr: WindowExpr { return value.to_array_of_size(record_batch.num_rows()); } let order_bys = get_orderby_values(self.order_by_columns(record_batch)?); - let most_recent_row_order_bys = most_recent_row - .map(|batch| self.order_by_columns(batch)) - .transpose()? - .map(get_orderby_values); // We iterate on each row to perform a running calculation. let length = values[0].len(); @@ -346,7 +364,7 @@ pub trait AggregateWindowExpr: WindowExpr { && !is_end_bound_safe( window_frame_ctx, &order_bys, - most_recent_row_order_bys.as_deref(), + most_recent_row_order_bys, self.order_by(), idx, )? @@ -604,25 +622,150 @@ pub enum WindowFn { /// PartitionKey would consist of unique `[a,b]` pairs pub type PartitionKey = Vec; +/// Stream-level context passed to [`WindowExpr::evaluate_stateful`]. +/// +/// This carries information that spans all partitions of the input, as +/// opposed to the per-partition state in [`PartitionBatches`] and +/// [`PartitionWindowAggStates`]. It is `non_exhaustive` so that fields can +/// be added without breaking implementors; construct it with +/// [`Default::default`] and the `with_*` builder methods. +#[derive(Debug, Clone, Copy, Default)] +#[non_exhaustive] +pub struct WindowEvalContext<'a> { + /// A single-row batch containing the most recent input row, whichever + /// partition that row belongs to. It is `Some` only when the input is + /// ordered by the first ORDER BY column across partitions (`Linear` + /// mode), in which case no future input row -- in any partition -- can + /// precede it in that column; implementations can use this bound to + /// decide whether pending window frames can be finalized before their + /// partition receives more data. + pub most_recent_row: Option<&'a RecordBatch>, +} + +impl<'a> WindowEvalContext<'a> { + /// Sets the most recent input row (see [`Self::most_recent_row`]). + pub fn with_most_recent_row(mut self, batch: Option<&'a RecordBatch>) -> Self { + self.most_recent_row = batch; + self + } +} + #[derive(Debug)] pub struct WindowState { pub state: WindowAggState, pub window_fn: WindowFn, + /// True once [`Self::aggregate_state`] has been called on this entry. + /// Guards against a second destructive [`Accumulator::state`] read: the + /// method itself errors on second call, and the observer loop in + /// `BoundedWindowAggStream::publish_finalized_states` uses this as an + /// early-skip so it doesn't attempt one. Independent of `state.is_end`, + /// which is a group-closed signal that the pruning path also reads. + pub published: bool, +} + +impl WindowState { + /// [`Accumulator::state`] if this window function is an aggregate, `None` + /// otherwise (built-in functions like `row_number`, `rank`, `lead`/`lag` + /// have no serializable accumulator state). + /// + /// [`Accumulator::state`] takes `&mut self` and its trait doc calls out + /// that "this function should not be called twice, otherwise it will + /// result in potentially non-deterministic behavior." Several built-in + /// impls (`median`, `percentile_cont`, `string_agg`, + /// `min_max_bytes`/`min_max_struct`) `std::mem::take` their internal + /// buffers on call — a second call returns *empty* state, not the same + /// state, so a downstream prefix-merge would silently lose every value + /// the accumulator had ingested. + /// + /// Enforced at this layer: on first call we set [`Self::published`] and + /// return the state; any later call errors rather than performing a + /// destructive re-read. + pub fn aggregate_state(&mut self) -> Result>> { + if self.published { + return exec_err!( + "WindowState::aggregate_state called more than once; \ + Accumulator::state is a destructive read for several \ + built-in aggregates and a second call would silently lose data" + ); + } + let state = match &mut self.window_fn { + WindowFn::Aggregate(accumulator) => Some(accumulator.state()?), + WindowFn::Builtin(_) => None, + }; + self.published = true; + Ok(state) + } } -pub type PartitionWindowAggStates = IndexMap; + +pub type PartitionWindowAggStates = IndexMap; /// The IndexMap (i.e. an ordered HashMap) where record batches are separated for each partition. -pub type PartitionBatches = IndexMap; +pub type PartitionBatches = IndexMap; #[cfg(test)] mod tests { use std::sync::Arc; - use crate::window::window_expr::is_row_ahead; + use crate::window::window_expr::{WindowFn, WindowState, is_row_ahead}; use arrow::array::{ArrayRef, Float64Array}; use arrow::compute::SortOptions; - use datafusion_common::Result; + use arrow::datatypes::DataType; + use datafusion_common::{Result, ScalarValue}; + use datafusion_expr::{Accumulator, window_state::WindowAggState}; + + /// Minimal [`Accumulator`] whose `state()` records how many times it was + /// called by returning the count as its single state element. Any second + /// call would surface (were it allowed to happen) as `[UInt64(2)]` + /// instead of `[UInt64(1)]`. + #[derive(Debug)] + struct CallCountingAccumulator { + calls: usize, + } + + impl Accumulator for CallCountingAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Null) + } + fn size(&self) -> usize { + size_of::() + } + fn state(&mut self) -> Result> { + self.calls += 1; + Ok(vec![ScalarValue::UInt64(Some(self.calls as u64))]) + } + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } + } + + #[test] + fn aggregate_state_errors_on_second_call() -> Result<()> { + // `Accumulator::state()` is a destructive read for several built-in + // aggregates (median, percentile_cont, string_agg, min_max_bytes/ + // min_max_struct all `mem::take` their internal buffers). Its trait + // doc says "should not be called twice"; `WindowState::aggregate_state` + // enforces that at this layer by returning an error rather than + // performing the second read. + let acc: Box = Box::new(CallCountingAccumulator { calls: 0 }); + let mut ws = WindowState { + state: WindowAggState::new(&DataType::UInt64)?, + window_fn: WindowFn::Aggregate(acc), + published: false, + }; + let first = ws.aggregate_state()?; + assert_eq!(first, Some(vec![ScalarValue::UInt64(Some(1))])); + assert!(ws.published, "published must flip on successful publish"); + let err = ws.aggregate_state().unwrap_err().to_string(); + assert!( + err.contains("called more than once"), + "expected second-call error, got: {err}" + ); + Ok(()) + } #[test] fn test_is_row_ahead() -> Result<()> { diff --git a/datafusion/physical-optimizer/Cargo.toml b/datafusion/physical-optimizer/Cargo.toml index 38c8a7c37211f..cb03303ac3c3f 100644 --- a/datafusion/physical-optimizer/Cargo.toml +++ b/datafusion/physical-optimizer/Cargo.toml @@ -50,6 +50,7 @@ datafusion-physical-expr = { workspace = true } datafusion-physical-expr-common = { workspace = true } datafusion-physical-plan = { workspace = true } datafusion-pruning = { workspace = true } +datafusion-session = { workspace = true } itertools = { workspace = true } recursive = { workspace = true, optional = true } diff --git a/datafusion/physical-optimizer/src/aggregate_statistics.rs b/datafusion/physical-optimizer/src/aggregate_statistics.rs index d0be53d59b3cf..43b1abb4b68a9 100644 --- a/datafusion/physical-optimizer/src/aggregate_statistics.rs +++ b/datafusion/physical-optimizer/src/aggregate_statistics.rs @@ -25,7 +25,10 @@ use datafusion_physical_plan::aggregates::{ }; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; -use datafusion_physical_plan::udaf::{AggregateFunctionExpr, StatisticsArgs}; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; +use datafusion_physical_plan::udaf::{ + AggregateFunctionExpr, StatisticsArgs as PlanStatisticsArgs, +}; use datafusion_physical_plan::{ExecutionPlan, expressions}; use std::sync::Arc; @@ -55,12 +58,13 @@ impl PhysicalOptimizerRule for AggregateStatistics { let partial_agg_exec = partial_agg_exec .downcast_ref::() .expect("take_optimizable() ensures that this is a AggregateExec"); - let stats = partial_agg_exec.input().partition_statistics(None)?; + let stats = StatisticsContext::new() + .compute(partial_agg_exec.input().as_ref(), &StatisticsArgs::new())?; let mut projections = vec![]; for expr in partial_agg_exec.aggr_expr() { let field = expr.field(); let args = expr.expressions(); - let statistics_args = StatisticsArgs { + let statistics_args = PlanStatisticsArgs { statistics: &stats, return_type: field.data_type(), is_distinct: expr.is_distinct(), @@ -148,7 +152,7 @@ fn take_optimizable(plan: &Arc) -> Option Option<(ScalarValue, String)> { let value = agg_expr.fun().value_from_stats(statistics_args); diff --git a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs index 74e938e75ed64..297a92c45a16d 100644 --- a/datafusion/physical-optimizer/src/combine_partial_final_agg.rs +++ b/datafusion/physical-optimizer/src/combine_partial_final_agg.rs @@ -35,7 +35,8 @@ use datafusion_physical_expr::{PhysicalExpr, physical_exprs_equal}; /// CombinePartialFinalAggregate optimizer rule combines the adjacent Partial and Final AggregateExecs /// into a Single AggregateExec if their grouping exprs and aggregate exprs equal. /// -/// This rule should be applied after the EnforceDistribution and EnforceSorting rules +/// This rule should be applied after the `EnsureRequirements` rule (which +/// handles both distribution and sorting enforcement). #[derive(Default, Debug)] pub struct CombinePartialFinalAggregate {} diff --git a/datafusion/physical-optimizer/src/ensure_coop.rs b/datafusion/physical-optimizer/src/ensure_coop.rs index 102e21a4853a4..93862df3b4236 100644 --- a/datafusion/physical-optimizer/src/ensure_coop.rs +++ b/datafusion/physical-optimizer/src/ensure_coop.rs @@ -130,7 +130,10 @@ impl PhysicalOptimizerRule for EnsureCooperative { #[cfg(test)] mod tests { use super::*; - use datafusion_physical_plan::{displayable, test::scan_partitioned}; + use datafusion_physical_plan::{ + ChildrenPropertiesMode, ReplaceChildrenOptions, displayable, + test::scan_partitioned, + }; use insta::assert_snapshot; #[tokio::test] @@ -328,9 +331,10 @@ mod tests { fn children(&self) -> Vec<&Arc> { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(DummyExec::new( &self.name, @@ -339,6 +343,15 @@ mod tests { self.evaluation_type, ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn execute( &self, _: usize, @@ -349,7 +362,7 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-optimizer/src/enforce_distribution.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs similarity index 72% rename from datafusion/physical-optimizer/src/enforce_distribution.rs rename to datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs index c522867c05196..07bc98b2db798 100644 --- a/datafusion/physical-optimizer/src/enforce_distribution.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_distribution.rs @@ -15,17 +15,23 @@ // specific language governing permissions and limitations // under the License. -//! EnforceDistribution optimizer rule inspects the physical plan with respect -//! to distribution requirements and adds [`RepartitionExec`]s to satisfy them -//! when necessary. If increasing parallelism is beneficial (and also desirable -//! according to the configuration), this rule increases partition counts in -//! the physical plan. +//! Distribution enforcement helpers. The standalone `EnforceDistribution` +//! rule that previously lived here has been retired in favour of +//! `EnsureRequirements` (which composes distribution and sorting +//! enforcement into a single idempotent pass). The helpers in this +//! module — `adjust_input_keys_ordering`, `reorder_join_keys_to_inputs`, +//! `DistributionContext`, `ensure_distribution`, … — are used directly +//! by `EnsureRequirements`. +//! +//! These helpers inspect the physical plan with respect to distribution +//! requirements and add [`RepartitionExec`]s to satisfy them when necessary. +//! If increasing parallelism is beneficial (and also desirable according to +//! configuration), they increase partition counts in the physical plan. use std::any::Any; use std::fmt::Debug; use std::sync::Arc; -use crate::optimizer::PhysicalOptimizerRule; use crate::output_requirements::OutputRequirementExec; use crate::utils::{ add_sort_above_with_check, is_coalesce_partitions, is_repartition, @@ -36,201 +42,45 @@ use arrow::compute::SortOptions; use datafusion_common::config::ConfigOptions; use datafusion_common::error::Result; use datafusion_common::stats::Precision; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_common::tree_node::Transformed; use datafusion_expr::logical_plan::{Aggregate, JoinType}; use datafusion_physical_expr::expressions::{Column, NoOp}; use datafusion_physical_expr::utils::map_columns_before_projection; use datafusion_physical_expr::{ - EquivalenceProperties, PhysicalExpr, PhysicalExprRef, physical_exprs_equal, + EquivalenceProperties, OrderingRequirements, PhysicalExpr, PhysicalExprRef, + physical_exprs_equal, }; use datafusion_physical_plan::ExecutionPlanProperties; use datafusion_physical_plan::aggregates::{ AggregateExec, AggregateMode, PhysicalGroupBy, }; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion_physical_plan::execution_plan::EmissionType; +use datafusion_physical_plan::execution_plan::{ + EmissionType, replace_children_if_necessary, +}; use datafusion_physical_plan::joins::{ CrossJoinExec, HashJoinExec, PartitionMode, SortMergeJoinExec, }; use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::union::{InterleaveExec, UnionExec, can_interleave}; use datafusion_physical_plan::windows::WindowAggExec; use datafusion_physical_plan::windows::{BoundedWindowAggExec, get_best_fitting_window}; -use datafusion_physical_plan::{Distribution, ExecutionPlan, Partitioning}; +use datafusion_physical_plan::{ + ChildSatisfactionOptions, Distribution, ExecutionPlan, InputDistributionRequirements, + Partitioning, +}; use itertools::izip; -/// The `EnforceDistribution` rule ensures that distribution requirements are -/// met. In doing so, this rule will increase the parallelism in the plan by -/// introducing repartitioning operators to the physical plan. -/// -/// For example, given an input such as: -/// -/// -/// ```text -/// ┌─────────────────────────────────┐ -/// │ │ -/// │ ExecutionPlan │ -/// │ │ -/// └─────────────────────────────────┘ -/// ▲ ▲ -/// │ │ -/// ┌─────┘ └─────┐ -/// │ │ -/// │ │ -/// │ │ -/// ┌───────────┐ ┌───────────┐ -/// │ │ │ │ -/// │ batch A1 │ │ batch B1 │ -/// │ │ │ │ -/// ├───────────┤ ├───────────┤ -/// │ │ │ │ -/// │ batch A2 │ │ batch B2 │ -/// │ │ │ │ -/// ├───────────┤ ├───────────┤ -/// │ │ │ │ -/// │ batch A3 │ │ batch B3 │ -/// │ │ │ │ -/// └───────────┘ └───────────┘ -/// -/// Input Input -/// A B -/// ``` -/// -/// This rule will attempt to add a `RepartitionExec` to increase parallelism -/// (to 3, in this case) and create the following arrangement: -/// -/// ```text -/// ┌─────────────────────────────────┐ -/// │ │ -/// │ ExecutionPlan │ -/// │ │ -/// └─────────────────────────────────┘ -/// ▲ ▲ ▲ Input now has 3 -/// │ │ │ partitions -/// ┌───────┘ │ └───────┐ -/// │ │ │ -/// │ │ │ -/// ┌───────────┐ ┌───────────┐ ┌───────────┐ -/// │ │ │ │ │ │ -/// │ batch A1 │ │ batch A3 │ │ batch B3 │ -/// │ │ │ │ │ │ -/// ├───────────┤ ├───────────┤ ├───────────┤ -/// │ │ │ │ │ │ -/// │ batch B2 │ │ batch B1 │ │ batch A2 │ -/// │ │ │ │ │ │ -/// └───────────┘ └───────────┘ └───────────┘ -/// ▲ ▲ ▲ -/// │ │ │ -/// └─────────┐ │ ┌──────────┘ -/// │ │ │ -/// │ │ │ -/// ┌─────────────────────────────────┐ batches are -/// │ RepartitionExec(3) │ repartitioned -/// │ RoundRobin │ -/// │ │ -/// └─────────────────────────────────┘ -/// ▲ ▲ -/// │ │ -/// ┌─────┘ └─────┐ -/// │ │ -/// │ │ -/// │ │ -/// ┌───────────┐ ┌───────────┐ -/// │ │ │ │ -/// │ batch A1 │ │ batch B1 │ -/// │ │ │ │ -/// ├───────────┤ ├───────────┤ -/// │ │ │ │ -/// │ batch A2 │ │ batch B2 │ -/// │ │ │ │ -/// ├───────────┤ ├───────────┤ -/// │ │ │ │ -/// │ batch A3 │ │ batch B3 │ -/// │ │ │ │ -/// └───────────┘ └───────────┘ -/// -/// -/// Input Input -/// A B -/// ``` -/// -/// The `EnforceDistribution` rule -/// - is idempotent; i.e. it can be applied multiple times, each time producing -/// the same result. -/// - always produces a valid plan in terms of distribution requirements. Its -/// input plan can be valid or invalid with respect to distribution requirements, -/// but the output plan will always be valid. -/// - produces a valid plan in terms of ordering requirements, *if* its input is -/// a valid plan in terms of ordering requirements. If the input plan is invalid, -/// this rule does not attempt to fix it as doing so is the responsibility of the -/// `EnforceSorting` rule. -/// -/// Note that distribution requirements are met in the strictest way. This may -/// result in more than strictly necessary [`RepartitionExec`]s in the plan, but -/// meeting the requirements in the strictest way may help avoid possible data -/// skew in joins. -/// -/// For example for a hash join with keys (a, b, c), the required Distribution(a, b, c) -/// can be satisfied by several alternative partitioning ways: (a, b, c), (a, b), -/// (a, c), (b, c), (a), (b), (c) and ( ). -/// -/// This rule only chooses the exact match and satisfies the Distribution(a, b, c) -/// by a HashPartition(a, b, c). -#[derive(Default, Debug)] -pub struct EnforceDistribution {} - -impl EnforceDistribution { - #[expect(missing_docs)] - pub fn new() -> Self { - Self {} - } -} - -impl PhysicalOptimizerRule for EnforceDistribution { - fn optimize( - &self, - plan: Arc, - config: &ConfigOptions, - ) -> Result> { - let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering; - - let adjusted = if top_down_join_key_reordering { - // Run a top-down process to adjust input key ordering recursively - let plan_requirements = PlanWithKeyRequirements::new_default(plan); - let adjusted = plan_requirements - .transform_down(adjust_input_keys_ordering) - .data()?; - adjusted.plan - } else { - // Run a bottom-up process - plan.transform_up(|plan| { - Ok(Transformed::yes(reorder_join_keys_to_inputs(plan)?)) - }) - .data()? - }; - - let distribution_context = DistributionContext::new_default(adjusted); - // Distribution enforcement needs to be applied bottom-up. - let distribution_context = distribution_context - .transform_up(|distribution_context| { - ensure_distribution(distribution_context, config) - }) - .data()?; - Ok(distribution_context.plan) - } - - fn name(&self) -> &str { - "EnforceDistribution" - } - - fn schema_check(&self) -> bool { - true - } -} +// The `EnforceDistribution` rule was retired in favour of `EnsureRequirements`, +// which composes distribution and sorting enforcement into a single idempotent +// pass. The helper functions below (`adjust_input_keys_ordering`, +// `reorder_join_keys_to_inputs`, `DistributionContext`, `ensure_distribution`, +// etc.) remain — `EnsureRequirements` calls into them directly. #[derive(Debug, Clone)] struct JoinKeyPairs { @@ -783,12 +633,9 @@ fn expected_expr_positions( let mut current = current.to_vec(); for expr in expected.iter() { // Find the position of the expected expr in the current expressions - if let Some(expected_position) = current.iter().position(|e| e.eq(expr)) { - current[expected_position] = Arc::new(NoOp::new()); - indexes.push(expected_position); - } else { - return None; - } + let expected_position = current.iter().position(|e| e.eq(expr))?; + current[expected_position] = Arc::new(NoOp::new()); + indexes.push(expected_position); } Some(indexes) } @@ -853,72 +700,34 @@ fn add_roundrobin_on_top( } } -/// Adds a hash repartition operator: -/// - to increase parallelism, and/or -/// - to satisfy requirements of the subsequent operators. -/// -/// Repartition(Hash) is added on top of operator `input`. -/// -/// # Arguments -/// -/// * `input`: Current node. -/// * `hash_exprs`: Stores Physical Exprs that are used during hashing. -/// * `n_target`: desired target partition number, if partition number of the -/// current executor is less than this value. Partition number will be increased. -/// * `allow_subset_satisfy_partitioning`: Whether to allow subset partitioning logic in satisfaction checks. -/// Set to `false` for partitioned hash joins to ensure exact hash matching. -/// -/// # Returns -/// -/// A [`Result`] object that contains new execution plan where the desired -/// distribution is satisfied by adding a Hash repartition. -fn add_hash_on_top( - input: DistributionContext, - hash_exprs: Vec>, - n_target: usize, +// Partial aggregates require unspecified input distribution, but their output +// may already satisfy the final aggregate's key distribution because partial +// aggregation preserves/projects input partitioning. Keep that reusable output +// partitioning intact when preserve_file_partitions would otherwise insert +// RoundRobin below the partial aggregate. +fn partial_aggregate_output_satisfies_final_partitioning( + plan: &Arc, allow_subset_satisfy_partitioning: bool, -) -> Result { - // Early return if hash repartition is unnecessary - // `RepartitionExec: partitioning=Hash([...], 1), input_partitions=1` is unnecessary. - if n_target == 1 && input.plan.output_partitioning().partition_count() == 1 { - return Ok(input); - } - - let dist = Distribution::HashPartitioned(hash_exprs); - let satisfaction = input.plan.output_partitioning().satisfaction( - &dist, - input.plan.equivalence_properties(), - allow_subset_satisfy_partitioning, - ); - - // Add hash repartitioning when: - // - When subset satisfaction is enabled (current >= threshold): only repartition if not satisfied - // - When below threshold (current < threshold): repartition if expressions don't match OR to increase parallelism - let needs_repartition = if allow_subset_satisfy_partitioning { - !satisfaction.is_satisfied() - } else { - !satisfaction.is_satisfied() - || n_target > input.plan.output_partitioning().partition_count() +) -> bool { + let Some(aggregate) = plan.downcast_ref::() else { + return false; }; - - if needs_repartition { - // When there is an existing ordering, we preserve ordering during - // repartition. This will be rolled back in the future if any of the - // following conditions is true: - // - Preserving ordering is not helpful in terms of satisfying ordering - // requirements. - // - Usage of order preserving variants is not desirable (per the flag - // `config.optimizer.prefer_existing_sort`). - let partitioning = dist.create_partitioning(n_target); - let repartition = - RepartitionExec::try_new(Arc::clone(&input.plan), partitioning)? - .with_preserve_order(); - let plan = Arc::new(repartition) as _; - - return Ok(DistributionContext::new(plan, true, vec![input])); + if aggregate.mode() != &AggregateMode::Partial + || aggregate.group_expr().is_empty() + || aggregate.group_expr().has_grouping_set() + { + return false; } - Ok(input) + let key_distribution = Distribution::KeyPartitioned(aggregate.output_group_expr()); + + plan.output_partitioning() + .satisfaction( + &key_distribution, + plan.equivalence_properties(), + allow_subset_satisfy_partitioning, + ) + .is_satisfied() } /// Adds a [`SortPreservingMergeExec`] or a [`CoalescePartitionsExec`] operator @@ -953,8 +762,10 @@ fn preserving_order_enables_streaming( return Ok(false); } // Build parent with the ordered child - let with_ordered = - Arc::clone(parent).with_new_children(vec![Arc::clone(ordered_child)])?; + let with_ordered = replace_children_if_necessary( + Arc::clone(parent), + vec![Arc::clone(ordered_child)], + )?; if with_ordered.pipeline_behavior() == EmissionType::Final { // Parent is blocking even with ordering — no benefit return Ok(false); @@ -962,7 +773,8 @@ fn preserving_order_enables_streaming( // Build parent with an unordered child via CoalescePartitionsExec. let unordered_child: Arc = Arc::new(CoalescePartitionsExec::new(Arc::clone(ordered_child))); - let without_ordered = Arc::clone(parent).with_new_children(vec![unordered_child])?; + let without_ordered = + replace_children_if_necessary(Arc::clone(parent), vec![unordered_child])?; Ok(without_ordered.pipeline_behavior() == EmissionType::Final) } @@ -970,7 +782,10 @@ fn preserving_order_enables_streaming( /// /// Updated node with an execution plan, where the desired single distribution /// requirement is satisfied. -fn add_merge_on_top(input: DistributionContext) -> DistributionContext { +fn add_merge_on_top( + input: DistributionContext, + fetch: Option, +) -> DistributionContext { // Apply only when the partition count is larger than one. if input.plan.output_partitioning().partition_count() > 1 { // When there is an existing ordering, we preserve ordering @@ -979,14 +794,20 @@ fn add_merge_on_top(input: DistributionContext) -> DistributionContext { // - Preserving ordering is not helpful in terms of satisfying ordering requirements // - Usage of order preserving variants is not desirable // (determined by flag `config.optimizer.prefer_existing_sort`) - let new_plan = if let Some(req) = input.plan.output_ordering() { - Arc::new(SortPreservingMergeExec::new( - req.clone(), - Arc::clone(&input.plan), - )) as _ + let new_plan: Arc = if let Some(req) = + input.plan.output_ordering() + { + let mut spm = + SortPreservingMergeExec::new(req.clone(), Arc::clone(&input.plan)); + if let Some(f) = fetch { + spm = spm.with_fetch(Some(f)); + } + Arc::new(spm) } else { // If there is no input order, we can simply coalesce partitions: - Arc::new(CoalescePartitionsExec::new(Arc::clone(&input.plan))) as _ + Arc::new( + CoalescePartitionsExec::new(Arc::clone(&input.plan)).with_fetch(fetch), + ) }; DistributionContext::new(new_plan, true, vec![input]) @@ -1012,20 +833,41 @@ fn add_merge_on_top(input: DistributionContext) -> DistributionContext { /// ```text /// "DataSourceExec: file_groups={2 groups: \[\[x], \[y]]}, projection=\[a, b, c, d, e], output_ordering=\[a@0 ASC], file_type=parquet", /// ``` +/// Returned by [`remove_dist_changing_operators`] to carry the fetch value +/// that may have been on a removed `SortPreservingMergeExec` or `CoalescePartitionsExec`. +struct RemovedDistOps { + context: DistributionContext, + /// The fetch value from the removed SPM/Coalesce, if any. + /// Must be re-applied when distribution operators are re-inserted. + removed_fetch: Option, +} + fn remove_dist_changing_operators( mut distribution_context: DistributionContext, -) -> Result { +) -> Result { + let mut removed_fetch = None; while is_repartition(&distribution_context.plan) || is_coalesce_partitions(&distribution_context.plan) || is_sort_preserving_merge(&distribution_context.plan) { + // Preserve fetch from SPM or CoalescePartitions before removing (#14150). + if let Some(fetch) = distribution_context.plan.fetch() { + removed_fetch = Some( + removed_fetch + .map(|existing: usize| existing.min(fetch)) + .unwrap_or(fetch), + ); + } // All of above operators have a single child. First child is only child. // Remove any distribution changing operators at the beginning: distribution_context = distribution_context.children.swap_remove(0); // Note that they will be re-inserted later on if necessary or helpful. } - Ok(distribution_context) + Ok(RemovedDistOps { + context: distribution_context, + removed_fetch, + }) } /// Updates the [`DistributionContext`] if preserving ordering while changing partitioning is not helpful or desirable. @@ -1094,6 +936,14 @@ struct RepartitionRequirementStatus { hash_necessary: bool, } +/// Per-child state while enforcing a parent's distribution requirements. +struct DistributionChildState { + context: DistributionContext, + required_input_ordering: Option, + maintains_input_order: bool, + requirement: Distribution, +} + /// Calculates the `RepartitionRequirementStatus` for each children to generate /// consistent and sensible (in terms of performance) distribution requirements. /// As an example, a hash join's left (build) child might produce @@ -1125,6 +975,10 @@ struct RepartitionRequirementStatus { /// hash_necessary: true /// } /// ``` +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] fn get_repartition_requirement_status( plan: &Arc, batch_size: usize, @@ -1133,20 +987,25 @@ fn get_repartition_requirement_status( let mut needs_alignment = false; let children = plan.children(); let rr_beneficial = plan.benefits_from_input_partitioning(); - let requirements = plan.required_input_distribution(); + let requirements = plan.input_distribution_requirements().into_per_child(); let mut repartition_status_flags = vec![]; for (child, requirement, roundrobin_beneficial) in izip!(children.into_iter(), requirements, rr_beneficial) { // Decide whether adding a round robin is beneficial depending on // the statistical information we have on the number of rows: - let roundrobin_beneficial_stats = match child.partition_statistics(None)?.num_rows + let roundrobin_beneficial_stats = match StatisticsContext::new() + .compute(child.as_ref(), &StatisticsArgs::new())? + .num_rows { Precision::Exact(n_rows) => n_rows > batch_size, Precision::Inexact(n_rows) => !should_use_estimates || (n_rows > batch_size), Precision::Absent => true, }; - let is_hash = matches!(requirement, Distribution::HashPartitioned(_)); + let is_hash = matches!( + requirement, + Distribution::HashPartitioned(_) | Distribution::KeyPartitioned(_) + ); // Hash re-partitioning is necessary when the input has more than one // partitions: let multi_partitions = child.output_partitioning().partition_count() > 1; @@ -1180,6 +1039,90 @@ fn get_repartition_requirement_status( .collect()) } +/// Enforce cross-child distribution relationships after each child has already +/// satisfied its own distribution requirement. +/// +/// See [`InputDistributionRequirements`] for the distinction between +/// independent per-child requirements and co-partitioned child relationships. +/// +/// Currently, unsatisfied co-partitioning is repaired by hash repartitioning +/// key-partitioned children and other relationship kinds are rejected. +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] +fn enforce_distribution_relationships( + plan_name: &str, + input_distributions: &InputDistributionRequirements, + children: &mut [DistributionChildState], + target_partitions: usize, +) -> Result<()> { + let mut repartitioned_for_relationship = vec![false; children.len()]; + + loop { + let child_plan_refs = children + .iter() + .map(|child| child.context.plan.as_ref()) + .collect::>(); + let unsatisfied_children = input_distributions + .unsatisfied_co_partitioned_children(plan_name, &child_plan_refs)?; + + if unsatisfied_children.is_empty() { + return Ok(()); + } + + let mut changed = false; + for child_idx in unsatisfied_children { + if repartitioned_for_relationship[child_idx] { + continue; + } + + let (Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs)) = &children[child_idx].requirement + else { + continue; + }; + + let already_target_hash = matches!( + children[child_idx].context.plan.output_partitioning(), + Partitioning::Hash(_, partition_count) if *partition_count == target_partitions + ) && input_distributions + .child_satisfaction( + child_idx, + children[child_idx].context.plan.as_ref(), + ChildSatisfactionOptions::new(), + )? + .is_satisfied(); + + if already_target_hash { + continue; + } + + let partitioning = Distribution::KeyPartitioned(exprs.to_vec()) + .create_partitioning(target_partitions); + let repartition = RepartitionExec::try_new( + Arc::clone(&children[child_idx].context.plan), + partitioning, + )? + .with_preserve_order(); + let plan = Arc::new(repartition) as _; + let original_child = std::mem::replace( + &mut children[child_idx].context, + DistributionContext::new(plan, true, vec![]), + ); + children[child_idx].context.children = vec![original_child]; + repartitioned_for_relationship[child_idx] = true; + changed = true; + } + + if !changed { + return datafusion_common::internal_err!( + "{plan_name} has distribution relationships that could not be enforced" + ); + } + } +} + /// This function checks whether we need to add additional data exchange /// operators to satisfy distribution requirements. Since this function /// takes care of such requirements, we should avoid manually adding data @@ -1188,6 +1131,10 @@ fn get_repartition_requirement_status( /// This function is intended to be used in a bottom up traversal, as it /// can first repartition (or newly partition) at the datasources -- these /// source partitions may be later repartitioned with additional data exchange operators. +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] pub fn ensure_distribution( dist_context: DistributionContext, config: &ConfigOptions, @@ -1202,7 +1149,7 @@ pub fn ensure_distribution( // When `false`, round robin repartition will not be added to increase parallelism let enable_round_robin = config.optimizer.enable_round_robin_repartition; let repartition_file_scans = config.optimizer.repartition_file_scans; - let batch_size = config.execution.batch_size; + let batch_size = config.execution.batch_size.get(); let should_use_estimates = config .execution .use_row_number_estimates_to_optimize_partitioning; @@ -1219,11 +1166,16 @@ pub fn ensure_distribution( let order_preserving_variants_desirable = unbounded_and_pipeline_friendly || config.optimizer.prefer_existing_sort; - // Remove unnecessary repartition from the physical plan if any - let DistributionContext { - mut plan, - data, - children, + // Remove unnecessary repartition from the physical plan if any. + // Preserve fetch from removed SPM/Coalesce (#14150). + let RemovedDistOps { + context: + DistributionContext { + mut plan, + data, + children, + }, + removed_fetch, } = remove_dist_changing_operators(dist_context)?; if let Some(exec) = plan.downcast_ref::() { @@ -1231,6 +1183,7 @@ pub fn ensure_distribution( exec.window_expr(), exec.input(), &exec.partition_keys(), + None, )? { plan = updated_window; } @@ -1239,6 +1192,7 @@ pub fn ensure_distribution( exec.window_expr(), exec.input(), &exec.partition_keys(), + exec.state_observer().cloned(), )? { plan = updated_window; @@ -1278,6 +1232,7 @@ pub fn ensure_distribution( .is_some_and(|join| join.mode == PartitionMode::Partitioned) || plan.is::(); + let input_distributions = plan.input_distribution_requirements(); let repartition_status_flags = get_repartition_requirement_status(&plan, batch_size, should_use_estimates)?; // This loop iterates over all the children to: @@ -1285,7 +1240,8 @@ pub fn ensure_distribution( // - Satisfy the distribution requirements of every child, if it is not // already satisfied. // We store the updated children in `new_children`. - let children = izip!( + let mut children = izip!( + 0..children.len(), children.into_iter(), plan.required_input_ordering(), plan.maintains_input_order(), @@ -1293,6 +1249,7 @@ pub fn ensure_distribution( ) .map( |( + child_idx, mut child, required_input_ordering, maintains, @@ -1303,27 +1260,25 @@ pub fn ensure_distribution( hash_necessary, }, )| { - let increases_partition_count = - child.plan.output_partitioning().partition_count() < target_partitions; - - let add_roundrobin = enable_round_robin - // Operator benefits from partitioning (e.g. filter): - && roundrobin_beneficial - && roundrobin_beneficial_stats - // Unless partitioning increases the partition count, it is not beneficial: - && increases_partition_count; - // Allow subset satisfaction when: // 1. Current partition count >= threshold // 2. Not a partitioned join since must use exact hash matching for joins // 3. Not a grouping set aggregate (requires exact hash including __grouping_id) + // + // Partitioned joins still require exact satisfaction. If that + // exact check already passes, preserve_file_partitions can skip + // repartitioning whose only purpose is increasing partition count. let current_partitions = child.plan.output_partitioning().partition_count(); + let preserve_file_partition_threshold_met = + config.optimizer.preserve_file_partitions > 0 + && current_partitions >= config.optimizer.preserve_file_partitions; // Check if the hash partitioning requirement includes __grouping_id column. // Grouping set aggregates (ROLLUP, CUBE, GROUPING SETS) require exact hash // partitioning on all group columns including __grouping_id to ensure partial // aggregates from different partitions are correctly combined. - let requires_grouping_id = matches!(&requirement, Distribution::HashPartitioned(exprs) + let requires_grouping_id = matches!(&requirement, + Distribution::HashPartitioned(exprs) | Distribution::KeyPartitioned(exprs) if exprs.iter().any(|expr| { (expr.as_ref() as &dyn Any) .downcast_ref::() @@ -1337,11 +1292,28 @@ pub fn ensure_distribution( // partitioning to the optimizer. Respect it when the only // reason to repartition would be to increase partition count // beyond the preserved file-group count. - || (config.optimizer.preserve_file_partitions > 0 + || (preserve_file_partition_threshold_met && current_partitions < target_partitions)) && !is_partitioned_join && !requires_grouping_id; + let increases_partition_count = current_partitions < target_partitions; + + let preserve_partial_aggregate_partitioning = + preserve_file_partition_threshold_met + && partial_aggregate_output_satisfies_final_partitioning( + &plan, + allow_subset_satisfy_partitioning, + ); + + let add_roundrobin = enable_round_robin + // Operator benefits from partitioning (e.g. filter): + && roundrobin_beneficial + && roundrobin_beneficial_stats + // Unless partitioning increases the partition count, it is not beneficial: + && increases_partition_count + && !preserve_partial_aggregate_partitioning; + // When `repartition_file_scans` is set, attempt to increase // parallelism at the source. // @@ -1359,18 +1331,58 @@ pub fn ensure_distribution( // Satisfy the distribution requirement if it is unmet. match &requirement { Distribution::SinglePartition => { - child = add_merge_on_top(child); + child = add_merge_on_top(child, removed_fetch); } - Distribution::HashPartitioned(exprs) => { + Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs) => { + let child_partitions = + child.plan.output_partitioning().partition_count(); + let partitioning_satisfied = input_distributions + .child_satisfaction( + child_idx, + child.plan.as_ref(), + ChildSatisfactionOptions::new() + .with_allow_subset(allow_subset_satisfy_partitioning), + )? + .is_satisfied(); + let preserve_satisfying_file_partitioning = + preserve_file_partition_threshold_met + && !requires_grouping_id + && partitioning_satisfied + && target_partitions > child_partitions; + + // When subset satisfaction is enabled, preserve an + // already-satisfying partitioning. Otherwise, hash + // repartition may also increase parallelism. + let needs_hash_repartition = if allow_subset_satisfy_partitioning { + !partitioning_satisfied + } else { + !partitioning_satisfied + || (target_partitions > child_partitions + && !preserve_satisfying_file_partitioning) + }; + let should_add_hash_repartition = + hash_necessary && needs_hash_repartition; + // See https://github.com/apache/datafusion/issues/18341#issuecomment-3503238325 for background // When inserting hash is necessary to satisfy hash requirement, insert hash repartition. - if hash_necessary { - child = add_hash_on_top( - child, - exprs.to_vec(), - target_partitions, - allow_subset_satisfy_partitioning, - )?; + if should_add_hash_repartition { + // When there is an existing ordering, we preserve ordering during + // repartition. This will be rolled back in the future if any of the + // following conditions is true: + // - Preserving ordering is not helpful in terms of satisfying ordering + // requirements. + // - Usage of order preserving variants is not desirable (per the flag + // `config.optimizer.prefer_existing_sort`). + let partitioning = Distribution::KeyPartitioned(exprs.to_vec()) + .create_partitioning(target_partitions); + let repartition = RepartitionExec::try_new( + Arc::clone(&child.plan), + partitioning, + )? + .with_preserve_order(); + let plan = Arc::new(repartition) as _; + child = DistributionContext::new(plan, true, vec![child]); } } Distribution::UnspecifiedDistribution => { @@ -1382,73 +1394,101 @@ pub fn ensure_distribution( } }; - let streaming_benefit = if child.data { - preserving_order_enables_streaming(&plan, &child.plan)? - } else { - false - }; + Ok(DistributionChildState { + context: child, + required_input_ordering, + maintains_input_order: maintains, + requirement, + }) + }, + ) + .collect::>>()?; - // There is an ordering requirement of the operator: - if let Some(required_input_ordering) = required_input_ordering { - // Either: - // - Ordering requirement cannot be satisfied by preserving ordering through repartitions, or - // - using order preserving variant is not desirable. - let sort_req = required_input_ordering.into_single(); - let ordering_satisfied = child - .plan - .equivalence_properties() - .ordering_satisfy_requirement(sort_req.clone())?; - - if (!ordering_satisfied || !order_preserving_variants_desirable) - && !streaming_benefit - && child.data - { - child = replace_order_preserving_variants(child)?; - // If ordering requirements were satisfied before repartitioning, - // make sure ordering requirements are still satisfied after. - if ordering_satisfied { - // Make sure to satisfy ordering requirement: - child = add_sort_above_with_check( - child, - sort_req, - plan.downcast_ref::() - .map(|output| output.fetch()) - .unwrap_or(None), - )?; - } - } - // Stop tracking distribution changing operators - child.data = false; - } else { - let streaming_benefit = if child.data { - preserving_order_enables_streaming(&plan, &child.plan)? + // This is called after each child satisfies its own distribution requirement. + // It enforces relationships between child partition layouts for multi-child + // operators that process matching partition indexes together. + enforce_distribution_relationships( + plan.name(), + &input_distributions, + &mut children, + target_partitions, + )?; + + let children = children + .into_iter() + .map( + |DistributionChildState { + mut context, + required_input_ordering, + maintains_input_order, + requirement, + }| { + let streaming_benefit = if context.data { + preserving_order_enables_streaming(&plan, &context.plan)? } else { false }; - // no ordering requirement - match requirement { - // Operator requires specific distribution. - Distribution::SinglePartition | Distribution::HashPartitioned(_) => { - // If the parent doesn't maintain input order, preserving - // ordering is pointless. However, if it does maintain - // input order, we keep order-preserving variants so - // ordering can flow through to ancestors that need it. - if !maintains && !streaming_benefit { - child = replace_order_preserving_variants(child)?; + + // There is an ordering requirement of the operator: + if let Some(required_input_ordering) = required_input_ordering { + // Either: + // - Ordering requirement cannot be satisfied by preserving ordering through repartitions, or + // - using order preserving variant is not desirable. + let sort_req = required_input_ordering.into_single(); + let ordering_satisfied = context + .plan + .equivalence_properties() + .ordering_satisfy_requirement(sort_req.clone())?; + + if (!ordering_satisfied || !order_preserving_variants_desirable) + && !streaming_benefit + && context.data + { + context = replace_order_preserving_variants(context)?; + // If ordering requirements were satisfied before repartitioning, + // make sure ordering requirements are still satisfied after. + if ordering_satisfied { + // Make sure to satisfy ordering requirement: + context = add_sort_above_with_check( + context, + sort_req, + plan.downcast_ref::() + .map(|output| output.fetch()) + .unwrap_or(None), + )?; } } - Distribution::UnspecifiedDistribution => { - // Since ordering is lost, trying to preserve ordering is pointless - if !maintains || plan.is::() { - child = replace_order_preserving_variants(child)?; + // Stop tracking distribution changing operators + context.data = false; + } else { + // no ordering requirement + match requirement { + // Operator requires specific distribution. + Distribution::SinglePartition + | Distribution::HashPartitioned(_) + | Distribution::KeyPartitioned(_) => { + // If the parent doesn't maintain input order, preserving + // ordering is pointless. However, if it does maintain + // input order, we keep order-preserving variants so + // ordering can flow through to ancestors that need it. + if !maintains_input_order && !streaming_benefit { + context = replace_order_preserving_variants(context)?; + } + } + Distribution::UnspecifiedDistribution => { + // Since ordering is lost, trying to preserve ordering is pointless + if !maintains_input_order + || plan.is::() + { + context = replace_order_preserving_variants(context)?; + } } } } - } - Ok(child) - }, - ) - .collect::>>()?; + Ok(context) + }, + ) + .collect::>>()?; let children_plans = children .iter() @@ -1484,7 +1524,16 @@ pub fn ensure_distribution( // Data Arc::new(InterleaveExec::try_new(children_plans)?) } else { - plan.with_new_children(children_plans)? + // Route through `replace_children_if_necessary` so the common + // case where no child was replaced above skips the expensive + // `replace_children` rebuild. For nodes like `ProjectionExec`, + // `replace_children` recomputes schema / equivalence properties / + // output ordering via `try_new` even when the input Arcs are + // identical, which dominates `ensure_distribution` time on deep + // projection stacks over plans where no distribution change + // applies (point queries with no join / aggregate / unmet + // ordering). + replace_children_if_necessary(plan, children_plans)? }; Ok(Transformed::yes(DistributionContext::new( @@ -1514,7 +1563,8 @@ fn update_children(mut dist_context: DistributionContext) -> Result Self { - Self {} - } -} - -/// This context object is used within the [`EnforceSorting`] rule to track the closest +/// Context object used by sort enforcement to track the closest /// [`SortExec`] descendant(s) for every child of a plan. The data attribute /// stores whether the plan is a `SortExec` or is connected to a `SortExec` /// via its children. @@ -135,7 +129,7 @@ fn update_sort_ctx_children_data( Ok(node_and_ctx) } -/// This object is used within the [`EnforceSorting`] rule to track the closest +/// Tracks the closest /// [`CoalescePartitionsExec`] descendant(s) for every child of a plan. The data /// attribute stores whether the plan is a `CoalescePartitionsExec` or is /// connected to a `CoalescePartitionsExec` via its children. @@ -184,85 +178,21 @@ fn update_coalesce_ctx_children( // and connected to some `CoalescePartitionsExec`: node.data && !matches!( - coalesce_context.plan.required_input_distribution()[idx], - Distribution::SinglePartition + coalesce_context + .plan + .input_distribution_requirements() + .child_distribution(idx), + Some(Distribution::SinglePartition) ) }) }; } -/// Performs optimizations based upon a series of subrules. -/// Refer to each subrule for detailed descriptions of the optimizations performed: -/// Subrule application is ordering dependent. -/// -/// Optimizer consists of 5 main parts which work sequentially -/// 1. [`ensure_sorting`] Works down-to-top to be able to remove unnecessary [`SortExec`]s, [`SortPreservingMergeExec`]s -/// add [`SortExec`]s if necessary by a requirement and adjusts window operators. -/// 2. [`parallelize_sorts`] (Optional, depends on the `repartition_sorts` configuration) -/// Responsible to identify and remove unnecessary partition unifier operators -/// such as [`SortPreservingMergeExec`], [`CoalescePartitionsExec`] follows [`SortExec`]s does possible simplifications. -/// 3. [`replace_with_order_preserving_variants()`] Replaces with alternative operators, for example can merge -/// a [`SortExec`] and a [`CoalescePartitionsExec`] into one [`SortPreservingMergeExec`] -/// or a [`SortExec`] + [`RepartitionExec`] combination into an order preserving [`RepartitionExec`] -/// 4. [`sort_pushdown`] Works top-down. Responsible to push down sort operators as deep as possible in the plan. -/// 5. `replace_with_partial_sort` Checks if it's possible to replace [`SortExec`]s with [`PartialSortExec`] operators -impl PhysicalOptimizerRule for EnforceSorting { - fn optimize( - &self, - plan: Arc, - config: &ConfigOptions, - ) -> Result> { - let plan_requirements = PlanWithCorrespondingSort::new_default(plan); - // Execute a bottom-up traversal to enforce sorting requirements, - // remove unnecessary sorts, and optimize sort-sensitive operators: - let adjusted = plan_requirements.transform_up(ensure_sorting)?.data; - let new_plan = if config.optimizer.repartition_sorts { - let plan_with_coalesce_partitions = - PlanWithCorrespondingCoalescePartitions::new_default(adjusted.plan); - let parallel = plan_with_coalesce_partitions - .transform_up(parallelize_sorts) - .data()?; - parallel.plan - } else { - adjusted.plan - }; - - let plan_with_pipeline_fixer = OrderPreservationContext::new_default(new_plan); - let updated_plan = plan_with_pipeline_fixer - .transform_up(|plan_with_pipeline_fixer| { - replace_with_order_preserving_variants( - plan_with_pipeline_fixer, - false, - true, - config, - ) - }) - .data()?; - // Execute a top-down traversal to exploit sort push-down opportunities - // missed by the bottom-up traversal: - let mut sort_pushdown = SortPushDown::new_default(updated_plan.plan); - assign_initial_requirements(&mut sort_pushdown); - let adjusted = pushdown_sorts(sort_pushdown)?; - adjusted - .plan - .transform_up(|plan| Ok(Transformed::yes(replace_with_partial_sort(plan)?))) - .data() - } - - fn name(&self) -> &str { - "EnforceSorting" - } - - fn schema_check(&self) -> bool { - true - } -} - /// Only interested with [`SortExec`]s and their unbounded children. /// If the plan is not a [`SortExec`] or its child is not unbounded, returns the original plan. /// Otherwise, by checking the requirement satisfaction searches for a replacement chance. /// If there's one replaces the [`SortExec`] plan with a [`PartialSortExec`] -fn replace_with_partial_sort( +pub fn replace_with_partial_sort( plan: Arc, ) -> Result> { let Some(sort_plan) = plan.downcast_ref::() else { @@ -531,13 +461,17 @@ pub fn ensure_sorting( } else if is_sort_preserving_merge(&requirements.plan) && child_node.plan.output_partitioning().partition_count() <= 1 { - // This `SortPreservingMergeExec` is unnecessary, input already has a - // single partition and no fetch is required. - let mut child_node = requirements.children.swap_remove(0); + // This `SortPreservingMergeExec` is unnecessary because its input has a + // single partition. + let child_node = requirements.children.swap_remove(0); if let Some(fetch) = requirements.plan.fetch() { - // Add the limit exec if the original SPM had a fetch: - child_node.plan = - Arc::new(LocalLimitExec::new(Arc::clone(&child_node.plan), fetch)); + let mut limit = LocalLimitExec::new(Arc::clone(&child_node.plan), fetch); + limit.set_required_ordering(requirements.plan.output_ordering().cloned()); + return Ok(Transformed::yes(PlanContext::new( + Arc::new(limit), + false, + vec![child_node], + ))); } return Ok(Transformed::yes(child_node)); } @@ -612,24 +546,43 @@ fn adjust_window_sort_removal( let child_node = remove_corresponding_sort_from_sub_plan( window_tree.children.swap_remove(0), matches!( - window_tree.plan.required_input_distribution()[0], - Distribution::SinglePartition + window_tree + .plan + .input_distribution_requirements() + .child_distribution(0), + Some(Distribution::SinglePartition) ), )?; window_tree.children.push(child_node); let child_plan = &window_tree.children[0].plan; + // Captured up-front so the fallback `BoundedWindowAggExec::try_new` below + // can reinstall the observer that was on the source exec. `None` when + // the source is a `WindowAggExec` (no observer) or when no observer was + // installed on the source `BoundedWindowAggExec`. + let state_observer = window_tree + .plan + .downcast_ref::() + .and_then(|exec| exec.state_observer().cloned()); let (window_expr, new_window) = if let Some(exec) = window_tree.plan.downcast_ref::() { let window_expr = exec.window_expr(); - let new_window = - get_best_fitting_window(window_expr, child_plan, &exec.partition_keys())?; + let new_window = get_best_fitting_window( + window_expr, + child_plan, + &exec.partition_keys(), + None, + )?; (window_expr, new_window) } else if let Some(exec) = window_tree.plan.downcast_ref::() { let window_expr = exec.window_expr(); - let new_window = - get_best_fitting_window(window_expr, child_plan, &exec.partition_keys())?; + let new_window = get_best_fitting_window( + window_expr, + child_plan, + &exec.partition_keys(), + state_observer.clone(), + )?; (window_expr, new_window) } else { return plan_err!("Expected WindowAggExec or BoundedWindowAggExec"); @@ -652,12 +605,15 @@ fn adjust_window_sort_removal( window_tree.children.push(child_node); if window_expr.iter().all(|e| e.uses_bounded_memory()) { - Arc::new(BoundedWindowAggExec::try_new( - window_expr.to_vec(), - child_plan, - InputOrderMode::Sorted, - !window_expr[0].partition_by().is_empty(), - )?) as _ + Arc::new( + BoundedWindowAggExec::try_new( + window_expr.to_vec(), + child_plan, + InputOrderMode::Sorted, + !window_expr[0].partition_by().is_empty(), + )? + .with_state_observer(state_observer)?, + ) as _ } else { Arc::new(WindowAggExec::try_new( window_expr.to_vec(), @@ -676,11 +632,49 @@ fn adjust_window_sort_removal( /// the plan, some of the remaining `RepartitionExec`s might become unnecessary. /// Removes such `RepartitionExec`s from the plan as well. fn remove_bottleneck_in_subplan( + requirements: PlanWithCorrespondingCoalescePartitions, +) -> Result { + // The root is the node `parallelize_sorts` is rewriting (a `SortExec`, + // `SortPreservingMergeExec` or `CoalescePartitionsExec`). Its own distribution + // requirement does not constrain the removal, because the caller drops the node and + // rebuilds the cascade around the result. + remove_bottleneck_in_subplan_impl(requirements, true) +} + +fn remove_bottleneck_in_subplan_impl( mut requirements: PlanWithCorrespondingCoalescePartitions, + is_root: bool, ) -> Result { let plan = &requirements.plan; + // Below the root, a `CoalescePartitionsExec` feeding a child that requires + // `Distribution::SinglePartition` is not an avoidable bottleneck: it is what satisfies + // that requirement. Removing it leaves the parent with a multi-partition input it cannot + // accept, and nothing re-runs distribution enforcement afterwards, so the plan reaches + // `SanityCheckPlan` invalid. The traversal reaches such a node because + // `update_coalesce_ctx_children` marks a node as connected when *any* child qualifies: + // a `CollectLeft` `HashJoinExec` whose probe side is connected is descended into even + // though its build side must stay single-partition. + // + // Only `SinglePartition` is protected. A `HashPartitioned` child is in principle in the + // same position — a single-partition input trivially satisfies a hash requirement, so a + // coalesce below one is also load-bearing — but nothing puts a coalesce there: + // `ensure_distribution` satisfies a hash requirement with a `RepartitionExec`, never a + // `CoalescePartitionsExec`. Widening the check would be dead code today. + let dist_reqs = plan.input_distribution_requirements(); + let removable = |idx: usize| { + is_root + || !matches!( + dist_reqs.child_distribution(idx), + Some(Distribution::SinglePartition) + ) + }; + let remove_from_first_child = requirements + .children + .first() + .is_some_and(|child| is_coalesce_partitions(&child.plan)) + && removable(0); let children = &mut requirements.children; - if is_coalesce_partitions(&children[0].plan) { + if remove_from_first_child { // We can safely use the 0th index since we have a `CoalescePartitionsExec`. let mut new_child_node = children[0].children.swap_remove(0); while new_child_node.plan.output_partitioning() == plan.output_partitioning() @@ -694,9 +688,14 @@ fn remove_bottleneck_in_subplan( requirements.children = requirements .children .into_iter() - .map(|node| { - if node.data { - remove_bottleneck_in_subplan(node) + .enumerate() + .map(|(idx, node)| { + // Deliberately conservative: not descending at all also skips legitimate + // cleanups *below* a protected child (a redundant second coalesce under the + // load-bearing one, say). This could later be narrowed to "descend, but + // protect only the topmost coalesce" if that turns out to matter. + if node.data && removable(idx) { + remove_bottleneck_in_subplan_impl(node, false) } else { Ok(node) } @@ -727,8 +726,10 @@ fn update_child_to_remove_unnecessary_sort( ) -> Result { if node.data { let requires_single_partition = matches!( - parent.required_input_distribution()[child_idx], - Distribution::SinglePartition + parent + .input_distribution_requirements() + .child_distribution(child_idx), + Some(Distribution::SinglePartition) ); node = remove_corresponding_sort_from_sub_plan(node, requires_single_partition)?; } @@ -749,7 +750,7 @@ fn remove_corresponding_sort_from_sub_plan( } } else { let mut any_connection = false; - let required_dist = node.plan.required_input_distribution(); + let required_dist = node.plan.input_distribution_requirements().into_per_child(); node.children = node .children .into_iter() diff --git a/datafusion/physical-optimizer/src/enforce_sorting/replace_with_order_preserving_variants.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs similarity index 100% rename from datafusion/physical-optimizer/src/enforce_sorting/replace_with_order_preserving_variants.rs rename to datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/replace_with_order_preserving_variants.rs diff --git a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs similarity index 65% rename from datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs rename to datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs index 400161a94cff4..5c17ffbd1e7db 100644 --- a/datafusion/physical-optimizer/src/enforce_sorting/sort_pushdown.rs +++ b/datafusion/physical-optimizer/src/ensure_requirements/enforce_sorting/sort_pushdown.rs @@ -19,7 +19,8 @@ use std::fmt::Debug; use std::sync::Arc; use crate::utils::{ - add_sort_above, is_sort, is_sort_preserving_merge, is_union, is_window, + add_sort_above_with_distribution, is_sort, is_sort_preserving_merge, is_union, + is_window, }; use arrow::datatypes::SchemaRef; @@ -29,7 +30,7 @@ use datafusion_expr::JoinType; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::utils::collect_columns; use datafusion_physical_expr::{ - EquivalenceProperties, add_offset_to_physical_sort_exprs, + Distribution, EquivalenceProperties, add_offset_to_physical_sort_exprs, }; use datafusion_physical_expr_common::sort_expr::{ LexOrdering, LexRequirement, OrderingRequirements, PhysicalSortExpr, @@ -42,23 +43,36 @@ use datafusion_physical_plan::joins::utils::{ ColumnIndex, calculate_join_output_ordering, }; use datafusion_physical_plan::joins::{HashJoinExec, SortMergeJoinExec}; -use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; use datafusion_physical_plan::repartition::RepartitionExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::tree_node::PlanContext; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; -/// This is a "data class" we use within the [`EnforceSorting`] rule to push -/// down [`SortExec`] in the plan. In some cases, we can reduce the total -/// computational cost by pushing down `SortExec`s through some executors. The -/// object carries the parent required ordering and the (optional) `fetch` value -/// of the parent node as its data. -/// -/// [`EnforceSorting`]: crate::enforce_sorting::EnforceSorting -#[derive(Default, Clone, Debug)] +/// "Data class" used by sort pushdown (now driven from `EnsureRequirements`) +/// to push down [`SortExec`] in the plan. In some cases the total +/// computational cost is reduced by pushing down `SortExec`s through certain +/// executors. The object carries the parent required ordering, the (optional) +/// `fetch` value of the parent node, and the parent's distribution requirement +/// (used by the distribution-aware pushdown path) as its data. +#[derive(Clone, Debug)] pub struct ParentRequirements { ordering_requirement: Option, fetch: Option, + /// The distribution required by the consumer above any SortExec we insert. + /// When this is `SinglePartition` and the input has multiple partitions, + /// `add_sort_above_with_distribution` wraps the sort in `SortPreservingMergeExec`. + distribution_requirement: Distribution, +} + +impl Default for ParentRequirements { + fn default() -> Self { + Self { + ordering_requirement: None, + fetch: None, + distribution_requirement: Distribution::UnspecifiedDistribution, + } + } } pub type SortPushDown = PlanContext; @@ -66,12 +80,20 @@ pub type SortPushDown = PlanContext; /// Assigns the ordering requirement of the root node to the its children. pub fn assign_initial_requirements(sort_push_down: &mut SortPushDown) { let reqs = sort_push_down.plan.required_input_ordering(); - for (child, requirement) in sort_push_down.children.iter_mut().zip(reqs) { + let dists = sort_push_down + .plan + .input_distribution_requirements() + .into_per_child(); + for (idx, (child, requirement)) in + sort_push_down.children.iter_mut().zip(reqs).enumerate() + { child.data = ParentRequirements { ordering_requirement: requirement, - // If the parent has a fetch value, assign it to the children - // Or use the fetch value of the child. fetch: child.plan.fetch(), + distribution_requirement: dists + .get(idx) + .cloned() + .unwrap_or(Distribution::UnspecifiedDistribution), }; } } @@ -92,11 +114,35 @@ fn min_fetch(f1: Option, f2: Option) -> Option { } } +/// Returns the stricter of two distribution requirements. +/// `SinglePartition` is the strictest. +#[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" +)] +fn stronger_distribution(a: &Distribution, b: &Distribution) -> Distribution { + match (a, b) { + (Distribution::SinglePartition, _) | (_, Distribution::SinglePartition) => { + Distribution::SinglePartition + } + (Distribution::HashPartitioned(exprs), _) + | (Distribution::KeyPartitioned(exprs), _) => { + Distribution::KeyPartitioned(exprs.clone()) + } + (_, Distribution::HashPartitioned(exprs)) + | (_, Distribution::KeyPartitioned(exprs)) => { + Distribution::KeyPartitioned(exprs.clone()) + } + _ => Distribution::UnspecifiedDistribution, + } +} + fn pushdown_sorts_helper( mut sort_push_down: SortPushDown, ) -> Result> { let plan = sort_push_down.plan; let parent_fetch = sort_push_down.data.fetch; + let parent_distribution = sort_push_down.data.distribution_requirement.clone(); let Some(parent_requirement) = sort_push_down.data.ordering_requirement.clone() else { @@ -121,6 +167,17 @@ fn pushdown_sorts_helper( return pushdown_sorts_helper(sort_push_down); } sort_push_down.plan = plan; + // No ordering is being pushed; use each child's own distribution requirement + let dists = sort_push_down + .plan + .input_distribution_requirements() + .into_per_child(); + for (idx, child) in sort_push_down.children.iter_mut().enumerate() { + child.data.distribution_requirement = dists + .get(idx) + .cloned() + .unwrap_or(Distribution::UnspecifiedDistribution); + } return Ok(Transformed::no(sort_push_down)); }; @@ -149,22 +206,29 @@ fn pushdown_sorts_helper( // The sort was imposing a different ordering than the one being // pushed down. Replace it with a sort that matches the pushed-down // ordering, and continue the pushdown. - // Add back the sort: - sort_push_down = add_sort_above( + // Add back the sort (distribution-aware): + sort_push_down = add_sort_above_with_distribution( sort_push_down, parent_requirement.into_single(), parent_fetch, + &parent_distribution, ); // Update pushdown requirements: sort_push_down.children[0].data = ParentRequirements { ordering_requirement: Some(OrderingRequirements::from(sort_ordering)), fetch: sort_fetch, + distribution_requirement: Distribution::UnspecifiedDistribution, }; return Ok(Transformed::yes(sort_push_down)); } else { // Sort was unnecessary, just propagate the stricter fetch and - // ordering requirements: + // ordering requirements. Reset distribution to Unspecified + // because the sort we're removing may have been below a + // partition-merging node (like SortPreservingMergeExec) that + // already satisfies SinglePartition. sort_push_down.data.fetch = min_fetch(sort_fetch, parent_fetch); + sort_push_down.data.distribution_requirement = + Distribution::UnspecifiedDistribution; let current_is_stricter = eqp.requirements_compatible( sort_ordering.clone().into(), parent_requirement.first().clone(), @@ -184,10 +248,31 @@ fn pushdown_sorts_helper( if satisfy_parent { // For non-sort operators which satisfy ordering: let reqs = sort_push_down.plan.required_input_ordering(); + let dists = sort_push_down + .plan + .input_distribution_requirements() + .into_per_child(); - for (child, order) in sort_push_down.children.iter_mut().zip(reqs) { + // If this node already outputs single partition, don't push SinglePartition + // requirement to children (they're below the merge point). + let effective_parent_dist = + if sort_push_down.plan.output_partitioning().partition_count() == 1 { + Distribution::UnspecifiedDistribution + } else { + parent_distribution.clone() + }; + + for (idx, (child, order)) in + sort_push_down.children.iter_mut().zip(reqs).enumerate() + { child.data.ordering_requirement = order; child.data.fetch = min_fetch(parent_fetch, child.data.fetch); + child.data.distribution_requirement = stronger_distribution( + &effective_parent_dist, + dists + .get(idx) + .unwrap_or(&Distribution::UnspecifiedDistribution), + ); } } else if let Some(adjusted) = pushdown_requirement_to_children( &sort_push_down.plan, @@ -195,19 +280,39 @@ fn pushdown_sorts_helper( parent_fetch, )? { // For operators that can take a sort pushdown, continue with updated - // requirements: + // requirements. If this node already outputs single partition (e.g. SPM), + // don't push SinglePartition to children. let current_fetch = sort_push_down.plan.fetch(); - for (child, order) in sort_push_down.children.iter_mut().zip(adjusted) { + let dists = sort_push_down + .plan + .input_distribution_requirements() + .into_per_child(); + let effective_dist = + if sort_push_down.plan.output_partitioning().partition_count() == 1 { + Distribution::UnspecifiedDistribution + } else { + parent_distribution.clone() + }; + for (idx, (child, order)) in + sort_push_down.children.iter_mut().zip(adjusted).enumerate() + { child.data.ordering_requirement = order; child.data.fetch = min_fetch(current_fetch, parent_fetch); + child.data.distribution_requirement = stronger_distribution( + &effective_dist, + dists + .get(idx) + .unwrap_or(&Distribution::UnspecifiedDistribution), + ); } sort_push_down.data.ordering_requirement = None; } else { - // Can not push down requirements, add new `SortExec`: - sort_push_down = add_sort_above( + // Can not push down requirements, add new `SortExec` (distribution-aware): + sort_push_down = add_sort_above_with_distribution( sort_push_down, parent_requirement.into_single(), parent_fetch, + &parent_distribution, ); assign_initial_requirements(&mut sort_push_down); } @@ -263,7 +368,20 @@ fn pushdown_requirement_to_children( return Ok(None); }; match determine_children_requirement(&parent_required, &child_req, child_plan) { - RequirementsCompatibility::Satisfy => Ok(Some(vec![Some(child_req)])), + RequirementsCompatibility::Satisfy => { + // Window input requirements may be empty or constant-only. + // Such requirements do not guarantee the parent's output ordering, so + // keep the sort above the window unless the window output is known + // to satisfy it. + if !plan + .equivalence_properties() + .ordering_satisfy_requirement(parent_required.first().clone())? + { + return Ok(None); + } + + Ok(Some(vec![Some(child_req)])) + } RequirementsCompatibility::Compatible(adjusted) => { // If parent requirements are more specific than output ordering // of the window plan, then we can deduce that the parent expects @@ -271,7 +389,7 @@ fn pushdown_requirement_to_children( // that's the case, we block the pushdown of sort operation. if !plan .equivalence_properties() - .ordering_satisfy_requirement(parent_required.into_single())? + .ordering_satisfy_requirement(parent_required.first().clone())? { return Ok(None); } @@ -304,14 +422,34 @@ fn pushdown_requirement_to_children( // Push down through operator with fetch when: // - requirement is aligned with output ordering // - it preserves ordering during execution + // + // A `ProjectionExec` reports a `fetch()` forwarded from its input and + // can renumber/reorder columns, so the requirement (expressed in the + // projection's output schema) must be remapped into the child schema + // before being pushed down — forwarding it unchanged would let a key + // such as `score@1` (valid in the output schema) refer to a different + // column in the child schema, producing a `SortExec` whose key points + // at the wrong column ("does not satisfy order requirements ... + // Child-0 order: []"). If a required column maps to a computed + // (non-`Column`) projection expression it cannot be expressed below the + // projection, so the sort is kept above it. + let child_required = + if let Some(projection) = plan.downcast_ref::() { + match remap_requirement_through_projection(projection, &parent_required) { + Some(remapped) => remapped, + None => return Ok(None), + } + } else { + parent_required.clone() + }; let Some(ordering) = plan.properties().output_ordering() else { - return Ok(Some(vec![Some(parent_required)])); + return Ok(Some(vec![Some(child_required)])); }; if plan.properties().eq_properties.requirements_compatible( parent_required.first().clone(), ordering.clone().into(), ) { - Ok(Some(vec![Some(parent_required)])) + Ok(Some(vec![Some(child_required)])) } else { Ok(None) } @@ -356,12 +494,12 @@ fn pushdown_requirement_to_children( } } else if let Some(aggregate_exec) = plan.downcast_ref::() { handle_aggregate_pushdown(aggregate_exec, parent_required) + } else if let Some(projection_exec) = plan.downcast_ref::() { + handle_projection_pushdown(projection_exec, &parent_required) } else if maintains_input_order.is_empty() || !maintains_input_order.iter().any(|o| *o) || plan.is::() || plan.is::() - // TODO: Add support for Projection push down - || plan.is::() || pushdown_would_violate_requirements(&parent_required, plan.as_ref()) { // If the current plan is a leaf node or can not maintain any of the input ordering, can not pushed down requirements. @@ -388,7 +526,51 @@ fn pushdown_requirement_to_children( } else { handle_custom_pushdown(plan, parent_required, &maintains_input_order) } - // TODO: Add support for Projection push down +} + +/// Remap an ordering requirement expressed in a [`ProjectionExec`]'s output +/// schema into its child (input) schema. +/// +/// Every alternative requirement is remapped independently, and the +/// hard/soft-ness of the original [`OrderingRequirements`] is preserved. An +/// alternative that references a computed (non-[`Column`]) projection +/// expression cannot be expressed in the child schema and is dropped; if every +/// alternative drops out, this returns `None` (and pushdown is declined, i.e. +/// the sort is kept above the projection). +fn remap_requirement_through_projection( + projection: &ProjectionExec, + parent_required: &OrderingRequirements, +) -> Option { + let exprs = projection.expr(); + let (alternatives, soft) = parent_required.clone().into_alternatives(); + let remapped = alternatives + .iter() + .filter_map(|req| remap_lex_requirement_through_projection(exprs, req)); + OrderingRequirements::new_alternatives(remapped, soft) +} + +/// Remap a single [`LexRequirement`] expressed in a [`ProjectionExec`]'s output +/// schema into its child (input) schema. +/// +/// Each requirement column at output index `i` is rewritten to the column the +/// projection produces at that index (`projection.expr()[i]`). Returns `None` +/// if any required column maps to a computed (non-[`Column`]) projection +/// expression, since that ordering cannot be expressed in the child schema. +fn remap_lex_requirement_through_projection( + exprs: &[ProjectionExpr], + req: &LexRequirement, +) -> Option { + let mut child_reqs = Vec::with_capacity(req.len()); + for sort_req in req.iter() { + let col = sort_req.expr.downcast_ref::()?; + let proj_expr = exprs.get(col.index())?; + let child_col = proj_expr.expr.downcast_ref::()?; + child_reqs.push(PhysicalSortRequirement::new( + Arc::new(child_col.clone()), + sort_req.options, + )); + } + LexRequirement::new(child_reqs) } /// Try to push sorting through [`AggregateExec`] @@ -801,12 +983,11 @@ fn handle_hash_join( } else { column_indices.iter().collect() }; - let len_of_left_fields = projected_indices - .iter() - .filter(|ci| ci.side == JoinSide::Left) - .count(); - - let all_from_right_child = all_indices.iter().all(|i| *i >= len_of_left_fields); + let all_from_right_child = all_indices.iter().all(|i| { + projected_indices + .get(*i) + .is_some_and(|ci| ci.side == JoinSide::Right) + }); let plan_children = plan.children(); @@ -881,3 +1062,224 @@ enum RequirementsCompatibility { /// Requirements not compatible NonCompatible, } + +/// Attempts to push parent ordering requirements through a [`ProjectionExec`]. +/// +/// This is safe when every required sort expression refers to a projected output +/// column that is backed by a simple input column. In that case, the requirement +/// can be remapped from the projection output schema to the projection input +/// schema while preserving the original sort options. +/// +/// For example, a parent requirement on `a@2` over: +/// +/// ```text +/// ProjectionExec: expr=[c@2 as c, b@1 as b, a@0 as a] +/// ``` +/// +/// is remapped to a child requirement on `a@0`. +/// +/// The implementation is intentionally conservative: computed projection +/// expressions and non-column sort expressions are not pushed down. Returning +/// `Ok(None)` leaves sorting above the projection, preserving correctness. +fn handle_projection_pushdown( + projection_exec: &ProjectionExec, + parent_required: &OrderingRequirements, +) -> Result>>> { + // Only push sorting through pure column projections. Source-dependent + // expressions must stay close enough to the scan to be rewritten + // by the source and cannot be evaluated by [`ProjectionExec`]. + if projection_exec + .expr() + .iter() + .any(|expr| !expr.expr.is::()) + { + return Ok(None); + } + + Ok( + remap_requirement_through_projection(projection_exec, parent_required) + .map(|requirements| vec![Some(requirements)]), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_expr::Operator; + use datafusion_physical_expr::PhysicalExpr; + use datafusion_physical_expr::expressions::{BinaryExpr, col}; + use datafusion_physical_plan::empty::EmptyExec; + + const DESC: SortOptions = SortOptions { + descending: true, + nulls_first: false, + }; + const ASC: SortOptions = SortOptions { + descending: false, + nulls_first: true, + }; + + /// Child (input) schema fed to the projections under test: `[a, b, c]`. + fn child_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + Field::new("c", DataType::Int32, true), + ])) + } + + /// A projection over `[a, b, c]` whose output is `[a@0, c@2 as score, + /// b@1 as value]` — i.e. it *reorders* (`c` moves index 2 -> 1) and renames. + fn reordering_projection() -> Arc { + let schema = child_schema(); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + Arc::new( + ProjectionExec::try_new( + vec![ + (col("a", &schema).unwrap(), "a".to_string()), + (col("c", &schema).unwrap(), "score".to_string()), + (col("b", &schema).unwrap(), "value".to_string()), + ], + input, + ) + .unwrap(), + ) + } + + /// A projection over `[a, b, c]` whose output is `[a@0, b + c as computed]`, + /// so output column index 1 maps to a *computed* (non-`Column`) expression. + fn computed_projection() -> Arc { + let schema = child_schema(); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let b_plus_c = Arc::new(BinaryExpr::new( + col("b", &schema).unwrap(), + Operator::Plus, + col("c", &schema).unwrap(), + )) as Arc; + Arc::new( + ProjectionExec::try_new( + vec![ + (col("a", &schema).unwrap(), "a".to_string()), + (b_plus_c, "computed".to_string()), + ], + input, + ) + .unwrap(), + ) + } + + /// `PhysicalSortRequirement` for `@ ` in `schema`. + fn req(name: &str, schema: &Schema, options: SortOptions) -> PhysicalSortRequirement { + PhysicalSortRequirement::new(col(name, schema).unwrap(), Some(options)) + } + + fn lex(reqs: impl IntoIterator) -> LexRequirement { + LexRequirement::new(reqs).unwrap() + } + + #[test] + fn remap_single_hard_requirement_through_reordering_projection() { + let projection = reordering_projection(); + let out = projection.schema(); + let child = child_schema(); + + // `score@1 DESC, a@0 ASC` in the output schema. + let required = OrderingRequirements::new(lex([ + req("score", &out, DESC), + req("a", &out, ASC), + ])); + + let remapped = + remap_requirement_through_projection(&projection, &required).unwrap(); + + // `score@1` -> `c@2`, `a@0` -> `a@0`; still a single hard requirement. + let expected = OrderingRequirements::new(lex([ + req("c", &child, DESC), + req("a", &child, ASC), + ])); + assert_eq!(remapped, expected); + } + + #[test] + fn remap_preserves_softness() { + let projection = reordering_projection(); + let out = projection.schema(); + let child = child_schema(); + + let required = OrderingRequirements::new_soft(lex([req("score", &out, DESC)])); + let remapped = + remap_requirement_through_projection(&projection, &required).unwrap(); + + let expected = OrderingRequirements::new_soft(lex([req("c", &child, DESC)])); + assert_eq!(remapped, expected); + // Hardness/softness is preserved through the remap. + assert!(matches!(remapped, OrderingRequirements::Soft(_))); + } + + #[test] + fn remap_preserves_all_hard_alternatives() { + let projection = reordering_projection(); + let out = projection.schema(); + let child = child_schema(); + + // Two alternatives: `score@1 DESC` or `a@0 ASC, value@2 ASC`. + let mut required = OrderingRequirements::new(lex([req("score", &out, DESC)])); + required.add_alternative(lex([req("a", &out, ASC), req("value", &out, ASC)])); + + let remapped = + remap_requirement_through_projection(&projection, &required).unwrap(); + + // Both alternatives survive and are remapped; hardness preserved. + let (alts, soft) = remapped.into_alternatives(); + assert!(!soft); + assert_eq!(alts.len(), 2); + assert_eq!(alts[0], lex([req("c", &child, DESC)])); + // `value@2` -> `b@1`, `a@0` -> `a@0`. + assert_eq!(alts[1], lex([req("a", &child, ASC), req("b", &child, ASC)])); + } + + #[test] + fn remap_drops_unsatisfiable_alternative_but_keeps_others() { + let projection = computed_projection(); + let out = projection.schema(); + let child = child_schema(); + + // Alt 1 (`a@0 ASC`) is expressible below the projection; alt 2 + // (`computed@1 DESC`) maps to `b + c` and is not. + let mut required = OrderingRequirements::new(lex([req("a", &out, ASC)])); + required.add_alternative(lex([req("computed", &out, DESC)])); + + let remapped = + remap_requirement_through_projection(&projection, &required).unwrap(); + + // Only the satisfiable alternative is kept; hardness preserved. + let (alts, soft) = remapped.into_alternatives(); + assert!(!soft); + assert_eq!(alts.len(), 1); + assert_eq!(alts[0], lex([req("a", &child, ASC)])); + } + + #[test] + fn remap_declines_when_required_column_is_computed() { + let projection = computed_projection(); + let out = projection.schema(); + + // The only required column maps to a computed expression -> decline. + let required = OrderingRequirements::new(lex([req("computed", &out, DESC)])); + assert!(remap_requirement_through_projection(&projection, &required).is_none()); + } + + #[test] + fn remap_declines_when_all_alternatives_are_computed() { + let projection = computed_projection(); + let out = projection.schema(); + + let mut required = OrderingRequirements::new(lex([req("computed", &out, DESC)])); + required.add_alternative(lex([req("computed", &out, ASC)])); + + assert!(remap_requirement_through_projection(&projection, &required).is_none()); + } +} diff --git a/datafusion/physical-optimizer/src/ensure_requirements/mod.rs b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs new file mode 100644 index 0000000000000..41a03bb031629 --- /dev/null +++ b/datafusion/physical-optimizer/src/ensure_requirements/mod.rs @@ -0,0 +1,259 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! [`EnsureRequirements`] optimizer rule that enforces distribution and +//! sorting requirements together so that the two never invalidate each other. +//! +//! This rule replaces the separate `EnforceDistribution` + `EnforceSorting` +//! rules with a unified approach inspired by Apache Spark's `EnsureRequirements` +//! and Presto/Trino's `AddExchanges`. +//! +//! # Motivation +//! +//! The previous two-rule design (`EnforceDistribution` then `EnforceSorting`) +//! suffers from non-idempotent composition: `EnforceSorting`'s `pushdown_sorts` +//! can break distribution invariants established by `EnforceDistribution`, +//! because `SortExec.preserve_partitioning` couples sorting and distribution +//! decisions. See for details. +//! +//! # Architecture +//! +//! `optimize` runs several tree traversals. The defining property of this +//! rule is **Phase 2**: a single combined bottom-up pass that resolves +//! distribution *and* sorting for each node together. The surrounding phases +//! are independent traversals (top-down join-key reorder, then several +//! follow-up sort/order rewrites). Some of those could be consolidated +//! further in a follow-up. +//! +//! ```text +//! EnsureRequirements::optimize(plan) +//! │ +//! ├─ Phase 1: top-down join-key reorder (adjust_input_keys_ordering) +//! │ +//! ├─ Phase 2: combined distribution + sorting (single bottom-up pass) +//! │ └─ For each node (bottom-up), for each child: +//! │ Step 1: ensure distribution requirement +//! │ └─ insert RepartitionExec / CoalescePartitionsExec / +//! │ SortPreservingMergeExec as needed +//! │ Step 2: ensure ordering requirement (distribution-aware) +//! │ └─ insert SortExec with the correct `preserve_partitioning`, +//! │ with SortPreservingMergeExec on top if needed +//! │ +//! └─ Phase 3: small follow-up passes (bottom-up unless noted) +//! ├─ parallelize_sorts +//! ├─ replace_with_order_preserving_variants +//! ├─ pushdown_sorts (recursive walk) +//! └─ replace_with_partial_sort +//! ``` +//! +//! # Key Properties +//! +//! - **Idempotent across the whole rule**: Running `EnsureRequirements` +//! twice produces the same plan. This is the property that fixes +//! , where the old +//! two-rule pipeline could regress a parallel sort plan into a serial one +//! on pass 2. +//! - **Distribution before sorting**: For each child, distribution is +//! resolved before ordering, so sorting decisions always have full +//! distribution context. +//! - **Sort pushdown is implicit**: Phase 2 only adds `SortExec` where the +//! child doesn't already satisfy the ordering requirement, so sorts land +//! at the deepest valid position without a separate destructive pass. +//! +//! # Behavior: parallelism via repartitioning +//! +//! Phase 2 Step 1 inserts `RepartitionExec` to satisfy distribution +//! requirements. When configuration allows, it also increases parallelism by +//! repartitioning over otherwise-serial inputs. For example, given two +//! 1-partition inputs feeding an operator that can run with more +//! parallelism: +//! +//! ```text +//! ┌─────────────────────────────────┐ +//! │ ExecutionPlan │ +//! └─────────────────────────────────┘ +//! ▲ ▲ +//! │ │ +//! ┌───────────┐ ┌───────────┐ +//! │ batch A │ │ batch B │ Input: 2 partitions +//! └───────────┘ └───────────┘ +//! ``` +//! +//! `EnsureRequirements` inserts a `RepartitionExec` so the operator runs +//! with three partitions: +//! +//! ```text +//! ┌─────────────────────────────────┐ +//! │ ExecutionPlan │ Input now has 3 partitions +//! └─────────────────────────────────┘ +//! ▲ ▲ ▲ +//! └──────┼───────┘ +//! │ +//! ┌─────────────────────────────────┐ +//! │ RepartitionExec(3) │ batches are repartitioned +//! │ RoundRobin │ +//! └─────────────────────────────────┘ +//! ▲ ▲ +//! ┌───────────┐ ┌───────────┐ +//! │ batch A │ │ batch B │ +//! └───────────┘ └───────────┘ +//! ``` +//! +//! # Behavior: joint distribution + sorting +//! +//! Resolving distribution and sorting together lets Phase 2 produce a +//! parallel sort plan in cases where the two-rule pipeline historically +//! risked a serial one. Given `Sort(DESC) ← Coalesce ← MultiPartitionSource`, +//! `EnsureRequirements` rewrites it into: +//! +//! ```text +//! SortPreservingMergeExec: [a DESC] (cheap k-way merge of sorted streams) +//! SortExec: [a DESC], preserve_partitioning=true (N sorts run in parallel) +//! MultiPartitionSource +//! ``` +//! +//! Each input partition is sorted in parallel, then a `SortPreservingMergeExec` +//! at the top performs a cheap merge of pre-sorted streams. For TopK queries +//! (`fetch=K`), each parallel sort only keeps K rows per partition, so total +//! memory is `N × K` rather than coalescing the entire stream first. +//! +//! # Behavior: strictest distribution match for joins +//! +//! Distribution requirements are met in the strictest way. For example, a +//! hash join with keys `(a, b, c)` requires `Distribution(a, b, c)`. This +//! can in principle be satisfied by partitioning on any superset of any +//! subset of `(a, b, c)`, but this rule always partitions on the exact key +//! tuple `(a, b, c)`. This is sometimes more aggressive than strictly +//! necessary, but the strictest match helps avoid data skew in joins. + +// Internal implementation modules. Re-exported from `crate` root for tests +// in `core/tests/physical_optimizer/{enforce_distribution,enforce_sorting}.rs`. +pub mod enforce_distribution; +pub mod enforce_sorting; + +use std::sync::Arc; + +use crate::PhysicalOptimizerRule; + +use datafusion_common::Result; +use datafusion_common::config::ConfigOptions; +use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_physical_plan::ExecutionPlan; + +/// Optimizer rule that enforces both distribution and sorting requirements. +/// +/// This rule combines the functionality of `EnforceDistribution` and +/// `EnforceSorting` into a coordinated sequence where distribution is +/// always settled before sorting for each operator, preventing the +/// non-idempotent interactions between the two separate rules. +/// +/// See [module level documentation](self) for more details. +#[derive(Default, Debug)] +pub struct EnsureRequirements {} + +impl EnsureRequirements { + /// Create a new `EnsureRequirements` optimizer rule. + pub fn new() -> Self { + Self {} + } +} + +impl PhysicalOptimizerRule for EnsureRequirements { + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result> { + // Phase 1: Join key reordering (top-down, from EnforceDistribution) + use super::enforce_distribution::{ + PlanWithKeyRequirements, adjust_input_keys_ordering, + }; + let top_down_join_key_reordering = config.optimizer.top_down_join_key_reordering; + let plan = if top_down_join_key_reordering { + let ctx = PlanWithKeyRequirements::new_default(plan); + ctx.transform_down(adjust_input_keys_ordering).data()?.plan + } else { + use super::enforce_distribution::reorder_join_keys_to_inputs; + plan.transform_up(|p| Ok(Transformed::yes(reorder_join_keys_to_inputs(p)?))) + .data()? + }; + + // Phase 2: Combined distribution + sorting enforcement (single bottom-up pass) + // For each node: distribution first, then sorting. + use super::enforce_distribution::{DistributionContext, ensure_distribution}; + use super::enforce_sorting::{PlanWithCorrespondingSort, ensure_sorting}; + + // Step 2a: Distribution enforcement (bottom-up) + let dist_ctx = DistributionContext::new_default(plan); + let dist_ctx = dist_ctx + .transform_up(|ctx| ensure_distribution(ctx, config)) + .data()?; + + // Step 2b: Sorting enforcement (bottom-up) — runs on distribution-fixed plan + let sort_ctx = PlanWithCorrespondingSort::new_default(dist_ctx.plan); + let sort_ctx = sort_ctx.transform_up(ensure_sorting)?.data; + + // Phase 3: Optimization passes + // 3a: Parallelize sorts (Coalesce+Sort → SPM+Sort) + use super::enforce_sorting::{ + PlanWithCorrespondingCoalescePartitions, parallelize_sorts, + replace_with_partial_sort, + }; + let plan = if config.optimizer.repartition_sorts { + let ctx = PlanWithCorrespondingCoalescePartitions::new_default(sort_ctx.plan); + ctx.transform_up(parallelize_sorts).data()?.plan + } else { + sort_ctx.plan + }; + + // 3b: Order-preserving variants + use super::enforce_sorting::replace_with_order_preserving_variants::{ + OrderPreservationContext, replace_with_order_preserving_variants, + }; + let ctx = OrderPreservationContext::new_default(plan); + let plan = ctx + .transform_up(|c| { + replace_with_order_preserving_variants(c, false, true, config) + }) + .data()? + .plan; + + // 3c: Sort pushdown (distribution-aware) + use super::enforce_sorting::sort_pushdown::{ + SortPushDown, assign_initial_requirements, pushdown_sorts, + }; + let mut sort_pushdown = SortPushDown::new_default(plan); + assign_initial_requirements(&mut sort_pushdown); + let adjusted = pushdown_sorts(sort_pushdown)?; + + // 3d: Partial sort + adjusted + .plan + .transform_up(|p| Ok(Transformed::yes(replace_with_partial_sort(p)?))) + .data() + } + + fn name(&self) -> &str { + "EnsureRequirements" + } + + fn schema_check(&self) -> bool { + true + } +} + +// See tests in datafusion/core/tests/physical_optimizer/ensure_requirements.rs diff --git a/datafusion/physical-optimizer/src/filter_pushdown.rs b/datafusion/physical-optimizer/src/filter_pushdown.rs index 28f8155002a50..18fe151000511 100644 --- a/datafusion/physical-optimizer/src/filter_pushdown.rs +++ b/datafusion/physical-optimizer/src/filter_pushdown.rs @@ -39,11 +39,12 @@ use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{Result, assert_eq_or_internal_err, config::ConfigOptions}; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr_common::physical_expr::is_volatile; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use datafusion_physical_plan::filter_pushdown::{ ChildFilterPushdownResult, ChildPushdownResult, FilterPushdownPhase, FilterPushdownPropagation, PushedDown, }; -use datafusion_physical_plan::{ExecutionPlan, with_new_children_if_necessary}; use itertools::{Itertools, izip}; @@ -486,6 +487,14 @@ fn push_down_filters( // currently. `self_filters` are the predicates which are provided by the current node, // and tried to be pushed down over the child similarly. + assert_eq_or_internal_err!( + parent_filters.len(), + parent_filtered.len(), + "Filter pushdown expected {} to return one parent filter result per input filter for child {}", + node.name(), + child_idx + ); + // Filter out self_filters that contain volatile expressions and track indices let self_filtered = FilteredVec::new(&self_filters, allow_pushdown_for_expr); @@ -565,7 +574,7 @@ fn push_down_filters( } // Re-create this node with new children - let updated_node = with_new_children_if_necessary(Arc::clone(node), new_children)?; + let updated_node = replace_children_if_necessary(Arc::clone(node), new_children)?; // TODO: by calling `handle_child_pushdown_result` we are assuming that the // `ExecutionPlan` implementation will not change the plan itself. diff --git a/datafusion/physical-optimizer/src/hash_join_buffering.rs b/datafusion/physical-optimizer/src/hash_join_buffering.rs index 7a198cac13fc9..dbdfd34a9a01e 100644 --- a/datafusion/physical-optimizer/src/hash_join_buffering.rs +++ b/datafusion/physical-optimizer/src/hash_join_buffering.rs @@ -21,6 +21,7 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; use datafusion_physical_plan::ExecutionPlan; use datafusion_physical_plan::buffer::BufferExec; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use datafusion_physical_plan::joins::HashJoinExec; use std::sync::Arc; @@ -74,19 +75,25 @@ impl PhysicalOptimizerRule for HashJoinBuffering { if node.left.is::() { return Ok(Transformed::no(plan)); } - plan.with_new_children(vec![ - Arc::new(BufferExec::new(Arc::clone(&node.left), capacity)), - Arc::clone(&node.right), - ])? + replace_children_if_necessary( + plan, + vec![ + Arc::new(BufferExec::new(Arc::clone(&node.left), capacity)), + Arc::clone(&node.right), + ], + )? } else { // Do not stack BufferExec nodes together. if node.right.is::() { return Ok(Transformed::no(plan)); } - plan.with_new_children(vec![ - Arc::clone(&node.left), - Arc::new(BufferExec::new(Arc::clone(&node.right), capacity)), - ])? + replace_children_if_necessary( + plan, + vec![ + Arc::clone(&node.left), + Arc::new(BufferExec::new(Arc::clone(&node.right), capacity)), + ], + )? }, )) }) diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs index 74c6cbb19aea9..42736f8205089 100644 --- a/datafusion/physical-optimizer/src/join_selection.rs +++ b/datafusion/physical-optimizer/src/join_selection.rs @@ -40,6 +40,7 @@ use datafusion_physical_plan::joins::{ StreamJoinPartitionMode, SymmetricHashJoinExec, }; use datafusion_physical_plan::operator_statistics::StatisticsRegistry; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use std::sync::Arc; @@ -65,7 +66,7 @@ fn get_stats( reg.compute(plan) .map(|s| Arc::::clone(s.base_arc())) } else { - plan.partition_statistics(None) + StatisticsContext::new().compute(plan, &StatisticsArgs::new()) } } diff --git a/datafusion/physical-optimizer/src/lib.rs b/datafusion/physical-optimizer/src/lib.rs index 5fac8948b7f04..b9eb248f6e843 100644 --- a/datafusion/physical-optimizer/src/lib.rs +++ b/datafusion/physical-optimizer/src/lib.rs @@ -27,9 +27,12 @@ pub mod aggregate_statistics; pub mod combine_partial_final_agg; -pub mod enforce_distribution; -pub mod enforce_sorting; pub mod ensure_coop; +pub mod ensure_requirements; +// `enforce_distribution` and `enforce_sorting` are now internal implementation +// details of `ensure_requirements`. Re-export at the crate root so external test +// modules keep their public paths. +pub use ensure_requirements::{enforce_distribution, enforce_sorting}; pub mod filter_pushdown; pub mod join_selection; pub mod limit_pushdown; diff --git a/datafusion/physical-optimizer/src/limit_pushdown.rs b/datafusion/physical-optimizer/src/limit_pushdown.rs index 6164d86e5342a..f88a2be14e984 100644 --- a/datafusion/physical-optimizer/src/limit_pushdown.rs +++ b/datafusion/physical-optimizer/src/limit_pushdown.rs @@ -72,10 +72,12 @@ use datafusion_common::tree_node::{Transformed, TreeNodeRecursion}; use datafusion_common::utils::combine_limit; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::empty::EmptyExec; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; use datafusion_physical_plan::projection::ProjectionExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; /// This rule inspects [`ExecutionPlan`]'s and pushes down the fetch limit from /// the parent to the child if applicable. @@ -351,7 +353,9 @@ fn limit_eliminable_exact_num_rows( } if matches!( - current.partition_statistics(None)?.num_rows, + StatisticsContext::new() + .compute(current.as_ref(), &StatisticsArgs::new())? + .num_rows, Precision::Exact(0) ) { return Ok(Some(0)); @@ -375,6 +379,14 @@ pub(crate) fn pushdown_limits( (new_node, global_state) = pushdown_limit_helper(new_node.data, global_state)?; } + // Once a limit has been materialized above the current node, child + // subtrees should not inherit its `skip`. Keep `fetch`, but clear + // `skip` before recursing so child-local limits are not merged with + // an `OFFSET` that has already been applied. + if global_state.satisfied { + global_state.skip = 0; + } + // Apply pushdown limits in children let children = new_node.data.children(); let mut changed = false; @@ -392,7 +404,7 @@ pub(crate) fn pushdown_limits( .collect::>()?; if changed { - new_node.data.with_new_children(new_children) + replace_children_if_necessary(new_node.data, new_children) } else { Ok(new_node.data) } diff --git a/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs b/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs index 852dc2a2a9434..192a139f36021 100644 --- a/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs +++ b/datafusion/physical-optimizer/src/limited_distinct_aggregation.rs @@ -72,7 +72,8 @@ impl LimitedDistinctAggregation { if let Some(local_limit) = plan.downcast_ref::() { limit = local_limit.fetch(); children = local_limit.children().into_iter().cloned().collect(); - } else if let Some(global_limit) = plan.downcast_ref::() { + } else { + let global_limit = plan.downcast_ref::()?; global_fetch = global_limit.fetch(); global_fetch?; global_skip = global_limit.skip(); @@ -80,8 +81,6 @@ impl LimitedDistinctAggregation { limit = global_fetch.unwrap() + global_skip; children = global_limit.children().into_iter().cloned().collect(); is_global_limit = true - } else { - return None; } let child = children.iter().exactly_one().ok()?; // ensure there is no output ordering; can this rule be relaxed? diff --git a/datafusion/physical-optimizer/src/optimizer.rs b/datafusion/physical-optimizer/src/optimizer.rs index 05df642f8446b..aed25546cd09b 100644 --- a/datafusion/physical-optimizer/src/optimizer.rs +++ b/datafusion/physical-optimizer/src/optimizer.rs @@ -22,9 +22,8 @@ use std::sync::Arc; use crate::aggregate_statistics::AggregateStatistics; use crate::combine_partial_final_agg::CombinePartialFinalAggregate; -use crate::enforce_distribution::EnforceDistribution; -use crate::enforce_sorting::EnforceSorting; use crate::ensure_coop::EnsureCooperative; +use crate::ensure_requirements::EnsureRequirements; use crate::filter_pushdown::FilterPushdown; use crate::join_selection::JoinSelection; use crate::limit_pushdown::LimitPushdown; @@ -40,29 +39,10 @@ use crate::hash_join_buffering::HashJoinBuffering; use crate::limit_pushdown_past_window::LimitPushPastWindows; use crate::pushdown_sort::PushdownSort; use crate::window_topn::WindowTopN; -use datafusion_common::Result; use datafusion_common::config::ConfigOptions; -use datafusion_physical_plan::ExecutionPlan; -use datafusion_physical_plan::operator_statistics::StatisticsRegistry; -/// Context available to physical optimizer rules. -/// -/// This trait provides access to configuration options and optional statistics -/// registry for enhanced statistics lookup. It allows optimizer rules to access -/// extended context without changing the core [`PhysicalOptimizerRule::optimize`] -/// signature. -pub trait PhysicalOptimizerContext: Send + Sync { - /// Returns the configuration options. - fn config_options(&self) -> &ConfigOptions; - - /// Returns the statistics registry for enhanced statistics lookup. - /// - /// Returns `None` if no registry is configured, in which case rules - /// should fall back to using `ExecutionPlan::partition_statistics()`. - fn statistics_registry(&self) -> Option<&StatisticsRegistry> { - None - } -} +// Re-export from this module for backwards compatibility. +pub use datafusion_session::{PhysicalOptimizerContext, PhysicalOptimizerRule}; /// Simple context wrapping [`ConfigOptions`] for backward compatibility. /// @@ -86,47 +66,6 @@ impl PhysicalOptimizerContext for ConfigOnlyContext<'_> { } } -/// `PhysicalOptimizerRule` transforms one ['ExecutionPlan'] into another which -/// computes the same results, but in a potentially more efficient way. -/// -/// Use [`SessionState::add_physical_optimizer_rule`] to register additional -/// `PhysicalOptimizerRule`s. -/// -/// [`SessionState::add_physical_optimizer_rule`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html#method.add_physical_optimizer_rule -pub trait PhysicalOptimizerRule: Debug + std::any::Any { - /// Rewrite `plan` to an optimized form. - /// - /// This is the primary optimization method. For rules that need access to - /// the statistics registry, override [`optimize_with_context`](Self::optimize_with_context) instead. - fn optimize( - &self, - plan: Arc, - config: &ConfigOptions, - ) -> Result>; - - /// Rewrite `plan` with access to extended context (statistics registry, etc.). - /// - /// Override this method if you need access to the statistics registry for - /// enhanced statistics lookup. The default implementation simply calls - /// [`optimize`](Self::optimize) with the config options from the context. - fn optimize_with_context( - &self, - plan: Arc, - context: &dyn PhysicalOptimizerContext, - ) -> Result> { - self.optimize(plan, context.config_options()) - } - - /// A human readable name for this optimizer rule - fn name(&self) -> &str; - - /// A flag to indicate whether the physical planner should validate that the rule will not - /// change the schema of the plan after the rewriting. - /// Some of the optimization rules might change the nullable properties of the schema - /// and should disable the schema check. - fn schema_check(&self) -> bool; -} - /// A rule-based physical optimizer. #[derive(Clone, Debug)] pub struct PhysicalOptimizer { @@ -156,11 +95,11 @@ impl PhysicalOptimizer { Arc::new(AggregateStatistics::new()), // Statistics-based join selection will change the Auto mode to a real join implementation, // like collect left, or hash join, or future sort merge join, which will influence the - // EnforceDistribution and EnforceSorting rules as they decide whether to add additional - // repartitioning and local sorting steps to meet distribution and ordering requirements. - // Therefore, it should run before EnforceDistribution and EnforceSorting. + // EnsureRequirements rule as it decides whether to add additional repartitioning and + // local sorting steps to meet distribution and ordering requirements. Therefore, it + // should run before EnsureRequirements. Arc::new(JoinSelection::new()), - // The LimitedDistinctAggregation rule should be applied before the EnforceDistribution rule, + // The LimitedDistinctAggregation rule should be applied before EnsureRequirements, // as that rule may inject other operations in between the different AggregateExecs. // Applying the rule early means only directly-connected AggregateExecs must be examined. Arc::new(LimitedDistinctAggregation::new()), @@ -170,25 +109,35 @@ impl PhysicalOptimizer { // those are handled by the later `FilterPushdown` rule. // See `FilterPushdownPhase` for more details. Arc::new(FilterPushdown::new()), - // The EnforceDistribution rule is for adding essential repartitioning to satisfy distribution - // requirements. Please make sure that the whole plan tree is determined before this rule. - // This rule increases parallelism if doing so is beneficial to the physical plan; i.e. at - // least one of the operators in the plan benefits from increased parallelism. - Arc::new(EnforceDistribution::new()), - // The CombinePartialFinalAggregate rule should be applied after the EnforceDistribution rule + // WindowTopN: replaces Filter(rn<=K) → Window(ROW_NUMBER) + // with Window(ROW_NUMBER) → PartitionedTopKExec(fetch=K). + // Must run before EnsureRequirements (so it can rewrite against the + // window's declared ordering without pattern-matching a SortExec) + // and before ProjectionPushdown (which embeds projections into FilterExec). + Arc::new(WindowTopN::new()), + // Ensures each input plan satisfies the distribution and ordering + // requirements declared by `ExecutionPlan::required_input_distribution` + // and `ExecutionPlan::required_input_ordering`. + // + // If the requirements are already satisfied, this rule leaves the plan + // unchanged. For example, it does not add sorting when the input is a + // file scan whose existing order already satisfies the required ordering. + // Otherwise, this rule inserts the necessary repartitioning and sorting + // operators. + // + // This used to be implemented as two separate rules: `EnforceDistribution` + // and `EnforceSorting`. It is now a single idempotent rule that decides + // distribution and sorting together in one bottom-up pass, so the + // `pushdown_sorts` step no longer breaks distribution invariants set + // earlier in the pipeline. See the module-level doc on + // [`EnsureRequirements`](crate::ensure_requirements) for the per-phase + // breakdown, and + // for the original failure mode. + Arc::new(EnsureRequirements::new()), + // The CombinePartialFinalAggregate rule should be applied after distribution enforcement Arc::new(CombinePartialFinalAggregate::new()), - // The EnforceSorting rule is for adding essential local sorting to satisfy the required - // ordering. Please make sure that the whole plan tree is determined before this rule. - // Note that one should always run this rule after running the EnforceDistribution rule - // as the latter may break local sorting requirements. - Arc::new(EnforceSorting::new()), // Run once after the local sorting requirement is changed Arc::new(OptimizeAggregateOrder::new()), - // WindowTopN: replaces Filter(rn<=K) → Window(ROW_NUMBER) → Sort - // with Window(ROW_NUMBER) → PartitionedTopKExec(fetch=K). - // Must run after EnforceSorting (which inserts SortExec) and before - // ProjectionPushdown (which embeds projections into FilterExec). - Arc::new(WindowTopN::new()), // TODO: `try_embed_to_hash_join` in the ProjectionPushdown rule would be block by the CoalesceBatches, so add it before CoalesceBatches. Maybe optimize it in the future. Arc::new(ProjectionPushdown::new()), // Remove the ancillary output requirement operator since we are done with the planning @@ -201,7 +150,7 @@ impl PhysicalOptimizer { Arc::new(TopKAggregation::new()), // Tries to push limits down through window functions, growing as appropriate // This can possibly be combined with [LimitPushdown] - // It needs to come after [EnforceSorting] + // It needs to come after [EnsureRequirements] (which handles sort enforcement) Arc::new(LimitPushPastWindows::new()), // The HashJoinBuffering rule adds a BufferExec node with the configured capacity // in the prob side of hash joins. That way, the probe side gets eagerly polled before diff --git a/datafusion/physical-optimizer/src/output_requirements.rs b/datafusion/physical-optimizer/src/output_requirements.rs index 81df6f943c15e..541981270169e 100644 --- a/datafusion/physical-optimizer/src/output_requirements.rs +++ b/datafusion/physical-optimizer/src/output_requirements.rs @@ -30,19 +30,23 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, }; -use datafusion_common::{Result, Statistics}; +use datafusion_common::{Result, Statistics, internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::Distribution; use datafusion_physical_expr_common::sort_expr::OrderingRequirements; -use datafusion_physical_plan::execution_plan::Boundedness; +use datafusion_physical_plan::execution_plan::{ + Boundedness, replace_children_if_necessary, +}; use datafusion_physical_plan::projection::{ ProjectionExec, make_with_child, update_expr, update_ordering_requirement, }; +use datafusion_physical_plan::scalar_subquery::ScalarSubqueryExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - SendableRecordBatchStream, + ChildStats, ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, StatisticsArgs, }; /// This rule either adds or removes [`OutputRequirements`]s to/from the physical @@ -63,7 +67,9 @@ impl OutputRequirements { /// Create a new rule which works in `Add` mode; i.e. it simply adds a /// top-level [`OutputRequirementExec`] into the physical plan to keep track /// of global ordering and distribution requirements if there are any. - /// Note that this rule should run at the beginning. + /// Note that this rule should run at the beginning. It is idempotent: when + /// invoked on a plan that already contains an `OutputRequirementExec` (at + /// the root or below it), it returns the plan unchanged. pub fn new_add_mode() -> Self { Self { mode: RuleMode::Add, @@ -207,7 +213,15 @@ impl ExecutionPlan for OutputRequirementExec { } fn required_input_distribution(&self) -> Vec { - vec![self.dist_requirement.clone()] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements( + &self, + ) -> datafusion_physical_plan::InputDistributionRequirements { + datafusion_physical_plan::InputDistributionRequirements::new(vec![ + self.dist_requirement.clone(), + ]) } fn maintains_input_order(&self) -> Vec { @@ -222,9 +236,10 @@ impl ExecutionPlan for OutputRequirementExec { vec![self.order_requirement.clone()] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self::new( children.remove(0), // has a single child @@ -234,6 +249,16 @@ impl ExecutionPlan for OutputRequirementExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -242,10 +267,22 @@ impl ExecutionPlan for OutputRequirementExec { unreachable!(); } - fn partition_statistics(&self, partition: Option) -> Result> { - self.input.partition_statistics(partition) + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) } + #[expect( + deprecated, + reason = "HashPartitioned is accepted during the KeyPartitioned migration" + )] fn try_swapping_with_projection( &self, projection: &ProjectionExec, @@ -270,8 +307,12 @@ impl ExecutionPlan for OutputRequirementExec { requirements = OrderingRequirements::new_alternatives(updated_reqs, soft); } - let dist_req = match &self.required_input_distribution()[0] { - Distribution::HashPartitioned(exprs) => { + let input_distributions = self.input_distribution_requirements(); + let dist_req = match input_distributions.child_distribution(0) { + Some( + Distribution::HashPartitioned(exprs) + | Distribution::KeyPartitioned(exprs), + ) => { let mut updated_exprs = vec![]; for expr in exprs { let Some(new_expr) = update_expr(expr, projection.expr(), false)? @@ -280,9 +321,14 @@ impl ExecutionPlan for OutputRequirementExec { }; updated_exprs.push(new_expr); } - Distribution::HashPartitioned(updated_exprs) + Distribution::KeyPartitioned(updated_exprs) + } + Some(dist) => dist.clone(), + None => { + return internal_err!( + "OutputRequirementExec missing input distribution requirement" + ); } - dist => dist.clone(), }; make_with_child(projection, &self.input()).map(|input| { @@ -297,32 +343,11 @@ impl ExecutionPlan for OutputRequirementExec { fn apply_expressions( &self, - f: &mut dyn FnMut( - &dyn datafusion_physical_expr_common::physical_expr::PhysicalExpr, + _f: &mut dyn FnMut( + &Arc, ) -> Result, ) -> Result { - // Visit expressions in order_requirement - let mut tnr = TreeNodeRecursion::Continue; - if let Some(order_reqs) = &self.order_requirement { - let lexes = match order_reqs { - OrderingRequirements::Hard(alternatives) => alternatives, - OrderingRequirements::Soft(alternatives) => alternatives, - }; - for lex in lexes { - for sort_expr in lex { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - } - - // Visit expressions in dist_requirement if it's HashPartitioned - if let Distribution::HashPartitioned(exprs) = &self.dist_requirement { - for expr in exprs { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - } - - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } } @@ -357,7 +382,15 @@ impl PhysicalOptimizerRule for OutputRequirements { /// This functions adds ancillary `OutputRequirementExec` to the physical plan, so that /// global requirements are not lost during optimization. +/// +/// Idempotent: re-running this rule (as adaptive execution in datafusion-ballista +/// AQE does after every completed stage, see datafusion-ballista#1359) does not +/// stack wrappers, whether the previously-added `OutputRequirementExec` sits at +/// the root (handled here) or below it (handled in `require_top_ordering_helper`). fn require_top_ordering(plan: Arc) -> Result> { + if plan.downcast_ref::().is_some() { + return Ok(plan); + } let (new_plan, is_changed) = require_top_ordering_helper(plan)?; if is_changed { Ok(new_plan) @@ -373,21 +406,43 @@ fn require_top_ordering(plan: Arc) -> Result Option { + if plan.children().len() == 1 { + Some(0) + } else if plan.downcast_ref::().is_some() { + // `ScalarSubqueryExec` is multi-child but order-transparent on child 0 + // (the main input); its other children are subquery plans that don't + // affect output ordering, so descend into child 0. Without this the + // search stops here and loses the query's global ORDER BY. + Some(0) + } else { + None + } +} + /// Helper function that adds an ancillary `OutputRequirementExec` to the given plan. /// First entry in the tuple is resulting plan, second entry indicates whether any /// `OutputRequirementExec` is added to the plan. fn require_top_ordering_helper( plan: Arc, ) -> Result<(Arc, bool)> { - let mut children = plan.children(); + // A previous run of this rule already captured the ordering requirement at + // this node. Report it as already handled. + if plan.downcast_ref::().is_some() { + return Ok((plan, true)); + } + // Global ordering defines desired ordering in the final result. - if children.len() != 1 { - Ok((plan, false)) - } else if let Some(sort_exec) = plan.downcast_ref::() { + if let Some(sort_exec) = plan.downcast_ref::() { // In case of constant columns, output ordering of the `SortExec` would // be an empty set. Therefore; we check the sort expression field to // assign the requirements. - let req_dist = sort_exec.required_input_distribution().swap_remove(0); + let req_dist = sort_exec + .input_distribution_requirements() + .into_per_child() + .swap_remove(0); let req_ordering = sort_exec.expr(); let reqs = OrderingRequirements::from(req_ordering.clone()); let fetch = sort_exec.fetch(); @@ -413,25 +468,27 @@ fn require_top_ordering_helper( )) as _, true, )) - } else if plan.maintains_input_order()[0] - && (plan.required_input_ordering()[0] - .as_ref() - .is_none_or(|o| matches!(o, OrderingRequirements::Soft(_)))) - { - // Keep searching for a `SortExec` as long as ordering is maintained, - // and on-the-way operators do not themselves require an ordering. - // When an operator requires an ordering, any `SortExec` below can not - // be responsible for (i.e. the originator of) the global ordering. - let (new_child, is_changed) = - require_top_ordering_helper(Arc::clone(children.swap_remove(0)))?; - - let plan = if is_changed { - plan.with_new_children(vec![new_child])? - } else { - plan - }; - - Ok((plan, is_changed)) + } else if let Some(idx) = output_requirement_child(plan.as_ref()) { + // Keep searching for a `SortExec` / `SortPreservingMergeExec` as long as + // ordering is maintained, and on-the-way operators do not themselves + // require an ordering. When an operator requires an ordering, any + // `SortExec` below can not be responsible for (i.e. the originator of) + // the global ordering. + if plan.maintains_input_order()[idx] + && plan.required_input_ordering()[idx] + .as_ref() + .is_none_or(|o| matches!(o, OrderingRequirements::Soft(_))) + { + let mut children: Vec> = + plan.children().into_iter().map(Arc::clone).collect(); + let (new_child, is_changed) = + require_top_ordering_helper(Arc::clone(&children[idx]))?; + if is_changed { + children[idx] = new_child; + return Ok((replace_children_if_necessary(plan, children)?, true)); + } + } + Ok((plan, false)) } else { // Stop searching, there is no global ordering desired for the query. Ok((plan, false)) diff --git a/datafusion/physical-optimizer/src/pushdown_sort.rs b/datafusion/physical-optimizer/src/pushdown_sort.rs index 40a6fe2c205c7..5dfe221ed24c0 100644 --- a/datafusion/physical-optimizer/src/pushdown_sort.rs +++ b/datafusion/physical-optimizer/src/pushdown_sort.rs @@ -57,13 +57,15 @@ use crate::PhysicalOptimizerRule; use datafusion_common::Result; use datafusion_common::config::ConfigOptions; -use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_physical_plan::ExecutionPlan; +use datafusion_common::tree_node::{ + Transformed, TransformedResult, TreeNode, TreeNodeRecursion, +}; use datafusion_physical_plan::SortOrderPushdownResult; use datafusion_physical_plan::buffer::BufferExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion_physical_plan::{ExecutionPlan, ExecutionPlanProperties}; use std::sync::Arc; /// A PhysicalOptimizerRule that attempts to push down sort requirements to data sources. @@ -133,7 +135,15 @@ impl PhysicalOptimizerRule for PushdownSort { Arc::new(new_sort), ) .with_fetch(spm.fetch()); - return Ok(Transformed::yes(Arc::new(new_spm))); + // The replacement already has the required + // `SortPreservingMergeExec` parent. Do not descend + // into its `SortExec` child and treat it as a + // standalone TopK. + return Ok(Transformed::new( + Arc::new(new_spm), + true, + TreeNodeRecursion::Jump, + )); } SortOrderPushdownResult::Unsupported => { return Ok(Transformed::no(plan)); @@ -172,13 +182,35 @@ impl PhysicalOptimizerRule for PushdownSort { // Data source is optimized for the ordering but not perfectly sorted // Keep the Sort operator but use the optimized input // Benefits: TopK queries can terminate early, better cache locality - Ok(Transformed::yes(Arc::new( + // A standalone multi-partition TopK still needs a global + // merge; otherwise a later coalesce can concatenate + // locally sorted partitions. + let preserve_partitioning = sort_exec.preserve_partitioning(); + let needs_global_topk = + preserve_partitioning && sort_exec.fetch().is_some(); + let input_partitions = inner.output_partitioning().partition_count(); + let new_sort: Arc = Arc::new( SortExec::new(required_ordering.clone(), inner) .with_fetch(sort_exec.fetch()) - .with_preserve_partitioning( - sort_exec.preserve_partitioning(), - ), - ))) + .with_preserve_partitioning(preserve_partitioning), + ); + if needs_global_topk && input_partitions > 1 { + let new_spm = SortPreservingMergeExec::new( + required_ordering.clone(), + new_sort, + ) + .with_fetch(sort_exec.fetch()); + // Do not descend into the newly inserted + // `SortExec`, or this standalone branch will wrap it + // in another `SortPreservingMergeExec`. + Ok(Transformed::new( + Arc::new(new_spm), + true, + TreeNodeRecursion::Jump, + )) + } else { + Ok(Transformed::yes(new_sort)) + } } SortOrderPushdownResult::Unsupported => { // Cannot optimize for this ordering - no change diff --git a/datafusion/physical-optimizer/src/sanity_checker.rs b/datafusion/physical-optimizer/src/sanity_checker.rs index 40c6245d894d4..713213b70612d 100644 --- a/datafusion/physical-optimizer/src/sanity_checker.rs +++ b/datafusion/physical-optimizer/src/sanity_checker.rs @@ -30,9 +30,13 @@ use datafusion_common::config::{ConfigOptions, OptimizerOptions}; use datafusion_common::plan_err; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_physical_expr::intervals::utils::{check_support, is_datatype_supported}; -use datafusion_physical_plan::execution_plan::{Boundedness, EmissionType}; +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, InvariantLevel, +}; use datafusion_physical_plan::joins::SymmetricHashJoinExec; -use datafusion_physical_plan::{ExecutionPlanProperties, get_plan_string}; +use datafusion_physical_plan::{ + ChildSatisfactionOptions, ExecutionPlanProperties, get_plan_string, +}; use crate::PhysicalOptimizerRule; use datafusion_physical_expr_common::sort_expr::format_physical_sort_requirement_list; @@ -141,11 +145,12 @@ pub fn check_plan_sanity( optimizer_options: &OptimizerOptions, ) -> Result<()> { check_finiteness_requirements(plan.as_ref(), optimizer_options)?; + let input_distributions = plan.input_distribution_requirements(); for ((idx, child), sort_req, dist_req) in izip!( plan.children().into_iter().enumerate(), plan.required_input_ordering(), - plan.required_input_distribution(), + input_distributions.per_child_distributions(), ) { let child_eq_props = child.equivalence_properties(); if let Some(sort_req) = sort_req { @@ -162,9 +167,12 @@ pub fn check_plan_sanity( } } - if !child - .output_partitioning() - .satisfaction(&dist_req, child_eq_props, true) + if !input_distributions + .child_satisfaction( + idx, + child.as_ref(), + ChildSatisfactionOptions::new().with_allow_subset(true), + )? .is_satisfied() { let plan_str = get_plan_string(plan); @@ -178,6 +186,8 @@ pub fn check_plan_sanity( } } + plan.check_invariants(InvariantLevel::Executable)?; + Ok(()) } diff --git a/datafusion/physical-optimizer/src/topk_aggregation.rs b/datafusion/physical-optimizer/src/topk_aggregation.rs index e1779c04a6a92..0eddb5d5507e4 100644 --- a/datafusion/physical-optimizer/src/topk_aggregation.rs +++ b/datafusion/physical-optimizer/src/topk_aggregation.rs @@ -46,6 +46,7 @@ impl TopKAggregation { aggr: &AggregateExec, order_by: &str, order_desc: bool, + nulls_first: bool, limit: usize, ) -> Option> { // Current only support single group key @@ -66,6 +67,26 @@ impl TopKAggregation { // Check if this is ordering by an aggregate function (MIN/MAX) if let Some((field, desc)) = aggr.get_minmax_desc() { + // A nullable MIN/MAX starts as NULL and becomes non-NULL when the + // group sees its first value. With NULLS FIRST that transition + // worsens the group's rank, so a bounded aggregation cannot safely + // discard other NULL groups. Use regular aggregation for exact + // results. Non-nullable inputs never take this transition and can + // still use TopK. + let input_nullable = aggr + .aggr_expr() + .iter() + .exactly_one() + .ok()? + .expressions() + .into_iter() + .exactly_one() + .ok()? + .nullable(aggr.input_schema.as_ref()) + .ok()?; + if nulls_first && input_nullable { + return None; + } // ensure the sort direction matches aggregate function if desc != order_desc { return None; @@ -100,6 +121,7 @@ impl TopKAggregation { let order = sort.properties().output_ordering()?; let order = order.iter().exactly_one().ok()?; let order_desc = order.options.descending; + let nulls_first = order.options.nulls_first; let order = order.expr.downcast_ref::()?; let mut cur_col_name = order.name().to_string(); let limit = sort.fetch()?; @@ -111,7 +133,13 @@ impl TopKAggregation { } if let Some(aggr) = plan.downcast_ref::() { // either we run into an Aggregate and transform it - match Self::transform_agg(aggr, &cur_col_name, order_desc, limit) { + match Self::transform_agg( + aggr, + &cur_col_name, + order_desc, + nulls_first, + limit, + ) { None => cardinality_preserved = false, Some(plan) => return Ok(Transformed::yes(plan)), } diff --git a/datafusion/physical-optimizer/src/topk_repartition.rs b/datafusion/physical-optimizer/src/topk_repartition.rs index 115bdc3cb535f..d8fa1ac986f90 100644 --- a/datafusion/physical-optimizer/src/topk_repartition.rs +++ b/datafusion/physical-optimizer/src/topk_repartition.rs @@ -48,6 +48,7 @@ use crate::PhysicalOptimizerRule; use datafusion_common::Result; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use std::sync::Arc; // CoalesceBatchesExec is deprecated on main (replaced by arrow-rs BatchCoalescer), // but older DataFusion versions may still insert it between SortExec and RepartitionExec. @@ -151,7 +152,7 @@ impl PhysicalOptimizerRule for TopKRepartition { // Rebuild the tree above the repartition let new_sort_input = if let Some(parent) = repart_parent { - parent.with_new_children(vec![new_repartition])? + replace_children_if_necessary(parent, vec![new_repartition])? } else { new_repartition }; diff --git a/datafusion/physical-optimizer/src/utils.rs b/datafusion/physical-optimizer/src/utils.rs index a6b01637c970e..04229e1cc2737 100644 --- a/datafusion/physical-optimizer/src/utils.rs +++ b/datafusion/physical-optimizer/src/utils.rs @@ -18,7 +18,7 @@ use std::sync::Arc; use datafusion_common::Result; -use datafusion_physical_expr::{LexOrdering, LexRequirement}; +use datafusion_physical_expr::{Distribution, LexOrdering, LexRequirement}; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::repartition::RepartitionExec; @@ -58,6 +58,56 @@ pub fn add_sort_above( PlanContext::new(Arc::new(new_sort), T::default(), vec![node]) } +/// Like [`add_sort_above`], but also inserts a [`SortPreservingMergeExec`] when +/// the parent distribution requires a single partition and the input has +/// multiple partitions. This prevents `SortExec(preserve_partitioning=true)` +/// from violating `SinglePartition` requirements. +pub fn add_sort_above_with_distribution( + node: PlanContext, + sort_requirements: LexRequirement, + fetch: Option, + required_distribution: &Distribution, +) -> PlanContext { + let mut sort_reqs: Vec<_> = sort_requirements.into(); + sort_reqs.retain(|sort_expr| { + node.plan + .equivalence_properties() + .is_expr_constant(&sort_expr.expr) + .is_none() + }); + let sort_exprs = sort_reqs.into_iter().map(Into::into).collect::>(); + let Some(ordering) = LexOrdering::new(sort_exprs) else { + return node; + }; + let input_has_multiple_partitions = + node.plan.output_partitioning().partition_count() > 1; + + let mut new_sort = + SortExec::new(ordering.clone(), Arc::clone(&node.plan)).with_fetch(fetch); + if input_has_multiple_partitions { + new_sort = new_sort.with_preserve_partitioning(true); + } + + let sort_node = PlanContext::new(Arc::new(new_sort), T::default(), vec![node]); + + // If the parent requires SinglePartition and the input has multiple partitions, + // wrap the partition-preserving sort in SortPreservingMergeExec. + if matches!(required_distribution, Distribution::SinglePartition) + && input_has_multiple_partitions + { + PlanContext::new( + Arc::new( + SortPreservingMergeExec::new(ordering, Arc::clone(&sort_node.plan)) + .with_fetch(fetch), + ), + T::default(), + vec![sort_node], + ) + } else { + sort_node + } +} + /// This utility function adds a `SortExec` above an operator according to the /// given ordering requirements while preserving the original partitioning. If /// requirement is already satisfied no `SortExec` is added. diff --git a/datafusion/physical-optimizer/src/window_topn.rs b/datafusion/physical-optimizer/src/window_topn.rs index 40dbddfbdf9fb..20bd8b0d38a1c 100644 --- a/datafusion/physical-optimizer/src/window_topn.rs +++ b/datafusion/physical-optimizer/src/window_topn.rs @@ -26,12 +26,27 @@ //! ) WHERE rn <= K; //! ``` //! -//! And replaces the `FilterExec → BoundedWindowAggExec → SortExec` pipeline -//! with `BoundedWindowAggExec → PartitionedTopKExec(fetch=K)`, removing both -//! the `FilterExec` and `SortExec`. +//! or with `RANK()` in place of `ROW_NUMBER()`: //! -//! See [`PartitionedTopKExec`] -//! for details on the replacement operator. +//! ```sql +//! SELECT * FROM ( +//! SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk +//! FROM t +//! ) WHERE rk <= K; +//! ``` +//! +//! And replaces the `FilterExec → BoundedWindowAggExec` pipeline with +//! `BoundedWindowAggExec → PartitionedTopKExec(fetch=K)`, removing the +//! `FilterExec` and inserting `PartitionedTopKExec` under the window. +//! +//! The appropriate [`WindowFnKind`] is forwarded to `PartitionedTopKExec`. +//! RANK requires a non-empty `ORDER BY` clause (otherwise all rows tie at +//! rank 1 and the optimization is degenerate). +//! +//! See [`PartitionedTopKExec`] for details on the replacement operator. +//! +//! [`PartitionedTopKExec`]: datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec +//! [`WindowFnKind`]: datafusion_physical_plan::sorts::partitioned_topk::WindowFnKind use std::sync::Arc; @@ -43,36 +58,40 @@ use datafusion_common::{Result, ScalarValue}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal}; use datafusion_physical_expr::window::StandardWindowExpr; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::execution_plan::replace_children_if_necessary; use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::projection::ProjectionExec; -use datafusion_physical_plan::sorts::partitioned_topk::PartitionedTopKExec; -use datafusion_physical_plan::sorts::sort::SortExec; +use datafusion_physical_plan::repartition::RepartitionExec; +use datafusion_physical_plan::sorts::partitioned_topk::{ + PartitionedTopKExec, WindowFnKind, +}; use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; -/// Physical optimizer rule that converts per-partition `ROW_NUMBER` top-K -/// queries into a more efficient plan using [`PartitionedTopKExec`]. +/// Physical optimizer rule that converts per-partition `ROW_NUMBER` and +/// `RANK` top-K queries into a more efficient plan using +/// [`PartitionedTopKExec`]. /// /// # Pattern Detected /// /// ```text -/// FilterExec(rn <= K) +/// FilterExec( <= K) /// [optional ProjectionExec] -/// BoundedWindowAggExec(ROW_NUMBER PARTITION BY ... ORDER BY ...) -/// SortExec(partition_keys, order_keys) +/// BoundedWindowAggExec( PARTITION BY ... ORDER BY ...) /// ``` /// /// # Replacement /// /// ```text /// [optional ProjectionExec] -/// BoundedWindowAggExec(ROW_NUMBER PARTITION BY ... ORDER BY ...) -/// PartitionedTopKExec(partition_keys, order_keys, fetch=K) +/// BoundedWindowAggExec( PARTITION BY ... ORDER BY ...) +/// PartitionedTopKExec(fn=, partition_keys, order_keys, fetch=K) /// ``` /// -/// The `FilterExec` is removed entirely (all output rows have `rn ∈ {1..K}`). -/// The `SortExec` is replaced by `PartitionedTopKExec` which maintains a -/// per-partition top-K heap instead of sorting the entire dataset. +/// The `FilterExec` is removed entirely. The child of `BoundedWindowAggExec` is now +/// `PartitionedTopKExec`, which maintains a per-partition top-K heap (and, +/// for `RANK`, a sibling ties `Vec`) instead of sorting the whole dataset. /// /// # Supported Predicates /// @@ -85,10 +104,13 @@ use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowUDFExpr}; /// /// All of the following must be true: /// - Config flag `enable_window_topn` is `true` -/// - The plan matches `FilterExec → [ProjectionExec] → BoundedWindowAggExec → SortExec` -/// - The window function is `ROW_NUMBER` (not `RANK`, `DENSE_RANK`, etc.) -/// - `ROW_NUMBER` has a `PARTITION BY` clause (global top-K is already -/// handled by `SortExec` with `fetch`) +/// - The plan matches `FilterExec → [ProjectionExec] → BoundedWindowAggExec` +/// - The window function is `ROW_NUMBER` or `RANK` (not `DENSE_RANK`) +/// - The window function has a `PARTITION BY` clause (global top-K is +/// already handled by `SortExec` with `fetch`) +/// - For `RANK`: a non-empty `ORDER BY` clause (otherwise all rows tie +/// at rank 1 — the optimization is useless and the boundary-tie storage +/// would be unbounded) /// - The filter predicate compares the window output column to an integer /// literal using `<=`, `<`, `>=`, or `>` /// @@ -104,7 +126,7 @@ impl WindowTopN { /// Attempt to transform a single plan node. /// /// Returns `Some(new_plan)` if the node matches the - /// `FilterExec → [ProjectionExec] → BoundedWindowAggExec → SortExec` + /// `FilterExec → [ProjectionExec] → BoundedWindowAggExec` /// pattern and can be rewritten, or `None` if the node should be /// left unchanged. fn try_transform(plan: &Arc) -> Option> { @@ -119,29 +141,24 @@ impl WindowTopN { // Step 2: Extract limit from predicate (rn <= K, rn < K, etc.) let (col_idx, limit_n) = extract_window_limit(filter.predicate())?; - // Step 3: Walk through optional ProjectionExec to find BoundedWindowAggExec + // Step 3: Walk through optional ProjectionExec and RepartitionExec to find BoundedWindowAggExec let child = filter.input(); - let (window_exec, proj_between) = find_window_below(child)?; + let (window_exec, intermediates) = find_window_below(child)?; - // Step 4: Verify col_idx references a ROW_NUMBER window output column - let input_field_count = window_exec.input().schema().fields().len(); + // Step 4: Verify col_idx references a supported window function output column + let window_exec_typed = window_exec.downcast_ref::()?; + let input_field_count = window_exec_typed.input().schema().fields().len(); if col_idx < input_field_count { return None; // Filter is on an input column, not a window column } let window_expr_idx = col_idx - input_field_count; - let window_exprs = window_exec.window_expr(); + let window_exprs = window_exec_typed.window_expr(); if window_expr_idx >= window_exprs.len() { return None; } - if !is_row_number(&window_exprs[window_expr_idx]) { - return None; - } + let fn_kind = supported_window_fn(&window_exprs[window_expr_idx])?; - // Step 5: Verify child of window is SortExec - let sort_exec = window_exec.input().downcast_ref::()?; - let sort_child = sort_exec.input(); - - // Step 6: Determine partition_prefix_len from the window expression + // Step 5: Validate PARTITION BY / ORDER BY and collect sort keys from the window expr let partition_by = window_exprs[window_expr_idx].partition_by(); let partition_prefix_len = partition_by.len(); @@ -151,38 +168,44 @@ impl WindowTopN { return None; } - // Step 7: Build PartitionedTopKExec using SortExec's expressions + // For RANK: an empty ORDER BY makes every row tie at rank 1 — + // the optimization is degenerate (we'd retain the entire input) + // and tie storage would be unbounded. + let order_by = window_exprs[window_expr_idx].order_by(); + if matches!(fn_kind, WindowFnKind::Rank) && order_by.is_empty() { + return None; + } + + // Step 6: Build PartitionedTopKExec from the window's partition/order keys + let expr_iterator = partition_by + .iter() + .map(|e| PhysicalSortExpr::new_default(Arc::clone(e))) + .chain(order_by.iter().cloned()); + let expr = LexOrdering::new(expr_iterator)?; + let partitioned_topk = PartitionedTopKExec::try_new( - Arc::clone(sort_child), - sort_exec.expr().clone(), + Arc::clone(window_exec_typed.input()), + expr, partition_prefix_len, limit_n, + fn_kind, ) .ok()?; - // Step 8: Rebuild window with new child - let new_window = Arc::clone(&child_as_arc(window_exec)) - .with_new_children(vec![Arc::new(partitioned_topk)]) - .ok()?; + // Step 7: Rebuild window with PartitionedTopKExec as its child + let mut result = + replace_children_if_necessary(window_exec, vec![Arc::new(partitioned_topk)]) + .ok()?; - // Step 9: If ProjectionExec was between Filter and Window, rebuild it - let result = match proj_between { - Some(proj) => Arc::clone(&child_as_arc(proj)) - .with_new_children(vec![new_window]) - .ok()?, - None => new_window, - }; + // Step 8: Rebuild intermediate nodes (ProjectionExec/RepartitionExec) + for node in intermediates.into_iter().rev() { + result = replace_children_if_necessary(node, vec![result]).ok()?; + } Some(result) } } -/// Helper to get an `Arc` from a reference. -/// We need this because `with_new_children` takes `Arc`. -fn child_as_arc(plan: &T) -> Arc { - Arc::new(plan.clone()) -} - impl PhysicalOptimizerRule for WindowTopN { fn optimize( &self, @@ -287,45 +310,52 @@ fn scalar_to_usize(value: &ScalarValue) -> Option { } } -/// Check if a window expression is `ROW_NUMBER`. +/// Identify which supported ranking window function `expr` is. /// /// Downcasts through `StandardWindowExpr` → `WindowUDFExpr` and checks -/// that the UDF name is `"row_number"`. Returns `false` for all other -/// window functions (e.g., `RANK`, `DENSE_RANK`, `SUM`). -fn is_row_number(expr: &Arc) -> bool { - let Some(swe) = expr.as_any().downcast_ref::() else { - return false; - }; +/// the UDF name. Returns: +/// - `Some(WindowFnKind::RowNumber)` for `"row_number"` +/// - `Some(WindowFnKind::Rank)` for `"rank"` +/// - `None` for everything else (e.g. `dense_rank`) +fn supported_window_fn( + expr: &Arc, +) -> Option { + let swe = expr.as_any().downcast_ref::()?; let swfe = swe.get_standard_func_expr(); - let Some(udf) = swfe.as_any().downcast_ref::() else { - return false; - }; - udf.fun().name() == "row_number" + let udf = swfe.as_any().downcast_ref::()?; + match udf.fun().name() { + "row_number" => Some(WindowFnKind::RowNumber), + "rank" => Some(WindowFnKind::Rank), + _ => None, + } } +type PlanAndIntermediates = (Arc, Vec>); + /// Walk below a plan node looking for a [`BoundedWindowAggExec`]. /// -/// Handles two cases: -/// - Direct child: `FilterExec → BoundedWindowAggExec` -/// - With projection: `FilterExec → ProjectionExec → BoundedWindowAggExec` +/// Handles sequences of `ProjectionExec` and `RepartitionExec`. +/// This is safe because `PartitionedTopKExec` can be pushed below them: +/// projections only provide aliases, and pushing the limit below repartitions +/// is safe because the limit is computed per-partition. /// -/// Returns the window exec and an optional `ProjectionExec` in between, -/// or `None` if no `BoundedWindowAggExec` is found within one or two levels. -fn find_window_below( - plan: &Arc, -) -> Option<(&BoundedWindowAggExec, Option<&ProjectionExec>)> { - // Direct child is BoundedWindowAggExec - if let Some(window) = plan.downcast_ref::() { - return Some((window, None)); - } +/// Returns the window exec and a list of intermediate nodes to rebuild, +/// or `None` if no `BoundedWindowAggExec` is found. +fn find_window_below(plan: &Arc) -> Option { + let mut current = Arc::clone(plan); + let mut intermediates = Vec::new(); - // Child is ProjectionExec with BoundedWindowAggExec below - if let Some(proj) = plan.downcast_ref::() { - let proj_child = proj.input(); - if let Some(window) = proj_child.downcast_ref::() { - return Some((window, Some(proj))); + loop { + if current.downcast_ref::().is_some() { + return Some((current, intermediates)); + } else if current.downcast_ref::().is_some() + || current.downcast_ref::().is_some() + { + let next = Arc::clone(current.children().first()?); + intermediates.push(current); + current = next; + } else { + return None; } } - - None } diff --git a/datafusion/physical-plan/Cargo.toml b/datafusion/physical-plan/Cargo.toml index c6710262776c7..0f72b74840d01 100644 --- a/datafusion/physical-plan/Cargo.toml +++ b/datafusion/physical-plan/Cargo.toml @@ -42,6 +42,15 @@ force_hash_collisions = [] test_utils = ["arrow/test_utils"] tokio_coop = [] tokio_coop_fallback = [] +# Enables `PhysicalExpr::try_to_proto` / `try_from_proto` hooks on the +# physical expressions defined in this crate (e.g. `HashExpr`). Off by +# default so consumers that never serialize plans pay nothing. +proto = [ + "dep:datafusion-proto-models", + "dep:datafusion-proto-common", + "datafusion-physical-expr/proto", + "datafusion-physical-expr-common/proto", +] [lib] name = "datafusion_physical_plan" @@ -56,6 +65,7 @@ arrow-ipc = { workspace = true, features = ["lz4", "zstd"] } arrow-ord = { workspace = true } arrow-schema = { workspace = true } async-trait = { workspace = true } +bytes = { workspace = true } datafusion-common = { workspace = true } datafusion-common-runtime = { workspace = true, default-features = true } datafusion-execution = { workspace = true } @@ -65,6 +75,8 @@ datafusion-functions-aggregate-common = { workspace = true } datafusion-functions-window-common = { workspace = true } datafusion-physical-expr = { workspace = true, default-features = true } datafusion-physical-expr-common = { workspace = true } +datafusion-proto-common = { workspace = true, optional = true } +datafusion-proto-models = { workspace = true, optional = true } futures = { workspace = true } half = { workspace = true } hashbrown = { workspace = true } @@ -73,7 +85,8 @@ itertools = { workspace = true, features = ["use_std"] } log = { workspace = true } num-traits = { workspace = true } parking_lot = { workspace = true } -pin-project-lite = "^0.2.7" +pin-project-lite = { workspace = true } +serde_json = { workspace = true, features = ["preserve_order"] } tokio = { workspace = true } [dev-dependencies] @@ -113,6 +126,24 @@ harness = false name = "aggregate_vectorized" required-features = ["test_utils"] +[[bench]] +harness = false +name = "compute_statistics" + [[bench]] harness = false name = "dictionary_group_values" + +[[bench]] +harness = false +name = "hash_join_semi_anti" +required-features = ["test_utils"] + +[[bench]] +harness = false +name = "multi_group_by" +required-features = ["test_utils"] + +[[bench]] +harness = false +name = "bounded_window" diff --git a/datafusion/physical-plan/benches/aggregate_vectorized.rs b/datafusion/physical-plan/benches/aggregate_vectorized.rs index 48ca76d80d2d3..488647d5f8315 100644 --- a/datafusion/physical-plan/benches/aggregate_vectorized.rs +++ b/datafusion/physical-plan/benches/aggregate_vectorized.rs @@ -21,7 +21,6 @@ use arrow::util::bench_util::{ create_primitive_array, create_string_view_array_with_len, create_string_view_array_with_max_len, }; -use arrow::util::test_util::seedable_rng; use arrow_schema::DataType; use criterion::measurement::WallTime; use criterion::{ @@ -30,7 +29,9 @@ use criterion::{ use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupColumn; use datafusion_physical_plan::aggregates::group_values::multi_group_by::bytes_view::ByteViewGroupValueBuilder; use datafusion_physical_plan::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder; +use rand::SeedableRng; use rand::distr::{Bernoulli, Distribution}; +use rand::rngs::StdRng; use std::hint::black_box; use std::sync::Arc; @@ -128,7 +129,7 @@ fn bytes_bench( input, "0.75 true", { - let mut rng = seedable_rng(); + let mut rng = StdRng::seed_from_u64(42); let d = Bernoulli::new(0.75).unwrap(); (0..size).map(|_| d.sample(&mut rng)).collect::>() }, @@ -141,7 +142,7 @@ fn bytes_bench( input, "0.5 true", { - let mut rng = seedable_rng(); + let mut rng = StdRng::seed_from_u64(42); let d = Bernoulli::new(0.5).unwrap(); (0..size).map(|_| d.sample(&mut rng)).collect::>() }, @@ -154,7 +155,7 @@ fn bytes_bench( input, "0.25 true", { - let mut rng = seedable_rng(); + let mut rng = StdRng::seed_from_u64(42); let d = Bernoulli::new(0.25).unwrap(); (0..size).map(|_| d.sample(&mut rng)).collect::>() }, @@ -236,7 +237,7 @@ fn bench_single_primitive( &input, "0.75 true", { - let mut rng = seedable_rng(); + let mut rng = StdRng::seed_from_u64(42); let d = Bernoulli::new(0.75).unwrap(); (0..size).map(|_| d.sample(&mut rng)).collect::>() }, @@ -249,7 +250,7 @@ fn bench_single_primitive( &input, "0.5 true", { - let mut rng = seedable_rng(); + let mut rng = StdRng::seed_from_u64(42); let d = Bernoulli::new(0.5).unwrap(); (0..size).map(|_| d.sample(&mut rng)).collect::>() }, @@ -262,7 +263,7 @@ fn bench_single_primitive( &input, "0.25 true", { - let mut rng = seedable_rng(); + let mut rng = StdRng::seed_from_u64(42); let d = Bernoulli::new(0.25).unwrap(); (0..size).map(|_| d.sample(&mut rng)).collect::>() }, diff --git a/datafusion/physical-plan/benches/bounded_window.rs b/datafusion/physical-plan/benches/bounded_window.rs new file mode 100644 index 0000000000000..56e195afbd4f2 --- /dev/null +++ b/datafusion/physical-plan/benches/bounded_window.rs @@ -0,0 +1,280 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for `BoundedWindowAggExec` with many partitions. +//! +//! The streaming window operator keeps per-partition state keyed by +//! `PartitionKey` (`Vec`) and, in `Linear` mode (input sorted +//! by the ORDER BY column but not by the partition columns), visits every +//! live partition on every batch while never retiring partitions until the +//! input is exhausted. The cases here stress that path in different ways: +//! +//! - `linear N partitions`: dense round-robin keys -- every partition +//! receives rows in every batch, so per-visit fixed costs dominate. +//! - `linear sparse N partitions`: keys are clustered in time, so each +//! batch touches only a small, fresh subset of keys while the set of live +//! partitions keeps growing -- per-batch work on quiet partitions +//! dominates. +//! - `linear rows N partitions`: the dense layout with a ROWS frame, whose +//! results can only be finalized as more rows of the same partition +//! arrive. +//! - `linear multi N partitions`: two window expressions over the dense +//! layout, doubling the per-partition evaluation sweeps. +//! - `sorted N partitions`: control; input sorted by partition key, so +//! finished partitions are pruned eagerly and the state maps stay small. + +use std::sync::Arc; + +use arrow::array::UInt64Array; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use criterion::{Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_execution::TaskContext; +use datafusion_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, +}; +use datafusion_functions_aggregate::count::count_udaf; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr}; +use datafusion_physical_plan::test::TestMemoryExec; +use datafusion_physical_plan::windows::{BoundedWindowAggExec, create_window_expr}; +use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, collect}; + +const BATCH_SIZE: usize = 8192; +const N_BATCHES: usize = 16; +/// Distinct partition keys per batch in the sparse layout. Each batch +/// introduces this many previously-unseen keys, so the total partition count +/// is `N_BATCHES * SPARSE_KEYS_PER_BATCH`. +const SPARSE_KEYS_PER_BATCH: usize = 2048; + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("pk", DataType::UInt64, false), + Field::new("ts", DataType::UInt64, false), + ])) +} + +/// Batches with `ts` ascending across the whole input and partition keys +/// chosen by `pk_of_row`. +fn make_batches(pk_of_row: impl Fn(usize) -> u64) -> Vec { + (0..N_BATCHES) + .map(|b| { + let start = b * BATCH_SIZE; + let pk: UInt64Array = (start..start + BATCH_SIZE) + .map(|i| Some(pk_of_row(i))) + .collect(); + let ts: UInt64Array = (start..start + BATCH_SIZE) + .map(|i| Some(i as u64)) + .collect(); + RecordBatch::try_new(schema(), vec![Arc::new(pk), Arc::new(ts)]).unwrap() + }) + .collect() +} + +/// Round-robin over `n_partitions`: every partition receives rows in every +/// batch (when `n_partitions <= BATCH_SIZE`). +fn dense_batches(n_partitions: usize) -> Vec { + make_batches(move |i| (i % n_partitions) as u64) +} + +/// Keys clustered in time: batch `b` only contains keys in +/// `[b * SPARSE_KEYS_PER_BATCH, (b + 1) * SPARSE_KEYS_PER_BATCH)`, cycled so +/// that consecutive rows belong to different partitions. Previously-seen +/// keys never recur, but `Linear` mode cannot know that, so the live +/// partition set grows for the whole run. +fn sparse_batches() -> Vec { + make_batches(|i| { + ((i / BATCH_SIZE) * SPARSE_KEYS_PER_BATCH + (i % SPARSE_KEYS_PER_BATCH)) as u64 + }) +} + +/// Input laid out partition-by-partition (the `Sorted` layout). +fn sorted_batches(n_partitions: usize) -> Vec { + let rows_per_partition = BATCH_SIZE * N_BATCHES / n_partitions; + make_batches(move |i| (i / rows_per_partition) as u64) +} + +fn sort_expr(name: &str) -> PhysicalSortExpr { + PhysicalSortExpr { + expr: col(name, &schema()).unwrap(), + options: Default::default(), + } +} + +/// `RANGE BETWEEN CURRENT ROW AND 10 FOLLOWING` +fn range_frame() -> WindowFrame { + WindowFrame::new_bounds( + WindowFrameUnits::Range, + WindowFrameBound::CurrentRow, + WindowFrameBound::Following(ScalarValue::UInt64(Some(10))), + ) +} + +/// `ROWS BETWEEN CURRENT ROW AND 2 FOLLOWING` +fn rows_frame() -> WindowFrame { + WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::CurrentRow, + WindowFrameBound::Following(ScalarValue::UInt64(Some(2))), + ) +} + +/// `(ts) OVER (PARTITION BY pk ORDER BY ts )` for each +/// aggregate in `aggregates`. +fn window_exec( + batches: Vec, + mode: InputOrderMode, + input_ordering: Vec, + window_frame: &WindowFrame, + aggregates: &[(WindowFunctionDefinition, &str)], +) -> Arc { + let schema = schema(); + let source = TestMemoryExec::try_new(&[batches], Arc::clone(&schema), None) + .expect("memory exec") + .try_with_sort_information(LexOrdering::new(input_ordering).into_iter().collect()) + .expect("sort information"); + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(source))); + let args = vec![col("ts", &schema).unwrap()]; + let partitionby_exprs = vec![col("pk", &schema).unwrap()]; + let orderby_exprs = vec![PhysicalSortExpr { + expr: col("ts", &schema).unwrap(), + options: Default::default(), + }]; + let window_expr = aggregates + .iter() + .map(|(fun, name)| { + create_window_expr( + fun, + name.to_string(), + &args, + &partitionby_exprs, + &orderby_exprs, + Arc::new(window_frame.clone()), + input.schema(), + false, + false, + None, + ) + .expect("window expr") + }) + .collect::>(); + Arc::new( + BoundedWindowAggExec::try_new(window_expr, input, mode, true) + .expect("bounded window exec"), + ) +} + +fn count() -> (WindowFunctionDefinition, &'static str) { + ( + WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count", + ) +} + +fn sum() -> (WindowFunctionDefinition, &'static str) { + (WindowFunctionDefinition::AggregateUDF(sum_udaf()), "sum") +} + +fn bounded_window_benchmark(c: &mut Criterion) { + let rt = tokio::runtime::Runtime::new().unwrap(); + let mut group = c.benchmark_group("bounded_window_partitions"); + group.sample_size(10); + + let mut run_case = |name: String, plan: Arc| { + group.bench_function(name, |b| { + b.iter(|| { + let task_ctx = Arc::new(TaskContext::default()); + let batches = rt + .block_on(collect(Arc::clone(&plan), task_ctx)) + .expect("execution"); + assert_eq!( + batches.iter().map(|b| b.num_rows()).sum::(), + BATCH_SIZE * N_BATCHES + ); + }) + }); + }; + + for n_partitions in [100, 10_000] { + run_case( + format!("linear {n_partitions} partitions"), + window_exec( + dense_batches(n_partitions), + InputOrderMode::Linear, + vec![sort_expr("ts")], + &range_frame(), + &[count()], + ), + ); + } + + run_case( + format!( + "linear sparse {} partitions", + N_BATCHES * SPARSE_KEYS_PER_BATCH + ), + window_exec( + sparse_batches(), + InputOrderMode::Linear, + vec![sort_expr("ts")], + &range_frame(), + &[count()], + ), + ); + + run_case( + "linear rows 10000 partitions".to_string(), + window_exec( + dense_batches(10_000), + InputOrderMode::Linear, + vec![sort_expr("ts")], + &rows_frame(), + &[count()], + ), + ); + + run_case( + "linear multi 10000 partitions".to_string(), + window_exec( + dense_batches(10_000), + InputOrderMode::Linear, + vec![sort_expr("ts")], + &range_frame(), + &[count(), sum()], + ), + ); + + // Control: the same query over partition-sorted input, where finished + // partitions are pruned eagerly and the state maps stay small. + run_case( + "sorted 10000 partitions".to_string(), + window_exec( + sorted_batches(10_000), + InputOrderMode::Sorted, + vec![sort_expr("pk"), sort_expr("ts")], + &range_frame(), + &[count()], + ), + ); + + group.finish(); +} + +criterion_group!(benches, bounded_window_benchmark); +criterion_main!(benches); diff --git a/datafusion/physical-plan/benches/compute_statistics.rs b/datafusion/physical-plan/benches/compute_statistics.rs new file mode 100644 index 0000000000000..cddf4c2396f42 --- /dev/null +++ b/datafusion/physical-plan/benches/compute_statistics.rs @@ -0,0 +1,354 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for `compute_statistics` with `StatsCache`. +//! +//! Demonstrates that caching eliminates redundant subtree walks in plans +//! containing partition-merging operators (CoalescePartitionsExec) and +//! binary join trees (CrossJoinExec). +//! +//! The plan shapes here mirror the reproducers from the planning-speed +//! EPIC (): +//! - Coalesce chain: deep linear plans (e.g. deeply nested subqueries) +//! - Cross-join tree: balanced binary trees from multi-way joins +//! (mirrors the `physical_many_self_joins` sql_planner benchmark) + +use std::fmt; +use std::sync::Arc; + +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::ScalarValue; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, Statistics}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::EquivalenceProperties; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::expressions::Literal; +use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion_physical_plan::execution_plan::{ + Boundedness, EmissionType, ExecutionPlan, PlanProperties, +}; +use datafusion_physical_plan::filter::FilterExec; +use datafusion_physical_plan::joins::CrossJoinExec; +use datafusion_physical_plan::statistics::StatisticsArgs; +use datafusion_physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Partitioning, + ReplaceChildrenOptions, SendableRecordBatchStream, StatisticsContext, +}; + +/// Minimal leaf node for benchmarking +#[derive(Debug)] +struct BenchLeaf { + schema: SchemaRef, + cache: Arc, +} + +impl BenchLeaf { + fn new(col_name: &str) -> Self { + let schema = Arc::new(Schema::new(vec![Field::new( + col_name, + DataType::Int32, + false, + )])); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + Partitioning::UnknownPartitioning(2), + EmissionType::Incremental, + Boundedness::Bounded, + )); + Self { schema, cache } + } +} + +impl DisplayAs for BenchLeaf { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "BenchLeaf") + } +} + +impl ExecutionPlan for BenchLeaf { + fn name(&self) -> &str { + "BenchLeaf" + } + + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::new(Statistics::new_unknown(&self.schema))) + } +} + +/// Build: CoalescePartitions^depth -> BenchLeaf +fn build_coalesce_chain(depth: usize) -> Arc { + let mut plan: Arc = Arc::new(BenchLeaf::new("a")); + for _ in 0..depth { + plan = Arc::new(CoalescePartitionsExec::new(plan)); + } + plan +} + +/// Build a balanced binary tree of CrossJoinExec with 2^depth leaves. +/// Mirrors the plan shape produced by multi-way self-joins like the +/// `physical_many_self_joins` benchmark in sql_planner.rs (#19795). +fn build_cross_join_tree(depth: usize, next_col: &mut usize) -> Arc { + if depth == 0 { + let col_name = format!("c{next_col}"); + *next_col += 1; + return Arc::new(BenchLeaf::new(&col_name)); + } + let left = build_cross_join_tree(depth - 1, next_col); + let right = build_cross_join_tree(depth - 1, next_col); + Arc::new(CrossJoinExec::new(left, right)) +} + +/// Build: Filter^depth -> BenchLeaf (always-true predicate). +fn build_filter_chain(depth: usize) -> Arc { + let mut plan: Arc = Arc::new(BenchLeaf::new("a")); + let predicate: Arc = + Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))); + for _ in 0..depth { + plan = Arc::new( + FilterExec::try_new(Arc::clone(&predicate), plan) + .expect("FilterExec::try_new failed"), + ); + } + plan +} + +/// Build a mixed chain alternating partition-merging and partition-preserving +/// operators: (Coalesce -> Filter -> Filter) repeated `groups` times -> BenchLeaf. +/// Exercises the cache with both None and Some(p) lookups in the same walk. +fn build_mixed_chain(groups: usize) -> Arc { + let mut plan: Arc = Arc::new(BenchLeaf::new("a")); + let predicate: Arc = + Arc::new(Literal::new(ScalarValue::Boolean(Some(true)))); + for _ in 0..groups { + // Two partition-preserving filters + for _ in 0..2 { + plan = Arc::new( + FilterExec::try_new(Arc::clone(&predicate), plan) + .expect("FilterExec::try_new failed"), + ); + } + // One partition-merging coalesce + plan = Arc::new(CoalescePartitionsExec::new(plan)); + } + plan +} + +/// Recursive walk without a shared cross-node cache, simulating pre-cache behavior. +/// Each node is computed with a fresh `StatisticsContext`, so every call triggers a +/// fresh subtree walk, resulting in O(n^2) total node visits for a chain of depth n. +/// +/// Note: each `StatisticsContext::compute` re-walk still benefits from its own +/// ephemeral cache; only the cross-node sharing is removed. +fn compute_statistics_without_shared_cache( + plan: &dyn ExecutionPlan, + partition: Option, +) -> Result> { + for child in plan.children() { + compute_statistics_without_shared_cache(child.as_ref(), None)?; + } + let args = StatisticsArgs::new().with_partition(partition); + StatisticsContext::new().compute(plan, &args) +} + +fn bench_compute_statistics(c: &mut Criterion) { + // --- Coalesce chain (linear plan) --- + // Deep linear plans arise from deeply nested subqueries, CTEs, etc. + let mut group = c.benchmark_group("compute_statistics_coalesce_chain"); + for depth in [10, 20, 50] { + let plan = build_coalesce_chain(depth); + group.bench_with_input(BenchmarkId::new("cached", depth), &plan, |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute(plan.as_ref(), &StatisticsArgs::new()) + .unwrap() + }); + }); + group.bench_with_input( + BenchmarkId::new("no_shared_cache", depth), + &plan, + |b, plan| { + b.iter(|| { + compute_statistics_without_shared_cache(plan.as_ref(), None).unwrap() + }); + }, + ); + } + group.finish(); + + // --- Cross-join tree (balanced binary plan) --- + // Binary trees arise from multi-way joins (e.g. physical_many_self_joins + // in sql_planner.rs, see #19795). CrossJoinExec calls + // StatisticsContext::compute for per-partition stats, re-walking the left + // subtree at each node. The gap between cached/uncached is smaller than + // the linear chain because only the left child triggers a re-walk. + let mut group = c.benchmark_group("compute_statistics_cross_join_tree"); + for depth in [3, 5, 7] { + let mut next_col = 0; + let plan = build_cross_join_tree(depth, &mut next_col); + let label = format!("depth={depth}_leaves={}", 1usize << depth); + group.bench_with_input(BenchmarkId::new("cached", &label), &plan, |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap() + }); + }); + group.bench_with_input( + BenchmarkId::new("no_shared_cache", &label), + &plan, + |b, plan| { + b.iter(|| { + compute_statistics_without_shared_cache(plan.as_ref(), Some(0)) + .unwrap() + }); + }, + ); + } + group.finish(); + + // --- Filter chain (partition-preserving linear plan) --- + // When called with Some(0), the framework first walks the entire tree + // computing None stats, then each filter requests Some(0) on demand. + // Both walks are cached, so the total cost is ~2n vs n node visits for None. + let mut group = c.benchmark_group("compute_statistics_filter_chain"); + for depth in [10, 20, 50] { + let plan = build_filter_chain(depth); + group.bench_with_input( + BenchmarkId::new("cached_partition", depth), + &plan, + |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap() + }); + }, + ); + group.bench_with_input( + BenchmarkId::new("cached_overall", depth), + &plan, + |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute(plan.as_ref(), &StatisticsArgs::new()) + .unwrap() + }); + }, + ); + group.bench_with_input( + BenchmarkId::new("no_shared_cache", depth), + &plan, + |b, plan| { + b.iter(|| { + compute_statistics_without_shared_cache(plan.as_ref(), Some(0)) + .unwrap() + }); + }, + ); + } + group.finish(); + + // --- Mixed chain (partition-preserving + partition-merging) --- + // Alternates Filter (preserving) and CoalescePartitions (merging) to + // exercise the cache with both None and Some(p) lookups in a single walk. + let mut group = c.benchmark_group("compute_statistics_mixed_chain"); + for groups in [3, 5, 10] { + let plan = build_mixed_chain(groups); + let depth = groups * 3; // 2 filters + 1 coalesce per group + group.bench_with_input(BenchmarkId::new("cached", depth), &plan, |b, plan| { + b.iter(|| { + StatisticsContext::new() + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap() + }); + }); + group.bench_with_input( + BenchmarkId::new("no_shared_cache", depth), + &plan, + |b, plan| { + b.iter(|| { + compute_statistics_without_shared_cache(plan.as_ref(), Some(0)) + .unwrap() + }); + }, + ); + } + group.finish(); +} + +criterion_group!(benches, bench_compute_statistics); +criterion_main!(benches); diff --git a/datafusion/physical-plan/benches/hash_join_semi_anti.rs b/datafusion/physical-plan/benches/hash_join_semi_anti.rs new file mode 100644 index 0000000000000..1e11da36be73c --- /dev/null +++ b/datafusion/physical-plan/benches/hash_join_semi_anti.rs @@ -0,0 +1,387 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Criterion benchmarks for Hash Join with RightSemi/RightAnti joins with Int32 keys. +//! +//! ## Key Benchmark Axes +//! +//! - **Density**: How tightly distinct keys pack into their numeric range. +//! `density = num_distinct_keys / (max_key - min_key + 1)`. +//! Examples for 5 distinct keys: +//! - `[0, 1, 2, 3, 4]` → 5/5 = 100% (fully packed) +//! - `[0, 2, 4, 6, 8]` → 5/9 ≈ 55% (every 2nd slot) +//! - `[0, 10, 20, 30, 40]` → 5/41 ≈ 12% (every 10th slot) +//! +//! Why it matters for this workload: future potential semi/anti-join +//! fast paths could exploit densely packed build keys to outperform the +//! general hash-table path, which is largely insensitive to density. +//! Varying density across benchmarks helps surface those potential gains +//! under different key distributions. Density describes only the +//! build-side key layout; the per-probe match count is tracked +//! separately as fanout. +//! +//! - **Hit Rate**: The percentage of probe rows that find a match in the build side. +//! This controls how often the join produces output rows. +//! +//! Semi/anti joins can short-circuit after finding the first match, so these +//! benchmarks help evaluate optimization strategies for existence checks. + +use std::sync::Arc; + +use arrow::array::{Int32Array, RecordBatch, StringArray}; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_common::{JoinType, NullEquality}; +use datafusion_execution::TaskContext; +use datafusion_physical_expr::expressions::col; +use datafusion_physical_plan::collect; +use datafusion_physical_plan::joins::{HashJoinExec, PartitionMode, utils::JoinOn}; +use datafusion_physical_plan::test::TestMemoryExec; +use tokio::runtime::Runtime; + +/// Build RecordBatches with Int32 keys. +/// +/// Schema: (key: Int32, data: Int32, payload: Utf8) +/// +/// `key_mod` controls distinct key count: key = row_index % key_mod. +/// `key_offset` shifts keys to control hit rate. +fn build_batches( + num_rows: usize, + key_mod: usize, + key_offset: i32, + schema: &SchemaRef, +) -> Vec { + let keys: Vec = (0..num_rows) + .map(|i| ((i % key_mod) as i32) + key_offset) + .collect(); + let data: Vec = (0..num_rows).map(|i| i as i32).collect(); + let payload: Vec = data.iter().map(|d| format!("val_{d}")).collect(); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(data)), + Arc::new(StringArray::from(payload)), + ], + ) + .unwrap(); + + let batch_size = 8192; + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(batch_size); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn make_exec( + batches: &[RecordBatch], + schema: &SchemaRef, +) -> Arc { + TestMemoryExec::try_new_exec(&[batches.to_vec()], Arc::clone(schema), None).unwrap() +} + +fn schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("data", DataType::Int32, false), + Field::new("payload", DataType::Utf8, false), + ])) +} + +fn do_hash_join( + left: Arc, + right: Arc, + join_type: JoinType, + rt: &Runtime, +) -> usize { + let on: JoinOn = vec![( + col("key", &left.schema()).unwrap(), + col("key", &right.schema()).unwrap(), + )]; + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &join_type, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + ) + .unwrap(); + + let task_ctx = Arc::new(TaskContext::default()); + rt.block_on(async { + let batches = collect(Arc::new(join), task_ctx).await.unwrap(); + batches.iter().map(|b| b.num_rows()).sum() + }) +} + +/// Build batches with sparse keys (key = row_index % key_mod * multiplier + key_offset). +/// The `multiplier` controls density: 1 = 100%, 2 = 50%, 10 = 10%. +fn build_batches_sparse( + num_rows: usize, + key_mod: usize, + key_offset: i32, + multiplier: i32, + schema: &SchemaRef, +) -> Vec { + let keys: Vec = (0..num_rows) + .map(|i| ((i % key_mod) as i32) * multiplier + key_offset) + .collect(); + let data: Vec = (0..num_rows).map(|i| i as i32).collect(); + let payload: Vec = data.iter().map(|d| format!("val_{d}")).collect(); + + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(keys)), + Arc::new(Int32Array::from(data)), + Arc::new(StringArray::from(payload)), + ], + ) + .unwrap(); + + let batch_size = 8192; + let mut batches = Vec::new(); + let mut offset = 0; + while offset < batch.num_rows() { + let len = (batch.num_rows() - offset).min(batch_size); + batches.push(batch.slice(offset, len)); + offset += len; + } + batches +} + +fn bench_hash_join_semi_anti(c: &mut Criterion) { + let rt = Runtime::new().unwrap(); + let s = schema(); + + let mut group = c.benchmark_group("hash_join_semi_anti"); + + // Build side: 100K rows, Probe side: 1M rows + // Matching ratio: 1:1 (build keys are unique, each probe matches at most 1 build row) + let build_rows = 100_000; + let probe_rows = 1_000_000; + + // ========================================================================= + // RightSemi Join benchmarks + // ========================================================================= + + // RightSemi - 100% Density, 100% hit rate + // Keys: 0..100K contiguous, all probe rows find a match + { + let left_batches = build_batches(build_rows, build_rows, 0, &s); + let right_batches = build_batches(probe_rows, build_rows, 0, &s); + group.bench_function(BenchmarkId::new("right_semi_d100_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 100% Density, 10% hit rate + // Keys: 0..100K contiguous, only 10% of probe rows find a match + { + let left_batches = build_batches(build_rows, build_rows, 0, &s); + let right_batches = build_batches(probe_rows, build_rows * 10, 0, &s); + group.bench_function(BenchmarkId::new("right_semi_d100_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 50% Density, 100% hit rate + // Keys: 0, 2, 4, ... (sparse, multiplier=2), all probe rows find a match + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 2, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows, 0, 2, &s); + group.bench_function(BenchmarkId::new("right_semi_d50_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 50% Density, 10% hit rate + // Keys: 0, 2, 4, ... (sparse), only 10% of probe rows find a match + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 2, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows * 10, 0, 2, &s); + group.bench_function(BenchmarkId::new("right_semi_d50_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 10% Density, 100% hit rate + // Keys: 0, 10, 20, ... (very sparse, multiplier=10), all probe rows find a match + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 10, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows, 0, 10, &s); + group.bench_function(BenchmarkId::new("right_semi_d10_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 10% Density, 10% hit rate + // Keys: 0, 10, 20, ... (very sparse), only 10% of probe rows find a match + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 10, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows * 10, 0, 10, &s); + group.bench_function(BenchmarkId::new("right_semi_d10_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }); + } + + // RightSemi - 100% Density, ~1% hit rate, fanout ~100 + // Build keys are duplicated: 100K rows over 1K distinct keys. Matching + // probe rows produce many duplicate probe indices before RightSemi + // deduplication. + { + let fanout_keys = 1_000; + let left_batches = build_batches(build_rows, fanout_keys, 0, &s); + let right_batches = build_batches(probe_rows, build_rows, 0, &s); + group.bench_function( + BenchmarkId::new("right_semi_fanout100_h1", probe_rows), + |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightSemi, &rt) + }) + }, + ); + } + + // ========================================================================= + // RightAnti Join benchmarks + // ========================================================================= + + // RightAnti - 100% Density, 100% hit rate (no output) + // Keys: 0..100K contiguous, all probe rows find a match -> no output + { + let left_batches = build_batches(build_rows, build_rows, 0, &s); + let right_batches = build_batches(probe_rows, build_rows, 0, &s); + group.bench_function(BenchmarkId::new("right_anti_d100_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 100% Density, 10% hit rate (90% output) + // Keys: 0..100K contiguous, only 10% of probe rows find a match -> 90% output + { + let left_batches = build_batches(build_rows, build_rows, 0, &s); + let right_batches = build_batches(probe_rows, build_rows * 10, 0, &s); + group.bench_function(BenchmarkId::new("right_anti_d100_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 50% Density, 100% hit rate (no output) + // Keys: 0, 2, 4, ... (sparse), all probe rows find a match -> no output + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 2, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows, 0, 2, &s); + group.bench_function(BenchmarkId::new("right_anti_d50_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 50% Density, 10% hit rate (90% output) + // Keys: 0, 2, 4, ... (sparse), only 10% of probe rows find a match -> 90% output + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 2, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows * 10, 0, 2, &s); + group.bench_function(BenchmarkId::new("right_anti_d50_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 10% Density, 100% hit rate (no output) + // Keys: 0, 10, 20, ... (very sparse), all probe rows find a match -> no output + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 10, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows, 0, 10, &s); + group.bench_function(BenchmarkId::new("right_anti_d10_h100", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + // RightAnti - 10% Density, 10% hit rate (90% output) + // Keys: 0, 10, 20, ... (very sparse), only 10% of probe rows find a match -> 90% output + { + let left_batches = build_batches_sparse(build_rows, build_rows, 0, 10, &s); + let right_batches = build_batches_sparse(probe_rows, build_rows * 10, 0, 10, &s); + group.bench_function(BenchmarkId::new("right_anti_d10_h10", probe_rows), |b| { + b.iter(|| { + let left = make_exec(&left_batches, &s); + let right = make_exec(&right_batches, &s); + do_hash_join(left, right, JoinType::RightAnti, &rt) + }) + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_hash_join_semi_anti); +criterion_main!(benches); diff --git a/datafusion/physical-plan/benches/multi_group_by.rs b/datafusion/physical-plan/benches/multi_group_by.rs new file mode 100644 index 0000000000000..0c689f9fcb6ce --- /dev/null +++ b/datafusion/physical-plan/benches/multi_group_by.rs @@ -0,0 +1,815 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Benchmarks for multi-column GROUP BY performance comparing vectorized +//! (`GroupValuesColumn`) vs row-based (`GroupValuesRows`) implementations. +//! +//! Motivated by which +//! showed vectorized can regress for low-cardinality, high-row-count scenarios. +//! +//! Uses the direct `GroupValues::intern()` API with identical data for both +//! implementations — a fair apples-to-apples comparison with the same hashing +//! and data layout. Most experiments use `Int32` columns; `bench_fixed_size_binary` +//! covers a `(FixedSizeBinary, Int32)` key to exercise the +//! `FixedSizeBinaryGroupValueBuilder`. + +use arrow::array::{ + ArrayRef, Decimal256Array, DurationMicrosecondArray, Float16Array, Int32Array, + IntervalMonthDayNanoArray, UInt32Array, +}; +use arrow::compute::take; +use arrow::datatypes::{ + DataType, Field, IntervalMonthDayNano, IntervalUnit, Schema, SchemaRef, TimeUnit, + i256, +}; +use arrow::util::bench_util::create_fsb_array; +use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main}; +use datafusion_physical_plan::aggregates::group_values::GroupValues; +use datafusion_physical_plan::aggregates::group_values::GroupValuesRows; +use datafusion_physical_plan::aggregates::group_values::multi_group_by::GroupValuesColumn; +use half::f16; +use std::hint::black_box; +use std::sync::Arc; + +const DEFAULT_BATCH_SIZE: usize = 8192; + +fn make_schema(num_cols: usize) -> SchemaRef { + let fields: Vec = (0..num_cols) + .map(|i| Field::new(format!("col_{i}"), DataType::Int32, false)) + .collect(); + Arc::new(Schema::new(fields)) +} + +fn generate_batches( + num_cols: usize, + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + let per_col_card = (num_distinct_groups as f64) + .powf(1.0 / num_cols as f64) + .ceil() as usize; + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + (0..num_cols) + .map(|col_idx| { + let values: Vec = (0..current_batch_size) + .map(|row| { + let global_row = batch_start + row; + let group_id = global_row % num_distinct_groups; + let divisor = per_col_card.pow(col_idx as u32); + ((group_id / divisor) % per_col_card) as i32 + }) + .collect(); + Arc::new(Int32Array::from(values)) as ArrayRef + }) + .collect() + }) + .collect() +} + +fn create_group_values(schema: &SchemaRef, vectorized: bool) -> Box { + if vectorized { + Box::new(GroupValuesColumn::::try_new(Arc::clone(schema)).unwrap()) + } else { + Box::new(GroupValuesRows::try_new(Arc::clone(schema)).unwrap()) + } +} + +fn bench_intern( + gv: &mut Box, + batches: &[Vec], + groups: &mut Vec, +) { + for batch in batches { + groups.clear(); + gv.intern(batch, groups).unwrap(); + } + black_box(&*groups); +} + +/// Experiment 1: Issue #17850 regression scenario. +/// 3 columns, 64 groups (4^3), scaling row count. +fn bench_issue_17850_regression(c: &mut Criterion) { + let mut group = c.benchmark_group("issue_17850_regression"); + group.sample_size(10); + + let num_cols = 3; + let num_groups = 64; + let schema = make_schema(num_cols); + + for num_rows in [1_000_000, 5_000_000, 10_000_000, 20_000_000, 50_000_000] { + let batches = + generate_batches(num_cols, num_groups, num_rows, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("{num_rows}_rows")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 2: Low cardinality sweep. +fn bench_low_cardinality(c: &mut Criterion) { + let mut group = c.benchmark_group("low_cardinality"); + group.sample_size(15); + + for (num_cols, per_col_card) in + [(3usize, 2usize), (3, 4), (3, 8), (4, 2), (4, 4), (4, 8)] + { + let num_groups = per_col_card.pow(num_cols as u32); + let schema = make_schema(num_cols); + let batches = + generate_batches(num_cols, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new( + label, + format!("cols_{num_cols}_card_{per_col_card}_grp_{num_groups}"), + ), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 3: Batch size sensitivity. +fn bench_batch_size_sensitivity(c: &mut Criterion) { + let mut group = c.benchmark_group("batch_size_sensitivity"); + group.sample_size(10); + + let num_cols = 3; + let num_groups = 64; + let schema = make_schema(num_cols); + + for batch_size in [1024, 4096, 8192, 16384, 32768] { + let batches = generate_batches(num_cols, num_groups, 1_000_000, batch_size); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("batch_{batch_size}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(batch_size), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 4: Column count scaling with low groups. +fn bench_column_scaling(c: &mut Criterion) { + let mut group = c.benchmark_group("column_scaling"); + group.sample_size(15); + + let cases: &[(usize, usize)] = + &[(2, 100), (3, 125), (4, 81), (6, 729), (8, 256), (10, 1024)]; + + for &(num_cols, num_groups) in cases { + let schema = make_schema(num_cols); + let batches = + generate_batches(num_cols, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("cols_{num_cols}_grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 5: High cardinality column scaling (~1M groups). +fn bench_high_cardinality_scaling(c: &mut Criterion) { + let mut group = c.benchmark_group("high_cardinality_scaling"); + group.sample_size(10); + + for num_cols in [2, 3, 4, 6, 8, 10] { + let num_groups = 1_000_000; + let schema = make_schema(num_cols); + let batches = + generate_batches(num_cols, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("cols_{num_cols}_grp_1M")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Experiment 6: Group count sweep with fixed 4 columns. +fn bench_group_count_sweep(c: &mut Criterion) { + let mut group = c.benchmark_group("group_count_sweep"); + group.sample_size(15); + + let num_cols = 4; + let schema = make_schema(num_cols); + + for num_groups in [ + 16, 64, 256, 1000, 5000, 10_000, 50_000, 100_000, 500_000, 1_000_000, + ] { + let batches = + generate_batches(num_cols, num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +/// Width in bytes of the FixedSizeBinary group column (UUID-sized). +const FSB_WIDTH: usize = 16; + +/// Schema for the FixedSizeBinary experiment: a `FixedSizeBinary` group column +/// paired with an `Int32` column, exercising a multi-column GROUP BY that +/// includes a fixed-width binary key (e.g. grouping on a UUID). +fn make_fsb_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("fsb", DataType::FixedSizeBinary(FSB_WIDTH as i32), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(FixedSizeBinary, Int32)` batches with exactly +/// `num_distinct_groups` distinct keys. +/// +/// The distinct FixedSizeBinary values come from arrow-rs's `create_fsb_array` +/// benchmark generator; rows cycle through that pool (mirroring how +/// `generate_batches` controls Int32 cardinality) so the group count is +/// controlled. The `Int32` column is keyed identically, keeping the combined +/// cardinality equal to `num_distinct_groups`. +fn generate_fsb_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + // Pool of distinct FixedSizeBinary values (fixed seed, no nulls). + let pool = create_fsb_array(num_distinct_groups, 0.0, FSB_WIDTH); + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let indices: UInt32Array = group_ids.clone().map(|g| g as u32).collect(); + let fsb = take(&pool, &indices, None).unwrap(); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![fsb, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 7: Group count sweep for a `(FixedSizeBinary, Int32)` key. +/// +/// Exercises the `FixedSizeBinaryGroupValueBuilder` used by multi-column +/// GROUP BY. Before FixedSizeBinary support, such a schema fell back to the +/// row-based `GroupValuesRows`; this compares the vectorized columnar path +/// (`vectorized`) against that baseline (`row_based`). +fn bench_fixed_size_binary(c: &mut Criterion) { + let mut group = c.benchmark_group("fixed_size_binary"); + group.sample_size(15); + + let schema = make_fsb_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = generate_fsb_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +fn make_f16_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("f16", DataType::Float16, false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Float16, Int32)` batches with `num_distinct_groups` distinct keys. +/// +/// `f16` has only ~63.5k finite values, so `num_distinct_groups` must stay well +/// under that (see `bench_float16`). Distinct keys are the low finite `f16` bit +/// patterns, skipping NaN and inf. The `Int32` column is keyed identically so +/// the combined cardinality equals `num_distinct_groups`. +fn generate_f16_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + let pool: Vec = (0u16..) + .map(f16::from_bits) + .filter(|v| v.is_finite()) + .take(num_distinct_groups) + .collect(); + + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let keys = Float16Array::from_iter_values(group_ids.clone().map(|g| pool[g])); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 8: Group count sweep for a `(Float16, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Float16` on the +/// multi-column path (previously such a schema fell back to `GroupValuesRows`). +/// Group counts are capped below `f16`'s ~63.5k distinct finite values. +fn bench_float16(c: &mut Criterion) { + let mut group = c.benchmark_group("float16"); + group.sample_size(15); + + let schema = make_f16_schema(); + + for num_groups in [1_000, 60_000] { + let batches = generate_f16_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +fn make_duration_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("dur", DataType::Duration(TimeUnit::Microsecond), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Duration(Microsecond), Int32)` batches with `num_distinct_groups` +/// distinct keys. +/// +/// Each distinct duration is `g` microseconds. The `Int32` column is keyed +/// identically so the combined cardinality equals `num_distinct_groups`. +fn generate_duration_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let keys = DurationMicrosecondArray::from_iter_values( + group_ids.clone().map(|g| g as i64), + ); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 9: Group count sweep for a `(Duration, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Duration` on the +/// multi-column path (previously such a schema fell back to `GroupValuesRows`). +fn bench_duration(c: &mut Criterion) { + let mut group = c.benchmark_group("duration"); + group.sample_size(15); + + let schema = make_duration_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_duration_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +fn make_interval_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("iv", DataType::Interval(IntervalUnit::MonthDayNano), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Interval(MonthDayNano), Int32)` batches with `num_distinct_groups` +/// distinct keys. +/// +/// Each distinct interval is `MonthDayNano(g, 0, 0)`. The `Int32` column is +/// keyed identically so the combined cardinality equals `num_distinct_groups`. +fn generate_interval_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let keys = IntervalMonthDayNanoArray::from_iter_values( + group_ids + .clone() + .map(|g| IntervalMonthDayNano::new(g as i32, 0, 0)), + ); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 10: Group count sweep for an `(Interval, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Interval` on the +/// multi-column path (previously such a schema fell back to `GroupValuesRows`). +fn bench_interval(c: &mut Criterion) { + let mut group = c.benchmark_group("interval"); + group.sample_size(15); + + let schema = make_interval_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_interval_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +fn make_decimal256_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("dec", DataType::Decimal256(50, 0), false), + Field::new("id", DataType::Int32, false), + ])) +} + +/// Generate `(Decimal256(50, 0), Int32)` batches with `num_distinct_groups` +/// distinct keys. +/// +/// Each distinct value is `i256::from_i128(g)`, and precision > 38 keeps it a +/// genuine `Decimal256`. The `Int32` column is keyed identically so the combined +/// cardinality equals `num_distinct_groups`. +fn generate_decimal256_batches( + num_distinct_groups: usize, + num_rows: usize, + batch_size: usize, +) -> Vec> { + let num_full_batches = num_rows / batch_size; + let remainder = num_rows % batch_size; + let num_batches = num_full_batches + if remainder > 0 { 1 } else { 0 }; + + (0..num_batches) + .map(|batch_idx| { + let batch_start = batch_idx * batch_size; + let current_batch_size = if batch_idx == num_batches - 1 && remainder > 0 { + remainder + } else { + batch_size + }; + + let group_ids = (0..current_batch_size) + .map(|row| (batch_start + row) % num_distinct_groups); + + let keys = Decimal256Array::from_iter_values( + group_ids.clone().map(|g| i256::from_i128(g as i128)), + ) + .with_precision_and_scale(50, 0) + .unwrap(); + let id: Int32Array = group_ids.map(|g| g as i32).collect(); + + vec![Arc::new(keys) as ArrayRef, Arc::new(id) as ArrayRef] + }) + .collect() +} + +/// Experiment 11: Group count sweep for a `(Decimal256, Int32)` key. +/// +/// Exercises the primitive `GroupColumn` builder for `Decimal256` (32-byte +/// `i256` native) on the multi-column path (previously such a schema fell back +/// to `GroupValuesRows`). +fn bench_decimal256(c: &mut Criterion) { + let mut group = c.benchmark_group("decimal256"); + group.sample_size(15); + + let schema = make_decimal256_schema(); + + for num_groups in [1_000, 1_000_000] { + let batches = + generate_decimal256_batches(num_groups, 1_000_000, DEFAULT_BATCH_SIZE); + + for vectorized in [true, false] { + let label = if vectorized { + "vectorized" + } else { + "row_based" + }; + group.bench_with_input( + BenchmarkId::new(label, format!("grp_{num_groups}")), + &batches, + |b, batches| { + b.iter_batched_ref( + || { + ( + create_group_values(&schema, vectorized), + Vec::::with_capacity(DEFAULT_BATCH_SIZE), + ) + }, + |(gv, groups)| bench_intern(gv, batches, groups), + criterion::BatchSize::LargeInput, + ); + }, + ); + } + } + group.finish(); +} + +criterion_group!( + benches, + bench_issue_17850_regression, + bench_low_cardinality, + bench_batch_size_sensitivity, + bench_column_scaling, + bench_high_cardinality_scaling, + bench_group_count_sweep, + bench_fixed_size_binary, + bench_float16, + bench_duration, + bench_interval, + bench_decimal256, +); +criterion_main!(benches); diff --git a/datafusion/physical-plan/benches/spill_io.rs b/datafusion/physical-plan/benches/spill_io.rs index fac2547a131b4..ddd83ca565533 100644 --- a/datafusion/physical-plan/benches/spill_io.rs +++ b/datafusion/physical-plan/benches/spill_io.rs @@ -547,7 +547,7 @@ fn benchmark_spill_batches_for_all_codec( let write_throughput = (mem_bytes as u128 / write_time.as_millis().max(1)) * 1000; // calculate compression ratio - let disk_bytes = std::fs::metadata(spill_file.path()) + let disk_bytes = std::fs::metadata(spill_file.path().unwrap()) .expect("metadata read fail") .len() as usize; let ratio = mem_bytes as f64 / disk_bytes.max(1) as f64; diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs new file mode 100644 index 0000000000000..91e9d6555c3e7 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common.rs @@ -0,0 +1,693 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::array::{ArrayRef, AsArray, new_null_array}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, internal_err}; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::{EmitTo, GroupsAccumulator}; +use datafusion_physical_expr::aggregate::AggregateFunctionExpr; + +use crate::PhysicalExpr; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::grouped_hash_stream::create_group_accumulator; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{ + AggregateExec, PhysicalGroupBy, aggregate_expressions, evaluate_group_by, +}; + +/// Marker for raw rows -> partial state aggregation. +pub(in crate::aggregates) struct PartialMarker; +/// Marker for raw rows -> final value aggregation. +pub(in crate::aggregates) struct SingleMarker; +/// Marker for partial state -> partial state aggregation. +pub(in crate::aggregates) struct PartialReduceMarker; +/// Marker for raw rows -> partial state conversion without aggregation. +pub(in crate::aggregates) struct PartialSkipMarker; +/// Marker for partial state -> final value aggregation. +pub(in crate::aggregates) struct FinalMarker; + +/// Grouped hash table shared by the partial and final paths. +/// +/// While building, it consumes input batches and updates group / accumulator +/// state. While outputting, it incrementally drains that state into output +/// batches. +/// +/// # Logical and Physical Model +/// +/// Logically, this is a hash table that maps { group keys -> accumulator states } +/// For example, `AVG(v) GROUP BY k` stores one entry per `k`, where each +/// entry owns the `sum(v)` and `count(v)` state needed to compute the final +/// average. +/// +/// Physically, the group keys and accumulators are backed by [`GroupValues`] and +/// [`GroupsAccumulator`]. Both use columnar storage so aggregation can stay +/// vectorized. +/// +/// # Marker Type +/// `AggrMode` selects the aggregate semantics. +/// +/// e.g. `AggregateHashTable::::new(...)` creates an aggregate hash table +/// for the partial hash aggregate stage, the input schema is raw rows and output +/// schema is intermediate states. +/// +/// It is a zero-sized compile-time marker, so each stage keeps its update logic +/// in a separate impl block, to make the behavior difference explicit. +pub(in crate::aggregates) struct AggregateHashTable { + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Raw input schema, used to evaluate expressions and synthesize empty + /// grouping-set rows. + pub(super) input_schema: SchemaRef, + + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Intermediate-state schema used when memory pressure requires the table + /// to spill its current state. + pub(super) state_schema: SchemaRef, + + /// Maximum rows per emitted output batch, from config `batch_size`. + pub(super) batch_size: usize, + + /// Lifecycle-specific state: building stage / outputting stage. + pub(super) state: AggregateHashTableState, + + pub(super) _mode: PhantomData, +} + +/// Methods shared by all aggregate hash table modes. +impl AggregateHashTable { + pub(super) fn new_with_filters( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + filters: Vec>>, + ) -> Result { + if batch_size == 0 { + return internal_err!("AggregateHashTable requires config batch_size >= 1"); + } + + let input_schema = agg.input().schema(); + let aggregate_arguments = aggregate_expressions( + &agg.aggr_expr, + &agg.mode, + agg.group_by.num_group_exprs(), + )?; + let accumulators: Vec<_> = agg + .aggr_expr + .iter() + .zip(aggregate_arguments) + .zip(filters) + .map(|((agg_expr, arguments), filter)| { + let accumulator = create_group_accumulator(agg_expr)?; + Ok(HashAggregateAccumulator::new( + Arc::clone(agg_expr), + arguments, + filter, + accumulator, + )) + }) + .collect::>()?; + + let group_schema = agg.group_by.group_schema(&input_schema)?; + let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + + Ok(Self { + group_by_metrics: GroupByMetrics::new(&agg.metrics, partition), + input_schema, + output_schema, + state_schema, + batch_size, + state: AggregateHashTableState::Building(AggregateHashTableBuffer { + group_by: Arc::clone(&agg.group_by), + group_values, + batch_group_indices: Default::default(), + accumulators, + }), + _mode: PhantomData, + }) + } + + /// See comments in [`EvaluatedAggregateBatch`] + pub(super) fn evaluate_batch( + &self, + batch: &RecordBatch, + ) -> Result { + let state = self.state.building(); + let timer = self.group_by_metrics.time_calculating_group_ids.timer(); + // outer vec: one per each grouping set + // inner vec: all group by exprs for the current grouping set + let grouping_set_args = evaluate_group_by(&state.group_by, batch)?; + drop(timer); + + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + // The evaluated args for each accumulator + let accumulator_args = self + .state + .building() + .accumulators + .iter() + .map(|acc| acc.evaluate_acc_args(batch)) + .collect::>>()?; + drop(timer); + + Ok(EvaluatedAggregateBatch { + grouping_set_args, + accumulator_args, + }) + } + + /// Aggregates one input batch after selecting the mode-specific accumulator + /// operation. + /// + /// Each aggregation mode chooses a different `aggregate_fn` according to its + /// semantics. For example, partial aggregation takes raw inputs, and update them + /// into stored partial states, so [`GroupsAccumulator::update_batch`] is used. + pub(super) fn aggregate_batch_inner( + &mut self, + batch: &RecordBatch, + aggregate_fn: AggregateBatchFn, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + let state = self.state.building_mut(); + + let _timer = self.group_by_metrics.aggregation_time.timer(); + for group_values in &evaluated_batch.grouping_set_args { + state + .group_values + .intern(group_values, &mut state.batch_group_indices)?; + let group_indices = &state.batch_group_indices; + let total_num_groups = state.group_values.len(); + + for (acc, values) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + aggregate_fn(acc, values, group_indices, total_num_groups)?; + } + } + + Ok(()) + } + + /// Materializes the full output once, then returns it downstream incrementally + /// by slicing it into `batch_size` chunks. + /// + /// Each aggregation mode chooses a different `materialize_accumulator_fn` + /// according to its semantics. For example, partial aggregation emits + /// partial states to feed the final stage, so it uses [`GroupsAccumulator::state`]. + /// + /// This is a temporary solution until blocked state management is implemented: + /// Issue: + pub(super) fn next_output_batch_inner( + &mut self, + materialize_accumulator_fn: MaterializeAccumulatorFn, + ) -> Result> { + let output_schema = Arc::clone(&self.output_schema); + let batch_size = self.batch_size; + + let mut output = + match std::mem::replace(&mut self.state, AggregateHashTableState::Done) { + AggregateHashTableState::Outputting(mut state) => { + if state.group_values.is_empty() { + return Ok(None); + } + + // Accumulator output consumes internal state. Materialize all + // groups once, then slice the materialized batch on later polls. + let emit_to = EmitTo::All; + let timer = self.group_by_metrics.emitting_time.timer(); + let mut columns = state.group_values.emit(emit_to)?; + for acc in state.accumulators.iter_mut() { + columns.extend(materialize_accumulator_fn(acc, emit_to)?); + } + drop(timer); + + let batch = RecordBatch::try_new(output_schema, columns)?; + debug_assert!(batch.num_rows() > 0); + MaterializedAggregateOutput::new(batch) + } + AggregateHashTableState::OutputtingMaterialized(output) => output, + AggregateHashTableState::Done => return Ok(None), + AggregateHashTableState::Building(_) => { + return internal_err!( + "next_output_batch must be called in the outputting state" + ); + } + }; + + let batch = output.next_batch(batch_size); + if output.is_exhausted() { + self.state = AggregateHashTableState::Done; + } else { + self.state = AggregateHashTableState::OutputtingMaterialized(output); + } + Ok(batch) + } + + pub(in crate::aggregates) fn memory_size(&self) -> usize { + match &self.state { + AggregateHashTableState::Building(state) + | AggregateHashTableState::Outputting(state) => { + let acc = state + .accumulators + .iter() + .map(|acc| acc.accumulator.size()) + .sum::(); + + acc + state.group_values.size() + + state.batch_group_indices.allocated_size() + } + AggregateHashTableState::OutputtingMaterialized(output) => { + output.memory_size() + } + AggregateHashTableState::Done => 0, + } + } + + pub(in crate::aggregates) fn group_by_metrics(&self) -> &GroupByMetrics { + &self.group_by_metrics + } + + /// Returns the number of distinct groups accumulated so far. + pub(in crate::aggregates) fn building_group_count(&self) -> usize { + self.state.building().group_values.len() + } + + /// Takes every intermediate aggregate state and resets the table so it can + /// continue accumulating raw input. + /// + /// Unlike normal single aggregation output, this materializes intermediate + /// states rather than final values. The states can therefore be merged after + /// spilling without finalizing the same group more than once. + pub(in crate::aggregates) fn take_state_batch( + &mut self, + ) -> Result> { + let state_schema = Arc::clone(&self.state_schema); + let state = self.state.building_mut(); + if state.group_values.is_empty() { + return Ok(None); + } + + let mut output = state.group_values.emit(EmitTo::All)?; + for acc in &mut state.accumulators { + output.extend(acc.state(EmitTo::All)?); + } + + let batch = RecordBatch::try_new(state_schema, output)?; + debug_assert!(batch.num_rows() > 0); + + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the + // key/index buffers too so the memory reservation can be released + // before the batch is sorted for spilling. + state.group_values.clear_shrink(0); + state.batch_group_indices.clear(); + state.batch_group_indices.shrink_to_fit(); + + Ok(Some(batch)) + } + + pub(in crate::aggregates) fn is_building(&self) -> bool { + matches!(self.state, AggregateHashTableState::Building(_)) + } + + pub(in crate::aggregates) fn is_done(&self) -> bool { + matches!(self.state, AggregateHashTableState::Done) + } + + pub(super) fn start_outputting(&mut self) { + let AggregateHashTableState::Building(mut state) = + std::mem::replace(&mut self.state, AggregateHashTableState::Done) + else { + unreachable!("hash aggregate table is not building") + }; + + state.batch_group_indices = Vec::new(); + self.state = AggregateHashTableState::Outputting(state); + } +} + +/// State and argument information for a single Aggregate +/// +/// For example, for `SELECT COUNT(x), SUM(y WHERE z > 10) ...` there would be two +/// `HashAggregateAccumulator`, one each for `COUNT(x)` and `SUM(y WHERE z > 10)` +pub(super) struct HashAggregateAccumulator { + /// Aggregate expression used to create a fresh accumulator for related + /// hash tables, such as the partial-skip table. + aggregate_expr: Arc, + + /// Arguments to pass to this accumulator. + /// + /// Example: `CORR(x, y)` stores two expressions here, while `SUM(x)` stores one. + arguments: Vec>, + + /// Optional `FILTER` expression for this accumulator. + /// + /// Example: `SUM(x) FILTER (WHERE x > 10)` stores the `x > 10` predicate. + filter: Option>, + + /// Accumulator state for all groups for one aggregate expression. + accumulator: Box, +} + +pub(super) type AggregateAccumulator = HashAggregateAccumulator; + +/// Function used by [`AggregateHashTable::aggregate_batch_inner`] to update one +/// accumulator with one evaluated input batch. +/// +/// Arguments: +/// * accumulator to update. +/// * accumulator's evaluated arguments and optional filter. +/// * one group index per input row, mapping each row to its interned group. +/// * total number of groups currently interned in that buffer, including newly +/// interned groups. +pub(super) type AggregateBatchFn = fn( + &mut AggregateAccumulator, + &EvaluatedAccumulatorArgs, + &[usize], + usize, +) -> Result<()>; + +/// Function used by [`AggregateHashTable::next_output_batch_inner`] to +/// materialize one accumulator's output columns. +/// +/// Arguments: +/// * accumulator to materialize. +/// * group range to emit from the accumulator. +pub(super) type MaterializeAccumulatorFn = + fn(&mut AggregateAccumulator, EmitTo) -> Result>; + +/// Evaluated aggregate arguments and filter for one input batch. +/// +/// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` +/// and `x > 0`. +/// +/// These arrays can be passed directly to [`GroupsAccumulator`]. +pub(super) struct EvaluatedAccumulatorArgs { + /// Evaluated argument arrays. Some aggregate functions take multiple arguments. + pub(super) arguments: Vec, + /// Evaluated filter array, `Some` if the aggregate has a `FILTER` expression. + pub(super) filter: Option, +} + +/// Evaluated all group by keys and accumulator args. +/// +/// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function evaluates +/// `k+1`, `v*v` +pub(super) struct EvaluatedAggregateBatch { + /// One entry per grouping set; each entry contains all evaluated group key + /// arrays for the current input batch. + pub(super) grouping_set_args: Vec>, + + /// Evaluated arguments and filters, one entry per aggregate expression. + pub(super) accumulator_args: Vec, +} + +/// Buffer for the aggregate hash table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits final results during the +/// outputting stage. +/// +/// [`GroupValues`] stores the physical group-key layout, while +/// [`GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct AggregateHashTableBuffer { + /// GROUP BY expressions evaluated for each input batch. + pub(super) group_by: Arc, + + /// Interned group keys. Accumulator state is stored separately by group index. + pub(super) group_values: Box, + + /// Group index for each row in the current input batch. + /// + /// Each value indexes into `group_values`, and the same index is used by every + /// accumulator to update that group's aggregate state. + pub(super) batch_group_indices: Vec, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec, +} + +pub(super) enum AggregateHashTableState { + /// Accumulating input rows into group keys and aggregate state. + Building(AggregateHashTableBuffer), + /// Emitting results directly from group keys and aggregate state. + Outputting(AggregateHashTableBuffer), + /// Materialize all the output results, and then incrementally output in the `OutputtingMaterialized` state. + /// + /// Note this is a temporary solution until the `GroupValues` issue is solved: + /// Issue: + OutputtingMaterialized(MaterializedAggregateOutput), + Done, +} + +/// Fully evaluated aggregate output and the next row offset to emit. +/// +/// Final aggregate evaluation consumes accumulator state, and partial terminal +/// output should not repeatedly renumber group values with `EmitTo::First`. +/// Materialize once and then slice to honor `batch_size` across output polls. +pub(super) struct MaterializedAggregateOutput { + batch: RecordBatch, + offset: usize, +} + +impl MaterializedAggregateOutput { + pub(super) fn new(batch: RecordBatch) -> Self { + Self { batch, offset: 0 } + } + + pub(super) fn next_batch(&mut self, batch_size: usize) -> Option { + debug_assert!(batch_size > 0); + if self.is_exhausted() { + return None; + } + + let length = batch_size.min(self.batch.num_rows() - self.offset); + let batch = self.batch.slice(self.offset, length); + self.offset += length; + Some(batch) + } + + pub(super) fn is_exhausted(&self) -> bool { + self.offset >= self.batch.num_rows() + } + + pub(super) fn memory_size(&self) -> usize { + self.batch.get_array_memory_size() + } +} + +impl HashAggregateAccumulator { + pub(super) fn new( + aggregate_expr: Arc, + arguments: Vec>, + filter: Option>, + accumulator: Box, + ) -> Self { + Self { + aggregate_expr, + arguments, + filter, + accumulator, + } + } + + /// Construct a new accumulator with the same definition, but with empty internal + /// state buffers (empty [`GroupsAccumulator`]). + pub(super) fn empty_like(&self) -> Result { + let accumulator = create_group_accumulator(&self.aggregate_expr)?; + Ok(Self::new( + Arc::clone(&self.aggregate_expr), + self.arguments.clone(), + self.filter.clone(), + accumulator, + )) + } + + /// Evaluate aggregate arguments and filter for one input batch. + /// + /// For example, `AVG(x + 1) FILTER (WHERE x > 0)` evaluates both `x + 1` + /// and `x > 0`. + /// + /// These arrays can be passed directly to [`GroupsAccumulator`] next. + pub(super) fn evaluate_acc_args( + &self, + batch: &RecordBatch, + ) -> Result { + let arguments = self + .arguments + .iter() + .map(|expr| { + expr.evaluate(batch) + .and_then(|value| value.into_array(batch.num_rows())) + }) + .collect::>()?; + + let filter = self + .filter + .as_ref() + .map(|filter| { + filter + .evaluate(batch) + .and_then(|value| value.into_array(batch.num_rows())) + }) + .transpose()?; + + Ok(EvaluatedAccumulatorArgs { arguments, filter }) + } + + pub(super) fn size(&self) -> usize { + self.accumulator.size() + } + + pub(super) fn update_batch( + &mut self, + values: &EvaluatedAccumulatorArgs, + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + let filter = values.filter.as_ref().map(|filter| filter.as_boolean()); + self.accumulator.update_batch( + &values.arguments, + group_indices, + filter, + total_num_groups, + ) + } + + pub(super) fn merge_batch( + &mut self, + values: &EvaluatedAccumulatorArgs, + group_indices: &[usize], + total_num_groups: usize, + ) -> Result<()> { + debug_assert!(values.filter.is_none()); + self.accumulator + .merge_batch(&values.arguments, group_indices, total_num_groups) + } + + /// Evaluating final aggregate results according to `EmitTo`, and reset inner + /// states. (e.g. after `evaluate(EmitTo::All)`, it returns all accumulated groups + /// , and clear the inner buffers) + pub(super) fn evaluate(&mut self, emit_to: EmitTo) -> Result { + self.accumulator.evaluate(emit_to) + } + + pub(super) fn evaluate_to_columns( + &mut self, + emit_to: EmitTo, + ) -> Result> { + Ok(vec![self.evaluate(emit_to)?]) + } + + /// Evaluating partial aggregate results according to `EmitTo`, and reset inner + /// states. (e.g. after `state(EmitTo::All)`, it returns all accumulated groups + /// , and clear the inner buffers) + pub(super) fn state(&mut self, emit_to: EmitTo) -> Result> { + self.accumulator.state(emit_to) + } + + pub(super) fn convert_to_state( + &mut self, + values: &EvaluatedAccumulatorArgs, + ) -> Result> { + let opt_filter = values.filter.as_ref().map(|filter| filter.as_boolean()); + self.accumulator + .convert_to_state(&values.arguments, opt_filter) + } + + pub(super) fn null_arguments( + &self, + input_schema: &SchemaRef, + ) -> Result> { + self.arguments + .iter() + .map(|expr| { + let data_type = expr.data_type(input_schema)?; + Ok(new_null_array(&data_type, 1)) + }) + .collect() + } +} + +impl AggregateHashTableState { + pub(super) fn building(&self) -> &AggregateHashTableBuffer { + let Self::Building(state) = self else { + unreachable!("hash aggregate table is not building") + }; + state + } + + pub(super) fn building_mut(&mut self) -> &mut AggregateHashTableBuffer { + let Self::Building(state) = self else { + unreachable!("hash aggregate table is not building") + }; + state + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use arrow::array::{Array, Int32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + + use super::*; + + #[test] + fn materialized_aggregate_output_slices_batches_until_exhausted() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "group_col", + DataType::Int32, + false, + )])); + let batch = RecordBatch::try_new( + schema, + vec![Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5]))], + )?; + let mut output = MaterializedAggregateOutput::new(batch); + + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![1, 2]); + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![3, 4]); + assert_eq!(int32_values(&output.next_batch(2).unwrap(), 0), vec![5]); + assert!(output.next_batch(2).is_none()); + assert!(output.is_exhausted()); + + Ok(()) + } + + fn int32_values(batch: &RecordBatch, column: usize) -> Vec { + let array = batch + .column(column) + .as_any() + .downcast_ref::() + .unwrap(); + (0..array.len()).map(|idx| array.value(idx)).collect() + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs new file mode 100644 index 0000000000000..2293e7b1b8e89 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/common_ordered.rs @@ -0,0 +1,410 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Common utilities for aggregate tables used in aggregations that inputs are ordered +//! by the groups. + +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_common::assert_or_internal_err; +use datafusion_execution::memory_pool::proxy::VecAllocExt; +use datafusion_expr::EmitTo; + +use crate::InputOrderMode; +use crate::PhysicalExpr; +use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; +use crate::aggregates::grouped_hash_stream::create_group_accumulator; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, aggregate_expressions, + evaluate_group_by, +}; + +use super::common::{AggregateAccumulator, EvaluatedAggregateBatch}; + +/// Aggregate table shared by the ordered partial and final paths. +/// +/// # Ordering optimization +/// +/// The table consumes input batches while `GroupOrdering` tracks which groups +/// are proven complete. Completed groups can be emitted before the input stream +/// ends, which keeps memory bounded by the active ordered key range. +/// +/// # Partial and final variant difference +/// +/// The partial and final aggregate tables implement the two stages of grouped +/// aggregation. See +/// [`OrderedPartialAggregateStream`](crate::aggregates::ordered_partial_stream::OrderedPartialAggregateStream) +/// for the high-level plan shape. +/// +/// Example: `AVG(v) FILTER (WHERE v>0) GROUP BY k` +/// +/// Partial table ([`AggregateMode::Partial`], with optional filter from query): +/// - Input rows: `k, v` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, sum(v), count(v)` +/// +/// Final table ([`AggregateMode::Final`], no filters): +/// - Input rows: `k, sum(v), count(v)` +/// - Table stores: `k, sum(v), count(v)` +/// - Output schema: `k, avg(v)` +/// +/// # Marker Type +/// +/// `OrderedAggrMode` selects the aggregate semantics. For example, +/// `OrderedAggregateTable::::new(...)` consumes raw rows +/// and emits partial states, while +/// `OrderedAggregateTable::::new_with_input_order(...)` +/// consumes partial states and emits final values. +/// +/// Shared methods live on `impl`; partial/final behavior lives on +/// marker-specific impls. +pub(in crate::aggregates) struct OrderedAggregateTable { + /// Output schema: group columns followed by aggregate state or final values. + pub(super) output_schema: SchemaRef, + + /// Intermediate-state schema used when memory pressure requires the table + /// to pass through or spill its current state. + pub(super) state_schema: SchemaRef, + + /// Maximum rows per emitted output batch, from config `batch_size`. + pub(super) batch_size: usize, + + /// Grouping and accumulator-specific timing metrics. + pub(super) group_by_metrics: GroupByMetrics, + + /// Group keys, ordering state, and accumulator states. + pub(super) buffer: OrderedAggregateTableBuffer, + + _mode: PhantomData, +} + +/// Buffer for the ordered aggregate table's group keys and accumulator states. +/// +/// It accumulates input during aggregation and emits output rows as soon as the +/// input ordering proves those groups are complete. +/// +/// [`GroupOrdering`] tracks when and how to do early emit. +/// [`GroupValues`] stores the physical group-key layout, while +/// [`datafusion_expr::GroupsAccumulator`] stores per-group aggregate state. +pub(super) struct OrderedAggregateTableBuffer { + /// GROUP BY expressions evaluated against input batches. + pub(super) group_by: Arc, + + /// Tracks how far ordered input allows this table to drain safely. + pub(super) group_ordering: GroupOrdering, + + /// Interned group keys, in the same group-id order used by accumulators. + pub(super) group_values: Box, + + /// Scratch group id vector for the current input batch. + pub(super) group_indices: Vec, + + /// One item per aggregate expression. + /// + /// Example: `COUNT(x), SUM(y)` creates two items. Each item owns the input + /// expressions, optional filter, and accumulator state for all groups. + pub(super) accumulators: Vec, +} + +/// Methods shared by all aggregate modes +impl OrderedAggregateTable { + #[expect( + clippy::too_many_arguments, + reason = "keeps ordered partial and final table construction explicit" + )] + pub(super) fn new_for_mode( + agg: &AggregateExec, + input_schema: &SchemaRef, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + input_order_mode: &InputOrderMode, + aggregate_mode: &AggregateMode, + filters: Vec>>, + group_by_metrics: GroupByMetrics, + ) -> Result { + assert_or_internal_err!( + batch_size > 0, + "OrderedAggregateTable requires config batch_size >= 1" + ); + + let group_ordering = GroupOrdering::try_new(input_order_mode)?; + let group_schema = agg.group_by.group_schema(input_schema)?; + let group_values = new_group_values(group_schema, &group_ordering)?; + let aggregate_arguments = aggregate_expressions( + &agg.aggr_expr, + aggregate_mode, + agg.group_by.num_group_exprs(), + )?; + let accumulators = agg + .aggr_expr + .iter() + .zip(aggregate_arguments) + .zip(filters) + .map(|((agg_expr, arguments), filter)| { + let accumulator = create_group_accumulator(agg_expr)?; + Ok(AggregateAccumulator::new( + Arc::clone(agg_expr), + arguments, + filter, + accumulator, + )) + }) + .collect::>()?; + + Ok(Self { + output_schema, + state_schema, + batch_size, + group_by_metrics, + buffer: OrderedAggregateTableBuffer { + group_by: Arc::clone(&agg.group_by), + group_ordering, + group_values, + group_indices: vec![], + accumulators, + }, + _mode: PhantomData, + }) + } + + /// Evaluates all group by keys and accumulator args. + /// + /// e.g., `select k+1, sum(v*v) from t group by (k+1)`, this function + /// evaluates `k+1`, `v*v`. + pub(super) fn evaluate_batch( + &self, + batch: &RecordBatch, + ) -> Result { + let timer = self.group_by_metrics.time_calculating_group_ids.timer(); + let grouping_set_args = evaluate_group_by(&self.buffer.group_by, batch)?; + drop(timer); + + let timer = self.group_by_metrics.aggregate_arguments_time.timer(); + let accumulator_args = self + .buffer + .accumulators + .iter() + .map(|acc| acc.evaluate_acc_args(batch)) + .collect::>>()?; + drop(timer); + + Ok(EvaluatedAggregateBatch { + grouping_set_args, + accumulator_args, + }) + } + + /// Called after the input stream is exhausted and the last batch has been + /// aggregated. + /// + /// Updates the internal `GroupOrdering` so it can continue emitting until + /// the buffer is empty. + pub(in crate::aggregates) fn input_done(&mut self) { + self.buffer.group_ordering.input_done(); + } + + /// Returns the ordering state used to decide how memory pressure is handled. + pub(in crate::aggregates) fn group_ordering(&self) -> &GroupOrdering { + &self.buffer.group_ordering + } + + /// Number of groups currently buffered. + pub(in crate::aggregates) fn num_groups(&self) -> usize { + self.buffer.group_values.len() + } + + /// Check if there is zero groups accumulated so far. + pub(in crate::aggregates) fn is_empty(&self) -> bool { + self.num_groups() == 0 + } + + /// All internal buffer's memory size. + pub(in crate::aggregates) fn memory_size(&self) -> usize { + self.buffer + .accumulators + .iter() + .map(|acc| acc.size()) + .sum::() + + self.buffer.group_values.size() + + self.buffer.group_ordering.size() + + self.buffer.group_indices.allocated_size() + } + + pub(in crate::aggregates) fn group_by_metrics(&self) -> GroupByMetrics { + self.group_by_metrics.clone() + } + + /// Takes every intermediate aggregate state and resets the table so it can + /// continue with a new ordered input segment. + /// + /// Unlike normal ordered emission, this operation is allowed to take the + /// active (incomplete) groups. Partial aggregation can pass those states to + /// its final stage, while final aggregation sorts and spills them before + /// replay. + pub(in crate::aggregates) fn take_state_batch( + &mut self, + ) -> Result> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let mut output = self.buffer.group_values.emit(EmitTo::All)?; + for acc in &mut self.buffer.accumulators { + output.extend(acc.state(EmitTo::All)?); + } + + let batch = RecordBatch::try_new(Arc::clone(&self.state_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + // `emit(EmitTo::All)` resets accumulator state. Explicitly shrink the + // key/index buffers too so the memory reservation can be released + // before the batch is passed downstream or sorted for spilling. + self.buffer.group_values.clear_shrink(0); + self.buffer.group_indices.clear(); + self.buffer.group_indices.shrink_to_fit(); + self.buffer.group_ordering.reset(); + + Ok(Some(batch)) + } + + /// Returns the [`EmitTo`], clamped to the specified batch size + /// + /// Returns `(emit_to, should_remove_groups)`, where `emit_to` is the number + /// of groups to emit from `GroupValues` / accumulators, and + /// `should_remove_groups` indicates whether `GroupOrdering` must also shift + /// its tracked indexes. + pub(super) fn clamp_emit_to( + &self, + group_count: usize, + emit_to: EmitTo, + ) -> (EmitTo, bool) { + match emit_to { + EmitTo::First(n) => (EmitTo::First(n.min(self.batch_size)), true), + EmitTo::All if group_count <= self.batch_size => (EmitTo::All, false), + EmitTo::All => (EmitTo::First(self.batch_size), false), + } + } + /// Aggregates one evaluated input batch. + /// + /// This common utility is used by ordered partial and ordered final aggregation. + /// + /// # Argument: `is_final` + /// + /// - `true`: merge partial aggregate states for final aggregation. + /// - `false`: update aggregate states from raw input for partial aggregation. + pub(super) fn aggregate_evaluated_batch( + &mut self, + evaluated_batch: &EvaluatedAggregateBatch, + is_final: bool, + ) -> Result<()> { + for group_values in &evaluated_batch.grouping_set_args { + let starting_num_groups = self.buffer.group_values.len(); + self.buffer + .group_values + .intern(group_values, &mut self.buffer.group_indices)?; + let total_num_groups = self.buffer.group_values.len(); + if total_num_groups > starting_num_groups { + self.buffer.group_ordering.new_groups( + group_values, + &self.buffer.group_indices, + total_num_groups, + )?; + } + + let timer = self.group_by_metrics.aggregation_time.timer(); + for (acc, values) in self + .buffer + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + if is_final { + acc.merge_batch( + values, + &self.buffer.group_indices, + total_num_groups, + )?; + } else { + acc.update_batch( + values, + &self.buffer.group_indices, + total_num_groups, + )?; + } + } + drop(timer); + } + + Ok(()) + } + + /// Emits groups allowed by `GroupOrdering`, leaving only the current + /// unfinished ordered-key range buffered. + /// + /// This common utility is used by ordered partial and ordered final aggregation. + /// + /// # Argument: `is_final` + /// + /// - `true`: output final aggregate values. + /// - `false`: output partial accumulator states. + pub(super) fn next_output_batch_for_mode( + &mut self, + is_final: bool, + ) -> Result> { + if self.buffer.group_values.is_empty() { + return Ok(None); + } + + let Some(emit_to) = self.buffer.group_ordering.emit_to() else { + return Ok(None); + }; + let (emit_to, should_remove_groups) = + self.clamp_emit_to(self.buffer.group_values.len(), emit_to); + + let timer = self.group_by_metrics.emitting_time.timer(); + let mut output = self.buffer.group_values.emit(emit_to)?; + if should_remove_groups { + match emit_to { + EmitTo::First(n) => self.buffer.group_ordering.remove_groups(n), + // `EmitTo::All` is only used after `input_done`, when all + // buffered groups are known complete and the ordering state is + // no longer needed. + EmitTo::All => {} + } + } + + for acc in &mut self.buffer.accumulators { + if is_final { + output.push(acc.evaluate(emit_to)?); + } else { + output.extend(acc.state(emit_to)?); + } + } + drop(timer); + + let batch = RecordBatch::try_new(Arc::clone(&self.output_schema), output)?; + debug_assert!(batch.num_rows() > 0); + + Ok(Some(batch)) + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs new file mode 100644 index 0000000000000..b80e15d7f8345 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/final_table.rs @@ -0,0 +1,77 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; + +use super::common::{AggregateHashTable, FinalMarker, HashAggregateAccumulator}; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +impl AggregateHashTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_with_filters( + agg, + partition, + output_schema, + Arc::clone(&agg.input().schema()), + batch_size, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Emits the next batch of aggregated group keys and final aggregate values. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns) + } + + /// Final aggregation consumes partial aggregate states and merges them into + /// the table's partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs new file mode 100644 index 0000000000000..2c7ec01654a63 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/mod.rs @@ -0,0 +1,31 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +mod common; +mod common_ordered; +mod final_table; +mod ordered_final_table; +mod ordered_partial_table; +mod partial_reduce_table; +mod partial_table; +mod single_table; + +pub(super) use common::{ + AggregateHashTable, FinalMarker, PartialMarker, PartialReduceMarker, + PartialSkipMarker, SingleMarker, +}; +pub(super) use common_ordered::OrderedAggregateTable; diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs new file mode 100644 index 0000000000000..fd064ebffec12 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_final_table.rs @@ -0,0 +1,85 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Aggregate table for final aggregation when partial-state input is ordered. +//! +//! See comments in [`super::ordered_partial_table`] for details. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::InputOrderMode; +use crate::aggregates::aggregate_hash_table::FinalMarker; +use crate::aggregates::group_values::GroupByMetrics; +use crate::aggregates::{AggregateExec, AggregateMode}; + +use super::common_ordered::OrderedAggregateTable; + +/// Implementation specific to final aggregation, where the table stores partial +/// aggregate states and the input rows are also partial states. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, sum(x), count(x)` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable { + pub(in crate::aggregates) fn new_with_input_order( + agg: &AggregateExec, + input_schema: &SchemaRef, + output_schema: SchemaRef, + batch_size: usize, + input_order_mode: &InputOrderMode, + group_by_metrics: GroupByMetrics, + ) -> Result { + Self::new_for_mode( + agg, + input_schema, + output_schema, + Arc::clone(input_schema), + batch_size, + input_order_mode, + &AggregateMode::Final, + vec![None; agg.aggr_expr.len()], + group_by_metrics, + ) + } + + /// Merges one partial-state input batch and updates ordering information for + /// any newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + // `PhysicalGroupBy::as_final()` removes grouping sets while planning + // final aggregation, so final ordered aggregation sees one grouping. + debug_assert_eq!(evaluated_batch.grouping_set_args.len(), 1); + self.aggregate_evaluated_batch(&evaluated_batch, true) + } + + /// See comments in `ordered_partial_stream::next_output_batch` + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_for_mode(true) + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs new file mode 100644 index 0000000000000..a04e4dda8fb39 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/ordered_partial_table.rs @@ -0,0 +1,106 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Aggregate table for partial aggregation when input is ordered by group keys. +//! +//! See the [`super::common_ordered`] comments for the high-level ideas. +//! +//! This operator handles input that is ordered by group keys: +//! - Fully ordered: `GROUP BY a, b`, input is `ORDER BY a, b` +//! - Partially ordered: `GROUP BY a, b`, input is `ORDER BY a` +//! +//! When a group key combination is exhausted, this table eagerly flushes the +//! completed groups to improve memory efficiency. +//! +//! The implementation is separated from other aggregate tables because this +//! execution path is likely to be optimized further in the future. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::{ + AggregateExec, AggregateMode, aggregate_hash_table::PartialMarker, + group_values::GroupByMetrics, +}; + +use super::common_ordered::OrderedAggregateTable; + +/// Implementation specific to partial aggregation, where the table stores +/// partial aggregate states and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, x` +/// +/// See comments at [`OrderedAggregateTable`] for details. +impl OrderedAggregateTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + let input_schema = agg.input().schema(); + let state_schema = Arc::clone(&output_schema); + let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + Self::new_for_mode( + agg, + &input_schema, + output_schema, + state_schema, + batch_size, + &agg.input_order_mode, + &AggregateMode::Partial, + agg.filter_expr.iter().cloned().collect(), + group_by_metrics, + ) + } + + /// Aggregates one raw input batch and updates ordering information for any + /// newly observed groups. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + let evaluated_batch = self.evaluate_batch(batch)?; + self.aggregate_evaluated_batch(&evaluated_batch, false) + } + + /// Emits the next batch of partial state rows for groups proven complete by + /// the input ordering. + /// + /// For example, when the query is `GROUP BY a` and the input is ordered by + /// `a`, seeing a latest input row with `a = 3` means all groups with `a < 3` + /// are complete and safe to emit. + /// + /// Key steps: + /// 1. Ask `group_ordering` to decide how many groups can be emitted eagerly. + /// 2. Remove the emitted groups from `group_ordering`, `GroupValues`, and + /// all `GroupsAccumulator`s. + /// + /// This may output small batches. Avoiding tiny batches is left to future + /// ordered-aggregation optimizations. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_for_mode(false) + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs new file mode 100644 index 0000000000000..4dfd6a74d18b8 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_reduce_table.rs @@ -0,0 +1,71 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; + +use super::common::{AggregateHashTable, HashAggregateAccumulator, PartialReduceMarker}; + +/// Methods specific to the aggregate hash table used in the partial-reduce stage. +impl AggregateHashTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_with_filters( + agg, + partition, + Arc::clone(&output_schema), + output_schema, + batch_size, + vec![None; agg.aggr_expr.len()], + ) + } + + /// Emits the next batch of aggregated group keys and aggregate states. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner(HashAggregateAccumulator::state) + } + + /// Partial-reduce aggregation consumes partial aggregate states and merges + /// them into the table's partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner(batch, HashAggregateAccumulator::merge_batch) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs new file mode 100644 index 0000000000000..a64fd32536eeb --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/partial_table.rs @@ -0,0 +1,219 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; +use std::marker::PhantomData; +use std::sync::Arc; + +use arrow::array::{ArrayRef, BooleanArray, new_null_array}; +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{Result, assert_eq_or_internal_err}; + +use crate::aggregates::group_values::new_group_values; +use crate::aggregates::order::GroupOrdering; +use crate::aggregates::{AggregateExec, group_id_array, max_duplicate_ordinal}; + +use super::common::{ + AggregateHashTable, AggregateHashTableBuffer, AggregateHashTableState, + EvaluatedAccumulatorArgs, HashAggregateAccumulator, PartialMarker, PartialSkipMarker, +}; + +/// Implementation specific to partial aggregation, where the table stores +/// partial aggregate states and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, sum(x), count(x)` +/// - Input rows: `k, x` +impl AggregateHashTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_with_filters( + agg, + partition, + Arc::clone(&output_schema), + output_schema, + batch_size, + agg.filter_expr.iter().cloned().collect(), + ) + } + + /// Emits the next batch of aggregated group keys and aggregate states. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner(HashAggregateAccumulator::state) + } + + /// In skip-partial-aggregation optimization, when a decision has been made to skip + /// partial stage, build a typed hash table only for aggregation state conversion + /// row-by-row. + pub(in crate::aggregates) fn partial_skip_table( + &self, + ) -> Result> { + let state = self.state.building(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let group_values = new_group_values(group_schema, &GroupOrdering::None)?; + let accumulators = state + .accumulators + .iter() + .map(HashAggregateAccumulator::empty_like) + .collect::>>()?; + + Ok(AggregateHashTable { + group_by_metrics: self.group_by_metrics.clone(), + input_schema: Arc::clone(&self.input_schema), + output_schema: Arc::clone(&self.output_schema), + state_schema: Arc::clone(&self.state_schema), + batch_size: self.batch_size, + state: AggregateHashTableState::Building(AggregateHashTableBuffer { + group_by: Arc::clone(&state.group_by), + group_values, + batch_group_indices: Default::default(), + accumulators, + }), + _mode: PhantomData, + }) + } + + /// Partial aggregation consumes raw input rows and updates the table's + /// partial-state accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.init_empty_grouping_sets()?; + self.start_outputting(); + Ok(()) + } + + /// Creates the required empty grouping-set rows when the input is empty. + /// + /// For example, this query must still produce one grand-total group even if + /// `t` has no rows: + /// + /// ```sql + /// SELECT COUNT(v) + /// FROM t + /// GROUP BY GROUPING SETS (()); + /// ``` + /// + /// The synthetic row is filtered out before accumulator update so aggregates + /// see the same state they would see for an empty input, rather than a real + /// null-valued row. + fn init_empty_grouping_sets(&mut self) -> Result<()> { + let state = self.state.building_mut(); + if !state.group_by.has_grouping_set() || !state.group_values.is_empty() { + return Ok(()); + } + + let max_ordinal = max_duplicate_ordinal(state.group_by.groups()); + let mut ordinals: HashMap<&[bool], usize> = HashMap::new(); + let group_schema = state.group_by.group_schema(&self.input_schema)?; + let n_expr = state.group_by.expr().len(); + let mut any_interned = false; + + for group in state.group_by.groups() { + let ordinal = { + let entry = ordinals.entry(group.as_slice()).or_insert(0); + let ordinal = *entry; + *entry += 1; + ordinal + }; + + if !group.iter().all(|&is_null| is_null) { + continue; + } + + let mut cols: Vec = group_schema + .fields() + .iter() + .take(n_expr) + .map(|field| new_null_array(field.data_type(), 1)) + .collect(); + cols.push(group_id_array(group, ordinal, max_ordinal, 1)?); + + state + .group_values + .intern(&cols, &mut state.batch_group_indices)?; + any_interned = true; + } + + if any_interned { + let total_groups = state.group_values.len(); + let false_filter = BooleanArray::from(vec![false]); + for acc in state.accumulators.iter_mut() { + let null_args = acc.null_arguments(&self.input_schema)?; + let values = EvaluatedAccumulatorArgs { + arguments: null_args, + filter: Some(Arc::new(false_filter.clone())), + }; + acc.update_batch(&values, &[0], total_groups)?; + } + } + + Ok(()) + } +} + +impl AggregateHashTable { + pub(in crate::aggregates) fn convert_batch_to_state( + &mut self, + batch: &RecordBatch, + ) -> Result { + let evaluated_batch = self.evaluate_batch(batch)?; + + assert_eq_or_internal_err!( + evaluated_batch.grouping_set_args.len(), + 1, + "group_values expected to have single element" + ); + let mut output = evaluated_batch + .grouping_set_args + .into_iter() + .next() + .unwrap_or_default(); + + let state = self.state.building_mut(); + for (acc, values) in state + .accumulators + .iter_mut() + .zip(evaluated_batch.accumulator_args.iter()) + { + output.extend(acc.convert_to_state(values)?); + } + + Ok(RecordBatch::try_new( + Arc::clone(&self.output_schema), + output, + )?) + } +} diff --git a/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs new file mode 100644 index 0000000000000..56d601c793206 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/aggregate_hash_table/single_table.rs @@ -0,0 +1,76 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; + +use crate::aggregates::AggregateExec; + +use super::common::{AggregateHashTable, HashAggregateAccumulator, SingleMarker}; + +/// Implementation specific to single aggregation, where the table stores final +/// aggregate values and the input rows are raw rows. +/// +/// Example: `AVG(x) GROUP BY k` +/// +/// - Aggregate table stores: `k, avg(x)` +/// - Input rows: `k, x` +impl AggregateHashTable { + pub(in crate::aggregates) fn new( + agg: &AggregateExec, + partition: usize, + output_schema: SchemaRef, + state_schema: SchemaRef, + batch_size: usize, + ) -> Result { + Self::new_with_filters( + agg, + partition, + output_schema, + state_schema, + batch_size, + agg.filter_expr.iter().cloned().collect(), + ) + } + + /// Emits the next batch of aggregated group keys and final aggregate values. + /// + /// The output batch size is determined by `self.batch_size`. + /// + /// Returns `Some(batch)` for each emitted batch, `None` when output is + /// exhausted, and an internal error if polled in the `Building` state. + pub(in crate::aggregates) fn next_output_batch( + &mut self, + ) -> Result> { + self.next_output_batch_inner(HashAggregateAccumulator::evaluate_to_columns) + } + + /// Single aggregation consumes raw input rows and updates the table's + /// final-value accumulators. + pub(in crate::aggregates) fn aggregate_batch( + &mut self, + batch: &RecordBatch, + ) -> Result<()> { + self.aggregate_batch_inner(batch, HashAggregateAccumulator::update_batch) + } + + pub(in crate::aggregates) fn start_output(&mut self) -> Result<()> { + self.start_outputting(); + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/no_grouping.rs b/datafusion/physical-plan/src/aggregates/aggregate_stream.rs similarity index 100% rename from datafusion/physical-plan/src/aggregates/no_grouping.rs rename to datafusion/physical-plan/src/aggregates/aggregate_stream.rs diff --git a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs index b6c32204e85f0..1c6285d793b88 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/metrics.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/metrics.rs @@ -19,6 +19,7 @@ use crate::metrics::{ExecutionPlanMetricsSet, MetricBuilder, Time}; +#[derive(Clone)] pub(crate) struct GroupByMetrics { /// Time spent calculating the group IDs from the evaluated grouping columns. pub(crate) time_calculating_group_ids: Time, @@ -59,6 +60,7 @@ mod tests { use arrow::record_batch::RecordBatch; use datafusion_common::Result; use datafusion_execution::TaskContext; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_functions_aggregate::count::count_udaf; use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_physical_expr::aggregate::AggregateExprBuilder; @@ -135,7 +137,13 @@ mod tests { schema, )?); - let task_ctx = Arc::new(TaskContext::default()); + // This test is for `GroupByMetrics`, which are maintained by + // `GroupedHashAggregateStream`. Use a finite memory pool so the partial + // aggregate does not take the initial-partial stream path. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(10 * 1024 * 1024, 1.0) + .build_arc()?; + let task_ctx = Arc::new(TaskContext::default().with_runtime(runtime)); let _result = collect(Arc::clone(&aggregate_exec) as _, Arc::clone(&task_ctx)).await?; diff --git a/datafusion/physical-plan/src/aggregates/group_values/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/mod.rs index 2f3b1a19e7d73..1101d535311e4 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/mod.rs @@ -31,10 +31,10 @@ use datafusion_expr::EmitTo; pub mod multi_group_by; mod row; +pub use row::GroupValuesRows; mod single_group_by; use datafusion_physical_expr::binary_map::OutputType; use multi_group_by::GroupValuesColumn; -use row::GroupValuesRows; pub(crate) use single_group_by::primitive::HashValue; @@ -99,7 +99,9 @@ pub trait GroupValues: Send { /// assigned. fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()>; - /// Returns the number of bytes of memory used by this [`GroupValues`] + /// Returns the number of bytes of memory used by this [`GroupValues`]. + /// + /// May be expensive; check the implementation before calling on hot paths. fn size(&self) -> usize; /// Returns true if this [`GroupValues`] is empty @@ -130,7 +132,7 @@ pub trait GroupValues: Send { /// /// `GroupColumn`: crate::aggregates::group_values::multi_group_by::GroupColumn /// `GroupValuesColumn`: crate::aggregates::group_values::multi_group_by::GroupValuesColumn -/// `GroupValuesRows`: crate::aggregates::group_values::row::GroupValuesRows +/// `GroupValuesRows`: crate::aggregates::group_values::GroupValuesRows pub fn new_group_values( schema: SchemaRef, group_ordering: &GroupOrdering, diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs index 350ec13712652..c83b1da4049bc 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes.rs @@ -16,7 +16,7 @@ // under the License. use crate::aggregates::group_values::multi_group_by::{ - GroupColumn, Nulls, nulls_equal_to, split_vec_min_alloc, + GroupColumn, Nulls, nulls_equal_to, }; use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; use arrow::array::{ @@ -26,6 +26,7 @@ use arrow::array::{ use arrow::buffer::{OffsetBuffer, ScalarBuffer}; use arrow::datatypes::{ByteArrayType, DataType, GenericBinaryType}; use datafusion_common::utils::proxy::VecAllocExt; +use datafusion_common::utils::split_vec_min_alloc; use datafusion_common::{Result, exec_datafusion_err}; use datafusion_physical_expr_common::binary_map::{INITIAL_BUFFER_CAPACITY, OutputType}; use std::mem::size_of; diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs index 9267cf4f27f35..8625772e2c995 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/bytes_view.rs @@ -21,11 +21,11 @@ use crate::aggregates::group_values::multi_group_by::{ use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; use arrow::array::{ Array, ArrayRef, AsArray, BooleanBufferBuilder, ByteView, GenericByteViewArray, - make_view, }; use arrow::buffer::{Buffer, ScalarBuffer}; use arrow::datatypes::ByteViewType; use datafusion_common::Result; +use datafusion_common::utils::split_vec_min_alloc; use std::marker::PhantomData; use std::mem::{replace, size_of}; use std::sync::Arc; @@ -144,7 +144,11 @@ impl ByteViewGroupValueBuilder { } } - fn vectorized_append_inner(&mut self, array: &ArrayRef, rows: &[usize]) { + fn vectorized_append_inner( + &mut self, + array: &ArrayRef, + rows: &[usize], + ) -> Result<()> { let arr = array.as_byte_view::(); let null_count = array.null_count(); let num_rows = array.len(); @@ -165,8 +169,24 @@ impl ByteViewGroupValueBuilder { Nulls::None => { self.nulls.append_n(rows.len(), false); - for &row in rows { - self.do_append_val_inner(arr, row); + if arr.data_buffers().is_empty() { + // Fast path: all strings are inline (≤12 bytes). + // The input array's u128 views are already in the correct format; + // copy them directly instead of going through value() → make_view(). + self.views.extend(rows.iter().map(|&row| arr.views()[row])); + } else { + // Slow path: some strings may be non-inline (>12 bytes). + // Pre-reserve and delegate to do_append_val_inner which + // reads raw views directly and reuses source prefixes. + self.views.try_reserve(rows.len()).map_err(|e| { + datafusion_common::exec_datafusion_err!( + "failed to reserve {0} views: {e}", + rows.len() + ) + })?; + for &row in rows { + self.do_append_val_inner(arr, row); + } } } @@ -176,31 +196,40 @@ impl ByteViewGroupValueBuilder { self.views.resize(new_len, 0); } } + Ok(()) } fn do_append_val_inner(&mut self, array: &GenericByteViewArray, row: usize) where B: ByteViewType, { - let value: &[u8] = array.value(row).as_ref(); + // SAFETY: the caller ensures `row` is valid + let view = unsafe { *array.views().get_unchecked(row) }; + let len = view as u32; - let value_len = value.len(); - let view = if value_len <= 12 { - make_view(value, 0, 0) + if len <= 12 { + // Inline value: the view is already self-contained, push as-is. + self.views.push(view); } else { - // Ensure big enough block to hold the value firstly - self.ensure_in_progress_big_enough(value_len); - - // Append value - let buffer_index = self.completed.len(); - let offset = self.in_progress.len(); - self.in_progress.extend_from_slice(value); - - make_view(value, buffer_index as u32, offset as u32) - }; - - // Append view - self.views.push(view); + // Non-inline value: copy the buffer data and construct a new view + // that points into our own buffers, reusing the source prefix. + let src = ByteView::from(view); + self.ensure_in_progress_big_enough(len as usize); + let new_buffer_index = self.completed.len() as u32; + let new_offset = self.in_progress.len() as u32; + let src_buf = &array.data_buffers()[src.buffer_index as usize]; + self.in_progress.extend_from_slice( + &src_buf[src.offset as usize..(src.offset + src.length) as usize], + ); + let new_view = ByteView { + length: src.length, + prefix: src.prefix, + buffer_index: new_buffer_index, + offset: new_offset, + } + .as_u128(); + self.views.push(new_view); + } } fn ensure_in_progress_big_enough(&mut self, value_len: usize) { @@ -363,7 +392,7 @@ impl ByteViewGroupValueBuilder { // // - Shift the `buffer index` of remaining non-inlined `views` // - let first_n_views = self.views.drain(0..n).collect::>(); + let first_n_views = split_vec_min_alloc(&mut self.views, n); let last_non_inlined_view = first_n_views .iter() @@ -547,8 +576,7 @@ impl GroupColumn for ByteViewGroupValueBuilder { } fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { - self.vectorized_append_inner(array, rows); - Ok(()) + self.vectorized_append_inner(array, rows) } fn len(&self) -> usize { diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs new file mode 100644 index 0000000000000..589083c8f7ce2 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/fixed_size_binary.rs @@ -0,0 +1,515 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use crate::aggregates::group_values::multi_group_by::{ + GroupColumn, Nulls, nulls_equal_to, +}; +use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBufferBuilder, FixedSizeBinaryArray, +}; +use arrow::buffer::{Buffer, NullBuffer}; +use datafusion_common::utils::proxy::VecAllocExt; +use datafusion_common::utils::split_vec_min_alloc; +use datafusion_common::{Result, exec_datafusion_err}; +use std::sync::Arc; + +/// An implementation of [`GroupColumn`] for `FixedSizeBinary` values +/// +/// Stores the group values in a single flat buffer, `byte_width` bytes per +/// value, in a way that allows: +/// +/// 1. Efficient comparison of incoming rows to existing rows +/// 2. Efficient construction of the final output array (the buffer is handed +/// to [`FixedSizeBinaryArray`] as-is, no offsets needed) +/// +/// Null values occupy `byte_width` zeroed bytes in the buffer so that the +/// value of row `i` is always stored at `i * byte_width..(i + 1) * byte_width`. +pub struct FixedSizeBinaryGroupValueBuilder { + /// The width in bytes of each value, from `DataType::FixedSizeBinary` + byte_width: usize, + /// The flattened group values, `byte_width` bytes per value + buffer: Vec, + /// The number of group values stored + /// + /// Tracked explicitly rather than derived from `buffer.len()` because + /// `byte_width` may be `0` + len: usize, + /// Null state (null rows still occupy `byte_width` bytes in `buffer`) + nulls: MaybeNullBufferBuilder, +} + +impl FixedSizeBinaryGroupValueBuilder { + /// Create a new builder for values of `byte_width` bytes each + /// + /// `byte_width` is the width carried by `DataType::FixedSizeBinary` and + /// must be non-negative (negative widths are rejected by the dispatch in + /// `make_group_column`) + pub fn new(byte_width: i32) -> Self { + debug_assert!(byte_width >= 0); + Self { + byte_width: byte_width as usize, + buffer: Vec::new(), + len: 0, + nulls: MaybeNullBufferBuilder::new(), + } + } + + fn do_append_val_inner(&mut self, array: &FixedSizeBinaryArray, row: usize) { + if array.is_null(row) { + self.nulls.append(true); + // Null rows still occupy `byte_width` (zeroed) bytes in the + // buffer so the value offset stays a function of the row index + self.buffer.resize(self.buffer.len() + self.byte_width, 0); + } else { + self.nulls.append(false); + self.buffer.extend_from_slice(array.value(row)); + } + self.len += 1; + } + + fn do_equal_to_inner( + &self, + lhs_row: usize, + array: &FixedSizeBinaryArray, + rhs_row: usize, + ) -> bool { + let exist_null = self.nulls.is_null(lhs_row); + let input_null = array.is_null(rhs_row); + if let Some(result) = nulls_equal_to(exist_null, input_null) { + return result; + } + // Otherwise, we need to check their values + self.value(lhs_row) == array.value(rhs_row) + } + + /// return the current value of the specified row irrespective of null + /// (null rows store `byte_width` zeroed bytes) + pub fn value(&self, row: usize) -> &[u8] { + let start = row * self.byte_width; + &self.buffer[start..start + self.byte_width] + } + + /// Assemble an output array from `values` + `nulls` parts + /// + /// Uses `try_new_with_len` rather than `try_new` because the length + /// cannot be derived from the values buffer when `byte_width == 0` + fn build_array( + byte_width: usize, + values: Vec, + nulls: Option, + len: usize, + ) -> ArrayRef { + let array = FixedSizeBinaryArray::try_new_with_len( + byte_width as i32, + Buffer::from(values), + nulls, + len, + ) + .expect("buffer, nulls and len kept consistent on append"); + Arc::new(array) + } +} + +impl GroupColumn for FixedSizeBinaryGroupValueBuilder { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + self.do_equal_to_inner(lhs_row, array.as_fixed_size_binary(), rhs_row) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let arr = array.as_fixed_size_binary(); + debug_assert_eq!(arr.value_size(), self.byte_width); + self.do_append_val_inner(arr, row); + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + let array = array.as_fixed_size_binary(); + + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + // Has found not equal to in previous column, don't need to check + if !equal_to_results.get_bit(idx) { + continue; + } + + if !self.do_equal_to_inner(lhs_row, array, rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + let arr = array.as_fixed_size_binary(); + debug_assert_eq!(arr.value_size(), self.byte_width); + + let reserve_bytes = rows.len() * self.byte_width; + self.buffer.try_reserve(reserve_bytes).map_err(|e| { + exec_datafusion_err!("failed to reserve {reserve_bytes} bytes: {e}") + })?; + + let null_count = array.null_count(); + let num_rows = array.len(); + let all_null_or_non_null = if null_count == 0 { + Nulls::None + } else if null_count == num_rows { + Nulls::All + } else { + Nulls::Some + }; + + match all_null_or_non_null { + Nulls::Some => { + for &row in rows { + self.do_append_val_inner(arr, row); + } + } + + Nulls::None => { + self.nulls.append_n(rows.len(), false); + for &row in rows { + self.buffer.extend_from_slice(arr.value(row)); + } + self.len += rows.len(); + } + + Nulls::All => { + self.nulls.append_n(rows.len(), true); + self.buffer + .resize(self.buffer.len() + rows.len() * self.byte_width, 0); + self.len += rows.len(); + } + } + + Ok(()) + } + + fn len(&self) -> usize { + self.len + } + + fn size(&self) -> usize { + self.buffer.allocated_size() + self.nulls.allocated_size() + } + + fn build(self: Box) -> ArrayRef { + let Self { + byte_width, + buffer, + len, + nulls, + } = *self; + + Self::build_array(byte_width, buffer, nulls.build(), len) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + debug_assert!(self.len >= n); + + let null_buffer = self.nulls.take_n(n); + let first_n = split_vec_min_alloc(&mut self.buffer, n * self.byte_width); + self.len -= n; + + Self::build_array(self.byte_width, first_n, null_buffer, n) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use crate::aggregates::group_values::multi_group_by::fixed_size_binary::FixedSizeBinaryGroupValueBuilder; + use arrow::array::{ArrayRef, BooleanBufferBuilder, FixedSizeBinaryArray}; + + use super::GroupColumn; + + fn make_true_buffer(n: usize) -> BooleanBufferBuilder { + let mut buf = BooleanBufferBuilder::new(n); + buf.append_n(n, true); + buf + } + + fn to_vec(buf: &BooleanBufferBuilder) -> Vec { + (0..buf.len()).map(|i| buf.get_bit(i)).collect() + } + + fn make_array(values: Vec>, byte_width: i32) -> ArrayRef { + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values.into_iter(), + byte_width, + ) + .unwrap(), + ) + } + + #[test] + fn test_fixed_size_binary_equal_to() { + let append = |builder: &mut FixedSizeBinaryGroupValueBuilder, + builder_array: &ArrayRef, + append_rows: &[usize]| { + for &index in append_rows { + builder.append_val(builder_array, index).unwrap(); + } + }; + + let equal_to = + |builder: &FixedSizeBinaryGroupValueBuilder, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + let iter = lhs_rows.iter().zip(rhs_rows.iter()); + for (idx, (&lhs_row, &rhs_row)) in iter.enumerate() { + equal_to_results + .set_bit(idx, builder.equal_to(lhs_row, input_array, rhs_row)); + } + }; + + test_fixed_size_binary_equal_to_internal(append, equal_to); + } + + #[test] + fn test_fixed_size_binary_vectorized_equal_to() { + let append = |builder: &mut FixedSizeBinaryGroupValueBuilder, + builder_array: &ArrayRef, + append_rows: &[usize]| { + builder + .vectorized_append(builder_array, append_rows) + .unwrap(); + }; + + let equal_to = + |builder: &FixedSizeBinaryGroupValueBuilder, + lhs_rows: &[usize], + input_array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder| { + builder.vectorized_equal_to( + lhs_rows, + input_array, + rhs_rows, + equal_to_results, + ); + }; + + test_fixed_size_binary_equal_to_internal(append, equal_to); + } + + fn test_fixed_size_binary_equal_to_internal(mut append: A, mut equal_to: E) + where + A: FnMut(&mut FixedSizeBinaryGroupValueBuilder, &ArrayRef, &[usize]), + E: FnMut( + &FixedSizeBinaryGroupValueBuilder, + &[usize], + &ArrayRef, + &[usize], + &mut BooleanBufferBuilder, + ), + { + // Will cover such cases: + // - exist null, input not null + // - exist null, input null; values not equal + // - exist null, input null; values equal + // - exist not null, input null + // - exist not null, input not null; values not equal + // - exist not null, input not null; values equal + + // Define FixedSizeBinaryGroupValueBuilder + let mut builder = FixedSizeBinaryGroupValueBuilder::new(3); + let builder_array = make_array( + vec![ + None, + None, + None, + Some(b"foo".as_slice()), + Some(b"bar".as_slice()), + Some(b"baz".as_slice()), + ], + 3, + ); + append(&mut builder, &builder_array, &[0, 1, 2, 3, 4, 5]); + + // Define input array; the value behind the null at row 3 happens to + // match the existing group value to make sure nulls win over values + let input_array = make_array( + vec![ + Some(b"foo".as_slice()), + None, + None, + None, + Some(b"foo".as_slice()), + Some(b"baz".as_slice()), + ], + 3, + ); + + // Check + let mut equal_to_results = make_true_buffer(builder.len()); + equal_to( + &builder, + &[0, 1, 2, 3, 4, 5], + &input_array, + &[0, 1, 2, 3, 4, 5], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(!results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(!results[3]); + assert!(!results[4]); + assert!(results[5]); + } + + #[test] + fn test_fixed_size_binary_vectorized_operation_special_case() { + // Test the special `all nulls` or `not nulls` input array case + // for vectorized append and equal to + + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + + // All nulls input array + let all_nulls_input_array = make_array(vec![None, None, None, None, None], 2); + builder + .vectorized_append(&all_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_nulls_input_array.len()); + builder.vectorized_equal_to( + &[0, 1, 2, 3, 4], + &all_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + + // All not nulls input array + let all_not_nulls_input_array = make_array( + vec![ + Some(b"v1".as_slice()), + Some(b"v2".as_slice()), + Some(b"v3".as_slice()), + Some(b"v4".as_slice()), + Some(b"v5".as_slice()), + ], + 2, + ); + builder + .vectorized_append(&all_not_nulls_input_array, &[0, 1, 2, 3, 4]) + .unwrap(); + + let mut equal_to_results = make_true_buffer(all_not_nulls_input_array.len()); + builder.vectorized_equal_to( + &[5, 6, 7, 8, 9], + &all_not_nulls_input_array, + &[0, 1, 2, 3, 4], + &mut equal_to_results, + ); + let results = to_vec(&equal_to_results); + + assert!(results[0]); + assert!(results[1]); + assert!(results[2]); + assert!(results[3]); + assert!(results[4]); + } + + #[test] + fn test_fixed_size_binary_take_n() { + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + let array = make_array(vec![Some(b"aa".as_slice()), None], 2); + // aa, null, null + builder.append_val(&array, 0).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 1).unwrap(); + + // (aa, null) remaining: null + let output = builder.take_n(2); + assert_eq!(&output, &array); + assert_eq!(builder.len(), 1); + + // null, aa, null, aa + builder.append_val(&array, 0).unwrap(); + builder.append_val(&array, 1).unwrap(); + builder.append_val(&array, 0).unwrap(); + + // (null, aa) remaining: (null, aa) + let output = builder.take_n(2); + let expected = make_array(vec![None, Some(b"aa".as_slice())], 2); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 2); + + // take the remaining (null, aa) + let output = builder.take_n(2); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 0); + } + + #[test] + fn test_fixed_size_binary_build() { + let mut builder = FixedSizeBinaryGroupValueBuilder::new(2); + let array = make_array( + vec![Some(b"aa".as_slice()), None, Some(b"bb".as_slice())], + 2, + ); + builder.vectorized_append(&array, &[0, 1, 2]).unwrap(); + assert_eq!(builder.len(), 3); + + let output = Box::new(builder).build(); + assert_eq!(&output, &array); + } + + #[test] + fn test_zero_width_fixed_size_binary() { + // A zero byte width is valid per the Arrow spec; the builder must + // track its length without relying on the (empty) values buffer + let mut builder = FixedSizeBinaryGroupValueBuilder::new(0); + let array = make_array(vec![Some(b"".as_slice()), None, Some(b"".as_slice())], 0); + + builder.vectorized_append(&array, &[0, 1, 2]).unwrap(); + assert_eq!(builder.len(), 3); + + // Empty values compare equal, null only equals null + assert!(builder.equal_to(0, &array, 2)); + assert!(builder.equal_to(1, &array, 1)); + assert!(!builder.equal_to(1, &array, 0)); + + let output = builder.take_n(2); + let expected = make_array(vec![Some(b"".as_slice()), None], 0); + assert_eq!(&output, &expected); + assert_eq!(builder.len(), 1); + + let output = Box::new(builder).build(); + let expected = make_array(vec![Some(b"".as_slice())], 0); + assert_eq!(&output, &expected); + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs index f603839bee271..5b474f3bae075 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/mod.rs @@ -20,22 +20,29 @@ mod boolean; mod bytes; pub mod bytes_view; +mod fixed_size_binary; pub mod primitive; +pub mod row_backed; use std::mem::{self, size_of}; use crate::aggregates::group_values::GroupValues; use crate::aggregates::group_values::multi_group_by::{ boolean::BooleanGroupValueBuilder, bytes::ByteGroupValueBuilder, - bytes_view::ByteViewGroupValueBuilder, primitive::PrimitiveGroupValueBuilder, + bytes_view::ByteViewGroupValueBuilder, + fixed_size_binary::FixedSizeBinaryGroupValueBuilder, + primitive::PrimitiveGroupValueBuilder, row_backed::RowsGroupColumn, }; use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; use arrow::compute::cast; use arrow::datatypes::{ - BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Float32Type, - Float64Type, Int8Type, Int16Type, Int32Type, Int64Type, Schema, SchemaRef, - StringViewType, Time32MillisecondType, Time32SecondType, Time64MicrosecondType, - Time64NanosecondType, TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, + BinaryViewType, DataType, Date32Type, Date64Type, Decimal128Type, Decimal256Type, + DurationMicrosecondType, DurationMillisecondType, DurationNanosecondType, + DurationSecondType, Field, Float16Type, Float32Type, Float64Type, Int8Type, + Int16Type, Int32Type, Int64Type, IntervalDayTimeType, IntervalMonthDayNanoType, + IntervalUnit, IntervalYearMonthType, Schema, SchemaRef, StringViewType, + Time32MillisecondType, Time32SecondType, Time64MicrosecondType, Time64NanosecondType, + TimeUnit, TimestampMicrosecondType, TimestampMillisecondType, TimestampNanosecondType, TimestampSecondType, UInt8Type, UInt16Type, UInt32Type, UInt64Type, }; @@ -107,19 +114,6 @@ pub trait GroupColumn: Send + Sync { fn take_n(&mut self, n: usize) -> ArrayRef; } -/// Splits `vec` at `n`, returning the first `n` elements and leaving the -/// remainder in `vec`. Allocates for whichever portion is smaller to minimize -/// peak memory: `drain+collect` when `n <= remaining`, `split_off+replace` -/// when `remaining < n`. -pub(super) fn split_vec_min_alloc(vec: &mut Vec, n: usize) -> Vec { - if n * 2 <= vec.len() { - vec.drain(0..n).collect() - } else { - let remaining = vec.split_off(n); - mem::replace(vec, remaining) - } -} - /// Determines if the nullability of the existing and new input array can be used /// to short-circuit the comparison of the two values. /// @@ -225,7 +219,7 @@ pub struct GroupValuesColumn { /// more general purpose [`GroupValuesRows`]. See the ticket for details: /// /// - /// [`GroupValuesRows`]: crate::aggregates::group_values::row::GroupValuesRows + /// [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows group_values: Vec>, /// reused buffer to store hashes @@ -285,6 +279,7 @@ impl GroupValuesColumn { /// Create a new instance of GroupValuesColumn if supported for the specified schema pub fn try_new(schema: SchemaRef) -> Result { let map = HashTable::with_capacity(0); + let group_values = Self::build_group_columns(&schema)?; Ok(Self { schema, map, @@ -292,12 +287,27 @@ impl GroupValuesColumn { emit_group_index_list_buffer: Vec::new(), vectorized_operation_buffers: VectorizedOperationBuffers::default(), map_size: 0, - group_values: vec![], + group_values, hashes_buffer: Default::default(), random_state: crate::aggregates::AGGREGATION_HASH_SEED, }) } + /// Build one fresh [`GroupColumn`] per field in the schema. + /// + /// Used at construction time (`try_new`) and to repopulate the column + /// vector after operations that drain it (`emit(EmitTo::All)`, + /// `clear_shrink`). Centralising it keeps the post-condition that + /// `self.group_values` always contains exactly one builder per schema + /// field outside of those transient drain points. + fn build_group_columns(schema: &Schema) -> Result>> { + let mut v: Vec> = Vec::with_capacity(schema.fields().len()); + for f in schema.fields().iter() { + v.push(make_group_column(f.as_ref())?); + } + Ok(v) + } + // ======================================================================== // Scalarized intern // ======================================================================== @@ -734,7 +744,7 @@ impl GroupValuesColumn { /// /// The hash collision may be not frequent, so the fallback will indeed hardly happen. /// In most situations, `scalarized_indices` will found to be empty after finishing to - /// preform `vectorized_equal_to`. + /// perform `vectorized_equal_to`. fn scalarized_intern_remaining( &mut self, cols: &[ArrayRef], @@ -911,172 +921,238 @@ macro_rules! instantiate_primitive { }; } -impl GroupValues for GroupValuesColumn { - fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { - if self.group_values.is_empty() { - let mut v = Vec::with_capacity(cols.len()); - - for f in self.schema.fields().iter() { - let nullable = f.is_nullable(); - let data_type = f.data_type(); - match data_type { - &DataType::Int8 => { - instantiate_primitive!(v, nullable, Int8Type, data_type) - } - &DataType::Int16 => { - instantiate_primitive!(v, nullable, Int16Type, data_type) - } - &DataType::Int32 => { - instantiate_primitive!(v, nullable, Int32Type, data_type) - } - &DataType::Int64 => { - instantiate_primitive!(v, nullable, Int64Type, data_type) - } - &DataType::UInt8 => { - instantiate_primitive!(v, nullable, UInt8Type, data_type) - } - &DataType::UInt16 => { - instantiate_primitive!(v, nullable, UInt16Type, data_type) - } - &DataType::UInt32 => { - instantiate_primitive!(v, nullable, UInt32Type, data_type) - } - &DataType::UInt64 => { - instantiate_primitive!(v, nullable, UInt64Type, data_type) - } - &DataType::Float32 => { - instantiate_primitive!(v, nullable, Float32Type, data_type) - } - &DataType::Float64 => { - instantiate_primitive!(v, nullable, Float64Type, data_type) - } - &DataType::Date32 => { - instantiate_primitive!(v, nullable, Date32Type, data_type) - } - &DataType::Date64 => { - instantiate_primitive!(v, nullable, Date64Type, data_type) - } - &DataType::Time32(t) => match t { - TimeUnit::Second => { - instantiate_primitive!( - v, - nullable, - Time32SecondType, - data_type - ) - } - TimeUnit::Millisecond => { - instantiate_primitive!( - v, - nullable, - Time32MillisecondType, - data_type - ) - } - _ => {} - }, - &DataType::Time64(t) => match t { - TimeUnit::Microsecond => { - instantiate_primitive!( - v, - nullable, - Time64MicrosecondType, - data_type - ) - } - TimeUnit::Nanosecond => { - instantiate_primitive!( - v, - nullable, - Time64NanosecondType, - data_type - ) - } - _ => {} - }, - &DataType::Timestamp(t, _) => match t { - TimeUnit::Second => { - instantiate_primitive!( - v, - nullable, - TimestampSecondType, - data_type - ) - } - TimeUnit::Millisecond => { - instantiate_primitive!( - v, - nullable, - TimestampMillisecondType, - data_type - ) - } - TimeUnit::Microsecond => { - instantiate_primitive!( - v, - nullable, - TimestampMicrosecondType, - data_type - ) - } - TimeUnit::Nanosecond => { - instantiate_primitive!( - v, - nullable, - TimestampNanosecondType, - data_type - ) - } - }, - &DataType::Decimal128(_, _) => { - instantiate_primitive! { - v, - nullable, - Decimal128Type, - data_type - } - } - &DataType::Utf8 => { - let b = ByteGroupValueBuilder::::new(OutputType::Utf8); - v.push(Box::new(b) as _) - } - &DataType::LargeUtf8 => { - let b = ByteGroupValueBuilder::::new(OutputType::Utf8); - v.push(Box::new(b) as _) - } - &DataType::Binary => { - let b = ByteGroupValueBuilder::::new(OutputType::Binary); - v.push(Box::new(b) as _) - } - &DataType::LargeBinary => { - let b = ByteGroupValueBuilder::::new(OutputType::Binary); - v.push(Box::new(b) as _) - } - &DataType::Utf8View => { - let b = ByteViewGroupValueBuilder::::new(); - v.push(Box::new(b) as _) - } - &DataType::BinaryView => { - let b = ByteViewGroupValueBuilder::::new(); - v.push(Box::new(b) as _) - } - &DataType::Boolean => { - if nullable { - let b = BooleanGroupValueBuilder::::new(); - v.push(Box::new(b) as _) - } else { - let b = BooleanGroupValueBuilder::::new(); - v.push(Box::new(b) as _) - } - } - dt => { - return not_impl_err!("{dt} not supported in GroupValuesColumn"); - } - } +/// Returns true if the specified data type has a specialized +/// [`GroupColumn`] builder in [`make_group_column`]. +/// +/// This is the allow-list that gates the `GroupValuesRows` fallback in +/// [`crate::aggregates::group_values::new_group_values`]: it must accept +/// exactly the set of types that [`make_group_column`] constructs a +/// builder for. The `group_column_supported_type_matches_make_group_column` +/// test below pins this biconditional. +fn group_column_supported_type(data_type: &DataType) -> bool { + // Nested types (Struct / List / LargeList / FixedSizeList, recursively) have + // no type-specialized `GroupColumn`; they are handled by the generic + // row-backed fallback in `make_group_column` whenever arrow's row format can + // encode them. Gate the fallback to nested types so intentionally-excluded + // scalar types (e.g. Float16, Decimal256) stay on `GroupValuesRows` and the + // `group_column_supported_type` ⇔ `make_group_column` invariant holds. + if data_type.is_nested() { + return RowsGroupColumn::supports_type(data_type); + } + matches!( + *data_type, + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::UInt8 + | DataType::UInt16 + | DataType::UInt32 + | DataType::UInt64 + | DataType::Float16 + | DataType::Float32 + | DataType::Float64 + | DataType::Decimal128(_, _) + | DataType::Decimal256(_, _) + | DataType::Utf8 + | DataType::LargeUtf8 + | DataType::Binary + | DataType::LargeBinary + // Only non-negative widths: a negative width is not a valid + // Arrow type (no array can be constructed for it), and the + // dispatcher in `make_group_column` rejects it. Keep the two + // in lockstep. + | DataType::FixedSizeBinary(0..) + | DataType::Date32 + | DataType::Date64 + // Only the semantically valid Time variants per the Arrow spec. + // The dispatcher in `make_group_column` returns NotImpl for the + // other unit combinations, so accepting them here would cause a + // schema to be routed into GroupValuesColumn and then fail at + // intern. Keep these two arms in lockstep with the dispatcher. + | DataType::Time32(TimeUnit::Second) + | DataType::Time32(TimeUnit::Millisecond) + | DataType::Time64(TimeUnit::Microsecond) + | DataType::Time64(TimeUnit::Nanosecond) + | DataType::Timestamp(_, _) + | DataType::Duration(_) + | DataType::Interval(_) + | DataType::Utf8View + | DataType::BinaryView + | DataType::Boolean + ) +} + +/// Build a [`GroupColumn`] for a single schema field. +/// +/// Extracted from the inline match that used to live in +/// [`GroupValuesColumn::intern`] so the per-field dispatch lives in one +/// place. This factory is the single source of truth for which Arrow types +/// map to which builder, and it is the function that future nested-type +/// specializations (e.g. `Struct`, `List`, `LargeList`) plug into without +/// having to enumerate every combination inline. +/// +/// Returns `Err(not_impl_err!(...))` for any type not in the supported set; +/// callers (`GroupValues::intern`) propagate that error so the +/// `GroupValuesRows` fallback can take over upstream of this builder. +/// +/// The allow-list that gates this dispatcher lives in +/// [`group_column_supported_type`] directly above. +fn make_group_column(field: &Field) -> Result> { + let nullable = field.is_nullable(); + let data_type = field.data_type(); + let mut v: Vec> = Vec::with_capacity(1); + match *data_type { + DataType::Int8 => instantiate_primitive!(v, nullable, Int8Type, data_type), + DataType::Int16 => instantiate_primitive!(v, nullable, Int16Type, data_type), + DataType::Int32 => instantiate_primitive!(v, nullable, Int32Type, data_type), + DataType::Int64 => instantiate_primitive!(v, nullable, Int64Type, data_type), + DataType::UInt8 => instantiate_primitive!(v, nullable, UInt8Type, data_type), + DataType::UInt16 => instantiate_primitive!(v, nullable, UInt16Type, data_type), + DataType::UInt32 => instantiate_primitive!(v, nullable, UInt32Type, data_type), + DataType::UInt64 => instantiate_primitive!(v, nullable, UInt64Type, data_type), + DataType::Float16 => { + instantiate_primitive!(v, nullable, Float16Type, data_type) + } + DataType::Float32 => { + instantiate_primitive!(v, nullable, Float32Type, data_type) + } + DataType::Float64 => { + instantiate_primitive!(v, nullable, Float64Type, data_type) + } + DataType::Date32 => instantiate_primitive!(v, nullable, Date32Type, data_type), + DataType::Date64 => instantiate_primitive!(v, nullable, Date64Type, data_type), + DataType::Time32(t) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, Time32SecondType, data_type) + } + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, Time32MillisecondType, data_type) + } + // Time32 with Microsecond / Nanosecond is not a valid Arrow type + // combination; reject explicitly so group_column_supported_type + // and this dispatcher stay in lockstep (see consistency fuzz below). + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + }, + DataType::Time64(t) => match t { + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, Time64MicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, Time64NanosecondType, data_type) + } + // Time64 with Second / Millisecond is not a valid Arrow type + // combination; reject explicitly. + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + }, + DataType::Timestamp(t, _) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, TimestampSecondType, data_type) } - self.group_values = v; + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, TimestampMillisecondType, data_type) + } + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, TimestampMicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, TimestampNanosecondType, data_type) + } + }, + DataType::Duration(t) => match t { + TimeUnit::Second => { + instantiate_primitive!(v, nullable, DurationSecondType, data_type) + } + TimeUnit::Millisecond => { + instantiate_primitive!(v, nullable, DurationMillisecondType, data_type) + } + TimeUnit::Microsecond => { + instantiate_primitive!(v, nullable, DurationMicrosecondType, data_type) + } + TimeUnit::Nanosecond => { + instantiate_primitive!(v, nullable, DurationNanosecondType, data_type) + } + }, + // `IntervalUnit` has exactly three variants, so this match is exhaustive + // with no fallback arm (unlike Time32 / Time64). + DataType::Interval(u) => match u { + IntervalUnit::YearMonth => { + instantiate_primitive!(v, nullable, IntervalYearMonthType, data_type) + } + IntervalUnit::DayTime => { + instantiate_primitive!(v, nullable, IntervalDayTimeType, data_type) + } + IntervalUnit::MonthDayNano => { + instantiate_primitive!(v, nullable, IntervalMonthDayNanoType, data_type) + } + }, + DataType::Decimal128(_, _) => { + instantiate_primitive!(v, nullable, Decimal128Type, data_type) + } + DataType::Decimal256(_, _) => { + instantiate_primitive!(v, nullable, Decimal256Type, data_type) + } + DataType::Utf8 => { + v.push(Box::new(ByteGroupValueBuilder::::new( + OutputType::Utf8, + ))); + } + DataType::LargeUtf8 => { + v.push(Box::new(ByteGroupValueBuilder::::new( + OutputType::Utf8, + ))); } + DataType::Binary => { + v.push(Box::new(ByteGroupValueBuilder::::new( + OutputType::Binary, + ))); + } + DataType::LargeBinary => { + v.push(Box::new(ByteGroupValueBuilder::::new( + OutputType::Binary, + ))); + } + // A negative width is not a valid Arrow type; it falls to the `_` + // arm below, matching `group_column_supported_type`. + DataType::FixedSizeBinary(byte_width @ 0..) => { + v.push(Box::new(FixedSizeBinaryGroupValueBuilder::new(byte_width))); + } + DataType::Utf8View => { + v.push(Box::new(ByteViewGroupValueBuilder::::new())); + } + DataType::BinaryView => { + v.push(Box::new(ByteViewGroupValueBuilder::::new())); + } + DataType::Boolean => { + if nullable { + v.push(Box::new(BooleanGroupValueBuilder::::new())); + } else { + v.push(Box::new(BooleanGroupValueBuilder::::new())); + } + } + // Generic fallback for nested types (Struct / List / LargeList / + // FixedSizeList, recursively) that lack a type-specialized builder but + // can be encoded by arrow's row format. This is what lets a mixed + // schema keep the column-wise fast path for its native columns instead + // of dropping the whole key onto `GroupValuesRows`. + ref dt if dt.is_nested() && RowsGroupColumn::supports_type(dt) => { + v.push(Box::new(RowsGroupColumn::try_new(dt.clone())?)); + } + _ => return not_impl_err!("{data_type} not supported in GroupValuesColumn"), + } + debug_assert_eq!( + v.len(), + 1, + "make_group_column must push exactly one builder" + ); + Ok(v.into_iter().next().unwrap()) +} +impl GroupValues for GroupValuesColumn { + fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + // `try_new` and the reset points in `emit` / `clear_shrink` keep + // `self.group_values` populated with one builder per schema field, + // so no lazy initialization is needed here. if !STREAMING { self.vectorized_intern(cols, groups) } else { @@ -1104,8 +1180,14 @@ impl GroupValues for GroupValuesColumn { fn emit(&mut self, emit_to: EmitTo) -> Result> { let mut output = match emit_to { EmitTo::All => { - let group_values = mem::take(&mut self.group_values); - debug_assert!(self.group_values.is_empty()); + // Replace the column builders with a fresh set so the + // aggregator is immediately reusable after the drain. + // Same `self.schema` was already validated by `try_new`, + // so `build_group_columns` would only error here if some + // out-of-band schema mutation occurred — propagate it as + // a real Result rather than panicking. + let fresh = Self::build_group_columns(&self.schema)?; + let group_values = mem::replace(&mut self.group_values, fresh); group_values .into_iter() @@ -1204,7 +1286,12 @@ impl GroupValues for GroupValuesColumn { } fn clear_shrink(&mut self, num_rows: usize) { - self.group_values.clear(); + // Reset to a fresh column-builder vector. The schema was validated + // in `try_new`, so rebuilding cannot fail unless something else + // mutated the schema out-of-band — surface that as a panic since + // `clear_shrink` is infallible by trait signature. + self.group_values = Self::build_group_columns(&self.schema) + .expect("schema previously validated in try_new"); self.map.clear(); self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared self.map_size = self.map.capacity() * size_of::<(u64, usize)>(); @@ -1226,39 +1313,7 @@ pub fn supported_schema(schema: &Schema) -> bool { .fields() .iter() .map(|f| f.data_type()) - .all(supported_type) -} - -/// Returns true if the specified data type is supported by [`GroupValuesColumn`] -/// -/// In order to be supported, there must be a specialized implementation of -/// [`GroupColumn`] for the data type, instantiated in [`GroupValuesColumn::intern`] -fn supported_type(data_type: &DataType) -> bool { - matches!( - *data_type, - DataType::Int8 - | DataType::Int16 - | DataType::Int32 - | DataType::Int64 - | DataType::UInt8 - | DataType::UInt16 - | DataType::UInt32 - | DataType::UInt64 - | DataType::Float32 - | DataType::Float64 - | DataType::Decimal128(_, _) - | DataType::Utf8 - | DataType::LargeUtf8 - | DataType::Binary - | DataType::LargeBinary - | DataType::Date32 - | DataType::Date64 - | DataType::Time32(_) - | DataType::Timestamp(_, _) - | DataType::Utf8View - | DataType::BinaryView - | DataType::Boolean - ) + .all(group_column_supported_type) } ///Shows how many `null`s there are in an array @@ -1275,7 +1330,11 @@ enum Nulls { mod tests { use std::{collections::HashMap, sync::Arc}; - use arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray, StringViewArray}; + use arrow::array::{ + Array, ArrayRef, DurationMicrosecondArray, FixedSizeBinaryArray, Float16Array, + Int32Array, Int64Array, PrimitiveArray, RecordBatch, StringArray, + StringViewArray, + }; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; use arrow::{compute::concat_batches, util::pretty::pretty_format_batches}; use datafusion_common::utils::proxy::HashTableAllocExt; @@ -1285,49 +1344,565 @@ mod tests { GroupValues, multi_group_by::GroupValuesColumn, }; - use super::{GroupIndexView, split_vec_min_alloc}; + use super::{ + GroupIndexView, group_column_supported_type, make_group_column, supported_schema, + }; + + /// A mixed group-by key of several native columns plus one nested column + /// that has no type-specialized `GroupColumn`. + /// + /// Before the generic row-backed fallback, `supported_schema` returned + /// `false` for this schema, so the *entire* key dropped to the row-wise + /// `GroupValuesRows`. Now only the nested column pays the row-encoding + /// cost; the native columns keep their compact column-wise storage. This + /// test proves both that (a) the results are identical and (b) the + /// column-wise path now uses less memory than the all-rows fallback. + #[test] + fn mixed_schema_column_path_uses_less_memory_than_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Int64Array}; + use arrow::datatypes::Int64Type; + + // 8 native Int64 columns + 1 FixedSizeList ("embedding"). + let fsl_field = Arc::new(Field::new("item", DataType::Int64, true)); + let mut fields: Vec = (0..8) + .map(|i| Field::new(format!("k{i}"), DataType::Int64, false)) + .collect(); + fields.push(Field::new( + "emb", + DataType::FixedSizeList(Arc::clone(&fsl_field), 4), + true, + )); + let schema: SchemaRef = Arc::new(Schema::new(fields)); + + // The whole schema must now be eligible for the column-wise path. + assert!( + supported_schema(schema.as_ref()), + "mixed native + nested schema should be column-supported now" + ); + + // Build `n_groups` distinct rows (each row is its own group). + let n_groups = 4000usize; + let mut cols: Vec = (0..8) + .map(|c| { + let vals: Vec = + (0..n_groups).map(|r| (r as i64) * 8 + c as i64).collect(); + Arc::new(Int64Array::from(vals)) as ArrayRef + }) + .collect(); + let emb: Vec>>> = (0..n_groups) + .map(|r| { + Some(vec![ + Some(r as i64), + Some(r as i64 + 1), + Some(r as i64 + 2), + Some(r as i64 + 3), + ]) + }) + .collect(); + cols.push( + Arc::new(FixedSizeListArray::from_iter_primitive::( + emb, 4, + )) as ArrayRef, + ); + + // Intern the same data into both implementations. + let mut column_path = GroupValuesColumn::::try_new(Arc::clone(&schema)) + .expect("column path"); + let mut rows_path = + GroupValuesRows::try_new(Arc::clone(&schema)).expect("rows path"); + + let mut g1 = vec![]; + let mut g2 = vec![]; + column_path.intern(&cols, &mut g1).unwrap(); + rows_path.intern(&cols, &mut g2).unwrap(); + + // (a) Correctness: same number of groups and identical group assignment. + assert_eq!(column_path.len(), n_groups); + assert_eq!(rows_path.len(), n_groups); + assert_eq!(g1, g2, "group assignment must match the rows fallback"); + + // (b) Memory: the column-wise path stores the 8 native columns compactly + // and only row-encodes the nested one, so it should be smaller than + // encoding every column into rows. + // + // The delta is only printed here — a hard `column_size < rows_size` + // assert would be brittle to future Arrow row-format or memory- + // accounting changes without reflecting a grouping-correctness + // regression. Track the memory improvement via benchmarks instead. + let column_size = column_path.size(); + let rows_size = rows_path.size(); + println!( + "mixed-schema group values size: column-wise = {column_size} bytes, \ + all-rows fallback = {rows_size} bytes \ + ({:.1}% of fallback)", + 100.0 * column_size as f64 / rows_size as f64 + ); + + // Emitted values must be equal too (compare via the rows fallback which + // is the established reference implementation). + let out_col = column_path.emit(EmitTo::All).unwrap(); + let out_row = rows_path.emit(EmitTo::All).unwrap(); + assert_eq!(out_col.len(), out_row.len()); + for (a, b) in out_col.iter().zip(out_row.iter()) { + assert_eq!(a.as_ref(), b.as_ref()); + } + } + + /// Relabel a group-index vector so labels are assigned in order of first + /// appearance. Two vectors are equivalent groupings iff their canonical + /// forms are equal — this ignores the (opaque, non-semantic) difference in + /// group-index numbering between the vectorized column path and the + /// sequential rows fallback. + /// + /// The [`GroupValues`] trait only guarantees that equal keys receive the + /// same group-id and that new keys receive a fresh id; the order in which + /// new ids are handed out is deliberately not part of the contract, and + /// can differ between correct implementations (e.g. because of internal + /// hash-map ordering). Canonicalizing before comparison is what lets us + /// assert equivalence across implementations. + fn canonical_grouping(groups: &[usize]) -> Vec { + let mut map = HashMap::new(); + let mut next = 0usize; + groups + .iter() + .map(|&g| { + *map.entry(g).or_insert_with(|| { + let v = next; + next += 1; + v + }) + }) + .collect() + } + + /// The generic row-backed column must be behavior-preserving: for the + /// nested columns it now handles, `GroupValuesColumn` must induce the same + /// grouping (partition of rows) as the established `GroupValuesRows` + /// fallback — including the float `-0.0` / `+0.0` / `NaN` edge cases decided + /// jointly by hashing and the row format. + #[test] + fn nested_float_edge_cases_match_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Float64Array}; + + let item = Arc::new(Field::new("item", DataType::Float64, true)); + let schema: SchemaRef = Arc::new(Schema::new(vec![Field::new( + "emb", + DataType::FixedSizeList(Arc::clone(&item), 2), + true, + )])); + assert!(supported_schema(schema.as_ref())); + + // Rows exercising +0.0 vs -0.0, two NaN bit patterns, and inner nulls. + let nan = f64::NAN; + let other_nan = f64::from_bits(0x7ff8_0000_0000_0001); + let values = Float64Array::from(vec![ + Some(0.0), + Some(1.0), // [ +0.0, 1.0 ] + Some(-0.0), + Some(1.0), // [ -0.0, 1.0 ] + Some(nan), + Some(2.0), // [ NaN, 2.0 ] + Some(other_nan), + Some(2.0), // [ NaN', 2.0 ] + Some(0.0), + Some(1.0), // [ +0.0, 1.0 ] (dup of row 0) + ]); + let field_ref = Arc::new(Field::new("item", DataType::Float64, true)); + let input: ArrayRef = Arc::new(FixedSizeListArray::new( + field_ref, + 2, + Arc::new(values), + None, + )); + + let cols = vec![input]; + + let mut column_path = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap(); + + let mut g1 = vec![]; + let mut g2 = vec![]; + column_path.intern(&cols, &mut g1).unwrap(); + rows_path.intern(&cols, &mut g2).unwrap(); + + assert_eq!( + canonical_grouping(&g1), + canonical_grouping(&g2), + "column-wise path must induce the same grouping as the rows fallback \ + on float edge cases (got column={g1:?}, rows={g2:?})" + ); + assert_eq!(column_path.len(), rows_path.len()); + } + + /// Equivalence across multiple `intern` batches and `EmitTo::First(n)`. + #[test] + fn multi_batch_and_emit_first_matches_rows_fallback() { + use crate::aggregates::group_values::GroupValuesRows; + use arrow::array::{FixedSizeListArray, Int32Array}; + use arrow::datatypes::Int32Type; + + let item = Arc::new(Field::new("item", DataType::Int32, true)); + let schema: SchemaRef = Arc::new(Schema::new(vec![ + Field::new("k", DataType::Int32, false), + Field::new("emb", DataType::FixedSizeList(Arc::clone(&item), 2), true), + ])); + + let make_batch = |base: i32| -> Vec { + let k = Arc::new(Int32Array::from(vec![base, base + 1, base])) as ArrayRef; + let emb: Vec>>> = vec![ + Some(vec![Some(base), Some(base)]), + Some(vec![Some(base + 1), None]), + Some(vec![Some(base), Some(base)]), // dup of row 0 + ]; + let emb = Arc::new( + FixedSizeListArray::from_iter_primitive::(emb, 2), + ) as ArrayRef; + vec![k, emb] + }; + + let mut column_path = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + let mut rows_path = GroupValuesRows::try_new(Arc::clone(&schema)).unwrap(); + + for base in [0, 10, 0] { + let cols = make_batch(base); + let (mut a, mut b) = (vec![], vec![]); + column_path.intern(&cols, &mut a).unwrap(); + rows_path.intern(&cols, &mut b).unwrap(); + // Same grouping (partition), even if the opaque group-index labels + // differ between the vectorized and sequential paths. + assert_eq!( + canonical_grouping(&a), + canonical_grouping(&b), + "grouping must match for batch base={base}" + ); + } + let total_groups = column_path.len(); + assert_eq!(total_groups, rows_path.len()); + + // `EmitTo::First(n)` then `EmitTo::All` on the nested column path must + // work and together emit exactly `total_groups` rows. (Cross-path value + // equality is covered by `mixed_schema_...` and the row_backed unit + // tests; group-index ordering differs here so we check counts.) + let col_first = column_path.emit(EmitTo::First(2)).unwrap(); + assert_eq!(col_first[0].len(), 2); + let col_rest = column_path.emit(EmitTo::All).unwrap(); + assert_eq!(col_first[0].len() + col_rest[0].len(), total_groups); + // Column count / schema preserved on both emits. + assert_eq!(col_first.len(), schema.fields().len()); + assert_eq!(col_rest.len(), schema.fields().len()); + } + + /// CRITICAL invariant: if `group_column_supported_type(t)` returns true + /// the dispatcher must accept that type at intern time, and conversely + /// if `group_column_supported_type(t)` returns false the planner must + /// NOT route it through `GroupValuesColumn`. A divergence here would + /// let the planner select `GroupValuesColumn` for a type whose + /// dispatcher arm is missing, producing a runtime `not_impl_err` after + /// the field reaches the builder factory. + /// + /// This test fuzzes a representative cross-section of types and asserts + /// both directions of the biconditional. When a new specialization is + /// added (`Float16`, `FixedSizeList`, `Struct`, ...) it should be added + /// to the supported_cases vector; when a type is intentionally rejected + /// it should be added to unsupported_cases. #[test] - fn test_split_vec_min_alloc_drain_branch() { - // n * 2 <= len → drain+collect branch (allocates n elements) - let mut v = vec![1, 2, 3, 4, 5, 6]; - let first = split_vec_min_alloc(&mut v, 2); - assert_eq!(first, vec![1, 2]); - assert_eq!(v, vec![3, 4, 5, 6]); + fn group_column_supported_type_matches_make_group_column() { + let supported_cases: Vec = vec![ + DataType::Int8, + DataType::Int64, + DataType::UInt64, + DataType::Float32, + DataType::Float64, + DataType::Float16, + DataType::Decimal128(38, 10), + DataType::Decimal256(76, 10), + DataType::Utf8, + DataType::LargeUtf8, + DataType::Utf8View, + DataType::Binary, + DataType::LargeBinary, + DataType::BinaryView, + DataType::FixedSizeBinary(16), + // Zero-width FixedSizeBinary is valid per the Arrow spec + DataType::FixedSizeBinary(0), + DataType::Boolean, + DataType::Date32, + DataType::Date64, + DataType::Time32(arrow::datatypes::TimeUnit::Second), + DataType::Time32(arrow::datatypes::TimeUnit::Millisecond), + DataType::Time64(arrow::datatypes::TimeUnit::Microsecond), + DataType::Time64(arrow::datatypes::TimeUnit::Nanosecond), + DataType::Timestamp(arrow::datatypes::TimeUnit::Nanosecond, None), + DataType::Duration(arrow::datatypes::TimeUnit::Second), + DataType::Duration(arrow::datatypes::TimeUnit::Millisecond), + DataType::Duration(arrow::datatypes::TimeUnit::Microsecond), + DataType::Duration(arrow::datatypes::TimeUnit::Nanosecond), + DataType::Interval(arrow::datatypes::IntervalUnit::YearMonth), + DataType::Interval(arrow::datatypes::IntervalUnit::DayTime), + DataType::Interval(arrow::datatypes::IntervalUnit::MonthDayNano), + ]; + + for dt in &supported_cases { + assert!( + group_column_supported_type(dt), + "expected group_column_supported_type=true for {dt:?}" + ); + let field = Field::new("col", dt.clone(), true); + make_group_column(&field).unwrap_or_else(|e| { + panic!( + "group_column_supported_type accepted {dt:?} but make_group_column rejected: {e}" + ) + }); + } + + let unsupported_cases: Vec = vec![ + // Invalid Time-unit combinations: Time32 is defined only for + // Second / Millisecond and Time64 only for Microsecond / + // Nanosecond. The TimeUnit enum allows constructing the other + // combinations programmatically, but they are not valid Arrow + // types and must be rejected by both group_column_supported_type + // and the dispatcher. + DataType::Time64(arrow::datatypes::TimeUnit::Second), + DataType::Time64(arrow::datatypes::TimeUnit::Millisecond), + DataType::Time32(arrow::datatypes::TimeUnit::Microsecond), + DataType::Time32(arrow::datatypes::TimeUnit::Nanosecond), + // A negative width is representable in the DataType but is not + // a valid Arrow type; no array can be constructed for it. + DataType::FixedSizeBinary(-5), + ]; + + for dt in &unsupported_cases { + assert!( + !group_column_supported_type(dt), + "expected group_column_supported_type=false for {dt:?}" + ); + let field = Field::new("col", dt.clone(), true); + assert!( + make_group_column(&field).is_err(), + "group_column_supported_type rejected {dt:?} but make_group_column accepted it" + ); + } } + // `Duration` group keys stay on the `GroupValuesColumn` fast path, dedup + // (including nulls), and round-trip with the `Duration` type preserved. #[test] - fn test_split_vec_min_alloc_split_off_branch() { - // remaining < n → split_off+replace branch (allocates remaining elements) - let mut v = vec![1, 2, 3, 4, 5, 6]; - let first = split_vec_min_alloc(&mut v, 4); - assert_eq!(first, vec![1, 2, 3, 4]); - assert_eq!(v, vec![5, 6]); + fn test_group_values_column_duration() { + use arrow::datatypes::TimeUnit; + + let schema = Arc::new(Schema::new(vec![ + Field::new("d", DataType::Duration(TimeUnit::Microsecond), true), + Field::new("i", DataType::Int64, true), + ])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + // (d, i) rows, where row 3 repeats row 0 and row 4 repeats the null pair. + let d: ArrayRef = Arc::new(DurationMicrosecondArray::from(vec![ + Some(10), + None, + Some(20), + Some(10), + None, + ])); + let i: ArrayRef = Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + Some(1), + None, + ])); + let mut groups = Vec::new(); + group_values.intern(&[d, i], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 2, 0, 1]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + // The Duration column round-trips as Duration on emit, not bare i64. + assert_eq!( + emitted[0].data_type(), + &DataType::Duration(TimeUnit::Microsecond) + ); + let actual = emitted[0] + .as_any() + .downcast_ref::() + .expect("emitted column should be a DurationMicrosecondArray"); + // Three groups in first-seen order: 10, null, 20. + assert_eq!(actual.len(), 3); + assert_eq!(actual.value(0), 10); + assert!(actual.is_null(1)); + assert_eq!(actual.value(2), 20); } + // `(Float16, Int32)` keys: ±0.0 collapse (stored as +0.0), NaNs collapse, and + // the Int32 key keeps `(0.0, 4)` distinct from `(±0.0, 3)`. #[test] - fn test_split_vec_min_alloc_exactly_half() { - // n * 2 == len → drain branch (boundary condition) - let mut v = vec![1, 2, 3, 4]; - let first = split_vec_min_alloc(&mut v, 2); - assert_eq!(first, vec![1, 2]); - assert_eq!(v, vec![3, 4]); + fn test_group_values_column_float16() { + use half::f16; + + let schema = Arc::new(Schema::new(vec![ + Field::new("f", DataType::Float16, true), + Field::new("i", DataType::Int32, true), + ])); + assert!(supported_schema(&schema)); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + let f: ArrayRef = Arc::new(Float16Array::from(vec![ + Some(f16::from_f32(1.0)), + Some(f16::from_f32(-0.0)), + Some(f16::from_f32(0.0)), + Some(f16::from_f32(0.0)), + Some(f16::NAN), + Some(f16::NAN), + None, + None, + ])); + let i: ArrayRef = Arc::new(Int32Array::from(vec![ + Some(3), + Some(3), + Some(3), + Some(4), + Some(3), + Some(3), + Some(3), + Some(3), + ])); + let mut groups = Vec::new(); + group_values.intern(&[f, i], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 1, 2, 3, 3, 4, 4]); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + assert_eq!(emitted[0].data_type(), &DataType::Float16); + let keys = emitted[0] + .as_any() + .downcast_ref::() + .expect("emitted column should be a Float16Array"); + assert_eq!(keys.len(), 5); + assert_eq!(keys.value(0), f16::from_f32(1.0)); + // The ±0.0 group is stored canonically as +0.0 (not -0.0). + assert_eq!(keys.value(1).to_bits(), f16::from_f32(0.0).to_bits()); + assert_eq!(keys.value(2).to_bits(), f16::from_f32(0.0).to_bits()); + assert!(keys.value(3).is_nan()); + assert!(keys.is_null(4)); + let ids = emitted[1] + .as_any() + .downcast_ref::() + .expect("emitted column should be an Int32Array"); + assert_eq!(ids.values().to_vec(), vec![3, 3, 4, 3, 3]); } + // `(Interval, Int32)` keys for each of the three interval units: null keys + // dedup, the Int32 key splits equal intervals, and emit gives back Interval. #[test] - fn test_split_vec_min_alloc_take_all() { - let mut v = vec![1, 2, 3]; - let first = split_vec_min_alloc(&mut v, 3); - assert_eq!(first, vec![1, 2, 3]); - assert!(v.is_empty()); + fn test_group_values_column_interval() { + use arrow::datatypes::{ + ArrowPrimitiveType, IntervalDayTime, IntervalDayTimeType, + IntervalMonthDayNano, IntervalMonthDayNanoType, IntervalUnit, + IntervalYearMonthType, + }; + + fn check(unit: IntervalUnit, value: T::Native) { + let schema = Arc::new(Schema::new(vec![ + Field::new("i", DataType::Interval(unit), true), + Field::new("n", DataType::Int32, true), + ])); + assert!(supported_schema(&schema), "{unit:?} schema not supported"); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + let i: ArrayRef = Arc::new(PrimitiveArray::::from_iter([ + Some(value), + None, + Some(value), + None, + Some(value), + ])); + let n: ArrayRef = Arc::new(Int32Array::from(vec![3, 3, 3, 3, 4])); + let mut groups = Vec::new(); + group_values.intern(&[i, n], &mut groups).unwrap(); + assert_eq!(groups, vec![0, 1, 0, 1, 2], "{unit:?}"); + + let emitted = group_values.emit(EmitTo::All).unwrap(); + assert_eq!(emitted.len(), 2); + // The emitted key keeps its Interval type, not the bare native. + assert_eq!(emitted[0].data_type(), &DataType::Interval(unit)); + let actual = emitted[0] + .as_any() + .downcast_ref::>() + .unwrap_or_else(|| panic!("emitted column should be a {unit:?} array")); + // Three groups in first-seen order: value, null, value (n=4). + assert_eq!(actual.len(), 3, "{unit:?}"); + assert_eq!(actual.value(0), value, "{unit:?}"); + assert!(actual.is_null(1), "{unit:?}"); + assert_eq!(actual.value(2), value, "{unit:?}"); + let ids = emitted[1] + .as_any() + .downcast_ref::() + .expect("emitted column should be an Int32Array"); + assert_eq!(ids.values().to_vec(), vec![3, 3, 4], "{unit:?}"); + } + + check::(IntervalUnit::YearMonth, 13); + check::(IntervalUnit::DayTime, IntervalDayTime::new(1, 500)); + check::( + IntervalUnit::MonthDayNano, + IntervalMonthDayNano::new(1, 0, 0), + ); } #[test] - fn test_split_vec_min_alloc_take_none() { - let mut v = vec![1, 2, 3]; - let first = split_vec_min_alloc(&mut v, 0); - assert!(first.is_empty()); - assert_eq!(v, vec![1, 2, 3]); + fn supported_schema_rejects_mix_of_supported_and_unsupported() { + // One unsupported column flips the whole schema to the GroupValuesRows + // fallback. Time64(Second) stays invalid as new primitive builders land. + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new( + "c", + DataType::Time64(arrow::datatypes::TimeUnit::Second), + true, + ), + ]); + assert!(!supported_schema(&schema)); + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + Field::new("c", DataType::Boolean, true), + ]); + assert!(supported_schema(&schema)); + } + + #[test] + fn try_new_returns_not_impl_for_unsupported_top_level_type() { + // `try_new` now eagerly constructs the per-field GroupColumn + // builders via `make_group_column`, so an unsupported schema is + // rejected at construction time rather than at first `intern`. + // `GroupValuesColumn` doesn't implement `Debug`, so explicit match + // instead of `unwrap_err`. + let schema = Arc::new(Schema::new(vec![Field::new( + "x", + DataType::Time64(arrow::datatypes::TimeUnit::Second), + true, + )])); + match GroupValuesColumn::::try_new(schema) { + Ok(_) => panic!("expected NotImpl error, but try_new succeeded"), + Err(e) => { + let msg = e.to_string(); + assert!( + msg.contains("not supported in GroupValuesColumn"), + "expected NotImpl error from dispatcher, got: {msg}" + ); + } + } } #[test] @@ -1343,6 +1918,78 @@ mod tests { check_result(&actual_batch, &data_set.expected_batch); } + #[test] + fn test_intern_for_fixed_size_binary_group_values() { + // Two-column group by `(FixedSizeBinary(2), Int64)` exercising the + // vectorized intern path end-to-end (hashing included), with nulls, + // within-batch repeats and across-batch repeats. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::FixedSizeBinary(2), true), + Field::new("b", DataType::Int64, true), + ])); + let mut group_values = + GroupValuesColumn::::try_new(Arc::clone(&schema)).unwrap(); + + fn fsb(values: Vec>) -> ArrayRef { + Arc::new( + FixedSizeBinaryArray::try_from_sparse_iter_with_size( + values.into_iter(), + 2, + ) + .unwrap(), + ) + } + + let batch1: Vec = vec![ + fsb(vec![Some(b"aa"), Some(b"aa"), None, None, Some(b"bb")]), + Arc::new(Int64Array::from(vec![ + Some(1), + Some(1), + None, + Some(2), + None, + ])), + ]; + // Mix of groups repeated from batch1 and new groups + let batch2: Vec = vec![ + fsb(vec![Some(b"aa"), Some(b"cc"), None, Some(b"bb")]), + Arc::new(Int64Array::from(vec![Some(1), Some(1), None, Some(3)])), + ]; + + group_values.intern(&batch1, &mut vec![]).unwrap(); + group_values.intern(&batch2, &mut vec![]).unwrap(); + + let actual_batch = group_values.emit(EmitTo::All).unwrap(); + let actual_batch = + RecordBatch::try_new(Arc::clone(&schema), actual_batch).unwrap(); + + let expected_batch = RecordBatch::try_new( + schema, + vec![ + fsb(vec![ + Some(b"aa"), + None, + None, + Some(b"bb"), + Some(b"cc"), + Some(b"bb"), + ]), + Arc::new(Int64Array::from(vec![ + Some(1), + None, + Some(2), + None, + Some(1), + Some(3), + ])), + ], + ) + .unwrap(); + + assert_eq!(actual_batch.num_rows(), expected_batch.num_rows()); + check_result(&actual_batch, &expected_batch); + } + #[test] fn test_emit_first_n_for_vectorized_group_values() { let data_set = VectorizedTestDataSet::new(); @@ -1400,6 +2047,17 @@ mod tests { let schema = Arc::new(Schema::new_with_metadata(vec![field], HashMap::new())); let mut group_values = GroupValuesColumn::::try_new(schema).unwrap(); + // Seed the column with 12 placeholder rows so the upcoming + // `emit(EmitTo::First(4))` calls can `take_n` without panicking. + // The hashmap entries below reference group indices 0..=11, so the + // single column builder needs at least 12 rows to back them. + let seed: ArrayRef = Arc::new(Int32Array::from(vec![0_i32; 12])); + for row in 0..12 { + group_values.group_values[0] + .append_val(&seed, row) + .expect("seed append"); + } + // Insert group index views and check if success to insert insert_inline_group_index_view(&mut group_values, 0, 0); insert_non_inline_group_index_view(&mut group_values, 1, vec![1, 2]); diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs index 4aae996f6811d..148c5697dea3b 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/primitive.rs @@ -15,8 +15,9 @@ // specific language governing permissions and limitations // under the License. +use crate::aggregates::group_values::HashValue; use crate::aggregates::group_values::multi_group_by::{ - GroupColumn, Nulls, nulls_equal_to, split_vec_min_alloc, + GroupColumn, Nulls, nulls_equal_to, }; use crate::aggregates::group_values::null_builder::MaybeNullBufferBuilder; use arrow::array::ArrowNativeTypeOp; @@ -28,6 +29,7 @@ use arrow::buffer::ScalarBuffer; use arrow::datatypes::DataType; use arrow::util::bit_util::apply_bitwise_binary_op; use datafusion_common::Result; +use datafusion_common::utils::split_vec_min_alloc; use datafusion_execution::memory_pool::proxy::VecAllocExt; use std::iter; use std::sync::Arc; @@ -50,6 +52,7 @@ pub struct PrimitiveGroupValueBuilder PrimitiveGroupValueBuilder where T: ArrowPrimitiveType, + T::Native: HashValue, { /// Create a new `PrimitiveGroupValueBuilder` pub fn new(data_type: DataType) -> Self { @@ -80,6 +83,9 @@ where for (i, (&lhs_row, &rhs_row)) in lhs_rows.iter().zip(rhs_rows.iter()).enumerate() { + if !equal_to_results.get_bit(i) { + continue; + } let left = if cfg!(debug_assertions) { self.group_values[lhs_row] } else { @@ -90,7 +96,9 @@ where } else { unsafe { *array_values.get_unchecked(rhs_row) } }; - if left.is_eq(right) { + // `left` was already canonicalized on append; canonicalize the + // input so ±0 (and any future equivalence class) compares equal. + if left.is_eq(right.canonicalize()) { cmp_buf[i / 8] |= 1 << (i % 8); } } @@ -122,7 +130,6 @@ where if !equal_to_results.get_bit(idx) { continue; } - let exist_null = self.nulls.is_null(lhs_row); let input_null = array.is_null(rhs_row); if let Some(result) = nulls_equal_to(exist_null, input_null) { @@ -132,7 +139,7 @@ where continue; } - if !self.group_values[lhs_row].is_eq(array.value(rhs_row)) { + if !self.group_values[lhs_row].is_eq(array.value(rhs_row).canonicalize()) { equal_to_results.set_bit(idx, false); } } @@ -141,6 +148,8 @@ where impl GroupColumn for PrimitiveGroupValueBuilder +where + T::Native: HashValue, { fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { // Perf: skip null check (by short circuit) if input is not nullable @@ -153,7 +162,8 @@ impl GroupColumn // Otherwise, we need to check their values } - self.group_values[lhs_row].is_eq(array.as_primitive::().value(rhs_row)) + self.group_values[lhs_row] + .is_eq(array.as_primitive::().value(rhs_row).canonicalize()) } fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { @@ -164,10 +174,12 @@ impl GroupColumn self.group_values.push(T::default_value()); } else { self.nulls.append(false); - self.group_values.push(array.as_primitive::().value(row)); + self.group_values + .push(array.as_primitive::().value(row).canonicalize()); } } else { - self.group_values.push(array.as_primitive::().value(row)); + self.group_values + .push(array.as_primitive::().value(row).canonicalize()); } Ok(()) @@ -213,7 +225,7 @@ impl GroupColumn self.group_values.push(T::default_value()); } else { self.nulls.append(false); - self.group_values.push(arr.value(row)); + self.group_values.push(arr.value(row).canonicalize()); } } } @@ -221,7 +233,7 @@ impl GroupColumn (true, Nulls::None) => { self.nulls.append_n(rows.len(), false); for &row in rows { - self.group_values.push(arr.value(row)); + self.group_values.push(arr.value(row).canonicalize()); } } @@ -233,7 +245,7 @@ impl GroupColumn (false, _) => { for &row in rows { - self.group_values.push(arr.value(row)); + self.group_values.push(arr.value(row).canonicalize()); } } } @@ -283,9 +295,10 @@ mod tests { use crate::aggregates::group_values::multi_group_by::primitive::PrimitiveGroupValueBuilder; use arrow::array::{ - ArrayRef, BooleanBufferBuilder, Float32Array, Int64Array, NullBufferBuilder, + ArrayRef, BooleanBufferBuilder, Float32Array, Int32Array, Int64Array, + NullBufferBuilder, }; - use arrow::datatypes::{DataType, Float32Type, Int64Type}; + use arrow::datatypes::{DataType, Float32Type, Int32Type, Int64Type}; use super::GroupColumn; @@ -584,6 +597,25 @@ mod tests { assert!(results[4]); } + // All bits false: every row must be skipped; accessing any lhs/rhs index would panic. + #[test] + fn test_vectorized_equal_to_skips_false_rows() { + let mut builder = + PrimitiveGroupValueBuilder::::new(DataType::Int32); + let array = Arc::new(Int32Array::from(vec![None::, None])) as ArrayRef; + builder.vectorized_append(&array, &[0, 1]).unwrap(); + + let mut results = BooleanBufferBuilder::new(2); + results.append_n(2, false); + + builder.vectorized_equal_to( + &[usize::MAX, usize::MAX], + &array, + &[usize::MAX, usize::MAX], + &mut results, + ); + } + #[test] fn test_primitive_take_n() { // drain branch: n * 2 <= len diff --git a/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs new file mode 100644 index 0000000000000..1445a81f2189b --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/group_values/multi_group_by/row_backed.rs @@ -0,0 +1,1129 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! A generic [`GroupColumn`] backed by the arrow row format. +//! +//! Unlike the type-specialized builders in this module (primitive, byte, +//! boolean, ...), [`RowsGroupColumn`] works for *any* data type that arrow's +//! [`RowConverter`] can encode — including nested types such as `Struct`, +//! `List`, `LargeList` and `FixedSizeList`. It stores one group value per row +//! in a single-column [`Rows`] buffer and compares group keys by their encoded +//! bytes. +//! +//! # Why this exists +//! +//! [`GroupValuesColumn`] can only be used when *every* column of the group-by +//! key has a [`GroupColumn`] implementation; otherwise the whole aggregation +//! falls back to the row-wise [`GroupValuesRows`], which is materially slower +//! and heavier for the columns that *would* have qualified for the column-wise +//! fast path. By providing a generic fallback `GroupColumn`, a schema like +//! `GROUP BY int_col, struct_col` keeps `int_col` on its fast native builder +//! and only pays the row-encoding cost on `struct_col`, instead of dragging both +//! columns onto `GroupValuesRows`. +//! +//! # Relationship to hashing +//! +//! This column does not hash anything itself: [`GroupValuesColumn`] hashes the +//! raw input columns via `create_hashes`, which already supports nested types. +//! Equality is decided here by comparing arrow-row bytes. For the two to agree +//! on group identity, values that this column considers equal must hash equal — +//! see the float `-0.0` / `NaN` note on [`RowsGroupColumn`]. +//! +//! [`GroupValuesColumn`]: crate::aggregates::group_values::multi_group_by::GroupValuesColumn +//! [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows + +use crate::aggregates::group_values::multi_group_by::GroupColumn; +use crate::aggregates::group_values::row::encode_array_if_necessary; + +use arrow::array::{Array, ArrayRef, BooleanBufferBuilder}; +use arrow::datatypes::DataType; +use arrow::row::{RowConverter, Rows, SortField}; +use datafusion_common::{DataFusionError, Result}; + +/// A [`GroupColumn`] that stores group values for a single column in the arrow +/// [row format], backed by a single-field [`RowConverter`]. +/// +/// # NULL semantics +/// +/// The [`GroupColumn`] contract treats two NULLs as equal. The row format +/// encodes NULL with a distinct sentinel, so `null`-row bytes compare equal to +/// each other and unequal to any non-null row — matching the contract without +/// special-casing. +/// +/// # Float `-0.0` / `NaN` +/// +/// Equality here is byte equality under arrow's IEEE-754 *totalOrder* row +/// encoding, which treats `-0.0` and `+0.0` as distinct and canonicalizes +/// `NaN`. Because hashing is performed separately (on the raw input array), a +/// caller must ensure the two agree — e.g. by normalizing `-0.0 → +0.0` on the +/// input columns before hashing when a float leaf is present (as +/// [`GroupValuesRows`] does). See the module docs. +/// +/// [row format]: arrow::row +/// [`GroupValuesRows`]: crate::aggregates::group_values::GroupValuesRows +pub struct RowsGroupColumn { + /// Single-field row converter for this column's data type. + row_converter: RowConverter, + /// Accumulated group values in row format; `group_values.row(i)` is the + /// group value for group index `i`. + group_values: Rows, + /// The column's expected output type. The row format decodes dictionary / + /// run-end encoded values to their plain value type, so emitted arrays are + /// re-encoded to this type in `build` / `take_n` (mirroring + /// `GroupValuesRows::emit`). + output_type: DataType, +} + +/// Walk `data_type`'s subtree and return `true` if it contains a +/// [`DataType::FixedSizeList`] whose descendant tree includes any +/// [`DataType::Dictionary`]. +/// +/// Two-state recursion: once we cross a `FixedSizeList`, `inside_fsl` +/// stays true for every descendant, so a `Dictionary` anywhere below +/// counts. Above that boundary, encountering a `Dictionary` is fine — +/// only nested containers propagate the risk. +/// +/// TODO: this guard works around +/// (`decode_fixed_size_list` panics instead of applying the +/// dictionary-flatten `corrected_type` step). Fixed upstream by +/// (merged 2026-07-24, not +/// yet in a release as of arrow 59.1.0). Once DataFusion upgrades to an +/// arrow release containing that fix, `FixedSizeList` will +/// decode like the other list-likes (flattened child, re-encoded by +/// `encode_array_if_necessary`'s existing `FixedSizeList` arm) — remove +/// this guard and its `supports_type` rejection at that point. +fn contains_fsl_with_dictionary(data_type: &DataType) -> bool { + fn walk(dt: &DataType, inside_fsl: bool) -> bool { + match dt { + DataType::Dictionary(_, _) => inside_fsl, + DataType::FixedSizeList(f, _) => walk(f.data_type(), true), + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) => walk(f.data_type(), inside_fsl), + DataType::Map(f, _) => walk(f.data_type(), inside_fsl), + DataType::Struct(fs) => fs.iter().any(|f| walk(f.data_type(), inside_fsl)), + DataType::RunEndEncoded(_, values) => walk(values.data_type(), inside_fsl), + DataType::Union(fs, _) => { + fs.iter().any(|(_, f)| walk(f.data_type(), inside_fsl)) + } + _ => false, + } + } + walk(data_type, false) +} + +/// Return `true` if `data_type` contains a [`DataType::Union`] or +/// [`DataType::RunEndEncoded`] anywhere in its subtree. +/// +/// These two nested variants can round-trip through `RowConverter` in +/// principle, but their arrow-row decoders have not been validated by +/// this crate's test matrix against the full range of leaf types (dict, +/// nested, etc.). Before this PR both were handled by `GroupValuesRows` +/// (they were not `is_nested`-eligible for `GroupValuesColumn`), so +/// reject them here to preserve the pre-PR routing rather than route +/// untested shapes through `RowsGroupColumn`. When we grow explicit +/// round-trip tests for these types, this blacklist can be removed. +fn contains_union_or_run_end_encoded(data_type: &DataType) -> bool { + match data_type { + DataType::Union(_, _) | DataType::RunEndEncoded(_, _) => true, + DataType::List(f) + | DataType::LargeList(f) + | DataType::ListView(f) + | DataType::LargeListView(f) + | DataType::FixedSizeList(f, _) => { + contains_union_or_run_end_encoded(f.data_type()) + } + DataType::Map(f, _) => contains_union_or_run_end_encoded(f.data_type()), + DataType::Struct(fs) => fs + .iter() + .any(|f| contains_union_or_run_end_encoded(f.data_type())), + _ => false, + } +} + +impl RowsGroupColumn { + /// Returns whether `data_type` can be handled by this generic column. + /// + /// This is stricter than [`RowConverter::supports_fields`]: the row + /// format also has to survive the `build` / `take_n` reverse trip + /// through [`RowConverter::convert_rows`], and arrow's + /// `decode_fixed_size_list` (arrow-row 59.1.0) skips the + /// dictionary-flatten correction that the other list-like decoders + /// apply, so any `FixedSizeList` containing a `Dictionary` leaf + /// panics on emit with `"FixedSizeListArray expected data type + /// Dictionary(...) got for \"item\""`. + /// + /// Reject those shapes here so `make_group_column` falls back to + /// `GroupValuesRows`. The other list-likes (`List`, `LargeList`, + /// `ListView`, `LargeListView`, `Map`) do carry the correction, so + /// they decode without panicking — but the correction *flattens* any + /// dictionary child to its value type, so `build` / `take_n` must + /// re-encode the emitted array back to `output_type` via + /// `encode_array_if_necessary` (which has a reconstruction arm for + /// each of these containers). + /// + /// Additionally, `Union` and `RunEndEncoded` are rejected because + /// they were routed to `GroupValuesRows` before this column existed + /// and their arrow-row round-trip has not been covered by this + /// crate's tests yet. Keeping them on the pre-PR path avoids + /// introducing an untested code path for those types. + pub fn supports_type(data_type: &DataType) -> bool { + if contains_fsl_with_dictionary(data_type) { + return false; + } + if contains_union_or_run_end_encoded(data_type) { + return false; + } + RowConverter::supports_fields(&[SortField::new(data_type.clone())]) + } + + /// Create an empty [`RowsGroupColumn`] for `data_type`. + pub fn try_new(data_type: DataType) -> Result { + let row_converter = RowConverter::new(vec![SortField::new(data_type.clone())])?; + let group_values = row_converter.empty_rows(0, 0); + Ok(Self { + row_converter, + group_values, + output_type: data_type, + }) + } + + /// Materialize `rows` into a single array of `self.output_type`, re-applying + /// dictionary / run-end encoding the row format strips on decode. + fn rows_to_array<'a>( + &self, + rows: impl IntoIterator>, + ) -> ArrayRef { + let mut arrays = self + .row_converter + .convert_rows(rows) + .expect("row conversion during emit"); + assert_eq!( + arrays.len(), + 1, + "Single field row converter must produce exactly one array, actual length is {}", + arrays.len() + ); + let array = arrays.pop().unwrap(); + encode_array_if_necessary(&array, &self.output_type) + .expect("dictionary re-encode during emit") + } + + /// Encode a whole incoming column into the row format. + fn convert(&self, array: &ArrayRef) -> Result { + self.row_converter + .convert_columns(std::slice::from_ref(array)) + .map_err(DataFusionError::from) + } +} + +impl GroupColumn for RowsGroupColumn { + fn equal_to(&self, lhs_row: usize, array: &ArrayRef, rhs_row: usize) -> bool { + // Scalar path (hash-collision remainder / streaming). Encode just the + // single incoming row rather than the whole column. The vectorized + // methods below encode the batch once; this path is expected to be rare. + let incoming = self + .convert(&array.slice(rhs_row, 1)) + .expect("row conversion during equal_to"); + self.group_values.row(lhs_row) == incoming.row(0) + } + + fn append_val(&mut self, array: &ArrayRef, row: usize) -> Result<()> { + let incoming = self.convert(&array.slice(row, 1))?; + self.group_values.push(incoming.row(0)); + Ok(()) + } + + fn vectorized_equal_to( + &self, + lhs_rows: &[usize], + array: &ArrayRef, + rhs_rows: &[usize], + equal_to_results: &mut BooleanBufferBuilder, + ) { + // Encode the incoming column once for the whole batch. + let incoming = self + .convert(array) + .expect("row conversion during vectorized_equal_to"); + for (idx, (&lhs_row, &rhs_row)) in + lhs_rows.iter().zip(rhs_rows.iter()).enumerate() + { + // Preserve the AND-accumulate contract: skip rows already false. + if !equal_to_results.get_bit(idx) { + continue; + } + if self.group_values.row(lhs_row) != incoming.row(rhs_row) { + equal_to_results.set_bit(idx, false); + } + } + } + + fn vectorized_append(&mut self, array: &ArrayRef, rows: &[usize]) -> Result<()> { + // Encode the incoming column once, then push the selected rows. + let incoming = self.convert(array)?; + for &row in rows { + self.group_values.push(incoming.row(row)); + } + Ok(()) + } + + fn len(&self) -> usize { + self.group_values.num_rows() + } + + fn size(&self) -> usize { + self.row_converter.size() + self.group_values.size() + } + + fn build(self: Box) -> ArrayRef { + self.rows_to_array(&self.group_values) + } + + fn take_n(&mut self, n: usize) -> ArrayRef { + debug_assert!(n <= self.group_values.num_rows()); + + // Materialize the first `n` group rows. + let output = self.rows_to_array(self.group_values.iter().take(n)); + + // Shift the remaining rows to the front by rebuilding the buffer. + // TODO: mirror the arrow-rs efficiency TODO in `GroupValuesRows::emit`. + let remaining_rows = self.group_values.num_rows() - n; + let remaining_bytes = self.group_values.lengths().skip(n).sum(); + let mut remaining = self + .row_converter + .empty_rows(remaining_rows, remaining_bytes); + for row in self.group_values.iter().skip(n) { + remaining.push(row); + } + self.group_values = remaining; + + output + } +} + +#[cfg(test)] +mod tests { + use super::*; + + use arrow::array::{ + Array, ArrayRef, FixedSizeListArray, Int32Array, StringArray, StructArray, + }; + use arrow::datatypes::{DataType, Field, Int32Type}; + use std::sync::Arc; + + fn fsl_i32(data: Vec>>>, list_len: i32) -> ArrayRef { + Arc::new(FixedSizeListArray::from_iter_primitive::( + data, list_len, + )) + } + + /// Build a `FixedSizeList` with `list_len == 1`. Each entry is one + /// row holding a single (optionally null) string, and an outer `None` + /// marks a null list. Variable-length string payloads give retained rows + /// distinct encoded lengths, which is what `take_n`'s byte preallocation + /// depends on. + fn fsl_utf8(rows: Vec>>) -> ArrayRef { + let child = StringArray::from( + rows.iter() + .map(|row| row.and_then(|inner| inner)) + .collect::>(), + ); + let outer_nulls = arrow::buffer::NullBuffer::from( + rows.iter().map(|row| row.is_some()).collect::>(), + ); + Arc::new(FixedSizeListArray::new( + Arc::new(Field::new("item", DataType::Utf8, true)), + 1, + Arc::new(child), + Some(outer_nulls), + )) + } + + /// The generic column must agree with a per-row reference for equality, + /// including inner-null and outer-null rows, on a `FixedSizeList`. + #[test] + fn fsl_append_equal_to_build_roundtrip() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 2, + ); + let mut col = Box::new(RowsGroupColumn::try_new(dt).unwrap()); + + // group values: [1,2], null-outer, [3, null-inner] + let input = fsl_i32( + vec![ + Some(vec![Some(1), Some(2)]), + None, + Some(vec![Some(3), None]), + ], + 2, + ); + + col.vectorized_append(&input, &[0, 1, 2]).unwrap(); + assert_eq!(col.len(), 3); + + // Probe with a fresh batch: row0 == group0, row1 (null) == group1, + // row2 differs from group0, row3 (inner null) == group2. + let probe = fsl_i32( + vec![ + Some(vec![Some(1), Some(2)]), // == g0 + None, // == g1 + Some(vec![Some(9), Some(9)]), // != g0 + Some(vec![Some(3), None]), // == g2 + ], + 2, + ); + + assert!(col.equal_to(0, &probe, 0)); + assert!(col.equal_to(1, &probe, 1)); + assert!(!col.equal_to(0, &probe, 2)); + assert!(col.equal_to(2, &probe, 3)); + + // Vectorized equal_to should match the scalar reference. + let mut results = BooleanBufferBuilder::new(3); + results.append_n(3, true); + col.vectorized_equal_to(&[0, 1, 2], &probe, &[0, 1, 3], &mut results); + assert!(results.get_bit(0)); + assert!(results.get_bit(1)); + assert!(results.get_bit(2)); + + // build() must reproduce the original group values. + let out = col.build(); + let out = out.as_any().downcast_ref::().unwrap(); + assert_eq!(out.len(), 3); + assert!(out.is_null(1)); + assert!(!out.is_null(0)); + } + + /// `take_n` must emit the first `n` rows and shift the rest to the front. + #[test] + fn fsl_take_n_shifts_remaining() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 1, + ); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + let input = fsl_i32( + vec![ + Some(vec![Some(10)]), + Some(vec![Some(20)]), + Some(vec![Some(30)]), + ], + 1, + ); + col.vectorized_append(&input, &[0, 1, 2]).unwrap(); + + let first = col.take_n(1); + let first = first.as_any().downcast_ref::().unwrap(); + let first_vals = first + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .clone(); + assert_eq!(first_vals.value(0), 10); + assert_eq!(col.len(), 2); + + // Remaining 20, 30 should now be at indices 0, 1. + let rest = Box::new(col).build(); + let rest = rest.as_any().downcast_ref::().unwrap(); + assert_eq!(rest.len(), 2); + let g0 = rest + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0); + assert_eq!(g0, 20); + } + + /// `take_n` preallocates the retained-row buffer from the known retained + /// row count and byte size + /// + /// To exercise the byte-sum path directly, the retained rows are + /// `FixedSizeList` values with deliberately unequal payload + /// lengths plus an inner-null. Here we assert every emitted and + /// every shifted-down value is byte-for-byte unchanged. + #[test] + fn take_n_preallocated_rebuild_preserves_variable_length_rows() { + let dt = DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Utf8, true)), + 1, + ); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + // Rows 0-2 are emitted; rows 3-6 are retained and shifted to the + // front. The retained rows intentionally have different encoded + // lengths so `lengths().skip(3).sum()` is not a simple row_count * k. + let input = fsl_utf8(vec![ + Some(Some("emit_a")), // 0: emitted + Some(None), // 1: emitted (inner-null) + None, // 2: emitted (outer-null) + Some(Some("")), // 3: retained, empty payload + Some(Some("xyz")), // 4: retained, short payload + Some(None), // 5: retained, inner-null + Some(Some("a_much_longer_payload_string")), // 6: retained, long payload + ]); + col.vectorized_append(&input, &[0, 1, 2, 3, 4, 5, 6]) + .unwrap(); + assert_eq!(col.len(), 7); + + // Emit the first three rows; four rows should remain. + let emitted = col.take_n(3); + let emitted = emitted + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(emitted.len(), 3); + assert_eq!( + emitted + .value(0) + .as_any() + .downcast_ref::() + .unwrap() + .value(0), + "emit_a" + ); + // Row 1 was an inner-null; row 2 was an outer-null. + assert!( + emitted + .value(1) + .as_any() + .downcast_ref::() + .unwrap() + .is_null(0) + ); + assert!(emitted.is_null(2)); + + assert_eq!(col.len(), 4); + + // The four retained rows must survive the rebuild intact, in order: + // "", "xyz", inner-null, "a_much_longer_payload_string". + let rest = Box::new(col).build(); + let rest = rest.as_any().downcast_ref::().unwrap(); + assert_eq!(rest.len(), 4); + + let value_at = |idx: usize| { + rest.value(idx) + .as_any() + .downcast_ref::() + .unwrap() + .clone() + }; + assert_eq!(value_at(0).value(0), ""); + assert_eq!(value_at(1).value(0), "xyz"); + assert!( + value_at(2).is_null(0), + "retained inner-null row must be preserved" + ); + assert_eq!(value_at(3).value(0), "a_much_longer_payload_string"); + } + + /// Works for `Struct` too — proves the column is type-generic. + #[test] + fn struct_roundtrip() { + let dt = DataType::Struct(vec![Field::new("a", DataType::Int32, true)].into()); + let mut col = RowsGroupColumn::try_new(dt).unwrap(); + + let a: ArrayRef = Arc::new(Int32Array::from(vec![Some(1), Some(2)])); + let input: ArrayRef = Arc::new(StructArray::new( + vec![Field::new("a", DataType::Int32, true)].into(), + vec![a], + None, + )); + col.vectorized_append(&input, &[0, 1]).unwrap(); + assert_eq!(col.len(), 2); + assert!(col.equal_to(0, &input, 0)); + assert!(!col.equal_to(0, &input, 1)); + } + + #[test] + fn supports_type_matches_row_converter_impl() { + assert!(RowsGroupColumn::supports_type(&DataType::FixedSizeList( + Arc::new(Field::new("item", DataType::Int32, true)), + 3 + ))); + assert!(RowsGroupColumn::supports_type(&DataType::Struct( + vec![Field::new("a", DataType::Int32, true)].into() + ))); + // Whether Map is encodable depends on the arrow-rs version. + // Just assert that our `supports_type` agrees with arrow's + // `RowConverter::supports_fields` — either both accept it or both + // reject it. Both are correct wrt the invariant. + let map_field = Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", DataType::Int32, true), + ] + .into(), + ), + false, + )); + let map_dt = DataType::Map(map_field, false); + let arrow_supports = + RowConverter::supports_fields(&[SortField::new(map_dt.clone())]); + assert_eq!(RowsGroupColumn::supports_type(&map_dt), arrow_supports); + } + + /// Regression test for the nested-container recursion in + /// [`crate::aggregates::group_values::row::encode_array_if_necessary`]. + /// `RowConverter` flattens dictionary values on the way in, so a + /// `List>` schema round-trips with `Utf8` values + /// unless the helper re-encodes the leaf. Without that recursion, + /// `build()` would emit an array whose data type does not match the + /// group column's declared type. + #[test] + fn build_preserves_list_of_dictionary_schema() { + use arrow::array::{DictionaryArray, ListArray, StringArray}; + use arrow::buffer::OffsetBuffer; + use arrow::datatypes::Int32Type; + + let dict_dt = + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)); + let item_field = Arc::new(Field::new("item", dict_dt.clone(), true)); + let outer_dt = DataType::List(Arc::clone(&item_field)); + + // Skip if this arrow-rs version rejects the nesting — the invariant we + // care about is `output().data_type() == declared type` conditional on + // supports_type saying yes. + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + // Build List> of one row = ["a", "b"]. + let values = Arc::new(StringArray::from(vec!["a", "b"])); + let keys = Int32Array::from(vec![0, 1]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = OffsetBuffer::from_lengths([2]); + let list = + ListArray::try_new(Arc::clone(&item_field), offsets, Arc::new(dict), None) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + col.vectorized_append(&input, &[0]).unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "build() must return the declared List data type, \ + not the RowConverter-flattened List", + ); + } + + // ---- FSL rejection ---------------------------------------- + // + // arrow-row 59.1.0's `decode_fixed_size_list` skips the + // dict-flatten correction that the generic `decode` path applies + // to `List` / `LargeList` / `ListView` / `LargeListView` / `Map`, + // so any `FixedSizeList` containing a `Dictionary` leaf panics on + // emit. `supports_type` must reject those shapes so + // `GroupValuesRows` fallback handles them instead. These tests pin + // the current shape of that black-list. + + fn dict_utf8() -> DataType { + DataType::Dictionary(Box::new(DataType::Int32), Box::new(DataType::Utf8)) + } + + fn fsl_of(inner: DataType) -> DataType { + DataType::FixedSizeList(Arc::new(Field::new("item", inner, true)), 2) + } + + #[test] + fn supports_type_rejects_fixed_size_list_of_dict() { + // Direct case: `FixedSizeList>`. + assert!(!RowsGroupColumn::supports_type(&fsl_of(dict_utf8()))); + } + + #[test] + fn supports_type_rejects_fsl_with_dict_nested_in_struct() { + // The dict is one level deep under a struct that is itself the + // FSL element. arrow-row still panics because `convert_raw` + // returns the struct with a decoded (Utf8) field while the + // FSL builder expects the declared struct-with-dict shape. + let struct_dt = DataType::Struct(vec![Field::new("d", dict_utf8(), true)].into()); + assert!(!RowsGroupColumn::supports_type(&fsl_of(struct_dt))); + } + + #[test] + fn supports_type_rejects_fsl_with_dict_nested_in_list() { + // `FixedSizeList>` — the inner `List` handles + // dicts correctly on its own, but the outer FSL wrapper still + // panics with the mismatched declared child type. + let list_of_dict = + DataType::List(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(!RowsGroupColumn::supports_type(&fsl_of(list_of_dict))); + } + + #[test] + fn supports_type_rejects_fsl_hidden_under_outer_list() { + // Sibling positioning: the outer container is a `List` (which is + // fine on its own), but its child is a `FixedSizeList`. + // The panic surface is at the inner FSL layer regardless of what + // wraps it, so this must still be rejected. + let outer = + DataType::List(Arc::new(Field::new("item", fsl_of(dict_utf8()), true))); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + #[test] + fn supports_type_rejects_fsl_hidden_under_outer_struct() { + // Same, but the outer wrapper is a struct. + let outer = + DataType::Struct(vec![Field::new("f", fsl_of(dict_utf8()), true)].into()); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + // ---- FSL without dicts is still fine ---------------------------- + + #[test] + fn supports_type_accepts_fsl_of_primitive() { + // Sanity: a plain FSL must not get caught by the + // dict-under-FSL blacklist. + assert!(RowsGroupColumn::supports_type(&fsl_of(DataType::Int32))); + } + + #[test] + fn supports_type_accepts_fsl_of_struct_without_dict() { + // FSL of struct where the struct's fields are all primitives. + let struct_dt = + DataType::Struct(vec![Field::new("a", DataType::Int32, true)].into()); + assert!(RowsGroupColumn::supports_type(&fsl_of(struct_dt))); + } + + // ---- Positive round-trip tests for non-FSL list-likes ----------- + // + // The other list-like decoders in arrow-row 59.1.0 + // (`GenericListArrayOrMap` path) apply the corrected_type fix, so + // `List`, `LargeList`, `ListView`, `LargeListView` + // and `Map<..., Dict>` all round-trip cleanly. These tests pin + // that they are (a) accepted by `supports_type` and (b) actually + // survive `vectorized_append` + `build()` without panicking, so a + // future arrow-rs regression there is caught here rather than in + // production. + + #[test] + fn supports_type_accepts_large_list_of_dict() { + let dt = DataType::LargeList(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_accepts_list_view_of_dict() { + let dt = DataType::ListView(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_accepts_large_list_view_of_dict() { + let dt = DataType::LargeListView(Arc::new(Field::new("item", dict_utf8(), true))); + assert!(RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_map_agrees_with_row_converter() { + // Map>. Whether arrow-row supports Map + // depends on the version; either way, our `supports_type` must + // agree with `RowConverter::supports_fields` — otherwise we'd + // pick a strategy the converter can't back. + let entries = Arc::new(Field::new( + "entries", + DataType::Struct( + vec![ + Field::new("keys", DataType::Int32, false), + Field::new("values", dict_utf8(), true), + ] + .into(), + ), + false, + )); + let map_dt = DataType::Map(entries, false); + let arrow_supports = + RowConverter::supports_fields(&[SortField::new(map_dt.clone())]); + assert_eq!(RowsGroupColumn::supports_type(&map_dt), arrow_supports); + } + + /// End-to-end regression: `LargeList>` must + /// actually survive `vectorized_append` + `build()` on the current + /// arrow-rs version, not just be accepted by `supports_type`. + #[test] + fn build_preserves_large_list_of_dictionary_schema() { + use arrow::array::{DictionaryArray, LargeListArray, StringArray}; + use arrow::buffer::OffsetBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::LargeList(Arc::clone(&item_field)); + + // Skip if this arrow-rs version rejects the nesting (defensive: + // the invariant we care about is `output().data_type() == declared` + // conditional on `supports_type` saying yes). + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + let values = Arc::new(StringArray::from(vec!["a", "b"])); + let keys = Int32Array::from(vec![0, 1]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = OffsetBuffer::::from_lengths([2]); + let list = LargeListArray::try_new( + Arc::clone(&item_field), + offsets, + Arc::new(dict), + None, + ) + .unwrap(); + + col.vectorized_append(&(Arc::new(list) as ArrayRef), &[0]) + .unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "LargeList: build() must preserve the declared type", + ); + } + + /// Build a two-row `ListView>` array with rows + /// `["a", "b"]` and `["c"]` — the shape from the review reproducer: + /// `arrow_cast(a, 'ListView(Dictionary(Int32, Utf8))')`. + fn list_view_of_dict_input() -> (DataType, ArrayRef) { + use arrow::array::{DictionaryArray, ListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::ListView(Arc::clone(&item_field)); + + let values = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::::from(vec![0, 2]); + let sizes = ScalarBuffer::::from(vec![2, 1]); + let list = ListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + (outer_dt, Arc::new(list) as ArrayRef) + } + + /// `ListView`: arrow-row's `decode_list_view` flattens the + /// dictionary child (`corrected_type`), so `build` must re-encode + /// the emitted array back to the declared type. Regression for the + /// review reproducer that failed with + /// `expected ListView(Dictionary(Int32, Utf8)) but found ListView(Utf8)`. + #[test] + fn build_preserves_list_view_of_dictionary_schema() { + let (outer_dt, input) = list_view_of_dict_input(); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + assert_eq!(col.len(), 2); + + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "ListView: build() must return the declared type, \ + not the RowConverter-flattened ListView", + ); + assert_eq!(built.len(), 2); + } + + /// Same regression through the `take_n` path (used by + /// `EmitTo::First(n)`), including the type of the *remaining* + /// values emitted by a subsequent `build`. + #[test] + fn take_n_preserves_list_view_of_dictionary_schema() { + let (outer_dt, input) = list_view_of_dict_input(); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + + let taken = col.take_n(1); + assert_eq!( + taken.data_type(), + &outer_dt, + "ListView: take_n() must return the declared type", + ); + assert_eq!(taken.len(), 1); + + let rest = col.build(); + assert_eq!( + rest.data_type(), + &outer_dt, + "ListView: build() after take_n must also preserve the type", + ); + assert_eq!(rest.len(), 1); + } + + /// `LargeListView` fails the same way as `ListView` + /// per the review; cover both `build` and `take_n`. + #[test] + fn build_and_take_n_preserve_large_list_view_of_dictionary_schema() { + use arrow::array::{DictionaryArray, LargeListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::LargeListView(Arc::clone(&item_field)); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let values = Arc::new(StringArray::from(vec!["a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::::from(vec![0, 2]); + let sizes = ScalarBuffer::::from(vec![2, 1]); + let list = LargeListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + col.vectorized_append(&input, &[0, 1]).unwrap(); + + let taken = col.take_n(1); + assert_eq!( + taken.data_type(), + &outer_dt, + "LargeListView: take_n() must return the declared type", + ); + + let rest = col.build(); + assert_eq!( + rest.data_type(), + &outer_dt, + "LargeListView: build() must return the declared type", + ); + assert_eq!(rest.len(), 1); + } + + /// Group-identity must survive the dictionary flatten + re-encode + /// round trip: appending the same logical list twice (with distinct + /// dictionary key mappings) must map to one group, a different list + /// to another. Mirrors the review reproducer's GROUP BY semantics + /// (2 distinct groups from 3 input rows). + #[test] + fn list_view_of_dict_groups_by_logical_value() { + use arrow::array::{DictionaryArray, ListViewArray, StringArray}; + use arrow::buffer::ScalarBuffer; + + let item_field = Arc::new(Field::new("item", dict_utf8(), true)); + let outer_dt = DataType::ListView(Arc::clone(&item_field)); + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + // Rows: ["a","b"], ["a","b"], ["c"] → 2 distinct groups. + let values = Arc::new(StringArray::from(vec!["a", "b", "a", "b", "c"])); + let keys = Int32Array::from(vec![0, 1, 2, 3, 4]); + let dict = DictionaryArray::::try_new(keys, values).unwrap(); + let offsets = ScalarBuffer::::from(vec![0, 2, 4]); + let sizes = ScalarBuffer::::from(vec![2, 2, 1]); + let list = ListViewArray::try_new( + Arc::clone(&item_field), + offsets, + sizes, + Arc::new(dict), + None, + ) + .unwrap(); + let input: ArrayRef = Arc::new(list); + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + // Append row 0 as group 0. + col.vectorized_append(&input, &[0]).unwrap(); + // Row 1 must compare equal to group 0 (same logical value). + assert!( + col.equal_to(0, &input, 1), + "identical logical lists must be equal regardless of dict keys", + ); + // Row 2 must not. + assert!( + !col.equal_to(0, &input, 2), + "different logical lists must not be equal", + ); + + col.vectorized_append(&input, &[2]).unwrap(); + assert_eq!(col.len(), 2, "3 input rows → 2 distinct groups"); + + let built = col.build(); + assert_eq!(built.data_type(), &outer_dt); + assert_eq!(built.len(), 2); + } + + /// End-to-end regression for `Map>` when + /// arrow-row supports it. Same intent as the LargeList test. + #[test] + fn build_preserves_map_of_dictionary_schema() { + use arrow::array::{ + DictionaryArray, Int32Array, MapArray, StringArray, StructArray, + }; + use arrow::buffer::OffsetBuffer; + + let key_field = Arc::new(Field::new("keys", DataType::Int32, false)); + let value_field = Arc::new(Field::new("values", dict_utf8(), true)); + let entries_field = Arc::new(Field::new( + "entries", + DataType::Struct(vec![(*key_field).clone(), (*value_field).clone()].into()), + false, + )); + let outer_dt = DataType::Map(Arc::clone(&entries_field), false); + + if !RowsGroupColumn::supports_type(&outer_dt) { + return; + } + + let mut col = Box::new(RowsGroupColumn::try_new(outer_dt.clone()).unwrap()); + + // One map entry: {1 -> "a"}. + let keys = Arc::new(Int32Array::from(vec![1])) as ArrayRef; + let values_arr = Arc::new(StringArray::from(vec!["a"])); + let value_keys = Int32Array::from(vec![0]); + let value_dict = + DictionaryArray::::try_new(value_keys, values_arr).unwrap(); + let entries = StructArray::try_new( + vec![(*key_field).clone(), (*value_field).clone()].into(), + vec![keys, Arc::new(value_dict)], + None, + ) + .unwrap(); + let offsets = OffsetBuffer::::from_lengths([1]); + let map = + MapArray::try_new(Arc::clone(&entries_field), offsets, entries, None, false) + .unwrap(); + + col.vectorized_append(&(Arc::new(map) as ArrayRef), &[0]) + .unwrap(); + let built = col.build(); + assert_eq!( + built.data_type(), + &outer_dt, + "Map<..., Dict>: build() must preserve the declared type", + ); + } + + // ---- Union / RunEndEncoded defensive rejection ----------------- + // + // Before this PR both types were routed to `GroupValuesRows` + // (`group_column_supported_type` didn't have a nested branch). This + // PR added `is_nested`-based dispatch to `RowsGroupColumn`, which + // would opt them in — but the arrow-row round-trip for these two + // families hasn't been covered by our tests. Reject them here so + // the pre-PR routing is preserved; drop the blacklist when the + // round-trip matrix grows to include them. + + #[test] + fn supports_type_rejects_union() { + use arrow::datatypes::UnionFields; + + let fields = UnionFields::try_new( + vec![0_i8, 1_i8], + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ], + ) + .unwrap(); + let dt = DataType::Union(fields, arrow::datatypes::UnionMode::Dense); + assert!( + !RowsGroupColumn::supports_type(&dt), + "Union must fall back to GroupValuesRows until arrow-row \ + round-trip is covered by our tests", + ); + } + + #[test] + fn supports_type_rejects_run_end_encoded_with_nested_values() { + // REE with `is_nested() = true` (nested values) is what this PR + // could otherwise opt into RowsGroupColumn; keep it on + // GroupValuesRows. + let list_of_i32 = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + let dt = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", list_of_i32, true)), + ); + assert!(!RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_rejects_run_end_encoded_with_scalar_values() { + // REE with scalar values is `is_nested() == false`, so + // `group_column_supported_type` never routes it to us via the + // nested branch anyway — but pin the invariant explicitly so a + // future refactor doesn't accidentally opt it in. + let dt = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ); + assert!(!RowsGroupColumn::supports_type(&dt)); + } + + #[test] + fn supports_type_rejects_ree_hidden_under_outer_wrapper() { + // REE buried under a struct or list: still rejected because + // the wrapper's decoder recurses through the REE branch we + // haven't validated. + let ree = DataType::RunEndEncoded( + Arc::new(Field::new("run_ends", DataType::Int32, false)), + Arc::new(Field::new("values", DataType::Utf8, true)), + ); + let outer = DataType::Struct(vec![Field::new("f", ree, true)].into()); + assert!(!RowsGroupColumn::supports_type(&outer)); + } + + #[test] + fn supports_type_accepts_plain_list_and_struct_still() { + // Sanity: the defensive Union/REE blacklist must not accidentally + // catch the well-tested list-likes / structs that this column + // exists to serve. + let list_of_int = + DataType::List(Arc::new(Field::new("item", DataType::Int32, true))); + assert!(RowsGroupColumn::supports_type(&list_of_int)); + + let struct_of_prims = DataType::Struct( + vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Utf8, true), + ] + .into(), + ); + assert!(RowsGroupColumn::supports_type(&struct_of_prims)); + } +} diff --git a/datafusion/physical-plan/src/aggregates/group_values/row.rs b/datafusion/physical-plan/src/aggregates/group_values/row.rs index a3bd31f76c233..cbd7a609c5caa 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/row.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/row.rs @@ -17,7 +17,8 @@ use crate::aggregates::group_values::GroupValues; use arrow::array::{ - Array, ArrayRef, ListArray, PrimitiveArray, RunArray, StructArray, + Array, ArrayRef, FixedSizeListArray, LargeListArray, LargeListViewArray, ListArray, + ListViewArray, MapArray, PrimitiveArray, RunArray, StructArray, downcast_run_end_index, }; use arrow::compute::cast; @@ -26,6 +27,7 @@ use arrow::row::{RowConverter, Rows, SortField}; use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; +use datafusion_common::utils::normalize_float_zero; use datafusion_execution::memory_pool::proxy::{HashTableAllocExt, VecAllocExt}; use datafusion_expr::EmitTo; use hashbrown::hash_table::HashTable; @@ -116,6 +118,13 @@ impl GroupValuesRows { impl GroupValues for GroupValuesRows { fn intern(&mut self, cols: &[ArrayRef], groups: &mut Vec) -> Result<()> { + // Normalize -0.0 → +0.0 so RowConverter (IEEE 754 totalOrder) and + // primitive hashing both group ±0 together. No-op for non-float + // columns. + let normalized_cols: Vec = + cols.iter().map(normalize_float_zero).collect(); + let cols = normalized_cols.as_slice(); + // Convert the group keys into the row format let group_rows = &mut self.rows_buffer; group_rows.clear(); @@ -239,7 +248,7 @@ impl GroupValues for GroupValuesRows { // https://github.com/apache/datafusion/issues/7647 for (field, array) in self.schema.fields.iter().zip(&mut output) { let expected = field.data_type(); - *array = dictionary_encode_if_necessary(array, expected)?; + *array = encode_array_if_necessary(array, expected)?; } self.group_values = Some(group_values); @@ -259,7 +268,17 @@ impl GroupValues for GroupValuesRows { } } -fn dictionary_encode_if_necessary( +/// Re-apply dictionary / run-end encoding to `array` so it matches `expected`. +/// +/// Arrow's [`RowConverter`] flattens dictionary and run-end-encoded values to +/// their plain value type during row encoding (at [`RowConverter::append`]), +/// so any group-value array produced from the row format is in that plain +/// type and must be re-encoded to match the schema's expected type before +/// being returned. Shared with the generic row-backed `GroupColumn`. +/// +/// [`RowConverter`]: arrow::row::RowConverter +/// [`RowConverter::append`]: arrow::row::RowConverter::append +pub(crate) fn encode_array_if_necessary( array: &ArrayRef, expected: &DataType, ) -> Result { @@ -270,7 +289,7 @@ fn dictionary_encode_if_necessary( .iter() .zip(struct_array.columns()) .map(|(expected_field, column)| { - dictionary_encode_if_necessary(column, expected_field.data_type()) + encode_array_if_necessary(column, expected_field.data_type()) }) .collect::>>()?; @@ -286,13 +305,82 @@ fn dictionary_encode_if_necessary( Ok(Arc::new(ListArray::try_new( Arc::::clone(expected_field), list.offsets().clone(), - dictionary_encode_if_necessary( - list.values(), - expected_field.data_type(), - )?, + encode_array_if_necessary(list.values(), expected_field.data_type())?, list.nulls().cloned(), )?)) } + (DataType::LargeList(expected_field), &DataType::LargeList(_)) => { + let list = array.as_any().downcast_ref::().unwrap(); + + Ok(Arc::new(LargeListArray::try_new( + Arc::::clone(expected_field), + list.offsets().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::ListView(expected_field), &DataType::ListView(_)) => { + // arrow-row's `decode_list_view` applies the dictionary-flatten + // `corrected_type` to the child, so a `ListView>` + // decodes as `ListView` and the child must be + // re-encoded here (same as `List` above, plus the `sizes` + // buffer that view-lists carry). + let list = array.as_any().downcast_ref::().unwrap(); + + Ok(Arc::new(ListViewArray::try_new( + Arc::::clone(expected_field), + list.offsets().clone(), + list.sizes().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::LargeListView(expected_field), &DataType::LargeListView(_)) => { + let list = array.as_any().downcast_ref::().unwrap(); + + Ok(Arc::new(LargeListViewArray::try_new( + Arc::::clone(expected_field), + list.offsets().clone(), + list.sizes().clone(), + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + ( + DataType::FixedSizeList(expected_field, expected_size), + &DataType::FixedSizeList(_, _), + ) => { + let list = array.as_any().downcast_ref::().unwrap(); + + Ok(Arc::new(FixedSizeListArray::try_new( + Arc::::clone(expected_field), + *expected_size, + encode_array_if_necessary(list.values(), expected_field.data_type())?, + list.nulls().cloned(), + )?)) + } + (DataType::Map(expected_entries_field, ordered), &DataType::Map(_, _)) => { + let map = array.as_any().downcast_ref::().unwrap(); + // Re-encode the entries `StructArray` (which holds key/value + // columns) against the expected entries field's struct type. + let entries_as_ref: ArrayRef = Arc::new(map.entries().clone()); + let entries = encode_array_if_necessary( + &entries_as_ref, + expected_entries_field.data_type(), + )?; + let entries = entries + .as_any() + .downcast_ref::() + .expect("Map entries recurse must yield a StructArray") + .clone(); + Ok(Arc::new(MapArray::try_new( + Arc::::clone(expected_entries_field), + map.offsets().clone(), + entries, + map.nulls().cloned(), + *ordered, + )?)) + } (DataType::Dictionary(_, _), _) => Ok(cast(array.as_ref(), expected)?), ( DataType::RunEndEncoded(run_ends_field, expected_values_field), @@ -304,7 +392,7 @@ fn dictionary_encode_if_necessary( .as_any() .downcast_ref::>() .unwrap(); - let values = dictionary_encode_if_necessary( + let values = encode_array_if_necessary( &(Arc::clone(run_array.values()) as ArrayRef), expected_values_field.data_type(), )?; diff --git a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs index efaf7eba0f1b5..e254aebcfd7ce 100644 --- a/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs +++ b/datafusion/physical-plan/src/aggregates/group_values/single_group_by/primitive.rs @@ -24,6 +24,7 @@ use arrow::array::{ use arrow::datatypes::{DataType, i256}; use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; +use datafusion_common::utils::split_vec_min_alloc; use datafusion_execution::memory_pool::proxy::VecAllocExt; use datafusion_expr::EmitTo; use half::f16; @@ -34,8 +35,21 @@ use std::mem::size_of; use std::sync::Arc; /// A trait to allow hashing of floating point numbers -pub(crate) trait HashValue { +pub trait HashValue { fn hash(&self, state: &RandomState) -> u64; + + /// Return a canonical representative whose bit pattern is identical for + /// all values that should be grouped together. Default is the identity; + /// floats override this to fold `-0.0` into `+0.0` so the bit-equal + /// `is_eq` check used during insertion treats them as the same group. + /// NaN payload bits are preserved. + #[inline] + fn canonicalize(self) -> Self + where + Self: Sized, + { + self + } } macro_rules! hash_integer { @@ -62,13 +76,20 @@ macro_rules! hash_float { $(impl HashValue for $t { #[cfg(not(feature = "force_hash_collisions"))] fn hash(&self, state: &RandomState) -> u64 { - state.hash_one(self.to_bits()) + state.hash_one(self.canonicalize().to_bits()) } #[cfg(feature = "force_hash_collisions")] fn hash(&self, _state: &RandomState) -> u64 { 0 } + + #[inline] + fn canonicalize(self) -> Self { + let bits = self.to_bits(); + let bits = if bits << 1 == 0 { 0 } else { bits }; + Self::from_bits(bits) + } })+ }; } @@ -126,6 +147,10 @@ where group_id }), Some(key) => { + // Fold equivalence-class duplicates (e.g. `-0.0` → `+0.0`) + // so the bit-equal `is_eq` matches and the stored value is + // the canonical representative. + let key = key.canonicalize(); let state = &self.random_state; let hash = key.hash(state); let insert = self.map.entry( @@ -207,9 +232,7 @@ where Some(_) => self.null_group.take(), None => None, }; - let mut split = self.values.split_off(n); - std::mem::swap(&mut self.values, &mut split); - build_primitive(split, null_group) + build_primitive(split_vec_min_alloc(&mut self.values, n), null_group) } }; @@ -223,3 +246,51 @@ where self.map.shrink_to(num_rows, |_| 0); // hasher does not matter since the map is cleared } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::types::Int32Type; + use arrow::array::{ArrayRef, Int32Array}; + use arrow::datatypes::DataType; + use datafusion_expr::EmitTo; + use std::sync::Arc; + + /// Mirror of the `EmitTo::take_needed` regression test, applied to the + /// concrete `GroupValuesPrimitive` accumulator. + /// + /// When `n` is small, the old `split_off(n) + swap` pattern used inside + /// `emit(EmitTo::First(n))` left `self.values` with a small fresh allocation + /// and returned the emitted prefix carrying the original large backing. + /// + /// With `split_vec_min_alloc` and `n * 2 <= len`, the drain branch is taken: + /// the emitted prefix gets a compact allocation and `self.values` retains the + /// original large one. + #[test] + fn emit_first_small_n_allocates_minimally() -> Result<()> { + let mut gv = GroupValuesPrimitive::::new(DataType::Int32); + + // Intern 20 distinct values; `new()` pre-allocates capacity 128 for `values`. + let arr: ArrayRef = Arc::new(Int32Array::from_iter_values(0..20i32)); + let mut groups = vec![]; + gv.intern(&[arr], &mut groups)?; + let capacity_before = gv.values.capacity(); // 128 + + // n=4, n*2=8 <= len=20 -> drain branch + let emitted = gv.emit(EmitTo::First(4))?; + + assert_eq!(emitted[0].len(), 4); + + // `self.values` must retain its original large allocation. + // Old split_off+swap left it with a fresh small allocation (~16). + assert_eq!( + gv.values.capacity(), + capacity_before, + "self.values capacity {} should equal original {} after small First(n) emit", + gv.values.capacity(), + capacity_before, + ); + + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/row_hash.rs b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs similarity index 86% rename from datafusion/physical-plan/src/aggregates/row_hash.rs rename to datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs index b4ac7d060576f..99c101199459f 100644 --- a/datafusion/physical-plan/src/aggregates/row_hash.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_hash_stream.rs @@ -22,6 +22,7 @@ use std::task::{Context, Poll}; use std::vec; use super::order::GroupOrdering; +use super::skip_partial::SkipAggregationProbe; use super::{AggregateExec, format_human_display}; use crate::aggregates::group_values::{GroupByMetrics, GroupValues, new_group_values}; use crate::aggregates::order::GroupOrderingFull; @@ -118,100 +119,6 @@ struct SpillState { // Metrics related to spilling are managed inside `spill_manager` } -/// Tracks if the aggregate should skip partial aggregations -/// -/// See "partial aggregation" discussion on [`GroupedHashAggregateStream`] -struct SkipAggregationProbe { - // ======================================================================== - // PROPERTIES: - // These fields are initialized at the start and remain constant throughout - // the execution. - // ======================================================================== - /// Aggregation ratio check performed when the number of input rows exceeds - /// this threshold (from `SessionConfig`) - probe_rows_threshold: usize, - /// Maximum ratio of `num_groups` to `input_rows` for continuing aggregation - /// (from `SessionConfig`). If the ratio exceeds this value, aggregation - /// is skipped and input rows are directly converted to output - probe_ratio_threshold: f64, - - // ======================================================================== - // STATES: - // Fields changes during execution. Can be buffer, or state flags that - // influence the execution in parent `GroupedHashAggregateStream` - // ======================================================================== - /// Number of processed input rows (updated during probing) - input_rows: usize, - /// Number of total group values for `input_rows` (updated during probing) - num_groups: usize, - - /// Flag indicating further data aggregation may be skipped (decision made - /// when probing complete) - should_skip: bool, - /// Flag indicating further updates of `SkipAggregationProbe` state won't - /// make any effect (set either while probing or on probing completion) - is_locked: bool, - - // ======================================================================== - // METRICS: - // ======================================================================== - /// Number of rows where state was output without aggregation. - /// - /// * If 0, all input rows were aggregated (should_skip was always false) - /// - /// * if greater than zero, the number of rows which were output directly - /// without aggregation - skipped_aggregation_rows: metrics::Count, -} - -impl SkipAggregationProbe { - fn new( - probe_rows_threshold: usize, - probe_ratio_threshold: f64, - skipped_aggregation_rows: metrics::Count, - ) -> Self { - Self { - input_rows: 0, - num_groups: 0, - probe_rows_threshold, - probe_ratio_threshold, - should_skip: false, - is_locked: false, - skipped_aggregation_rows, - } - } - - /// Updates `SkipAggregationProbe` state: - /// - increments the number of input rows - /// - replaces the number of groups with the new value - /// - on `probe_rows_threshold` exceeded calculates - /// aggregation ratio and sets `should_skip` flag - /// - if `should_skip` is set, locks further state updates - fn update_state(&mut self, input_rows: usize, num_groups: usize) { - if self.is_locked { - return; - } - self.input_rows += input_rows; - self.num_groups = num_groups; - if self.input_rows >= self.probe_rows_threshold { - self.should_skip = self.num_groups as f64 / self.input_rows as f64 - >= self.probe_ratio_threshold; - // Set is_locked to true only if we have decided to skip, otherwise we can try to skip - // during processing the next record_batch. - self.is_locked = self.should_skip; - } - } - - fn should_skip(&self) -> bool { - self.should_skip - } - - /// Record the number of rows that were output directly without aggregation - fn record_skipped(&mut self, batch: &RecordBatch) { - self.skipped_aggregation_rows.add(batch.num_rows()); - } -} - /// Controls the behavior when an out-of-memory condition occurs. #[derive(PartialEq, Debug)] enum OutOfMemoryMode { @@ -225,6 +132,16 @@ enum OutOfMemoryMode { /// HashTable based Grouping Aggregator /// +/// # Development Note +/// +/// This implementation is being incrementally refactored. See the tracking issue +/// for details. +/// +/// New features and improvements should go directly into the new implementation. +/// Please coordinate through the tracking issue. +/// +/// Issue: +/// /// # Design Goals /// /// This structure is designed so that updating the aggregates can be @@ -262,7 +179,7 @@ enum OutOfMemoryMode { /// /// group_values accumulators /// -/// ``` +/// ``` /// /// For example, given a query like `COUNT(x), SUM(y) ... GROUP BY z`, /// [`group_values`] will store the distinct values of `z`. There will @@ -300,8 +217,7 @@ enum OutOfMemoryMode { /// aggregator must store the intermediate state for each group. /// /// If the ratio of the number of groups to the number of input rows exceeds a -/// threshold, and [`GroupsAccumulator::supports_convert_to_state`] is -/// supported, this operator will stop applying Partial aggregation and directly +/// threshold, this operator will stop applying Partial aggregation and directly /// pass the input rows to the next aggregation phase. /// /// [`Accumulator::state`]: datafusion_expr::Accumulator::state @@ -620,8 +536,7 @@ impl GroupedHashAggregateStream { merging_aggregate_arguments, merging_group_by: PhysicalGroupBy::new_single(merging_group_by_expr), peak_mem_used: MetricBuilder::new(&agg.metrics) - .with_category(MetricCategory::Bytes) - .gauge("peak_mem_used", partition), + .peak_memory_usage("peak_mem_used", partition), spill_manager, }; @@ -629,14 +544,9 @@ impl GroupedHashAggregateStream { // - aggregation mode is Partial // - input is not ordered by GROUP BY expressions, // since Final mode expects unique group values as its input - // - all accumulators support input batch to intermediate - // aggregate state conversion // - there is only one GROUP BY expressions set let skip_aggregation_probe = if agg.mode == AggregateMode::Partial && matches!(group_ordering, GroupOrdering::None) - && accumulators - .iter() - .all(|acc| acc.supports_convert_to_state()) && agg_group_by.is_single() { let options = &context.session_config().options().execution; @@ -644,14 +554,20 @@ impl GroupedHashAggregateStream { options.skip_partial_aggregation_probe_rows_threshold; let probe_ratio_threshold = options.skip_partial_aggregation_probe_ratio_threshold; - let skipped_aggregation_rows = MetricBuilder::new(&agg.metrics) - .with_category(MetricCategory::Rows) - .counter("skipped_aggregation_rows", partition); - Some(SkipAggregationProbe::new( - probe_rows_threshold, - probe_ratio_threshold, - skipped_aggregation_rows, - )) + // A threshold >= 1.0 means the ratio (num_groups / input_rows) can + // never exceed it, so the feature is effectively disabled. + if probe_ratio_threshold >= 1.0 { + None + } else { + let skipped_aggregation_rows = MetricBuilder::new(&agg.metrics) + .with_category(MetricCategory::Rows) + .counter("skipped_aggregation_rows", partition); + Some(SkipAggregationProbe::new( + probe_rows_threshold, + probe_ratio_threshold, + skipped_aggregation_rows, + )) + } } else { None }; @@ -1010,7 +926,7 @@ impl GroupedHashAggregateStream { // if aggregation is over intermediate states, // use merge - acc.merge_batch(values, group_indices, None, total_num_groups)?; + acc.merge_batch(values, group_indices, total_num_groups)?; } self.group_by_metrics .aggregation_time @@ -1198,17 +1114,31 @@ impl GroupedHashAggregateStream { // Prime each accumulator for the registered group count with no data. // // We build 1-row null arrays for each aggregate argument and pass them - // with an all-false filter. The filter ensures no row is accumulated - // into any group, which keeps every group in its "zero" initial state - // (NULL for SUM/AVG/MIN/MAX, 0 for COUNT). + // with an all-false filter to update_batch. The filter ensures no row + // is accumulated into any group, which keeps every group in its "zero" + // initial state (NULL for SUM/AVG/MIN/MAX, 0 for COUNT). // // Using a 1-row batch rather than 0 rows is required to avoid a fast // path in `NullState::accumulate` that treats "0 nulls in a 0-row // array" as "all groups have been seen", which would cause SUM to // return 0 instead of NULL. // - // Argument types are inferred directly from the expression metadata so - // we never need to construct a full `RecordBatch`. + // This path always runs in a Raw input mode, so `update_batch` (not + // `merge_batch`) is the right entry point: + // + // - `has_grouping_set()` can only be true for the Partial / Single / + // SinglePartitioned modes, whose `input_mode()` is `Raw`. The final + // modes rebuild their group-by via `PhysicalGroupBy::as_final()`, + // which clears `has_grouping_set`, so this method returns early for + // them and never reaches here. + // + // Since every row is filtered out, the actual data content never + // matters. The assert documents and guards the invariant above. + debug_assert_eq!( + self.mode.input_mode(), + AggregateInputMode::Raw, + "init_empty_grouping_sets must only run in a Raw input mode" + ); let total_groups = self.group_values.len(); let null_args: Vec> = self .aggregate_arguments @@ -1224,11 +1154,7 @@ impl GroupedHashAggregateStream { .collect::>>()?; let false_filter = BooleanArray::from(vec![false]); for (acc, args) in self.accumulators.iter_mut().zip(null_args.iter()) { - if self.mode.input_mode() == AggregateInputMode::Raw { - acc.update_batch(args, &[0], Some(&false_filter), total_groups)?; - } else { - acc.merge_batch(args, &[0], Some(&false_filter), total_groups)?; - } + acc.update_batch(args, &[0], Some(&false_filter), total_groups)?; } } @@ -1473,7 +1399,6 @@ impl GroupedHashAggregateStream { mod tests { use super::*; use crate::InputOrderMode; - use crate::execution_plan::ExecutionPlan; use crate::test::TestMemoryExec; use arrow::array::{Int32Array, Int64Array}; use arrow::datatypes::{DataType, Field, Schema}; @@ -1482,6 +1407,8 @@ mod tests { use datafusion_physical_expr::aggregate::AggregateExprBuilder; use datafusion_physical_expr::expressions::col; + // Migrated to PartialHashAggregateStream coverage in hash_stream.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] async fn test_double_emission_race_condition_bug() -> Result<()> { // Fix for https://github.com/apache/datafusion/issues/18701 @@ -1588,152 +1515,8 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_skip_aggregation_probe_not_locked_until_skip() -> Result<()> { - // Test that the probe is not locked until we actually decide to skip. - // This allows us to continue evaluating the skip condition across multiple batches. - // - // Scenario: - // - Batch 1: Hits rows threshold but NOT ratio threshold (low cardinality) -> don't skip - // - Batch 2: Now hits ratio threshold (high cardinality) -> skip - // - // Without the fix, the probe would be locked after batch 1, preventing the skip - // decision from being made on batch 2. - - let schema = Arc::new(Schema::new(vec![ - Field::new("group_col", DataType::Int32, false), - Field::new("value_col", DataType::Int32, false), - ])); - - // Configure thresholds: - // - probe_rows_threshold: 100 rows - // - probe_ratio_threshold: 0.8 (80%) - let probe_rows_threshold = 100; - let probe_ratio_threshold = 0.8; - - // Batch 1: 100 rows with only 10 unique groups - // Ratio: 10/100 = 0.1 (10%) < 0.8 -> should NOT skip - // This will hit the rows threshold but not the ratio threshold - let batch1_rows = 100; - let batch1_groups = 10; - let mut group_ids_batch1 = Vec::new(); - for i in 0..batch1_rows { - group_ids_batch1.push((i % batch1_groups) as i32); - } - let values_batch1: Vec = vec![1; batch1_rows]; - - let batch1 = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(group_ids_batch1)), - Arc::new(Int32Array::from(values_batch1)), - ], - )?; - - // Batch 2: 350 rows with 350 unique NEW groups (starting from group 10) - // After batch 2, total: 450 rows, 360 groups - // Ratio: 360/450 = 0.8 (80%) >= 0.8 -> SHOULD decide to skip - let batch2_rows = 350; - let batch2_groups = 350; - let group_ids_batch2: Vec = (batch1_groups..(batch1_groups + batch2_groups)) - .map(|x| x as i32) - .collect(); - let values_batch2: Vec = vec![1; batch2_rows]; - - let batch2 = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(group_ids_batch2)), - Arc::new(Int32Array::from(values_batch2)), - ], - )?; - - // Batch 3: This batch should be skipped since we decided to skip after batch 2 - // 100 rows with 100 unique groups (continuing from where batch 2 left off) - let batch3_rows = 100; - let batch3_groups = 100; - let batch3_start_group = batch1_groups + batch2_groups; - let group_ids_batch3: Vec = (batch3_start_group - ..(batch3_start_group + batch3_groups)) - .map(|x| x as i32) - .collect(); - let values_batch3: Vec = vec![1; batch3_rows]; - - let batch3 = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(Int32Array::from(group_ids_batch3)), - Arc::new(Int32Array::from(values_batch3)), - ], - )?; - - let input_partitions = vec![vec![batch1, batch2, batch3]]; - - let runtime = RuntimeEnvBuilder::default().build_arc()?; - let mut task_ctx = TaskContext::default().with_runtime(runtime); - - // Configure skip aggregation settings - let mut session_config = task_ctx.session_config().clone(); - session_config = session_config.set( - "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", - &datafusion_common::ScalarValue::UInt64(Some(probe_rows_threshold)), - ); - session_config = session_config.set( - "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", - &datafusion_common::ScalarValue::Float64(Some(probe_ratio_threshold)), - ); - task_ctx = task_ctx.with_session_config(session_config); - let task_ctx = Arc::new(task_ctx); - - // Create aggregate: COUNT(*) GROUP BY group_col - let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; - let aggr_expr = vec![Arc::new( - AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("count_value") - .build()?, - )]; - - let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; - let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); - - // Use Partial mode - let aggregate_exec = AggregateExec::try_new( - AggregateMode::Partial, - PhysicalGroupBy::new_single(group_expr), - aggr_expr, - vec![None], - exec, - Arc::clone(&schema), - )?; - - // Execute and collect results - let mut stream = - GroupedHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; - let mut results = Vec::new(); - - while let Some(result) = stream.next().await { - let batch = result?; - results.push(batch); - } - - // Check that skip aggregation actually happened - // The key metric is skipped_aggregation_rows - let metrics = aggregate_exec.metrics().unwrap(); - let skipped_rows = metrics - .sum_by_name("skipped_aggregation_rows") - .map(|m| m.as_usize()) - .unwrap_or(0); - - // We expect batch 3's rows to be skipped (100 rows) - assert_eq!( - skipped_rows, batch3_rows, - "Expected batch 3's rows ({batch3_rows}) to be skipped", - ); - - Ok(()) - } - + // Migrated to OrderedPartialAggregateStream coverage in aggregates/mod.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. #[tokio::test] async fn test_emit_early_with_partially_sorted() -> Result<()> { // Reproducer for #20445: EmitEarly with PartiallySorted panics in diff --git a/datafusion/physical-plan/src/aggregates/topk_stream.rs b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs similarity index 76% rename from datafusion/physical-plan/src/aggregates/topk_stream.rs rename to datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs index 9128844f1d1ef..193fdba4b0198 100644 --- a/datafusion/physical-plan/src/aggregates/topk_stream.rs +++ b/datafusion/physical-plan/src/aggregates/grouped_topk_stream.rs @@ -28,7 +28,8 @@ use crate::aggregates::{ use crate::metrics::BaselineMetrics; use crate::stream::EmptyRecordBatchStream; use crate::{RecordBatchStream, SendableRecordBatchStream}; -use arrow::array::{Array, ArrayRef, RecordBatch}; +use arrow::array::{Array, ArrayRef, RecordBatch, new_null_array}; +use arrow::compute::concat; use arrow::datatypes::SchemaRef; use arrow::util::pretty::print_batches; use datafusion_common::Result; @@ -46,6 +47,7 @@ pub struct GroupedTopKAggregateStream { partition: usize, row_count: usize, started: bool, + done: bool, schema: SchemaRef, input: SendableRecordBatchStream, baseline_metrics: BaselineMetrics, @@ -53,6 +55,8 @@ pub struct GroupedTopKAggregateStream { aggregate_arguments: Vec>>, group_by: Arc, priority_map: PriorityMap, + /// Whether a NULL group key has been seen for a group-by-only aggregation. + null_group_seen: bool, } impl GroupedTopKAggregateStream { @@ -109,6 +113,7 @@ impl GroupedTopKAggregateStream { Ok(GroupedTopKAggregateStream { partition, started: false, + done: false, row_count: 0, schema: agg_schema, input, @@ -117,6 +122,7 @@ impl GroupedTopKAggregateStream { aggregate_arguments, group_by, priority_map, + null_group_seen: false, }) } } @@ -128,6 +134,10 @@ impl RecordBatchStream for GroupedTopKAggregateStream { } impl GroupedTopKAggregateStream { + fn is_group_by_only(&self) -> bool { + self.aggregate_arguments.is_empty() + } + fn intern(&mut self, ids: &ArrayRef, vals: &ArrayRef) -> Result<()> { let _timer = self.group_by_metrics.time_calculating_group_ids.timer(); @@ -136,11 +146,62 @@ impl GroupedTopKAggregateStream { .set_batch(Arc::clone(ids), Arc::clone(vals)); let has_nulls = vals.null_count() > 0; + if has_nulls && self.is_group_by_only() { + self.null_group_seen = true; + } + // Keep the common no-NULL path free of NULL bookkeeping. Once a NULL + // group exists, use the NULL-aware path until it has been resolved. + let track_null_groups = !self.is_group_by_only() + && (has_nulls || self.priority_map.has_null_groups()); for row_idx in 0..len { if has_nulls && vals.is_null(row_idx) { + // MIN/MAX ignore NULL inputs, but a group whose values are all + // NULL must still be emitted with a NULL aggregate value, so + // track it. (GROUP BY-only aggregations handle NULL group keys + // via `null_group_seen` instead.) + if !self.is_group_by_only() { + self.priority_map.insert_null(row_idx); + } continue; } - self.priority_map.insert(row_idx)?; + if track_null_groups { + self.priority_map.insert_with_null_groups(row_idx)?; + } else { + self.priority_map.insert(row_idx)?; + } + } + Ok(()) + } + + fn emit_columns(&mut self) -> Result> { + let mut cols = if self.priority_map.is_empty() { + vec![] + } else { + self.priority_map.emit()? + }; + + // GROUP BY-only aggregation covers DISTINCT-like queries. The group + // key and heap value are the same column, but the output schema has + // only the group key. + if self.is_group_by_only() { + cols.truncate(1); + if self.null_group_seen { + self.append_null_group(&mut cols)?; + } + } + + Ok(cols) + } + + fn append_null_group(&self, cols: &mut Vec) -> Result<()> { + let dt = self.schema.field(0).data_type(); + let null_arr = new_null_array(dt, 1); + if cols.is_empty() { + cols.push(null_arr); + } else { + // NULL group keys are tracked outside the heap, so append a + // one-row NULL array to the emitted non-NULL group key column. + cols[0] = concat(&[cols[0].as_ref(), null_arr.as_ref()])?; } Ok(()) } @@ -153,6 +214,9 @@ impl Stream for GroupedTopKAggregateStream { mut self: Pin<&mut Self>, cx: &mut Context<'_>, ) -> Poll> { + if self.done { + return Poll::Ready(None); + } let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); let emitting_time = self.group_by_metrics.emitting_time.clone(); while let Poll::Ready(res) = self.input.poll_next_unpin(cx) { @@ -185,8 +249,8 @@ impl Stream for GroupedTopKAggregateStream { "Exactly 1 group value required" ); let group_by_values = Arc::clone(&group_by_values[0][0]); - let input_values = if self.aggregate_arguments.is_empty() { - // DISTINCT case: use group key as both key and value + let input_values = if self.is_group_by_only() { + // GROUP BY-only case: use group key as both key and value Arc::clone(&group_by_values) } else { // MIN/MAX case: evaluate aggregate expressions @@ -209,18 +273,14 @@ impl Stream for GroupedTopKAggregateStream { // Release the input pipeline's resources before emitting. let input_schema = self.input.schema(); self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); - if self.priority_map.is_empty() { + if self.priority_map.is_empty() && !self.null_group_seen { trace!("partition {} emit None", self.partition); + self.done = true; return Poll::Ready(None); } let batch = { let _timer = emitting_time.timer(); - let mut cols = self.priority_map.emit()?; - // For DISTINCT case (no aggregate expressions), only use the group key column - // since the schema only has one field and key/value are the same - if self.aggregate_arguments.is_empty() { - cols.truncate(1); - } + let cols = self.emit_columns()?; RecordBatch::try_new(Arc::clone(&self.schema), cols)? }; let batch = batch.record_output(&self.baseline_metrics); @@ -232,6 +292,7 @@ impl Stream for GroupedTopKAggregateStream { if log::log_enabled!(Level::Trace) { print_batches(std::slice::from_ref(&batch))?; } + self.done = true; return Poll::Ready(Some(Ok(batch))); } // inner had error, return to caller diff --git a/datafusion/physical-plan/src/aggregates/hash_stream.rs b/datafusion/physical-plan/src/aggregates/hash_stream.rs new file mode 100644 index 0000000000000..f697e5a394f65 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/hash_stream.rs @@ -0,0 +1,1838 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! 2-stage hash aggregation stream implementation. +//! +//! See comments in [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`] +//! for details. +//! +//! Note these streams are an incremental migration of the existing +//! [`crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! See issue for details: + +use std::mem::size_of; +use std::ops::ControlFlow; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{ + AggregateHashTable, FinalMarker, PartialMarker, PartialSkipMarker, +}; +use super::group_values::GroupByMetrics; +use super::ordered_final_stream::OrderedFinalAggregateStream; +use super::skip_partial::SkipAggregationProbe; +use crate::metrics::{ + BaselineMetrics, MetricBuilder, MetricCategory, RecordOutput, SpillMetrics, +}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::SpillManager; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream, metrics}; + +/// Hash aggregation is implemented in two stages: partial and final. This +/// stream implements the partial stage. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=final) +/// -- RepartitionExec(hash(k)) +/// ---- AggregateExec(stage=partial) +/// +/// ## Partial Stage Behavior +/// Input: raw rows +/// Output: partial states for all groups (for example, `AVG(x)` emits `SUM(x)` +/// and `COUNT(x)`) +/// +/// ## Final Stage Behavior +/// Input: partial states +/// Output: results for all groups (for example, `AVG(x)` calculated from the +/// state) +/// +/// # Optimization: DISTINCT LIMIT Soft Limit +/// +/// This optimization applies to both [`PartialHashAggregateStream`] and +/// [`FinalHashAggregateStream`]. +/// +/// Unordered distinct queries such as: +/// +/// ```sql +/// SELECT DISTINCT x FROM t LIMIT 10; +/// ``` +/// +/// are optimized into a two-stage aggregate like: +/// +/// ```txt +/// LimitExec, limit=10 +/// --AggregateExec(Final), group_by=[x], aggr=[], soft_limit=10 +/// ---- RepartitionExec, partitioning=hash(x) +/// ------ AggregateExec(Partial), group_by=[x], aggr=[], soft_limit=10 +/// -------- Scan(t) +/// ``` +/// +/// After each input batch, the stream checks whether the soft limit has been +/// reached. If so, it emits the accumulated groups and stops reading input. +/// +/// This operator does not guarantee an exact limit because a single batch can +/// cross the threshold. The downstream limit operator enforces the exact result +/// size. +/// +/// # Optimization: Partial Aggregation Skip +/// +/// Partial aggregation can be counterproductive for high-cardinality inputs, +/// where most rows create distinct groups. The stream probes the ratio of +/// accumulated groups to input rows while it is still aggregating. If the ratio +/// crosses the configured threshold and all aggregate accumulators can convert +/// raw inputs directly to partial state, the stream emits any already +/// accumulated groups, then switches to a skip state. In that state, each +/// remaining input batch is converted directly to partial aggregate state rows +/// without inserting the rows into the grouped hash table. +/// +/// # Feature: Memory-limited Execution +/// +/// ## Partial Aggregation +/// +/// Partial aggregation can emit incomplete results because the final stage merges +/// all intermediate states for the same group. If the memory reservation exceeds +/// its limit after aggregating an input batch, this stream emits all accumulated +/// states and continues aggregating the remaining input with an empty table. +/// +/// ## Final Aggregation +/// +/// During final aggregation, group keys and states accumulate. If memory usage +/// exceeds the budget, spilling is triggered as follows: +/// 1. After aggregating a new input batch, if the memory reservation exceeds its +/// limit, spill all accumulated groups and states. +/// - Sort all groups by the group keys before spilling. +/// 2. Repeat until the input is exhausted. +/// 3. Perform a sort-preserving merge of all spill files and feed the merged output +/// into an ordered streaming aggregation, which ensures bounded memory usage and +/// evaluates the final result. +/// - [`OrderedFinalAggregateStream`] is reused for the streaming aggregation. +pub(crate) struct PartialHashAggregateStream { + /// Output schema: group columns followed by partial aggregate state columns. + schema: SchemaRef, + + /// Input batches containing raw rows, not partial aggregate state. + input: SendableRecordBatchStream, + + /// Target output batch size from configuration. + batch_size: usize, + + /// Memory reservation for group keys and accumulators. + reservation: MemoryReservation, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// Tracks partial aggregation row reduction, matching `GroupedHashAggregateStream`. + reduction_factor: metrics::RatioMetrics, + + /// Tracks whether partial aggregation should switch to direct state conversion. + skip_aggregation_probe: Option, + + /// Optional soft limit on the number of groups to accumulate before output. + /// + /// Invariant: when this is `Some(..)`, the accumulators inside `hash_table` must + /// be empty. See struct comments for details. + group_values_soft_limit: Option, + + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// state for emitting output batches. + state: Option, +} + +/// States for partial hash aggregation processing. +enum PartialHashAggregateState { + ReadingInput { + hash_table: AggregateHashTable, + }, + /// A fully materialized partial-state batch being emitted incrementally. + EmittingOnMemoryPressure { + hash_table: AggregateHashTable, + // After each incremental emitting step, the `remaining_groups` will be updated + // with batch slicing. + remaining_groups: RecordBatch, + }, + ProducingOutput { + hash_table: AggregateHashTable, + /// If `None`, partial skip was never triggered and this state will + /// finish in `Done`. If `Some`, partial skip has triggered and the + /// stream will move to `SkippingAggregation` after these accumulated + /// groups are emitted. + skip_hash_table: Option>, + }, + SkippingAggregation { + hash_table: AggregateHashTable, + }, + Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, +} + +type PartialHashAggregatePoll = Poll>>; +type PartialHashAggregateStateTransition = ControlFlow< + (PartialHashAggregatePoll, PartialHashAggregateState), + PartialHashAggregateState, +>; + +/// Spill configuration and accumulated runs for final hash aggregation. +/// +/// Each spill event drains all currently buffered groups, sorts their intermediate +/// states by the full group key, and writes them to one spill file. All files are +/// merged and replayed after the original input ends. +struct FinalSpillContext { + /// Aggregate configuration used to construct the final replay stream. + final_agg: AggregateExec, + /// Task context. + context: Arc, + /// Original partition index. + partition: usize, + /// Target batch size from configuration. + batch_size: usize, + /// Full group-key ordering kept by every spill file and the merged input. + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Spill runs waiting to be merged, they're all sorted by full group-by keys. + spills: Vec, +} + +/// Hash aggregation is implemented in two stages: partial and final. This +/// stream implements the final stage. +/// +/// See [`PartialHashAggregateStream`] for details. +pub(crate) struct FinalHashAggregateStream { + /// Output schema: group columns followed by final aggregate value columns. + schema: SchemaRef, + + /// Input batches containing partial aggregate state rows. + input: SendableRecordBatchStream, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// Memory reservation for group keys, accumulators, and spill sorting. + reservation: MemoryReservation, + + /// See comments for the same variable in [`PartialHashAggregateStream`]. + group_values_soft_limit: Option, + + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// state for emitting output batches. + state: Option, +} + +/// States for final hash aggregation processing. +// The typestate pattern is used in case the inner logic becomes more complex in +// the future. +enum FinalHashAggregateState { + ReadingInput { + hash_table: AggregateHashTable, + /// `None` if spilling is not supported by the configured `DiskManager`. + spill_context: Option>, + }, + Spilling { + hash_table: AggregateHashTable, + spill_context: Box, + }, + ProducingOutput { + hash_table: AggregateHashTable, + }, + PreparingMergeInput { + hash_table: AggregateHashTable, + spill_context: Box, + }, + MergingSpills { + stream: SendableRecordBatchStream, + }, + Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, +} + +type FinalHashAggregatePoll = Poll>>; +type FinalHashAggregateStateTransition = ControlFlow< + (FinalHashAggregatePoll, FinalHashAggregateState), + FinalHashAggregateState, +>; + +impl FinalSpillContext { + fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + batch_size: usize, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result { + let group_schema = agg.group_by.group_schema(&agg.input().schema())?; + let output_ordering = agg.cache.output_ordering(); + let spill_sort_exprs = + group_schema + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + let output_expr = Column::new(field.name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Final hash aggregate spill expression is empty"); + }; + + let spill_manager = SpillManager::new( + context.runtime_env(), + spill_metrics, + Arc::clone(spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + let mut final_agg = agg.clone(); + final_agg.input_order_mode = InputOrderMode::Sorted; + + Ok(Self { + final_agg, + context: Arc::clone(context), + partition, + batch_size, + spill_expr, + spill_manager, + spills: vec![], + }) + } + + fn has_spills(&self) -> bool { + !self.spills.is_empty() + } + + /// Sorts and spills the aggregated groups. Memory reservation should be updated + /// by the caller. + /// + /// Individual spill files are ordered by the `group by` keys. + /// + /// See [`FinalHashAggregateStream`] for spilling details. + fn spill_table( + &mut self, + hash_table: &mut AggregateHashTable, + ) -> Result<()> { + let Some(batch) = hash_table.take_state_batch()? else { + return Ok(()); + }; + + let sorted_iter = + IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "FinalHashAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Final hash aggregation produced an empty spill"); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// Merges every sorted run, and do the aggregate evaluation with + /// [`OrderedFinalAggregateStream`] + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result { + let Self { + final_agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &final_agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) + } +} + +impl PartialHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert_eq!(agg.mode, super::AggregateMode::Partial); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let reduction_factor = MetricBuilder::new(&agg.metrics) + .with_type(metrics::MetricType::Summary) + .ratio_metrics("reduction_factor", partition); + + let hash_table = AggregateHashTable::::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + let skip_aggregation_probe = if agg.group_by.is_single() { + let options = &context.session_config().options().execution; + let probe_ratio_threshold = + options.skip_partial_aggregation_probe_ratio_threshold; + // A threshold >= 1.0 means the ratio (num_groups / input_rows) can + // never exceed it, so the feature is effectively disabled. + if probe_ratio_threshold >= 1.0 { + None + } else { + let skipped_aggregation_rows = MetricBuilder::new(&agg.metrics) + .with_category(MetricCategory::Rows) + .counter("skipped_aggregation_rows", partition); + Some(SkipAggregationProbe::new( + options.skip_partial_aggregation_probe_rows_threshold, + probe_ratio_threshold, + skipped_aggregation_rows, + )) + } + } else { + None + }; + + let reservation = + MemoryConsumer::new(format!("PartialHashAggregateStream[{partition}]")) + .with_can_spill(true) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + batch_size, + baseline_metrics, + reservation, + reduction_factor, + skip_aggregation_probe, + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), + state: Some(PartialHashAggregateState::ReadingInput { hash_table }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> PartialHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + PartialHashAggregateState::Error, + )) + } + + fn break_with_internal_err( + message: impl std::fmt::Display, + ) -> PartialHashAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + + /// See comments in [`Self::group_values_soft_limit`] for details. + fn hit_soft_group_limit( + &self, + hash_table: &AggregateHashTable, + ) -> bool { + self.group_values_soft_limit + .is_some_and(|limit| limit <= hash_table.building_group_count()) + } + + /// Updates skip aggregation probe state. + fn update_skip_aggregation_probe(&mut self, input_rows: usize, num_groups: usize) { + if let Some(probe) = self.skip_aggregation_probe.as_mut() { + probe.update_state(input_rows, num_groups); + } + } + + /// Returns true if the aggregation probe indicates that aggregation + /// should be skipped. + fn should_skip_aggregation(&self) -> bool { + self.skip_aggregation_probe + .as_ref() + .is_some_and(|probe| probe.should_skip()) + } + + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable, + close_input: bool, + ) -> Result<()> { + if close_input { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + hash_table.start_output() + } + + /// Handle ReadingInput state - aggregate input batches into the hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + original_state: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::ReadingInput { mut hash_table } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected ReadingInput state", + ); + }; + debug_assert!(hash_table.is_building()); + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + PartialHashAggregateState::ReadingInput { hash_table }, + )), + Poll::Ready(Some(Ok(batch))) => { + // ---------------------------------- + // Step 1: Aggregate the input batch + // ---------------------------------- + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let input_rows = batch.num_rows(); + self.reduction_factor.add_total(input_rows); + let result = hash_table.aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + // -------------------------------- + // Step 2: Soft limit optimization + // -------------------------------- + if self.hit_soft_group_limit(&hash_table) { + let timer = elapsed_compute.timer(); + let result = self.start_output(&mut hash_table, true); + timer.done(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + return ControlFlow::Continue( + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table: None, + }, + ); + } + + // ---------------------------------------------- + // Step 3: Skip partial aggregation optimization + // ---------------------------------------------- + self.update_skip_aggregation_probe( + input_rows, + hash_table.building_group_count(), + ); + + // True branch: a decision has been made to skip partial aggregation. + if self.should_skip_aggregation() { + let timer = elapsed_compute.timer(); + let result = match hash_table.partial_skip_table() { + Ok(skip_hash_table) => self + .start_output(&mut hash_table, false) + .map(|()| skip_hash_table), + Err(e) => Err(e), + }; + timer.done(); + + match result { + Ok(skip_hash_table) => { + // Move to `ProducingOutput` first. Its `skip_hash_table` + // field moves the stream to skip-partial aggregation after + // the accumulated batches have been output. + return ControlFlow::Continue( + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table: Some(skip_hash_table), + }, + ); + } + Err(e) => return Self::break_with_err(e), + } + } + + // ------------------------------------------------- + // Step 4: Larger-than-memory execution (early emit) + // ------------------------------------------------- + let timer = elapsed_compute.timer(); + let resize_result = self.reservation.try_resize(hash_table.memory_size()); + timer.done(); + match resize_result { + Ok(()) => {} + Err(DataFusionError::ResourcesExhausted(_)) => { + let elapsed_compute = + self.baseline_metrics.elapsed_compute().clone(); + // Stops on drop + let _timer = elapsed_compute.timer(); + let state_batch_result = hash_table.take_state_batch(); + + // Emitting clears the aggregate table and releases its + // accumulated memory. Update the reservation accordingly. + let resize_result = + self.reservation.try_resize(hash_table.memory_size()); + + if let Err(e) = resize_result { + return Self::break_with_err(e); + } + + let materialized_group_states = match state_batch_result { + Ok(Some(batch)) => batch, + Ok(None) => { + return Self::break_with_err(internal_datafusion_err!( + "Partial hash aggregate ran out of memory with no aggregated groups" + )); + } + Err(e) => return Self::break_with_err(e), + }; + + return ControlFlow::Continue( + PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: materialized_group_states, + }, + ); + } + Err(e) => return Self::break_with_err(e), + } + + ControlFlow::Continue(PartialHashAggregateState::ReadingInput { + hash_table, + }) + } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = self.start_output(&mut hash_table, true); + timer.done(); + + match result { + Ok(()) => ControlFlow::Continue( + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table: None, + }, + ), + Err(e) => Self::break_with_err(e), + } + } + } + } + + /// Handle EmittingOnMemoryPressure state - emit a materialized partial-state + /// batch in `batch_size`(from configuration) slices, then resume reading input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_emitting_on_memory_pressure( + &mut self, + original_state: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: batch, + } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected EmittingOnMemoryPressure state", + ); + }; + + let (output_batch, next_state) = if batch.num_rows() <= self.batch_size { + // Last batch to output, go back to `ReadingInput` + ( + batch, + PartialHashAggregateState::ReadingInput { hash_table }, + ) + } else { + // More batch to output, continue in the current state. + let remaining = + batch.slice(self.batch_size, batch.num_rows() - self.batch_size); + let output = batch.slice(0, self.batch_size); + ( + output, + PartialHashAggregateState::EmittingOnMemoryPressure { + hash_table, + remaining_groups: remaining, + }, + ) + }; + + self.reduction_factor.add_part(output_batch.num_rows()); + debug_assert!(output_batch.num_rows() > 0); + ControlFlow::Break(( + Poll::Ready(Some(Ok(output_batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + + /// Handle ProducingOutput state - emit partial aggregate state batches. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + original_state: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::ProducingOutput { + mut hash_table, + skip_hash_table, + } = original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected ProducingOutput state", + ); + }; + debug_assert!(!hash_table.is_building()); + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.next_output_batch(); + timer.done(); + + match result { + Ok(Some(batch)) => { + let _ = self.reservation.try_resize(hash_table.memory_size()); + self.reduction_factor.add_part(batch.num_rows()); + debug_assert!(batch.num_rows() > 0); + let next_state = if hash_table.is_done() { + match skip_hash_table { + Some(hash_table) => { + PartialHashAggregateState::SkippingAggregation { hash_table } + } + None => PartialHashAggregateState::Done, + } + } else { + PartialHashAggregateState::ProducingOutput { + hash_table, + skip_hash_table, + } + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Ok(None) => { + let _ = self.reservation.try_resize(0); + // If the previous `Aggregating` stage decided to skip partial + // aggregation, go to the `SkippingAggregation` stage; otherwise finish. + let next_state = match skip_hash_table { + Some(hash_table) => { + PartialHashAggregateState::SkippingAggregation { hash_table } + } + None => PartialHashAggregateState::Done, + }; + ControlFlow::Continue(next_state) + } + Err(e) => Self::break_with_err(e), + } + } + + /// Handle SkippingAggregation state - convert raw input directly to partial states. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_skipping_aggregation( + &mut self, + cx: &mut Context<'_>, + original_state: PartialHashAggregateState, + ) -> PartialHashAggregateStateTransition { + let PartialHashAggregateState::SkippingAggregation { mut hash_table } = + original_state + else { + return Self::break_with_internal_err( + "Partial hash aggregate stream expected SkippingAggregation state", + ); + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + PartialHashAggregateState::SkippingAggregation { hash_table }, + )), + Poll::Ready(Some(Ok(batch))) => { + if let Some(probe) = self.skip_aggregation_probe.as_mut() { + probe.record_skipped(&batch); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.convert_batch_to_state(&batch); + timer.done(); + + match result { + Ok(batch) => ControlFlow::Break(( + Poll::Ready(Some( + Ok(batch.record_output(&self.baseline_metrics)), + )), + PartialHashAggregateState::SkippingAggregation { hash_table }, + )), + Err(e) => Self::break_with_err(e), + } + } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + ControlFlow::Continue(PartialHashAggregateState::Done) + } + } + } +} + +impl Stream for PartialHashAggregateStream { + type Item = Result; + + /// Entry point for the partial hash aggregate state machine. + /// + /// See comments in [`PartialHashAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling input and aggregating batches into the + /// in-memory hash table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one batch, update the inner aggregate hash table, and + /// continue with the next input batch. + /// -> EmittingOnMemoryPressure + /// The table cannot reserve enough memory. Materialize all accumulated + /// partial states and begin emitting them incrementally. + /// -> ProducingOutput(skip=None) + /// Input was exhausted, or the soft group limit was reached. Move to + /// the next state to start outputting. + /// -> ProducingOutput(skip=Some) + /// Partial skip aggregation was triggered. First move to the + /// `ProducingOutput` state to drain the accumulated state, then move to + /// the `SkippingAggregation` state to convert input directly to partial + /// state without aggregation. + /// + /// EmittingOnMemoryPressure + /// -> EmittingOnMemoryPressure + /// One batch-sized slice was yielded; repeat until all materialized + /// partial states are emitted. + /// -> ReadingInput + /// The materialized states were emitted; continue with the empty table. + /// + /// ProducingOutput(skip=None) + /// -> ProducingOutput(skip=None) + /// One accumulated output batch was yielded, repeat to continue producing + /// output incrementally. + /// -> Done + /// All accumulated output was emitted. + /// + /// ProducingOutput(skip=Some) + /// -> ProducingOutput(skip=Some) + /// One accumulated output batch was yielded, repeat to continue producing + /// output incrementally. + /// -> SkippingAggregation + /// All accumulated output was emitted. Continue by converting raw + /// input batches directly to partial aggregate state. + /// + /// SkippingAggregation + /// -> SkippingAggregation + /// One `convert_to_state` batch was yielded; repeat to continue + /// processing. + /// -> Done + /// Input was exhausted. + /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + let cur_state = self + .state + .take() + .expect("PartialHashAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ PartialHashAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ PartialHashAggregateState::EmittingOnMemoryPressure { .. } => { + self.handle_emitting_on_memory_pressure(state) + } + state @ PartialHashAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ PartialHashAggregateState::SkippingAggregation { .. } => { + self.handle_skipping_aggregation(cx, state) + } + state @ PartialHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } + state @ PartialHashAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!(next_state, PartialHashAggregateState::Error)); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(PartialHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for PartialHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +impl FinalHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert!(matches!( + agg.mode, + super::AggregateMode::Final | super::AggregateMode::FinalPartitioned + )); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let input_schema = input.schema(); + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let spill_metrics = SpillMetrics::new(&agg.metrics, partition); + + let hash_table = AggregateHashTable::::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + + let can_spill = context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + Some(Box::new(FinalSpillContext::new( + agg, + context, + partition, + batch_size, + &input_schema, + spill_metrics, + )?)) + } else { + None + }; + + let reservation = + MemoryConsumer::new(format!("FinalHashAggregateStream[{partition}]")) + .with_can_spill(can_spill) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + group_values_soft_limit: agg.limit_options().map(|config| config.limit()), + state: Some(FinalHashAggregateState::ReadingInput { + hash_table, + spill_context, + }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> FinalHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + FinalHashAggregateState::Error, + )) + } + + fn break_with_internal_err( + message: impl std::fmt::Display, + ) -> FinalHashAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + + /// See comments in [`Self::group_values_soft_limit`] for details. + fn hit_soft_group_limit(&self, hash_table: &AggregateHashTable) -> bool { + self.group_values_soft_limit + .is_some_and(|limit| limit <= hash_table.building_group_count()) + } + + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable, + ) -> Result<()> { + self.close_input(); + hash_table.start_output() + } + + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + hash_table: &AggregateHashTable, + spill_context: Option<&FinalSpillContext>, + ) -> usize { + let table_size = hash_table.memory_size(); + if spill_context.is_some() { + // Count extra space needed for in-memory sorting and spilling. Only + // count memory for indices, the payload will be materialize incrementally + // in smaller chunks. + table_size.saturating_add( + hash_table + .building_group_count() + .saturating_mul(size_of::()), + ) + } else { + table_size + } + } + + /// Handle ReadingInput state - aggregate partial state batches into the hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + original_state: FinalHashAggregateState, + ) -> FinalHashAggregateStateTransition { + let FinalHashAggregateState::ReadingInput { + mut hash_table, + spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected ReadingInput state", + ); + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + FinalHashAggregateState::ReadingInput { + hash_table, + spill_context, + }, + )), + Poll::Ready(Some(Ok(batch))) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + // Soft group limits are usually small and rarely coincide with + // spilling. Once spilling has occurred, skip this optimization to + // make the internal logic simpler. + let spilled = spill_context + .as_ref() + .is_some_and(|context| context.has_spills()); + if self.hit_soft_group_limit(&hash_table) && !spilled { + let timer = elapsed_compute.timer(); + let result = self.start_output(&mut hash_table); + timer.done(); + + return match result { + Ok(()) => ControlFlow::Continue( + FinalHashAggregateState::ProducingOutput { hash_table }, + ), + Err(e) => Self::break_with_err(e), + }; + } + + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + self.reservation + .try_resize(Self::reservation_size_for_table( + &hash_table, + spill_context.as_deref(), + )); + timer.done(); + match resize_result { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + // OOM and don't support spilling from configuration + let Some(spill_context) = spill_context else { + return Self::break_with_err(e.context( + "Final hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + )); + }; + // Sanity check: impossible to OOM when there is no group aggregated. + if hash_table.building_group_count() == 0 { + return Self::break_with_internal_err( + "Final hash aggregate ran out of memory with no aggregated groups", + ); + } + // Go to the next state to perform spilling the aggregated + // groups so far. + return ControlFlow::Continue( + FinalHashAggregateState::Spilling { + hash_table, + spill_context, + }, + ); + } + Err(e) => return Self::break_with_err(e), + } + + ControlFlow::Continue(FinalHashAggregateState::ReadingInput { + hash_table, + spill_context, + }) + } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + // Input done, move to next state: + // - If spilled before, perform merging spill runs + // - If not spilled, start producing outputs + Poll::Ready(None) => { + self.close_input(); + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue( + FinalHashAggregateState::PreparingMergeInput { + hash_table, + spill_context, + }, + ) + } + _ => { + let elapsed_compute = + self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.start_output(); + timer.done(); + + match result { + Ok(()) => ControlFlow::Continue( + FinalHashAggregateState::ProducingOutput { hash_table }, + ), + Err(e) => Self::break_with_err(e), + } + } + } + } + } + } + + /// Sorts and spills one complete in-memory state run, then resumes input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_spilling( + &mut self, + original_state: FinalHashAggregateState, + ) -> FinalHashAggregateStateTransition { + let FinalHashAggregateState::Spilling { + mut hash_table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected Spilling state", + ); + }; + + // Sanity check: it is impossible to OOM when the table is empty. + if hash_table.building_group_count() == 0 { + return Self::break_with_internal_err( + "Final hash aggregation entered Spilling with an empty table", + ); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let mut result = spill_context.spill_table(&mut hash_table); + + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. + if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) { + result = + Err(e.context("Decreasing allocation after spilling should succeed")); + } + + timer.done(); + + match result { + // Finished spilling the aggregate table, continue aggregating from input. + Ok(()) => ControlFlow::Continue(FinalHashAggregateState::ReadingInput { + hash_table, + spill_context: Some(spill_context), + }), + Err(e) => Self::break_with_err(e), + } + } + + /// 1. Spills the last in-memory run. + /// 2. Constructs a globally ordered input stream by applying a sort-preserving + /// merge to all spills. + /// 3. Constructs a replay stream: an ordered final aggregate stream over the + /// fully ordered input constructed from the spills. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_preparing_merge_input( + &mut self, + original_state: FinalHashAggregateState, + ) -> FinalHashAggregateStateTransition { + let FinalHashAggregateState::PreparingMergeInput { + mut hash_table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected PreparingMergeInput state", + ); + }; + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let replay = match spill_context.spill_table(&mut hash_table) { + Ok(()) => { + let group_by_metrics = hash_table.group_by_metrics().clone(); + drop(hash_table); + match self.reservation.try_resize(0) { + Ok(()) => (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_metrics, + self.reservation.new_empty(), + ), + Err(e) => Err(e), + } + } + Err(e) => Err(e), + }; + timer.done(); + + match replay { + Ok(stream) => { + ControlFlow::Continue(FinalHashAggregateState::MergingSpills { stream }) + } + Err(e) => Self::break_with_err(e), + } + } + + /// Forwards output from the fully ordered stream that consumes the merged + /// spill runs. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_merging_spills( + &mut self, + cx: &mut Context<'_>, + original_state: FinalHashAggregateState, + ) -> FinalHashAggregateStateTransition { + let FinalHashAggregateState::MergingSpills { mut stream } = original_state else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected MergingSpills state", + ); + }; + + match stream.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + FinalHashAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( + Poll::Ready(Some(Ok(batch))), + FinalHashAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => ControlFlow::Continue(FinalHashAggregateState::Done), + } + } + + /// Handle ProducingOutput state - emit final aggregate value batches. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + original_state: FinalHashAggregateState, + ) -> FinalHashAggregateStateTransition { + let FinalHashAggregateState::ProducingOutput { mut hash_table } = original_state + else { + return Self::break_with_internal_err( + "Final hash aggregate stream expected ProducingOutput state", + ); + }; + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.next_output_batch(); + timer.done(); + + match result { + Ok(Some(batch)) => { + let next_state = if hash_table.is_done() { + drop(hash_table); + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + FinalHashAggregateState::Done + } else { + if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) + { + return Self::break_with_err(e); + } + FinalHashAggregateState::ProducingOutput { hash_table } + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Err(e) => Self::break_with_err(e), + Ok(None) => { + drop(hash_table); + let next_state = FinalHashAggregateState::Done; + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + ControlFlow::Continue(next_state) + } + } + } +} + +impl Stream for FinalHashAggregateStream { + type Item = Result; + + /// Entry point for the final hash aggregate state machine. + /// + /// See comments in [`FinalHashAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling partial-state input and aggregating + /// those states into the final hash table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one partial-state input batch. If it fits in memory, + /// continue with the next input batch. + /// -> Spilling + /// The table cannot reserve enough memory. Move all current states into + /// one fully group-key-sorted spill run. + /// -> ProducingOutput + /// Input was exhausted without spilling, or the soft group limit was + /// reached. Start outputting final aggregate values. + /// -> PreparingMergeInput + /// Input was exhausted after spilling. Spill the last in-memory run and + /// construct the ordered input used to merge all spill files. + /// + /// Spilling + /// -> ReadingInput + /// One sorted run was written; resume reading the original input. + /// + /// PreparingMergeInput + /// Spill the final in-memory run and build the input ordered replay stream. + /// -> MergingSpills + /// The final run was spilled and the ordered replay stream was built. + /// + /// MergingSpills + /// Aggregate the merged spill runs and emit final results. + /// -> MergingSpills + /// Forward one result batch from the fully ordered replay stream that + /// consumes the sort-preserving merge. + /// -> Done + /// The merged spill input was fully aggregated. + /// + /// ProducingOutput + /// -> ProducingOutput + /// One final output batch was yielded; repeat to continue producing + /// output incrementally. + /// -> Done + /// All final output was emitted. + /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + let cur_state = self + .state + .take() + .expect("FinalHashAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ FinalHashAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ FinalHashAggregateState::Spilling { .. } => { + self.handle_spilling(state) + } + state @ FinalHashAggregateState::PreparingMergeInput { .. } => { + self.handle_preparing_merge_input(state) + } + state @ FinalHashAggregateState::MergingSpills { .. } => { + self.handle_merging_spills(cx, state) + } + state @ FinalHashAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ FinalHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } + state @ FinalHashAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!(next_state, FinalHashAggregateState::Error)); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(FinalHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for FinalHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use super::*; + use crate::aggregates::{AggregateMode, PhysicalGroupBy}; + use crate::execution_plan::ExecutionPlan; + use crate::test::TestMemoryExec; + + use arrow::array::{Int32Array, Int64Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::Result; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use futures::StreamExt; + + #[tokio::test] + async fn test_partial_hash_stream_double_emission_race_condition_bug() -> Result<()> { + // Fix for https://github.com/apache/datafusion/issues/18701 + // This test specifically proves that we have fixed double emission race condition + // where emit_early_if_necessary() and switch_to_skip_aggregation() + // both emit in the same loop iteration, causing data loss + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + // Create data that will trigger BOTH conditions in the same iteration: + // 1. More groups than batch_size (triggers early emission when memory pressure hits) + // 2. High cardinality ratio (triggers skip aggregation) + let batch_size = 1024; // We'll set this in session config + let num_groups = batch_size + 100; // Slightly more than batch_size (1124 groups) + + // Create exactly 1 row per group = 100% cardinality ratio + let group_ids: Vec = (0..num_groups as i32).collect(); + let values: Vec = vec![1; num_groups]; + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids)), + Arc::new(Int64Array::from(values)), + ], + )?; + let input_partitions = vec![vec![batch]]; + + // Create constrained memory to trigger early emission but not completely fail + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(1024, 1.0) // small enough to start but will trigger pressure + .build_arc()?; + + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure to trigger BOTH conditions: + // 1. Low probe threshold (triggers skip probe after few rows) + // 2. Low ratio threshold (triggers skip aggregation immediately) + // 3. Set batch_size to 1024 so our 1124 groups will trigger early emission + // This creates the race condition where both emit paths are triggered + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.batch_size", + &datafusion_common::ScalarValue::UInt64(Some(1024)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(50)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(0.8)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode where the race condition occurs + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Count total groups emitted + let mut total_output_groups = 0; + for batch in &results { + total_output_groups += batch.num_rows(); + } + + assert_eq!( + total_output_groups, num_groups, + "Unexpected number of groups", + ); + + Ok(()) + } + + #[tokio::test] + async fn test_partial_hash_stream_skip_aggregation_probe_not_locked_until_skip() + -> Result<()> { + // Test that the probe is not locked until we actually decide to skip. + // This allows us to continue evaluating the skip condition across multiple batches. + // + // Scenario: + // - Batch 1: Hits rows threshold but NOT ratio threshold (low cardinality) -> don't skip + // - Batch 2: Now hits ratio threshold (high cardinality) -> skip + // + // Without the fix, the probe would be locked after batch 1, preventing the skip + // decision from being made on batch 2. + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int32, false), + ])); + + // Configure thresholds: + // - probe_rows_threshold: 100 rows + // - probe_ratio_threshold: 0.8 (80%) + let probe_rows_threshold = 100; + let probe_ratio_threshold = 0.8; + + // Batch 1: 100 rows with only 10 unique groups + // Ratio: 10/100 = 0.1 (10%) < 0.8 -> should NOT skip + // This will hit the rows threshold but not the ratio threshold + let batch1_rows = 100; + let batch1_groups = 10; + let mut group_ids_batch1 = Vec::new(); + for i in 0..batch1_rows { + group_ids_batch1.push((i % batch1_groups) as i32); + } + let values_batch1: Vec = vec![1; batch1_rows]; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch1)), + Arc::new(Int32Array::from(values_batch1)), + ], + )?; + + // Batch 2: 360 rows with 360 unique NEW groups (starting from group 10) + // After batch 2, total: 460 rows, 370 groups + // Ratio: 370/460 is about 0.804 (80.4%) > 0.8 -> SHOULD decide to skip + let batch2_rows = 360; + let batch2_groups = 360; + let group_ids_batch2: Vec = (batch1_groups..(batch1_groups + batch2_groups)) + .map(|x| x as i32) + .collect(); + let values_batch2: Vec = vec![1; batch2_rows]; + + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch2)), + Arc::new(Int32Array::from(values_batch2)), + ], + )?; + + // Batch 3: This batch should be skipped since we decided to skip after batch 2 + // 100 rows with 100 unique groups (continuing from where batch 2 left off) + let batch3_rows = 100; + let batch3_groups = 100; + let batch3_start_group = batch1_groups + batch2_groups; + let group_ids_batch3: Vec = (batch3_start_group + ..(batch3_start_group + batch3_groups)) + .map(|x| x as i32) + .collect(); + let values_batch3: Vec = vec![1; batch3_rows]; + + let batch3 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch3)), + Arc::new(Int32Array::from(values_batch3)), + ], + )?; + + let input_partitions = vec![vec![batch1, batch2, batch3]]; + + let runtime = RuntimeEnvBuilder::default().build_arc()?; + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure skip aggregation settings + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(probe_rows_threshold)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(probe_ratio_threshold)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + PartialHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Check that skip aggregation actually happened. + // The key metric is skipped_aggregation_rows. + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + + // We expect batch 3's rows to be skipped (100 rows) + assert_eq!( + skipped_rows, batch3_rows, + "Expected batch 3's rows ({batch3_rows}) to be skipped", + ); + + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/mod.rs b/datafusion/physical-plan/src/aggregates/mod.rs index d1498e4a3ea55..a39c6f34862d0 100644 --- a/datafusion/physical-plan/src/aggregates/mod.rs +++ b/datafusion/physical-plan/src/aggregates/mod.rs @@ -15,28 +15,162 @@ // specific language governing permissions and limitations // under the License. -//! Aggregates functionalities +//! Aggregate functionality +//! +//! # Aggregate planning +//! +//! DataFusion selects different aggregate implementations (streams) based on the +//! query shape and configuration. This section provides an overview of the +//! available stream variants. +//! +//! See each stream's documentation for details. +//! +//! ## 1. Two-stage hash aggregation +//! +//! Two-stage hash aggregation is used for regular parallel execution. +//! +//! The input passes through three execution operators to produce the final +//! aggregation result: +//! +//! 1. Partial aggregation reads the input and produces partial states. It +//! aggregates independently within each partition, which usually reduces +//! cardinality before the later shuffle. +//! 2. Hash repartitioning on the group keys sends all partial states for each +//! group to the same output partition for final aggregation. +//! 3. Final aggregation reads the partial states, combines them, and emits the +//! final results. +//! +//! ```text +//! AggregateExec (final) +//! RepartitionExec (hash by group keys) +//! AggregateExec (partial) +//! ``` +//! +//! See [`PartialHashAggregateStream`] and [`FinalHashAggregateStream`] for details. +//! +//! ### Ordering optimization +//! +//! When the input is ordered by the group key, an ordered fast path is used. It +//! uses a similar two-stage hash aggregation with an early-emission optimization. +//! +//! ```text +//! AggregateExec (final, ordered) +//! RepartitionExec (hash by group keys, order-preserving) +//! AggregateExec (partial, ordered) +//! ``` +//! +//! See [`OrderedPartialAggregateStream`] and [`OrderedFinalAggregateStream`] for +//! details. +//! +//! Related configuration: +//! +//! - [`datafusion.execution.target_partitions`](datafusion_common::config::ExecutionOptions::target_partitions) +//! - [`datafusion.optimizer.repartition_aggregations`](datafusion_common::config::OptimizerOptions::repartition_aggregations) +//! - [`datafusion.optimizer.prefer_existing_sort`](datafusion_common::config::OptimizerOptions::prefer_existing_sort) +//! +//! ## 2. Single-stage hash aggregation +//! +//! When there is a single partition, or the aggregation input is already +//! key-partitioned (e.g., a data source has existing range partitioning), +//! `Single` mode aggregation is used. +//! +//! It takes raw input and directly produces the final result. +//! +//! ```text +//! AggregateExec (mode=Single or SinglePartitioned) +//! input +//! ``` +//! +//! See [`SingleHashAggregateStream`] for details. +//! +//! Related configuration: +//! +//! - [`datafusion.execution.target_partitions`](datafusion_common::config::ExecutionOptions::target_partitions) +//! - [`datafusion.optimizer.repartition_aggregations`](datafusion_common::config::OptimizerOptions::repartition_aggregations) +//! +//! ## 3. Aggregation without grouping expressions +//! +//! A global aggregate maintains one accumulator set per input partition rather +//! than a hash table of groups. Partial stages compute local states and a final +//! stage combines them into one output row: +//! +//! ```text +//! AggregateExec (final, no-grouping) +//! CoalescePartitionsExec +//! AggregateExec (partial, no-grouping) +//! ``` +//! +//! Every stage without grouping expressions uses [`AggregateStream`]. This path +//! is selected before the grouped-stream migration setting is considered. +//! +//! ## 4. Grouped TopK aggregation +//! +//! When a query only needs the best `N` groups, retaining every group in a hash +//! table and sorting them afterward does unnecessary work. The optimizer pushes +//! the sort limit and direction into the aggregate: +//! +//! ```text +//! SortExec (fetch=N) +//! AggregateExec (limit=N, order=...) +//! input +//! ``` +//! +//! [`GroupedTopKAggregateStream`] keeps a bounded priority map for a single group +//! key. It supports group-by-only queries and compatible `MIN` or `MAX` +//! aggregates. An unordered group-by-only soft limit instead stays on the normal +//! hash aggregation path. +//! +//! Related configuration: +//! +//! - [`datafusion.optimizer.enable_topk_aggregation`](datafusion_common::config::OptimizerOptions::enable_topk_aggregation) +//! - [`datafusion.optimizer.enable_distinct_aggregation_soft_limit`](datafusion_common::config::OptimizerOptions::enable_distinct_aggregation_soft_limit) +//! +//! ## 5. Partial-reduce hash aggregation +//! +//! This implementation will not be planned by DataFusion SQL interface, it must be +//! manually constructed at [`ExecutionPlan`] level. +//! +//! This mode is useful in a distributed setting. +//! +//! See [`PartialReduceHashAggregateStream`] for details. +//! +//! ## 6. Fallback grouped hash aggregation +//! +//! [`GroupedHashAggregateStream`] is the legacy implementation for several of the +//! stream types above. It is being incrementally migrated to separate streams. +//! +//! See the issue for details: +#![expect(rustdoc::private_intra_doc_links)] use std::borrow::Cow; use std::sync::Arc; use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; use crate::aggregates::{ - no_grouping::AggregateStream, row_hash::GroupedHashAggregateStream, - topk_stream::GroupedTopKAggregateStream, + aggregate_stream::AggregateStream, + grouped_hash_stream::GroupedHashAggregateStream, + grouped_topk_stream::GroupedTopKAggregateStream, + hash_stream::{FinalHashAggregateStream, PartialHashAggregateStream}, + ordered_final_stream::OrderedFinalAggregateStream, + ordered_partial_stream::OrderedPartialAggregateStream, + partial_reduce_stream::PartialReduceHashAggregateStream, + single_stream::SingleHashAggregateStream, +}; +use crate::execution_plan::{ + CardinalityEffect, EmissionType, plan_contains_expression_id, }; -use crate::execution_plan::{CardinalityEffect, EmissionType}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, - FilterPushdownPropagation, PushedDownPredicate, + FilterPushdownPropagation, }; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, InputOrderMode, - SendableRecordBatchStream, Statistics, check_if_same_properties, + DisplayFormatType, Distribution, ExecutionPlan, InputDistributionRequirements, + InputOrderMode, SendableRecordBatchStream, Statistics, }; use datafusion_common::config::ConfigOptions; -use datafusion_physical_expr::utils::collect_columns; use parking_lot::Mutex; use std::collections::{HashMap, HashSet}; @@ -47,10 +181,11 @@ use arrow_schema::FieldRef; use datafusion_common::stats::Precision; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ - Constraint, Constraints, Result, ScalarValue, assert_eq_or_internal_err, - internal_err, not_impl_err, + ColumnStatistics, Constraint, Constraints, Result, ScalarValue, + assert_eq_or_internal_err, internal_err, not_impl_err, }; use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::MemoryLimit; use datafusion_expr::{Accumulator, Aggregate}; use datafusion_physical_expr::aggregate::AggregateFunctionExpr; use datafusion_physical_expr::equivalence::ProjectionMapping; @@ -69,12 +204,19 @@ use itertools::Itertools; use topk::hash_table::is_supported_hash_key_type; use topk::heap::is_supported_heap_type; +mod aggregate_hash_table; +mod aggregate_stream; pub mod group_values; -mod no_grouping; +mod grouped_hash_stream; +mod grouped_topk_stream; +mod hash_stream; pub mod order; -mod row_hash; +mod ordered_final_stream; +mod ordered_partial_stream; +mod partial_reduce_stream; +mod single_stream; +mod skip_partial; mod topk; -mod topk_stream; /// Returns true if TopK aggregation data structures support the provided key and value types. /// @@ -207,6 +349,12 @@ pub enum AggregateMode { /// / \ / \ /// Partial Partial Partial Partial /// ``` + /// + /// # Motivation + /// + /// This reduces shuffling traffic in a distributed setting. See + /// + /// for details. PartialReduce, } @@ -497,10 +645,51 @@ impl PartialEq for PhysicalGroupBy { } } +/// Streams used by [`AggregateExec`]. +/// +/// # Stream Variant Schema Notation +/// For example, `SELECT g, AVG(x) FROM t GROUP BY g` uses these schemas: +/// +/// ```text +/// initial input: [g, x] +/// partial state: [g, AVG(x) state columns, e.g. sum/count] +/// final result: [g, AVG(x)] +/// ``` #[expect(clippy::large_enum_variant)] enum StreamType { + /// Single group (no group by) aggregate stream. + /// Input output scheme: initial input -> final result AggregateStream(AggregateStream), + /// Partial stage of the hash aggregation + /// Input output scheme: initial input -> partial state + PartialHash(PartialHashAggregateStream), + /// Partial-reduce stage of the hash aggregation + /// Input output scheme: partial state -> partial state + PartialReduceHash(PartialReduceHashAggregateStream), + /// Final stage of the hash aggregation + /// Input output scheme: partial state -> final result + FinalHash(FinalHashAggregateStream), + /// Single stage of the hash aggregation + /// Input output scheme: initial input -> final result + SingleHash(SingleHashAggregateStream), + /// Partial stage of aggregation for ordered input. + OrderedPartialAggregate(OrderedPartialAggregateStream), + /// Final stage of aggregation for ordered input. + OrderedFinalAggregate(OrderedFinalAggregateStream), + /// Hash aggregation reused for multiple stages + /// + /// Note this is being incrementally migrated to dedicated streams like + /// [`StreamType::PartialHash`], [`StreamType::FinalHash`], + /// [`StreamType::OrderedPartialAggregate`], and + /// [`StreamType::OrderedFinalAggregate`] + /// + /// See issue for details: GroupedHash(GroupedHashAggregateStream), + /// Grouped TopK aggregate stream. + /// Input output scheme: initial input -> final result + /// + /// Used for grouped aggregation with LIMIT / ordering, where the stream keeps + /// only the top groups required by the query. GroupedPriorityQueue(GroupedTopKAggregateStream), } @@ -508,6 +697,12 @@ impl From for SendableRecordBatchStream { fn from(stream: StreamType) -> Self { match stream { StreamType::AggregateStream(stream) => Box::pin(stream), + StreamType::PartialHash(stream) => Box::pin(stream), + StreamType::PartialReduceHash(stream) => Box::pin(stream), + StreamType::FinalHash(stream) => Box::pin(stream), + StreamType::SingleHash(stream) => Box::pin(stream), + StreamType::OrderedPartialAggregate(stream) => stream.into_stream(), + StreamType::OrderedFinalAggregate(stream) => Box::pin(stream), StreamType::GroupedHash(stream) => Box::pin(stream), StreamType::GroupedPriorityQueue(stream) => Box::pin(stream), } @@ -551,7 +746,7 @@ impl From for SendableRecordBatchStream { /// The filter is kept in the `DataSourceExec`, and it will gets update during execution, /// the reader will interpret it as "the upstream only needs rows that such filter /// predicate is evaluated to true", and certain scanner implementation like `parquet` -/// can evalaute column statistics on those dynamic filters, to decide if they can +/// can evaluate column statistics on those dynamic filters, to decide if they can /// prune a whole range. /// /// ### Examples @@ -832,6 +1027,7 @@ impl AggregateExec { &input, Arc::clone(&schema), &group_expr_mapping, + group_by.is_true_no_grouping(), &mode, &input_order_mode, aggr_expr.as_ref(), @@ -895,6 +1091,10 @@ impl AggregateExec { } /// Returns the dynamic filter expression for this aggregate, if set. + #[deprecated( + since = "55.0.0", + note = "Use ExecutionPlan::dynamic_expressions_produced instead" + )] pub fn dynamic_filter_expr(&self) -> Option<&Arc> { self.dynamic_filter.as_ref().map(|df| &df.filter) } @@ -945,6 +1145,8 @@ impl AggregateExec { Arc::clone(&self.input_schema) } + /// Aggregation has multiple specialized implementations optimized for + /// different workloads. This function picks the best available path. fn execute_typed( &self, partition: usize, @@ -965,12 +1167,135 @@ impl AggregateExec { )); } - // grouping by something else and we need to just materialize all results + // Select the stream type based on the query shape and configuration. + // For an overview, see the `Aggregate planning` section in this file's + // documentation. + // + // # Implementation Note + // + // `GroupedHashAggregateStream` is being incrementally refactored. See the + // tracking issue for details. + // + // New features and improvements should go directly into the new implementation. + // Please coordinate through the tracking issue. + // + // Issue: + if context + .session_config() + .options() + .execution + .enable_migration_aggregate + { + if self.should_use_ordered_partial_aggregate_stream(context) { + return Ok(StreamType::OrderedPartialAggregate( + OrderedPartialAggregateStream::new(self, context, partition)?, + )); + } + + if self.should_use_partial_hash_stream(context) { + return Ok(StreamType::PartialHash(PartialHashAggregateStream::new( + self, context, partition, + )?)); + } + + if self.should_use_partial_reduce_hash_stream(context) { + return Ok(StreamType::PartialReduceHash( + PartialReduceHashAggregateStream::new(self, context, partition)?, + )); + } + + if self.should_use_ordered_final_aggregate_stream(context) { + return Ok(StreamType::OrderedFinalAggregate( + OrderedFinalAggregateStream::new(self, context, partition)?, + )); + } + + if self.should_use_final_hash_stream(context) { + return Ok(StreamType::FinalHash(FinalHashAggregateStream::new( + self, context, partition, + )?)); + } + + if self.should_use_single_hash_stream(context) { + return Ok(StreamType::SingleHash(SingleHashAggregateStream::new( + self, context, partition, + )?)); + } + } + + // Execution paths that have not been migrated use the fallback implementation Ok(StreamType::GroupedHash(GroupedHashAggregateStream::new( self, context, partition, )?)) } + fn should_use_partial_hash_stream(&self, _context: &TaskContext) -> bool { + self.mode == AggregateMode::Partial + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + && self.limit_options_supported_by_hash_stream() + } + + fn should_use_ordered_partial_aggregate_stream( + &self, + _context: &TaskContext, + ) -> bool { + self.mode == AggregateMode::Partial + && self.input_order_mode != InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + && self.limit_options_supported_by_hash_stream() + } + + fn should_use_final_hash_stream(&self, _context: &TaskContext) -> bool { + matches!( + self.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) && self.limit_options_supported_by_hash_stream() + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + + fn should_use_partial_reduce_hash_stream(&self, context: &TaskContext) -> bool { + // TODO: implement memory-limited path and remove this limitation + if matches!(context.memory_pool().memory_limit(), MemoryLimit::Finite(_)) { + return false; + } + + self.mode == AggregateMode::PartialReduce + && self.limit_options.is_none() + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + + fn should_use_single_hash_stream(&self, _context: &TaskContext) -> bool { + matches!( + self.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + ) && self.limit_options.is_none() + && self.input_order_mode == InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + + fn should_use_ordered_final_aggregate_stream(&self, _context: &TaskContext) -> bool { + matches!( + self.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + ) && self.limit_options_supported_by_hash_stream() + && self.input_order_mode != InputOrderMode::Linear + && !self.group_by.is_true_no_grouping() + && self.group_by.is_single() + } + + /// See comments in `PartialHashAggregateStream` limit optimization section + fn limit_options_supported_by_hash_stream(&self) -> bool { + self.limit_options.is_none() || self.is_unordered_unfiltered_group_by_distinct() + } + /// Finds the DataType and SortDirection for this Aggregate, if there is one pub fn get_minmax_desc(&self) -> Option<(FieldRef, bool)> { let agg_expr = self.aggr_expr.iter().exactly_one().ok()?; @@ -1022,6 +1347,7 @@ impl AggregateExec { input: &Arc, schema: SchemaRef, group_expr_mapping: &ProjectionMapping, + is_true_no_grouping: bool, mode: &AggregateMode, input_order_mode: &InputOrderMode, aggr_exprs: &[Arc], @@ -1031,9 +1357,12 @@ impl AggregateExec { .equivalence_properties() .project(group_expr_mapping, schema); - // If the group by is empty, then we ensure that the operator will produce - // only one row, and mark the generated result as a constant value. - if group_expr_mapping.is_empty() { + // True no-group aggregates produce only one row in each output + // partition, so aggregate outputs are constants within the partition. + // Grouping sets with empty grouping expressions are not covered here: + // their output schema can include grouping-set columns before the + // aggregate columns, so this aggregate-column mapping does not apply. + if is_true_no_grouping { let new_constants = aggr_exprs.iter().enumerate().map(|(idx, func)| { let column = Arc::new(Column::new(func.name(), idx)); ConstExpr::from(column as Arc) @@ -1092,6 +1421,11 @@ impl AggregateExec { /// Estimates output statistics for this aggregate node. /// + /// For aggregations without group-by expressions, row count follows the + /// number of logical aggregate rows and the aggregate output mode. True + /// no-group aggregates have one logical row; empty grouping sets have one + /// logical row per grouping-set occurrence. + /// /// For grouped aggregations with known input row count > 1, the output row /// count is estimated as: /// @@ -1129,7 +1463,11 @@ impl AggregateExec { /// - Per-set products are summed across all grouping sets /// - Requires NDV stats for ALL active group-by columns; if any lacks stats, /// falls back to `input_rows` (or `Absent` if that is also unknown) - fn statistics_inner(&self, child_statistics: &Statistics) -> Result { + fn statistics_inner( + &self, + child_statistics: &Statistics, + partition: Option, + ) -> Result { // TODO stats: group expressions: // - once expressions will be able to compute their own stats, use it here // - case where we group by on a column for which with have the `distinct` stat @@ -1153,21 +1491,24 @@ impl AggregateExec { column_statistics }; - match self.mode { - AggregateMode::Final | AggregateMode::FinalPartitioned - if self.group_by.expr.is_empty() => - { + match self.exact_output_rows_without_group_exprs(partition) { + Some(output_rows) => { let total_byte_size = - Self::calculate_scaled_byte_size(child_statistics, 1); + Self::calculate_scaled_byte_size(child_statistics, output_rows); Ok(Statistics { - num_rows: Precision::Exact(1), + num_rows: Precision::Exact(output_rows), column_statistics, total_byte_size, }) } - _ => { - let num_rows = self.estimate_num_rows(child_statistics); + None => { + let num_rows = self.estimate_num_rows(child_statistics, partition); + let column_statistics = self.nullify_group_columns_for_empty_input( + column_statistics, + child_statistics, + &num_rows, + ); let total_byte_size = num_rows .get_value() @@ -1187,9 +1528,124 @@ impl AggregateExec { } } + /// Exact physical output row count for aggregates without group-by + /// expressions. + /// + /// `partition` follows [`ExecutionPlan::partition_statistics`]: `Some(_)` + /// requests one output partition, while `None` requests the entire plan. + /// Partial-state output contains the logical rows in each output partition; + /// final-value output contains the global logical rows once. + /// This mirrors execution, where partial aggregation without group-by + /// expressions emits its logical rows from every output partition, including + /// empty input partitions. + /// + /// Returns `None` when grouping expressions are present and grouped + /// cardinality estimation should be used instead. + fn exact_output_rows_without_group_exprs( + &self, + partition: Option, + ) -> Option { + let logical_rows = self.logical_rows_without_group_exprs()?; + + Some(self.scale_logical_rows(logical_rows, partition)) + } + + /// Scales a logical aggregate row count to the rows this operator emits, + /// which for partial aggregation is once per output partition. + fn scale_logical_rows(&self, logical_rows: usize, partition: Option) -> usize { + match (self.mode.output_mode(), partition) { + (AggregateOutputMode::Final, _) => logical_rows, + (AggregateOutputMode::Partial, Some(_)) => logical_rows, + (AggregateOutputMode::Partial, None) => { + logical_rows * self.cache.output_partitioning().partition_count() + } + } + } + + /// Number of rows a grouped aggregate emits for an empty input. + /// + /// Grouping expressions yield no groups, so the only rows are the + /// grand-total rows of the empty grouping sets that `GROUPING SETS(())`, + /// `ROLLUP` and `CUBE` introduce alongside the non-empty ones. + fn output_rows_for_empty_input(&self, partition: Option) -> usize { + let empty_grouping_sets = self + .group_by + .groups + .iter() + .filter(|nulls| nulls.iter().all(|is_null| *is_null)) + .count(); + + self.scale_logical_rows(empty_grouping_sets, partition) + } + + /// Reports the grouping columns of an empty input as all NULL. + /// + /// The only rows such an input produces are grand-total rows, which hold + /// NULL in every grouping column, so the values copied from the child do not + /// describe the output. Rules that answer `MIN`/`MAX` from statistics read + /// these values, so an input value here becomes a wrong query result. + /// + /// The bounds are typed nulls rather than [`Precision::Absent`], both + /// because NULL is the `MIN`/`MAX` of such a column and because the data + /// type lets downstream interval analysis keep intersecting intervals of + /// that type, as `FilterExec` does for a column with no rows. + fn nullify_group_columns_for_empty_input( + &self, + mut column_statistics: Vec, + child_statistics: &Statistics, + num_rows: &Precision, + ) -> Vec { + let empty_input = child_statistics.num_rows.get_value() == Some(&0); + let emits_rows = num_rows.get_value().is_some_and(|&rows| rows > 0); + if !empty_input || !emits_rows { + return column_statistics; + } + + let schema = self.schema(); + for (idx, column_stats) in column_statistics + .iter_mut() + .take(self.group_by.expr.len()) + .enumerate() + { + let typed_null = ScalarValue::try_from(schema.field(idx).data_type()) + .unwrap_or(ScalarValue::Null); + let mut null_bound = Precision::Exact(typed_null); + if matches!(num_rows, Precision::Inexact(_)) { + null_bound = null_bound.to_inexact(); + } + column_stats.min_value = null_bound.clone(); + column_stats.max_value = null_bound; + column_stats.distinct_count = num_rows.map(|_| 0); + column_stats.null_count = *num_rows; + } + + column_statistics + } + + /// Exact number of logical aggregate rows for aggregates without group-by + /// expressions. + /// + /// A true no-group aggregate has one logical aggregate row. Empty grouping + /// sets have one logical aggregate row per grouping-set occurrence, even + /// when there are duplicate empty grouping sets. Returns `None` when there + /// are grouping expressions. + fn logical_rows_without_group_exprs(&self) -> Option { + if self.group_by.is_true_no_grouping() { + Some(1) + } else if self.group_by.expr.is_empty() { + Some(self.group_by.groups.len()) + } else { + None + } + } + /// Estimates the output row count for grouped aggregations, combining NDV, /// input row count, and TopK limit into a single [`Precision`]. - fn estimate_num_rows(&self, child_statistics: &Statistics) -> Precision { + fn estimate_num_rows( + &self, + child_statistics: &Statistics, + partition: Option, + ) -> Precision { let ndv = if !self.group_by.expr.is_empty() { self.compute_group_ndv(child_statistics) } else { @@ -1208,7 +1664,11 @@ impl AggregateExec { } num_rows } else if value == 0 { - child_statistics.num_rows + // The limit bounds groups built from input rows, not the rows + // the empty grouping sets contribute. + child_statistics + .num_rows + .map(|_| self.output_rows_for_empty_input(partition)) } else { let grouping_set_num = self.group_by.groups.len(); let mut num_rows = @@ -1372,17 +1832,6 @@ impl AggregateExec { _ => Precision::Absent, } } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for AggregateExec { @@ -1545,17 +1994,21 @@ impl ExecutionPlan for AggregateExec { } fn required_input_distribution(&self) -> Vec { - match &self.mode { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + InputDistributionRequirements::new(match &self.mode { AggregateMode::Partial | AggregateMode::PartialReduce => { vec![Distribution::UnspecifiedDistribution] } AggregateMode::FinalPartitioned | AggregateMode::SinglePartitioned => { - vec![Distribution::HashPartitioned(self.group_by.input_exprs())] + vec![Distribution::KeyPartitioned(self.group_by.input_exprs())] } AggregateMode::Final | AggregateMode::Single => { vec![Distribution::SinglePartition] } - } + }) } fn required_input_ordering(&self) -> Vec> { @@ -1579,55 +2032,90 @@ impl ExecutionPlan for AggregateExec { vec![&self.input] } - fn apply_expressions( - &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, - ) -> Result { - // Apply to group by expressions - let mut tnr = TreeNodeRecursion::Continue; - for expr in self.group_by.input_exprs() { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - - // Apply to aggregate expressions - for aggr in self.aggr_expr.iter() { - for expr in aggr.expressions() { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut me = AggregateExec::try_new_with_schema( + self.mode, + Arc::clone(&self.group_by), + self.aggr_expr.to_vec(), + Arc::clone(&self.filter_expr), + Arc::clone(&children[0]), + Arc::clone(&self.input_schema), + Arc::clone(&self.schema), + )?; + me.limit_options = self.limit_options; + me.dynamic_filter.clone_from(&self.dynamic_filter); + Ok(Arc::new(me)) } } - - // Apply to filter expressions (FILTER WHERE clauses) - for filter in self.filter_expr.iter().flatten() { - tnr = tnr.visit_sibling(|| f(filter.as_ref()))?; - } - - // Apply to dynamic filter expression if present - if let Some(dyn_filter) = &self.dynamic_filter { - tnr = tnr.visit_sibling(|| f(dyn_filter.filter.as_ref()))?; - } - - Ok(tnr) } fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } - let mut me = AggregateExec::try_new_with_schema( - self.mode, - Arc::clone(&self.group_by), - self.aggr_expr.to_vec(), - Arc::clone(&self.filter_expr), - Arc::clone(&children[0]), - Arc::clone(&self.input_schema), - Arc::clone(&self.schema), - )?; - me.limit_options = self.limit_options; - me.dynamic_filter.clone_from(&self.dynamic_filter); + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + let group_by = self.group_by.input_exprs(); + let aggregates = self.aggr_expr.iter().flat_map(|aggr| { + let expressions = aggr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.order_by_exprs) + }); + let filters = self.filter_expr.iter().flatten().cloned(); + let dynamic_filter = self.dynamic_filter.iter().map(|dynamic_filter| { + Arc::::clone(&dynamic_filter.filter) + as Arc + }); + crate::apply_expression_roots( + group_by + .into_iter() + .chain(aggregates) + .chain(filters) + .chain(dynamic_filter), + f, + ) + } + + fn dynamic_expressions_produced(&self) -> Vec> { + self.dynamic_filter + .iter() + .map(|dynamic_filter| { + Arc::::clone(&dynamic_filter.filter) + as Arc + }) + .collect() + } - Ok(Arc::new(me)) + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -1643,9 +2131,19 @@ impl ExecutionPlan for AggregateExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let child_statistics = self.input().partition_statistics(partition)?; - Ok(Arc::new(self.statistics_inner(&child_statistics)?)) + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + let child_statistics = Arc::clone(&input_stats[0]); + Ok(Arc::new( + self.statistics_inner(&child_statistics, args.partition())?, + )) } fn cardinality_effect(&self) -> CardinalityEffect { @@ -1668,70 +2166,35 @@ impl ExecutionPlan for AggregateExec { // This optimization is NOT safe for filters on aggregated columns (like filtering on // the result of SUM or COUNT), as those require computing all groups first. - // Build grouping columns using output indices because parent filters reference the - // AggregateExec's output schema where grouping columns in the output schema. The - // grouping expressions reference input columns which may not match the output schema. - // - // It is safe to assume that the output_schema contains group by columns in the same order - // as the group by expression. See [`create_schema`] and [`AggregateExec`]. - let output_schema = self.schema(); - let grouping_columns: HashSet<_> = (0..self.group_by.expr().len()) - .map(|i| Column::new(output_schema.field(i).name(), i)) - .collect(); - - // Analyze each filter separately to determine if it can be pushed down - let mut safe_filters = Vec::new(); - let mut unsafe_filters = Vec::new(); - - for filter in parent_filters { - let filter_columns: HashSet<_> = - collect_columns(&filter).into_iter().collect(); - - // Check if this filter references non-grouping columns - let references_non_grouping = !grouping_columns.is_empty() - && !filter_columns.is_subset(&grouping_columns); - - if references_non_grouping { - unsafe_filters.push(filter); - continue; - } - - // For GROUPING SETS, verify this filter's columns appear in all grouping sets - if self.group_by.groups().len() > 1 { - let filter_column_indices: Vec = filter_columns - .iter() - .filter_map(|filter_col| { - grouping_columns.get(filter_col).map(|col| col.index()) - }) - .collect(); - - // Check if any of this filter's columns are missing from any grouping set - let has_missing_column = self.group_by.groups().iter().any(|null_mask| { - filter_column_indices - .iter() - .any(|&idx| null_mask.get(idx) == Some(&true)) - }); - - if has_missing_column { - unsafe_filters.push(filter); - continue; - } - } - - // This filter is safe to push down - safe_filters.push(filter); + // Grouping columns are output before aggregate columns, in the same order + // as the grouping expressions. A grouping-set null mask marks grouping + // columns that are not available in that set. + let mut allowed_indices: HashSet = + (0..self.group_by.expr().len()).collect(); + for null_mask in self.group_by.groups() { + allowed_indices.retain(|idx| null_mask.get(*idx) != Some(&true)); } - // Build child filter description with both safe and unsafe filters let child = self.children()[0]; - let mut child_desc = ChildFilterDescription::from_child(&safe_filters, child)?; - - // Add unsafe filters as unsupported - child_desc.parent_filters.extend( - unsafe_filters - .into_iter() - .map(PushedDownPredicate::unsupported), - ); + // Global aggregates and grouping sets containing an empty grouping set + // emit a row even when their input is empty. Parent filters therefore + // cannot be pushed below them, including filters without column + // references. + let may_emit_on_empty_input = self.group_by.is_true_no_grouping() + || self + .group_by + .groups() + .iter() + .any(|null_mask| null_mask.iter().all(|is_null| *is_null)); + let mut child_desc = if may_emit_on_empty_input { + ChildFilterDescription::all_unsupported(&parent_filters) + } else { + ChildFilterDescription::from_child_with_allowed_indices( + &parent_filters, + allowed_indices, + child, + )? + }; // Include self dynamic filter when it's possible if phase == FilterPushdownPhase::Post @@ -1760,27 +2223,12 @@ impl ExecutionPlan for AggregateExec { if phase == FilterPushdownPhase::Post && let Some(dyn_filter) = &self.dynamic_filter { - // let child_accepts_dyn_filter = child_pushdown_result - // .self_filters - // .first() - // .map(|filters| { - // assert_eq_or_internal_err!( - // filters.len(), - // 1, - // "Aggregate only pushdown one self dynamic filter" - // ); - // let filter = filters.get(0).unwrap(); // Asserted above - // Ok(matches!(filter.discriminant, PushedDown::Yes)) - // }) - // .unwrap_or_else(|| internal_err!("The length of self filters equals to the number of child of this ExecutionPlan, so it must be 1"))?; - - // HACK: The above snippet should be used, however, now the child reply - // `PushDown::No` can indicate they're not able to push down row-level - // filter, but still keep the filter for statistics pruning. - // So here, we try to use ref count to determine if the dynamic filter - // has actually be pushed down. - // Issue: - let child_accepts_dyn_filter = Arc::strong_count(dyn_filter) > 1; + let child_accepts_dyn_filter = dyn_filter + .filter + .expression_id() + .map(|id| plan_contains_expression_id(&self.input, id)) + .transpose()? + .unwrap_or(false); if !child_accepts_dyn_filter { // Child can't consume the self dynamic filter, so disable it by setting @@ -1795,72 +2243,477 @@ impl ExecutionPlan for AggregateExec { Ok(result) } -} - -/// Creates the output schema for an [`AggregateExec`] containing the group by columns followed -/// by the aggregate columns. -fn create_schema( - input_schema: &Schema, - group_by: &PhysicalGroupBy, - aggr_expr: &[Arc], - mode: AggregateMode, -) -> Result { - let mut fields = Vec::with_capacity(group_by.num_output_exprs() + aggr_expr.len()); - fields.extend(group_by.output_fields(input_schema)?); - match mode.output_mode() { - AggregateOutputMode::Final => { - // in final mode, the field with the final result of the accumulator - for expr in aggr_expr { - fields.push(expr.field()) + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `AggregateExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + mode, + group_by, + aggr_expr, + filter_expr, + limit_options, + input, + // Derived at construction by `create_schema` from `input_schema`, + // `group_by`, `aggr_expr` and `mode`. + schema: _, + input_schema, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + // Derived at construction from the input ordering and `group_by`. + required_input_ordering: _, + // Derived at construction from the input ordering and `group_by`. + input_order_mode: _, + // Derived at construction by `Self::compute_properties`. + cache: _, + dynamic_filter, + } = self; + + let input = ctx.encode_child(input)?; + let group_expr = + ctx.encode_expressions(group_by.expr().iter().map(|(expr, _)| expr))?; + let group_expr_name = group_by + .expr() + .iter() + .map(|(_, name)| name.to_owned()) + .collect(); + let null_expr = + ctx.encode_expressions(group_by.null_expr().iter().map(|(expr, _)| expr))?; + let groups = group_by.groups().iter().flatten().copied().collect(); + let aggr_expr_name = aggr_expr + .iter() + .map(|expr| expr.name().to_string()) + .collect(); + let aggr_expr = aggr_expr + .iter() + .map(|expr| encode_aggregate_expr(expr, ctx)) + .collect::>>()?; + let filter_expr = filter_expr + .iter() + .map(|filter| { + Ok(protobuf::MaybeFilter { + expr: filter + .as_ref() + .map(|expr| ctx.encode_expr(expr)) + .transpose()?, + }) + }) + .collect::>>()?; + // Match by name because the protobuf and execution enums use different + // discriminants, so a numeric cast would corrupt the wire format. + let mode = match mode { + AggregateMode::Partial => protobuf::AggregateMode::Partial, + AggregateMode::Final => protobuf::AggregateMode::Final, + AggregateMode::FinalPartitioned => protobuf::AggregateMode::FinalPartitioned, + AggregateMode::Single => protobuf::AggregateMode::Single, + AggregateMode::SinglePartitioned => { + protobuf::AggregateMode::SinglePartitioned } - } - AggregateOutputMode::Partial => { - // in partial mode, the fields of the accumulator's state - for expr in aggr_expr { - fields.extend(expr.state_fields()?.iter().cloned()); + AggregateMode::PartialReduce => protobuf::AggregateMode::PartialReduce, + }; + let limit = limit_options.map(|options| protobuf::AggLimit { + limit: options.limit() as u64, + descending: options.descending(), + }); + // Only the shared `filter` expr is on the wire; the accumulator bounds + // in `AggrDynFilter` are runtime state repopulated during execution. + let dynamic_filter = match dynamic_filter { + Some(dynamic_filter) => { + let expr: Arc = + Arc::clone(&dynamic_filter.filter) as Arc; + Some(ctx.encode_expr(&expr)?) } - } + None => None, + }; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Aggregate(Box::new( + protobuf::AggregateExecNode { + group_expr, + group_expr_name, + aggr_expr, + filter_expr, + aggr_expr_name, + mode: mode as i32, + input: Some(Box::new(input)), + input_schema: Some(input_schema.as_ref().try_into()?), + null_expr, + groups, + limit, + has_grouping_set: group_by.has_grouping_set(), + dynamic_filter, + schema: Some(self.schema.as_ref().try_into()?), + }, + )), + ), + })) } +} - Ok(Schema::new_with_metadata( - fields, - input_schema.metadata().clone(), - )) +/// Keep this marker byte-identical to the copy used by the deprecated +/// aggregate serializer in `datafusion-proto` until that path is removed. +#[cfg(feature = "proto")] +const HUMAN_DISPLAY_ALIAS_PREFIX: &str = "\u{1f}datafusion_human_display_alias_v1:"; + +#[cfg(feature = "proto")] +fn encode_human_display_alias(human_display: &str, alias: &str) -> String { + format!( + "{HUMAN_DISPLAY_ALIAS_PREFIX}{}:{alias}{human_display}", + alias.len() + ) } -/// Determines the lexical ordering requirement for an aggregate expression. -/// -/// # Parameters -/// -/// - `aggr_expr`: A reference to an `AggregateFunctionExpr` representing the -/// aggregate expression. -/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the -/// physical GROUP BY expression. -/// - `agg_mode`: A reference to an `AggregateMode` instance representing the -/// mode of aggregation. -/// - `include_soft_requirement`: When `false`, only hard requirements are -/// considered, as indicated by [`AggregateFunctionExpr::order_sensitivity`] -/// returning [`AggregateOrderSensitivity::HardRequirement`]. -/// Otherwise, also soft requirements ([`AggregateOrderSensitivity::SoftRequirement`]) -/// are considered. -/// -/// # Returns -/// -/// A `LexOrdering` instance indicating the lexical ordering requirement for -/// the aggregate expression. -fn get_aggregate_expr_req( - aggr_expr: &AggregateFunctionExpr, - group_by: &PhysicalGroupBy, - agg_mode: &AggregateMode, - include_soft_requirement: bool, -) -> Option { - // If the aggregation is performing a "second stage" calculation, - // then ignore the ordering requirement. Ordering requirement applies - // only to the aggregation input data. - if agg_mode.input_mode() == AggregateInputMode::Partial { - return None; - } +#[cfg(feature = "proto")] +fn split_human_display_alias<'a>( + human_display: &'a str, + name: &'a str, +) -> (&'a str, Option<&'a str>) { + if let Some(encoded) = human_display.strip_prefix(HUMAN_DISPLAY_ALIAS_PREFIX) + && let Some((alias_len, encoded)) = encoded.split_once(':') + && let Ok(alias_len) = alias_len.parse::() + && let Some(alias) = encoded.get(..alias_len) + && let Some(human_display) = encoded.get(alias_len..) + && alias == name + && !human_display.is_empty() + { + return (human_display, Some(alias)); + } + + (human_display, None) +} + +#[cfg(feature = "proto")] +fn encode_aggregate_expr( + aggr_expr: &Arc, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, +) -> Result { + use datafusion_proto_models::protobuf; + + let expressions = aggr_expr.expressions(); + let expr = ctx.encode_expressions(expressions.iter())?; + let ordering_req = + datafusion_physical_expr_common::sort_expr::sort_exprs_try_to_proto( + aggr_expr.order_bys(), + &ctx.expr_ctx(), + )?; + let name = aggr_expr.fun().name().to_string(); + // The context already applies `(!buf.is_empty()).then_some(buf)`. + let fun_definition = ctx.encode_udaf(aggr_expr.fun())?; + let human_display = match (aggr_expr.human_display(), aggr_expr.human_display_alias()) + { + (Some(display), Some(alias)) => encode_human_display_alias(display, alias), + (Some(display), None) => display.to_string(), + (None, _) => String::new(), + }; + + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::AggregateExpr( + protobuf::PhysicalAggregateExprNode { + aggregate_function: Some( + protobuf::physical_aggregate_expr_node::AggregateFunction::UserDefinedAggrFunction(name), + ), + expr, + ordering_req, + distinct: aggr_expr.is_distinct(), + ignore_nulls: aggr_expr.ignore_nulls(), + fun_definition, + human_display, + is_reversed: aggr_expr.is_reversed(), + }, + )), + }) +} + +#[cfg(feature = "proto")] +impl AggregateExec { + /// Reconstruct an [`AggregateExec`] from its protobuf representation. + /// + /// Grouping expressions are decoded against the child schema. Aggregate + /// arguments, ordering, filters, and the dynamic filter are decoded against + /// the aggregate input schema carried in the protobuf node. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_proto_models::protobuf; + use protobuf::physical_aggregate_expr_node::AggregateFunction; + use protobuf::physical_expr_node::ExprType; + + let hash_agg = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Aggregate, + "AggregateExec", + ); + // Exhaustive destructure: a new field on `AggregateExecNode` is a + // compile error here rather than a silently ignored wire field. + let protobuf::AggregateExecNode { + group_expr, + aggr_expr, + mode, + input, + group_expr_name, + aggr_expr_name, + input_schema, + null_expr, + groups, + filter_expr, + limit, + has_grouping_set, + dynamic_filter, + schema, + } = hash_agg.as_ref(); + + let input = + ctx.decode_required_child(input.as_deref(), "AggregateExec", "input")?; + // Match by name because the protobuf and execution enums use different + // discriminants, so a numeric cast would corrupt the wire format. + let mode = protobuf::AggregateMode::try_from(*mode).map_err(|_| { + datafusion_common::internal_datafusion_err!( + "Received an AggregateNode message with unknown AggregateMode {mode}" + ) + })?; + let mode = match mode { + protobuf::AggregateMode::Partial => AggregateMode::Partial, + protobuf::AggregateMode::Final => AggregateMode::Final, + protobuf::AggregateMode::FinalPartitioned => AggregateMode::FinalPartitioned, + protobuf::AggregateMode::Single => AggregateMode::Single, + protobuf::AggregateMode::SinglePartitioned => { + AggregateMode::SinglePartitioned + } + protobuf::AggregateMode::PartialReduce => AggregateMode::PartialReduce, + }; + let num_expr = group_expr.len(); + // Grouping expressions refer to the child plan's output schema. + let child_schema = input.schema(); + let group_expr = group_expr + .iter() + .zip(group_expr_name.iter()) + .map(|(expr, name)| { + Ok(( + ctx.decode_expr(expr, child_schema.as_ref())?, + name.to_string(), + )) + }) + .collect::>>()?; + let null_expr = null_expr + .iter() + .zip(group_expr_name.iter()) + .map(|(expr, name)| { + Ok(( + ctx.decode_expr(expr, child_schema.as_ref())?, + name.to_string(), + )) + }) + .collect::>>()?; + let groups = if groups.is_empty() { + vec![] + } else { + groups + .chunks(num_expr) + .map(|group| group.to_vec()) + .collect() + }; + // Aggregate arguments, ordering, filters, and dynamic filters refer to + // the aggregate input schema carried in the protobuf node. + let input_schema = input_schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "input_schema in AggregateNode is missing." + ) + })?; + let input_schema: SchemaRef = SchemaRef::new(input_schema.try_into()?); + let filter_expr = filter_expr + .iter() + .map(|filter| { + filter + .expr + .as_ref() + .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) + .transpose() + }) + .collect::>>()?; + let aggr_expr = aggr_expr + .iter() + .zip(aggr_expr_name.iter()) + .map(|(expr, name)| { + let expr_type = expr.expr_type.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "Unexpected empty aggregate physical expression" + ) + })?; + let ExprType::AggregateExpr(aggregate) = expr_type else { + return internal_err!( + "Invalid aggregate expression for AggregateExec" + ); + }; + let args = aggregate + .expr + .iter() + .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) + .collect::>>()?; + let order_by = + datafusion_physical_expr_common::sort_expr::sort_exprs_try_from_proto( + &aggregate.ordering_req, + &ctx.expr_ctx(input_schema.as_ref()), + )?; + let Some(AggregateFunction::UserDefinedAggrFunction(udaf_name)) = + aggregate.aggregate_function.as_ref() + else { + return internal_err!( + "Invalid AggregateExpr, missing aggregate_function" + ); + }; + // The context owns the payload-to-codec and + // registry-to-codec fallback order. + let udaf = + ctx.decode_udaf(udaf_name, aggregate.fun_definition.as_deref())?; + let (human_display, human_display_alias) = + split_human_display_alias(&aggregate.human_display, name); + let builder = AggregateExprBuilder::new(udaf, args) + .schema(Arc::clone(&input_schema)) + .alias(name) + .with_ignore_nulls(aggregate.ignore_nulls) + .with_distinct(aggregate.distinct) + .order_by(order_by) + .with_reversed(aggregate.is_reversed) + .human_display(human_display); + let builder = if let Some(alias) = human_display_alias { + builder.human_display_alias(alias) + } else { + builder + }; + builder.build().map(Arc::new) + }) + .collect::>>()?; + let group_by = + PhysicalGroupBy::new(group_expr, null_expr, groups, *has_grouping_set); + let aggregate = if let Some(schema) = schema { + let schema = SchemaRef::new(schema.try_into()?); + AggregateExec::try_new_with_schema( + mode, + group_by, + aggr_expr, + filter_expr, + input, + Arc::clone(&input_schema), + schema, + ) + } else { + AggregateExec::try_new( + mode, + group_by, + aggr_expr, + filter_expr, + input, + Arc::clone(&input_schema), + ) + }?; + let aggregate = if let Some(limit) = limit { + let options = match limit.descending { + Some(descending) => { + LimitOptions::new_with_order(limit.limit as usize, descending) + } + None => LimitOptions::new(limit.limit as usize), + }; + aggregate.with_limit_options(Some(options)) + } else { + aggregate + }; + let aggregate = if let Some(dynamic_filter) = dynamic_filter { + let dynamic_filter = + ctx.decode_expr(dynamic_filter, input_schema.as_ref())?; + let dynamic_filter = (dynamic_filter + as Arc) + .downcast::() + .map_err(|_| { + datafusion_common::internal_datafusion_err!( + "AggregateExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + aggregate.with_dynamic_filter_expr(dynamic_filter)? + } else { + let mut aggregate = aggregate; + aggregate.dynamic_filter = None; + aggregate + }; + + Ok(Arc::new(aggregate)) + } +} + +/// Creates the output schema for an [`AggregateExec`] containing the group by columns followed +/// by the aggregate columns. +fn create_schema( + input_schema: &Schema, + group_by: &PhysicalGroupBy, + aggr_expr: &[Arc], + mode: AggregateMode, +) -> Result { + let mut fields = Vec::with_capacity(group_by.num_output_exprs() + aggr_expr.len()); + fields.extend(group_by.output_fields(input_schema)?); + + match mode.output_mode() { + AggregateOutputMode::Final => { + // in final mode, the field with the final result of the accumulator + for expr in aggr_expr { + fields.push(expr.field()) + } + } + AggregateOutputMode::Partial => { + // in partial mode, the fields of the accumulator's state + for expr in aggr_expr { + fields.extend(expr.state_fields()?.iter().cloned()); + } + } + } + + Ok(Schema::new_with_metadata( + fields, + input_schema.metadata().clone(), + )) +} + +/// Determines the lexical ordering requirement for an aggregate expression. +/// +/// # Parameters +/// +/// - `aggr_expr`: A reference to an `AggregateFunctionExpr` representing the +/// aggregate expression. +/// - `group_by`: A reference to a `PhysicalGroupBy` instance representing the +/// physical GROUP BY expression. +/// - `agg_mode`: A reference to an `AggregateMode` instance representing the +/// mode of aggregation. +/// - `include_soft_requirement`: When `false`, only hard requirements are +/// considered, as indicated by [`AggregateFunctionExpr::order_sensitivity`] +/// returning [`AggregateOrderSensitivity::HardRequirement`]. +/// Otherwise, also soft requirements ([`AggregateOrderSensitivity::SoftRequirement`]) +/// are considered. +/// +/// # Returns +/// +/// A `LexOrdering` instance indicating the lexical ordering requirement for +/// the aggregate expression. +fn get_aggregate_expr_req( + aggr_expr: &AggregateFunctionExpr, + group_by: &PhysicalGroupBy, + agg_mode: &AggregateMode, + include_soft_requirement: bool, +) -> Option { + // If the aggregation is performing a "second stage" calculation, + // then ignore the ordering requirement. Ordering requirement applies + // only to the aggregation input data. + if agg_mode.input_mode() == AggregateInputMode::Partial { + return None; + } match aggr_expr.order_sensitivity() { AggregateOrderSensitivity::Insensitive => return None, @@ -2211,6 +3064,32 @@ pub(crate) fn max_duplicate_ordinal(groups: &[Vec]) -> usize { /// The outer Vec appears to be for grouping sets /// The inner Vec contains the results per expression /// The inner-inner Array contains the results per row +/// +/// For example, for `GROUP BY GROUPING SETS ((a, b), (a))` with input: +/// +/// ```text +/// a b +/// 1 1 +/// 1 2 +/// 2 1 +/// ``` +/// +/// The output is: +/// +/// ```text +/// [ +/// [ +/// a: [1, 1, 2] +/// b: [1, 2, 1] +/// grouping_id: [0, 0, 0] +/// ], +/// [ +/// a: [1, 1, 2] +/// b: [NULL, NULL, NULL] +/// grouping_id: [1, 1, 1] +/// ] +/// ] +/// ``` pub fn evaluate_group_by( group_by: &PhysicalGroupBy, batch: &RecordBatch, @@ -2267,7 +3146,9 @@ mod tests { use crate::empty::EmptyExec; use crate::execution_plan::Boundedness; use crate::expressions::col; + use crate::filter::FilterExecBuilder; use crate::metrics::MetricValue; + use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test::TestMemoryExec; use crate::test::assert_is_pending; use crate::test::exec::{ @@ -2275,8 +3156,8 @@ mod tests { }; use arrow::array::{ - DictionaryArray, Float32Array, Float64Array, Int32Array, Int64Array, StructArray, - UInt32Array, UInt64Array, + BooleanArray, DictionaryArray, Float32Array, Float64Array, Int32Array, + Int64Array, StructArray, UInt32Array, UInt64Array, }; use arrow::compute::{SortOptions, concat_batches}; use arrow::datatypes::Int32Type; @@ -2285,6 +3166,12 @@ mod tests { use datafusion_execution::config::SessionConfig; use datafusion_execution::memory_pool::FairSpillPool; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; + use datafusion_expr::{ + Accumulator, AggregateUDF, AggregateUDFImpl, EmitTo, GroupsAccumulator, + Signature, Volatility, + }; + use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; use datafusion_functions_aggregate::array_agg::array_agg_udaf; use datafusion_functions_aggregate::average::avg_udaf; use datafusion_functions_aggregate::count::count_udaf; @@ -2295,13 +3182,36 @@ mod tests { use datafusion_physical_expr::Partitioning; use datafusion_physical_expr::PhysicalSortExpr; use datafusion_physical_expr::aggregate::AggregateExprBuilder; - use datafusion_physical_expr::expressions::Literal; + use datafusion_physical_expr::expressions::{Literal, NotExpr}; use crate::projection::ProjectionExec; + use crate::repartition::RepartitionExec; use datafusion_physical_expr::projection::ProjectionExpr; - use futures::{FutureExt, Stream}; + use futures::{FutureExt, Stream, StreamExt}; use insta::{allow_duplicates, assert_snapshot}; + #[cfg(feature = "proto")] + #[test] + fn split_human_display_alias_ignores_mismatched_alias() { + let encoded = encode_human_display_alias("sum(value)", "revenue"); + + assert_eq!( + split_human_display_alias(&encoded, "other"), + (encoded.as_str(), None) + ); + } + + #[cfg(feature = "proto")] + #[test] + fn split_human_display_alias_keeps_malformed_prefix_literal() { + let display = format!("{HUMAN_DISPLAY_ALIAS_PREFIX}not-an-encoding"); + + assert_eq!( + split_human_display_alias(&display, "agg"), + (display.as_str(), None) + ); + } + // Generate a schema which consists of 5 columns (a, b, c, d, e) fn create_test_schema() -> Result { let a = Field::new("a", DataType::Int32, true); @@ -2409,6 +3319,34 @@ mod tests { Arc::new(task_ctx) } + fn migrated_hash_session_config(batch_size: usize) -> SessionConfig { + SessionConfig::new() + .with_batch_size(batch_size) + .set_bool("datafusion.execution.enable_migration_aggregate", true) + } + + fn new_migrated_hash_ctx(batch_size: usize) -> Arc { + Arc::new( + TaskContext::default() + .with_session_config(migrated_hash_session_config(batch_size)), + ) + } + + fn new_finite_memory_migrated_hash_ctx( + batch_size: usize, + max_memory: usize, + ) -> Result> { + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(max_memory, 1.0) + .build_arc()?; + + Ok(Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(migrated_hash_session_config(batch_size)), + )) + } + async fn check_grouping_sets( input: Arc, spill: bool, @@ -2618,7 +3556,8 @@ mod tests { | 2 | 1 | 1.0 | | 3 | 1 | 2.0 | | 3 | 2 | 5.0 | - | 4 | 3 | 11.0 | + | 4 | 1 | 4.0 | + | 4 | 2 | 7.0 | +---+---------------+-------------+ "); } @@ -2650,12 +3589,13 @@ mod tests { )?); // Verify statistics are preserved proportionally through aggregation - let final_stats = merged_aggregate.partition_statistics(None)?; + let final_stats = StatisticsContext::new() + .compute(merged_aggregate.as_ref(), &StatisticsArgs::new())?; assert!(final_stats.total_byte_size.get_value().is_some()); let task_ctx = if spill { // enlarge memory limit to let the final aggregation finish - new_spill_ctx(2, 2600) + new_spill_ctx(2, 4640) } else { Arc::clone(&task_ctx) }; @@ -2684,17 +3624,12 @@ mod tests { let spilled_bytes = metrics.spilled_bytes().unwrap(); let spilled_rows = metrics.spilled_rows().unwrap(); + assert_eq!(3, output_rows); if spill { - // When spilling, the output rows metrics become partial output size + final output size - // This is because final aggregation starts while partial aggregation is still emitting - assert_eq!(8, output_rows); - assert!(spill_count > 0); assert!(spilled_bytes > 0); assert!(spilled_rows > 0); } else { - assert_eq!(3, output_rows); - assert_eq!(0, spill_count); assert_eq!(0, spilled_bytes); assert_eq!(0, spilled_rows); @@ -2764,18 +3699,29 @@ mod tests { vec![] } - fn apply_expressions( - &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, - ) -> Result { - Ok(TreeNodeRecursion::Continue) + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + internal_err!("Children cannot be replaced in {self:?}") } fn with_new_children( self: Arc, - _: Vec>, + children: Vec>, ) -> Result> { - internal_err!("Children cannot be replaced in {self:?}") + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) } fn execute( @@ -2792,11 +3738,12 @@ mod tests { Ok(Box::pin(stream)) } - fn partition_statistics( + fn statistics_from_inputs( &self, - partition: Option, + _input_stats: &[Arc], + args: &StatisticsArgs, ) -> Result> { - if partition.is_some() { + if args.partition().is_some() { return Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref()))); } let (_, batches) = some_data(); @@ -2937,7 +3884,7 @@ mod tests { let aggregates_v0: Vec> = vec![Arc::new(test_median_agg_expr(Arc::clone(&input_schema))?)]; - // use fast-path in `row_hash.rs`. + // Use the fast path in `single_stream.rs`. let aggregates_v2: Vec> = vec![Arc::new( AggregateExprBuilder::new(avg_udaf(), vec![col("b", &input_schema)?]) .schema(Arc::clone(&input_schema)) @@ -2970,7 +3917,7 @@ mod tests { assert!(matches!(stream, StreamType::GroupedHash(_))); } 2 => { - assert!(matches!(stream, StreamType::GroupedHash(_))); + assert!(matches!(stream, StreamType::SingleHash(_))); } _ => panic!("Unknown version: {version}"), } @@ -2990,62 +3937,722 @@ mod tests { } #[tokio::test] - async fn test_drop_cancel_without_groups() -> Result<()> { - let task_ctx = Arc::new(TaskContext::default()); - let schema = - Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)])); - - let groups = PhysicalGroupBy::default(); - + async fn partial_grouped_aggregate_uses_raw_partial_stream() -> Result<()> { + let (schema, batches) = some_data(); + let input = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64], + vec![DataType::Int32], + DataType::Int64, + ))); let aggregates: Vec> = vec![Arc::new( - AggregateExprBuilder::new(avg_udaf(), vec![col("a", &schema)?]) + AggregateExprBuilder::new(udaf, vec![col("b", &schema)?]) .schema(Arc::clone(&schema)) - .alias("AVG(a)") + .alias("input_type_asserting(b)") .build()?, )]; - let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); - let refs = blocking_exec.refs(); - let aggregate_exec = Arc::new(AggregateExec::try_new( + let partial_aggregate = Arc::new(AggregateExec::try_new( AggregateMode::Partial, - groups.clone(), + group_by.clone(), aggregates.clone(), vec![None], - blocking_exec, - schema, + input, + Arc::clone(&schema), )?); + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(2) + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); - let fut = crate::collect(aggregate_exec, task_ctx); - let mut fut = fut.boxed(); + let partial_stream = partial_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(partial_stream, StreamType::PartialHash(_))); - assert_is_pending(&mut fut); - drop(fut); - assert_strong_count_converges_to_zero(refs).await; + let fallback_task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(2) + .set_bool("datafusion.execution.enable_migration_aggregate", false), + ), + ); + let stream = partial_aggregate.execute_typed(0, &fallback_task_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + + let stream: SendableRecordBatchStream = partial_stream.into(); + let batches = collect(stream).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 1] + ); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + + let merge = Arc::new(CoalescePartitionsExec::new(partial_aggregate)); + let final_aggregate = AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + aggregates, + vec![None], + merge, + Arc::clone(&schema), + )?; + + let final_stream = final_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(final_stream, StreamType::FinalHash(_))); + + let stream = final_aggregate.execute_typed(0, &fallback_task_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + + let stream: SendableRecordBatchStream = final_stream.into(); + let batches = collect(stream).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 1] + ); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); Ok(()) } #[tokio::test] - async fn test_drop_cancel_with_groups() -> Result<()> { - let task_ctx = Arc::new(TaskContext::default()); + async fn partial_grouped_aggregate_materializes_before_slicing() -> Result<()> { let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Float64, true), - Field::new("b", DataType::Float64, true), + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int32, false), ])); - - let groups = - PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); - + let input_batches = vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![10, 20, 30])), + ], + )?]; + let input = + TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let udaf = Arc::new(AggregateUDF::from(NoFirstEmitUdaf::new())); let aggregates: Vec> = vec![Arc::new( - AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) + AggregateExprBuilder::new(udaf, vec![col("value", &schema)?]) .schema(Arc::clone(&schema)) - .alias("AVG(b)") + .alias("no_first_emit(value)") .build()?, )]; - - let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); - let refs = blocking_exec.refs(); - let aggregate_exec = Arc::new(AggregateExec::try_new( + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggregates, + vec![None], + input, + Arc::clone(&schema), + )?); + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(2) + .set_bool("datafusion.execution.enable_migration_aggregate", true) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(2.0)), + ), + ), + ); + + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::PartialHash(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let batches = collect(stream).await?; + assert_eq!( + batches + .iter() + .map(RecordBatch::num_rows) + .collect::>(), + vec![2, 1] + ); + assert_eq!(batches.iter().map(RecordBatch::num_rows).sum::(), 3); + assert_snapshot!(batches_to_sort_string(&batches), @r" + +-----+-----------------------------+ + | key | no_first_emit(value)[count] | + +-----+-----------------------------+ + | 1 | 1 | + | 2 | 1 | + | 3 | 1 | + +-----+-----------------------------+ + "); + + Ok(()) + } + + #[tokio::test] + async fn limited_distinct_aggregate_uses_migrated_hash_streams() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, false)])); + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![1, 2, 1]))], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![3, 4]))], + )?, + ]; + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); + + let partial_input = TestMemoryExec::try_new_exec( + std::slice::from_ref(&input_batches), + Arc::clone(&schema), + None, + )?; + let partial_aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + vec![], + vec![], + partial_input, + Arc::clone(&schema), + )? + .with_limit_options(Some(LimitOptions::new(2))), + ); + + let partial_stream = partial_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(partial_stream, StreamType::PartialHash(_))); + let stream: SendableRecordBatchStream = partial_stream.into(); + let partial_output = collect(stream).await?; + assert_eq!( + partial_output + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 2 + ); + assert_snapshot!(batches_to_sort_string(&partial_output), @r" ++---+ +| a | ++---+ +| 1 | +| 2 | ++---+ +"); + + let final_input = + TestMemoryExec::try_new_exec(&[input_batches], Arc::clone(&schema), None)?; + let final_aggregate = Arc::new( + AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + vec![], + vec![], + final_input, + Arc::clone(&schema), + )? + .with_limit_options(Some(LimitOptions::new(2))), + ); + + let final_stream = final_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(final_stream, StreamType::FinalHash(_))); + let stream: SendableRecordBatchStream = final_stream.into(); + let final_output = collect(stream).await?; + assert_eq!( + final_output + .iter() + .map(RecordBatch::num_rows) + .sum::(), + 2 + ); + assert_snapshot!(batches_to_sort_string(&final_output), @r" ++---+ +| a | ++---+ +| 1 | +| 2 | ++---+ +"); + + Ok(()) + } + + fn single_test_aggregate() -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + let input_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 1, 3])), + Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])), + ], + )?; + let input = TestMemoryExec::try_new_exec( + &[vec![input_batch]], + Arc::clone(&schema), + None, + )?; + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let aggregates: Vec> = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )]; + + AggregateExec::try_new( + AggregateMode::Single, + group_by, + aggregates, + vec![None], + input, + schema, + ) + } + + /// For single aggregation, ensures `SingleHashAggregateStream` is used when + /// enabled by migration config. + #[tokio::test] + async fn single_aggregate_planning() -> Result<()> { + let single = single_test_aggregate()?; + let task_ctx = new_migrated_hash_ctx(2); + + let stream = single.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 3); + assert_snapshot!(batches_to_sort_string(&output), @r" ++---+--------+ +| a | SUM(b) | ++---+--------+ +| 1 | 50.0 | +| 2 | 20.0 | +| 3 | 30.0 | ++---+--------+ +"); + + Ok(()) + } + + /// Single hash aggregation supports finite memory. + #[tokio::test] + async fn single_aggregate_with_memory_limit_planning() -> Result<()> { + let single = single_test_aggregate()?; + let task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?; + + let stream = single.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::SingleHash(_))); + + Ok(()) + } + + fn partial_reduce_test_aggregate() -> Result { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let aggregates: Vec> = vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("SUM(b)") + .build()?, + )]; + + let empty_input = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggregates.clone(), + vec![None], + empty_input, + Arc::clone(&schema), + )?; + let partial_schema = partial.schema(); + let partial_state_batch = RecordBatch::try_new( + Arc::clone(&partial_schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 1, 3])), + Arc::new(Float64Array::from(vec![10.0, 20.0, 40.0, 30.0])), + ], + )?; + let partial_reduce_input = TestMemoryExec::try_new_exec( + &[vec![partial_state_batch]], + Arc::clone(&partial_schema), + None, + )?; + + AggregateExec::try_new( + AggregateMode::PartialReduce, + group_by, + aggregates, + vec![None], + partial_reduce_input, + partial_schema, + ) + } + + /// For partial-reduce aggregation, ensures `PartialReduceHashAggregateStream` + /// is used when enabled by migration config. + #[tokio::test] + async fn partial_reduce_aggregate_planning() -> Result<()> { + let partial_reduce = partial_reduce_test_aggregate()?; + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .set_bool("datafusion.execution.enable_migration_aggregate", true), + ), + ); + + let stream = partial_reduce.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::PartialReduceHash(_))); + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_eq!(output.iter().map(RecordBatch::num_rows).sum::(), 3); + + Ok(()) + } + + /// Spilling behavior is not implemented for partial-reduce stream yet, so fall + /// back to the existing `GroupedHashAggregateStream` + #[tokio::test] + async fn partial_reduce_aggregate_with_memory_limit_planning() -> Result<()> { + let partial_reduce = partial_reduce_test_aggregate()?; + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(1, 1.0) + .build_arc()?; + let task_ctx = + Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::new().set_bool( + "datafusion.execution.enable_migration_aggregate", + true, + )) + .with_runtime(runtime), + ); + + let stream = partial_reduce.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::GroupedHash(_))); + + Ok(()) + } + + /// Ensures for ordered input, `OrderedPartialAggregateStream` is used. + #[tokio::test] + async fn ordered_partial_aggregate_planning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("sort_col", DataType::Int32, false), + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + let input_batches = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 1])), + Arc::new(Int32Array::from(vec![10, 11, 10])), + Arc::new(Int64Array::from(vec![1, 1, 1])), + ], + )?, + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 2])), + Arc::new(Int32Array::from(vec![20, 21])), + Arc::new(Int64Array::from(vec![1, 1])), + ], + )?, + ]; + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("sort_col", 0), + ))]) + .unwrap(); + let input = TestMemoryExec::try_new(&[input_batches], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + + let group_by = PhysicalGroupBy::new_single(vec![ + (col("sort_col", &schema)?, "sort_col".to_string()), + (col("group_col", &schema)?, "group_col".to_string()), + ]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value_col)") + .build()?, + )]; + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + input, + Arc::clone(&schema), + )?; + assert!(matches!( + aggregate.input_order_mode(), + InputOrderMode::PartiallySorted(_) + )); + + let task_ctx = new_migrated_hash_ctx(2); + let stream = aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::OrderedPartialAggregate(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++----------+-----------+-------------------------+ +| sort_col | group_col | COUNT(value_col)[count] | ++----------+-----------+-------------------------+ +| 1 | 10 | 2 | +| 1 | 11 | 1 | +| 2 | 20 | 1 | +| 2 | 21 | 1 | ++----------+-----------+-------------------------+ +"); + + // Ordered partial aggregation supports finite memory. + let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?; + let stream = aggregate.execute_typed(0, &finite_memory_task_ctx)?; + assert!(matches!(stream, StreamType::OrderedPartialAggregate(_))); + + Ok(()) + } + + /// Ensures for ordered input, `OrderedFinalAggregateStream` is used. + #[tokio::test] + async fn ordered_final_aggregate_planning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, false), + Field::new("value", DataType::Int64, false), + ])); + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("COUNT(value)") + .build()?, + )]; + + let empty_input = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema), None)?; + let partial_aggregate = AggregateExec::try_new( + AggregateMode::Partial, + group_by.clone(), + aggr_expr.clone(), + vec![None], + empty_input, + Arc::clone(&schema), + )?; + let partial_schema = partial_aggregate.schema(); + let partial_state_batch = RecordBatch::try_new( + Arc::clone(&partial_schema), + vec![ + Arc::new(Int32Array::from(vec![1, 1, 2, 3])), + Arc::new(Int64Array::from(vec![2, 3, 5, 7])), + ], + )?; + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("key", 0), + ))]) + .unwrap(); + let final_input = + TestMemoryExec::try_new(&[vec![partial_state_batch]], partial_schema, None)? + .try_with_sort_information(vec![ordering])?; + let final_input = Arc::new(TestMemoryExec::update_cache(&Arc::new(final_input))); + + let final_aggregate = AggregateExec::try_new( + AggregateMode::Final, + group_by.as_final(), + aggr_expr, + vec![None], + final_input, + Arc::clone(&schema), + )?; + assert_eq!(final_aggregate.input_order_mode(), &InputOrderMode::Sorted); + + let task_ctx = new_migrated_hash_ctx(2); + let stream = final_aggregate.execute_typed(0, &task_ctx)?; + assert!(matches!(stream, StreamType::OrderedFinalAggregate(_))); + + let stream: SendableRecordBatchStream = stream.into(); + let output = collect(stream).await?; + assert_snapshot!(batches_to_sort_string(&output), @r" ++-----+--------------+ +| key | COUNT(value) | ++-----+--------------+ +| 1 | 5 | +| 2 | 5 | +| 3 | 7 | ++-----+--------------+ +"); + + // Ordered final aggregation supports finite memory. + let finite_memory_task_ctx = new_finite_memory_migrated_hash_ctx(2, 1024 * 1024)?; + let stream = final_aggregate.execute_typed(0, &finite_memory_task_ctx)?; + assert!(matches!(stream, StreamType::OrderedFinalAggregate(_))); + + Ok(()) + } + + #[tokio::test] + async fn ordered_partial_aggregate_partially_sorted_no_emit_panic() -> Result<()> { + // Reproducer for #20445: emitting from PartiallySorted input must not + // drain more groups than the completed sort boundary allows. + let schema = Arc::new(Schema::new(vec![ + Field::new("sort_col", DataType::Int32, false), + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int64, false), + ])); + + // All rows share sort_col=1, so there is no completed sort boundary + // inside this batch even though there are many distinct groups. + let n = 256; + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1; n])), + Arc::new(Int32Array::from((0..n as i32).collect::>())), + Arc::new(Int64Array::from(vec![1; n])), + ], + )?; + + let ordering = LexOrdering::new([PhysicalSortExpr::new_default(Arc::new( + Column::new("sort_col", 0), + ))]) + .unwrap(); + let input = TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let input = Arc::new(TestMemoryExec::update_cache(&Arc::new(input))); + + let aggregate = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![ + (col("sort_col", &schema)?, "sort_col".to_string()), + (col("group_col", &schema)?, "group_col".to_string()), + ]), + vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )], + vec![None], + input, + Arc::clone(&schema), + )?; + assert!(matches!( + aggregate.input_order_mode(), + InputOrderMode::PartiallySorted(_) + )); + + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(4096, 1.0) + .build_arc()?; + let session_config = SessionConfig::new().with_batch_size(128).set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::UInt64(Some(u64::MAX)), + ); + let task_ctx = Arc::new( + TaskContext::default() + .with_runtime(runtime) + .with_session_config(session_config), + ); + + let mut stream: SendableRecordBatchStream = + OrderedPartialAggregateStream::new(&aggregate, &task_ctx, 0)?.into_stream(); + + while let Some(result) = stream.next().await { + if let Err(e) = result { + if e.to_string().contains("Resources exhausted") { + break; + } + return Err(e); + } + } + + Ok(()) + } + + #[tokio::test] + async fn test_drop_cancel_without_groups() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Float64, true)])); + + let groups = PhysicalGroupBy::default(); + + let aggregates: Vec> = vec![Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("a", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("AVG(a)") + .build()?, + )]; + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let refs = blocking_exec.refs(); + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups.clone(), + aggregates.clone(), + vec![None], + blocking_exec, + schema, + )?); + + let fut = crate::collect(aggregate_exec, task_ctx); + let mut fut = fut.boxed(); + + assert_is_pending(&mut fut); + drop(fut); + assert_strong_count_converges_to_zero(refs).await; + + Ok(()) + } + + #[tokio::test] + async fn test_drop_cancel_with_groups() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, true), + Field::new("b", DataType::Float64, true), + ])); + + let groups = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + + let aggregates: Vec> = vec![Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("AVG(b)") + .build()?, + )]; + + let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); + let refs = blocking_exec.refs(); + let aggregate_exec = Arc::new(AggregateExec::try_new( AggregateMode::Partial, groups, aggregates.clone(), @@ -3068,7 +4675,7 @@ mod tests { async fn run_first_last_multi_partitions() -> Result<()> { for is_first_acc in [false, true] { for spill in [false, true] { - first_last_multi_partitions(is_first_acc, spill, 4200).await? + first_last_multi_partitions(is_first_acc, spill, 5000).await? } } Ok(()) @@ -3451,8 +5058,10 @@ mod tests { Arc::clone(&blocking_exec) as Arc, schema, )?); - let new_agg = - Arc::clone(&aggregate_exec).with_new_children(vec![blocking_exec])?; + let new_agg = Arc::clone(&aggregate_exec).replace_children( + vec![blocking_exec], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; assert_eq!(new_agg.schema(), aggregate_exec.schema()); Ok(()) } @@ -3627,40 +5236,217 @@ mod tests { .map(Arc::new)?, ]; - let input = TestMemoryExec::try_new_exec( - &[vec![batch.clone()]], - Arc::::clone(&batch.schema()), - None, - )?; + let input = TestMemoryExec::try_new_exec( + &[vec![batch.clone()]], + Arc::::clone(&batch.schema()), + None, + )?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::FinalPartitioned, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc, + batch.schema(), + )?); + + let session_config = SessionConfig::default(); + let ctx = TaskContext::default().with_session_config(session_config); + let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_string(&output), @r" + +--------------+------------+ + | labels | SUM(value) | + +--------------+------------+ + | {a: a, b: b} | 2 | + | {a: , b: c} | 1 | + +--------------+------------+ + "); + } + + Ok(()) + } + + // Migrated to PartialHashAggregateStream coverage below; + // kept here for the legacy GroupedHashAggregateStream implementation. + #[tokio::test] + async fn test_skip_aggregation_after_first_batch() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + ])); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(String::from("COUNT(val)")) + .build() + .map(Arc::new)?, + ]; + + let input_data = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + ]; + + let input = + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc, + schema, + )?); + + let mut session_config = SessionConfig::default(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(2)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(0.1)), + ); + + let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let stream: SendableRecordBatchStream = Box::pin( + GroupedHashAggregateStream::new(aggregate_exec.as_ref(), &ctx, 0)?, + ); + let output = collect(stream).await?; + + allow_duplicates! { + assert_snapshot!(batches_to_string(&output), @r" + +-----+-------------------+ + | key | COUNT(val)[count] | + +-----+-------------------+ + | 1 | 1 | + | 2 | 1 | + | 3 | 1 | + | 2 | 1 | + | 3 | 1 | + | 4 | 1 | + +-----+-------------------+ + "); + } + + Ok(()) + } + + // Migrated to PartialHashAggregateStream coverage below; + // kept here for the legacy GroupedHashAggregateStream implementation. + #[tokio::test] + async fn test_skip_aggregation_after_threshold() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + ])); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(String::from("COUNT(val)")) + .build() + .map(Arc::new)?, + ]; + + let input_data = vec![ + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![2, 3, 4])), + Arc::new(Int32Array::from(vec![0, 0, 0])), + ], + ) + .unwrap(), + ]; + + let input = + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; let aggregate_exec = Arc::new(AggregateExec::try_new( - AggregateMode::FinalPartitioned, + AggregateMode::Partial, group_by, aggr_expr, vec![None], Arc::clone(&input) as Arc, - batch.schema(), + schema, )?); - let session_config = SessionConfig::default(); - let ctx = TaskContext::default().with_session_config(session_config); - let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?; + let mut session_config = SessionConfig::default(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(5)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(0.1)), + ); + + let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let stream: SendableRecordBatchStream = Box::pin( + GroupedHashAggregateStream::new(aggregate_exec.as_ref(), &ctx, 0)?, + ); + let output = collect(stream).await?; allow_duplicates! { - assert_snapshot!(batches_to_string(&output), @r" - +--------------+------------+ - | labels | SUM(value) | - +--------------+------------+ - | {a: a, b: b} | 2 | - | {a: , b: c} | 1 | - +--------------+------------+ - "); + assert_snapshot!(batches_to_string(&output), @r" + +-----+-------------------+ + | key | COUNT(val)[count] | + +-----+-------------------+ + | 1 | 1 | + | 2 | 2 | + | 3 | 2 | + | 4 | 1 | + | 2 | 1 | + | 3 | 1 | + | 4 | 1 | + +-----+-------------------+ + "); } Ok(()) } #[tokio::test] - async fn test_skip_aggregation_after_first_batch() -> Result<()> { + async fn test_partial_hash_stream_skip_aggregation_after_first_batch() -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Int32, true), Field::new("val", DataType::Int32, true), @@ -3707,39 +5493,47 @@ mod tests { schema, )?); - let mut session_config = SessionConfig::default(); - session_config = session_config.set( - "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", - &ScalarValue::Int64(Some(2)), - ); - session_config = session_config.set( - "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", - &ScalarValue::Float64(Some(0.1)), - ); + let session_config = SessionConfig::default() + .set_bool("datafusion.execution.enable_migration_aggregate", true) + .set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(2)), + ) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(0.1)), + ); - let ctx = TaskContext::default().with_session_config(session_config); - let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?; + let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let output = collect(aggregate_exec.execute(0, Arc::clone(&ctx))?).await?; allow_duplicates! { - assert_snapshot!(batches_to_string(&output), @r" + assert_snapshot!(batches_to_sort_string(&output), @r" +-----+-------------------+ | key | COUNT(val)[count] | +-----+-------------------+ | 1 | 1 | | 2 | 1 | - | 3 | 1 | | 2 | 1 | | 3 | 1 | + | 3 | 1 | | 4 | 1 | +-----+-------------------+ "); } + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert_eq!(skipped_rows, 3); + Ok(()) } #[tokio::test] - async fn test_skip_aggregation_after_threshold() -> Result<()> { + async fn test_partial_hash_stream_skip_aggregation_after_threshold() -> Result<()> { let schema = Arc::new(Schema::new(vec![ Field::new("key", DataType::Int32, true), Field::new("val", DataType::Int32, true), @@ -3794,35 +5588,127 @@ mod tests { schema, )?); - let mut session_config = SessionConfig::default(); - session_config = session_config.set( - "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", - &ScalarValue::Int64(Some(5)), - ); - session_config = session_config.set( - "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", - &ScalarValue::Float64(Some(0.1)), - ); + let session_config = SessionConfig::default() + .set_bool("datafusion.execution.enable_migration_aggregate", true) + .set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(5)), + ) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(0.1)), + ); - let ctx = TaskContext::default().with_session_config(session_config); - let output = collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?; + let ctx = Arc::new(TaskContext::default().with_session_config(session_config)); + let output = collect(aggregate_exec.execute(0, Arc::clone(&ctx))?).await?; allow_duplicates! { - assert_snapshot!(batches_to_string(&output), @r" - +-----+-------------------+ - | key | COUNT(val)[count] | - +-----+-------------------+ - | 1 | 1 | - | 2 | 2 | - | 3 | 2 | - | 4 | 1 | - | 2 | 1 | - | 3 | 1 | - | 4 | 1 | - +-----+-------------------+ - "); + assert_snapshot!(batches_to_sort_string(&output), @r" + +-----+-------------------+ + | key | COUNT(val)[count] | + +-----+-------------------+ + | 1 | 1 | + | 2 | 1 | + | 2 | 2 | + | 3 | 1 | + | 3 | 2 | + | 4 | 1 | + | 4 | 1 | + +-----+-------------------+ + "); } + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert_eq!(skipped_rows, 3); + + Ok(()) + } + + /// When `skip_partial_aggregation_probe_ratio_threshold` is set to 1.0, + /// the feature must be effectively disabled: even with 100% cardinality + /// (every row is a unique group), no rows should be skipped. + #[tokio::test] + async fn test_skip_aggregation_disabled_at_threshold_one() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("key", DataType::Int32, true), + Field::new("val", DataType::Int32, true), + ])); + + let group_by = + PhysicalGroupBy::new_single(vec![(col("key", &schema)?, "key".to_string())]); + + let aggr_expr = vec![ + AggregateExprBuilder::new(count_udaf(), vec![col("val", &schema)?]) + .schema(Arc::clone(&schema)) + .alias(String::from("COUNT(val)")) + .build() + .map(Arc::new)?, + ]; + + // Two batches are required: batch 1 triggers the probe threshold so the + // skip decision is evaluated; batch 2 is what would be skipped on main + // (where >= caused threshold=1.0 to still skip at 100% cardinality). + // All rows have unique keys => ratio = 1.0 (100% cardinality). + let input_data = vec![ + // Batch 1: fires the probe check (ratio = 5/5 = 1.0) + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5])), + Arc::new(Int32Array::from(vec![0, 0, 0, 0, 0])), + ], + ) + .unwrap(), + // Batch 2: would be skipped if threshold=1.0 did not disable the feature + RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![6, 7, 8, 9, 10])), + Arc::new(Int32Array::from(vec![0, 0, 0, 0, 0])), + ], + ) + .unwrap(), + ]; + + let input = + TestMemoryExec::try_new_exec(&[input_data], Arc::clone(&schema), None)?; + let aggregate_exec = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + group_by, + aggr_expr, + vec![None], + Arc::clone(&input) as Arc, + schema, + )?); + + let session_config = SessionConfig::default() + .set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &ScalarValue::Int64(Some(1)), + ) + .set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &ScalarValue::Float64(Some(1.0)), + ); + + let ctx = TaskContext::default().with_session_config(session_config); + collect(aggregate_exec.execute(0, Arc::new(ctx))?).await?; + + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + + assert_eq!( + skipped_rows, 0, + "threshold=1.0 should disable skip aggregation, but {skipped_rows} rows were skipped" + ); + Ok(()) } @@ -3895,9 +5781,11 @@ mod tests { Field::new("b", DataType::Float64, false), ])); + let group_keys = [2, 3, 4, 4].repeat(1_000); + let values = [1.0, 2.0, 3.0, 4.0].repeat(1_000); let batches = vec![ - create_record_batch(&schema, (vec![2, 3, 4, 4], vec![1.0, 2.0, 3.0, 4.0]))?, - create_record_batch(&schema, (vec![2, 3, 4, 4], vec![1.0, 2.0, 3.0, 4.0]))?, + create_record_batch(&schema, (group_keys.clone(), values.clone()))?, + create_record_batch(&schema, (group_keys, values))?, ]; let plan: Arc = TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; @@ -3997,9 +5885,9 @@ mod tests { #[tokio::test] async fn test_aggregate_with_spill_if_necessary() -> Result<()> { // test with spill - run_test_with_spill_pool_if_necessary(2_000, true).await?; + run_test_with_spill_pool_if_necessary(20_000, true).await?; // test without spill - run_test_with_spill_pool_if_necessary(20_000, false).await?; + run_test_with_spill_pool_if_necessary(200_000, false).await?; Ok(()) } @@ -4134,7 +6022,7 @@ mod tests { PhysicalGroupBy::default(), None, )?; - let stats = agg.partition_statistics(None)?; + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; assert_eq!(stats.total_byte_size, Precision::Absent); let zero_row_stats = Statistics { @@ -4151,30 +6039,226 @@ mod tests { PhysicalGroupBy::default(), None, )?; - let stats_zero = agg_zero.partition_statistics(None)?; + let stats_zero = + StatisticsContext::new().compute(&agg_zero, &StatisticsArgs::new())?; assert_eq!(stats_zero.total_byte_size, Precision::Absent); + let single_input = + Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; + let single_agg_zero = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::default(), + vec![count_a_aggregate(&schema)?], + vec![None], + single_input, + Arc::clone(&schema), + )?; + assert_eq!( + single_agg_zero + .properties() + .output_partitioning() + .partition_count(), + 1 + ); + let single_stats_zero = + StatisticsContext::new().compute(&single_agg_zero, &StatisticsArgs::new())?; + assert_eq!(single_stats_zero.num_rows, Precision::Exact(1)); + + Ok(()) + } + + #[tokio::test] + async fn test_aggregate_statistics_empty_input_with_grouping_sets() -> Result<()> { + let schema = empty_grouping_sets_test_schema(); + + // `GROUP BY a` produces no groups for an empty input. + let grouped = build_test_aggregate( + &schema, + empty_input_statistics(), + simple_group_by(&schema, &["a"]), + None, + )?; + let stats = StatisticsContext::new().compute(&grouped, &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(0)); + + // `GROUPING SETS((a), ())`, as ROLLUP and CUBE produce, still emits the + // grand-total row of the empty grouping set on an empty input. + let with_empty_set = build_test_aggregate( + &schema, + empty_input_statistics(), + grouping_sets_with_empty(&schema, 1)?, + None, + )?; + let stats = + StatisticsContext::new().compute(&with_empty_set, &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(1)); + + // `GROUPING SETS((a), (), ())` emits one grand-total row per empty + // grouping set, because execution gives each duplicate its own ordinal. + let with_duplicate_empty_sets = build_test_aggregate( + &schema, + empty_input_statistics(), + grouping_sets_with_empty(&schema, 2)?, + None, + )?; + let stats = StatisticsContext::new() + .compute(&with_duplicate_empty_sets, &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(2)); + + Ok(()) + } + + /// Partial aggregation emits the grand-total row from every output + /// partition, so the whole-plan estimate scales with the partition count + /// while a single-partition request does not. + #[tokio::test] + async fn test_aggregate_statistics_empty_input_partial_mode_scaling() -> Result<()> { + let schema = empty_grouping_sets_test_schema(); + let input = Arc::new(RepartitionExec::try_new( + Arc::new(StatisticsExec::new( + empty_input_statistics(), + (*schema).clone(), + )), + Partitioning::RoundRobinBatch(4), + )?) as Arc; + + let agg = AggregateExec::try_new( + AggregateMode::Partial, + grouping_sets_with_empty(&schema, 1)?, + vec![count_a_aggregate(&schema)?], + vec![None], + input, + Arc::clone(&schema), + )?; + assert_eq!(agg.properties().output_partitioning().partition_count(), 4); + + let context = StatisticsContext::new(); + assert_eq!( + context.compute(&agg, &StatisticsArgs::new())?.num_rows, + Precision::Exact(4) + ); + // Inexact because a repartition only estimates its per-partition row + // count. The grouping column statistics carry that same precision. + let partition_statistics = + context.compute(&agg, &StatisticsArgs::new().with_partition(Some(0)))?; + assert_eq!(partition_statistics.num_rows, Precision::Inexact(1)); + let group_column = &partition_statistics.column_statistics[0]; + let typed_null = Precision::Inexact(ScalarValue::Int32(None)); + assert_eq!(group_column.min_value, typed_null); + assert_eq!(group_column.max_value, typed_null); + assert_eq!(group_column.distinct_count, Precision::Inexact(0)); + assert_eq!(group_column.null_count, Precision::Inexact(1)); + + Ok(()) + } + + /// The input's min, max and distinct values must not reach the output + /// column statistics. See `nullify_group_columns_for_empty_input`. + #[tokio::test] + async fn test_aggregate_statistics_empty_input_nullifies_group_columns() -> Result<()> + { + let schema = empty_grouping_sets_test_schema(); + let mut input_statistics = empty_input_statistics(); + input_statistics.column_statistics[0] = ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Exact(ScalarValue::Int32(Some(5))), + min_value: Precision::Exact(ScalarValue::Int32(Some(5))), + sum_value: Precision::Absent, + distinct_count: Precision::Exact(1), + byte_size: Precision::Absent, + }; + + let agg = build_test_aggregate( + &schema, + input_statistics, + grouping_sets_with_empty(&schema, 1)?, + None, + )?; + + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; + assert_eq!(stats.num_rows, Precision::Exact(1)); + let group_column = &stats.column_statistics[0]; + let typed_null = Precision::Exact(ScalarValue::Int32(None)); + assert_eq!(group_column.min_value, typed_null); + assert_eq!(group_column.max_value, typed_null); + assert_eq!(group_column.distinct_count, Precision::Exact(0)); + assert_eq!(group_column.null_count, Precision::Exact(1)); + Ok(()) } + fn empty_grouping_sets_test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Float64, false), + ])) + } + + fn empty_input_statistics() -> Statistics { + Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ + ColumnStatistics::new_unknown(), + ColumnStatistics::new_unknown(), + ], + } + } + + /// `GROUPING SETS((a), (), ...)` with `empty_sets` empty grouping sets, as + /// `ROLLUP(a)` and `CUBE(a)` produce with one. + fn grouping_sets_with_empty( + schema: &SchemaRef, + empty_sets: usize, + ) -> Result { + let mut groups = vec![vec![false]]; + groups.resize(1 + empty_sets, vec![true]); + Ok(PhysicalGroupBy::new( + vec![(col("a", schema)?, "a".to_string())], + vec![(lit(ScalarValue::Int32(None)), "a".to_string())], + groups, + true, + )) + } + fn build_test_aggregate( schema: &SchemaRef, stats: Statistics, group_by: PhysicalGroupBy, limit: Option, + ) -> Result { + build_test_aggregate_with_mode( + schema, + stats, + group_by, + limit, + AggregateMode::Final, + ) + } + + fn count_a_aggregate(schema: &SchemaRef) -> Result> { + Ok(Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("a", schema)?]) + .schema(Arc::clone(schema)) + .alias("COUNT(a)") + .build()?, + )) + } + + fn build_test_aggregate_with_mode( + schema: &SchemaRef, + stats: Statistics, + group_by: PhysicalGroupBy, + limit: Option, + mode: AggregateMode, ) -> Result { let input = Arc::new(StatisticsExec::new(stats, (**schema).clone())) as Arc; let mut agg = AggregateExec::try_new( - AggregateMode::Final, + mode, group_by, - vec![Arc::new( - AggregateExprBuilder::new(count_udaf(), vec![col("a", schema)?]) - .schema(Arc::clone(schema)) - .alias("COUNT(a)") - .build()?, - )], + vec![count_a_aggregate(schema)?], vec![None], input, Arc::clone(schema), @@ -4504,7 +6588,7 @@ mod tests { let agg = build_test_aggregate(&schema, input_stats, group_by, case.limit_options)?; - let stats = agg.partition_statistics(None)?; + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; assert_eq!( stats.num_rows, case.expected_num_rows, "FAILED: '{}' — expected {:?}, got {:?}", @@ -4543,7 +6627,7 @@ mod tests { None, )?; - let stats = agg.partition_statistics(None)?; + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; assert_eq!( stats.column_statistics[0].distinct_count, Precision::Exact(100), @@ -4597,7 +6681,7 @@ mod tests { let agg = build_test_aggregate(&schema, input_stats, grouping_set, None)?; - let stats = agg.partition_statistics(None)?; + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; // Per-set NDV: (a,NULL)=100, (NULL,b)=50, (a,b)=100*50=5000 // Total = 100 + 50 + 5000 = 5150 assert_eq!( @@ -4609,6 +6693,75 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_aggregate_stats_duplicate_empty_grouping_sets() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + let duplicate_empty_grouping_sets = + PhysicalGroupBy::new(vec![], vec![], vec![vec![], vec![]], true); + + let single_input = + Arc::new(EmptyExec::new(Arc::clone(&schema))) as Arc; + let single_agg = AggregateExec::try_new( + AggregateMode::Single, + duplicate_empty_grouping_sets.clone(), + vec![count_a_aggregate(&schema)?], + vec![None], + single_input, + Arc::clone(&schema), + )?; + assert_eq!( + StatisticsContext::new() + .compute(&single_agg, &StatisticsArgs::new())? + .num_rows, + Precision::Exact(2) + ); + + let partial_input = + Arc::new(EmptyExec::new(Arc::clone(&schema)).with_partitions(2)) + as Arc; + let partial_agg = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + duplicate_empty_grouping_sets, + vec![count_a_aggregate(&schema)?], + vec![None], + partial_input, + Arc::clone(&schema), + )?); + + assert_eq!( + partial_agg + .properties() + .output_partitioning() + .partition_count(), + 2 + ); + let task_ctx = Arc::new(TaskContext::default()); + for partition in 0..2 { + assert_eq!( + StatisticsContext::new() + .compute( + partial_agg.as_ref(), + &StatisticsArgs::new().with_partition(Some(partition)), + )? + .num_rows, + Precision::Exact(2) + ); + let result = + collect(partial_agg.execute(partition, Arc::clone(&task_ctx))?).await?; + assert_eq!(result.iter().map(RecordBatch::num_rows).sum::(), 2); + } + + assert_eq!( + StatisticsContext::new() + .compute(partial_agg.as_ref(), &StatisticsArgs::new())? + .num_rows, + Precision::Exact(4) + ); + + Ok(()) + } + #[test] fn test_aggregate_stats_non_column_expr_bails_out() -> Result<()> { use datafusion_common::ColumnStatistics; @@ -4646,7 +6799,7 @@ mod tests { PhysicalGroupBy::new_single(vec![(expr_a_plus_b, "a+b".to_string())]); let agg = build_test_aggregate(&schema, input_stats, group_by, None)?; - let stats = agg.partition_statistics(None)?; + let stats = StatisticsContext::new().compute(&agg, &StatisticsArgs::new())?; assert_eq!( stats.num_rows, Precision::Inexact(1_000_000), @@ -4851,159 +7004,831 @@ mod tests { Arc::clone(&schema), )?); - // Pool must be large enough for accumulation to start but too small for - // sort_memory after clearing. - let task_ctx = new_spill_ctx(1, 500); - let result = collect(aggr.execute(0, Arc::clone(&task_ctx))?).await; + // Pool must be large enough for accumulation to start but too small for + // sort_memory after clearing. + let task_ctx = new_spill_ctx(1, 500); + let result = collect(aggr.execute(0, Arc::clone(&task_ctx))?).await; + + match &result { + Ok(_) => panic!("Expected ResourcesExhausted error but query succeeded"), + Err(e) => { + let root = e.find_root(); + assert!( + matches!(root, DataFusionError::ResourcesExhausted(_)), + "Expected ResourcesExhausted, got: {root}", + ); + } + } + + Ok(()) + } + + /// Tests that PartialReduce mode: + /// 1. Accepts state as input (like Final) + /// 2. Produces state as output (like Partial) + /// 3. Can be followed by a Final stage to get the correct result + /// + /// This simulates a tree-reduce pattern: + /// Partial -> PartialReduce -> Final + async fn evaluate_partial_reduce( + groups: PhysicalGroupBy, + aggregates: Vec>, + partition_1_and_2_batches: [Vec; 2], + ) -> Result> { + let schema = partition_1_and_2_batches + .iter() + .flatten() + .next() + .expect("Must have at least 1 batch") + .schema(); + + let [partition_1, partition_2] = partition_1_and_2_batches; + + // Step 1: Partial aggregation on partition 1 + let input1 = + TestMemoryExec::try_new_exec(&[partition_1], Arc::clone(&schema), None)?; + let partial1 = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + input1, + Arc::clone(&schema), + )?); + + // Step 2: Partial aggregation on partition 2 + let input2 = + TestMemoryExec::try_new_exec(&[partition_2], Arc::clone(&schema), None)?; + let partial2 = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + groups.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + input2, + Arc::clone(&schema), + )?); + + // Collect partial results + let task_ctx = Arc::new(TaskContext::default()); + let partial_result1 = + crate::collect(Arc::clone(&partial1) as _, Arc::clone(&task_ctx)).await?; + let partial_result2 = + crate::collect(Arc::clone(&partial2) as _, Arc::clone(&task_ctx)).await?; + + // The partial results have state schema (group cols + accumulator state) + let partial_schema = partial1.schema(); + + // Step 3: PartialReduce — combine partial results, still producing state + let combined_input = TestMemoryExec::try_new_exec( + &[partial_result1, partial_result2], + Arc::clone(&partial_schema), + None, + )?; + // Coalesce into a single partition for the PartialReduce + let coalesced = Arc::new(CoalescePartitionsExec::new(combined_input)); + + let partial_reduce = Arc::new(AggregateExec::try_new( + AggregateMode::PartialReduce, + groups.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + coalesced, + Arc::clone(&partial_schema), + )?); + + // Verify PartialReduce output schema matches Partial output schema + // (both produce state, not final values) + assert_eq!(partial_reduce.schema(), partial_schema); + + // Collect PartialReduce results + let reduce_result = + crate::collect(Arc::clone(&partial_reduce) as _, Arc::clone(&task_ctx)) + .await?; + + // Step 4: Final aggregation on the PartialReduce output + let final_input = TestMemoryExec::try_new_exec( + &[reduce_result], + Arc::clone(&partial_schema), + None, + )?; + let final_agg = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + groups.clone(), + aggregates.clone(), + vec![None; aggregates.len()], + final_input, + Arc::clone(&partial_schema), + )?); + + let result = crate::collect(final_agg, Arc::clone(&task_ctx)).await?; + + Ok(result) + } + + /// Builds the shared `Partial -> PartialReduce -> Final` fixture used by + /// the `test_partial_reduce_*` tests below and runs the pipeline against + /// the aggregate produced by `build_aggregates`. + /// + /// Each test only needs to supply the UDAF/alias under test, so the test + /// body stays focused on which aggregate shape is being exercised. + async fn run_partial_reduce_pipeline( + build_aggregates: F, + ) -> Result> + where + F: FnOnce(&Arc) -> Result>>, + { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::Float64, false), + ])); + + // Two partitions of input data so the Partial stage produces multiple + // partial states that PartialReduce must combine. + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 3])), + Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0])), + ], + )?; + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![1, 2, 3])), + Arc::new(Float64Array::from(vec![40.0, 50.0, 60.0])), + ], + )?; + + let groups = + PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); + let aggregates = build_aggregates(&schema)?; + + evaluate_partial_reduce(groups, aggregates, [vec![batch1], vec![batch2]]).await + } + + // ------------------------------------------------------------------- + // PartialReduce regression coverage. + // + // Each shape (single state field / single input arg, multi-state / + // single-input, more-state-than-input) is covered twice: + // * once against a real UDAF, to round-trip an actual aggregate end + // to end through `Partial -> PartialReduce -> Final`; and + // * once against [`InputTypeAssertingUdaf`], whose input / state / + // output types are deliberately pairwise-disjoint within each test + // so a regression that swapped state-field types for input-field + // types (or vice versa) fails the assertion instead of slipping + // through on a coincidental type match. + // + // The stub variants do the heavy lifting on the contract; the real + // ones make sure no real aggregate is broken by it. + // ------------------------------------------------------------------- + + /// Real-UDAF round-trip: aggregate with a single state field and a + /// single input argument (`SUM(b)` — state and input are both `Float64`). + #[tokio::test] + async fn test_partial_reduce_with_single_state_field_and_single_input_arg() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + Ok(vec![Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("SUM(b)") + .build()?, + )]) + }) + .await?; + + // Expected: group 1 -> 10+40=50, group 2 -> 20+50=70, group 3 -> 30+60=90 + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+--------+ + | a | SUM(b) | + +---+--------+ + | 1 | 50.0 | + | 2 | 70.0 | + | 3 | 90.0 | + +---+--------+ + "); + + Ok(()) + } + + /// Real-UDAF round-trip: aggregate with multiple state fields and a + /// single input argument (`AVG(b)` — state is `[sum: Float64, count: + /// UInt64]`). + #[tokio::test] + async fn test_partial_reduce_with_multiple_state_fields_and_single_input_arg() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + Ok(vec![Arc::new( + AggregateExprBuilder::new(avg_udaf(), vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("AVG(b)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+--------+ + | a | AVG(b) | + +---+--------+ + | 1 | 25.0 | + | 2 | 35.0 | + | 3 | 45.0 | + +---+--------+ + "); + + Ok(()) + } + + /// Real-UDAF round-trip: aggregate whose state has more fields than the + /// input has arguments (`approx_percentile_cont` carries a t-digest). + #[tokio::test] + async fn test_partial_reduce_with_more_state_fields_than_input_args() -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + Ok(vec![Arc::new( + AggregateExprBuilder::new( + approx_percentile_cont_udaf(), + vec![col("b", schema)?, lit(0.75f32)], + ) + .schema(Arc::clone(schema)) + .alias("approx_percentile_cont(b, 0.75)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+---------------------------------+ + | a | approx_percentile_cont(b, 0.75) | + +---+---------------------------------+ + | 1 | 40.0 | + | 2 | 50.0 | + | 3 | 60.0 | + +---+---------------------------------+ + "); + + Ok(()) + } + + /// Stub variant of + /// [`test_partial_reduce_with_single_state_field_and_single_input_arg`] + /// with disjoint input / state / output types. + /// + /// - input: `Float64` + /// - state: `Int32` + /// - output: `Int64` + /// + /// Any mode that accidentally forwarded state-field types in place of + /// input-field types would fail the assertion in + /// [`InputTypeAssertingUdaf`] instead of being masked by a coincidental + /// type match. + #[tokio::test] + async fn test_partial_reduce_with_single_state_field_and_single_input_arg_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64], + vec![DataType::Int32], + DataType::Int64, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b)") + .build()?, + )]) + }) + .await?; + + // Pipeline completing without error is the real assertion. The + // snapshot guards against silent regressions in the row shape. + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+-------------------------+ + | a | input_type_asserting(b) | + +---+-------------------------+ + | 1 | 0 | + | 2 | 0 | + | 3 | 0 | + +---+-------------------------+ + "); + + Ok(()) + } + + /// Stub variant of + /// [`test_partial_reduce_with_multiple_state_fields_and_single_input_arg`] + /// with disjoint input / state / output types. + /// + /// - input: `Float64` + /// - state: `[Int32, Utf8]` + /// - output: `Int64` + #[tokio::test] + async fn test_partial_reduce_with_multiple_state_fields_and_single_input_arg_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64], + vec![DataType::Int32, DataType::Utf8], + DataType::Int64, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new(udaf, vec![col("b", schema)?]) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+-------------------------+ + | a | input_type_asserting(b) | + +---+-------------------------+ + | 1 | 0 | + | 2 | 0 | + | 3 | 0 | + +---+-------------------------+ + "); + + Ok(()) + } + + /// Stub variant of + /// [`test_partial_reduce_with_more_state_fields_than_input_args`] with + /// disjoint input / state / output types — and with multiple input + /// arguments to exercise the multi-arg path explicitly. + /// + /// - input: `[Float64, Date32]` + /// - state: `[Int32, Utf8, Boolean]` + /// - output: `Int64` + #[tokio::test] + async fn test_partial_reduce_with_more_state_fields_than_input_args_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64, DataType::Date32], + vec![DataType::Int32, DataType::Utf8, DataType::Boolean], + DataType::Int64, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new( + udaf, + vec![col("b", schema)?, lit(ScalarValue::Date32(Some(1)))], + ) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b, lit)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+------------------------------+ + | a | input_type_asserting(b, lit) | + +---+------------------------------+ + | 1 | 0 | + | 2 | 0 | + | 3 | 0 | + +---+------------------------------+ + "); + + Ok(()) + } + + /// Stub test: many input args, few state fields (5 inputs / 2 state). + /// + /// All eight types involved are pairwise-disjoint: + /// - input: `[Float64, Date32, UInt16, Boolean, Int32]` + /// - state: `[Utf8, Int64]` + /// - output: `Float32` + #[tokio::test] + async fn test_partial_reduce_with_5_input_args_and_2_state_fields_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![ + DataType::Float64, + DataType::Date32, + DataType::UInt16, + DataType::Boolean, + DataType::Int32, + ], + vec![DataType::Utf8, DataType::Int64], + DataType::Float32, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new( + udaf, + vec![ + col("b", schema)?, + lit(ScalarValue::Date32(Some(1))), + lit(ScalarValue::UInt16(Some(1))), + lit(ScalarValue::Boolean(Some(false))), + lit(ScalarValue::Int32(Some(1))), + ], + ) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b, l1, l2, l3, l4)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+-----------------------------------------+ + | a | input_type_asserting(b, l1, l2, l3, l4) | + +---+-----------------------------------------+ + | 1 | 0.0 | + | 2 | 0.0 | + | 3 | 0.0 | + +---+-----------------------------------------+ + "); - match &result { - Ok(_) => panic!("Expected ResourcesExhausted error but query succeeded"), - Err(e) => { - let root = e.find_root(); - assert!( - matches!(root, DataFusionError::ResourcesExhausted(_)), - "Expected ResourcesExhausted, got: {root}", - ); - let msg = root.to_string(); - assert!( - msg.contains("Failed to reserve memory for sort during spill"), - "Expected sort reservation error, got: {msg}", - ); + Ok(()) + } + + /// Stub test: few input args, many state fields (2 inputs / 5 state). + /// + /// All eight types involved are pairwise-disjoint: + /// - input: `[Float64, Date32]` + /// - state: `[Boolean, Int32, Utf8, Int64, UInt16]` + /// - output: `Float32` + #[tokio::test] + async fn test_partial_reduce_with_2_input_args_and_5_state_fields_using_unique_types() + -> Result<()> { + let result = run_partial_reduce_pipeline(|schema| { + let udaf = Arc::new(AggregateUDF::from(InputTypeAssertingUdaf::new( + vec![DataType::Float64, DataType::Date32], + vec![ + DataType::Boolean, + DataType::Int32, + DataType::Utf8, + DataType::Int64, + DataType::UInt16, + ], + DataType::Float32, + ))); + Ok(vec![Arc::new( + AggregateExprBuilder::new( + udaf, + vec![col("b", schema)?, lit(ScalarValue::Date32(Some(1)))], + ) + .schema(Arc::clone(schema)) + .alias("input_type_asserting(b, lit)") + .build()?, + )]) + }) + .await?; + + assert_snapshot!(batches_to_sort_string(&result), @r" + +---+------------------------------+ + | a | input_type_asserting(b, lit) | + +---+------------------------------+ + | 1 | 0.0 | + | 2 | 0.0 | + | 3 | 0.0 | + +---+------------------------------+ + "); + + Ok(()) + } + + /// Test-only aggregate whose `return_type`, `state_fields`, and + /// `accumulator` hooks all assert that they receive the originally- + /// declared input types; the companion accumulator further asserts + /// `update_batch` sees inputs and `merge_batch` sees state. + /// + /// Each test instantiates it with input / state / output types that + /// are pairwise-disjoint, so a regression that forwarded the wrong + /// types fails on type mismatch rather than passing by accident. + #[derive(Debug, PartialEq, Eq, Hash)] + struct InputTypeAssertingUdaf { + signature: Signature, + input_types: Vec, + state_types: Vec, + output_type: DataType, + } + + fn assert_data_types( + what: &str, + expected: &[DataType], + actual: &[DataType], + ) -> Result<()> { + if actual != expected { + return internal_err!( + "InputTypeAssertingUdaf: {} expected types {:?} but got {:?} — a regression is leaking the wrong types into the accumulator contract", + what, + expected, + actual + ); + } + Ok(()) + } + + /// Produce a zeroed [`ScalarValue`] for `dt`. Only the data types the + /// tests above plug into [`InputTypeAssertingUdaf`] are listed; adding + /// a new type to a test requires extending this match. + fn zero_scalar_for(dt: &DataType) -> Result { + match dt { + DataType::Boolean => Ok(ScalarValue::Boolean(Some(false))), + DataType::Int32 => Ok(ScalarValue::Int32(Some(0))), + DataType::Int64 => Ok(ScalarValue::Int64(Some(0))), + DataType::UInt16 => Ok(ScalarValue::UInt16(Some(0))), + DataType::Float32 => Ok(ScalarValue::Float32(Some(0.0))), + DataType::Utf8 => Ok(ScalarValue::Utf8(Some(String::new()))), + other => internal_err!( + "InputTypeAssertingUdaf: no zero ScalarValue registered for {other:?} \ + — extend `zero_scalar_for` when adding a new state/output type" + ), + } + } + + impl InputTypeAssertingUdaf { + fn new( + input_types: Vec, + state_types: Vec, + output_type: DataType, + ) -> Self { + // Within-test type-disjointness is enforced by construction so + // a future test author can't quietly reintroduce overlap. + assert!( + all_pairwise_distinct(&input_types, &state_types, &output_type), + "InputTypeAssertingUdaf::new: input ({input_types:?}), state \ + ({state_types:?}), and output ({output_type:?}) types must be \ + pairwise-disjoint to avoid accidental passes", + ); + Self { + signature: Signature::exact(input_types.clone(), Volatility::Immutable), + input_types, + state_types, + output_type, } } + } - Ok(()) + /// True iff every type in `inputs ∪ states ∪ {output}` is unique. + fn all_pairwise_distinct( + inputs: &[DataType], + states: &[DataType], + output: &DataType, + ) -> bool { + let mut seen = HashSet::new(); + for dt in inputs + .iter() + .chain(states.iter()) + .chain(std::iter::once(output)) + { + if !seen.insert(dt) { + return false; + } + } + true } - /// Tests that PartialReduce mode: - /// 1. Accepts state as input (like Final) - /// 2. Produces state as output (like Partial) - /// 3. Can be followed by a Final stage to get the correct result + impl AggregateUDFImpl for InputTypeAssertingUdaf { + fn name(&self) -> &str { + "input_type_asserting" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + assert_data_types("return_type(arg_types)", &self.input_types, arg_types)?; + Ok(self.output_type.clone()) + } + + fn state_fields(&self, args: StateFieldsArgs) -> Result> { + let actual: Vec = args + .input_fields + .iter() + .map(|f| f.data_type().clone()) + .collect(); + assert_data_types( + "state_fields(args.input_fields)", + &self.input_types, + &actual, + )?; + Ok(self + .state_types + .iter() + .enumerate() + .map(|(i, dt)| { + Field::new(format!("{}[s{i}]", args.name), dt.clone(), true).into() + }) + .collect()) + } + + fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { + let actual: Vec = acc_args + .expr_fields + .iter() + .map(|f| f.data_type().clone()) + .collect(); + assert_data_types( + "accumulator(acc_args.expr_fields)", + &self.input_types, + &actual, + )?; + Ok(Box::new(InputTypeAssertingAccumulator { + input_types: self.input_types.clone(), + state_types: self.state_types.clone(), + output_type: self.output_type.clone(), + })) + } + } + + /// Companion accumulator for [`InputTypeAssertingUdaf`]. /// - /// This simulates a tree-reduce pattern: - /// Partial -> PartialReduce -> Final - #[tokio::test] - async fn test_partial_reduce_mode() -> Result<()> { - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::UInt32, false), - Field::new("b", DataType::Float64, false), - ])); + /// - `update_batch` must always receive arrays of the original input + /// types. + /// - `merge_batch` must always receive arrays of the declared state + /// types. + /// + /// Anything else means a non-input mode is calling the wrong path. + #[derive(Debug)] + struct InputTypeAssertingAccumulator { + input_types: Vec, + state_types: Vec, + output_type: DataType, + } - // Produce two partitions of input data - let batch1 = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(UInt32Array::from(vec![1, 2, 3])), - Arc::new(Float64Array::from(vec![10.0, 20.0, 30.0])), - ], - )?; - let batch2 = RecordBatch::try_new( - Arc::clone(&schema), - vec![ - Arc::new(UInt32Array::from(vec![1, 2, 3])), - Arc::new(Float64Array::from(vec![40.0, 50.0, 60.0])), - ], - )?; + impl Accumulator for InputTypeAssertingAccumulator { + fn update_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + let actual: Vec = + values.iter().map(|a| a.data_type().clone()).collect(); + assert_data_types("update_batch(values)", &self.input_types, &actual) + } - let groups = - PhysicalGroupBy::new_single(vec![(col("a", &schema)?, "a".to_string())]); - let aggregates: Vec> = vec![Arc::new( - AggregateExprBuilder::new(sum_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("SUM(b)") - .build()?, - )]; + fn evaluate(&mut self) -> Result { + zero_scalar_for(&self.output_type) + } - // Step 1: Partial aggregation on partition 1 - let input1 = - TestMemoryExec::try_new_exec(&[vec![batch1]], Arc::clone(&schema), None)?; - let partial1 = Arc::new(AggregateExec::try_new( - AggregateMode::Partial, - groups.clone(), - aggregates.clone(), - vec![None], - input1, - Arc::clone(&schema), - )?); + fn size(&self) -> usize { + size_of_val(self) + } - // Step 2: Partial aggregation on partition 2 - let input2 = - TestMemoryExec::try_new_exec(&[vec![batch2]], Arc::clone(&schema), None)?; - let partial2 = Arc::new(AggregateExec::try_new( - AggregateMode::Partial, - groups.clone(), - aggregates.clone(), - vec![None], - input2, - Arc::clone(&schema), - )?); + fn state(&mut self) -> Result> { + self.state_types.iter().map(zero_scalar_for).collect() + } - // Collect partial results - let task_ctx = Arc::new(TaskContext::default()); - let partial_result1 = - crate::collect(Arc::clone(&partial1) as _, Arc::clone(&task_ctx)).await?; - let partial_result2 = - crate::collect(Arc::clone(&partial2) as _, Arc::clone(&task_ctx)).await?; + fn merge_batch(&mut self, states: &[ArrayRef]) -> Result<()> { + let actual: Vec = + states.iter().map(|a| a.data_type().clone()).collect(); + assert_data_types("merge_batch(states)", &self.state_types, &actual) + } + } - // The partial results have state schema (group cols + accumulator state) - let partial_schema = partial1.schema(); + #[derive(Debug, PartialEq, Eq, Hash)] + struct NoFirstEmitUdaf { + signature: Signature, + } - // Step 3: PartialReduce — combine partial results, still producing state - let combined_input = TestMemoryExec::try_new_exec( - &[partial_result1, partial_result2], - Arc::clone(&partial_schema), - None, - )?; - // Coalesce into a single partition for the PartialReduce - let coalesced = Arc::new(CoalescePartitionsExec::new(combined_input)); + impl NoFirstEmitUdaf { + fn new() -> Self { + Self { + signature: Signature::exact(vec![DataType::Int32], Volatility::Immutable), + } + } + } - let partial_reduce = Arc::new(AggregateExec::try_new( - AggregateMode::PartialReduce, - groups.clone(), - aggregates.clone(), - vec![None], - coalesced, - Arc::clone(&partial_schema), - )?); + impl AggregateUDFImpl for NoFirstEmitUdaf { + fn name(&self) -> &str { + "no_first_emit" + } - // Verify PartialReduce output schema matches Partial output schema - // (both produce state, not final values) - assert_eq!(partial_reduce.schema(), partial_schema); + fn signature(&self) -> &Signature { + &self.signature + } - // Collect PartialReduce results - let reduce_result = - crate::collect(Arc::clone(&partial_reduce) as _, Arc::clone(&task_ctx)) - .await?; + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } - // Step 4: Final aggregation on the PartialReduce output - let final_input = TestMemoryExec::try_new_exec( - &[reduce_result], - Arc::clone(&partial_schema), - None, - )?; - let final_agg = Arc::new(AggregateExec::try_new( - AggregateMode::Final, - groups.clone(), - aggregates.clone(), - vec![None], - final_input, - Arc::clone(&partial_schema), - )?); + fn state_fields(&self, args: StateFieldsArgs) -> Result> { + Ok(vec![Arc::new(Field::new( + format!("{}[count]", args.name), + DataType::Int64, + false, + ))]) + } - let result = crate::collect(final_agg, Arc::clone(&task_ctx)).await?; + fn accumulator( + &self, + _acc_args: AccumulatorArgs, + ) -> Result> { + Ok(Box::new(NoFirstEmitAccumulator)) + } - // Expected: group 1 -> 10+40=50, group 2 -> 20+50=70, group 3 -> 30+60=90 - assert_snapshot!(batches_to_sort_string(&result), @r" - +---+--------+ - | a | SUM(b) | - +---+--------+ - | 1 | 50.0 | - | 2 | 70.0 | - | 3 | 90.0 | - +---+--------+ - "); + fn groups_accumulator_supported(&self, _args: AccumulatorArgs) -> bool { + true + } - Ok(()) + fn create_groups_accumulator( + &self, + _args: AccumulatorArgs, + ) -> Result> { + Ok(Box::new(NoFirstEmitGroupsAccumulator { counts: vec![] })) + } + } + + #[derive(Debug)] + struct NoFirstEmitAccumulator; + + impl Accumulator for NoFirstEmitAccumulator { + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(Some(0))) + } + + fn size(&self) -> usize { + size_of_val(self) + } + + fn state(&mut self) -> Result> { + Ok(vec![ScalarValue::Int64(Some(0))]) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } + } + + #[derive(Debug)] + struct NoFirstEmitGroupsAccumulator { + counts: Vec, + } + + impl NoFirstEmitGroupsAccumulator { + fn emit_counts(&mut self, emit_to: EmitTo) -> Result { + match emit_to { + EmitTo::All => { + let counts = std::mem::take(&mut self.counts); + Ok(Arc::new(Int64Array::from(counts))) + } + EmitTo::First(_) => internal_err!( + "partial grouped aggregate output must materialize with EmitTo::All before slicing" + ), + } + } + } + + impl GroupsAccumulator for NoFirstEmitGroupsAccumulator { + fn update_batch( + &mut self, + _values: &[ArrayRef], + group_indices: &[usize], + _opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + self.counts.resize(total_num_groups, 0); + for group_index in group_indices { + self.counts[*group_index] += 1; + } + Ok(()) + } + + fn evaluate(&mut self, emit_to: EmitTo) -> Result { + self.emit_counts(emit_to) + } + + fn state(&mut self, emit_to: EmitTo) -> Result> { + Ok(vec![self.emit_counts(emit_to)?]) + } + + fn convert_to_state( + &self, + values: &[ArrayRef], + opt_filter: Option<&BooleanArray>, + ) -> Result> { + assert_eq!(values.len(), 1, "one argument to convert_to_state"); + let counts = match opt_filter { + Some(filter) => filter + .iter() + .map(|value| i64::from(value.unwrap_or(false))) + .collect::>(), + None => vec![1; values[0].len()], + }; + Ok(vec![Arc::new(Int64Array::from(counts))]) + } + + fn merge_batch( + &mut self, + _values: &[ArrayRef], + _group_indices: &[usize], + _total_num_groups: usize, + ) -> Result<()> { + Ok(()) + } + + fn size(&self) -> usize { + size_of_val(self) + self.counts.capacity() * size_of::() + } } /// Test that [`AggregateExec::with_dynamic_filter_expr`] overrides the existing dynamic filter @@ -5034,11 +7859,14 @@ mod tests { lit(false), )); let agg = agg.with_dynamic_filter_expr(Arc::clone(&new_df))?; + let produced = agg.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + assert_eq!(produced[0].expression_id(), new_df.expression_id()); // The aggregate's filter should now resolve to the new inner expression. - let swapped = agg - .dynamic_filter_expr() - .expect("should still have dynamic filter") + let swapped = produced[0] + .downcast_ref::() + .expect("produced expression should be a DynamicFilterPhysicalExpr") .current()?; assert_eq!(format!("{swapped}"), format!("{}", lit(false))); @@ -5059,6 +7887,38 @@ mod tests { Ok(()) } + #[test] + fn test_plan_contains_expression_id_recurses_plans_and_expressions() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let empty: Arc = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![col("a", &schema)?], + lit(true), + )); + let expression_id = dynamic_filter + .expression_id() + .expect("dynamic filters always have an expression ID"); + + assert!(!plan_contains_expression_id(&empty, expression_id)?); + + let dynamic_filter_expr: Arc = + Arc::::clone(&dynamic_filter); + let predicate: Arc = + Arc::new(NotExpr::new(dynamic_filter_expr)); + let filter: Arc = + Arc::new(FilterExecBuilder::new(predicate, empty).build()?); + let projection: Arc = Arc::new(ProjectionExec::try_new( + [ProjectionExpr::new_from_expression( + col("a", &schema)?, + &schema, + )?], + filter, + )?); + + assert!(plan_contains_expression_id(&projection, expression_id)?); + Ok(()) + } + /// Test that [`AggregateExec::with_dynamic_filter_expr`] errors when the aggregate does not support dynamic filtering #[test] fn test_with_dynamic_filter_error_unsupported() -> Result<()> { @@ -5082,7 +7942,7 @@ mod tests { child, Arc::clone(&schema), )?; - assert!(agg.dynamic_filter_expr().is_none()); + assert!(agg.dynamic_expressions_produced().is_empty()); let df = Arc::new(DynamicFilterPhysicalExpr::new( vec![col("a", &schema)?], diff --git a/datafusion/physical-plan/src/aggregates/order/full.rs b/datafusion/physical-plan/src/aggregates/order/full.rs index eb98611f79dfb..ca818d6a2d598 100644 --- a/datafusion/physical-plan/src/aggregates/order/full.rs +++ b/datafusion/physical-plan/src/aggregates/order/full.rs @@ -115,6 +115,11 @@ impl GroupOrderingFull { self.state = State::Complete; } + /// Starts tracking a new fully ordered input segment. + pub fn reset(&mut self) { + self.state = State::Start; + } + /// Called when new groups are added in a batch. See documentation /// on [`super::GroupOrdering::new_groups`] pub fn new_groups(&mut self, total_num_groups: usize) { diff --git a/datafusion/physical-plan/src/aggregates/order/mod.rs b/datafusion/physical-plan/src/aggregates/order/mod.rs index 97fbd519c825c..259411b00b697 100644 --- a/datafusion/physical-plan/src/aggregates/order/mod.rs +++ b/datafusion/physical-plan/src/aggregates/order/mod.rs @@ -93,6 +93,20 @@ impl GroupOrdering { } } + /// Resets the ordering state while preserving the configured ordering mode. + /// + /// Ordered partial aggregation uses this after passing intermediate states + /// downstream, and ordered final aggregation uses it after spilling a run. + /// In both cases the hash table is empty and can start tracking the next + /// input batch from a fresh ordering state. + pub fn reset(&mut self) { + match self { + GroupOrdering::None => {} + GroupOrdering::Partial(partial) => partial.reset(), + GroupOrdering::Full(full) => full.reset(), + } + } + /// Removes the first `n` groups from the internal state, shifting all /// existing indexes down by `n`. pub fn remove_groups(&mut self, n: usize) { diff --git a/datafusion/physical-plan/src/aggregates/order/partial.rs b/datafusion/physical-plan/src/aggregates/order/partial.rs index 476551a7ca210..1603bb6d079be 100644 --- a/datafusion/physical-plan/src/aggregates/order/partial.rs +++ b/datafusion/physical-plan/src/aggregates/order/partial.rs @@ -186,6 +186,12 @@ impl GroupOrderingPartial { }; } + /// Starts tracking a new ordered input segment with the same sort-key + /// columns. + pub fn reset(&mut self) { + self.state = State::Start; + } + fn updated_sort_key( current_sort: usize, sort_key: Option>, diff --git a/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs new file mode 100644 index 0000000000000..19deedc258c46 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/ordered_final_stream.rs @@ -0,0 +1,902 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Final aggregate stream for ordered partial-state input. + +use std::ops::ControlFlow; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{FinalMarker, OrderedAggregateTable}; +use super::group_values::GroupByMetrics; +use crate::aggregates::AggregateMode; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::SpillManager; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Final aggregate stream for `InputOrderMode::Sorted` and +/// `InputOrderMode::PartiallySorted`. +/// +/// See comments at [`super::ordered_partial_stream::OrderedPartialAggregateStream`] for details. +/// +/// # Spilling +/// +/// This section is only for implementation notes, for background, see [`super::ordered_partial_stream::OrderedPartialAggregateStream`] +/// +/// For partially sorted input, spilling works as follows: +/// +/// - Reserve the table footprint plus one `u32` sort index per buffered group. The +/// extra index array is used in later sorting before spilling. +/// - On memory pressure, materialize all group states into one batch. +/// - Use [`IncrementalSortIterator`] to compute the full-batch index, then +/// materialize and write one sorted `batch_size` slice at a time. The original +/// batch and full index remain live until the run is written. +/// - After input ends, merge the sorted runs and replay them through a fully +/// ordered final aggregate stream. +pub(crate) struct OrderedFinalAggregateStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + state: Option, +} + +/// Spill configuration and accumulated runs for partially ordered final +/// aggregation. +/// +/// Each spill event drains all currently buffered groups, sorts their intermediate +/// states by the full group key, and writes them to one spill file. All files are +/// merged and replayed after the original input ends. +struct OrderedFinalSpillContext { + /// Aggregate configuration + agg: AggregateExec, + /// Task context + context: Arc, + /// Original partition index + partition: usize, + /// Target batch size from configuration + batch_size: usize, + /// Full group-key ordering, such ordering with be kept in: a) individual spill + /// files, b) order after final merging and streaming aggregate + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Fully sorted spill runs waiting to be merged. + spills: Vec, +} + +/// See comments at `poll_next()` for details. +enum OrderedFinalAggregateState { + ReadingInput { + table: OrderedAggregateTable, + /// None if either + /// - Disk Manager doesn't enable temporary file creation + /// - The group keys are fully ordered, it's expected to use bounded memory + spill_context: Option>, + }, + Spilling { + table: OrderedAggregateTable, + spill_context: Box, + }, + ProducingOutput { + table: OrderedAggregateTable, + }, + PreparingMergeInput { + table: OrderedAggregateTable, + spill_context: Box, + }, + MergingSpills { + stream: SendableRecordBatchStream, + }, + Done, +} + +type OrderedFinalAggregatePoll = Poll>>; +type OrderedFinalAggregateStateTransition = ControlFlow< + (OrderedFinalAggregatePoll, OrderedFinalAggregateState), + OrderedFinalAggregateState, +>; + +impl OrderedFinalSpillContext { + fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + batch_size: usize, + input_order_mode: &InputOrderMode, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result { + let group_schema = agg.group_by.group_schema(spill_schema)?; + let output_ordering = agg.cache.output_ordering(); + let InputOrderMode::PartiallySorted(order_indices) = input_order_mode else { + return internal_err!("Ordered final spill requires partially ordered input"); + }; + let spill_indices = order_indices.iter().copied().chain( + (0..group_schema.fields().len()).filter(|idx| !order_indices.contains(idx)), + ); + let spill_sort_exprs = spill_indices.map(|idx| { + let field = group_schema.field(idx); + let output_expr = Column::new(field.name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Ordered final spill expression is empty"); + }; + + let spill_manager = SpillManager::new( + context.runtime_env(), + spill_metrics, + Arc::clone(spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + Ok(Self { + agg: agg.clone(), + context: Arc::clone(context), + partition, + batch_size, + spill_expr, + spill_manager, + spills: vec![], + }) + } + + fn has_spills(&self) -> bool { + !self.spills.is_empty() + } + + /// Sorts and spills the aggregated groups. Memory reservation should be updated + /// by the caller. + /// + /// Individual spill files are ordered by the `group by` keys. + /// + /// See [`OrderedFinalAggregateStream`] for spilling details. + fn spill_table( + &mut self, + table: &mut OrderedAggregateTable, + ) -> Result<()> { + let Some(batch) = table.take_state_batch()? else { + return Ok(()); + }; + + let sorted_iter = + IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "OrderedFinalAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Ordered final aggregation produced an empty spill"); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// Merges every sorted run and finalizes it through the fully ordered path. + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result { + let Self { + agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) + } +} + +impl OrderedFinalAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert!(matches!( + agg.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + )); + debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + + let input = agg.input.execute(partition, Arc::clone(context))?; + Self::new_with_input(agg, context, partition, input, &agg.input_order_mode) + } + + pub(in crate::aggregates) fn new_with_input( + agg: &AggregateExec, + context: &Arc, + partition: usize, + input: SendableRecordBatchStream, + input_order_mode: &InputOrderMode, + ) -> Result { + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let group_by_metrics = GroupByMetrics::new(&agg.metrics, partition); + let spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let reservation = + MemoryConsumer::new(format!("OrderedFinalAggregateStream[{partition}]")) + // HACK: Technically, fully ordered aggregate is a non-spillable + // consumer, since it uses bounded memory. There is a known race + // condition bug, and we set it to spillable to let it have larger + // memory budget to suppress the bug. + // Bug issue: https://github.com/apache/datafusion/issues/17334 + .with_can_spill(true) + .register(context.memory_pool()); + Self::new_with_input_and_metrics( + agg, + context, + partition, + input, + input_order_mode, + baseline_metrics, + group_by_metrics, + Some(spill_metrics), + reservation, + ) + } + + #[expect( + clippy::too_many_arguments, + reason = "keeps replay metric reuse explicit" + )] + /// Builds the stream with the reservation of its logical aggregate operator. + /// Replay callers pass a sibling of the reservation used by the merge input, + /// keeping both components under one memory-consumer registration. + pub(in crate::aggregates) fn new_with_input_and_metrics( + agg: &AggregateExec, + context: &Arc, + partition: usize, + input: SendableRecordBatchStream, + input_order_mode: &InputOrderMode, + baseline_metrics: BaselineMetrics, + group_by_metrics: GroupByMetrics, + spill_metrics: Option, + reservation: MemoryReservation, + ) -> Result { + debug_assert!(matches!( + agg.mode, + AggregateMode::Final | AggregateMode::FinalPartitioned + )); + debug_assert_ne!(*input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input_schema = input.schema(); + let batch_size = context.session_config().batch_size(); + + let can_spill = matches!(input_order_mode, InputOrderMode::PartiallySorted(_)) + && context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + let Some(spill_metrics) = spill_metrics else { + return internal_err!("Spillable ordered final stream requires metrics"); + }; + Some(Box::new(OrderedFinalSpillContext::new( + agg, + context, + partition, + batch_size, + input_order_mode, + &input_schema, + spill_metrics, + )?)) + } else { + None + }; + + let table = OrderedAggregateTable::::new_with_input_order( + agg, + &input_schema, + Arc::clone(&schema), + batch_size, + input_order_mode, + group_by_metrics, + )?; + Ok(Self { + schema, + input, + reservation, + baseline_metrics, + state: Some(OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_internal_err(message: &str) -> OrderedFinalAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(internal_err!("{message}"))), + OrderedFinalAggregateState::Done, + )) + } + + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + table: &OrderedAggregateTable, + spill_context: Option<&OrderedFinalSpillContext>, + ) -> usize { + let table_size = table.memory_size(); + if spill_context.is_some() { + // See `OrderedFinalAggregateStream` comments for how is it estimated + table_size.saturating_add(table.num_groups().saturating_mul(size_of::())) + } else { + table_size + } + } + + /// Consumes one ordered partial-state input batch, then immediately emits + /// finalized groups if the ordering proves any group is ready. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::ReadingInput { + mut table, + spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected ReadingInput state", + ); + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, + )), + Poll::Ready(Some(Ok(batch))) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = table.aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, + )); + } + + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + self.reservation + .try_resize(Self::reservation_size_for_table( + &table, + spill_context.as_deref(), + )); + timer.done(); + match resize_result { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + let Some(spill_context) = spill_context else { + // `None` means spilling is not supported, see comments + // at `OrderedFinalAggregateState` for details. + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + }; + if table.is_empty() { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + return ControlFlow::Continue( + OrderedFinalAggregateState::Spilling { + table, + spill_context, + }, + ); + } + Err(e) => { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + } + + let result = if spill_context + .as_ref() + .is_some_and(|spill_context| spill_context.has_spills()) + { + // Once one incomplete run is spilled, every remaining state + // must participate in replay so no group is finalized twice. + Ok(None) + } else { + let timer = elapsed_compute.timer(); + let result = table.next_output_batch(); + timer.done(); + result + }; + + match result { + // Some finalized groups can be emitted. Yield them, then + // continue aggregating input in the current state. + Ok(Some(batch)) => { + if let Err(e) = + self.reservation + .try_resize(Self::reservation_size_for_table( + &table, + spill_context.as_deref(), + )) + { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + let next_state = OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok( + batch.record_output(&self.baseline_metrics) + ))), + next_state, + )) + } + // Can't do early emit, continue aggregating. + Ok(None) => { + ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, + )), + } + } + Poll::Ready(Some(Err(e))) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ReadingInput { + table, + spill_context, + }, + )), + Poll::Ready(None) => { + self.close_input(); + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue( + OrderedFinalAggregateState::PreparingMergeInput { + table, + spill_context, + }, + ) + } + _ => { + table.input_done(); + ControlFlow::Continue( + OrderedFinalAggregateState::ProducingOutput { table }, + ) + } + } + } + } + } + + /// Sorts and spills one complete in-memory state run, then resumes input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_spilling( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::Spilling { + mut table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected Spilling state", + ); + }; + + // Sanity check: it's impossible to OOM when the table is empty + if table.is_empty() { + return ControlFlow::Break(( + Poll::Ready(Some(internal_err!( + "Ordered final aggregation entered Spilling with an empty table" + ))), + OrderedFinalAggregateState::Done, + )); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let mut result = spill_context.spill_table(&mut table); + + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + result = + Err(e.context("Decreasing allocation after spilling should succeed")); + } + + timer.done(); + + match result { + // Finished spilling the aggregate table, continue aggregating from input + Ok(()) => ControlFlow::Continue(OrderedFinalAggregateState::ReadingInput { + table, + spill_context: Some(spill_context), + }), + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )), + } + } + + /// 1. Spills the last in-memory run. + /// 2. Constructs a globally ordered input stream by applying a sort-preserving + /// merge to all spills. + /// 3. Constructs a replay stream: an ordered aggregate stream over the fully + /// ordered input constructed from the spills. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_preparing_merge_input( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::PreparingMergeInput { + mut table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected PreparingMergeInput state", + ); + }; + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let replay = match spill_context.spill_table(&mut table) { + Ok(()) => { + let group_by_metrics = table.group_by_metrics(); + drop(table); + match self.reservation.try_resize(0) { + Ok(()) => (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_metrics, + self.reservation.new_empty(), + ), + Err(e) => Err(e), + } + } + Err(e) => Err(e), + }; + timer.done(); + + match replay { + Ok(stream) => { + ControlFlow::Continue(OrderedFinalAggregateState::MergingSpills { + stream, + }) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )), + } + } + + /// Forwards output from the fully ordered stream that consumes the merged + /// spill runs. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_merging_spills( + &mut self, + cx: &mut Context<'_>, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::MergingSpills { mut stream } = original_state + else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected MergingSpills state", + ); + }; + + match stream.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + OrderedFinalAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( + Poll::Ready(Some(Ok(batch))), + OrderedFinalAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Err(e))) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )), + Poll::Ready(None) => ControlFlow::Continue(OrderedFinalAggregateState::Done), + } + } + + /// Emits one batch after input is exhausted. + /// + /// `table.input_done()` has already made every remaining group safe to emit, + /// so this state keeps draining until the table is empty. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + original_state: OrderedFinalAggregateState, + ) -> OrderedFinalAggregateStateTransition { + let OrderedFinalAggregateState::ProducingOutput { table } = original_state else { + return Self::break_with_internal_err( + "Ordered final aggregate stream expected ProducingOutput state", + ); + }; + + let mut table = table; + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = table.next_output_batch(); + timer.done(); + + match result { + Ok(Some(batch)) => { + let next_state = if table.is_empty() { + drop(table); + if let Err(e) = self.reservation.try_resize(0) { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::Done, + )); + } + OrderedFinalAggregateState::Done + } else { + if let Err(e) = self.reservation.try_resize(table.memory_size()) { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ProducingOutput { table }, + )); + } + OrderedFinalAggregateState::ProducingOutput { table } + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Err(e) => ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + OrderedFinalAggregateState::ProducingOutput { table }, + )), + Ok(None) => { + drop(table); + let next_state = OrderedFinalAggregateState::Done; + if let Err(e) = self.reservation.try_resize(0) { + return ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)); + } + ControlFlow::Continue(next_state) + } + } + } +} + +impl Stream for OrderedFinalAggregateStream { + type Item = Result; + + /// Entry point for the ordered final aggregate state machine. + /// + /// See comments in [`OrderedFinalAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered partial-state input and merging + /// those states into the ordered final aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Merge one input batch. If it fits in memory, optionally yield groups + /// proven complete by the input ordering, then read the next batch. + /// -> Spilling + /// The table cannot reserve enough memory. Move all current states into + /// one fully group-key-sorted spill run. + /// -> ProducingOutput + /// Input was exhausted without spilling. Mark every remaining group as + /// complete and produce its final result. + /// -> PreparingMergeInput + /// Input was exhausted after spilling. Spill the last in-memory run and + /// construct the ordered input used to merge all spill files. + /// + /// Spilling + /// -> ReadingInput + /// One sorted run was written; resume reading the original input. + /// + /// PreparingMergeInput + /// Spill the final in-memory run and build the input ordered replay stream. + /// -> MergingSpills + /// The final run was spilled and the ordered replay stream was built. + /// + /// MergingSpills + /// Aggregate the merged spill runs and emit final results. + /// -> MergingSpills + /// Forward one result batch from the fully ordered replay stream that + /// consumes the sort-preserving merge. + /// -> Done + /// The merged spill input was fully aggregated. + /// + /// ProducingOutput + /// -> ProducingOutput + /// One remaining final aggregate batch was yielded; repeat to continue + /// draining the table. + /// -> Done + /// All remaining groups were emitted. + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + let cur_state = self + .state + .take() + .expect("OrderedFinalAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ OrderedFinalAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ OrderedFinalAggregateState::Spilling { .. } => { + self.handle_spilling(state) + } + state @ OrderedFinalAggregateState::PreparingMergeInput { .. } => { + self.handle_preparing_merge_input(state) + } + state @ OrderedFinalAggregateState::MergingSpills { .. } => { + self.handle_merging_spills(cx, state) + } + state @ OrderedFinalAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ OrderedFinalAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + // Errors are terminal: discard all operator state and release + // its upstream input and memory reservation before returning. + drop(next_state); + self.close_input(); + self.reservation.free(); + self.state = Some(OrderedFinalAggregateState::Done); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for OrderedFinalAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} diff --git a/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs new file mode 100644 index 0000000000000..9e93a111a6466 --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/ordered_partial_stream.rs @@ -0,0 +1,352 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Partial aggregate stream for ordered group input. + +use std::sync::Arc; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result}; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{TaskContext, TryEmitter, async_try_stream}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{OrderedAggregateTable, PartialMarker}; +use crate::aggregates::AggregateMode; +use crate::aggregates::order::GroupOrdering; +use crate::metrics::{BaselineMetrics, MetricBuilder, SpillMetrics}; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; +use crate::{InputOrderMode, SendableRecordBatchStream, metrics}; + +/// Partial aggregate stream for `InputOrderMode::Sorted` and +/// `InputOrderMode::PartiallySorted`. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// If the input is ordered by `k`, the aggregate can use ordered partial and +/// final stages: +/// +/// ## Plan +/// AggregateExec(stage=final, ordered) +/// -- RepartitionExec(hash(k), preserves_order=true) +/// ---- AggregateExec(stage=partial, ordered) +/// +/// ## Partial Stage Behavior +/// Input: raw rows +/// Output: partial states for all groups (for example, `AVG(x)` emits `SUM(x)` +/// and `COUNT(x)`) +/// +/// ## Final Stage Behavior +/// Input: partial states +/// Output: results for all groups (for example, `AVG(x)` calculated from the +/// state) +/// +/// # Order-based Optimization +/// +/// For the aggregation work, the hash aggregation implementation is reused. +/// +/// After each input batch, check whether any groups can be emitted eagerly to +/// improve memory efficiency. For example, if the last group key seen is +/// `k = 100`, it is safe to emit all groups with keys less than 100 because the +/// input is ordered. +/// +/// # Memory Pressure and Spilling +/// +/// ## Fully ordered case +/// +/// If the input is ordered by every group key, for example: +/// +/// - Input order: `a, b` +/// - `GROUP BY`: `a, b` +/// +/// Completed groups can be emitted as soon as the next group is observed. Thus, +/// only the current group remains active after completed groups are emitted, and +/// memory usage does not grow with the total number of groups. +/// +/// If a memory reservation nevertheless fails, the stream returns the error +/// directly, indicating an unexpected behavior. +/// +/// ## Partially ordered case +/// +/// If the input is ordered by only a subset of the group keys, for example: +/// +/// - Input order: `a` +/// - `GROUP BY`: `a, b` +/// +/// If one `a` value contains many distinct `b` values, the table may accumulate +/// enough groups to exceed the memory limit. +/// +/// - `OrderedPartialAggregateStream`: On reservation failure, it emits all current +/// intermediate states downstream and resets the table. The final stage can +/// merge repeated `(a, b)` state rows, so no disk spill is required. +/// - `OrderedFinalAggregateStream`: It cannot emit incomplete final results. On +/// reservation failure, it sorts the current intermediate states by the complete +/// group key and spills them as one run. After the input ends, it spills any +/// remaining states, performs a sort-preserving merge of all runs, and feeds the +/// merged input into a fully ordered final aggregate stream. +/// +/// ## Implementation Note +/// +/// This is intentionally kept simple and closely maps to +/// `GroupedHashAggregateStream` to finish the refactor sooner. +/// +/// See issue for details: +/// +pub(crate) struct OrderedPartialAggregateStream { + schema: SchemaRef, + input: SendableRecordBatchStream, + reservation: MemoryReservation, + baseline_metrics: BaselineMetrics, + reduction_factor: metrics::RatioMetrics, + table: Option>, +} + +impl OrderedPartialAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert_eq!(agg.mode, AggregateMode::Partial); + debug_assert_ne!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let reduction_factor = MetricBuilder::new(&agg.metrics) + .with_type(metrics::MetricType::Summary) + .ratio_metrics("reduction_factor", partition); + + let table = OrderedAggregateTable::::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + let reservation = + MemoryConsumer::new(format!("OrderedPartialAggregateStream[{partition}]")) + .with_can_spill(matches!( + table.group_ordering(), + GroupOrdering::Partial(_) + )) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + reservation, + baseline_metrics, + reduction_factor, + table: Some(table), + }) + } + + pub(crate) fn into_stream(self) -> SendableRecordBatchStream { + let schema_clone = Arc::clone(&self.schema); + + let cloned_metrics = self.baseline_metrics.clone(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), + )); + + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + } + + /// Entry point for the ordered partial aggregate state machine. + /// + /// See comments in [`OrderedPartialAggregateStream`] for high-level ideas. + /// + /// State transitions are implemented using the generator pattern; see the comments in [`async_try_stream`]. + /// + /// Conceptual state-transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling ordered input and aggregating batches + /// into the ordered partial aggregate table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one input batch. If the ordering proves some groups are + /// complete, yield one partial-state batch immediately, then continue + /// reading input. Otherwise continue directly with the next input batch. + /// -> DrainingFinal + /// Input was exhausted. Mark the table input as done so every remaining + /// group is safe to emit. + /// + /// DrainingFinal + /// -> DrainingFinal + /// One remaining partial-state batch was yielded; repeat to continue + /// draining the table. + /// -> Done + /// All remaining groups were emitted. + /// + /// Done + /// -> (end) + /// ``` + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|mut emitter| async move { + let mut table = self + .table + .take() + .expect("OrderedPartialAggregateStream state should not be None"); + + self.handle_reading_input(&mut table, &mut emitter).await?; + + // Input has exhausted, move to the final draining stage. + self.close_input(); + table.input_done(); + + self.handle_draining_final(table, &mut emitter).await?; + + Ok(()) + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + /// Consumes one ordered input batch, then immediately emits completed groups + /// if the ordering proves any group is ready. + /// + /// See comments at [`Self::create_stream`] for details. + async fn handle_reading_input( + &mut self, + table: &mut OrderedAggregateTable, + emitter: &mut TryEmitter, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + + while let Some(batch) = self.input.next().await.transpose()? { + let input_rows = batch.num_rows(); + self.reduction_factor.add_total(input_rows); + + let timer = elapsed_compute.timer(); + + table.aggregate_batch(&batch)?; + + // Check memory reservation. See function comments for details. + if let Some(batch) = self.resize_or_take_state_batch(table)? { + self.reduction_factor.add_part(batch.num_rows()); + drop(timer); + emitter.emit(batch).await; + continue; + } + + let Some(batch) = table.next_output_batch()? else { + // Can't do early emit, continue aggregating. + continue; + }; + + self.reduction_factor.add_part(batch.num_rows()); + self.reservation.try_resize(table.memory_size())?; + + drop(timer); + emitter.emit(batch).await; + } + + Ok(()) + } + + /// Update the memory reservation, and: + /// - If memory reservation succeed, returns `Ok(None)` + /// - If memory reservation failed, + /// - If input is partially ordered, materialize all the output, and + /// directly send them to the final aggregation stage. + /// Returns `Ok(Some(batch))` + /// - If input is fully ordered, directly return error. It's not + /// expected to use more than constant memory. + /// Returns `Err(..)` + /// + /// # Implementation Note + /// Incrementally output it after the blocked state management is ready, keep + /// it simple for now. + /// + /// Issue: + fn resize_or_take_state_batch( + &mut self, + table: &mut OrderedAggregateTable, + ) -> Result> { + let oom = match self.reservation.try_resize(table.memory_size()) { + Ok(()) => return Ok(None), + Err(e @ DataFusionError::ResourcesExhausted(_)) => e, + Err(e) => return Err(e), + }; + + if matches!(table.group_ordering(), GroupOrdering::Full(_)) { + return Err(oom); + } + + let Some(batch) = table.take_state_batch()? else { + return Err(oom); + }; + self.reservation.try_resize(table.memory_size())?; + Ok(Some(batch)) + } + + /// Emits one batch after input is exhausted. + /// + /// `table.input_done()` has already made every remaining group safe to emit, + /// so this state keeps draining until the table is empty. + /// + /// See comments at [`Self::create_stream`] for details. + /// + async fn handle_draining_final( + &mut self, + mut table: OrderedAggregateTable, + emitter: &mut TryEmitter, + ) -> Result<()> { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + while let Some(batch) = table.next_output_batch()? { + self.reduction_factor.add_part(batch.num_rows()); + + if table.is_empty() { + // Clear memory before emitting last batch so we don't have to wait for next poll to clear + drop(table); + let _ = self.reservation.try_resize(0); + drop(timer); + + emitter.emit(batch).await; + + return Ok(()); + } + + self.reservation.try_resize(table.memory_size())?; + + timer.done(); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + // was empty + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs new file mode 100644 index 0000000000000..2f4535e66f4ef --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/partial_reduce_stream.rs @@ -0,0 +1,385 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Partial-reduce hash aggregation stream implementation. +//! +//! This stream is part of the incremental migration from +//! [`crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! See issue for details: + +use std::ops::ControlFlow; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::Result; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::stream::{Stream, StreamExt}; + +use super::AggregateExec; +use super::aggregate_hash_table::{AggregateHashTable, PartialReduceMarker}; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Hash aggregation can combine multiple partial stages before final +/// evaluation. This stream implements the partial-reduce stage. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=final) +/// -- RepartitionExec(hash(k)) +/// ---- AggregateExec(stage=partial_reduce) +/// ------ RepartitionExec(hash(k)) +/// -------- AggregateExec(stage=partial) +/// +/// Note: the example plan is only intended to demonstrate this stream's semantics; +/// the default DataFusion SQL planner does not produce plans in this shape. +/// +/// This stream implements the middle partial-reduce aggregation in the plan above. +/// +/// The motivation is to reduce shuffling traffic in a distributed setting. See +/// +/// +/// ## Partial-Reduce Stage Behavior +/// Input: partial aggregate state rows +/// Output: merged partial aggregate state rows +/// +/// This stage is useful for tree-reduce plans. It consumes the same schema as +/// a final aggregate stage, but emits the same schema as a partial aggregate +/// stage. +pub(crate) struct PartialReduceHashAggregateStream { + /// Output schema: group columns followed by partial aggregate state columns. + schema: SchemaRef, + + /// Input batches containing partial aggregate state rows. + input: SendableRecordBatchStream, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// Memory reservation for group keys and accumulators. + reservation: MemoryReservation, + + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// state for emitting output batches. + state: Option, +} + +/// States for partial-reduce hash aggregation processing. +// The typestate pattern mirrors the final stream and keeps the input/output +// semantics explicit for this mode. +enum PartialReduceHashAggregateState { + ReadingInput { + hash_table: AggregateHashTable, + }, + ProducingOutput { + hash_table: AggregateHashTable, + }, + Done, +} + +type PartialReduceHashAggregatePoll = Poll>>; +type PartialReduceHashAggregateStateTransition = ControlFlow< + ( + PartialReduceHashAggregatePoll, + PartialReduceHashAggregateState, + ), + PartialReduceHashAggregateState, +>; + +impl PartialReduceHashAggregateState { + fn hash_table(&self) -> &AggregateHashTable { + match self { + Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { + hash_table + } + Self::Done => unreachable!("Done state does not hold a hash table"), + } + } + + fn hash_table_mut(&mut self) -> &mut AggregateHashTable { + match self { + Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { + hash_table + } + Self::Done => unreachable!("Done state does not hold a hash table"), + } + } + + fn into_hash_table(self) -> AggregateHashTable { + match self { + Self::ReadingInput { hash_table } | Self::ProducingOutput { hash_table } => { + hash_table + } + Self::Done => unreachable!("Done state does not hold a hash table"), + } + } + + fn into_producing_output(self) -> Self { + Self::ProducingOutput { + hash_table: self.into_hash_table(), + } + } + + fn into_done(self) -> Self { + Self::Done + } +} + +impl PartialReduceHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert_eq!(agg.mode, super::AggregateMode::PartialReduce); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + + // Preserve the existing aggregate metric surface for this plan node. + let _spill_metrics = SpillMetrics::new(&agg.metrics, partition); + + let hash_table = AggregateHashTable::::new( + agg, + partition, + Arc::clone(&schema), + batch_size, + )?; + + let reservation = + MemoryConsumer::new(format!("PartialReduceHashAggregateStream[{partition}]")) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + state: Some(PartialReduceHashAggregateState::ReadingInput { hash_table }), + }) + } + + fn start_output( + &mut self, + hash_table: &mut AggregateHashTable, + ) -> Result<()> { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + hash_table.start_output() + } + + /// Handle ReadingInput state - aggregate partial state batches into the hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + mut original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + debug_assert!(matches!( + &original_state, + PartialReduceHashAggregateState::ReadingInput { .. } + )); + debug_assert!(original_state.hash_table().is_building()); + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break((Poll::Pending, original_state)), + // Get a new input batch, aggregate it in the hash table + Poll::Ready(Some(Ok(batch))) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = original_state.hash_table_mut().aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + original_state, + )); + } + + if let Err(e) = self + .reservation + .try_resize(original_state.hash_table().memory_size()) + { + return ControlFlow::Break(( + Poll::Ready(Some(Err(e))), + original_state, + )); + } + + ControlFlow::Continue(original_state) + } + Poll::Ready(Some(Err(e))) => { + ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + } + // Input ends, move to output state + Poll::Ready(None) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = self.start_output(original_state.hash_table_mut()); + timer.done(); + + match result { + Ok(()) => { + ControlFlow::Continue(original_state.into_producing_output()) + } + Err(e) => { + ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)) + } + } + } + } + } + + /// Handle ProducingOutput state - emit merged partial aggregate state batches. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + mut original_state: PartialReduceHashAggregateState, + ) -> PartialReduceHashAggregateStateTransition { + debug_assert!(matches!( + &original_state, + PartialReduceHashAggregateState::ProducingOutput { .. } + )); + debug_assert!(!original_state.hash_table().is_building()); + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = original_state.hash_table_mut().next_output_batch(); + timer.done(); + + match result { + Ok(Some(batch)) => { + let _ = self + .reservation + .try_resize(original_state.hash_table().memory_size()); + debug_assert!(batch.num_rows() > 0); + let next_state = if original_state.hash_table().is_done() { + original_state.into_done() + } else { + original_state + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Ok(None) => { + let _ = self.reservation.try_resize(0); + ControlFlow::Continue(original_state.into_done()) + } + Err(e) => ControlFlow::Break((Poll::Ready(Some(Err(e))), original_state)), + } + } +} + +impl Stream for PartialReduceHashAggregateStream { + type Item = Result; + + /// Entry point for the partial-reduce hash aggregate state machine. + /// + /// See comments in [`PartialReduceHashAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling partial-state input and merging those + /// states into the partial-reduce hash table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one partial-state input batch, update the inner aggregate + /// hash table, and continue with the next input batch. + /// + /// -> ProducingOutput + /// Input was exhausted. Move to the next state to start outputting + /// merged partial aggregate states. + /// + /// ProducingOutput + /// -> ProducingOutput + /// One merged partial-state output batch was yielded; repeat to + /// continue producing output incrementally. + /// + /// -> Done + /// All merged partial-state output was emitted. + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + let cur_state = self + .state + .take() + .expect("PartialReduceHashAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ PartialReduceHashAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ PartialReduceHashAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ PartialReduceHashAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for PartialReduceHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} diff --git a/datafusion/physical-plan/src/aggregates/single_stream.rs b/datafusion/physical-plan/src/aggregates/single_stream.rs new file mode 100644 index 0000000000000..c6f25dc2cf28b --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/single_stream.rs @@ -0,0 +1,833 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Single-stage hash aggregation stream implementation. +//! +//! This stream is part of the incremental migration from +//! [`crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream`]. +//! +//! See issue for details: + +use std::ops::ControlFlow; +use std::sync::Arc; +use std::task::{Context, Poll}; + +use arrow::datatypes::SchemaRef; +use arrow::record_batch::RecordBatch; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_execution::TaskContext; +use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_physical_expr::PhysicalSortExpr; +use datafusion_physical_expr::expressions::Column; +use datafusion_physical_expr_common::sort_expr::LexOrdering; +use futures::stream::{Stream, StreamExt}; + +use super::aggregate_hash_table::{AggregateHashTable, SingleMarker}; +use super::group_values::GroupByMetrics; +use super::ordered_final_stream::OrderedFinalAggregateStream; +use super::{AggregateExec, create_schema}; +use crate::aggregates::AggregateMode; +use crate::metrics::{BaselineMetrics, RecordOutput, SpillMetrics}; +use crate::sorts::IncrementalSortIterator; +use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; +use crate::spill::spill_manager::SpillManager; +use crate::stream::EmptyRecordBatchStream; +use crate::{InputOrderMode, RecordBatchStream, SendableRecordBatchStream}; + +/// Hash aggregation can run the full logical aggregation in one operator. This +/// stream implements the single stage for grouped hash aggregation. +/// +/// This aggregation variant is useful when: +/// - There is only one partition (config `target_partitions` is set to 1) +/// - When input is already partitioned (`t` is backed by Parquet files, that is range/hash +/// partitioned on the group keys), the single aggregation mode is the most efficient +/// approach to use. +/// +/// # Example +/// +/// SELECT k, AVG(v) FROM t GROUP BY k; +/// +/// ## Plan +/// AggregateExec(stage=single) +/// -- DataSourceExec(t) +/// +/// ## Single Stage Behavior +/// Input: raw rows +/// Output: final aggregate values for all groups (for example, `AVG(x)`) +/// +/// This stream implements the complete aggregation without a partial/final +/// split. It consumes raw input rows and emits final aggregate values. +/// +/// # Spilling +/// +/// During aggregation, group keys and states accumulate. If memory usage exceeds +/// the budget, spilling is triggered as follows: +/// 1. After aggregating a new input batch, if the memory reservation exceeds its +/// limit, spill all accumulated groups and states. +/// - Sort all groups by the group keys before spilling. +/// 2. Repeat until the input is exhausted. +/// 3. Perform a sort-preserving merge of all spill files and feed the merged output +/// into an ordered streaming aggregation, which ensures bounded memory usage and +/// evaluates the final result. +/// - [`OrderedFinalAggregateStream`] is reused for the streaming aggregation. +pub(crate) struct SingleHashAggregateStream { + /// Output schema: group columns followed by final aggregate value columns. + schema: SchemaRef, + + /// Input batches containing raw rows, not partial aggregate state. + input: SendableRecordBatchStream, + + /// Execution metrics shared with the aggregate plan node. + baseline_metrics: BaselineMetrics, + + /// Memory reservation for group keys, accumulators, and spill sorting. + reservation: MemoryReservation, + + /// Tracks the high-level stream lifecycle. The hash table owns the lower-level + /// state for emitting output batches. + state: Option, +} + +/// Spill configuration and accumulated runs for single hash aggregation. +/// +/// Each spill event drains all currently buffered groups, sorts their intermediate +/// states by the full group key, and writes them to one spill file. All files are +/// merged and replayed after the original input ends. +struct SingleSpillContext { + /// Aggregate configuration used to construct the final replay stream. + /// + /// Spilled rows already contain evaluated group keys and intermediate + /// aggregate states. Replay must therefore use final aggregation semantics + /// and column-based group expressions rather than evaluating the raw input + /// expressions a second time. After the spill files are merged into ordered + /// input, this configuration is used to construct an + /// [`OrderedFinalAggregateStream`], and perform the final evaluation step. + final_agg: AggregateExec, + /// Task context. + context: Arc, + /// Original partition index. + partition: usize, + /// Target batch size from configuration. + batch_size: usize, + /// Full group-key ordering kept by every spill file and the merged input. + spill_expr: LexOrdering, + /// Spill I/O and metrics manager. + spill_manager: SpillManager, + /// Spill runs waiting to be merged, they're all sorted by full group-by keys. + spills: Vec, +} + +/// See comments at `poll_next()` for details. +enum SingleHashAggregateState { + ReadingInput { + hash_table: AggregateHashTable, + spill_context: Option>, + }, + Spilling { + hash_table: AggregateHashTable, + spill_context: Box, + }, + ProducingOutput { + hash_table: AggregateHashTable, + }, + PreparingMergeInput { + hash_table: AggregateHashTable, + spill_context: Box, + }, + MergingSpills { + stream: SendableRecordBatchStream, + }, + Done, + /// Sentinel state to use when returning error from any other states, because: + /// - It explicitly releases state-owned resources immediately + /// - More defensive against accidentally resuming execution after error + Error, +} + +type SingleHashAggregatePoll = Poll>>; +type SingleHashAggregateStateTransition = ControlFlow< + (SingleHashAggregatePoll, SingleHashAggregateState), + SingleHashAggregateState, +>; + +impl SingleSpillContext { + fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + batch_size: usize, + spill_schema: &SchemaRef, + spill_metrics: SpillMetrics, + ) -> Result { + let group_schema = agg.group_by.group_schema(&agg.input().schema())?; + let output_ordering = agg.cache.output_ordering(); + let spill_sort_exprs = + group_schema + .fields() + .iter() + .enumerate() + .map(|(idx, field)| { + let output_expr = Column::new(field.name(), idx); + let sort_options = output_ordering + .and_then(|ordering| ordering.get_sort_options(&output_expr)) + .unwrap_or_default(); + PhysicalSortExpr::new(Arc::new(output_expr), sort_options) + }); + let Some(spill_expr) = LexOrdering::new(spill_sort_exprs) else { + return internal_err!("Single hash aggregate spill expression is empty"); + }; + + let spill_manager = SpillManager::new( + context.runtime_env(), + spill_metrics, + Arc::clone(spill_schema), + ) + .with_compression_type(context.session_config().spill_compression()); + + // See `SingleSpillContext::final_agg` comments for `final_agg`'s usage + let mut final_agg = agg.clone(); + final_agg.mode = match agg.mode { + AggregateMode::Single => AggregateMode::Final, + AggregateMode::SinglePartitioned => AggregateMode::FinalPartitioned, + mode => { + return internal_err!( + "Single hash aggregate spill cannot replay aggregate mode {mode:?}" + ); + } + }; + final_agg.group_by = Arc::new(agg.group_by.as_final()); + final_agg.input_order_mode = InputOrderMode::Sorted; + + Ok(Self { + final_agg, + context: Arc::clone(context), + partition, + batch_size, + spill_expr, + spill_manager, + spills: vec![], + }) + } + + fn has_spills(&self) -> bool { + !self.spills.is_empty() + } + + /// Sorts and spills the aggregated groups. Memory reservation should be updated + /// by the caller. + /// + /// Individual spill files are ordered by the `group by` keys. + /// + /// See [`SingleHashAggregateStream`] for spilling details. + fn spill_table( + &mut self, + hash_table: &mut AggregateHashTable, + ) -> Result<()> { + let Some(batch) = hash_table.take_state_batch()? else { + return Ok(()); + }; + + let sorted_iter = + IncrementalSortIterator::new(batch, self.spill_expr.clone(), self.batch_size); + let spill_file = self + .spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + sorted_iter, + "SingleHashAggregateSpill", + )?; + + let Some((file, max_record_batch_memory)) = spill_file else { + return internal_err!("Single hash aggregation produced an empty spill"); + }; + + self.spills.push(SortedSpillFile { + file, + max_record_batch_memory, + }); + + Ok(()) + } + + /// Merges every sorted run, and do the aggregate evaluation with + /// [`OrderedFinalAggregateStream`] + fn into_replay_stream( + self, + baseline_metrics: &BaselineMetrics, + group_by_metrics: GroupByMetrics, + reservation: MemoryReservation, + ) -> Result { + let Self { + final_agg, + context, + partition, + batch_size, + spill_expr, + spill_manager, + spills, + } = self; + + let spill_schema = Arc::clone(spill_manager.schema()); + // The merge and replay table are two components of the same aggregate + // operator. Keep them under one consumer registration so a fair memory + // pool does not divide this operator's quota between its own phases. + let merge_reservation = reservation.new_empty(); + let merged = StreamingMergeBuilder::new() + .with_schema(spill_schema) + .with_spill_manager(spill_manager) + .with_sorted_spill_files(spills) + .with_expressions(&spill_expr) + .with_metrics(baseline_metrics.intermediate()) + .with_batch_size(batch_size) + .with_reservation(merge_reservation) + .build()?; + let replay = OrderedFinalAggregateStream::new_with_input_and_metrics( + &final_agg, + &context, + partition, + merged, + &InputOrderMode::Sorted, + baseline_metrics.clone(), + group_by_metrics, + None, + reservation, + )?; + Ok(Box::pin(replay)) + } +} + +impl SingleHashAggregateStream { + pub fn new( + agg: &AggregateExec, + context: &Arc, + partition: usize, + ) -> Result { + debug_assert!(matches!( + agg.mode, + AggregateMode::Single | AggregateMode::SinglePartitioned + )); + debug_assert_eq!(agg.input_order_mode, InputOrderMode::Linear); + + let schema = Arc::clone(&agg.schema); + let input = agg.input.execute(partition, Arc::clone(context))?; + let input_schema = input.schema(); + let batch_size = context.session_config().batch_size(); + let baseline_metrics = BaselineMetrics::new(&agg.metrics, partition); + let spill_metrics = SpillMetrics::new(&agg.metrics, partition); + let state_schema = Arc::new(create_schema( + input_schema.as_ref(), + &agg.group_by, + &agg.aggr_expr, + AggregateMode::Partial, + )?); + + let hash_table = AggregateHashTable::::new( + agg, + partition, + Arc::clone(&schema), + Arc::clone(&state_schema), + batch_size, + )?; + + let can_spill = context.runtime_env().disk_manager.tmp_files_enabled(); + let spill_context = if can_spill { + Some(Box::new(SingleSpillContext::new( + agg, + context, + partition, + batch_size, + &state_schema, + spill_metrics, + )?)) + } else { + None + }; + + let reservation = + MemoryConsumer::new(format!("SingleHashAggregateStream[{partition}]")) + .with_can_spill(can_spill) + .register(context.memory_pool()); + + Ok(Self { + schema, + input, + baseline_metrics, + reservation, + state: Some(SingleHashAggregateState::ReadingInput { + hash_table, + spill_context, + }), + }) + } + + fn close_input(&mut self) { + let input_schema = self.input.schema(); + self.input = Box::pin(EmptyRecordBatchStream::new(input_schema)); + } + + fn break_with_err(error: DataFusionError) -> SingleHashAggregateStateTransition { + ControlFlow::Break(( + Poll::Ready(Some(Err(error))), + SingleHashAggregateState::Error, + )) + } + + fn break_with_internal_err(message: &str) -> SingleHashAggregateStateTransition { + Self::break_with_err(internal_datafusion_err!("{message}")) + } + + /// Reserve memory for the current aggregate table. + fn reservation_size_for_table( + hash_table: &AggregateHashTable, + spill_context: Option<&SingleSpillContext>, + ) -> usize { + let table_size = hash_table.memory_size(); + if spill_context.is_some() { + // See `SingleHashAggregateStream` comments for how this is estimated. + table_size.saturating_add( + hash_table + .building_group_count() + .saturating_mul(size_of::()), + ) + } else { + table_size + } + } + + /// Consumes one raw input batch and updates the single-stage hash table. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_reading_input( + &mut self, + cx: &mut Context<'_>, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::ReadingInput { + mut hash_table, + spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Single hash aggregate stream expected ReadingInput state", + ); + }; + + match self.input.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + SingleHashAggregateState::ReadingInput { + hash_table, + spill_context, + }, + )), + Poll::Ready(Some(Ok(batch))) => { + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.aggregate_batch(&batch); + timer.done(); + + if let Err(e) = result { + return Self::break_with_err(e); + } + + // Check memory reservation, and potentially spill. + let timer = elapsed_compute.timer(); + let resize_result = + self.reservation + .try_resize(Self::reservation_size_for_table( + &hash_table, + spill_context.as_deref(), + )); + timer.done(); + match resize_result { + Ok(()) => {} + Err(e @ DataFusionError::ResourcesExhausted(_)) => { + let Some(spill_context) = spill_context else { + return Self::break_with_err(e.context( + "Single hash aggregate cannot spill because temporary files are not enabled in the DiskManager", + )); + }; + if hash_table.building_group_count() == 0 { + return Self::break_with_internal_err( + "Single hash aggregate ran out of memory with no aggregated groups", + ); + } + return ControlFlow::Continue( + SingleHashAggregateState::Spilling { + hash_table, + spill_context, + }, + ); + } + Err(e) => { + return Self::break_with_err(e); + } + } + + ControlFlow::Continue(SingleHashAggregateState::ReadingInput { + hash_table, + spill_context, + }) + } + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => { + self.close_input(); + match spill_context { + Some(spill_context) if spill_context.has_spills() => { + ControlFlow::Continue( + SingleHashAggregateState::PreparingMergeInput { + hash_table, + spill_context, + }, + ) + } + _ => { + let elapsed_compute = + self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.start_output(); + timer.done(); + + match result { + Ok(()) => ControlFlow::Continue( + SingleHashAggregateState::ProducingOutput { hash_table }, + ), + Err(e) => Self::break_with_err(e), + } + } + } + } + } + } + + /// Sorts and spills one complete in-memory state run, then resumes input. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_spilling( + &mut self, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::Spilling { + mut hash_table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Single hash aggregate stream expected Spilling state", + ); + }; + + // Sanity check: it is impossible to OOM when the table is empty. + if hash_table.building_group_count() == 0 { + return Self::break_with_internal_err( + "Single hash aggregation entered Spilling with an empty table", + ); + } + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let mut result = spill_context.spill_table(&mut hash_table); + + // Spilling shrinks the aggregate table and releases its accumulated + // memory. Update the reservation accordingly. + if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) { + result = + Err(e.context("Decreasing allocation after spilling should succeed")); + } + + timer.done(); + + match result { + // Finished spilling the aggregate table, continue aggregating from input. + Ok(()) => ControlFlow::Continue(SingleHashAggregateState::ReadingInput { + hash_table, + spill_context: Some(spill_context), + }), + Err(e) => Self::break_with_err(e), + } + } + + /// 1. Spills the last in-memory run. + /// 2. Constructs a globally ordered input stream by applying a sort-preserving + /// merge to all spills. + /// 3. Constructs a replay stream: an ordered final aggregate stream over the + /// fully ordered input constructed from the spills. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_preparing_merge_input( + &mut self, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::PreparingMergeInput { + mut hash_table, + mut spill_context, + } = original_state + else { + return Self::break_with_internal_err( + "Single hash aggregate stream expected PreparingMergeInput state", + ); + }; + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let replay = match spill_context.spill_table(&mut hash_table) { + Ok(()) => { + let group_by_metrics = hash_table.group_by_metrics().clone(); + drop(hash_table); + match self.reservation.try_resize(0) { + Ok(()) => (*spill_context).into_replay_stream( + &self.baseline_metrics, + group_by_metrics, + self.reservation.new_empty(), + ), + Err(e) => Err(e), + } + } + Err(e) => Err(e), + }; + timer.done(); + + match replay { + Ok(stream) => { + ControlFlow::Continue(SingleHashAggregateState::MergingSpills { stream }) + } + Err(e) => Self::break_with_err(e), + } + } + + /// Forwards output from the fully ordered stream that consumes the merged + /// spill runs. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_merging_spills( + &mut self, + cx: &mut Context<'_>, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::MergingSpills { mut stream } = original_state + else { + return Self::break_with_internal_err( + "Single hash aggregate stream expected MergingSpills state", + ); + }; + + match stream.poll_next_unpin(cx) { + Poll::Pending => ControlFlow::Break(( + Poll::Pending, + SingleHashAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Ok(batch))) => ControlFlow::Break(( + Poll::Ready(Some(Ok(batch))), + SingleHashAggregateState::MergingSpills { stream }, + )), + Poll::Ready(Some(Err(e))) => Self::break_with_err(e), + Poll::Ready(None) => ControlFlow::Continue(SingleHashAggregateState::Done), + } + } + + /// Emits one batch after input is exhausted. + /// + /// See comments at `poll_next()` for details. + /// + /// Returns the next operator state with control flow decision. + fn handle_producing_output( + &mut self, + original_state: SingleHashAggregateState, + ) -> SingleHashAggregateStateTransition { + let SingleHashAggregateState::ProducingOutput { mut hash_table } = original_state + else { + return Self::break_with_internal_err( + "Single hash aggregate stream expected ProducingOutput state", + ); + }; + + let elapsed_compute = self.baseline_metrics.elapsed_compute().clone(); + let timer = elapsed_compute.timer(); + let result = hash_table.next_output_batch(); + timer.done(); + + match result { + Ok(Some(batch)) => { + let next_state = if hash_table.is_done() { + drop(hash_table); + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + SingleHashAggregateState::Done + } else { + if let Err(e) = self.reservation.try_resize(hash_table.memory_size()) + { + return Self::break_with_err(e); + } + SingleHashAggregateState::ProducingOutput { hash_table } + }; + + ControlFlow::Break(( + Poll::Ready(Some(Ok(batch.record_output(&self.baseline_metrics)))), + next_state, + )) + } + Err(e) => Self::break_with_err(e), + Ok(None) => { + drop(hash_table); + let next_state = SingleHashAggregateState::Done; + if let Err(e) = self.reservation.try_resize(0) { + return Self::break_with_err(e); + } + ControlFlow::Continue(next_state) + } + } + } +} + +impl Stream for SingleHashAggregateStream { + type Item = Result; + + /// Entry point for the single hash aggregate state machine. + /// + /// See comments in [`SingleHashAggregateStream`] for high-level ideas. + /// + /// State transition graph: + /// + /// ```text + /// (start) + /// -> ReadingInput + /// The stream starts by polling raw input rows and aggregating those + /// rows into the single-stage hash table. + /// + /// ReadingInput + /// -> ReadingInput + /// Aggregate one raw input batch. If it fits in memory, continue with + /// the next input batch. + /// -> Spilling + /// The table cannot reserve enough memory. Move all current states into + /// one fully group-key-sorted spill run. + /// -> ProducingOutput + /// Input was exhausted without spilling. Start outputting final values. + /// -> PreparingMergeInput + /// Input was exhausted after spilling. Spill the last in-memory run and + /// construct the ordered input used to merge all spill files. + /// + /// Spilling + /// -> ReadingInput + /// One sorted run was written; resume reading the original input. + /// + /// PreparingMergeInput + /// Spill the final in-memory run and build the input ordered replay stream. + /// -> MergingSpills + /// The final run was spilled and the ordered replay stream was built. + /// + /// MergingSpills + /// Aggregate the merged spill runs and emit final results. + /// -> MergingSpills + /// Forward one result batch from the fully ordered replay stream that + /// consumes the sort-preserving merge. + /// -> Done + /// The merged spill input was fully aggregated. + /// + /// ProducingOutput + /// -> ProducingOutput + /// One final output batch was yielded; repeat to continue producing + /// output incrementally. + /// -> Done + /// All final output was emitted. + /// + /// Any active state + /// -> Error + /// An error drops state-owned resources before it is returned. + /// + /// Error + /// -> (end) + /// + /// Done + /// -> (end) + /// ``` + fn poll_next( + mut self: std::pin::Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + loop { + let cur_state = self + .state + .take() + .expect("SingleHashAggregateStream state should not be None"); + + let next_state = match cur_state { + state @ SingleHashAggregateState::ReadingInput { .. } => { + self.handle_reading_input(cx, state) + } + state @ SingleHashAggregateState::Spilling { .. } => { + self.handle_spilling(state) + } + state @ SingleHashAggregateState::PreparingMergeInput { .. } => { + self.handle_preparing_merge_input(state) + } + state @ SingleHashAggregateState::MergingSpills { .. } => { + self.handle_merging_spills(cx, state) + } + state @ SingleHashAggregateState::ProducingOutput { .. } => { + self.handle_producing_output(state) + } + state @ SingleHashAggregateState::Error => { + self.close_input(); + self.reservation.free(); + self.state = Some(state); + return Poll::Ready(None); + } + state @ SingleHashAggregateState::Done => { + let _ = self.reservation.try_resize(0); + self.state = Some(state); + return Poll::Ready(None); + } + }; + + match next_state { + ControlFlow::Continue(next_state) => { + self.state = Some(next_state); + continue; + } + ControlFlow::Break((Poll::Ready(Some(Err(e))), next_state)) => { + debug_assert!(matches!(next_state, SingleHashAggregateState::Error)); + + // The handler has already discarded its state-owned resources. + // Release the remaining stream-owned resources before returning. + self.close_input(); + self.reservation.free(); + self.state = Some(SingleHashAggregateState::Error); + return Poll::Ready(Some(Err(e))); + } + ControlFlow::Break((poll, next_state)) => { + self.state = Some(next_state); + return poll; + } + } + } + } +} + +impl RecordBatchStream for SingleHashAggregateStream { + fn schema(&self) -> SchemaRef { + Arc::clone(&self.schema) + } +} diff --git a/datafusion/physical-plan/src/aggregates/skip_partial.rs b/datafusion/physical-plan/src/aggregates/skip_partial.rs new file mode 100644 index 0000000000000..20e17d2b2790e --- /dev/null +++ b/datafusion/physical-plan/src/aggregates/skip_partial.rs @@ -0,0 +1,305 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::record_batch::RecordBatch; + +use crate::metrics; + +/// Tracks if the aggregate should skip partial aggregations +/// +/// See "partial aggregation" discussion on +/// [`crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream`]. +pub(super) struct SkipAggregationProbe { + // ======================================================================== + // PROPERTIES: + // These fields are initialized at the start and remain constant throughout + // the execution. + // ======================================================================== + /// Aggregation ratio check performed when the number of input rows exceeds + /// this threshold (from `SessionConfig`) + probe_rows_threshold: usize, + /// Maximum ratio of `num_groups` to `input_rows` for continuing aggregation + /// (from `SessionConfig`). If the ratio exceeds this value, aggregation + /// is skipped and input rows are directly converted to output + probe_ratio_threshold: f64, + + // ======================================================================== + // STATES: + // Fields changes during execution. Can be buffer, or state flags that + // influence the execution in parent `GroupedHashAggregateStream` + // ======================================================================== + /// Number of processed input rows (updated during probing) + input_rows: usize, + /// Number of total group values for `input_rows` (updated during probing) + num_groups: usize, + + /// Flag indicating further data aggregation may be skipped (decision made + /// when probing complete) + should_skip: bool, + /// Flag indicating further updates of `SkipAggregationProbe` state won't + /// make any effect (set either while probing or on probing completion) + is_locked: bool, + + // ======================================================================== + // METRICS: + // ======================================================================== + /// Number of rows where state was output without aggregation. + /// + /// * If 0, all input rows were aggregated (should_skip was always false) + /// + /// * if greater than zero, the number of rows which were output directly + /// without aggregation + skipped_aggregation_rows: metrics::Count, +} + +impl SkipAggregationProbe { + pub(super) fn new( + probe_rows_threshold: usize, + probe_ratio_threshold: f64, + skipped_aggregation_rows: metrics::Count, + ) -> Self { + Self { + input_rows: 0, + num_groups: 0, + probe_rows_threshold, + probe_ratio_threshold, + should_skip: false, + is_locked: false, + skipped_aggregation_rows, + } + } + + /// Updates `SkipAggregationProbe` state: + /// - increments the number of input rows + /// - replaces the number of groups with the new value + /// - on `probe_rows_threshold` exceeded calculates + /// aggregation ratio and sets `should_skip` flag + /// - if `should_skip` is set, locks further state updates + pub(super) fn update_state(&mut self, input_rows: usize, num_groups: usize) { + if self.is_locked { + return; + } + self.input_rows += input_rows; + self.num_groups = num_groups; + if self.input_rows >= self.probe_rows_threshold { + self.should_skip = self.num_groups as f64 / self.input_rows as f64 + > self.probe_ratio_threshold; + // Set is_locked to true only if we have decided to skip, otherwise we can try to skip + // during processing the next record_batch. + self.is_locked = self.should_skip; + } + } + + pub(super) fn should_skip(&self) -> bool { + self.should_skip + } + + /// Record the number of rows that were output directly without aggregation + pub(super) fn record_skipped(&mut self, batch: &RecordBatch) { + self.skipped_aggregation_rows.add(batch.num_rows()); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::aggregates::grouped_hash_stream::GroupedHashAggregateStream; + use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; + use crate::execution_plan::ExecutionPlan; + use crate::test::TestMemoryExec; + + use std::sync::Arc; + + use arrow::array::Int32Array; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::Result; + use datafusion_execution::TaskContext; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_functions_aggregate::count::count_udaf; + use datafusion_physical_expr::aggregate::AggregateExprBuilder; + use datafusion_physical_expr::expressions::col; + use futures::StreamExt; + + // Migrated to PartialHashAggregateStream coverage in hash_stream.rs; + // kept here for the legacy GroupedHashAggregateStream implementation. + #[tokio::test] + async fn test_skip_aggregation_probe_not_locked_until_skip() -> Result<()> { + // Test that the probe is not locked until we actually decide to skip. + // This allows us to continue evaluating the skip condition across multiple batches. + // + // Scenario: + // - Batch 1: Hits rows threshold but NOT ratio threshold (low cardinality) -> don't skip + // - Batch 2: Now hits ratio threshold (high cardinality) -> skip + // + // Without the fix, the probe would be locked after batch 1, preventing the skip + // decision from being made on batch 2. + + let schema = Arc::new(Schema::new(vec![ + Field::new("group_col", DataType::Int32, false), + Field::new("value_col", DataType::Int32, false), + ])); + + // Configure thresholds: + // - probe_rows_threshold: 100 rows + // - probe_ratio_threshold: 0.8 (80%) + let probe_rows_threshold = 100; + let probe_ratio_threshold = 0.8; + + // Batch 1: 100 rows with only 10 unique groups + // Ratio: 10/100 = 0.1 (10%) < 0.8 -> should NOT skip + // This will hit the rows threshold but not the ratio threshold + let batch1_rows = 100; + let batch1_groups = 10; + let mut group_ids_batch1 = Vec::new(); + for i in 0..batch1_rows { + group_ids_batch1.push((i % batch1_groups) as i32); + } + let values_batch1: Vec = vec![1; batch1_rows]; + + let batch1 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch1)), + Arc::new(Int32Array::from(values_batch1)), + ], + )?; + + // Batch 2: 360 rows with 360 unique NEW groups (starting from group 10) + // After batch 2, total: 460 rows, 370 groups + // Ratio: 370/460 is about 0.804 (80.4%) > 0.8 -> SHOULD decide to skip + let batch2_rows = 360; + let batch2_groups = 360; + let group_ids_batch2: Vec = (batch1_groups..(batch1_groups + batch2_groups)) + .map(|x| x as i32) + .collect(); + let values_batch2: Vec = vec![1; batch2_rows]; + + let batch2 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch2)), + Arc::new(Int32Array::from(values_batch2)), + ], + )?; + + // Batch 3: This batch should be skipped since we decided to skip after batch 2 + // 100 rows with 100 unique groups (continuing from where batch 2 left off) + let batch3_rows = 100; + let batch3_groups = 100; + let batch3_start_group = batch1_groups + batch2_groups; + let group_ids_batch3: Vec = (batch3_start_group + ..(batch3_start_group + batch3_groups)) + .map(|x| x as i32) + .collect(); + let values_batch3: Vec = vec![1; batch3_rows]; + + let batch3 = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(group_ids_batch3)), + Arc::new(Int32Array::from(values_batch3)), + ], + )?; + + let input_partitions = vec![vec![batch1, batch2, batch3]]; + + let runtime = RuntimeEnvBuilder::default().build_arc()?; + let mut task_ctx = TaskContext::default().with_runtime(runtime); + + // Configure skip aggregation settings + let mut session_config = task_ctx.session_config().clone(); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_rows_threshold", + &datafusion_common::ScalarValue::UInt64(Some(probe_rows_threshold)), + ); + session_config = session_config.set( + "datafusion.execution.skip_partial_aggregation_probe_ratio_threshold", + &datafusion_common::ScalarValue::Float64(Some(probe_ratio_threshold)), + ); + task_ctx = task_ctx.with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // Create aggregate: COUNT(*) GROUP BY group_col + let group_expr = vec![(col("group_col", &schema)?, "group_col".to_string())]; + let aggr_expr = vec![Arc::new( + AggregateExprBuilder::new(count_udaf(), vec![col("value_col", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_value") + .build()?, + )]; + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + + // Use Partial mode + let aggregate_exec = AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(group_expr), + aggr_expr, + vec![None], + exec, + Arc::clone(&schema), + )?; + + // Execute and collect results + let mut stream = + GroupedHashAggregateStream::new(&aggregate_exec, &Arc::clone(&task_ctx), 0)?; + let mut results = Vec::new(); + + while let Some(result) = stream.next().await { + let batch = result?; + results.push(batch); + } + + // Check that skip aggregation actually happened. + // The key metric is skipped_aggregation_rows. + let metrics = aggregate_exec.metrics().unwrap(); + let skipped_rows = metrics + .sum_by_name("skipped_aggregation_rows") + .map(|m| m.as_usize()) + .unwrap_or(0); + + // We expect batch 3's rows to be skipped (100 rows) + assert_eq!( + skipped_rows, batch3_rows, + "Expected batch 3's rows ({batch3_rows}) to be skipped", + ); + + Ok(()) + } + + #[test] + fn test_skip_aggregation_probe_equality_does_not_skip() { + // When num_groups / input_rows == probe_ratio_threshold, the `>` boundary + // means we must NOT skip: equality is not sufficient to trigger skip. + let threshold_ratio = 0.5_f64; + let threshold_rows = 10_usize; + let mut probe = SkipAggregationProbe::new( + threshold_rows, + threshold_ratio, + metrics::Count::new(), + ); + + // 10 rows, 5 groups: ratio = 5/10 = 0.5 exactly equals threshold + probe.update_state(10, 5); + + assert!( + !probe.should_skip(), + "ratio == threshold should not trigger skip (boundary is exclusive)" + ); + } +} diff --git a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs index 694780f08547f..adc8f8c315b32 100644 --- a/datafusion/physical-plan/src/aggregates/topk/hash_table.rs +++ b/datafusion/physical-plan/src/aggregates/topk/hash_table.rs @@ -39,6 +39,11 @@ pub trait KeyType: Clone + Comparable + Debug {} impl KeyType for T where T: Clone + Comparable + Debug {} +/// `heap_idx` assigned to groups whose aggregate values are all NULL. Such +/// groups are tracked in the hash table only (they never enter the heap), so +/// they can be emitted with a NULL aggregate value at the end. +const NULL_HEAP_IDX: usize = usize::MAX; + /// An entry in our hash table that: /// 1. memoizes the hash /// 2. contains the key (ID) @@ -57,10 +62,25 @@ struct TopKHashTable { map: HashTable, // Store the actual items separately to allow for index-based access store: Vec>>, - // Free index in the store for reuse - free_index: Option, + // Free indexes in the store for reuse + free_indices: Vec, // The maximum number of entries allowed limit: usize, + // Number of entries registered as all-NULL (heap_idx == NULL_HEAP_IDX) + null_count: usize, +} + +/// Outcome of [`ArrowHashTable::find_or_insert`], letting the caller keep its +/// own all-NULL group accounting in sync without an extra lookup. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum InsertKind { + /// The group already existed as a valued group + Existing, + /// The group was newly inserted as a valued group + New, + /// The group was registered as all-NULL and has now been converted into a + /// valued group + ReplacedNull, } /// An interface to hide the generic type signature of TopKHashTable behind arrow arrays @@ -70,7 +90,20 @@ pub trait ArrowHashTable { fn update_heap_idx(&mut self, mapper: &[(usize, usize)]); fn heap_idx_at(&self, map_idx: usize) -> usize; fn take_all(&mut self, indexes: Vec) -> ArrayRef; - fn find_or_insert(&mut self, row_idx: usize, replace_idx: usize) -> (usize, bool); + fn find_or_insert( + &mut self, + row_idx: usize, + replace_idx: usize, + ) -> (usize, InsertKind); + /// Register the group at `row_idx` as all-NULL. Returns true if it was + /// newly registered; false if the group is already tracked or the NULL + /// group limit has been reached. + fn insert_null(&mut self, row_idx: usize) -> bool; + /// Remove the group at `row_idx` if it is registered as all-NULL. Returns + /// true if a NULL registration was removed. + fn remove_if_null(&mut self, row_idx: usize) -> bool; + /// Store indexes of all groups registered as all-NULL + fn null_map_idxs(&self) -> Vec; } /// Returns true if the given data type can be used as a top-K aggregation hash key. @@ -150,6 +183,13 @@ impl StringHashTable { Some(value.to_string()) } } + + /// Computes the id and its hash for the given row, for hash table lookups + fn id_and_hash(&self, row_idx: usize) -> (Option, u64) { + let id = self.extract_string_value(row_idx); + let hash = self.rnd.hash_one(id.as_deref()); + (id, hash) + } } impl ArrowHashTable for StringHashTable { @@ -179,7 +219,11 @@ impl ArrowHashTable for StringHashTable { } } - fn find_or_insert(&mut self, row_idx: usize, replace_idx: usize) -> (usize, bool) { + fn find_or_insert( + &mut self, + row_idx: usize, + replace_idx: usize, + ) -> (usize, InsertKind) { let id = self.extract_string_value(row_idx); // Compute hash and create equality closure for hash table lookup. @@ -190,6 +234,23 @@ impl ArrowHashTable for StringHashTable { // Use entry API to avoid double lookup self.map.find_or_insert(hash, id, replace_idx, eq) } + + fn insert_null(&mut self, row_idx: usize) -> bool { + let (id, hash) = self.id_and_hash(row_idx); + let id_for_eq = id.clone(); + let eq = move |mi: &Option| id_for_eq.as_deref() == mi.as_deref(); + self.map.insert_null(hash, id, eq) + } + + fn remove_if_null(&mut self, row_idx: usize) -> bool { + let (id, hash) = self.id_and_hash(row_idx); + let eq = move |mi: &Option| id.as_deref() == mi.as_deref(); + self.map.remove_if_null(hash, eq) + } + + fn null_map_idxs(&self) -> Vec { + self.map.null_map_idxs() + } } impl PrimitiveHashTable @@ -210,6 +271,18 @@ where kt, } } + + /// Computes the id and its hash for the given row, for hash table lookups + fn id_and_hash(&self, row_idx: usize) -> (Option, u64) { + let ids = self.owned.as_primitive::(); + let id: Option = if ids.is_null(row_idx) { + None + } else { + Some(ids.value(row_idx)) + }; + let hash: u64 = id.hash(&self.rnd); + (id, hash) + } } impl ArrowHashTable for PrimitiveHashTable @@ -247,7 +320,11 @@ where Arc::new(ids) } - fn find_or_insert(&mut self, row_idx: usize, replace_idx: usize) -> (usize, bool) { + fn find_or_insert( + &mut self, + row_idx: usize, + replace_idx: usize, + ) -> (usize, InsertKind) { let ids = self.owned.as_primitive::(); let id: Option = if ids.is_null(row_idx) { None @@ -261,6 +338,22 @@ where // Use entry API to avoid double lookup self.map.find_or_insert(hash, id, replace_idx, eq) } + + fn insert_null(&mut self, row_idx: usize) -> bool { + let (id, hash) = self.id_and_hash(row_idx); + let eq = move |mi: &Option| id == *mi; + self.map.insert_null(hash, id, eq) + } + + fn remove_if_null(&mut self, row_idx: usize) -> bool { + let (id, hash) = self.id_and_hash(row_idx); + let eq = move |mi: &Option| id == *mi; + self.map.remove_if_null(hash, eq) + } + + fn null_map_idxs(&self) -> Vec { + self.map.null_map_idxs() + } } use hashbrown::hash_table::Entry; @@ -269,8 +362,9 @@ impl TopKHashTable { Self { map: HashTable::with_capacity(capacity), store: Vec::with_capacity(capacity), - free_index: None, + free_indices: Vec::new(), limit, + null_count: 0, } } @@ -278,25 +372,33 @@ impl TopKHashTable { self.store[map_idx].as_ref().unwrap().heap_idx } - pub fn remove_if_full(&mut self, replace_idx: usize) -> usize { - if self.map.len() >= self.limit { - let item_to_remove = self.store[replace_idx].as_ref().unwrap(); - let hash = item_to_remove.hash; - let id_to_remove = &item_to_remove.id; - - let eq = |&idx: &usize| self.store[idx].as_ref().unwrap().id == *id_to_remove; - let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; - match self.map.entry(hash, eq, hasher) { - Entry::Occupied(entry) => { - let (removed_idx, _) = entry.remove(); - self.store[removed_idx] = None; - self.free_index = Some(removed_idx); - } - Entry::Vacant(_) => unreachable!(), + /// Remove the entry stored at `map_idx`, freeing its store slot for reuse + fn remove_at(&mut self, map_idx: usize) { + let item_to_remove = self.store[map_idx].as_ref().unwrap(); + let hash = item_to_remove.hash; + let id_to_remove = &item_to_remove.id; + + let eq = |&idx: &usize| self.store[idx].as_ref().unwrap().id == *id_to_remove; + let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; + match self.map.entry(hash, eq, hasher) { + Entry::Occupied(entry) => { + let (removed_idx, _) = entry.remove(); + self.store[removed_idx] = None; + self.free_indices.push(removed_idx); } + Entry::Vacant(_) => unreachable!(), + } + } + + pub fn remove_if_full(&mut self, replace_idx: usize) -> usize { + // All-NULL groups are tracked outside the heap, so only valued + // groups count towards the limit here + let valued_len = self.map.len() - self.null_count; + if valued_len >= self.limit { + self.remove_at(replace_idx); 0 // if full, always replace top node } else { - self.map.len() // if we're not full, always append to end + valued_len // if we're not full, always append to end } } @@ -307,7 +409,8 @@ impl TopKHashTable { } /// Find an existing entry or insert a new one, avoiding double hash table lookup. - /// Returns (map_idx, is_new) where is_new indicates if this was a new insertion. + /// Returns (map_idx, kind) where kind describes whether the group already + /// existed, was newly inserted, or was converted from an all-NULL group. /// If inserting a new entry and the table is full, replaces the entry at replace_idx. pub fn find_or_insert( &mut self, @@ -315,19 +418,28 @@ impl TopKHashTable { id: ID, replace_idx: usize, mut eq: impl FnMut(&ID) -> bool, - ) -> (usize, bool) { + ) -> (usize, InsertKind) { // Check if entry exists - this is the only hash table lookup + let mut replaced_null = false; { let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); if let Some(&map_idx) = self.map.find(hash, eq_fn) { - return (map_idx, false); + if self.store[map_idx].as_ref().unwrap().heap_idx == NULL_HEAP_IDX { + // This group was registered as all-NULL but now produced a + // value: unregister it so it is inserted as a valued group + self.remove_at(map_idx); + self.null_count -= 1; + replaced_null = true; + } else { + return (map_idx, InsertKind::Existing); + } } } // Entry doesn't exist - compute heap_idx and prepare item let heap_idx = self.remove_if_full(replace_idx); let mi = HashTableItem::new(hash, id, heap_idx); - let store_idx = if let Some(idx) = self.free_index.take() { + let store_idx = if let Some(idx) = self.free_indices.pop() { self.store[idx] = Some(mi); idx } else { @@ -343,7 +455,80 @@ impl TopKHashTable { // Insert without checking again since we already confirmed it doesn't exist self.map.insert_unique(hash, store_idx, hasher); - (store_idx, true) + let kind = if replaced_null { + InsertKind::ReplacedNull + } else { + InsertKind::New + }; + (store_idx, kind) + } + + /// Register a group whose aggregate values are all NULL, unless it is + /// already tracked. NULL groups are stored with a sentinel `heap_idx` and + /// never enter the heap. At most `limit` NULL groups are tracked: they all + /// tie on the sort key, so any `limit` of them is a valid top-k superset. + /// Returns true if the group was newly registered. + pub fn insert_null( + &mut self, + hash: u64, + id: ID, + mut eq: impl FnMut(&ID) -> bool, + ) -> bool { + { + let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); + if self.map.find(hash, eq_fn).is_some() { + return false; + } + } + if self.null_count >= self.limit { + return false; + } + + let mi = HashTableItem::new(hash, id, NULL_HEAP_IDX); + let store_idx = if let Some(idx) = self.free_indices.pop() { + self.store[idx] = Some(mi); + idx + } else { + self.store.push(Some(mi)); + self.store.len() - 1 + }; + + let hasher = |idx: &usize| self.store[*idx].as_ref().unwrap().hash; + if self.map.len() == self.map.capacity() { + self.map.reserve(self.limit, hasher); + } + self.map.insert_unique(hash, store_idx, hasher); + self.null_count += 1; + true + } + + /// Remove the given group if it is registered as all-NULL. Used when an + /// all-NULL group produces a value that loses to the current top-k: the + /// group can no longer reach the top-k, but it must not be emitted with a + /// NULL value either. Returns true if a NULL registration was removed. + pub fn remove_if_null(&mut self, hash: u64, mut eq: impl FnMut(&ID) -> bool) -> bool { + let eq_fn = |idx: &usize| eq(&self.store[*idx].as_ref().unwrap().id); + if let Some(&map_idx) = self.map.find(hash, eq_fn) + && self.store[map_idx].as_ref().unwrap().heap_idx == NULL_HEAP_IDX + { + self.remove_at(map_idx); + self.null_count -= 1; + return true; + } + false + } + + /// Store indexes of all groups registered as all-NULL + pub fn null_map_idxs(&self) -> Vec { + self.store + .iter() + .enumerate() + .filter_map(|(idx, item)| { + item.as_ref() + .filter(|item| item.heap_idx == NULL_HEAP_IDX) + .map(|_| idx) + }) + .collect() } pub fn len(&self) -> usize { @@ -357,7 +542,8 @@ impl TopKHashTable { .collect(); self.map.clear(); self.store.clear(); - self.free_index = None; + self.free_indices.clear(); + self.null_count = 0; ids } } @@ -453,9 +639,9 @@ mod tests { for (heap_idx, id) in ["1", "2", "3", "4", "5"].iter().enumerate() { let value = Some(id.to_string()); let hash = heap_idx as u64; - let (map_idx, is_new) = + let (map_idx, kind) = map.find_or_insert(hash, value.clone(), heap_idx, |v| *v == value); - assert!(is_new, "Entry should be new"); + assert_eq!(kind, InsertKind::New, "Entry should be new"); heap_to_map.insert(heap_idx, map_idx); } @@ -477,4 +663,65 @@ mod tests { Ok(()) } + + #[test] + fn should_track_null_groups() -> Result<()> { + let mut map = TopKHashTable::>::new(2, 10); + + let a = Some("a".to_string()); + let b = Some("b".to_string()); + let c = Some("c".to_string()); + + // register two all-NULL groups; the third exceeds the NULL group limit + assert!(map.insert_null(100, a.clone(), |v| *v == a)); + assert!(map.insert_null(200, b.clone(), |v| *v == b)); + assert!(!map.insert_null(300, c.clone(), |v| *v == c)); + // re-registering an existing NULL group is a no-op + assert!(!map.insert_null(100, a.clone(), |v| *v == a)); + assert_eq!(map.null_count, 2); + assert_eq!(map.null_map_idxs(), vec![0, 1]); + + // a valued insert for a NULL group converts it to a valued group + let (map_idx, kind) = map.find_or_insert(200, b.clone(), 0, |v| *v == b); + assert_eq!(kind, InsertKind::ReplacedNull, "NULL group should convert"); + assert_eq!(map.heap_idx_at(map_idx), 0, "Heap should append at 0"); + assert_eq!(map.null_count, 1); + assert_eq!(map.null_map_idxs(), vec![0]); + + // remove the remaining NULL group; removing twice is a no-op + map.remove_if_null(100, |v| *v == a); + assert_eq!(map.null_count, 0); + assert!(map.null_map_idxs().is_empty()); + map.remove_if_null(100, |v| *v == a); + // removing a valued group via remove_if_null is a no-op + map.remove_if_null(200, |v| *v == b); + assert_eq!(map.len(), 1); + + Ok(()) + } + + #[test] + fn should_reuse_all_freed_store_slots() -> Result<()> { + let mut map = TopKHashTable::>::new(1, 10); + + let a = Some("a".to_string()); + let b = Some("b".to_string()); + let c = Some("c".to_string()); + + let (b_idx, kind) = map.find_or_insert(100, b.clone(), 0, |v| *v == b); + assert_eq!(kind, InsertKind::New); + assert!(map.insert_null(200, a.clone(), |v| *v == a)); + + // Converting a NULL group while the valued heap is full frees two + // slots: the NULL registration and the evicted valued group. + let (_, kind) = map.find_or_insert(200, a.clone(), b_idx, |v| *v == a); + assert_eq!(kind, InsertKind::ReplacedNull); + + // Both freed slots must remain reusable. Otherwise repeated + // conversions make the backing store grow without bound. + assert!(map.insert_null(300, c.clone(), |v| *v == c)); + assert_eq!(map.store.len(), 2); + + Ok(()) + } } diff --git a/datafusion/physical-plan/src/aggregates/topk/heap.rs b/datafusion/physical-plan/src/aggregates/topk/heap.rs index 889fe04bf830a..ca321cdf99784 100644 --- a/datafusion/physical-plan/src/aggregates/topk/heap.rs +++ b/datafusion/physical-plan/src/aggregates/topk/heap.rs @@ -334,10 +334,7 @@ impl TopKHeap { pub fn worst_val(&self) -> Option<&VAL> { let root = self.heap.first()?; - let hi = match root { - None => return None, - Some(hi) => hi, - }; + let hi = root.as_ref()?; Some(&hi.val) } diff --git a/datafusion/physical-plan/src/aggregates/topk/priority_map.rs b/datafusion/physical-plan/src/aggregates/topk/priority_map.rs index c74b648d373ce..f46cb22a7a63c 100644 --- a/datafusion/physical-plan/src/aggregates/topk/priority_map.rs +++ b/datafusion/physical-plan/src/aggregates/topk/priority_map.rs @@ -17,9 +17,10 @@ //! A `Map` / `PriorityQueue` combo that evicts the worst values after reaching `capacity` -use crate::aggregates::topk::hash_table::{ArrowHashTable, new_hash_table}; +use crate::aggregates::topk::hash_table::{ArrowHashTable, InsertKind, new_hash_table}; use crate::aggregates::topk::heap::{ArrowHeap, new_heap}; -use arrow::array::ArrayRef; +use arrow::array::{ArrayRef, new_null_array}; +use arrow::compute::concat; use arrow::datatypes::DataType; use datafusion_common::Result; @@ -29,6 +30,11 @@ pub struct PriorityMap { heap: Box, capacity: usize, mapper: Vec<(usize, usize)>, + val_type: DataType, + /// Mirror of the map's all-NULL group count, kept as a plain field so the + /// per-row `insert` path can check it without a `dyn` call (measured to + /// regress the topk_aggregate benchmarks when read through the trait) + null_count: usize, } impl PriorityMap { @@ -40,9 +46,11 @@ impl PriorityMap { ) -> Result { Ok(Self { map: new_hash_table(capacity, key_type)?, - heap: new_heap(capacity, descending, val_type)?, + heap: new_heap(capacity, descending, val_type.clone())?, capacity, mapper: Vec::with_capacity(capacity), + val_type, + null_count: 0, }) } @@ -53,19 +61,47 @@ impl PriorityMap { pub fn insert(&mut self, row_idx: usize) -> Result<()> { assert!(self.map.len() <= self.capacity, "Overflow"); + debug_assert_eq!(self.null_count, 0); // if we're full, and the new val is worse than all our values, just bail if self.heap.is_worse(row_idx) { return Ok(()); } + self.insert_eligible(row_idx) + } + + /// Insert a value while all-NULL groups are being tracked. This is kept + /// separate from [`Self::insert`] so the common no-NULL path does not pay + /// for NULL bookkeeping on every row. + pub fn insert_with_null_groups(&mut self, row_idx: usize) -> Result<()> { + // valued groups are capped at `capacity`; up to `capacity` additional + // all-NULL groups may be tracked alongside them + assert!(self.map.len() <= 2 * self.capacity, "Overflow"); + + if self.heap.is_worse(row_idx) { + // A group that was registered as all-NULL now has a value that + // loses to the current top-k: it can no longer reach the top-k, + // but it must not be emitted with a NULL value either + if self.null_count > 0 && self.map.remove_if_null(row_idx) { + self.null_count -= 1; + } + return Ok(()); + } + self.insert_eligible(row_idx) + } + + fn insert_eligible(&mut self, row_idx: usize) -> Result<()> { let map = &mut self.mapper; // handle new groups we haven't seen yet map.clear(); let replace_idx = self.heap.worst_map_idx(); - let (map_idx, did_insert) = self.map.find_or_insert(row_idx, replace_idx); - if did_insert { + let (map_idx, kind) = self.map.find_or_insert(row_idx, replace_idx); + if kind == InsertKind::ReplacedNull { + self.null_count -= 1; + } + if kind != InsertKind::Existing { self.heap.insert(row_idx, map_idx, map); self.map.update_heap_idx(map); return Ok(()); @@ -80,9 +116,35 @@ impl PriorityMap { Ok(()) } + pub fn has_null_groups(&self) -> bool { + self.null_count > 0 + } + + /// Track a group whose aggregate values are all NULL, so it can be emitted + /// with a NULL value. MIN/MAX ignore NULL inputs, but an all-NULL group + /// must still appear in the aggregation output; such groups all tie on the + /// sort key, so tracking up to `capacity` of them preserves top-k semantics. + pub fn insert_null(&mut self, row_idx: usize) { + assert!(self.map.len() <= 2 * self.capacity, "Overflow"); + if self.map.insert_null(row_idx) { + self.null_count += 1; + } + } + pub fn emit(&mut self) -> Result> { - let (vals, map_idxs) = self.heap.drain(); + let (vals, mut map_idxs) = self.heap.drain(); + // Groups whose values are all NULL are tracked in the map only; + // append them with a NULL value so they are not lost from the output + let null_idxs = self.map.null_map_idxs(); + let vals = if null_idxs.is_empty() { + vals + } else { + map_idxs.extend(null_idxs.iter().copied()); + let nulls = new_null_array(&self.val_type, null_idxs.len()); + concat(&[vals.as_ref(), nulls.as_ref()])? + }; let ids = self.map.take_all(map_idxs); + self.null_count = 0; Ok(vec![ids, vals]) } @@ -495,6 +557,224 @@ mod tests { Ok(()) } + #[test] + fn should_emit_all_null_groups() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None, None])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?; + agg.set_batch(ids, vals); + agg.insert_null(0); + agg.insert_null(1); + // re-registering an existing NULL group is a no-op + agg.insert_null(0); + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | | + | 2 | | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_emit_null_groups_alongside_valued_groups() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2", "3"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![Some(7), None, Some(3)])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 3, true)?; + agg.set_batch(ids, vals); + agg.insert(0)?; + agg.insert_null(1); + agg.insert_with_null_groups(2)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 7 | + | 3 | 3 | + | 2 | | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_cap_null_groups_at_limit() -> Result<()> { + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1", "2", "3", "4", "5"])); + let vals: ArrayRef = + Arc::new(Int64Array::from(vec![None, None, None, None, None])); + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, false)?; + agg.set_batch(ids, vals); + for row_idx in 0..5 { + agg.insert_null(row_idx); + } + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | | + | 2 | | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_convert_null_group_to_valued() -> Result<()> { + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, true)?; + + // group "1" only produces NULLs in the first batch + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None])); + agg.set_batch(ids, vals); + agg.insert_null(0); + + // group "1" produces a value in a later batch + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![5])); + agg.set_batch(ids, vals); + agg.insert_with_null_groups(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 5 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_not_duplicate_valued_group_as_null() -> Result<()> { + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 2, false)?; + + // group "1" produces a value in the first batch + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![5])); + agg.set_batch(ids, vals); + agg.insert(0)?; + + // group "1" only produces NULLs in a later batch + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None])); + agg.set_batch(ids, vals); + agg.insert_null(0); + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 5 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_evict_worst_when_converting_null_group() -> Result<()> { + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, true)?; + + // group "2" holds the single top-k slot + let ids: ArrayRef = Arc::new(StringArray::from(vec!["2"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![10])); + agg.set_batch(ids, vals); + agg.insert(0)?; + + // group "1" starts out all-NULL + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None])); + agg.set_batch(ids, vals); + agg.insert_null(0); + + // group "1" produces a better value and evicts group "2" + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![20])); + agg.set_batch(ids, vals); + agg.insert_with_null_groups(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 1 | 20 | + +----------+--------------+ + " + ); + + Ok(()) + } + + #[test] + fn should_drop_null_group_that_loses_to_topk() -> Result<()> { + let mut agg = PriorityMap::new(DataType::Utf8, DataType::Int64, 1, true)?; + + // group "1" starts out all-NULL + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![None])); + agg.set_batch(ids, vals); + agg.insert_null(0); + + // group "2" fills the single top-k slot + let ids: ArrayRef = Arc::new(StringArray::from(vec!["2"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![10])); + agg.set_batch(ids, vals); + agg.insert_with_null_groups(0)?; + + // group "1" produces a value that loses to the current top-k: the + // group can no longer reach the top-k and must not be emitted as NULL + let ids: ArrayRef = Arc::new(StringArray::from(vec!["1"])); + let vals: ArrayRef = Arc::new(Int64Array::from(vec![5])); + agg.set_batch(ids, vals); + agg.insert_with_null_groups(0)?; + + let cols = agg.emit()?; + let batch = RecordBatch::try_new(test_schema(), cols)?; + let actual = format!("{}", pretty_format_batches(&[batch])?); + assert_snapshot!(actual, @r" + +----------+--------------+ + | trace_id | timestamp_ms | + +----------+--------------+ + | 2 | 10 | + +----------+--------------+ + " + ); + + Ok(()) + } + fn test_schema() -> SchemaRef { Arc::new(Schema::new(vec![ Field::new("trace_id", DataType::Utf8, true), diff --git a/datafusion/physical-plan/src/analyze.rs b/datafusion/physical-plan/src/analyze.rs index ea3abf439e4c1..d1519828c24be 100644 --- a/datafusion/physical-plan/src/analyze.rs +++ b/datafusion/physical-plan/src/analyze.rs @@ -25,13 +25,20 @@ use super::{ SendableRecordBatchStream, }; use crate::display::DisplayableExecutionPlan; +use crate::execution_plan::EvaluationType; use crate::metrics::{MetricCategory, MetricType}; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, +}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; +use datafusion_common::format::ExplainFormat; use datafusion_common::instant::Instant; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{DataFusionError, Result, assert_eq_or_internal_err}; +use datafusion_common::{ + DataFusionError, Result, assert_eq_or_internal_err, internal_err, +}; use datafusion_execution::TaskContext; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_expr::PhysicalExpr; @@ -50,6 +57,8 @@ pub struct AnalyzeExec { metric_types: Vec, /// Optional filter by semantic category (rows / bytes / timing). metric_categories: Option>, + /// Output format for the rendered plan + metrics. + format: ExplainFormat, /// The input plan (the plan being analyzed) pub(crate) input: Arc, /// The output schema for RecordBatches of this exec node @@ -57,27 +66,81 @@ pub struct AnalyzeExec { cache: Arc, } -impl AnalyzeExec { - /// Create a new AnalyzeExec +/// Builder for [`AnalyzeExec`]. +/// +/// Builder for [AnalyzeExec]. +pub struct AnalyzeExecBuilder { + verbose: bool, + show_statistics: bool, + input: Arc, + schema: SchemaRef, + metric_types: Vec, + metric_categories: Option>, + format: ExplainFormat, +} + +impl AnalyzeExecBuilder { pub fn new( verbose: bool, show_statistics: bool, - metric_types: Vec, - metric_categories: Option>, input: Arc, schema: SchemaRef, ) -> Self { - let cache = Self::compute_properties(&input, Arc::clone(&schema)); - AnalyzeExec { + Self { verbose, show_statistics, - metric_types, - metric_categories, input, schema, + metric_types: vec![MetricType::Summary, MetricType::Dev], + metric_categories: None, + format: ExplainFormat::Indent, + } + } + + pub fn with_metric_types(mut self, metric_types: Vec) -> Self { + self.metric_types = metric_types; + self + } + + pub fn with_metric_categories( + mut self, + metric_categories: Option>, + ) -> Self { + self.metric_categories = metric_categories; + self + } + + pub fn with_format(mut self, format: ExplainFormat) -> Self { + self.format = format; + self + } + + pub fn build(self) -> AnalyzeExec { + let cache = + AnalyzeExec::compute_properties(&self.input, Arc::clone(&self.schema)); + AnalyzeExec { + verbose: self.verbose, + show_statistics: self.show_statistics, + metric_types: self.metric_types, + metric_categories: self.metric_categories, + format: self.format, + input: self.input, + schema: self.schema, cache: Arc::new(cache), } } +} + +impl AnalyzeExec { + /// Returns a builder for constructing an [`AnalyzeExec`]. + pub fn builder( + verbose: bool, + show_statistics: bool, + input: Arc, + schema: SchemaRef, + ) -> AnalyzeExecBuilder { + AnalyzeExecBuilder::new(verbose, show_statistics, input, schema) + } /// Access to verbose pub fn verbose(&self) -> bool { @@ -94,6 +157,11 @@ impl AnalyzeExec { self.metric_categories.as_deref() } + /// Access to format + pub fn format(&self) -> &ExplainFormat { + &self.format + } + /// The input plan pub fn input(&self) -> &Arc { &self.input @@ -110,6 +178,7 @@ impl AnalyzeExec { input.pipeline_behavior(), input.boundedness(), ) + .with_evaluation_type(EvaluationType::Eager) } } @@ -146,28 +215,49 @@ impl ExecutionPlan for AnalyzeExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { - Ok(Arc::new(Self::new( - self.verbose, - self.show_statistics, - self.metric_types.clone(), - self.metric_categories.clone(), - children.pop().unwrap(), - Arc::clone(&self.schema), - ))) + Ok(Arc::new( + AnalyzeExec::builder( + self.verbose, + self.show_statistics, + children.pop().unwrap(), + Arc::clone(&self.schema), + ) + .with_metric_types(self.metric_types.clone()) + .with_metric_categories(self.metric_categories.clone()) + .with_format(self.format.clone()) + .build(), + )) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn execute( @@ -204,6 +294,7 @@ impl ExecutionPlan for AnalyzeExec { let show_statistics = self.show_statistics; let metric_types = self.metric_types.clone(); let metric_categories = self.metric_categories.clone(); + let format = self.format.clone(); // future that gathers the results from all the tasks in the // JoinSet that computes the overall row count and final @@ -226,6 +317,7 @@ impl ExecutionPlan for AnalyzeExec { &captured_schema, &metric_types, metric_categories.as_deref(), + &format, ) }; @@ -234,6 +326,129 @@ impl ExecutionPlan for AnalyzeExec { futures::stream::once(output), ))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `AnalyzeExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + verbose, + show_statistics, + // TODO: not on the wire. `AnalyzeExecBuilder` always resets this to + // `[Summary, Dev]`, so a non-default selection is lost on + // round-trip. Fixing it needs a new proto field. + metric_types: _, + metric_categories, + format, + input, + schema, + // Derived at construction from `input` and `schema`. + cache: _, + } = self; + + let input = ctx.encode_child(input)?; + let (has_metric_categories, metric_categories) = match metric_categories { + Some(categories) => { + (true, categories.iter().map(ToString::to_string).collect()) + } + None => (false, vec![]), + }; + let format = match format { + ExplainFormat::Indent => protobuf::ExplainFormat::Indent, + ExplainFormat::Tree => protobuf::ExplainFormat::Tree, + ExplainFormat::PostgresJSON => protobuf::ExplainFormat::Pgjson, + ExplainFormat::Graphviz => protobuf::ExplainFormat::Graphviz, + } as i32; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Analyze(Box::new( + protobuf::AnalyzeExecNode { + verbose: *verbose, + show_statistics: *show_statistics, + input: Some(Box::new(input)), + schema: Some(schema.as_ref().try_into()?), + has_metric_categories, + metric_categories, + format, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl AnalyzeExec { + /// Reconstruct an [`AnalyzeExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let analyze = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Analyze, + "AnalyzeExec", + ); + // Exhaustive destructure: a new field on `AnalyzeExecNode` is a compile + // error here rather than a silently ignored wire field. + let protobuf::AnalyzeExecNode { + verbose, + show_statistics, + input, + schema, + has_metric_categories, + metric_categories, + format, + } = analyze.as_ref(); + + let input = + ctx.decode_required_child(input.as_deref(), "AnalyzeExec", "input")?; + let metric_categories = if *has_metric_categories { + Some( + metric_categories + .iter() + .map(|category| category.parse::()) + .collect::>>()?, + ) + } else { + None + }; + let proto_format = protobuf::ExplainFormat::try_from(*format).map_err(|_| { + DataFusionError::Internal(format!( + "Received an AnalyzeExecNode message with unknown ExplainFormat {format}" + )) + })?; + let format = match proto_format { + protobuf::ExplainFormat::Indent => ExplainFormat::Indent, + protobuf::ExplainFormat::Tree => ExplainFormat::Tree, + protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON, + protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz, + }; + let schema = schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "AnalyzeExec is missing required field 'schema'" + ) + })?; + Ok(Arc::new( + AnalyzeExec::builder( + *verbose, + *show_statistics, + input, + Arc::new(arrow::datatypes::Schema::try_from(schema)?), + ) + .with_metric_categories(metric_categories) + .with_format(format) + .build(), + )) + } } /// Creates the output of AnalyzeExec as a RecordBatch @@ -247,39 +462,61 @@ fn create_output_batch( schema: &SchemaRef, metric_types: &[MetricType], metric_categories: Option<&[MetricCategory]>, + format: &ExplainFormat, ) -> Result { let mut type_builder = StringBuilder::with_capacity(1, 1024); let mut plan_builder = StringBuilder::with_capacity(1, 1024); - // TODO use some sort of enum rather than strings? - type_builder.append_value("Plan with Metrics"); - - let annotated_plan = DisplayableExecutionPlan::with_metrics(input.as_ref()) - .set_metric_types(metric_types.to_vec()) - .set_metric_categories(metric_categories.map(|c| c.to_vec())) - .set_show_statistics(show_statistics) - .indent(verbose) - .to_string(); - plan_builder.append_value(annotated_plan); - - // Verbose output - // TODO make this more sophisticated - if verbose { - type_builder.append_value("Plan with Full Metrics"); - - let annotated_plan = DisplayableExecutionPlan::with_full_metrics(input.as_ref()) - .set_metric_types(metric_types.to_vec()) - .set_metric_categories(metric_categories.map(|c| c.to_vec())) - .set_show_statistics(show_statistics) - .indent(verbose) - .to_string(); - plan_builder.append_value(annotated_plan); - - type_builder.append_value("Output Rows"); - plan_builder.append_value(total_rows.to_string()); - - type_builder.append_value("Duration"); - plan_builder.append_value(format!("{duration:?}")); + match format { + ExplainFormat::Indent => { + // TODO use some sort of enum rather than strings? + type_builder.append_value("Plan with Metrics"); + let annotated_plan = DisplayableExecutionPlan::with_metrics(input.as_ref()) + .set_metric_types(metric_types.to_vec()) + .set_metric_categories(metric_categories.map(|c| c.to_vec())) + .set_show_statistics(show_statistics) + .indent(verbose) + .to_string(); + plan_builder.append_value(annotated_plan); + // Verbose output + // TODO make this more sophisticated + if verbose { + type_builder.append_value("Plan with Full Metrics"); + let annotated_plan = + DisplayableExecutionPlan::with_full_metrics(input.as_ref()) + .set_metric_types(metric_types.to_vec()) + .set_metric_categories(metric_categories.map(|c| c.to_vec())) + .set_show_statistics(show_statistics) + .indent(verbose) + .to_string(); + plan_builder.append_value(annotated_plan); + type_builder.append_value("Output Rows"); + plan_builder.append_value(total_rows.to_string()); + type_builder.append_value("Duration"); + plan_builder.append_value(format!("{duration:?}")); + } + } + ExplainFormat::PostgresJSON => { + // `show_statistics` is intentionally not forwarded here: the pgjson + // renderer does not emit statistics, and the planner rejects the + // `show_statistics` + pgjson combination up front. + type_builder.append_value("Plan with Metrics"); + let mut displayable = if verbose { + DisplayableExecutionPlan::with_full_metrics(input.as_ref()) + } else { + DisplayableExecutionPlan::with_metrics(input.as_ref()) + }; + displayable = displayable + .set_metric_types(metric_types.to_vec()) + .set_metric_categories(metric_categories.map(|c| c.to_vec())); + if verbose { + displayable = displayable.set_summary(Some(total_rows), Some(duration)); + } + plan_builder.append_value(displayable.pgjson(verbose).to_string()); + } + ExplainFormat::Tree | ExplainFormat::Graphviz => { + return internal_err!("AnalyzeExec does not support {format} output format"); + } } RecordBatch::try_new( @@ -314,14 +551,8 @@ mod tests { let blocking_exec = Arc::new(BlockingExec::new(Arc::clone(&schema), 1)); let refs = blocking_exec.refs(); - let analyze_exec = Arc::new(AnalyzeExec::new( - true, - false, - vec![MetricType::Summary, MetricType::Dev], - None, - blocking_exec, - schema, - )); + let analyze_exec = + Arc::new(AnalyzeExec::builder(true, false, blocking_exec, schema).build()); let fut = collect(analyze_exec, task_ctx); let mut fut = fut.boxed(); diff --git a/datafusion/physical-plan/src/async_func.rs b/datafusion/physical-plan/src/async_func.rs index 8ad4ecb096962..f3ef13d4fd3a8 100644 --- a/datafusion/physical-plan/src/async_func.rs +++ b/datafusion/physical-plan/src/async_func.rs @@ -19,14 +19,14 @@ use crate::coalesce::LimitedBatchCoalescer; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, ReplaceChildrenOptions, + validate_child_count, }; use arrow::array::RecordBatch; -use arrow_schema::{Fields, Schema, SchemaRef}; -use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::tree_node::{Transformed, TreeNode}; -use datafusion_common::{Result, assert_eq_or_internal_err}; +use arrow_schema::{FieldRef, Fields, Schema, SchemaRef}; +use datafusion_common::Result; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::ScalarFunctionExpr; use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; @@ -62,8 +62,8 @@ impl AsyncFuncExec { ) -> Result { let async_fields = async_exprs .iter() - .map(|async_expr| async_expr.field(input.schema().as_ref())) - .collect::>>()?; + .map(|async_expr| async_expr.return_field(input.schema().as_ref())) + .collect::>>()?; // compute the output schema: input schema then async expressions let fields: Fields = input @@ -71,7 +71,7 @@ impl AsyncFuncExec { .fields() .iter() .cloned() - .chain(async_fields.into_iter().map(Arc::new)) + .chain(async_fields) .collect(); let schema = Arc::new(Schema::new(fields)); @@ -107,6 +107,10 @@ impl AsyncFuncExec { )) } + #[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `AsyncFuncExec` serializes itself via `AsyncFuncExec::try_to_proto`, which reads the field directly. There is no replacement; please open an issue if you have a use case for it." + )] pub fn async_exprs(&self) -> &[Arc] { &self.async_exprs } @@ -114,17 +118,6 @@ impl AsyncFuncExec { pub fn input(&self) -> &Arc { &self.input } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for AsyncFuncExec { @@ -167,25 +160,54 @@ impl ExecutionPlan for AsyncFuncExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - Ok(TreeNodeRecursion::Continue) + crate::apply_expression_roots( + self.async_exprs + .iter() + .cloned() + .map(|expr| expr as Arc), + f, + ) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - assert_eq_or_internal_err!( - children.len(), - 1, - "AsyncFuncExec wrong number of children" - ); - check_if_same_properties!(self, children); - Ok(Arc::new(AsyncFuncExec::try_new( - self.async_exprs.clone(), - children.swap_remove(0), - )?)) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(AsyncFuncExec::try_new( + self.async_exprs.clone(), + children.swap_remove(0), + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -216,7 +238,7 @@ impl ExecutionPlan for AsyncFuncExec { input_stream, batch_coalescer: LimitedBatchCoalescer::new( Arc::clone(&self.input.schema()), - config_options_ref.execution.batch_size, + config_options_ref.execution.batch_size.get(), None, ), }; @@ -254,6 +276,96 @@ impl ExecutionPlan for AsyncFuncExec { fn metrics(&self) -> Option { Some(self.metrics.clone_inner()) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `AsyncFuncExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + async_exprs, + input, + // Derived at construction by `AsyncFuncExec::compute_properties`. + cache: _, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + } = self; + + let input = ctx.encode_child(input)?; + let async_expr_names = async_exprs.iter().map(|e| e.name().to_string()).collect(); + let async_exprs = ctx.encode_expressions(async_exprs.iter().map(|e| &e.func))?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::AsyncFunc(Box::new( + protobuf::AsyncFuncExecNode { + input: Some(Box::new(input)), + async_exprs, + async_expr_names, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl AsyncFuncExec { + /// Reconstruct an [`AsyncFuncExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. Child plans and expressions are decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::assert_eq_or_internal_err; + use datafusion_proto_models::protobuf; + let async_func = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::AsyncFunc, + "AsyncFuncExec", + ); + // Exhaustive destructure: a new field on `AsyncFuncExecNode` is a + // compile error here rather than a silently ignored wire field. + let protobuf::AsyncFuncExecNode { + input, + async_exprs, + async_expr_names, + } = async_func.as_ref(); + + let input = + ctx.decode_required_child(input.as_deref(), "AsyncFuncExec", "input")?; + let input_schema = input.schema(); + assert_eq_or_internal_err!( + async_exprs.len(), + async_expr_names.len(), + "AsyncFuncExecNode async_exprs length does not match async_expr_names" + ); + let async_exprs = async_exprs + .iter() + .zip(async_expr_names.iter()) + .map(|(expr, name)| { + let physical_expr = ctx.decode_expr(expr, input_schema.as_ref())?; + Ok(Arc::new(AsyncFuncExpr::try_new( + name.clone(), + physical_expr, + input_schema.as_ref(), + )?)) + }) + .collect::>>()?; + Ok(Arc::new(AsyncFuncExec::try_new(async_exprs, input)?)) + } } struct CoalesceInputStream { diff --git a/datafusion/physical-plan/src/buffer.rs b/datafusion/physical-plan/src/buffer.rs index 0cc4a1d71814e..24cca6b0b17f4 100644 --- a/datafusion/physical-plan/src/buffer.rs +++ b/datafusion/physical-plan/src/buffer.rs @@ -18,21 +18,24 @@ //! [`BufferExec`] decouples production and consumption on messages by buffering the input in the //! background up to a certain capacity. -use crate::execution_plan::{CardinalityEffect, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, }; use crate::projection::ProjectionExec; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, SortOrderPushdownResult, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, SortOrderPushdownResult, validate_child_count, }; use arrow::array::RecordBatch; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{Result, Statistics, internal_err, plan_err}; +use datafusion_common::{Result, Statistics, internal_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; @@ -41,9 +44,10 @@ use datafusion_physical_expr_common::metrics::{ }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; -use futures::{Stream, StreamExt, TryStreamExt}; +use futures::{FutureExt, Stream, StreamExt, TryStreamExt}; use pin_project_lite::pin_project; use std::fmt; +use std::panic::AssertUnwindSafe; use std::pin::Pin; use std::sync::Arc; use std::sync::atomic::{AtomicUsize, Ordering}; @@ -102,7 +106,8 @@ impl BufferExec { /// Builds a new [BufferExec] with the provided capacity in bytes. pub fn new(input: Arc, capacity: usize) -> Self { let properties = PlanProperties::clone(input.properties()) - .with_scheduling_type(SchedulingType::Cooperative); + .with_scheduling_type(SchedulingType::Cooperative) + .with_evaluation_type(EvaluationType::Eager); Self { input, @@ -121,17 +126,6 @@ impl BufferExec { pub fn capacity(&self) -> usize { self.capacity } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for BufferExec { @@ -170,20 +164,47 @@ impl ExecutionPlan for BufferExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - if children.len() != 1 { - return plan_err!("BufferExec can only have one child"); + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity))) + } } - Ok(Arc::new(Self::new(children.swap_remove(0), self.capacity))) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -200,8 +221,7 @@ impl ExecutionPlan for BufferExec { let curr_mem_out = Arc::clone(&curr_mem_in); let mut max_mem_in = 0; let max_mem = MetricBuilder::new(&self.metrics) - .with_category(MetricCategory::Bytes) - .gauge("max_mem_used", partition); + .peak_memory_usage("max_mem_used", partition); let curr_queued_in = Arc::new(AtomicUsize::new(0)); let curr_queued_out = Arc::clone(&curr_queued_in); @@ -244,8 +264,16 @@ impl ExecutionPlan for BufferExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - self.input.partition_statistics(partition) + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) } fn supports_limit_pushdown(&self) -> bool { @@ -261,9 +289,10 @@ impl ExecutionPlan for BufferExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } @@ -296,6 +325,48 @@ impl ExecutionPlan for BufferExec { Ok(Arc::new(Self::new(new_input, self.capacity)) as Arc) }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Buffer(Box::new( + protobuf::BufferExecNode { + input: Some(Box::new(input)), + capacity: self.capacity() as u64, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl BufferExec { + /// Reconstruct a [`BufferExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let buffer = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Buffer, + "BufferExec", + ); + let input = + ctx.decode_required_child(buffer.input.as_deref(), "BufferExec", "input")?; + Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) + } } /// Represents anything that occupies a capacity in a [MemoryBufferedStream]. @@ -344,11 +415,24 @@ impl MemoryBufferedStream { let item_or_err = tokio::select! { biased; _ = batch_tx.closed() => break, - item_or_err = input.next() => { - let Some(item_or_err) = item_or_err else { - break; // stream finished - }; - item_or_err + // Catch a panic in the input poll so it surfaces as a stream error + // instead of dropping `batch_tx` and looking like a clean EOF. + polled = AssertUnwindSafe(input.next()).catch_unwind() => { + match polled { + Ok(Some(item_or_err)) => item_or_err, + Ok(None) => break, // stream finished + Err(panic) => { + let msg = panic + .downcast_ref::<&str>() + .map(|s| s.to_string()) + .or_else(|| panic.downcast_ref::().cloned()) + .unwrap_or_else(|| "unknown panic".to_string()); + let _ = batch_tx.send(internal_err!( + "BufferExec input stream panicked: {msg}" + )); + break; + } + } } }; @@ -561,6 +645,29 @@ mod tests { Ok(()) } + #[tokio::test] + async fn panic_in_input_is_propagated() -> Result<(), Box> { + // A panic while polling the input must surface as a stream error, not a + // silent end-of-stream that drops the rest of the partition's output. + let input = futures::stream::iter([1, 2, 3, 4]).map(|v| { + if v == 3 { + panic!("boom on 3"); + } + Ok(v) + }); + let (_, res) = memory_pool_and_reservation(); + + let mut buffered = MemoryBufferedStream::new(input, 10, res); + wait_for_buffering().await; + + pull_ok_msg(&mut buffered).await?; + pull_ok_msg(&mut buffered).await?; + let err = pull_err_msg(&mut buffered).await?; + assert_contains!(err.to_string(), "panicked"); + + Ok(()) + } + #[tokio::test] async fn memory_gets_released_if_stream_drops() -> Result<(), Box> { let input = futures::stream::iter([1, 2, 3, 4]).map(Ok); diff --git a/datafusion/physical-plan/src/coalesce_batches.rs b/datafusion/physical-plan/src/coalesce_batches.rs index 34cd770260915..cb0f9b2ce4b36 100644 --- a/datafusion/physical-plan/src/coalesce_batches.rs +++ b/datafusion/physical-plan/src/coalesce_batches.rs @@ -24,10 +24,11 @@ use std::task::{Context, Poll}; use super::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use super::{DisplayAs, ExecutionPlanProperties, PlanProperties, Statistics}; use crate::projection::ProjectionExec; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayFormatType, ExecutionPlan, RecordBatchStream, SendableRecordBatchStream, - check_if_same_properties, + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, RecordBatchStream, + ReplaceChildrenOptions, SendableRecordBatchStream, validate_child_count, }; use arrow::datatypes::SchemaRef; @@ -38,7 +39,7 @@ use datafusion_execution::TaskContext; use datafusion_physical_expr::PhysicalExpr; use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; -use crate::execution_plan::CardinalityEffect; +use crate::execution_plan::{CardinalityEffect, replace_children_if_necessary}; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -117,17 +118,6 @@ impl CoalesceBatchesExec { input.boundedness(), ) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } #[expect(deprecated)] @@ -186,20 +176,48 @@ impl ExecutionPlan for CoalesceBatchesExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) - .with_fetch(self.fetch), - )) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new( + CoalesceBatchesExec::new(children.swap_remove(0), self.target_batch_size) + .with_fetch(self.fetch), + )), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -223,8 +241,16 @@ impl ExecutionPlan for CoalesceBatchesExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?); + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let stats = input_stats[0].as_ref().clone(); Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) } @@ -251,9 +277,10 @@ impl ExecutionPlan for CoalesceBatchesExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } @@ -289,6 +316,61 @@ impl ExecutionPlan for CoalesceBatchesExec { ) as Arc) }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches( + Box::new(protobuf::CoalesceBatchesExecNode { + input: Some(Box::new(input)), + target_batch_size: self.target_batch_size() as u32, + fetch: self.fetch().map(|n| n as u32), + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +#[expect(deprecated)] +impl CoalesceBatchesExec { + /// Reconstruct a [`CoalesceBatchesExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. The child plan is decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let coalesce_batches = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::CoalesceBatches, + "CoalesceBatchesExec", + ); + let input = ctx.decode_required_child( + coalesce_batches.input.as_deref(), + "CoalesceBatchesExec", + "input", + )?; + Ok(Arc::new( + CoalesceBatchesExec::new(input, coalesce_batches.target_batch_size as usize) + .with_fetch(coalesce_batches.fetch.map(|f| f as usize)), + )) + } } /// Stream for [`CoalesceBatchesExec`]. See [`CoalesceBatchesExec`] for more details. diff --git a/datafusion/physical-plan/src/coalesce_partitions.rs b/datafusion/physical-plan/src/coalesce_partitions.rs index 3399554612431..6f58eb2f1e6be 100644 --- a/datafusion/physical-plan/src/coalesce_partitions.rs +++ b/datafusion/physical-plan/src/coalesce_partitions.rs @@ -26,11 +26,17 @@ use super::{ DisplayAs, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, Statistics, }; -use crate::execution_plan::{CardinalityEffect, EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; use crate::filter_pushdown::{FilterDescription, FilterPushdownPhase}; use crate::projection::{ProjectionExec, make_with_child}; use crate::sort_pushdown::SortOrderPushdownResult; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning, check_if_same_properties}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, validate_child_count, +}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use datafusion_common::config::ConfigOptions; @@ -100,17 +106,6 @@ impl CoalescePartitionsExec { .with_evaluation_type(drive) .with_scheduling_type(scheduling) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for CoalescePartitionsExec { @@ -156,19 +151,49 @@ impl ExecutionPlan for CoalescePartitionsExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); - plan.fetch = self.fetch; - Ok(Arc::new(plan)) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut plan = CoalescePartitionsExec::new(children.swap_remove(0)); + plan.fetch = self.fetch; + Ok(Arc::new(plan)) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -239,8 +264,16 @@ impl ExecutionPlan for CoalescePartitionsExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, _partition: Option) -> Result> { - let stats = Arc::unwrap_or_clone(self.input.partition_statistics(None)?); + fn child_stats_requests(&self, _partition: Option) -> Vec { + vec![ChildStats::At(None)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let stats = input_stats[0].as_ref().clone(); Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) } @@ -295,8 +328,7 @@ impl ExecutionPlan for CoalescePartitionsExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } @@ -345,6 +377,56 @@ impl ExecutionPlan for CoalescePartitionsExec { } }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Merge(Box::new( + protobuf::CoalescePartitionsExecNode { + input: Some(Box::new(input)), + fetch: self.fetch().map(|f| f as u32), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CoalescePartitionsExec { + /// Reconstruct a [`CoalescePartitionsExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. Note the protobuf + /// variant is named `Merge` (node [`CoalescePartitionsExecNode`]). + /// + /// [`CoalescePartitionsExecNode`]: datafusion_proto_models::protobuf::CoalescePartitionsExecNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let merge = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Merge, + "CoalescePartitionsExec", + ); + let input = ctx.decode_required_child( + merge.input.as_deref(), + "CoalescePartitionsExec", + "input", + )?; + Ok(Arc::new( + CoalescePartitionsExec::new(input) + .with_fetch(merge.fetch.map(|f| f as usize)), + )) + } } #[cfg(test)] diff --git a/datafusion/physical-plan/src/common.rs b/datafusion/physical-plan/src/common.rs index 0dafcf6bd3390..734ec96debc85 100644 --- a/datafusion/physical-plan/src/common.rs +++ b/datafusion/physical-plan/src/common.rs @@ -181,7 +181,8 @@ pub fn project_plan_to_schema( } /// If running in a tokio context spawns the execution of `stream` to a separate task -/// allowing it to execute in parallel with an intermediate buffer of size `buffer` +/// allowing it to execute in parallel with an intermediate buffer of size `buffer`. +/// At most `buffer` record batches will be produced ahead of the consumer. pub fn spawn_buffered( mut input: SendableRecordBatchStream, buffer: usize, @@ -196,11 +197,22 @@ pub fn spawn_buffered( let sender = builder.tx(); builder.spawn(async move { - while let Some(item) = input.next().await { - if sender.send(item).await.is_err() { - // Receiver dropped when query is shutdown early (e.g., limit) or error, - // no need to return propagate the send error. - return Ok(()); + // We call `reserve` (which waits until there's room for at least 1 message in the + // channel buffer) **before** polling from input to ensure we hold a maximum of + // `buffer` record batches in memory. + // Polling from input and then calling send() would block when the channel is full + // so it would essentially hold `buffer` + 1 record batches: + // * `buffer`: this many elements would live inside the channel, since this is the + // channel's capacity + // * 1 extra RecordBatch which was produced, but there was no room for it in the + // channel, so it's being owned by the send() future, which keeps the batch in + // memory while it waits for a slot to free up + while let Ok(permit) = sender.reserve().await { + // Receiver dropped when query is shutdown early (e.g., limit) or error, + // no need to return propagate the send error. + match input.next().await { + Some(item) => permit.send(item), + None => break, } } @@ -298,10 +310,13 @@ mod tests { use crate::empty::EmptyExec; use crate::projection::ProjectionExec; + use crate::stream::RecordBatchStreamAdapter; + use futures::stream; use std::collections::HashMap; + use std::sync::atomic::{AtomicUsize, Ordering}; use arrow::{ - array::{Float32Array, Float64Array, UInt64Array}, + array::{Float32Array, Float64Array, Int32Array, UInt64Array}, datatypes::{DataType, Field, Schema}, }; @@ -554,4 +569,55 @@ mod tests { let err = project_plan_to_schema(input, &expected_schema).unwrap_err(); assert!(err.to_string().contains("schema metadata differ")); } + + /// Verifies that `spawn_buffered` holds exactly `buffer` record batches in memory + /// when no receiver is polling + async fn spawn_buffered_max_in_flight_batches(buffer_size: usize) { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let num_batches = 10; + + let produced_count = Arc::new(AtomicUsize::new(0)); + let produced_clone = Arc::clone(&produced_count); + let schema_clone = Arc::clone(&schema); + + // Stream increments the counter each time a batch is pulled by the producer. + let input_stream = stream::unfold(0usize, move |i| { + let schema = Arc::clone(&schema_clone); + let counter = Arc::clone(&produced_clone); + async move { + if i >= num_batches { + return None; + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![i as i32]))], + ) + .unwrap(); + counter.fetch_add(1, Ordering::SeqCst); + Some((Ok(batch), i + 1)) + } + }); + + let input = Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + input_stream, + )); + // Drop the returned stream immediately so no receiver is ever polled. + let _buffered = spawn_buffered(input, buffer_size); + + // Give the producer task time to fill the channel and stall on send(). + tokio::time::sleep(std::time::Duration::from_millis(100)).await; + + assert_eq!( + produced_count.load(Ordering::SeqCst), + buffer_size, + "expected exactly {buffer_size} batch(es) in memory with no receiver polling" + ); + } + + #[tokio::test(flavor = "multi_thread")] + async fn test_spawn_buffered_max_in_flight_batches() { + spawn_buffered_max_in_flight_batches(1).await; + spawn_buffered_max_in_flight_batches(2).await; + } } diff --git a/datafusion/physical-plan/src/coop.rs b/datafusion/physical-plan/src/coop.rs index fe6a3bc3d5678..9e27b26d6e7c9 100644 --- a/datafusion/physical-plan/src/coop.rs +++ b/datafusion/physical-plan/src/coop.rs @@ -85,16 +85,18 @@ use crate::filter_pushdown::{ FilterPushdownPropagation, }; use crate::projection::ProjectionExec; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, SortOrderPushdownResult, check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + SortOrderPushdownResult, validate_child_count, }; use arrow::record_batch::RecordBatch; use arrow_schema::Schema; -use datafusion_common::{Result, Statistics, assert_eq_or_internal_err}; +use datafusion_common::{Result, Statistics}; use datafusion_execution::TaskContext; -use crate::execution_plan::SchedulingType; +use crate::execution_plan::{SchedulingType, replace_children_if_necessary}; use crate::stream::RecordBatchStreamAdapter; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; use futures::{Stream, StreamExt}; @@ -234,16 +236,6 @@ impl CooperativeExec { pub fn input(&self) -> &Arc { &self.input } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - ..Self::clone(self) - } - } } impl DisplayAs for CooperativeExec { @@ -279,22 +271,46 @@ impl ExecutionPlan for CooperativeExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - assert_eq_or_internal_err!( - children.len(), - 1, - "CooperativeExec requires exactly one child" - ); - check_if_same_properties!(self, children); - Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(CooperativeExec::new(children.swap_remove(0)))) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -306,8 +322,16 @@ impl ExecutionPlan for CooperativeExec { Ok(make_cooperative(child_stream)) } - fn partition_statistics(&self, partition: Option) -> Result> { - self.input.partition_statistics(partition) + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) } fn supports_limit_pushdown(&self) -> bool { @@ -323,9 +347,10 @@ impl ExecutionPlan for CooperativeExec { projection: &ProjectionExec, ) -> Result>> { match self.input.try_swapping_with_projection(projection)? { - Some(new_input) => Ok(Some( - Arc::new(self.clone()).with_new_children(vec![new_input])?, - )), + Some(new_input) => Ok(Some(replace_children_if_necessary( + Arc::new(self.clone()), + vec![new_input], + )?)), None => Ok(None), } } @@ -356,11 +381,13 @@ impl ExecutionPlan for CooperativeExec { match child.try_pushdown_sort(order)? { SortOrderPushdownResult::Exact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Exact { inner: new_exec }) } SortOrderPushdownResult::Inexact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Inexact { inner: new_exec }) } SortOrderPushdownResult::Unsupported => { @@ -368,6 +395,50 @@ impl ExecutionPlan for CooperativeExec { } } } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Cooperative(Box::new( + protobuf::CooperativeExecNode { + input: Some(Box::new(input)), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CooperativeExec { + /// Reconstruct a [`CooperativeExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let cooperative = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Cooperative, + "CooperativeExec", + ); + let input = ctx.decode_required_child( + cooperative.input.as_deref(), + "CooperativeExec", + "input", + )?; + Ok(Arc::new(CooperativeExec::new(input))) + } } /// Creates a [`CooperativeStream`] wrapper around the given [`RecordBatchStream`]. diff --git a/datafusion/physical-plan/src/display.rs b/datafusion/physical-plan/src/display.rs index 756a68b1a958d..d2bdcef2e97a3 100644 --- a/datafusion/physical-plan/src/display.rs +++ b/datafusion/physical-plan/src/display.rs @@ -21,6 +21,7 @@ use std::collections::{BTreeMap, HashMap}; use std::fmt; use std::fmt::Formatter; +use std::time::Duration; use arrow::datatypes::SchemaRef; @@ -28,9 +29,11 @@ use datafusion_common::display::{GraphvizBuilder, PlanType, StringifiedPlan}; use datafusion_expr::display_schema; use datafusion_physical_expr::LexOrdering; -use crate::metrics::{MetricCategory, MetricType}; +use crate::metrics::{MetricCategory, MetricType, MetricValue}; use crate::render_tree::RenderTree; +use crate::statistics::{StatisticsArgs, StatisticsContext}; + use super::{ExecutionPlan, ExecutionPlanVisitor, accept}; /// Options for controlling how each [`ExecutionPlan`] should format itself @@ -75,7 +78,7 @@ pub enum DisplayFormatType { /// │ partition_sizes: [1] │ /// │ Parquet │ /// └───────────────────────────┘ - /// ``` + /// ``` TreeRender, } @@ -126,8 +129,22 @@ pub struct DisplayableExecutionPlan<'a> { /// Optional filter by semantic category (rows / bytes / timing). /// `None` means show all categories; `Some(vec![])` means plan-only. metric_categories: Option>, + /// Optional filter by metric names. Only metric names in this list + /// will be rendered. + metric_names: Option>, // (TreeRender) Maximum total width of the rendered tree tree_maximum_render_width: usize, + /// Optional summary totals (currently only used by `pgjson`) — the total + /// row count and wall-clock duration of the `AnalyzeExec` execution. + summary: Option, +} + +/// Summary information attached to the root of an `EXPLAIN ANALYZE` +/// pgjson render. +#[derive(Debug, Clone, Copy)] +struct AnalyzeSummary { + total_rows: Option, + duration: Option, } impl<'a> DisplayableExecutionPlan<'a> { @@ -145,7 +162,9 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, + summary: None, } } @@ -160,7 +179,9 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, + summary: None, } } @@ -175,7 +196,9 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: false, metric_types: Self::default_metric_types(), metric_categories: None, + metric_names: None, tree_maximum_render_width: 240, + summary: None, } } @@ -217,12 +240,39 @@ impl<'a> DisplayableExecutionPlan<'a> { self } + /// Specify which metric names to include. + /// + /// - An empty vector means plan-only — suppress all metrics. + /// - `vec!["metric_1"]` means show only the metric named `metric_1`. + /// + /// Name filtering is intersected with other types of filters, like metric + /// category and metric type. + pub fn set_metric_names(mut self, metric_names: Vec) -> Self { + self.metric_names = Some(metric_names); + self + } + /// Set the maximum render width for the tree format pub fn set_tree_maximum_render_width(mut self, width: usize) -> Self { self.tree_maximum_render_width = width; self } + /// Attach an `EXPLAIN ANALYZE` summary (total output rows and duration) + /// to the rendered output. Currently only used by [`Self::pgjson`], which + /// serializes the summary alongside the root plan object. + pub fn set_summary( + mut self, + total_rows: Option, + duration: Option, + ) -> Self { + self.summary = Some(AnalyzeSummary { + total_rows, + duration, + }); + self + } + /// Return a `format`able structure that produces a single line /// per node. /// @@ -247,6 +297,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -259,6 +310,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), }; accept(self.plan, &mut visitor) } @@ -271,6 +323,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -292,6 +345,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { fn fmt(&self, f: &mut Formatter) -> fmt::Result { @@ -304,6 +358,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: self.show_statistics, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), graphviz_builder: GraphvizBuilder::default(), parents: Vec::new(), }; @@ -323,6 +378,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_statistics: self.show_statistics, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -349,6 +405,78 @@ impl<'a> DisplayableExecutionPlan<'a> { } } + /// Returns a `format`able structure that produces PostgreSQL-style JSON + /// output, mirroring the logical-plan pgjson format. + /// + /// Each node is rendered as a JSON object with: + /// - `"Node Type"` — `ExecutionPlan::name()` + /// - `"Details"` — the one-line `DisplayAs::Default` rendering + /// - `"Output"` — schema column names (when `set_show_schema(true)`) + /// - `"Actual Rows"` / `"Actual Total Time"` — PG-canonical metric keys + /// populated from `output_rows` / `elapsed_compute` when available + /// - `"Extras"` — remaining metrics keyed by DataFusion metric name + /// - `"Plans"` — array of child nodes + /// + /// When a summary has been set via [`Self::set_summary`], `"Total Rows"` + /// and `"Duration"` fields are attached at the root. + pub fn pgjson(&self, verbose: bool) -> impl fmt::Display + 'a { + struct Wrapper<'a> { + plan: &'a dyn ExecutionPlan, + verbose: bool, + show_metrics: ShowMetrics, + show_schema: bool, + metric_types: Vec, + metric_categories: Option>, + metric_names: Option>, + summary: Option, + } + impl fmt::Display for Wrapper<'_> { + fn fmt(&self, f: &mut Formatter) -> fmt::Result { + let mut visitor = PgJsonExecutionPlanVisitor { + verbose: self.verbose, + show_metrics: self.show_metrics, + show_schema: self.show_schema, + metric_types: &self.metric_types, + metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), + objects: HashMap::new(), + parent_ids: Vec::new(), + next_id: 0, + root: None, + }; + accept(self.plan, &mut visitor).map_err(|_| fmt::Error)?; + let root = visitor.root.ok_or(fmt::Error)?; + let mut root_entry = serde_json::json!({ "Plan": root }); + if let Some(summary) = self.summary { + if let Some(total_rows) = summary.total_rows { + root_entry["Total Rows"] = serde_json::Value::from(total_rows); + } + if let Some(duration) = summary.duration { + root_entry["Duration"] = + serde_json::Value::from(format!("{duration:?}")); + } + } + let doc = serde_json::Value::Array(vec![root_entry]); + write!( + f, + "{}", + serde_json::to_string_pretty(&doc).map_err(|_| fmt::Error)? + ) + } + } + + Wrapper { + plan: self.inner, + verbose, + show_metrics: self.show_metrics, + show_schema: self.show_schema, + metric_types: self.metric_types.clone(), + metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), + summary: self.summary, + } + } + /// Return a single-line summary of the root of the plan /// Example: `ProjectionExec: expr=[a@0 as a]`. pub fn one_line(&self) -> impl fmt::Display + 'a { @@ -359,6 +487,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: bool, metric_types: Vec, metric_categories: Option>, + metric_names: Option>, } impl fmt::Display for Wrapper<'_> { @@ -372,6 +501,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: &self.metric_types, metric_categories: self.metric_categories.as_deref(), + metric_names: self.metric_names.as_deref(), }; visitor.pre_visit(self.plan)?; Ok(()) @@ -385,6 +515,7 @@ impl<'a> DisplayableExecutionPlan<'a> { show_schema: self.show_schema, metric_types: self.metric_types.clone(), metric_categories: self.metric_categories.clone(), + metric_names: self.metric_names.clone(), } } @@ -443,6 +574,8 @@ struct IndentVisitor<'a, 'b> { metric_types: &'a [MetricType], /// Optional filter by semantic category (rows / bytes / timing). metric_categories: Option<&'a [MetricCategory]>, + /// Optional filter by metric name. + metric_names: Option<&'a [String]>, } impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { @@ -462,6 +595,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } write!(self.f, ", metrics=[{metrics}]")?; } else { write!(self.f, ", metrics=[]")?; @@ -473,6 +609,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } write!(self.f, ", metrics=[{metrics}]")?; } else { write!(self.f, ", metrics=[]")?; @@ -480,7 +619,9 @@ impl ExecutionPlanVisitor for IndentVisitor<'_, '_> { } } if self.show_statistics { - let stats = plan.partition_statistics(None).map_err(|_e| fmt::Error)?; + let stats = StatisticsContext::new() + .compute(plan, &StatisticsArgs::new()) + .map_err(|_e| fmt::Error)?; write!(self.f, ", statistics=[{stats}]")?; } if self.show_schema { @@ -513,6 +654,8 @@ struct GraphvizVisitor<'a, 'b> { metric_types: &'a [MetricType], /// Optional filter by semantic category metric_categories: Option<&'a [MetricCategory]>, + /// Optional filter by metric name. + metric_names: Option<&'a [String]>, graphviz_builder: GraphvizBuilder, /// Used to record parent node ids when visiting a plan. @@ -557,6 +700,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } format!("metrics=[{metrics}]") } else { "metrics=[]".to_string() @@ -568,6 +714,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { if let Some(cats) = self.metric_categories { metrics = metrics.filter_by_categories(cats); } + if let Some(names) = self.metric_names { + metrics = metrics.filter_by_names(names); + } format!("metrics=[{metrics}]") } else { "metrics=[]".to_string() @@ -576,7 +725,9 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { }; let statistics = if self.show_statistics { - let stats = plan.partition_statistics(None).map_err(|_e| fmt::Error)?; + let stats = StatisticsContext::new() + .compute(plan, &StatisticsArgs::new()) + .map_err(|_e| fmt::Error)?; format!("statistics=[{stats}]") } else { "".to_string() @@ -611,6 +762,192 @@ impl ExecutionPlanVisitor for GraphvizVisitor<'_, '_> { } } +/// Formats physical plans into PostgreSQL-style JSON output with live +/// per-operator metrics. +/// +/// This visitor mirrors the logical-plan `PgJsonVisitor` in +/// `datafusion-expr`: during `pre_visit` it assembles a JSON object for the +/// current node; during `post_visit` it attaches that object into its +/// parent's `"Plans"` array (or stores it as the root). +struct PgJsonExecutionPlanVisitor<'a> { + verbose: bool, + show_metrics: ShowMetrics, + show_schema: bool, + metric_types: &'a [MetricType], + metric_categories: Option<&'a [MetricCategory]>, + metric_names: Option<&'a [String]>, + objects: HashMap, + parent_ids: Vec, + next_id: u32, + root: Option, +} + +impl PgJsonExecutionPlanVisitor<'_> { + /// Produce the one-line `DisplayAs::Default` rendering of a node. + fn one_line_details(plan: &dyn ExecutionPlan) -> String { + struct One<'b>(&'b dyn ExecutionPlan); + impl fmt::Display for One<'_> { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + self.0.fmt_as(DisplayFormatType::Default, f) + } + } + // Some operators include internal newlines; collapse them so the + // rendered JSON value stays on a single line. + format!("{}", One(plan)) + .replace('\n', " ") + .trim() + .to_string() + } + + /// Render the given `MetricValue` into the most natural `serde_json::Value` + /// we can produce: a number for simple counts/gauges/times, a float-ms for + /// `ElapsedCompute`, and a string fallback for anything else. + fn metric_value_to_json(value: &MetricValue) -> serde_json::Value { + match value { + MetricValue::OutputRows(c) => serde_json::Value::from(c.value()), + MetricValue::SpillCount(c) + | MetricValue::OutputBatches(c) + | MetricValue::SpilledRows(c) => serde_json::Value::from(c.value()), + MetricValue::SpilledBytes(c) | MetricValue::OutputBytes(c) => { + serde_json::Value::from(c.value()) + } + MetricValue::CurrentMemoryUsage(g) => serde_json::Value::from(g.value()), + MetricValue::ElapsedCompute(t) => { + // Emit as float milliseconds to align with PG's + // `"Actual Total Time"` convention. DataFusion tracks compute + // time (summed across partitions), not wall time — visualizers + // should be read with that caveat in mind. + let ms = (t.value() as f64) / 1_000_000.0; + serde_json::Value::from(ms) + } + MetricValue::Count { count, .. } => serde_json::Value::from(count.value()), + MetricValue::Gauge { gauge, .. } => serde_json::Value::from(gauge.value()), + MetricValue::PeakMemoryUsage { gauge, .. } => { + serde_json::Value::from(gauge.value()) + } + MetricValue::Time { time, .. } => { + let ms = (time.value() as f64) / 1_000_000.0; + serde_json::Value::from(ms) + } + // Timestamps, PruningMetrics, Ratio, Custom: fall back to Display. + other => serde_json::Value::String(format!("{other}")), + } + } + + /// Populate `"Actual Rows"`, `"Actual Total Time"`, and `"Extras"` for + /// the given node from its aggregated `MetricsSet`, honoring the same + /// filtering pipeline used by `IndentVisitor`. + fn attach_metrics(&self, plan: &dyn ExecutionPlan, object: &mut serde_json::Value) { + if matches!(self.show_metrics, ShowMetrics::None) { + return; + } + let Some(metrics) = plan.metrics() else { + return; + }; + + let metrics = match self.show_metrics { + ShowMetrics::None => return, + ShowMetrics::Aggregated => metrics + .filter_by_metric_types(self.metric_types) + .aggregate_by_name() + .sorted_for_display() + .timestamps_removed(), + ShowMetrics::Full => metrics.filter_by_metric_types(self.metric_types), + }; + let metrics = if let Some(cats) = self.metric_categories { + metrics.filter_by_categories(cats) + } else { + metrics + }; + + let metrics = if let Some(names) = self.metric_names { + metrics.filter_by_names(names) + } else { + metrics + }; + + // Build the Extras bucket, while extracting PG-canonical keys to the + // top level. + let mut extras = serde_json::Map::new(); + for metric in metrics.iter() { + let value = metric.value(); + match value { + MetricValue::OutputRows(c) => { + object["Actual Rows"] = serde_json::Value::from(c.value()); + } + MetricValue::ElapsedCompute(t) => { + let ms = (t.value() as f64) / 1_000_000.0; + object["Actual Total Time"] = serde_json::Value::from(ms); + } + _ => { + extras.insert( + value.name().to_string(), + Self::metric_value_to_json(value), + ); + } + } + } + if !extras.is_empty() { + object["Extras"] = serde_json::Value::Object(extras); + } + } +} + +impl ExecutionPlanVisitor for PgJsonExecutionPlanVisitor<'_> { + type Error = fmt::Error; + + fn pre_visit(&mut self, plan: &dyn ExecutionPlan) -> Result { + let id = self.next_id; + self.next_id += 1; + + // Build fields in reading order: Node Type, Details, (schema), + // (metrics), Plans last — so the JSON output reads top-down like a + // PostgreSQL plan. + let mut object = serde_json::json!({ + "Node Type": plan.name(), + "Details": Self::one_line_details(plan), + }); + + if self.show_schema || self.verbose { + // Always include output columns when a caller asked for schema; + // also include them in verbose mode so the pgjson output mirrors + // the extra context shown by indent's verbose flag. + let columns: Vec = plan + .schema() + .fields() + .iter() + .map(|f| serde_json::Value::String(f.name().to_string())) + .collect(); + object["Output"] = serde_json::Value::Array(columns); + } + + self.attach_metrics(plan, &mut object); + + object["Plans"] = serde_json::Value::Array(vec![]); + + self.objects.insert(id, object); + self.parent_ids.push(id); + Ok(true) + } + + fn post_visit(&mut self, _plan: &dyn ExecutionPlan) -> Result { + let id = self.parent_ids.pop().ok_or(fmt::Error)?; + let current = self.objects.remove(&id).ok_or(fmt::Error)?; + + if let Some(parent_id) = self.parent_ids.last() { + let parent = self.objects.get_mut(parent_id).ok_or(fmt::Error)?; + let plans = parent + .get_mut("Plans") + .and_then(|p| p.as_array_mut()) + .ok_or(fmt::Error)?; + plans.push(current); + } else { + self.root = Some(current); + } + Ok(true) + } +} + /// This module implements a tree-like art renderer for execution plans, /// based on DuckDB's implementation: /// @@ -727,7 +1064,7 @@ impl TreeRenderVisitor<'_, '_> { continue; } // there are nodes next to this, fill the space - write!(self.f, "{}", &" ".repeat(Self::NODE_RENDER_WIDTH))?; + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; } } writeln!(self.f)?; @@ -917,13 +1254,13 @@ impl TreeRenderVisitor<'_, '_> { )?; write!(self.f, "{}", Self::RDCORNER)?; } else if root.has_node(x, y + 1) { - write!(self.f, "{}", &" ".repeat(Self::NODE_RENDER_WIDTH / 2))?; + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH / 2))?; write!(self.f, "{}", Self::VERTICAL)?; if has_adjacent_nodes || Self::should_render_whitespace(root, x, y) { - write!(self.f, "{}", &" ".repeat(Self::NODE_RENDER_WIDTH / 2))?; + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH / 2))?; } } else if has_adjacent_nodes || Self::should_render_whitespace(root, x, y) { - write!(self.f, "{}", &" ".repeat(Self::NODE_RENDER_WIDTH))?; + write!(self.f, "{}", " ".repeat(Self::NODE_RENDER_WIDTH))?; } } writeln!(self.f)?; @@ -1173,7 +1510,11 @@ mod tests { use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_physical_expr::PhysicalExpr; - use crate::{DisplayAs, ExecutionPlan, PlanProperties}; + use crate::statistics::StatisticsArgs; + use crate::{ + ChildrenPropertiesMode, DisplayAs, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, + }; use super::DisplayableExecutionPlan; @@ -1207,20 +1548,31 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _: usize, @@ -1229,11 +1581,12 @@ mod tests { todo!() } - fn partition_statistics( + fn statistics_from_inputs( &self, - partition: Option, + _input_stats: &[Arc], + args: &StatisticsArgs, ) -> Result> { - if partition.is_some() { + if args.partition().is_some() { return Ok(Arc::new(Statistics::new_unknown(self.schema().as_ref()))); } match self { @@ -1285,4 +1638,249 @@ mod tests { fn test_display_when_stats_ok_with_show_stats() { test_stats_display(TestStatsExecPlan::Ok, false); } + + mod pgjson { + use std::sync::Arc; + use std::time::Duration; + + use arrow::datatypes::{DataType, Field, Schema}; + use insta::assert_snapshot; + + use super::super::DisplayableExecutionPlan; + use crate::empty::EmptyExec; + use crate::filter::FilterExec; + use crate::projection::ProjectionExec; + use crate::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; + use datafusion_physical_expr::expressions::{binary, col, lit}; + use datafusion_physical_expr::{Partitioning, PhysicalExpr}; + + fn sample_plan() -> Arc { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let empty = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let predicate = binary( + col("a", &schema).unwrap(), + datafusion_expr::Operator::Gt, + lit(5i32), + &schema, + ) + .unwrap(); + let filter = Arc::new(FilterExec::try_new(predicate, empty).unwrap()); + let proj_expr: Vec<(Arc, String)> = + vec![(col("a", &schema).unwrap(), "a".to_string())]; + let _ = Partitioning::UnknownPartitioning(1); + Arc::new(ProjectionExec::try_new(proj_expr, filter).unwrap()) + } + + #[test] + fn pgjson_renders_plan_without_metrics() { + let plan = sample_plan(); + let out = DisplayableExecutionPlan::new(plan.as_ref()) + .pgjson(false) + .to_string(); + let value: serde_json::Value = serde_json::from_str(&out).unwrap(); + // Root is an array with one {"Plan": ...} entry. + let root = value + .as_array() + .expect("root array") + .first() + .expect("root entry") + .get("Plan") + .expect("plan object"); + assert_eq!(root["Node Type"].as_str(), Some("ProjectionExec")); + assert!(root.get("Actual Rows").is_none()); + assert!(root.get("Extras").is_none()); + let plans = root["Plans"].as_array().expect("Plans array"); + assert_eq!(plans.len(), 1); + assert_eq!(plans[0]["Node Type"].as_str(), Some("FilterExec")); + } + + #[test] + fn pgjson_emits_pg_canonical_metric_keys() { + use crate::metrics::{Count, Metric, MetricValue, MetricsSet, Time}; + use crate::{DisplayFormatType, ExecutionPlan, PlanProperties}; + use datafusion_common::Result; + use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + + // Wrap `sample_plan()` with an adapter node that exposes a + // hand-crafted `MetricsSet` so we can assert the PG key mapping + // without running anything. + #[derive(Debug)] + struct WithMetrics { + inner: Arc, + metrics: MetricsSet, + } + impl crate::DisplayAs for WithMetrics { + fn fmt_as( + &self, + _t: DisplayFormatType, + f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + write!(f, "WithMetrics") + } + } + impl ExecutionPlan for WithMetrics { + fn name(&self) -> &'static str { + "WithMetrics" + } + fn properties(&self) -> &Arc { + self.inner.properties() + } + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut( + &Arc, + ) -> Result< + datafusion_common::tree_node::TreeNodeRecursion, + >, + ) -> Result + { + Ok(datafusion_common::tree_node::TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + unimplemented!() + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( + &self, + _: usize, + _: Arc, + ) -> Result { + unimplemented!() + } + fn metrics(&self) -> Option { + Some(self.metrics.clone()) + } + } + + let mut metrics = MetricsSet::new(); + let rows = Count::new(); + rows.add(42); + metrics.push(Arc::new(Metric::new(MetricValue::OutputRows(rows), None))); + let elapsed = Time::new(); + elapsed.add_duration(Duration::from_millis(5)); + metrics.push(Arc::new(Metric::new( + MetricValue::ElapsedCompute(elapsed), + None, + ))); + let batches = Count::new(); + batches.add(7); + metrics.push(Arc::new(Metric::new( + MetricValue::OutputBatches(batches), + None, + ))); + + let plan: Arc = Arc::new(WithMetrics { + inner: sample_plan(), + metrics, + }); + + let out = DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .pgjson(false) + .to_string(); + let value: serde_json::Value = serde_json::from_str(&out).unwrap(); + let root = value[0].get("Plan").expect("plan"); + assert_eq!(root["Actual Rows"].as_u64(), Some(42)); + assert_eq!(root["Actual Total Time"].as_f64(), Some(5.0)); + assert_eq!(root["Extras"]["output_batches"].as_u64(), Some(7)); + + let metric_names = vec!["output_rows".to_string()]; + for rendered in [ + DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_metric_names(metric_names.clone()) + .indent(false) + .to_string(), + DisplayableExecutionPlan::with_full_metrics(plan.as_ref()) + .set_metric_names(metric_names.clone()) + .indent(false) + .to_string(), + DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_metric_names(metric_names.clone()) + .graphviz() + .to_string(), + DisplayableExecutionPlan::with_full_metrics(plan.as_ref()) + .set_metric_names(metric_names.clone()) + .graphviz() + .to_string(), + ] { + assert!(rendered.contains("output_rows")); + assert!(!rendered.contains("elapsed_compute")); + assert!(!rendered.contains("output_batches")); + } + + let out = DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_metric_names(metric_names) + .pgjson(false) + .to_string(); + let value: serde_json::Value = serde_json::from_str(&out).unwrap(); + let root = value[0].get("Plan").expect("plan"); + assert_eq!(root["Actual Rows"].as_u64(), Some(42)); + assert!(root.get("Actual Total Time").is_none()); + assert!(root.get("Extras").is_none()); + } + + #[test] + fn pgjson_includes_summary_when_set() { + let plan = sample_plan(); + let out = DisplayableExecutionPlan::with_metrics(plan.as_ref()) + .set_summary(Some(42), Some(Duration::from_millis(7))) + .pgjson(false) + .to_string(); + let value: serde_json::Value = serde_json::from_str(&out).unwrap(); + let entry = &value.as_array().unwrap()[0]; + assert_eq!(entry["Total Rows"].as_u64(), Some(42)); + assert!(entry["Duration"].is_string()); + } + + #[test] + fn pgjson_snapshot_of_sample_plan() { + let plan = sample_plan(); + let out = DisplayableExecutionPlan::new(plan.as_ref()) + .pgjson(false) + .to_string(); + // This snapshot assumes `serde_json` is built with the + // `preserve_order` feature (enabled via this crate's dev-deps). + assert_snapshot!(out, @r#" + [ + { + "Plan": { + "Node Type": "ProjectionExec", + "Details": "ProjectionExec: expr=[a@0 as a]", + "Plans": [ + { + "Node Type": "FilterExec", + "Details": "FilterExec: a@0 > 5", + "Plans": [ + { + "Node Type": "EmptyExec", + "Details": "EmptyExec", + "Plans": [] + } + ] + } + ] + } + } + ] + "#); + } + } } diff --git a/datafusion/physical-plan/src/distribution_requirements.rs b/datafusion/physical-plan/src/distribution_requirements.rs new file mode 100644 index 0000000000000..6405b1f121ef7 --- /dev/null +++ b/datafusion/physical-plan/src/distribution_requirements.rs @@ -0,0 +1,359 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Input distribution requirements for physical execution plans. + +use datafusion_common::{Result, internal_err}; +use datafusion_physical_expr::{Distribution, Partitioning, PartitioningSatisfaction}; + +use crate::execution_plan::{ExecutionPlan, ExecutionPlanProperties, InvariantLevel}; + +/// Distribution requirements for an [`ExecutionPlan`]'s inputs. +/// +/// [`InputDistributionRequirements`] describes what distribution an operator +/// requires from each child. +/// +/// - [`Self::new`] describes independent per-child requirements. +/// - [`Self::co_partitioned`] additionally requires child partitions with the +/// same index to cover compatible key ranges. +/// +/// For a single-input aggregate: +/// +/// ```text +/// AggregateExec +/// child 0 requirement: KeyPartitioned(group_exprs) +/// ``` +/// +/// each input partition can aggregate its own key domain independently. +/// +/// For a partitioned join: +/// +/// ```text +/// HashJoinExec +/// child 0 requirement: KeyPartitioned(left_keys) +/// child 1 requirement: KeyPartitioned(right_keys) +/// +/// partition 0: join(left partition 0, right partition 0) +/// partition 1: join(left partition 1, right partition 1) +/// partition 2: join(left partition 2, right partition 2) +/// ``` +/// +/// each child must satisfy its own key requirement. In addition, matching +/// partition indexes must be safe to process together. +#[non_exhaustive] +#[derive(Debug, Clone)] +pub struct InputDistributionRequirements { + /// Per-child distribution requirements, indexed by child position. + children: Vec, + /// Child indexes that must also have compatible partition layouts. + co_partitioned: Option>, +} + +/// Options for checking child distribution satisfaction. +#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] +pub struct ChildSatisfactionOptions { + allow_subset: bool, +} + +impl ChildSatisfactionOptions { + /// Create default satisfaction options. + pub fn new() -> Self { + Self::default() + } + + /// Allow a child partitioning whose key expressions are a subset of the + /// required key expressions to satisfy the requirement. + pub fn with_allow_subset(mut self, allow_subset: bool) -> Self { + self.allow_subset = allow_subset; + self + } + + /// Whether subset satisfaction is enabled. + pub fn allow_subset(&self) -> bool { + self.allow_subset + } +} + +impl InputDistributionRequirements { + /// Create independent per-child requirements. + pub fn new(per_child: Vec) -> Self { + let children = per_child + .into_iter() + .map(|distribution| ChildDistributionRequirement { distribution }) + .collect(); + + Self { + children, + co_partitioned: None, + } + } + + /// Create a requirement that all children are co-partitioned. + /// + /// Each child must satisfy its own [`Distribution`]. Matching partition + /// indexes are processed together: + /// + /// ```text + /// left: Range(left.a ASC, split_points=[10, 20]) + /// right: Range(right.x ASC, split_points=[10, 20]) + /// + /// partition 0 from both sides contains keys before 10 + /// partition 1 from both sides contains keys in [10, 20) + /// partition 2 from both sides contains keys at/after 20 + /// ``` + /// + /// If the split points differ, partition `i` from one side no longer covers + /// the same key range as partition `i` from the other side. + pub fn co_partitioned(per_child: Vec) -> Self { + debug_assert!( + per_child.len() >= 2, + "co-partitioned distribution requirements need at least two children" + ); + let co_partitioned = (0..per_child.len()).collect(); + let mut result = Self::new(per_child); + result.co_partitioned = Some(co_partitioned); + result + } + + /// Return the per-child distribution requirements. + pub fn per_child_distributions( + &self, + ) -> impl ExactSizeIterator + '_ { + self.children.iter().map(|child| &child.distribution) + } + + /// Return the distribution requirement for a child. + pub fn child_distribution(&self, child_idx: usize) -> Option<&Distribution> { + self.children + .get(child_idx) + .map(|child| &child.distribution) + } + + /// Return the per-child distribution requirements. + /// + /// WARNING: This intentionally drops any grouped relationship. + pub fn into_per_child(self) -> Vec { + self.children + .into_iter() + .map(|child| child.distribution) + .collect() + } + + /// Returns how a child satisfies its distribution requirement. + /// + /// This preserves the requirement set's satisfaction policy. + pub fn child_satisfaction( + &self, + child_idx: usize, + child: &dyn ExecutionPlan, + options: ChildSatisfactionOptions, + ) -> Result { + let Some(requirement) = self.children.get(child_idx) else { + return internal_err!( + "missing distribution requirement for child {child_idx}" + ); + }; + + Ok(child.output_partitioning().satisfaction( + &requirement.distribution, + child.equivalence_properties(), + options.allow_subset(), + )) + } + + /// Return child indexes whose co-partitioning requirements are + /// unsatisfied by the provided candidate children. + /// + /// Independent per-child requirements are intentionally ignored here, use + /// [`Self::child_satisfaction`] for those checks. An empty result means all + /// co-partitioning requirements are satisfied. + #[doc(hidden)] + pub fn unsatisfied_co_partitioned_children( + &self, + plan_name: &str, + children: &[&dyn ExecutionPlan], + ) -> Result> { + self.validate_shape(plan_name, children.len())?; + + let Some(co_partitioned) = &self.co_partitioned else { + return Ok(vec![]); + }; + if self.co_partitioning_satisfied(co_partitioned, children) { + return Ok(vec![]); + } + + Ok(co_partitioned.clone()) + } + + /// Validate the requirements against a plan's children. + pub(crate) fn check_invariants( + &self, + plan: &P, + check: InvariantLevel, + ) -> Result<()> { + let children = plan.children(); + self.validate_shape(plan.name(), children.len())?; + + let children = children + .into_iter() + .map(|child| child.as_ref()) + .collect::>(); + if matches!(check, InvariantLevel::Executable) + && let Some(co_partitioned) = &self.co_partitioned + && !self.co_partitioning_satisfied(co_partitioned, &children) + { + return internal_err!( + "{} requires children {:?} to be co-partitioned", + plan.name(), + co_partitioned + ); + } + + Ok(()) + } + + fn validate_shape(&self, plan_name: &str, children_len: usize) -> Result<()> { + if self.children.len() != children_len { + return internal_err!( + "{plan_name}::input_distribution_requirements returned incorrect child count: {} != {}", + self.children.len(), + children_len + ); + } + + if let Some(co_partitioned) = &self.co_partitioned { + if co_partitioned.len() < 2 { + return internal_err!( + "{plan_name} has invalid co-partitioning requirement: at least two children are required" + ); + } + let mut seen = vec![false; self.children.len()]; + for &child in co_partitioned { + validate_child_index(plan_name, child, self.children.len(), &mut seen)?; + if matches!( + self.children[child].distribution, + Distribution::UnspecifiedDistribution + ) { + return internal_err!( + "{plan_name} has invalid co-partitioning requirement: child {child} has unspecified distribution" + ); + } + } + } + + Ok(()) + } + + fn co_partitioning_satisfied( + &self, + co_partitioned: &[usize], + children: &[&dyn ExecutionPlan], + ) -> bool { + let first_idx = co_partitioned[0]; + let first_requirement = &self.children[first_idx]; + let first = children[first_idx]; + let first_partitioning = first.output_partitioning(); + + if !first_partitioning + .satisfaction( + &first_requirement.distribution, + first.equivalence_properties(), + false, + ) + .is_satisfied() + { + return false; + } + + for &child_idx in co_partitioned.iter().skip(1) { + let requirement = &self.children[child_idx]; + let child = children[child_idx]; + if !child + .output_partitioning() + .satisfaction( + &requirement.distribution, + child.equivalence_properties(), + false, + ) + .is_satisfied() + || !compatible_co_partitioning_layout( + first_partitioning, + child.output_partitioning(), + ) + { + return false; + } + } + + true + } +} + +/// A distribution requirement for a single child. +#[derive(Debug, Clone)] +struct ChildDistributionRequirement { + distribution: Distribution, +} + +fn validate_child_index( + plan_name: &str, + child_idx: usize, + child_count: usize, + seen: &mut [bool], +) -> Result<()> { + if child_idx >= child_count { + return internal_err!( + "{plan_name} has invalid distribution requirement: child index {child_idx} out of bounds" + ); + } + if seen[child_idx] { + return internal_err!( + "{plan_name} has invalid distribution requirement: child {child_idx} appears more than once" + ); + } + seen[child_idx] = true; + Ok(()) +} + +fn compatible_co_partitioning_layout( + first_partitioning: &Partitioning, + other_partitioning: &Partitioning, +) -> bool { + if first_partitioning.partition_count() == 1 + && other_partitioning.partition_count() == 1 + { + return true; + } + + if first_partitioning.partition_count() != other_partitioning.partition_count() { + return false; + } + + match (first_partitioning, other_partitioning) { + (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) => true, + (Partitioning::Range(left), Partitioning::Range(right)) => { + left.split_points() == right.split_points() + && left.ordering().len() == right.ordering().len() + && left + .ordering() + .iter() + .zip(right.ordering()) + .all(|(left, right)| left.options == right.options) + } + _ => false, + } +} diff --git a/datafusion/physical-plan/src/empty.rs b/datafusion/physical-plan/src/empty.rs index 8103695ad08fa..dd08ff36a9d88 100644 --- a/datafusion/physical-plan/src/empty.rs +++ b/datafusion/physical-plan/src/empty.rs @@ -20,7 +20,10 @@ use std::sync::Arc; use crate::memory::MemoryStream; -use crate::{DisplayAs, PlanProperties, SendableRecordBatchStream, Statistics}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, +}; use crate::{ DisplayFormatType, ExecutionPlan, Partitioning, execution_plan::{Boundedness, EmissionType}, @@ -35,6 +38,7 @@ use datafusion_execution::TaskContext; use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; use crate::execution_plan::SchedulingType; +use crate::statistics::StatisticsArgs; use log::trace; /// Execution plan for empty relation with produce_one_row=false @@ -121,18 +125,29 @@ impl ExecutionPlan for EmptyExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -159,8 +174,12 @@ impl ExecutionPlan for EmptyExec { )?)) } - fn partition_statistics(&self, partition: Option) -> Result> { - if let Some(partition) = partition { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + if let Some(partition) = args.partition() { assert_or_internal_err!( partition < self.partitions, "EmptyExec invalid partition {} (expected less than {})", @@ -188,14 +207,62 @@ impl ExecutionPlan for EmptyExec { Ok(Arc::new(stats)) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Empty( + protobuf::EmptyExecNode { + schema: Some(schema), + partitions: self + .properties() + .output_partitioning() + .partition_count() as u32, + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl EmptyExec { + /// Reconstruct an [`EmptyExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let empty = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Empty, + "EmptyExec", + ); + let schema = empty.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "EmptyExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(arrow::datatypes::Schema::try_from(schema)?); + // A zero (absent) partition count comes from a plan encoded before the + // field existed, which always meant a single partition. + let partitions = empty.partitions.max(1) as usize; + Ok(Arc::new(EmptyExec::new(schema).with_partitions(partitions))) + } } #[cfg(test)] mod tests { use super::*; use crate::common; + use crate::execution_plan::replace_children_if_necessary; use crate::test; - use crate::with_new_children_if_necessary; #[tokio::test] async fn empty() -> Result<()> { @@ -218,7 +285,7 @@ mod tests { let schema = test::aggr_test_schema(); let empty = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let empty2 = with_new_children_if_necessary( + let empty2 = replace_children_if_necessary( Arc::clone(&empty) as Arc, vec![], )?; @@ -226,7 +293,7 @@ mod tests { let too_many_kids = vec![empty2]; assert!( - with_new_children_if_necessary(empty, too_many_kids).is_err(), + replace_children_if_necessary(empty, too_many_kids).is_err(), "expected error when providing list of kids" ); Ok(()) diff --git a/datafusion/physical-plan/src/execution_plan.rs b/datafusion/physical-plan/src/execution_plan.rs index 1a67ea0ded11b..a4d081b3d9e75 100644 --- a/datafusion/physical-plan/src/execution_plan.rs +++ b/datafusion/physical-plan/src/execution_plan.rs @@ -16,6 +16,7 @@ // under the License. pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; +use crate::distribution_requirements::InputDistributionRequirements; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, @@ -34,12 +35,14 @@ pub use datafusion_common::utils::project_schema; pub use datafusion_common::{ColumnStatistics, Statistics, internal_err}; pub use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; pub use datafusion_expr::{Accumulator, ColumnarValue}; +use datafusion_physical_expr::projection::ProjectionExpr; pub use datafusion_physical_expr::window::WindowExpr; pub use datafusion_physical_expr::{ Distribution, Partitioning, PhysicalExpr, expressions, }; use std::any::Any; +use std::collections::HashSet; use std::fmt::Debug; use std::sync::{Arc, LazyLock}; @@ -47,6 +50,9 @@ use crate::coalesce_partitions::CoalescePartitionsExec; use crate::display::DisplayableExecutionPlan; use crate::metrics::MetricsSet; use crate::projection::ProjectionExec; +use crate::repartition::RepartitionExec; +use crate::sorts::sort_preserving_merge::SortPreservingMergeExec; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; use arrow::array::{Array, RecordBatch}; @@ -117,6 +123,30 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } } + /// Returns the plan that provides this plan's public + /// [`ExecutionPlan`] downcast identity. + /// + /// This hook is for wrapper nodes that delegate their public downcast + /// identity to another plan while adding cross-cutting behavior such as + /// instrumentation. The default implementation returns `None`, meaning this + /// plan's concrete type is used for type introspection. + /// + /// Most `ExecutionPlan` implementations should use the default `None`; + /// override this only for wrapper plans that intentionally delegate their + /// public downcast identity to another plan. + /// + /// The `is` and `downcast_ref` helpers follow the returned delegate instead + /// of checking the current concrete type, making intermediate delegating + /// wrappers invisible to normal downcast-based inspection. + /// + /// Implementations that opt in should return the delegate plan, not `self`. + /// + /// This is independent from [`Self::children`] and should not be used for + /// plan traversal or optimizer rewrites. + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { + None + } + /// Get the schema for this execution plan fn schema(&self) -> SchemaRef { Arc::clone(self.properties().schema()) @@ -139,12 +169,46 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { check_default_invariants(self, check) } - /// Specifies the data distribution requirements for all the - /// children for this `ExecutionPlan`, By default it's [[Distribution::UnspecifiedDistribution]] for each child, + /// Returns the dynamic expressions produced by this plan node. + /// + /// A dynamic expression is produced when this node updates or completes its + /// runtime state during execution. Expressions that this node only consumes + /// must not be returned. This method is shallow and does not include dynamic + /// expressions produced by child plans. + /// + /// Each returned expression must have a [`PhysicalExpr::expression_id`] + /// since all dynamic expressions such as [`DynamicFilterPhysicalExpr`] + /// have an expression id. + /// + /// [`DynamicFilterPhysicalExpr`]: datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr + fn dynamic_expressions_produced(&self) -> Vec> { + Vec::new() + } + + /// Specifies simple per-child input distribution requirements. + /// + /// Deprecated: override [`Self::input_distribution_requirements`] instead. + /// + /// By default, each child has [`Distribution::UnspecifiedDistribution`]. + #[deprecated(since = "55.0.0", note = "Use input_distribution_requirements")] fn required_input_distribution(&self) -> Vec { vec![Distribution::UnspecifiedDistribution; self.children().len()] } + /// Specifies the input distribution requirements for this plan. + /// + /// The default implementation wraps [`Self::required_input_distribution`]. + /// Override this method for richer requirements, such as allowing alternate + /// satisfaction policies or requiring multiple children to be co-partitioned. + /// See [`InputDistributionRequirements`] for details. + fn input_distribution_requirements(&self) -> InputDistributionRequirements { + #[expect( + deprecated, + reason = "compatibility shim for external ExecutionPlan implementations" + )] + InputDistributionRequirements::new(self.required_input_distribution()) + } + /// Specifies the ordering required for all of the children of this /// `ExecutionPlan`. /// @@ -191,8 +255,8 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { fn benefits_from_input_partitioning(&self) -> Vec { // By default try to maximize parallelism with more CPUs if // possible - self.required_input_distribution() - .into_iter() + self.input_distribution_requirements() + .per_child_distributions() .map(|dist| !matches!(dist, Distribution::SinglePartition)) .collect() } @@ -203,18 +267,44 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// joins). fn children(&self) -> Vec<&Arc>; - /// Apply a closure `f` to each expression (non-recursively) in the current - /// physical plan node. This does not include expressions in any children. + /// Returns a clone of the existing plan with the children replaced, + /// skipping recomputation of plan properties when the options indicate + /// the new children's properties are unchanged. + /// + /// Callers should typically call [`replace_children_if_necessary`] and + /// not invoke this method directly. + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + #[expect(deprecated)] + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.with_new_children_and_same_properties(children) + } + ChildrenPropertiesMode::Recompute => self.with_new_children(children), + } + } + + /// Apply a closure `f` to each root expression that this node owns and uses + /// during execution, either by evaluating it or updating it dynamically. /// - /// The closure `f` is applied to expressions in the order they appear in the plan. - /// The closure can return `TreeNodeRecursion::Continue` to continue visiting, - /// `TreeNodeRecursion::Stop` to stop visiting immediately, or `TreeNodeRecursion::Jump` - /// to skip any remaining expressions (though typically all expressions are visited). + /// An expression must not be visited solely because it describes an input or + /// output property, such as cached ordering, partitioning, or equivalence + /// metadata. However, these may be traversed indirectly. For example, + /// `RepartitionExec` visits the partitioning expressions it evaluates and + /// `SortExec` visits the sort expressions it evaluates to order rows. + /// + /// This method is shallow: it must not visit expression children or expressions + /// owned by child execution plans. + /// + /// Similarly to other [`TreeNode`] APIs, the closure can return + /// [`TreeNodeRecursion::Stop`] to stop iteration, otherwise iteration + /// should continue. Note that [`TreeNodeRecursion::Continue`] and + /// [`TreeNodeRecursion::Jump`] are equivalent because this method is not + /// recursive. /// - /// The expressions visited do not necessarily represent or even contribute - /// to the output schema of this node. For example, `FilterExec` visits the - /// filter predicate even though the output of a Filter has the same columns - /// as the input. /// /// # Example Usage /// ``` @@ -234,64 +324,140 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// /// # Implementation Examples /// - /// ## Node with no expressions (e.g., EmptyExec, MemoryExec) - /// ```ignore - /// fn apply_expressions( - /// &self, - /// _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, - /// ) -> Result { - /// Ok(TreeNodeRecursion::Continue) - /// } - /// ``` + /// ## Node with expressions (e.g., FilterExec, ProjectionExec) /// - /// ## Node with a single expression (e.g., FilterExec) + /// Use [`apply_expression_roots`] to implement this method. It abstracts away the + /// [`TreeNodeRecursion`] iteration from implementors. /// ```ignore /// fn apply_expressions( /// &self, - /// f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + /// f: &mut dyn FnMut(&Arc) -> Result, /// ) -> Result { - /// f(self.predicate.as_ref()) + /// apply_expression_roots([&self.predicate], f) /// } /// ``` /// - /// ## Node with multiple expressions (e.g., ProjectionExec, JoinExec) - /// - /// Use [`TreeNodeRecursion::visit_sibling`] when iterating over multiple - /// expressions. This correctly propagates [`TreeNodeRecursion::Stop`]: if - /// `f` returns `Stop` for an earlier expression, `visit_sibling` short-circuits - /// and skips the remaining ones. + /// ## Node with no expressions (e.g., EmptyExec, MemoryExec) /// ```ignore /// fn apply_expressions( /// &self, - /// f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + /// _f: &mut dyn FnMut(&Arc) -> Result, /// ) -> Result { - /// let mut tnr = TreeNodeRecursion::Continue; - /// for expr in &self.expressions { - /// tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - /// } - /// Ok(tnr) + /// Ok(TreeNodeRecursion::Continue) /// } /// ``` fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result; - /// Returns a new `ExecutionPlan` where all existing children were replaced - /// by the `children`, in order + /// Deprecated. + /// + /// DataFusion will remove this method in the future in favor of + /// [`ExecutionPlan::replace_children`]. + /// + /// Note that this method is still required by the trait; implementations + /// should delegate to [`ExecutionPlan::replace_children`] with + /// [`ChildrenPropertiesMode::Recompute`]. + /// + /// # Example Implementation + /// ``` + /// # #![allow(deprecated)] + /// # use std::fmt; + /// # use std::sync::Arc; + /// # use datafusion_common::Result; + /// # use datafusion_common::tree_node::TreeNodeRecursion; + /// # use datafusion_execution::{SendableRecordBatchStream, TaskContext}; + /// # use datafusion_physical_expr::PhysicalExpr; + /// # use datafusion_physical_plan::{ + /// # ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, + /// # PlanProperties, ReplaceChildrenOptions, + /// # }; + /// # #[derive(Debug)] + /// # struct MyExec { + /// # input: Arc, + /// # } + /// # impl DisplayAs for MyExec { + /// # fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + /// # write!(f, "MyExec") + /// # } + /// # } + /// impl ExecutionPlan for MyExec { + /// // ... + /// # fn name(&self) -> &'static str { + /// # "MyExec" + /// # } + /// # fn properties(&self) -> &Arc { + /// # self.input.properties() + /// # } + /// # fn children(&self) -> Vec<&Arc> { + /// # vec![&self.input] + /// # } + /// # fn apply_expressions( + /// # &self, + /// # _f: &mut dyn FnMut(&Arc) -> Result, + /// # ) -> Result { + /// # Ok(TreeNodeRecursion::Continue) + /// # } + /// # fn execute( + /// # &self, + /// # _partition: usize, + /// # _context: Arc, + /// # ) -> Result { + /// # unimplemented!() + /// # } + /// fn replace_children( + /// self: Arc, + /// mut children: Vec>, + /// _options: ReplaceChildrenOptions, + /// ) -> Result> { + /// Ok(Arc::new(MyExec { + /// input: children.swap_remove(0), + /// })) + /// } + /// + /// fn with_new_children( + /// self: Arc, + /// children: Vec>, + /// ) -> Result> { + /// // call into `replace_children` with `ReplaceChildrenOptions` + /// self.replace_children( + /// children, + /// ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + /// ) + /// } + /// } + /// ``` + #[deprecated( + since = "55.0.0", + note = "Use `ExecutionPlan::replace_children` with `ReplaceChildrenOptions`" + )] fn with_new_children( self: Arc, children: Vec>, ) -> Result>; + /// Deprecated. Implement [`ExecutionPlan::replace_children`] instead. + #[deprecated( + since = "55.0.0", + note = "Use `ExecutionPlan::replace_children` with `ReplaceChildrenOptions`" + )] + #[expect(deprecated)] + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.with_new_children(children) + } + /// Reset any internal state within this [`ExecutionPlan`]. /// /// This method is called when an [`ExecutionPlan`] needs to be re-executed, - /// such as in recursive queries. Unlike [`ExecutionPlan::with_new_children`], this method + /// such as in recursive queries. Unlike [`ExecutionPlan::replace_children`], this method /// ensures that any stateful components (e.g., [`DynamicFilterPhysicalExpr`]) /// are reset to their initial state. /// - /// The default implementation simply calls [`ExecutionPlan::with_new_children`] with the existing children, + /// The default implementation simply calls [`ExecutionPlan::replace_children`] with the existing children, /// effectively creating a new instance of the [`ExecutionPlan`] with the same children but without /// necessarily resetting any internal state. Implementations that require resetting of some /// internal state should override this method to provide the necessary logic. @@ -300,13 +466,16 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// it will be called from within a walk of the execution plan tree so that it will be called on each child later /// or was already called on each child. /// - /// Note to implementers: unlike [`ExecutionPlan::with_new_children`] this method does not accept new children as an argument, + /// Note to implementers: unlike [`ExecutionPlan::replace_children`] this method does not accept new children as an argument, /// thus it is expected that any cached plan properties will remain valid after the reset. /// /// [`DynamicFilterPhysicalExpr`]: datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr fn reset_state(self: Arc) -> Result> { let children = self.children().into_iter().cloned().collect(); - self.with_new_children(children) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } /// If supported, attempt to increase the partitioning of this `ExecutionPlan` to @@ -315,7 +484,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// If the `ExecutionPlan` does not support changing its partitioning, /// returns `Ok(None)` (the default). /// - /// It is the `ExecutionPlan` can increase its partitioning, but not to the + /// If the `ExecutionPlan` can increase its partitioning, but not to /// `target_partitions`, it may return an ExecutionPlan with fewer /// partitions. This might happen, for example, if each new partition would /// be too small to be efficiently processed individually. @@ -436,7 +605,7 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// partition: usize, /// context: Arc, /// ) -> Result { - /// // use functions from futures crate convert the batch into a stream + /// // use functions from futures crate to convert the batch into a stream /// let fut = futures::future::ready(Ok(self.batch.clone())); /// let stream = futures::stream::once(fut); /// Ok(Box::pin(RecordBatchStreamAdapter::new( @@ -546,9 +715,11 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { } /// Returns statistics for a specific partition of this `ExecutionPlan` node. - /// If statistics are not available, should return [`Statistics::new_unknown`] - /// (the default), not an error. - /// If `partition` is `None`, it returns statistics for the entire plan. + /// + /// Deprecated: use [`StatisticsContext::compute`] instead. + /// + /// [`StatisticsContext::compute`]: crate::statistics::StatisticsContext::compute + #[deprecated(since = "55.0.0", note = "Use StatisticsContext::compute instead")] fn partition_statistics(&self, partition: Option) -> Result> { if let Some(idx) = partition { // Validate partition index @@ -563,6 +734,47 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { Ok(Arc::new(Statistics::new_unknown(&self.schema()))) } + /// Returns statistics for a specific partition of this `ExecutionPlan` node, + /// given pre-computed child statistics. + /// + /// If statistics are not available, should return [`Statistics::new_unknown`] + /// (the default), not an error. + /// If `args.partition()` is `None`, it returns statistics for all partitions. + /// + /// Implementations should not call [`StatisticsContext::compute`] from within + /// this method; child statistics are provided via `input_stats`. + /// + /// Use [`StatisticsContext::compute`] to initiate a full plan-tree walk. + /// + /// [`StatisticsContext::compute`]: crate::statistics::StatisticsContext::compute + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + #[expect(deprecated)] + self.partition_statistics(args.partition()) + } + + /// Returns, per child, which statistics the [`StatisticsContext`] should resolve + /// before calling [`Self::statistics_from_inputs`]. + /// + /// One entry per child (same order as [`Self::children`]): [`ChildStats::At`] + /// requests the child's statistics at a partition (`None` = overall); + /// [`ChildStats::Skip`] omits a child whose statistics this node does not need + /// (a `Statistics::new_unknown` placeholder fills its `input_stats` slot). + /// + /// The default skips every child, so a node that derives nothing from its + /// children (for example one that only overrides the deprecated + /// [`Self::partition_statistics`]) triggers no child traversal. A node that reads + /// `input_stats` in [`Self::statistics_from_inputs`] must override this to declare + /// the children it uses. + /// + /// [`StatisticsContext`]: crate::statistics::StatisticsContext + fn child_stats_requests(&self, _partition: Option) -> Vec { + self.children().iter().map(|_| ChildStats::Skip).collect() + } + /// Returns `true` if a limit can be safely pushed down through this /// `ExecutionPlan` node. /// @@ -625,11 +837,17 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { /// up the plan that `DataSourceExec` can actually bind the filters. /// /// The default implementation bars all parent filters from being pushed down and adds no new filters. - /// This is the safest option, making filter pushdown opt-in on a per-node pasis. + /// This is the safest option, making filter pushdown opt-in on a per-node basis. /// /// There are two different phases in filter pushdown, which some operators may handle the same and some differently. /// Depending on the phase the operator may or may not be allowed to modify the plan. /// See [`FilterPushdownPhase`] for more details. + /// + /// Implementations must preserve the order of `parent_filters` in the + /// returned child [`FilterDescription`]: each child parent-filter result is + /// matched back to the corresponding input parent filter by position. + /// Unsupported filters should therefore be marked unsupported in place, + /// rather than removed or appended after supported filters. fn gather_filters_for_pushdown( &self, _phase: FilterPushdownPhase, @@ -789,25 +1007,176 @@ pub trait ExecutionPlan: Any + Debug + DisplayAs + Send + Sync { ) -> Option> { None } + + /// Serialize this plan to its protobuf representation, if it knows how. + /// + /// This is the `ExecutionPlan` analog of + /// [`PhysicalExpr::try_to_proto`]. + /// + /// * `Ok(None)` (the default) — "I don't serialize myself"; the caller + /// (`datafusion-proto`) falls back to the central downcast chain. Every + /// un-migrated plan keeps its existing behavior. + /// * `Ok(Some(node))` — fully serialized; the caller must not fall back. + /// * `Err(_)` — a real failure (e.g. a child failed to serialize). + /// + /// Only *self-contained* plans should override this — see [`crate::proto`] + /// for the session-dependency boundary. + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + Ok(None) + } +} + +/// Options for [`ExecutionPlan::replace_children`] +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct ReplaceChildrenOptions { + /// Describes how plan properties should be handled for the replacement + /// children. + pub children_properties: ChildrenPropertiesMode, +} + +impl ReplaceChildrenOptions { + /// Create new options for [`ExecutionPlan::replace_children`]. + pub const fn new(children_properties: ChildrenPropertiesMode) -> Self { + Self { + children_properties, + } + } +} + +/// Indicates whether the plan properties of the new children must be recomputed. +/// +/// Part of [`ReplaceChildrenOptions`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChildrenPropertiesMode { + /// The plan properties of the new children are identical to the properties + /// of the existing children, so we can skip recomputation. + Keep, + /// The plan properties of the new children are different from the properties + /// of the existing children, so we must recompute the properties from scratch. + Recompute, +} + +/// Allows a type to be treated as a reference to an +/// [`Arc`]. +/// +/// Used by [`apply_expression_roots`]. +pub trait AsPhysicalExprRef { + /// Returns the referenced physical expression. + fn as_physical_expr_ref(&self) -> &Arc; +} + +/// Allows an [`Arc`] to be treated as a reference to itself. +/// +/// This is needed because `Arc` does not implement +/// `AsRef>`. +impl AsPhysicalExprRef for Arc { + fn as_physical_expr_ref(&self) -> &Arc { + self + } +} + +/// Allows a [`ProjectionExpr`] to be treated as a reference to its +/// [`Arc`]. +impl AsPhysicalExprRef for ProjectionExpr { + fn as_physical_expr_ref(&self) -> &Arc { + self.as_ref() + } +} + +impl AsPhysicalExprRef for &T +where + T: AsPhysicalExprRef + ?Sized, +{ + fn as_physical_expr_ref(&self) -> &Arc { + (*self).as_physical_expr_ref() + } +} + +/// Applies `f` to a shallow sequence of physical expression roots. +/// +/// [`TreeNodeRecursion::Stop`] stops iteration and is returned immediately. +/// [`TreeNodeRecursion::Jump`] is normalized to [`TreeNodeRecursion::Continue`] +/// because this function does not visit expression children. +pub fn apply_expression_roots( + roots: I, + f: &mut dyn FnMut(&Arc) -> Result, +) -> Result +where + I: IntoIterator, + I::Item: AsPhysicalExprRef, +{ + for root in roots { + match f(root.as_physical_expr_ref())? { + TreeNodeRecursion::Stop => return Ok(TreeNodeRecursion::Stop), + TreeNodeRecursion::Continue | TreeNodeRecursion::Jump => {} + } + } + Ok(TreeNodeRecursion::Continue) +} + +/// Returns whether `plan` contains a physical expression with `expression_id`. +/// +/// This traverses both the execution plan and the children of each expression root +/// reported by [`ExecutionPlan::apply_expressions`]. +pub(crate) fn plan_contains_expression_id( + plan: &Arc, + expression_id: u64, +) -> Result { + let mut found = false; + plan.apply(|node| { + node.apply_expressions(&mut |root| { + root.apply(|expr| { + if expr.expression_id() == Some(expression_id) { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + Ok(found) } impl dyn ExecutionPlan { /// Returns `true` if the plan is of type `T`. /// + /// If this plan provides a [`ExecutionPlan::downcast_delegate`], delegates + /// to it. + /// /// Prefer this over `downcast_ref::().is_some()`. Works correctly when /// called on `Arc` via auto-deref. pub fn is(&self) -> bool { - (self as &dyn Any).is::() + match self.downcast_delegate() { + Some(delegate) => delegate.is::(), + None => (self as &dyn Any).is::(), + } } /// Attempts to downcast this plan to a concrete type `T`, returning `None` /// if the plan is not of that type. /// + /// If this plan provides a [`ExecutionPlan::downcast_delegate`], delegates + /// to it. + /// /// Works correctly when called on `Arc` via auto-deref, /// unlike `(&arc as &dyn Any).downcast_ref::()` which would attempt to /// downcast the `Arc` itself. pub fn downcast_ref(&self) -> Option<&T> { - (self as &dyn Any).downcast_ref() + match self.downcast_delegate() { + Some(delegate) => delegate.downcast_ref::(), + None => (self as &dyn Any).downcast_ref(), + } } } @@ -1002,25 +1371,36 @@ pub enum SchedulingType { Cooperative, } -/// Represents how an operator's `Stream` implementation generates `RecordBatch`es. +/// Represents how an operator's stream drives [`RecordBatch`] production +/// relative to downstream demand. /// -/// Most operators in DataFusion generate `RecordBatch`es when asked to do so by a call to -/// `Stream::poll_next`. This is known as demand-driven or lazy evaluation. -/// -/// Some operators like `Repartition` need to drive `RecordBatch` generation themselves though. This -/// is known as data-driven or eager evaluation. +/// This is execution-topology metadata for optimizers. It distinguishes streams +/// whose batch production is driven directly by downstream calls to +/// `Stream::poll_next` from streams that may also drive input or output +/// production independently, such as by spawning tasks or buffering batches +/// ahead of demand. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum EvaluationType { - /// The stream generated by [`execute`](ExecutionPlan::execute) only generates `RecordBatch` - /// instances when it is demanded by invoking `Stream::poll_next`. - /// Filter, projection, and join are examples of such lazy operators. + /// The stream generated by [`execute`](ExecutionPlan::execute) is + /// demand-driven: it produces [`RecordBatch`]es in response to downstream + /// calls to `Stream::poll_next`. + /// + /// Filter, projection, and join operators are examples of lazy operators. /// /// Lazy operators are also known as demand-driven operators. Lazy, - /// The stream generated by [`execute`](ExecutionPlan::execute) eagerly generates `RecordBatch` - /// in one or more spawned Tokio tasks. Eager evaluation is only started the first time - /// `Stream::poll_next` is called. - /// Examples of eager operators are repartition, coalesce partitions, and sort preserving merge. + /// The stream generated by [`execute`](ExecutionPlan::execute) may drive + /// input or output [`RecordBatch`] production ahead of, or independently + /// from, downstream calls to `Stream::poll_next`. + /// + /// Eager operators commonly poll input streams from spawned Tokio tasks, + /// buffer batches ahead of demand, or otherwise create an independent + /// child-polling pipeline. Eager work may start when `execute` creates the + /// stream or when the returned stream is first polled; that timing is an + /// implementation detail. + /// + /// Repartition, coalesce partitions, sort-preserving merge, buffer, and + /// analyze operators are examples of eager operators. /// /// Eager operators are also known as a data-driven operators. Eager, @@ -1098,12 +1478,10 @@ pub(crate) fn emission_type_from_children<'a>( } } -/// Stores certain, often expensive to compute, plan properties used in query -/// optimization. +/// Stores plan properties used in query optimization. /// -/// These properties are stored a single structure to permit this information to -/// be computed once and then those cached results used multiple times without -/// recomputation (aka a cache) +/// Serves as a cache for these properties, which are often +/// expensive to compute. #[derive(Debug, Clone)] pub struct PlanProperties { /// See [ExecutionPlanProperties::equivalence_properties] @@ -1233,36 +1611,93 @@ macro_rules! check_len { }; } +/// All dynamic expressions must have an expression id. +fn check_dynamic_expression_invariants( + plan: &P, +) -> Result<()> { + let mut produced_ids = HashSet::new(); + for expr in plan.dynamic_expressions_produced() { + let Some(expression_id) = expr.expression_id() else { + return internal_err!( + "{}::dynamic_expressions_produced returned an expression without an expression ID", + plan.name() + ); + }; + assert_or_internal_err!( + produced_ids.insert(expression_id), + "{}::dynamic_expressions_produced returned duplicate expression ID {expression_id}", + plan.name() + ); + } + Ok(()) +} + /// Checks a set of invariants that apply to all ExecutionPlan implementations. /// Returns an error if the given node does not conform. pub fn check_default_invariants( plan: &P, - _check: InvariantLevel, + check: InvariantLevel, ) -> Result<(), DataFusionError> { let children_len = plan.children().len(); check_len!(plan, maintains_input_order, children_len); check_len!(plan, required_input_ordering, children_len); - check_len!(plan, required_input_distribution, children_len); check_len!(plan, benefits_from_input_partitioning, children_len); + plan.input_distribution_requirements() + .check_invariants(plan, check)?; + check_dynamic_expression_invariants(plan)?; Ok(()) } -/// Indicate whether a data exchange is needed for the input of `plan`, which will be very helpful -/// especially for the distributed engine to judge whether need to deal with shuffling. -/// Currently, there are 3 kinds of execution plan which needs data exchange -/// 1. RepartitionExec for changing the partition number between two `ExecutionPlan`s -/// 2. CoalescePartitionsExec for collapsing all of the partitions into one without ordering guarantee -/// 3. SortPreservingMergeExec for collapsing all of the sorted partitions into one with ordering guarantee +/// Indicate whether a data exchange is needed for the input of `plan`. +/// +/// This identifies physical operators that redistribute child partitions or +/// gather multiple child partitions into one output partition: +/// +/// 1. RepartitionExec for non-round-robin repartitioning +/// 2. CoalescePartitionsExec for collapsing multiple partitions into one without ordering guarantee +/// 3. SortPreservingMergeExec for collapsing multiple sorted partitions into one with ordering guarantee #[expect(clippy::needless_pass_by_value)] pub fn need_data_exchange(plan: Arc) -> bool { - plan.properties().evaluation_type == EvaluationType::Eager + if let Some(repartition) = plan.downcast_ref::() { + !matches!(repartition.partitioning(), Partitioning::RoundRobinBatch(_)) + } else if let Some(coalesce) = plan.downcast_ref::() { + coalesce.input().output_partitioning().partition_count() > 1 + } else if let Some(sort_preserving_merge) = + plan.downcast_ref::() + { + sort_preserving_merge + .input() + .output_partitioning() + .partition_count() + > 1 + } else { + false + } } -/// Returns a copy of this plan if we change any child according to the pointer comparison. +/// Returns a plan with the given children, skipping as much work as possible. +/// +/// This helper is the single entry point for "rebuild a plan from new +/// children" and applies three layers of short-circuits, from cheapest to +/// most expensive: +/// +/// 1. **Same child pointers** — if every `children[i]` is `Arc::ptr_eq` to the +/// corresponding existing child, the original `plan` is returned +/// unchanged (no allocation, no [`ExecutionPlan::replace_children`] +/// call). +/// 2. **Same child properties** — if the children's `PlanProperties` Arcs +/// match (via [`has_same_children_properties`]), the plan's own +/// `PlanProperties` cache can be reused. This calls +/// [`ExecutionPlan::replace_children`] with [`ChildrenPropertiesMode::Keep`], +/// which swaps the child pointers without recomputing `PlanProperties`. +/// 3. **Full recompute** — otherwise, delegate to +/// [`ExecutionPlan::replace_children`] with [`ChildrenPropertiesMode::Recompute`], +/// which recomputes `PlanProperties` from scratch. +/// /// The size of `children` must be equal to the size of `ExecutionPlan::children()`. -pub fn with_new_children_if_necessary( +pub fn replace_children_if_necessary( plan: Arc, children: Vec>, ) -> Result> { @@ -1272,16 +1707,36 @@ pub fn with_new_children_if_necessary( old_children.len(), "Wrong number of children" ); - if children.is_empty() - || children + if !children.is_empty() { + // Layer 1: same child pointers → return the plan unchanged. + if children .iter() .zip(old_children.iter()) - .any(|(c1, c2)| !Arc::ptr_eq(c1, c2)) - { - plan.with_new_children(children) - } else { - Ok(plan) + .all(|(c1, c2)| Arc::ptr_eq(c1, c2)) + { + return Ok(plan); + } + // Layer 2: same child properties → reuse `PlanProperties` cache. + if has_same_children_properties(plan.as_ref(), &children)? { + return plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ); + } } + // Layer 3: full recompute. + plan.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) +} + +#[deprecated(since = "55.0.0", note = "Use `replace_children_if_necessary`")] +pub fn with_new_children_if_necessary( + plan: Arc, + children: Vec>, +) -> Result> { + replace_children_if_necessary(plan, children) } /// Return a [`DisplayableExecutionPlan`] wrapper around an @@ -1336,6 +1791,13 @@ pub async fn collect_partitioned( plan: Arc, context: Arc, ) -> Result>> { + // Avoid `JoinSet::spawn` for single partition + if plan.output_partitioning().partition_count() == 1 { + let stream = plan.execute(0, context)?; + let batches: Vec = stream.try_collect().await?; + return Ok(vec![batches]); + } + let streams = execute_stream_partitioned(plan, context)?; let mut join_set = JoinSet::new(); @@ -1524,7 +1986,7 @@ pub fn reset_plan_states(plan: Arc) -> Result], ) -> Result { let old_children = plan.children(); @@ -1544,6 +2006,11 @@ pub fn has_same_children_properties( /// Helper macro to avoid properties re-computation if passed children properties /// the same as plan already has. Could be used to implement fast-path for method /// [`ExecutionPlan::with_new_children`]. +/// +/// New call sites should route through [`replace_children_if_necessary`], +/// which applies this check together with the child-pointer short-circuit +/// (see [`replace_children_if_necessary`] for the layered policy). This +/// macro remains for direct-caller sites that have not been migrated yet. #[macro_export] macro_rules! check_if_same_properties { ($plan: expr, $children: expr) => { @@ -1551,12 +2018,28 @@ macro_rules! check_if_same_properties { $plan.as_ref(), &$children, )? { - let plan = $plan.with_new_children_and_same_properties($children); - return Ok(::std::sync::Arc::new(plan)); + return ::std::sync::Arc::clone(&$plan) + .with_new_children_and_same_properties($children); } }; } +/// Helper macro to validate that replacement children match a plan's existing +/// child count. +/// +/// This is useful for [`ExecutionPlan::replace_children`] implementations that +/// need to preserve the same child-count validation behavior. +#[macro_export] +macro_rules! validate_child_count { + ($plan: expr, $children: expr) => { + datafusion_common::assert_eq_or_internal_err!( + $children.len(), + $plan.children().len(), + "Wrong number of children" + ); + }; +} + /// Utility function yielding a string representation of the given [`ExecutionPlan`]. pub fn get_plan_string(plan: &Arc) -> Vec { let formatted = displayable(plan.as_ref()).indent(true).to_string(); @@ -1596,17 +2079,32 @@ pub(crate) fn stub_properties() -> Arc { mod tests { use super::*; + use crate::buffer::BufferExec; + use crate::test::exec::MockExec; use crate::{DisplayAs, DisplayFormatType, ExecutionPlan}; use arrow::array::{DictionaryArray, Int32Array, NullArray, RunArray}; use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; #[derive(Debug)] - pub struct EmptyExec; + pub struct EmptyExec { + dynamic_expressions: Vec>, + } impl EmptyExec { pub fn new(_schema: SchemaRef) -> Self { - Self + Self { + dynamic_expressions: vec![], + } + } + + fn with_dynamic_expressions( + mut self, + dynamic_expressions: Vec>, + ) -> Self { + self.dynamic_expressions = dynamic_expressions; + self } } @@ -1633,20 +2131,35 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn dynamic_expressions_produced(&self) -> Vec> { + self.dynamic_expressions.iter().map(Arc::clone).collect() + } + fn execute( &self, _partition: usize, @@ -1655,14 +2168,41 @@ mod tests { unimplemented!() } - fn partition_statistics( + fn statistics_from_inputs( &self, - _partition: Option, + _input_stats: &[Arc], + _args: &StatisticsArgs, ) -> Result> { unimplemented!() } } + #[test] + fn test_dynamic_expression_invariants() -> Result<()> { + let schema = Arc::new(Schema::empty()); + let dynamic: Arc = + Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let valid = EmptyExec::new(Arc::clone(&schema)) + .with_dynamic_expressions(vec![Arc::clone(&dynamic)]); + check_default_invariants(&valid, InvariantLevel::Always)?; + + let missing_id = + EmptyExec::new(Arc::clone(&schema)).with_dynamic_expressions(vec![lit(true)]); + let error = check_default_invariants(&missing_id, InvariantLevel::Always) + .unwrap_err() + .strip_backtrace(); + assert!(error.contains("without an expression ID"), "{error}"); + + let duplicate = EmptyExec::new(schema) + .with_dynamic_expressions(vec![Arc::clone(&dynamic), dynamic]); + let error = check_default_invariants(&duplicate, InvariantLevel::Always) + .unwrap_err() + .strip_backtrace(); + assert!(error.contains("duplicate expression ID"), "{error}"); + + Ok(()) + } + #[derive(Debug)] pub struct RenamedEmptyExec; @@ -1704,18 +2244,101 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + unimplemented!() + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + unimplemented!() + } + } + + #[derive(Debug)] + struct DowncastDelegatingExec(Arc); + + impl DisplayAs for DowncastDelegatingExec { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for DowncastDelegatingExec { + fn name(&self) -> &'static str { + Self::static_name() + } + + fn properties(&self) -> &Arc { + unimplemented!() + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.0.apply_expressions(f) + } + + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { + Some(self.0.as_ref()) + } + fn execute( &self, _partition: usize, @@ -1731,12 +2354,400 @@ mod tests { unimplemented!() } } + /// Test leaf plan with a real [`PlanProperties`] cache. Different instances + /// can share the same cache Arc by cloning `cache`. + #[derive(Debug, Clone)] + struct WithChildrenTestLeaf { + cache: Arc, + } + + impl WithChildrenTestLeaf { + fn new(cache: Arc) -> Self { + Self { cache } + } + } + + impl DisplayAs for WithChildrenTestLeaf { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestLeaf { + fn name(&self) -> &'static str { + "WithChildrenTestLeaf" + } + fn properties(&self) -> &Arc { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + } + + /// Test unary plan that counts which of `with_new_children` (full + /// recompute) vs `with_new_children_and_same_properties` (fast path) is + /// taken. + #[derive(Debug, Clone)] + struct WithChildrenTestParent { + input: Arc, + cache: Arc, + recompute_calls: Arc, + fast_path_calls: Arc, + } + + impl WithChildrenTestParent { + fn new(input: Arc) -> Self { + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Self { + input, + cache, + recompute_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + fast_path_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + } + + impl DisplayAs for WithChildrenTestParent { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestParent { + fn name(&self) -> &'static str { + "WithChildrenTestParent" + } + fn properties(&self) -> &Arc { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.fast_path_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + Ok(Arc::new(Self { + input: children.swap_remove(0), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => { + self.recompute_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + // Full recompute: allocate a fresh `PlanProperties` Arc so this + // path is observable via `Arc::ptr_eq` on properties. + let new_input = children.swap_remove(0); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Ok(Arc::new(Self { + input: new_input, + cache, + recompute_calls: Arc::clone(&self.recompute_calls), + fast_path_calls: Arc::clone(&self.fast_path_calls), + })) + } + } + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) + } + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + } + + /// Test unary plan that does **not** override + /// `with_new_children_and_same_properties`. Used to verify the default + /// trait fallback still routes through `with_new_children` (which is + /// the semantics-preserving path for downstream / external + /// `ExecutionPlan` implementations that haven't opted into the + /// fast path yet). + #[derive(Debug, Clone)] + struct WithChildrenTestParentDefault { + input: Arc, + cache: Arc, + recompute_calls: Arc, + } + + impl WithChildrenTestParentDefault { + fn new(input: Arc) -> Self { + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Self { + input, + cache, + recompute_calls: Arc::new(std::sync::atomic::AtomicUsize::new(0)), + } + } + } + + impl DisplayAs for WithChildrenTestParentDefault { + fn fmt_as( + &self, + _t: DisplayFormatType, + _f: &mut std::fmt::Formatter, + ) -> std::fmt::Result { + unimplemented!() + } + } + + impl ExecutionPlan for WithChildrenTestParentDefault { + fn name(&self) -> &'static str { + "WithChildrenTestParentDefault" + } + fn properties(&self) -> &Arc { + &self.cache + } + fn children(&self) -> Vec<&Arc> { + vec![&self.input] + } + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + fn with_new_children( + self: Arc, + mut children: Vec>, + ) -> Result> { + self.recompute_calls + .fetch_add(1, std::sync::atomic::Ordering::SeqCst); + let new_input = children.swap_remove(0); + let cache = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + Ok(Arc::new(Self { + input: new_input, + cache, + recompute_calls: Arc::clone(&self.recompute_calls), + })) + } + // Intentionally does **not** override + // `with_new_children_and_same_properties` — relies on the trait + // default that falls back to `with_new_children`. + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unimplemented!() + } + } + + /// Cover the three short-circuit layers of + /// [`replace_children_if_necessary`]. + #[test] + fn test_replace_children_if_necessary_layers() -> Result<()> { + use std::sync::atomic::Ordering; + + // Two leaves that share the same `PlanProperties` Arc but sit behind + // distinct `Arc` pointers. + let leaf_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_a: Arc = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + let leaf_b: Arc = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + // A third leaf with a *different* `PlanProperties` Arc — for layer 3. + let leaf_c_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_c: Arc = + Arc::new(WithChildrenTestLeaf::new(leaf_c_props)); + + let parent = Arc::new(WithChildrenTestParent::new(Arc::clone(&leaf_a))); + let parent_dyn: Arc = Arc::clone(&parent) as _; + let orig_props = Arc::clone(parent.properties()); + + // Layer 1: same child pointer → returns the original plan Arc verbatim. + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_a)], + )?; + assert!(Arc::ptr_eq(&out, &parent_dyn)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 0); + + // Layer 2: distinct child Arc, but children share the same + // `PlanProperties` Arc → fast path, parent's `PlanProperties` cache + // Arc is reused (not reallocated). + assert!(!Arc::ptr_eq(&leaf_a, &leaf_b)); + assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties())); + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_b)], + )?; + assert!(Arc::ptr_eq(out.properties(), &orig_props)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 1); + + // Layer 3: child's `PlanProperties` Arc differs → full recompute. + assert!(!Arc::ptr_eq(leaf_a.properties(), leaf_c.properties())); + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_c)], + )?; + assert!(!Arc::ptr_eq(out.properties(), &orig_props)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 1); + assert_eq!(parent.fast_path_calls.load(Ordering::SeqCst), 1); + + Ok(()) + } + + /// A plan that does not override `with_new_children_and_same_properties` + /// (per @kosiew's review on #23332) must still be routed through + /// `with_new_children` when the helper hits the "same properties" + /// branch. The default trait implementation forwards to + /// `with_new_children`, so downstream / external `ExecutionPlan` + /// implementations keep the semantics-preserving path. + #[test] + fn test_replace_children_if_necessary_default_fallback() -> Result<()> { + use std::sync::atomic::Ordering; + + let leaf_props = Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::new(Schema::empty())), + Partitioning::UnknownPartitioning(1), + EmissionType::Final, + Boundedness::Bounded, + )); + let leaf_a: Arc = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + let leaf_b: Arc = + Arc::new(WithChildrenTestLeaf::new(Arc::clone(&leaf_props))); + assert!(!Arc::ptr_eq(&leaf_a, &leaf_b)); + assert!(Arc::ptr_eq(leaf_a.properties(), leaf_b.properties())); + + let parent = Arc::new(WithChildrenTestParentDefault::new(Arc::clone(&leaf_a))); + let parent_dyn: Arc = Arc::clone(&parent) as _; + + // Using the same child means we return the original plan Arc verbatim, so even when + // the `replace_children` `ChildrenPropertiesMode::Keep` path is not defined, + // we do not recompute. + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_a)], + )?; + assert!(Arc::ptr_eq(&out, &parent_dyn)); + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 0); + + // Using a distinct child but the same `PlanProperties` Arc means the helper + // attempts to enter the Keep branch. If it does not exist, we fall back + // to recomputation. + let out = replace_children_if_necessary( + Arc::clone(&parent_dyn), + vec![Arc::clone(&leaf_b)], + )?; + // `with_new_children` was invoked exactly once via the default. + assert_eq!(parent.recompute_calls.load(Ordering::SeqCst), 1); + // The returned plan has a freshly-recomputed `PlanProperties` Arc, + // so it differs from the parent's original cache. This confirms + // the fallback ran and did not short-circuit. + assert!(!Arc::ptr_eq(out.properties(), parent.properties())); + + Ok(()) + } /// A test node that holds a fixed list of expressions, used to test /// `apply_expressions` behavior. #[derive(Debug)] struct MultiExprExec { exprs: Vec>, + children: Vec>, } impl DisplayAs for MultiExprExec { @@ -1759,7 +2770,7 @@ mod tests { } fn children(&self) -> Vec<&Arc> { - vec![] + self.children.iter().collect() } fn with_new_children( @@ -1771,13 +2782,9 @@ mod tests { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for expr in &self.exprs { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - Ok(tnr) + apply_expression_roots(&self.exprs, f) } fn execute( @@ -1809,6 +2816,7 @@ mod tests { fn test_apply_expressions_continue_visits_all() -> Result<()> { let plan = MultiExprExec { exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], }; let mut visited = 0usize; plan.apply_expressions(&mut |_expr| { @@ -1823,6 +2831,7 @@ mod tests { fn test_apply_expressions_stop_halts_early() -> Result<()> { let plan = MultiExprExec { exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], }; let mut visited = 0usize; let tnr = plan.apply_expressions(&mut |_expr| { @@ -1835,6 +2844,69 @@ mod tests { Ok(()) } + #[test] + fn test_apply_expressions_jump_visits_next_root() -> Result<()> { + let plan = MultiExprExec { + exprs: vec![lit_expr(1), lit_expr(2), lit_expr(3)], + children: vec![], + }; + let mut visited = 0usize; + let tnr = plan.apply_expressions(&mut |_expr| { + visited += 1; + Ok(TreeNodeRecursion::Jump) + })?; + assert_eq!(visited, 3); + assert_eq!(tnr, TreeNodeRecursion::Continue); + Ok(()) + } + + #[test] + fn test_apply_expressions_does_not_recurse() -> Result<()> { + use datafusion_physical_expr::expressions::NegativeExpr; + + let child: Arc = Arc::new(MultiExprExec { + exprs: vec![lit_expr(2)], + children: vec![], + }); + let nested: Arc = Arc::new(NegativeExpr::new(lit_expr(1))); + let plan = MultiExprExec { + exprs: vec![nested], + children: vec![child], + }; + + let mut visited = 0; + plan.apply_expressions(&mut |expr| { + visited += 1; + assert!(expr.is::()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(visited, 1); + Ok(()) + } + + #[test] + fn test_apply_expressions_callback_can_retain_arc() -> Result<()> { + let expected = lit_expr(1); + let plan = MultiExprExec { + exprs: vec![Arc::clone(&expected)], + children: vec![], + }; + let mut retained = None; + plan.apply_expressions(&mut |expr| { + retained = Some(Arc::clone(expr)); + Ok(TreeNodeRecursion::Continue) + })?; + drop(plan); + + assert!(Arc::ptr_eq( + &expected, + retained + .as_ref() + .expect("callback should retain expression") + )); + Ok(()) + } + #[test] fn test_execution_plan_name() { let schema1 = Arc::new(Schema::empty()); @@ -1847,6 +2919,24 @@ mod tests { assert_eq!(RenamedEmptyExec::static_name(), "MyRenamedEmptyExec"); } + #[test] + fn test_execution_plan_downcast_delegates_to_downcast_delegate() { + let schema = Arc::new(Schema::empty()); + let inner: Arc = Arc::new(EmptyExec::new(schema)); + let wrapped: Arc = Arc::new(DowncastDelegatingExec(inner)); + let nested: Arc = + Arc::new(DowncastDelegatingExec(Arc::clone(&wrapped))); + + for plan in [wrapped.as_ref(), nested.as_ref()] { + assert!(!plan.is::()); + assert!(plan.downcast_ref::().is_none()); + assert!(plan.is::()); + assert!(plan.downcast_ref::().is_some()); + assert!(!plan.is::()); + assert!(plan.downcast_ref::().is_none()); + } + } + /// A compilation test to ensure that the `ExecutionPlan::name()` method can /// be called from a trait object. /// Related ticket: https://github.com/apache/datafusion/pull/11047 @@ -1855,6 +2945,15 @@ mod tests { let _ = plan.name(); } + #[test] + fn buffer_exec_does_not_need_data_exchange() { + let schema = Arc::new(Schema::empty()); + let input: Arc = Arc::new(MockExec::new(vec![], schema)); + let buffer: Arc = Arc::new(BufferExec::new(input, 1024)); + + assert!(!need_data_exchange(buffer)); + } + #[test] fn test_check_not_null_constraints_accept_non_null() -> Result<()> { check_not_null_constraints( diff --git a/datafusion/physical-plan/src/explain.rs b/datafusion/physical-plan/src/explain.rs index 617a1a6cdaf53..3b31ee748b736 100644 --- a/datafusion/physical-plan/src/explain.rs +++ b/datafusion/physical-plan/src/explain.rs @@ -22,7 +22,10 @@ use std::sync::Arc; use super::{DisplayAs, PlanProperties, SendableRecordBatchStream}; use crate::execution_plan::{Boundedness, EmissionType}; use crate::stream::RecordBatchStreamAdapter; -use crate::{DisplayFormatType, ExecutionPlan, Partitioning}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, +}; use arrow::{array::StringBuilder, datatypes::SchemaRef, record_batch::RecordBatch}; use datafusion_common::display::StringifiedPlan; @@ -119,18 +122,29 @@ impl ExecutionPlan for ExplainExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -193,6 +207,188 @@ impl ExecutionPlan for ExplainExec { futures::stream::iter(vec![Ok(record_batch)]), ))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Explain( + protobuf::ExplainExecNode { + schema: Some(self.schema().as_ref().try_into()?), + stringified_plans: self + .stringified_plans() + .iter() + .map(stringified_plan_to_proto) + .collect(), + verbose: self.verbose(), + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ExplainExec { + /// Reconstruct an [`ExplainExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let explain = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Explain, + "ExplainExec", + ); + let schema = explain.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ExplainExec is missing required field 'schema'" + ) + })?; + Ok(Arc::new(ExplainExec::new( + Arc::new(arrow::datatypes::Schema::try_from(schema)?), + explain + .stringified_plans + .iter() + .map(stringified_plan_from_proto) + .collect(), + explain.verbose, + ))) + } +} + +#[cfg(feature = "proto")] +fn stringified_plan_to_proto( + stringified_plan: &StringifiedPlan, +) -> datafusion_proto_models::protobuf::StringifiedPlan { + use datafusion_common::display::PlanType; + use datafusion_proto_models::datafusion_common::EmptyMessage; + use datafusion_proto_models::protobuf; + use protobuf::plan_type::PlanTypeEnum::{ + AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, + FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, + InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, + InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, + PhysicalPlanError, + }; + + protobuf::StringifiedPlan { + plan_type: match stringified_plan.clone().plan_type { + PlanType::InitialLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(InitialLogicalPlan(EmptyMessage {})), + }), + PlanType::AnalyzedLogicalPlan { analyzer_name } => Some(protobuf::PlanType { + plan_type_enum: Some(AnalyzedLogicalPlan( + protobuf::AnalyzedLogicalPlanType { analyzer_name }, + )), + }), + PlanType::FinalAnalyzedLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalAnalyzedLogicalPlan(EmptyMessage {})), + }), + PlanType::OptimizedLogicalPlan { optimizer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(OptimizedLogicalPlan( + protobuf::OptimizedLogicalPlanType { optimizer_name }, + )), + }) + } + PlanType::FinalLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalLogicalPlan(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlan(EmptyMessage {})), + }), + PlanType::OptimizedPhysicalPlan { optimizer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(OptimizedPhysicalPlan( + protobuf::OptimizedPhysicalPlanType { optimizer_name }, + )), + }) + } + PlanType::FinalPhysicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlan(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlanWithStats => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlanWithStats(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlanWithSchema => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlanWithSchema(EmptyMessage {})), + }), + PlanType::FinalPhysicalPlanWithStats => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlanWithStats(EmptyMessage {})), + }), + PlanType::FinalPhysicalPlanWithSchema => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlanWithSchema(EmptyMessage {})), + }), + PlanType::PhysicalPlanError => Some(protobuf::PlanType { + plan_type_enum: Some(PhysicalPlanError(EmptyMessage {})), + }), + }, + plan: stringified_plan.plan.to_string(), + } +} + +#[cfg(feature = "proto")] +fn stringified_plan_from_proto( + stringified_plan: &datafusion_proto_models::protobuf::StringifiedPlan, +) -> StringifiedPlan { + use datafusion_common::display::PlanType; + use datafusion_proto_models::protobuf::plan_type::PlanTypeEnum::{ + AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, + FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, + InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, + InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, + PhysicalPlanError, + }; + use datafusion_proto_models::protobuf::{ + AnalyzedLogicalPlanType, OptimizedLogicalPlanType, OptimizedPhysicalPlanType, + }; + + StringifiedPlan { + plan_type: match stringified_plan + .plan_type + .as_ref() + .and_then(|plan_type| plan_type.plan_type_enum.as_ref()) + .unwrap_or_else(|| { + panic!( + "Cannot create protobuf::StringifiedPlan from {stringified_plan:?}" + ) + }) { + InitialLogicalPlan(_) => PlanType::InitialLogicalPlan, + AnalyzedLogicalPlan(AnalyzedLogicalPlanType { analyzer_name }) => { + PlanType::AnalyzedLogicalPlan { + analyzer_name: analyzer_name.clone(), + } + } + FinalAnalyzedLogicalPlan(_) => PlanType::FinalAnalyzedLogicalPlan, + OptimizedLogicalPlan(OptimizedLogicalPlanType { optimizer_name }) => { + PlanType::OptimizedLogicalPlan { + optimizer_name: optimizer_name.clone(), + } + } + FinalLogicalPlan(_) => PlanType::FinalLogicalPlan, + InitialPhysicalPlan(_) => PlanType::InitialPhysicalPlan, + InitialPhysicalPlanWithStats(_) => PlanType::InitialPhysicalPlanWithStats, + InitialPhysicalPlanWithSchema(_) => PlanType::InitialPhysicalPlanWithSchema, + OptimizedPhysicalPlan(OptimizedPhysicalPlanType { optimizer_name }) => { + PlanType::OptimizedPhysicalPlan { + optimizer_name: optimizer_name.clone(), + } + } + FinalPhysicalPlan(_) => PlanType::FinalPhysicalPlan, + FinalPhysicalPlanWithStats(_) => PlanType::FinalPhysicalPlanWithStats, + FinalPhysicalPlanWithSchema(_) => PlanType::FinalPhysicalPlanWithSchema, + PhysicalPlanError(_) => PlanType::PhysicalPlanError, + }, + plan: Arc::new(stringified_plan.plan.clone()), + } } /// If this plan should be shown, given the previous plan that was diff --git a/datafusion/physical-plan/src/filter.rs b/datafusion/physical-plan/src/filter.rs index c485e181f3826..5df5482fb75de 100644 --- a/datafusion/physical-plan/src/filter.rs +++ b/datafusion/physical-plan/src/filter.rs @@ -28,10 +28,9 @@ use super::{ ColumnStatistics, DisplayAs, ExecutionPlanProperties, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, }; -use crate::check_if_same_properties; use crate::coalesce::{LimitedBatchCoalescer, PushBatchStatus}; use crate::common::can_project; -use crate::execution_plan::CardinalityEffect; +use crate::execution_plan::{CardinalityEffect, replace_children_if_necessary}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, PushedDown, @@ -42,7 +41,9 @@ use crate::projection::{ EmbeddedProjection, ProjectionExec, ProjectionExpr, make_with_child, try_embed_projection, update_expr, }; +use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; use crate::stream::EmptyRecordBatchStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use crate::{ DisplayFormatType, ExecutionPlan, metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RatioMetrics}, @@ -61,7 +62,9 @@ use datafusion_common::{ use datafusion_execution::TaskContext; use datafusion_expr::Operator; use datafusion_physical_expr::equivalence::ProjectionMapping; -use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, lit}; +use datafusion_physical_expr::expressions::{ + BinaryExpr, Column, IsNotNullExpr, Literal, lit, +}; use datafusion_physical_expr::intervals::utils::check_support; use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns}; use datafusion_physical_expr::{ @@ -311,11 +314,24 @@ impl FilterExec { &self.projection } - /// Calculates `Statistics` for `FilterExec`, by applying selectivity - /// (either default, or estimated) to input statistics. + /// Calculates `Statistics` for `FilterExec` by applying the filter's + /// selectivity (default, or estimated from interval analysis) to the input + /// statistics. + /// + /// The estimated output row count is used to keep the per-column statistics + /// consistent with it: + /// - null and distinct counts are capped at the estimated row count; + /// - byte sizes (per column and total) are scaled by the selectivity, and + /// are an exact zero when the row count is an exact zero; + /// - a column constrained to a single value (`col = literal`, or an + /// interval that collapses to one point) gets a distinct count of 1; + /// - a column in a null-rejecting conjunct gets a null count of 0. + /// + /// When interval analysis applies, min/max are also tightened to the + /// surviving value range. /// - /// Equality predicates (`col = literal`) set NDV to `Exact(1)`, or - /// `Exact(0)` when the predicate is contradictory (e.g. `a = 1 AND a = 2`). + /// A contradictory predicate (e.g. `a = 1 AND a = 2`) yields zero rows and + /// empty-column statistics. pub(crate) fn statistics_helper( schema: &SchemaRef, input_stats: Statistics, @@ -328,8 +344,8 @@ impl FilterExec { let input_total_byte_size = input_stats.total_byte_size; let (selectivity, num_rows, column_statistics) = if is_infeasible { - // Contradictory predicate: zero rows, and null/min/max are - // undefined on an empty column. + // Contradictory predicate: no rows survive. Row-bounded counts are + // zero; value statistics are undefined on an empty column. let mut cs = input_stats.to_inexact().column_statistics; for col_stat in &mut cs { col_stat.distinct_count = Precision::Exact(0); @@ -340,47 +356,58 @@ impl FilterExec { col_stat.byte_size = Precision::Exact(0); } (0.0, Precision::Exact(0), cs) - } else if !check_support(predicate, schema) { - // Interval analysis is not applicable; fall back to the default - // selectivity but still pin NDV=1 for every `col = literal` column. - let selectivity = default_selectivity as f64 / 100.0; - let mut cs = input_stats.to_inexact().column_statistics; - for &idx in &eq_columns { - if idx < cs.len() && cs[idx].distinct_count != Precision::Exact(0) { - cs[idx].distinct_count = Precision::Exact(1); + } else { + let null_rejecting_columns = collect_null_rejecting_columns(predicate); + + if check_support(predicate, schema) { + let input_analysis_ctx = AnalysisContext::try_from_statistics( + schema, + &input_stats.column_statistics, + )?; + let analysis_ctx = analyze(predicate, input_analysis_ctx, schema)?; + let selectivity = analysis_ctx.selectivity.unwrap_or(1.0); + let filtered_num_rows = + input_num_rows.with_estimated_selectivity(selectivity); + let cs = collect_new_statistics( + schema, + &input_stats.column_statistics, + analysis_ctx.boundaries, + selectivity, + &null_rejecting_columns, + filtered_num_rows, + ); + (selectivity, filtered_num_rows, cs) + } else { + // Without interval boundaries, use the default selectivity and + // apply the row-count constraints that still follow from the + // filter predicate. + let selectivity = default_selectivity as f64 / 100.0; + let filtered_num_rows = + input_num_rows.with_estimated_selectivity(selectivity); + let mut cs = input_stats.to_inexact().column_statistics; + for (idx, col_stat) in cs.iter_mut().enumerate() { + col_stat.byte_size = scale_byte_size_at_rows( + col_stat.byte_size, + selectivity, + filtered_num_rows, + ); + col_stat.null_count = if null_rejecting_columns.contains(&idx) { + Precision::Exact(0) + } else { + cap_at_rows(col_stat.null_count, filtered_num_rows) + }; + col_stat.distinct_count = if eq_columns.contains(&idx) { + distinct_count_for_singleton_domain(filtered_num_rows) + } else { + cap_at_rows(col_stat.distinct_count, filtered_num_rows) + }; } + (selectivity, filtered_num_rows, cs) } - ( - selectivity, - input_num_rows.with_estimated_selectivity(selectivity), - cs, - ) - } else { - // Interval-analysis path. `collect_new_statistics` already sets - // distinct_count = Exact(1) when an interval collapses to a single - // value, so no post-fix is needed here. - let input_analysis_ctx = AnalysisContext::try_from_statistics( - schema, - &input_stats.column_statistics, - )?; - let analysis_ctx = analyze(predicate, input_analysis_ctx, schema)?; - let selectivity = analysis_ctx.selectivity.unwrap_or(1.0); - let filtered_num_rows = - input_num_rows.with_estimated_selectivity(selectivity); - let cs = collect_new_statistics( - schema, - &input_stats.column_statistics, - analysis_ctx.boundaries, - match &filtered_num_rows { - Precision::Absent => None, - p => Some(*p), - }, - ); - (selectivity, filtered_num_rows, cs) }; let total_byte_size = - input_total_byte_size.with_estimated_selectivity(selectivity); + scale_byte_size_at_rows(input_total_byte_size, selectivity, num_rows); Ok(Statistics { num_rows, @@ -401,7 +428,10 @@ impl FilterExec { let schema = input.schema(); let stats = Self::statistics_helper( &schema, - Arc::unwrap_or_clone(input.partition_statistics(None)?), + Arc::unwrap_or_clone( + StatisticsContext::new() + .compute(input.as_ref(), &StatisticsArgs::new())?, + ), predicate, default_selectivity, )?; @@ -449,17 +479,6 @@ impl FilterExec { input.boundedness(), )) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for FilterExec { @@ -523,9 +542,9 @@ impl ExecutionPlan for FilterExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - f(self.predicate.as_ref()) + crate::apply_expression_roots([&self.predicate], f) } fn maintains_input_order(&self) -> Vec { @@ -533,16 +552,46 @@ impl ExecutionPlan for FilterExec { vec![true] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - let new_input = children.swap_remove(0); - FilterExecBuilder::from(&*self) - .with_input(new_input) - .build() - .map(|e| Arc::new(e) as _) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new_input = children.swap_remove(0); + FilterExecBuilder::from(&*self) + .with_input(new_input) + .build() + .map(|e| Arc::new(e) as _) + } + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -575,11 +624,18 @@ impl ExecutionPlan for FilterExec { Some(self.metrics.clone_inner()) } + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + /// The output statistics of a filtering operation can be estimated if the /// predicate's selectivity value can be determined for the incoming data. - fn partition_statistics(&self, partition: Option) -> Result> { - let input_stats = - Arc::unwrap_or_clone(self.input.partition_statistics(partition)?); + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let input_stats = input_stats[0].as_ref().clone(); let stats = Self::statistics_helper( &self.input.schema(), input_stats, @@ -783,11 +839,105 @@ impl ExecutionPlan for FilterExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = ctx.encode_expr(self.predicate())?; + // Preserve the exact wire format: `None` (full projection) is serialized + // as the identity projection `[0, 1, ..., num_fields - 1]` so that it is + // distinguishable from an explicit projection on decode. + let projection = if let Some(v) = self.projection() { + v.iter().map(|x| *x as u32).collect() + } else { + (0..self.input().schema().fields().len()) + .map(|i| i as u32) + .collect() + }; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Filter(Box::new( + protobuf::FilterExecNode { + input: Some(Box::new(input)), + expr: Some(expr), + default_filter_selectivity: self.default_selectivity() as u32, + projection, + batch_size: self.batch_size() as u32, + fetch: self.fetch().map(|f| f as u32), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl FilterExec { + /// Reconstruct a [`FilterExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one signature. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let filter = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Filter, + "FilterExec", + ); + let input = + ctx.decode_required_child(filter.input.as_deref(), "FilterExec", "input")?; + let predicate = ctx.decode_required_expr( + filter.expr.as_ref(), + input.schema().as_ref(), + "FilterExec", + "expr", + )?; + let filter_selectivity = filter.default_filter_selectivity.try_into(); + + // `None` is encoded as the full identity projection. Reconstruct it only + // when all input columns are present in order, leaving an empty list as + // `Some(vec![])`. + let num_fields = input.schema().fields().len(); + let mut is_full_projection = filter.projection.len() == num_fields; + let mut projection_vec: Vec = Vec::with_capacity(filter.projection.len()); + for (i, idx) in filter.projection.iter().enumerate() { + let idx = *idx as usize; + is_full_projection &= idx == i; + projection_vec.push(idx); + } + let projection = if is_full_projection { + None + } else { + Some(projection_vec) + }; + let filter = FilterExecBuilder::new(predicate, input) + .apply_projection(projection)? + .with_batch_size(filter.batch_size as usize) + .with_fetch(filter.fetch.map(|f| f as usize)) + .build()?; + match filter_selectivity { + Ok(filter_selectivity) => Ok(Arc::new( + filter.with_default_selectivity(filter_selectivity)?, + )), + Err(_) => Err(datafusion_common::internal_datafusion_err!( + "filter_selectivity in PhysicalPlanNode is invalid" + )), + } + } } impl EmbeddedProjection for FilterExec { @@ -854,6 +1004,48 @@ fn collect_equality_columns(predicate: &Arc) -> (HashSet) -> HashSet { + let mut columns = HashSet::new(); + + for expr in split_conjunction(predicate) { + // `col IS NOT NULL` keeps only rows where `col` is non-null. + if let Some(is_not_null) = expr.downcast_ref::() { + if let Some(col) = is_not_null.arg().downcast_ref::() { + columns.insert(col.index()); + } + continue; + } + + // A binary operator that returns NULL on NULL input rejects rows where + // a direct column operand is NULL. + if let Some(binary) = expr.downcast_ref::() { + if !binary.op().returns_null_on_null() { + continue; + } + if let Some(col) = binary.left().downcast_ref::() { + columns.insert(col.index()); + } + if let Some(col) = binary.right().downcast_ref::() { + columns.insert(col.index()); + } + } + } + + columns +} + /// Converts an interval bound to a [`Precision`] value. NULL bounds (which /// represent "unbounded" in the interval type) map to [`Precision::Absent`]. fn interval_bound_to_precision( @@ -869,15 +1061,67 @@ fn interval_bound_to_precision( } } -/// This function ensures that all bounds in the `ExprBoundaries` vector are -/// converted to closed bounds. If a lower/upper bound is initially open, it -/// is adjusted by using the next/previous value for its data type to convert -/// it into a closed bound. +/// Caps a row-bounded column statistic (a null count or distinct count) at the +/// filtered row count, since a column cannot have more nulls or distinct values +/// than it has rows. Known counts are demoted to inexact because a +/// filter-derived row bound is normally an estimate, the exception being an +/// exact zero, which proves the column is empty. +fn cap_at_rows( + value: Precision, + filtered_num_rows: Precision, +) -> Precision { + match filtered_num_rows { + Precision::Absent => value.to_inexact(), + Precision::Exact(0) => Precision::Exact(0), + rows => value.to_inexact().min(&rows), + } +} + +/// Scales a byte size by the filter selectivity. An exact zero row count means +/// the output is exactly empty, so the byte size is an exact zero too. +fn scale_byte_size_at_rows( + byte_size: Precision, + selectivity: f64, + filtered_num_rows: Precision, +) -> Precision { + if filtered_num_rows == Precision::Exact(0) { + Precision::Exact(0) + } else { + byte_size.with_estimated_selectivity(selectivity) + } +} + +/// Returns the NDV for a column constrained to one non-null value (e.g. +/// `column = literal` or a singleton interval), derived from the filtered row +/// estimate: zero rows means zero distinct values, a known positive row count +/// means exactly one, and an unknown row count means an inexact one (the column +/// could still be empty). +/// +/// The caller is responsible for proving the singleton domain. +fn distinct_count_for_singleton_domain( + filtered_num_rows: Precision, +) -> Precision { + match filtered_num_rows { + Precision::Exact(0) | Precision::Inexact(0) => filtered_num_rows, + // The row count is unknown, so the column could still be empty (zero + // distinct values); report an inexact one rather than overstating it. + Precision::Absent => Precision::Inexact(1), + _ => Precision::Exact(1), + } +} + +/// Builds output column statistics from interval-analysis boundaries. +/// +/// The interval bounds become min/max values, singleton intervals become +/// singleton NDV, and row-bounded counts are kept consistent with the filtered +/// row estimate. fn collect_new_statistics( schema: &SchemaRef, input_column_stats: &[ColumnStatistics], analysis_boundaries: Vec, - filtered_num_rows: Option>, + selectivity: f64, + null_rejecting_columns: &HashSet, + filtered_num_rows: Precision, ) -> Vec { analysis_boundaries .into_iter() @@ -912,24 +1156,32 @@ fn collect_new_statistics( !lower.is_null() && !upper.is_null() && lower == upper; let min_value = interval_bound_to_precision(lower, is_single_value); let max_value = interval_bound_to_precision(upper, is_single_value); - // When the interval collapses to a single value (equality - // predicate), the column has exactly 1 distinct value. - // Otherwise, cap NDV at the filtered row count. + + // Distinct and null counts cannot exceed the number of rows + // that survive the filter. Singleton intervals and + // null-rejecting predicates provide tighter bounds. let capped_distinct_count = if is_single_value { - Precision::Exact(1) + distinct_count_for_singleton_domain(filtered_num_rows) } else { - match filtered_num_rows { - Some(rows) => distinct_count.to_inexact().min(&rows), - None => distinct_count.to_inexact(), - } + cap_at_rows(distinct_count, filtered_num_rows) }; + let capped_null_count = if null_rejecting_columns.contains(&idx) { + Precision::Exact(0) + } else { + cap_at_rows(input_column_stats[idx].null_count, filtered_num_rows) + }; + let byte_size = scale_byte_size_at_rows( + input_column_stats[idx].byte_size, + selectivity, + filtered_num_rows, + ); ColumnStatistics { - null_count: input_column_stats[idx].null_count.to_inexact(), + null_count: capped_null_count, max_value, min_value, sum_value: Precision::Absent, distinct_count: capped_distinct_count, - byte_size: input_column_stats[idx].byte_size, + byte_size, } }, ) @@ -1160,6 +1412,7 @@ mod tests { use super::*; use crate::empty::EmptyExec; use crate::expressions::*; + use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; use crate::test::exec::StatisticsExec; use arrow::datatypes::{Field, Schema, UnionFields, UnionMode}; @@ -1236,7 +1489,8 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Inexact(25)); assert_eq!( statistics.total_byte_size, @@ -1245,6 +1499,8 @@ mod tests { assert_eq!( statistics.column_statistics, vec![ColumnStatistics { + // `a <= 25` rejects nulls, so the column has no surviving nulls. + null_count: Precision::Exact(0), min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), max_value: Precision::Inexact(ScalarValue::Int32(Some(25))), ..Default::default() @@ -1286,11 +1542,14 @@ mod tests { sub_filter, )?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Inexact(16)); assert_eq!( statistics.column_statistics, vec![ColumnStatistics { + // `a <= 25 AND a >= 10` rejects nulls in `a`. + null_count: Precision::Exact(0), min_value: Precision::Inexact(ScalarValue::Int32(Some(10))), max_value: Precision::Inexact(ScalarValue::Int32(Some(25))), ..Default::default() @@ -1346,7 +1605,8 @@ mod tests { binary(col("a", &schema)?, Operator::GtEq, lit(10i32), &schema)?, b_gt_5, )?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; // On a uniform distribution, only fifteen rows will satisfy the // filter that 'a' proposed (a >= 10 AND a <= 25) (15/100) and only // 5 rows will satisfy the filter that 'b' proposed (b > 45) (5/50). @@ -1358,11 +1618,16 @@ mod tests { statistics.column_statistics, vec![ ColumnStatistics { + // `a <= 25 AND a >= 10` rejects nulls in `a`. + null_count: Precision::Exact(0), min_value: Precision::Inexact(ScalarValue::Int32(Some(10))), max_value: Precision::Inexact(ScalarValue::Int32(Some(25))), ..Default::default() }, ColumnStatistics { + // `b > 45` in the upstream filter zeroes b's nulls; the outer + // filter then caps the (already zero) count, demoting to inexact. + null_count: Precision::Inexact(0), min_value: Precision::Inexact(ScalarValue::Int32(Some(46))), max_value: Precision::Inexact(ScalarValue::Int32(Some(50))), ..Default::default() @@ -1391,7 +1656,8 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Absent); Ok(()) @@ -1464,7 +1730,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; // 0.5 (from a) * 0.333333... (from b) * 0.798387... (from c) ≈ 0.1330... // num_rows after ceil => 133.0... => 134 // total_byte_size after ceil => 532.0... => 533 @@ -1559,11 +1826,20 @@ mod tests { Arc::new(Column::new("b", 1)), )), )); - // Since filter predicate passes all entries, statistics after filter shouldn't change. - let expected = input.partition_statistics(None)?.column_statistics.clone(); + // The filter predicate passes all (non-null) entries, so min/max/NDV + // are unchanged. `a < 200` and `1 <= b` are null-rejecting, though, so + // both columns lose any nulls regardless of selectivity. + let mut expected = StatisticsContext::new() + .compute(input.as_ref(), &StatisticsArgs::new())? + .column_statistics + .clone(); + for col in &mut expected { + col.null_count = Precision::Exact(0); + } let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Inexact(1000)); assert_eq!(statistics.total_byte_size, Precision::Inexact(4000)); @@ -1616,7 +1892,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Inexact(0)); assert_eq!(statistics.total_byte_size, Precision::Inexact(0)); @@ -1703,7 +1980,8 @@ mod tests { Arc::new(FilterExec::try_new(outer_predicate, inner_filter)?); // Should succeed without error - let statistics = outer_filter.partition_statistics(None)?; + let statistics = StatisticsContext::new() + .compute(outer_filter.as_ref(), &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Inexact(0)); Ok(()) @@ -1742,7 +2020,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Inexact(490)); assert_eq!(statistics.total_byte_size, Precision::Inexact(1960)); @@ -1750,10 +2029,14 @@ mod tests { statistics.column_statistics, vec![ ColumnStatistics { + // `a < 50` rejects nulls in `a`. + null_count: Precision::Exact(0), min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), max_value: Precision::Inexact(ScalarValue::Int32(Some(49))), ..Default::default() }, + // `b` is not referenced by the predicate, so its stats are + // unchanged (null count stays unknown). ColumnStatistics { min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), @@ -1792,13 +2075,16 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let filter_statistics = filter.partition_statistics(None)?; + let filter_statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; let expected_filter_statistics = Statistics { num_rows: Precision::Absent, total_byte_size: Precision::Absent, column_statistics: vec![ColumnStatistics { - null_count: Precision::Absent, + // `a <= 10` rejects nulls, so `a` has no surviving nulls even + // though the input statistics are entirely unknown. + null_count: Precision::Exact(0), min_value: Precision::Inexact(ScalarValue::Int32(Some(5))), max_value: Precision::Inexact(ScalarValue::Int32(Some(10))), sum_value: Precision::Absent, @@ -1827,7 +2113,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let filter_statistics = filter.partition_statistics(None)?; + let filter_statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; // First column is "a", and it is a column with only one value after the filter. assert!(filter_statistics.column_statistics[0].is_singleton()); @@ -1874,11 +2161,13 @@ mod tests { Arc::new(Literal::new(ScalarValue::Decimal128(Some(10), 10, 10))), )); let filter = FilterExec::try_new(predicate, input)?; - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(&filter, &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Inexact(200)); assert_eq!(statistics.total_byte_size, Precision::Inexact(800)); let filter = filter.with_default_selectivity(40)?; - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(&filter, &StatisticsArgs::new())?; assert_eq!(statistics.num_rows, Precision::Inexact(400)); assert_eq!(statistics.total_byte_size, Precision::Inexact(1600)); Ok(()) @@ -1913,7 +2202,9 @@ mod tests { Arc::new(EmptyExec::new(Arc::clone(&schema))), )?; - exec.partition_statistics(None).unwrap(); + StatisticsContext::new() + .compute(&exec, &StatisticsArgs::new()) + .unwrap(); Ok(()) } @@ -2069,8 +2360,10 @@ mod tests { assert_eq!(filter1.projection(), filter2.projection()); // Verify statistics are the same - let stats1 = filter1.partition_statistics(None)?; - let stats2 = filter2.partition_statistics(None)?; + let stats1 = + StatisticsContext::new().compute(&filter1, &StatisticsArgs::new())?; + let stats2 = + StatisticsContext::new().compute(&filter2, &StatisticsArgs::new())?; assert_eq!(stats1.num_rows, stats2.num_rows); assert_eq!(stats1.total_byte_size, stats2.total_byte_size); @@ -2123,7 +2416,8 @@ mod tests { .unwrap() .build()?; - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(&filter, &StatisticsArgs::new())?; // Verify statistics reflect both filtering and projection assert!(matches!(statistics.num_rows, Precision::Inexact(_))); @@ -2354,7 +2648,8 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; let col_b_stats = &statistics.column_statistics[1]; assert_eq!(col_b_stats.min_value, Precision::Absent); assert_eq!(col_b_stats.max_value, Precision::Absent); @@ -2433,7 +2728,7 @@ mod tests { vec![Precision::Exact(1)], ), ( - "OR preserves original NDV", + "OR is not collapsed to NDV=1, but NDV is capped at filtered rows", vec![Field::new("name", DataType::Utf8, false)], vec![ColumnStatistics { distinct_count: Precision::Inexact(50), @@ -2452,7 +2747,9 @@ mod tests { Arc::new(Literal::new(ScalarValue::Utf8(Some("b".to_string())))), )), )), - vec![Precision::Inexact(50)], + // Input NDV is 50, but the 20% default selectivity on 100 rows + // estimates 20 output rows, so NDV is capped at 20. + vec![Precision::Inexact(20)], ), ( "AND with mixed types (Utf8 + Int32)", @@ -2639,7 +2936,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = StatisticsContext::new() + .compute(filter.as_ref(), &StatisticsArgs::new())?; for (i, expected) in expected_ndvs.iter().enumerate() { assert_eq!( @@ -2651,11 +2949,211 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_filter_statistics_preserves_exactly_empty_input() -> Result<()> { + // A satisfiable predicate over an exactly empty input: the filter cannot + // produce rows, so the whole estimate stays exact. Column `b` is not + // mentioned by the predicate, so its null and distinct counts go through + // the generic row cap. + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int32, true), + ]); + let input_stats = Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ + ColumnStatistics { + null_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + ..Default::default() + }, + ColumnStatistics { + null_count: Precision::Exact(3), + distinct_count: Precision::Exact(7), + byte_size: Precision::Exact(0), + ..Default::default() + }, + ], + }; + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )); + + let input = Arc::new(StatisticsExec::new(input_stats, schema.clone())); + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Exact(0)); + assert_eq!(statistics.total_byte_size, Precision::Exact(0)); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Exact(0) + ); + assert_eq!( + statistics.column_statistics[1].null_count, + Precision::Exact(0) + ); + assert_eq!( + statistics.column_statistics[1].distinct_count, + Precision::Exact(0) + ); + + // A contradictory predicate (`a = 1 AND a = 2`) discards all rows, the + // output is empty independently of the input. + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(1000), + total_byte_size: Precision::Inexact(8000), + column_statistics: vec![ColumnStatistics::new_unknown(); 2], + }, + schema, + )); + let contradiction = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(1)))), + )), + Operator::And, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(2)))), + )), + )); + let filter: Arc = + Arc::new(FilterExec::try_new(contradiction, input)?); + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Exact(0)); + assert_eq!(statistics.total_byte_size, Precision::Exact(0)); + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_exact_empty_input_zeroes_byte_size() -> Result<()> { + let cases = [ + ("absent", Precision::Absent, Precision::Absent), + ("inexact", Precision::Inexact(8000), Precision::Inexact(400)), + ]; + + for (desc, input_total_byte_size, input_byte_size) in cases { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); + let input_stats = Statistics { + num_rows: Precision::Exact(0), + total_byte_size: input_total_byte_size, + column_statistics: vec![ColumnStatistics { + byte_size: input_byte_size, + ..Default::default() + }], + }; + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Gt, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )); + + let input = Arc::new(StatisticsExec::new(input_stats, schema)); + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = StatisticsContext::new() + .compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!( + statistics.num_rows, + Precision::Exact(0), + "case '{desc}': num_rows mismatch" + ); + assert_eq!( + statistics.total_byte_size, + Precision::Exact(0), + "case '{desc}': total_byte_size mismatch" + ); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Exact(0), + "case '{desc}': byte_size mismatch" + ); + } + + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_empty_input_equality_ndv_zero() -> Result<()> { + let cases: Vec<(&str, Schema, Statistics, Arc)> = vec![ + ( + "fallback string equality", + Schema::new(vec![Field::new("name", DataType::Utf8, true)]), + Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ColumnStatistics { + distinct_count: Precision::Exact(0), + null_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + ..Default::default() + }], + }, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("name", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Utf8(Some("x".to_string())))), + )), + ), + ( + "interval numeric equality", + Schema::new(vec![Field::new("a", DataType::Int32, true)]), + Statistics { + num_rows: Precision::Exact(0), + total_byte_size: Precision::Exact(0), + column_statistics: vec![ColumnStatistics { + min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), + max_value: Precision::Inexact(ScalarValue::Int32(Some(10))), + distinct_count: Precision::Exact(0), + null_count: Precision::Exact(0), + byte_size: Precision::Exact(0), + ..Default::default() + }], + }, + Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::Eq, + Arc::new(Literal::new(ScalarValue::Int32(Some(5)))), + )), + ), + ]; + + for (desc, schema, input_stats, predicate) in cases { + let input = Arc::new(StatisticsExec::new(input_stats, schema)); + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, input)?); + let statistics = StatisticsContext::new() + .compute(filter.as_ref(), &StatisticsArgs::new())?; + + assert_eq!( + statistics.num_rows, + Precision::Exact(0), + "case '{desc}': row count mismatch" + ); + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Exact(0), + "case '{desc}': NDV should be capped at zero rows" + ); + } + Ok(()) + } + #[tokio::test] async fn test_filter_statistics_and_equality_ndv() -> Result<()> { - // a: min=1, max=100, ndv=80 - // b: min=1, max=50, ndv=40 - // c: min=1, max=200, ndv=150 let schema = Schema::new(vec![ Field::new("a", DataType::Int32, false), Field::new("b", DataType::Int32, false), @@ -2669,6 +3167,7 @@ mod tests { ColumnStatistics { min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + null_count: Precision::Inexact(80), distinct_count: Precision::Inexact(80), ..Default::default() }, @@ -2681,6 +3180,7 @@ mod tests { ColumnStatistics { min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), max_value: Precision::Inexact(ScalarValue::Int32(Some(200))), + null_count: Precision::Inexact(90), distinct_count: Precision::Inexact(150), ..Default::default() }, @@ -2713,12 +3213,17 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; - // a = 42 collapses to single value + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + // Equality predicates collapse NDV and reject nulls for their columns. assert_eq!( statistics.column_statistics[0].distinct_count, Precision::Exact(1) ); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); // b > 10 narrows to [11, 50] but doesn't collapse to a single value. // The combined selectivity of a=42 (1/80) and c=7 (1/150) on 100 rows // computes num_rows = 1, so NDV is capped at the row count: min(40, 1) = 1. @@ -2726,11 +3231,14 @@ mod tests { statistics.column_statistics[1].distinct_count, Precision::Inexact(1) ); - // c = 7 collapses to single value assert_eq!( statistics.column_statistics[2].distinct_count, Precision::Exact(1) ); + assert_eq!( + statistics.column_statistics[2].null_count, + Precision::Exact(0) + ); Ok(()) } @@ -2750,8 +3258,8 @@ mod tests { schema.clone(), )); - // a = 42: even without known bounds, interval analysis resolves - // the equality to [42, 42], so NDV is correctly set to Exact(1) + // Even without input bounds, interval analysis can derive singleton + // bounds from the equality itself. let predicate = Arc::new(BinaryExpr::new( Arc::new(Column::new("a", 0)), Operator::Eq, @@ -2759,7 +3267,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!( statistics.column_statistics[0].distinct_count, Precision::Exact(1) @@ -2792,7 +3301,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!( statistics.column_statistics[0].distinct_count, Precision::Exact(1) @@ -2825,7 +3335,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!( statistics.column_statistics[0].distinct_count, Precision::Exact(1) @@ -2858,7 +3369,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!( statistics.column_statistics[0].distinct_count, Precision::Exact(1) @@ -2892,7 +3404,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!( statistics.column_statistics[0].distinct_count, Precision::Exact(1) @@ -2938,7 +3451,8 @@ mod tests { )); let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; assert_eq!( statistics.column_statistics[0].distinct_count, Precision::Exact(1) @@ -3216,16 +3730,17 @@ mod tests { #[tokio::test] async fn test_filter_statistics_ndv_capped_at_row_count() -> Result<()> { - // Table: a: min=1, max=100, distinct_count=80, 100 rows - let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let schema = Schema::new(vec![Field::new("a", DataType::Int32, true)]); let input = Arc::new(StatisticsExec::new( Statistics { num_rows: Precision::Inexact(100), - total_byte_size: Precision::Inexact(400), + total_byte_size: Precision::Inexact(1000), column_statistics: vec![ColumnStatistics { min_value: Precision::Inexact(ScalarValue::Int32(Some(1))), max_value: Precision::Inexact(ScalarValue::Int32(Some(100))), + null_count: Precision::Inexact(80), distinct_count: Precision::Inexact(80), + byte_size: Precision::Exact(1000), ..Default::default() }], }, @@ -3239,15 +3754,154 @@ mod tests { let filter: Arc = Arc::new(FilterExec::try_new(predicate, input)?); - let statistics = filter.partition_statistics(None)?; + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; // Filter estimates ~10 rows (selectivity = 10/100) assert_eq!(statistics.num_rows, Precision::Inexact(10)); - // NDV should be capped at the filtered row count (10), not the original 80 let ndv = &statistics.column_statistics[0].distinct_count; assert!( ndv.get_value().copied() <= Some(10), "Expected NDV <= 10 (filtered row count), got {ndv:?}" ); + // `a <= 10` rejects nulls, so the 80 input nulls drop to exactly zero. + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + // byte_size follows the same 10% selectivity estimate. + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Inexact(100) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_default_selectivity_column_stats() -> Result<()> { + let schema = Schema::new(vec![Field::new("name", DataType::Utf8, true)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1000), + column_statistics: vec![ColumnStatistics { + null_count: Precision::Inexact(80), + distinct_count: Precision::Inexact(60), + byte_size: Precision::Exact(1000), + ..Default::default() + }], + }, + schema.clone(), + )); + + // Utf8 interval analysis is unsupported, so this exercises the default + // selectivity path. The predicate rejects nulls but does not constrain + // the column to one value. + let predicate: Arc = + binary(col("name", &schema)?, Operator::Gt, lit("m"), &schema)?; + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(20)); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Inexact(200) + ); + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Inexact(20) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_or_does_not_reject_nulls() -> Result<()> { + let schema = Schema::new(vec![Field::new("name", DataType::Utf8, true)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1000), + column_statistics: vec![ColumnStatistics { + null_count: Precision::Inexact(80), + distinct_count: Precision::Inexact(60), + byte_size: Precision::Exact(1000), + ..Default::default() + }], + }, + schema.clone(), + )); + + let predicate: Arc = binary( + binary(col("name", &schema)?, Operator::Gt, lit("m"), &schema)?, + Operator::Or, + is_null(col("name", &schema)?)?, + &schema, + )?; + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(20)); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Inexact(20) + ); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Inexact(200) + ); + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Inexact(20) + ); + Ok(()) + } + + #[tokio::test] + async fn test_filter_statistics_is_not_null_rejects_nulls() -> Result<()> { + let schema = Schema::new(vec![Field::new("name", DataType::Utf8, true)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1000), + column_statistics: vec![ColumnStatistics { + null_count: Precision::Inexact(80), + distinct_count: Precision::Inexact(60), + byte_size: Precision::Exact(1000), + ..Default::default() + }], + }, + schema.clone(), + )); + + // `name IS NOT NULL` keeps only non-null rows, so the surviving null + // count is exactly zero. Utf8 interval analysis is unsupported, so this + // also exercises the default-selectivity path. + let predicate: Arc = is_not_null(col("name", &schema)?)?; + let filter: Arc = + Arc::new(FilterExec::try_new(predicate, input)?); + + let statistics = + StatisticsContext::new().compute(filter.as_ref(), &StatisticsArgs::new())?; + assert_eq!(statistics.num_rows, Precision::Inexact(20)); + assert_eq!( + statistics.column_statistics[0].null_count, + Precision::Exact(0) + ); + assert_eq!( + statistics.column_statistics[0].byte_size, + Precision::Inexact(200) + ); + assert_eq!( + statistics.column_statistics[0].distinct_count, + Precision::Inexact(20) + ); Ok(()) } } diff --git a/datafusion/physical-plan/src/filter_pushdown.rs b/datafusion/physical-plan/src/filter_pushdown.rs index 810f9ffcbcdb1..382967c7ee1ef 100644 --- a/datafusion/physical-plan/src/filter_pushdown.rs +++ b/datafusion/physical-plan/src/filter_pushdown.rs @@ -302,6 +302,9 @@ pub struct ChildFilterDescription { /// Description of which parent filters can be pushed down into this node. /// Since we need to transmit filter pushdown results back to this node's parent /// we need to track each parent filter for each child, even those that are unsupported / won't be pushed down. + /// The entries must stay in the same order as the input parent filters: the + /// filter pushdown optimizer maps child results back to parent filters by + /// position. pub(crate) parent_filters: Vec, /// Description of which filters this node is pushing down to its children. /// Since this is not transmitted back to the parents we can have variable sized inner arrays diff --git a/datafusion/physical-plan/src/joins/array_map.rs b/datafusion/physical-plan/src/joins/array_map.rs index ad40d6776df4f..4e56cf013c8f7 100644 --- a/datafusion/physical-plan/src/joins/array_map.rs +++ b/datafusion/physical-plan/src/joins/array_map.rs @@ -89,7 +89,7 @@ macro_rules! downcast_supported_integer { /// ``` /// The resulting `range` (10) correctly represents the size of the interval `[-5, 5]`. /// -/// **2. Index Lookup (in `get_matched_indices`)** +/// **2. Index Lookup (in `get_matched_indices_with_limit_offset`)** /// /// For a probe value of `0` (which is stored as `0u64`): /// ```text @@ -157,13 +157,23 @@ impl ArrayMap { max_val.wrapping_sub(min_val) } + #[inline] + fn key_to_index(key: u64, offset: u64, data_len: usize) -> Option { + let idx = key.wrapping_sub(offset); + if idx < data_len as u64 { + Some(idx as usize) + } else { + None + } + } + /// Creates a new [`ArrayMap`] from the given array of join keys. /// /// Note: This function processes only the non-null values in the input `array`, /// ignoring any rows where the key is `NULL`. /// pub(crate) fn try_new(array: &ArrayRef, min_val: u64, max_val: u64) -> Result { - let range = max_val.wrapping_sub(min_val); + let range = Self::calculate_range(min_val, max_val); if range >= usize::MAX as u64 { return internal_err!("ArrayMap key range is too large to be allocated."); } @@ -207,10 +217,9 @@ impl ArrayMap { for (i, val) in arr.iter().enumerate().rev() { if let Some(val) = val { let key: u64 = val.as_(); - let idx = key.wrapping_sub(offset_val) as usize; - if idx >= data.len() { + let Some(idx) = Self::key_to_index(key, offset_val, data.len()) else { return internal_err!("failed build Array idx >= data.len()"); - } + }; if data[idx] != 0 { if next.is_empty() { @@ -264,6 +273,16 @@ impl ArrayMap { ) } + /// Looks up `key` (a raw probe value cast to `u64`) in the build side, + /// returning the 1-based build-side slot if the key maps to a non-empty + /// bucket, or `None` otherwise. + #[inline] + fn get_value(&self, key: u64) -> Option { + let idx = Self::key_to_index(key, self.offset, self.data.len())?; + let value = self.data[idx]; + (value != 0).then_some(value) + } + fn lookup_and_get_indices( &self, array: &ArrayRef, @@ -294,14 +313,10 @@ impl ArrayMap { } // SAFETY: prob_idx is guaranteed to be within bounds by the loop range. let prob_val: u64 = unsafe { arr.value_unchecked(prob_idx) }.as_(); - let idx_in_build_side = prob_val.wrapping_sub(self.offset) as usize; - - if idx_in_build_side >= self.data.len() - || self.data[idx_in_build_side] == 0 - { + let Some(build_value) = self.get_value(prob_val) else { continue; - } - build_indices.push((self.data[idx_in_build_side] - 1) as u64); + }; + build_indices.push((build_value - 1) as u64); probe_indices.push(prob_idx as u32); } Ok(None) @@ -337,7 +352,7 @@ impl ArrayMap { return Ok(Some((prob_side_idx, None))); } - if arr.is_null(prob_side_idx) { + if have_null && arr.is_null(prob_side_idx) { continue; } @@ -345,14 +360,9 @@ impl ArrayMap { // SAFETY: prob_idx is guaranteed to be within bounds by the loop range. let prob_val: u64 = unsafe { arr.value_unchecked(prob_side_idx) }.as_(); - let idx_in_build_side = prob_val.wrapping_sub(self.offset) as usize; - if idx_in_build_side >= self.data.len() - || self.data[idx_in_build_side] == 0 - { + let Some(build_idx) = self.get_value(prob_val) else { continue; - } - - let build_idx = self.data[idx_in_build_side]; + }; if let Some(offset) = traverse_chain( &self.next, @@ -381,14 +391,14 @@ impl ArrayMap { downcast_supported_integer!( array.data_type() => ( - contain_hashes_helper, + contain_keys_helper, self, array ) ) } - fn contain_hashes_helper( + fn contain_keys_helper( &self, array: &ArrayRef, ) -> Result @@ -402,8 +412,7 @@ impl ArrayMap { } // SAFETY: i is within bounds [0, arr.len()) let key: u64 = unsafe { arr.value_unchecked(i) }.as_(); - let idx = key.wrapping_sub(self.offset) as usize; - idx < self.data.len() && self.data[idx] != 0 + self.get_value(key).is_some() }); Ok(BooleanArray::new(buffer, None)) } @@ -414,6 +423,7 @@ mod tests { use super::*; use arrow::array::Int32Array; use arrow::array::Int64Array; + use arrow::array::UInt64Array; use std::sync::Arc; #[test] @@ -506,6 +516,50 @@ mod tests { Ok(()) } + #[test] + fn test_array_map_rejects_large_out_of_range_probe_key() -> Result<()> { + let build: ArrayRef = + Arc::new(UInt64Array::from(vec![0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])); + let map = ArrayMap::try_new(&build, 0, 10)?; + + assert_eq!(ArrayMap::key_to_index(3, 0, 11), Some(3)); + + // Pick a key for which the computed bucket offset is larger than + // u32::MAX but has low 32 bits equal to 3. It must be bounds-checked + // before casting to usize, otherwise 32-bit targets can truncate it + // into range. + let out_of_range_key = (1_u64 << 32) + 3; + assert_eq!(ArrayMap::key_to_index(out_of_range_key, 0, 11), None); + + let probe = [Arc::new(UInt64Array::from(vec![ + Some(3), + Some(out_of_range_key), + Some(11), + None, + ])) as ArrayRef]; + + let mut matched_probe_indices = vec![]; + let mut matched_build_indices = vec![]; + let next = map.get_matched_indices_with_limit_offset( + &probe, + 10, + (0, None), + &mut matched_probe_indices, + &mut matched_build_indices, + )?; + assert_eq!(matched_probe_indices, vec![0]); + assert_eq!(matched_build_indices, vec![3]); + assert!(next.is_none()); + + let contains = map.contain_keys(&probe)?; + assert!(contains.value(0)); + assert!(!contains.value(1)); + assert!(!contains.value(2)); + assert!(!contains.value(3)); + + Ok(()) + } + #[test] fn test_array_map_i64_with_negative_and_positive_numbers() -> Result<()> { // Build array with a mix of negative and positive i64 values, no duplicates diff --git a/datafusion/physical-plan/src/joins/cross_join.rs b/datafusion/physical-plan/src/joins/cross_join.rs index ab66955dc6034..8a477c1021d1b 100644 --- a/datafusion/physical-plan/src/joins/cross_join.rs +++ b/datafusion/physical-plan/src/joins/cross_join.rs @@ -31,11 +31,13 @@ use crate::projection::{ ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children, physical_to_column_exprs, }; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, check_if_same_properties, handle_state, + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, PlanProperties, RecordBatchStream, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, handle_state, + validate_child_count, }; use arrow::array::{RecordBatch, RecordBatchOptions}; @@ -195,23 +197,6 @@ impl CrossJoinExec { &self.right.schema(), ) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - - Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - left_fut: Default::default(), - cache: Arc::clone(&self.cache), - schema: Arc::clone(&self.schema), - } - } } /// Asynchronously collect the result of the left child @@ -286,21 +271,57 @@ impl ExecutionPlan for CrossJoinExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // CrossJoin has no join conditions or expressions Ok(TreeNodeRecursion::Continue) } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + left_fut: Default::default(), + cache: Arc::clone(&self.cache), + schema: Arc::clone(&self.schema), + })) + } + ChildrenPropertiesMode::Recompute => Ok(Arc::new(CrossJoinExec::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + ))), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(CrossJoinExec::new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - ))) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn reset_state(self: Arc) -> Result> { @@ -316,10 +337,14 @@ impl ExecutionPlan for CrossJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn execute( @@ -381,11 +406,19 @@ impl ExecutionPlan for CrossJoinExec { } } - fn partition_statistics(&self, partition: Option) -> Result> { - // Get the all partitions statistics of the left - let left_stats = Arc::unwrap_or_clone(self.left.partition_statistics(None)?); - let right_stats = - Arc::unwrap_or_clone(self.right.partition_statistics(partition)?); + fn child_stats_requests(&self, partition: Option) -> Vec { + // Left side is always broadcast, so it always needs overall stats. + // Right side is partitioned, so it needs per-partition stats. + vec![ChildStats::At(None), ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let left_stats = input_stats[0].as_ref().clone(); + let right_stats = input_stats[1].as_ref().clone(); Ok(Arc::new(stats_cartesian_product(left_stats, right_stats))) } @@ -430,6 +463,56 @@ impl ExecutionPlan for CrossJoinExec { Arc::new(new_right), )))) } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::CrossJoin(Box::new( + protobuf::CrossJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl CrossJoinExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let crossjoin = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::CrossJoin, + "CrossJoinExec", + ); + + let left = ctx.decode_required_child( + crossjoin.left.as_deref(), + "CrossJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + crossjoin.right.as_deref(), + "CrossJoinExec", + "right", + )?; + + Ok(Arc::new(CrossJoinExec::new(left, right))) + } } /// [left/right]_col_count are required in case the column statistics are None @@ -440,13 +523,14 @@ fn stats_cartesian_product( let left_row_count = left_stats.num_rows; let right_row_count = right_stats.num_rows; - // calculate global stats + // Calculate global stats let num_rows = left_row_count.multiply(&right_row_count); - // the result size is two times a*b because you have the columns of both left and right - let total_byte_size = left_stats - .total_byte_size - .multiply(&right_stats.total_byte_size) - .multiply(&Precision::Exact(2)); + + // Each output row includes every left and right column, so the left side is + // repeated once per right row and the right side once per left row. + let left_byte_size = left_stats.total_byte_size.multiply(&right_row_count); + let right_byte_size = right_stats.total_byte_size.multiply(&left_row_count); + let total_byte_size = left_byte_size.add(&right_byte_size); let left_col_stats = left_stats.column_statistics; let right_col_stats = right_stats.column_statistics; @@ -504,7 +588,7 @@ fn stats_cartesian_product( } } -/// A stream that issues [RecordBatch]es as they arrive from the right of the join. +/// A stream that issues [RecordBatch]es as they arrive from the right of the join. struct CrossJoinStream { /// Input schema schema: Arc, @@ -765,7 +849,9 @@ mod tests { let expected = Statistics { num_rows: Precision::Exact(left_row_count * right_row_count), - total_byte_size: Precision::Exact(2 * left_bytes * right_bytes), + total_byte_size: Precision::Exact( + left_bytes * right_row_count + right_bytes * left_row_count, + ), column_statistics: vec![ ColumnStatistics { distinct_count: Precision::Exact(5), diff --git a/datafusion/physical-plan/src/joins/hash_join/exec.rs b/datafusion/physical-plan/src/joins/hash_join/exec.rs index 3cdd60d7ab3c8..08d209003ad91 100644 --- a/datafusion/physical-plan/src/joins/hash_join/exec.rs +++ b/datafusion/physical-plan/src/joins/hash_join/exec.rs @@ -22,10 +22,9 @@ use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; use std::sync::{Arc, OnceLock}; use std::vec; -use crate::ExecutionPlanProperties; use crate::execution_plan::{ EmissionType, boundedness_from_children, has_same_children_properties, - stub_properties, + plan_contains_expression_id, stub_properties, }; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, @@ -49,13 +48,18 @@ use crate::joins::{JoinOn, JoinOnRef, PartitionMode, SharedBitmapBuilder}; use crate::metrics::{Count, MetricBuilder, MetricCategory}; use crate::projection::{ EmbeddedProjection, JoinData, ProjectionExec, try_embed_projection, - try_pushdown_through_join, + try_pushdown_through_join_with_column_indices, }; use crate::repartition::REPARTITION_RANDOM_STATE; -use crate::spill::get_record_batch_memory_size; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, - PlanProperties, SendableRecordBatchStream, Statistics, + ChildrenPropertiesMode, ExecutionPlanProperties, ReplaceChildrenOptions, + validate_child_count, +}; +use crate::{ + DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + InputDistributionRequirements, Partitioning, PlanProperties, + SendableRecordBatchStream, Statistics, common::can_project, joins::utils::{ BuildProbeJoinMetrics, ColumnIndex, JoinFilter, JoinHashMapType, @@ -73,7 +77,7 @@ use arrow::util::bit_util; use arrow_schema::{DataType, Schema}; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::utils::memory::estimate_memory_size; +use datafusion_common::utils::memory::{RecordBatchMemoryCounter, estimate_memory_size}; use datafusion_common::{ JoinSide, JoinType, NullEquality, Result, assert_or_internal_err, internal_err, plan_err, project_schema, @@ -230,6 +234,23 @@ impl JoinLeftData { &self.batch } + /// Returns `true` if the build side physically contains rows. + /// + /// This is distinct from [`Self::has_matchable_build_rows`]: a build side + /// can hold rows while its hash map is empty (see that method). + pub(super) fn has_build_rows(&self) -> bool { + self.batch().num_rows() > 0 + } + + /// Returns `true` if the build-side hash map has any matchable entries. + /// + /// Under [`NullEquality::NullEqualsNothing`] build rows whose join key is + /// NULL are omitted from the map, so this can be `false` even when + /// [`Self::has_build_rows`] is `true`. + pub(super) fn has_matchable_build_rows(&self) -> bool { + !self.map().is_empty() + } + /// returns a reference to the build side expressions values pub(super) fn values(&self) -> &[ArrayRef] { &self.values @@ -846,14 +867,42 @@ impl HashJoinExec { return false; } - // `preserve_file_partitions` can report Hash partitioning for Hive-style - // file groups, but those partitions are not actually hash-distributed. - // Partitioned dynamic filters rely on hash routing, so disable them in - // this mode to avoid incorrect results. Follow-up work: enable dynamic - // filtering for preserve_file_partitioned scans (issue #20195). + // A null-aware anti join emits a build-side NULL only when the probe + // is truly empty. The pushed filter can empty the probe by pruning + // every row, which would surface that NULL wrongly. A NOT NULL build + // key cannot produce such a NULL, so the filter stays there. + if self.null_aware + && self.on.iter().any(|(build_key, _)| { + build_key.nullable(&self.left.schema()).unwrap_or(true) + }) + { + return false; + } + + // `preserve_file_partitions` can report Hive-style file groups as Hash + // partitioned even though their partition indexes do not follow the + // hash router used by partitioned dynamic filters. Reject Hash inputs + // because the metadata cannot distinguish those scans from a real hash + // repartition. Compatible Range inputs remain safe because matching + // ordering and split points align each build filter with its probe + // partition. Other unsupported layouts are rejected. + // Follow-up work: enable dynamic filtering for preserve_file_partitioned scans (issue #20195). // https://github.com/apache/datafusion/issues/20195 if config.optimizer.preserve_file_partitions > 0 && self.mode == PartitionMode::Partitioned + && matches!( + ( + self.left.output_partitioning(), + self.right.output_partitioning() + ), + (Partitioning::Hash(_, _), Partitioning::Hash(_, _)) + ) + { + return false; + } + + if self.mode == PartitionMode::Partitioned + && !self.has_partitioned_dynamic_filter_routing() { return false; } @@ -861,6 +910,30 @@ impl HashJoinExec { true } + fn has_partitioned_dynamic_filter_routing(&self) -> bool { + match ( + self.left.output_partitioning(), + self.right.output_partitioning(), + ) { + ( + Partitioning::Hash(_, left_partition_count), + Partitioning::Hash(_, right_partition_count), + ) => left_partition_count == right_partition_count, + (Partitioning::Range(_), Partitioning::Range(_)) => { + let children = [self.left.as_ref(), self.right.as_ref()]; + matches!( + self.input_distribution_requirements() + .unsatisfied_co_partitioned_children(self.name(), &children), + Ok(unsatisfied) if unsatisfied.is_empty() + ) + } + (left_partitioning, right_partitioning) => { + left_partitioning.partition_count() == 1 + && right_partitioning.partition_count() == 1 + } + } + } + /// left (build) side which gets hashed pub fn left(&self) -> &Arc { &self.left @@ -902,8 +975,11 @@ impl HashJoinExec { self.null_equality } - /// Get the dynamic filter expression for testing purposes. - /// Returns the dynamic filter expression for this hash join, if set. + /// Returns the dynamic filter expression produced by this hash join, if set. + #[deprecated( + since = "55.0.0", + note = "Use ExecutionPlan::dynamic_expressions_produced instead" + )] pub fn dynamic_filter_expr(&self) -> Option<&Arc> { self.dynamic_filter.as_ref().map(|df| &df.filter) } @@ -1068,6 +1144,12 @@ impl HashJoinExec { &self, partition_mode: PartitionMode, ) -> Result> { + assert_or_internal_err!( + self.dynamic_filter.is_none(), + "Cannot swap HashJoinExec inputs after dynamic filters have been constructed. \ + Optimizer rules that reorder join inputs must run before optimizer rules `FilterPushdown::new_post_optimization()`" + ); + let left = self.left(); let right = self.right(); let new_join = self @@ -1142,6 +1224,8 @@ impl DisplayAs for HashJoinExec { let display_fetch = self .fetch .map_or_else(String::new, |f| format!(", fetch={f}")); + let display_null_aware = + if self.null_aware { ", null_aware" } else { "" }; let on = self .on .iter() @@ -1150,7 +1234,7 @@ impl DisplayAs for HashJoinExec { .join(", "); write!( f, - "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}", + "HashJoinExec: mode={:?}, join_type={:?}, on=[{}]{}{}{}{}{}", self.mode, self.join_type, on, @@ -1158,6 +1242,7 @@ impl DisplayAs for HashJoinExec { display_projections, display_null_equality, display_fetch, + display_null_aware, ) } DisplayFormatType::TreeRender => { @@ -1180,6 +1265,10 @@ impl DisplayAs for HashJoinExec { writeln!(f, "NullsEqual: true")?; } + if self.null_aware { + writeln!(f, "null_aware")?; + } + if let Some(filter) = self.filter.as_ref() { writeln!(f, "filter={filter}")?; } @@ -1204,26 +1293,30 @@ impl ExecutionPlan for HashJoinExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { match self.mode { - PartitionMode::CollectLeft => vec![ - Distribution::SinglePartition, - Distribution::UnspecifiedDistribution, - ], PartitionMode::Partitioned => { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), - ] + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) } - PartitionMode::Auto => vec![ + PartitionMode::CollectLeft => InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, Distribution::UnspecifiedDistribution, + ]), + PartitionMode::Auto => InputDistributionRequirements::new(vec![ Distribution::UnspecifiedDistribution, - ], + Distribution::UnspecifiedDistribution, + ]), } } @@ -1253,26 +1346,31 @@ impl ExecutionPlan for HashJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to join key expressions from both sides - let mut tnr = TreeNodeRecursion::Continue; - for (left, right) in &self.on { - tnr = tnr.visit_sibling(|| f(left.as_ref()))?; - tnr = tnr.visit_sibling(|| f(right.as_ref()))?; - } - - // Apply to join filter expression if present - if let Some(filter) = &self.filter { - tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; - } - - // Apply to dynamic filter expression if present - if let Some(df) = &self.dynamic_filter { - tnr = tnr.visit_sibling(|| f(df.filter.as_ref()))?; - } + let join_keys = self + .on + .iter() + .flat_map(|(left, right)| [Arc::clone(left), Arc::clone(right)]); + let filter = self + .filter + .iter() + .map(|filter| Arc::clone(filter.expression())); + let dynamic_filter = self.dynamic_filter.iter().map(|dynamic_filter| { + Arc::::clone(&dynamic_filter.filter) + as Arc + }); + crate::apply_expression_roots(join_keys.chain(filter).chain(dynamic_filter), f) + } - Ok(tnr) + fn dynamic_expressions_produced(&self) -> Vec> { + self.dynamic_filter + .iter() + .map(|dynamic_filter| { + Arc::::clone(&dynamic_filter.filter) + as Arc + }) + .collect() } /// Creates a new HashJoinExec with different children while preserving configuration. @@ -1280,11 +1378,32 @@ impl ExecutionPlan for HashJoinExec { /// This method is called during query optimization when the optimizer creates new /// plan nodes. Importantly, it creates a fresh bounds_accumulator via `try_new` /// rather than cloning the existing one because partitioning may have changed. + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + self.builder().with_new_children(children)?.build_exec() + } + ChildrenPropertiesMode::Recompute => self + .builder() + .recompute_properties() + .with_new_children(children)? + .build_exec(), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - self.builder().with_new_children(children)?.build_exec() + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn reset_state(self: Arc) -> Result> { @@ -1317,18 +1436,20 @@ impl ExecutionPlan for HashJoinExec { consider using CoalescePartitionsExec or the EnforceDistribution rule" ); - // Only enable dynamic filter pushdown if: - // - The session config enables dynamic filter pushdown - // - A dynamic filter exists - // - At least one consumer is holding a reference to it, this avoids expensive filter - // computation when disabled or when no consumer will use it. - let enable_dynamic_filter_pushdown = self + // Only compute a dynamic filter when the probe subtree contains a consumer. + // Searching from `self` would always find the producer expression owned by this join. + let enable_dynamic_filter_pushdown = if self .allow_join_dynamic_filter_pushdown(context.session_config().options()) - && self - .dynamic_filter + { + self.dynamic_filter .as_ref() - .map(|df| df.filter.is_used()) - .unwrap_or(false); + .and_then(|df| df.filter.expression_id()) + .map(|id| plan_contains_expression_id(&self.right, id)) + .transpose()? + .unwrap_or(false) + } else { + false + }; let join_metrics = BuildProbeJoinMetrics::new(partition, &self.metrics); @@ -1356,6 +1477,8 @@ impl ExecutionPlan for HashJoinExec { filter, on_right, repartition_random_state, + self.null_equality, + self.null_aware, )) }))) }) @@ -1460,56 +1583,43 @@ impl ExecutionPlan for HashJoinExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let stats = match (partition, self.mode) { - // For CollectLeft mode, the left side is collected into a single partition, - // so all left partitions are available to each output partition. - // For the right side, we need the specific partition statistics. - (Some(partition), PartitionMode::CollectLeft) => { - let left_stats = self.left.partition_statistics(None)?; - let right_stats = self.right.partition_statistics(Some(partition))?; - - estimate_join_statistics( - Arc::unwrap_or_clone(left_stats), - Arc::unwrap_or_clone(right_stats), - &self.on, - &self.join_type, - &self.join_schema, - )? + fn child_stats_requests(&self, partition: Option) -> Vec { + match (partition, self.mode) { + // Left side is broadcast, so it always needs overall stats + // Right side is partitioned, so it needs per-partition stats + (Some(_), PartitionMode::CollectLeft) => { + vec![ChildStats::At(None), ChildStats::At(partition)] } - - // For Partitioned mode, both sides are partitioned, so each output partition - // only has access to the corresponding partition from both sides. - (Some(partition), PartitionMode::Partitioned) => { - let left_stats = self.left.partition_statistics(Some(partition))?; - let right_stats = self.right.partition_statistics(Some(partition))?; - - estimate_join_statistics( - Arc::unwrap_or_clone(left_stats), - Arc::unwrap_or_clone(right_stats), - &self.on, - &self.join_type, - &self.join_schema, - )? + // For Partitioned mode, both sides are hash-partitioned symmetrically, + // so each output partition uses the matching partition from both sides. + (Some(_), PartitionMode::Partitioned) => { + vec![ChildStats::At(partition), ChildStats::At(partition)] } - - // For Auto mode or when no specific partition is requested, fall back to - // the current behavior of getting all partition statistics. - (None, _) | (Some(_), PartitionMode::Auto) => { - // TODO stats: it is not possible in general to know the output size of joins - // There are some special cases though, for example: - // - `A LEFT JOIN B ON A.col=B.col` with `COUNT_DISTINCT(B.col)=COUNT(B.col)` - let left_stats = self.left.partition_statistics(None)?; - let right_stats = self.right.partition_statistics(None)?; - estimate_join_statistics( - Arc::unwrap_or_clone(left_stats), - Arc::unwrap_or_clone(right_stats), - &self.on, - &self.join_type, - &self.join_schema, - )? + // Overall stats requested, look up overall child stats. + (None, _) => vec![ChildStats::At(None), ChildStats::At(None)], + // Auto mode hasn't decided partitioning yet, so it needs + // overall stats from both sides. + (Some(_), PartitionMode::Auto) => { + vec![ChildStats::At(None), ChildStats::At(None)] } - }; + } + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let left_stats = Arc::clone(&input_stats[0]); + let right_stats = Arc::clone(&input_stats[1]); + let stats = estimate_join_statistics( + Arc::unwrap_or_clone(left_stats), + Arc::unwrap_or_clone(right_stats), + &self.on, + self.null_equality, + &self.join_type, + &self.join_schema, + )?; // Project statistics if there is a projection let stats = stats.project(self.projection.as_ref()); // Apply fetch limit to statistics @@ -1534,13 +1644,14 @@ impl ExecutionPlan for HashJoinExec { projected_right_child, join_filter, join_on, - }) = try_pushdown_through_join( + }) = try_pushdown_through_join_with_column_indices( projection, self.left(), self.right(), self.on(), &schema, self.filter(), + self.column_indices.as_slice(), )? { self.builder() .with_new_children(vec![ @@ -1601,14 +1712,11 @@ impl ExecutionPlan for HashJoinExec { }; }); - // For semi/anti joins, the non-preserved side's columns are not in the - // output, but filters on join key columns can still be pushed there. - // We find output columns that are join keys on the preserved side and - // add their output indices to the non-preserved side's allowed set. - // The name-based remap in FilterRemapper will then match them to the - // corresponding column in the non-preserved child's schema. + // For semi joins, filters on output join keys can also be pushed to the + // non-output side: every emitted row has an equal key there. This is not + // true for anti joins, whose emitted rows have no match. match self.join_type { - JoinType::LeftSemi | JoinType::LeftAnti => { + JoinType::LeftSemi => { let left_key_indices: HashSet = self .on .iter() @@ -1622,7 +1730,7 @@ impl ExecutionPlan for HashJoinExec { } } } - JoinType::RightSemi | JoinType::RightAnti => { + JoinType::RightSemi => { let right_key_indices: HashSet = self .on .iter() @@ -1660,8 +1768,14 @@ impl ExecutionPlan for HashJoinExec { ChildFilterDescription::all_unsupported(&parent_filters) }; - // Add dynamic filters in Post phase if enabled + // Add dynamic filters in Post phase if enabled. Skip when this join + // already carries a dynamic filter from a previous pass — the shared + // `Arc` is still wired into the probe-side + // scan's predicate, and re-creating it would AND a fresh duplicate + // onto every Post-phase invocation (apache/datafusion-ballista#1359 + // surfaces this in AQE replan loops). if phase == FilterPushdownPhase::Post + && self.dynamic_filter.is_none() && self.allow_join_dynamic_filter_pushdown(config) { // Add actual dynamic filter to right side (probe side) @@ -1723,26 +1837,231 @@ impl ExecutionPlan for HashJoinExec { .ok() .map(|exec| Arc::new(exec) as _) } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + let on = self + .on() + .iter() + .map(|(l, r)| -> Result { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(l)?), + right: Some(ctx.encode_expr(r)?), + }) + }) + .collect::>>()?; + + let join_type = crate::joins::proto::join_type_to_proto(*self.join_type()); + let null_equality = + crate::joins::proto::null_equality_to_proto(self.null_equality()); + // `PartitionMode` is specific to `HashJoinExec`, so its conversion stays + // inline (by-name on purpose: the enums are numbered differently). + let partition_mode = match self.partition_mode() { + PartitionMode::CollectLeft => protobuf::PartitionMode::CollectLeft, + PartitionMode::Partitioned => protobuf::PartitionMode::Partitioned, + PartitionMode::Auto => protobuf::PartitionMode::Auto, + }; + + let filter = self + .filter() + .map(|f| crate::joins::proto::join_filter_to_proto(f, ctx)) + .transpose()?; + + let dynamic_filter = self + .dynamic_expressions_produced() + .into_iter() + .next() + .map(|expr| ctx.encode_expr(&expr)) + .transpose()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::HashJoin(Box::new( + protobuf::HashJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + partition_mode: partition_mode.into(), + null_equality: null_equality.into(), + filter, + // Proto3 `repeated` cannot distinguish `None` from + // `Some(vec![])`. `Some(vec![])` (reachable via + // `try_embed_projection` for e.g. `SELECT count(1) … JOIN …`) + // changes the output schema, so it is encoded with the + // single-element sentinel `[u32::MAX]` (never a valid column + // index); every other state is sent as-is. See + // `try_from_proto` for the matching decoder. + projection: match self.projection.as_ref() { + None => Vec::new(), + Some(v) if v.is_empty() => vec![u32::MAX], + Some(v) => v.iter().map(|x| *x as u32).collect(), + }, + null_aware: self.null_aware, + dynamic_filter, + fetch: self.fetch.map(|f| f as u64), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl HashJoinExec { + /// Reconstruct a [`HashJoinExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::{internal_datafusion_err, plan_datafusion_err}; + use datafusion_proto_models::protobuf; + use std::any::Any; + + let hashjoin = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::HashJoin, + "HashJoinExec", + ); + + let left = + ctx.decode_required_child(hashjoin.left.as_deref(), "HashJoinExec", "left")?; + let right = ctx.decode_required_child( + hashjoin.right.as_deref(), + "HashJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + + let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = hashjoin + .on + .iter() + .map(|col| { + let l = ctx.decode_required_expr( + col.left.as_ref(), + left_schema.as_ref(), + "HashJoinExec", + "on.left", + )?; + let r = ctx.decode_required_expr( + col.right.as_ref(), + right_schema.as_ref(), + "HashJoinExec", + "on.right", + )?; + Ok((l, r)) + }) + .collect::>()?; + + let join_type = crate::joins::proto::join_type_from_proto( + hashjoin.join_type, + "HashJoinExec", + )?; + let null_equality = crate::joins::proto::null_equality_from_proto( + hashjoin.null_equality, + "HashJoinExec", + )?; + // `PartitionMode` is specific to `HashJoinExec`, so its conversion stays + // inline (by-name on purpose: the enums are numbered differently). + let partition_mode = match protobuf::PartitionMode::try_from( + hashjoin.partition_mode, + ) + .map_err(|_| { + internal_datafusion_err!( + "HashJoinExec: unknown PartitionMode {}", + hashjoin.partition_mode + ) + })? { + protobuf::PartitionMode::CollectLeft => PartitionMode::CollectLeft, + protobuf::PartitionMode::Partitioned => PartitionMode::Partitioned, + protobuf::PartitionMode::Auto => PartitionMode::Auto, + }; + + let filter = hashjoin + .filter + .as_ref() + .map(|f| crate::joins::proto::join_filter_from_proto(f, ctx, "HashJoinExec")) + .transpose()?; + + // Preserve the empty-projection sentinel written by `try_to_proto`. + let projection = match hashjoin.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), + }; + + // Restore the row limit that `limit_pushdown` may have pushed into the + // join. The field is presence-tracked, so a message written before it + // existed decodes to `None` (no limit) rather than to `Some(0)`. + // + // The conversion is checked, not `as usize`: `fetch` is a `u64` on the + // wire but a `usize` in the plan, and on a 32-bit target `as usize` + // truncates. A fetch of `1 << 32` would become `0` -- not merely a + // wrong limit but the worst one, silently turning the query into an + // empty result. Report the out-of-range value instead. Please do not + // "simplify" this back to `as usize`. + let fetch = hashjoin + .fetch + .map(|f| { + usize::try_from(f).map_err(|_| { + plan_datafusion_err!( + "HashJoinExec: fetch value {f} cannot be represented as usize on this target" + ) + }) + }) + .transpose()?; + + let mut hash_join = HashJoinExecBuilder::new(left, right, on, join_type) + .with_filter(filter) + .with_projection(projection) + .with_partition_mode(partition_mode) + .with_null_equality(null_equality) + .with_null_aware(hashjoin.null_aware) + .with_fetch(fetch) + .build()?; + + if let Some(dynamic_filter_proto) = &hashjoin.dynamic_filter { + // The dynamic filter is a `DynamicFilterPhysicalExpr` over the probe + // (right) side; decode against the right schema then downcast. + let dynamic_filter_expr = + ctx.decode_expr(dynamic_filter_proto, right_schema.as_ref())?; + let df = (dynamic_filter_expr as Arc) + .downcast::() + .map_err(|_| { + internal_datafusion_err!( + "HashJoinExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + hash_join = hash_join.with_dynamic_filter_expr(df)?; + } + + Ok(Arc::new(hash_join)) + } } /// Determines which sides of a join are "preserved" for filter pushdown. /// /// A preserved side means filters on that side's columns can be safely pushed -/// below the join. This mirrors the logic in the logical optimizer's -/// `lr_is_preserved` in `datafusion/optimizer/src/push_down_filter.rs`. +/// below the join. This mostly mirrors the logical optimizer's `lr_is_preserved`; +/// semi joins additionally allow join-key filters on the non-output side. fn lr_is_preserved(join_type: JoinType) -> (bool, bool) { match join_type { JoinType::Inner => (true, true), JoinType::Left => (true, false), JoinType::Right => (false, true), JoinType::Full => (false, false), - // Filters in semi/anti joins are either on the preserved side, or on join keys, - // as all output columns come from the preserved side. Join key filters can be - // safely pushed down into the other side. - JoinType::LeftSemi | JoinType::LeftAnti => (true, true), - JoinType::RightSemi | JoinType::RightAnti => (true, true), - JoinType::LeftMark => (true, false), - JoinType::RightMark => (false, true), + // Callers restrict the non-output side of semi joins to join-key columns. + JoinType::LeftSemi | JoinType::RightSemi => (true, true), + JoinType::LeftAnti | JoinType::LeftMark => (true, false), + JoinType::RightAnti | JoinType::RightMark => (false, true), } } @@ -1833,6 +2152,10 @@ struct BuildSideState { metrics: BuildProbeJoinMetrics, reservation: MemoryReservation, bounds_accumulators: Option>, + /// Counts the memory of `batches` for `reservation`. Batches can share + /// underlying buffers (e.g. when the input emits zero-copy slices of one + /// larger batch), so each buffer must be reserved only once. + memory_counter: RecordBatchMemoryCounter, } impl BuildSideState { @@ -1849,6 +2172,7 @@ impl BuildSideState { num_rows: 0, metrics, reservation, + memory_counter: RecordBatchMemoryCounter::new(), bounds_accumulators: should_compute_dynamic_filters .then(|| { on_left @@ -1939,7 +2263,7 @@ async fn collect_left_input( } // Decide if we spill or not - let batch_size = get_record_batch_memory_size(&batch); + let batch_size = state.memory_counter.count_batch(&batch); // Reserve memory for incoming batch state.reservation.try_grow(batch_size)?; // Update metrics @@ -1961,6 +2285,7 @@ async fn collect_left_input( metrics, mut reservation, bounds_accumulators, + memory_counter: _, } = state; // Compute bounds @@ -2031,6 +2356,7 @@ async fn collect_left_input( &mut hashes_buffer, 0, true, + null_equality, )?; offset += batch.num_rows(); } @@ -2145,8 +2471,11 @@ mod tests { } use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::execution_plan::Boundedness; + use crate::filter::FilterExecBuilder; use crate::joins::hash_join::stream::lookup_join_hashmap; use crate::test::{TestMemoryExec, assert_join_metrics}; + use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use crate::{ common, expressions::Column, repartition::RepartitionExec, test::build_table_i32, test::exec::MockExec, @@ -2167,11 +2496,85 @@ mod tests { use datafusion_execution::runtime_env::RuntimeEnvBuilder; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::{BinaryExpr, Literal}; + use datafusion_physical_expr::{ + EquivalenceProperties, PhysicalSortExpr, RangePartitioning, SplitPoint, + }; use hashbrown::HashTable; use insta::{allow_duplicates, assert_snapshot}; use rstest::*; use rstest_reuse::*; + #[derive(Debug)] + struct PartitionedTestExec { + cache: Arc, + } + + impl PartitionedTestExec { + fn try_new(schema: SchemaRef, partitioning: Partitioning) -> Result { + Ok(Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new(Arc::clone(&schema)), + partitioning, + EmissionType::Incremental, + Boundedness::Bounded, + )), + }) + } + } + + impl DisplayAs for PartitionedTestExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "PartitionedTestExec") + } + } + + impl ExecutionPlan for PartitionedTestExec { + fn name(&self) -> &'static str { + "PartitionedTestExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } + } + fn div_ceil(a: usize, b: usize) -> usize { a.div_ceil(b) } @@ -2304,7 +2707,12 @@ mod tests { Arc::new(Column::new_with_schema("b1", &right_schema).unwrap()) as _, )]; let right: Arc = Arc::new( - MockExec::new(vec![Ok(right_batch), err], right_schema).with_use_task(false), + MockExec::new(vec![Ok(right_batch), err], right_schema) + .with_use_task(false) + // The planted error must only surface if the probe side is + // polled, not when a parent node computes statistics during + // planning. + .with_unknown_statistics(), ); (left, right, on) @@ -2384,6 +2792,8 @@ mod tests { mode: PartitionMode, ) -> Result<(HashJoinExec, Arc)> { let dynamic_filter = HashJoinExec::create_dynamic_filter(&on); + let consumer: Arc = Arc::clone(&dynamic_filter) as _; + let right = Arc::new(FilterExecBuilder::new(consumer, right).build()?); let mut join = HashJoinExec::try_new( left, right, @@ -3345,6 +3755,171 @@ mod tests { Ok(()) } + /// Under NullEqualsNothing, NULL join keys are not inserted into the hash + /// map, so a build side whose keys are all NULL produces an empty map even + /// though it contains rows. Join types that emit unmatched build rows must + /// still produce them from the visited bitmap. + #[rstest] + #[tokio::test] + async fn join_all_null_build_keys( + #[values(PartitionMode::CollectLeft, PartitionMode::Partitioned)] + partition_mode: PartitionMode, + ) -> Result<()> { + let left = build_table_two_cols( + ("a1", &vec![Some(1), Some(2)]), + ("b1", &vec![None, None]), // all build-side join keys are NULL + ); + let right = build_table_two_cols( + ("a2", &vec![Some(10), Some(20), Some(30)]), + ("b1", &vec![Some(4), None, Some(6)]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::LeftAnti, + JoinType::RightSemi, + JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, + ] { + let (_, batches, metrics) = join_collect_with_partition_mode( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + &join_type, + partition_mode, + NullEquality::NullEqualsNothing, + Arc::new(TaskContext::default()), + ) + .await?; + + // For join types whose output requires a build-side match, an + // empty map guarantees an empty result, so `state_after_build_ready` + // completes the stream without ever fetching a probe batch (probe + // `input_rows` stays 0). All other join types must still scan the + // probe side. `input_rows` is summed across every partition. + let probe_rows = metrics + .sum_by_name("input_rows") + .map(|v| v.as_usize()) + .unwrap_or(0); + if join_type.empty_map_produces_empty_result() { + assert_eq!( + probe_rows, 0, + "{join_type} should skip the probe side for an all-NULL build" + ); + } else { + assert!(probe_rows > 0, "{join_type} must scan the probe side"); + } + + match join_type { + JoinType::Inner | JoinType::LeftSemi | JoinType::RightSemi => { + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, 0, "unexpected rows for {join_type}"); + } + JoinType::Left => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b1 | + +----+----+----+----+ + | 1 | | | | + | 2 | | | | + +----+----+----+----+ + "); + } + } + JoinType::Right => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b1 | + +----+----+----+----+ + | | | 10 | 4 | + | | | 20 | | + | | | 30 | 6 | + +----+----+----+----+ + "); + } + } + JoinType::Full => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+----+----+ + | a1 | b1 | a2 | b1 | + +----+----+----+----+ + | | | 10 | 4 | + | | | 20 | | + | | | 30 | 6 | + | 1 | | | | + | 2 | | | | + +----+----+----+----+ + "); + } + } + JoinType::LeftAnti => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+ + | a1 | b1 | + +----+----+ + | 1 | | + | 2 | | + +----+----+ + "); + } + } + JoinType::RightAnti => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+ + | a2 | b1 | + +----+----+ + | 10 | 4 | + | 20 | | + | 30 | 6 | + +----+----+ + "); + } + } + JoinType::LeftMark => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-------+ + | a1 | b1 | mark | + +----+----+-------+ + | 1 | | false | + | 2 | | false | + +----+----+-------+ + "); + } + } + JoinType::RightMark => { + allow_duplicates! { + assert_snapshot!(batches_to_sort_string(&batches), @r" + +----+----+-------+ + | a2 | b1 | mark | + +----+----+-------+ + | 10 | 4 | false | + | 20 | | false | + | 30 | 6 | false | + +----+----+-------+ + "); + } + } + } + } + + Ok(()) + } + #[apply(hash_join_exec_configs)] #[tokio::test] async fn partitioned_join_left_one( @@ -4475,6 +5050,7 @@ mod tests { &[right_keys_values], NullEquality::NullEqualsNothing, &hashes_buffer, + None, 8192, (0, None), &mut probe_indices_buffer, @@ -4536,6 +5112,7 @@ mod tests { &[right_keys_values], NullEquality::NullEqualsNothing, &hashes_buffer, + None, 8192, (0, None), &mut probe_indices_buffer, @@ -5385,6 +5962,61 @@ mod tests { Ok(()) } + #[tokio::test] + async fn build_side_sliced_batches_memory_accounting() -> Result<()> { + // The build side emits zero-copy slices of one large batch, as e.g. an + // aggregate emitting its output in batch_size chunks does. The buffers + // shared by the slices must be reserved once in total, not once per + // slice: per-slice accounting reserves number_of_slices x parent size + // and aborts queries that fit in memory with room to spare. + let n = 4096; + let v: Vec = (0..n).collect(); + let parent = build_table_i32(("a1", &v), ("b1", &v), ("c1", &v)); + let slices: Vec = + (0..16).map(|i| parent.slice(i * 256, 256)).collect(); + let left = + TestMemoryExec::try_new_exec(&[slices], parent.schema(), None).unwrap(); + + let right_batch = build_table_i32( + ("a2", &vec![10, 11]), + ("b2", &vec![0, 1]), + ("c2", &vec![14, 15]), + ); + let right = TestMemoryExec::try_new_exec( + &[vec![right_batch.clone()]], + right_batch.schema(), + None, + ) + .unwrap(); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &parent.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right_batch.schema())?) as _, + )]; + + // Enough for the parent batch (~48KB) plus the join hash table, but far + // below the ~768KB that per-slice accounting would reserve + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(400_000, 1.0) + .build_arc()?; + let task_ctx = TaskContext::default().with_runtime(runtime); + let task_ctx = Arc::new(task_ctx); + + let join = join( + left, + right, + on, + &JoinType::Inner, + NullEquality::NullEqualsNothing, + )?; + + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + let num_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(num_rows, 2); + + Ok(()) + } + #[tokio::test] async fn partitioned_join_overallocation() -> Result<()> { // Prepare partitioned inputs for HashJoinExec @@ -6321,10 +6953,10 @@ mod tests { assert_eq!(lr_is_preserved(JoinType::Right), (false, true)); assert_eq!(lr_is_preserved(JoinType::Full), (false, false)); assert_eq!(lr_is_preserved(JoinType::LeftSemi), (true, true)); - assert_eq!(lr_is_preserved(JoinType::LeftAnti), (true, true)); + assert_eq!(lr_is_preserved(JoinType::LeftAnti), (true, false)); assert_eq!(lr_is_preserved(JoinType::LeftMark), (true, false)); assert_eq!(lr_is_preserved(JoinType::RightSemi), (true, true)); - assert_eq!(lr_is_preserved(JoinType::RightAnti), (true, true)); + assert_eq!(lr_is_preserved(JoinType::RightAnti), (false, true)); assert_eq!(lr_is_preserved(JoinType::RightMark), (false, true)); } @@ -6345,7 +6977,7 @@ mod tests { NullEquality::NullEqualsNothing, false, )?; - assert!(join.dynamic_filter_expr().is_none()); + assert!(join.dynamic_expressions_produced().is_empty()); let df = Arc::new(DynamicFilterPhysicalExpr::new( vec![Arc::new(Column::new("b1", 1)) as _], @@ -6353,11 +6985,10 @@ mod tests { )); let join = join.with_dynamic_filter_expr(Arc::clone(&df))?; - let restored = join - .dynamic_filter_expr() - .expect("should have dynamic filter"); + let produced = join.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); assert_eq!( - restored + produced[0] .expression_id() .expect("DynamicFilterPhysicalExpr always has an expression_id"), df.expression_id() @@ -6366,6 +6997,261 @@ mod tests { Ok(()) } + #[test] + fn test_swap_inputs_rejects_dynamic_filter() -> Result<()> { + let left = build_table( + ("l_key", &vec![1]), + ("l_payload", &vec![10]), + ("l_other", &vec![100]), + ); + let right = build_table( + ("r_payload", &vec![20]), + ("r_key", &vec![1]), + ("r_other", &vec![200]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("l_key", &left.schema())?) as _, + Arc::new(Column::new_with_schema("r_key", &right.schema())?) as _, + )]; + + let dynamic_filter = HashJoinExec::create_dynamic_filter(&on); + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftSemi, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )? + .with_dynamic_filter_expr(dynamic_filter)?; + + let err = join.swap_inputs(PartitionMode::CollectLeft).unwrap_err(); + assert_contains!( + err.to_string(), + "Cannot swap HashJoinExec inputs after dynamic filters have been constructed" + ); + Ok(()) + } + + #[test] + fn test_dynamic_filter_pushdown_allowed_for_null_equal_join() -> Result<()> { + let (_, _, on) = build_schema_and_on()?; + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![1])); + let right = build_table(("a2", &vec![1]), ("b1", &vec![1]), ("c2", &vec![1])); + + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::RightSemi, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNull, + false, + )?; + + // Null-equal joins keep dynamic filter pushdown: the pushed predicate carries an + // `IS NULL` disjunct so a probe-side NULL still reaches the join. + assert!(join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn test_dynamic_filter_pushdown_rejects_null_aware_nullable_build_key() -> Result<()> + { + let left = build_table_two_cols( + ("a1", &vec![Some(1), None]), + ("b1", &vec![Some(1), Some(2)]), + ); + let right = build_table_two_cols( + ("a2", &vec![Some(2), Some(3)]), + ("b2", &vec![Some(1), Some(2)]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + assert!(!join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn test_dynamic_filter_pushdown_allows_null_aware_non_null_build_key() -> Result<()> { + // A NOT NULL build key cannot surface a build-side NULL, so the + // pushdown must stay enabled. + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![1])); + let right = build_table(("a2", &vec![2]), ("b2", &vec![2]), ("c2", &vec![2])); + let on = vec![( + Arc::new(Column::new_with_schema("a1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("a2", &right.schema())?) as _, + )]; + + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + let join = HashJoinExec::try_new( + left, + right, + on, + None, + &JoinType::LeftAnti, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + true, + )?; + + assert!(join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + fn range_partitioned_dynamic_filter_test_join( + left_split: i32, + right_split: i32, + ) -> Result<(HashJoinExec, JoinOn)> { + let (left_schema, right_schema, on) = build_schema_and_on()?; + let left_partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr { + expr: Arc::clone(&on[0].0), + options: Default::default(), + }] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(left_split))])], + )?); + let right_partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr { + expr: Arc::clone(&on[0].1), + options: Default::default(), + }] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(right_split))])], + )?); + let left = Arc::new(PartitionedTestExec::try_new( + left_schema, + left_partitioning, + )?); + let right = Arc::new(PartitionedTestExec::try_new( + right_schema, + right_partitioning, + )?); + + let join = HashJoinExec::try_new( + left, + right, + on.clone(), + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + )?; + Ok((join, on)) + } + + fn with_hash_partitioned_children( + join: &HashJoinExec, + on: &JoinOn, + ) -> Result { + join.builder() + .with_new_children(vec![ + Arc::new(PartitionedTestExec::try_new( + join.left().schema(), + Partitioning::Hash(vec![Arc::clone(&on[0].0)], 2), + )?), + Arc::new(PartitionedTestExec::try_new( + join.right().schema(), + Partitioning::Hash(vec![Arc::clone(&on[0].1)], 2), + )?), + ])? + .build() + } + + #[test] + fn test_partitioned_dynamic_filter_pushdown_allows_supported_partitioning() + -> Result<()> { + let (range_join, on) = range_partitioned_dynamic_filter_test_join(10, 10)?; + let hash_join = with_hash_partitioned_children(&range_join, &on)?; + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + assert!(range_join.allow_join_dynamic_filter_pushdown(session_config.options())); + assert!(hash_join.allow_join_dynamic_filter_pushdown(session_config.options())); + + session_config + .options_mut() + .optimizer + .preserve_file_partitions = 1; + assert!(range_join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + + #[test] + fn test_partitioned_dynamic_filter_pushdown_rejects_unsupported_partitioning() + -> Result<()> { + let (range_join, on) = range_partitioned_dynamic_filter_test_join(10, 10)?; + let hash_join = with_hash_partitioned_children(&range_join, &on)?; + let (mismatched_range_join, _) = + range_partitioned_dynamic_filter_test_join(10, 11)?; + let mut session_config = SessionConfig::default(); + session_config + .options_mut() + .optimizer + .enable_join_dynamic_filter_pushdown = true; + + assert!( + !mismatched_range_join + .allow_join_dynamic_filter_pushdown(session_config.options()) + ); + + session_config + .options_mut() + .optimizer + .preserve_file_partitions = 1; + assert!(!hash_join.allow_join_dynamic_filter_pushdown(session_config.options())); + + Ok(()) + } + #[test] fn test_with_dynamic_filter_rejects_invalid_columns() -> Result<()> { let (_, _, on) = build_schema_and_on()?; diff --git a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs index 0daac0bb86a75..60a25fc2efcff 100644 --- a/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs +++ b/datafusion/physical-plan/src/joins/hash_join/partitioned_hash_eval.rs @@ -27,6 +27,8 @@ use arrow::{ use datafusion_common::Result; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::{create_hashes, with_hashes}; +#[cfg(feature = "proto")] +use datafusion_common::internal_err; use datafusion_expr::ColumnarValue; use datafusion_physical_expr_common::physical_expr::{ DynHash, PhysicalExpr, PhysicalExprRef, @@ -199,6 +201,55 @@ impl PhysicalExpr for HashExpr { fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.description) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let on_columns = ctx.encode_children_expressions(&self.on_columns)?; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::HashExpr( + protobuf::PhysicalHashExprNode { + on_columns, + seed0: self.seed(), + description: self.description.clone(), + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl HashExpr { + /// Reconstruct a [`HashExpr`] from its protobuf representation. + /// + /// Takes the whole [`PhysicalExprNode`], the exact inverse of what + /// [`PhysicalExpr::try_to_proto`] produces, so every expression's + /// `try_from_proto` shares one signature. Child sub-expressions are + /// decoded recursively via [`PhysicalExprDecodeCtx::decode`]. + /// + /// [`PhysicalExprNode`]: datafusion_proto_models::protobuf::PhysicalExprNode + /// [`PhysicalExpr::try_to_proto`]: datafusion_physical_expr_common::physical_expr::PhysicalExpr::try_to_proto + /// [`PhysicalExprDecodeCtx::decode`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::decode + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let hash_expr = match &node.expr_type { + Some(protobuf::physical_expr_node::ExprType::HashExpr(h)) => h, + _ => return internal_err!("PhysicalExprNode is not a HashExpr"), + }; + let on_columns = ctx.decode_children_expressions(&hash_expr.on_columns)?; + Ok(Arc::new(HashExpr::new( + on_columns, + SeededRandomState::with_seed(hash_expr.seed0), + hash_expr.description.clone(), + ))) + } } /// Physical expression that checks join keys in a [`Map`] (hash table or array map). @@ -215,7 +266,6 @@ pub struct HashTableLookupExpr { /// Description for display description: String, } - impl HashTableLookupExpr { /// Create a new HashTableLookupExpr /// @@ -241,7 +291,6 @@ impl HashTableLookupExpr { } } } - impl std::fmt::Debug for HashTableLookupExpr { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { let cols = self @@ -337,7 +386,38 @@ impl PhysicalExpr for HashTableLookupExpr { } } } - + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use datafusion_proto_models::protobuf::physical_expr_node::ExprType; + + // HashTableLookupExpr holds a runtime Arc (the build-side hash + // table) that cannot be serialized, so it is replaced with lit(true). + // + // Dynamic filtering is a performance optimisation only — replacing the + // lookup with lit(true) preserves correctness by allowing all rows + // through. + // + // If a plan is serialized before execution, HashTableLookupExpr is not + // yet present in the dynamic filter expression. + // + // If a plan is serialized after execution, any runtime-created + // HashTableLookupExpr is replaced during serialization. Re-executing + // the plan requires reset_state(), after which HashJoinExec rebuilds + // fresh dynamic filters at runtime. + let value = datafusion_proto_common::ScalarValue { + value: Some(datafusion_proto_common::scalar_value::Value::BoolValue( + true, + )), + }; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(ExprType::Literal(value)), + })) + } fn fmt_sql(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { write!(f, "{}", self.description) } @@ -469,6 +549,172 @@ mod tests { assert_eq!(compute_hash(&expr1), compute_hash(&expr2)); } + #[cfg(feature = "proto")] + mod proto_tests { + use super::*; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::internal_datafusion_err; + use datafusion_physical_expr_common::physical_expr::proto_decode::{ + PhysicalExprDecode, PhysicalExprDecodeCtx, + }; + use datafusion_physical_expr_common::physical_expr::proto_encode::{ + PhysicalExprEncode, PhysicalExprEncodeCtx, + }; + use datafusion_proto_models::protobuf; + + struct TestEncoder; + + impl PhysicalExprEncode for TestEncoder { + fn encode( + &self, + expr: &Arc, + ) -> Result { + let ctx = PhysicalExprEncodeCtx::new(self); + expr.try_to_proto(&ctx)?.ok_or_else(|| { + internal_datafusion_err!("test encoder cannot encode {expr:?}") + }) + } + } + + struct TestDecoder; + + impl PhysicalExprDecode for TestDecoder { + fn decode( + &self, + node: &protobuf::PhysicalExprNode, + schema: &Schema, + ) -> Result> { + let ctx = PhysicalExprDecodeCtx::new(schema, self); + match &node.expr_type { + Some(protobuf::physical_expr_node::ExprType::Column(_)) => { + Column::try_from_proto(node, &ctx) + } + _ => internal_err!("test decoder cannot decode {node:?}"), + } + } + } + + fn test_decode_ctx<'a>( + schema: &'a Schema, + decoder: &'a TestDecoder, + ) -> PhysicalExprDecodeCtx<'a> { + PhysicalExprDecodeCtx::new(schema, decoder) + } + + #[test] + fn hash_expr_try_to_proto() { + let expr = HashExpr::new( + vec![Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1))], + SeededRandomState::with_seed(42), + "hash_join".to_string(), + ); + let encoder = TestEncoder; + let ctx = PhysicalExprEncodeCtx::new(&encoder); + + let proto = expr.try_to_proto(&ctx).unwrap().unwrap(); + + assert_eq!(proto.expr_id, None); + let hash_expr = match proto.expr_type.unwrap() { + protobuf::physical_expr_node::ExprType::HashExpr(hash_expr) => hash_expr, + other => panic!("expected HashExpr, got {other:?}"), + }; + assert_eq!(hash_expr.seed0, 42); + assert_eq!(hash_expr.description, "hash_join"); + assert_eq!(hash_expr.on_columns.len(), 2); + assert!( + hash_expr + .on_columns + .iter() + .all(|expr| expr.expr_id.is_none()) + ); + } + + #[test] + fn hash_expr_try_from_proto() { + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Utf8, true), + ]); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + let proto = protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::HashExpr( + protobuf::PhysicalHashExprNode { + on_columns: vec![ + protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some( + protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: "a".to_string(), + index: 0, + }, + ), + ), + }, + protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some( + protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: "b".to_string(), + index: 1, + }, + ), + ), + }, + ], + seed0: 42, + description: "hash_join".to_string(), + }, + )), + }; + + let expr = HashExpr::try_from_proto(&proto, &ctx).unwrap(); + let expr = expr.downcast_ref::().unwrap(); + + assert_eq!(expr.seed(), 42); + assert_eq!(expr.description(), "hash_join"); + assert_eq!(expr.on_columns().len(), 2); + assert_eq!( + expr.on_columns()[0] + .downcast_ref::() + .map(|col| (col.name(), col.index())), + Some(("a", 0)) + ); + assert_eq!( + expr.on_columns()[1] + .downcast_ref::() + .map(|col| (col.name(), col.index())), + Some(("b", 1)) + ); + } + + #[test] + fn hash_expr_try_from_proto_rejects_wrong_node_type() { + let schema = Schema::empty(); + let decoder = TestDecoder; + let ctx = test_decode_ctx(&schema, &decoder); + let proto = protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Column( + protobuf::PhysicalColumn { + name: "a".to_string(), + index: 0, + }, + )), + }; + + let err = HashExpr::try_from_proto(&proto, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalExprNode is not a HashExpr"), + "{err}" + ); + } + } + #[test] fn test_hash_table_lookup_expr_eq_same() { let col_a: PhysicalExprRef = Arc::new(Column::new("a", 0)); diff --git a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs index fba6b2c2db2e2..94ec4565a4cef 100644 --- a/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs +++ b/datafusion/physical-plan/src/joins/hash_join/shared_bounds.rs @@ -23,6 +23,7 @@ use std::sync::Arc; use crate::ExecutionPlan; use crate::ExecutionPlanProperties; +use crate::Partitioning; use crate::joins::Map; use crate::joins::PartitionMode; use crate::joins::hash_join::exec::HASH_JOIN_SEED; @@ -30,16 +31,22 @@ use crate::joins::hash_join::inlist_builder::build_struct_fields; use crate::joins::hash_join::partitioned_hash_eval::{ HashExpr, HashTableLookupExpr, SeededRandomState, }; +use crate::repartition::RangeExpr; use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::config::ConfigOptions; -use datafusion_common::{DataFusionError, Result, ScalarValue, SharedResult}; +use datafusion_common::{ + DataFusionError, NullEquality, Result, ScalarValue, SharedResult, + assert_or_internal_err, +}; use datafusion_expr::Operator; use datafusion_functions::core::r#struct as struct_func; use datafusion_physical_expr::expressions::{ - BinaryExpr, CaseExpr, DynamicFilterPhysicalExpr, InListExpr, lit, + BinaryExpr, CaseExpr, DynamicFilterPhysicalExpr, InListExpr, IsNullExpr, lit, +}; +use datafusion_physical_expr::{ + PhysicalExpr, PhysicalExprRef, RangePartitioning, ScalarFunctionExpr, }; -use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef, ScalarFunctionExpr}; use parking_lot::Mutex; use tokio::sync::Notify; @@ -255,6 +262,14 @@ pub(crate) struct SharedBuildAccumulator { repartition_random_state: SeededRandomState, /// Schema of the probe (right) side for evaluating filter expressions probe_schema: Arc, + /// Probe-side Range routing metadata for partitioned dynamic filters. + probe_range_partitioning: Option, + /// Null equality of the join. Under `NullEqualsNull` a probe-side NULL can match a + /// build-side NULL, so the pushed filter must keep NULL rows here too. + null_equality: NullEquality, + /// Null-aware anti join (`NOT IN`). A probe-side NULL must reach the join so its + /// three-valued logic can collapse the result, so the pushed filter keeps NULL rows. + null_aware: bool, } /// Strategy for filter pushdown (decided at collection time) @@ -274,10 +289,12 @@ pub(crate) enum PartitionBuildData { partition_id: usize, pushdown: PushdownStrategy, bounds: PartitionBounds, + keys_have_null: bool, }, CollectLeft { pushdown: PushdownStrategy, bounds: PartitionBounds, + keys_have_null: bool, }, } @@ -286,6 +303,9 @@ pub(crate) enum PartitionBuildData { struct PartitionData { bounds: PartitionBounds, pushdown: PushdownStrategy, + /// Whether any build key of this partition is NULL. Decides whether the pushed + /// filter must keep probe-side NULL rows for a null-equal join to match them. + keys_have_null: bool, } /// Build-side data organized by partition mode @@ -351,6 +371,7 @@ impl SharedBuildAccumulator { /// We cannot build a partial filter from some partitions - it would incorrectly eliminate /// valid join results. We must wait until we have complete information from ALL /// relevant partitions before updating the dynamic filter. + #[expect(clippy::too_many_arguments)] pub(crate) fn new_from_partition_mode( partition_mode: PartitionMode, left_child: &dyn ExecutionPlan, @@ -358,6 +379,8 @@ impl SharedBuildAccumulator { dynamic_filter: Arc, on_right: Vec, repartition_random_state: SeededRandomState, + null_equality: NullEquality, + null_aware: bool, ) -> Self { // Troubleshooting: If partition counts are incorrect, verify this logic matches // the actual execution pattern in collect_build_side() @@ -394,6 +417,14 @@ impl SharedBuildAccumulator { ), }; + let probe_range_partitioning = + match (partition_mode, right_child.output_partitioning()) { + (PartitionMode::Partitioned, Partitioning::Range(range)) => { + Some(range.clone()) + } + _ => None, + }; + Self { inner: Mutex::new(AccumulatorState { data: mode_data, @@ -404,6 +435,9 @@ impl SharedBuildAccumulator { on_right, repartition_random_state, probe_schema: right_child.schema(), + probe_range_partitioning, + null_equality, + null_aware, } } @@ -456,6 +490,7 @@ impl SharedBuildAccumulator { partition_id, pushdown, bounds, + keys_have_null, }, AccumulatedBuildData::Partitioned { partitions, @@ -465,11 +500,18 @@ impl SharedBuildAccumulator { if matches!(partitions[partition_id], PartitionStatus::Pending) { *completed_partitions += 1; } - partitions[partition_id] = - PartitionStatus::Reported(PartitionData { pushdown, bounds }); + partitions[partition_id] = PartitionStatus::Reported(PartitionData { + pushdown, + bounds, + keys_have_null, + }); } ( - PartitionBuildData::CollectLeft { pushdown, bounds }, + PartitionBuildData::CollectLeft { + pushdown, + bounds, + keys_have_null, + }, AccumulatedBuildData::CollectLeft { data, reported_count, @@ -477,7 +519,11 @@ impl SharedBuildAccumulator { }, ) => { if matches!(data, PartitionStatus::Pending) { - *data = PartitionStatus::Reported(PartitionData { pushdown, bounds }); + *data = PartitionStatus::Reported(PartitionData { + pushdown, + bounds, + keys_have_null, + }); } *reported_count += 1; } @@ -579,7 +625,10 @@ impl SharedBuildAccumulator { if let Some(filter_expr) = combine_membership_and_bounds(membership_expr, bounds_expr) { - self.dynamic_filter.update(filter_expr)?; + self.dynamic_filter.update(self.preserve_probe_nulls( + filter_expr, + partition_data.keys_have_null, + )?)?; } } PartitionStatus::Pending => { @@ -595,21 +644,11 @@ impl SharedBuildAccumulator { }, FinalizeInput::Partitioned(partitions) => { let num_partitions = partitions.len(); - let routing_hash_expr = Arc::new(HashExpr::new( - self.on_right.clone(), - self.repartition_random_state.clone(), - "hash_repartition".to_string(), - )) as Arc; - - let modulo_expr = Arc::new(BinaryExpr::new( - routing_hash_expr, - Operator::Modulo, - lit(ScalarValue::UInt64(Some(num_partitions as u64))), - )) as Arc; - - let mut real_branches = Vec::new(); + let mut partition_filters = Vec::with_capacity(num_partitions); + let mut real_partition_ids = Vec::new(); let mut empty_partition_ids = Vec::new(); let mut has_canceled_unknown = false; + let mut keys_have_null = false; for (partition_id, partition) in partitions.iter().enumerate() { match partition { @@ -617,8 +656,11 @@ impl SharedBuildAccumulator { if matches!(partition.pushdown, PushdownStrategy::Empty) => { empty_partition_ids.push(partition_id); + partition_filters.push(lit(false)); } PartitionStatus::Reported(partition) => { + real_partition_ids.push(partition_id); + keys_have_null |= partition.keys_have_null; let membership_expr = create_membership_predicate( &self.on_right, partition.pushdown.clone(), @@ -634,13 +676,14 @@ impl SharedBuildAccumulator { bounds_expr, ) .unwrap_or_else(|| lit(true)); - real_branches.push(( - lit(ScalarValue::UInt64(Some(partition_id as u64))), - then_expr, - )); + partition_filters.push(then_expr); } PartitionStatus::CanceledUnknown => { has_canceled_unknown = true; + partition_filters.push(lit(true)); + // A canceled partition's build content is unknown, so it + // may hold a NULL key. + keys_have_null = true; } PartitionStatus::Pending => { return datafusion_common::internal_err!( @@ -650,47 +693,155 @@ impl SharedBuildAccumulator { } } - let filter_expr = if has_canceled_unknown { - let mut when_then_branches = empty_partition_ids + let filter_expr = if has_canceled_unknown + && real_partition_ids.is_empty() + && empty_partition_ids.is_empty() + { + lit(true) + } else if !has_canceled_unknown && real_partition_ids.is_empty() { + lit(false) + } else if !has_canceled_unknown + && real_partition_ids.len() == 1 + && empty_partition_ids.len() + 1 == num_partitions + { + Arc::clone(&partition_filters[real_partition_ids[0]]) + } else if let Some(range_partitioning) = &self.probe_range_partitioning { + // Range partitioning + assert_or_internal_err!( + partition_filters.len() == range_partitioning.partition_count(), + "Dynamic filter partition count {} does not match Range partition count {}", + partition_filters.len(), + range_partitioning.partition_count() + ); + let routing_range_expr = Arc::new(RangeExpr::try_new( + self.on_right.clone(), + range_partitioning, + )?) + as Arc; + let else_expr = partition_filters + .pop() + .expect("Range partitioning always has at least one partition"); + + // CASE range_partition(key) + // WHEN 0 THEN F0 + // WHEN 1 THEN F1 + // ... + // ELSE Fn + // END + let when_then_expr = partition_filters .into_iter() - .map(|partition_id| { + .enumerate() + .map(|(partition_id, then_expr)| { ( lit(ScalarValue::UInt64(Some(partition_id as u64))), - lit(false), + then_expr, ) }) - .collect::>(); - when_then_branches.extend(real_branches); + .collect(); - if when_then_branches.is_empty() { - lit(true) - } else { - Arc::new(CaseExpr::try_new( - Some(modulo_expr), - when_then_branches, - Some(lit(true)), - )?) as Arc - } - } else if real_branches.is_empty() { - lit(false) - } else if real_branches.len() == 1 - && empty_partition_ids.len() + 1 == num_partitions - { - Arc::clone(&real_branches[0].1) + Arc::new(CaseExpr::try_new( + Some(routing_range_expr), + when_then_expr, + Some(else_expr), + )?) as Arc } else { + // Hash partitioning + let routing_hash_expr = Arc::new(HashExpr::new( + self.on_right.clone(), + self.repartition_random_state.clone(), + "hash_repartition".to_string(), + )) + as Arc; + let modulo_expr = Arc::new(BinaryExpr::new( + routing_hash_expr, + Operator::Modulo, + lit(ScalarValue::UInt64(Some(num_partitions as u64))), + )) as Arc; + + let mut when_then_branches = if has_canceled_unknown { + empty_partition_ids + .into_iter() + .map(|partition_id| { + ( + lit(ScalarValue::UInt64(Some(partition_id as u64))), + lit(false), + ) + }) + .collect::>() + } else { + vec![] + }; + when_then_branches.extend(real_partition_ids.into_iter().map( + |partition_id| { + ( + lit(ScalarValue::UInt64(Some(partition_id as u64))), + Arc::clone(&partition_filters[partition_id]), + ) + }, + )); + Arc::new(CaseExpr::try_new( Some(modulo_expr), - real_branches, - Some(lit(false)), + when_then_branches, + Some(lit(has_canceled_unknown)), )?) as Arc }; - self.dynamic_filter.update(filter_expr)?; + self.dynamic_filter + .update(self.preserve_probe_nulls(filter_expr, keys_have_null)?)?; } } Ok(()) } + + /// Keeps probe rows with a NULL key when the join semantics need them. + /// + /// The build-side predicate drops probe rows whose key is NULL. A null-aware anti join + /// (`NOT IN`) needs that NULL to reach the join so three-valued logic can collapse the + /// result, and a null-equal join needs it to match a build-side NULL. OR-ing `key IS NULL` + /// keeps those rows while preserving the filter's selectivity for the rest; the join refines + /// whatever the widened filter lets through. + fn preserve_probe_nulls( + &self, + filter_expr: Arc, + build_keys_have_null: bool, + ) -> Result> { + // A null-aware anti join needs every probe NULL no matter what the build holds: one + // probe NULL makes `NOT IN` unknown for every build row. A null-equal join needs probe + // NULLs only to match an actual build-side NULL, so a NULL-free build keeps the filter + // at full selectivity. + let needs_probe_nulls = self.null_aware + || (self.null_equality == NullEquality::NullEqualsNull + && build_keys_have_null); + if !needs_probe_nulls { + return Ok(filter_expr); + } + // Only a key that can actually be NULL needs the disjunct; a NOT NULL key never widens. + // Null-aware joins are single-key; null-equal joins can be multi-key, so OR every nullable + // key. If every key is NOT NULL the filter is left untouched, at full selectivity. + let mut any_key_is_null: Option> = None; + for key in &self.on_right { + // `nullable` fails only when a key is out of sync with the probe schema. That is + // a construction bug, so surface it instead of widening around it. + if !key.nullable(&self.probe_schema)? { + continue; + } + let is_null = + Arc::new(IsNullExpr::new(Arc::clone(key))) as Arc; + any_key_is_null = Some(match any_key_is_null { + Some(acc) => Arc::new(BinaryExpr::new(acc, Operator::Or, is_null)) as _, + None => is_null, + }); + } + // Cheap null check first short-circuits before the costlier dynamic filter. + Ok(match any_key_is_null { + Some(any_key_is_null) => { + Arc::new(BinaryExpr::new(any_key_is_null, Operator::Or, filter_expr)) + } + None => filter_expr, + }) + } } impl fmt::Debug for SharedBuildAccumulator { @@ -699,33 +850,208 @@ impl fmt::Debug for SharedBuildAccumulator { } } +#[cfg(test)] +pub(super) fn make_partitioned_accumulator_for_test( + num_partitions: usize, +) -> SharedBuildAccumulator { + let probe_schema = Arc::new(Schema::new(vec![Field::new( + "probe_key", + DataType::Int32, + false, + )])); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + SharedBuildAccumulator { + inner: Mutex::new(AccumulatorState { + data: AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; num_partitions], + completed_partitions: 0, + }, + completion: CompletionState::Pending, + }), + completion_notify: Notify::new(), + dynamic_filter, + on_right: vec![], + repartition_random_state: SeededRandomState::with_seed(1), + probe_schema, + probe_range_partitioning: None, + null_equality: NullEquality::NullEqualsNothing, + null_aware: false, + } +} + +#[cfg(test)] +pub(super) fn completed_partitions_for_test(acc: &SharedBuildAccumulator) -> usize { + let guard = acc.inner.lock(); + let AccumulatedBuildData::Partitioned { + completed_partitions, + .. + } = &guard.data + else { + panic!("expected partitioned accumulator"); + }; + *completed_partitions +} + #[cfg(test)] mod tests { use super::*; - fn make_partitioned_accumulator(num_partitions: usize) -> SharedBuildAccumulator { - let probe_schema = Arc::new(Schema::new(vec![Field::new( + use arrow::array::{ArrayRef, BooleanArray, Float64Array, Int32Array}; + use arrow::compute::SortOptions; + use arrow::record_batch::RecordBatch; + use datafusion_common::SplitPoint; + use datafusion_physical_expr::{ + PhysicalSortExpr, + expressions::{Column, Literal}, + }; + + fn test_on_right() -> Vec { + vec![Arc::new(Column::new("probe_key", 0))] + } + + fn test_probe_schema() -> Arc { + Arc::new(Schema::new(vec![Field::new( "probe_key", DataType::Int32, false, - )])); - let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + )])) + } + + fn test_dynamic_filter( + on_right: &[PhysicalExprRef], + ) -> Arc { + Arc::new(DynamicFilterPhysicalExpr::new(on_right.to_vec(), lit(true))) + } + + fn make_accumulator_for_test( + data: AccumulatedBuildData, + on_right: Vec, + ) -> SharedBuildAccumulator { + let dynamic_filter = test_dynamic_filter(&on_right); SharedBuildAccumulator { inner: Mutex::new(AccumulatorState { - data: AccumulatedBuildData::Partitioned { - partitions: vec![PartitionStatus::Pending; num_partitions], - completed_partitions: 0, - }, + data, completion: CompletionState::Pending, }), completion_notify: Notify::new(), dynamic_filter, - on_right: vec![], + on_right, repartition_random_state: SeededRandomState::with_seed(1), - probe_schema, + probe_schema: test_probe_schema(), + probe_range_partitioning: None, + null_equality: NullEquality::NullEqualsNothing, + null_aware: false, } } + fn make_collect_left_accumulator_for_test() -> SharedBuildAccumulator { + make_accumulator_for_test( + AccumulatedBuildData::CollectLeft { + data: PartitionStatus::Pending, + reported_count: 0, + expected_reports: 1, + }, + test_on_right(), + ) + } + + fn make_partitioned_expr_accumulator_for_test( + num_partitions: usize, + ) -> SharedBuildAccumulator { + make_accumulator_for_test( + AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; num_partitions], + completed_partitions: 0, + }, + test_on_right(), + ) + } + + fn in_list(values: &[i32]) -> PushdownStrategy { + PushdownStrategy::InList(Arc::new(Int32Array::from(values.to_vec())) as ArrayRef) + } + + fn bounds(min: i32, max: i32) -> PartitionBounds { + PartitionBounds::new(vec![ColumnBounds::new( + ScalarValue::Int32(Some(min)), + ScalarValue::Int32(Some(max)), + )]) + } + + fn no_bounds() -> PartitionBounds { + PartitionBounds::new(vec![]) + } + + fn reported(pushdown: PushdownStrategy, bounds: PartitionBounds) -> PartitionStatus { + PartitionStatus::Reported(PartitionData { + pushdown, + bounds, + keys_have_null: false, + }) + } + + fn current_expr(acc: &SharedBuildAccumulator) -> PhysicalExprRef { + acc.dynamic_filter + .current() + .expect("dynamic filter current expression should be available") + } + + fn in_list_expr(expr: &PhysicalExprRef) -> &InListExpr { + expr.downcast_ref::() + .expect("expected InListExpr dynamic filter") + } + + fn assert_in_list_column_values( + expr: &PhysicalExprRef, + expected_column_name: &str, + expected_column_index: usize, + expected_values: &[i32], + ) { + let in_list = in_list_expr(expr); + let column = in_list + .expr() + .downcast_ref::() + .expect("expected InListExpr child column"); + assert_eq!(column.name(), expected_column_name); + assert_eq!(column.index(), expected_column_index); + + let actual_values = in_list + .list() + .iter() + .map(|expr| { + let literal = expr + .downcast_ref::() + .expect("expected InListExpr literal value"); + match literal.value() { + ScalarValue::Int32(Some(value)) => *value, + value => panic!("expected Int32 in-list value, got {value:?}"), + } + }) + .collect::>(); + assert_eq!(actual_values, expected_values); + } + + fn binary_expr(expr: &PhysicalExprRef) -> &BinaryExpr { + expr.downcast_ref::() + .expect("expected BinaryExpr dynamic filter") + } + + fn case_expr(expr: &PhysicalExprRef) -> &CaseExpr { + expr.downcast_ref::() + .expect("expected CaseExpr dynamic filter") + } + + fn assert_literal_bool(expr: &PhysicalExprRef, expected: bool) { + let literal = expr + .downcast_ref::() + .expect("expected literal bool dynamic filter"); + assert_eq!(literal.value(), &ScalarValue::Boolean(Some(expected))); + } + + fn assert_top_binary_op(expr: &PhysicalExprRef, expected: Operator) { + assert_eq!(binary_expr(expr).op(), &expected); + } + fn partitioned_state(acc: &SharedBuildAccumulator) -> (Vec, usize) { let guard = acc.inner.lock(); let AccumulatedBuildData::Partitioned { @@ -738,6 +1064,282 @@ mod tests { (partitions.clone(), *completed_partitions) } + #[test] + fn collect_left_updates_with_membership_only() { + let acc = make_collect_left_accumulator_for_test(); + + acc.build_filter(FinalizeInput::CollectLeft(reported( + in_list(&[1, 2, 3]), + no_bounds(), + ))) + .unwrap(); + + let expr = current_expr(&acc); + assert_in_list_column_values(&expr, "probe_key", 0, &[1, 2, 3]); + } + + #[test] + fn collect_left_updates_with_bounds_only() { + let acc = make_collect_left_accumulator_for_test(); + + acc.build_filter(FinalizeInput::CollectLeft(reported( + PushdownStrategy::Empty, + bounds(10, 20), + ))) + .unwrap(); + + let expr = current_expr(&acc); + assert_top_binary_op(&expr, Operator::And); + } + + #[test] + fn collect_left_empty_build_data_does_not_update_filter() { + let acc = make_collect_left_accumulator_for_test(); + let initial_generation = acc.dynamic_filter.snapshot_generation(); + + acc.build_filter(FinalizeInput::CollectLeft(reported( + PushdownStrategy::Empty, + no_bounds(), + ))) + .unwrap(); + + assert_eq!( + acc.dynamic_filter.snapshot_generation(), + initial_generation, + "empty CollectLeft input must not update with a no-op filter" + ); + let expr = current_expr(&acc); + assert_literal_bool(&expr, true); + } + + #[test] + fn partitioned_one_real_partition_with_rest_empty_skips_case() { + let acc = make_partitioned_expr_accumulator_for_test(3); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(PushdownStrategy::Empty, no_bounds()), + reported(in_list(&[2]), no_bounds()), + reported(PushdownStrategy::Empty, no_bounds()), + ])) + .unwrap(); + + let expr = current_expr(&acc); + in_list_expr(&expr); + assert!(expr.downcast_ref::().is_none()); + } + + #[test] + fn partitioned_canceled_unknown_partitions_keep_unknown_routes_permissive() { + let acc = make_partitioned_expr_accumulator_for_test(2); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + PartitionStatus::CanceledUnknown, + reported(PushdownStrategy::Empty, no_bounds()), + ])) + .unwrap(); + + let expr = current_expr(&acc); + let case = case_expr(&expr); + assert_eq!(case.when_then_expr().len(), 1); + assert_literal_bool(&case.when_then_expr()[0].1, false); + assert_literal_bool( + case.else_expr().expect("expected permissive fallback"), + true, + ); + } + + #[test] + fn partitioned_range_dynamic_filter_routes_with_range_expr() -> Result<()> { + let mut acc = make_partitioned_expr_accumulator_for_test(4); + acc.probe_range_partitioning = Some(RangePartitioning::try_new( + [PhysicalSortExpr::new( + Arc::clone(&acc.on_right[0]), + Default::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + )?); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(PushdownStrategy::Empty, no_bounds()), + PartitionStatus::CanceledUnknown, + reported(in_list(&[20, 29]), no_bounds()), + reported(in_list(&[30]), no_bounds()), + ]))?; + + let expr = current_expr(&acc); + let case = case_expr(&expr); + assert!( + case.expr() + .and_then(|expr| expr.downcast_ref::()) + .is_some(), + "Range routing must use RangeExpr" + ); + assert_eq!(case.when_then_expr().len(), 3); + + let batch = RecordBatch::try_new( + test_probe_schema(), + vec![Arc::new(Int32Array::from(vec![ + 9, 10, 19, 20, 21, 29, 30, 31, + ]))], + )?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = result + .as_any() + .downcast_ref::() + .expect("dynamic filter should evaluate to BooleanArray"); + assert_eq!( + result, + &BooleanArray::from(vec![false, true, true, true, false, true, true, false,]) + ); + + Ok(()) + } + + #[test] + fn partitioned_range_dynamic_filter_routes_compound_nullable_keys() -> Result<()> { + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("probe_key", DataType::Int32, true), + Field::new("probe_tie", DataType::Int32, true), + ])); + let on_right: Vec = vec![ + Arc::new(Column::new("probe_key", 0)), + Arc::new(Column::new("probe_tie", 1)), + ]; + let mut acc = make_accumulator_for_test( + AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; 4], + completed_partitions: 0, + }, + on_right, + ); + acc.probe_schema = Arc::clone(&probe_schema); + acc.probe_range_partitioning = Some(RangePartitioning::try_new( + [ + PhysicalSortExpr::new( + Arc::clone(&acc.on_right[0]), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::clone(&acc.on_right[1]), + SortOptions::new(false, false), + ), + ] + .into(), + vec![ + SplitPoint::new(vec![ + ScalarValue::Int32(None), + ScalarValue::Int32(Some(10)), + ]), + SplitPoint::new(vec![ScalarValue::Int32(None), ScalarValue::Int32(None)]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(10)), + ScalarValue::Int32(None), + ]), + ], + )?); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + reported(PushdownStrategy::Empty, no_bounds()), + PartitionStatus::CanceledUnknown, + reported(PushdownStrategy::Empty, no_bounds()), + PartitionStatus::CanceledUnknown, + ]))?; + + let expr = current_expr(&acc); + let case = case_expr(&expr); + assert!(case.expr().is_some()); + assert_eq!(case.when_then_expr().len(), 3); + + let batch = RecordBatch::try_new( + probe_schema, + vec![ + Arc::new(Int32Array::from(vec![ + None, + None, + None, + None, + Some(9), + Some(10), + Some(10), + Some(11), + ])), + Arc::new(Int32Array::from(vec![ + Some(9), + Some(10), + Some(11), + None, + None, + Some(9), + None, + None, + ])), + ], + )?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = result + .as_any() + .downcast_ref::() + .expect("dynamic filter should evaluate to BooleanArray"); + assert_eq!( + result, + &BooleanArray::from( + vec![false, true, true, false, false, false, true, true,] + ) + ); + + Ok(()) + } + + #[test] + fn partitioned_range_dynamic_filter_preserves_signed_zero_routing() -> Result<()> { + let probe_schema = Arc::new(Schema::new(vec![Field::new( + "probe_key", + DataType::Float64, + false, + )])); + let on_right: Vec = vec![Arc::new(Column::new("probe_key", 0))]; + let mut acc = make_accumulator_for_test( + AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; 2], + completed_partitions: 0, + }, + on_right, + ); + acc.probe_schema = Arc::clone(&probe_schema); + acc.probe_range_partitioning = Some(RangePartitioning::try_new( + [PhysicalSortExpr::new( + Arc::clone(&acc.on_right[0]), + SortOptions::default(), + )] + .into(), + vec![SplitPoint::new(vec![ScalarValue::Float64(Some(0.0))])], + )?); + + acc.build_filter(FinalizeInput::Partitioned(vec![ + PartitionStatus::CanceledUnknown, + reported(PushdownStrategy::Empty, no_bounds()), + ]))?; + + let expr = current_expr(&acc); + let batch = RecordBatch::try_new( + probe_schema, + vec![Arc::new(Float64Array::from(vec![-0.0, 0.0]))], + )?; + let result = expr.evaluate(&batch)?.into_array(batch.num_rows())?; + let result = result + .as_any() + .downcast_ref::() + .expect("dynamic filter should evaluate to BooleanArray"); + assert_eq!(result, &BooleanArray::from(vec![true, false])); + + Ok(()) + } + // Regression guard for the build-report lifecycle fix: on `Drop`, a stream // in `BuildReportState::ReportScheduled` still calls `report_canceled_partition` // because it cannot tell whether the coordinator has already observed the @@ -748,7 +1350,7 @@ mod tests { // `Reported`. This test pins that invariant. #[test] fn report_canceled_partition_is_noop_after_report() { - let acc = make_partitioned_accumulator(2); + let acc = make_partitioned_accumulator_for_test(2); { let mut guard = acc.inner.lock(); @@ -758,6 +1360,7 @@ mod tests { partition_id: 0, pushdown: PushdownStrategy::Empty, bounds: PartitionBounds::new(vec![]), + keys_have_null: false, }, ) .unwrap(); @@ -780,7 +1383,7 @@ mod tests { // which is what unblocks sibling partitions waiting on the coordinator. #[test] fn report_canceled_partition_marks_pending_partition_canceled() { - let acc = make_partitioned_accumulator(2); + let acc = make_partitioned_accumulator_for_test(2); acc.report_canceled_partition(0); let (partitions, completed) = partitioned_state(&acc); @@ -794,4 +1397,120 @@ mod tests { assert!(matches!(partitions[0], PartitionStatus::CanceledUnknown)); assert_eq!(completed, 1); } + + fn null_semantics_accumulator( + probe_schema: Arc, + on_right: Vec, + null_equality: NullEquality, + null_aware: bool, + ) -> SharedBuildAccumulator { + SharedBuildAccumulator { + inner: Mutex::new(AccumulatorState { + data: AccumulatedBuildData::Partitioned { + partitions: vec![PartitionStatus::Pending; 1], + completed_partitions: 0, + }, + completion: CompletionState::Pending, + }), + completion_notify: Notify::new(), + dynamic_filter: Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))), + on_right, + repartition_random_state: SeededRandomState::with_seed(1), + probe_schema, + probe_range_partitioning: None, + null_equality, + null_aware, + } + } + + fn null_equal_accumulator( + probe_schema: Arc, + on_right: Vec, + ) -> SharedBuildAccumulator { + null_semantics_accumulator( + probe_schema, + on_right, + NullEquality::NullEqualsNull, + false, + ) + } + + #[test] + fn preserve_probe_nulls_only_widens_nullable_keys() { + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("k_nullable", DataType::Int32, true), + Field::new("k_not_null", DataType::Int32, false), + ])); + let on_right: Vec = vec![ + Arc::new(Column::new("k_nullable", 0)), + Arc::new(Column::new("k_not_null", 1)), + ]; + let acc = null_equal_accumulator(probe_schema, on_right); + + // Only the nullable key earns an IS NULL disjunct; the NOT NULL key is left out. + let widened = acc.preserve_probe_nulls(lit(true), true).unwrap(); + assert_eq!(format!("{widened}").matches("IS NULL").count(), 1); + } + + #[test] + fn preserve_probe_nulls_leaves_all_not_null_keys_untouched() { + let probe_schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + ])); + let on_right: Vec = + vec![Arc::new(Column::new("a", 0)), Arc::new(Column::new("b", 1))]; + let acc = null_equal_accumulator(probe_schema, on_right); + + // Every key is NOT NULL, so there is nothing to OR in and the filter is returned as-is. + let filter = lit(true); + let result = acc.preserve_probe_nulls(Arc::clone(&filter), true).unwrap(); + assert_eq!(format!("{result}"), format!("{filter}")); + } + + #[test] + fn preserve_probe_nulls_rejects_out_of_sync_key() { + let probe_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + // The key's column index points past the probe schema: a construction bug that + // must surface as an error, not get widened around. + let on_right: Vec = vec![Arc::new(Column::new("b", 1))]; + let acc = null_equal_accumulator(probe_schema, on_right); + + assert!(acc.preserve_probe_nulls(lit(true), true).is_err()); + } + + #[test] + fn preserve_probe_nulls_skips_wrap_when_build_has_no_nulls() { + let probe_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let on_right: Vec = vec![Arc::new(Column::new("a", 0))]; + let acc = null_equal_accumulator(probe_schema, on_right); + + // A NULL-free build has nothing for a probe NULL to null-match, so the + // filter keeps its full selectivity. + let filter = lit(true); + let result = acc + .preserve_probe_nulls(Arc::clone(&filter), false) + .unwrap(); + assert_eq!(format!("{result}"), format!("{filter}")); + } + + #[test] + fn preserve_probe_nulls_wraps_null_aware_regardless_of_build() { + let probe_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let on_right: Vec = vec![Arc::new(Column::new("a", 0))]; + let acc = null_semantics_accumulator( + probe_schema, + on_right, + NullEquality::NullEqualsNothing, + true, + ); + + // One probe NULL collapses `NOT IN` for every build row, so the wrap must not + // depend on the build content. + let widened = acc.preserve_probe_nulls(lit(true), false).unwrap(); + assert_eq!(format!("{widened}").matches("IS NULL").count(), 1); + } } diff --git a/datafusion/physical-plan/src/joins/hash_join/stream.rs b/datafusion/physical-plan/src/joins/hash_join/stream.rs index 040470c9be12b..686939537e73e 100644 --- a/datafusion/physical-plan/src/joins/hash_join/stream.rs +++ b/datafusion/physical-plan/src/joins/hash_join/stream.rs @@ -33,7 +33,7 @@ use crate::joins::hash_join::shared_bounds::{ PartitionBounds, PartitionBuildData, SharedBuildAccumulator, }; use crate::joins::utils::{ - OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap, + OnceFut, equal_rows_arr, get_final_indices_from_shared_bitmap, matchable_join_keys, }; use crate::stream::EmptyRecordBatchStream; use crate::{ @@ -48,6 +48,7 @@ use crate::{ }; use arrow::array::{Array, ArrayRef, UInt32Array, UInt64Array}; +use arrow::buffer::NullBuffer; use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::{ @@ -156,6 +157,10 @@ pub(super) struct ProcessProbeBatchState { batch: RecordBatch, /// Probe-side on expressions values values: Vec, + /// Combined validity of the probe-side key columns, set when NULL keys + /// exist and cannot match (`NullEquality::NullEqualsNothing`); NULL rows + /// are skipped during JoinHashMap lookups + valid_keys: Option, /// Starting offset for JoinHashMap lookups offset: MapOffset, /// Max joined probe-side index from current batch @@ -173,15 +178,109 @@ impl ProcessProbeBatchState { /// Lifecycle of this partition's build-data report to the shared coordinator. /// -/// `ReportScheduled` means the reporting `OnceFut` has been constructed but is -/// lazy: the coordinator has not yet observed the report. Only `ReportDelivered` -/// guarantees the coordinator saw it, so `Drop` must still cancel the partition -/// when the state is `ReportScheduled` — otherwise sibling partitions wait -/// forever for a report that never runs. +/// `Scheduled` means the reporting `OnceFut` has been constructed but is lazy: +/// the coordinator has not necessarily observed the report. Only `Delivered` +/// guarantees the coordinator saw it, so `Drop` must still cancel a `Scheduled` +/// partition — otherwise sibling partitions can wait forever for a report that +/// never runs. +#[derive(Debug, PartialEq, Eq)] enum BuildReportState { NotReported, - ReportScheduled, - ReportDelivered, + Scheduled, + Delivered, + Canceled, + Finalized, +} + +/// Owns the stream-side lifecycle for one partition's build-data report. +struct BuildReportHandle { + partition: usize, + mode: PartitionMode, + build_accumulator: Option>, + waiter: Option>, + state: BuildReportState, +} + +impl BuildReportHandle { + fn new( + partition: usize, + mode: PartitionMode, + build_accumulator: Option>, + ) -> Self { + Self { + partition, + mode, + build_accumulator, + waiter: None, + state: BuildReportState::NotReported, + } + } + + fn has_accumulator(&self) -> bool { + self.build_accumulator.is_some() + } + + fn schedule(&mut self, build_data: PartitionBuildData) { + let Some(build_accumulator) = &self.build_accumulator else { + // Defensive no-op terminal state; current callers avoid scheduling + // unless an accumulator is present. + self.finalize(); + return; + }; + + debug_assert!(matches!(self.state, BuildReportState::NotReported)); + let acc = Arc::clone(build_accumulator); + self.waiter = Some(OnceFut::new(async move { + acc.report_build_data(build_data).await + })); + self.state = BuildReportState::Scheduled; + } + + fn poll_delivery(&mut self, cx: &mut std::task::Context<'_>) -> Poll> { + if let Some(ref mut fut) = self.waiter { + ready!(fut.get_shared(cx))?; + if !matches!(self.state, BuildReportState::Delivered) { + debug_assert!(matches!(self.state, BuildReportState::Scheduled)); + self.state = BuildReportState::Delivered; + } + } + Poll::Ready(Ok(())) + } + + fn cancel_pending(&mut self) { + if matches!( + self.state, + BuildReportState::Delivered + | BuildReportState::Canceled + | BuildReportState::Finalized + ) { + return; + } + + if self.mode == PartitionMode::Partitioned + && let Some(build_accumulator) = &self.build_accumulator + { + build_accumulator.report_canceled_partition(self.partition); + self.state = BuildReportState::Canceled; + } else { + self.finalize(); + } + } + + fn finalize(&mut self) { + self.state = BuildReportState::Finalized; + } + + #[cfg(test)] + fn state(&self) -> &BuildReportState { + &self.state + } +} + +impl Drop for BuildReportHandle { + fn drop(&mut self) { + self.cancel_pending(); + } } /// [`Stream`] for [`super::HashJoinExec`] that does the actual join. @@ -228,13 +327,8 @@ pub(super) struct HashJoinStream { build_indices_buffer: Vec, /// Specifies whether the right side has an ordering to potentially preserve right_side_ordered: bool, - /// Shared build accumulator for coordinating dynamic filter updates (collects hash maps and/or bounds, optional) - build_accumulator: Option>, - /// Optional future to signal when build information has been reported by all partitions - /// and the dynamic filter has been updated - build_waiter: Option>, - /// Tracks where this partition is in the build-data reporting lifecycle. - build_report_state: BuildReportState, + /// Owns this partition's build-data report lifecycle. + build_report: BuildReportHandle, /// Partitioning mode to use mode: PartitionMode, /// Output buffer for coalescing small batches into larger ones with optional fetch limit. @@ -305,6 +399,7 @@ pub(super) fn lookup_join_hashmap( probe_side_values: &[ArrayRef], null_equality: NullEquality, hashes_buffer: &[u64], + valid_keys: Option<&NullBuffer>, limit: usize, offset: MapOffset, probe_indices_buffer: &mut Vec, @@ -312,6 +407,7 @@ pub(super) fn lookup_join_hashmap( ) -> Result<(UInt64Array, UInt32Array, Option)> { let next_offset = build_hashmap.get_matched_indices_with_limit_offset( hashes_buffer, + valid_keys, limit, offset, probe_indices_buffer, @@ -414,9 +510,7 @@ impl HashJoinStream { probe_indices_buffer: Vec::with_capacity(batch_size), build_indices_buffer: Vec::with_capacity(batch_size), right_side_ordered, - build_accumulator, - build_waiter: None, - build_report_state: BuildReportState::NotReported, + build_report: BuildReportHandle::new(partition, mode, build_accumulator), mode, output_buffer, null_aware, @@ -429,8 +523,15 @@ impl HashJoinStream { join_type: JoinType, left_data: &JoinLeftData, ) -> HashJoinStreamState { - if left_data.map().is_empty() - && join_type.empty_build_side_produces_empty_result() + let build_empty = !left_data.has_build_rows(); + // The map can be empty even when the build side has rows: under + // `NullEqualsNothing`, build rows with a NULL join key are omitted. For + // join types whose every output row requires a build match, that still + // guarantees an empty result, so we can skip scanning the probe side. + let map_empty = !left_data.has_matchable_build_rows(); + + if (build_empty && join_type.empty_build_side_produces_empty_result()) + || (map_empty && join_type.empty_map_produces_empty_result()) { HashJoinStreamState::Completed } else { @@ -449,35 +550,39 @@ impl HashJoinStream { &mut self, left_data: &Arc, ) -> HashJoinStreamState { - let Some(build_accumulator) = self.build_accumulator.as_ref() else { + if !self.build_report.has_accumulator() { return Self::state_after_build_ready(self.join_type, left_data.as_ref()); - }; + } let pushdown = left_data.membership().clone(); let bounds = left_data .bounds .clone() .unwrap_or_else(|| PartitionBounds::new(vec![])); + // Arrow tracks null counts per array, so this costs no data scan. + let keys_have_null = left_data + .values() + .iter() + .any(|array| array.null_count() > 0); let build_data = match self.mode { PartitionMode::Partitioned => PartitionBuildData::Partitioned { partition_id: self.partition, pushdown, bounds, + keys_have_null, + }, + PartitionMode::CollectLeft => PartitionBuildData::CollectLeft { + pushdown, + bounds, + keys_have_null, }, - PartitionMode::CollectLeft => { - PartitionBuildData::CollectLeft { pushdown, bounds } - } PartitionMode::Auto => unreachable!( "PartitionMode::Auto should not be present at execution time. This is a bug in DataFusion, please report it!" ), }; - let acc = Arc::clone(build_accumulator); - self.build_waiter = Some(OnceFut::new(async move { - acc.report_build_data(build_data).await - })); - self.build_report_state = BuildReportState::ReportScheduled; + self.build_report.schedule(build_data); HashJoinStreamState::WaitPartitionBoundsReport } @@ -541,10 +646,7 @@ impl HashJoinStream { &mut self, cx: &mut std::task::Context<'_>, ) -> Poll>>> { - if let Some(ref mut fut) = self.build_waiter { - ready!(fut.get_shared(cx))?; - self.build_report_state = BuildReportState::ReportDelivered; - } + ready!(self.build_report.poll_delivery(cx))?; let build_side = self.build_side.try_as_ready()?; self.state = Self::state_after_build_ready(self.join_type, build_side.left_data.as_ref()); @@ -599,7 +701,9 @@ impl HashJoinStream { // Precalculate hash values for fetched batch let keys_values = evaluate_expressions_to_arrays(&self.on_right, &batch)?; - if let Map::HashMap(_) = self.build_side.try_as_ready()?.left_data.map() { + let valid_keys = if let Map::HashMap(_) = + self.build_side.try_as_ready()?.left_data.map() + { self.hashes_buffer.clear(); self.hashes_buffer.resize(batch.num_rows(), 0); create_hashes( @@ -607,7 +711,10 @@ impl HashJoinStream { &self.random_state, &mut self.hashes_buffer, )?; - } + matchable_join_keys(&keys_values, self.null_equality) + } else { + None + }; self.join_metrics.input_batches.add(1); self.join_metrics.input_rows.add(batch.num_rows()); @@ -616,6 +723,7 @@ impl HashJoinStream { HashJoinStreamState::ProcessProbeBatch(ProcessProbeBatchState { batch, values: keys_values, + valid_keys, offset: (0, None), joined_probe_idx: None, }); @@ -679,14 +787,9 @@ impl HashJoinStream { } } - // If the build side is empty, this stream only reaches ProcessProbeBatch for - // join types whose output still depends on probe rows. - let is_empty = build_side.left_data.map().is_empty(); + let is_empty = !build_side.left_data.has_matchable_build_rows(); if is_empty { - // Invariant: state_after_build_ready should have already completed - // join types whose result is fixed to empty when the build side is empty. - debug_assert!(!self.join_type.empty_build_side_produces_empty_result()); let result = build_batch_empty_build_side( &self.schema, build_side.left_data.batch(), @@ -710,6 +813,7 @@ impl HashJoinStream { &state.values, self.null_equality, &self.hashes_buffer, + state.valid_keys.as_ref(), self.batch_size, state.offset, &mut self.probe_indices_buffer, @@ -966,14 +1070,75 @@ impl Stream for HashJoinStream { } } -impl Drop for HashJoinStream { - fn drop(&mut self) { - if self.mode == PartitionMode::Partitioned - && !matches!(self.build_report_state, BuildReportState::ReportDelivered) - && let Some(build_accumulator) = &self.build_accumulator +#[cfg(test)] +mod tests { + use super::*; + use crate::joins::hash_join::shared_bounds::{ + PushdownStrategy, completed_partitions_for_test, + make_partitioned_accumulator_for_test, + }; + + fn empty_build_data(partition_id: usize) -> PartitionBuildData { + PartitionBuildData::Partitioned { + partition_id, + pushdown: PushdownStrategy::Empty, + bounds: PartitionBounds::new(vec![]), + keys_have_null: false, + } + } + + fn partitioned_handle(acc: &Arc) -> BuildReportHandle { + BuildReportHandle::new(0, PartitionMode::Partitioned, Some(Arc::clone(acc))) + } + + #[test] + fn build_report_handle_cancels_scheduled_partition_on_drop() { + let acc = Arc::new(make_partitioned_accumulator_for_test(2)); + { - build_accumulator.report_canceled_partition(self.partition); - self.build_report_state = BuildReportState::ReportDelivered; + let mut handle = partitioned_handle(&acc); + handle.schedule(empty_build_data(0)); + assert_eq!(handle.state(), &BuildReportState::Scheduled); + } + + assert_eq!(completed_partitions_for_test(&acc), 1); + } + + #[test] + fn build_report_handle_does_not_cancel_delivered_partition_on_drop() { + let acc = Arc::new(make_partitioned_accumulator_for_test(1)); + + { + let mut handle = partitioned_handle(&acc); + handle.schedule(empty_build_data(0)); + let mut cx = std::task::Context::from_waker(futures::task::noop_waker_ref()); + assert!(matches!(handle.poll_delivery(&mut cx), Poll::Ready(Ok(())))); + assert_eq!(handle.state(), &BuildReportState::Delivered); } + + assert_eq!(completed_partitions_for_test(&acc), 1); + } + + #[test] + fn build_report_handle_cancel_pending_is_idempotent() { + let acc = Arc::new(make_partitioned_accumulator_for_test(2)); + let mut handle = partitioned_handle(&acc); + handle.schedule(empty_build_data(0)); + + handle.cancel_pending(); + handle.cancel_pending(); + + assert_eq!(handle.state(), &BuildReportState::Canceled); + assert_eq!(completed_partitions_for_test(&acc), 1); + } + + #[test] + fn build_report_handle_no_accumulator_finalizes() { + let mut handle = BuildReportHandle::new(0, PartitionMode::Partitioned, None); + + handle.schedule(empty_build_data(0)); + handle.cancel_pending(); + + assert_eq!(handle.state(), &BuildReportState::Finalized); } } diff --git a/datafusion/physical-plan/src/joins/join_hash_map.rs b/datafusion/physical-plan/src/joins/join_hash_map.rs index 8f0fb66b64fbf..454cc916aeb12 100644 --- a/datafusion/physical-plan/src/joins/join_hash_map.rs +++ b/datafusion/physical-plan/src/joins/join_hash_map.rs @@ -23,7 +23,7 @@ use std::fmt::{self, Debug}; use std::ops::Sub; use arrow::array::BooleanArray; -use arrow::buffer::BooleanBuffer; +use arrow::buffer::{BooleanBuffer, NullBuffer}; use arrow::datatypes::ArrowNativeType; use hashbrown::HashTable; use hashbrown::hash_table::Entry::{Occupied, Vacant}; @@ -117,9 +117,14 @@ pub trait JoinHashMapType: Send + Sync { deleted_offset: Option, ) -> (Vec, Vec); + /// Probe rows marked NULL in `valid_keys` are skipped without a lookup: + /// their key contains a NULL, which cannot match any build row under + /// `NullEquality::NullEqualsNothing`. Pass `None` when every probe key is + /// matchable. fn get_matched_indices_with_limit_offset( &self, hash_values: &[u64], + valid_keys: Option<&NullBuffer>, limit: usize, offset: MapOffset, input_indices: &mut Vec, @@ -185,6 +190,7 @@ impl JoinHashMapType for JoinHashMapU32 { fn get_matched_indices_with_limit_offset( &self, hash_values: &[u64], + valid_keys: Option<&NullBuffer>, limit: usize, offset: MapOffset, input_indices: &mut Vec, @@ -194,6 +200,7 @@ impl JoinHashMapType for JoinHashMapU32 { &self.map, &self.next, hash_values, + valid_keys, limit, offset, input_indices, @@ -263,6 +270,7 @@ impl JoinHashMapType for JoinHashMapU64 { fn get_matched_indices_with_limit_offset( &self, hash_values: &[u64], + valid_keys: Option<&NullBuffer>, limit: usize, offset: MapOffset, input_indices: &mut Vec, @@ -272,6 +280,7 @@ impl JoinHashMapType for JoinHashMapU64 { &self.map, &self.next, hash_values, + valid_keys, limit, offset, input_indices, @@ -376,10 +385,12 @@ where (input_indices, match_indices) } +#[expect(clippy::too_many_arguments)] pub fn get_matched_indices_with_limit_offset( map: &HashTable<(u64, T)>, next_chain: &[T], hash_values: &[u64], + valid_keys: Option<&NullBuffer>, limit: usize, offset: MapOffset, input_indices: &mut Vec, @@ -401,6 +412,10 @@ where let start = offset.0; let end = (start + limit).min(hash_values.len()); for (i, &hash) in hash_values[start..end].iter().enumerate() { + // NULL keys cannot match any build row + if valid_keys.is_some_and(|valid| valid.is_null(start + i)) { + continue; + } if let Some((_, idx)) = map.find(hash, |(h, _)| hash == *h) { input_indices.push(start as u32 + i as u32); match_indices.push((*idx - one).into()); @@ -445,6 +460,10 @@ where let hash_values_len = hash_values.len(); for (i, &hash) in hash_values[to_skip..].iter().enumerate() { let row_idx = to_skip + i; + // NULL keys cannot match any build row + if valid_keys.is_some_and(|valid| valid.is_null(row_idx)) { + continue; + } if let Some((_, idx)) = map.find(hash, |(h, _)| hash == *h) { let idx: T = *idx; let is_last = row_idx == hash_values_len - 1; @@ -494,4 +513,60 @@ mod tests { } } } + + #[test] + fn test_get_matched_indices_skips_invalid_keys() { + let mut hash_map = JoinHashMapU32::with_capacity(3); + hash_map.update_from_iter(Box::new([10u64, 20u64, 30u64].iter().enumerate()), 0); + + let probe_hashes = vec![10, 20, 30]; + // The probe row for hash 20 has a NULL key and must not match. + let valid_keys = NullBuffer::from(vec![true, false, true]); + + let mut input_indices = vec![]; + let mut match_indices = vec![]; + let next_offset = hash_map.get_matched_indices_with_limit_offset( + &probe_hashes, + Some(&valid_keys), + 8192, + (0, None), + &mut input_indices, + &mut match_indices, + ); + + assert_eq!(next_offset, None); + assert_eq!(input_indices, vec![0, 2]); + assert_eq!(match_indices, vec![0, 2]); + } + + #[test] + fn test_get_matched_indices_skips_invalid_keys_with_duplicates() { + // Duplicate build keys chain multiple rows under one hash value. + let mut hash_map = JoinHashMapU32::with_capacity(4); + hash_map.update_from_iter( + Box::new([10u64, 20u64, 10u64, 20u64].iter().enumerate()), + 0, + ); + + let probe_hashes = vec![10, 20]; + // The probe row for hash 10 has a NULL key: none of the build rows in + // its chain may match, while the valid probe row for hash 20 must + // still match its entire chain. + let valid_keys = NullBuffer::from(vec![false, true]); + + let mut input_indices = vec![]; + let mut match_indices = vec![]; + let next_offset = hash_map.get_matched_indices_with_limit_offset( + &probe_hashes, + Some(&valid_keys), + 8192, + (0, None), + &mut input_indices, + &mut match_indices, + ); + + assert_eq!(next_offset, None); + assert_eq!(input_indices, vec![1, 1]); + assert_eq!(match_indices, vec![3, 1]); + } } diff --git a/datafusion/physical-plan/src/joins/mod.rs b/datafusion/physical-plan/src/joins/mod.rs index 2cdfa1e6ac020..e4f7e2e123e0e 100644 --- a/datafusion/physical-plan/src/joins/mod.rs +++ b/datafusion/physical-plan/src/joins/mod.rs @@ -34,6 +34,8 @@ mod cross_join; mod hash_join; mod nested_loop_join; mod piecewise_merge_join; +#[cfg(feature = "proto")] +mod proto; mod sort_merge_join; mod stream_join_utils; mod symmetric_hash_join; @@ -50,6 +52,15 @@ pub mod join_hash_map; use array_map::ArrayMap; use utils::JoinHashMapType; +/// The build-side map of a hash join, indexing build rows by join key. +/// +/// Under [`NullEquality::NullEqualsNothing`], build rows with a NULL in any +/// join key column can never match a probe row and are omitted from the map. +/// [`Map::is_empty`] and [`Map::num_of_distinct_key`] therefore reflect the +/// *matchable* build rows: the map can be empty even when the build side +/// contains rows. +/// +/// [`NullEquality::NullEqualsNothing`]: datafusion_common::NullEquality::NullEqualsNothing pub enum Map { HashMap(Box), ArrayMap(ArrayMap), diff --git a/datafusion/physical-plan/src/joins/nested_loop_join.rs b/datafusion/physical-plan/src/joins/nested_loop_join.rs index feaf344200ac1..eb1df638c7dc5 100644 --- a/datafusion/physical-plan/src/joins/nested_loop_join.rs +++ b/datafusion/physical-plan/src/joins/nested_loop_join.rs @@ -40,12 +40,13 @@ use crate::metrics::{ }; use crate::projection::{ EmbeddedProjection, JoinData, ProjectionExec, try_embed_projection, - try_pushdown_through_join, + try_pushdown_through_join_with_column_indices, }; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, + SendableRecordBatchStream, validate_child_count, }; use arrow::array::{ @@ -62,12 +63,12 @@ use arrow_schema::DataType; use datafusion_common::cast::as_boolean_array; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ - JoinSide, Result, ScalarValue, Statistics, arrow_err, assert_eq_or_internal_err, - internal_datafusion_err, internal_err, project_schema, unwrap_or_internal_err, + JoinSide, NullEquality, Result, ScalarValue, Statistics, arrow_err, + assert_eq_or_internal_err, internal_datafusion_err, internal_err, project_schema, + unwrap_or_internal_err, }; -use datafusion_execution::TaskContext; -use datafusion_execution::disk_manager::RefCountedTempFile; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use datafusion_execution::{SpillFile, TaskContext}; use datafusion_expr::JoinType; use datafusion_physical_expr::equivalence::{ ProjectionMapping, join_equivalence_properties, @@ -489,28 +490,6 @@ impl NestedLoopJoinExec { Ok(plan) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - - Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - build_side_data: Default::default(), - left_spill_data: Arc::new(OnceAsync::default()), - cache: Arc::clone(&self.cache), - filter: self.filter.clone(), - join_type: self.join_type, - join_schema: Arc::clone(&self.join_schema), - column_indices: self.column_indices.clone(), - projection: self.projection.clone(), - } - } } impl DisplayAs for NestedLoopJoinExec { @@ -566,10 +545,14 @@ impl ExecutionPlan for NestedLoopJoinExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn maintains_input_order(&self) -> Vec { @@ -582,30 +565,70 @@ impl ExecutionPlan for NestedLoopJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn crate::PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // Apply to join filter expressions if present - if let Some(filter) = &self.filter { - f(filter.expression().as_ref())?; + crate::apply_expression_roots( + self.filter.iter().map(|filter| filter.expression()), + f, + ) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + build_side_data: Default::default(), + left_spill_data: Arc::new(OnceAsync::default()), + cache: Arc::clone(&self.cache), + filter: self.filter.clone(), + join_type: self.join_type, + join_schema: Arc::clone(&self.join_schema), + column_indices: self.column_indices.clone(), + projection: self.projection.clone(), + })) + } + ChildrenPropertiesMode::Recompute => Ok(Arc::new( + NestedLoopJoinExecBuilder::new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.join_type, + ) + .with_filter(self.filter.clone()) + .with_projection_ref(self.projection.clone()) + .build()?, + )), } - Ok(TreeNodeRecursion::Continue) } fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - NestedLoopJoinExecBuilder::new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - self.join_type, - ) - .with_filter(self.filter.clone()) - .with_projection_ref(self.projection.clone()) - .build()?, - )) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -701,7 +724,17 @@ impl ExecutionPlan for NestedLoopJoinExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { + fn child_stats_requests(&self, partition: Option) -> Vec { + // Left side is always broadcast, so it always needs overall stats. + // Right side is partitioned, so it needs per-partition stats. + vec![ChildStats::At(None), ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { // NestedLoopJoinExec is designed for joins without equijoin keys in the // ON clause (e.g., `t1 JOIN t2 ON (t1.v1 + t2.v1) % 2 = 0`). Any join // predicates are stored in `self.filter`, but `estimate_join_statistics` @@ -711,20 +744,14 @@ impl ExecutionPlan for NestedLoopJoinExec { // unknown row counts. let join_columns = Vec::new(); - // Left side is always a single partition (Distribution::SinglePartition), - // so we always request overall stats with `None`. Right side can have - // multiple partitions, so we forward the partition parameter to get - // partition-specific statistics when requested. - let left_stats = Arc::unwrap_or_clone(self.left.partition_statistics(None)?); - let right_stats = Arc::unwrap_or_clone(match partition { - Some(partition) => self.right.partition_statistics(Some(partition))?, - None => self.right.partition_statistics(None)?, - }); + let left_stats = input_stats[0].as_ref().clone(); + let right_stats = input_stats[1].as_ref().clone(); let stats = estimate_join_statistics( left_stats, right_stats, &join_columns, + NullEquality::NullEqualsNothing, &self.join_type, &self.join_schema, )?; @@ -750,13 +777,14 @@ impl ExecutionPlan for NestedLoopJoinExec { projected_right_child, join_filter, .. - }) = try_pushdown_through_join( + }) = try_pushdown_through_join_with_column_indices( projection, self.left(), self.right(), &[], &schema, self.filter(), + self.column_indices.as_slice(), )? { Ok(Some(Arc::new(NestedLoopJoinExec::try_new( Arc::new(projected_left_child), @@ -770,6 +798,91 @@ impl ExecutionPlan for NestedLoopJoinExec { try_embed_projection(projection, self) } } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + + let join_type = crate::joins::proto::join_type_to_proto(*self.join_type()); + + let filter = self + .filter() + .map(|f| crate::joins::proto::join_filter_to_proto(f, ctx)) + .transpose()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::NestedLoopJoin(Box::new( + protobuf::NestedLoopJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + join_type: join_type.into(), + filter, + projection: match self.projection.as_ref() { + None => Vec::new(), + Some(v) if v.is_empty() => vec![u32::MAX], + Some(v) => v.iter().map(|x| *x as u32).collect(), + }, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl NestedLoopJoinExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::NestedLoopJoin, + "NestedLoopJoinExec", + ); + + let left = ctx.decode_required_child( + join.left.as_deref(), + "NestedLoopJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + join.right.as_deref(), + "NestedLoopJoinExec", + "right", + )?; + + let join_type = crate::joins::proto::join_type_from_proto( + join.join_type, + "NestedLoopJoinExec", + )?; + + let filter = join + .filter + .as_ref() + .map(|f| { + crate::joins::proto::join_filter_from_proto(f, ctx, "NestedLoopJoinExec") + }) + .transpose()?; + + let projection = match join.projection.as_slice() { + [] => None, + [u32::MAX] => Some(Vec::new()), + indices => Some(indices.iter().map(|i| *i as usize).collect()), + }; + + Ok(Arc::new(NestedLoopJoinExec::try_new( + left, right, filter, &join_type, projection, + )?)) + } } impl EmbeddedProjection for NestedLoopJoinExec { @@ -884,6 +997,16 @@ enum NLJState { FetchingRight, ProbeRight, EmitRightUnmatched, + /// Entered exactly once per left chunk, when the probe (right) side is + /// exhausted and probing for the current chunk is finished. This state + /// owns the single [`JoinLeftData::report_probe_completed`] call that + /// decrements the shared probe-threads counter, and records in + /// `is_unmatched_left_emitter` whether this stream is the one responsible + /// for emitting unmatched-left rows. Splitting this decision out of + /// `EmitLeftUnmatched` makes "decrement exactly once" a structural + /// property of the state graph, so the (re-enterable) emit state no longer + /// has to guard against decrementing twice. + ProbeEnd, EmitLeftUnmatched, /// Emit unmatched right rows using the global bitmap accumulated across /// all left chunks. Only used in memory-limited mode for join types that @@ -901,7 +1024,7 @@ pub(crate) struct LeftSpillData { /// SpillManager used to read the spill file (has the left schema) spill_manager: SpillManager, /// The spill file containing all left-side batches - spill_file: RefCountedTempFile, + spill_file: Arc, /// Left-side schema schema: SchemaRef, } @@ -1074,6 +1197,18 @@ pub(crate) struct NestedLoopJoinStream { /// Memory-limited spill fallback state. See [`SpillState`] for details. spill_state: SpillState, + + /// Whether this stream is the one responsible for emitting unmatched-left + /// rows for the current left chunk. Set in the [`NLJState::ProbeEnd`] state, + /// which is entered exactly once per chunk and owns the single + /// [`JoinLeftData::report_probe_completed`] call: the stream that drives the + /// shared probe-threads counter to zero (the last to finish probing) becomes + /// the emitter. Because the decrement happens once in `ProbeEnd` rather than + /// in the re-enterable `EmitLeftUnmatched` state, the counter can never be + /// decremented twice, so it cannot reach zero before all partitions finish + /// probing (which would otherwise let a partition emit spurious NULL-padded + /// unmatched-left rows early). + is_unmatched_left_emitter: bool, } pub(crate) struct NestedLoopJoinMetrics { @@ -1117,7 +1252,7 @@ impl Stream for NestedLoopJoinStream { /// BufferingLeft → FetchingRight /// /// FetchingRight → ProbeRight (if right batch available) - /// FetchingRight → EmitLeftUnmatched (if right exhausted) + /// FetchingRight → ProbeEnd (if right exhausted) /// /// ProbeRight → ProbeRight (next left row or after yielding output) /// ProbeRight → EmitRightUnmatched (for special join types like right join) @@ -1125,6 +1260,9 @@ impl Stream for NestedLoopJoinStream { /// /// EmitRightUnmatched → FetchingRight /// + /// ProbeEnd → EmitLeftUnmatched (records whether this stream is the + /// unmatched-left emitter, then always continues to EmitLeftUnmatched) + /// /// EmitLeftUnmatched → EmitLeftUnmatched (only process 1 chunk for each /// iteration) /// EmitLeftUnmatched → Done (if finished) @@ -1160,8 +1298,8 @@ impl Stream for NestedLoopJoinStream { // 1. --> ProbeRight // Start processing the join for the newly fetched right // batch. - // 2. --> EmitLeftUnmatched: When the right side input is exhausted, (maybe) emit - // unmatched left side rows. + // 2. --> ProbeEnd: When the right side input is exhausted, + // probing for the current left chunk is finished. // // After fetching a new batch from the right side, it will // process all rows from the buffered left data: @@ -1175,9 +1313,10 @@ impl Stream for NestedLoopJoinStream { // at once in memory. // // So after the right side input is exhausted, the join phase - // for the current buffered left data is finished. We can go to - // the next `EmitLeftUnmatched` phase to check if there is any - // special handling (e.g., in cases like left join). + // for the current buffered left data is finished. We go to the + // `ProbeEnd` state, which records probe completion before the + // `EmitLeftUnmatched` phase checks if there is any special + // handling (e.g., in cases like left join). NLJState::FetchingRight => { debug!("[NLJState] Entering: {:?}", self.state); // stop on drop @@ -1240,6 +1379,28 @@ impl Stream for NestedLoopJoinStream { } } + // NLJState transitions: + // 1. --> EmitLeftUnmatched + // Probing for the current left chunk is finished. Report + // probe completion exactly once (decrementing the shared + // probe-threads counter) and record whether this stream is + // the unmatched-left emitter, then always advance to + // `EmitLeftUnmatched`. + NLJState::ProbeEnd => { + debug!("[NLJState] Entering: {:?}", self.state); + + // stop on drop + let join_metric = self.metrics.join_metrics.join_time.clone(); + let _join_timer = join_metric.timer(); + + match self.handle_probe_end() { + ControlFlow::Continue(()) => continue, + ControlFlow::Break(poll) => { + return self.metrics.join_metrics.baseline.record_poll(poll); + } + } + } + // NLJState transitions: // 1. --> EmitLeftUnmatched(1) // If we have already buffered enough output to yield, it @@ -1347,6 +1508,7 @@ impl NestedLoopJoinStream { handled_empty_output: false, should_track_unmatched_right: need_produce_right_in_final(join_type), spill_state, + is_unmatched_left_emitter: false, } } @@ -1540,7 +1702,7 @@ impl NestedLoopJoinStream { Poll::Ready(Ok(spill_data)) => { match spill_data .spill_manager - .read_spill_as_stream(spill_data.spill_file.clone(), None) + .read_spill_as_stream(Arc::clone(&spill_data.spill_file), None) { Ok(stream) => { active.left_schema = Some(Arc::clone(&spill_data.schema)); @@ -1722,7 +1884,10 @@ impl NestedLoopJoinStream { } Some(Err(e)) => ControlFlow::Break(Poll::Ready(Some(Err(e)))), None => { - self.state = NLJState::EmitLeftUnmatched; + // Right side exhausted: probing for the current left chunk + // is finished. `ProbeEnd` reports probe completion before + // emitting unmatched-left rows. + self.state = NLJState::ProbeEnd; ControlFlow::Continue(()) } }, @@ -1835,6 +2000,34 @@ impl NestedLoopJoinStream { } } + /// Handle ProbeEnd state - record probe completion for the current chunk. + /// + /// Entered exactly once per left chunk, when the right side is exhausted. + /// This is the single place that decrements the shared probe-threads counter + /// via [`JoinLeftData::report_probe_completed`]: the stream that drives the + /// counter to zero (the last to finish probing) is the one responsible for + /// emitting unmatched-left rows, recorded in `is_unmatched_left_emitter`. + /// + /// Owning the decrement here — rather than in the re-enterable + /// `EmitLeftUnmatched` state — makes "decrement exactly once per stream" a + /// structural property of the state graph, so the counter cannot reach zero + /// before all partitions finish probing (which would let a partition emit + /// spurious NULL-padded unmatched-left rows early). + /// + /// Always transitions to `EmitLeftUnmatched`. + fn handle_probe_end(&mut self) -> ControlFlow>>> { + // Decrement the shared counter exactly once for this stream/chunk. The + // last stream to finish probing (the one that drives the counter to + // zero) becomes the unmatched-left emitter. + let is_emitter = match self.get_left_data() { + Ok(left_data) => left_data.report_probe_completed(), + Err(e) => return ControlFlow::Break(Poll::Ready(Some(Err(e)))), + }; + self.is_unmatched_left_emitter = is_emitter; + self.state = NLJState::EmitLeftUnmatched; + ControlFlow::Continue(()) + } + /// Handle EmitLeftUnmatched state - emit unmatched left rows. /// /// In memory-limited mode, after processing all unmatched rows for the @@ -1873,6 +2066,10 @@ impl NestedLoopJoinStream { self.buffered_left_data = None; self.left_probe_idx = 0; self.left_emit_idx = 0; + // Each memory-limited chunk gets a fresh per-chunk + // `JoinLeftData`/counter; `is_unmatched_left_emitter` is + // recomputed when `ProbeEnd` is re-entered for the next + // chunk, so it does not need to be reset here. self.state = NLJState::BufferingLeft; } else if self.is_memory_limited() && self.should_track_unmatched_right @@ -2360,13 +2557,14 @@ impl NestedLoopJoinStream { // Early return if join type can't have unmatched rows let join_type_no_produce_left = !need_produce_result_in_final(self.join_type); - // Early return if another thread is already processing unmatched rows - let handled_by_other_partition = - self.left_emit_idx == 0 && !left_data.report_probe_completed(); // Stop processing unmatched rows, the caller will go to the next state let finished = self.left_emit_idx >= left_batch.num_rows(); - if join_type_no_produce_left || handled_by_other_partition || finished { + // `ProbeEnd` already recorded whether this stream emits unmatched-left + // rows. Every probe partition passes through this state, but only the + // one that finished probing last is the emitter, so this flag is false + // for the others. + if join_type_no_produce_left || !self.is_unmatched_left_emitter || finished { return Ok(false); } @@ -2986,6 +3184,7 @@ fn build_unmatched_batch( #[cfg(test)] pub(crate) mod tests { use super::*; + use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test::{TestMemoryExec, assert_join_metrics}; use crate::{ common, expressions::Column, repartition::RepartitionExec, test::build_table_i32, @@ -3365,7 +3564,8 @@ pub(crate) mod tests { &JoinType::Left, Some(vec![1, 2]), )?; - let stats = nested_loop_join.partition_statistics(None)?; + let stats = StatisticsContext::new() + .compute(&nested_loop_join, &StatisticsArgs::new())?; assert_eq!( nested_loop_join.schema().fields().len(), stats.column_statistics.len(), diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs index 36a043cc7d16b..50ef78f18bf65 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/classic_join.rs @@ -125,7 +125,7 @@ impl RecordBatchStream for ClassicPWMJStream { // Classic Joins // 1. `WaitBufferedSide` - Load in the buffered side data into memory. // 2. `FetchStreamBatch` - Fetch + sort incoming stream batches. We switch the state to -// `Completed` if there are are still remaining partitions to process. It is only switched to +// `Completed` if there are still remaining partitions to process. It is only switched to // `ExhaustedStreamBatch` if all partitions have been processed. // 3. `ProcessStreamBatch` - Compare stream batch row values against the buffered side data. // 4. `ExhaustedStreamBatch` - If the join type is Left or Inner we will return state as diff --git a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs index 2b20089f8e221..c42ec67ef80d5 100644 --- a/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/piecewise_merge_join/exec.rs @@ -53,7 +53,8 @@ use crate::joins::piecewise_merge_join::utils::{ use crate::joins::utils::asymmetric_join_output_partitioning; use crate::metrics::MetricsSet; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlanProperties, check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlanProperties, + ReplaceChildrenOptions, validate_child_count, }; use crate::{ ExecutionPlan, PlanProperties, @@ -468,31 +469,6 @@ impl PiecewiseMergeJoinExec { pub fn swap_inputs(&self) -> Result> { todo!() } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let buffered = children.swap_remove(0); - let streamed = children.swap_remove(0); - Self { - buffered, - streamed, - on: self.on.clone(), - operator: self.operator, - join_type: self.join_type, - schema: Arc::clone(&self.schema), - left_child_plan_required_order: self.left_child_plan_required_order.clone(), - right_batch_required_orders: self.right_batch_required_orders.clone(), - sort_options: self.sort_options, - cache: Arc::clone(&self.cache), - num_partitions: self.num_partitions, - - // Re-set state. - metrics: ExecutionPlanMetricsSet::new(), - buffered_fut: Default::default(), - } - } } impl ExecutionPlan for PiecewiseMergeJoinExec { @@ -510,17 +486,21 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { // Apply to the two expressions being compared in the range predicate - f(self.on.0.as_ref())?.visit_sibling(|| f(self.on.1.as_ref())) + crate::apply_expression_roots([&self.on.0, &self.on.1], f) } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ Distribution::SinglePartition, Distribution::UnspecifiedDistribution, - ] + ]) } fn required_input_ordering(&self) -> Vec> { @@ -538,32 +518,80 @@ impl ExecutionPlan for PiecewiseMergeJoinExec { } } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let buffered = children.swap_remove(0); + let streamed = children.swap_remove(0); + Ok(Arc::new(Self { + buffered, + streamed, + on: self.on.clone(), + operator: self.operator, + join_type: self.join_type, + schema: Arc::clone(&self.schema), + left_child_plan_required_order: self + .left_child_plan_required_order + .clone(), + right_batch_required_orders: self.right_batch_required_orders.clone(), + sort_options: self.sort_options, + cache: Arc::clone(&self.cache), + num_partitions: self.num_partitions, + + // Re-set state. + metrics: ExecutionPlanMetricsSet::new(), + buffered_fut: Default::default(), + })) + } + ChildrenPropertiesMode::Recompute => match &children[..] { + [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.operator, + self.join_type, + self.num_partitions, + )?)), + _ => internal_err!( + "PiecewiseMergeJoin should have 2 children, found {}", + children.len() + ), + }, + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - match &children[..] { - [left, right] => Ok(Arc::new(PiecewiseMergeJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - self.on.clone(), - self.operator, - self.join_type, - self.num_partitions, - )?)), - _ => internal_err!( - "PiecewiseMergeJoin should have 2 children, found {}", - children.len() - ), - } + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn reset_state(self: Arc) -> Result> { - Ok(Arc::new(self.with_new_children_and_same_properties(vec![ - Arc::clone(&self.buffered), - Arc::clone(&self.streamed), - ]))) + let buffered = Arc::clone(&self.buffered); + let streamed = Arc::clone(&self.streamed); + self.replace_children( + vec![buffered, streamed], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( diff --git a/datafusion/physical-plan/src/joins/proto.rs b/datafusion/physical-plan/src/joins/proto.rs new file mode 100644 index 0000000000000..2272828b690b2 --- /dev/null +++ b/datafusion/physical-plan/src/joins/proto.rs @@ -0,0 +1,161 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Protobuf conversions shared by the join operators' `try_to_proto` / +//! `try_from_proto` implementations. +//! +//! The enum conversions are by-name exhaustive matches on purpose: the proto +//! enums and the `datafusion_common` enums are numbered differently, so a +//! numeric cast would silently corrupt them. + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{ + JoinSide, JoinType, NullEquality, Result, internal_datafusion_err, +}; +use datafusion_proto_models::protobuf; + +use crate::joins::utils::{ColumnIndex, JoinFilter}; +use crate::proto::{ExecutionPlanDecodeCtx, ExecutionPlanEncodeCtx}; + +pub(crate) fn join_type_to_proto(join_type: JoinType) -> protobuf::JoinType { + match join_type { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + } +} + +pub(crate) fn join_type_from_proto(value: i32, plan_name: &str) -> Result { + let join_type = protobuf::JoinType::try_from(value) + .map_err(|_| internal_datafusion_err!("{plan_name}: unknown JoinType {value}"))?; + Ok(match join_type { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }) +} + +pub(crate) fn join_side_to_proto(side: JoinSide) -> protobuf::JoinSide { + match side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + } +} + +pub(crate) fn join_side_from_proto(value: i32, plan_name: &str) -> Result { + let side = protobuf::JoinSide::try_from(value) + .map_err(|_| internal_datafusion_err!("{plan_name}: unknown JoinSide {value}"))?; + Ok(match side { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }) +} + +pub(crate) fn null_equality_to_proto( + null_equality: NullEquality, +) -> protobuf::NullEquality { + match null_equality { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + } +} + +pub(crate) fn null_equality_from_proto( + value: i32, + plan_name: &str, +) -> Result { + let null_equality = protobuf::NullEquality::try_from(value).map_err(|_| { + internal_datafusion_err!("{plan_name}: unknown NullEquality {value}") + })?; + Ok(match null_equality { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }) +} + +pub(crate) fn join_filter_to_proto( + filter: &JoinFilter, + ctx: &ExecutionPlanEncodeCtx<'_>, +) -> Result { + let expression = ctx.encode_expr(filter.expression())?; + let column_indices = filter + .column_indices() + .iter() + .map(|column_index| protobuf::ColumnIndex { + index: column_index.index as u32, + side: join_side_to_proto(column_index.side).into(), + }) + .collect(); + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(filter.schema().as_ref().try_into()?), + }) +} + +pub(crate) fn join_filter_from_proto( + filter: &protobuf::JoinFilter, + ctx: &ExecutionPlanDecodeCtx<'_>, + plan_name: &str, +) -> Result { + let schema: Schema = filter + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!("{plan_name}: JoinFilter missing schema") + })? + .try_into()?; + let expression = ctx.decode_required_expr( + filter.expression.as_ref(), + &schema, + plan_name, + "filter.expression", + )?; + let column_indices = filter + .column_indices + .iter() + .map(|column_index| { + Ok(ColumnIndex { + index: column_index.index as usize, + side: join_side_from_proto(column_index.side, plan_name)?, + }) + }) + .collect::>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) +} diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs index ad7312426bd18..1b90f24b96acc 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/bitwise_stream.rs @@ -84,8 +84,9 @@ //! //! Key groups can span batch boundaries on either side. The stream handles //! this by detecting when a group extends to the end of a batch, loading the -//! next batch, and continuing if the key matches. The [`PendingBoundary`] enum -//! preserves loop context across async `Poll::Pending` re-entries. +//! next batch, and continuing if the key matches. The generator-based stream +//! suspends in place at `await` points, so no explicit re-entry state is +//! needed. //! //! # Memory //! @@ -119,33 +120,32 @@ //! factor than the pair-materialization approach. use std::cmp::Ordering; -use std::fs::File; -use std::io::BufReader; -use std::pin::Pin; use std::sync::Arc; -use std::task::{Context, Poll}; +use crate::EmptyRecordBatchStream; use crate::joins::utils::{JoinFilter, JoinKeyComparator, compare_join_arrays}; use crate::metrics::{ - BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, + BaselineMetrics, Count, ExecutionPlanMetricsSet, Gauge, MetricBuilder, Time, }; +use crate::spill::in_progress_spill_file::InProgressSpillFile; use crate::spill::spill_manager::SpillManager; -use crate::{EmptyRecordBatchStream, RecordBatchStream}; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::array::{Array, ArrayRef, BooleanArray, BooleanBufferBuilder, RecordBatch}; use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch, not}; use arrow::datatypes::SchemaRef; -use arrow::ipc::reader::StreamReader; use arrow::util::bit_chunk_iterator::UnalignedBitChunk; use arrow::util::bit_util::apply_bitwise_binary_op; +use datafusion_common::instant::Instant; use datafusion_common::{ - JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err, + DataFusionError, JoinSide, JoinType, NullEquality, Result, ScalarValue, internal_err, }; -use datafusion_execution::SendableRecordBatchStream; -use datafusion_execution::disk_manager::RefCountedTempFile; use datafusion_execution::memory_pool::MemoryReservation; +use datafusion_execution::{ + SendableRecordBatchStream, SpillFile, TryEmitter, async_try_stream, +}; use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; -use futures::{Stream, StreamExt, ready}; +use futures::StreamExt; /// Evaluates join key expressions against a batch, returning one array per key. fn evaluate_join_keys( @@ -198,26 +198,6 @@ fn find_key_group_end(cmp: &JoinKeyComparator, from: usize, len: usize) -> usize lo } -/// When an outer key group spans a batch boundary, the boundary loop emits -/// the current batch, then polls for the next. If that poll returns Pending, -/// `ready!` exits `poll_join` and we re-enter from the top on the next call. -/// Without this state, the new batch would be processed fresh by the -/// merge-scan — but inner already advanced past this key, so the matching -/// outer rows would be skipped via `Ordering::Less` and never marked. -/// -/// This enum carries the last key (as single-row sliced arrays) from the -/// previous batch so we can check whether the next batch continues the same -/// key group. Stored as `Option`: `None` means normal -/// processing. -#[derive(Debug)] -enum PendingBoundary { - /// Resuming a no-filter boundary loop. - NoFilter { saved_keys: Vec }, - /// Resuming a filtered boundary loop. Inner key data remains in the - /// buffer (or spill file) for the resumed loop. - Filtered { saved_keys: Vec }, -} - /// Sort-Merge join stream for Semi/Anti/Mark joins. /// /// Named "bitwise" because it tracks outer-row matches via a per-batch @@ -255,18 +235,9 @@ pub(crate) struct BitwiseSortMergeJoinStream { // Inner key group buffer: all inner rows sharing the current join key. // Only populated when a filter is present. Unbounded — a single key // with many inner rows will buffer them all. See "Degenerate cases" - // in exec.rs. Spilled to disk when memory reservation fails. + // in exec.rs. On memory pool overflow the buffered slices move to a + // per-group spill file (see [`Self::buffer_inner_key_group`]). inner_key_buffer: Vec, - inner_key_spill: Option, - - // True when buffer_inner_key_group returned Pending after partially - // filling inner_key_buffer. On re-entry, buffer_inner_key_group - // must skip clear() and resume from poll_next_inner_batch (the - // current inner_batch was already sliced and pushed before Pending). - buffering_inner_pending: bool, - - // Boundary re-entry state — see PendingBoundary doc comment. - pending_boundary: Option, // Join ON expressions, evaluated against each new batch to produce // the key arrays used for sorted key comparisons. @@ -283,12 +254,18 @@ pub(crate) struct BitwiseSortMergeJoinStream { coalescer: BatchCoalescer, schema: SchemaRef, - // Metrics - join_time: crate::metrics::Time, + // Metrics — output rows/batches and end time are recorded by the + // ObservedStream wrapper in try_new, not here. input_batches: Count, input_rows: Count, - baseline_metrics: BaselineMetrics, peak_mem_used: Gauge, + /// Time spent doing the join's own work (including spill write and + /// read-back). The clock is stopped while awaiting the child inputs or + /// the consumer taking an emitted batch — see [`Self::stop_join_time`]. + join_time: Time, + /// Start of the currently running `join_time` span; `None` while the + /// clock is stopped. + join_time_start: Option, // Memory / spill — only the inner key buffer is tracked via reservation, // matching existing SMJ (which tracks only the buffered side). The outer @@ -305,14 +282,6 @@ pub(crate) struct BitwiseSortMergeJoinStream { outer_self_cmp: Option, /// Comparator for inner self-comparison (find_key_group_end on inner) inner_self_cmp: Option, - - // True once the current outer batch has been emitted. The Equal - // branch's inner loops call emit then `ready!(poll_next_outer_batch)`. - // If that poll returns Pending, poll_join re-enters from the top - // on the next poll — with outer_batch still Some and outer_offset - // past the end. The main loop's step 3 would re-emit without this - // guard. Cleared when poll_next_outer_batch loads a new batch. - batch_emitted: bool, } impl BitwiseSortMergeJoinStream { @@ -333,7 +302,7 @@ impl BitwiseSortMergeJoinStream { reservation: MemoryReservation, spill_manager: SpillManager, runtime_env: Arc, - ) -> Result { + ) -> Result { debug_assert!( matches!( join_type, @@ -356,9 +325,10 @@ impl BitwiseSortMergeJoinStream { MetricBuilder::new(metrics).counter("input_batches", partition); let input_rows = MetricBuilder::new(metrics).counter("input_rows", partition); let baseline_metrics = BaselineMetrics::new(metrics, partition); - let peak_mem_used = MetricBuilder::new(metrics).gauge("peak_mem_used", partition); + let peak_mem_used = + MetricBuilder::new(metrics).peak_memory_usage("peak_mem_used", partition); - Ok(Self { + let mut state = Self { join_type, outer, inner, @@ -370,9 +340,6 @@ impl BitwiseSortMergeJoinStream { inner_key_arrays: vec![], matched: BooleanBufferBuilder::new(0), inner_key_buffer: vec![], - inner_key_spill: None, - buffering_inner_pending: false, - pending_boundary: None, on_outer, on_inner, filter, @@ -381,12 +348,12 @@ impl BitwiseSortMergeJoinStream { outer_is_left, coalescer: BatchCoalescer::new(Arc::clone(&schema), batch_size) .with_biggest_coalesce_batch_size(Some(batch_size / 2)), - schema, - join_time, + schema: Arc::clone(&schema), input_batches, input_rows, - baseline_metrics, peak_mem_used, + join_time, + join_time_start: None, reservation, spill_manager, runtime_env, @@ -394,8 +361,39 @@ impl BitwiseSortMergeJoinStream { outer_inner_cmp: None, outer_self_cmp: None, inner_self_cmp: None, - batch_emitted: false, - }) + }; + + let stream = async_try_stream(|mut emitter| async move { + state.start_join_time(); + let result = state.join(&mut emitter).await; + state.stop_join_time(); + result + }); + // ObservedStream records the baseline metrics (output rows/batches, + // end time) exactly as the former hand-written poll_next did. + Ok(Box::pin(ObservedStream::new( + Box::pin(RecordBatchStreamAdapter::new(schema, stream)), + baseline_metrics, + None, + ))) + } + + /// Start (resume) the `join_time` clock. + fn start_join_time(&mut self) { + debug_assert!(self.join_time_start.is_none(), "join_time already running"); + self.join_time_start = Some(Instant::now()); + } + + /// Stop (pause) the `join_time` clock, accumulating the elapsed span. + /// + /// Called around awaits whose duration is not the join's own work: the + /// child input streams' `next()` and `emitter.emit()` (where the + /// consumer processes the batch). The join's own spill read-back is NOT + /// excluded — that time is join work. + fn stop_join_time(&mut self) { + if let Some(start) = self.join_time_start.take() { + self.join_time.add_elapsed(start); + } } /// Resize the memory reservation to match current tracked usage. @@ -445,18 +443,24 @@ impl BitwiseSortMergeJoinStream { Ok(self.inner_self_cmp.as_ref().unwrap()) } - /// Spill the in-memory inner key buffer to disk and clear it. - fn spill_inner_key_buffer(&mut self) -> Result<()> { - let spill_file = self - .spill_manager - .spill_record_batch_and_finish( - &self.inner_key_buffer, - "semi_anti_smj_inner_key_spill", - )? - .expect("inner_key_buffer is non-empty when spilling"); - self.inner_key_buffer.clear(); + /// Spill the in-memory inner key buffer to disk and clear it. One key + /// group can spill repeatedly; every call appends to `writer` — the + /// group's single open spill file — creating it on first use. + fn spill_inner_key_buffer( + &mut self, + writer: &mut Option, + ) -> Result<()> { + if writer.is_none() { + *writer = Some( + self.spill_manager + .create_in_progress_file("semi_anti_smj_inner_key_spill")?, + ); + } + let writer = writer.as_mut().unwrap(); + for batch in self.inner_key_buffer.drain(..) { + writer.append_batch(&batch)?; + } self.inner_buffer_size = 0; - self.inner_key_spill = Some(spill_file); // Should succeed now — inner buffer has been spilled. self.try_resize_reservation() } @@ -467,21 +471,24 @@ impl BitwiseSortMergeJoinStream { /// pool interactions (see apache/datafusion#20729). fn clear_inner_key_group(&mut self) { self.inner_key_buffer.clear(); - self.inner_key_spill = None; self.inner_buffer_size = 0; } - /// Poll for the next outer batch. Returns true if a batch was loaded. - fn poll_next_outer_batch(&mut self, cx: &mut Context<'_>) -> Poll> { + /// Fetch the next outer batch. Returns true if a batch was loaded. + async fn next_outer_batch(&mut self) -> Result { loop { - match ready!(self.outer.poll_next_unpin(cx)) { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.outer.next().await; + self.start_join_time(); + match item { None => { // Release the outer input pipeline's resources. let outer_schema = self.outer.schema(); self.outer = Box::pin(EmptyRecordBatchStream::new(outer_schema)); - return Poll::Ready(Ok(false)); + return Ok(false); } - Some(Err(e)) => return Poll::Ready(Err(e)), + Some(Err(e)) => return Err(e), Some(Ok(batch)) => { let batch_num_rows = batch.num_rows(); self.input_batches.add(1); @@ -495,26 +502,29 @@ impl BitwiseSortMergeJoinStream { self.outer_key_arrays = keys; self.outer_inner_cmp = None; self.outer_self_cmp = None; - self.batch_emitted = false; self.matched = BooleanBufferBuilder::new(batch_num_rows); self.matched.append_n(batch_num_rows, false); - return Poll::Ready(Ok(true)); + return Ok(true); } } } } - /// Poll for the next inner batch. Returns true if a batch was loaded. - fn poll_next_inner_batch(&mut self, cx: &mut Context<'_>) -> Poll> { + /// Fetch the next inner batch. Returns true if a batch was loaded. + async fn next_inner_batch(&mut self) -> Result { loop { - match ready!(self.inner.poll_next_unpin(cx)) { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.inner.next().await; + self.start_join_time(); + match item { None => { // Release the inner input pipeline's resources. let inner_schema = self.inner.schema(); self.inner = Box::pin(EmptyRecordBatchStream::new(inner_schema)); - return Poll::Ready(Ok(false)); + return Ok(false); } - Some(Err(e)) => return Poll::Ready(Err(e)), + Some(Err(e)) => return Err(e), Some(Ok(batch)) => { let batch_num_rows = batch.num_rows(); self.input_batches.add(1); @@ -528,22 +538,17 @@ impl BitwiseSortMergeJoinStream { self.inner_key_arrays = keys; self.outer_inner_cmp = None; self.inner_self_cmp = None; - return Poll::Ready(Ok(true)); + return Ok(true); } } } } - /// Emit the current outer batch through the coalescer, applying the - /// matched bitset as a selection mask. No-op if already emitted - /// (see `batch_emitted` field). + /// Push the current outer batch into the coalescer, applying the matched + /// bitset as a selection mask. Consumes the batch (`outer_batch` becomes + /// `None`). fn emit_outer_batch(&mut self) -> Result<()> { - if self.batch_emitted { - return Ok(()); - } - self.batch_emitted = true; - - let batch = self.outer_batch.as_ref().unwrap(); + let batch = self.outer_batch.take().unwrap(); // finish() converts the bit-packed builder directly to a // BooleanBuffer — no iteration or repacking needed. @@ -566,14 +571,14 @@ impl BitwiseSortMergeJoinStream { } JoinType::LeftSemi | JoinType::RightSemi => { let selection = BooleanArray::new(matched_buf, None); - let filtered = filter_record_batch(batch, &selection)?; + let filtered = filter_record_batch(&batch, &selection)?; if filtered.num_rows() > 0 { self.coalescer.push_batch(filtered)?; } } JoinType::LeftAnti | JoinType::RightAnti => { let selection = not(&BooleanArray::new(matched_buf, None))?; - let filtered = filter_record_batch(batch, &selection)?; + let filtered = filter_record_batch(&batch, &selection)?; if filtered.num_rows() > 0 { self.coalescer.push_batch(filtered)?; } @@ -583,181 +588,139 @@ impl BitwiseSortMergeJoinStream { Ok(()) } - /// Process a key match between outer and inner sides (no filter). - /// Sets matched bits for all outer rows sharing the current key. - fn process_key_match_no_filter(&mut self) -> Result<()> { - let outer_batch = self.outer_batch.as_ref().unwrap(); - let num_outer = outer_batch.num_rows(); + /// Mark all outer rows in the current key group as matched and advance + /// the outer cursor past the group (within the current batch). + fn mark_outer_key_group_matched(&mut self) -> Result<()> { + let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); + let from = self.outer_offset; + let group_end = find_key_group_end(self.get_outer_self_cmp()?, from, num_outer); - self.get_outer_self_cmp()?; - let outer_group_end = find_key_group_end( - self.outer_self_cmp.as_ref().unwrap(), - self.outer_offset, - num_outer, - ); - - for i in self.outer_offset..outer_group_end { + for i in from..group_end { self.matched.set_bit(i, true); } - self.outer_offset = outer_group_end; + self.outer_offset = group_end; Ok(()) } - /// Advance inner past the current key group. Returns Ok(true) if inner + /// Advance the inner cursor past the current key group. The group may + /// span multiple inner batches. Sets `inner_batch` to `None` if inner /// is exhausted. - fn advance_inner_past_key_group( - &mut self, - cx: &mut Context<'_>, - ) -> Poll> { + async fn advance_inner_past_key_group(&mut self) -> Result<()> { loop { - let inner_batch = match &self.inner_batch { - Some(b) => b, - None => return Poll::Ready(Ok(true)), + let Some(inner_batch) = &self.inner_batch else { + return Ok(()); }; let num_inner = inner_batch.num_rows(); - - self.get_inner_self_cmp()?; - let group_end = find_key_group_end( - self.inner_self_cmp.as_ref().unwrap(), - self.inner_offset, - num_inner, - ); + let from = self.inner_offset; + let group_end = + find_key_group_end(self.get_inner_self_cmp()?, from, num_inner); if group_end < num_inner { self.inner_offset = group_end; - return Poll::Ready(Ok(false)); + return Ok(()); } - // Key group extends to end of batch — need to check next batch + // Key group extends to the end of the batch — it may continue + // into the next one; save the last key so we can check. let saved_inner_keys = slice_keys(&self.inner_key_arrays, num_inner - 1); - match ready!(self.poll_next_inner_batch(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(false) => { - return Poll::Ready(Ok(true)); - } - Ok(true) => { - if keys_match( - &saved_inner_keys, - &self.inner_key_arrays, - &self.sort_options, - self.null_equality, - )? { - continue; - } else { - return Poll::Ready(Ok(false)); - } - } + if !self.next_inner_batch().await? { + self.inner_batch = None; + return Ok(()); + } + if !keys_match( + &saved_inner_keys, + &self.inner_key_arrays, + &self.sort_options, + self.null_equality, + )? { + return Ok(()); } } } - /// Buffer inner key group for filter evaluation. Collects all inner rows - /// with the current key across batch boundaries. + /// Buffer the inner key group for filter evaluation, advancing the inner + /// cursor past the group. Collects all inner rows with the current key + /// across batch boundaries. Sets `inner_batch` to `None` if inner is + /// exhausted. /// - /// If poll_next_inner_batch returns Pending, we save progress via - /// buffering_inner_pending. On re-entry (from the Equal branch in - /// poll_join), we skip clear() and the slice+push for the current - /// batch (which was already buffered before Pending), and go directly - /// to polling for the next inner batch. - fn buffer_inner_key_group(&mut self, cx: &mut Context<'_>) -> Poll> { - // On re-entry after Pending: don't clear the partially-filled - // buffer. The current inner_batch was already sliced and pushed - // before Pending, so jump to polling for the next batch. - let mut resume_from_poll = false; - if self.buffering_inner_pending { - self.buffering_inner_pending = false; - resume_from_poll = true; - } else { - self.clear_inner_key_group(); - } - - loop { - if self.inner_batch.is_none() { - return Poll::Ready(Ok(true)); - } - let num_inner = self.inner_batch.as_ref().unwrap().num_rows(); - self.get_inner_self_cmp()?; - let group_end = find_key_group_end( - self.inner_self_cmp.as_ref().unwrap(), - self.inner_offset, - num_inner, - ); - - if !resume_from_poll { - let inner_batch = self.inner_batch.as_ref().unwrap(); - let slice = - inner_batch.slice(self.inner_offset, group_end - self.inner_offset); - self.inner_buffer_size += slice.get_array_memory_size(); - self.inner_key_buffer.push(slice); - - // Reserve memory for the newly buffered slice. If the pool - // is exhausted, spill the entire buffer to disk. - if self.try_resize_reservation().is_err() { - if self.runtime_env.disk_manager.tmp_files_enabled() { - self.spill_inner_key_buffer()?; - } else { - // Re-attempt to get the error message - self.try_resize_reservation().map_err(|e| { - datafusion_common::DataFusionError::Execution(format!( - "{e}. Disk spilling disabled." - )) - })?; - } + /// Slices that overflow the memory pool are appended to a single spill + /// file, returned finished — ready for reading — once the whole group + /// has been buffered. `None` means the group fit in memory. + async fn buffer_inner_key_group(&mut self) -> Result>> { + self.clear_inner_key_group(); + let mut writer: Option = None; + + while let Some(inner_batch) = &self.inner_batch { + let num_inner = inner_batch.num_rows(); + let from = self.inner_offset; + let group_end = + find_key_group_end(self.get_inner_self_cmp()?, from, num_inner); + + let inner_batch = self.inner_batch.as_ref().unwrap(); + let slice = inner_batch.slice(from, group_end - from); + self.inner_buffer_size += slice.get_array_memory_size(); + self.inner_key_buffer.push(slice); + + // Reserve memory for the newly buffered slice. If the pool + // is exhausted, spill the entire buffer to disk. + if self.try_resize_reservation().is_err() { + if self.runtime_env.disk_manager.tmp_files_enabled() { + self.spill_inner_key_buffer(&mut writer)?; + } else { + // Re-attempt to get the error message + self.try_resize_reservation().map_err(|e| { + DataFusionError::Execution(format!( + "{e}. Disk spilling disabled." + )) + })?; } + } - if group_end < num_inner { - self.inner_offset = group_end; - return Poll::Ready(Ok(false)); - } + if group_end < num_inner { + self.inner_offset = group_end; + break; } - resume_from_poll = false; - // Key group extends to end of batch — check next + // Key group extends to the end of the batch — it may continue + // into the next one; save the last key so we can check. let saved_inner_keys = slice_keys(&self.inner_key_arrays, num_inner - 1); - // If poll returns Pending, the current batch is already - // in inner_key_buffer. - self.buffering_inner_pending = true; - match ready!(self.poll_next_inner_batch(cx)) { - Err(e) => { - self.buffering_inner_pending = false; - return Poll::Ready(Err(e)); - } - Ok(false) => { - self.buffering_inner_pending = false; - return Poll::Ready(Ok(true)); - } - Ok(true) => { - self.buffering_inner_pending = false; - if keys_match( - &saved_inner_keys, - &self.inner_key_arrays, - &self.sort_options, - self.null_equality, - )? { - continue; - } else { - return Poll::Ready(Ok(false)); - } - } + if !self.next_inner_batch().await? { + self.inner_batch = None; + break; } + if !keys_match( + &saved_inner_keys, + &self.inner_key_arrays, + &self.sort_options, + self.null_equality, + )? { + break; + } + } + + match writer { + Some(mut writer) => writer.finish(), + None => Ok(None), } } /// Process a key match with a filter. For each inner row in the buffered - /// key group, evaluates the filter against the outer key group and ORs - /// the results into the matched bitset using u64-chunked bitwise ops. - fn process_key_match_with_filter(&mut self) -> Result<()> { - self.get_outer_self_cmp()?; - let filter = self.filter.as_ref().unwrap(); - let outer_batch = self.outer_batch.as_ref().unwrap(); - let num_outer = outer_batch.num_rows(); + /// key group — the spilled slices in `spill` plus the in-memory + /// `inner_key_buffer` — evaluates the filter against the outer key group + /// and ORs the results into the matched bitset using u64-chunked bitwise + /// ops. + async fn process_key_match_with_filter( + &mut self, + spill: Option<&Arc>, + ) -> Result<()> { + let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); // buffer_inner_key_group must be called before this function debug_assert!( - !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(), + !self.inner_key_buffer.is_empty() || spill.is_some(), "process_key_match_with_filter called with no inner key data" ); debug_assert!( @@ -769,40 +732,57 @@ impl BitwiseSortMergeJoinStream { "matched vector must be sized for the current outer batch" ); - let outer_group_end = find_key_group_end( - self.outer_self_cmp.as_ref().unwrap(), - self.outer_offset, - num_outer, - ); - let outer_group_len = outer_group_end - self.outer_offset; - let outer_slice = outer_batch.slice(self.outer_offset, outer_group_len); + let outer_group_start = self.outer_offset; + let outer_group_end = + find_key_group_end(self.get_outer_self_cmp()?, outer_group_start, num_outer); + let outer_group_len = outer_group_end - outer_group_start; + + let filter = self.filter.as_ref().unwrap(); + let outer_batch = self.outer_batch.as_ref().unwrap(); + let outer_slice = outer_batch.slice(outer_group_start, outer_group_len); // Count already-matched bits using popcnt on u64 chunks (zero-copy). let mut matched_count = UnalignedBitChunk::new( self.matched.as_slice(), - self.outer_offset, + outer_group_start, outer_group_len, ) .count_ones(); - // Process spilled inner batches first (read back from disk). - if let Some(spill_file) = &self.inner_key_spill { - let file = BufReader::new(File::open(spill_file.path())?); - let reader = StreamReader::try_new(file, None)?; - for batch_result in reader { - let inner_slice = batch_result?; - matched_count = eval_filter_for_inner_slice( - self.outer_is_left, - filter, - &outer_slice, - &inner_slice, - &mut self.matched, - self.outer_offset, - outer_group_len, - matched_count, - )?; - if matched_count == outer_group_len { - break; + // Process spilled inner batches first asynchronously. + if matched_count < outer_group_len + && let Some(spill_file) = spill + { + let mut spill_stream = self + .spill_manager + .read_spill_as_stream(Arc::clone(spill_file), None)?; + let mut spill_stream_has_data = false; + + // Note: the clock keeps running across the spill reads — the + // spill file is the join's own data, so reading it back is + // join work (unlike the child inputs' `next()`). + while matched_count < outer_group_len { + match spill_stream.next().await { + Some(Ok(inner_slice)) => { + spill_stream_has_data = true; + matched_count = eval_filter_for_inner_slice( + self.outer_is_left, + filter, + &outer_slice, + &inner_slice, + &mut self.matched, + outer_group_start, + outer_group_len, + matched_count, + )?; + } + Some(Err(e)) => return Err(e), + None => { + if !spill_stream_has_data { + return internal_err!("Spill file was empty"); + } + break; + } } } } @@ -819,7 +799,7 @@ impl BitwiseSortMergeJoinStream { &outer_slice, inner_slice, &mut self.matched, - self.outer_offset, + outer_group_start, outer_group_len, matched_count, )?; @@ -830,337 +810,296 @@ impl BitwiseSortMergeJoinStream { } self.outer_offset = outer_group_end; + Ok(()) } - /// Continue processing an outer key group that spans multiple outer - /// batches. Returns `true` if this outer batch was fully consumed - /// by the key group and the caller should load another. - fn resume_boundary(&mut self) -> Result { - debug_assert!( - self.outer_batch.is_some(), - "caller must load outer_batch first" - ); - match self.pending_boundary.take() { - Some(PendingBoundary::NoFilter { saved_keys }) => { - let same_key = keys_match( - &saved_keys, - &self.outer_key_arrays, - &self.sort_options, - self.null_equality, - )?; - if same_key { - self.process_key_match_no_filter()?; - let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); - if self.outer_offset >= num_outer { - self.pending_boundary = Some(PendingBoundary::NoFilter { - saved_keys: slice_keys(&self.outer_key_arrays, num_outer - 1), - }); - self.emit_outer_batch()?; - self.outer_batch = None; - return Ok(true); - } - } + /// Evaluate the filter for the buffered inner key group against the + /// outer key group. If the outer key group continues into subsequent + /// outer batches, keep evaluating there too. Dropping `spill` on return + /// deletes the group's temp file. + async fn process_filtered_match_loop( + &mut self, + spill: Option>, + ) -> Result<()> { + loop { + self.process_key_match_with_filter(spill.as_ref()).await?; + + let outer_batch = self.outer_batch.as_ref().unwrap(); + if self.outer_offset < outer_batch.num_rows() { + break; } - Some(PendingBoundary::Filtered { saved_keys }) => { - debug_assert!( - !self.inner_key_buffer.is_empty() || self.inner_key_spill.is_some(), - "Filtered pending boundary entered but no inner key data exists" - ); - let same_key = keys_match( - &saved_keys, - &self.outer_key_arrays, - &self.sort_options, - self.null_equality, - )?; - if same_key { - self.process_key_match_with_filter()?; - let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); - if self.outer_offset >= num_outer { - self.pending_boundary = Some(PendingBoundary::Filtered { - saved_keys: slice_keys(&self.outer_key_arrays, num_outer - 1), - }); - self.emit_outer_batch()?; - self.outer_batch = None; - return Ok(true); - } - } - self.clear_inner_key_group(); + + // The outer key group may continue into the next outer batch; + // save the last key so we can check. + let saved_keys = + slice_keys(&self.outer_key_arrays, outer_batch.num_rows() - 1); + + self.emit_outer_batch()?; + + if !self.next_outer_batch().await? { + break; + } + if !keys_match( + &saved_keys, + &self.outer_key_arrays, + &self.sort_options, + self.null_equality, + )? { + break; } - None => {} } - Ok(false) - } - /// Main loop: drive the merge-scan to produce output batches. - fn poll_join(&mut self, cx: &mut Context<'_>) -> Poll>> { - let join_time = self.join_time.clone(); - let _timer = join_time.timer(); + self.clear_inner_key_group(); + Ok(()) + } + /// Mark the outer key group as matched. If the outer key group continues + /// into subsequent outer batches, keep marking there too. + async fn process_unfiltered_match_loop(&mut self) -> Result<()> { loop { - // 1. Ensure we have an outer batch - if self.outer_batch.is_none() { - match ready!(self.poll_next_outer_batch(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(false) => { - // Outer exhausted — flush coalescer - self.pending_boundary = None; - self.coalescer.finish_buffered_batch()?; - if let Some(batch) = self.coalescer.next_completed_batch() { - return Poll::Ready(Ok(Some(batch))); - } - return Poll::Ready(Ok(None)); - } - Ok(true) => { - if self.resume_boundary()? { - continue; - } - } - } + self.mark_outer_key_group_matched()?; + + let outer_batch = self.outer_batch.as_ref().unwrap(); + if self.outer_offset < outer_batch.num_rows() { + return Ok(()); } - // 2. Ensure we have an inner batch (unless inner is exhausted). - // Skip this when resuming a pending boundary — inner was already - // advanced past the key group before the boundary loop started. - if self.inner_batch.is_none() && self.pending_boundary.is_none() { - match ready!(self.poll_next_inner_batch(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(false) => { - // Inner exhausted — emit remaining outer batches. - // For semi: no more matches possible. - // For anti: all remaining outer rows are unmatched. - self.emit_outer_batch()?; - self.outer_batch = None; - - loop { - match ready!(self.poll_next_outer_batch(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(false) => break, - Ok(true) => { - self.emit_outer_batch()?; - self.outer_batch = None; - } - } - } + // The outer key group may continue into the next outer batch; + // save the last key so we can check. + let saved_keys = + slice_keys(&self.outer_key_arrays, outer_batch.num_rows() - 1); - self.coalescer.finish_buffered_batch()?; - if let Some(batch) = self.coalescer.next_completed_batch() { - return Poll::Ready(Ok(Some(batch))); - } - return Poll::Ready(Ok(None)); - } - Ok(true) => {} - } + self.emit_outer_batch()?; + + if !self.next_outer_batch().await? { + return Ok(()); + } + if !keys_match( + &saved_keys, + &self.outer_key_arrays, + &self.sort_options, + self.null_equality, + )? { + return Ok(()); } + } + } - // 3. Main merge-scan loop - let outer_batch = self.outer_batch.as_ref().unwrap(); - let num_outer = outer_batch.num_rows(); + /// Keys at both cursors are equal: determine which outer rows in the key + /// group have a match. Both key groups may span batch boundaries. + async fn process_key_match(&mut self) -> Result<()> { + if self.filter.is_some() { + // Buffer the inner key group so each inner row can be evaluated + // against the outer key group, OR-ing filter results into the + // matched bitset. + let spill = self.buffer_inner_key_group().await?; + self.process_filtered_match_loop(spill).await + } else { + // Without a filter, key equality alone means every outer row in + // the group matches; the inner rows themselves are not needed. + self.advance_inner_past_key_group().await?; + self.process_unfiltered_match_loop().await + } + } - if self.outer_offset >= num_outer { - self.emit_outer_batch()?; - self.outer_batch = None; + /// Compare the join keys at the outer and inner cursors, returning the + /// ordering of the outer key relative to the inner key (e.g. `Greater` + /// means outer key > inner key, per the sort options). + fn compare_current_keys(&mut self) -> Result { + let (outer_idx, inner_idx) = (self.outer_offset, self.inner_offset); + Ok(self.get_outer_inner_cmp()?.compare(outer_idx, inner_idx)) + } - if let Some(batch) = self.coalescer.next_completed_batch() { - return Poll::Ready(Ok(Some(batch))); - } - continue; - } + /// Outer key is unmatched: advance the outer cursor past its key group + /// (within the current batch). If the group continues into the next + /// batch, those rows compare Less again and are skipped the same way. + fn skip_outer_key_group(&mut self) -> Result<()> { + let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); + let from = self.outer_offset; + self.outer_offset = + find_key_group_end(self.get_outer_self_cmp()?, from, num_outer); + Ok(()) + } - let inner_batch = match &self.inner_batch { - Some(b) => b, - None => { + /// Sync fast path for `Ordering::Greater`: skip the inner key group when + /// it ends within the current batch. Returns false — leaving all state + /// unchanged — when the group reaches the batch boundary, in which case + /// the caller must take [`Self::advance_inner_past_key_group`]. + fn try_skip_inner_key_group(&mut self) -> Result { + let num_inner = self.inner_batch.as_ref().unwrap().num_rows(); + let from = self.inner_offset; + let group_end = find_key_group_end(self.get_inner_self_cmp()?, from, num_inner); + if group_end >= num_inner { + return Ok(false); + } + self.inner_offset = group_end; + Ok(true) + } + + /// Sync fast path for `Ordering::Equal` without a filter: when both key + /// groups end within their current batches (the common case — a group + /// only reaches a batch boundary once per batch), mark the outer group + /// matched and advance both cursors without any async machinery. + /// Returns false — leaving all state unchanged — when a filter is + /// present or either group reaches a batch boundary, in which case the + /// caller must take [`Self::process_key_match`]. + fn try_process_key_match(&mut self) -> Result { + if self.filter.is_some() { + return Ok(false); + } + + let num_inner = self.inner_batch.as_ref().unwrap().num_rows(); + let inner_from = self.inner_offset; + let inner_group_end = + find_key_group_end(self.get_inner_self_cmp()?, inner_from, num_inner); + if inner_group_end >= num_inner { + return Ok(false); + } + + let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); + let outer_from = self.outer_offset; + let outer_group_end = + find_key_group_end(self.get_outer_self_cmp()?, outer_from, num_outer); + if outer_group_end >= num_outer { + return Ok(false); + } + + for i in outer_from..outer_group_end { + self.matched.set_bit(i, true); + } + self.outer_offset = outer_group_end; + self.inner_offset = inner_group_end; + Ok(true) + } + + /// True when the outer cursor already points at an unprocessed row: the + /// sync fast path of [`Self::advance_outer_row`]. Checked inline in the + /// hot loop so the async helper (and its state machine) is only entered + /// at batch boundaries — same pattern as `sorts/merge.rs`. + fn has_current_outer_row(&self) -> bool { + self.outer_batch + .as_ref() + .is_some_and(|batch| self.outer_offset < batch.num_rows()) + } + + /// True when the inner cursor already points at an unprocessed row: the + /// sync fast path of [`Self::advance_inner_row`]. + fn has_current_inner_row(&self) -> bool { + self.inner_batch + .as_ref() + .is_some_and(|batch| self.inner_offset < batch.num_rows()) + } + + /// Ensure the outer cursor points at an unprocessed row, emitting + /// finished outer batches and loading new ones as needed. Returns false + /// when outer is exhausted. + async fn advance_outer_row( + &mut self, + emitter: &mut TryEmitter, + ) -> Result { + loop { + match &self.outer_batch { + Some(batch) if self.outer_offset < batch.num_rows() => { + return Ok(true); + } + Some(_) => { + // Current batch fully scanned — emit it and load the next. self.emit_outer_batch()?; - self.outer_batch = None; - continue; + self.emit_completed_batches(emitter).await; } - }; - let num_inner = inner_batch.num_rows(); - - if self.inner_offset >= num_inner { - match ready!(self.poll_next_inner_batch(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(false) => { - self.inner_batch = None; - continue; + None => { + if !self.next_outer_batch().await? { + return Ok(false); } - Ok(true) => continue, } } + } + } - // 4. Compare keys at current positions - self.get_outer_inner_cmp()?; - let cmp = self - .outer_inner_cmp - .as_ref() - .unwrap() - .compare(self.outer_offset, self.inner_offset); - - match cmp { - Ordering::Less => { - self.get_outer_self_cmp()?; - let group_end = find_key_group_end( - self.outer_self_cmp.as_ref().unwrap(), - self.outer_offset, - num_outer, - ); - self.outer_offset = group_end; - } + /// Ensure the inner cursor points at an unprocessed row, loading new + /// inner batches as needed. Returns false when inner is exhausted. + async fn advance_inner_row(&mut self) -> Result { + loop { + if let Some(batch) = &self.inner_batch + && self.inner_offset < batch.num_rows() + { + return Ok(true); + } + if !self.next_inner_batch().await? { + self.inner_batch = None; + return Ok(false); + } + } + } + + /// Inner is exhausted, so no further matches are possible: emit the + /// current outer batch and all remaining ones with their current matched + /// bits (semi drops unmatched rows, anti emits them, mark emits them + /// with mark=false). + async fn drain_outer(&mut self) -> Result<()> { + self.emit_outer_batch()?; + while self.next_outer_batch().await? { + self.emit_outer_batch()?; + } + Ok(()) + } + + /// Emit all completed coalescer batches to the stream consumer. + async fn emit_completed_batches( + &mut self, + emitter: &mut TryEmitter, + ) { + while let Some(batch) = self.coalescer.next_completed_batch() { + // While the emitted batch is in the consumer's hands the join + // isn't doing any work. + self.stop_join_time(); + emitter.emit(batch).await; + self.start_join_time(); + } + } + + /// Main loop: a classic merge-scan over the two sorted inputs, emitting + /// output batches as they complete. + async fn join( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + // The `has_current_*` / `has_completed_batch` fast paths keep async + // state machinery out of the per-key-group hot path; the awaiting + // helpers are only entered at batch boundaries. + while self.has_current_outer_row() || self.advance_outer_row(emitter).await? { + if !(self.has_current_inner_row() || self.advance_inner_row().await?) { + self.drain_outer().await?; + break; + } + + // Each arm handles the common case synchronously (`try_*`); the + // async continuations only run when a key group reaches a batch + // boundary or a filter must be evaluated. + match self.compare_current_keys()? { + Ordering::Less => self.skip_outer_key_group()?, Ordering::Greater => { - self.get_inner_self_cmp()?; - let group_end = find_key_group_end( - self.inner_self_cmp.as_ref().unwrap(), - self.inner_offset, - num_inner, - ); - if group_end >= num_inner { - let saved_keys = - slice_keys(&self.inner_key_arrays, num_inner - 1); - match ready!(self.poll_next_inner_batch(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(false) => { - self.inner_batch = None; - continue; - } - Ok(true) => { - if keys_match( - &saved_keys, - &self.inner_key_arrays, - &self.sort_options, - self.null_equality, - )? { - match ready!(self.advance_inner_past_key_group(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(_) => continue, - } - } - continue; - } - } - } else { - self.inner_offset = group_end; + if !self.try_skip_inner_key_group()? { + self.advance_inner_past_key_group().await?; } } Ordering::Equal => { - if self.filter.is_some() { - // Buffer inner key group (may span batches) - match ready!(self.buffer_inner_key_group(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(_inner_exhausted) => {} - } - - // Process outer rows against buffered inner group - // (may need to handle outer batch boundary) - loop { - self.process_key_match_with_filter()?; - - let outer_batch = self.outer_batch.as_ref().unwrap(); - if self.outer_offset >= outer_batch.num_rows() { - let saved_keys = slice_keys( - &self.outer_key_arrays, - outer_batch.num_rows() - 1, - ); - - self.emit_outer_batch()?; - debug_assert!( - !self.inner_key_buffer.is_empty() - || self.inner_key_spill.is_some(), - "Filtered pending boundary requires inner key data in buffer or spill" - ); - self.pending_boundary = - Some(PendingBoundary::Filtered { saved_keys }); - - match ready!(self.poll_next_outer_batch(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(false) => { - self.pending_boundary = None; - self.outer_batch = None; - break; - } - Ok(true) => { - let Some(PendingBoundary::Filtered { - saved_keys, - }) = self.pending_boundary.take() - else { - unreachable!() - }; - let same = keys_match( - &saved_keys, - &self.outer_key_arrays, - &self.sort_options, - self.null_equality, - )?; - if same { - continue; - } - break; - } - } - } else { - break; - } - } - - self.clear_inner_key_group(); - } else { - // No filter: advance inner past key group, then - // mark all outer rows with this key as matched. - match ready!(self.advance_inner_past_key_group(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(_inner_exhausted) => {} - } - - loop { - self.process_key_match_no_filter()?; - - let num_outer = self.outer_batch.as_ref().unwrap().num_rows(); - if self.outer_offset >= num_outer { - let saved_keys = - slice_keys(&self.outer_key_arrays, num_outer - 1); - - self.emit_outer_batch()?; - self.pending_boundary = - Some(PendingBoundary::NoFilter { saved_keys }); - - match ready!(self.poll_next_outer_batch(cx)) { - Err(e) => return Poll::Ready(Err(e)), - Ok(false) => { - self.pending_boundary = None; - self.outer_batch = None; - break; - } - Ok(true) => { - let Some(PendingBoundary::NoFilter { - saved_keys, - }) = self.pending_boundary.take() - else { - unreachable!() - }; - let same_key = keys_match( - &saved_keys, - &self.outer_key_arrays, - &self.sort_options, - self.null_equality, - )?; - if same_key { - continue; - } - break; - } - } - } else { - break; - } - } + if !self.try_process_key_match()? { + self.process_key_match().await?; } } } - // Check for completed coalescer batch - if let Some(batch) = self.coalescer.next_completed_batch() { - return Poll::Ready(Ok(Some(batch))); + if self.coalescer.has_completed_batch() { + self.emit_completed_batches(emitter).await; } } + + // Flush whatever is still buffered in the coalescer. + self.coalescer.finish_buffered_batch()?; + self.emit_completed_batches(emitter).await; + Ok(()) } } @@ -1313,7 +1252,7 @@ fn evaluate_filter_for_inner_row( .as_any() .downcast_ref::() .ok_or_else(|| { - datafusion_common::DataFusionError::Internal( + DataFusionError::Internal( "Filter expression did not return BooleanArray".to_string(), ) })?; @@ -1324,21 +1263,3 @@ fn evaluate_filter_for_inner_row( Ok(bool_arr.clone()) } } - -impl Stream for BitwiseSortMergeJoinStream { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - let poll = self.poll_join(cx).map(|result| result.transpose()); - self.baseline_metrics.record_poll(poll) - } -} - -impl RecordBatchStream for BitwiseSortMergeJoinStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs index 3f309431614a4..b48905500d546 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/exec.rs @@ -38,9 +38,11 @@ use crate::projection::{ physical_to_column_exprs, update_join_on, }; use crate::spill::spill_manager::SpillManager; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, SendableRecordBatchStream, Statistics, check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, InputDistributionRequirements, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, validate_child_count, }; use arrow::compute::SortOptions; @@ -80,8 +82,7 @@ use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequiremen /// on the output batch size of the execution plan. There is no spilling support for streamed input. /// The comparisons are performed from values of join keys in streamed input with the values of /// join keys in buffered input. One row in streamed record batch could be matched with multiple rows in -/// buffered input batches. The streamed input is managed through the states in `StreamedState` -/// and streamed input batches are represented by `StreamedBatch`. +/// buffered input batches. Streamed input batches are represented by `StreamedBatch`. /// /// Buffered input is buffered for all record batches having the same value of join key. /// If the memory limit increases beyond the specified value and spilling is enabled, @@ -91,8 +92,7 @@ use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequiremen /// memory/disk depends on the number of rows of buffered input having the same value /// of join key as that of streamed input rows currently present in memory. Due to pre-sorted inputs, /// the algorithm understands when it is not needed anymore, and releases the buffered batches -/// from memory/disk. The buffered input is managed through the states in `BufferedState` -/// and buffered input batches are represented by `BufferedBatch`. +/// from memory/disk. Buffered input batches are represented by `BufferedBatch`. /// /// Depending on the type of join, left or right input may be selected as streamed or buffered /// respectively. For example, in a left-outer join, the left execution plan will be selected as @@ -344,20 +344,6 @@ impl SortMergeJoinExec { reorder_output_after_swap(Arc::new(new_join), &left.schema(), &right.schema()) } } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for SortMergeJoinExec { @@ -424,15 +410,19 @@ impl ExecutionPlan for SortMergeJoinExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { let (left_expr, right_expr) = self .on .iter() .map(|(l, r)| (Arc::clone(l), Arc::clone(r))) .unzip(); - vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), - ] + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) } fn required_input_ordering(&self) -> Vec> { @@ -452,38 +442,63 @@ impl ExecutionPlan for SortMergeJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn crate::PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to join keys from both sides - let mut tnr = TreeNodeRecursion::Continue; - for (left, right) in &self.on { - tnr = tnr.visit_sibling(|| f(left.as_ref()))?; - tnr = tnr.visit_sibling(|| f(right.as_ref()))?; - } - // Apply to join filter expressions if present - if let Some(filter) = &self.filter { - tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + let filter = self.filter.iter().map(|filter| filter.expression()); + crate::apply_expression_roots(join_keys.chain(filter), f) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => match &children[..] { + [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( + Arc::clone(left), + Arc::clone(right), + self.on.clone(), + self.filter.clone(), + self.join_type, + self.sort_options.clone(), + self.null_equality, + )?)), + _ => internal_err!("SortMergeJoin wrong number of children"), + }, } - Ok(tnr) } fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - match &children[..] { - [left, right] => Ok(Arc::new(SortMergeJoinExec::try_new( - Arc::clone(left), - Arc::clone(right), - self.on.clone(), - self.filter.clone(), - self.join_type, - self.sort_options.clone(), - self.null_equality, - )?)), - _ => internal_err!("SortMergeJoin wrong number of children"), - } + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -540,7 +555,7 @@ impl ExecutionPlan for SortMergeJoinExec { | JoinType::LeftMark | JoinType::RightMark ) { - Ok(Box::pin(BitwiseSortMergeJoinStream::try_new( + BitwiseSortMergeJoinStream::try_new( Arc::clone(&self.schema), self.sort_options.clone(), self.null_equality, @@ -556,9 +571,9 @@ impl ExecutionPlan for SortMergeJoinExec { reservation, spill_manager, context.runtime_env(), - )?)) + ) } else { - Ok(Box::pin(MaterializingSortMergeJoinStream::try_new( + MaterializingSortMergeJoinStream::try_new( Arc::clone(&self.schema), self.sort_options.clone(), self.null_equality, @@ -573,7 +588,7 @@ impl ExecutionPlan for SortMergeJoinExec { reservation, spill_manager, context.runtime_env(), - )?)) + ) } } @@ -581,25 +596,29 @@ impl ExecutionPlan for SortMergeJoinExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition), ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { // SortMergeJoinExec uses symmetric hash partitioning where both left and right // inputs are hash-partitioned on the join keys. This means partition `i` of the // left input is joined with partition `i` of the right input. // - // Therefore, partition-specific statistics can be computed by getting the - // partition-specific statistics from both children and combining them via - // `estimate_join_statistics`. - // // TODO stats: it is not possible in general to know the output size of joins // There are some special cases though, for example: // - `A LEFT JOIN B ON A.col=B.col` with `COUNT_DISTINCT(B.col)=COUNT(B.col)` - let left_stats = Arc::unwrap_or_clone(self.left.partition_statistics(partition)?); - let right_stats = - Arc::unwrap_or_clone(self.right.partition_statistics(partition)?); + let left_stats = input_stats[0].as_ref().clone(); + let right_stats = input_stats[1].as_ref().clone(); Ok(Arc::new(estimate_join_statistics( left_stats, right_stats, &self.on, + self.null_equality, &self.join_type, &self.schema, )?)) @@ -659,4 +678,149 @@ impl ExecutionPlan for SortMergeJoinExec { self.null_equality, )?))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + let on = self + .on() + .iter() + .map(|(left, right)| { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(left)?), + right: Some(ctx.encode_expr(right)?), + }) + }) + .collect::>>()?; + + let join_type = crate::joins::proto::join_type_to_proto(self.join_type()); + let null_equality = + crate::joins::proto::null_equality_to_proto(self.null_equality()); + let filter = self + .filter() + .as_ref() + .map(|filter| crate::joins::proto::join_filter_to_proto(filter, ctx)) + .transpose()?; + let sort_options = self + .sort_options() + .iter() + .map(|options| protobuf::SortExprNode { + expr: None, + asc: !options.descending, + nulls_first: options.nulls_first, + }) + .collect(); + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin(Box::new( + protobuf::SortMergeJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + filter, + sort_options, + null_equality: null_equality.into(), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortMergeJoinExec { + /// Reconstruct a [`SortMergeJoinExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let sort_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SortMergeJoin, + "SortMergeJoinExec", + ); + let left = ctx.decode_required_child( + sort_join.left.as_deref(), + "SortMergeJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + sort_join.right.as_deref(), + "SortMergeJoinExec", + "right", + )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let on = sort_join + .on + .iter() + .map(|columns| { + let left = ctx.decode_required_expr( + columns.left.as_ref(), + left_schema.as_ref(), + "SortMergeJoinExec", + "on.left", + )?; + let right = ctx.decode_required_expr( + columns.right.as_ref(), + right_schema.as_ref(), + "SortMergeJoinExec", + "on.right", + )?; + Ok((left, right)) + }) + .collect::>()?; + + let join_type = crate::joins::proto::join_type_from_proto( + sort_join.join_type, + "SortMergeJoinExec", + )?; + let null_equality = crate::joins::proto::null_equality_from_proto( + sort_join.null_equality, + "SortMergeJoinExec", + )?; + let filter = sort_join + .filter + .as_ref() + .map(|filter| { + crate::joins::proto::join_filter_from_proto( + filter, + ctx, + "SortMergeJoinExec", + ) + }) + .transpose()?; + let sort_options = sort_join + .sort_options + .iter() + .map(|options| SortOptions { + descending: !options.asc, + nulls_first: options.nulls_first, + }) + .collect(); + + Ok(Arc::new(Self::try_new( + left, + right, + on, + filter, + join_type, + sort_options, + null_equality, + )?)) + } } diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs index 9bcc749c23dce..3baa0c4a3e792 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/materializing_stream.rs @@ -17,21 +17,17 @@ //! Sort-Merge Join execution //! -//! This module implements the runtime state machine for the Sort-Merge Join -//! operator. It drives two sorted input streams (the *streamed* side and the -//! *buffered* side), compares join keys, and produces joined `RecordBatch`es. +//! This module implements the Sort-Merge Join operator as an async +//! generator running a merge scan: it drives two sorted input streams (the +//! *streamed* side and the *buffered* side), compares join keys, and +//! produces joined `RecordBatch`es. use std::cmp::Ordering; use std::collections::{HashMap, VecDeque}; -use std::fs::File; -use std::io::BufReader; +use std::fmt::Debug; use std::mem::size_of; use std::ops::Range; -use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering::Relaxed; -use std::task::{Context, Poll}; use crate::joins::sort_merge_join::filter::{ FilterMetadata, filter_record_batch_by_join_type, get_corrected_filter_mask, @@ -39,69 +35,28 @@ use crate::joins::sort_merge_join::filter::{ }; use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; use crate::joins::utils::{JoinFilter, JoinKeyComparator}; -use crate::metrics::RecordOutput; +use crate::metrics::Time; use crate::spill::spill_manager::SpillManager; -use crate::stream::EmptyRecordBatchStream; -use crate::{PhysicalExpr, RecordBatchStream, SendableRecordBatchStream}; +use crate::stream::{EmptyRecordBatchStream, ObservedStream, RecordBatchStreamAdapter}; +use crate::{PhysicalExpr, SendableRecordBatchStream}; use arrow::array::{types::UInt64Type, *}; use arrow::compute::{ self, BatchCoalescer, SortOptions, concat_batches, filter_record_batch, interleave, - take, take_arrays, + take_arrays, }; use arrow::datatypes::SchemaRef; -use arrow::ipc::reader::StreamReader; use datafusion_common::cast::as_uint64_array; -use datafusion_common::{JoinType, NullEquality, Result, exec_err, internal_err}; -use datafusion_execution::disk_manager::RefCountedTempFile; +use datafusion_common::instant::Instant; +use datafusion_common::{ + DataFusionError, JoinType, NullEquality, Result, exec_err, internal_err, +}; use datafusion_execution::memory_pool::MemoryReservation; use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_execution::{SpillFile, TryEmitter, async_try_stream}; use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; -use futures::{Stream, StreamExt}; - -/// State of SMJ stream -#[derive(Debug, PartialEq, Eq)] -pub(super) enum SortMergeJoinState { - /// Init joining with a new streamed row or a new buffered batches - Init, - /// Polling one streamed row or one buffered batch, or both - Polling, - /// Joining polled data and making output - JoinOutput, - /// Emit ready data if have any and then go back to [`Self::Init`] state - EmitReadyThenInit, - /// No more output - Exhausted, -} - -/// State of streamed data stream -#[derive(Debug, PartialEq, Eq)] -pub(super) enum StreamedState { - /// Init polling - Init, - /// Polling one streamed row - Polling, - /// Ready to produce one streamed row - Ready, - /// No more streamed row - Exhausted, -} - -/// State of buffered data stream -#[derive(Debug, PartialEq, Eq)] -pub(super) enum BufferedState { - /// Init polling - Init, - /// Polling first row in the next batch - PollingFirst, - /// Polling rest rows in the next batch - PollingRest, - /// Ready to produce one batch - Ready, - /// No more buffered batches - Exhausted, -} +use futures::StreamExt; /// Represents a chunk of joined data from streamed and buffered side pub(super) struct StreamedJoinedChunk { @@ -223,7 +178,7 @@ pub(super) enum FilterState { /// A buffered batch that contains contiguous rows with same join key /// -/// `BufferedBatch` can exist as either an in-memory `RecordBatch` or a `RefCountedTempFile` on disk. +/// `BufferedBatch` can exist as either an in-memory `RecordBatch` or a `SpillFile`. #[derive(Debug)] pub(super) struct BufferedBatch { /// Represents in memory or spilled record batch @@ -299,14 +254,23 @@ impl BufferedBatch { // TODO: Spill join arrays (https://github.com/apache/datafusion/pull/17429) // Used to represent whether the buffered data is currently in memory or written to disk -#[derive(Debug)] pub(super) enum BufferedBatchState { // In memory record batch InMemory(RecordBatch), // Spilled temp file - Spilled(RefCountedTempFile), + Spilled(Arc), } +impl Debug for BufferedBatchState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::InMemory(batch) => f.debug_tuple("InMemory").field(batch).finish(), + Self::Spilled(_) => { + write!(f, "Spilled(Custom_Backend)") + } + } + } +} /// Sort-Merge join stream for Inner/Left/Right/Full joins. /// /// Named "materializing" because it builds explicit `(streamed, buffered)` row @@ -328,6 +292,9 @@ pub(super) struct MaterializingSortMergeJoinStream { pub filter: Option, /// How the join is performed pub join_type: JoinType, + /// Cached `needs_deferred_filtering(filter, join_type)` — both inputs + /// are fixed at construction time. + pub deferred_filtering: bool, /// Target output batch size pub batch_size: usize, @@ -341,10 +308,8 @@ pub(super) struct MaterializingSortMergeJoinStream { pub streamed: SendableRecordBatchStream, /// Current processing record batch of streamed pub streamed_batch: StreamedBatch, - /// (used in outer join) Is current streamed row joined at least once? - pub streamed_joined: bool, - /// State of streamed - pub streamed_state: StreamedState, + /// True once the streamed input has no more rows + pub streamed_exhausted: bool, /// Join key columns of streamed pub on_streamed: Vec, @@ -358,10 +323,11 @@ pub(super) struct MaterializingSortMergeJoinStream { pub buffered: SendableRecordBatchStream, /// Current buffered data pub buffered_data: BufferedData, - /// (used in outer join) Is current buffered batches joined at least once? - pub buffered_joined: bool, - /// State of buffered - pub buffered_state: BufferedState, + /// Has any streamed row matched the current buffered key group? + /// (FULL join: an unmatched group is emitted null-joined when passed.) + pub buffered_group_matched: bool, + /// True once the buffered input has no more rows and no group remains + pub buffered_exhausted: bool, /// Join key columns of buffered pub on_buffered: Vec, @@ -370,18 +336,26 @@ pub(super) struct MaterializingSortMergeJoinStream { // These fields track the execution state of merge join and are updated // during the execution. // ======================================================================== - /// Current state of the stream - pub state: SortMergeJoinState, /// Staging output array builders pub joined_record_batches: JoinedRecordBatches, /// Output buffer. Currently used by filtering as it requires double buffering - /// to avoid small/empty batches. Non-filtered join outputs directly from `staging_output_record_batches.batches` + /// to avoid small/empty batches. Non-filtered joins output directly from + /// `joined_record_batches.joined_batches` pub output: BatchCoalescer, - /// The comparison result of current streamed row and buffered batches - pub current_ordering: Ordering, /// Manages the process of spilling and reading back intermediate data pub spill_manager: SpillManager, + /// Tracks the number of batches currently spilled + pub spilled_batch_count: usize, + + /// Time spent doing the join's own work (including spill write and + /// read-back). The clock is stopped while awaiting the child inputs or + /// the consumer taking an emitted batch — see [`Self::stop_join_time`]. + pub join_time: Time, + /// Start of the currently running `join_time` span; `None` while the + /// clock is stopped. + pub join_time_start: Option, + // ======================================================================== // CACHED COMPARATORS: // Pre-built comparators to avoid per-row type dispatch in hot loops. @@ -401,8 +375,9 @@ pub(super) struct MaterializingSortMergeJoinStream { pub reservation: MemoryReservation, /// Runtime env pub runtime_env: Arc, - /// A unique number for each batch - pub streamed_batch_counter: AtomicUsize, + /// A unique id per streamed batch, tagging deferred-filter metadata so + /// `get_corrected_filter_mask` can group output rows by input batch. + pub streamed_batch_counter: usize, } /// Staging area for joined data before output @@ -548,240 +523,6 @@ impl JoinedRecordBatches { self.debug_assert_empty_consistency(); } } -impl RecordBatchStream for MaterializingSortMergeJoinStream { - fn schema(&self) -> SchemaRef { - Arc::clone(&self.schema) - } -} - -impl Stream for MaterializingSortMergeJoinStream { - type Item = Result; - - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - let join_time = self.join_metrics.join_time().clone(); - let _timer = join_time.timer(); - loop { - match &self.state { - SortMergeJoinState::Init => { - let streamed_exhausted = - self.streamed_state == StreamedState::Exhausted; - let buffered_exhausted = - self.buffered_state == BufferedState::Exhausted; - self.state = if streamed_exhausted && buffered_exhausted { - SortMergeJoinState::Exhausted - } else { - match self.current_ordering { - Ordering::Less | Ordering::Equal => { - if !streamed_exhausted { - // Batch deferred filtering: process_filtered_batches() - // only when >= batch_size rows have accumulated. - // Without this gate, unique keys cause per-row pipeline - // execution (concat + correct_mask + filter_by_type), - // which dominates runtime. - // - // Accumulated rows are bounded to ~2*batch_size: - // one batch_size worth from freeze_dequeuing_buffered() - // (when an input batch is fully consumed), plus up to - // batch_size pairs accumulating toward the next freeze. - // This does not reintroduce the unbounded buffering - // fixed by PR #20482. Exhausted state flushes remainder. - if needs_deferred_filtering( - &self.filter, - self.join_type, - ) { - let accumulated = self.num_unfrozen_pairs() - + self - .joined_record_batches - .filter_metadata - .filter_mask - .len(); - if accumulated >= self.batch_size { - match self.process_filtered_batches()? { - Poll::Ready(Some(batch)) => { - return Poll::Ready(Some(Ok(batch))); - } - Poll::Ready(None) | Poll::Pending => {} - } - } - } - - self.streamed_joined = false; - self.streamed_state = StreamedState::Init; - } - } - Ordering::Greater => { - if !buffered_exhausted { - self.buffered_joined = false; - self.buffered_state = BufferedState::Init; - } - } - } - SortMergeJoinState::Polling - }; - } - SortMergeJoinState::Polling => { - if ![StreamedState::Exhausted, StreamedState::Ready] - .contains(&self.streamed_state) - { - match self.poll_streamed_row(cx)? { - Poll::Ready(_) => {} - Poll::Pending => return Poll::Pending, - } - } - - if ![BufferedState::Exhausted, BufferedState::Ready] - .contains(&self.buffered_state) - { - match self.poll_buffered_batches(cx)? { - Poll::Ready(_) => {} - Poll::Pending => return Poll::Pending, - } - } - let streamed_exhausted = - self.streamed_state == StreamedState::Exhausted; - let buffered_exhausted = - self.buffered_state == BufferedState::Exhausted; - if streamed_exhausted && buffered_exhausted { - self.state = SortMergeJoinState::Exhausted; - continue; - } - self.current_ordering = self.compare_streamed_buffered()?; - self.state = SortMergeJoinState::JoinOutput; - } - SortMergeJoinState::EmitReadyThenInit => { - // If have data to emit, emit it and if no more, change to next - - // Verify metadata alignment before checking if we have batches to output - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); - - // For filtered joins, skip output and let Init state handle it - if needs_deferred_filtering(&self.filter, self.join_type) { - self.state = SortMergeJoinState::Init; - continue; - } - - // For non-filtered joins, only output if we have a completed batch - // (opportunistic output when target batch size is reached) - if self - .joined_record_batches - .joined_batches - .has_completed_batch() - { - let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); - } - self.state = SortMergeJoinState::Init; - } - SortMergeJoinState::JoinOutput => { - self.join_partial()?; - - if self.num_unfrozen_pairs() < self.batch_size { - if self.buffered_data.scanning_finished() { - self.buffered_data.scanning_reset(); - self.state = SortMergeJoinState::EmitReadyThenInit; - } - } else { - self.freeze_all()?; - - // Verify metadata alignment before checking if we have batches to output - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); - - // For filtered joins, skip output and let Init state handle it - if needs_deferred_filtering(&self.filter, self.join_type) { - continue; - } - - // For non-filtered joins, only output if we have a completed batch - // (opportunistic output when target batch size is reached) - if self - .joined_record_batches - .joined_batches - .has_completed_batch() - { - let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); - } - // Otherwise keep buffering (don't output yet) - } - } - SortMergeJoinState::Exhausted => { - self.freeze_all()?; - - // Verify metadata alignment before final output - self.joined_record_batches - .filter_metadata - .debug_assert_metadata_aligned(); - - // For filtered joins, must concat and filter ALL data at once - if needs_deferred_filtering(&self.filter, self.join_type) - && !self.joined_record_batches.joined_batches.is_empty() - { - let record_batch = self.filter_joined_batch()?; - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); - } - - // For non-filtered joins, finish buffered data first - if !self.joined_record_batches.joined_batches.is_empty() { - self.joined_record_batches - .joined_batches - .finish_buffered_batch()?; - } - - // Output one completed batch at a time (stay in Exhausted until empty) - if self - .joined_record_batches - .joined_batches - .has_completed_batch() - { - let record_batch = self - .joined_record_batches - .joined_batches - .next_completed_batch() - .expect("has_completed_batch was true"); - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); - } - - // Finally check self.output BatchCoalescer (used by filtered joins) - return if !self.output.is_empty() { - self.output.finish_buffered_batch()?; - let record_batch = self - .output - .next_completed_batch() - .expect("Failed to get last batch"); - (&record_batch) - .record_output(&self.join_metrics.baseline_metrics()); - Poll::Ready(Some(Ok(record_batch))) - } else { - Poll::Ready(None) - }; - } - } - } - } -} impl MaterializingSortMergeJoinStream { #[expect(clippy::too_many_arguments)] @@ -800,7 +541,7 @@ impl MaterializingSortMergeJoinStream { reservation: MemoryReservation, spill_manager: SpillManager, runtime_env: Arc, - ) -> Result { + ) -> Result { let streamed_schema = streamed.schema(); let buffered_schema = buffered.schema(); debug_assert!( @@ -811,8 +552,8 @@ impl MaterializingSortMergeJoinStream { "MaterializingSortMergeJoinStream does not handle {join_type:?}; \ semi/anti/mark joins use BitwiseSortMergeJoinStream" ); - Ok(Self { - state: SortMergeJoinState::Init, + let join_time = join_metrics.join_time(); + let mut this = Self { sort_options, null_equality, schema: Arc::clone(&schema), @@ -822,13 +563,12 @@ impl MaterializingSortMergeJoinStream { buffered, streamed_batch: StreamedBatch::new_empty(streamed_schema), buffered_data: BufferedData::default(), - streamed_joined: false, - buffered_joined: false, - streamed_state: StreamedState::Init, - buffered_state: BufferedState::Init, - current_ordering: Ordering::Equal, + buffered_group_matched: false, + streamed_exhausted: false, + buffered_exhausted: false, on_streamed, on_buffered, + deferred_filtering: needs_deferred_filtering(&filter, join_type), filter, joined_record_batches: JoinedRecordBatches { joined_batches: BatchCoalescer::new(Arc::clone(&schema), batch_size) @@ -843,10 +583,299 @@ impl MaterializingSortMergeJoinStream { reservation, runtime_env, spill_manager, + spilled_batch_count: 0, + join_time, + join_time_start: None, streamed_buffered_cmp: None, buffered_equality_cmp: None, - streamed_batch_counter: AtomicUsize::new(0), - }) + streamed_batch_counter: 0, + }; + + let schema = Arc::clone(&this.schema); + let baseline_metrics = this.join_metrics.baseline_metrics(); + + let stream = async_try_stream(|mut emitter| async move { + this.start_join_time(); + let result = this.join(&mut emitter).await; + this.stop_join_time(); + result + }); + // ObservedStream records the baseline metrics (output rows/batches, + // end time). + Ok(Box::pin(ObservedStream::new( + Box::pin(RecordBatchStreamAdapter::new(schema, stream)), + baseline_metrics, + None, + ))) + } + + /// Main loop: the textbook sort-merge join. + /// + /// Both inputs arrive sorted on the join keys. The streamed side is + /// consumed one row at a time; the buffered side one key *group* (all + /// contiguous rows sharing a key) at a time + async fn join( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + // 1. Load the first streamed row and the first buffered key group. + self.load_next_streamed_batch().await?; + self.advance_buffered_group().await?; + + // 2. Merge-scan while either input still has rows. + while !(self.streamed_exhausted && self.buffered_exhausted) { + // Flush the deferred-filtering pipeline once a full batch of + // rows accumulated (filtered outer joins output through it). + if self.deferred_filtering + && self.deferred_rows_accumulated() >= self.batch_size + { + self.emit_deferred_output(emitter).await?; + } + + // 3. Compare the join keys at both cursors. An exhausted side + // compares as the larger one, so the other side keeps + // draining through its own arm. + match self.compare_streamed_buffered()? { + // 3a. The streamed row can never match: null-join it (outer + // joins emit it; inner joins drop it), then advance. + Ordering::Less => { + self.null_join_streamed_row(); + if self.num_unfrozen_pairs() >= self.batch_size { + self.freeze_and_emit(emitter).await?; + } + if !self.try_advance_streamed_row() { + self.load_next_streamed_batch().await?; + } + } + // 3b. The buffered group can never match again: null-join + // it if nothing matched it (FULL join), then advance to + // the next key group. + Ordering::Greater => { + self.null_join_buffered_group(); + if !self.try_advance_buffered_group()? { + self.advance_buffered_group().await?; + } + } + // 3c. Match: pair the streamed row with the whole group — + // materializing ("freezing") mid-scan whenever a full + // batch of pairs accumulates — then advance streamed. + // The group stays for the next streamed row. + Ordering::Equal => { + while !self.pair_streamed_row_with_group() { + self.freeze_and_emit(emitter).await?; + } + if !self.try_advance_streamed_row() { + self.load_next_streamed_batch().await?; + } + } + } + + // 4. Emit completed output batches (filtered joins emit + // through the deferred-filtering pipeline above instead). + if !self.deferred_filtering + && self + .joined_record_batches + .joined_batches + .has_completed_batch() + { + self.emit_completed_joined_batches(emitter).await; + } + } + + // 5. Flush everything that remains. + self.on_children_exhausted(emitter).await + } + + /// `Equal`: pair the current streamed row with every row of the + /// buffered key group, and mark the group as matched. + /// + /// Returns false when a full batch of pairs has accumulated (the scan + /// may or may not be complete): the caller must materialize + /// (`freeze_and_emit`) and call again, which resumes the scan where it + /// paused. Returns true when the group scan is complete and there is + /// room for more pairs. + fn pair_streamed_row_with_group(&mut self) -> bool { + while !self.buffered_data.scanning_finished() + && self.num_unfrozen_pairs() < self.batch_size + { + let scanning_idx = self.buffered_data.scanning_idx(); + self.streamed_batch.append_output_pair( + Some(self.buffered_data.scanning_batch_idx), + Some(scanning_idx), + self.batch_size, + ); + self.buffered_data.scanning_advance(); + } + if self.num_unfrozen_pairs() >= self.batch_size { + return false; + } + + self.buffered_group_matched = true; + self.buffered_data.scanning_reset(); + true + } + + /// `Less` (outer joins): no buffered row matches the current streamed + /// row — emit it joined to NULLs. Inner joins emit nothing. + fn null_join_streamed_row(&mut self) { + if matches!( + self.join_type, + JoinType::Left | JoinType::Right | JoinType::Full + ) { + let scanning_batch_idx = if self.buffered_data.scanning_finished() { + None + } else { + Some(self.buffered_data.scanning_batch_idx) + }; + self.streamed_batch.append_output_pair( + scanning_batch_idx, + None, + self.batch_size, + ); + } + self.buffered_data.scanning_reset(); + } + + /// `Greater` (FULL join): the buffered group can never match a streamed + /// row anymore — if nothing matched it, mark all its rows for + /// null-joined output (produced when the group's batches are dequeued). + fn null_join_buffered_group(&mut self) { + if self.join_type == JoinType::Full && !self.buffered_group_matched { + while !self.buffered_data.scanning_finished() { + let scanning_idx = self.buffered_data.scanning_idx(); + self.buffered_data + .scanning_batch_mut() + .null_joined + .push(scanning_idx); + self.buffered_data.scanning_advance(); + } + } + self.buffered_data.scanning_reset(); + } + + /// Start (resume) the `join_time` clock. + fn start_join_time(&mut self) { + debug_assert!(self.join_time_start.is_none(), "join_time already running"); + self.join_time_start = Some(Instant::now()); + } + + /// Stop (pause) the `join_time` clock, accumulating the elapsed span. + /// + /// Called around awaits whose duration is not the join's own work: the + /// child input streams' `next()` and `emitter.emit()` (where the + /// consumer processes the batch). The join's own spill write and + /// read-back are NOT excluded — that time is join work. + fn stop_join_time(&mut self) { + if let Some(start) = self.join_time_start.take() { + self.join_time.add_elapsed(start); + } + } + + /// Number of rows currently waiting in the deferred-filtering pipeline. + /// + /// Typically bounded to ~2*batch_size: one batch_size worth from + /// freeze_dequeuing_buffered() (when an input batch is fully consumed), + /// plus up to batch_size pairs accumulating toward the next freeze. A + /// single streamed row matching a very large key group can exceed that + /// (its pairs freeze into the pipeline before the gate runs again — same + /// as the pre-generator design). This does not reintroduce the unbounded + /// buffering fixed by PR #20482; `on_children_exhausted` flushes the + /// remainder. + fn deferred_rows_accumulated(&self) -> usize { + self.num_unfrozen_pairs() + + self.joined_record_batches.filter_metadata.filter_mask.len() + } + + /// Run the deferred-filtering pipeline over everything accumulated so + /// far and emit its completed output, if any. Clears the accumulation + /// it processed. + /// + /// The caller gates this on `deferred_rows_accumulated() >= batch_size`: + /// running the pipeline per row instead (concat + correct_mask + + /// filter_by_type) would dominate runtime for unique keys. + async fn emit_deferred_output( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + // Ensure required spilled batches are restored to memory before + // processing, as this path invokes freeze_all(). + self.restore_spilled_batches_for_freeze().await?; + if let Some(batch) = self.process_filtered_batches()? { + // While the emitted batch is in the consumer's hands the join + // isn't doing any work. + self.stop_join_time(); + emitter.emit(batch).await; + self.start_join_time(); + } + Ok(()) + } + + /// Restore every spilled buffered batch that the next freeze needs. + async fn restore_spilled_batches_for_freeze(&mut self) -> Result<()> { + let needed = self.get_required_batch_indices(self.buffered_data.batches.len()); + self.restore_spilled_batches(&needed).await + } + + /// Emit all completed joined batches to the stream consumer. + async fn emit_completed_joined_batches( + &mut self, + emitter: &mut TryEmitter, + ) { + while let Some(record_batch) = self + .joined_record_batches + .joined_batches + .next_completed_batch() + { + // While the emitted batch is in the consumer's hands the join + // isn't doing any work. + self.stop_join_time(); + emitter.emit(record_batch).await; + self.start_join_time(); + } + } + + /// Flush everything that remains once both inputs are exhausted. + async fn on_children_exhausted( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + // Freeze the remaining pairs, restoring any spilled batches needed. + self.restore_spilled_batches_for_freeze().await?; + self.freeze_all()?; + + // Verify metadata alignment before final output + self.joined_record_batches + .filter_metadata + .debug_assert_metadata_aligned(); + + if self.deferred_filtering { + // Filtered joins must concat and filter ALL remaining data at once + if !self.joined_record_batches.joined_batches.is_empty() { + let record_batch = self.filter_joined_batch()?; + self.stop_join_time(); + emitter.emit(record_batch).await; + self.start_join_time(); + } + } else if !self.joined_record_batches.joined_batches.is_empty() { + // For non-filtered joins, finish buffered data first, then emit + // every completed batch. + self.joined_record_batches + .joined_batches + .finish_buffered_batch()?; + self.emit_completed_joined_batches(emitter).await; + } + + // Drain the double-buffering coalescer used by filtered joins. + if !self.output.is_empty() { + self.output.finish_buffered_batch()?; + while let Some(record_batch) = self.output.next_completed_batch() { + self.stop_join_time(); + emitter.emit(record_batch).await; + self.start_join_time(); + } + } + + Ok(()) } /// Build a comparator for streamed vs buffered head batch keys. @@ -889,9 +918,9 @@ impl MaterializingSortMergeJoinStream { /// Process accumulated batches for filtered joins /// - /// Freezes unfrozen pairs, applies deferred filtering, and outputs if ready. - /// Returns Poll::Ready with a batch if one is available, otherwise Poll::Pending. - fn process_filtered_batches(&mut self) -> Poll>> { + /// Freezes unfrozen pairs, applies deferred filtering, and returns a + /// completed output batch if one is ready. + fn process_filtered_batches(&mut self) -> Result> { self.freeze_all()?; self.joined_record_batches @@ -909,60 +938,128 @@ impl MaterializingSortMergeJoinStream { .output .next_completed_batch() .expect("Failed to get output batch"); - (&record_batch).record_output(&self.join_metrics.baseline_metrics()); - return Poll::Ready(Some(Ok(record_batch))); + return Ok(Some(record_batch)); } } - Poll::Pending + Ok(None) } - /// Poll next streamed row - fn poll_streamed_row(&mut self, cx: &mut Context) -> Poll>> { - loop { - match &self.streamed_state { - StreamedState::Init => { - if self.streamed_batch.idx + 1 < self.streamed_batch.batch.num_rows() - { - self.streamed_batch.idx += 1; - self.streamed_state = StreamedState::Ready; - return Poll::Ready(Some(Ok(()))); - } else { - self.streamed_state = StreamedState::Polling; - } - } - StreamedState::Polling => match self.streamed.poll_next_unpin(cx)? { - Poll::Pending => { - return Poll::Pending; - } - Poll::Ready(None) => { - // Release the streamed input pipeline's resources. - let streamed_schema = self.streamed.schema(); - self.streamed = - Box::pin(EmptyRecordBatchStream::new(streamed_schema)); - self.streamed_state = StreamedState::Exhausted; + /// Identifies which buffered batches are needed for the upcoming freeze operation + fn get_required_batch_indices(&self, buffered_freeze_count: usize) -> Vec { + let mut needed = vec![]; + // Avoid scanning if no spilled batches exist + if self.spilled_batch_count == 0 { + return needed; + } + // We need all batches that matched with streamed rows + for chunk in &self.streamed_batch.output_indices { + if let Some(idx) = chunk.buffered_batch_idx { + needed.push(idx); + } + } + + // Full Joins need to emit null-joined rows, so we need batches up to freeze_count + if self.join_type == JoinType::Full { + needed.extend(0..buffered_freeze_count); + } + + needed.sort_unstable(); + needed.dedup(); + needed + } + + /// Asynchronously reads spilled batches back into memory. + /// Only processes the required indices to avoid OOMs. + async fn restore_spilled_batches( + &mut self, + required_indices: &[usize], + ) -> Result<()> { + for &idx in required_indices { + // Guard against indices that might be out of bounds if the queue was cleared + if idx >= self.buffered_data.batches.len() { + continue; + } + + let bb = &mut self.buffered_data.batches[idx]; + + if let BufferedBatchState::Spilled(spill_file) = &bb.batch { + let mut spill_stream = self + .spill_manager + .read_spill_as_stream(Arc::clone(spill_file), None)?; + + match spill_stream.next().await.transpose()? { + Some(batch) => { + // Transition the batch back to InMemory + bb.batch = BufferedBatchState::InMemory(batch); + self.spilled_batch_count -= 1; + // The batch is back in memory, so we must account for its size. + let newly_allocated = + bb.size_estimation.saturating_sub(bb.reserved_amount); + self.reservation.grow(newly_allocated); + bb.reserved_amount = bb.size_estimation; + + self.join_metrics + .peak_mem_used() + .set_max(self.reservation.size()); } - Poll::Ready(Some(batch)) => { - if batch.num_rows() > 0 { - self.freeze_streamed()?; - self.join_metrics.input_batches().add(1); - self.join_metrics.input_rows().add(batch.num_rows()); - self.streamed_batch = - StreamedBatch::new(batch, &self.on_streamed); - self.rebuild_streamed_buffered_cmp()?; - // Every incoming streaming batch should have its unique id - // Check `JoinedRecordBatches.self.streamed_batch_counter` documentation - self.streamed_batch_counter - .fetch_add(1, std::sync::atomic::Ordering::SeqCst); - self.streamed_state = StreamedState::Ready; - } + None => { + return internal_err!("Spill file was empty"); } - }, - StreamedState::Ready => { - return Poll::Ready(Some(Ok(()))); } - StreamedState::Exhausted => { - return Poll::Ready(None); + } + } + + Ok(()) + } + + /// Sync fast path of advancing the streamed cursor: move to the next row + /// of the current batch. Returns false at the batch boundary, where the + /// caller must load the next batch via + /// [`Self::load_next_streamed_batch`]. + fn try_advance_streamed_row(&mut self) -> bool { + if self.streamed_batch.idx + 1 < self.streamed_batch.batch.num_rows() { + self.streamed_batch.idx += 1; + return true; + } + false + } + + /// Load the next streamed batch (freezing the finished one) and point + /// the streamed cursor at its first row. Sets `streamed_exhausted` when + /// the streamed input has no more rows. + async fn load_next_streamed_batch(&mut self) -> Result<()> { + loop { + // Loading a new streamed batch freezes the current one, which + // materializes buffered columns — restore any spilled buffered + // batches it needs first. + self.restore_spilled_batches_for_freeze().await?; + + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.streamed.next().await.transpose(); + self.start_join_time(); + match item? { + None => { + // Release the streamed input pipeline's resources. + let streamed_schema = self.streamed.schema(); + self.streamed = + Box::pin(EmptyRecordBatchStream::new(streamed_schema)); + self.streamed_exhausted = true; + return Ok(()); + } + Some(batch) => { + if batch.num_rows() > 0 { + self.freeze_streamed()?; + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + self.streamed_batch = + StreamedBatch::new(batch, &self.on_streamed); + self.rebuild_streamed_buffered_cmp()?; + // Every incoming streamed batch gets a unique id. + self.streamed_batch_counter += 1; + return Ok(()); + } } } } @@ -997,6 +1094,7 @@ impl MaterializingSortMergeJoinStream { .unwrap(); // Operation only return None if no batches are spilled, here we ensure that at least one batch is spilled buffered_batch.batch = BufferedBatchState::Spilled(spill_file); + self.spilled_batch_count += 1; // Join key arrays remain in memory after the batch is // spilled — the comparator needs them for key boundary @@ -1025,125 +1123,185 @@ impl MaterializingSortMergeJoinStream { Ok(()) } - /// Poll next buffered batches - fn poll_buffered_batches(&mut self, cx: &mut Context) -> Poll>> { + /// Sync fast path of [`Self::advance_buffered_group`]: when the next + /// group starts in the single remaining buffered batch and provably ends + /// within it (the common case — a group only reaches a batch boundary + /// once per batch), advance entirely synchronously. Returns false — + /// leaving all state unchanged — when the async path must run instead. + fn try_advance_buffered_group(&mut self) -> Result { + if self.buffered_data.batches.len() != 1 { + return Ok(false); + } + let head_batch = self.buffered_data.head_batch(); + if head_batch.range.end == head_batch.num_rows { + // Fully consumed — needs dequeuing (and loading the next batch). + return Ok(false); + } + + if self.buffered_equality_cmp.is_none() { + self.rebuild_buffered_equality_cmp()?; + } + let cmp = self.buffered_equality_cmp.as_ref().unwrap(); + + // Scan the next group's extent before committing any state, so a + // bail-out (the group may span into the next batch) leaves + // everything untouched for the async path. + let batch = self.buffered_data.head_batch(); + let group_start = batch.range.end; + let mut group_end = group_start + 1; + while group_end < batch.num_rows && cmp.is_equal(group_start, group_end) { + group_end += 1; + } + if group_end == batch.num_rows { + return Ok(false); + } + + let batch = self.buffered_data.tail_batch_mut(); + batch.range.start = group_start; + batch.range.end = group_end; + self.buffered_group_matched = false; + Ok(true) + } + + /// Advance the buffered side to the next key group: dequeue batches + /// fully consumed by the previous group, then collect all contiguous + /// rows sharing the next join key (the group may span multiple buffered + /// batches). Sets `buffered_exhausted` when no group remains. + async fn advance_buffered_group(&mut self) -> Result<()> { + self.buffered_group_matched = false; + self.dequeue_consumed_buffered_batches().await?; + + if self.buffered_data.batches.is_empty() { + // Load the batch holding the first row of the next group. + if !self.load_next_buffered_batch().await? { + self.buffered_exhausted = true; + return Ok(()); + } + } else { + // Seed the next group at the first unconsumed row of the + // remaining batch. + let tail_batch = self.buffered_data.tail_batch_mut(); + tail_batch.range.start = tail_batch.range.end; + tail_batch.range.end += 1; + } + + self.extend_buffered_group().await + } + + /// Dequeue buffered batches fully consumed by the previous group, + /// producing their pending output (e.g. Full-join null-joined rows). + async fn dequeue_consumed_buffered_batches(&mut self) -> Result<()> { + let mut head_changed = false; + while !self.buffered_data.batches.is_empty() { + let head_batch = self.buffered_data.head_batch(); + if head_batch.range.end != head_batch.num_rows { + // The next group starts within the head batch: streamed rows + // will be joined with the head batch in the next step. + break; + } + // load the spilled head batch before dequeuing + let needed = self.get_required_batch_indices(1); + self.restore_spilled_batches(&needed).await?; + + self.freeze_dequeuing_buffered()?; + if let Some(mut buffered_batch) = self.buffered_data.batches.pop_front() { + self.produce_buffered_not_matched(&mut buffered_batch)?; + self.free_reservation(&buffered_batch); + if matches!(buffered_batch.batch, BufferedBatchState::Spilled(_)) { + self.spilled_batch_count -= 1; + } + head_changed = true; + } + } + if head_changed { + self.streamed_buffered_cmp = None; + self.buffered_equality_cmp = None; + } + Ok(()) + } + + /// Load the next non-empty buffered batch and seed a new group with its + /// first row. Returns false when the buffered input is exhausted. + async fn load_next_buffered_batch(&mut self) -> Result { loop { - match &self.buffered_state { - BufferedState::Init => { - // pop previous buffered batches - let mut head_changed = false; - while !self.buffered_data.batches.is_empty() { - let head_batch = self.buffered_data.head_batch(); - // If the head batch is fully processed, dequeue it and produce output of it. - if head_batch.range.end == head_batch.num_rows { - self.freeze_dequeuing_buffered()?; - if let Some(mut buffered_batch) = - self.buffered_data.batches.pop_front() - { - self.produce_buffered_not_matched(&mut buffered_batch)?; - self.free_reservation(&buffered_batch); - head_changed = true; - } - } else { - // If the head batch is not fully processed, break the loop. - // Streamed batch will be joined with the head batch in the next step. - break; - } - } - if head_changed { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.buffered.next().await.transpose(); + self.start_join_time(); + match item? { + None => { + // Release the buffered input pipeline's resources. + let buffered_schema = self.buffered.schema(); + self.buffered = + Box::pin(EmptyRecordBatchStream::new(buffered_schema)); + return Ok(false); + } + Some(batch) => { + self.join_metrics.input_batches().add(1); + self.join_metrics.input_rows().add(batch.num_rows()); + + if batch.num_rows() > 0 { + let buffered_batch = + BufferedBatch::new(batch, 0..1, &self.on_buffered); + self.allocate_reservation(buffered_batch)?; self.streamed_buffered_cmp = None; - self.buffered_equality_cmp = None; + return Ok(true); } - if self.buffered_data.batches.is_empty() { - self.buffered_state = BufferedState::PollingFirst; + } + } + } + } + + /// Extend the current group with every following row that shares its + /// key, loading more buffered batches as needed. + async fn extend_buffered_group(&mut self) -> Result<()> { + loop { + if self.buffered_data.tail_batch().range.end + < self.buffered_data.tail_batch().num_rows + { + if self.buffered_equality_cmp.is_none() { + self.rebuild_buffered_equality_cmp()?; + } + while self.buffered_data.tail_batch().range.end + < self.buffered_data.tail_batch().num_rows + { + if self.buffered_equality_cmp.as_ref().unwrap().is_equal( + self.buffered_data.head_batch().range.start, + self.buffered_data.tail_batch().range.end, + ) { + self.buffered_data.tail_batch_mut().range.end += 1; } else { - let tail_batch = self.buffered_data.tail_batch_mut(); - tail_batch.range.start = tail_batch.range.end; - tail_batch.range.end += 1; - self.buffered_state = BufferedState::PollingRest; + // Group complete within the current batch. + return Ok(()); } } - BufferedState::PollingFirst => match self.buffered.poll_next_unpin(cx)? { - Poll::Pending => { - return Poll::Pending; - } - Poll::Ready(None) => { + } else { + // The child's execution time is its own, not join_time. + self.stop_join_time(); + let item = self.buffered.next().await.transpose(); + self.start_join_time(); + match item? { + None => { + // Group complete; the input is done but the group is + // still valid — `buffered_exhausted` is only set once + // it has been fully consumed and dequeued. // Release the buffered input pipeline's resources. let buffered_schema = self.buffered.schema(); self.buffered = Box::pin(EmptyRecordBatchStream::new(buffered_schema)); - self.buffered_state = BufferedState::Exhausted; - return Poll::Ready(None); + return Ok(()); } - Poll::Ready(Some(batch)) => { + Some(batch) => { + // Polling batches coming concurrently as multiple partitions self.join_metrics.input_batches().add(1); self.join_metrics.input_rows().add(batch.num_rows()); - if batch.num_rows() > 0 { let buffered_batch = - BufferedBatch::new(batch, 0..1, &self.on_buffered); - + BufferedBatch::new(batch, 0..0, &self.on_buffered); self.allocate_reservation(buffered_batch)?; - self.streamed_buffered_cmp = None; - self.buffered_state = BufferedState::PollingRest; + self.buffered_equality_cmp = None; } } - }, - BufferedState::PollingRest => { - if self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows - { - if self.buffered_equality_cmp.is_none() { - self.rebuild_buffered_equality_cmp()?; - } - while self.buffered_data.tail_batch().range.end - < self.buffered_data.tail_batch().num_rows - { - if self.buffered_equality_cmp.as_ref().unwrap().is_equal( - self.buffered_data.head_batch().range.start, - self.buffered_data.tail_batch().range.end, - ) { - self.buffered_data.tail_batch_mut().range.end += 1; - } else { - self.buffered_state = BufferedState::Ready; - return Poll::Ready(Some(Ok(()))); - } - } - } else { - match self.buffered.poll_next_unpin(cx)? { - Poll::Pending => { - return Poll::Pending; - } - Poll::Ready(None) => { - // Release the buffered input pipeline's resources. - let buffered_schema = self.buffered.schema(); - self.buffered = Box::pin(EmptyRecordBatchStream::new( - buffered_schema, - )); - self.buffered_state = BufferedState::Ready; - } - Poll::Ready(Some(batch)) => { - // Polling batches coming concurrently as multiple partitions - self.join_metrics.input_batches().add(1); - self.join_metrics.input_rows().add(batch.num_rows()); - if batch.num_rows() > 0 { - let buffered_batch = BufferedBatch::new( - batch, - 0..0, - &self.on_buffered, - ); - self.allocate_reservation(buffered_batch)?; - self.buffered_equality_cmp = None; - } - } - } - } - } - BufferedState::Ready => { - return Poll::Ready(Some(Ok(()))); - } - BufferedState::Exhausted => { - return Poll::Ready(None); } } } @@ -1151,7 +1309,7 @@ impl MaterializingSortMergeJoinStream { /// Get comparison result of streamed row and buffered batches fn compare_streamed_buffered(&mut self) -> Result { - if self.streamed_state == StreamedState::Exhausted { + if self.streamed_exhausted { return Ok(Ordering::Greater); } if !self.buffered_data.has_buffered_rows() { @@ -1167,81 +1325,23 @@ impl MaterializingSortMergeJoinStream { )) } - /// Produce join and fill output buffer until reaching target batch size - /// or the join is finished - fn join_partial(&mut self) -> Result<()> { - // Whether to join streamed rows - let mut join_streamed = false; - // Whether to join buffered rows - let mut join_buffered = false; - - // determine whether we need to join streamed/buffered rows - match self.current_ordering { - Ordering::Less => { - if matches!( - self.join_type, - JoinType::Left | JoinType::Right | JoinType::Full - ) { - join_streamed = !self.streamed_joined; - } - } - Ordering::Equal => { - join_streamed = true; - join_buffered = true; - } - Ordering::Greater => { - if self.join_type == JoinType::Full { - join_buffered = !self.buffered_joined; - }; - } - } - if !join_streamed && !join_buffered { - // no joined data - self.buffered_data.scanning_finish(); - return Ok(()); - } - - if join_buffered { - // joining streamed/nulls and buffered - while !self.buffered_data.scanning_finished() - && self.num_unfrozen_pairs() < self.batch_size - { - let scanning_idx = self.buffered_data.scanning_idx(); - if join_streamed { - // Join streamed row and buffered row - self.streamed_batch.append_output_pair( - Some(self.buffered_data.scanning_batch_idx), - Some(scanning_idx), - self.batch_size, - ); - } else { - // Join nulls and buffered row for FULL join - self.buffered_data - .scanning_batch_mut() - .null_joined - .push(scanning_idx); - } - self.buffered_data.scanning_advance(); + /// Materialize ("freeze") the accumulated pairs — restoring any spilled + /// batches they reference first — and emit completed output batches + /// (filtered joins emit through the deferred-filtering gate instead). + async fn freeze_and_emit( + &mut self, + emitter: &mut TryEmitter, + ) -> Result<()> { + self.restore_spilled_batches_for_freeze().await?; + self.freeze_all()?; - if self.buffered_data.scanning_finished() { - self.streamed_joined = join_streamed; - self.buffered_joined = true; - } - } - } else { - // joining streamed and nulls - let scanning_batch_idx = if self.buffered_data.scanning_finished() { - None - } else { - Some(self.buffered_data.scanning_batch_idx) - }; - self.streamed_batch.append_output_pair( - scanning_batch_idx, - None, - self.batch_size, - ); - self.buffered_data.scanning_finish(); - self.streamed_joined = true; + if !self.deferred_filtering + && self + .joined_record_batches + .joined_batches + .has_completed_batch() + { + self.emit_completed_joined_batches(emitter).await; } Ok(()) } @@ -1375,7 +1475,7 @@ impl MaterializingSortMergeJoinStream { // but must flow through the same pipeline as matched rows to // preserve output ordering. Use null metadata as a sentinel so // get_corrected_filter_mask() passes them through unchanged. - if needs_deferred_filtering(&self.filter, self.join_type) { + if self.deferred_filtering { self.joined_record_batches .push_batch_with_null_metadata(batch, self.join_type); } else { @@ -1478,12 +1578,12 @@ impl MaterializingSortMergeJoinStream { filter_result_mask.clone() }; - if needs_deferred_filtering(&self.filter, self.join_type) { + if self.deferred_filtering { self.joined_record_batches.push_batch_with_filter_metadata( output_batch, &combined_left_indices, &mask, - self.streamed_batch_counter.load(Relaxed), + self.streamed_batch_counter, self.join_type, ); } else { @@ -1556,18 +1656,6 @@ impl MaterializingSortMergeJoinStream { as_uint64_array(&compute::concat(&refs)?)?.clone() }; - let spill_reservation = self.reservation.new_empty(); - if matches!( - &self.buffered_data.batches[first_batch_idx].batch, - BufferedBatchState::Spilled(_) - ) { - spill_reservation - .grow(self.buffered_data.batches[first_batch_idx].size_estimation); - self.join_metrics - .peak_mem_used() - .set_max(self.reservation.size() + spill_reservation.size()); - } - return fetch_right_columns_by_idxs( &self.buffered_data, first_batch_idx, @@ -1603,29 +1691,20 @@ impl MaterializingSortMergeJoinStream { let num_right_cols = self.buffered_schema.fields().len(); // Read each source batch once (spilled batches require disk I/O). - // Track memory for each spilled batch at the point of deserialization - // so the pool reflects actual usage as it grows. - let spill_reservation = self.reservation.new_empty(); - let mut source_data: Vec> = - Vec::with_capacity(source_batches.len()); - for &idx in &source_batches { - let bb = &self.buffered_data.batches[idx]; - match &bb.batch { - BufferedBatchState::InMemory(batch) => { - source_data.push(Some(batch.clone())); - } - BufferedBatchState::Spilled(spill_file) => { - spill_reservation.grow(bb.size_estimation); - self.join_metrics - .peak_mem_used() - .set_max(self.reservation.size() + spill_reservation.size()); - - let file = BufReader::new(File::open(spill_file.path())?); - let reader = StreamReader::try_new(file, None)?; - source_data.push(reader.into_iter().next().transpose()?); + let source_data_result: Result> = source_batches + .iter() + .map(|&idx| { + let bb = &self.buffered_data.batches[idx]; + match &bb.batch { + BufferedBatchState::InMemory(batch) => Ok(batch.clone()), + BufferedBatchState::Spilled(_) => { + internal_err!("Buffered batch should have been unspilled before fetching columns") + } } - } - } + }) + .collect(); + + let source_data = source_data_result?; let mut right_columns = Vec::with_capacity(num_right_cols); for col_idx in 0..num_right_cols { @@ -1637,14 +1716,7 @@ impl MaterializingSortMergeJoinStream { source_arrays.push(null_array.as_ref()); for data in &source_data { - match data { - Some(batch) => source_arrays.push(batch.column(col_idx).as_ref()), - None => { - return internal_err!( - "Failed to read spilled buffered batch during interleave" - ); - } - } + source_arrays.push(data.column(col_idx).as_ref()); } right_columns.push(interleave(&source_arrays, &interleave_indices)?); } @@ -1838,32 +1910,17 @@ fn fetch_right_columns_from_batch_by_idxs( buffered_indices: &UInt64Array, ) -> Result> { match &buffered_batch.batch { - // In memory batch - // In memory batch BufferedBatchState::InMemory(batch) => { - // When indices form a contiguous range (common in SMJ since the - // buffered side is scanned sequentially), use zero-copy slice. if let Some(range) = is_contiguous_range(buffered_indices) { Ok(batch.slice(range.start, range.len()).columns().to_vec()) } else { Ok(take_arrays(batch.columns(), buffered_indices, None)?) } } - // If the batch was spilled to disk, less likely - BufferedBatchState::Spilled(spill_file) => { - let mut buffered_cols: Vec = - Vec::with_capacity(buffered_indices.len()); - - let file = BufReader::new(File::open(spill_file.path())?); - let reader = StreamReader::try_new(file, None)?; - - for batch in reader { - batch?.columns().iter().for_each(|column| { - buffered_cols.extend(take(column, &buffered_indices, None)) - }); - } - - Ok(buffered_cols) + BufferedBatchState::Spilled(_) => { + internal_err!( + "Buffered batch should have been unspilled before fetching columns" + ) } } } @@ -1873,9 +1930,9 @@ fn fetch_right_columns_from_batch_by_idxs( pub(super) struct BufferedData { /// Buffered batches with the same key pub batches: VecDeque, - /// current scanning batch index used in join_partial() + /// current scanning batch index used by the group-scan phase pub scanning_batch_idx: usize, - /// current scanning offset used in join_partial() + /// current scanning offset used by the group-scan phase pub scanning_offset: usize, } @@ -1928,11 +1985,6 @@ impl BufferedData { pub fn scanning_finished(&self) -> bool { self.scanning_batch_idx == self.batches.len() } - - pub fn scanning_finish(&mut self) { - self.scanning_batch_idx = self.batches.len(); - self.scanning_offset = 0; - } } /// Get join array refs of given batch and join columns diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs b/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs index 62efb77f877ab..6f52a2234b3dc 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/metrics.rs @@ -46,9 +46,8 @@ impl SortMergeJoinMetrics { let input_rows = MetricBuilder::new(metrics) .with_category(MetricCategory::Rows) .counter("input_rows", partition); - let peak_mem_used = MetricBuilder::new(metrics) - .with_category(MetricCategory::Bytes) - .gauge("peak_mem_used", partition); + let peak_mem_used = + MetricBuilder::new(metrics).peak_memory_usage("peak_mem_used", partition); let baseline_metrics = BaselineMetrics::new(metrics, partition); diff --git a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs index c4377b3189ff7..175a9c0ea7198 100644 --- a/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs +++ b/datafusion/physical-plan/src/joins/sort_merge_join/tests.rs @@ -27,6 +27,7 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use std::time::Duration; use super::bitwise_stream::BitwiseSortMergeJoinStream; use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn}; @@ -49,7 +50,9 @@ use arrow::compute::{BatchCoalescer, SortOptions, filter_record_batch}; use arrow::datatypes::{DataType, Field, Schema}; use arrow_ord::sort::SortColumn; use arrow_schema::SchemaRef; +use bytes::Bytes; use datafusion_common::JoinType::*; +use datafusion_common::instant::Instant; use datafusion_common::{ JoinSide, internal_err, test_util::{batches_to_sort_string, batches_to_string}, @@ -59,9 +62,12 @@ use datafusion_common::{ }; use datafusion_common_runtime::JoinSet; use datafusion_execution::config::SessionConfig; -use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; +use datafusion_execution::disk_manager::{ + DiskManager, DiskManagerBuilder, DiskManagerMode, +}; use datafusion_execution::memory_pool::MemoryConsumer; use datafusion_execution::runtime_env::RuntimeEnvBuilder; +use datafusion_execution::spill_file::{SpillFile, SpillWriter, TempFileFactory}; use datafusion_execution::{SendableRecordBatchStream, TaskContext}; use datafusion_expr::Operator; use datafusion_physical_expr::expressions::BinaryExpr; @@ -70,6 +76,7 @@ use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use futures::{Stream, StreamExt}; use insta::assert_snapshot; use itertools::Itertools; +use std::collections::VecDeque; fn build_table( a: (&str, &Vec), @@ -175,7 +182,7 @@ fn build_fixed_size_binary_table( let batch = RecordBatch::try_new( Arc::new(schema), vec![ - Arc::new(FixedSizeBinaryArray::from(a.1.clone())), + Arc::new(FixedSizeBinaryArray::try_from_iter(a.1.iter().copied()).unwrap()), Arc::new(Int32Array::from(b.1.clone())), Arc::new(Int32Array::from(c.1.clone())), ], @@ -2459,6 +2466,24 @@ async fn overallocation_multi_batch_spill() -> Result<()> { assert!(join.metrics().unwrap().spilled_bytes().unwrap() > 0); assert!(join.metrics().unwrap().spilled_rows().unwrap() > 0); + // For Full joins, get_required_batch_indices extends 0..batches.len(), so + // poll_spilled_batches can restore all spilled batches at once via infallible + // grow(). Verify accounting tracked the transient spike and cleaned up. + let peak_mem = join + .metrics() + .and_then(|m| m.sum_by_name("peak_mem_used")) + .map(|m| m.as_usize()) + .unwrap_or(0); + assert!( + peak_mem > 0, + "peak_mem_used should be > 0 for {join_type:?} batch_size={batch_size}" + ); + assert_eq!( + runtime.memory_pool.reserved(), + 0, + "memory should be fully released after {join_type:?} completes + (batch_size={batch_size}): infallible grow during restore must be balanced" + ); // Run the test with no spill configuration as let task_ctx_no_spill = TaskContext::default().with_session_config(session_config.clone()); @@ -3365,7 +3390,7 @@ async fn test_left_outer_join_filtered_mask() -> Result<()> { #[test] fn test_partition_statistics() -> Result<()> { - use crate::ExecutionPlan; + use crate::statistics::{StatisticsArgs, StatisticsContext}; use datafusion_common::stats::Precision; let left = build_table( @@ -3402,7 +3427,8 @@ fn test_partition_statistics() -> Result<()> { // Test aggregate statistics (partition = None) // Should return meaningful statistics computed from both inputs - let stats = join_exec.partition_statistics(None)?; + let stats = + StatisticsContext::new().compute(&join_exec, &StatisticsArgs::new())?; assert_eq!( stats.column_statistics.len(), expected_cols, @@ -3420,7 +3446,8 @@ fn test_partition_statistics() -> Result<()> { // Since the child TestMemoryExec returns unknown stats for specific partitions, // the join output will also have Absent num_rows. This is expected behavior // as the statistics depend on what the children can provide. - let partition_stats = join_exec.partition_statistics(Some(0))?; + let partition_stats = StatisticsContext::new() + .compute(&join_exec, &StatisticsArgs::new().with_partition(Some(0)))?; assert_eq!( partition_stats.column_statistics.len(), expected_cols, @@ -3789,7 +3816,7 @@ async fn consume_stream_until_finish_barrier_reached( let mut after_finish_barrier_reached = vec![]; let mut background_task = JoinSet::new(); - let mut start_time_since_last_ready = datafusion_common::instant::Instant::now(); + let mut start_time_since_last_ready = Instant::now(); loop { let next_item = output_stream.next(); @@ -3809,7 +3836,7 @@ async fn consume_stream_until_finish_barrier_reached( } else { output_batched.push(batch); } - start_time_since_last_ready = datafusion_common::instant::Instant::now(); + start_time_since_last_ready = Instant::now(); } Poll::Ready(Some(Err(e))) => return Err(e), Poll::Ready(None) if !switch_to_finish_barrier => { @@ -3836,9 +3863,7 @@ async fn consume_stream_until_finish_barrier_reached( } // Make sure the test doesn't run forever - if start_time_since_last_ready.elapsed() - > std::time::Duration::from_secs(5) - { + if start_time_since_last_ready.elapsed() > Duration::from_secs(5) { return internal_err!( "Stream should have emitted data by now, but it's still pending. Output batches so far: {}", output_batched.len() @@ -4007,7 +4032,7 @@ fn columns(schema: &Schema) -> Vec { // ==================== BitwiseSortMergeJoinStream direct tests ==================== // // These tests construct a BitwiseSortMergeJoinStream directly (bypassing exec) -// to exercise async re-entry and spill edge cases using PendingStream. +// to exercise waiting on inputs and spill edge cases using PendingStream. /// Create test memory/spill resources for stream-level tests. fn test_stream_resources( @@ -4087,18 +4112,353 @@ impl RecordBatchStream for PendingStream { } /// Helper: collect all output from a BitwiseSortMergeJoinStream. -async fn collect_stream(stream: BitwiseSortMergeJoinStream) -> Result> { - common::collect(Box::pin(stream)).await +async fn collect_stream(stream: SendableRecordBatchStream) -> Result> { + common::collect(stream).await +} + +// ==================== join_time metric tests ==================== +// +// These verify that `join_time` measures only the join's own work: waiting +// for either child input or for the consumer to take an emitted batch must +// not be counted. + +/// Stream that sleeps `delay` before yielding each batch, to simulate a +/// slow input. +fn delayed_stream( + batches: Vec, + delay: Duration, +) -> SendableRecordBatchStream { + let schema = batches[0].schema(); + Box::pin(crate::stream::RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(batches.into_iter().map(Ok)).then(move |item| async move { + tokio::time::sleep(delay).await; + item + }), + )) } -/// Reproduces the buffer_inner_key_group re-entry bug: +/// Three 2-row batches with unique matching keys. +fn join_time_batches() -> Vec { + vec![ + build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 2]), + ("c1", &vec![7, 8]), + ), + build_table_i32( + ("a1", &vec![2, 3]), + ("b1", &vec![3, 4]), + ("c1", &vec![7, 8]), + ), + build_table_i32( + ("a1", &vec![4, 5]), + ("b1", &vec![5, 6]), + ("c1", &vec![7, 8]), + ), + ] +} + +/// Build a no-filter LeftSemi bitwise stream over the given input streams. +/// The small batch size makes each outer batch surface as its own output +/// batch, so a slow consumer test sees multiple emits. +fn join_time_test_join( + outer: SendableRecordBatchStream, + inner: SendableRecordBatchStream, +) -> (SendableRecordBatchStream, ExecutionPlanMetricsSet) { + let metrics = ExecutionPlanMetricsSet::new(); + let outer_schema = outer.schema(); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(inner.schema(), &metrics); + let stream = BitwiseSortMergeJoinStream::try_new( + outer_schema, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + outer, + inner, + vec![Arc::new(Column::new("b1", 1)) as PhysicalExprRef], + vec![Arc::new(Column::new("b1", 1)) as PhysicalExprRef], + None, + LeftSemi, + 2, + 0, + &metrics, + reservation, + spill_manager, + runtime_env, + ) + .unwrap(); + (stream, metrics) +} + +fn join_time_of(metrics: &ExecutionPlanMetricsSet) -> Duration { + Duration::from_nanos( + metrics + .clone_inner() + .sum_by_name("join_time") + .map(|m| m.as_usize()) + .unwrap_or(0) as u64, + ) +} + +/// Run a join with the given injected `delay`, retrying with 4x the delay +/// (up to 3 attempts) when `join_time < delay` fails. +/// +/// This de-flakes the check without masking real bugs: a genuine exclusion +/// bug makes `join_time` absorb the injected waits, so it scales with the +/// delay and fails at every escalation level. Only a fixed-size disturbance +/// (e.g. the OS preempting the test thread while the join_time clock is +/// running) is filtered out, since it cannot grow 4x with the delay. /// -/// When buffer_inner_key_group buffers inner rows across batch boundaries -/// and poll_next_inner_batch returns Pending mid-way, the ready! macro -/// exits poll_join. On re-entry, the merge-scan reaches Equal again and -/// calls buffer_inner_key_group a second time -- which starts with -/// clear(), destroying the partially collected inner rows. Previously -/// consumed batches are gone, so re-buffering misses them. +/// `run` returns `(join_time, wall)` for one join execution. Deterministic +/// invariants (row counts, wall-time lower bounds) stay as asserts inside +/// `run` — deliberately: a panic there fails the test immediately without +/// retrying, since those cannot flake and escalation would only mask a real +/// bug. Likewise `Err` from `run` (join execution failure) propagates +/// immediately. Only the preemption-sensitive `join_time` check is retried. +async fn check_join_time_excluded(mut run: F) -> Result<()> +where + F: FnMut(Duration) -> Fut, + Fut: Future>, +{ + let mut delay = Duration::from_millis(50); + for attempt in 0..3 { + let (join_time, wall) = run(delay).await?; + if join_time < delay { + return Ok(()); + } + assert!( + attempt < 2, + "join_time ({join_time:?}) should be well below the injected \ + delay ({delay:?}) even after escalating retries; wall {wall:?}" + ); + delay *= 4; + } + unreachable!() +} + +/// join_time must not include time spent waiting for the outer input. +#[tokio::test] +async fn join_time_excludes_outer_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let outer = delayed_stream(join_time_batches(), delay); + let inner = delayed_stream(join_time_batches(), Duration::ZERO); + let (stream, metrics) = join_time_test_join(outer, inner); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all outer rows should match"); + assert!( + wall >= delay * 3, + "outer delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time spent waiting for the inner input. +#[tokio::test] +async fn join_time_excludes_inner_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let outer = delayed_stream(join_time_batches(), Duration::ZERO); + let inner = delayed_stream(join_time_batches(), delay); + let (stream, metrics) = join_time_test_join(outer, inner); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all outer rows should match"); + assert!( + wall >= delay * 3, + "inner delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time the consumer spends holding an emitted +/// batch (the generator is suspended inside `emitter.emit` meanwhile). +#[tokio::test] +async fn join_time_excludes_consumer_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let outer = delayed_stream(join_time_batches(), Duration::ZERO); + let inner = delayed_stream(join_time_batches(), Duration::ZERO); + let (mut stream, metrics) = join_time_test_join(outer, inner); + + let start = Instant::now(); + let mut output_batches = 0u32; + while let Some(batch) = stream.next().await { + batch?; + output_batches += 1; + // Simulate a slow consumer between emitted batches. + tokio::time::sleep(delay).await; + } + let wall = start.elapsed(); + + assert!( + output_batches >= 3, + "expected multiple emitted batches, got {output_batches}" + ); + assert!( + wall >= delay * output_batches, + "consumer delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// Three 2-row batches with unique matching keys, right-side column names. +fn join_time_batches_right() -> Vec { + vec![ + build_table_i32( + ("a2", &vec![0, 1]), + ("b2", &vec![1, 2]), + ("c2", &vec![7, 8]), + ), + build_table_i32( + ("a2", &vec![2, 3]), + ("b2", &vec![3, 4]), + ("c2", &vec![7, 8]), + ), + build_table_i32( + ("a2", &vec![4, 5]), + ("b2", &vec![5, 6]), + ("c2", &vec![7, 8]), + ), + ] +} + +/// Build a no-filter Inner materializing join over the given input streams. +/// The small batch size makes the output surface as multiple batches, so a +/// slow consumer test sees multiple emits. +fn materializing_join_time_test_join( + streamed: SendableRecordBatchStream, + buffered: SendableRecordBatchStream, +) -> (SendableRecordBatchStream, ExecutionPlanMetricsSet) { + use crate::joins::sort_merge_join::materializing_stream::MaterializingSortMergeJoinStream; + use crate::joins::sort_merge_join::metrics::SortMergeJoinMetrics; + + let metrics = ExecutionPlanMetricsSet::new(); + let out_schema = Arc::new(Schema::new( + streamed + .schema() + .fields() + .iter() + .chain(buffered.schema().fields().iter()) + .map(|f| f.as_ref().clone()) + .collect::>(), + )); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(buffered.schema(), &metrics); + let stream = MaterializingSortMergeJoinStream::try_new( + out_schema, + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + streamed, + buffered, + vec![Arc::new(Column::new("b1", 1)) as _], + vec![Arc::new(Column::new("b2", 1)) as _], + None, + Inner, + 2, + SortMergeJoinMetrics::new(0, &metrics), + reservation, + spill_manager, + runtime_env, + ) + .unwrap(); + (stream, metrics) +} + +/// join_time must not include time spent waiting for the streamed input. +#[tokio::test] +async fn materializing_join_time_excludes_streamed_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), delay); + let buffered = delayed_stream(join_time_batches_right(), Duration::ZERO); + let (stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all rows should match"); + assert!( + wall >= delay * 3, + "streamed delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time spent waiting for the buffered input. +#[tokio::test] +async fn materializing_join_time_excludes_buffered_input_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), Duration::ZERO); + let buffered = delayed_stream(join_time_batches_right(), delay); + let (stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let batches = collect_stream(stream).await?; + let wall = start.elapsed(); + + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(rows, 6, "all rows should match"); + assert!( + wall >= delay * 3, + "buffered delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// join_time must not include time the consumer spends holding an emitted +/// batch (the generator is suspended inside `emitter.emit` meanwhile). +#[tokio::test] +async fn materializing_join_time_excludes_consumer_wait() -> Result<()> { + check_join_time_excluded(|delay| async move { + let streamed = delayed_stream(join_time_batches(), Duration::ZERO); + let buffered = delayed_stream(join_time_batches_right(), Duration::ZERO); + let (mut stream, metrics) = materializing_join_time_test_join(streamed, buffered); + + let start = Instant::now(); + let mut output_batches = 0u32; + while let Some(batch) = stream.next().await { + batch?; + output_batches += 1; + // Simulate a slow consumer between emitted batches. + tokio::time::sleep(delay).await; + } + let wall = start.elapsed(); + + assert!( + output_batches >= 3, + "expected multiple emitted batches, got {output_batches}" + ); + assert!( + wall >= delay * output_batches, + "consumer delays should dominate wall time, got {wall:?}" + ); + Ok((join_time_of(&metrics), wall)) + }) + .await +} + +/// An inner key group spanning multiple inner batches must survive the inner +/// input returning Pending mid-way: inner rows delivered before the Pending +/// still take part in the filter evaluation. /// /// Setup: /// - Inner: 3 single-row batches, all with key=1, filter values c2=[10, 20, 30] @@ -4106,8 +4466,7 @@ async fn collect_stream(stream: BitwiseSortMergeJoinStream) -> Result Result<()> { let left_schema = Arc::new(Schema::new(vec![ @@ -4224,22 +4583,17 @@ async fn filter_buffer_pending_loses_inner_rows() -> Result<()> { Ok(()) } -/// Reproduces the no-filter boundary Pending re-entry bug: -/// -/// When an outer key group spans a batch boundary, the no-filter path -/// emits the current batch, then polls for the next outer batch. If -/// poll returns Pending, poll_join exits. On re-entry, without the -/// PendingBoundary fix, the new batch is processed fresh by the -/// merge-scan. Since inner already advanced past this key, the outer -/// rows with the matching key are skipped via Ordering::Less. +/// A matched outer key group spanning a batch boundary must survive the outer +/// input returning Pending at that boundary: the rows continuing the key group +/// still count as matched, even though the inner side has already advanced +/// past the key. /// /// Setup: /// - Outer: 2 single-row batches, both with key=1 (key group spans boundary) /// - Inner: 1 row with key=1 /// - Pending injected on outer before 2nd batch /// -/// Without fix: only first outer row emitted (second lost on re-entry) -/// With fix: both outer rows emitted +/// Expected: both outer rows emitted #[tokio::test] async fn no_filter_boundary_pending_loses_outer_rows() -> Result<()> { let left_schema = Arc::new(Schema::new(vec![ @@ -4323,8 +4677,106 @@ async fn no_filter_boundary_pending_loses_outer_rows() -> Result<()> { Ok(()) } -/// Tests the filtered boundary Pending re-entry: outer key group spans -/// batches with a filter, and poll_next_outer_batch returns Pending. +/// Verifies no-filter semi/anti joins when a matching outer key group spans +/// multiple batches and the next outer batch is temporarily unavailable. +/// +/// The outer input has an unmatched prefix row followed by a matching key +/// group that continues in the next batch. Both rows with key=1 should be +/// treated as matched. Returning `Pending` before the second batch makes the +/// join wait for the continuation while the key group is still open. +#[tokio::test] +async fn no_filter_boundary_pending_with_unmatched_prefix() -> Result<()> { + let left_schema = Arc::new(Schema::new(vec![ + Field::new("a1", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c1", DataType::Int32, false), + ])); + let right_schema = Arc::new(Schema::new(vec![ + Field::new("a2", DataType::Int32, false), + Field::new("b1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])); + + // Key=0 is unmatched. Key=1 matches inner and spans the batch boundary. + let outer_batch1 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![0, 1])), + Arc::new(Int32Array::from(vec![0, 1])), + Arc::new(Int32Array::from(vec![0, 10])), + ], + )?; + let outer_batch2 = RecordBatch::try_new( + Arc::clone(&left_schema), + vec![ + Arc::new(Int32Array::from(vec![2])), + Arc::new(Int32Array::from(vec![1])), // same key + Arc::new(Int32Array::from(vec![20])), + ], + )?; + + // Key=1 matches two outer rows. Key=2 keeps the inner input non-exhausted. + let inner_batch = RecordBatch::try_new( + Arc::clone(&right_schema), + vec![ + Arc::new(Int32Array::from(vec![100, 200])), + Arc::new(Int32Array::from(vec![1, 2])), + Arc::new(Int32Array::from(vec![50, 60])), + ], + )?; + + let on_outer: Vec = vec![Arc::new(Column::new("b1", 1))]; + let on_inner: Vec = vec![Arc::new(Column::new("b1", 1))]; + + for (join_type, expected_a1) in [(LeftSemi, vec![1, 2]), (LeftAnti, vec![0])] { + let outer: SendableRecordBatchStream = Box::pin(PendingStream::new( + vec![outer_batch1.clone(), outer_batch2.clone()], + vec![false, true], // Pending before 2nd outer batch + )); + let inner: SendableRecordBatchStream = + Box::pin(PendingStream::new(vec![inner_batch.clone()], vec![false])); + + let metrics = ExecutionPlanMetricsSet::new(); + let inner_schema = inner.schema(); + let (reservation, spill_manager, runtime_env) = + test_stream_resources(inner_schema, &metrics); + let stream = BitwiseSortMergeJoinStream::try_new( + Arc::clone(&left_schema), + vec![SortOptions::default()], + NullEquality::NullEqualsNothing, + outer, + inner, + on_outer.clone(), + on_inner.clone(), + None, // no filter + join_type, + 8192, + 0, + &metrics, + reservation, + spill_manager, + runtime_env, + )?; + + let batches = collect_stream(stream).await?; + let actual_a1 = batches + .iter() + .flat_map(|batch| { + let values = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + (0..batch.num_rows()).map(|row| values.value(row)) + }) + .collect::>(); + assert_eq!(actual_a1, expected_a1, "{join_type:?}"); + } + Ok(()) +} + +/// Same as the no-filter boundary case, with a filter: the outer key group +/// spans batches and the outer input returns Pending at the boundary. /// /// Setup: /// - Outer: 2 single-row batches, both key=1, c1=[10, 20] @@ -4532,6 +4984,21 @@ async fn bitwise_spill_with_filter() -> Result<()> { metrics.spilled_rows().unwrap() > 0, "expected spilled_rows > 0 for {join_type:?}, batch_size={batch_size}" ); + let join_time = metrics + .sum_by_name("join_time") + .map(|m| m.as_usize()) + .unwrap_or(0); + assert!( + join_time > 0, + "expected join_time > 0 for {join_type:?}, batch_size={batch_size}" + ); + let output_rows = metrics.output_rows().unwrap_or(0); + let collected_rows: usize = spilled_result.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + output_rows, collected_rows, + "output_rows metric should match collected rows for \ + {join_type:?}, batch_size={batch_size}" + ); // Run without spilling and compare results let task_ctx_no_spill = Arc::new( @@ -4566,22 +5033,90 @@ async fn bitwise_spill_with_filter() -> Result<()> { Ok(()) } -/// Reproduces a bug where `resume_boundary` for the Filtered pending case -/// only checks `inner_key_buffer.is_empty()` but ignores `inner_key_spill`. -/// After spilling, the in-memory buffer is cleared while the spill file -/// holds the data. If the outer key group spans a batch boundary, the -/// second outer batch's rows are never evaluated against the inner group. +/// A single inner key group spanning several inner batches can spill more +/// than once under memory pressure. Every spilled slice must still be +/// evaluated against the outer rows — an earlier spill file must not be +/// dropped when a later slice of the same group spills. +#[tokio::test] +async fn bitwise_multi_spill_inner_key_group() -> Result<()> { + // Outer: one row with key 1, c1 = 5. + let left = build_table(("a1", &vec![1]), ("b1", &vec![1]), ("c1", &vec![5])); + + // Inner: one key group (b2 = 1) spanning two batches. Only the first + // batch satisfies the filter c1 < c2 (5 < 10); the second (5 < 0) does + // not, so dropping the first spilled slice flips the semi-join result. + let right_batches = vec![ + build_table_i32(("a2", &vec![10]), ("b2", &vec![1]), ("c2", &vec![10])), + build_table_i32(("a2", &vec![20]), ("b2", &vec![1]), ("c2", &vec![0])), + ]; + let right = build_table_from_batches(right_batches); + + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + let filter = build_c1_lt_c2_filter(left.schema().as_ref(), right.schema().as_ref()); + + // 100-byte pool: every buffered slice fails its reservation, so each + // inner batch of the key group spills separately. + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder( + DiskManagerBuilder::default().with_mode(DiskManagerMode::OsTmpDirectory), + ) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(SessionConfig::default().with_batch_size(1)) + .with_runtime(runtime), + ); + + let join = SortMergeJoinExec::try_new( + left, + right, + on, + Some(filter), + LeftSemi, + sort_options, + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + + let output_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + output_rows, 1, + "left row must match the group's first (spilled) inner slice", + ); + + let metrics = join.metrics().expect("must have metrics"); + assert_eq!( + metrics.spill_count(), + Some(1), + "all overflows of one key group must share a single spill file", + ); + assert_eq!( + metrics.spilled_rows(), + Some(2), + "both inner slices of the group must be spilled", + ); + Ok(()) +} + +/// Once the inner key group has spilled, an outer key group spanning a batch +/// boundary must still be evaluated against the spilled inner rows — the +/// second outer batch's rows must not be treated as having no inner group to +/// match against. /// /// Setup: /// - Outer: 2 single-row batches, both key=1, c1=[10, 10] /// - Inner: 1 batch with many rows all key=1 (enough to trigger spill) /// - Filter: c1 == c2 (matches when c2=10) /// - Memory limit: tiny (100 bytes) to force spilling -/// - Pending before 2nd outer batch to trigger boundary re-entry +/// - Pending before 2nd outer batch, while the key group is still open /// /// Expected: both outer rows match (semi=2 rows, anti=0 rows) -/// Bug: second outer row is skipped because resume_boundary sees empty -/// inner_key_buffer and skips re-evaluation. #[tokio::test] async fn spill_filtered_boundary_loses_outer_rows() -> Result<()> { let left_schema = Arc::new(Schema::new(vec![ @@ -4947,3 +5482,302 @@ async fn spill_read_back_single_source() -> Result<()> { Ok(()) } + +/// Small chunk size so even tiny test spill files are split into several +/// pieces, forcing multiple genuine suspend/resume cycles instead of one. +const PENDING_CHUNK_SIZE: usize = 16; + +/// Splits real spill bytes into fixed-size chunks and yields `Poll::Pending` +/// before every chunk +struct PendingChunkedStream { + chunks: VecDeque, + yield_pending: bool, +} + +impl PendingChunkedStream { + fn new(bytes: Bytes) -> Self { + let mut chunks = VecDeque::new(); + if bytes.is_empty() { + chunks.push_back(bytes); + } else { + let mut remaining = bytes; + while !remaining.is_empty() { + let take = PENDING_CHUNK_SIZE.min(remaining.len()); + chunks.push_back(remaining.split_to(take)); + } + } + Self { + chunks, + yield_pending: true, + } + } +} + +impl Stream for PendingChunkedStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + cx: &mut Context<'_>, + ) -> Poll> { + if self.yield_pending { + self.yield_pending = false; + cx.waker().wake_by_ref(); + return Poll::Pending; + } + // Pending before every subsequent chunk as well. + self.yield_pending = true; + match self.chunks.pop_front() { + Some(chunk) => Poll::Ready(Some(Ok(chunk))), + None => Poll::Ready(None), + } + } +} + +/// A `SpillFile` that delegates everything to a real local spill file, +/// except `read_stream`, which is forced through `PendingChunkedStream`. +struct PendingSpillFile { + inner: Arc, +} + +impl SpillFile for PendingSpillFile { + fn path(&self) -> Option<&std::path::Path> { + self.inner.path() + } + + fn size(&self) -> Option { + self.inner.size() + } + + fn read_stream(&self) -> Result> + Send>>> { + let path = self + .inner + .path() + .expect("PendingSpillFile only wraps local files") + .to_owned(); + + let stream = futures::stream::once(async move { + tokio::fs::read(&path) + .await + .map(Bytes::from) + .map_err(datafusion_common::DataFusionError::IoError) + }) + .flat_map( + |read_result| -> Pin> + Send>> { + match read_result { + Ok(bytes) => Box::pin(PendingChunkedStream::new(bytes)), + Err(e) => Box::pin(futures::stream::once(async move { Err(e) })), + } + }, + ); + + Ok(Box::pin(stream)) + } + + fn open_writer(&self) -> Result> { + self.inner.open_writer() + } +} + +/// Wraps the default `OsTmpDirectory` factory so every spill file it +/// creates is a [`PendingSpillFile`]. +struct PendingTempFileFactory { + inner: Arc, +} + +impl TempFileFactory for PendingTempFileFactory { + fn create_temp_file(&self, description: &str) -> Result> { + Ok(Arc::new(PendingSpillFile { + inner: self.inner.create_tmp_file(description)?, + })) + } +} + +fn pending_disk_manager_builder() -> DiskManagerBuilder { + let inner = Arc::new( + DiskManagerBuilder::default() + .with_mode(DiskManagerMode::OsTmpDirectory) + .build() + .unwrap(), + ); + DiskManagerBuilder::default().with_mode(DiskManagerMode::Custom(Arc::new( + PendingTempFileFactory { inner }, + ))) +} + +/// Materializing-side (Inner/Left/Right/Full) coverage: identical to +/// `overallocation_multi_batch_spill`, but every spill read goes through +/// `PendingSpillFile`, so `poll_spilled_batches` must actually hit and +/// recover from `Poll::Pending` mid-read. +#[tokio::test] +async fn materializing_spill_pending_stream() -> Result<()> { + let left_batch_1 = build_table_i32( + ("a1", &vec![0, 1]), + ("b1", &vec![1, 1]), + ("c1", &vec![4, 5]), + ); + let left_batch_2 = build_table_i32( + ("a1", &vec![2, 3]), + ("b1", &vec![1, 1]), + ("c1", &vec![6, 7]), + ); + let right_batch_1 = build_table_i32( + ("a2", &vec![0, 10]), + ("b2", &vec![1, 1]), + ("c2", &vec![50, 60]), + ); + let right_batch_2 = build_table_i32( + ("a2", &vec![20, 30]), + ("b2", &vec![1, 1]), + ("c2", &vec![70, 80]), + ); + let left = build_table_from_batches(vec![left_batch_1, left_batch_2]); + let right = build_table_from_batches(vec![right_batch_1, right_batch_2]); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b2", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(500, 1.0) + .with_disk_manager_builder(pending_disk_manager_builder()) + .build_arc()?; + + for join_type in [Inner, Left, Right, Full] { + let task_ctx = + Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let join = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "expected spill_count > 0 for {join_type:?}" + ); + + // Compare against a no-spill run to make sure waiting on the + // spill reads didn't corrupt or drop any data. + let task_ctx_no_spill = Arc::new(TaskContext::default()); + let join_no_spill = join_with_options( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + assert_eq!( + spilled_result, no_spill_result, + "Pending-forced spill read produced different results for {join_type:?}" + ); + } + + Ok(()) +} + +/// Bitwise-side (Semi/Anti) coverage: identical to `bitwise_spill_with_filter`, +/// but every spill read goes through `PendingSpillFile`, so reading the +/// spilled inner rows back must actually hit and recover from `Poll::Pending` +/// mid-read. +#[tokio::test] +async fn bitwise_spill_pending_stream() -> Result<()> { + let left = build_table( + ("a1", &vec![1, 2, 3, 4, 5, 6]), + ("b1", &vec![1, 2, 3, 4, 5, 6]), + ("c1", &vec![4, 5, 6, 7, 8, 9]), + ); + let right = build_table( + ("a2", &vec![10, 20, 30, 40, 50]), + ("b1", &vec![1, 3, 4, 6, 8]), + ("c2", &vec![50, 60, 70, 80, 90]), + ); + let on = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?) as _, + Arc::new(Column::new_with_schema("b1", &right.schema())?) as _, + )]; + let sort_options = vec![SortOptions::default(); on.len()]; + + // c1 < c2 is always true for matching keys — same filter as + // bitwise_spill_with_filter, so the inner key group is buffered + // (and spilled) rather than short-circuited. + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("c1", 0)), + Operator::Lt, + Arc::new(Column::new("c2", 1)), + )), + vec![ + ColumnIndex { + index: 2, + side: JoinSide::Left, + }, + ColumnIndex { + index: 2, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![ + Field::new("c1", DataType::Int32, false), + Field::new("c2", DataType::Int32, false), + ])), + ); + + let runtime = RuntimeEnvBuilder::new() + .with_memory_limit(100, 1.0) + .with_disk_manager_builder(pending_disk_manager_builder()) + .build_arc()?; + + for join_type in [LeftSemi, LeftAnti, RightSemi, RightAnti] { + let task_ctx = + Arc::new(TaskContext::default().with_runtime(Arc::clone(&runtime))); + let join = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Some(filter.clone()), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join.execute(0, task_ctx)?; + let spilled_result = common::collect(stream).await.unwrap(); + + let metrics = join.metrics().unwrap(); + assert!( + metrics.spill_count().unwrap() > 0, + "expected spill_count > 0 for {join_type:?}" + ); + + let task_ctx_no_spill = Arc::new(TaskContext::default()); + let join_no_spill = SortMergeJoinExec::try_new( + Arc::clone(&left), + Arc::clone(&right), + on.clone(), + Some(filter.clone()), + join_type, + sort_options.clone(), + NullEquality::NullEqualsNothing, + )?; + let stream = join_no_spill.execute(0, task_ctx_no_spill)?; + let no_spill_result = common::collect(stream).await.unwrap(); + + assert_eq!( + spilled_result, no_spill_result, + "Pending-forced spill read produced different results for {join_type:?}" + ); + } + + Ok(()) +} diff --git a/datafusion/physical-plan/src/joins/stream_join_utils.rs b/datafusion/physical-plan/src/joins/stream_join_utils.rs index 571c199abb448..05a56d241102e 100644 --- a/datafusion/physical-plan/src/joins/stream_join_utils.rs +++ b/datafusion/physical-plan/src/joins/stream_join_utils.rs @@ -37,6 +37,7 @@ use arrow::array::{ ArrowPrimitiveType, BooleanArray, BooleanBufferBuilder, NativeAdapter, PrimitiveArray, RecordBatch, }; +use arrow::buffer::NullBuffer; use arrow::compute::concat_batches; use arrow::datatypes::{ArrowNativeType, Schema, SchemaRef}; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; @@ -80,6 +81,7 @@ impl JoinHashMapType for PruningJoinHashMap { fn get_matched_indices_with_limit_offset( &self, hash_values: &[u64], + valid_keys: Option<&NullBuffer>, limit: usize, offset: MapOffset, input_indices: &mut Vec, @@ -91,6 +93,7 @@ impl JoinHashMapType for PruningJoinHashMap { &self.map, &next, hash_values, + valid_keys, limit, offset, input_indices, diff --git a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs index 34af88ea4027b..0c6e84b36cc55 100644 --- a/datafusion/physical-plan/src/joins/symmetric_hash_join.rs +++ b/datafusion/physical-plan/src/joins/symmetric_hash_join.rs @@ -31,7 +31,6 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::vec; -use crate::check_if_same_properties; use crate::common::SharedMemoryReservation; use crate::execution_plan::{boundedness_from_children, emission_type_from_children}; use crate::joins::stream_join_utils::{ @@ -44,16 +43,17 @@ use crate::joins::utils::{ BatchSplitter, BatchTransformer, ColumnIndex, JoinFilter, JoinHashMapType, JoinOn, JoinOnRef, NoopBatchTransformer, StatefulStreamResult, apply_join_filter_to_indices, build_batch_from_indices, build_join_schema, check_join_is_valid, equal_rows_arr, - symmetric_join_output_partitioning, update_hash, + matchable_join_keys, symmetric_join_output_partitioning, update_hash, }; use crate::projection::{ - ProjectionExec, join_allows_pushdown, join_table_borders, new_join_children, - physical_to_column_exprs, update_join_filter, update_join_on, + JoinData, ProjectionExec, try_pushdown_through_join_with_column_indices, }; use crate::stream::EmptyRecordBatchStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - PlanProperties, RecordBatchStream, SendableRecordBatchStream, + InputDistributionRequirements, PlanProperties, RecordBatchStream, + SendableRecordBatchStream, joins::StreamJoinPartitionMode, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; @@ -361,20 +361,6 @@ impl SymmetricHashJoinExec { } Ok(false) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - let left = children.swap_remove(0); - let right = children.swap_remove(0); - Self { - left, - right, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for SymmetricHashJoinExec { @@ -427,6 +413,10 @@ impl ExecutionPlan for SymmetricHashJoinExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { match self.mode { StreamJoinPartitionMode::Partitioned => { let (left_expr, right_expr) = self @@ -434,13 +424,16 @@ impl ExecutionPlan for SymmetricHashJoinExec { .iter() .map(|(l, r)| (Arc::clone(l) as _, Arc::clone(r) as _)) .unzip(); - vec![ - Distribution::HashPartitioned(left_expr), - Distribution::HashPartitioned(right_expr), - ] + InputDistributionRequirements::co_partitioned(vec![ + Distribution::KeyPartitioned(left_expr), + Distribution::KeyPartitioned(right_expr), + ]) } StreamJoinPartitionMode::SinglePartition => { - vec![Distribution::SinglePartition, Distribution::SinglePartition] + InputDistributionRequirements::new(vec![ + Distribution::SinglePartition, + Distribution::SinglePartition, + ]) } } } @@ -462,37 +455,64 @@ impl ExecutionPlan for SymmetricHashJoinExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn crate::PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to join keys from both sides - let mut tnr = TreeNodeRecursion::Continue; - for (left, right) in &self.on { - tnr = tnr.visit_sibling(|| f(left.as_ref()))?; - tnr = tnr.visit_sibling(|| f(right.as_ref()))?; - } - // Apply to join filter expressions if present - if let Some(filter) = &self.filter { - tnr = tnr.visit_sibling(|| f(filter.expression().as_ref()))?; + let join_keys = self.on.iter().flat_map(|(left, right)| [left, right]); + let filter = self.filter.iter().map(|filter| filter.expression()); + crate::apply_expression_roots(join_keys.chain(filter), f) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => { + let left = children.swap_remove(0); + let right = children.swap_remove(0); + Ok(Arc::new(Self { + left, + right, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })) + } + ChildrenPropertiesMode::Recompute => { + Ok(Arc::new(SymmetricHashJoinExec::try_new( + Arc::clone(&children[0]), + Arc::clone(&children[1]), + self.on.clone(), + self.filter.clone(), + &self.join_type, + self.null_equality, + self.left_sort_exprs.clone(), + self.right_sort_exprs.clone(), + self.mode, + )?)) + } } - Ok(tnr) } fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(SymmetricHashJoinExec::try_new( - Arc::clone(&children[0]), - Arc::clone(&children[1]), - self.on.clone(), - self.filter.clone(), - &self.join_type, - self.null_equality, - self.left_sort_exprs.clone(), - self.right_sort_exprs.clone(), - self.mode, - )?)) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn metrics(&self) -> Option { @@ -607,69 +627,305 @@ impl ExecutionPlan for SymmetricHashJoinExec { &self, projection: &ProjectionExec, ) -> Result>> { - // Convert projected PhysicalExpr's to columns. If not possible, we cannot proceed. - let Some(projection_as_columns) = physical_to_column_exprs(projection.expr()) - else { - return Ok(None); - }; - - let (far_right_left_col_ind, far_left_right_col_ind) = join_table_borders( - self.left().schema().fields().len(), - &projection_as_columns, - ); - - if !join_allows_pushdown( - &projection_as_columns, - &self.schema(), - far_right_left_col_ind, - far_left_right_col_ind, - ) { - return Ok(None); + let schema = self.schema(); + if let Some(JoinData { + projected_left_child, + projected_right_child, + join_filter, + join_on, + }) = try_pushdown_through_join_with_column_indices( + projection, + self.left(), + self.right(), + self.on(), + &schema, + self.filter(), + self.column_indices.as_slice(), + )? { + SymmetricHashJoinExec::try_new( + Arc::new(projected_left_child), + Arc::new(projected_right_child), + join_on, + join_filter, + self.join_type(), + self.null_equality(), + self.right().output_ordering().cloned(), + self.left().output_ordering().cloned(), + self.partition_mode(), + ) + .map(|e| Some(Arc::new(e) as _)) + } else { + Ok(None) } + } - let Some(new_on) = update_join_on( - &projection_as_columns[0..=far_right_left_col_ind as _], - &projection_as_columns[far_left_right_col_ind as _..], - self.on(), - self.left().schema().fields().len(), - ) else { - return Ok(None); + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let left = ctx.encode_child(self.left())?; + let right = ctx.encode_child(self.right())?; + let on = self + .on() + .iter() + .map(|(left, right)| { + Ok(protobuf::JoinOn { + left: Some(ctx.encode_expr(left)?), + right: Some(ctx.encode_expr(right)?), + }) + }) + .collect::>>()?; + + let join_type = match self.join_type() { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, }; - - let new_filter = if let Some(filter) = self.filter() { - match update_join_filter( - &projection_as_columns[0..=far_right_left_col_ind as _], - &projection_as_columns[far_left_right_col_ind as _..], - filter, - self.left().schema().fields().len(), - ) { - Some(updated_filter) => Some(updated_filter), - None => return Ok(None), + let null_equality = match self.null_equality() { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + }; + let partition_mode = match self.partition_mode() { + StreamJoinPartitionMode::SinglePartition => { + protobuf::StreamPartitionMode::SinglePartition + } + StreamJoinPartitionMode::Partitioned => { + protobuf::StreamPartitionMode::PartitionedExec } - } else { - None }; + let filter = self + .filter() + .map(|filter| -> Result { + let expression = ctx.encode_expr(filter.expression())?; + let column_indices = filter + .column_indices() + .iter() + .map(|column_index| { + let side = match column_index.side { + JoinSide::Left => protobuf::JoinSide::LeftSide, + JoinSide::Right => protobuf::JoinSide::RightSide, + JoinSide::None => protobuf::JoinSide::None, + }; + protobuf::ColumnIndex { + index: column_index.index as u32, + side: side.into(), + } + }) + .collect(); + Ok(protobuf::JoinFilter { + expression: Some(expression), + column_indices, + schema: Some(filter.schema().as_ref().try_into()?), + }) + }) + .transpose()?; + let expr_ctx = ctx.expr_ctx(); + let left_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto( + self.left_sort_exprs(), + &expr_ctx, + )?; + let right_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto( + self.right_sort_exprs(), + &expr_ctx, + )?; - let (new_left, new_right) = new_join_children( - &projection_as_columns, - far_right_left_col_ind, - far_left_right_col_ind, - self.left(), - self.right(), + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin( + Box::new(protobuf::SymmetricHashJoinExecNode { + left: Some(Box::new(left)), + right: Some(Box::new(right)), + on, + join_type: join_type.into(), + partition_mode: partition_mode.into(), + null_equality: null_equality.into(), + filter, + left_sort_exprs, + right_sort_exprs, + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SymmetricHashJoinExec { + /// Reconstruct a [`SymmetricHashJoinExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_common::internal_datafusion_err; + use datafusion_proto_models::protobuf; + + let sym_join = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SymmetricHashJoin, + "SymmetricHashJoinExec", + ); + let left = ctx.decode_required_child( + sym_join.left.as_deref(), + "SymmetricHashJoinExec", + "left", + )?; + let right = ctx.decode_required_child( + sym_join.right.as_deref(), + "SymmetricHashJoinExec", + "right", )?; + let left_schema = left.schema(); + let right_schema = right.schema(); + let on = sym_join + .on + .iter() + .map(|columns| { + let left = ctx.decode_required_expr( + columns.left.as_ref(), + left_schema.as_ref(), + "SymmetricHashJoinExec", + "on.left", + )?; + let right = ctx.decode_required_expr( + columns.right.as_ref(), + right_schema.as_ref(), + "SymmetricHashJoinExec", + "on.right", + )?; + Ok((left, right)) + }) + .collect::>()?; - SymmetricHashJoinExec::try_new( - Arc::new(new_left), - Arc::new(new_right), - new_on, - new_filter, - self.join_type(), - self.null_equality(), - self.right().output_ordering().cloned(), - self.left().output_ordering().cloned(), - self.partition_mode(), + let join_type = + match protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown JoinType {}", + sym_join.join_type + ) + })? { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + }; + let null_equality = match protobuf::NullEquality::try_from(sym_join.null_equality) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown NullEquality {}", + sym_join.null_equality + ) + })? { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + }; + let partition_mode = + match protobuf::StreamPartitionMode::try_from(sym_join.partition_mode) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown StreamPartitionMode {}", + sym_join.partition_mode + ) + })? { + protobuf::StreamPartitionMode::SinglePartition => { + StreamJoinPartitionMode::SinglePartition + } + protobuf::StreamPartitionMode::PartitionedExec => { + StreamJoinPartitionMode::Partitioned + } + }; + let filter = sym_join + .filter + .as_ref() + .map(|filter| -> Result { + let schema: Schema = filter + .schema + .as_ref() + .ok_or_else(|| { + internal_datafusion_err!( + "SymmetricHashJoinExec: JoinFilter missing schema" + ) + })? + .try_into()?; + let expression = ctx.decode_required_expr( + filter.expression.as_ref(), + &schema, + "SymmetricHashJoinExec", + "filter.expression", + )?; + let column_indices = filter + .column_indices + .iter() + .map(|column_index| { + let side = protobuf::JoinSide::try_from(column_index.side) + .map_err(|_| { + internal_datafusion_err!( + "SymmetricHashJoinExec: unknown JoinSide {}", + column_index.side + ) + })?; + let side = match side { + protobuf::JoinSide::LeftSide => JoinSide::Left, + protobuf::JoinSide::RightSide => JoinSide::Right, + protobuf::JoinSide::None => JoinSide::None, + }; + Ok(ColumnIndex { + index: column_index.index as usize, + side, + }) + }) + .collect::>>()?; + Ok(JoinFilter::new( + expression, + column_indices, + Arc::new(schema), + )) + }) + .transpose()?; + let left_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto( + &sym_join.left_sort_exprs, + &ctx.expr_ctx(left_schema.as_ref()), + )?; + let right_sort_exprs = + datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto( + &sym_join.right_sort_exprs, + &ctx.expr_ctx(right_schema.as_ref()), + )?; + + Self::try_new( + left, + right, + on, + filter, + &join_type, + null_equality, + left_sort_exprs, + right_sort_exprs, + partition_mode, ) - .map(|e| Some(Arc::new(e) as _)) + .map(|exec| Arc::new(exec) as _) } } @@ -1131,8 +1387,20 @@ fn lookup_join_hashmap( // (5,1) // // With this approach, the lexicographic order on both the probe side and the build side is preserved. + // + // Probe rows whose key contains a NULL cannot match any build row and are + // skipped without a map lookup. + let valid_keys = matchable_join_keys(&keys_values, null_equality); let (mut matched_probe, mut matched_build) = build_hashmap.get_matched_indices( - Box::new(hash_values.iter().enumerate().rev()), + Box::new( + hash_values + .iter() + .enumerate() + .filter(|(i, _)| { + valid_keys.as_ref().is_none_or(|valid| valid.is_valid(*i)) + }) + .rev(), + ), deleted_offset, ); @@ -1209,6 +1477,7 @@ impl OneSideHashJoiner { /// /// * `batch` - The incoming [RecordBatch] to be merged with the internal input buffer /// * `random_state` - The random state used to hash values + /// * `null_equality` - Null semantics to use /// /// # Returns /// @@ -1217,6 +1486,7 @@ impl OneSideHashJoiner { &mut self, batch: &RecordBatch, random_state: &RandomState, + null_equality: NullEquality, ) -> Result<()> { // Merge the incoming batch with the existing input buffer: self.input_buffer = concat_batches(&batch.schema(), [&self.input_buffer, batch])?; @@ -1233,6 +1503,7 @@ impl OneSideHashJoiner { &mut self.hashes_buffer, self.deleted_offset, false, + null_equality, )?; Ok(()) } @@ -1706,7 +1977,11 @@ impl SymmetricHashJoinStream { probe_side_metrics.input_batches.add(1); probe_side_metrics.input_rows.add(probe_batch.num_rows()); // Update the internal state of the hash joiner for the build side: - probe_hash_joiner.update_internal_state(probe_batch, &self.random_state)?; + probe_hash_joiner.update_internal_state( + probe_batch, + &self.random_state, + self.null_equality, + )?; // Join the two sides: let equal_result = join_with_probe_batch( build_hash_joiner, diff --git a/datafusion/physical-plan/src/joins/utils.rs b/datafusion/physical-plan/src/joins/utils.rs index b4aa295562b67..20467a7ec5e33 100644 --- a/datafusion/physical-plan/src/joins/utils.rs +++ b/datafusion/physical-plan/src/joins/utils.rs @@ -33,7 +33,8 @@ use crate::metrics::{ }; use crate::projection::{ProjectionExec, ProjectionExpr}; use crate::{ - ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, Statistics, + ColumnStatistics, ExecutionPlan, ExecutionPlanProperties, Partitioning, + RangePartitioning, Statistics, }; // compatibility pub use super::join_filter::JoinFilter; @@ -53,23 +54,21 @@ use arrow::array::{ TimestampNanosecondArray, TimestampSecondArray, UInt8Array, UInt16Array, }; use arrow::buffer::{BooleanBuffer, NullBuffer}; -use arrow::compute::kernels::cmp::eq; -use arrow::compute::{self, FilterBuilder, and, take}; +use arrow::compute::{self, take}; use arrow::datatypes::{ ArrowNativeType, Field, Schema, SchemaBuilder, UInt32Type, UInt64Type, }; -use arrow_ord::cmp::not_distinct; use arrow_ord::ord::{DynComparator, make_comparator}; -use arrow_schema::{ArrowError, DataType, SortOptions, TimeUnit}; +use arrow_schema::{DataType, SortOptions, TimeUnit}; use datafusion_common::cast::as_boolean_array; use datafusion_common::hash_utils::RandomState; use datafusion_common::hash_utils::create_hashes; use datafusion_common::stats::Precision; +use datafusion_common::utils::normalize_float_zero; use datafusion_common::{ DataFusionError, JoinSide, JoinType, NullEquality, Result, SharedResult, - not_impl_err, plan_err, + internal_datafusion_err, not_impl_err, plan_err, }; -use datafusion_expr::Operator; use datafusion_expr::interval_arithmetic::Interval; use datafusion_physical_expr::expressions::Column; use datafusion_physical_expr::utils::collect_columns; @@ -78,7 +77,6 @@ use datafusion_physical_expr::{ add_offset_to_physical_sort_exprs, }; -use datafusion_physical_expr_common::datum::compare_op_for_nested; use datafusion_physical_expr_common::utils::evaluate_expressions_to_arrays; use futures::future::{BoxFuture, Shared}; use futures::{FutureExt, ready}; @@ -144,6 +142,21 @@ pub fn adjust_right_output_partitioning( .collect::>()?; Partitioning::Hash(new_exprs, *size) } + Partitioning::Range(range) => { + let ordering = add_offset_to_physical_sort_exprs( + range.ordering().iter().cloned(), + left_columns_len as _, + )?; + let ordering = LexOrdering::new(ordering).ok_or_else(|| { + internal_datafusion_err!( + "Offsetting range partitioning produced an empty ordering" + ) + })?; + Partitioning::Range(RangePartitioning::new( + ordering, + range.split_points().to_vec(), + )) + } result => result.clone(), }; Ok(result) @@ -410,6 +423,7 @@ impl Clone for OnceFut { #[derive(Clone, Debug, Default)] struct PartialJoinStatistics { pub num_rows: usize, + pub total_byte_size: Precision, pub column_statistics: Vec, } @@ -423,9 +437,11 @@ struct PartialJoinStatistics { /// column-level statistics (distinct counts, min/max values) of the join keys. /// - **Column statistics**: Combines column statistics from both inputs. For join types /// that preserve all columns (Inner, Left, Right, Full), statistics from both sides -/// are concatenated. For semi/anti joins, only the relevant side's statistics are kept. -/// - **Byte size**: Always returns `Precision::Absent` as join output size is difficult -/// to estimate without knowing the actual data. +/// are concatenated. For semi/anti joins, the preserved side's statistics are +/// normalized as subset estimates. +/// - **Byte size**: For semi/anti joins, sums normalized column byte-size estimates +/// when every output column has one. Other join types return `Precision::Absent` +/// because join output size is difficult to estimate without knowing the actual data. /// /// # The `on` Parameter /// @@ -439,24 +455,34 @@ struct PartialJoinStatistics { /// - Does not account for selectivity of arbitrary join filter expressions /// (e.g., `(t1.v1 + t2.v1) % 2 = 0`). Such filters, common in NestedLoopJoinExec, /// are not factored into the cardinality estimation. -/// - Column statistics for the output are simply combined from inputs without -/// adjusting for join selectivity (acknowledged in the code as needing -/// "filter selectivity analysis"). +/// - Column statistics for inner/outer joins are simply combined from inputs +/// without adjusting for join selectivity (acknowledged in the code as +/// needing "filter selectivity analysis"). pub(crate) fn estimate_join_statistics( left_stats: Statistics, right_stats: Statistics, on: &JoinOn, + null_equality: NullEquality, join_type: &JoinType, schema: &Schema, ) -> Result { - let join_stats = estimate_join_cardinality(join_type, left_stats, right_stats, on); - let (num_rows, column_statistics) = match join_stats { - Some(stats) => (Precision::Inexact(stats.num_rows), stats.column_statistics), - None => (Precision::Absent, Statistics::unknown_column(schema)), + let join_stats = + estimate_join_cardinality(join_type, left_stats, right_stats, on, null_equality); + let (num_rows, total_byte_size, column_statistics) = match join_stats { + Some(stats) => ( + Precision::Inexact(stats.num_rows), + stats.total_byte_size, + stats.column_statistics, + ), + None => ( + Precision::Absent, + Precision::Absent, + Statistics::unknown_column(schema), + ), }; Ok(Statistics { num_rows, - total_byte_size: Precision::Absent, + total_byte_size, column_statistics, }) } @@ -467,23 +493,24 @@ fn estimate_join_cardinality( left_stats: Statistics, right_stats: Statistics, on: &JoinOn, + null_equality: NullEquality, ) -> Option { - let (left_col_stats, right_col_stats) = on + let on_column_indices = on .iter() - .map(|(left, right)| { - match ( - left.downcast_ref::(), - right.downcast_ref::(), - ) { - (Some(left), Some(right)) => ( - left_stats.column_statistics[left.index()].clone(), - right_stats.column_statistics[right.index()].clone(), - ), - _ => ( - ColumnStatistics::new_unknown(), - ColumnStatistics::new_unknown(), - ), - } + .map(|(left, right)| equijoin_column_indices(left, right)) + .collect::>(); + + let (left_key_stats, right_key_stats) = on_column_indices + .iter() + .map(|indices| match indices { + Some((left_index, right_index)) => ( + left_stats.column_statistics[*left_index].clone(), + right_stats.column_statistics[*right_index].clone(), + ), + None => ( + ColumnStatistics::new_unknown(), + ColumnStatistics::new_unknown(), + ), }) .unzip::<_, _, Vec<_>, Vec<_>>(); @@ -493,12 +520,12 @@ fn estimate_join_cardinality( Statistics { num_rows: left_stats.num_rows, total_byte_size: Precision::Absent, - column_statistics: left_col_stats, + column_statistics: left_key_stats, }, Statistics { num_rows: right_stats.num_rows, total_byte_size: Precision::Absent, - column_statistics: right_col_stats, + column_statistics: right_key_stats, }, )?; @@ -519,6 +546,7 @@ fn estimate_join_cardinality( Some(PartialJoinStatistics { num_rows: *cardinality.get_value()?, + total_byte_size: Precision::Absent, // We don't do anything specific here, just combine the existing // statistics which might yield subpar results (although it is // true, esp regarding min/max). For a better estimation, we need @@ -538,42 +566,81 @@ fn estimate_join_cardinality( let is_left = matches!(join_type, JoinType::LeftSemi | JoinType::LeftAnti); let is_anti = matches!(join_type, JoinType::LeftAnti | JoinType::RightAnti); - let ((outer_stats, inner_stats), (outer_col_stats, inner_col_stats)) = - if is_left { - ( - (&left_stats, &right_stats), - (&left_col_stats, &right_col_stats), - ) - } else { - ( - (&right_stats, &left_stats), - (&right_col_stats, &left_col_stats), - ) - }; + let (outer_stats, inner_stats, outer_key_stats, inner_key_stats) = if is_left + { + (left_stats, right_stats, left_key_stats, right_key_stats) + } else { + (right_stats, left_stats, right_key_stats, left_key_stats) + }; let outer_rows = *outer_stats.num_rows.get_value()?; - let cardinality = - if estimate_disjoint_inputs(outer_stats, inner_stats).is_some() { - // Disjoint inputs: semi produces 0, anti keeps all rows. - if is_anti { outer_rows } else { 0 } + let outer_join_key_stats = Statistics { + num_rows: outer_stats.num_rows, + total_byte_size: Precision::Absent, + column_statistics: outer_key_stats.clone(), + }; + let inner_join_key_stats = Statistics { + num_rows: inner_stats.num_rows, + total_byte_size: Precision::Absent, + column_statistics: inner_key_stats.clone(), + }; + + let semi_cardinality = + if estimate_disjoint_inputs(&outer_join_key_stats, &inner_join_key_stats) + .is_some() + { + // If join keys are disjoint, no rows will match + Some(0) } else { - match estimate_semi_join_cardinality( + estimate_semi_join_cardinality( &outer_stats.num_rows, &inner_stats.num_rows, - outer_col_stats, - inner_col_stats, - ) { - Some(semi) if is_anti => outer_rows.saturating_sub(semi), - Some(semi) => semi, - None => outer_rows, - } + &outer_key_stats, + &inner_key_stats, + null_equality, + ) }; - let outer_stats = if is_left { left_stats } else { right_stats }; + // Semi joins keep the matching rows; anti joins keep the rest. When no + // estimate is available, conservatively assume all outer rows pass. + let cardinality = match (semi_cardinality, is_anti) { + (Some(semi), true) => outer_rows.saturating_sub(semi), + (Some(semi), false) => semi, + (None, _) => outer_rows, + }; + + // The outer side is the one whose columns a semi/anti join emits, so + // its statistics are the ones to normalize into the subset estimate. + let Statistics { + num_rows: preserved_num_rows, + column_statistics: preserved_column_statistics, + .. + } = outer_stats; + let preserved_join_key_indices = on_column_indices + .iter() + .filter_map(|&indices| { + indices.map( + |(left_index, right_index)| { + if is_left { left_index } else { right_index } + }, + ) + }) + .collect::>(); + let column_statistics = normalize_semi_anti_join_column_statistics( + preserved_column_statistics, + &preserved_num_rows, + cardinality, + &preserved_join_key_indices, + is_anti, + null_equality, + ); + let total_byte_size = + total_byte_size_from_column_statistics(&column_statistics); Some(PartialJoinStatistics { num_rows: cardinality, - column_statistics: outer_stats.column_statistics, + total_byte_size, + column_statistics, }) } @@ -583,6 +650,7 @@ fn estimate_join_cardinality( column_statistics.push(ColumnStatistics::new_unknown()); Some(PartialJoinStatistics { num_rows, + total_byte_size: Precision::Absent, column_statistics, }) } @@ -592,12 +660,132 @@ fn estimate_join_cardinality( column_statistics.push(ColumnStatistics::new_unknown()); Some(PartialJoinStatistics { num_rows, + total_byte_size: Precision::Absent, column_statistics, }) } } } +fn equijoin_column_indices( + left: &PhysicalExprRef, + right: &PhysicalExprRef, +) -> Option<(usize, usize)> { + Some(( + left.downcast_ref::()?.index(), + right.downcast_ref::()?.index(), + )) +} + +/// Adjusts the preserved input's column statistics to describe the subset of +/// rows a semi or anti join emits. Most values become estimates (marked +/// inexact) bounded by the smaller output row count: +/// +/// - `null_count` and `byte_size` are scaled by the output/input row ratio. +/// - `distinct_count` is capped at the number of non-null output rows. +/// - `sum_value` is dropped, since the input sum does not apply to the subset. +/// +/// Join-key columns are the exception for `null_count`: under regular SQL +/// equality, null keys never match, so a semi join keeps none of those rows and +/// an anti join keeps all of them. Under null-equal joins, null keys can match +/// and are treated like the rest of the subset. +fn normalize_semi_anti_join_column_statistics( + column_statistics: Vec, + input_num_rows: &Precision, + output_num_rows: usize, + join_key_indices: &[usize], + is_anti: bool, + null_equality: NullEquality, +) -> Vec { + let input_num_rows = input_num_rows.get_value().copied().unwrap_or(0); + + column_statistics + .into_iter() + .enumerate() + .map(|(idx, stats)| { + let mut stats = stats.to_inexact(); + stats.null_count = if join_key_indices.contains(&idx) { + normalize_semi_anti_join_key_null_count( + stats.null_count, + input_num_rows, + output_num_rows, + is_anti, + null_equality, + ) + } else { + scale_subset_count(stats.null_count, input_num_rows, output_num_rows) + .min(&Precision::Inexact(output_num_rows)) + }; + let max_distinct_count = stats + .null_count + .get_value() + .map(|null_count| output_num_rows.saturating_sub(*null_count)) + .unwrap_or(output_num_rows); + stats.distinct_count = stats + .distinct_count + .min(&Precision::Inexact(max_distinct_count)); + stats.byte_size = + scale_subset_count(stats.byte_size, input_num_rows, output_num_rows); + stats.sum_value = Precision::Absent; + stats + }) + .collect() +} + +fn normalize_semi_anti_join_key_null_count( + null_count: Precision, + input_num_rows: usize, + output_num_rows: usize, + is_anti: bool, + null_equality: NullEquality, +) -> Precision { + match (is_anti, null_equality) { + (false, NullEquality::NullEqualsNothing) => Precision::Exact(0), + (true, NullEquality::NullEqualsNothing) => null_count + .to_inexact() + .min(&Precision::Inexact(output_num_rows)), + (_, NullEquality::NullEqualsNull) => { + scale_subset_count(null_count, input_num_rows, output_num_rows) + .min(&Precision::Inexact(output_num_rows)) + } + } +} + +// Scale a column-level count to an estimated row subset. Rounding up keeps a +// small non-zero count from disappearing solely because the subset is small. +fn scale_subset_count( + count: Precision, + input_num_rows: usize, + output_num_rows: usize, +) -> Precision { + let scaled = match count { + Precision::Exact(count) | Precision::Inexact(count) => { + if input_num_rows == 0 { + 0 + } else { + (count as u128 * output_num_rows as u128).div_ceil(input_num_rows as u128) + as usize + } + } + Precision::Absent => return Precision::Absent, + }; + + Precision::Inexact(scaled) +} + +fn total_byte_size_from_column_statistics( + column_statistics: &[ColumnStatistics], +) -> Precision { + column_statistics + .iter() + .map(|stats| stats.byte_size.get_value().copied()) + .try_fold(0usize, |acc, byte_size| { + byte_size.map(|byte_size| acc.saturating_add(byte_size)) + }) + .map(Precision::Inexact) + .unwrap_or(Precision::Absent) +} + /// Estimate the inner join cardinality by using the basic building blocks of /// column-level statistics and the total row count. This is a very naive and /// a very conservative implementation that can quickly give up if there is not @@ -622,8 +810,15 @@ fn estimate_inner_join_cardinality( .. } = right_stats; - // The algorithm here is partly based on the non-histogram selectivity estimation - // from Spark's Catalyst optimizer. + if left_num_rows == Precision::Exact(0) || right_num_rows == Precision::Exact(0) { + return Some(Precision::Exact(0)); + } + if left_num_rows == Precision::Inexact(0) || right_num_rows == Precision::Inexact(0) { + return Some(Precision::Inexact(0)); + } + + // Follow Spark Catalyst's conservative NDV join estimate: for multi-key + // joins, use the most selective key instead of multiplying all key denominators. let mut join_selectivity = Precision::Absent; for (left_stat, right_stat) in left_column_statistics .iter() @@ -636,22 +831,33 @@ fn estimate_inner_join_cardinality( // Seems like there are a few implementations of this algorithm that implement // exponential decay for the selectivity (like Hive's Optiq Optimizer). Needs // further exploration. - join_selectivity = max_distinct; + join_selectivity = if join_selectivity.get_value().is_some() { + join_selectivity.max(&max_distinct) + } else { + max_distinct + }; } } // With the assumption that the smaller input's domain is generally represented in the bigger // input's domain, we can estimate the inner join's cardinality by taking the cartesian product // of the two inputs and normalizing it by the selectivity factor. - let left_num_rows = left_stats.num_rows.get_value()?; - let right_num_rows = right_stats.num_rows.get_value()?; + let left_num_rows = *left_stats.num_rows.get_value()?; + let right_num_rows = *right_stats.num_rows.get_value()?; + // Widen before multiplying so the intermediate Cartesian product does not + // overflow when the normalized cardinality is still representable as usize. + let cartesian_product = (left_num_rows as u128) * (right_num_rows as u128); + let normalized_cardinality = + |value: usize| usize::try_from(cartesian_product / value as u128); match join_selectivity { - Precision::Exact(value) if value > 0 => { - Some(Precision::Exact((left_num_rows * right_num_rows) / value)) - } - Precision::Inexact(value) if value > 0 => { - Some(Precision::Inexact((left_num_rows * right_num_rows) / value)) - } + Precision::Exact(value) if value > 0 => Some( + normalized_cardinality(value) + .map(Precision::Exact) + .unwrap_or(Precision::Inexact(usize::MAX)), + ), + Precision::Inexact(value) if value > 0 => Some(Precision::Inexact( + normalized_cardinality(value).unwrap_or(usize::MAX), + )), // Since we don't have any information about the selectivity (which is derived // from the number of distinct rows information) we can give up here for now. // And let other passes handle this (otherwise we would need to produce an @@ -721,8 +927,8 @@ fn estimate_disjoint_inputs( /// Under the uniformity assumption (each distinct value contributes /// equally to row counts), the surviving fraction of outer rows is: /// -/// Null rows cannot match, so each column's selectivity is further -/// reduced by the outer null fraction: +/// Under regular SQL equality, null rows cannot match, so each column's +/// selectivity is further reduced by the outer null fraction: /// /// ```text /// null_frac_i = outer_null_count_i / outer_rows @@ -739,7 +945,7 @@ fn estimate_disjoint_inputs( /// Anti join cardinality is derived as the complement: /// `outer_rows - semi_cardinality`. /// -/// Boundary cases: +/// With `NullEqualsNothing`, boundary cases are: /// * `inner_ndv >= outer_ndv` → selectivity = `1.0 - null_frac` /// * `null_frac = 1.0` → selectivity = 0.0 (no non-null rows can match) /// * Missing NDV statistics → returns `None` (fallback to `outer_rows`) @@ -752,8 +958,9 @@ fn estimate_disjoint_inputs( fn estimate_semi_join_cardinality( outer_num_rows: &Precision, inner_num_rows: &Precision, - outer_col_stats: &[ColumnStatistics], - inner_col_stats: &[ColumnStatistics], + outer_key_stats: &[ColumnStatistics], + inner_key_stats: &[ColumnStatistics], + null_equality: NullEquality, ) -> Option { let outer_rows = *outer_num_rows.get_value()?; if outer_rows == 0 { @@ -767,7 +974,7 @@ fn estimate_semi_join_cardinality( let mut selectivity = 1.0_f64; let mut has_selectivity_estimate = false; - for (outer_stat, inner_stat) in outer_col_stats.iter().zip(inner_col_stats.iter()) { + for (outer_stat, inner_stat) in outer_key_stats.iter().zip(inner_key_stats.iter()) { let outer_has_stats = outer_stat.distinct_count.get_value().is_some() || (outer_stat.min_value.get_value().is_some() && outer_stat.max_value.get_value().is_some()); @@ -784,11 +991,21 @@ fn estimate_semi_join_cardinality( if let (Some(&o), Some(&i)) = (outer_ndv.get_value(), inner_ndv.get_value()) && o > 0 { - let null_frac = outer_stat - .null_count - .get_value() - .map(|&nc| nc as f64 / outer_rows as f64) - .unwrap_or(0.0); + let null_frac = if null_equality == NullEquality::NullEqualsNothing { + outer_stat + .null_count + .get_value() + .map(|&nc| { + if nc > outer_rows { + 0.0 + } else { + nc as f64 / outer_rows as f64 + } + }) + .unwrap_or(0.0) + } else { + 0.0 + }; selectivity *= (o.min(i) as f64) / (o as f64) * (1.0 - null_frac); has_selectivity_estimate = true; } @@ -1168,7 +1385,9 @@ pub(crate) fn build_batch_from_indices( Ok(RecordBatch::try_new(Arc::new(schema.clone()), columns)?) } -/// Returns a new [RecordBatch] resulting of a join where the build/left side is empty. +/// Returns a new [RecordBatch] for a probe batch when no probe row can find a +/// match: the build-side map is empty, either because the build side has no +/// rows or because none of its rows has a matchable (non-NULL) join key. /// The resulting batch has [Schema] `schema`. pub(crate) fn build_batch_empty_build_side( schema: &Schema, @@ -1351,7 +1570,9 @@ pub(crate) fn append_right_indices( } } -/// Returns `range` indices which are not present in `input_indices` +/// Returns `range` indices which are not present in `input_indices`. +/// +/// `input_indices` must be sorted ascending and contain no nulls. pub(crate) fn get_anti_indices( range: Range, input_indices: &PrimitiveArray, @@ -1359,18 +1580,51 @@ pub(crate) fn get_anti_indices( where NativeAdapter: From<::Native>, { - let bitmap = build_range_bitmap(&range, input_indices); - let offset = range.start; + debug_assert_eq!( + input_indices.null_count(), + 0, + "get_anti_indices requires non-null input_indices" + ); + debug_assert!( + input_indices + .values() + .windows(2) + .all(|w| w[0].as_usize() <= w[1].as_usize()), + "get_anti_indices requires ascending input_indices" + ); - // get the anti index - (range) - .filter_map(|idx| { - (!bitmap.get_bit(idx - offset)).then_some(T::Native::from_usize(idx)) - }) - .collect() + let mut next_unmatched_idx = range.start; + let mut output: Vec = Vec::with_capacity(range.len()); + + for &v in input_indices.values() { + let idx = v.as_usize(); + + if idx < range.start { + continue; + } + if idx >= range.end { + break; + } + + if next_unmatched_idx < idx { + output.extend((next_unmatched_idx..idx).map(|idx| { + T::Native::from_usize(idx).expect("join index exceeds output index type") + })); + } + next_unmatched_idx = idx + 1; + } + + if next_unmatched_idx < range.end { + output.extend((next_unmatched_idx..range.end).map(|idx| { + T::Native::from_usize(idx).expect("join index exceeds output index type") + })); + } + PrimitiveArray::::new(output.into(), None) } -/// Returns intersection of `range` and `input_indices` omitting duplicates +/// Returns the intersection of `range` and `input_indices`, omitting duplicates. +/// +/// `input_indices` must be sorted ascending and contain no nulls. pub(crate) fn get_semi_indices( range: Range, input_indices: &PrimitiveArray, @@ -1378,14 +1632,38 @@ pub(crate) fn get_semi_indices( where NativeAdapter: From<::Native>, { - let bitmap = build_range_bitmap(&range, input_indices); - let offset = range.start; - // get the semi index - (range) - .filter_map(|idx| { - (bitmap.get_bit(idx - offset)).then_some(T::Native::from_usize(idx)) - }) - .collect() + debug_assert_eq!( + input_indices.null_count(), + 0, + "get_semi_indices requires non-null input_indices" + ); + debug_assert!( + input_indices + .values() + .windows(2) + .all(|w| w[0].as_usize() <= w[1].as_usize()), + "get_semi_indices requires ascending input_indices" + ); + + let mut prev_idx: Option = None; + let mut output = Vec::with_capacity(input_indices.len().min(range.len())); + + for &v in input_indices.values() { + let idx = v.as_usize(); + + if idx < range.start { + continue; + } + if idx >= range.end { + break; + } + + if prev_idx.replace(idx) != Some(idx) { + output.push(v); + } + } + + PrimitiveArray::::new(output.into(), None) } pub(crate) fn get_mark_indices( @@ -1532,9 +1810,8 @@ impl BuildProbeJoinMetrics { .with_category(MetricCategory::Rows) .counter("build_input_rows", partition); - let build_mem_used = MetricBuilder::new(metrics) - .with_category(MetricCategory::Bytes) - .gauge("build_mem_used", partition); + let build_mem_used = + MetricBuilder::new(metrics).peak_memory_usage("build_mem_used", partition); let input_batches = MetricBuilder::new(metrics) .with_category(MetricCategory::Rows) @@ -1840,6 +2117,9 @@ pub fn swap_join_projection( /// `fifo_hashmap` sets the order of iteration over `batch` rows while updating hashmap, /// which allows to keep either first (if set to true) or last (if set to false) row index /// as a chain head for rows with equal hash values. +/// +/// Under [`NullEquality::NullEqualsNothing`], rows with a NULL in any key +/// column can never match a probe row, so they are not inserted into the map. #[expect(clippy::too_many_arguments)] pub fn update_hash( on: &[PhysicalExprRef], @@ -1850,6 +2130,7 @@ pub fn update_hash( hashes_buffer: &mut [u64], deleted_offset: usize, fifo_hashmap: bool, + null_equality: NullEquality, ) -> Result<()> { // evaluate the keys let keys_values = evaluate_expressions_to_arrays(on, batch)?; @@ -1860,10 +2141,14 @@ pub fn update_hash( // For usual JoinHashmap, the implementation is void. hash_map.extend_zero(batch.num_rows()); + // Unmatchable NULL-key rows are filtered out below. + let valid_keys = matchable_join_keys(&keys_values, null_equality); + // Updating JoinHashMap from hash values iterator let hash_values_iter = hash_values .iter() .enumerate() + .filter(|(i, _)| valid_keys.as_ref().is_none_or(|nulls| nulls.is_valid(*i))) .map(|(i, val)| (i + offset, val)); if fifo_hashmap { @@ -1875,6 +2160,31 @@ pub fn update_hash( Ok(()) } +/// Returns the combined validity of the join key columns `join_key_arrays`: a row +/// is valid only if every key column is non-NULL at that row. +/// +/// Returns `None` when no rows need to be filtered: either every row has +/// fully non-NULL keys, or `null_equality` is +/// [`NullEquality::NullEqualsNull`], where NULL keys are matchable. +pub(crate) fn matchable_join_keys( + join_key_arrays: &[ArrayRef], + null_equality: NullEquality, +) -> Option { + match null_equality { + NullEquality::NullEqualsNothing => { + let logical_nulls: Vec<_> = join_key_arrays + .iter() + .map(|values| values.logical_nulls()) + .collect(); + NullBuffer::union_many(logical_nulls.iter().map(Option::as_ref)) + // An all-valid array can still have a validity buffer; return + // `None` in that case, since there is nothing to filter. + .filter(|nulls| nulls.null_count() > 0) + } + NullEquality::NullEqualsNull => None, + } +} + pub(super) fn equal_rows_arr( indices_left: &UInt64Array, indices_right: &UInt32Array, @@ -1882,58 +2192,146 @@ pub(super) fn equal_rows_arr( right_arrays: &[ArrayRef], null_equality: NullEquality, ) -> Result<(UInt64Array, UInt32Array)> { - let mut iter = left_arrays.iter().zip(right_arrays.iter()); + if indices_left.len() != indices_right.len() { + return Err(internal_datafusion_err!( + "Cannot compare join indices with different lengths: left={}, right={}", + indices_left.len(), + indices_right.len() + )); + } + + if left_arrays.len() != right_arrays.len() { + return Err(internal_datafusion_err!( + "Cannot compare join keys with different column counts: left={}, right={}", + left_arrays.len(), + right_arrays.len() + )); + } - let Some((first_left, first_right)) = iter.next() else { + if left_arrays.is_empty() { return Ok((Vec::::new().into(), Vec::::new().into())); - }; + } - let arr_left = take(first_left.as_ref(), indices_left, None)?; - let arr_right = take(first_right.as_ref(), indices_right, None)?; + // Fast path: single-column keys of a specialized type run a monomorphized + // equality loop, avoiding the per-pair boxed `DynComparator` dispatch and + // `Ordering` computation of the general `JoinKeyComparator` path. Falls + // through to the general path for multi-column keys and unspecialized + // types (e.g. floats, dictionaries, nested). + let single_col_fast_path = if left_arrays.len() == 1 { + equal_rows_single_col( + indices_left, + indices_right, + left_arrays[0].as_ref(), + right_arrays[0].as_ref(), + null_equality, + ) + } else { + None + }; + if let Some(res) = single_col_fast_path { + return Ok(res); + } - let mut equal: BooleanArray = eq_dyn_null(&arr_left, &arr_right, null_equality)?; + let sort_options = vec![SortOptions::default(); left_arrays.len()]; + let comparator = + JoinKeyComparator::new(left_arrays, right_arrays, &sort_options, null_equality)?; - // Use map and try_fold to iterate over the remaining pairs of arrays. - // In each iteration, take is used on the pair of arrays and their equality is determined. - // The results are then folded (combined) using the and function to get a final equality result. - equal = iter - .map(|(left, right)| { - let arr_left = take(left.as_ref(), indices_left, None)?; - let arr_right = take(right.as_ref(), indices_right, None)?; - eq_dyn_null(arr_left.as_ref(), arr_right.as_ref(), null_equality) - }) - .try_fold(equal, |acc, equal2| and(&acc, &equal2?))?; + let mut left_filtered = Vec::with_capacity(indices_left.len()); + let mut right_filtered = Vec::with_capacity(indices_right.len()); - let filter_builder = FilterBuilder::new(&equal).optimize().build(); + for (left, right) in indices_left.values().iter().zip(indices_right.values()) { + let left_idx = usize::try_from(*left).map_err(|_| { + internal_datafusion_err!("Join index {left} can not be represented as usize") + })?; + let right_idx = *right as usize; - let left_filtered = filter_builder.filter(indices_left)?; - let right_filtered = filter_builder.filter(indices_right)?; + if comparator.is_equal(left_idx, right_idx) { + left_filtered.push(*left); + right_filtered.push(*right); + } + } - Ok(( - downcast_array(left_filtered.as_ref()), - downcast_array(right_filtered.as_ref()), - )) + Ok((left_filtered.into(), right_filtered.into())) } -// version of eq_dyn supporting equality on null arrays -fn eq_dyn_null( +/// Specialized single-column equi-join key filtering. +/// +/// Dispatches once on the key column's type and runs a monomorphized equality +/// loop with typed value comparison. This avoids the per-pair boxed +/// `DynComparator` call and the three-way `Ordering` computation used by the +/// general [`JoinKeyComparator`] path, which dominates for high-fanout +/// single-column joins (e.g. long string keys with near-100% match rates). +/// +/// Returns `None` for types it does not specialize (including when the left and +/// right key types differ, handled by the failed downcast) so the caller falls +/// back to the general path. Floats are intentionally excluded so their `-0.0` / +/// `NaN` semantics stay on the exact same code path as before. +fn equal_rows_single_col( + indices_left: &UInt64Array, + indices_right: &UInt32Array, left: &dyn Array, right: &dyn Array, null_equality: NullEquality, -) -> Result { - // Nested datatypes cannot use the underlying not_distinct/eq function and must use a special - // implementation - // - if left.data_type().is_nested() { - let op = match null_equality { - NullEquality::NullEqualsNothing => Operator::Eq, - NullEquality::NullEqualsNull => Operator::IsNotDistinctFrom, - }; - return Ok(compare_op_for_nested(op, &left, &right)?); +) -> Option<(UInt64Array, UInt32Array)> { + let null_equals_null = matches!(null_equality, NullEquality::NullEqualsNull); + + macro_rules! eq_loop { + ($T:ty) => {{ + let l = left.as_any().downcast_ref::<$T>()?; + let r = right.as_any().downcast_ref::<$T>()?; + + let mut left_filtered = Vec::with_capacity(indices_left.len()); + let mut right_filtered = Vec::with_capacity(indices_right.len()); + + for (left_idx, right_idx) in + indices_left.values().iter().zip(indices_right.values()) + { + let i = *left_idx as usize; + let j = *right_idx as usize; + + let is_equal = match (l.is_null(i), r.is_null(j)) { + (false, false) => l.value(i) == r.value(j), + (true, true) => null_equals_null, + _ => false, + }; + + if is_equal { + left_filtered.push(*left_idx); + right_filtered.push(*right_idx); + } + } + + return Some((left_filtered.into(), right_filtered.into())); + }}; } - match null_equality { - NullEquality::NullEqualsNothing => eq(&left, &right), - NullEquality::NullEqualsNull => not_distinct(&left, &right), + + match left.data_type() { + DataType::Boolean => eq_loop!(BooleanArray), + DataType::Int8 => eq_loop!(Int8Array), + DataType::Int16 => eq_loop!(Int16Array), + DataType::Int32 => eq_loop!(Int32Array), + DataType::Int64 => eq_loop!(Int64Array), + DataType::UInt8 => eq_loop!(UInt8Array), + DataType::UInt16 => eq_loop!(UInt16Array), + DataType::UInt32 => eq_loop!(UInt32Array), + DataType::UInt64 => eq_loop!(UInt64Array), + DataType::Decimal128(..) => eq_loop!(Decimal128Array), + DataType::Binary => eq_loop!(BinaryArray), + DataType::LargeBinary => eq_loop!(LargeBinaryArray), + DataType::BinaryView => eq_loop!(BinaryViewArray), + DataType::FixedSizeBinary(_) => eq_loop!(FixedSizeBinaryArray), + DataType::Utf8 => eq_loop!(StringArray), + DataType::LargeUtf8 => eq_loop!(LargeStringArray), + DataType::Utf8View => eq_loop!(StringViewArray), + DataType::Date32 => eq_loop!(Date32Array), + DataType::Date64 => eq_loop!(Date64Array), + DataType::Timestamp(time_unit, _) => match time_unit { + TimeUnit::Second => eq_loop!(TimestampSecondArray), + TimeUnit::Millisecond => eq_loop!(TimestampMillisecondArray), + TimeUnit::Microsecond => eq_loop!(TimestampMicrosecondArray), + TimeUnit::Nanosecond => eq_loop!(TimestampNanosecondArray), + }, + _ => None, } } @@ -1979,7 +2377,16 @@ impl JoinKeyComparator { .zip(right_arrays.iter()) .zip(sort_options.iter()) .map(|((l, r), opts)| { - let inner = make_comparator(l.as_ref(), r.as_ref(), *opts)?; + // `make_comparator` uses IEEE 754 totalOrder for floats and + // treats `-0.0` / `+0.0` as distinct. Normalize float arrays + // so SMJ / piecewise-merge equi-keys honor SQL equality; + // no-op (Arc::clone) for non-floats and for float arrays + // that contain no `-0.0`. `normalize_float_zero` preserves + // null positions, so the original null masks below remain + // valid. + let l_norm = normalize_float_zero(l); + let r_norm = normalize_float_zero(r); + let inner = make_comparator(l_norm.as_ref(), r_norm.as_ref(), *opts)?; if null_equality == NullEquality::NullEqualsNothing { let ln = l.logical_nulls().filter(|n| n.null_count() > 0); let rn = r.logical_nulls().filter(|n| n.null_count() > 0); @@ -2144,11 +2551,164 @@ mod tests { use arrow::datatypes::{DataType, Fields}; use arrow::error::{ArrowError, Result as ArrowResult}; use datafusion_common::stats::Precision::{Absent, Exact, Inexact}; - use datafusion_common::{ScalarValue, arrow_datafusion_err, arrow_err}; + use datafusion_common::{ScalarValue, SplitPoint, arrow_datafusion_err, arrow_err}; use datafusion_physical_expr::PhysicalSortExpr; use rstest::rstest; + fn assert_u32_values(array: &UInt32Array, expected: &[u32]) { + assert_eq!(array.values().as_ref(), expected); + } + + #[test] + fn get_anti_indices_returns_unmatched_range_indices() { + let input = UInt32Array::from(vec![3, 5, 5]); + + let result = get_anti_indices(2..8, &input); + + assert_u32_values(&result, &[2, 4, 6, 7]); + } + + #[test] + fn get_anti_indices_ignores_out_of_range_indices() { + let input = UInt32Array::from(vec![0, 1, 3, 5, 8, 12]); + + let result = get_anti_indices(2..8, &input); + + assert_u32_values(&result, &[2, 4, 6, 7]); + } + + #[test] + fn update_hash_skips_null_keys_for_null_equals_nothing() -> Result<()> { + use crate::joins::join_hash_map::JoinHashMapU32; + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![ + Some(1), + None, + Some(2), + None, + Some(1), + ]))], + )?; + let on: Vec = vec![Arc::new(Column::new("a", 0))]; + let random_state = RandomState::with_seed(42); + let mut hashes_buffer = vec![0; batch.num_rows()]; + create_hashes([batch.column(0)], &random_state, &mut hashes_buffer)?; + + let matched_build_indices = + |map: &JoinHashMapU32, hashes_buffer: &[u64]| -> Vec { + let mut input_indices = vec![]; + let mut match_indices = vec![]; + map.get_matched_indices_with_limit_offset( + hashes_buffer, + None, + 8192, + (0, None), + &mut input_indices, + &mut match_indices, + ); + match_indices.sort_unstable(); + match_indices.dedup(); + match_indices + }; + + let mut map = JoinHashMapU32::with_capacity(batch.num_rows()); + update_hash( + &on, + &batch, + &mut map, + 0, + &random_state, + &mut hashes_buffer, + 0, + true, + NullEquality::NullEqualsNothing, + )?; + // NULL keys can never match under NullEqualsNothing, so they must not + // be inserted into the map. Assert row indices rather than map length: + // with forced hash collisions, multiple logical keys can share one + // hash table entry. + assert_eq!(matched_build_indices(&map, &hashes_buffer), vec![0, 2, 4]); + + let mut map = JoinHashMapU32::with_capacity(batch.num_rows()); + update_hash( + &on, + &batch, + &mut map, + 0, + &random_state, + &mut hashes_buffer, + 0, + true, + NullEquality::NullEqualsNull, + )?; + // Under NullEqualsNull, NULL keys can match, so the build-side NULL + // rows must be present in the map. + assert_eq!( + matched_build_indices(&map, &hashes_buffer), + vec![0, 1, 2, 3, 4] + ); + + Ok(()) + } + + #[test] + fn get_anti_indices_handles_dense_matches() { + let input = UInt32Array::from(vec![2, 3, 4, 5]); + + let result = get_anti_indices(2..6, &input); + + assert!(result.is_empty()); + } + + #[test] + fn get_anti_indices_handles_sparse_matches() { + let input = UInt32Array::from(vec![0, 8]); + + let result = get_anti_indices(2..6, &input); + + assert_u32_values(&result, &[2, 3, 4, 5]); + } + + #[test] + fn get_semi_indices_returns_distinct_matches_in_range() { + let input = UInt32Array::from(vec![1, 3, 3, 3, 5, 8]); + + let result = get_semi_indices(2..7, &input); + + assert_u32_values(&result, &[3, 5]); + } + + #[test] + fn get_semi_indices_ignores_out_of_range_indices() { + let input = UInt32Array::from(vec![0, 1, 3, 5, 8, 12]); + + let result = get_semi_indices(2..8, &input); + + assert_u32_values(&result, &[3, 5]); + } + + #[test] + fn get_semi_indices_handles_dense_matches() { + let input = UInt32Array::from(vec![2, 3, 4, 5]); + + let result = get_semi_indices(2..6, &input); + + assert_u32_values(&result, &[2, 3, 4, 5]); + } + + #[test] + fn get_semi_indices_handles_empty_input() { + let input = UInt32Array::from(Vec::::new()); + + let result = get_semi_indices(2..6, &input); + + assert!(result.is_empty()); + } + fn check( left: &[Column], right: &[Column], @@ -2568,6 +3128,7 @@ mod tests { create_stats(Some(left_num_rows), left_col_stats.clone(), false), create_stats(Some(right_num_rows), right_col_stats.clone(), false), &join_on, + NullEquality::NullEqualsNothing, ); assert_eq!( @@ -2582,6 +3143,46 @@ mod tests { Ok(()) } + #[test] + fn test_inner_join_cardinality_multiplication_overflow() { + let statistics = |num_rows, distinct_count| Statistics { + num_rows, + total_byte_size: Absent, + column_statistics: vec![ColumnStatistics { + distinct_count, + ..Default::default() + }], + }; + let large_row_count = usize::MAX / 2 + 1; + + // The Cartesian product overflows usize, but applying the NDV divisor + // produces a representable cardinality. + assert_eq!( + estimate_inner_join_cardinality( + statistics(Inexact(large_row_count), Inexact(1)), + statistics(Inexact(3), Inexact(3)), + ), + Some(Inexact(large_row_count)) + ); + assert_eq!( + estimate_inner_join_cardinality( + statistics(Exact(large_row_count), Exact(1)), + statistics(Exact(3), Exact(3)), + ), + Some(Exact(large_row_count)) + ); + + // If the normalized result itself cannot fit in usize, cap the + // estimate and mark it as inexact. + assert_eq!( + estimate_inner_join_cardinality( + statistics(Exact(usize::MAX), Exact(1)), + statistics(Exact(2), Exact(1)), + ), + Some(Inexact(usize::MAX)) + ); + } + #[test] fn test_inner_join_cardinality_multiple_column() -> Result<()> { let left_col_stats = vec![ @@ -2700,6 +3301,7 @@ mod tests { create_stats(Some(1000), left_col_stats.clone(), false), create_stats(Some(2000), right_col_stats.clone(), false), &join_on, + NullEquality::NullEqualsNothing, ) .unwrap(); assert_eq!(partial_join_stats.num_rows, expected_num_rows); @@ -2713,14 +3315,78 @@ mod tests { } #[test] - fn test_join_cardinality_when_one_column_is_disjoint() -> Result<()> { - // Left table (rows=1000) - // a: min=0, max=100, distinct=100 - // b: min=0, max=500, distinct=500 - // x: min=1000, max=10000, distinct=None - // - // Right table (rows=2000) - // c: min=0, max=100, distinct=50 + fn test_join_cardinality_key_order() -> Result<()> { + // Reversing join key order should not change estimated cardinality + let left_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(100), Absent), + create_column_stats(Inexact(0), Inexact(500), Inexact(500), Absent), + create_column_stats(Inexact(1000), Inexact(10000), Absent, Absent), + ]; + + let right_col_stats = vec![ + create_column_stats(Inexact(0), Inexact(100), Inexact(50), Absent), + create_column_stats(Inexact(0), Inexact(2000), Inexact(2500), Absent), + create_column_stats(Inexact(0), Inexact(100), Absent, Absent), + ]; + + let join_on_ab = vec![ + ( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("c", 0)) as _, + ), + ( + Arc::new(Column::new("b", 1)) as _, + Arc::new(Column::new("d", 1)) as _, + ), + ]; + let join_on_ba = vec![ + ( + Arc::new(Column::new("b", 1)) as _, + Arc::new(Column::new("d", 1)) as _, + ), + ( + Arc::new(Column::new("a", 0)) as _, + Arc::new(Column::new("c", 0)) as _, + ), + ]; + + let stats_ab = estimate_join_cardinality( + &JoinType::Inner, + create_stats(Some(1000), left_col_stats.clone(), false), + create_stats(Some(2000), right_col_stats.clone(), false), + &join_on_ab, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + let stats_ba = estimate_join_cardinality( + &JoinType::Inner, + create_stats(Some(1000), left_col_stats.clone(), false), + create_stats(Some(2000), right_col_stats.clone(), false), + &join_on_ba, + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + assert_eq!(stats_ab.num_rows, 1000); + assert_eq!(stats_ba.num_rows, stats_ab.num_rows); + assert_eq!(stats_ba.column_statistics, stats_ab.column_statistics); + assert_eq!( + stats_ab.column_statistics, + [left_col_stats, right_col_stats].concat() + ); + + Ok(()) + } + + #[test] + fn test_join_cardinality_when_one_column_is_disjoint() -> Result<()> { + // Left table (rows=1000) + // a: min=0, max=100, distinct=100 + // b: min=0, max=500, distinct=500 + // x: min=1000, max=10000, distinct=None + // + // Right table (rows=2000) + // c: min=0, max=100, distinct=50 // d: min=0, max=2000, distinct=2500 (how? some inexact statistics) // y: min=0, max=100, distinct=None // @@ -2771,6 +3437,7 @@ mod tests { create_stats(Some(1000), left_col_stats.clone(), true), create_stats(Some(2000), right_col_stats.clone(), true), &join_on, + NullEquality::NullEqualsNothing, ) .unwrap(); assert_eq!(partial_join_stats.num_rows, expected_num_rows); @@ -3006,6 +3673,7 @@ mod tests { column_statistics: inner_col_stats, }, &join_on, + NullEquality::NullEqualsNothing, ) .map(|cardinality| cardinality.num_rows); @@ -3040,6 +3708,7 @@ mod tests { column_statistics: dummy_column_stats.clone(), }, &join_on, + NullEquality::NullEqualsNothing, ); assert!( absent_outer_estimation.is_none(), @@ -3059,6 +3728,7 @@ mod tests { column_statistics: dummy_column_stats.clone(), }, &join_on, + NullEquality::NullEqualsNothing, ).expect("Expected non-empty PartialJoinStatistics for SemiJoin with absent inner num_rows"); assert_eq!( @@ -3079,6 +3749,7 @@ mod tests { column_statistics: dummy_column_stats, }, &join_on, + NullEquality::NullEqualsNothing, ); assert!( absent_inner_estimation.is_none(), @@ -3125,6 +3796,7 @@ mod tests { ], }, &join_on, + NullEquality::NullEqualsNothing, ) .map(|c| c.num_rows); assert_eq!(result, Some(13), "multi-column semi join"); @@ -3149,6 +3821,7 @@ mod tests { ], }, &join_on, + NullEquality::NullEqualsNothing, ) .map(|c| c.num_rows); assert_eq!(result, Some(87), "multi-column anti join"); @@ -3176,6 +3849,7 @@ mod tests { ], }, &join_on, + NullEquality::NullEqualsNothing, ) .map(|c| c.num_rows); assert_eq!(result, Some(50), "mixed stats: col1 skipped"); @@ -3200,6 +3874,7 @@ mod tests { ], }, &join_on, + NullEquality::NullEqualsNothing, ) .map(|c| c.num_rows); assert_eq!(result, Some(100), "no column has stats on both sides"); @@ -3228,6 +3903,7 @@ mod tests { ], }, &join_on, + NullEquality::NullEqualsNothing, ) .map(|c| c.num_rows); assert_eq!( @@ -3239,6 +3915,392 @@ mod tests { Ok(()) } + #[test] + fn test_semi_anti_join_disjoint_check_uses_only_join_keys() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + // Ranges for the join key overlap; ranges for the other column are disjoint + let left_stats = Statistics { + num_rows: Inexact(50), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Inexact(1), Inexact(10), Absent, Absent), + create_column_stats(Inexact(100), Inexact(200), Absent, Absent), + ], + }; + let right_stats = Statistics { + num_rows: Inexact(10), + total_byte_size: Absent, + column_statistics: vec![ + create_column_stats(Inexact(1), Inexact(10), Absent, Absent), + create_column_stats(Inexact(1000), Inexact(2000), Absent, Absent), + ], + }; + + let left_semi = estimate_join_cardinality( + &JoinType::LeftSemi, + left_stats.clone(), + right_stats.clone(), + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!(left_semi, Some(50)); + + let left_anti = estimate_join_cardinality( + &JoinType::LeftAnti, + left_stats, + right_stats, + &join_on, + NullEquality::NullEqualsNothing, + ) + .map(|c| c.num_rows); + assert_eq!(left_anti, Some(0)); + } + + #[test] + fn test_semi_join_scales_preserved_column_statistics() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(432_187), + total_byte_size: Absent, + column_statistics: vec![ + ColumnStatistics { + null_count: Exact(7_196), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(432_187_i64)), + sum_value: Absent, + distinct_count: Absent, + byte_size: Exact(3_457_496), + }, + ColumnStatistics { + null_count: Exact(7_196), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(432_187_i64)), + sum_value: Exact(ScalarValue::from(1_000_000_i64)), + distinct_count: Exact(500_000), + byte_size: Exact(3_457_496), + }, + ], + }, + Statistics { + num_rows: Inexact(32), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Inexact(1), + Inexact(32), + Absent, + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .expect("semi join cardinality should be estimated"); + + assert_eq!(result.num_rows, 32); + assert_eq!(result.total_byte_size, Inexact(512)); + assert_eq!(result.column_statistics[0].null_count, Exact(0)); + assert_eq!(result.column_statistics[0].distinct_count, Absent); + assert_eq!( + result.column_statistics[0].min_value, + Inexact(ScalarValue::from(1_i64)) + ); + assert_eq!( + result.column_statistics[0].max_value, + Inexact(ScalarValue::from(432_187_i64)) + ); + assert_eq!(result.column_statistics[0].byte_size, Inexact(256)); + assert_eq!(result.column_statistics[1].null_count, Inexact(1)); + // distinct_count is capped at the non-null output rows (32 - 1). + assert_eq!(result.column_statistics[1].distinct_count, Inexact(31)); + assert_eq!(result.column_statistics[1].sum_value, Absent); + assert_eq!(result.column_statistics[1].byte_size, Inexact(256)); + } + + #[test] + fn test_semi_join_null_equals_null_scales_join_key_nulls() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(100), + Exact(20), + )], + }, + Statistics { + num_rows: Inexact(10), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(10), + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNull, + ) + .expect("semi join cardinality should be estimated"); + + assert_eq!(result.num_rows, 10); + assert_eq!(result.column_statistics[0].null_count, Inexact(2)); + assert_eq!(result.column_statistics[0].distinct_count, Inexact(8)); + } + + #[test] + fn test_semi_join_total_byte_size_absent_if_any_column_byte_size_absent() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftSemi, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![ + ColumnStatistics { + null_count: Exact(0), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(100_i64)), + sum_value: Absent, + distinct_count: Absent, + byte_size: Exact(800), + }, + ColumnStatistics { + null_count: Exact(0), + min_value: Absent, + max_value: Absent, + sum_value: Absent, + distinct_count: Absent, + byte_size: Absent, + }, + ], + }, + Statistics { + num_rows: Inexact(10), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Inexact(1), + Inexact(10), + Absent, + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .expect("semi join cardinality should be estimated"); + + assert_eq!(result.num_rows, 10); + assert_eq!(result.total_byte_size, Absent); + } + + #[test] + fn test_anti_join_preserves_join_key_nulls() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftAnti, + Statistics { + num_rows: Inexact(1_000_000), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(900_000), + Exact(100_000), + )], + }, + Statistics { + num_rows: Inexact(900_000), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(900_000), + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .expect("anti join cardinality should be estimated"); + + assert_eq!(result.num_rows, 100_000); + assert_eq!(result.column_statistics[0].null_count, Inexact(100_000)); + assert_eq!(result.column_statistics[0].distinct_count, Inexact(0)); + } + + #[test] + fn test_anti_join_null_equals_null_scales_join_key_nulls() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + let result = estimate_join_cardinality( + &JoinType::LeftAnti, + Statistics { + num_rows: Inexact(100), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(100), + Exact(20), + )], + }, + Statistics { + num_rows: Inexact(10), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Absent, + Absent, + Inexact(10), + Absent, + )], + }, + &join_on, + NullEquality::NullEqualsNull, + ) + .expect("anti join cardinality should be estimated"); + + assert_eq!(result.num_rows, 90); + assert_eq!(result.column_statistics[0].null_count, Inexact(18)); + assert_eq!(result.column_statistics[0].distinct_count, Inexact(72)); + } + + #[test] + fn test_right_semi_join_scales_preserved_column_statistics() { + let join_on = vec![( + Arc::new(Column::new("l_key", 0)) as _, + Arc::new(Column::new("r_key", 0)) as _, + )]; + + // For a right semi join the right input is preserved, so its column + // statistics (and right join-key index) are the ones normalized. + let result = estimate_join_cardinality( + &JoinType::RightSemi, + Statistics { + num_rows: Inexact(32), + total_byte_size: Absent, + column_statistics: vec![create_column_stats( + Inexact(1), + Inexact(32), + Absent, + Absent, + )], + }, + Statistics { + num_rows: Inexact(432_187), + total_byte_size: Absent, + column_statistics: vec![ + ColumnStatistics { + null_count: Exact(7_196), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(432_187_i64)), + sum_value: Absent, + distinct_count: Absent, + byte_size: Exact(3_457_496), + }, + ColumnStatistics { + null_count: Exact(7_196), + min_value: Exact(ScalarValue::from(1_i64)), + max_value: Exact(ScalarValue::from(432_187_i64)), + sum_value: Exact(ScalarValue::from(1_000_000_i64)), + distinct_count: Exact(500_000), + byte_size: Exact(3_457_496), + }, + ], + }, + &join_on, + NullEquality::NullEqualsNothing, + ) + .expect("right semi join cardinality should be estimated"); + + assert_eq!(result.num_rows, 32); + // Join-key column: null counts collapse to exact zero (null keys never match). + assert_eq!(result.column_statistics[0].null_count, Exact(0)); + assert_eq!(result.column_statistics[0].byte_size, Inexact(256)); + // Non-key column: counts scaled to the subset, sum dropped, distinct + // capped at the non-null output rows (32 - 1). + assert_eq!(result.column_statistics[1].null_count, Inexact(1)); + assert_eq!(result.column_statistics[1].distinct_count, Inexact(31)); + assert_eq!(result.column_statistics[1].sum_value, Absent); + assert_eq!(result.column_statistics[1].byte_size, Inexact(256)); + } + + #[test] + fn test_adjust_right_output_partitioning_preserves_range() -> Result<()> { + let split_points = vec![ + SplitPoint::new(vec![ + ScalarValue::Int32(Some(10)), + ScalarValue::Int32(Some(100)), + ]), + SplitPoint::new(vec![ + ScalarValue::Int32(Some(20)), + ScalarValue::Int32(Some(50)), + ]), + ]; + let range = RangePartitioning::try_new( + LexOrdering::new([ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 0)), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 2)), + SortOptions::new(true, false), + ), + ]) + .unwrap(), + split_points.clone(), + )?; + + let adjusted = adjust_right_output_partitioning(&Partitioning::Range(range), 3)?; + let expected = Partitioning::Range(RangePartitioning::new( + LexOrdering::new([ + PhysicalSortExpr::new( + Arc::new(Column::new("a", 3)), + SortOptions::new(false, true), + ), + PhysicalSortExpr::new( + Arc::new(Column::new("b", 5)), + SortOptions::new(true, false), + ), + ]) + .unwrap(), + split_points, + )); + + assert_eq!(adjusted, expected); + Ok(()) + } + #[test] fn test_calculate_join_output_ordering() -> Result<()> { let left_ordering = LexOrdering::new(vec![ @@ -3556,6 +4618,282 @@ mod tests { assert_eq!(cmp_nl.compare(1, 1), Ordering::Less); } + #[test] + fn test_equal_rows_arr_filters_candidate_pairs() { + let left_a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 2, 3])); + let left_b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d"])); + let right_a: ArrayRef = Arc::new(Int32Array::from(vec![2, 2, 3, 4])); + let right_b: ArrayRef = Arc::new(StringArray::from(vec!["b", "d", "d", "a"])); + + let left_indices = UInt64Array::from(vec![0, 1, 2, 3]); + let right_indices = UInt32Array::from(vec![0, 0, 1, 2]); + + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[left_a, left_b], + &[right_a, right_b], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + assert_eq!(left_filtered, UInt64Array::from(vec![1, 3])); + assert_eq!(right_filtered, UInt32Array::from(vec![0, 2])); + } + + #[test] + fn test_equal_rows_arr_empty_keys_returns_empty() { + let left_indices = UInt64Array::from(vec![0, 1, 2]); + let right_indices = UInt32Array::from(vec![0, 1, 2]); + + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[], + &[], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + + assert_eq!(left_filtered.len(), 0); + assert_eq!(right_filtered.len(), 0); + } + + #[test] + fn test_equal_rows_arr_respects_null_equality() { + let left: ArrayRef = + Arc::new(Int32Array::from(vec![Some(1), None, Some(2), None])); + let right: ArrayRef = + Arc::new(Int32Array::from(vec![None, Some(1), Some(2), None])); + let left_indices = UInt64Array::from(vec![0, 1, 2, 3]); + let right_indices = UInt32Array::from(vec![1, 0, 2, 3]); + + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[Arc::clone(&left)], + &[Arc::clone(&right)], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0, 2])); + assert_eq!(right_filtered, UInt32Array::from(vec![1, 2])); + + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[left], + &[right], + NullEquality::NullEqualsNull, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0, 1, 2, 3])); + assert_eq!(right_filtered, UInt32Array::from(vec![1, 0, 2, 3])); + } + + #[test] + fn test_equal_rows_arr_single_string_col_fast_path() { + // Single-column string keys exercise the specialized fast path, + // including null handling under both null-equality modes. + let left: ArrayRef = Arc::new(StringArray::from(vec![ + Some("long_shared_join_key_value"), + None, + Some("long_shared_join_key_value"), + Some("other"), + ])); + let right: ArrayRef = Arc::new(StringArray::from(vec![ + Some("long_shared_join_key_value"), + None, + Some("mismatch"), + None, + ])); + let left_indices = UInt64Array::from(vec![0, 1, 2, 3]); + let right_indices = UInt32Array::from(vec![0, 1, 2, 3]); + + // NullEqualsNothing: only the (0,0) value pair matches; both-null drops. + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[Arc::clone(&left)], + &[Arc::clone(&right)], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0])); + assert_eq!(right_filtered, UInt32Array::from(vec![0])); + + // NullEqualsNull: the both-null (1,1) pair now also matches. + let (left_filtered, right_filtered) = equal_rows_arr( + &left_indices, + &right_indices, + &[left], + &[right], + NullEquality::NullEqualsNull, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0, 1])); + assert_eq!(right_filtered, UInt32Array::from(vec![0, 1])); + } + + #[test] + fn test_equal_rows_arr_single_col_covers_all_specialized_types() { + // Drive every specialized single-column fast-path arm. Each case has a + // matching pair at index 0 and a non-matching pair at index 1, so a + // correct arm keeps exactly the first pair. + fn check(left: ArrayRef, right: ArrayRef) { + let (left_filtered, right_filtered) = equal_rows_arr( + &UInt64Array::from(vec![0, 1]), + &UInt32Array::from(vec![0, 1]), + &[left], + &[right], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0])); + assert_eq!(right_filtered, UInt32Array::from(vec![0])); + } + + check( + Arc::new(BooleanArray::from(vec![true, false])), + Arc::new(BooleanArray::from(vec![true, true])), + ); + check( + Arc::new(Int8Array::from(vec![1, 2])), + Arc::new(Int8Array::from(vec![1, 3])), + ); + check( + Arc::new(Int16Array::from(vec![1, 2])), + Arc::new(Int16Array::from(vec![1, 3])), + ); + check( + Arc::new(Int64Array::from(vec![1, 2])), + Arc::new(Int64Array::from(vec![1, 3])), + ); + check( + Arc::new(UInt8Array::from(vec![1, 2])), + Arc::new(UInt8Array::from(vec![1, 3])), + ); + check( + Arc::new(UInt16Array::from(vec![1, 2])), + Arc::new(UInt16Array::from(vec![1, 3])), + ); + check( + Arc::new(UInt32Array::from(vec![1, 2])), + Arc::new(UInt32Array::from(vec![1, 3])), + ); + check( + Arc::new(UInt64Array::from(vec![1, 2])), + Arc::new(UInt64Array::from(vec![1, 3])), + ); + check( + Arc::new(Decimal128Array::from(vec![1i128, 2])), + Arc::new(Decimal128Array::from(vec![1i128, 3])), + ); + check( + Arc::new(BinaryArray::from_iter_values([b"a".as_ref(), b"b"])), + Arc::new(BinaryArray::from_iter_values([b"a".as_ref(), b"c"])), + ); + check( + Arc::new(LargeBinaryArray::from_iter_values([b"a".as_ref(), b"b"])), + Arc::new(LargeBinaryArray::from_iter_values([b"a".as_ref(), b"c"])), + ); + check( + Arc::new(BinaryViewArray::from_iter_values([b"a".as_ref(), b"b"])), + Arc::new(BinaryViewArray::from_iter_values([b"a".as_ref(), b"c"])), + ); + check( + Arc::new( + FixedSizeBinaryArray::try_from_iter([[1u8], [2u8]].into_iter()).unwrap(), + ), + Arc::new( + FixedSizeBinaryArray::try_from_iter([[1u8], [3u8]].into_iter()).unwrap(), + ), + ); + check( + Arc::new(LargeStringArray::from(vec!["a", "b"])), + Arc::new(LargeStringArray::from(vec!["a", "c"])), + ); + check( + Arc::new(StringViewArray::from(vec!["a", "b"])), + Arc::new(StringViewArray::from(vec!["a", "c"])), + ); + check( + Arc::new(Date32Array::from(vec![1, 2])), + Arc::new(Date32Array::from(vec![1, 3])), + ); + check( + Arc::new(Date64Array::from(vec![1, 2])), + Arc::new(Date64Array::from(vec![1, 3])), + ); + check( + Arc::new(TimestampSecondArray::from(vec![1, 2])), + Arc::new(TimestampSecondArray::from(vec![1, 3])), + ); + check( + Arc::new(TimestampMillisecondArray::from(vec![1, 2])), + Arc::new(TimestampMillisecondArray::from(vec![1, 3])), + ); + check( + Arc::new(TimestampMicrosecondArray::from(vec![1, 2])), + Arc::new(TimestampMicrosecondArray::from(vec![1, 3])), + ); + check( + Arc::new(TimestampNanosecondArray::from(vec![1, 2])), + Arc::new(TimestampNanosecondArray::from(vec![1, 3])), + ); + } + + #[test] + fn test_equal_rows_arr_single_float_col_uses_general_path() { + // Floats are intentionally not specialized: the fast path returns + // `None` and the general comparator handles them (covers the + // fall-through arm). + let left: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 2.0])); + let right: ArrayRef = Arc::new(Float64Array::from(vec![1.0, 3.0])); + let (left_filtered, right_filtered) = equal_rows_arr( + &UInt64Array::from(vec![0, 1]), + &UInt32Array::from(vec![0, 1]), + &[left], + &[right], + NullEquality::NullEqualsNothing, + ) + .unwrap(); + assert_eq!(left_filtered, UInt64Array::from(vec![0])); + assert_eq!(right_filtered, UInt32Array::from(vec![0])); + } + + #[test] + fn test_equal_rows_arr_rejects_mismatched_inputs() { + let left: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + let right: ArrayRef = Arc::new(Int32Array::from(vec![1, 2])); + + let err = equal_rows_arr( + &UInt64Array::from(vec![0, 1]), + &UInt32Array::from(vec![0]), + &[Arc::clone(&left)], + &[Arc::clone(&right)], + NullEquality::NullEqualsNothing, + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("Cannot compare join indices with different lengths") + ); + + let err = equal_rows_arr( + &UInt64Array::from(vec![0, 1]), + &UInt32Array::from(vec![0, 1]), + &[left, Arc::new(Int32Array::from(vec![3, 4]))], + &[right], + NullEquality::NullEqualsNothing, + ) + .unwrap_err(); + assert!( + err.to_string() + .contains("Cannot compare join keys with different column counts") + ); + } + #[test] fn test_max_distinct_count_preserves_precision_when_not_capped() { assert_eq!( diff --git a/datafusion/physical-plan/src/lib.rs b/datafusion/physical-plan/src/lib.rs index 3005e975424b4..9e50a93b2163f 100644 --- a/datafusion/physical-plan/src/lib.rs +++ b/datafusion/physical-plan/src/lib.rs @@ -37,18 +37,25 @@ pub use datafusion_expr::{Accumulator, ColumnarValue}; use datafusion_physical_expr::PhysicalSortExpr; pub use datafusion_physical_expr::window::WindowExpr; pub use datafusion_physical_expr::{ - Distribution, Partitioning, PhysicalExpr, expressions, + Distribution, Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, expressions, }; pub use crate::display::{DefaultDisplay, DisplayAs, DisplayFormatType, VerboseDisplay}; +pub use crate::distribution_requirements::{ + ChildSatisfactionOptions, InputDistributionRequirements, +}; +#[expect(deprecated)] pub use crate::execution_plan::{ - ExecutionPlan, ExecutionPlanProperties, PlanProperties, collect, collect_partitioned, - displayable, execute_input_stream, execute_stream, execute_stream_partitioned, - get_plan_string, with_new_children_if_necessary, + AsPhysicalExprRef, ChildrenPropertiesMode, ExecutionPlan, ExecutionPlanProperties, + PlanProperties, ReplaceChildrenOptions, apply_expression_roots, collect, + collect_partitioned, displayable, execute_input_stream, execute_stream, + execute_stream_partitioned, get_plan_string, replace_children_if_necessary, + with_new_children_if_necessary, }; pub use crate::metrics::Metric; pub use crate::ordering::InputOrderMode; pub use crate::sort_pushdown::SortOrderPushdownResult; +pub use crate::statistics::{ChildStats, StatisticsArgs, StatisticsContext}; pub use crate::stream::EmptyRecordBatchStream; pub use crate::topk::TopK; pub use crate::visitor::{ExecutionPlanVisitor, accept, visit_execution_plan}; @@ -71,6 +78,7 @@ pub mod column_rewriter; pub mod common; pub mod coop; pub mod display; +pub mod distribution_requirements; pub mod empty; pub mod execution_plan; pub mod explain; @@ -83,12 +91,15 @@ pub mod metrics; pub mod operator_statistics; pub mod placeholder_row; pub mod projection; +#[cfg(feature = "proto")] +pub mod proto; pub mod recursive_query; pub mod repartition; pub mod scalar_subquery; pub mod sort_pushdown; pub mod sorts; pub mod spill; +pub mod statistics; pub mod stream; pub mod streaming; pub mod tree_node; diff --git a/datafusion/physical-plan/src/limit.rs b/datafusion/physical-plan/src/limit.rs index 223a476493b39..dd62c93d1cfe0 100644 --- a/datafusion/physical-plan/src/limit.rs +++ b/datafusion/physical-plan/src/limit.rs @@ -27,15 +27,16 @@ use super::{ SendableRecordBatchStream, Statistics, }; use crate::execution_plan::{Boundedness, CardinalityEffect}; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, Partitioning, - check_if_same_properties, + ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan, Partitioning, + ReplaceChildrenOptions, validate_child_count, }; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{Result, assert_eq_or_internal_err, internal_err}; +use datafusion_common::{Result, assert_eq_or_internal_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; @@ -54,8 +55,8 @@ pub struct GlobalLimitExec { fetch: Option, /// Execution metrics metrics: ExecutionPlanMetricsSet, - /// Does the limit have to preserve the order of its input, and if so what is it? - /// Some optimizations may reorder the input if no particular sort is required + /// Input ordering that must be preserved so limit pushdown does not change + /// which rows are returned. required_ordering: Option, cache: Arc, } @@ -109,17 +110,6 @@ impl GlobalLimitExec { pub fn set_required_ordering(&mut self, required_ordering: Option) { self.required_ordering = required_ordering; } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for GlobalLimitExec { @@ -163,7 +153,11 @@ impl ExecutionPlan for GlobalLimitExec { } fn required_input_distribution(&self) -> Vec { - vec![Distribution::SinglePartition] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } fn maintains_input_order(&self) -> Vec { @@ -176,28 +170,50 @@ impl ExecutionPlan for GlobalLimitExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to required ordering expressions if present - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = &self.required_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut new_limit = + GlobalLimitExec::new(children.swap_remove(0), self.skip, self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) } } - Ok(tnr) } fn with_new_children( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(GlobalLimitExec::new( - children.swap_remove(0), - self.skip, - self.fetch, - ))) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -234,8 +250,16 @@ impl ExecutionPlan for GlobalLimitExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?); + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let stats = input_stats[0].as_ref().clone(); Ok(Arc::new(stats.with_fetch(self.fetch, self.skip, 1)?)) } @@ -246,6 +270,68 @@ impl ExecutionPlan for GlobalLimitExec { fn supports_limit_pushdown(&self) -> bool { true } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto; + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let required_ordering = optional_ordering_try_to_proto( + self.required_ordering.as_ref(), + &ctx.expr_ctx(), + )?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit(Box::new( + protobuf::GlobalLimitExecNode { + input: Some(Box::new(input)), + skip: self.skip() as u32, + fetch: match self.fetch() { + Some(n) => n as i64, + _ => -1, // no limit + }, + required_ordering, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl GlobalLimitExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto; + use datafusion_proto_models::protobuf; + let limit = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::GlobalLimit, + "GlobalLimitExec", + ); + let input = ctx.decode_required_child( + limit.input.as_deref(), + "GlobalLimitExec", + "input", + )?; + let fetch = if limit.fetch >= 0 { + Some(limit.fetch as usize) + } else { + None + }; + let required_ordering = optional_ordering_try_from_proto( + &limit.required_ordering, + &ctx.expr_ctx(input.schema().as_ref()), + )?; + let mut exec = GlobalLimitExec::new(input, limit.skip as usize, fetch); + exec.set_required_ordering(required_ordering); + Ok(Arc::new(exec)) + } } /// LocalLimitExec applies a limit to a single partition @@ -257,8 +343,8 @@ pub struct LocalLimitExec { fetch: usize, /// Execution metrics metrics: ExecutionPlanMetricsSet, - /// If the child plan is a sort node, after the sort node is removed during - /// physical optimization, we should add the required ordering to the limit node + /// Input ordering that must be preserved so limit pushdown does not change + /// which rows are returned. required_ordering: Option, cache: Arc, } @@ -306,17 +392,6 @@ impl LocalLimitExec { pub fn set_required_ordering(&mut self, required_ordering: Option) { self.required_ordering = required_ordering; } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for LocalLimitExec { @@ -360,30 +435,50 @@ impl ExecutionPlan for LocalLimitExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to required ordering expressions if present - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = &self.required_ordering { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut new_limit = + LocalLimitExec::new(children.swap_remove(0), self.fetch); + new_limit.set_required_ordering(self.required_ordering.clone()); + Ok(Arc::new(new_limit)) } } - Ok(tnr) } fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - match children.len() { - 1 => Ok(Arc::new(LocalLimitExec::new( - Arc::clone(&children[0]), - self.fetch, - ))), - _ => internal_err!("LocalLimitExec wrong number of children"), - } + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -411,8 +506,16 @@ impl ExecutionPlan for LocalLimitExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let stats = Arc::unwrap_or_clone(self.input.partition_statistics(partition)?); + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let stats = input_stats[0].as_ref().clone(); Ok(Arc::new(stats.with_fetch(Some(self.fetch), 0, 1)?)) } @@ -427,6 +530,56 @@ impl ExecutionPlan for LocalLimitExec { fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::LowerEqual } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_to_proto; + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let required_ordering = optional_ordering_try_to_proto( + self.required_ordering.as_ref(), + &ctx.expr_ctx(), + )?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::LocalLimit(Box::new( + protobuf::LocalLimitExecNode { + input: Some(Box::new(input)), + fetch: self.fetch() as u32, + required_ordering, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl LocalLimitExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_physical_expr_common::sort_expr::optional_ordering_try_from_proto; + use datafusion_proto_models::protobuf; + let limit = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::LocalLimit, + "LocalLimitExec", + ); + let input = + ctx.decode_required_child(limit.input.as_deref(), "LocalLimitExec", "input")?; + let required_ordering = optional_ordering_try_from_proto( + &limit.required_ordering, + &ctx.expr_ctx(input.schema().as_ref()), + )?; + let mut exec = LocalLimitExec::new(input, limit.fetch as usize); + exec.set_required_ordering(required_ordering); + Ok(Arc::new(exec)) + } } /// A Limit stream skips `skip` rows, and then fetch up to `fetch` rows. @@ -559,13 +712,16 @@ mod tests { use super::*; use crate::coalesce_partitions::CoalescePartitionsExec; use crate::common::collect; + use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; use crate::aggregates::{AggregateExec, AggregateMode, PhysicalGroupBy}; use arrow::array::RecordBatchOptions; + use arrow::compute::SortOptions; use arrow::datatypes::Schema; use datafusion_common::stats::Precision; use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr::{PhysicalExpr, PhysicalSortExpr}; #[tokio::test] async fn limit() -> Result<()> { @@ -755,80 +911,106 @@ mod tests { Ok(()) } - #[tokio::test] - async fn test_row_number_statistics_for_global_limit() -> Result<()> { - let row_count = row_number_statistics_for_global_limit(0, Some(10)).await?; + #[test] + fn replace_children_preserves_required_ordering() -> Result<()> { + let source = test::scan_partitioned(1); + let schema = source.schema(); + let ordering = LexOrdering::new(vec![PhysicalSortExpr { + expr: col("i", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }]); + + let mut global = GlobalLimitExec::new(Arc::clone(&source), 0, Some(10)); + global.set_required_ordering(ordering.clone()); + let rebuilt = Arc::new(global).replace_children( + vec![test::scan_partitioned(1)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let rebuilt = rebuilt.downcast_ref::().unwrap(); + assert_eq!(rebuilt.required_ordering(), &ordering); + + let mut local = LocalLimitExec::new(source, 10); + local.set_required_ordering(ordering.clone()); + let rebuilt = Arc::new(local).replace_children( + vec![test::scan_partitioned(1)], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + let rebuilt = rebuilt.downcast_ref::().unwrap(); + assert_eq!(rebuilt.required_ordering(), &ordering); + + Ok(()) + } + + #[test] + fn test_row_number_statistics_for_global_limit() -> Result<()> { + let row_count = row_number_statistics_for_global_limit(0, Some(10))?; assert_eq!(row_count, Precision::Exact(10)); - let row_count = row_number_statistics_for_global_limit(5, Some(10)).await?; + let row_count = row_number_statistics_for_global_limit(5, Some(10))?; assert_eq!(row_count, Precision::Exact(10)); - let row_count = row_number_statistics_for_global_limit(400, Some(10)).await?; + let row_count = row_number_statistics_for_global_limit(400, Some(10))?; assert_eq!(row_count, Precision::Exact(0)); - let row_count = row_number_statistics_for_global_limit(398, Some(10)).await?; + let row_count = row_number_statistics_for_global_limit(398, Some(10))?; assert_eq!(row_count, Precision::Exact(2)); - let row_count = row_number_statistics_for_global_limit(398, Some(1)).await?; + let row_count = row_number_statistics_for_global_limit(398, Some(1))?; assert_eq!(row_count, Precision::Exact(1)); - let row_count = row_number_statistics_for_global_limit(398, None).await?; + let row_count = row_number_statistics_for_global_limit(398, None)?; assert_eq!(row_count, Precision::Exact(2)); - let row_count = - row_number_statistics_for_global_limit(0, Some(usize::MAX)).await?; + let row_count = row_number_statistics_for_global_limit(0, Some(usize::MAX))?; assert_eq!(row_count, Precision::Exact(400)); - let row_count = - row_number_statistics_for_global_limit(398, Some(usize::MAX)).await?; + let row_count = row_number_statistics_for_global_limit(398, Some(usize::MAX))?; assert_eq!(row_count, Precision::Exact(2)); - let row_count = - row_number_inexact_statistics_for_global_limit(0, Some(10)).await?; + let row_count = row_number_inexact_statistics_for_global_limit(0, Some(10))?; assert_eq!(row_count, Precision::Inexact(10)); - let row_count = - row_number_inexact_statistics_for_global_limit(5, Some(10)).await?; + let row_count = row_number_inexact_statistics_for_global_limit(5, Some(10))?; assert_eq!(row_count, Precision::Inexact(10)); // Input was Inexact, so an `nr <= skip` outcome must remain Inexact: // the inexact estimate could be wrong, so we cannot promote 0 to // Exact. - let row_count = - row_number_inexact_statistics_for_global_limit(400, Some(10)).await?; + let row_count = row_number_inexact_statistics_for_global_limit(400, Some(10))?; assert_eq!(row_count, Precision::Inexact(0)); - let row_count = - row_number_inexact_statistics_for_global_limit(398, Some(10)).await?; + let row_count = row_number_inexact_statistics_for_global_limit(398, Some(10))?; assert_eq!(row_count, Precision::Inexact(2)); - let row_count = - row_number_inexact_statistics_for_global_limit(398, Some(1)).await?; + let row_count = row_number_inexact_statistics_for_global_limit(398, Some(1))?; assert_eq!(row_count, Precision::Inexact(1)); - let row_count = row_number_inexact_statistics_for_global_limit(398, None).await?; + let row_count = row_number_inexact_statistics_for_global_limit(398, None)?; assert_eq!(row_count, Precision::Inexact(2)); let row_count = - row_number_inexact_statistics_for_global_limit(0, Some(usize::MAX)).await?; + row_number_inexact_statistics_for_global_limit(0, Some(usize::MAX))?; assert_eq!(row_count, Precision::Inexact(400)); let row_count = - row_number_inexact_statistics_for_global_limit(398, Some(usize::MAX)).await?; + row_number_inexact_statistics_for_global_limit(398, Some(usize::MAX))?; assert_eq!(row_count, Precision::Inexact(2)); Ok(()) } - #[tokio::test] - async fn test_row_number_statistics_for_local_limit() -> Result<()> { - let row_count = row_number_statistics_for_local_limit(4, 10).await?; + #[test] + fn test_row_number_statistics_for_local_limit() -> Result<()> { + let row_count = row_number_statistics_for_local_limit(4, 10)?; assert_eq!(row_count, Precision::Exact(10)); Ok(()) } - async fn row_number_statistics_for_global_limit( + fn row_number_statistics_for_global_limit( skip: usize, fetch: Option, ) -> Result> { @@ -840,7 +1022,9 @@ mod tests { let offset = GlobalLimitExec::new(Arc::new(CoalescePartitionsExec::new(csv)), skip, fetch); - Ok(offset.partition_statistics(None)?.num_rows) + Ok(StatisticsContext::new() + .compute(&offset, &StatisticsArgs::new())? + .num_rows) } pub fn build_group_by( @@ -854,7 +1038,7 @@ mod tests { PhysicalGroupBy::new_single(group_by_expr.clone()) } - async fn row_number_inexact_statistics_for_global_limit( + fn row_number_inexact_statistics_for_global_limit( skip: usize, fetch: Option, ) -> Result> { @@ -880,10 +1064,12 @@ mod tests { fetch, ); - Ok(offset.partition_statistics(None)?.num_rows) + Ok(StatisticsContext::new() + .compute(&offset, &StatisticsArgs::new())? + .num_rows) } - async fn row_number_statistics_for_local_limit( + fn row_number_statistics_for_local_limit( num_partitions: usize, fetch: usize, ) -> Result> { @@ -893,7 +1079,9 @@ mod tests { let offset = LocalLimitExec::new(csv, fetch); - Ok(offset.partition_statistics(None)?.num_rows) + Ok(StatisticsContext::new() + .compute(&offset, &StatisticsArgs::new())? + .num_rows) } /// Return a RecordBatch with a single array with row_count sz diff --git a/datafusion/physical-plan/src/memory.rs b/datafusion/physical-plan/src/memory.rs index e172ef4463ec4..efe42c7ebc5f0 100644 --- a/datafusion/physical-plan/src/memory.rs +++ b/datafusion/physical-plan/src/memory.rs @@ -26,8 +26,8 @@ use crate::coop::cooperative; use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - RecordBatchStream, SendableRecordBatchStream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, }; use arrow::array::RecordBatch; @@ -314,14 +314,15 @@ impl ExecutionPlan for LazyMemoryExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_or_internal_err!( children.is_empty(), @@ -330,6 +331,16 @@ impl ExecutionPlan for LazyMemoryExec { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/operator_statistics/mod.rs b/datafusion/physical-plan/src/operator_statistics/mod.rs index eca017cde9d0c..16b89e9eca926 100644 --- a/datafusion/physical-plan/src/operator_statistics/mod.rs +++ b/datafusion/physical-plan/src/operator_statistics/mod.rs @@ -94,6 +94,7 @@ use datafusion_common::stats::Precision; use datafusion_common::{Result, Statistics}; use crate::ExecutionPlan; +use crate::statistics::{StatisticsArgs, StatisticsContext}; // ============================================================================ // ExtendedStatistics: Statistics with type-safe extensions @@ -266,7 +267,7 @@ impl StatisticsProvider for DefaultStatisticsProvider { plan: &dyn ExecutionPlan, _child_stats: &[ExtendedStatistics], ) -> Result { - let base = plan.partition_statistics(None)?; + let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; Ok(StatisticsResult::Computed(ExtendedStatistics::new_arc( base, ))) @@ -358,7 +359,7 @@ impl StatisticsRegistry { pub fn compute(&self, plan: &dyn ExecutionPlan) -> Result { // Fast path: no providers registered, skip the walk entirely if self.providers.is_empty() { - let base = plan.partition_statistics(None)?; + let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; return Ok(ExtendedStatistics::new_arc(base)); } @@ -382,7 +383,7 @@ impl StatisticsRegistry { } } // Fallback: use plan's built-in stats - let base = plan.partition_statistics(None)?; + let base = StatisticsContext::new().compute(plan, &StatisticsArgs::new())?; Ok(ExtendedStatistics::new_arc(base)) } @@ -505,7 +506,9 @@ fn computed_with_row_count( plan: &dyn ExecutionPlan, num_rows: Precision, ) -> Result { - let mut base = Arc::unwrap_or_clone(plan.partition_statistics(None)?); + let mut base = Arc::unwrap_or_clone( + StatisticsContext::new().compute(plan, &StatisticsArgs::new())?, + ); rescale_byte_size(&mut base, num_rows); Ok(StatisticsResult::Computed(ExtendedStatistics::new(base))) } @@ -1023,7 +1026,11 @@ mod tests { use super::*; use crate::filter::FilterExec; use crate::projection::ProjectionExec; - use crate::{DisplayAs, DisplayFormatType, PlanProperties}; + use crate::statistics::StatisticsArgs; + use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, PlanProperties, + ReplaceChildrenOptions, + }; use arrow::datatypes::{DataType, Field, Schema}; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; @@ -1103,20 +1110,31 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn properties(&self) -> &Arc { &self.cache } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -1129,9 +1147,10 @@ mod tests { unimplemented!() } - fn partition_statistics( + fn statistics_from_inputs( &self, - _partition: Option, + _input_stats: &[Arc], + _args: &StatisticsArgs, ) -> Result> { Ok(Arc::new(self.stats.clone())) } @@ -1209,22 +1228,33 @@ mod tests { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(CustomExec { input: Arc::clone(&children[0]), })) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn properties(&self) -> &Arc { self.input.properties() } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } diff --git a/datafusion/physical-plan/src/placeholder_row.rs b/datafusion/physical-plan/src/placeholder_row.rs index ae8e73cd74ade..67c063b65cbc6 100644 --- a/datafusion/physical-plan/src/placeholder_row.rs +++ b/datafusion/physical-plan/src/placeholder_row.rs @@ -23,8 +23,9 @@ use crate::coop::cooperative; use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::memory::MemoryStream; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, - SendableRecordBatchStream, Statistics, common, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, + common, }; use arrow::array::{ArrayRef, NullArray, RecordBatch, RecordBatchOptions}; @@ -35,6 +36,7 @@ use datafusion_execution::TaskContext; use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_expr::PhysicalExpr; +use crate::statistics::StatisticsArgs; use log::trace; /// Execution plan for empty relation with produce_one_row=true @@ -139,18 +141,29 @@ impl ExecutionPlan for PlaceholderRowExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -173,12 +186,16 @@ impl ExecutionPlan for PlaceholderRowExec { Ok(Box::pin(cooperative(ms))) } - fn partition_statistics(&self, partition: Option) -> Result> { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { let batches = self .data() .expect("Create single row placeholder RecordBatch should not fail"); - let batches = match partition { + let batches = match args.partition() { Some(_) => vec![batches], // entire plan None => vec![batches; self.partitions], @@ -190,21 +207,70 @@ impl ExecutionPlan for PlaceholderRowExec { None, ))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + _ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let schema = self.schema().as_ref().try_into()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow( + protobuf::PlaceholderRowExecNode { + schema: Some(schema), + partitions: self + .properties() + .output_partitioning() + .partition_count() as u32, + }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl PlaceholderRowExec { + /// Reconstruct a [`PlaceholderRowExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + _ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let placeholder = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow, + "PlaceholderRowExec", + ); + let schema = placeholder.schema.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "PlaceholderRowExec is missing required field 'schema'" + ) + })?; + let schema = Arc::new(Schema::try_from(schema)?); + // A zero (absent) partition count comes from a plan encoded before the + // field existed, which always meant a single partition. + let partitions = placeholder.partitions.max(1) as usize; + Ok(Arc::new( + PlaceholderRowExec::new(schema).with_partitions(partitions), + )) + } } #[cfg(test)] mod tests { use super::*; - use crate::test; - use crate::with_new_children_if_necessary; + use crate::{execution_plan::replace_children_if_necessary, test}; #[test] - fn with_new_children() -> Result<()> { + fn replace_children() -> Result<()> { let schema = test::aggr_test_schema(); let placeholder = Arc::new(PlaceholderRowExec::new(schema)); - let placeholder_2 = with_new_children_if_necessary( + let placeholder_2 = replace_children_if_necessary( Arc::clone(&placeholder) as Arc, vec![], )?; @@ -212,7 +278,7 @@ mod tests { let too_many_kids = vec![placeholder_2]; assert!( - with_new_children_if_necessary(placeholder, too_many_kids).is_err(), + replace_children_if_necessary(placeholder, too_many_kids).is_err(), "expected error when providing list of kids" ); Ok(()) diff --git a/datafusion/physical-plan/src/projection.rs b/datafusion/physical-plan/src/projection.rs index e5b91fbb1c5d4..cf362cdee55d3 100644 --- a/datafusion/physical-plan/src/projection.rs +++ b/datafusion/physical-plan/src/projection.rs @@ -27,25 +27,29 @@ use super::{ SendableRecordBatchStream, SortOrderPushdownResult, Statistics, }; use crate::column_rewriter::PhysicalColumnRewriter; -use crate::execution_plan::CardinalityEffect; +use crate::execution_plan::{CardinalityEffect, replace_children_if_necessary}; use crate::filter_pushdown::{ ChildFilterDescription, ChildPushdownResult, FilterDescription, FilterPushdownPhase, FilterPushdownPropagation, FilterRemapper, PushedDownPredicate, }; use crate::joins::utils::{ColumnIndex, JoinFilter, JoinOn, JoinOnRef}; -use crate::{DisplayFormatType, ExecutionPlan, PhysicalExpr, check_if_same_properties}; +use crate::statistics::{ChildStats, StatisticsArgs}; +use crate::{ + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, PhysicalExpr, + ReplaceChildrenOptions, validate_child_count, +}; use std::collections::HashMap; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; -use arrow::datatypes::SchemaRef; +use arrow::datatypes::{Schema, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::config::ConfigOptions; use datafusion_common::tree_node::{ Transformed, TransformedResult, TreeNode, TreeNodeRecursion, }; -use datafusion_common::{DataFusionError, JoinSide, Result, internal_err}; +use datafusion_common::{DataFusionError, JoinSide, Result, internal_err, plan_err}; use datafusion_execution::TaskContext; use datafusion_expr::ExpressionPlacement; use datafusion_physical_expr::equivalence::ProjectionMapping; @@ -142,6 +146,34 @@ impl ProjectionExec { Self::try_from_projector(projector, input) } + /// Create a projection using field and schema metadata from + /// `projected_schema`. + /// + /// Field names, data types, and nullability are still derived from the physical + /// projection expressions and the input plan; only field and schema metadata are + /// taken from `projected_schema`. + /// + /// # Errors + /// + /// Returns an error if the projection cannot be applied to the input plan, or if + /// `projected_schema` has a different number of fields than the projection. + pub fn try_new_with_schema_metadata( + expr: I, + input: Arc, + projected_schema: &Schema, + ) -> Result + where + I: IntoIterator, + E: Into, + { + let input_schema = input.schema(); + let expr_arc = expr.into_iter().map(Into::into).collect::>(); + let projection = ProjectionExprs::from_expressions(expr_arc); + let projector = projection + .make_projector_with_schema_metadata(&input_schema, projected_schema)?; + Self::try_from_projector(projector, input) + } + fn try_from_projector( projector: Projector, input: Arc, @@ -221,17 +253,6 @@ impl ProjectionExec { } Ok(alias_map) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for ProjectionExec { @@ -314,25 +335,49 @@ impl ExecutionPlan for ProjectionExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for proj_expr in self.projector.projection().as_ref().iter() { - tnr = tnr.visit_sibling(|| f(proj_expr.expr.as_ref()))?; + crate::apply_expression_roots(self.projector.projection().as_ref().iter(), f) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => ProjectionExec::try_from_projector( + self.projector.clone(), + children.swap_remove(0), + ) + .map(|p| Arc::new(p) as _), } - Ok(tnr) } fn with_new_children( self: Arc, - mut children: Vec>, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - ProjectionExec::try_from_projector( - self.projector.clone(), - children.swap_remove(0), + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), ) - .map(|p| Arc::new(p) as _) } fn execute( @@ -359,9 +404,16 @@ impl ExecutionPlan for ProjectionExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let input_stats = - Arc::unwrap_or_clone(self.input.partition_statistics(partition)?); + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let input_stats = input_stats[0].as_ref().clone(); let output_schema = self.schema(); Ok(Arc::new( self.projector @@ -382,12 +434,9 @@ impl ExecutionPlan for ProjectionExec { &self, projection: &ProjectionExec, ) -> Result>> { - let maybe_unified = try_unifying_projections(projection, self)?; - if let Some(new_plan) = maybe_unified { - // To unify 3 or more sequential projections: - remove_unnecessary_projections(new_plan).data().map(Some) - } else { - Ok(Some(Arc::new(projection.clone()))) + match try_collapse_projection_chain(projection)? { + Some(plan) => Ok(Some(plan)), + None => Ok(Some(Arc::new(projection.clone()))), } } @@ -481,11 +530,13 @@ impl ExecutionPlan for ProjectionExec { // Recursively push down to child node match child.try_pushdown_sort(&child_order)? { SortOrderPushdownResult::Exact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Exact { inner: new_exec }) } SortOrderPushdownResult::Inexact { inner } => { - let new_exec = Arc::new(self.clone()).with_new_children(vec![inner])?; + let new_exec = + replace_children_if_necessary(Arc::new(self.clone()), vec![inner])?; Ok(SortOrderPushdownResult::Inexact { inner: new_exec }) } SortOrderPushdownResult::Unsupported => { @@ -501,11 +552,75 @@ impl ExecutionPlan for ProjectionExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = ctx.encode_expressions(self.expr().iter().map(|p| &p.expr))?; + let expr_name = self.expr().iter().map(|p| p.alias.clone()).collect(); + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Projection(Box::new( + protobuf::ProjectionExecNode { + input: Some(Box::new(input)), + expr, + expr_name, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ProjectionExec { + /// Reconstruct a [`ProjectionExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]: it takes the whole + /// [`PhysicalPlanNode`] so every plan's `try_from_proto` shares one + /// signature. Child plans and expressions are decoded recursively via the + /// [`ExecutionPlanDecodeCtx`]. + /// + /// [`PhysicalPlanNode`]: datafusion_proto_models::protobuf::PhysicalPlanNode + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + /// [`ExecutionPlanDecodeCtx`]: crate::proto::ExecutionPlanDecodeCtx + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let projection = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Projection, + "ProjectionExec", + ); + let input = ctx.decode_required_child( + projection.input.as_deref(), + "ProjectionExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = projection + .expr + .iter() + .zip(projection.expr_name.iter()) + .map(|(expr, name)| { + Ok(ProjectionExpr { + expr: ctx.decode_expr(expr, input_schema.as_ref())?, + alias: name.to_string(), + }) + }) + .collect::>>()?; + Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) + } } impl ProjectionStream { @@ -654,6 +769,10 @@ pub struct JoinData { pub join_on: JoinOn, } +#[deprecated( + since = "55.0.0", + note = "Use try_pushdown_through_join_with_column_indices instead" +)] pub fn try_pushdown_through_join( projection: &ProjectionExec, join_left: &Arc, @@ -662,53 +781,149 @@ pub fn try_pushdown_through_join( schema: &SchemaRef, filter: Option<&JoinFilter>, ) -> Result> { + let left_field_count = join_left.schema().fields().len(); + let column_indices = schema + .fields() + .iter() + .enumerate() + .map(|(index, _)| { + if index < left_field_count { + ColumnIndex { + index, + side: JoinSide::Left, + } + } else { + ColumnIndex { + index: index - left_field_count, + side: JoinSide::Right, + } + } + }) + .collect::>(); + + try_pushdown_through_join_with_column_indices( + projection, + join_left, + join_right, + join_on, + schema, + filter, + &column_indices, + ) +} + +/// Attempts to move a projection below a join by mapping each join output +/// column to the child column that produced it. +/// +/// `schema` is the complete output schema of the join, not either child's +/// schema. `column_indices` must contain one entry for each field in `schema`. +/// Each [`JoinSide::Left`] or [`JoinSide::Right`] entry identifies the source +/// child and uses an index relative to that child's schema. +/// +/// [`JoinSide::None`] identifies a column produced by the join itself, such as +/// a mark column. If `projection` references such a column, this function +/// returns `Ok(None)` because neither child can produce it. +/// +/// Returns `Ok(None)` when the projection cannot be pushed down safely. +/// +/// # Errors +/// +/// Returns an error if `column_indices` does not match `schema` or contains an +/// index outside the corresponding child schema. +pub fn try_pushdown_through_join_with_column_indices( + projection: &ProjectionExec, + join_left: &Arc, + join_right: &Arc, + join_on: JoinOnRef, + schema: &SchemaRef, + filter: Option<&JoinFilter>, + column_indices: &[ColumnIndex], +) -> Result> { + if column_indices.len() != schema.fields().len() { + return plan_err!( + "Column index mapping has {} entries but join schema has {} fields", + column_indices.len(), + schema.fields().len() + ); + } + // Validate each output-to-child mapping before using it to rewrite the + // projection. Synthetic outputs have no child index to validate. + for (output_index, column_index) in column_indices.iter().enumerate() { + let (side, child_field_count) = match column_index.side { + JoinSide::Left => ("left", join_left.schema().fields().len()), + JoinSide::Right => ("right", join_right.schema().fields().len()), + JoinSide::None => continue, + }; + if column_index.index >= child_field_count { + return plan_err!( + "Join output column {output_index} maps to {side} child column {}, but the child has {child_field_count} fields", + column_index.index + ); + } + } + // Convert projected expressions to columns. We can not proceed if this is not possible. let Some(projection_as_columns) = physical_to_column_exprs(projection.expr()) else { return Ok(None); }; - let (far_right_left_col_ind, far_left_right_col_ind) = - join_table_borders(join_left.schema().fields().len(), &projection_as_columns); + if projection_as_columns.len() >= schema.fields().len() { + return Ok(None); + } + let mut left_proj: Vec<(Column, String)> = Vec::new(); + let mut right_proj: Vec<(Column, String)> = Vec::new(); + let mut seen_right = false; + for (col, alias) in &projection_as_columns { + let Some(origin) = column_indices.get(col.index()) else { + return plan_err!( + "Projection column {} is outside the {}-entry column index mapping", + col.index(), + column_indices.len() + ); + }; + match origin.side { + // Keep the "left block before right block" contiguity the current + // pushdown supports; a left column after a right one is "mixed". + JoinSide::Left => { + if seen_right { + return Ok(None); + } + left_proj.push((Column::new(col.name(), origin.index), alias.clone())); + } + JoinSide::Right => { + seen_right = true; + right_proj.push((Column::new(col.name(), origin.index), alias.clone())); + } + // Synthetic column (e.g. mark): belongs to neither child. + // Phase 2 declines; Phase 3 keeps it at the join output instead. + JoinSide::None => return Ok(None), + } + } - if !join_allows_pushdown( - &projection_as_columns, - schema, - far_right_left_col_ind, - far_left_right_col_ind, - ) { + // Parity: neither side fully dropped. + if left_proj.is_empty() || right_proj.is_empty() { return Ok(None); } + // `left_proj` / `right_proj` carry *child* indices (from `column_indices`), + // so the shared `update_join_*` helpers must use a 0 column-index offset for + // both sides (the offset bridges child -> join-output index, which is the + // identity here). let new_filter = if let Some(filter) = filter { - match update_join_filter( - &projection_as_columns[0..=far_right_left_col_ind as _], - &projection_as_columns[far_left_right_col_ind as _..], - filter, - join_left.schema().fields().len(), - ) { - Some(updated_filter) => Some(updated_filter), + match update_join_filter(&left_proj, &right_proj, filter, 0) { + Some(updated) => Some(updated), None => return Ok(None), } } else { None }; - let Some(new_on) = update_join_on( - &projection_as_columns[0..=far_right_left_col_ind as _], - &projection_as_columns[far_left_right_col_ind as _..], - join_on, - join_left.schema().fields().len(), - ) else { + let Some(new_on) = update_join_on(&left_proj, &right_proj, join_on, 0) else { return Ok(None); }; - let (new_left, new_right) = new_join_children( - &projection_as_columns, - far_right_left_col_ind, - far_left_right_col_ind, - join_left, - join_right, - )?; + let (new_left, new_right) = + new_join_children_from_groups(&left_proj, &right_proj, join_left, join_right)?; Ok(Some(JoinData { projected_left_child: new_left, @@ -892,6 +1107,34 @@ pub fn new_join_children( Ok((new_left, new_right)) } +/// Build the projected left and right children from side-grouped projection +/// columns whose indices are already *child*-relative (e.g. derived from a +/// join's `ColumnIndex`). Unlike [`new_join_children`], this does not infer +/// child ownership from output position, so it is safe for join schemas whose +/// output is not a plain `left ++ right` (used by the schema-aware +/// `try_pushdown_through_join_with_column_indices`). +fn new_join_children_from_groups( + left_proj: &[(Column, String)], + right_proj: &[(Column, String)], + left_child: &Arc, + right_child: &Arc, +) -> Result<(ProjectionExec, ProjectionExec)> { + let build = |cols: &[(Column, String)], child: &Arc| { + ProjectionExec::try_new( + cols.iter().map(|(col, alias)| ProjectionExpr { + expr: Arc::new(Column::new(col.name(), col.index())) as _, + alias: alias.clone(), + }), + Arc::clone(child), + ) + }; + + Ok(( + build(left_proj, left_child)?, + build(right_proj, right_child)?, + )) +} + /// Checks three conditions for pushing a projection down through a join: /// - Projection must narrow the join output schema. /// - Columns coming from left/right tables must be collected at the left/right @@ -958,14 +1201,10 @@ pub fn update_join_on( .map(|(left, right)| (left, right)) .unzip(); - let new_left_columns = new_columns_for_join_on(&left_idx, proj_left_exprs, 0); - let new_right_columns = - new_columns_for_join_on(&right_idx, proj_right_exprs, left_field_size); - - match (new_left_columns, new_right_columns) { - (Some(left), Some(right)) => Some(left.into_iter().zip(right).collect()), - _ => None, - } + let new_left = new_columns_for_join_on(&left_idx, proj_left_exprs, 0)?; + let new_right = + new_columns_for_join_on(&right_idx, proj_right_exprs, left_field_size)?; + Some(new_left.into_iter().zip(new_right).collect()) } /// Tries to update the column indices of a [`JoinFilter`] as if the input of @@ -1014,55 +1253,69 @@ pub fn update_join_filter( }) } -/// Unifies `projection` with its input (which is also a [`ProjectionExec`]). -fn try_unifying_projections( - projection: &ProjectionExec, - child: &ProjectionExec, +/// Collapse a chain of consecutive [`ProjectionExec`]s into one. Returns +/// `None` if nothing could be merged. +fn try_collapse_projection_chain( + outer: &ProjectionExec, ) -> Result>> { - let mut projected_exprs = vec![]; + let mut current_exprs: Vec = outer.expr().to_vec(); + let mut current_input: Arc = Arc::clone(outer.input()); let mut column_ref_map: HashMap = HashMap::new(); + let mut collapsed_any = false; + + 'outer: while let Some(inner_proj) = current_input.downcast_ref::() { + // Collect the column references usage in the outer projection. + column_ref_map.clear(); + for proj_expr in ¤t_exprs { + proj_expr.expr.apply(|expr| { + if let Some(column) = expr.downcast_ref::() { + *column_ref_map.entry(column.clone()).or_default() += 1; + } + Ok(TreeNodeRecursion::Continue) + })?; + } + let inner_exprs = inner_proj.expr(); + // Merging these projections is not beneficial, e.g + // If an expression is not trivial (KeepInPlace) and it is referred more than 1, unifies projections will be + // beneficial as caching mechanism for non-trivial computations. + // See discussion in: https://github.com/apache/datafusion/issues/8296 + let blocked = column_ref_map.iter().any(|(column, count)| { + *count > 1 + && !inner_exprs[column.index()] + .expr + .placement() + .should_push_to_leaves() + }); + if blocked { + break; + } - // Collect the column references usage in the outer projection. - projection.expr().iter().for_each(|proj_expr| { - proj_expr - .expr - .apply(|expr| { - Ok({ - if let Some(column) = expr.downcast_ref::() { - *column_ref_map.entry(column.clone()).or_default() += 1; - } - TreeNodeRecursion::Continue - }) - }) - .unwrap(); - }); - // Merging these projections is not beneficial, e.g - // If an expression is not trivial (KeepInPlace) and it is referred more than 1, unifies projections will be - // beneficial as caching mechanism for non-trivial computations. - // See discussion in: https://github.com/apache/datafusion/issues/8296 - if column_ref_map.iter().any(|(column, count)| { - *count > 1 - && !child.expr()[column.index()] - .expr - .placement() - .should_push_to_leaves() - }) { - return Ok(None); + let mut new_phys: Vec> = + Vec::with_capacity(current_exprs.len()); + for proj_expr in ¤t_exprs { + // If there is no match in the input projection, we cannot unify these + // projections. This case will arise if the projection expression contains + // a `PhysicalExpr` variant `update_expr` doesn't support. + let Some(expr) = update_expr(&proj_expr.expr, inner_exprs, true)? else { + break 'outer; + }; + new_phys.push(expr); + } + for (proj_expr, expr) in current_exprs.iter_mut().zip(new_phys) { + proj_expr.expr = expr; + } + current_input = Arc::clone(inner_proj.input()); + collapsed_any = true; } - for proj_expr in projection.expr() { - // If there is no match in the input projection, we cannot unify these - // projections. This case will arise if the projection expression contains - // a `PhysicalExpr` variant `update_expr` doesn't support. - let Some(expr) = update_expr(&proj_expr.expr, child.expr(), true)? else { - return Ok(None); - }; - projected_exprs.push(ProjectionExpr { - expr, - alias: proj_expr.alias.clone(), - }); + + if !collapsed_any { + return Ok(None); } - ProjectionExec::try_new(projected_exprs, Arc::clone(child.input())) - .map(|e| Some(Arc::new(e) as _)) + + // To unify 3 or more sequential projections: + let unified: Arc = + Arc::new(ProjectionExec::try_new(current_exprs, current_input)?); + remove_unnecessary_projections(unified).data().map(Some) } /// Collect all column indices from the given projection expressions. @@ -1182,8 +1435,10 @@ mod tests { use super::*; use crate::common::collect; + use crate::empty::EmptyExec; use crate::filter_pushdown::PushedDown; + use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test; use crate::test::exec::StatisticsExec; @@ -1196,6 +1451,46 @@ mod tests { BinaryExpr, Column, DynamicFilterPhysicalExpr, Literal, binary, col, lit, }; + #[test] + fn test_try_new_with_schema_metadata_only_replaces_metadata() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "input", + DataType::Int32, + false, + )])); + let input: Arc = Arc::new(EmptyExec::new(input_schema)); + let field_metadata = + HashMap::from([("field-key".to_string(), "field-value".to_string())]); + let schema_metadata = + HashMap::from([("schema-key".to_string(), "schema-value".to_string())]); + let metadata_schema = Schema::new_with_metadata( + vec![ + Field::new("ignored", DataType::Utf8, true) + .with_metadata(field_metadata.clone()), + ], + schema_metadata.clone(), + ); + + let projection = ProjectionExec::try_new_with_schema_metadata( + [ProjectionExpr { + expr: Arc::new(Column::new("input", 0)), + alias: "output".to_string(), + }], + input, + &metadata_schema, + )?; + + let expected_schema = Arc::new(Schema::new_with_metadata( + vec![ + Field::new("output", DataType::Int32, false) + .with_metadata(field_metadata), + ], + schema_metadata, + )); + assert_eq!(projection.schema(), expected_schema); + Ok(()) + } + #[test] fn test_collect_column_indices() -> Result<()> { let expr = Arc::new(BinaryExpr::new( @@ -1216,6 +1511,113 @@ mod tests { Ok(()) } + #[test] + fn test_try_pushdown_through_join_validates_column_indices() -> Result<()> { + let child_schema = + Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let left: Arc = + Arc::new(EmptyExec::new(Arc::clone(&child_schema))); + let right: Arc = Arc::new(EmptyExec::new(child_schema)); + let join_schema = Arc::new(Schema::new(vec![ + Field::new("left_i", DataType::Int32, false), + Field::new("right_i", DataType::Int32, false), + ])); + let join: Arc = + Arc::new(EmptyExec::new(Arc::clone(&join_schema))); + let projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(Column::new("left_i", 0)), + alias: "left_i".to_string(), + }], + join, + )?; + + let Err(error) = try_pushdown_through_join_with_column_indices( + &projection, + &left, + &right, + &[], + &join_schema, + None, + &[], + ) else { + panic!("expected a mismatched mapping length to return an error"); + }; + assert!( + error.to_string().contains( + "Column index mapping has 0 entries but join schema has 2 fields" + ) + ); + + let invalid_child_index = [ + ColumnIndex { + index: 1, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let Err(error) = try_pushdown_through_join_with_column_indices( + &projection, + &left, + &right, + &[], + &join_schema, + None, + &invalid_child_index, + ) else { + panic!("expected an invalid child index to return an error"); + }; + assert!(error.to_string().contains( + "Join output column 0 maps to left child column 1, but the child has 1 fields" + )); + + let wider_join_schema = Arc::new(Schema::new(vec![ + Field::new("left_i", DataType::Int32, false), + Field::new("right_i", DataType::Int32, false), + Field::new("extra", DataType::Int32, false), + ])); + let wider_join: Arc = + Arc::new(EmptyExec::new(wider_join_schema)); + let out_of_mapping_projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(Column::new("extra", 2)), + alias: "extra".to_string(), + }], + wider_join, + )?; + let valid_child_indices = [ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ]; + let Err(error) = try_pushdown_through_join_with_column_indices( + &out_of_mapping_projection, + &left, + &right, + &[], + &join_schema, + None, + &valid_child_indices, + ) else { + panic!("expected an out-of-mapping projection to return an error"); + }; + assert!( + error.to_string().contains( + "Projection column 2 is outside the 2-entry column index mapping" + ) + ); + + Ok(()) + } + #[test] fn test_join_table_borders() -> Result<()> { let projections = vec![ @@ -1374,7 +1776,9 @@ mod tests { let projection = ProjectionExec::try_new(exprs, input).unwrap(); - let stats = projection.partition_statistics(None).unwrap(); + let stats = StatisticsContext::new() + .compute(&projection, &StatisticsArgs::new()) + .unwrap(); assert_eq!(stats.num_rows, Precision::Exact(10)); assert_eq!( diff --git a/datafusion/physical-plan/src/proto.rs b/datafusion/physical-plan/src/proto.rs new file mode 100644 index 0000000000000..7640d76c3e010 --- /dev/null +++ b/datafusion/physical-plan/src/proto.rs @@ -0,0 +1,386 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Serialization hooks for [`ExecutionPlan`], mirroring the +//! `try_to_proto`/`try_from_proto` pattern used for `PhysicalExpr`. +//! +//! # Why the indirection +//! +//! An `ExecutionPlan` must be able to (de)serialize its child plans and its +//! child physical expressions recursively. The concrete recursion lives in +//! `datafusion-proto` (it owns the extension codec, the session context and the +//! central converter), but `datafusion-proto` sits *above* `datafusion-physical-plan` +//! in the crate graph. To let a plan drive that recursion without a dependency +//! cycle, this module defines: +//! +//! * [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] — the stable, +//! concrete context types a plan author interacts with. New capabilities can +//! be added here without changing every plan's hook signature. +//! * [`ExecutionPlanEncode`] / [`ExecutionPlanDecode`] — internal dispatch +//! traits, *defined* here but *implemented* in `datafusion-proto`, that the +//! context types delegate to. This is the dependency inversion that keeps the +//! proto types flowing in one direction only. They are `#[doc(hidden)]`: not +//! public API, `pub` only because their implementors live in another crate. +//! +//! `datafusion-physical-plan` depends on the pure prost types in +//! `datafusion-proto-models` (feature `proto`), never on `datafusion-proto`. +//! +//! # Function-carrying plans +//! +//! Plans that reference UD(A/W)Fs (`AggregateExec`, the window execs, …) also +//! ride the hook: the context exposes typed, *bytes-only* function serde — +//! [`encode_udaf`](ExecutionPlanEncodeCtx::encode_udaf) / +//! [`decode_udaf`](ExecutionPlanDecodeCtx::decode_udaf) and the udf/udwf +//! siblings. These take/return `datafusion-expr` types plus `Vec` and never +//! name a proto type, so the `PhysicalExtensionCodec` (which only +//! `datafusion-proto` can name) stays fully encapsulated behind the adapter that +//! backs these traits. The lookup-order policy (payload → codec; else registry → +//! codec fallback) lives once, in that adapter, rather than in every plan. +//! +//! This is possible because `datafusion-physical-plan` sits *above* +//! `datafusion-expr` in the crate graph; the expression-side ctx (in +//! `physical-expr-common`, *below* `datafusion-expr`) cannot do this, which is +//! why `ScalarFunctionExpr` remains special-cased there. +//! +//! [`ExecutionPlan`]: crate::ExecutionPlan + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{Result, internal_datafusion_err}; +use datafusion_execution::TaskContext; +use datafusion_expr::physical_planning_context::ScalarSubqueryResults; +use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr_common::physical_expr::proto_decode::{ + PhysicalExprDecode, PhysicalExprDecodeCtx, +}; +use datafusion_physical_expr_common::physical_expr::proto_encode::{ + PhysicalExprEncode, PhysicalExprEncodeCtx, +}; +use datafusion_proto_models::protobuf::{PhysicalExprNode, PhysicalPlanNode}; + +use crate::ExecutionPlan; + +/// Internal dispatch trait backing [`ExecutionPlanEncodeCtx`]. +/// +/// Implemented by `datafusion-proto`. Plan authors never name this trait; they +/// call methods on [`ExecutionPlanEncodeCtx`] instead. +/// +/// **Not public API.** `pub` only because the implementors live in another +/// crate; `#[doc(hidden)]` records that, so encoding primitives can be added +/// here as the serialization hooks grow without breaking downstream code. +#[doc(hidden)] +pub trait ExecutionPlanEncode { + /// Serialize a child execution plan (recursing through the central + /// serializer, so the child's own `try_to_proto` hook is honored). + fn encode_plan(&self, plan: &Arc) -> Result; + + /// Serialize a physical expression owned by the plan. + fn encode_expr(&self, expr: &Arc) -> Result; + + /// Serialize a scalar UDF to an opaque payload. `None` means "decodable by + /// name alone" (built-ins). Bytes-only: no proto types cross this boundary. + fn encode_udf(&self, udf: &ScalarUDF) -> Result>>; + + /// Serialize an aggregate UDF to an opaque payload. `None` means "decodable + /// by name alone". + fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>>; + + /// Serialize a window UDF to an opaque payload. `None` means "decodable by + /// name alone". + fn encode_udwf(&self, udwf: &WindowUDF) -> Result>>; +} + +/// Internal dispatch trait backing [`ExecutionPlanDecodeCtx`]. +/// +/// Implemented by `datafusion-proto`. Plan authors never name this trait; they +/// call methods on [`ExecutionPlanDecodeCtx`] instead. +/// +/// **Not public API.** `pub` only because the implementors live in another +/// crate; `#[doc(hidden)]` records that, so decoding primitives can be added +/// here as the serialization hooks grow without breaking downstream code. +#[doc(hidden)] +pub trait ExecutionPlanDecode { + /// Deserialize a child execution plan (recursing through the central + /// deserializer, so the child's own `try_from_proto` is honored). + fn decode_plan(&self, node: &PhysicalPlanNode) -> Result>; + + /// Deserialize a child plan with `results` active for scalar subquery + /// expressions in that plan's subtree. + fn decode_plan_with_scalar_subquery_results( + &self, + node: &PhysicalPlanNode, + results: ScalarSubqueryResults, + ) -> Result>; + + /// Deserialize a physical expression against `input_schema`. + fn decode_expr( + &self, + node: &PhysicalExprNode, + input_schema: &Schema, + ) -> Result>; + + /// The session task context, used by plans that need the function registry + /// or session configuration. Never exposes the proto extension codec. + fn task_ctx(&self) -> &TaskContext; + + /// Reconstruct a scalar UDF from its name and optional payload. Encapsulates + /// the lookup-order policy (payload → codec; else registry → codec fallback) + /// so no plan re-derives it. Bytes-only: no proto types cross this boundary. + fn decode_udf(&self, name: &str, payload: Option<&[u8]>) -> Result>; + + /// Reconstruct an aggregate UDF from its name and optional payload. + fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result>; + + /// Reconstruct a window UDF from its name and optional payload. + fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result>; +} + +/// Context handed to [`ExecutionPlan::try_to_proto`]. +/// +/// +/// Provides the primitives a plan needs to serialize its children and +/// expressions without naming `datafusion-proto`. +pub struct ExecutionPlanEncodeCtx<'a> { + encoder: &'a dyn ExecutionPlanEncode, +} + +impl<'a> ExecutionPlanEncodeCtx<'a> { + /// Create a new encode context wrapping an [`ExecutionPlanEncode`] + /// implementation (supplied by `datafusion-proto`). + pub fn new(encoder: &'a dyn ExecutionPlanEncode) -> Self { + Self { encoder } + } + + /// Serialize a single child plan. + pub fn encode_child( + &self, + plan: &Arc, + ) -> Result { + self.encoder.encode_plan(plan) + } + + /// Serialize an iterator of child plans. + pub fn encode_children<'b, I>(&self, plans: I) -> Result> + where + I: IntoIterator>, + { + plans.into_iter().map(|p| self.encode_child(p)).collect() + } + + /// Serialize a single physical expression. + pub fn encode_expr(&self, expr: &Arc) -> Result { + self.encoder.encode_expr(expr) + } + + /// Serialize an iterator of physical expressions. + pub fn encode_expressions<'b, I>(&self, exprs: I) -> Result> + where + I: IntoIterator>, + { + exprs.into_iter().map(|e| self.encode_expr(e)).collect() + } + + /// Serialize a scalar UDF to an opaque payload (`None` = built-in, decodable + /// by name). No proto types cross this boundary. + pub fn encode_udf(&self, udf: &ScalarUDF) -> Result>> { + self.encoder.encode_udf(udf) + } + + /// Serialize an aggregate UDF to an opaque payload (`None` = decodable by + /// name). + pub fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>> { + self.encoder.encode_udaf(udaf) + } + + /// Serialize a window UDF to an opaque payload (`None` = decodable by name). + pub fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { + self.encoder.encode_udwf(udwf) + } + + /// An expression-level encode context backed by this plan context. + /// + /// Lets a plan hand `ctx` to expression-level conversions that own their own + /// wire logic — e.g. + /// [`Partitioning::try_to_proto`](datafusion_physical_expr::Partitioning::try_to_proto) + /// and + /// [`PhysicalSortExpr::try_to_proto`](datafusion_physical_expr::PhysicalSortExpr::try_to_proto). + pub fn expr_ctx(&self) -> PhysicalExprEncodeCtx<'_> { + PhysicalExprEncodeCtx::new(self) + } +} + +/// Lets [`ExecutionPlanEncodeCtx`] back a [`PhysicalExprEncodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprEncode for ExecutionPlanEncodeCtx<'_> { + fn encode(&self, expr: &Arc) -> Result { + self.encode_expr(expr) + } +} + +/// Context handed to a plan's `try_from_proto` associated function. +/// +/// Provides the primitives a plan needs to deserialize its children and +/// expressions without naming `datafusion-proto`. +pub struct ExecutionPlanDecodeCtx<'a> { + decoder: &'a dyn ExecutionPlanDecode, +} + +impl<'a> ExecutionPlanDecodeCtx<'a> { + /// Create a new decode context wrapping an [`ExecutionPlanDecode`] + /// implementation (supplied by `datafusion-proto`). + pub fn new(decoder: &'a dyn ExecutionPlanDecode) -> Self { + Self { decoder } + } + + /// Deserialize a single child plan. + pub fn decode_child( + &self, + node: &PhysicalPlanNode, + ) -> Result> { + self.decoder.decode_plan(node) + } + + /// Deserialize a child plan with `results` active for scalar subquery + /// expressions in that plan's subtree. + pub fn decode_child_with_scalar_subquery_results( + &self, + node: &PhysicalPlanNode, + results: ScalarSubqueryResults, + ) -> Result> { + self.decoder + .decode_plan_with_scalar_subquery_results(node, results) + } + + /// Deserialize a required child plan, producing a uniform "missing required + /// field" error when the optional wire field is absent. + pub fn decode_required_child( + &self, + node: Option<&PhysicalPlanNode>, + plan_name: &str, + field: &str, + ) -> Result> { + let node = node.ok_or_else(|| { + internal_datafusion_err!("{plan_name} is missing required field '{field}'") + })?; + self.decode_child(node) + } + + /// Deserialize a physical expression against `input_schema`. + pub fn decode_expr( + &self, + node: &PhysicalExprNode, + input_schema: &Schema, + ) -> Result> { + self.decoder.decode_expr(node, input_schema) + } + + /// Deserialize a required physical expression against `input_schema`. + pub fn decode_required_expr( + &self, + node: Option<&PhysicalExprNode>, + input_schema: &Schema, + plan_name: &str, + field: &str, + ) -> Result> { + let node = node.ok_or_else(|| { + internal_datafusion_err!("{plan_name} is missing required field '{field}'") + })?; + self.decode_expr(node, input_schema) + } + + /// The session task context (function registry + session config). Never + /// exposes the proto extension codec. + pub fn task_ctx(&self) -> &TaskContext { + self.decoder.task_ctx() + } + + /// Reconstruct a scalar UDF from its name and optional payload. The + /// lookup-order policy is owned by `datafusion-proto`; no proto types cross + /// this boundary. + pub fn decode_udf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udf(name, payload) + } + + /// Reconstruct an aggregate UDF from its name and optional payload. + pub fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udaf(name, payload) + } + + /// Reconstruct a window UDF from its name and optional payload. + pub fn decode_udwf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + self.decoder.decode_udwf(name, payload) + } + + /// An expression-level decode context backed by this plan context, bound to + /// `input_schema`. + /// + /// The decode counterpart of + /// [`ExecutionPlanEncodeCtx::expr_ctx`], for calling conversions such as + /// [`Partitioning::try_from_proto`](datafusion_physical_expr::Partitioning::try_from_proto). + pub fn expr_ctx<'s>(&'s self, input_schema: &'s Schema) -> PhysicalExprDecodeCtx<'s> { + PhysicalExprDecodeCtx::new(input_schema, self) + } +} + +/// Lets [`ExecutionPlanDecodeCtx`] back a [`PhysicalExprDecodeCtx`], so +/// expression-level conversions can be reused from plan hooks. +impl PhysicalExprDecode for ExecutionPlanDecodeCtx<'_> { + fn decode( + &self, + node: &PhysicalExprNode, + schema: &Schema, + ) -> Result> { + self.decode_expr(node, schema) + } +} + +/// Assert that a [`PhysicalPlanNode`] carries the expected `PhysicalPlanType` +/// variant, returning a reference to the inner payload, else an `internal_err!`. +/// Mirrors `expect_expr_variant!` on the expression side. Field access on the +/// result auto-derefs through the `Box` that boxed variants use. +#[macro_export] +macro_rules! expect_plan_variant { + ($node:expr, $variant:path, $plan_name:literal $(,)?) => {{ + match &$node.physical_plan_type { + Some($variant(inner)) => inner, + _ => { + return ::datafusion_common::internal_err!(concat!( + "PhysicalPlanNode is not a ", + $plan_name + )); + } + } + }}; +} diff --git a/datafusion/physical-plan/src/recursive_query.rs b/datafusion/physical-plan/src/recursive_query.rs index c160f9a0dc763..0a56488de84dd 100644 --- a/datafusion/physical-plan/src/recursive_query.rs +++ b/datafusion/physical-plan/src/recursive_query.rs @@ -30,16 +30,18 @@ use crate::metrics::{ BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet, RecordOutput, }; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, }; use arrow::array::{BooleanArray, BooleanBuilder}; use arrow::compute::filter_record_batch; -use arrow::datatypes::{Field, Schema, SchemaRef}; +use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::tree_node::{Transformed, TransformedResult, TreeNode}; -use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; +use datafusion_common::{ + Result, exec_datafusion_err, internal_datafusion_err, not_impl_err, +}; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::{MemoryConsumer, MemoryReservation}; use datafusion_physical_expr::PhysicalExpr; @@ -84,6 +86,7 @@ impl RecursiveQueryExec { /// Create a new RecursiveQueryExec pub fn try_new( name: String, + output_schema: SchemaRef, static_term: Arc, recursive_term: Arc, is_distinct: bool, @@ -91,8 +94,6 @@ impl RecursiveQueryExec { // Each recursive query needs its own work table let work_table = Arc::new(WorkTable::new(name.clone())); // Use the same work table for both the WorkTableExec and the recursive term - let output_schema = - recursive_output_schema(&static_term.schema(), &recursive_term.schema()); let static_term = project_plan_to_schema(static_term, &output_schema)?; let recursive_term = assign_work_table(recursive_term, &work_table)?; let recursive_term = project_plan_to_schema(recursive_term, &output_schema)?; @@ -156,7 +157,7 @@ impl ExecutionPlan for RecursiveQueryExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -172,18 +173,24 @@ impl ExecutionPlan for RecursiveQueryExec { } fn required_input_distribution(&self) -> Vec { - vec![ + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ crate::Distribution::SinglePartition, crate::Distribution::SinglePartition, - ] + ]) } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { RecursiveQueryExec::try_new( self.name.clone(), + self.schema(), Arc::clone(&children[0]), Arc::clone(&children[1]), self.is_distinct, @@ -191,6 +198,16 @@ impl ExecutionPlan for RecursiveQueryExec { .map(|e| Arc::new(e) as _) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -370,30 +387,6 @@ impl RecursiveQueryStream { } } -fn recursive_output_schema( - static_schema: &SchemaRef, - recursive_schema: &SchemaRef, -) -> SchemaRef { - let fields = static_schema - .fields() - .iter() - .zip(recursive_schema.fields()) - .map(|(static_field, recursive_field)| { - Field::new( - static_field.name(), - static_field.data_type().clone(), - static_field.is_nullable() || recursive_field.is_nullable(), - ) - .with_metadata(static_field.metadata().clone()) - }) - .collect::>(); - - Arc::new(Schema::new_with_metadata( - fields, - static_schema.metadata().clone(), - )) -} - fn assign_work_table( plan: Arc, work_table: &Arc, @@ -489,7 +482,14 @@ impl DistinctDeduplicator { /// We also detect duplicates by enforcing that group ids are increasing. fn deduplicate(&mut self, batch: &RecordBatch) -> Result { let size_before = self.group_values.len(); - self.intern_output_buffer.reserve(batch.num_rows()); + let additional = batch.num_rows(); + self.intern_output_buffer + .try_reserve(additional) + .map_err(|e| { + exec_datafusion_err!( + "failed to reserve {additional} recursive query group ids: {e}" + ) + })?; self.group_values .intern(batch.columns(), &mut self.intern_output_buffer)?; let mask = new_groups_mask(&self.intern_output_buffer, size_before); @@ -537,6 +537,7 @@ mod tests { let exec = RecursiveQueryExec::try_new( "numbers".to_string(), + static_term.schema(), Arc::clone(&static_term), Arc::clone(&recursive_term), false, @@ -558,9 +559,15 @@ mod tests { let static_term = empty_exec(vec![Field::new("value", DataType::Int32, false)]); let recursive_term = empty_exec(vec![Field::new("value + Int32(1)", DataType::Int32, true)]); + let output_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int32, + true, + )])); let exec = RecursiveQueryExec::try_new( "numbers".to_string(), + Arc::clone(&output_schema), static_term, recursive_term, false, diff --git a/datafusion/physical-plan/src/repartition/mod.rs b/datafusion/physical-plan/src/repartition/mod.rs index 5d87836ba518b..063954a72a094 100644 --- a/datafusion/physical-plan/src/repartition/mod.rs +++ b/datafusion/physical-plan/src/repartition/mod.rs @@ -15,14 +15,15 @@ // specific language governing permissions and limitations // under the License. -//! This file implements the [`RepartitionExec`] operator, which maps N input +//! This file implements the [`RepartitionExec`] operator, which maps N input //! partitions to M output partitions based on a partitioning scheme, optionally //! maintaining the order of the input rows in the output. -use std::fmt::{Debug, Formatter}; +use std::cmp::Ordering; +use std::fmt::{Debug, Display, Formatter}; use std::pin::Pin; use std::sync::Arc; -use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::task::{Context, Poll}; use std::vec; @@ -38,29 +39,41 @@ use crate::metrics::{BaselineMetrics, SpillMetrics}; use crate::projection::{ProjectionExec, all_columns, make_with_child, update_expr}; use crate::sorts::streaming_merge::StreamingMergeBuilder; use crate::spill::spill_manager::SpillManager; -use crate::spill::spill_pool::{self, SpillPoolWriter}; +use crate::spill::spill_pool::{self, SpillPoolSink, SpillPoolWriter}; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::{EmptyRecordBatchStream, RecordBatchStreamAdapter}; use crate::{ - DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, Statistics, - check_if_same_properties, + ChildrenPropertiesMode, DisplayFormatType, ExecutionPlan, Partitioning, + PlanProperties, ReplaceChildrenOptions, Statistics, validate_child_count, }; -use arrow::array::{PrimitiveArray, RecordBatch, RecordBatchOptions}; +use arrow::array::{Array, PrimitiveArray, RecordBatch, RecordBatchOptions, UInt64Array}; use arrow::compute::take_arrays; -use arrow::datatypes::{SchemaRef, UInt32Type}; +use arrow::datatypes::{DataType, Schema, SchemaRef, UInt32Type}; +use arrow_schema::SortOptions; use datafusion_common::config::ConfigOptions; use datafusion_common::stats::Precision; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::utils::transpose; +use datafusion_common::utils::{compare_rows, extract_row_at_idx_to_buf, transpose}; use datafusion_common::{ - ColumnStatistics, DataFusionError, HashMap, assert_or_internal_err, internal_err, + ColumnStatistics, DataFusionError, HashMap, ScalarValue, SplitPoint, + assert_or_internal_err, internal_datafusion_err, internal_err, + validate_range_split_points, }; use datafusion_common::{Result, not_impl_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::TaskContext; use datafusion_execution::memory_pool::MemoryConsumer; -use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr}; +use datafusion_expr::ColumnarValue; +use datafusion_physical_expr::{EquivalenceProperties, PhysicalExpr, RangePartitioning}; +use datafusion_physical_expr_common::physical_expr::PhysicalExprRef; use datafusion_physical_expr_common::sort_expr::LexOrdering; +#[cfg(feature = "proto")] +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; +#[cfg(feature = "proto")] +use datafusion_proto_models::protobuf; use crate::filter_pushdown::{ ChildPushdownResult, FilterDescription, FilterPushdownPhase, @@ -161,10 +174,40 @@ type InputPartitionsToCurrentPartitionReceiver = Vec, reservation: SharedMemoryReservation, - spill_writer: SpillPoolWriter, + spill_writer: SpillPoolSink, shared_coalescer: Option, } +/// The set of spill-pool writers for a single output partition, before they are handed to the +/// per-input tasks. The variant encodes the repartition mode so the wrong writer topology cannot +/// be constructed for a given mode. +enum PartitionSpillWriters { + /// `preserve_order`: one single-producer FIFO writer per input partition. Each is `take`n + /// exactly once (moved into the matching input task), so the pool always has one writer. + PerInput(Vec>), + /// Non-preserve-order: one shared writer, cloned into every input task. + Shared(SpillPoolWriter), +} + +impl PartitionSpillWriters { + /// Hand out the writer for input partition `input`. + /// + /// In `PerInput` mode this moves the dedicated writer out (it must only be requested once per + /// input); in `Shared` mode it clones the shared writer. + fn take_for_input(&mut self, input: usize) -> Result { + match self { + PartitionSpillWriters::PerInput(writers) => { + writers[input].take().ok_or_else(|| { + internal_datafusion_err!( + "spill writer for input partition requested more than once" + ) + }) + } + PartitionSpillWriters::Shared(writer) => Ok(writer.new_sink()), + } + } +} + impl OutputChannel { fn coalesce(&mut self, batch: RecordBatch) -> Result> { match &self.shared_coalescer { @@ -256,7 +299,7 @@ impl SharedCoalescer { /// sender, finalize the coalescer and return its residual batches; if /// other senders are still active, return `Ok(None)`. fn finalize(&self) -> Result> { - let was_last = self.active_senders.fetch_sub(1, Ordering::AcqRel) == 1; + let was_last = self.active_senders.fetch_sub(1, AtomicOrdering::AcqRel) == 1; if !was_last { return Ok(vec![]); } @@ -290,7 +333,7 @@ impl SharedCoalescer { /// /// See [`RepartitionExec`] for the overall N×M architecture. /// -/// [`spill_pool::channel`]: crate::spill::spill_pool::channel +/// [`spill_pool::channel`]: crate::spill::spill_pool::spsc_channel struct PartitionChannels { /// Senders for each input partition to send data to this output partition tx: InputPartitionsToCurrentPartitionSender, @@ -302,9 +345,11 @@ struct PartitionChannels { /// partition. `None` in preserve-order mode (downstream /// `StreamingMergeBuilder` handles batching). shared_coalescer: Option, - /// Spill writers for writing spilled data. - /// SpillPoolWriter is Clone, so multiple writers can share state in non-preserve-order mode. - spill_writers: Vec, + /// Spill writers for writing spilled data, before they are handed to the per-input tasks. + /// The variant is chosen by the repartition mode (see [`PartitionSpillWriters`]): a dedicated + /// single-producer FIFO writer per input in preserve-order mode, or one shared writer in + /// non-preserve-order mode. + spill_writers: PartitionSpillWriters, /// Spill readers for reading spilled data - one per input partition (FIFO semantics). /// Each (input, output) pair gets its own reader to maintain proper ordering. spill_readers: Vec, @@ -420,6 +465,7 @@ impl RepartitionExecState { let num_input_partitions = streams_and_metrics.len(); let num_output_partitions = partitioning.partition_count(); + let coalesce_batches = !preserve_order && !input.boundedness().is_unbounded(); let spill_manager = Arc::new(spill_manager); @@ -460,22 +506,38 @@ impl RepartitionExecState { .session_config() .options() .execution - .max_spill_file_size_bytes; - let num_spill_channels = if preserve_order { - num_input_partitions + .max_spill_file_size_bytes + .get(); + + let (spill_writers, spill_readers) = if preserve_order { + // preserve_order: one dedicated single-producer FIFO pool per input partition. + // Each writer is moved into exactly one input task (never cloned), so the ordering + // the downstream merge relies on is preserved across the spill boundary. + let mut writers = Vec::with_capacity(num_input_partitions); + let mut readers = Vec::with_capacity(num_input_partitions); + for _ in 0..num_input_partitions { + let (writer, reader) = spill_pool::spsc_channel( + max_file_size, + Arc::clone(&spill_manager), + ); + writers.push(Some(writer)); + readers.push(reader); + } + (PartitionSpillWriters::PerInput(writers), readers) } else { - 1 + // non-preserve-order: one shared multi-producer pool per output partition, since + // all inputs share the same receiver and the output is an unordered multiset. + let (writer, reader) = + spill_pool::mpsc_channel(max_file_size, Arc::clone(&spill_manager)); + (PartitionSpillWriters::Shared(writer), vec![reader]) }; - let (spill_writers, spill_readers): (Vec<_>, Vec<_>) = (0 - ..num_spill_channels) - .map(|_| spill_pool::channel(max_file_size, Arc::clone(&spill_manager))) - .unzip(); // Coalesce on the producer side, before the channel's gate, so // the consumer never sees the per-input-task small batches. - // Skip in preserve-order mode: each input has its own dedicated - // channel and `StreamingMergeBuilder` handles batching. - let shared_coalescer = (!preserve_order).then(|| { + // Skip in preserve-order mode, where `StreamingMergeBuilder` + // handles batching, and for unbounded inputs, where a residual + // batch could otherwise be withheld indefinitely. + let shared_coalescer = coalesce_batches.then(|| { SharedCoalescer::new( input.schema(), context.session_config().batch_size(), @@ -502,23 +564,22 @@ impl RepartitionExecState { std::mem::take(streams_and_metrics).into_iter().enumerate() { let txs: HashMap<_, _> = channels - .iter() + .iter_mut() .map(|(partition, channels)| { - // In preserve_order mode: each input gets its own spill writer (index i) - // In non-preserve-order mode: all inputs share spill writer 0 via clone - let spill_writer_idx = if preserve_order { i } else { 0 }; - ( + // Hand this input task its spill writer: in preserve_order mode this moves + // the input's dedicated FIFO writer out; otherwise it clones the shared + // writer. See [`PartitionSpillWriters::take_for_input`]. + Ok(( *partition, OutputChannel { sender: channels.tx[i].clone(), reservation: Arc::clone(&channels.reservation), - spill_writer: channels.spill_writers[spill_writer_idx] - .clone(), + spill_writer: channels.spill_writers.take_for_input(i)?, shared_coalescer: channels.shared_coalescer.clone(), }, - ) + )) }) - .collect(); + .collect::>>()?; // Extract senders for wait_for_task before moving txs let senders: HashMap<_, _> = txs @@ -570,12 +631,237 @@ enum BatchPartitionerState { num_partitions: usize, next_idx: usize, }, + Range { + /// Ordered partitioning key. + ordering: LexOrdering, + /// Sort options from the `LexOrdering` + sort_options: Vec, + /// Boundaries between adjacent partitions. + split_points: Vec, + /// Row indices grouped by output partition + indices: Vec>, + /// Buffer of `ScalarValue` used to represent the values for a row - based on the `LexOrdering` ordering - to compare against split points + partition_buffer: Vec, + }, } /// Fixed RandomState used for hash repartitioning to ensure consistent behavior across /// executions and runs. pub const REPARTITION_RANDOM_STATE: SeededRandomState = SeededRandomState::with_seed(0); +/// Physical expression that returns the Range partition for each input row. +/// +/// This uses the same routing function as [`BatchPartitioner`], so dynamic +/// filtering and repartitioning agree for every [`ScalarValue`] comparison. +#[derive(Debug, Hash, PartialEq, Eq)] +pub struct RangeExpr { + on_columns: Vec, + split_points: Vec, + sort_options: Vec, +} + +impl RangeExpr { + /// Creates a Range expression for `on_columns` using the supplied routing + /// metadata. + pub fn try_new( + on_columns: Vec, + range_partitioning: &RangePartitioning, + ) -> Result { + let sort_options = range_partitioning + .ordering() + .iter() + .map(|expr| expr.options) + .collect(); + Self::try_new_parts( + on_columns, + range_partitioning.split_points().to_vec(), + sort_options, + ) + } + + fn try_new_parts( + on_columns: Vec, + split_points: Vec, + sort_options: Vec, + ) -> Result { + assert_or_internal_err!(!on_columns.is_empty(), "RangeExpr requires a key"); + assert_or_internal_err!( + on_columns.len() == sort_options.len(), + "RangeExpr key count must match sort options" + ); + validate_range_split_points(&split_points, &sort_options)?; + Ok(Self { + on_columns, + split_points, + sort_options, + }) + } + + /// Get the columns used to compute Range partition IDs. + pub fn on_columns(&self) -> &[PhysicalExprRef] { + &self.on_columns + } + + /// Returns the Range split points used for routing. + pub fn split_points(&self) -> &[SplitPoint] { + &self.split_points + } + + /// Returns the per-key sort options used for routing. + pub fn sort_options(&self) -> &[SortOptions] { + &self.sort_options + } +} + +impl Display for RangeExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "range_partition") + } +} + +impl PhysicalExpr for RangeExpr { + fn children(&self) -> Vec<&PhysicalExprRef> { + self.on_columns.iter().collect() + } + + fn with_new_children( + self: Arc, + children: Vec, + ) -> Result { + assert_or_internal_err!( + children.len() == self.on_columns.len(), + "RangeExpr expected {} children, got {}", + self.on_columns.len(), + children.len() + ); + Ok(Arc::new(Self::try_new_parts( + children, + self.split_points.clone(), + self.sort_options.clone(), + )?)) + } + + fn data_type(&self, _input_schema: &Schema) -> Result { + Ok(DataType::UInt64) + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + Ok(false) + } + + fn evaluate(&self, batch: &RecordBatch) -> Result { + let arrays = evaluate_expressions_to_arrays(self.on_columns.iter(), batch)?; + let mut row_key_buffer = Vec::with_capacity(arrays.len()); + let mut partition_ids = Vec::with_capacity(batch.num_rows()); + for row_idx in 0..batch.num_rows() { + extract_row_at_idx_to_buf(&arrays, row_idx, &mut row_key_buffer)?; + partition_ids.push(range_partition_id( + &row_key_buffer, + &self.split_points, + &self.sort_options, + )? as u64); + } + Ok(ColumnarValue::Array(Arc::new(UInt64Array::from( + partition_ids, + )))) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "range_partition") + } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx<'_>, + ) -> Result> { + // Encode the raw ordered children: rebuilding a `LexOrdering` would + // deduplicate equivalent children after dynamic-filter remapping. + let sort_exprs = self + .on_columns + .iter() + .zip(&self.sort_options) + .map(|(expr, options)| PhysicalSortExpr::new(Arc::clone(expr), *options)) + .collect::>(); + let sort_expr = sort_exprs_try_to_proto(&sort_exprs, ctx)?; + let split_point = self + .split_points + .iter() + .map(|split_point| { + let value = split_point + .values() + .iter() + .map(|value| value.try_into().map_err(Into::into)) + .collect::>>()?; + Ok(protobuf::PhysicalRangeSplitPoint { value }) + }) + .collect::>>()?; + Ok(Some(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::RangeExpr( + protobuf::PhysicalRangeExprNode { + sort_expr, + split_point, + }, + )), + })) + } +} + +#[cfg(feature = "proto")] +impl RangeExpr { + /// Reconstructs a [`RangeExpr`] from its protobuf representation. + pub fn try_from_proto( + node: &protobuf::PhysicalExprNode, + ctx: &datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx<'_>, + ) -> Result { + // Decode the raw ordered children for the same reason as `try_to_proto`. + let range_expr = match &node.expr_type { + Some(protobuf::physical_expr_node::ExprType::RangeExpr(expr)) => expr, + _ => return internal_err!("PhysicalExprNode is not a RangeExpr"), + }; + let sort_exprs = sort_exprs_try_from_proto(&range_expr.sort_expr, ctx)?; + let (on_columns, sort_options) = sort_exprs + .into_iter() + .map(|sort_expr| (sort_expr.expr, sort_expr.options)) + .unzip(); + let split_points = range_expr + .split_point + .iter() + .map(|split_point| { + let values = split_point + .value + .iter() + .map(|value| ScalarValue::try_from(value).map_err(Into::into)) + .collect::>>()?; + Ok(SplitPoint::new(values)) + }) + .collect::>>()?; + Ok(Arc::new(Self::try_new_parts( + on_columns, + split_points, + sort_options, + )?)) + } +} + +fn range_partition_id( + row_key: &[ScalarValue], + split_points: &[SplitPoint], + sort_options: &[SortOptions], +) -> Result { + let mut low = 0; + let mut high = split_points.len(); + while low < high { + let mid = low + (high - low) / 2; + match compare_rows(row_key, split_points[mid].values(), sort_options)? { + Ordering::Less => high = mid, + Ordering::Equal | Ordering::Greater => low = mid + 1, + } + } + Ok(low) +} + /// Computes `value % divisor` without division in the hot loop when `divisor` /// is fixed for many values. /// @@ -706,13 +992,40 @@ impl BatchPartitioner { timer, } } + + /// Create a new [`BatchPartitioner`] for range-based repartitioning. + /// + /// # Parameters + /// - `range_partitioning`: `RangePartitioning` struct used for ordering, split points, and number of partitions + /// - `timer`: Metric used to record time spent during repartitioning. + pub fn new_range_partitioner( + range_partitioning: &RangePartitioning, + timer: metrics::Time, + ) -> Self { + let ordering = range_partitioning.ordering().clone(); + let split_points = range_partitioning.split_points().to_vec(); + let num_partitions = range_partitioning.partition_count(); + let sort_options: Vec = ordering.iter().map(|e| e.options).collect(); + + Self { + state: BatchPartitionerState::Range { + partition_buffer: Vec::with_capacity(ordering.len()), + ordering, + sort_options, + split_points, + indices: vec![vec![]; num_partitions], + }, + timer, + } + } + /// Create a new [`BatchPartitioner`] based on the provided [`Partitioning`] scheme. /// /// This is a convenience constructor that delegates to the specialized - /// hash or round-robin constructors depending on the partitioning variant. + /// hash, round-robin, or range constructors depending on the partitioning variant. /// /// # Parameters - /// - `partitioning`: Partitioning scheme to apply (hash or round-robin). + /// - `partitioning`: Partitioning scheme to apply (hash, round-robin, or range). /// - `timer`: Metric used to record time spent during repartitioning. /// - `input_partition`: Index of the current input partition. /// - `num_input_partitions`: Total number of input partitions. @@ -738,6 +1051,9 @@ impl BatchPartitioner { num_input_partitions, )) } + Partitioning::Range(range_repartitioning) => { + Ok(Self::new_range_partitioner(&range_repartitioning, timer)) + } other => { not_impl_err!("Unsupported repartitioning scheme {other:?}") } @@ -824,22 +1140,82 @@ impl BatchPartitioner { Box::new(partitioned_batches.into_iter()) } + BatchPartitionerState::Range { + ordering, + sort_options, + split_points, + indices, + partition_buffer, + } => { + // Tracking time required for distributing indexes across output partitions + let timer = self.timer.timer(); + if split_points.is_empty() { + timer.done(); + Box::new(std::iter::once(Ok((0, batch)))) + } else { + let arrays = evaluate_expressions_to_arrays( + ordering.iter().map(|e| &e.expr), + &batch, + )?; + + indices.iter_mut().for_each(|v| v.clear()); + + Self::partition_range_indices( + &arrays, + split_points, + sort_options, + partition_buffer, + indices, + )?; + + // Finished building index-arrays for output partitions + timer.done(); + + let partitioned_batches = + Self::partition_grouped_take(&batch, indices, &self.timer)?; + + Box::new(partitioned_batches.into_iter()) + } + } }; Ok(it) } + /// Groups input row indices by range partition. This populates `indices[p]` with the + /// row indices from `arrays` that belong in output partition `p` according to `split_points` and `sort_options`. + fn partition_range_indices( + arrays: &[Arc], + split_points: &[SplitPoint], + sort_options: &[SortOptions], + row_key_buffer: &mut Vec, + indices: &mut [Vec], + ) -> Result<()> { + let num_rows = arrays.first().map(|a| a.len()).unwrap_or(0); + for row_idx in 0..num_rows { + // Note that `extract_row_at_idx_to_buf` clears the `row_key_buffer` on each invocation, creating a new row key for comparison for each row + extract_row_at_idx_to_buf(arrays, row_idx, row_key_buffer)?; + + let partition = + range_partition_id(row_key_buffer, split_points, sort_options)?; + indices[partition].push(row_idx as u32) + } + + Ok(()) + } + // return the number of output partitions fn num_partitions(&self) -> usize { match &self.state { BatchPartitionerState::RoundRobin { num_partitions, .. } => *num_partitions, - BatchPartitionerState::Hash { indices, .. } => indices.len(), + BatchPartitionerState::Hash { indices, .. } + | BatchPartitionerState::Range { indices, .. } => indices.len(), } } - /// Build repartitioned hash output batches using one `take` per input batch. + /// Build repartitioned hash/range output batches using one `take` per input batch. /// - /// The hash router first fills one index vector per output partition. This method + /// The routers first fills one index vector per output partition. This method /// concatenates those index vectors, performs one grouped `take_arrays`, and /// then returns each output partition as a slice of the reordered batch. /// @@ -917,7 +1293,7 @@ impl BatchPartitioner { /// used to get 3 even streams of `RecordBatch`es /// /// -///```text +/// ```text /// ▲ ▲ ▲ /// │ │ │ /// │ │ │ @@ -962,7 +1338,8 @@ impl BatchPartitioner { /// Repartitioning one [`RecordBatch`] implies creating multiple smaller batches, potentially /// as many as the number of output partitions. [`RepartitionExec`] makes sure that the returned /// batches adhere to the configured `datafusion.execution.batch_size` for efficient operations, -/// and for that, it will automatically coalesce batches right after repartitioning. +/// and for that, it will automatically coalesce batches right after repartitioning for bounded +/// inputs. Coalescing is skipped for unbounded inputs so partial batches are emitted promptly. /// /// For this, one shared [`LimitedBatchCoalescer`] per output partition is used: /// @@ -1110,18 +1487,6 @@ impl RepartitionExec { pub fn name(&self) -> &str { "RepartitionExec" } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - state: Default::default(), - ..Self::clone(self) - } - } } impl DisplayAs for RepartitionExec { @@ -1187,32 +1552,62 @@ impl ExecutionPlan for RepartitionExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to hash partition expressions if this is a hash repartition - if let Partitioning::Hash(exprs, _) = self.partitioning() { - let mut tnr = TreeNodeRecursion::Continue; - for expr in exprs { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - return Ok(tnr); + match self.partitioning() { + Partitioning::Hash(exprs, _) => crate::apply_expression_roots(exprs, f), + Partitioning::Range(range) => crate::apply_expression_roots( + range.ordering().iter().map(|sort_expr| &sort_expr.expr), + f, + ), + _ => Ok(TreeNodeRecursion::Continue), } - Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - let mut repartition = RepartitionExec::try_new( - children.swap_remove(0), - self.partitioning().clone(), - )?; - if self.preserve_order { - repartition = repartition.with_preserve_order(); + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + state: Default::default(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let mut repartition = RepartitionExec::try_new( + children.swap_remove(0), + self.partitioning().clone(), + )?; + if self.preserve_order { + repartition = repartition.with_preserve_order(); + } + Ok(Arc::new(repartition)) + } } - Ok(Arc::new(repartition)) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn benefits_from_input_partitioning(&self) -> Vec { @@ -1307,6 +1702,12 @@ impl ExecutionPlan for RepartitionExec { if preserve_order { // Store streams from all the input partitions: // Each input partition gets its own spill reader to maintain proper FIFO ordering + // + // Pass None for metrics here — these intermediate streams feed into + // StreamingMerge which is the actual output. Only the merge's + // BaselineMetrics should contribute to the operator's reported + // output_rows. Without this, every row would be counted twice + // (once by PerPartitionStream, once by StreamingMerge). let input_streams = rx .into_iter() .zip(spill_readers) @@ -1319,7 +1720,7 @@ impl ExecutionPlan for RepartitionExec { Arc::clone(&reservation), spill_stream, 1, // Each receiver handles one input partition - BaselineMetrics::new(&metrics, partition), + None, )) as SendableRecordBatchStream }) .collect::>(); @@ -1357,7 +1758,7 @@ impl ExecutionPlan for RepartitionExec { reservation, spill_stream, num_input_partitions, - BaselineMetrics::new(&metrics, partition), + Some(BaselineMetrics::new(&metrics, partition)), )) as SendableRecordBatchStream) } }) @@ -1370,21 +1771,27 @@ impl ExecutionPlan for RepartitionExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - if let Some(partition) = partition { - let partition_count = self.partitioning().partition_count(); - if partition_count == 0 { - return Ok(Arc::new(Statistics::new_unknown(&self.schema()))); - } + fn child_stats_requests(&self, _partition: Option) -> Vec { + vec![ChildStats::At(None)] + } + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + if args.partition().is_some() { + let partition_count = self.partitioning().partition_count(); + // `StatisticsContext::compute` validates the partition index against + // this same count before calling, so it is non-zero here; guard + // defensively against a direct call so the division below cannot + // divide by zero assert_or_internal_err!( - partition < partition_count, - "RepartitionExec invalid partition {} (expected less than {})", - partition, - partition_count + partition_count > 0, + "RepartitionExec statistics requested for a partition but the partition count is 0" ); - let mut stats = Arc::unwrap_or_clone(self.input.partition_statistics(None)?); + let mut stats = input_stats[0].as_ref().clone(); // Distribute statistics across partitions stats.num_rows = stats @@ -1407,7 +1814,7 @@ impl ExecutionPlan for RepartitionExec { Ok(Arc::new(stats)) } else { - self.input.partition_statistics(None) + Ok(Arc::clone(&input_stats[0])) } } @@ -1446,6 +1853,30 @@ impl ExecutionPlan for RepartitionExec { } Partitioning::Hash(new_partitions, *size) } + Partitioning::Range(range_partitioning) => { + // Rewrite range key expressions through the projection. + let mut sort_exprs = + Vec::with_capacity(range_partitioning.ordering().len()); + for sort_expr in range_partitioning.ordering() { + let Some(new_expr) = + update_expr(&sort_expr.expr, projection.expr(), false)? + else { + return Ok(None); + }; + sort_exprs.push(PhysicalSortExpr::new(new_expr, sort_expr.options)); + } + + let Some(ordering) = LexOrdering::new(sort_exprs) else { + return internal_err!( + "failed to create LexOrdering for range partitioning" + ); + }; + + Partitioning::Range(RangePartitioning::try_new( + ordering, + range_partitioning.split_points().to_vec(), + )?) + } others => others.clone(), }; @@ -1504,6 +1935,10 @@ impl ExecutionPlan for RepartitionExec { new_properties.partitioning = match new_properties.partitioning { RoundRobinBatch(_) => RoundRobinBatch(target_partitions), Hash(hash, _) => Hash(hash, target_partitions), + Range(_) => { + // Number of partitions is constrained by the split points and cannot be changed + return Ok(None); + } UnknownPartitioning(_) => UnknownPartitioning(target_partitions), }; Ok(Some(Arc::new(Self { @@ -1514,6 +1949,72 @@ impl ExecutionPlan for RepartitionExec { cache: new_properties.into(), }))) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + let input = ctx.encode_child(self.input())?; + + let partitioning = self.partitioning().try_to_proto(&ctx.expr_ctx())?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Repartition(Box::new( + protobuf::RepartitionExecNode { + input: Some(Box::new(input)), + partitioning: Some(partitioning), + preserve_order: self.preserve_order(), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl RepartitionExec { + /// Reconstruct a [`RepartitionExec`] from its protobuf representation. + pub fn try_from_proto( + node: &protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + let repart = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Repartition, + "RepartitionExec", + ); + let input = ctx.decode_required_child( + repart.input.as_deref(), + "RepartitionExec", + "input", + )?; + let input_schema = input.schema(); + + let partitioning = repart + .partitioning + .as_ref() + .map(|partitioning| { + Partitioning::try_from_proto( + partitioning, + &ctx.expr_ctx(input_schema.as_ref()), + ) + }) + .transpose()? + .flatten() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "RepartitionExec is missing required field 'partitioning'" + ) + })?; + + let mut repart_exec = RepartitionExec::try_new(input, partitioning)?; + if repart.preserve_order { + repart_exec = repart_exec.with_preserve_order(); + } + Ok(Arc::new(repart_exec)) + } } impl RepartitionExec { @@ -1617,26 +2118,12 @@ impl RepartitionExec { input_partition: usize, num_input_partitions: usize, ) -> Result<()> { - let mut partitioner = match &partitioning { - Partitioning::Hash(exprs, num_partitions) => { - BatchPartitioner::new_hash_partitioner( - exprs.clone(), - *num_partitions, - metrics.repartition_time.clone(), - )? - } - Partitioning::RoundRobinBatch(num_partitions) => { - BatchPartitioner::new_round_robin_partitioner( - *num_partitions, - metrics.repartition_time.clone(), - input_partition, - num_input_partitions, - ) - } - other => { - return not_impl_err!("Unsupported repartitioning scheme {other:?}"); - } - }; + let mut partitioner = BatchPartitioner::try_new( + partitioning, + metrics.repartition_time.clone(), + input_partition, + num_input_partitions, + )?; // While there are still outputs to send to, keep pulling inputs let mut batches_until_yield = partitioner.num_partitions(); @@ -1836,8 +2323,8 @@ struct PerPartitionStream { /// each sending None when complete. We must wait for all of them. remaining_partitions: usize, - /// Execution metrics - baseline_metrics: BaselineMetrics, + /// Execution metrics (None in preserve-order mode where StreamingMerge owns the metrics) + baseline_metrics: Option, } impl PerPartitionStream { @@ -1848,7 +2335,7 @@ impl PerPartitionStream { reservation: SharedMemoryReservation, spill_stream: SendableRecordBatchStream, num_input_partitions: usize, - baseline_metrics: BaselineMetrics, + baseline_metrics: Option, ) -> Self { Self { schema, @@ -1867,8 +2354,11 @@ impl PerPartitionStream { cx: &mut Context<'_>, ) -> Poll>> { use futures::StreamExt; - let cloned_time = self.baseline_metrics.elapsed_compute().clone(); - let _timer = cloned_time.timer(); + let elapsed = self + .baseline_metrics + .as_ref() + .map(|m| m.elapsed_compute().clone()); + let _timer = elapsed.as_ref().map(|t| t.timer()); loop { match self.state { @@ -1954,7 +2444,11 @@ impl Stream for PerPartitionStream { cx: &mut Context<'_>, ) -> Poll> { let poll = self.poll_next_inner(cx); - self.baseline_metrics.record_poll(poll) + if let Some(metrics) = &self.baseline_metrics { + metrics.record_poll(poll) + } else { + poll + } } } @@ -1970,6 +2464,9 @@ mod tests { use std::collections::HashSet; use super::*; + use crate::empty::EmptyExec; + use crate::projection::ProjectionExpr; + use crate::streaming::{PartitionStream, StreamingTableExec}; use crate::test::TestMemoryExec; use crate::{ test::{ @@ -1984,14 +2481,78 @@ mod tests { use arrow::array::{ArrayRef, StringArray, UInt32Array}; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::cast::as_string_array; + use datafusion_common::ScalarValue; + use datafusion_common::cast::{as_string_array, as_uint32_array}; use datafusion_common::exec_err; use datafusion_common::test_util::batches_to_sort_string; use datafusion_common_runtime::JoinSet; use datafusion_execution::config::SessionConfig; use datafusion_execution::runtime_env::RuntimeEnvBuilder; + use datafusion_physical_expr::{PhysicalSortExpr, RangePartitioning, SplitPoint}; use insta::assert_snapshot; + #[derive(Debug)] + struct UnboundedTestPartition { + schema: SchemaRef, + batch: RecordBatch, + } + + impl PartitionStream for UnboundedTestPartition { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + fn execute(&self, _ctx: Arc) -> SendableRecordBatchStream { + let stream = futures::stream::iter([Ok(self.batch.clone())]) + .chain(futures::stream::pending()); + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + stream, + )) + } + } + + #[test] + fn range_expr_preserves_duplicate_remapped_children() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ])); + let sort_options = [SortOptions::new(false, false), SortOptions::new(true, true)]; + let split_points = vec![SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(20)), + ])]; + let range_partitioning = RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, sort_options[0]), + PhysicalSortExpr::new(col("b", &schema)?, sort_options[1]), + ] + .into(), + split_points.clone(), + )?; + let expr = Arc::new(RangeExpr::try_new( + vec![col("a", &schema)?, col("b", &schema)?], + &range_partitioning, + )?); + let remapped = col("a", &schema)?; + let rewritten = + expr.with_new_children(vec![Arc::clone(&remapped), Arc::clone(&remapped)])?; + + let rewritten = rewritten + .downcast_ref::() + .expect("rewritten expression should remain a RangeExpr"); + assert_eq!(rewritten.on_columns().len(), 2); + assert!(Arc::ptr_eq( + &rewritten.on_columns()[0], + &rewritten.on_columns()[1] + )); + assert_eq!(rewritten.sort_options(), sort_options); + assert_eq!(rewritten.split_points(), split_points); + + Ok(()) + } + #[test] fn strength_reduced_u64_remainder_matches_modulo() { let divisors = [ @@ -2087,7 +2648,7 @@ mod tests { #[tokio::test] async fn one_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition]; @@ -2110,7 +2671,7 @@ mod tests { #[tokio::test] async fn many_to_one_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2127,7 +2688,7 @@ mod tests { #[tokio::test] async fn many_to_many_round_robin() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2148,7 +2709,7 @@ mod tests { #[tokio::test] async fn many_to_many_hash_partition() -> Result<()> { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2171,43 +2732,693 @@ mod tests { } #[tokio::test] - async fn test_repartition_with_coalescing() -> Result<()> { - let schema = test_schema(); - // create 50 batches, each having 8 rows + async fn many_to_many_range_partition() -> Result<()> { + let schema = test_schema(false); let partition = create_vec_batches(50); - let partitions = vec![partition.clone(), partition.clone()]; - let partitioning = Partitioning::RoundRobinBatch(1); + let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; - let session_config = SessionConfig::new().with_batch_size(200); - let task_ctx = TaskContext::default().with_session_config(session_config); - let task_ctx = Arc::new(task_ctx); + // create_batch values are [1, 2, 3, 4, 5, 6, 7, 8]; split at 3 and 6 yields + // 2, 3, and 3 rows per batch respectively + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![3, 6])?; - // create physical plan - let exec = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; - let exec = RepartitionExec::try_new(exec, partitioning)?; + let output_partitions = repartition(&schema, partitions, partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!(300, partition_row_count(&output_partitions[0])); + assert_eq!(450, partition_row_count(&output_partitions[1])); + assert_eq!(450, partition_row_count(&output_partitions[2])); + assert_eq!( + collect_partition_u32_values(&output_partitions[0]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([1, 2]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[1]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([3, 4, 5]) + ); + assert_eq!( + collect_partition_u32_values(&output_partitions[2]) + .into_iter() + .flatten() + .collect::>(), + HashSet::from([6, 7, 8]) + ); - for i in 0..exec.partitioning().partition_count() { - let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; - while let Some(result) = stream.next().await { - let batch = result?; - assert_eq!(200, batch.num_rows()); - } - } Ok(()) } - fn test_schema() -> Arc { - Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) - } + #[tokio::test] + async fn range_repartition_routes_compound_keys() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(UInt32Array::from(vec![5, 10, 10, 10, 10, 15])), + Arc::new(UInt32Array::from(vec![1, 1, 3, 5, 7, 0])), + ], + )?; + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, SortOptions::default()), + PhysicalSortExpr::new(col("b", &schema)?, SortOptions::default()), + ] + .into(), + vec![ + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(1)), + ]), + SplitPoint::new(vec![ + ScalarValue::UInt32(Some(10)), + ScalarValue::UInt32(Some(5)), + ]), + ], + )?); - async fn repartition( - schema: &SchemaRef, - input_partitions: Vec>, - partitioning: Partitioning, - ) -> Result>> { - let task_ctx = Arc::new(TaskContext::default()); - // create physical plan - let exec = + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![(5, 1)], + collect_partition_u32_pairs(&output_partitions[0]) + ); + assert_eq!( + vec![(10, 1), (10, 3)], + collect_partition_u32_pairs(&output_partitions[1]) + ); + assert_eq!( + vec![(10, 5), (10, 7), (15, 0)], + collect_partition_u32_pairs(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_last() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, false), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![None, Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_nulls_asc_nulls_first() -> Result<()> { + let schema = test_schema(true); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![ + None, + Some(5), + Some(10), + Some(15), + ]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(false, true), vec![10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(2, output_partitions.len()); + assert_eq!( + vec![None, Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_asc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::default(), vec![10, 20])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(5)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(10), Some(15)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_rows_desc() -> Result<()> { + let schema = test_schema(false); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(UInt32Array::from(vec![5, 10, 15, 20, 25]))], + )?; + let partitioning = + u32_range_partitioning(&schema, SortOptions::new(true, false), vec![20, 10])?; + + let output_partitions = + repartition(&schema, vec![vec![batch]], partitioning).await?; + + assert_eq!(3, output_partitions.len()); + assert_eq!( + vec![Some(25)], + collect_partition_u32_values(&output_partitions[0]) + ); + assert_eq!( + vec![Some(15), Some(20)], + collect_partition_u32_values(&output_partitions[1]) + ); + assert_eq!( + vec![Some(5), Some(10)], + collect_partition_u32_values(&output_partitions[2]) + ); + + Ok(()) + } + + #[tokio::test] + async fn range_repartition_routes_string_rows() -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let batch = RecordBatch::try_from_iter(vec![( + "my_awesome_field", + Arc::new(StringArray::from(vec!["bar", "baz", "foo", "qux"])) as ArrayRef, + )])?; + + let schema = batch.schema(); + let expr = col("my_awesome_field", &schema)?; + let input = MockExec::new(vec![Ok(batch)], Arc::clone(&schema)); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new_default(expr)].into(), + vec![SplitPoint::new(vec![ScalarValue::Utf8(Some( + "foo".to_string(), + ))])], + )?); + let exec = RepartitionExec::try_new(Arc::new(input), partitioning)?; + + let mut partition_0 = Vec::new(); + let mut stream = exec.execute(0, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + partition_0.push(result?); + } + + let mut partition_1 = Vec::new(); + let mut stream = exec.execute(1, task_ctx)?; + while let Some(result) = stream.next().await { + partition_1.push(result?); + } + + assert_eq!( + vec!["bar", "baz"], + collect_partition_string_values(&partition_0) + ); + assert_eq!( + vec!["foo", "qux"], + collect_partition_string_values(&partition_1) + ); + + Ok(()) + } + + #[test] + fn range_repartition_swaps_with_projection_rewrites_key_index() -> Result<()> { + // Three columns so the projection both narrows the schema (required for + // swap) and moves the range key from @0 to @1. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("region", DataType::Utf8, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload", "id"])?; + + let swapped = repartition + .try_swapping_with_projection(&projection)? + .expect("swap should succeed when projection keeps the range key"); + let swapped_repartition = swapped + .downcast_ref::() + .expect("top node should be RepartitionExec"); + + assert!(swapped_repartition.input().is::()); + let range = expect_range_partitioning(swapped_repartition.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@1 ASC"); + assert_eq!( + range.split_points(), + &[SplitPoint::new(vec![ScalarValue::UInt32(Some(10))])] + ); + + Ok(()) + } + + #[test] + fn range_repartition_does_not_swap_when_projection_drops_key() -> Result<()> { + // Drop a simple range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::UInt32, false), + Field::new("payload", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["payload"])?; + assert!( + repartition + .try_swapping_with_projection(&projection)? + .is_none() + ); + + // Drop part of a compound range key. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::UInt32, false), + Field::new("b", DataType::UInt32, false), + Field::new("c", DataType::UInt32, false), + ])); + let repartition = Arc::new(RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema))), + range_partitioning_on_columns(&schema, &["a", "b"], vec![vec![10, 1]])?, + )?); + let projection = + projection_on_columns(&(Arc::clone(&repartition) as _), &["a", "c"])?; + assert!( + repartition + .try_swapping_with_projection(&projection)? + .is_none() + ); + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_when_maintains_order() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source with preserve_order: Range maintains input order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let repartition = Arc::new( + RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )? + .with_preserve_order(), + ); + assert!(repartition.maintains_input_order()[0]); + + match repartition.try_pushdown_sort(ordering.as_ref())? { + SortOrderPushdownResult::Exact { inner } => { + let pushed = inner + .downcast_ref::() + .expect("pushdown should keep RepartitionExec"); + + assert!(pushed.preserve_order()); + assert!(pushed.maintains_input_order()[0]); + + let range = expect_range_partitioning(pushed.partitioning()); + assert_eq!(range.ordering()[0].to_string(), "id@0 ASC"); + assert_eq!( + inner.properties().output_ordering().map(|o| o.to_string()), + Some(ordering.to_string()), + "pushed repartition output ordering should match the requested sort" + ); + } + other => panic!("expected Exact sort pushdown, got {other:?}"), + } + + Ok(()) + } + + #[test] + fn range_repartition_try_pushdown_sort_unsupported_without_order_maintenance() + -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::UInt32, false)])); + let ordering = LexOrdering::new([PhysicalSortExpr::new( + col("id", &schema)?, + SortOptions::default(), + )]) + .expect("ordering must not be empty"); + + // Multi-partition source without preserve_order: Range does not maintain order. + let source = Arc::new(ExactSortPushdownExec::new( + Arc::clone(&schema), + 2, + ordering.clone(), + )); + let repartition = Arc::new(RepartitionExec::try_new( + source, + range_partitioning_on_columns(&schema, &["id"], vec![vec![10]])?, + )?); + assert!(!repartition.maintains_input_order()[0]); + + assert!(matches!( + repartition.try_pushdown_sort(ordering.as_ref())?, + SortOrderPushdownResult::Unsupported + )); + + Ok(()) + } + + fn range_partitioning_on_columns( + schema: &SchemaRef, + key_columns: &[&str], + split_points: Vec>, + ) -> Result { + let Some(ordering) = LexOrdering::new( + key_columns + .iter() + .map(|name| { + Ok(PhysicalSortExpr::new( + col(name, schema)?, + SortOptions::default(), + )) + }) + .collect::>>()?, + ) else { + return exec_err!("range ordering must not be empty"); + }; + Ok(Partitioning::Range(RangePartitioning::try_new( + ordering, + split_points + .into_iter() + .map(|values| { + SplitPoint::new( + values + .into_iter() + .map(|value| ScalarValue::UInt32(Some(value))) + .collect(), + ) + }) + .collect(), + )?)) + } + + fn projection_on_columns( + input: &Arc, + names: &[&str], + ) -> Result { + let exprs = names + .iter() + .map(|name| { + Ok(ProjectionExpr { + expr: col(name, &input.schema())?, + alias: (*name).to_string(), + }) + }) + .collect::>>()?; + ProjectionExec::try_new(exprs, Arc::clone(input)) + } + + fn expect_range_partitioning(partitioning: &Partitioning) -> &RangePartitioning { + match partitioning { + Partitioning::Range(range) => range, + other => panic!("expected Range partitioning, got {other:?}"), + } + } + + /// Test source that claims Exact support for any sort pushdown request. + #[derive(Debug, Clone)] + struct ExactSortPushdownExec { + cache: Arc, + } + + impl ExactSortPushdownExec { + fn new(schema: SchemaRef, num_partitions: usize, ordering: LexOrdering) -> Self { + use crate::execution_plan::{Boundedness, EmissionType}; + Self { + cache: Arc::new(PlanProperties::new( + EquivalenceProperties::new_with_orderings(schema, [ordering]), + Partitioning::UnknownPartitioning(num_partitions), + EmissionType::Incremental, + Boundedness::Bounded, + )), + } + } + } + + impl DisplayAs for ExactSortPushdownExec { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ExactSortPushdownExec") + } + } + + impl ExecutionPlan for ExactSortPushdownExec { + fn name(&self) -> &str { + "ExactSortPushdownExec" + } + + fn properties(&self) -> &Arc { + &self.cache + } + + fn children(&self) -> Vec<&Arc> { + vec![] + } + + fn apply_expressions( + &self, + _f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + Ok(TreeNodeRecursion::Continue) + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + Ok(self) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + Ok(Box::pin(EmptyRecordBatchStream::new(self.schema()))) + } + + fn try_pushdown_sort( + &self, + _order: &[PhysicalSortExpr], + ) -> Result>> { + Ok(SortOrderPushdownResult::Exact { + inner: Arc::new(self.clone()), + }) + } + } + + #[tokio::test] + async fn test_repartition_with_coalescing() -> Result<()> { + let schema = test_schema(false); + // create 50 batches, each having 8 rows + let partition = create_vec_batches(50); + let partitions = vec![partition.clone(), partition.clone()]; + let partitioning = Partitioning::RoundRobinBatch(1); + + let session_config = SessionConfig::new().with_batch_size(200); + let task_ctx = TaskContext::default().with_session_config(session_config); + let task_ctx = Arc::new(task_ctx); + + // create physical plan + let exec = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let exec = RepartitionExec::try_new(exec, partitioning)?; + + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + let batch = result?; + assert_eq!(200, batch.num_rows()); + } + } + Ok(()) + } + + #[tokio::test] + async fn unbounded_input_emits_before_batch_size() -> Result<()> { + let schema = test_schema(false); + let batch = create_batch(); + let source = Arc::new(StreamingTableExec::try_new( + Arc::clone(&schema), + vec![Arc::new(UnboundedTestPartition { + schema: Arc::clone(&schema), + batch: batch.clone(), + })], + None, + vec![], + true, + None, + )?); + let exec = RepartitionExec::try_new(source, Partitioning::RoundRobinBatch(1))?; + let session_config = SessionConfig::new().with_batch_size(batch.num_rows() * 2); + let task_ctx = + Arc::new(TaskContext::default().with_session_config(session_config)); + + let mut stream = exec.execute(0, task_ctx)?; + let output = + tokio::time::timeout(std::time::Duration::from_secs(5), stream.next()) + .await + .expect("unbounded repartition withheld a partial batch") + .expect("unbounded input ended unexpectedly")?; + + assert_eq!(batch, output); + Ok(()) + } + + fn test_schema(nullable: bool) -> Arc { + Arc::new(Schema::new(vec![Field::new( + "c0", + DataType::UInt32, + nullable, + )])) + } + + fn u32_range_partitioning( + schema: &SchemaRef, + sort_options: SortOptions, + split_values: Vec, + ) -> Result { + let expr = col("c0", schema)?; + Ok(Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new(expr, sort_options)].into(), + split_values + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::UInt32(Some(value))])) + .collect(), + )?)) + } + + fn partition_row_count(batches: &[RecordBatch]) -> usize { + batches.iter().map(|batch| batch.num_rows()).sum() + } + + fn collect_partition_u32_values(batches: &[RecordBatch]) -> Vec> { + batches + .iter() + .flat_map(|batch| { + let array = + as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + (0..array.len()) + .map(|idx| { + if array.is_null(idx) { + None + } else { + Some(array.value(idx)) + } + }) + .collect::>() + }) + .collect() + } + + fn collect_partition_u32_pairs(batches: &[RecordBatch]) -> Vec<(u32, u32)> { + batches + .iter() + .flat_map(|batch| { + let a = as_uint32_array(batch.column(0)).expect("expected UInt32 column"); + let b = as_uint32_array(batch.column(1)).expect("expected UInt32 column"); + (0..a.len()) + .map(|idx| (a.value(idx), b.value(idx))) + .collect::>() + }) + .collect() + } + + fn collect_partition_string_values(batches: &[RecordBatch]) -> Vec<&str> { + batches + .iter() + .flat_map(|batch| { + let array = + as_string_array(batch.column(0)).expect("expected Utf8 column"); + (0..array.len()) + .map(|idx| array.value(idx)) + .collect::>() + }) + .collect() + } + + async fn repartition( + schema: &SchemaRef, + input_partitions: Vec>, + partitioning: Partitioning, + ) -> Result>> { + let task_ctx = Arc::new(TaskContext::default()); + // create physical plan + let exec = TestMemoryExec::try_new_exec(&input_partitions, Arc::clone(schema), None)?; let exec = RepartitionExec::try_new(exec, partitioning)?; @@ -2230,7 +3441,7 @@ mod tests { let handle: SpawnedTask>>> = SpawnedTask::spawn(async move { // define input partitions - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let partitions = vec![partition.clone(), partition.clone(), partition.clone()]; @@ -2590,7 +3801,7 @@ mod tests { #[tokio::test] async fn repartition_with_spilling() -> Result<()> { // Test that repartition successfully spills to disk when memory is constrained - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2652,7 +3863,7 @@ mod tests { #[tokio::test] async fn repartition_with_partial_spilling() -> Result<()> { // Test that repartition can handle partial spilling (some batches in memory, some spilled) - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2722,7 +3933,7 @@ mod tests { #[tokio::test] async fn repartition_without_spilling() -> Result<()> { // Test that repartition does not spill when there's ample memory - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2784,7 +3995,7 @@ mod tests { use datafusion_execution::disk_manager::{DiskManagerBuilder, DiskManagerMode}; // Test that repartition fails with OOM when disk manager is disabled - let schema = test_schema(); + let schema = test_schema(false); let partition = create_vec_batches(50); let input_partitions = vec![partition]; let partitioning = Partitioning::RoundRobinBatch(4); @@ -2827,7 +4038,7 @@ mod tests { /// Create batch fn create_batch() -> RecordBatch { - let schema = test_schema(); + let schema = test_schema(false); RecordBatch::try_new( schema, vec![Arc::new(UInt32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8]))], @@ -2837,7 +4048,7 @@ mod tests { /// Create batches with sequential values for ordering tests fn create_ordered_batches(num_batches: usize) -> Vec { - let schema = test_schema(); + let schema = test_schema(false); (0..num_batches) .map(|i| { let start = (i * 8) as u32; @@ -2858,7 +4069,7 @@ mod tests { // This tests the state machine fix where we must block on spill_stream // when a Spilled marker is received, rather than continuing to poll the channel - let schema = test_schema(); + let schema = test_schema(false); // Create batches with sequential values: batch 0 has [0,1,2,3,4,5,6,7], // batch 1 has [8,9,10,11,12,13,14,15], etc. let partition = create_ordered_batches(20); @@ -2929,14 +4140,14 @@ mod tests { #[cfg(test)] mod test { - use arrow::array::record_batch; - use arrow::compute::SortOptions; - use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common::assert_batches_eq; - use super::*; use crate::test::TestMemoryExec; use crate::union::UnionExec; + use arrow::array::{UInt32Array, record_batch}; + use arrow::compute::SortOptions; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::assert_batches_eq; + use datafusion_common::config::ConfigNonZeroUsize; use datafusion_physical_expr::expressions::col; @@ -3038,9 +4249,9 @@ mod test { let input_partitions = vec![partition1, partition2]; // Set up context with tight memory limit to force spilling - // Sorting needs some non-spillable memory, so 64 bytes should force spilling while still allowing the query to complete + // Sorting needs some non-spillable memory, so 608 bytes should force spilling while still allowing the query to complete let runtime = RuntimeEnvBuilder::default() - .with_memory_limit(64, 1.0) + .with_memory_limit(608, 1.0) .build_arc()?; let task_ctx = TaskContext::default().with_runtime(runtime); @@ -3105,38 +4316,102 @@ mod test { assert_batches_eq!(expected, std::slice::from_ref(batch)); } - // We should have spilled ~ all of the data. - // - We spill data during the repartitioning phase - // - We may also spill during the final merge sort - let all_batches = [batch1, batch2, batch3, batch4, batch5, batch6]; + // We should have spilled let metrics = exec.metrics().unwrap(); assert!( - metrics.spill_count().unwrap() > input_partitions.len(), - "Expected spill_count > {} for order-preserving repartition, but got {:?}", - input_partitions.len(), - metrics.spill_count() + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur for order-preserving repartition at this \ + memory limit. If this fails, the memory limit may need adjustment." ); - assert!( - metrics.spilled_bytes().unwrap() - > all_batches - .iter() - .map(|b| b.get_array_memory_size()) - .sum::(), - "Expected spilled_bytes > {} for order-preserving repartition, got {}", - all_batches - .iter() - .map(|b| b.get_array_memory_size()) - .sum::(), - metrics.spilled_bytes().unwrap() + Ok(()) + } + + /// Regression test for order preservation across spill *file rotation*. + /// + /// A `preserve_order` repartition relies on each per-(input, output) spill pool delivering + /// batches in strict FIFO order (see [`spill_pool::spsc_channel`] / [`SpillPoolSink`]). This uses + /// the same memory profile as [`Self::test_preserve_order_with_spilling`] — which is tuned to + /// force spilling while still completing — but additionally sets `max_spill_file_size_bytes` + /// to 1 so every spilled batch lands in its own file. That exercises the FIFO-across-rotation + /// path: if ordering were lost across rotated files (e.g. by feeding an ordered pool with a + /// shared multi-producer writer), the downstream `StreamingMerge` would emit out-of-order rows + /// and the sortedness assertion below would fail. + #[tokio::test] + async fn test_preserve_order_with_spill_file_rotation() -> Result<()> { + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnvBuilder; + + // Same sorted input as `test_preserve_order_with_spilling`: + // Partition1: [1,3], [5,7], [9,11]; Partition2: [2,4], [6,8], [10,12] + let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap(); + let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap(); + let batch3 = record_batch!(("c0", UInt32, [5, 7])).unwrap(); + let batch4 = record_batch!(("c0", UInt32, [6, 8])).unwrap(); + let batch5 = record_batch!(("c0", UInt32, [9, 11])).unwrap(); + let batch6 = record_batch!(("c0", UInt32, [10, 12])).unwrap(); + let schema = batch1.schema(); + let sort_exprs = LexOrdering::new([PhysicalSortExpr { + expr: col("c0", &schema).unwrap(), + options: SortOptions::default().asc(), + }]) + .unwrap(); + let partition1 = vec![batch1, batch3, batch5]; + let partition2 = vec![batch2, batch4, batch6]; + let input_partitions = vec![partition1, partition2]; + + // Force a new spill file per spilled batch to exercise FIFO across rotation. + let mut session_config = SessionConfig::new(); + session_config + .options_mut() + .execution + .max_spill_file_size_bytes = ConfigNonZeroUsize::try_new(1).unwrap(); + // Same tight limit as `test_preserve_order_with_spilling`: forces spilling while leaving + // the merge enough non-spillable headroom to complete. + let runtime = RuntimeEnvBuilder::default() + .with_memory_limit(608, 1.0) + .build_arc()?; + let task_ctx = Arc::new( + TaskContext::default() + .with_session_config(session_config) + .with_runtime(runtime), ); + + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?; + let exec = Arc::new(TestMemoryExec::update_cache(&Arc::new(exec))); + let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))? + .with_preserve_order(); + + // Each output partition merges sorted substreams, so its rows must be non-decreasing. + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + let mut last: Option = None; + while let Some(result) = stream.next().await { + let batch = result?; + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + for r in 0..col.len() { + let v = col.value(r); + if let Some(prev) = last { + assert!( + prev <= v, + "output partition {i} not sorted: {prev} came before {v}" + ); + } + last = Some(v); + } + } + } + + let metrics = exec.metrics().unwrap(); assert!( - metrics.spilled_rows().unwrap() - >= all_batches.iter().map(|b| b.num_rows()).sum::(), - "Expected spilled_rows > {} for order-preserving repartition, got {}", - all_batches.iter().map(|b| b.num_rows()).sum::(), - metrics.spilled_rows().unwrap() + metrics.spill_count().unwrap() > 0, + "Expected spilling to occur for order-preserving repartition at this \ + memory limit. If this fails, the memory limit may need adjustment." ); - Ok(()) } @@ -3230,6 +4505,40 @@ mod test { Ok(()) } + #[test] + fn test_range_repartitioned_returns_none() -> Result<()> { + let schema = test_schema(); + let source = memory_exec(&schema); + let partitioning = Partitioning::Range(RangePartitioning::try_new( + [PhysicalSortExpr::new( + col("c0", &schema)?, + SortOptions::default(), + )] + .into(), + vec![ + SplitPoint::new(vec![ScalarValue::UInt32(Some(10))]), + SplitPoint::new(vec![ScalarValue::UInt32(Some(20))]), + ], + )?); + let exec = RepartitionExec::try_new(source, partitioning)?; + + let mut expressions = vec![]; + exec.apply_expressions(&mut |expr| { + expressions.push(expr.to_string()); + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(expressions, ["c0@0"]); + + // Range partition count is fixed by split points, so repartitioned() + // cannot change it to an arbitrary target. + let result = exec.repartitioned(10, &Default::default())?; + assert!( + result.is_none(), + "range repartitioning should not support changing partition count" + ); + Ok(()) + } + fn test_schema() -> Arc { Arc::new(Schema::new(vec![Field::new("c0", DataType::UInt32, false)])) } @@ -3257,4 +4566,47 @@ mod test { let exec = Arc::new(exec); Arc::new(TestMemoryExec::update_cache(&exec)) } + + /// preserve_order repartition should not double-count + /// output rows. + #[tokio::test] + async fn test_preserve_order_output_rows_not_double_counted() -> Result<()> { + use datafusion_execution::TaskContext; + + // Two sorted input partitions, 2 rows each (4 total) + let batch1 = record_batch!(("c0", UInt32, [1, 3])).unwrap(); + let batch2 = record_batch!(("c0", UInt32, [2, 4])).unwrap(); + let schema = batch1.schema(); + let sort_exprs = sort_exprs(&schema); + + let input_partitions = vec![vec![batch1], vec![batch2]]; + let exec = TestMemoryExec::try_new(&input_partitions, Arc::clone(&schema), None)? + .try_with_sort_information(vec![sort_exprs.clone(), sort_exprs])?; + let exec = Arc::new(exec); + let exec = Arc::new(TestMemoryExec::update_cache(&exec)); + + let exec = RepartitionExec::try_new(exec, Partitioning::RoundRobinBatch(3))? + .with_preserve_order(); + + let task_ctx = Arc::new(TaskContext::default()); + let mut total_rows = 0; + for i in 0..exec.partitioning().partition_count() { + let mut stream = exec.execute(i, Arc::clone(&task_ctx))?; + while let Some(result) = stream.next().await { + total_rows += result?.num_rows(); + } + } + + assert_eq!(total_rows, 4, "actual rows collected should be 4"); + + let metrics = exec.metrics().unwrap(); + let reported_output_rows = metrics.output_rows().unwrap(); + assert_eq!( + reported_output_rows, total_rows, + "metrics output_rows ({reported_output_rows}) should match \ + actual rows collected ({total_rows}), not double-count" + ); + + Ok(()) + } } diff --git a/datafusion/physical-plan/src/scalar_subquery.rs b/datafusion/physical-plan/src/scalar_subquery.rs index 82421d66dee9e..f2b7c5e0b53e9 100644 --- a/datafusion/physical-plan/src/scalar_subquery.rs +++ b/datafusion/physical-plan/src/scalar_subquery.rs @@ -30,13 +30,17 @@ use std::sync::Arc; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, ScalarValue, Statistics, exec_err, internal_err}; use datafusion_execution::TaskContext; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; use datafusion_physical_expr::PhysicalExpr; use crate::execution_plan::{CardinalityEffect, ExecutionPlan, PlanProperties}; use crate::joins::utils::{OnceAsync, OnceFut}; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::RecordBatchStreamAdapter; -use crate::{DisplayAs, DisplayFormatType, SendableRecordBatchStream}; +use crate::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ReplaceChildrenOptions, + SendableRecordBatchStream, +}; use futures::StreamExt; use futures::TryStreamExt; @@ -163,9 +167,10 @@ impl ExecutionPlan for ScalarSubqueryExec { children } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { // First child is the main input, the rest are subquery plans. let input = children.remove(0); @@ -185,6 +190,16 @@ impl ExecutionPlan for ScalarSubqueryExec { ))) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn reset_state(self: Arc) -> Result> { self.results.clear(); Ok(Arc::new(ScalarSubqueryExec { @@ -203,9 +218,9 @@ impl ExecutionPlan for ScalarSubqueryExec { ) -> Result { let subqueries = self.subqueries.clone(); let results = self.results.clone(); - let subquery_ctx = Arc::clone(&context); + let planning_ctx = Arc::clone(&context); let mut subquery_future = self.subquery_future.try_once(move || { - Ok(async move { execute_subqueries(subqueries, results, subquery_ctx).await }) + Ok(async move { execute_subqueries(subqueries, results, planning_ctx).await }) })?; let input = Arc::clone(&self.input); let schema = self.schema(); @@ -227,7 +242,7 @@ impl ExecutionPlan for ScalarSubqueryExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } @@ -244,13 +259,86 @@ impl ExecutionPlan for ScalarSubqueryExec { vec![false; self.subqueries.len() + 1] } - fn partition_statistics(&self, partition: Option) -> Result> { - self.input.partition_statistics(partition) + fn child_stats_requests(&self, partition: Option) -> Vec { + // Only `self.input` (child 0) is used; the subqueries are skipped. + let mut requests = vec![ChildStats::Skip; 1 + self.subqueries.len()]; + requests[0] = ChildStats::At(partition); + requests + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) } fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::Equal } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let input = ctx.encode_child(self.input())?; + // Subquery indices are positional and recovered during decoding. + let subqueries = + ctx.encode_children(self.subqueries().iter().map(|subquery| &subquery.plan))?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::ScalarSubquery(Box::new( + protobuf::ScalarSubqueryExecNode { + input: Some(Box::new(input)), + subqueries, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl ScalarSubqueryExec { + /// Reconstruct a [`ScalarSubqueryExec`] from its protobuf representation. + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let scalar_subquery = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::ScalarSubquery, + "ScalarSubqueryExec", + ); + let results = ScalarSubqueryResults::new(scalar_subquery.subqueries.len()); + let input_node = scalar_subquery.input.as_deref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "ScalarSubqueryExec is missing required field 'input'" + ) + })?; + // The input's ScalarSubqueryExpr nodes must share this results container. + let input = + ctx.decode_child_with_scalar_subquery_results(input_node, results.clone())?; + let subqueries = scalar_subquery + .subqueries + .iter() + .enumerate() + .map(|(index, plan)| { + Ok(ScalarSubqueryLink { + plan: ctx.decode_child(plan)?, + index: SubqueryIndex::new(index), + }) + }) + .collect::>>()?; + + Ok(Arc::new(Self::new(input, subqueries, results))) + } } /// Wait for the subquery execution future to complete. @@ -378,9 +466,10 @@ mod tests { vec![&self.inner] } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::new(Self::new( children.remove(0), @@ -390,11 +479,21 @@ mod tests { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/sorts/cursor.rs b/datafusion/physical-plan/src/sorts/cursor.rs index 288ec4cee1594..d71eaad663410 100644 --- a/datafusion/physical-plan/src/sorts/cursor.rs +++ b/datafusion/physical-plan/src/sorts/cursor.rs @@ -16,6 +16,7 @@ // under the License. use std::cmp::Ordering; +use std::fmt::Debug; use std::sync::Arc; use arrow::array::{ @@ -32,7 +33,7 @@ use datafusion_execution::memory_pool::MemoryReservation; /// /// This is a trait as there are several specialized implementations, such as for /// single columns or for normalized multi column keys ([`Rows`]) -pub trait CursorValues { +pub trait CursorValues: Debug + Sync + Send { fn len(&self) -> usize; /// Returns true if `l[l_idx] == r[r_idx]` @@ -44,6 +45,14 @@ pub trait CursorValues { /// Returns comparison of `l[l_idx]` and `r[r_idx]` fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering; + + /// Notifies the values that the owning [`Cursor`] moved to `offset` (always + /// `< len()`), so caching implementations can refresh the value(s) read by + /// the hot comparisons. Default no-op (e.g. byte/row cursors don't benefit). + #[inline] + fn set_offset(&mut self, offset: usize) { + let _ = offset; + } } /// A comparable cursor, used by sort operations @@ -68,14 +77,10 @@ pub trait CursorValues { /// │ │ /// │ CursorValues │ /// └───────────────────────┘ +/// ``` /// -/// -/// Store logical rows using -/// one of several formats, -/// with specialized -/// implementations -/// depending on the column -/// types +/// Store logical rows using one of several formats, with specialized +/// implementations depending on the column types #[derive(Debug)] pub struct Cursor { offset: usize, @@ -89,14 +94,22 @@ impl Cursor { } /// Returns true if there are no more rows in this cursor + #[inline] pub fn is_finished(&self) -> bool { self.offset == self.values.len() } /// Advance the cursor, returning the previous row index + #[inline] pub fn advance(&mut self) -> usize { let t = self.offset; self.offset += 1; + // Refresh the cache for the new position. The guard keeps `set_offset` + // in bounds; a finished cursor's stale cache is never read (it is taken + // before the next comparison). + if self.offset < self.values.len() { + self.values.set_offset(self.offset); + } t } @@ -112,6 +125,7 @@ impl Cursor { } impl PartialEq for Cursor { + #[inline] fn eq(&self, other: &Self) -> bool { T::eq(&self.values, self.offset, &other.values, other.offset) } @@ -142,6 +156,7 @@ impl PartialOrd for Cursor { } impl Ord for Cursor { + #[inline] fn cmp(&self, other: &Self) -> Ordering { T::compare(&self.values, self.offset, &other.values, other.offset) } @@ -180,10 +195,14 @@ impl RowValues { } impl CursorValues for RowValues { + #[inline] fn len(&self) -> usize { self.rows.num_rows() } + // No inline hint on purpose: for the heavyweight `Rows` byte comparison the + // compiler's own choice wins — both `#[inline]` and `#[inline(never)]` + // measurably regress the multi-column merge path. fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { l.rows.row(l_idx) == r.rows.row(r_idx) } @@ -209,38 +228,85 @@ impl CursorArray for PrimitiveArray { type Values = PrimitiveValues; fn values(&self) -> Self::Values { - PrimitiveValues(self.values().clone()) + PrimitiveValues::new(self.values().clone()) } } +/// [`CursorValues`] for a primitive column. +/// +/// Caches the value at the current (and previous) offset, refreshed once per +/// [`Cursor::advance`] via [`CursorValues::set_offset`], so the hot loser-tree +/// comparisons read a cached field instead of indexing the buffer each time. #[derive(Debug)] -pub struct PrimitiveValues(ScalarBuffer); +pub struct PrimitiveValues { + values: ScalarBuffer, + /// Cached `values[offset]`. + current: T, + /// Cached `values[offset - 1]` (read by `eq_to_previous`, only past offset 0). + previous: T, + /// Current offset; used only to `debug_assert!` the cache is read in sync. + offset: usize, +} + +impl PrimitiveValues { + fn new(values: ScalarBuffer) -> Self { + // Non-empty in practice; `unwrap_or_default` just avoids a panic. + let first = values.first().copied().unwrap_or_default(); + Self { + values, + current: first, + previous: first, + offset: 0, + } + } +} impl CursorValues for PrimitiveValues { + #[inline(always)] fn len(&self) -> usize { - self.0.len() + self.values.len() } + #[inline(always)] fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { - l.0[l_idx].is_eq(r.0[r_idx]) + // Arbitrary indices (cross-batch comparison), so index directly. + l.values[l_idx].is_eq(r.values[r_idx]) } + #[inline(always)] fn eq_to_previous(cursor: &Self, idx: usize) -> bool { assert!(idx > 0); - cursor.0[idx].is_eq(cursor.0[idx - 1]) + debug_assert_eq!(idx, cursor.offset); + cursor.current.is_eq(cursor.previous) } + #[inline(always)] fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { - l.0[l_idx].compare(r.0[r_idx]) + debug_assert_eq!(l_idx, l.offset); + debug_assert_eq!(r_idx, r.offset); + l.current.compare(r.current) + } + + #[inline(always)] + fn set_offset(&mut self, offset: usize) { + // The caller (`Cursor::advance`) guarantees `offset < len`; inlined, that + // guard dominates the index below so its bounds check is elided — the + // length is checked once per row, not per comparison. The old `current` + // is `values[offset - 1]`, so it becomes `previous`. + self.previous = self.current; + self.current = self.values[offset]; + self.offset = offset; } } +#[derive(Debug)] pub struct ByteArrayValues { offsets: OffsetBuffer, values: Buffer, } impl ByteArrayValues { + #[inline] fn value(&self, idx: usize) -> &[u8] { assert!(idx < self.len()); // Safety: offsets are valid and checked bounds above @@ -253,19 +319,23 @@ impl ByteArrayValues { } impl CursorValues for ByteArrayValues { + #[inline] fn len(&self) -> usize { self.offsets.len() - 1 } + #[inline] fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { l.value(l_idx) == r.value(r_idx) } + #[inline] fn eq_to_previous(cursor: &Self, idx: usize) -> bool { assert!(idx > 0); cursor.value(idx) == cursor.value(idx - 1) } + #[inline] fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { l.value(l_idx).cmp(r.value(r_idx)) } @@ -394,16 +464,19 @@ impl ArrayValues { } } + #[inline(always)] fn is_null(&self, idx: usize) -> bool { (idx < self.null_threshold) == self.options.nulls_first } } impl CursorValues for ArrayValues { + #[inline(always)] fn len(&self) -> usize { self.values.len() } + #[inline(always)] fn eq(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> bool { match (l.is_null(l_idx), r.is_null(r_idx)) { (true, true) => true, @@ -412,15 +485,19 @@ impl CursorValues for ArrayValues { } } + #[inline(always)] fn eq_to_previous(cursor: &Self, idx: usize) -> bool { assert!(idx > 0); match (cursor.is_null(idx), cursor.is_null(idx - 1)) { (true, true) => true, - (false, false) => T::eq(&cursor.values, idx, &cursor.values, idx - 1), + // Delegate to inner `eq_to_previous` so a caching cursor can answer + // without indexing. + (false, false) => T::eq_to_previous(&cursor.values, idx), _ => false, } } + #[inline(always)] fn compare(l: &Self, l_idx: usize, r: &Self, r_idx: usize) -> Ordering { match (l.is_null(l_idx), r.is_null(r_idx)) { (true, true) => Ordering::Equal, @@ -438,6 +515,12 @@ impl CursorValues for ArrayValues { }, } } + + #[inline(always)] + fn set_offset(&mut self, offset: usize) { + // Forward to the wrapped values (e.g. caching `PrimitiveValues`). + self.values.set_offset(offset); + } } #[cfg(test)] @@ -463,7 +546,7 @@ mod tests { let reservation = consumer.register(&memory_pool); let values = ArrayValues { - values: PrimitiveValues(values), + values: PrimitiveValues::new(values), null_threshold, options, _reservation: reservation, diff --git a/datafusion/physical-plan/src/sorts/index.rs b/datafusion/physical-plan/src/sorts/index.rs deleted file mode 100644 index 29441e3f1fc59..0000000000000 --- a/datafusion/physical-plan/src/sorts/index.rs +++ /dev/null @@ -1,61 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -/// A `RowIndex` identifies a specific row in a logical stream. -/// -/// Each stream is identified by an `stream_idx` and is formed from a -/// sequence of RecordBatches batches, each of which is identified by -/// a unique `batch_idx` within that stream. -/// -/// This is used by `SortPreservingMergeStream` to identify which -/// the order of the tuples in the final sorted output stream. -/// -/// ```text -/// ┌────┐ ┌────┐ ┌────┐ RecordBatch -/// │ │ │ │ │ │ -/// │ C1 │ │... │ │ CN │◀─────── (batch_idx = 0) -/// │ │ │ │ │ │ -/// └────┘ └────┘ └────┘ -/// ┌────┐ ┌────┐ ┌────┐ RecordBatch -/// │ │ │ │ │ │ -/// │ C1 │ │... │ │ CN │◀─────── (batch_idx = 1) -/// │ │ │ │ │ │ -/// └────┘ └────┘ └────┘ -/// ┌────┐ -/// │ │ ... -/// │ C1 │ -/// │ │ ┌────┐ RecordBatch -/// └────┘ │ │ -/// │ CN │◀────── (batch_idx = M-1) -/// │ │ -/// └────┘ -/// -///"Stream"s each with Stream N has M -/// a potentially RecordBatches -///different number of -/// RecordBatches -/// ``` -#[derive(Debug, Clone)] -#[deprecated(since = "46.0.0", note = "unused and will be removed in the future")] -pub struct RowIndex { - /// The index of the stream (uniquely identifies the stream) - pub stream_idx: usize, - /// The index of the batch within the stream's VecDequeue. - pub batch_idx: usize, - /// The row index within the batch - pub row_idx: usize, -} diff --git a/datafusion/physical-plan/src/sorts/merge.rs b/datafusion/physical-plan/src/sorts/merge.rs index c29933535adc5..647649038766d 100644 --- a/datafusion/physical-plan/src/sorts/merge.rs +++ b/datafusion/physical-plan/src/sorts/merge.rs @@ -18,21 +18,23 @@ //! Merge that deals with an arbitrary size of streaming inputs. //! This is an order-preserving merge. -use std::pin::Pin; +use std::fmt::Debug; +use std::future::poll_fn; use std::sync::Arc; -use std::task::{Context, Poll, ready}; +use std::task::{Context, Poll}; -use crate::RecordBatchStream; +use crate::SendableRecordBatchStream; use crate::metrics::BaselineMetrics; use crate::sorts::builder::BatchBuilder; use crate::sorts::cursor::{Cursor, CursorValues}; use crate::sorts::stream::PartitionedStream; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; -use datafusion_common::Result; +use datafusion_common::{DataFusionError, Result, assert_or_internal_err, internal_err}; use datafusion_execution::memory_pool::MemoryReservation; - +use datafusion_execution::{TryEmitter, async_try_stream}; use futures::Stream; /// A fallible [`PartitionedStream`] of [`Cursor`] and [`RecordBatch`] @@ -49,18 +51,6 @@ pub(crate) struct SortPreservingMergeStream { /// used to record execution metrics metrics: BaselineMetrics, - /// If the stream has encountered an error or reaches the - /// `fetch` limit. - done: bool, - - /// Whether buffered rows should be drained after `done` is set. - /// - /// This is enabled when we stop because the `fetch` limit has been - /// reached, allowing partial batches left over after overflow handling to - /// be emitted on subsequent polls. It remains disabled for terminal - /// errors so the stream does not yield data after returning `Err`. - drain_in_progress_on_done: bool, - /// A loser tree that always produces the minimum cursor /// /// Node 0 stores the top winner, Nodes 1..num_streams store @@ -93,41 +83,12 @@ pub(crate) struct SortPreservingMergeStream { /// reference: loser_tree: Vec, - /// If the most recently yielded overall winner has been replaced - /// within the loser tree. A value of `false` indicates that the - /// overall winner has been yielded but the loser tree has not - /// been updated - loser_tree_adjusted: bool, - /// Target batch size batch_size: usize, /// Cursors for each input partition. `None` means the input is exhausted cursors: Vec>>, - /// Configuration parameter to enable round-robin selection of tied winners of loser tree. - /// - /// This option controls the tie-breaker strategy and attempts to avoid the - /// issue of unbalanced polling between partitions - /// - /// If `true`, when multiple partitions have the same value, the partition - /// that has the fewest poll counts is selected. This strategy ensures that - /// multiple partitions with the same value are chosen equally, distributing - /// the polling load in a round-robin fashion. This approach balances the - /// workload more effectively across partitions and avoids excessive buffer - /// growth. - /// - /// if `false`, partitions with smaller indices are consistently chosen as - /// the winners, which can lead to an uneven distribution of polling and potentially - /// causing upstream operator buffers for the other partitions to grow - /// excessively, as they continued receiving data without consuming it. - /// - /// For example, an upstream operator like `RepartitionExec` execution would - /// keep sending data to certain partitions, but those partitions wouldn't - /// consume the data if they weren't selected as winners. This resulted in - /// inefficient buffer usage. - enable_round_robin_tie_breaker: bool, - /// Flag indicating whether we are in the mode of round-robin /// tie breaker for the loser tree winners. round_robin_tie_breaker_mode: bool, @@ -142,17 +103,15 @@ pub(crate) struct SortPreservingMergeStream { /// Current reset count current_reset_epoch: usize, - /// Stores the previous value of each partitions for tracking the poll counts on the same value. - prev_cursors: Vec>>, + /// Stores the previous value of each partitions for tracking the poll counts on the same value + /// Used if and only if round robin tie breaker is enabled, otherwise None + prev_cursors: Option>>>, /// Optional number of rows to fetch fetch: Option, /// number of rows produced produced: usize, - - /// This vector contains the indices of the partitions that have not started emitting yet. - uninitiated_partitions: Vec, } impl SortPreservingMergeStream { @@ -165,30 +124,47 @@ impl SortPreservingMergeStream { reservation: MemoryReservation, enable_round_robin_tie_breaker: bool, ) -> Self { + assert_ne!(batch_size, 0, "batch size cannot be 0"); + assert_ne!(fetch, Some(0), "fetch must not be Some(0)"); + let stream_count = streams.partitions(); Self { in_progress: BatchBuilder::new(schema, stream_count, batch_size, reservation), streams, metrics, - done: false, - drain_in_progress_on_done: false, cursors: (0..stream_count).map(|_| None).collect(), - prev_cursors: (0..stream_count).map(|_| None).collect(), + prev_cursors: if enable_round_robin_tie_breaker { + Some((0..stream_count).map(|_| None).collect()) + } else { + None + }, round_robin_tie_breaker_mode: false, num_of_polled_with_same_value: vec![0; stream_count], current_reset_epoch: 0, poll_reset_epochs: vec![0; stream_count], loser_tree: vec![], - loser_tree_adjusted: false, batch_size, fetch, produced: 0, - uninitiated_partitions: (0..stream_count).collect(), - enable_round_robin_tie_breaker, } } + pub(crate) fn into_stream(self) -> SendableRecordBatchStream + where + C: 'static, + { + let schema_clone = Arc::clone(self.in_progress.schema()); + + let cloned_metrics = self.metrics.clone(); + let stream = Box::pin(RecordBatchStreamAdapter::new( + schema_clone, + self.create_stream(), + )); + + Box::pin(ObservedStream::new(stream, cloned_metrics, None)) + } + /// If the stream at the given index is not exhausted, and the last cursor for the /// stream is finished, poll the stream for the next RecordBatch and create a new /// cursor for the stream from the returned result @@ -219,103 +195,172 @@ impl SortPreservingMergeStream { result } - fn poll_next_inner( + async fn flush_in_progress( &mut self, - cx: &mut Context<'_>, - ) -> Poll>> { - if self.done { - // When `build_record_batch()` hits an i32 offset overflow (e.g. - // combined string offsets exceed 2 GB), it emits a partial batch - // and keeps the remaining rows in `self.in_progress.indices`. - // Drain those leftover rows before terminating the stream, - // otherwise they would be silently dropped. - // Repeated overflows are fine — each poll emits another partial - // batch until `in_progress` is fully drained. - if self.drain_in_progress_on_done && !self.in_progress.is_empty() { - return Poll::Ready(self.emit_in_progress_batch().transpose()); - } - return Poll::Ready(None); + mut emitter: TryEmitter, + ) -> Result<()> { + if self.in_progress.is_empty() { + return Ok(()); } - // Once all partitions have set their corresponding cursors for the loser tree, - // we skip the following block. Until then, this function may be called multiple - // times and can return Poll::Pending if any partition returns Poll::Pending. - - if self.loser_tree.is_empty() { - // Manual indexing since we're iterating over the vector and shrinking it in the loop - let mut idx = 0; - while idx < self.uninitiated_partitions.len() { - let partition_idx = self.uninitiated_partitions[idx]; - match self.maybe_poll_stream(cx, partition_idx) { - Poll::Ready(Err(e)) => { - self.done = true; - return Poll::Ready(Some(Err(e))); - } - Poll::Pending => { - // The polled stream is pending which means we're already set up to - // be woken when necessary - // Try the next stream - idx += 1; - } - _ => { - // The polled stream is ready - // Remove it from uninitiated_partitions - // Don't bump idx here, since a new element will have taken its - // place which we'll try in the next loop iteration - // swap_remove will change the partition poll order, but that shouldn't - // make a difference since we're waiting for all streams to be ready. - self.uninitiated_partitions.swap_remove(idx); - } - } - } - if self.uninitiated_partitions.is_empty() { - // If there are no more uninitiated partitions, set up the loser tree and continue - // to the next phase. + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + // When `build_record_batch()` hits an i32 offset overflow (e.g. + // combined string offsets exceed 2 GB), it emits a partial batch + // and keeps the remaining rows in `self.in_progress.indices`. + // Drain those leftover rows before terminating the stream, + // otherwise they would be silently dropped. + // Repeated overflows are fine — each poll emits another partial + // batch until `in_progress` is fully drained. + while let Some(batch) = self.emit_in_progress_batch()? { + drop(timer); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } - // Claim the memory for the uninitiated partitions - self.uninitiated_partitions.shrink_to_fit(); - self.init_loser_tree(); - } else { - // There are still uninitiated partitions so return pending. - // We only get here if we've polled all uninitiated streams and at least one of them - // returned pending itself. That means we will be woken as soon as one of the - // streams would like to be polled again. - // There is no need to reschedule ourselves eagerly. - return Poll::Pending; + Ok(()) + } + + fn create_stream(mut self) -> impl Stream> { + async_try_stream(|mut emitter| async move { + // 1. Make sure we have data from each stream so we can initialize the loser tree + { + // This vector contains the indices of the partitions that have not started emitting yet. + let mut uninitiated_partitions = + (0..self.streams.partitions()).collect::>(); + + poll_fn(|cx| { + self.initialize_all_partitions(&mut uninitiated_partitions, cx) + }) + .await?; + + assert_eq!(uninitiated_partitions.len(), 0); } - } - // NB timer records time taken on drop, so there are no - // calls to `timer.done()` below. - let elapsed_compute = self.metrics.elapsed_compute().clone(); - let _timer = elapsed_compute.timer(); - - loop { - // Adjust the loser tree if necessary, returning control if needed - if !self.loser_tree_adjusted { - let winner = self.loser_tree[0]; - if let Err(e) = ready!(self.maybe_poll_stream(cx, winner)) { - self.done = true; - return Poll::Ready(Some(Err(e))); + let elapsed_compute = self.metrics.elapsed_compute().clone(); + let mut timer = elapsed_compute.timer(); + + // 2. Init loser tree + self.init_loser_tree(); + + // 3. loop until all streams have been exhausted + while !self.is_exhausted() { + // 3.1. add loser_tree[0] (minimum) stream to pending record batch + let winner_stream = self.loser_tree[0]; + self.in_progress.push_row(winner_stream); + + // 3.2. If the new row reached the limit + if self.fetch_reached() { + break; } + + // 3.3. if there is enough to emit for a full record batch + if self.in_progress.len() >= self.batch_size { + // 3.3.1 build pending record batch and reset builder + let Some(batch) = self.emit_in_progress_batch()? else { + return internal_err!("must have batch in progress to emit"); + }; + + // 3.3.2 emit pending record batch + drop(timer); + emitter.emit(batch).await; + timer = elapsed_compute.timer(); + } + + // 3.4. advance cursor for the winner stream + { + let should_poll_next_batch_for_stream = + self.advance_cursors(winner_stream); + + // Fast path: skip the `maybe_poll_stream` call (and its `Poll` + // plumbing) unless the winner's cursor is exhausted and needs a + // fresh batch — it is live for almost every row. + if should_poll_next_batch_for_stream { + assert_or_internal_err!( + self.cursors[winner_stream].is_none(), + "cursor should be exhausted" + ); + + drop(timer); + poll_fn(|cx| self.maybe_poll_stream(cx, winner_stream)).await?; + timer = elapsed_compute.timer(); + } + } + + // 3.5. Adjusting the loser tree if necessary self.update_loser_tree(); } - let stream_idx = self.loser_tree[0]; - if self.advance_cursors(stream_idx) { - self.loser_tree_adjusted = false; - self.in_progress.push_row(stream_idx); + // 4. Flush any remaining rows in `self.in_progress` + self.flush_in_progress(emitter).await?; - // stop sorting if fetch has been reached - if self.fetch_reached() { - self.done = true; - self.drain_in_progress_on_done = true; - } else if self.in_progress.len() < self.batch_size { - continue; + Ok(()) + }) + } + + /// Returns `true` once every input stream is exhausted. + /// + /// Should only be called for valid adjusted tree, i.e. the initial tree or after [`Self::update_loser_tree`] call + fn is_exhausted(&self) -> bool { + let winner = self.loser_tree[0]; + + // Checking only the tree root suffices for valid tree + // since the winner of the tree cannot be an exhausted stream for a valid tree + // as what value is winning over the non exhausted stream? + self.cursors[winner].is_none() + } + + /// Initialize all partitions, return `Poll::Pending` if any partition returns `Poll::Pending` + /// + /// This DOES NOT return `Poll::Pending` as soon as the first uninitiated partition returns `Poll::Pending` + /// so we can continue to initialize the remaining partitions + fn initialize_all_partitions( + &mut self, + uninitiated_partitions: &mut Vec, + cx: &mut Context, + ) -> Poll> { + assert_eq!( + self.loser_tree.len(), + 0, + "loser tree must be empty when initializing" + ); + + // Manual indexing since we're iterating over the vector and shrinking it in the loop + let mut idx = 0; + while idx < uninitiated_partitions.len() { + let partition_idx = uninitiated_partitions[idx]; + match self.maybe_poll_stream(cx, partition_idx) { + Poll::Ready(Err(e)) => { + return Poll::Ready(Err(e)); + } + Poll::Pending => { + // The polled stream is pending which means we're already set up to + // be woken when necessary + // Try the next stream + idx += 1; + } + _ => { + // The polled stream is ready + // Remove it from uninitiated_partitions + // Don't bump idx here, since a new element will have taken its + // place which we'll try in the next loop iteration + // swap_remove will change the partition poll order, but that shouldn't + // make a difference since we're waiting for all streams to be ready. + uninitiated_partitions.swap_remove(idx); } } + } - return Poll::Ready(self.emit_in_progress_batch().transpose()); + if uninitiated_partitions.is_empty() { + Poll::Ready(Ok(())) + } else { + // There are still uninitiated partitions so return pending. + // We only get here if we've polled all uninitiated streams and at least one of them + // returned pending itself. That means we will be woken as soon as one of the + // streams would like to be polled again. + // There is no need to reschedule ourselves eagerly. + Poll::Pending } } @@ -332,7 +377,13 @@ impl SortPreservingMergeStream { if let Some(c) = cursor.as_mut() { // Compare with the last row in the previous batch - let prev_cursor = &self.prev_cursors[partition_idx]; + let prev_cursor = self + .prev_cursors + .as_ref() + .map(|v| &v[partition_idx]) + .expect( + "prev_cursor should be set when round robin tie breaker is enabled", + ); if c.is_eq_to_prev_one(prev_cursor.as_ref()) { self.num_of_polled_with_same_value[partition_idx] += 1; } else { @@ -341,6 +392,31 @@ impl SortPreservingMergeStream { } } + /// Whether round-robin selection of tied winners of loser tree is enabled. + /// + /// This option controls the tie-breaker strategy and attempts to avoid the + /// issue of unbalanced polling between partitions + /// + /// If `true`, when multiple partitions have the same value, the partition + /// that has the fewest poll counts is selected. This strategy ensures that + /// multiple partitions with the same value are chosen equally, distributing + /// the polling load in a round-robin fashion. This approach balances the + /// workload more effectively across partitions and avoids excessive buffer + /// growth. + /// + /// if `false`, partitions with smaller indices are consistently chosen as + /// the winners, which can lead to an uneven distribution of polling and potentially + /// causing upstream operator buffers for the other partitions to grow + /// excessively, as they continued receiving data without consuming it. + /// + /// For example, an upstream operator like `RepartitionExec` execution would + /// keep sending data to certain partitions, but those partitions wouldn't + /// consume the data if they weren't selected as winners. This resulted in + /// inefficient buffer usage. + fn round_robin_tie_breaker_enabled(&self) -> bool { + self.prev_cursors.is_some() + } + fn fetch_reached(&mut self) -> bool { self.fetch .map(|fetch| self.produced + self.in_progress.len() >= fetch) @@ -350,18 +426,23 @@ impl SortPreservingMergeStream { /// Advances the actual cursor. If it reaches its end, update the /// previous cursor with it. /// - /// If the given partition is not exhausted, the function returns `true`. + /// If the given partition batch is exhausted, return `true` to signal a poll is needed fn advance_cursors(&mut self, stream_idx: usize) -> bool { if let Some(cursor) = &mut self.cursors[stream_idx] { let _ = cursor.advance(); - if cursor.is_finished() { + let finished = cursor.is_finished(); + if finished { // Take the current cursor, leaving `None` in its place - self.prev_cursors[stream_idx] = self.cursors[stream_idx].take(); + let taken = self.cursors[stream_idx].take(); + if let Some(prev_cursors) = &mut self.prev_cursors { + prev_cursors[stream_idx] = taken; + } } - true - } else { - false + return finished; } + + // the entire stream is exhausted, so return true (poll won't help here anyway) + true } /// Returns `true` if the cursor at index `a` is greater than at index `b`. @@ -445,7 +526,6 @@ impl SortPreservingMergeStream { } self.loser_tree[cmp_node] = winner; } - self.loser_tree_adjusted = true; } /// Resets the poll count by incrementing the reset epoch. @@ -512,22 +592,33 @@ impl SortPreservingMergeStream { let mut cmp_node = self.lt_leaf_node_index(winner); // Traverse up the tree to adjust comparisons until reaching the root. - while cmp_node != 0 { + while cmp_node > 1 { let challenger = self.loser_tree[cmp_node]; + if self.is_gt(winner, challenger) { + self.update_winner(cmp_node, &mut winner, challenger); + } + cmp_node = self.lt_parent_node_index(cmp_node); + } + + if cmp_node == 1 { + let challenger = self.loser_tree[1]; // If round-robin tie-breaker is enabled and we're at the final comparison (cmp_node == 1) - if self.enable_round_robin_tie_breaker && cmp_node == 1 { + if self.round_robin_tie_breaker_enabled() { match (&self.cursors[winner], &self.cursors[challenger]) { - (Some(ac), Some(bc)) => { - if ac == bc { + (Some(ac), Some(bc)) => match ac.cmp(bc) { + std::cmp::Ordering::Equal => { self.handle_tie(cmp_node, &mut winner, challenger); - } else { + } + std::cmp::Ordering::Greater => { // Ends of tie breaker self.round_robin_tie_breaker_mode = false; - if ac > bc { - self.update_winner(cmp_node, &mut winner, challenger); - } + self.update_winner(cmp_node, &mut winner, challenger); } - } + std::cmp::Ordering::Less => { + // Ends of tie breaker + self.round_robin_tie_breaker_mode = false; + } + }, (None, _) => { // Challenger wins, update winner // Ends of tie breaker @@ -543,28 +634,9 @@ impl SortPreservingMergeStream { } else if self.is_gt(winner, challenger) { self.update_winner(cmp_node, &mut winner, challenger); } - cmp_node = self.lt_parent_node_index(cmp_node); } - self.loser_tree[0] = winner; - self.loser_tree_adjusted = true; - } -} - -impl Stream for SortPreservingMergeStream { - type Item = Result; - fn poll_next( - mut self: Pin<&mut Self>, - cx: &mut Context<'_>, - ) -> Poll> { - let poll = self.poll_next_inner(cx); - self.metrics.record_poll(poll) - } -} - -impl RecordBatchStream for SortPreservingMergeStream { - fn schema(&self) -> SchemaRef { - Arc::clone(self.in_progress.schema()) + self.loser_tree[0] = winner; } } @@ -578,7 +650,7 @@ mod tests { use datafusion_execution::memory_pool::{ MemoryConsumer, MemoryPool, UnboundedMemoryPool, }; - use futures::task::noop_waker_ref; + use futures::TryStreamExt; use std::cmp::Ordering; #[derive(Debug)] @@ -621,8 +693,8 @@ mod tests { } } - #[test] - fn test_done_drains_buffered_rows() { + #[tokio::test] + async fn test_done_drains_buffered_rows() { let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); let pool: Arc = Arc::new(UnboundedMemoryPool::default()); let reservation = MemoryConsumer::new("test").register(&pool); @@ -638,24 +710,20 @@ mod tests { true, ); + // Simulate rows left buffered in `in_progress` (as happens when + // `build_record_batch` emits a partial batch on offset overflow). With + // an empty input stream the merge loop breaks immediately, so the only + // way these rows reach the consumer is the generator's final drain loop. let batch = RecordBatch::try_new(schema, vec![Arc::new(Int32Array::from(vec![1]))]) .unwrap(); stream.in_progress.push_batch(0, batch).unwrap(); stream.in_progress.push_row(0); - stream.done = true; - stream.drain_in_progress_on_done = true; - let waker = noop_waker_ref(); - let mut cx = Context::from_waker(waker); + // Drive the actual stream and confirm the buffered row is drained. + let batches: Vec = stream.into_stream().try_collect().await.unwrap(); - match stream.poll_next_inner(&mut cx) { - Poll::Ready(Some(Ok(batch))) => assert_eq!(batch.num_rows(), 1), - other => { - panic!("expected buffered rows to be drained after done, got {other:?}") - } - } - assert!(stream.in_progress.is_empty()); - assert!(matches!(stream.poll_next_inner(&mut cx), Poll::Ready(None))); + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].num_rows(), 1); } } diff --git a/datafusion/physical-plan/src/sorts/multi_level_merge.rs b/datafusion/physical-plan/src/sorts/multi_level_merge.rs index 8985e1d8c70ee..3ec52cc70c0a9 100644 --- a/datafusion/physical-plan/src/sorts/multi_level_merge.rs +++ b/datafusion/physical-plan/src/sorts/multi_level_merge.rs @@ -27,13 +27,13 @@ use std::sync::Arc; use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; -use datafusion_common::Result; +use datafusion_common::{Result, internal_err, resources_err}; use datafusion_execution::memory_pool::MemoryReservation; use crate::sorts::builder::try_grow_reservation_to_at_least; use crate::sorts::sort::get_reserved_bytes_for_record_batch_size; use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; -use crate::stream::RecordBatchStreamAdapter; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::TryStreamExt; @@ -119,17 +119,35 @@ use futures::{Stream, StreamExt}; /// ## Memory Management Strategy /// /// This multi-level merge make sure that we can handle any amount of data to sort as long as -/// we have enough memory to merge at least 2 streams at a time. +/// we have enough memory to merge at least 2 streams at a time, even when individual record +/// batches are skewed (very wide). /// /// 1. **Worst-Case Memory Reservation**: Reserves memory based on the largest /// batch size encountered in each spill file to merge, ensuring sufficient memory is always /// available during merge operations. /// 2. **Adaptive Buffer Sizing**: Reduces buffer sizes when memory is constrained /// 3. **Spill-to-Disk**: Spill to disk when we cannot merge all files in memory +/// 4. **Re-spilling Skewed Runs**: If even at the smallest read-buffer size we still cannot +/// reserve memory for the minimum of 2 streams - because a single run's largest batch is so +/// wide that two streams' worth of reservation exceeds the budget - the larger of the two +/// runs is re-spilled with each batch sliced in half. This shrinks its largest batch, +/// lowering the per-stream reservation, and the merge pass is retried. The re-spilled run +/// is tracked alongside a per-run batch-size limit equal to half the batch size it was +/// written with, so any later merge that includes it caps its output batch size to match - +/// otherwise the merged run could rebuild a full-size batch and reintroduce the skew. +/// Crucially the global merge batch size is *not* lowered, so re-spilling more than one run +/// does not compound the reduction. If a batch cannot be split any further (a single row +/// wider than the budget), the merge surfaces `ResourcesExhausted` instead of looping +/// forever. pub(crate) struct MultiLevelMergeBuilder { spill_manager: SpillManager, schema: SchemaRef, - sorted_spill_files: Vec, + /// Sorted runs still to be merged. Each run is paired with the batch-size limit a + /// merge consuming it must cap its output at. Runs written at the full batch size + /// carry `batch_size`. A run re-spilled smaller to resolve skew carries its halved + /// limit (see [`Self::split_spill_file_in_half`]). Tracking it here keeps this limit + /// out of the public [`SortedSpillFile`], so no external caller has to set it. + sorted_spill_files: Vec<(SortedSpillFile, usize)>, sorted_streams: Vec, expr: LexOrdering, metrics: BaselineMetrics, @@ -162,7 +180,12 @@ impl MultiLevelMergeBuilder { Self { spill_manager, schema, - sorted_spill_files, + // Initial runs are written at the full batch size, so they impose no cap + // on later merges - record `batch_size` as their (unconstrained) limit. + sorted_spill_files: sorted_spill_files + .into_iter() + .map(|file| (file, batch_size)) + .collect(), sorted_streams, expr, metrics, @@ -182,7 +205,22 @@ impl MultiLevelMergeBuilder { async fn create_stream(mut self) -> Result { loop { - let mut stream = self.merge_sorted_runs_within_mem_limit()?; + let (mut stream, batch_size_limit) = + match self.merge_sorted_runs_within_mem_limit()? { + MergeStep::Stream { + stream, + batch_size_limit, + } => (stream, batch_size_limit), + MergeStep::SplitThenRetry(index) => { + // Couldn't reserve memory for the minimum of 2 streams. Re-spill + // the larger of the two we're trying to merge with half its batch + // size so its largest batch shrinks, lowering the per-stream + // reservation, then retry. Makes the merge resilient to skewed + // (very wide) rows. + self.split_spill_file_in_half(index).await?; + continue; + } + }; // TODO - add a threshold for number of files to disk even if empty and reading from disk so // we can avoid the memory reservation @@ -210,46 +248,74 @@ impl MultiLevelMergeBuilder { continue; }; - // Add the spill file - self.sorted_spill_files.push(SortedSpillFile { - file: spill_file, - max_record_batch_memory, - }); + // Add the spill file paired with the batch-size limit of the merge that + // produced it: if that merge consumed a shrunk (skew-resolved) run, its + // output was capped and this intermediate run is likewise capped, so a + // later pass that re-merges it won't rebuild an oversized batch. + self.sorted_spill_files.push(( + SortedSpillFile { + file: spill_file, + max_record_batch_memory, + }, + batch_size_limit, + )); } } /// This tries to create a stream that merges the most sorted streams and sorted spill files /// as possible within the memory limit. - fn merge_sorted_runs_within_mem_limit( - &mut self, - ) -> Result { + fn merge_sorted_runs_within_mem_limit(&mut self) -> Result { match (self.sorted_spill_files.len(), self.sorted_streams.len()) { // No data so empty batch - (0, 0) => Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone( - &self.schema, - )))), + (0, 0) => { + let empty_stream = + Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); + Ok(MergeStep::Stream { + stream: self.observe_output(empty_stream), + batch_size_limit: self.batch_size, + }) + } // Only in-memory stream, return that - (0, 1) => Ok(self.sorted_streams.remove(0)), + (0, 1) => { + let output_stream = self.sorted_streams.remove(0); + Ok(MergeStep::Stream { + stream: self.observe_output(output_stream), + batch_size_limit: self.batch_size, + }) + } // Only single sorted spill file so return it (1, 0) => { - let spill_file = self.sorted_spill_files.remove(0); + let (spill_file, batch_size) = self.sorted_spill_files.remove(0); // Not reserving any memory for this disk as we are not holding it in memory - self.spill_manager - .read_spill_as_stream(spill_file.file, None) + let output_stream = self + .spill_manager + .read_spill_as_stream(spill_file.file, None)?; + + Ok(MergeStep::Stream { + stream: self.observe_output(output_stream), + batch_size_limit: batch_size, + }) } - // Only in memory streams, so merge them all in a single pass + // Only in memory streams, so merge them all in a single pass. In-memory + // runs are never shrunk for skew, so this merge runs at the full batch + // size and its output carries no limit. (0, _) => { let sorted_stream = mem::take(&mut self.sorted_streams); - self.create_new_merge_sort( - sorted_stream, - // If we have no sorted spill files left, this is the last run - true, - true, - ) + // No need to wrap with observed stream since merge sort will update the observed metrics + Ok(MergeStep::Stream { + stream: self.create_new_merge_sort( + sorted_stream, + // If we have no sorted spill files left, this is the last run + true, + true, + self.batch_size, + )?, + batch_size_limit: self.batch_size, + }) } // Need to merge multiple streams @@ -261,17 +327,33 @@ impl MultiLevelMergeBuilder { // allocation. let mut memory_reservation = self.reservation.take(); - // Don't account for existing streams memory - // as we are not holding the memory for them - let mut sorted_streams = mem::take(&mut self.sorted_streams); + // Compute the minimum before taking the in-memory streams so that, if we + // need to re-spill and retry, `self.sorted_streams` is left untouched. + let minimum_number_of_required_streams = + 2_usize.saturating_sub(self.sorted_streams.len()); - let (sorted_spill_files, buffer_size) = self + let (sorted_spill_files, buffer_size) = match self .get_sorted_spill_files_to_merge( 2, // we must have at least 2 streams to merge - 2_usize.saturating_sub(sorted_streams.len()), + minimum_number_of_required_streams, &mut memory_reservation, - )?; + )? { + SpillFilesToMerge::Ready(sorted_spill_files, buffer_size) => { + (sorted_spill_files, buffer_size) + } + // Not enough memory to seat 2 streams. Re-spill the blocking file + // smaller and retry. `get_sorted_spill_files_to_merge` already freed + // the reservation and `self.sorted_streams` is untouched, so the + // retry starts clean. + SpillFilesToMerge::SplitThenRetry(index) => { + return Ok(MergeStep::SplitThenRetry(index)); + } + }; + + // Don't account for existing streams memory + // as we are not holding the memory for them + let mut sorted_streams = mem::take(&mut self.sorted_streams); let is_only_merging_memory_streams = sorted_spill_files.is_empty(); @@ -284,7 +366,15 @@ impl MultiLevelMergeBuilder { mem::swap(&mut self.reservation, &mut memory_reservation); } - for spill in sorted_spill_files { + // Cap the merge output at the smallest limit among the runs we're + // about to merge. Runs that were shrunk for skew carry a smaller limit, + // if none do, every run carries `self.batch_size` and the merge runs at + // the full batch size. The output stream is tagged with the same limit + // (see the `MergeStep::Stream` returns below) so a re-spilled + // intermediate run stays shrunk and won't rebuild an oversized batch on + // a later pass. + let mut output_batch_size = self.batch_size; + for (spill, batch_size_limit) in sorted_spill_files { let stream = self .spill_manager .clone() @@ -293,6 +383,7 @@ impl MultiLevelMergeBuilder { spill.file, Some(spill.max_record_batch_memory), )?; + output_batch_size = output_batch_size.min(batch_size_limit); sorted_streams.push(stream); } let merge_sort_stream = self.create_new_merge_sort( @@ -300,6 +391,7 @@ impl MultiLevelMergeBuilder { // If we have no sorted spill files left, this is the last run self.sorted_spill_files.is_empty(), is_only_merging_memory_streams, + output_batch_size, )?; // If we're only merging memory streams, we don't need to attach the memory reservation @@ -311,14 +403,20 @@ impl MultiLevelMergeBuilder { "when only merging memory streams, we should not have any memory reservation and let the merge sort handle the memory" ); - Ok(merge_sort_stream) + Ok(MergeStep::Stream { + stream: merge_sort_stream, + batch_size_limit: output_batch_size, + }) } else { // Attach the memory reservation to the stream to make sure we have enough memory // throughout the merge process as we bypassed the memory pool for the merge sort stream - Ok(Box::pin(StreamAttachedReservation::new( - merge_sort_stream, - memory_reservation, - ))) + Ok(MergeStep::Stream { + stream: Box::pin(StreamAttachedReservation::new( + merge_sort_stream, + memory_reservation, + )), + batch_size_limit: output_batch_size, + }) } } } @@ -329,11 +427,12 @@ impl MultiLevelMergeBuilder { streams: Vec, is_output: bool, all_in_memory: bool, + output_batch_size: usize, ) -> Result { let mut builder = StreamingMergeBuilder::new() .with_schema(Arc::clone(&self.schema)) .with_expressions(&self.expr) - .with_batch_size(self.batch_size) + .with_batch_size(output_batch_size) .with_fetch(self.fetch) .with_metrics(if is_output { // Only add the metrics to the last run @@ -370,16 +469,26 @@ impl MultiLevelMergeBuilder { buffer_len: usize, minimum_number_of_required_streams: usize, reservation: &mut MemoryReservation, - ) -> Result<(Vec, usize)> { + ) -> Result { assert_ne!(buffer_len, 0, "Buffer length must be greater than 0"); let mut number_of_spills_to_read_for_current_phase = 0; + let configured_fan_in = self + .spill_manager + .env() + .disk_manager + .max_spill_merge_fan_in(); + let max_spill_files = effective_spill_merge_fan_in(configured_fan_in); // Track total memory needed for spill file buffers. When the // reservation has pre-reserved bytes (from sort_spill_reservation_bytes), // those bytes cover the first N spill files without additional pool // allocation, preventing starvation under memory pressure. let mut total_needed: usize = 0; - for spill in &self.sorted_spill_files { + for (spill, _) in &self.sorted_spill_files { + if number_of_spills_to_read_for_current_phase >= max_spill_files { + break; + } + let per_spill = get_reserved_bytes_for_record_batch_size( spill.max_record_batch_memory, // Size will be the same as the sliced size, bc it is a spilled batch. @@ -412,7 +521,24 @@ impl MultiLevelMergeBuilder { ); } - return Err(err); + // buffer_len == 1 and we still can't seat the minimum of 2 streams. + if number_of_spills_to_read_for_current_phase == 0 { + // We couldn't even reserve a single stream - one record batch + // is larger than the whole merge budget. That's the lone-batch + // case, not the 2-stream merge skew we rescue here - surface it. + return Err(err); + } + + // We seated one stream (index 0) but not the second (index 1, the + // batch that just failed to reserve). Those are by definition the + // only two streams we are trying to merge, so re-spill the larger + // of them with a smaller batch size and retry, the smaller max + // batch lowers the per-stream reservation enough to seat both. + let split_index = usize::from( + self.sorted_spill_files[1].0.max_record_batch_memory + > self.sorted_spill_files[0].0.max_record_batch_memory, + ); + return Ok(SpillFilesToMerge::SplitThenRetry(split_index)); } // We reached the maximum amount of memory we can use @@ -427,7 +553,161 @@ impl MultiLevelMergeBuilder { .drain(..number_of_spills_to_read_for_current_phase) .collect::>(); - Ok((spills, buffer_len)) + Ok(SpillFilesToMerge::Ready(spills, buffer_len)) + } + + /// Re-spill the spill file at `index` with half its batch size, putting it back + /// at the same position. We read the file back and re-spill it through the normal + /// spill API (which owns batch layout), slicing every batch in two, which halves + /// the largest written batch and so lowers the per-stream merge reservation enough + /// for the next attempt to seat both streams. One stream's worth of memory is + /// reserved for the duration and freed afterwards. Makes the merge resilient to skew. + /// + /// Instead of halving the *global* merge batch size (which would compound when more + /// than one run is re-spilled), the shrunk run records its own smaller batch-size + /// limit (tracked alongside the run in `sorted_spill_files`), so only merges that + /// actually consume it pay the reduced batch size. + async fn split_spill_file_in_half(&mut self, index: usize) -> Result<()> { + log::debug!( + "2 spilled streams could not be loaded into memory for merge \ + (requires 2x of the largest batch from both), re-spilling the larger of the two with half \ + the batch size to reduce memory needs for the next merge attempt. the shrunk run carries \ + a halved batch-size limit so only merges consuming it use the smaller batch size" + ); + + // Extract the target in O(1) instead of `remove(index)`, which would shift + // every following spill file. Swap it to the back and pop it; the matching + // swap after re-spilling restores the original order, so the vec ends up + // exactly as it started, just with the target file shrunk. + // `old_batch_size` is the batch size this run was written with (the full merge + // batch size unless it was already shrunk once). Halving it caps the next merge + // that reads this run so the merged output can't rebuild a full-size batch. + let last = self.sorted_spill_files.len() - 1; + self.sorted_spill_files.swap(index, last); + let (target, old_batch_size) = self + .sorted_spill_files + .pop() + .expect("index is in bounds, so the vec is non-empty"); + let old_max = target.max_record_batch_memory; + + // Reserve enough to hold a single stream of this file while we re-spill it. + let reservation = self.reservation.new_empty(); + reservation + .try_grow(get_reserved_bytes_for_record_batch_size(old_max, old_max))?; + + let source = self + .spill_manager + .read_spill_as_stream(target.file, Some(old_max))?; + // Re-spill with half the batch size: slice every batch in two. The spill + // writer owns the batch layout, we only change how many rows per batch. + let mut halved: SendableRecordBatchStream = + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&self.schema), + source.flat_map(|batch| { + futures::stream::iter(match batch { + Ok(batch) => split_batch_in_half(batch) + .into_iter() + .map(Ok) + .collect::>(), + Err(e) => vec![Err(e)], + }) + }), + )); + + let result = self + .spill_manager + .spill_record_batch_stream_and_return_max_batch_memory( + &mut halved, + "MultiLevelMergeBuilder split skewed spill", + ) + .await?; + + reservation.free(); + + let Some((file, new_max)) = result else { + return internal_err!("re-spilling a skewed spill file produced no data"); + }; + + // If halving could not reduce the largest batch (e.g. a single row that is + // itself wider than the budget), there is nothing more we can do - surface + // the out-of-memory condition instead of looping forever. + if new_max >= old_max { + return resources_err!( + "Cannot merge sorted runs: a single record batch of {old_max} bytes \ + exceeds the available merge memory and cannot be split further" + ); + } + + // Record the halved batch size as a *per-run* limit rather than lowering the + // global batch size. Merges that don't touch this run keep the full batch + // size. a merge that reads it caps its output at this limit so the merged run + // can't rebuild a full-size batch and reintroduce the skew. + let new_batch_size_limit = (old_batch_size / 2).max(1); + + // Push the re-spilled (smaller) file and swap it back into `index`, undoing + // the swap-to-back above so the order is preserved. + self.sorted_spill_files.push(( + SortedSpillFile { + file, + max_record_batch_memory: new_max, + }, + new_batch_size_limit, + )); + let last = self.sorted_spill_files.len() - 1; + self.sorted_spill_files.swap(index, last); + + Ok(()) + } + + fn observe_output( + &self, + stream: SendableRecordBatchStream, + ) -> SendableRecordBatchStream { + Box::pin(ObservedStream::new(stream, self.metrics.clone(), None)) + } +} + +/// Outcome of trying to reserve memory for one multi-level merge pass. +enum SpillFilesToMerge { + /// Enough memory: the spill files to read this pass (each paired with its + /// batch-size limit) and the read-ahead buffer size. + Ready(Vec<(SortedSpillFile, usize)>, usize), + /// Could not seat the minimum of 2 streams. Re-spill the spill file at this index + /// with a smaller (halved) batch size, then retry the pass. + SplitThenRetry(usize), +} + +/// What one iteration of the multi-level merge loop should do next. +enum MergeStep { + /// A merged stream is ready to be consumed (and possibly spilled back). + Stream { + stream: SendableRecordBatchStream, + /// The batch-size limit to stamp on the run if this stream is re-spilled as an + /// intermediate result: the batch size its merge ran at. It equals the full + /// merge batch size unless the merge consumed a skew-resolved run, in which + /// case it is that run's smaller limit so the re-spilled result stays capped + /// and can't rebuild an oversized batch. + batch_size_limit: usize, + }, + /// Re-spill the spill file at this index smaller, then retry the merge step. + SplitThenRetry(usize), +} + +/// Slice `batch` into two row-halves so a re-spill writes batches half the size. +fn split_batch_in_half(batch: RecordBatch) -> Vec { + let num_rows = batch.num_rows(); + if num_rows <= 1 { + return vec![batch]; + } + let mid = num_rows / 2; + vec![batch.slice(0, mid), batch.slice(mid, num_rows - mid)] +} + +fn effective_spill_merge_fan_in(configured_fan_in: usize) -> usize { + if configured_fan_in == 0 { + usize::MAX + } else { + configured_fan_in.max(2) } } @@ -481,3 +761,340 @@ impl RecordBatchStream for StreamAttachedReservation { self.stream.schema() } } + +#[cfg(test)] +mod tests { + use super::*; + + use crate::expressions::PhysicalSortExpr; + use arrow::array::{AsArray, Int64Array}; + use arrow::compute::concat_batches; + use arrow::datatypes::{DataType, Field, Int64Type, Schema}; + use datafusion_execution::memory_pool::{ + GreedyMemoryPool, MemoryConsumer, MemoryPool, + }; + use datafusion_execution::runtime_env::{RuntimeEnv, RuntimeEnvBuilder}; + use datafusion_physical_expr::expressions::{Column, col}; + use datafusion_physical_expr_common::metrics::{ + ExecutionPlanMetricsSet, SpillMetrics, + }; + + fn test_schema() -> SchemaRef { + Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, false)])) + } + + fn build_spill_manager(env: &Arc, schema: &SchemaRef) -> SpillManager { + SpillManager::new( + Arc::clone(env), + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(schema), + ) + } + + /// Spill `values` (which must already be sorted) as a single sorted run and + /// return it as a `SortedSpillFile` carrying its recorded largest-batch memory. + fn make_sorted_spill_file( + spill_manager: &SpillManager, + schema: &SchemaRef, + values: Vec, + ) -> SortedSpillFile { + let batch = RecordBatch::try_new( + Arc::clone(schema), + vec![Arc::new(Int64Array::from(values))], + ) + .unwrap(); + let batches: Vec> = vec![Ok(batch)]; + let (file, max_record_batch_memory) = spill_manager + .spill_record_batch_iter_and_return_max_batch_memory( + batches.into_iter(), + "test input run", + ) + .unwrap() + .expect("spill should produce a file"); + SortedSpillFile { + file, + max_record_batch_memory, + } + } + + fn build_merge_builder( + spill_manager: SpillManager, + schema: SchemaRef, + sorted_spill_files: Vec, + pool: &Arc, + batch_size: usize, + ) -> MultiLevelMergeBuilder { + let reservation = MemoryConsumer::new("test merge").register(pool); + let expr: LexOrdering = + [PhysicalSortExpr::new_default(Arc::new(Column::new("x", 0)))].into(); + MultiLevelMergeBuilder::new( + spill_manager, + schema, + sorted_spill_files, + vec![], + expr, + BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + batch_size, + reservation, + None, + false, + ) + } + + /// Two sorted runs whose largest batches are too big to both + /// be seated in the merge budget at once are re-spilled (halved) until they + /// fit, and the merge then completes with fully sorted, complete output. + #[tokio::test] + async fn skewed_runs_are_respilled_so_the_merge_fits() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + + let n: i64 = 16384; + let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); + + // Seating two streams needs ~4*m (2*m each), which does NOT fit, but the + // budget is large enough once a run is halved. The rescue keeps halving + // the blocking run until two streams fit (here, after one halving). + let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 7 / 2)); + + let builder = build_merge_builder( + spill_manager, + Arc::clone(&schema), + vec![f0, f1], + &pool, + 8192, + ); + let stream = builder.create_spillable_merge_stream(); + let batches: Vec = stream.try_collect().await?; + + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + total_rows, + (2 * n) as usize, + "the merge must emit every input row" + ); + + let merged = concat_batches(&schema, &batches)?; + let col = merged.column(0).as_primitive::(); + for i in 1..col.len() { + assert!( + col.value(i - 1) <= col.value(i), + "merge output must be sorted: {} > {} at {i}", + col.value(i - 1), + col.value(i), + ); + } + + Ok(()) + } + + /// Tests the `new_max >= old_max` guard: a single-row run cannot be split + /// any smaller, so re-spilling it does not shrink the largest batch and the + /// rescue surfaces `ResourcesExhausted` rather than looping forever. + #[tokio::test] + async fn respilling_an_unsplittable_run_surfaces_resources_exhausted() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + + // A one-row run: `split_batch_in_half` returns it unchanged, so the + // re-spilled file's largest batch cannot drop below the original. + let f0 = make_sorted_spill_file(&spill_manager, &schema, vec![42]); + + // Ample budget so the only possible failure is the un-splittable guard, + // not the single-stream reservation itself. + let pool: Arc = Arc::new(GreedyMemoryPool::new(1024 * 1024)); + let mut builder = + build_merge_builder(spill_manager, schema, vec![f0], &pool, 1024); + + let err = builder + .split_spill_file_in_half(0) + .await + .expect_err("re-spilling a one-row run cannot shrink it"); + assert!( + err.to_string().contains("cannot be split further"), + "expected the un-splittable guard error, got: {err}" + ); + + Ok(()) + } + + /// Proves the re-spill also halves the merge output batch size: after one + /// re-spill the merged run is emitted in 4096-row batches (not the original + /// 8192), so it cannot rebuild a full-size batch and reintroduce the skew. + #[tokio::test] + async fn respill_halves_the_merge_output_batch_size() -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + + let n: i64 = 16384; + let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); + + // 3.5*m forces exactly one re-spill (split one run, then both fit), which + // halves the merge output batch size. + let initial_batch_size = 8192; + let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 7 / 2)); + + let builder = build_merge_builder( + spill_manager, + Arc::clone(&schema), + vec![f0, f1], + &pool, + initial_batch_size, + ); + let stream = builder.create_spillable_merge_stream(); + let batches: Vec = stream.try_collect().await?; + + // All rows are still present. + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, (2 * n) as usize); + + // The largest emitted batch is the halved size, not the original 8192: the + // shrunk run carries a halved batch-size limit, and the final pass consumes + // it, so the merge output is capped there. Without the per-run limit the merge + // would rebuild 8192-row batches. + let expected_batch_size = initial_batch_size / 2; + let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); + assert_eq!( + max_batch_rows, expected_batch_size, + "after one re-spill the merge must emit {expected_batch_size}-row \ + batches, got a largest batch of {max_batch_rows} rows" + ); + + Ok(()) + } + + /// Same as [`respill_halves_the_merge_output_batch_size`], but under a budget tight + /// enough that *both* runs must be re-spilled before the merge fits - the scenario + /// where the batch-size reduction could compound. Because the reduction is tracked + /// per-run (each run capped at half) rather than by halving the global batch size on + /// every split, the merged output is emitted in 4096-row batches - half, not a + /// quarter. A global-halving implementation would have halved once per re-spill and + /// emitted 2048-row batches. + #[tokio::test] + async fn respilling_two_skewed_runs_halves_the_output_without_compounding() + -> Result<()> { + let env = Arc::new(RuntimeEnv::default()); + let schema = test_schema(); + let spill_manager = build_spill_manager(&env, &schema); + + let n: i64 = 16384; + let f0 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let f1 = make_sorted_spill_file(&spill_manager, &schema, (0..n).collect()); + let m = f0.max_record_batch_memory.max(f1.max_record_batch_memory); + + // 2.5*m is tight enough that even after halving one run the two still don't + // fit, so *both* runs are re-spilled once before the merge succeeds. (3.5*m, + // as in the single-split test, would let the pair fit after one split.) This + // is exactly the scenario where a compounding, global-halving implementation + // would drive the output batch size down to a quarter. + let initial_batch_size = 8192; + let pool: Arc = Arc::new(GreedyMemoryPool::new(m * 5 / 2)); + + let builder = build_merge_builder( + spill_manager, + Arc::clone(&schema), + vec![f0, f1], + &pool, + initial_batch_size, + ); + let stream = builder.create_spillable_merge_stream(); + let batches: Vec = stream.try_collect().await?; + + // All rows are still present. + let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total_rows, (2 * n) as usize); + + // Each run was re-spilled once, so each is capped at half the original batch + // size and the merge caps its output at that half - NOT a quarter. A global + // halving-per-split implementation would have emitted 2048-row batches here. + let expected_batch_size = initial_batch_size / 2; + let max_batch_rows = batches.iter().map(|b| b.num_rows()).max().unwrap_or(0); + assert_eq!( + max_batch_rows, expected_batch_size, + "two re-spills must halve (not quarter) the output: expected \ + {expected_batch_size}-row batches, got a largest batch of \ + {max_batch_rows} rows" + ); + + Ok(()) + } + + #[test] + fn spill_merge_fan_in_is_unlimited_by_default() { + assert_eq!(effective_spill_merge_fan_in(0), usize::MAX); + } + + #[test] + fn spill_merge_fan_in_preserves_merge_progress() { + assert_eq!(effective_spill_merge_fan_in(1), 2); + assert_eq!(effective_spill_merge_fan_in(2), 2); + assert_eq!(effective_spill_merge_fan_in(8), 8); + } + + #[test] + fn spill_merge_phase_respects_configured_fan_in() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let runtime = RuntimeEnvBuilder::new() + .with_max_spill_merge_fan_in(2) + .build_arc()?; + let spill_manager = SpillManager::new( + Arc::clone(&runtime), + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(&schema), + ); + let sorted_spill_files = (0..4) + .map(|idx| { + Ok(SortedSpillFile { + file: runtime + .disk_manager + .create_tmp_file(&format!("spill fan-in test {idx}"))?, + max_record_batch_memory: 1, + }) + }) + .collect::>>()?; + let expr = LexOrdering::new([PhysicalSortExpr::new_default(col("a", &schema)?)]) + .unwrap(); + let reservation = + MemoryConsumer::new("spill_merge_phase_respects_configured_fan_in") + .register(&runtime.memory_pool); + let metrics = BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let mut builder = MultiLevelMergeBuilder::new( + spill_manager, + schema, + sorted_spill_files, + vec![], + expr, + metrics, + 1024, + reservation, + None, + false, + ); + let mut merge_reservation = MemoryConsumer::new("spill_merge_fan_in_phase") + .register(&runtime.memory_pool); + + let (spills, buffer_len) = match builder.get_sorted_spill_files_to_merge( + 1, + 2, + &mut merge_reservation, + )? { + SpillFilesToMerge::Ready(spills, buffer_len) => (spills, buffer_len), + SpillFilesToMerge::SplitThenRetry(index) => { + panic!("expected ready spill files, got retry for index {index}") + } + }; + + assert_eq!(spills.len(), 2); + assert_eq!(buffer_len, 1); + assert_eq!(builder.sorted_spill_files.len(), 2); + + Ok(()) + } +} diff --git a/datafusion/physical-plan/src/sorts/partial_sort.rs b/datafusion/physical-plan/src/sorts/partial_sort.rs index abd9ebb142a66..478ac14e119d2 100644 --- a/datafusion/physical-plan/src/sorts/partial_sort.rs +++ b/datafusion/physical-plan/src/sorts/partial_sort.rs @@ -58,11 +58,12 @@ use std::task::{Context, Poll}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::sorts::sort::sort_batch; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, }; use arrow::compute::concat_batches; @@ -77,7 +78,104 @@ use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; use futures::{Stream, StreamExt, ready}; use log::trace; -/// Partial Sort execution plan. +/// Sort execution plan for inputs that are already partially sorted. +/// +/// This operator takes input ordered by a prefix of the required ordering, and +/// produces output ordered by the required ordering, emitting rows sooner +/// (streaming) and using less peak memory than [`SortExec`] which must buffer +/// all rows before producing any output. +/// +/// [`PartialSortExec`] relies on the property that rows with the same sort +/// prefix are contiguous, so it can sort one prefix group at a time, emitting +/// completed groups without reading (and buffering) the entire input. +/// +/// For example, if the required output is `(a, b, c)`, but the input is only +/// ordered by `(a, b)`, `PartialSortExec` sorts only within each `(a, b)` +/// group to produce output ordered by `(a, b, c)`. +/// +/// ```text +/// input ordered by a, b output ordered by a, b, c +/// +/// +---+---+---+ +---+---+---+ +/// | a | b | c | | a | b | c | +/// +---+---+---+ +---+---+---+ +/// | 0 | 0 | 3 | -- new group --> | 0 | 0 | 1 | +/// | 0 | 0 | 2 | | 0 | 0 | 2 | +/// | 0 | 0 | 1 | | 0 | 0 | 3 | +/// | 0 | 1 | 1 | -- new group --> | 0 | 1 | 1 | +/// | 0 | 2 | 4 | -- new group --> | 0 | 2 | 0 | +/// | 0 | 2 | 0 | | 0 | 2 | 4 | +/// | 1 | 0 | 5 | -- new group --> | 1 | 0 | 5 | +/// +---+---+---+ +---+---+---+ +/// ``` +/// +/// # Buffering and Emitting Rows +/// +/// [`PartialSortExec`] buffers rows only until it can *prove* a prefix group +/// will never be seen again, then sorts and emits buffered rows. A group is +/// guaranteed to never be seen again once a row with a *different* prefix +/// value arrives. This relies on the input's existing ordering guarantees. +/// +/// Using the example from above, rows accumulate in the in-memory buffer in +/// batches. As long as the `(a, b)` prefix keeps repeating, more rows are +/// buffered. +/// +/// ```text +/// Buffer +/// +---+---+---+ +/// | a | b | c | +/// +---+---+---+ +/// | 0 | 0 | 3 | +/// | 0 | 0 | 2 | +/// | 0 | 0 | 1 | +/// +---+---+---+ +/// ``` +/// +/// Once a batch arrives that contains a new `(a, b)` prefix, e.g. `(0, 2)`: +/// every buffered row for previous prefixes may be emitted: +/// +/// ```text +/// Buffer +/// +---+---+---+ +/// | a | b | c | +/// +---+---+---+ +/// | 0 | 0 | 3 | +/// | 0 | 0 | 2 | +/// | 0 | 0 | 1 | +/// | 0 | 1 | 1 | <-- first row of new batch, new prefix +/// | 0 | 2 | 4 | <-- new prefix +/// | 0 | 2 | 0 | +/// | 1 | 0 | 5 | <-- last row of new batch, new prefix +/// +---+---+---+ +/// ``` +/// +/// Once known complete, the buffered rows are sorted by the full `(a, b, c)` +/// ordering and emitted as a [`RecordBatch`]; Any rows from the most recently +/// seen prefix remain buffered (as more rows with the same prefix may arrive in +/// future batches. +/// +/// ```text +/// Emitted <-- fully sorted on (a, b, c) +/// +---+---+---+ +/// | a | b | c | +/// +---+---+---+ +/// | 0 | 0 | 1 | <-- completed group +/// | 0 | 0 | 2 | +/// | 0 | 0 | 3 | +/// | 0 | 2 | 0 | <-- completed group +/// | 0 | 2 | 4 | +/// | 0 | 1 | 1 | <-- completed group +/// +---+---+---+ +/// +/// Buffer +/// +---+---+---+ +/// | a | b | c | +/// +---+---+---+ +/// | 1 | 0 | 5 | <-- (possibly) in progress group +/// +---+---+---+ +/// ``` +/// +/// [`SortExec`]: crate::sorts::sort::SortExec #[derive(Debug, Clone)] pub struct PartialSortExec { /// Input schema @@ -205,17 +303,6 @@ impl PartialSortExec { input.boundedness(), )) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics_set: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for PartialSortExec { @@ -269,11 +356,15 @@ impl ExecutionPlan for PartialSortExec { } fn required_input_distribution(&self) -> Vec { - if self.preserve_partitioning { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.preserve_partitioning { vec![Distribution::UnspecifiedDistribution] } else { vec![Distribution::SinglePartition] - } + }) } fn benefits_from_input_partitioning(&self) -> Vec { @@ -286,29 +377,58 @@ impl ExecutionPlan for PartialSortExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for sort_expr in &self.expr { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics_set: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new_partial_sort = PartialSortExec::new( + self.expr.clone(), + Arc::clone(&children[0]), + self.common_prefix_length, + ) + .with_fetch(self.fetch) + .with_preserve_partitioning(self.preserve_partitioning); + + Ok(Arc::new(new_partial_sort)) + } } - Ok(tnr) } fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - let new_partial_sort = PartialSortExec::new( - self.expr.clone(), - Arc::clone(&children[0]), - self.common_prefix_length, + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), ) - .with_fetch(self.fetch) - .with_preserve_partitioning(self.preserve_partitioning); + } - Ok(Arc::new(new_partial_sort)) + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -346,8 +466,16 @@ impl ExecutionPlan for PartialSortExec { Some(self.metrics_set.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - self.input.partition_statistics(partition) + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::clone(&input_stats[0])) } } diff --git a/datafusion/physical-plan/src/sorts/partitioned_topk.rs b/datafusion/physical-plan/src/sorts/partitioned_topk.rs index f4c2585ea790d..41ccfab6833d5 100644 --- a/datafusion/physical-plan/src/sorts/partitioned_topk.rs +++ b/datafusion/physical-plan/src/sorts/partitioned_topk.rs @@ -23,36 +23,52 @@ //! FROM t WHERE rn <= N //! ``` //! -//! Instead of sorting the entire dataset, this operator maintains a -//! [`TopK`] heap per partition (reusing the existing TopK implementation) -//! and emits only the top-K rows per partition in sorted order -//! `(partition_keys, order_keys)`. +//! Instead of sorting the entire dataset, this operator delegates to a +//! per-partition heap-of-K implementation (one variant for `ROW_NUMBER` +//! and a sibling variant for `RANK`), both of which maintain one heap per +//! distinct partition key while sharing a single [`arrow::row::RowConverter`], +//! [`MemoryReservation`](datafusion_execution::memory_pool::MemoryReservation), +//! and metrics set across all partitions, and emit only the top-K rows +//! per partition in sorted order `(partition_keys, order_keys)`. use std::fmt::{self, Formatter}; use std::sync::Arc; -use arrow::array::{RecordBatch, UInt32Array}; -use arrow::compute::{BatchCoalescer, take_record_batch}; use arrow::datatypes::SchemaRef; -use arrow::row::{OwnedRow, RowConverter}; +use arrow::row::SortField; +use datafusion_common::Result; use datafusion_common::tree_node::TreeNodeRecursion; -use datafusion_common::{HashMap, Result}; use datafusion_execution::TaskContext; +use datafusion_execution::runtime_env::RuntimeEnv; use datafusion_physical_expr::PhysicalExpr; -use datafusion_physical_expr::expressions::{DynamicFilterPhysicalExpr, lit}; use datafusion_physical_expr_common::sort_expr::LexOrdering; use futures::StreamExt; use futures::TryStreamExt; -use parking_lot::RwLock; use crate::execution_plan::{Boundedness, EmissionType}; use crate::metrics::ExecutionPlanMetricsSet; -use crate::topk::{TopK, TopKDynamicFilters, build_sort_fields}; +use crate::topk::{PartitionedTopK, PartitionedTopKRank, build_sort_fields}; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use crate::{ DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, PlanProperties, SendableRecordBatchStream, stream::RecordBatchStreamAdapter, }; +/// Which window function `PartitionedTopKExec` is optimizing. +/// +/// Different ranking functions have different per-partition retention rules: +/// - [`RowNumber`](Self::RowNumber): exactly K rows per partition. +/// - [`Rank`](Self::Rank): K rows plus any rows tied at the boundary +/// ORDER BY value (RANK semantics — `WHERE rk <= K` may keep more +/// than K rows when ties straddle the boundary). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum WindowFnKind { + /// `ROW_NUMBER()` — keep exactly K rows per partition. + RowNumber, + /// `RANK()` — keep K rows plus any rows tied at the boundary. + Rank, +} + /// Per-partition Top-K operator for window function queries. /// /// # Background @@ -93,9 +109,14 @@ use crate::{ /// DataSourceExec /// ``` /// -/// Instead of sorting the entire dataset, this operator reads unsorted input, -/// maintains a [`TopK`] heap per distinct partition key, and emits only the -/// top-K rows per partition in sorted order `(partition_keys, order_keys)`. +/// Instead of sorting the entire dataset, this operator reads unsorted input +/// and delegates to a per-partition heap-of-K implementation (`PartitionedTopK` +/// for `ROW_NUMBER` and `PartitionedTopKRank` for `RANK`), each maintaining +/// one heap per distinct partition key while sharing a single +/// [`arrow::row::RowConverter`] / +/// [`MemoryReservation`](datafusion_execution::memory_pool::MemoryReservation) +/// across all partitions, and emits only the top-K rows per partition in +/// sorted order `(partition_keys, order_keys)`. /// /// Cost: O(N log K) time instead of O(N log N), and O(K × P × row_size) /// memory where K = fetch, P = number of distinct partitions. @@ -143,9 +164,11 @@ use crate::{ /// /// # Limitations /// -/// - Only activated when the window function is `ROW_NUMBER` with a -/// `PARTITION BY` clause. Global top-K (no `PARTITION BY`) is already -/// handled efficiently by `SortExec` with `fetch`. +/// - Only activated when the window function is `ROW_NUMBER` or `RANK` with +/// a `PARTITION BY` clause. `RANK` additionally requires a non-empty +/// `ORDER BY` (with an empty `ORDER BY`, every row ties at rank 1 and the +/// heap-of-K rewrite doesn't apply). Global top-K (no `PARTITION BY`) is +/// already handled efficiently by `SortExec` with `fetch`. /// - For very high cardinality partition keys (millions of distinct values), /// both memory usage and runtime overhead can become significant. In such /// cases, the sort-based plan is more robust. Therefore, this optimization @@ -168,6 +191,9 @@ pub struct PartitionedTopKExec { /// Derived from the filter predicate: `rn <= 3` → `fetch = 3`, /// `rn < 3` → `fetch = 2`. fetch: usize, + /// Which window function this operator is optimizing. Selects the + /// per-partition retention policy (see [`WindowFnKind`]). + fn_kind: WindowFnKind, /// Execution metrics metrics_set: ExecutionPlanMetricsSet, /// Cached plan properties (output ordering, partitioning, etc.) @@ -185,6 +211,8 @@ impl PartitionedTopKExec { /// * `partition_prefix_len` - Number of leading expressions in `expr` /// that form the partition key. Must be >= 1. /// * `fetch` - Maximum rows to retain per partition (the K in "top-K"). + /// * `fn_kind` - Which ranking window function this operator optimizes + /// ([`WindowFnKind::RowNumber`] or [`WindowFnKind::Rank`]). /// /// # Example /// @@ -195,6 +223,7 @@ impl PartitionedTopKExec { /// LexOrdering([store ASC, revenue DESC]), /// 1, // partition_prefix_len: 1 partition column (store) /// 5, // fetch: keep top 5 per partition + /// WindowFnKind::RowNumber, /// ) /// ``` pub fn try_new( @@ -202,6 +231,7 @@ impl PartitionedTopKExec { expr: LexOrdering, partition_prefix_len: usize, fetch: usize, + fn_kind: WindowFnKind, ) -> Result { let cache = Self::compute_properties(&input, expr.clone())?; Ok(Self { @@ -209,6 +239,7 @@ impl PartitionedTopKExec { expr, partition_prefix_len, fetch, + fn_kind, metrics_set: ExecutionPlanMetricsSet::new(), cache: Arc::new(cache), }) @@ -235,6 +266,11 @@ impl PartitionedTopKExec { self.fetch } + /// Returns which window function this operator is optimizing. + pub fn fn_kind(&self) -> WindowFnKind { + self.fn_kind + } + /// Compute [`PlanProperties`] for this operator. /// /// The output is sorted by `sort_exprs` (partition keys then order keys), @@ -258,6 +294,10 @@ impl PartitionedTopKExec { impl DisplayAs for PartitionedTopKExec { fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> fmt::Result { + let fn_label = match self.fn_kind { + WindowFnKind::RowNumber => "row_number", + WindowFnKind::Rank => "rank", + }; match t { DisplayFormatType::Default | DisplayFormatType::Verbose => { let partition_exprs: Vec = self.expr[..self.partition_prefix_len] @@ -270,7 +310,8 @@ impl DisplayAs for PartitionedTopKExec { .collect(); write!( f, - "PartitionedTopKExec: fetch={}, partition=[{}], order=[{}]", + "PartitionedTopKExec: fn={}, fetch={}, partition=[{}], order=[{}]", + fn_label, self.fetch, partition_exprs.join(", "), order_exprs.join(", "), @@ -285,6 +326,7 @@ impl DisplayAs for PartitionedTopKExec { .iter() .map(|e| format!("{e}")) .collect(); + writeln!(f, "fn={fn_label}")?; writeln!(f, "fetch={}", self.fetch)?; writeln!(f, "partition=[{}]", partition_exprs.join(", "))?; writeln!(f, "order=[{}]", order_exprs.join(", ")) @@ -303,12 +345,18 @@ impl ExecutionPlan for PartitionedTopKExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { let partition_exprs: Vec> = self.expr [..self.partition_prefix_len] .iter() .map(|e| Arc::clone(&e.expr)) .collect(); - vec![Distribution::HashPartitioned(partition_exprs)] + crate::InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + partition_exprs, + )]) } fn maintains_input_order(&self) -> Vec { @@ -319,9 +367,10 @@ impl ExecutionPlan for PartitionedTopKExec { vec![&self.input] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert_eq!(children.len(), 1); Ok(Arc::new(PartitionedTopKExec::try_new( @@ -329,18 +378,28 @@ impl ExecutionPlan for PartitionedTopKExec { self.expr.clone(), self.partition_prefix_len, self.fetch, + self.fn_kind, )?)) } fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for sort_expr in &self.expr { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - Ok(tnr) + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) } fn execute( @@ -354,8 +413,6 @@ impl ExecutionPlan for PartitionedTopKExec { let partition_sort_fields = build_sort_fields(&self.expr[..self.partition_prefix_len], &schema)?; - let partition_converter = RowConverter::new(partition_sort_fields)?; - let partition_exprs: Vec> = self.expr [..self.partition_prefix_len] .iter() @@ -365,18 +422,21 @@ impl ExecutionPlan for PartitionedTopKExec { LexOrdering::new(self.expr[self.partition_prefix_len..].iter().cloned()) .expect("PartitionedTopKExec requires at least one order-by expression"); let fetch = self.fetch; + let fn_kind = self.fn_kind; let batch_size = context.session_config().batch_size(); let runtime = Arc::clone(&context.runtime_env()); let metrics_set = self.metrics_set.clone(); let stream = futures::stream::once(async move { do_partitioned_topk( + partition, input, schema, - partition_converter, partition_exprs, + partition_sort_fields, order_expr, fetch, + fn_kind, batch_size, runtime, metrics_set, @@ -392,136 +452,79 @@ impl ExecutionPlan for PartitionedTopKExec { } } -/// Create a no-op [`TopKDynamicFilters`] for a per-partition [`TopK`]. -/// -/// In normal `SortExec` top-K mode, dynamic filters push predicates down to -/// the data source (e.g., telling Parquet to skip rows worse than the current -/// K-th best). For per-partition heaps the data is already in memory and split -/// by partition key, so there is no data source to push filters to. We pass -/// `lit(true)` (accept everything) so the filter never rejects any row. -fn create_noop_dynamic_filter() -> Arc> { - Arc::new(RwLock::new(TopKDynamicFilters::new(Arc::new( - DynamicFilterPhysicalExpr::new(vec![], lit(true)), - )))) -} - -/// Read all input, split batches by partition key, feed each sub-batch -/// to a per-partition [`TopK`], then emit results in partition-key order. +/// Read all input, feed each batch into a per-partition top-K state +/// (either [`PartitionedTopK`] for `ROW_NUMBER` or +/// [`PartitionedTopKRank`] for `RANK`), then emit results ordered by +/// `(partition_keys, order_keys)`. /// /// # Phases /// -/// 1. **Accumulation** — For each input batch: -/// - Evaluate partition key expressions to get partition column arrays -/// - Convert partition columns to binary [`arrow::row::Row`] format -/// - Group row indices by partition key -/// - Extract sub-batches via [`take_record_batch`] and insert into -/// the partition's [`TopK`] heap +/// 1. **Accumulation** — forward each input `RecordBatch` to the +/// per-partition state's `insert_batch`. The `RowConverter` for +/// ORDER BY columns, the operator's `MemoryReservation`, and the +/// `TopKMetrics` are shared across all distinct partition keys for +/// this operator instance. /// -/// 2. **Emission** — After all input is consumed: -/// - Sort partition keys so output is ordered by partition key -/// - For each partition in sorted order, call [`TopK::emit`] to get -/// rows sorted by order-by key -/// - Return all batches as a single stream +/// 2. **Emission** — `emit` drains all per-partition heaps in sorted +/// partition-key order, returning a coalesced batch stream. For +/// `RANK`, boundary-tied rows are materialized and emitted after +/// each partition's heap rows. /// /// # Cost /// /// - Time: O(N log K) where N = total rows, K = fetch /// - Memory: O(K × P × row_size) where P = number of distinct partitions +/// plus, for RANK, the boundary ties' rows #[expect(clippy::too_many_arguments)] async fn do_partitioned_topk( + partition_id: usize, mut input: SendableRecordBatchStream, schema: SchemaRef, - partition_converter: RowConverter, partition_exprs: Vec>, + partition_sort_fields: Vec, order_expr: LexOrdering, fetch: usize, + fn_kind: WindowFnKind, batch_size: usize, - runtime: Arc, + runtime: Arc, metrics_set: ExecutionPlanMetricsSet, ) -> Result { - let mut partitions: HashMap = HashMap::new(); - let mut partition_counter: usize = 0; - - // Macro-like helper: create a new TopK for a partition - macro_rules! new_topk { - () => {{ - let id = partition_counter; - partition_counter += 1; - TopK::try_new( - id, - Arc::clone(&schema), - vec![], - order_expr.clone(), + match fn_kind { + WindowFnKind::RowNumber => { + let mut state = PartitionedTopK::try_new( + partition_id, + schema, + partition_exprs, + partition_sort_fields, + order_expr, fetch, batch_size, - Arc::clone(&runtime), + &runtime, &metrics_set, - create_noop_dynamic_filter(), - ) - }}; - } - - // ---------- Accumulation phase ---------- - while let Some(batch) = input.next().await { - let batch = batch?; - let num_rows = batch.num_rows(); - if num_rows == 0 { - continue; - } - - // Evaluate partition key columns - let pk_arrays: Vec<_> = partition_exprs - .iter() - .map(|e| e.evaluate(&batch).and_then(|v| v.into_array(num_rows))) - .collect::>>()?; - - let pk_rows = partition_converter.convert_columns(&pk_arrays)?; - - // Group row indices by partition key - let mut groups: HashMap> = HashMap::new(); - for row_idx in 0..num_rows { - let pk = pk_rows.row(row_idx).owned(); - groups.entry(pk).or_default().push(row_idx as u32); - } - - // For each partition group, create a sub-batch and feed to TopK - for (pk, indices) in groups { - if !partitions.contains_key(&pk) { - partitions.insert(pk.clone(), new_topk!()?); + )?; + while let Some(batch) = input.next().await { + state.insert_batch(&batch?)?; } - let topk = partitions.get_mut(&pk).unwrap(); - let indices_array = UInt32Array::from(indices); - let sub_batch = take_record_batch(&batch, &indices_array)?; - topk.insert_batch(sub_batch)?; + drop(input); + state.emit() } - } - // Release the input pipeline now that accumulation is complete. - drop(input); - - // ---------- Emit phase ---------- - // Sort partition keys so output is ordered by (partition_keys, order_keys). - let mut sorted_pks: Vec = partitions.keys().cloned().collect(); - sorted_pks.sort(); - - let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); - - for pk in sorted_pks { - if let Some(topk) = partitions.remove(&pk) { - // TopK::emit() returns a stream of sorted batches - let mut stream = topk.emit()?; - while let Some(batch) = stream.next().await { - coalescer.push_batch(batch?)?; + WindowFnKind::Rank => { + let mut state = PartitionedTopKRank::try_new( + partition_id, + schema, + partition_exprs, + partition_sort_fields, + order_expr, + fetch, + batch_size, + &runtime, + &metrics_set, + )?; + while let Some(batch) = input.next().await { + state.insert_batch(&batch?)?; } + drop(input); + state.emit() } } - coalescer.finish_buffered_batch()?; - let mut output_batches: Vec = Vec::new(); - while let Some(batch) = coalescer.next_completed_batch() { - output_batches.push(batch); - } - - Ok(Box::pin(RecordBatchStreamAdapter::new( - schema, - futures::stream::iter(output_batches.into_iter().map(Ok)), - ))) } diff --git a/datafusion/physical-plan/src/sorts/sort.rs b/datafusion/physical-plan/src/sorts/sort.rs index 90d4b5ec12f91..6c782f5134484 100644 --- a/datafusion/physical-plan/src/sorts/sort.rs +++ b/datafusion/physical-plan/src/sorts/sort.rs @@ -28,6 +28,7 @@ use parking_lot::RwLock; use crate::common::spawn_buffered; use crate::execution_plan::{ Boundedness, CardinalityEffect, EmissionType, has_same_children_properties, + replace_children_if_necessary, }; use crate::expressions::PhysicalSortExpr; use crate::filter::FilterExec; @@ -45,14 +46,15 @@ use crate::sorts::streaming_merge::{SortedSpillFile, StreamingMergeBuilder}; use crate::spill::get_record_batch_memory_size; use crate::spill::in_progress_spill_file::InProgressSpillFile; use crate::spill::spill_manager::{GetSlicedSize, SpillManager}; -use crate::stream::RecordBatchStreamAdapter; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::ReservationStream; +use crate::stream::{ObservedStream, RecordBatchStreamAdapter}; use crate::topk::TopK; use crate::topk::TopKDynamicFilters; use crate::{ - DisplayAs, DisplayFormatType, Distribution, EmptyRecordBatchStream, ExecutionPlan, - ExecutionPlanProperties, Partitioning, PlanProperties, SendableRecordBatchStream, - Statistics, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, + EmptyRecordBatchStream, ExecutionPlan, ExecutionPlanProperties, Partitioning, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, }; use arrow::array::{RecordBatch, RecordBatchOptions}; @@ -374,7 +376,7 @@ impl ExternalSorter { // allocation. Only needed for the non-spill path; the spill // path transfers the reservation to the merge stream instead. self.merge_reservation.free(); - self.in_mem_sort_stream(self.metrics.baseline.clone()) + self.in_mem_sort_stream(true, true) } } @@ -406,7 +408,7 @@ impl ExternalSorter { /// Appending globally sorted batches to the in-progress spill file, and clears /// the `globally_sorted_batches` (also its memory reservation) afterwards. - async fn consume_and_spill_append( + fn consume_and_spill_append( &mut self, globally_sorted_batches: &mut Vec, ) -> Result<()> { @@ -445,7 +447,7 @@ impl ExternalSorter { } /// Finishes the in-progress spill file and moves it to the finished spill files. - async fn spill_finish(&mut self) -> Result<()> { + fn spill_finish(&mut self) -> Result<()> { let (mut in_progress_file, max_record_batch_memory) = self.in_progress_spill_file.take().ok_or_else(|| { internal_datafusion_err!("Should be called after `spill_append`") @@ -476,8 +478,11 @@ impl ExternalSorter { // reserved again for the next spill. self.merge_reservation.free(); - let mut sorted_stream = - self.in_mem_sort_stream(self.metrics.baseline.intermediate())?; + let mut sorted_stream = self.in_mem_sort_stream( + false, + // No coalescing on the spill path: it raises per-run peak memory. + false, + )?; // After `in_mem_sort_stream()` is constructed, all `in_mem_batches` is taken // to construct a globally sorted stream. assert_or_internal_err!( @@ -497,8 +502,7 @@ impl ExternalSorter { // already in memory, so it's okay to combine it with previously // sorted batches, and spill together. globally_sorted_batches.push(batch); - self.consume_and_spill_append(&mut globally_sorted_batches) - .await?; // reservation is freed in spill() + self.consume_and_spill_append(&mut globally_sorted_batches)?; // reservation is freed in spill() } else { globally_sorted_batches.push(batch); } @@ -508,9 +512,8 @@ impl ExternalSorter { // upcoming `self.reserve_memory_for_merge()` may fail due to insufficient memory. drop(sorted_stream); - self.consume_and_spill_append(&mut globally_sorted_batches) - .await?; - self.spill_finish().await?; + self.consume_and_spill_append(&mut globally_sorted_batches)?; + self.spill_finish()?; // Sanity check after spilling let buffers_cleared_property = @@ -584,19 +587,22 @@ impl ExternalSorter { /// /// in_mem_batches /// ``` + /// `coalesce_runs` merges buffered batches into fewer, larger sorted runs to + /// reduce merge fan-in. Disabled on the spill path to keep peak memory low. fn in_mem_sort_stream( &mut self, - metrics: BaselineMetrics, + is_output_stream: bool, + coalesce_runs: bool, ) -> Result { if self.in_mem_batches.is_empty() { - return Ok(Box::pin(EmptyRecordBatchStream::new(Arc::clone( - &self.schema, - )))); + let empty_stream = + Box::pin(EmptyRecordBatchStream::new(Arc::clone(&self.schema))); + return Ok(self.observe_if_output(empty_stream, is_output_stream)); } // The elapsed compute timer is updated when the value is dropped. // There is no need for an explicit call to drop. - let elapsed_compute = metrics.elapsed_compute().clone(); + let elapsed_compute = self.metrics.baseline.elapsed_compute().clone(); let _timer = elapsed_compute.timer(); // Please pay attention that any operation inside of `in_mem_sort_stream` will @@ -608,7 +614,8 @@ impl ExternalSorter { if self.in_mem_batches.len() == 1 { let batch = self.in_mem_batches.swap_remove(0); let reservation = self.reservation.take(); - return self.sort_batch_stream(batch, &metrics, reservation); + let sorted_stream = self.sort_batch_stream(batch, reservation)?; + return Ok(self.observe_if_output(sorted_stream, is_output_stream)); } // If less than sort_in_place_threshold_bytes, concatenate and sort in place @@ -620,17 +627,29 @@ impl ExternalSorter { .try_resize(get_reserved_bytes_for_record_batch(&batch)?) .map_err(Self::err_with_oom_context)?; let reservation = self.reservation.take(); - return self.sort_batch_stream(batch, &metrics, reservation); - } + let sorted_stream = self.sort_batch_stream(batch, reservation)?; + return Ok(self.observe_if_output(sorted_stream, is_output_stream)); + } + + // For single-column sorts, coalesce the buffered batches into fewer, + // larger runs to cut the merge fan-in (where the cheap per-key compare is + // dominated by per-stream cursor/merge overhead). Multi-column sorts are + // left as one run per batch: the row-format merge of many small runs + // beats sorting a few large runs with the lexicographic comparator. + let batches = std::mem::take(&mut self.in_mem_batches); + let runs = if coalesce_runs && self.expr.len() == 1 { + self.coalesce_in_mem_batches_into_runs(batches)? + } else { + batches + }; - let streams = std::mem::take(&mut self.in_mem_batches) + let streams = runs .into_iter() .map(|batch| { - let metrics = self.metrics.baseline.intermediate(); let reservation = self .reservation .split(get_reserved_bytes_for_record_batch(&batch)?); - let input = self.sort_batch_stream(batch, &metrics, reservation)?; + let input = self.sort_batch_stream(batch, reservation)?; Ok(spawn_buffered(input, 1)) }) .collect::>()?; @@ -639,13 +658,69 @@ impl ExternalSorter { .with_streams(streams) .with_schema(Arc::clone(&self.schema)) .with_expressions(&self.expr.clone()) - .with_metrics(metrics) + .with_metrics(if is_output_stream { + self.metrics.baseline.clone() + } else { + self.metrics.baseline.intermediate() + }) .with_batch_size(self.batch_size) .with_fetch(None) .with_reservation(self.merge_reservation.new_empty()) .build() } + /// Concatenates `batches` into fewer, larger runs, each bounded by + /// `sort_in_place_threshold_bytes`, to reduce merge fan-in. `self.reservation` + /// is resized to the coalesced footprint so the caller's per-run splits stay + /// exact. + fn coalesce_in_mem_batches_into_runs( + &mut self, + batches: Vec, + ) -> Result> { + let target = self.sort_in_place_threshold_bytes.max(1); + let mut runs: Vec = Vec::new(); + let mut group: Vec = Vec::new(); + let mut group_bytes = 0usize; + + // Flush a group into a run, skipping the copy for a single-batch group. + let flush = |group: &mut Vec, + runs: &mut Vec, + schema: &SchemaRef| + -> Result<()> { + match group.len() { + 0 => {} + 1 => runs.push(group.pop().unwrap()), + _ => { + runs.push(concat_batches(schema, group.iter())?); + group.clear(); + } + } + Ok(()) + }; + + for batch in batches { + let bytes = get_reserved_bytes_for_record_batch(&batch)?; + if !group.is_empty() && group_bytes.saturating_add(bytes) > target { + flush(&mut group, &mut runs, &self.schema)?; + group_bytes = 0; + } + group_bytes += bytes; + group.push(batch); + } + flush(&mut group, &mut runs, &self.schema)?; + + // Realign the reservation: concatenation may shift the footprint slightly. + let total: usize = runs + .iter() + .map(get_reserved_bytes_for_record_batch) + .sum::>()?; + self.reservation + .try_resize(total) + .map_err(Self::err_with_oom_context)?; + + Ok(runs) + } + /// Sorts a single `RecordBatch` into a single stream. /// /// This may output multiple batches depending on the size of the @@ -658,7 +733,6 @@ impl ExternalSorter { fn sort_batch_stream( &self, batch: RecordBatch, - metrics: &BaselineMetrics, reservation: MemoryReservation, ) -> Result { assert_eq!( @@ -669,7 +743,6 @@ impl ExternalSorter { let schema = batch.schema(); let expressions = self.expr.clone(); let batch_size = self.batch_size; - let output_row_metrics = metrics.output_rows().clone(); let stream = futures::stream::once(async move { let schema = batch.schema(); @@ -699,14 +772,7 @@ impl ExternalSorter { reservation, )) as SendableRecordBatchStream) }) - .try_flatten() - .map(move |batch| match batch { - Ok(batch) => { - output_row_metrics.add(batch.num_rows()); - Ok(batch) - } - Err(e) => Err(e), - }); + .try_flatten(); Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) } @@ -765,6 +831,22 @@ impl ExternalSorter { _ => e, } } + + fn observe_if_output( + &self, + mut stream: SendableRecordBatchStream, + wrap: bool, + ) -> SendableRecordBatchStream { + if wrap { + stream = Box::pin(ObservedStream::new( + stream, + self.metrics.baseline.clone(), + None, + )) + } + + stream + } } /// Estimate how much memory is needed to sort a `RecordBatch`. @@ -905,19 +987,51 @@ impl SortExec { self.preserve_partitioning = preserve_partitioning; Arc::make_mut(&mut self.cache).partitioning = Self::output_partitioning_helper(&self.input, self.preserve_partitioning); + if self.fetch.is_some() { + self.rebuild_filter_for_current_partitioning(); + } self } - /// Add or reset `self.filter` to a new `TopKDynamicFilters`. + fn topk_emitter_count(&self) -> usize { + self.cache.output_partitioning().partition_count() + } + + /// Build a new shared TopK dynamic filter wrapper for this `SortExec`. fn create_filter(&self) -> Arc> { let children = self .expr .iter() .map(|sort_expr| Arc::clone(&sort_expr.expr)) .collect::>(); - Arc::new(RwLock::new(TopKDynamicFilters::new(Arc::new( - DynamicFilterPhysicalExpr::new(children, lit(true)), - )))) + self.create_filter_with_expr(Arc::new(DynamicFilterPhysicalExpr::new( + children, + lit(true), + ))) + } + + fn create_filter_with_expr( + &self, + expr: Arc, + ) -> Arc> { + Arc::new(RwLock::new( + TopKDynamicFilters::new_with_topk_emitter_count( + expr, + self.topk_emitter_count(), + ), + )) + } + + /// Rebuild the shared TopK filter wrapper for the current output partitioning. + /// + /// The dynamic filter expression is preserved, but wrapper state such as the + /// shared threshold and remaining emitter count is reset for the new + /// partitioning. + fn rebuild_filter_for_current_partitioning(&mut self) { + let filter_expr = self.filter.as_ref().map(|filter| filter.read().expr()); + if let Some(filter_expr) = filter_expr { + self.filter = Some(self.create_filter_with_expr(filter_expr)); + } } fn cloned(&self) -> Self { @@ -952,14 +1066,20 @@ impl SortExec { if fetch.is_some() && is_pipeline_friendly { cache = cache.with_boundedness(Boundedness::Bounded); } - let filter = fetch.is_some().then(|| { - // If we already have a filter, keep it. Otherwise, create a new one. - self.filter.clone().unwrap_or_else(|| self.create_filter()) - }); let mut new_sort = self.cloned(); new_sort.fetch = fetch; new_sort.cache = cache.into(); - new_sort.filter = filter; + if fetch.is_some() { + if new_sort.filter.is_some() { + // Keep the dynamic filter expression, but reset wrapper state + // such as the shared threshold and expected emitter count. + new_sort.rebuild_filter_for_current_partitioning(); + } else { + new_sort.filter = Some(new_sort.create_filter()); + } + } else { + new_sort.filter = None; + } new_sort } @@ -979,6 +1099,10 @@ impl SortExec { } /// Returns the dynamic filter expression for this sort (TopK), if set. + #[deprecated( + since = "55.0.0", + note = "Use ExecutionPlan::dynamic_expressions_produced instead" + )] pub fn dynamic_filter_expr(&self) -> Option> { self.filter.as_ref().map(|f| f.read().expr()) } @@ -998,7 +1122,7 @@ impl SortExec { for child in filter.children() { child.data_type(&input_schema)?; } - self.filter = Some(Arc::new(RwLock::new(TopKDynamicFilters::new(filter)))); + self.filter = Some(self.create_filter_with_expr(filter)); Ok(self) } @@ -1138,13 +1262,18 @@ impl ExecutionPlan for SortExec { } fn required_input_distribution(&self) -> Vec { - if self.preserve_partitioning { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(if self.preserve_partitioning { vec![Distribution::UnspecifiedDistribution] } else { // global sort - // TODO support RangePartition and OrderedDistribution + // TODO support range partitioning and OrderedDistribution. + // See https://github.com/apache/datafusion/issues/22395 vec![Distribution::SinglePartition] - } + }) } fn children(&self) -> Vec<&Arc> { @@ -1153,37 +1282,43 @@ impl ExecutionPlan for SortExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to sort expressions - let mut tnr = TreeNodeRecursion::Continue; - for sort_expr in &self.expr { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - - // Apply to dynamic filter expression if present (when fetch is Some, TopK mode) - if let Some(filter) = &self.filter { - let filter_guard = filter.read(); - tnr = tnr.visit_sibling(|| f(filter_guard.expr().as_ref()))?; - } + let dynamic_filter = self + .filter + .as_ref() + .map(|filter| filter.read().expr() as Arc); + crate::apply_expression_roots( + self.expr + .iter() + .map(|sort_expr| &sort_expr.expr) + .chain(dynamic_filter.iter()), + f, + ) + } - Ok(tnr) + fn dynamic_expressions_produced(&self) -> Vec> { + self.filter + .iter() + .map(|filter| filter.read().expr() as Arc) + .collect() } fn benefits_from_input_partitioning(&self) -> Vec { vec![false] } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { let mut new_sort = self.cloned(); assert_eq!(children.len(), 1, "SortExec should have exactly one child"); new_sort.input = Arc::clone(&children[0]); - if !has_same_children_properties(self.as_ref(), &children)? { - // Recompute the properties based on the new input since they may have changed + if options.children_properties == ChildrenPropertiesMode::Recompute { + // Recompute the properties based on the new input since they may have changed. let (cache, sort_prefix) = Self::compute_properties( &new_sort.input, new_sort.expr.clone(), @@ -1191,17 +1326,36 @@ impl ExecutionPlan for SortExec { )?; new_sort.cache = Arc::new(cache); new_sort.common_sort_prefix = sort_prefix; + if new_sort.fetch.is_some() { + new_sort.rebuild_filter_for_current_partitioning(); + } } Ok(Arc::new(new_sort)) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + match has_same_children_properties(self.as_ref(), &children)? { + true => self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ), + false => self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ), + } + } + fn reset_state(self: Arc) -> Result> { let children = self.children().into_iter().cloned().collect(); - let new_sort = self.with_new_children(children)?; + let new_sort = replace_children_if_necessary(self, children)?; let mut new_sort = new_sort .downcast_ref::() - .expect("cloned 1 lines above this line, we know the type") + .expect("rebuilt SortExec with new children") .clone(); // Our dynamic filter and execution metrics are the state we need to reset. new_sort.filter = Some(new_sort.create_filter()); @@ -1302,13 +1456,21 @@ impl ExecutionPlan for SortExec { Some(self.metrics_set.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let p = if !self.preserve_partitioning() { - None - } else { + fn child_stats_requests(&self, partition: Option) -> Vec { + let child_partition = if self.preserve_partitioning() { partition + } else { + None }; - let stats = Arc::unwrap_or_clone(self.input.partition_statistics(p)?); + vec![ChildStats::At(child_partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let stats = input_stats[0].as_ref().clone(); Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) } @@ -1436,6 +1598,119 @@ impl ExecutionPlan for SortExec { updated_node: Some(new_sort), }) } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = self + .expr() + .iter() + .map(|sort_expr| { + let sort_node = Box::new(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&sort_expr.expr)?)), + asc: !sort_expr.options.descending, + nulls_first: sort_expr.options.nulls_first, + }); + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( + sort_node, + )), + }) + }) + .collect::>>()?; + let dynamic_filter = self + .dynamic_expressions_produced() + .into_iter() + .next() + .map(|expr| ctx.encode_expr(&expr)) + .transpose()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Sort(Box::new( + protobuf::SortExecNode { + input: Some(Box::new(input)), + expr, + fetch: match self.fetch() { + Some(n) => n as i64, + None => -1, + }, + preserve_partitioning: self.preserve_partitioning(), + dynamic_filter, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + use protobuf::physical_expr_node::ExprType; + let sort = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Sort, + "SortExec", + ); + let input = + ctx.decode_required_child(sort.input.as_deref(), "SortExec", "input")?; + let input_schema = input.schema(); + let exprs = sort + .expr + .iter() + .map(|expr| { + let Some(ExprType::Sort(sort_expr)) = expr.expr_type.as_ref() else { + return datafusion_common::internal_err!( + "SortExec expr must be a sort expression" + ); + }; + let expr_node = sort_expr.expr.as_deref().ok_or_else(|| { + internal_datafusion_err!( + "SortExec sort expression is missing its inner expr" + ) + })?; + Ok(PhysicalSortExpr { + expr: ctx.decode_expr(expr_node, input_schema.as_ref())?, + options: arrow::compute::SortOptions { + descending: !sort_expr.asc, + nulls_first: sort_expr.nulls_first, + }, + }) + }) + .collect::>>()?; + let Some(ordering) = LexOrdering::new(exprs) else { + return datafusion_common::internal_err!("SortExec requires an ordering"); + }; + let fetch = (sort.fetch >= 0).then_some(sort.fetch as usize); + let new_sort = SortExec::new(ordering, input) + .with_fetch(fetch) + .with_preserve_partitioning(sort.preserve_partitioning); + + let new_sort = if let Some(df_proto) = &sort.dynamic_filter { + let df_expr = + ctx.decode_expr(df_proto, new_sort.input().schema().as_ref())?; + let df = (df_expr as Arc) + .downcast::() + .map_err(|_| { + internal_datafusion_err!( + "SortExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" + ) + })?; + new_sort.with_dynamic_filter_expr(df)? + } else { + new_sort + }; + + Ok(Arc::new(new_sort)) + } } #[cfg(test)] @@ -1469,9 +1744,10 @@ mod tests { GreedyMemoryPool, MemoryConsumer, MemoryPool, }; use datafusion_execution::runtime_env::RuntimeEnvBuilder; - use datafusion_physical_expr::EquivalenceProperties; use datafusion_physical_expr::expressions::{Column, Literal}; + use datafusion_physical_expr::{DynamicFilterTracking, EquivalenceProperties}; + use datafusion_physical_expr_common::metrics::MetricValue; use futures::{FutureExt, Stream, TryStreamExt}; use insta::assert_snapshot; @@ -1523,20 +1799,31 @@ mod tests { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -1622,6 +1909,104 @@ mod tests { Ok(()) } + /// Single-column run coalescing: many small batches above a tiny in-place + /// threshold (with ample memory, so no spill) must still produce a correct + /// total order, including NULLs. + #[tokio::test] + async fn test_in_mem_sort_coalesced_runs() -> Result<()> { + // Tiny in-place threshold forces the sort-then-merge path and, for a + // single column, the coalescing branch. Ample memory => no spill. + let task_ctx = Arc::new( + TaskContext::default().with_session_config( + SessionConfig::new() + .with_batch_size(64) + .with_sort_in_place_threshold_bytes(1024), + ), + ); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + + // Build many small batches of shuffled values with interspersed NULLs, + // so coalescing produces several multi-row runs that must be merged. + let num_batches = 40; + let rows_per_batch = 50; + let mut all_values: Vec> = Vec::new(); + let mut batches = Vec::with_capacity(num_batches); + for b in 0..num_batches { + let mut col_values: Vec> = Vec::with_capacity(rows_per_batch); + for r in 0..rows_per_batch { + let idx = (b * rows_per_batch + r) as i64; + // Deterministic scramble to avoid any pre-existing ordering. + let scrambled = ((idx.wrapping_mul(2_654_435_761)) % 1000) as i32; + let v = if idx % 7 == 0 { None } else { Some(scrambled) }; + col_values.push(v); + all_values.push(v); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(col_values))], + )?; + batches.push(batch); + } + let total_rows = num_batches * rows_per_batch; + + let options = SortOptions::default(); + let sort_exec = Arc::new(SortExec::new( + [PhysicalSortExpr { + expr: col("a", &schema)?, + options, + }] + .into(), + TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )?, + )); + + let result = collect( + Arc::clone(&sort_exec) as Arc, + Arc::clone(&task_ctx), + ) + .await?; + + // Flatten the sorted output. + let mut got: Vec> = Vec::with_capacity(total_rows); + for batch in &result { + let arr = as_primitive_array::(batch.column(0))?; + for i in 0..arr.len() { + got.push(if arr.is_null(i) { + None + } else { + Some(arr.value(i)) + }); + } + } + assert_eq!(got.len(), total_rows, "row count must be preserved"); + + // Reference: sort the original values with the same semantics + // (ascending, NULLs first per SortOptions::default()). + let mut expected = all_values.clone(); + expected.sort_by(|a, b| match (a, b) { + (None, None) => std::cmp::Ordering::Equal, + (None, Some(_)) => std::cmp::Ordering::Less, // nulls_first + (Some(_), None) => std::cmp::Ordering::Greater, + (Some(x), Some(y)) => x.cmp(y), + }); + + assert_eq!( + got, expected, + "coalesced-run sort output must be totally ordered" + ); + assert_eq!( + task_ctx.runtime_env().memory_pool.reserved(), + 0, + "The sort should have returned all memory used back to the memory manager" + ); + + Ok(()) + } + #[tokio::test] async fn test_sort_spill() -> Result<()> { // trigger spill w/ 100 batches @@ -2323,7 +2708,8 @@ mod tests { } #[tokio::test] - async fn should_return_stream_with_batches_in_the_requested_size() -> Result<()> { + async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics() + -> Result<()> { let batch_size = 100; let create_task_ctx = |_: &[RecordBatch]| { @@ -2335,19 +2721,22 @@ mod tests { }; // Smaller than batch size and require more than a single batch to get the requested batch size - test_sort_output_batch_size(10, batch_size / 4, create_task_ctx).await?; + test_sort_output_batch_size_and_base_metrics(10, batch_size / 4, create_task_ctx) + .await?; // Not evenly divisible by batch size - test_sort_output_batch_size(10, batch_size + 7, create_task_ctx).await?; + test_sort_output_batch_size_and_base_metrics(10, batch_size + 7, create_task_ctx) + .await?; // Evenly divisible by batch size and is larger than 2 output batches - test_sort_output_batch_size(10, batch_size * 3, create_task_ctx).await?; + test_sort_output_batch_size_and_base_metrics(10, batch_size * 3, create_task_ctx) + .await?; Ok(()) } #[tokio::test] - async fn should_return_stream_with_batches_in_the_requested_size_when_sorting_in_place() + async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_sorting_in_place() -> Result<()> { let batch_size = 100; @@ -2361,8 +2750,12 @@ mod tests { // Smaller than batch size and require more than a single batch to get the requested batch size { - let metrics = - test_sort_output_batch_size(10, batch_size / 4, create_task_ctx).await?; + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size / 4, + create_task_ctx, + ) + .await?; assert_eq!( metrics.spill_count(), @@ -2373,8 +2766,12 @@ mod tests { // Not evenly divisible by batch size { - let metrics = - test_sort_output_batch_size(10, batch_size + 7, create_task_ctx).await?; + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size + 7, + create_task_ctx, + ) + .await?; assert_eq!( metrics.spill_count(), @@ -2385,8 +2782,12 @@ mod tests { // Evenly divisible by batch size and is larger than 2 output batches { - let metrics = - test_sort_output_batch_size(10, batch_size * 3, create_task_ctx).await?; + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size * 3, + create_task_ctx, + ) + .await?; assert_eq!( metrics.spill_count(), @@ -2399,7 +2800,7 @@ mod tests { } #[tokio::test] - async fn should_return_stream_with_batches_in_the_requested_size_when_having_a_single_batch() + async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_having_a_single_batch() -> Result<()> { let batch_size = 100; @@ -2410,7 +2811,7 @@ mod tests { // Smaller than batch size and require more than a single batch to get the requested batch size { - let metrics = test_sort_output_batch_size( + let metrics = test_sort_output_batch_size_and_base_metrics( // Single batch 1, batch_size / 4, @@ -2427,7 +2828,7 @@ mod tests { // Not evenly divisible by batch size { - let metrics = test_sort_output_batch_size( + let metrics = test_sort_output_batch_size_and_base_metrics( // Single batch 1, batch_size + 7, @@ -2444,7 +2845,7 @@ mod tests { // Evenly divisible by batch size and is larger than 2 output batches { - let metrics = test_sort_output_batch_size( + let metrics = test_sort_output_batch_size_and_base_metrics( // Single batch 1, batch_size * 3, @@ -2463,7 +2864,7 @@ mod tests { } #[tokio::test] - async fn should_return_stream_with_batches_in_the_requested_size_when_having_to_spill() + async fn should_return_stream_with_batches_in_the_requested_size_and_update_metrics_when_having_to_spill() -> Result<()> { let batch_size = 100; @@ -2491,24 +2892,36 @@ mod tests { // Smaller than batch size and require more than a single batch to get the requested batch size { - let metrics = - test_sort_output_batch_size(10, batch_size / 4, create_task_ctx).await?; + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size / 4, + create_task_ctx, + ) + .await?; assert_ne!(metrics.spill_count().unwrap(), 0, "expected to spill"); } // Not evenly divisible by batch size { - let metrics = - test_sort_output_batch_size(10, batch_size + 7, create_task_ctx).await?; + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size + 7, + create_task_ctx, + ) + .await?; assert_ne!(metrics.spill_count().unwrap(), 0, "expected to spill"); } // Evenly divisible by batch size and is larger than 2 batches { - let metrics = - test_sort_output_batch_size(10, batch_size * 3, create_task_ctx).await?; + let metrics = test_sort_output_batch_size_and_base_metrics( + 10, + batch_size * 3, + create_task_ctx, + ) + .await?; assert_ne!(metrics.spill_count().unwrap(), 0, "expected to spill"); } @@ -2516,7 +2929,7 @@ mod tests { Ok(()) } - async fn test_sort_output_batch_size( + async fn test_sort_output_batch_size_and_base_metrics( number_of_batches: usize, batch_size_to_generate: usize, create_task_ctx: impl Fn(&[RecordBatch]) -> TaskContext, @@ -2526,10 +2939,13 @@ mod tests { .collect::>(); let task_ctx = create_task_ctx(batches.as_slice()); + let output_rows = batches.iter().map(|item| item.num_rows()).sum(); + let expected_batch_size = task_ctx.session_config().batch_size(); + let schema = batches[0].schema(); let (mut output_batches, metrics) = - run_sort_on_input(task_ctx, "i", batches).await?; + run_sort_on_input(task_ctx, "i", batches, schema).await?; let last_batch = output_batches.pop().unwrap(); @@ -2544,18 +2960,87 @@ mod tests { } assert_eq!(last_batch.num_rows(), last_expected_batch_size); + assert_baseline_metrics_for_non_empty_output( + &metrics, + output_rows, + expected_batch_size, + ); + Ok(metrics) } + #[tokio::test] + async fn empty_sort_stream_should_report_end_time() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("i", DataType::Int32, false)])); + let task_ctx = TaskContext::default(); + + let (_, metrics) = run_sort_on_input(task_ctx, "i", vec![], schema).await?; + + let end_time = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::EndTimestamp(end) => Some(end), + _ => None, + }) + .expect("Must have end time metric since it exists in the baseline"); + + assert_eq!( + metrics.spill_count().unwrap_or_default(), + 0, + "expected to not have spills" + ); + assert_ne!(end_time.value(), None); + + Ok(()) + } + + fn assert_baseline_metrics_for_non_empty_output( + metrics: &MetricsSet, + output_rows: usize, + batch_size: usize, + ) { + let end_time = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::EndTimestamp(end) => Some(end), + _ => None, + }) + .expect("Must have end time metric since it exists in the baseline"); + + assert_ne!(end_time.value(), None); + + assert_eq!(metrics.output_rows(), Some(output_rows)); + + let output_bytes = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::OutputBytes(total) => Some(total), + _ => None, + }) + .expect("Must have output_bytes metric since it exists in the baseline"); + + assert_ne!(output_bytes.value(), 0_usize); + + let output_batches = metrics + .iter() + .find_map(|item| match item.value() { + MetricValue::OutputBatches(total) => Some(total), + _ => None, + }) + .expect("Must have output_batches metric since it exists in the baseline"); + + assert_eq!(output_batches.value(), output_rows.div_ceil(batch_size)); + } + async fn run_sort_on_input( task_ctx: TaskContext, order_by_col: &str, batches: Vec, + schema: SchemaRef, ) -> Result<(Vec, MetricsSet)> { let task_ctx = Arc::new(task_ctx); // let task_ctx = env. - let schema = batches[0].schema(); let ordering: LexOrdering = [PhysicalSortExpr { expr: col(order_by_col, &schema)?, options: SortOptions { @@ -2566,7 +3051,11 @@ mod tests { .into(); let sort_exec: Arc = Arc::new(SortExec::new( ordering.clone(), - TestMemoryExec::try_new_exec(std::slice::from_ref(&batches), schema, None)?, + TestMemoryExec::try_new_exec( + std::slice::from_ref(&batches), + Arc::clone(&schema), + None, + )?, )); let sorted_batches = @@ -2576,11 +3065,10 @@ mod tests { // assert output { - let input_batches_concat = concat_batches(batches[0].schema_ref(), &batches)?; + let input_batches_concat = concat_batches(&schema, &batches)?; let sorted_input_batch = sort_batch(&input_batches_concat, &ordering, None)?; - let sorted_batches_concat = - concat_batches(sorted_batches[0].schema_ref(), &sorted_batches)?; + let sorted_batches_concat = concat_batches(&schema, &sorted_batches)?; assert_eq!(sorted_input_batch, sorted_batches_concat); } @@ -2765,9 +3253,9 @@ mod tests { .with_fetch(Some(10)); // SortExec with fetch creates a dynamic filter automatically. - let original_id = sort - .dynamic_filter_expr() - .expect("should have dynamic filter with fetch") + let produced = sort.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + let original_id = produced[0] .expression_id() .expect("DynamicFilterPhysicalExpr always has an expression_id"); @@ -2780,9 +3268,9 @@ mod tests { .expression_id() .expect("DynamicFilterPhysicalExpr always has an expression_id"); let sort = sort.with_dynamic_filter_expr(Arc::clone(&new_df))?; - let restored_id = sort - .dynamic_filter_expr() - .expect("should still have dynamic filter") + let produced = sort.dynamic_expressions_produced(); + assert_eq!(produced.len(), 1); + let restored_id = produced[0] .expression_id() .expect("DynamicFilterPhysicalExpr always has an expression_id"); assert_eq!(restored_id, new_id); @@ -2790,6 +3278,136 @@ mod tests { Ok(()) } + async fn emit_sort_partition( + sort: &Arc, + partition: usize, + task_ctx: Arc, + ) -> Result<()> { + let _batches: Vec = + sort.execute(partition, task_ctx)?.try_collect().await?; + Ok(()) + } + + fn assert_filter_still_waiting(filter: &Arc) { + let dynamic_filter_expr: Arc = + Arc::::clone(filter); + assert!( + matches!( + DynamicFilterTracking::classify(&dynamic_filter_expr), + DynamicFilterTracking::Watching(_) + ), + "the shared filter should remain watchable until every partition emits" + ); + } + + fn dynamic_filter_produced( + plan: &dyn ExecutionPlan, + ) -> Arc { + let expr = plan + .dynamic_expressions_produced() + .into_iter() + .next() + .expect("plan should produce a dynamic filter"); + (expr as Arc) + .downcast::() + .expect("produced expression should be a DynamicFilterPhysicalExpr") + } + + #[tokio::test] + async fn test_preserved_topk_filter_waits_for_all_sort_partitions() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let partitions = vec![ + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![3, 1, 2]))], + )?], + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![6, 4, 5]))], + )?], + ]; + let input = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let sort = SortExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(), + input, + ) + // `with_fetch` creates the TopK filter; preserving partitioning after + // that must rebuild it with one emitter per output partition. + .with_fetch(Some(2)) + .with_preserve_partitioning(true); + + let dynamic_filter = dynamic_filter_produced(&sort); + let sort = Arc::new(sort); + let task_ctx = Arc::new(TaskContext::default()); + + emit_sort_partition(&sort, 0, Arc::clone(&task_ctx)).await?; + assert_filter_still_waiting(&dynamic_filter); + + emit_sort_partition(&sort, 1, task_ctx).await?; + tokio::time::timeout( + std::time::Duration::from_secs(1), + dynamic_filter.wait_complete(), + ) + .await + .expect("the final preserved SortExec partition should complete the filter"); + + Ok(()) + } + + #[tokio::test] + async fn test_with_fetch_rebuilds_existing_topk_filter() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let partitions = vec![ + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![3, 1, 2]))], + )?], + vec![RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(Int32Array::from(vec![6, 4, 5]))], + )?], + ]; + let input = TestMemoryExec::try_new_exec(&partitions, Arc::clone(&schema), None)?; + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0))], + lit(true), + )); + let dynamic_filter_id = dynamic_filter + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"); + let sort = SortExec::new( + [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(), + input, + ) + .with_dynamic_filter_expr(dynamic_filter)? + .with_preserve_partitioning(true) + .with_fetch(Some(2)); + + let dynamic_filter = dynamic_filter_produced(&sort); + assert_eq!( + dynamic_filter + .expression_id() + .expect("DynamicFilterPhysicalExpr always has an expression_id"), + dynamic_filter_id + ); + + let sort = Arc::new(sort); + let task_ctx = Arc::new(TaskContext::default()); + + emit_sort_partition(&sort, 0, Arc::clone(&task_ctx)).await?; + assert_filter_still_waiting(&dynamic_filter); + + emit_sort_partition(&sort, 1, task_ctx).await?; + tokio::time::timeout( + std::time::Duration::from_secs(1), + dynamic_filter.wait_complete(), + ) + .await + .expect("the final preserved SortExec partition should complete the filter"); + + Ok(()) + } + #[test] fn test_with_dynamic_filter_rejects_invalid_columns() -> Result<()> { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); diff --git a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs index 13c28ccb10991..ad17f2c2136af 100644 --- a/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs +++ b/datafusion/physical-plan/src/sorts/sort_preserving_merge.rs @@ -24,10 +24,11 @@ use crate::limit::LimitStream; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; use crate::projection::{ProjectionExec, make_with_child, update_ordering}; use crate::sorts::streaming_merge::StreamingMergeBuilder; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::{ - DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, ExecutionPlanProperties, - Partitioning, PlanProperties, SendableRecordBatchStream, Statistics, - check_if_same_properties, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, + ExecutionPlanProperties, Partitioning, PlanProperties, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, validate_child_count, }; use datafusion_common::tree_node::TreeNodeRecursion; @@ -37,7 +38,9 @@ use datafusion_execution::memory_pool::MemoryConsumer; use datafusion_physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{LexOrdering, OrderingRequirements}; -use crate::execution_plan::{EvaluationType, SchedulingType}; +use crate::execution_plan::{ + CardinalityEffect, EvaluationType, SchedulingType, replace_children_if_necessary, +}; use log::{debug, trace}; /// Sort preserving merge execution plan @@ -182,17 +185,6 @@ impl SortPreservingMergeExec { .with_evaluation_type(drive) .with_scheduling_type(scheduling) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for SortPreservingMergeExec { @@ -261,14 +253,19 @@ impl ExecutionPlan for SortPreservingMergeExec { self.input .with_preserve_order(preserve_order) .and_then(|new_input| { - Arc::new(self.clone()) - .with_new_children(vec![new_input]) + replace_children_if_necessary(Arc::new(self.clone()), vec![new_input]) .ok() }) } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn benefits_from_input_partitioning(&self) -> Vec { @@ -289,24 +286,51 @@ impl ExecutionPlan for SortPreservingMergeExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for sort_expr in &self.expr { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; + crate::apply_expression_roots( + self.expr.iter().map(|sort_expr| &sort_expr.expr), + f, + ) + } + + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new( + SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) + .with_fetch(self.fetch), + )), } - Ok(tnr) } fn with_new_children( self: Arc, - mut children: Vec>, + children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new( - SortPreservingMergeExec::new(self.expr.clone(), children.swap_remove(0)) - .with_fetch(self.fetch), - )) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -393,8 +417,25 @@ impl ExecutionPlan for SortPreservingMergeExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, _partition: Option) -> Result> { - self.input.partition_statistics(None) + fn child_stats_requests(&self, _partition: Option) -> Vec { + vec![ChildStats::At(None)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let stats = input_stats[0].as_ref().clone(); + Ok(Arc::new(stats.with_fetch(self.fetch, 0, 1)?)) + } + + fn cardinality_effect(&self) -> CardinalityEffect { + if self.fetch.is_none() { + CardinalityEffect::Equal + } else { + CardinalityEffect::LowerEqual + } } fn supports_limit_pushdown(&self) -> bool { @@ -426,6 +467,98 @@ impl ExecutionPlan for SortPreservingMergeExec { .with_fetch(self.fetch()), ))) } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let input = ctx.encode_child(self.input())?; + let expr = self + .expr() + .iter() + .map(|e| { + Ok(protobuf::PhysicalExprNode { + expr_id: None, + expr_type: Some(protobuf::physical_expr_node::ExprType::Sort( + Box::new(protobuf::PhysicalSortExprNode { + expr: Some(Box::new(ctx.encode_expr(&e.expr)?)), + asc: !e.options.descending, + nulls_first: e.options.nulls_first, + }), + )), + }) + }) + .collect::>>()?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge( + Box::new(protobuf::SortPreservingMergeExecNode { + input: Some(Box::new(input)), + expr, + fetch: self.fetch().map(|f| f as i64).unwrap_or(-1), + }), + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl SortPreservingMergeExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use arrow::compute::SortOptions; + use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; + use datafusion_proto_models::protobuf; + let spm = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::SortPreservingMerge, + "SortPreservingMergeExec", + ); + let input = ctx.decode_required_child( + spm.input.as_deref(), + "SortPreservingMergeExec", + "input", + )?; + let input_schema = input.schema(); + let exprs = spm + .expr + .iter() + .map(|e| { + let sort = match &e.expr_type { + Some(protobuf::physical_expr_node::ExprType::Sort(s)) => s, + _ => { + return internal_err!( + "SortPreservingMergeExec expression is not a sort expression" + ); + } + }; + let expr = ctx.decode_required_expr( + sort.expr.as_deref(), + input_schema.as_ref(), + "SortPreservingMergeExec", + "sort expression", + )?; + Ok(PhysicalSortExpr { + expr, + options: SortOptions { + descending: !sort.asc, + nulls_first: sort.nulls_first, + }, + }) + }) + .collect::>>()?; + let Some(ordering) = LexOrdering::new(exprs) else { + return internal_err!("SortPreservingMergeExec requires an ordering"); + }; + let fetch = (spm.fetch >= 0).then_some(spm.fetch as usize); + Ok(Arc::new( + SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), + )) + } } #[cfg(test)] @@ -444,9 +577,12 @@ mod tests { use crate::metrics::{MetricValue, Timestamp}; use crate::repartition::RepartitionExec; use crate::sorts::sort::SortExec; + use crate::statistics::StatisticsContext; use crate::stream::RecordBatchReceiverStream; use crate::test::TestMemoryExec; - use crate::test::exec::{BlockingExec, assert_strong_count_converges_to_zero}; + use crate::test::exec::{ + BlockingExec, StatisticsExec, assert_strong_count_converges_to_zero, + }; use crate::test::{self, assert_is_pending, make_partition}; use crate::{collect, common}; @@ -456,8 +592,9 @@ mod tests { }; use arrow::compute::SortOptions; use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; + use datafusion_common::stats::Precision; use datafusion_common::test_util::batches_to_string; - use datafusion_common::{assert_batches_eq, exec_err}; + use datafusion_common::{ColumnStatistics, assert_batches_eq, exec_err}; use datafusion_common_runtime::SpawnedTask; use datafusion_execution::RecordBatchStream; use datafusion_execution::config::SessionConfig; @@ -480,7 +617,8 @@ mod tests { .with_memory_limit(20_000_000, 1.0) .build_arc()?; let mut config = SessionConfig::new(); - config.options_mut().execution.batch_size = target_batch_size; + config.options_mut().execution.batch_size = + datafusion_common::config::ConfigNonZeroUsize::try_new(target_batch_size)?; let task_ctx = TaskContext::default() .with_runtime(runtime) .with_session_config(config); @@ -521,6 +659,52 @@ mod tests { Ok(Arc::new(spm)) } + #[test] + fn test_fetch_caps_statistics() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input = Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Exact(1_000), + total_byte_size: Precision::Exact(8_000), + column_statistics: vec![ColumnStatistics::new_unknown()], + }, + schema.clone(), + )); + let sort = [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + + let spm = SortPreservingMergeExec::new(sort, input).with_fetch(Some(1)); + let statistics = + StatisticsContext::new().compute(&spm, &StatisticsArgs::new())?; + + assert_eq!(statistics.num_rows, Precision::Exact(1)); + assert_eq!(statistics.total_byte_size, Precision::Inexact(8)); + assert!(matches!( + spm.cardinality_effect(), + CardinalityEffect::LowerEqual + )); + Ok(()) + } + + #[test] + fn test_no_fetch_preserves_statistics() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let input_stats = Statistics { + num_rows: Precision::Absent, + total_byte_size: Precision::Exact(8_000), + column_statistics: vec![ColumnStatistics::new_unknown()], + }; + let input = Arc::new(StatisticsExec::new(input_stats.clone(), schema.clone())); + let sort = [PhysicalSortExpr::new_default(Arc::new(Column::new("a", 0)))].into(); + + let spm = SortPreservingMergeExec::new(sort, input); + let statistics = + StatisticsContext::new().compute(&spm, &StatisticsArgs::new())?; + + assert_eq!(*statistics, input_stats); + assert!(matches!(spm.cardinality_effect(), CardinalityEffect::Equal)); + Ok(()) + } + /// This test verifies that memory usage stays within limits when the tie breaker is enabled. /// Any errors here could indicate unintended changes in tie breaker logic. /// @@ -1421,16 +1605,27 @@ mod tests { } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } fn execute( &self, partition: usize, @@ -1505,11 +1700,7 @@ mod tests { let task_ctx = Arc::new(TaskContext::default()); let schema = Schema::new(vec![Field::new("c1", DataType::UInt64, false)]); let properties = CongestedExec::compute_properties(Arc::new(schema.clone())); - let &partition_count = match properties.output_partitioning() { - Partitioning::RoundRobinBatch(partitions) => partitions, - Partitioning::Hash(_, partitions) => partitions, - Partitioning::UnknownPartitioning(partitions) => partitions, - }; + let partition_count = properties.output_partitioning().partition_count(); let source = CongestedExec { schema: schema.clone(), cache: Arc::new(properties), diff --git a/datafusion/physical-plan/src/sorts/stream.rs b/datafusion/physical-plan/src/sorts/stream.rs index ff7f259dd1347..107631074ed3d 100644 --- a/datafusion/physical-plan/src/sorts/stream.rs +++ b/datafusion/physical-plan/src/sorts/stream.rs @@ -15,8 +15,8 @@ // specific language governing permissions and limitations // under the License. -use crate::SendableRecordBatchStream; use crate::sorts::cursor::{ArrayValues, CursorArray, RowValues}; +use crate::{EmptyRecordBatchStream, SendableRecordBatchStream}; use crate::{PhysicalExpr, PhysicalSortExpr}; use arrow::array::{Array, UInt32Array}; use arrow::compute::take_record_batch; @@ -73,9 +73,21 @@ impl FusedStreams { stream_idx: usize, ) -> Poll>> { loop { - match ready!(self.0[stream_idx].poll_next_unpin(cx)) { - Some(Ok(b)) if b.num_rows() == 0 => continue, - r => return Poll::Ready(r), + let poll_result = self.0[stream_idx].poll_next_unpin(cx); + match &poll_result { + Poll::Pending => return Poll::Pending, + Poll::Ready(Some(Ok(b))) if b.num_rows() == 0 => continue, + Poll::Ready(Some(Ok(_))) => return poll_result, + Poll::Ready(None) | Poll::Ready(Some(Err(_))) => { + let stream_schema = self.0[stream_idx].get_ref().schema(); + + // Replace the stream with an empty stream, so we can drop memory usage + let empty_stream: SendableRecordBatchStream = + Box::pin(EmptyRecordBatchStream::new(stream_schema)); + self.0[stream_idx] = empty_stream.fuse(); + + return poll_result; + } } } } @@ -388,8 +400,12 @@ mod tests { use super::*; use arrow::array::{AsArray, Int32Array}; use arrow::datatypes::{DataType, Field, Int32Type}; + use arrow_schema::SchemaRef; use datafusion_common::DataFusionError; + use datafusion_execution::RecordBatchStream; use datafusion_physical_expr::expressions::col; + use futures::Stream; + use std::pin::Pin; /// Verifies that `take_record_batch` in `IncrementalSortIterator` actually /// copies the data into a new allocation rather than returning a zero-copy @@ -435,4 +451,93 @@ mod tests { assert_eq!(total_rows, original_len); Ok(()) } + + #[test] + fn test_fused_stream_drop_finished_streams() { + #[derive(Clone)] + struct SingleItemManualStream { + // Held only so its `Arc` strong count reveals when the stream is dropped. + #[expect(dead_code)] + hold_ref: Arc<()>, + record_batch: RecordBatch, + should_finish: bool, + } + + impl Stream for SingleItemManualStream { + type Item = Result; + + fn poll_next( + mut self: Pin<&mut Self>, + _cx: &mut Context<'_>, + ) -> Poll> { + if !self.should_finish { + self.should_finish = true; + return Poll::Ready(Some(Ok(self.record_batch.clone()))); + } + + Poll::Ready(None) + } + } + + impl RecordBatchStream for SingleItemManualStream { + fn schema(&self) -> SchemaRef { + self.record_batch.schema() + } + } + + let hold_ref = Arc::new(()); + let record_batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])), + vec![Arc::new(Int32Array::from(vec![1]))], + ) + .unwrap(); + + let stream_1 = SingleItemManualStream { + hold_ref: Arc::clone(&hold_ref), + should_finish: false, + record_batch: record_batch.clone(), + }; + let stream_2 = stream_1.clone(); + + let stream_1: SendableRecordBatchStream = Box::pin(stream_1); + let stream_2: SendableRecordBatchStream = Box::pin(stream_2); + + let mut fused_stream = FusedStreams(vec![stream_1.fuse(), stream_2.fuse()]); + + let waker = futures::task::noop_waker(); + let mut cx = Context::from_waker(&waker); + + // The original plus one clone held by each of the two streams. + assert_eq!(Arc::strong_count(&hold_ref), 3); + + // First fetch from stream 0 yields its single batch. + // the stream is not finished yet, so nothing is dropped. + let poll = fused_stream.poll_next(&mut cx, 0); + assert!(matches!(poll, Poll::Ready(Some(Ok(_))))); + assert_eq!(Arc::strong_count(&hold_ref), 3); + + // Second fetch from stream 0 returns `None`, so it is replaced with an + // empty stream and dropped, releasing its `hold_ref` clone. + // running 3 times to make sure the stream is fused correctly + for _ in 0..3 { + let poll = fused_stream.poll_next(&mut cx, 0); + assert!(matches!(poll, Poll::Ready(None))); + assert_eq!(Arc::strong_count(&hold_ref), 2); + } + + // First fetch from stream 1 yields its single batch + // the stream is not finished yet, so nothing is dropped. + let poll = fused_stream.poll_next(&mut cx, 1); + assert!(matches!(poll, Poll::Ready(Some(Ok(_))))); + assert_eq!(Arc::strong_count(&hold_ref), 2); + + // Second fetch from stream 1 returns `None`, so it is replaced with an + // empty stream and dropped, releasing its `hold_ref` clone. + // running 3 times to make sure the stream is fused correctly + for _ in 0..3 { + let poll = fused_stream.poll_next(&mut cx, 1); + assert!(matches!(poll, Poll::Ready(None))); + assert_eq!(Arc::strong_count(&hold_ref), 1); + } + } } diff --git a/datafusion/physical-plan/src/sorts/streaming_merge.rs b/datafusion/physical-plan/src/sorts/streaming_merge.rs index 8129c3d8f695d..81adad8e9ec84 100644 --- a/datafusion/physical-plan/src/sorts/streaming_merge.rs +++ b/datafusion/physical-plan/src/sorts/streaming_merge.rs @@ -24,12 +24,12 @@ use crate::sorts::{ merge::SortPreservingMergeStream, stream::{FieldCursorStream, RowCursorStream}, }; -use crate::{SendableRecordBatchStream, SpillManager}; +use crate::{EmptyRecordBatchStream, SendableRecordBatchStream, SpillManager}; use arrow::array::*; use arrow::datatypes::{DataType, SchemaRef}; use datafusion_common::human_readable_size; use datafusion_common::{Result, assert_or_internal_err, internal_err}; -use datafusion_execution::disk_manager::RefCountedTempFile; +use datafusion_execution::SpillFile; use datafusion_execution::memory_pool::{ MemoryConsumer, MemoryPool, MemoryReservation, UnboundedMemoryPool, }; @@ -46,7 +46,7 @@ macro_rules! merge_helper { ($t:ty, $sort:ident, $streams:ident, $schema:ident, $tracking_metrics:ident, $batch_size:ident, $fetch:ident, $reservation:ident, $enable_round_robin_tie_breaker:ident) => {{ let streams = FieldCursorStream::<$t>::new($sort, $streams, $reservation.new_empty()); - return Ok(Box::pin(SortPreservingMergeStream::new( + return Ok(SortPreservingMergeStream::new( Box::new(streams), $schema, $tracking_metrics, @@ -54,12 +54,13 @@ macro_rules! merge_helper { $fetch, $reservation, $enable_round_robin_tie_breaker, - ))); + ) + .into_stream()); }}; } pub struct SortedSpillFile { - pub file: RefCountedTempFile, + pub file: Arc, /// how much memory the largest memory batch is taking pub max_record_batch_memory: usize, @@ -67,12 +68,19 @@ pub struct SortedSpillFile { impl std::fmt::Debug for SortedSpillFile { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!( - f, - "SortedSpillFile({:?}) takes {}", - self.file.path(), - human_readable_size(self.max_record_batch_memory) - ) + match self.file.path() { + Some(path) => write!( + f, + "SortedSpillFile({:?}) takes {}", + path, + human_readable_size(self.max_record_batch_memory) + ), + None => write!( + f, + "SortedSpillFile() takes {}", + human_readable_size(self.max_record_batch_memory) + ), + } } } @@ -187,13 +195,22 @@ impl<'a> StreamingMergeBuilder<'a> { let Some(expressions) = expressions else { return internal_err!("Sort expressions cannot be empty for streaming merge"); }; + let schema = schema.expect("Schema cannot be empty for streaming merge"); + + if fetch.is_some_and(|fetch| fetch == 0) { + return Ok(Box::pin(EmptyRecordBatchStream::new(schema))); + } + + let batch_size = + batch_size.expect("Batch size cannot be empty for streaming merge"); + + if batch_size == 0 { + return internal_err!("Batch size cannot be zero for streaming merge"); + } if !sorted_spill_files.is_empty() { // Unwrapping mandatory fields - let schema = schema.expect("Schema cannot be empty for streaming merge"); let metrics = metrics.expect("Metrics cannot be empty for streaming merge"); - let batch_size = - batch_size.expect("Batch size cannot be empty for streaming merge"); let reservation = reservation.expect("Reservation cannot be empty for streaming merge"); @@ -219,10 +236,7 @@ impl<'a> StreamingMergeBuilder<'a> { ); // Unwrapping mandatory fields - let schema = schema.expect("Schema cannot be empty for streaming merge"); let metrics = metrics.expect("Metrics cannot be empty for streaming merge"); - let batch_size = - batch_size.expect("Batch size cannot be empty for streaming merge"); let reservation = reservation.expect("Reservation cannot be empty for streaming merge"); @@ -247,7 +261,7 @@ impl<'a> StreamingMergeBuilder<'a> { streams, reservation.new_empty(), )?; - Ok(Box::pin(SortPreservingMergeStream::new( + Ok(SortPreservingMergeStream::new( Box::new(streams), schema, metrics, @@ -255,6 +269,114 @@ impl<'a> StreamingMergeBuilder<'a> { fetch, reservation, enable_round_robin_tie_breaker, - ))) + ) + .into_stream()) + } +} + +#[cfg(test)] +mod tests { + use crate::{common::collect, stream::RecordBatchStreamAdapter}; + use std::sync::Arc; + + use super::*; + + use arrow::array::{ArrayRef, RecordBatch}; + use arrow_schema::SortOptions; + use datafusion_common::Result; + use datafusion_execution::TaskContext; + use datafusion_physical_expr::{PhysicalSortExpr, expressions::col}; + use datafusion_physical_expr_common::metrics::{ + ExecutionPlanMetricsSet, SpillMetrics, + }; + + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_only_1_stream() { + test_fetch_0_should_output_0_rows(1, 0).await.unwrap(); + } + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_2_streams() { + test_fetch_0_should_output_0_rows(2, 0).await.unwrap(); + } + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_only_1_spill_file() { + test_fetch_0_should_output_0_rows(0, 1).await.unwrap(); + } + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_2_spill_files() { + test_fetch_0_should_output_0_rows(0, 2).await.unwrap(); + } + #[tokio::test] + async fn test_sort_merge_fetch_zero_with_1_stream_and_1_spill_file() { + test_fetch_0_should_output_0_rows(1, 1).await.unwrap(); + } + + async fn test_fetch_0_should_output_0_rows( + number_of_streams: usize, + number_of_spilled_files: usize, + ) -> Result<()> { + let task_ctx = Arc::new(TaskContext::default()); + let a: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 7, 9, 3])); + let b: ArrayRef = Arc::new(StringArray::from(vec!["a", "b", "c", "d", "e"])); + let batch = RecordBatch::try_from_iter(vec![("a", a), ("b", b)]).unwrap(); + let schema = batch.schema(); + + let sort: LexOrdering = [PhysicalSortExpr { + expr: col("b", &schema).unwrap(), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + + let streams = (0..number_of_streams) + .map(|_| { + Box::pin(RecordBatchStreamAdapter::new( + Arc::clone(&schema), + futures::stream::iter(vec![Ok(batch.clone())]), + )) as SendableRecordBatchStream + }) + .collect::>(); + + let spill_manager = SpillManager::new( + task_ctx.runtime_env(), + SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0), + Arc::clone(&schema), + ); + + let mut sorted_spill_files: Vec = vec![]; + + for _ in 0..number_of_spilled_files { + let file = spill_manager + .spill_record_batch_and_finish(std::slice::from_ref(&batch), "spill") + .unwrap() + .unwrap(); + sorted_spill_files.push(SortedSpillFile { + file, + max_record_batch_memory: batch.get_array_memory_size(), + }); + } + + let sorted_output_stream = StreamingMergeBuilder::new() + .with_batch_size(100) + .with_metrics(BaselineMetrics::new(&ExecutionPlanMetricsSet::new(), 0)) + // Just to avoid having to provide memory pool + .with_bypass_mempool() + .with_schema(schema) + .with_streams(streams) + .with_sorted_spill_files(sorted_spill_files) + .with_spill_manager(spill_manager) + .with_expressions(&sort) + // The whole point of the test - fetch is 0 + .with_fetch(Some(0)) + .build() + .unwrap(); + + let collected = collect(sorted_output_stream).await.unwrap(); + let total: usize = collected.iter().map(|b| b.num_rows()).sum(); + assert_eq!(total, 0, "fetch=Some(0) must emit zero rows, got {total}"); + + Ok(()) } } diff --git a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs index e0548bd5bf860..71d7cce1bcc7d 100644 --- a/datafusion/physical-plan/src/spill/in_progress_spill_file.rs +++ b/datafusion/physical-plan/src/spill/in_progress_spill_file.rs @@ -22,7 +22,7 @@ use std::sync::Arc; use arrow::array::RecordBatch; use datafusion_common::exec_datafusion_err; -use datafusion_execution::disk_manager::RefCountedTempFile; +use datafusion_execution::spill_file::SpillFile; use super::{ IPCStreamWriter, gc_view_arrays, @@ -37,13 +37,13 @@ pub struct InProgressSpillFile { /// Lazily initialized writer writer: Option, /// Lazily initialized in-progress file, it will be moved out when the `finish` method is invoked - in_progress_file: Option, + in_progress_file: Option>, } impl InProgressSpillFile { pub fn new( spill_writer: Arc, - in_progress_file: RefCountedTempFile, + in_progress_file: Arc, ) -> Self { Self { spill_writer, @@ -79,40 +79,27 @@ impl InProgressSpillFile { // when they come from different branches of a UnionExec. The SpillManager's // schema represents the canonical schema that all batches should conform to. let schema = self.spill_writer.schema(); - if let Some(in_progress_file) = &mut self.in_progress_file { + if let Some(in_progress_file) = &self.in_progress_file { + let spill_writer = in_progress_file.open_writer()?; + self.writer = Some(IPCStreamWriter::new( - in_progress_file.path(), + spill_writer, schema.as_ref(), self.spill_writer.compression, )?); // Update metrics self.spill_writer.metrics.spill_file_count.add(1); - - // Update initial size (schema/header) - in_progress_file.update_disk_usage()?; - let initial_size = in_progress_file.current_disk_usage(); - self.spill_writer - .metrics - .spilled_bytes - .add(initial_size as usize); + let header_bytes = self.writer.as_ref().unwrap().bytes_written(); + self.spill_writer.metrics.spilled_bytes.add(header_bytes); } } if let Some(writer) = &mut self.writer { - let (spilled_rows, _) = writer.write(&gc_batch)?; - if let Some(in_progress_file) = &mut self.in_progress_file { - let pre_size = in_progress_file.current_disk_usage(); - in_progress_file.update_disk_usage()?; - let post_size = in_progress_file.current_disk_usage(); - - self.spill_writer.metrics.spilled_rows.add(spilled_rows); - self.spill_writer - .metrics - .spilled_bytes - .add((post_size - pre_size) as usize); - } else { - unreachable!() // Already checked inside current function - } + // The writer calculates how many serialized bytes were emitted + let (spilled_rows, delta_bytes) = writer.write(&gc_batch)?; + + self.spill_writer.metrics.spilled_rows.add(spilled_rows); + self.spill_writer.metrics.spilled_bytes.add(delta_bytes); } gc_batch.get_sliced_size() } @@ -126,31 +113,26 @@ impl InProgressSpillFile { /// Returns a reference to the in-progress file, if it exists. /// This can be used to get the file path for creating readers before the file is finished. - pub fn file(&self) -> Option<&RefCountedTempFile> { + pub fn file(&self) -> Option<&Arc> { self.in_progress_file.as_ref() } - /// Finalizes the file, returning the completed file reference. + /// Finalizes the write process, returning the completed `SpillFile`. /// If there are no batches spilled before, it returns `None`. - pub fn finish(&mut self) -> Result> { - if let Some(writer) = &mut self.writer { - writer.finish()?; + pub fn finish(&mut self) -> Result>> { + if self.in_progress_file.is_none() && self.writer.is_none() { + return Err(exec_datafusion_err!( + "Finish operation failed: file has already been finalized." + )); + } + if let Some(mut writer) = self.writer.take() { + // Finish the writer and capture any final trailing bytes emitted + let delta_bytes = writer.finish()?; + self.spill_writer.metrics.spilled_bytes.add(delta_bytes); } else { return Ok(None); } - // Since spill files are append-only, add the file size to spilled_bytes - if let Some(in_progress_file) = &mut self.in_progress_file { - // Since writer.finish() writes continuation marker and message length at the end - let pre_size = in_progress_file.current_disk_usage(); - in_progress_file.update_disk_usage()?; - let post_size = in_progress_file.current_disk_usage(); - self.spill_writer - .metrics - .spilled_bytes - .add((post_size - pre_size) as usize); - } - Ok(self.in_progress_file.take()) } } diff --git a/datafusion/physical-plan/src/spill/mod.rs b/datafusion/physical-plan/src/spill/mod.rs index 3c95a1da5b33c..addcf78d2df84 100644 --- a/datafusion/physical-plan/src/spill/mod.rs +++ b/datafusion/physical-plan/src/spill/mod.rs @@ -21,16 +21,13 @@ pub(crate) mod in_progress_spill_file; pub(crate) mod replayable_spill_input; pub(crate) mod spill_manager; pub mod spill_pool; - +use datafusion_execution::spill_file::SpillWriter; // Moved for refactor, re-export to keep the public API stable pub use datafusion_common::utils::memory::get_record_batch_memory_size; // Re-export SpillManager for doctests only (hidden from public docs) #[doc(hidden)] pub use spill_manager::SpillManager; -use std::fs::File; -use std::io::BufReader; -use std::path::{Path, PathBuf}; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; @@ -39,254 +36,216 @@ use arrow::array::{ Array, ArrayRef, BinaryViewArray, BufferSpec, GenericByteViewArray, StringViewArray, layout, make_array, }; +use arrow::buffer::Buffer; use arrow::datatypes::DataType; use arrow::datatypes::{ByteViewType, Schema, SchemaRef}; use arrow::ipc::{ MetadataVersion, - reader::StreamReader, + reader::StreamDecoder, writer::{IpcWriteOptions, StreamWriter}, }; use arrow::record_batch::RecordBatch; use arrow_data::ArrayDataBuilder; use arrow_ipc::CompressionType; +use datafusion_common::Result; use datafusion_common::config::SpillCompression; -use datafusion_common::{DataFusionError, Result, exec_datafusion_err, exec_err}; -use datafusion_common_runtime::SpawnedTask; use datafusion_execution::RecordBatchStream; -use datafusion_execution::disk_manager::RefCountedTempFile; -use futures::{FutureExt as _, Stream}; +use datafusion_execution::spill_file::SpillFile; +use futures::Stream; use log::debug; -/// Stream that reads spill files from disk where each batch is read in a spawned blocking task -/// It will read one batch at a time and will not do any buffering, to buffer data use [`crate::common::spawn_buffered`] -/// -/// A simpler solution would be spawning a long-running blocking task for each -/// file read (instead of each batch). This approach does not work because when -/// the number of concurrent reads exceeds the Tokio thread pool limit, -/// deadlocks can occur and block progress. +/// Stream that reads spill files from a [`SpillFile`] backend as a stream of [`RecordBatch`]es. +/// Uses [`StreamDecoder`] to decode IPC bytes received from the backend's async byte stream. +/// Backends handle their own threading concerns internally - OS files use +/// `tokio::fs::File` which performs blocking IO per-syscall without holding a thread +/// for the file's lifetime, avoiding deadlocks when concurrent reads exceed thread pool limits. struct SpillReaderStream { schema: SchemaRef, - state: SpillReaderStreamState, + decoder: StreamDecoder, + byte_stream: Pin> + Send>>, + is_done: bool, + /// Maximum memory size observed among spilling sorted record batches. /// This is used for validation purposes during reading each RecordBatch from spill. /// For context on why this value is recorded and validated, /// see `physical_plan/sort/multi_level_merge.rs`. max_record_batch_memory: Option, -} - -// Small margin allowed to accommodate slight memory accounting variation -const SPILL_BATCH_MEMORY_MARGIN: usize = 4096; - -/// When we poll for the next batch, we will get back both the batch and the reader, -/// so we can call `next` again. -type NextRecordBatchResult = Result<(StreamReader>, Option)>; - -enum SpillReaderStreamState { - /// Initial state: the stream was not initialized yet - /// and the file was not opened - Uninitialized(RefCountedTempFile), - /// A read is in progress in a spawned blocking task for which we hold the handle. - ReadInProgress(SpawnedTask), + /// Holds leftover bytes from a chunk when a batch is yielded early + current_buffer: Buffer, - /// A read has finished and we wait for being polled again in order to start reading the next batch. - Waiting(StreamReader>), + /// Keeps the file alive until the stream is dropped + _spill_file: Arc, - /// The stream has finished, successfully or not. - Done, + schema_validated: bool, } +// Small margin allowed to accommodate slight memory accounting variation +const SPILL_BATCH_MEMORY_MARGIN: usize = 4096; + impl SpillReaderStream { fn new( schema: SchemaRef, - spill_file: RefCountedTempFile, + spill_file: Arc, max_record_batch_memory: Option, - ) -> Self { - Self { + ) -> Result { + let byte_stream = spill_file.read_stream()?; + // DataFusion controls what it writes so it can trust its own IPC output, + // matching the behavior of the previous StreamReader-based implementation. + let decoder = unsafe { StreamDecoder::new().with_skip_validation(true) }; + Ok(Self { schema, - state: SpillReaderStreamState::Uninitialized(spill_file), + decoder, + byte_stream, max_record_batch_memory, - } + is_done: false, + current_buffer: Buffer::from(&[]), + _spill_file: spill_file, + schema_validated: false, + }) } +} - fn poll_next_inner( - &mut self, - cx: &mut Context<'_>, - ) -> Poll>> { - match &mut self.state { - SpillReaderStreamState::Uninitialized(_) => { - // Temporarily replace with `Done` to be able to pass the file to the task. - let SpillReaderStreamState::Uninitialized(spill_file) = - std::mem::replace(&mut self.state, SpillReaderStreamState::Done) - else { - unreachable!() - }; - - let expected_schema = Arc::clone(&self.schema); - let task = SpawnedTask::spawn_blocking(move || { - let file = BufReader::new(File::open(spill_file.path())?); - // SAFETY: DataFusion's spill writer strictly follows Arrow IPC specifications - // with validated schemas and buffers. Skip redundant validation during read - // to speedup read operation. This is safe for DataFusion as input guaranteed to be correct when written. - let mut reader = unsafe { - StreamReader::try_new(file, None)?.with_skip_validation(true) - }; - - // Validate the schema read from Arrow IPC file is the same as the - // schema of the current `SpillManager` - let actual_schema = reader.schema(); - - if actual_schema != expected_schema { - return exec_err!( - "Spill file schema mismatch: expected {}, got {}. \ - The caller must use the same SpillManager that created the spill file to read it.", - expected_schema, - actual_schema - ); - } - - // TODO: Same-schema reads from a different SpillManager still pass today. - // Add a SpillManager UID to IPC metadata and validate it here as well. - let next_batch = reader.next().transpose()?; - - Ok((reader, next_batch)) - }); +impl Stream for SpillReaderStream { + type Item = Result; - self.state = SpillReaderStreamState::ReadInProgress(task); + fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { + let this = self.get_mut(); - // Poll again immediately so the inner task is polled and the waker is - // registered. - self.poll_next_inner(cx) - } + if this.is_done { + return Poll::Ready(None); + } - SpillReaderStreamState::ReadInProgress(task) => { - let result = futures::ready!(task.poll_unpin(cx)) - .unwrap_or_else(|err| Err(DataFusionError::External(Box::new(err)))); - - match result { - Ok((reader, batch)) => { - match batch { - Some(batch) => { - if let Some(max_record_batch_memory) = - self.max_record_batch_memory - { - let actual_size = - get_record_batch_memory_size(&batch); - if actual_size - > max_record_batch_memory - + SPILL_BATCH_MEMORY_MARGIN - { - debug!( - "Record batch memory usage ({actual_size} bytes) exceeds the expected limit ({max_record_batch_memory} bytes) \n\ - by more than the allowed tolerance ({SPILL_BATCH_MEMORY_MARGIN} bytes).\n\ - This likely indicates a bug in memory accounting during spilling.\n\ - Please report this issue in https://github.com/apache/datafusion/issues/17340." - ); - } - } - self.state = SpillReaderStreamState::Waiting(reader); - - Poll::Ready(Some(Ok(batch))) + loop { + if !this.current_buffer.is_empty() { + match this.decoder.decode(&mut this.current_buffer) { + Ok(Some(batch)) => { + // One-time schema validation on the first decoded batch. + // The IPC stream embeds the writer's schema in its header; + // StreamDecoder surfaces it via the first batch's schema. + // We check here rather than in new() because schema bytes + // only arrive after decoding the IPC header from the stream. + if !this.schema_validated { + this.schema_validated = true; + let actual = batch.schema(); + if actual != this.schema { + this.is_done = true; + return Poll::Ready(Some(Err( + datafusion_common::exec_datafusion_err!( + "Spill file schema mismatch: expected {}, got {}. \ + The caller must use the same SpillManager that created \ + the spill file to read it.", + this.schema, + actual + ), + ))); } - None => { - // Stream is done - self.state = SpillReaderStreamState::Done; - - Poll::Ready(None) + } + if let Some(max_record_batch_memory) = + this.max_record_batch_memory + { + let actual_size = get_record_batch_memory_size(&batch); + if actual_size + > max_record_batch_memory + SPILL_BATCH_MEMORY_MARGIN + { + debug!( + "Record batch memory usage ({actual_size} bytes) exceeds the expected limit ({max_record_batch_memory} bytes) \n\ + by more than the allowed tolerance ({SPILL_BATCH_MEMORY_MARGIN} bytes).\n\ + This likely indicates a bug in memory accounting during spilling." + ); } } + return Poll::Ready(Some(Ok(batch))); } - Err(err) => { - self.state = SpillReaderStreamState::Done; - - Poll::Ready(Some(Err(err))) + Ok(None) => { + // The chunk didn't form a complete message. Arrow consumed the partial bytes + // into its internal scratch pad, leaving our current_buffer completely empty. + // We do nothing and fall through to fetch more data. + } + Err(e) => { + this.is_done = true; + return Poll::Ready(Some(Err(e.into()))); } } } - SpillReaderStreamState::Waiting(_) => { - // Temporarily replace with `Done` to be able to pass the file to the task. - let SpillReaderStreamState::Waiting(mut reader) = - std::mem::replace(&mut self.state, SpillReaderStreamState::Done) - else { - unreachable!() - }; - - let task = SpawnedTask::spawn_blocking(move || { - let next_batch = reader.next().transpose()?; - - Ok((reader, next_batch)) - }); - - self.state = SpillReaderStreamState::ReadInProgress(task); + match futures::ready!(this.byte_stream.as_mut().poll_next(cx)) { + Some(Ok(chunk)) => { + this.current_buffer = Buffer::from(chunk); + } + Some(Err(e)) => { + this.is_done = true; + return Poll::Ready(Some(Err(e))); + } + None => { + this.is_done = true; - // Poll again immediately so the inner task is polled and the waker is - // registered. - self.poll_next_inner(cx) + if let Err(e) = this.decoder.finish() { + return Poll::Ready(Some(Err(e.into()))); + } + return Poll::Ready(None); + } } - - SpillReaderStreamState::Done => Poll::Ready(None), } } } -impl Stream for SpillReaderStream { - type Item = Result; - - fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll> { - self.get_mut().poll_next_inner(cx) - } -} - impl RecordBatchStream for SpillReaderStream { fn schema(&self) -> SchemaRef { Arc::clone(&self.schema) } } -/// Spill the `RecordBatch` to disk as smaller batches -/// split by `batch_size_rows` -#[deprecated( - since = "46.0.0", - note = "This method is deprecated. Use `SpillManager::spill_record_batch_by_size` instead." -)] -#[expect(clippy::needless_pass_by_value)] -pub fn spill_record_batch_by_size( - batch: &RecordBatch, - path: PathBuf, - schema: SchemaRef, - batch_size_rows: usize, -) -> Result<()> { - let mut offset = 0; - let total_rows = batch.num_rows(); - let mut writer = - IPCStreamWriter::new(&path, schema.as_ref(), SpillCompression::Uncompressed)?; - - while offset < total_rows { - let length = std::cmp::min(total_rows - offset, batch_size_rows); - let batch = batch.slice(offset, length); - offset += batch.num_rows(); - writer.write(&batch)?; +/// A wrapper that counts the exact compressed IPC bytes written by Arrow. +/// +/// Arrow's `StreamWriter` does not return the number of bytes written during its +/// `write()` calls. To accurately track the `spilled_bytes` metrics (especially +/// when LZ4/ZSTD compression is applied), we must intercept the `std::io::Write` +/// trait boundary to count the final serialized payload size. +pub(crate) struct TrackingSpillWriter { + inner: Box, + pub(crate) total_bytes_written: usize, +} + +impl TrackingSpillWriter { + pub fn new(inner: Box) -> Self { + Self { + inner, + total_bytes_written: 0, + } } - writer.finish()?; - Ok(()) + pub fn finish(mut self) -> Result<()> { + self.inner.finish() + } } -/// Write in Arrow IPC Stream format to a file. -/// -/// Stream format is used for spill because it supports dictionary replacement, and the random -/// access of IPC File format is not needed (IPC File format doesn't support dictionary replacement). +impl std::io::Write for TrackingSpillWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let n = self.inner.write(buf)?; + + self.total_bytes_written += n; + + Ok(n) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + +/// Write in Arrow IPC Stream format to an underlying `SpillWriter` backend. +/// Stream format also supports dictionary replacement. struct IPCStreamWriter { /// Inner writer - pub writer: StreamWriter, + writer: Option>, /// Batches written - pub num_batches: usize, + num_batches: usize, /// Rows written - pub num_rows: usize, + num_rows: usize, /// Bytes written - pub num_bytes: usize, + num_bytes: usize, } impl IPCStreamWriter { @@ -303,14 +262,10 @@ impl IPCStreamWriter { /// rather than relying solely on workspace-level feature unification; /// see #21917. pub fn new( - path: &Path, + spill_writer: Box, schema: &Schema, spill_compression: SpillCompression, ) -> Result { - let file = File::create(path).map_err(|e| { - exec_datafusion_err!("(Hint: you may increase the file descriptor limit with shell command 'ulimit -n 4096') Failed to create partition file at {path:?}: {e:?}") - })?; - let metadata_version = MetadataVersion::V5; // Depending on the schema, some array types such as StringViewArray require larger (16 byte in this case) alignment. // If the actual buffer layout after IPC read does not satisfy the alignment requirement, @@ -320,15 +275,18 @@ impl IPCStreamWriter { let alignment = get_max_alignment_for_schema(schema); let mut write_options = IpcWriteOptions::try_new(alignment, false, metadata_version)?; + let compression_type = Option::::from(spill_compression); write_options = write_options.try_with_compression(compression_type)?; - let writer = StreamWriter::try_new_with_options(file, schema, write_options)?; + let adapter = TrackingSpillWriter::new(spill_writer); + let writer = StreamWriter::try_new_with_options(adapter, schema, write_options)?; + Ok(Self { num_batches: 0, num_rows: 0, num_bytes: 0, - writer, + writer: Some(writer), }) } @@ -336,23 +294,50 @@ impl IPCStreamWriter { /// /// Returns a tuple containing the change in the number of rows and bytes written. pub fn write(&mut self, batch: &RecordBatch) -> Result<(usize, usize)> { - self.writer.write(batch)?; + let writer = self.writer.as_mut().unwrap(); + + let bytes_before = writer.get_ref().total_bytes_written; + writer.write(batch)?; + let bytes_after = writer.get_ref().total_bytes_written; self.num_batches += 1; let delta_num_rows = batch.num_rows(); self.num_rows += delta_num_rows; - let delta_num_bytes: usize = batch.get_array_memory_size(); + let delta_num_bytes = bytes_after - bytes_before; self.num_bytes += delta_num_bytes; Ok((delta_num_rows, delta_num_bytes)) } pub fn flush(&mut self) -> Result<()> { - self.writer.flush()?; + use std::io::Write; + if let Some(writer) = &mut self.writer { + writer.get_mut().flush()?; + } Ok(()) } - /// Finish the writer - pub fn finish(&mut self) -> Result<()> { - self.writer.finish().map_err(Into::into) + /// Finish the writer. + /// + /// Returns the number of trailing bytes written during the finish operation + /// (e.g., IPC metadata and footers). + pub fn finish(&mut self) -> Result { + let mut writer = self.writer.take().unwrap(); + + let bytes_before = writer.get_ref().total_bytes_written; + writer.finish()?; // Writes IPC tail + + // Extract the adapter and flush the final bytes + let adapter = writer.into_inner()?; + let bytes_after = adapter.total_bytes_written; + adapter.finish()?; + + Ok(bytes_after - bytes_before) + } + /// Returns the total number of bytes written so far + pub fn bytes_written(&self) -> usize { + self.writer + .as_ref() + .map(|w| w.get_ref().total_bytes_written) + .unwrap_or(0) } } @@ -529,7 +514,7 @@ fn calculate_string_view_waste_ratio(array: &StringViewArray) -> f64 { #[cfg(test)] fn calculate_view_waste_ratio( len: usize, - data_buffers: &[arrow::buffer::Buffer], + data_buffers: &[Buffer], get_value_size: F, ) -> f64 where @@ -587,7 +572,7 @@ mod tests { let spill_file = spill_manager .spill_record_batch_and_finish(&[batch1, batch2], "Test")? .unwrap(); - assert!(spill_file.path().exists()); + assert!(spill_file.path().unwrap().exists()); let spilled_rows = spill_manager.metrics.spilled_rows.value(); assert_eq!(spilled_rows, num_rows); @@ -684,7 +669,7 @@ mod tests { "Test Spill", )? .unwrap(); - assert!(spill_file.path().exists()); + assert!(spill_file.path().unwrap().exists()); assert!(max_batch_mem > 0); let stream = spill_manager.read_spill_as_stream(spill_file, None)?; @@ -714,7 +699,7 @@ mod tests { async fn validate( spill_manager: &SpillManager, - spill_file: RefCountedTempFile, + spill_file: Arc, num_rows: usize, schema: SchemaRef, batch_count: usize, @@ -764,14 +749,14 @@ mod tests { let zstd_spill_file = zstd_spill_manager .spill_record_batch_and_finish(&batches, "ZSTD_Test")? .unwrap(); - assert!(uncompressed_spill_file.path().exists()); - assert!(lz4_spill_file.path().exists()); - assert!(zstd_spill_file.path().exists()); + assert!(uncompressed_spill_file.path().unwrap().exists()); + assert!(lz4_spill_file.path().unwrap().exists()); + assert!(zstd_spill_file.path().unwrap().exists()); - let lz4_spill_size = std::fs::metadata(lz4_spill_file.path())?.len(); - let zstd_spill_size = std::fs::metadata(zstd_spill_file.path())?.len(); + let lz4_spill_size = std::fs::metadata(lz4_spill_file.path().unwrap())?.len(); + let zstd_spill_size = std::fs::metadata(zstd_spill_file.path().unwrap())?.len(); let uncompressed_spill_size = - std::fs::metadata(uncompressed_spill_file.path())?.len(); + std::fs::metadata(uncompressed_spill_file.path().unwrap())?.len(); assert!(uncompressed_spill_size > lz4_spill_size); assert!(uncompressed_spill_size > zstd_spill_size); @@ -826,7 +811,7 @@ mod tests { let temp_file = spill_manager.spill_record_batch_and_finish(&[batch], "Test")?; assert!(temp_file.is_some()); - assert!(temp_file.unwrap().path().exists()); + assert!(temp_file.unwrap().path().unwrap().exists()); Ok(()) } @@ -900,7 +885,7 @@ mod tests { let completed_file = in_progress_file.finish()?; assert!(completed_file.is_some()); - assert!(completed_file.unwrap().path().exists()); + assert!(completed_file.unwrap().path().unwrap().exists()); verify_metrics(&in_progress_file, 1, 712, 6)?; // Double finish produce error let result = in_progress_file.finish(); @@ -1314,7 +1299,7 @@ mod tests { } let spill_file = in_progress_file.finish()?.unwrap(); - let file_size = fs::metadata(spill_file.path())?.len() as usize; + let file_size = fs::metadata(spill_file.path().unwrap())?.len() as usize; let theoretical_without_gc = total_buffer_size * sliced_batches.len(); let reduction_percent = ((theoretical_without_gc - file_size) as f64 @@ -1484,7 +1469,7 @@ mod tests { .unwrap(); // 4. Check file size on disk - let file_size = fs::metadata(spill_file.path())?.len(); + let file_size = fs::metadata(spill_file.path().unwrap())?.len(); // The original buffer size is around 70KB. // Without GC, the spill file would be > 70KB. @@ -1528,7 +1513,7 @@ mod tests { .unwrap(); // 4. Check file size on disk - let file_size = fs::metadata(spill_file.path())?.len(); + let file_size = fs::metadata(spill_file.path().unwrap())?.len(); // Original buffer is 100KB. // With GC, it should be much smaller. diff --git a/datafusion/physical-plan/src/spill/replayable_spill_input.rs b/datafusion/physical-plan/src/spill/replayable_spill_input.rs index fea998d268c59..94a0aef7dcc6f 100644 --- a/datafusion/physical-plan/src/spill/replayable_spill_input.rs +++ b/datafusion/physical-plan/src/spill/replayable_spill_input.rs @@ -26,9 +26,8 @@ use std::task::{Context, Poll}; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::{Result, internal_err}; -use datafusion_execution::RecordBatchStream; use datafusion_execution::SendableRecordBatchStream; -use datafusion_execution::disk_manager::RefCountedTempFile; +use datafusion_execution::{RecordBatchStream, SpillFile}; use futures::Stream; use parking_lot::Mutex; @@ -91,7 +90,7 @@ pub(crate) struct ReplayableStreamSource { /// Inner state exclusively owned by either [`ReplayableStreamSource`] or one [`ReplayableSpillStream`] enum StateInner { Unopened, - Replayable(Option), + Replayable(Option>), Poisoned, } @@ -222,10 +221,10 @@ impl ReplayableSpillStream { schema: SchemaRef, spill_manager: &SpillManager, shared_state: Arc>>, - spill_file: Option, + spill_file: Option>, ) -> Result { let inner = if let Some(file) = spill_file.as_ref() { - spill_manager.read_spill_as_stream(file.clone(), None)? + spill_manager.read_spill_as_stream(Arc::clone(file), None)? } else { Box::pin(EmptyRecordBatchStream::new(Arc::clone(&schema))) }; diff --git a/datafusion/physical-plan/src/spill/spill_manager.rs b/datafusion/physical-plan/src/spill/spill_manager.rs index 365a9f977eace..aee9e917c755d 100644 --- a/datafusion/physical-plan/src/spill/spill_manager.rs +++ b/datafusion/physical-plan/src/spill/spill_manager.rs @@ -25,8 +25,8 @@ use arrow::datatypes::{ByteViewType, SchemaRef}; use arrow::record_batch::RecordBatch; use datafusion_common::{DataFusionError, Result, config::SpillCompression}; use datafusion_execution::SendableRecordBatchStream; -use datafusion_execution::disk_manager::RefCountedTempFile; use datafusion_execution::runtime_env::RuntimeEnv; +use datafusion_execution::spill_file::SpillFile; use std::borrow::Borrow; use std::sync::Arc; @@ -76,6 +76,10 @@ impl SpillManager { &self.schema } + pub(crate) fn env(&self) -> &RuntimeEnv { + &self.env + } + /// Creates a temporary file for in-progress operations, returning an error /// message if file creation fails. The file can be used to append batches /// incrementally and then finish the file when done. @@ -99,7 +103,7 @@ impl SpillManager { &self, batches: &[RecordBatch], request_msg: &str, - ) -> Result> { + ) -> Result>> { let mut in_progress_file = self.create_in_progress_file(request_msg)?; for batch in batches { @@ -115,7 +119,7 @@ impl SpillManager { &self, mut iter: impl Iterator>>, request_description: &str, - ) -> Result> { + ) -> Result, usize)>> { let mut in_progress_file = self.create_in_progress_file(request_description)?; let mut max_record_batch_size = 0; @@ -141,7 +145,7 @@ impl SpillManager { &self, stream: &mut SendableRecordBatchStream, request_description: &str, - ) -> Result> { + ) -> Result, usize)>> { use futures::StreamExt; let mut in_progress_file = self.create_in_progress_file(request_description)?; @@ -178,14 +182,14 @@ impl SpillManager { /// the merge degree when merging multiple sorted runs. pub fn read_spill_as_stream( &self, - spill_file_path: RefCountedTempFile, + spill_file_path: Arc, max_record_batch_memory: Option, ) -> Result { let stream = Box::pin(cooperative(SpillReaderStream::new( Arc::clone(&self.schema), spill_file_path, max_record_batch_memory, - ))); + )?)); Ok(spawn_buffered(stream, self.batch_read_buffer_capacity)) } @@ -193,14 +197,14 @@ impl SpillManager { /// Same as `read_spill_as_stream`, but without buffering. pub fn read_spill_as_stream_unbuffered( &self, - spill_file_path: RefCountedTempFile, + spill_file_path: Arc, max_record_batch_memory: Option, ) -> Result { Ok(Box::pin(cooperative(SpillReaderStream::new( Arc::clone(&self.schema), spill_file_path, max_record_batch_memory, - )))) + )?))) } } diff --git a/datafusion/physical-plan/src/spill/spill_pool.rs b/datafusion/physical-plan/src/spill/spill_pool.rs index 2639188a2609d..6e964d7a6497b 100644 --- a/datafusion/physical-plan/src/spill/spill_pool.rs +++ b/datafusion/physical-plan/src/spill/spill_pool.rs @@ -17,6 +17,7 @@ use futures::{Stream, StreamExt}; use std::collections::VecDeque; +use std::mem; use std::sync::Arc; use std::task::Waker; @@ -25,8 +26,7 @@ use parking_lot::Mutex; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::Result; -use datafusion_execution::disk_manager::RefCountedTempFile; -use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream}; +use datafusion_execution::{RecordBatchStream, SendableRecordBatchStream, SpillFile}; use super::in_progress_spill_file::InProgressSpillFile; use super::spill_manager::SpillManager; @@ -48,7 +48,7 @@ use super::spill_manager::SpillManager; /// **Lock ordering discipline**: Never hold both locks simultaneously to prevent deadlock. /// Always: acquire outer lock → release outer lock → acquire inner lock (if needed). struct SpillPoolShared { - /// Queue of ALL files (including the current write file if it exists). + /// Queue of ALL files (including the current write files if any exist). /// Readers always read from the front of this queue (FIFO). /// Each file has its own lock to enable concurrent reader/writer access. files: VecDeque>>, @@ -56,15 +56,14 @@ struct SpillPoolShared { spill_manager: Arc, /// Pool-level waker to notify when new files are available (single reader) waker: Option, - /// Whether the writer has been dropped (no more files will be added) - writer_dropped: bool, - /// Writer's reference to the current file (shared by all cloned writers). - /// Has its own lock to allow I/O without blocking queue access. - current_write_file: Option>>, - /// Number of active writer clones. Only when this reaches zero should - /// `writer_dropped` be set to true. This prevents premature EOF signaling - /// when one writer clone is dropped while others are still active. - active_writer_count: usize, + /// FIFO queue of open write files. The queue may contain multiple items when multiple + /// writers concurrently write to the pool. + /// Each write file has its own lock to allow I/O without blocking queue access. + open_write_files: VecDeque>>, + /// Number of `SpillPoolWriter` instances that have not been dropped yet. As long as this value + /// is greater than zero, readers should assume batches may still be pushed. This prevents + /// premature EOF signaling. + remaining_writer_count: usize, } impl SpillPoolShared { @@ -74,9 +73,8 @@ impl SpillPoolShared { files: VecDeque::new(), spill_manager, waker: None, - writer_dropped: false, - current_write_file: None, - active_writer_count: 1, + open_write_files: VecDeque::new(), + remaining_writer_count: 1, } } @@ -93,69 +91,112 @@ impl SpillPoolShared { } } -/// Writer for a spill pool. Provides coordinated write access with FIFO semantics. +/// Writer for a spill pool that can be cloned to produce additional writers. /// -/// Created by [`channel`]. See that function for architecture diagrams and usage examples. -/// -/// The writer is `Clone`, allowing multiple writers to coordinate on the same pool. -/// All clones share the same current write file and coordinate file rotation. -/// The writer automatically manages file rotation based on the `max_file_size_bytes` -/// configured in [`channel`]. When the last writer clone is dropped, it finalizes the -/// current file so readers can access all written data. +/// Created by [`mpsc_channel`]. See that function for architecture diagrams and usage +/// examples. pub struct SpillPoolWriter { - /// Maximum size in bytes before rotating to a new file. - /// Typically set from configuration `datafusion.execution.max_spill_file_size_bytes`. - max_file_size_bytes: usize, - /// Shared state with readers (includes current_write_file for coordination) - shared: Arc>, + /// The underlying shared writer. Kept private and never cloned, so this pool always has + /// exactly one writer. + inner: SpillPoolSink, +} + +impl SpillPoolWriter { + /// Spills a batch to the pool, rotating files when necessary. + /// + /// See [`mpsc_channel`] for the rotation semantics. + /// + /// # Errors + /// + /// Returns an error if disk I/O fails or disk quota is exceeded. + pub fn push_batch(&self, batch: &RecordBatch) -> Result<()> { + self.inner.push_batch(batch) + } +} + +impl SpillPoolWriter { + /// Returns a new sink that can be used to spill batches to the pool. + /// + /// As an alternative to this function, it is also possible to clone the writer. The benefit + /// of this method is that the output type matches the type used by [`spsc_channel`]. This + /// enables cost-free abstraction for producers over SPSC and MPSC channels. + pub fn new_sink(&self) -> SpillPoolSink { + // Increment `remaining_writer_count`. The corresponding decrement is done in the `Drop` + // implementation of `SpillPoolWriter`. + self.inner.shared.lock().remaining_writer_count += 1; + SpillPoolSink { + max_file_size_bytes: self.inner.max_file_size_bytes, + shared: Arc::clone(&self.inner.shared), + } + } } impl Clone for SpillPoolWriter { fn clone(&self) -> Self { - // Increment the active writer count so that `writer_dropped` is only - // set to true when the *last* clone is dropped. - self.shared.lock().active_writer_count += 1; Self { - max_file_size_bytes: self.max_file_size_bytes, - shared: Arc::clone(&self.shared), + inner: self.new_sink(), } } } -impl SpillPoolWriter { +impl Drop for SpillPoolSink { + fn drop(&mut self) { + let mut shared = self.shared.lock(); + + shared.remaining_writer_count -= 1; + let is_last_writer = shared.remaining_writer_count == 0; + + if !is_last_writer { + // Other writer clones are still active; do not finalize or + // signal EOF to readers. + return; + } + + // Finalize any spill files that were not finished yet + if !shared.open_write_files.is_empty() { + let files = mem::take(&mut shared.open_write_files); + drop(shared); + + for file in files { + let mut file_shared = file.lock(); + + // Finish the current writer if it exists + if let Some(mut writer) = file_shared.writer.take() { + // Ignore errors on drop - we're in destructor + let _ = writer.finish(); + } + + // Mark as finished so readers know not to wait for more data + file_shared.writer_finished = true; + + // Wake reader waiting on this file (it's now finished) + file_shared.wake(); + drop(file_shared); + } + + shared = self.shared.lock(); + } + + // Wake pool-level readers + shared.wake(); + } +} + +/// Single writer for a spill pool that cannot be cloned. +/// +/// Created by [`spsc_channel`] and [`SpillPoolWriter::new_sink`]. +pub struct SpillPoolSink { + /// Maximum size in bytes before rotating to a new file. + /// Typically set from configuration `datafusion.execution.max_spill_file_size_bytes`. + max_file_size_bytes: usize, + /// Shared state with readers (includes current_write_file for coordination) + shared: Arc>, +} + +impl SpillPoolSink { /// Spills a batch to the pool, rotating files when necessary. /// - /// If the current file would exceed `max_file_size_bytes` after adding - /// this batch, the file is finalized and a new one is started. - /// - /// See [`channel`] for overall architecture and examples. - /// - /// # File Rotation Logic - /// - /// ```text - /// push_batch() - /// │ - /// ▼ - /// Current file exists? - /// │ - /// ├─ No ──▶ Create new file ──▶ Add to shared queue - /// │ Wake readers - /// ▼ - /// Write batch to current file - /// │ - /// ▼ - /// estimated_size > max_file_size_bytes? - /// │ - /// ├─ No ──▶ Keep current file for next batch - /// │ - /// ▼ - /// Yes: finish() current file - /// Mark writer_finished = true - /// Wake readers - /// │ - /// ▼ - /// Next push_batch() creates new file - /// ``` + /// See [`spsc_channel`] for overall architecture and examples. /// /// # Errors /// @@ -171,15 +212,19 @@ impl SpillPoolWriter { // Fine-grained locking: Lock shared state briefly for queue access let mut shared = self.shared.lock(); - // Create new file if we don't have one yet - if shared.current_write_file.is_none() { + // Create new file if there is none available to append to + let write_file = if !shared.open_write_files.is_empty() { + shared.open_write_files.pop_front().unwrap() + } else { let spill_manager = Arc::clone(&shared.spill_manager); // Release shared lock before disk I/O (fine-grained locking) drop(shared); let writer = spill_manager.create_in_progress_file("SpillPool")?; // Clone the file so readers can access it immediately - let file = writer.file().expect("InProgressSpillFile should always have a file when it is first created").clone(); + let file = Arc::clone(writer.file().expect( + "InProgressSpillFile should always have a file when it is first created", + )); let file_shared = Arc::new(Mutex::new(ActiveSpillFileShared { writer: Some(writer), @@ -193,107 +238,62 @@ impl SpillPoolWriter { // Re-acquire lock and push to shared queue shared = self.shared.lock(); shared.files.push_back(Arc::clone(&file_shared)); - shared.current_write_file = Some(file_shared); shared.wake(); // Wake readers waiting for new files - } + file_shared + }; - let current_write_file = shared.current_write_file.take(); // Release shared lock before file I/O (fine-grained locking) // This allows readers to access the queue while we do disk I/O drop(shared); // Write batch to current file - lock only the specific file - if let Some(current_file) = current_write_file { - // Now lock just this file for I/O (separate from shared lock) - let mut file_shared = current_file.lock(); - - // Append the batch - if let Some(ref mut writer) = file_shared.writer { - writer.append_batch(batch)?; - // make sure we flush the writer for readers - writer.flush()?; - file_shared.batches_written += 1; - file_shared.estimated_size += batch_size; - } - - // Wake reader waiting on this specific file - file_shared.wake(); - - // Check if we need to rotate - let needs_rotation = file_shared.estimated_size > self.max_file_size_bytes; - - if needs_rotation { - // Finish the IPC writer - if let Some(mut writer) = file_shared.writer.take() { - writer.finish()?; - } - // Mark as finished so readers know not to wait for more data - file_shared.writer_finished = true; - // Wake reader waiting on this file (it's now finished) - file_shared.wake(); - // Don't put back current_write_file - let it rotate - } else { - // Release file lock - drop(file_shared); - // Put back the current file for further writing - let mut shared = self.shared.lock(); - shared.current_write_file = Some(current_file); - } - } - - Ok(()) - } -} - -impl Drop for SpillPoolWriter { - fn drop(&mut self) { - let mut shared = self.shared.lock(); - - shared.active_writer_count -= 1; - let is_last_writer = shared.active_writer_count == 0; - - if !is_last_writer { - // Other writer clones are still active; do not finalize or - // signal EOF to readers. - return; + let mut file_shared = write_file.lock(); + + // Append the batch + if let Some(ref mut writer) = file_shared.writer { + writer.append_batch(batch)?; + // make sure we flush the writer for readers + writer.flush()?; + file_shared.batches_written += 1; + file_shared.estimated_size += batch_size; } - // Finalize the current file when the last writer is dropped - if let Some(current_file) = shared.current_write_file.take() { - // Release shared lock before locking file - drop(shared); + // Wake reader waiting on this specific file + file_shared.wake(); - let mut file_shared = current_file.lock(); + let max_file_size_reached = file_shared.estimated_size > self.max_file_size_bytes; - // Finish the current writer if it exists + if max_file_size_reached { + // Finish the IPC writer if let Some(mut writer) = file_shared.writer.take() { - // Ignore errors on drop - we're in destructor - let _ = writer.finish(); + writer.finish()?; } - // Mark as finished so readers know not to wait for more data file_shared.writer_finished = true; - // Wake reader waiting on this file (it's now finished) file_shared.wake(); + // Don't place `write_file` back in the `open_write_files` queue so we don't + // try writing to it again + } else { + // Release file lock drop(file_shared); - shared = self.shared.lock(); + // Put back the current file for further writing + let mut shared = self.shared.lock(); + shared.open_write_files.push_back(write_file); } - // Mark writer as dropped and wake pool-level readers - shared.writer_dropped = true; - shared.wake(); + Ok(()) } } -/// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, single-consumer) -/// semantics. +/// Creates a paired writer and reader for a spill pool with SPSC (single-producer, +/// single-consumer) semantics and strict FIFO ordering. +/// +/// If you need a spill pool that supports several producers, use [`mpsc_channel`] instead. /// -/// This is the recommended way to create a spill pool. The writer is `Clone`, allowing -/// multiple producers to coordinate writes to the same pool. The reader can consume batches -/// in FIFO order. The reader can start reading immediately after a writer appends a batch -/// to the spill file, without waiting for the file to be sealed, while writers continue to +/// The reader can start reading immediately after the writer appends a batch +/// to the spill file, without waiting for the file to be sealed, while the writer continues to /// write more data. /// /// Internally this coordinates rotating spill files based on size limits, and @@ -320,18 +320,18 @@ impl Drop for SpillPoolWriter { /// │ Writer Side Shared State Reader Side │ /// │ ─────────── ──────────── ─────────── │ /// │ │ -/// │ SpillPoolWriter ┌────────────────────┐ SpillPoolReader │ +/// │ SpillPoolSink ┌────────────────────┐ RecordBatchStream │ /// │ │ │ VecDeque │ │ │ /// │ │ │ ┌────┐┌────┐ │ │ │ /// │ push_batch() │ │ F1 ││ F2 │ ... │ next().await │ /// │ │ │ └────┘└────┘ │ │ │ -/// │ ▼ │ (FIFO order) │ ▼ │ +/// │ ▼ │ │ ▼ │ /// │ ┌─────────┐ │ │ ┌──────────┐ │ /// │ │Current │───────▶│ Coordination: │◀───│ Current │ │ /// │ │Write │ │ - Wakers │ │ Read │ │ /// │ │File │ │ - Batch counts │ │ File │ │ /// │ └─────────┘ │ - Writer status │ └──────────┘ │ -/// │ │ └────────────────────┘ │ │ +/// │ │ └────────────────────┘ │ │ /// │ │ │ │ /// │ Size > limit? Read all batches? │ /// │ │ │ │ @@ -339,7 +339,7 @@ impl Drop for SpillPoolWriter { /// │ Rotate to new file Pop from queue │ /// └─────────────────────────────────────────────────────────────────────────┘ /// -/// Writer produces → Shared FIFO queue → Reader consumes +/// Writer produces → Shared queue → Reader consumes /// ``` /// /// # File State Machine @@ -382,7 +382,7 @@ impl Drop for SpillPoolWriter { /// /// # Returns /// -/// A tuple of `(SpillPoolWriter, SendableRecordBatchStream)` that share the same +/// A tuple of `(SpillPoolSink, SendableRecordBatchStream)` that share the same /// underlying pool. The reader is returned as a stream for immediate use with /// async stream combinators. /// @@ -409,7 +409,7 @@ impl Drop for SpillPoolWriter { /// # let spill_manager = Arc::new(SpillManager::new(env, metrics, schema.clone())); /// # /// // Create channel with 1MB file size limit -/// let (writer, mut reader) = spill_pool::channel(1024 * 1024, spill_manager); +/// let (writer, mut reader) = spill_pool::spsc_channel(1024 * 1024, spill_manager); /// /// // Spawn writer and reader concurrently; writer wakes reader via wakers /// let writer_task = tokio::spawn(async move { @@ -458,14 +458,14 @@ impl Drop for SpillPoolWriter { /// If instead we use file rotation, and as long as the readers can keep up with the writer, /// then we can ensure that once a file is fully read by all readers it can be deleted, /// thus bounding the maximum disk usage to roughly `max_file_size_bytes`. -pub fn channel( +pub fn spsc_channel( max_file_size_bytes: usize, spill_manager: Arc, -) -> (SpillPoolWriter, SendableRecordBatchStream) { +) -> (SpillPoolSink, SendableRecordBatchStream) { let schema = Arc::clone(spill_manager.schema()); let shared = Arc::new(Mutex::new(SpillPoolShared::new(spill_manager))); - let writer = SpillPoolWriter { + let writer = SpillPoolSink { max_file_size_bytes, shared: Arc::clone(&shared), }; @@ -475,6 +475,51 @@ pub fn channel( (writer, Box::pin(reader)) } +/// Alias for [`mpsc_channel`]. +#[deprecated(note = "Use mpsc_channel instead")] +pub fn channel( + max_file_size_bytes: usize, + spill_manager: Arc, +) -> (SpillPoolWriter, SendableRecordBatchStream) { + mpsc_channel(max_file_size_bytes, spill_manager) +} + +/// Creates a paired writer and reader for a spill pool with MPSC (multi-producer, +/// single-consumer) semantics. See [`spsc_channel`] for the general architecture description +/// of the spill pool. +/// +/// Additional writers can be created by cloning the returned [`SpillPoolWriter`]. +/// +/// In contrast to [`spsc_channel`], this implementation provides no guarantees regarding +/// the read order of the returned [`SendableRecordBatchStream`]. +/// +/// If you need strict end-to-end FIFO (a single writer whose batches are read back in exact +/// write order), use [`spsc_channel`] instead. +/// +/// # File Management +/// +/// The shared channel uses the same size-based rotation trigger as the [single producer channel](spsc_channel). +/// All writers share the same pool of write files and coordinate file rotation. The number of open +/// files is kept as small as possible. When more writes occur concurrently than there are open write +/// files an additional file will be opened to write to. This prevents multiple writers from blocking +/// each other. +/// +/// When the last writer clone is dropped, it finalizes any remaining open write files so that all +/// written data can be accessed by the reader. +/// +/// # Returns +/// +/// A tuple of `(SpillPoolWriter, SendableRecordBatchStream)` that share the same +/// underlying pool. The reader is returned as a stream for immediate use with +/// async stream combinators. The writer can be cloned to create additional writers. +pub fn mpsc_channel( + max_file_size_bytes: usize, + spill_manager: Arc, +) -> (SpillPoolWriter, SendableRecordBatchStream) { + let (inner, reader) = spsc_channel(max_file_size_bytes, spill_manager); + (SpillPoolWriter { inner }, reader) +} + /// Shared state between writer and readers for an active spill file. /// Protected by a Mutex to coordinate between concurrent readers and the writer. struct ActiveSpillFileShared { @@ -482,7 +527,7 @@ struct ActiveSpillFileShared { writer: Option, /// The spill file, set when the writer finishes. /// Taken by the reader when creating a stream (the file stays open via file handles). - file: Option, + file: Option>, /// Total number of batches written to this file batches_written: usize, /// Estimated size in bytes of data written to this file @@ -507,25 +552,25 @@ impl ActiveSpillFileShared { } } -/// Reader state for a SpillFile (owned by individual SpillFile instances). +/// Reader state for a SpillPoolFile (owned by individual SpillPoolFile instances). /// This is kept separate from the shared state to avoid holding locks during I/O. -struct SpillFileReader { +struct SpillPoolFileReader { /// The actual stream reading from disk stream: SendableRecordBatchStream, /// Number of batches this reader has consumed batches_read: usize, } -struct SpillFile { +struct SpillPoolFile { /// Shared coordination state (contains writer and batch counts) shared: Arc>, - /// Reader state (lazy-initialized, owned by this SpillFile) - reader: Option, + /// Reader state (lazy-initialized, owned by this SpillPoolFile) + reader: Option, /// Spill manager for creating readers spill_manager: Arc, } -impl Stream for SpillFile { +impl Stream for SpillPoolFile { type Item = Result; fn poll_next( @@ -568,7 +613,7 @@ impl Stream for SpillFile { .read_spill_as_stream_unbuffered(file, None) { Ok(stream) => { - self.reader = Some(SpillFileReader { + self.reader = Some(SpillPoolFileReader { stream, batches_read: 0, }); @@ -607,9 +652,9 @@ impl Stream for SpillFile { } } -/// A stream that reads from a SpillPool in FIFO order. +/// A stream that reads from a SpillPool. The reader guarantees FIFO order if a single writer is used. /// -/// Created by [`channel`]. See that function for architecture diagrams and usage examples. +/// Created by [`spsc_channel`]. See that function for architecture diagrams and usage examples. /// /// The stream automatically handles file rotation and reads from completed files. /// When no data is available, it returns `Poll::Pending` and registers a waker to @@ -627,8 +672,8 @@ impl Stream for SpillFile { pub struct SpillPoolReader { /// Shared reference to the spill pool shared: Arc>, - /// Current SpillFile we're reading from - current_file: Option, + /// Current SpillPoolFile we're reading from + current_file: Option, /// Schema of the spilled data schema: SchemaRef, } @@ -636,7 +681,7 @@ pub struct SpillPoolReader { impl SpillPoolReader { /// Creates a new reader from shared pool state. /// - /// This is private - use the `channel()` function to create a reader/writer pair. + /// This is private - use the [`spsc_channel`] function to create a reader/writer pair. /// /// # Arguments /// @@ -706,12 +751,12 @@ impl Stream for SpillPoolReader { // Peek at the front of the queue (don't pop yet) if let Some(file_shared) = shared.files.front() { - // Create a SpillFile from the shared state + // Create a SpillPoolFile from the shared state let spill_manager = Arc::clone(&shared.spill_manager); let file_shared = Arc::clone(file_shared); - drop(shared); // Release lock before creating SpillFile + drop(shared); // Release lock before creating SpillPoolFile - self.current_file = Some(SpillFile { + self.current_file = Some(SpillPoolFile { shared: file_shared, reader: None, spill_manager, @@ -722,7 +767,7 @@ impl Stream for SpillPoolReader { } // No files in queue - check if writer is done - if shared.writer_dropped { + if shared.remaining_writer_count == 0 { // Writer is done and no more files will be added - EOF return Poll::Ready(None); } @@ -746,7 +791,7 @@ mod tests { use crate::metrics::{ExecutionPlanMetricsSet, SpillMetrics}; use arrow::array::{ArrayRef, Int32Array}; use arrow::datatypes::{DataType, Field, Schema}; - use datafusion_common_runtime::SpawnedTask; + use datafusion_common_runtime::{JoinSet, SpawnedTask}; use datafusion_execution::runtime_env::RuntimeEnv; fn create_test_schema() -> SchemaRef { @@ -763,24 +808,35 @@ mod tests { fn create_spill_channel( max_file_size: usize, + ) -> (SpillPoolSink, SendableRecordBatchStream) { + let env = Arc::new(RuntimeEnv::default()); + let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); + let schema = create_test_schema(); + let spill_manager = Arc::new(SpillManager::new(env, metrics, schema)); + + spsc_channel(max_file_size, spill_manager) + } + + fn create_shared_spill_channel( + max_file_size: usize, ) -> (SpillPoolWriter, SendableRecordBatchStream) { let env = Arc::new(RuntimeEnv::default()); let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); let schema = create_test_schema(); let spill_manager = Arc::new(SpillManager::new(env, metrics, schema)); - channel(max_file_size, spill_manager) + mpsc_channel(max_file_size, spill_manager) } fn create_spill_channel_with_metrics( max_file_size: usize, - ) -> (SpillPoolWriter, SendableRecordBatchStream, SpillMetrics) { + ) -> (SpillPoolSink, SendableRecordBatchStream, SpillMetrics) { let env = Arc::new(RuntimeEnv::default()); let metrics = SpillMetrics::new(&ExecutionPlanMetricsSet::new(), 0); let schema = create_test_schema(); let spill_manager = Arc::new(SpillManager::new(env, metrics.clone(), schema)); - let (writer, reader) = channel(max_file_size, spill_manager); + let (writer, reader) = spsc_channel(max_file_size, spill_manager); (writer, reader, metrics) } @@ -1204,6 +1260,57 @@ mod tests { Ok(()) } + #[tokio::test(flavor = "multi_thread", worker_threads = 10)] + async fn test_concurrent_writers() -> Result<()> { + let (writer, mut reader) = create_shared_spill_channel(1024 * 1024); + + // Spawn writer tasks + let mut writer_join_set = JoinSet::new(); + for w in 0..10 { + let writer = writer.clone(); + writer_join_set.spawn(async move { + for b in 0..10 { + let batch = create_test_batch((w * 100) + (b * 10), 10); + writer.push_batch(&batch).unwrap(); + } + }); + } + drop(writer); + + // Reader task (runs concurrently) + let reader_handle = SpawnedTask::spawn(async move { + let mut batch_order = vec![]; + loop { + match reader.next().await { + None => break, + Some(batch) => { + let batch = batch.unwrap(); + + assert_eq!(batch.num_rows(), 10); + + let col = batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + batch_order.push(col.value(0) / 10); + } + } + } + batch_order + }); + + // Wait for both to complete + writer_join_set.join_all().await; + let mut batch_order = reader_handle.await.unwrap(); + + // When used with multiple writers, order is not guaranteed + batch_order.sort(); + assert_eq!(batch_order, (0i32..100i32).collect::>()); + + Ok(()) + } + #[tokio::test] async fn test_reader_catches_up_to_writer() -> Result<()> { let (writer, mut reader) = create_spill_channel(1024 * 1024); @@ -1322,7 +1429,7 @@ mod tests { let spill_manager = Arc::new(SpillManager::new(Arc::clone(&env), metrics.clone(), schema)); - let (writer, mut reader) = channel(1024 * 1024, spill_manager); + let (writer, mut reader) = spsc_channel(1024 * 1024, spill_manager); // Write some batches for i in 0..5 { @@ -1384,7 +1491,7 @@ mod tests { /// 5. EOF is only signalled after writer2 is also dropped. #[tokio::test] async fn test_clone_drop_does_not_signal_eof_prematurely() -> Result<()> { - let (writer1, mut reader) = create_spill_channel(1024 * 1024); + let (writer1, mut reader) = create_shared_spill_channel(1024 * 1024); let writer2 = writer1.clone(); // Synchronization: tell writer2 when it may proceed. @@ -1463,7 +1570,7 @@ mod tests { let schema = create_test_schema(); let spill_manager = Arc::new(SpillManager::new(runtime, metrics.clone(), schema)); - let (writer, mut reader) = channel(batch_size, spill_manager); + let (writer, mut reader) = spsc_channel(batch_size - 1, spill_manager); // Step 3: Write NUM_BATCHES batches to create approximately NUM_BATCHES files for i in 0..NUM_BATCHES { @@ -1474,10 +1581,8 @@ mod tests { // Check how many files were created (should be at least a few due to file rotation) let file_count = metrics.spill_file_count.value(); assert_eq!( - file_count, - NUM_BATCHES - 1, - "Expected at {} files with rotation, got {file_count}", - NUM_BATCHES - 1 + file_count, NUM_BATCHES, + "Expected at {NUM_BATCHES} files with rotation, got {file_count}" ); // Step 4: Verify initial disk usage reflects all files diff --git a/datafusion/physical-plan/src/statistics.rs b/datafusion/physical-plan/src/statistics.rs new file mode 100644 index 0000000000000..9246d7d9f5a9c --- /dev/null +++ b/datafusion/physical-plan/src/statistics.rs @@ -0,0 +1,274 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Statistics computation for physical plans. +//! +//! [`StatisticsArgs`] provides external context to +//! [`ExecutionPlan::statistics_from_inputs`]. + +use crate::ExecutionPlan; +use datafusion_common::{ + Result, Statistics, assert_eq_or_internal_err, assert_or_internal_err, +}; +use std::cell::RefCell; +use std::collections::HashMap; +use std::rc::Rc; +use std::sync::Arc; + +/// Per-call memoization cache for statistics computation. +/// +/// Keyed by `(plan node pointer address, partition)`. Shared across +/// a single statistics walk via [`StatisticsContext`]. +/// +/// The pointer-based key is safe within a single synchronous walk: +/// all `Arc` nodes are held by the plan tree for +/// the duration of the walk, so addresses cannot be reused. +#[derive(Debug, Default)] +struct StatsCache(HashMap<(usize, Option), Arc>); + +impl StatsCache { + fn get( + &self, + plan: &dyn ExecutionPlan, + partition: Option, + ) -> Option<&Arc> { + let key = ( + plan as *const dyn ExecutionPlan as *const () as usize, + partition, + ); + self.0.get(&key) + } + + fn insert( + &mut self, + plan: &dyn ExecutionPlan, + partition: Option, + stats: Arc, + ) { + let key = ( + plan as *const dyn ExecutionPlan as *const () as usize, + partition, + ); + self.0.insert(key, stats); + } +} + +/// Arguments passed to [`ExecutionPlan::statistics_from_inputs`] carrying +/// external information that operators can use when computing their +/// statistics. +#[derive(Debug, Default, Clone)] +pub struct StatisticsArgs { + partition: Option, +} + +impl StatisticsArgs { + /// Creates new statistics arguments. + /// + /// By default the partition is set to `None` (statistics should be computed + /// for the entire plan). + pub fn new() -> Self { + Default::default() + } + + /// Set the partition to compute statistics + /// + /// * `None` means statistics should be computed for the entire plan. + /// * `Some(idx)` means statistics should be computed for the specified + /// partition index. + pub fn set_partition(&mut self, partition: Option) { + self.partition = partition; + } + + /// Builder Style API for [`Self::set_partition`] + pub fn with_partition(mut self, partition: Option) -> Self { + self.set_partition(partition); + self + } + + /// Return the partition to compute statistics + pub fn partition(&self) -> Option { + self.partition + } +} + +/// Directive returned by [`ExecutionPlan::child_stats_requests`] describing +/// how the [`StatisticsContext`] should obtain each child's statistics. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ChildStats { + /// Compute the child's statistics at this partition (`None` = overall). + At(Option), + /// Skip this child; the parent does not need its statistics. A placeholder + /// [`Statistics::new_unknown`] is supplied in its slot. + Skip, +} + +/// Owns the bottom-up traversal and per-walk memoization cache for statistics +/// computation. Call [`StatisticsContext::compute`] to walk a plan tree. +pub struct StatisticsContext { + cache: Rc>, +} + +impl Default for StatisticsContext { + fn default() -> Self { + Self::new() + } +} + +impl StatisticsContext { + /// Creates a context with an empty cache. + pub fn new() -> Self { + Self { + cache: Rc::new(RefCell::new(StatsCache::default())), + } + } + + /// Clears the memoization cache. + /// + /// The cache is keyed by raw plan-node pointers, which are only stable + /// while the current plan tree is alive. Reset between optimizer passes + /// (which rewrite the plan) when reusing one context across them, so stale + /// pointer keys cannot collide. + pub fn reset_cache(&self) { + self.cache.borrow_mut().0.clear(); + } + + /// Computes statistics for `plan`, resolving children first and passing + /// the results to [`ExecutionPlan::statistics_from_inputs`]. + /// + /// When `args.partition()` is `Some(idx)`, `idx` is validated against the + /// plan's partition count. + pub fn compute( + &self, + plan: &dyn ExecutionPlan, + args: &StatisticsArgs, + ) -> Result> { + let partition = args.partition(); + + if let Some(idx) = partition { + let partition_count = plan.properties().partitioning.partition_count(); + assert_or_internal_err!( + idx < partition_count, + "Invalid partition index: {}, the partition count is {}", + idx, + partition_count + ); + } + + if let Some(cached) = self.cache.borrow().get(plan, partition) { + return Ok(Arc::clone(cached)); + } + + let children = plan.children(); + let requests = plan.child_stats_requests(partition); + assert_eq_or_internal_err!( + requests.len(), + children.len(), + "{} child_stats_requests returned {} entries for {} children", + plan.name(), + requests.len(), + children.len() + ); + let child_stats = children + .iter() + .zip(requests) + .map(|(child, directive)| match directive { + ChildStats::At(p) => { + self.compute(child.as_ref(), &StatisticsArgs::new().with_partition(p)) + } + ChildStats::Skip => { + Ok(Arc::new(Statistics::new_unknown(child.schema().as_ref()))) + } + }) + .collect::>>()?; + + let result = plan.statistics_from_inputs(&child_stats, args)?; + self.cache + .borrow_mut() + .insert(plan, partition, Arc::clone(&result)); + Ok(result) + } +} + +#[cfg(all(test, feature = "test_utils"))] +mod tests { + use super::*; + use crate::coalesce_partitions::CoalescePartitionsExec; + use crate::test::exec::StatisticsExec; + use arrow::datatypes::{DataType, Field, Schema}; + use datafusion_common::{ColumnStatistics, stats::Precision}; + + fn make_stats_leaf(num_rows: usize) -> Arc { + let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]); + let col_stats = vec![ColumnStatistics { + null_count: Precision::Exact(0), + max_value: Precision::Absent, + min_value: Precision::Absent, + sum_value: Precision::Absent, + distinct_count: Precision::Absent, + byte_size: Precision::Absent, + }]; + Arc::new(StatisticsExec::new( + Statistics { + num_rows: Precision::Exact(num_rows), + total_byte_size: Precision::Absent, + column_statistics: col_stats, + }, + schema, + )) + } + + #[test] + fn coalesce_returns_overall_stats_for_any_partition() { + let leaf = make_stats_leaf(100); + let plan: Arc = Arc::new(CoalescePartitionsExec::new(leaf)); + + let ctx = StatisticsContext::new(); + let stats = ctx + .compute( + plan.as_ref(), + &StatisticsArgs::new().with_partition(Some(0)), + ) + .unwrap(); + assert_eq!(stats.num_rows, Precision::Exact(100)); + + let stats_none = ctx.compute(plan.as_ref(), &StatisticsArgs::new()).unwrap(); + assert_eq!(stats_none.num_rows, Precision::Exact(100)); + } + + #[test] + fn context_caches_within_walk() { + let leaf = make_stats_leaf(42); + let ctx = StatisticsContext::new(); + let args = StatisticsArgs::new(); + + let s1 = ctx.compute(leaf.as_ref(), &args).unwrap(); + assert!(!ctx.cache.borrow().0.is_empty()); + + let s2 = ctx.compute(leaf.as_ref(), &args).unwrap(); + assert!(Arc::ptr_eq(&s1, &s2)); + } + + #[test] + fn reset_cache_clears_entries() { + let leaf = make_stats_leaf(10); + let ctx = StatisticsContext::new(); + let _ = ctx.compute(leaf.as_ref(), &StatisticsArgs::new()).unwrap(); + assert!(!ctx.cache.borrow().0.is_empty()); + ctx.reset_cache(); + assert!(ctx.cache.borrow().0.is_empty()); + } +} diff --git a/datafusion/physical-plan/src/streaming.rs b/datafusion/physical-plan/src/streaming.rs index 250eb59f19b87..7b0058e79887c 100644 --- a/datafusion/physical-plan/src/streaming.rs +++ b/datafusion/physical-plan/src/streaming.rs @@ -30,13 +30,17 @@ use crate::projection::{ ProjectionExec, all_alias_free_columns, new_projections_for_columns, update_ordering, }; use crate::stream::RecordBatchStreamAdapter; -use crate::{ExecutionPlan, Partitioning, SendableRecordBatchStream}; +use crate::{ + ChildrenPropertiesMode, ExecutionPlan, Partitioning, ReplaceChildrenOptions, + SendableRecordBatchStream, +}; use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{Result, internal_err, plan_err}; use datafusion_execution::TaskContext; use datafusion_physical_expr::PhysicalExpr; +use datafusion_physical_expr::projection::ProjectionMapping; use datafusion_physical_expr::{EquivalenceProperties, LexOrdering}; use async_trait::async_trait; @@ -102,7 +106,7 @@ impl StreamingTableExec { let cache = Self::compute_properties( Arc::clone(&projected_schema), projected_output_ordering.clone(), - &partitions, + Partitioning::UnknownPartitioning(partitions.len()), infinite, ); Ok(Self { @@ -117,6 +121,25 @@ impl StreamingTableExec { }) } + /// Declares the output partitioning of this stream. + /// + /// `output_partitioning` must describe this plan's current output and have + /// the same number of partitions as the stream. + pub fn with_output_partitioning( + mut self, + output_partitioning: Partitioning, + ) -> Result { + if output_partitioning.partition_count() != self.partitions.len() { + return plan_err!( + "Output partitioning has {} partitions but stream has {} partitions", + output_partitioning.partition_count(), + self.partitions.len() + ); + } + Arc::make_mut(&mut self.cache).partitioning = output_partitioning; + Ok(self) + } + pub fn partitions(&self) -> &Vec> { &self.partitions } @@ -149,14 +172,12 @@ impl StreamingTableExec { fn compute_properties( schema: SchemaRef, orderings: Vec, - partitions: &[Arc], + output_partitioning: Partitioning, infinite: bool, ) -> PlanProperties { // Calculate equivalence properties: let eq_properties = EquivalenceProperties::new_with_orderings(schema, orderings); - // Get output partitioning: - let output_partitioning = Partitioning::UnknownPartitioning(partitions.len()); let boundedness = if infinite { Boundedness::Unbounded { requires_infinite_memory: false, @@ -206,6 +227,16 @@ impl DisplayAs for StreamingTableExec { if let Some(fetch) = self.limit { write!(f, ", fetch={fetch}")?; } + if !matches!( + self.cache.output_partitioning(), + Partitioning::UnknownPartitioning(_) + ) { + write!( + f, + ", output_partitioning={}", + self.cache.output_partitioning() + )?; + } display_orderings(f, &self.projected_output_ordering)?; @@ -247,14 +278,15 @@ impl ExecutionPlan for StreamingTableExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { if children.is_empty() { Ok(self) @@ -263,6 +295,16 @@ impl ExecutionPlan for StreamingTableExec { } } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, @@ -315,6 +357,17 @@ impl ExecutionPlan for StreamingTableExec { }; lex_orderings.push(ordering); } + let projection_mapping = ProjectionMapping::try_new( + projection + .expr() + .iter() + .map(|expr| (Arc::clone(&expr.expr), expr.alias.clone())), + &self.schema(), + )?; + let output_partitioning = self + .cache + .output_partitioning() + .project(&projection_mapping, self.cache.equivalence_properties()); StreamingTableExec::try_new( Arc::clone(self.partition_schema()), @@ -324,6 +377,7 @@ impl ExecutionPlan for StreamingTableExec { self.is_infinite(), self.limit(), ) + .and_then(|exec| exec.with_output_partitioning(output_partitioning)) .map(|e| Some(Arc::new(e) as _)) } diff --git a/datafusion/physical-plan/src/test.rs b/datafusion/physical-plan/src/test.rs index 4c4724e4dcc4f..b38a46d160755 100644 --- a/datafusion/physical-plan/src/test.rs +++ b/datafusion/physical-plan/src/test.rs @@ -24,13 +24,14 @@ use std::pin::Pin; use std::sync::Arc; use std::task::Context; -use crate::ExecutionPlan; use crate::common; use crate::execution_plan::{Boundedness, EmissionType}; use crate::memory::MemoryStream; use crate::metrics::MetricsSet; +use crate::statistics::StatisticsArgs; use crate::stream::RecordBatchStreamAdapter; use crate::streaming::PartitionStream; +use crate::{ChildrenPropertiesMode, ExecutionPlan, ReplaceChildrenOptions}; use crate::{DisplayAs, DisplayFormatType, PlanProperties}; use arrow::array::{Array, ArrayRef, Int32Array, RecordBatch}; @@ -142,25 +143,29 @@ impl ExecutionPlan for TestMemoryExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - // Apply to all sort information orderings - let mut tnr = TreeNodeRecursion::Continue; - for ordering in &self.sort_information { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) + Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn repartitioned( &self, _target_partitions: usize, @@ -181,8 +186,12 @@ impl ExecutionPlan for TestMemoryExec { unimplemented!() } - fn partition_statistics(&self, partition: Option) -> Result> { - if partition.is_some() { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + if args.partition().is_some() { Ok(Arc::new(Statistics::new_unknown(&self.schema))) } else { Ok(Arc::new(self.statistics_inner()?)) diff --git a/datafusion/physical-plan/src/test/exec.rs b/datafusion/physical-plan/src/test/exec.rs index 200223b9b660a..1e2005e908fbf 100644 --- a/datafusion/physical-plan/src/test/exec.rs +++ b/datafusion/physical-plan/src/test/exec.rs @@ -17,10 +17,11 @@ //! Simple iterator over batches for use in testing +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions}; use crate::{ DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, common, - execution_plan::Boundedness, + execution_plan::Boundedness, statistics::StatisticsArgs, }; use crate::{ execution_plan::EmissionType, @@ -125,6 +126,9 @@ pub struct MockExec { /// if true (the default), sends data using a separate task to ensure the /// batches are not available without this stream yielding first use_task: bool, + /// if true, report unknown statistics instead of deriving them from + /// `data` (which propagates any planted errors at planning time) + unknown_statistics: bool, cache: Arc, } @@ -142,6 +146,7 @@ impl MockExec { data, schema, use_task: true, + unknown_statistics: false, cache: Arc::new(cache), } } @@ -154,6 +159,17 @@ impl MockExec { self } + /// Report unknown statistics rather than computing them from `data`. + /// + /// By default statistics are derived from `data`, which propagates any + /// planted errors when statistics are requested during planning (for + /// example when a parent node computes its properties). Use this when a + /// planted error should only surface at execution time. + pub fn with_unknown_statistics(mut self) -> Self { + self.unknown_statistics = true; + self + } + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. fn compute_properties(schema: SchemaRef) -> PlanProperties { PlanProperties::new( @@ -198,18 +214,29 @@ impl ExecutionPlan for MockExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Returns a stream which yields data fn execute( &self, @@ -256,9 +283,14 @@ impl ExecutionPlan for MockExec { } } - // Panics if one of the batches is an error - fn partition_statistics(&self, partition: Option) -> Result> { - if partition.is_some() { + // Errors if one of the batches is an error, unless + // `with_unknown_statistics` was used + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + if self.unknown_statistics || args.partition().is_some() { return Ok(Arc::new(Statistics::new_unknown(&self.schema))); } let data: Result> = self @@ -429,20 +461,31 @@ impl ExecutionPlan for BarrierExec { unimplemented!() } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Returns a stream which yields data fn execute( &self, @@ -489,8 +532,12 @@ impl ExecutionPlan for BarrierExec { Ok(builder.build()) } - fn partition_statistics(&self, partition: Option) -> Result> { - if partition.is_some() { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + if args.partition().is_some() { return Ok(Arc::new(Statistics::new_unknown(&self.schema))); } Ok(Arc::new(common::compute_record_batch_statistics( @@ -568,20 +615,31 @@ impl ExecutionPlan for ErrorExec { unimplemented!() } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { unimplemented!() } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Returns a stream which yields data fn execute( &self, @@ -663,18 +721,29 @@ impl ExecutionPlan for StatisticsExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -683,8 +752,12 @@ impl ExecutionPlan for StatisticsExec { unimplemented!("This plan only serves for testing statistics") } - fn partition_statistics(&self, partition: Option) -> Result> { - Ok(Arc::new(if partition.is_some() { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + Ok(Arc::new(if args.partition().is_some() { Statistics::new_unknown(&self.schema) } else { self.stats.clone() @@ -768,20 +841,31 @@ impl ExecutionPlan for BlockingExec { vec![] } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { internal_err!("Children cannot be replaced in {self:?}") } fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, _partition: usize, @@ -919,18 +1003,29 @@ impl ExecutionPlan for PanicExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { internal_err!("Children cannot be replaced in {:?}", self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + fn execute( &self, partition: usize, diff --git a/datafusion/physical-plan/src/topk/mod.rs b/datafusion/physical-plan/src/topk/mod.rs index 9da606dc90db2..1e3efff36b1d8 100644 --- a/datafusion/physical-plan/src/topk/mod.rs +++ b/datafusion/physical-plan/src/topk/mod.rs @@ -19,11 +19,15 @@ use arrow::{ array::{Array, AsArray}, - compute::{FilterBuilder, interleave_record_batch, prep_null_mask_filter}, - row::{RowConverter, Rows, SortField}, + compute::{ + BatchCoalescer, FilterBuilder, interleave_record_batch, prep_null_mask_filter, + take_record_batch, + }, + row::{OwnedRow, RowConverter, Rows, SortField}, }; use datafusion_expr::{ColumnarValue, Operator}; use std::mem::size_of; +use std::sync::atomic::{AtomicUsize, Ordering as AtomicOrdering}; use std::{cmp::Ordering, collections::BinaryHeap, sync::Arc}; use super::metrics::{ @@ -33,7 +37,7 @@ use super::metrics::{ use crate::spill::get_record_batch_memory_size; use crate::{SendableRecordBatchStream, stream::RecordBatchStreamAdapter}; -use arrow::array::{ArrayRef, RecordBatch}; +use arrow::array::{ArrayRef, RecordBatch, UInt32Array}; use arrow::datatypes::SchemaRef; use datafusion_common::{ HashMap, Result, ScalarValue, internal_datafusion_err, internal_err, @@ -49,7 +53,7 @@ use datafusion_physical_expr::{ use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; use parking_lot::RwLock; -/// Global TopK +/// TopK /// /// # Background /// @@ -84,10 +88,13 @@ use parking_lot::RwLock; /// # Partial Sort Optimization /// /// This implementation additionally optimizes queries where the input is already -/// partially sorted by a common prefix of the requested ordering. Once the top K -/// heap is full, if subsequent rows are guaranteed to be strictly greater (in sort -/// order) on this prefix than the largest row currently stored, the operator -/// safely terminates early. +/// partially sorted by a common prefix of the requested ordering. If subsequent +/// rows are guaranteed to be strictly greater (in sort order) than a known TopK +/// boundary on this prefix, the operator safely terminates early. +/// +/// For a local TopK, that boundary comes from the local heap once it has K rows. +/// For a partitioned `SortExec`, a shared dynamic-filter threshold can provide +/// the same prefix boundary before a lagging partition has filled its local heap. /// /// ## Example /// @@ -135,34 +142,189 @@ pub struct TopK { /// For more background, please also see the [Dynamic Filters: Passing Information Between Operators During Execution for 25x Faster Queries blog] /// /// [Dynamic Filters: Passing Information Between Operators During Execution for 25x Faster Queries blog]: https://datafusion.apache.org/blog/2025/09/10/dynamic-filters -#[derive(Debug, Clone)] +#[derive(Debug)] pub struct TopKDynamicFilters { - /// The current *global* threshold for the dynamic filter. - /// This is shared across all partitions and is updated by any of them. - /// Stored as row bytes for efficient comparison. - threshold_row: Option>, + /// The current threshold shared by all TopK emitters that use this dynamic + /// filter. Any emitter may tighten it. + /// + /// The full sort-key row and common-prefix row are stored together so they + /// always describe the same heap row. + shared_threshold: Option, /// The expression used to evaluate the dynamic filter /// Only updated when lock held for the duration of the update expr: Arc, + /// Number of local TopK emitters that have not called `emit` yet. + /// + /// A partition-preserving `SortExec` creates one local TopK per output + /// partition. The shared dynamic filter is complete only after every local + /// TopK has emitted. + /// + /// `emit` only needs a read guard on the shared filter wrapper, so + /// concurrent emitters use this atomic counter instead of taking an + /// exclusive lock just to mark their partition done. + remaining_topk_emitters: AtomicUsize, +} + +#[derive(Debug, Clone)] +struct TopKThreshold { + /// The full sort-key row bytes for efficient comparison. + full_sort_key_row: Vec, + /// The same heap row encoded with the common-prefix converter, when the + /// input ordering shares a prefix with the TopK ordering. + /// + /// This lets each partition stop from a shared TopK threshold even if its + /// local heap has not filled yet. + common_prefix_row: Option>, +} + +impl TopKThreshold { + fn new(full_sort_key_row: Vec, common_prefix_row: Option>) -> Self { + Self { + full_sort_key_row, + common_prefix_row, + } + } + + fn full_sort_key_row(&self) -> &[u8] { + self.full_sort_key_row.as_slice() + } + + fn common_prefix_row(&self) -> Option<&[u8]> { + self.common_prefix_row.as_deref() + } + + fn is_more_selective_than(&self, current: &Self) -> bool { + self.full_sort_key_row() < current.full_sort_key_row() + } +} + +#[derive(Clone, Copy)] +struct TopKHeapBoundaryRow<'a> { + row: &'a TopKRow, +} + +impl<'a> TopKHeapBoundaryRow<'a> { + fn new(row: &'a TopKRow) -> Self { + Self { row } + } + + fn full_sort_key_row(&self) -> &[u8] { + self.row.row() + } + + fn is_more_selective_than(&self, current: Option<&TopKThreshold>) -> bool { + current + .map(|current| self.full_sort_key_row() < current.full_sort_key_row()) + .unwrap_or(true) + } +} + +#[derive(Clone, Copy)] +struct TopKHeapBoundary<'a> { + row: &'a TopKRow, + batch: &'a RecordBatch, +} + +impl<'a> TopKHeapBoundary<'a> { + fn new(row: &'a TopKRow, batch: &'a RecordBatch) -> Self { + Self { row, batch } + } + + fn threshold_values( + &self, + sort_exprs: &[PhysicalSortExpr], + ) -> Result> { + let mut scalar_values = Vec::with_capacity(sort_exprs.len()); + for sort_expr in sort_exprs { + let value = sort_expr + .expr + .evaluate(&self.batch.slice(self.row.index, 1))?; + + let scalar = match value { + ColumnarValue::Scalar(scalar) => scalar, + ColumnarValue::Array(array) if array.len() == 1 => { + ScalarValue::try_from_array(&array, 0)? + } + array => { + return internal_err!("Expected a scalar value, got {:?}", array); + } + }; + scalar_values.push(scalar); + } + + Ok(scalar_values) + } + + fn threshold(&self, common_prefix_row: Option>) -> TopKThreshold { + TopKThreshold::new(self.row.row().to_vec(), common_prefix_row) + } } impl TopKDynamicFilters { /// Create a new `TopKDynamicFilters` with the given expression pub fn new(expr: Arc) -> Self { + Self::new_with_topk_emitter_count(expr, 1) + } + + /// Create a new `TopKDynamicFilters` with the expected number of local + /// TopK emitters that share it. + pub fn new_with_topk_emitter_count( + expr: Arc, + topk_emitter_count: usize, + ) -> Self { + debug_assert!(topk_emitter_count > 0); Self { - threshold_row: None, + shared_threshold: None, expr, + remaining_topk_emitters: AtomicUsize::new(topk_emitter_count), } } pub fn expr(&self) -> Arc { Arc::clone(&self.expr) } + + fn mark_topk_emitted(&self) { + let previous = self + .remaining_topk_emitters + .fetch_update( + AtomicOrdering::AcqRel, + AtomicOrdering::Acquire, + |remaining| remaining.checked_sub(1), + ) + .unwrap_or(0); + debug_assert!( + previous > 0, + "TopK dynamic filter emitter completed more times than expected" + ); + + if previous == 1 { + self.expr.mark_complete(); + } + } } // Guesstimate for memory allocation: estimated number of bytes used per row in the RowConverter const ESTIMATED_BYTES_PER_ROW: usize = 20; +/// Owned data of a row that was just evicted from a [`TopKHeap`]. +/// +/// Returned by [`TopKHeap::add`] so that callers (e.g. rank-aware +/// wrappers that retain boundary ties) can decide whether to retain +/// the evicted row externally. The underlying batch is captured +/// before the heap's internal `RecordBatchStore` decrements the +/// batch's use count, so the data remains accessible even if the +/// heap drops its internal reference to the batch. +#[derive(Debug, Clone)] +pub(crate) struct EvictedRow { + /// The record batch the evicted row came from. + pub batch: RecordBatch, + /// Row index within `batch`. + pub index: usize, + /// Encoded ORDER BY tuple for the evicted row, in [`arrow::row`] format. + pub row_bytes: Vec, +} + pub(crate) fn build_sort_fields( ordering: &[PhysicalSortExpr], schema: &SchemaRef, @@ -206,7 +368,7 @@ impl TopK { let scratch_rows = row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); - let prefix_row_converter = if common_sort_prefix.is_empty() { + let common_prefix_row_converter = if common_sort_prefix.is_empty() { None } else { let input_sort_fields = build_sort_fields(&common_sort_prefix, &schema)?; @@ -222,7 +384,7 @@ impl TopK { row_converter, scratch_rows, heap: TopKHeap::new(k), - common_sort_prefix_converter: prefix_row_converter, + common_sort_prefix_converter: common_prefix_row_converter, common_sort_prefix: Arc::from(common_sort_prefix), finished: false, filter, @@ -255,7 +417,9 @@ impl TopK { let array = filtered.into_array(num_rows)?; let mut filter = array.as_boolean().clone(); if !filter.has_true() { - // nothing to filter, so no need to update + // The heap is unchanged, but a fully rejected batch can still prove + // that the shared sort prefix has passed the heap boundary. + self.attempt_early_completion(&batch)?; return Ok(()); } // only update the keys / rows if the filter does not match all rows @@ -312,6 +476,10 @@ impl TopK { // update the filter representation of our TopK heap self.update_filter()?; + } else { + // The heap did not change, but this batch's prefix may still prove + // that no later rows can enter the TopK. + self.attempt_early_completion(&batch)?; } Ok(()) @@ -339,6 +507,28 @@ impl TopK { replacements } + fn current_heap_boundary_row(&self) -> Option> { + self.heap.max().map(TopKHeapBoundaryRow::new) + } + + fn current_heap_boundary(&self) -> Result>> { + let Some(row) = self.heap.max() else { + return Ok(None); + }; + + self.heap_boundary(row).map(Some) + } + + fn heap_boundary<'a>(&'a self, row: &'a TopKRow) -> Result> { + let batch_entry = self + .heap + .store + .get(row.batch_id) + .ok_or_else(|| internal_datafusion_err!("Invalid batch ID in TopKRow"))?; + + Ok(TopKHeapBoundary::new(row, &batch_entry.batch)) + } + /// Update the filter representation of our TopK heap. /// For example, given the sort expression `ORDER BY a DESC, b ASC LIMIT 3`, /// and the current heap values `[(1, 5), (1, 4), (2, 3)]`, @@ -351,65 +541,46 @@ impl TopK { /// ``` fn update_filter(&mut self) -> Result<()> { // If the heap doesn't have k elements yet, we can't create thresholds - let Some(max_row) = self.heap.max() else { + let Some(boundary_row) = self.current_heap_boundary_row() else { return Ok(()); }; - let new_threshold_row = &max_row.row; - // Fast path: check if the current value in topk is better than what is // currently set in the filter with a read only lock - let needs_update = self - .filter - .read() - .threshold_row - .as_ref() - .map(|current_row| { - // new < current means new threshold is more selective - new_threshold_row < current_row - }) - .unwrap_or(true); // No current threshold, so we need to set one + let needs_update = { + let filter = self.filter.read(); + boundary_row.is_more_selective_than(filter.shared_threshold.as_ref()) + }; // exit early if the current values are better if !needs_update { return Ok(()); } + let boundary = self.heap_boundary(boundary_row.row)?; + // Extract scalar values BEFORE acquiring lock to reduce critical section - let thresholds = match self.heap.get_threshold_values(&self.expr)? { - Some(t) => t, - None => return Ok(()), - }; + let thresholds = boundary.threshold_values(&self.expr)?; // Build the filter expression OUTSIDE any synchronization let predicate = Self::build_filter_expression(&self.expr, &thresholds)?; - let new_threshold = new_threshold_row.to_vec(); + let new_threshold = + boundary.threshold(self.encode_topk_common_prefix_row(boundary)?); // update the threshold. Since there was a lock gap, we must check if it is still the best // may have changed while we were building the expression without the lock let mut filter = self.filter.write(); - let old_threshold = filter.threshold_row.take(); - - // Update filter if we successfully updated the threshold - // (or if there was no previous threshold and we're the first) - match old_threshold { - Some(old_threshold) => { - // new threshold is still better than the old one - if new_threshold.as_slice() < old_threshold.as_slice() { - filter.threshold_row = Some(new_threshold); - } else { - // some other thread updated the threshold to a better - // one while we were building so there is no need to - // update the filter - filter.threshold_row = Some(old_threshold); - return Ok(()); - } - } - None => { - // No previous threshold, so we can set the new one - filter.threshold_row = Some(new_threshold); - } - }; + let still_needs_update = filter + .shared_threshold + .as_ref() + .map(|current| new_threshold.is_more_selective_than(current)) + .unwrap_or(true); + if !still_needs_update { + // some other thread updated the threshold to a better one while we + // were building so there is no need to update the filter + return Ok(()); + } + filter.shared_threshold = Some(new_threshold); // Update the filter expression if let Some(pred) = predicate @@ -507,78 +678,110 @@ impl TopK { Ok(dynamic_predicate) } - /// If input ordering shares a common sort prefix with the TopK, and if the TopK's heap is full, + /// If input ordering shares a common sort prefix with the TopK, /// check if the computation can be finished early. - /// This is the case if the last row of the current batch is strictly greater than the max row in the heap, - /// comparing only on the shared prefix columns. + /// + /// This is the case if the last row of the current batch is strictly + /// greater than either the shared dynamic-filter threshold prefix or the max + /// row in the local heap, comparing only on the shared prefix columns. fn attempt_early_completion(&mut self, batch: &RecordBatch) -> Result<()> { // Early exit if the batch is empty as there is no last row to extract from it. if batch.num_rows() == 0 { return Ok(()); } - // prefix_row_converter is only `Some` if the input ordering has a common prefix with the TopK, + // common_prefix_row_converter is only `Some` if the input ordering has a common prefix with the TopK, // so early exit if it is `None`. let Some(prefix_converter) = &self.common_sort_prefix_converter else { return Ok(()); }; - // Early exit if the heap is not full (`heap.max()` only returns `Some` if the heap is full). - let Some(max_topk_row) = self.heap.max() else { - return Ok(()); - }; - // Evaluate the prefix for the last row of the current batch. let last_row_idx = batch.num_rows() - 1; let mut batch_prefix_scratch = prefix_converter.empty_rows(1, ESTIMATED_BYTES_PER_ROW); // 1 row with capacity ESTIMATED_BYTES_PER_ROW - self.compute_common_sort_prefix(batch, last_row_idx, &mut batch_prefix_scratch)?; - - // Retrieve the max row from the heap. - let store_entry = self - .heap - .store - .get(max_topk_row.batch_id) - .ok_or(internal_datafusion_err!("Invalid batch id in topK heap"))?; - let max_batch = &store_entry.batch; - let mut heap_prefix_scratch = - prefix_converter.empty_rows(1, ESTIMATED_BYTES_PER_ROW); // 1 row with capacity ESTIMATED_BYTES_PER_ROW - self.compute_common_sort_prefix( - max_batch, - max_topk_row.index, - &mut heap_prefix_scratch, + self.append_common_prefix_row( + prefix_converter, + batch, + last_row_idx, + &mut batch_prefix_scratch, )?; + let batch_common_prefix_row = batch_prefix_scratch.row(0); + let batch_common_prefix = batch_common_prefix_row.as_ref(); + + let finished_by_shared_threshold = self + .filter + .read() + .shared_threshold + .as_ref() + .and_then(TopKThreshold::common_prefix_row) + .map(|common_prefix_row| batch_common_prefix > common_prefix_row) + .unwrap_or(false); + if finished_by_shared_threshold { + self.finished = true; + return Ok(()); + } + + // Early exit only from the local heap once it has a full boundary row. + let Some(boundary) = self.current_heap_boundary()? else { + return Ok(()); + }; - // If the last row's prefix is strictly greater than the max prefix, mark as finished. - if batch_prefix_scratch.row(0).as_ref() > heap_prefix_scratch.row(0).as_ref() { + if self.batch_prefix_exceeds_heap_boundary(batch_common_prefix, boundary)? { self.finished = true; } Ok(()) } - // Helper function to compute the prefix for a given batch and row index, storing the result in scratch. - fn compute_common_sort_prefix( + fn batch_prefix_exceeds_heap_boundary( + &self, + batch_common_prefix: &[u8], + boundary: TopKHeapBoundary<'_>, + ) -> Result { + let Some(heap_common_prefix_row) = + self.encode_topk_common_prefix_row(boundary)? + else { + return Ok(false); + }; + + Ok(batch_common_prefix > heap_common_prefix_row.as_slice()) + } + + fn encode_topk_common_prefix_row( + &self, + boundary: TopKHeapBoundary<'_>, + ) -> Result>> { + let Some(prefix_converter) = &self.common_sort_prefix_converter else { + return Ok(None); + }; + + let mut scratch = prefix_converter.empty_rows(1, ESTIMATED_BYTES_PER_ROW); + self.append_common_prefix_row( + prefix_converter, + boundary.batch, + boundary.row.index, + &mut scratch, + )?; + Ok(Some(scratch.row(0).as_ref().to_vec())) + } + + fn append_common_prefix_row( &self, + prefix_converter: &RowConverter, batch: &RecordBatch, - last_row_idx: usize, + row_idx: usize, scratch: &mut Rows, ) -> Result<()> { - let last_row: Vec = self + let row = batch.slice(row_idx, 1); + let prefix_columns: Vec = self .common_sort_prefix .iter() - .map(|expr| { - expr.expr - .evaluate(&batch.slice(last_row_idx, 1))? - .into_array(1) - }) + .map(|expr| expr.expr.evaluate(&row)?.into_array(1)) .collect::>()?; - self.common_sort_prefix_converter - .as_ref() - .unwrap() - .append(scratch, &last_row)?; + prefix_converter.append(scratch, &prefix_columns)?; Ok(()) } @@ -600,8 +803,9 @@ impl TopK { } = self; let _timer = metrics.baseline.elapsed_compute().timer(); // time updated on drop - // Mark the dynamic filter as complete now that TopK processing is finished. - filter.read().expr().mark_complete(); + // Mark this local TopK as emitted. For shared filters, the final + // local emitter marks the dynamic filter complete. + filter.read().mark_topk_emitted(); // break into record batches as needed let mut batches = vec![]; @@ -709,12 +913,16 @@ impl TopKHeap { /// Adds `row` to this heap. If inserting this new item would /// increase the size past `k`, removes the previously smallest /// item. + /// + /// Returns `Some(EvictedRow)` if an existing row was evicted to + /// make room for `row`, or `None` if the row was inserted into a + /// non-full heap. fn add( &mut self, batch_entry: &mut RecordBatchEntry, row: impl AsRef<[u8]>, index: usize, - ) { + ) -> Option { let batch_id = batch_entry.id; batch_entry.uses += 1; @@ -725,6 +933,26 @@ impl TopKHeap { if self.inner.len() == self.k { let mut prev_min = self.inner.peek_mut().unwrap(); + // Capture evicted row data before `unuse` (which may GC the + // batch from the store) and `replace_with` (which overwrites + // `prev_min` in place). The batch comes from `self.store` for + // cross-batch evictions, or directly from `batch_entry` when + // a row evicts another row from the same in-flight batch + // (entry not yet registered in the store). + let evicted_batch = if prev_min.batch_id == batch_entry.id { + batch_entry.batch.clone() + } else { + self.store + .get(prev_min.batch_id) + .map(|entry| entry.batch.clone()) + .expect("evicted row's batch must be present in the store") + }; + let evicted = EvictedRow { + batch: evicted_batch, + index: prev_min.index, + row_bytes: prev_min.row.clone(), + }; + // Update batch use if prev_min.batch_id == batch_entry.id { batch_entry.uses -= 1; @@ -738,12 +966,15 @@ impl TopKHeap { prev_min.replace_with(row, batch_id, index); self.owned_bytes += prev_min.owned_size(); + + Some(evicted) } else { let new_row = TopKRow::new(row, batch_id, index); self.owned_bytes += new_row.owned_size(); // put the new row into the heap self.inner.push(new_row); - }; + None + } } /// Returns the values stored in this heap, from values low to @@ -755,7 +986,7 @@ impl TopKHeap { /// Returns the values stored in this heap, from values low to /// high, as a single [`RecordBatch`], and a sorted vec of the /// current heap's contents - pub fn emit_with_state(&mut self) -> Result<(Option, Vec)> { + fn emit_with_state(&mut self) -> Result<(Option, Vec)> { // generate sorted rows let topk_rows = std::mem::take(&mut self.inner).into_sorted_vec(); @@ -842,47 +1073,6 @@ impl TopKHeap { + self.store.size() + self.owned_bytes } - - fn get_threshold_values( - &self, - sort_exprs: &[PhysicalSortExpr], - ) -> Result>> { - // If the heap doesn't have k elements yet, we can't create thresholds - let max_row = match self.max() { - Some(row) => row, - None => return Ok(None), - }; - - // Get the batch that contains the max row - let batch_entry = match self.store.get(max_row.batch_id) { - Some(entry) => entry, - None => return internal_err!("Invalid batch ID in TopKRow"), - }; - - // Extract threshold values for each sort expression - let mut scalar_values = Vec::with_capacity(sort_exprs.len()); - for sort_expr in sort_exprs { - // Extract the value for this column from the max row - let expr = Arc::clone(&sort_expr.expr); - let value = expr.evaluate(&batch_entry.batch.slice(max_row.index, 1))?; - - // Convert to scalar value - should be a single value since we're evaluating on a single row batch - let scalar = match value { - ColumnarValue::Scalar(scalar) => scalar, - ColumnarValue::Array(array) if array.len() == 1 => { - // Extract the first (and only) value from the array - ScalarValue::try_from_array(&array, 0)? - } - array => { - return internal_err!("Expected a scalar value, got {:?}", array); - } - }; - - scalar_values.push(scalar); - } - - Ok(Some(scalar_values)) - } } /// Represents one of the top K rows held in this heap. Orders @@ -1057,128 +1247,716 @@ impl RecordBatchStore { } } -#[cfg(test)] -mod tests { - use super::*; - use arrow::array::{BooleanArray, Float64Array, Int32Array}; - use arrow::datatypes::{DataType, Field, Schema}; - use arrow_schema::SortOptions; - use datafusion_common::assert_batches_eq; - use datafusion_physical_expr::expressions::col; - use futures::TryStreamExt; +/// Top-K-per-partition operator state. +/// +/// Sibling to [`TopK`]. Where `TopK` maintains a single global heap, +/// `PartitionedTopK` maintains one [`TopKHeap`] per distinct partition +/// key while sharing a single [`RowConverter`], [`MemoryReservation`], +/// scratch [`Rows`] buffer, and [`TopKMetrics`] across all partitions. +/// +/// This sharing is the point of the type: with N distinct partition +/// keys, a naive `HashMap<_, TopK>` pays N × constant overhead for +/// `RowConverter::new`, `MemoryConsumer::register`, and metric +/// counter setup. `PartitionedTopK` pays it once. +pub(crate) struct PartitionedTopK { + schema: SchemaRef, + metrics: TopKMetrics, + reservation: MemoryReservation, + /// ORDER BY expressions (excludes PARTITION BY). + expr: LexOrdering, + /// Encoder for ORDER BY columns. Reused across partitions. + row_converter: RowConverter, + /// Scratch row buffer reused across `insert_batch` calls. + scratch_rows: Rows, + /// PARTITION BY expressions. + partition_exprs: Vec>, + /// Encoder for the partition key. + partition_converter: RowConverter, + /// One heap per distinct partition key seen so far. + heaps: HashMap, + k: usize, + batch_size: usize, +} - /// This test ensures the size calculation is correct for RecordBatches with multiple columns. - #[test] - fn test_record_batch_store_size() { - // given - let schema = Arc::new(Schema::new(vec![ - Field::new("ints", DataType::Int32, true), - Field::new("float64", DataType::Float64, false), - ])); - let mut record_batch_store = RecordBatchStore::new(); - let int_array = - Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]); // 5 * 4 = 20 - let float64_array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]); // 5 * 8 = 40 +impl PartitionedTopK { + #[expect(clippy::too_many_arguments)] + pub(crate) fn try_new( + partition_id: usize, + schema: SchemaRef, + partition_exprs: Vec>, + partition_sort_fields: Vec, + order_expr: LexOrdering, + k: usize, + batch_size: usize, + runtime: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result { + assert!(k > 0, "PartitionedTopK requires k > 0"); + let reservation = MemoryConsumer::new(format!("PartitionedTopK[{partition_id}]")) + .register(&runtime.memory_pool); - let record_batch_entry = RecordBatchEntry { - id: 0, - batch: RecordBatch::try_new( - schema, - vec![Arc::new(int_array), Arc::new(float64_array)], - ) - .unwrap(), - uses: 1, - }; + let order_sort_fields = build_sort_fields(&order_expr, &schema)?; + let row_converter = RowConverter::new(order_sort_fields)?; + let scratch_rows = + row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); - // when insert record batch entry - record_batch_store.insert(record_batch_entry); - assert_eq!(record_batch_store.batches_size, 60); + let partition_converter = RowConverter::new(partition_sort_fields)?; - // when unuse record batch entry - record_batch_store.unuse(0); - assert_eq!(record_batch_store.batches_size, 0); + Ok(Self { + schema, + metrics: TopKMetrics::new(metrics, partition_id), + reservation, + expr: order_expr, + row_converter, + scratch_rows, + partition_exprs, + partition_converter, + heaps: HashMap::new(), + k, + batch_size, + }) } - /// This test validates that the `try_finish` method marks the TopK operator as finished - /// when the prefix (on column "a") of the last row in the current batch is strictly greater - /// than the max top‑k row. - /// The full sort expression is defined on both columns ("a", "b"), but the input ordering is only on "a". - #[tokio::test] - async fn test_try_finish_marks_finished_with_prefix() -> Result<()> { - // Create a schema with two columns. - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int32, false), - Field::new("b", DataType::Float64, false), - ])); - - // Create sort expressions. - // Full sort: first by "a", then by "b". - let sort_expr_a = PhysicalSortExpr { - expr: col("a", schema.as_ref())?, - options: SortOptions::default(), - }; - let sort_expr_b = PhysicalSortExpr { - expr: col("b", schema.as_ref())?, - options: SortOptions::default(), - }; + /// Demultiplex `batch` rows by partition key, encode the ORDER BY + /// columns once for the whole batch, and feed each partition's + /// rows into its dedicated [`TopKHeap`]. + pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { + let baseline = self.metrics.baseline.clone(); + let _timer = baseline.elapsed_compute().timer(); - // Input ordering uses only column "a" (a prefix of the full sort). - let prefix = vec![sort_expr_a.clone()]; - let full_expr = LexOrdering::from([sort_expr_a, sort_expr_b]); + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(()); + } - // Create a dummy runtime environment and metrics. - let runtime = Arc::new(RuntimeEnv::default()); - let metrics = ExecutionPlanMetricsSet::new(); + // 1. Evaluate + encode partition columns. + let pk_arrays: Vec = self + .partition_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + let pk_rows = self.partition_converter.convert_columns(&pk_arrays)?; + + // 2. Demultiplex row indices by partition key (per-batch). + let mut groups: HashMap> = HashMap::new(); + for i in 0..num_rows { + groups + .entry(pk_rows.row(i).owned()) + .or_default() + .push(i as u32); + } - // Create a TopK instance with k = 3 and batch_size = 2. - let mut topk = TopK::try_new( - 0, - Arc::clone(&schema), - prefix, - full_expr, + // 3. Evaluate ORDER BY columns on the full batch and encode ONCE. + let ob_arrays: Vec = self + .expr + .iter() + .map(|e| e.expr.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + self.scratch_rows.clear(); + self.row_converter + .append(&mut self.scratch_rows, &ob_arrays)?; + + // 4. Per-partition: take the sub-batch, walk indices, dispatch + // qualifying rows into the partition's heap. + let k = self.k; + let mut replacements: usize = 0; + for (pk, indices) in groups { + let heap = self.heaps.entry(pk).or_insert_with(|| TopKHeap::new(k)); + + // Once a heap is full, most rows at high partition cardinality + // are rejected. Skip the gather + batch registration entirely + // when nothing in this partition group can improve the heap. + let any_qualify = indices.iter().any(|&orig_idx| { + let bytes = self.scratch_rows.row(orig_idx as usize); + match heap.max() { + Some(max_row) => bytes.as_ref() < max_row.row(), + None => true, + } + }); + if !any_qualify { + continue; + } + + let indices_arr = UInt32Array::from(indices); + let sub_batch = take_record_batch(batch, &indices_arr)?; + let mut entry = heap.register_batch(sub_batch); + + for (sub_idx, &orig_idx) in indices_arr.values().iter().enumerate() { + let row = self.scratch_rows.row(orig_idx as usize); + match heap.max() { + Some(max_row) if row.as_ref() >= max_row.row() => {} + None | Some(_) => { + heap.add(&mut entry, row, sub_idx); + replacements += 1; + } + } + } + + heap.insert_batch_entry(entry); + heap.maybe_compact()?; + } + + if replacements > 0 { + self.metrics.row_replacements.add(replacements); + } + self.reservation.try_resize(self.size())?; + Ok(()) + } + + /// Drain all heaps in partition-key order and return the rows as + /// a stream of coalesced `RecordBatch`es ordered by + /// `(partition_keys, order_keys)`. + pub(crate) fn emit(self) -> Result { + let Self { + schema, + metrics, + reservation: _, + expr: _, + row_converter: _, + scratch_rows: _, + partition_exprs: _, + partition_converter: _, + mut heaps, + k: _, + batch_size, + } = self; + let _timer = metrics.baseline.elapsed_compute().timer(); + + let mut sorted_pks: Vec = heaps.keys().cloned().collect(); + sorted_pks.sort(); + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + + for pk in sorted_pks { + let mut heap = heaps.remove(&pk).expect("key from heaps.keys()"); + if let Some(batch) = heap.emit()? { + (&batch).record_output(&metrics.baseline); + coalescer.push_batch(batch)?; + } + } + coalescer.finish_buffered_batch()?; + + let mut out: Vec> = Vec::new(); + while let Some(b) = coalescer.next_completed_batch() { + out.push(Ok(b)); + } + + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(out), + ))) + } + + /// Total memory currently held by this operator, including all + /// per-partition heaps. + fn size(&self) -> usize { + size_of::() + + self.row_converter.size() + + self.partition_converter.size() + + self.scratch_rows.size() + + self.heaps.values().map(|h| h.size()).sum::() + + self.heaps.capacity() * (size_of::() + size_of::()) + } +} + +/// A run of rows from a single source [`RecordBatch`] that tied at the +/// boundary when inserted. Stored as `(batch, indices)` and materialized +/// at emit time via [`take_record_batch`]. +#[derive(Debug)] +struct TieEntry { + batch: RecordBatch, + /// Indices into `batch` of the rows tied at the (then-current) + /// boundary. Always non-empty by construction. + row_indices: Vec, + /// `get_record_batch_memory_size(&batch)` captured at push time so + /// `RankPartitionState::size()` doesn't recurse through `batch`'s + /// columns on every `try_resize` call. + batch_bytes: usize, +} + +/// Per-partition state for `RANK()` semantics. +/// +/// Composes [`TopKHeap`] as the K-bounded core plus a sibling +/// `Vec` for boundary-tied rows. `RANK ≤ K` keeps the K +/// best rows by ORDER BY plus every row tied at the K-th-best +/// ORDER BY value — the boundary. So the total retained rows can +/// exceed K when ties straddle the boundary. +struct RankPartitionState { + heap: TopKHeap, + ties: Vec, +} + +impl RankPartitionState { + fn size(&self) -> usize { + let ties_buffer = self.ties.capacity() * size_of::(); + let ties_contents: usize = self + .ties + .iter() + .map(|t| t.row_indices.capacity() * size_of::() + t.batch_bytes) + .sum(); + self.heap.size() + ties_buffer + ties_contents + } +} + +/// Sibling to [`PartitionedTopK`] implementing `RANK()` semantics. +/// +/// Per partition, retains the K-best rows plus every row tied at the +/// K-th-best ORDER BY value (so `WHERE rk <= K` may keep more than K +/// rows when ties straddle the boundary). Like [`PartitionedTopK`], +/// the [`RowConverter`], [`MemoryReservation`], scratch [`Rows`] +/// buffer, and [`TopKMetrics`] are shared across all partitions for +/// this operator instance. +/// +/// # Algorithm (per row) +/// +/// For each incoming row, compare its encoded ORDER BY bytes against +/// `heap.max()` — the K-th-best row, which is by definition the +/// admission boundary. `heap.max()` is `None` until the heap fills +/// to K rows: +/// +/// - heap not full (`max() == None`) → forward to the heap +/// - row's ob `==` max → push to ties (no heap call) +/// - row's ob `>` max → drop +/// - row's ob `<` max → forward to heap; on eviction, compare the +/// new `heap.max()` to the evicted row's bytes: if equal, push +/// evicted to ties (still tied at the new boundary's rank); else +/// clear ties (boundary moved up, old ties no longer satisfy +/// `rk ≤ K`) +pub(crate) struct PartitionedTopKRank { + schema: SchemaRef, + metrics: TopKMetrics, + reservation: MemoryReservation, + /// ORDER BY expressions (excludes PARTITION BY). + expr: LexOrdering, + /// Encoder for ORDER BY columns. Reused across partitions. + row_converter: RowConverter, + /// Scratch row buffer reused across `insert_batch` calls. + scratch_rows: Rows, + /// PARTITION BY expressions. + partition_exprs: Vec>, + /// Encoder for the partition key. + partition_converter: RowConverter, + /// Scratch row buffer for partition-key encoding. Reused across + /// `insert_batch` calls (cleared + appended each batch) so we + /// avoid allocating a fresh `Rows` buffer every batch. + partition_scratch_rows: Rows, + /// One rank state per distinct partition key seen so far. + states: HashMap, + k: usize, + batch_size: usize, +} + +impl PartitionedTopKRank { + #[expect(clippy::too_many_arguments)] + pub(crate) fn try_new( + partition_id: usize, + schema: SchemaRef, + partition_exprs: Vec>, + partition_sort_fields: Vec, + order_expr: LexOrdering, + k: usize, + batch_size: usize, + runtime: &Arc, + metrics: &ExecutionPlanMetricsSet, + ) -> Result { + assert!(k > 0, "PartitionedTopKRank requires k > 0"); + let reservation = + MemoryConsumer::new(format!("PartitionedTopKRank[{partition_id}]")) + .register(&runtime.memory_pool); + + let order_sort_fields = build_sort_fields(&order_expr, &schema)?; + let row_converter = RowConverter::new(order_sort_fields)?; + let scratch_rows = + row_converter.empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); + + let partition_converter = RowConverter::new(partition_sort_fields)?; + let partition_scratch_rows = partition_converter + .empty_rows(batch_size, ESTIMATED_BYTES_PER_ROW * batch_size); + + Ok(Self { + schema, + metrics: TopKMetrics::new(metrics, partition_id), + reservation, + expr: order_expr, + row_converter, + scratch_rows, + partition_exprs, + partition_converter, + partition_scratch_rows, + states: HashMap::new(), + k, + batch_size, + }) + } + + /// Demultiplex `batch` rows by partition key, encode the ORDER BY + /// columns once for the whole batch, and feed each partition's + /// rows through the rank classifier into its dedicated heap and + /// ties Vec. + pub(crate) fn insert_batch(&mut self, batch: &RecordBatch) -> Result<()> { + let baseline = self.metrics.baseline.clone(); + let _timer = baseline.elapsed_compute().timer(); + + let num_rows = batch.num_rows(); + if num_rows == 0 { + return Ok(()); + } + + // Captured once so the per-tie push from this batch can reuse + // it (computing `get_record_batch_memory_size` is O(cols × + // buffer walk) and we'd otherwise pay it per push and again + // per `try_resize` call). + let input_batch_bytes = get_record_batch_memory_size(batch); + + // 1. Evaluate + encode partition columns into the reusable + // scratch (cleared then appended). + let pk_arrays: Vec = self + .partition_exprs + .iter() + .map(|e| e.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + self.partition_scratch_rows.clear(); + self.partition_converter + .append(&mut self.partition_scratch_rows, &pk_arrays)?; + let pk_rows = &self.partition_scratch_rows; + + // 2. Demultiplex row indices by partition key (per-batch). + let mut groups: HashMap> = HashMap::new(); + for i in 0..num_rows { + groups + .entry(pk_rows.row(i).owned()) + .or_default() + .push(i as u32); + } + + // 3. Evaluate ORDER BY columns on the full batch and encode ONCE. + let ob_arrays: Vec = self + .expr + .iter() + .map(|e| e.expr.evaluate(batch).and_then(|v| v.into_array(num_rows))) + .collect::>()?; + self.scratch_rows.clear(); + self.row_converter + .append(&mut self.scratch_rows, &ob_arrays)?; + + // 4. Per-partition: classify each row and dispatch. + let k = self.k; + let mut replacements: usize = 0; + + for (pk, indices) in groups { + let state = self.states.entry(pk).or_insert_with(|| RankPartitionState { + heap: TopKHeap::new(k), + ties: Vec::new(), + }); + + // Equal indices for THIS batch only. Coalesced into a single + // tie entry at the end of the partition's loop. Discarded if + // the boundary moves up mid-loop (those rows were tied to the + // old boundary, which is now strictly worse than the new K-th). + let mut equal_indices: Vec = Vec::new(); + // Lazy-registered: only attached if at least one row reaches + // the heap from this batch in this partition. + let mut entry: Option = None; + + for &orig_idx in &indices { + let row = self.scratch_rows.row(orig_idx as usize); + + // Classify against the current K-th-best (the heap top). + // `heap.max()` returns `None` while the heap is filling, + // so unclassified rows fall through to the heap path. + let classification = state + .heap + .max() + .map(|max_row| row.as_ref().cmp(max_row.row())); + + match classification { + Some(Ordering::Equal) => { + equal_indices.push(orig_idx); + continue; + } + Some(Ordering::Greater) => continue, + Some(Ordering::Less) | None => { + // Heap path: heap not yet full, or row strictly + // better than the current boundary. + let entry_ref = entry.get_or_insert_with(|| { + state.heap.register_batch(batch.clone()) + }); + if let Some(EvictedRow { + batch: evicted_batch, + index: evicted_index, + row_bytes: evicted_bytes, + }) = state.heap.add(entry_ref, row, orig_idx as usize) + { + // Compare the new boundary (post-eviction heap + // top) against the evicted row's bytes — both + // already in encoded form, no clones needed. + let boundary_changed = state + .heap + .max() + .expect("heap was full to evict; must still be full") + .row() + != evicted_bytes.as_slice(); + if boundary_changed { + // Boundary moved up — prior ties (across + // all prior batches) and equal_indices + // accumulated earlier in THIS batch were + // tied to the old boundary, now strictly + // worse than the new K-th-best. Discard. + state.ties.clear(); + equal_indices.clear(); + } else { + // Boundary unchanged — evicted row is tied + // at the (unchanged) boundary; push as a + // single-row entry. + let batch_bytes = + get_record_batch_memory_size(&evicted_batch); + state.ties.push(TieEntry { + batch: evicted_batch, + row_indices: vec![evicted_index as u32], + batch_bytes, + }); + } + } + replacements += 1; + } + } + } + + if let Some(e) = entry { + state.heap.insert_batch_entry(e); + state.heap.maybe_compact()?; + } + + // Commit this batch's ties as a single entry. + if !equal_indices.is_empty() { + state.ties.push(TieEntry { + batch: batch.clone(), + row_indices: equal_indices, + batch_bytes: input_batch_bytes, + }); + } + } + + if replacements > 0 { + self.metrics.row_replacements.add(replacements); + } + self.reservation.try_resize(self.size())?; + Ok(()) + } + + /// Drain all heaps and ties in partition-key order and return the + /// rows as a stream of coalesced [`RecordBatch`]es ordered by + /// `(partition_keys, order_keys)`. Within a partition, heap rows + /// come first (sorted by ob), then tie rows (all sharing the + /// boundary ob). + pub(crate) fn emit(self) -> Result { + let Self { + schema, + metrics, + reservation: _, + expr: _, + row_converter: _, + scratch_rows: _, + partition_exprs: _, + partition_converter: _, + partition_scratch_rows: _, + mut states, + k: _, + batch_size, + } = self; + let _timer = metrics.baseline.elapsed_compute().timer(); + + let mut sorted_pks: Vec = states.keys().cloned().collect(); + sorted_pks.sort(); + + let mut coalescer = BatchCoalescer::new(Arc::clone(&schema), batch_size); + + for pk in sorted_pks { + let RankPartitionState { mut heap, ties, .. } = + states.remove(&pk).expect("key from states.keys()"); + if let Some(batch) = heap.emit()? { + (&batch).record_output(&metrics.baseline); + coalescer.push_batch(batch)?; + } + for tie in ties { + let indices = UInt32Array::from(tie.row_indices); + let tie_batch = take_record_batch(&tie.batch, &indices)?; + (&tie_batch).record_output(&metrics.baseline); + coalescer.push_batch(tie_batch)?; + } + } + coalescer.finish_buffered_batch()?; + + let mut out: Vec> = Vec::new(); + while let Some(b) = coalescer.next_completed_batch() { + out.push(Ok(b)); + } + + Ok(Box::pin(RecordBatchStreamAdapter::new( + schema, + futures::stream::iter(out), + ))) + } + + /// Total memory currently held, including all per-partition states. + fn size(&self) -> usize { + size_of::() + + self.row_converter.size() + + self.partition_converter.size() + + self.scratch_rows.size() + + self.partition_scratch_rows.size() + + self.states.values().map(|s| s.size()).sum::() + + self.states.capacity() + * (size_of::() + size_of::()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{BooleanArray, Float64Array, Int32Array}; + use arrow::datatypes::{DataType, Field, Schema}; + use arrow_schema::SortOptions; + use datafusion_common::assert_batches_eq; + use datafusion_physical_expr::{DynamicFilterTracking, expressions::col}; + use futures::TryStreamExt; + + /// This test ensures the size calculation is correct for RecordBatches with multiple columns. + #[test] + fn test_record_batch_store_size() { + // given + let schema = Arc::new(Schema::new(vec![ + Field::new("ints", DataType::Int32, true), + Field::new("float64", DataType::Float64, false), + ])); + let mut record_batch_store = RecordBatchStore::new(); + let int_array = + Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4), Some(5)]); // 5 * 4 = 20 + let float64_array = Float64Array::from(vec![1.0, 2.0, 3.0, 4.0, 5.0]); // 5 * 8 = 40 + + let record_batch_entry = RecordBatchEntry { + id: 0, + batch: RecordBatch::try_new( + schema, + vec![Arc::new(int_array), Arc::new(float64_array)], + ) + .unwrap(), + uses: 1, + }; + + // when insert record batch entry + record_batch_store.insert(record_batch_entry); + assert_eq!(record_batch_store.batches_size, 60); + + // when unuse record batch entry + record_batch_store.unuse(0); + assert_eq!(record_batch_store.batches_size, 0); + } + + fn make_ab_schema() -> SchemaRef { + make_ab_schema_with_nullable_a(false) + } + + fn make_ab_schema_with_nullable_a(a_nullable: bool) -> SchemaRef { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, a_nullable), + Field::new("b", DataType::Float64, false), + ])) + } + + // Local TopK tests use one emitter; shared-filter cases pass the partition count explicitly. + fn make_topk_filter() -> Arc> { + make_shared_topk_filter(1) + } + + fn make_shared_topk_filter( + topk_emitter_count: usize, + ) -> Arc> { + Arc::new(RwLock::new( + TopKDynamicFilters::new_with_topk_emitter_count( + Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))), + topk_emitter_count, + ), + )) + } + + /// Builds the `(a, b)` fixture used by prefix-completion tests: + /// full sort `(a, b)`, input prefix `[a]`, `k = 3`, and batch size 2. + fn make_ab_topk( + schema: SchemaRef, + filter: Arc>, + ) -> Result { + make_ab_topk_with_options(0, schema, filter, SortOptions::default()) + } + + fn make_ab_topk_with_options( + partition_id: usize, + schema: SchemaRef, + filter: Arc>, + a_options: SortOptions, + ) -> Result { + let sort_expr_a = PhysicalSortExpr { + expr: col("a", schema.as_ref())?, + options: a_options, + }; + let sort_expr_b = PhysicalSortExpr { + expr: col("b", schema.as_ref())?, + options: SortOptions::default(), + }; + + TopK::try_new( + partition_id, + schema, + vec![sort_expr_a.clone()], + LexOrdering::from([sort_expr_a, sort_expr_b]), 3, 2, - runtime, - &metrics, - Arc::new(RwLock::new(TopKDynamicFilters::new(Arc::new( - DynamicFilterPhysicalExpr::new(vec![], lit(true)), - )))), - )?; + Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + filter, + ) + } - // Create the first batch with two columns: - // Column "a": [1, 1, 2], Column "b": [20.0, 15.0, 30.0]. - let array_a1: ArrayRef = - Arc::new(Int32Array::from(vec![Some(1), Some(1), Some(2)])); - let array_b1: ArrayRef = Arc::new(Float64Array::from(vec![20.0, 15.0, 30.0])); - let batch1 = RecordBatch::try_new(Arc::clone(&schema), vec![array_a1, array_b1])?; + fn make_ab_batch( + schema: SchemaRef, + a: &[Option], + b: &[f64], + ) -> Result { + Ok(RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from(a.to_vec())) as ArrayRef, + Arc::new(Float64Array::from(b.to_vec())) as ArrayRef, + ], + )?) + } - // Insert the first batch. - // At this point the heap is not yet “finished” because the prefix of the last row of the batch - // is not strictly greater than the prefix of the max top‑k row (both being `2`). - topk.insert_batch(batch1)?; - assert!( - !topk.finished, - "Expected 'finished' to be false after the first batch." - ); + type AbRow = (Option, f64); - // Create the second batch with two columns: - // Column "a": [2, 3], Column "b": [10.0, 20.0]. - let array_a2: ArrayRef = Arc::new(Int32Array::from(vec![Some(2), Some(3)])); - let array_b2: ArrayRef = Arc::new(Float64Array::from(vec![10.0, 20.0])); - let batch2 = RecordBatch::try_new(Arc::clone(&schema), vec![array_a2, array_b2])?; + fn make_ab_rows_batch(schema: SchemaRef, rows: &[AbRow]) -> Result { + let (a, b): (Vec<_>, Vec<_>) = rows.iter().copied().unzip(); + make_ab_batch(schema, &a, &b) + } - // Insert the second batch. - // The last row in this batch has a prefix value of `3`, - // which is strictly greater than the max top‑k row (with value `2`), - // so try_finish should mark the TopK as finished. - topk.insert_batch(batch2)?; - assert!( - topk.finished, - "Expected 'finished' to be true after the second batch." - ); + #[tokio::test] + async fn test_early_completion_marks_finished_with_prefix() -> Result<()> { + let schema = make_ab_schema(); + let mut topk = make_ab_topk(Arc::clone(&schema), make_topk_filter())?; + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(1), Some(1), Some(2)], + &[20.0, 15.0, 30.0], + )?)?; + assert!(!topk.finished); + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(2), Some(3)], + &[10.0, 20.0], + )?)?; + assert!(topk.finished); - // Verify the TopK correctly emits the top k rows from both batches - // (the value 10.0 for b is from the second batch). let results: Vec<_> = topk.emit()?.try_collect().await?; assert_batches_eq!( &[ @@ -1196,50 +1974,306 @@ mod tests { Ok(()) } - /// This test verifies that the dynamic filter is marked as complete after TopK processing finishes. + /// Regression test for #22849: a batch whose rows are entirely rejected by the + /// heap's dynamic filter must still trigger `attempt_early_completion` when its + /// last row's prefix is worse than the heap's worst. #[tokio::test] - async fn test_topk_marks_filter_complete() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + async fn test_early_completion_fires_when_filter_rejects_entire_batch() -> Result<()> + { + let schema = make_ab_schema(); + let mut topk = make_ab_topk(Arc::clone(&schema), make_topk_filter())?; + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(1), Some(1), Some(2)], + &[20.0, 15.0, 30.0], + )?)?; + assert!(!topk.finished); + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(3), Some(3)], + &[10.0, 20.0], + )?)?; + assert!(topk.finished); + + let results: Vec<_> = topk.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+---+------+", + "| a | b |", + "+---+------+", + "| 1 | 15.0 |", + "| 1 | 20.0 |", + "| 2 | 30.0 |", + "+---+------+", + ], + &results + ); + + Ok(()) + } + + #[tokio::test] + async fn test_early_completion_fires_when_batch_makes_no_replacements() -> Result<()> + { + let schema = make_ab_schema(); + let filter = make_topk_filter(); + let mut topk = make_ab_topk(Arc::clone(&schema), Arc::clone(&filter))?; + + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(1), Some(1), Some(2)], + &[20.0, 15.0, 30.0], + )?)?; + assert!(!topk.finished); + + let replacements_before = topk.metrics.row_replacements.value(); + + // Keep the dynamic filter permissive so the second batch reaches + // `find_new_topk_items`; all of its rows are worse than the heap max, + // so this specifically exercises the `replacements == 0` path. + filter.read().expr().update(lit(true))?; + topk.insert_batch(make_ab_batch( + Arc::clone(&schema), + &[Some(3), Some(3)], + &[10.0, 20.0], + )?)?; + assert_eq!(topk.metrics.row_replacements.value(), replacements_before); + assert!(topk.finished); + + let results: Vec<_> = topk.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+---+------+", + "| a | b |", + "+---+------+", + "| 1 | 15.0 |", + "| 1 | 20.0 |", + "| 2 | 30.0 |", + "+---+------+", + ], + &results + ); + + Ok(()) + } + struct SharedPrefixCase { + name: &'static str, + a_nullable: bool, + a_options: SortOptions, + threshold_source_rows: &'static [AbRow], + lagging_partition_rows: &'static [AbRow], + expected_finished: bool, + } + + fn assert_shared_prefix_case(case: SharedPrefixCase) -> Result<()> { + let schema = make_ab_schema_with_nullable_a(case.a_nullable); + let filter = make_shared_topk_filter(2); + + let mut threshold_source = make_ab_topk_with_options( + 0, + Arc::clone(&schema), + Arc::clone(&filter), + case.a_options, + )?; + threshold_source.insert_batch(make_ab_rows_batch( + Arc::clone(&schema), + case.threshold_source_rows, + )?)?; + assert!( + filter + .read() + .shared_threshold + .as_ref() + .and_then(TopKThreshold::common_prefix_row) + .is_some(), + "{}: threshold-source partition should establish the shared prefix threshold", + case.name + ); + + let mut lagging_partition = make_ab_topk_with_options( + 1, + Arc::clone(&schema), + Arc::clone(&filter), + case.a_options, + )?; + lagging_partition + .insert_batch(make_ab_rows_batch(schema, case.lagging_partition_rows)?)?; + + assert!( + lagging_partition.heap.inner.is_empty(), + "{}: lagging partition's local heap should remain empty", + case.name + ); + assert_eq!( + lagging_partition.finished, case.expected_finished, + "{}", + case.name + ); + + Ok(()) + } + + #[test] + fn test_shared_filter_can_finish_partition_before_local_heap_is_full() -> Result<()> { + assert_shared_prefix_case(SharedPrefixCase { + name: "shared threshold should finish lagging partition", + a_nullable: false, + a_options: SortOptions::default(), + threshold_source_rows: &[(Some(1), 20.0), (Some(1), 15.0), (Some(2), 30.0)], + lagging_partition_rows: &[(Some(3), 10.0), (Some(3), 20.0)], + expected_finished: true, + }) + } + + #[test] + fn test_shared_prefix_threshold_boundary_cases() -> Result<()> { + for case in [ + SharedPrefixCase { + name: "equal prefix cannot prove completion", + a_nullable: false, + a_options: SortOptions::default(), + threshold_source_rows: &[ + (Some(1), 20.0), + (Some(1), 15.0), + (Some(2), 30.0), + ], + lagging_partition_rows: &[(Some(2), 40.0), (Some(2), 50.0)], + expected_finished: false, + }, + SharedPrefixCase { + name: "descending prefix uses sort-order row encoding", + a_nullable: false, + a_options: SortOptions { + descending: true, + nulls_first: true, + }, + threshold_source_rows: &[ + (Some(10), 1.0), + (Some(10), 2.0), + (Some(9), 3.0), + ], + lagging_partition_rows: &[(Some(8), 1.0), (Some(8), 2.0)], + expected_finished: true, + }, + SharedPrefixCase { + name: "NULLS LAST prefix uses sort-order row encoding", + a_nullable: true, + a_options: SortOptions { + descending: false, + nulls_first: false, + }, + threshold_source_rows: &[ + (Some(1), 20.0), + (Some(1), 15.0), + (Some(2), 30.0), + ], + lagging_partition_rows: &[(None, 10.0), (None, 20.0)], + expected_finished: true, + }, + ] { + assert_shared_prefix_case(case)?; + } + Ok(()) + } + + fn make_single_column_topk( + dynamic_filter: Arc, + ) -> Result<(SchemaRef, TopK)> { + make_single_column_topk_with_filter( + 0, + Arc::new(RwLock::new(TopKDynamicFilters::new(dynamic_filter))), + ) + } + + fn make_single_column_topk_with_filter( + partition_id: usize, + filter: Arc>, + ) -> Result<(SchemaRef, TopK)> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); let sort_expr = PhysicalSortExpr { expr: col("a", schema.as_ref())?, options: SortOptions::default(), }; - let full_expr = LexOrdering::from([sort_expr.clone()]); - let prefix = vec![sort_expr]; + let topk = TopK::try_new( + partition_id, + Arc::clone(&schema), + vec![sort_expr.clone()], + LexOrdering::from([sort_expr]), + 2, + 10, + Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + filter, + )?; + + Ok((schema, topk)) + } + + #[tokio::test] + async fn test_topk_marks_filter_complete() -> Result<()> { + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); + let dynamic_filter_clone = Arc::clone(&dynamic_filter); + let (schema, mut topk) = make_single_column_topk(dynamic_filter)?; + + let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), Some(1), Some(2)])); + let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array])?; + topk.insert_batch(batch)?; + + let _results: Vec<_> = topk.emit()?.try_collect().await?; - // Create a dummy runtime environment and metrics - let runtime = Arc::new(RuntimeEnv::default()); - let metrics = ExecutionPlanMetricsSet::new(); + tokio::time::timeout( + std::time::Duration::from_secs(1), + dynamic_filter_clone.wait_complete(), + ) + .await + .expect("single-emitter TopK should mark the dynamic filter complete"); + + Ok(()) + } - // Create a dynamic filter that we'll check for completion + #[tokio::test] + async fn test_shared_topk_filter_completes_after_last_emitter() -> Result<()> { let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new(vec![], lit(true))); let dynamic_filter_clone = Arc::clone(&dynamic_filter); + let shared_filter = Arc::new(RwLock::new( + TopKDynamicFilters::new_with_topk_emitter_count(dynamic_filter, 2), + )); - // Create a TopK instance - let mut topk = TopK::try_new( - 0, - Arc::clone(&schema), - prefix, - full_expr, - 2, - 10, - runtime, - &metrics, - Arc::new(RwLock::new(TopKDynamicFilters::new(dynamic_filter))), - )?; + let (schema, mut topk_0) = + make_single_column_topk_with_filter(0, Arc::clone(&shared_filter))?; + let (_, mut topk_1) = + make_single_column_topk_with_filter(1, Arc::clone(&shared_filter))?; let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(3), Some(1), Some(2)])); let batch = RecordBatch::try_new(Arc::clone(&schema), vec![array])?; - topk.insert_batch(batch)?; + topk_0.insert_batch(batch)?; + let _results: Vec<_> = topk_0.emit()?.try_collect().await?; - // Call emit to finish TopK processing - let _results: Vec<_> = topk.emit()?.try_collect().await?; + let dynamic_filter_expr: Arc = + Arc::::clone(&dynamic_filter_clone); + assert!( + matches!( + DynamicFilterTracking::classify(&dynamic_filter_expr), + DynamicFilterTracking::Watching(_) + ), + "the shared filter should remain watchable until every TopK emits" + ); + + let array: ArrayRef = Arc::new(Int32Array::from(vec![Some(6), Some(4), Some(5)])); + let batch = RecordBatch::try_new(schema, vec![array])?; + topk_1.insert_batch(batch)?; + let _results: Vec<_> = topk_1.emit()?.try_collect().await?; - // After emit is called, the dynamic filter should be marked as complete - // wait_complete() should return immediately - dynamic_filter_clone.wait_complete().await; + tokio::time::timeout( + std::time::Duration::from_secs(1), + dynamic_filter_clone.wait_complete(), + ) + .await + .expect("the final shared TopK emitter should mark the dynamic filter complete"); Ok(()) } @@ -1423,4 +2457,713 @@ mod tests { Ok(()) } + + /// Builds a `(pk Int32, val Int32)` schema and a `PartitionedTopK` + /// partitioned by `pk` with order `val ASC`. Helper for the + /// `PartitionedTopK` tests below. + fn build_partitioned_topk(k: usize) -> Result<(Arc, PartitionedTopK)> { + build_partitioned_topk_with_opts(k, SortOptions::default(), false) + } + + /// Variant of [`build_partitioned_topk`] that lets the test pick the + /// `val` column's `SortOptions` (direction, null ordering) and + /// nullability. Used by tests that exercise the shared encoder + /// across `ASC`/`DESC` and `NULLS FIRST/LAST` paths. + fn build_partitioned_topk_with_opts( + k: usize, + val_sort_options: SortOptions, + val_nullable: bool, + ) -> Result<(Arc, PartitionedTopK)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new("val", DataType::Int32, val_nullable), + ])); + + let pk_expr: Arc = col("pk", schema.as_ref())?; + let pk_sort_expr = PhysicalSortExpr { + expr: Arc::clone(&pk_expr), + options: SortOptions::default(), + }; + let val_sort_expr = PhysicalSortExpr { + expr: col("val", schema.as_ref())?, + options: val_sort_options, + }; + + let partition_sort_fields = build_sort_fields(&[pk_sort_expr], &schema)?; + let order_expr = LexOrdering::from([val_sort_expr]); + + let state = PartitionedTopK::try_new( + 0, + Arc::clone(&schema), + vec![pk_expr], + partition_sort_fields, + order_expr, + k, + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + )?; + Ok((schema, state)) + } + + fn pk_val_batch( + schema: &Arc, + pks: Vec, + vals: Vec, + ) -> Result { + Ok(RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(pks)), + Arc::new(Int32Array::from(vals)), + ], + )?) + } + + /// Variant of [`pk_val_batch`] that accepts nullable `val`s. Used by + /// tests that exercise null-ordering through the shared encoder. + fn nullable_pk_val_batch( + schema: &Arc, + pks: Vec, + vals: Vec>, + ) -> Result { + Ok(RecordBatch::try_new( + Arc::clone(schema), + vec![ + Arc::new(Int32Array::from(pks)), + Arc::new(Int32Array::from(vals)), + ], + )?) + } + + /// Multiple distinct partition keys interleaved within a single + /// input batch — the per-batch demux, per-partition heap eviction, + /// and partition-key-ordered emit must all behave correctly. + #[tokio::test] + async fn test_partitioned_topk_multi_partition_within_batch() -> Result<()> { + let (schema, mut state) = build_partitioned_topk(2)?; + + // pk=1 vals: 10, 5, 8 → top-2 ASC = [5, 8] + // pk=2 vals: 20, 15 → top-2 ASC = [15, 20] + // pk=3 vals: 7 → top-2 ASC = [7] + let batch = + pk_val_batch(&schema, vec![1, 2, 1, 2, 1, 3], vec![10, 20, 5, 15, 8, 7])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 8 |", + "| 2 | 15 |", + "| 2 | 20 |", + "| 3 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// State must accumulate across `insert_batch` calls: a partition + /// key seen in batch 1 should still own its heap when batch 2 + /// arrives, and a row in batch 2 that beats the existing K-th + /// best should evict the loser. + #[tokio::test] + async fn test_partitioned_topk_cross_batch_eviction() -> Result<()> { + let (schema, mut state) = build_partitioned_topk(2)?; + + // Batch 1: pk=1 fills the heap with [50, 40]. + state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![50, 40])?)?; + + // Batch 2: pk=1 sees a smaller value (10) — it must evict 50. + // pk=2 appears for the first time mid-stream. + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 2, 1], + vec![10, 99, 60], // 60 > 40 stays on top, gets discarded + )?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 10 |", + "| 1 | 40 |", + "| 2 | 99 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// Empty input must produce an empty output stream, not panic. + #[tokio::test] + async fn test_partitioned_topk_empty_input() -> Result<()> { + let (_schema, state) = build_partitioned_topk(3)?; + let results: Vec<_> = state.emit()?.try_collect().await?; + assert!(results.is_empty(), "empty input → empty output"); + Ok(()) + } + + /// `fetch = 1` is a common case (rn = 1 filter). The heap should + /// hold exactly one row per partition: the partition's minimum. + #[tokio::test] + async fn test_partitioned_topk_fetch_one() -> Result<()> { + let (schema, mut state) = build_partitioned_topk(1)?; + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 1, 2, 2, 3], + vec![3, 1, 9, 4, 7], + )?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 1 |", + "| 2 | 4 |", + "| 3 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ORDER BY val DESC` exercises the shared encoder's sort-direction + /// handling: the row converter must flip the sort sign for `val` so + /// that larger values compare smaller in row-encoded form. Each + /// partition should keep its top-K *largest* values. + #[tokio::test] + async fn test_partitioned_topk_desc_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_with_opts( + 2, + SortOptions { + descending: true, + nulls_first: false, + }, + false, + )?; + + // pk=1 vals: 10, 5, 8, 12 → top-2 DESC = [12, 10] + // pk=2 vals: 20, 15, 25 → top-2 DESC = [25, 20] + let batch = pk_val_batch( + &schema, + vec![1, 2, 1, 2, 1, 1, 2], + vec![10, 20, 5, 15, 8, 12, 25], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 12 |", + "| 1 | 10 |", + "| 2 | 25 |", + "| 2 | 20 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// NULL sort values exercise the shared encoder's null-ordering + /// handling. With `ASC NULLS LAST`, NULLs sort *after* every + /// non-NULL value, so a partition whose only non-NULL value beats + /// a NULL must evict the NULL when `K = 1`. A partition that holds + /// only NULLs must still emit them. + #[tokio::test] + async fn test_partitioned_topk_nulls_last_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_with_opts( + 1, + SortOptions { + descending: false, + nulls_first: false, + }, + true, + )?; + + // pk=1 vals: NULL, 7, NULL → top-1 ASC NULLS LAST = [7] + // pk=2 vals: NULL → top-1 = [NULL] + // pk=3 vals: NULL, 4, 2 → top-1 = [2] + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 1, 3, 3, 3], + vec![None, None, Some(7), None, None, Some(4), Some(2)], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 7 |", + "| 2 | |", + "| 3 | 2 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ASC NULLS FIRST` (the `SortOptions::default()`) sorts NULLs + /// *before* every non-NULL value, so under `fetch = K` a partition's + /// NULLs are kept preferentially over larger non-NULL values. + #[tokio::test] + async fn test_partitioned_topk_nulls_first_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_with_opts( + 2, + SortOptions { + descending: false, + nulls_first: true, + }, + true, + )?; + + // pk=1 vals: NULL, 5, NULL, 8 → top-2 ASC NULLS FIRST = [NULL, NULL] + // pk=2 vals: 7, NULL → top-2 = [NULL, 7] + // pk=3 vals: 3, 1 → top-2 = [1, 3] + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 3, 1, 2, 1, 3], + vec![ + None, + Some(7), + Some(5), + Some(3), + None, + None, + Some(8), + Some(1), + ], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | |", + "| 1 | |", + "| 2 | |", + "| 2 | 7 |", + "| 3 | 1 |", + "| 3 | 3 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + // ==================================================================== + // PartitionedTopKRank operator tests + // + // These mirror the PartitionedTopK tests above plus three RANK-specific + // cases for the Equal / boundary-shift / boundary-unchanged-eviction + // arms in `PartitionedTopKRank::insert_batch`. + // ==================================================================== + + /// Builds a `(pk Int32, val Int32)` schema and a `PartitionedTopKRank` + /// keyed on `pk ASC` (partition) and `val ASC` (ORDER BY). + fn build_partitioned_topk_rank( + k: usize, + ) -> Result<(Arc, PartitionedTopKRank)> { + build_partitioned_topk_rank_with_opts(k, SortOptions::default(), false) + } + + /// Variant of [`build_partitioned_topk_rank`] that lets the test pick + /// the `val` column's `SortOptions` (direction, null ordering) and + /// nullability. + fn build_partitioned_topk_rank_with_opts( + k: usize, + val_sort_options: SortOptions, + val_nullable: bool, + ) -> Result<(Arc, PartitionedTopKRank)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::Int32, false), + Field::new("val", DataType::Int32, val_nullable), + ])); + + let pk_expr: Arc = col("pk", schema.as_ref())?; + let pk_sort_expr = PhysicalSortExpr { + expr: Arc::clone(&pk_expr), + options: SortOptions::default(), + }; + let val_sort_expr = PhysicalSortExpr { + expr: col("val", schema.as_ref())?, + options: val_sort_options, + }; + + let partition_sort_fields = build_sort_fields(&[pk_sort_expr], &schema)?; + let order_expr = LexOrdering::from([val_sort_expr]); + + let state = PartitionedTopKRank::try_new( + 0, + Arc::clone(&schema), + vec![pk_expr], + partition_sort_fields, + order_expr, + k, + 8, // batch_size + &Arc::new(RuntimeEnv::default()), + &ExecutionPlanMetricsSet::new(), + )?; + Ok((schema, state)) + } + + /// Multiple distinct partition keys interleaved within a single + /// input batch — the per-batch demux, per-partition heap eviction, + /// and partition-key-ordered emit must all behave correctly. No + /// ties: result should match a `ROW_NUMBER` top-K under the same K. + #[tokio::test] + async fn test_partitioned_topk_rank_multi_partition_within_batch() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 10, 5, 8 → top-2 ASC = [5, 8] + // pk=2 vals: 20, 15 → top-2 ASC = [15, 20] + // pk=3 vals: 7 → top-2 ASC = [7] + let batch = + pk_val_batch(&schema, vec![1, 2, 1, 2, 1, 3], vec![10, 20, 5, 15, 8, 7])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 8 |", + "| 2 | 15 |", + "| 2 | 20 |", + "| 3 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// State must accumulate across `insert_batch` calls. A row in + /// batch 2 that's strictly better than the existing K-th must + /// evict it; an evicted row whose bytes match the new boundary + /// becomes a `TieEntry` pinned to the prior batch. + #[tokio::test] + async fn test_partitioned_topk_rank_cross_batch_eviction() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // Batch 1: pk=1 fills the heap with [50, 40]. + state.insert_batch(&pk_val_batch(&schema, vec![1, 1], vec![50, 40])?)?; + + // Batch 2: pk=1 sees a smaller value (10) — it must evict 50; + // 60 > 40 so it's dropped. pk=2 appears mid-stream. + state.insert_batch(&pk_val_batch(&schema, vec![1, 2, 1], vec![10, 99, 60])?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 10 |", + "| 1 | 40 |", + "| 2 | 99 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// Empty input must produce an empty output stream, not panic. + #[tokio::test] + async fn test_partitioned_topk_rank_empty_input() -> Result<()> { + let (_schema, state) = build_partitioned_topk_rank(3)?; + let results: Vec<_> = state.emit()?.try_collect().await?; + assert!(results.is_empty(), "empty input → empty output"); + Ok(()) + } + + /// `fetch = 1` is a common case (rk = 1 filter) and exercises the + /// boundary-defined-immediately path: after the first admission per + /// partition, `heap.max()` is `Some`, so every subsequent row goes + /// through full Equal/Greater/Less classification. + #[tokio::test] + async fn test_partitioned_topk_rank_fetch_one() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(1)?; + state.insert_batch(&pk_val_batch( + &schema, + vec![1, 1, 2, 2, 3], + vec![3, 1, 9, 4, 7], + )?)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 1 |", + "| 2 | 4 |", + "| 3 | 7 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ORDER BY val DESC` exercises the shared encoder's sort-direction + /// handling: the row converter flips the sort sign for `val` so + /// larger values compare smaller in row-encoded form. Each + /// partition keeps its top-K *largest* values. + #[tokio::test] + async fn test_partitioned_topk_rank_desc_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank_with_opts( + 2, + SortOptions { + descending: true, + nulls_first: false, + }, + false, + )?; + + // pk=1 vals: 10, 5, 8, 12 → top-2 DESC = [12, 10] + // pk=2 vals: 20, 15, 25 → top-2 DESC = [25, 20] + let batch = pk_val_batch( + &schema, + vec![1, 2, 1, 2, 1, 1, 2], + vec![10, 20, 5, 15, 8, 12, 25], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 12 |", + "| 1 | 10 |", + "| 2 | 25 |", + "| 2 | 20 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// NULL sort values exercise the shared encoder's null-ordering + /// handling. With `ASC NULLS LAST`, NULLs sort *after* every + /// non-NULL value, so a partition whose only non-NULL value beats + /// a NULL must evict the NULL when `K = 1`. A partition that holds + /// only NULLs must still emit them. + #[tokio::test] + async fn test_partitioned_topk_rank_nulls_last_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank_with_opts( + 1, + SortOptions { + descending: false, + nulls_first: false, + }, + true, + )?; + + // pk=1 vals: NULL, 7, NULL → top-1 ASC NULLS LAST = [7] + // pk=2 vals: NULL → top-1 = [NULL] + // pk=3 vals: NULL, 4, 2 → top-1 = [2] + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 1, 3, 3, 3], + vec![None, None, Some(7), None, None, Some(4), Some(2)], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 7 |", + "| 2 | |", + "| 3 | 2 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// `ASC NULLS FIRST` (the `SortOptions::default()`) sorts NULLs + /// *before* every non-NULL value, so under `fetch = K` a partition's + /// NULLs are kept preferentially over larger non-NULL values. + #[tokio::test] + async fn test_partitioned_topk_rank_nulls_first_ordering() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank_with_opts( + 2, + SortOptions { + descending: false, + nulls_first: true, + }, + true, + )?; + + // pk=1 vals: NULL, 5, NULL, 8 → top-2 ASC NULLS FIRST = [NULL, NULL] + // pk=2 vals: 7, NULL → top-2 = [NULL, 7] + // pk=3 vals: 3, 1 → top-2 = [1, 3] + let batch = nullable_pk_val_batch( + &schema, + vec![1, 2, 1, 3, 1, 2, 1, 3], + vec![ + None, + Some(7), + Some(5), + Some(3), + None, + None, + Some(8), + Some(1), + ], + )?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | |", + "| 1 | |", + "| 2 | |", + "| 2 | 7 |", + "| 3 | 1 |", + "| 3 | 3 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// RANK-specific: heap fills with K rows tied at the same OB value, + /// then more rows at that same value arrive. They take the Equal arm + /// (heap is full, `heap.max() == row`) and accumulate as ties, while + /// strictly-greater rows are dropped. All retained rows have rank 1. + #[tokio::test] + async fn test_partitioned_topk_rank_boundary_ties_retained() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 5, 5, 10, 5 + // - first two 5s fill the heap (max=None until heap reaches K=2) + // - third row 10 > 5 → drop (Greater) + // - fourth row 5 == 5 → push to ties (Equal) + // Sorted RANKs: 5→1, 5→1, 5→1, 10→4. WHERE rk ≤ 2 keeps the three 5s. + let batch = pk_val_batch(&schema, vec![1, 1, 1, 1], vec![5, 5, 10, 5])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 5 |", + "| 1 | 5 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// RANK-specific: heap fills with K rows tied at value V, equal_indices + /// accumulate at V, then a strictly-better row arrives whose admission + /// shifts the boundary strictly below V. The boundary-changed branch + /// must clear both `state.ties` and the in-flight `equal_indices` — + /// otherwise the now-rank-> K rows at value V would leak into output. + #[tokio::test] + async fn test_partitioned_topk_rank_boundary_shifts_clears_ties() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 10, 10, 10, 5, 3 + // - first two 10s fill heap (max=10) + // - third 10 → Equal → equal_indices=[2] + // - 5 < 10 → admit, evict 10 → heap={5,10}, max=10 (unchanged). + // Push evicted to ties: ties=[10@curr_batch[ev_idx]]. + // - 3 < 10 → admit, evict 10 → heap={3,5}, max=5 (CHANGED). + // Clear ties AND equal_indices. + // Sorted RANKs: 3→1, 5→2, 10→3, 10→3, 10→3. WHERE rk ≤ 2 → [3, 5]. + let batch = pk_val_batch(&schema, vec![1, 1, 1, 1, 1], vec![10, 10, 10, 5, 3])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 3 |", + "| 1 | 5 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } + + /// RANK-specific: heap has multiple rows at boundary value V, then a + /// strictly-better row arrives. The heap evicts one V (popping + /// `prev_min`), but `heap.max()` is still V — boundary unchanged. + /// The evicted V row must be pushed as a `TieEntry`; without that + /// branch a `rk <= K` query would silently lose a tied row. + #[tokio::test] + async fn test_partitioned_topk_rank_eviction_at_unchanged_boundary() -> Result<()> { + let (schema, mut state) = build_partitioned_topk_rank(2)?; + + // pk=1 vals: 10, 10, 5 + // - first two 10s fill the heap (max=10) + // - 5 < 10 → admit, evict 10. New heap={5,10}, max=10 (unchanged). + // Push the evicted 10 to ties. + // Sorted RANKs: 5→1, 10→2, 10→2. WHERE rk ≤ 2 → all 3 rows. + let batch = pk_val_batch(&schema, vec![1, 1, 1], vec![10, 10, 5])?; + state.insert_batch(&batch)?; + + let results: Vec<_> = state.emit()?.try_collect().await?; + assert_batches_eq!( + &[ + "+----+-----+", + "| pk | val |", + "+----+-----+", + "| 1 | 5 |", + "| 1 | 10 |", + "| 1 | 10 |", + "+----+-----+", + ], + &results + ); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/tree_node.rs b/datafusion/physical-plan/src/tree_node.rs index aa4f144f91898..dcdceff8693e3 100644 --- a/datafusion/physical-plan/src/tree_node.rs +++ b/datafusion/physical-plan/src/tree_node.rs @@ -20,7 +20,8 @@ use std::fmt::{self, Display, Formatter}; use std::sync::Arc; -use crate::{ExecutionPlan, displayable, with_new_children_if_necessary}; +use crate::execution_plan::replace_children_if_necessary; +use crate::{ExecutionPlan, displayable}; use datafusion_common::Result; use datafusion_common::tree_node::{ConcreteTreeNode, DynTreeNode}; @@ -35,7 +36,7 @@ impl DynTreeNode for dyn ExecutionPlan { arc_self: Arc, new_children: Vec>, ) -> Result> { - with_new_children_if_necessary(arc_self, new_children) + replace_children_if_necessary(arc_self, new_children) } } @@ -73,7 +74,7 @@ impl PlanContext { /// if the `PlanContext.children` have been changed. pub fn update_plan_from_children(mut self) -> Result { let children_plans = self.children.iter().map(|c| Arc::clone(&c.plan)).collect(); - self.plan = with_new_children_if_necessary(self.plan, children_plans)?; + self.plan = replace_children_if_necessary(self.plan, children_plans)?; Ok(self) } diff --git a/datafusion/physical-plan/src/union.rs b/datafusion/physical-plan/src/union.rs index ec9ea376e0b6d..c1cc5da31abaf 100644 --- a/datafusion/physical-plan/src/union.rs +++ b/datafusion/physical-plan/src/union.rs @@ -31,7 +31,6 @@ use super::{ PlanProperties, RecordBatchStream, SendableRecordBatchStream, Statistics, metrics::{ExecutionPlanMetricsSet, MetricsSet}, }; -use crate::check_if_same_properties; use crate::execution_plan::{ CardinalityEffect, InvariantLevel, boundedness_from_children, check_default_invariants, emission_type_from_children, @@ -42,8 +41,10 @@ use crate::filter_pushdown::{ FilterPushdownPropagation, PushedDown, }; use crate::metrics::BaselineMetrics; -use crate::projection::{ProjectionExec, make_with_child}; +use crate::projection::{ProjectionExec, ProjectionExpr, make_with_child}; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::ObservedStream; +use crate::{ChildrenPropertiesMode, ReplaceChildrenOptions, validate_child_count}; use arrow::datatypes::{Field, Schema, SchemaRef}; use arrow::record_batch::RecordBatch; @@ -51,9 +52,10 @@ use datafusion_common::config::ConfigOptions; use datafusion_common::stats::NdvFallback; use datafusion_common::tree_node::TreeNodeRecursion; use datafusion_common::{ - Result, assert_or_internal_err, exec_err, internal_datafusion_err, + Result, assert_or_internal_err, exec_err, internal_datafusion_err, plan_err, }; use datafusion_execution::TaskContext; +use datafusion_physical_expr::expressions::{CastExpr, Column}; use datafusion_physical_expr::{ EquivalenceProperties, PhysicalExpr, calculate_union, conjunction, }; @@ -63,6 +65,71 @@ use itertools::Itertools; use log::{debug, trace, warn}; use tokio::macros::support::thread_rng_n; +/// Coerces `input`'s output schema to exactly `schema` via a `ProjectionExec` +/// that re-stamps each column with the union's merged field (same +/// `DataType`, but the union's merged nullability/name/metadata), or returns +/// `input` unchanged if its schema already matches. [`UnionExec::try_new`] +/// and [`InterleaveExec::try_new`] call this on every child, so the coercion +/// is visible in the plan tree (e.g. in `EXPLAIN`) instead of happening +/// invisibly inside the union operator's own `execute()`. +/// +/// A column whose `DataType` doesn't already match the union's is a genuine +/// data type mismatch (as opposed to a nullability/name/metadata-only one), +/// and is rejected eagerly here rather than silently cast or deferred to a +/// runtime failure -- this only ever changes a column's declared schema, +/// never its values. +/// +/// Casting a column to its own `DataType` (only the `Field`'s nullability, +/// name, or metadata changes) is a zero-copy relabeling: the cast kernel's +/// same-type fast path (`cast_array_by_name`) just clones the `Arc`, so this carries no runtime overhead over the schema it replaces. +/// +/// See . +fn coerce_schema( + input: Arc, + schema: &SchemaRef, +) -> Result> { + let input_schema = input.schema(); + if &input_schema == schema { + return Ok(input); + } + + let exprs = input_schema + .fields() + .iter() + .zip(schema.fields()) + .enumerate() + .map(|(i, (input_field, target_field))| { + if input_field.data_type() != target_field.data_type() { + return plan_err!( + "UnionExec/InterleaveExec requires all inputs to have the same \ + data type per column; column {i} has type {} in one input, but \ + the union schema expects {}", + input_field.data_type(), + target_field.data_type() + ); + } + let column: Arc = + Arc::new(Column::new(input_field.name(), i)); + let expr = if input_field == target_field { + column + } else { + Arc::new(CastExpr::new_with_target_field( + column, + Arc::clone(target_field), + None, + )) as Arc + }; + Ok(ProjectionExpr { + expr, + alias: target_field.name().clone(), + }) + }) + .collect::>>()?; + + Ok(Arc::new(ProjectionExec::try_new(exprs, input)?)) +} + /// `UnionExec`: `UNION ALL` execution plan. /// /// `UnionExec` combines multiple inputs with the same schema by @@ -111,24 +178,6 @@ pub struct UnionExec { } impl UnionExec { - /// Create a new UnionExec - #[deprecated(since = "44.0.0", note = "Use UnionExec::try_new instead")] - pub fn new(inputs: Vec>) -> Self { - let schema = - union_schema(&inputs).expect("UnionExec::new called with empty inputs"); - // The schema of the inputs and the union schema is consistent when: - // - They have the same number of fields, and - // - Their fields have same types at the same indices. - // Here, we know that schemas are consistent and the call below can - // not return an error. - let cache = Self::compute_properties(&inputs, schema).unwrap(); - UnionExec { - inputs, - metrics: ExecutionPlanMetricsSet::new(), - cache: Arc::new(cache), - } - } - /// Try to create a new UnionExec. /// /// # Errors @@ -148,9 +197,11 @@ impl UnionExec { // The schema of the inputs and the union schema is consistent when: // - They have the same number of fields, and // - Their fields have same types at the same indices. - // Here, we know that schemas are consistent and the call below can - // not return an error. - let cache = Self::compute_properties(&inputs, schema).unwrap(); + let inputs = inputs + .into_iter() + .map(|input| coerce_schema(input, &schema)) + .collect::>>()?; + let cache = Self::compute_properties(&inputs, schema)?; Ok(Arc::new(UnionExec { inputs, metrics: ExecutionPlanMetricsSet::new(), @@ -165,6 +216,20 @@ impl UnionExec { &self.inputs } + /// Maps a global output partition index to the `(input index, local + /// partition index)` of the input that owns it, or `None` if out of range. + fn owning_input(&self, partition: usize) -> Option<(usize, usize)> { + let mut remaining = partition; + for (i, input) in self.inputs.iter().enumerate() { + let count = input.output_partitioning().partition_count(); + if remaining < count { + return Some((i, remaining)); + } + remaining -= count; + } + None + } + /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. fn compute_properties( inputs: &[Arc], @@ -190,17 +255,6 @@ impl UnionExec { boundedness_from_children(inputs), )) } - - fn with_new_children_and_same_properties( - &self, - children: Vec>, - ) -> Self { - Self { - inputs: children, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for UnionExec { @@ -271,17 +325,45 @@ impl ExecutionPlan for UnionExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => UnionExec::try_new(children), + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - UnionExec::try_new(children) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -326,26 +408,46 @@ impl ExecutionPlan for UnionExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { + fn child_stats_requests(&self, partition: Option) -> Vec { if let Some(partition_idx) = partition { + // For a specific partition, compute stats only for the input that + // owns it; the other inputs are not needed and are skipped. + let targeted = self.owning_input(partition_idx); + self.inputs + .iter() + .enumerate() + .map(|(i, _)| match targeted { + Some((target_i, target_partition)) if i == target_i => { + ChildStats::At(Some(target_partition)) + } + _ => ChildStats::Skip, + }) + .collect() + } else { + vec![ChildStats::At(None); self.inputs.len()] + } + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + args: &StatisticsArgs, + ) -> Result> { + if let Some(partition_idx) = args.partition() { // For a specific partition, find which input it belongs to - let mut remaining_idx = partition_idx; - for input in &self.inputs { - let input_partition_count = input.output_partitioning().partition_count(); - if remaining_idx < input_partition_count { - // This partition belongs to this input - return input.partition_statistics(Some(remaining_idx)); - } - remaining_idx -= input_partition_count; + if let Some((target_i, _)) = self.owning_input(partition_idx) { + // This partition belongs to this input - return its stats + return Ok(Arc::clone(&input_stats[target_i])); } // If we get here, the partition index is out of bounds Ok(Arc::new(Statistics::new_unknown(&self.schema()))) } else { - let schema = self.schema(); - Ok(Arc::new(merge_input_statistics( - &self.inputs, - None, - schema.as_ref(), + let stats_refs = input_stats.iter().map(|s| s.as_ref()).collect::>(); + + Ok(Arc::new(Statistics::try_merge_iter_with_ndv_fallback( + stats_refs, + self.schema().as_ref(), + NdvFallback::Sum, )?)) } } @@ -466,11 +568,50 @@ impl ExecutionPlan for UnionExec { // on all children (either pushed down or via FilterExec) Ok(propagation) } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let inputs = ctx.encode_children(self.inputs())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Union( + protobuf::UnionExecNode { inputs }, + ), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl UnionExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let union = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Union, + "UnionExec", + ); + let inputs = union + .inputs + .iter() + .map(|input| ctx.decode_child(input)) + .collect::>>()?; + UnionExec::try_new(inputs) + } } /// Combines multiple input streams by interleaving them. /// -/// This only works if all inputs have the same hash-partitioning. +/// All inputs must share an identical [`Partitioning::Hash`] or [`Partitioning::Range`] so that +/// partition `k` covers the same data across every input. Each output partition is the +/// interleaving of the same-indexed partition from all inputs: +/// `output[k] = input[0][k] + input[1][k] + ... + input[n-1][k]` /// /// # Data Flow /// ```text @@ -515,9 +656,14 @@ impl InterleaveExec { pub fn try_new(inputs: Vec>) -> Result { assert_or_internal_err!( can_interleave(inputs.iter()), - "Not all InterleaveExec children have a consistent hash partitioning" + "Not all InterleaveExec children have a consistent hash or range partitioning" ); - let cache = Self::compute_properties(&inputs)?; + let schema = union_schema(&inputs)?; + let inputs = inputs + .into_iter() + .map(|input| coerce_schema(input, &schema)) + .collect::>>()?; + let cache = Self::compute_properties(&inputs, schema)?; Ok(InterleaveExec { inputs, metrics: ExecutionPlanMetricsSet::new(), @@ -531,8 +677,10 @@ impl InterleaveExec { } /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc. - fn compute_properties(inputs: &[Arc]) -> Result { - let schema = union_schema(inputs)?; + fn compute_properties( + inputs: &[Arc], + schema: SchemaRef, + ) -> Result { let eq_properties = EquivalenceProperties::new(schema); // Get output partitioning: let output_partitioning = inputs[0].output_partitioning().clone(); @@ -543,17 +691,6 @@ impl InterleaveExec { boundedness_from_children(inputs), )) } - - fn with_new_children_and_same_properties( - &self, - children: Vec>, - ) -> Self { - Self { - inputs: children, - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for InterleaveExec { @@ -591,22 +728,52 @@ impl ExecutionPlan for InterleaveExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } + fn replace_children( + self: Arc, + children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + inputs: children, + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + // New children are no longer interleavable, which might be a bug of optimization rewrite. + assert_or_internal_err!( + can_interleave(children.iter()), + "Can not create InterleaveExec: new children can not be interleaved" + ); + Ok(Arc::new(InterleaveExec::try_new(children)?)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - // New children are no longer interleavable, which might be a bug of optimization rewrite. - assert_or_internal_err!( - can_interleave(children.iter()), - "Can not create InterleaveExec: new children can not be interleaved" - ); - check_if_same_properties!(self, children); - Ok(Arc::new(InterleaveExec::try_new(children)?)) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -629,7 +796,8 @@ impl ExecutionPlan for InterleaveExec { let mut input_stream_vec = vec![]; for input in self.inputs.iter() { if partition < input.output_partitioning().partition_count() { - input_stream_vec.push(input.execute(partition, Arc::clone(&context))?); + let stream = input.execute(partition, Arc::clone(&context))?; + input_stream_vec.push(stream); } else { // Do not find a partition to execute break; @@ -656,22 +824,74 @@ impl ExecutionPlan for InterleaveExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let schema = self.schema(); - Ok(Arc::new(merge_input_statistics( - &self.inputs, - partition, - schema.as_ref(), + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition); self.inputs.len()] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let stats = input_stats + .iter() + .map(|s| s.as_ref().clone()) + .collect::>(); + + Ok(Arc::new(Statistics::try_merge_iter_with_ndv_fallback( + stats.iter(), + self.schema().as_ref(), + NdvFallback::Sum, )?)) } fn benefits_from_input_partitioning(&self) -> Vec { vec![false; self.children().len()] } + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let inputs = ctx.encode_children(self.inputs())?; + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Interleave( + protobuf::InterleaveExecNode { inputs }, + ), + ), + })) + } } -/// If all the input partitions have the same Hash partition spec with the first_input_partition -/// The InterleaveExec is partition aware. +#[cfg(feature = "proto")] +impl InterleaveExec { + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + let interleave = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Interleave, + "InterleaveExec", + ); + let inputs = interleave + .inputs + .iter() + .map(|input| ctx.decode_child(input)) + .collect::>>()?; + Ok(Arc::new(InterleaveExec::try_new(inputs)?)) + } +} + +/// Returns true if all inputs have the same [`Partitioning::Hash`] or [`Partitioning::Range`] +/// spec, making them safe to interleave. Two inputs are interleave-compatible when partition +/// `k` covers the identical key range or hash bucket across every input. +/// +/// Note: compatibility is checked sequentially against the first input, so +/// `InputDistributionRequirements::co_partitioned` is not needed here. /// /// It might be too strict here in the case that the input partition specs are compatible but not exactly the same. /// For example one input partition has the partition spec Hash('a','b','c') and @@ -684,7 +904,7 @@ pub fn can_interleave>>( }; let reference = first.borrow().output_partitioning(); - matches!(reference, Partitioning::Hash(_, _)) + matches!(reference, Partitioning::Hash(_, _) | Partitioning::Range(_)) && inputs .map(|plan| plan.borrow().output_partitioning().clone()) .all(|partition| partition == *reference) @@ -821,37 +1041,24 @@ impl Stream for CombinedRecordBatchStream { } } -fn merge_input_statistics( - inputs: &[Arc], - partition: Option, - schema: &Schema, -) -> Result { - let stats = inputs - .iter() - .map(|input| { - input - .partition_statistics(partition) - .map(Arc::unwrap_or_clone) - }) - .collect::>>()?; - - Statistics::try_merge_iter_with_ndv_fallback(stats.iter(), schema, NdvFallback::Sum) -} - #[cfg(test)] mod tests { use super::*; use crate::collect; use crate::repartition::RepartitionExec; + use crate::statistics::{StatisticsArgs, StatisticsContext}; use crate::test::exec::StatisticsExec; use crate::test::{self, TestMemoryExec}; use arrow::compute::SortOptions; use arrow::datatypes::DataType; + use datafusion_common::SplitPoint; use datafusion_common::stats::Precision; use datafusion_common::{ColumnStatistics, ScalarValue}; + use datafusion_physical_expr::RangePartitioning; use datafusion_physical_expr::equivalence::convert_to_orderings; use datafusion_physical_expr::expressions::col; + use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr}; // Generate a schema which consists of 7 columns (a, b, c, d, e, f, g) fn create_test_schema() -> Result { @@ -904,6 +1111,52 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_interleave_conforms_batch_schema() -> Result<()> { + // Two inputs agree on the column's type but disagree on nullability; + // InterleaveExec's declared schema ORs nullability across inputs, so + // every yielded batch must be re-stamped with that schema. See + // . + let task_ctx = Arc::new(TaskContext::default()); + + let schema_not_null = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch_not_null = RecordBatch::try_new( + Arc::clone(&schema_not_null), + vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2]))], + )?; + + let schema_nullable = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let batch_nullable = RecordBatch::try_new( + Arc::clone(&schema_nullable), + vec![Arc::new(arrow::array::Int32Array::from(vec![3, 4]))], + )?; + + let hash_expr = vec![col("a", schema_not_null.as_ref())?]; + let left: Arc = Arc::new(RepartitionExec::try_new( + TestMemoryExec::try_new_exec(&[vec![batch_not_null]], schema_not_null, None)?, + Partitioning::Hash(hash_expr.clone(), 1), + )?); + let right: Arc = Arc::new(RepartitionExec::try_new( + TestMemoryExec::try_new_exec(&[vec![batch_nullable]], schema_nullable, None)?, + Partitioning::Hash(hash_expr, 1), + )?); + + let interleave: Arc = + Arc::new(InterleaveExec::try_new(vec![left, right])?); + let interleave_schema = interleave.schema(); + assert!(interleave_schema.field(0).is_nullable()); + + let batches = collect(interleave, task_ctx).await?; + assert!(!batches.is_empty()); + for batch in &batches { + assert_eq!(batch.schema(), interleave_schema); + } + + Ok(()) + } + fn stats_merge_inputs() -> (SchemaRef, Statistics, Statistics, Statistics) { let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::UInt32, true)])); @@ -1033,7 +1286,8 @@ mod tests { Arc::new(StatisticsExec::new(right, schema.as_ref().clone())); let union = UnionExec::try_new(vec![left, right])?; - let stats = union.partition_statistics(None)?; + let stats = + StatisticsContext::new().compute(union.as_ref(), &StatisticsArgs::new())?; assert_eq!(stats.as_ref(), &expected); Ok(()) @@ -1050,12 +1304,114 @@ mod tests { Arc::new(StatisticsExec::new(right, schema.as_ref().clone())); let union = UnionExec::try_new(vec![left, right])?; - let stats = union.partition_statistics(None)?; + let stats = + StatisticsContext::new().compute(union.as_ref(), &StatisticsArgs::new())?; + + assert_eq!(stats.as_ref(), &expected); + Ok(()) + } + + #[test] + fn test_union_partition_statistics_with_mismatched_nullability() -> Result<()> { + // Regression test for the `ProjectionExec` wrapper `UnionExec::try_new` + // inserts above the non-nullable leg here (via `coerce_schema`): + // exact column statistics (min/max/null/distinct/sum/byte_size) must + // still make it through the wrapper's same-type `CastExpr`, not get + // poisoned into `Absent` the way a generic (type-changing) cast's + // statistics would be. + let (_, left, right, expected) = stats_merge_inputs(); + + // `total_byte_size` differs from the plain-merge fixture (52): the + // wrapper is a `ProjectionExec`, whose `statistics_from_inputs` + // recomputes `total_byte_size` from the (unchanged) schema's row + // width times row count, rather than trusting the wrapped leg's own + // self-reported total -- still `Exact`, just derived differently. + // left: 5 rows * 4 bytes (UInt32) = 20 (was 23); right is untouched + // (already nullable, so `coerce_schema` doesn't wrap it): 20 + 29 = 49. + let expected = expected.with_total_byte_size(Precision::Exact(49)); + + let non_nullable_schema = + Schema::new(vec![Field::new("a", DataType::UInt32, false)]); + let nullable_schema = Schema::new(vec![Field::new("a", DataType::UInt32, true)]); + + let left: Arc = + Arc::new(StatisticsExec::new(left, non_nullable_schema)); + let right: Arc = + Arc::new(StatisticsExec::new(right, nullable_schema)); + + let union = UnionExec::try_new(vec![left, right])?; + let stats = + StatisticsContext::new().compute(union.as_ref(), &StatisticsArgs::new())?; assert_eq!(stats.as_ref(), &expected); Ok(()) } + #[tokio::test] + async fn test_coerce_schema_no_op_when_already_matching() -> Result<()> { + let schema_not_null = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input: Arc = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema_not_null), None)?; + + let coerced = coerce_schema(Arc::clone(&input), &schema_not_null)?; + assert!(Arc::ptr_eq(&coerced, &input)); + + Ok(()) + } + + #[tokio::test] + async fn test_coerce_schema_casts_only_nullability() -> Result<()> { + // Mismatched nullability: the input gets wrapped in a `ProjectionExec` + // whose `CastExpr` re-stamps the column with the target's `Field` + // (same `DataType`, so this is a zero-copy relabeling, not a real cast). + let schema_not_null = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let batch_not_null = RecordBatch::try_new( + Arc::clone(&schema_not_null), + vec![Arc::new(arrow::array::Int32Array::from(vec![1, 2]))], + )?; + let input: Arc = TestMemoryExec::try_new_exec( + &[vec![batch_not_null]], + Arc::clone(&schema_not_null), + None, + )?; + + let nullable_schema = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let coerced = coerce_schema(Arc::clone(&input), &nullable_schema)?; + assert_eq!(&coerced.schema(), &nullable_schema); + let plan_str = crate::displayable(coerced.as_ref()) + .indent(true) + .to_string(); + assert!( + plan_str.contains("CAST"), + "expected a CAST in the coerced plan:\n{plan_str}" + ); + + let task_ctx = Arc::new(TaskContext::default()); + let batches = collect(coerced, task_ctx).await?; + assert_eq!(batches.len(), 1); + assert_eq!(batches[0].schema(), nullable_schema); + + Ok(()) + } + + #[test] + fn test_coerce_schema_rejects_genuine_type_mismatch() -> Result<()> { + let schema_int = + Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, false)])); + let input: Arc = + TestMemoryExec::try_new_exec(&[vec![]], Arc::clone(&schema_int), None)?; + + let schema_utf8 = + Arc::new(Schema::new(vec![Field::new("a", DataType::Utf8, false)])); + let err = coerce_schema(input, &schema_utf8).unwrap_err(); + assert!(err.to_string().contains("same data type per column")); + + Ok(()) + } + #[test] fn test_interleave_partition_statistics_uses_shared_statistics_merge() -> Result<()> { let (schema, left, right, expected) = stats_merge_inputs(); @@ -1071,7 +1427,8 @@ mod tests { )?); let interleave = InterleaveExec::try_new(vec![left, right])?; - let stats = interleave.partition_statistics(None)?; + let stats = + StatisticsContext::new().compute(&interleave, &StatisticsArgs::new())?; assert_eq!(stats.as_ref(), &expected); Ok(()) @@ -1093,7 +1450,8 @@ mod tests { )?); let interleave = InterleaveExec::try_new(vec![left, right])?; - let stats = interleave.partition_statistics(Some(0))?; + let stats = StatisticsContext::new() + .compute(&interleave, &StatisticsArgs::new().with_partition(Some(0)))?; let expected = Statistics::default() .with_num_rows(Precision::Inexact(5)) @@ -1288,6 +1646,124 @@ mod tests { ); } + fn make_hash_exec( + schema: &SchemaRef, + hash_cols: Vec<&str>, + buckets: usize, + ) -> Result> { + let exprs = hash_cols + .iter() + .map(|c| col(c, schema)) + .collect::>>()?; + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Hash(exprs, buckets), + )?)) + } + + fn make_range_exec( + schema: &SchemaRef, + split_values: Vec, + sort_options: SortOptions, + ) -> Result> { + let sort_expr = + PhysicalSortExpr::new(col(schema.field(0).name(), schema)?, sort_options); + let ordering = LexOrdering::new(vec![sort_expr]).unwrap(); + let split_points = split_values + .into_iter() + .map(|v| SplitPoint::new(vec![ScalarValue::Int32(Some(v))])) + .collect(); + let base = Arc::new(TestMemoryExec::try_new(&[], Arc::clone(schema), None)?); + Ok(Arc::new(RepartitionExec::try_new( + base, + Partitioning::Range(RangePartitioning::try_new(ordering, split_points)?), + )?)) + } + + #[test] + fn test_can_interleave_matrix() -> Result<()> { + let name_column = "name"; + let age_column = "age"; + let schema = Arc::new(Schema::new(vec![ + Field::new(name_column, DataType::Int32, true), + Field::new(age_column, DataType::Int32, true), + ])); + + let ascending = SortOptions { + descending: false, + nulls_first: false, + }; + struct Case { + inputs: Vec>, + expected: bool, + label: &'static str, + } + + let cases = vec![ + // compatible + Case { + label: "matching hash on single column", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column], 3)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + Case { + label: "matching hash on multiple columns", + expected: true, + inputs: vec![ + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + make_hash_exec(&schema, vec![name_column, age_column], 3)?, + ], + }, + Case { + label: "matching range same splits and order", + expected: true, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 20], ascending)?, + ], + }, + // incompatible + Case { + label: "subset range partition", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 15], ascending)?, + ], + }, + Case { + label: "range different split points", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_range_exec(&schema, vec![10, 30], ascending)?, + ], + }, + Case { + label: "mixed range and hash", + expected: false, + inputs: vec![ + make_range_exec(&schema, vec![10, 20], ascending)?, + make_hash_exec(&schema, vec![name_column], 3)?, + ], + }, + ]; + + for case in cases { + assert_eq!( + can_interleave(case.inputs.iter()), + case.expected, + "{}", + case.label + ); + } + Ok(()) + } + #[test] fn test_union_cardinality_effect() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/physical-plan/src/unnest.rs b/datafusion/physical-plan/src/unnest.rs index 3a4b9d7232f4d..3fa274b27a7bd 100644 --- a/datafusion/physical-plan/src/unnest.rs +++ b/datafusion/physical-plan/src/unnest.rs @@ -28,8 +28,9 @@ use super::metrics::{ use super::{DisplayAs, ExecutionPlanProperties, PlanProperties}; use crate::stream::EmptyRecordBatchStream; use crate::{ - DisplayFormatType, Distribution, ExecutionPlan, RecordBatchStream, - SendableRecordBatchStream, check_if_same_properties, + ChildrenPropertiesMode, DisplayFormatType, Distribution, ExecutionPlan, + RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + validate_child_count, }; use arrow::array::{ @@ -196,17 +197,6 @@ impl UnnestExec { pub fn options(&self) -> &UnnestOptions { &self.options } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for UnnestExec { @@ -241,27 +231,61 @@ impl ExecutionPlan for UnnestExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(UnnestExec::new( - children.swap_remove(0), - self.list_column_indices.clone(), - self.struct_column_indices.clone(), - Arc::clone(&self.schema), - self.options.clone(), - )?)) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(UnnestExec::new( + children.swap_remove(0), + self.list_column_indices.clone(), + self.struct_column_indices.clone(), + Arc::clone(&self.schema), + self.options.clone(), + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn required_input_distribution(&self) -> Vec { - vec![Distribution::UnspecifiedDistribution] + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> crate::InputDistributionRequirements { + crate::InputDistributionRequirements::new(vec![ + Distribution::UnspecifiedDistribution, + ]) } fn execute( @@ -285,6 +309,170 @@ impl ExecutionPlan for UnnestExec { fn metrics(&self) -> Option { Some(self.metrics.clone_inner()) } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `UnnestExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + input, + schema, + list_column_indices, + struct_column_indices, + options, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + // Derived at construction by `UnnestExec::compute_properties`. + cache: _, + } = self; + + let input = ctx.encode_child(input)?; + let schema = schema.as_ref().try_into()?; + let list_type_columns = list_column_indices + .iter() + .map(|column| protobuf::ListUnnest { + index_in_input_schema: column.index_in_input_schema as _, + depth: column.depth as _, + }) + .collect(); + let struct_type_columns = struct_column_indices + .iter() + .map(|index| *index as _) + .collect(); + let null_handling = { + use datafusion_common::NullHandling; + use protobuf::unnest_options::NullHandling as ProtoNullHandling; + match options.null_handling { + NullHandling::Preserve => ProtoNullHandling::Preserve, + NullHandling::Drop => ProtoNullHandling::Drop, + NullHandling::PreserveAndExpandEmpty => { + ProtoNullHandling::PreserveAndExpandEmpty + } + } + } as i32; + let options = protobuf::UnnestOptions { + null_handling, + recursions: options + .recursions + .iter() + .map(|recursion| protobuf::RecursionUnnestOption { + input_column: Some((&recursion.input_column).into()), + output_column: Some((&recursion.output_column).into()), + depth: recursion.depth as _, + }) + .collect(), + }; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Unnest(Box::new( + protobuf::UnnestExecNode { + input: Some(Box::new(input)), + schema: Some(schema), + list_type_columns, + struct_type_columns, + options: Some(options), + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl UnnestExec { + /// Reconstruct an [`UnnestExec`] from its protobuf representation. + /// + /// The exact inverse of [`ExecutionPlan::try_to_proto`]. + /// + /// [`ExecutionPlan::try_to_proto`]: crate::ExecutionPlan::try_to_proto + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + let unnest = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Unnest, + "UnnestExec", + ); + // Exhaustive destructure: a new field on `UnnestExecNode` is a compile + // error here rather than a silently ignored wire field. + let protobuf::UnnestExecNode { + input, + schema, + list_type_columns, + struct_type_columns, + options, + } = unnest.as_ref(); + + let input = ctx.decode_required_child(input.as_deref(), "UnnestExec", "input")?; + let schema: Schema = schema + .as_ref() + .ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "UnnestExec is missing required field 'schema'" + ) + })? + .try_into()?; + let list_column_indices = list_type_columns + .iter() + .map(|column| ListUnnest { + index_in_input_schema: column.index_in_input_schema as _, + depth: column.depth as _, + }) + .collect(); + let struct_column_indices = struct_type_columns + .iter() + .map(|index| *index as _) + .collect(); + let options = options.as_ref().ok_or_else(|| { + datafusion_common::internal_datafusion_err!( + "UnnestExec is missing required field 'options'" + ) + })?; + let null_handling = { + use datafusion_common::NullHandling; + use protobuf::unnest_options::NullHandling as ProtoNullHandling; + match ProtoNullHandling::try_from(options.null_handling) { + Ok(ProtoNullHandling::Preserve) => NullHandling::Preserve, + Ok(ProtoNullHandling::Drop) => NullHandling::Drop, + Ok(ProtoNullHandling::PreserveAndExpandEmpty) => { + NullHandling::PreserveAndExpandEmpty + } + // Unknown enum values fall back to the default (Preserve), + // matching DataFusion's historical behavior. + Err(_) => NullHandling::Preserve, + } + }; + let options = UnnestOptions { + null_handling, + recursions: options + .recursions + .iter() + .map(|recursion| datafusion_common::RecursionUnnestOption { + input_column: recursion.input_column.as_ref().unwrap().into(), + output_column: recursion.output_column.as_ref().unwrap().into(), + depth: recursion.depth as _, + }) + .collect(), + }; + + Ok(Arc::new(UnnestExec::new( + input, + list_column_indices, + struct_column_indices, + Arc::new(schema), + options, + )?)) + } } #[derive(Clone, Debug)] @@ -769,14 +957,21 @@ fn build_batch( /// l2: [4,5], [], null, [6, 7] /// ``` /// -/// If `preserve_nulls` is false, the longest length array will be: +/// With [`datafusion_common::NullHandling::Drop`], the longest length array will be: /// /// ```ignore /// longest_length: [3, 0, 0, 2] /// ``` /// -/// whereas if `preserve_nulls` is true, the longest length array will be: +/// With [`datafusion_common::NullHandling::Preserve`] (the default), the longest length array +/// will be: /// +/// ```ignore +/// longest_length: [3, 1, 1, 2] +/// ``` +/// +/// With [`datafusion_common::NullHandling::PreserveAndExpandEmpty`], empty input lists are +/// also bumped to length 1 so they produce a single `NULL` output row: /// /// ```ignore /// longest_length: [3, 1, 1, 2] @@ -785,12 +980,16 @@ fn find_longest_length( list_arrays: &[ArrayRef], options: &UnnestOptions, ) -> Result { - // The length of a NULL list - let null_length = if options.preserve_nulls { + // The length to substitute for a NULL input list. + let null_length = if options.preserve_nulls() { Scalar::new(Int64Array::from_value(1, 1)) } else { Scalar::new(Int64Array::from_value(0, 1)) }; + let expand_empty = options.expand_empty_as_null(); + // Reused scalars for the empty-list rewrite when expand_empty is set. + let zero = Scalar::new(Int64Array::from_value(0, 1)); + let one = Scalar::new(Int64Array::from_value(1, 1)); let list_lengths: Vec = list_arrays .iter() .map(|list_array| { @@ -799,6 +998,12 @@ fn find_longest_length( length_array = cast(&length_array, &DataType::Int64)?; length_array = zip(&is_not_null(&length_array)?, &length_array, &null_length)?; + if expand_empty { + // Bump empty lists (length 0) to length 1 so they + // produce a single output row padded with NULL. + let is_zero = arrow_ord::cmp::eq(&length_array, &zero)?; + length_array = zip(&is_zero, &one, &length_array)?; + } Ok(length_array) }) .collect::>()?; @@ -1068,6 +1273,7 @@ mod tests { }; use arrow::buffer::{NullBuffer, OffsetBuffer}; use arrow::datatypes::{Field, Int32Type}; + use datafusion_common::NullHandling; use datafusion_common::test_util::batches_to_string; use insta::assert_snapshot; @@ -1250,12 +1456,375 @@ mod tests { list_type_columns.as_ref(), &HashSet::default(), &UnnestOptions { - preserve_nulls: true, + null_handling: NullHandling::Preserve, + recursions: vec![], + }, + )? + .unwrap(); + + assert_snapshot!(batches_to_string(&[ret]), + @r" + +---------------------------------+---------------------------------+---------------------------------+ + | col1_unnest_placeholder_depth_1 | col1_unnest_placeholder_depth_2 | col2_unnest_placeholder_depth_1 | + +---------------------------------+---------------------------------+---------------------------------+ + | [1, 2, 3] | 1 | a | + | | 2 | b | + | [4, 5] | 3 | | + | [1, 2, 3] | | a | + | | | b | + | [4, 5] | | | + | [1, 2, 3] | 4 | a | + | | 5 | b | + | [4, 5] | | | + | [7, 8, 9, 10] | 7 | c | + | | 8 | d | + | [11, 12, 13] | 9 | | + | | 10 | | + | [7, 8, 9, 10] | | c | + | | | d | + | [11, 12, 13] | | | + | [7, 8, 9, 10] | 11 | c | + | | 12 | d | + | [11, 12, 13] | 13 | | + | | | e | + +---------------------------------+---------------------------------+---------------------------------+ + "); + Ok(()) + } + + #[test] + fn test_build_batch_preserve_and_expand_empty() -> Result<()> { + // c1: [A, B, C], [], NULL, [D], NULL, [NULL, F] c2: 1, 2, 3, 4, 5, 6 + // Expected for `NullHandling::PreserveAndExpandEmpty`: + // [A, B, C] -> three rows with c2 = 1, 1, 1 + // [] -> one row with c2 = 2 and unnested value NULL + // NULL -> one row with c2 = 3 and unnested value NULL + // [D] -> one row with c2 = 4 + // NULL -> one row with c2 = 5 and unnested value NULL + // [NULL, F] -> two rows with c2 = 6, 6 + let list_array = Arc::new(make_generic_array::()) as ArrayRef; + let other = + Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "c1", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + Field::new("c2", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("c1_unnested", DataType::Utf8, true), + Field::new("c2", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![Arc::clone(&list_array), Arc::clone(&other)], + )?; + let list_type_columns = vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + )? + .unwrap(); + + assert_snapshot!(batches_to_string(&[ret]), + @r" + +-------------+----+ + | c1_unnested | c2 | + +-------------+----+ + | A | 1 | + | B | 1 | + | C | 1 | + | | 2 | + | | 3 | + | D | 4 | + | | 5 | + | | 6 | + | F | 6 | + +-------------+----+ + "); + Ok(()) + } + + // PreserveAndExpandEmpty must work for LargeListArray (i64 offsets) too, + // not just the i32-offset ListArray exercised above. + #[test] + fn test_build_batch_preserve_and_expand_empty_largelist() -> Result<()> { + let list_array = Arc::new(make_generic_array::()) as ArrayRef; + let other = + Arc::new(arrow::array::Int32Array::from(vec![1, 2, 3, 4, 5, 6])) as ArrayRef; + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "c1", + DataType::LargeList(Arc::new(Field::new_list_field( + DataType::Utf8, + true, + ))), + true, + ), + Field::new("c2", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("c1_unnested", DataType::Utf8, true), + Field::new("c2", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![Arc::clone(&list_array), Arc::clone(&other)], + )?; + let list_type_columns = vec![ListUnnest { + index_in_input_schema: 0, + depth: 1, + }]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + )? + .unwrap(); + + // Same expected shape as the ListArray case — exercises the LargeList + // code path in unnest_list_array. + assert_snapshot!(batches_to_string(&[ret]), + @r" + +-------------+----+ + | c1_unnested | c2 | + +-------------+----+ + | A | 1 | + | B | 1 | + | C | 1 | + | | 2 | + | | 3 | + | D | 4 | + | | 5 | + | | 6 | + | F | 6 | + +-------------+----+ + "); + Ok(()) + } + + // When two list columns are unnested together, `find_longest_length` + // takes the per-row max. PreserveAndExpandEmpty must bump zeros to ones + // in each input column independently, then the row-wise max picks up + // the right value. + #[test] + fn test_build_batch_preserve_and_expand_empty_multi_column() -> Result<()> { + // col_a: [1, 2], [], NULL, [3] + // col_b: ['x'], ['y'],['z'], NULL + let col_a = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2)]), + Some(vec![]), + None, + Some(vec![Some(3)]), + ]); + let col_b = { + let mut b = + arrow::array::ListBuilder::new(arrow::array::StringBuilder::new()); + b.values().append_value("x"); + b.append(true); + b.values().append_value("y"); + b.append(true); + b.values().append_value("z"); + b.append(true); + b.append(false); + b.finish() + }; + let id = + Arc::new(arrow::array::Int32Array::from(vec![10, 20, 30, 40])) as ArrayRef; + + let in_schema = Arc::new(Schema::new(vec![ + Field::new( + "a", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new( + "b", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ), + Field::new("id", DataType::Int32, true), + ])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new("a_unnested", DataType::Int32, true), + Field::new("b_unnested", DataType::Utf8, true), + Field::new("id", DataType::Int32, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&in_schema), + vec![ + Arc::new(col_a) as ArrayRef, + Arc::new(col_b) as ArrayRef, + Arc::clone(&id), + ], + )?; + let list_type_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, recursions: vec![], }, )? .unwrap(); + // Row 0: longest = max(len([1,2])=2, len(['x'])=1) = 2 → a=[1,2], b=['x',NULL] + // Row 1: a=[] bumped to len 1, b=['y'] len 1 → a=[NULL], b=['y'] + // Row 2: a=NULL bumped to len 1, b=['z'] len 1 → a=[NULL], b=['z'] + // Row 3: a=[3] len 1, b=NULL bumped to len 1 → a=[3], b=[NULL] + assert_snapshot!(batches_to_string(&[ret]), + @r" + +------------+------------+----+ + | a_unnested | b_unnested | id | + +------------+------------+----+ + | 1 | x | 10 | + | 2 | | 10 | + | | y | 20 | + | | z | 30 | + | 3 | | 40 | + +------------+------------+----+ + "); + Ok(()) + } + + // PreserveAndExpandEmpty must propagate through recursive depth-2 + // unnesting: an outer NULL or empty produces one NULL output row at + // each level. Adapted from `test_build_batch_list_arr_recursive`. + #[test] + fn test_build_batch_preserve_and_expand_empty_recursive() -> Result<()> { + // col1 | col2 + // [[1,2,3],null,[4,5]] | ['a','b'] + // [[7,8,9,10], null, [11,12,13]] | ['c','d'] + // null | ['e'] + let list_arr1 = ListArray::from_iter_primitive::(vec![ + Some(vec![Some(1), Some(2), Some(3)]), + None, + Some(vec![Some(4), Some(5)]), + Some(vec![Some(7), Some(8), Some(9), Some(10)]), + None, + Some(vec![Some(11), Some(12), Some(13)]), + ]); + let list_arr1_ref = Arc::new(list_arr1) as ArrayRef; + let offsets = OffsetBuffer::from_lengths([3, 3, 0]); + let mut nulls = NullBufferBuilder::new(3); + nulls.append_non_null(); + nulls.append_non_null(); + nulls.append_null(); + let col1_field = Field::new_list_field( + DataType::List(Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + ))), + true, + ); + let col1 = ListArray::new( + Arc::new(Field::new_list_field( + list_arr1_ref.data_type().to_owned(), + true, + )), + offsets, + list_arr1_ref, + nulls.finish(), + ); + + let list_arr2 = StringArray::from(vec![ + Some("a"), + Some("b"), + Some("c"), + Some("d"), + Some("e"), + ]); + let offsets = OffsetBuffer::from_lengths([2, 2, 1]); + let mut nulls = NullBufferBuilder::new(3); + nulls.append_n_non_nulls(3); + let col2_field = Field::new( + "col2", + DataType::List(Arc::new(Field::new_list_field(DataType::Utf8, true))), + true, + ); + let col2 = GenericListArray::::new( + Arc::new(Field::new_list_field(DataType::Utf8, true)), + OffsetBuffer::new(offsets.into()), + Arc::new(list_arr2), + nulls.finish(), + ); + let schema = Arc::new(Schema::new(vec![col1_field, col2_field])); + let out_schema = Arc::new(Schema::new(vec![ + Field::new( + "col1_unnest_placeholder_depth_1", + DataType::List(Arc::new(Field::new_list_field(DataType::Int32, true))), + true, + ), + Field::new("col1_unnest_placeholder_depth_2", DataType::Int32, true), + Field::new("col2_unnest_placeholder_depth_1", DataType::Utf8, true), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(col1) as ArrayRef, Arc::new(col2) as ArrayRef], + )?; + let list_type_columns = vec![ + ListUnnest { + index_in_input_schema: 0, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 0, + depth: 2, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ]; + + let ret = build_batch( + &batch, + &out_schema, + &list_type_columns, + &HashSet::default(), + &UnnestOptions { + null_handling: NullHandling::PreserveAndExpandEmpty, + recursions: vec![], + }, + )? + .unwrap(); + + // The third input row (col1 = null, col2 = ['e']) now produces a + // NULL row for the depth-1 col1 placeholder *and* the depth-2 one, + // instead of being dropped at depth 1 and again at depth 2 the way + // it would be under `Drop`. Inner NULLs inside [...null...] sub- + // lists are still padded with NULL as before. assert_snapshot!(batches_to_string(&[ret]), @r" +---------------------------------+---------------------------------+---------------------------------+ @@ -1333,11 +1902,11 @@ mod tests { fn verify_longest_length( list_arrays: &[ArrayRef], - preserve_nulls: bool, + null_handling: NullHandling, expected: Vec, ) -> Result<()> { let options = UnnestOptions { - preserve_nulls, + null_handling, recursions: vec![], }; let longest_length = find_longest_length(list_arrays, &options)?; @@ -1357,20 +1926,55 @@ mod tests { // Test with single ListArray // [A, B, C], [], NULL, [D], NULL, [NULL, F] let list_array = Arc::new(make_generic_array::()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![3, 0, 0, 1, 0, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![3, 0, 1, 1, 1, 2])?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![3, 0, 0, 1, 0, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![3, 0, 1, 1, 1, 2], + )?; + // PreserveAndExpandEmpty also treats empty lists as a NULL row. + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 1, 1, 1, 2], + )?; // Test with single LargeListArray // [A, B, C], [], NULL, [D], NULL, [NULL, F] let list_array = Arc::new(make_generic_array::()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![3, 0, 0, 1, 0, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![3, 0, 1, 1, 1, 2])?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![3, 0, 0, 1, 0, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![3, 0, 1, 1, 1, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 1, 1, 1, 2], + )?; // Test with single FixedSizeListArray // [A, B], NULL, [C, D], NULL, [NULL, F], [NULL, NULL] let list_array = Arc::new(make_fixed_list()) as ArrayRef; - verify_longest_length(&[Arc::clone(&list_array)], false, vec![2, 0, 2, 0, 2, 2])?; - verify_longest_length(&[Arc::clone(&list_array)], true, vec![2, 1, 2, 1, 2, 2])?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Drop, + vec![2, 0, 2, 0, 2, 2], + )?; + verify_longest_length( + &[Arc::clone(&list_array)], + NullHandling::Preserve, + vec![2, 1, 2, 1, 2, 2], + )?; // Test with multiple list arrays // [A, B, C], [], NULL, [D], NULL, [NULL, F] @@ -1378,8 +1982,17 @@ mod tests { let list1 = Arc::new(make_generic_array::()) as ArrayRef; let list2 = Arc::new(make_fixed_list()) as ArrayRef; let list_arrays = vec![Arc::clone(&list1), Arc::clone(&list2)]; - verify_longest_length(&list_arrays, false, vec![3, 0, 2, 1, 2, 2])?; - verify_longest_length(&list_arrays, true, vec![3, 1, 2, 1, 2, 2])?; + verify_longest_length(&list_arrays, NullHandling::Drop, vec![3, 0, 2, 1, 2, 2])?; + verify_longest_length( + &list_arrays, + NullHandling::Preserve, + vec![3, 1, 2, 1, 2, 2], + )?; + verify_longest_length( + &list_arrays, + NullHandling::PreserveAndExpandEmpty, + vec![3, 1, 2, 1, 2, 2], + )?; Ok(()) } diff --git a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs index f442bcea94be2..d4c98009ba70d 100644 --- a/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/bounded_window_agg_exec.rs @@ -28,20 +28,22 @@ use std::task::{Context, Poll}; use super::utils::create_schema; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::windows::{ calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, window_equivalence_properties, }; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, InputOrderMode, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, + InputOrderMode, PlanProperties, RecordBatchStream, ReplaceChildrenOptions, + SendableRecordBatchStream, Statistics, WindowExpr, validate_child_count, }; use arrow::compute::take_record_batch; use arrow::{ - array::{Array, ArrayRef, RecordBatchOptions, UInt32Builder}, + array::{Array, ArrayRef, RecordBatchOptions, UInt32Array, UInt32Builder}, compute::{concat, concat_batches, sort_to_indices, take_arrays}, datatypes::SchemaRef, record_batch::RecordBatch, @@ -53,13 +55,14 @@ use datafusion_common::utils::{ evaluate_partition_ranges, get_at_indices, get_row_at_idx, }; use datafusion_common::{ - HashMap, Result, arrow_datafusion_err, exec_datafusion_err, exec_err, + HashMap, Result, ScalarValue, arrow_datafusion_err, exec_datafusion_err, exec_err, }; use datafusion_execution::TaskContext; use datafusion_expr::ColumnarValue; use datafusion_expr::window_state::{PartitionBatchState, WindowAggState}; use datafusion_physical_expr::window::{ - PartitionBatches, PartitionKey, PartitionWindowAggStates, WindowState, + PartitionBatches, PartitionKey, PartitionWindowAggStates, WindowEvalContext, + WindowState, }; use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_expr_common::sort_expr::{ @@ -74,8 +77,47 @@ use hashbrown::hash_table::HashTable; use indexmap::IndexMap; use log::debug; +/// Callback receiver for per-partition window state. +/// +/// `state` is the result of [`Accumulator::state`], which is a `&mut self` +/// call whose trait doc states "this function should not be called twice." +/// Several built-in aggregates (`median`, `percentile_cont`, `string_agg`, +/// `min_max_bytes`/`min_max_struct`) `std::mem::take` their internal +/// buffers to build that state — so `state` is a destructive read, not a +/// snapshot. The exec fires this at most once per group; a callee that +/// needs the value beyond the callback must retain it (e.g. clone into +/// owned storage). +/// +/// [`Accumulator::state`]: datafusion_expr::Accumulator::state +pub trait WindowStateObserver: Send + Sync { + /// Invoked once per (output-partition-index, window-expression, + /// PARTITION BY tuple) as each PARTITION BY group closes, for every + /// aggregate window expression on the exec. Non-aggregate window + /// functions (e.g. `row_number`, `rank`, `lead`/`lag`) do not fire this + /// callback. + /// + /// # Arguments + /// + /// * `partition_idx` - Output partition index of the [`BoundedWindowAggExec`] + /// stream firing this callback. + /// * `window_expr` - The window expression whose state just closed. + /// * `partition_key` - The PARTITION BY tuple that just closed. + /// * `state` - [`Accumulator::state`] for the closed group of + /// `window_expr`. See the trait-level doc for the destructive-read + /// contract. + /// + /// [`Accumulator::state`]: datafusion_expr::Accumulator::state + fn finalize_window_aggregate( + &self, + partition_idx: usize, + window_expr: &Arc, + partition_key: &PartitionKey, + state: Vec, + ) -> Result<()>; +} + /// Window execution plan -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct BoundedWindowAggExec { /// Input plan input: Arc, @@ -98,6 +140,32 @@ pub struct BoundedWindowAggExec { cache: Arc, /// If `can_rerepartition` is false, partition_keys is always empty. can_repartition: bool, + /// Invoked at partition-close to publish finalized per-partition window + /// state. Storage and multi-group handling are the caller's; the exec is + /// a pure event source. + state_observer: Option>, +} + +impl std::fmt::Debug for BoundedWindowAggExec { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("BoundedWindowAggExec") + .field("input", &self.input) + .field("window_expr", &self.window_expr) + .field("schema", &self.schema) + .field("metrics", &self.metrics) + .field("input_order_mode", &self.input_order_mode) + .field( + "ordered_partition_by_indices", + &self.ordered_partition_by_indices, + ) + .field("cache", &self.cache) + .field("can_repartition", &self.can_repartition) + .field( + "state_observer", + &self.state_observer.as_ref().map(|_| "..."), + ) + .finish() + } } impl BoundedWindowAggExec { @@ -138,9 +206,50 @@ impl BoundedWindowAggExec { ordered_partition_by_indices, cache: Arc::new(cache), can_repartition, + state_observer: None, }) } + /// Install (or clear) a [`WindowStateObserver`] that receives each + /// PARTITION BY group's finalized window state at partition close. + /// + /// Errors when `observer` is `Some` and any window expression on this + /// exec has a non-ever-expanding frame (i.e. its start bound is not + /// `UNBOUNDED PRECEDING`). Those frames use `SlidingAggregateWindowExpr` + /// under the hood, whose accumulator calls `retract_batch` — at + /// partition close the accumulator holds only the last frame's rows, + /// not the partition aggregate, so the observed state would silently + /// misrepresent the group. + pub fn with_state_observer( + mut self, + observer: Option>, + ) -> Result { + if observer.is_some() { + for expr in &self.window_expr { + if !expr.get_window_frame().is_ever_expanding() { + return exec_err!( + "cannot install WindowStateObserver on BoundedWindowAggExec \ + with a sliding aggregate window frame (start != \ + UNBOUNDED PRECEDING) for `{}`; sliding accumulator state \ + is frame-only, not the partition aggregate", + expr.name() + ); + } + } + } + self.state_observer = observer; + Ok(self) + } + + /// The currently-installed [`WindowStateObserver`], if any. Optimizer + /// rules that rebuild this exec via + /// [`crate::windows::get_best_fitting_window`] or a direct `try_new` + /// call must read this and reinstall it on the new exec, otherwise a + /// caller-installed observer is silently dropped by the rewrite. + pub fn state_observer(&self) -> Option<&Arc> { + self.state_observer.as_ref() + } + /// Window expressions pub fn window_expr(&self) -> &[Arc] { &self.window_expr @@ -250,17 +359,6 @@ impl BoundedWindowAggExec { total_byte_size: Precision::Absent, }) } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for BoundedWindowAggExec { @@ -323,15 +421,17 @@ impl ExecutionPlan for BoundedWindowAggExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for window_expr in &self.window_expr { - for expr in window_expr.expressions() { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - } - Ok(tnr) + let expressions = self.window_expr.iter().flat_map(|window_expr| { + let expressions = window_expr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.partition_by_exprs) + .chain(expressions.order_by_exprs) + }); + crate::apply_expression_roots(expressions, f) } fn required_input_ordering(&self) -> Vec> { @@ -345,11 +445,17 @@ impl ExecutionPlan for BoundedWindowAggExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { if self.partition_keys().is_empty() { debug!("No partition defined for BoundedWindowAggExec!!!"); - vec![Distribution::SinglePartition] + InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } else { - vec![Distribution::HashPartitioned(self.partition_keys().clone())] + InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + self.partition_keys(), + )]) } } @@ -357,17 +463,49 @@ impl ExecutionPlan for BoundedWindowAggExec { vec![true] } + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new = BoundedWindowAggExec::try_new( + self.window_expr.clone(), + Arc::clone(&children[0]), + self.input_order_mode.clone(), + self.can_repartition, + )? + .with_state_observer(self.state_observer.clone())?; + Ok(Arc::new(new)) + } + } + } + fn with_new_children( self: Arc, children: Vec>, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(BoundedWindowAggExec::try_new( - self.window_expr.clone(), - Arc::clone(&children[0]), - self.input_order_mode.clone(), - self.can_repartition, - )?)) + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -383,6 +521,8 @@ impl ExecutionPlan for BoundedWindowAggExec { input, BaselineMetrics::new(&self.metrics, partition), search_mode, + partition, + self.state_observer.clone(), )?); Ok(stream) } @@ -391,15 +531,97 @@ impl ExecutionPlan for BoundedWindowAggExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let input_stat = - Arc::unwrap_or_clone(self.input.partition_statistics(partition)?); + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let input_stat = input_stats[0].as_ref().clone(); Ok(Arc::new(self.statistics_helper(input_stat)?)) } fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::Equal } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use super::proto::encode_physical_window_expr; + use datafusion_proto_common::protobuf_common::EmptyMessage; + use datafusion_proto_models::protobuf; + use protobuf::window_agg_exec_node::InputOrderMode as ProtoInputOrderMode; + + // Exhaustive destructure: adding a field to `BoundedWindowAggExec` + // without deciding how it is serialized is a compile error, not a + // silent round-trip gap. + let Self { + input, + window_expr, + // Derived at construction by `create_schema` from the input schema + // and the window expressions. + schema: _, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + input_order_mode, + // Derived at construction from `input_order_mode` and the window + // expressions' PARTITION BY. + ordered_partition_by_indices: _, + // Derived at construction by `Self::compute_properties`. + cache: _, + // No wire field of its own; it is folded into `partition_keys` + // below, since `partition_keys()` returns an empty vec when this is + // false and the decoder recovers it as `!partition_keys.is_empty()`. + can_repartition: _, + // Runtime callback installed after planning; not part of the wire + // format. Any decoder that needs it must reinstall via + // `with_state_observer`. + state_observer: _, + } = self; + + let input = ctx.encode_child(input)?; + let window_expr = window_expr + .iter() + .map(|expr| encode_physical_window_expr(expr, ctx)) + .collect::>>()?; + let partition_keys = self + .partition_keys() + .iter() + .map(|expr| ctx.encode_expr(expr)) + .collect::>>()?; + // A `Some(input_order_mode)` is what tells the shared `Window` decode + // arm to rebuild a `BoundedWindowAggExec` rather than a `WindowAggExec`. + let input_order_mode = match input_order_mode { + InputOrderMode::Linear => ProtoInputOrderMode::Linear(EmptyMessage {}), + InputOrderMode::PartiallySorted(columns) => { + ProtoInputOrderMode::PartiallySorted( + protobuf::PartiallySortedInputOrderMode { + columns: columns.iter().map(|column| *column as u64).collect(), + }, + ) + } + InputOrderMode::Sorted => ProtoInputOrderMode::Sorted(EmptyMessage {}), + }; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Window(Box::new( + protobuf::WindowAggExecNode { + input: Some(Box::new(input)), + window_expr, + partition_keys, + input_order_mode: Some(input_order_mode), + }, + )), + ), + })) + } } /// Trait that specifies how we search for (or calculate) partitions. It has two @@ -476,25 +698,6 @@ trait PartitionSearcher: Send { } } - if self.is_mode_linear() { - // In `Linear` mode, it is guaranteed that the first ORDER BY column - // is sorted across partitions. Note that only the first ORDER BY - // column is guaranteed to be ordered. As a counter example, consider - // the case, `PARTITION BY b, ORDER BY a, c` when the input is sorted - // by `[a, b, c]`. In this case, `BoundedWindowAggExec` mode will be - // `Linear`. However, we cannot guarantee that the last row of the - // input data will be the "last" data in terms of the ordering requirement - // `[a, c]` -- it will be the "last" data in terms of `[a, b, c]`. - // Hence, only column `a` should be used as a guarantee of the "last" - // data across partitions. For other modes (`Sorted`, `PartiallySorted`), - // we do not need to keep track of the most recent row guarantee across - // partitions. Since leading ordering separates partitions, guaranteed - // by the most recent row, already prune the previous partitions completely. - let last_row = get_last_row_batch(&record_batch)?; - for (_, partition_batch) in partition_buffers.iter_mut() { - partition_batch.set_most_recent_row(last_row.clone()); - } - } self.mark_partition_end(partition_buffers); *input_buffer = if input_buffer.num_rows() == 0 { @@ -632,17 +835,25 @@ impl PartitionSearcher for LinearSearch { evaluate_partition_by_column_values(record_batch, window_expr)?; // NOTE: In Linear or PartiallySorted modes, we are sure that // `partition_bys` are not empty. - // Calculate indices for each partition and construct a new record - // batch from the rows at these indices for each partition: - self.get_per_partition_indices(&partition_bys, record_batch)? + let (mut keys, permutation, bounds) = + self.compute_partition_permutation(&partition_bys, record_batch)?; + if keys.len() == 1 { + // The batch contains a single partition, so the gather below + // would be an identity permutation; use the batch as-is. + let key = keys.remove(0); + return Ok(vec![(key, record_batch.clone())]); + } + // Reorder the batch with a single `take` so that each partition's + // rows become contiguous, then hand each partition a zero-copy slice + // of the result. The slices share the gathered batch's buffers; + // `PartitionBatchState::extend` copies out of them the next time the + // partition receives rows. + let gathered = take_record_batch(record_batch, &UInt32Array::from(permutation))?; + Ok(keys .into_iter() - .map(|(row, indices)| { - let mut new_indices = UInt32Builder::with_capacity(indices.len()); - new_indices.append_slice(&indices); - let indices = new_indices.finish(); - Ok((row, take_record_batch(record_batch, &indices)?)) - }) - .collect() + .zip(bounds.windows(2)) + .map(|(key, bound)| (key, gathered.slice(bound[0], bound[1] - bound[0]))) + .collect()) } fn prune(&mut self, n_out: usize) { @@ -696,42 +907,66 @@ impl LinearSearch { } } - /// Calculate indices of each partition (according to PARTITION BY expression) - /// `columns` contain partition by expression results. - fn get_per_partition_indices( + /// Splits the rows of `batch` by partition, according to the PARTITION BY + /// expression results in `columns`. Returns the distinct partition keys + /// in first-appearance order, a permutation of the row indices of + /// `batch` that groups each partition's rows together, and the + /// boundaries of each partition's run of rows within that permutation: + /// partition `p` occupies `permutation[bounds[p]..bounds[p + 1]]`, and + /// its indices are in ascending (stream) order. + fn compute_partition_permutation( &mut self, columns: &[ArrayRef], batch: &RecordBatch, - ) -> Result)>> { - let mut batch_hashes = vec![0; batch.num_rows()]; + ) -> Result<(Vec, Vec, Vec)> { + let num_rows = batch.num_rows(); + let mut batch_hashes = vec![0; num_rows]; create_hashes(columns, &self.random_state, &mut batch_hashes)?; self.input_buffer_hashes.extend(&batch_hashes); // reset row_map for new calculation self.row_map_batch.clear(); - // res stores PartitionKey and row indices (indices where these partition occurs in the `batch`) for each partition. - let mut result: Vec<(PartitionKey, Vec)> = vec![]; + let mut keys: Vec = vec![]; + // Partition id of each row, in row order: + let mut row_partition_ids = Vec::with_capacity(num_rows); + // Number of rows in each partition: + let mut counts: Vec = vec![]; for (hash, row_idx) in batch_hashes.into_iter().zip(0u32..) { let entry = self.row_map_batch.find_mut(hash, |(_, group_idx)| { - // We can safely get the first index of the partition indices - // since partition indices has one element during initialization. let row = get_row_at_idx(columns, row_idx as usize).unwrap(); - // Handle hash collusions with an equality check: - row.eq(&result[*group_idx].0) + // Handle hash collisions with an equality check: + row == keys[*group_idx] }); - if let Some((_, group_idx)) = entry { - result[*group_idx].1.push(row_idx) + let group_idx = if let Some((_, group_idx)) = entry { + *group_idx } else { - self.row_map_batch.insert_unique( - hash, - (hash, result.len()), - |(hash, _)| *hash, - ); - let row = get_row_at_idx(columns, row_idx as usize)?; - // This is a new partition its only index is row_idx for now. - result.push((row, vec![row_idx])); - } + let group_idx = keys.len(); + self.row_map_batch + .insert_unique(hash, (hash, group_idx), |(hash, _)| *hash); + keys.push(get_row_at_idx(columns, row_idx as usize)?); + counts.push(0); + group_idx + }; + row_partition_ids.push(group_idx); + counts[group_idx] += 1; + } + // A prefix sum over the counts gives each partition's run boundaries + // in the permutation. + let mut bounds = Vec::with_capacity(counts.len() + 1); + let mut total = 0; + bounds.push(0); + for count in counts { + total += count; + bounds.push(total); + } + // Scatter each row's index into its partition's run. Visiting rows + // in ascending order keeps each run in ascending row order. + let mut cursors: Vec = bounds[..bounds.len() - 1].to_vec(); + let mut permutation = vec![0u32; num_rows]; + for (row_idx, group_idx) in row_partition_ids.into_iter().enumerate() { + permutation[cursors[group_idx]] = row_idx as u32; + cursors[group_idx] += 1; } - Ok(result) + Ok((keys, permutation, bounds)) } /// Calculates partition keys and result indices for each partition. @@ -941,14 +1176,9 @@ pub struct BoundedWindowAggStream { /// The record batch executor receives as input (i.e. the columns needed /// while calculating aggregation results). input_buffer: RecordBatch, - /// We separate `input_buffer` based on partitions (as - /// determined by PARTITION BY columns) and store them per partition - /// in `partition_batches`. We use this variable when calculating results - /// for each window expression. This enables us to use the same batch for - /// different window expressions without copying. - // Note that we could keep record batches for each window expression in - // `PartitionWindowAggStates`. However, this would use more memory (as - // many times as the number of window expressions). + /// Each partition's rows, accumulated across input batches. All window + /// expressions calculate their results against these shared rows without + /// copying. partition_buffers: PartitionBatches, /// An executor can run multiple window expressions if the PARTITION BY /// and ORDER BY sections are same. We keep state of the each window @@ -960,20 +1190,77 @@ pub struct BoundedWindowAggStream { /// Search mode for partition columns. This determines the algorithm with /// which we group each partition. search_mode: Box, + /// In `Linear` mode, a single-row batch containing the most recent input + /// row (whichever partition that row belongs to); `None` in other modes + /// and before the first non-empty batch arrives. Since in `Linear` mode + /// the input is sorted by the first ORDER BY column, no future input row + /// -- in any partition -- can precede this row in that column. Every + /// partition's evaluation consults this bound to decide whether pending + /// window frames can be finalized before the partition receives more + /// data (which in turn allows buffered state to be pruned). Note that + /// only the first ORDER BY column provides this guarantee. As a counter + /// example, consider `PARTITION BY b, ORDER BY a, c` when the input is + /// sorted by `[a, b, c]`: the mode will be `Linear`, but the last row of + /// the input is the "last" data in terms of `[a, b, c]`, not in terms of + /// the ordering requirement `[a, c]`. Hence, only column `a` can serve + /// as a guarantee of the "last" data across partitions. In the `Sorted` + /// and `PartiallySorted` modes, the leading ordering separates + /// partitions, so finished partitions are pruned eagerly instead and no + /// such bound is needed. + most_recent_row: Option, + /// Output partition index this stream serves; passed as the first + /// argument to [`WindowStateObserver::finalize_window_aggregate`]. + partition_idx: usize, + /// If set, invoked from [`Self::publish_finalized_states`] with the + /// finalized per-window-expression state for every partition key that is + /// about to be dropped. + state_observer: Option>, } impl BoundedWindowAggStream { + /// Fire `observer` once per (window expression, partition key) for every + /// group whose [`WindowAggState::is_end`] is true. Always mutates when + /// called: [`datafusion_expr::Accumulator::state`] requires `&mut`, which + /// propagates up here. The caller is responsible for deciding whether to + /// fire (i.e. checking whether an observer is installed). + /// + /// Exactly-once per group is enforced by [`WindowState::aggregate_state`], + /// which errors on second call; the `published` early-skip below avoids reaching the error. + fn publish_finalized_states( + &mut self, + observer: &dyn WindowStateObserver, + ) -> Result<()> { + let partition_idx = self.partition_idx; + for (expr_idx, per_expr) in self.window_agg_states.iter_mut().enumerate() { + let window_expr = &self.window_expr[expr_idx]; + for (key, ws) in per_expr.iter_mut() { + if ws.published || !ws.state.is_end { + continue; + } + if let Some(state) = ws.aggregate_state()? { + observer.finalize_window_aggregate( + partition_idx, + window_expr, + key, + state, + )?; + } + } + } + Ok(()) + } + /// Prunes sections of the state that are no longer needed when calculating /// results (as determined by window frame boundaries and number of results generated). // For instance, if first `n` (not necessarily same with `n_out`) elements are no longer needed to // calculate window expression result (outside the window frame boundary) we retract first `n` elements - // from `self.partition_batches` in corresponding partition. + // from the corresponding partition's batch in `self.partition_buffers`. // For instance, if `n_out` number of rows are calculated, we can remove // first `n_out` rows from `self.input_buffer`. fn prune_state(&mut self, n_out: usize) -> Result<()> { // Prune `self.window_agg_states`: self.prune_out_columns(); - // Prune `self.partition_batches`: + // Prune `self.partition_buffers`: self.prune_partition_batches(); // Prune `self.input_buffer`: self.prune_input_batch(n_out)?; @@ -1003,28 +1290,47 @@ impl BoundedWindowAggStream { input: SendableRecordBatchStream, baseline_metrics: BaselineMetrics, search_mode: Box, + partition_idx: usize, + state_observer: Option>, ) -> Result { - let state = window_expr.iter().map(|_| IndexMap::new()).collect(); + let state = window_expr.iter().map(|_| IndexMap::default()).collect(); let empty_batch = RecordBatch::new_empty(Arc::clone(&schema)); Ok(Self { schema, input, input_buffer: empty_batch, - partition_buffers: IndexMap::new(), + partition_buffers: IndexMap::default(), window_agg_states: state, finished: false, window_expr, baseline_metrics, search_mode, + most_recent_row: None, + partition_idx, + state_observer, }) } fn compute_aggregates(&mut self) -> Result> { // calculate window cols + let eval_ctx = WindowEvalContext::default() + .with_most_recent_row(self.most_recent_row.as_ref()); for (cur_window_expr, state) in self.window_expr.iter().zip(&mut self.window_agg_states) { - cur_window_expr.evaluate_stateful(&self.partition_buffers, state)?; + cur_window_expr.evaluate_stateful( + &self.partition_buffers, + state, + &eval_ctx, + )?; + } + + // Fire before `calculate_out_columns`: on causal frames every row + // already streamed out, so at EOS that call returns `None` and the + // prune path is skipped — the final partition would otherwise be + // dropped unobserved. + if let Some(observer) = self.state_observer.clone() { + self.publish_finalized_states(observer.as_ref())?; } let schema = Arc::clone(&self.schema); @@ -1068,6 +1374,9 @@ impl BoundedWindowAggStream { // stopped when dropped. let _timer = elapsed_compute.timer(); + if self.search_mode.is_mode_linear() && batch.num_rows() > 0 { + self.most_recent_row = Some(get_last_row_batch(&batch)?); + } self.search_mode.update_partition_batch( &mut self.input_buffer, batch, @@ -1099,49 +1408,76 @@ impl BoundedWindowAggStream { } } - /// Prunes the sections of the record batch (for each partition) - /// that we no longer need to calculate the window function result. + /// Removes partitions that have ended. For the remaining partitions, + /// drops buffered rows that no window expression will need again. fn prune_partition_batches(&mut self) { + // Check that per-state and per-partition end-flags are consistent; + // otherwise, the pruning code below might produce inconsistent state. + #[cfg(debug_assertions)] + for window_agg_state in self.window_agg_states.iter() { + for (partition_row, WindowState { state, .. }) in window_agg_state.iter() { + debug_assert_eq!( + state.is_end, self.partition_buffers[partition_row].is_end, + "window state's recorded end flag is out of sync with its partition" + ); + } + } + // Remove partitions which we know already ended (is_end flag is true). // Since the retain method preserves insertion order, we still have // ordering in between partitions after removal. self.partition_buffers .retain(|_, partition_batch_state| !partition_batch_state.is_end); - - // The data in `self.partition_batches` is used by all window expressions. - // Therefore, when removing from `self.partition_batches`, we need to remove - // from the earliest range boundary among all window expressions. Variable - // `n_prune_each_partition` fill the earliest range boundary information for - // each partition. This way, we can delete the no-longer-needed sections from - // `self.partition_batches`. - // For instance, if window frame one uses [10, 20] and window frame two uses - // [5, 15]; we only prune the first 5 elements from the corresponding record - // batch in `self.partition_batches`. - - // Calculate how many elements to prune for each partition batch - let mut n_prune_each_partition = HashMap::new(); + // Likewise, drop per-window-expression state for ended partitions. for window_agg_state in self.window_agg_states.iter_mut() { window_agg_state.retain(|_, WindowState { state, .. }| !state.is_end); - for (partition_row, WindowState { state: value, .. }) in window_agg_state { + } + + // Calculate how many rows to prune from each partition's batch. For a + // single window expression, rows before min(window_frame_range.start, + // last_calculated_index) are prunable: their results are already + // calculated, and frame boundaries never move backwards, so no future + // frame can include them. All window expressions share the partition + // batch, so a row can only be pruned once every expression is done with + // it: the count to prune is the minimum across expressions. A partition + // missing from the map has nothing to prune. + let mut n_prune_each_partition = HashMap::new(); + if let Some((first, rest)) = self.window_agg_states.split_first() { + // First window expression seeds the prune-count map + for (partition_row, WindowState { state, .. }) in first.iter() { let n_prune = - min(value.window_frame_range.start, value.last_calculated_index); - if let Some(current) = n_prune_each_partition.get_mut(partition_row) { - if n_prune < *current { - *current = n_prune; - } - } else { + min(state.window_frame_range.start, state.last_calculated_index); + if n_prune > 0 { n_prune_each_partition.insert(partition_row.clone(), n_prune); } } + // Take the per-partition min of the prune-count for each + // additional window expression + for window_agg_state in rest { + n_prune_each_partition.retain(|partition_row, current| { + let Some(WindowState { state, .. }) = + window_agg_state.get(partition_row) + else { + return false; + }; + let n_prune = + min(state.window_frame_range.start, state.last_calculated_index); + *current = min(*current, n_prune); + *current > 0 + }); + } } - // Retract no longer needed parts during window calculations from partition batch: + // Drop the prunable prefix of each partition's buffered batch: for (partition_row, n_prune) in n_prune_each_partition.iter() { + debug_assert!( + *n_prune > 0, + "prune-count map must only contain positive entries" + ); let pb_state = &mut self.partition_buffers[partition_row]; let batch = &pb_state.record_batch; pb_state.record_batch = batch.slice(*n_prune, batch.num_rows() - n_prune); - pb_state.n_out_row = 0; // Update state indices since we have pruned some rows from the beginning: for window_agg_state in self.window_agg_states.iter_mut() { @@ -1175,23 +1511,29 @@ impl BoundedWindowAggStream { // field of `WindowAggState`. Given how many rows are emitted, we remove // these sections from state. for partition_window_agg_states in self.window_agg_states.iter_mut() { - // Remove `n_out` entries from the `out_col` field of `WindowAggState`. - // `n_out` is stored in `self.partition_buffers` for each partition. - // If `is_end` is set, directly remove them; this shrinks the hash map. + // If `is_end` is set, directly remove the entry; this shrinks the + // hash map. partition_window_agg_states .retain(|_, partition_batch_state| !partition_batch_state.state.is_end); - for ( - partition_key, - WindowState { - state: WindowAggState { out_col, .. }, - .. - }, - ) in partition_window_agg_states - { - let partition_batch = &mut self.partition_buffers[partition_key]; - let n_to_del = partition_batch.n_out_row; - let n_to_keep = out_col.len() - n_to_del; - *out_col = out_col.slice(n_to_del, n_to_keep); + } + // Only partitions that emitted rows since the previous pruning pass + // have output columns to shrink. Their emitted-row counts are + // consumed and reset here, so partitions that emitted nothing keep + // a count of zero and are passed over without any hash lookups. + for (partition_key, partition_batch) in self.partition_buffers.iter_mut() { + let n_emitted = partition_batch.n_out_row; + if n_emitted == 0 { + continue; + } + partition_batch.n_out_row = 0; + for partition_window_agg_states in self.window_agg_states.iter_mut() { + if let Some(WindowState { state, .. }) = + partition_window_agg_states.get_mut(partition_key) + { + let out_col = &mut state.out_col; + let n_to_keep = out_col.len() - n_emitted; + *out_col = out_col.slice(n_emitted, n_to_keep); + } } } } @@ -1277,10 +1619,11 @@ mod tests { use crate::projection::{ProjectionExec, ProjectionExpr}; use crate::streaming::{PartitionStream, StreamingTableExec}; use crate::test::TestMemoryExec; + use crate::windows::bounded_window_agg_exec::WindowStateObserver; use crate::windows::{ BoundedWindowAggExec, InputOrderMode, create_udwf_window_expr, create_window_expr, }; - use crate::{ExecutionPlan, displayable, execute_stream}; + use crate::{ExecutionPlan, WindowExpr, displayable, execute_stream}; use arrow::array::{ RecordBatch, @@ -1298,10 +1641,11 @@ mod tests { WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, }; use datafusion_functions_aggregate::count::count_udaf; + use datafusion_functions_aggregate::sum::sum_udaf; use datafusion_functions_window::nth_value::last_value_udwf; use datafusion_functions_window::nth_value::nth_value_udwf; use datafusion_physical_expr::expressions::{Column, Literal, col}; - use datafusion_physical_expr::window::StandardWindowExpr; + use datafusion_physical_expr::window::{PartitionKey, StandardWindowExpr}; use datafusion_physical_expr::{LexOrdering, PhysicalExpr}; use futures::future::Shared; @@ -1737,6 +2081,108 @@ mod tests { Ok(()) } + // In `Linear` mode, a partition may receive no new rows for several + // input batches while other partitions keep growing. Once all of a + // partition's buffered rows have results, the evaluation sweep skips + // it until it receives rows again, so this test drives a partition + // through quiet batches and then resumes it: the results after the + // gap must continue from the retained accumulator state. Both frames + // are causal, so results finalize in the batch their row arrives in + // and the quiet partition is fully calculated while it waits. + #[tokio::test] + async fn bounded_window_linear_quiet_partition_resume() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("pk", DataType::UInt64, false), + Field::new("ts", DataType::UInt64, false), + ])); + let make_batch = |rows: &[(u64, u64)]| -> Result { + let mut pk = UInt64Builder::with_capacity(rows.len()); + let mut ts = UInt64Builder::with_capacity(rows.len()); + for (p, t) in rows { + pk.append_value(*p); + ts.append_value(*t); + } + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(pk.finish()), Arc::new(ts.finish())], + )?) + }; + // `ts` ascends globally; partition 0 is absent from the middle batches. + let batches = vec![ + make_batch(&[(0, 0), (0, 1), (1, 2)])?, + make_batch(&[(1, 3), (1, 4)])?, + make_batch(&[(1, 5)])?, + make_batch(&[(0, 6), (1, 7)])?, + ]; + let memory_exec = + TestMemoryExec::try_new_exec(&[batches], Arc::clone(&schema), None)?; + + let partition_by = vec![col("pk", &schema)?]; + let order_by = [PhysicalSortExpr { + expr: col("ts", &schema)?, + options: SortOptions::default(), + }]; + // A running COUNT (plain aggregate) and a SUM over the previous and + // current row (sliding aggregate). + let count_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count".to_string(), + &[col("ts", &schema)?], + &partition_by, + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + Arc::clone(&schema), + false, + false, + None, + )?; + let sum_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(sum_udaf()), + "sum".to_string(), + &[col("ts", &schema)?], + &partition_by, + &order_by, + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))), + WindowFrameBound::CurrentRow, + )), + Arc::clone(&schema), + false, + false, + None, + )?; + let physical_plan = BoundedWindowAggExec::try_new( + vec![count_expr, sum_expr], + memory_exec, + InputOrderMode::Linear, + true, + ) + .map(|e| Arc::new(e) as Arc)?; + + let batches = collect(physical_plan.execute(0, task_context())?).await?; + + assert_snapshot!(batches_to_string(&batches), @r" + +----+----+-------+-----+ + | pk | ts | count | sum | + +----+----+-------+-----+ + | 0 | 0 | 1 | 0 | + | 0 | 1 | 2 | 1 | + | 1 | 2 | 1 | 2 | + | 1 | 3 | 2 | 5 | + | 1 | 4 | 3 | 7 | + | 1 | 5 | 4 | 9 | + | 0 | 6 | 3 | 7 | + | 1 | 7 | 5 | 12 | + +----+----+-------+-----+ + "); + Ok(()) + } + // This test, tests whether most recent row guarantee by the input batch of the `BoundedWindowAggExec` // helps `BoundedWindowAggExec` to generate low latency result in the `Linear` mode. // Input data generated at the source is @@ -1858,6 +2304,629 @@ mod tests { Ok(()) } + type Observation = (usize, PartitionKey, Vec); + + /// Test [`WindowStateObserver`] that records every callback into a shared + /// `Vec` for later assertion. + struct RecordingObserver { + sink: Arc>>, + } + + impl WindowStateObserver for RecordingObserver { + fn finalize_window_aggregate( + &self, + partition_idx: usize, + _window_expr: &Arc, + partition_key: &PartitionKey, + state: Vec, + ) -> Result<()> { + self.sink + .lock() + .unwrap() + .push((partition_idx, partition_key.clone(), state)); + Ok(()) + } + } + + /// Build a `BoundedWindowAggExec` for `count(sn) OVER (PARTITION BY hash + /// ORDER BY sn )` over a fixed two-group source (hash=1 × 3, + /// hash=2 × 3, sorted by (hash, sn)). Returns the plan pre-observer so + /// callers can decide how to install it. + fn build_partition_close_plan(frame: WindowFrame) -> Result { + let schema = test_schema(); + + let mut sn_b = UInt64Builder::with_capacity(6); + let mut hash_b = Int64Builder::with_capacity(6); + for (sn, hash) in [(1u64, 1i64), (2, 1), (3, 1), (4, 2), (5, 2), (6, 2)] { + sn_b.append_value(sn); + hash_b.append_value(hash); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?; + let ordering: LexOrdering = [ + PhysicalSortExpr { + expr: col("hash", &schema)?, + options: SortOptions::default(), + }, + PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }, + ] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "cnt".to_string(), + &[col("sn", &schema)?], + &[col("hash", &schema)?], + &[PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }], + Arc::new(frame), + source.schema(), + false, + false, + None, + )?; + + BoundedWindowAggExec::try_new(vec![expr], source, InputOrderMode::Sorted, false) + } + + // Two PARTITION BY groups: hash=1 [sn=1,2,3] then hash=2 [sn=4,5,6]. + // Input is sorted by (hash, sn) so we can run in Sorted mode; in that + // mode `mark_partition_end` closes the leading group mid-stream and + // EOS closes the tail — both fire the observer for an ever-expanding + // frame. Sliding frames are rejected at install time. + + #[tokio::test] + async fn test_state_observer_rejects_sliding_frame() -> Result<()> { + // `CURRENT ROW → UNBOUNDED FOLLOWING` is not ever-expanding, so this + // maps to `SlidingAggregateWindowExpr` whose accumulator retracts as + // rows leave the frame — at partition close the accumulator holds + // only the last frame's rows, not the partition aggregate. + // `with_state_observer` refuses this configuration. + use std::sync::Mutex; + + let plan = build_partition_close_plan(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::CurrentRow, + WindowFrameBound::Following(ScalarValue::UInt64(None)), + ))?; + let observer: Arc = Arc::new(RecordingObserver { + sink: Arc::new(Mutex::new(vec![])), + }); + let err = plan.with_state_observer(Some(observer)).unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("sliding aggregate window frame"), + "expected sliding-frame rejection, got: {msg}" + ); + Ok(()) + } + + #[tokio::test] + async fn test_finalized_state_observer_fires_on_causal_frame() -> Result<()> { + // `ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW` — ever-expanding, + // `PlainAggregateWindowExpr` under the hood. At partition close the + // accumulator holds the partition aggregate. Both mid-stream close + // (hash=1 as hash=2 rows arrive) and EOS (hash=2 at drain) fire. + use std::sync::Mutex; + + let task_ctx = Arc::new(TaskContext::default()); + let plan = build_partition_close_plan(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + ))?; + + let observations: Arc>> = Arc::new(Mutex::new(vec![])); + let observer: Arc = Arc::new(RecordingObserver { + sink: Arc::clone(&observations), + }); + let plan = plan.with_state_observer(Some(observer))?; + + let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + // count(sn) over each of hash=1 (3 rows) and hash=2 (3 rows), in + // close order — hash=1 first (mid-stream close), hash=2 second (EOS). + let observed: Vec<(usize, i64, Vec)> = observations + .lock() + .unwrap() + .iter() + .map(|(idx, key, state)| { + let hash = match &key[0] { + ScalarValue::Int64(Some(v)) => *v, + other => panic!("unexpected partition-key element: {other:?}"), + }; + (*idx, hash, state.clone()) + }) + .collect(); + assert_eq!( + observed, + vec![ + (0, 1, vec![ScalarValue::Int64(Some(3))]), + (0, 2, vec![ScalarValue::Int64(Some(3))]), + ] + ); + Ok(()) + } + + #[tokio::test] + async fn test_finalized_state_observer_fires_exactly_once_across_batches() + -> Result<()> { + // Regression guard for the exactly-once observer contract when + // partition close and pruning happen on different `compute_aggregates` + // calls. + // + // The observer fires from `publish_finalized_states`, called at the + // top of every `compute_aggregates`. Entries are only cleared by + // `prune_state`, which runs only when `calculate_out_columns` returns + // `Some`. Nothing in the type system ties the two together, so a + // group whose state was published on batch N must not be re-published + // on batch N+1 or at EOS. + // + // Layout: three PARTITION BY groups streamed across two input + // batches, so each group closes on a distinct `compute_aggregates` + // call: + // batch 1 = [hash=1 × 2] — no close (single group). + // batch 2 = [hash=2 × 2, hash=3 × 2] — `mark_partition_end` + // closes hash=1 and hash=2. + // EOS — closes hash=3. + // + // Assertion: each key appears exactly once across all observations. + use std::sync::Mutex; + + let task_ctx = Arc::new(TaskContext::default()); + let schema = test_schema(); + + // Two batches, same output partition. + let make_batch = |rows: &[(u64, i64)]| -> Result { + let mut sn_b = UInt64Builder::with_capacity(rows.len()); + let mut hash_b = Int64Builder::with_capacity(rows.len()); + for &(sn, hash) in rows { + sn_b.append_value(sn); + hash_b.append_value(hash); + } + Ok(RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?) + }; + let batch1 = make_batch(&[(1, 1), (2, 1)])?; + let batch2 = make_batch(&[(3, 2), (4, 2), (5, 3), (6, 3)])?; + + let ordering: LexOrdering = [ + PhysicalSortExpr { + expr: col("hash", &schema)?, + options: SortOptions::default(), + }, + PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }, + ] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch1, batch2]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "cnt".to_string(), + &[col("sn", &schema)?], + &[col("hash", &schema)?], + &[PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + source.schema(), + false, + false, + None, + )?; + + let observations: Arc>> = Arc::new(Mutex::new(vec![])); + let observer: Arc = Arc::new(RecordingObserver { + sink: Arc::clone(&observations), + }); + + let plan = BoundedWindowAggExec::try_new( + vec![expr], + source, + InputOrderMode::Sorted, + false, + )? + .with_state_observer(Some(observer))?; + + let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + let fired: Vec = observations + .lock() + .unwrap() + .iter() + .map(|(_, key, _)| match &key[0] { + ScalarValue::Int64(Some(v)) => *v, + other => panic!("unexpected partition-key element: {other:?}"), + }) + .collect(); + // Each group closes on a distinct `compute_aggregates` call — hash=1 + // and hash=2 on batch 2's `mark_partition_end`, hash=3 at EOS — and + // each appears exactly once, in close order. + assert_eq!(fired, vec![1, 2, 3]); + Ok(()) + } + + /// Run one task's local BWAG for `SUM(sn) OVER (ORDER BY sn ROWS + /// UNBOUNDED PRECEDING TO CURRENT ROW)` with no PARTITION BY, over + /// `input` sorted ascending. Returns the per-row output values and the + /// observed finalized state total (which the caller uses as a carry-in + /// for the next task). + async fn run_running_sum_task( + input: &[u64], + task_ctx: Arc, + ) -> Result<(Vec, u64)> { + use arrow::array::UInt64Array; + use datafusion_functions_aggregate::sum::sum_udaf; + use std::sync::Mutex; + + /// Observer for `run_running_sum_task`: captures the single running + /// SUM total published at EOS. Asserts exactly-one fire and rejects + /// non-empty partition keys (this helper is no-PARTITION-BY only). + struct RunningSumObserver { + sink: Arc>>, + } + + impl WindowStateObserver for RunningSumObserver { + fn finalize_window_aggregate( + &self, + _partition_idx: usize, + _window_expr: &Arc, + partition_key: &PartitionKey, + state: Vec, + ) -> Result<()> { + assert!( + partition_key.is_empty(), + "empty PartitionKey for no-PARTITION-BY plan" + ); + let total = match &state[0] { + ScalarValue::UInt64(Some(v)) => *v, + ScalarValue::Int64(Some(v)) => *v as u64, + other => panic!("unexpected sum state element: {other:?}"), + }; + let prev = self.sink.lock().unwrap().replace(total); + assert!(prev.is_none(), "observer must fire exactly once per task"); + Ok(()) + } + } + + let schema = test_schema(); + let mut sn_b = UInt64Builder::with_capacity(input.len()); + let mut hash_b = Int64Builder::with_capacity(input.len()); + for &sn in input { + sn_b.append_value(sn); + hash_b.append_value(0); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?; + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let window_fn = WindowFunctionDefinition::AggregateUDF(sum_udaf()); + let args = vec![col("sn", &schema)?]; + let partition_by: Vec> = vec![]; + let order_by = vec![PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }]; + let frame = WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + ); + let expr = create_window_expr( + &window_fn, + "running_sum".to_string(), + &args, + &partition_by, + &order_by, + Arc::new(frame), + source.schema(), + false, + false, + None, + )?; + + let total_sink: Arc>> = Arc::new(Mutex::new(None)); + let observer: Arc = Arc::new(RunningSumObserver { + sink: Arc::clone(&total_sink), + }); + + let plan = BoundedWindowAggExec::try_new( + vec![expr], + source, + InputOrderMode::Sorted, + false, + )? + .with_state_observer(Some(observer))?; + let batches = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + let mut out = Vec::with_capacity(input.len()); + for batch in &batches { + let col = batch + .column_by_name("running_sum") + .expect("running_sum column present"); + let arr = col + .as_any() + .downcast_ref::() + .expect("SUM(UInt64) → UInt64Array"); + for i in 0..arr.len() { + out.push(arr.value(i)); + } + } + let total = total_sink + .lock() + .unwrap() + .expect("observer must have fired at EOS"); + Ok((out, total)) + } + + /// Run one task's local BWAG for `approx_distinct(sn) OVER (ORDER BY sn + /// ROWS UNBOUNDED PRECEDING TO CURRENT ROW)` with no PARTITION BY, and + /// return the single EOS-observed [`Accumulator::state`] Vec. + async fn run_approx_distinct_task( + input: &[u64], + task_ctx: Arc, + ) -> Result> { + use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf; + use std::sync::Mutex; + + /// Observer for `run_approx_distinct_task`: capture the single EOS + /// state. Asserts exactly-one fire and rejects non-empty partition + /// keys (helper is no-PARTITION-BY only). + struct ApproxDistinctObserver { + sink: Arc>>>, + } + + impl WindowStateObserver for ApproxDistinctObserver { + fn finalize_window_aggregate( + &self, + _partition_idx: usize, + _window_expr: &Arc, + partition_key: &PartitionKey, + state: Vec, + ) -> Result<()> { + assert!( + partition_key.is_empty(), + "empty PartitionKey for no-PARTITION-BY plan" + ); + let prev = self.sink.lock().unwrap().replace(state); + assert!(prev.is_none(), "observer must fire exactly once per task"); + Ok(()) + } + } + + let schema = test_schema(); + let mut sn_b = UInt64Builder::with_capacity(input.len()); + let mut hash_b = Int64Builder::with_capacity(input.len()); + for &sn in input { + sn_b.append_value(sn); + hash_b.append_value(0); + } + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![Arc::new(sn_b.finish()), Arc::new(hash_b.finish())], + )?; + let ordering: LexOrdering = [PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }] + .into(); + let source_raw = + TestMemoryExec::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .try_with_sort_information(vec![ordering])?; + let source: Arc = + Arc::new(TestMemoryExec::update_cache(&Arc::new(source_raw))); + + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(approx_distinct_udaf()), + "approx_distinct_sn".to_string(), + &[col("sn", &schema)?], + &[], + &[PhysicalSortExpr { + expr: col("sn", &schema)?, + options: SortOptions::default(), + }], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + source.schema(), + false, + false, + None, + )?; + + let state_sink: Arc>>> = Arc::new(Mutex::new(None)); + let observer: Arc = Arc::new(ApproxDistinctObserver { + sink: Arc::clone(&state_sink), + }); + + let plan = BoundedWindowAggExec::try_new( + vec![expr], + source, + InputOrderMode::Sorted, + false, + )? + .with_state_observer(Some(observer))?; + let _ = collect(Arc::new(plan).execute(0, task_ctx)?).await?; + + state_sink + .lock() + .unwrap() + .take() + .ok_or_else(|| exec_datafusion_err!("observer never fired")) + } + + #[tokio::test] + async fn test_prefix_scan_across_tasks_matches_single_bwag() -> Result<()> { + // Demonstrates the parallel-window shape reviewers asked about: + // range-shuffle `SUM(sn) OVER (ORDER BY sn UNBOUNDED PRECEDING TO + // CURRENT ROW)` across two tasks, then prefix-scan each task's + // finalized state (from the observer) to carry-in the next task's + // rows. Result must match a single BWAG over the concatenated input. + let task_ctx = Arc::new(TaskContext::default()); + + // Two tasks under range partition on sn: + let (task1_out, task1_total) = + run_running_sum_task(&[1, 1, 2, 2, 3, 3, 4, 4], Arc::clone(&task_ctx)) + .await?; + let (task2_out, task2_total) = + run_running_sum_task(&[5, 5, 6, 6, 7, 7, 8, 8], Arc::clone(&task_ctx)) + .await?; + + // Local (uncorrected) outputs and totals — first pass. + assert_eq!(task1_out, vec![1, 2, 4, 6, 9, 12, 16, 20]); + assert_eq!(task1_total, 20); + assert_eq!(task2_out, vec![5, 10, 16, 22, 29, 36, 44, 52]); + assert_eq!(task2_total, 52); + + // Prefix scan over per-task totals → carry-in for each task. Task 0's + // carry-in is 0; task N's carry-in is the sum of tasks [0, N). + let carry_ins = [0u64, task1_total]; + + // Second pass: shift each task's local values by its carry-in. + let task1_final: Vec = task1_out.iter().map(|v| v + carry_ins[0]).collect(); + let task2_final: Vec = task2_out.iter().map(|v| v + carry_ins[1]).collect(); + let parallel_result: Vec = task1_final + .iter() + .chain(task2_final.iter()) + .copied() + .collect(); + + // Oracle: single BWAG over the full concatenated input. + let (single_result, single_total) = run_running_sum_task( + &[1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7, 8, 8], + task_ctx, + ) + .await?; + + assert_eq!( + parallel_result, single_result, + "two-task prefix-scan must match single-BWAG oracle" + ); + // And matches the sequence in the design discussion. + assert_eq!( + single_result, + vec![1, 2, 4, 6, 9, 12, 16, 20, 25, 30, 36, 42, 49, 56, 64, 72] + ); + assert_eq!(single_total, 72); + Ok(()) + } + + #[tokio::test] + async fn test_prefix_merge_across_tasks_approx_distinct() -> Result<()> { + // Load-bearing contract for the parallel-window use case: the state + // exposed by `WindowStateObserver::finalize_window_aggregate` must be + // compatible with `Accumulator::merge_batch` on a fresh accumulator + // of the same UDAF. This is what allows non-decomposable aggregates + // like `approx_distinct` (HLL sketch state) to be prefix-merged + // across shard tasks — the reason we exposed accumulator state at + // all. If this ever breaks, downstream parallel-window work has to + // wait for a public API change. + use arrow::array::{ArrayRef, BinaryArray}; + use arrow::datatypes::FieldRef; + use datafusion_expr::function::AccumulatorArgs; + use datafusion_functions_aggregate::approx_distinct::approx_distinct_udaf; + + let task_ctx = Arc::new(TaskContext::default()); + + // Two tasks with overlapping inputs; concatenated distinct universe + // is {1,2,3,4,5}. + let state1 = + run_approx_distinct_task(&[1, 1, 2, 3], Arc::clone(&task_ctx)).await?; + let state2 = run_approx_distinct_task(&[3, 4, 5], Arc::clone(&task_ctx)).await?; + let state_single = + run_approx_distinct_task(&[1, 1, 2, 3, 3, 4, 5], Arc::clone(&task_ctx)) + .await?; + + // approx_distinct state is a single serialized-HLL Binary field. + assert_eq!(state1.len(), 1, "single state field"); + assert_eq!(state2.len(), 1, "single state field"); + assert_eq!(state_single.len(), 1, "single state field"); + + // Seed a fresh accumulator with the given serialized HLL states via + // `merge_batch` and return its distinct-count evaluation. + fn evaluate_merged(states: &[&ScalarValue]) -> Result { + let udaf = approx_distinct_udaf(); + let input_schema = + Arc::new(Schema::new(vec![Field::new("sn", DataType::UInt64, true)])); + let return_field: FieldRef = + Arc::new(Field::new("approx_distinct_sn", DataType::UInt64, true)); + let expr_field: FieldRef = Arc::new(Field::new("sn", DataType::UInt64, true)); + let physical_col: Arc = col("sn", &input_schema)?; + let args = AccumulatorArgs { + return_field: Arc::clone(&return_field), + schema: &input_schema, + ignore_nulls: false, + order_bys: &[], + is_reversed: false, + name: "approx_distinct", + is_distinct: false, + exprs: std::slice::from_ref(&physical_col), + expr_fields: std::slice::from_ref(&expr_field), + }; + let mut acc = udaf.accumulator(args)?; + let byte_slices: Vec<&[u8]> = states + .iter() + .map(|s| match s { + ScalarValue::Binary(Some(v)) => v.as_slice(), + other => panic!("expected Binary state, got {other:?}"), + }) + .collect(); + let bin: ArrayRef = Arc::new(BinaryArray::from_iter_values(byte_slices)); + acc.merge_batch(std::slice::from_ref(&bin))?; + acc.evaluate() + } + + let merged = evaluate_merged(&[&state1[0], &state2[0]])?; + let oracle = evaluate_merged(&[&state_single[0]])?; + + assert_eq!( + merged, oracle, + "merged task states must match single-BWAG oracle — parallel prefix-merge contract" + ); + // HLL is approximate but exact for a 5-element universe. + assert_eq!(merged, ScalarValue::UInt64(Some(5))); + Ok(()) + } + #[test] fn test_bounded_window_agg_cardinality_effect() -> Result<()> { let schema = test_schema(); @@ -1874,4 +2943,88 @@ mod tests { )); Ok(()) } + + /// Checks the per-partition batches that `LinearSearch` splits an input + /// batch into: partitions appear in first-appearance order, rows within a + /// partition keep their stream order, NULL keys form their own partition, + /// and a single-partition batch is passed through without copying. + #[test] + fn test_linear_search_evaluate_partition_batches() -> Result<()> { + use super::{LinearSearch, PartitionSearcher}; + use arrow::array::{Int32Array, Int64Array}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int32, true), + Field::new("b", DataType::Int64, false), + ])); + let window_expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "count".to_string(), + &[col("b", &schema)?], + &[col("a", &schema)?], + &[], + Arc::new(WindowFrame::new(None)), + Arc::clone(&schema), + false, + false, + None, + )?; + let mut searcher = LinearSearch::new(vec![], Arc::clone(&schema)); + + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![ + Some(1), + Some(2), + Some(1), + None, + Some(2), + Some(1), + ])), + Arc::new(Int64Array::from(vec![10, 20, 11, 30, 21, 12])), + ], + )?; + let result = + searcher.evaluate_partition_batches(&batch, &[Arc::clone(&window_expr)])?; + assert_eq!(result.len(), 3); + let expected = [ + ( + ScalarValue::Int32(Some(1)), + vec![Some(1); 3], + vec![10i64, 11, 12], + ), + (ScalarValue::Int32(Some(2)), vec![Some(2); 2], vec![20, 21]), + (ScalarValue::Int32(None), vec![None], vec![30]), + ]; + for ((key, partition_batch), (exp_key, exp_a, exp_b)) in + result.iter().zip(expected) + { + assert_eq!(key, &vec![exp_key]); + let exp_batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(exp_a)), + Arc::new(Int64Array::from(exp_b)), + ], + )?; + assert_eq!(partition_batch, &exp_batch); + } + + let single = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(Int32Array::from(vec![Some(7), Some(7)])), + Arc::new(Int64Array::from(vec![70, 71])), + ], + )?; + let result = searcher.evaluate_partition_batches(&single, &[window_expr])?; + assert_eq!(result.len(), 1); + assert_eq!(result[0].0, vec![ScalarValue::Int32(Some(7))]); + assert_eq!(result[0].1, single); + // The whole batch belongs to one partition, so its columns are reused + // rather than gathered into a new batch. + assert!(Arc::ptr_eq(result[0].1.column(0), single.column(0))); + Ok(()) + } } diff --git a/datafusion/physical-plan/src/windows/mod.rs b/datafusion/physical-plan/src/windows/mod.rs index b72a65cf996be..089bdc23ee2c4 100644 --- a/datafusion/physical-plan/src/windows/mod.rs +++ b/datafusion/physical-plan/src/windows/mod.rs @@ -18,6 +18,8 @@ //! Physical expressions for window functions mod bounded_window_agg_exec; +#[cfg(feature = "proto")] +mod proto; mod utils; mod window_agg_exec; @@ -52,7 +54,7 @@ use datafusion_physical_expr_common::sort_expr::{ use itertools::Itertools; // Public interface: -pub use bounded_window_agg_exec::BoundedWindowAggExec; +pub use bounded_window_agg_exec::{BoundedWindowAggExec, WindowStateObserver}; pub use datafusion_physical_expr::window::{ PlainAggregateWindowExpr, StandardWindowExpr, WindowExpr, }; @@ -594,6 +596,12 @@ pub fn get_best_fitting_window( // They are either the same with `window_expr`'s PARTITION BY columns, // or it is empty if partitioning is not desirable for this windowing operator. physical_partition_keys: &[Arc], + // A [`WindowStateObserver`] installed on the source + // [`BoundedWindowAggExec`] (via [`BoundedWindowAggExec::with_state_observer`]) + // that must survive the rebuild. Ignored when the rebuilt exec is a + // [`WindowAggExec`], which does not carry an observer. `None` when the + // source is a [`WindowAggExec`] or has no observer installed. + state_observer: Option>, ) -> Result>> { // Contains at least one window expr and all of the partition by and order by sections // of the window_exprs are same. @@ -633,12 +641,15 @@ pub fn get_best_fitting_window( // If all window expressions can run with bounded memory, choose the // bounded window variant: if window_expr.iter().all(|e| e.uses_bounded_memory()) { - Ok(Some(Arc::new(BoundedWindowAggExec::try_new( - window_expr, - Arc::clone(input), - input_order_mode, - !physical_partition_keys.is_empty(), - )?) as _)) + Ok(Some(Arc::new( + BoundedWindowAggExec::try_new( + window_expr, + Arc::clone(input), + input_order_mode, + !physical_partition_keys.is_empty(), + )? + .with_state_observer(state_observer)?, + ) as _)) } else if input_order_mode != InputOrderMode::Sorted { // For `WindowAggExec` to work correctly PARTITION BY columns should be sorted. // Hence, if `input_order_mode` is not `Sorted` we should convert @@ -937,6 +948,79 @@ mod tests { Ok(()) } + #[tokio::test] + async fn get_best_fitting_window_preserves_state_observer() -> Result<()> { + // `EnforceSorting`/`EnforceDistribution` call `get_best_fitting_window` + // on a source `BoundedWindowAggExec` and replace it with the returned + // exec. Without observer propagation, a `WindowStateObserver` + // installed on the source is silently dropped by the rebuild. + use datafusion_common::ScalarValue; + use datafusion_expr::{WindowFrameBound, WindowFrameUnits}; + + struct NoopObserver; + impl WindowStateObserver for NoopObserver { + fn finalize_window_aggregate( + &self, + _partition_idx: usize, + _window_expr: &Arc, + _partition_key: &datafusion_physical_expr::window::PartitionKey, + _state: Vec, + ) -> Result<()> { + Ok(()) + } + } + + let schema = create_test_schema()?; + let sort = sort_expr("nullable_col", &schema); + let ordering: LexOrdering = [sort.clone()].into(); + let source = streaming_table_exec(&schema, ordering, false)?; + + let expr = create_window_expr( + &WindowFunctionDefinition::AggregateUDF(count_udaf()), + "cnt".to_string(), + &[col("nullable_col", &schema)?], + &[], + &[sort], + Arc::new(WindowFrame::new_bounds( + WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(None)), + WindowFrameBound::CurrentRow, + )), + source.schema(), + false, + false, + None, + )?; + + let observer: Arc = Arc::new(NoopObserver); + let bounded = BoundedWindowAggExec::try_new( + vec![expr], + Arc::clone(&source), + Sorted, + false, + )? + .with_state_observer(Some(Arc::clone(&observer)))?; + + let rebuilt = get_best_fitting_window( + bounded.window_expr(), + bounded.input(), + &bounded.partition_keys(), + bounded.state_observer().cloned(), + )? + .expect("rebuild should produce a plan"); + let bwag = rebuilt + .downcast_ref::() + .expect("rebuild yielded BoundedWindowAggExec"); + let installed = bwag + .state_observer() + .expect("observer preserved through rebuild"); + assert!( + Arc::ptr_eq(installed, &observer), + "observer identity preserved through rebuild", + ); + Ok(()) + } + #[tokio::test] async fn test_satisfy_nullable() -> Result<()> { let schema = create_test_schema()?; diff --git a/datafusion/physical-plan/src/windows/proto.rs b/datafusion/physical-plan/src/windows/proto.rs new file mode 100644 index 0000000000000..e96b0a9fb1087 --- /dev/null +++ b/datafusion/physical-plan/src/windows/proto.rs @@ -0,0 +1,263 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Protobuf conversions shared by window execution plans. + +use std::sync::Arc; + +use arrow::datatypes::Schema; +use datafusion_common::{ + Result, ScalarValue, internal_datafusion_err, internal_err, not_impl_err, +}; +use datafusion_expr::{ + WindowFrame, WindowFrameBound, WindowFrameUnits, WindowFunctionDefinition, +}; +use datafusion_physical_expr::window::SlidingAggregateWindowExpr; +use datafusion_physical_expr_common::sort_expr::{ + sort_exprs_try_from_proto, sort_exprs_try_to_proto, +}; +use datafusion_proto_common::protobuf_common; +use datafusion_proto_models::protobuf::{self, physical_window_expr_node}; + +use super::{ + PlainAggregateWindowExpr, StandardWindowExpr, WindowExpr, WindowUDFExpr, + create_window_expr, schema_add_window_field, +}; + +pub(super) fn encode_physical_window_expr( + window_expr: &Arc, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, +) -> Result { + let expr = window_expr.as_any(); + let mut args = window_expr.expressions().to_vec(); + let window_frame = window_expr.get_window_frame(); + let (window_function, fun_definition, ignore_nulls, distinct) = + if let Some(plain) = expr.downcast_ref::() { + let aggregate_expr = plain.get_aggregate_expr(); + ( + physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + aggregate_expr.fun().name().to_string(), + ), + ctx.encode_udaf(aggregate_expr.fun())?, + aggregate_expr.ignore_nulls(), + aggregate_expr.is_distinct(), + ) + } else if let Some(sliding) = expr.downcast_ref::() { + let aggregate_expr = sliding.get_aggregate_expr(); + ( + physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + aggregate_expr.fun().name().to_string(), + ), + ctx.encode_udaf(aggregate_expr.fun())?, + aggregate_expr.ignore_nulls(), + aggregate_expr.is_distinct(), + ) + } else if let Some(standard) = expr.downcast_ref::() { + if let Some(window_udf) = standard + .get_standard_func_expr() + .as_any() + .downcast_ref::() + { + // `WindowUDFExpr::args` returns the full, unfiltered argument list so + // every argument survives the round-trip. + args = window_udf.args().to_vec(); + ( + physical_window_expr_node::WindowFunction::UserDefinedWindowFunction( + window_udf.fun().name().to_string(), + ), + ctx.encode_udwf(window_udf.fun().as_ref())?, + false, + false, + ) + } else { + return not_impl_err!( + "User-defined window function not supported: {window_expr:?}" + ); + } + } else { + return not_impl_err!("WindowExpr not supported: {window_expr:?}"); + }; + + let args = ctx.encode_expressions(&args)?; + let partition_by = ctx.encode_expressions(window_expr.partition_by())?; + let order_by = sort_exprs_try_to_proto(window_expr.order_by(), &ctx.expr_ctx())?; + + Ok(protobuf::PhysicalWindowExprNode { + args, + partition_by, + order_by, + window_frame: Some(encode_window_frame(window_frame.as_ref())?), + window_function: Some(window_function), + name: window_expr.name().to_string(), + fun_definition, + ignore_nulls, + distinct, + }) +} + +pub(super) fn decode_physical_window_expr( + proto: &protobuf::PhysicalWindowExprNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + input_schema: &Schema, +) -> Result> { + let args = proto + .args + .iter() + .map(|expr| ctx.decode_expr(expr, input_schema)) + .collect::>>()?; + let partition_by = proto + .partition_by + .iter() + .map(|expr| ctx.decode_expr(expr, input_schema)) + .collect::>>()?; + let order_by = + sort_exprs_try_from_proto(&proto.order_by, &ctx.expr_ctx(input_schema))?; + let window_frame = proto + .window_frame + .as_ref() + .map(decode_window_frame) + .transpose()? + .ok_or_else(|| { + internal_datafusion_err!("Missing required field 'window_frame' in protobuf") + })?; + let function = match proto.window_function.as_ref() { + Some(physical_window_expr_node::WindowFunction::UserDefinedAggrFunction( + name, + )) => WindowFunctionDefinition::AggregateUDF( + ctx.decode_udaf(name, proto.fun_definition.as_deref())?, + ), + Some(physical_window_expr_node::WindowFunction::UserDefinedWindowFunction( + name, + )) => WindowFunctionDefinition::WindowUDF( + ctx.decode_udwf(name, proto.fun_definition.as_deref())?, + ), + None => { + return internal_err!("Missing required field 'window_function' in protobuf"); + } + }; + + let name = proto.name.clone(); + // TODO: Remove extended_schema if functions are all UDAF + let extended_schema = schema_add_window_field(&args, input_schema, &function, &name)?; + create_window_expr( + &function, + name, + &args, + &partition_by, + &order_by, + Arc::new(window_frame), + extended_schema, + proto.ignore_nulls, + proto.distinct, + None, + ) +} + +fn encode_window_frame(window_frame: &WindowFrame) -> Result { + let units = match window_frame.units { + WindowFrameUnits::Rows => protobuf::WindowFrameUnits::Rows, + WindowFrameUnits::Range => protobuf::WindowFrameUnits::Range, + WindowFrameUnits::Groups => protobuf::WindowFrameUnits::Groups, + }; + Ok(protobuf::WindowFrame { + window_frame_units: units.into(), + start_bound: Some(encode_window_frame_bound(&window_frame.start_bound)?), + end_bound: Some(protobuf::window_frame::EndBound::Bound( + encode_window_frame_bound(&window_frame.end_bound)?, + )), + }) +} + +fn encode_window_frame_bound( + bound: &WindowFrameBound, +) -> Result { + let encode_value = |value: &ScalarValue| -> Result { + Ok(value.try_into()?) + }; + Ok(match bound { + WindowFrameBound::CurrentRow => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::CurrentRow.into(), + bound_value: None, + }, + WindowFrameBound::Preceding(value) => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::Preceding.into(), + bound_value: Some(encode_value(value)?), + }, + WindowFrameBound::Following(value) => protobuf::WindowFrameBound { + window_frame_bound_type: protobuf::WindowFrameBoundType::Following.into(), + bound_value: Some(encode_value(value)?), + }, + }) +} + +fn decode_window_frame(window_frame: &protobuf::WindowFrame) -> Result { + let units = protobuf::WindowFrameUnits::try_from(window_frame.window_frame_units) + .map_err(|_| { + internal_datafusion_err!( + "Received a WindowFrame message with unknown WindowFrameUnits {}", + window_frame.window_frame_units + ) + })?; + let units = match units { + protobuf::WindowFrameUnits::Rows => WindowFrameUnits::Rows, + protobuf::WindowFrameUnits::Range => WindowFrameUnits::Range, + protobuf::WindowFrameUnits::Groups => WindowFrameUnits::Groups, + }; + let start_bound = + decode_window_frame_bound(window_frame.start_bound.as_ref().ok_or_else( + || internal_datafusion_err!("Missing start_bound in WindowFrame"), + )?)?; + let end_bound = window_frame + .end_bound + .as_ref() + .map(|end_bound| match end_bound { + protobuf::window_frame::EndBound::Bound(bound) => { + decode_window_frame_bound(bound) + } + }) + .transpose()? + .unwrap_or(WindowFrameBound::CurrentRow); + Ok(WindowFrame::new_bounds(units, start_bound, end_bound)) +} + +fn decode_window_frame_bound( + bound: &protobuf::WindowFrameBound, +) -> Result { + let decode_value = |value: &protobuf_common::ScalarValue| -> Result { + Ok(ScalarValue::try_from(value)?) + }; + let bound_type = protobuf::WindowFrameBoundType::try_from( + bound.window_frame_bound_type, + ) + .map_err(|_| { + internal_datafusion_err!( + "Received a WindowFrameBound message with unknown WindowFrameBoundType {}", + bound.window_frame_bound_type + ) + })?; + match bound_type { + protobuf::WindowFrameBoundType::CurrentRow => Ok(WindowFrameBound::CurrentRow), + protobuf::WindowFrameBoundType::Preceding => match &bound.bound_value { + Some(value) => Ok(WindowFrameBound::Preceding(decode_value(value)?)), + None => Ok(WindowFrameBound::Preceding(ScalarValue::UInt64(None))), + }, + protobuf::WindowFrameBoundType::Following => match &bound.bound_value { + Some(value) => Ok(WindowFrameBound::Following(decode_value(value)?)), + None => Ok(WindowFrameBound::Following(ScalarValue::UInt64(None))), + }, + } +} diff --git a/datafusion/physical-plan/src/windows/window_agg_exec.rs b/datafusion/physical-plan/src/windows/window_agg_exec.rs index 9e8fc8a6ebb62..d794e7df9d0a9 100644 --- a/datafusion/physical-plan/src/windows/window_agg_exec.rs +++ b/datafusion/physical-plan/src/windows/window_agg_exec.rs @@ -21,18 +21,22 @@ use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +#[cfg(feature = "proto")] +use super::proto::{decode_physical_window_expr, encode_physical_window_expr}; use super::utils::create_schema; use crate::execution_plan::{CardinalityEffect, EmissionType}; use crate::metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet}; +use crate::statistics::{ChildStats, StatisticsArgs}; use crate::stream::EmptyRecordBatchStream; use crate::windows::{ calc_requirements, get_ordered_partition_by_indices, get_partition_by_sort_exprs, window_equivalence_properties, }; use crate::{ - ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, ExecutionPlan, - ExecutionPlanProperties, PhysicalExpr, PlanProperties, RecordBatchStream, - SendableRecordBatchStream, Statistics, WindowExpr, check_if_same_properties, + ChildrenPropertiesMode, ColumnStatistics, DisplayAs, DisplayFormatType, Distribution, + ExecutionPlan, ExecutionPlanProperties, InputDistributionRequirements, PhysicalExpr, + PlanProperties, RecordBatchStream, ReplaceChildrenOptions, SendableRecordBatchStream, + Statistics, WindowExpr, validate_child_count, }; use arrow::array::ArrayRef; @@ -159,17 +163,6 @@ impl WindowAggExec { .unwrap_or_else(Vec::new) } } - - fn with_new_children_and_same_properties( - &self, - mut children: Vec>, - ) -> Self { - Self { - input: children.swap_remove(0), - metrics: ExecutionPlanMetricsSet::new(), - ..Self::clone(self) - } - } } impl DisplayAs for WindowAggExec { @@ -224,15 +217,17 @@ impl ExecutionPlan for WindowAggExec { fn apply_expressions( &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for window_expr in &self.window_expr { - for expr in window_expr.expressions() { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - } - Ok(tnr) + let expressions = self.window_expr.iter().flat_map(|window_expr| { + let expressions = window_expr.all_expressions(); + expressions + .args + .into_iter() + .chain(expressions.partition_by_exprs) + .chain(expressions.order_by_exprs) + }); + crate::apply_expression_roots(expressions, f) } fn maintains_input_order(&self) -> Vec { @@ -254,23 +249,57 @@ impl ExecutionPlan for WindowAggExec { } fn required_input_distribution(&self) -> Vec { + self.input_distribution_requirements().into_per_child() + } + + fn input_distribution_requirements(&self) -> InputDistributionRequirements { if self.partition_keys().is_empty() { - vec![Distribution::SinglePartition] + InputDistributionRequirements::new(vec![Distribution::SinglePartition]) } else { - vec![Distribution::HashPartitioned(self.partition_keys())] + InputDistributionRequirements::new(vec![Distribution::KeyPartitioned( + self.partition_keys(), + )]) } } - fn with_new_children( + fn replace_children( self: Arc, mut children: Vec>, + options: ReplaceChildrenOptions, ) -> Result> { - check_if_same_properties!(self, children); - Ok(Arc::new(WindowAggExec::try_new( - self.window_expr.clone(), - children.swap_remove(0), - true, - )?)) + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => Ok(Arc::new(WindowAggExec::try_new( + self.window_expr.clone(), + children.swap_remove(0), + true, + )?)), + } + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn with_new_children_and_same_properties( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Keep), + ) } fn execute( @@ -294,9 +323,16 @@ impl ExecutionPlan for WindowAggExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, partition: Option) -> Result> { - let input_stat = - Arc::unwrap_or_clone(self.input.partition_statistics(partition)?); + fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] + } + + fn statistics_from_inputs( + &self, + input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { + let input_stat = input_stats[0].as_ref().clone(); let win_cols = self.window_expr.len(); let input_cols = self.input.schema().fields().len(); // TODO stats: some windowing function will maintain invariants such as min, max... @@ -316,6 +352,133 @@ impl ExecutionPlan for WindowAggExec { fn cardinality_effect(&self) -> CardinalityEffect { CardinalityEffect::Equal } + + #[cfg(feature = "proto")] + fn try_to_proto( + &self, + ctx: &crate::proto::ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + use datafusion_proto_models::protobuf; + + // Exhaustive destructure: adding a field to `WindowAggExec` without + // deciding how it is serialized is a compile error, not a silent + // round-trip gap. + let Self { + input, + window_expr, + // Derived at construction by `create_schema` from the input schema + // and the window expressions. + schema: _, + // Runtime execution state, rebuilt empty on decode. + metrics: _, + // Derived at construction by `get_ordered_partition_by_indices`. + ordered_partition_by_indices: _, + // Derived at construction by `Self::compute_properties`. + cache: _, + // No wire field of its own; it is folded into `partition_keys` + // below, since `partition_keys()` returns an empty vec when this is + // false and the decoder recovers it as `!partition_keys.is_empty()`. + can_repartition: _, + } = self; + + let input = ctx.encode_child(input)?; + let window_expr = window_expr + .iter() + .map(|expr| encode_physical_window_expr(expr, ctx)) + .collect::>>()?; + let partition_keys = self + .partition_keys() + .iter() + .map(|expr| ctx.encode_expr(expr)) + .collect::>>()?; + + Ok(Some(protobuf::PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Window(Box::new( + protobuf::WindowAggExecNode { + input: Some(Box::new(input)), + window_expr, + partition_keys, + // `None` distinguishes a `WindowAggExec` from a + // `BoundedWindowAggExec` on the shared `Window` variant. + input_order_mode: None, + }, + )), + ), + })) + } +} + +#[cfg(feature = "proto")] +impl WindowAggExec { + /// Reconstruct a window plan from its protobuf representation. + /// + /// This returns a [`WindowAggExec`] when `input_order_mode` is absent and a + /// [`BoundedWindowAggExec`] when it is present. + /// + /// [`BoundedWindowAggExec`]: crate::windows::BoundedWindowAggExec + pub fn try_from_proto( + node: &datafusion_proto_models::protobuf::PhysicalPlanNode, + ctx: &crate::proto::ExecutionPlanDecodeCtx<'_>, + ) -> Result> { + use super::BoundedWindowAggExec; + use crate::InputOrderMode; + use datafusion_proto_models::protobuf; + use protobuf::window_agg_exec_node::InputOrderMode as ProtoInputOrderMode; + + let window_agg = crate::expect_plan_variant!( + node, + protobuf::physical_plan_node::PhysicalPlanType::Window, + "WindowAggExec", + ); + // Exhaustive destructure: a new field on `WindowAggExecNode` is a + // compile error here rather than a silently ignored wire field. + let protobuf::WindowAggExecNode { + input, + window_expr, + partition_keys, + input_order_mode, + } = window_agg.as_ref(); + + let input = + ctx.decode_required_child(input.as_deref(), "WindowAggExec", "input")?; + let input_schema = input.schema(); + let window_expr = window_expr + .iter() + .map(|expr| decode_physical_window_expr(expr, ctx, input_schema.as_ref())) + .collect::>>()?; + let partition_keys = partition_keys + .iter() + .map(|expr| ctx.decode_expr(expr, input_schema.as_ref())) + .collect::>>()?; + + if let Some(input_order_mode) = input_order_mode.as_ref() { + let input_order_mode = match input_order_mode { + ProtoInputOrderMode::Linear(_) => InputOrderMode::Linear, + ProtoInputOrderMode::PartiallySorted( + protobuf::PartiallySortedInputOrderMode { columns }, + ) => InputOrderMode::PartiallySorted( + columns.iter().map(|column| *column as usize).collect(), + ), + ProtoInputOrderMode::Sorted(_) => InputOrderMode::Sorted, + }; + Ok(Arc::new(BoundedWindowAggExec::try_new( + window_expr, + input, + input_order_mode, + // `can_repartition` has no wire field: the encoder writes an + // empty `partition_keys` when it is false. + !partition_keys.is_empty(), + )?)) + } else { + Ok(Arc::new(WindowAggExec::try_new( + window_expr, + input, + // See above: `can_repartition` is recovered from `partition_keys`. + !partition_keys.is_empty(), + )?)) + } + } } /// Compute the window aggregate columns diff --git a/datafusion/physical-plan/src/work_table.rs b/datafusion/physical-plan/src/work_table.rs index 0855dbf2fd635..b5d6fd47bc465 100644 --- a/datafusion/physical-plan/src/work_table.rs +++ b/datafusion/physical-plan/src/work_table.rs @@ -25,10 +25,11 @@ use crate::execution_plan::{Boundedness, EmissionType, SchedulingType}; use crate::memory::MemoryStream; use crate::metrics::{ExecutionPlanMetricsSet, MetricsSet}; use crate::{ - DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, - SendableRecordBatchStream, Statistics, + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PlanProperties, + ReplaceChildrenOptions, SendableRecordBatchStream, Statistics, }; +use crate::statistics::StatisticsArgs; use arrow::datatypes::SchemaRef; use arrow::record_batch::RecordBatch; use datafusion_common::tree_node::TreeNodeRecursion; @@ -188,18 +189,29 @@ impl ExecutionPlan for WorkTableExec { fn apply_expressions( &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, + _f: &mut dyn FnMut(&Arc) -> Result, ) -> Result { Ok(TreeNodeRecursion::Continue) } - fn with_new_children( + fn replace_children( self: Arc, _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(Arc::clone(&self) as Arc) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + /// Stream the batches that were written to the work table. fn execute( &self, @@ -235,7 +247,11 @@ impl ExecutionPlan for WorkTableExec { Some(self.metrics.clone_inner()) } - fn partition_statistics(&self, _partition: Option) -> Result> { + fn statistics_from_inputs( + &self, + _input_stats: &[Arc], + _args: &StatisticsArgs, + ) -> Result> { Ok(Arc::new(Statistics::new_unknown(&self.schema()))) } diff --git a/datafusion/proto-common/Cargo.toml b/datafusion/proto-common/Cargo.toml index 46dae36ba40ed..0670d7cbf757f 100644 --- a/datafusion/proto-common/Cargo.toml +++ b/datafusion/proto-common/Cargo.toml @@ -31,6 +31,12 @@ rust-version = { workspace = true } [package.metadata.docs.rs] all-features = true +# Note: add additional linter rules in lib.rs. +# Rust does not support workspace + new linter rules in subcrates yet +# https://github.com/rust-lang/cargo/issues/13157 +[lints] +workspace = true + [lib] name = "datafusion_proto_common" diff --git a/datafusion/proto-common/gen/Cargo.toml b/datafusion/proto-common/gen/Cargo.toml index f0e60819d42a8..0cbba311b2c3f 100644 --- a/datafusion/proto-common/gen/Cargo.toml +++ b/datafusion/proto-common/gen/Cargo.toml @@ -38,4 +38,4 @@ workspace = true [dependencies] # Pin these dependencies so that the generated output is deterministic pbjson-build = "=0.9.0" -prost-build = "=0.14.3" +prost-build = "=0.14.4" diff --git a/datafusion/proto-common/gen/src/main.rs b/datafusion/proto-common/gen/src/main.rs index 02e1ecf00bab8..d672832d43897 100644 --- a/datafusion/proto-common/gen/src/main.rs +++ b/datafusion/proto-common/gen/src/main.rs @@ -33,14 +33,12 @@ fn main() -> Result<(), String> { .map_err(|e| format!("protobuf compilation failed: {e}"))?; let descriptor_set = std::fs::read(&descriptor_path) - .unwrap_or_else(|e| panic!("Cannot read {:?}: {}", &descriptor_path, e)); + .unwrap_or_else(|e| panic!("Cannot read {descriptor_path:?}: {e}")); pbjson_build::Builder::new() .out_dir("src") .register_descriptors(&descriptor_set) - .unwrap_or_else(|e| { - panic!("Cannot register descriptors {:?}: {}", &descriptor_set, e) - }) + .unwrap_or_else(|e| panic!("Cannot register descriptors {descriptor_set:?}: {e}")) .build(&[".datafusion_common"]) .map_err(|e| format!("pbjson compilation failed: {e}"))?; diff --git a/datafusion/proto-common/proto/datafusion_common.proto b/datafusion/proto-common/proto/datafusion_common.proto index 684d9a2612408..27d1101036d9b 100644 --- a/datafusion/proto-common/proto/datafusion_common.proto +++ b/datafusion/proto-common/proto/datafusion_common.proto @@ -617,6 +617,8 @@ message ParquetOptions { uint64 max_row_group_size = 15; + uint64 max_in_list_size = 38; + string created_by = 16; oneof coerce_int96_opt { @@ -627,7 +629,11 @@ message ParquetOptions { uint64 max_predicate_cache_size = 33; } - CdcOptions content_defined_chunking = 35; + oneof max_row_group_bytes_opt { + uint64 max_row_group_bytes = 37; + } + + ParquetCdcOptions content_defined_chunking = 35; // Optional timezone applied to INT96-coerced timestamps when `coerce_int96` // is set. When `Some`, INT96 columns coerce to @@ -638,10 +644,12 @@ message ParquetOptions { } } -message CdcOptions { - uint64 min_chunk_size = 1; - uint64 max_chunk_size = 2; - int32 norm_level = 3; +// Content-defined chunking (CDC) options for writing parquet files. +message ParquetCdcOptions { + bool enabled = 1; + uint64 min_chunk_size = 2; + uint64 max_chunk_size = 3; + int32 norm_level = 4; } enum JoinSide { @@ -683,4 +691,30 @@ enum ExplainFormat { EXPLAIN_FORMAT_TREE = 1; EXPLAIN_FORMAT_PGJSON = 2; EXPLAIN_FORMAT_GRAPHVIZ = 3; +} + +// Verbosity level for `EXPLAIN ANALYZE`. Mirrors +// `datafusion_common::format::MetricType`. +enum MetricType { + METRIC_TYPE_SUMMARY = 0; + METRIC_TYPE_DEV = 1; +} + +// Category of an `EXPLAIN ANALYZE` metric. Mirrors +// `datafusion_common::format::MetricCategory`. +enum MetricCategory { + METRIC_CATEGORY_ROWS = 0; + METRIC_CATEGORY_BYTES = 1; + METRIC_CATEGORY_TIMING = 2; + METRIC_CATEGORY_UNCATEGORIZED = 3; +} + +// Wire encoding for `datafusion_common::format::ExplainAnalyzeCategories`. +// +// If `all` is true, every category is shown (the `only` list is ignored). +// If `all` is false, only the categories listed in `only` are shown — an +// empty `only` means "plan only", i.e. suppress all metrics. +message ExplainAnalyzeCategoriesNode { + bool all = 1; + repeated MetricCategory only = 2; } \ No newline at end of file diff --git a/datafusion/proto-common/src/from_proto/mod.rs b/datafusion/proto-common/src/from_proto/mod.rs index 94a06bcc13bbd..169ff7f3d9ff2 100644 --- a/datafusion/proto-common/src/from_proto/mod.rs +++ b/datafusion/proto-common/src/from_proto/mod.rs @@ -39,8 +39,8 @@ use datafusion_common::{ DataFusionError, JoinSide, ScalarValue, Statistics, TableReference, arrow_datafusion_err, config::{ - CdcOptions, CsvOptions, JsonOptions, ParquetColumnOptions, ParquetOptions, - TableParquetOptions, + CsvOptions, JsonOptions, MaxRowGroupBytes, ParquetCdcOptions, + ParquetColumnOptions, ParquetOptions, TableParquetOptions, }, file_options::{csv_writer::CsvWriterOptions, json_writer::JsonWriterOptions}, parsers::CompressionTypeVariant, @@ -1081,6 +1081,7 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions { }) .unwrap_or(None), max_row_group_size: value.max_row_group_size as usize, + max_in_list_size: value.max_in_list_size as usize, created_by: value.created_by.clone(), column_index_truncate_length: value .column_index_truncate_length_opt.as_ref() @@ -1130,21 +1131,25 @@ impl TryFrom<&protobuf::ParquetOptions> for ParquetOptions { max_predicate_cache_size: value.max_predicate_cache_size_opt.map(|opt| match opt { protobuf::parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(v) => Some(v as usize), }).unwrap_or(None), - use_content_defined_chunking: value.content_defined_chunking.map(|cdc| { - let defaults = CdcOptions::default(); - CdcOptions { - // proto3 uses 0 as the wire default for uint64; a zero chunk size is - // invalid, so treat it as "field not set" and fall back to the default. - min_chunk_size: if cdc.min_chunk_size != 0 { cdc.min_chunk_size as usize } else { defaults.min_chunk_size }, - max_chunk_size: if cdc.max_chunk_size != 0 { cdc.max_chunk_size as usize } else { defaults.max_chunk_size }, - // norm_level = 0 is a valid value (and the default), so pass it through directly. - norm_level: cdc.norm_level, - } + max_row_group_bytes: value.max_row_group_bytes_opt.and_then(|opt| match opt { + protobuf::parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(v) => MaxRowGroupBytes::try_new(v as usize).ok(), }), + content_defined_chunking: value.content_defined_chunking.map(ParquetCdcOptions::from).unwrap_or_default(), }) } } +impl From for ParquetCdcOptions { + fn from(value: protobuf::ParquetCdcOptions) -> Self { + ParquetCdcOptions { + enabled: value.enabled, + min_chunk_size: value.min_chunk_size as usize, + max_chunk_size: value.max_chunk_size as usize, + norm_level: value.norm_level, + } + } +} + impl TryFrom<&protobuf::ParquetColumnOptions> for ParquetColumnOptions { type Error = DataFusionError; fn try_from( @@ -1255,9 +1260,7 @@ fn vec_to_array(v: Vec) -> [T; N] { } /// Converts a vector of `protobuf::Field`s to `Arc`s. -pub fn parse_proto_fields_to_fields<'a, I>( - fields: I, -) -> std::result::Result, Error> +pub fn parse_proto_fields_to_fields<'a, I>(fields: I) -> Result, Error> where I: IntoIterator, { @@ -1329,7 +1332,9 @@ pub(crate) fn csv_writer_options_from_proto( #[cfg(test)] mod tests { - use datafusion_common::config::{CdcOptions, ParquetOptions, TableParquetOptions}; + use datafusion_common::config::{ + MaxRowGroupBytes, ParquetCdcOptions, ParquetOptions, TableParquetOptions, + }; fn parquet_options_proto_round_trip(opts: ParquetOptions) -> ParquetOptions { let proto: crate::protobuf_common::ParquetOptions = @@ -1348,7 +1353,7 @@ mod tests { #[test] fn test_parquet_options_cdc_disabled_round_trip() { let opts = ParquetOptions::default(); - assert!(opts.use_content_defined_chunking.is_none()); + assert!(!opts.content_defined_chunking.enabled); let recovered = parquet_options_proto_round_trip(opts.clone()); assert_eq!(opts, recovered); } @@ -1373,6 +1378,21 @@ mod tests { assert_eq!(recovered.coerce_int96_tz, Some("UTC".to_string())); } + #[test] + fn test_parquet_options_max_row_group_bytes_round_trip() { + let opts = ParquetOptions { + max_row_group_bytes: Some( + MaxRowGroupBytes::try_new(64 * 1024 * 1024).unwrap(), + ), + ..ParquetOptions::default() + }; + let recovered = parquet_options_proto_round_trip(opts.clone()); + assert_eq!( + recovered.max_row_group_bytes.map(|v| v.get()), + Some(64 * 1024 * 1024) + ); + } + #[test] fn test_table_parquet_options_coerce_int96_tz_round_trip() { let mut opts = TableParquetOptions::default(); @@ -1389,15 +1409,17 @@ mod tests { #[test] fn test_parquet_options_cdc_enabled_round_trip() { let opts = ParquetOptions { - use_content_defined_chunking: Some(CdcOptions { + content_defined_chunking: ParquetCdcOptions { + enabled: true, min_chunk_size: 128 * 1024, max_chunk_size: 512 * 1024, norm_level: 2, - }), + }, ..ParquetOptions::default() }; let recovered = parquet_options_proto_round_trip(opts.clone()); - let cdc = recovered.use_content_defined_chunking.unwrap(); + let cdc = recovered.content_defined_chunking; + assert!(cdc.enabled); assert_eq!(cdc.min_chunk_size, 128 * 1024); assert_eq!(cdc.max_chunk_size, 512 * 1024); assert_eq!(cdc.norm_level, 2); @@ -1406,30 +1428,30 @@ mod tests { #[test] fn test_parquet_options_cdc_negative_norm_level_round_trip() { let opts = ParquetOptions { - use_content_defined_chunking: Some(CdcOptions { + content_defined_chunking: ParquetCdcOptions { + enabled: true, norm_level: -3, - ..CdcOptions::default() - }), + ..ParquetCdcOptions::default() + }, ..ParquetOptions::default() }; let recovered = parquet_options_proto_round_trip(opts); - assert_eq!( - recovered.use_content_defined_chunking.unwrap().norm_level, - -3 - ); + assert_eq!(recovered.content_defined_chunking.norm_level, -3); } #[test] fn test_table_parquet_options_cdc_round_trip() { let mut opts = TableParquetOptions::default(); - opts.global.use_content_defined_chunking = Some(CdcOptions { + opts.global.content_defined_chunking = ParquetCdcOptions { + enabled: true, min_chunk_size: 64 * 1024, max_chunk_size: 2 * 1024 * 1024, norm_level: -1, - }); + }; let recovered = table_parquet_options_proto_round_trip(opts.clone()); - let cdc = recovered.global.use_content_defined_chunking.unwrap(); + let cdc = recovered.global.content_defined_chunking; + assert!(cdc.enabled); assert_eq!(cdc.min_chunk_size, 64 * 1024); assert_eq!(cdc.max_chunk_size, 2 * 1024 * 1024); assert_eq!(cdc.norm_level, -1); @@ -1438,8 +1460,8 @@ mod tests { #[test] fn test_table_parquet_options_cdc_disabled_round_trip() { let opts = TableParquetOptions::default(); - assert!(opts.global.use_content_defined_chunking.is_none()); + assert!(!opts.global.content_defined_chunking.enabled); let recovered = table_parquet_options_proto_round_trip(opts.clone()); - assert!(recovered.global.use_content_defined_chunking.is_none()); + assert!(!recovered.global.content_defined_chunking.enabled); } } diff --git a/datafusion/proto-common/src/generated/mod.rs b/datafusion/proto-common/src/generated/mod.rs index 9c2ca9385aa5e..49d09bf3f432b 100644 --- a/datafusion/proto-common/src/generated/mod.rs +++ b/datafusion/proto-common/src/generated/mod.rs @@ -18,6 +18,8 @@ // This code is generated so we don't want to fix any lint violations manually #[allow(clippy::allow_attributes)] #[allow(clippy::all)] +#[allow(clippy::uninlined_format_args)] +#[allow(unused_qualifications)] #[rustfmt::skip] pub mod datafusion_proto_common { include!("prost.rs"); diff --git a/datafusion/proto-common/src/generated/pbjson.rs b/datafusion/proto-common/src/generated/pbjson.rs index 0568982e97a44..c222cd1cb8687 100644 --- a/datafusion/proto-common/src/generated/pbjson.rs +++ b/datafusion/proto-common/src/generated/pbjson.rs @@ -911,144 +911,6 @@ impl<'de> serde::Deserialize<'de> for AvroOptions { deserializer.deserialize_struct("datafusion_common.AvroOptions", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for CdcOptions { - #[allow(deprecated)] - fn serialize(&self, serializer: S) -> std::result::Result - where - S: serde::Serializer, - { - use serde::ser::SerializeStruct; - let mut len = 0; - if self.min_chunk_size != 0 { - len += 1; - } - if self.max_chunk_size != 0 { - len += 1; - } - if self.norm_level != 0 { - len += 1; - } - let mut struct_ser = serializer.serialize_struct("datafusion_common.CdcOptions", len)?; - if self.min_chunk_size != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("minChunkSize", ToString::to_string(&self.min_chunk_size).as_str())?; - } - if self.max_chunk_size != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("maxChunkSize", ToString::to_string(&self.max_chunk_size).as_str())?; - } - if self.norm_level != 0 { - struct_ser.serialize_field("normLevel", &self.norm_level)?; - } - struct_ser.end() - } -} -impl<'de> serde::Deserialize<'de> for CdcOptions { - #[allow(deprecated)] - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - const FIELDS: &[&str] = &[ - "min_chunk_size", - "minChunkSize", - "max_chunk_size", - "maxChunkSize", - "norm_level", - "normLevel", - ]; - - #[allow(clippy::enum_variant_names)] - enum GeneratedField { - MinChunkSize, - MaxChunkSize, - NormLevel, - } - impl<'de> serde::Deserialize<'de> for GeneratedField { - fn deserialize(deserializer: D) -> std::result::Result - where - D: serde::Deserializer<'de>, - { - struct GeneratedVisitor; - - impl serde::de::Visitor<'_> for GeneratedVisitor { - type Value = GeneratedField; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(formatter, "expected one of: {:?}", &FIELDS) - } - - #[allow(unused_variables)] - fn visit_str(self, value: &str) -> std::result::Result - where - E: serde::de::Error, - { - match value { - "minChunkSize" | "min_chunk_size" => Ok(GeneratedField::MinChunkSize), - "maxChunkSize" | "max_chunk_size" => Ok(GeneratedField::MaxChunkSize), - "normLevel" | "norm_level" => Ok(GeneratedField::NormLevel), - _ => Err(serde::de::Error::unknown_field(value, FIELDS)), - } - } - } - deserializer.deserialize_identifier(GeneratedVisitor) - } - } - struct GeneratedVisitor; - impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = CdcOptions; - - fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion_common.CdcOptions") - } - - fn visit_map(self, mut map_: V) -> std::result::Result - where - V: serde::de::MapAccess<'de>, - { - let mut min_chunk_size__ = None; - let mut max_chunk_size__ = None; - let mut norm_level__ = None; - while let Some(k) = map_.next_key()? { - match k { - GeneratedField::MinChunkSize => { - if min_chunk_size__.is_some() { - return Err(serde::de::Error::duplicate_field("minChunkSize")); - } - min_chunk_size__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; - } - GeneratedField::MaxChunkSize => { - if max_chunk_size__.is_some() { - return Err(serde::de::Error::duplicate_field("maxChunkSize")); - } - max_chunk_size__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; - } - GeneratedField::NormLevel => { - if norm_level__.is_some() { - return Err(serde::de::Error::duplicate_field("normLevel")); - } - norm_level__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; - } - } - } - Ok(CdcOptions { - min_chunk_size: min_chunk_size__.unwrap_or_default(), - max_chunk_size: max_chunk_size__.unwrap_or_default(), - norm_level: norm_level__.unwrap_or_default(), - }) - } - } - deserializer.deserialize_struct("datafusion_common.CdcOptions", FIELDS, GeneratedVisitor) - } -} impl serde::Serialize for Column { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -4116,6 +3978,118 @@ impl<'de> serde::Deserialize<'de> for EmptyMessage { deserializer.deserialize_struct("datafusion_common.EmptyMessage", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for ExplainAnalyzeCategoriesNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.all { + len += 1; + } + if !self.only.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion_common.ExplainAnalyzeCategoriesNode", len)?; + if self.all { + struct_ser.serialize_field("all", &self.all)?; + } + if !self.only.is_empty() { + let v = self.only.iter().cloned().map(|v| { + MetricCategory::try_from(v) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", v))) + }).collect::, _>>()?; + struct_ser.serialize_field("only", &v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ExplainAnalyzeCategoriesNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "all", + "only", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + All, + Only, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "all" => Ok(GeneratedField::All), + "only" => Ok(GeneratedField::Only), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ExplainAnalyzeCategoriesNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion_common.ExplainAnalyzeCategoriesNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut all__ = None; + let mut only__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::All => { + if all__.is_some() { + return Err(serde::de::Error::duplicate_field("all")); + } + all__ = Some(map_.next_value()?); + } + GeneratedField::Only => { + if only__.is_some() { + return Err(serde::de::Error::duplicate_field("only")); + } + only__ = Some(map_.next_value::>()?.into_iter().map(|x| x as i32).collect()); + } + } + } + Ok(ExplainAnalyzeCategoriesNode { + all: all__.unwrap_or_default(), + only: only__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion_common.ExplainAnalyzeCategoriesNode", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for ExplainFormat { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -5474,6 +5448,154 @@ impl<'de> serde::Deserialize<'de> for Map { deserializer.deserialize_struct("datafusion_common.Map", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for MetricCategory { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Rows => "METRIC_CATEGORY_ROWS", + Self::Bytes => "METRIC_CATEGORY_BYTES", + Self::Timing => "METRIC_CATEGORY_TIMING", + Self::Uncategorized => "METRIC_CATEGORY_UNCATEGORIZED", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for MetricCategory { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "METRIC_CATEGORY_ROWS", + "METRIC_CATEGORY_BYTES", + "METRIC_CATEGORY_TIMING", + "METRIC_CATEGORY_UNCATEGORIZED", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = MetricCategory; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "METRIC_CATEGORY_ROWS" => Ok(MetricCategory::Rows), + "METRIC_CATEGORY_BYTES" => Ok(MetricCategory::Bytes), + "METRIC_CATEGORY_TIMING" => Ok(MetricCategory::Timing), + "METRIC_CATEGORY_UNCATEGORIZED" => Ok(MetricCategory::Uncategorized), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} +impl serde::Serialize for MetricType { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Summary => "METRIC_TYPE_SUMMARY", + Self::Dev => "METRIC_TYPE_DEV", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for MetricType { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "METRIC_TYPE_SUMMARY", + "METRIC_TYPE_DEV", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = MetricType; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "METRIC_TYPE_SUMMARY" => Ok(MetricType::Summary), + "METRIC_TYPE_DEV" => Ok(MetricType::Dev), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} impl serde::Serialize for NdJsonFormat { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -5636,6 +5758,161 @@ impl<'de> serde::Deserialize<'de> for NullEquality { deserializer.deserialize_any(GeneratedVisitor) } } +impl serde::Serialize for ParquetCdcOptions { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.enabled { + len += 1; + } + if self.min_chunk_size != 0 { + len += 1; + } + if self.max_chunk_size != 0 { + len += 1; + } + if self.norm_level != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion_common.ParquetCdcOptions", len)?; + if self.enabled { + struct_ser.serialize_field("enabled", &self.enabled)?; + } + if self.min_chunk_size != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("minChunkSize", ToString::to_string(&self.min_chunk_size).as_str())?; + } + if self.max_chunk_size != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("maxChunkSize", ToString::to_string(&self.max_chunk_size).as_str())?; + } + if self.norm_level != 0 { + struct_ser.serialize_field("normLevel", &self.norm_level)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ParquetCdcOptions { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "enabled", + "min_chunk_size", + "minChunkSize", + "max_chunk_size", + "maxChunkSize", + "norm_level", + "normLevel", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Enabled, + MinChunkSize, + MaxChunkSize, + NormLevel, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "enabled" => Ok(GeneratedField::Enabled), + "minChunkSize" | "min_chunk_size" => Ok(GeneratedField::MinChunkSize), + "maxChunkSize" | "max_chunk_size" => Ok(GeneratedField::MaxChunkSize), + "normLevel" | "norm_level" => Ok(GeneratedField::NormLevel), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ParquetCdcOptions; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion_common.ParquetCdcOptions") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut enabled__ = None; + let mut min_chunk_size__ = None; + let mut max_chunk_size__ = None; + let mut norm_level__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Enabled => { + if enabled__.is_some() { + return Err(serde::de::Error::duplicate_field("enabled")); + } + enabled__ = Some(map_.next_value()?); + } + GeneratedField::MinChunkSize => { + if min_chunk_size__.is_some() { + return Err(serde::de::Error::duplicate_field("minChunkSize")); + } + min_chunk_size__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::MaxChunkSize => { + if max_chunk_size__.is_some() { + return Err(serde::de::Error::duplicate_field("maxChunkSize")); + } + max_chunk_size__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::NormLevel => { + if norm_level__.is_some() { + return Err(serde::de::Error::duplicate_field("normLevel")); + } + norm_level__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + } + } + Ok(ParquetCdcOptions { + enabled: enabled__.unwrap_or_default(), + min_chunk_size: min_chunk_size__.unwrap_or_default(), + max_chunk_size: max_chunk_size__.unwrap_or_default(), + norm_level: norm_level__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion_common.ParquetCdcOptions", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for ParquetColumnOptions { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -6132,6 +6409,9 @@ impl serde::Serialize for ParquetOptions { if self.max_row_group_size != 0 { len += 1; } + if self.max_in_list_size != 0 { + len += 1; + } if !self.created_by.is_empty() { len += 1; } @@ -6171,6 +6451,9 @@ impl serde::Serialize for ParquetOptions { if self.max_predicate_cache_size_opt.is_some() { len += 1; } + if self.max_row_group_bytes_opt.is_some() { + len += 1; + } if self.coerce_int96_tz_opt.is_some() { len += 1; } @@ -6249,6 +6532,11 @@ impl serde::Serialize for ParquetOptions { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("maxRowGroupSize", ToString::to_string(&self.max_row_group_size).as_str())?; } + if self.max_in_list_size != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("maxInListSize", ToString::to_string(&self.max_in_list_size).as_str())?; + } if !self.created_by.is_empty() { struct_ser.serialize_field("createdBy", &self.created_by)?; } @@ -6342,6 +6630,15 @@ impl serde::Serialize for ParquetOptions { } } } + if let Some(v) = self.max_row_group_bytes_opt.as_ref() { + match v { + parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(v) => { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("maxRowGroupBytes", ToString::to_string(&v).as_str())?; + } + } + } if let Some(v) = self.coerce_int96_tz_opt.as_ref() { match v { parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(v) => { @@ -6398,6 +6695,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "dataPageRowCountLimit", "max_row_group_size", "maxRowGroupSize", + "max_in_list_size", + "maxInListSize", "created_by", "createdBy", "content_defined_chunking", @@ -6422,6 +6721,8 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "coerceInt96", "max_predicate_cache_size", "maxPredicateCacheSize", + "max_row_group_bytes", + "maxRowGroupBytes", "coerce_int96_tz", "coerceInt96Tz", ]; @@ -6448,6 +6749,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { DictionaryPageSizeLimit, DataPageRowCountLimit, MaxRowGroupSize, + MaxInListSize, CreatedBy, ContentDefinedChunking, MetadataSizeHint, @@ -6461,6 +6763,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { BloomFilterNdv, CoerceInt96, MaxPredicateCacheSize, + MaxRowGroupBytes, CoerceInt96Tz, } impl<'de> serde::Deserialize<'de> for GeneratedField { @@ -6503,6 +6806,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "dictionaryPageSizeLimit" | "dictionary_page_size_limit" => Ok(GeneratedField::DictionaryPageSizeLimit), "dataPageRowCountLimit" | "data_page_row_count_limit" => Ok(GeneratedField::DataPageRowCountLimit), "maxRowGroupSize" | "max_row_group_size" => Ok(GeneratedField::MaxRowGroupSize), + "maxInListSize" | "max_in_list_size" => Ok(GeneratedField::MaxInListSize), "createdBy" | "created_by" => Ok(GeneratedField::CreatedBy), "contentDefinedChunking" | "content_defined_chunking" => Ok(GeneratedField::ContentDefinedChunking), "metadataSizeHint" | "metadata_size_hint" => Ok(GeneratedField::MetadataSizeHint), @@ -6516,6 +6820,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { "bloomFilterNdv" | "bloom_filter_ndv" => Ok(GeneratedField::BloomFilterNdv), "coerceInt96" | "coerce_int96" => Ok(GeneratedField::CoerceInt96), "maxPredicateCacheSize" | "max_predicate_cache_size" => Ok(GeneratedField::MaxPredicateCacheSize), + "maxRowGroupBytes" | "max_row_group_bytes" => Ok(GeneratedField::MaxRowGroupBytes), "coerceInt96Tz" | "coerce_int96_tz" => Ok(GeneratedField::CoerceInt96Tz), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } @@ -6556,6 +6861,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut dictionary_page_size_limit__ = None; let mut data_page_row_count_limit__ = None; let mut max_row_group_size__ = None; + let mut max_in_list_size__ = None; let mut created_by__ = None; let mut content_defined_chunking__ = None; let mut metadata_size_hint_opt__ = None; @@ -6569,6 +6875,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { let mut bloom_filter_ndv_opt__ = None; let mut coerce_int96_opt__ = None; let mut max_predicate_cache_size_opt__ = None; + let mut max_row_group_bytes_opt__ = None; let mut coerce_int96_tz_opt__ = None; while let Some(k) = map_.next_key()? { match k { @@ -6706,6 +7013,14 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::MaxInListSize => { + if max_in_list_size__.is_some() { + return Err(serde::de::Error::duplicate_field("maxInListSize")); + } + max_in_list_size__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } GeneratedField::CreatedBy => { if created_by__.is_some() { return Err(serde::de::Error::duplicate_field("createdBy")); @@ -6784,6 +7099,12 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { } max_predicate_cache_size_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(x.0)); } + GeneratedField::MaxRowGroupBytes => { + if max_row_group_bytes_opt__.is_some() { + return Err(serde::de::Error::duplicate_field("maxRowGroupBytes")); + } + max_row_group_bytes_opt__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(x.0)); + } GeneratedField::CoerceInt96Tz => { if coerce_int96_tz_opt__.is_some() { return Err(serde::de::Error::duplicate_field("coerceInt96Tz")); @@ -6813,6 +7134,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { dictionary_page_size_limit: dictionary_page_size_limit__.unwrap_or_default(), data_page_row_count_limit: data_page_row_count_limit__.unwrap_or_default(), max_row_group_size: max_row_group_size__.unwrap_or_default(), + max_in_list_size: max_in_list_size__.unwrap_or_default(), created_by: created_by__.unwrap_or_default(), content_defined_chunking: content_defined_chunking__, metadata_size_hint_opt: metadata_size_hint_opt__, @@ -6826,6 +7148,7 @@ impl<'de> serde::Deserialize<'de> for ParquetOptions { bloom_filter_ndv_opt: bloom_filter_ndv_opt__, coerce_int96_opt: coerce_int96_opt__, max_predicate_cache_size_opt: max_predicate_cache_size_opt__, + max_row_group_bytes_opt: max_row_group_bytes_opt__, coerce_int96_tz_opt: coerce_int96_tz_opt__, }) } diff --git a/datafusion/proto-common/src/generated/prost.rs b/datafusion/proto-common/src/generated/prost.rs index 632b16929faa6..bdbe38538e1d7 100644 --- a/datafusion/proto-common/src/generated/prost.rs +++ b/datafusion/proto-common/src/generated/prost.rs @@ -862,10 +862,12 @@ pub struct ParquetOptions { pub data_page_row_count_limit: u64, #[prost(uint64, tag = "15")] pub max_row_group_size: u64, + #[prost(uint64, tag = "38")] + pub max_in_list_size: u64, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] - pub content_defined_chunking: ::core::option::Option, + pub content_defined_chunking: ::core::option::Option, #[prost(oneof = "parquet_options::MetadataSizeHintOpt", tags = "4")] pub metadata_size_hint_opt: ::core::option::Option< parquet_options::MetadataSizeHintOpt, @@ -900,6 +902,10 @@ pub struct ParquetOptions { pub max_predicate_cache_size_opt: ::core::option::Option< parquet_options::MaxPredicateCacheSizeOpt, >, + #[prost(oneof = "parquet_options::MaxRowGroupBytesOpt", tags = "37")] + pub max_row_group_bytes_opt: ::core::option::Option< + parquet_options::MaxRowGroupBytesOpt, + >, /// Optional timezone applied to INT96-coerced timestamps when `coerce_int96` /// is set. When `Some`, INT96 columns coerce to /// `Timestamp(, Some())` instead of the default @@ -964,6 +970,11 @@ pub mod parquet_options { #[prost(uint64, tag = "33")] MaxPredicateCacheSize(u64), } + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum MaxRowGroupBytesOpt { + #[prost(uint64, tag = "37")] + MaxRowGroupBytes(u64), + } /// Optional timezone applied to INT96-coerced timestamps when `coerce_int96` /// is set. When `Some`, INT96 columns coerce to /// `Timestamp(, Some())` instead of the default @@ -974,13 +985,16 @@ pub mod parquet_options { CoerceInt96Tz(::prost::alloc::string::String), } } +/// Content-defined chunking (CDC) options for writing parquet files. #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] -pub struct CdcOptions { - #[prost(uint64, tag = "1")] - pub min_chunk_size: u64, +pub struct ParquetCdcOptions { + #[prost(bool, tag = "1")] + pub enabled: bool, #[prost(uint64, tag = "2")] + pub min_chunk_size: u64, + #[prost(uint64, tag = "3")] pub max_chunk_size: u64, - #[prost(int32, tag = "3")] + #[prost(int32, tag = "4")] pub norm_level: i32, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1014,6 +1028,18 @@ pub struct ColumnStats { #[prost(message, optional, tag = "6")] pub byte_size: ::core::option::Option, } +/// Wire encoding for `datafusion_common::format::ExplainAnalyzeCategories`. +/// +/// If `all` is true, every category is shown (the `only` list is ignored). +/// If `all` is false, only the categories listed in `only` are shown — an +/// empty `only` means "plan only", i.e. suppress all metrics. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ExplainAnalyzeCategoriesNode { + #[prost(bool, tag = "1")] + pub all: bool, + #[prost(enumeration = "MetricCategory", repeated, tag = "2")] + pub only: ::prost::alloc::vec::Vec, +} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum JoinType { @@ -1360,3 +1386,65 @@ impl ExplainFormat { } } } +/// Verbosity level for `EXPLAIN ANALYZE`. Mirrors +/// `datafusion_common::format::MetricType`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum MetricType { + Summary = 0, + Dev = 1, +} +impl MetricType { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Summary => "METRIC_TYPE_SUMMARY", + Self::Dev => "METRIC_TYPE_DEV", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "METRIC_TYPE_SUMMARY" => Some(Self::Summary), + "METRIC_TYPE_DEV" => Some(Self::Dev), + _ => None, + } + } +} +/// Category of an `EXPLAIN ANALYZE` metric. Mirrors +/// `datafusion_common::format::MetricCategory`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum MetricCategory { + Rows = 0, + Bytes = 1, + Timing = 2, + Uncategorized = 3, +} +impl MetricCategory { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Rows => "METRIC_CATEGORY_ROWS", + Self::Bytes => "METRIC_CATEGORY_BYTES", + Self::Timing => "METRIC_CATEGORY_TIMING", + Self::Uncategorized => "METRIC_CATEGORY_UNCATEGORIZED", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "METRIC_CATEGORY_ROWS" => Some(Self::Rows), + "METRIC_CATEGORY_BYTES" => Some(Self::Bytes), + "METRIC_CATEGORY_TIMING" => Some(Self::Timing), + "METRIC_CATEGORY_UNCATEGORIZED" => Some(Self::Uncategorized), + _ => None, + } + } +} diff --git a/datafusion/proto-common/src/to_proto/mod.rs b/datafusion/proto-common/src/to_proto/mod.rs index 940679b836ff1..360981746585b 100644 --- a/datafusion/proto-common/src/to_proto/mod.rs +++ b/datafusion/proto-common/src/to_proto/mod.rs @@ -29,14 +29,14 @@ use arrow::datatypes::{ SchemaRef, TimeUnit, UnionMode, }; use arrow::ipc::writer::{ - CompressionContext, DictionaryTracker, IpcDataGenerator, IpcWriteOptions, + DictionaryTracker, IpcDataGenerator, IpcWriteContext, IpcWriteOptions, }; use datafusion_common::parsers::CsvQuoteStyle; use datafusion_common::{ Column, ColumnStatistics, Constraint, Constraints, DFSchema, DFSchemaRef, DataFusionError, JoinSide, ScalarValue, Statistics, config::{ - CsvOptions, JsonOptions, ParquetColumnOptions, ParquetOptions, + CsvOptions, JsonOptions, ParquetCdcOptions, ParquetColumnOptions, ParquetOptions, TableParquetOptions, }, file_options::{csv_writer::CsvWriterOptions, json_writer::JsonWriterOptions}, @@ -115,7 +115,7 @@ impl TryFrom<&DataType> for protobuf::ArrowType { } } -impl TryFrom<&DataType> for protobuf::arrow_type::ArrowTypeEnum { +impl TryFrom<&DataType> for ArrowTypeEnum { type Error = Error; fn try_from(val: &DataType) -> Result { @@ -439,9 +439,7 @@ impl TryFrom<&ScalarValue> for protobuf::ScalarValue { }) } None => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::NullValue( - (&data_type).try_into()?, - )), + value: Some(Value::NullValue((&data_type).try_into()?)), }), }, ScalarValue::Decimal64(val, p, s) => match *val { @@ -457,9 +455,7 @@ impl TryFrom<&ScalarValue> for protobuf::ScalarValue { }) } None => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::NullValue( - (&data_type).try_into()?, - )), + value: Some(Value::NullValue((&data_type).try_into()?)), }), }, ScalarValue::Decimal128(val, p, s) => match *val { @@ -475,9 +471,7 @@ impl TryFrom<&ScalarValue> for protobuf::ScalarValue { }) } None => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::NullValue( - (&data_type).try_into()?, - )), + value: Some(Value::NullValue((&data_type).try_into()?)), }), }, ScalarValue::Decimal256(val, p, s) => match *val { @@ -493,9 +487,7 @@ impl TryFrom<&ScalarValue> for protobuf::ScalarValue { }) } None => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::NullValue( - (&data_type).try_into()?, - )), + value: Some(Value::NullValue((&data_type).try_into()?)), }), }, ScalarValue::Date64(val) => { @@ -788,8 +780,8 @@ impl From<&Precision> for protobuf::Precision { } } -impl From<&Precision> for protobuf::Precision { - fn from(s: &Precision) -> protobuf::Precision { +impl From<&Precision> for protobuf::Precision { + fn from(s: &Precision) -> protobuf::Precision { match s { Precision::Exact(val) => protobuf::Precision { precision_info: protobuf::PrecisionInfo::Exact.into(), @@ -920,6 +912,7 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions { dictionary_page_size_limit: value.dictionary_page_size_limit as u64, statistics_enabled_opt: value.statistics_enabled.clone().map(protobuf::parquet_options::StatisticsEnabledOpt::StatisticsEnabled), max_row_group_size: value.max_row_group_size as u64, + max_in_list_size: value.max_in_list_size as u64, created_by: value.created_by.clone(), column_index_truncate_length_opt: value.column_index_truncate_length.map(|v| protobuf::parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(v as u64)), statistics_truncate_length_opt: value.statistics_truncate_length.map(|v| protobuf::parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(v as u64)), @@ -938,17 +931,23 @@ impl TryFrom<&ParquetOptions> for protobuf::ParquetOptions { coerce_int96_opt: value.coerce_int96.clone().map(protobuf::parquet_options::CoerceInt96Opt::CoerceInt96), coerce_int96_tz_opt: value.coerce_int96_tz.clone().map(protobuf::parquet_options::CoerceInt96TzOpt::CoerceInt96Tz), max_predicate_cache_size_opt: value.max_predicate_cache_size.map(|v| protobuf::parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(v as u64)), - content_defined_chunking: value.use_content_defined_chunking.as_ref().map(|cdc| - protobuf::CdcOptions { - min_chunk_size: cdc.min_chunk_size as u64, - max_chunk_size: cdc.max_chunk_size as u64, - norm_level: cdc.norm_level, - } - ), + max_row_group_bytes_opt: value.max_row_group_bytes.map(|v| protobuf::parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(v.get() as u64)), + content_defined_chunking: Some((&value.content_defined_chunking).into()), }) } } +impl From<&ParquetCdcOptions> for protobuf::ParquetCdcOptions { + fn from(value: &ParquetCdcOptions) -> Self { + protobuf::ParquetCdcOptions { + enabled: value.enabled, + min_chunk_size: value.min_chunk_size as u64, + max_chunk_size: value.max_chunk_size as u64, + norm_level: value.norm_level, + } + } +} + impl TryFrom<&ParquetColumnOptions> for protobuf::ParquetColumnOptions { type Error = DataFusionError; @@ -1070,16 +1069,14 @@ impl TryFrom<&JsonOptions> for protobuf::JsonOptions { /// Creates a scalar protobuf value from an optional value (T), and /// encoding None as the appropriate datatype -fn create_proto_scalar protobuf::scalar_value::Value>( +fn create_proto_scalar Value>( v: Option<&I>, null_arrow_type: &DataType, constructor: T, ) -> Result { let value = v .map(constructor) - .unwrap_or(protobuf::scalar_value::Value::NullValue( - null_arrow_type.try_into()?, - )); + .unwrap_or(Value::NullValue(null_arrow_type.try_into()?)); Ok(protobuf::ScalarValue { value: Some(value) }) } @@ -1106,7 +1103,7 @@ fn encode_scalar_nested_value( &mut dict_tracker, &write_options, ); - let mut compression_context = CompressionContext::default(); + let mut compression_context = IpcWriteContext::default(); let (encoded_dictionaries, encoded_message) = ipc_gen .encode( &batch, @@ -1135,35 +1132,25 @@ fn encode_scalar_nested_value( match val { ScalarValue::List(_) => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::ListValue(scalar_list_value)), + value: Some(Value::ListValue(scalar_list_value)), }), ScalarValue::LargeList(_) => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::LargeListValue( - scalar_list_value, - )), + value: Some(Value::LargeListValue(scalar_list_value)), }), ScalarValue::FixedSizeList(_) => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::FixedSizeListValue( - scalar_list_value, - )), + value: Some(Value::FixedSizeListValue(scalar_list_value)), }), ScalarValue::ListView(_) => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::ListViewValue( - scalar_list_value, - )), + value: Some(Value::ListViewValue(scalar_list_value)), }), ScalarValue::LargeListView(_) => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::LargeListViewValue( - scalar_list_value, - )), + value: Some(Value::LargeListViewValue(scalar_list_value)), }), ScalarValue::Struct(_) => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::StructValue( - scalar_list_value, - )), + value: Some(Value::StructValue(scalar_list_value)), }), ScalarValue::Map(_) => Ok(protobuf::ScalarValue { - value: Some(protobuf::scalar_value::Value::MapValue(scalar_list_value)), + value: Some(Value::MapValue(scalar_list_value)), }), _ => unreachable!(), } diff --git a/datafusion/proto-models/.gitignore b/datafusion/proto-models/.gitignore new file mode 100644 index 0000000000000..662b95f238c24 --- /dev/null +++ b/datafusion/proto-models/.gitignore @@ -0,0 +1,5 @@ +# Files generated by regen.sh +proto/proto_descriptor.bin +src/datafusion.rs +src/datafusion.serde.rs +src/datafusion_common.rs diff --git a/datafusion/proto-models/Cargo.toml b/datafusion/proto-models/Cargo.toml new file mode 100644 index 0000000000000..83b1a202c24ed --- /dev/null +++ b/datafusion/proto-models/Cargo.toml @@ -0,0 +1,52 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +[package] +name = "datafusion-proto-models" +description = "Protobuf-generated model types for DataFusion logical and physical plans" +keywords = ["arrow", "query", "sql"] +readme = "README.md" +version = { workspace = true } +edition = { workspace = true } +homepage = { workspace = true } +repository = { workspace = true } +license = { workspace = true } +authors = { workspace = true } +rust-version = { workspace = true } + +[package.metadata.docs.rs] +all-features = true + +# Note: add additional linter rules in lib.rs. +# Rust does not support workspace + new linter rules in subcrates yet +# https://github.com/rust-lang/cargo/issues/13157 +[lints] +workspace = true + +[lib] +name = "datafusion_proto_models" + +[features] +default = [] +json = ["serde", "pbjson", "datafusion-proto-common/json"] + +[dependencies] +datafusion-common = { workspace = true } +datafusion-proto-common = { workspace = true } +pbjson = { workspace = true, optional = true } +prost = { workspace = true } +serde = { version = "1.0", optional = true } diff --git a/datafusion/proto-models/LICENSE.txt b/datafusion/proto-models/LICENSE.txt new file mode 100644 index 0000000000000..d74c6b599d2ae --- /dev/null +++ b/datafusion/proto-models/LICENSE.txt @@ -0,0 +1,212 @@ + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. + + +This project includes code from Apache Aurora. + +* dev/release/{release,changelog,release-candidate} are based on the scripts from + Apache Aurora + +Copyright: 2016 The Apache Software Foundation. +Home page: https://aurora.apache.org/ +License: http://www.apache.org/licenses/LICENSE-2.0 diff --git a/datafusion/proto-models/NOTICE.txt b/datafusion/proto-models/NOTICE.txt new file mode 100644 index 0000000000000..0bd2d52368fea --- /dev/null +++ b/datafusion/proto-models/NOTICE.txt @@ -0,0 +1,5 @@ +Apache DataFusion +Copyright 2019-2026 The Apache Software Foundation + +This product includes software developed at +The Apache Software Foundation (http://www.apache.org/). diff --git a/datafusion/proto-models/README.md b/datafusion/proto-models/README.md new file mode 100644 index 0000000000000..34adae9c1aaef --- /dev/null +++ b/datafusion/proto-models/README.md @@ -0,0 +1,43 @@ + + +# Apache DataFusion Protobuf Models + +[Apache DataFusion] is an extensible query execution framework, written in Rust, that uses [Apache Arrow] as its in-memory format. + +This crate contains the [prost]-generated Rust types for DataFusion's logical +and physical plan protobuf schemas. It is intentionally kept narrow: it has no +DataFusion dependencies beyond [`datafusion-proto-common`] and exposes only the +generated structs (and optional [pbjson]/[serde] support). + +This crate is consumed by [`datafusion-proto`] and may also be depended on +directly by other DataFusion crates that need to refer to the proto schema +types without pulling in the full [`datafusion-proto`] surface. + +Most projects should use the [`datafusion-proto`] crate directly, which +re-exports this module. If you are already using the [`datafusion-proto`] +crate, there is no reason to use this crate directly in your project as well. + +[apache arrow]: https://arrow.apache.org/ +[apache datafusion]: https://datafusion.apache.org/ +[prost]: https://docs.rs/prost/latest/prost/ +[pbjson]: https://docs.rs/pbjson/latest/pbjson/ +[serde]: https://serde.rs/ +[`datafusion-proto`]: https://crates.io/crates/datafusion-proto +[`datafusion-proto-common`]: https://crates.io/crates/datafusion-proto-common diff --git a/datafusion/proto/gen/Cargo.toml b/datafusion/proto-models/gen/Cargo.toml similarity index 98% rename from datafusion/proto/gen/Cargo.toml rename to datafusion/proto-models/gen/Cargo.toml index 8b48dfe70e6c7..9724b63cccf3c 100644 --- a/datafusion/proto/gen/Cargo.toml +++ b/datafusion/proto-models/gen/Cargo.toml @@ -38,4 +38,4 @@ workspace = true [dependencies] # Pin these dependencies so that the generated output is deterministic pbjson-build = "=0.9.0" -prost-build = "=0.14.3" +prost-build = "=0.14.4" diff --git a/datafusion/proto/gen/src/main.rs b/datafusion/proto-models/gen/src/main.rs similarity index 85% rename from datafusion/proto/gen/src/main.rs rename to datafusion/proto-models/gen/src/main.rs index 7f163162035c8..b9cbf81bb11c8 100644 --- a/datafusion/proto/gen/src/main.rs +++ b/datafusion/proto-models/gen/src/main.rs @@ -18,9 +18,9 @@ use std::path::Path; fn main() -> Result<(), String> { - let proto_dir = Path::new("datafusion/proto"); - let proto_path = Path::new("datafusion/proto/proto/datafusion.proto"); - let out_dir = Path::new("datafusion/proto/src"); + let proto_dir = Path::new("datafusion/proto-models"); + let proto_path = Path::new("datafusion/proto-models/proto/datafusion.proto"); + let out_dir = Path::new("datafusion/proto-models/src"); // proto definitions has to be there let descriptor_path = proto_dir.join("proto/proto_descriptor.bin"); @@ -35,14 +35,12 @@ fn main() -> Result<(), String> { .map_err(|e| format!("protobuf compilation failed: {e}"))?; let descriptor_set = std::fs::read(&descriptor_path) - .unwrap_or_else(|e| panic!("Cannot read {:?}: {}", &descriptor_path, e)); + .unwrap_or_else(|e| panic!("Cannot read {descriptor_path:?}: {e}")); pbjson_build::Builder::new() .out_dir(out_dir) .register_descriptors(&descriptor_set) - .unwrap_or_else(|e| { - panic!("Cannot register descriptors {:?}: {}", &descriptor_set, e) - }) + .unwrap_or_else(|e| panic!("Cannot register descriptors {descriptor_set:?}: {e}")) .build(&[".datafusion"]) .map_err(|e| format!("pbjson compilation failed: {e}"))?; diff --git a/datafusion/proto/proto/datafusion.proto b/datafusion/proto-models/proto/datafusion.proto similarity index 84% rename from datafusion/proto/proto/datafusion.proto rename to datafusion/proto-models/proto/datafusion.proto index d34acf36c54f2..43a90264c2b1f 100644 --- a/datafusion/proto/proto/datafusion.proto +++ b/datafusion/proto-models/proto/datafusion.proto @@ -92,8 +92,8 @@ message ListingTableScanNode { datafusion_common.Schema schema = 5; repeated LogicalExprNode filters = 6; repeated PartitionColumn table_partition_cols = 7; - bool collect_stat = 8; - uint32 target_partitions = 9; + reserved 8; // was bool collect_stat + reserved 9; // was uint32 target_partitions oneof FileFormatType { datafusion_common.CsvFormat csv = 10; datafusion_common.ParquetFormat parquet = 11; @@ -148,9 +148,19 @@ message RepartitionNode { oneof partition_method { uint64 round_robin = 2; HashRepartition hash = 3; + RangeRepartition range = 4; } } +message RangeSplitPoint { + repeated datafusion_common.ScalarValue value = 1; +} + +message RangeRepartition { + repeated SortExprNode sort_expr = 1; + repeated RangeSplitPoint split_point = 2; +} + message HashRepartition { repeated LogicalExprNode hash_expr = 1; uint64 partition_count = 2; @@ -163,7 +173,8 @@ message EmptyRelationNode { message CreateExternalTableNode { reserved 1; // was string name TableReference name = 9; - string location = 2; + string location = 2; // deprecated; use repeated locations + repeated string locations = 16; string file_type = 3; datafusion_common.DfSchema schema = 4; repeated string table_partition_cols = 5; @@ -224,12 +235,22 @@ message ValuesNode { message AnalyzeNode { LogicalPlanNode input = 1; bool verbose = 2; + // Statement-level override for `datafusion.explain.analyze_level`. + // Absent means "fall back to session config". + optional datafusion_common.MetricType analyze_level = 3; + // Statement-level override for `datafusion.explain.analyze_categories`. + // Absent means "fall back to session config". + optional datafusion_common.ExplainAnalyzeCategoriesNode analyze_categories = 4; + datafusion_common.ExplainFormat format = 5; } message ExplainNode { LogicalPlanNode input = 1; bool verbose = 2; datafusion_common.ExplainFormat format = 3; + // Statement-level override for `datafusion.explain.show_statistics`. + // Absent means "fall back to session config". + optional bool show_statistics = 4; } message AggregateNode { @@ -252,6 +273,7 @@ message JoinNode { repeated LogicalExprNode right_join_key = 6; datafusion_common.NullEquality null_equality = 7; LogicalExprNode filter = 8; + bool null_aware = 9; } message DistinctNode { @@ -292,7 +314,7 @@ message FileFormatProto { } message DmlNode{ - enum Type { + enum Type { UPDATE = 0; DELETE = 1; CTAS = 2; @@ -300,13 +322,63 @@ message DmlNode{ INSERT_OVERWRITE = 4; INSERT_REPLACE = 5; TRUNCATE = 6; + MERGE_INTO = 7; } Type dml_type = 1; LogicalPlanNode input = 2; TableReference table_name = 3; LogicalPlanNode target = 5; + // Populated only when dml_type == MERGE_INTO. + MergeIntoOpNode merge_into = 6; +} + +// Carries the ON condition and WHEN clauses of a MERGE INTO operation. +message MergeIntoOpNode { + LogicalExprNode on = 1; + repeated MergeIntoClauseNode clauses = 2; +} + +// A single WHEN clause within a MERGE INTO statement. +message MergeIntoClauseNode { + enum Kind { + MATCHED = 0; + NOT_MATCHED = 1; + NOT_MATCHED_BY_TARGET = 2; + NOT_MATCHED_BY_SOURCE = 3; + } + Kind kind = 1; + // Optional `AND ` predicate. Absent when the clause has no predicate. + LogicalExprNode predicate = 2; + MergeIntoActionNode action = 3; +} + +// The action for a single WHEN clause. +message MergeIntoActionNode { + oneof action { + MergeUpdateAction update = 1; + MergeInsertAction insert = 2; + MergeDeleteAction delete = 3; + } +} + +message MergeUpdateAction { + repeated MergeAssignment assignments = 1; +} + +message MergeAssignment { + string column = 1; + LogicalExprNode value = 2; +} + +message MergeInsertAction { + // May be empty (meaning all columns). + repeated string columns = 1; + // One expression per inserted column. + repeated LogicalExprNode values = 2; } +message MergeDeleteAction {} + message UnnestNode { LogicalPlanNode input = 1; repeated datafusion_common.Column exec_columns = 2; @@ -331,7 +403,22 @@ message ColumnUnnestListRecursion { } message UnnestOptions { - bool preserve_nulls = 1; + // Reserved for the historical `bool preserve_nulls = 1;` field. + // Use `null_handling` instead. + reserved 1; + reserved "preserve_nulls"; + + enum NullHandling { + // Preserve nulls; empty lists produce no rows. The historical default. + PRESERVE = 0; + // Drop both null and empty lists from the output. + DROP = 1; + // Preserve nulls, and additionally expand empty lists into a single + // NULL output row (outer-unnest semantics). + PRESERVE_AND_EXPAND_EMPTY = 2; + } + + NullHandling null_handling = 3; repeated RecursionUnnestOption recursions = 2; } @@ -430,6 +517,10 @@ message LogicalExprNode { // Subquery expressions ScalarSubqueryExprNode scalar_subquery_expr = 36; + + HigherOrderUDFExprNode higher_order_udf_expr = 37; + Lambda lambda = 38; + LambdaVariable lambda_variable = 39; } } @@ -542,6 +633,9 @@ message NegativeNode { message Unnest { repeated LogicalExprNode exprs = 1; + // When true, this Unnest expression has outer-unnest semantics: NULL and + // empty input lists both produce a single NULL output row. + bool outer = 2; } message InListNode { @@ -567,6 +661,22 @@ message ScalarUDFExprNode { optional bytes fun_definition = 3; } +message HigherOrderUDFExprNode { + string fun_name = 1; + repeated LogicalExprNode args = 2; + optional bytes fun_definition = 3; +} + +message Lambda { + repeated string params = 1; + LogicalExprNode body = 2; +} + +message LambdaVariable { + string name = 1; + datafusion_common.Field field = 2; +} + message WindowExprNode { oneof window_function { // BuiltInWindowFunction built_in_function = 2; @@ -936,6 +1046,11 @@ message PhysicalExprNode { PhysicalScalarSubqueryExprNode scalar_subquery = 22; PhysicalDynamicFilterNode dynamic_filter = 23; + + PhysicalHigherOrderUdfNode higher_order_udf = 24; + PhysicalLambdaExprNode lambda = 25; + PhysicalLambdaVariableExprNode lambda_variable = 26; + PhysicalRangeExprNode range_expr = 27; } } @@ -956,6 +1071,22 @@ message PhysicalScalarUdfNode { string return_field_name = 6; } +message PhysicalHigherOrderUdfNode { + string name = 1; + repeated PhysicalExprNode args = 2; + optional bytes fun_definition = 3; +} + +message PhysicalLambdaExprNode { + repeated string params = 1; + PhysicalExprNode body = 2; +} + +message PhysicalLambdaVariableExprNode { + uint32 index = 1; + datafusion_common.Field field = 2; +} + message PhysicalAggregateExprNode { oneof AggregateFunction { string user_defined_aggr_function = 4; @@ -966,6 +1097,7 @@ message PhysicalAggregateExprNode { bool ignore_nulls = 6; optional bytes fun_definition = 7; string human_display = 8; + bool is_reversed = 9; } message PhysicalWindowExprNode { @@ -1071,6 +1203,11 @@ message PhysicalHashExprNode { string description = 6; } +message PhysicalRangeExprNode { + repeated PhysicalSortExprNode sort_expr = 1; + repeated PhysicalRangeSplitPoint split_point = 2; +} + message FilterExecNode { PhysicalPlanNode input = 1; PhysicalExprNode expr = 2; @@ -1119,6 +1256,10 @@ message FileScanExecConf { optional uint64 batch_size = 12; optional ProjectionExprs projection_exprs = 13; + // Was optional bool partitioned_by_file_group = 14. + reserved 14; + reserved "partitioned_by_file_group"; + optional Partitioning output_partitioning = 15; } message ParquetScanExecNode { @@ -1190,6 +1331,14 @@ message HashJoinExecNode { bool null_aware = 10; // Optional dynamic filter expression for pushing down to the probe side. PhysicalExprNode dynamic_filter = 11; + // Optional row limit pushed into the join by the `limit_pushdown` rule. + // + // This is presence-tracked (`optional`) on purpose: messages produced by + // versions predating this field carry no `fetch` at all, and a plain proto3 + // scalar would decode that absence as `0`, i.e. "fetch 0 rows", silently + // turning old plans into empty results. With `optional`, absent decodes to + // `None`, which is the correct reading of an older message. + optional uint64 fetch = 12; } enum StreamPartitionMode { @@ -1232,6 +1381,7 @@ message AnalyzeExecNode { // Empty means "plan only". Absent (has_metric_categories=false) means "all". bool has_metric_categories = 5; repeated string metric_categories = 6; + datafusion_common.ExplainFormat format = 7; } message CrossJoinExecNode { @@ -1255,10 +1405,16 @@ message JoinOn { message EmptyExecNode { datafusion_common.Schema schema = 1; + // Number of output partitions. Absent (0) means a single partition, so that + // plans encoded before this field existed decode to the previous default. + uint32 partitions = 2; } message PlaceholderRowExecNode { datafusion_common.Schema schema = 1; + // Number of output partitions. Absent (0) means a single partition, so that + // plans encoded before this field existed decode to the previous default. + uint32 partitions = 2; } message ProjectionExecNode { @@ -1323,6 +1479,8 @@ message AggregateExecNode { bool has_grouping_set = 12; // Optional dynamic filter expression for pushing down to the child. PhysicalExprNode dynamic_filter = 13; + // Output schema preserved by physical optimizer rewrites. + datafusion_common.Schema schema = 14; } message GlobalLimitExecNode { @@ -1331,11 +1489,15 @@ message GlobalLimitExecNode { uint32 skip = 2; // Maximum number of rows to fetch; negative means no limit int64 fetch = 3; + // Ordering the limit must preserve; empty means none + repeated PhysicalSortExprNode required_ordering = 4; } message LocalLimitExecNode { PhysicalPlanNode input = 1; uint32 fetch = 2; + // Ordering the limit must preserve; empty means none + repeated PhysicalSortExprNode required_ordering = 3; } message SortExecNode { @@ -1379,13 +1541,22 @@ message PhysicalHashRepartition { uint64 partition_count = 2; } +message PhysicalRangePartitioning { + repeated PhysicalSortExprNode sort_expr = 1; + repeated PhysicalRangeSplitPoint split_point = 2; +} + +message PhysicalRangeSplitPoint { + repeated datafusion_common.ScalarValue value = 1; +} + message RepartitionExecNode{ PhysicalPlanNode input = 1; - // oneof partition_method { + // Legacy direct partitioning fields: // uint64 round_robin = 2; // PhysicalHashRepartition hash = 3; // uint64 unknown = 4; - // } + // New partitioning variants are stored in `partitioning`. Partitioning partitioning = 5; bool preserve_order = 6; } @@ -1395,6 +1566,7 @@ message Partitioning { uint64 round_robin = 1; PhysicalHashRepartition hash = 2; uint64 unknown = 3; + PhysicalRangePartitioning range = 4; } } @@ -1416,6 +1588,7 @@ message PartitionedFile { repeated datafusion_common.ScalarValue partition_values = 4; FileRange range = 5; datafusion_common.Statistics statistics = 6; + datafusion_common.Schema arrow_schema = 7; } message FileRange { @@ -1438,15 +1611,15 @@ message RecursiveQueryNode { } message CteWorkTableScanNode { - string name = 1; - datafusion_common.Schema schema = 2; + string name = 1; + datafusion_common.Schema schema = 2; } message EmptyTableScanNode { - TableReference table_name = 1; - datafusion_common.Schema schema = 2; - ProjectionColumns projection = 3; - repeated LogicalExprNode filters = 4; + TableReference table_name = 1; + datafusion_common.Schema schema = 2; + ProjectionColumns projection = 3; + repeated LogicalExprNode filters = 4; } enum GenerateSeriesName { @@ -1455,44 +1628,44 @@ enum GenerateSeriesName { } message GenerateSeriesArgsContainsNull { - GenerateSeriesName name = 1; + GenerateSeriesName name = 1; } message GenerateSeriesArgsInt64 { - int64 start = 1; - int64 end = 2; - int64 step = 3; - bool include_end = 4; - GenerateSeriesName name = 5; + int64 start = 1; + int64 end = 2; + int64 step = 3; + bool include_end = 4; + GenerateSeriesName name = 5; } message GenerateSeriesArgsTimestamp { - int64 start = 1; - int64 end = 2; - datafusion_common.IntervalMonthDayNanoValue step = 3; - optional string tz = 4; - bool include_end = 5; - GenerateSeriesName name = 6; + int64 start = 1; + int64 end = 2; + datafusion_common.IntervalMonthDayNanoValue step = 3; + optional string tz = 4; + bool include_end = 5; + GenerateSeriesName name = 6; } message GenerateSeriesArgsDate { - int64 start = 1; - int64 end = 2; - datafusion_common.IntervalMonthDayNanoValue step = 3; - bool include_end = 4; - GenerateSeriesName name = 5; + int64 start = 1; + int64 end = 2; + datafusion_common.IntervalMonthDayNanoValue step = 3; + bool include_end = 4; + GenerateSeriesName name = 5; } message GenerateSeriesNode { - datafusion_common.Schema schema = 1; - uint32 target_batch_size = 2; - - oneof args { - GenerateSeriesArgsContainsNull contains_null = 3; - GenerateSeriesArgsInt64 int64_args = 4; - GenerateSeriesArgsTimestamp timestamp_args = 5; - GenerateSeriesArgsDate date_args = 6; - } + datafusion_common.Schema schema = 1; + uint32 target_batch_size = 2; + + oneof args { + GenerateSeriesArgsContainsNull contains_null = 3; + GenerateSeriesArgsInt64 int64_args = 4; + GenerateSeriesArgsTimestamp timestamp_args = 5; + GenerateSeriesArgsDate date_args = 6; + } } message SortMergeJoinExecNode { diff --git a/datafusion/proto-models/regen.sh b/datafusion/proto-models/regen.sh new file mode 100755 index 0000000000000..4bb07a1f32228 --- /dev/null +++ b/datafusion/proto-models/regen.sh @@ -0,0 +1,21 @@ +#!/usr/bin/env bash + +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +repo_root=$(git rev-parse --show-toplevel) +cd "$repo_root" && cargo run --manifest-path datafusion/proto-models/gen/Cargo.toml diff --git a/datafusion/proto-models/src/from_proto.rs b/datafusion/proto-models/src/from_proto.rs new file mode 100644 index 0000000000000..74ead8c52049b --- /dev/null +++ b/datafusion/proto-models/src/from_proto.rs @@ -0,0 +1,519 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Conversions from the protobuf messages in this crate to their +//! `datafusion-common` counterparts. +//! +//! The DataFusion side of these conversions lives *below* this crate in the +//! dependency graph, so it cannot host the impls itself. They live here +//! instead, on the local proto type — the same arrangement +//! `datafusion-proto-common` uses for `ScalarValue` and `Statistics`. + +use std::sync::Arc; + +use datafusion_common::config::{ + CsvOptions, JsonOptions, MaxRowGroupBytes, ParquetCdcOptions, ParquetColumnOptions, + ParquetOptions, TableParquetOptions, +}; +use datafusion_common::display::{PlanType, StringifiedPlan}; +use datafusion_common::parsers::{CompressionTypeVariant, CsvQuoteStyle}; +use datafusion_common::{ + JoinConstraint, JoinType, NullEquality, RecursionUnnestOption, TableReference, + UnnestOptions, +}; +use datafusion_proto_common::FromProtoError as Error; + +use crate::protobuf::{ + self, AnalyzedLogicalPlanType, CsvOptions as CsvOptionsProto, + CsvQuoteStyle as CsvQuoteStyleProto, JsonOptions as JsonOptionsProto, + OptimizedLogicalPlanType, OptimizedPhysicalPlanType, + ParquetCdcOptions as ParquetCdcOptionsProto, + ParquetColumnOptions as ParquetColumnOptionsProto, + ParquetOptions as ParquetOptionsProto, + TableParquetOptions as TableParquetOptionsProto, parquet_column_options, + parquet_options, + plan_type::PlanTypeEnum::{ + AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, + FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, + InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, + InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, + PhysicalPlanError, + }, +}; + +impl From<&protobuf::UnnestOptions> for UnnestOptions { + fn from(opts: &protobuf::UnnestOptions) -> Self { + use datafusion_common::NullHandling; + use protobuf::unnest_options::NullHandling as ProtoNullHandling; + let null_handling = match ProtoNullHandling::try_from(opts.null_handling) { + Ok(ProtoNullHandling::Preserve) => NullHandling::Preserve, + Ok(ProtoNullHandling::Drop) => NullHandling::Drop, + Ok(ProtoNullHandling::PreserveAndExpandEmpty) => { + NullHandling::PreserveAndExpandEmpty + } + // Unknown enum values fall back to the default (Preserve), which + // matches DataFusion's historical behavior. + Err(_) => NullHandling::Preserve, + }; + Self { + null_handling, + recursions: opts + .recursions + .iter() + .map(|r| RecursionUnnestOption { + input_column: r.input_column.as_ref().unwrap().into(), + output_column: r.output_column.as_ref().unwrap().into(), + depth: r.depth as usize, + }) + .collect::>(), + } + } +} + +impl TryFrom for TableReference { + type Error = Error; + + fn try_from(value: protobuf::TableReference) -> Result { + use protobuf::table_reference::TableReferenceEnum; + let table_reference_enum = value + .table_reference_enum + .ok_or_else(|| Error::required("table_reference_enum"))?; + + match table_reference_enum { + TableReferenceEnum::Bare(protobuf::BareTableReference { table }) => { + Ok(TableReference::bare(table)) + } + TableReferenceEnum::Partial(protobuf::PartialTableReference { + schema, + table, + }) => Ok(TableReference::partial(schema, table)), + TableReferenceEnum::Full(protobuf::FullTableReference { + catalog, + schema, + table, + }) => Ok(TableReference::full(catalog, schema, table)), + } + } +} + +impl From<&protobuf::StringifiedPlan> for StringifiedPlan { + fn from(stringified_plan: &protobuf::StringifiedPlan) -> Self { + Self { + plan_type: match stringified_plan + .plan_type + .as_ref() + .and_then(|pt| pt.plan_type_enum.as_ref()) + .unwrap_or_else(|| { + panic!( + "Cannot create protobuf::StringifiedPlan from {stringified_plan:?}" + ) + }) { + InitialLogicalPlan(_) => PlanType::InitialLogicalPlan, + AnalyzedLogicalPlan(AnalyzedLogicalPlanType { analyzer_name }) => { + PlanType::AnalyzedLogicalPlan { + analyzer_name:analyzer_name.clone() + } + } + FinalAnalyzedLogicalPlan(_) => PlanType::FinalAnalyzedLogicalPlan, + OptimizedLogicalPlan(OptimizedLogicalPlanType { optimizer_name }) => { + PlanType::OptimizedLogicalPlan { + optimizer_name: optimizer_name.clone(), + } + } + FinalLogicalPlan(_) => PlanType::FinalLogicalPlan, + InitialPhysicalPlan(_) => PlanType::InitialPhysicalPlan, + InitialPhysicalPlanWithStats(_) => PlanType::InitialPhysicalPlanWithStats, + InitialPhysicalPlanWithSchema(_) => PlanType::InitialPhysicalPlanWithSchema, + OptimizedPhysicalPlan(OptimizedPhysicalPlanType { optimizer_name }) => { + PlanType::OptimizedPhysicalPlan { + optimizer_name: optimizer_name.clone(), + } + } + FinalPhysicalPlan(_) => PlanType::FinalPhysicalPlan, + FinalPhysicalPlanWithStats(_) => PlanType::FinalPhysicalPlanWithStats, + FinalPhysicalPlanWithSchema(_) => PlanType::FinalPhysicalPlanWithSchema, + PhysicalPlanError(_) => PlanType::PhysicalPlanError, + }, + plan: Arc::new(stringified_plan.plan.clone()), + } + } +} + +impl From for JoinType { + fn from(t: protobuf::JoinType) -> Self { + match t { + protobuf::JoinType::Inner => JoinType::Inner, + protobuf::JoinType::Left => JoinType::Left, + protobuf::JoinType::Right => JoinType::Right, + protobuf::JoinType::Full => JoinType::Full, + protobuf::JoinType::Leftsemi => JoinType::LeftSemi, + protobuf::JoinType::Rightsemi => JoinType::RightSemi, + protobuf::JoinType::Leftanti => JoinType::LeftAnti, + protobuf::JoinType::Rightanti => JoinType::RightAnti, + protobuf::JoinType::Leftmark => JoinType::LeftMark, + protobuf::JoinType::Rightmark => JoinType::RightMark, + } + } +} + +impl From for JoinConstraint { + fn from(t: protobuf::JoinConstraint) -> Self { + match t { + protobuf::JoinConstraint::On => JoinConstraint::On, + protobuf::JoinConstraint::Using => JoinConstraint::Using, + } + } +} + +impl From for NullEquality { + fn from(t: protobuf::NullEquality) -> Self { + match t { + protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, + protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, + } + } +} + +impl From<&CsvOptionsProto> for CsvOptions { + fn from(proto: &CsvOptionsProto) -> Self { + CsvOptions { + has_header: if !proto.has_header.is_empty() { + Some(proto.has_header[0] != 0) + } else { + None + }, + delimiter: proto.delimiter.first().copied().unwrap_or(b','), + quote: proto.quote.first().copied().unwrap_or(b'"'), + terminator: if !proto.terminator.is_empty() { + Some(proto.terminator[0]) + } else { + None + }, + escape: if !proto.escape.is_empty() { + Some(proto.escape[0]) + } else { + None + }, + double_quote: if !proto.double_quote.is_empty() { + Some(proto.double_quote[0] != 0) + } else { + None + }, + compression: match proto.compression { + 0 => CompressionTypeVariant::GZIP, + 1 => CompressionTypeVariant::BZIP2, + 2 => CompressionTypeVariant::XZ, + 3 => CompressionTypeVariant::ZSTD, + _ => CompressionTypeVariant::UNCOMPRESSED, + }, + schema_infer_max_rec: proto.schema_infer_max_rec.map(|v| v as usize), + date_format: if proto.date_format.is_empty() { + None + } else { + Some(proto.date_format.clone()) + }, + datetime_format: if proto.datetime_format.is_empty() { + None + } else { + Some(proto.datetime_format.clone()) + }, + timestamp_format: if proto.timestamp_format.is_empty() { + None + } else { + Some(proto.timestamp_format.clone()) + }, + timestamp_tz_format: if proto.timestamp_tz_format.is_empty() { + None + } else { + Some(proto.timestamp_tz_format.clone()) + }, + time_format: if proto.time_format.is_empty() { + None + } else { + Some(proto.time_format.clone()) + }, + null_value: if proto.null_value.is_empty() { + None + } else { + Some(proto.null_value.clone()) + }, + null_regex: if proto.null_regex.is_empty() { + None + } else { + Some(proto.null_regex.clone()) + }, + comment: if !proto.comment.is_empty() { + Some(proto.comment[0]) + } else { + None + }, + newlines_in_values: if proto.newlines_in_values.is_empty() { + None + } else { + Some(proto.newlines_in_values[0] != 0) + }, + truncated_rows: if proto.truncated_rows.is_empty() { + None + } else { + Some(proto.truncated_rows[0] != 0) + }, + compression_level: proto.compression_level, + quote_style: match CsvQuoteStyleProto::try_from(proto.quote_style) { + Ok(CsvQuoteStyleProto::Always) => CsvQuoteStyle::Always, + Ok(CsvQuoteStyleProto::NonNumeric) => CsvQuoteStyle::NonNumeric, + Ok(CsvQuoteStyleProto::Never) => CsvQuoteStyle::Never, + Ok(CsvQuoteStyleProto::Necessary) => CsvQuoteStyle::Necessary, + _ => CsvQuoteStyle::Necessary, + }, + ignore_leading_whitespace: if proto.ignore_leading_whitespace.is_empty() { + None + } else { + Some(proto.ignore_leading_whitespace[0] != 0) + }, + ignore_trailing_whitespace: if proto.ignore_trailing_whitespace.is_empty() { + None + } else { + Some(proto.ignore_trailing_whitespace[0] != 0) + }, + } + } +} + +impl From<&JsonOptionsProto> for JsonOptions { + fn from(proto: &JsonOptionsProto) -> Self { + JsonOptions { + compression: match proto.compression { + 0 => CompressionTypeVariant::GZIP, + 1 => CompressionTypeVariant::BZIP2, + 2 => CompressionTypeVariant::XZ, + 3 => CompressionTypeVariant::ZSTD, + _ => CompressionTypeVariant::UNCOMPRESSED, + }, + schema_infer_max_rec: proto.schema_infer_max_rec.map(|v| v as usize), + compression_level: proto.compression_level, + newline_delimited: proto.newline_delimited.unwrap_or(true), + } + } +} + +impl From for ParquetCdcOptions { + fn from(value: ParquetCdcOptionsProto) -> Self { + ParquetCdcOptions { + enabled: value.enabled, + min_chunk_size: value.min_chunk_size as usize, + max_chunk_size: value.max_chunk_size as usize, + norm_level: value.norm_level, + } + } +} + +impl TryFrom<&ParquetOptionsProto> for ParquetOptions { + type Error = datafusion_common::DataFusionError; + + fn try_from( + proto: &ParquetOptionsProto, + ) -> datafusion_common::Result { + let writer_version = match proto.writer_version.as_str() { + // Proto3 decodes an omitted string field as the empty string. The + // schema documents writer_version's logical default as "1.0", so + // preserve that default when the field is absent on the wire. + "" => ParquetOptions::default().writer_version, + version => version.parse()?, + }; + + Ok(ParquetOptions { + enable_page_index: proto.enable_page_index, + pruning: proto.pruning, + skip_metadata: proto.skip_metadata, + metadata_size_hint: proto + .metadata_size_hint_opt + .as_ref() + .map(|opt| match opt { + parquet_options::MetadataSizeHintOpt::MetadataSizeHint(size) => { + *size as usize + } + }), + pushdown_filters: proto.pushdown_filters, + reorder_filters: proto.reorder_filters, + force_filter_selections: proto.force_filter_selections, + data_pagesize_limit: proto.data_pagesize_limit as usize, + write_batch_size: proto.write_batch_size as usize, + writer_version, + compression: proto.compression_opt.as_ref().map(|opt| match opt { + parquet_options::CompressionOpt::Compression(compression) => { + compression.clone() + } + }), + dictionary_enabled: proto.dictionary_enabled_opt.as_ref().map(|opt| { + match opt { + parquet_options::DictionaryEnabledOpt::DictionaryEnabled( + enabled, + ) => *enabled, + } + }), + dictionary_page_size_limit: proto.dictionary_page_size_limit as usize, + statistics_enabled: proto.statistics_enabled_opt.as_ref().map( + |opt| match opt { + parquet_options::StatisticsEnabledOpt::StatisticsEnabled( + statistics, + ) => statistics.clone(), + }, + ), + max_row_group_size: proto.max_row_group_size as usize, + max_in_list_size: proto.max_in_list_size as usize, + created_by: proto.created_by.clone(), + column_index_truncate_length: proto + .column_index_truncate_length_opt + .as_ref() + .map(|opt| match opt { + parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length) => *length as usize, + }), + statistics_truncate_length: proto + .statistics_truncate_length_opt + .as_ref() + .map(|opt| match opt { + parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(length) => *length as usize, + }), + data_page_row_count_limit: proto.data_page_row_count_limit as usize, + encoding: proto.encoding_opt.as_ref().map(|opt| match opt { + parquet_options::EncodingOpt::Encoding(encoding) => { + encoding.clone() + } + }), + bloom_filter_on_read: proto.bloom_filter_on_read, + bloom_filter_on_write: proto.bloom_filter_on_write, + bloom_filter_fpp: proto + .bloom_filter_fpp_opt + .as_ref() + .map(|opt| match opt { + parquet_options::BloomFilterFppOpt::BloomFilterFpp(fpp) => *fpp, + }), + bloom_filter_ndv: proto + .bloom_filter_ndv_opt + .as_ref() + .map(|opt| match opt { + parquet_options::BloomFilterNdvOpt::BloomFilterNdv(ndv) => *ndv, + }), + allow_single_file_parallelism: proto.allow_single_file_parallelism, + maximum_parallel_row_group_writers: proto + .maximum_parallel_row_group_writers + as usize, + maximum_buffered_record_batches_per_stream: proto + .maximum_buffered_record_batches_per_stream + as usize, + schema_force_view_types: proto.schema_force_view_types, + binary_as_string: proto.binary_as_string, + skip_arrow_metadata: proto.skip_arrow_metadata, + coerce_int96: proto.coerce_int96_opt.as_ref().map(|opt| match opt { + parquet_options::CoerceInt96Opt::CoerceInt96(coerce_int96) => { + coerce_int96.clone() + } + }), + coerce_int96_tz: proto + .coerce_int96_tz_opt + .as_ref() + .map(|opt| match opt { + parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(tz) => { + tz.clone() + } + }), + max_predicate_cache_size: proto + .max_predicate_cache_size_opt + .as_ref() + .map(|opt| match opt { + parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize( + size, + ) => *size as usize, + }), + max_row_group_bytes: proto + .max_row_group_bytes_opt + .as_ref() + .and_then(|opt| match opt { + parquet_options::MaxRowGroupBytesOpt::MaxRowGroupBytes(size) => { + MaxRowGroupBytes::try_new(*size as usize).ok() + } + }), + content_defined_chunking: proto + .content_defined_chunking + .map(ParquetCdcOptions::from) + .unwrap_or_default(), + }) + } +} + +impl From for ParquetColumnOptions { + fn from(proto: ParquetColumnOptionsProto) -> Self { + ParquetColumnOptions { + bloom_filter_enabled: proto.bloom_filter_enabled_opt.map( + |parquet_column_options::BloomFilterEnabledOpt::BloomFilterEnabled(v)| v, + ), + encoding: proto + .encoding_opt + .map(|parquet_column_options::EncodingOpt::Encoding(v)| v), + dictionary_enabled: proto.dictionary_enabled_opt.map( + |parquet_column_options::DictionaryEnabledOpt::DictionaryEnabled(v)| v, + ), + compression: proto + .compression_opt + .map(|parquet_column_options::CompressionOpt::Compression(v)| v), + statistics_enabled: proto.statistics_enabled_opt.map( + |parquet_column_options::StatisticsEnabledOpt::StatisticsEnabled(v)| v, + ), + bloom_filter_fpp: proto + .bloom_filter_fpp_opt + .map(|parquet_column_options::BloomFilterFppOpt::BloomFilterFpp(v)| v), + bloom_filter_ndv: proto + .bloom_filter_ndv_opt + .map(|parquet_column_options::BloomFilterNdvOpt::BloomFilterNdv(v)| v), + } + } +} + +impl TryFrom<&TableParquetOptionsProto> for TableParquetOptions { + type Error = datafusion_common::DataFusionError; + + fn try_from( + proto: &TableParquetOptionsProto, + ) -> datafusion_common::Result { + Ok(TableParquetOptions { + global: proto + .global + .as_ref() + .map(ParquetOptions::try_from) + .transpose()? + .unwrap_or_default(), + column_specific_options: proto + .column_specific_options + .iter() + .map(|parquet_column_options| { + ( + parquet_column_options.column_name.clone(), + ParquetColumnOptions::from( + parquet_column_options.options.clone().unwrap_or_default(), + ), + ) + }) + .collect(), + key_value_metadata: proto + .key_value_metadata + .iter() + .map(|(k, v)| (k.clone(), Some(v.clone()))) + .collect(), + ..Default::default() + }) + } +} diff --git a/datafusion/proto/src/generated/datafusion_proto_common.rs b/datafusion/proto-models/src/generated/datafusion_proto_common.rs similarity index 93% rename from datafusion/proto/src/generated/datafusion_proto_common.rs rename to datafusion/proto-models/src/generated/datafusion_proto_common.rs index 632b16929faa6..bdbe38538e1d7 100644 --- a/datafusion/proto/src/generated/datafusion_proto_common.rs +++ b/datafusion/proto-models/src/generated/datafusion_proto_common.rs @@ -862,10 +862,12 @@ pub struct ParquetOptions { pub data_page_row_count_limit: u64, #[prost(uint64, tag = "15")] pub max_row_group_size: u64, + #[prost(uint64, tag = "38")] + pub max_in_list_size: u64, #[prost(string, tag = "16")] pub created_by: ::prost::alloc::string::String, #[prost(message, optional, tag = "35")] - pub content_defined_chunking: ::core::option::Option, + pub content_defined_chunking: ::core::option::Option, #[prost(oneof = "parquet_options::MetadataSizeHintOpt", tags = "4")] pub metadata_size_hint_opt: ::core::option::Option< parquet_options::MetadataSizeHintOpt, @@ -900,6 +902,10 @@ pub struct ParquetOptions { pub max_predicate_cache_size_opt: ::core::option::Option< parquet_options::MaxPredicateCacheSizeOpt, >, + #[prost(oneof = "parquet_options::MaxRowGroupBytesOpt", tags = "37")] + pub max_row_group_bytes_opt: ::core::option::Option< + parquet_options::MaxRowGroupBytesOpt, + >, /// Optional timezone applied to INT96-coerced timestamps when `coerce_int96` /// is set. When `Some`, INT96 columns coerce to /// `Timestamp(, Some())` instead of the default @@ -964,6 +970,11 @@ pub mod parquet_options { #[prost(uint64, tag = "33")] MaxPredicateCacheSize(u64), } + #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Oneof)] + pub enum MaxRowGroupBytesOpt { + #[prost(uint64, tag = "37")] + MaxRowGroupBytes(u64), + } /// Optional timezone applied to INT96-coerced timestamps when `coerce_int96` /// is set. When `Some`, INT96 columns coerce to /// `Timestamp(, Some())` instead of the default @@ -974,13 +985,16 @@ pub mod parquet_options { CoerceInt96Tz(::prost::alloc::string::String), } } +/// Content-defined chunking (CDC) options for writing parquet files. #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] -pub struct CdcOptions { - #[prost(uint64, tag = "1")] - pub min_chunk_size: u64, +pub struct ParquetCdcOptions { + #[prost(bool, tag = "1")] + pub enabled: bool, #[prost(uint64, tag = "2")] + pub min_chunk_size: u64, + #[prost(uint64, tag = "3")] pub max_chunk_size: u64, - #[prost(int32, tag = "3")] + #[prost(int32, tag = "4")] pub norm_level: i32, } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1014,6 +1028,18 @@ pub struct ColumnStats { #[prost(message, optional, tag = "6")] pub byte_size: ::core::option::Option, } +/// Wire encoding for `datafusion_common::format::ExplainAnalyzeCategories`. +/// +/// If `all` is true, every category is shown (the `only` list is ignored). +/// If `all` is false, only the categories listed in `only` are shown — an +/// empty `only` means "plan only", i.e. suppress all metrics. +#[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] +pub struct ExplainAnalyzeCategoriesNode { + #[prost(bool, tag = "1")] + pub all: bool, + #[prost(enumeration = "MetricCategory", repeated, tag = "2")] + pub only: ::prost::alloc::vec::Vec, +} #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] #[repr(i32)] pub enum JoinType { @@ -1360,3 +1386,65 @@ impl ExplainFormat { } } } +/// Verbosity level for `EXPLAIN ANALYZE`. Mirrors +/// `datafusion_common::format::MetricType`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum MetricType { + Summary = 0, + Dev = 1, +} +impl MetricType { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Summary => "METRIC_TYPE_SUMMARY", + Self::Dev => "METRIC_TYPE_DEV", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "METRIC_TYPE_SUMMARY" => Some(Self::Summary), + "METRIC_TYPE_DEV" => Some(Self::Dev), + _ => None, + } + } +} +/// Category of an `EXPLAIN ANALYZE` metric. Mirrors +/// `datafusion_common::format::MetricCategory`. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord, ::prost::Enumeration)] +#[repr(i32)] +pub enum MetricCategory { + Rows = 0, + Bytes = 1, + Timing = 2, + Uncategorized = 3, +} +impl MetricCategory { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Rows => "METRIC_CATEGORY_ROWS", + Self::Bytes => "METRIC_CATEGORY_BYTES", + Self::Timing => "METRIC_CATEGORY_TIMING", + Self::Uncategorized => "METRIC_CATEGORY_UNCATEGORIZED", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "METRIC_CATEGORY_ROWS" => Some(Self::Rows), + "METRIC_CATEGORY_BYTES" => Some(Self::Bytes), + "METRIC_CATEGORY_TIMING" => Some(Self::Timing), + "METRIC_CATEGORY_UNCATEGORIZED" => Some(Self::Uncategorized), + _ => None, + } + } +} diff --git a/datafusion/proto/src/generated/mod.rs b/datafusion/proto-models/src/generated/mod.rs similarity index 97% rename from datafusion/proto/src/generated/mod.rs rename to datafusion/proto-models/src/generated/mod.rs index ca32b1500d57b..4362b741d93a9 100644 --- a/datafusion/proto/src/generated/mod.rs +++ b/datafusion/proto-models/src/generated/mod.rs @@ -18,6 +18,7 @@ // This code is generated so we don't want to fix any lint violations manually #[allow(clippy::allow_attributes)] #[allow(clippy::all)] +#[allow(unused_qualifications)] #[rustfmt::skip] pub mod datafusion { include!("prost.rs"); diff --git a/datafusion/proto/src/generated/pbjson.rs b/datafusion/proto-models/src/generated/pbjson.rs similarity index 91% rename from datafusion/proto/src/generated/pbjson.rs rename to datafusion/proto-models/src/generated/pbjson.rs index f71fabbdaca67..908f9752b7f18 100644 --- a/datafusion/proto/src/generated/pbjson.rs +++ b/datafusion/proto-models/src/generated/pbjson.rs @@ -157,6 +157,9 @@ impl serde::Serialize for AggregateExecNode { if self.dynamic_filter.is_some() { len += 1; } + if self.schema.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.AggregateExecNode", len)?; if !self.group_expr.is_empty() { struct_ser.serialize_field("groupExpr", &self.group_expr)?; @@ -199,6 +202,9 @@ impl serde::Serialize for AggregateExecNode { if let Some(v) = self.dynamic_filter.as_ref() { struct_ser.serialize_field("dynamicFilter", v)?; } + if let Some(v) = self.schema.as_ref() { + struct_ser.serialize_field("schema", v)?; + } struct_ser.end() } } @@ -231,6 +237,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { "hasGroupingSet", "dynamic_filter", "dynamicFilter", + "schema", ]; #[allow(clippy::enum_variant_names)] @@ -248,6 +255,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { Limit, HasGroupingSet, DynamicFilter, + Schema, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -282,6 +290,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { "limit" => Ok(GeneratedField::Limit), "hasGroupingSet" | "has_grouping_set" => Ok(GeneratedField::HasGroupingSet), "dynamicFilter" | "dynamic_filter" => Ok(GeneratedField::DynamicFilter), + "schema" => Ok(GeneratedField::Schema), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -314,6 +323,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { let mut limit__ = None; let mut has_grouping_set__ = None; let mut dynamic_filter__ = None; + let mut schema__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::GroupExpr => { @@ -394,6 +404,12 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { } dynamic_filter__ = map_.next_value()?; } + GeneratedField::Schema => { + if schema__.is_some() { + return Err(serde::de::Error::duplicate_field("schema")); + } + schema__ = map_.next_value()?; + } } } Ok(AggregateExecNode { @@ -410,6 +426,7 @@ impl<'de> serde::Deserialize<'de> for AggregateExecNode { limit: limit__, has_grouping_set: has_grouping_set__.unwrap_or_default(), dynamic_filter: dynamic_filter__, + schema: schema__, }) } } @@ -999,6 +1016,9 @@ impl serde::Serialize for AnalyzeExecNode { if !self.metric_categories.is_empty() { len += 1; } + if self.format != 0 { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.AnalyzeExecNode", len)?; if self.verbose { struct_ser.serialize_field("verbose", &self.verbose)?; @@ -1018,6 +1038,11 @@ impl serde::Serialize for AnalyzeExecNode { if !self.metric_categories.is_empty() { struct_ser.serialize_field("metricCategories", &self.metric_categories)?; } + if self.format != 0 { + let v = super::datafusion_common::ExplainFormat::try_from(self.format) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.format)))?; + struct_ser.serialize_field("format", &v)?; + } struct_ser.end() } } @@ -1037,6 +1062,7 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { "hasMetricCategories", "metric_categories", "metricCategories", + "format", ]; #[allow(clippy::enum_variant_names)] @@ -1047,6 +1073,7 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { Schema, HasMetricCategories, MetricCategories, + Format, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -1074,6 +1101,7 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { "schema" => Ok(GeneratedField::Schema), "hasMetricCategories" | "has_metric_categories" => Ok(GeneratedField::HasMetricCategories), "metricCategories" | "metric_categories" => Ok(GeneratedField::MetricCategories), + "format" => Ok(GeneratedField::Format), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -1099,6 +1127,7 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { let mut schema__ = None; let mut has_metric_categories__ = None; let mut metric_categories__ = None; + let mut format__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Verbose => { @@ -1137,6 +1166,12 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { } metric_categories__ = Some(map_.next_value()?); } + GeneratedField::Format => { + if format__.is_some() { + return Err(serde::de::Error::duplicate_field("format")); + } + format__ = Some(map_.next_value::()? as i32); + } } } Ok(AnalyzeExecNode { @@ -1146,6 +1181,7 @@ impl<'de> serde::Deserialize<'de> for AnalyzeExecNode { schema: schema__, has_metric_categories: has_metric_categories__.unwrap_or_default(), metric_categories: metric_categories__.unwrap_or_default(), + format: format__.unwrap_or_default(), }) } } @@ -1166,6 +1202,15 @@ impl serde::Serialize for AnalyzeNode { if self.verbose { len += 1; } + if self.analyze_level.is_some() { + len += 1; + } + if self.analyze_categories.is_some() { + len += 1; + } + if self.format != 0 { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.AnalyzeNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; @@ -1173,6 +1218,19 @@ impl serde::Serialize for AnalyzeNode { if self.verbose { struct_ser.serialize_field("verbose", &self.verbose)?; } + if let Some(v) = self.analyze_level.as_ref() { + let v = super::datafusion_common::MetricType::try_from(*v) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", *v)))?; + struct_ser.serialize_field("analyzeLevel", &v)?; + } + if let Some(v) = self.analyze_categories.as_ref() { + struct_ser.serialize_field("analyzeCategories", v)?; + } + if self.format != 0 { + let v = super::datafusion_common::ExplainFormat::try_from(self.format) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.format)))?; + struct_ser.serialize_field("format", &v)?; + } struct_ser.end() } } @@ -1185,12 +1243,20 @@ impl<'de> serde::Deserialize<'de> for AnalyzeNode { const FIELDS: &[&str] = &[ "input", "verbose", + "analyze_level", + "analyzeLevel", + "analyze_categories", + "analyzeCategories", + "format", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Input, Verbose, + AnalyzeLevel, + AnalyzeCategories, + Format, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -1214,6 +1280,9 @@ impl<'de> serde::Deserialize<'de> for AnalyzeNode { match value { "input" => Ok(GeneratedField::Input), "verbose" => Ok(GeneratedField::Verbose), + "analyzeLevel" | "analyze_level" => Ok(GeneratedField::AnalyzeLevel), + "analyzeCategories" | "analyze_categories" => Ok(GeneratedField::AnalyzeCategories), + "format" => Ok(GeneratedField::Format), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -1235,6 +1304,9 @@ impl<'de> serde::Deserialize<'de> for AnalyzeNode { { let mut input__ = None; let mut verbose__ = None; + let mut analyze_level__ = None; + let mut analyze_categories__ = None; + let mut format__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Input => { @@ -1249,11 +1321,32 @@ impl<'de> serde::Deserialize<'de> for AnalyzeNode { } verbose__ = Some(map_.next_value()?); } + GeneratedField::AnalyzeLevel => { + if analyze_level__.is_some() { + return Err(serde::de::Error::duplicate_field("analyzeLevel")); + } + analyze_level__ = map_.next_value::<::std::option::Option>()?.map(|x| x as i32); + } + GeneratedField::AnalyzeCategories => { + if analyze_categories__.is_some() { + return Err(serde::de::Error::duplicate_field("analyzeCategories")); + } + analyze_categories__ = map_.next_value()?; + } + GeneratedField::Format => { + if format__.is_some() { + return Err(serde::de::Error::duplicate_field("format")); + } + format__ = Some(map_.next_value::()? as i32); + } } } Ok(AnalyzeNode { input: input__, verbose: verbose__.unwrap_or_default(), + analyze_level: analyze_level__, + analyze_categories: analyze_categories__, + format: format__.unwrap_or_default(), }) } } @@ -3561,6 +3654,9 @@ impl serde::Serialize for CreateExternalTableNode { if !self.location.is_empty() { len += 1; } + if !self.locations.is_empty() { + len += 1; + } if !self.file_type.is_empty() { len += 1; } @@ -3604,6 +3700,9 @@ impl serde::Serialize for CreateExternalTableNode { if !self.location.is_empty() { struct_ser.serialize_field("location", &self.location)?; } + if !self.locations.is_empty() { + struct_ser.serialize_field("locations", &self.locations)?; + } if !self.file_type.is_empty() { struct_ser.serialize_field("fileType", &self.file_type)?; } @@ -3652,6 +3751,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { const FIELDS: &[&str] = &[ "name", "location", + "locations", "file_type", "fileType", "schema", @@ -3676,6 +3776,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { enum GeneratedField { Name, Location, + Locations, FileType, Schema, TablePartitionCols, @@ -3711,6 +3812,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { match value { "name" => Ok(GeneratedField::Name), "location" => Ok(GeneratedField::Location), + "locations" => Ok(GeneratedField::Locations), "fileType" | "file_type" => Ok(GeneratedField::FileType), "schema" => Ok(GeneratedField::Schema), "tablePartitionCols" | "table_partition_cols" => Ok(GeneratedField::TablePartitionCols), @@ -3744,6 +3846,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { { let mut name__ = None; let mut location__ = None; + let mut locations__ = None; let mut file_type__ = None; let mut schema__ = None; let mut table_partition_cols__ = None; @@ -3770,6 +3873,12 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { } location__ = Some(map_.next_value()?); } + GeneratedField::Locations => { + if locations__.is_some() { + return Err(serde::de::Error::duplicate_field("locations")); + } + locations__ = Some(map_.next_value()?); + } GeneratedField::FileType => { if file_type__.is_some() { return Err(serde::de::Error::duplicate_field("fileType")); @@ -3851,6 +3960,7 @@ impl<'de> serde::Deserialize<'de> for CreateExternalTableNode { Ok(CreateExternalTableNode { name: name__, location: location__.unwrap_or_default(), + locations: locations__.unwrap_or_default(), file_type: file_type__.unwrap_or_default(), schema: schema__, table_partition_cols: table_partition_cols__.unwrap_or_default(), @@ -5411,6 +5521,9 @@ impl serde::Serialize for DmlNode { if self.target.is_some() { len += 1; } + if self.merge_into.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.DmlNode", len)?; if self.dml_type != 0 { let v = dml_node::Type::try_from(self.dml_type) @@ -5426,6 +5539,9 @@ impl serde::Serialize for DmlNode { if let Some(v) = self.target.as_ref() { struct_ser.serialize_field("target", v)?; } + if let Some(v) = self.merge_into.as_ref() { + struct_ser.serialize_field("mergeInto", v)?; + } struct_ser.end() } } @@ -5442,6 +5558,8 @@ impl<'de> serde::Deserialize<'de> for DmlNode { "table_name", "tableName", "target", + "merge_into", + "mergeInto", ]; #[allow(clippy::enum_variant_names)] @@ -5450,6 +5568,7 @@ impl<'de> serde::Deserialize<'de> for DmlNode { Input, TableName, Target, + MergeInto, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -5475,6 +5594,7 @@ impl<'de> serde::Deserialize<'de> for DmlNode { "input" => Ok(GeneratedField::Input), "tableName" | "table_name" => Ok(GeneratedField::TableName), "target" => Ok(GeneratedField::Target), + "mergeInto" | "merge_into" => Ok(GeneratedField::MergeInto), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -5498,6 +5618,7 @@ impl<'de> serde::Deserialize<'de> for DmlNode { let mut input__ = None; let mut table_name__ = None; let mut target__ = None; + let mut merge_into__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::DmlType => { @@ -5524,6 +5645,12 @@ impl<'de> serde::Deserialize<'de> for DmlNode { } target__ = map_.next_value()?; } + GeneratedField::MergeInto => { + if merge_into__.is_some() { + return Err(serde::de::Error::duplicate_field("mergeInto")); + } + merge_into__ = map_.next_value()?; + } } } Ok(DmlNode { @@ -5531,6 +5658,7 @@ impl<'de> serde::Deserialize<'de> for DmlNode { input: input__, table_name: table_name__, target: target__, + merge_into: merge_into__, }) } } @@ -5551,6 +5679,7 @@ impl serde::Serialize for dml_node::Type { Self::InsertOverwrite => "INSERT_OVERWRITE", Self::InsertReplace => "INSERT_REPLACE", Self::Truncate => "TRUNCATE", + Self::MergeInto => "MERGE_INTO", }; serializer.serialize_str(variant) } @@ -5569,6 +5698,7 @@ impl<'de> serde::Deserialize<'de> for dml_node::Type { "INSERT_OVERWRITE", "INSERT_REPLACE", "TRUNCATE", + "MERGE_INTO", ]; struct GeneratedVisitor; @@ -5616,6 +5746,7 @@ impl<'de> serde::Deserialize<'de> for dml_node::Type { "INSERT_OVERWRITE" => Ok(dml_node::Type::InsertOverwrite), "INSERT_REPLACE" => Ok(dml_node::Type::InsertReplace), "TRUNCATE" => Ok(dml_node::Type::Truncate), + "MERGE_INTO" => Ok(dml_node::Type::MergeInto), _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), } } @@ -5760,10 +5891,16 @@ impl serde::Serialize for EmptyExecNode { if self.schema.is_some() { len += 1; } + if self.partitions != 0 { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.EmptyExecNode", len)?; if let Some(v) = self.schema.as_ref() { struct_ser.serialize_field("schema", v)?; } + if self.partitions != 0 { + struct_ser.serialize_field("partitions", &self.partitions)?; + } struct_ser.end() } } @@ -5775,11 +5912,13 @@ impl<'de> serde::Deserialize<'de> for EmptyExecNode { { const FIELDS: &[&str] = &[ "schema", + "partitions", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Schema, + Partitions, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -5802,6 +5941,7 @@ impl<'de> serde::Deserialize<'de> for EmptyExecNode { { match value { "schema" => Ok(GeneratedField::Schema), + "partitions" => Ok(GeneratedField::Partitions), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -5822,6 +5962,7 @@ impl<'de> serde::Deserialize<'de> for EmptyExecNode { V: serde::de::MapAccess<'de>, { let mut schema__ = None; + let mut partitions__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Schema => { @@ -5830,10 +5971,19 @@ impl<'de> serde::Deserialize<'de> for EmptyExecNode { } schema__ = map_.next_value()?; } + GeneratedField::Partitions => { + if partitions__.is_some() { + return Err(serde::de::Error::duplicate_field("partitions")); + } + partitions__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } } } Ok(EmptyExecNode { schema: schema__, + partitions: partitions__.unwrap_or_default(), }) } } @@ -6218,6 +6368,9 @@ impl serde::Serialize for ExplainNode { if self.format != 0 { len += 1; } + if self.show_statistics.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.ExplainNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; @@ -6230,6 +6383,9 @@ impl serde::Serialize for ExplainNode { .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.format)))?; struct_ser.serialize_field("format", &v)?; } + if let Some(v) = self.show_statistics.as_ref() { + struct_ser.serialize_field("showStatistics", v)?; + } struct_ser.end() } } @@ -6243,6 +6399,8 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { "input", "verbose", "format", + "show_statistics", + "showStatistics", ]; #[allow(clippy::enum_variant_names)] @@ -6250,6 +6408,7 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { Input, Verbose, Format, + ShowStatistics, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -6274,6 +6433,7 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { "input" => Ok(GeneratedField::Input), "verbose" => Ok(GeneratedField::Verbose), "format" => Ok(GeneratedField::Format), + "showStatistics" | "show_statistics" => Ok(GeneratedField::ShowStatistics), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -6296,6 +6456,7 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { let mut input__ = None; let mut verbose__ = None; let mut format__ = None; + let mut show_statistics__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Input => { @@ -6316,12 +6477,19 @@ impl<'de> serde::Deserialize<'de> for ExplainNode { } format__ = Some(map_.next_value::()? as i32); } + GeneratedField::ShowStatistics => { + if show_statistics__.is_some() { + return Err(serde::de::Error::duplicate_field("showStatistics")); + } + show_statistics__ = map_.next_value()?; + } } } Ok(ExplainNode { input: input__, verbose: verbose__.unwrap_or_default(), format: format__.unwrap_or_default(), + show_statistics: show_statistics__, }) } } @@ -6848,6 +7016,9 @@ impl serde::Serialize for FileScanExecConf { if self.projection_exprs.is_some() { len += 1; } + if self.output_partitioning.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.FileScanExecConf", len)?; if !self.file_groups.is_empty() { struct_ser.serialize_field("fileGroups", &self.file_groups)?; @@ -6884,6 +7055,9 @@ impl serde::Serialize for FileScanExecConf { if let Some(v) = self.projection_exprs.as_ref() { struct_ser.serialize_field("projectionExprs", v)?; } + if let Some(v) = self.output_partitioning.as_ref() { + struct_ser.serialize_field("outputPartitioning", v)?; + } struct_ser.end() } } @@ -6911,6 +7085,8 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "batchSize", "projection_exprs", "projectionExprs", + "output_partitioning", + "outputPartitioning", ]; #[allow(clippy::enum_variant_names)] @@ -6926,6 +7102,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { Constraints, BatchSize, ProjectionExprs, + OutputPartitioning, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -6958,6 +7135,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { "constraints" => Ok(GeneratedField::Constraints), "batchSize" | "batch_size" => Ok(GeneratedField::BatchSize), "projectionExprs" | "projection_exprs" => Ok(GeneratedField::ProjectionExprs), + "outputPartitioning" | "output_partitioning" => Ok(GeneratedField::OutputPartitioning), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -6988,6 +7166,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { let mut constraints__ = None; let mut batch_size__ = None; let mut projection_exprs__ = None; + let mut output_partitioning__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::FileGroups => { @@ -7061,6 +7240,12 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { } projection_exprs__ = map_.next_value()?; } + GeneratedField::OutputPartitioning => { + if output_partitioning__.is_some() { + return Err(serde::de::Error::duplicate_field("outputPartitioning")); + } + output_partitioning__ = map_.next_value()?; + } } } Ok(FileScanExecConf { @@ -7075,6 +7260,7 @@ impl<'de> serde::Deserialize<'de> for FileScanExecConf { constraints: constraints__, batch_size: batch_size__, projection_exprs: projection_exprs__, + output_partitioning: output_partitioning__, }) } } @@ -8614,6 +8800,9 @@ impl serde::Serialize for GlobalLimitExecNode { if self.fetch != 0 { len += 1; } + if !self.required_ordering.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.GlobalLimitExecNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; @@ -8626,6 +8815,9 @@ impl serde::Serialize for GlobalLimitExecNode { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("fetch", ToString::to_string(&self.fetch).as_str())?; } + if !self.required_ordering.is_empty() { + struct_ser.serialize_field("requiredOrdering", &self.required_ordering)?; + } struct_ser.end() } } @@ -8639,6 +8831,8 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { "input", "skip", "fetch", + "required_ordering", + "requiredOrdering", ]; #[allow(clippy::enum_variant_names)] @@ -8646,6 +8840,7 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { Input, Skip, Fetch, + RequiredOrdering, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -8670,6 +8865,7 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { "input" => Ok(GeneratedField::Input), "skip" => Ok(GeneratedField::Skip), "fetch" => Ok(GeneratedField::Fetch), + "requiredOrdering" | "required_ordering" => Ok(GeneratedField::RequiredOrdering), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -8692,6 +8888,7 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { let mut input__ = None; let mut skip__ = None; let mut fetch__ = None; + let mut required_ordering__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Input => { @@ -8716,12 +8913,19 @@ impl<'de> serde::Deserialize<'de> for GlobalLimitExecNode { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::RequiredOrdering => { + if required_ordering__.is_some() { + return Err(serde::de::Error::duplicate_field("requiredOrdering")); + } + required_ordering__ = Some(map_.next_value()?); + } } } Ok(GlobalLimitExecNode { input: input__, skip: skip__.unwrap_or_default(), fetch: fetch__.unwrap_or_default(), + required_ordering: required_ordering__.unwrap_or_default(), }) } } @@ -8857,6 +9061,9 @@ impl serde::Serialize for HashJoinExecNode { if self.dynamic_filter.is_some() { len += 1; } + if self.fetch.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.HashJoinExecNode", len)?; if let Some(v) = self.left.as_ref() { struct_ser.serialize_field("left", v)?; @@ -8894,6 +9101,11 @@ impl serde::Serialize for HashJoinExecNode { if let Some(v) = self.dynamic_filter.as_ref() { struct_ser.serialize_field("dynamicFilter", v)?; } + if let Some(v) = self.fetch.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("fetch", ToString::to_string(&v).as_str())?; + } struct_ser.end() } } @@ -8919,6 +9131,7 @@ impl<'de> serde::Deserialize<'de> for HashJoinExecNode { "nullAware", "dynamic_filter", "dynamicFilter", + "fetch", ]; #[allow(clippy::enum_variant_names)] @@ -8933,6 +9146,7 @@ impl<'de> serde::Deserialize<'de> for HashJoinExecNode { Projection, NullAware, DynamicFilter, + Fetch, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -8964,6 +9178,7 @@ impl<'de> serde::Deserialize<'de> for HashJoinExecNode { "projection" => Ok(GeneratedField::Projection), "nullAware" | "null_aware" => Ok(GeneratedField::NullAware), "dynamicFilter" | "dynamic_filter" => Ok(GeneratedField::DynamicFilter), + "fetch" => Ok(GeneratedField::Fetch), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -8993,6 +9208,7 @@ impl<'de> serde::Deserialize<'de> for HashJoinExecNode { let mut projection__ = None; let mut null_aware__ = None; let mut dynamic_filter__ = None; + let mut fetch__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Left => { @@ -9058,6 +9274,14 @@ impl<'de> serde::Deserialize<'de> for HashJoinExecNode { } dynamic_filter__ = map_.next_value()?; } + GeneratedField::Fetch => { + if fetch__.is_some() { + return Err(serde::de::Error::duplicate_field("fetch")); + } + fetch__ = + map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| x.0) + ; + } } } Ok(HashJoinExecNode { @@ -9071,6 +9295,7 @@ impl<'de> serde::Deserialize<'de> for HashJoinExecNode { projection: projection__.unwrap_or_default(), null_aware: null_aware__.unwrap_or_default(), dynamic_filter: dynamic_filter__, + fetch: fetch__, }) } } @@ -9191,7 +9416,7 @@ impl<'de> serde::Deserialize<'de> for HashRepartition { deserializer.deserialize_struct("datafusion.HashRepartition", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for ILikeNode { +impl serde::Serialize for HigherOrderUdfExprNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -9199,54 +9424,49 @@ impl serde::Serialize for ILikeNode { { use serde::ser::SerializeStruct; let mut len = 0; - if self.negated { - len += 1; - } - if self.expr.is_some() { + if !self.fun_name.is_empty() { len += 1; } - if self.pattern.is_some() { + if !self.args.is_empty() { len += 1; } - if !self.escape_char.is_empty() { + if self.fun_definition.is_some() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.ILikeNode", len)?; - if self.negated { - struct_ser.serialize_field("negated", &self.negated)?; - } - if let Some(v) = self.expr.as_ref() { - struct_ser.serialize_field("expr", v)?; + let mut struct_ser = serializer.serialize_struct("datafusion.HigherOrderUDFExprNode", len)?; + if !self.fun_name.is_empty() { + struct_ser.serialize_field("funName", &self.fun_name)?; } - if let Some(v) = self.pattern.as_ref() { - struct_ser.serialize_field("pattern", v)?; + if !self.args.is_empty() { + struct_ser.serialize_field("args", &self.args)?; } - if !self.escape_char.is_empty() { - struct_ser.serialize_field("escapeChar", &self.escape_char)?; + if let Some(v) = self.fun_definition.as_ref() { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("funDefinition", pbjson::private::base64::encode(&v).as_str())?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for ILikeNode { +impl<'de> serde::Deserialize<'de> for HigherOrderUdfExprNode { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "negated", - "expr", - "pattern", - "escape_char", - "escapeChar", + "fun_name", + "funName", + "args", + "fun_definition", + "funDefinition", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - Negated, - Expr, - Pattern, - EscapeChar, + FunName, + Args, + FunDefinition, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -9268,10 +9488,9 @@ impl<'de> serde::Deserialize<'de> for ILikeNode { E: serde::de::Error, { match value { - "negated" => Ok(GeneratedField::Negated), - "expr" => Ok(GeneratedField::Expr), - "pattern" => Ok(GeneratedField::Pattern), - "escapeChar" | "escape_char" => Ok(GeneratedField::EscapeChar), + "funName" | "fun_name" => Ok(GeneratedField::FunName), + "args" => Ok(GeneratedField::Args), + "funDefinition" | "fun_definition" => Ok(GeneratedField::FunDefinition), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -9281,60 +9500,54 @@ impl<'de> serde::Deserialize<'de> for ILikeNode { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = ILikeNode; + type Value = HigherOrderUdfExprNode; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.ILikeNode") + formatter.write_str("struct datafusion.HigherOrderUDFExprNode") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut negated__ = None; - let mut expr__ = None; - let mut pattern__ = None; - let mut escape_char__ = None; + let mut fun_name__ = None; + let mut args__ = None; + let mut fun_definition__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::Negated => { - if negated__.is_some() { - return Err(serde::de::Error::duplicate_field("negated")); - } - negated__ = Some(map_.next_value()?); - } - GeneratedField::Expr => { - if expr__.is_some() { - return Err(serde::de::Error::duplicate_field("expr")); + GeneratedField::FunName => { + if fun_name__.is_some() { + return Err(serde::de::Error::duplicate_field("funName")); } - expr__ = map_.next_value()?; + fun_name__ = Some(map_.next_value()?); } - GeneratedField::Pattern => { - if pattern__.is_some() { - return Err(serde::de::Error::duplicate_field("pattern")); + GeneratedField::Args => { + if args__.is_some() { + return Err(serde::de::Error::duplicate_field("args")); } - pattern__ = map_.next_value()?; + args__ = Some(map_.next_value()?); } - GeneratedField::EscapeChar => { - if escape_char__.is_some() { - return Err(serde::de::Error::duplicate_field("escapeChar")); + GeneratedField::FunDefinition => { + if fun_definition__.is_some() { + return Err(serde::de::Error::duplicate_field("funDefinition")); } - escape_char__ = Some(map_.next_value()?); + fun_definition__ = + map_.next_value::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| x.0) + ; } } } - Ok(ILikeNode { - negated: negated__.unwrap_or_default(), - expr: expr__, - pattern: pattern__, - escape_char: escape_char__.unwrap_or_default(), + Ok(HigherOrderUdfExprNode { + fun_name: fun_name__.unwrap_or_default(), + args: args__.unwrap_or_default(), + fun_definition: fun_definition__, }) } } - deserializer.deserialize_struct("datafusion.ILikeNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.HigherOrderUDFExprNode", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for InListNode { +impl serde::Serialize for ILikeNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -9342,21 +9555,164 @@ impl serde::Serialize for InListNode { { use serde::ser::SerializeStruct; let mut len = 0; - if self.expr.is_some() { + if self.negated { len += 1; } - if !self.list.is_empty() { + if self.expr.is_some() { len += 1; } - if self.negated { + if self.pattern.is_some() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.InListNode", len)?; - if let Some(v) = self.expr.as_ref() { - struct_ser.serialize_field("expr", v)?; - } - if !self.list.is_empty() { - struct_ser.serialize_field("list", &self.list)?; + if !self.escape_char.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.ILikeNode", len)?; + if self.negated { + struct_ser.serialize_field("negated", &self.negated)?; + } + if let Some(v) = self.expr.as_ref() { + struct_ser.serialize_field("expr", v)?; + } + if let Some(v) = self.pattern.as_ref() { + struct_ser.serialize_field("pattern", v)?; + } + if !self.escape_char.is_empty() { + struct_ser.serialize_field("escapeChar", &self.escape_char)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for ILikeNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "negated", + "expr", + "pattern", + "escape_char", + "escapeChar", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Negated, + Expr, + Pattern, + EscapeChar, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "negated" => Ok(GeneratedField::Negated), + "expr" => Ok(GeneratedField::Expr), + "pattern" => Ok(GeneratedField::Pattern), + "escapeChar" | "escape_char" => Ok(GeneratedField::EscapeChar), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = ILikeNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.ILikeNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut negated__ = None; + let mut expr__ = None; + let mut pattern__ = None; + let mut escape_char__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Negated => { + if negated__.is_some() { + return Err(serde::de::Error::duplicate_field("negated")); + } + negated__ = Some(map_.next_value()?); + } + GeneratedField::Expr => { + if expr__.is_some() { + return Err(serde::de::Error::duplicate_field("expr")); + } + expr__ = map_.next_value()?; + } + GeneratedField::Pattern => { + if pattern__.is_some() { + return Err(serde::de::Error::duplicate_field("pattern")); + } + pattern__ = map_.next_value()?; + } + GeneratedField::EscapeChar => { + if escape_char__.is_some() { + return Err(serde::de::Error::duplicate_field("escapeChar")); + } + escape_char__ = Some(map_.next_value()?); + } + } + } + Ok(ILikeNode { + negated: negated__.unwrap_or_default(), + expr: expr__, + pattern: pattern__, + escape_char: escape_char__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.ILikeNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for InListNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.expr.is_some() { + len += 1; + } + if !self.list.is_empty() { + len += 1; + } + if self.negated { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.InListNode", len)?; + if let Some(v) = self.expr.as_ref() { + struct_ser.serialize_field("expr", v)?; + } + if !self.list.is_empty() { + struct_ser.serialize_field("list", &self.list)?; } if self.negated { struct_ser.serialize_field("negated", &self.negated)?; @@ -10510,6 +10866,9 @@ impl serde::Serialize for JoinNode { if self.filter.is_some() { len += 1; } + if self.null_aware { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.JoinNode", len)?; if let Some(v) = self.left.as_ref() { struct_ser.serialize_field("left", v)?; @@ -10541,6 +10900,9 @@ impl serde::Serialize for JoinNode { if let Some(v) = self.filter.as_ref() { struct_ser.serialize_field("filter", v)?; } + if self.null_aware { + struct_ser.serialize_field("nullAware", &self.null_aware)?; + } struct_ser.end() } } @@ -10564,6 +10926,8 @@ impl<'de> serde::Deserialize<'de> for JoinNode { "null_equality", "nullEquality", "filter", + "null_aware", + "nullAware", ]; #[allow(clippy::enum_variant_names)] @@ -10576,6 +10940,7 @@ impl<'de> serde::Deserialize<'de> for JoinNode { RightJoinKey, NullEquality, Filter, + NullAware, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -10605,6 +10970,7 @@ impl<'de> serde::Deserialize<'de> for JoinNode { "rightJoinKey" | "right_join_key" => Ok(GeneratedField::RightJoinKey), "nullEquality" | "null_equality" => Ok(GeneratedField::NullEquality), "filter" => Ok(GeneratedField::Filter), + "nullAware" | "null_aware" => Ok(GeneratedField::NullAware), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -10632,6 +10998,7 @@ impl<'de> serde::Deserialize<'de> for JoinNode { let mut right_join_key__ = None; let mut null_equality__ = None; let mut filter__ = None; + let mut null_aware__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Left => { @@ -10682,6 +11049,12 @@ impl<'de> serde::Deserialize<'de> for JoinNode { } filter__ = map_.next_value()?; } + GeneratedField::NullAware => { + if null_aware__.is_some() { + return Err(serde::de::Error::duplicate_field("nullAware")); + } + null_aware__ = Some(map_.next_value()?); + } } } Ok(JoinNode { @@ -10693,6 +11066,7 @@ impl<'de> serde::Deserialize<'de> for JoinNode { right_join_key: right_join_key__.unwrap_or_default(), null_equality: null_equality__.unwrap_or_default(), filter: filter__, + null_aware: null_aware__.unwrap_or_default(), }) } } @@ -11152,7 +11526,7 @@ impl<'de> serde::Deserialize<'de> for JsonSinkExecNode { deserializer.deserialize_struct("datafusion.JsonSinkExecNode", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for LikeNode { +impl serde::Serialize for Lambda { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -11160,54 +11534,37 @@ impl serde::Serialize for LikeNode { { use serde::ser::SerializeStruct; let mut len = 0; - if self.negated { - len += 1; - } - if self.expr.is_some() { - len += 1; - } - if self.pattern.is_some() { + if !self.params.is_empty() { len += 1; } - if !self.escape_char.is_empty() { + if self.body.is_some() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.LikeNode", len)?; - if self.negated { - struct_ser.serialize_field("negated", &self.negated)?; - } - if let Some(v) = self.expr.as_ref() { - struct_ser.serialize_field("expr", v)?; - } - if let Some(v) = self.pattern.as_ref() { - struct_ser.serialize_field("pattern", v)?; + let mut struct_ser = serializer.serialize_struct("datafusion.Lambda", len)?; + if !self.params.is_empty() { + struct_ser.serialize_field("params", &self.params)?; } - if !self.escape_char.is_empty() { - struct_ser.serialize_field("escapeChar", &self.escape_char)?; + if let Some(v) = self.body.as_ref() { + struct_ser.serialize_field("body", v)?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for LikeNode { +impl<'de> serde::Deserialize<'de> for Lambda { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "negated", - "expr", - "pattern", - "escape_char", - "escapeChar", + "params", + "body", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - Negated, - Expr, - Pattern, - EscapeChar, + Params, + Body, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -11229,10 +11586,8 @@ impl<'de> serde::Deserialize<'de> for LikeNode { E: serde::de::Error, { match value { - "negated" => Ok(GeneratedField::Negated), - "expr" => Ok(GeneratedField::Expr), - "pattern" => Ok(GeneratedField::Pattern), - "escapeChar" | "escape_char" => Ok(GeneratedField::EscapeChar), + "params" => Ok(GeneratedField::Params), + "body" => Ok(GeneratedField::Body), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -11242,60 +11597,44 @@ impl<'de> serde::Deserialize<'de> for LikeNode { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = LikeNode; + type Value = Lambda; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.LikeNode") + formatter.write_str("struct datafusion.Lambda") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut negated__ = None; - let mut expr__ = None; - let mut pattern__ = None; - let mut escape_char__ = None; + let mut params__ = None; + let mut body__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::Negated => { - if negated__.is_some() { - return Err(serde::de::Error::duplicate_field("negated")); + GeneratedField::Params => { + if params__.is_some() { + return Err(serde::de::Error::duplicate_field("params")); } - negated__ = Some(map_.next_value()?); - } - GeneratedField::Expr => { - if expr__.is_some() { - return Err(serde::de::Error::duplicate_field("expr")); - } - expr__ = map_.next_value()?; - } - GeneratedField::Pattern => { - if pattern__.is_some() { - return Err(serde::de::Error::duplicate_field("pattern")); - } - pattern__ = map_.next_value()?; + params__ = Some(map_.next_value()?); } - GeneratedField::EscapeChar => { - if escape_char__.is_some() { - return Err(serde::de::Error::duplicate_field("escapeChar")); + GeneratedField::Body => { + if body__.is_some() { + return Err(serde::de::Error::duplicate_field("body")); } - escape_char__ = Some(map_.next_value()?); + body__ = map_.next_value()?; } } } - Ok(LikeNode { - negated: negated__.unwrap_or_default(), - expr: expr__, - pattern: pattern__, - escape_char: escape_char__.unwrap_or_default(), + Ok(Lambda { + params: params__.unwrap_or_default(), + body: body__, }) } } - deserializer.deserialize_struct("datafusion.LikeNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.Lambda", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for LimitNode { +impl serde::Serialize for LambdaVariable { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -11303,42 +11642,293 @@ impl serde::Serialize for LimitNode { { use serde::ser::SerializeStruct; let mut len = 0; - if self.input.is_some() { - len += 1; - } - if self.skip != 0 { + if !self.name.is_empty() { len += 1; } - if self.fetch != 0 { + if self.field.is_some() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.LimitNode", len)?; - if let Some(v) = self.input.as_ref() { - struct_ser.serialize_field("input", v)?; - } - if self.skip != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("skip", ToString::to_string(&self.skip).as_str())?; + let mut struct_ser = serializer.serialize_struct("datafusion.LambdaVariable", len)?; + if !self.name.is_empty() { + struct_ser.serialize_field("name", &self.name)?; } - if self.fetch != 0 { - #[allow(clippy::needless_borrow)] - #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("fetch", ToString::to_string(&self.fetch).as_str())?; + if let Some(v) = self.field.as_ref() { + struct_ser.serialize_field("field", v)?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for LimitNode { +impl<'de> serde::Deserialize<'de> for LambdaVariable { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "input", - "skip", - "fetch", + "name", + "field", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Name, + Field, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "name" => Ok(GeneratedField::Name), + "field" => Ok(GeneratedField::Field), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = LambdaVariable; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.LambdaVariable") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut name__ = None; + let mut field__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Name => { + if name__.is_some() { + return Err(serde::de::Error::duplicate_field("name")); + } + name__ = Some(map_.next_value()?); + } + GeneratedField::Field => { + if field__.is_some() { + return Err(serde::de::Error::duplicate_field("field")); + } + field__ = map_.next_value()?; + } + } + } + Ok(LambdaVariable { + name: name__.unwrap_or_default(), + field: field__, + }) + } + } + deserializer.deserialize_struct("datafusion.LambdaVariable", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for LikeNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.negated { + len += 1; + } + if self.expr.is_some() { + len += 1; + } + if self.pattern.is_some() { + len += 1; + } + if !self.escape_char.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.LikeNode", len)?; + if self.negated { + struct_ser.serialize_field("negated", &self.negated)?; + } + if let Some(v) = self.expr.as_ref() { + struct_ser.serialize_field("expr", v)?; + } + if let Some(v) = self.pattern.as_ref() { + struct_ser.serialize_field("pattern", v)?; + } + if !self.escape_char.is_empty() { + struct_ser.serialize_field("escapeChar", &self.escape_char)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for LikeNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "negated", + "expr", + "pattern", + "escape_char", + "escapeChar", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Negated, + Expr, + Pattern, + EscapeChar, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "negated" => Ok(GeneratedField::Negated), + "expr" => Ok(GeneratedField::Expr), + "pattern" => Ok(GeneratedField::Pattern), + "escapeChar" | "escape_char" => Ok(GeneratedField::EscapeChar), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = LikeNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.LikeNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut negated__ = None; + let mut expr__ = None; + let mut pattern__ = None; + let mut escape_char__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Negated => { + if negated__.is_some() { + return Err(serde::de::Error::duplicate_field("negated")); + } + negated__ = Some(map_.next_value()?); + } + GeneratedField::Expr => { + if expr__.is_some() { + return Err(serde::de::Error::duplicate_field("expr")); + } + expr__ = map_.next_value()?; + } + GeneratedField::Pattern => { + if pattern__.is_some() { + return Err(serde::de::Error::duplicate_field("pattern")); + } + pattern__ = map_.next_value()?; + } + GeneratedField::EscapeChar => { + if escape_char__.is_some() { + return Err(serde::de::Error::duplicate_field("escapeChar")); + } + escape_char__ = Some(map_.next_value()?); + } + } + } + Ok(LikeNode { + negated: negated__.unwrap_or_default(), + expr: expr__, + pattern: pattern__, + escape_char: escape_char__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.LikeNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for LimitNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.input.is_some() { + len += 1; + } + if self.skip != 0 { + len += 1; + } + if self.fetch != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.LimitNode", len)?; + if let Some(v) = self.input.as_ref() { + struct_ser.serialize_field("input", v)?; + } + if self.skip != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("skip", ToString::to_string(&self.skip).as_str())?; + } + if self.fetch != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("fetch", ToString::to_string(&self.fetch).as_str())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for LimitNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "input", + "skip", + "fetch", ]; #[allow(clippy::enum_variant_names)] @@ -11786,12 +12376,6 @@ impl serde::Serialize for ListingTableScanNode { if !self.table_partition_cols.is_empty() { len += 1; } - if self.collect_stat { - len += 1; - } - if self.target_partitions != 0 { - len += 1; - } if !self.file_sort_order.is_empty() { len += 1; } @@ -11820,12 +12404,6 @@ impl serde::Serialize for ListingTableScanNode { if !self.table_partition_cols.is_empty() { struct_ser.serialize_field("tablePartitionCols", &self.table_partition_cols)?; } - if self.collect_stat { - struct_ser.serialize_field("collectStat", &self.collect_stat)?; - } - if self.target_partitions != 0 { - struct_ser.serialize_field("targetPartitions", &self.target_partitions)?; - } if !self.file_sort_order.is_empty() { struct_ser.serialize_field("fileSortOrder", &self.file_sort_order)?; } @@ -11868,10 +12446,6 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { "filters", "table_partition_cols", "tablePartitionCols", - "collect_stat", - "collectStat", - "target_partitions", - "targetPartitions", "file_sort_order", "fileSortOrder", "csv", @@ -11890,8 +12464,6 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { Schema, Filters, TablePartitionCols, - CollectStat, - TargetPartitions, FileSortOrder, Csv, Parquet, @@ -11926,8 +12498,6 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { "schema" => Ok(GeneratedField::Schema), "filters" => Ok(GeneratedField::Filters), "tablePartitionCols" | "table_partition_cols" => Ok(GeneratedField::TablePartitionCols), - "collectStat" | "collect_stat" => Ok(GeneratedField::CollectStat), - "targetPartitions" | "target_partitions" => Ok(GeneratedField::TargetPartitions), "fileSortOrder" | "file_sort_order" => Ok(GeneratedField::FileSortOrder), "csv" => Ok(GeneratedField::Csv), "parquet" => Ok(GeneratedField::Parquet), @@ -11960,8 +12530,6 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { let mut schema__ = None; let mut filters__ = None; let mut table_partition_cols__ = None; - let mut collect_stat__ = None; - let mut target_partitions__ = None; let mut file_sort_order__ = None; let mut file_format_type__ = None; while let Some(k) = map_.next_key()? { @@ -12008,20 +12576,6 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { } table_partition_cols__ = Some(map_.next_value()?); } - GeneratedField::CollectStat => { - if collect_stat__.is_some() { - return Err(serde::de::Error::duplicate_field("collectStat")); - } - collect_stat__ = Some(map_.next_value()?); - } - GeneratedField::TargetPartitions => { - if target_partitions__.is_some() { - return Err(serde::de::Error::duplicate_field("targetPartitions")); - } - target_partitions__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) - ; - } GeneratedField::FileSortOrder => { if file_sort_order__.is_some() { return Err(serde::de::Error::duplicate_field("fileSortOrder")); @@ -12073,8 +12627,6 @@ impl<'de> serde::Deserialize<'de> for ListingTableScanNode { schema: schema__, filters: filters__.unwrap_or_default(), table_partition_cols: table_partition_cols__.unwrap_or_default(), - collect_stat: collect_stat__.unwrap_or_default(), - target_partitions: target_partitions__.unwrap_or_default(), file_sort_order: file_sort_order__.unwrap_or_default(), file_format_type: file_format_type__, }) @@ -12097,6 +12649,9 @@ impl serde::Serialize for LocalLimitExecNode { if self.fetch != 0 { len += 1; } + if !self.required_ordering.is_empty() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.LocalLimitExecNode", len)?; if let Some(v) = self.input.as_ref() { struct_ser.serialize_field("input", v)?; @@ -12104,6 +12659,9 @@ impl serde::Serialize for LocalLimitExecNode { if self.fetch != 0 { struct_ser.serialize_field("fetch", &self.fetch)?; } + if !self.required_ordering.is_empty() { + struct_ser.serialize_field("requiredOrdering", &self.required_ordering)?; + } struct_ser.end() } } @@ -12116,12 +12674,15 @@ impl<'de> serde::Deserialize<'de> for LocalLimitExecNode { const FIELDS: &[&str] = &[ "input", "fetch", + "required_ordering", + "requiredOrdering", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Input, Fetch, + RequiredOrdering, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -12145,6 +12706,7 @@ impl<'de> serde::Deserialize<'de> for LocalLimitExecNode { match value { "input" => Ok(GeneratedField::Input), "fetch" => Ok(GeneratedField::Fetch), + "requiredOrdering" | "required_ordering" => Ok(GeneratedField::RequiredOrdering), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -12166,6 +12728,7 @@ impl<'de> serde::Deserialize<'de> for LocalLimitExecNode { { let mut input__ = None; let mut fetch__ = None; + let mut required_ordering__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Input => { @@ -12182,11 +12745,18 @@ impl<'de> serde::Deserialize<'de> for LocalLimitExecNode { Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) ; } + GeneratedField::RequiredOrdering => { + if required_ordering__.is_some() { + return Err(serde::de::Error::duplicate_field("requiredOrdering")); + } + required_ordering__ = Some(map_.next_value()?); + } } } Ok(LocalLimitExecNode { input: input__, fetch: fetch__.unwrap_or_default(), + required_ordering: required_ordering__.unwrap_or_default(), }) } } @@ -12394,6 +12964,15 @@ impl serde::Serialize for LogicalExprNode { logical_expr_node::ExprType::ScalarSubqueryExpr(v) => { struct_ser.serialize_field("scalarSubqueryExpr", v)?; } + logical_expr_node::ExprType::HigherOrderUdfExpr(v) => { + struct_ser.serialize_field("higherOrderUdfExpr", v)?; + } + logical_expr_node::ExprType::Lambda(v) => { + struct_ser.serialize_field("lambda", v)?; + } + logical_expr_node::ExprType::LambdaVariable(v) => { + struct_ser.serialize_field("lambdaVariable", v)?; + } } } struct_ser.end() @@ -12457,6 +13036,11 @@ impl<'de> serde::Deserialize<'de> for LogicalExprNode { "unnest", "scalar_subquery_expr", "scalarSubqueryExpr", + "higher_order_udf_expr", + "higherOrderUdfExpr", + "lambda", + "lambda_variable", + "lambdaVariable", ]; #[allow(clippy::enum_variant_names)] @@ -12493,6 +13077,9 @@ impl<'de> serde::Deserialize<'de> for LogicalExprNode { Placeholder, Unnest, ScalarSubqueryExpr, + HigherOrderUdfExpr, + Lambda, + LambdaVariable, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -12546,6 +13133,9 @@ impl<'de> serde::Deserialize<'de> for LogicalExprNode { "placeholder" => Ok(GeneratedField::Placeholder), "unnest" => Ok(GeneratedField::Unnest), "scalarSubqueryExpr" | "scalar_subquery_expr" => Ok(GeneratedField::ScalarSubqueryExpr), + "higherOrderUdfExpr" | "higher_order_udf_expr" => Ok(GeneratedField::HigherOrderUdfExpr), + "lambda" => Ok(GeneratedField::Lambda), + "lambdaVariable" | "lambda_variable" => Ok(GeneratedField::LambdaVariable), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -12790,6 +13380,27 @@ impl<'de> serde::Deserialize<'de> for LogicalExprNode { return Err(serde::de::Error::duplicate_field("scalarSubqueryExpr")); } expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_expr_node::ExprType::ScalarSubqueryExpr) +; + } + GeneratedField::HigherOrderUdfExpr => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("higherOrderUdfExpr")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_expr_node::ExprType::HigherOrderUdfExpr) +; + } + GeneratedField::Lambda => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("lambda")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_expr_node::ExprType::Lambda) +; + } + GeneratedField::LambdaVariable => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("lambdaVariable")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(logical_expr_node::ExprType::LambdaVariable) ; } } @@ -13903,7 +14514,7 @@ impl<'de> serde::Deserialize<'de> for MemoryScanExecNode { deserializer.deserialize_struct("datafusion.MemoryScanExecNode", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for NamedStructField { +impl serde::Serialize for MergeAssignment { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -13911,29 +14522,37 @@ impl serde::Serialize for NamedStructField { { use serde::ser::SerializeStruct; let mut len = 0; - if self.name.is_some() { + if !self.column.is_empty() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.NamedStructField", len)?; - if let Some(v) = self.name.as_ref() { - struct_ser.serialize_field("name", v)?; + if self.value.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.MergeAssignment", len)?; + if !self.column.is_empty() { + struct_ser.serialize_field("column", &self.column)?; + } + if let Some(v) = self.value.as_ref() { + struct_ser.serialize_field("value", v)?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for NamedStructField { +impl<'de> serde::Deserialize<'de> for MergeAssignment { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "name", + "column", + "value", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - Name, + Column, + Value, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -13955,7 +14574,8 @@ impl<'de> serde::Deserialize<'de> for NamedStructField { E: serde::de::Error, { match value { - "name" => Ok(GeneratedField::Name), + "column" => Ok(GeneratedField::Column), + "value" => Ok(GeneratedField::Value), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -13965,66 +14585,66 @@ impl<'de> serde::Deserialize<'de> for NamedStructField { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = NamedStructField; + type Value = MergeAssignment; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.NamedStructField") + formatter.write_str("struct datafusion.MergeAssignment") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut name__ = None; + let mut column__ = None; + let mut value__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::Name => { - if name__.is_some() { - return Err(serde::de::Error::duplicate_field("name")); + GeneratedField::Column => { + if column__.is_some() { + return Err(serde::de::Error::duplicate_field("column")); } - name__ = map_.next_value()?; + column__ = Some(map_.next_value()?); + } + GeneratedField::Value => { + if value__.is_some() { + return Err(serde::de::Error::duplicate_field("value")); + } + value__ = map_.next_value()?; } } } - Ok(NamedStructField { - name: name__, + Ok(MergeAssignment { + column: column__.unwrap_or_default(), + value: value__, }) } } - deserializer.deserialize_struct("datafusion.NamedStructField", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.MergeAssignment", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for NegativeNode { +impl serde::Serialize for MergeDeleteAction { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where S: serde::Serializer, { use serde::ser::SerializeStruct; - let mut len = 0; - if self.expr.is_some() { - len += 1; - } - let mut struct_ser = serializer.serialize_struct("datafusion.NegativeNode", len)?; - if let Some(v) = self.expr.as_ref() { - struct_ser.serialize_field("expr", v)?; - } + let len = 0; + let struct_ser = serializer.serialize_struct("datafusion.MergeDeleteAction", len)?; struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for NegativeNode { +impl<'de> serde::Deserialize<'de> for MergeDeleteAction { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "expr", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - Expr, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -14045,10 +14665,7 @@ impl<'de> serde::Deserialize<'de> for NegativeNode { where E: serde::de::Error, { - match value { - "expr" => Ok(GeneratedField::Expr), - _ => Err(serde::de::Error::unknown_field(value, FIELDS)), - } + Err(serde::de::Error::unknown_field(value, FIELDS)) } } deserializer.deserialize_identifier(GeneratedVisitor) @@ -14056,36 +14673,27 @@ impl<'de> serde::Deserialize<'de> for NegativeNode { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = NegativeNode; + type Value = MergeDeleteAction; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.NegativeNode") + formatter.write_str("struct datafusion.MergeDeleteAction") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut expr__ = None; - while let Some(k) = map_.next_key()? { - match k { - GeneratedField::Expr => { - if expr__.is_some() { - return Err(serde::de::Error::duplicate_field("expr")); - } - expr__ = map_.next_value()?; - } - } + while map_.next_key::()?.is_some() { + let _ = map_.next_value::()?; } - Ok(NegativeNode { - expr: expr__, + Ok(MergeDeleteAction { }) } } - deserializer.deserialize_struct("datafusion.NegativeNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.MergeDeleteAction", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for NestedLoopJoinExecNode { +impl serde::Serialize for MergeInsertAction { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -14093,64 +14701,37 @@ impl serde::Serialize for NestedLoopJoinExecNode { { use serde::ser::SerializeStruct; let mut len = 0; - if self.left.is_some() { - len += 1; - } - if self.right.is_some() { - len += 1; - } - if self.join_type != 0 { - len += 1; - } - if self.filter.is_some() { + if !self.columns.is_empty() { len += 1; } - if !self.projection.is_empty() { + if !self.values.is_empty() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.NestedLoopJoinExecNode", len)?; - if let Some(v) = self.left.as_ref() { - struct_ser.serialize_field("left", v)?; - } - if let Some(v) = self.right.as_ref() { - struct_ser.serialize_field("right", v)?; - } - if self.join_type != 0 { - let v = super::datafusion_common::JoinType::try_from(self.join_type) - .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_type)))?; - struct_ser.serialize_field("joinType", &v)?; - } - if let Some(v) = self.filter.as_ref() { - struct_ser.serialize_field("filter", v)?; + let mut struct_ser = serializer.serialize_struct("datafusion.MergeInsertAction", len)?; + if !self.columns.is_empty() { + struct_ser.serialize_field("columns", &self.columns)?; } - if !self.projection.is_empty() { - struct_ser.serialize_field("projection", &self.projection)?; + if !self.values.is_empty() { + struct_ser.serialize_field("values", &self.values)?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for NestedLoopJoinExecNode { +impl<'de> serde::Deserialize<'de> for MergeInsertAction { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "left", - "right", - "join_type", - "joinType", - "filter", - "projection", + "columns", + "values", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - Left, - Right, - JoinType, - Filter, - Projection, + Columns, + Values, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -14172,11 +14753,8 @@ impl<'de> serde::Deserialize<'de> for NestedLoopJoinExecNode { E: serde::de::Error, { match value { - "left" => Ok(GeneratedField::Left), - "right" => Ok(GeneratedField::Right), - "joinType" | "join_type" => Ok(GeneratedField::JoinType), - "filter" => Ok(GeneratedField::Filter), - "projection" => Ok(GeneratedField::Projection), + "columns" => Ok(GeneratedField::Columns), + "values" => Ok(GeneratedField::Values), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -14186,71 +14764,44 @@ impl<'de> serde::Deserialize<'de> for NestedLoopJoinExecNode { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = NestedLoopJoinExecNode; + type Value = MergeInsertAction; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.NestedLoopJoinExecNode") + formatter.write_str("struct datafusion.MergeInsertAction") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut left__ = None; - let mut right__ = None; - let mut join_type__ = None; - let mut filter__ = None; - let mut projection__ = None; + let mut columns__ = None; + let mut values__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::Left => { - if left__.is_some() { - return Err(serde::de::Error::duplicate_field("left")); - } - left__ = map_.next_value()?; - } - GeneratedField::Right => { - if right__.is_some() { - return Err(serde::de::Error::duplicate_field("right")); - } - right__ = map_.next_value()?; - } - GeneratedField::JoinType => { - if join_type__.is_some() { - return Err(serde::de::Error::duplicate_field("joinType")); - } - join_type__ = Some(map_.next_value::()? as i32); - } - GeneratedField::Filter => { - if filter__.is_some() { - return Err(serde::de::Error::duplicate_field("filter")); + GeneratedField::Columns => { + if columns__.is_some() { + return Err(serde::de::Error::duplicate_field("columns")); } - filter__ = map_.next_value()?; + columns__ = Some(map_.next_value()?); } - GeneratedField::Projection => { - if projection__.is_some() { - return Err(serde::de::Error::duplicate_field("projection")); + GeneratedField::Values => { + if values__.is_some() { + return Err(serde::de::Error::duplicate_field("values")); } - projection__ = - Some(map_.next_value::>>()? - .into_iter().map(|x| x.0).collect()) - ; + values__ = Some(map_.next_value()?); } } } - Ok(NestedLoopJoinExecNode { - left: left__, - right: right__, - join_type: join_type__.unwrap_or_default(), - filter: filter__, - projection: projection__.unwrap_or_default(), + Ok(MergeInsertAction { + columns: columns__.unwrap_or_default(), + values: values__.unwrap_or_default(), }) } } - deserializer.deserialize_struct("datafusion.NestedLoopJoinExecNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.MergeInsertAction", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for Not { +impl serde::Serialize for MergeIntoActionNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -14258,29 +14809,43 @@ impl serde::Serialize for Not { { use serde::ser::SerializeStruct; let mut len = 0; - if self.expr.is_some() { + if self.action.is_some() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.Not", len)?; - if let Some(v) = self.expr.as_ref() { - struct_ser.serialize_field("expr", v)?; + let mut struct_ser = serializer.serialize_struct("datafusion.MergeIntoActionNode", len)?; + if let Some(v) = self.action.as_ref() { + match v { + merge_into_action_node::Action::Update(v) => { + struct_ser.serialize_field("update", v)?; + } + merge_into_action_node::Action::Insert(v) => { + struct_ser.serialize_field("insert", v)?; + } + merge_into_action_node::Action::Delete(v) => { + struct_ser.serialize_field("delete", v)?; + } + } } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for Not { +impl<'de> serde::Deserialize<'de> for MergeIntoActionNode { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "expr", + "update", + "insert", + "delete", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - Expr, + Update, + Insert, + Delete, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -14302,7 +14867,9 @@ impl<'de> serde::Deserialize<'de> for Not { E: serde::de::Error, { match value { - "expr" => Ok(GeneratedField::Expr), + "update" => Ok(GeneratedField::Update), + "insert" => Ok(GeneratedField::Insert), + "delete" => Ok(GeneratedField::Delete), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -14312,57 +14879,913 @@ impl<'de> serde::Deserialize<'de> for Not { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = Not; + type Value = MergeIntoActionNode; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.Not") + formatter.write_str("struct datafusion.MergeIntoActionNode") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut expr__ = None; + let mut action__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::Expr => { - if expr__.is_some() { - return Err(serde::de::Error::duplicate_field("expr")); + GeneratedField::Update => { + if action__.is_some() { + return Err(serde::de::Error::duplicate_field("update")); } - expr__ = map_.next_value()?; + action__ = map_.next_value::<::std::option::Option<_>>()?.map(merge_into_action_node::Action::Update) +; + } + GeneratedField::Insert => { + if action__.is_some() { + return Err(serde::de::Error::duplicate_field("insert")); + } + action__ = map_.next_value::<::std::option::Option<_>>()?.map(merge_into_action_node::Action::Insert) +; + } + GeneratedField::Delete => { + if action__.is_some() { + return Err(serde::de::Error::duplicate_field("delete")); + } + action__ = map_.next_value::<::std::option::Option<_>>()?.map(merge_into_action_node::Action::Delete) +; } } } - Ok(Not { - expr: expr__, + Ok(MergeIntoActionNode { + action: action__, }) } } - deserializer.deserialize_struct("datafusion.Not", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.MergeIntoActionNode", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for NullTreatment { +impl serde::Serialize for MergeIntoClauseNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where S: serde::Serializer, { - let variant = match self { - Self::RespectNulls => "RESPECT_NULLS", - Self::IgnoreNulls => "IGNORE_NULLS", - }; - serializer.serialize_str(variant) + use serde::ser::SerializeStruct; + let mut len = 0; + if self.kind != 0 { + len += 1; + } + if self.predicate.is_some() { + len += 1; + } + if self.action.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.MergeIntoClauseNode", len)?; + if self.kind != 0 { + let v = merge_into_clause_node::Kind::try_from(self.kind) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.kind)))?; + struct_ser.serialize_field("kind", &v)?; + } + if let Some(v) = self.predicate.as_ref() { + struct_ser.serialize_field("predicate", v)?; + } + if let Some(v) = self.action.as_ref() { + struct_ser.serialize_field("action", v)?; + } + struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for NullTreatment { +impl<'de> serde::Deserialize<'de> for MergeIntoClauseNode { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "RESPECT_NULLS", - "IGNORE_NULLS", + "kind", + "predicate", + "action", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Kind, + Predicate, + Action, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "kind" => Ok(GeneratedField::Kind), + "predicate" => Ok(GeneratedField::Predicate), + "action" => Ok(GeneratedField::Action), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = MergeIntoClauseNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.MergeIntoClauseNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut kind__ = None; + let mut predicate__ = None; + let mut action__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Kind => { + if kind__.is_some() { + return Err(serde::de::Error::duplicate_field("kind")); + } + kind__ = Some(map_.next_value::()? as i32); + } + GeneratedField::Predicate => { + if predicate__.is_some() { + return Err(serde::de::Error::duplicate_field("predicate")); + } + predicate__ = map_.next_value()?; + } + GeneratedField::Action => { + if action__.is_some() { + return Err(serde::de::Error::duplicate_field("action")); + } + action__ = map_.next_value()?; + } + } + } + Ok(MergeIntoClauseNode { + kind: kind__.unwrap_or_default(), + predicate: predicate__, + action: action__, + }) + } + } + deserializer.deserialize_struct("datafusion.MergeIntoClauseNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for merge_into_clause_node::Kind { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Matched => "MATCHED", + Self::NotMatched => "NOT_MATCHED", + Self::NotMatchedByTarget => "NOT_MATCHED_BY_TARGET", + Self::NotMatchedBySource => "NOT_MATCHED_BY_SOURCE", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for merge_into_clause_node::Kind { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "MATCHED", + "NOT_MATCHED", + "NOT_MATCHED_BY_TARGET", + "NOT_MATCHED_BY_SOURCE", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = merge_into_clause_node::Kind; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "MATCHED" => Ok(merge_into_clause_node::Kind::Matched), + "NOT_MATCHED" => Ok(merge_into_clause_node::Kind::NotMatched), + "NOT_MATCHED_BY_TARGET" => Ok(merge_into_clause_node::Kind::NotMatchedByTarget), + "NOT_MATCHED_BY_SOURCE" => Ok(merge_into_clause_node::Kind::NotMatchedBySource), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} +impl serde::Serialize for MergeIntoOpNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.on.is_some() { + len += 1; + } + if !self.clauses.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.MergeIntoOpNode", len)?; + if let Some(v) = self.on.as_ref() { + struct_ser.serialize_field("on", v)?; + } + if !self.clauses.is_empty() { + struct_ser.serialize_field("clauses", &self.clauses)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for MergeIntoOpNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "on", + "clauses", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + On, + Clauses, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "on" => Ok(GeneratedField::On), + "clauses" => Ok(GeneratedField::Clauses), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = MergeIntoOpNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.MergeIntoOpNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut on__ = None; + let mut clauses__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::On => { + if on__.is_some() { + return Err(serde::de::Error::duplicate_field("on")); + } + on__ = map_.next_value()?; + } + GeneratedField::Clauses => { + if clauses__.is_some() { + return Err(serde::de::Error::duplicate_field("clauses")); + } + clauses__ = Some(map_.next_value()?); + } + } + } + Ok(MergeIntoOpNode { + on: on__, + clauses: clauses__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.MergeIntoOpNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for MergeUpdateAction { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.assignments.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.MergeUpdateAction", len)?; + if !self.assignments.is_empty() { + struct_ser.serialize_field("assignments", &self.assignments)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for MergeUpdateAction { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "assignments", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Assignments, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "assignments" => Ok(GeneratedField::Assignments), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = MergeUpdateAction; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.MergeUpdateAction") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut assignments__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Assignments => { + if assignments__.is_some() { + return Err(serde::de::Error::duplicate_field("assignments")); + } + assignments__ = Some(map_.next_value()?); + } + } + } + Ok(MergeUpdateAction { + assignments: assignments__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.MergeUpdateAction", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for NamedStructField { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.name.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.NamedStructField", len)?; + if let Some(v) = self.name.as_ref() { + struct_ser.serialize_field("name", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for NamedStructField { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "name", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Name, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "name" => Ok(GeneratedField::Name), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = NamedStructField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.NamedStructField") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut name__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Name => { + if name__.is_some() { + return Err(serde::de::Error::duplicate_field("name")); + } + name__ = map_.next_value()?; + } + } + } + Ok(NamedStructField { + name: name__, + }) + } + } + deserializer.deserialize_struct("datafusion.NamedStructField", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for NegativeNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.expr.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.NegativeNode", len)?; + if let Some(v) = self.expr.as_ref() { + struct_ser.serialize_field("expr", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for NegativeNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "expr", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Expr, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "expr" => Ok(GeneratedField::Expr), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = NegativeNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.NegativeNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut expr__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Expr => { + if expr__.is_some() { + return Err(serde::de::Error::duplicate_field("expr")); + } + expr__ = map_.next_value()?; + } + } + } + Ok(NegativeNode { + expr: expr__, + }) + } + } + deserializer.deserialize_struct("datafusion.NegativeNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for NestedLoopJoinExecNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.left.is_some() { + len += 1; + } + if self.right.is_some() { + len += 1; + } + if self.join_type != 0 { + len += 1; + } + if self.filter.is_some() { + len += 1; + } + if !self.projection.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.NestedLoopJoinExecNode", len)?; + if let Some(v) = self.left.as_ref() { + struct_ser.serialize_field("left", v)?; + } + if let Some(v) = self.right.as_ref() { + struct_ser.serialize_field("right", v)?; + } + if self.join_type != 0 { + let v = super::datafusion_common::JoinType::try_from(self.join_type) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.join_type)))?; + struct_ser.serialize_field("joinType", &v)?; + } + if let Some(v) = self.filter.as_ref() { + struct_ser.serialize_field("filter", v)?; + } + if !self.projection.is_empty() { + struct_ser.serialize_field("projection", &self.projection)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for NestedLoopJoinExecNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "left", + "right", + "join_type", + "joinType", + "filter", + "projection", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Left, + Right, + JoinType, + Filter, + Projection, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "left" => Ok(GeneratedField::Left), + "right" => Ok(GeneratedField::Right), + "joinType" | "join_type" => Ok(GeneratedField::JoinType), + "filter" => Ok(GeneratedField::Filter), + "projection" => Ok(GeneratedField::Projection), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = NestedLoopJoinExecNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.NestedLoopJoinExecNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut left__ = None; + let mut right__ = None; + let mut join_type__ = None; + let mut filter__ = None; + let mut projection__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Left => { + if left__.is_some() { + return Err(serde::de::Error::duplicate_field("left")); + } + left__ = map_.next_value()?; + } + GeneratedField::Right => { + if right__.is_some() { + return Err(serde::de::Error::duplicate_field("right")); + } + right__ = map_.next_value()?; + } + GeneratedField::JoinType => { + if join_type__.is_some() { + return Err(serde::de::Error::duplicate_field("joinType")); + } + join_type__ = Some(map_.next_value::()? as i32); + } + GeneratedField::Filter => { + if filter__.is_some() { + return Err(serde::de::Error::duplicate_field("filter")); + } + filter__ = map_.next_value()?; + } + GeneratedField::Projection => { + if projection__.is_some() { + return Err(serde::de::Error::duplicate_field("projection")); + } + projection__ = + Some(map_.next_value::>>()? + .into_iter().map(|x| x.0).collect()) + ; + } + } + } + Ok(NestedLoopJoinExecNode { + left: left__, + right: right__, + join_type: join_type__.unwrap_or_default(), + filter: filter__, + projection: projection__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.NestedLoopJoinExecNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for Not { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.expr.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.Not", len)?; + if let Some(v) = self.expr.as_ref() { + struct_ser.serialize_field("expr", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for Not { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "expr", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Expr, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "expr" => Ok(GeneratedField::Expr), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = Not; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.Not") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut expr__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Expr => { + if expr__.is_some() { + return Err(serde::de::Error::duplicate_field("expr")); + } + expr__ = map_.next_value()?; + } + } + } + Ok(Not { + expr: expr__, + }) + } + } + deserializer.deserialize_struct("datafusion.Not", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for NullTreatment { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::RespectNulls => "RESPECT_NULLS", + Self::IgnoreNulls => "IGNORE_NULLS", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for NullTreatment { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "RESPECT_NULLS", + "IGNORE_NULLS", ]; struct GeneratedVisitor; @@ -15545,6 +16968,9 @@ impl serde::Serialize for PartitionedFile { if self.statistics.is_some() { len += 1; } + if self.arrow_schema.is_some() { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PartitionedFile", len)?; if !self.path.is_empty() { struct_ser.serialize_field("path", &self.path)?; @@ -15568,6 +16994,9 @@ impl serde::Serialize for PartitionedFile { if let Some(v) = self.statistics.as_ref() { struct_ser.serialize_field("statistics", v)?; } + if let Some(v) = self.arrow_schema.as_ref() { + struct_ser.serialize_field("arrowSchema", v)?; + } struct_ser.end() } } @@ -15586,6 +17015,8 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { "partitionValues", "range", "statistics", + "arrow_schema", + "arrowSchema", ]; #[allow(clippy::enum_variant_names)] @@ -15596,6 +17027,7 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { PartitionValues, Range, Statistics, + ArrowSchema, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -15623,6 +17055,7 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { "partitionValues" | "partition_values" => Ok(GeneratedField::PartitionValues), "range" => Ok(GeneratedField::Range), "statistics" => Ok(GeneratedField::Statistics), + "arrowSchema" | "arrow_schema" => Ok(GeneratedField::ArrowSchema), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -15648,6 +17081,7 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { let mut partition_values__ = None; let mut range__ = None; let mut statistics__ = None; + let mut arrow_schema__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Path => { @@ -15690,6 +17124,12 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { } statistics__ = map_.next_value()?; } + GeneratedField::ArrowSchema => { + if arrow_schema__.is_some() { + return Err(serde::de::Error::duplicate_field("arrowSchema")); + } + arrow_schema__ = map_.next_value()?; + } } } Ok(PartitionedFile { @@ -15699,6 +17139,7 @@ impl<'de> serde::Deserialize<'de> for PartitionedFile { partition_values: partition_values__.unwrap_or_default(), range: range__, statistics: statistics__, + arrow_schema: arrow_schema__, }) } } @@ -15732,6 +17173,9 @@ impl serde::Serialize for Partitioning { #[allow(clippy::needless_borrows_for_generic_args)] struct_ser.serialize_field("unknown", ToString::to_string(&v).as_str())?; } + partitioning::PartitionMethod::Range(v) => { + struct_ser.serialize_field("range", v)?; + } } } struct_ser.end() @@ -15748,6 +17192,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { "roundRobin", "hash", "unknown", + "range", ]; #[allow(clippy::enum_variant_names)] @@ -15755,6 +17200,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { RoundRobin, Hash, Unknown, + Range, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -15779,6 +17225,7 @@ impl<'de> serde::Deserialize<'de> for Partitioning { "roundRobin" | "round_robin" => Ok(GeneratedField::RoundRobin), "hash" => Ok(GeneratedField::Hash), "unknown" => Ok(GeneratedField::Unknown), + "range" => Ok(GeneratedField::Range), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -15820,6 +17267,13 @@ impl<'de> serde::Deserialize<'de> for Partitioning { } partition_method__ = map_.next_value::<::std::option::Option<::pbjson::private::NumberDeserialize<_>>>()?.map(|x| partitioning::PartitionMethod::Unknown(x.0)); } + GeneratedField::Range => { + if partition_method__.is_some() { + return Err(serde::de::Error::duplicate_field("range")); + } + partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(partitioning::PartitionMethod::Range) +; + } } } Ok(Partitioning { @@ -15856,6 +17310,9 @@ impl serde::Serialize for PhysicalAggregateExprNode { if !self.human_display.is_empty() { len += 1; } + if self.is_reversed { + len += 1; + } if self.aggregate_function.is_some() { len += 1; } @@ -15880,6 +17337,9 @@ impl serde::Serialize for PhysicalAggregateExprNode { if !self.human_display.is_empty() { struct_ser.serialize_field("humanDisplay", &self.human_display)?; } + if self.is_reversed { + struct_ser.serialize_field("isReversed", &self.is_reversed)?; + } if let Some(v) = self.aggregate_function.as_ref() { match v { physical_aggregate_expr_node::AggregateFunction::UserDefinedAggrFunction(v) => { @@ -15907,6 +17367,8 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode { "funDefinition", "human_display", "humanDisplay", + "is_reversed", + "isReversed", "user_defined_aggr_function", "userDefinedAggrFunction", ]; @@ -15919,6 +17381,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode { IgnoreNulls, FunDefinition, HumanDisplay, + IsReversed, UserDefinedAggrFunction, } impl<'de> serde::Deserialize<'de> for GeneratedField { @@ -15947,6 +17410,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode { "ignoreNulls" | "ignore_nulls" => Ok(GeneratedField::IgnoreNulls), "funDefinition" | "fun_definition" => Ok(GeneratedField::FunDefinition), "humanDisplay" | "human_display" => Ok(GeneratedField::HumanDisplay), + "isReversed" | "is_reversed" => Ok(GeneratedField::IsReversed), "userDefinedAggrFunction" | "user_defined_aggr_function" => Ok(GeneratedField::UserDefinedAggrFunction), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } @@ -15973,6 +17437,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode { let mut ignore_nulls__ = None; let mut fun_definition__ = None; let mut human_display__ = None; + let mut is_reversed__ = None; let mut aggregate_function__ = None; while let Some(k) = map_.next_key()? { match k { @@ -16014,6 +17479,12 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode { } human_display__ = Some(map_.next_value()?); } + GeneratedField::IsReversed => { + if is_reversed__.is_some() { + return Err(serde::de::Error::duplicate_field("isReversed")); + } + is_reversed__ = Some(map_.next_value()?); + } GeneratedField::UserDefinedAggrFunction => { if aggregate_function__.is_some() { return Err(serde::de::Error::duplicate_field("userDefinedAggrFunction")); @@ -16029,6 +17500,7 @@ impl<'de> serde::Deserialize<'de> for PhysicalAggregateExprNode { ignore_nulls: ignore_nulls__.unwrap_or_default(), fun_definition: fun_definition__, human_display: human_display__.unwrap_or_default(), + is_reversed: is_reversed__.unwrap_or_default(), aggregate_function: aggregate_function__, }) } @@ -17008,6 +18480,18 @@ impl serde::Serialize for PhysicalExprNode { physical_expr_node::ExprType::DynamicFilter(v) => { struct_ser.serialize_field("dynamicFilter", v)?; } + physical_expr_node::ExprType::HigherOrderUdf(v) => { + struct_ser.serialize_field("higherOrderUdf", v)?; + } + physical_expr_node::ExprType::Lambda(v) => { + struct_ser.serialize_field("lambda", v)?; + } + physical_expr_node::ExprType::LambdaVariable(v) => { + struct_ser.serialize_field("lambdaVariable", v)?; + } + physical_expr_node::ExprType::RangeExpr(v) => { + struct_ser.serialize_field("rangeExpr", v)?; + } } } struct_ser.end() @@ -17058,6 +18542,13 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "scalarSubquery", "dynamic_filter", "dynamicFilter", + "higher_order_udf", + "higherOrderUdf", + "lambda", + "lambda_variable", + "lambdaVariable", + "range_expr", + "rangeExpr", ]; #[allow(clippy::enum_variant_names)] @@ -17084,6 +18575,10 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { HashExpr, ScalarSubquery, DynamicFilter, + HigherOrderUdf, + Lambda, + LambdaVariable, + RangeExpr, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -17127,6 +18622,10 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { "hashExpr" | "hash_expr" => Ok(GeneratedField::HashExpr), "scalarSubquery" | "scalar_subquery" => Ok(GeneratedField::ScalarSubquery), "dynamicFilter" | "dynamic_filter" => Ok(GeneratedField::DynamicFilter), + "higherOrderUdf" | "higher_order_udf" => Ok(GeneratedField::HigherOrderUdf), + "lambda" => Ok(GeneratedField::Lambda), + "lambdaVariable" | "lambda_variable" => Ok(GeneratedField::LambdaVariable), + "rangeExpr" | "range_expr" => Ok(GeneratedField::RangeExpr), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -17303,6 +18802,34 @@ impl<'de> serde::Deserialize<'de> for PhysicalExprNode { return Err(serde::de::Error::duplicate_field("dynamicFilter")); } expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::DynamicFilter) +; + } + GeneratedField::HigherOrderUdf => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("higherOrderUdf")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::HigherOrderUdf) +; + } + GeneratedField::Lambda => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("lambda")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::Lambda) +; + } + GeneratedField::LambdaVariable => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("lambdaVariable")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::LambdaVariable) +; + } + GeneratedField::RangeExpr => { + if expr_type__.is_some() { + return Err(serde::de::Error::duplicate_field("rangeExpr")); + } + expr_type__ = map_.next_value::<::std::option::Option<_>>()?.map(physical_expr_node::ExprType::RangeExpr) ; } } @@ -17660,17 +19187,131 @@ impl<'de> serde::Deserialize<'de> for PhysicalHashExprNode { } } } - Ok(PhysicalHashExprNode { - on_columns: on_columns__.unwrap_or_default(), - seed0: seed0__.unwrap_or_default(), - description: description__.unwrap_or_default(), + Ok(PhysicalHashExprNode { + on_columns: on_columns__.unwrap_or_default(), + seed0: seed0__.unwrap_or_default(), + description: description__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalHashExprNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalHashRepartition { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.hash_expr.is_empty() { + len += 1; + } + if self.partition_count != 0 { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalHashRepartition", len)?; + if !self.hash_expr.is_empty() { + struct_ser.serialize_field("hashExpr", &self.hash_expr)?; + } + if self.partition_count != 0 { + #[allow(clippy::needless_borrow)] + #[allow(clippy::needless_borrows_for_generic_args)] + struct_ser.serialize_field("partitionCount", ToString::to_string(&self.partition_count).as_str())?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "hash_expr", + "hashExpr", + "partition_count", + "partitionCount", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + HashExpr, + PartitionCount, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "hashExpr" | "hash_expr" => Ok(GeneratedField::HashExpr), + "partitionCount" | "partition_count" => Ok(GeneratedField::PartitionCount), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalHashRepartition; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalHashRepartition") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut hash_expr__ = None; + let mut partition_count__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::HashExpr => { + if hash_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("hashExpr")); + } + hash_expr__ = Some(map_.next_value()?); + } + GeneratedField::PartitionCount => { + if partition_count__.is_some() { + return Err(serde::de::Error::duplicate_field("partitionCount")); + } + partition_count__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + } + } + Ok(PhysicalHashRepartition { + hash_expr: hash_expr__.unwrap_or_default(), + partition_count: partition_count__.unwrap_or_default(), }) } } - deserializer.deserialize_struct("datafusion.PhysicalHashExprNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.PhysicalHashRepartition", FIELDS, GeneratedVisitor) } } -impl serde::Serialize for PhysicalHashRepartition { +impl serde::Serialize for PhysicalHigherOrderUdfNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result where @@ -17678,41 +19319,48 @@ impl serde::Serialize for PhysicalHashRepartition { { use serde::ser::SerializeStruct; let mut len = 0; - if !self.hash_expr.is_empty() { + if !self.name.is_empty() { len += 1; } - if self.partition_count != 0 { + if !self.args.is_empty() { len += 1; } - let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalHashRepartition", len)?; - if !self.hash_expr.is_empty() { - struct_ser.serialize_field("hashExpr", &self.hash_expr)?; + if self.fun_definition.is_some() { + len += 1; } - if self.partition_count != 0 { + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalHigherOrderUdfNode", len)?; + if !self.name.is_empty() { + struct_ser.serialize_field("name", &self.name)?; + } + if !self.args.is_empty() { + struct_ser.serialize_field("args", &self.args)?; + } + if let Some(v) = self.fun_definition.as_ref() { #[allow(clippy::needless_borrow)] #[allow(clippy::needless_borrows_for_generic_args)] - struct_ser.serialize_field("partitionCount", ToString::to_string(&self.partition_count).as_str())?; + struct_ser.serialize_field("funDefinition", pbjson::private::base64::encode(&v).as_str())?; } struct_ser.end() } } -impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { +impl<'de> serde::Deserialize<'de> for PhysicalHigherOrderUdfNode { #[allow(deprecated)] fn deserialize(deserializer: D) -> std::result::Result where D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "hash_expr", - "hashExpr", - "partition_count", - "partitionCount", + "name", + "args", + "fun_definition", + "funDefinition", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - HashExpr, - PartitionCount, + Name, + Args, + FunDefinition, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -17734,8 +19382,9 @@ impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { E: serde::de::Error, { match value { - "hashExpr" | "hash_expr" => Ok(GeneratedField::HashExpr), - "partitionCount" | "partition_count" => Ok(GeneratedField::PartitionCount), + "name" => Ok(GeneratedField::Name), + "args" => Ok(GeneratedField::Args), + "funDefinition" | "fun_definition" => Ok(GeneratedField::FunDefinition), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -17745,43 +19394,51 @@ impl<'de> serde::Deserialize<'de> for PhysicalHashRepartition { } struct GeneratedVisitor; impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { - type Value = PhysicalHashRepartition; + type Value = PhysicalHigherOrderUdfNode; fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - formatter.write_str("struct datafusion.PhysicalHashRepartition") + formatter.write_str("struct datafusion.PhysicalHigherOrderUdfNode") } - fn visit_map(self, mut map_: V) -> std::result::Result + fn visit_map(self, mut map_: V) -> std::result::Result where V: serde::de::MapAccess<'de>, { - let mut hash_expr__ = None; - let mut partition_count__ = None; + let mut name__ = None; + let mut args__ = None; + let mut fun_definition__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::HashExpr => { - if hash_expr__.is_some() { - return Err(serde::de::Error::duplicate_field("hashExpr")); + GeneratedField::Name => { + if name__.is_some() { + return Err(serde::de::Error::duplicate_field("name")); } - hash_expr__ = Some(map_.next_value()?); + name__ = Some(map_.next_value()?); } - GeneratedField::PartitionCount => { - if partition_count__.is_some() { - return Err(serde::de::Error::duplicate_field("partitionCount")); + GeneratedField::Args => { + if args__.is_some() { + return Err(serde::de::Error::duplicate_field("args")); } - partition_count__ = - Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + args__ = Some(map_.next_value()?); + } + GeneratedField::FunDefinition => { + if fun_definition__.is_some() { + return Err(serde::de::Error::duplicate_field("funDefinition")); + } + fun_definition__ = + map_.next_value::<::std::option::Option<::pbjson::private::BytesDeserialize<_>>>()?.map(|x| x.0) ; } } } - Ok(PhysicalHashRepartition { - hash_expr: hash_expr__.unwrap_or_default(), - partition_count: partition_count__.unwrap_or_default(), + Ok(PhysicalHigherOrderUdfNode { + name: name__.unwrap_or_default(), + args: args__.unwrap_or_default(), + fun_definition: fun_definition__, }) } } - deserializer.deserialize_struct("datafusion.PhysicalHashRepartition", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.PhysicalHigherOrderUdfNode", FIELDS, GeneratedVisitor) } } impl serde::Serialize for PhysicalInListNode { @@ -18083,12 +19740,230 @@ impl<'de> serde::Deserialize<'de> for PhysicalIsNull { } } } - Ok(PhysicalIsNull { - expr: expr__, + Ok(PhysicalIsNull { + expr: expr__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalIsNull", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalLambdaExprNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.params.is_empty() { + len += 1; + } + if self.body.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalLambdaExprNode", len)?; + if !self.params.is_empty() { + struct_ser.serialize_field("params", &self.params)?; + } + if let Some(v) = self.body.as_ref() { + struct_ser.serialize_field("body", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalLambdaExprNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "params", + "body", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Params, + Body, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "params" => Ok(GeneratedField::Params), + "body" => Ok(GeneratedField::Body), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalLambdaExprNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalLambdaExprNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut params__ = None; + let mut body__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Params => { + if params__.is_some() { + return Err(serde::de::Error::duplicate_field("params")); + } + params__ = Some(map_.next_value()?); + } + GeneratedField::Body => { + if body__.is_some() { + return Err(serde::de::Error::duplicate_field("body")); + } + body__ = map_.next_value()?; + } + } + } + Ok(PhysicalLambdaExprNode { + params: params__.unwrap_or_default(), + body: body__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalLambdaExprNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalLambdaVariableExprNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if self.index != 0 { + len += 1; + } + if self.field.is_some() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalLambdaVariableExprNode", len)?; + if self.index != 0 { + struct_ser.serialize_field("index", &self.index)?; + } + if let Some(v) = self.field.as_ref() { + struct_ser.serialize_field("field", v)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalLambdaVariableExprNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "index", + "field", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Index, + Field, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "index" => Ok(GeneratedField::Index), + "field" => Ok(GeneratedField::Field), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalLambdaVariableExprNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalLambdaVariableExprNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut index__ = None; + let mut field__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Index => { + if index__.is_some() { + return Err(serde::de::Error::duplicate_field("index")); + } + index__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } + GeneratedField::Field => { + if field__.is_some() { + return Err(serde::de::Error::duplicate_field("field")); + } + field__ = map_.next_value()?; + } + } + } + Ok(PhysicalLambdaVariableExprNode { + index: index__.unwrap_or_default(), + field: field__, }) } } - deserializer.deserialize_struct("datafusion.PhysicalIsNull", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.PhysicalLambdaVariableExprNode", FIELDS, GeneratedVisitor) } } impl serde::Serialize for PhysicalLikeExprNode { @@ -19007,12 +20882,323 @@ impl<'de> serde::Deserialize<'de> for PhysicalPlanNode { } } } - Ok(PhysicalPlanNode { - physical_plan_type: physical_plan_type__, + Ok(PhysicalPlanNode { + physical_plan_type: physical_plan_type__, + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalPlanNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalRangeExprNode { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.sort_expr.is_empty() { + len += 1; + } + if !self.split_point.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalRangeExprNode", len)?; + if !self.sort_expr.is_empty() { + struct_ser.serialize_field("sortExpr", &self.sort_expr)?; + } + if !self.split_point.is_empty() { + struct_ser.serialize_field("splitPoint", &self.split_point)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalRangeExprNode { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "sort_expr", + "sortExpr", + "split_point", + "splitPoint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + SortExpr, + SplitPoint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "sortExpr" | "sort_expr" => Ok(GeneratedField::SortExpr), + "splitPoint" | "split_point" => Ok(GeneratedField::SplitPoint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalRangeExprNode; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalRangeExprNode") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut sort_expr__ = None; + let mut split_point__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::SortExpr => { + if sort_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("sortExpr")); + } + sort_expr__ = Some(map_.next_value()?); + } + GeneratedField::SplitPoint => { + if split_point__.is_some() { + return Err(serde::de::Error::duplicate_field("splitPoint")); + } + split_point__ = Some(map_.next_value()?); + } + } + } + Ok(PhysicalRangeExprNode { + sort_expr: sort_expr__.unwrap_or_default(), + split_point: split_point__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalRangeExprNode", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalRangePartitioning { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.sort_expr.is_empty() { + len += 1; + } + if !self.split_point.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalRangePartitioning", len)?; + if !self.sort_expr.is_empty() { + struct_ser.serialize_field("sortExpr", &self.sort_expr)?; + } + if !self.split_point.is_empty() { + struct_ser.serialize_field("splitPoint", &self.split_point)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalRangePartitioning { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "sort_expr", + "sortExpr", + "split_point", + "splitPoint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + SortExpr, + SplitPoint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "sortExpr" | "sort_expr" => Ok(GeneratedField::SortExpr), + "splitPoint" | "split_point" => Ok(GeneratedField::SplitPoint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalRangePartitioning; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalRangePartitioning") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut sort_expr__ = None; + let mut split_point__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::SortExpr => { + if sort_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("sortExpr")); + } + sort_expr__ = Some(map_.next_value()?); + } + GeneratedField::SplitPoint => { + if split_point__.is_some() { + return Err(serde::de::Error::duplicate_field("splitPoint")); + } + split_point__ = Some(map_.next_value()?); + } + } + } + Ok(PhysicalRangePartitioning { + sort_expr: sort_expr__.unwrap_or_default(), + split_point: split_point__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.PhysicalRangePartitioning", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for PhysicalRangeSplitPoint { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.value.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.PhysicalRangeSplitPoint", len)?; + if !self.value.is_empty() { + struct_ser.serialize_field("value", &self.value)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for PhysicalRangeSplitPoint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "value", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Value, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "value" => Ok(GeneratedField::Value), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = PhysicalRangeSplitPoint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.PhysicalRangeSplitPoint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut value__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Value => { + if value__.is_some() { + return Err(serde::de::Error::duplicate_field("value")); + } + value__ = Some(map_.next_value()?); + } + } + } + Ok(PhysicalRangeSplitPoint { + value: value__.unwrap_or_default(), }) } } - deserializer.deserialize_struct("datafusion.PhysicalPlanNode", FIELDS, GeneratedVisitor) + deserializer.deserialize_struct("datafusion.PhysicalRangeSplitPoint", FIELDS, GeneratedVisitor) } } impl serde::Serialize for PhysicalScalarSubqueryExprNode { @@ -20173,10 +22359,16 @@ impl serde::Serialize for PlaceholderRowExecNode { if self.schema.is_some() { len += 1; } + if self.partitions != 0 { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.PlaceholderRowExecNode", len)?; if let Some(v) = self.schema.as_ref() { struct_ser.serialize_field("schema", v)?; } + if self.partitions != 0 { + struct_ser.serialize_field("partitions", &self.partitions)?; + } struct_ser.end() } } @@ -20188,11 +22380,13 @@ impl<'de> serde::Deserialize<'de> for PlaceholderRowExecNode { { const FIELDS: &[&str] = &[ "schema", + "partitions", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Schema, + Partitions, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -20215,6 +22409,7 @@ impl<'de> serde::Deserialize<'de> for PlaceholderRowExecNode { { match value { "schema" => Ok(GeneratedField::Schema), + "partitions" => Ok(GeneratedField::Partitions), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -20235,6 +22430,7 @@ impl<'de> serde::Deserialize<'de> for PlaceholderRowExecNode { V: serde::de::MapAccess<'de>, { let mut schema__ = None; + let mut partitions__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Schema => { @@ -20243,10 +22439,19 @@ impl<'de> serde::Deserialize<'de> for PlaceholderRowExecNode { } schema__ = map_.next_value()?; } + GeneratedField::Partitions => { + if partitions__.is_some() { + return Err(serde::de::Error::duplicate_field("partitions")); + } + partitions__ = + Some(map_.next_value::<::pbjson::private::NumberDeserialize<_>>()?.0) + ; + } } } Ok(PlaceholderRowExecNode { schema: schema__, + partitions: partitions__.unwrap_or_default(), }) } } @@ -21193,6 +23398,207 @@ impl<'de> serde::Deserialize<'de> for ProjectionNode { deserializer.deserialize_struct("datafusion.ProjectionNode", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for RangeRepartition { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.sort_expr.is_empty() { + len += 1; + } + if !self.split_point.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.RangeRepartition", len)?; + if !self.sort_expr.is_empty() { + struct_ser.serialize_field("sortExpr", &self.sort_expr)?; + } + if !self.split_point.is_empty() { + struct_ser.serialize_field("splitPoint", &self.split_point)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for RangeRepartition { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "sort_expr", + "sortExpr", + "split_point", + "splitPoint", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + SortExpr, + SplitPoint, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "sortExpr" | "sort_expr" => Ok(GeneratedField::SortExpr), + "splitPoint" | "split_point" => Ok(GeneratedField::SplitPoint), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = RangeRepartition; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.RangeRepartition") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut sort_expr__ = None; + let mut split_point__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::SortExpr => { + if sort_expr__.is_some() { + return Err(serde::de::Error::duplicate_field("sortExpr")); + } + sort_expr__ = Some(map_.next_value()?); + } + GeneratedField::SplitPoint => { + if split_point__.is_some() { + return Err(serde::de::Error::duplicate_field("splitPoint")); + } + split_point__ = Some(map_.next_value()?); + } + } + } + Ok(RangeRepartition { + sort_expr: sort_expr__.unwrap_or_default(), + split_point: split_point__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.RangeRepartition", FIELDS, GeneratedVisitor) + } +} +impl serde::Serialize for RangeSplitPoint { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + use serde::ser::SerializeStruct; + let mut len = 0; + if !self.value.is_empty() { + len += 1; + } + let mut struct_ser = serializer.serialize_struct("datafusion.RangeSplitPoint", len)?; + if !self.value.is_empty() { + struct_ser.serialize_field("value", &self.value)?; + } + struct_ser.end() + } +} +impl<'de> serde::Deserialize<'de> for RangeSplitPoint { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "value", + ]; + + #[allow(clippy::enum_variant_names)] + enum GeneratedField { + Value, + } + impl<'de> serde::Deserialize<'de> for GeneratedField { + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = GeneratedField; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + #[allow(unused_variables)] + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "value" => Ok(GeneratedField::Value), + _ => Err(serde::de::Error::unknown_field(value, FIELDS)), + } + } + } + deserializer.deserialize_identifier(GeneratedVisitor) + } + } + struct GeneratedVisitor; + impl<'de> serde::de::Visitor<'de> for GeneratedVisitor { + type Value = RangeSplitPoint; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter.write_str("struct datafusion.RangeSplitPoint") + } + + fn visit_map(self, mut map_: V) -> std::result::Result + where + V: serde::de::MapAccess<'de>, + { + let mut value__ = None; + while let Some(k) = map_.next_key()? { + match k { + GeneratedField::Value => { + if value__.is_some() { + return Err(serde::de::Error::duplicate_field("value")); + } + value__ = Some(map_.next_value()?); + } + } + } + Ok(RangeSplitPoint { + value: value__.unwrap_or_default(), + }) + } + } + deserializer.deserialize_struct("datafusion.RangeSplitPoint", FIELDS, GeneratedVisitor) + } +} impl serde::Serialize for RecursionUnnestOption { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result @@ -21621,6 +24027,9 @@ impl serde::Serialize for RepartitionNode { repartition_node::PartitionMethod::Hash(v) => { struct_ser.serialize_field("hash", v)?; } + repartition_node::PartitionMethod::Range(v) => { + struct_ser.serialize_field("range", v)?; + } } } struct_ser.end() @@ -21637,6 +24046,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { "round_robin", "roundRobin", "hash", + "range", ]; #[allow(clippy::enum_variant_names)] @@ -21644,6 +24054,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { Input, RoundRobin, Hash, + Range, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -21668,6 +24079,7 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { "input" => Ok(GeneratedField::Input), "roundRobin" | "round_robin" => Ok(GeneratedField::RoundRobin), "hash" => Ok(GeneratedField::Hash), + "range" => Ok(GeneratedField::Range), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -21708,6 +24120,13 @@ impl<'de> serde::Deserialize<'de> for RepartitionNode { return Err(serde::de::Error::duplicate_field("hash")); } partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(repartition_node::PartitionMethod::Hash) +; + } + GeneratedField::Range => { + if partition_method__.is_some() { + return Err(serde::de::Error::duplicate_field("range")); + } + partition_method__ = map_.next_value::<::std::option::Option<_>>()?.map(repartition_node::PartitionMethod::Range) ; } } @@ -24604,10 +27023,16 @@ impl serde::Serialize for Unnest { if !self.exprs.is_empty() { len += 1; } + if self.outer { + len += 1; + } let mut struct_ser = serializer.serialize_struct("datafusion.Unnest", len)?; if !self.exprs.is_empty() { struct_ser.serialize_field("exprs", &self.exprs)?; } + if self.outer { + struct_ser.serialize_field("outer", &self.outer)?; + } struct_ser.end() } } @@ -24619,11 +27044,13 @@ impl<'de> serde::Deserialize<'de> for Unnest { { const FIELDS: &[&str] = &[ "exprs", + "outer", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { Exprs, + Outer, } impl<'de> serde::Deserialize<'de> for GeneratedField { fn deserialize(deserializer: D) -> std::result::Result @@ -24646,6 +27073,7 @@ impl<'de> serde::Deserialize<'de> for Unnest { { match value { "exprs" => Ok(GeneratedField::Exprs), + "outer" => Ok(GeneratedField::Outer), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } } @@ -24666,6 +27094,7 @@ impl<'de> serde::Deserialize<'de> for Unnest { V: serde::de::MapAccess<'de>, { let mut exprs__ = None; + let mut outer__ = None; while let Some(k) = map_.next_key()? { match k { GeneratedField::Exprs => { @@ -24674,10 +27103,17 @@ impl<'de> serde::Deserialize<'de> for Unnest { } exprs__ = Some(map_.next_value()?); } + GeneratedField::Outer => { + if outer__.is_some() { + return Err(serde::de::Error::duplicate_field("outer")); + } + outer__ = Some(map_.next_value()?); + } } } Ok(Unnest { exprs: exprs__.unwrap_or_default(), + outer: outer__.unwrap_or_default(), }) } } @@ -25059,15 +27495,17 @@ impl serde::Serialize for UnnestOptions { { use serde::ser::SerializeStruct; let mut len = 0; - if self.preserve_nulls { + if self.null_handling != 0 { len += 1; } if !self.recursions.is_empty() { len += 1; } let mut struct_ser = serializer.serialize_struct("datafusion.UnnestOptions", len)?; - if self.preserve_nulls { - struct_ser.serialize_field("preserveNulls", &self.preserve_nulls)?; + if self.null_handling != 0 { + let v = unnest_options::NullHandling::try_from(self.null_handling) + .map_err(|_| serde::ser::Error::custom(format!("Invalid variant {}", self.null_handling)))?; + struct_ser.serialize_field("nullHandling", &v)?; } if !self.recursions.is_empty() { struct_ser.serialize_field("recursions", &self.recursions)?; @@ -25082,14 +27520,14 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { D: serde::Deserializer<'de>, { const FIELDS: &[&str] = &[ - "preserve_nulls", - "preserveNulls", + "null_handling", + "nullHandling", "recursions", ]; #[allow(clippy::enum_variant_names)] enum GeneratedField { - PreserveNulls, + NullHandling, Recursions, } impl<'de> serde::Deserialize<'de> for GeneratedField { @@ -25112,7 +27550,7 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { E: serde::de::Error, { match value { - "preserveNulls" | "preserve_nulls" => Ok(GeneratedField::PreserveNulls), + "nullHandling" | "null_handling" => Ok(GeneratedField::NullHandling), "recursions" => Ok(GeneratedField::Recursions), _ => Err(serde::de::Error::unknown_field(value, FIELDS)), } @@ -25133,15 +27571,15 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { where V: serde::de::MapAccess<'de>, { - let mut preserve_nulls__ = None; + let mut null_handling__ = None; let mut recursions__ = None; while let Some(k) = map_.next_key()? { match k { - GeneratedField::PreserveNulls => { - if preserve_nulls__.is_some() { - return Err(serde::de::Error::duplicate_field("preserveNulls")); + GeneratedField::NullHandling => { + if null_handling__.is_some() { + return Err(serde::de::Error::duplicate_field("nullHandling")); } - preserve_nulls__ = Some(map_.next_value()?); + null_handling__ = Some(map_.next_value::()? as i32); } GeneratedField::Recursions => { if recursions__.is_some() { @@ -25152,7 +27590,7 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { } } Ok(UnnestOptions { - preserve_nulls: preserve_nulls__.unwrap_or_default(), + null_handling: null_handling__.unwrap_or_default(), recursions: recursions__.unwrap_or_default(), }) } @@ -25160,6 +27598,80 @@ impl<'de> serde::Deserialize<'de> for UnnestOptions { deserializer.deserialize_struct("datafusion.UnnestOptions", FIELDS, GeneratedVisitor) } } +impl serde::Serialize for unnest_options::NullHandling { + #[allow(deprecated)] + fn serialize(&self, serializer: S) -> std::result::Result + where + S: serde::Serializer, + { + let variant = match self { + Self::Preserve => "PRESERVE", + Self::Drop => "DROP", + Self::PreserveAndExpandEmpty => "PRESERVE_AND_EXPAND_EMPTY", + }; + serializer.serialize_str(variant) + } +} +impl<'de> serde::Deserialize<'de> for unnest_options::NullHandling { + #[allow(deprecated)] + fn deserialize(deserializer: D) -> std::result::Result + where + D: serde::Deserializer<'de>, + { + const FIELDS: &[&str] = &[ + "PRESERVE", + "DROP", + "PRESERVE_AND_EXPAND_EMPTY", + ]; + + struct GeneratedVisitor; + + impl serde::de::Visitor<'_> for GeneratedVisitor { + type Value = unnest_options::NullHandling; + + fn expecting(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(formatter, "expected one of: {:?}", &FIELDS) + } + + fn visit_i64(self, v: i64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Signed(v), &self) + }) + } + + fn visit_u64(self, v: u64) -> std::result::Result + where + E: serde::de::Error, + { + i32::try_from(v) + .ok() + .and_then(|x| x.try_into().ok()) + .ok_or_else(|| { + serde::de::Error::invalid_value(serde::de::Unexpected::Unsigned(v), &self) + }) + } + + fn visit_str(self, value: &str) -> std::result::Result + where + E: serde::de::Error, + { + match value { + "PRESERVE" => Ok(unnest_options::NullHandling::Preserve), + "DROP" => Ok(unnest_options::NullHandling::Drop), + "PRESERVE_AND_EXPAND_EMPTY" => Ok(unnest_options::NullHandling::PreserveAndExpandEmpty), + _ => Err(serde::de::Error::unknown_variant(value, FIELDS)), + } + } + } + deserializer.deserialize_any(GeneratedVisitor) + } +} impl serde::Serialize for ValuesNode { #[allow(deprecated)] fn serialize(&self, serializer: S) -> std::result::Result diff --git a/datafusion/proto/src/generated/prost.rs b/datafusion/proto-models/src/generated/prost.rs similarity index 88% rename from datafusion/proto/src/generated/prost.rs rename to datafusion/proto-models/src/generated/prost.rs index 24bea5cae9b66..ba00577ab9a1b 100644 --- a/datafusion/proto/src/generated/prost.rs +++ b/datafusion/proto-models/src/generated/prost.rs @@ -119,10 +119,6 @@ pub struct ListingTableScanNode { pub filters: ::prost::alloc::vec::Vec, #[prost(message, repeated, tag = "7")] pub table_partition_cols: ::prost::alloc::vec::Vec, - #[prost(bool, tag = "8")] - pub collect_stat: bool, - #[prost(uint32, tag = "9")] - pub target_partitions: u32, #[prost(message, repeated, tag = "13")] pub file_sort_order: ::prost::alloc::vec::Vec, #[prost( @@ -214,7 +210,7 @@ pub struct SortNode { pub struct RepartitionNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - #[prost(oneof = "repartition_node::PartitionMethod", tags = "2, 3")] + #[prost(oneof = "repartition_node::PartitionMethod", tags = "2, 3, 4")] pub partition_method: ::core::option::Option, } /// Nested message and enum types in `RepartitionNode`. @@ -225,9 +221,23 @@ pub mod repartition_node { RoundRobin(u64), #[prost(message, tag = "3")] Hash(super::HashRepartition), + #[prost(message, tag = "4")] + Range(super::RangeRepartition), } } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct RangeSplitPoint { + #[prost(message, repeated, tag = "1")] + pub value: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct RangeRepartition { + #[prost(message, repeated, tag = "1")] + pub sort_expr: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "2")] + pub split_point: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct HashRepartition { #[prost(message, repeated, tag = "1")] pub hash_expr: ::prost::alloc::vec::Vec, @@ -243,8 +253,11 @@ pub struct EmptyRelationNode { pub struct CreateExternalTableNode { #[prost(message, optional, tag = "9")] pub name: ::core::option::Option, + /// deprecated; use repeated locations #[prost(string, tag = "2")] pub location: ::prost::alloc::string::String, + #[prost(string, repeated, tag = "16")] + pub locations: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, #[prost(string, tag = "3")] pub file_type: ::prost::alloc::string::String, #[prost(message, optional, tag = "4")] @@ -344,6 +357,18 @@ pub struct AnalyzeNode { pub input: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(bool, tag = "2")] pub verbose: bool, + /// Statement-level override for `datafusion.explain.analyze_level`. + /// Absent means "fall back to session config". + #[prost(enumeration = "super::datafusion_common::MetricType", optional, tag = "3")] + pub analyze_level: ::core::option::Option, + /// Statement-level override for `datafusion.explain.analyze_categories`. + /// Absent means "fall back to session config". + #[prost(message, optional, tag = "4")] + pub analyze_categories: ::core::option::Option< + super::datafusion_common::ExplainAnalyzeCategoriesNode, + >, + #[prost(enumeration = "super::datafusion_common::ExplainFormat", tag = "5")] + pub format: i32, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ExplainNode { @@ -353,6 +378,10 @@ pub struct ExplainNode { pub verbose: bool, #[prost(enumeration = "super::datafusion_common::ExplainFormat", tag = "3")] pub format: i32, + /// Statement-level override for `datafusion.explain.show_statistics`. + /// Absent means "fall back to session config". + #[prost(bool, optional, tag = "4")] + pub show_statistics: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct AggregateNode { @@ -388,6 +417,8 @@ pub struct JoinNode { pub null_equality: i32, #[prost(message, optional, boxed, tag = "8")] pub filter: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(bool, tag = "9")] + pub null_aware: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct DistinctNode { @@ -435,6 +466,9 @@ pub struct DmlNode { pub table_name: ::core::option::Option, #[prost(message, optional, boxed, tag = "5")] pub target: ::core::option::Option<::prost::alloc::boxed::Box>, + /// Populated only when dml_type == MERGE_INTO. + #[prost(message, optional, boxed, tag = "6")] + pub merge_into: ::core::option::Option<::prost::alloc::boxed::Box>, } /// Nested message and enum types in `DmlNode`. pub mod dml_node { @@ -458,6 +492,7 @@ pub mod dml_node { InsertOverwrite = 4, InsertReplace = 5, Truncate = 6, + MergeInto = 7, } impl Type { /// String value of the enum field names used in the ProtoBuf definition. @@ -473,6 +508,7 @@ pub mod dml_node { Self::InsertOverwrite => "INSERT_OVERWRITE", Self::InsertReplace => "INSERT_REPLACE", Self::Truncate => "TRUNCATE", + Self::MergeInto => "MERGE_INTO", } } /// Creates an enum from field names used in the ProtoBuf definition. @@ -485,11 +521,117 @@ pub mod dml_node { "INSERT_OVERWRITE" => Some(Self::InsertOverwrite), "INSERT_REPLACE" => Some(Self::InsertReplace), "TRUNCATE" => Some(Self::Truncate), + "MERGE_INTO" => Some(Self::MergeInto), + _ => None, + } + } + } +} +/// Carries the ON condition and WHEN clauses of a MERGE INTO operation. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct MergeIntoOpNode { + #[prost(message, optional, boxed, tag = "1")] + pub on: ::core::option::Option<::prost::alloc::boxed::Box>, + #[prost(message, repeated, tag = "2")] + pub clauses: ::prost::alloc::vec::Vec, +} +/// A single WHEN clause within a MERGE INTO statement. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct MergeIntoClauseNode { + #[prost(enumeration = "merge_into_clause_node::Kind", tag = "1")] + pub kind: i32, + /// Optional `AND ` predicate. Absent when the clause has no predicate. + #[prost(message, optional, tag = "2")] + pub predicate: ::core::option::Option, + #[prost(message, optional, tag = "3")] + pub action: ::core::option::Option, +} +/// Nested message and enum types in `MergeIntoClauseNode`. +pub mod merge_into_clause_node { + #[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + PartialOrd, + Ord, + ::prost::Enumeration + )] + #[repr(i32)] + pub enum Kind { + Matched = 0, + NotMatched = 1, + NotMatchedByTarget = 2, + NotMatchedBySource = 3, + } + impl Kind { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Matched => "MATCHED", + Self::NotMatched => "NOT_MATCHED", + Self::NotMatchedByTarget => "NOT_MATCHED_BY_TARGET", + Self::NotMatchedBySource => "NOT_MATCHED_BY_SOURCE", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "MATCHED" => Some(Self::Matched), + "NOT_MATCHED" => Some(Self::NotMatched), + "NOT_MATCHED_BY_TARGET" => Some(Self::NotMatchedByTarget), + "NOT_MATCHED_BY_SOURCE" => Some(Self::NotMatchedBySource), _ => None, } } } } +/// The action for a single WHEN clause. +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct MergeIntoActionNode { + #[prost(oneof = "merge_into_action_node::Action", tags = "1, 2, 3")] + pub action: ::core::option::Option, +} +/// Nested message and enum types in `MergeIntoActionNode`. +pub mod merge_into_action_node { + #[derive(Clone, PartialEq, ::prost::Oneof)] + pub enum Action { + #[prost(message, tag = "1")] + Update(super::MergeUpdateAction), + #[prost(message, tag = "2")] + Insert(super::MergeInsertAction), + #[prost(message, tag = "3")] + Delete(super::MergeDeleteAction), + } +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct MergeUpdateAction { + #[prost(message, repeated, tag = "1")] + pub assignments: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct MergeAssignment { + #[prost(string, tag = "1")] + pub column: ::prost::alloc::string::String, + #[prost(message, optional, tag = "2")] + pub value: ::core::option::Option, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct MergeInsertAction { + /// May be empty (meaning all columns). + #[prost(string, repeated, tag = "1")] + pub columns: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + /// One expression per inserted column. + #[prost(message, repeated, tag = "2")] + pub values: ::prost::alloc::vec::Vec, +} +#[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] +pub struct MergeDeleteAction {} #[derive(Clone, PartialEq, ::prost::Message)] pub struct UnnestNode { #[prost(message, optional, boxed, tag = "1")] @@ -528,11 +670,57 @@ pub struct ColumnUnnestListRecursion { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct UnnestOptions { - #[prost(bool, tag = "1")] - pub preserve_nulls: bool, + #[prost(enumeration = "unnest_options::NullHandling", tag = "3")] + pub null_handling: i32, #[prost(message, repeated, tag = "2")] pub recursions: ::prost::alloc::vec::Vec, } +/// Nested message and enum types in `UnnestOptions`. +pub mod unnest_options { + #[derive( + Clone, + Copy, + Debug, + PartialEq, + Eq, + Hash, + PartialOrd, + Ord, + ::prost::Enumeration + )] + #[repr(i32)] + pub enum NullHandling { + /// Preserve nulls; empty lists produce no rows. The historical default. + Preserve = 0, + /// Drop both null and empty lists from the output. + Drop = 1, + /// Preserve nulls, and additionally expand empty lists into a single + /// NULL output row (outer-unnest semantics). + PreserveAndExpandEmpty = 2, + } + impl NullHandling { + /// String value of the enum field names used in the ProtoBuf definition. + /// + /// The values are not transformed in any way and thus are considered stable + /// (if the ProtoBuf definition does not change) and safe for programmatic use. + pub fn as_str_name(&self) -> &'static str { + match self { + Self::Preserve => "PRESERVE", + Self::Drop => "DROP", + Self::PreserveAndExpandEmpty => "PRESERVE_AND_EXPAND_EMPTY", + } + } + /// Creates an enum from field names used in the ProtoBuf definition. + pub fn from_str_name(value: &str) -> ::core::option::Option { + match value { + "PRESERVE" => Some(Self::Preserve), + "DROP" => Some(Self::Drop), + "PRESERVE_AND_EXPAND_EMPTY" => Some(Self::PreserveAndExpandEmpty), + _ => None, + } + } + } +} #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] pub struct RecursionUnnestOption { #[prost(message, optional, tag = "1")] @@ -582,7 +770,7 @@ pub struct SubqueryAliasNode { pub struct LogicalExprNode { #[prost( oneof = "logical_expr_node::ExprType", - tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 13, 14, 15, 17, 18, 19, 20, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36" + tags = "1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 13, 14, 15, 17, 18, 19, 20, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39" )] pub expr_type: ::core::option::Option, } @@ -663,6 +851,12 @@ pub mod logical_expr_node { /// Subquery expressions #[prost(message, tag = "36")] ScalarSubqueryExpr(::prost::alloc::boxed::Box), + #[prost(message, tag = "37")] + HigherOrderUdfExpr(super::HigherOrderUdfExprNode), + #[prost(message, tag = "38")] + Lambda(::prost::alloc::boxed::Box), + #[prost(message, tag = "39")] + LambdaVariable(super::LambdaVariable), } } #[derive(Clone, PartialEq, Eq, Hash, ::prost::Message)] @@ -815,6 +1009,10 @@ pub struct NegativeNode { pub struct Unnest { #[prost(message, repeated, tag = "1")] pub exprs: ::prost::alloc::vec::Vec, + /// When true, this Unnest expression has outer-unnest semantics: NULL and + /// empty input lists both produce a single NULL output row. + #[prost(bool, tag = "2")] + pub outer: bool, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct InListNode { @@ -852,6 +1050,29 @@ pub struct ScalarUdfExprNode { pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec>, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct HigherOrderUdfExprNode { + #[prost(string, tag = "1")] + pub fun_name: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "2")] + pub args: ::prost::alloc::vec::Vec, + #[prost(bytes = "vec", optional, tag = "3")] + pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct Lambda { + #[prost(string, repeated, tag = "1")] + pub params: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(message, optional, boxed, tag = "2")] + pub body: ::core::option::Option<::prost::alloc::boxed::Box>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct LambdaVariable { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + #[prost(message, optional, tag = "2")] + pub field: ::core::option::Option, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct WindowExprNode { #[prost(message, repeated, tag = "4")] pub exprs: ::prost::alloc::vec::Vec, @@ -1336,7 +1557,7 @@ pub struct PhysicalExprNode { pub expr_id: ::core::option::Option, #[prost( oneof = "physical_expr_node::ExprType", - tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23" + tags = "1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 14, 15, 16, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27" )] pub expr_type: ::core::option::Option, } @@ -1393,6 +1614,14 @@ pub mod physical_expr_node { ScalarSubquery(super::PhysicalScalarSubqueryExprNode), #[prost(message, tag = "23")] DynamicFilter(::prost::alloc::boxed::Box), + #[prost(message, tag = "24")] + HigherOrderUdf(super::PhysicalHigherOrderUdfNode), + #[prost(message, tag = "25")] + Lambda(::prost::alloc::boxed::Box), + #[prost(message, tag = "26")] + LambdaVariable(super::PhysicalLambdaVariableExprNode), + #[prost(message, tag = "27")] + RangeExpr(super::PhysicalRangeExprNode), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -1424,6 +1653,29 @@ pub struct PhysicalScalarUdfNode { pub return_field_name: ::prost::alloc::string::String, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalHigherOrderUdfNode { + #[prost(string, tag = "1")] + pub name: ::prost::alloc::string::String, + #[prost(message, repeated, tag = "2")] + pub args: ::prost::alloc::vec::Vec, + #[prost(bytes = "vec", optional, tag = "3")] + pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalLambdaExprNode { + #[prost(string, repeated, tag = "1")] + pub params: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(message, optional, boxed, tag = "2")] + pub body: ::core::option::Option<::prost::alloc::boxed::Box>, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalLambdaVariableExprNode { + #[prost(uint32, tag = "1")] + pub index: u32, + #[prost(message, optional, tag = "2")] + pub field: ::core::option::Option, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct PhysicalAggregateExprNode { #[prost(message, repeated, tag = "2")] pub expr: ::prost::alloc::vec::Vec, @@ -1437,6 +1689,8 @@ pub struct PhysicalAggregateExprNode { pub fun_definition: ::core::option::Option<::prost::alloc::vec::Vec>, #[prost(string, tag = "8")] pub human_display: ::prost::alloc::string::String, + #[prost(bool, tag = "9")] + pub is_reversed: bool, #[prost(oneof = "physical_aggregate_expr_node::AggregateFunction", tags = "4")] pub aggregate_function: ::core::option::Option< physical_aggregate_expr_node::AggregateFunction, @@ -1609,6 +1863,13 @@ pub struct PhysicalHashExprNode { pub description: ::prost::alloc::string::String, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalRangeExprNode { + #[prost(message, repeated, tag = "1")] + pub sort_expr: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "2")] + pub split_point: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct FilterExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, @@ -1675,6 +1936,8 @@ pub struct FileScanExecConf { pub batch_size: ::core::option::Option, #[prost(message, optional, tag = "13")] pub projection_exprs: ::core::option::Option, + #[prost(message, optional, tag = "15")] + pub output_partitioning: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ParquetScanExecNode { @@ -1777,6 +2040,15 @@ pub struct HashJoinExecNode { /// Optional dynamic filter expression for pushing down to the probe side. #[prost(message, optional, tag = "11")] pub dynamic_filter: ::core::option::Option, + /// Optional row limit pushed into the join by the `limit_pushdown` rule. + /// + /// This is presence-tracked (`optional`) on purpose: messages produced by + /// versions predating this field carry no `fetch` at all, and a plain proto3 + /// scalar would decode that absence as `0`, i.e. "fetch 0 rows", silently + /// turning old plans into empty results. With `optional`, absent decodes to + /// `None`, which is the correct reading of an older message. + #[prost(uint64, optional, tag = "12")] + pub fetch: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct SymmetricHashJoinExecNode { @@ -1834,6 +2106,8 @@ pub struct AnalyzeExecNode { pub has_metric_categories: bool, #[prost(string, repeated, tag = "6")] pub metric_categories: ::prost::alloc::vec::Vec<::prost::alloc::string::String>, + #[prost(enumeration = "super::datafusion_common::ExplainFormat", tag = "7")] + pub format: i32, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct CrossJoinExecNode { @@ -1865,11 +2139,19 @@ pub struct JoinOn { pub struct EmptyExecNode { #[prost(message, optional, tag = "1")] pub schema: ::core::option::Option, + /// Number of output partitions. Absent (0) means a single partition, so that + /// plans encoded before this field existed decode to the previous default. + #[prost(uint32, tag = "2")] + pub partitions: u32, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct PlaceholderRowExecNode { #[prost(message, optional, tag = "1")] pub schema: ::core::option::Option, + /// Number of output partitions. Absent (0) means a single partition, so that + /// plans encoded before this field existed decode to the previous default. + #[prost(uint32, tag = "2")] + pub partitions: u32, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct ProjectionExecNode { @@ -1959,6 +2241,9 @@ pub struct AggregateExecNode { /// Optional dynamic filter expression for pushing down to the child. #[prost(message, optional, tag = "13")] pub dynamic_filter: ::core::option::Option, + /// Output schema preserved by physical optimizer rewrites. + #[prost(message, optional, tag = "14")] + pub schema: ::core::option::Option, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct GlobalLimitExecNode { @@ -1970,6 +2255,9 @@ pub struct GlobalLimitExecNode { /// Maximum number of rows to fetch; negative means no limit #[prost(int64, tag = "3")] pub fetch: i64, + /// Ordering the limit must preserve; empty means none + #[prost(message, repeated, tag = "4")] + pub required_ordering: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct LocalLimitExecNode { @@ -1977,6 +2265,9 @@ pub struct LocalLimitExecNode { pub input: ::core::option::Option<::prost::alloc::boxed::Box>, #[prost(uint32, tag = "2")] pub fetch: u32, + /// Ordering the limit must preserve; empty means none + #[prost(message, repeated, tag = "3")] + pub required_ordering: ::prost::alloc::vec::Vec, } #[derive(Clone, PartialEq, ::prost::Message)] pub struct SortExecNode { @@ -2040,14 +2331,26 @@ pub struct PhysicalHashRepartition { pub partition_count: u64, } #[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalRangePartitioning { + #[prost(message, repeated, tag = "1")] + pub sort_expr: ::prost::alloc::vec::Vec, + #[prost(message, repeated, tag = "2")] + pub split_point: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] +pub struct PhysicalRangeSplitPoint { + #[prost(message, repeated, tag = "1")] + pub value: ::prost::alloc::vec::Vec, +} +#[derive(Clone, PartialEq, ::prost::Message)] pub struct RepartitionExecNode { #[prost(message, optional, boxed, tag = "1")] pub input: ::core::option::Option<::prost::alloc::boxed::Box>, - /// oneof partition_method { + /// Legacy direct partitioning fields: /// uint64 round_robin = 2; /// PhysicalHashRepartition hash = 3; /// uint64 unknown = 4; - /// } + /// New partitioning variants are stored in `partitioning`. #[prost(message, optional, tag = "5")] pub partitioning: ::core::option::Option, #[prost(bool, tag = "6")] @@ -2055,7 +2358,7 @@ pub struct RepartitionExecNode { } #[derive(Clone, PartialEq, ::prost::Message)] pub struct Partitioning { - #[prost(oneof = "partitioning::PartitionMethod", tags = "1, 2, 3")] + #[prost(oneof = "partitioning::PartitionMethod", tags = "1, 2, 3, 4")] pub partition_method: ::core::option::Option, } /// Nested message and enum types in `Partitioning`. @@ -2068,6 +2371,8 @@ pub mod partitioning { Hash(super::PhysicalHashRepartition), #[prost(uint64, tag = "3")] Unknown(u64), + #[prost(message, tag = "4")] + Range(super::PhysicalRangePartitioning), } } #[derive(Clone, PartialEq, ::prost::Message)] @@ -2102,6 +2407,8 @@ pub struct PartitionedFile { pub range: ::core::option::Option, #[prost(message, optional, tag = "6")] pub statistics: ::core::option::Option, + #[prost(message, optional, tag = "7")] + pub arrow_schema: ::core::option::Option, } #[derive(Clone, Copy, PartialEq, Eq, Hash, ::prost::Message)] pub struct FileRange { diff --git a/datafusion/proto-models/src/lib.rs b/datafusion/proto-models/src/lib.rs new file mode 100644 index 0000000000000..3276c0811e2c5 --- /dev/null +++ b/datafusion/proto-models/src/lib.rs @@ -0,0 +1,59 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +#![doc( + html_logo_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg", + html_favicon_url = "https://raw.githubusercontent.com/apache/datafusion/19fe44cf2f30cbdd63d4a4f52c74055163c6cc38/docs/logos/standalone_logo/logo_original.svg" +)] +#![cfg_attr(docsrs, feature(doc_cfg))] +// Make sure fast / cheap clones on Arc are explicit: +// https://github.com/apache/datafusion/issues/11143 +#![deny(clippy::clone_on_ref_ptr)] + +//! `prost`-generated DataFusion protobuf model types. +//! +//! This crate contains the generated structs for DataFusion's logical and +//! physical plan protobuf schemas (see `proto/datafusion.proto`), plus the +//! [`From`] / [`TryFrom`] conversions between them and the `datafusion-common` +//! types they mirror. Those conversions live here because their DataFusion side +//! sits *below* this crate in the dependency graph and so cannot host the impls +//! itself — see [`from_proto`] and [`to_proto`]. It is the schema source of +//! truth for [`datafusion-proto`]. +//! +//! Most users should depend on [`datafusion-proto`] instead, which re-exports +//! these types under [`datafusion_proto::protobuf`]. +//! +//! [`datafusion-proto`]: https://crates.io/crates/datafusion-proto +//! [`datafusion-proto-common`]: https://crates.io/crates/datafusion-proto-common +//! [`datafusion_proto::protobuf`]: https://docs.rs/datafusion-proto/latest/datafusion_proto/protobuf/index.html + +pub mod from_proto; +pub mod generated; +pub mod to_proto; + +/// All DataFusion protobuf model types. +/// +/// Includes both the types declared in `datafusion.proto` and the +/// `datafusion_proto_common` types it imports, in a single flat namespace +/// so consumers can `use datafusion_proto_models::protobuf::*;`. +pub mod protobuf { + pub use crate::generated::datafusion::*; +} + +/// Re-export of the `datafusion_proto_common` types as exposed through this +/// crate's generated module, for callers that want the common-only namespace. +pub use generated::datafusion_common; diff --git a/datafusion/proto-models/src/to_proto.rs b/datafusion/proto-models/src/to_proto.rs new file mode 100644 index 0000000000000..d1c857a7c5cba --- /dev/null +++ b/datafusion/proto-models/src/to_proto.rs @@ -0,0 +1,325 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Conversions from `datafusion-common` types to the protobuf messages in this +//! crate. +//! +//! See [`crate::from_proto`] for why the impls live here rather than next to +//! the DataFusion types. + +use datafusion_common::DataFusionError; +use datafusion_common::display::{PlanType, StringifiedPlan}; +use datafusion_common::{ + JoinConstraint, JoinType, NullEquality, TableReference, UnnestOptions, +}; + +use crate::generated::datafusion_common::EmptyMessage; +use crate::protobuf::{ + self, AnalyzedLogicalPlanType, OptimizedLogicalPlanType, OptimizedPhysicalPlanType, + RecursionUnnestOption, + plan_type::PlanTypeEnum::{ + AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, + FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, + InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, + InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, + PhysicalPlanError, + }, +}; + +impl From<&UnnestOptions> for protobuf::UnnestOptions { + fn from(opts: &UnnestOptions) -> Self { + use datafusion_common::NullHandling; + use protobuf::unnest_options::NullHandling as ProtoNullHandling; + let null_handling = match opts.null_handling { + NullHandling::Preserve => ProtoNullHandling::Preserve, + NullHandling::Drop => ProtoNullHandling::Drop, + NullHandling::PreserveAndExpandEmpty => { + ProtoNullHandling::PreserveAndExpandEmpty + } + } as i32; + Self { + null_handling, + recursions: opts + .recursions + .iter() + .map(|r| RecursionUnnestOption { + input_column: Some((&r.input_column).into()), + output_column: Some((&r.output_column).into()), + depth: r.depth as u32, + }) + .collect(), + } + } +} + +impl From<&StringifiedPlan> for protobuf::StringifiedPlan { + fn from(stringified_plan: &StringifiedPlan) -> Self { + Self { + plan_type: match stringified_plan.clone().plan_type { + PlanType::InitialLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(InitialLogicalPlan(EmptyMessage {})), + }), + PlanType::AnalyzedLogicalPlan { analyzer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(AnalyzedLogicalPlan( + AnalyzedLogicalPlanType { analyzer_name }, + )), + }) + } + PlanType::FinalAnalyzedLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalAnalyzedLogicalPlan(EmptyMessage {})), + }), + PlanType::OptimizedLogicalPlan { optimizer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(OptimizedLogicalPlan( + OptimizedLogicalPlanType { optimizer_name }, + )), + }) + } + PlanType::FinalLogicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalLogicalPlan(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlan(EmptyMessage {})), + }), + PlanType::OptimizedPhysicalPlan { optimizer_name } => { + Some(protobuf::PlanType { + plan_type_enum: Some(OptimizedPhysicalPlan( + OptimizedPhysicalPlanType { optimizer_name }, + )), + }) + } + PlanType::FinalPhysicalPlan => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlan(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlanWithStats => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlanWithStats(EmptyMessage {})), + }), + PlanType::InitialPhysicalPlanWithSchema => Some(protobuf::PlanType { + plan_type_enum: Some(InitialPhysicalPlanWithSchema(EmptyMessage {})), + }), + PlanType::FinalPhysicalPlanWithStats => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlanWithStats(EmptyMessage {})), + }), + PlanType::FinalPhysicalPlanWithSchema => Some(protobuf::PlanType { + plan_type_enum: Some(FinalPhysicalPlanWithSchema(EmptyMessage {})), + }), + PlanType::PhysicalPlanError => Some(protobuf::PlanType { + plan_type_enum: Some(PhysicalPlanError(EmptyMessage {})), + }), + }, + plan: stringified_plan.plan.to_string(), + } + } +} + +impl From for protobuf::TableReference { + fn from(t: TableReference) -> Self { + use protobuf::table_reference::TableReferenceEnum; + let table_reference_enum = match t { + TableReference::Bare { table } => { + TableReferenceEnum::Bare(protobuf::BareTableReference { + table: table.to_string(), + }) + } + TableReference::Partial { schema, table } => { + TableReferenceEnum::Partial(protobuf::PartialTableReference { + schema: schema.to_string(), + table: table.to_string(), + }) + } + TableReference::Full { + catalog, + schema, + table, + } => TableReferenceEnum::Full(protobuf::FullTableReference { + catalog: catalog.to_string(), + schema: schema.to_string(), + table: table.to_string(), + }), + }; + + protobuf::TableReference { + table_reference_enum: Some(table_reference_enum), + } + } +} + +impl From for protobuf::JoinType { + fn from(t: JoinType) -> Self { + match t { + JoinType::Inner => protobuf::JoinType::Inner, + JoinType::Left => protobuf::JoinType::Left, + JoinType::Right => protobuf::JoinType::Right, + JoinType::Full => protobuf::JoinType::Full, + JoinType::LeftSemi => protobuf::JoinType::Leftsemi, + JoinType::RightSemi => protobuf::JoinType::Rightsemi, + JoinType::LeftAnti => protobuf::JoinType::Leftanti, + JoinType::RightAnti => protobuf::JoinType::Rightanti, + JoinType::LeftMark => protobuf::JoinType::Leftmark, + JoinType::RightMark => protobuf::JoinType::Rightmark, + } + } +} + +impl From for protobuf::JoinConstraint { + fn from(t: JoinConstraint) -> Self { + match t { + JoinConstraint::On => protobuf::JoinConstraint::On, + JoinConstraint::Using => protobuf::JoinConstraint::Using, + } + } +} + +impl From for protobuf::NullEquality { + fn from(t: NullEquality) -> Self { + match t { + NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, + } + } +} + +/// Encode any slice of file-like values as a [`protobuf::FileGroup`]. +/// +/// `datafusion-datasource` cannot host this impl: `&T` is `#[fundamental]` but +/// `[T]` is not, so `&[PartitionedFile]` counts as foreign there and the orphan +/// rule rejects it. Here the *self* type is local, which is all the orphan rule +/// needs — and staying generic over the element means this crate never has to +/// name `PartitionedFile`, which lives above it in the dependency graph. +/// +/// The element bound is satisfied by +/// `impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile` in +/// `datafusion-datasource`, so `protobuf::FileGroup::try_from(&files[..])` +/// resolves for callers exactly as it did before the proto types were split out. +impl TryFrom<&[T]> for protobuf::FileGroup +where + for<'a> &'a T: TryInto, +{ + type Error = DataFusionError; + + fn try_from(files: &[T]) -> Result { + Ok(protobuf::FileGroup { + files: files + .iter() + .map(TryInto::try_into) + .collect::, _>>()?, + }) + } +} + +#[cfg(test)] +mod tests { + use datafusion_common::{NullHandling, RecursionUnnestOption}; + + use super::*; + + #[test] + fn table_reference_roundtrip() { + for reference in [ + TableReference::bare("t"), + TableReference::partial("s", "t"), + TableReference::full("c", "s", "t"), + ] { + let encoded = protobuf::TableReference::from(reference.clone()); + let decoded = TableReference::try_from(encoded).unwrap(); + assert_eq!(decoded, reference); + } + } + + #[test] + fn table_reference_from_proto_rejects_missing_oneof() { + let proto = protobuf::TableReference { + table_reference_enum: None, + }; + let err = TableReference::try_from(proto).unwrap_err(); + assert!( + err.to_string().contains("table_reference_enum"), + "unexpected error: {err}" + ); + } + + #[test] + fn join_enums_roundtrip() { + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftSemi, + JoinType::RightSemi, + JoinType::LeftAnti, + JoinType::RightAnti, + JoinType::LeftMark, + JoinType::RightMark, + ] { + assert_eq!( + JoinType::from(protobuf::JoinType::from(join_type)), + join_type + ); + } + for constraint in [JoinConstraint::On, JoinConstraint::Using] { + assert_eq!( + JoinConstraint::from(protobuf::JoinConstraint::from(constraint)), + constraint + ); + } + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + assert_eq!( + NullEquality::from(protobuf::NullEquality::from(null_equality)), + null_equality + ); + } + } + + #[test] + fn unnest_options_roundtrip() { + let options = UnnestOptions { + null_handling: NullHandling::Drop, + recursions: vec![RecursionUnnestOption { + input_column: "a".into(), + output_column: "b".into(), + depth: 2, + }], + }; + + let encoded = protobuf::UnnestOptions::from(&options); + let decoded = UnnestOptions::from(&encoded); + + assert_eq!(decoded.null_handling, options.null_handling); + assert_eq!(decoded.recursions, options.recursions); + } + + #[test] + fn stringified_plan_roundtrip() { + let plan = StringifiedPlan::new( + PlanType::OptimizedLogicalPlan { + optimizer_name: "push_down_filter".to_string(), + }, + "some plan", + ); + + let encoded = protobuf::StringifiedPlan::from(&plan); + let decoded = StringifiedPlan::from(&encoded); + + assert_eq!(decoded.plan_type, plan.plan_type); + assert_eq!(decoded.plan, plan.plan); + } +} diff --git a/datafusion/proto/Cargo.toml b/datafusion/proto/Cargo.toml index 4484846813296..b6e9d258681e8 100644 --- a/datafusion/proto/Cargo.toml +++ b/datafusion/proto/Cargo.toml @@ -31,12 +31,23 @@ rust-version = { workspace = true } [package.metadata.docs.rs] all-features = true +# Note: add additional linter rules in lib.rs. +# Rust does not support workspace + new linter rules in subcrates yet +# https://github.com/rust-lang/cargo/issues/13157 +[lints] +workspace = true + [lib] name = "datafusion_proto" [features] default = ["parquet"] -json = ["pbjson", "serde", "serde_json", "datafusion-proto-common/json"] +recursive_protection = ["dep:recursive"] +json = [ + "serde_json", + "datafusion-proto-common/json", + "datafusion-proto-models/json", +] parquet = ["datafusion-datasource-parquet", "datafusion-common/parquet", "datafusion/parquet"] avro = ["datafusion-datasource-avro"] @@ -46,27 +57,26 @@ avro = ["datafusion-datasource-avro"] [dependencies] arrow = { workspace = true } -chrono = { workspace = true } datafusion-catalog = { workspace = true } datafusion-catalog-listing = { workspace = true } datafusion-common = { workspace = true } -datafusion-datasource = { workspace = true } -datafusion-datasource-arrow = { workspace = true } -datafusion-datasource-avro = { workspace = true, optional = true } -datafusion-datasource-csv = { workspace = true } -datafusion-datasource-json = { workspace = true } -datafusion-datasource-parquet = { workspace = true, optional = true } +datafusion-datasource = { workspace = true, features = ["proto"] } +datafusion-datasource-arrow = { workspace = true, features = ["proto"] } +datafusion-datasource-avro = { workspace = true, optional = true, features = ["proto"] } +datafusion-datasource-csv = { workspace = true, features = ["proto"] } +datafusion-datasource-json = { workspace = true, features = ["proto"] } +datafusion-datasource-parquet = { workspace = true, optional = true, features = ["proto"] } datafusion-execution = { workspace = true } -datafusion-expr = { workspace = true } +datafusion-expr = { workspace = true, features = ["proto"] } datafusion-functions-table = { workspace = true } -datafusion-physical-expr = { workspace = true } -datafusion-physical-expr-common = { workspace = true } -datafusion-physical-plan = { workspace = true } +datafusion-physical-expr = { workspace = true, features = ["proto"] } +datafusion-physical-expr-common = { workspace = true, features = ["proto"] } +datafusion-physical-plan = { workspace = true, features = ["proto"] } datafusion-proto-common = { workspace = true } +datafusion-proto-models = { workspace = true } object_store = { workspace = true } -pbjson = { workspace = true, optional = true } prost = { workspace = true } -serde = { version = "1.0", optional = true } +recursive = { workspace = true, optional = true } serde_json = { workspace = true, optional = true } [dev-dependencies] diff --git a/datafusion/proto/regen.sh b/datafusion/proto/regen.sh index 02970a90add47..c4bcea9ff5408 100755 --- a/datafusion/proto/regen.sh +++ b/datafusion/proto/regen.sh @@ -17,5 +17,7 @@ # specific language governing permissions and limitations # under the License. +# The proto schema and code generation now live in `datafusion-proto-models`. +# This script is kept as a convenience wrapper. repo_root=$(git rev-parse --show-toplevel) -cd "$repo_root" && cargo run --manifest-path datafusion/proto/gen/Cargo.toml +exec "$repo_root/datafusion/proto-models/regen.sh" diff --git a/datafusion/proto/src/bytes/mod.rs b/datafusion/proto/src/bytes/mod.rs index 2b7d7ed8e849b..388e373c3fdff 100644 --- a/datafusion/proto/src/bytes/mod.rs +++ b/datafusion/proto/src/bytes/mod.rs @@ -192,6 +192,10 @@ pub fn physical_plan_to_bytes(plan: Arc) -> Result { /// Serialize a PhysicalPlan as JSON #[cfg(feature = "json")] +#[expect( + clippy::needless_pass_by_value, + reason = "Preserve the existing public API" +)] pub fn physical_plan_to_json(plan: Arc) -> Result { let extension_codec = DefaultPhysicalExtensionCodec {}; let proto_converter = DefaultPhysicalProtoConverter {}; @@ -213,6 +217,7 @@ pub fn physical_plan_to_bytes_with_extension_codec( /// Serialize a PhysicalPlan as bytes, using the provided extension codec /// and protobuf converter. +#[expect(clippy::needless_pass_by_value)] // Taking the plan by value is part of the public API pub fn physical_plan_to_bytes_with_proto_converter( plan: Arc, extension_codec: &dyn PhysicalExtensionCodec, diff --git a/datafusion/proto/src/common.rs b/datafusion/proto/src/common.rs index 22ded708d8c71..dd3dc752e1892 100644 --- a/datafusion/proto/src/common.rs +++ b/datafusion/proto/src/common.rs @@ -15,27 +15,6 @@ // specific language governing permissions and limitations // under the License. -use datafusion_common::{Result, assert_eq_or_internal_err, internal_datafusion_err}; - -pub(crate) fn str_to_byte(s: &String, description: &str) -> Result { - assert_eq_or_internal_err!( - s.len(), - 1, - "Invalid CSV {description}: expected single character, got {s}" - ); - Ok(s.as_bytes()[0]) -} - -pub(crate) fn byte_to_string(b: u8, description: &str) -> Result { - let b = &[b]; - let b = std::str::from_utf8(b).map_err(|_| { - internal_datafusion_err!( - "Invalid CSV {description}: can not represent {b:0x?} as utf8" - ) - })?; - Ok(b.to_owned()) -} - #[macro_export] macro_rules! convert_required { ($PB:expr) => {{ diff --git a/datafusion/proto/src/generated/datafusion.rs b/datafusion/proto/src/generated/datafusion.rs deleted file mode 100644 index 8b137891791fe..0000000000000 --- a/datafusion/proto/src/generated/datafusion.rs +++ /dev/null @@ -1 +0,0 @@ - diff --git a/datafusion/proto/src/lib.rs b/datafusion/proto/src/lib.rs index 7ddc930fa257e..71feae506dc6f 100644 --- a/datafusion/proto/src/lib.rs +++ b/datafusion/proto/src/lib.rs @@ -123,12 +123,10 @@ //! ``` pub mod bytes; pub mod common; -pub mod generated; pub mod logical_plan; pub mod physical_plan; pub mod protobuf { - pub use crate::generated::datafusion::*; pub use datafusion_proto_common::common::proto_error; pub use datafusion_proto_common::protobuf_common::{ ArrowFormat, ArrowOptions, ArrowType, AvroFormat, AvroOptions, CsvFormat, @@ -136,6 +134,40 @@ pub mod protobuf { ScalarValue, Schema, }; pub use datafusion_proto_common::{FromProtoError, ToProtoError}; + // Re-export every type from `datafusion-proto-models`'s generated module + // so the existing `datafusion_proto::protobuf::Foo` paths keep resolving. + // Going through the deeper `generated::datafusion` path (rather than + // `datafusion_proto_models::protobuf`, which is itself a `pub use ::*`) + // avoids a double wildcard re-export that some tools (cargo-semver-checks) + // don't follow. + pub use datafusion_proto_models::generated::datafusion::*; +} + +/// Backwards-compatible re-export of the moved generated types. +/// +/// The prost-generated structs now live in `datafusion-proto-models`; +/// this module preserves the legacy `datafusion_proto::generated::*` paths +/// for downstream callers. Prefer the [`protobuf`] module (or +/// [`datafusion_proto_models`] directly) in new code. +#[deprecated( + since = "53.1.0", + note = "use `datafusion_proto::protobuf` (or `datafusion_proto_models::protobuf`) instead" +)] +pub mod generated { + /// Re-export of the prost-generated types defined in `datafusion.proto`. + #[deprecated( + since = "53.1.0", + note = "use `datafusion_proto::protobuf` (or `datafusion_proto_models::protobuf`) instead" + )] + pub use datafusion_proto_models::generated::datafusion; + + /// Re-export of the prost-generated common types defined in + /// `datafusion_common.proto`. + #[deprecated( + since = "53.1.0", + note = "use `datafusion_proto_common::protobuf_common` instead" + )] + pub use datafusion_proto_models::generated::datafusion_common; } #[cfg(doctest)] diff --git a/datafusion/proto/src/logical_plan/file_formats.rs b/datafusion/proto/src/logical_plan/file_formats.rs index 683b6a612a53f..10c54cf55c5e7 100644 --- a/datafusion/proto/src/logical_plan/file_formats.rs +++ b/datafusion/proto/src/logical_plan/file_formats.rs @@ -18,15 +18,9 @@ use std::sync::Arc; use super::LogicalExtensionCodec; -use crate::protobuf::{ - CsvOptions as CsvOptionsProto, CsvQuoteStyle as CsvQuoteStyleProto, - JsonOptions as JsonOptionsProto, -}; +use crate::protobuf::{CsvOptions as CsvOptionsProto, JsonOptions as JsonOptionsProto}; use datafusion_common::config::{CsvOptions, JsonOptions}; -use datafusion_common::{ - TableReference, exec_datafusion_err, exec_err, not_impl_err, - parsers::{CompressionTypeVariant, CsvQuoteStyle}, -}; +use datafusion_common::{TableReference, exec_datafusion_err, exec_err, not_impl_err}; use datafusion_datasource::file_format::FileFormatFactory; use datafusion_datasource_arrow::file_format::ArrowFormatFactory; use datafusion_datasource_csv::file_format::CsvFormatFactory; @@ -37,153 +31,6 @@ use prost::Message; #[derive(Debug)] pub struct CsvLogicalExtensionCodec; -impl CsvOptionsProto { - fn from_factory(factory: &CsvFormatFactory) -> Self { - if let Some(options) = &factory.options { - CsvOptionsProto { - has_header: options.has_header.map_or(vec![], |v| vec![v as u8]), - delimiter: vec![options.delimiter], - quote: vec![options.quote], - terminator: options.terminator.map_or(vec![], |v| vec![v]), - escape: options.escape.map_or(vec![], |v| vec![v]), - double_quote: options.double_quote.map_or(vec![], |v| vec![v as u8]), - compression: options.compression as i32, - schema_infer_max_rec: options.schema_infer_max_rec.map(|v| v as u64), - date_format: options.date_format.clone().unwrap_or_default(), - datetime_format: options.datetime_format.clone().unwrap_or_default(), - timestamp_format: options.timestamp_format.clone().unwrap_or_default(), - timestamp_tz_format: options - .timestamp_tz_format - .clone() - .unwrap_or_default(), - time_format: options.time_format.clone().unwrap_or_default(), - null_value: options.null_value.clone().unwrap_or_default(), - null_regex: options.null_regex.clone().unwrap_or_default(), - comment: options.comment.map_or(vec![], |v| vec![v]), - newlines_in_values: options - .newlines_in_values - .map_or(vec![], |v| vec![v as u8]), - truncated_rows: options.truncated_rows.map_or(vec![], |v| vec![v as u8]), - compression_level: options.compression_level, - quote_style: options.quote_style as i32, - ignore_leading_whitespace: options - .ignore_leading_whitespace - .map_or(vec![], |v| vec![v as u8]), - ignore_trailing_whitespace: options - .ignore_trailing_whitespace - .map_or(vec![], |v| vec![v as u8]), - } - } else { - CsvOptionsProto::default() - } - } -} - -impl From<&CsvOptionsProto> for CsvOptions { - fn from(proto: &CsvOptionsProto) -> Self { - CsvOptions { - has_header: if !proto.has_header.is_empty() { - Some(proto.has_header[0] != 0) - } else { - None - }, - delimiter: proto.delimiter.first().copied().unwrap_or(b','), - quote: proto.quote.first().copied().unwrap_or(b'"'), - terminator: if !proto.terminator.is_empty() { - Some(proto.terminator[0]) - } else { - None - }, - escape: if !proto.escape.is_empty() { - Some(proto.escape[0]) - } else { - None - }, - double_quote: if !proto.double_quote.is_empty() { - Some(proto.double_quote[0] != 0) - } else { - None - }, - compression: match proto.compression { - 0 => CompressionTypeVariant::GZIP, - 1 => CompressionTypeVariant::BZIP2, - 2 => CompressionTypeVariant::XZ, - 3 => CompressionTypeVariant::ZSTD, - _ => CompressionTypeVariant::UNCOMPRESSED, - }, - schema_infer_max_rec: proto.schema_infer_max_rec.map(|v| v as usize), - date_format: if proto.date_format.is_empty() { - None - } else { - Some(proto.date_format.clone()) - }, - datetime_format: if proto.datetime_format.is_empty() { - None - } else { - Some(proto.datetime_format.clone()) - }, - timestamp_format: if proto.timestamp_format.is_empty() { - None - } else { - Some(proto.timestamp_format.clone()) - }, - timestamp_tz_format: if proto.timestamp_tz_format.is_empty() { - None - } else { - Some(proto.timestamp_tz_format.clone()) - }, - time_format: if proto.time_format.is_empty() { - None - } else { - Some(proto.time_format.clone()) - }, - null_value: if proto.null_value.is_empty() { - None - } else { - Some(proto.null_value.clone()) - }, - null_regex: if proto.null_regex.is_empty() { - None - } else { - Some(proto.null_regex.clone()) - }, - comment: if !proto.comment.is_empty() { - Some(proto.comment[0]) - } else { - None - }, - newlines_in_values: if proto.newlines_in_values.is_empty() { - None - } else { - Some(proto.newlines_in_values[0] != 0) - }, - truncated_rows: if proto.truncated_rows.is_empty() { - None - } else { - Some(proto.truncated_rows[0] != 0) - }, - compression_level: proto.compression_level, - quote_style: match CsvQuoteStyleProto::try_from(proto.quote_style) { - Ok(CsvQuoteStyleProto::Always) => CsvQuoteStyle::Always, - Ok(CsvQuoteStyleProto::NonNumeric) => CsvQuoteStyle::NonNumeric, - Ok(CsvQuoteStyleProto::Never) => CsvQuoteStyle::Never, - Ok(CsvQuoteStyleProto::Necessary) => CsvQuoteStyle::Necessary, - _ => CsvQuoteStyle::Necessary, - }, - ignore_leading_whitespace: if proto.ignore_leading_whitespace.is_empty() { - None - } else { - Some(proto.ignore_leading_whitespace[0] != 0) - }, - ignore_trailing_whitespace: if proto.ignore_trailing_whitespace.is_empty() { - None - } else { - Some(proto.ignore_trailing_whitespace[0] != 0) - }, - } - } -} - // TODO! This is a placeholder for now and needs to be implemented for real. impl LogicalExtensionCodec for CsvLogicalExtensionCodec { fn try_decode( @@ -230,7 +77,7 @@ impl LogicalExtensionCodec for CsvLogicalExtensionCodec { let proto = CsvOptionsProto::decode(buf).map_err(|e| { exec_datafusion_err!("Failed to decode CsvOptionsProto: {e:?}") })?; - let options: CsvOptions = (&proto).into(); + let options = CsvOptions::from(&proto); Ok(Arc::new(CsvFormatFactory { options: Some(options), })) @@ -247,7 +94,7 @@ impl LogicalExtensionCodec for CsvLogicalExtensionCodec { return exec_err!("{}", "Unsupported FileFormatFactory type".to_string()); }; - let proto = CsvOptionsProto::from_factory(&CsvFormatFactory { + let proto = CsvOptionsProto::from(&CsvFormatFactory { options: Some(options), }); @@ -259,38 +106,6 @@ impl LogicalExtensionCodec for CsvLogicalExtensionCodec { } } -impl JsonOptionsProto { - fn from_factory(factory: &JsonFormatFactory) -> Self { - if let Some(options) = &factory.options { - JsonOptionsProto { - compression: options.compression as i32, - schema_infer_max_rec: options.schema_infer_max_rec.map(|v| v as u64), - compression_level: options.compression_level, - newline_delimited: Some(options.newline_delimited), - } - } else { - JsonOptionsProto::default() - } - } -} - -impl From<&JsonOptionsProto> for JsonOptions { - fn from(proto: &JsonOptionsProto) -> Self { - JsonOptions { - compression: match proto.compression { - 0 => CompressionTypeVariant::GZIP, - 1 => CompressionTypeVariant::BZIP2, - 2 => CompressionTypeVariant::XZ, - 3 => CompressionTypeVariant::ZSTD, - _ => CompressionTypeVariant::UNCOMPRESSED, - }, - schema_infer_max_rec: proto.schema_infer_max_rec.map(|v| v as usize), - compression_level: proto.compression_level, - newline_delimited: proto.newline_delimited.unwrap_or(true), - } - } -} - #[derive(Debug)] pub struct JsonLogicalExtensionCodec; @@ -340,7 +155,7 @@ impl LogicalExtensionCodec for JsonLogicalExtensionCodec { let proto = JsonOptionsProto::decode(buf).map_err(|e| { exec_datafusion_err!("Failed to decode JsonOptionsProto: {e:?}") })?; - let options: JsonOptions = (&proto).into(); + let options = JsonOptions::from(&proto); Ok(Arc::new(JsonFormatFactory { options: Some(options), })) @@ -358,7 +173,7 @@ impl LogicalExtensionCodec for JsonLogicalExtensionCodec { return exec_err!("Unsupported FileFormatFactory type"); }; - let proto = JsonOptionsProto::from_factory(&JsonFormatFactory { + let proto = JsonOptionsProto::from(&JsonFormatFactory { options: Some(options), }); @@ -374,270 +189,10 @@ impl LogicalExtensionCodec for JsonLogicalExtensionCodec { mod parquet { use super::*; - use crate::protobuf::{ - CdcOptions as CdcOptionsProto, ParquetColumnOptions as ParquetColumnOptionsProto, - ParquetColumnSpecificOptions, ParquetOptions as ParquetOptionsProto, - TableParquetOptions as TableParquetOptionsProto, parquet_column_options, - parquet_options, - }; - use datafusion_common::config::{ - CdcOptions, ParquetColumnOptions, ParquetOptions, TableParquetOptions, - }; + use crate::protobuf::TableParquetOptions as TableParquetOptionsProto; + use datafusion_common::config::TableParquetOptions; use datafusion_datasource_parquet::file_format::ParquetFormatFactory; - impl TableParquetOptionsProto { - fn from_factory(factory: &ParquetFormatFactory) -> Self { - let global_options = if let Some(ref options) = factory.options { - options.clone() - } else { - return TableParquetOptionsProto::default(); - }; - - let column_specific_options = global_options.column_specific_options; - TableParquetOptionsProto { - global: Some(ParquetOptionsProto { - enable_page_index: global_options.global.enable_page_index, - pruning: global_options.global.pruning, - skip_metadata: global_options.global.skip_metadata, - metadata_size_hint_opt: global_options.global.metadata_size_hint.map(|size| { - parquet_options::MetadataSizeHintOpt::MetadataSizeHint(size as u64) - }), - pushdown_filters: global_options.global.pushdown_filters, - reorder_filters: global_options.global.reorder_filters, - force_filter_selections: global_options.global.force_filter_selections, - data_pagesize_limit: global_options.global.data_pagesize_limit as u64, - write_batch_size: global_options.global.write_batch_size as u64, - writer_version: global_options.global.writer_version.to_string(), - compression_opt: global_options.global.compression.map(|compression| { - parquet_options::CompressionOpt::Compression(compression) - }), - dictionary_enabled_opt: global_options.global.dictionary_enabled.map(|enabled| { - parquet_options::DictionaryEnabledOpt::DictionaryEnabled(enabled) - }), - dictionary_page_size_limit: global_options.global.dictionary_page_size_limit as u64, - statistics_enabled_opt: global_options.global.statistics_enabled.map(|enabled| { - parquet_options::StatisticsEnabledOpt::StatisticsEnabled(enabled) - }), - max_row_group_size: global_options.global.max_row_group_size as u64, - created_by: global_options.global.created_by.clone(), - column_index_truncate_length_opt: global_options.global.column_index_truncate_length.map(|length| { - parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length as u64) - }), - statistics_truncate_length_opt: global_options.global.statistics_truncate_length.map(|length| { - parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(length as u64) - }), - data_page_row_count_limit: global_options.global.data_page_row_count_limit as u64, - encoding_opt: global_options.global.encoding.map(|encoding| { - parquet_options::EncodingOpt::Encoding(encoding) - }), - bloom_filter_on_read: global_options.global.bloom_filter_on_read, - bloom_filter_on_write: global_options.global.bloom_filter_on_write, - bloom_filter_fpp_opt: global_options.global.bloom_filter_fpp.map(|fpp| { - parquet_options::BloomFilterFppOpt::BloomFilterFpp(fpp) - }), - bloom_filter_ndv_opt: global_options.global.bloom_filter_ndv.map(|ndv| { - parquet_options::BloomFilterNdvOpt::BloomFilterNdv(ndv) - }), - allow_single_file_parallelism: global_options.global.allow_single_file_parallelism, - maximum_parallel_row_group_writers: global_options.global.maximum_parallel_row_group_writers as u64, - maximum_buffered_record_batches_per_stream: global_options.global.maximum_buffered_record_batches_per_stream as u64, - schema_force_view_types: global_options.global.schema_force_view_types, - binary_as_string: global_options.global.binary_as_string, - skip_arrow_metadata: global_options.global.skip_arrow_metadata, - coerce_int96_opt: global_options.global.coerce_int96.map(|compression| { - parquet_options::CoerceInt96Opt::CoerceInt96(compression) - }), - coerce_int96_tz_opt: global_options.global.coerce_int96_tz.map(|tz| { - parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(tz) - }), - max_predicate_cache_size_opt: global_options.global.max_predicate_cache_size.map(|size| { - parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(size as u64) - }), - content_defined_chunking: global_options.global.use_content_defined_chunking.as_ref().map(|cdc| { - CdcOptionsProto { - min_chunk_size: cdc.min_chunk_size as u64, - max_chunk_size: cdc.max_chunk_size as u64, - norm_level: cdc.norm_level, - } - }), - }), - column_specific_options: column_specific_options.into_iter().map(|(column_name, options)| { - ParquetColumnSpecificOptions { - column_name, - options: Some(ParquetColumnOptionsProto { - bloom_filter_enabled_opt: options.bloom_filter_enabled.map(|enabled| { - parquet_column_options::BloomFilterEnabledOpt::BloomFilterEnabled(enabled) - }), - encoding_opt: options.encoding.map(|encoding| { - parquet_column_options::EncodingOpt::Encoding(encoding) - }), - dictionary_enabled_opt: options.dictionary_enabled.map(|enabled| { - parquet_column_options::DictionaryEnabledOpt::DictionaryEnabled(enabled) - }), - compression_opt: options.compression.map(|compression| { - parquet_column_options::CompressionOpt::Compression(compression) - }), - statistics_enabled_opt: options.statistics_enabled.map(|enabled| { - parquet_column_options::StatisticsEnabledOpt::StatisticsEnabled(enabled) - }), - bloom_filter_fpp_opt: options.bloom_filter_fpp.map(|fpp| { - parquet_column_options::BloomFilterFppOpt::BloomFilterFpp(fpp) - }), - bloom_filter_ndv_opt: options.bloom_filter_ndv.map(|ndv| { - parquet_column_options::BloomFilterNdvOpt::BloomFilterNdv(ndv) - }), - }) - } - }).collect(), - key_value_metadata: global_options.key_value_metadata - .iter() - .filter_map(|(key, value)| { - value.as_ref().map(|v| (key.clone(), v.clone())) - }) - .collect(), - } - } - } - - impl From<&ParquetOptionsProto> for ParquetOptions { - fn from(proto: &ParquetOptionsProto) -> Self { - ParquetOptions { - enable_page_index: proto.enable_page_index, - pruning: proto.pruning, - skip_metadata: proto.skip_metadata, - metadata_size_hint: proto.metadata_size_hint_opt.as_ref().map(|opt| match opt { - parquet_options::MetadataSizeHintOpt::MetadataSizeHint(size) => *size as usize, - }), - pushdown_filters: proto.pushdown_filters, - reorder_filters: proto.reorder_filters, - force_filter_selections: proto.force_filter_selections, - data_pagesize_limit: proto.data_pagesize_limit as usize, - write_batch_size: proto.write_batch_size as usize, - // TODO: Consider changing to TryFrom to avoid panic on invalid proto data - writer_version: proto.writer_version.parse().expect(" - Invalid parquet writer version in proto, expected '1.0' or '2.0' - "), - compression: proto.compression_opt.as_ref().map(|opt| match opt { - parquet_options::CompressionOpt::Compression(compression) => compression.clone(), - }), - dictionary_enabled: proto.dictionary_enabled_opt.as_ref().map(|opt| match opt { - parquet_options::DictionaryEnabledOpt::DictionaryEnabled(enabled) => *enabled, - }), - dictionary_page_size_limit: proto.dictionary_page_size_limit as usize, - statistics_enabled: proto.statistics_enabled_opt.as_ref().map(|opt| match opt { - parquet_options::StatisticsEnabledOpt::StatisticsEnabled(statistics) => statistics.clone(), - }), - max_row_group_size: proto.max_row_group_size as usize, - created_by: proto.created_by.clone(), - column_index_truncate_length: proto.column_index_truncate_length_opt.as_ref().map(|opt| match opt { - parquet_options::ColumnIndexTruncateLengthOpt::ColumnIndexTruncateLength(length) => *length as usize, - }), - statistics_truncate_length: proto.statistics_truncate_length_opt.as_ref().map(|opt| match opt { - parquet_options::StatisticsTruncateLengthOpt::StatisticsTruncateLength(length) => *length as usize, - }), - data_page_row_count_limit: proto.data_page_row_count_limit as usize, - encoding: proto.encoding_opt.as_ref().map(|opt| match opt { - parquet_options::EncodingOpt::Encoding(encoding) => encoding.clone(), - }), - bloom_filter_on_read: proto.bloom_filter_on_read, - bloom_filter_on_write: proto.bloom_filter_on_write, - bloom_filter_fpp: proto.bloom_filter_fpp_opt.as_ref().map(|opt| match opt { - parquet_options::BloomFilterFppOpt::BloomFilterFpp(fpp) => *fpp, - }), - bloom_filter_ndv: proto.bloom_filter_ndv_opt.as_ref().map(|opt| match opt { - parquet_options::BloomFilterNdvOpt::BloomFilterNdv(ndv) => *ndv, - }), - allow_single_file_parallelism: proto.allow_single_file_parallelism, - maximum_parallel_row_group_writers: proto.maximum_parallel_row_group_writers as usize, - maximum_buffered_record_batches_per_stream: proto.maximum_buffered_record_batches_per_stream as usize, - schema_force_view_types: proto.schema_force_view_types, - binary_as_string: proto.binary_as_string, - skip_arrow_metadata: proto.skip_arrow_metadata, - coerce_int96: proto.coerce_int96_opt.as_ref().map(|opt| match opt { - parquet_options::CoerceInt96Opt::CoerceInt96(coerce_int96) => coerce_int96.clone(), - }), - coerce_int96_tz: proto.coerce_int96_tz_opt.as_ref().map(|opt| match opt { - parquet_options::CoerceInt96TzOpt::CoerceInt96Tz(tz) => tz.clone(), - }), - max_predicate_cache_size: proto.max_predicate_cache_size_opt.as_ref().map(|opt| match opt { - parquet_options::MaxPredicateCacheSizeOpt::MaxPredicateCacheSize(size) => *size as usize, - }), - use_content_defined_chunking: proto.content_defined_chunking.map(|cdc| { - let defaults = CdcOptions::default(); - CdcOptions { - // proto3 uses 0 as the wire default for uint64; a zero chunk size is - // invalid, so treat it as "field not set" and fall back to the default. - min_chunk_size: if cdc.min_chunk_size != 0 { cdc.min_chunk_size as usize } else { defaults.min_chunk_size }, - max_chunk_size: if cdc.max_chunk_size != 0 { cdc.max_chunk_size as usize } else { defaults.max_chunk_size }, - // norm_level = 0 is a valid value (and the default), so pass it through directly. - norm_level: cdc.norm_level, - } - }), - } - } - } - - impl From for ParquetColumnOptions { - fn from(proto: ParquetColumnOptionsProto) -> Self { - ParquetColumnOptions { - bloom_filter_enabled: proto.bloom_filter_enabled_opt.map( - |parquet_column_options::BloomFilterEnabledOpt::BloomFilterEnabled(v)| v, - ), - encoding: proto - .encoding_opt - .map(|parquet_column_options::EncodingOpt::Encoding(v)| v), - dictionary_enabled: proto.dictionary_enabled_opt.map( - |parquet_column_options::DictionaryEnabledOpt::DictionaryEnabled(v)| v, - ), - compression: proto - .compression_opt - .map(|parquet_column_options::CompressionOpt::Compression(v)| v), - statistics_enabled: proto.statistics_enabled_opt.map( - |parquet_column_options::StatisticsEnabledOpt::StatisticsEnabled(v)| v, - ), - bloom_filter_fpp: proto - .bloom_filter_fpp_opt - .map(|parquet_column_options::BloomFilterFppOpt::BloomFilterFpp(v)| v), - bloom_filter_ndv: proto - .bloom_filter_ndv_opt - .map(|parquet_column_options::BloomFilterNdvOpt::BloomFilterNdv(v)| v), - } - } - } - - impl From<&TableParquetOptionsProto> for TableParquetOptions { - fn from(proto: &TableParquetOptionsProto) -> Self { - TableParquetOptions { - global: proto - .global - .as_ref() - .map(ParquetOptions::from) - .unwrap_or_default(), - column_specific_options: proto - .column_specific_options - .iter() - .map(|parquet_column_options| { - ( - parquet_column_options.column_name.clone(), - ParquetColumnOptions::from( - parquet_column_options - .options - .clone() - .unwrap_or_default(), - ), - ) - }) - .collect(), - key_value_metadata: proto - .key_value_metadata - .iter() - .map(|(k, v)| (k.clone(), Some(v.clone()))) - .collect(), - ..Default::default() - } - } - } - #[derive(Debug)] pub struct ParquetLogicalExtensionCodec; @@ -688,12 +243,10 @@ mod parquet { let proto = TableParquetOptionsProto::decode(buf).map_err(|e| { exec_datafusion_err!("Failed to decode TableParquetOptionsProto: {e:?}") })?; - let options: TableParquetOptions = (&proto).into(); - Ok(Arc::new( - datafusion_datasource_parquet::file_format::ParquetFormatFactory { - options: Some(options), - }, - )) + let options = TableParquetOptions::try_from(&proto)?; + Ok(Arc::new(ParquetFormatFactory { + options: Some(options), + })) } fn try_encode_file_format( @@ -711,7 +264,7 @@ mod parquet { return exec_err!("Unsupported FileFormatFactory type"); }; - let proto = TableParquetOptionsProto::from_factory(&ParquetFormatFactory { + let proto = TableParquetOptionsProto::from(&ParquetFormatFactory { options: Some(options), }); @@ -722,6 +275,66 @@ mod parquet { Ok(()) } } + + #[cfg(test)] + mod tests { + use super::*; + use crate::protobuf::ParquetOptions as ParquetOptionsProto; + use datafusion_common::config::ParquetOptions; + + fn encode_table_options(proto: TableParquetOptionsProto) -> Vec { + let mut buf = Vec::new(); + proto.encode(&mut buf).expect("encode parquet options"); + buf + } + + #[test] + fn try_decode_file_format_errors_on_invalid_writer_version() { + let proto = TableParquetOptionsProto { + global: Some(ParquetOptionsProto { + writer_version: "3.0".to_string(), + ..Default::default() + }), + ..Default::default() + }; + + let result = ParquetLogicalExtensionCodec.try_decode_file_format( + &encode_table_options(proto), + &TaskContext::default(), + ); + + let err = result.expect_err("invalid writer version should error"); + assert!( + err.to_string() + .contains("Invalid parquet writer version: 3.0"), + "{err}" + ); + } + + #[test] + fn try_decode_file_format_defaults_empty_writer_version() { + let proto = TableParquetOptionsProto { + global: Some(ParquetOptionsProto::default()), + ..Default::default() + }; + + let factory = ParquetLogicalExtensionCodec + .try_decode_file_format( + &encode_table_options(proto), + &TaskContext::default(), + ) + .expect("decode parquet options"); + let parquet_factory = factory + .downcast_ref::() + .expect("parquet format factory"); + let options = parquet_factory.options.as_ref().expect("parquet options"); + + assert_eq!( + options.global.writer_version, + ParquetOptions::default().writer_version + ); + } + } } #[cfg(feature = "parquet")] pub use parquet::ParquetLogicalExtensionCodec; diff --git a/datafusion/proto/src/logical_plan/from_proto.rs b/datafusion/proto/src/logical_plan/from_proto.rs index 78ffd362c8e48..d4d0ea7292ffe 100644 --- a/datafusion/proto/src/logical_plan/from_proto.rs +++ b/datafusion/proto/src/logical_plan/from_proto.rs @@ -20,240 +20,138 @@ use std::sync::Arc; use arrow::datatypes::{DataType, Field}; use datafusion_common::datatype::DataTypeExt; use datafusion_common::{ - NullEquality, RecursionUnnestOption, Result, ScalarValue, TableReference, - UnnestOptions, exec_datafusion_err, internal_err, plan_datafusion_err, + Result, ScalarValue, SplitPoint, TableReference, exec_datafusion_err, internal_err, + plan_datafusion_err, }; use datafusion_execution::TaskContext; use datafusion_execution::registry::FunctionRegistry; -use datafusion_expr::dml::InsertOp; -use datafusion_expr::expr::{Alias, NullTreatment, Placeholder, Sort}; +use datafusion_expr::dml::{ + InsertOp, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, +}; +use datafusion_expr::expr::{ + Alias, Lambda, LambdaVariable, NullTreatment, Placeholder, Sort, +}; use datafusion_expr::expr::{Unnest, WildcardOptions}; use datafusion_expr::logical_plan::Subquery; use datafusion_expr::{ Between, BinaryExpr, Case, Cast, Expr, GroupingSet, GroupingSet::GroupingSets, - JoinConstraint, JoinType, Like, Operator, TryCast, WindowFrame, WindowFrameBound, - WindowFrameUnits, + Like, Operator, TryCast, WindowFrame, expr::{self, InList, WindowFunction}, - logical_plan::{PlanType, StringifiedPlan}, }; use datafusion_expr::{ExprFunctionExt, WriteOp}; use datafusion_proto_common::{FromProtoError as Error, from_proto::FromOptionalField}; -use crate::protobuf::plan_type::PlanTypeEnum::{ - FinalPhysicalPlanWithSchema, InitialPhysicalPlanWithSchema, -}; -use crate::protobuf::{ - self, AnalyzedLogicalPlanType, CubeNode, GroupingSetNode, OptimizedLogicalPlanType, - OptimizedPhysicalPlanType, PlaceholderNode, RollupNode, - plan_type::PlanTypeEnum::{ - AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, - FinalPhysicalPlan, FinalPhysicalPlanWithStats, InitialLogicalPlan, - InitialPhysicalPlan, InitialPhysicalPlanWithStats, OptimizedLogicalPlan, - OptimizedPhysicalPlan, PhysicalPlanError, - }, -}; +use crate::protobuf::{self, CubeNode, GroupingSetNode, PlaceholderNode, RollupNode}; use super::{AsLogicalPlan, LogicalExtensionCodec}; -impl From<&protobuf::UnnestOptions> for UnnestOptions { - fn from(opts: &protobuf::UnnestOptions) -> Self { - Self { - preserve_nulls: opts.preserve_nulls, - recursions: opts - .recursions - .iter() - .map(|r| RecursionUnnestOption { - input_column: r.input_column.as_ref().unwrap().into(), - output_column: r.output_column.as_ref().unwrap().into(), - depth: r.depth as usize, - }) - .collect::>(), - } - } -} - -impl From for WindowFrameUnits { - fn from(units: protobuf::WindowFrameUnits) -> Self { - match units { - protobuf::WindowFrameUnits::Rows => Self::Rows, - protobuf::WindowFrameUnits::Range => Self::Range, - protobuf::WindowFrameUnits::Groups => Self::Groups, - } - } -} - -impl TryFrom for TableReference { - type Error = Error; - - fn try_from(value: protobuf::TableReference) -> Result { - use protobuf::table_reference::TableReferenceEnum; - let table_reference_enum = value - .table_reference_enum - .ok_or_else(|| Error::required("table_reference_enum"))?; - - match table_reference_enum { - TableReferenceEnum::Bare(protobuf::BareTableReference { table }) => { - Ok(TableReference::bare(table)) - } - TableReferenceEnum::Partial(protobuf::PartialTableReference { - schema, - table, - }) => Ok(TableReference::partial(schema, table)), - TableReferenceEnum::Full(protobuf::FullTableReference { - catalog, - schema, - table, - }) => Ok(TableReference::full(catalog, schema, table)), - } - } -} - -impl From<&protobuf::StringifiedPlan> for StringifiedPlan { - fn from(stringified_plan: &protobuf::StringifiedPlan) -> Self { - Self { - plan_type: match stringified_plan - .plan_type - .as_ref() - .and_then(|pt| pt.plan_type_enum.as_ref()) - .unwrap_or_else(|| { - panic!( - "Cannot create protobuf::StringifiedPlan from {stringified_plan:?}" - ) - }) { - InitialLogicalPlan(_) => PlanType::InitialLogicalPlan, - AnalyzedLogicalPlan(AnalyzedLogicalPlanType { analyzer_name }) => { - PlanType::AnalyzedLogicalPlan { - analyzer_name:analyzer_name.clone() - } - } - FinalAnalyzedLogicalPlan(_) => PlanType::FinalAnalyzedLogicalPlan, - OptimizedLogicalPlan(OptimizedLogicalPlanType { optimizer_name }) => { - PlanType::OptimizedLogicalPlan { - optimizer_name: optimizer_name.clone(), - } - } - FinalLogicalPlan(_) => PlanType::FinalLogicalPlan, - InitialPhysicalPlan(_) => PlanType::InitialPhysicalPlan, - InitialPhysicalPlanWithStats(_) => PlanType::InitialPhysicalPlanWithStats, - InitialPhysicalPlanWithSchema(_) => PlanType::InitialPhysicalPlanWithSchema, - OptimizedPhysicalPlan(OptimizedPhysicalPlanType { optimizer_name }) => { - PlanType::OptimizedPhysicalPlan { - optimizer_name: optimizer_name.clone(), - } - } - FinalPhysicalPlan(_) => PlanType::FinalPhysicalPlan, - FinalPhysicalPlanWithStats(_) => PlanType::FinalPhysicalPlanWithStats, - FinalPhysicalPlanWithSchema(_) => PlanType::FinalPhysicalPlanWithSchema, - PhysicalPlanError(_) => PlanType::PhysicalPlanError, - }, - plan: Arc::new(stringified_plan.plan.clone()), - } - } -} - -impl TryFrom for WindowFrame { - type Error = Error; - - fn try_from(window: protobuf::WindowFrame) -> Result { - let units = protobuf::WindowFrameUnits::try_from(window.window_frame_units) - .map_err(|_| Error::unknown("WindowFrameUnits", window.window_frame_units))? - .into(); - let start_bound = window.start_bound.required("start_bound")?; - let end_bound = window - .end_bound - .map(|end_bound| match end_bound { - protobuf::window_frame::EndBound::Bound(end_bound) => { - end_bound.try_into() - } - }) - .transpose()? - .unwrap_or(WindowFrameBound::CurrentRow); - Ok(WindowFrame::new_bounds(units, start_bound, end_bound)) - } -} - -impl TryFrom for WindowFrameBound { - type Error = Error; - - fn try_from(bound: protobuf::WindowFrameBound) -> Result { - let bound_type = - protobuf::WindowFrameBoundType::try_from(bound.window_frame_bound_type) - .map_err(|_| { - Error::unknown("WindowFrameBoundType", bound.window_frame_bound_type) - })?; - match bound_type { - protobuf::WindowFrameBoundType::CurrentRow => Ok(Self::CurrentRow), - protobuf::WindowFrameBoundType::Preceding => match bound.bound_value { - Some(x) => Ok(Self::Preceding(ScalarValue::try_from(&x)?)), - None => Ok(Self::Preceding(ScalarValue::UInt64(None))), - }, - protobuf::WindowFrameBoundType::Following => match bound.bound_value { - Some(x) => Ok(Self::Following(ScalarValue::try_from(&x)?)), - None => Ok(Self::Following(ScalarValue::UInt64(None))), - }, - } - } -} - -impl From for JoinType { - fn from(t: protobuf::JoinType) -> Self { - match t { - protobuf::JoinType::Inner => JoinType::Inner, - protobuf::JoinType::Left => JoinType::Left, - protobuf::JoinType::Right => JoinType::Right, - protobuf::JoinType::Full => JoinType::Full, - protobuf::JoinType::Leftsemi => JoinType::LeftSemi, - protobuf::JoinType::Rightsemi => JoinType::RightSemi, - protobuf::JoinType::Leftanti => JoinType::LeftAnti, - protobuf::JoinType::Rightanti => JoinType::RightAnti, - protobuf::JoinType::Leftmark => JoinType::LeftMark, - protobuf::JoinType::Rightmark => JoinType::RightMark, - } - } -} - -impl From for JoinConstraint { - fn from(t: protobuf::JoinConstraint) -> Self { - match t { - protobuf::JoinConstraint::On => JoinConstraint::On, - protobuf::JoinConstraint::Using => JoinConstraint::Using, +/// Reconstruct a [`WriteOp`] from a [`protobuf::DmlNode`], reading the +/// `merge_into` payload when the type tag is `MergeInto`. +pub fn parse_write_op( + node: &protobuf::DmlNode, + ctx: &TaskContext, + codec: &dyn LogicalExtensionCodec, +) -> Result { + let typ = node.dml_type(); + Ok(match typ { + protobuf::dml_node::Type::Update => WriteOp::Update, + protobuf::dml_node::Type::Delete => WriteOp::Delete, + protobuf::dml_node::Type::InsertAppend => WriteOp::Insert(InsertOp::Append), + protobuf::dml_node::Type::InsertOverwrite => WriteOp::Insert(InsertOp::Overwrite), + protobuf::dml_node::Type::InsertReplace => WriteOp::Insert(InsertOp::Replace), + protobuf::dml_node::Type::Ctas => WriteOp::Ctas, + protobuf::dml_node::Type::Truncate => WriteOp::Truncate, + protobuf::dml_node::Type::MergeInto => { + let merge_into = node.merge_into.as_deref().ok_or_else(|| { + Error::General( + "DmlNode with MERGE_INTO type is missing the merge_into payload" + .to_string(), + ) + })?; + WriteOp::MergeInto(Box::new(parse_merge_into_op(merge_into, ctx, codec)?)) } - } + }) } -impl From for NullEquality { - fn from(t: protobuf::NullEquality) -> Self { - match t { - protobuf::NullEquality::NullEqualsNothing => NullEquality::NullEqualsNothing, - protobuf::NullEquality::NullEqualsNull => NullEquality::NullEqualsNull, - } - } +fn parse_merge_into_op( + op: &protobuf::MergeIntoOpNode, + ctx: &TaskContext, + codec: &dyn LogicalExtensionCodec, +) -> Result { + let on = op.on.as_ref().ok_or_else(|| { + Error::General("MergeIntoOpNode is missing required `on` expression".to_string()) + })?; + let on = parse_expr(on, ctx, codec)?; + let clauses = op + .clauses + .iter() + .map(|c| parse_merge_into_clause(c, ctx, codec)) + .collect::, Error>>()?; + Ok(MergeIntoOp { on, clauses }) } -impl From for WriteOp { - fn from(t: protobuf::dml_node::Type) -> Self { - match t { - protobuf::dml_node::Type::Update => WriteOp::Update, - protobuf::dml_node::Type::Delete => WriteOp::Delete, - protobuf::dml_node::Type::InsertAppend => WriteOp::Insert(InsertOp::Append), - protobuf::dml_node::Type::InsertOverwrite => { - WriteOp::Insert(InsertOp::Overwrite) - } - protobuf::dml_node::Type::InsertReplace => WriteOp::Insert(InsertOp::Replace), - protobuf::dml_node::Type::Ctas => WriteOp::Ctas, - protobuf::dml_node::Type::Truncate => WriteOp::Truncate, - } - } +fn parse_merge_into_clause( + clause: &protobuf::MergeIntoClauseNode, + ctx: &TaskContext, + codec: &dyn LogicalExtensionCodec, +) -> Result { + let kind = protobuf::merge_into_clause_node::Kind::try_from(clause.kind) + .map_err(|_| { + Error::General(format!( + "MergeIntoClauseNode has unknown kind tag {}", + clause.kind + )) + }) + .map(MergeIntoClauseKind::from)?; + let predicate = clause + .predicate + .as_ref() + .map(|e| parse_expr(e, ctx, codec)) + .transpose()?; + let action = clause.action.as_ref().ok_or_else(|| { + Error::General("MergeIntoClauseNode is missing required `action`".to_string()) + })?; + let action = parse_merge_into_action(action, ctx, codec)?; + Ok(MergeIntoClause { + kind, + predicate, + action, + }) } -impl From for NullTreatment { - fn from(t: protobuf::NullTreatment) -> Self { - match t { - protobuf::NullTreatment::RespectNulls => NullTreatment::RespectNulls, - protobuf::NullTreatment::IgnoreNulls => NullTreatment::IgnoreNulls, +fn parse_merge_into_action( + action: &protobuf::MergeIntoActionNode, + ctx: &TaskContext, + codec: &dyn LogicalExtensionCodec, +) -> Result { + use protobuf::merge_into_action_node::Action; + let action = action.action.as_ref().ok_or_else(|| { + Error::General("MergeIntoActionNode is missing the `action` oneof".to_string()) + })?; + Ok(match action { + Action::Update(update) => { + let assignments = update + .assignments + .iter() + .map(|a| { + let value = a.value.as_ref().ok_or_else(|| { + Error::General(format!( + "MergeAssignment for column `{}` is missing its value", + a.column + )) + })?; + Ok((a.column.clone(), parse_expr(value, ctx, codec)?)) + }) + .collect::, Error>>()?; + MergeIntoAction::Update(assignments) } - } + Action::Insert(insert) => MergeIntoAction::Insert { + columns: insert.columns.clone(), + values: parse_exprs(&insert.values, ctx, codec)?, + }, + Action::Delete(_) => MergeIntoAction::Delete, + }) } pub fn parse_expr( @@ -304,7 +202,7 @@ pub fn parse_expr( .window_frame .as_ref() .map::, _>(|window_frame| { - let window_frame: WindowFrame = window_frame.clone().try_into()?; + let window_frame = WindowFrame::try_from(window_frame.clone())?; window_frame .regularize_order_bys(&mut order_by) .map(|_| window_frame) @@ -558,7 +456,10 @@ pub fn parse_expr( if exprs.len() != 1 { return Err(proto_error("Unnest must have exactly one expression")); } - Ok(Expr::Unnest(Unnest::new(exprs.swap_remove(0)))) + Ok(Expr::Unnest(Unnest { + expr: Box::new(exprs.swap_remove(0)), + outer: unnest.outer, + })) } ExprType::InList(in_list) => Ok(Expr::InList(InList::new( Box::new(parse_required_expr( @@ -571,7 +472,10 @@ pub fn parse_expr( in_list.negated, ))), ExprType::Wildcard(protobuf::Wildcard { qualifier }) => { - let qualifier = qualifier.to_owned().map(|x| x.try_into()).transpose()?; + let qualifier = qualifier + .to_owned() + .map(TableReference::try_from) + .transpose()?; #[expect(deprecated)] Ok(Expr::Wildcard { qualifier, @@ -594,6 +498,22 @@ pub fn parse_expr( parse_exprs(args, ctx, codec)?, ))) } + ExprType::HigherOrderUdfExpr(protobuf::HigherOrderUdfExprNode { + fun_name, + args, + fun_definition, + }) => { + let hof_fn = match fun_definition { + Some(buf) => codec.try_decode_higher_order_function(fun_name, buf)?, + None => ctx + .higher_order_function(fun_name.as_str()) + .or_else(|_| codec.try_decode_higher_order_function(fun_name, &[]))?, + }; + Ok(Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + hof_fn, + parse_exprs(args, ctx, codec)?, + ))) + } ExprType::AggregateUdfExpr(pb) => { let agg_fn = match &pb.fun_definition { Some(buf) => codec.try_decode_udaf(&pb.fun_name, buf)?, @@ -667,6 +587,16 @@ pub fn parse_expr( )?; Ok(Expr::ScalarSubquery(subquery)) } + ExprType::Lambda(lambda) => Ok(Expr::Lambda(Lambda::new( + lambda.params.clone(), + parse_required_expr(lambda.body.as_deref(), ctx, "body", codec)?, + ))), + ExprType::LambdaVariable(lambda_variable) => { + Ok(Expr::LambdaVariable(LambdaVariable::new( + lambda_variable.name.clone(), + lambda_variable.field.as_ref().optional()?.map(Arc::new), + ))) + } } } @@ -742,42 +672,11 @@ fn parse_escape_char(s: &str) -> Result> { } pub fn from_proto_binary_op(op: &str) -> Result { - match op { - "And" => Ok(Operator::And), - "Or" => Ok(Operator::Or), - "Eq" => Ok(Operator::Eq), - "NotEq" => Ok(Operator::NotEq), - "LtEq" => Ok(Operator::LtEq), - "Lt" => Ok(Operator::Lt), - "Gt" => Ok(Operator::Gt), - "GtEq" => Ok(Operator::GtEq), - "Plus" => Ok(Operator::Plus), - "Minus" => Ok(Operator::Minus), - "Multiply" => Ok(Operator::Multiply), - "Divide" => Ok(Operator::Divide), - "Modulo" => Ok(Operator::Modulo), - "IsDistinctFrom" => Ok(Operator::IsDistinctFrom), - "IsNotDistinctFrom" => Ok(Operator::IsNotDistinctFrom), - "BitwiseAnd" => Ok(Operator::BitwiseAnd), - "BitwiseOr" => Ok(Operator::BitwiseOr), - "BitwiseXor" => Ok(Operator::BitwiseXor), - "BitwiseShiftLeft" => Ok(Operator::BitwiseShiftLeft), - "BitwiseShiftRight" => Ok(Operator::BitwiseShiftRight), - "RegexIMatch" => Ok(Operator::RegexIMatch), - "RegexMatch" => Ok(Operator::RegexMatch), - "RegexNotIMatch" => Ok(Operator::RegexNotIMatch), - "RegexNotMatch" => Ok(Operator::RegexNotMatch), - "LikeMatch" => Ok(Operator::LikeMatch), - "ILikeMatch" => Ok(Operator::ILikeMatch), - "NotLikeMatch" => Ok(Operator::NotLikeMatch), - "NotILikeMatch" => Ok(Operator::NotILikeMatch), - "StringConcat" => Ok(Operator::StringConcat), - "AtArrow" => Ok(Operator::AtArrow), - "ArrowAt" => Ok(Operator::ArrowAt), - other => Err(proto_error(format!( - "Unsupported binary operator '{other:?}'" - ))), - } + // The proto-string <-> `Operator` mapping is canonically owned by + // `datafusion-expr-common` so `datafusion-proto` (logical plans) and + // `PhysicalExpr` decoders (e.g. `BinaryExpr`) share one source of truth. + Operator::from_proto_name(op) + .ok_or_else(|| proto_error(format!("Unsupported binary operator '{op:?}'"))) } fn parse_optional_expr( @@ -806,3 +705,14 @@ fn parse_required_expr( fn proto_error>(message: S) -> Error { Error::General(message.into()) } + +pub(super) fn parse_protobuf_range_split_point( + split_point: &protobuf::RangeSplitPoint, +) -> Result { + let values = split_point + .value + .iter() + .map(ScalarValue::try_from) + .collect::>()?; + Ok(SplitPoint::new(values)) +} diff --git a/datafusion/proto/src/logical_plan/mod.rs b/datafusion/proto/src/logical_plan/mod.rs index 8228e8e6f2ff0..2efff1cbde793 100644 --- a/datafusion/proto/src/logical_plan/mod.rs +++ b/datafusion/proto/src/logical_plan/mod.rs @@ -25,7 +25,7 @@ use crate::protobuf::{ CustomTableScanNode, DmlNode, SortExprNodeCollection, dml_node, }; use crate::{ - convert_required, into_required, + convert_required, protobuf::{ self, LogicalExtensionNode, LogicalPlanNode, listing_table_scan_node::FileFormatType, logical_plan_node::LogicalPlanType, @@ -37,9 +37,11 @@ use arrow::datatypes::{DataType, Field, Schema, SchemaBuilder, SchemaRef}; use datafusion_catalog::cte_worktable::CteWorkTable; use datafusion_catalog::empty::EmptyTable; use datafusion_common::file_options::file_type::FileType; -use datafusion_common::format::ExplainFormat; +use datafusion_common::format::{ + ExplainAnalyzeCategories, ExplainFormat, MetricCategory, MetricType, +}; use datafusion_common::{ - Result, TableReference, ToDFSchema, assert_or_internal_err, context, + NullEquality, Result, TableReference, assert_or_internal_err, context, internal_datafusion_err, internal_err, not_impl_err, plan_err, }; use datafusion_datasource::file_format::FileFormat; @@ -55,22 +57,25 @@ use datafusion_datasource_json::file_format::{ }; #[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::{ParquetFormat, ParquetFormatFactory}; +use datafusion_expr::dml::InsertOp; use datafusion_expr::{ - AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RecursiveQuery, SkipType, - TableSource, Unnest, + AggregateUDF, DmlStatement, FetchType, HigherOrderUDF, RangePartitioning, + RecursiveQuery, SkipType, TableSource, Unnest, WriteOp, }; use datafusion_expr::{ - DistinctOn, DropView, Expr, LogicalPlan, LogicalPlanBuilder, ScalarUDF, SortExpr, - Statement, WindowUDF, dml, + DistinctOn, DropView, Expr, JoinConstraint, LogicalPlan, LogicalPlanBuilder, + ScalarUDF, SortExpr, Statement, WindowUDF, dml, logical_plan::{ Aggregate, CreateCatalog, CreateCatalogSchema, CreateExternalTable, CreateView, - DdlStatement, Distinct, EmptyRelation, Extension, Join, JoinConstraint, Prepare, - Projection, Repartition, Sort, SubqueryAlias, TableScan, Values, Window, + DdlStatement, Distinct, EmptyRelation, Extension, Join, Prepare, Projection, + Repartition, Sort, SubqueryAlias, TableScan, TableScanBuilder, Values, Window, builder::project, }, }; +use datafusion_proto_common::protobuf_common; use self::to_proto::{serialize_expr, serialize_exprs}; +use crate::logical_plan::to_proto::serialize_range_split_point; use crate::logical_plan::to_proto::serialize_sorts; use datafusion_catalog::TableProvider; use datafusion_catalog::default_table_source::{provider_as_source, source_as_provider}; @@ -109,6 +114,29 @@ pub trait AsLogicalPlan: Debug + Send + Sync + Clone { Self: Sized; } +// In debug builds, keep each serializer arm's local temporaries out of the +// recursive dispatcher frame. Without this call boundary, they inflate the +// frame of every recursive invocation. +#[cfg_attr(debug_assertions, inline(never))] +fn serialize_logical_plan_arm(serializer: F) -> Result +where + F: FnOnce() -> Result, +{ + serializer() +} + +macro_rules! dispatch_logical_plan { + ($plan:expr, { $($pattern:pat => $body:expr $(,)?)+ }) => { + match $plan { + $( + $pattern => serialize_logical_plan_arm(|| -> Result { + $body + }), + )+ + } + }; +} + pub trait LogicalExtensionCodec: Debug + Send + Sync + std::any::Any { fn try_decode( &self, @@ -162,7 +190,7 @@ pub trait LogicalExtensionCodec: Debug + Send + Sync + std::any::Any { &self, name: &str, _buf: &[u8], - ) -> Result> { + ) -> Result> { not_impl_err!( "LogicalExtensionCodec is not provided for higher order function {name}" ) @@ -170,7 +198,7 @@ pub trait LogicalExtensionCodec: Debug + Send + Sync + std::any::Any { fn try_encode_higher_order_function( &self, - _node: &dyn HigherOrderUDF, + _node: &HigherOrderUDF, _buf: &mut Vec, ) -> Result<()> { Ok(()) @@ -342,7 +370,7 @@ fn from_table_reference( ) })?; - Ok(table_ref.clone().try_into()?) + Ok(TableReference::try_from(table_ref.clone())?) } /// Converts [LogicalPlan::TableScan] to [TableSource] @@ -371,19 +399,85 @@ fn from_table_source( target: Arc, extension_codec: &dyn LogicalExtensionCodec, ) -> Result { - let projected_schema = target.schema().to_dfschema_ref()?; - let r = LogicalPlan::TableScan(TableScan { - table_name, - source: target, - projection: None, - projected_schema, - filters: vec![], - fetch: None, - }); + let r = LogicalPlan::TableScan(TableScanBuilder::new(table_name, target).build()?); LogicalPlanNode::try_from_logical_plan(&r, extension_codec) } +fn metric_type_from_proto(value: i32) -> Result { + let pb = protobuf_common::MetricType::try_from(value) + .map_err(|_| proto_error(format!("Unknown MetricType discriminant: {value}")))?; + Ok(match pb { + protobuf_common::MetricType::Summary => MetricType::Summary, + protobuf_common::MetricType::Dev => MetricType::Dev, + }) +} + +fn metric_type_to_proto(value: MetricType) -> protobuf_common::MetricType { + match value { + MetricType::Summary => protobuf_common::MetricType::Summary, + MetricType::Dev => protobuf_common::MetricType::Dev, + } +} + +fn metric_category_from_proto(value: i32) -> Result { + let pb = protobuf_common::MetricCategory::try_from(value).map_err(|_| { + proto_error(format!("Unknown MetricCategory discriminant: {value}")) + })?; + Ok(match pb { + protobuf_common::MetricCategory::Rows => MetricCategory::Rows, + protobuf_common::MetricCategory::Bytes => MetricCategory::Bytes, + protobuf_common::MetricCategory::Timing => MetricCategory::Timing, + protobuf_common::MetricCategory::Uncategorized => MetricCategory::Uncategorized, + }) +} + +fn metric_category_to_proto(value: MetricCategory) -> protobuf_common::MetricCategory { + match value { + MetricCategory::Rows => protobuf_common::MetricCategory::Rows, + MetricCategory::Bytes => protobuf_common::MetricCategory::Bytes, + MetricCategory::Timing => protobuf_common::MetricCategory::Timing, + MetricCategory::Uncategorized => protobuf_common::MetricCategory::Uncategorized, + } +} + +fn explain_analyze_categories_from_proto( + node: &protobuf_common::ExplainAnalyzeCategoriesNode, +) -> Result { + if node.all { + Ok(ExplainAnalyzeCategories::All) + } else { + let cats = node + .only + .iter() + .copied() + .map(metric_category_from_proto) + .collect::>>()?; + Ok(ExplainAnalyzeCategories::Only(cats)) + } +} + +fn explain_analyze_categories_to_proto( + value: &ExplainAnalyzeCategories, +) -> protobuf_common::ExplainAnalyzeCategoriesNode { + match value { + ExplainAnalyzeCategories::All => protobuf_common::ExplainAnalyzeCategoriesNode { + all: true, + only: vec![], + }, + ExplainAnalyzeCategories::Only(cats) => { + protobuf_common::ExplainAnalyzeCategoriesNode { + all: false, + only: cats + .iter() + .copied() + .map(|c| metric_category_to_proto(c) as i32) + .collect(), + } + } + } +} + impl AsLogicalPlan for LogicalPlanNode { fn try_decode(buf: &[u8]) -> Result where @@ -576,8 +670,6 @@ impl AsLogicalPlan for LogicalPlanNode { let options = ListingOptions::new(file_format) .with_file_extension(&scan.file_extension) .with_table_partition_cols(partition_columns) - .with_collect_stat(scan.collect_stat) - .with_target_partitions(scan.target_partitions as usize) .with_file_sort_order(all_sort_orders); let config = @@ -610,7 +702,7 @@ impl AsLogicalPlan for LogicalPlanNode { )? .build() } - LogicalPlanType::CustomScan(scan) => { + CustomScan(scan) => { let schema: Schema = convert_required!(scan.schema)?; let schema = Arc::new(schema); let mut projection = None; @@ -676,6 +768,16 @@ impl AsLogicalPlan for LogicalPlanNode { PartitionMethod::RoundRobin(partition_count) => { Partitioning::RoundRobinBatch(*partition_count as usize) } + PartitionMethod::Range(protobuf::RangeRepartition { + sort_expr: pb_sort_expr, + split_point, + }) => Partitioning::Range(RangePartitioning::try_new( + from_proto::parse_sorts(pb_sort_expr, ctx, extension_codec)?, + split_point + .iter() + .map(from_proto::parse_protobuf_range_split_point) + .collect::, _>>()?, + )?), }; LogicalPlanBuilder::from(input) @@ -719,27 +821,43 @@ impl AsLogicalPlan for LogicalPlanNode { column_defaults.insert(col_name.clone(), expr); } + let locations = if !create_extern_table.locations.is_empty() { + create_extern_table.locations.clone() + } else if !create_extern_table.location.is_empty() { + vec![create_extern_table.location.clone()] + } else { + return Err(proto_error( + "CreateExternalTableNode requires at least one location", + )); + }; + let location = locations[0].clone(); + Ok(LogicalPlan::Ddl(DdlStatement::CreateExternalTable( - CreateExternalTable::builder( - from_table_reference( - create_extern_table.name.as_ref(), - "CreateExternalTable", - )?, - create_extern_table.location.clone(), - create_extern_table.file_type.clone(), - pb_schema.try_into()?, - ) - .with_partition_cols(create_extern_table.table_partition_cols.clone()) - .with_order_exprs(order_exprs) - .with_if_not_exists(create_extern_table.if_not_exists) - .with_or_replace(create_extern_table.or_replace) - .with_temporary(create_extern_table.temporary) - .with_definition(definition) - .with_unbounded(create_extern_table.unbounded) - .with_options(create_extern_table.options.clone()) - .with_constraints(constraints.into()) - .with_column_defaults(column_defaults) - .build(), + Box::new( + CreateExternalTable::builder( + from_table_reference( + create_extern_table.name.as_ref(), + "CreateExternalTable", + )?, + location, + create_extern_table.file_type.clone(), + pb_schema.try_into()?, + ) + .with_locations(locations) + .with_partition_cols( + create_extern_table.table_partition_cols.clone(), + ) + .with_order_exprs(order_exprs) + .with_if_not_exists(create_extern_table.if_not_exists) + .with_or_replace(create_extern_table.or_replace) + .with_temporary(create_extern_table.temporary) + .with_definition(definition) + .with_unbounded(create_extern_table.unbounded) + .with_options(create_extern_table.options.clone()) + .with_constraints(constraints.into()) + .with_column_defaults(column_defaults) + .build(), + ), ))) } LogicalPlanType::CreateView(create_view) => { @@ -795,8 +913,37 @@ impl AsLogicalPlan for LogicalPlanNode { LogicalPlanType::Analyze(analyze) => { let input: LogicalPlan = into_logical_plan!(analyze.input, ctx, extension_codec)?; + let analyze_level = analyze + .analyze_level + .map(metric_type_from_proto) + .transpose()?; + let analyze_categories = analyze + .analyze_categories + .as_ref() + .map(explain_analyze_categories_from_proto) + .transpose()?; + let pb_format = protobuf::ExplainFormat::try_from(analyze.format) + .map_err(|_| { + proto_error(format!( + "Received an AnalyzeNode message with unknown ExplainFormat {}", + analyze.format + )) + })?; + let analyze_format = match pb_format { + protobuf::ExplainFormat::Indent => ExplainFormat::Indent, + protobuf::ExplainFormat::Tree => ExplainFormat::Tree, + protobuf::ExplainFormat::Pgjson => ExplainFormat::PostgresJSON, + protobuf::ExplainFormat::Graphviz => ExplainFormat::Graphviz, + }; + let explain_option = + datafusion_expr::logical_plan::ExplainOption::default() + .with_verbose(analyze.verbose) + .with_analyze(true) + .with_analyze_level(analyze_level) + .with_analyze_categories(analyze_categories) + .with_format(analyze_format); LogicalPlanBuilder::from(input) - .explain(analyze.verbose, true)? + .explain_option_format(explain_option)? .build() } LogicalPlanType::Explain(explain) => { @@ -818,7 +965,8 @@ impl AsLogicalPlan for LogicalPlanNode { let explain_option = datafusion_expr::logical_plan::ExplainOption::default() .with_verbose(explain.verbose) - .with_format(explain_format); + .with_format(explain_format) + .with_show_statistics(explain.show_statistics); LogicalPlanBuilder::from(input) .explain_option_format(explain_option)? .build() @@ -850,6 +998,13 @@ impl AsLogicalPlan for LogicalPlanNode { from_proto::parse_exprs(&join.left_join_key, ctx, extension_codec)?; let right_keys: Vec = from_proto::parse_exprs(&join.right_join_key, ctx, extension_codec)?; + if left_keys.len() != right_keys.len() { + return Err(proto_error(format!( + "Received a JoinNode message with left_join_key and right_join_key of different lengths: {} and {}", + left_keys.len(), + right_keys.len() + ))); + } let join_type = protobuf::JoinType::try_from(join.join_type).map_err(|_| { proto_error(format!( @@ -866,44 +1021,39 @@ impl AsLogicalPlan for LogicalPlanNode { join.join_constraint )) })?; + let null_equality = protobuf::NullEquality::try_from(join.null_equality) + .map_err(|_| { + proto_error(format!( + "Received a JoinNode message with unknown NullEquality {}", + join.null_equality + )) + })?; let filter: Option = join .filter .as_ref() .map(|expr| from_proto::parse_expr(expr, ctx, extension_codec)) .map_or(Ok(None), |v| v.map(Some))?; - - let builder = LogicalPlanBuilder::from(into_logical_plan!( - join.left, - ctx, - extension_codec - )?); - let builder = match join_constraint.into() { - JoinConstraint::On => builder.join_with_expr_keys( - into_logical_plan!(join.right, ctx, extension_codec)?, - join_type.into(), - (left_keys, right_keys), - filter, - )?, - JoinConstraint::Using => { - // The equijoin keys in using-join must be column. - let using_keys = left_keys - .into_iter() - .map(|key| { - key.try_as_col().cloned() - .ok_or_else(|| internal_datafusion_err!( - "Using join keys must be column references, got: {key:?}" - )) - }) - .collect::, _>>()?; - builder.join_using( - into_logical_plan!(join.right, ctx, extension_codec)?, - join_type.into(), - using_keys, - )? - } - }; - - builder.build() + let left = into_logical_plan!(join.left, ctx, extension_codec)?; + let right = into_logical_plan!(join.right, ctx, extension_codec)?; + let on: Vec<(Expr, Expr)> = + left_keys.into_iter().zip(right_keys).collect(); + + // Construct the Join directly instead of going through + // LogicalPlanBuilder. The builder methods hardcode + // `null_equality` and `null_aware`, so a round trip through + // them silently loses both fields. Both sides of the round + // trip should already have validated keys, so we don't need + // the builder's normalization / equijoin-pair checks. + Ok(LogicalPlan::Join(Join::try_new( + Arc::new(left), + Arc::new(right), + on, + filter, + datafusion_expr::JoinType::from(join_type), + JoinConstraint::from(join_constraint), + NullEquality::from(null_equality), + join.null_aware, + )?)) } LogicalPlanType::Union(union) => { assert_or_internal_err!( @@ -1064,7 +1214,13 @@ impl AsLogicalPlan for LogicalPlanNode { LogicalPlanBuilder::from(input) .unnest_columns_with_options( unnest.exec_columns.iter().map(|c| c.into()).collect(), - into_required!(unnest.options)?, + unnest + .options + .as_ref() + .map(datafusion_common::UnnestOptions::from) + .ok_or_else(|| { + proto_error("Missing required field in protobuf") + })?, )? .build() } @@ -1085,12 +1241,15 @@ impl AsLogicalPlan for LogicalPlanNode { ))? .try_into_logical_plan(ctx, extension_codec)?; - Ok(LogicalPlan::RecursiveQuery(RecursiveQuery { - name: recursive_query_node.name.clone(), - static_term: Arc::new(static_term), - recursive_term: Arc::new(recursive_term), - is_distinct: recursive_query_node.is_distinct, - })) + // The output schema is derived state, so decoding goes through + // the constructor after restoring the child terms. + RecursiveQuery::try_new( + recursive_query_node.name.clone(), + Arc::new(static_term), + Arc::new(recursive_term), + recursive_query_node.is_distinct, + ) + .map(LogicalPlan::RecursiveQuery) } LogicalPlanType::CteWorkTableScan(cte_work_table_scan_node) => { let CteWorkTableScanNode { name, schema } = cte_work_table_scan_node; @@ -1133,16 +1292,22 @@ impl AsLogicalPlan for LogicalPlanNode { .build() } LogicalPlanType::Dml(dml_node) => { - Ok(LogicalPlan::Dml(datafusion_expr::DmlStatement::new( - from_table_reference(dml_node.table_name.as_ref(), "DML ")?, - to_table_source(&dml_node.target, ctx, extension_codec)?, - dml_node.dml_type().into(), + let table_name = + from_table_reference(dml_node.table_name.as_ref(), "DML ")?; + let target = to_table_source(&dml_node.target, ctx, extension_codec)?; + let write_op = + from_proto::parse_write_op(dml_node, ctx, extension_codec)?; + Ok(LogicalPlan::Dml(DmlStatement::new( + table_name, + target, + write_op, Arc::new(into_logical_plan!(dml_node.input, ctx, extension_codec)?), ))) } } } + #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn try_from_logical_plan( plan: &LogicalPlan, extension_codec: &dyn LogicalExtensionCodec, @@ -1150,7 +1315,7 @@ impl AsLogicalPlan for LogicalPlanNode { where Self: Sized, { - match plan { + dispatch_logical_plan!(plan, { LogicalPlan::Values(Values { values, .. }) => { let n_cols = if values.is_empty() { 0 @@ -1291,8 +1456,9 @@ impl AsLogicalPlan for LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::ListingScan( protobuf::ListingTableScanNode { file_format_type: Some(file_format_type), - table_name: Some(table_name.clone().into()), - collect_stat: options.collect_stat, + table_name: Some(protobuf::TableReference::from( + table_name.clone(), + )), file_extension: options.file_extension.clone(), table_partition_cols: partition_columns, paths: listing_table @@ -1303,7 +1469,6 @@ impl AsLogicalPlan for LogicalPlanNode { schema: Some(schema), projection, filters, - target_partitions: options.target_partitions as u32, file_sort_order: exprs_vec, }, )), @@ -1313,7 +1478,9 @@ impl AsLogicalPlan for LogicalPlanNode { Ok(LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::ViewScan(Box::new( protobuf::ViewTableScanNode { - table_name: Some(table_name.clone().into()), + table_name: Some(protobuf::TableReference::from( + table_name.clone(), + )), input: Some(Box::new( LogicalPlanNode::try_from_logical_plan( view_table.logical_plan(), @@ -1338,7 +1505,7 @@ impl AsLogicalPlan for LogicalPlanNode { Ok(LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::CteWorkTableScan( - protobuf::CteWorkTableScanNode { + CteWorkTableScanNode { name, schema: Some(schema), }, @@ -1350,7 +1517,9 @@ impl AsLogicalPlan for LogicalPlanNode { Ok(LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::EmptyTableScan( protobuf::EmptyTableScanNode { - table_name: Some(table_name.clone().into()), + table_name: Some(protobuf::TableReference::from( + table_name.clone(), + )), schema: Some(schema), projection, filters, @@ -1364,7 +1533,9 @@ impl AsLogicalPlan for LogicalPlanNode { .try_encode_table_provider(table_name, provider, &mut bytes) .map_err(|e| context!("Error serializing custom table", e))?; let scan = CustomScan(CustomTableScanNode { - table_name: Some(table_name.clone().into()), + table_name: Some(protobuf::TableReference::from( + table_name.clone(), + )), projection, schema: Some(schema), filters, @@ -1492,7 +1663,9 @@ impl AsLogicalPlan for LogicalPlanNode { join_type, join_constraint, null_equality, - .. + null_aware, + // Not encoded; recomputed by `Join::try_new` on decode. + schema: _, }) => { let left: LogicalPlanNode = LogicalPlanNode::try_from_logical_plan( left.as_ref(), @@ -1513,11 +1686,11 @@ impl AsLogicalPlan for LogicalPlanNode { .collect::, ToProtoError>>()? .into_iter() .unzip(); - let join_type: protobuf::JoinType = join_type.to_owned().into(); - let join_constraint: protobuf::JoinConstraint = - join_constraint.to_owned().into(); - let null_equality: protobuf::NullEquality = - null_equality.to_owned().into(); + let join_type = protobuf::JoinType::from(join_type.to_owned()); + let join_constraint = + protobuf::JoinConstraint::from(join_constraint.to_owned()); + let null_equality = + protobuf::NullEquality::from(null_equality.to_owned()); let filter = filter .as_ref() .map(|e| serialize_expr(e, extension_codec).map(Box::new)) @@ -1533,6 +1706,7 @@ impl AsLogicalPlan for LogicalPlanNode { right_join_key, null_equality: null_equality.into(), filter, + null_aware: *null_aware, }, ))), }) @@ -1555,7 +1729,7 @@ impl AsLogicalPlan for LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::SubqueryAlias(Box::new( protobuf::SubqueryAliasNode { input: Some(Box::new(input)), - alias: Some((*alias).clone().into()), + alias: Some(protobuf::TableReference::from((*alias).clone())), }, ))), }) @@ -1627,6 +1801,19 @@ impl AsLogicalPlan for LogicalPlanNode { Partitioning::RoundRobinBatch(partition_count) => { PartitionMethod::RoundRobin(*partition_count as u64) } + Partitioning::Range(range_partitioning) => { + let ordering = range_partitioning.ordering(); + let split_point = range_partitioning + .split_points() + .iter() + .map(serialize_range_split_point) + .collect::, _>>()?; + + PartitionMethod::Range(protobuf::RangeRepartition { + sort_expr: serialize_sorts(ordering, extension_codec)?, + split_point, + }) + } Partitioning::DistributeBy(_) => { return not_impl_err!("DistributeBy"); } @@ -1650,10 +1837,10 @@ impl AsLogicalPlan for LogicalPlanNode { }, )), }), - LogicalPlan::Ddl(DdlStatement::CreateExternalTable( - CreateExternalTable { + LogicalPlan::Ddl(DdlStatement::CreateExternalTable(ce)) => { + let CreateExternalTable { name, - location, + locations, file_type, schema: df_schema, table_partition_cols, @@ -1666,8 +1853,7 @@ impl AsLogicalPlan for LogicalPlanNode { constraints, column_defaults, temporary, - }, - )) => { + } = ce.as_ref(); let mut converted_order_exprs: Vec = vec![]; for order in order_exprs { let temp = SortExprNodeCollection { @@ -1682,12 +1868,17 @@ impl AsLogicalPlan for LogicalPlanNode { converted_column_defaults .insert(col_name.clone(), serialize_expr(expr, extension_codec)?); } + let (legacy_location, proto_locations) = match locations.as_slice() { + [location] => (location.clone(), vec![]), + _ => (String::new(), locations.clone()), + }; Ok(LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::CreateExternalTable( protobuf::CreateExternalTableNode { - name: Some(name.clone().into()), - location: location.clone(), + name: Some(protobuf::TableReference::from(name.clone())), + location: legacy_location, + locations: proto_locations, file_type: file_type.clone(), schema: Some(df_schema.try_into()?), table_partition_cols: table_partition_cols.clone(), @@ -1713,7 +1904,7 @@ impl AsLogicalPlan for LogicalPlanNode { })) => Ok(LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::CreateView(Box::new( protobuf::CreateViewNode { - name: Some(name.clone().into()), + name: Some(protobuf::TableReference::from(name.clone())), input: Some(Box::new(LogicalPlanNode::try_from_logical_plan( input, extension_codec, @@ -1762,6 +1953,23 @@ impl AsLogicalPlan for LogicalPlanNode { protobuf::AnalyzeNode { input: Some(Box::new(input)), verbose: a.verbose, + analyze_level: a + .analyze_level + .map(|m| metric_type_to_proto(m) as i32), + analyze_categories: a + .analyze_categories + .as_ref() + .map(explain_analyze_categories_to_proto), + format: match &a.format { + ExplainFormat::Indent => protobuf::ExplainFormat::Indent, + ExplainFormat::Tree => protobuf::ExplainFormat::Tree, + ExplainFormat::PostgresJSON => { + protobuf::ExplainFormat::Pgjson + } + ExplainFormat::Graphviz => { + protobuf::ExplainFormat::Graphviz + } + } as i32, }, ))), }) @@ -1787,6 +1995,7 @@ impl AsLogicalPlan for LogicalPlanNode { } } .into(), + show_statistics: a.show_statistics, }, ))), }) @@ -1885,7 +2094,7 @@ impl AsLogicalPlan for LogicalPlanNode { .map(|c| *c as u64) .collect(), schema: Some(schema.try_into()?), - options: Some(options.into()), + options: Some(protobuf::UnnestOptions::from(options)), }, ))), }) @@ -1906,7 +2115,7 @@ impl AsLogicalPlan for LogicalPlanNode { })) => Ok(LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::DropView( protobuf::DropViewNode { - name: Some(name.clone().into()), + name: Some(protobuf::TableReference::from(name.clone())), if_exists: *if_exists, schema: Some(schema.try_into()?), }, @@ -1933,7 +2142,33 @@ impl AsLogicalPlan for LogicalPlanNode { }) => { let input = LogicalPlanNode::try_from_logical_plan(input, extension_codec)?; - let dml_type: dml_node::Type = op.into(); + let (dml_type, merge_into) = match op { + WriteOp::Insert(InsertOp::Append) => { + (dml_node::Type::InsertAppend, None) + } + WriteOp::Insert(InsertOp::Overwrite) => { + (dml_node::Type::InsertOverwrite, None) + } + WriteOp::Insert(InsertOp::Replace) => { + (dml_node::Type::InsertReplace, None) + } + WriteOp::Delete => (dml_node::Type::Delete, None), + WriteOp::Update => (dml_node::Type::Update, None), + WriteOp::Ctas => (dml_node::Type::Ctas, None), + WriteOp::Truncate => (dml_node::Type::Truncate, None), + WriteOp::MergeInto(merge_op) => ( + dml_node::Type::MergeInto, + Some(Box::new(to_proto::serialize_merge_into_op( + merge_op, + extension_codec, + )?)), + ), + other => { + return Err(proto_error(format!( + "WriteOp variant has no DmlNode encoding: {other}" + ))); + } + }; Ok(LogicalPlanNode { logical_plan_type: Some(LogicalPlanType::Dml(Box::new(DmlNode { input: Some(Box::new(input)), @@ -1942,8 +2177,11 @@ impl AsLogicalPlan for LogicalPlanNode { Arc::clone(target), extension_codec, )?)), - table_name: Some(table_name.clone().into()), + table_name: Some(protobuf::TableReference::from( + table_name.clone(), + )), dml_type: dml_type.into(), + merge_into, }))), }) } @@ -1995,6 +2233,6 @@ impl AsLogicalPlan for LogicalPlanNode { ))), }) } - } + }) } } diff --git a/datafusion/proto/src/logical_plan/to_proto.rs b/datafusion/proto/src/logical_plan/to_proto.rs index d79107d1d0f2b..16c3468465541 100644 --- a/datafusion/proto/src/logical_plan/to_proto.rs +++ b/datafusion/proto/src/logical_plan/to_proto.rs @@ -19,163 +19,24 @@ //! DataFusion logical plans to be serialized and transmitted between //! processes. -use std::collections::HashMap; - -use datafusion_common::{NullEquality, TableReference, UnnestOptions}; -use datafusion_expr::WriteOp; -use datafusion_expr::dml::InsertOp; +use datafusion_common::SplitPoint; +use datafusion_expr::dml::{MergeIntoAction, MergeIntoClause, MergeIntoOp}; use datafusion_expr::expr::{ - self, AggregateFunctionParams, Alias, Between, BinaryExpr, Cast, GroupingSet, InList, - Like, NullTreatment, Placeholder, ScalarFunction, Unnest, + self, AggregateFunctionParams, Alias, Between, BinaryExpr, Cast, GroupingSet, + HigherOrderFunction, InList, Lambda, LambdaVariable, Like, Placeholder, + ScalarFunction, Unnest, }; use datafusion_expr::logical_plan::Subquery; -use datafusion_expr::{ - Expr, JoinConstraint, JoinType, SortExpr, TryCast, WindowFrame, WindowFrameBound, - WindowFrameUnits, WindowFunctionDefinition, logical_plan::PlanType, - logical_plan::StringifiedPlan, -}; +use datafusion_expr::{Expr, SortExpr, TryCast, WindowFunctionDefinition}; -use crate::protobuf::RecursionUnnestOption; use crate::protobuf::{ - self, AnalyzedLogicalPlanType, CubeNode, EmptyMessage, GroupingSetNode, - LogicalExprList, OptimizedLogicalPlanType, OptimizedPhysicalPlanType, - PlaceholderNode, RollupNode, ToProtoError as Error, - plan_type::PlanTypeEnum::{ - AnalyzedLogicalPlan, FinalAnalyzedLogicalPlan, FinalLogicalPlan, - FinalPhysicalPlan, FinalPhysicalPlanWithSchema, FinalPhysicalPlanWithStats, - InitialLogicalPlan, InitialPhysicalPlan, InitialPhysicalPlanWithSchema, - InitialPhysicalPlanWithStats, OptimizedLogicalPlan, OptimizedPhysicalPlan, - PhysicalPlanError, - }, + self, CubeNode, GroupingSetNode, LogicalExprList, PlaceholderNode, RollupNode, + ToProtoError as Error, }; use super::{AsLogicalPlan, LogicalExtensionCodec}; use crate::protobuf::LogicalPlanNode; -impl From<&UnnestOptions> for protobuf::UnnestOptions { - fn from(opts: &UnnestOptions) -> Self { - Self { - preserve_nulls: opts.preserve_nulls, - recursions: opts - .recursions - .iter() - .map(|r| RecursionUnnestOption { - input_column: Some((&r.input_column).into()), - output_column: Some((&r.output_column).into()), - depth: r.depth as u32, - }) - .collect(), - } - } -} - -impl From<&StringifiedPlan> for protobuf::StringifiedPlan { - fn from(stringified_plan: &StringifiedPlan) -> Self { - Self { - plan_type: match stringified_plan.clone().plan_type { - PlanType::InitialLogicalPlan => Some(protobuf::PlanType { - plan_type_enum: Some(InitialLogicalPlan(EmptyMessage {})), - }), - PlanType::AnalyzedLogicalPlan { analyzer_name } => { - Some(protobuf::PlanType { - plan_type_enum: Some(AnalyzedLogicalPlan( - AnalyzedLogicalPlanType { analyzer_name }, - )), - }) - } - PlanType::FinalAnalyzedLogicalPlan => Some(protobuf::PlanType { - plan_type_enum: Some(FinalAnalyzedLogicalPlan(EmptyMessage {})), - }), - PlanType::OptimizedLogicalPlan { optimizer_name } => { - Some(protobuf::PlanType { - plan_type_enum: Some(OptimizedLogicalPlan( - OptimizedLogicalPlanType { optimizer_name }, - )), - }) - } - PlanType::FinalLogicalPlan => Some(protobuf::PlanType { - plan_type_enum: Some(FinalLogicalPlan(EmptyMessage {})), - }), - PlanType::InitialPhysicalPlan => Some(protobuf::PlanType { - plan_type_enum: Some(InitialPhysicalPlan(EmptyMessage {})), - }), - PlanType::OptimizedPhysicalPlan { optimizer_name } => { - Some(protobuf::PlanType { - plan_type_enum: Some(OptimizedPhysicalPlan( - OptimizedPhysicalPlanType { optimizer_name }, - )), - }) - } - PlanType::FinalPhysicalPlan => Some(protobuf::PlanType { - plan_type_enum: Some(FinalPhysicalPlan(EmptyMessage {})), - }), - PlanType::InitialPhysicalPlanWithStats => Some(protobuf::PlanType { - plan_type_enum: Some(InitialPhysicalPlanWithStats(EmptyMessage {})), - }), - PlanType::InitialPhysicalPlanWithSchema => Some(protobuf::PlanType { - plan_type_enum: Some(InitialPhysicalPlanWithSchema(EmptyMessage {})), - }), - PlanType::FinalPhysicalPlanWithStats => Some(protobuf::PlanType { - plan_type_enum: Some(FinalPhysicalPlanWithStats(EmptyMessage {})), - }), - PlanType::FinalPhysicalPlanWithSchema => Some(protobuf::PlanType { - plan_type_enum: Some(FinalPhysicalPlanWithSchema(EmptyMessage {})), - }), - PlanType::PhysicalPlanError => Some(protobuf::PlanType { - plan_type_enum: Some(PhysicalPlanError(EmptyMessage {})), - }), - }, - plan: stringified_plan.plan.to_string(), - } - } -} - -impl From for protobuf::WindowFrameUnits { - fn from(units: WindowFrameUnits) -> Self { - match units { - WindowFrameUnits::Rows => Self::Rows, - WindowFrameUnits::Range => Self::Range, - WindowFrameUnits::Groups => Self::Groups, - } - } -} - -impl TryFrom<&WindowFrameBound> for protobuf::WindowFrameBound { - type Error = Error; - - fn try_from(bound: &WindowFrameBound) -> Result { - Ok(match bound { - WindowFrameBound::CurrentRow => Self { - window_frame_bound_type: protobuf::WindowFrameBoundType::CurrentRow - .into(), - bound_value: None, - }, - WindowFrameBound::Preceding(v) => Self { - window_frame_bound_type: protobuf::WindowFrameBoundType::Preceding.into(), - bound_value: Some(v.try_into()?), - }, - WindowFrameBound::Following(v) => Self { - window_frame_bound_type: protobuf::WindowFrameBoundType::Following.into(), - bound_value: Some(v.try_into()?), - }, - }) - } -} - -impl TryFrom<&WindowFrame> for protobuf::WindowFrame { - type Error = Error; - - fn try_from(window: &WindowFrame) -> Result { - Ok(Self { - window_frame_units: protobuf::WindowFrameUnits::from(window.units).into(), - start_bound: Some((&window.start_bound).try_into()?), - end_bound: Some(protobuf::window_frame::EndBound::Bound( - (&window.end_bound).try_into()?, - )), - }) - } -} - pub fn serialize_exprs<'a, I>( exprs: I, codec: &dyn LogicalExtensionCodec, @@ -209,13 +70,13 @@ pub fn serialize_expr( expr: Some(Box::new(serialize_expr(expr.as_ref(), codec)?)), relation: relation .to_owned() - .map(|r| vec![r.into()]) + .map(|r| vec![protobuf::TableReference::from(r)]) .unwrap_or(vec![]), alias: name.to_owned(), metadata: metadata .as_ref() .map(|m| m.to_hashmap()) - .unwrap_or(HashMap::new()), + .unwrap_or_default(), }); protobuf::LogicalExprNode { expr_type: Some(ExprType::Alias(alias)), @@ -339,8 +200,7 @@ pub fn serialize_expr( let partition_by = serialize_exprs(partition_by, codec)?; let order_by = serialize_sorts(order_by, codec)?; - let window_frame: Option = - Some(window_frame.try_into()?); + let window_frame = Some(protobuf::WindowFrame::try_from(window_frame)?); let window_expr = protobuf::WindowExprNode { exprs: serialize_exprs(args, codec)?, @@ -409,6 +269,19 @@ pub fn serialize_expr( })), } } + Expr::HigherOrderFunction(HigherOrderFunction { func, args }) => { + let mut buf = Vec::new(); + let _ = codec.try_encode_higher_order_function(func.as_ref(), &mut buf); + protobuf::LogicalExprNode { + expr_type: Some(ExprType::HigherOrderUdfExpr( + protobuf::HigherOrderUdfExprNode { + fun_name: func.name().to_string(), + fun_definition: (!buf.is_empty()).then_some(buf), + args: serialize_exprs(args, codec)?, + }, + )), + } + } Expr::Not(expr) => { let expr = Box::new(protobuf::Not { expr: Some(Box::new(serialize_expr(expr.as_ref(), codec)?)), @@ -553,9 +426,10 @@ pub fn serialize_expr( expr_type: Some(ExprType::Negative(expr)), } } - Expr::Unnest(Unnest { expr }) => { + Expr::Unnest(Unnest { expr, outer }) => { let expr = protobuf::Unnest { exprs: vec![serialize_expr(expr.as_ref(), codec)?], + outer: *outer, }; protobuf::LogicalExprNode { expr_type: Some(ExprType::Unnest(expr)), @@ -578,7 +452,7 @@ pub fn serialize_expr( #[expect(deprecated)] Expr::Wildcard { qualifier, .. } => protobuf::LogicalExprNode { expr_type: Some(ExprType::Wildcard(protobuf::Wildcard { - qualifier: qualifier.to_owned().map(|x| x.into()), + qualifier: qualifier.to_owned().map(protobuf::TableReference::from), })), }, Expr::ScalarSubquery(subquery) => protobuf::LogicalExprNode { @@ -631,14 +505,25 @@ pub fn serialize_expr( metadata: field .as_ref() .map(|f| f.metadata().clone()) - .unwrap_or(HashMap::new()), + .unwrap_or_default(), + })), + }, + Expr::Lambda(Lambda { params, body }) => protobuf::LogicalExprNode { + expr_type: Some(ExprType::Lambda(Box::new(protobuf::Lambda { + params: params.clone(), + body: Some(Box::new(serialize_expr(body, codec)?)), + }))), + }, + Expr::LambdaVariable(LambdaVariable { + name, + field, + spans: _, + }) => protobuf::LogicalExprNode { + expr_type: Some(ExprType::LambdaVariable(protobuf::LambdaVariable { + name: name.clone(), + field: field.as_deref().map(|v| v.try_into()).transpose()?, })), }, - Expr::HigherOrderFunction(_) | Expr::Lambda(_) | Expr::LambdaVariable(_) => { - return Err(Error::General( - "Proto serialization error: Lambda not implemented".to_string(), - )); - } }; Ok(expr_node) @@ -681,94 +566,81 @@ where .collect::, Error>>() } -impl From for protobuf::TableReference { - fn from(t: TableReference) -> Self { - use protobuf::table_reference::TableReferenceEnum; - let table_reference_enum = match t { - TableReference::Bare { table } => { - TableReferenceEnum::Bare(protobuf::BareTableReference { - table: table.to_string(), - }) - } - TableReference::Partial { schema, table } => { - TableReferenceEnum::Partial(protobuf::PartialTableReference { - schema: schema.to_string(), - table: table.to_string(), - }) - } - TableReference::Full { - catalog, - schema, - table, - } => TableReferenceEnum::Full(protobuf::FullTableReference { - catalog: catalog.to_string(), - schema: schema.to_string(), - table: table.to_string(), - }), - }; - - protobuf::TableReference { - table_reference_enum: Some(table_reference_enum), - } - } -} - -impl From for protobuf::JoinType { - fn from(t: JoinType) -> Self { - match t { - JoinType::Inner => protobuf::JoinType::Inner, - JoinType::Left => protobuf::JoinType::Left, - JoinType::Right => protobuf::JoinType::Right, - JoinType::Full => protobuf::JoinType::Full, - JoinType::LeftSemi => protobuf::JoinType::Leftsemi, - JoinType::RightSemi => protobuf::JoinType::Rightsemi, - JoinType::LeftAnti => protobuf::JoinType::Leftanti, - JoinType::RightAnti => protobuf::JoinType::Rightanti, - JoinType::LeftMark => protobuf::JoinType::Leftmark, - JoinType::RightMark => protobuf::JoinType::Rightmark, - } - } -} - -impl From for protobuf::JoinConstraint { - fn from(t: JoinConstraint) -> Self { - match t { - JoinConstraint::On => protobuf::JoinConstraint::On, - JoinConstraint::Using => protobuf::JoinConstraint::Using, - } - } +pub(super) fn serialize_range_split_point( + split_point: &SplitPoint, +) -> Result { + Ok(protobuf::RangeSplitPoint { + value: split_point + .values() + .iter() + .map(TryInto::::try_into) + .collect::>()?, + }) } -impl From for protobuf::NullEquality { - fn from(t: NullEquality) -> Self { - match t { - NullEquality::NullEqualsNothing => protobuf::NullEquality::NullEqualsNothing, - NullEquality::NullEqualsNull => protobuf::NullEquality::NullEqualsNull, - } - } +pub fn serialize_merge_into_op( + op: &MergeIntoOp, + codec: &dyn LogicalExtensionCodec, +) -> Result { + Ok(protobuf::MergeIntoOpNode { + on: Some(Box::new(serialize_expr(&op.on, codec)?)), + clauses: op + .clauses + .iter() + .map(|c| serialize_merge_into_clause(c, codec)) + .collect::, Error>>()?, + }) } -impl From<&WriteOp> for protobuf::dml_node::Type { - fn from(t: &WriteOp) -> Self { - match t { - WriteOp::Insert(InsertOp::Append) => protobuf::dml_node::Type::InsertAppend, - WriteOp::Insert(InsertOp::Overwrite) => { - protobuf::dml_node::Type::InsertOverwrite - } - WriteOp::Insert(InsertOp::Replace) => protobuf::dml_node::Type::InsertReplace, - WriteOp::Delete => protobuf::dml_node::Type::Delete, - WriteOp::Update => protobuf::dml_node::Type::Update, - WriteOp::Ctas => protobuf::dml_node::Type::Ctas, - WriteOp::Truncate => protobuf::dml_node::Type::Truncate, - } - } +fn serialize_merge_into_clause( + clause: &MergeIntoClause, + codec: &dyn LogicalExtensionCodec, +) -> Result { + let kind = protobuf::merge_into_clause_node::Kind::from(clause.kind); + let predicate = clause + .predicate + .as_ref() + .map(|e| serialize_expr(e, codec)) + .transpose()?; + Ok(protobuf::MergeIntoClauseNode { + kind: kind.into(), + predicate, + action: Some(serialize_merge_into_action(&clause.action, codec)?), + }) } -impl From for protobuf::NullTreatment { - fn from(t: NullTreatment) -> Self { - match t { - NullTreatment::RespectNulls => protobuf::NullTreatment::RespectNulls, - NullTreatment::IgnoreNulls => protobuf::NullTreatment::IgnoreNulls, +fn serialize_merge_into_action( + action: &MergeIntoAction, + codec: &dyn LogicalExtensionCodec, +) -> Result { + let action = match action { + MergeIntoAction::Update(assignments) => { + let assignments = assignments + .iter() + .map(|(column, value)| { + Ok(protobuf::MergeAssignment { + column: column.clone(), + value: Some(serialize_expr(value, codec)?), + }) + }) + .collect::, Error>>()?; + protobuf::merge_into_action_node::Action::Update( + protobuf::MergeUpdateAction { assignments }, + ) + } + MergeIntoAction::Insert { columns, values } => { + protobuf::merge_into_action_node::Action::Insert( + protobuf::MergeInsertAction { + columns: columns.clone(), + values: serialize_exprs(values, codec)?, + }, + ) } - } + MergeIntoAction::Delete => protobuf::merge_into_action_node::Action::Delete( + protobuf::MergeDeleteAction {}, + ), + }; + Ok(protobuf::MergeIntoActionNode { + action: Some(action), + }) } diff --git a/datafusion/proto/src/physical_plan/from_proto.rs b/datafusion/proto/src/physical_plan/from_proto.rs index 43ebf0474320a..06105be806cfc 100644 --- a/datafusion/proto/src/physical_plan/from_proto.rs +++ b/datafusion/proto/src/physical_plan/from_proto.rs @@ -23,52 +23,35 @@ use arrow::array::RecordBatch; use arrow::compute::SortOptions; use arrow::datatypes::{Field, Schema}; use arrow::ipc::reader::StreamReader; -use chrono::{TimeZone, Utc}; -use datafusion_common::{DataFusionError, Result, internal_datafusion_err, not_impl_err}; +use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; +use datafusion_datasource::TableSchema; use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_groups::FileGroup; -use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; -use datafusion_datasource::file_sink_config::FileSinkConfig; -use datafusion_datasource::{FileRange, ListingTableUrl, PartitionedFile, TableSchema}; -use datafusion_datasource_csv::file_format::CsvSink; -use datafusion_datasource_json::file_format::JsonSink; -#[cfg(feature = "parquet")] -use datafusion_datasource_parquet::file_format::ParquetSink; -use datafusion_execution::object_store::ObjectStoreUrl; +use datafusion_datasource::file_scan_config::FileScanConfig; use datafusion_execution::{FunctionRegistry, TaskContext}; use datafusion_expr::WindowFunctionDefinition; -use datafusion_expr::dml::InsertOp; -use datafusion_expr::execution_props::SubqueryIndex; -use datafusion_physical_expr::projection::{ProjectionExpr, ProjectionExprs}; +use datafusion_physical_expr::expressions::{LambdaExpr, LambdaVariable}; use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; -use datafusion_physical_expr::{LexOrdering, PhysicalSortExpr, ScalarFunctionExpr}; +use datafusion_physical_expr::{ + HigherOrderFunctionExpr, PhysicalSortExpr, ScalarFunctionExpr, +}; use datafusion_physical_plan::expressions::{ - BinaryExpr, CaseExpr, CastExpr, Column, IsNotNullExpr, IsNullExpr, LikeExpr, Literal, - NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn, in_list, + BinaryExpr, CaseExpr, CastExpr, Column, InListExpr, IsNotNullExpr, IsNullExpr, + LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr, UnKnownColumn, }; -use datafusion_physical_plan::joins::{HashExpr, SeededRandomState}; +use datafusion_physical_plan::joins::HashExpr; +use datafusion_physical_plan::proto::ExecutionPlanDecodeCtx; +use datafusion_physical_plan::repartition::RangeExpr; use datafusion_physical_plan::windows::{create_window_expr, schema_add_window_field}; use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use datafusion_proto_common::common::proto_error; -use object_store::ObjectMeta; -use object_store::path::Path; use super::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, - PhysicalProtoConverterExtension, + ConverterPlanDecoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalPlanDecodeContext, PhysicalProtoConverterExtension, }; -use crate::logical_plan::{self}; use crate::protobuf::physical_expr_node::ExprType; use crate::{convert_required, protobuf}; -use datafusion_physical_expr::expressions::{ - DynamicFilterInner, DynamicFilterPhysicalExpr, -}; - -impl From<&protobuf::PhysicalColumn> for Column { - fn from(c: &protobuf::PhysicalColumn) -> Column { - Column::new(&c.name, c.index as usize) - } -} +use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; /// Parses a physical sort expression from a protobuf. /// @@ -154,7 +137,7 @@ pub fn parse_physical_window_expr( let window_frame = proto .window_frame .as_ref() - .map(|wf| wf.clone().try_into()) + .map(|wf| datafusion_expr::WindowFrame::try_from(wf.clone())) .transpose() .map_err(|e| internal_datafusion_err!("{e}"))? .ok_or_else(|| { @@ -266,57 +249,27 @@ pub fn parse_physical_expr_with_converter( .as_ref() .ok_or_else(|| proto_error("Unexpected empty physical expression"))?; - let pexpr: Arc = match expr_type { - ExprType::Column(c) => { - let pcol: Column = c.into(); - Arc::new(pcol) - } - ExprType::UnknownColumn(c) => Arc::new(UnKnownColumn::new(&c.name)), - ExprType::Literal(scalar) => Arc::new(Literal::new(scalar.try_into()?)), - ExprType::BinaryExpr(binary_expr) => { - let op = logical_plan::from_proto::from_proto_binary_op(&binary_expr.op)?; - if !binary_expr.operands.is_empty() { - // New linearized format: reduce the flat operands list back into - // a nested binary expression tree. - let operands: Vec> = binary_expr - .operands - .iter() - .map(|e| proto_converter.proto_to_physical_expr(e, input_schema, ctx)) - .collect::>>()?; - - if operands.len() < 2 { - return Err(proto_error( - "A binary expression must always have at least 2 operands", - )); - } + // Decoder context handed to per-expression `try_from_proto` constructors. + // This is the new shape the codebase is migrating toward (see #21835); + // the remaining `ExprType` variants stay matched inline until they migrate. + let decoder = ConverterDecoder { + ctx, + proto_converter, + }; + let decode_ctx = + datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::new( + input_schema, + &decoder, + ); - operands - .into_iter() - .reduce(|left, right| Arc::new(BinaryExpr::new(left, op, right))) - .expect( - "Binary expression could not be reduced to a single expression.", - ) - } else { - // Legacy format with l/r fields - Arc::new(BinaryExpr::new( - parse_required_physical_expr( - binary_expr.l.as_deref(), - ctx, - "left", - input_schema, - proto_converter, - )?, - op, - parse_required_physical_expr( - binary_expr.r.as_deref(), - ctx, - "right", - input_schema, - proto_converter, - )?, - )) - } - } + let pexpr: Arc = match expr_type { + // Migrated expressions take the whole `PhysicalExprNode` and unwrap + // their own `ExprType` variant — see #21835. This match only routes + // to the right constructor. + ExprType::Column(_) => Column::try_from_proto(proto, &decode_ctx)?, + ExprType::UnknownColumn(_) => UnKnownColumn::try_from_proto(proto, &decode_ctx)?, + ExprType::Literal(_) => Literal::try_from_proto(proto, &decode_ctx)?, + ExprType::BinaryExpr(_) => BinaryExpr::try_from_proto(proto, &decode_ctx)?, ExprType::AggregateExpr(_) => { return not_impl_err!( "Cannot convert aggregate expr node to physical expression" @@ -330,108 +283,14 @@ pub fn parse_physical_expr_with_converter( ExprType::Sort(_) => { return not_impl_err!("Cannot convert sort expr node to physical expression"); } - ExprType::IsNullExpr(e) => { - Arc::new(IsNullExpr::new(parse_required_physical_expr( - e.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?)) - } - ExprType::IsNotNullExpr(e) => { - Arc::new(IsNotNullExpr::new(parse_required_physical_expr( - e.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?)) - } - ExprType::NotExpr(e) => Arc::new(NotExpr::new(parse_required_physical_expr( - e.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?)), - ExprType::Negative(e) => { - Arc::new(NegativeExpr::new(parse_required_physical_expr( - e.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?)) - } - ExprType::InList(e) => in_list( - parse_required_physical_expr( - e.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?, - parse_physical_exprs(&e.list, ctx, input_schema, proto_converter)?, - &e.negated, - input_schema, - )?, - ExprType::Case(e) => Arc::new(CaseExpr::try_new( - e.expr - .as_ref() - .map(|e| { - proto_converter.proto_to_physical_expr(e.as_ref(), input_schema, ctx) - }) - .transpose()?, - e.when_then_expr - .iter() - .map(|e| { - Ok(( - parse_required_physical_expr( - e.when_expr.as_ref(), - ctx, - "when_expr", - input_schema, - proto_converter, - )?, - parse_required_physical_expr( - e.then_expr.as_ref(), - ctx, - "then_expr", - input_schema, - proto_converter, - )?, - )) - }) - .collect::>>()?, - e.else_expr - .as_ref() - .map(|e| { - proto_converter.proto_to_physical_expr(e.as_ref(), input_schema, ctx) - }) - .transpose()?, - )?), - ExprType::Cast(e) => Arc::new(CastExpr::new( - parse_required_physical_expr( - e.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?, - convert_required!(e.arrow_type)?, - None, - )), - ExprType::TryCast(e) => Arc::new(TryCastExpr::new( - parse_required_physical_expr( - e.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?, - convert_required!(e.arrow_type)?, - )), + ExprType::IsNullExpr(_) => IsNullExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::IsNotNullExpr(_) => IsNotNullExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::NotExpr(_) => NotExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::Negative(_) => NegativeExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::InList(_) => InListExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::Case(_) => CaseExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::Cast(_) => CastExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::TryCast(_) => TryCastExpr::try_from_proto(proto, &decode_ctx)?, ExprType::ScalarUdf(e) => { let udf = match &e.fun_definition { Some(buf) => ctx.codec().try_decode_udf(&e.name, buf)?, @@ -462,104 +321,45 @@ pub fn parse_physical_expr_with_converter( .with_nullable(e.nullable), ) } - ExprType::LikeExpr(like_expr) => Arc::new(LikeExpr::new( - like_expr.negated, - like_expr.case_insensitive, - parse_required_physical_expr( - like_expr.expr.as_deref(), - ctx, - "expr", - input_schema, - proto_converter, - )?, - parse_required_physical_expr( - like_expr.pattern.as_deref(), - ctx, - "pattern", - input_schema, - proto_converter, - )?, - )), - ExprType::HashExpr(hash_expr) => { - let on_columns = parse_physical_exprs( - &hash_expr.on_columns, - ctx, + ExprType::HigherOrderUdf(e) => { + let func = match &e.fun_definition { + Some(buf) => { + ctx.codec().try_decode_higher_order_function(&e.name, buf)? + } + None => ctx + .task_ctx() + .higher_order_function(e.name.as_str()) + .or_else(|_| { + ctx.codec().try_decode_higher_order_function(&e.name, &[]) + })?, + }; + let func_def = Arc::clone(&func); + + let args = parse_physical_exprs(&e.args, ctx, input_schema, proto_converter)?; + + let config_options = Arc::clone(ctx.task_ctx().session_config().options()); + + Arc::new(HigherOrderFunctionExpr::try_new_with_schema( + func_def, + args, input_schema, - proto_converter, - )?; - Arc::new(HashExpr::new( - on_columns, - SeededRandomState::with_seed(hash_expr.seed0), - hash_expr.description.clone(), - )) + config_options, + )?) } - ExprType::ScalarSubquery(sq) => { - let data_type: arrow::datatypes::DataType = sq - .data_type - .as_ref() - .ok_or_else(|| { - proto_error("Missing data_type in PhysicalScalarSubqueryExprNode") - })? - .try_into()?; + ExprType::LikeExpr(_) => LikeExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::HashExpr(_) => HashExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::RangeExpr(_) => RangeExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::ScalarSubquery(_) => { let results = ctx.scalar_subquery_results().ok_or_else(|| { proto_error( "ScalarSubqueryExpr can only be deserialized as part \ of a surrounding ScalarSubqueryExec", ) })?; - Arc::new(ScalarSubqueryExpr::new( - data_type, - sq.nullable, - SubqueryIndex::new(sq.index as usize), - results.clone(), - )) + ScalarSubqueryExpr::try_from_proto(proto, &decode_ctx, results)? } - ExprType::DynamicFilter(dynamic_filter) => { - let children = parse_physical_exprs( - &dynamic_filter.children, - ctx, - input_schema, - proto_converter, - )?; - - let remapped_children = if !dynamic_filter.remapped_children.is_empty() { - Some(parse_physical_exprs( - &dynamic_filter.remapped_children, - ctx, - input_schema, - proto_converter, - )?) - } else { - None - }; - - let inner_expr = parse_required_physical_expr( - dynamic_filter.inner_expr.as_deref(), - ctx, - "inner_expr", - input_schema, - proto_converter, - )?; - - let expression_id = proto.expr_id.ok_or_else(|| { - proto_error( - "DynamicFilterPhysicalExpr requires PhysicalExprNode.expr_id \ - to be set by the serializer", - ) - })?; - - let base_filter: Arc = - Arc::new(DynamicFilterPhysicalExpr::from_parts( - children, - remapped_children, - DynamicFilterInner { - expression_id, - generation: dynamic_filter.generation, - expr: inner_expr, - is_complete: dynamic_filter.is_complete, - }, - )); - base_filter + ExprType::DynamicFilter(_) => { + DynamicFilterPhysicalExpr::try_from_proto(proto, &decode_ctx)? } ExprType::Extension(extension) => { let inputs: Vec> = extension @@ -567,48 +367,37 @@ pub fn parse_physical_expr_with_converter( .iter() .map(|e| proto_converter.proto_to_physical_expr(e, input_schema, ctx)) .collect::>()?; - ctx.codec() - .try_decode_expr(extension.expr.as_slice(), &inputs)? as _ + ctx.codec().try_decode_expr( + extension.expr.as_slice(), + &inputs, + &decode_ctx, + )? as _ + } + ExprType::Lambda(_) => LambdaExpr::try_from_proto(proto, &decode_ctx)?, + ExprType::LambdaVariable(_) => { + LambdaVariable::try_from_proto(proto, &decode_ctx)? } }; Ok(pexpr) } -fn parse_required_physical_expr( - expr: Option<&protobuf::PhysicalExprNode>, - ctx: &PhysicalPlanDecodeContext<'_>, - field: &str, - input_schema: &Schema, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result> { - expr.map(|e| proto_converter.proto_to_physical_expr(e, input_schema, ctx)) - .transpose()? - .ok_or_else(|| internal_datafusion_err!("Missing required field {field:?}")) -} - pub fn parse_protobuf_hash_partitioning( partitioning: Option<&protobuf::PhysicalHashRepartition>, ctx: &PhysicalPlanDecodeContext<'_>, input_schema: &Schema, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - match partitioning { - Some(hash_part) => { - let expr = parse_physical_exprs( - &hash_part.hash_expr, - ctx, - input_schema, - proto_converter, - )?; - - Ok(Some(Partitioning::Hash( - expr, - hash_part.partition_count.try_into().unwrap(), - ))) - } - None => Ok(None), - } + // Delegate to the shared decoder rather than keep a second copy of the hash + // wire format: a partition count that does not fit in `usize` (a 32-bit + // target reading a plan written on a 64-bit one) is then an error here too + // instead of a panic. + let hash = partitioning.map(|hash_part| protobuf::Partitioning { + partition_method: Some(protobuf::partitioning::PartitionMethod::Hash( + hash_part.clone(), + )), + }); + parse_protobuf_partitioning(hash.as_ref(), ctx, input_schema, proto_converter) } pub fn parse_protobuf_partitioning( @@ -617,32 +406,24 @@ pub fn parse_protobuf_partitioning( input_schema: &Schema, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - match partitioning { - Some(protobuf::Partitioning { partition_method }) => match partition_method { - Some(protobuf::partitioning::PartitionMethod::RoundRobin( - partition_count, - )) => Ok(Some(Partitioning::RoundRobinBatch( - *partition_count as usize, - ))), - Some(protobuf::partitioning::PartitionMethod::Hash(hash_repartition)) => { - parse_protobuf_hash_partitioning( - Some(hash_repartition), - ctx, - input_schema, - proto_converter, - ) - } - Some(protobuf::partitioning::PartitionMethod::Unknown(partition_count)) => { - Ok(Some(Partitioning::UnknownPartitioning( - *partition_count as usize, - ))) - } - None => Ok(None), - }, - None => Ok(None), - } + let decoder = ConverterDecoder { + ctx, + proto_converter, + }; + let decode_ctx = + datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx::new( + input_schema, + &decoder, + ); + partitioning + .map(|partitioning| Partitioning::try_from_proto(partitioning, &decode_ctx)) + .transpose() + .map(Option::flatten) } - +#[deprecated( + since = "55.0.0", + note = "unused by DataFusion; use `FileScanConfig::parse_table_schema_from_proto` to reconstruct the full table schema" +)] pub fn parse_protobuf_file_scan_schema( proto: &protobuf::FileScanExecConf, ) -> Result> { @@ -653,31 +434,7 @@ pub fn parse_protobuf_file_scan_schema( pub fn parse_table_schema_from_proto( proto: &protobuf::FileScanExecConf, ) -> Result { - let schema: Arc = parse_protobuf_file_scan_schema(proto)?; - - // Reacquire the partition column types from the schema before removing them below. - let table_partition_cols = proto - .table_partition_cols - .iter() - .map(|col| Ok(Arc::new(schema.field_with_name(col)?.clone()))) - .collect::>>()?; - - // Remove partition columns from the schema after recreating table_partition_cols - // because the partition columns are not in the file. They are present to allow - // the partition column types to be reconstructed after serde. - let file_schema = Arc::new( - Schema::new( - schema - .fields() - .iter() - .filter(|field| !table_partition_cols.contains(field)) - .cloned() - .collect::>(), - ) - .with_metadata(schema.metadata.clone()), - ); - - Ok(TableSchema::new(file_schema, table_partition_cols)) + FileScanConfig::parse_table_schema_from_proto(proto) } pub fn parse_protobuf_file_scan_config( @@ -686,71 +443,21 @@ pub fn parse_protobuf_file_scan_config( proto_converter: &dyn PhysicalProtoConverterExtension, file_source: Arc, ) -> Result { - let schema: Arc = parse_protobuf_file_scan_schema(proto)?; - - let constraints = convert_required!(proto.constraints)?; - let statistics = convert_required!(proto.statistics)?; - - let file_groups = proto - .file_groups - .iter() - .map(|f| f.try_into()) - .collect::, _>>()?; - - let object_store_url = match proto.object_store_url.is_empty() { - false => ObjectStoreUrl::parse(&proto.object_store_url)?, - true => ObjectStoreUrl::local_filesystem(), + let decoder = ConverterPlanDecoder { + ctx, + proto_converter, }; - - let mut output_ordering = vec![]; - for node_collection in &proto.output_ordering { - let sort_exprs = parse_physical_sort_exprs( - &node_collection.physical_sort_expr_nodes, - ctx, - &schema, - proto_converter, - )?; - output_ordering.extend(LexOrdering::new(sort_exprs)); - } - - // Parse projection expressions if present and apply to file source - let file_source = if let Some(proto_projection_exprs) = &proto.projection_exprs { - let projection_exprs: Vec = proto_projection_exprs - .projections - .iter() - .map(|proto_expr| { - let expr = proto_converter.proto_to_physical_expr( - proto_expr.expr.as_ref().ok_or_else(|| { - internal_datafusion_err!("ProjectionExpr missing expr field") - })?, - &schema, - ctx, - )?; - Ok(ProjectionExpr::new(expr, proto_expr.alias.clone())) - }) - .collect::>>()?; - - let projection_exprs = ProjectionExprs::new(projection_exprs); - - // Apply projection to file source - file_source - .try_pushdown_projection(&projection_exprs)? - .unwrap_or(file_source) - } else { - file_source - }; - - let config = FileScanConfigBuilder::new(object_store_url, file_source) - .with_file_groups(file_groups) - .with_constraints(constraints) - .with_statistics(statistics) - .with_limit(proto.limit.as_ref().map(|sl| sl.limit as usize)) - .with_output_ordering(output_ordering) - .with_batch_size(proto.batch_size.map(|s| s as usize)) - .build(); - Ok(config) + FileScanConfig::try_from_proto( + proto, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) } +#[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `MemorySourceConfig` deserializes its record batches itself via `MemorySourceConfig::try_from_proto`" +)] pub fn parse_record_batches(buf: &[u8]) -> Result> { if buf.is_empty() { return Ok(vec![]); @@ -763,185 +470,29 @@ pub fn parse_record_batches(buf: &[u8]) -> Result> { Ok(batches) } -impl TryFrom<&protobuf::PartitionedFile> for PartitionedFile { - type Error = DataFusionError; - - fn try_from(val: &protobuf::PartitionedFile) -> Result { - let mut pf = PartitionedFile::new_from_meta(ObjectMeta { - location: Path::parse(val.path.as_str()) - .map_err(|e| proto_error(format!("Invalid object_store path: {e}")))?, - last_modified: Utc.timestamp_nanos(val.last_modified_ns as i64), - size: val.size, - e_tag: None, - version: None, - }) - .with_partition_values( - val.partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - ); - if let Some(range) = val.range.as_ref() { - let file_range: FileRange = range.try_into()?; - pf = pf.with_range(file_range.start, file_range.end); - } - if let Some(proto_stats) = val.statistics.as_ref() { - pf = pf.with_statistics(Arc::new(proto_stats.try_into()?)); - } - Ok(pf) - } -} - -impl TryFrom<&protobuf::FileRange> for FileRange { - type Error = DataFusionError; - - fn try_from(value: &protobuf::FileRange) -> Result { - Ok(FileRange { - start: value.start, - end: value.end, - }) - } -} - -impl TryFrom<&protobuf::FileGroup> for FileGroup { - type Error = DataFusionError; - - fn try_from(val: &protobuf::FileGroup) -> Result { - let files = val - .files - .iter() - .map(|f| f.try_into()) - .collect::, _>>()?; - Ok(FileGroup::new(files)) - } -} - -impl TryFrom<&protobuf::JsonSink> for JsonSink { - type Error = DataFusionError; - - fn try_from(value: &protobuf::JsonSink) -> Result { - Ok(Self::new( - convert_required!(value.config)?, - convert_required!(value.writer_options)?, - )) - } -} - -#[cfg(feature = "parquet")] -impl TryFrom<&protobuf::ParquetSink> for ParquetSink { - type Error = DataFusionError; - - fn try_from(value: &protobuf::ParquetSink) -> Result { - Ok(Self::new( - convert_required!(value.config)?, - convert_required!(value.parquet_options)?, - )) - } -} - -impl TryFrom<&protobuf::CsvSink> for CsvSink { - type Error = DataFusionError; - - fn try_from(value: &protobuf::CsvSink) -> Result { - Ok(Self::new( - convert_required!(value.config)?, - convert_required!(value.writer_options)?, - )) - } -} - -impl TryFrom<&protobuf::FileSinkConfig> for FileSinkConfig { - type Error = DataFusionError; - - fn try_from(conf: &protobuf::FileSinkConfig) -> Result { - let file_group = FileGroup::new( - conf.file_groups - .iter() - .map(|f| f.try_into()) - .collect::>>()?, - ); - let table_paths = conf - .table_paths - .iter() - .map(ListingTableUrl::parse) - .collect::>>()?; - let table_partition_cols = conf - .table_partition_cols - .iter() - .map(|protobuf::PartitionColumn { name, arrow_type }| { - let data_type = convert_required!(arrow_type)?; - Ok((name.clone(), data_type)) - }) - .collect::>>()?; - let insert_op = match conf.insert_op() { - protobuf::InsertOp::Append => InsertOp::Append, - protobuf::InsertOp::Overwrite => InsertOp::Overwrite, - protobuf::InsertOp::Replace => InsertOp::Replace, - }; - let file_output_mode = match conf.file_output_mode() { - protobuf::FileOutputMode::Automatic => { - datafusion_datasource::file_sink_config::FileOutputMode::Automatic - } - protobuf::FileOutputMode::SingleFile => { - datafusion_datasource::file_sink_config::FileOutputMode::SingleFile - } - protobuf::FileOutputMode::Directory => { - datafusion_datasource::file_sink_config::FileOutputMode::Directory - } - }; - Ok(Self { - original_url: String::default(), - object_store_url: ObjectStoreUrl::parse(&conf.object_store_url)?, - file_group, - table_paths, - output_schema: Arc::new(convert_required!(conf.output_schema)?), - table_partition_cols, - insert_op, - keep_partition_by_columns: conf.keep_partition_by_columns, - file_extension: conf.file_extension.clone(), - file_output_mode, - }) - } +/// Concrete [`PhysicalExprDecode`] driver that backs +/// [`PhysicalExprDecodeCtx`] inside `parse_physical_expr_with_converter`. +/// +/// Today this is a thin wrapper that re-enters the central match through +/// `proto_to_physical_expr`; once more expressions migrate, the central match +/// shrinks and a future builder-style decoder can take over. +/// +/// [`PhysicalExprDecode`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecode +/// [`PhysicalExprDecodeCtx`]: datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx +struct ConverterDecoder<'a, 'b> { + ctx: &'a PhysicalPlanDecodeContext<'b>, + proto_converter: &'a dyn PhysicalProtoConverterExtension, } -#[cfg(test)] -mod tests { - - use super::*; - - #[test] - fn partitioned_file_path_roundtrip_percent_encoded() { - let path_str = "foo/foo%2Fbar/baz%252Fqux"; - let pf = PartitionedFile::new_from_meta(ObjectMeta { - location: Path::parse(path_str).unwrap(), - last_modified: Utc.timestamp_nanos(1_000), - size: 42, - e_tag: None, - version: None, - }); - - let proto = protobuf::PartitionedFile::try_from(&pf).unwrap(); - assert_eq!(proto.path, path_str); - - let pf2 = PartitionedFile::try_from(&proto).unwrap(); - assert_eq!(pf2.object_meta.location.as_ref(), path_str); - assert_eq!(pf2.object_meta.location, pf.object_meta.location); - assert_eq!(pf2.object_meta.size, pf.object_meta.size); - assert_eq!(pf2.object_meta.last_modified, pf.object_meta.last_modified); - } - - #[test] - fn partitioned_file_from_proto_invalid_path() { - let proto = protobuf::PartitionedFile { - path: "foo//bar".to_string(), - size: 1, - last_modified_ns: 0, - partition_values: vec![], - range: None, - statistics: None, - }; - - let err = PartitionedFile::try_from(&proto).unwrap_err(); - assert!(err.to_string().contains("Invalid object_store path")); +impl datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecode + for ConverterDecoder<'_, '_> +{ + fn decode( + &self, + node: &protobuf::PhysicalExprNode, + schema: &Schema, + ) -> Result> { + self.proto_converter + .proto_to_physical_expr(node, schema, self.ctx) } } diff --git a/datafusion/proto/src/physical_plan/mod.rs b/datafusion/proto/src/physical_plan/mod.rs index 9a5489177319d..222901aff5211 100644 --- a/datafusion/proto/src/physical_plan/mod.rs +++ b/datafusion/proto/src/physical_plan/mod.rs @@ -21,19 +21,11 @@ use std::collections::HashMap; use std::fmt::Debug; use std::sync::Arc; -use arrow::compute::SortOptions; use arrow::datatypes::{IntervalMonthDayNanoType, Schema, SchemaRef}; use datafusion_catalog::memory::MemorySourceConfig; -use datafusion_common::config::CsvOptions; use datafusion_common::{ DataFusionError, Result, internal_datafusion_err, internal_err, not_impl_err, }; -#[cfg(feature = "parquet")] -use datafusion_datasource::file::FileSource; -use datafusion_datasource::file_compression_type::FileCompressionType; -use datafusion_datasource::file_scan_config::{FileScanConfig, FileScanConfigBuilder}; -use datafusion_datasource::sink::DataSinkExec; -use datafusion_datasource::source::{DataSource, DataSourceExec}; use datafusion_datasource_arrow::source::ArrowSource; #[cfg(feature = "avro")] use datafusion_datasource_avro::source::AvroSource; @@ -42,79 +34,59 @@ use datafusion_datasource_csv::source::CsvSource; use datafusion_datasource_json::file_format::JsonSink; use datafusion_datasource_json::source::JsonSource; #[cfg(feature = "parquet")] -use datafusion_datasource_parquet::CachedParquetFileReaderFactory; -#[cfg(feature = "parquet")] use datafusion_datasource_parquet::file_format::ParquetSink; #[cfg(feature = "parquet")] use datafusion_datasource_parquet::source::ParquetSource; -#[cfg(feature = "parquet")] -use datafusion_execution::object_store::ObjectStoreUrl; use datafusion_execution::{FunctionRegistry, TaskContext}; -use datafusion_expr::execution_props::{ScalarSubqueryResults, SubqueryIndex}; -use datafusion_expr::{AggregateUDF, ScalarUDF, WindowUDF}; +use datafusion_expr::physical_planning_context::ScalarSubqueryResults; +use datafusion_expr::{AggregateUDF, HigherOrderUDF, ScalarUDF, WindowUDF}; use datafusion_functions_table::generate_series::{ Empty, GenSeriesArgs, GenerateSeriesTable, GenericSeriesState, TimestampValue, }; -use datafusion_physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; -use datafusion_physical_expr::async_scalar_function::AsyncFuncExpr; -use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; -use datafusion_physical_expr::{LexOrdering, LexRequirement, PhysicalExprRef}; -use datafusion_physical_plan::aggregates::{ - AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, -}; +use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; +use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; +use datafusion_physical_plan::aggregates::AggregateExec; use datafusion_physical_plan::analyze::AnalyzeExec; use datafusion_physical_plan::async_func::AsyncFuncExec; use datafusion_physical_plan::buffer::BufferExec; -#[expect(deprecated)] +#[expect( + deprecated, + reason = "`CoalesceBatchesExec` remains supported for protobuf compatibility" +)] use datafusion_physical_plan::coalesce_batches::CoalesceBatchesExec; use datafusion_physical_plan::coalesce_partitions::CoalescePartitionsExec; use datafusion_physical_plan::coop::CooperativeExec; use datafusion_physical_plan::empty::EmptyExec; use datafusion_physical_plan::explain::ExplainExec; -use datafusion_physical_plan::expressions::PhysicalSortExpr; -use datafusion_physical_plan::filter::{FilterExec, FilterExecBuilder}; -use datafusion_physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion_physical_plan::filter::FilterExec; use datafusion_physical_plan::joins::{ - CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, + CrossJoinExec, HashJoinExec, NestedLoopJoinExec, SortMergeJoinExec, + SymmetricHashJoinExec, }; use datafusion_physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; use datafusion_physical_plan::memory::LazyMemoryExec; -use datafusion_physical_plan::metrics::{MetricCategory, MetricType}; use datafusion_physical_plan::placeholder_row::PlaceholderRowExec; -use datafusion_physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion_physical_plan::projection::ProjectionExec; +use datafusion_physical_plan::proto::{ + ExecutionPlanDecode, ExecutionPlanDecodeCtx, ExecutionPlanEncode, + ExecutionPlanEncodeCtx, +}; use datafusion_physical_plan::repartition::RepartitionExec; -use datafusion_physical_plan::scalar_subquery::{ScalarSubqueryExec, ScalarSubqueryLink}; +use datafusion_physical_plan::scalar_subquery::ScalarSubqueryExec; use datafusion_physical_plan::sorts::sort::SortExec; use datafusion_physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; use datafusion_physical_plan::union::{InterleaveExec, UnionExec}; -use datafusion_physical_plan::unnest::{ListUnnest, UnnestExec}; -use datafusion_physical_plan::windows::{BoundedWindowAggExec, WindowAggExec}; -use datafusion_physical_plan::{ExecutionPlan, InputOrderMode, PhysicalExpr, WindowExpr}; +use datafusion_physical_plan::unnest::UnnestExec; +use datafusion_physical_plan::windows::WindowAggExec; +use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; use prost::Message; use prost::bytes::BufMut; -use self::from_proto::parse_protobuf_partitioning; -use self::to_proto::serialize_partitioning; -use crate::common::{byte_to_string, str_to_byte}; -use crate::physical_plan::from_proto::{ - parse_physical_expr_with_converter, parse_physical_sort_expr, - parse_physical_sort_exprs, parse_physical_window_expr, - parse_protobuf_file_scan_config, parse_record_batches, parse_table_schema_from_proto, -}; -use crate::physical_plan::to_proto::{ - serialize_file_scan_config, serialize_maybe_filter, serialize_physical_aggr_expr, - serialize_physical_expr_with_converter, serialize_physical_sort_exprs, - serialize_physical_window_expr, serialize_record_batches, -}; -use crate::protobuf::physical_aggregate_expr_node::AggregateFunction; -use crate::protobuf::physical_expr_node::ExprType; +use crate::convert_required; +use crate::physical_plan::from_proto::parse_physical_expr_with_converter; +use crate::physical_plan::to_proto::serialize_physical_expr_with_converter; use crate::protobuf::physical_plan_node::PhysicalPlanType; -use crate::protobuf::{ - self, ListUnnest as ProtoListUnnest, SortExprNode, SortMergeJoinExecNode, - proto_error, window_agg_exec_node, -}; -use crate::{convert_required, into_required}; +use crate::protobuf::{self, proto_error}; pub mod from_proto; pub mod to_proto; @@ -128,46 +100,829 @@ fn encode_human_display_alias(human_display: &str, alias: &str) -> String { ) } -fn split_human_display_alias<'a>( - human_display: &'a str, - name: &'a str, -) -> (&'a str, Option<&'a str>) { - if let Some(encoded) = human_display.strip_prefix(HUMAN_DISPLAY_ALIAS_PREFIX) - && let Some((alias_len, encoded)) = encoded.split_once(':') - && let Ok(alias_len) = alias_len.parse::() - && let Some(alias) = encoded.get(..alias_len) - && let Some(human_display) = encoded.get(alias_len..) - && alias == name - && !human_display.is_empty() - { - return (human_display, Some(alias)); +#[cfg(test)] +mod file_scan_config_serde { + use super::*; + use arrow::datatypes::{DataType, Field}; + use datafusion_common::{Constraint, Constraints, ScalarValue, Statistics}; + use datafusion_datasource::file::FileSource; + use datafusion_datasource::file_groups::FileGroup; + use datafusion_datasource::file_scan_config::{ + FileScanConfig, FileScanConfigBuilder, + }; + use datafusion_datasource::file_stream::FileOpener; + use datafusion_datasource::{PartitionedFile, TableSchema}; + use datafusion_execution::object_store::ObjectStoreUrl; + use datafusion_physical_expr::expressions::Column; + use datafusion_physical_expr::projection::{ + ProjectionExpr as FileProjectionExpr, ProjectionExprs as FileProjectionExprs, + }; + use datafusion_physical_expr::{ + LexOrdering, Partitioning, PhysicalSortExpr, RangePartitioning, SplitPoint, + }; + use datafusion_physical_plan::metrics::ExecutionPlanMetricsSet; + use object_store::ObjectStore; + + #[derive(Clone)] + struct SerdeTestSource { + metrics: ExecutionPlanMetricsSet, + table_schema: TableSchema, + projection: Option, + } + + impl SerdeTestSource { + fn new( + table_schema: TableSchema, + projection: Option, + ) -> Self { + Self { + metrics: ExecutionPlanMetricsSet::new(), + table_schema, + projection, + } + } } - (human_display, None) -} + impl FileSource for SerdeTestSource { + fn create_file_opener( + &self, + _object_store: Arc, + _base_config: &FileScanConfig, + _partition: usize, + ) -> Result> { + internal_err!("not needed for FileScanConfig serde tests") + } -#[cfg(test)] -mod tests { - use super::*; + fn table_schema(&self) -> &TableSchema { + &self.table_schema + } + + fn with_batch_size(&self, _batch_size: usize) -> Arc { + Arc::new(self.clone()) + } + + fn metrics(&self) -> &ExecutionPlanMetricsSet { + &self.metrics + } + + fn file_type(&self) -> &str { + "serde-test" + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut( + &Arc, + ) + -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots( + self.projection.iter().flatten(), + f, + ) + } + + fn try_pushdown_projection( + &self, + projection: &FileProjectionExprs, + ) -> Result>> { + Ok(Some(Arc::new(Self { + projection: Some(projection.clone()), + ..self.clone() + }))) + } + + fn projection(&self) -> Option<&FileProjectionExprs> { + self.projection.as_ref() + } + } + + fn populated_projection() -> FileProjectionExprs { + FileProjectionExprs::new(vec![FileProjectionExpr::new( + Arc::new(Column::new("value", 0)), + "projected_value", + )]) + } + + fn test_config(output_partitioning: Option) -> FileScanConfig { + test_config_with_projection(output_partitioning, Some(populated_projection())) + } + + fn test_config_with_projection( + output_partitioning: Option, + projection: Option, + ) -> FileScanConfig { + let file_schema = Arc::new( + Schema::new(vec![ + Field::new("value", DataType::Int32, false), + Field::new("label", DataType::Utf8, true), + ]) + .with_metadata(HashMap::from([( + "serde_test_key".to_string(), + "serde_test_value".to_string(), + )])), + ); + let table_schema = TableSchema::builder(Arc::clone(&file_schema)) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part", + DataType::Utf8, + false, + ))]) + .build(); + let table_statistics = Statistics::new_unknown(table_schema.table_schema()); + let source = Arc::new(SerdeTestSource::new(table_schema, projection)); + let first_file = PartitionedFile::new("data/part=a/file.arrow", 1024) + .with_partition_values(vec![ScalarValue::Utf8(Some("a".to_string()))]) + .with_range(10, 900) + .with_arrow_schema(Arc::clone(&file_schema)) + .with_statistics(Arc::new(table_statistics.clone())); + let second_file = PartitionedFile::new("data/part=b/file.arrow", 2048) + .with_partition_values(vec![ScalarValue::Utf8(Some("b".to_string()))]); + let third_file = PartitionedFile::new("data/part=c/file.arrow", 4096) + .with_partition_values(vec![ScalarValue::Utf8(Some("c".to_string()))]) + .with_arrow_schema(Arc::clone(&file_schema)); + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("value", 0), + ))]) + .expect("single expression ordering"); + + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file_groups(vec![ + FileGroup::new(vec![first_file, second_file]), + FileGroup::new(vec![third_file]), + ]) + .with_constraints(Constraints::new_unverified(vec![Constraint::PrimaryKey( + vec![0], + )])) + .with_statistics(table_statistics) + .with_limit(Some(17)) + .with_batch_size(Some(256)) + .with_output_ordering(vec![ordering]) + .with_output_partitioning(output_partitioning) + .build() + } + + fn hash_partitioning() -> Partitioning { + Partitioning::Hash(vec![Arc::new(Column::new("value", 0))], 3) + } + + fn range_partitioning() -> Partitioning { + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new( + Column::new("value", 0), + ))]) + .expect("single expression ordering"); + Partitioning::Range(RangePartitioning::new( + ordering, + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )) + } + + fn decode_source(conf: &protobuf::FileScanExecConf) -> Result> { + Ok(Arc::new(SerdeTestSource::new( + FileScanConfig::parse_table_schema_from_proto(conf)?, + None, + ))) + } + + struct FileScanSerdeHarness { + codec: DefaultPhysicalExtensionCodec, + converter: DefaultPhysicalProtoConverter, + task_ctx: TaskContext, + } + + impl FileScanSerdeHarness { + fn new() -> Self { + Self { + codec: DefaultPhysicalExtensionCodec {}, + converter: DefaultPhysicalProtoConverter {}, + task_ctx: TaskContext::default(), + } + } + + fn encode(&self, config: &FileScanConfig) -> Result { + let encoder = ConverterPlanEncoder { + codec: &self.codec, + proto_converter: &self.converter, + }; + config.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) + } + + fn decode(&self, conf: &protobuf::FileScanExecConf) -> Result { + self.decode_with_source(conf, decode_source(conf)?) + } + + fn decode_with_source( + &self, + conf: &protobuf::FileScanExecConf, + file_source: Arc, + ) -> Result { + let physical_decode_ctx = + PhysicalPlanDecodeContext::new(&self.task_ctx, &self.codec); + let decoder = ConverterPlanDecoder { + ctx: &physical_decode_ctx, + proto_converter: &self.converter, + }; + FileScanConfig::try_from_proto( + conf, + &ExecutionPlanDecodeCtx::new(&decoder), + file_source, + ) + } + } + + #[test] + fn new_file_scan_config_serde_roundtrips_all_partitioning_variants() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + + for config in [ + test_config(None), + test_config(Some(Partitioning::RoundRobinBatch(2))), + test_config(Some(hash_partitioning())), + test_config(Some(range_partitioning())), + test_config(Some(Partitioning::UnknownPartitioning(4))), + ] { + let encoded = serde.encode(&config)?; + let reencoded = serde.encode(&serde.decode(&encoded)?)?; + assert_eq!(reencoded.output_partitioning, encoded.output_partitioning); + } + + Ok(()) + } #[test] - fn split_human_display_alias_ignores_mismatched_alias() { - let encoded = encode_human_display_alias("sum(value)", "revenue"); + fn new_file_scan_config_serde_preserves_complete_fixture() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let config = test_config(None); + let decoded = serde.decode(&serde.encode(&config)?)?; + assert_eq!(decoded.constraints, config.constraints); assert_eq!( - split_human_display_alias(&encoded, "other"), - (encoded.as_str(), None) + decoded.file_schema().metadata, + config.file_schema().metadata ); + assert_eq!(decoded.file_groups.len(), 2); + assert_eq!(decoded.file_groups[0].len(), 2); + assert_eq!(decoded.file_groups[1].len(), 1); + assert!(decoded.file_groups[0].files()[0].arrow_schema.is_some()); + assert!(decoded.file_groups[0].files()[1].arrow_schema.is_none()); + + Ok(()) } #[test] - fn split_human_display_alias_keeps_malformed_prefix_literal() { - let display = format!("{HUMAN_DISPLAY_ALIAS_PREFIX}not-an-encoding"); + fn new_file_scan_config_serde_preserves_projection_presence() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + + let absent = serde.encode(&test_config_with_projection(None, None))?; + assert!(absent.projection_exprs.is_none()); + assert!(serde.decode(&absent)?.file_source().projection().is_none()); + + let empty = serde.encode(&test_config_with_projection( + None, + Some(FileProjectionExprs::new(vec![])), + ))?; + assert!( + empty + .projection_exprs + .as_ref() + .is_some_and(|projection| projection.projections.is_empty()) + ); + assert!( + serde + .decode(&empty)? + .file_source() + .projection() + .is_some_and(|projection| projection.as_ref().is_empty()) + ); - assert_eq!( - split_human_display_alias(&display, "agg"), - (display.as_str(), None) + Ok(()) + } + + #[test] + fn new_file_scan_config_decode_rejects_malformed_required_fields() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let valid = serde.encode(&test_config(None))?; + let file_source = decode_source(&valid)?; + + for (field, malformed) in [ + ( + "schema", + protobuf::FileScanExecConf { + schema: None, + ..valid.clone() + }, + ), + ( + "constraints", + protobuf::FileScanExecConf { + constraints: None, + ..valid.clone() + }, + ), + ( + "statistics", + protobuf::FileScanExecConf { + statistics: None, + ..valid.clone() + }, + ), + ] { + let err = serde + .decode_with_source(&malformed, Arc::clone(&file_source)) + .expect_err("missing required field must fail"); + assert!(err.to_string().contains(field), "unexpected error: {err}"); + } + + let mut missing_projection_expr = valid.clone(); + missing_projection_expr + .projection_exprs + .as_mut() + .expect("test config has projection expressions") + .projections[0] + .expr = None; + let err = serde + .decode_with_source(&missing_projection_expr, file_source) + .expect_err("missing projection expression must fail"); + assert!( + err.to_string() + .contains("ProjectionExpr missing expr field"), + "unexpected error: {err}" + ); + + Ok(()) + } + + #[test] + fn new_file_scan_config_decode_rejects_invalid_range_ordering() -> Result<()> { + let serde = FileScanSerdeHarness::new(); + let mut proto = serde.encode(&test_config(Some(range_partitioning())))?; + + let mut duplicate_ordering = proto.clone(); + let range = match duplicate_ordering + .output_partitioning + .as_mut() + .and_then(|p| p.partition_method.as_mut()) + { + Some(protobuf::partitioning::PartitionMethod::Range(range)) => range, + other => panic!("expected range partitioning, got {other:?}"), + }; + range.sort_expr.push(range.sort_expr[0].clone()); + + let err = serde + .decode(&duplicate_ordering) + .expect_err("duplicate range ordering must fail"); + assert!( + err.to_string().contains("duplicate expressions"), + "unexpected error: {err}" + ); + + let range = match proto + .output_partitioning + .as_mut() + .and_then(|p| p.partition_method.as_mut()) + { + Some(protobuf::partitioning::PartitionMethod::Range(range)) => range, + other => panic!("expected range partitioning, got {other:?}"), + }; + range.sort_expr.clear(); + + let err = serde + .decode(&proto) + .expect_err("empty range ordering must fail"); + assert!( + err.to_string().contains("requires non-empty ordering"), + "unexpected error: {err}" ); + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Unit tests for the bytes-only function serde exposed on + /// [`ExecutionPlanEncodeCtx`] / [`ExecutionPlanDecodeCtx`] and backed by + /// [`ConverterPlanEncoder`] / [`ConverterPlanDecoder`]. Function-carrying + /// plans migrate in follow-up PRs, so these paths have no in-tree plan + /// caller yet; the tests pin the payload semantics (`None` == encode by + /// name) and the decode lookup order (payload → codec; else registry → + /// codec fallback with an empty buffer) that those migrations rely on. + mod function_serde { + use super::*; + use arrow::datatypes::{DataType, Field, FieldRef}; + use datafusion_common::plan_err; + use datafusion_execution::config::SessionConfig; + use datafusion_execution::runtime_env::RuntimeEnv; + use datafusion_expr::function::AccumulatorArgs; + use datafusion_expr::{ + Accumulator, AggregateUDFImpl, ColumnarValue, PartitionEvaluator, + ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, WindowUDFImpl, + }; + use datafusion_functions_window_common::field::WindowUDFFieldArgs; + use datafusion_functions_window_common::partition::PartitionEvaluatorArgs; + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdf { + signature: Signature, + } + + impl TestUdf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64], + Volatility::Immutable, + ), + } + } + } + + impl ScalarUDFImpl for TestUdf { + fn name(&self) -> &str { + "test_udf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _args: &[DataType]) -> Result { + Ok(DataType::Int64) + } + fn invoke_with_args( + &self, + _args: ScalarFunctionArgs, + ) -> Result { + plan_err!("test only") + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdaf { + signature: Signature, + } + + impl TestUdaf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64], + Volatility::Immutable, + ), + } + } + } + + impl AggregateUDFImpl for TestUdaf { + fn name(&self) -> &str { + "test_udaf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } + fn accumulator( + &self, + _acc_args: AccumulatorArgs, + ) -> Result> { + plan_err!("test only") + } + } + + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestUdwf { + signature: Signature, + } + + impl TestUdwf { + fn new() -> Self { + Self { + signature: Signature::exact( + vec![DataType::Int64], + Volatility::Immutable, + ), + } + } + } + + impl WindowUDFImpl for TestUdwf { + fn name(&self) -> &str { + "test_udwf" + } + fn signature(&self) -> &Signature { + &self.signature + } + fn partition_evaluator( + &self, + _partition_evaluator_args: PartitionEvaluatorArgs, + ) -> Result> { + plan_err!("test only") + } + fn field(&self, field_args: WindowUDFFieldArgs) -> Result { + Ok(Field::new(field_args.name(), DataType::Int64, true).into()) + } + } + + /// Codec that encodes every function as its name bytes and decodes by + /// checking the payload it receives, so tests can observe exactly what + /// crosses the bytes-only boundary. + #[derive(Debug)] + struct PayloadCodec; + + impl PhysicalExtensionCodec for PayloadCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("not needed for these tests") + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + assert_eq!(name, "test_udf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(ScalarUDF::from(TestUdf::new()))) + } + + fn try_encode_udaf( + &self, + node: &AggregateUDF, + buf: &mut Vec, + ) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udaf( + &self, + name: &str, + buf: &[u8], + ) -> Result> { + assert_eq!(name, "test_udaf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(AggregateUDF::from(TestUdaf::new()))) + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + buf.extend_from_slice(node.name().as_bytes()); + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + assert_eq!(name, "test_udwf"); + assert_eq!(buf, name.as_bytes()); + Ok(Arc::new(WindowUDF::from(TestUdwf::new()))) + } + } + + /// Codec whose decode hooks only accept an empty payload, to pin the + /// by-name decode fallback (registry miss → codec with `&[]`). + #[derive(Debug)] + struct EmptyPayloadOnlyCodec; + + impl PhysicalExtensionCodec for EmptyPayloadOnlyCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not needed for these tests") + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("not needed for these tests") + } + + fn try_decode_udf(&self, _name: &str, buf: &[u8]) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(ScalarUDF::from(TestUdf::new()))) + } + + fn try_decode_udaf( + &self, + _name: &str, + buf: &[u8], + ) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(AggregateUDF::from(TestUdaf::new()))) + } + + fn try_decode_udwf(&self, _name: &str, buf: &[u8]) -> Result> { + assert!(buf.is_empty()); + Ok(Arc::new(WindowUDF::from(TestUdwf::new()))) + } + } + + fn encode_ctx_over<'a>( + codec: &'a dyn PhysicalExtensionCodec, + proto_converter: &'a dyn PhysicalProtoConverterExtension, + ) -> ConverterPlanEncoder<'a> { + ConverterPlanEncoder { + codec, + proto_converter, + } + } + + #[test] + fn encode_by_name_functions_produce_no_payload() -> Result<()> { + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DefaultPhysicalProtoConverter {}; + let encoder = encode_ctx_over(&codec, &converter); + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert!(ctx.encode_udf(&ScalarUDF::from(TestUdf::new()))?.is_none()); + assert!( + ctx.encode_udaf(&AggregateUDF::from(TestUdaf::new()))? + .is_none() + ); + assert!( + ctx.encode_udwf(&WindowUDF::from(TestUdwf::new()))? + .is_none() + ); + Ok(()) + } + + #[test] + fn encode_functions_surface_codec_payload() -> Result<()> { + let codec = PayloadCodec; + let converter = DefaultPhysicalProtoConverter {}; + let encoder = encode_ctx_over(&codec, &converter); + let ctx = ExecutionPlanEncodeCtx::new(&encoder); + + assert_eq!( + ctx.encode_udf(&ScalarUDF::from(TestUdf::new()))?.as_deref(), + Some(b"test_udf".as_slice()) + ); + assert_eq!( + ctx.encode_udaf(&AggregateUDF::from(TestUdaf::new()))? + .as_deref(), + Some(b"test_udaf".as_slice()) + ); + assert_eq!( + ctx.encode_udwf(&WindowUDF::from(TestUdwf::new()))? + .as_deref(), + Some(b"test_udwf".as_slice()) + ); + Ok(()) + } + + #[test] + fn decode_functions_prefer_explicit_payload() -> Result<()> { + let task_ctx = TaskContext::default(); + let codec = PayloadCodec; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert_eq!( + ctx.decode_udf("test_udf", Some(b"test_udf"))?.name(), + "test_udf" + ); + assert_eq!( + ctx.decode_udaf("test_udaf", Some(b"test_udaf"))?.name(), + "test_udaf" + ); + assert_eq!( + ctx.decode_udwf("test_udwf", Some(b"test_udwf"))?.name(), + "test_udwf" + ); + Ok(()) + } + + #[test] + fn decode_functions_by_name_resolve_from_registry() -> Result<()> { + let udf = Arc::new(ScalarUDF::from(TestUdf::new())); + let udaf = Arc::new(AggregateUDF::from(TestUdaf::new())); + let udwf = Arc::new(WindowUDF::from(TestUdwf::new())); + let task_ctx = TaskContext::new( + None, + "test".to_string(), + SessionConfig::new(), + HashMap::from([("test_udf".to_string(), Arc::clone(&udf))]), + HashMap::new(), + HashMap::from([("test_udaf".to_string(), Arc::clone(&udaf))]), + HashMap::from([("test_udwf".to_string(), Arc::clone(&udwf))]), + Arc::new(RuntimeEnv::default()), + ); + // The default codec fails any decode, so a success proves the + // registry satisfied the lookup without a codec fallback. + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert!(Arc::ptr_eq(&ctx.decode_udf("test_udf", None)?, &udf)); + assert!(Arc::ptr_eq(&ctx.decode_udaf("test_udaf", None)?, &udaf)); + assert!(Arc::ptr_eq(&ctx.decode_udwf("test_udwf", None)?, &udwf)); + assert_eq!(ctx.task_ctx().session_id(), "test"); + Ok(()) + } + + #[test] + fn decode_functions_by_name_fall_back_to_codec_on_registry_miss() -> Result<()> { + let task_ctx = TaskContext::default(); + let codec = EmptyPayloadOnlyCodec; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + assert_eq!(ctx.decode_udf("test_udf", None)?.name(), "test_udf"); + assert_eq!(ctx.decode_udaf("test_udaf", None)?.name(), "test_udaf"); + assert_eq!(ctx.decode_udwf("test_udwf", None)?.name(), "test_udwf"); + Ok(()) + } + + #[test] + fn decode_required_helpers_error_on_missing_fields() { + let task_ctx = TaskContext::default(); + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + let err = ctx + .decode_required_child(None, "FooExec", "input") + .unwrap_err(); + assert!( + err.to_string() + .contains("FooExec is missing required field 'input'"), + "unexpected error: {err}" + ); + + let schema = Schema::empty(); + let err = ctx + .decode_required_expr(None, &schema, "FooExec", "predicate") + .unwrap_err(); + assert!( + err.to_string() + .contains("FooExec is missing required field 'predicate'"), + "unexpected error: {err}" + ); + } + + #[test] + fn try_from_proto_rejects_wrong_plan_variant() { + let task_ctx = TaskContext::default(); + let codec = DefaultPhysicalExtensionCodec {}; + let decode_context = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let converter = DefaultPhysicalProtoConverter {}; + let decoder = ConverterPlanDecoder { + ctx: &decode_context, + proto_converter: &converter, + }; + let ctx = ExecutionPlanDecodeCtx::new(&decoder); + + let node = protobuf::PhysicalPlanNode { + physical_plan_type: None, + }; + let err = ProjectionExec::try_from_proto(&node, &ctx).unwrap_err(); + assert!( + err.to_string() + .contains("PhysicalPlanNode is not a ProjectionExec"), + "unexpected error: {err}" + ); + } } } @@ -270,8 +1025,24 @@ impl AsExecutionPlan for protobuf::PhysicalPlanNode { } } -impl protobuf::PhysicalPlanNode { - pub fn try_into_physical_plan_with_converter( +/// Extension methods on [`protobuf::PhysicalPlanNode`]. +/// +/// The prost-generated `PhysicalPlanNode` struct lives in +/// `datafusion-proto-models`, which is foreign to this crate, so the orphan +/// rule forbids inherent `impl` blocks here. Instead, all (de)serialization +/// helpers are exposed through this trait. Callers can bring it in scope with +/// `use datafusion_proto::physical_plan::PhysicalPlanNodeExt;`. +/// +/// Method bodies live in the default trait implementation. To make the trait +/// usable as if it were inherent (i.e. let bodies access fields on `self`), +/// implementors provide [`PhysicalPlanNodeExt::node`] returning a reference +/// back to the concrete `protobuf::PhysicalPlanNode`. Default method bodies +/// then go through `self.node()` to read fields. +pub trait PhysicalPlanNodeExt: Sized { + /// Returns a reference to the underlying [`protobuf::PhysicalPlanNode`]. + fn node(&self) -> &protobuf::PhysicalPlanNode; + + fn try_into_physical_plan_with_converter( &self, ctx: &TaskContext, codec: &dyn PhysicalExtensionCodec, @@ -281,3344 +1052,329 @@ impl protobuf::PhysicalPlanNode { self.try_into_physical_plan_with_context(&decode_ctx, proto_converter) } - pub(crate) fn try_into_physical_plan_with_context( + fn try_into_physical_plan_with_context( &self, ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - let plan = self.physical_plan_type.as_ref().ok_or_else(|| { + let plan = self.node().physical_plan_type.as_ref().ok_or_else(|| { proto_error(format!( - "physical_plan::from_proto() Unsupported physical plan '{self:?}'" + "physical_plan::from_proto() Unsupported physical plan '{:?}'", + self.node(), )) })?; + // Decode context for plans migrated to the `try_from_proto` pattern + // (#22419). Arms for migrated plans are one-liners delegating to the + // plan's own crate; un-migrated arms keep their inline bodies. + let plan_decoder = ConverterPlanDecoder { + ctx, + proto_converter, + }; + let decode_ctx = ExecutionPlanDecodeCtx::new(&plan_decoder); match plan { - PhysicalPlanType::Explain(explain) => { - self.try_into_explain_physical_plan(explain, ctx, proto_converter) + PhysicalPlanType::Explain(_) => { + ExplainExec::try_from_proto(self.node(), &decode_ctx) + } + PhysicalPlanType::Projection(_) => { + ProjectionExec::try_from_proto(self.node(), &decode_ctx) + } + PhysicalPlanType::Filter(_) => { + FilterExec::try_from_proto(self.node(), &decode_ctx) + } + PhysicalPlanType::CsvScan(_) => { + CsvSource::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Projection(projection) => { - self.try_into_projection_physical_plan(projection, ctx, proto_converter) + PhysicalPlanType::JsonScan(_) => { + JsonSource::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Filter(filter) => { - self.try_into_filter_physical_plan(filter, ctx, proto_converter) + PhysicalPlanType::ParquetScan(_) => { + #[cfg(feature = "parquet")] + { + ParquetSource::try_from_proto(self.node(), &decode_ctx) + } + #[cfg(not(feature = "parquet"))] + not_impl_err!( + "Unable to process a Parquet PhysicalPlan when the `parquet` feature is not enabled" + ) } - PhysicalPlanType::CsvScan(scan) => { - self.try_into_csv_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::AvroScan(_) => { + #[cfg(feature = "avro")] + { + AvroSource::try_from_proto(self.node(), &decode_ctx) + } + #[cfg(not(feature = "avro"))] + panic!( + "Unable to process a Avro PhysicalPlan when `avro` feature is not enabled" + ) } - PhysicalPlanType::JsonScan(scan) => { - self.try_into_json_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::MemoryScan(_) => { + MemorySourceConfig::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::ParquetScan(scan) => { - self.try_into_parquet_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::ArrowScan(_) => { + ArrowSource::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::AvroScan(scan) => { - self.try_into_avro_scan_physical_plan(scan, ctx, proto_converter) + #[expect( + deprecated, + reason = "`CoalesceBatchesExec` remains supported for protobuf compatibility" + )] + PhysicalPlanType::CoalesceBatches(_) => { + CoalesceBatchesExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::MemoryScan(scan) => { - self.try_into_memory_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::Merge(_) => { + CoalescePartitionsExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::ArrowScan(scan) => { - self.try_into_arrow_scan_physical_plan(scan, ctx, proto_converter) + PhysicalPlanType::Repartition(_) => { + RepartitionExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::CoalesceBatches(coalesce_batches) => self - .try_into_coalesce_batches_physical_plan( - coalesce_batches, - ctx, - proto_converter, - ), - PhysicalPlanType::Merge(merge) => { - self.try_into_merge_physical_plan(merge, ctx, proto_converter) + PhysicalPlanType::GlobalLimit(_) => { + GlobalLimitExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Repartition(repart) => { - self.try_into_repartition_physical_plan(repart, ctx, proto_converter) + PhysicalPlanType::LocalLimit(_) => { + LocalLimitExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::GlobalLimit(limit) => { - self.try_into_global_limit_physical_plan(limit, ctx, proto_converter) + PhysicalPlanType::Window(_) => { + WindowAggExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::LocalLimit(limit) => { - self.try_into_local_limit_physical_plan(limit, ctx, proto_converter) + PhysicalPlanType::Aggregate(_) => { + AggregateExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Window(window_agg) => { - self.try_into_window_physical_plan(window_agg, ctx, proto_converter) + PhysicalPlanType::HashJoin(_) => { + HashJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Aggregate(hash_agg) => { - self.try_into_aggregate_physical_plan(hash_agg, ctx, proto_converter) + PhysicalPlanType::SymmetricHashJoin(_) => { + SymmetricHashJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::HashJoin(hashjoin) => { - self.try_into_hash_join_physical_plan(hashjoin, ctx, proto_converter) + PhysicalPlanType::Union(_) => { + UnionExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::SymmetricHashJoin(sym_join) => self - .try_into_symmetric_hash_join_physical_plan( - sym_join, - ctx, - proto_converter, - ), - PhysicalPlanType::Union(union) => { - self.try_into_union_physical_plan(union, ctx, proto_converter) + PhysicalPlanType::Interleave(_) => { + InterleaveExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Interleave(interleave) => { - self.try_into_interleave_physical_plan(interleave, ctx, proto_converter) + PhysicalPlanType::CrossJoin(_) => { + CrossJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::CrossJoin(crossjoin) => { - self.try_into_cross_join_physical_plan(crossjoin, ctx, proto_converter) + PhysicalPlanType::Empty(_) => { + EmptyExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Empty(empty) => { - self.try_into_empty_physical_plan(empty, ctx, proto_converter) + PhysicalPlanType::PlaceholderRow(_) => { + PlaceholderRowExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::PlaceholderRow(placeholder) => { - self.try_into_placeholder_row_physical_plan(placeholder, ctx) + PhysicalPlanType::Sort(_) => { + SortExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Sort(sort) => { - self.try_into_sort_physical_plan(sort, ctx, proto_converter) + PhysicalPlanType::SortPreservingMerge(_) => { + SortPreservingMergeExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::SortPreservingMerge(sort) => self - .try_into_sort_preserving_merge_physical_plan(sort, ctx, proto_converter), PhysicalPlanType::Extension(extension) => { self.try_into_extension_physical_plan(extension, ctx, proto_converter) } - PhysicalPlanType::NestedLoopJoin(join) => { - self.try_into_nested_loop_join_physical_plan(join, ctx, proto_converter) + PhysicalPlanType::NestedLoopJoin(_) => { + NestedLoopJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Analyze(analyze) => { - self.try_into_analyze_physical_plan(analyze, ctx, proto_converter) + PhysicalPlanType::Analyze(_) => { + AnalyzeExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::JsonSink(sink) => { - self.try_into_json_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::JsonSink(_) => { + JsonSink::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::CsvSink(sink) => { - self.try_into_csv_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::CsvSink(_) => { + CsvSink::try_from_proto(self.node(), &decode_ctx) } - #[cfg_attr(not(feature = "parquet"), allow(unused_variables))] - PhysicalPlanType::ParquetSink(sink) => { - self.try_into_parquet_sink_physical_plan(sink, ctx, proto_converter) + PhysicalPlanType::ParquetSink(_) => { + #[cfg(feature = "parquet")] + { + ParquetSink::try_from_proto(self.node(), &decode_ctx) + } + #[cfg(not(feature = "parquet"))] + not_impl_err!("ParquetSink requires the `parquet` feature") } - PhysicalPlanType::Unnest(unnest) => { - self.try_into_unnest_physical_plan(unnest, ctx, proto_converter) + PhysicalPlanType::Unnest(_) => { + UnnestExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Cooperative(cooperative) => { - self.try_into_cooperative_physical_plan(cooperative, ctx, proto_converter) + PhysicalPlanType::Cooperative(_) => { + CooperativeExec::try_from_proto(self.node(), &decode_ctx) } PhysicalPlanType::GenerateSeries(generate_series) => { self.try_into_generate_series_physical_plan(generate_series) } - PhysicalPlanType::SortMergeJoin(sort_join) => { - self.try_into_sort_join(sort_join, ctx, proto_converter) - } - PhysicalPlanType::AsyncFunc(async_func) => { - self.try_into_async_func_physical_plan(async_func, ctx, proto_converter) + PhysicalPlanType::SortMergeJoin(_) => { + SortMergeJoinExec::try_from_proto(self.node(), &decode_ctx) } - PhysicalPlanType::Buffer(buffer) => { - self.try_into_buffer_physical_plan(buffer, ctx, proto_converter) - } - PhysicalPlanType::ScalarSubquery(sq) => { - self.try_into_scalar_subquery_physical_plan(sq, ctx, proto_converter) - } - } - } - - pub fn try_from_physical_plan_with_converter( - plan: Arc, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result - where - Self: Sized, - { - let plan_clone = Arc::clone(&plan); - let plan = plan.as_ref() as &dyn Any; - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_explain_exec(exec, codec); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_projection_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_analyze_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_filter_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(limit) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_global_limit_exec( - limit, - codec, - proto_converter, - ); - } - - if let Some(limit) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_local_limit_exec( - limit, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_hash_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_symmetric_hash_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_merge_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_cross_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_aggregate_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(empty) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_empty_exec(empty, codec); - } - - if let Some(empty) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_placeholder_row_exec( - empty, codec, - ); - } - - #[expect(deprecated)] - if let Some(coalesce_batches) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_coalesce_batches_exec( - coalesce_batches, - codec, - proto_converter, - ); - } - - if let Some(data_source_exec) = plan.downcast_ref::() - && let Some(node) = protobuf::PhysicalPlanNode::try_from_data_source_exec( - data_source_exec, - codec, - proto_converter, - )? - { - return Ok(node); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_coalesce_partitions_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_repartition_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(union) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_union_exec( - union, - codec, - proto_converter, - ); - } - - if let Some(interleave) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_interleave_exec( - interleave, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_sort_preserving_merge_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_nested_loop_join_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_window_agg_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_bounded_window_agg_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() - && let Some(node) = protobuf::PhysicalPlanNode::try_from_data_sink_exec( - exec, - codec, - proto_converter, - )? - { - return Ok(node); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_unnest_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_cooperative_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() - && let Some(node) = - protobuf::PhysicalPlanNode::try_from_lazy_memory_exec(exec)? - { - return Ok(node); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_async_func_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_buffer_exec( - exec, - codec, - proto_converter, - ); - } - - if let Some(exec) = plan.downcast_ref::() { - return protobuf::PhysicalPlanNode::try_from_scalar_subquery_exec( - exec, - codec, - proto_converter, - ); - } - - let mut buf: Vec = vec![]; - match codec.try_encode(Arc::clone(&plan_clone), &mut buf) { - Ok(_) => { - let inputs: Vec = plan_clone - .children() - .into_iter() - .cloned() - .map(|i| { - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - i, - codec, - proto_converter, - ) - }) - .collect::>()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Extension( - protobuf::PhysicalExtensionNode { node: buf, inputs }, - )), - }) - } - Err(e) => internal_err!( - "Unsupported plan and extension codec failed with [{e}]. Plan: {plan_clone:?}" - ), - } - } -} - -impl protobuf::PhysicalPlanNode { - fn try_into_explain_physical_plan( - &self, - explain: &protobuf::ExplainExecNode, - _ctx: &PhysicalPlanDecodeContext<'_>, - _proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - Ok(Arc::new(ExplainExec::new( - Arc::new(explain.schema.as_ref().unwrap().try_into()?), - explain - .stringified_plans - .iter() - .map(|plan| plan.into()) - .collect(), - explain.verbose, - ))) - } - - fn try_into_projection_physical_plan( - &self, - projection: &protobuf::ProjectionExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&projection.input, ctx, proto_converter)?; - let exprs = projection - .expr - .iter() - .zip(projection.expr_name.iter()) - .map(|(expr, name)| { - Ok(( - proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - name.to_string(), - )) - }) - .collect::, String)>>>()?; - let proj_exprs: Vec = exprs - .into_iter() - .map(|(expr, alias)| ProjectionExpr { expr, alias }) - .collect(); - Ok(Arc::new(ProjectionExec::try_new(proj_exprs, input)?)) - } - - fn try_into_filter_physical_plan( - &self, - filter: &protobuf::FilterExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&filter.input, ctx, proto_converter)?; - - let predicate = filter - .expr - .as_ref() - .map(|expr| { - proto_converter.proto_to_physical_expr(expr, input.schema().as_ref(), ctx) - }) - .transpose()? - .ok_or_else(|| { - internal_datafusion_err!( - "filter (FilterExecNode) in PhysicalPlanNode is missing." - ) - })?; - - let filter_selectivity = filter.default_filter_selectivity.try_into(); - // Preserve the `None` state across proto boundaries. Proto cannot distinguish - // between `None` (full projection) and `Some(vec![])` (empty projection) since - // both serialize as an empty list. If all columns are included, we reconstruct - // `None` to avoid losing this semantic distinction on deserialization. - let num_fields = input.schema().fields().len(); - let mut is_full_projection = filter.projection.len() == num_fields; - let mut projection_vec: Vec = Vec::with_capacity(filter.projection.len()); - for (i, idx) in filter.projection.iter().enumerate() { - let idx = *idx as usize; - is_full_projection &= idx == i; - projection_vec.push(idx); - } - let projection = if is_full_projection { - None - } else { - Some(projection_vec) - }; - let filter = FilterExecBuilder::new(predicate, input) - .apply_projection(projection)? - .with_batch_size(filter.batch_size as usize) - .with_fetch(filter.fetch.map(|f| f as usize)) - .build()?; - match filter_selectivity { - Ok(filter_selectivity) => Ok(Arc::new( - filter.with_default_selectivity(filter_selectivity)?, - )), - Err(_) => Err(internal_datafusion_err!( - "filter_selectivity in PhysicalPlanNode is invalid " - )), - } - } - - fn try_into_csv_scan_physical_plan( - &self, - scan: &protobuf::CsvScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let escape = - if let Some(protobuf::csv_scan_exec_node::OptionalEscape::Escape(escape)) = - &scan.optional_escape - { - Some(str_to_byte(escape, "escape")?) - } else { - None - }; - - let comment = if let Some( - protobuf::csv_scan_exec_node::OptionalComment::Comment(comment), - ) = &scan.optional_comment - { - Some(str_to_byte(comment, "comment")?) - } else { - None - }; - - // Parse table schema with partition columns - let table_schema = - parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?; - - let csv_options = CsvOptions { - has_header: Some(scan.has_header), - delimiter: str_to_byte(&scan.delimiter, "delimiter")?, - quote: str_to_byte(&scan.quote, "quote")?, - newlines_in_values: Some(scan.newlines_in_values), - ..Default::default() - }; - let source = Arc::new( - CsvSource::new(table_schema) - .with_csv_options(csv_options) - .with_escape(escape) - .with_comment(comment), - ); - - let conf = FileScanConfigBuilder::from(parse_protobuf_file_scan_config( - scan.base_conf.as_ref().unwrap(), - ctx, - proto_converter, - source, - )?) - .with_file_compression_type(FileCompressionType::UNCOMPRESSED) - .build(); - Ok(DataSourceExec::from_data_source(conf)) - } - - fn try_into_json_scan_physical_plan( - &self, - scan: &protobuf::JsonScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let base_conf = scan.base_conf.as_ref().unwrap(); - let table_schema = parse_table_schema_from_proto(base_conf)?; - let scan_conf = parse_protobuf_file_scan_config( - base_conf, - ctx, - proto_converter, - Arc::new(JsonSource::new(table_schema)), - )?; - Ok(DataSourceExec::from_data_source(scan_conf)) - } - - fn try_into_arrow_scan_physical_plan( - &self, - scan: &protobuf::ArrowScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let base_conf = scan.base_conf.as_ref().ok_or_else(|| { - internal_datafusion_err!("base_conf in ArrowScanExecNode is missing.") - })?; - let table_schema = parse_table_schema_from_proto(base_conf)?; - let scan_conf = parse_protobuf_file_scan_config( - base_conf, - ctx, - proto_converter, - Arc::new(ArrowSource::new_file_source(table_schema)), - )?; - Ok(DataSourceExec::from_data_source(scan_conf)) - } - - #[cfg_attr(not(feature = "parquet"), expect(unused_variables))] - fn try_into_parquet_scan_physical_plan( - &self, - scan: &protobuf::ParquetScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - #[cfg(feature = "parquet")] - { - let schema = from_proto::parse_protobuf_file_scan_schema( - scan.base_conf.as_ref().unwrap(), - )?; - - // Check if there's a projection and use projected schema for predicate parsing - let base_conf = scan.base_conf.as_ref().unwrap(); - let predicate_schema = if !base_conf.projection.is_empty() { - // Create projected schema for parsing the predicate - let projected_fields: Vec<_> = base_conf - .projection - .iter() - .map(|&i| schema.field(i as usize).clone()) - .collect(); - Arc::new(Schema::new(projected_fields)) - } else { - schema - }; - - let predicate = scan - .predicate - .as_ref() - .map(|expr| { - proto_converter.proto_to_physical_expr( - expr, - predicate_schema.as_ref(), - ctx, - ) - }) - .transpose()?; - let mut options = datafusion_common::config::TableParquetOptions::default(); - - if let Some(table_options) = scan.parquet_options.as_ref() { - options = table_options.try_into()?; - } - - // Parse table schema with partition columns - let table_schema = parse_table_schema_from_proto(base_conf)?; - let object_store_url = match base_conf.object_store_url.is_empty() { - false => ObjectStoreUrl::parse(&base_conf.object_store_url)?, - true => ObjectStoreUrl::local_filesystem(), - }; - let store = ctx - .task_ctx() - .runtime_env() - .object_store(object_store_url)?; - let metadata_cache = ctx - .task_ctx() - .runtime_env() - .cache_manager - .get_file_metadata_cache(); - let reader_factory = - Arc::new(CachedParquetFileReaderFactory::new(store, metadata_cache)); - - let mut source = ParquetSource::new(table_schema) - .with_parquet_file_reader_factory(reader_factory) - .with_table_parquet_options(options); - - if let Some(predicate) = predicate { - source = source.with_predicate(predicate); - } - let base_config = parse_protobuf_file_scan_config( - base_conf, - ctx, - proto_converter, - Arc::new(source), - )?; - Ok(DataSourceExec::from_data_source(base_config)) - } - #[cfg(not(feature = "parquet"))] - panic!( - "Unable to process a Parquet PhysicalPlan when `parquet` feature is not enabled" - ) - } - - #[cfg_attr(not(feature = "avro"), expect(unused_variables))] - fn try_into_avro_scan_physical_plan( - &self, - scan: &protobuf::AvroScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - #[cfg(feature = "avro")] - { - let table_schema = - parse_table_schema_from_proto(scan.base_conf.as_ref().unwrap())?; - let conf = parse_protobuf_file_scan_config( - scan.base_conf.as_ref().unwrap(), - ctx, - proto_converter, - Arc::new(AvroSource::new(table_schema)), - )?; - Ok(DataSourceExec::from_data_source(conf)) - } - - #[cfg(not(feature = "avro"))] - panic!("Unable to process a Avro PhysicalPlan when `avro` feature is not enabled") - } - - fn try_into_memory_scan_physical_plan( - &self, - scan: &protobuf::MemoryScanExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let partitions = scan - .partitions - .iter() - .map(|p| parse_record_batches(p)) - .collect::>>()?; - - let proto_schema = scan.schema.as_ref().ok_or_else(|| { - internal_datafusion_err!("schema in MemoryScanExecNode is missing.") - })?; - let schema: SchemaRef = SchemaRef::new(proto_schema.try_into()?); - - let projection = if !scan.projection.is_empty() { - Some( - scan.projection - .iter() - .map(|i| *i as usize) - .collect::>(), - ) - } else { - None - }; - - let mut sort_information = vec![]; - for ordering in &scan.sort_information { - let sort_exprs = parse_physical_sort_exprs( - &ordering.physical_sort_expr_nodes, - ctx, - &schema, - proto_converter, - )?; - sort_information.extend(LexOrdering::new(sort_exprs)); - } - - let source = MemorySourceConfig::try_new(&partitions, schema, projection)? - .with_limit(scan.fetch.map(|f| f as usize)) - .with_show_sizes(scan.show_sizes); - - let source = source.try_with_sort_information(sort_information)?; - - Ok(DataSourceExec::from_data_source(source)) - } - - fn try_into_coalesce_batches_physical_plan( - &self, - coalesce_batches: &protobuf::CoalesceBatchesExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&coalesce_batches.input, ctx, proto_converter)?; - Ok(Arc::new( - #[expect(deprecated)] - CoalesceBatchesExec::new(input, coalesce_batches.target_batch_size as usize) - .with_fetch(coalesce_batches.fetch.map(|f| f as usize)), - )) - } - - fn try_into_merge_physical_plan( - &self, - merge: &protobuf::CoalescePartitionsExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&merge.input, ctx, proto_converter)?; - Ok(Arc::new( - CoalescePartitionsExec::new(input) - .with_fetch(merge.fetch.map(|f| f as usize)), - )) - } - - fn try_into_repartition_physical_plan( - &self, - repart: &protobuf::RepartitionExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&repart.input, ctx, proto_converter)?; - let partitioning = parse_protobuf_partitioning( - repart.partitioning.as_ref(), - ctx, - input.schema().as_ref(), - proto_converter, - )?; - let mut repart_exec = RepartitionExec::try_new(input, partitioning.unwrap())?; - if repart.preserve_order { - repart_exec = repart_exec.with_preserve_order(); - } - Ok(Arc::new(repart_exec)) - } - - fn try_into_global_limit_physical_plan( - &self, - limit: &protobuf::GlobalLimitExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&limit.input, ctx, proto_converter)?; - let fetch = if limit.fetch >= 0 { - Some(limit.fetch as usize) - } else { - None - }; - Ok(Arc::new(GlobalLimitExec::new( - input, - limit.skip as usize, - fetch, - ))) - } - - fn try_into_local_limit_physical_plan( - &self, - limit: &protobuf::LocalLimitExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&limit.input, ctx, proto_converter)?; - Ok(Arc::new(LocalLimitExec::new(input, limit.fetch as usize))) - } - - fn try_into_window_physical_plan( - &self, - window_agg: &protobuf::WindowAggExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&window_agg.input, ctx, proto_converter)?; - let input_schema = input.schema(); - - let physical_window_expr: Vec> = window_agg - .window_expr - .iter() - .map(|window_expr| { - parse_physical_window_expr( - window_expr, - ctx, - input_schema.as_ref(), - proto_converter, - ) - }) - .collect::, _>>()?; - - let partition_keys = window_agg - .partition_keys - .iter() - .map(|expr| { - proto_converter.proto_to_physical_expr(expr, input.schema().as_ref(), ctx) - }) - .collect::>>>()?; - - if let Some(input_order_mode) = window_agg.input_order_mode.as_ref() { - let input_order_mode = match input_order_mode { - window_agg_exec_node::InputOrderMode::Linear(_) => InputOrderMode::Linear, - window_agg_exec_node::InputOrderMode::PartiallySorted( - protobuf::PartiallySortedInputOrderMode { columns }, - ) => InputOrderMode::PartiallySorted( - columns.iter().map(|c| *c as usize).collect(), - ), - window_agg_exec_node::InputOrderMode::Sorted(_) => InputOrderMode::Sorted, - }; - - Ok(Arc::new(BoundedWindowAggExec::try_new( - physical_window_expr, - input, - input_order_mode, - !partition_keys.is_empty(), - )?)) - } else { - Ok(Arc::new(WindowAggExec::try_new( - physical_window_expr, - input, - !partition_keys.is_empty(), - )?)) - } - } - - fn try_into_aggregate_physical_plan( - &self, - hash_agg: &protobuf::AggregateExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&hash_agg.input, ctx, proto_converter)?; - let mode = protobuf::AggregateMode::try_from(hash_agg.mode).map_err(|_| { - proto_error(format!( - "Received a AggregateNode message with unknown AggregateMode {}", - hash_agg.mode - )) - })?; - let agg_mode: AggregateMode = match mode { - protobuf::AggregateMode::Partial => AggregateMode::Partial, - protobuf::AggregateMode::Final => AggregateMode::Final, - protobuf::AggregateMode::FinalPartitioned => AggregateMode::FinalPartitioned, - protobuf::AggregateMode::Single => AggregateMode::Single, - protobuf::AggregateMode::SinglePartitioned => { - AggregateMode::SinglePartitioned - } - protobuf::AggregateMode::PartialReduce => AggregateMode::PartialReduce, - }; - - let num_expr = hash_agg.group_expr.len(); - - let group_expr = hash_agg - .group_expr - .iter() - .zip(hash_agg.group_expr_name.iter()) - .map(|(expr, name)| { - proto_converter - .proto_to_physical_expr(expr, input.schema().as_ref(), ctx) - .map(|expr| (expr, name.to_string())) - }) - .collect::, _>>()?; - - let null_expr = hash_agg - .null_expr - .iter() - .zip(hash_agg.group_expr_name.iter()) - .map(|(expr, name)| { - proto_converter - .proto_to_physical_expr(expr, input.schema().as_ref(), ctx) - .map(|expr| (expr, name.to_string())) - }) - .collect::, _>>()?; - - let groups: Vec> = if !hash_agg.groups.is_empty() { - hash_agg - .groups - .chunks(num_expr) - .map(|g| g.to_vec()) - .collect::>>() - } else { - vec![] - }; - - let has_grouping_set = hash_agg.has_grouping_set; - - let input_schema = hash_agg.input_schema.as_ref().ok_or_else(|| { - internal_datafusion_err!("input_schema in AggregateNode is missing.") - })?; - let physical_schema: SchemaRef = SchemaRef::new(input_schema.try_into()?); - - let physical_filter_expr = hash_agg - .filter_expr - .iter() - .map(|expr| { - expr.expr - .as_ref() - .map(|e| { - proto_converter.proto_to_physical_expr(e, &physical_schema, ctx) - }) - .transpose() - }) - .collect::, _>>()?; - - let physical_aggr_expr: Vec> = hash_agg - .aggr_expr - .iter() - .zip(hash_agg.aggr_expr_name.iter()) - .map(|(expr, name)| { - let expr_type = expr.expr_type.as_ref().ok_or_else(|| { - proto_error("Unexpected empty aggregate physical expression") - })?; - - match expr_type { - ExprType::AggregateExpr(agg_node) => { - let input_phy_expr: Vec> = agg_node - .expr - .iter() - .map(|e| { - proto_converter.proto_to_physical_expr( - e, - &physical_schema, - ctx, - ) - }) - .collect::>>()?; - let order_bys = agg_node - .ordering_req - .iter() - .map(|e| { - parse_physical_sort_expr( - e, - ctx, - &physical_schema, - proto_converter, - ) - }) - .collect::>()?; - agg_node - .aggregate_function - .as_ref() - .map(|func| match func { - AggregateFunction::UserDefinedAggrFunction(udaf_name) => { - let agg_udf = match &agg_node.fun_definition { - Some(buf) => { - ctx.codec().try_decode_udaf(udaf_name, buf)? - } - None => ctx.task_ctx().udaf(udaf_name).or_else( - |_| { - ctx.codec() - .try_decode_udaf(udaf_name, &[]) - }, - )?, - }; - - let (human_display, human_display_alias) = - split_human_display_alias( - &agg_node.human_display, - name, - ); - let builder = AggregateExprBuilder::new( - agg_udf, - input_phy_expr, - ) - .schema(Arc::clone(&physical_schema)) - .alias(name) - .with_ignore_nulls(agg_node.ignore_nulls) - .with_distinct(agg_node.distinct) - .order_by(order_bys) - .human_display(human_display); - let builder = if let Some(alias) = human_display_alias - { - builder.human_display_alias(alias) - } else { - builder - }; - builder.build().map(Arc::new) - } - }) - .transpose()? - .ok_or_else(|| { - proto_error( - "Invalid AggregateExpr, missing aggregate_function", - ) - }) - } - _ => internal_err!("Invalid aggregate expression for AggregateExec"), - } - }) - .collect::, _>>()?; - - let physical_schema_ref = Arc::clone(&physical_schema); - let agg = AggregateExec::try_new( - agg_mode, - PhysicalGroupBy::new(group_expr, null_expr, groups, has_grouping_set), - physical_aggr_expr, - physical_filter_expr, - input, - physical_schema, - )?; - - let agg = if let Some(limit_proto) = &hash_agg.limit { - let limit = limit_proto.limit as usize; - let limit_options = match limit_proto.descending { - Some(descending) => LimitOptions::new_with_order(limit, descending), - None => LimitOptions::new(limit), - }; - agg.with_limit_options(Some(limit_options)) - } else { - agg - }; - - let agg = if let Some(dynamic_filter_proto) = &hash_agg.dynamic_filter { - let dynamic_filter_expr = proto_converter.proto_to_physical_expr( - dynamic_filter_proto, - physical_schema_ref.as_ref(), - ctx, - )?; - let df = (dynamic_filter_expr as Arc) - .downcast::() - .map_err(|_| { - internal_datafusion_err!( - "AggregateExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" - ) - })?; - agg.with_dynamic_filter_expr(df)? - } else { - agg - }; - - Ok(Arc::new(agg)) - } - - fn try_into_hash_join_physical_plan( - &self, - hashjoin: &protobuf::HashJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left: Arc = - into_physical_plan(&hashjoin.left, ctx, proto_converter)?; - let right: Arc = - into_physical_plan(&hashjoin.right, ctx, proto_converter)?; - let left_schema = left.schema(); - let right_schema = right.schema(); - let on: Vec<(PhysicalExprRef, PhysicalExprRef)> = hashjoin - .on - .iter() - .map(|col| { - let left = proto_converter.proto_to_physical_expr( - &col.left.clone().unwrap(), - left_schema.as_ref(), - ctx, - )?; - let right = proto_converter.proto_to_physical_expr( - &col.right.clone().unwrap(), - right_schema.as_ref(), - ctx, - )?; - Ok((left, right)) - }) - .collect::>()?; - let join_type = - protobuf::JoinType::try_from(hashjoin.join_type).map_err(|_| { - proto_error(format!( - "Received a HashJoinNode message with unknown JoinType {}", - hashjoin.join_type - )) - })?; - let null_equality = protobuf::NullEquality::try_from(hashjoin.null_equality) - .map_err(|_| { - proto_error(format!( - "Received a HashJoinNode message with unknown NullEquality {}", - hashjoin.null_equality - )) - })?; - let filter = hashjoin - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter.proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f.column_indices - .iter() - .map(|i| { - let side = protobuf::JoinSide::try_from(i.side) - .map_err(|_| proto_error(format!( - "Received a HashJoinNode message with JoinSide in Filter {}", - i.side)) - )?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>>()?; - - Ok(JoinFilter::new(expression, column_indices, Arc::new(schema))) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let partition_mode = protobuf::PartitionMode::try_from(hashjoin.partition_mode) - .map_err(|_| { - proto_error(format!( - "Received a HashJoinNode message with unknown PartitionMode {}", - hashjoin.partition_mode - )) - })?; - let partition_mode = match partition_mode { - protobuf::PartitionMode::CollectLeft => PartitionMode::CollectLeft, - protobuf::PartitionMode::Partitioned => PartitionMode::Partitioned, - protobuf::PartitionMode::Auto => PartitionMode::Auto, - }; - let projection = if !hashjoin.projection.is_empty() { - Some( - hashjoin - .projection - .iter() - .map(|i| *i as usize) - .collect::>(), - ) - } else { - None - }; - let mut hash_join = HashJoinExec::try_new( - left, - right, - on, - filter, - &join_type.into(), - projection, - partition_mode, - null_equality.into(), - hashjoin.null_aware, - )?; - - if let Some(dynamic_filter_proto) = &hashjoin.dynamic_filter { - let dynamic_filter_expr = proto_converter.proto_to_physical_expr( - dynamic_filter_proto, - right_schema.as_ref(), - ctx, - )?; - let df = (dynamic_filter_expr as Arc) - .downcast::() - .map_err(|_| { - internal_datafusion_err!( - "HashJoinExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" - ) - })?; - hash_join = hash_join.with_dynamic_filter_expr(df)?; - } - - Ok(Arc::new(hash_join)) - } - - fn try_into_symmetric_hash_join_physical_plan( - &self, - sym_join: &protobuf::SymmetricHashJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left = into_physical_plan(&sym_join.left, ctx, proto_converter)?; - let right = into_physical_plan(&sym_join.right, ctx, proto_converter)?; - let left_schema = left.schema(); - let right_schema = right.schema(); - let on = sym_join - .on - .iter() - .map(|col| { - let left = proto_converter.proto_to_physical_expr( - &col.left.clone().unwrap(), - left_schema.as_ref(), - ctx, - )?; - let right = proto_converter.proto_to_physical_expr( - &col.right.clone().unwrap(), - right_schema.as_ref(), - ctx, - )?; - Ok((left, right)) - }) - .collect::>()?; - let join_type = - protobuf::JoinType::try_from(sym_join.join_type).map_err(|_| { - proto_error(format!( - "Received a SymmetricHashJoin message with unknown JoinType {}", - sym_join.join_type - )) - })?; - let null_equality = protobuf::NullEquality::try_from(sym_join.null_equality) - .map_err(|_| { - proto_error(format!( - "Received a SymmetricHashJoin message with unknown NullEquality {}", - sym_join.null_equality - )) - })?; - let filter = sym_join - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter.proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f.column_indices - .iter() - .map(|i| { - let side = protobuf::JoinSide::try_from(i.side) - .map_err(|_| proto_error(format!( - "Received a HashJoinNode message with JoinSide in Filter {}", - i.side)) - )?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>()?; - - Ok(JoinFilter::new(expression, column_indices, Arc::new(schema))) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let left_sort_exprs = parse_physical_sort_exprs( - &sym_join.left_sort_exprs, - ctx, - &left_schema, - proto_converter, - )?; - let left_sort_exprs = LexOrdering::new(left_sort_exprs); - - let right_sort_exprs = parse_physical_sort_exprs( - &sym_join.right_sort_exprs, - ctx, - &right_schema, - proto_converter, - )?; - let right_sort_exprs = LexOrdering::new(right_sort_exprs); - - let partition_mode = protobuf::StreamPartitionMode::try_from( - sym_join.partition_mode, - ) - .map_err(|_| { - proto_error(format!( - "Received a SymmetricHashJoin message with unknown PartitionMode {}", - sym_join.partition_mode - )) - })?; - let partition_mode = match partition_mode { - protobuf::StreamPartitionMode::SinglePartition => { - StreamJoinPartitionMode::SinglePartition - } - protobuf::StreamPartitionMode::PartitionedExec => { - StreamJoinPartitionMode::Partitioned - } - }; - SymmetricHashJoinExec::try_new( - left, - right, - on, - filter, - &join_type.into(), - null_equality.into(), - left_sort_exprs, - right_sort_exprs, - partition_mode, - ) - .map(|e| Arc::new(e) as _) - } - - fn try_into_union_physical_plan( - &self, - union: &protobuf::UnionExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let mut inputs: Vec> = vec![]; - for input in &union.inputs { - inputs.push(proto_converter.proto_to_execution_plan(input, ctx)?); - } - UnionExec::try_new(inputs) - } - - fn try_into_interleave_physical_plan( - &self, - interleave: &protobuf::InterleaveExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let mut inputs: Vec> = vec![]; - for input in &interleave.inputs { - inputs.push(proto_converter.proto_to_execution_plan(input, ctx)?); - } - Ok(Arc::new(InterleaveExec::try_new(inputs)?)) - } - - fn try_into_cross_join_physical_plan( - &self, - crossjoin: &protobuf::CrossJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left: Arc = - into_physical_plan(&crossjoin.left, ctx, proto_converter)?; - let right: Arc = - into_physical_plan(&crossjoin.right, ctx, proto_converter)?; - Ok(Arc::new(CrossJoinExec::new(left, right))) - } - - fn try_into_empty_physical_plan( - &self, - empty: &protobuf::EmptyExecNode, - _ctx: &PhysicalPlanDecodeContext<'_>, - _proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let schema = Arc::new(convert_required!(empty.schema)?); - Ok(Arc::new(EmptyExec::new(schema))) - } - - fn try_into_placeholder_row_physical_plan( - &self, - placeholder: &protobuf::PlaceholderRowExecNode, - _ctx: &PhysicalPlanDecodeContext<'_>, - ) -> Result> { - let schema = Arc::new(convert_required!(placeholder.schema)?); - Ok(Arc::new(PlaceholderRowExec::new(schema))) - } - - fn try_into_sort_physical_plan( - &self, - sort: &protobuf::SortExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&sort.input, ctx, proto_converter)?; - let exprs = sort - .expr - .iter() - .map(|expr| { - let expr = expr.expr_type.as_ref().ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected expr {self:?}" - )) - })?; - if let ExprType::Sort(sort_expr) = expr { - let expr = sort_expr - .expr - .as_ref() - .ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected sort expr {self:?}" - )) - })? - .as_ref(); - Ok(PhysicalSortExpr { - expr: proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - } else { - internal_err!( - "physical_plan::from_proto() {self:?}" - ) - } - }) - .collect::>>()?; - let Some(ordering) = LexOrdering::new(exprs) else { - return internal_err!("SortExec requires an ordering"); - }; - let fetch = (sort.fetch >= 0).then_some(sort.fetch as _); - let new_sort = SortExec::new(ordering, input) - .with_fetch(fetch) - .with_preserve_partitioning(sort.preserve_partitioning); - - let new_sort = if let Some(dynamic_filter_proto) = &sort.dynamic_filter { - let dynamic_filter_expr = proto_converter.proto_to_physical_expr( - dynamic_filter_proto, - new_sort.input().schema().as_ref(), - ctx, - )?; - let df = (dynamic_filter_expr as Arc) - .downcast::() - .map_err(|_| { - internal_datafusion_err!( - "SortExec dynamic_filter did not decode to a DynamicFilterPhysicalExpr" - ) - })?; - new_sort.with_dynamic_filter_expr(df)? - } else { - new_sort - }; - - Ok(Arc::new(new_sort)) - } - - fn try_into_sort_preserving_merge_physical_plan( - &self, - sort: &protobuf::SortPreservingMergeExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&sort.input, ctx, proto_converter)?; - let exprs = sort - .expr - .iter() - .map(|expr| { - let expr = expr.expr_type.as_ref().ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected expr {self:?}" - )) - })?; - if let ExprType::Sort(sort_expr) = expr { - let expr = sort_expr - .expr - .as_ref() - .ok_or_else(|| { - proto_error(format!( - "physical_plan::from_proto() Unexpected sort expr {self:?}" - )) - })? - .as_ref(); - Ok(PhysicalSortExpr { - expr: proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?, - options: SortOptions { - descending: !sort_expr.asc, - nulls_first: sort_expr.nulls_first, - }, - }) - } else { - internal_err!("physical_plan::from_proto() {self:?}") - } - }) - .collect::>>()?; - let Some(ordering) = LexOrdering::new(exprs) else { - return internal_err!("SortExec requires an ordering"); - }; - let fetch = (sort.fetch >= 0).then_some(sort.fetch as _); - Ok(Arc::new( - SortPreservingMergeExec::new(ordering, input).with_fetch(fetch), - )) - } - - fn try_into_extension_physical_plan( - &self, - extension: &protobuf::PhysicalExtensionNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let inputs: Vec> = extension - .inputs - .iter() - .map(|i| proto_converter.proto_to_execution_plan(i, ctx)) - .collect::>()?; - - let extension_node = - ctx.codec() - .try_decode(extension.node.as_slice(), &inputs, ctx.task_ctx())?; - - Ok(extension_node) - } - - fn try_into_nested_loop_join_physical_plan( - &self, - join: &protobuf::NestedLoopJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left: Arc = - into_physical_plan(&join.left, ctx, proto_converter)?; - let right: Arc = - into_physical_plan(&join.right, ctx, proto_converter)?; - let join_type = protobuf::JoinType::try_from(join.join_type).map_err(|_| { - proto_error(format!( - "Received a NestedLoopJoinExecNode message with unknown JoinType {}", - join.join_type - )) - })?; - let filter = join - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter - .proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f.column_indices - .iter() - .map(|i| { - let side = protobuf::JoinSide::try_from(i.side) - .map_err(|_| proto_error(format!( - "Received a NestedLoopJoinExecNode message with JoinSide in Filter {}", - i.side)) - )?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>>()?; - - Ok(JoinFilter::new(expression, column_indices, Arc::new(schema))) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let projection = if !join.projection.is_empty() { - Some( - join.projection - .iter() - .map(|i| *i as usize) - .collect::>(), - ) - } else { - None - }; - - Ok(Arc::new(NestedLoopJoinExec::try_new( - left, - right, - filter, - &join_type.into(), - projection, - )?)) - } - - fn try_into_analyze_physical_plan( - &self, - analyze: &protobuf::AnalyzeExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&analyze.input, ctx, proto_converter)?; - let metric_categories = if analyze.has_metric_categories { - let cats: Result> = analyze - .metric_categories - .iter() - .map(|s| s.parse::()) - .collect(); - Some(cats?) - } else { - None - }; - Ok(Arc::new(AnalyzeExec::new( - analyze.verbose, - analyze.show_statistics, - vec![MetricType::Summary, MetricType::Dev], - metric_categories, - input, - Arc::new(convert_required!(analyze.schema)?), - ))) - } - - fn try_into_json_sink_physical_plan( - &self, - sink: &protobuf::JsonSinkExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink: JsonSink = sink - .sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))? - .try_into()?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) - } - - fn try_into_csv_sink_physical_plan( - &self, - sink: &protobuf::CsvSinkExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink: CsvSink = sink - .sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))? - .try_into()?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) - } - - #[cfg_attr(not(feature = "parquet"), expect(unused_variables))] - fn try_into_parquet_sink_physical_plan( - &self, - sink: &protobuf::ParquetSinkExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - #[cfg(feature = "parquet")] - { - let input = into_physical_plan(&sink.input, ctx, proto_converter)?; - - let data_sink: ParquetSink = sink - .sink - .as_ref() - .ok_or_else(|| proto_error("Missing required field in protobuf"))? - .try_into()?; - let sink_schema = input.schema(); - let sort_order = sink - .sort_order - .as_ref() - .map(|collection| { - parse_physical_sort_exprs( - &collection.physical_sort_expr_nodes, - ctx, - &sink_schema, - proto_converter, - ) - .map(|sort_exprs| { - LexRequirement::new(sort_exprs.into_iter().map(Into::into)) - }) - }) - .transpose()? - .flatten(); - Ok(Arc::new(DataSinkExec::new( - input, - Arc::new(data_sink), - sort_order, - ))) - } - #[cfg(not(feature = "parquet"))] - panic!("Trying to use ParquetSink without `parquet` feature enabled"); - } - - fn try_into_unnest_physical_plan( - &self, - unnest: &protobuf::UnnestExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&unnest.input, ctx, proto_converter)?; - - Ok(Arc::new(UnnestExec::new( - input, - unnest - .list_type_columns - .iter() - .map(|c| ListUnnest { - index_in_input_schema: c.index_in_input_schema as _, - depth: c.depth as _, - }) - .collect(), - unnest.struct_type_columns.iter().map(|c| *c as _).collect(), - Arc::new(convert_required!(unnest.schema)?), - into_required!(unnest.options)?, - )?)) - } - - fn generate_series_name_to_str(name: protobuf::GenerateSeriesName) -> &'static str { - match name { - protobuf::GenerateSeriesName::GsGenerateSeries => "generate_series", - protobuf::GenerateSeriesName::GsRange => "range", - } - } - fn try_into_sort_join( - &self, - sort_join: &SortMergeJoinExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let left = into_physical_plan(&sort_join.left, ctx, proto_converter)?; - let left_schema = left.schema(); - let right = into_physical_plan(&sort_join.right, ctx, proto_converter)?; - let right_schema = right.schema(); - - let filter = sort_join - .filter - .as_ref() - .map(|f| { - let schema = f - .schema - .as_ref() - .ok_or_else(|| proto_error("Missing JoinFilter schema"))? - .try_into()?; - - let expression = proto_converter.proto_to_physical_expr( - f.expression.as_ref().ok_or_else(|| { - proto_error("Unexpected empty filter expression") - })?, - &schema, - ctx, - )?; - let column_indices = f - .column_indices - .iter() - .map(|i| { - let side = - protobuf::JoinSide::try_from(i.side).map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with JoinSide in Filter {}", - i.side - )) - })?; - - Ok(ColumnIndex { - index: i.index as usize, - side: side.into(), - }) - }) - .collect::>>()?; - - Ok(JoinFilter::new( - expression, - column_indices, - Arc::new(schema), - )) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let join_type = - protobuf::JoinType::try_from(sort_join.join_type).map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with unknown JoinType {}", - sort_join.join_type - )) - })?; - - let null_equality = protobuf::NullEquality::try_from(sort_join.null_equality) - .map_err(|_| { - proto_error(format!( - "Received a SortMergeJoinExecNode message with unknown NullEquality {}", - sort_join.null_equality - )) - })?; - - let sort_options = sort_join - .sort_options - .iter() - .map(|e| SortOptions { - descending: !e.asc, - nulls_first: e.nulls_first, - }) - .collect(); - let on = sort_join - .on - .iter() - .map(|col| { - let left = proto_converter.proto_to_physical_expr( - &col.left.clone().unwrap(), - left_schema.as_ref(), - ctx, - )?; - let right = proto_converter.proto_to_physical_expr( - &col.right.clone().unwrap(), - right_schema.as_ref(), - ctx, - )?; - Ok((left, right)) - }) - .collect::>()?; - - Ok(Arc::new(SortMergeJoinExec::try_new( - left, - right, - on, - filter, - join_type.into(), - sort_options, - null_equality.into(), - )?)) - } - - fn try_into_generate_series_physical_plan( - &self, - generate_series: &protobuf::GenerateSeriesNode, - ) -> Result> { - let schema: SchemaRef = Arc::new(convert_required!(generate_series.schema)?); - - let args = match &generate_series.args { - Some(protobuf::generate_series_node::Args::ContainsNull(args)) => { - GenSeriesArgs::ContainsNull { - name: Self::generate_series_name_to_str(args.name()), - } - } - Some(protobuf::generate_series_node::Args::Int64Args(args)) => { - GenSeriesArgs::Int64Args { - start: args.start, - end: args.end, - step: args.step, - include_end: args.include_end, - name: Self::generate_series_name_to_str(args.name()), - } - } - Some(protobuf::generate_series_node::Args::TimestampArgs(args)) => { - let step_proto = args.step.as_ref().ok_or_else(|| { - internal_datafusion_err!("Missing step in TimestampArgs") - })?; - let step = IntervalMonthDayNanoType::make_value( - step_proto.months, - step_proto.days, - step_proto.nanos, - ); - GenSeriesArgs::TimestampArgs { - start: args.start, - end: args.end, - step, - tz: args.tz.as_ref().map(|s| Arc::from(s.as_str())), - include_end: args.include_end, - name: Self::generate_series_name_to_str(args.name()), - } - } - Some(protobuf::generate_series_node::Args::DateArgs(args)) => { - let step_proto = args.step.as_ref().ok_or_else(|| { - internal_datafusion_err!("Missing step in DateArgs") - })?; - let step = IntervalMonthDayNanoType::make_value( - step_proto.months, - step_proto.days, - step_proto.nanos, - ); - GenSeriesArgs::DateArgs { - start: args.start, - end: args.end, - step, - include_end: args.include_end, - name: Self::generate_series_name_to_str(args.name()), - } - } - None => return internal_err!("Missing args in GenerateSeriesNode"), - }; - - let table = GenerateSeriesTable::new(Arc::clone(&schema), args); - let generator = table.as_generator(generate_series.target_batch_size as usize)?; - - Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?)) - } - - fn try_into_cooperative_physical_plan( - &self, - field_stream: &protobuf::CooperativeExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input = into_physical_plan(&field_stream.input, ctx, proto_converter)?; - Ok(Arc::new(CooperativeExec::new(input))) - } - - fn try_into_async_func_physical_plan( - &self, - async_func: &protobuf::AsyncFuncExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&async_func.input, ctx, proto_converter)?; - - if async_func.async_exprs.len() != async_func.async_expr_names.len() { - return internal_err!( - "AsyncFuncExecNode async_exprs length does not match async_expr_names" - ); - } - - let async_exprs = async_func - .async_exprs - .iter() - .zip(async_func.async_expr_names.iter()) - .map(|(expr, name)| { - let physical_expr = proto_converter.proto_to_physical_expr( - expr, - input.schema().as_ref(), - ctx, - )?; - - Ok(Arc::new(AsyncFuncExpr::try_new( - name.clone(), - physical_expr, - input.schema().as_ref(), - )?)) - }) - .collect::>>()?; - - Ok(Arc::new(AsyncFuncExec::try_new(async_exprs, input)?)) - } - - fn try_into_buffer_physical_plan( - &self, - buffer: &protobuf::BufferExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: Arc = - into_physical_plan(&buffer.input, ctx, proto_converter)?; - - Ok(Arc::new(BufferExec::new(input, buffer.capacity as usize))) - } - - fn try_into_scalar_subquery_physical_plan( - &self, - sq: &protobuf::ScalarSubqueryExecNode, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - // First, deserialize the main input plan. We set up the subquery results - // container first, so that ScalarSubqueryExpr nodes can reference it. - let subquery_results = ScalarSubqueryResults::new(sq.subqueries.len()); - let input_ctx = ctx.with_scalar_subquery_results(subquery_results.clone()); - let input = into_physical_plan(&sq.input, &input_ctx, proto_converter)?; - - // Now deserialize the subquery children. - let subqueries: Vec = sq - .subqueries - .iter() - .enumerate() - .map(|(index, sq_plan)| { - let plan = - sq_plan.try_into_physical_plan_with_context(ctx, proto_converter)?; - Ok(ScalarSubqueryLink { - plan, - index: SubqueryIndex::new(index), - }) - }) - .collect::>>()?; - - Ok(Arc::new(ScalarSubqueryExec::new( - input, - subqueries, - subquery_results, - ))) - } - - fn try_from_explain_exec( - exec: &ExplainExec, - _codec: &dyn PhysicalExtensionCodec, - ) -> Result { - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Explain( - protobuf::ExplainExecNode { - schema: Some(exec.schema().as_ref().try_into()?), - stringified_plans: exec - .stringified_plans() - .iter() - .map(|plan| plan.into()) - .collect(), - verbose: exec.verbose(), - }, - )), - }) - } - - fn try_from_projection_exec( - exec: &ProjectionExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let expr = exec - .expr() - .iter() - .map(|proj_expr| { - proto_converter.physical_expr_to_proto(&proj_expr.expr, codec) - }) - .collect::>>()?; - let expr_name = exec - .expr() - .iter() - .map(|proj_expr| proj_expr.alias.clone()) - .collect(); - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Projection(Box::new( - protobuf::ProjectionExecNode { - input: Some(Box::new(input)), - expr, - expr_name, - }, - ))), - }) - } - - fn try_from_analyze_exec( - exec: &AnalyzeExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let (has_metric_categories, metric_categories) = match exec.metric_categories() { - Some(cats) => (true, cats.iter().map(|c| c.to_string()).collect()), - None => (false, vec![]), - }; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Analyze(Box::new( - protobuf::AnalyzeExecNode { - verbose: exec.verbose(), - show_statistics: exec.show_statistics(), - input: Some(Box::new(input)), - schema: Some(exec.schema().as_ref().try_into()?), - has_metric_categories, - metric_categories, - }, - ))), - }) - } - - fn try_from_filter_exec( - exec: &FilterExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Filter(Box::new( - protobuf::FilterExecNode { - input: Some(Box::new(input)), - expr: Some( - proto_converter - .physical_expr_to_proto(exec.predicate(), codec)?, - ), - default_filter_selectivity: exec.default_selectivity() as u32, - projection: match exec.projection() { - None => (0..exec.input().schema().fields().len()) - .map(|i| i as u32) - .collect(), - Some(v) => v.iter().map(|x| *x as u32).collect(), - }, - batch_size: exec.batch_size() as u32, - fetch: exec.fetch().map(|f| f as u32), - }, - ))), - }) - } - - fn try_from_global_limit_exec( - limit: &GlobalLimitExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - limit.input().to_owned(), - codec, - proto_converter, - )?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::GlobalLimit(Box::new( - protobuf::GlobalLimitExecNode { - input: Some(Box::new(input)), - skip: limit.skip() as u32, - fetch: match limit.fetch() { - Some(n) => n as i64, - _ => -1, // no limit - }, - }, - ))), - }) - } - - fn try_from_local_limit_exec( - limit: &LocalLimitExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - limit.input().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::LocalLimit(Box::new( - protobuf::LocalLimitExecNode { - input: Some(Box::new(input)), - fetch: limit.fetch() as u32, - }, - ))), - }) - } - - fn try_from_hash_join_exec( - exec: &HashJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - let on: Vec = exec - .on() - .iter() - .map(|tuple| { - let l = proto_converter.physical_expr_to_proto(&tuple.0, codec)?; - let r = proto_converter.physical_expr_to_proto(&tuple.1, codec)?; - Ok::<_, DataFusionError>(protobuf::JoinOn { - left: Some(l), - right: Some(r), - }) - }) - .collect::>()?; - let join_type: protobuf::JoinType = exec.join_type().to_owned().into(); - let null_equality: protobuf::NullEquality = exec.null_equality().into(); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let partition_mode = match exec.partition_mode() { - PartitionMode::CollectLeft => protobuf::PartitionMode::CollectLeft, - PartitionMode::Partitioned => protobuf::PartitionMode::Partitioned, - PartitionMode::Auto => protobuf::PartitionMode::Auto, - }; - - let dynamic_filter = exec - .dynamic_filter_expr() - .map(|df| { - let df_expr: Arc = - Arc::clone(df) as Arc; - proto_converter.physical_expr_to_proto(&df_expr, codec) - }) - .transpose()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::HashJoin(Box::new( - protobuf::HashJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - partition_mode: partition_mode.into(), - null_equality: null_equality.into(), - filter, - projection: exec.projection.as_ref().map_or_else(Vec::new, |v| { - v.iter().map(|x| *x as u32).collect::>() - }), - null_aware: exec.null_aware, - dynamic_filter, - }, - ))), - }) - } - - fn try_from_symmetric_hash_join_exec( - exec: &SymmetricHashJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - let on = exec - .on() - .iter() - .map(|tuple| { - let l = proto_converter.physical_expr_to_proto(&tuple.0, codec)?; - let r = proto_converter.physical_expr_to_proto(&tuple.1, codec)?; - Ok::<_, DataFusionError>(protobuf::JoinOn { - left: Some(l), - right: Some(r), - }) - }) - .collect::>()?; - let join_type: protobuf::JoinType = exec.join_type().to_owned().into(); - let null_equality: protobuf::NullEquality = exec.null_equality().into(); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let partition_mode = match exec.partition_mode() { - StreamJoinPartitionMode::SinglePartition => { - protobuf::StreamPartitionMode::SinglePartition - } - StreamJoinPartitionMode::Partitioned => { - protobuf::StreamPartitionMode::PartitionedExec - } - }; - - let left_sort_exprs = exec - .left_sort_exprs() - .map(|exprs| { - exprs - .iter() - .map(|expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }) - }) - .collect::>>() - }) - .transpose()? - .unwrap_or(vec![]); - - let right_sort_exprs = exec - .right_sort_exprs() - .map(|exprs| { - exprs - .iter() - .map(|expr| { - Ok(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }) - }) - .collect::>>() - }) - .transpose()? - .unwrap_or(vec![]); - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SymmetricHashJoin(Box::new( - protobuf::SymmetricHashJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - partition_mode: partition_mode.into(), - null_equality: null_equality.into(), - left_sort_exprs, - right_sort_exprs, - filter, - }, - ))), - }) - } - - fn try_from_sort_merge_join_exec( - exec: &SortMergeJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - let on = exec - .on() - .iter() - .map(|tuple| { - let l = proto_converter.physical_expr_to_proto(&tuple.0, codec)?; - let r = proto_converter.physical_expr_to_proto(&tuple.1, codec)?; - Ok::<_, DataFusionError>(protobuf::JoinOn { - left: Some(l), - right: Some(r), - }) - }) - .collect::>()?; - let join_type: protobuf::JoinType = exec.join_type().to_owned().into(); - let null_equality: protobuf::NullEquality = exec.null_equality().into(); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - let sort_options = exec - .sort_options() - .iter() - .map( - |SortOptions { - descending, - nulls_first, - }| { - SortExprNode { - expr: None, - asc: !*descending, - nulls_first: *nulls_first, - } - }, - ) - .collect(); - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SortMergeJoin(Box::new( - SortMergeJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - on, - join_type: join_type.into(), - null_equality: null_equality.into(), - filter, - sort_options, - }, - ))), - }) - } - - fn try_from_cross_join_exec( - exec: &CrossJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CrossJoin(Box::new( - protobuf::CrossJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - }, - ))), - }) - } - - fn try_from_aggregate_exec( - exec: &AggregateExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let groups: Vec = exec - .group_expr() - .groups() - .iter() - .flatten() - .copied() - .collect(); - - let group_names = exec - .group_expr() - .expr() - .iter() - .map(|expr| expr.1.to_owned()) - .collect(); - - let filter = exec - .filter_expr() - .iter() - .map(|expr| serialize_maybe_filter(expr.to_owned(), codec, proto_converter)) - .collect::>>()?; - - let agg = exec - .aggr_expr() - .iter() - .map(|expr| { - serialize_physical_aggr_expr(expr.to_owned(), codec, proto_converter) - }) - .collect::>>()?; - - let agg_names = exec - .aggr_expr() - .iter() - .map(|expr| expr.name().to_string()) - .collect::>(); - - let agg_mode = match exec.mode() { - AggregateMode::Partial => protobuf::AggregateMode::Partial, - AggregateMode::Final => protobuf::AggregateMode::Final, - AggregateMode::FinalPartitioned => protobuf::AggregateMode::FinalPartitioned, - AggregateMode::Single => protobuf::AggregateMode::Single, - AggregateMode::SinglePartitioned => { - protobuf::AggregateMode::SinglePartitioned - } - AggregateMode::PartialReduce => protobuf::AggregateMode::PartialReduce, - }; - let input_schema = exec.input_schema(); - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - - let null_expr = exec - .group_expr() - .null_expr() - .iter() - .map(|expr| proto_converter.physical_expr_to_proto(&expr.0, codec)) - .collect::>>()?; - - let group_expr = exec - .group_expr() - .expr() - .iter() - .map(|expr| proto_converter.physical_expr_to_proto(&expr.0, codec)) - .collect::>>()?; - - let limit = exec.limit_options().map(|config| protobuf::AggLimit { - limit: config.limit() as u64, - descending: config.descending(), - }); - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Aggregate(Box::new( - protobuf::AggregateExecNode { - group_expr, - group_expr_name: group_names, - aggr_expr: agg, - filter_expr: filter, - aggr_expr_name: agg_names, - mode: agg_mode as i32, - input: Some(Box::new(input)), - input_schema: Some(input_schema.as_ref().try_into()?), - null_expr, - groups, - limit, - has_grouping_set: exec.group_expr().has_grouping_set(), - dynamic_filter: exec - .dynamic_filter_expr() - .map(|df| { - let df_expr: Arc = - Arc::clone(df) as Arc; - proto_converter.physical_expr_to_proto(&df_expr, codec) - }) - .transpose()?, - }, - ))), - }) - } - - fn try_from_empty_exec( - empty: &EmptyExec, - _codec: &dyn PhysicalExtensionCodec, - ) -> Result { - let schema = empty.schema().as_ref().try_into()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Empty(protobuf::EmptyExecNode { - schema: Some(schema), - })), - }) - } - - fn try_from_placeholder_row_exec( - empty: &PlaceholderRowExec, - _codec: &dyn PhysicalExtensionCodec, - ) -> Result { - let schema = empty.schema().as_ref().try_into()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::PlaceholderRow( - protobuf::PlaceholderRowExecNode { - schema: Some(schema), - }, - )), - }) - } - - #[expect(deprecated)] - fn try_from_coalesce_batches_exec( - coalesce_batches: &CoalesceBatchesExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - coalesce_batches.input().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CoalesceBatches(Box::new( - protobuf::CoalesceBatchesExecNode { - input: Some(Box::new(input)), - target_batch_size: coalesce_batches.target_batch_size() as u32, - fetch: coalesce_batches.fetch().map(|n| n as u32), - }, - ))), - }) - } - - fn try_from_data_source_exec( - data_source_exec: &DataSourceExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let data_source = data_source_exec.data_source(); - if let Some(maybe_csv) = data_source.downcast_ref::() { - let source = maybe_csv.file_source(); - if let Some(csv_config) = source.downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CsvScan( - protobuf::CsvScanExecNode { - base_conf: Some(serialize_file_scan_config( - maybe_csv, - codec, - proto_converter, - )?), - has_header: csv_config.has_header(), - delimiter: byte_to_string( - csv_config.delimiter(), - "delimiter", - )?, - quote: byte_to_string(csv_config.quote(), "quote")?, - optional_escape: if let Some(escape) = csv_config.escape() { - Some( - protobuf::csv_scan_exec_node::OptionalEscape::Escape( - byte_to_string(escape, "escape")?, - ), - ) - } else { - None - }, - optional_comment: if let Some(comment) = csv_config.comment() - { - Some(protobuf::csv_scan_exec_node::OptionalComment::Comment( - byte_to_string(comment, "comment")?, - )) - } else { - None - }, - newlines_in_values: csv_config.newlines_in_values(), - truncate_rows: csv_config.truncate_rows(), - }, - )), - })); - } - } - - if let Some(scan_conf) = data_source.downcast_ref::() { - let source = scan_conf.file_source(); - if let Some(_json_source) = source.downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::JsonScan( - protobuf::JsonScanExecNode { - base_conf: Some(serialize_file_scan_config( - scan_conf, - codec, - proto_converter, - )?), - }, - )), - })); - } - } - - if let Some(scan_conf) = data_source.downcast_ref::() { - let source = scan_conf.file_source(); - if let Some(_arrow_source) = source.downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ArrowScan( - protobuf::ArrowScanExecNode { - base_conf: Some(serialize_file_scan_config( - scan_conf, - codec, - proto_converter, - )?), - }, - )), - })); - } - } - - #[cfg(feature = "parquet")] - if let Some((maybe_parquet, conf)) = - data_source_exec.downcast_to_file_source::() - { - let predicate = conf - .filter() - .map(|pred| proto_converter.physical_expr_to_proto(&pred, codec)) - .transpose()?; - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ParquetScan( - protobuf::ParquetScanExecNode { - base_conf: Some(serialize_file_scan_config( - maybe_parquet, - codec, - proto_converter, - )?), - predicate, - parquet_options: Some(conf.table_parquet_options().try_into()?), - }, - )), - })); - } - - #[cfg(feature = "avro")] - if let Some(maybe_avro) = data_source.downcast_ref::() { - let source = maybe_avro.file_source(); - if source.downcast_ref::().is_some() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::AvroScan( - protobuf::AvroScanExecNode { - base_conf: Some(serialize_file_scan_config( - maybe_avro, - codec, - proto_converter, - )?), - }, - )), - })); - } - } - - if let Some(source_conf) = data_source.downcast_ref::() { - let proto_partitions = source_conf - .partitions() - .iter() - .map(|p| serialize_record_batches(p)) - .collect::>>()?; - - let proto_schema: protobuf::Schema = - source_conf.original_schema().as_ref().try_into()?; - - let proto_projection = source_conf - .projection() - .as_ref() - .map_or_else(Vec::new, |v| { - v.iter().map(|x| *x as u32).collect::>() - }); - - let proto_sort_information = source_conf - .sort_information() - .iter() - .map(|ordering| { - let sort_exprs = serialize_physical_sort_exprs( - ordering.to_owned(), - codec, - proto_converter, - )?; - Ok::<_, DataFusionError>(protobuf::PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: sort_exprs, - }) - }) - .collect::, _>>()?; - - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::MemoryScan( - protobuf::MemoryScanExecNode { - partitions: proto_partitions, - schema: Some(proto_schema), - projection: proto_projection, - sort_information: proto_sort_information, - show_sizes: source_conf.show_sizes(), - fetch: source_conf.fetch().map(|f| f as u32), - }, - )), - })); - } - - Ok(None) - } - - fn try_from_coalesce_partitions_exec( - exec: &CoalescePartitionsExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Merge(Box::new( - protobuf::CoalescePartitionsExecNode { - input: Some(Box::new(input)), - fetch: exec.fetch().map(|f| f as u32), - }, - ))), - }) - } - - fn try_from_repartition_exec( - exec: &RepartitionExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - - let pb_partitioning = - serialize_partitioning(exec.partitioning(), codec, proto_converter)?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Repartition(Box::new( - protobuf::RepartitionExecNode { - input: Some(Box::new(input)), - partitioning: Some(pb_partitioning), - preserve_order: exec.preserve_order(), - }, - ))), - }) - } - - fn try_from_sort_exec( - exec: &SortExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = proto_converter.execution_plan_to_proto(exec.input(), codec)?; - let expr = exec - .expr() - .iter() - .map(|expr| { - let sort_expr = Box::new(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }); - Ok(protobuf::PhysicalExprNode { - expr_id: None, - expr_type: Some(ExprType::Sort(sort_expr)), - }) - }) - .collect::>>()?; - let dynamic_filter = exec - .dynamic_filter_expr() - .map(|df| { - let df_expr: Arc = df as Arc; - proto_converter.physical_expr_to_proto(&df_expr, codec) - }) - .transpose()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Sort(Box::new( - protobuf::SortExecNode { - input: Some(Box::new(input)), - expr, - fetch: match exec.fetch() { - Some(n) => n as i64, - _ => -1, - }, - preserve_partitioning: exec.preserve_partitioning(), - dynamic_filter, - }, - ))), - }) - } - - fn try_from_union_exec( - union: &UnionExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let mut inputs: Vec = vec![]; - for input in union.inputs() { - inputs.push( - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - input.to_owned(), - codec, - proto_converter, - )?, - ); - } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Union(protobuf::UnionExecNode { - inputs, - })), - }) - } - - fn try_from_interleave_exec( - interleave: &InterleaveExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let mut inputs: Vec = vec![]; - for input in interleave.inputs() { - inputs.push( - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - input.to_owned(), - codec, - proto_converter, - )?, - ); - } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Interleave( - protobuf::InterleaveExecNode { inputs }, - )), - }) - } - - fn try_from_sort_preserving_merge_exec( - exec: &SortPreservingMergeExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let expr = exec - .expr() - .iter() - .map(|expr| { - let sort_expr = Box::new(protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }); - Ok(protobuf::PhysicalExprNode { - expr_id: None, - expr_type: Some(ExprType::Sort(sort_expr)), - }) - }) - .collect::>>()?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::SortPreservingMerge(Box::new( - protobuf::SortPreservingMergeExecNode { - input: Some(Box::new(input)), - expr, - fetch: exec.fetch().map(|f| f as i64).unwrap_or(-1), - }, - ))), - }) - } - - fn try_from_nested_loop_join_exec( - exec: &NestedLoopJoinExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let left = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.left().to_owned(), - codec, - proto_converter, - )?; - let right = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.right().to_owned(), - codec, - proto_converter, - )?; - - let join_type: protobuf::JoinType = exec.join_type().to_owned().into(); - let filter = exec - .filter() - .as_ref() - .map(|f| { - let expression = - proto_converter.physical_expr_to_proto(f.expression(), codec)?; - let column_indices = f - .column_indices() - .iter() - .map(|i| { - let side: protobuf::JoinSide = i.side.to_owned().into(); - protobuf::ColumnIndex { - index: i.index as u32, - side: side.into(), - } - }) - .collect(); - let schema = f.schema().as_ref().try_into()?; - Ok(protobuf::JoinFilter { - expression: Some(expression), - column_indices, - schema: Some(schema), - }) - }) - .map_or(Ok(None), |v: Result| v.map(Some))?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::NestedLoopJoin(Box::new( - protobuf::NestedLoopJoinExecNode { - left: Some(Box::new(left)), - right: Some(Box::new(right)), - join_type: join_type.into(), - filter, - projection: exec.projection().as_ref().map_or_else(Vec::new, |v| { - v.iter().map(|x| *x as u32).collect::>() - }), - }, - ))), - }) - } - - fn try_from_window_agg_exec( - exec: &WindowAggExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - - let window_expr = exec - .window_expr() - .iter() - .map(|e| serialize_physical_window_expr(e, codec, proto_converter)) - .collect::>>()?; - - let partition_keys = exec - .partition_keys() - .iter() - .map(|e| proto_converter.physical_expr_to_proto(e, codec)) - .collect::>>()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Window(Box::new( - protobuf::WindowAggExecNode { - input: Some(Box::new(input)), - window_expr, - partition_keys, - input_order_mode: None, - }, - ))), - }) - } - - fn try_from_bounded_window_agg_exec( - exec: &BoundedWindowAggExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - - let window_expr = exec - .window_expr() - .iter() - .map(|e| serialize_physical_window_expr(e, codec, proto_converter)) - .collect::>>()?; - - let partition_keys = exec - .partition_keys() - .iter() - .map(|e| proto_converter.physical_expr_to_proto(e, codec)) - .collect::>>()?; - - let input_order_mode = match &exec.input_order_mode { - InputOrderMode::Linear => { - window_agg_exec_node::InputOrderMode::Linear(protobuf::EmptyMessage {}) - } - InputOrderMode::PartiallySorted(columns) => { - window_agg_exec_node::InputOrderMode::PartiallySorted( - protobuf::PartiallySortedInputOrderMode { - columns: columns.iter().map(|c| *c as u64).collect(), - }, - ) + PhysicalPlanType::AsyncFunc(_) => { + AsyncFuncExec::try_from_proto(self.node(), &decode_ctx) } - InputOrderMode::Sorted => { - window_agg_exec_node::InputOrderMode::Sorted(protobuf::EmptyMessage {}) + PhysicalPlanType::Buffer(_) => { + BufferExec::try_from_proto(self.node(), &decode_ctx) } - }; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Window(Box::new( - protobuf::WindowAggExecNode { - input: Some(Box::new(input)), - window_expr, - partition_keys, - input_order_mode: Some(input_order_mode), - }, - ))), - }) + PhysicalPlanType::ScalarSubquery(_) => { + ScalarSubqueryExec::try_from_proto(self.node(), &decode_ctx) + } + } } - fn try_from_data_sink_exec( - exec: &DataSinkExec, + fn try_from_physical_plan_with_converter( + plan: Arc, codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result> { - let input: protobuf::PhysicalPlanNode = - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; - let sort_order = match exec.sort_order() { - Some(requirements) => { - let expr = requirements - .iter() - .map(|requirement| { - let expr: PhysicalSortExpr = requirement.to_owned().into(); - let sort_expr = protobuf::PhysicalSortExprNode { - expr: Some(Box::new( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - )), - asc: !expr.options.descending, - nulls_first: expr.options.nulls_first, - }; - Ok(sort_expr) - }) - .collect::>>()?; - Some(protobuf::PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: expr, - }) - } - None => None, + ) -> Result { + let plan_clone = Arc::clone(&plan); + let mut plan = plan.as_ref(); + // Resolve the downcast identity first so wrapper plans serialize as + // their delegate, matching how the `downcast_ref` chain below sees + // them. Without this a wrapper around a migrated plan would hit the + // wrapper's default `try_to_proto` (`Ok(None)`) and find no fallback + // arm for the delegate. + while let Some(delegate) = plan.downcast_delegate() { + plan = delegate; + } + + // Self-serializing plans handle themselves via the `try_to_proto` hook + // (#22419). `Ok(None)` means "not migrated" and falls through to the + // central downcast chain below. + let encoder = ConverterPlanEncoder { + codec, + proto_converter, }; - - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::JsonSink(Box::new( - protobuf::JsonSinkExecNode { - input: Some(Box::new(input)), - sink: Some(sink.try_into()?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); + let encode_ctx = ExecutionPlanEncodeCtx::new(&encoder); + if let Some(node) = plan.try_to_proto(&encode_ctx)? { + return Ok(node); } - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::CsvSink(Box::new( - protobuf::CsvSinkExecNode { - input: Some(Box::new(input)), - sink: Some(sink.try_into()?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); + if let Some(exec) = plan.downcast_ref::() + && let Some(node) = + protobuf::PhysicalPlanNode::try_from_lazy_memory_exec(exec)? + { + return Ok(node); } - #[cfg(feature = "parquet")] - if let Some(sink) = exec.sink().downcast_ref::() { - return Ok(Some(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ParquetSink(Box::new( - protobuf::ParquetSinkExecNode { - input: Some(Box::new(input)), - sink: Some(sink.try_into()?), - sink_schema: Some(exec.schema().as_ref().try_into()?), - sort_order, - }, - ))), - })); - } + let mut buf: Vec = vec![]; + match codec.try_encode(Arc::clone(&plan_clone), &mut buf, proto_converter) { + Ok(_) => { + let inputs: Vec = plan_clone + .children() + .into_iter() + .cloned() + .map(|i| { + protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( + i, + codec, + proto_converter, + ) + }) + .collect::>()?; - // If unknown DataSink then let extension handle it - Ok(None) + Ok(protobuf::PhysicalPlanNode { + physical_plan_type: Some(PhysicalPlanType::Extension( + protobuf::PhysicalExtensionNode { node: buf, inputs }, + )), + }) + } + Err(e) => internal_err!( + "Unsupported plan and extension codec failed with [{e}]. Plan: {plan_clone:?}" + ), + } } - fn try_from_unnest_exec( - exec: &UnnestExec, - codec: &dyn PhysicalExtensionCodec, + fn try_into_extension_physical_plan( + &self, + extension: &protobuf::PhysicalExtensionNode, + ctx: &PhysicalPlanDecodeContext<'_>, proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, + ) -> Result> { + let inputs: Vec> = extension + .inputs + .iter() + .map(|i| proto_converter.proto_to_execution_plan(i, ctx)) + .collect::>()?; + + let extension_node = ctx.codec().try_decode( + extension.node.as_slice(), + &inputs, + ctx.task_ctx(), proto_converter, )?; - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Unnest(Box::new( - protobuf::UnnestExecNode { - input: Some(Box::new(input)), - schema: Some(exec.schema().try_into()?), - list_type_columns: exec - .list_column_indices() - .iter() - .map(|c| ProtoListUnnest { - index_in_input_schema: c.index_in_input_schema as _, - depth: c.depth as _, - }) - .collect(), - struct_type_columns: exec - .struct_column_indices() - .iter() - .map(|c| *c as _) - .collect(), - options: Some(exec.options().into()), - }, - ))), - }) + Ok(extension_node) } - fn try_from_cooperative_exec( - exec: &CooperativeExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - exec.input().to_owned(), - codec, - proto_converter, - )?; + fn generate_series_name_to_str(name: protobuf::GenerateSeriesName) -> &'static str { + match name { + protobuf::GenerateSeriesName::GsGenerateSeries => "generate_series", + protobuf::GenerateSeriesName::GsRange => "range", + } + } - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Cooperative(Box::new( - protobuf::CooperativeExecNode { - input: Some(Box::new(input)), - }, - ))), - }) + fn try_into_generate_series_physical_plan( + &self, + generate_series: &protobuf::GenerateSeriesNode, + ) -> Result> { + let schema: SchemaRef = Arc::new(convert_required!(generate_series.schema)?); + + let args = match &generate_series.args { + Some(protobuf::generate_series_node::Args::ContainsNull(args)) => { + GenSeriesArgs::ContainsNull { + name: protobuf::PhysicalPlanNode::generate_series_name_to_str( + args.name(), + ), + } + } + Some(protobuf::generate_series_node::Args::Int64Args(args)) => { + GenSeriesArgs::Int64Args { + start: args.start, + end: args.end, + step: args.step, + include_end: args.include_end, + name: protobuf::PhysicalPlanNode::generate_series_name_to_str( + args.name(), + ), + } + } + Some(protobuf::generate_series_node::Args::TimestampArgs(args)) => { + let step_proto = args.step.as_ref().ok_or_else(|| { + internal_datafusion_err!("Missing step in TimestampArgs") + })?; + let step = IntervalMonthDayNanoType::make_value( + step_proto.months, + step_proto.days, + step_proto.nanos, + ); + GenSeriesArgs::TimestampArgs { + start: args.start, + end: args.end, + step, + tz: args.tz.as_ref().map(|s| Arc::from(s.as_str())), + include_end: args.include_end, + name: protobuf::PhysicalPlanNode::generate_series_name_to_str( + args.name(), + ), + } + } + Some(protobuf::generate_series_node::Args::DateArgs(args)) => { + let step_proto = args.step.as_ref().ok_or_else(|| { + internal_datafusion_err!("Missing step in DateArgs") + })?; + let step = IntervalMonthDayNanoType::make_value( + step_proto.months, + step_proto.days, + step_proto.nanos, + ); + GenSeriesArgs::DateArgs { + start: args.start, + end: args.end, + step, + include_end: args.include_end, + name: protobuf::PhysicalPlanNode::generate_series_name_to_str( + args.name(), + ), + } + } + None => return internal_err!("Missing args in GenerateSeriesNode"), + }; + + let table = GenerateSeriesTable::new(Arc::clone(&schema), args); + let generator = table.as_generator(generate_series.target_batch_size as usize)?; + + Ok(Arc::new(LazyMemoryExec::try_new(schema, vec![generator])?)) } fn str_to_generate_series_name(name: &str) -> Result { @@ -3629,7 +1385,9 @@ impl protobuf::PhysicalPlanNode { } } - fn try_from_lazy_memory_exec(exec: &LazyMemoryExec) -> Result> { + fn try_from_lazy_memory_exec( + exec: &LazyMemoryExec, + ) -> Result> { let generators = exec.generators(); // ensure we only have one generator @@ -3647,7 +1405,9 @@ impl protobuf::PhysicalPlanNode { target_batch_size: 8192, // Default batch size args: Some(protobuf::generate_series_node::Args::ContainsNull( protobuf::GenerateSeriesArgsContainsNull { - name: Self::str_to_generate_series_name(empty_gen.name())? as i32, + name: protobuf::PhysicalPlanNode::str_to_generate_series_name( + empty_gen.name(), + )? as i32, }, )), }; @@ -3671,7 +1431,9 @@ impl protobuf::PhysicalPlanNode { end: *int_64.end(), step: *int_64.step(), include_end: int_64.include_end(), - name: Self::str_to_generate_series_name(int_64.name())? as i32, + name: protobuf::PhysicalPlanNode::str_to_generate_series_name( + int_64.name(), + )? as i32, }, )), }; @@ -3698,7 +1460,9 @@ impl protobuf::PhysicalPlanNode { nanos: step_value.nanoseconds, }); let include_end = timestamp_args.include_end(); - let name = Self::str_to_generate_series_name(timestamp_args.name())? as i32; + let name = protobuf::PhysicalPlanNode::str_to_generate_series_name( + timestamp_args.name(), + )? as i32; let args = match timestamp_args.current().tz_str() { Some(tz) => protobuf::generate_series_node::Args::TimestampArgs( @@ -3735,89 +1499,11 @@ impl protobuf::PhysicalPlanNode { Ok(None) } +} - fn try_from_async_func_exec( - exec: &AsyncFuncExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - codec, - proto_converter, - )?; - - let mut async_exprs = vec![]; - let mut async_expr_names = vec![]; - - for async_expr in exec.async_exprs() { - async_exprs - .push(proto_converter.physical_expr_to_proto(&async_expr.func, codec)?); - async_expr_names.push(async_expr.name.clone()) - } - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::AsyncFunc(Box::new( - protobuf::AsyncFuncExecNode { - input: Some(Box::new(input)), - async_exprs, - async_expr_names, - }, - ))), - }) - } - - fn try_from_buffer_exec( - exec: &BufferExec, - extension_codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - extension_codec, - proto_converter, - )?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::Buffer(Box::new( - protobuf::BufferExecNode { - input: Some(Box::new(input)), - capacity: exec.capacity() as u64, - }, - ))), - }) - } - - fn try_from_scalar_subquery_exec( - exec: &ScalarSubqueryExec, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, - ) -> Result { - let input = protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(exec.input()), - codec, - proto_converter, - )?; - let subqueries = exec - .subqueries() - .iter() - .map(|sq| { - protobuf::PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(&sq.plan), - codec, - proto_converter, - ) - }) - .collect::>>()?; - - Ok(protobuf::PhysicalPlanNode { - physical_plan_type: Some(PhysicalPlanType::ScalarSubquery(Box::new( - protobuf::ScalarSubqueryExecNode { - input: Some(Box::new(input)), - subqueries, - }, - ))), - }) +impl PhysicalPlanNodeExt for protobuf::PhysicalPlanNode { + fn node(&self) -> &protobuf::PhysicalPlanNode { + self } } @@ -3852,9 +1538,15 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync + Any { buf: &[u8], inputs: &[Arc], ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result>; - fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()>; + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()>; fn try_decode_udf(&self, name: &str, _buf: &[u8]) -> Result> { not_impl_err!("PhysicalExtensionCodec is not provided for scalar function {name}") @@ -3864,18 +1556,62 @@ pub trait PhysicalExtensionCodec: Debug + Send + Sync + Any { Ok(()) } + fn try_decode_higher_order_function( + &self, + name: &str, + _buf: &[u8], + ) -> Result> { + not_impl_err!( + "PhysicalExtensionCodec is not provided for higher order function {name}" + ) + } + + fn try_encode_higher_order_function( + &self, + _node: &HigherOrderUDF, + _buf: &mut Vec, + ) -> Result<()> { + Ok(()) + } + + /// Decode a custom extension expression from `buf`. + /// + /// `inputs` holds the already-decoded children carried in the + /// `PhysicalExtensionExprNode.inputs` field. If the codec instead embeds + /// nested `PhysicalExprNode`s *inside* `buf`, decode them through + /// `ctx.decode(..)` (equivalently [`PhysicalExprDecodeCtx::decode`]) rather + /// than the free [`parse_physical_expr`] function: `ctx` carries the active + /// schema and task context (so UDF/column references resolve against the + /// real registry) and routes through any active `DeduplicatingDeserializer`, + /// so a shared inner expression (e.g. a `DynamicFilterPhysicalExpr` + /// referenced both from a `SortExec.filter` and from inside this blob) + /// cache-hits on its `expr_id` and re-shares one `Arc`. + /// + /// [`parse_physical_expr`]: crate::physical_plan::from_proto::parse_physical_expr fn try_decode_expr( &self, _buf: &[u8], _inputs: &[Arc], + _ctx: &PhysicalExprDecodeCtx<'_>, ) -> Result> { not_impl_err!("PhysicalExtensionCodec is not provided") } + /// Encode a custom extension expression into `buf`. + /// + /// If the codec embeds nested `PhysicalExprNode`s inside `buf`, encode them + /// through `ctx.encode_child(..)` (equivalently + /// [`PhysicalExprEncodeCtx::encode_child`]) rather than the free + /// [`serialize_physical_expr`] function, so an active + /// `DeduplicatingProtoConverter` stamps matching `expr_id`s for shared + /// inner expressions. See [`Self::try_decode_expr`]. + /// + /// [`serialize_physical_expr`]: crate::physical_plan::to_proto::serialize_physical_expr fn try_encode_expr( &self, _node: &Arc, _buf: &mut Vec, + _ctx: &PhysicalExprEncodeCtx<'_>, ) -> Result<()> { not_impl_err!("PhysicalExtensionCodec is not provided") } @@ -3908,6 +1644,7 @@ impl PhysicalExtensionCodec for DefaultPhysicalExtensionCodec { _buf: &[u8], _inputs: &[Arc], _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { not_impl_err!("PhysicalExtensionCodec is not provided") } @@ -3916,6 +1653,7 @@ impl PhysicalExtensionCodec for DefaultPhysicalExtensionCodec { &self, _node: Arc, _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result<()> { not_impl_err!("PhysicalExtensionCodec is not provided") } @@ -4239,12 +1977,22 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { buf: &[u8], inputs: &[Arc], ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result> { - self.decode_protobuf(buf, |codec, data| codec.try_decode(data, inputs, ctx)) + self.decode_protobuf(buf, |codec, data| { + codec.try_decode(data, inputs, ctx, proto_converter) + }) } - fn try_encode(&self, node: Arc, buf: &mut Vec) -> Result<()> { - self.encode_protobuf(buf, |codec, data| codec.try_encode(Arc::clone(&node), data)) + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + self.encode_protobuf(buf, |codec, data| { + codec.try_encode(Arc::clone(&node), data, proto_converter) + }) } fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { @@ -4264,14 +2012,129 @@ impl PhysicalExtensionCodec for ComposedPhysicalExtensionCodec { } } -fn into_physical_plan( - node: &Option>, - ctx: &PhysicalPlanDecodeContext<'_>, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result> { - if let Some(field) = node { - proto_converter.proto_to_execution_plan(field, ctx) - } else { - Err(proto_error("Missing required field in protobuf")) +/// Adapter backing [`ExecutionPlanEncodeCtx`] for plans migrated to the +/// `try_to_proto` hook (#22419). Routes child-plan and child-expr encoding back +/// through the central converter so nested plans honor their own hooks. +struct ConverterPlanEncoder<'a> { + codec: &'a dyn PhysicalExtensionCodec, + proto_converter: &'a dyn PhysicalProtoConverterExtension, +} + +impl ExecutionPlanEncode for ConverterPlanEncoder<'_> { + fn encode_plan( + &self, + plan: &Arc, + ) -> Result { + self.proto_converter + .execution_plan_to_proto(plan, self.codec) + } + + fn encode_expr( + &self, + expr: &Arc, + ) -> Result { + self.proto_converter + .physical_expr_to_proto(expr, self.codec) + } + + // Bytes-only function serde. `(!buf.is_empty()).then_some(buf)` preserves the + // existing `fun_definition` wire semantics (empty payload == encode-by-name). + fn encode_udf(&self, udf: &ScalarUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udf(udf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } + + fn encode_udaf(&self, udaf: &AggregateUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udaf(udaf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } + + fn encode_udwf(&self, udwf: &WindowUDF) -> Result>> { + let mut buf = vec![]; + self.codec.try_encode_udwf(udwf, &mut buf)?; + Ok((!buf.is_empty()).then_some(buf)) + } +} + +/// Adapter backing [`ExecutionPlanDecodeCtx`] for plans migrated to the +/// `try_from_proto` pattern (#22419). Routes child-plan and child-expr decoding +/// back through the central converter, and exposes the session task context +/// (never the extension codec). +struct ConverterPlanDecoder<'a, 'ctx> { + ctx: &'a PhysicalPlanDecodeContext<'ctx>, + proto_converter: &'a dyn PhysicalProtoConverterExtension, +} + +impl ExecutionPlanDecode for ConverterPlanDecoder<'_, '_> { + fn decode_plan( + &self, + node: &protobuf::PhysicalPlanNode, + ) -> Result> { + self.proto_converter.proto_to_execution_plan(node, self.ctx) + } + + fn decode_plan_with_scalar_subquery_results( + &self, + node: &protobuf::PhysicalPlanNode, + results: ScalarSubqueryResults, + ) -> Result> { + let scoped_ctx = self.ctx.with_scalar_subquery_results(results); + self.proto_converter + .proto_to_execution_plan(node, &scoped_ctx) + } + + fn decode_expr( + &self, + node: &protobuf::PhysicalExprNode, + input_schema: &Schema, + ) -> Result> { + self.proto_converter + .proto_to_physical_expr(node, input_schema, self.ctx) + } + + fn task_ctx(&self) -> &TaskContext { + self.ctx.task_ctx() + } + + // Lookup-order policy, owned here so no plan re-derives it: an explicit + // payload is decoded by the codec; otherwise resolve by name from the + // registry, falling back to the codec with an empty buffer. + fn decode_udf(&self, name: &str, payload: Option<&[u8]>) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udf(name, buf), + None => self + .ctx + .task_ctx() + .udf(name) + .or_else(|_| self.ctx.codec().try_decode_udf(name, &[])), + } + } + + fn decode_udaf( + &self, + name: &str, + payload: Option<&[u8]>, + ) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udaf(name, buf), + None => self + .ctx + .task_ctx() + .udaf(name) + .or_else(|_| self.ctx.codec().try_decode_udaf(name, &[])), + } + } + + fn decode_udwf(&self, name: &str, payload: Option<&[u8]>) -> Result> { + match payload { + Some(buf) => self.ctx.codec().try_decode_udwf(name, buf), + None => self + .ctx + .task_ctx() + .udwf(name) + .or_else(|_| self.ctx.codec().try_decode_udwf(name, &[])), + } } } diff --git a/datafusion/proto/src/physical_plan/to_proto.rs b/datafusion/proto/src/physical_plan/to_proto.rs index ec8e16817813b..5ae57752de676 100644 --- a/datafusion/proto/src/physical_plan/to_proto.rs +++ b/datafusion/proto/src/physical_plan/to_proto.rs @@ -20,38 +20,23 @@ use std::sync::Arc; use arrow::array::RecordBatch; use arrow::datatypes::Schema; use arrow::ipc::writer::StreamWriter; -use datafusion_common::{ - DataFusionError, Result, internal_datafusion_err, internal_err, not_impl_err, -}; +use datafusion_common::{Result, internal_datafusion_err, internal_err, not_impl_err}; use datafusion_datasource::file_scan_config::FileScanConfig; -use datafusion_datasource::file_sink_config::{FileSink, FileSinkConfig}; -use datafusion_datasource::{FileRange, PartitionedFile}; -use datafusion_datasource_csv::file_format::CsvSink; -use datafusion_datasource_json::file_format::JsonSink; -#[cfg(feature = "parquet")] -use datafusion_datasource_parquet::file_format::ParquetSink; use datafusion_expr::WindowFrame; -use datafusion_physical_expr::ScalarFunctionExpr; -use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; use datafusion_physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; +use datafusion_physical_expr::{HigherOrderFunctionExpr, ScalarFunctionExpr}; use datafusion_physical_expr_common::sort_expr::PhysicalSortExpr; -use datafusion_physical_plan::expressions::{ - BinaryExpr, CaseExpr, CastExpr, Column, DynamicFilterPhysicalExpr, InListExpr, - IsNotNullExpr, IsNullExpr, LikeExpr, Literal, NegativeExpr, NotExpr, TryCastExpr, - UnKnownColumn, -}; -use datafusion_physical_plan::joins::{HashExpr, HashTableLookupExpr}; +use datafusion_physical_plan::proto::ExecutionPlanEncodeCtx; use datafusion_physical_plan::udaf::AggregateFunctionExpr; use datafusion_physical_plan::windows::{PlainAggregateWindowExpr, WindowUDFExpr}; use datafusion_physical_plan::{Partitioning, PhysicalExpr, WindowExpr}; use super::{ - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + ConverterPlanEncoder, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalProtoConverterExtension, encode_human_display_alias, }; use crate::protobuf::{ - self, PhysicalSortExprNode, PhysicalSortExprNodeCollection, - physical_aggregate_expr_node, physical_window_expr_node, + self, PhysicalSortExprNode, physical_aggregate_expr_node, physical_window_expr_node, }; #[expect(clippy::needless_pass_by_value)] @@ -88,6 +73,7 @@ pub fn serialize_physical_aggr_expr( ignore_nulls: aggr_expr.ignore_nulls(), fun_definition: (!buf.is_empty()).then_some(buf), human_display, + is_reversed: aggr_expr.is_reversed(), }, )), }) @@ -178,9 +164,7 @@ pub fn serialize_physical_window_expr( codec, proto_converter, )?; - let window_frame: protobuf::WindowFrame = window_frame - .as_ref() - .try_into() + let window_frame = protobuf::WindowFrame::try_from(window_frame.as_ref()) .map_err(|e| internal_datafusion_err!("{e}"))?; Ok(protobuf::PhysicalWindowExprNode { @@ -253,6 +237,29 @@ pub fn serialize_physical_expr( ) } +/// Concrete [`PhysicalExprEncode`] driver used to back +/// [`PhysicalExprEncodeCtx`] when expressions invoke `PhysicalExpr::to_proto`. +/// +/// Wraps the existing extension codec + converter pair so individual +/// expressions can recurse into children without depending on +/// `datafusion-proto` directly. +/// +/// [`PhysicalExprEncode`]: datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncode +/// [`PhysicalExprEncodeCtx`]: datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx +struct ConverterEncoder<'a> { + codec: &'a dyn PhysicalExtensionCodec, + proto_converter: &'a dyn PhysicalProtoConverterExtension, +} + +impl datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncode + for ConverterEncoder<'_> +{ + fn encode(&self, expr: &Arc) -> Result { + self.proto_converter + .physical_expr_to_proto(expr, self.codec) + } +} + /// Serialize a `PhysicalExpr` to default protobuf representation. /// /// If required, a [`PhysicalExtensionCodec`] can be provided which can handle @@ -266,222 +273,22 @@ pub fn serialize_physical_expr_with_converter( ) -> Result { let expr = value.as_ref(); let expr_id = value.expression_id(); - // HashTableLookupExpr is used for dynamic filter pushdown in hash joins. - // It contains an Arc (the build-side hash table) which - // cannot be serialized - the hash table is a runtime structure built during - // execution on the build side. - // - // We replace it with lit(true) which is safe because: - // 1. The filter is a performance optimization, not a correctness requirement - // 2. lit(true) passes all rows, so no valid rows are incorrectly filtered out - // 3. The join itself will still produce correct results, just without the - // benefit of early filtering on the probe side - // - // In distributed execution, the remote worker won't have access to the hash - // table anyway, so the best we can do is skip this optimization. - if expr.downcast_ref::().is_some() { - let value = datafusion_proto_common::ScalarValue { - value: Some(datafusion_proto_common::scalar_value::Value::BoolValue( - true, - )), - }; - return Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::Literal(value)), - }); - } - if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::Column( - protobuf::PhysicalColumn { - name: expr.name().to_string(), - index: expr.index() as u32, - }, - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::UnknownColumn( - protobuf::UnknownColumn { - name: expr.name().to_string(), - }, - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - // Linearize a nested binary expression tree of the same operator - // into a flat vector of operands to avoid deep recursion in proto. - let op = expr.op(); - let mut operand_refs: Vec<&Arc> = vec![expr.right()]; - let mut current_expr: &BinaryExpr = expr; - loop { - match current_expr.left().downcast_ref::() { - Some(bin) if bin.op() == op => { - operand_refs.push(bin.right()); - current_expr = bin; - } - _ => { - operand_refs.push(current_expr.left()); - break; - } - } - } - - // Reverse so operands are ordered from left innermost to right outermost - operand_refs.reverse(); - - let operands = operand_refs - .iter() - .map(|e| proto_converter.physical_expr_to_proto(e, codec)) - .collect::>>()?; - - let binary_expr = Box::new(protobuf::PhysicalBinaryExprNode { - l: None, - r: None, - op: format!("{:?}", op), - operands, - }); + // Give the expression a chance to serialize itself first. Returning + // `Ok(Some(node))` lets expressions with private state (e.g. + // `DynamicFilterPhysicalExpr`) avoid exposing pub-for-proto accessors. + // `Ok(None)` falls through to the downcast chain below — that's the + // default for built-in expressions which haven't been migrated yet. + let encoder = ConverterEncoder { + codec, + proto_converter, + }; + let ctx = datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx::new(&encoder); + if let Some(node) = expr.try_to_proto(&ctx)? { + return Ok(node); + } - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::BinaryExpr( - binary_expr, - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some( - protobuf::physical_expr_node::ExprType::Case( - Box::new( - protobuf::PhysicalCaseNode { - expr: expr - .expr() - .map(|exp| { - proto_converter - .physical_expr_to_proto(exp, codec) - .map(Box::new) - }) - .transpose()?, - when_then_expr: expr - .when_then_expr() - .iter() - .map(|(when_expr, then_expr)| { - serialize_when_then_expr( - when_expr, - then_expr, - codec, - proto_converter, - ) - }) - .collect::, - DataFusionError, - >>()?, - else_expr: expr - .else_expr() - .map(|a| { - proto_converter - .physical_expr_to_proto(a, codec) - .map(Box::new) - }) - .transpose()?, - }, - ), - ), - ), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::NotExpr(Box::new( - protobuf::PhysicalNot { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(expr.arg(), codec)?, - )), - }, - ))), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::IsNullExpr( - Box::new(protobuf::PhysicalIsNull { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(expr.arg(), codec)?, - )), - }), - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::IsNotNullExpr( - Box::new(protobuf::PhysicalIsNotNull { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(expr.arg(), codec)?, - )), - }), - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::InList(Box::new( - protobuf::PhysicalInListNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(expr.expr(), codec)?, - )), - list: serialize_physical_exprs(expr.list(), codec, proto_converter)?, - negated: expr.negated(), - }, - ))), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::Negative(Box::new( - protobuf::PhysicalNegativeNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(expr.arg(), codec)?, - )), - }, - ))), - }) - } else if let Some(lit) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::Literal( - lit.value().try_into()?, - )), - }) - } else if let Some(cast) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::Cast(Box::new( - protobuf::PhysicalCastNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(cast.expr(), codec)?, - )), - arrow_type: Some(cast.cast_type().try_into()?), - }, - ))), - }) - } else if let Some(cast) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::TryCast(Box::new( - protobuf::PhysicalTryCastNode { - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(cast.expr(), codec)?, - )), - arrow_type: Some(cast.cast_type().try_into()?), - }, - ))), - }) - } else if let Some(expr) = expr.downcast_ref::() { + if let Some(expr) = expr.downcast_ref::() { let mut buf = Vec::new(); codec.try_encode_udf(expr.fun(), &mut buf)?; Ok(protobuf::PhysicalExprNode { @@ -500,84 +307,22 @@ pub fn serialize_physical_expr_with_converter( }, )), }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::LikeExpr(Box::new( - protobuf::PhysicalLikeExprNode { - negated: expr.negated(), - case_insensitive: expr.case_insensitive(), - expr: Some(Box::new( - proto_converter.physical_expr_to_proto(expr.expr(), codec)?, - )), - pattern: Some(Box::new( - proto_converter.physical_expr_to_proto(expr.pattern(), codec)?, - )), - }, - ))), - }) - } else if let Some(expr) = expr.downcast_ref::() { - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::HashExpr( - protobuf::PhysicalHashExprNode { - on_columns: serialize_physical_exprs( - expr.on_columns(), - codec, - proto_converter, - )?, - seed0: expr.seed(), - description: expr.description().to_string(), - }, - )), - }) - } else if let Some(expr) = expr.downcast_ref::() { + } else if let Some(expr) = expr.downcast_ref::() { + let mut buf = Vec::new(); + codec.try_encode_higher_order_function(expr.fun(), &mut buf)?; Ok(protobuf::PhysicalExprNode { expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::ScalarSubquery( - protobuf::PhysicalScalarSubqueryExprNode { - data_type: Some(expr.data_type().try_into()?), - nullable: expr.nullable(), - index: expr.index().as_usize() as u32, + expr_type: Some(protobuf::physical_expr_node::ExprType::HigherOrderUdf( + protobuf::PhysicalHigherOrderUdfNode { + name: expr.name().to_string(), + args: serialize_physical_exprs(expr.args(), codec, proto_converter)?, + fun_definition: (!buf.is_empty()).then_some(buf), }, )), }) - } else if let Some(df) = expr.downcast_ref::() { - let children = df - .original_children() - .iter() - .map(|child| proto_converter.physical_expr_to_proto(child, codec)) - .collect::>>()?; - - let remapped_children = if let Some(remapped) = df.remapped_children() { - remapped - .iter() - .map(|child| proto_converter.physical_expr_to_proto(child, codec)) - .collect::>>()? - } else { - vec![] - }; - - // Atomic snapshot of inner state. - let inner = df.inner(); - let inner_expr = - Box::new(proto_converter.physical_expr_to_proto(&inner.expr, codec)?); - - Ok(protobuf::PhysicalExprNode { - expr_id, - expr_type: Some(protobuf::physical_expr_node::ExprType::DynamicFilter( - Box::new(protobuf::PhysicalDynamicFilterNode { - children, - remapped_children, - generation: inner.generation, - inner_expr: Some(inner_expr), - is_complete: inner.is_complete, - }), - )), - }) } else { let mut buf: Vec = vec![]; - match codec.try_encode_expr(value, &mut buf) { + match codec.try_encode_expr(value, &mut buf, &ctx) { Ok(_) => { let inputs: Vec = value .children() @@ -603,92 +348,13 @@ pub fn serialize_partitioning( codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let serialized_partitioning = match partitioning { - Partitioning::RoundRobinBatch(partition_count) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::RoundRobin( - *partition_count as u64, - )), - }, - Partitioning::Hash(exprs, partition_count) => { - let serialized_exprs = - serialize_physical_exprs(exprs, codec, proto_converter)?; - protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Hash( - protobuf::PhysicalHashRepartition { - hash_expr: serialized_exprs, - partition_count: *partition_count as u64, - }, - )), - } - } - Partitioning::UnknownPartitioning(partition_count) => protobuf::Partitioning { - partition_method: Some(protobuf::partitioning::PartitionMethod::Unknown( - *partition_count as u64, - )), - }, + let encoder = ConverterEncoder { + codec, + proto_converter, }; - Ok(serialized_partitioning) -} - -fn serialize_when_then_expr( - when_expr: &Arc, - then_expr: &Arc, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result { - Ok(protobuf::PhysicalWhenThen { - when_expr: Some(proto_converter.physical_expr_to_proto(when_expr, codec)?), - then_expr: Some(proto_converter.physical_expr_to_proto(then_expr, codec)?), - }) -} - -impl TryFrom<&PartitionedFile> for protobuf::PartitionedFile { - type Error = DataFusionError; - - fn try_from(pf: &PartitionedFile) -> Result { - let last_modified = pf.object_meta.last_modified; - let last_modified_ns = last_modified.timestamp_nanos_opt().ok_or_else(|| { - DataFusionError::Plan(format!( - "Invalid timestamp on PartitionedFile::ObjectMeta: {last_modified}" - )) - })? as u64; - Ok(protobuf::PartitionedFile { - path: pf.object_meta.location.as_ref().to_owned(), - size: pf.object_meta.size, - last_modified_ns, - partition_values: pf - .partition_values - .iter() - .map(|v| v.try_into()) - .collect::, _>>()?, - range: pf.range.as_ref().map(|r| r.try_into()).transpose()?, - statistics: pf.statistics.as_ref().map(|s| s.as_ref().into()), - }) - } -} - -impl TryFrom<&FileRange> for protobuf::FileRange { - type Error = DataFusionError; - - fn try_from(value: &FileRange) -> Result { - Ok(protobuf::FileRange { - start: value.start, - end: value.end, - }) - } -} - -impl TryFrom<&[PartitionedFile]> for protobuf::FileGroup { - type Error = DataFusionError; - - fn try_from(gr: &[PartitionedFile]) -> Result { - Ok(protobuf::FileGroup { - files: gr - .iter() - .map(|f| f.try_into()) - .collect::, _>>()?, - }) - } + partitioning.try_to_proto( + &datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx::new(&encoder), + ) } pub fn serialize_file_scan_config( @@ -696,78 +362,11 @@ pub fn serialize_file_scan_config( codec: &dyn PhysicalExtensionCodec, proto_converter: &dyn PhysicalProtoConverterExtension, ) -> Result { - let file_groups = conf - .file_groups - .iter() - .map(|p| p.files().try_into()) - .collect::, _>>()?; - - let mut output_orderings = vec![]; - for order in &conf.output_ordering { - let ordering = - serialize_physical_sort_exprs(order.to_vec(), codec, proto_converter)?; - output_orderings.push(ordering) - } - - // Fields must be added to the schema so that they can persist in the protobuf, - // and then they are to be removed from the schema in `parse_protobuf_file_scan_config` - let mut fields = conf - .file_schema() - .fields() - .iter() - .cloned() - .collect::>(); - fields.extend(conf.table_partition_cols().iter().cloned()); - - let schema = Arc::new( - Schema::new(fields.clone()).with_metadata(conf.file_schema().metadata.clone()), - ); - - let projection_exprs = conf - .file_source - .projection() - .as_ref() - .map(|projection_exprs| { - let projections = projection_exprs.iter().cloned().collect::>(); - Ok::<_, DataFusionError>(protobuf::ProjectionExprs { - projections: projections - .into_iter() - .map(|expr| { - Ok(protobuf::ProjectionExpr { - alias: expr.alias.to_string(), - expr: Some( - proto_converter - .physical_expr_to_proto(&expr.expr, codec)?, - ), - }) - }) - .collect::>>()?, - }) - }) - .transpose()?; - - Ok(protobuf::FileScanExecConf { - file_groups, - statistics: Some((&conf.statistics()).into()), - limit: conf.limit.map(|l| protobuf::ScanLimit { limit: l as u32 }), - projection: vec![], - schema: Some(schema.as_ref().try_into()?), - table_partition_cols: conf - .table_partition_cols() - .iter() - .map(|x| x.name().clone()) - .collect::>(), - object_store_url: conf.object_store_url.to_string(), - output_ordering: output_orderings - .into_iter() - .map(|e| PhysicalSortExprNodeCollection { - physical_sort_expr_nodes: e, - }) - .collect::>(), - constraints: Some(conf.constraints.clone().into()), - batch_size: conf.batch_size.map(|s| s as u64), - projection_exprs, - }) + let encoder = ConverterPlanEncoder { + codec, + proto_converter, + }; + conf.try_to_proto(&ExecutionPlanEncodeCtx::new(&encoder)) } pub fn serialize_maybe_filter( @@ -783,6 +382,10 @@ pub fn serialize_maybe_filter( } } +#[deprecated( + since = "55.0.0", + note = "unused by DataFusion; `MemorySourceConfig` serializes its record batches itself via `DataSource::try_to_proto`" +)] pub fn serialize_record_batches(batches: &[RecordBatch]) -> Result> { if batches.is_empty() { return Ok(vec![]); @@ -796,86 +399,3 @@ pub fn serialize_record_batches(batches: &[RecordBatch]) -> Result> { writer.finish()?; Ok(buf) } - -impl TryFrom<&JsonSink> for protobuf::JsonSink { - type Error = DataFusionError; - - fn try_from(value: &JsonSink) -> Result { - Ok(Self { - config: Some(value.config().try_into()?), - writer_options: Some(value.writer_options().try_into()?), - }) - } -} - -impl TryFrom<&CsvSink> for protobuf::CsvSink { - type Error = DataFusionError; - - fn try_from(value: &CsvSink) -> Result { - Ok(Self { - config: Some(value.config().try_into()?), - writer_options: Some(value.writer_options().try_into()?), - }) - } -} - -#[cfg(feature = "parquet")] -impl TryFrom<&ParquetSink> for protobuf::ParquetSink { - type Error = DataFusionError; - - fn try_from(value: &ParquetSink) -> Result { - Ok(Self { - config: Some(value.config().try_into()?), - parquet_options: Some(value.parquet_options().try_into()?), - }) - } -} - -impl TryFrom<&FileSinkConfig> for protobuf::FileSinkConfig { - type Error = DataFusionError; - - fn try_from(conf: &FileSinkConfig) -> Result { - let file_groups = conf - .file_group - .iter() - .map(TryInto::try_into) - .collect::>>()?; - let table_paths = conf - .table_paths - .iter() - .map(ToString::to_string) - .collect::>(); - let table_partition_cols = conf - .table_partition_cols - .iter() - .map(|(name, data_type)| { - Ok(protobuf::PartitionColumn { - name: name.to_owned(), - arrow_type: Some(data_type.try_into()?), - }) - }) - .collect::>>()?; - let file_output_mode = match conf.file_output_mode { - datafusion_datasource::file_sink_config::FileOutputMode::Automatic => { - protobuf::FileOutputMode::Automatic - } - datafusion_datasource::file_sink_config::FileOutputMode::SingleFile => { - protobuf::FileOutputMode::SingleFile - } - datafusion_datasource::file_sink_config::FileOutputMode::Directory => { - protobuf::FileOutputMode::Directory - } - }; - Ok(Self { - object_store_url: conf.object_store_url.to_string(), - file_groups, - table_paths, - output_schema: Some(conf.output_schema.as_ref().try_into()?), - table_partition_cols, - keep_partition_by_columns: conf.keep_partition_by_columns, - insert_op: conf.insert_op as i32, - file_extension: conf.file_extension.to_string(), - file_output_mode: file_output_mode.into(), - }) - } -} diff --git a/datafusion/proto/tests/cases/mod.rs b/datafusion/proto/tests/cases/mod.rs index 3abbaccf79673..3f62fe223bdfd 100644 --- a/datafusion/proto/tests/cases/mod.rs +++ b/datafusion/proto/tests/cases/mod.rs @@ -21,8 +21,10 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion_common::plan_err; use datafusion_expr::function::AccumulatorArgs; use datafusion_expr::{ - Accumulator, AggregateUDFImpl, LimitEffect, PartitionEvaluator, ScalarFunctionArgs, - ScalarUDFImpl, Signature, Volatility, WindowUDFImpl, + Accumulator, AggregateUDFImpl, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDFImpl, LambdaParametersProgress, LimitEffect, + PartitionEvaluator, ScalarFunctionArgs, ScalarUDFImpl, Signature, ValueOrLambda, + Volatility, WindowUDFImpl, }; use datafusion_functions_window_common::field::WindowUDFFieldArgs; use datafusion_functions_window_common::partition::PartitionEvaluatorArgs; @@ -30,9 +32,11 @@ use std::fmt::Debug; use std::hash::Hash; use std::sync::Arc; +mod plans; +mod public_conversions; mod roundtrip_logical_plan; -mod roundtrip_physical_plan; mod serialize; +mod stack_safety; #[derive(Debug, PartialEq, Eq, Hash)] struct MyRegexUdf { @@ -180,3 +184,66 @@ pub(in crate::cases) struct CustomUDWFNode { #[prost(string, tag = "1")] pub payload: String, } + +#[derive(Debug, PartialEq, Eq, Hash)] +pub(in crate::cases) struct MyHigherOrderUDF { + signature: HigherOrderSignature, + pub payload: String, +} + +impl MyHigherOrderUDF { + pub fn new(payload: String) -> Self { + Self { + signature: HigherOrderSignature::any(2, Volatility::Immutable), + payload, + } + } +} + +impl HigherOrderUDFImpl for MyHigherOrderUDF { + fn name(&self) -> &str { + "higher_order_udf" + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn lambda_parameters( + &self, + _step: usize, + fields: &[ValueOrLambda>], + ) -> datafusion_common::Result { + let list = match fields.first() { + Some(ValueOrLambda::Value(field)) => field, + _ => return plan_err!("higher_order_udf expects a list as first argument"), + }; + let element = match list.data_type() { + DataType::List(field) | DataType::LargeList(field) => Arc::clone(field), + other => { + return plan_err!("higher_order_udf expected a list, got {other}"); + } + }; + Ok(LambdaParametersProgress::Complete(vec![vec![element]])) + } + + fn return_field_from_args( + &self, + _args: HigherOrderReturnFieldArgs, + ) -> datafusion_common::Result { + Ok(Arc::new(Field::new("", DataType::Int64, true))) + } + + fn invoke_with_args( + &self, + _args: HigherOrderFunctionArgs, + ) -> datafusion_common::Result { + unimplemented!() + } +} + +#[derive(Clone, PartialEq, ::prost::Message)] +pub(in crate::cases) struct MyHigherOrderUdfNode { + #[prost(string, tag = "1")] + pub payload: String, +} diff --git a/datafusion/proto/tests/cases/plans/aggregates.rs b/datafusion/proto/tests/cases/plans/aggregates.rs new file mode 100644 index 0000000000000..e57ac9fb5045b --- /dev/null +++ b/datafusion/proto/tests/cases/plans/aggregates.rs @@ -0,0 +1,361 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `AggregateExec`. + +use super::{roundtrip_test, roundtrip_test_with_context}; +use datafusion::arrow::array::ArrayRef; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::logical_expr::Volatility; +use datafusion::physical_expr::LexOrdering; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_optimizer::update_aggr_exprs::OptimizeAggregateOrder; +use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, +}; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{PhysicalSortExpr, col, lit}; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::{ExecutionPlan, PhysicalExpr}; +use datafusion::prelude::SessionContext; +use datafusion::scalar::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{DataFusionError, Result}; +use datafusion_expr::{ + Accumulator, AccumulatorFactoryFunction, AggregateUDF, Signature, SimpleAggregateUDF, +}; +use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; +use datafusion_functions_aggregate::array_agg::array_agg_udaf; +use datafusion_functions_aggregate::average::avg_udaf; +use datafusion_functions_aggregate::first_last::first_value_udaf; +use datafusion_functions_aggregate::nth_value::nth_value_udaf; +use datafusion_functions_aggregate::string_agg::string_agg_udaf; +use datafusion_functions_aggregate::sum::sum_udaf; +use datafusion_proto::physical_plan::{AsExecutionPlan, DefaultPhysicalExtensionCodec}; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::PhysicalPlanNode; +use prost::Message; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_aggregate() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let groups: Vec<(Arc, String)> = + vec![(col("a", &schema)?, "unused".to_string())]; + + let avg_expr = AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("AVG(b)") + .build()?; + let nth_expr = + AggregateExprBuilder::new(nth_value_udaf(), vec![col("b", &schema)?, lit(1u64)]) + .schema(Arc::clone(&schema)) + .alias("NTH_VALUE(b, 1)") + .build()?; + let str_agg_expr = + AggregateExprBuilder::new(string_agg_udaf(), vec![col("b", &schema)?, lit(1u64)]) + .schema(Arc::clone(&schema)) + .alias("NTH_VALUE(b, 1)") + .build()?; + + let test_cases = vec![ + // AVG + vec![Arc::new(avg_expr)], + // NTH_VALUE + vec![Arc::new(nth_expr)], + // STRING_AGG + vec![Arc::new(str_agg_expr)], + ]; + + for aggregates in test_cases { + let schema = schema.clone(); + roundtrip_test(Arc::new(AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new_single(groups.clone()), + aggregates, + vec![None], + Arc::new(EmptyExec::new(schema.clone())), + schema, + )?))?; + } + + Ok(()) +} + +#[test] +fn roundtrip_aggregate_preserves_optimizer_schema_and_reversed_state() -> Result<()> { + let input_schema = + Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); + let input_ordering = LexOrdering::new(vec![PhysicalSortExpr { + expr: col("b", &input_schema)?, + options: SortOptions::new(true, true), + }]) + .expect("single sort expression should form an ordering"); + let input: Arc = Arc::new(SortExec::new( + input_ordering, + Arc::new(EmptyExec::new(Arc::clone(&input_schema))), + )); + let original_name = "first_value(b) ORDER BY [b ASC NULLS LAST]"; + let aggregate_expr = Arc::new( + AggregateExprBuilder::new(first_value_udaf(), vec![col("b", &input_schema)?]) + .order_by(vec![PhysicalSortExpr { + expr: col("b", &input_schema)?, + options: SortOptions::new(false, false), + }]) + .schema(Arc::clone(&input_schema)) + .alias(original_name) + .build()?, + ); + let aggregate = AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new(vec![], vec![], vec![], false), + vec![aggregate_expr], + vec![None], + input, + input_schema, + )?; + + let optimized = OptimizeAggregateOrder::new() + .optimize(Arc::new(aggregate), &ConfigOptions::new())?; + let optimized_aggregate = optimized + .downcast_ref::() + .expect("expected optimized AggregateExec"); + assert_eq!(optimized.schema().field(0).name(), original_name); + assert_eq!( + optimized_aggregate.aggr_expr()[0].name(), + "last_value(b) ORDER BY [b DESC NULLS FIRST]" + ); + assert!(optimized_aggregate.aggr_expr()[0].is_reversed()); + + let codec = DefaultPhysicalExtensionCodec {}; + let node = PhysicalPlanNode::try_from_physical_plan(Arc::clone(&optimized), &codec)?; + let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let ctx = SessionContext::new(); + let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + let decoded_aggregate = decoded + .downcast_ref::() + .expect("expected decoded AggregateExec"); + + assert_eq!(optimized.schema(), decoded.schema()); + assert!(decoded_aggregate.aggr_expr()[0].is_reversed()); + Ok(()) +} + +#[test] +fn decode_aggregate_without_output_schema() -> Result<()> { + let input_schema = + Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); + let aggregate_expr = Arc::new( + AggregateExprBuilder::new(sum_udaf(), vec![col("b", &input_schema)?]) + .schema(Arc::clone(&input_schema)) + .alias("SUM(b)") + .build()?, + ); + let plan: Arc = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new(vec![], vec![], vec![], false), + vec![aggregate_expr], + vec![None], + Arc::new(EmptyExec::new(Arc::clone(&input_schema))), + input_schema, + )?); + + let codec = DefaultPhysicalExtensionCodec {}; + let mut node = PhysicalPlanNode::try_from_physical_plan(Arc::clone(&plan), &codec)?; + let Some(protobuf::physical_plan_node::PhysicalPlanType::Aggregate(aggregate)) = + node.physical_plan_type.as_mut() + else { + panic!("expected AggregateExecNode"); + }; + assert!(aggregate.schema.take().is_some()); + + let ctx = SessionContext::new(); + let decoded = node.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + assert_eq!(plan.schema(), decoded.schema()); + Ok(()) +} + +#[test] +fn roundtrip_aggregate_with_limit() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let groups: Vec<(Arc, String)> = + vec![(col("a", &schema)?, "unused".to_string())]; + + let aggregates = vec![ + AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("AVG(b)") + .build() + .map(Arc::new)?, + ]; + + let agg = AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new_single(groups.clone()), + aggregates, + vec![None], + Arc::new(EmptyExec::new(schema.clone())), + schema, + )?; + let agg = agg.with_limit_options(Some(LimitOptions::new_with_order(12, false))); + roundtrip_test(Arc::new(agg)) +} + +#[test] +fn roundtrip_aggregate_with_approx_pencentile_cont() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let groups: Vec<(Arc, String)> = + vec![(col("a", &schema)?, "unused".to_string())]; + + let aggregates = vec![ + AggregateExprBuilder::new( + approx_percentile_cont_udaf(), + vec![col("b", &schema)?, lit(0.5)], + ) + .schema(Arc::clone(&schema)) + .alias("APPROX_PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY b)") + .build() + .map(Arc::new)?, + ]; + + let agg = AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new_single(groups.clone()), + aggregates, + vec![None], + Arc::new(EmptyExec::new(schema.clone())), + schema, + )?; + roundtrip_test(Arc::new(agg)) +} + +#[test] +fn roundtrip_aggregate_with_sort() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let groups: Vec<(Arc, String)> = + vec![(col("a", &schema)?, "unused".to_string())]; + let sort_exprs = vec![PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }]; + + let aggregates = vec![ + AggregateExprBuilder::new(array_agg_udaf(), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("ARRAY_AGG(b)") + .order_by(sort_exprs) + .build() + .map(Arc::new)?, + ]; + + let agg = AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new_single(groups.clone()), + aggregates, + vec![None], + Arc::new(EmptyExec::new(schema.clone())), + schema, + )?; + roundtrip_test(Arc::new(agg)) +} + +#[test] +fn roundtrip_aggregate_udaf() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + #[derive(Debug)] + struct Example; + impl Accumulator for Example { + fn state(&mut self) -> Result> { + Ok(vec![ScalarValue::Int64(Some(0))]) + } + + fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { + Ok(()) + } + + fn evaluate(&mut self) -> Result { + Ok(ScalarValue::Int64(Some(0))) + } + + fn size(&self) -> usize { + 0 + } + } + + let return_type = DataType::Int64; + let accumulator: AccumulatorFactoryFunction = Arc::new(|_| Ok(Box::new(Example))); + + let udaf = AggregateUDF::from(SimpleAggregateUDF::new_with_signature( + "example", + Signature::exact(vec![DataType::Int64], Volatility::Immutable), + return_type, + accumulator, + vec![Field::new("value", DataType::Int64, true).into()], + )); + + let ctx = SessionContext::new(); + ctx.register_udaf(udaf.clone()); + + let groups: Vec<(Arc, String)> = + vec![(col("a", &schema)?, "unused".to_string())]; + + let aggregates = vec![ + AggregateExprBuilder::new(Arc::new(udaf), vec![col("b", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("example_agg") + .build() + .map(Arc::new)?, + ]; + + roundtrip_test_with_context( + Arc::new(AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new_single(groups.clone()), + aggregates, + vec![None], + Arc::new(EmptyExec::new(schema.clone())), + schema, + )?), + &ctx, + ) +} diff --git a/datafusion/proto/tests/cases/plans/dispatch.rs b/datafusion/proto/tests/cases/plans/dispatch.rs new file mode 100644 index 0000000000000..75f299107e358 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/dispatch.rs @@ -0,0 +1,322 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Serde dispatch itself: which hook the central (de)serializer reaches, +//! and how a custom converter or a deprecated shim participates. + +use super::roundtrip_test_and_return; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::execution::TaskContext; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{PhysicalSortExpr, col}; +use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use datafusion::prelude::SessionContext; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{Result, exec_datafusion_err}; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, PhysicalExtensionCodec, + PhysicalPlanDecodeContext, PhysicalPlanNodeExt, PhysicalProtoConverterExtension, +}; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; +use std::fmt::Formatter; +use std::sync::{Arc, RwLock}; +use std::vec; + +#[derive(Debug)] +struct DowncastDelegatingExec { + inner: Arc, +} + +impl DowncastDelegatingExec { + fn new(inner: Arc) -> Self { + Self { inner } + } +} + +impl DisplayAs for DowncastDelegatingExec { + fn fmt_as(&self, t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + self.inner.fmt_as(t, f) + } +} + +impl ExecutionPlan for DowncastDelegatingExec { + fn name(&self) -> &str { + self.inner.name() + } + + fn properties(&self) -> &Arc { + self.inner.properties() + } + + fn children(&self) -> Vec<&Arc> { + self.inner.children() + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + self.inner.apply_expressions(f) + } + + fn replace_children( + self: Arc, + children: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + let inner = Arc::clone(&self.inner).replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + Ok(Arc::new(Self::new(inner))) + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn downcast_delegate(&self) -> Option<&dyn ExecutionPlan> { + Some(self.inner.as_ref()) + } + + fn execute( + &self, + partition: usize, + context: Arc, + ) -> Result { + self.inner.execute(partition, context) + } +} + +#[test] +fn serialize_uses_downcast_delegate() -> Result<()> { + let inner: Arc = + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))); + let plan: Arc = Arc::new(DowncastDelegatingExec::new(inner)); + let codec = DefaultPhysicalExtensionCodec {}; + + let proto = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?; + + assert!(matches!( + proto.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::Empty(_)) + )); + + Ok(()) +} + +/// A wrapper delegating to a plan that serializes itself via the +/// `try_to_proto` hook must serialize as its delegate: the wrapper's default +/// hook returns `Ok(None)` and the delegate has no downcast-chain fallback. +#[test] +fn serialize_uses_downcast_delegate_for_self_serializing_plan() -> Result<()> { + let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]); + let input = Arc::new(EmptyExec::new(Arc::new(schema.clone()))); + let inner: Arc = Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: col("a", &schema)?, + alias: "a".to_string(), + }], + input, + )?); + let plan: Arc = Arc::new(DowncastDelegatingExec::new(inner)); + let codec = DefaultPhysicalExtensionCodec {}; + + let proto = PhysicalPlanNode::try_from_physical_plan(plan, &codec)?; + + assert!(matches!( + proto.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::Projection( + _ + )) + )); + + Ok(()) +} + +#[tokio::test] +async fn roundtrip_physical_plan_node() { + use datafusion::prelude::*; + use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, + }; + use datafusion_proto::protobuf::PhysicalPlanNode; + + let ctx = SessionContext::new(); + + ctx.register_parquet( + "pt", + &format!( + "{}/alltypes_plain.snappy.parquet", + datafusion_common::test_util::parquet_test_data() + ), + ParquetReadOptions::default(), + ) + .await + .unwrap(); + + let plan = ctx + .sql("select id, string_col, timestamp_col from pt where id > 4 order by string_col") + .await + .unwrap() + .create_physical_plan() + .await + .unwrap(); + + let node: PhysicalPlanNode = + PhysicalPlanNode::try_from_physical_plan(plan, &DefaultPhysicalExtensionCodec {}) + .unwrap(); + + let plan = node + .try_into_physical_plan(&ctx.task_ctx(), &DefaultPhysicalExtensionCodec {}) + .unwrap(); + + let _ = plan.execute(0, ctx.task_ctx()).unwrap(); +} + +#[test] +fn custom_proto_converter_intercepts() -> Result<()> { + #[derive(Default)] + struct CustomConverterInterceptor { + num_proto_plans: RwLock, + num_physical_plans: RwLock, + num_proto_exprs: RwLock, + num_physical_exprs: RwLock, + } + + impl PhysicalProtoConverterExtension for CustomConverterInterceptor { + fn proto_to_execution_plan( + &self, + proto: &PhysicalPlanNode, + ctx: &PhysicalPlanDecodeContext<'_>, + ) -> Result> { + { + let mut counter = self + .num_proto_plans + .write() + .map_err(|err| exec_datafusion_err!("{err}"))?; + *counter += 1; + } + self.default_proto_to_execution_plan(proto, ctx) + } + + fn execution_plan_to_proto( + &self, + plan: &Arc, + codec: &dyn PhysicalExtensionCodec, + ) -> Result + where + Self: Sized, + { + { + let mut counter = self + .num_physical_plans + .write() + .map_err(|err| exec_datafusion_err!("{err}"))?; + *counter += 1; + } + PhysicalPlanNode::try_from_physical_plan_with_converter( + Arc::clone(plan), + codec, + self, + ) + } + + fn proto_to_physical_expr( + &self, + proto: &PhysicalExprNode, + input_schema: &Schema, + ctx: &PhysicalPlanDecodeContext<'_>, + ) -> Result> + where + Self: Sized, + { + { + let mut counter = self + .num_proto_exprs + .write() + .map_err(|err| exec_datafusion_err!("{err}"))?; + *counter += 1; + } + self.default_proto_to_physical_expr(proto, input_schema, ctx) + } + + fn physical_expr_to_proto( + &self, + expr: &Arc, + codec: &dyn PhysicalExtensionCodec, + ) -> Result { + { + let mut counter = self + .num_physical_exprs + .write() + .map_err(|err| exec_datafusion_err!("{err}"))?; + *counter += 1; + } + serialize_physical_expr_with_converter(expr, codec, self) + } + } + + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let sort_exprs = [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + ] + .into(); + + let exec_plan = Arc::new(SortExec::new(sort_exprs, Arc::new(EmptyExec::new(schema)))); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = CustomConverterInterceptor::default(); + roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + + assert_eq!(*proto_converter.num_proto_exprs.read().unwrap(), 2); + assert_eq!(*proto_converter.num_physical_exprs.read().unwrap(), 2); + assert_eq!(*proto_converter.num_proto_plans.read().unwrap(), 2); + assert_eq!(*proto_converter.num_physical_plans.read().unwrap(), 2); + + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/dynamic_filters.rs b/datafusion/proto/tests/cases/plans/dynamic_filters.rs new file mode 100644 index 0000000000000..ee0ff9d8b1faf --- /dev/null +++ b/datafusion/proto/tests/cases/plans/dynamic_filters.rs @@ -0,0 +1,1126 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Dynamic filter expressions: their deduplication across a plan, and the +//! plans that produce them. + +use super::{roundtrip_test_and_return, roundtrip_test_sql_with_context}; +use arrow::array::RecordBatch; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::datasource::empty::EmptyTable; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::datasource::physical_plan::{ + FileGroup, FileScanConfig, FileScanConfigBuilder, ParquetSource, +}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::{JoinType, Operator}; +use datafusion::physical_expr::LexOrdering; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_optimizer::filter_pushdown::FilterPushdown; +use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{ + BinaryExpr, Column, DynamicFilterPhysicalExpr, PhysicalSortExpr, lit, +}; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode}; +use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::{ + ChildrenPropertiesMode, DisplayAs, DisplayFormatType, ExecutionPlan, PhysicalExpr, + PlanProperties, ReplaceChildrenOptions, SendableRecordBatchStream, +}; +use datafusion::prelude::SessionContext; +use datafusion_common::config::{ConfigOptions, TableParquetOptions}; +use datafusion_common::tree_node::TreeNodeRecursion; +use datafusion_common::{NullEquality, Result, internal_datafusion_err, internal_err}; +use datafusion_datasource::file::FileSource; +use datafusion_expr::ColumnarValue; +use datafusion_physical_expr::utils::reassign_expr_columns; +use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; +use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; +use datafusion_proto::bytes::{ + physical_plan_from_bytes_with_proto_converter, + physical_plan_to_bytes_with_proto_converter, +}; +use datafusion_proto::physical_plan::{ + DeduplicatingProtoConverter, DefaultPhysicalExtensionCodec, + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, + PhysicalProtoConverterExtension, +}; +use datafusion_proto::protobuf::PhysicalExprNode; +use prost::Message; +use std::fmt::{Display, Formatter}; +use std::sync::Arc; +use std::vec; + +/// Create a [`DynamicFilterPhysicalExpr`] with child column expression "a" @ index 0. +fn make_dynamic_filter() -> Arc { + Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0)) as Arc], + lit(true), + )) as Arc +} + +/// Update a [`DynamicFilterPhysicalExpr`]'s children to support child schema "b" @ 0, "a" @ 1. +fn make_reassigned_dynamic_filter( + filter: Arc, +) -> Result<(Arc, Arc)> { + let schema = Arc::new(Schema::new(vec![ + Field::new("b", DataType::Int64, false), + Field::new("a", DataType::Int64, false), + ])); + let reassigned = reassign_expr_columns(filter, &schema)?; + Ok((schema, reassigned)) +} + +/// Extract the expression id from a [`PhysicalExpr`] proto. Populated by the +/// default serializer from `PhysicalExpr::expression_id`. +fn proto_expression_id(expr: &PhysicalExprNode) -> u64 { + expr.expr_id + .expect("expected PhysicalExprNode.expr_id to be populated") +} + +/// Roundtrip a single physical expression shaped like so: +/// +/// ```text +/// BinaryExpr(AND) +/// / \ +/// filter_expr_1 filter_expr_2 +/// ``` +/// +/// Returns filter_expr_1 and filter_expr_2 after deserialization. +fn roundtrip_dynamic_filter_expr_pair( + filter_expr_1: Arc, + filter_expr_2: Arc, + schema: Arc, +) -> Result<(Arc, Arc)> { + let pair_expr = Arc::new(BinaryExpr::new( + Arc::clone(&filter_expr_1), + Operator::And, + Arc::clone(&filter_expr_2), + )) as Arc; + + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DeduplicatingProtoConverter {}; + let proto = converter.physical_expr_to_proto(&pair_expr, &codec)?; + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), &codec); + let deserialized_expr = + converter.proto_to_physical_expr(&proto, &schema, &decode_ctx)?; + + let binary = deserialized_expr + .downcast_ref::() + .expect("Expected BinaryExpr"); + + Ok((Arc::clone(binary.left()), Arc::clone(binary.right()))) +} + +/// Roundtrip an execution plan shaped like so: +/// +/// ```text +/// FilterExec(dynamic_filter_1 on a@0) +/// ProjectionExec(a := Column("a", source_index)) +/// DataSourceExec +/// ParquetSource(predicate = dynamic_filter_2) +/// ``` +/// +/// `dynamic_filter_1` and `dynamic_filter_2` are the same dynamic filter, except with +/// different children. +/// +/// Returns +/// - `dynamic_filter_1` before serialization +/// - `dynamic_filter_2` before serialization +/// - `dynamic_filter_1` after serialization +/// - `dynamic_filter_2` after serialization +#[expect(clippy::type_complexity)] +fn roundtrip_dynamic_filter_plan_pair() -> Result<( + Arc, + Arc, + Arc, + Arc, +)> { + let filter_expr_1 = make_dynamic_filter(); + let (data_source_schema, filter_expr_2) = + make_reassigned_dynamic_filter(Arc::clone(&filter_expr_1))?; + let left_before = Arc::clone(&filter_expr_1); + let right_before = Arc::clone(&filter_expr_2); + let file_source = Arc::new( + ParquetSource::new(Arc::clone(&data_source_schema)) + .with_predicate(Arc::clone(&filter_expr_2)), + ); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .build(); + let data_source_exec = + DataSourceExec::from_data_source(scan_config) as Arc; + + let projection_exec = Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(Column::new("a", 1)) as Arc, + alias: "a".to_string(), + }], + data_source_exec, + )?) as Arc; + let filter_exec = Arc::new(FilterExec::try_new( + Arc::clone(&filter_expr_1), + projection_exec, + )?) as Arc; + + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DeduplicatingProtoConverter {}; + let proto = converter.execution_plan_to_proto(&filter_exec, &codec)?; + + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), &codec); + let deserialized_plan = converter.proto_to_execution_plan(&proto, &decode_ctx)?; + + let outer_filter = deserialized_plan + .downcast_ref::() + .expect("Expected outer FilterExec"); + let left_filter = Arc::clone(outer_filter.predicate()); + let projection = outer_filter.children()[0] + .downcast_ref::() + .expect("Expected ProjectionExec"); + let data_source = projection + .input() + .downcast_ref::() + .expect("Expected DataSourceExec"); + let scan_config = data_source + .data_source() + .downcast_ref::() + .expect("Expected FileScanConfig"); + let right_filter = scan_config + .file_source() + .filter() + .expect("Expected pushed-down predicate"); + + Ok((left_before, right_before, left_filter, right_filter)) +} + +/// Takes two [`DynamicFilterPhysicalExpr`] and asserts that updates to one are visible +/// via the other. This helps assert that referential integrity is maintained after +/// deserializing. +fn assert_dynamic_filter_update_is_visible( + left_filter: &Arc, + right_filter: &Arc, +) -> Result<()> { + let left_filter = left_filter + .downcast_ref::() + .expect("Expected dynamic filter"); + let right_filter = right_filter + .downcast_ref::() + .expect("Expected dynamic filter"); + + // Sanity check that the filters have the same generation. + let original_generation = left_filter.snapshot_generation(); + assert_eq!(original_generation, right_filter.snapshot_generation(),); + + left_filter.update(lit(123_i64))?; + + // Assert that both generations updated. + assert_eq!(original_generation + 1, right_filter.snapshot_generation(),); + assert_eq!( + left_filter.snapshot_generation(), + right_filter.snapshot_generation(), + ); + + // Ensure both filters have the updated expr. + let expected_current = r#"Literal { value: Int64(123), field: Field { name: "lit", data_type: Int64 } }"#; + assert_eq!(expected_current, format!("{:?}", left_filter.current()?),); + assert_eq!(expected_current, format!("{:?}", right_filter.current()?),); + + Ok(()) +} + +/// Extract the dynamic-filter predicate that was pushed down to the parquet +/// scan at the bottom of the plan tree. +fn parquet_source_predicate(child: &Arc) -> Arc { + let data_source = child + .downcast_ref::() + .expect("Child should be DataSourceExec"); + let (_, parquet_source) = data_source + .downcast_to_file_source::() + .expect("Should be ParquetSource"); + parquet_source + .filter() + .expect("ParquetSource should have a predicate after roundtrip") +} + +/// Assert that two dynamic filters are equal both structurally (Debug output) +/// and by identity (`expression_id`). +fn assert_dynamic_filters_equal( + expected: &Arc, + actual: &Arc, +) { + // Structural. + let expected_dbg = format!("{expected:?}"); + let actual_dbg = format!("{actual:?}"); + if expected_dbg == actual_dbg { + return; + } + + // Note that the `DeduplicatingDeserializer` routes every cache hit through + // `with_new_children`. This produces an equivalent expression, but with + // remapped children that are equal to the original. Handle that case here. + let rewritten = Arc::clone(expected) + .with_new_children(expected.children().iter().map(|c| Arc::clone(c)).collect()) + .expect("with_new_children on a dynamic filter should not fail"); + assert_eq!(format!("{rewritten:?}"), actual_dbg); +} + +// Two clones of a dynamic filter expression should be deduped to the exact same expression. +#[test] +fn test_dynamic_filter_roundtrip_dedupe() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let filter_expr_1 = make_dynamic_filter(); + let filter_expr_2 = Arc::clone(&filter_expr_1); + + let (filter_expr_1_after_roundtrip, filter_expr_2_after_roundtrip) = + roundtrip_dynamic_filter_expr_pair( + Arc::clone(&filter_expr_1), + Arc::clone(&filter_expr_2), + schema, + )?; + + // Assert the filters are not modified during roundtrip. + assert_dynamic_filters_equal(&filter_expr_1, &filter_expr_1_after_roundtrip); + assert_dynamic_filters_equal(&filter_expr_2, &filter_expr_2_after_roundtrip); + assert_dynamic_filters_equal( + &filter_expr_1_after_roundtrip, + &filter_expr_2_after_roundtrip, + ); + + // Assert referential integrity. + assert_dynamic_filter_update_is_visible( + &filter_expr_1_after_roundtrip, + &filter_expr_2_after_roundtrip, + )?; + + Ok(()) +} + +/// Roundtrip test for an execution plan where there are multiple instances of a dynamic filter +/// with different children. +#[test] +fn test_dynamic_filter_plan_roundtrip_dedupe() -> Result<()> { + let ( + filter_expr_1, + filter_expr_2, + filter_expr_1_after_roundtrip, + filter_expr_2_after_roundtrip, + ) = roundtrip_dynamic_filter_plan_pair()?; + + // Assert the filters are not modified during roundtrip. + assert_dynamic_filters_equal(&filter_expr_1, &filter_expr_1_after_roundtrip); + assert_dynamic_filters_equal(&filter_expr_2, &filter_expr_2_after_roundtrip); + + // Assert referential integrity. + assert_dynamic_filter_update_is_visible( + &filter_expr_1_after_roundtrip, + &filter_expr_2_after_roundtrip, + )?; + + Ok(()) +} + +#[test] +fn test_dynamic_filter_expression_id_is_stable_between_serializations() -> Result<()> { + let filter_expr = make_dynamic_filter(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DeduplicatingProtoConverter {}; + + let proto1 = proto_converter.physical_expr_to_proto(&filter_expr, &codec)?; + let expr_id1 = proto_expression_id(&proto1); + + let proto2 = proto_converter.physical_expr_to_proto(&filter_expr, &codec)?; + let expr_id2 = proto_expression_id(&proto2); + + assert_eq!( + expr_id1, expr_id2, + "Expected the same dynamic filter expression id across serializations" + ); + + Ok(()) +} + +/// Create a DataSourceExec backed by a ParquetSource that accepts filter pushdown, +/// along with a ConfigOptions that enables all dynamic filter pushdown options. +fn datasource_for_dynamic_filter_pushdown( + schema: &Arc, +) -> (Arc, ConfigOptions) { + let mut parquet_options = TableParquetOptions::new(); + parquet_options.global.pushdown_filters = true; + let source = Arc::new( + ParquetSource::new(Arc::clone(schema)) + .with_table_parquet_options(parquet_options), + ); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) + .with_file(PartitionedFile::new("/path/to/file.parquet", 1024)) + .build(); + + let mut config = ConfigOptions::default(); + config.execution.parquet.pushdown_filters = true; + config.optimizer.enable_join_dynamic_filter_pushdown = true; + config.optimizer.enable_aggregate_dynamic_filter_pushdown = true; + config.optimizer.enable_topk_dynamic_filter_pushdown = true; + + (DataSourceExec::from_data_source(scan_config), config) +} + +/// Test that plan containing a HashJoinExec with dynamic filter pushdown +/// can be serialized and deserialized while preserving references to the dynamic filter. +#[test] +fn test_hash_join_with_dynamic_filter_roundtrip() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); + + let left_child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let (right_child, config) = datasource_for_dynamic_filter_pushdown(&schema); + + let on: Vec<(Arc, Arc)> = vec![( + Arc::new(Column::new("col", 0)), + Arc::new(Column::new("col", 0)), + )]; + + let hash_join = Arc::new(HashJoinExec::try_new( + left_child, + right_child, + on, + None, + &JoinType::Inner, + None, + PartitionMode::CollectLeft, + NullEquality::NullEqualsNothing, + false, + )?) as Arc; + + // Run the optimizer rule for filter pushdown. + let optimizer = FilterPushdown::new_post_optimization(); + let plan = optimizer.optimize(hash_join, &config)?; + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DeduplicatingProtoConverter {}; + let deserialized = roundtrip_test_and_return(plan, &ctx, &codec, &converter)?; + + // Extract the deserialized HashJoinExec and its dynamic filter. + let deserialized_join = deserialized + .downcast_ref::() + .expect("Should be HashJoinExec"); + let deserialized_hash_join_df = deserialized_join + .dynamic_expressions_produced() + .into_iter() + .next() + .expect("HashJoinExec should have a dynamic filter after roundtrip"); + + // Extract the dynamic filter pushed down to the probe side's ParquetSource. + let deserialized_predicate = parquet_source_predicate(deserialized_join.right()); + + // The HashJoinExec's dynamic filter and the probe side's predicate should + // refer to the same underlying expression. + let plan_df = deserialized_hash_join_df; + assert_dynamic_filters_equal(&plan_df, &deserialized_predicate); + assert_dynamic_filter_update_is_visible(&plan_df, &deserialized_predicate)?; + + Ok(()) +} + +/// returns a SessionContext with an empty `netflow` table registered +fn netflow_context() -> Result { + let ctx = SessionContext::new(); + let schema = Arc::new(Schema::new(vec![ + Field::new("dst_geo_country_name", DataType::Utf8, true), + Field::new("dst_geo_city_name", DataType::Utf8, true), + Field::new("packets", DataType::UInt64, true), + Field::new("src_addr", DataType::Utf8, true), + Field::new("dst_addr", DataType::Utf8, true), + ])); + + ctx.register_table("netflow", Arc::new(EmptyTable::new(schema)))?; + + Ok(ctx) +} + +/// Regression test for issue #18602: +/// https://github.com/apache/datafusion/issues/18602 +/// +/// The physical filter expression here contains a long chain of `AND` predicates. +/// Before linearizing `PhysicalBinaryExprNode`, encoding then decoding the protobuf +/// could fail with `DecodeError: recursion limit reached`. +#[tokio::test] +async fn roundtrip_issue_18602_complex_filter_decode_recursion() -> Result<()> { + let ctx = netflow_context()?; + let sql = "SELECT \ + dst_geo_country_name AS x_axis_1, \ + dst_geo_city_name AS x_axis_2, \ + sum(packets) AS y_axis_1 \ + FROM netflow \ + WHERE dst_geo_country_name IS NOT NULL \ + AND src_addr NOT LIKE '10.201.%' \ + AND dst_addr NOT LIKE '10.201.%' \ + AND src_addr NOT LIKE '10.202.%' \ + AND dst_addr NOT LIKE '10.202.%' \ + AND src_addr NOT LIKE '10.203.%' \ + AND dst_addr NOT LIKE '10.203.%' \ + AND src_addr NOT LIKE '10.204.%' \ + AND dst_addr NOT LIKE '10.204.%' \ + AND src_addr NOT LIKE '172.16.186.%' \ + AND dst_addr NOT LIKE '172.16.186.%' \ + AND src_addr NOT LIKE '172.16.187.%' \ + AND dst_addr NOT LIKE '172.16.187.%' \ + AND src_addr NOT LIKE '172.16.188.%' \ + AND dst_addr NOT LIKE '172.16.188.%' \ + AND src_addr NOT LIKE '10.102.45.%' \ + AND dst_addr NOT LIKE '10.102.45.%' \ + AND src_addr NOT LIKE '172.25.210.%' \ + AND dst_addr NOT LIKE '172.25.210.%' \ + AND src_addr NOT LIKE '172.25.211.%' \ + AND dst_addr NOT LIKE '172.25.211.%' \ + AND src_addr NOT LIKE '141.226.101.%' \ + AND dst_addr NOT LIKE '141.226.101.%' \ + AND src_addr NOT LIKE '167.86.40.%' \ + AND dst_addr NOT LIKE '167.86.40.%' \ + AND src_addr NOT LIKE '66.22.38.%' \ + AND dst_addr NOT LIKE '66.22.38.%' \ + AND src_addr != '168.143.191.55' \ + AND dst_addr != '168.143.191.55' \ + AND src_addr != '82.112.107.142' \ + AND dst_addr != '82.112.107.142' \ + AND src_addr != '20.76.39.176' \ + AND dst_addr != '20.76.39.176' \ + AND src_addr != '162.159.129.83' \ + AND dst_addr != '162.159.129.83' \ + AND src_addr != '34.201.223.155' \ + AND dst_addr != '34.201.223.155' \ + AND src_addr != '34.201.223.156' \ + AND dst_addr != '34.201.223.156' \ + AND src_addr != '34.201.223.157' \ + AND dst_addr != '34.201.223.157' \ + AND src_addr != '134.201.223.157' \ + AND dst_addr != '134.201.223.157' \ + AND src_addr != '341.201.223.157' \ + AND dst_addr != '341.201.223.157' \ + GROUP BY x_axis_1, x_axis_2 \ + ORDER BY y_axis_1 DESC \ + LIMIT 20"; + + roundtrip_test_sql_with_context(sql, &ctx).await +} + +/// Test that plan containing a AggregateExec with dynamic filter pushdown +/// can be serialized and deserialized while preserving references to the dynamic filter. +#[test] +fn test_aggregate_with_dynamic_filter_roundtrip() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let col_a: Arc = Arc::new(Column::new("a", 0)); + + let (child, config) = datasource_for_dynamic_filter_pushdown(&schema); + + let agg = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![]), + vec![ + AggregateExprBuilder::new( + datafusion::functions_aggregate::min_max::min_udaf(), + vec![Arc::clone(&col_a)], + ) + .schema(Arc::clone(&schema)) + .alias("min_a") + .build() + .map(Arc::new)?, + ], + vec![None], + child, + Arc::clone(&schema), + )?) as Arc; + + // Run the optimizer rule for filter pushdown. + let optimizer = FilterPushdown::new_post_optimization(); + let plan = optimizer.optimize(agg, &config)?; + + // Roundtrip with deduplication. + // + // Note: We don't use `roundtrip_test_and_return` here because there's a + // pre-existing issue with PhysicalGroupBy serialization where empty groups + // `[[]]` become `[]` after roundtrip. This behavior is unrelated to this test. + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DeduplicatingProtoConverter {}; + let bytes = physical_plan_to_bytes_with_proto_converter( + Arc::clone(&plan), + &codec, + &converter, + )?; + let deserialized = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + &codec, + &converter, + )?; + + // Extract the deserialized AggregateExec and its dynamic filter. + let deserialized_agg = deserialized + .downcast_ref::() + .expect("Should be AggregateExec"); + let deserialized_agg_df = deserialized_agg + .dynamic_expressions_produced() + .into_iter() + .next() + .expect("AggregateExec should have a dynamic filter after roundtrip"); + + // Extract the dynamic filter pushed down to the child ParquetSource. + let deserialized_predicate = parquet_source_predicate(deserialized_agg.input()); + + // The AggregateExec's dynamic filter and the child's predicate should + // refer to the same underlying expression. + let plan_df = deserialized_agg_df; + assert_dynamic_filters_equal(&plan_df, &deserialized_predicate); + assert_dynamic_filter_update_is_visible(&plan_df, &deserialized_predicate)?; + + Ok(()) +} + +#[test] +fn test_aggregate_without_dynamic_filter_roundtrip() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let col_a: Arc = Arc::new(Column::new("a", 0)); + let child = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Partial, + PhysicalGroupBy::new_single(vec![]), + vec![ + AggregateExprBuilder::new( + datafusion::functions_aggregate::min_max::min_udaf(), + vec![col_a], + ) + .schema(Arc::clone(&schema)) + .alias("min_a") + .build() + .map(Arc::new)?, + ], + vec![None], + child, + Arc::clone(&schema), + )?) as Arc; + + let mut config = ConfigOptions::default(); + config.optimizer.enable_aggregate_dynamic_filter_pushdown = true; + let optimizer = FilterPushdown::new_post_optimization(); + let plan = optimizer.optimize(aggregate, &config)?; + assert!( + plan.downcast_ref::() + .expect("Should be AggregateExec") + .dynamic_expressions_produced() + .is_empty() + ); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DefaultPhysicalProtoConverter {}; + let bytes = physical_plan_to_bytes_with_proto_converter(plan, &codec, &converter)?; + let deserialized = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + &codec, + &converter, + )?; + + assert!( + deserialized + .downcast_ref::() + .expect("Should be AggregateExec") + .dynamic_expressions_produced() + .is_empty() + ); + Ok(()) +} + +/// Test that plan containing a SortExec with dynamic filter pushdown +/// can be serialized and deserialized while preserving references to the dynamic filter. +#[test] +fn test_sort_topk_with_dynamic_filter_roundtrip() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let col_a: Arc = Arc::new(Column::new("a", 0)); + + let (child, config) = datasource_for_dynamic_filter_pushdown(&schema); + + let sort = Arc::new( + SortExec::new( + LexOrdering::new(vec![PhysicalSortExpr { + expr: Arc::clone(&col_a), + options: SortOptions::default(), + }]) + .unwrap(), + child, + ) + .with_fetch(Some(10)), + ) as Arc; + + // Verify the optimizer kept the dynamic filter on the SortExec. + let optimizer = FilterPushdown::new_post_optimization(); + let plan = optimizer.optimize(sort, &config)?; + + // Roundtrip with deduplication. + // + // Note: We don't use `roundtrip_test_and_return` here because + // `DeduplicatingDeserializer` rewrites cache hits via `with_new_children`, + // which sets `remapped_children: Some(...)` on the second encounter of a + // shared `DynamicFilterPhysicalExpr`. SortExec's `Debug` includes its + // dynamic filter, so the original-vs-deserialized structural equality check + // would fail purely on this artifact. + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DeduplicatingProtoConverter {}; + let bytes = physical_plan_to_bytes_with_proto_converter( + Arc::clone(&plan), + &codec, + &converter, + )?; + let deserialized = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + &codec, + &converter, + )?; + + // Extract the deserialized SortExec and its dynamic filter. + let deserialized_sort = deserialized + .downcast_ref::() + .expect("Should be SortExec"); + let deserialized_sort_df = deserialized_sort + .dynamic_expressions_produced() + .into_iter() + .next() + .expect("SortExec should have a dynamic filter after roundtrip"); + + // Extract the dynamic filter pushed down to the child ParquetSource. + let deserialized_predicate = parquet_source_predicate(deserialized_sort.input()); + + // The SortExec's dynamic filter and the child's predicate should + // refer to the same underlying expression. + let plan_df = deserialized_sort_df; + assert_dynamic_filters_equal(&plan_df, &deserialized_predicate); + assert_dynamic_filter_update_is_visible(&plan_df, &deserialized_predicate)?; + + Ok(()) +} + +/// A custom [`ExecutionPlan`] which stores [`PhysicalExpr`]s. +struct CustomExecWithExprs { + exprs: Vec>, + child: Arc, +} + +#[derive(Clone, PartialEq, Message)] +struct CustomExecWithExprsProto { + #[prost(message, repeated, tag = "1")] + exprs: Vec, +} + +impl std::fmt::Debug for CustomExecWithExprs { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + f.debug_struct("CustomExecWithExprs") + .field("exprs", &self.exprs) + .field("child", &self.child) + .finish() + } +} + +impl CustomExecWithExprs { + fn new(exprs: Vec>, child: Arc) -> Self { + Self { exprs, child } + } +} + +impl DisplayAs for CustomExecWithExprs { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "CustomExecWithExprs") + } +} + +impl ExecutionPlan for CustomExecWithExprs { + fn name(&self) -> &str { + "CustomExecWithExprs" + } + + fn schema(&self) -> SchemaRef { + self.child.schema() + } + + fn properties(&self) -> &Arc { + self.child.properties() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.child] + } + + fn apply_expressions( + &self, + f: &mut dyn FnMut(&Arc) -> Result, + ) -> Result { + datafusion_physical_plan::apply_expression_roots(&self.exprs, f) + } + + fn replace_children( + self: Arc, + _: Vec>, + _: ReplaceChildrenOptions, + ) -> Result> { + unreachable!() + } + + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children( + children, + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + ) + } + + fn execute( + &self, + _partition: usize, + _context: Arc, + ) -> Result { + unreachable!() + } +} + +/// A [`PhysicalExtensionCodec`] for [`CustomExecWithExprs`]. +#[derive(Debug)] +struct CustomExecWithExprsCodec {} + +impl PhysicalExtensionCodec for CustomExecWithExprsCodec { + fn try_decode( + &self, + buf: &[u8], + inputs: &[Arc], + ctx: &TaskContext, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + let decode_ctx = PhysicalPlanDecodeContext::new(ctx, self); + let input_schema = inputs[0].schema(); + let proto = CustomExecWithExprsProto::decode(buf) + .map_err(|e| internal_datafusion_err!("Failed to decode custom exec: {e}"))?; + let exprs = proto + .exprs + .iter() + .map(|expr_proto| { + proto_converter.proto_to_physical_expr( + expr_proto, + input_schema.as_ref(), + &decode_ctx, + ) + }) + .collect::>>()?; + + Ok(Arc::new(CustomExecWithExprs::new(exprs, inputs[0].clone()))) + } + + fn try_encode( + &self, + node: Arc, + buf: &mut Vec, + proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + let custom = node + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected CustomExecWithExprs"))?; + let proto = CustomExecWithExprsProto { + exprs: custom + .exprs + .iter() + .map(|expr| proto_converter.physical_expr_to_proto(expr, self)) + .collect::>>()?, + }; + proto + .encode(buf) + .map_err(|e| internal_datafusion_err!("Failed to encode custom exec: {e}"))?; + + Ok(()) + } +} + +/// Tests that a custom [`ExecutionPlan`] with [`PhysicalExpr`] can +/// dedupe dynamic filters by using the proto converter in its +/// [`PhysicalExtensionCodec`] implementation. +#[test] +fn test_custom_node_with_dynamic_filter_dedup_roundtrip() -> Result<()> { + // Create the plan: + // + // FilterExec(dynamic_filter) + // -> CustomExecWithExprs(exprs: [dynamic_filter]) + // -> EmptyExec + // + // The same dynamic filter expression is saved in both the FilterExec and CustomExecWithExprs. + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let dynamic_filter = Arc::new(DynamicFilterPhysicalExpr::new( + vec![Arc::new(Column::new("a", 0)) as Arc], + lit(true), + )); + let dynamic_filter_expr: Arc = dynamic_filter; + + let empty = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let custom_exec = Arc::new(CustomExecWithExprs::new( + vec![Arc::clone(&dynamic_filter_expr)], + empty, + )); + let filter_exec = Arc::new(FilterExec::try_new( + Arc::clone(&dynamic_filter_expr), + custom_exec, + )?) as Arc; + + // Roundtrip with DeduplicatingProtoConverter + let codec = CustomExecWithExprsCodec {}; + let converter = DeduplicatingProtoConverter {}; + + let bytes = physical_plan_to_bytes_with_proto_converter( + Arc::clone(&filter_exec), + &codec, + &converter, + )?; + + let ctx = SessionContext::new(); + let deser_converter = DeduplicatingProtoConverter {}; + let deserialized = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + &codec, + &deser_converter, + )?; + + // Extract the deserialized FilterExec's dynamic filter + let deser_filter = deserialized + .downcast_ref::() + .expect("Top-level should be FilterExec"); + let deser_filter_df = deser_filter.predicate(); + + // Extract the deserialized custom node's dynamic filter + let deser_custom = deser_filter + .input() + .downcast_ref::() + .expect("FilterExec child should be CustomExecWithExprs"); + assert_eq!(deser_custom.exprs.len(), 1, "Should have one expression"); + let [deser_custom_df] = deser_custom.exprs.as_slice() else { + return internal_err!("Custom node should have one expression"); + }; + + // Pass the un-remapped filter first so the helper's `with_new_children` + // rewrite can reconstruct the remapped form on the other side. + assert_dynamic_filters_equal(deser_custom_df, deser_filter_df); + assert_dynamic_filter_update_is_visible(deser_custom_df, deser_filter_df)?; + + Ok(()) +} + +/// A custom `PhysicalExpr` whose extension codec embeds a nested +/// `PhysicalExprNode` *inside its own blob* (rather than the standard +/// `PhysicalExtensionExprNode.inputs` field). This is the case that only +/// works if the expr-level codec methods receive the encode/decode context. +#[derive(Debug)] +struct WrapperExpr { + inner: Arc, +} + +impl Display for WrapperExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "WrapperExpr({})", self.inner) + } +} + +impl PartialEq for WrapperExpr { + fn eq(&self, other: &Self) -> bool { + self.inner.eq(&other.inner) + } +} + +impl Eq for WrapperExpr {} + +impl std::hash::Hash for WrapperExpr { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + } +} + +impl PhysicalExpr for WrapperExpr { + fn data_type(&self, input_schema: &Schema) -> Result { + self.inner.data_type(input_schema) + } + fn nullable(&self, input_schema: &Schema) -> Result { + self.inner.nullable(input_schema) + } + fn evaluate(&self, _batch: &RecordBatch) -> Result { + internal_err!("WrapperExpr is not executable in this test") + } + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + Ok(Arc::new(WrapperExpr { + inner: Arc::clone(&children[0]), + })) + } + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } +} + +/// Wire layout for [`WrapperExpr`]: a single nested `PhysicalExprNode`. +#[derive(Clone, PartialEq, prost::Message)] +struct WrapperExprProto { + #[prost(message, optional, boxed, tag = "1")] + inner: Option>, +} + +#[derive(Debug)] +struct WrapperCodec; + +impl PhysicalExtensionCodec for WrapperCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + internal_err!("not used") + } + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + internal_err!("not used") + } + fn try_decode_expr( + &self, + buf: &[u8], + _inputs: &[Arc], + ctx: &PhysicalExprDecodeCtx<'_>, + ) -> Result> { + let proto = WrapperExprProto::decode(buf) + .map_err(|e| internal_datafusion_err!("decode WrapperExprProto: {e}"))?; + let inner_proto = proto + .inner + .ok_or_else(|| internal_datafusion_err!("missing inner"))?; + // Decode the nested expr through the context so it resolves against + // the real schema/registry AND participates in dedup — no fabricated + // `SessionContext` or hard-coded schema required. + let inner = ctx.decode(&inner_proto)?; + Ok(Arc::new(WrapperExpr { inner })) + } + fn try_encode_expr( + &self, + node: &Arc, + buf: &mut Vec, + ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result<()> { + let wrapper = node + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("not WrapperExpr"))?; + // Encode the nested expr through the context so an active + // `DeduplicatingProtoConverter` stamps a matching `expr_id`. + let inner_proto = ctx.encode_child(&wrapper.inner)?; + let proto = WrapperExprProto { + inner: Some(Box::new(inner_proto)), + }; + proto + .encode(buf) + .map_err(|e| internal_datafusion_err!("encode WrapperExprProto: {e}"))?; + Ok(()) + } +} + +/// A `DynamicFilterPhysicalExpr` referenced both as a bare expression and +/// nested inside a custom expression's codec blob must reconstruct to a +/// single shared `Inner` after roundtrip. +/// +/// This exercises the expr-level codec hooks receiving the encode/decode +/// context: `try_encode_expr` routes its nested `PhysicalExprNode` through +/// `ctx.encode_child` and `try_decode_expr` through `ctx.decode`, so the +/// nested filter picks up the same `DeduplicatingProtoConverter` / +/// `DeduplicatingDeserializer` cache as the bare reference. Without the +/// context the nested expr would serialize with `expr_id: None` and decode +/// into a distinct `Inner`, breaking heap-max propagation across the +/// extension boundary in distributed execution. +#[test] +fn extension_codec_expr_participates_in_deduplication() -> Result<()> { + use prost::Message; + + // A single composite expression holding TWO references to the same + // dynamic filter: bare on the left of an AND, wrapped on the right. + let dyn_filter = make_dynamic_filter(); + let wrapper: Arc = Arc::new(WrapperExpr { + inner: Arc::clone(&dyn_filter), + }); + let composite: Arc = Arc::new(BinaryExpr::new( + Arc::clone(&dyn_filter), + Operator::And, + Arc::clone(&wrapper), + )); + + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let codec = WrapperCodec; + let converter = DeduplicatingProtoConverter {}; + + // Encode, then round-trip through prost bytes to mimic the wire. + let proto = converter.physical_expr_to_proto(&composite, &codec)?; + let bytes = proto.encode_to_vec(); + let decoded_proto = PhysicalExprNode::decode(bytes.as_slice()).unwrap(); + + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), &codec); + let decoded = + converter.proto_to_physical_expr(&decoded_proto, &schema, &decode_ctx)?; + + let binary = decoded + .downcast_ref::() + .expect("must decode back to BinaryExpr"); + let decoded_left = Arc::clone(binary.left()); + let decoded_right = Arc::clone(binary.right()); + let decoded_wrapper = decoded_right + .downcast_ref::() + .expect("right side must decode back to WrapperExpr"); + + // The load-bearing check: an `update()` on the bare-side filter must be + // observable from the wrapped-side filter, proving both refs back the + // same `Inner`. + assert_dynamic_filter_update_is_visible(&decoded_left, &decoded_wrapper.inner)?; + + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/exprs.rs b/datafusion/proto/tests/cases/plans/exprs.rs new file mode 100644 index 0000000000000..518b4a62ce072 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/exprs.rs @@ -0,0 +1,455 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Physical expressions embedded in plans, including the binary +//! expression linearization. + +use super::{roundtrip_test, roundtrip_test_and_return}; +use arrow::datatypes::Fields; +use datafusion::arrow::compute::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema}; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::expressions::Literal; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{ + BinaryExpr, Column, PhysicalSortExpr, binary, col, like, lit, +}; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion::physical_plan::repartition::RangeExpr; +use datafusion::physical_plan::{ + ExecutionPlan, PhysicalExpr, RangePartitioning, SplitPoint, +}; +use datafusion::prelude::SessionContext; +use datafusion::scalar::ScalarValue; +use datafusion_common::Result; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + PhysicalProtoConverterExtension, +}; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::PhysicalPlanNode; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_date_time_interval() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("some_date", DataType::Date32, false), + Field::new( + "some_interval", + DataType::Interval(IntervalUnit::DayTime), + false, + ), + ]); + let input = Arc::new(EmptyExec::new(Arc::new(schema.clone()))); + let date_expr = col("some_date", &schema)?; + let literal_expr = col("some_interval", &schema)?; + let date_time_interval_expr = + binary(date_expr, Operator::Plus, literal_expr, &schema)?; + let plan = Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: date_time_interval_expr, + alias: "result".to_string(), + }], + input, + )?); + roundtrip_test(plan) +} + +#[test] +fn roundtrip_like() -> Result<()> { + let schema = Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Utf8, false), + ]); + let input = Arc::new(EmptyExec::new(Arc::new(schema.clone()))); + let like_expr = like( + false, + false, + col("a", &schema)?, + col("b", &schema)?, + &schema, + )?; + let plan = Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: like_expr, + alias: "result".to_string(), + }], + input, + )?); + roundtrip_test(plan) +} + +/// Test that HashTableLookupExpr serializes to lit(true) +/// +/// HashTableLookupExpr contains a runtime hash table that cannot be serialized. +/// The serialization code replaces it with lit(true) which is safe because +/// it's a performance optimization filter, not a correctness requirement. +#[test] +fn roundtrip_hash_table_lookup_expr_to_lit() -> Result<()> { + use datafusion::physical_plan::joins::join_hash_map::JoinHashMapU32; + use datafusion::physical_plan::joins::{HashTableLookupExpr, Map}; + + // Create a simple schema and input plan + let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(schema.clone())); + + // Create a HashTableLookupExpr - it will be replaced with lit(true) during serialization + let hash_map = Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(0)))); + let on_columns = vec![col("col", &schema)?]; + let lookup_expr: Arc = Arc::new(HashTableLookupExpr::new( + on_columns, + datafusion::physical_plan::joins::SeededRandomState::with_seed(0), + hash_map, + "test_lookup".to_string(), + )); + + // Create a filter with the lookup expression + let filter = Arc::new(FilterExec::try_new(lookup_expr, input)?); + + // Serialize + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + + let proto: PhysicalPlanNode = + PhysicalPlanNode::try_from_physical_plan(filter.clone(), &codec) + .expect("serialization should succeed"); + + // Deserialize + let result: Arc = proto + .try_into_physical_plan(&ctx.task_ctx(), &codec) + .expect("deserialization should succeed"); + + // The deserialized plan should have lit(true) instead of HashTableLookupExpr + // Verify the filter predicate is a Literal(true) + let result_filter = result.downcast_ref::().unwrap(); + let predicate = result_filter.predicate(); + let literal = predicate.downcast_ref::().unwrap(); + assert_eq!(*literal.value(), ScalarValue::Boolean(Some(true))); + + Ok(()) +} + +#[test] +fn roundtrip_hash_expr() -> Result<()> { + use datafusion::physical_plan::joins::{HashExpr, SeededRandomState}; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Utf8, false), + ])); + + // Create a HashExpr with test columns and seeds + let on_columns = vec![col("a", &schema)?, col("b", &schema)?]; + let hash_expr: Arc = Arc::new(HashExpr::new( + on_columns, + SeededRandomState::with_seed(0), // arbitrary random seed for testing + "test_hash".to_string(), + )); + + // Wrap in a filter by comparing hash value to a literal + // hash_expr > 0 is always boolean + let filter_expr = binary(hash_expr, Operator::Gt, lit(0u64), &schema)?; + let filter = Arc::new(FilterExec::try_new( + filter_expr, + Arc::new(EmptyExec::new(schema)), + )?); + + // Confirm that the debug string contains the random state seeds + assert!( + format!("{filter:?}").contains("test_hash(a@0, b@1, [0])"), + "Debug string missing seeds: {filter:?}" + ); + roundtrip_test(filter) +} + +#[test] +fn roundtrip_range_expr() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Float64, false), + Field::new("b", DataType::Float64, false), + ])); + let options = [SortOptions::new(true, true), SortOptions::new(false, false)]; + let range_partitioning = RangePartitioning::try_new( + [ + PhysicalSortExpr::new(col("a", &schema)?, options[0]), + PhysicalSortExpr::new(col("b", &schema)?, options[1]), + ] + .into(), + vec![SplitPoint::new(vec![ + ScalarValue::Float64(Some(0.0)), + ScalarValue::Float64(Some(1.0)), + ])], + )?; + let range_expr: Arc = Arc::new(RangeExpr::try_new( + // Expression remapping may produce duplicate children. Preserve both + // so their sort options stay aligned with the split-point values. + vec![col("a", &schema)?, col("a", &schema)?], + &range_partitioning, + )?); + let filter_expr = binary(range_expr, Operator::Eq, lit(0u64), &schema)?; + let plan = Arc::new(FilterExec::try_new( + filter_expr, + Arc::new(EmptyExec::new(Arc::clone(&schema))), + )?); + + let ctx = SessionContext::new(); + let result = roundtrip_test_and_return( + plan, + &ctx, + &DefaultPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + let filter = result.downcast_ref::().unwrap(); + let binary = filter.predicate().downcast_ref::().unwrap(); + let range_expr = binary.left().downcast_ref::().unwrap(); + assert_eq!(range_expr.split_points(), range_partitioning.split_points()); + assert_eq!(range_expr.sort_options(), &options); + let children = range_expr.on_columns(); + assert_eq!(children.len(), 2); + for child in children { + let column = child.downcast_ref::().unwrap(); + assert_eq!((column.name(), column.index()), ("a", 0)); + } + + Ok(()) +} + +#[test] +fn roundtrip_call_null_scalar_struct_dict() -> Result<()> { + let data_type = DataType::Struct(Fields::from(vec![Field::new( + "item", + DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), + true, + )])); + + let schema = Arc::new(Schema::new(vec![Field::new("a", data_type.clone(), true)])); + let scan = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let scalar = lit(ScalarValue::try_from(data_type)?); + let filter = Arc::new(FilterExec::try_new( + Arc::new(BinaryExpr::new(scalar, Operator::Eq, col("a", &schema)?)), + scan, + )?); + + roundtrip_test(filter) +} + +/// Test that a chain of the same operator (a AND b AND c) is linearized +/// and roundtrips correctly. +#[test] +fn roundtrip_binary_expr_chain_same_op() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Boolean, false); + let field_c = Field::new("c", DataType::Boolean, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c])); + let ab = binary( + col("a", &schema)?, + Operator::And, + col("b", &schema)?, + &schema, + )?; + let abc = binary(ab, Operator::And, col("c", &schema)?, &schema)?; + roundtrip_test(Arc::new(FilterExec::try_new( + abc, + Arc::new(EmptyExec::new(schema)), + )?)) +} + +/// Test that mixed operators (a AND b OR c) are NOT linearized together — +/// only chains of the same operator are flattened. +#[test] +fn roundtrip_binary_expr_mixed_ops() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Boolean, false); + let field_c = Field::new("c", DataType::Boolean, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c])); + // (a AND b) OR c — AND and OR are different operators, so linearization stops + let a_and_b = binary( + col("a", &schema)?, + Operator::And, + col("b", &schema)?, + &schema, + )?; + let expr = binary(a_and_b, Operator::Or, col("c", &schema)?, &schema)?; + roundtrip_test(Arc::new(FilterExec::try_new( + expr, + Arc::new(EmptyExec::new(schema)), + )?)) +} + +/// Test that a deeply nested chain of AND expressions (like many WHERE conditions) +/// roundtrips correctly. This is the scenario from issue #18602. +#[test] +fn roundtrip_binary_expr_deeply_nested_and_chain() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let schema = Arc::new(Schema::new(vec![field_a])); + + // Build a chain: a AND a AND a AND ... (100 times) + let col_a = col("a", &schema)?; + let mut expr = Arc::clone(&col_a); + for _ in 0..99 { + expr = binary(expr, Operator::And, Arc::clone(&col_a), &schema)?; + } + + roundtrip_test(Arc::new(FilterExec::try_new( + expr, + Arc::new(EmptyExec::new(schema)), + )?)) +} + +/// Test that a deeply nested chain of OR expressions roundtrips correctly. +#[test] +fn roundtrip_binary_expr_deeply_nested_or_chain() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let schema = Arc::new(Schema::new(vec![field_a])); + + let col_a = col("a", &schema)?; + let mut expr = Arc::clone(&col_a); + for _ in 0..99 { + expr = binary(expr, Operator::Or, Arc::clone(&col_a), &schema)?; + } + + roundtrip_test(Arc::new(FilterExec::try_new( + expr, + Arc::new(EmptyExec::new(schema)), + )?)) +} + +/// Test that alternating AND/OR operators produce correct results — +/// each sub-chain gets linearized independently. +#[test] +fn roundtrip_binary_expr_alternating_and_or() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Boolean, false); + let field_c = Field::new("c", DataType::Boolean, false); + let field_d = Field::new("d", DataType::Boolean, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c, field_d])); + + // (a AND b) OR (c AND d) + let a_and_b = binary( + col("a", &schema)?, + Operator::And, + col("b", &schema)?, + &schema, + )?; + let c_and_d = binary( + col("c", &schema)?, + Operator::And, + col("d", &schema)?, + &schema, + )?; + let expr = binary(a_and_b, Operator::Or, c_and_d, &schema)?; + + roundtrip_test(Arc::new(FilterExec::try_new( + expr, + Arc::new(EmptyExec::new(schema)), + )?)) +} + +/// Verify that the linearized proto format has a flat operands list +/// rather than deeply nested l/r fields. +#[test] +fn test_linearization_produces_flat_operands() -> Result<()> { + // Build: a AND a AND a AND a (4 operands, 3 levels of nesting) + let col_a: Arc = Arc::new(Column::new("a", 0)); + let expr: Arc = Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::new(BinaryExpr::new( + Arc::clone(&col_a), + Operator::And, + Arc::clone(&col_a), + )), + Operator::And, + Arc::clone(&col_a), + )), + Operator::And, + Arc::clone(&col_a), + )); + + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let proto = proto_converter.physical_expr_to_proto(&expr, &codec)?; + + // The top-level should use the operands field with 4 entries + match &proto.expr_type { + Some(protobuf::physical_expr_node::ExprType::BinaryExpr(b)) => { + assert!( + b.l.is_none(), + "l should be None when using linearized operands" + ); + assert!( + b.r.is_none(), + "r should be None when using linearized operands" + ); + assert_eq!( + b.operands.len(), + 4, + "Expected 4 linearized operands for a AND a AND a AND a" + ); + assert_eq!(b.op, "And"); + } + other => panic!("Expected BinaryExpr, got {other:?}"), + } + + Ok(()) +} + +/// Test that linearization stops when encountering a different operator. +/// For (a AND b) OR c, only the top-level OR should be represented, and +/// the left-hand AND subtree should be a separate nested BinaryExpr. +#[test] +fn test_linearization_stops_at_different_op() -> Result<()> { + // (a AND b) OR c + let a_and_b: Arc = Arc::new(BinaryExpr::new( + Arc::new(Column::new("a", 0)), + Operator::And, + Arc::new(Column::new("b", 1)), + )); + let expr: Arc = Arc::new(BinaryExpr::new( + a_and_b, + Operator::Or, + Arc::new(Column::new("c", 2)), + )); + + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let proto = proto_converter.physical_expr_to_proto(&expr, &codec)?; + + // The top-level OR should have only 2 operands (can't linearize through AND) + match &proto.expr_type { + Some(protobuf::physical_expr_node::ExprType::BinaryExpr(b)) => { + assert_eq!( + b.operands.len(), + 2, + "Expected 2 operands for (a AND b) OR c" + ); + assert_eq!(b.op, "Or"); + // The first operand should be a nested AND BinaryExpr + match &b.operands[0].expr_type { + Some(protobuf::physical_expr_node::ExprType::BinaryExpr(inner)) => { + assert_eq!(inner.op, "And"); + assert_eq!(inner.operands.len(), 2); + } + other => panic!("Expected inner BinaryExpr(AND), got {other:?}"), + } + } + other => panic!("Expected BinaryExpr, got {other:?}"), + } + + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/filters.rs b/datafusion/proto/tests/cases/plans/filters.rs new file mode 100644 index 0000000000000..296923fcc7f93 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/filters.rs @@ -0,0 +1,103 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `FilterExec`. + +use super::{roundtrip_test, roundtrip_test_and_return}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::logical_expr::Operator; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{NotExpr, binary, col, in_list, lit}; +use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; +use datafusion::prelude::SessionContext; +use datafusion::scalar::ScalarValue; +use datafusion_common::Result; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, +}; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_filter_with_not_and_in_list() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let field_c = Field::new("c", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c])); + let not = Arc::new(NotExpr::new(col("a", &schema)?)); + let in_list = in_list( + col("b", &schema)?, + vec![ + lit(ScalarValue::Int64(Some(1))), + lit(ScalarValue::Int64(Some(2))), + ], + &false, + schema.as_ref(), + )?; + let and = binary(not, Operator::And, in_list, &schema)?; + roundtrip_test(Arc::new(FilterExec::try_new( + and, + Arc::new(EmptyExec::new(schema.clone())), + )?)) +} + +#[test] +fn roundtrip_filter_with_fetch() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let predicate = col("a", &schema)?; + let filter = FilterExecBuilder::new(predicate, Arc::new(EmptyExec::new(schema))) + .with_fetch(Some(10)) + .build()?; + assert_eq!(filter.fetch(), Some(10)); + roundtrip_test(Arc::new(filter)) +} + +#[test] +fn roundtrip_filter_projection_states() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Boolean, false), + Field::new("b", DataType::Int64, false), + ])); + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + for projection in [None, Some(vec![]), Some(vec![0])] { + let filter = FilterExecBuilder::new( + col("a", &schema)?, + Arc::new(EmptyExec::new(Arc::clone(&schema))), + ) + .apply_projection(projection.clone())? + .with_default_selectivity(37) + .with_batch_size(1024) + .with_fetch(Some(5)) + .build()?; + + let result = + roundtrip_test_and_return(Arc::new(filter), &ctx, &codec, &proto_converter)?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.projection().as_deref(), projection.as_deref()); + assert_eq!(result.default_selectivity(), 37); + assert_eq!(result.batch_size(), 1024); + assert_eq!(result.fetch(), Some(5)); + } + + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/joins.rs b/datafusion/proto/tests/cases/plans/joins.rs new file mode 100644 index 0000000000000..941e8832952c6 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/joins.rs @@ -0,0 +1,458 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The join execs. + +use super::{roundtrip_test, roundtrip_test_and_return}; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::logical_expr::{JoinType, Operator}; +use datafusion::physical_expr::LexOrdering; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{BinaryExpr, Column, PhysicalSortExpr}; +use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter}; +use datafusion::physical_plan::joins::{ + HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, + StreamJoinPartitionMode, SymmetricHashJoinExec, +}; +use datafusion::prelude::SessionContext; +use datafusion_common::{JoinSide, NullEquality, Result}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, +}; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_hash_join() -> Result<()> { + let field_a = Field::new("col", DataType::Int64, false); + let schema_left = Schema::new(vec![field_a.clone()]); + let schema_right = Schema::new(vec![field_a]); + let on = vec![( + Arc::new(Column::new("col", schema_left.index_of("col")?)) as _, + Arc::new(Column::new("col", schema_right.index_of("col")?)) as _, + )]; + + let schema_left = Arc::new(schema_left); + let schema_right = Arc::new(schema_right); + for join_type in &[ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftAnti, + JoinType::RightAnti, + JoinType::LeftSemi, + JoinType::RightSemi, + ] { + for partition_mode in &[PartitionMode::Partitioned, PartitionMode::CollectLeft] { + roundtrip_test(Arc::new(HashJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + on.clone(), + None, + join_type, + None, + *partition_mode, + NullEquality::NullEqualsNothing, + false, + )?))?; + } + } + Ok(()) +} + +#[test] +fn roundtrip_nested_loop_join() -> Result<()> { + let field_a = Field::new("col", DataType::Int64, false); + let schema_left = Schema::new(vec![field_a.clone()]); + let schema_right = Schema::new(vec![field_a]); + + let schema_left = Arc::new(schema_left); + let schema_right = Arc::new(schema_right); + for join_type in &[ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftAnti, + JoinType::RightAnti, + JoinType::LeftSemi, + JoinType::RightSemi, + ] { + roundtrip_test(Arc::new(NestedLoopJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + None, + join_type, + Some(vec![0]), + )?))?; + } + Ok(()) +} + +/// Regression: proto3 `repeated` fields cannot distinguish "absent" from "empty", +/// so a naive encoding collapses `Some(vec![])` and `None` into the same wire +/// representation. `try_embed_projection` (DataFusion 53+) produces +/// `HashJoinExec.projection = Some(vec![])` for `SELECT count(1) … JOIN …`, +/// which previously round-tripped to `None` and caused downstream consumers (e.g. +/// distributed Flight executors) to receive a different number of output +/// columns than the planner declared. Verify all three states preserve. +#[test] +fn roundtrip_hash_join_projection_states() -> Result<()> { + let field_a = Field::new("col", DataType::Int64, false); + let schema_left = Arc::new(Schema::new(vec![field_a.clone()])); + let schema_right = Arc::new(Schema::new(vec![field_a])); + let on = vec![( + Arc::new(Column::new("col", schema_left.index_of("col")?)) as _, + Arc::new(Column::new("col", schema_right.index_of("col")?)) as _, + )]; + + for projection in [None, Some(vec![]), Some(vec![0]), Some(vec![1])] { + roundtrip_test(Arc::new(HashJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + on.clone(), + None, + &JoinType::Inner, + projection, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + )?))?; + } + Ok(()) +} + +/// Regression: `HashJoinExecNode` had no `fetch` field, so the row limit that +/// the `limit_pushdown` physical optimizer rule pushes into the join via +/// `ExecutionPlan::with_fetch` was silently dropped by serde. Because that rule +/// also removes the enclosing `GlobalLimitExec` once the join absorbs the limit, +/// a round-tripped plan had no limit left at all and a distributed executor +/// returned more rows than the query asked for. +/// +/// Note this cannot be covered by `roundtrip_test`: that helper compares +/// `format!("{plan:?}")`, and `HashJoinExec`'s `Debug` output does not include +/// `fetch`, so the before/after strings match even when the value is lost. The +/// assertions below therefore inspect `fetch()` directly. +#[test] +fn roundtrip_hash_join_fetch() -> Result<()> { + let field_a = Field::new("col", DataType::Int64, false); + let schema_left = Arc::new(Schema::new(vec![field_a.clone()])); + let schema_right = Arc::new(Schema::new(vec![field_a])); + let on = vec![( + Arc::new(Column::new("col", schema_left.index_of("col")?)) as _, + Arc::new(Column::new("col", schema_right.index_of("col")?)) as _, + )]; + + // `usize::MAX` and `u32::MAX as usize` pin the decode-side `u64 -> usize` + // conversion: it is a checked `usize::try_from`, and a large fetch must + // survive the round trip exactly rather than being truncated or clamped. + // Both are representable on every target (on a 32-bit target `usize::MAX` + // is simply `u32::MAX`), so this stays portable. The truncating case + // itself -- a `u64` fetch above `usize::MAX` -- is only reachable on a + // 32-bit target and so is not exercised by this test on a 64-bit host. + for fetch in [None, Some(7), Some(u32::MAX as usize), Some(usize::MAX)] { + let join = HashJoinExec::try_new( + Arc::new(EmptyExec::new(Arc::clone(&schema_left))), + Arc::new(EmptyExec::new(Arc::clone(&schema_right))), + on.clone(), + None, + &JoinType::Inner, + None, + PartitionMode::Partitioned, + NullEquality::NullEqualsNothing, + false, + )?; + + let plan: Arc = match fetch { + // This is how `limit_pushdown` installs the limit. + Some(fetch) => join + .with_fetch(Some(fetch)) + .expect("HashJoinExec supports fetch"), + None => Arc::new(join), + }; + assert_eq!(plan.fetch(), fetch); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let deserialized = + roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; + + let deserialized_join = deserialized + .downcast_ref::() + .expect("should be a HashJoinExec"); + assert_eq!(deserialized_join.fetch(), fetch); + } + Ok(()) +} + +/// Same regression coverage for `NestedLoopJoinExec`, which shares the +/// `repeated uint32 projection` proto field shape with `HashJoinExec`. +#[test] +fn roundtrip_nested_loop_join_projection_states() -> Result<()> { + let field_a = Field::new("col", DataType::Int64, false); + let schema_left = Arc::new(Schema::new(vec![field_a.clone()])); + let schema_right = Arc::new(Schema::new(vec![field_a])); + + for projection in [None, Some(vec![]), Some(vec![0]), Some(vec![1])] { + roundtrip_test(Arc::new(NestedLoopJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + None, + &JoinType::Inner, + projection, + )?))?; + } + Ok(()) +} + +#[test] +fn roundtrip_sym_hash_join() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let field_a = Field::new("col_a", DataType::Int64, false); + let field_b = Field::new("col_b", DataType::Int64, false); + let schema_left = Schema::new(vec![field_a.clone()]); + let schema_right = Schema::new(vec![field_b.clone()]); + let on = vec![( + Arc::new(Column::new("col_a", schema_left.index_of("col_a")?)) as _, + Arc::new(Column::new("col_b", schema_right.index_of("col_b")?)) as _, + )]; + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("col_a", 0)), + Operator::Gt, + Arc::new(Column::new("col_b", 1)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![field_a, field_b])), + ); + + let schema_left = Arc::new(schema_left); + let schema_right = Arc::new(schema_right); + let left_order: LexOrdering = [PhysicalSortExpr { + expr: Arc::new(Column::new("col_a", schema_left.index_of("col_a")?)), + options: SortOptions { + descending: true, + nulls_first: false, + }, + }] + .into(); + let right_order: LexOrdering = [PhysicalSortExpr { + expr: Arc::new(Column::new("col_b", schema_right.index_of("col_b")?)), + options: SortOptions { + descending: false, + nulls_first: true, + }, + }] + .into(); + let ordering_cases = [ + (None, None), + (Some(left_order.clone()), None), + (None, Some(right_order.clone())), + (Some(left_order), Some(right_order)), + ]; + let ordering_options = |ordering: Option<&LexOrdering>| { + ordering + .map(|ordering| ordering.iter().map(|expr| expr.options).collect::>()) + }; + + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftAnti, + JoinType::RightAnti, + JoinType::LeftSemi, + JoinType::RightSemi, + JoinType::LeftMark, + JoinType::RightMark, + ] { + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + for filter in [None, Some(filter.clone())] { + for partition_mode in [ + StreamJoinPartitionMode::Partitioned, + StreamJoinPartitionMode::SinglePartition, + ] { + for (left_order, right_order) in &ordering_cases { + let result = roundtrip_test_and_return( + Arc::new(SymmetricHashJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + on.clone(), + filter.clone(), + &join_type, + null_equality, + left_order.clone(), + right_order.clone(), + partition_mode, + )?), + &ctx, + &codec, + &proto_converter, + )?; + let result = + result.downcast_ref::().unwrap(); + assert_eq!(result.join_type(), &join_type); + assert_eq!(result.null_equality(), null_equality); + assert_eq!(result.partition_mode(), partition_mode); + assert_eq!( + ordering_options(result.left_sort_exprs()), + ordering_options(left_order.as_ref()) + ); + assert_eq!( + ordering_options(result.right_sort_exprs()), + ordering_options(right_order.as_ref()) + ); + assert_eq!( + result.filter().map(JoinFilter::column_indices), + filter.as_ref().map(JoinFilter::column_indices) + ); + } + } + } + } + } + Ok(()) +} + +#[test] +fn roundtrip_sort_merge_join() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let field_a = Field::new("col_a", DataType::Int64, false); + let field_b = Field::new("col_b", DataType::Int64, false); + let schema_left = Schema::new(vec![field_a.clone()]); + let schema_right = Schema::new(vec![field_b.clone()]); + let on = vec![( + Arc::new(Column::new("col_a", schema_left.index_of("col_a")?)) as _, + Arc::new(Column::new("col_b", schema_right.index_of("col_b")?)) as _, + )]; + + let filter = JoinFilter::new( + Arc::new(BinaryExpr::new( + Arc::new(Column::new("col_a", 1)), + Operator::Gt, + Arc::new(Column::new("col_b", 0)), + )), + vec![ + ColumnIndex { + index: 0, + side: JoinSide::Left, + }, + ColumnIndex { + index: 0, + side: JoinSide::Right, + }, + ], + Arc::new(Schema::new(vec![field_a, field_b])), + ); + + let schema_left = Arc::new(schema_left); + let schema_right = Arc::new(schema_right); + let sort_options = vec![SortOptions { + descending: true, + nulls_first: false, + }]; + for null_equality in [ + NullEquality::NullEqualsNothing, + NullEquality::NullEqualsNull, + ] { + for filter in [None, Some(filter.clone())] { + for join_type in [ + JoinType::Inner, + JoinType::Left, + JoinType::Right, + JoinType::Full, + JoinType::LeftAnti, + JoinType::RightAnti, + JoinType::LeftSemi, + JoinType::RightSemi, + JoinType::LeftMark, + JoinType::RightMark, + ] { + let result = roundtrip_test_and_return( + Arc::new(SortMergeJoinExec::try_new( + Arc::new(EmptyExec::new(schema_left.clone())), + Arc::new(EmptyExec::new(schema_right.clone())), + on.clone(), + filter.clone(), + join_type, + sort_options.clone(), + null_equality, + )?), + &ctx, + &codec, + &proto_converter, + )?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.join_type(), join_type); + assert_eq!(result.null_equality(), null_equality); + assert_eq!(result.sort_options(), sort_options); + assert_eq!( + result.filter().as_ref().map(|f| f.column_indices()), + filter.as_ref().map(|f| f.column_indices()) + ); + } + } + } + Ok(()) +} + +#[tokio::test] +async fn roundtrip_logical_plan_sort_merge_join() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_csv( + "t0", + "tests/testdata/test.csv", + datafusion::prelude::CsvReadOptions::default().has_header(true), + ) + .await?; + ctx.register_csv( + "t1", + "tests/testdata/test.csv", + datafusion::prelude::CsvReadOptions::default().has_header(true), + ) + .await?; + + ctx.sql("SET datafusion.optimizer.prefer_hash_join = false") + .await? + .show() + .await?; + + let query = "SELECT t1.* FROM t0 join t1 on t0.a = t1.a"; + let plan = ctx.sql(query).await?.create_physical_plan().await?; + roundtrip_test(plan) +} diff --git a/datafusion/proto/tests/cases/plans/leaves.rs b/datafusion/proto/tests/cases/plans/leaves.rs new file mode 100644 index 0000000000000..afcab2dda24bc --- /dev/null +++ b/datafusion/proto/tests/cases/plans/leaves.rs @@ -0,0 +1,89 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Leaf plans: `EmptyExec` and `PlaceholderRowExec`. + +use super::{roundtrip_test, roundtrip_test_and_return}; +use datafusion::arrow::datatypes::Schema; +use datafusion::physical_plan::ExecutionPlanProperties; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion::prelude::SessionContext; +use datafusion_common::Result; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, +}; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::PhysicalPlanNode; +use std::sync::Arc; + +#[test] +fn roundtrip_empty() -> Result<()> { + roundtrip_test(Arc::new(EmptyExec::new(Arc::new(Schema::empty())))) +} + +#[test] +fn roundtrip_empty_with_partitions() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let plan = Arc::new(EmptyExec::new(Arc::new(Schema::empty())).with_partitions(4)); + let plan = roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; + assert_eq!(plan.output_partitioning().partition_count(), 4); + Ok(()) +} + +#[test] +fn roundtrip_placeholder_row_with_partitions() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let plan = + Arc::new(PlaceholderRowExec::new(Arc::new(Schema::empty())).with_partitions(4)); + let plan = roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; + assert_eq!(plan.output_partitioning().partition_count(), 4); + Ok(()) +} + +/// Plans encoded before `partitions` was added carry no value for it, which +/// decodes as zero and must be treated as the previous default of one. +#[test] +fn decode_empty_and_placeholder_row_without_partitions() -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let schema: protobuf::Schema = (&Schema::empty()).try_into()?; + + for physical_plan_type in [ + protobuf::physical_plan_node::PhysicalPlanType::Empty(protobuf::EmptyExecNode { + schema: Some(schema.clone()), + partitions: 0, + }), + protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow( + protobuf::PlaceholderRowExecNode { + schema: Some(schema.clone()), + partitions: 0, + }, + ), + ] { + let node = PhysicalPlanNode { + physical_plan_type: Some(physical_plan_type), + }; + let plan = node.try_into_physical_plan(ctx.task_ctx().as_ref(), &codec)?; + assert_eq!(plan.output_partitioning().partition_count(), 1); + } + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/limits.rs b/datafusion/proto/tests/cases/plans/limits.rs new file mode 100644 index 0000000000000..a832d46d53152 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/limits.rs @@ -0,0 +1,240 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Plans that carry a row limit or shape the batch pipeline: the limit +//! execs, the coalescing execs, `BufferExec` and `CooperativeExec`. + +use super::{roundtrip_test, roundtrip_test_and_return}; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::listing::PartitionedFile; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::datasource::physical_plan::{ + FileGroup, FileScanConfig, FileScanConfigBuilder, ParquetSource, +}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::physical_expr::LexOrdering; +use datafusion::physical_optimizer::PhysicalOptimizerRule; +use datafusion::physical_optimizer::limit_pushdown::LimitPushdown; +use datafusion::physical_plan::buffer::BufferExec; +#[expect(deprecated)] +use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; +use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; +use datafusion::physical_plan::coop::CooperativeExec; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{PhysicalSortExpr, col}; +use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; +use datafusion::physical_plan::{ChildrenPropertiesMode, ReplaceChildrenOptions}; +use datafusion::prelude::SessionContext; +use datafusion_common::Result; +use datafusion_common::config::ConfigOptions; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, +}; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_local_limit() -> Result<()> { + roundtrip_test(Arc::new(LocalLimitExec::new( + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), + 25, + ))) +} + +#[test] +fn roundtrip_global_limit() -> Result<()> { + roundtrip_test(Arc::new(GlobalLimitExec::new( + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), + 0, + Some(25), + ))) +} + +#[test] +fn roundtrip_global_skip_no_limit() -> Result<()> { + roundtrip_test(Arc::new(GlobalLimitExec::new( + Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), + 10, + None, // no limit + ))) +} + +/// Sort key at index 1, so a decoder that misbinds column name vs index +/// cannot pass. +fn limit_test_schema() -> Arc { + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Int64, false), + Field::new("b", DataType::Int64, false), + ])) +} + +/// Non-default sort options, so a decode that falls back to defaults cannot +/// pass. +fn limit_required_ordering(schema: &Schema) -> Result> { + Ok(LexOrdering::new(vec![PhysicalSortExpr { + expr: col("b", schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }])) +} + +#[test] +fn roundtrip_limit_with_required_ordering() -> Result<()> { + let schema = limit_test_schema(); + let required_ordering = limit_required_ordering(&schema)?; + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + let mut global = + GlobalLimitExec::new(Arc::new(EmptyExec::new(Arc::clone(&schema))), 3, Some(25)); + global.set_required_ordering(required_ordering.clone()); + let decoded = + roundtrip_test_and_return(Arc::new(global), &ctx, &codec, &proto_converter)?; + let decoded = decoded + .downcast_ref::() + .expect("expected GlobalLimitExec"); + assert_eq!(decoded.required_ordering(), &required_ordering); + + let mut local = LocalLimitExec::new(Arc::new(EmptyExec::new(schema)), 25); + local.set_required_ordering(required_ordering.clone()); + let decoded = + roundtrip_test_and_return(Arc::new(local), &ctx, &codec, &proto_converter)?; + let decoded = decoded + .downcast_ref::() + .expect("expected LocalLimitExec"); + assert_eq!(decoded.required_ordering(), &required_ordering); + Ok(()) +} + +/// A limit's `required_ordering` is the only record that an `ORDER BY ... LIMIT` +/// whose sort node was optimized away is order-sensitive, so it must survive +/// serde all the way into the scan's `preserve_order` flag. +#[test] +fn roundtrip_limit_required_ordering_reaches_data_source() -> Result<()> { + let file_schema = limit_test_schema(); + let make_scan = || { + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .build(); + DataSourceExec::from_data_source(scan_config) + }; + let scan_after_limit_pushdown = |limit: GlobalLimitExec| -> Result { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let decoded = + roundtrip_test_and_return(Arc::new(limit), &ctx, &codec, &proto_converter)?; + + // Child replacement must not erase the decoded ordering before pushdown. + let rebuilt = decoded.replace_children( + vec![make_scan()], + ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute), + )?; + + let optimized = + LimitPushdown::new().optimize(rebuilt, &ConfigOptions::default())?; + let scan = optimized + .downcast_ref::() + .expect("limit should be absorbed into the scan"); + Ok(scan + .data_source() + .downcast_ref::() + .expect("expected FileScanConfig") + .clone()) + }; + + let mut limit = GlobalLimitExec::new(make_scan(), 0, Some(10)); + limit.set_required_ordering(limit_required_ordering(&file_schema)?); + let scan_config = scan_after_limit_pushdown(limit)?; + assert_eq!(scan_config.limit, Some(10)); + assert!(scan_config.preserve_order); + + let scan_config = + scan_after_limit_pushdown(GlobalLimitExec::new(make_scan(), 0, Some(10)))?; + assert_eq!(scan_config.limit, Some(10)); + assert!(!scan_config.preserve_order); + Ok(()) +} + +#[test] +fn roundtrip_coalesce_batches_with_fetch() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + #[expect(deprecated)] + roundtrip_test(Arc::new(CoalesceBatchesExec::new( + Arc::new(EmptyExec::new(schema.clone())), + 8096, + )))?; + + #[expect(deprecated)] + roundtrip_test(Arc::new( + CoalesceBatchesExec::new(Arc::new(EmptyExec::new(schema)), 8096) + .with_fetch(Some(10)), + )) +} + +#[test] +fn roundtrip_coalesce_partitions_with_fetch() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + roundtrip_test(Arc::new(CoalescePartitionsExec::new(Arc::new( + EmptyExec::new(schema.clone()), + ))))?; + + roundtrip_test(Arc::new( + CoalescePartitionsExec::new(Arc::new(EmptyExec::new(schema))) + .with_fetch(Some(10)), + )) +} + +#[test] +fn roundtrip_cooperative() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Boolean, false)])); + roundtrip_test(Arc::new(CooperativeExec::new(Arc::new(EmptyExec::new( + schema, + ))))) +} + +#[test] +fn roundtrip_buffer() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Boolean, false)])); + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let result = roundtrip_test_and_return( + Arc::new(BufferExec::new(Arc::new(EmptyExec::new(schema)), 4096)), + &ctx, + &codec, + &proto_converter, + )?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.capacity(), 4096); + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/misc.rs b/datafusion/proto/tests/cases/plans/misc.rs new file mode 100644 index 0000000000000..41bb051c28730 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/misc.rs @@ -0,0 +1,520 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Plans that do not (yet) warrant a file of their own: unions, unnest, +//! repartitioning and the analyze/explain execs. + +use super::{roundtrip_test, roundtrip_test_and_return}; +use arrow::datatypes::{Fields, TimeUnit}; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::physical_expr::LexOrdering; +use datafusion::physical_plan::analyze::AnalyzeExec; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::explain::ExplainExec; +use datafusion::physical_plan::expressions::{PhysicalSortExpr, col, lit}; +use datafusion::physical_plan::metrics::MetricCategory; +use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion::physical_plan::repartition::RepartitionExec; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::union::{InterleaveExec, UnionExec}; +use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; +use datafusion::physical_plan::{ + ExecutionPlan, Partitioning, RangePartitioning, SplitPoint, +}; +use datafusion::prelude::SessionContext; +use datafusion::scalar::ScalarValue; +use datafusion_common::display::{PlanType, StringifiedPlan}; +use datafusion_common::format::ExplainFormat; +use datafusion_common::{DataFusionError, Result, UnnestOptions}; +use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + PhysicalPlanDecodeContext, +}; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::PhysicalPlanNode; +use prost::Message; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_analyze() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("plan_type", DataType::Utf8, false), + Field::new("plan", DataType::Utf8, false), + ])); + let input = Arc::new(PlaceholderRowExec::new(Arc::clone(&schema))); + let metric_categories = vec![MetricCategory::Rows, MetricCategory::Timing]; + let analyze = Arc::new( + AnalyzeExec::builder(true, true, input, Arc::clone(&schema)) + .with_metric_categories(Some(metric_categories.clone())) + .with_format(ExplainFormat::Tree) + .build(), + ); + + let ctx = SessionContext::new(); + let roundtripped = roundtrip_test_and_return( + analyze, + &ctx, + &DefaultPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + let roundtripped = roundtripped.downcast_ref::().unwrap(); + + assert_eq!(roundtripped.schema(), schema); + assert!(roundtripped.verbose()); + assert!(roundtripped.show_statistics()); + assert_eq!( + roundtripped.metric_categories(), + Some(metric_categories.as_slice()) + ); + assert_eq!(roundtripped.format(), &ExplainFormat::Tree); + assert!( + roundtripped + .input() + .downcast_ref::() + .is_some() + ); + Ok(()) +} + +#[test] +fn roundtrip_explain() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("plan_type", DataType::Utf8, false), + Field::new("plan", DataType::Utf8, false), + ])); + let stringified_plans = vec![ + StringifiedPlan::new(PlanType::InitialLogicalPlan, "initial logical"), + StringifiedPlan::new( + PlanType::AnalyzedLogicalPlan { + analyzer_name: "analyzer".to_string(), + }, + "analyzed logical", + ), + StringifiedPlan::new(PlanType::FinalAnalyzedLogicalPlan, "final analyzed"), + StringifiedPlan::new( + PlanType::OptimizedLogicalPlan { + optimizer_name: "logical optimizer".to_string(), + }, + "optimized logical", + ), + StringifiedPlan::new(PlanType::FinalLogicalPlan, "final logical"), + StringifiedPlan::new(PlanType::InitialPhysicalPlan, "initial physical"), + StringifiedPlan::new( + PlanType::InitialPhysicalPlanWithStats, + "initial physical with stats", + ), + StringifiedPlan::new( + PlanType::InitialPhysicalPlanWithSchema, + "initial physical with schema", + ), + StringifiedPlan::new( + PlanType::OptimizedPhysicalPlan { + optimizer_name: "physical optimizer".to_string(), + }, + "optimized physical", + ), + StringifiedPlan::new(PlanType::FinalPhysicalPlan, "final physical"), + StringifiedPlan::new( + PlanType::FinalPhysicalPlanWithStats, + "final physical with stats", + ), + StringifiedPlan::new( + PlanType::FinalPhysicalPlanWithSchema, + "final physical with schema", + ), + StringifiedPlan::new(PlanType::PhysicalPlanError, "physical plan error"), + ]; + let explain = Arc::new(ExplainExec::new( + Arc::clone(&schema), + stringified_plans.clone(), + true, + )); + + let ctx = SessionContext::new(); + let roundtripped = roundtrip_test_and_return( + explain, + &ctx, + &DefaultPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + let roundtripped = roundtripped.downcast_ref::().unwrap(); + + assert_eq!(roundtripped.schema(), schema); + assert_eq!(roundtripped.stringified_plans(), stringified_plans); + assert!(roundtripped.verbose()); + Ok(()) +} + +#[test] +fn roundtrip_union() -> Result<()> { + let field_a = Field::new("col", DataType::Int64, false); + let schema_left = Schema::new(vec![field_a.clone()]); + let schema_right = Schema::new(vec![field_a]); + let left = EmptyExec::new(Arc::new(schema_left)); + let right = EmptyExec::new(Arc::new(schema_right)); + let inputs: Vec> = vec![Arc::new(left), Arc::new(right)]; + let union = UnionExec::try_new(inputs)?; + roundtrip_test(union) +} + +/// `UnionExec::try_new` coerces a nullability-mismatched leg by wrapping it +/// in a `ProjectionExec` with a same-type `CastExpr` (see `coerce_schema` in +/// `datafusion-physical-plan`'s `union` module) -- a zero-copy relabeling, +/// not a real cast. `ProjectionExec` has an ordinary protobuf message, so +/// unlike the node this replaced, there's no wrapper-erasure trick to verify; +/// just that the decoded plan still contains the coercion and that its +/// emitted batches expose the union's nullable schema. +#[tokio::test] +async fn roundtrip_union_with_mismatched_nullability_executes() -> Result<()> { + let literal_leg = |value: ScalarValue| -> Result> { + Ok(Arc::new(ProjectionExec::try_new( + vec![ProjectionExpr { + expr: lit(value), + alias: "a".to_string(), + }], + Arc::new(PlaceholderRowExec::new(Arc::new(Schema::empty()))), + )?)) + }; + let non_nullable_leg = literal_leg(ScalarValue::Int64(Some(1)))?; + let nullable_leg = literal_leg(ScalarValue::Int64(None))?; + + let union: Arc = + UnionExec::try_new(vec![non_nullable_leg, nullable_leg])?; + assert!(union.schema().field(0).is_nullable()); + assert!( + format!("{union:?}").contains("CastExpr"), + "expected a coercing CastExpr in plan:\n{union:?}" + ); + + let ctx = SessionContext::new(); + let bytes = datafusion_proto::bytes::physical_plan_to_bytes(Arc::clone(&union))?; + let roundtripped = datafusion_proto::bytes::physical_plan_from_bytes( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + )?; + assert!(roundtripped.schema().field(0).is_nullable()); + assert!( + format!("{roundtripped:?}").contains("CastExpr"), + "expected a coercing CastExpr after roundtrip:\n{roundtripped:?}" + ); + + let batches = + datafusion::physical_plan::collect(roundtripped, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + for batch in &batches { + assert!(batch.schema().field(0).is_nullable()); + } + + Ok(()) +} + +#[test] +fn roundtrip_repartition_preserve_order() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a])); + let sort_exprs: LexOrdering = [PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions::default(), + }] + .into(); + + // Create two sorted single-partition inputs, then union them to get + // a sorted input with 2 partitions. + let source1 = SortExec::new( + sort_exprs.clone(), + Arc::new(EmptyExec::new(Arc::clone(&schema))), + ); + let source2 = SortExec::new(sort_exprs, Arc::new(EmptyExec::new(schema))); + let union = UnionExec::try_new(vec![ + Arc::new(source1) as Arc, + Arc::new(source2) as Arc, + ])?; + + let repartition = RepartitionExec::try_new(union, Partitioning::RoundRobinBatch(10))? + .with_preserve_order(); + assert!(repartition.preserve_order()); + + roundtrip_test(Arc::new(repartition)) +} + +#[test] +fn roundtrip_range_partitioning() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let input = Arc::new(EmptyExec::new(Arc::clone(&schema))); + let range_partitioning = Partitioning::Range(RangePartitioning::new( + [PhysicalSortExpr::new_default(col("a", &schema)?)].into(), + vec![SplitPoint::new(vec![ScalarValue::Int64(Some(10))])], + )); + // RepartitionExec is used only to carry the partitioning through proto. + // Executing range repartitioning is intentionally unsupported. + let repartition = RepartitionExec::try_new(input, range_partitioning)?; + + roundtrip_test(Arc::new(repartition)) +} + +/// `parse_protobuf_hash_partitioning` has no in-tree callers left; it delegates +/// to the shared `Partitioning::try_from_proto`, so pin that it still decodes +/// the hash message it is handed. +#[test] +fn parse_hash_partitioning_delegates_to_shared_decoder() -> Result<()> { + use datafusion_proto::physical_plan::from_proto::parse_protobuf_hash_partitioning; + + let schema = Schema::new(vec![Field::new("a", DataType::Int64, false)]); + let ctx = SessionContext::new(); + let task_ctx = ctx.task_ctx(); + let codec = DefaultPhysicalExtensionCodec {}; + let decode_ctx = PhysicalPlanDecodeContext::new(&task_ctx, &codec); + let proto_converter = DefaultPhysicalProtoConverter {}; + + let hash_expr = serialize_physical_expr_with_converter( + &col("a", &schema)?, + &codec, + &proto_converter, + )?; + let hash = protobuf::PhysicalHashRepartition { + hash_expr: vec![hash_expr], + partition_count: 4, + }; + + let partitioning = parse_protobuf_hash_partitioning( + Some(&hash), + &decode_ctx, + &schema, + &proto_converter, + )?; + let Some(Partitioning::Hash(exprs, count)) = partitioning else { + panic!("expected hash partitioning, got {partitioning:?}"); + }; + assert_eq!(count, 4); + assert_eq!(exprs.len(), 1); + assert_eq!(exprs[0].to_string(), col("a", &schema)?.to_string()); + + // No message means no partitioning, as before. + assert!( + parse_protobuf_hash_partitioning(None, &decode_ctx, &schema, &proto_converter)? + .is_none() + ); + + // The count is a `u64` on the wire and a `usize` in memory, so decoding + // narrows it. A count that does not fit is the case that motivated routing + // this through the shared decoder: it used to `unwrap()` and panic, and now + // reports an error. Only a target narrower than 64 bits can reach that arm + // -- on a 64-bit target every `u64` fits, and the assertion there is that + // the largest possible count survives whole rather than being truncated. + let oversized = protobuf::PhysicalHashRepartition { + hash_expr: vec![serialize_physical_expr_with_converter( + &col("a", &schema)?, + &codec, + &proto_converter, + )?], + partition_count: u64::MAX, + }; + let decoded = parse_protobuf_hash_partitioning( + Some(&oversized), + &decode_ctx, + &schema, + &proto_converter, + ); + + #[cfg(target_pointer_width = "64")] + { + let Some(Partitioning::Hash(_, count)) = decoded? else { + panic!("expected hash partitioning"); + }; + assert_eq!(count, usize::MAX); + } + + #[cfg(not(target_pointer_width = "64"))] + assert!( + decoded + .unwrap_err() + .to_string() + .contains("Partition count 18446744073709551615 exceeds usize::MAX") + ); + + Ok(()) +} + +#[test] +fn roundtrip_interleave() -> Result<()> { + let field_a = Field::new("col", DataType::Int64, false); + let schema_left = Schema::new(vec![field_a.clone()]); + let schema_right = Schema::new(vec![field_a]); + let partition = Partitioning::Hash(vec![], 3); + let left = RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::new(schema_left))), + partition.clone(), + )?; + let right = RepartitionExec::try_new( + Arc::new(EmptyExec::new(Arc::new(schema_right))), + partition, + )?; + let inputs: Vec> = vec![Arc::new(left), Arc::new(right)]; + let interleave = InterleaveExec::try_new(inputs)?; + roundtrip_test(Arc::new(interleave)) +} + +/// See [`roundtrip_union_with_mismatched_nullability_executes`]: the same +/// wrapper-reinsertion behavior applies to `InterleaveExec::try_from_proto`. +#[tokio::test] +async fn roundtrip_interleave_with_mismatched_nullability_executes() -> Result<()> { + let partition = Partitioning::Hash(vec![], 3); + let literal_leg = |value: ScalarValue| -> Result> { + let projection = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: lit(value), + alias: "a".to_string(), + }], + Arc::new(PlaceholderRowExec::new(Arc::new(Schema::empty()))), + )?; + Ok(Arc::new(RepartitionExec::try_new( + Arc::new(projection), + partition.clone(), + )?)) + }; + let non_nullable_leg = literal_leg(ScalarValue::Int64(Some(1)))?; + let nullable_leg = literal_leg(ScalarValue::Int64(None))?; + + let interleave: Arc = Arc::new(InterleaveExec::try_new(vec![ + non_nullable_leg, + nullable_leg, + ])?); + assert!(interleave.schema().field(0).is_nullable()); + assert!( + format!("{interleave:?}").contains("CastExpr"), + "expected a coercing CastExpr in plan:\n{interleave:?}" + ); + + let ctx = SessionContext::new(); + let bytes = datafusion_proto::bytes::physical_plan_to_bytes(Arc::clone(&interleave))?; + let roundtripped = datafusion_proto::bytes::physical_plan_from_bytes( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + )?; + assert!(roundtripped.schema().field(0).is_nullable()); + assert!( + format!("{roundtripped:?}").contains("CastExpr"), + "expected a coercing CastExpr after roundtrip:\n{roundtripped:?}" + ); + + let batches = + datafusion::physical_plan::collect(roundtripped, ctx.task_ctx()).await?; + assert_eq!(batches.iter().map(|b| b.num_rows()).sum::(), 2); + for batch in &batches { + assert!(batch.schema().field(0).is_nullable()); + } + + Ok(()) +} + +#[test] +fn roundtrip_unnest() -> Result<()> { + let fa = Field::new("a", DataType::Int64, true); + let fb0 = Field::new_list_field(DataType::Utf8, true); + let fb = Field::new_list("b", fb0.clone(), false); + let fc1 = Field::new("c1", DataType::Boolean, false); + let fc2 = Field::new("c2", DataType::Date64, true); + let fc = Field::new_struct("c", Fields::from(vec![fc1.clone(), fc2.clone()]), true); + let fd0 = Field::new_list_field(DataType::Float32, false); + let fd = Field::new_list("d", fd0.clone(), true); + let fe1 = Field::new("e1", DataType::UInt16, false); + let fe2 = Field::new("e2", DataType::Duration(TimeUnit::Millisecond), true); + let fe3 = Field::new("e3", DataType::Timestamp(TimeUnit::Millisecond, None), true); + let fe_fields = Fields::from(vec![fe1.clone(), fe2.clone(), fe3.clone()]); + let fe = Field::new_struct("e", fe_fields, false); + + let fb0 = fb0.with_name("b"); + let fd0 = fd0.with_name("d"); + let input_schema = Arc::new(Schema::new(vec![fa.clone(), fb, fc, fd, fe])); + let output_schema = + Arc::new(Schema::new(vec![fa, fb0, fc1, fc2, fd0, fe1, fe2, fe3])); + let input = Arc::new(EmptyExec::new(input_schema)); + let options = UnnestOptions { + null_handling: datafusion_common::NullHandling::Drop, + recursions: vec![datafusion_common::RecursionUnnestOption { + input_column: datafusion_common::Column::new_unqualified("b"), + output_column: datafusion_common::Column::new_unqualified("b"), + depth: 2, + }], + }; + let unnest = UnnestExec::new( + input, + vec![ + ListUnnest { + index_in_input_schema: 1, + depth: 1, + }, + ListUnnest { + index_in_input_schema: 1, + depth: 2, + }, + ListUnnest { + index_in_input_schema: 3, + depth: 2, + }, + ], + vec![2, 4], + output_schema, + options.clone(), + )?; + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let result = + roundtrip_test_and_return(Arc::new(unnest), &ctx, &codec, &proto_converter)?; + let result = result.downcast_ref::().unwrap(); + assert_eq!(result.options(), &options); + + Ok(()) +} + +#[tokio::test] +/// Tests that we can serialize an unoptimized "analyze" plan and it will work on the other end +async fn analyze_roundtrip_unoptimized() -> Result<()> { + let ctx = SessionContext::new(); + + // No optimizations + let session_state = + datafusion::execution::SessionStateBuilder::new_from_existing(ctx.state()) + .with_physical_optimizer_rules(vec![]) + .build(); + + let logical_plan = session_state + .create_logical_plan("explain analyze select 1") + .await?; + let plan = session_state.create_physical_plan(&logical_plan).await?; + + let node = PhysicalPlanNode::try_from_physical_plan( + plan.clone(), + &DefaultPhysicalExtensionCodec {}, + )?; + + let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + + let unoptimized = + node.try_into_physical_plan(&ctx.task_ctx(), &DefaultPhysicalExtensionCodec {})?; + + let physical_planner = + datafusion::physical_planner::DefaultPhysicalPlanner::default(); + physical_planner.optimize_physical_plan(unoptimized, &session_state, |_, _| {})?; + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/mod.rs b/datafusion/proto/tests/cases/plans/mod.rs new file mode 100644 index 0000000000000..8dad1eff67032 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/mod.rs @@ -0,0 +1,135 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Round trip tests for the physical plan protobuf representation. +//! +//! The tests are grouped by the kind of plan they cover; the shared +//! round trip helpers live here. + +use datafusion::physical_plan::ExecutionPlan; +use datafusion::prelude::{ParquetReadOptions, SessionContext}; +use datafusion_common::Result; +use datafusion_proto::bytes::{ + physical_plan_from_bytes_with_proto_converter, + physical_plan_to_bytes_with_proto_converter, +}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalProtoConverterExtension, +}; +use std::sync::Arc; + +mod aggregates; +mod dispatch; +mod dynamic_filters; +mod exprs; +mod filters; +mod joins; +mod leaves; +mod limits; +mod misc; +mod scalar_subquery; +mod sinks; +mod sorts; +mod sources; +mod tpch; +mod udfs; +mod windows; + +/// Perform a serde roundtrip and assert that the string representation of the before and after plans +/// are identical. Note that this often isn't sufficient to guarantee that no information is +/// lost during serde because the string representation of a plan often only shows a subset of state. +fn roundtrip_test(exec_plan: Arc) -> Result<()> { + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + Ok(()) +} + +/// Perform a serde roundtrip and assert that the string representation of the before and after plans +/// are identical. Note that this often isn't sufficient to guarantee that no information is +/// lost during serde because the string representation of a plan often only shows a subset of state. +/// +/// This version of the roundtrip_test method returns the final plan after serde so that it can be inspected +/// farther in tests. +fn roundtrip_test_and_return( + exec_plan: Arc, + ctx: &SessionContext, + codec: &dyn PhysicalExtensionCodec, + proto_converter: &dyn PhysicalProtoConverterExtension, +) -> Result> { + let bytes = physical_plan_to_bytes_with_proto_converter( + Arc::clone(&exec_plan), + codec, + proto_converter, + )?; + let result_exec_plan = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + codec, + proto_converter, + )?; + + pretty_assertions::assert_eq!( + format!("{exec_plan:?}"), + format!("{result_exec_plan:?}") + ); + Ok(result_exec_plan) +} + +/// Perform a serde roundtrip and assert that the string representation of the before and after plans +/// are identical. Note that this often isn't sufficient to guarantee that no information is +/// lost during serde because the string representation of a plan often only shows a subset of state. +/// +/// This version of the roundtrip_test function accepts a SessionContext, which is required when +/// performing serde on some plans. +fn roundtrip_test_with_context( + exec_plan: Arc, + ctx: &SessionContext, +) -> Result<()> { + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + roundtrip_test_and_return(exec_plan, ctx, &codec, &proto_converter)?; + Ok(()) +} + +/// Perform a serde roundtrip for the specified sql query, and assert that +/// query results are identical. +async fn roundtrip_test_sql_with_context(sql: &str, ctx: &SessionContext) -> Result<()> { + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let initial_plan = ctx.sql(sql).await?.create_physical_plan().await?; + + roundtrip_test_and_return(initial_plan, ctx, &codec, &proto_converter)?; + Ok(()) +} + +/// returns a SessionContext with `alltypes_plain` registered +async fn all_types_context() -> Result { + let ctx = SessionContext::new(); + + let testdata = datafusion::test_util::parquet_test_data(); + ctx.register_parquet( + "alltypes_plain", + &format!("{testdata}/alltypes_plain.parquet"), + ParquetReadOptions::default(), + ) + .await?; + + Ok(ctx) +} diff --git a/datafusion/proto/tests/cases/plans/scalar_subquery.rs b/datafusion/proto/tests/cases/plans/scalar_subquery.rs new file mode 100644 index 0000000000000..34d30aa03dece --- /dev/null +++ b/datafusion/proto/tests/cases/plans/scalar_subquery.rs @@ -0,0 +1,278 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `ScalarSubqueryExec` and the results it scopes to its subtree. + +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::logical_expr::Operator; +use datafusion::physical_plan::ExecutionPlan; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{BinaryExpr, binary, col}; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::scalar_subquery::{ + ScalarSubqueryExec, ScalarSubqueryLink, +}; +use datafusion::prelude::SessionContext; +use datafusion_common::Result; +use datafusion_expr::physical_planning_context::{ScalarSubqueryResults, SubqueryIndex}; +use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; +use datafusion_proto::bytes::{ + physical_plan_from_bytes_with_proto_converter, + physical_plan_to_bytes_with_proto_converter, +}; +use datafusion_proto::physical_plan::{ + DeduplicatingProtoConverter, DefaultPhysicalExtensionCodec, +}; +use std::sync::Arc; +use std::vec; + +/// Verify that ScalarSubqueryExpr nodes in the input plan are connected to the +/// same shared results container as ScalarSubqueryExec after a proto round-trip. +#[test] +fn roundtrip_scalar_subquery_exec() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let results = ScalarSubqueryResults::new(1); + + // Build the input plan: a filter whose predicate references the + // scalar subquery result via ScalarSubqueryExpr. + let sq_expr = Arc::new(ScalarSubqueryExpr::new( + DataType::Int64, + true, + SubqueryIndex::new(0), + results.clone(), + )); + let predicate = binary(col("a", &schema)?, Operator::Eq, sq_expr, &schema)?; + let filter = + FilterExec::try_new(predicate, Arc::new(EmptyExec::new(schema.clone())))?; + + // Build a trivial subquery plan. + let subquery_plan = + Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( + "x", + DataType::Int64, + true, + )])))); + + let exec: Arc = Arc::new(ScalarSubqueryExec::new( + Arc::new(filter), + vec![ScalarSubqueryLink { + plan: subquery_plan, + index: SubqueryIndex::new(0), + }], + results, + )); + + // Perform the round-trip using DeduplicatingProtoConverter, which + // creates a DeduplicatingDeserializer that threads scalar subquery + // results through expression deserialization. + let codec = DefaultPhysicalExtensionCodec {}; + let converter = DeduplicatingProtoConverter {}; + let bytes = physical_plan_to_bytes_with_proto_converter( + Arc::clone(&exec), + &codec, + &converter, + )?; + let ctx = SessionContext::new(); + let deserialized = physical_plan_from_bytes_with_proto_converter( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + &codec, + &converter, + )?; + + // Verify the deserialized ScalarSubqueryExec's results container is + // shared with the ScalarSubqueryExpr in the input plan. + let sq_exec = deserialized + .downcast_ref::() + .expect("expected ScalarSubqueryExec"); + let exec_results = sq_exec.results(); + + // Walk the input plan to find the ScalarSubqueryExpr and verify it + // points to the same results container. + let filter_exec = sq_exec + .input() + .downcast_ref::() + .expect("expected FilterExec"); + let binary_expr = filter_exec + .predicate() + .downcast_ref::() + .expect("expected BinaryExpr"); + let deserialized_sq_expr = binary_expr + .right() + .downcast_ref::() + .expect("expected ScalarSubqueryExpr"); + + assert!( + ScalarSubqueryResults::ptr_eq(exec_results, deserialized_sq_expr.results()), + "ScalarSubqueryExpr should share the same results container as ScalarSubqueryExec" + ); + Ok(()) +} + +/// Verify that nested ScalarSubqueryExec nodes deserialize with distinct +/// scoped results containers, and that each ScalarSubqueryExpr is wired to the +/// container for its own surrounding ScalarSubqueryExec. +#[test] +fn roundtrip_nested_scalar_subquery_exec_scopes_results() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); + let subquery_schema = + Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); + + let inner_results = ScalarSubqueryResults::new(1); + let inner_sq_expr = Arc::new(ScalarSubqueryExpr::new( + DataType::Int64, + true, + SubqueryIndex::new(0), + inner_results.clone(), + )); + let inner_predicate = + binary(col("a", &schema)?, Operator::Eq, inner_sq_expr, &schema)?; + let inner_filter = Arc::new(FilterExec::try_new( + inner_predicate, + Arc::new(EmptyExec::new(schema.clone())), + )?); + let inner_exec: Arc = Arc::new(ScalarSubqueryExec::new( + inner_filter, + vec![ScalarSubqueryLink { + plan: Arc::new(EmptyExec::new(subquery_schema.clone())), + index: SubqueryIndex::new(0), + }], + inner_results, + )); + + let outer_results = ScalarSubqueryResults::new(1); + let outer_sq_expr = Arc::new(ScalarSubqueryExpr::new( + DataType::Int64, + true, + SubqueryIndex::new(0), + outer_results.clone(), + )); + let outer_predicate = + binary(col("a", &schema)?, Operator::Eq, outer_sq_expr, &schema)?; + let outer_filter = Arc::new(FilterExec::try_new(outer_predicate, inner_exec)?); + let outer_exec: Arc = Arc::new(ScalarSubqueryExec::new( + outer_filter, + vec![ScalarSubqueryLink { + plan: Arc::new(EmptyExec::new(subquery_schema)), + index: SubqueryIndex::new(0), + }], + outer_results, + )); + + let bytes = datafusion_proto::bytes::physical_plan_to_bytes(Arc::clone(&outer_exec))?; + let ctx = SessionContext::new(); + let deserialized = datafusion_proto::bytes::physical_plan_from_bytes( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + )?; + + let outer_exec = deserialized + .downcast_ref::() + .expect("expected outer ScalarSubqueryExec"); + let outer_results = outer_exec.results(); + let outer_filter = outer_exec + .input() + .downcast_ref::() + .expect("expected outer FilterExec"); + let outer_binary = outer_filter + .predicate() + .downcast_ref::() + .expect("expected outer BinaryExpr"); + let outer_sq_expr = outer_binary + .right() + .downcast_ref::() + .expect("expected outer ScalarSubqueryExpr"); + + let inner_exec = outer_filter + .input() + .downcast_ref::() + .expect("expected inner ScalarSubqueryExec"); + let inner_results = inner_exec.results(); + let inner_filter = inner_exec + .input() + .downcast_ref::() + .expect("expected inner FilterExec"); + let inner_binary = inner_filter + .predicate() + .downcast_ref::() + .expect("expected inner BinaryExpr"); + let inner_sq_expr = inner_binary + .right() + .downcast_ref::() + .expect("expected inner ScalarSubqueryExpr"); + + assert!( + ScalarSubqueryResults::ptr_eq(outer_results, outer_sq_expr.results()), + "outer ScalarSubqueryExpr should use outer ScalarSubqueryExec results" + ); + assert!( + ScalarSubqueryResults::ptr_eq(inner_results, inner_sq_expr.results()), + "inner ScalarSubqueryExpr should use inner ScalarSubqueryExec results" + ); + assert!( + !ScalarSubqueryResults::ptr_eq(outer_results, inner_results), + "nested ScalarSubqueryExec nodes should not share results containers" + ); + assert!( + !ScalarSubqueryResults::ptr_eq(outer_results, inner_sq_expr.results()), + "inner ScalarSubqueryExpr must not read from outer results" + ); + assert!( + !ScalarSubqueryResults::ptr_eq(inner_results, outer_sq_expr.results()), + "outer ScalarSubqueryExpr must not read from inner results" + ); + + Ok(()) +} + +/// Verify that the default physical plan bytes round-trip preserves executable +/// scalar subquery plans. +#[tokio::test] +async fn roundtrip_scalar_subquery_exec_with_default_converter_executes() -> Result<()> { + let ctx = SessionContext::new(); + let sql = "SELECT x + (SELECT max(y) FROM (VALUES (10), (20)) AS u(y)) AS s \ + FROM (VALUES (2), (1)) AS t(x) \ + ORDER BY s"; + + let initial_plan = ctx.sql(sql).await?.create_physical_plan().await?; + assert!( + format!("{initial_plan:?}").contains("ScalarSubqueryExec"), + "expected ScalarSubqueryExec in plan:\n{initial_plan:?}" + ); + + let bytes = + datafusion_proto::bytes::physical_plan_to_bytes(Arc::clone(&initial_plan))?; + let roundtripped = datafusion_proto::bytes::physical_plan_from_bytes( + bytes.as_ref(), + ctx.task_ctx().as_ref(), + )?; + assert!( + format!("{roundtripped:?}").contains("ScalarSubqueryExec"), + "expected ScalarSubqueryExec after roundtrip:\n{roundtripped:?}" + ); + + let batches = datafusion::physical_plan::common::collect( + roundtripped.execute(0, ctx.task_ctx())?, + ) + .await?; + datafusion::assert_batches_eq!( + &["+----+", "| s |", "+----+", "| 21 |", "| 22 |", "+----+",], + &batches + ); + + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/sinks.rs b/datafusion/proto/tests/cases/plans/sinks.rs new file mode 100644 index 0000000000000..517659e9e9826 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/sinks.rs @@ -0,0 +1,329 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Data sinks and their file sink configurations. + +use super::{roundtrip_test, roundtrip_test_and_return}; +use arrow::csv::WriterBuilder; +use async_trait::async_trait; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use datafusion::datasource::file_format::csv::CsvSink; +use datafusion::datasource::file_format::json::JsonSink; +use datafusion::datasource::file_format::parquet::ParquetSink; +use datafusion::datasource::listing::{ListingTableUrl, PartitionedFile}; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::datasource::physical_plan::{FileGroup, FileOutputMode, FileSinkConfig}; +use datafusion::datasource::sink::{DataSink, DataSinkExec}; +use datafusion::execution::TaskContext; +use datafusion::physical_expr::PhysicalSortRequirement; +use datafusion::physical_plan::expressions::Column; +use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; +use datafusion::physical_plan::proto::ExecutionPlanEncodeCtx; +use datafusion::physical_plan::{ + DisplayAs, DisplayFormatType, ExecutionPlan, SendableRecordBatchStream, +}; +use datafusion::prelude::SessionContext; +use datafusion_common::Result; +use datafusion_common::config::TableParquetOptions; +use datafusion_common::file_options::csv_writer::CsvWriterOptions; +use datafusion_common::file_options::json_writer::JsonWriterOptions; +use datafusion_common::parsers::CompressionTypeVariant; +use datafusion_expr::dml::InsertOp; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, +}; +use datafusion_proto::protobuf; +use datafusion_proto::protobuf::PhysicalPlanNode; +use std::fmt::Formatter; +use std::sync::Arc; +use std::vec; + +#[derive(Debug)] +struct ProtoHookSink { + schema: SchemaRef, +} + +impl DisplayAs for ProtoHookSink { + fn fmt_as(&self, _t: DisplayFormatType, f: &mut Formatter) -> std::fmt::Result { + write!(f, "ProtoHookSink") + } +} + +#[async_trait] +impl DataSink for ProtoHookSink { + fn schema(&self) -> &SchemaRef { + &self.schema + } + + async fn write_all( + &self, + _data: SendableRecordBatchStream, + _context: &Arc, + ) -> Result { + unreachable!("serialization test does not execute the sink") + } + + fn try_to_proto( + &self, + exec: &DataSinkExec, + ctx: &ExecutionPlanEncodeCtx<'_>, + ) -> Result> { + let input = ctx.encode_child(exec.input())?; + let sort_order = exec.encode_sort_order(ctx)?; + assert!(matches!( + input.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::PlaceholderRow(_)) + )); + assert_eq!( + sort_order + .as_ref() + .map(|ordering| ordering.physical_sort_expr_nodes.len()), + Some(1) + ); + assert_eq!(exec.schema().fields().len(), 1); + + Ok(Some(PhysicalPlanNode { + physical_plan_type: Some( + protobuf::physical_plan_node::PhysicalPlanType::Empty( + protobuf::EmptyExecNode { + schema: Some(exec.schema().as_ref().try_into()?), + partitions: 1, + }, + ), + ), + })) + } +} + +#[test] +fn data_sink_exec_delegates_to_sink_proto_hook() -> Result<()> { + let input_schema = Arc::new(Schema::new(vec![Field::new( + "value", + DataType::Int64, + false, + )])); + let input = Arc::new(PlaceholderRowExec::new(Arc::clone(&input_schema))); + let sink = Arc::new(ProtoHookSink { + schema: Arc::clone(&input_schema), + }); + let sort_order = [PhysicalSortRequirement::new( + Arc::new(Column::new("value", 0)), + Some(SortOptions::default()), + )] + .into(); + let plan = Arc::new(DataSinkExec::new(input, sink, Some(sort_order))); + + let node = PhysicalPlanNode::try_from_physical_plan( + plan, + &DefaultPhysicalExtensionCodec {}, + )?; + + assert!(matches!( + node.physical_plan_type, + Some(protobuf::physical_plan_node::PhysicalPlanType::Empty(_)) + )); + Ok(()) +} + +#[test] +fn file_sink_config_roundtrip_preserves_fields() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "partition", + DataType::Utf8, + false, + )])); + let config = FileSinkConfig { + original_url: "file:///tmp/output".to_string(), + object_store_url: ObjectStoreUrl::local_filesystem(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp/output", 1)]), + table_paths: vec![ListingTableUrl::parse("file:///tmp/output")?], + output_schema: schema, + table_partition_cols: vec![("partition".to_string(), DataType::Utf8)], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: true, + file_extension: "parquet".to_string(), + file_output_mode: FileOutputMode::Directory, + }; + + let encoded = protobuf::FileSinkConfig::try_from(&config)?; + assert_eq!(encoded.insert_op(), protobuf::InsertOp::Overwrite); + assert_eq!( + encoded.file_output_mode(), + protobuf::FileOutputMode::Directory + ); + + let decoded = FileSinkConfig::try_from(&encoded)?; + assert_eq!(decoded.object_store_url, config.object_store_url); + assert_eq!(decoded.table_paths, config.table_paths); + assert_eq!( + decoded.output_schema.as_ref(), + config.output_schema.as_ref() + ); + assert_eq!(decoded.table_partition_cols, config.table_partition_cols); + assert_eq!(decoded.insert_op, config.insert_op); + assert_eq!( + decoded.keep_partition_by_columns, + config.keep_partition_by_columns + ); + assert_eq!(decoded.file_extension, config.file_extension); + assert_eq!(decoded.file_output_mode, config.file_output_mode); + + let [decoded_file] = decoded.file_group.files() else { + panic!("expected one decoded output file"); + }; + let [config_file] = config.file_group.files() else { + panic!("expected one configured output file"); + }; + assert_eq!( + decoded_file.object_meta.location, + config_file.object_meta.location + ); + assert_eq!(decoded_file.object_meta.size, config_file.object_meta.size); + Ok(()) +} + +#[test] +fn roundtrip_json_sink() -> Result<()> { + let field_a = Field::new("plan_type", DataType::Utf8, false); + let field_b = Field::new("plan", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let input = Arc::new(PlaceholderRowExec::new(schema.clone())); + + let file_sink_config = FileSinkConfig { + original_url: String::default(), + object_store_url: ObjectStoreUrl::local_filesystem(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), + table_paths: vec![ListingTableUrl::parse("file:///")?], + output_schema: schema.clone(), + table_partition_cols: vec![("plan_type".to_string(), DataType::Utf8)], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: true, + file_extension: "json".into(), + file_output_mode: FileOutputMode::SingleFile, + }; + let data_sink = Arc::new(JsonSink::new( + file_sink_config, + JsonWriterOptions::new(CompressionTypeVariant::UNCOMPRESSED), + )); + let sort_order = [PhysicalSortRequirement::new( + Arc::new(Column::new("plan_type", 0)), + Some(SortOptions { + descending: true, + nulls_first: false, + }), + )] + .into(); + + roundtrip_test(Arc::new(DataSinkExec::new( + input, + data_sink, + Some(sort_order), + ))) +} + +#[test] +fn roundtrip_csv_sink() -> Result<()> { + let field_a = Field::new("plan_type", DataType::Utf8, false); + let field_b = Field::new("plan", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let input = Arc::new(PlaceholderRowExec::new(schema.clone())); + + let file_sink_config = FileSinkConfig { + original_url: String::default(), + object_store_url: ObjectStoreUrl::local_filesystem(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), + table_paths: vec![ListingTableUrl::parse("file:///")?], + output_schema: schema.clone(), + table_partition_cols: vec![("plan_type".to_string(), DataType::Utf8)], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: true, + file_extension: "csv".into(), + file_output_mode: FileOutputMode::Directory, + }; + let data_sink = Arc::new(CsvSink::new( + file_sink_config, + CsvWriterOptions::new(WriterBuilder::default(), CompressionTypeVariant::ZSTD), + )); + let sort_order = [PhysicalSortRequirement::new( + Arc::new(Column::new("plan_type", 0)), + Some(SortOptions { + descending: true, + nulls_first: false, + }), + )] + .into(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + let roundtrip_plan = roundtrip_test_and_return( + Arc::new(DataSinkExec::new(input, data_sink, Some(sort_order))), + &ctx, + &codec, + &proto_converter, + )?; + + let roundtrip_plan = roundtrip_plan.downcast_ref::().unwrap(); + let csv_sink = roundtrip_plan.sink().downcast_ref::().unwrap(); + assert_eq!( + CompressionTypeVariant::ZSTD, + csv_sink.writer_options().compression + ); + + Ok(()) +} + +#[test] +fn roundtrip_parquet_sink() -> Result<()> { + let field_a = Field::new("plan_type", DataType::Utf8, false); + let field_b = Field::new("plan", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let input = Arc::new(PlaceholderRowExec::new(schema.clone())); + + let file_sink_config = FileSinkConfig { + original_url: String::default(), + object_store_url: ObjectStoreUrl::local_filesystem(), + file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), + table_paths: vec![ListingTableUrl::parse("file:///")?], + output_schema: schema.clone(), + table_partition_cols: vec![("plan_type".to_string(), DataType::Utf8)], + insert_op: InsertOp::Overwrite, + keep_partition_by_columns: true, + file_extension: "parquet".into(), + file_output_mode: FileOutputMode::Automatic, + }; + let data_sink = Arc::new(ParquetSink::new( + file_sink_config, + TableParquetOptions::default(), + )); + let sort_order = [PhysicalSortRequirement::new( + Arc::new(Column::new("plan_type", 0)), + Some(SortOptions { + descending: true, + nulls_first: false, + }), + )] + .into(); + + roundtrip_test(Arc::new(DataSinkExec::new( + input, + data_sink, + Some(sort_order), + ))) +} diff --git a/datafusion/proto/tests/cases/plans/sorts.rs b/datafusion/proto/tests/cases/plans/sorts.rs new file mode 100644 index 0000000000000..1172775b1bad7 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/sorts.rs @@ -0,0 +1,251 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! `SortExec` and `SortPreservingMergeExec`. + +use super::{roundtrip_test, roundtrip_test_and_return}; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::physical_expr::LexOrdering; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{PhysicalSortExpr, col}; +use datafusion::physical_plan::sorts::sort::SortExec; +use datafusion::physical_plan::sorts::sort_preserving_merge::SortPreservingMergeExec; +use datafusion::prelude::SessionContext; +use datafusion_common::Result; +use datafusion_proto::physical_plan::{ + DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, +}; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_sort() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let sort_exprs = [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + ] + .into(); + roundtrip_test(Arc::new(SortExec::new( + sort_exprs, + Arc::new(EmptyExec::new(schema)), + ))) +} + +#[test] +fn roundtrip_sort_preserve_partitioning() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let sort_exprs: LexOrdering = [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + ] + .into(); + + roundtrip_test(Arc::new(SortExec::new( + sort_exprs.clone(), + Arc::new(EmptyExec::new(schema.clone())), + )))?; + + roundtrip_test(Arc::new( + SortExec::new(sort_exprs, Arc::new(EmptyExec::new(schema))) + .with_preserve_partitioning(true), + )) +} + +/// `SortExec::fetch` turns a sort into a top-k sort. Losing it during serde +/// would silently widen the result set, so exercise the `Some(..)` state +/// explicitly (`roundtrip_sort` only covers `None`). +/// +/// `SortExec` currently derives `Debug`, so `roundtrip_test`'s +/// `format!("{plan:?}")` comparison does observe `fetch`. The assertions below +/// go through the accessor instead so that this coverage does not silently +/// disappear if `SortExec` ever grows a hand-written `Debug` impl. +#[test] +fn roundtrip_sort_with_fetch() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let sort_exprs: LexOrdering = [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + ] + .into(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + let roundtripped = roundtrip_test_and_return( + Arc::new( + SortExec::new( + sort_exprs.clone(), + Arc::new(EmptyExec::new(Arc::clone(&schema))), + ) + .with_fetch(Some(7)), + ), + &ctx, + &codec, + &proto_converter, + )?; + let roundtripped = roundtripped + .downcast_ref::() + .expect("should decode back into a SortExec"); + assert_eq!(roundtripped.fetch(), Some(7)); + assert_eq!(roundtripped.expr(), &sort_exprs); + + // `fetch` combined with `preserve_partitioning`, since both share the same + // proto node. + let roundtripped = roundtrip_test_and_return( + Arc::new( + SortExec::new(sort_exprs.clone(), Arc::new(EmptyExec::new(schema))) + .with_fetch(Some(3)) + .with_preserve_partitioning(true), + ), + &ctx, + &codec, + &proto_converter, + )?; + let roundtripped = roundtripped + .downcast_ref::() + .expect("should decode back into a SortExec"); + assert_eq!(roundtripped.fetch(), Some(3)); + assert!(roundtripped.preserve_partitioning()); + Ok(()) +} + +/// Round trip a [`SortPreservingMergeExec`], which had no dedicated round trip +/// test at all. +/// +/// Covers everything that is actually on the wire for this plan: the input, the +/// sort expressions and `fetch` in both its `None` and `Some(..)` states. +/// +/// Note that `SortPreservingMergeExec::enable_round_robin_repartition` is +/// deliberately *not* asserted on here: it has no field in +/// `SortPreservingMergeExecNode`, so it is not serialized and decoding always +/// restores the `true` default from `SortPreservingMergeExec::new`. Asserting +/// round trip equality on it would give a false sense of coverage. +/// +/// `SortPreservingMergeExec` derives `Debug`, so `roundtrip_test`'s +/// `format!("{plan:?}")` comparison does observe `expr` and `fetch`. The +/// assertions below use the accessors so the coverage survives a future +/// hand-written `Debug` impl. +#[test] +fn roundtrip_sort_preserving_merge() -> Result<()> { + let field_a = Field::new("a", DataType::Boolean, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + let sort_exprs: LexOrdering = [ + PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: true, + nulls_first: false, + }, + }, + PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions { + descending: false, + nulls_first: true, + }, + }, + ] + .into(); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + + // No fetch: `fetch` is encoded as -1 and must decode back to `None`. + let roundtripped = roundtrip_test_and_return( + Arc::new(SortPreservingMergeExec::new( + sort_exprs.clone(), + Arc::new(EmptyExec::new(Arc::clone(&schema))), + )), + &ctx, + &codec, + &proto_converter, + )?; + let roundtripped = roundtripped + .downcast_ref::() + .expect("should decode back into a SortPreservingMergeExec"); + assert_eq!(roundtripped.fetch(), None); + assert_eq!(roundtripped.expr(), &sort_exprs); + assert_eq!(roundtripped.input().schema(), schema); + + // With a fetch: dropping it would turn a bounded merge into an unbounded + // one and change the query result. + let roundtripped = roundtrip_test_and_return( + Arc::new( + SortPreservingMergeExec::new( + sort_exprs.clone(), + Arc::new(EmptyExec::new(Arc::clone(&schema))), + ) + .with_fetch(Some(11)), + ), + &ctx, + &codec, + &proto_converter, + )?; + let roundtripped = roundtripped + .downcast_ref::() + .expect("should decode back into a SortPreservingMergeExec"); + assert_eq!(roundtripped.fetch(), Some(11)); + assert_eq!(roundtripped.expr(), &sort_exprs); + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/sources.rs b/datafusion/proto/tests/cases/plans/sources.rs new file mode 100644 index 0000000000000..04708dec6439c --- /dev/null +++ b/datafusion/proto/tests/cases/plans/sources.rs @@ -0,0 +1,790 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Scans and data sources: file formats, `FileScanConfig`, listing tables +//! and memory sources. + +use super::{ + all_types_context, roundtrip_test, roundtrip_test_and_return, + roundtrip_test_sql_with_context, +}; +use arrow::array::RecordBatch; +use arrow::datatypes::Fields; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::datasource::empty::EmptyTable; +use datafusion::datasource::file_format::json::JsonFormat; +use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, PartitionedFile, +}; +use datafusion::datasource::object_store::ObjectStoreUrl; +use datafusion::datasource::physical_plan::{ + ArrowSource, CsvSource, FileGroup, FileScanConfig, FileScanConfigBuilder, JsonSource, + ParquetSource, wrap_partition_type_in_dict, wrap_partition_value_in_dict, +}; +use datafusion::datasource::source::DataSourceExec; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::Operator; +use datafusion::physical_expr::LexOrdering; +use datafusion::physical_plan::expressions::{ + BinaryExpr, Column, PhysicalSortExpr, col, lit, +}; +use datafusion::physical_plan::filter::FilterExecBuilder; +use datafusion::physical_plan::{ + ExecutionPlan, Partitioning, PhysicalExpr, RangePartitioning, SplitPoint, Statistics, + displayable, +}; +use datafusion::prelude::SessionContext; +use datafusion::scalar::ScalarValue; +use datafusion_common::config::TableParquetOptions; +use datafusion_common::stats::Precision; +use datafusion_common::{DataFusionError, Result, internal_datafusion_err, internal_err}; +use datafusion_datasource::{TableSchema, TableSchemaBuilder}; +use datafusion_expr::ColumnarValue; +use datafusion_physical_expr_common::physical_expr::proto_decode::PhysicalExprDecodeCtx; +use datafusion_physical_expr_common::physical_expr::proto_encode::PhysicalExprEncodeCtx; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, + PhysicalExtensionCodec, PhysicalProtoConverterExtension, +}; +use datafusion_proto::protobuf::PhysicalPlanNode; +use prost::Message; +use std::collections::HashMap; +use std::fmt::{Display, Formatter}; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + + let predicate = Arc::new(BinaryExpr::new( + Arc::new(Column::new("col", 1)), + Operator::Eq, + lit("1"), + )); + + let mut options = TableParquetOptions::new(); + options.global.pushdown_filters = true; + + let file_source = Arc::new( + ParquetSource::new(Arc::clone(&file_schema)) + .with_table_parquet_options(options) + .with_predicate(predicate), + ); + + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_statistics(Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1024), + column_statistics: Statistics::unknown_column(&Arc::new(Schema::new( + vec![Field::new("col", DataType::Utf8, false)], + ))), + }) + .build(); + + roundtrip_test(DataSourceExec::from_data_source(scan_config)) +} + +#[test] +fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_statistics(Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1024), + column_statistics: Statistics::unknown_column(&file_schema), + }) + .build(); + let exec_plan = DataSourceExec::from_data_source(scan_config); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let roundtripped = + roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + + let data_source = roundtripped + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("Expected DataSourceExec after roundtrip") + })?; + let file_scan = data_source + .data_source() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("Expected FileScanConfig after roundtrip") + })?; + let parquet_source = file_scan + .file_source() + .downcast_ref::() + .ok_or_else(|| { + internal_datafusion_err!("Expected ParquetSource after roundtrip") + })?; + + assert!( + parquet_source.parquet_file_reader_factory().is_some(), + "Parquet reader factory should be attached after decoding from protobuf" + ); + Ok(()) +} + +#[test] +fn roundtrip_arrow_scan() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + + let table_schema = TableSchema::from(&file_schema); + let file_source = Arc::new(ArrowSource::new_file_source(table_schema)); + + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.arrow".to_string(), + 1024, + )])]) + .with_statistics(Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1024), + column_statistics: Statistics::unknown_column(&file_schema), + }) + .build(); + + roundtrip_test(DataSourceExec::from_data_source(scan_config)) +} + +#[test] +fn roundtrip_json_scan() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(JsonSource::new(TableSchema::from(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.json".to_string(), + 1024, + )])]) + .build(); + roundtrip_test(DataSourceExec::from_data_source(scan_config)) +} + +#[cfg(feature = "avro")] +#[test] +fn roundtrip_avro_scan() -> Result<()> { + use datafusion_datasource_avro::source::AvroSource; + + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(AvroSource::new(TableSchema::from(&file_schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.avro".to_string(), + 1024, + )])]) + .build(); + roundtrip_test(DataSourceExec::from_data_source(scan_config)) +} + +#[test] +fn roundtrip_csv_scan_preserves_format_options() -> Result<()> { + use datafusion::common::config::CsvOptions; + + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let table_schema = TableSchema::from(&file_schema); + let file_source = + Arc::new(CsvSource::new(table_schema).with_csv_options(CsvOptions { + has_header: Some(false), + delimiter: b'|', + quote: b'\'', + escape: Some(b'\\'), + comment: Some(b'#'), + newlines_in_values: Some(true), + truncated_rows: Some(true), + ..Default::default() + })); + + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.csv".to_string(), + 1024, + )])]) + .build(); + + let ctx = SessionContext::new(); + let roundtripped = roundtrip_test_and_return( + DataSourceExec::from_data_source(scan_config), + &ctx, + &DefaultPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + let data_source = roundtripped + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected DataSourceExec"))?; + let file_scan = data_source + .data_source() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected FileScanConfig"))?; + let csv_source = file_scan + .file_source() + .downcast_ref::() + .ok_or_else(|| internal_datafusion_err!("Expected CsvSource"))?; + + assert!(!csv_source.has_header()); + assert_eq!(csv_source.delimiter(), b'|'); + assert_eq!(csv_source.quote(), b'\''); + assert_eq!(csv_source.escape(), Some(b'\\')); + assert_eq!(csv_source.comment(), Some(b'#')); + assert!(csv_source.newlines_in_values()); + assert!(csv_source.truncate_rows()); + Ok(()) +} + +#[tokio::test] +async fn roundtrip_parquet_exec_with_table_partition_cols() -> Result<()> { + let mut file_group = + PartitionedFile::new("/path/to/part=0/file.parquet".to_string(), 1024); + file_group.partition_values = + vec![wrap_partition_value_in_dict(ScalarValue::Int64(Some(0)))]; + let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + + let table_schema = TableSchemaBuilder::from(&schema) + .with_table_partition_cols(vec![Arc::new(Field::new( + "part".to_string(), + wrap_partition_type_in_dict(DataType::Int16), + false, + ))]) + .build(); + + let file_source = Arc::new(ParquetSource::new(table_schema.clone())); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_projection_indices(Some(vec![0, 1]))? + .with_file_group(FileGroup::new(vec![file_group])) + .build(); + + roundtrip_test(DataSourceExec::from_data_source(scan_config)) +} + +#[test] +fn roundtrip_parquet_exec_with_custom_predicate_expr() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + + let custom_predicate_expr = Arc::new(CustomPredicateExpr { + inner: Arc::new(Column::new("col", 1)), + }); + + let file_source = Arc::new( + ParquetSource::new(Arc::clone(&file_schema)) + .with_predicate(custom_predicate_expr), + ); + + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_statistics(Statistics { + num_rows: Precision::Inexact(100), + total_byte_size: Precision::Inexact(1024), + column_statistics: Statistics::unknown_column(&Arc::new(Schema::new( + vec![Field::new("col", DataType::Utf8, false)], + ))), + }) + .build(); + + #[derive(Debug, Clone, Eq)] + struct CustomPredicateExpr { + inner: Arc, + } + + // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 + impl PartialEq for CustomPredicateExpr { + fn eq(&self, other: &Self) -> bool { + self.inner.eq(&other.inner) + } + } + + impl std::hash::Hash for CustomPredicateExpr { + fn hash(&self, state: &mut H) { + self.inner.hash(state); + } + } + + impl Display for CustomPredicateExpr { + fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + write!(f, "CustomPredicateExpr") + } + } + + impl PhysicalExpr for CustomPredicateExpr { + fn data_type(&self, _input_schema: &Schema) -> Result { + unreachable!() + } + + fn nullable(&self, _input_schema: &Schema) -> Result { + unreachable!() + } + + fn evaluate(&self, _batch: &RecordBatch) -> Result { + unreachable!() + } + + fn children(&self) -> Vec<&Arc> { + vec![&self.inner] + } + + fn with_new_children( + self: Arc, + _children: Vec>, + ) -> Result> { + Ok(self) + } + + fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { + Display::fmt(self, f) + } + } + + #[derive(Debug)] + struct CustomPhysicalExtensionCodec; + impl PhysicalExtensionCodec for CustomPhysicalExtensionCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + unreachable!() + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + unreachable!() + } + + fn try_decode_expr( + &self, + buf: &[u8], + inputs: &[Arc], + _ctx: &PhysicalExprDecodeCtx<'_>, + ) -> Result> { + if buf == "CustomPredicateExpr".as_bytes() { + Ok(Arc::new(CustomPredicateExpr { + inner: inputs[0].clone(), + })) + } else { + internal_err!("Not supported") + } + } + + fn try_encode_expr( + &self, + node: &Arc, + buf: &mut Vec, + _ctx: &PhysicalExprEncodeCtx<'_>, + ) -> Result<()> { + if node.downcast_ref::().is_some() { + buf.extend_from_slice("CustomPredicateExpr".as_bytes()); + Ok(()) + } else { + internal_err!("Not supported") + } + } + } + + let exec_plan = DataSourceExec::from_data_source(scan_config); + + let ctx = SessionContext::new(); + roundtrip_test_and_return( + exec_plan, + &ctx, + &CustomPhysicalExtensionCodec {}, + &DefaultPhysicalProtoConverter {}, + )?; + Ok(()) +} + +#[tokio::test] +async fn roundtrip_json_source() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_json("t1", "../core/tests/data/1.json", Default::default()) + .await?; + let plan = ctx.table("t1").await?.create_physical_plan().await?; + roundtrip_test(plan) +} + +#[tokio::test] +async fn roundtrip_coalesce() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_table( + "t", + Arc::new(EmptyTable::new(Arc::new(Schema::new(Fields::from([ + Arc::new(Field::new("f", DataType::Int64, false)), + ]))))), + )?; + let df = ctx.sql("select coalesce(f) as f from t").await?; + let plan = df.create_physical_plan().await?; + + let node = PhysicalPlanNode::try_from_physical_plan( + plan.clone(), + &DefaultPhysicalExtensionCodec {}, + )?; + let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let restored = + node.try_into_physical_plan(&ctx.task_ctx(), &DefaultPhysicalExtensionCodec {})?; + + assert_eq!( + plan.schema(), + restored.schema(), + "Schema mismatch for plans:\n>> initial:\n{}>> final: \n{}", + displayable(plan.as_ref()) + .set_show_schema(true) + .indent(true), + displayable(restored.as_ref()) + .set_show_schema(true) + .indent(true), + ); + + Ok(()) +} + +#[tokio::test] +async fn roundtrip_generate_series() -> Result<()> { + let ctx = SessionContext::new(); + ctx.register_table( + "t", + Arc::new(EmptyTable::new(Arc::new(Schema::new(Fields::from([ + Arc::new(Field::new("f", DataType::Int64, false)), + ]))))), + )?; + let df = ctx.sql("select * from generate_series(1, 10000)").await?; + let plan = df.create_physical_plan().await?; + + let node = PhysicalPlanNode::try_from_physical_plan( + plan.clone(), + &DefaultPhysicalExtensionCodec {}, + )?; + let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) + .map_err(|e| DataFusionError::External(Box::new(e)))?; + let restored = + node.try_into_physical_plan(&ctx.task_ctx(), &DefaultPhysicalExtensionCodec {})?; + + assert_eq!( + plan.schema(), + restored.schema(), + "Schema mismatch for plans:\n>> initial:\n{}>> final: \n{}", + displayable(plan.as_ref()) + .set_show_schema(true) + .indent(true), + displayable(restored.as_ref()) + .set_show_schema(true) + .indent(true), + ); + + Ok(()) +} + +#[tokio::test] +async fn roundtrip_projection_source() -> Result<()> { + let schema = Arc::new(Schema::new(Fields::from([ + Arc::new(Field::new("a", DataType::Utf8, false)), + Arc::new(Field::new("b", DataType::Utf8, false)), + Arc::new(Field::new("c", DataType::Int32, false)), + Arc::new(Field::new("d", DataType::Int32, false)), + ]))); + + let statistics = Statistics::new_unknown(&schema); + + let file_source = Arc::new(ParquetSource::new(Arc::clone(&schema))); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_statistics(statistics) + .with_projection_indices(Some(vec![0, 1, 2]))? + .build(); + + let filter = Arc::new( + FilterExecBuilder::new( + Arc::new(BinaryExpr::new(col("c", &schema)?, Operator::Eq, lit(1))), + DataSourceExec::from_data_source(scan_config), + ) + .apply_projection(Some(vec![0, 1]))? + .build()?, + ); + + roundtrip_test(filter) +} + +#[tokio::test] +async fn roundtrip_parquet_select_star() -> Result<()> { + let ctx = all_types_context().await?; + let sql = "select * from alltypes_plain"; + roundtrip_test_sql_with_context(sql, &ctx).await +} + +#[tokio::test] +async fn roundtrip_parquet_select_projection() -> Result<()> { + let ctx = all_types_context().await?; + let sql = "select string_col, timestamp_col from alltypes_plain"; + roundtrip_test_sql_with_context(sql, &ctx).await +} + +#[tokio::test] +async fn roundtrip_parquet_select_star_predicate() -> Result<()> { + let ctx = all_types_context().await?; + let sql = "select * from alltypes_plain where id > 4"; + roundtrip_test_sql_with_context(sql, &ctx).await +} + +#[tokio::test] +async fn roundtrip_parquet_select_projection_predicate() -> Result<()> { + let ctx = all_types_context().await?; + let sql = "select string_col, timestamp_col from alltypes_plain where id > 4"; + roundtrip_test_sql_with_context(sql, &ctx).await +} + +#[tokio::test] +async fn roundtrip_empty_projection() -> Result<()> { + let ctx = all_types_context().await?; + let sql = "select 1 from alltypes_plain"; + roundtrip_test_sql_with_context(sql, &ctx).await +} + +#[tokio::test] +async fn roundtrip_memory_source_empty_projection() -> Result<()> { + // Memory scan: `Some(vec![])` must not decode back as `None` + let ctx = SessionContext::new(); + let batch = RecordBatch::try_new( + Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Int64, false), + ])), + vec![ + Arc::new(arrow::array::StringArray::from(vec!["Tom"])), + Arc::new(arrow::array::Int64Array::from(vec![18i64])), + ], + )?; + ctx.register_batch("tmem", batch)?; + let sql = "select 1 from tmem"; + roundtrip_test_sql_with_context(sql, &ctx).await +} + +#[tokio::test] +async fn roundtrip_memory_source() -> Result<()> { + let ctx = SessionContext::new(); + let plan = ctx + .sql("select * from values ('Tom', 18)") + .await? + .create_physical_plan() + .await?; + roundtrip_test(plan) +} + +#[tokio::test] +async fn roundtrip_memory_source_sort_information_and_fetch() -> Result<()> { + use datafusion::datasource::memory::MemorySourceConfig; + use datafusion::datasource::source::DataSource as _; + + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, false), + Field::new("b", DataType::Int64, false), + ])); + let batch = RecordBatch::try_new( + Arc::clone(&schema), + vec![ + Arc::new(arrow::array::StringArray::from(vec!["Tom", "Bob"])), + Arc::new(arrow::array::Int64Array::from(vec![18i64, 21i64])), + ], + )?; + let ordering = LexOrdering::new(vec![PhysicalSortExpr::new( + col("b", &schema)?, + SortOptions { + descending: true, + nulls_first: false, + }, + )]) + .unwrap(); + let source = MemorySourceConfig::try_new(&[vec![batch]], Arc::clone(&schema), None)? + .with_limit(Some(1)) + .with_show_sizes(false) + .try_with_sort_information(vec![ordering])?; + let exec_plan = DataSourceExec::from_data_source(source.clone()); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let decoded = roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + + // The string representation does not include every field; check the + // decoded source directly. + let decoded = decoded + .downcast_ref::() + .expect("expected DataSourceExec"); + let decoded_source = decoded + .data_source() + .downcast_ref::() + .expect("expected MemorySourceConfig"); + assert_eq!(decoded_source.partitions(), source.partitions()); + assert_eq!(decoded_source.original_schema(), source.original_schema()); + assert_eq!(decoded_source.projection(), source.projection()); + assert_eq!(decoded_source.sort_information(), source.sort_information()); + assert_eq!(decoded_source.fetch(), Some(1)); + assert!(!decoded_source.show_sizes()); + Ok(()) +} + +#[tokio::test] +async fn roundtrip_listing_table_with_schema_metadata() -> Result<()> { + let ctx = SessionContext::new(); + let file_format = JsonFormat::default(); + let table_partition_cols = vec![("part".to_owned(), DataType::Int64)]; + let data = "../core/tests/data/partitioned_table_json"; + let listing_table_url = ListingTableUrl::parse(data)?; + let listing_options = ListingOptions::new(Arc::new(file_format)) + .with_table_partition_cols(table_partition_cols); + + let config = ListingTableConfig::new(listing_table_url) + .with_listing_options(listing_options) + .infer_schema(&ctx.state()) + .await?; + + // Decorate metadata onto the inferred ListingTable schema + let schema_with_meta = config + .file_schema + .clone() + .map(|s| { + let mut meta: HashMap = HashMap::new(); + meta.insert("foo.bar".to_string(), "baz".to_string()); + s.as_ref().clone().with_metadata(meta) + }) + .expect("Must decorate metadata"); + + let config = config.with_schema(Arc::new(schema_with_meta)); + ctx.register_table("hive_style", Arc::new(ListingTable::try_new(config)?))?; + + let plan = ctx + .sql("select * from hive_style limit 1") + .await? + .create_physical_plan() + .await?; + + roundtrip_test(plan) +} + +fn roundtrip_file_scan_config(scan_config: FileScanConfig) -> Result { + let exec_plan: Arc = DataSourceExec::from_data_source(scan_config); + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let result_plan = + roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; + + let data_source_exec = result_plan + .downcast_ref::() + .expect("Expected DataSourceExec"); + let file_scan_config = data_source_exec + .data_source() + .downcast_ref::() + .expect("Expected FileScanConfig"); + Ok(file_scan_config.clone()) +} + +#[test] +fn roundtrip_parquet_exec_output_partitioning() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let output_partitioning = + Partitioning::Hash(vec![Arc::new(Column::new("col", 0))], 1); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( + "/path/to/file.parquet".to_string(), + 1024, + )])]) + .with_output_partitioning(Some(output_partitioning.clone())) + .build(); + + assert_eq!( + roundtrip_file_scan_config(scan_config)?.output_partitioning, + Some(output_partitioning) + ); + + Ok(()) +} + +#[test] +fn roundtrip_parquet_exec_range_output_partitioning() -> Result<()> { + let file_schema = + Arc::new(Schema::new(vec![Field::new("col", DataType::Int32, false)])); + let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); + let output_partitioning = Partitioning::Range(RangePartitioning::new( + LexOrdering::new(vec![PhysicalSortExpr::new_default(Arc::new(Column::new( + "col", 0, + )))]) + .unwrap(), + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )); + let scan_config = + FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) + .with_file_groups(vec![ + FileGroup::new(vec![PartitionedFile::new( + "/path/to/file-1.parquet".to_string(), + 1024, + )]), + FileGroup::new(vec![PartitionedFile::new( + "/path/to/file-2.parquet".to_string(), + 1024, + )]), + ]) + .with_output_partitioning(Some(output_partitioning.clone())) + .build(); + + assert_eq!( + roundtrip_file_scan_config(scan_config)?.output_partitioning, + Some(output_partitioning) + ); + + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/tpch.rs b/datafusion/proto/tests/cases/plans/tpch.rs new file mode 100644 index 0000000000000..d24150b3e01ea --- /dev/null +++ b/datafusion/proto/tests/cases/plans/tpch.rs @@ -0,0 +1,335 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! End to end round trips of the TPC-H queries, plus the human readable +//! display of the plans they produce. + +use super::{roundtrip_test_and_return, roundtrip_test_sql_with_context}; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::functions_aggregate::first_last::first_value_udaf; +use datafusion::physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; +use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{PhysicalSortExpr, col}; +use datafusion::prelude::SessionContext; +use datafusion_common::{DataFusionError, Result}; +use datafusion_proto::physical_plan::{ + AsExecutionPlan, DefaultPhysicalExtensionCodec, DefaultPhysicalProtoConverter, +}; +use datafusion_proto::protobuf::PhysicalPlanNode; +use std::sync::Arc; +use std::vec; + +/// Helper function to create a SessionContext with all TPC-H tables registered as external tables +async fn tpch_context() -> Result { + use datafusion_common::test_util::datafusion_test_data; + + let ctx = SessionContext::new(); + let test_data = datafusion_test_data(); + + // TPC-H table names + let tables = [ + "part", "supplier", "partsupp", "customer", "orders", "lineitem", "nation", + "region", + ]; + + // Create external tables for all TPC-H tables + for table in &tables { + let table_sql = format!( + "CREATE EXTERNAL TABLE {table} STORED AS PARQUET LOCATION '{test_data}/tpch_{table}_small.parquet'" + ); + ctx.sql(&table_sql).await.map_err(|e| { + DataFusionError::External( + format!("Failed to create {table} table: {e}").into(), + ) + })?; + } + + Ok(ctx) +} + +/// Helper function to get TPC-H query SQL +fn get_tpch_query_sql(query: usize) -> Result> { + use std::fs; + + if !(1..=22).contains(&query) { + return Err(DataFusionError::External( + format!("Invalid TPC-H query number: {query}").into(), + )); + } + + let filename = format!("../../benchmarks/queries/q{query}.sql"); + let contents = fs::read_to_string(&filename).map_err(|e| { + DataFusionError::External( + format!("Failed to read query file {filename}: {e}").into(), + ) + })?; + + Ok(contents + .split(';') + .map(|s| s.trim()) + .filter(|s| !s.is_empty()) + .map(|s| s.to_string()) + .collect()) +} + +#[tokio::test] +async fn test_serialize_deserialize_tpch_queries() -> Result<()> { + // Create context with TPC-H tables + let ctx = tpch_context().await?; + + // repeat to run all 22 queries + for query in 1..=22 { + // run all statements in the query + let sql = get_tpch_query_sql(query)?; + for stmt in sql { + let logical_plan = ctx.sql(&stmt).await?.into_unoptimized_plan(); + let optimized_plan = ctx.state().optimize(&logical_plan)?; + let physical_plan = ctx.state().create_physical_plan(&optimized_plan).await?; + + // serialize the physical plan + let codec = DefaultPhysicalExtensionCodec {}; + + let proto = + PhysicalPlanNode::try_from_physical_plan(physical_plan.clone(), &codec)?; + + // deserialize the physical plan + let _deserialized_plan = + proto.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + } + } + + Ok(()) +} + +// Bugs: https://github.com/apache/datafusion/issues/16772 +#[tokio::test] +async fn test_round_trip_tpch_queries() -> Result<()> { + // Create context with TPC-H tables + let ctx = tpch_context().await?; + + // repeat to run all 22 queries + for query in 1..=22 { + // run all statements in the query + let sql = get_tpch_query_sql(query)?; + for stmt in sql { + roundtrip_test_sql_with_context(&stmt, &ctx).await?; + } + } + + Ok(()) +} + +// Bug 1 of https://github.com/apache/datafusion/issues/16772 +/// Test that AggregateFunctionExpr human_display field is correctly preserved +/// during serialization/deserialization roundtrip. +/// +/// Test for issue where the human_display field (used for EXPLAIN output) +/// was not being serialized to protobuf, causing it to be lost during roundtrip +/// and resulting in empty or incorrect display strings in query plans. +#[tokio::test] +async fn test_round_trip_human_display() -> Result<()> { + // Create context with TPC-H tables + let ctx = tpch_context().await?; + + let sql = "select r_name, count(1) from region group by r_name"; + roundtrip_test_sql_with_context(sql, &ctx).await?; + + let sql = "select r_name, count(*) from region group by r_name"; + roundtrip_test_sql_with_context(sql, &ctx).await?; + + let sql = "select r_name, count(r_name) from region group by r_name"; + roundtrip_test_sql_with_context(sql, &ctx).await?; + + let sql = "select count(*) as count_star from region"; + roundtrip_test_sql_with_context(sql, &ctx).await?; + + Ok(()) +} + +#[test] +fn test_round_trip_aliased_reverse_human_display() -> Result<()> { + let aggregate_expr = roundtrip_first_value_aggregate( + "agg", + "first_value(b) ORDER BY [b ASC NULLS LAST]", + Some("agg"), + )?; + let reversed = aggregate_expr + .reverse_expr() + .expect("expected reverse expr"); + + assert_eq!(reversed.name(), "agg"); + assert_eq!(reversed.human_display_alias(), Some("agg")); + assert_eq!( + reversed.human_display(), + Some("last_value(b) ORDER BY [b DESC NULLS FIRST]") + ); + + Ok(()) +} + +#[test] +fn test_round_trip_human_display_alias_with_colon() -> Result<()> { + let aggregate_expr = roundtrip_first_value_aggregate( + "agg:one", + "first_value(b) ORDER BY [b ASC NULLS LAST]", + Some("agg:one"), + )?; + + assert_eq!(aggregate_expr.name(), "agg:one"); + assert_eq!(aggregate_expr.human_display_alias(), Some("agg:one")); + assert_eq!( + aggregate_expr.human_display(), + Some("first_value(b) ORDER BY [b ASC NULLS LAST]") + ); + + Ok(()) +} + +#[test] +fn test_round_trip_non_aliased_human_display_ending_like_alias() -> Result<()> { + let aggregate_expr = + roundtrip_first_value_aggregate("agg", "first_value(b) as agg", None)?; + + assert_eq!(aggregate_expr.name(), "agg"); + assert_eq!( + aggregate_expr.human_display(), + Some("first_value(b) as agg") + ); + assert_eq!(aggregate_expr.human_display_alias(), None); + + Ok(()) +} + +fn roundtrip_first_value_aggregate( + alias: &str, + human_display: &str, + human_display_alias: Option<&str>, +) -> Result> { + let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); + let mut builder = + AggregateExprBuilder::new(first_value_udaf(), vec![col("b", &schema)?]) + .order_by(vec![PhysicalSortExpr { + expr: col("b", &schema)?, + options: SortOptions::new(false, false), + }]) + .schema(Arc::clone(&schema)) + .alias(alias) + .human_display(human_display); + if let Some(human_display_alias) = human_display_alias { + builder = builder.human_display_alias(human_display_alias); + } + let agg_expr = builder.build().map(Arc::new)?; + + let plan = Arc::new(AggregateExec::try_new( + AggregateMode::Single, + PhysicalGroupBy::new(vec![], vec![], vec![], false), + vec![agg_expr], + vec![None], + Arc::new(EmptyExec::new(Arc::clone(&schema))), + schema, + )?); + + let ctx = SessionContext::new(); + let codec = DefaultPhysicalExtensionCodec {}; + let proto_converter = DefaultPhysicalProtoConverter {}; + let roundtrip_plan = roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; + let aggregate = roundtrip_plan + .as_ref() + .downcast_ref::() + .expect("expected AggregateExec after roundtrip"); + + Ok(Arc::clone(&aggregate.aggr_expr()[0])) +} + +// Bug 2 of https://github.com/apache/datafusion/issues/16772 +/// Test that PhysicalGroupBy groups field is correctly serialized/deserialized +/// for simple aggregates (no GROUP BY clause). +/// +/// Test for issue where simple aggregates like "SELECT SUM(col1 * col2) FROM table" +/// would incorrectly serialize groups as [[]] instead of [] during roundtrip serialization. +/// The groups field should be empty ([]) when there are no GROUP BY expressions. +#[tokio::test] +async fn test_round_trip_groups_display() -> Result<()> { + // Create context with TPC-H tables + let ctx = tpch_context().await?; + + let sql = "select sum(l_extendedprice * l_discount) as revenue from lineitem;"; + roundtrip_test_sql_with_context(sql, &ctx).await?; + + let sql = "select sum(l_extendedprice) as revenue from lineitem;"; + roundtrip_test_sql_with_context(sql, &ctx).await?; + + Ok(()) +} + +// Bug 3 of https://github.com/apache/datafusion/issues/16772 +/// Test that ScalarFunctionExpr return_field name is correctly preserved +/// during serialization/deserialization roundtrip. +/// +/// Test for issue where the return_field.name for scalar functions +/// was not being serialized to protobuf, causing it to be lost during roundtrip +/// and defaulting to a generic name like "f" instead of the proper function name. +#[tokio::test] +async fn test_round_trip_date_part_display() -> Result<()> { + // Create context with TPC-H tables + let ctx = tpch_context().await?; + + let sql = "select extract(year from l_shipdate) as l_year from lineitem "; + roundtrip_test_sql_with_context(sql, &ctx).await?; + + let sql = "select extract(month from l_shipdate) as l_year from lineitem "; + roundtrip_test_sql_with_context(sql, &ctx).await?; + + Ok(()) +} + +#[tokio::test] +async fn test_tpch_part_in_list_query_with_real_parquet_data() -> Result<()> { + use datafusion_common::test_util::datafusion_test_data; + + let ctx = SessionContext::new(); + + // Register the TPC-H part table using the local test data + let test_data = datafusion_test_data(); + let table_sql = format!( + "CREATE EXTERNAL TABLE part STORED AS PARQUET LOCATION '{test_data}/tpch_part_small.parquet'" + ); + ctx.sql(&table_sql).await.map_err(|e| { + DataFusionError::External(format!("Failed to create part table: {e}").into()) + })?; + + // Test the exact problematic query + let sql = + "SELECT p_size FROM part WHERE p_size IN (14, 6, 5, 31) and p_partkey > 1000"; + + let logical_plan = ctx.sql(sql).await?.into_unoptimized_plan(); + let optimized_plan = ctx.state().optimize(&logical_plan)?; + let physical_plan = ctx.state().create_physical_plan(&optimized_plan).await?; + + // Serialize the physical plan - bug may happen here already but not necessarily manifests + let codec = DefaultPhysicalExtensionCodec {}; + + let proto = PhysicalPlanNode::try_from_physical_plan(physical_plan.clone(), &codec)?; + + // This will fail with the bug, but should succeed when fixed + let _deserialized_plan = proto.try_into_physical_plan(&ctx.task_ctx(), &codec)?; + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/udfs.rs b/datafusion/proto/tests/cases/plans/udfs.rs new file mode 100644 index 0000000000000..08d00030e04e1 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/udfs.rs @@ -0,0 +1,573 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Plans carrying user defined functions, and the extension codec that +//! (de)serializes them. + +use super::{roundtrip_test_and_return, roundtrip_test_with_context}; +use crate::cases::{ + CustomUDWF, CustomUDWFNode, MyAggregateUDF, MyAggregateUdfNode, MyHigherOrderUDF, + MyHigherOrderUdfNode, MyRegexUdf, MyRegexUdfNode, +}; +use datafusion::arrow::array::ArrayRef; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::execution::TaskContext; +use datafusion::logical_expr::{Operator, Volatility, create_udf}; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_expr::expressions::Literal; +use datafusion::physical_expr::window::StandardWindowExpr; +use datafusion::physical_expr::{HigherOrderFunctionExpr, ScalarFunctionExpr}; +use datafusion::physical_plan::aggregates::{ + AggregateExec, AggregateMode, PhysicalGroupBy, +}; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{BinaryExpr, PhysicalSortExpr, col, lit}; +use datafusion::physical_plan::filter::FilterExec; +use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; +use datafusion::physical_plan::windows::{ + BoundedWindowAggExec, PlainAggregateWindowExpr, WindowAggExec, + create_udwf_window_expr, +}; +use datafusion::physical_plan::{ExecutionPlan, InputOrderMode, PhysicalExpr}; +use datafusion::prelude::SessionContext; +use datafusion::scalar::ScalarValue; +use datafusion_common::config::ConfigOptions; +use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; +use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; +use datafusion_expr::{ + AggregateUDF, ColumnarValue, HigherOrderUDF, ScalarFunctionArgs, ScalarUDF, + ScalarUDFImpl, Signature, WindowFrame, WindowFrameBound, WindowUDF, +}; +use datafusion_functions_aggregate::min_max::max_udaf; +use datafusion_physical_expr::expressions::{LambdaVariable, is_not_null, lambda}; +use datafusion_proto::physical_plan::{ + DefaultPhysicalProtoConverter, PhysicalExtensionCodec, + PhysicalProtoConverterExtension, +}; +use prost::Message; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_scalar_udf() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let input = Arc::new(EmptyExec::new(schema.clone())); + + let scalar_fn = Arc::new(|args: &[ColumnarValue]| { + let ColumnarValue::Array(array) = &args[0] else { + panic!("should be array") + }; + Ok(ColumnarValue::from(Arc::new(array.clone()) as ArrayRef)) + }); + + let udf = create_udf( + "dummy", + vec![DataType::Int64], + DataType::Int64, + Volatility::Immutable, + scalar_fn.clone(), + ); + + let fun_def = Arc::new(udf.clone()); + + let expr = ScalarFunctionExpr::new( + "dummy", + fun_def, + vec![col("a", &schema)?], + Field::new("f", DataType::Int64, true).into(), + Arc::new(ConfigOptions::default()), + ); + + let project = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(expr), + alias: "a".to_string(), + }], + input, + )?; + + let ctx = SessionContext::new(); + + ctx.register_udf(udf); + + roundtrip_test_with_context(Arc::new(project), &ctx) +} + +#[derive(Debug)] +struct UDFExtensionCodec; + +impl PhysicalExtensionCodec for UDFExtensionCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[Arc], + _ctx: &TaskContext, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result> { + not_impl_err!("No extension codec provided") + } + + fn try_encode( + &self, + _node: Arc, + _buf: &mut Vec, + _proto_converter: &dyn PhysicalProtoConverterExtension, + ) -> Result<()> { + not_impl_err!("No extension codec provided") + } + + fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { + if name == "regex_udf" { + let proto = MyRegexUdfNode::decode(buf).map_err(|err| { + internal_datafusion_err!("failed to decode regex_udf: {err}") + })?; + + Ok(Arc::new(ScalarUDF::from(MyRegexUdf::new(proto.pattern)))) + } else { + not_impl_err!("unrecognized scalar UDF implementation, cannot decode") + } + } + + fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { + let binding = node.inner(); + if let Some(udf) = binding.downcast_ref::() { + let proto = MyRegexUdfNode { + pattern: udf.pattern.clone(), + }; + proto + .encode(buf) + .map_err(|err| internal_datafusion_err!("failed to encode udf: {err}"))?; + } + Ok(()) + } + + fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { + if name == "aggregate_udf" { + let proto = MyAggregateUdfNode::decode(buf).map_err(|err| { + internal_datafusion_err!("failed to decode aggregate_udf: {err}") + })?; + + Ok(Arc::new(AggregateUDF::from(MyAggregateUDF::new( + proto.result, + )))) + } else { + not_impl_err!("unrecognized scalar UDF implementation, cannot decode") + } + } + + fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { + let binding = node.inner(); + if let Some(udf) = binding.downcast_ref::() { + let proto = MyAggregateUdfNode { + result: udf.result.clone(), + }; + proto.encode(buf).map_err(|err| { + internal_datafusion_err!("failed to encode udf: {err:?}") + })?; + } + Ok(()) + } + + fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { + if name == "custom_udwf" { + let proto = CustomUDWFNode::decode(buf).map_err(|err| { + internal_datafusion_err!("failed to decode custom_udwf: {err}") + })?; + + Ok(Arc::new(WindowUDF::from(CustomUDWF::new(proto.payload)))) + } else { + not_impl_err!( + "unrecognized user-defined window function implementation, cannot decode" + ) + } + } + + fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { + let binding = node.inner(); + if let Some(udwf) = binding.downcast_ref::() { + let proto = CustomUDWFNode { + payload: udwf.payload.clone(), + }; + proto.encode(buf).map_err(|err| { + internal_datafusion_err!("failed to encode udwf: {err:?}") + })?; + } + Ok(()) + } + + fn try_decode_higher_order_function( + &self, + name: &str, + buf: &[u8], + ) -> Result> { + if name == "higher_order_udf" { + let proto = MyHigherOrderUdfNode::decode(buf).map_err(|err| { + internal_datafusion_err!("failed to decode higher_order_udf: {err}") + })?; + + Ok(Arc::new(HigherOrderUDF::new_from_impl( + MyHigherOrderUDF::new(proto.payload), + ))) + } else { + not_impl_err!("unrecognized higher order UDF implementation, cannot decode") + } + } + + fn try_encode_higher_order_function( + &self, + node: &HigherOrderUDF, + buf: &mut Vec, + ) -> Result<()> { + if let Some(hof) = (node.inner().as_ref() as &dyn std::any::Any) + .downcast_ref::() + { + let proto = MyHigherOrderUdfNode { + payload: hof.payload.clone(), + }; + proto.encode(buf).map_err(|err| { + internal_datafusion_err!("failed to encode hof: {err:?}") + })?; + } + Ok(()) + } +} + +#[test] +fn roundtrip_scalar_udf_extension_codec() -> Result<()> { + let field_text = Field::new("text", DataType::Utf8, true); + let field_published = Field::new("published", DataType::Boolean, false); + let field_author = Field::new("author", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_text, field_published, field_author])); + let input = Arc::new(EmptyExec::new(schema.clone())); + + let udf_expr = Arc::new(ScalarFunctionExpr::new( + "regex_udf", + Arc::new(ScalarUDF::from(MyRegexUdf::new(".*".to_string()))), + vec![col("text", &schema)?], + Field::new("f", DataType::Int64, true).into(), + Arc::new(ConfigOptions::default()), + )); + + let filter = Arc::new(FilterExec::try_new( + Arc::new(BinaryExpr::new( + col("published", &schema)?, + Operator::And, + Arc::new(BinaryExpr::new(udf_expr.clone(), Operator::Gt, lit(0))), + )), + input, + )?); + let aggr_expr = + AggregateExprBuilder::new(max_udaf(), vec![udf_expr as Arc]) + .schema(schema.clone()) + .alias("max") + .build() + .map(Arc::new)?; + + let window = Arc::new(WindowAggExec::try_new( + vec![Arc::new(PlainAggregateWindowExpr::new( + aggr_expr.clone(), + &[col("author", &schema)?], + &[], + Arc::new(WindowFrame::new(None)), + None, + ))], + filter, + true, + )?); + + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new(vec![], vec![], vec![], false), + vec![aggr_expr], + vec![None], + window, + schema, + )?); + + let ctx = SessionContext::new(); + let proto_converter = DefaultPhysicalProtoConverter {}; + roundtrip_test_and_return(aggregate, &ctx, &UDFExtensionCodec, &proto_converter)?; + Ok(()) +} + +#[test] +fn roundtrip_higher_order_udf() -> Result<()> { + let element_field = Arc::new(Field::new("v", DataType::Int32, true)); + let list_field = Field::new( + "list_col", + DataType::List(Arc::clone(&element_field)), + false, + ); + let schema = Arc::new(Schema::new(vec![list_field])); + + let input = Arc::new(EmptyExec::new(schema.clone())); + + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + + let expr = HigherOrderFunctionExpr::try_new_with_schema( + Arc::clone(&hof), + vec![ + col("list_col", &schema)?, + lambda( + ["v"], + is_not_null(Arc::new(LambdaVariable::new(1, element_field)))?, + )?, + ], + &schema, + Arc::new(ConfigOptions::default()), + )?; + + let project = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: Arc::new(expr), + alias: "a".to_string(), + }], + input, + )?; + + let ctx = SessionContext::new(); + ctx.register_higher_order_function(hof); + + roundtrip_test_with_context(Arc::new(project), &ctx) +} + +#[test] +fn roundtrip_higher_order_udf_extension_codec() -> Result<()> { + let element_field = Arc::new(Field::new("v", DataType::Int32, true)); + let list_field = Field::new( + "list_col", + DataType::List(Arc::clone(&element_field)), + false, + ); + let schema = Arc::new(Schema::new(vec![list_field])); + let input = Arc::new(EmptyExec::new(schema.clone())); + + let lambda_body = Arc::new(LambdaVariable::new(1, Arc::clone(&element_field))); + let lambda_expr = lambda(["v"], lambda_body)?; + + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + let hof_expr = Arc::new(HigherOrderFunctionExpr::try_new_with_schema( + hof, + vec![col("list_col", &schema)?, lambda_expr], + &schema, + Arc::new(ConfigOptions::default()), + )?); + + let project = ProjectionExec::try_new( + vec![ProjectionExpr { + expr: hof_expr, + alias: "out".to_string(), + }], + input, + )?; + + let ctx = SessionContext::new(); + let proto_converter = DefaultPhysicalProtoConverter {}; + roundtrip_test_and_return( + Arc::new(project), + &ctx, + &UDFExtensionCodec, + &proto_converter, + )?; + Ok(()) +} + +#[test] +fn roundtrip_udwf_extension_codec() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let custom_udwf = Arc::new(WindowUDF::from(CustomUDWF::new("payload".to_string()))); + let udwf = create_udwf_window_expr( + &custom_udwf, + &[col("a", &schema)?], + schema.as_ref(), + "custom_udwf(a) PARTITION BY [b] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW".to_string(), + false, + )?; + + let window_frame = WindowFrame::new_bounds( + datafusion_expr::WindowFrameUnits::Range, + WindowFrameBound::Preceding(ScalarValue::Int64(None)), + WindowFrameBound::CurrentRow, + ); + + let udwf_expr = Arc::new(StandardWindowExpr::new( + udwf, + &[col("b", &schema)?], + &[PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }], + Arc::new(window_frame), + )); + + let input = Arc::new(EmptyExec::new(schema.clone())); + let window = Arc::new(BoundedWindowAggExec::try_new( + vec![udwf_expr], + input, + InputOrderMode::Sorted, + true, + )?); + + let ctx = SessionContext::new(); + let proto_converter = DefaultPhysicalProtoConverter {}; + roundtrip_test_and_return(window, &ctx, &UDFExtensionCodec, &proto_converter)?; + Ok(()) +} + +#[test] +fn roundtrip_aggregate_udf_extension_codec() -> Result<()> { + let field_text = Field::new("text", DataType::Utf8, true); + let field_published = Field::new("published", DataType::Boolean, false); + let field_author = Field::new("author", DataType::Utf8, false); + let schema = Arc::new(Schema::new(vec![field_text, field_published, field_author])); + let input = Arc::new(EmptyExec::new(schema.clone())); + + let udf_expr = Arc::new(ScalarFunctionExpr::new( + "regex_udf", + Arc::new(ScalarUDF::from(MyRegexUdf::new(".*".to_string()))), + vec![col("text", &schema)?], + Field::new("f", DataType::Int64, true).into(), + Arc::new(ConfigOptions::default()), + )); + + let udaf = Arc::new(AggregateUDF::from(MyAggregateUDF::new( + "result".to_string(), + ))); + let aggr_args: Vec> = + vec![Arc::new(Literal::new(ScalarValue::from(42)))]; + + let aggr_expr = AggregateExprBuilder::new(Arc::clone(&udaf), aggr_args.clone()) + .schema(Arc::clone(&schema)) + .alias("aggregate_udf") + .build() + .map(Arc::new)?; + + let filter = Arc::new(FilterExec::try_new( + Arc::new(BinaryExpr::new( + col("published", &schema)?, + Operator::And, + Arc::new(BinaryExpr::new(udf_expr, Operator::Gt, lit(0))), + )), + input, + )?); + + let window = Arc::new(WindowAggExec::try_new( + vec![Arc::new(PlainAggregateWindowExpr::new( + aggr_expr, + &[col("author", &schema)?], + &[], + Arc::new(WindowFrame::new(None)), + None, + ))], + filter, + true, + )?); + + let aggr_expr = AggregateExprBuilder::new(udaf, aggr_args.clone()) + .schema(Arc::clone(&schema)) + .alias("aggregate_udf") + .distinct() + .ignore_nulls() + .build() + .map(Arc::new)?; + + let aggregate = Arc::new(AggregateExec::try_new( + AggregateMode::Final, + PhysicalGroupBy::new(vec![], vec![], vec![], false), + vec![aggr_expr], + vec![None], + window, + schema, + )?); + + let ctx = SessionContext::new(); + let proto_converter = DefaultPhysicalProtoConverter {}; + roundtrip_test_and_return(aggregate, &ctx, &UDFExtensionCodec, &proto_converter)?; + Ok(()) +} + +#[tokio::test] +async fn roundtrip_async_func_exec() -> Result<()> { + #[derive(Debug, PartialEq, Eq, Hash)] + struct TestAsyncUDF { + signature: Signature, + } + + impl TestAsyncUDF { + fn new() -> Self { + Self { + signature: Signature::exact(vec![DataType::Int64], Volatility::Volatile), + } + } + } + + impl ScalarUDFImpl for TestAsyncUDF { + fn name(&self) -> &str { + "test_async_udf" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Int64) + } + + fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { + not_impl_err!("Must call from `invoke_async_with_args`") + } + } + + #[async_trait::async_trait] + impl AsyncScalarUDFImpl for TestAsyncUDF { + async fn invoke_async_with_args( + &self, + args: ScalarFunctionArgs, + ) -> Result { + Ok(args.args[0].clone()) + } + } + + let ctx = SessionContext::new(); + let async_udf = AsyncScalarUDF::new(Arc::new(TestAsyncUDF::new())); + ctx.register_udf(async_udf.into_scalar_udf()); + + let physical_plan = ctx + .sql("select test_async_udf(1)") + .await? + .create_physical_plan() + .await?; + + roundtrip_test_with_context(physical_plan, &ctx)?; + + Ok(()) +} diff --git a/datafusion/proto/tests/cases/plans/windows.rs b/datafusion/proto/tests/cases/plans/windows.rs new file mode 100644 index 0000000000000..cfd2532a88745 --- /dev/null +++ b/datafusion/proto/tests/cases/plans/windows.rs @@ -0,0 +1,308 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The window execs and their window functions. + +use super::roundtrip_test; +use datafusion::arrow::compute::kernels::sort::SortOptions; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::functions_aggregate::count::count_udaf; +use datafusion::functions_aggregate::sum::sum_udaf; +use datafusion::functions_window::nth_value::nth_value_udwf; +use datafusion::functions_window::row_number::row_number_udwf; +use datafusion::physical_expr::aggregate::AggregateExprBuilder; +use datafusion::physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; +use datafusion::physical_plan::InputOrderMode; +use datafusion::physical_plan::empty::EmptyExec; +use datafusion::physical_plan::expressions::{PhysicalSortExpr, cast, col, lit}; +use datafusion::physical_plan::windows::{ + BoundedWindowAggExec, PlainAggregateWindowExpr, WindowAggExec, + create_udwf_window_expr, +}; +use datafusion::scalar::ScalarValue; +use datafusion_common::Result; +use datafusion_expr::{WindowFrame, WindowFrameBound}; +use datafusion_functions_aggregate::average::avg_udaf; +use std::sync::Arc; +use std::vec; + +#[test] +fn roundtrip_udwf() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let udwf_expr = Arc::new(StandardWindowExpr::new( + create_udwf_window_expr( + &row_number_udwf(), + &[], + &schema, + "row_number() PARTITION BY [a] ORDER BY [b] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW".to_string(), + false, + )?, + &[ + col("a", &schema)? + ], + &[ + PhysicalSortExpr::new(col("b", &schema)?, SortOptions::new(true, true)) + ], + Arc::new(WindowFrame::new(None)), + )); + + let input = Arc::new(EmptyExec::new(schema.clone())); + + roundtrip_test(Arc::new(BoundedWindowAggExec::try_new( + vec![udwf_expr], + input, + InputOrderMode::Sorted, + true, + )?)) +} + +#[test] +fn roundtrip_window() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + let window_frame = WindowFrame::new_bounds( + datafusion_expr::WindowFrameUnits::Range, + WindowFrameBound::Preceding(ScalarValue::Int64(None)), + WindowFrameBound::CurrentRow, + ); + + let nth_value_window = + create_udwf_window_expr( + &nth_value_udwf(), + &[col("a", &schema)?, + lit(2)], schema.as_ref(), + "NTH_VALUE(a, 2) PARTITION BY [b] ORDER BY [a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW".to_string(), + false, + )?; + let udwf_expr = Arc::new(StandardWindowExpr::new( + nth_value_window, + &[col("b", &schema)?], + &[PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }], + Arc::new(window_frame), + )); + + let plain_aggr_window_expr = Arc::new(PlainAggregateWindowExpr::new( + AggregateExprBuilder::new( + avg_udaf(), + vec![cast(col("b", &schema)?, &schema, DataType::Float64)?], + ) + .schema(Arc::clone(&schema)) + .alias("avg(b)") + .build() + .map(Arc::new)?, + &[], + &[], + Arc::new(WindowFrame::new(None)), + None, + )); + + let window_frame = WindowFrame::new_bounds( + datafusion_expr::WindowFrameUnits::Range, + WindowFrameBound::CurrentRow, + WindowFrameBound::Preceding(ScalarValue::Int64(None)), + ); + + let args = vec![cast(col("a", &schema)?, &schema, DataType::Float64)?]; + let sum_expr = AggregateExprBuilder::new(sum_udaf(), args) + .schema(Arc::clone(&schema)) + .alias("SUM(a) RANGE BETWEEN CURRENT ROW AND UNBOUNDED PRECEDING") + .build() + .map(Arc::new)?; + + let sliding_aggr_window_expr = Arc::new(SlidingAggregateWindowExpr::new( + sum_expr, + &[], + &[], + Arc::new(window_frame), + None, + )); + + let input = Arc::new(EmptyExec::new(schema.clone())); + + roundtrip_test(Arc::new(WindowAggExec::try_new( + vec![plain_aggr_window_expr, sliding_aggr_window_expr, udwf_expr], + input, + false, + )?)) +} + +#[test] +fn roundtrip_window_distinct() -> Result<()> { + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + // Create a distinct count window expression with unbounded frame (becomes PlainAggregateWindowExpr) + let distinct_count_expr = Arc::new(PlainAggregateWindowExpr::new( + AggregateExprBuilder::new(count_udaf(), vec![col("a", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count(DISTINCT a)") + .distinct() // Enable distinct + .build() + .map(Arc::new)?, + &[col("b", &schema)?], // partition by b + &[], // no order by + Arc::new(WindowFrame::new(None)), // unbounded frame + None, + )); + + // Create a distinct sum window expression with bounded frame (becomes SlidingAggregateWindowExpr) + let bounded_frame = WindowFrame::new_bounds( + datafusion_expr::WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))), + WindowFrameBound::CurrentRow, + ); + + let distinct_sum_expr = Arc::new(SlidingAggregateWindowExpr::new( + AggregateExprBuilder::new( + sum_udaf(), + vec![cast(col("a", &schema)?, &schema, DataType::Float64)?], + ) + .schema(Arc::clone(&schema)) + .alias("sum(DISTINCT a)") + .distinct() // Enable distinct + .with_ignore_nulls(true) // Enable ignore nulls + .build() + .map(Arc::new)?, + &[], // no partition by + &[], // no order by + Arc::new(bounded_frame), // bounded frame + None, + )); + + let input = Arc::new(EmptyExec::new(schema.clone())); + + roundtrip_test(Arc::new(WindowAggExec::try_new( + vec![distinct_count_expr, distinct_sum_expr], + input, + false, + )?)) +} + +#[test] +fn test_distinct_window_serialization_end_to_end() -> Result<()> { + // Create a more comprehensive test that verifies distinct window functions + // work properly through the entire serialization/deserialization pipeline + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + // Test 1: DISTINCT COUNT with IGNORE NULLS + let distinct_count_ignore_nulls = Arc::new(PlainAggregateWindowExpr::new( + AggregateExprBuilder::new(count_udaf(), vec![col("a", &schema)?]) + .schema(Arc::clone(&schema)) + .alias("count_distinct_ignore_nulls") + .distinct() + .with_ignore_nulls(true) + .build() + .map(Arc::new)?, + &[col("b", &schema)?], + &[], + Arc::new(WindowFrame::new(None)), + None, + )); + + // Test 2: DISTINCT SUM (without ignore nulls) + let bounded_frame = WindowFrame::new_bounds( + datafusion_expr::WindowFrameUnits::Rows, + WindowFrameBound::Preceding(ScalarValue::UInt64(Some(2))), + WindowFrameBound::CurrentRow, + ); + + let distinct_sum = Arc::new(SlidingAggregateWindowExpr::new( + AggregateExprBuilder::new( + sum_udaf(), + vec![cast(col("a", &schema)?, &schema, DataType::Float64)?], + ) + .schema(Arc::clone(&schema)) + .alias("sum_distinct") + .distinct() + .build() + .map(Arc::new)?, + &[], + &[], + Arc::new(bounded_frame), + None, + )); + + let input = Arc::new(EmptyExec::new(schema.clone())); + + let window_exec = Arc::new(WindowAggExec::try_new( + vec![distinct_count_ignore_nulls, distinct_sum], + input, + false, + )?); + + // Perform the roundtrip test + roundtrip_test(window_exec) +} + +/// Tests that `lead` window function with offset and default value args +/// survives a protobuf round-trip. This is a regression test for a bug +/// where `expressions()` (used during serialization) returns only the +/// column expression for lead/lag, silently dropping the offset and +/// default value literal args. +#[test] +fn roundtrip_lead_with_default_value() -> Result<()> { + use datafusion::functions_window::lead_lag::lead_udwf; + + let field_a = Field::new("a", DataType::Int64, false); + let field_b = Field::new("b", DataType::Int64, false); + let schema = Arc::new(Schema::new(vec![field_a, field_b])); + + // lead(a, 2, 42) — column a, offset 2, default value 42 + let lead_window = create_udwf_window_expr( + &lead_udwf(), + &[col("a", &schema)?, lit(2i64), lit(42i64)], + schema.as_ref(), + "test lead with default".to_string(), + false, + )?; + + let udwf_expr = Arc::new(StandardWindowExpr::new( + lead_window, + &[col("b", &schema)?], + &[PhysicalSortExpr { + expr: col("a", &schema)?, + options: SortOptions { + descending: false, + nulls_first: false, + }, + }], + Arc::new(WindowFrame::new(None)), + )); + + let input = Arc::new(EmptyExec::new(schema.clone())); + + roundtrip_test(Arc::new(BoundedWindowAggExec::try_new( + vec![udwf_expr], + input, + InputOrderMode::Sorted, + true, + )?)) +} diff --git a/datafusion/proto/tests/cases/public_conversions.rs b/datafusion/proto/tests/cases/public_conversions.rs new file mode 100644 index 0000000000000..1cda8e01a0765 --- /dev/null +++ b/datafusion/proto/tests/cases/public_conversions.rs @@ -0,0 +1,128 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Compile-time guard for the `From` / `TryFrom` conversions between DataFusion +//! types and `datafusion_proto::protobuf` messages that downstream crates call. +//! +//! These impls were silently dropped once (see +//! ): they were replaced by +//! crate-local conversion traits as a stopgap during the `datafusion-proto-models` +//! extraction, and `cargo-semver-checks` has no lint for a removed hand-written +//! trait impl, so nothing caught the break. Coercing each conversion to a `fn` +//! pointer here does — moving an impl between crates is fine, removing one stops +//! compiling. +//! +//! Only the spelling is asserted. Behaviour is covered by the round-trip tests +//! next to each impl. + +use datafusion_common::config::{ + CsvOptions, JsonOptions, ParquetCdcOptions, ParquetColumnOptions, ParquetOptions, + TableParquetOptions, +}; +use datafusion_common::display::StringifiedPlan; +use datafusion_common::{ + JoinConstraint, JoinType, NullEquality, TableReference, UnnestOptions, +}; +use datafusion_datasource::file_groups::FileGroup; +use datafusion_datasource::file_sink_config::FileSinkConfig; +use datafusion_datasource::{FileRange, PartitionedFile}; +use datafusion_datasource_csv::file_format::{CsvFormatFactory, CsvSink}; +use datafusion_datasource_json::file_format::{JsonFormatFactory, JsonSink}; +use datafusion_datasource_parquet::file_format::{ParquetFormatFactory, ParquetSink}; +use datafusion_expr::dml::MergeIntoClauseKind; +use datafusion_expr::expr::NullTreatment; +use datafusion_expr::{WindowFrame, WindowFrameBound, WindowFrameUnits}; +use datafusion_physical_expr::expressions::Column; +use datafusion_proto::protobuf; + +/// Asserts `T: From` by naming the conversion. +fn assert_from>() { + let _: fn(F) -> T = From::from; +} + +/// Asserts `T: TryFrom` by naming the conversion. +fn assert_try_from>() { + let _: fn(F) -> Result = TryFrom::try_from; +} + +#[test] +fn file_scan_conversions_are_std_traits() { + assert_try_from::<&protobuf::PartitionedFile, PartitionedFile>(); + assert_try_from::<&PartitionedFile, protobuf::PartitionedFile>(); + assert_try_from::<&protobuf::FileRange, FileRange>(); + assert_try_from::<&FileRange, protobuf::FileRange>(); + assert_try_from::<&protobuf::FileGroup, FileGroup>(); + assert_try_from::<&FileGroup, protobuf::FileGroup>(); + assert_from::<&protobuf::PhysicalColumn, Column>(); + assert_from::<&Column, protobuf::PhysicalColumn>(); + assert_try_from::<&[PartitionedFile], protobuf::FileGroup>(); +} + +#[test] +fn file_sink_conversions_are_std_traits() { + assert_try_from::<&protobuf::FileSinkConfig, FileSinkConfig>(); + assert_try_from::<&FileSinkConfig, protobuf::FileSinkConfig>(); + assert_try_from::<&protobuf::JsonSink, JsonSink>(); + assert_try_from::<&JsonSink, protobuf::JsonSink>(); + assert_try_from::<&protobuf::CsvSink, CsvSink>(); + assert_try_from::<&CsvSink, protobuf::CsvSink>(); + assert_try_from::<&protobuf::ParquetSink, ParquetSink>(); + assert_try_from::<&ParquetSink, protobuf::ParquetSink>(); +} + +#[test] +fn window_frame_conversions_are_std_traits() { + assert_try_from::(); + assert_try_from::<&WindowFrame, protobuf::WindowFrame>(); + assert_try_from::(); + assert_try_from::<&WindowFrameBound, protobuf::WindowFrameBound>(); + assert_from::(); + assert_from::(); + assert_from::(); + assert_from::(); + assert_from::(); + assert_from::(); +} + +#[test] +fn common_type_conversions_are_std_traits() { + assert_from::<&protobuf::UnnestOptions, UnnestOptions>(); + assert_from::<&UnnestOptions, protobuf::UnnestOptions>(); + assert_try_from::(); + assert_from::(); + assert_from::<&protobuf::StringifiedPlan, StringifiedPlan>(); + assert_from::<&StringifiedPlan, protobuf::StringifiedPlan>(); + assert_from::(); + assert_from::(); + assert_from::(); + assert_from::(); + assert_from::(); + assert_from::(); +} + +#[test] +fn file_format_option_conversions_are_std_traits() { + assert_from::<&protobuf::CsvOptions, CsvOptions>(); + assert_from::<&protobuf::JsonOptions, JsonOptions>(); + assert_try_from::<&protobuf::ParquetOptions, ParquetOptions>(); + assert_from::(); + assert_from::(); + assert_try_from::<&protobuf::TableParquetOptions, TableParquetOptions>(); + assert_from::<&CsvFormatFactory, protobuf::CsvOptions>(); + assert_from::<&JsonFormatFactory, protobuf::JsonOptions>(); + assert_from::<&ParquetFormatFactory, protobuf::TableParquetOptions>(); +} diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs index 3e79ddab723eb..a450f7a7e888f 100644 --- a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs @@ -45,6 +45,7 @@ use std::vec; use datafusion::catalog::{TableProvider, TableProviderFactory}; use datafusion::datasource::DefaultTableSource; +use datafusion::datasource::empty::EmptyTable; use datafusion::datasource::file_format::arrow::ArrowFormatFactory; use datafusion::datasource::file_format::csv::CsvFormatFactory; use datafusion::datasource::file_format::parquet::ParquetFormatFactory; @@ -68,24 +69,32 @@ use datafusion::physical_expr::PhysicalExpr; use datafusion::prelude::*; use datafusion::test_util::{TestTableFactory, TestTableProvider}; use datafusion_common::config::TableOptions; -use datafusion_common::format::ExplainFormat; +use datafusion_common::format::{ + ExplainAnalyzeCategories, ExplainFormat, MetricCategory, MetricType, +}; use datafusion_common::scalar::ScalarStructBuilder; use datafusion_common::{ - DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, TableReference, - internal_datafusion_err, internal_err, not_impl_err, plan_err, + Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SplitPoint, + TableReference, internal_datafusion_err, internal_err, not_impl_err, plan_err, }; use datafusion_execution::TaskContext; use datafusion_expr::dml::CopyTo; +use datafusion_expr::dml::{ + MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, +}; use datafusion_expr::expr::{ - self, Between, BinaryExpr, Case, Cast, GroupingSet, InList, Like, NullTreatment, - ScalarFunction, Unnest, WildcardOptions, + self, Between, BinaryExpr, Case, Cast, GroupingSet, InList, LambdaVariable, Like, + NullTreatment, ScalarFunction, Unnest, WildcardOptions, +}; +use datafusion_expr::logical_plan::{ + ExplainOption, Extension, UserDefinedLogicalNodeCore, }; -use datafusion_expr::logical_plan::{Extension, UserDefinedLogicalNodeCore}; use datafusion_expr::{ - Accumulator, AggregateUDF, ColumnarValue, ExprFunctionExt, ExprSchemable, - LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, Operator, PartitionEvaluator, - ScalarUDF, Signature, TryCast, Volatility, WindowFrame, WindowFrameBound, - WindowFrameUnits, WindowFunctionDefinition, WindowUDF, WindowUDFImpl, + Accumulator, AggregateUDF, ColumnarValue, DmlStatement, ExprFunctionExt, + ExprSchemable, HigherOrderUDF, LimitEffect, Literal, LogicalPlan, LogicalPlanBuilder, + Operator, PartitionEvaluator, RangePartitioning, Repartition, ScalarUDF, Signature, + TryCast, Volatility, WindowFrame, WindowFrameBound, WindowFrameUnits, + WindowFunctionDefinition, WindowUDF, WindowUDFImpl, WriteOp, }; use datafusion_functions_aggregate::average::avg_udaf; use datafusion_functions_aggregate::expr_fn::{ @@ -109,7 +118,10 @@ use datafusion_proto::logical_plan::{ }; use datafusion_proto::protobuf; -use crate::cases::{MyAggregateUDF, MyAggregateUdfNode, MyRegexUdf, MyRegexUdfNode}; +use crate::cases::{ + MyAggregateUDF, MyAggregateUdfNode, MyHigherOrderUDF, MyHigherOrderUdfNode, + MyRegexUdf, MyRegexUdfNode, +}; #[cfg(feature = "json")] fn roundtrip_json_test(proto: &protobuf::LogicalExprNode) { @@ -138,7 +150,7 @@ fn roundtrip_expr_test_with_codec( let round_trip: Expr = from_proto::parse_expr(&proto, ctx.task_ctx().as_ref(), codec).unwrap(); - assert_eq!(format!("{:?}", &initial_struct), format!("{round_trip:?}")); + assert_eq!(format!("{initial_struct:?}"), format!("{round_trip:?}")); roundtrip_json_test(&proto); } @@ -299,6 +311,74 @@ async fn roundtrip_explain_format_tree() -> Result<()> { Ok(()) } +/// Build an `EXPLAIN`/`EXPLAIN ANALYZE` plan with statement-level overrides +/// set directly via the builder, then assert the proto round-trip preserves +/// every field. Going through the builder avoids depending on parser support +/// for the parenthesized option syntax in this test crate. +async fn assert_explain_roundtrip(option: ExplainOption) -> Result<()> { + let ctx = SessionContext::new(); + let input = ctx.sql("SELECT 1 AS x").await?.into_optimized_plan()?; + let plan = LogicalPlanBuilder::from(input) + .explain_option_format(option)? + .build()?; + + let bytes = logical_plan_to_bytes(&plan)?; + let round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(plan, round_trip); + Ok(()) +} + +#[tokio::test] +async fn roundtrip_explain_show_statistics_override() -> Result<()> { + for show_statistics in [None, Some(true), Some(false)] { + assert_explain_roundtrip( + ExplainOption::default() + .with_format(ExplainFormat::Indent) + .with_show_statistics(show_statistics), + ) + .await?; + } + Ok(()) +} + +#[tokio::test] +async fn roundtrip_analyze_level_override() -> Result<()> { + for analyze_level in [None, Some(MetricType::Summary), Some(MetricType::Dev)] { + assert_explain_roundtrip( + ExplainOption::default() + .with_analyze(true) + .with_analyze_level(analyze_level), + ) + .await?; + } + Ok(()) +} + +#[tokio::test] +async fn roundtrip_analyze_categories_override() -> Result<()> { + let cases = [ + None, + Some(ExplainAnalyzeCategories::All), + Some(ExplainAnalyzeCategories::Only(vec![])), + Some(ExplainAnalyzeCategories::Only(vec![MetricCategory::Rows])), + Some(ExplainAnalyzeCategories::Only(vec![ + MetricCategory::Rows, + MetricCategory::Bytes, + MetricCategory::Timing, + MetricCategory::Uncategorized, + ])), + ]; + for analyze_categories in cases { + assert_explain_roundtrip( + ExplainOption::default() + .with_analyze(true) + .with_analyze_categories(analyze_categories), + ) + .await?; + } + Ok(()) +} + #[tokio::test] async fn roundtrip_custom_listing_tables() -> Result<()> { let ctx = SessionContext::new(); @@ -328,6 +408,129 @@ async fn roundtrip_custom_listing_tables() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_create_external_table_multiple_locations() -> Result<()> { + let ctx = SessionContext::new(); + + // Planning a CREATE EXTERNAL TABLE does not read the referenced files, so + // the paths need not exist. Multiple locations must survive the round-trip + // through the `repeated locations` proto field. + let query = "CREATE EXTERNAL TABLE t (a INTEGER, b INTEGER) + STORED AS CSV + LOCATION ('file_a.csv', 'file_b.csv') + OPTIONS ('format.has_header' 'true')"; + + let plan = ctx.state().create_logical_plan(query).await?; + let bytes = logical_plan_to_bytes(&plan)?; + let protobuf_plan = protobuf::LogicalPlanNode::decode(bytes.as_ref()) + .expect("failed to decode CreateExternalTable proto"); + #[cfg(feature = "json")] + { + let json = serde_json::to_string(&protobuf_plan).unwrap(); + assert!(!json.contains("\"location\":")); + assert!(json.contains("\"locations\":[\"file_a.csv\",\"file_b.csv\"]")); + } + let Some(protobuf::logical_plan_node::LogicalPlanType::CreateExternalTable( + create_external_table, + )) = protobuf_plan.logical_plan_type + else { + panic!("expected a CreateExternalTable proto"); + }; + assert!(create_external_table.location.is_empty()); + assert_eq!( + create_external_table.locations, + vec!["file_a.csv".to_string(), "file_b.csv".to_string()] + ); + + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(plan, logical_round_trip); + + let LogicalPlan::Ddl(datafusion_expr::DdlStatement::CreateExternalTable(rt)) = + logical_round_trip + else { + panic!("expected a CreateExternalTable plan"); + }; + assert_eq!( + rt.locations, + vec!["file_a.csv".to_string(), "file_b.csv".to_string()] + ); + + Ok(()) +} + +#[tokio::test] +async fn roundtrip_create_external_table_single_location_legacy_field() -> Result<()> { + let ctx = SessionContext::new(); + let query = "CREATE EXTERNAL TABLE t (a INTEGER) + STORED AS CSV + LOCATION 'file.csv'"; + + let plan = ctx.state().create_logical_plan(query).await?; + let bytes = logical_plan_to_bytes(&plan)?; + let protobuf_plan = protobuf::LogicalPlanNode::decode(bytes.as_ref()) + .expect("failed to decode CreateExternalTable proto"); + #[cfg(feature = "json")] + { + let json = serde_json::to_string(&protobuf_plan).unwrap(); + assert!(json.contains("\"location\":\"file.csv\"")); + assert!(!json.contains("\"locations\"")); + } + let Some(protobuf::logical_plan_node::LogicalPlanType::CreateExternalTable( + create_external_table, + )) = protobuf_plan.logical_plan_type + else { + panic!("expected a CreateExternalTable proto"); + }; + assert_eq!(create_external_table.location, "file.csv"); + assert!(create_external_table.locations.is_empty()); + + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(plan, logical_round_trip); + + Ok(()) +} + +#[tokio::test] +async fn roundtrip_create_external_table_legacy_location() -> Result<()> { + let ctx = SessionContext::new(); + let schema = DFSchema::empty(); + let create_external_table = protobuf::CreateExternalTableNode { + name: Some(protobuf::TableReference::from(TableReference::bare("t"))), + location: "legacy.csv".to_string(), + locations: vec![], + file_type: "CSV".to_string(), + schema: Some((&schema).try_into()?), + table_partition_cols: vec![], + if_not_exists: false, + or_replace: false, + temporary: false, + definition: String::new(), + order_exprs: vec![], + unbounded: false, + options: HashMap::new(), + constraints: Some(Constraints::default().into()), + column_defaults: HashMap::new(), + }; + let protobuf_plan = protobuf::LogicalPlanNode { + logical_plan_type: Some( + protobuf::logical_plan_node::LogicalPlanType::CreateExternalTable( + create_external_table, + ), + ), + }; + let bytes = protobuf_plan.encode_to_vec(); + + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + let LogicalPlan::Ddl(datafusion_expr::DdlStatement::CreateExternalTable(rt)) = + logical_round_trip + else { + panic!("expected a CreateExternalTable plan"); + }; + assert_eq!(rt.locations, vec!["legacy.csv".to_string()]); + + Ok(()) +} + #[tokio::test] async fn roundtrip_logical_plan_aggregation_with_pk() -> Result<()> { let ctx = SessionContext::new(); @@ -452,6 +655,181 @@ async fn roundtrip_logical_plan_dml() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_logical_plan_dml_merge_into() -> Result<()> { + let ctx = SessionContext::new(); + let schema = Schema::new(vec![ + Field::new("a", DataType::Int64, true), + Field::new("b", DataType::Decimal128(15, 2), true), + ]); + ctx.register_csv( + "t1", + "tests/testdata/test.csv", + CsvReadOptions::default().schema(&schema), + ) + .await?; + + let scan = ctx.table("t1").await?.into_optimized_plan()?; + let target = match &scan { + LogicalPlan::TableScan(t) => Arc::clone(&t.source), + other => panic!("expected TableScan, got {other:?}"), + }; + + let merge = WriteOp::MergeInto(Box::new(MergeIntoOp { + on: col("a").eq(lit(1_i64)), + clauses: vec![ + MergeIntoClause { + kind: MergeIntoClauseKind::Matched, + predicate: Some(col("b").gt(lit(ScalarValue::Decimal128( + Some(0), + 15, + 2, + )))), + action: MergeIntoAction::Update(vec![("b".to_string(), col("b"))]), + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatched, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec!["a".to_string(), "b".to_string()], + values: vec![col("a"), col("b")], + }, + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatchedByTarget, + predicate: None, + action: MergeIntoAction::Insert { + columns: vec![], + values: vec![col("a"), col("b")], + }, + }, + MergeIntoClause { + kind: MergeIntoClauseKind::NotMatchedBySource, + predicate: Some(col("a").eq(lit(2_i64))), + action: MergeIntoAction::Delete, + }, + ], + })); + + let plan = LogicalPlan::Dml(DmlStatement::new( + "t1".into(), + target, + merge, + Arc::new(scan), + )); + + let bytes = logical_plan_to_bytes(&plan)?; + let round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan}"), format!("{round_trip}")); + Ok(()) +} + +#[test] +fn parse_write_op_merge_into_without_payload_errors() { + let ctx = SessionContext::new(); + let codec = DefaultLogicalExtensionCodec {}; + let node = protobuf::DmlNode { + dml_type: protobuf::dml_node::Type::MergeInto.into(), + ..Default::default() + }; + let err = from_proto::parse_write_op(&node, ctx.task_ctx().as_ref(), &codec) + .expect_err("MergeInto tag without payload must fail"); + assert!( + err.to_string().contains("merge_into"), + "unexpected error: {err}" + ); +} + +/// Build a `DmlNode` whose `merge_into` payload is exactly the supplied +/// `MergeIntoOpNode`. Used by the error-path tests below. +fn dml_node_with_merge_payload(payload: protobuf::MergeIntoOpNode) -> protobuf::DmlNode { + protobuf::DmlNode { + dml_type: protobuf::dml_node::Type::MergeInto.into(), + merge_into: Some(Box::new(payload)), + ..Default::default() + } +} + +#[test] +fn parse_merge_into_op_missing_on_errors() { + let ctx = SessionContext::new(); + let codec = DefaultLogicalExtensionCodec {}; + let node = dml_node_with_merge_payload(protobuf::MergeIntoOpNode { + on: None, + clauses: vec![], + }); + let err = from_proto::parse_write_op(&node, ctx.task_ctx().as_ref(), &codec) + .expect_err("missing `on` must fail"); + assert!(err.to_string().contains("`on`"), "unexpected error: {err}"); +} + +#[test] +fn parse_merge_into_clause_unknown_kind_errors() { + let ctx = SessionContext::new(); + let codec = DefaultLogicalExtensionCodec {}; + let on = serialize_expr(&lit(true), &codec).unwrap(); + let node = dml_node_with_merge_payload(protobuf::MergeIntoOpNode { + on: Some(Box::new(on)), + clauses: vec![protobuf::MergeIntoClauseNode { + kind: 999, // unknown enum tag + predicate: None, + action: Some(protobuf::MergeIntoActionNode { + action: Some(protobuf::merge_into_action_node::Action::Delete( + protobuf::MergeDeleteAction {}, + )), + }), + }], + }); + let err = from_proto::parse_write_op(&node, ctx.task_ctx().as_ref(), &codec) + .expect_err("unknown clause kind tag must fail"); + assert!( + err.to_string().contains("unknown kind tag"), + "unexpected error: {err}" + ); +} + +#[test] +fn parse_merge_into_clause_missing_action_errors() { + let ctx = SessionContext::new(); + let codec = DefaultLogicalExtensionCodec {}; + let on = serialize_expr(&lit(true), &codec).unwrap(); + let node = dml_node_with_merge_payload(protobuf::MergeIntoOpNode { + on: Some(Box::new(on)), + clauses: vec![protobuf::MergeIntoClauseNode { + kind: protobuf::merge_into_clause_node::Kind::Matched.into(), + predicate: None, + action: None, + }], + }); + let err = from_proto::parse_write_op(&node, ctx.task_ctx().as_ref(), &codec) + .expect_err("missing clause `action` must fail"); + assert!( + err.to_string().contains("missing required `action`"), + "unexpected error: {err}" + ); +} + +#[test] +fn parse_merge_into_action_missing_oneof_errors() { + let ctx = SessionContext::new(); + let codec = DefaultLogicalExtensionCodec {}; + let on = serialize_expr(&lit(true), &codec).unwrap(); + let node = dml_node_with_merge_payload(protobuf::MergeIntoOpNode { + on: Some(Box::new(on)), + clauses: vec![protobuf::MergeIntoClauseNode { + kind: protobuf::merge_into_clause_node::Kind::Matched.into(), + predicate: None, + action: Some(protobuf::MergeIntoActionNode { action: None }), + }], + }); + let err = from_proto::parse_write_op(&node, ctx.task_ctx().as_ref(), &codec) + .expect_err("missing action oneof must fail"); + assert!( + err.to_string().contains("missing the `action` oneof"), + "unexpected error: {err}" + ); +} + #[tokio::test] async fn roundtrip_logical_plan_copy_to_sql_options() -> Result<()> { let ctx = SessionContext::new(); @@ -1324,7 +1702,7 @@ pub mod proto { pub expr: Option, } - #[allow(dead_code)] + #[expect(dead_code)] #[derive(Clone, PartialEq, Eq, ::prost::Message)] pub struct TopKExecProto { #[prost(uint64, tag = "1")] @@ -1554,6 +1932,41 @@ impl LogicalExtensionCodec for UDFExtensionCodec { .map_err(|err| internal_datafusion_err!("failed to encode udf: {err}"))?; Ok(()) } + + fn try_decode_higher_order_function( + &self, + name: &str, + buf: &[u8], + ) -> Result> { + if name == "higher_order_udf" { + let proto = MyHigherOrderUdfNode::decode(buf).map_err(|err| { + internal_datafusion_err!("failed to decode higher_order_udf: {err}") + })?; + + Ok(Arc::new(HigherOrderUDF::new_from_impl( + MyHigherOrderUDF::new(proto.payload), + ))) + } else { + not_impl_err!("unrecognized higher order UDF implementation, cannot decode") + } + } + + fn try_encode_higher_order_function( + &self, + node: &HigherOrderUDF, + buf: &mut Vec, + ) -> Result<()> { + let hof = (node.inner().as_ref() as &dyn Any) + .downcast_ref::() + .unwrap(); + let proto = MyHigherOrderUdfNode { + payload: hof.payload.clone(), + }; + proto + .encode(buf) + .map_err(|err| internal_datafusion_err!("failed to encode hof: {err}"))?; + Ok(()) + } } #[test] @@ -1801,7 +2214,7 @@ fn round_trip_scalar_values_and_data_types() { Arc::new(Field::new( "entries", DataType::Struct(Fields::from(vec![ - Field::new("key", DataType::Int32, true), + Field::new("key", DataType::Int32, false), Field::new("value", DataType::Utf8, false), ])), false, @@ -1813,7 +2226,7 @@ fn round_trip_scalar_values_and_data_types() { Arc::new(Field::new( "entries", DataType::Struct(Fields::from(vec![ - Field::new("key", DataType::Int32, true), + Field::new("key", DataType::Int32, false), Field::new("value", DataType::Utf8, true), ])), false, @@ -2102,7 +2515,7 @@ fn roundtrip_null_scalar_values() { for test_case in test_types.into_iter() { let proto_scalar: protobuf::ScalarValue = (&test_case).try_into().unwrap(); let returned_scalar: ScalarValue = (&proto_scalar).try_into().unwrap(); - assert_eq!(format!("{:?}", &test_case), format!("{returned_scalar:?}")); + assert_eq!(format!("{test_case:?}"), format!("{returned_scalar:?}")); } } @@ -2321,6 +2734,18 @@ fn roundtrip_inlist() { fn roundtrip_unnest() { let test_expr = Expr::Unnest(Unnest { expr: Box::new(col("col")), + outer: false, + }); + + let ctx = SessionContext::new(); + roundtrip_expr_test(test_expr, ctx); +} + +#[test] +fn roundtrip_unnest_outer() { + let test_expr = Expr::Unnest(Unnest { + expr: Box::new(col("col")), + outer: true, }); let ctx = SessionContext::new(); @@ -2597,7 +3022,7 @@ fn roundtrip_scalar_udf_extension_codec() { from_proto::parse_expr(&proto, ctx.task_ctx().as_ref(), &UDFExtensionCodec) .expect("parse expr"); - assert_eq!(format!("{:?}", &test_expr), format!("{round_trip:?}")); + assert_eq!(format!("{test_expr:?}"), format!("{round_trip:?}")); roundtrip_json_test(&proto); } @@ -2611,7 +3036,116 @@ fn roundtrip_aggregate_udf_extension_codec() { from_proto::parse_expr(&proto, ctx.task_ctx().as_ref(), &UDFExtensionCodec) .expect("parse expr"); - assert_eq!(format!("{:?}", &test_expr), format!("{round_trip:?}")); + assert_eq!(format!("{test_expr:?}"), format!("{round_trip:?}")); + roundtrip_json_test(&proto); +} + +fn dummy_higher_order_function_args() -> Vec { + let list = ScalarValue::List(ScalarValue::new_list_nullable( + &[ScalarValue::Int32(Some(1))], + &DataType::Int32, + )); + let lambda_var_with_field = Expr::LambdaVariable(LambdaVariable::new( + "x".to_string(), + Some(Arc::new(Field::new("x", DataType::Int32, true))), + )); + let lambda_var_without_field = + Expr::LambdaVariable(LambdaVariable::new("x".into(), None)); + let lambda = lambda(["x"], lambda_var_with_field + lambda_var_without_field); + vec![Expr::Literal(list, None), lambda] +} + +#[test] +fn roundtrip_higher_order_function() { + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + + let test_expr = Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + Arc::clone(&hof), + dummy_higher_order_function_args(), + )); + + let ctx = SessionContext::new(); + ctx.register_higher_order_function(hof); + + roundtrip_expr_test(test_expr.clone(), ctx); + + // Now test loading the HOF without registering it in the context, but rather creating it + // in the extension codec. + #[derive(Debug)] + struct DummyHigherOrderUDFExtensionCodec; + + impl LogicalExtensionCodec for DummyHigherOrderUDFExtensionCodec { + fn try_decode( + &self, + _buf: &[u8], + _inputs: &[LogicalPlan], + _ctx: &TaskContext, + ) -> Result { + not_impl_err!("LogicalExtensionCodec is not provided") + } + + fn try_encode(&self, _node: &Extension, _buf: &mut Vec) -> Result<()> { + not_impl_err!("LogicalExtensionCodec is not provided") + } + + fn try_decode_table_provider( + &self, + _buf: &[u8], + _table_ref: &TableReference, + _schema: SchemaRef, + _ctx: &TaskContext, + ) -> Result> { + not_impl_err!("LogicalExtensionCodec is not provided") + } + + fn try_encode_table_provider( + &self, + _table_ref: &TableReference, + _node: Arc, + _buf: &mut Vec, + ) -> Result<()> { + not_impl_err!("LogicalExtensionCodec is not provided") + } + + fn try_decode_higher_order_function( + &self, + name: &str, + _buf: &[u8], + ) -> Result> { + if name == "higher_order_udf" { + Ok(Arc::new(HigherOrderUDF::new_from_impl( + MyHigherOrderUDF::new("payload".to_string()), + ))) + } else { + Err(internal_datafusion_err!("HOF {name} not found")) + } + } + } + + let ctx = SessionContext::new(); + roundtrip_expr_test_with_codec(test_expr, ctx, &DummyHigherOrderUDFExtensionCodec) +} + +#[test] +fn roundtrip_higher_order_udf_extension_codec() { + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + + let test_expr = Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + hof, + dummy_higher_order_function_args(), + )); + + let ctx = SessionContext::new(); + let proto = serialize_expr(&test_expr, &UDFExtensionCodec).expect("serialize expr"); + let round_trip = + from_proto::parse_expr(&proto, ctx.task_ctx().as_ref(), &UDFExtensionCodec) + .expect("parse expr"); + + assert_eq!(format!("{test_expr:?}"), format!("{round_trip:?}")); roundtrip_json_test(&proto); } @@ -3128,9 +3662,7 @@ async fn roundtrip_empty_table_scan() -> Result<()> { Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); - let table = Arc::new(datafusion::datasource::empty::EmptyTable::new(Arc::clone( - &schema, - ))); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); let ctx = SessionContext::new(); ctx.register_table("empty", table)?; @@ -3152,9 +3684,7 @@ async fn roundtrip_empty_table_scan_with_projection() -> Result<()> { Field::new("id", DataType::Int32, false), Field::new("name", DataType::Utf8, true), ])); - let table = Arc::new(datafusion::datasource::empty::EmptyTable::new(Arc::clone( - &schema, - ))); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); let ctx = SessionContext::new(); ctx.register_table("empty", table)?; @@ -3173,3 +3703,170 @@ async fn roundtrip_empty_table_scan_with_projection() -> Result<()> { ); Ok(()) } + +// Regression test for https://github.com/apache/datafusion/issues/22065: +// the decoder must preserve `null_aware = true` (NOT IN semantics) +// across a to_proto -> from_proto round trip. `null_equality` is at +// its default (`NullEqualsNothing`). +#[tokio::test] +async fn roundtrip_join_null_aware() -> Result<()> { + use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion_expr::JoinType; + + let ctx = SessionContext::new(); + let sql = " + SELECT id + FROM (VALUES (1), (2), (3)) AS t1(id) + WHERE id NOT IN ( + SELECT bad_id + FROM (VALUES (CAST(1 AS INT)), (CAST(NULL AS INT))) AS excludes(bad_id) + ) + "; + + let df = ctx.sql(sql).await?; + let plan = ctx.state().optimize(df.logical_plan())?; + + let mut found_null_aware = false; + plan.apply(|n| { + if let LogicalPlan::Join(j) = n + && j.join_type == JoinType::LeftAnti + && j.null_aware + { + found_null_aware = true; + } + Ok(TreeNodeRecursion::Continue) + })?; + assert!(found_null_aware); + + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + + Ok(()) +} + +// Regression test for `null_equality` round-trip (related to #22065): +// the decoder must preserve a non-default `null_equality` +// (`NullEqualsNull`) across a to_proto -> from_proto round trip. +// `null_aware` is at its default (`false`). +#[tokio::test] +async fn roundtrip_join_null_equality() -> Result<()> { + use datafusion_common::NullEquality; + use datafusion_expr::JoinType; + use datafusion_expr::logical_plan::{Join, JoinConstraint}; + + let ctx = SessionContext::new(); + + let left_schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int32, true)])); + let right_schema = + Arc::new(Schema::new(vec![Field::new("b", DataType::Int32, true)])); + ctx.register_table("t1", Arc::new(EmptyTable::new(left_schema)))?; + ctx.register_table("t2", Arc::new(EmptyTable::new(right_schema)))?; + let left = ctx.table("t1").await?.into_optimized_plan()?; + let right = ctx.table("t2").await?.into_optimized_plan()?; + + let join = LogicalPlan::Join(Join::try_new( + Arc::new(left), + Arc::new(right), + vec![(col("t1.a"), col("t2.b"))], + None, + JoinType::Inner, + JoinConstraint::On, + NullEquality::NullEqualsNull, + false, + )?); + + let bytes = logical_plan_to_bytes(&join)?; + let rt = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{join:?}"), format!("{rt:?}")); + + Ok(()) +} + +// Single column, single split point range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_single_col() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("id", DataType::Int32, false), + Field::new("name", DataType::Utf8, true), + ])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("id").sort(true, true)], + vec![SplitPoint::new(vec![ScalarValue::Int32(Some(10))])], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} + +// Multi-column compound key with multiple split points for range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_multi_col() -> Result<()> { + let schema = Arc::new(Schema::new(vec![ + Field::new("ts", DataType::Int64, false), + Field::new("region", DataType::Utf8, false), + ])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("ts").sort(true, true), col("region").sort(true, true)], + vec![ + SplitPoint::new(vec![ + ScalarValue::Int64(Some(1000)), + ScalarValue::Utf8(Some("east".to_string())), + ]), + SplitPoint::new(vec![ + ScalarValue::Int64(Some(2000)), + ScalarValue::Utf8(Some("west".to_string())), + ]), + ], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} + +// Non-default sort options: descending with nulls last for range partitioning +#[tokio::test] +async fn roundtrip_range_partitioning_desc_nulls_last() -> Result<()> { + let schema = Arc::new(Schema::new(vec![Field::new( + "score", + DataType::Float64, + true, + )])); + let table = Arc::new(EmptyTable::new(Arc::clone(&schema))); + + let ctx = SessionContext::new(); + ctx.register_table("empty", table)?; + + let scan_plan = ctx.table("empty").await?.into_optimized_plan()?; + let plan = LogicalPlan::Repartition(Repartition { + input: Arc::new(scan_plan), + partitioning_scheme: Partitioning::Range(RangePartitioning::try_new( + vec![col("score").sort(false, false)], + vec![SplitPoint::new(vec![ScalarValue::Float64(Some(50.0))])], + )?), + }); + let bytes = logical_plan_to_bytes(&plan)?; + let logical_round_trip = logical_plan_from_bytes(&bytes, &ctx.task_ctx())?; + assert_eq!(format!("{plan:?}"), format!("{logical_round_trip:?}")); + Ok(()) +} diff --git a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs b/datafusion/proto/tests/cases/roundtrip_physical_plan.rs deleted file mode 100644 index 0cb6068af3b29..0000000000000 --- a/datafusion/proto/tests/cases/roundtrip_physical_plan.rs +++ /dev/null @@ -1,3874 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::HashMap; -use std::fmt::{Display, Formatter}; -use std::sync::{Arc, RwLock}; -use std::vec; - -use arrow::array::RecordBatch; -use arrow::csv::WriterBuilder; -use arrow::datatypes::{Fields, TimeUnit}; -use datafusion::arrow::array::ArrayRef; -use datafusion::arrow::compute::kernels::sort::SortOptions; -use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema}; -use datafusion::datasource::empty::EmptyTable; -use datafusion::datasource::file_format::csv::CsvSink; -use datafusion::datasource::file_format::json::{JsonFormat, JsonSink}; -use datafusion::datasource::file_format::parquet::ParquetSink; -use datafusion::datasource::listing::{ - ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, PartitionedFile, -}; -use datafusion::datasource::object_store::ObjectStoreUrl; -use datafusion::datasource::physical_plan::{ - ArrowSource, FileGroup, FileOutputMode, FileScanConfig, FileScanConfigBuilder, - FileSinkConfig, ParquetSource, wrap_partition_type_in_dict, - wrap_partition_value_in_dict, -}; -use datafusion::datasource::sink::DataSinkExec; -use datafusion::datasource::source::DataSourceExec; -use datafusion::execution::TaskContext; -use datafusion::functions_aggregate::count::count_udaf; -use datafusion::functions_aggregate::first_last::first_value_udaf; -use datafusion::functions_aggregate::sum::sum_udaf; -use datafusion::functions_window::nth_value::nth_value_udwf; -use datafusion::functions_window::row_number::row_number_udwf; -use datafusion::logical_expr::{JoinType, Operator, Volatility, create_udf}; -use datafusion::physical_expr::aggregate::{AggregateExprBuilder, AggregateFunctionExpr}; -use datafusion::physical_expr::expressions::Literal; -use datafusion::physical_expr::window::{SlidingAggregateWindowExpr, StandardWindowExpr}; -use datafusion::physical_expr::{ - LexOrdering, PhysicalSortRequirement, ScalarFunctionExpr, -}; -use datafusion::physical_optimizer::PhysicalOptimizerRule; -use datafusion::physical_optimizer::filter_pushdown::FilterPushdown; -use datafusion::physical_plan::aggregates::{ - AggregateExec, AggregateMode, LimitOptions, PhysicalGroupBy, -}; -use datafusion::physical_plan::analyze::AnalyzeExec; -#[expect(deprecated)] -use datafusion::physical_plan::coalesce_batches::CoalesceBatchesExec; -use datafusion::physical_plan::coalesce_partitions::CoalescePartitionsExec; -use datafusion::physical_plan::empty::EmptyExec; -use datafusion::physical_plan::expressions::{ - BinaryExpr, Column, NotExpr, PhysicalSortExpr, binary, cast, col, in_list, like, lit, -}; -use datafusion::physical_plan::filter::{FilterExec, FilterExecBuilder}; -use datafusion::physical_plan::joins::{ - HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec, - StreamJoinPartitionMode, SymmetricHashJoinExec, -}; -use datafusion::physical_plan::limit::{GlobalLimitExec, LocalLimitExec}; -use datafusion::physical_plan::metrics::MetricType; -use datafusion::physical_plan::placeholder_row::PlaceholderRowExec; -use datafusion::physical_plan::projection::{ProjectionExec, ProjectionExpr}; -use datafusion::physical_plan::repartition::RepartitionExec; -use datafusion::physical_plan::scalar_subquery::{ - ScalarSubqueryExec, ScalarSubqueryLink, -}; -use datafusion::physical_plan::sorts::sort::SortExec; -use datafusion::physical_plan::union::{InterleaveExec, UnionExec}; -use datafusion::physical_plan::unnest::{ListUnnest, UnnestExec}; -use datafusion::physical_plan::windows::{ - BoundedWindowAggExec, PlainAggregateWindowExpr, WindowAggExec, - create_udwf_window_expr, -}; -use datafusion::physical_plan::{ - ExecutionPlan, InputOrderMode, Partitioning, PhysicalExpr, Statistics, displayable, -}; -use datafusion::prelude::{ParquetReadOptions, SessionContext}; -use datafusion::scalar::ScalarValue; -use datafusion_common::config::{ConfigOptions, TableParquetOptions}; -use datafusion_common::file_options::csv_writer::CsvWriterOptions; -use datafusion_common::file_options::json_writer::JsonWriterOptions; -use datafusion_common::parsers::CompressionTypeVariant; -use datafusion_common::stats::Precision; -use datafusion_common::{ - DataFusionError, NullEquality, Result, UnnestOptions, exec_datafusion_err, - internal_datafusion_err, internal_err, not_impl_err, -}; -use datafusion_datasource::TableSchema; -use datafusion_datasource::file::FileSource; -use datafusion_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; -use datafusion_expr::dml::InsertOp; -use datafusion_expr::{ - Accumulator, AccumulatorFactoryFunction, AggregateUDF, ColumnarValue, - ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, SimpleAggregateUDF, - WindowFrame, WindowFrameBound, WindowUDF, - execution_props::{ScalarSubqueryResults, SubqueryIndex}, -}; -use datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf; -use datafusion_functions_aggregate::array_agg::array_agg_udaf; -use datafusion_functions_aggregate::average::avg_udaf; -use datafusion_functions_aggregate::min_max::max_udaf; -use datafusion_functions_aggregate::nth_value::nth_value_udaf; -use datafusion_functions_aggregate::string_agg::string_agg_udaf; -use datafusion_physical_expr::scalar_subquery::ScalarSubqueryExpr; -use datafusion_proto::bytes::{ - physical_plan_from_bytes_with_proto_converter, - physical_plan_to_bytes_with_proto_converter, -}; -use datafusion_proto::physical_plan::to_proto::serialize_physical_expr_with_converter; -use datafusion_proto::physical_plan::{ - AsExecutionPlan, DeduplicatingProtoConverter, DefaultPhysicalExtensionCodec, - DefaultPhysicalProtoConverter, PhysicalExtensionCodec, PhysicalPlanDecodeContext, - PhysicalProtoConverterExtension, -}; -use datafusion_proto::protobuf; -use datafusion_proto::protobuf::{PhysicalExprNode, PhysicalPlanNode}; -use prost::Message; - -use crate::cases::{ - CustomUDWF, CustomUDWFNode, MyAggregateUDF, MyAggregateUdfNode, MyRegexUdf, - MyRegexUdfNode, -}; -use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; -use datafusion_physical_expr::utils::reassign_expr_columns; - -/// Perform a serde roundtrip and assert that the string representation of the before and after plans -/// are identical. Note that this often isn't sufficient to guarantee that no information is -/// lost during serde because the string representation of a plan often only shows a subset of state. -fn roundtrip_test(exec_plan: Arc) -> Result<()> { - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; - Ok(()) -} - -/// Perform a serde roundtrip and assert that the string representation of the before and after plans -/// are identical. Note that this often isn't sufficient to guarantee that no information is -/// lost during serde because the string representation of a plan often only shows a subset of state. -/// -/// This version of the roundtrip_test method returns the final plan after serde so that it can be inspected -/// farther in tests. -fn roundtrip_test_and_return( - exec_plan: Arc, - ctx: &SessionContext, - codec: &dyn PhysicalExtensionCodec, - proto_converter: &dyn PhysicalProtoConverterExtension, -) -> Result> { - let bytes = physical_plan_to_bytes_with_proto_converter( - Arc::clone(&exec_plan), - codec, - proto_converter, - )?; - let result_exec_plan = physical_plan_from_bytes_with_proto_converter( - bytes.as_ref(), - ctx.task_ctx().as_ref(), - codec, - proto_converter, - )?; - - pretty_assertions::assert_eq!( - format!("{exec_plan:?}"), - format!("{result_exec_plan:?}") - ); - Ok(result_exec_plan) -} - -/// Perform a serde roundtrip and assert that the string representation of the before and after plans -/// are identical. Note that this often isn't sufficient to guarantee that no information is -/// lost during serde because the string representation of a plan often only shows a subset of state. -/// -/// This version of the roundtrip_test function accepts a SessionContext, which is required when -/// performing serde on some plans. -fn roundtrip_test_with_context( - exec_plan: Arc, - ctx: &SessionContext, -) -> Result<()> { - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - roundtrip_test_and_return(exec_plan, ctx, &codec, &proto_converter)?; - Ok(()) -} - -/// Perform a serde roundtrip for the specified sql query, and assert that -/// query results are identical. -async fn roundtrip_test_sql_with_context(sql: &str, ctx: &SessionContext) -> Result<()> { - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - let initial_plan = ctx.sql(sql).await?.create_physical_plan().await?; - - roundtrip_test_and_return(initial_plan, ctx, &codec, &proto_converter)?; - Ok(()) -} - -/// returns a SessionContext with `alltypes_plain` registered -async fn all_types_context() -> Result { - let ctx = SessionContext::new(); - - let testdata = datafusion::test_util::parquet_test_data(); - ctx.register_parquet( - "alltypes_plain", - &format!("{testdata}/alltypes_plain.parquet"), - ParquetReadOptions::default(), - ) - .await?; - - Ok(ctx) -} - -#[test] -fn roundtrip_empty() -> Result<()> { - roundtrip_test(Arc::new(EmptyExec::new(Arc::new(Schema::empty())))) -} - -#[test] -fn roundtrip_date_time_interval() -> Result<()> { - let schema = Schema::new(vec![ - Field::new("some_date", DataType::Date32, false), - Field::new( - "some_interval", - DataType::Interval(IntervalUnit::DayTime), - false, - ), - ]); - let input = Arc::new(EmptyExec::new(Arc::new(schema.clone()))); - let date_expr = col("some_date", &schema)?; - let literal_expr = col("some_interval", &schema)?; - let date_time_interval_expr = - binary(date_expr, Operator::Plus, literal_expr, &schema)?; - let plan = Arc::new(ProjectionExec::try_new( - vec![ProjectionExpr { - expr: date_time_interval_expr, - alias: "result".to_string(), - }], - input, - )?); - roundtrip_test(plan) -} - -#[test] -fn roundtrip_local_limit() -> Result<()> { - roundtrip_test(Arc::new(LocalLimitExec::new( - Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), - 25, - ))) -} - -#[test] -fn roundtrip_global_limit() -> Result<()> { - roundtrip_test(Arc::new(GlobalLimitExec::new( - Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), - 0, - Some(25), - ))) -} - -#[test] -fn roundtrip_global_skip_no_limit() -> Result<()> { - roundtrip_test(Arc::new(GlobalLimitExec::new( - Arc::new(EmptyExec::new(Arc::new(Schema::empty()))), - 10, - None, // no limit - ))) -} - -#[test] -fn roundtrip_hash_join() -> Result<()> { - let field_a = Field::new("col", DataType::Int64, false); - let schema_left = Schema::new(vec![field_a.clone()]); - let schema_right = Schema::new(vec![field_a]); - let on = vec![( - Arc::new(Column::new("col", schema_left.index_of("col")?)) as _, - Arc::new(Column::new("col", schema_right.index_of("col")?)) as _, - )]; - - let schema_left = Arc::new(schema_left); - let schema_right = Arc::new(schema_right); - for join_type in &[ - JoinType::Inner, - JoinType::Left, - JoinType::Right, - JoinType::Full, - JoinType::LeftAnti, - JoinType::RightAnti, - JoinType::LeftSemi, - JoinType::RightSemi, - ] { - for partition_mode in &[PartitionMode::Partitioned, PartitionMode::CollectLeft] { - roundtrip_test(Arc::new(HashJoinExec::try_new( - Arc::new(EmptyExec::new(schema_left.clone())), - Arc::new(EmptyExec::new(schema_right.clone())), - on.clone(), - None, - join_type, - None, - *partition_mode, - NullEquality::NullEqualsNothing, - false, - )?))?; - } - } - Ok(()) -} - -#[test] -fn roundtrip_nested_loop_join() -> Result<()> { - let field_a = Field::new("col", DataType::Int64, false); - let schema_left = Schema::new(vec![field_a.clone()]); - let schema_right = Schema::new(vec![field_a]); - - let schema_left = Arc::new(schema_left); - let schema_right = Arc::new(schema_right); - for join_type in &[ - JoinType::Inner, - JoinType::Left, - JoinType::Right, - JoinType::Full, - JoinType::LeftAnti, - JoinType::RightAnti, - JoinType::LeftSemi, - JoinType::RightSemi, - ] { - roundtrip_test(Arc::new(NestedLoopJoinExec::try_new( - Arc::new(EmptyExec::new(schema_left.clone())), - Arc::new(EmptyExec::new(schema_right.clone())), - None, - join_type, - Some(vec![0]), - )?))?; - } - Ok(()) -} - -#[test] -fn roundtrip_udwf() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let udwf_expr = Arc::new(StandardWindowExpr::new( - create_udwf_window_expr( - &row_number_udwf(), - &[], - &schema, - "row_number() PARTITION BY [a] ORDER BY [b] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW".to_string(), - false, - )?, - &[ - col("a", &schema)? - ], - &[ - PhysicalSortExpr::new(col("b", &schema)?, SortOptions::new(true, true)) - ], - Arc::new(WindowFrame::new(None)), - )); - - let input = Arc::new(EmptyExec::new(schema.clone())); - - roundtrip_test(Arc::new(BoundedWindowAggExec::try_new( - vec![udwf_expr], - input, - InputOrderMode::Sorted, - true, - )?)) -} - -#[test] -fn roundtrip_window() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let window_frame = WindowFrame::new_bounds( - datafusion_expr::WindowFrameUnits::Range, - WindowFrameBound::Preceding(ScalarValue::Int64(None)), - WindowFrameBound::CurrentRow, - ); - - let nth_value_window = - create_udwf_window_expr( - &nth_value_udwf(), - &[col("a", &schema)?, - lit(2)], schema.as_ref(), - "NTH_VALUE(a, 2) PARTITION BY [b] ORDER BY [a ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW".to_string(), - false, - )?; - let udwf_expr = Arc::new(StandardWindowExpr::new( - nth_value_window, - &[col("b", &schema)?], - &[PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions { - descending: false, - nulls_first: false, - }, - }], - Arc::new(window_frame), - )); - - let plain_aggr_window_expr = Arc::new(PlainAggregateWindowExpr::new( - AggregateExprBuilder::new( - avg_udaf(), - vec![cast(col("b", &schema)?, &schema, DataType::Float64)?], - ) - .schema(Arc::clone(&schema)) - .alias("avg(b)") - .build() - .map(Arc::new)?, - &[], - &[], - Arc::new(WindowFrame::new(None)), - None, - )); - - let window_frame = WindowFrame::new_bounds( - datafusion_expr::WindowFrameUnits::Range, - WindowFrameBound::CurrentRow, - WindowFrameBound::Preceding(ScalarValue::Int64(None)), - ); - - let args = vec![cast(col("a", &schema)?, &schema, DataType::Float64)?]; - let sum_expr = AggregateExprBuilder::new(sum_udaf(), args) - .schema(Arc::clone(&schema)) - .alias("SUM(a) RANGE BETWEEN CURRENT ROW AND UNBOUNDED PRECEDING") - .build() - .map(Arc::new)?; - - let sliding_aggr_window_expr = Arc::new(SlidingAggregateWindowExpr::new( - sum_expr, - &[], - &[], - Arc::new(window_frame), - None, - )); - - let input = Arc::new(EmptyExec::new(schema.clone())); - - roundtrip_test(Arc::new(WindowAggExec::try_new( - vec![plain_aggr_window_expr, sliding_aggr_window_expr, udwf_expr], - input, - false, - )?)) -} - -#[test] -fn roundtrip_window_distinct() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - // Create a distinct count window expression with unbounded frame (becomes PlainAggregateWindowExpr) - let distinct_count_expr = Arc::new(PlainAggregateWindowExpr::new( - AggregateExprBuilder::new(count_udaf(), vec![col("a", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("count(DISTINCT a)") - .distinct() // Enable distinct - .build() - .map(Arc::new)?, - &[col("b", &schema)?], // partition by b - &[], // no order by - Arc::new(WindowFrame::new(None)), // unbounded frame - None, - )); - - // Create a distinct sum window expression with bounded frame (becomes SlidingAggregateWindowExpr) - let bounded_frame = WindowFrame::new_bounds( - datafusion_expr::WindowFrameUnits::Rows, - WindowFrameBound::Preceding(ScalarValue::UInt64(Some(1))), - WindowFrameBound::CurrentRow, - ); - - let distinct_sum_expr = Arc::new(SlidingAggregateWindowExpr::new( - AggregateExprBuilder::new( - sum_udaf(), - vec![cast(col("a", &schema)?, &schema, DataType::Float64)?], - ) - .schema(Arc::clone(&schema)) - .alias("sum(DISTINCT a)") - .distinct() // Enable distinct - .with_ignore_nulls(true) // Enable ignore nulls - .build() - .map(Arc::new)?, - &[], // no partition by - &[], // no order by - Arc::new(bounded_frame), // bounded frame - None, - )); - - let input = Arc::new(EmptyExec::new(schema.clone())); - - roundtrip_test(Arc::new(WindowAggExec::try_new( - vec![distinct_count_expr, distinct_sum_expr], - input, - false, - )?)) -} - -#[test] -fn test_distinct_window_serialization_end_to_end() -> Result<()> { - // Create a more comprehensive test that verifies distinct window functions - // work properly through the entire serialization/deserialization pipeline - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - // Test 1: DISTINCT COUNT with IGNORE NULLS - let distinct_count_ignore_nulls = Arc::new(PlainAggregateWindowExpr::new( - AggregateExprBuilder::new(count_udaf(), vec![col("a", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("count_distinct_ignore_nulls") - .distinct() - .with_ignore_nulls(true) - .build() - .map(Arc::new)?, - &[col("b", &schema)?], - &[], - Arc::new(WindowFrame::new(None)), - None, - )); - - // Test 2: DISTINCT SUM (without ignore nulls) - let bounded_frame = WindowFrame::new_bounds( - datafusion_expr::WindowFrameUnits::Rows, - WindowFrameBound::Preceding(ScalarValue::UInt64(Some(2))), - WindowFrameBound::CurrentRow, - ); - - let distinct_sum = Arc::new(SlidingAggregateWindowExpr::new( - AggregateExprBuilder::new( - sum_udaf(), - vec![cast(col("a", &schema)?, &schema, DataType::Float64)?], - ) - .schema(Arc::clone(&schema)) - .alias("sum_distinct") - .distinct() - .build() - .map(Arc::new)?, - &[], - &[], - Arc::new(bounded_frame), - None, - )); - - let input = Arc::new(EmptyExec::new(schema.clone())); - - let window_exec = Arc::new(WindowAggExec::try_new( - vec![distinct_count_ignore_nulls, distinct_sum], - input, - false, - )?); - - // Perform the roundtrip test - roundtrip_test(window_exec) -} - -#[test] -fn roundtrip_aggregate() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let groups: Vec<(Arc, String)> = - vec![(col("a", &schema)?, "unused".to_string())]; - - let avg_expr = AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("AVG(b)") - .build()?; - let nth_expr = - AggregateExprBuilder::new(nth_value_udaf(), vec![col("b", &schema)?, lit(1u64)]) - .schema(Arc::clone(&schema)) - .alias("NTH_VALUE(b, 1)") - .build()?; - let str_agg_expr = - AggregateExprBuilder::new(string_agg_udaf(), vec![col("b", &schema)?, lit(1u64)]) - .schema(Arc::clone(&schema)) - .alias("NTH_VALUE(b, 1)") - .build()?; - - let test_cases = vec![ - // AVG - vec![Arc::new(avg_expr)], - // NTH_VALUE - vec![Arc::new(nth_expr)], - // STRING_AGG - vec![Arc::new(str_agg_expr)], - ]; - - for aggregates in test_cases { - let schema = schema.clone(); - roundtrip_test(Arc::new(AggregateExec::try_new( - AggregateMode::Final, - PhysicalGroupBy::new_single(groups.clone()), - aggregates, - vec![None], - Arc::new(EmptyExec::new(schema.clone())), - schema, - )?))?; - } - - Ok(()) -} - -#[test] -fn roundtrip_aggregate_with_limit() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let groups: Vec<(Arc, String)> = - vec![(col("a", &schema)?, "unused".to_string())]; - - let aggregates = vec![ - AggregateExprBuilder::new(avg_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("AVG(b)") - .build() - .map(Arc::new)?, - ]; - - let agg = AggregateExec::try_new( - AggregateMode::Final, - PhysicalGroupBy::new_single(groups.clone()), - aggregates, - vec![None], - Arc::new(EmptyExec::new(schema.clone())), - schema, - )?; - let agg = agg.with_limit_options(Some(LimitOptions::new_with_order(12, false))); - roundtrip_test(Arc::new(agg)) -} - -#[test] -fn roundtrip_aggregate_with_approx_pencentile_cont() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let groups: Vec<(Arc, String)> = - vec![(col("a", &schema)?, "unused".to_string())]; - - let aggregates = vec![ - AggregateExprBuilder::new( - approx_percentile_cont_udaf(), - vec![col("b", &schema)?, lit(0.5)], - ) - .schema(Arc::clone(&schema)) - .alias("APPROX_PERCENTILE_CONT(0.5) WITHIN GROUP (ORDER BY b)") - .build() - .map(Arc::new)?, - ]; - - let agg = AggregateExec::try_new( - AggregateMode::Final, - PhysicalGroupBy::new_single(groups.clone()), - aggregates, - vec![None], - Arc::new(EmptyExec::new(schema.clone())), - schema, - )?; - roundtrip_test(Arc::new(agg)) -} - -#[test] -fn roundtrip_aggregate_with_sort() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let groups: Vec<(Arc, String)> = - vec![(col("a", &schema)?, "unused".to_string())]; - let sort_exprs = vec![PhysicalSortExpr { - expr: col("b", &schema)?, - options: SortOptions { - descending: false, - nulls_first: true, - }, - }]; - - let aggregates = vec![ - AggregateExprBuilder::new(array_agg_udaf(), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("ARRAY_AGG(b)") - .order_by(sort_exprs) - .build() - .map(Arc::new)?, - ]; - - let agg = AggregateExec::try_new( - AggregateMode::Final, - PhysicalGroupBy::new_single(groups.clone()), - aggregates, - vec![None], - Arc::new(EmptyExec::new(schema.clone())), - schema, - )?; - roundtrip_test(Arc::new(agg)) -} - -#[test] -fn roundtrip_aggregate_udaf() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - #[derive(Debug)] - struct Example; - impl Accumulator for Example { - fn state(&mut self) -> Result> { - Ok(vec![ScalarValue::Int64(Some(0))]) - } - - fn update_batch(&mut self, _values: &[ArrayRef]) -> Result<()> { - Ok(()) - } - - fn merge_batch(&mut self, _states: &[ArrayRef]) -> Result<()> { - Ok(()) - } - - fn evaluate(&mut self) -> Result { - Ok(ScalarValue::Int64(Some(0))) - } - - fn size(&self) -> usize { - 0 - } - } - - let return_type = DataType::Int64; - let accumulator: AccumulatorFactoryFunction = Arc::new(|_| Ok(Box::new(Example))); - - let udaf = AggregateUDF::from(SimpleAggregateUDF::new_with_signature( - "example", - Signature::exact(vec![DataType::Int64], Volatility::Immutable), - return_type, - accumulator, - vec![Field::new("value", DataType::Int64, true).into()], - )); - - let ctx = SessionContext::new(); - ctx.register_udaf(udaf.clone()); - - let groups: Vec<(Arc, String)> = - vec![(col("a", &schema)?, "unused".to_string())]; - - let aggregates = vec![ - AggregateExprBuilder::new(Arc::new(udaf), vec![col("b", &schema)?]) - .schema(Arc::clone(&schema)) - .alias("example_agg") - .build() - .map(Arc::new)?, - ]; - - roundtrip_test_with_context( - Arc::new(AggregateExec::try_new( - AggregateMode::Final, - PhysicalGroupBy::new_single(groups.clone()), - aggregates, - vec![None], - Arc::new(EmptyExec::new(schema.clone())), - schema, - )?), - &ctx, - ) -} - -#[test] -fn roundtrip_filter_with_not_and_in_list() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let field_c = Field::new("c", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c])); - let not = Arc::new(NotExpr::new(col("a", &schema)?)); - let in_list = in_list( - col("b", &schema)?, - vec![ - lit(ScalarValue::Int64(Some(1))), - lit(ScalarValue::Int64(Some(2))), - ], - &false, - schema.as_ref(), - )?; - let and = binary(not, Operator::And, in_list, &schema)?; - roundtrip_test(Arc::new(FilterExec::try_new( - and, - Arc::new(EmptyExec::new(schema.clone())), - )?)) -} - -#[test] -fn roundtrip_filter_with_fetch() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let predicate = col("a", &schema)?; - let filter = FilterExecBuilder::new(predicate, Arc::new(EmptyExec::new(schema))) - .with_fetch(Some(10)) - .build()?; - assert_eq!(filter.fetch(), Some(10)); - roundtrip_test(Arc::new(filter)) -} - -#[test] -fn roundtrip_sort() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let sort_exprs = [ - PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions { - descending: true, - nulls_first: false, - }, - }, - PhysicalSortExpr { - expr: col("b", &schema)?, - options: SortOptions { - descending: false, - nulls_first: true, - }, - }, - ] - .into(); - roundtrip_test(Arc::new(SortExec::new( - sort_exprs, - Arc::new(EmptyExec::new(schema)), - ))) -} - -#[test] -fn roundtrip_sort_preserve_partitioning() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let sort_exprs: LexOrdering = [ - PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions { - descending: true, - nulls_first: false, - }, - }, - PhysicalSortExpr { - expr: col("b", &schema)?, - options: SortOptions { - descending: false, - nulls_first: true, - }, - }, - ] - .into(); - - roundtrip_test(Arc::new(SortExec::new( - sort_exprs.clone(), - Arc::new(EmptyExec::new(schema.clone())), - )))?; - - roundtrip_test(Arc::new( - SortExec::new(sort_exprs, Arc::new(EmptyExec::new(schema))) - .with_preserve_partitioning(true), - )) -} - -#[test] -fn roundtrip_coalesce_batches_with_fetch() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - #[expect(deprecated)] - roundtrip_test(Arc::new(CoalesceBatchesExec::new( - Arc::new(EmptyExec::new(schema.clone())), - 8096, - )))?; - - #[expect(deprecated)] - roundtrip_test(Arc::new( - CoalesceBatchesExec::new(Arc::new(EmptyExec::new(schema)), 8096) - .with_fetch(Some(10)), - )) -} - -#[test] -fn roundtrip_coalesce_partitions_with_fetch() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - roundtrip_test(Arc::new(CoalescePartitionsExec::new(Arc::new( - EmptyExec::new(schema.clone()), - ))))?; - - roundtrip_test(Arc::new( - CoalescePartitionsExec::new(Arc::new(EmptyExec::new(schema))) - .with_fetch(Some(10)), - )) -} - -#[test] -fn roundtrip_parquet_exec_with_pruning_predicate() -> Result<()> { - let file_schema = - Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); - - let predicate = Arc::new(BinaryExpr::new( - Arc::new(Column::new("col", 1)), - Operator::Eq, - lit("1"), - )); - - let mut options = TableParquetOptions::new(); - options.global.pushdown_filters = true; - - let file_source = Arc::new( - ParquetSource::new(Arc::clone(&file_schema)) - .with_table_parquet_options(options) - .with_predicate(predicate), - ); - - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) - .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )])]) - .with_statistics(Statistics { - num_rows: Precision::Inexact(100), - total_byte_size: Precision::Inexact(1024), - column_statistics: Statistics::unknown_column(&Arc::new(Schema::new( - vec![Field::new("col", DataType::Utf8, false)], - ))), - }) - .build(); - - roundtrip_test(DataSourceExec::from_data_source(scan_config)) -} - -#[test] -fn roundtrip_parquet_exec_attaches_cached_reader_factory_after_roundtrip() -> Result<()> { - let file_schema = - Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); - let file_source = Arc::new(ParquetSource::new(Arc::clone(&file_schema))); - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) - .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )])]) - .with_statistics(Statistics { - num_rows: Precision::Inexact(100), - total_byte_size: Precision::Inexact(1024), - column_statistics: Statistics::unknown_column(&file_schema), - }) - .build(); - let exec_plan = DataSourceExec::from_data_source(scan_config); - - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - let roundtripped = - roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; - - let data_source = roundtripped - .downcast_ref::() - .ok_or_else(|| { - internal_datafusion_err!("Expected DataSourceExec after roundtrip") - })?; - let file_scan = data_source - .data_source() - .downcast_ref::() - .ok_or_else(|| { - internal_datafusion_err!("Expected FileScanConfig after roundtrip") - })?; - let parquet_source = file_scan - .file_source() - .downcast_ref::() - .ok_or_else(|| { - internal_datafusion_err!("Expected ParquetSource after roundtrip") - })?; - - assert!( - parquet_source.parquet_file_reader_factory().is_some(), - "Parquet reader factory should be attached after decoding from protobuf" - ); - Ok(()) -} - -#[test] -fn roundtrip_arrow_scan() -> Result<()> { - let file_schema = - Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); - - let table_schema = TableSchema::new(file_schema.clone(), vec![]); - let file_source = Arc::new(ArrowSource::new_file_source(table_schema)); - - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) - .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( - "/path/to/file.arrow".to_string(), - 1024, - )])]) - .with_statistics(Statistics { - num_rows: Precision::Inexact(100), - total_byte_size: Precision::Inexact(1024), - column_statistics: Statistics::unknown_column(&file_schema), - }) - .build(); - - roundtrip_test(DataSourceExec::from_data_source(scan_config)) -} - -#[tokio::test] -async fn roundtrip_parquet_exec_with_table_partition_cols() -> Result<()> { - let mut file_group = - PartitionedFile::new("/path/to/part=0/file.parquet".to_string(), 1024); - file_group.partition_values = - vec![wrap_partition_value_in_dict(ScalarValue::Int64(Some(0)))]; - let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); - - let table_schema = TableSchema::new( - schema.clone(), - vec![Arc::new(Field::new( - "part".to_string(), - wrap_partition_type_in_dict(DataType::Int16), - false, - ))], - ); - - let file_source = Arc::new(ParquetSource::new(table_schema.clone())); - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) - .with_projection_indices(Some(vec![0, 1]))? - .with_file_group(FileGroup::new(vec![file_group])) - .build(); - - roundtrip_test(DataSourceExec::from_data_source(scan_config)) -} - -#[test] -fn roundtrip_parquet_exec_with_custom_predicate_expr() -> Result<()> { - let file_schema = - Arc::new(Schema::new(vec![Field::new("col", DataType::Utf8, false)])); - - let custom_predicate_expr = Arc::new(CustomPredicateExpr { - inner: Arc::new(Column::new("col", 1)), - }); - - let file_source = Arc::new( - ParquetSource::new(Arc::clone(&file_schema)) - .with_predicate(custom_predicate_expr), - ); - - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) - .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )])]) - .with_statistics(Statistics { - num_rows: Precision::Inexact(100), - total_byte_size: Precision::Inexact(1024), - column_statistics: Statistics::unknown_column(&Arc::new(Schema::new( - vec![Field::new("col", DataType::Utf8, false)], - ))), - }) - .build(); - - #[derive(Debug, Clone, Eq)] - struct CustomPredicateExpr { - inner: Arc, - } - - // Manually derive PartialEq and Hash to work around https://github.com/rust-lang/rust/issues/78808 - impl PartialEq for CustomPredicateExpr { - fn eq(&self, other: &Self) -> bool { - self.inner.eq(&other.inner) - } - } - - impl std::hash::Hash for CustomPredicateExpr { - fn hash(&self, state: &mut H) { - self.inner.hash(state); - } - } - - impl Display for CustomPredicateExpr { - fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - write!(f, "CustomPredicateExpr") - } - } - - impl PhysicalExpr for CustomPredicateExpr { - fn data_type(&self, _input_schema: &Schema) -> Result { - unreachable!() - } - - fn nullable(&self, _input_schema: &Schema) -> Result { - unreachable!() - } - - fn evaluate(&self, _batch: &RecordBatch) -> Result { - unreachable!() - } - - fn children(&self) -> Vec<&Arc> { - vec![&self.inner] - } - - fn with_new_children( - self: Arc, - _children: Vec>, - ) -> Result> { - Ok(self) - } - - fn fmt_sql(&self, f: &mut Formatter<'_>) -> std::fmt::Result { - std::fmt::Display::fmt(self, f) - } - } - - #[derive(Debug)] - struct CustomPhysicalExtensionCodec; - impl PhysicalExtensionCodec for CustomPhysicalExtensionCodec { - fn try_decode( - &self, - _buf: &[u8], - _inputs: &[Arc], - _ctx: &TaskContext, - ) -> Result> { - unreachable!() - } - - fn try_encode( - &self, - _node: Arc, - _buf: &mut Vec, - ) -> Result<()> { - unreachable!() - } - - fn try_decode_expr( - &self, - buf: &[u8], - inputs: &[Arc], - ) -> Result> { - if buf == "CustomPredicateExpr".as_bytes() { - Ok(Arc::new(CustomPredicateExpr { - inner: inputs[0].clone(), - })) - } else { - internal_err!("Not supported") - } - } - - fn try_encode_expr( - &self, - node: &Arc, - buf: &mut Vec, - ) -> Result<()> { - if node.downcast_ref::().is_some() { - buf.extend_from_slice("CustomPredicateExpr".as_bytes()); - Ok(()) - } else { - internal_err!("Not supported") - } - } - } - - let exec_plan = DataSourceExec::from_data_source(scan_config); - - let ctx = SessionContext::new(); - roundtrip_test_and_return( - exec_plan, - &ctx, - &CustomPhysicalExtensionCodec {}, - &DefaultPhysicalProtoConverter {}, - )?; - Ok(()) -} - -#[test] -fn roundtrip_scalar_udf() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let input = Arc::new(EmptyExec::new(schema.clone())); - - let scalar_fn = Arc::new(|args: &[ColumnarValue]| { - let ColumnarValue::Array(array) = &args[0] else { - panic!("should be array") - }; - Ok(ColumnarValue::from(Arc::new(array.clone()) as ArrayRef)) - }); - - let udf = create_udf( - "dummy", - vec![DataType::Int64], - DataType::Int64, - Volatility::Immutable, - scalar_fn.clone(), - ); - - let fun_def = Arc::new(udf.clone()); - - let expr = ScalarFunctionExpr::new( - "dummy", - fun_def, - vec![col("a", &schema)?], - Field::new("f", DataType::Int64, true).into(), - Arc::new(ConfigOptions::default()), - ); - - let project = ProjectionExec::try_new( - vec![ProjectionExpr { - expr: Arc::new(expr), - alias: "a".to_string(), - }], - input, - )?; - - let ctx = SessionContext::new(); - - ctx.register_udf(udf); - - roundtrip_test_with_context(Arc::new(project), &ctx) -} - -#[derive(Debug)] -struct UDFExtensionCodec; - -impl PhysicalExtensionCodec for UDFExtensionCodec { - fn try_decode( - &self, - _buf: &[u8], - _inputs: &[Arc], - _ctx: &TaskContext, - ) -> Result> { - not_impl_err!("No extension codec provided") - } - - fn try_encode( - &self, - _node: Arc, - _buf: &mut Vec, - ) -> Result<()> { - not_impl_err!("No extension codec provided") - } - - fn try_decode_udf(&self, name: &str, buf: &[u8]) -> Result> { - if name == "regex_udf" { - let proto = MyRegexUdfNode::decode(buf).map_err(|err| { - internal_datafusion_err!("failed to decode regex_udf: {err}") - })?; - - Ok(Arc::new(ScalarUDF::from(MyRegexUdf::new(proto.pattern)))) - } else { - not_impl_err!("unrecognized scalar UDF implementation, cannot decode") - } - } - - fn try_encode_udf(&self, node: &ScalarUDF, buf: &mut Vec) -> Result<()> { - let binding = node.inner(); - if let Some(udf) = binding.downcast_ref::() { - let proto = MyRegexUdfNode { - pattern: udf.pattern.clone(), - }; - proto - .encode(buf) - .map_err(|err| internal_datafusion_err!("failed to encode udf: {err}"))?; - } - Ok(()) - } - - fn try_decode_udaf(&self, name: &str, buf: &[u8]) -> Result> { - if name == "aggregate_udf" { - let proto = MyAggregateUdfNode::decode(buf).map_err(|err| { - internal_datafusion_err!("failed to decode aggregate_udf: {err}") - })?; - - Ok(Arc::new(AggregateUDF::from(MyAggregateUDF::new( - proto.result, - )))) - } else { - not_impl_err!("unrecognized scalar UDF implementation, cannot decode") - } - } - - fn try_encode_udaf(&self, node: &AggregateUDF, buf: &mut Vec) -> Result<()> { - let binding = node.inner(); - if let Some(udf) = binding.downcast_ref::() { - let proto = MyAggregateUdfNode { - result: udf.result.clone(), - }; - proto.encode(buf).map_err(|err| { - internal_datafusion_err!("failed to encode udf: {err:?}") - })?; - } - Ok(()) - } - - fn try_decode_udwf(&self, name: &str, buf: &[u8]) -> Result> { - if name == "custom_udwf" { - let proto = CustomUDWFNode::decode(buf).map_err(|err| { - internal_datafusion_err!("failed to decode custom_udwf: {err}") - })?; - - Ok(Arc::new(WindowUDF::from(CustomUDWF::new(proto.payload)))) - } else { - not_impl_err!( - "unrecognized user-defined window function implementation, cannot decode" - ) - } - } - - fn try_encode_udwf(&self, node: &WindowUDF, buf: &mut Vec) -> Result<()> { - let binding = node.inner(); - if let Some(udwf) = binding.downcast_ref::() { - let proto = CustomUDWFNode { - payload: udwf.payload.clone(), - }; - proto.encode(buf).map_err(|err| { - internal_datafusion_err!("failed to encode udwf: {err:?}") - })?; - } - Ok(()) - } -} - -#[test] -fn roundtrip_scalar_udf_extension_codec() -> Result<()> { - let field_text = Field::new("text", DataType::Utf8, true); - let field_published = Field::new("published", DataType::Boolean, false); - let field_author = Field::new("author", DataType::Utf8, false); - let schema = Arc::new(Schema::new(vec![field_text, field_published, field_author])); - let input = Arc::new(EmptyExec::new(schema.clone())); - - let udf_expr = Arc::new(ScalarFunctionExpr::new( - "regex_udf", - Arc::new(ScalarUDF::from(MyRegexUdf::new(".*".to_string()))), - vec![col("text", &schema)?], - Field::new("f", DataType::Int64, true).into(), - Arc::new(ConfigOptions::default()), - )); - - let filter = Arc::new(FilterExec::try_new( - Arc::new(BinaryExpr::new( - col("published", &schema)?, - Operator::And, - Arc::new(BinaryExpr::new(udf_expr.clone(), Operator::Gt, lit(0))), - )), - input, - )?); - let aggr_expr = - AggregateExprBuilder::new(max_udaf(), vec![udf_expr as Arc]) - .schema(schema.clone()) - .alias("max") - .build() - .map(Arc::new)?; - - let window = Arc::new(WindowAggExec::try_new( - vec![Arc::new(PlainAggregateWindowExpr::new( - aggr_expr.clone(), - &[col("author", &schema)?], - &[], - Arc::new(WindowFrame::new(None)), - None, - ))], - filter, - true, - )?); - - let aggregate = Arc::new(AggregateExec::try_new( - AggregateMode::Final, - PhysicalGroupBy::new(vec![], vec![], vec![], false), - vec![aggr_expr], - vec![None], - window, - schema, - )?); - - let ctx = SessionContext::new(); - let proto_converter = DefaultPhysicalProtoConverter {}; - roundtrip_test_and_return(aggregate, &ctx, &UDFExtensionCodec, &proto_converter)?; - Ok(()) -} - -#[test] -fn roundtrip_udwf_extension_codec() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - let custom_udwf = Arc::new(WindowUDF::from(CustomUDWF::new("payload".to_string()))); - let udwf = create_udwf_window_expr( - &custom_udwf, - &[col("a", &schema)?], - schema.as_ref(), - "custom_udwf(a) PARTITION BY [b] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW".to_string(), - false, - )?; - - let window_frame = WindowFrame::new_bounds( - datafusion_expr::WindowFrameUnits::Range, - WindowFrameBound::Preceding(ScalarValue::Int64(None)), - WindowFrameBound::CurrentRow, - ); - - let udwf_expr = Arc::new(StandardWindowExpr::new( - udwf, - &[col("b", &schema)?], - &[PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions { - descending: false, - nulls_first: false, - }, - }], - Arc::new(window_frame), - )); - - let input = Arc::new(EmptyExec::new(schema.clone())); - let window = Arc::new(BoundedWindowAggExec::try_new( - vec![udwf_expr], - input, - InputOrderMode::Sorted, - true, - )?); - - let ctx = SessionContext::new(); - let proto_converter = DefaultPhysicalProtoConverter {}; - roundtrip_test_and_return(window, &ctx, &UDFExtensionCodec, &proto_converter)?; - Ok(()) -} - -#[test] -fn roundtrip_aggregate_udf_extension_codec() -> Result<()> { - let field_text = Field::new("text", DataType::Utf8, true); - let field_published = Field::new("published", DataType::Boolean, false); - let field_author = Field::new("author", DataType::Utf8, false); - let schema = Arc::new(Schema::new(vec![field_text, field_published, field_author])); - let input = Arc::new(EmptyExec::new(schema.clone())); - - let udf_expr = Arc::new(ScalarFunctionExpr::new( - "regex_udf", - Arc::new(ScalarUDF::from(MyRegexUdf::new(".*".to_string()))), - vec![col("text", &schema)?], - Field::new("f", DataType::Int64, true).into(), - Arc::new(ConfigOptions::default()), - )); - - let udaf = Arc::new(AggregateUDF::from(MyAggregateUDF::new( - "result".to_string(), - ))); - let aggr_args: Vec> = - vec![Arc::new(Literal::new(ScalarValue::from(42)))]; - - let aggr_expr = AggregateExprBuilder::new(Arc::clone(&udaf), aggr_args.clone()) - .schema(Arc::clone(&schema)) - .alias("aggregate_udf") - .build() - .map(Arc::new)?; - - let filter = Arc::new(FilterExec::try_new( - Arc::new(BinaryExpr::new( - col("published", &schema)?, - Operator::And, - Arc::new(BinaryExpr::new(udf_expr, Operator::Gt, lit(0))), - )), - input, - )?); - - let window = Arc::new(WindowAggExec::try_new( - vec![Arc::new(PlainAggregateWindowExpr::new( - aggr_expr, - &[col("author", &schema)?], - &[], - Arc::new(WindowFrame::new(None)), - None, - ))], - filter, - true, - )?); - - let aggr_expr = AggregateExprBuilder::new(udaf, aggr_args.clone()) - .schema(Arc::clone(&schema)) - .alias("aggregate_udf") - .distinct() - .ignore_nulls() - .build() - .map(Arc::new)?; - - let aggregate = Arc::new(AggregateExec::try_new( - AggregateMode::Final, - PhysicalGroupBy::new(vec![], vec![], vec![], false), - vec![aggr_expr], - vec![None], - window, - schema, - )?); - - let ctx = SessionContext::new(); - let proto_converter = DefaultPhysicalProtoConverter {}; - roundtrip_test_and_return(aggregate, &ctx, &UDFExtensionCodec, &proto_converter)?; - Ok(()) -} - -#[test] -fn roundtrip_like() -> Result<()> { - let schema = Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Utf8, false), - ]); - let input = Arc::new(EmptyExec::new(Arc::new(schema.clone()))); - let like_expr = like( - false, - false, - col("a", &schema)?, - col("b", &schema)?, - &schema, - )?; - let plan = Arc::new(ProjectionExec::try_new( - vec![ProjectionExpr { - expr: like_expr, - alias: "result".to_string(), - }], - input, - )?); - roundtrip_test(plan) -} - -#[test] -fn roundtrip_analyze() -> Result<()> { - let field_a = Field::new("plan_type", DataType::Utf8, false); - let field_b = Field::new("plan", DataType::Utf8, false); - let schema = Schema::new(vec![field_a, field_b]); - let input = Arc::new(PlaceholderRowExec::new(Arc::new(schema.clone()))); - - roundtrip_test(Arc::new(AnalyzeExec::new( - false, - false, - vec![MetricType::Summary, MetricType::Dev], - None, - input, - Arc::new(schema), - ))) -} - -#[tokio::test] -async fn roundtrip_json_source() -> Result<()> { - let ctx = SessionContext::new(); - ctx.register_json("t1", "../core/tests/data/1.json", Default::default()) - .await?; - let plan = ctx.table("t1").await?.create_physical_plan().await?; - roundtrip_test(plan) -} - -#[test] -fn roundtrip_json_sink() -> Result<()> { - let field_a = Field::new("plan_type", DataType::Utf8, false); - let field_b = Field::new("plan", DataType::Utf8, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let input = Arc::new(PlaceholderRowExec::new(schema.clone())); - - let file_sink_config = FileSinkConfig { - original_url: String::default(), - object_store_url: ObjectStoreUrl::local_filesystem(), - file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), - table_paths: vec![ListingTableUrl::parse("file:///")?], - output_schema: schema.clone(), - table_partition_cols: vec![("plan_type".to_string(), DataType::Utf8)], - insert_op: InsertOp::Overwrite, - keep_partition_by_columns: true, - file_extension: "json".into(), - file_output_mode: FileOutputMode::SingleFile, - }; - let data_sink = Arc::new(JsonSink::new( - file_sink_config, - JsonWriterOptions::new(CompressionTypeVariant::UNCOMPRESSED), - )); - let sort_order = [PhysicalSortRequirement::new( - Arc::new(Column::new("plan_type", 0)), - Some(SortOptions { - descending: true, - nulls_first: false, - }), - )] - .into(); - - roundtrip_test(Arc::new(DataSinkExec::new( - input, - data_sink, - Some(sort_order), - ))) -} - -#[test] -fn roundtrip_csv_sink() -> Result<()> { - let field_a = Field::new("plan_type", DataType::Utf8, false); - let field_b = Field::new("plan", DataType::Utf8, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let input = Arc::new(PlaceholderRowExec::new(schema.clone())); - - let file_sink_config = FileSinkConfig { - original_url: String::default(), - object_store_url: ObjectStoreUrl::local_filesystem(), - file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), - table_paths: vec![ListingTableUrl::parse("file:///")?], - output_schema: schema.clone(), - table_partition_cols: vec![("plan_type".to_string(), DataType::Utf8)], - insert_op: InsertOp::Overwrite, - keep_partition_by_columns: true, - file_extension: "csv".into(), - file_output_mode: FileOutputMode::Directory, - }; - let data_sink = Arc::new(CsvSink::new( - file_sink_config, - CsvWriterOptions::new(WriterBuilder::default(), CompressionTypeVariant::ZSTD), - )); - let sort_order = [PhysicalSortRequirement::new( - Arc::new(Column::new("plan_type", 0)), - Some(SortOptions { - descending: true, - nulls_first: false, - }), - )] - .into(); - - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - - let roundtrip_plan = roundtrip_test_and_return( - Arc::new(DataSinkExec::new(input, data_sink, Some(sort_order))), - &ctx, - &codec, - &proto_converter, - )?; - - let roundtrip_plan = roundtrip_plan.downcast_ref::().unwrap(); - let csv_sink = roundtrip_plan.sink().downcast_ref::().unwrap(); - assert_eq!( - CompressionTypeVariant::ZSTD, - csv_sink.writer_options().compression - ); - - Ok(()) -} - -#[test] -fn roundtrip_parquet_sink() -> Result<()> { - let field_a = Field::new("plan_type", DataType::Utf8, false); - let field_b = Field::new("plan", DataType::Utf8, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let input = Arc::new(PlaceholderRowExec::new(schema.clone())); - - let file_sink_config = FileSinkConfig { - original_url: String::default(), - object_store_url: ObjectStoreUrl::local_filesystem(), - file_group: FileGroup::new(vec![PartitionedFile::new("/tmp".to_string(), 1)]), - table_paths: vec![ListingTableUrl::parse("file:///")?], - output_schema: schema.clone(), - table_partition_cols: vec![("plan_type".to_string(), DataType::Utf8)], - insert_op: InsertOp::Overwrite, - keep_partition_by_columns: true, - file_extension: "parquet".into(), - file_output_mode: FileOutputMode::Automatic, - }; - let data_sink = Arc::new(ParquetSink::new( - file_sink_config, - TableParquetOptions::default(), - )); - let sort_order = [PhysicalSortRequirement::new( - Arc::new(Column::new("plan_type", 0)), - Some(SortOptions { - descending: true, - nulls_first: false, - }), - )] - .into(); - - roundtrip_test(Arc::new(DataSinkExec::new( - input, - data_sink, - Some(sort_order), - ))) -} - -#[test] -fn roundtrip_sym_hash_join() -> Result<()> { - let field_a = Field::new("col", DataType::Int64, false); - let schema_left = Schema::new(vec![field_a.clone()]); - let schema_right = Schema::new(vec![field_a]); - let on = vec![( - Arc::new(Column::new("col", schema_left.index_of("col")?)) as _, - Arc::new(Column::new("col", schema_right.index_of("col")?)) as _, - )]; - - let schema_left = Arc::new(schema_left); - let schema_right = Arc::new(schema_right); - for join_type in &[ - JoinType::Inner, - JoinType::Left, - JoinType::Right, - JoinType::Full, - JoinType::LeftAnti, - JoinType::RightAnti, - JoinType::LeftSemi, - JoinType::RightSemi, - ] { - for partition_mode in &[ - StreamJoinPartitionMode::Partitioned, - StreamJoinPartitionMode::SinglePartition, - ] { - for left_order in &[ - None, - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::new(Column::new("col", schema_left.index_of("col")?)), - options: Default::default(), - }]), - ] { - for right_order in [ - None, - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::new(Column::new("col", schema_right.index_of("col")?)), - options: Default::default(), - }]), - ] { - roundtrip_test(Arc::new(SymmetricHashJoinExec::try_new( - Arc::new(EmptyExec::new(schema_left.clone())), - Arc::new(EmptyExec::new(schema_right.clone())), - on.clone(), - None, - join_type, - NullEquality::NullEqualsNothing, - left_order.clone(), - right_order, - *partition_mode, - )?))?; - } - } - } - } - Ok(()) -} - -#[test] -fn roundtrip_union() -> Result<()> { - let field_a = Field::new("col", DataType::Int64, false); - let schema_left = Schema::new(vec![field_a.clone()]); - let schema_right = Schema::new(vec![field_a]); - let left = EmptyExec::new(Arc::new(schema_left)); - let right = EmptyExec::new(Arc::new(schema_right)); - let inputs: Vec> = vec![Arc::new(left), Arc::new(right)]; - let union = UnionExec::try_new(inputs)?; - roundtrip_test(union) -} - -#[test] -fn roundtrip_repartition_preserve_order() -> Result<()> { - let field_a = Field::new("a", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a])); - let sort_exprs: LexOrdering = [PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions::default(), - }] - .into(); - - // Create two sorted single-partition inputs, then union them to get - // a sorted input with 2 partitions. - let source1 = SortExec::new( - sort_exprs.clone(), - Arc::new(EmptyExec::new(Arc::clone(&schema))), - ); - let source2 = SortExec::new(sort_exprs, Arc::new(EmptyExec::new(schema))); - let union = UnionExec::try_new(vec![ - Arc::new(source1) as Arc, - Arc::new(source2) as Arc, - ])?; - - let repartition = RepartitionExec::try_new(union, Partitioning::RoundRobinBatch(10))? - .with_preserve_order(); - assert!(repartition.preserve_order()); - - roundtrip_test(Arc::new(repartition)) -} - -#[test] -fn roundtrip_interleave() -> Result<()> { - let field_a = Field::new("col", DataType::Int64, false); - let schema_left = Schema::new(vec![field_a.clone()]); - let schema_right = Schema::new(vec![field_a]); - let partition = Partitioning::Hash(vec![], 3); - let left = RepartitionExec::try_new( - Arc::new(EmptyExec::new(Arc::new(schema_left))), - partition.clone(), - )?; - let right = RepartitionExec::try_new( - Arc::new(EmptyExec::new(Arc::new(schema_right))), - partition, - )?; - let inputs: Vec> = vec![Arc::new(left), Arc::new(right)]; - let interleave = InterleaveExec::try_new(inputs)?; - roundtrip_test(Arc::new(interleave)) -} - -#[test] -fn roundtrip_unnest() -> Result<()> { - let fa = Field::new("a", DataType::Int64, true); - let fb0 = Field::new_list_field(DataType::Utf8, true); - let fb = Field::new_list("b", fb0.clone(), false); - let fc1 = Field::new("c1", DataType::Boolean, false); - let fc2 = Field::new("c2", DataType::Date64, true); - let fc = Field::new_struct("c", Fields::from(vec![fc1.clone(), fc2.clone()]), true); - let fd0 = Field::new_list_field(DataType::Float32, false); - let fd = Field::new_list("d", fd0.clone(), true); - let fe1 = Field::new("e1", DataType::UInt16, false); - let fe2 = Field::new("e2", DataType::Duration(TimeUnit::Millisecond), true); - let fe3 = Field::new("e3", DataType::Timestamp(TimeUnit::Millisecond, None), true); - let fe_fields = Fields::from(vec![fe1.clone(), fe2.clone(), fe3.clone()]); - let fe = Field::new_struct("e", fe_fields, false); - - let fb0 = fb0.with_name("b"); - let fd0 = fd0.with_name("d"); - let input_schema = Arc::new(Schema::new(vec![fa.clone(), fb, fc, fd, fe])); - let output_schema = - Arc::new(Schema::new(vec![fa, fb0, fc1, fc2, fd0, fe1, fe2, fe3])); - let input = Arc::new(EmptyExec::new(input_schema)); - let options = UnnestOptions::default(); - let unnest = UnnestExec::new( - input, - vec![ - ListUnnest { - index_in_input_schema: 1, - depth: 1, - }, - ListUnnest { - index_in_input_schema: 1, - depth: 2, - }, - ListUnnest { - index_in_input_schema: 3, - depth: 2, - }, - ], - vec![2, 4], - output_schema, - options, - )?; - roundtrip_test(Arc::new(unnest)) -} - -#[tokio::test] -async fn roundtrip_coalesce() -> Result<()> { - let ctx = SessionContext::new(); - ctx.register_table( - "t", - Arc::new(EmptyTable::new(Arc::new(Schema::new(Fields::from([ - Arc::new(Field::new("f", DataType::Int64, false)), - ]))))), - )?; - let df = ctx.sql("select coalesce(f) as f from t").await?; - let plan = df.create_physical_plan().await?; - - let node = PhysicalPlanNode::try_from_physical_plan( - plan.clone(), - &DefaultPhysicalExtensionCodec {}, - )?; - let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - let restored = - node.try_into_physical_plan(&ctx.task_ctx(), &DefaultPhysicalExtensionCodec {})?; - - assert_eq!( - plan.schema(), - restored.schema(), - "Schema mismatch for plans:\n>> initial:\n{}>> final: \n{}", - displayable(plan.as_ref()) - .set_show_schema(true) - .indent(true), - displayable(restored.as_ref()) - .set_show_schema(true) - .indent(true), - ); - - Ok(()) -} - -#[tokio::test] -async fn roundtrip_generate_series() -> Result<()> { - let ctx = SessionContext::new(); - ctx.register_table( - "t", - Arc::new(EmptyTable::new(Arc::new(Schema::new(Fields::from([ - Arc::new(Field::new("f", DataType::Int64, false)), - ]))))), - )?; - let df = ctx.sql("select * from generate_series(1, 10000)").await?; - let plan = df.create_physical_plan().await?; - - let node = PhysicalPlanNode::try_from_physical_plan( - plan.clone(), - &DefaultPhysicalExtensionCodec {}, - )?; - let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - let restored = - node.try_into_physical_plan(&ctx.task_ctx(), &DefaultPhysicalExtensionCodec {})?; - - assert_eq!( - plan.schema(), - restored.schema(), - "Schema mismatch for plans:\n>> initial:\n{}>> final: \n{}", - displayable(plan.as_ref()) - .set_show_schema(true) - .indent(true), - displayable(restored.as_ref()) - .set_show_schema(true) - .indent(true), - ); - - Ok(()) -} - -#[tokio::test] -async fn roundtrip_projection_source() -> Result<()> { - let schema = Arc::new(Schema::new(Fields::from([ - Arc::new(Field::new("a", DataType::Utf8, false)), - Arc::new(Field::new("b", DataType::Utf8, false)), - Arc::new(Field::new("c", DataType::Int32, false)), - Arc::new(Field::new("d", DataType::Int32, false)), - ]))); - - let statistics = Statistics::new_unknown(&schema); - - let file_source = Arc::new(ParquetSource::new(Arc::clone(&schema))); - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) - .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )])]) - .with_statistics(statistics) - .with_projection_indices(Some(vec![0, 1, 2]))? - .build(); - - let filter = Arc::new( - FilterExecBuilder::new( - Arc::new(BinaryExpr::new(col("c", &schema)?, Operator::Eq, lit(1))), - DataSourceExec::from_data_source(scan_config), - ) - .apply_projection(Some(vec![0, 1]))? - .build()?, - ); - - roundtrip_test(filter) -} - -#[tokio::test] -async fn roundtrip_parquet_select_star() -> Result<()> { - let ctx = all_types_context().await?; - let sql = "select * from alltypes_plain"; - roundtrip_test_sql_with_context(sql, &ctx).await -} - -#[tokio::test] -async fn roundtrip_parquet_select_projection() -> Result<()> { - let ctx = all_types_context().await?; - let sql = "select string_col, timestamp_col from alltypes_plain"; - roundtrip_test_sql_with_context(sql, &ctx).await -} - -#[tokio::test] -async fn roundtrip_parquet_select_star_predicate() -> Result<()> { - let ctx = all_types_context().await?; - let sql = "select * from alltypes_plain where id > 4"; - roundtrip_test_sql_with_context(sql, &ctx).await -} - -#[tokio::test] -async fn roundtrip_parquet_select_projection_predicate() -> Result<()> { - let ctx = all_types_context().await?; - let sql = "select string_col, timestamp_col from alltypes_plain where id > 4"; - roundtrip_test_sql_with_context(sql, &ctx).await -} - -#[tokio::test] -async fn roundtrip_empty_projection() -> Result<()> { - let ctx = all_types_context().await?; - let sql = "select 1 from alltypes_plain"; - roundtrip_test_sql_with_context(sql, &ctx).await -} - -#[tokio::test] -async fn roundtrip_physical_plan_node() { - use datafusion::prelude::*; - use datafusion_proto::physical_plan::{ - AsExecutionPlan, DefaultPhysicalExtensionCodec, - }; - use datafusion_proto::protobuf::PhysicalPlanNode; - - let ctx = SessionContext::new(); - - ctx.register_parquet( - "pt", - &format!( - "{}/alltypes_plain.snappy.parquet", - datafusion_common::test_util::parquet_test_data() - ), - ParquetReadOptions::default(), - ) - .await - .unwrap(); - - let plan = ctx - .sql("select id, string_col, timestamp_col from pt where id > 4 order by string_col") - .await - .unwrap() - .create_physical_plan() - .await - .unwrap(); - - let node: PhysicalPlanNode = - PhysicalPlanNode::try_from_physical_plan(plan, &DefaultPhysicalExtensionCodec {}) - .unwrap(); - - let plan = node - .try_into_physical_plan(&ctx.task_ctx(), &DefaultPhysicalExtensionCodec {}) - .unwrap(); - - let _ = plan.execute(0, ctx.task_ctx()).unwrap(); -} - -/// Helper function to create a SessionContext with all TPC-H tables registered as external tables -async fn tpch_context() -> Result { - use datafusion_common::test_util::datafusion_test_data; - - let ctx = SessionContext::new(); - let test_data = datafusion_test_data(); - - // TPC-H table names - let tables = [ - "part", "supplier", "partsupp", "customer", "orders", "lineitem", "nation", - "region", - ]; - - // Create external tables for all TPC-H tables - for table in &tables { - let table_sql = format!( - "CREATE EXTERNAL TABLE {table} STORED AS PARQUET LOCATION '{test_data}/tpch_{table}_small.parquet'" - ); - ctx.sql(&table_sql).await.map_err(|e| { - DataFusionError::External( - format!("Failed to create {table} table: {e}").into(), - ) - })?; - } - - Ok(ctx) -} - -/// Helper function to get TPC-H query SQL -fn get_tpch_query_sql(query: usize) -> Result> { - use std::fs; - - if !(1..=22).contains(&query) { - return Err(DataFusionError::External( - format!("Invalid TPC-H query number: {query}").into(), - )); - } - - let filename = format!("../../benchmarks/queries/q{query}.sql"); - let contents = fs::read_to_string(&filename).map_err(|e| { - DataFusionError::External( - format!("Failed to read query file {filename}: {e}").into(), - ) - })?; - - Ok(contents - .split(';') - .map(|s| s.trim()) - .filter(|s| !s.is_empty()) - .map(|s| s.to_string()) - .collect()) -} - -#[tokio::test] -async fn test_serialize_deserialize_tpch_queries() -> Result<()> { - // Create context with TPC-H tables - let ctx = tpch_context().await?; - - // repeat to run all 22 queries - for query in 1..=22 { - // run all statements in the query - let sql = get_tpch_query_sql(query)?; - for stmt in sql { - let logical_plan = ctx.sql(&stmt).await?.into_unoptimized_plan(); - let optimized_plan = ctx.state().optimize(&logical_plan)?; - let physical_plan = ctx.state().create_physical_plan(&optimized_plan).await?; - - // serialize the physical plan - let codec = DefaultPhysicalExtensionCodec {}; - - let proto = - PhysicalPlanNode::try_from_physical_plan(physical_plan.clone(), &codec)?; - - // deserialize the physical plan - let _deserialized_plan = - proto.try_into_physical_plan(&ctx.task_ctx(), &codec)?; - } - } - - Ok(()) -} - -// Bugs: https://github.com/apache/datafusion/issues/16772 -#[tokio::test] -async fn test_round_trip_tpch_queries() -> Result<()> { - // Create context with TPC-H tables - let ctx = tpch_context().await?; - - // repeat to run all 22 queries - for query in 1..=22 { - // run all statements in the query - let sql = get_tpch_query_sql(query)?; - for stmt in sql { - roundtrip_test_sql_with_context(&stmt, &ctx).await?; - } - } - - Ok(()) -} - -// Bug 1 of https://github.com/apache/datafusion/issues/16772 -/// Test that AggregateFunctionExpr human_display field is correctly preserved -/// during serialization/deserialization roundtrip. -/// -/// Test for issue where the human_display field (used for EXPLAIN output) -/// was not being serialized to protobuf, causing it to be lost during roundtrip -/// and resulting in empty or incorrect display strings in query plans. -#[tokio::test] -async fn test_round_trip_human_display() -> Result<()> { - // Create context with TPC-H tables - let ctx = tpch_context().await?; - - let sql = "select r_name, count(1) from region group by r_name"; - roundtrip_test_sql_with_context(sql, &ctx).await?; - - let sql = "select r_name, count(*) from region group by r_name"; - roundtrip_test_sql_with_context(sql, &ctx).await?; - - let sql = "select r_name, count(r_name) from region group by r_name"; - roundtrip_test_sql_with_context(sql, &ctx).await?; - - let sql = "select count(*) as count_star from region"; - roundtrip_test_sql_with_context(sql, &ctx).await?; - - Ok(()) -} - -#[test] -fn test_round_trip_aliased_reverse_human_display() -> Result<()> { - let aggregate_expr = roundtrip_first_value_aggregate( - "agg", - "first_value(b) ORDER BY [b ASC NULLS LAST]", - Some("agg"), - )?; - let reversed = aggregate_expr - .reverse_expr() - .expect("expected reverse expr"); - - assert_eq!(reversed.name(), "agg"); - assert_eq!(reversed.human_display_alias(), Some("agg")); - assert_eq!( - reversed.human_display(), - Some("last_value(b) ORDER BY [b DESC NULLS FIRST]") - ); - - Ok(()) -} - -#[test] -fn test_round_trip_human_display_alias_with_colon() -> Result<()> { - let aggregate_expr = roundtrip_first_value_aggregate( - "agg:one", - "first_value(b) ORDER BY [b ASC NULLS LAST]", - Some("agg:one"), - )?; - - assert_eq!(aggregate_expr.name(), "agg:one"); - assert_eq!(aggregate_expr.human_display_alias(), Some("agg:one")); - assert_eq!( - aggregate_expr.human_display(), - Some("first_value(b) ORDER BY [b ASC NULLS LAST]") - ); - - Ok(()) -} - -#[test] -fn test_round_trip_non_aliased_human_display_ending_like_alias() -> Result<()> { - let aggregate_expr = - roundtrip_first_value_aggregate("agg", "first_value(b) as agg", None)?; - - assert_eq!(aggregate_expr.name(), "agg"); - assert_eq!( - aggregate_expr.human_display(), - Some("first_value(b) as agg") - ); - assert_eq!(aggregate_expr.human_display_alias(), None); - - Ok(()) -} - -fn roundtrip_first_value_aggregate( - alias: &str, - human_display: &str, - human_display_alias: Option<&str>, -) -> Result> { - let schema = Arc::new(Schema::new(vec![Field::new("b", DataType::Int64, true)])); - let mut builder = - AggregateExprBuilder::new(first_value_udaf(), vec![col("b", &schema)?]) - .order_by(vec![PhysicalSortExpr { - expr: col("b", &schema)?, - options: SortOptions::new(false, false), - }]) - .schema(Arc::clone(&schema)) - .alias(alias) - .human_display(human_display); - if let Some(human_display_alias) = human_display_alias { - builder = builder.human_display_alias(human_display_alias); - } - let agg_expr = builder.build().map(Arc::new)?; - - let plan = Arc::new(AggregateExec::try_new( - AggregateMode::Single, - PhysicalGroupBy::new(vec![], vec![], vec![], false), - vec![agg_expr], - vec![None], - Arc::new(EmptyExec::new(Arc::clone(&schema))), - schema, - )?); - - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - let roundtrip_plan = roundtrip_test_and_return(plan, &ctx, &codec, &proto_converter)?; - let aggregate = roundtrip_plan - .as_ref() - .downcast_ref::() - .expect("expected AggregateExec after roundtrip"); - - Ok(Arc::clone(&aggregate.aggr_expr()[0])) -} - -// Bug 2 of https://github.com/apache/datafusion/issues/16772 -/// Test that PhysicalGroupBy groups field is correctly serialized/deserialized -/// for simple aggregates (no GROUP BY clause). -/// -/// Test for issue where simple aggregates like "SELECT SUM(col1 * col2) FROM table" -/// would incorrectly serialize groups as [[]] instead of [] during roundtrip serialization. -/// The groups field should be empty ([]) when there are no GROUP BY expressions. -#[tokio::test] -async fn test_round_trip_groups_display() -> Result<()> { - // Create context with TPC-H tables - let ctx = tpch_context().await?; - - let sql = "select sum(l_extendedprice * l_discount) as revenue from lineitem;"; - roundtrip_test_sql_with_context(sql, &ctx).await?; - - let sql = "select sum(l_extendedprice) as revenue from lineitem;"; - roundtrip_test_sql_with_context(sql, &ctx).await?; - - Ok(()) -} - -// Bug 3 of https://github.com/apache/datafusion/issues/16772 -/// Test that ScalarFunctionExpr return_field name is correctly preserved -/// during serialization/deserialization roundtrip. -/// -/// Test for issue where the return_field.name for scalar functions -/// was not being serialized to protobuf, causing it to be lost during roundtrip -/// and defaulting to a generic name like "f" instead of the proper function name. -#[tokio::test] -async fn test_round_trip_date_part_display() -> Result<()> { - // Create context with TPC-H tables - let ctx = tpch_context().await?; - - let sql = "select extract(year from l_shipdate) as l_year from lineitem "; - roundtrip_test_sql_with_context(sql, &ctx).await?; - - let sql = "select extract(month from l_shipdate) as l_year from lineitem "; - roundtrip_test_sql_with_context(sql, &ctx).await?; - - Ok(()) -} - -#[tokio::test] -async fn test_tpch_part_in_list_query_with_real_parquet_data() -> Result<()> { - use datafusion_common::test_util::datafusion_test_data; - - let ctx = SessionContext::new(); - - // Register the TPC-H part table using the local test data - let test_data = datafusion_test_data(); - let table_sql = format!( - "CREATE EXTERNAL TABLE part STORED AS PARQUET LOCATION '{test_data}/tpch_part_small.parquet'" - ); - ctx.sql(&table_sql).await.map_err(|e| { - DataFusionError::External(format!("Failed to create part table: {e}").into()) - })?; - - // Test the exact problematic query - let sql = - "SELECT p_size FROM part WHERE p_size IN (14, 6, 5, 31) and p_partkey > 1000"; - - let logical_plan = ctx.sql(sql).await?.into_unoptimized_plan(); - let optimized_plan = ctx.state().optimize(&logical_plan)?; - let physical_plan = ctx.state().create_physical_plan(&optimized_plan).await?; - - // Serialize the physical plan - bug may happen here already but not necessarily manifests - let codec = DefaultPhysicalExtensionCodec {}; - - let proto = PhysicalPlanNode::try_from_physical_plan(physical_plan.clone(), &codec)?; - - // This will fail with the bug, but should succeed when fixed - let _deserialized_plan = proto.try_into_physical_plan(&ctx.task_ctx(), &codec)?; - Ok(()) -} - -#[tokio::test] -/// Tests that we can serialize an unoptimized "analyze" plan and it will work on the other end -async fn analyze_roundtrip_unoptimized() -> Result<()> { - let ctx = SessionContext::new(); - - // No optimizations - let session_state = - datafusion::execution::SessionStateBuilder::new_from_existing(ctx.state()) - .with_physical_optimizer_rules(vec![]) - .build(); - - let logical_plan = session_state - .create_logical_plan("explain analyze select 1") - .await?; - let plan = session_state.create_physical_plan(&logical_plan).await?; - - let node = PhysicalPlanNode::try_from_physical_plan( - plan.clone(), - &DefaultPhysicalExtensionCodec {}, - )?; - - let node = PhysicalPlanNode::decode(node.encode_to_vec().as_slice()) - .map_err(|e| DataFusionError::External(Box::new(e)))?; - - let unoptimized = - node.try_into_physical_plan(&ctx.task_ctx(), &DefaultPhysicalExtensionCodec {})?; - - let physical_planner = - datafusion::physical_planner::DefaultPhysicalPlanner::default(); - physical_planner.optimize_physical_plan(unoptimized, &session_state, |_, _| {})?; - Ok(()) -} - -#[test] -fn roundtrip_sort_merge_join() -> Result<()> { - let field_a = Field::new("col_a", DataType::Int64, false); - let field_b = Field::new("col_b", DataType::Int64, false); - let schema_left = Schema::new(vec![field_a.clone()]); - let schema_right = Schema::new(vec![field_b.clone()]); - let on = vec![( - Arc::new(Column::new("col_a", schema_left.index_of("col_a")?)) as _, - Arc::new(Column::new("col_b", schema_right.index_of("col_b")?)) as _, - )]; - - let filter = datafusion::physical_plan::joins::utils::JoinFilter::new( - Arc::new(BinaryExpr::new( - Arc::new(Column::new("col_a", 1)), - Operator::Gt, - Arc::new(Column::new("col_b", 0)), - )), - vec![ - datafusion::physical_plan::joins::utils::ColumnIndex { - index: 0, - side: datafusion_common::JoinSide::Left, - }, - datafusion::physical_plan::joins::utils::ColumnIndex { - index: 0, - side: datafusion_common::JoinSide::Right, - }, - ], - Arc::new(Schema::new(vec![field_a, field_b])), - ); - - let schema_left = Arc::new(schema_left); - let schema_right = Arc::new(schema_right); - for filter in [None, Some(filter)] { - for join_type in [ - JoinType::Inner, - JoinType::Left, - JoinType::Right, - JoinType::Full, - JoinType::LeftAnti, - JoinType::RightAnti, - JoinType::LeftSemi, - JoinType::RightSemi, - ] { - roundtrip_test(Arc::new(SortMergeJoinExec::try_new( - Arc::new(EmptyExec::new(schema_left.clone())), - Arc::new(EmptyExec::new(schema_right.clone())), - on.clone(), - filter.clone(), - join_type, - vec![Default::default()], - NullEquality::NullEqualsNothing, - )?))?; - } - } - Ok(()) -} - -#[tokio::test] -async fn roundtrip_logical_plan_sort_merge_join() -> Result<()> { - let ctx = SessionContext::new(); - ctx.register_csv( - "t0", - "tests/testdata/test.csv", - datafusion::prelude::CsvReadOptions::default().has_header(true), - ) - .await?; - ctx.register_csv( - "t1", - "tests/testdata/test.csv", - datafusion::prelude::CsvReadOptions::default().has_header(true), - ) - .await?; - - ctx.sql("SET datafusion.optimizer.prefer_hash_join = false") - .await? - .show() - .await?; - - let query = "SELECT t1.* FROM t0 join t1 on t0.a = t1.a"; - let plan = ctx.sql(query).await?.create_physical_plan().await?; - roundtrip_test(plan) -} - -#[tokio::test] -async fn roundtrip_memory_source() -> Result<()> { - let ctx = SessionContext::new(); - let plan = ctx - .sql("select * from values ('Tom', 18)") - .await? - .create_physical_plan() - .await?; - roundtrip_test(plan) -} - -#[tokio::test] -async fn roundtrip_listing_table_with_schema_metadata() -> Result<()> { - let ctx = SessionContext::new(); - let file_format = JsonFormat::default(); - let table_partition_cols = vec![("part".to_owned(), DataType::Int64)]; - let data = "../core/tests/data/partitioned_table_json"; - let listing_table_url = ListingTableUrl::parse(data)?; - let listing_options = ListingOptions::new(Arc::new(file_format)) - .with_table_partition_cols(table_partition_cols); - - let config = ListingTableConfig::new(listing_table_url) - .with_listing_options(listing_options) - .infer_schema(&ctx.state()) - .await?; - - // Decorate metadata onto the inferred ListingTable schema - let schema_with_meta = config - .file_schema - .clone() - .map(|s| { - let mut meta: HashMap = HashMap::new(); - meta.insert("foo.bar".to_string(), "baz".to_string()); - s.as_ref().clone().with_metadata(meta) - }) - .expect("Must decorate metadata"); - - let config = config.with_schema(Arc::new(schema_with_meta)); - ctx.register_table("hive_style", Arc::new(ListingTable::try_new(config)?))?; - - let plan = ctx - .sql("select * from hive_style limit 1") - .await? - .create_physical_plan() - .await?; - - roundtrip_test(plan) -} - -#[tokio::test] -async fn roundtrip_async_func_exec() -> Result<()> { - #[derive(Debug, PartialEq, Eq, Hash)] - struct TestAsyncUDF { - signature: Signature, - } - - impl TestAsyncUDF { - fn new() -> Self { - Self { - signature: Signature::exact(vec![DataType::Int64], Volatility::Volatile), - } - } - } - - impl ScalarUDFImpl for TestAsyncUDF { - fn name(&self) -> &str { - "test_async_udf" - } - - fn signature(&self) -> &Signature { - &self.signature - } - - fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(DataType::Int64) - } - - fn invoke_with_args(&self, _args: ScalarFunctionArgs) -> Result { - not_impl_err!("Must call from `invoke_async_with_args`") - } - } - - #[async_trait::async_trait] - impl AsyncScalarUDFImpl for TestAsyncUDF { - async fn invoke_async_with_args( - &self, - args: ScalarFunctionArgs, - ) -> Result { - Ok(args.args[0].clone()) - } - } - - let ctx = SessionContext::new(); - let async_udf = AsyncScalarUDF::new(Arc::new(TestAsyncUDF::new())); - ctx.register_udf(async_udf.into_scalar_udf()); - - let physical_plan = ctx - .sql("select test_async_udf(1)") - .await? - .create_physical_plan() - .await?; - - roundtrip_test_with_context(physical_plan, &ctx)?; - - Ok(()) -} - -/// Test that HashTableLookupExpr serializes to lit(true) -/// -/// HashTableLookupExpr contains a runtime hash table that cannot be serialized. -/// The serialization code replaces it with lit(true) which is safe because -/// it's a performance optimization filter, not a correctness requirement. -#[test] -fn roundtrip_hash_table_lookup_expr_to_lit() -> Result<()> { - use datafusion::physical_plan::joins::join_hash_map::JoinHashMapU32; - use datafusion::physical_plan::joins::{HashTableLookupExpr, Map}; - - // Create a simple schema and input plan - let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); - let input = Arc::new(EmptyExec::new(schema.clone())); - - // Create a HashTableLookupExpr - it will be replaced with lit(true) during serialization - let hash_map = Arc::new(Map::HashMap(Box::new(JoinHashMapU32::with_capacity(0)))); - let on_columns = vec![datafusion::physical_plan::expressions::col("col", &schema)?]; - let lookup_expr: Arc = Arc::new(HashTableLookupExpr::new( - on_columns, - datafusion::physical_plan::joins::SeededRandomState::with_seed(0), - hash_map, - "test_lookup".to_string(), - )); - - // Create a filter with the lookup expression - let filter = Arc::new(FilterExec::try_new(lookup_expr, input)?); - - // Serialize - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - - let proto: PhysicalPlanNode = - PhysicalPlanNode::try_from_physical_plan(filter.clone(), &codec) - .expect("serialization should succeed"); - - // Deserialize - let result: Arc = proto - .try_into_physical_plan(&ctx.task_ctx(), &codec) - .expect("deserialization should succeed"); - - // The deserialized plan should have lit(true) instead of HashTableLookupExpr - // Verify the filter predicate is a Literal(true) - let result_filter = result.downcast_ref::().unwrap(); - let predicate = result_filter.predicate(); - let literal = predicate.downcast_ref::().unwrap(); - assert_eq!(*literal.value(), ScalarValue::Boolean(Some(true))); - - Ok(()) -} - -#[test] -fn roundtrip_hash_expr() -> Result<()> { - use datafusion::physical_plan::joins::{HashExpr, SeededRandomState}; - - let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Int64, false), - Field::new("b", DataType::Utf8, false), - ])); - - // Create a HashExpr with test columns and seeds - let on_columns = vec![col("a", &schema)?, col("b", &schema)?]; - let hash_expr: Arc = Arc::new(HashExpr::new( - on_columns, - SeededRandomState::with_seed(0), // arbitrary random seed for testing - "test_hash".to_string(), - )); - - // Wrap in a filter by comparing hash value to a literal - // hash_expr > 0 is always boolean - let filter_expr = binary(hash_expr, Operator::Gt, lit(0u64), &schema)?; - let filter = Arc::new(FilterExec::try_new( - filter_expr, - Arc::new(EmptyExec::new(schema)), - )?); - - // Confirm that the debug string contains the random state seeds - assert!( - format!("{filter:?}").contains("test_hash(a@0, b@1, [0])"), - "Debug string missing seeds: {filter:?}" - ); - roundtrip_test(filter) -} - -#[test] -fn custom_proto_converter_intercepts() -> Result<()> { - #[derive(Default)] - struct CustomConverterInterceptor { - num_proto_plans: RwLock, - num_physical_plans: RwLock, - num_proto_exprs: RwLock, - num_physical_exprs: RwLock, - } - - impl PhysicalProtoConverterExtension for CustomConverterInterceptor { - fn proto_to_execution_plan( - &self, - proto: &protobuf::PhysicalPlanNode, - ctx: &PhysicalPlanDecodeContext<'_>, - ) -> Result> { - { - let mut counter = self - .num_proto_plans - .write() - .map_err(|err| exec_datafusion_err!("{err}"))?; - *counter += 1; - } - self.default_proto_to_execution_plan(proto, ctx) - } - - fn execution_plan_to_proto( - &self, - plan: &Arc, - codec: &dyn PhysicalExtensionCodec, - ) -> Result - where - Self: Sized, - { - { - let mut counter = self - .num_physical_plans - .write() - .map_err(|err| exec_datafusion_err!("{err}"))?; - *counter += 1; - } - PhysicalPlanNode::try_from_physical_plan_with_converter( - Arc::clone(plan), - codec, - self, - ) - } - - fn proto_to_physical_expr( - &self, - proto: &PhysicalExprNode, - input_schema: &Schema, - ctx: &PhysicalPlanDecodeContext<'_>, - ) -> Result> - where - Self: Sized, - { - { - let mut counter = self - .num_proto_exprs - .write() - .map_err(|err| exec_datafusion_err!("{err}"))?; - *counter += 1; - } - self.default_proto_to_physical_expr(proto, input_schema, ctx) - } - - fn physical_expr_to_proto( - &self, - expr: &Arc, - codec: &dyn PhysicalExtensionCodec, - ) -> Result { - { - let mut counter = self - .num_physical_exprs - .write() - .map_err(|err| exec_datafusion_err!("{err}"))?; - *counter += 1; - } - serialize_physical_expr_with_converter(expr, codec, self) - } - } - - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - let sort_exprs = [ - PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions { - descending: true, - nulls_first: false, - }, - }, - PhysicalSortExpr { - expr: col("b", &schema)?, - options: SortOptions { - descending: false, - nulls_first: true, - }, - }, - ] - .into(); - - let exec_plan = Arc::new(SortExec::new(sort_exprs, Arc::new(EmptyExec::new(schema)))); - - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = CustomConverterInterceptor::default(); - roundtrip_test_and_return(exec_plan, &ctx, &codec, &proto_converter)?; - - assert_eq!(*proto_converter.num_proto_exprs.read().unwrap(), 2); - assert_eq!(*proto_converter.num_physical_exprs.read().unwrap(), 2); - assert_eq!(*proto_converter.num_proto_plans.read().unwrap(), 2); - assert_eq!(*proto_converter.num_physical_plans.read().unwrap(), 2); - - Ok(()) -} - -#[test] -fn roundtrip_call_null_scalar_struct_dict() -> Result<()> { - let data_type = DataType::Struct(Fields::from(vec![Field::new( - "item", - DataType::Dictionary(Box::new(DataType::UInt32), Box::new(DataType::Utf8)), - true, - )])); - - let schema = Arc::new(Schema::new(vec![Field::new("a", data_type.clone(), true)])); - let scan = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let scalar = lit(ScalarValue::try_from(data_type)?); - let filter = Arc::new(FilterExec::try_new( - Arc::new(BinaryExpr::new(scalar, Operator::Eq, col("a", &schema)?)), - scan, - )?); - - roundtrip_test(filter) -} - -/// Create a [`DynamicFilterPhysicalExpr`] with child column expression "a" @ index 0. -fn make_dynamic_filter() -> Arc { - Arc::new(DynamicFilterPhysicalExpr::new( - vec![Arc::new(Column::new("a", 0)) as Arc], - lit(true), - )) as Arc -} - -/// Update a [`DynamicFilterPhysicalExpr`]'s children to support child schema "b" @ 0, "a" @ 1. -fn make_reassigned_dynamic_filter( - filter: Arc, -) -> Result<(Arc, Arc)> { - let schema = Arc::new(Schema::new(vec![ - Field::new("b", DataType::Int64, false), - Field::new("a", DataType::Int64, false), - ])); - let reassigned = reassign_expr_columns(filter, &schema)?; - Ok((schema, reassigned)) -} - -/// Extract the expression id from a [`PhysicalExpr`] proto. Populated by the -/// default serializer from `PhysicalExpr::expression_id`. -fn proto_expression_id(expr: &PhysicalExprNode) -> u64 { - expr.expr_id - .expect("expected PhysicalExprNode.expr_id to be populated") -} - -/// Roundtrip a single physical expression shaped like so: -/// -/// ```text -/// BinaryExpr(AND) -/// / \ -/// filter_expr_1 filter_expr_2 -/// ``` -/// -/// Returns filter_expr_1 and filter_expr_2 after deserialization. -fn roundtrip_dynamic_filter_expr_pair( - filter_expr_1: Arc, - filter_expr_2: Arc, - schema: Arc, -) -> Result<(Arc, Arc)> { - let pair_expr = Arc::new(BinaryExpr::new( - Arc::clone(&filter_expr_1), - Operator::And, - Arc::clone(&filter_expr_2), - )) as Arc; - - let codec = DefaultPhysicalExtensionCodec {}; - let converter = DeduplicatingProtoConverter {}; - let proto = converter.physical_expr_to_proto(&pair_expr, &codec)?; - let ctx = SessionContext::new(); - let task_ctx = ctx.task_ctx(); - let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), &codec); - let deserialized_expr = - converter.proto_to_physical_expr(&proto, &schema, &decode_ctx)?; - - let binary = deserialized_expr - .downcast_ref::() - .expect("Expected BinaryExpr"); - - Ok((Arc::clone(binary.left()), Arc::clone(binary.right()))) -} - -/// Roundtrip an execution plan shaped like so: -/// -/// ```text -/// FilterExec(dynamic_filter_1 on a@0) -/// ProjectionExec(a := Column("a", source_index)) -/// DataSourceExec -/// ParquetSource(predicate = dynamic_filter_2) -/// ``` -/// -/// `dynamic_filter_1` and `dynamic_filter_2` are the same dynamic filter, except with -/// different children. -/// -/// Returns -/// - `dynamic_filter_1` before serialization -/// - `dynamic_filter_2` before serialization -/// - `dynamic_filter_1` after serialization -/// - `dynamic_filter_2` after serialization -#[allow(clippy::type_complexity)] -fn roundtrip_dynamic_filter_plan_pair() -> Result<( - Arc, - Arc, - Arc, - Arc, -)> { - let filter_expr_1 = make_dynamic_filter(); - let (data_source_schema, filter_expr_2) = - make_reassigned_dynamic_filter(Arc::clone(&filter_expr_1))?; - let left_before = Arc::clone(&filter_expr_1); - let right_before = Arc::clone(&filter_expr_2); - let file_source = Arc::new( - ParquetSource::new(Arc::clone(&data_source_schema)) - .with_predicate(Arc::clone(&filter_expr_2)), - ); - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), file_source) - .with_file_groups(vec![FileGroup::new(vec![PartitionedFile::new( - "/path/to/file.parquet".to_string(), - 1024, - )])]) - .build(); - let data_source_exec = - DataSourceExec::from_data_source(scan_config) as Arc; - - let projection_exec = Arc::new(ProjectionExec::try_new( - vec![ProjectionExpr { - expr: Arc::new(Column::new("a", 1)) as Arc, - alias: "a".to_string(), - }], - data_source_exec, - )?) as Arc; - let filter_exec = Arc::new(FilterExec::try_new( - Arc::clone(&filter_expr_1), - projection_exec, - )?) as Arc; - - let codec = DefaultPhysicalExtensionCodec {}; - let converter = DeduplicatingProtoConverter {}; - let proto = converter.execution_plan_to_proto(&filter_exec, &codec)?; - - let ctx = SessionContext::new(); - let task_ctx = ctx.task_ctx(); - let decode_ctx = PhysicalPlanDecodeContext::new(task_ctx.as_ref(), &codec); - let deserialized_plan = converter.proto_to_execution_plan(&proto, &decode_ctx)?; - - let outer_filter = deserialized_plan - .downcast_ref::() - .expect("Expected outer FilterExec"); - let left_filter = Arc::clone(outer_filter.predicate()); - let projection = outer_filter.children()[0] - .downcast_ref::() - .expect("Expected ProjectionExec"); - let data_source = projection - .input() - .downcast_ref::() - .expect("Expected DataSourceExec"); - let scan_config = data_source - .data_source() - .downcast_ref::() - .expect("Expected FileScanConfig"); - let right_filter = scan_config - .file_source() - .filter() - .expect("Expected pushed-down predicate"); - - Ok((left_before, right_before, left_filter, right_filter)) -} - -/// Takes two [`DynamicFilterPhysicalExpr`] and asserts that updates to one are visible -/// via the other. This helps assert that referential integrity is maintained after -/// deserializing. -fn assert_dynamic_filter_update_is_visible( - left_filter: &Arc, - right_filter: &Arc, -) -> Result<()> { - let left_filter = left_filter - .downcast_ref::() - .expect("Expected dynamic filter"); - let right_filter = right_filter - .downcast_ref::() - .expect("Expected dynamic filter"); - - // Sanity check that the filters have the same generation. - let original_generation = left_filter.snapshot_generation(); - assert_eq!(original_generation, right_filter.snapshot_generation(),); - - left_filter.update(lit(123_i64))?; - - // Assert that both generations updated. - assert_eq!(original_generation + 1, right_filter.snapshot_generation(),); - assert_eq!( - left_filter.snapshot_generation(), - right_filter.snapshot_generation(), - ); - - // Ensure both filters have the updated expr. - let expected_current = r#"Literal { value: Int64(123), field: Field { name: "lit", data_type: Int64 } }"#; - assert_eq!(expected_current, format!("{:?}", left_filter.current()?),); - assert_eq!(expected_current, format!("{:?}", right_filter.current()?),); - - Ok(()) -} - -/// Extract the dynamic-filter predicate that was pushed down to the parquet -/// scan at the bottom of the plan tree. -fn parquet_source_predicate(child: &Arc) -> Arc { - let data_source = child - .downcast_ref::() - .expect("Child should be DataSourceExec"); - let (_, parquet_source) = data_source - .downcast_to_file_source::() - .expect("Should be ParquetSource"); - parquet_source - .filter() - .expect("ParquetSource should have a predicate after roundtrip") -} - -/// Assert that two dynamic filters are equal both structurally (Debug output) -/// and by identity (`expression_id`). -fn assert_dynamic_filters_equal( - expected: &Arc, - actual: &Arc, -) { - // Structural. - let expected_dbg = format!("{expected:?}"); - let actual_dbg = format!("{actual:?}"); - if expected_dbg == actual_dbg { - return; - } - - // Note that the `DeduplicatingDeserializer` routes every cache hit through - // `with_new_children`. This produces an equivalent expression, but with - // remapped children that are equal to the original. Handle that case here. - let rewritten = Arc::clone(expected) - .with_new_children(expected.children().iter().map(|c| Arc::clone(c)).collect()) - .expect("with_new_children on a dynamic filter should not fail"); - assert_eq!(format!("{rewritten:?}"), actual_dbg); -} - -// Two clones of a dynamic filter expression should be deduped to the exact same expression. -#[test] -fn test_dynamic_filter_roundtrip_dedupe() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - let filter_expr_1 = make_dynamic_filter(); - let filter_expr_2 = Arc::clone(&filter_expr_1); - - let (filter_expr_1_after_roundtrip, filter_expr_2_after_roundtrip) = - roundtrip_dynamic_filter_expr_pair( - Arc::clone(&filter_expr_1), - Arc::clone(&filter_expr_2), - schema, - )?; - - // Assert the filters are not modified during roundtrip. - assert_dynamic_filters_equal(&filter_expr_1, &filter_expr_1_after_roundtrip); - assert_dynamic_filters_equal(&filter_expr_2, &filter_expr_2_after_roundtrip); - assert_dynamic_filters_equal( - &filter_expr_1_after_roundtrip, - &filter_expr_2_after_roundtrip, - ); - - // Assert referential integrity. - assert_dynamic_filter_update_is_visible( - &filter_expr_1_after_roundtrip, - &filter_expr_2_after_roundtrip, - )?; - - Ok(()) -} - -/// Roundtrip test for an execution plan where there are multiple instances of a dynamic filter -/// with different children. -#[test] -fn test_dynamic_filter_plan_roundtrip_dedupe() -> Result<()> { - let ( - filter_expr_1, - filter_expr_2, - filter_expr_1_after_roundtrip, - filter_expr_2_after_roundtrip, - ) = roundtrip_dynamic_filter_plan_pair()?; - - // Assert the filters are not modified during roundtrip. - assert_dynamic_filters_equal(&filter_expr_1, &filter_expr_1_after_roundtrip); - assert_dynamic_filters_equal(&filter_expr_2, &filter_expr_2_after_roundtrip); - - // Assert referential integrity. - assert_dynamic_filter_update_is_visible( - &filter_expr_1_after_roundtrip, - &filter_expr_2_after_roundtrip, - )?; - - Ok(()) -} - -#[test] -fn test_dynamic_filter_expression_id_is_stable_between_serializations() -> Result<()> { - let filter_expr = make_dynamic_filter(); - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DeduplicatingProtoConverter {}; - - let proto1 = proto_converter.physical_expr_to_proto(&filter_expr, &codec)?; - let expr_id1 = proto_expression_id(&proto1); - - let proto2 = proto_converter.physical_expr_to_proto(&filter_expr, &codec)?; - let expr_id2 = proto_expression_id(&proto2); - - assert_eq!( - expr_id1, expr_id2, - "Expected the same dynamic filter expression id across serializations" - ); - - Ok(()) -} - -/// Tests that `lead` window function with offset and default value args -/// survives a protobuf round-trip. This is a regression test for a bug -/// where `expressions()` (used during serialization) returns only the -/// column expression for lead/lag, silently dropping the offset and -/// default value literal args. -#[test] -fn roundtrip_lead_with_default_value() -> Result<()> { - use datafusion::functions_window::lead_lag::lead_udwf; - - let field_a = Field::new("a", DataType::Int64, false); - let field_b = Field::new("b", DataType::Int64, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b])); - - // lead(a, 2, 42) — column a, offset 2, default value 42 - let lead_window = create_udwf_window_expr( - &lead_udwf(), - &[col("a", &schema)?, lit(2i64), lit(42i64)], - schema.as_ref(), - "test lead with default".to_string(), - false, - )?; - - let udwf_expr = Arc::new(StandardWindowExpr::new( - lead_window, - &[col("b", &schema)?], - &[PhysicalSortExpr { - expr: col("a", &schema)?, - options: SortOptions { - descending: false, - nulls_first: false, - }, - }], - Arc::new(WindowFrame::new(None)), - )); - - let input = Arc::new(EmptyExec::new(schema.clone())); - - roundtrip_test(Arc::new(BoundedWindowAggExec::try_new( - vec![udwf_expr], - input, - InputOrderMode::Sorted, - true, - )?)) -} - -/// Verify that ScalarSubqueryExpr nodes in the input plan are connected to the -/// same shared results container as ScalarSubqueryExec after a proto round-trip. -#[test] -fn roundtrip_scalar_subquery_exec() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - let results = ScalarSubqueryResults::new(1); - - // Build the input plan: a filter whose predicate references the - // scalar subquery result via ScalarSubqueryExpr. - let sq_expr = Arc::new(ScalarSubqueryExpr::new( - DataType::Int64, - true, - SubqueryIndex::new(0), - results.clone(), - )); - let predicate = binary(col("a", &schema)?, Operator::Eq, sq_expr, &schema)?; - let filter = - FilterExec::try_new(predicate, Arc::new(EmptyExec::new(schema.clone())))?; - - // Build a trivial subquery plan. - let subquery_plan = - Arc::new(EmptyExec::new(Arc::new(Schema::new(vec![Field::new( - "x", - DataType::Int64, - true, - )])))); - - let exec: Arc = Arc::new(ScalarSubqueryExec::new( - Arc::new(filter), - vec![ScalarSubqueryLink { - plan: subquery_plan, - index: SubqueryIndex::new(0), - }], - results, - )); - - // Perform the round-trip using DeduplicatingProtoConverter, which - // creates a DeduplicatingDeserializer that threads scalar subquery - // results through expression deserialization. - let codec = DefaultPhysicalExtensionCodec {}; - let converter = DeduplicatingProtoConverter {}; - let bytes = physical_plan_to_bytes_with_proto_converter( - Arc::clone(&exec), - &codec, - &converter, - )?; - let ctx = SessionContext::new(); - let deserialized = physical_plan_from_bytes_with_proto_converter( - bytes.as_ref(), - ctx.task_ctx().as_ref(), - &codec, - &converter, - )?; - - // Verify the deserialized ScalarSubqueryExec's results container is - // shared with the ScalarSubqueryExpr in the input plan. - let sq_exec = deserialized - .downcast_ref::() - .expect("expected ScalarSubqueryExec"); - let exec_results = sq_exec.results(); - - // Walk the input plan to find the ScalarSubqueryExpr and verify it - // points to the same results container. - let filter_exec = sq_exec - .input() - .downcast_ref::() - .expect("expected FilterExec"); - let binary_expr = filter_exec - .predicate() - .downcast_ref::() - .expect("expected BinaryExpr"); - let deserialized_sq_expr = binary_expr - .right() - .downcast_ref::() - .expect("expected ScalarSubqueryExpr"); - - assert!( - ScalarSubqueryResults::ptr_eq(exec_results, deserialized_sq_expr.results()), - "ScalarSubqueryExpr should share the same results container as ScalarSubqueryExec" - ); - Ok(()) -} - -/// Verify that nested ScalarSubqueryExec nodes deserialize with distinct -/// scoped results containers, and that each ScalarSubqueryExpr is wired to the -/// container for its own surrounding ScalarSubqueryExec. -#[test] -fn roundtrip_nested_scalar_subquery_exec_scopes_results() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - let subquery_schema = - Arc::new(Schema::new(vec![Field::new("x", DataType::Int64, true)])); - - let inner_results = ScalarSubqueryResults::new(1); - let inner_sq_expr = Arc::new(ScalarSubqueryExpr::new( - DataType::Int64, - true, - SubqueryIndex::new(0), - inner_results.clone(), - )); - let inner_predicate = - binary(col("a", &schema)?, Operator::Eq, inner_sq_expr, &schema)?; - let inner_filter = Arc::new(FilterExec::try_new( - inner_predicate, - Arc::new(EmptyExec::new(schema.clone())), - )?); - let inner_exec: Arc = Arc::new(ScalarSubqueryExec::new( - inner_filter, - vec![ScalarSubqueryLink { - plan: Arc::new(EmptyExec::new(subquery_schema.clone())), - index: SubqueryIndex::new(0), - }], - inner_results, - )); - - let outer_results = ScalarSubqueryResults::new(1); - let outer_sq_expr = Arc::new(ScalarSubqueryExpr::new( - DataType::Int64, - true, - SubqueryIndex::new(0), - outer_results.clone(), - )); - let outer_predicate = - binary(col("a", &schema)?, Operator::Eq, outer_sq_expr, &schema)?; - let outer_filter = Arc::new(FilterExec::try_new(outer_predicate, inner_exec)?); - let outer_exec: Arc = Arc::new(ScalarSubqueryExec::new( - outer_filter, - vec![ScalarSubqueryLink { - plan: Arc::new(EmptyExec::new(subquery_schema)), - index: SubqueryIndex::new(0), - }], - outer_results, - )); - - let bytes = datafusion_proto::bytes::physical_plan_to_bytes(Arc::clone(&outer_exec))?; - let ctx = SessionContext::new(); - let deserialized = datafusion_proto::bytes::physical_plan_from_bytes( - bytes.as_ref(), - ctx.task_ctx().as_ref(), - )?; - - let outer_exec = deserialized - .downcast_ref::() - .expect("expected outer ScalarSubqueryExec"); - let outer_results = outer_exec.results(); - let outer_filter = outer_exec - .input() - .downcast_ref::() - .expect("expected outer FilterExec"); - let outer_binary = outer_filter - .predicate() - .downcast_ref::() - .expect("expected outer BinaryExpr"); - let outer_sq_expr = outer_binary - .right() - .downcast_ref::() - .expect("expected outer ScalarSubqueryExpr"); - - let inner_exec = outer_filter - .input() - .downcast_ref::() - .expect("expected inner ScalarSubqueryExec"); - let inner_results = inner_exec.results(); - let inner_filter = inner_exec - .input() - .downcast_ref::() - .expect("expected inner FilterExec"); - let inner_binary = inner_filter - .predicate() - .downcast_ref::() - .expect("expected inner BinaryExpr"); - let inner_sq_expr = inner_binary - .right() - .downcast_ref::() - .expect("expected inner ScalarSubqueryExpr"); - - assert!( - ScalarSubqueryResults::ptr_eq(outer_results, outer_sq_expr.results()), - "outer ScalarSubqueryExpr should use outer ScalarSubqueryExec results" - ); - assert!( - ScalarSubqueryResults::ptr_eq(inner_results, inner_sq_expr.results()), - "inner ScalarSubqueryExpr should use inner ScalarSubqueryExec results" - ); - assert!( - !ScalarSubqueryResults::ptr_eq(outer_results, inner_results), - "nested ScalarSubqueryExec nodes should not share results containers" - ); - assert!( - !ScalarSubqueryResults::ptr_eq(outer_results, inner_sq_expr.results()), - "inner ScalarSubqueryExpr must not read from outer results" - ); - assert!( - !ScalarSubqueryResults::ptr_eq(inner_results, outer_sq_expr.results()), - "outer ScalarSubqueryExpr must not read from inner results" - ); - - Ok(()) -} - -/// Verify that the default physical plan bytes round-trip preserves executable -/// scalar subquery plans. -#[tokio::test] -async fn roundtrip_scalar_subquery_exec_with_default_converter_executes() -> Result<()> { - let ctx = SessionContext::new(); - let sql = "SELECT x + (SELECT max(y) FROM (VALUES (10), (20)) AS u(y)) AS s \ - FROM (VALUES (2), (1)) AS t(x) \ - ORDER BY s"; - - let initial_plan = ctx.sql(sql).await?.create_physical_plan().await?; - assert!( - format!("{initial_plan:?}").contains("ScalarSubqueryExec"), - "expected ScalarSubqueryExec in plan:\n{initial_plan:?}" - ); - - let bytes = - datafusion_proto::bytes::physical_plan_to_bytes(Arc::clone(&initial_plan))?; - let roundtripped = datafusion_proto::bytes::physical_plan_from_bytes( - bytes.as_ref(), - ctx.task_ctx().as_ref(), - )?; - assert!( - format!("{roundtripped:?}").contains("ScalarSubqueryExec"), - "expected ScalarSubqueryExec after roundtrip:\n{roundtripped:?}" - ); - - let batches = datafusion::physical_plan::common::collect( - roundtripped.execute(0, ctx.task_ctx())?, - ) - .await?; - datafusion::assert_batches_eq!( - &["+----+", "| s |", "+----+", "| 21 |", "| 22 |", "+----+",], - &batches - ); - - Ok(()) -} - -/// Test that a chain of the same operator (a AND b AND c) is linearized -/// and roundtrips correctly. -#[test] -fn roundtrip_binary_expr_chain_same_op() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Boolean, false); - let field_c = Field::new("c", DataType::Boolean, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c])); - let ab = binary( - col("a", &schema)?, - Operator::And, - col("b", &schema)?, - &schema, - )?; - let abc = binary(ab, Operator::And, col("c", &schema)?, &schema)?; - roundtrip_test(Arc::new(FilterExec::try_new( - abc, - Arc::new(EmptyExec::new(schema)), - )?)) -} - -/// Test that mixed operators (a AND b OR c) are NOT linearized together — -/// only chains of the same operator are flattened. -#[test] -fn roundtrip_binary_expr_mixed_ops() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Boolean, false); - let field_c = Field::new("c", DataType::Boolean, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c])); - // (a AND b) OR c — AND and OR are different operators, so linearization stops - let a_and_b = binary( - col("a", &schema)?, - Operator::And, - col("b", &schema)?, - &schema, - )?; - let expr = binary(a_and_b, Operator::Or, col("c", &schema)?, &schema)?; - roundtrip_test(Arc::new(FilterExec::try_new( - expr, - Arc::new(EmptyExec::new(schema)), - )?)) -} - -/// Test that a deeply nested chain of AND expressions (like many WHERE conditions) -/// roundtrips correctly. This is the scenario from issue #18602. -#[test] -fn roundtrip_binary_expr_deeply_nested_and_chain() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let schema = Arc::new(Schema::new(vec![field_a])); - - // Build a chain: a AND a AND a AND ... (100 times) - let col_a = col("a", &schema)?; - let mut expr = Arc::clone(&col_a); - for _ in 0..99 { - expr = binary(expr, Operator::And, Arc::clone(&col_a), &schema)?; - } - - roundtrip_test(Arc::new(FilterExec::try_new( - expr, - Arc::new(EmptyExec::new(schema)), - )?)) -} - -/// Test that a deeply nested chain of OR expressions roundtrips correctly. -#[test] -fn roundtrip_binary_expr_deeply_nested_or_chain() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let schema = Arc::new(Schema::new(vec![field_a])); - - let col_a = col("a", &schema)?; - let mut expr = Arc::clone(&col_a); - for _ in 0..99 { - expr = binary(expr, Operator::Or, Arc::clone(&col_a), &schema)?; - } - - roundtrip_test(Arc::new(FilterExec::try_new( - expr, - Arc::new(EmptyExec::new(schema)), - )?)) -} - -/// Test that alternating AND/OR operators produce correct results — -/// each sub-chain gets linearized independently. -#[test] -fn roundtrip_binary_expr_alternating_and_or() -> Result<()> { - let field_a = Field::new("a", DataType::Boolean, false); - let field_b = Field::new("b", DataType::Boolean, false); - let field_c = Field::new("c", DataType::Boolean, false); - let field_d = Field::new("d", DataType::Boolean, false); - let schema = Arc::new(Schema::new(vec![field_a, field_b, field_c, field_d])); - - // (a AND b) OR (c AND d) - let a_and_b = binary( - col("a", &schema)?, - Operator::And, - col("b", &schema)?, - &schema, - )?; - let c_and_d = binary( - col("c", &schema)?, - Operator::And, - col("d", &schema)?, - &schema, - )?; - let expr = binary(a_and_b, Operator::Or, c_and_d, &schema)?; - - roundtrip_test(Arc::new(FilterExec::try_new( - expr, - Arc::new(EmptyExec::new(schema)), - )?)) -} - -/// Verify that the linearized proto format has a flat operands list -/// rather than deeply nested l/r fields. -#[test] -fn test_linearization_produces_flat_operands() -> Result<()> { - // Build: a AND a AND a AND a (4 operands, 3 levels of nesting) - let col_a: Arc = Arc::new(Column::new("a", 0)); - let expr: Arc = Arc::new(BinaryExpr::new( - Arc::new(BinaryExpr::new( - Arc::new(BinaryExpr::new( - Arc::clone(&col_a), - Operator::And, - Arc::clone(&col_a), - )), - Operator::And, - Arc::clone(&col_a), - )), - Operator::And, - Arc::clone(&col_a), - )); - - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - let proto = proto_converter.physical_expr_to_proto(&expr, &codec)?; - - // The top-level should use the operands field with 4 entries - match &proto.expr_type { - Some(protobuf::physical_expr_node::ExprType::BinaryExpr(b)) => { - assert!( - b.l.is_none(), - "l should be None when using linearized operands" - ); - assert!( - b.r.is_none(), - "r should be None when using linearized operands" - ); - assert_eq!( - b.operands.len(), - 4, - "Expected 4 linearized operands for a AND a AND a AND a" - ); - assert_eq!(b.op, "And"); - } - other => panic!("Expected BinaryExpr, got {other:?}"), - } - - Ok(()) -} - -/// Test that linearization stops when encountering a different operator. -/// For (a AND b) OR c, only the top-level OR should be represented, and -/// the left-hand AND subtree should be a separate nested BinaryExpr. -#[test] -fn test_linearization_stops_at_different_op() -> Result<()> { - // (a AND b) OR c - let a_and_b: Arc = Arc::new(BinaryExpr::new( - Arc::new(Column::new("a", 0)), - Operator::And, - Arc::new(Column::new("b", 1)), - )); - let expr: Arc = Arc::new(BinaryExpr::new( - a_and_b, - Operator::Or, - Arc::new(Column::new("c", 2)), - )); - - let codec = DefaultPhysicalExtensionCodec {}; - let proto_converter = DefaultPhysicalProtoConverter {}; - let proto = proto_converter.physical_expr_to_proto(&expr, &codec)?; - - // The top-level OR should have only 2 operands (can't linearize through AND) - match &proto.expr_type { - Some(protobuf::physical_expr_node::ExprType::BinaryExpr(b)) => { - assert_eq!( - b.operands.len(), - 2, - "Expected 2 operands for (a AND b) OR c" - ); - assert_eq!(b.op, "Or"); - // The first operand should be a nested AND BinaryExpr - match &b.operands[0].expr_type { - Some(protobuf::physical_expr_node::ExprType::BinaryExpr(inner)) => { - assert_eq!(inner.op, "And"); - assert_eq!(inner.operands.len(), 2); - } - other => panic!("Expected inner BinaryExpr(AND), got {other:?}"), - } - } - other => panic!("Expected BinaryExpr, got {other:?}"), - } - - Ok(()) -} - -/// Create a DataSourceExec backed by a ParquetSource that accepts filter pushdown, -/// along with a ConfigOptions that enables all dynamic filter pushdown options. -fn datasource_for_dynamic_filter_pushdown( - schema: &Arc, -) -> (Arc, ConfigOptions) { - let mut parquet_options = TableParquetOptions::new(); - parquet_options.global.pushdown_filters = true; - let source = Arc::new( - ParquetSource::new(Arc::clone(schema)) - .with_table_parquet_options(parquet_options), - ); - let scan_config = - FileScanConfigBuilder::new(ObjectStoreUrl::local_filesystem(), source) - .with_file(PartitionedFile::new("/path/to/file.parquet", 1024)) - .build(); - - let mut config = ConfigOptions::default(); - config.execution.parquet.pushdown_filters = true; - config.optimizer.enable_join_dynamic_filter_pushdown = true; - config.optimizer.enable_aggregate_dynamic_filter_pushdown = true; - config.optimizer.enable_topk_dynamic_filter_pushdown = true; - - (DataSourceExec::from_data_source(scan_config), config) -} - -/// Test that plan containing a HashJoinExec with dynamic filter pushdown -/// can be serialized and deserialized while preserving references to the dynamic filter. -#[test] -fn test_hash_join_with_dynamic_filter_roundtrip() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("col", DataType::Int64, false)])); - - let left_child = Arc::new(EmptyExec::new(Arc::clone(&schema))); - let (right_child, config) = datasource_for_dynamic_filter_pushdown(&schema); - - let on: Vec<(Arc, Arc)> = vec![( - Arc::new(Column::new("col", 0)), - Arc::new(Column::new("col", 0)), - )]; - - let hash_join = Arc::new(HashJoinExec::try_new( - left_child, - right_child, - on, - None, - &JoinType::Inner, - None, - PartitionMode::CollectLeft, - NullEquality::NullEqualsNothing, - false, - )?) as Arc; - - // Run the optimizer rule for filter pushdown. - let optimizer = FilterPushdown::new_post_optimization(); - let plan = optimizer.optimize(hash_join, &config)?; - - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let converter = DeduplicatingProtoConverter {}; - let deserialized = roundtrip_test_and_return(plan, &ctx, &codec, &converter)?; - - // Extract the deserialized HashJoinExec and its dynamic filter. - let deserialized_join = deserialized - .downcast_ref::() - .expect("Should be HashJoinExec"); - let deserialized_hash_join_df = deserialized_join - .dynamic_filter_expr() - .expect("HashJoinExec should have a dynamic filter after roundtrip"); - - // Extract the dynamic filter pushed down to the probe side's ParquetSource. - let deserialized_predicate = parquet_source_predicate(deserialized_join.right()); - - // The HashJoinExec's dynamic filter and the probe side's predicate should - // refer to the same underlying expression. - let plan_df: Arc = deserialized_hash_join_df.clone(); - assert_dynamic_filters_equal(&plan_df, &deserialized_predicate); - assert_dynamic_filter_update_is_visible(&plan_df, &deserialized_predicate)?; - - Ok(()) -} - -/// returns a SessionContext with an empty `netflow` table registered -fn netflow_context() -> Result { - let ctx = SessionContext::new(); - let schema = Arc::new(Schema::new(vec![ - Field::new("dst_geo_country_name", DataType::Utf8, true), - Field::new("dst_geo_city_name", DataType::Utf8, true), - Field::new("packets", DataType::UInt64, true), - Field::new("src_addr", DataType::Utf8, true), - Field::new("dst_addr", DataType::Utf8, true), - ])); - - ctx.register_table("netflow", Arc::new(EmptyTable::new(schema)))?; - - Ok(ctx) -} - -/// Regression test for issue #18602: -/// https://github.com/apache/datafusion/issues/18602 -/// -/// The physical filter expression here contains a long chain of `AND` predicates. -/// Before linearizing `PhysicalBinaryExprNode`, encoding then decoding the protobuf -/// could fail with `DecodeError: recursion limit reached`. -#[tokio::test] -async fn roundtrip_issue_18602_complex_filter_decode_recursion() -> Result<()> { - let ctx = netflow_context()?; - let sql = "SELECT \ - dst_geo_country_name AS x_axis_1, \ - dst_geo_city_name AS x_axis_2, \ - sum(packets) AS y_axis_1 \ - FROM netflow \ - WHERE dst_geo_country_name IS NOT NULL \ - AND src_addr NOT LIKE '10.201.%' \ - AND dst_addr NOT LIKE '10.201.%' \ - AND src_addr NOT LIKE '10.202.%' \ - AND dst_addr NOT LIKE '10.202.%' \ - AND src_addr NOT LIKE '10.203.%' \ - AND dst_addr NOT LIKE '10.203.%' \ - AND src_addr NOT LIKE '10.204.%' \ - AND dst_addr NOT LIKE '10.204.%' \ - AND src_addr NOT LIKE '172.16.186.%' \ - AND dst_addr NOT LIKE '172.16.186.%' \ - AND src_addr NOT LIKE '172.16.187.%' \ - AND dst_addr NOT LIKE '172.16.187.%' \ - AND src_addr NOT LIKE '172.16.188.%' \ - AND dst_addr NOT LIKE '172.16.188.%' \ - AND src_addr NOT LIKE '10.102.45.%' \ - AND dst_addr NOT LIKE '10.102.45.%' \ - AND src_addr NOT LIKE '172.25.210.%' \ - AND dst_addr NOT LIKE '172.25.210.%' \ - AND src_addr NOT LIKE '172.25.211.%' \ - AND dst_addr NOT LIKE '172.25.211.%' \ - AND src_addr NOT LIKE '141.226.101.%' \ - AND dst_addr NOT LIKE '141.226.101.%' \ - AND src_addr NOT LIKE '167.86.40.%' \ - AND dst_addr NOT LIKE '167.86.40.%' \ - AND src_addr NOT LIKE '66.22.38.%' \ - AND dst_addr NOT LIKE '66.22.38.%' \ - AND src_addr != '168.143.191.55' \ - AND dst_addr != '168.143.191.55' \ - AND src_addr != '82.112.107.142' \ - AND dst_addr != '82.112.107.142' \ - AND src_addr != '20.76.39.176' \ - AND dst_addr != '20.76.39.176' \ - AND src_addr != '162.159.129.83' \ - AND dst_addr != '162.159.129.83' \ - AND src_addr != '34.201.223.155' \ - AND dst_addr != '34.201.223.155' \ - AND src_addr != '34.201.223.156' \ - AND dst_addr != '34.201.223.156' \ - AND src_addr != '34.201.223.157' \ - AND dst_addr != '34.201.223.157' \ - AND src_addr != '134.201.223.157' \ - AND dst_addr != '134.201.223.157' \ - AND src_addr != '341.201.223.157' \ - AND dst_addr != '341.201.223.157' \ - GROUP BY x_axis_1, x_axis_2 \ - ORDER BY y_axis_1 DESC \ - LIMIT 20"; - - roundtrip_test_sql_with_context(sql, &ctx).await -} - -/// Test that plan containing a AggregateExec with dynamic filter pushdown -/// can be serialized and deserialized while preserving references to the dynamic filter. -#[test] -fn test_aggregate_with_dynamic_filter_roundtrip() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - let col_a: Arc = Arc::new(Column::new("a", 0)); - - let (child, config) = datasource_for_dynamic_filter_pushdown(&schema); - - let agg = Arc::new(AggregateExec::try_new( - AggregateMode::Partial, - PhysicalGroupBy::new_single(vec![]), - vec![ - AggregateExprBuilder::new( - datafusion::functions_aggregate::min_max::min_udaf(), - vec![Arc::clone(&col_a)], - ) - .schema(Arc::clone(&schema)) - .alias("min_a") - .build() - .map(Arc::new)?, - ], - vec![None], - child, - Arc::clone(&schema), - )?) as Arc; - - // Run the optimizer rule for filter pushdown. - let optimizer = FilterPushdown::new_post_optimization(); - let plan = optimizer.optimize(agg, &config)?; - - // Roundtrip with deduplication. - // - // Note: We don't use `roundtrip_test_and_return` here because there's a - // pre-existing issue with PhysicalGroupBy serialization where empty groups - // `[[]]` become `[]` after roundtrip. This behavior is unrelated to this test. - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let converter = DeduplicatingProtoConverter {}; - let bytes = physical_plan_to_bytes_with_proto_converter( - Arc::clone(&plan), - &codec, - &converter, - )?; - let deserialized = physical_plan_from_bytes_with_proto_converter( - bytes.as_ref(), - ctx.task_ctx().as_ref(), - &codec, - &converter, - )?; - - // Extract the deserialized AggregateExec and its dynamic filter. - let deserialized_agg = deserialized - .downcast_ref::() - .expect("Should be AggregateExec"); - let deserialized_agg_df = deserialized_agg - .dynamic_filter_expr() - .expect("AggregateExec should have a dynamic filter after roundtrip"); - - // Extract the dynamic filter pushed down to the child ParquetSource. - let deserialized_predicate = parquet_source_predicate(deserialized_agg.input()); - - // The AggregateExec's dynamic filter and the child's predicate should - // refer to the same underlying expression. - let plan_df: Arc = deserialized_agg_df.clone(); - assert_dynamic_filters_equal(&plan_df, &deserialized_predicate); - assert_dynamic_filter_update_is_visible(&plan_df, &deserialized_predicate)?; - - Ok(()) -} - -/// Test that plan containing a SortExec with dynamic filter pushdown -/// can be serialized and deserialized while preserving references to the dynamic filter. -#[test] -fn test_sort_topk_with_dynamic_filter_roundtrip() -> Result<()> { - let schema = Arc::new(Schema::new(vec![Field::new("a", DataType::Int64, false)])); - let col_a: Arc = Arc::new(Column::new("a", 0)); - - let (child, config) = datasource_for_dynamic_filter_pushdown(&schema); - - let sort = Arc::new( - SortExec::new( - LexOrdering::new(vec![PhysicalSortExpr { - expr: Arc::clone(&col_a), - options: SortOptions::default(), - }]) - .unwrap(), - child, - ) - .with_fetch(Some(10)), - ) as Arc; - - // Verify the optimizer kept the dynamic filter on the SortExec. - let optimizer = FilterPushdown::new_post_optimization(); - let plan = optimizer.optimize(sort, &config)?; - - // Roundtrip with deduplication. - // - // Note: We don't use `roundtrip_test_and_return` here because - // `DeduplicatingDeserializer` rewrites cache hits via `with_new_children`, - // which sets `remapped_children: Some(...)` on the second encounter of a - // shared `DynamicFilterPhysicalExpr`. SortExec's `Debug` includes its - // dynamic filter, so the original-vs-deserialized structural equality check - // would fail purely on this artifact. - let ctx = SessionContext::new(); - let codec = DefaultPhysicalExtensionCodec {}; - let converter = DeduplicatingProtoConverter {}; - let bytes = physical_plan_to_bytes_with_proto_converter( - Arc::clone(&plan), - &codec, - &converter, - )?; - let deserialized = physical_plan_from_bytes_with_proto_converter( - bytes.as_ref(), - ctx.task_ctx().as_ref(), - &codec, - &converter, - )?; - - // Extract the deserialized SortExec and its dynamic filter. - let deserialized_sort = deserialized - .downcast_ref::() - .expect("Should be SortExec"); - let deserialized_sort_df = deserialized_sort - .dynamic_filter_expr() - .expect("SortExec should have a dynamic filter after roundtrip"); - - // Extract the dynamic filter pushed down to the child ParquetSource. - let deserialized_predicate = parquet_source_predicate(deserialized_sort.input()); - - // The SortExec's dynamic filter and the child's predicate should - // refer to the same underlying expression. - let plan_df: Arc = deserialized_sort_df; - assert_dynamic_filters_equal(&plan_df, &deserialized_predicate); - assert_dynamic_filter_update_is_visible(&plan_df, &deserialized_predicate)?; - - Ok(()) -} diff --git a/datafusion/proto/tests/cases/serialize.rs b/datafusion/proto/tests/cases/serialize.rs index 850fd42ce131b..a0a917e3239c2 100644 --- a/datafusion/proto/tests/cases/serialize.rs +++ b/datafusion/proto/tests/cases/serialize.rs @@ -22,14 +22,18 @@ use arrow::datatypes::{DataType, Field}; use datafusion::execution::FunctionRegistry; use datafusion::prelude::SessionContext; -use datafusion_expr::expr::Placeholder; -use datafusion_expr::{ColumnarValue, col, create_udf, lit}; +use datafusion_common::ScalarValue; +use datafusion_expr::expr::{HigherOrderFunction, LambdaVariable, Placeholder}; +use datafusion_expr::{ColumnarValue, HigherOrderUDF, col, create_udf, lambda, lit}; use datafusion_expr::{Expr, Volatility}; use datafusion_functions::string; use datafusion_proto::bytes::Serializeable; use datafusion_proto::logical_plan::DefaultLogicalExtensionCodec; +use datafusion_proto::logical_plan::from_proto::parse_expr; use datafusion_proto::logical_plan::to_proto::serialize_expr; +use crate::cases::MyHigherOrderUDF; + #[test] #[should_panic( expected = "Error decoding expr as protobuf: failed to decode Protobuf message" @@ -298,3 +302,92 @@ fn test_expression_serialization_roundtrip() { name.split('(').next().unwrap().to_string() } } + +/// return a `SessionContext` with `MyHigherOrderUDF` registered as a higher-order UDF +fn context_with_higher_order_function() -> SessionContext { + let ctx = SessionContext::new(); + let hof = Arc::new(HigherOrderUDF::new_from_impl(MyHigherOrderUDF::new( + "payload".to_string(), + ))); + ctx.register_higher_order_function(hof); + ctx +} + +fn dummy_higher_order_function_call(hof: Arc) -> Expr { + let list = ScalarValue::List(ScalarValue::new_list_nullable( + &[ScalarValue::Int32(Some(1))], + &DataType::Int32, + )); + let lambda_var_with_field = Expr::LambdaVariable(LambdaVariable::new( + "x".to_string(), + Some(Arc::new(Field::new("x", DataType::Int32, true))), + )); + let lambda_var_without_field = + Expr::LambdaVariable(LambdaVariable::new("x".into(), None)); + let lambda = lambda(["x"], lambda_var_with_field + lambda_var_without_field); + Expr::HigherOrderFunction(HigherOrderFunction::new( + hof, + vec![Expr::Literal(list, None), lambda], + )) +} + +#[test] +fn hof_roundtrip_with_registry() { + let ctx = context_with_higher_order_function(); + let hof = ctx + .higher_order_function("higher_order_udf") + .expect("could not find higher order udf"); + + let expr = dummy_higher_order_function_call(hof); + + let bytes = expr.to_bytes().unwrap(); + let deserialized_expr = + Expr::from_bytes_with_ctx(&bytes, ctx.task_ctx().as_ref()).unwrap(); + + assert_eq!(expr, deserialized_expr); +} + +#[test] +#[should_panic( + expected = "LogicalExtensionCodec is not provided for higher order function higher_order_udf" +)] +fn hof_roundtrip_without_registry() { + let ctx = context_with_higher_order_function(); + let hof = ctx + .higher_order_function("higher_order_udf") + .expect("could not find higher order udf"); + + let expr = dummy_higher_order_function_call(hof); + + let bytes = expr.to_bytes().unwrap(); + Expr::from_bytes(&bytes).unwrap(); +} + +#[test] +fn test_higher_order_serialization_roundtrip() { + let ctx = SessionContext::new(); + let list = ScalarValue::List(ScalarValue::new_list_nullable( + &[ScalarValue::Int32(Some(1))], + &DataType::Int32, + )); + let lambda_var_with_field = Expr::LambdaVariable(LambdaVariable::new( + "x".to_string(), + Some(Arc::new(Field::new("x", DataType::Int32, true))), + )); + let lambda_var_without_field = + Expr::LambdaVariable(LambdaVariable::new("x".into(), None)); + let lambda = lambda(["x"], lambda_var_with_field + lambda_var_without_field); + let args = vec![Expr::Literal(list, None), lambda]; + + for function in datafusion::functions_nested::all_default_higher_order_functions() { + let expr = + Expr::HigherOrderFunction(HigherOrderFunction::new(function, args.clone())); + + let extension_codec = DefaultLogicalExtensionCodec {}; + let proto = serialize_expr(&expr, &extension_codec).unwrap(); + let deserialize = + parse_expr(&proto, ctx.task_ctx().as_ref(), &extension_codec).unwrap(); + + assert_eq!(deserialize, expr); + } +} diff --git a/datafusion/proto/tests/cases/stack_safety.rs b/datafusion/proto/tests/cases/stack_safety.rs new file mode 100644 index 0000000000000..5caf4119a7186 --- /dev/null +++ b/datafusion/proto/tests/cases/stack_safety.rs @@ -0,0 +1,99 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::process::Command; +use std::sync::Arc; + +use datafusion_common::DFSchema; +use datafusion_expr::logical_plan::{EmptyRelation, LogicalPlan, LogicalPlanBuilder}; +use datafusion_proto::bytes::logical_plan_to_bytes; + +const CHILD_ENV: &str = "DATAFUSION_PROTO_ISSUE_23823_CHILD"; +const ALIAS_DEPTH_ENV: &str = "DATAFUSION_PROTO_ISSUE_23823_ALIAS_DEPTH"; +const TWO_MIB_TEST_NAME: &str = + "cases::stack_safety::logical_plan_serialization_fits_a_two_mib_stack"; +#[cfg(feature = "recursive_protection")] +const GROWABLE_STACK_TEST_NAME: &str = + "cases::stack_safety::deeply_nested_logical_plan_serialization_uses_a_growable_stack"; + +fn deeply_aliased_plan(alias_depth: usize) -> LogicalPlan { + let mut plan = LogicalPlan::EmptyRelation(EmptyRelation { + produce_one_row: false, + schema: Arc::new(DFSchema::empty()), + }); + + for level in 0..alias_depth { + plan = LogicalPlanBuilder::from(plan) + .alias(format!("level_{level}")) + .unwrap() + .build() + .unwrap(); + } + + plan +} + +fn serialize_on_two_mib_stack(alias_depth: usize) { + let plan = deeply_aliased_plan(alias_depth); + std::thread::Builder::new() + .name("two-megabyte-stack".into()) + .stack_size(2 * 1024 * 1024) + .spawn(move || logical_plan_to_bytes(&plan).unwrap()) + .unwrap() + .join() + .unwrap(); +} + +fn run_in_child(test_name: &str, alias_depth: usize) { + if std::env::var_os(CHILD_ENV).is_some() { + let alias_depth = std::env::var(ALIAS_DEPTH_ENV).unwrap().parse().unwrap(); + serialize_on_two_mib_stack(alias_depth); + return; + } + + // A native stack overflow aborts the process. Re-run this exact test in a + // child process so a regression produces a normal test failure. + let output = Command::new(std::env::current_exe().unwrap()) + .args(["--exact", test_name, "--nocapture"]) + .env(CHILD_ENV, "1") + .env(ALIAS_DEPTH_ENV, alias_depth.to_string()) + .output() + .unwrap(); + + assert!( + output.status.success(), + "child process failed with status {}\nstdout:\n{}\nstderr:\n{}", + output.status, + String::from_utf8_lossy(&output.stdout), + String::from_utf8_lossy(&output.stderr), + ); +} + +#[test] +fn logical_plan_serialization_fits_a_two_mib_stack() { + // Ten aliases reproduce #23823. Use 100 to provide a safety margin while + // verifying the dispatcher reduction without runtime stack growth. + run_in_child(TWO_MIB_TEST_NAME, 100); +} + +#[cfg(feature = "recursive_protection")] +#[test] +fn deeply_nested_logical_plan_serialization_uses_a_growable_stack() { + // This depth exceeds the 2 MiB thread stack without recursive protection, + // exercising the `recursive` stack-growth checkpoint. + run_in_child(GROWABLE_STACK_TEST_NAME, 2_000); +} diff --git a/datafusion/proto/tests/proto_integration.rs b/datafusion/proto/tests/proto_integration.rs index 6ce41c9de71a8..07a72f13ffb82 100644 --- a/datafusion/proto/tests/proto_integration.rs +++ b/datafusion/proto/tests/proto_integration.rs @@ -15,5 +15,9 @@ // specific language governing permissions and limitations // under the License. +// Test helpers take owned values for convenience, matching the `#![cfg_attr(test, ...)]` +// exemption the DataFusion crates apply to their own unit tests. +#![cfg_attr(test, allow(clippy::needless_pass_by_value))] + /// Run all tests that are found in the `cases` directory mod cases; diff --git a/datafusion/pruning/src/file_pruner.rs b/datafusion/pruning/src/file_pruner.rs index f850e0c0114fb..661832915c40f 100644 --- a/datafusion/pruning/src/file_pruner.rs +++ b/datafusion/pruning/src/file_pruner.rs @@ -22,7 +22,8 @@ use std::sync::Arc; use arrow::datatypes::{FieldRef, SchemaRef}; use datafusion_common::{Result, internal_datafusion_err, pruning::PrunableStatistics}; use datafusion_datasource::PartitionedFile; -use datafusion_physical_expr_common::physical_expr::{PhysicalExpr, snapshot_generation}; +use datafusion_physical_expr::DynamicFilterTracking; +use datafusion_physical_expr_common::physical_expr::PhysicalExpr; use datafusion_physical_plan::metrics::Count; use log::debug; @@ -34,8 +35,14 @@ use crate::build_pruning_predicate; /// which substitutes partition column references with their literal values before /// the predicate reaches this pruner. pub struct FilePruner { - predicate_generation: Option, predicate: Arc, + /// Tracks the dynamic filters inside `predicate` so we only rebuild the + /// pruning predicate when one of them has actually moved. + tracking: DynamicFilterTracking, + /// Whether [`Self::should_prune`] has built+evaluated the pruning predicate + /// at least once. The first check always runs; subsequent checks only run + /// when a watched dynamic filter changed. + checked_once: bool, /// Schema used for pruning (the logical file schema). file_schema: SchemaRef, file_stats_pruning: PrunableStatistics, @@ -69,42 +76,72 @@ impl FilePruner { }) } - /// Create a new file pruner if statistics are available. - /// Returns None if this file does not have statistics. + /// Create a file pruner for this file, or `None` when pruning it cannot + /// help. + /// + /// Returns `None` when the file has no statistics struct to evaluate a + /// pruning predicate against, or when the predicate is purely static and the + /// file has no usable column statistics — in that case planning already did + /// everything such a pruner could. A predicate carrying a dynamic filter is + /// always accepted (given a statistics struct), since it may prune via + /// partition-value folding even without column statistics. pub fn try_new( predicate: Arc, file_schema: &SchemaRef, partitioned_file: &PartitionedFile, predicate_creation_errors: Count, ) -> Option { + // A pruning predicate is evaluated against a statistics struct, so one + // must exist (its columns may all be `Absent`). let file_stats = partitioned_file.statistics.as_ref()?; + let tracking = DynamicFilterTracking::classify(&predicate); + // Only build a pruner when it could prune something planning didn't + // already: the file has real column statistics, or the predicate carries + // a dynamic filter (whose value, or folded partition columns, can prune + // even without column statistics). For a purely static predicate with no + // usable stats there is nothing to gain. + if !partitioned_file.has_statistics() && !tracking.contains_dynamic_filter() { + return None; + } let file_stats_pruning = PrunableStatistics::new(vec![file_stats.clone()], Arc::clone(file_schema)); Some(Self { - predicate_generation: None, predicate, + tracking, + checked_once: false, file_schema: Arc::clone(file_schema), file_stats_pruning, predicate_creation_errors, }) } + /// Returns `true` if this pruner watches a dynamic filter that can still + /// change, meaning [`Self::should_prune`] is worth re-checking as the scan + /// progresses. When `false`, the predicate is effectively static for the + /// remainder of the scan and the caller can avoid wrapping the stream in a + /// per-batch re-pruning adapter. + pub fn is_watching(&self) -> bool { + matches!(self.tracking, DynamicFilterTracking::Watching(_)) + } + pub fn should_prune(&mut self) -> Result { - // Check if the predicate has changed since last invocation by tracking - // its "generation". Dynamic filter expressions can change their values - // during query execution, so we use generation tracking to detect when - // the predicate has been updated and needs to be rebuilt. + // Building the pruning predicate is expensive (it involves expression + // analysis), so we only do it on the first check and whenever a dynamic + // filter inside the predicate has actually moved. // - // If the generation hasn't changed, we can skip rebuilding the pruning - // predicate, which is an expensive operation involving expression analysis. - let new_generation = snapshot_generation(&self.predicate); - if let Some(current_generation) = self.predicate_generation.as_mut() { - if *current_generation == new_generation { - return Ok(false); - } - *current_generation = new_generation; + // Dynamic filter expressions can change their values during query + // execution; `DynamicFilterTracking` watches the still-incomplete + // filters and reports a change at most once per update. A purely static + // predicate (or one whose dynamic filters have all completed) is checked + // exactly once. + let should_build = if self.checked_once { + self.tracking.watcher().is_some_and(|w| w.changed()) } else { - self.predicate_generation = Some(new_generation); + self.checked_once = true; + true + }; + if !should_build { + return Ok(false); } let pruning_predicate = build_pruning_predicate( Arc::clone(&self.predicate), diff --git a/datafusion/pruning/src/lib.rs b/datafusion/pruning/src/lib.rs index be17f29eaafa0..2b334d2847980 100644 --- a/datafusion/pruning/src/lib.rs +++ b/datafusion/pruning/src/lib.rs @@ -22,6 +22,6 @@ mod pruning_predicate; pub use file_pruner::FilePruner; pub use pruning_predicate::{ - PredicateRewriter, PruningPredicate, PruningStatistics, RequiredColumns, - UnhandledPredicateHook, build_pruning_predicate, + MAX_IN_LIST_SIZE, PredicateRewriter, PruningPredicate, PruningPredicateBuilder, + PruningStatistics, RequiredColumns, UnhandledPredicateHook, build_pruning_predicate, }; diff --git a/datafusion/pruning/src/pruning_predicate.rs b/datafusion/pruning/src/pruning_predicate.rs index 28d4fe9028760..3a63451495e4c 100644 --- a/datafusion/pruning/src/pruning_predicate.rs +++ b/datafusion/pruning/src/pruning_predicate.rs @@ -36,11 +36,14 @@ use log::{debug, trace}; use datafusion_common::error::Result; use datafusion_common::tree_node::{TransformedResult, TreeNodeRecursion}; -use datafusion_common::{Column, DFSchema, assert_eq_or_internal_err}; +use datafusion_common::{ + _internal_datafusion_err, Column, DFSchema, assert_eq_or_internal_err, +}; use datafusion_common::{ ScalarValue, internal_datafusion_err, plan_datafusion_err, plan_err, tree_node::{Transformed, TreeNode}, }; +use datafusion_expr_common::casts::try_cast_literal_to_type; use datafusion_expr_common::operator::Operator; use datafusion_physical_expr::utils::{Guarantee, LiteralGuarantee}; use datafusion_physical_expr::{PhysicalExprRef, expressions as phys_expr}; @@ -105,7 +108,7 @@ use datafusion_physical_plan::{ColumnarValue, PhysicalExpr}; /// C: true (rows might match x = 5) /// ``` /// -/// See [`PruningPredicate::try_new`] and [`PruningPredicate::prune`] for more information. +/// See [`PruningPredicateBuilder`] and [`PruningPredicate::prune`] for more information. /// /// # Background /// @@ -387,18 +390,147 @@ pub fn build_pruning_predicate( file_schema: &SchemaRef, predicate_creation_errors: &Count, ) -> Option> { - match PruningPredicate::try_new(predicate, Arc::clone(file_schema)) { - Ok(pruning_predicate) => { - if !pruning_predicate.always_true() { - return Some(Arc::new(pruning_predicate)); + PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(file_schema)) + .with_error_counter(predicate_creation_errors) + .build(predicate) +} + +/// Builder for a [`PruningPredicate`]. Groups optional configuration — +/// `IN (...)` rewrite cap, error counter — so future additions do not +/// churn the top-level API. +/// +/// The two entry points are: +/// - [`Self::build`]: convenience for scan sites that already track a +/// `predicate_creation_errors` counter. Returns `Some(Arc<..>)` when the +/// resulting predicate can actually prune, `None` when it is trivially +/// true or when construction failed (in which case the error counter is +/// incremented if one was supplied). +/// - [`Self::try_build`]: returns a raw `Result` for +/// callers that want to surface errors themselves. +/// +#[derive(Default)] +pub struct PruningPredicateBuilder<'a> { + file_schema: Option, + error_counter: Option<&'a Count>, + max_in_list_size: usize, +} + +impl<'a> PruningPredicateBuilder<'a> { + /// Create a new builder with the default pruning predicate configuration. + pub fn new() -> Self { + Self { + file_schema: None, + error_counter: None, + max_in_list_size: MAX_IN_LIST_SIZE, + } + } + + /// Set the schema of the container that will be pruned (typically the + /// parquet file schema). + pub fn with_file_schema(mut self, file_schema: SchemaRef) -> Self { + self.file_schema = Some(file_schema); + self + } + + /// Metric counter incremented once per predicate that fails to build. + /// Only consulted by [`Self::build`]; [`Self::try_build`] surfaces the + /// error directly. + pub fn with_error_counter(mut self, error_counter: &'a Count) -> Self { + self.error_counter = Some(error_counter); + self + } + + /// Cap on the size of `IN (...)` lists that will be rewritten into per- + /// value min/max statistics checks. Lists longer than this fall back to + /// the unhandled-predicate hook (typically "keep the container"). + /// + /// Query engines typically pass + /// `datafusion.execution.parquet.max_in_list_size` here. + pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self { + self.max_in_list_size = max_in_list_size; + self + } + + /// Build a [`PruningPredicate`] wrapped in `Some(Arc<..>)` when it can + /// prune, `None` when it is trivially true or when construction fails. + /// If [`Self::with_error_counter`] was set, construction failures are + /// recorded there. + pub fn build( + self, + predicate: Arc, + ) -> Option> { + let error_counter = self.error_counter; + match self.try_build(predicate) { + Ok(pruning_predicate) => { + if !pruning_predicate.always_true() { + return Some(Arc::new(pruning_predicate)); + } + } + Err(e) => { + debug!("Could not create pruning predicate for: {e}"); + if let Some(counter) = error_counter { + counter.add(1); + } } } - Err(e) => { - debug!("Could not create pruning predicate for: {e}"); - predicate_creation_errors.add(1); + None + } + + /// Build a [`PruningPredicate`], returning the construction error + /// directly. Callers that want the always-true predicate elided or + /// errors folded into a counter should use [`Self::build`] instead. + pub fn try_build( + self, + mut predicate: Arc, + ) -> Result { + let file_schema = self.file_schema.ok_or_else(|| { + _internal_datafusion_err!( + "PruningPredicateBuilder requires a file schema (call `with_file_schema`)" + ) + })?; + + // Get a (simpler) snapshot of the physical expr here to use with `PruningPredicate`. + // In particular this unravels any `DynamicFilterPhysicalExpr`s by snapshotting them + // so that PruningPredicate can work with a static expression. + let tf = snapshot_physical_expr_opt(predicate)?; + if tf.transformed { + // If we had an expression such as Dynamic(part_col < 5 and col < 10) + // (this could come from something like `select * from t order by part_col, col, limit 10`) + // after snapshotting and because `DynamicFilterPhysicalExpr` applies child replacements to its + // children after snapshotting and previously `replace_columns_with_literals` may have been called with partition values + // the expression we have now is `8 < 5 and col < 10`. + // Thus we need as simplifier pass to get `false and col < 10` => `false` here. + let simplifier = PhysicalExprSimplifier::new(&file_schema); + predicate = simplifier.simplify(tf.data)?; + } else { + predicate = tf.data; } + let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _; + + // build predicate expression once + let mut required_columns = RequiredColumns::new(); + let predicate_expr = build_predicate_expression( + &predicate, + &file_schema, + &mut required_columns, + &unhandled_hook, + self.max_in_list_size, + ); + let predicate_schema = required_columns.schema(); + // Simplify the newly created predicate to get rid of redundant casts, comparisons, etc. + let predicate_expr = + PhysicalExprSimplifier::new(&predicate_schema).simplify(predicate_expr)?; + let literal_guarantees = LiteralGuarantee::analyze(&predicate); + + Ok(PruningPredicate { + schema: file_schema, + predicate_expr, + required_columns, + orig_expr: predicate, + literal_guarantees, + }) } - None } /// Rewrites predicates that [`PredicateRewriter`] can not handle, e.g. certain @@ -460,46 +592,13 @@ impl PruningPredicate { /// returns a new expression. /// It is recommended that you pass the expressions through [`PhysicalExprSimplifier`] /// before calling this method to make sure the expressions can be used for pruning. - pub fn try_new(mut expr: Arc, schema: SchemaRef) -> Result { - // Get a (simpler) snapshot of the physical expr here to use with `PruningPredicate`. - // In particular this unravels any `DynamicFilterPhysicalExpr`s by snapshotting them - // so that PruningPredicate can work with a static expression. - let tf = snapshot_physical_expr_opt(expr)?; - if tf.transformed { - // If we had an expression such as Dynamic(part_col < 5 and col < 10) - // (this could come from something like `select * from t order by part_col, col, limit 10`) - // after snapshotting and because `DynamicFilterPhysicalExpr` applies child replacements to its - // children after snapshotting and previously `replace_columns_with_literals` may have been called with partition values - // the expression we have now is `8 < 5 and col < 10`. - // Thus we need as simplifier pass to get `false and col < 10` => `false` here. - let simplifier = PhysicalExprSimplifier::new(&schema); - expr = simplifier.simplify(tf.data)?; - } else { - expr = tf.data; - } - let unhandled_hook = Arc::new(ConstantUnhandledPredicateHook::default()) as _; - - // build predicate expression once - let mut required_columns = RequiredColumns::new(); - let predicate_expr = build_predicate_expression( - &expr, - &schema, - &mut required_columns, - &unhandled_hook, - ); - let predicate_schema = required_columns.schema(); - // Simplify the newly created predicate to get rid of redundant casts, comparisons, etc. - let predicate_expr = - PhysicalExprSimplifier::new(&predicate_schema).simplify(predicate_expr)?; - let literal_guarantees = LiteralGuarantee::analyze(&expr); - - Ok(Self { - schema, - predicate_expr, - required_columns, - orig_expr: expr, - literal_guarantees, - }) + /// + /// Use [`PruningPredicateBuilder`] to construct new pruning predicates. + #[deprecated(since = "55.0.0", note = "Use PruningPredicateBuilder instead")] + pub fn try_new(expr: Arc, schema: SchemaRef) -> Result { + PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr) } /// For each set of statistics, evaluates the pruning predicate @@ -1359,20 +1458,26 @@ fn build_is_null_column_expr( } } -/// The maximum number of entries in an `InList` that might be rewritten into -/// an OR chain -const MAX_LIST_VALUE_SIZE_REWRITE: usize = 20; +/// Default maximum number of entries in an `IN (...)` list that will be +/// rewritten into a chain of per-value min/max checks by +/// `build_predicate_expression`. Callers threading a [`PredicateRewriter`] +/// can override this via [`PredicateRewriter::with_max_in_list_size`], and +/// query engines can wire it from the +/// `datafusion.execution.parquet.max_in_list_size` config option. +pub const MAX_IN_LIST_SIZE: usize = 20; /// Rewrite a predicate expression in terms of statistics (min/max/null_counts) /// for use as a [`PruningPredicate`]. pub struct PredicateRewriter { unhandled_hook: Arc, + max_in_list_size: usize, } impl Default for PredicateRewriter { fn default() -> Self { Self { unhandled_hook: Arc::new(ConstantUnhandledPredicateHook::default()), + max_in_list_size: MAX_IN_LIST_SIZE, } } } @@ -1385,10 +1490,24 @@ impl PredicateRewriter { /// Set the unhandled hook to be used when a predicate can not be rewritten pub fn with_unhandled_hook( - self, + mut self, unhandled_hook: Arc, ) -> Self { - Self { unhandled_hook } + self.unhandled_hook = unhandled_hook; + self + } + + /// Set the maximum size of an `IN (...)` list that will be rewritten into a + /// chain of per-value statistics checks. Lists longer than this fall back + /// to the unhandled-predicate hook (typically "keep the container"), + /// effectively skipping container-level pruning for large IN lists. + /// + /// The default (see [`MAX_IN_LIST_SIZE`]) preserves the + /// historical behaviour. Callers wiring config through can override via + /// `datafusion.execution.max_in_list_size`. + pub fn with_max_in_list_size(mut self, max_in_list_size: usize) -> Self { + self.max_in_list_size = max_in_list_size; + self } /// Translate logical filter expression into pruning predicate @@ -1399,7 +1518,8 @@ impl PredicateRewriter { /// /// Returns the pruning predicate as an [`PhysicalExpr`] /// - /// Notice: Does not handle [`phys_expr::InListExpr`] greater than 20, which will fall back to calling `unhandled_hook` + /// Notice: `IN (...)` lists longer than `max_in_list_size` (default + /// [`MAX_IN_LIST_SIZE`]) fall back to calling `unhandled_hook`. pub fn rewrite_predicate_to_statistics_predicate( &self, expr: &Arc, @@ -1411,6 +1531,7 @@ impl PredicateRewriter { &Arc::new(schema.clone()), &mut required_columns, &self.unhandled_hook, + self.max_in_list_size, ) } } @@ -1423,12 +1544,15 @@ impl PredicateRewriter { /// /// Returns the pruning predicate as an [`PhysicalExpr`] /// -/// Notice: Does not handle [`phys_expr::InListExpr`] greater than 20, which will fall back to calling `unhandled_hook` +/// `max_in_list_size` is the largest `IN (...)` list that will be rewritten +/// into a chain of per-value statistics checks; longer lists fall back to +/// `unhandled_hook`. fn build_predicate_expression( expr: &Arc, schema: &SchemaRef, required_columns: &mut RequiredColumns, unhandled_hook: &Arc, + max_in_list_size: usize, ) -> Arc { if is_always_false(expr) { // Shouldn't return `unhandled_hook.handle(expr)` @@ -1463,9 +1587,7 @@ fn build_predicate_expression( } } if let Some(in_list) = expr.downcast_ref::() { - if !in_list.list().is_empty() - && in_list.list().len() <= MAX_LIST_VALUE_SIZE_REWRITE - { + if !in_list.list().is_empty() && in_list.list().len() <= max_in_list_size { let eq_op = if in_list.negated() { Operator::NotEq } else { @@ -1493,6 +1615,7 @@ fn build_predicate_expression( schema, required_columns, unhandled_hook, + max_in_list_size, ); } else { return unhandled_hook.handle(expr); @@ -1527,10 +1650,20 @@ fn build_predicate_expression( }; if op == Operator::And || op == Operator::Or { - let left_expr = - build_predicate_expression(&left, schema, required_columns, unhandled_hook); - let right_expr = - build_predicate_expression(&right, schema, required_columns, unhandled_hook); + let left_expr = build_predicate_expression( + &left, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + ); + let right_expr = build_predicate_expression( + &right, + schema, + required_columns, + unhandled_hook, + max_in_list_size, + ); // simplify boolean expression if applicable let expr = match (&left_expr, op, &right_expr) { (left, Operator::And, right) @@ -1816,6 +1949,13 @@ fn extract_string_literal(expr: &Arc) -> Option<&str> { None } +/// Wrap a string in a `Literal` whose `ScalarValue` matches `target_type` +fn string_literal_as(value: String, target_type: &DataType) -> Arc { + let utf8 = ScalarValue::Utf8(Some(value)); + let scalar = try_cast_literal_to_type(&utf8, target_type).unwrap_or(utf8); + Arc::new(phys_expr::Literal::new(scalar)) +} + /// Convert `column LIKE literal` where P is a constant prefix of the literal /// to a range check on the column: `P <= column && column < P'`, where P' is the /// lowest string after all P* strings. @@ -1835,6 +1975,8 @@ fn build_like_match( let min_column_expr = expr_builder.min_column_expr().ok()?; let max_column_expr = expr_builder.max_column_expr().ok()?; let scalar_expr = expr_builder.scalar_expr(); + // Synthesized bounds must match the column type (e.g. `Utf8View`). + let target_type = expr_builder.field.data_type(); // check that the scalar is a string literal let s = extract_string_literal(scalar_expr)?; // ANSI SQL specifies two wildcards: % and _. % matches zero or more characters, _ matches exactly one character. @@ -1846,18 +1988,12 @@ fn build_like_match( } let (lower_bound, upper_bound) = if has_wildcard { let incremented_prefix = increment_utf8(&decoded_prefix)?; - let lower_bound_lit = Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some( - decoded_prefix, - )))); - let upper_bound_lit = Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some( - incremented_prefix, - )))); + let lower_bound_lit = string_literal_as(decoded_prefix, target_type); + let upper_bound_lit = string_literal_as(incremented_prefix, target_type); (lower_bound_lit, upper_bound_lit) } else { // the like expression is a literal and can be converted into a comparison - let bound = Arc::new(phys_expr::Literal::new(ScalarValue::Utf8(Some( - decoded_prefix, - )))); + let bound = string_literal_as(decoded_prefix, target_type); (Arc::clone(&bound), bound) }; let lower_bound_expr = Arc::new(phys_expr::BinaryExpr::new( @@ -2452,7 +2588,10 @@ mod tests { ])); let expr = col("c1").eq(lit(100)).and(col("c2").eq(lit(200))); let expr = logical2physical(&expr, &schema); - let p = PruningPredicate::try_new(expr, Arc::clone(&schema)).unwrap(); + let p = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(expr) + .unwrap(); // note pruning expression refers to row_count twice assert_eq!( "c1_null_count@2 != row_count@3 AND c1_min@0 <= 100 AND 100 <= c1_max@1 AND c2_null_count@6 != row_count@3 AND c2_min@4 <= 200 AND 200 <= c2_max@5", @@ -3092,8 +3231,10 @@ mod tests { dynamic_phys_expr.with_new_children(remapped_expr).unwrap(); // After substitution the expression is c1 > 5 AND part = "B" which should prune the file since the partition value is "A" let expected = &[false]; - let p = - PruningPredicate::try_new(dynamic_filter_expr, Arc::clone(&schema)).unwrap(); + let p = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(dynamic_filter_expr) + .unwrap(); let result = p.prune(&statistics).unwrap(); assert_eq!(result, expected); } @@ -3310,7 +3451,7 @@ mod tests { fn row_group_predicate_in_list_to_many_values() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); // test c1 in(1..21) - // in pruning.rs has MAX_LIST_VALUE_SIZE_REWRITE = 20, more than this value will be rewrite + // in pruning.rs has MAX_IN_LIST_SIZE = 20, more than this value will be rewrite // always true let expr = col("c1").in_list((1..=21).map(lit).collect(), false); @@ -3322,6 +3463,123 @@ mod tests { Ok(()) } + // With the configurable cap, a caller that raises + // `max_in_list_size` above the default gets the IN list rewritten + // into a per-value min/max chain instead of falling through to `true`. + // This verifies both `PredicateRewriter::with_max_in_list_size` and the + // recursive OR path inside `build_predicate_expression`. + #[test] + fn row_group_predicate_in_list_rewritten_at_raised_cap() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + // 25 items — above the default 20, below a raised cap of 32. + let expr = col("c1").in_list((1..=25).map(lit).collect(), false); + let physical = logical2physical(&expr, &schema); + let rewriter = PredicateRewriter::new().with_max_in_list_size(32); + let predicate_expr = + rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema); + // At the raised cap, IN is rewritten into per-value min/max checks + // OR'd together; the resulting predicate must not collapse to + // `true` (which is what the default cap produces). + assert_ne!( + predicate_expr.to_string(), + "true", + "IN(25) with raised cap must rewrite into a statistics-based predicate, not fall through to `true`" + ); + // Sanity: the rewritten predicate references per-value literals. + assert!( + predicate_expr.to_string().contains(" <= 1 ") + && predicate_expr.to_string().contains(" <= 25 "), + "rewritten predicate should include per-value bounds for each IN entry, got: {predicate_expr}" + ); + Ok(()) + } + + // Guard: when the cap is 0 (opt-out) the IN branch is skipped entirely + // regardless of list length, so even a small IN falls through to the + // unhandled hook. + #[test] + fn row_group_predicate_in_list_disabled_at_zero_cap() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + let expr = col("c1").in_list(vec![lit(1), lit(2), lit(3)], false); + let physical = logical2physical(&expr, &schema); + let rewriter = PredicateRewriter::new().with_max_in_list_size(0); + let predicate_expr = + rewriter.rewrite_predicate_to_statistics_predicate(&physical, &schema); + assert_eq!( + predicate_expr.to_string(), + "true", + "cap=0 must skip IN rewrite even for small lists" + ); + Ok(()) + } + + // The high-level [`PruningPredicateBuilder`] should thread + // `max_in_list_size` all the way through: a 25-item IN with the default + // cap must fall through to the unhandled hook (`predicate_expr = true`), + // while a raised cap produces a real per-value statistics predicate. + #[test] + fn pruning_predicate_builder_threads_max_in_list_size() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + let expr = col("c1").in_list((1..=25).map(lit).collect(), false); + let physical = logical2physical(&expr, &schema); + + // With the default cap the IN branch bails out and the pruning + // predicate expression collapses to `true` (i.e., no container + // pruning based on stats). + let default_pp = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .try_build(Arc::clone(&physical))?; + assert_eq!( + default_pp.predicate_expr().to_string(), + "true", + "default cap must fall through to `true` for 25-item IN" + ); + + // Raising the cap produces a real statistics predicate with per- + // value bounds. + let raised_pp = PruningPredicateBuilder::new() + .with_file_schema(Arc::clone(&schema)) + .with_max_in_list_size(32) + .try_build(physical)?; + let raised_expr = raised_pp.predicate_expr().to_string(); + assert_ne!( + raised_expr, "true", + "raised cap must produce a real statistics predicate for 25-item IN" + ); + assert!( + raised_expr.contains(" <= 1 ") && raised_expr.contains(" <= 25 "), + "raised-cap predicate should include per-value bounds, got: {raised_expr}" + ); + Ok(()) + } + + #[test] + #[expect(deprecated)] + fn deprecated_try_new_delegates_to_builder() -> Result<()> { + let schema = + Arc::new(Schema::new(vec![Field::new("c1", DataType::Int32, false)])); + let expr = logical2physical(&col("c1").eq(lit(1)), &schema); + + let deprecated = + PruningPredicate::try_new(Arc::clone(&expr), Arc::clone(&schema))?; + let builder = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr)?; + + assert_eq!( + deprecated.predicate_expr().to_string(), + builder.predicate_expr().to_string() + ); + assert_eq!( + deprecated.required_columns().schema(), + builder.required_columns().schema() + ); + Ok(()) + } + #[test] fn row_group_predicate_cast_int_int() -> Result<()> { let schema = Schema::new(vec![Field::new("c1", DataType::Int32, false)]); @@ -5708,7 +5966,10 @@ mod tests { ) { println!("Pruning with expr: {expr}"); let expr = logical2physical(&expr, schema); - let p = PruningPredicate::try_new(expr, Arc::::clone(schema)).unwrap(); + let p = PruningPredicateBuilder::new() + .with_file_schema(Arc::::clone(schema)) + .try_build(expr) + .unwrap(); let result = p.prune(statistics).unwrap(); assert_eq!(result, expected); } @@ -5723,7 +5984,10 @@ mod tests { let expr = logical2physical(&expr, schema); let simplifier = PhysicalExprSimplifier::new(schema); let expr = simplifier.simplify(expr).unwrap(); - let p = PruningPredicate::try_new(expr, Arc::::clone(schema)).unwrap(); + let p = PruningPredicateBuilder::new() + .with_file_schema(Arc::::clone(schema)) + .try_build(expr) + .unwrap(); let result = p.prune(statistics).unwrap(); assert_eq!(result, expected); } @@ -5756,6 +6020,7 @@ mod tests { &Arc::new(schema.clone()), required_columns, &unhandled_hook, + MAX_IN_LIST_SIZE, ) } diff --git a/datafusion/session/Cargo.toml b/datafusion/session/Cargo.toml index 230e26d1fc9fc..2bbbdd20df1b8 100644 --- a/datafusion/session/Cargo.toml +++ b/datafusion/session/Cargo.toml @@ -31,6 +31,7 @@ version.workspace = true all-features = true [dependencies] +arrow-schema = { workspace = true } async-trait = { workspace = true } datafusion-common = { workspace = true } datafusion-execution = { workspace = true } diff --git a/datafusion/session/README.md b/datafusion/session/README.md index 4bb605b1e199c..72a693e81deb4 100644 --- a/datafusion/session/README.md +++ b/datafusion/session/README.md @@ -21,7 +21,7 @@ [Apache DataFusion] is an extensible query execution framework, written in Rust, that uses [Apache Arrow] as its in-memory format. -This crate provides **session-related abstractions** used in the DataFusion query engine. A _session_ represents the runtime context for query execution, including configuration, runtime environment, function registry, and planning. +This crate defines the **session-related APIs and extension points** used in the DataFusion query engine. A _session_ represents the runtime context for query execution, including configuration, runtime environment, function registry, and planning. This crate focuses on shared interfaces; concrete query-engine implementations live in higher-level DataFusion crates. Most projects should use the [`datafusion`] crate directly, which re-exports this module. If you are already using the [`datafusion`] crate, there is no diff --git a/datafusion/session/src/catalog.rs b/datafusion/session/src/catalog.rs new file mode 100644 index 0000000000000..bd9eb781abe77 --- /dev/null +++ b/datafusion/session/src/catalog.rs @@ -0,0 +1,246 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::any::Any; +use std::fmt::Debug; +use std::sync::Arc; + +pub use crate::schema::SchemaProvider; +use datafusion_common::Result; +use datafusion_common::not_impl_err; + +/// A catalog list that contains no catalogs. +/// +/// [`Session`](crate::Session) implementations that do not provide catalog +/// access can return this list explicitly. +#[derive(Debug, Default)] +pub struct EmptyCatalogProviderList; + +impl CatalogProviderList for EmptyCatalogProviderList { + fn register_catalog( + &self, + _name: String, + _catalog: Arc, + ) -> Option> { + None + } + + fn catalog_names(&self) -> Vec { + vec![] + } + + fn catalog(&self, _name: &str) -> Option> { + None + } +} + +/// Represents a catalog, comprising a number of named schemas. +/// +/// # Catalog Overview +/// +/// To plan and execute queries, DataFusion needs a "Catalog" that provides +/// metadata such as which schemas and tables exist, their columns and data +/// types, and how to access the data. +/// +/// The Catalog API consists: +/// * [`CatalogProviderList`]: a collection of `CatalogProvider`s +/// * [`CatalogProvider`]: a collection of `SchemaProvider`s (sometimes called a "database" in other systems) +/// * [`SchemaProvider`]: a collection of `TableProvider`s (often called a "schema" in other systems) +/// * [`TableProvider`]: individual tables +/// +/// # Implementing Catalogs +/// +/// To implement a catalog, you implement at least one of the [`CatalogProviderList`], +/// [`CatalogProvider`] and [`SchemaProvider`] traits and register them +/// appropriately in the `SessionContext`. +/// +/// DataFusion comes with a simple in-memory catalog implementation, +/// `MemoryCatalogProvider`, that is used by default and has no persistence. +/// DataFusion does not include more complex Catalog implementations because +/// catalog management is a key design choice for most data systems, and thus +/// it is unlikely that any general-purpose catalog implementation will work +/// well across many use cases. +/// +/// # Implementing "Remote" catalogs +/// +/// See [`remote_catalog`] for an end to end example of how to implement a +/// remote catalog. +/// +/// Sometimes catalog information is stored remotely and requires a network call +/// to retrieve. For example, the [Delta Lake] table format stores table +/// metadata in files on S3 that must be first downloaded to discover what +/// schemas and tables exist. +/// +/// [Delta Lake]: https://delta.io/ +/// [`remote_catalog`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/data_io/remote_catalog.rs +/// +/// The [`CatalogProvider`] can support this use case, but it takes some care. +/// The planning APIs in DataFusion are not `async` and thus network IO can not +/// be performed "lazily" / "on demand" during query planning. The rationale for +/// this design is that using remote procedure calls for all catalog accesses +/// required for query planning would likely result in multiple network calls +/// per plan, resulting in very poor planning performance. +/// +/// To implement [`CatalogProvider`] and [`SchemaProvider`] for remote catalogs, +/// you need to provide an in memory snapshot of the required metadata. Most +/// systems typically either already have this information cached locally or can +/// batch access to the remote catalog to retrieve multiple schemas and tables +/// in a single network call. +/// +/// Note that [`SchemaProvider::table`] **is** an `async` function in order to +/// simplify implementing simple [`SchemaProvider`]s. For many table formats it +/// is easy to list all available tables but there is additional non trivial +/// access required to read table details (e.g. statistics). +/// +/// The pattern that DataFusion itself uses to plan SQL queries is to walk over +/// the query to find all table references, performing required remote catalog +/// lookups in parallel, storing the results in a cached snapshot, and then plans +/// the query using that snapshot. +/// +/// # Example Catalog Implementations +/// +/// Here are some examples of how to implement custom catalogs: +/// +/// * [`datafusion-cli`]: [`DynamicFileCatalogProvider`] catalog provider +/// that treats files and directories on a filesystem as tables. +/// +/// * The [`catalog.rs`]: a simple directory based catalog. +/// +/// * [delta-rs]: [`UnityCatalogProvider`] implementation that can +/// read from Delta Lake tables +/// +/// [`datafusion-cli`]: https://datafusion.apache.org/user-guide/cli/index.html +/// [`DynamicFileCatalogProvider`]: https://github.com/apache/datafusion/blob/31b9b48b08592b7d293f46e75707aad7dadd7cbc/datafusion-cli/src/catalog.rs#L75 +/// [`catalog.rs`]: https://github.com/apache/datafusion/blob/main/datafusion-examples/examples/data_io/catalog.rs +/// [delta-rs]: https://github.com/delta-io/delta-rs +/// [`UnityCatalogProvider`]: https://github.com/delta-io/delta-rs/blob/951436ecec476ce65b5ed3b58b50fb0846ca7b91/crates/deltalake-core/src/data_catalog/unity/datafusion.rs#L111-L123 +/// +/// [`TableProvider`]: crate::TableProvider +pub trait CatalogProvider: Any + Debug + Sync + Send { + /// Retrieves the list of available schema names in this catalog. + fn schema_names(&self) -> Vec; + + /// Retrieves a specific schema from the catalog by name, provided it exists. + fn schema(&self, name: &str) -> Option>; + + /// Adds a new schema to this catalog. + /// + /// If a schema of the same name existed before, it is replaced in + /// the catalog and returned. + /// + /// By default returns a "Not Implemented" error + fn register_schema( + &self, + name: &str, + schema: Arc, + ) -> Result>> { + // use variables to avoid unused variable warnings + let _ = name; + let _ = schema; + not_impl_err!("Registering new schemas is not supported") + } + + /// Removes a schema from this catalog. Implementations of this method should return + /// errors if the schema exists but cannot be dropped. For example, in DataFusion's + /// default in-memory catalog, `MemoryCatalogProvider`, a non-empty schema + /// will only be successfully dropped when `cascade` is true. + /// This is equivalent to how DROP SCHEMA works in PostgreSQL. + /// + /// Implementations of this method should return None if schema with `name` + /// does not exist. + /// + /// By default returns a "Not Implemented" error + fn deregister_schema( + &self, + _name: &str, + _cascade: bool, + ) -> Result>> { + not_impl_err!("Deregistering new schemas is not supported") + } +} + +impl dyn CatalogProvider { + /// Returns `true` if the catalog provider is of type `T`. + /// + /// Prefer this over `downcast_ref::().is_some()`. Works correctly when + /// called on `Arc` via auto-deref. + pub fn is(&self) -> bool { + (self as &dyn Any).is::() + } + + /// Attempts to downcast this catalog provider to a concrete type `T`, + /// returning `None` if the provider is not of that type. + /// + /// Works correctly when called on `Arc` via auto-deref, + /// unlike `(&arc as &dyn Any).downcast_ref::()` which would attempt to + /// downcast the `Arc` itself. + pub fn downcast_ref(&self) -> Option<&T> { + (self as &dyn Any).downcast_ref() + } +} + +/// Represent a list of named [`CatalogProvider`]s. +/// +/// Please see the documentation on [`CatalogProvider`] for details of +/// implementing a custom catalog. +pub trait CatalogProviderList: Any + Debug + Sync + Send { + /// Adds a new catalog to this catalog list + /// If a catalog of the same name existed before, it is replaced in the list and returned. + fn register_catalog( + &self, + name: String, + catalog: Arc, + ) -> Option>; + + /// Retrieves the list of available catalog names + fn catalog_names(&self) -> Vec; + + /// Retrieves a specific catalog by name, provided it exists. + fn catalog(&self, name: &str) -> Option>; +} + +impl dyn CatalogProviderList { + /// Returns `true` if the catalog provider list is of type `T`. + /// + /// Prefer this over `downcast_ref::().is_some()`. Works correctly when + /// called on `Arc` via auto-deref. + pub fn is(&self) -> bool { + (self as &dyn Any).is::() + } + + /// Attempts to downcast this catalog provider list to a concrete type `T`, + /// returning `None` if the provider list is not of that type. + /// + /// Works correctly when called on `Arc` via + /// auto-deref, unlike `(&arc as &dyn Any).downcast_ref::()` which would + /// attempt to downcast the `Arc` itself. + pub fn downcast_ref(&self) -> Option<&T> { + (self as &dyn Any).downcast_ref() + } +} + +#[cfg(test)] +mod tests { + use super::{CatalogProviderList, EmptyCatalogProviderList}; + + #[test] + fn empty_catalog_provider_list_has_no_catalogs() { + let catalogs = EmptyCatalogProviderList; + assert!(catalogs.catalog_names().is_empty()); + assert!(catalogs.catalog("missing").is_none()); + } +} diff --git a/datafusion/session/src/lib.rs b/datafusion/session/src/lib.rs index 11f734e757452..6f7cfb7792c73 100644 --- a/datafusion/session/src/lib.rs +++ b/datafusion/session/src/lib.rs @@ -15,18 +15,27 @@ // specific language governing permissions and limitations // under the License. +// Make sure fast / cheap clones on Arc are explicit: +// https://github.com/apache/datafusion/issues/11143 +#![cfg_attr(not(test), deny(clippy::clone_on_ref_ptr))] #![cfg_attr(test, allow(clippy::needless_pass_by_value))] -//! Session management for DataFusion query execution environment +//! Session APIs for the DataFusion query execution environment //! -//! This module provides the core session management functionality for DataFusion, -//! handling both Catalog (Table) and Datasource (File) configurations. It defines -//! the fundamental interfaces and implementations for maintaining query execution -//! state and configurations. +//! This crate defines shared interfaces for session-related APIs and extension +//! points. Concrete query-engine implementations are provided by higher-level +//! DataFusion crates. //! //! Key components: -//! * [`Session`] - Manages query execution context, including configurations, +//! * [`Session`] - Describes a query execution context, including configurations, //! catalogs, and runtime state +//! * [`CatalogProviderList`], [`CatalogProvider`], and [`SchemaProvider`] - +//! Describe catalog hierarchies +//! * [`TableProvider`] - Provides data for query planning and execution +//! * [`QueryPlanner`], [`PhysicalPlanner`], and [`ExtensionPlanner`] - Query and +//! physical planning contracts +//! * [`PhysicalOptimizerRule`] and [`PhysicalOptimizerContext`] - Physical +//! optimization contracts //! * [`SessionStore`] - Handles session persistence and retrieval //! //! The session system enables: @@ -36,6 +45,23 @@ //! * Runtime environment configuration //! * Query state persistence +pub mod catalog; +pub mod physical_optimizer; +pub mod planner; +pub mod schema; pub mod session; +pub mod table; +pub use crate::catalog::{ + CatalogProvider, CatalogProviderList, EmptyCatalogProviderList, +}; +pub use crate::physical_optimizer::{PhysicalOptimizerContext, PhysicalOptimizerRule}; +pub use crate::planner::{ + ExtensionPlanner, PhysicalPlanner, QueryPlanner, UnsupportedQueryPlanner, +}; +pub use crate::schema::SchemaProvider; pub use crate::session::{Session, SessionStore}; +pub use crate::table::{ + ScanArgs, ScanResult, TableFunction, TableFunctionArgs, TableFunctionImpl, + TableProvider, TableProviderFactory, +}; diff --git a/datafusion/session/src/physical_optimizer.rs b/datafusion/session/src/physical_optimizer.rs new file mode 100644 index 0000000000000..751a8e12d93ed --- /dev/null +++ b/datafusion/session/src/physical_optimizer.rs @@ -0,0 +1,84 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Physical optimizer interfaces. + +use std::fmt::Debug; +use std::sync::Arc; + +use datafusion_common::Result; +use datafusion_common::config::ConfigOptions; +use datafusion_physical_plan::ExecutionPlan; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; + +/// Context available to physical optimizer rules. +/// +/// This trait provides access to configuration options and an optional statistics +/// registry for enhanced statistics lookup. +pub trait PhysicalOptimizerContext: Send + Sync { + /// Returns the configuration options. + fn config_options(&self) -> &ConfigOptions; + + /// Returns the statistics registry for enhanced statistics lookup. + /// + /// Returns `None` if no registry is configured, in which case rules + /// should fall back to using [`ExecutionPlan::partition_statistics`]. + fn statistics_registry(&self) -> Option<&StatisticsRegistry> { + None + } +} + +/// `PhysicalOptimizerRule` transforms one [`ExecutionPlan`] into another which +/// computes the same results, but in a potentially more efficient way. +/// +/// Use [`SessionState::add_physical_optimizer_rule`] to register additional +/// `PhysicalOptimizerRule`s. +/// +/// [`SessionState::add_physical_optimizer_rule`]: https://docs.rs/datafusion/latest/datafusion/execution/session_state/struct.SessionState.html#method.add_physical_optimizer_rule +pub trait PhysicalOptimizerRule: Debug + std::any::Any { + /// Rewrite `plan` to an optimized form. + /// + /// This is the primary optimization method. For rules that need access to + /// the statistics registry, override [`optimize_with_context`](Self::optimize_with_context) instead. + fn optimize( + &self, + plan: Arc, + config: &ConfigOptions, + ) -> Result>; + + /// Rewrite `plan` with access to extended context (statistics registry, etc.). + /// + /// Override this method if you need access to the statistics registry for + /// enhanced statistics lookup. The default implementation simply calls + /// [`optimize`](Self::optimize) with the config options from the context. + fn optimize_with_context( + &self, + plan: Arc, + context: &dyn PhysicalOptimizerContext, + ) -> Result> { + self.optimize(plan, context.config_options()) + } + + /// A human readable name for this optimizer rule + fn name(&self) -> &str; + + /// A flag to indicate whether the physical planner should validate that the rule will not + /// change the schema of the plan after the rewriting. + /// Some of the optimization rules might change the nullable properties of the schema + /// and should disable the schema check. + fn schema_check(&self) -> bool; +} diff --git a/datafusion/session/src/planner.rs b/datafusion/session/src/planner.rs new file mode 100644 index 0000000000000..37726009f0f4d --- /dev/null +++ b/datafusion/session/src/planner.rs @@ -0,0 +1,198 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Query planner interfaces. + +use std::any::Any; +use std::fmt::Debug; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion_common::{DFSchema, Result, not_impl_err}; +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion_expr::{Expr, LogicalPlan, TableScan, UserDefinedLogicalNode}; +use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; + +use crate::Session; + +/// A planner that creates a physical plan for a query. +#[async_trait] +pub trait QueryPlanner: Any + Debug { + /// Given a [`LogicalPlan`], create an [`ExecutionPlan`] suitable for execution + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result>; +} + +/// A query planner that reports that planning is not implemented. +/// +/// [`Session`] implementations that do not expose a query planner can return +/// this planner explicitly. +#[derive(Debug, Default)] +pub struct UnsupportedQueryPlanner; + +#[async_trait] +impl QueryPlanner for UnsupportedQueryPlanner { + async fn create_physical_plan( + &self, + _logical_plan: &LogicalPlan, + _session: &dyn Session, + ) -> Result> { + not_impl_err!("This session does not expose its query planner") + } +} + +/// Physical query planner that converts a [`LogicalPlan`] to an +/// [`ExecutionPlan`] suitable for execution. +#[async_trait] +pub trait PhysicalPlanner: Send + Sync { + /// Create a physical plan from a logical plan + async fn create_physical_plan( + &self, + logical_plan: &LogicalPlan, + session: &dyn Session, + ) -> Result>; + + /// Create a physical expression from a logical expression + /// suitable for evaluation + /// + /// `expr`: the expression to convert + /// + /// `input_dfschema`: the logical plan schema for evaluating `expr` + /// + /// `planning_ctx`: the [`PhysicalPlanningContext`] used to resolve + /// `Expr::ScalarSubquery` nodes. During physical planning the planner + /// threads the context of the plan currently being converted to a physical + /// plan (for example into [`ExtensionPlanner::plan_extension`], which + /// should forward it here). Callers creating physical expressions outside + /// of a plan should pass `&PhysicalPlanningContext::default()`. + fn create_physical_expr( + &self, + expr: &Expr, + input_dfschema: &DFSchema, + session: &dyn Session, + planning_ctx: &PhysicalPlanningContext, + ) -> Result>; +} + +/// This trait exposes the ability to plan an [`ExecutionPlan`] out of a [`LogicalPlan`]. +#[async_trait] +pub trait ExtensionPlanner { + /// Create a physical plan for a [`UserDefinedLogicalNode`]. + /// + /// `input_dfschema`: the logical plan schema for the inputs to this node + /// + /// Returns an error when the planner knows how to plan the concrete + /// implementation of `node` but errors while doing so. + /// + /// Returns `None` when the planner does not know how to plan the + /// `node` and wants to delegate the planning to another + /// [`ExtensionPlanner`]. + /// + /// `planning_ctx` is the [`PhysicalPlanningContext`] of the plan subtree + /// currently being converted to a physical plan. Forward it to + /// [`PhysicalPlanner::create_physical_expr`] when creating this node's + /// physical expressions so that scalar subqueries resolve against the same + /// subquery state as the rest of the plan. + async fn plan_extension( + &self, + planner: &dyn PhysicalPlanner, + node: &dyn UserDefinedLogicalNode, + logical_inputs: &[&LogicalPlan], + physical_inputs: &[Arc], + session: &dyn Session, + planning_ctx: &PhysicalPlanningContext, + ) -> Result>>; + + /// Create a physical plan for a [`LogicalPlan::TableScan`]. + /// + /// This is useful for planning valid [`TableSource`]s that are not `TableProvider`s. + /// + /// Returns: + /// * `Ok(Some(plan))` if the planner knows how to plan the `scan` + /// * `Ok(None)` if the planner does not know how to plan the `scan` and wants to delegate the planning to another [`ExtensionPlanner`] + /// * `Err` if the planner knows how to plan the `scan` but errors while doing so + /// + /// # Example + /// + /// ```rust,ignore + /// use std::sync::Arc; + /// use datafusion::physical_plan::ExecutionPlan; + /// use datafusion::logical_expr::TableScan; + /// use datafusion::catalog::Session; + /// use datafusion::error::Result; + /// use datafusion_session::{ExtensionPlanner, PhysicalPlanner}; + /// use async_trait::async_trait; + /// + /// // Your custom table source type + /// struct MyCustomTableSource { /* ... */ } + /// + /// // Your custom execution plan + /// struct MyCustomExec { /* ... */ } + /// + /// struct MyExtensionPlanner; + /// + /// #[async_trait] + /// impl ExtensionPlanner for MyExtensionPlanner { + /// async fn plan_extension( + /// &self, + /// _planner: &dyn PhysicalPlanner, + /// _node: &dyn UserDefinedLogicalNode, + /// _logical_inputs: &[&LogicalPlan], + /// _physical_inputs: &[Arc], + /// _session: &dyn Session, + /// _planning_ctx: &PhysicalPlanningContext, + /// ) -> Result>> { + /// Ok(None) + /// } + /// + /// async fn plan_table_scan( + /// &self, + /// _planner: &dyn PhysicalPlanner, + /// scan: &TableScan, + /// _session: &dyn Session, + /// _planning_ctx: &PhysicalPlanningContext, + /// ) -> Result>> { + /// // Check if this is your custom table source + /// if scan.source.is::() { + /// // Create a custom execution plan for your table source + /// let exec = MyCustomExec::new( + /// scan.table_name.clone(), + /// Arc::clone(scan.projected_schema.inner()), + /// ); + /// Ok(Some(Arc::new(exec))) + /// } else { + /// // Return None to let other extension planners handle it + /// Ok(None) + /// } + /// } + /// } + /// ``` + /// + /// [`TableSource`]: datafusion_expr::TableSource + async fn plan_table_scan( + &self, + _planner: &dyn PhysicalPlanner, + _scan: &TableScan, + _session: &dyn Session, + _planning_ctx: &PhysicalPlanningContext, + ) -> Result>> { + Ok(None) + } +} diff --git a/datafusion/session/src/schema.rs b/datafusion/session/src/schema.rs new file mode 100644 index 0000000000000..7a66072bb4d8a --- /dev/null +++ b/datafusion/session/src/schema.rs @@ -0,0 +1,107 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Describes the interface and built-in implementations of schemas, +//! representing collections of named tables. + +use async_trait::async_trait; +use datafusion_common::{DataFusionError, exec_err}; +use std::any::Any; +use std::fmt::Debug; +use std::sync::Arc; + +use crate::table::TableProvider; +use datafusion_common::Result; +use datafusion_expr::TableType; + +/// Represents a schema, comprising a number of named tables. +/// +/// Please see [`CatalogProvider`] for details of implementing a custom catalog. +/// +/// [`CatalogProvider`]: super::CatalogProvider +#[async_trait] +pub trait SchemaProvider: Any + Debug + Sync + Send { + /// Returns the owner of the Schema, default is None. This value is reported + /// as part of `information_schema.schemata`. + fn owner_name(&self) -> Option<&str> { + None + } + + /// Retrieves the list of available table names in this schema. + fn table_names(&self) -> Vec; + + /// Retrieves a specific table from the schema by name, if it exists, + /// otherwise returns `None`. + async fn table( + &self, + name: &str, + ) -> Result>, DataFusionError>; + + /// Retrieves the type of a specific table from the schema by name, if it exists, otherwise + /// returns `None`. Implementations for which this operation is cheap but [Self::table] is + /// expensive can override this to improve operations that only need the type, e.g. + /// `SELECT * FROM information_schema.tables`. + async fn table_type(&self, name: &str) -> Result> { + self.table(name).await.map(|o| o.map(|t| t.table_type())) + } + + /// If supported by the implementation, adds a new table named `name` to + /// this schema. + /// + /// If a table of the same name was already registered, returns "Table + /// already exists" error. + #[expect(unused_variables)] + fn register_table( + &self, + name: String, + table: Arc, + ) -> Result>> { + exec_err!("schema provider does not support registering tables") + } + + /// If supported by the implementation, removes the `name` table from this + /// schema and returns the previously registered [`TableProvider`], if any. + /// + /// If no `name` table exists, returns Ok(None). + #[expect(unused_variables)] + fn deregister_table(&self, name: &str) -> Result>> { + exec_err!("schema provider does not support deregistering tables") + } + + /// Returns true if table exist in the schema provider, false otherwise. + fn table_exist(&self, name: &str) -> bool; +} + +impl dyn SchemaProvider { + /// Returns `true` if the schema provider is of type `T`. + /// + /// Prefer this over `downcast_ref::().is_some()`. Works correctly when + /// called on `Arc` via auto-deref. + pub fn is(&self) -> bool { + (self as &dyn Any).is::() + } + + /// Attempts to downcast this schema provider to a concrete type `T`, + /// returning `None` if the provider is not of that type. + /// + /// Works correctly when called on `Arc` via auto-deref, + /// unlike `(&arc as &dyn Any).downcast_ref::()` which would attempt to + /// downcast the `Arc` itself. + pub fn downcast_ref(&self) -> Option<&T> { + (self as &dyn Any).downcast_ref() + } +} diff --git a/datafusion/session/src/session.rs b/datafusion/session/src/session.rs index 82dda6655f8e2..f6143cc4a4d1d 100644 --- a/datafusion/session/src/session.rs +++ b/datafusion/session/src/session.rs @@ -26,12 +26,17 @@ use datafusion_expr::registry::ExtensionTypeRegistryRef; use datafusion_expr::{ AggregateUDF, Expr, HigherOrderUDF, LogicalPlan, ScalarUDF, WindowUDF, }; +use datafusion_physical_plan::operator_statistics::StatisticsRegistry; use datafusion_physical_plan::{ExecutionPlan, PhysicalExpr}; + +use crate::CatalogProviderList; use parking_lot::{Mutex, RwLock}; use std::any::Any; use std::collections::HashMap; use std::sync::{Arc, Weak}; +use crate::{PhysicalOptimizerRule, QueryPlanner, UnsupportedQueryPlanner}; + /// Interface for accessing [`SessionState`] from the catalog and data source. /// /// This trait provides access to the information needed to plan and execute @@ -79,11 +84,62 @@ pub trait Session: Send + Sync { /// Return the [`SessionConfig`] fn config(&self) -> &SessionConfig; + /// Return the catalogs registered with this session. + fn catalog_list(&self) -> Arc; + /// return the [`ConfigOptions`] fn config_options(&self) -> &ConfigOptions { self.config().options() } + /// Return the query planner for this session. + /// + /// # Warning + /// + /// The default implementation returns an [`UnsupportedQueryPlanner`], so + /// [`Session::create_physical_plan`] will fail. Sessions that support + /// physical planning should override this method (for example by returning + /// `SessionState::query_planner`). + fn query_planner(&self) -> Arc { + Arc::new(UnsupportedQueryPlanner) + } + + /// Optimize a logical plan. + /// + /// # Warning + /// + /// The default implementation returns the plan **unchanged**, applying no + /// logical optimizations whatsoever. This is almost never what you want: + /// without optimization, queries execute in their naive, unoptimized form + /// and may be dramatically slower or fail to run at all. The default exists + /// only so this crate need not depend on the optimizer; any real session + /// should override this method (for example by delegating to + /// `SessionState::optimize`). + fn optimize(&self, plan: &LogicalPlan) -> Result { + Ok(plan.clone()) + } + + /// Return the physical optimizer rules for this session. + /// + /// # Warning + /// + /// The default implementation returns **no rules**. This is almost never + /// what you want: DataFusion relies on physical optimizer rules for + /// correctness-critical rewrites (such as inserting the repartitioning and + /// coalescing needed for parallel and multi-partition execution), so a + /// session with no rules will produce plans that are inefficient or that + /// fail to execute. The default exists only so this crate need not depend + /// on the optimizer; any real session should override this method (for + /// example by returning `SessionState::physical_optimizers`). + fn physical_optimizers(&self) -> &[Arc] { + &[] + } + + /// Return the optional statistics registry used during physical optimization. + fn statistics_registry(&self) -> Option<&StatisticsRegistry> { + None + } + /// Creates a physical [`ExecutionPlan`] plan from a [`LogicalPlan`]. /// /// Note: this will optimize the provided plan first. @@ -114,7 +170,7 @@ pub trait Session: Send + Sync { fn scalar_functions(&self) -> &HashMap>; /// Return reference to higher_order_functions - fn higher_order_functions(&self) -> &HashMap>; + fn higher_order_functions(&self) -> &HashMap>; /// Return reference to aggregate_functions fn aggregate_functions(&self) -> &HashMap>; diff --git a/datafusion/session/src/table.rs b/datafusion/session/src/table.rs new file mode 100644 index 0000000000000..69e7e731b32c7 --- /dev/null +++ b/datafusion/session/src/table.rs @@ -0,0 +1,662 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::any::Any; +use std::borrow::Cow; +use std::fmt::Debug; +use std::sync::Arc; + +use crate::session::Session; +use arrow_schema::SchemaRef; +use async_trait::async_trait; +use datafusion_common::{Constraints, Statistics, not_impl_err}; +use datafusion_common::{DFSchemaRef, Result, internal_err}; +use datafusion_expr::Expr; +use datafusion_expr::statistics::StatisticsRequest; + +use datafusion_expr::dml::{InsertOp, MergeIntoClause}; +use datafusion_expr::{ + CreateExternalTable, LogicalPlan, TableProviderFilterPushDown, TableType, +}; +use datafusion_physical_plan::ExecutionPlan; + +/// A table which can be queried and modified. +/// +/// Please see [`CatalogProvider`] for details of implementing a custom catalog. +/// +/// [`TableProvider`] represents a source of data which can provide data as +/// Apache Arrow [`RecordBatch`]es. Implementations of this trait provide +/// important information for planning such as: +/// +/// 1. [`Self::schema`]: The schema (columns and their types) of the table +/// 2. [`Self::supports_filters_pushdown`]: Should filters be pushed into this scan +/// 2. [`Self::scan`]: An [`ExecutionPlan`] that can read data +/// +/// [`RecordBatch`]: https://docs.rs/arrow/latest/arrow/record_batch/struct.RecordBatch.html +/// [`CatalogProvider`]: super::CatalogProvider +#[async_trait] +pub trait TableProvider: Any + Debug + Sync + Send { + /// Get a reference to the schema for this table + fn schema(&self) -> SchemaRef; + + /// Get a reference to the constraints of the table. + /// Returns: + /// - `None` for tables that do not support constraints. + /// - `Some(&Constraints)` for tables supporting constraints. + /// Therefore, a `Some(&Constraints::empty())` return value indicates that + /// this table supports constraints, but there are no constraints. + fn constraints(&self) -> Option<&Constraints> { + None + } + + /// Get the type of this table for metadata/catalog purposes. + fn table_type(&self) -> TableType; + + /// Get the create statement used to create this table, if available. + fn get_table_definition(&self) -> Option<&str> { + None + } + + /// Get the [`LogicalPlan`] of this table, if available. + fn get_logical_plan(&'_ self) -> Option> { + None + } + + /// Get the default value for a column, if available. + fn get_column_default(&self, _column: &str) -> Option<&Expr> { + None + } + + /// Create an [`ExecutionPlan`] for scanning the table with optional + /// `projection`, `filter`, and `limit`, described below. + /// + /// The returned `ExecutionPlan` is responsible for scanning the datasource's + /// partitions in a streaming, parallelized fashion. + /// + /// # Projection + /// + /// If specified, only a subset of columns should be returned, in the order + /// specified. The projection is a set of indexes of the fields in + /// [`Self::schema`]. + /// + /// DataFusion provides the projection so the scan reads only the columns + /// actually used in the query, an optimization called "Projection + /// Pushdown". Some datasources, such as Parquet, can use this information + /// to go significantly faster when only a subset of columns is required. + /// + /// # Filters + /// + /// A list of boolean filter [`Expr`]s to evaluate *during* the scan, in the + /// manner specified by [`Self::supports_filters_pushdown`]. Only rows for + /// which *all* of the `Expr`s evaluate to `true` must be returned (that is, + /// the expressions are `AND`ed together). + /// + /// To enable filter pushdown, override + /// [`Self::supports_filters_pushdown`]. The default implementation does not + /// push down filters, and `filters` will be empty. + /// + /// DataFusion pushes filters into scans whenever possible ("Filter + /// Pushdown"). Depending on the data format and implementation, evaluating + /// predicates during the scan can significantly improve performance. + /// + /// ## Note: Some columns may appear *only* in Filters + /// + /// In some cases, a query may use a column only in a filter and the + /// projection will not contain all columns referenced by the filter + /// expressions. + /// + /// For example, given the query `SELECT t.a FROM t WHERE t.b > 5`, + /// + /// ```text + /// ┌────────────────────┐ + /// │ Projection(t.a) │ + /// └────────────────────┘ + /// ▲ + /// │ + /// │ + /// ┌────────────────────┐ Filter ┌────────────────────┐ Projection ┌────────────────────┐ + /// │ Filter(t.b > 5) │────Pushdown──▶ │ Projection(t.a) │ ───Pushdown───▶ │ Projection(t.a) │ + /// └────────────────────┘ └────────────────────┘ └────────────────────┘ + /// ▲ ▲ ▲ + /// │ │ │ + /// │ │ ┌────────────────────┐ + /// ┌────────────────────┐ ┌────────────────────┐ │ Scan │ + /// │ Scan │ │ Scan │ │ filter=(t.b > 5) │ + /// └────────────────────┘ │ filter=(t.b > 5) │ │ projection=(t.a) │ + /// └────────────────────┘ └────────────────────┘ + /// + /// Initial Plan If `TableProviderFilterPushDown` Projection pushdown notes that + /// returns true, filter pushdown the scan only needs t.a + /// pushes the filter into the scan + /// BUT internally evaluating the + /// predicate still requires t.b + /// ``` + /// + /// # Limit + /// + /// If `limit` is specified, the scan must produce *at least* this many + /// rows, though it may return more. Like Projection Pushdown and Filter + /// Pushdown, DataFusion pushes `LIMIT`s as far down in the plan as + /// possible. This is called "Limit Pushdown", and some sources can use the + /// information to improve performance. + /// + /// Note: If any pushed-down filters are `Inexact`, the `LIMIT` cannot be + /// pushed down. Inexact filters do not guarantee that every filtered row is + /// removed, so applying the limit could leave too few rows to return in the + /// final result. + /// + /// # Evaluation Order + /// + /// The logical evaluation order is `filters`, then `limit`, then + /// `projection`. + /// + /// Note that `limit` applies to the filtered result, not to the unfiltered + /// input, and `projection` affects only which columns are returned, not + /// which rows qualify. + /// + /// For example, if a scan receives: + /// + /// - `projection = [a]` + /// - `filters = [b > 5]` + /// - `limit = Some(3)` + /// + /// It must logically produce results equivalent to: + /// + /// ```text + /// PROJECTION a (LIMIT 3 (SCAN WHERE b > 5)) + /// ``` + /// + /// As noted above, columns referenced only by pushed-down filters may be + /// absent from `projection`. + async fn scan( + &self, + state: &dyn Session, + projection: Option<&Vec>, + filters: &[Expr], + limit: Option, + ) -> Result>; + + /// Create an [`ExecutionPlan`] for scanning the table using structured arguments. + /// + /// This method uses [`ScanArgs`] to pass scan parameters in a structured way + /// and returns a [`ScanResult`] containing the execution plan. + /// + /// Table providers can override this method to take advantage of additional + /// parameters like the upcoming `preferred_ordering` that may not be available through + /// other scan methods. + /// + /// # Arguments + /// * `state` - The session state containing configuration and context + /// * `args` - Structured scan arguments including projection, filters, limit, and ordering preferences + /// + /// # Returns + /// A [`ScanResult`] containing the [`ExecutionPlan`] for scanning the table + /// + /// See [`Self::scan`] for detailed documentation about projection, filters, and limits. + async fn scan_with_args<'a>( + &self, + state: &dyn Session, + args: ScanArgs<'a>, + ) -> Result { + let filters = args.filters().unwrap_or(&[]); + let projection = args.projection().map(|p| p.to_vec()); + let limit = args.limit(); + let plan = self + .scan(state, projection.as_ref(), filters, limit) + .await?; + Ok(plan.into()) + } + + /// Specify if DataFusion should provide filter expressions to the + /// TableProvider to apply *during* the scan. + /// + /// Some TableProviders can evaluate filters more efficiently than the + /// `Filter` operator in DataFusion, for example by using an index. + /// + /// # Parameters and Return Value + /// + /// The return `Vec` must have one element for each element of the `filters` + /// argument. The value of each element indicates if the TableProvider can + /// apply the corresponding filter during the scan. The position in the return + /// value corresponds to the expression in the `filters` parameter. + /// + /// If the length of the resulting `Vec` does not match the `filters` input + /// an error will be thrown. + /// + /// Each element in the resulting `Vec` is one of the following: + /// * [`Exact`] or [`Inexact`]: The TableProvider can apply the filter + /// during scan + /// * [`Unsupported`]: The TableProvider cannot apply the filter during scan + /// + /// By default, this function returns [`Unsupported`] for all filters, + /// meaning no filters will be provided to [`Self::scan`]. + /// + /// [`Unsupported`]: TableProviderFilterPushDown::Unsupported + /// [`Exact`]: TableProviderFilterPushDown::Exact + /// [`Inexact`]: TableProviderFilterPushDown::Inexact + /// # Example + /// + /// ```rust + /// # use std::any::Any; + /// # use std::sync::Arc; + /// # use arrow_schema::SchemaRef; + /// # use async_trait::async_trait; + /// # use datafusion_session::{TableProvider, Session}; + /// # use datafusion_common::Result; + /// # use datafusion_expr::{Expr, TableProviderFilterPushDown, TableType}; + /// # use datafusion_physical_plan::ExecutionPlan; + /// // Define a struct that implements the TableProvider trait + /// #[derive(Debug)] + /// struct TestDataSource {} + /// + /// #[async_trait] + /// impl TableProvider for TestDataSource { + /// # fn schema(&self) -> SchemaRef { todo!() } + /// # fn table_type(&self) -> TableType { todo!() } + /// # async fn scan(&self, s: &dyn Session, p: Option<&Vec>, f: &[Expr], l: Option) -> Result> { + /// todo!() + /// # } + /// // Override the supports_filters_pushdown to evaluate which expressions + /// // to accept as pushdown predicates. + /// fn supports_filters_pushdown(&self, filters: &[&Expr]) -> Result> { + /// // Process each filter + /// let support: Vec<_> = filters.iter().map(|expr| { + /// match expr { + /// // This example only supports a between expr with a single column named "c1". + /// Expr::Between(between_expr) => { + /// between_expr.expr + /// .try_as_col() + /// .map(|column| { + /// if column.name == "c1" { + /// TableProviderFilterPushDown::Exact + /// } else { + /// TableProviderFilterPushDown::Unsupported + /// } + /// }) + /// // If there is no column in the expr set the filter to unsupported. + /// .unwrap_or(TableProviderFilterPushDown::Unsupported) + /// } + /// _ => { + /// // For all other cases return Unsupported. + /// TableProviderFilterPushDown::Unsupported + /// } + /// } + /// }).collect(); + /// Ok(support) + /// } + /// } + /// ``` + fn supports_filters_pushdown( + &self, + filters: &[&Expr], + ) -> Result> { + Ok(vec![ + TableProviderFilterPushDown::Unsupported; + filters.len() + ]) + } + + /// Get statistics for this table, if available + /// Although not presently used in mainline DataFusion, this allows implementation specific + /// behavior for downstream repositories, in conjunction with specialized optimizer rules to + /// perform operations such as re-ordering of joins. + fn statistics(&self) -> Option { + None + } + + /// Return an [`ExecutionPlan`] to insert data into this table, if + /// supported. + /// + /// The returned plan should return a single row in a UInt64 + /// column called "count" such as the following + /// + /// ```text + /// +-------+, + /// | count |, + /// +-------+, + /// | 6 |, + /// +-------+, + /// ``` + /// + /// # See Also + /// + /// See [`DataSinkExec`] for the common pattern of inserting a + /// streams of `RecordBatch`es as files to an ObjectStore. + /// + /// [`DataSinkExec`]: https://docs.rs/datafusion-datasource/latest/datafusion_datasource/sink/struct.DataSinkExec.html + async fn insert_into( + &self, + _state: &dyn Session, + _input: Arc, + _insert_op: InsertOp, + ) -> Result> { + not_impl_err!("Insert into not implemented for this table") + } + + /// Delete rows matching the filter predicates. + /// + /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). + /// Empty `filters` deletes all rows. + async fn delete_from( + &self, + _state: &dyn Session, + _filters: Vec, + ) -> Result> { + not_impl_err!("DELETE not supported for {} table", self.table_type()) + } + + /// Update rows matching the filter predicates. + /// + /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). + /// Empty `filters` updates all rows. + async fn update( + &self, + _state: &dyn Session, + _assignments: Vec<(String, Expr)>, + _filters: Vec, + ) -> Result> { + not_impl_err!("UPDATE not supported for {} table", self.table_type()) + } + + /// Remove all rows from the table. + /// + /// Should return an [ExecutionPlan] producing a single row with count (UInt64), + /// representing the number of rows removed. + async fn truncate(&self, _state: &dyn Session) -> Result> { + not_impl_err!("TRUNCATE not supported for {} table", self.table_type()) + } + + /// Merge rows from a source into this table. + /// + /// The `source` is an [`ExecutionPlan`] representing the USING clause. + /// The `merge_schema` contains the target columns followed by the source + /// columns, preserving their logical qualifiers. Providers can use this + /// schema to resolve the logical expressions against the combined rows + /// they construct while executing the merge. + /// The `on` condition is the join predicate from the ON clause. + /// The `clauses` describe the WHEN MATCHED / WHEN NOT MATCHED actions. + /// + /// Returns an [`ExecutionPlan`] producing a single row with `count` (UInt64). + async fn merge_into( + &self, + _state: &dyn Session, + _source: Arc, + _merge_schema: DFSchemaRef, + _on: Expr, + _clauses: Vec, + ) -> Result> { + not_impl_err!("MERGE INTO not supported for {} table", self.table_type()) + } +} + +impl dyn TableProvider { + /// Returns `true` if the table provider is of type `T`. + /// + /// Prefer this over `downcast_ref::().is_some()`. Works correctly when + /// called on `Arc` via auto-deref. + pub fn is(&self) -> bool { + (self as &dyn Any).is::() + } + + /// Attempts to downcast this table provider to a concrete type `T`, + /// returning `None` if the provider is not of that type. + /// + /// Works correctly when called on `Arc` via auto-deref, + /// unlike `(&arc as &dyn Any).downcast_ref::()` which would attempt to + /// downcast the `Arc` itself. + pub fn downcast_ref(&self) -> Option<&T> { + (self as &dyn Any).downcast_ref() + } +} + +/// Arguments for scanning a table with [`TableProvider::scan_with_args`]. +#[derive(Debug, Clone, Default)] +pub struct ScanArgs<'a> { + filters: Option<&'a [Expr]>, + projection: Option<&'a [usize]>, + limit: Option, + statistics_requests: &'a [StatisticsRequest], +} + +impl<'a> ScanArgs<'a> { + /// Set the column projection for the scan. + /// + /// The projection is a list of column indices from [`TableProvider::schema`] + /// that should be included in the scan results. If `None`, all columns are included. + /// + /// # Arguments + /// * `projection` - Optional slice of column indices to project + pub fn with_projection(mut self, projection: Option<&'a [usize]>) -> Self { + self.projection = projection; + self + } + + /// Get the column projection for the scan. + /// + /// Returns a reference to the projection column indices, or `None` if + /// no projection was specified (meaning all columns should be included). + pub fn projection(&self) -> Option<&'a [usize]> { + self.projection + } + + /// Set the filter expressions for the scan. + /// + /// Filters are boolean expressions that should be evaluated during the scan + /// to reduce the number of rows returned. All expressions are combined with AND logic. + /// Whether filters are actually pushed down depends on [`TableProvider::supports_filters_pushdown`]. + /// + /// # Arguments + /// * `filters` - Optional slice of filter expressions + pub fn with_filters(mut self, filters: Option<&'a [Expr]>) -> Self { + self.filters = filters; + self + } + + /// Get the filter expressions for the scan. + /// + /// Returns a reference to the filter expressions, or `None` if no filters were specified. + pub fn filters(&self) -> Option<&'a [Expr]> { + self.filters + } + + /// Set the maximum number of rows to return from the scan. + /// + /// If specified, the scan should return at most this many rows. This is typically + /// used to optimize queries with `LIMIT` clauses. + /// + /// # Arguments + /// * `limit` - Optional maximum number of rows to return + pub fn with_limit(mut self, limit: Option) -> Self { + self.limit = limit; + self + } + + /// Get the maximum number of rows to return from the scan. + /// + /// Returns the row limit, or `None` if no limit was specified. + pub fn limit(&self) -> Option { + self.limit + } + + /// Specifies the statistics the caller may use when optimizing the query. + /// + /// This is intended to allow the `TableProvider` to cheaply provide + /// statistics that may help, such as those it has in an in-memory catalog + /// or from some other metadata source. + /// + /// `TableProvider`s read these via [`Self::statistics_requests()`]; anything + /// a `TableProvider` cannot answer cheaply it simply ignores. DataFusion's + /// own `TableProvider`s ignore this field — it exists so a request can be + /// threaded from a custom optimizer rule (which annotates + /// `TableScan::statistics_requests`) through to a custom `TableProvider`. + pub fn with_statistics_requests( + mut self, + statistics_requests: &'a [StatisticsRequest], + ) -> Self { + self.statistics_requests = statistics_requests; + self + } + + /// Get the statistics requests for the scan. Empty if none were set. + /// + /// See [`Self::with_statistics_requests`] for more details + pub fn statistics_requests(&self) -> &'a [StatisticsRequest] { + self.statistics_requests + } +} + +/// Result of a table scan operation from [`TableProvider::scan_with_args`]. +#[derive(Debug, Clone)] +pub struct ScanResult { + /// The ExecutionPlan to run. + plan: Arc, +} + +impl ScanResult { + /// Create a new `ScanResult` with the given execution plan. + /// + /// # Arguments + /// * `plan` - The execution plan that will perform the table scan + pub fn new(plan: Arc) -> Self { + Self { plan } + } + + /// Get a reference to the execution plan for this scan result. + /// + /// Returns a reference to the [`ExecutionPlan`] that will perform + /// the actual table scanning and data retrieval. + pub fn plan(&self) -> &Arc { + &self.plan + } + + /// Consume this ScanResult and return the execution plan. + /// + /// Returns the owned [`ExecutionPlan`] that will perform + /// the actual table scanning and data retrieval. + pub fn into_inner(self) -> Arc { + self.plan + } +} + +impl From> for ScanResult { + fn from(plan: Arc) -> Self { + Self::new(plan) + } +} + +/// A factory which creates [`TableProvider`]s at runtime given a URL. +/// +/// For example, this can be used to create a table "on the fly" +/// from a directory of files only when that name is referenced. +#[async_trait] +pub trait TableProviderFactory: Debug + Sync + Send { + /// Create a TableProvider with the given url + async fn create( + &self, + state: &dyn Session, + cmd: &CreateExternalTable, + ) -> Result>; +} + +/// Describes arguments provided to the table function call. +pub struct TableFunctionArgs<'e, 's> { + /// Call arguments. + exprs: &'e [Expr], + /// Session within which the function is called. + session: &'s dyn Session, +} + +impl<'e, 's> TableFunctionArgs<'e, 's> { + /// Make a new [`TableFunctionArgs`]. + pub fn new(exprs: &'e [Expr], session: &'s dyn Session) -> Self { + Self { exprs, session } + } + + /// Get expressions passed as the called function arguments. + pub fn exprs(&self) -> &'e [Expr] { + self.exprs + } + + /// Get a session where the table function is called. + pub fn session(&self) -> &'s dyn Session { + self.session + } +} + +/// A trait for table function implementations +pub trait TableFunctionImpl: Debug + Sync + Send + Any { + /// Create a table provider + #[deprecated( + since = "53.0.0", + note = "Implement `TableFunctionImpl::call_with_args` instead" + )] + fn call(&self, _exprs: &[Expr]) -> Result> { + internal_err!( + "TableFunctionImpl::call is not implemented. Implement TableFunctionImpl::call_with_args instead." + ) + } + + /// Create a table provider + fn call_with_args(&self, args: TableFunctionArgs) -> Result> { + #[expect(deprecated)] + self.call(args.exprs) + } +} + +/// A table that uses a function to generate data +#[derive(Clone, Debug)] +pub struct TableFunction { + /// Name of the table function + name: String, + /// Function implementation + fun: Arc, +} + +impl TableFunction { + /// Create a new table function + pub fn new(name: String, fun: Arc) -> Self { + Self { name, fun } + } + + /// Get the name of the table function + pub fn name(&self) -> &str { + &self.name + } + + /// Get the implementation of the table function + pub fn function(&self) -> &Arc { + &self.fun + } + + /// Get the function implementation and generate a table + #[deprecated( + since = "53.0.0", + note = "Use `TableFunction::create_table_provider_with_args` instead" + )] + pub fn create_table_provider(&self, args: &[Expr]) -> Result> { + #[expect(deprecated)] + self.fun.call(args) + } + + /// Get the function implementation and generate a table + pub fn create_table_provider_with_args( + &self, + args: TableFunctionArgs, + ) -> Result> { + self.fun.call_with_args(args) + } +} diff --git a/datafusion/spark/Cargo.toml b/datafusion/spark/Cargo.toml index 14f9396d7656e..93987b553f2f5 100644 --- a/datafusion/spark/Cargo.toml +++ b/datafusion/spark/Cargo.toml @@ -71,7 +71,8 @@ url = { workspace = true } arrow = { workspace = true, features = ["test_utils"] } criterion = { workspace = true } # for SessionStateBuilderSpark tests -datafusion = { workspace = true, default-features = false } +datafusion = { workspace = true, default-features = false, features = ["sql"] } +tokio = { workspace = true, features = ["rt"] } [[bench]] harness = false diff --git a/datafusion/spark/benches/hex.rs b/datafusion/spark/benches/hex.rs index 9785371cc5827..38a59cb944e50 100644 --- a/datafusion/spark/benches/hex.rs +++ b/datafusion/spark/benches/hex.rs @@ -135,11 +135,21 @@ fn criterion_benchmark(c: &mut Criterion) { run_benchmark(c, "hex_utf8", size, Arc::new(data)); } + for &size in &sizes { + let data = generate_utf8_data(size, 0.0); + run_benchmark(c, "hex_utf8_no_nulls", size, Arc::new(data)); + } + for &size in &sizes { let data = generate_binary_data(size, null_density); run_benchmark(c, "hex_binary", size, Arc::new(data)); } + for &size in &sizes { + let data = generate_binary_data(size, 0.0); + run_benchmark(c, "hex_binary_no_nulls", size, Arc::new(data)); + } + for &size in &sizes { let data = generate_int64_dict_data(size, null_density); run_benchmark(c, "hex_int64_dict", size, Arc::new(data)); diff --git a/datafusion/spark/src/function/aggregate/avg.rs b/datafusion/spark/src/function/aggregate/avg.rs index 5f4d2c253a2dc..46e63013dbafb 100644 --- a/datafusion/spark/src/function/aggregate/avg.rs +++ b/datafusion/spark/src/function/aggregate/avg.rs @@ -289,7 +289,6 @@ where &mut self, values: &[ArrayRef], group_indices: &[usize], - _opt_filter: Option<&BooleanArray>, total_num_groups: usize, ) -> Result<()> { assert_eq!(values.len(), 2, "two arguments to merge_batch"); @@ -368,11 +367,6 @@ where Arc::new(counts) as ArrayRef, ]) } - - fn supports_convert_to_state(&self) -> bool { - true - } - fn size(&self) -> usize { self.counts.capacity() * size_of::() + self.sums.capacity() * size_of::() } @@ -388,12 +382,6 @@ mod tests { Ok(sum / count as f64) }) } - - #[test] - fn supports_convert_to_state() { - assert!(make_acc().supports_convert_to_state()); - } - #[test] fn convert_to_state_basic() { let acc = make_acc(); @@ -464,7 +452,6 @@ mod tests { acc.merge_batch( &state, &[0, 0, 0], - None, 1, // single group ) .unwrap(); @@ -486,7 +473,7 @@ mod tests { Some(3.0), ]))]; let state = acc.convert_to_state(&input, None).unwrap(); - acc.merge_batch(&state, &[0, 0, 0], None, 1).unwrap(); + acc.merge_batch(&state, &[0, 0, 0], 1).unwrap(); let result = acc.evaluate(EmitTo::All).unwrap(); let result = result.as_primitive::(); diff --git a/datafusion/spark/src/function/aggregate/collect.rs b/datafusion/spark/src/function/aggregate/collect.rs index 5af0fd39cca07..310bc1c890657 100644 --- a/datafusion/spark/src/function/aggregate/collect.rs +++ b/datafusion/spark/src/function/aggregate/collect.rs @@ -18,7 +18,7 @@ use arrow::array::ArrayRef; use arrow::datatypes::{DataType, Field, FieldRef}; use datafusion_common::utils::SingleRowListArrayBuilder; -use datafusion_common::{Result, ScalarValue}; +use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::function::{AccumulatorArgs, StateFieldsArgs}; use datafusion_expr::utils::format_state_name; use datafusion_expr::{Accumulator, AggregateUDFImpl, Signature, Volatility}; @@ -33,6 +33,19 @@ use std::sync::Arc; // - returns an empty list when all inputs are NULL // - does not support ordering +/// Build an empty list `ScalarValue` for a `List(element_type)` data type. +/// Used as the result for empty window frames and for groups whose inputs +/// were all NULL, matching Spark's `collect_list` / `collect_set` semantics. +fn empty_list_scalar(list_type: &DataType) -> Result { + let DataType::List(field) = list_type else { + return internal_err!( + "collect_list/collect_set expected List return type, got {list_type:?}" + ); + }; + let empty = arrow::array::new_empty_array(field.data_type()); + Ok(SingleRowListArrayBuilder::new(empty).build_list_scalar()) +} + // #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkCollectList { @@ -81,14 +94,17 @@ impl AggregateUDFImpl for SparkCollectList { } fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { - let field = &acc_args.expr_fields[0]; - let data_type = field.data_type().clone(); + let element_type = acc_args.expr_fields[0].data_type().clone(); let ignore_nulls = true; Ok(Box::new(NullToEmptyListAccumulator::new( - ArrayAggAccumulator::try_new(&data_type, ignore_nulls)?, - data_type, + ArrayAggAccumulator::try_new(&element_type, ignore_nulls)?, + acc_args.return_type().clone(), ))) } + + fn default_value(&self, data_type: &DataType) -> Result { + empty_list_scalar(data_type) + } } // @@ -139,14 +155,17 @@ impl AggregateUDFImpl for SparkCollectSet { } fn accumulator(&self, acc_args: AccumulatorArgs) -> Result> { - let field = &acc_args.expr_fields[0]; - let data_type = field.data_type().clone(); + let element_type = acc_args.expr_fields[0].data_type().clone(); let ignore_nulls = true; Ok(Box::new(NullToEmptyListAccumulator::new( - DistinctArrayAggAccumulator::try_new(&data_type, None, ignore_nulls)?, - data_type, + DistinctArrayAggAccumulator::try_new(&element_type, None, ignore_nulls)?, + acc_args.return_type().clone(), ))) } + + fn default_value(&self, data_type: &DataType) -> Result { + empty_list_scalar(data_type) + } } /// Wrapper accumulator that returns an empty list instead of NULL when all inputs are NULL. @@ -154,12 +173,12 @@ impl AggregateUDFImpl for SparkCollectSet { #[derive(Debug)] struct NullToEmptyListAccumulator { inner: T, - data_type: DataType, + list_type: DataType, } impl NullToEmptyListAccumulator { - pub fn new(inner: T, data_type: DataType) -> Self { - Self { inner, data_type } + pub fn new(inner: T, list_type: DataType) -> Self { + Self { inner, list_type } } } @@ -179,14 +198,21 @@ impl Accumulator for NullToEmptyListAccumulator { fn evaluate(&mut self) -> Result { let result = self.inner.evaluate()?; if result.is_null() { - let empty_array = arrow::array::new_empty_array(&self.data_type); - Ok(SingleRowListArrayBuilder::new(empty_array).build_list_scalar()) + empty_list_scalar(&self.list_type) } else { Ok(result) } } + fn retract_batch(&mut self, values: &[ArrayRef]) -> Result<()> { + self.inner.retract_batch(values) + } + + fn supports_retract_batch(&self) -> bool { + self.inner.supports_retract_batch() + } + fn size(&self) -> usize { - self.inner.size() + self.data_type.size() + self.inner.size() + self.list_type.size() } } diff --git a/datafusion/spark/src/function/aggregate/try_sum.rs b/datafusion/spark/src/function/aggregate/try_sum.rs index 3918dea0f5072..d1f99f4ebc0c3 100644 --- a/datafusion/spark/src/function/aggregate/try_sum.rs +++ b/datafusion/spark/src/function/aggregate/try_sum.rs @@ -190,7 +190,7 @@ fn update_decimal128( acc: &mut TrySumAccumulator, array: &PrimitiveArray, ) -> Result<()> { - let precision = acc.dec_precision.unwrap_or(38); + let precision = acc.dec_precision.unwrap_or(DECIMAL128_MAX_PRECISION); for v in array.iter().flatten() { let v_i128 = unsafe { std::mem::transmute_copy::(&v) }; diff --git a/datafusion/spark/src/function/array/repeat.rs b/datafusion/spark/src/function/array/repeat.rs index da9b19a768680..6effdf9a50f9a 100644 --- a/datafusion/spark/src/function/array/repeat.rs +++ b/datafusion/spark/src/function/array/repeat.rs @@ -74,9 +74,11 @@ impl ScalarUDFImpl for SparkArrayRepeat { // Coerce the second argument to Int64/UInt64 if it's a numeric type let second = match second_type { - DataType::Int8 | DataType::Int16 | DataType::Int32 | DataType::Int64 => { - DataType::Int64 - } + DataType::Int8 + | DataType::Int16 + | DataType::Int32 + | DataType::Int64 + | DataType::Null => DataType::Int64, DataType::UInt8 | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => { DataType::UInt64 } diff --git a/datafusion/spark/src/function/array/shuffle.rs b/datafusion/spark/src/function/array/shuffle.rs index 031dd17177577..2673c9155fe08 100644 --- a/datafusion/spark/src/function/array/shuffle.rs +++ b/datafusion/spark/src/function/array/shuffle.rs @@ -185,7 +185,7 @@ fn general_array_shuffle( if array.is_null(row_index) { nulls.push(false); offsets.push(offsets[row_index] + O::one()); - mutable.extend(0, 0, 1); + mutable.try_extend(0, 0, 1)?; continue; } nulls.push(true); @@ -200,7 +200,7 @@ fn general_array_shuffle( // Add shuffled elements for &index in &indices { - mutable.extend(0, index, index + 1); + mutable.try_extend(0, index, index + 1)?; } offsets.push(offsets[row_index] + O::usize_as(length)); @@ -239,7 +239,7 @@ fn fixed_size_array_shuffle( // skip the null value if array.is_null(row_index) { nulls.push(false); - mutable.extend(0, 0, value_length); + mutable.try_extend(0, 0, value_length)?; continue; } nulls.push(true); @@ -253,7 +253,7 @@ fn fixed_size_array_shuffle( // Add shuffled elements for &index in &indices { - mutable.extend(0, index, index + 1); + mutable.try_extend(0, index, index + 1)?; } } diff --git a/datafusion/spark/src/function/array/slice.rs b/datafusion/spark/src/function/array/slice.rs index bcd10a1bf7d79..f471565c4062a 100644 --- a/datafusion/spark/src/function/array/slice.rs +++ b/datafusion/spark/src/function/array/slice.rs @@ -157,7 +157,7 @@ fn calculate_start_end(args: &[ArrayRef]) -> Result<(ArrayRef, ArrayRef)> { } let start = start.value(row); let length = length.value(row); - let value_length = values.value(row).len() as i64; + let value_length = values.value_length(row) as i64; if start == 0 { return exec_err!("Start index must not be zero"); @@ -172,6 +172,15 @@ fn calculate_start_end(args: &[ArrayRef]) -> Result<(ArrayRef, ArrayRef)> { start }; + // Spark returns an empty array when the adjusted start lands before + // position 1 (e.g. slice([1], -2, 2)). array_slice would otherwise + // treat 0 the same as 1 and return the first element. + if adjusted_start_value < 1 { + adjusted_start.append_value(1); + end.append_value(0); + continue; + } + adjusted_start.append_value(adjusted_start_value); end.append_value(adjusted_start_value + (length - 1)); } diff --git a/datafusion/spark/src/function/bitmap/bitmap_count.rs b/datafusion/spark/src/function/bitmap/bitmap_count.rs index 89bea101afbe7..18d584868830b 100644 --- a/datafusion/spark/src/function/bitmap/bitmap_count.rs +++ b/datafusion/spark/src/function/bitmap/bitmap_count.rs @@ -28,8 +28,8 @@ use arrow::datatypes::{DataType, FieldRef, Int8Type, Int16Type, Int32Type, Int64 use datafusion_common::utils::take_function_args; use datafusion_common::{Result, internal_err}; use datafusion_expr::{ - Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, EncodingPreservation, ScalarFunctionArgs, ScalarUDFImpl, + Signature, TypeSignatureClass, Volatility, }; use datafusion_functions::downcast_arg; use datafusion_functions::utils::make_scalar_function; @@ -49,7 +49,10 @@ impl BitmapCount { pub fn new() -> Self { Self { signature: Signature::coercible( - vec![Coercion::new_exact(TypeSignatureClass::Binary)], + vec![ + Coercion::new_exact(TypeSignatureClass::Binary) + .with_encoding_preservation(EncodingPreservation::dictionary()), + ], Volatility::Immutable, ), } diff --git a/datafusion/spark/src/function/datetime/mod.rs b/datafusion/spark/src/function/datetime/mod.rs index 3133ed7337f25..70ab024c329aa 100644 --- a/datafusion/spark/src/function/datetime/mod.rs +++ b/datafusion/spark/src/function/datetime/mod.rs @@ -26,11 +26,13 @@ pub mod from_utc_timestamp; pub mod last_day; pub mod make_dt_interval; pub mod make_interval; +pub mod monthname; pub mod next_day; pub mod time_trunc; pub mod to_utc_timestamp; pub mod trunc; pub mod unix; +pub mod weekday; use datafusion_expr::ScalarUDF; use datafusion_functions::make_udf_function; @@ -52,11 +54,13 @@ make_udf_function!(extract::SparkSecond, second); make_udf_function!(last_day::SparkLastDay, last_day); make_udf_function!(make_dt_interval::SparkMakeDtInterval, make_dt_interval); make_udf_function!(make_interval::SparkMakeInterval, make_interval); +make_udf_function!(monthname::SparkMonthName, monthname); make_udf_function!(next_day::SparkNextDay, next_day); make_udf_function!(time_trunc::SparkTimeTrunc, time_trunc); make_udf_function!(to_utc_timestamp::SparkToUtcTimestamp, to_utc_timestamp); make_udf_function!(trunc::SparkTrunc, trunc); make_udf_function!(unix::SparkUnixDate, unix_date); +make_udf_function!(weekday::SparkWeekDay, weekday); make_udf_function!( unix::SparkUnixTimestamp, unix_micros, @@ -117,6 +121,11 @@ pub mod expr_fn { "Make interval from years, months, weeks, days, hours, mins and secs.", years months weeks days hours mins secs )); + export_functions!(( + monthname, + "Returns the three-letter abbreviated month name from a date or timestamp.", + arg1 + )); // TODO: add once ANSI support is added: // "When both of the input parameters are not NULL and day_of_week is an invalid input, the function throws SparkIllegalArgumentException if spark.sql.ansi.enabled is set to true, otherwise NULL." export_functions!(( @@ -179,6 +188,11 @@ pub mod expr_fn { "Returns the number of seconds since epoch (1970-01-01 00:00:00 UTC) for the given timestamp `ts`.", ts )); + export_functions!(( + weekday, + "Returns the day of the week for date/timestamp as an integer where Monday = 0, Tuesday = 1, ..., Sunday = 6.", + arg1 + )); } pub fn functions() -> Vec> { @@ -195,6 +209,7 @@ pub fn functions() -> Vec> { make_dt_interval(), make_interval(), minute(), + monthname(), next_day(), second(), time_trunc(), @@ -204,5 +219,6 @@ pub fn functions() -> Vec> { unix_micros(), unix_millis(), unix_seconds(), + weekday(), ] } diff --git a/datafusion/spark/src/function/datetime/monthname.rs b/datafusion/spark/src/function/datetime/monthname.rs new file mode 100644 index 0000000000000..6cfa9c0a9212e --- /dev/null +++ b/datafusion/spark/src/function/datetime/monthname.rs @@ -0,0 +1,115 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{AsArray, StringArray}; +use arrow::compute::{DatePart, date_part}; +use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::types::{NativeType, logical_date}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, ScalarValue, internal_err}; +use datafusion_expr::{ + Coercion, ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, + Signature, TypeSignatureClass, Volatility, +}; + +const MONTH_NAMES: [&str; 12] = [ + "Jan", "Feb", "Mar", "Apr", "May", "Jun", "Jul", "Aug", "Sep", "Oct", "Nov", "Dec", +]; + +fn month_number_to_name(month: i32) -> Option<&'static str> { + MONTH_NAMES.get((month - 1) as usize).copied() +} + +/// Spark-compatible `monthname` expression. +/// Returns the three-letter abbreviated month name from a date or timestamp. +/// +/// +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkMonthName { + signature: Signature, +} + +impl Default for SparkMonthName { + fn default() -> Self { + Self::new() + } +} + +impl SparkMonthName { + pub fn new() -> Self { + Self { + signature: Signature::coercible( + vec![Coercion::new_implicit( + TypeSignatureClass::Native(logical_date()), + vec![TypeSignatureClass::Timestamp], + NativeType::Date, + )], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for SparkMonthName { + fn name(&self) -> &str { + "monthname" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); + Ok(Arc::new(Field::new(self.name(), DataType::Utf8, nullable))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [arg] = take_function_args(self.name(), args.args)?; + match arg { + ColumnarValue::Scalar(scalar) => { + if scalar.is_null() { + return Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))); + } + let arr = scalar.to_array_of_size(1)?; + let month_arr = date_part(&arr, DatePart::Month)?; + let month_val = month_arr + .as_primitive::() + .value(0); + let name = month_number_to_name(month_val).map(|s| s.to_string()); + Ok(ColumnarValue::Scalar(ScalarValue::Utf8(name))) + } + ColumnarValue::Array(arr) => { + let month_arr = date_part(&arr, DatePart::Month)?; + let int_arr = month_arr.as_primitive::(); + + let result: StringArray = int_arr + .iter() + .map(|maybe_month| maybe_month.and_then(month_number_to_name)) + .collect(); + + Ok(ColumnarValue::Array(Arc::new(result))) + } + } + } +} diff --git a/datafusion/spark/src/function/datetime/next_day.rs b/datafusion/spark/src/function/datetime/next_day.rs index 2241043d44cd7..09d7de7b4a4de 100644 --- a/datafusion/spark/src/function/datetime/next_day.rs +++ b/datafusion/spark/src/function/datetime/next_day.rs @@ -19,7 +19,7 @@ use std::sync::Arc; use arrow::array::{ArrayRef, AsArray, Date32Array, StringArrayType}; use arrow::datatypes::{DataType, Date32Type, Field, FieldRef}; -use chrono::{Datelike, Duration, Weekday}; +use chrono::{Datelike, Weekday}; use datafusion_common::{Result, ScalarValue, exec_err, internal_err}; use datafusion_expr::{ ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, @@ -210,7 +210,7 @@ where fn spark_next_day(days: i32, day_of_week: &str) -> Option { let date = Date32Type::to_naive_date_opt(days)?; - let day_of_week = day_of_week.trim().to_uppercase(); + let day_of_week = day_of_week.to_uppercase(); let day_of_week = match day_of_week.as_str() { "MO" | "MON" | "MONDAY" => Some("MONDAY"), "TU" | "TUE" | "TUESDAY" => Some("TUESDAY"), @@ -229,11 +229,17 @@ fn spark_next_day(days: i32, day_of_week: &str) -> Option { if let Some(day_of_week) = day_of_week { let day_of_week = day_of_week.parse::(); match day_of_week { - Ok(day_of_week) => Some(Date32Type::from_naive_date( - date + Duration::days( - (7 - date.weekday().days_since(day_of_week)) as i64, - ), - )), + Ok(day_of_week) => { + // Advance 1..=7 days from `days` to the next occurrence of + // `day_of_week`. Compute the result on the epoch day directly + // instead of constructing a `NaiveDate`: the result can land + // past `NaiveDate::MAX` (epoch day 95026236), and building that + // date panics (`NaiveDate + TimeDelta overflowed`). Spark's + // `DateTimeUtils.getNextDateForDayOfWeek` is pure `Int` + // arithmetic and keeps producing a value up to `Int.MaxValue`. + let delta = 7 - date.weekday().days_since(day_of_week) as i32; + days.checked_add(delta) + } Err(_) => { // TODO: if spark.sql.ansi.enabled is false, // returns NULL instead of an error for a malformed dayOfWeek. @@ -279,4 +285,23 @@ mod tests { assert_eq!(field.data_type(), &DataType::Date32); assert!(field.is_nullable()); } + + #[test] + fn next_day_rejects_whitespace_padded_day_names() { + let monday = 19723; // 2024-01-01 + assert_eq!(spark_next_day(monday, " MO "), None); + } + + #[test] + fn next_day_handles_far_future_start_dates() { + // Regression for #23891: for start dates near the end of the + // representable `Date32` range, the next occurrence can land past + // `chrono::NaiveDate::MAX` (epoch day 95026236). Computing the result + // on the epoch day directly (as Spark does) must return a value rather + // than panicking with `NaiveDate + TimeDelta overflowed`. + // + // 95026236 is a Monday, so `next_day(.., "Mon")` advances a full week. + assert_eq!(spark_next_day(95026236, "Mon"), Some(95026243)); + assert_eq!(spark_next_day(95026230, "Tue"), Some(95026237)); + } } diff --git a/datafusion/spark/src/function/datetime/weekday.rs b/datafusion/spark/src/function/datetime/weekday.rs new file mode 100644 index 0000000000000..b9ac7e43750ba --- /dev/null +++ b/datafusion/spark/src/function/datetime/weekday.rs @@ -0,0 +1,191 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::AsArray; +use arrow::compute::{DatePart, date_part}; +use arrow::datatypes::{DataType, Field, FieldRef, Int32Type}; +use datafusion_common::types::{NativeType, logical_date}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, ScalarValue, internal_err}; +use datafusion_expr::{ + Coercion, ColumnarValue, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, + Signature, TypeSignatureClass, Volatility, +}; + +/// Spark-compatible `weekday` expression. +/// Returns the day of the week for a date or timestamp as an integer index where +/// Monday = 0, Tuesday = 1, ..., Sunday = 6. +/// +/// Note: this differs from `dayofweek`, which is 1-indexed with Sunday = 1. +/// +/// +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkWeekDay { + signature: Signature, +} + +impl Default for SparkWeekDay { + fn default() -> Self { + Self::new() + } +} + +impl SparkWeekDay { + pub fn new() -> Self { + Self { + signature: Signature::coercible( + vec![Coercion::new_implicit( + TypeSignatureClass::Native(logical_date()), + vec![TypeSignatureClass::Timestamp], + NativeType::Date, + )], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for SparkWeekDay { + fn name(&self) -> &str { + "weekday" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + internal_err!("return_field_from_args should be used instead") + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); + Ok(Arc::new(Field::new(self.name(), DataType::Int32, nullable))) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let [arg] = take_function_args(self.name(), args.args)?; + match arg { + ColumnarValue::Scalar(scalar) => { + if scalar.is_null() { + return Ok(ColumnarValue::Scalar(ScalarValue::Int32(None))); + } + let arr = scalar.to_array_of_size(1)?; + // `DayOfWeekMonday0` returns 0..=6 with Monday = 0, which + // matches Spark `weekday` semantics exactly. + let weekday_arr = date_part(&arr, DatePart::DayOfWeekMonday0)?; + let value = weekday_arr.as_primitive::().value(0); + Ok(ColumnarValue::Scalar(ScalarValue::Int32(Some(value)))) + } + ColumnarValue::Array(arr) => { + let weekday_arr = date_part(&arr, DatePart::DayOfWeekMonday0)?; + Ok(ColumnarValue::Array(weekday_arr)) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Date32Array, Int32Array}; + + #[test] + fn test_weekday_return_field_nullability_matches_input() { + let func = SparkWeekDay::new(); + + let non_nullable_arg = Arc::new(Field::new("arg", DataType::Date32, false)); + let nullable_arg = Arc::new(Field::new("arg", DataType::Date32, true)); + + let non_nullable_out = func + .return_field_from_args(ReturnFieldArgs { + arg_fields: &[Arc::clone(&non_nullable_arg)], + scalar_arguments: &[None], + }) + .expect("non-nullable arg should succeed"); + assert_eq!(non_nullable_out.data_type(), &DataType::Int32); + assert!(!non_nullable_out.is_nullable()); + + let nullable_out = func + .return_field_from_args(ReturnFieldArgs { + arg_fields: &[Arc::clone(&nullable_arg)], + scalar_arguments: &[None], + }) + .expect("nullable arg should succeed"); + assert_eq!(nullable_out.data_type(), &DataType::Int32); + assert!(nullable_out.is_nullable()); + } + + #[test] + fn test_weekday_scalar() -> Result<()> { + let func = SparkWeekDay::new(); + + // 2024-03-15 is a Friday -> Spark weekday = 4 (Mon=0). + let result = func.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Scalar(ScalarValue::Date32(Some(19797)))], + arg_fields: vec![Arc::new(Field::new("arg", DataType::Date32, true))], + number_rows: 1, + return_field: Arc::new(Field::new("weekday", DataType::Int32, true)), + config_options: Arc::new(Default::default()), + })?; + match result { + ColumnarValue::Scalar(ScalarValue::Int32(Some(v))) => assert_eq!(v, 4), + other => panic!("unexpected result: {other:?}"), + } + + // NULL input -> NULL output. + let result = func.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Scalar(ScalarValue::Date32(None))], + arg_fields: vec![Arc::new(Field::new("arg", DataType::Date32, true))], + number_rows: 1, + return_field: Arc::new(Field::new("weekday", DataType::Int32, true)), + config_options: Arc::new(Default::default()), + })?; + match result { + ColumnarValue::Scalar(ScalarValue::Int32(None)) => {} + other => panic!("unexpected result: {other:?}"), + } + + Ok(()) + } + + #[test] + fn test_weekday_array() -> Result<()> { + let func = SparkWeekDay::new(); + + // 2024-01-01 Mon(0), 2024-01-06 Sat(5), 2024-01-07 Sun(6), NULL. + let input = Date32Array::from(vec![Some(19723), Some(19728), Some(19729), None]); + let result = func.invoke_with_args(ScalarFunctionArgs { + args: vec![ColumnarValue::Array(Arc::new(input))], + arg_fields: vec![Arc::new(Field::new("arg", DataType::Date32, true))], + number_rows: 4, + return_field: Arc::new(Field::new("weekday", DataType::Int32, true)), + config_options: Arc::new(Default::default()), + })?; + match result { + ColumnarValue::Array(arr) => { + let expected = Int32Array::from(vec![Some(0), Some(5), Some(6), None]); + assert_eq!(arr.as_primitive::(), &expected); + } + other => panic!("unexpected result: {other:?}"), + } + + Ok(()) + } +} diff --git a/datafusion/spark/src/function/hash/sha1.rs b/datafusion/spark/src/function/hash/sha1.rs index dd9009eb8233f..05a224f33f25a 100644 --- a/datafusion/spark/src/function/hash/sha1.rs +++ b/datafusion/spark/src/function/hash/sha1.rs @@ -24,6 +24,7 @@ use datafusion_common::cast::{ as_large_binary_array, }; use datafusion_common::types::{NativeType, logical_string}; +use datafusion_common::utils::hex::{HexCase, encode_bytes}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, internal_err}; use datafusion_expr::{ @@ -89,18 +90,9 @@ impl ScalarUDFImpl for SparkSha1 { } } -/// Hex encoding lookup table for fast byte-to-hex conversion -const HEX_CHARS_LOWER: &[u8; 16] = b"0123456789abcdef"; - #[inline] fn spark_sha1_digest(value: &[u8]) -> String { - let result = Sha1::digest(value); - let mut s = String::with_capacity(result.len() * 2); - for &b in result.as_slice() { - s.push(HEX_CHARS_LOWER[(b >> 4) as usize] as char); - s.push(HEX_CHARS_LOWER[(b & 0x0f) as usize] as char); - } - s + encode_bytes(&Sha1::digest(value), HexCase::Lower) } fn spark_sha1_impl<'a>(input: impl Iterator>) -> ArrayRef { diff --git a/datafusion/spark/src/function/hash/sha2.rs b/datafusion/spark/src/function/hash/sha2.rs index 38fa0cc643751..541df2957669e 100644 --- a/datafusion/spark/src/function/hash/sha2.rs +++ b/datafusion/spark/src/function/hash/sha2.rs @@ -20,6 +20,7 @@ use arrow::datatypes::{DataType, Int32Type}; use datafusion_common::types::{ NativeType, logical_binary, logical_int32, logical_string, }; +use datafusion_common::utils::hex::{HexCase, encode_bytes}; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, ScalarValue, internal_err}; use datafusion_expr::{ @@ -112,22 +113,22 @@ impl ScalarUDFImpl for SparkSha2 { 224 => { let mut digest = sha2::Sha224::default(); digest.update(bytes); - Some(hex_encode(digest.finalize())) + Some(encode_bytes(&digest.finalize(), HexCase::Lower)) } 0 | 256 => { let mut digest = sha2::Sha256::default(); digest.update(bytes); - Some(hex_encode(digest.finalize())) + Some(encode_bytes(&digest.finalize(), HexCase::Lower)) } 384 => { let mut digest = sha2::Sha384::default(); digest.update(bytes); - Some(hex_encode(digest.finalize())) + Some(encode_bytes(&digest.finalize(), HexCase::Lower)) } 512 => { let mut digest = sha2::Sha512::default(); digest.update(bytes); - Some(hex_encode(digest.finalize())) + Some(encode_bytes(&digest.finalize(), HexCase::Lower)) } _ => None, }; @@ -222,22 +223,22 @@ where (Some(value), Some(224)) => { let mut digest = sha2::Sha224::default(); digest.update(value); - Some(hex_encode(digest.finalize())) + Some(encode_bytes(&digest.finalize(), HexCase::Lower)) } (Some(value), Some(0 | 256)) => { let mut digest = sha2::Sha256::default(); digest.update(value); - Some(hex_encode(digest.finalize())) + Some(encode_bytes(&digest.finalize(), HexCase::Lower)) } (Some(value), Some(384)) => { let mut digest = sha2::Sha384::default(); digest.update(value); - Some(hex_encode(digest.finalize())) + Some(encode_bytes(&digest.finalize(), HexCase::Lower)) } (Some(value), Some(512)) => { let mut digest = sha2::Sha512::default(); digest.update(value); - Some(hex_encode(digest.finalize())) + Some(encode_bytes(&digest.finalize(), HexCase::Lower)) } // Unknown bit-lengths go to null, same as in Spark _ => None, @@ -245,19 +246,3 @@ where .collect::(); Arc::new(array) } - -const HEX_CHARS: [u8; 16] = *b"0123456789abcdef"; - -#[inline] -fn hex_encode>(data: T) -> String { - let bytes = data.as_ref(); - let mut out = Vec::with_capacity(bytes.len() * 2); - for &b in bytes { - let hi = b >> 4; - let lo = b & 0x0F; - out.push(HEX_CHARS[hi as usize]); - out.push(HEX_CHARS[lo as usize]); - } - // SAFETY: out contains only ASCII - unsafe { String::from_utf8_unchecked(out) } -} diff --git a/datafusion/spark/src/function/hash/xxhash64.rs b/datafusion/spark/src/function/hash/xxhash64.rs index 5dca47bcb8984..9d02a51b2217e 100644 --- a/datafusion/spark/src/function/hash/xxhash64.rs +++ b/datafusion/spark/src/function/hash/xxhash64.rs @@ -363,12 +363,17 @@ mod tests { #[test] fn test_xxhash64_fixed_size_binary() { - let array = FixedSizeBinaryArray::from(vec![ - Some(&[0x01, 0x02, 0x03, 0x04][..]), - Some(&[0x05, 0x06, 0x07, 0x08][..]), - None, - Some(&[0x00, 0x00, 0x00, 0x00][..]), - ]); + let array = FixedSizeBinaryArray::try_from_sparse_iter_with_size( + vec![ + Some(&[0x01, 0x02, 0x03, 0x04][..]), + Some(&[0x05, 0x06, 0x07, 0x08][..]), + None, + Some(&[0x00, 0x00, 0x00, 0x00][..]), + ] + .into_iter(), + 4, + ) + .unwrap(); let array_ref: ArrayRef = Arc::new(array); let mut hashes = vec![DEFAULT_SEED; 4]; diff --git a/datafusion/spark/src/function/map/map_from_arrays.rs b/datafusion/spark/src/function/map/map_from_arrays.rs index 692e837d00f5e..92dea2720fbfc 100644 --- a/datafusion/spark/src/function/map/map_from_arrays.rs +++ b/datafusion/spark/src/function/map/map_from_arrays.rs @@ -22,6 +22,7 @@ use crate::function::map::utils::{ use arrow::array::{Array, ArrayRef, NullArray}; use arrow::compute::kernels::cast; use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::config::MapKeyDedupPolicy; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, internal_err}; use datafusion_expr::{ @@ -81,11 +82,16 @@ impl ScalarUDFImpl for MapFromArrays { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(map_from_arrays_inner, vec![])(&args.args) + let last_value_wins = + args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin; + make_scalar_function( + move |args: &[ArrayRef]| map_from_arrays_inner(args, last_value_wins), + vec![], + )(&args.args) } } -fn map_from_arrays_inner(args: &[ArrayRef]) -> Result { +fn map_from_arrays_inner(args: &[ArrayRef], last_value_wins: bool) -> Result { let [keys, values] = take_function_args("map_from_arrays", args)?; if *keys.data_type() == DataType::Null || *values.data_type() == DataType::Null { @@ -105,6 +111,7 @@ fn map_from_arrays_inner(args: &[ArrayRef]) -> Result { &get_list_offsets(values)?, keys.nulls(), values.nulls(), + last_value_wins, ) } diff --git a/datafusion/spark/src/function/map/map_from_entries.rs b/datafusion/spark/src/function/map/map_from_entries.rs index facf9f8c53473..69ce352694bd1 100644 --- a/datafusion/spark/src/function/map/map_from_entries.rs +++ b/datafusion/spark/src/function/map/map_from_entries.rs @@ -24,6 +24,7 @@ use crate::function::map::utils::{ use arrow::array::{Array, ArrayRef, NullBufferBuilder, StructArray}; use arrow::buffer::NullBuffer; use arrow::datatypes::{DataType, Field, FieldRef}; +use datafusion_common::config::MapKeyDedupPolicy; use datafusion_common::utils::take_function_args; use datafusion_common::{Result, exec_err, internal_err}; use datafusion_expr::{ @@ -101,11 +102,16 @@ impl ScalarUDFImpl for MapFromEntries { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(map_from_entries_inner, vec![])(&args.args) + let last_value_wins = + args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin; + make_scalar_function( + move |args: &[ArrayRef]| map_from_entries_inner(args, last_value_wins), + vec![], + )(&args.args) } } -fn map_from_entries_inner(args: &[ArrayRef]) -> Result { +fn map_from_entries_inner(args: &[ArrayRef], last_value_wins: bool) -> Result { let [entries] = take_function_args("map_from_entries", args)?; let entries_offsets = get_list_offsets(entries)?; let entries_values = get_list_values(entries)?; @@ -148,6 +154,7 @@ fn map_from_entries_inner(args: &[ArrayRef]) -> Result { &entries_offsets, None, res_nulls.as_ref(), + last_value_wins, ) } diff --git a/datafusion/spark/src/function/map/str_to_map.rs b/datafusion/spark/src/function/map/str_to_map.rs index d0f4cf03cd432..abb4bd04762a3 100644 --- a/datafusion/spark/src/function/map/str_to_map.rs +++ b/datafusion/spark/src/function/map/str_to_map.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashSet; +use std::collections::{HashMap, HashSet}; use std::sync::Arc; use arrow::array::{ @@ -33,6 +33,7 @@ use datafusion_expr::{ }; use crate::function::map::utils::map_type_from_key_value_types; +use datafusion_common::config::MapKeyDedupPolicy; const DEFAULT_PAIR_DELIM: &str = ","; const DEFAULT_KV_DELIM: &str = ":"; @@ -48,11 +49,10 @@ const DEFAULT_KV_DELIM: &str = ":"; /// - keyValueDelim: Delimiter between key and value (default: ':') /// /// # Duplicate Key Handling -/// Uses EXCEPTION behavior (Spark 3.0+ default): errors on duplicate keys. -/// See `spark.sql.mapKeyDedupPolicy`: -/// -/// -/// TODO: Support configurable `spark.sql.mapKeyDedupPolicy` (LAST_WIN) in a follow-up PR. +/// Mirrors Spark's [`spark.sql.mapKeyDedupPolicy`](https://github.com/apache/spark/blob/v4.0.0/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4502-L4511), +/// wired through DataFusion's `datafusion.spark.map_key_dedup_policy`: +/// - `EXCEPTION` (default): error on duplicate keys. +/// - `LAST_WIN`: keep the last occurrence of each duplicate key. #[derive(Debug, PartialEq, Eq, Hash)] pub struct SparkStrToMap { signature: Signature, @@ -102,22 +102,32 @@ impl ScalarUDFImpl for SparkStrToMap { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + let last_value_wins = + args.config_options.spark.map_key_dedup_policy == MapKeyDedupPolicy::LastWin; let arrays: Vec = ColumnarValue::values_to_arrays(&args.args)?; - let result = str_to_map_inner(&arrays)?; + let result = str_to_map_inner(&arrays, last_value_wins)?; Ok(ColumnarValue::Array(result)) } } -fn str_to_map_inner(args: &[ArrayRef]) -> Result { +fn str_to_map_inner(args: &[ArrayRef], last_value_wins: bool) -> Result { match args.len() { 1 => match args[0].data_type() { - DataType::Utf8 => str_to_map_impl(as_string_array(&args[0])?, None, None), - DataType::LargeUtf8 => { - str_to_map_impl(as_large_string_array(&args[0])?, None, None) - } - DataType::Utf8View => { - str_to_map_impl(as_string_view_array(&args[0])?, None, None) + DataType::Utf8 => { + str_to_map_impl(as_string_array(&args[0])?, None, None, last_value_wins) } + DataType::LargeUtf8 => str_to_map_impl( + as_large_string_array(&args[0])?, + None, + None, + last_value_wins, + ), + DataType::Utf8View => str_to_map_impl( + as_string_view_array(&args[0])?, + None, + None, + last_value_wins, + ), other => exec_err!( "Unsupported data type {other:?} for str_to_map, \ expected Utf8, LargeUtf8, or Utf8View" @@ -128,16 +138,19 @@ fn str_to_map_inner(args: &[ArrayRef]) -> Result { as_string_array(&args[0])?, Some(as_string_array(&args[1])?), None, + last_value_wins, ), (DataType::LargeUtf8, DataType::LargeUtf8) => str_to_map_impl( as_large_string_array(&args[0])?, Some(as_large_string_array(&args[1])?), None, + last_value_wins, ), (DataType::Utf8View, DataType::Utf8View) => str_to_map_impl( as_string_view_array(&args[0])?, Some(as_string_view_array(&args[1])?), None, + last_value_wins, ), (t1, t2) => exec_err!( "Unsupported data types ({t1:?}, {t2:?}) for str_to_map, \ @@ -153,12 +166,14 @@ fn str_to_map_inner(args: &[ArrayRef]) -> Result { as_string_array(&args[0])?, Some(as_string_array(&args[1])?), Some(as_string_array(&args[2])?), + last_value_wins, ), (DataType::LargeUtf8, DataType::LargeUtf8, DataType::LargeUtf8) => { str_to_map_impl( as_large_string_array(&args[0])?, Some(as_large_string_array(&args[1])?), Some(as_large_string_array(&args[2])?), + last_value_wins, ) } (DataType::Utf8View, DataType::Utf8View, DataType::Utf8View) => { @@ -166,6 +181,7 @@ fn str_to_map_inner(args: &[ArrayRef]) -> Result { as_string_view_array(&args[0])?, Some(as_string_view_array(&args[1])?), Some(as_string_view_array(&args[2])?), + last_value_wins, ) } (t1, t2, t3) => exec_err!( @@ -181,6 +197,7 @@ fn str_to_map_impl<'a, V: StringArrayType<'a> + Copy>( text_array: V, pair_delim_array: Option, kv_delim_array: Option, + last_value_wins: bool, ) -> Result { let num_rows = text_array.len(); @@ -206,6 +223,10 @@ fn str_to_map_impl<'a, V: StringArrayType<'a> + Copy>( ); let mut seen_keys = HashSet::new(); + // LAST_WIN buffers pairs to support in-place value overwrite at the key's + // first-seen position — matches Spark's `ArrayBasedMapBuilder`. + let mut pairs: Vec<(&str, Option<&str>)> = Vec::new(); + let mut key_positions: HashMap<&str, usize> = HashMap::new(); for row_idx in 0..num_rows { if combined_nulls.as_ref().is_some_and(|n| n.is_null(row_idx)) { map_builder.append(false)?; @@ -226,31 +247,56 @@ fn str_to_map_impl<'a, V: StringArrayType<'a> + Copy>( continue; } - seen_keys.clear(); - for pair in text.split(pair_delim) { - if pair.is_empty() { - continue; + if last_value_wins { + pairs.clear(); + key_positions.clear(); + for pair in text.split(pair_delim) { + if pair.is_empty() { + continue; + } + let mut kv_iter = pair.splitn(2, kv_delim); + let key = kv_iter.next().unwrap_or(""); + let value = kv_iter.next(); + match key_positions.get(key) { + Some(&idx) => pairs[idx].1 = value, + None => { + key_positions.insert(key, pairs.len()); + pairs.push((key, value)); + } + } + } + for (key, value) in &pairs { + map_builder.keys().append_value(key); + match value { + Some(v) => map_builder.values().append_value(v), + None => map_builder.values().append_null(), + } } + } else { + seen_keys.clear(); + for pair in text.split(pair_delim) { + if pair.is_empty() { + continue; + } - let mut kv_iter = pair.splitn(2, kv_delim); - let key = kv_iter.next().unwrap_or(""); - let value = kv_iter.next(); + let mut kv_iter = pair.splitn(2, kv_delim); + let key = kv_iter.next().unwrap_or(""); + let value = kv_iter.next(); - // TODO: Support LAST_WIN policy via spark.sql.mapKeyDedupPolicy config - // EXCEPTION policy: error on duplicate keys (Spark 3.0+ default) - if !seen_keys.insert(key) { - return exec_err!( - "Duplicate map key '{key}' was found, please check the input data. \ - If you want to remove the duplicated keys, you can set \ - spark.sql.mapKeyDedupPolicy to \"LAST_WIN\" so that the key \ - inserted at last takes precedence." - ); - } + if !seen_keys.insert(key) { + return exec_err!( + "[DUPLICATED_MAP_KEY] Duplicate map key '{key}' was found, \ + please check the input data. To allow duplicate keys with \ + last-value-wins semantics, set \ + `datafusion.spark.map_key_dedup_policy` to `LAST_WIN`." + ); + } - map_builder.keys().append_value(key); - match value { - Some(v) => map_builder.values().append_value(v), - None => map_builder.values().append_null(), + map_builder.keys().append_value(key); + match value { + Some(v) => map_builder.values().append_value(v), + None => map_builder.values().append_null(), + } } } map_builder.append(true)?; diff --git a/datafusion/spark/src/function/map/utils.rs b/datafusion/spark/src/function/map/utils.rs index f5fff0c4b4c46..fa6b2a960dabb 100644 --- a/datafusion/spark/src/function/map/utils.rs +++ b/datafusion/spark/src/function/map/utils.rs @@ -16,12 +16,14 @@ // under the License. use std::borrow::Cow; -use std::collections::HashSet; +use std::collections::HashMap; use std::sync::Arc; -use arrow::array::{Array, ArrayRef, AsArray, BooleanBuilder, MapArray, StructArray}; +use arrow::array::{ + Array, ArrayRef, AsArray, BooleanBuilder, Int32Array, MapArray, StructArray, +}; use arrow::buffer::{NullBuffer, OffsetBuffer}; -use arrow::compute::filter; +use arrow::compute::{filter, take}; use arrow::datatypes::{DataType, Field, Fields}; use datafusion_common::{Result, ScalarValue, exec_err}; @@ -111,13 +113,13 @@ pub fn map_type_from_key_value_types( /// So the inputs can be [`ListArray`](`arrow::array::ListArray`)/[`LargeListArray`](`arrow::array::LargeListArray`)/[`FixedSizeListArray`](`arrow::array::FixedSizeListArray`)
/// To preserve the row info, [`offsets`](arrow::array::ListArray::offsets) and [`nulls`](arrow::array::ListArray::nulls) for both keys and values need to be provided
/// [`FixedSizeListArray`](`arrow::array::FixedSizeListArray`) has no `offsets`, so they can be generated as a cumulative sum of it's `Size` -/// 2. Spark provides [spark.sql.mapKeyDedupPolicy](https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961) -/// to handle duplicate keys
-/// For now, configurable functions are not supported by Datafusion
-/// So more permissive `LAST_WIN` option is used in this implementation (instead of `EXCEPTION`)
-/// `EXCEPTION` behaviour can still be achieved externally in cost of performance:
-/// `when(array_length(array_distinct(keys)) == array_length(keys), constructed_map)`
-/// `.otherwise(raise_error("duplicate keys occurred during map construction"))` +/// 2. Duplicate-key handling mirrors Spark's +/// [spark.sql.mapKeyDedupPolicy](https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961) +/// and is driven by `last_value_wins`: +/// - `false` (Spark's default `EXCEPTION`): raise `[DUPLICATED_MAP_KEY]` on any duplicate. +/// - `true` (`LAST_WIN`): keep the last occurrence of each duplicate key. +/// +/// Callers wire this from `datafusion.spark.map_key_dedup_policy`. pub fn map_from_keys_values_offsets_nulls( flat_keys: &ArrayRef, flat_values: &ArrayRef, @@ -125,6 +127,7 @@ pub fn map_from_keys_values_offsets_nulls( values_offsets: &[i32], keys_nulls: Option<&NullBuffer>, values_nulls: Option<&NullBuffer>, + last_value_wins: bool, ) -> Result { let (keys, values, offsets) = map_deduplicate_keys( flat_keys, @@ -133,6 +136,7 @@ pub fn map_from_keys_values_offsets_nulls( values_offsets, keys_nulls, values_nulls, + last_value_wins, )?; let nulls = NullBuffer::union(keys_nulls, values_nulls); @@ -155,6 +159,7 @@ fn map_deduplicate_keys( values_offsets: &[i32], keys_nulls: Option<&NullBuffer>, values_nulls: Option<&NullBuffer>, + last_value_wins: bool, ) -> Result<(ArrayRef, ArrayRef, OffsetBuffer)> { let offsets_len = keys_offsets.len(); let mut new_offsets = Vec::with_capacity(offsets_len); @@ -171,8 +176,14 @@ fn map_deduplicate_keys( let mut new_last_offset = 0; new_offsets.push(new_last_offset); + // Mirror Spark's `ArrayBasedMapBuilder`: the first occurrence of a key + // fixes its position in the output; under LAST_WIN a later duplicate + // overwrites that slot's value. `keys_mask` selects the first-seen keys, + // `value_indices` records the source index in `flat_values` to materialize + // for each output slot (updated in place on overwrite). let mut keys_mask_builder = BooleanBuilder::new(); - let mut values_mask_builder = BooleanBuilder::new(); + let mut value_indices: Vec = Vec::new(); + let mut key_to_output_idx: HashMap = HashMap::new(); for (row_idx, (next_keys_offset, next_values_offset)) in keys_offsets .iter() .zip(values_offsets.iter()) @@ -182,9 +193,6 @@ fn map_deduplicate_keys( let num_keys_entries = *next_keys_offset as usize - cur_keys_offset; let num_values_entries = *next_values_offset as usize - cur_values_offset; - let mut keys_mask_one = vec![false; num_keys_entries]; - let mut values_mask_one = vec![false; num_values_entries]; - let key_is_valid = keys_nulls.is_none_or(|buf| buf.is_valid(row_idx)); let value_is_valid = values_nulls.is_none_or(|buf| buf.is_valid(row_idx)); @@ -193,43 +201,175 @@ fn map_deduplicate_keys( return exec_err!( "map_deduplicate_keys: keys and values lists in the same row must have equal lengths" ); - } else if num_keys_entries != 0 { - let mut seen_keys = HashSet::new(); - - for cur_entry_idx in (0..num_keys_entries).rev() { - let key = ScalarValue::try_from_array( - &flat_keys, - cur_keys_offset + cur_entry_idx, - )? - .compacted(); - if seen_keys.contains(&key) { - // TODO: implement configuration and logic for spark.sql.mapKeyDedupPolicy=EXCEPTION (this is default spark-config) - // exec_err!("invalid argument: duplicate keys in map") - // https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961 - } else { - // This code implements deduplication logic for spark.sql.mapKeyDedupPolicy=LAST_WIN (this is NOT default spark-config) - keys_mask_one[cur_entry_idx] = true; - values_mask_one[cur_entry_idx] = true; - seen_keys.insert(key); - new_last_offset += 1; + } + key_to_output_idx.clear(); + for cur_entry_idx in 0..num_keys_entries { + let key = ScalarValue::try_from_array( + &flat_keys, + cur_keys_offset + cur_entry_idx, + )? + .compacted(); + let abs_value_idx = (cur_values_offset + cur_entry_idx) as i32; + + if let Some(&output_idx) = key_to_output_idx.get(&key) { + if last_value_wins { + value_indices[output_idx] = abs_value_idx; + keys_mask_builder.append_value(false); + continue; } + return exec_err!( + "[DUPLICATED_MAP_KEY] Duplicate map key {key} was found, \ + please check the input data. To allow duplicate keys with \ + last-value-wins semantics, set \ + `datafusion.spark.map_key_dedup_policy` to `LAST_WIN`." + ); } + keys_mask_builder.append_value(true); + key_to_output_idx.insert(key, value_indices.len()); + value_indices.push(abs_value_idx); + new_last_offset += 1; } } else { - // the result entry is NULL - // both current row offsets are skipped - // keys or values in the current row are marked false in the masks + // The result entry is NULL — no keys/values emitted. Still pad the + // mask so it stays aligned with `flat_keys`. + keys_mask_builder.append_n(num_keys_entries, false); } - keys_mask_builder.append_array(&keys_mask_one.into()); - values_mask_builder.append_array(&values_mask_one.into()); new_offsets.push(new_last_offset); cur_keys_offset += num_keys_entries; cur_values_offset += num_values_entries; } let keys_mask = keys_mask_builder.finish(); - let values_mask = values_mask_builder.finish(); let needed_keys = filter(&flat_keys, &keys_mask)?; - let needed_values = filter(&flat_values, &values_mask)?; + let value_indices_array = Int32Array::from(value_indices); + let needed_values = take(&flat_values, &value_indices_array, None)?; let offsets = OffsetBuffer::new(new_offsets.into()); Ok((needed_keys, needed_values, offsets)) } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{Int32Array, StringArray}; + + fn int32_utf8_inputs( + keys: Vec, + values: Vec>, + ) -> (ArrayRef, ArrayRef) { + let keys: ArrayRef = Arc::new(Int32Array::from(keys)); + let values: ArrayRef = Arc::new(StringArray::from(values)); + (keys, values) + } + + #[test] + fn happy_path_two_rows_no_duplicates() { + let (keys, values) = + int32_utf8_inputs(vec![1, 2, 3], vec![Some("a"), Some("b"), Some("c")]); + let offsets = [0i32, 2, 3]; + + let result = map_from_keys_values_offsets_nulls( + &keys, &values, &offsets, &offsets, None, None, false, + ) + .unwrap(); + + let map = result.as_map(); + assert_eq!(map.len(), 2); + assert_eq!(map.value_offsets(), &[0, 2, 3]); + } + + #[test] + fn single_row_duplicate_errors_under_exception() { + let (keys, values) = + int32_utf8_inputs(vec![1, 2, 1], vec![Some("a"), Some("b"), Some("c")]); + let offsets = [0i32, 3]; + + let err = map_from_keys_values_offsets_nulls( + &keys, &values, &offsets, &offsets, None, None, false, + ) + .unwrap_err() + .to_string(); + + assert!(err.contains("[DUPLICATED_MAP_KEY]"), "{err}"); + assert!(err.contains("map_key_dedup_policy"), "{err}"); + } + + #[test] + fn last_win_keeps_final_occurrence() { + let (keys, values) = int32_utf8_inputs( + vec![1, 2, 1, 3, 2], + vec![Some("a"), Some("b"), Some("c"), Some("d"), Some("e")], + ); + let offsets = [0i32, 5]; + + let result = map_from_keys_values_offsets_nulls( + &keys, &values, &offsets, &offsets, None, None, true, + ) + .unwrap(); + + let map = result.as_map(); + assert_eq!(map.len(), 1); + // 5 entries in, 3 unique keys -> offsets [0, 3] + assert_eq!(map.value_offsets(), &[0, 3]); + } + + #[test] + fn duplicate_in_later_row_still_errors() { + let (keys, values) = int32_utf8_inputs( + vec![1, 2, 1, 1], + vec![Some("a"), Some("b"), Some("x"), Some("y")], + ); + let offsets = [0i32, 2, 4]; + + let err = map_from_keys_values_offsets_nulls( + &keys, &values, &offsets, &offsets, None, None, false, + ) + .unwrap_err() + .to_string(); + + assert!(err.contains("[DUPLICATED_MAP_KEY]"), "{err}"); + } + + #[test] + fn empty_row_does_not_trigger_dedup() { + let (keys, values) = int32_utf8_inputs(vec![], vec![]); + let offsets = [0i32, 0]; + + let result = map_from_keys_values_offsets_nulls( + &keys, &values, &offsets, &offsets, None, None, false, + ) + .unwrap(); + + let map = result.as_map(); + assert_eq!(map.len(), 1); + assert_eq!(map.value_offsets(), &[0, 0]); + } + + #[test] + fn null_row_is_skipped_and_not_checked() { + // Row 0 is NULL (keys null). Its duplicate keys should be ignored; + // row 1 is a clean row. + let (keys, values) = int32_utf8_inputs( + vec![1, 1, 2, 3], + vec![Some("dup-a"), Some("dup-b"), Some("x"), Some("y")], + ); + let offsets = [0i32, 2, 4]; + let keys_nulls = NullBuffer::from(vec![false, true]); + + let result = map_from_keys_values_offsets_nulls( + &keys, + &values, + &offsets, + &offsets, + Some(&keys_nulls), + None, + false, + ) + .unwrap(); + + let map = result.as_map(); + assert_eq!(map.len(), 2); + // First row is NULL (no entries emitted), second row keeps both entries. + assert_eq!(map.value_offsets(), &[0, 0, 2]); + assert!(map.is_null(0)); + assert!(!map.is_null(1)); + } +} diff --git a/datafusion/spark/src/function/math/atan2.rs b/datafusion/spark/src/function/math/atan2.rs new file mode 100644 index 0000000000000..70cc1ffeb25a1 --- /dev/null +++ b/datafusion/spark/src/function/math/atan2.rs @@ -0,0 +1,84 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, AsArray, Float64Array}; +use arrow::compute::kernels::arity::binary; +use arrow::datatypes::{DataType, Float64Type}; +use datafusion_common::Result; +use datafusion_common::utils::take_function_args; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use datafusion_functions::utils::make_scalar_function; + +/// Spark-compatible `atan2` function. +/// +/// +/// +/// `atan2(exprY, exprX)` returns the angle in radians between the positive +/// x-axis and the point given by the coordinates (exprX, exprY). +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkAtan2 { + signature: Signature, +} + +impl Default for SparkAtan2 { + fn default() -> Self { + Self::new() + } +} + +impl SparkAtan2 { + pub fn new() -> Self { + Self { + // Spark only defines atan2 over doubles + signature: Signature::exact( + vec![DataType::Float64, DataType::Float64], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for SparkAtan2 { + fn name(&self) -> &str { + "atan2" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(spark_atan2, vec![])(&args.args) + } +} + +fn spark_atan2(args: &[ArrayRef]) -> Result { + // Spark arg order is atan2(exprY, exprX); Rust computes y.atan2(x). + let [y, x] = take_function_args("atan2", args)?; + let y = y.as_primitive::(); + let x = x.as_primitive::(); + let result: Float64Array = binary(y, x, |y, x| y.atan2(x))?; + Ok(Arc::new(result)) +} diff --git a/datafusion/spark/src/function/math/bin.rs b/datafusion/spark/src/function/math/bin.rs index 82afd48e8dc9f..e6a0e1a7359ef 100644 --- a/datafusion/spark/src/function/math/bin.rs +++ b/datafusion/spark/src/function/math/bin.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::array::{ArrayRef, AsArray, StringArray}; +use arrow::array::{Array, ArrayRef, AsArray, StringBuilder}; use arrow::datatypes::{DataType, Field, FieldRef, Int64Type}; use datafusion_common::types::{NativeType, logical_int64}; use datafusion_common::utils::take_function_args; @@ -88,12 +88,20 @@ fn spark_bin_inner(arg: &[ArrayRef]) -> Result { let [array] = take_function_args("bin", arg)?; match &array.data_type() { DataType::Int64 => { - let result: StringArray = array - .as_primitive::() - .iter() - .map(|opt| opt.map(spark_bin)) - .collect(); - Ok(Arc::new(result)) + let array = array.as_primitive::(); + let len = array.len(); + // Most values are small, so 8 digits per row is a reasonable estimate; + // the buffer grows on its own for wider ones. + let mut builder = StringBuilder::with_capacity(len, len * 8); + // Digits are rendered into this stack buffer, so no row allocates. + let mut digits = [0u8; MAX_BIN_DIGITS]; + for value in array.iter() { + match value { + Some(value) => builder.append_value(spark_bin(value, &mut digits)), + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish())) } data_type => { internal_err!("bin does not support: {data_type}") @@ -101,6 +109,24 @@ fn spark_bin_inner(arg: &[ArrayRef]) -> Result { } } -fn spark_bin(value: i64) -> String { - format!("{value:b}") +/// An `i64` renders as at most 64 binary digits. +const MAX_BIN_DIGITS: usize = 64; + +/// Renders `value` as binary, right-aligned in `digits`, and returns the digits written. +/// +/// Negative values render as their two's-complement bit pattern, matching `{:b}`. +fn spark_bin(value: i64, digits: &mut [u8; MAX_BIN_DIGITS]) -> &str { + let mut pos = MAX_BIN_DIGITS; + let mut remaining = value as u64; + // `while` alone would produce an empty string for zero. + loop { + pos -= 1; + digits[pos] = b'0' + (remaining & 1) as u8; + remaining >>= 1; + if remaining == 0 { + break; + } + } + // SAFETY: every byte written above is an ASCII '0' or '1'. + unsafe { std::str::from_utf8_unchecked(&digits[pos..]) } } diff --git a/datafusion/spark/src/function/math/hex.rs b/datafusion/spark/src/function/math/hex.rs index 22e0b5b0786ea..aa32100dd42de 100644 --- a/datafusion/spark/src/function/math/hex.rs +++ b/datafusion/spark/src/function/math/hex.rs @@ -18,7 +18,8 @@ use std::str::from_utf8_unchecked; use std::sync::Arc; -use arrow::array::{Array, ArrayRef, StringBuilder}; +use arrow::array::{Array, ArrayAccessor, ArrayRef, StringArray, StringBuilder}; +use arrow::buffer::{Buffer, OffsetBuffer}; use arrow::datatypes::DataType; use arrow::{ array::{as_dictionary_array, as_largestring_array, as_string_array}, @@ -27,15 +28,16 @@ use arrow::{ use datafusion_common::cast::as_large_binary_array; use datafusion_common::cast::as_string_view_array; use datafusion_common::types::{NativeType, logical_int64, logical_string}; +use datafusion_common::utils::hex::{HexCase, ToHex, encode_bytes_into}; use datafusion_common::utils::take_function_args; use datafusion_common::{ DataFusionError, cast::{as_binary_array, as_fixed_size_binary_array, as_int64_array}, - exec_err, + exec_datafusion_err, exec_err, }; use datafusion_expr::{ - Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignature, - TypeSignatureClass, Volatility, + Coercion, ColumnarValue, EncodingPreservation, ScalarFunctionArgs, ScalarUDFImpl, + Signature, TypeSignature, TypeSignatureClass, Volatility, }; /// #[derive(Debug, PartialEq, Eq, Hash)] @@ -60,7 +62,8 @@ impl SparkHex { let string = Coercion::new_exact(TypeSignatureClass::Native(logical_string())); - let binary = Coercion::new_exact(TypeSignatureClass::Binary); + let binary = Coercion::new_exact(TypeSignatureClass::Binary) + .with_encoding_preservation(EncodingPreservation::dictionary()); let variants = vec![ // accepts numeric types @@ -108,90 +111,81 @@ impl ScalarUDFImpl for SparkHex { } } -/// Hex encoding lookup tables for fast byte-to-hex conversion. -/// -/// Each entry maps a full byte to its two-character hex encoding so the -/// hot loop becomes one load + one two-byte extend per input byte instead -/// of two nibble lookups and two pushes. -const HEX_CHARS_UPPER_NIBBLES: &[u8; 16] = b"0123456789ABCDEF"; -const HEX_CHARS_LOWER_NIBBLES: &[u8; 16] = b"0123456789abcdef"; - -const HEX_LOOKUP_UPPER: [[u8; 2]; 256] = build_hex_lookup(HEX_CHARS_UPPER_NIBBLES); -const HEX_LOOKUP_LOWER: [[u8; 2]; 256] = build_hex_lookup(HEX_CHARS_LOWER_NIBBLES); - -const fn build_hex_lookup(nibbles: &[u8; 16]) -> [[u8; 2]; 256] { - let mut table = [[0u8; 2]; 256]; - let mut i = 0; - while i < 256 { - table[i][0] = nibbles[(i >> 4) & 0xF]; - table[i][1] = nibbles[i & 0xF]; - i += 1; - } - table -} - #[inline] -fn hex_int64(num: i64, buffer: &mut [u8; 16]) -> &[u8] { - if num == 0 { - return b"0"; - } - - // Walk the value two nibbles (one full byte) at a time. The buffer is - // filled from the right so the high-order nibbles end up first; the - // returned slice trims leading zeros automatically. - let mut n = num as u64; - let mut i = 16; - while n >= 0x10 { - i -= 2; - let pair = HEX_LOOKUP_UPPER[(n & 0xFF) as usize]; - buffer[i] = pair[0]; - buffer[i + 1] = pair[1]; - n >>= 8; - } - if n > 0 { - // Single remaining high nibble (value 0x1..=0xF). - i -= 1; - buffer[i] = HEX_CHARS_UPPER_NIBBLES[n as usize]; - } - &buffer[i..] +fn append_hex_bytes( + values: &mut Vec, + bytes: &[u8], + case: HexCase, +) -> Result { + let additional = bytes + .len() + .checked_mul(2) + .ok_or_else(|| exec_datafusion_err!("hex output size overflow"))?; + values.try_reserve(additional).map_err(|e| { + exec_datafusion_err!("failed to reserve {additional} bytes for hex output: {e}") + })?; + encode_bytes_into(bytes, case, values); + i32::try_from(values.len()) + .map_err(|_| exec_datafusion_err!("hex output exceeds i32 offset range")) } /// Generic hex encoding for byte array types -fn hex_encode_bytes<'a, I, T>( - iter: I, +fn hex_encode_bytes<'a, A, T>( + array: &A, lowercase: bool, - len: usize, ) -> Result where - I: Iterator>, - T: AsRef<[u8]> + 'a, + A: ArrayAccessor, + T: AsRef<[u8]> + ?Sized + 'a, { - let mut builder = StringBuilder::with_capacity(len, len * 64); - let mut buffer = Vec::with_capacity(64); - let lookup = if lowercase { - &HEX_LOOKUP_LOWER + let case = if lowercase { + HexCase::Lower } else { - &HEX_LOOKUP_UPPER + HexCase::Upper }; - - for v in iter { - if let Some(b) = v { - let bytes = b.as_ref(); - buffer.clear(); - buffer.reserve(bytes.len() * 2); - for &byte in bytes { - buffer.extend_from_slice(&lookup[byte as usize]); + let len = array.len(); + let nulls = array.nulls().cloned(); + + // Write hex digits directly into one growing value buffer, tracking offsets + // ourselves. Each input byte becomes exactly two output bytes, so there is + // no per-row `String`/`StringBuilder` copy — the hex digits are written once + // into the final buffer. + let mut values: Vec = Vec::with_capacity(len * 64); + let mut offsets: Vec = Vec::with_capacity(len + 1); + offsets.push(0); + + if let Some(ref nulls) = nulls { + for i in 0..len { + if nulls.is_valid(i) { + // SAFETY: `i` is in bounds and the validity buffer marks it valid. + let bytes = unsafe { array.value_unchecked(i) }.as_ref(); + offsets.push(append_hex_bytes(&mut values, bytes, case)?); + } else { + offsets.push(i32::try_from(values.len()).map_err(|_| { + exec_datafusion_err!("hex output exceeds i32 offset range") + })?); } - // SAFETY: buffer contains only ASCII hex digits, which are valid UTF-8. - unsafe { - builder.append_value(from_utf8_unchecked(&buffer)); - } - } else { - builder.append_null(); + } + } else { + for i in 0..len { + // SAFETY: `i` is in bounds and no null buffer means every value is valid. + let bytes = unsafe { array.value_unchecked(i) }.as_ref(); + offsets.push(append_hex_bytes(&mut values, bytes, case)?); } } - Ok(Arc::new(builder.finish())) + // SAFETY: the value buffer contains only ASCII hex digits (valid UTF-8) and + // the offsets are monotonically increasing and end at `values.len()`, so the + // array invariants hold. This mirrors the previous `from_utf8_unchecked` + // path and avoids a redundant UTF-8 validation pass over the whole buffer. + let array = unsafe { + StringArray::new_unchecked( + OffsetBuffer::new(offsets.into()), + Buffer::from_vec(values), + nulls, + ) + }; + Ok(Arc::new(array)) } /// Generic hex encoding for int64 type @@ -204,7 +198,7 @@ fn hex_encode_int64( for v in iter { if let Some(num) = v { let mut temp = [0u8; 16]; - let slice = hex_int64(num, &mut temp); + let slice = num.write_hex(HexCase::Upper, &mut temp); // SAFETY: slice contains only ASCII hex digests, which are valid UTF-8 unsafe { builder.append_value(from_utf8_unchecked(slice)); @@ -247,51 +241,27 @@ pub fn compute_hex( } DataType::Utf8 => { let array = as_string_array(array); - Ok(ColumnarValue::Array(hex_encode_bytes( - array.iter(), - lowercase, - array.len(), - )?)) + Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?)) } DataType::Utf8View => { let array = as_string_view_array(array)?; - Ok(ColumnarValue::Array(hex_encode_bytes( - array.iter(), - lowercase, - array.len(), - )?)) + Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?)) } DataType::LargeUtf8 => { let array = as_largestring_array(array); - Ok(ColumnarValue::Array(hex_encode_bytes( - array.iter(), - lowercase, - array.len(), - )?)) + Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?)) } DataType::Binary => { let array = as_binary_array(array)?; - Ok(ColumnarValue::Array(hex_encode_bytes( - array.iter(), - lowercase, - array.len(), - )?)) + Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?)) } DataType::LargeBinary => { let array = as_large_binary_array(array)?; - Ok(ColumnarValue::Array(hex_encode_bytes( - array.iter(), - lowercase, - array.len(), - )?)) + Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?)) } DataType::FixedSizeBinary(_) => { let array = as_fixed_size_binary_array(array)?; - Ok(ColumnarValue::Array(hex_encode_bytes( - array.iter(), - lowercase, - array.len(), - )?)) + Ok(ColumnarValue::Array(hex_encode_bytes(&array, lowercase)?)) } DataType::Dictionary(key_type, _) => { if **key_type != DataType::Int32 { @@ -311,27 +281,27 @@ pub fn compute_hex( } DataType::Utf8 => { let arr = as_string_array(dict_values); - hex_encode_bytes(arr.iter(), lowercase, arr.len())? + hex_encode_bytes(&arr, lowercase)? } DataType::LargeUtf8 => { let arr = as_largestring_array(dict_values); - hex_encode_bytes(arr.iter(), lowercase, arr.len())? + hex_encode_bytes(&arr, lowercase)? } DataType::Utf8View => { let arr = as_string_view_array(dict_values)?; - hex_encode_bytes(arr.iter(), lowercase, arr.len())? + hex_encode_bytes(&arr, lowercase)? } DataType::Binary => { let arr = as_binary_array(dict_values)?; - hex_encode_bytes(arr.iter(), lowercase, arr.len())? + hex_encode_bytes(&arr, lowercase)? } DataType::LargeBinary => { let arr = as_large_binary_array(dict_values)?; - hex_encode_bytes(arr.iter(), lowercase, arr.len())? + hex_encode_bytes(&arr, lowercase)? } DataType::FixedSizeBinary(_) => { let arr = as_fixed_size_binary_array(dict_values)?; - hex_encode_bytes(arr.iter(), lowercase, arr.len())? + hex_encode_bytes(&arr, lowercase)? } _ => { return exec_err!( @@ -352,11 +322,10 @@ pub fn compute_hex( #[cfg(test)] mod test { - use std::str::from_utf8_unchecked; use std::sync::Arc; use arrow::array::{ - BinaryArray, DictionaryArray, Int32Array, Int64Array, StringArray, + Array, BinaryArray, DictionaryArray, Int32Array, Int64Array, StringArray, }; use arrow::{ array::{ @@ -457,7 +426,7 @@ mod test { #[test] fn test_hex_int64() { - let test_cases = vec![ + let cases = vec![ (0_i64, "0"), (1, "1"), (15, "F"), @@ -470,37 +439,29 @@ mod test { (-1, "FFFFFFFFFFFFFFFF"), ]; - for (num, expected) in test_cases { - let mut cache = [0u8; 16]; - let slice = super::hex_int64(num, &mut cache); - - unsafe { - let result = from_utf8_unchecked(slice); - assert_eq!(expected, result, "hex_int64({num}) mismatch"); - } + let arr = + super::hex_encode_int64(cases.iter().map(|(n, _)| Some(*n)), cases.len()) + .unwrap(); + let arr = as_string_array(&arr); + for (i, (num, expected)) in cases.iter().enumerate() { + assert_eq!(*expected, arr.value(i), "hex({num})"); } } #[test] - fn test_hex_lookup_table_covers_all_bytes() { - // Cross-check the precomputed table against an independent encoder - // for every possible byte value and both casings. - for byte in 0u8..=255 { - let upper = format!("{byte:02X}"); - let lower = format!("{byte:02x}"); - let upper_pair = super::HEX_LOOKUP_UPPER[byte as usize]; - let lower_pair = super::HEX_LOOKUP_LOWER[byte as usize]; - assert_eq!( - upper.as_bytes(), - &upper_pair, - "upper encoding mismatch for byte 0x{byte:02X}" - ); - assert_eq!( - lower.as_bytes(), - &lower_pair, - "lower encoding mismatch for byte 0x{byte:02X}" - ); - } + fn test_hex_encode_bytes_lowercase() { + // Every in-repo caller of `hex_encode_bytes` goes through `spark_hex`, + // which always passes `lowercase = false`. The `lowercase = true` path + // is reachable only via `spark_sha2_hex`, which has no in-workspace + // caller, so it otherwise has no coverage. Drive it directly here. + let input = StringArray::from(vec![Some("hi"), Some("bye"), None, Some("rust")]); + let input_ref = &input; + let result = super::hex_encode_bytes(&input_ref, true).unwrap(); + let result = as_string_array(&result); + + let expected = + StringArray::from(vec![Some("6869"), Some("627965"), None, Some("72757374")]); + assert_eq!(result, &expected); } #[test] @@ -525,6 +486,56 @@ mod test { assert_eq!(strings.value(0), expected); } + #[test] + fn test_spark_hex_binary_no_nulls() { + let input = BinaryArray::from(vec![ + b"".as_slice(), + b"\x00\x7f\x80\xff".as_slice(), + b"DataFusion".as_slice(), + ]); + + let result = super::spark_hex(&[ColumnarValue::Array(Arc::new(input))]).unwrap(); + let array = match result { + ColumnarValue::Array(array) => array, + _ => panic!("Expected array"), + }; + let strings = as_string_array(&array); + + assert_eq!(strings.nulls(), None); + assert_eq!( + strings, + &StringArray::from(vec!["", "007F80FF", "44617461467573696F6E"]) + ); + } + + #[test] + fn test_spark_hex_binary_reuses_input_nulls() { + let input = BinaryArray::from(vec![ + Some(b"skip".as_slice()), + None, + Some(b"\x00\xff".as_slice()), + Some(b"hex".as_slice()), + None, + ]) + .slice(1, 4); + let input_nulls = input.nulls().unwrap().clone(); + + let result = super::spark_hex(&[ColumnarValue::Array(Arc::new(input))]).unwrap(); + let array = match result { + ColumnarValue::Array(array) => array, + _ => panic!("Expected array"), + }; + let strings = as_string_array(&array); + let output_nulls = strings.nulls().unwrap(); + + assert_eq!(output_nulls, &input_nulls); + assert!(output_nulls.inner().ptr_eq(input_nulls.inner())); + assert_eq!( + strings, + &StringArray::from(vec![None, Some("00FF"), Some("686578"), None]) + ); + } + #[test] fn test_spark_hex_int64() { let int_array = Int64Array::from(vec![Some(1), Some(2), None, Some(3)]); @@ -570,4 +581,25 @@ mod test { assert_eq!(&expected, result); } + + #[test] + fn test_dict_binary_values_null() { + let keys = Int32Array::from(vec![Some(0), None, Some(1)]); + let vals = BinaryArray::from(vec![Some(b"hi".as_slice()), None]); + // [b"hi", null, null] + let dict = DictionaryArray::new(keys, Arc::new(vals)); + + let result = super::spark_hex(&[ColumnarValue::Array(Arc::new(dict))]).unwrap(); + let result = match result { + ColumnarValue::Array(array) => array, + _ => panic!("Expected array"), + }; + let result = as_dictionary_array(&result).unwrap(); + + let keys = Int32Array::from(vec![Some(0), None, Some(1)]); + let vals = StringArray::from(vec![Some("6869"), None]); + let expected = DictionaryArray::new(keys, Arc::new(vals)); + + assert_eq!(&expected, result); + } } diff --git a/datafusion/spark/src/function/math/hypot.rs b/datafusion/spark/src/function/math/hypot.rs new file mode 100644 index 0000000000000..a1e30a7e4abe2 --- /dev/null +++ b/datafusion/spark/src/function/math/hypot.rs @@ -0,0 +1,84 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use arrow::array::{ArrayRef, AsArray, Float64Array}; +use arrow::compute::kernels::arity::binary; +use arrow::datatypes::{DataType, Float64Type}; +use datafusion_common::Result; +use datafusion_common::utils::take_function_args; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; +use datafusion_functions::utils::make_scalar_function; + +/// Spark-compatible `hypot` function. +/// +/// +/// +/// Returns `sqrt(expr1^2 + expr2^2)` computed without intermediate overflow or +/// underflow, matching Spark's use of `java.lang.Math.hypot`. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkHypot { + signature: Signature, +} + +impl Default for SparkHypot { + fn default() -> Self { + Self::new() + } +} + +impl SparkHypot { + pub fn new() -> Self { + Self { + // Spark only defines hypot over doubles + signature: Signature::exact( + vec![DataType::Float64, DataType::Float64], + Volatility::Immutable, + ), + } + } +} + +impl ScalarUDFImpl for SparkHypot { + fn name(&self) -> &str { + "hypot" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Float64) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(spark_hypot, vec![])(&args.args) + } +} + +fn spark_hypot(args: &[ArrayRef]) -> Result { + let [x, y] = take_function_args("hypot", args)?; + + let x = x.as_primitive::(); + let y = y.as_primitive::(); + let result: Float64Array = binary(x, y, |a, b| a.hypot(b))?; + Ok(Arc::new(result)) +} diff --git a/datafusion/spark/src/function/math/mod.rs b/datafusion/spark/src/function/math/mod.rs index 896eedd03387e..53cedaef9147c 100644 --- a/datafusion/spark/src/function/math/mod.rs +++ b/datafusion/spark/src/function/math/mod.rs @@ -16,14 +16,17 @@ // under the License. pub mod abs; +pub mod atan2; pub mod bin; pub mod ceil; pub mod expm1; pub mod factorial; pub mod floor; pub mod hex; +pub mod hypot; pub mod modulus; pub mod negative; +pub mod pow; pub mod rint; pub mod round; pub mod trigonometry; @@ -35,13 +38,16 @@ use datafusion_functions::make_udf_function; use std::sync::Arc; make_udf_function!(abs::SparkAbs, abs); +make_udf_function!(atan2::SparkAtan2, atan2); make_udf_function!(ceil::SparkCeil, ceil); make_udf_function!(expm1::SparkExpm1, expm1); make_udf_function!(factorial::SparkFactorial, factorial); make_udf_function!(floor::SparkFloor, floor); make_udf_function!(hex::SparkHex, hex); +make_udf_function!(hypot::SparkHypot, hypot); make_udf_function!(modulus::SparkMod, modulus); make_udf_function!(modulus::SparkPmod, pmod); +make_udf_function!(pow::SparkPow, pow); make_udf_function!(rint::SparkRint, rint); make_udf_function!(round::SparkRound, round); make_udf_function!(unhex::SparkUnhex, unhex); @@ -55,6 +61,7 @@ pub mod expr_fn { use datafusion_functions::export_functions; export_functions!((abs, "Returns abs(expr)", arg1)); + export_functions!((atan2, "Returns the angle in radians between the positive x-axis and the point (exprX, exprY).", arg1 arg2)); export_functions!((ceil, "Returns the ceiling of expr.", arg1)); export_functions!((expm1, "Returns exp(expr) - 1 as a Float64.", arg1)); export_functions!(( @@ -64,8 +71,14 @@ pub mod expr_fn { )); export_functions!((floor, "Returns floor of expr.", arg1)); export_functions!((hex, "Computes hex value of the given column.", arg1)); + export_functions!((hypot, "Returns sqrt(a^2 + b^2) without intermediate overflow or underflow.", arg1 arg2)); export_functions!((modulus, "Returns the remainder of division of the first argument by the second argument.", arg1 arg2)); export_functions!((pmod, "Returns the positive remainder of division of the first argument by the second argument.", arg1 arg2)); + export_functions!(( + pow, + "Returns base raised to the power of exponent. Returns Infinity for pow(0, negative).", + arg1 arg2 + )); export_functions!(( rint, "Returns the double value that is closest in value to the argument and is equal to a mathematical integer.", @@ -95,13 +108,16 @@ pub mod expr_fn { pub fn functions() -> Vec> { vec![ abs(), + atan2(), ceil(), expm1(), factorial(), floor(), hex(), + hypot(), modulus(), pmod(), + pow(), rint(), round(), unhex(), diff --git a/datafusion/spark/src/function/math/pow.rs b/datafusion/spark/src/function/math/pow.rs new file mode 100644 index 0000000000000..8655d71e42c9a --- /dev/null +++ b/datafusion/spark/src/function/math/pow.rs @@ -0,0 +1,152 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible `pow` / `power` function. +//! +//! Unlike the default DataFusion (PostgreSQL) implementation, Spark returns +//! `Infinity` for `pow(0, )` rather than raising an error. + +use std::sync::Arc; + +use arrow::array::{Array, ArrayRef, Float64Array}; +use arrow::datatypes::DataType; + +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{ + ColumnarValue, Documentation, ScalarFunctionArgs, ScalarUDFImpl, Signature, +}; +use datafusion_functions::math::power::PowerFunc; + +/// Spark-compatible implementation of `pow` / `power`. +/// +/// Behavioural difference from the DataFusion default: +/// - `pow(0, )` → `Infinity` (IEEE 754 / Spark semantics) +/// The default raises `"zero raised to a negative power is undefined"` to +/// match PostgreSQL. +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkPow { + inner: PowerFunc, + aliases: Vec, +} + +impl Default for SparkPow { + fn default() -> Self { + Self::new() + } +} + +impl SparkPow { + pub fn new() -> Self { + Self { + inner: PowerFunc::new(), + // SparkPow is named "pow"; expose "power" as an alias so that + // both names resolve to Spark semantics when this crate is active. + aliases: vec!["power".to_string()], + } + } +} + +impl ScalarUDFImpl for SparkPow { + fn name(&self) -> &str { + "pow" + } + + fn aliases(&self) -> &[String] { + &self.aliases + } + + fn signature(&self) -> &Signature { + self.inner.signature() + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + self.inner.return_type(arg_types) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + // Only Float64 × Float64 needs the Spark override. + // Decimal / integer / mixed-type paths are delegated to the standard + // PowerFunc which already handles them correctly (decimal can't + // represent Infinity anyway). + match args.args.as_slice() { + [base, exponent] + if matches!(base.data_type(), DataType::Float64) + && matches!(exponent.data_type(), DataType::Float64) => {} + _ => return self.inner.invoke_with_args(args), + } + + let num_rows = args.number_rows; + + // ── Scalar × Scalar fast path ──────────────────────────────────────── + // Pattern-match on the slice to avoid any ownership issues. + if let [ + ColumnarValue::Scalar(ScalarValue::Float64(base)), + ColumnarValue::Scalar(ScalarValue::Float64(exp)), + ] = args.args.as_slice() + { + // base and exp are &Option; Option is Copy. + let result = (*base).zip(*exp).map(|(base, exp)| { + if base == 0.0 && exp < 0.0 { + f64::INFINITY + } else { + base.powf(exp) + } + }); + return Ok(ColumnarValue::Scalar(ScalarValue::Float64(result))); + } + + // ── Array path ─────────────────────────────────────────────────────── + let [base, exponent] = take_function_args(self.name(), &args.args)?; + + let base_arr: ArrayRef = base.to_array(num_rows)?; + let exp_arr: ArrayRef = exponent.to_array(num_rows)?; + + let base_f64 = base_arr + .as_any() + .downcast_ref::() + .expect("base must be Float64Array"); + let exp_f64 = exp_arr + .as_any() + .downcast_ref::() + .expect("exponent must be Float64Array"); + + // Spark: 0^negative = +Infinity (covers both 0.0 and -0.0) + // IEEE 754: 0.0^-1.0 = +Infinity, -0.0^-1.0 = -Infinity + // Thus we need an explicit guard for base == 0.0 to ensure +Infinity. + let result: Float64Array = base_f64 + .iter() + .zip(exp_f64.iter()) + .map(|(base, exp)| match (base, exp) { + (Some(base), Some(exp)) => { + if base == 0.0 && exp < 0.0 { + Some(f64::INFINITY) + } else { + Some(base.powf(exp)) + } + } + _ => None, + }) + .collect(); + + Ok(ColumnarValue::Array(Arc::new(result))) + } + + fn documentation(&self) -> Option<&Documentation> { + self.inner.documentation() + } +} diff --git a/datafusion/spark/src/function/math/round.rs b/datafusion/spark/src/function/math/round.rs index 05745666183d3..471d38d804cac 100644 --- a/datafusion/spark/src/function/math/round.rs +++ b/datafusion/spark/src/function/math/round.rs @@ -462,18 +462,7 @@ fn spark_round(args: &[ColumnarValue], enable_ansi_mode: bool) -> Result { - let array = array.as_primitive::(); - let result: PrimitiveArray = array.try_unary(|x| { - let v_i64 = i64::try_from(x).map_err(|_| { - (exec_err!( - "round: UInt64 value {x} exceeds i64::MAX and cannot be rounded" - ) as Result<(), _>) - .unwrap_err() - })?; - round_integer(v_i64, scale, enable_ansi_mode) - .map(|v| v as u64) - })?; - Ok(ColumnarValue::Array(Arc::new(result))) + impl_integer_array_round!(array, UInt64Type, scale, enable_ansi_mode) } // Float types @@ -588,16 +577,20 @@ fn spark_round(args: &[ColumnarValue], enable_ansi_mode: bool) -> Result { - let v_i64 = i64::try_from(*v).map_err(|_| { - (exec_err!( - "round: UInt64 value {v} exceeds i64::MAX and cannot be rounded" - ) as Result<(), _>) - .unwrap_err() - })?; - let result = round_integer(v_i64, scale, enable_ansi_mode)?; - Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some( - result as u64, - )))) + if scale >= 0 { + Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some(*v)))) + } else { + let v_i64 = i64::try_from(*v).map_err(|_| { + (exec_err!( + "round: UInt64 value {v} exceeds i64::MAX and cannot be rounded" + ) as Result<(), _>) + .unwrap_err() + })?; + let result = round_integer(v_i64, scale, enable_ansi_mode)?; + Ok(ColumnarValue::Scalar(ScalarValue::UInt64(Some( + result as u64, + )))) + } } // Float scalars diff --git a/datafusion/spark/src/function/math/unhex.rs b/datafusion/spark/src/function/math/unhex.rs index f6c9e2fa27a67..6739e6a15c582 100644 --- a/datafusion/spark/src/function/math/unhex.rs +++ b/datafusion/spark/src/function/math/unhex.rs @@ -22,7 +22,9 @@ use datafusion_common::cast::{ }; use datafusion_common::types::logical_string; use datafusion_common::utils::take_function_args; -use datafusion_common::{DataFusionError, Result, ScalarValue, exec_err}; +use datafusion_common::{ + DataFusionError, Result, ScalarValue, exec_datafusion_err, exec_err, +}; use datafusion_expr::{ Coercion, ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, TypeSignatureClass, Volatility, @@ -125,7 +127,12 @@ where for v in iter { if let Some(s) = v { buffer.clear(); - buffer.reserve(s.as_ref().len().div_ceil(2)); + let additional = s.as_ref().len().div_ceil(2); + buffer.try_reserve(additional).map_err(|e| { + exec_datafusion_err!( + "failed to reserve {additional} bytes for unhex output: {e}" + ) + })?; if unhex_common(s.as_ref().as_bytes(), &mut buffer) { builder.append_value(&buffer); } else { diff --git a/datafusion/spark/src/function/math/width_bucket.rs b/datafusion/spark/src/function/math/width_bucket.rs index 79da924116d2e..93be47a4ca719 100644 --- a/datafusion/spark/src/function/math/width_bucket.rs +++ b/datafusion/spark/src/function/math/width_bucket.rs @@ -21,8 +21,7 @@ use arrow::array::{ Array, ArrayRef, DurationMicrosecondArray, Float64Array, IntervalMonthDayNanoArray, IntervalYearMonthArray, }; -use arrow::datatypes::DataType; -use arrow::datatypes::DataType::{Duration, Float64, Int32, Interval}; +use arrow::datatypes::DataType::{self, Duration, Float64, Int64, Interval}; use arrow::datatypes::IntervalUnit::{MonthDayNano, YearMonth}; use datafusion_common::cast::{ as_duration_microsecond_array, as_float64_array, as_int64_array, @@ -40,7 +39,7 @@ use datafusion_expr::{ }; use datafusion_functions::utils::make_scalar_function; -use arrow::array::{Int32Array, Int32Builder, Int64Array}; +use arrow::array::{Int64Array, Int64Builder}; use arrow::datatypes::TimeUnit::Microsecond; use datafusion_expr::Coercion; use datafusion_expr::Volatility::Immutable; @@ -125,7 +124,7 @@ impl ScalarUDFImpl for SparkWidthBucket { } fn return_type(&self, _arg_types: &[DataType]) -> Result { - Ok(Int32) + Ok(Int64) } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { @@ -199,9 +198,9 @@ macro_rules! width_bucket_kernel_impl { min: &$arr_ty, max: &$arr_ty, n_bucket: &Int64Array, - ) -> Int32Array { + ) -> Int64Array { let len = v.len(); - let mut b = Int32Builder::with_capacity(len); + let mut b = Int64Builder::with_capacity(len); for i in 0..len { if v.is_null(i) || min.is_null(i) || max.is_null(i) || n_bucket.is_null(i) @@ -218,7 +217,7 @@ macro_rules! width_bucket_kernel_impl { b.append_null(); continue; } - let next_bucket = (buckets + 1) as i32; + let next_bucket = (buckets + 1) as i64; if $check_nan { if !x.is_finite() || !l.is_finite() || !h.is_finite() { b.append_null(); @@ -264,7 +263,7 @@ macro_rules! width_bucket_kernel_impl { b.append_null(); continue; } - let mut bucket = ((x - l) / width).floor() as i32 + 1; + let mut bucket = ((x - l) / width).floor() as i64 + 1; if bucket < 1 { bucket = 1; } @@ -306,9 +305,9 @@ pub(crate) fn width_bucket_interval_mdn_exact( lo: &IntervalMonthDayNanoArray, hi: &IntervalMonthDayNanoArray, n: &Int64Array, -) -> Int32Array { +) -> Int64Array { let len = v.len(); - let mut b = Int32Builder::with_capacity(len); + let mut b = Int64Builder::with_capacity(len); for i in 0..len { if v.is_null(i) || lo.is_null(i) || hi.is_null(i) || n.is_null(i) { @@ -320,7 +319,7 @@ pub(crate) fn width_bucket_interval_mdn_exact( b.append_null(); continue; } - let next_bucket = (buckets + 1) as i32; + let next_bucket = buckets + 1; let x = v.value(i); let l = lo.value(i); @@ -366,7 +365,7 @@ pub(crate) fn width_bucket_interval_mdn_exact( continue; } - let mut bucket = ((x_m - l_m) / width).floor() as i32 + 1; + let mut bucket = ((x_m - l_m) / width).floor() as i64 + 1; if bucket < 1 { bucket = 1; } @@ -417,7 +416,7 @@ pub(crate) fn width_bucket_interval_mdn_exact( continue; } - let mut bucket = ((x_f - l_f) / width).floor() as i32 + 1; + let mut bucket = ((x_f - l_f) / width).floor() as i64 + 1; if bucket < 1 { bucket = 1; } @@ -437,10 +436,11 @@ pub(crate) fn width_bucket_interval_mdn_exact( #[cfg(test)] mod tests { use super::*; + use arrow::datatypes::Int64Type; use arrow::array::{ - ArrayRef, DurationMicrosecondArray, Float64Array, Int32Array, Int64Array, - IntervalYearMonthArray, + ArrayRef, AsArray, DurationMicrosecondArray, Float64Array, Int32Array, + Int64Array, IntervalYearMonthArray, }; use arrow::datatypes::IntervalMonthDayNano; @@ -466,10 +466,6 @@ mod tests { Arc::new(IntervalYearMonthArray::from(vals.to_vec())) } - fn downcast_i32(arr: &ArrayRef) -> &Int32Array { - arr.as_any().downcast_ref::().unwrap() - } - fn mdn_array(vals: &[(i32, i32, i64)]) -> Arc { let data: Vec = vals .iter() @@ -488,7 +484,7 @@ mod tests { let n = i64_array_all(5, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[1, 2, 10, 0, 11]); } @@ -500,7 +496,7 @@ mod tests { let n = i64_array_all(5, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[1, 1, 11, 11, 0]); } @@ -512,7 +508,7 @@ mod tests { let n = i64_array_all(3, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[1, 10, 11]); } @@ -524,7 +520,7 @@ mod tests { let n = i64_array_all(3, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[1, 11, 11]); } @@ -535,7 +531,7 @@ mod tests { let hi = f64_array(&[10.0, 10.0, 10.0]); let n = Arc::new(Int64Array::from(vec![0, -1, 10])); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert!(out.is_null(0)); assert!(out.is_null(1)); assert_eq!(out.value(2), 10); @@ -545,7 +541,7 @@ mod tests { let hi = f64_array(&[5.0]); let n = i64_array_all(1, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert!(out.is_null(0)); let v = f64_array_opt(&[Some(f64::NAN)]); @@ -553,7 +549,7 @@ mod tests { let hi = f64_array(&[10.0]); let n = i64_array_all(1, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert!(out.is_null(0)); } @@ -565,7 +561,7 @@ mod tests { let n = i64_array_all(4, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert!(out.is_null(0)); assert_eq!(out.value(1), 2); assert_eq!(out.value(2), 3); @@ -576,7 +572,7 @@ mod tests { let hi = f64_array(&[10.0]); let n = i64_array_all(1, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert!(out.is_null(0)); } @@ -590,7 +586,7 @@ mod tests { let n = i64_array_all(3, 2); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[2, 1, 0]); } @@ -601,7 +597,7 @@ mod tests { let hi = dur_us_array(&[1]); let n = i64_array_all(1, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - assert!(downcast_i32(&out).is_null(0)); + assert!(out.as_primitive::().is_null(0)); } // --- Interval(YearMonth) ------------------------------------------------ @@ -614,7 +610,7 @@ mod tests { let n = i64_array_all(5, 12); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[1, 6, 12, 13, 13]); } @@ -626,7 +622,7 @@ mod tests { let n = i64_array_all(5, 12); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[2, 1, 13, 13, 0]); } @@ -640,7 +636,7 @@ mod tests { let n = i64_array_all(5, 12); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[1, 6, 12, 13, 13]); } @@ -652,7 +648,7 @@ mod tests { let n = i64_array_all(5, 12); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); // Mismo patrón que YM descendente assert_eq!(out.values(), &[2, 1, 13, 13, 0]); } @@ -672,7 +668,7 @@ mod tests { let n = i64_array_all(6, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); // x==hi -> n+1, x 0, x>hi -> n+1 assert_eq!(out.values(), &[1, 6, 10, 11, 0, 11]); } @@ -685,7 +681,7 @@ mod tests { let n = i64_array_all(5, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[2, 1, 11, 11, 0]); } @@ -697,7 +693,7 @@ mod tests { let n = i64_array_all(5, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert_eq!(out.values(), &[1, 1, 11, 11, 0]); } @@ -710,7 +706,7 @@ mod tests { let n = i64_array_all(1, 4); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert!(out.is_null(0)); } @@ -722,7 +718,7 @@ mod tests { let n = i64_array_all(1, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - assert!(downcast_i32(&out).is_null(0)); + assert!(out.as_primitive::().is_null(0)); } #[test] @@ -733,7 +729,7 @@ mod tests { let n = Arc::new(Int64Array::from(vec![0])); // n <= 0 let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - assert!(downcast_i32(&out).is_null(0)); + assert!(out.as_primitive::().is_null(0)); } #[test] @@ -747,7 +743,7 @@ mod tests { let n = i64_array_all(2, 10); let out = width_bucket_kern(&[v, lo, hi, n]).unwrap(); - let out = downcast_i32(&out); + let out = out.as_primitive::(); assert!(out.is_null(0)); assert_eq!(out.value(1), 6); } diff --git a/datafusion/spark/src/function/string/char.rs b/datafusion/spark/src/function/string/char.rs index 15b00ee98f5c7..5d6de3ae368e3 100644 --- a/datafusion/spark/src/function/string/char.rs +++ b/datafusion/spark/src/function/string/char.rs @@ -112,6 +112,8 @@ fn chr(args: &[ArrayRef]) -> Result { integer_array.len(), ); + // Each character encodes into this stack buffer, so no row allocates a `String`. + let mut encoded = [0u8; 4]; for integer_opt in integer_array { match integer_opt { Some(integer) => { @@ -119,7 +121,7 @@ fn chr(args: &[ArrayRef]) -> Result { builder.append_value(""); // empty string for negative numbers. } else { match core::char::from_u32((integer % 256) as u32) { - Some(ch) => builder.append_value(ch.to_string()), + Some(ch) => builder.append_value(ch.encode_utf8(&mut encoded)), None => { return exec_err!( "requested character not compatible for encoding." diff --git a/datafusion/spark/src/function/string/concat.rs b/datafusion/spark/src/function/string/concat.rs index 57fd6cadd9dde..be5ced2edfbf0 100644 --- a/datafusion/spark/src/function/string/concat.rs +++ b/datafusion/spark/src/function/string/concat.rs @@ -71,8 +71,13 @@ impl ScalarUDFImpl for SparkConcat { } fn coerce_types(&self, arg_types: &[DataType]) -> Result> { - // Accept any string types, including zero arguments - Ok(arg_types.to_vec()) + if arg_types.is_empty() { + // Spark semantics: allow concat with zero arguments + Ok(vec![]) + } else { + // Use concat coercion rules + ConcatFunc::new().coerce_types(arg_types) + } } fn return_type(&self, _arg_types: &[DataType]) -> Result { datafusion_common::internal_err!( @@ -80,19 +85,15 @@ impl ScalarUDFImpl for SparkConcat { ) } fn return_field_from_args(&self, args: ReturnFieldArgs<'_>) -> Result { - use DataType::*; - // Spark semantics: concat returns NULL if ANY input is NULL let nullable = args.arg_fields.iter().any(|f| f.is_nullable()); - // Determine return type: Utf8View > LargeUtf8 > Utf8 - let mut dt = &Utf8; - for field in args.arg_fields { - let data_type = field.data_type(); - if data_type == &Utf8View || (data_type == &LargeUtf8 && dt != &Utf8View) { - dt = data_type; - } - } + let arg_types: Vec = args + .arg_fields + .iter() + .map(|f| f.data_type().clone()) + .collect(); + let dt = ConcatFunc::new().return_type(&arg_types)?; Ok(Arc::new(Field::new("concat", dt.clone(), nullable))) } @@ -113,17 +114,9 @@ fn spark_concat(args: ScalarFunctionArgs) -> Result { // Handle zero-argument case: return empty string if arg_values.is_empty() { let return_type = return_field.data_type(); - return match return_type { - DataType::Utf8View => Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(Some( - String::new(), - )))), - DataType::LargeUtf8 => Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8( - Some(String::new()), - ))), - _ => Ok(ColumnarValue::Scalar(ScalarValue::Utf8( - Some(String::new()), - ))), - }; + return Ok(ColumnarValue::Scalar(ScalarValue::new_default( + return_type, + )?)); } // Step 1: Check for NULL mask in incoming args @@ -132,13 +125,9 @@ fn spark_concat(args: ScalarFunctionArgs) -> Result { // If all scalars and any is NULL, return NULL immediately if matches!(null_mask, NullMaskResolution::ReturnNull) { let return_type = return_field.data_type(); - return match return_type { - DataType::Utf8View => Ok(ColumnarValue::Scalar(ScalarValue::Utf8View(None))), - DataType::LargeUtf8 => { - Ok(ColumnarValue::Scalar(ScalarValue::LargeUtf8(None))) - } - _ => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))), - }; + return Ok(ColumnarValue::Scalar(ScalarValue::try_new_null( + return_type, + )?)); } // Step 2: Delegate to DataFusion's concat diff --git a/datafusion/spark/src/function/string/concat_ws.rs b/datafusion/spark/src/function/string/concat_ws.rs new file mode 100644 index 0000000000000..c9ed1369a51a7 --- /dev/null +++ b/datafusion/spark/src/function/string/concat_ws.rs @@ -0,0 +1,297 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Spark-compatible `concat_ws`: joins strings (and array elements) with a separator. +//! +//! Null scalar args and null array elements are skipped; a null separator yields a +//! null row. Non-string args are coerced to STRING; list args (`List`, `LargeList`, +//! `ListView`, `LargeListView`, `FixedSizeList`) expand their elements. +//! +//! Differences with DataFusion core `concat_ws`: +//! - Accepts list arguments and expands their elements +//! - Always returns Utf8 (Spark's `STRING` type) +//! - Coerces non-string scalars (numbers, booleans, dates, ...) to Utf8 + +use std::fmt::Write as _; +use std::sync::Arc; + +use arrow::array::{ + Array, ArrayRef, AsArray, GenericListArray, LargeStringArray, OffsetSizeTrait, + StringArray, StringBuilder, StringViewArray, +}; +use arrow::datatypes::{DataType, Field}; +use datafusion_common::{Result, ScalarValue}; +use datafusion_expr::{ + ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, +}; + +use crate::function::error_utils::{ + invalid_arg_count_exec_err, unsupported_data_type_exec_err, +}; + +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkConcatWs { + signature: Signature, +} + +impl Default for SparkConcatWs { + fn default() -> Self { + Self::new() + } +} + +impl SparkConcatWs { + pub fn new() -> Self { + Self { + signature: Signature::user_defined(Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkConcatWs { + fn name(&self) -> &str { + "concat_ws" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Ok(DataType::Utf8) + } + + fn coerce_types(&self, arg_types: &[DataType]) -> Result> { + if arg_types.is_empty() { + return Err(invalid_arg_count_exec_err("concat_ws", (1, i32::MAX), 0)); + } + Ok(arg_types + .iter() + .enumerate() + .map(|(i, dt)| match dt { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => dt.clone(), + // Non-separator list args expand their elements at runtime. + // Normalize the list variant so the kernel only sees + // List/LargeList, AND force the element type to Utf8 so the + // planner inserts a cast for non-string children (Spark + // coerces them to STRING the same way it does for scalars). + DataType::List(f) + | DataType::ListView(f) + | DataType::FixedSizeList(f, _) + if i > 0 => + { + DataType::List(Arc::new(Field::new( + f.name(), + DataType::Utf8, + f.is_nullable(), + ))) + } + DataType::LargeList(f) | DataType::LargeListView(f) if i > 0 => { + DataType::LargeList(Arc::new(Field::new( + f.name(), + DataType::Utf8, + f.is_nullable(), + ))) + } + // Spark casts everything else (numbers, booleans, dates, + // binary, null...) to STRING. + _ => DataType::Utf8, + }) + .collect()) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + // Only separator provided → empty string (or NULL if separator is null). + // Arg-count validation happens in coerce_types at planning time. + if args.args.len() == 1 { + return only_separator(&args.args[0]); + } + + spark_concat_ws(&args.args, args.number_rows) + } +} + +fn only_separator(sep: &ColumnarValue) -> Result { + match sep { + ColumnarValue::Scalar(s) if s.is_null() => { + Ok(ColumnarValue::Scalar(ScalarValue::Utf8(None))) + } + ColumnarValue::Scalar(_) => Ok(ColumnarValue::Scalar(ScalarValue::Utf8(Some( + String::new(), + )))), + ColumnarValue::Array(arr) => { + let mut builder = StringBuilder::with_capacity(arr.len(), 0); + for row_idx in 0..arr.len() { + if arr.is_null(row_idx) { + builder.append_null(); + } else { + builder.append_value(""); + } + } + Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef)) + } + } +} + +fn spark_concat_ws(args: &[ColumnarValue], num_rows: usize) -> Result { + let arrays = ColumnarValue::values_to_arrays(args)?; + let sep_view = StringView::try_new(&arrays[0])?; + let arg_views: Vec = arrays[1..] + .iter() + .map(ArgView::try_new) + .collect::>()?; + + let mut builder = StringBuilder::with_capacity(num_rows, num_rows * 16); + + for row_idx in 0..num_rows { + if sep_view.is_null(row_idx) { + builder.append_null(); + continue; + } + + // Write parts directly into the builder via its `fmt::Write` impl; + // `append_value("")` then finalises the row (offset + validity) with + // no extra copy from an intermediate `String`. + let separator = sep_view.value(row_idx); + let mut first = true; + for view in &arg_views { + view.write_row(row_idx, separator, &mut builder, &mut first)?; + } + builder.append_value(""); + } + + Ok(ColumnarValue::Array(Arc::new(builder.finish()) as ArrayRef)) +} + +/// Typed view over a string array that downcasts once and exposes +/// per-row access without further dispatch. +enum StringView<'a> { + Utf8(&'a StringArray), + LargeUtf8(&'a LargeStringArray), + Utf8View(&'a StringViewArray), +} + +impl<'a> StringView<'a> { + fn try_new(arr: &'a ArrayRef) -> Result { + match arr.data_type() { + DataType::Utf8 => Ok(Self::Utf8(arr.as_string::())), + DataType::LargeUtf8 => Ok(Self::LargeUtf8(arr.as_string::())), + DataType::Utf8View => Ok(Self::Utf8View(arr.as_string_view())), + other => Err(unsupported_data_type_exec_err("concat_ws", "STRING", other)), + } + } + + fn value(&self, idx: usize) -> &str { + match self { + Self::Utf8(a) => a.value(idx), + Self::LargeUtf8(a) => a.value(idx), + Self::Utf8View(a) => a.value(idx), + } + } + + fn is_null(&self, idx: usize) -> bool { + match self { + Self::Utf8(a) => a.is_null(idx), + Self::LargeUtf8(a) => a.is_null(idx), + Self::Utf8View(a) => a.is_null(idx), + } + } +} + +/// Per-argument view: a string array or a list of strings. The downcast +/// happens once at construction time. `DataType::Null` cannot appear here — +/// `coerce_types` rewrites it to `Utf8` before invocation. +enum ArgView<'a> { + Str(StringView<'a>), + List(&'a GenericListArray), + LargeList(&'a GenericListArray), +} + +impl<'a> ArgView<'a> { + fn try_new(arr: &'a ArrayRef) -> Result { + match arr.data_type() { + DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { + Ok(Self::Str(StringView::try_new(arr)?)) + } + DataType::List(_) => Ok(Self::List(arr.as_list::())), + DataType::LargeList(_) => Ok(Self::LargeList(arr.as_list::())), + other => Err(unsupported_data_type_exec_err( + "concat_ws", + "STRING or ARRAY", + other, + )), + } + } + + fn write_row( + &self, + row_idx: usize, + sep: &str, + builder: &mut StringBuilder, + first: &mut bool, + ) -> Result<()> { + match self { + Self::Str(view) => { + if !view.is_null(row_idx) { + push_part(builder, view.value(row_idx), sep, first); + } + } + Self::List(list) => write_list_row(*list, row_idx, sep, builder, first)?, + Self::LargeList(list) => write_list_row(*list, row_idx, sep, builder, first)?, + } + Ok(()) + } +} + +fn write_list_row( + list: &GenericListArray, + row_idx: usize, + sep: &str, + builder: &mut StringBuilder, + first: &mut bool, +) -> Result<()> { + if list.is_null(row_idx) { + return Ok(()); + } + let values = list.value(row_idx); + // An empty array (e.g. `array()`) contributes nothing — Spark renders it + // as the empty string, not an error. + if values.is_empty() { + return Ok(()); + } + let view = StringView::try_new(&values)?; + for i in 0..values.len() { + if !view.is_null(i) { + push_part(builder, view.value(i), sep, first); + } + } + Ok(()) +} + +// `StringBuilder::write_str` only does `extend_from_slice` and never errors; +// the `.expect(..)` is a documentation hint, not a real failure path. +fn push_part(builder: &mut StringBuilder, part: &str, sep: &str, first: &mut bool) { + if !*first { + builder + .write_str(sep) + .expect("StringBuilder::write_str is infallible"); + } + *first = false; + builder + .write_str(part) + .expect("StringBuilder::write_str is infallible"); +} diff --git a/datafusion/spark/src/function/string/elt.rs b/datafusion/spark/src/function/string/elt.rs index c37ecd1d3fc39..b88477a7720f3 100644 --- a/datafusion/spark/src/function/string/elt.rs +++ b/datafusion/spark/src/function/string/elt.rs @@ -24,7 +24,7 @@ use arrow::compute::{can_cast_types, cast}; use arrow::datatypes::DataType::{Int64, Utf8}; use arrow::datatypes::{DataType, Int64Type}; use datafusion_common::cast::as_string_array; -use datafusion_common::{DataFusionError, Result, plan_datafusion_err}; +use datafusion_common::{DataFusionError, Result, exec_err, plan_datafusion_err}; use datafusion_expr::{ ColumnarValue, ScalarFunctionArgs, ScalarUDFImpl, Signature, Volatility, }; @@ -63,15 +63,19 @@ impl ScalarUDFImpl for SparkElt { } fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { - make_scalar_function(elt, vec![])(&args.args) + let enable_ansi_mode = args.config_options.execution.enable_ansi_mode; + make_scalar_function( + move |arrays: &[ArrayRef]| elt(arrays, enable_ansi_mode), + vec![], + )(&args.args) } fn coerce_types(&self, arg_types: &[DataType]) -> Result> { let length = arg_types.len(); if length < 2 { - plan_datafusion_err!( + return Err(plan_datafusion_err!( "ELT function expects at least 2 arguments: index, value1" - ); + )); } let idx_dt: &DataType = &arg_types[0]; @@ -80,18 +84,13 @@ impl ScalarUDFImpl for SparkElt { "ELT index must be Int64 (or castable to Int64), got {idx_dt:?}" ))); } - let mut coerced = Vec::with_capacity(arg_types.len()); - coerced.push(Int64); - - for _ in 1..length { - coerced.push(Utf8); - } - + let mut coerced = vec![Utf8; length]; + coerced[0] = Int64; Ok(coerced) } } -fn elt(args: &[ArrayRef]) -> Result { +fn elt(args: &[ArrayRef], enable_ansi_mode: bool) -> Result { let n_rows = args[0].len(); let idx: &PrimitiveArray = @@ -103,11 +102,10 @@ fn elt(args: &[ArrayRef]) -> Result { })?; let num_values = args.len() - 1; - let mut cols: Vec> = Vec::with_capacity(num_values); + let mut cols: Vec = Vec::with_capacity(num_values); for a in args.iter().skip(1) { let casted = cast(a, &Utf8)?; - let sa = as_string_array(&casted)?; - cols.push(Arc::new(sa.clone())); + cols.push(as_string_array(&casted)?.clone()); } let mut builder = StringBuilder::new(); @@ -120,10 +118,12 @@ fn elt(args: &[ArrayRef]) -> Result { let index = idx.value(i); - // TODO: if spark.sql.ansi.enabled is true, - // throw ArrayIndexOutOfBoundsException for invalid indices; - // if false, return NULL instead (current behavior). if index < 1 || (index as usize) > num_values { + if enable_ansi_mode { + return exec_err!( + "The index {index} is out of bounds. The array has {num_values} elements." + ); + } builder.append_null(); continue; } @@ -146,13 +146,13 @@ mod tests { use super::*; use arrow::array::Int64Array; - fn run_elt_arrays(arrs: Vec) -> Result> { - let arr = elt(&arrs)?; - let string_array = arr - .as_any() - .downcast_ref::() - .ok_or_else(|| DataFusionError::Internal("expected Utf8".into()))?; - Ok(Arc::new(string_array.clone())) + fn run_elt_arrays(arrs: Vec) -> Result { + run_elt_arrays_with(arrs, false) + } + + fn run_elt_arrays_with(arrs: Vec, ansi: bool) -> Result { + let arr = elt(&arrs, ansi)?; + Ok(as_string_array(&arr)?.clone()) } #[test] diff --git a/datafusion/spark/src/function/string/format_string.rs b/datafusion/spark/src/function/string/format_string.rs index 51e4ebfa7b465..6a65164a18318 100644 --- a/datafusion/spark/src/function/string/format_string.rs +++ b/datafusion/spark/src/function/string/format_string.rs @@ -891,23 +891,85 @@ fn unsigned_to_char(value: u64) -> Result { codepoint_to_char(codepoint) } -/// Convert a non-null integer scalar to a [`char`] for the `%c` conversion. -fn integer_scalar_to_char(scalar: &ScalarValue) -> Result { - match scalar { - ScalarValue::Int8(Some(value)) => signed_to_char(*value as i64), - ScalarValue::Int16(Some(value)) => signed_to_char(*value as i64), - ScalarValue::Int32(Some(value)) => signed_to_char(*value as i64), - ScalarValue::Int64(Some(value)) => signed_to_char(*value), - ScalarValue::UInt8(Some(value)) => unsigned_to_char(*value as u64), - ScalarValue::UInt16(Some(value)) => unsigned_to_char(*value as u64), - ScalarValue::UInt32(Some(value)) => unsigned_to_char(*value as u64), - ScalarValue::UInt64(Some(value)) => unsigned_to_char(*value), - _ => datafusion_common::internal_err!( - "integer_scalar_to_char expects a non-null integer scalar, got {scalar:?}" - ), - } +/// Formatting operations that differ between signed and unsigned integer +/// primitives. Signed values format as decimal for `%d` / `%s` / `%c`, but use +/// their original bit width for `%x` / `%o` via `unsigned_bits`. +trait IntegerFormatValue { + fn unsigned_bits(self) -> u64; + + fn to_char(self) -> Result; + + fn format_decimal( + self, + spec: &ConversionSpecifier, + writer: &mut String, + ) -> Result<()>; + + fn decimal_string(self) -> String; +} + +macro_rules! signed_integer_value { + ($source:ty, $unsigned:ty) => { + impl IntegerFormatValue for $source { + fn unsigned_bits(self) -> u64 { + (self as $unsigned) as u64 + } + + fn to_char(self) -> Result { + signed_to_char(self as i64) + } + + fn format_decimal( + self, + spec: &ConversionSpecifier, + writer: &mut String, + ) -> Result<()> { + spec.format_signed(writer, self as i64) + } + + fn decimal_string(self) -> String { + self.to_string() + } + } + }; } +signed_integer_value!(i8, u8); +signed_integer_value!(i16, u16); +signed_integer_value!(i32, u32); +signed_integer_value!(i64, u64); + +macro_rules! unsigned_integer_value { + ($source:ty) => { + impl IntegerFormatValue for $source { + fn unsigned_bits(self) -> u64 { + self as u64 + } + + fn to_char(self) -> Result { + unsigned_to_char(self as u64) + } + + fn format_decimal( + self, + spec: &ConversionSpecifier, + writer: &mut String, + ) -> Result<()> { + spec.format_unsigned(writer, self as u64) + } + + fn decimal_string(self) -> String { + self.to_string() + } + } + }; +} + +unsigned_integer_value!(u8); +unsigned_integer_value!(u16); +unsigned_integer_value!(u32); +unsigned_integer_value!(u64); + impl ConversionSpecifier { /// Validates that the grouping separator flag is not used with scientific /// notation conversions, matching Java/Spark behavior which throws @@ -940,189 +1002,14 @@ impl ConversionSpecifier { _ => self.format_boolean(string, value), }, - ScalarValue::Int8(Some(_)) - | ScalarValue::Int16(Some(_)) - | ScalarValue::Int32(Some(_)) - | ScalarValue::Int64(Some(_)) - | ScalarValue::UInt8(Some(_)) - | ScalarValue::UInt16(Some(_)) - | ScalarValue::UInt32(Some(_)) - | ScalarValue::UInt64(Some(_)) - if matches!( - self.conversion_type, - ConversionType::CharLower | ConversionType::CharUpper - ) => - { - self.format_char(string, integer_scalar_to_char(value)?) - } - ScalarValue::Int8(value) => match (self.conversion_type, value) { - (ConversionType::DecInt, Some(value)) => { - self.format_signed(string, *value as i64) - } - ( - ConversionType::HexIntLower - | ConversionType::HexIntUpper - | ConversionType::OctInt, - Some(value), - ) => self.format_unsigned(string, (*value as u8) as u64), - ( - ConversionType::StringLower | ConversionType::StringUpper, - Some(value), - ) => self.format_string(string, &value.to_string()), - (t, None) if t.supports_integer() => self.format_string(string, "null"), - _ => { - exec_err!( - "Invalid conversion type: {:?} for Int8", - self.conversion_type - ) - } - }, - ScalarValue::Int16(value) => match (self.conversion_type, value) { - (ConversionType::DecInt, Some(value)) => { - self.format_signed(string, *value as i64) - } - ( - ConversionType::HexIntLower - | ConversionType::HexIntUpper - | ConversionType::OctInt, - Some(value), - ) => self.format_unsigned(string, (*value as u16) as u64), - ( - ConversionType::StringLower | ConversionType::StringUpper, - Some(value), - ) => self.format_string(string, &value.to_string()), - (t, None) if t.supports_integer() => self.format_string(string, "null"), - _ => { - exec_err!( - "Invalid conversion type: {:?} for Int16", - self.conversion_type - ) - } - }, - ScalarValue::Int32(value) => match (self.conversion_type, value) { - (ConversionType::DecInt, Some(value)) => { - self.format_signed(string, *value as i64) - } - ( - ConversionType::HexIntLower - | ConversionType::HexIntUpper - | ConversionType::OctInt, - Some(value), - ) => self.format_unsigned(string, (*value as u32) as u64), - ( - ConversionType::StringLower | ConversionType::StringUpper, - Some(value), - ) => self.format_string(string, &value.to_string()), - (t, None) if t.supports_integer() => self.format_string(string, "null"), - _ => { - exec_err!( - "Invalid conversion type: {:?} for Int32", - self.conversion_type - ) - } - }, - ScalarValue::Int64(value) => match (self.conversion_type, value) { - (ConversionType::DecInt, Some(value)) => { - self.format_signed(string, *value) - } - ( - ConversionType::HexIntLower - | ConversionType::HexIntUpper - | ConversionType::OctInt, - Some(value), - ) => self.format_unsigned(string, *value as u64), - ( - ConversionType::StringLower | ConversionType::StringUpper, - Some(value), - ) => self.format_string(string, &value.to_string()), - (t, None) if t.supports_integer() => self.format_string(string, "null"), - _ => { - exec_err!( - "Invalid conversion type: {:?} for Int64", - self.conversion_type - ) - } - }, - ScalarValue::UInt8(value) => match (self.conversion_type, value) { - ( - ConversionType::DecInt - | ConversionType::HexIntLower - | ConversionType::HexIntUpper - | ConversionType::OctInt, - Some(value), - ) => self.format_unsigned(string, *value as u64), - ( - ConversionType::StringLower | ConversionType::StringUpper, - Some(value), - ) => self.format_string(string, &value.to_string()), - (t, None) if t.supports_integer() => self.format_string(string, "null"), - _ => { - exec_err!( - "Invalid conversion type: {:?} for UInt8", - self.conversion_type - ) - } - }, - ScalarValue::UInt16(value) => match (self.conversion_type, value) { - ( - ConversionType::DecInt - | ConversionType::HexIntLower - | ConversionType::HexIntUpper - | ConversionType::OctInt, - Some(value), - ) => self.format_unsigned(string, *value as u64), - ( - ConversionType::StringLower | ConversionType::StringUpper, - Some(value), - ) => self.format_string(string, &value.to_string()), - (t, None) if t.supports_integer() => self.format_string(string, "null"), - _ => { - exec_err!( - "Invalid conversion type: {:?} for UInt16", - self.conversion_type - ) - } - }, - ScalarValue::UInt32(value) => match (self.conversion_type, value) { - ( - ConversionType::DecInt - | ConversionType::HexIntLower - | ConversionType::HexIntUpper - | ConversionType::OctInt, - Some(value), - ) => self.format_unsigned(string, *value as u64), - ( - ConversionType::StringLower | ConversionType::StringUpper, - Some(value), - ) => self.format_string(string, &value.to_string()), - (t, None) if t.supports_integer() => self.format_string(string, "null"), - _ => { - exec_err!( - "Invalid conversion type: {:?} for UInt32", - self.conversion_type - ) - } - }, - ScalarValue::UInt64(value) => match (self.conversion_type, value) { - ( - ConversionType::DecInt - | ConversionType::HexIntLower - | ConversionType::HexIntUpper - | ConversionType::OctInt, - Some(value), - ) => self.format_unsigned(string, *value), - ( - ConversionType::StringLower | ConversionType::StringUpper, - Some(value), - ) => self.format_string(string, &value.to_string()), - (t, None) if t.supports_integer() => self.format_string(string, "null"), - _ => { - exec_err!( - "Invalid conversion type: {:?} for UInt64", - self.conversion_type - ) - } - }, + ScalarValue::Int8(value) => self.format_integer(string, value, "Int8"), + ScalarValue::Int16(value) => self.format_integer(string, value, "Int16"), + ScalarValue::Int32(value) => self.format_integer(string, value, "Int32"), + ScalarValue::Int64(value) => self.format_integer(string, value, "Int64"), + ScalarValue::UInt8(value) => self.format_integer(string, value, "UInt8"), + ScalarValue::UInt16(value) => self.format_integer(string, value, "UInt16"), + ScalarValue::UInt32(value) => self.format_integer(string, value, "UInt32"), + ScalarValue::UInt64(value) => self.format_integer(string, value, "UInt64"), ScalarValue::Float16(value) => match (self.conversion_type, value) { ( ConversionType::DecFloatLower @@ -1484,6 +1371,48 @@ impl ConversionSpecifier { } } + fn format_integer( + &self, + writer: &mut String, + value: &Option, + type_name: &str, + ) -> Result<()> + where + T: Copy + IntegerFormatValue, + { + let Some(value) = *value else { + return if self.conversion_type.supports_integer() { + self.format_string(writer, "null") + } else { + self.invalid_integer_conversion(type_name) + }; + }; + + match self.conversion_type { + ConversionType::DecInt => value.format_decimal(self, writer), + ConversionType::HexIntLower + | ConversionType::HexIntUpper + | ConversionType::OctInt => { + self.format_unsigned(writer, value.unsigned_bits()) + } + ConversionType::CharLower | ConversionType::CharUpper => { + self.format_char(writer, value.to_char()?) + } + ConversionType::StringLower | ConversionType::StringUpper => { + self.format_string(writer, &value.decimal_string()) + } + _ => self.invalid_integer_conversion(type_name), + } + } + + fn invalid_integer_conversion(&self, type_name: &str) -> Result { + exec_err!( + "Invalid conversion type: {:?} for {}", + self.conversion_type, + type_name + ) + } + fn format_hex_float(&self, writer: &mut String, value: f64) -> Result<()> { // Handle special cases first let (sign, raw_exponent, mantissa) = value.to_parts(); @@ -1860,34 +1789,8 @@ impl ConversionSpecifier { } } } - // Take care of padding - let NumericParam::Literal(width) = self.width else { - writer.push_str(&prefix); - writer.push_str(&number); - writer.push_str(&suffix); - return Ok(()); - }; - if self.left_adj { - let mut full_num = prefix + &number + &suffix; - while full_num.len() < width as usize { - full_num.push(' '); - } - writer.push_str(&full_num); - } else if self.zero_pad && value.is_finite() { - while prefix.len() + number.len() + suffix.len() < width as usize { - prefix.push('0'); - } - writer.push_str(&prefix); - writer.push_str(&number); - writer.push_str(&suffix); - } else { - let mut full_num = prefix + &number + &suffix; - while full_num.len() < width as usize { - full_num = " ".to_owned() + &full_num; - } - writer.push_str(&full_num); - }; + self.write_numeric_parts(writer, prefix, &number, &suffix, value.is_finite()); Ok(()) } @@ -2053,6 +1956,7 @@ impl ConversionSpecifier { self.validate_grouping_separator()?; let mut prefix = String::new(); + let mut suffix = String::new(); let upper = self.conversion_type.is_upper(); // Parse as BigDecimal @@ -2062,15 +1966,16 @@ impl ConversionSpecifier { let decimal = BigDecimal::from_bigint(decimal, scale); // Handle sign - // TODO: `negative_in_parentheses` (the `(` flag) is not implemented here. - // Java/Spark wrap negative values in parentheses when this flag is set - // (e.g. `%(,.2f` with -1234.5 → "(1,234.50)"), but this path always - // uses a minus sign. See `format_float` for the correct implementation. let is_negative = decimal.sign() == Sign::Minus; let abs_decimal = decimal.abs(); if is_negative { - prefix.push('-'); + if self.negative_in_parentheses { + prefix.push('('); + suffix.push(')'); + } else { + prefix.push('-'); + } } else if self.space_sign { prefix.push(' '); } else if self.force_sign { @@ -2145,33 +2050,7 @@ impl ConversionSpecifier { } }; - // Handle padding - let NumericParam::Literal(width) = self.width else { - writer.push_str(&prefix); - writer.push_str(&number); - return Ok(()); - }; - - if self.left_adj { - let mut full_num = prefix + &number; - while full_num.len() < width as usize { - full_num.push(' '); - } - writer.push_str(&full_num); - } else if self.zero_pad { - while prefix.len() + number.len() < width as usize { - prefix.push('0'); - } - writer.push_str(&prefix); - writer.push_str(&number); - } else { - let mut full_num = prefix + &number; - while full_num.len() < width as usize { - full_num = " ".to_owned() + &full_num; - } - writer.push_str(&full_num); - } - + self.write_numeric_parts(writer, prefix, &number, &suffix, true); Ok(()) } @@ -2335,6 +2214,44 @@ impl ConversionSpecifier { TimeFormat::CLower => Ok(dt.format("%a %b %d %H:%M:%S UTC %Y").to_string()), } } + + fn write_numeric_parts( + &self, + writer: &mut String, + mut prefix: String, + number: &str, + suffix: &str, + zero_pad_allowed: bool, + ) { + // Handle padding + let NumericParam::Literal(width) = self.width else { + writer.push_str(&prefix); + writer.push_str(number); + writer.push_str(suffix); + return; + }; + + if self.left_adj { + let mut full_num = prefix + number + suffix; + while full_num.len() < width as usize { + full_num.push(' '); + } + writer.push_str(&full_num); + } else if self.zero_pad && zero_pad_allowed { + while prefix.len() + number.len() + suffix.len() < width as usize { + prefix.push('0'); + } + writer.push_str(&prefix); + writer.push_str(number); + writer.push_str(suffix); + } else { + let mut full_num = prefix + number + suffix; + while full_num.len() < width as usize { + full_num = " ".to_owned() + &full_num; + } + writer.push_str(&full_num); + } + } } trait FloatFormattable: std::fmt::Display { @@ -2443,7 +2360,7 @@ mod tests { use super::*; use crate::function::utils::test::test_scalar_function; use arrow::array::StringArray; - use arrow::datatypes::DataType::Utf8; + use arrow::datatypes::{DataType::Utf8, i256}; #[test] fn test_format_string_nullability() -> Result<()> { @@ -2588,6 +2505,74 @@ mod tests { ); } + #[test] + fn test_integer_formatting_across_widths() -> Result<()> { + let cases = [ + ( + ScalarValue::Int8(Some(-1)), + "%d|%x|%o|%s", + 4, + "-1|ff|377|-1", + ), + ( + ScalarValue::Int16(Some(-1)), + "%d|%x|%o|%s", + 4, + "-1|ffff|177777|-1", + ), + ( + ScalarValue::Int32(Some(-1)), + "%d|%x|%o|%s", + 4, + "-1|ffffffff|37777777777|-1", + ), + ( + ScalarValue::Int64(Some(-1)), + "%d|%x|%o|%s", + 4, + "-1|ffffffffffffffff|1777777777777777777777|-1", + ), + ( + ScalarValue::UInt8(Some(255)), + "%d|%x|%o|%s|%c", + 5, + "255|ff|377|255|ÿ", + ), + ( + ScalarValue::UInt16(Some(65535)), + "%d|%x|%o|%s", + 4, + "65535|ffff|177777|65535", + ), + ( + ScalarValue::UInt32(Some(u32::MAX)), + "%d|%x|%o|%s", + 4, + "4294967295|ffffffff|37777777777|4294967295", + ), + ( + ScalarValue::UInt64(Some(u64::MAX)), + "%d|%x|%o|%s", + 4, + "18446744073709551615|ffffffffffffffff|1777777777777777777777|18446744073709551615", + ), + ( + ScalarValue::Int32(None), + "%d|%x|%o|%s|%c", + 5, + "null|null|null|null|null", + ), + ]; + + for (value, fmt, arg_count, expected) in cases { + let data_types = vec![value.data_type(); arg_count]; + let formatter = Formatter::parse(fmt, &data_types)?; + let args = vec![value; arg_count]; + assert_eq!(formatter.format(&args)?, expected, "{fmt}"); + } + Ok(()) + } + #[test] fn test_insert_thousands_separator() { assert_eq!(insert_thousands_separator("1234567.89"), "1,234,567.89"); @@ -2899,17 +2884,90 @@ mod tests { #[test] fn test_grouping_separator_parentheses_decimal() -> Result<()> { - // %(,15.2f on negative decimal — format_decimal ignores negative_in_parentheses, - // always uses '-'. Check TODO in fn format_decimal + test_scalar_function!( + FormatStringFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some("%(,.2f".to_string()))), + ColumnarValue::Scalar(ScalarValue::Decimal128(Some(-123450), 10, 2)), + ], + Ok(Some("(1,234.50)")), + &str, + Utf8, + StringArray + ); + + test_scalar_function!( + FormatStringFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some("%(,.2f".to_string()))), + ColumnarValue::Scalar(ScalarValue::Decimal256( + Some(i256::from(-123450)), + 10, + 2, + )), + ], + Ok(Some("(1,234.50)")), + &str, + Utf8, + StringArray + ); + // Java: String.format("%(,15.2f", -1234.5) → " (1,234.50)" - // Ours: " -1,234.50" (minus sign, no parens) test_scalar_function!( FormatStringFunc::new(), vec![ ColumnarValue::Scalar(ScalarValue::Utf8(Some("%(,15.2f".to_string()))), ColumnarValue::Scalar(ScalarValue::Decimal128(Some(-123450), 10, 2)), ], - Ok(Some(" -1,234.50")), + Ok(Some(" (1,234.50)")), + &str, + Utf8, + StringArray + ); + Ok(()) + } + + #[test] + fn test_grouping_separator_ignore_zero_padding_for_float_nan() -> Result<()> { + test_scalar_function!( + FormatStringFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some("%010.2f".to_string()))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(f64::NAN))), + ], + Ok(Some(" NaN")), + &str, + Utf8, + StringArray + ); + Ok(()) + } + + #[test] + fn test_grouping_separator_ignore_zero_padding_for_float_inf() -> Result<()> { + test_scalar_function!( + FormatStringFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some("%010.2f".to_string()))), + ColumnarValue::Scalar(ScalarValue::Float64(Some(f64::INFINITY))), + ], + Ok(Some(" Infinity")), + &str, + Utf8, + StringArray + ); + Ok(()) + } + + #[test] + fn test_grouping_separator_parentheses_zero_padding_decimal() -> Result<()> { + test_scalar_function!( + FormatStringFunc::new(), + vec![ + ColumnarValue::Scalar(ScalarValue::Utf8(Some("%(0,15.2f".to_string()))), + ColumnarValue::Scalar(ScalarValue::Decimal128(Some(-123450), 2, 2)), + ], + Ok(Some("(000001,234.50)")), &str, Utf8, StringArray diff --git a/datafusion/spark/src/function/string/mod.rs b/datafusion/spark/src/function/string/mod.rs index 64d603cb8bb67..bc94c27732c91 100644 --- a/datafusion/spark/src/function/string/mod.rs +++ b/datafusion/spark/src/function/string/mod.rs @@ -19,6 +19,7 @@ pub mod ascii; pub mod base64; pub mod char; pub mod concat; +pub mod concat_ws; pub mod elt; pub mod format_string; pub mod ilike; @@ -27,6 +28,7 @@ pub mod length; pub mod like; pub mod luhn_check; pub mod make_valid_utf8; +pub mod quote; pub mod soundex; pub mod space; pub mod substring; @@ -39,6 +41,7 @@ make_udf_function!(ascii::SparkAscii, ascii); make_udf_function!(base64::SparkBase64, base64); make_udf_function!(char::CharFunc, char); make_udf_function!(concat::SparkConcat, concat); +make_udf_function!(concat_ws::SparkConcatWs, concat_ws); make_udf_function!(ilike::SparkILike, ilike); make_udf_function!(length::SparkLengthFunc, length); make_udf_function!(elt::SparkElt, elt); @@ -51,6 +54,7 @@ make_udf_function!(base64::SparkUnBase64, unbase64); make_udf_function!(soundex::SparkSoundex, soundex); make_udf_function!(make_valid_utf8::SparkMakeValidUtf8, make_valid_utf8); make_udf_function!(is_valid_utf8::SparkIsValidUtf8, is_valid_utf8); +make_udf_function!(quote::SparkQuote, quote); pub mod expr_fn { use datafusion_functions::export_functions; @@ -75,6 +79,11 @@ pub mod expr_fn { "Concatenates multiple input strings into a single string. Returns NULL if any input is NULL.", args )); + export_functions!(( + concat_ws, + "Concatenates strings with separator. Supports arrays. Null values are skipped.", + sep args + )); export_functions!(( elt, "Returns the n-th input (1-indexed), e.g. returns 2nd input when n is 2. The function returns NULL if the index is 0 or exceeds the length of the array.", @@ -127,6 +136,11 @@ pub mod expr_fn { "Returns the original string if str is a valid UTF-8 string, otherwise returns a new string whose invalid UTF8 byte sequences are replaced using the UNICODE replacement character U+FFFD.", str )); + export_functions!(( + quote, + "Returns str enclosed by single quotes and each instance of single quote in it is preceded by a backslash", + str + )); } pub fn functions() -> Vec> { @@ -135,6 +149,7 @@ pub fn functions() -> Vec> { base64(), char(), concat(), + concat_ws(), elt(), ilike(), length(), @@ -147,5 +162,6 @@ pub fn functions() -> Vec> { soundex(), make_valid_utf8(), is_valid_utf8(), + quote(), ] } diff --git a/datafusion/spark/src/function/string/quote.rs b/datafusion/spark/src/function/string/quote.rs new file mode 100644 index 0000000000000..39ad8bf841764 --- /dev/null +++ b/datafusion/spark/src/function/string/quote.rs @@ -0,0 +1,121 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use arrow::array::{ArrayRef, OffsetSizeTrait, StringArray}; +use arrow::datatypes::DataType; +use datafusion::logical_expr::{Coercion, ColumnarValue, Signature, TypeSignatureClass}; +use datafusion_common::cast::{as_generic_string_array, as_string_view_array}; +use datafusion_common::types::{NativeType, logical_string}; +use datafusion_common::utils::take_function_args; +use datafusion_common::{Result, exec_err}; +use datafusion_expr::{ScalarFunctionArgs, ScalarUDFImpl, Volatility}; +use datafusion_functions::utils::make_scalar_function; + +use std::sync::Arc; + +/// Spark-compatible `quote` expression +/// +#[derive(Debug, PartialEq, Eq, Hash)] +pub struct SparkQuote { + signature: Signature, +} + +impl Default for SparkQuote { + fn default() -> Self { + Self::new() + } +} + +impl SparkQuote { + pub fn new() -> Self { + let str_coercion = Coercion::new_implicit( + TypeSignatureClass::Native(logical_string()), + vec![TypeSignatureClass::Any], + NativeType::String, + ); + Self { + signature: Signature::coercible(vec![str_coercion], Volatility::Immutable), + } + } +} + +impl ScalarUDFImpl for SparkQuote { + fn name(&self) -> &str { + "quote" + } + + fn signature(&self) -> &Signature { + &self.signature + } + + fn return_type(&self, arg_types: &[DataType]) -> Result { + match &arg_types[0] { + DataType::LargeUtf8 => Ok(DataType::LargeUtf8), + _ => Ok(DataType::Utf8), + } + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + make_scalar_function(spark_quote_inner, vec![])(&args.args) + } +} + +fn spark_quote_inner(arg: &[ArrayRef]) -> Result { + let [array] = take_function_args("quote", arg)?; + match &array.data_type() { + DataType::Utf8 => quote_array::(array), + DataType::LargeUtf8 => quote_array::(array), + DataType::Utf8View => quote_view(array), + other => { + exec_err!("unsupported data type {other:?} for function `quote`") + } + } +} + +fn quote_array(array: &ArrayRef) -> Result { + let str_array = as_generic_string_array::(array)?; + let result = str_array + .iter() + .map(|s| s.map(compute_quote)) + .collect::(); + Ok(Arc::new(result)) +} + +fn quote_view(str_view: &ArrayRef) -> Result { + let str_array = as_string_view_array(str_view)?; + let result = str_array + .iter() + .map(|opt_str| opt_str.map(compute_quote)) + .collect::(); + Ok(Arc::new(result) as ArrayRef) +} + +const QUOTE_CHAR: char = '\''; +const ESCAPE_CHAR: char = '\\'; + +fn compute_quote(s: &str) -> String { + let mut quoted = String::with_capacity(s.len() + 2); + quoted.push(QUOTE_CHAR); + for c in s.chars() { + if c == QUOTE_CHAR { + quoted.push(ESCAPE_CHAR); + } + quoted.push(c); + } + quoted.push(QUOTE_CHAR); + quoted +} diff --git a/datafusion/spark/src/function/url/parse_url.rs b/datafusion/spark/src/function/url/parse_url.rs index 18f0bb1e0d78b..9ceed8b155bbd 100644 --- a/datafusion/spark/src/function/url/parse_url.rs +++ b/datafusion/spark/src/function/url/parse_url.rs @@ -18,8 +18,8 @@ use std::sync::Arc; use arrow::array::{ - Array, ArrayRef, GenericStringBuilder, LargeStringArray, StringArray, - StringArrayType, StringViewArray, + Array, ArrayRef, AsArray, LargeStringArray, StringArray, StringArrayType, + StringViewArray, new_null_array, }; use arrow::datatypes::DataType; use datafusion_common::cast::{ @@ -272,20 +272,18 @@ pub fn spark_handled_parse_url( ), } } else { - // The 'key' argument is omitted, assume all values are null - // Create 'null' string array for 'key' argument - let mut builder: GenericStringBuilder = GenericStringBuilder::new(); - for _ in 0..args[0].len() { - builder.append_null(); - } - let key = builder.finish(); + // The 'key' argument is omitted, assume all values are null. + // `new_null_array` allocates the null array outright, rather than + // appending one null per row through a builder. + let key_array = new_null_array(&DataType::Utf8, args[0].len()); + let key = key_array.as_string::(); match (url.data_type(), part.data_type()) { (DataType::Utf8, DataType::Utf8) => { process_parse_url::<_, _, _, StringArray>( as_string_array(url)?, as_string_array(part)?, - &key, + key, handler_err, false, ) @@ -294,7 +292,7 @@ pub fn spark_handled_parse_url( process_parse_url::<_, _, _, StringViewArray>( as_string_view_array(url)?, as_string_view_array(part)?, - &key, + key, handler_err, false, ) @@ -303,7 +301,7 @@ pub fn spark_handled_parse_url( process_parse_url::<_, _, _, LargeStringArray>( as_large_string_array(url)?, as_large_string_array(part)?, - &key, + key, handler_err, false, ) diff --git a/datafusion/spark/src/function/url/try_url_decode.rs b/datafusion/spark/src/function/url/try_url_decode.rs index 78968288fc2f5..1acbd5e13b988 100644 --- a/datafusion/spark/src/function/url/try_url_decode.rs +++ b/datafusion/spark/src/function/url/try_url_decode.rs @@ -24,7 +24,9 @@ use datafusion_expr::{ }; use datafusion_functions::utils::make_scalar_function; -use crate::function::url::url_decode::{UrlDecode, spark_handled_url_decode}; +use crate::function::url::url_decode::{ + OnDecodeError, UrlDecode, spark_handled_url_decode, +}; #[derive(Debug, PartialEq, Eq, Hash)] pub struct TryUrlDecode { @@ -67,10 +69,7 @@ impl ScalarUDFImpl for TryUrlDecode { } fn spark_try_url_decode(args: &[ArrayRef]) -> Result { - spark_handled_url_decode(args, |x| match x { - Err(_) => Ok(None), - result => result, - }) + spark_handled_url_decode(args, OnDecodeError::Null) } #[cfg(test)] diff --git a/datafusion/spark/src/function/url/url_decode.rs b/datafusion/spark/src/function/url/url_decode.rs index 0966cc380e497..0e527068b7cb3 100644 --- a/datafusion/spark/src/function/url/url_decode.rs +++ b/datafusion/spark/src/function/url/url_decode.rs @@ -18,7 +18,9 @@ use std::borrow::Cow; use std::sync::Arc; -use arrow::array::{ArrayRef, LargeStringArray, StringArray, StringViewArray}; +use arrow::array::{ + Array, ArrayRef, LargeStringBuilder, StringBuilder, StringViewBuilder, +}; use arrow::datatypes::DataType; use datafusion_common::cast::{ as_large_string_array, as_string_array, as_string_view_array, @@ -61,18 +63,25 @@ impl UrlDecode { /// /// # Returns /// - /// * `Ok(String)` - The decoded string + /// * `Ok(Cow)` - The decoded string, borrowed from `value` when there + /// was nothing to rewrite and owned otherwise /// * `Err(DataFusionError)` - If the input is malformed or contains invalid UTF-8 - /// - fn decode(value: &str) -> Result { + fn decode(value: &str) -> Result> { // Check if the string has valid percent encoding Self::validate_percent_encoding(value)?; - let replaced = Self::replace_plus(value.as_bytes()); - percent_decode(&replaced) - .decode_utf8() - .map_err(|e| exec_datafusion_err!("Invalid UTF-8 sequence: {e}")) - .map(|parsed| parsed.into_owned()) + match Self::replace_plus(value.as_bytes()) { + // No '+' was rewritten, so the decode can borrow from `value` itself. + Cow::Borrowed(bytes) => percent_decode(bytes) + .decode_utf8() + .map_err(|e| exec_datafusion_err!("Invalid UTF-8 sequence: {e}")), + // Rewriting '+' already allocated, so owning the decoded form here + // costs nothing beyond what has been spent. + Cow::Owned(bytes) => percent_decode(&bytes) + .decode_utf8() + .map(|decoded| Cow::Owned(decoded.into_owned())) + .map_err(|e| exec_datafusion_err!("Invalid UTF-8 sequence: {e}")), + } } /// Replace b'+' with b' ' @@ -155,6 +164,15 @@ impl ScalarUDFImpl for UrlDecode { } } +/// How [`spark_handled_url_decode`] reacts to a malformed input value. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum OnDecodeError { + /// Propagate the error, as `url_decode` does. + Fail, + /// Return NULL for that row, as `try_url_decode` does. + Null, +} + /// Core implementation of URL decoding function. /// /// # Arguments @@ -165,38 +183,59 @@ impl ScalarUDFImpl for UrlDecode { /// /// * `Ok(ArrayRef)` - A new array of the same type containing decoded strings /// * `Err(DataFusionError)` - If validation fails or invalid arguments are provided -/// fn spark_url_decode(args: &[ArrayRef]) -> Result { - spark_handled_url_decode(args, |x| x) + spark_handled_url_decode(args, OnDecodeError::Fail) } pub fn spark_handled_url_decode( args: &[ArrayRef], - err_handle_fn: impl Fn(Result>) -> Result>, + on_error: OnDecodeError, ) -> Result { if args.len() != 1 { return exec_err!("`url_decode` expects 1 argument"); } + // Decoded values go straight into the builder, so a row that needs no + // unescaping is copied once rather than materialised as its own `String`. + macro_rules! decode_all { + ($array:expr, $builder:expr) => {{ + let array = $array; + let mut builder = $builder; + for value in array.iter() { + let Some(value) = value else { + builder.append_null(); + continue; + }; + match UrlDecode::decode(value) { + Ok(decoded) => builder.append_value(&decoded), + Err(e) => match on_error { + OnDecodeError::Fail => return Err(e), + OnDecodeError::Null => builder.append_null(), + }, + } + } + Ok(Arc::new(builder.finish()) as ArrayRef) + }}; + } + match &args[0].data_type() { - DataType::Utf8 => as_string_array(&args[0])? - .iter() - .map(|x| x.map(UrlDecode::decode).transpose()) - .map(&err_handle_fn) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), - DataType::LargeUtf8 => as_large_string_array(&args[0])? - .iter() - .map(|x| x.map(UrlDecode::decode).transpose()) - .map(&err_handle_fn) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), - DataType::Utf8View => as_string_view_array(&args[0])? - .iter() - .map(|x| x.map(UrlDecode::decode).transpose()) - .map(&err_handle_fn) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), + DataType::Utf8 => { + let array = as_string_array(&args[0])?; + let builder = + StringBuilder::with_capacity(array.len(), array.value_data().len()); + decode_all!(array, builder) + } + DataType::LargeUtf8 => { + let array = as_large_string_array(&args[0])?; + let builder = + LargeStringBuilder::with_capacity(array.len(), array.value_data().len()); + decode_all!(array, builder) + } + DataType::Utf8View => { + let array = as_string_view_array(&args[0])?; + let builder = StringViewBuilder::with_capacity(array.len()); + decode_all!(array, builder) + } other => exec_err!("`url_decode`: Expr must be STRING, got {other:?}"), } } @@ -205,49 +244,77 @@ pub fn spark_handled_url_decode( mod tests { use super::*; + use arrow::array::{LargeStringArray, StringArray, StringViewArray}; + + const INPUT: [Option<&str>; 7] = [ + Some("https%3A%2F%2Fspark.apache.org"), + Some("inva+lid://user:pass@host/file\\;param?query\\;p2"), + Some("inva lid://user:pass@host/file\\;param?query\\;p2"), + Some("%7E%21%40%23%24%25%5E%26%2A%28%29%5F%2B"), + Some("%E4%BD%A0%E5%A5%BD"), + Some(""), + None, + ]; + + const EXPECTED: [Option<&str>; 7] = [ + Some("https://spark.apache.org"), + Some("inva lid://user:pass@host/file\\;param?query\\;p2"), + Some("inva lid://user:pass@host/file\\;param?query\\;p2"), + Some("~!@#$%^&*()_+"), + Some("你好"), + Some(""), + None, + ]; + + // '%2s' is not a valid percent encoded character + const MALFORMED_INPUT: [Option<&str>; 3] = [ + Some("http%3A%2F%2spark.apache.org"), + // Valid cases + Some("https%3A%2F%2Fspark.apache.org"), + None, + ]; #[test] - fn test_decode() -> Result<()> { - let input = Arc::new(StringArray::from(vec![ - Some("https%3A%2F%2Fspark.apache.org"), - Some("inva+lid://user:pass@host/file\\;param?query\\;p2"), - Some("inva lid://user:pass@host/file\\;param?query\\;p2"), - Some("%7E%21%40%23%24%25%5E%26%2A%28%29%5F%2B"), - Some("%E4%BD%A0%E5%A5%BD"), - Some(""), - None, - ])); - let expected = StringArray::from(vec![ - Some("https://spark.apache.org"), - Some("inva lid://user:pass@host/file\\;param?query\\;p2"), - Some("inva lid://user:pass@host/file\\;param?query\\;p2"), - Some("~!@#$%^&*()_+"), - Some("你好"), - Some(""), - None, - ]); - - let result = spark_url_decode(&[input as ArrayRef])?; + fn test_decode_utf8() -> Result<()> { + let input = Arc::new(StringArray::from(INPUT.to_vec())) as ArrayRef; + let result = spark_url_decode(&[input])?; let result = as_string_array(&result)?; + assert_eq!(&StringArray::from(EXPECTED.to_vec()), result); + Ok(()) + } - assert_eq!(&expected, result); + #[test] + fn test_decode_large_utf8() -> Result<()> { + let input = Arc::new(LargeStringArray::from(INPUT.to_vec())) as ArrayRef; + let result = spark_url_decode(&[input])?; + let result = as_large_string_array(&result)?; + assert_eq!(&LargeStringArray::from(EXPECTED.to_vec()), result); + Ok(()) + } + #[test] + fn test_decode_utf8_view() -> Result<()> { + let input = Arc::new(StringViewArray::from(INPUT.to_vec())) as ArrayRef; + let result = spark_url_decode(&[input])?; + let result = as_string_view_array(&result)?; + assert_eq!(&StringViewArray::from(EXPECTED.to_vec()), result); Ok(()) } #[test] fn test_decode_error() -> Result<()> { - let input = Arc::new(StringArray::from(vec![ - Some("http%3A%2F%2spark.apache.org"), // '%2s' is not a valid percent encoded character - // Valid cases - Some("https%3A%2F%2Fspark.apache.org"), - None, - ])); - - let result = spark_url_decode(&[input]); - assert!( - result.is_err_and(|e| e.to_string().contains("Invalid percent-encoding")) - ); + let inputs: [ArrayRef; 3] = [ + Arc::new(StringArray::from(MALFORMED_INPUT.to_vec())), + Arc::new(LargeStringArray::from(MALFORMED_INPUT.to_vec())), + Arc::new(StringViewArray::from(MALFORMED_INPUT.to_vec())), + ]; + + for input in inputs { + let result = spark_url_decode(&[input]); + assert!( + result.is_err_and(|e| e.to_string().contains("Invalid percent-encoding")) + ); + } Ok(()) } diff --git a/datafusion/spark/src/function/url/url_encode.rs b/datafusion/spark/src/function/url/url_encode.rs index 1ad2a111851ee..87a70af4ac5b6 100644 --- a/datafusion/spark/src/function/url/url_encode.rs +++ b/datafusion/spark/src/function/url/url_encode.rs @@ -17,7 +17,9 @@ use std::sync::Arc; -use arrow::array::{ArrayRef, LargeStringArray, StringArray, StringViewArray}; +use arrow::array::{ + Array, ArrayRef, LargeStringBuilder, StringBuilder, StringViewBuilder, +}; use arrow::datatypes::DataType; use datafusion_common::cast::{ as_large_string_array, as_string_array, as_string_view_array, @@ -46,20 +48,6 @@ impl UrlEncode { signature: Signature::string(1, Volatility::Immutable), } } - - /// Encode a string to application/x-www-form-urlencoded format. - /// - /// # Arguments - /// - /// * `value` - The string to encode - /// - /// # Returns - /// - /// * `Ok(String)` - The encoded string - /// - fn encode(value: &str) -> Result { - Ok(byte_serialize(value.as_bytes()).collect::()) - } } impl ScalarUDFImpl for UrlEncode { @@ -105,22 +93,94 @@ fn spark_url_encode(args: &[ArrayRef]) -> Result { return exec_err!("`url_encode` expects 1 argument"); } + // The percent-encoded form of each value is assembled in a single scratch buffer + // reused across rows, rather than allocating a `String` per row. + macro_rules! encode_all { + ($array:expr, $builder:expr) => {{ + let array = $array; + let mut builder = $builder; + let mut encoded = String::new(); + for value in array.iter() { + match value { + Some(value) => { + encoded.clear(); + encoded.extend(byte_serialize(value.as_bytes())); + builder.append_value(&encoded); + } + None => builder.append_null(), + } + } + Ok(Arc::new(builder.finish()) as ArrayRef) + }}; + } + match &args[0].data_type() { - DataType::Utf8 => as_string_array(&args[0])? - .iter() - .map(|x| x.map(UrlEncode::encode).transpose()) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), - DataType::LargeUtf8 => as_large_string_array(&args[0])? - .iter() - .map(|x| x.map(UrlEncode::encode).transpose()) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), - DataType::Utf8View => as_string_view_array(&args[0])? - .iter() - .map(|x| x.map(UrlEncode::encode).transpose()) - .collect::>() - .map(|array| Arc::new(array) as ArrayRef), + DataType::Utf8 => { + let array = as_string_array(&args[0])?; + let builder = + StringBuilder::with_capacity(array.len(), array.value_data().len()); + encode_all!(array, builder) + } + DataType::LargeUtf8 => { + let array = as_large_string_array(&args[0])?; + let builder = + LargeStringBuilder::with_capacity(array.len(), array.value_data().len()); + encode_all!(array, builder) + } + DataType::Utf8View => { + let array = as_string_view_array(&args[0])?; + let builder = StringViewBuilder::with_capacity(array.len()); + encode_all!(array, builder) + } other => exec_err!("`url_encode`: Expr must be STRING, got {other:?}"), } } + +#[cfg(test)] +mod tests { + use super::*; + use arrow::array::{LargeStringArray, StringArray, StringViewArray}; + + const INPUT: [Option<&str>; 5] = [ + Some("https://spark.apache.org"), + Some("inva lid://user:pass@host/file\\;param?query\\;p2"), + Some("你好"), + Some(""), + None, + ]; + + const EXPECTED: [Option<&str>; 5] = [ + Some("https%3A%2F%2Fspark.apache.org"), + Some("inva+lid%3A%2F%2Fuser%3Apass%40host%2Ffile%5C%3Bparam%3Fquery%5C%3Bp2"), + Some("%E4%BD%A0%E5%A5%BD"), + Some(""), + None, + ]; + + #[test] + fn test_encode_utf8() -> Result<()> { + let input = Arc::new(StringArray::from(INPUT.to_vec())) as ArrayRef; + let result = spark_url_encode(&[input])?; + let result = as_string_array(&result)?; + assert_eq!(&StringArray::from(EXPECTED.to_vec()), result); + Ok(()) + } + + #[test] + fn test_encode_large_utf8() -> Result<()> { + let input = Arc::new(LargeStringArray::from(INPUT.to_vec())) as ArrayRef; + let result = spark_url_encode(&[input])?; + let result = as_large_string_array(&result)?; + assert_eq!(&LargeStringArray::from(EXPECTED.to_vec()), result); + Ok(()) + } + + #[test] + fn test_encode_utf8_view() -> Result<()> { + let input = Arc::new(StringViewArray::from(INPUT.to_vec())) as ArrayRef; + let result = spark_url_encode(&[input])?; + let result = as_string_view_array(&result)?; + assert_eq!(&StringViewArray::from(EXPECTED.to_vec()), result); + Ok(()) + } +} diff --git a/datafusion/spark/src/lib.rs b/datafusion/spark/src/lib.rs index 2eee94c52ef78..6cd4678da7560 100644 --- a/datafusion/spark/src/lib.rs +++ b/datafusion/spark/src/lib.rs @@ -59,7 +59,7 @@ //! # fn udafs(&self) -> HashSet { unimplemented!() } //! # fn udwfs(&self) -> HashSet { unimplemented!() } //! # fn udf(&self, _name: &str) -> Result> { unimplemented!() } -//! # fn higher_order_function(&self, name: &str) -> Result> { unimplemented!() } +//! # fn higher_order_function(&self, name: &str) -> Result> { unimplemented!() } //! # fn udaf(&self, name: &str) -> Result> {unimplemented!() } //! # fn udwf(&self, name: &str) -> Result> { unimplemented!() } //! # fn expr_planners(&self) -> Vec> { unimplemented!() } diff --git a/datafusion/spark/src/session_state.rs b/datafusion/spark/src/session_state.rs index e39de3a5888ea..839487772a9b2 100644 --- a/datafusion/spark/src/session_state.rs +++ b/datafusion/spark/src/session_state.rs @@ -88,6 +88,9 @@ impl SessionStateBuilderSpark for SessionStateBuilder { #[cfg(test)] mod tests { use super::*; + use datafusion::common::config::Dialect; + use datafusion::prelude::SessionConfig; + use datafusion::prelude::SessionContext; #[test] fn test_session_state_with_spark_features() { @@ -108,4 +111,37 @@ mod tests { "Apache Spark expr planners should be registered" ); } + + #[tokio::test] + async fn test_spark_dialect_with_spark_functions() { + let query = "SELECT sha2('abc', 256), CAST(1 AS LONG)"; + + let mut config = SessionConfig::new(); + config.options_mut().sql_parser.dialect = Dialect::Spark; + let state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .with_spark_features() + .build(); + let ctx = SessionContext::new_with_state(state); + + let result = ctx.sql(query).await.unwrap().collect().await.unwrap(); + assert_eq!(result.len(), 1); + assert_eq!(result[0].num_rows(), 1); + + let mut config = SessionConfig::new(); + config.options_mut().sql_parser.dialect = Dialect::Generic; + let state = SessionStateBuilder::new() + .with_config(config) + .with_default_features() + .with_spark_features() + .build(); + let ctx = SessionContext::new_with_state(state); + + let err = ctx.sql(query).await.unwrap_err().to_string(); + assert!( + err.contains("Unsupported SQL type LONG"), + "unexpected error: {err}" + ); + } } diff --git a/datafusion/sql/Cargo.toml b/datafusion/sql/Cargo.toml index cc299ce507099..318f8934639aa 100644 --- a/datafusion/sql/Cargo.toml +++ b/datafusion/sql/Cargo.toml @@ -44,7 +44,7 @@ name = "datafusion_sql" default = ["unicode_expressions", "unparser"] unicode_expressions = [] unparser = [] -recursive_protection = ["dep:recursive"] +recursive_protection = ["dep:recursive", "dep:stacker"] # Note the sql planner should not depend directly on the datafusion-function packages # so that it can be used in a standalone manner with other function implementations. @@ -62,6 +62,7 @@ log = { workspace = true } recursive = { workspace = true, optional = true } regex = { workspace = true } sqlparser = { workspace = true } +stacker = { workspace = true, optional = true } [dev-dependencies] ctor = { workspace = true } diff --git a/datafusion/sql/examples/sql.rs b/datafusion/sql/examples/sql.rs index dc49b4460fec5..883439fbf1e09 100644 --- a/datafusion/sql/examples/sql.rs +++ b/datafusion/sql/examples/sql.rs @@ -138,7 +138,7 @@ impl ContextProvider for MyContextProvider { None } - fn get_higher_order_meta(&self, _name: &str) -> Option> { + fn get_higher_order_meta(&self, _name: &str) -> Option> { None } diff --git a/datafusion/sql/src/cte.rs b/datafusion/sql/src/cte.rs index 18766d7056355..f735b336018cc 100644 --- a/datafusion/sql/src/cte.rs +++ b/datafusion/sql/src/cte.rs @@ -19,12 +19,13 @@ use std::sync::Arc; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; +use arrow::datatypes::{Schema, SchemaRef}; use datafusion_common::{ - Result, not_impl_err, plan_err, + Result, TableReference, not_impl_err, plan_err, tree_node::{TreeNode, TreeNodeRecursion}, }; use datafusion_expr::{LogicalPlan, LogicalPlanBuilder, TableSource}; -use sqlparser::ast::{Query, SetExpr, SetOperator, With}; +use sqlparser::ast::{Ident, Query, SetExpr, SetOperator, With}; impl SqlToRel<'_, S> { pub(super) fn plan_with_clause( @@ -45,14 +46,24 @@ impl SqlToRel<'_, S> { // Create a logical plan for the CTE let cte_plan = if is_recursive { - self.recursive_cte(&cte_name, *cte.query, planner_context)? + let columns = cte.alias.columns.iter().map(|c| c.name.clone()).collect(); + self.recursive_cte(&cte_name, columns, *cte.query, planner_context)? } else { self.non_recursive_cte(*cte.query, planner_context)? }; - // Each `WITH` block can change the column names in the last - // projection (e.g. "WITH table(t1, t2) AS SELECT 1, 2"). - let final_plan = self.apply_table_alias(cte_plan, cte.alias)?; + // Each `WITH` block can change the column names in the last projection + // (e.g. "WITH table(t1, t2) AS SELECT 1, 2"). Recursive CTEs apply those + // to the static term in recursive_cte(), so only the relation name here. + let final_plan = if is_recursive { + LogicalPlanBuilder::from(cte_plan) + .alias(TableReference::bare( + self.ident_normalizer.normalize(cte.alias.name), + ))? + .build()? + } else { + self.apply_table_alias(cte_plan, cte.alias)? + }; // Export the CTE to the outer query planner_context.insert_cte(cte_name, final_plan); } @@ -70,6 +81,7 @@ impl SqlToRel<'_, S> { fn recursive_cte( &self, cte_name: &str, + columns: Vec, mut cte_query: Query, planner_context: &mut PlannerContext, ) -> Result { @@ -90,9 +102,11 @@ impl SqlToRel<'_, S> { set_quantifier, } => (left, right, set_quantifier), other => { - // If the query is not a UNION, then it is not a recursive CTE + // Not a UNION, so not actually a recursive CTE. The caller adds only + // the relation name for recursive CTEs, so apply the column aliases here. *cte_query.body = other; - return self.non_recursive_cte(cte_query, planner_context); + let plan = self.non_recursive_cte(cte_query, planner_context)?; + return self.apply_expr_alias(plan, columns); } }; @@ -110,6 +124,10 @@ impl SqlToRel<'_, S> { // ---------- Step 1: Compile the static term ------------------ let static_plan = self.set_expr_to_plan(*left_expr, planner_context)?; + // Apply the declared column-list aliases (e.g. `t(n)`) to the static term, so + // the work table built from its schema below exposes the declared names. + let static_plan = self.apply_expr_alias(static_plan, columns)?; + // Since the recursive CTEs include a component that references a // table with its name, like the example below: // @@ -127,15 +145,19 @@ impl SqlToRel<'_, S> { // in the case of DataFusion). // // Since we can't simply register a table during planning stage (it is - // an execution problem), we'll use a relation object that preserves the - // schema of the input perfectly and also knows which recursive CTE it is - // bound to. + // an execution problem), we'll use a relation object that knows which + // recursive CTE it is bound to. // ---------- Step 2: Create a temporary relation ------------------ // Step 2.1: Create a table source for the temporary relation - let work_table_source = self - .context_provider - .create_cte_work_table(cte_name, Arc::clone(static_plan.schema().inner()))?; + // Recursive self-references must expose conservative (nullable) + // columns. Deriving them from the static term's possibly non-nullable + // schema would let the recursive term treat values from previous + // iterations as non-nullable. + let work_table_source = self.context_provider.create_cte_work_table( + cte_name, + nullable_schema(static_plan.schema().inner()), + )?; // Step 2.2: Create a temporary relation logical plan that will be used // as the input to the recursive term @@ -184,6 +206,19 @@ impl SqlToRel<'_, S> { } } +/// Return a copy of `schema` with every field marked nullable, preserving field +/// and schema metadata. +fn nullable_schema(schema: &Schema) -> SchemaRef { + Arc::new(Schema::new_with_metadata( + schema + .fields() + .iter() + .map(|field| field.as_ref().clone().with_nullable(true)) + .collect::>(), + schema.metadata().clone(), + )) +} + fn has_work_table_reference( plan: &LogicalPlan, work_table_source: &Arc, diff --git a/datafusion/sql/src/expr/binary_op.rs b/datafusion/sql/src/expr/binary_op.rs index 4e9025e02e0c7..c3e7939370e85 100644 --- a/datafusion/sql/src/expr/binary_op.rs +++ b/datafusion/sql/src/expr/binary_op.rs @@ -16,8 +16,10 @@ // under the License. use crate::planner::{ContextProvider, SqlToRel}; -use datafusion_common::{Result, not_impl_err}; +use datafusion_common::{Result, internal_datafusion_err, not_impl_err}; use datafusion_expr::Operator; +use datafusion_expr::expr::ScalarFunction; +use datafusion_expr::{BinaryExpr, Expr}; use sqlparser::ast::BinaryOperator; impl SqlToRel<'_, S> { @@ -72,4 +74,34 @@ impl SqlToRel<'_, S> { _ => not_impl_err!("Unsupported binary operator: {:?}", op), } } + + pub(crate) fn build_binary_expr( + &self, + op: &BinaryOperator, + left: Expr, + right: Expr, + ) -> Result { + if matches!(op, BinaryOperator::PGExp) { + let fun_name = "power"; + let fun = self + .context_provider + .get_function_meta(fun_name) + .ok_or_else(|| { + internal_datafusion_err!( + "Unable to find expected '{fun_name}' function" + ) + })?; + + return Ok(Expr::ScalarFunction(ScalarFunction::new_udf( + fun, + vec![left, right], + ))); + } + + Ok(Expr::BinaryExpr(BinaryExpr::new( + Box::new(left), + self.parse_sql_binary_op(op)?, + Box::new(right), + ))) + } } diff --git a/datafusion/sql/src/expr/function.rs b/datafusion/sql/src/expr/function.rs index 67abb8b822063..e6bee31fbf106 100644 --- a/datafusion/sql/src/expr/function.rs +++ b/datafusion/sql/src/expr/function.rs @@ -370,7 +370,7 @@ impl SqlToRel<'_, S> { if let Some(fm) = self.context_provider.get_higher_order_meta(&name) { // plan non-lambda arguments first so we can get theirs datatype and call - // HigherOrderUDF::lambda_parameters to then plan the lambda arguments with + // HigherOrderUDFImpl::lambda_parameters to then plan the lambda arguments with // resolved lambda variables enum ExprOrLambda { Expr(Expr), @@ -546,15 +546,25 @@ impl SqlToRel<'_, S> { } } - // Build Unnest expression - if name.eq("unnest") { + // Build Unnest expression. + // + // `unnest(col)` drops `NULL` and empty input lists (default SQL + // semantics, matching DuckDB/PostgreSQL). `unnest_outer(col)` sets + // `outer = true` so the downstream planner picks + // `NullHandling::PreserveAndExpandEmpty`, which preserves `NULL` + // and empty input lists as a single `NULL` output row. + if name.eq("unnest") || name.eq("unnest_outer") { + let outer = name.eq("unnest_outer"); let mut exprs = self.function_args_to_expr(args, schema, planner_context)?; if exprs.len() != 1 { - return plan_err!("unnest() requires exactly one argument"); + return plan_err!("{name}() requires exactly one argument"); } let expr = exprs.swap_remove(0); Self::check_unnest_arg(&expr, schema)?; - return Ok(Expr::Unnest(Unnest::new(expr))); + return Ok(Expr::Unnest(Unnest { + expr: Box::new(expr), + outer, + })); } if !order_by.is_empty() && is_function_window { diff --git a/datafusion/sql/src/expr/mod.rs b/datafusion/sql/src/expr/mod.rs index ba7811acd8f3c..b1de4e95fd8a2 100644 --- a/datafusion/sql/src/expr/mod.rs +++ b/datafusion/sql/src/expr/mod.rs @@ -15,6 +15,8 @@ // specific language governing permissions and limitations // under the License. +use std::ops::ControlFlow; + use arrow::datatypes::{DataType, TimeUnit}; use datafusion_expr::planner::{ PlannerResult, RawBinaryExpr, RawDictionaryExpr, RawFieldAccessExpr, @@ -22,13 +24,14 @@ use datafusion_expr::planner::{ use sqlparser::ast::{ AccessExpr, BinaryOperator, CastFormat, CastKind, CeilFloorKind, DataType as SQLDataType, DateTimeField, DictionaryField, Expr as SQLExpr, - ExprWithAlias as SQLExprWithAlias, JsonPath, MapEntry, StructField, Subscript, - TrimWhereField, TypedString, Value, ValueWithSpan, + ExprWithAlias as SQLExprWithAlias, JsonPath, MapEntry, Spanned, StructField, + Subscript, TrimWhereField, TypedString, Value, ValueWithSpan, }; +use sqlparser::ast::{Query, Visit, Visitor}; use datafusion_common::{ - DFSchema, Result, ScalarValue, internal_datafusion_err, internal_err, not_impl_err, - plan_err, + DFSchema, Diagnostic, Result, ScalarValue, Span, internal_datafusion_err, + internal_err, not_impl_err, plan_err, }; use datafusion_expr::expr::ScalarFunction; @@ -54,7 +57,86 @@ mod substring; mod unary_op; mod value; +fn null_value_span(expr: &SQLExpr) -> Option> { + if let SQLExpr::Value(ValueWithSpan { + value: Value::Null, + span, + }) = expr + { + Some(Span::try_from_sqlparser_span(*span)) + } else { + None + } +} + +fn null_equality_warning(expr: &SQLExpr) -> Option { + let SQLExpr::BinaryOp { left, op, right } = expr else { + return None; + }; + + let null_span = null_value_span(left).or_else(|| null_value_span(right))?; + + let (message, help) = match op { + BinaryOperator::Eq => ( + "comparison with NULL using `=` always evaluates to NULL", + "use `IS NULL` to check for NULL values", + ), + BinaryOperator::NotEq => ( + "comparison with NULL using `<>` always evaluates to NULL", + "use `IS NOT NULL` to check for non-NULL values", + ), + _ => return None, + }; + + Some( + Diagnostic::new_warning(message, Span::try_from_sqlparser_span(expr.span())) + .with_help(help, null_span), + ) +} + +struct NullEqualityPredicateVisitor<'a, 'b, S: ContextProvider> { + sql_to_rel: &'a SqlToRel<'b, S>, + subquery_depth: usize, +} + +impl<'a, 'b, S: ContextProvider> NullEqualityPredicateVisitor<'a, 'b, S> { + fn new(sql_to_rel: &'a SqlToRel<'b, S>) -> Self { + Self { + sql_to_rel, + subquery_depth: 0, + } + } +} + +impl Visitor for NullEqualityPredicateVisitor<'_, '_, S> { + type Break = (); + + fn pre_visit_query(&mut self, _query: &Query) -> ControlFlow { + self.subquery_depth += 1; + ControlFlow::Continue(()) + } + + fn post_visit_query(&mut self, _query: &Query) -> ControlFlow { + self.subquery_depth -= 1; + ControlFlow::Continue(()) + } + + fn pre_visit_expr(&mut self, expr: &SQLExpr) -> ControlFlow { + if self.subquery_depth == 0 + && let Some(warning) = null_equality_warning(expr) + { + self.sql_to_rel.add_warning(warning); + } + ControlFlow::Continue(()) + } +} + impl SqlToRel<'_, S> { + pub(crate) fn warn_on_null_equality_predicate(&self, predicate: &SQLExpr) { + let mut visitor = NullEqualityPredicateVisitor::new(self); + let _ = predicate.visit(&mut visitor); + } + pub(crate) fn sql_expr_to_logical_expr_with_alias( &self, sql: SQLExprWithAlias, @@ -142,11 +224,7 @@ impl SqlToRel<'_, S> { } let RawBinaryExpr { op, left, right } = binary_expr; - Ok(Expr::BinaryExpr(BinaryExpr::new( - Box::new(left), - self.parse_sql_binary_op(&op)?, - Box::new(right), - ))) + self.build_binary_expr(&op, left, right) } pub fn sql_to_expr_with_alias( @@ -930,10 +1008,6 @@ impl SqlToRel<'_, S> { planner_context: &mut PlannerContext, ) -> Result { let pattern = self.sql_expr_to_logical_expr(pattern, schema, planner_context)?; - let pattern_type = pattern.get_type(schema)?; - if pattern_type != DataType::Utf8 && pattern_type != DataType::Null { - return plan_err!("Invalid pattern in SIMILAR TO expression"); - } let escape_char = match escape_char.map(|v| v.value) { Some(Value::SingleQuotedString(char)) if char.len() == 1 => { Some(char.chars().next().unwrap()) @@ -1415,7 +1489,7 @@ mod tests { None } - fn get_higher_order_meta(&self, _name: &str) -> Option> { + fn get_higher_order_meta(&self, _name: &str) -> Option> { None } @@ -1519,4 +1593,32 @@ mod tests { assert!(matches!(expr, Expr::Alias(_))); } + + #[test] + fn test_parse_numbers_with_underscores() { + use datafusion_common::ScalarValue::*; + + let context_provider = TestContextProvider::new(); + let sql_to_rel = SqlToRel::new(&context_provider); + + // (input, positive result, negative result) + let test_cases = [ + ("1_000", Int64(Some(1000)), Int64(Some(-1000))), + ("100_000", Int64(Some(100000)), Int64(Some(-100000))), + ("1_2_3_4", Int64(Some(1234)), Int64(Some(-1234))), + ("0_0", Int64(Some(0)), Int64(Some(-0))), + ("1_23.4_56", Float64(Some(123.456)), Float64(Some(-123.456))), + ]; + + for (literal, out_positive, out_negative) in test_cases { + assert_eq!( + sql_to_rel.parse_sql_number(literal, false).unwrap(), + Expr::Literal(out_positive, None) + ); + assert_eq!( + sql_to_rel.parse_sql_number(literal, true).unwrap(), + Expr::Literal(out_negative, None) + ); + } + } } diff --git a/datafusion/sql/src/expr/order_by.rs b/datafusion/sql/src/expr/order_by.rs index faecfbcfecc05..0067a1ebd708c 100644 --- a/datafusion/sql/src/expr/order_by.rs +++ b/datafusion/sql/src/expr/order_by.rs @@ -109,7 +109,13 @@ impl SqlToRel<'_, S> { )) } e => { - self.sql_expr_to_logical_expr(e, order_by_schema, planner_context)? + let expr = self.sql_expr_to_logical_expr( + e, + order_by_schema, + planner_context, + )?; + let (expr, _) = expr.infer_placeholder_types(order_by_schema)?; + expr } }; sort_expr_vec.push(make_sort_expr(expr, asc, nulls_first)); diff --git a/datafusion/sql/src/expr/unary_op.rs b/datafusion/sql/src/expr/unary_op.rs index cd118c0fdd5c5..b7683898fa5ba 100644 --- a/datafusion/sql/src/expr/unary_op.rs +++ b/datafusion/sql/src/expr/unary_op.rs @@ -15,11 +15,14 @@ // specific language governing permissions and limitations // under the License. +use arrow::datatypes::DataType; + use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; use datafusion_common::{DFSchema, Diagnostic, Result, not_impl_err, plan_err}; use datafusion_expr::{ - Expr, ExprSchemable, - type_coercion::{is_interval, is_timestamp}, + Expr, ExprSchemable, Operator, + binary::BinaryTypeCoercer, + type_coercion::{is_interval, is_signed_numeric, is_timestamp}, }; use sqlparser::ast::{Expr as SQLExpr, UnaryOperator, Value, ValueWithSpan}; @@ -32,9 +35,37 @@ impl SqlToRel<'_, S> { planner_context: &mut PlannerContext, ) -> Result { match op { - UnaryOperator::Not => Ok(Expr::Not(Box::new( - self.sql_expr_to_logical_expr(expr, schema, planner_context)?, - ))), + UnaryOperator::Not => { + let operand = + self.sql_expr_to_logical_expr(expr, schema, planner_context)?; + let field = operand.to_field(schema)?.1; + let data_type = field.data_type(); + let bool_coercible = BinaryTypeCoercer::new( + data_type, + &Operator::IsDistinctFrom, + &DataType::Boolean, + ) + .get_input_types() + .is_ok(); + if bool_coercible { + Ok(Expr::Not(Box::new(operand))) + } else { + let span = operand.spans().and_then(|s| s.first()); + let mut diagnostic = Diagnostic::new_error( + format!("NOT cannot be used with {data_type}"), + span, + ); + diagnostic + .add_note("NOT can only be used with boolean expressions", None); + diagnostic + .add_help(format!("perhaps you need to cast {operand}"), None); + plan_err!( + "Unary operator 'NOT' requires a boolean expression, \ + got {data_type}"; + diagnostic = diagnostic + ) + } + } UnaryOperator::Plus => { let operand = self.sql_expr_to_logical_expr(expr, schema, planner_context)?; @@ -72,11 +103,38 @@ impl SqlToRel<'_, S> { self.sql_interval_to_expr(true, interval) } // Not a literal, apply negative operator on expression - _ => Ok(Expr::Negative(Box::new(self.sql_expr_to_logical_expr( - expr, - schema, - planner_context, - )?))), + _ => { + let operand = + self.sql_expr_to_logical_expr(expr, schema, planner_context)?; + let field = operand.to_field(schema)?.1; + let data_type = field.data_type(); + if data_type.is_null() + || is_signed_numeric(data_type) + || is_interval(data_type) + || is_timestamp(data_type) + { + Ok(Expr::Negative(Box::new(operand))) + } else { + let span = operand.spans().and_then(|s| s.first()); + let mut diagnostic = Diagnostic::new_error( + format!("- cannot be used with {data_type}"), + span, + ); + diagnostic.add_note( + "- can only be used with signed numeric types, intervals, and timestamps", + None, + ); + diagnostic.add_help( + format!("perhaps you need to cast {operand}"), + None, + ); + plan_err!( + "Unary operator '-' only supports signed numeric, \ + interval and timestamp types"; + diagnostic = diagnostic + ) + } + } } } _ => not_impl_err!("Unsupported SQL unary operator {op:?}"), diff --git a/datafusion/sql/src/expr/value.rs b/datafusion/sql/src/expr/value.rs index 13a47f545cf7e..1307e917e4251 100644 --- a/datafusion/sql/src/expr/value.rs +++ b/datafusion/sql/src/expr/value.rs @@ -74,10 +74,27 @@ impl SqlToRel<'_, S> { unsigned_number: &str, negative: bool, ) -> Result { - let signed_number: Cow = if negative { - Cow::Owned(format!("-{unsigned_number}")) - } else { + // remove underscores, since the Rust parser used here does not support them + let signed_number = if !negative && !unsigned_number.contains('_') { Cow::Borrowed(unsigned_number) + } else { + let mut signed_number = + String::with_capacity(unsigned_number.len() + usize::from(negative)); + if negative { + signed_number.push('-'); + } + unsigned_number.bytes().for_each(|b| { + if b != b'_' { + signed_number.push(b as char); + } + }); + Cow::Owned(signed_number) + }; + + let unsigned_number = if negative { + &signed_number[1..] + } else { + &signed_number }; // Try to parse as i64 first, then u64 if negative is false, then decimal or f64 diff --git a/datafusion/sql/src/lib.rs b/datafusion/sql/src/lib.rs index 7fef670933f9a..95601acac2542 100644 --- a/datafusion/sql/src/lib.rs +++ b/datafusion/sql/src/lib.rs @@ -57,9 +57,4 @@ mod statement; pub mod unparser; pub mod utils; mod values; -#[deprecated( - since = "46.0.0", - note = "use datafusion_common::{ResolvedTableReference, TableReference}" -)] -pub use datafusion_common::{ResolvedTableReference, TableReference}; pub use sqlparser; diff --git a/datafusion/sql/src/parser.rs b/datafusion/sql/src/parser.rs index ba37a2d7026a3..86a00ca767a4c 100644 --- a/datafusion/sql/src/parser.rs +++ b/datafusion/sql/src/parser.rs @@ -21,7 +21,8 @@ //! `CREATE EXTERNAL TABLE` use datafusion_common::DataFusionError; -use datafusion_common::config::SqlParserOptions; +use datafusion_common::config::{ConfigNonZeroUsize, SqlParserOptions}; +use datafusion_common::format::{ExplainFormat, ExplainStatementOptions}; use datafusion_common::{Diagnostic, Span, sql_err}; use sqlparser::ast::{ExprWithAlias, Ident, OrderByOptions}; use sqlparser::tokenizer::TokenWithSpan; @@ -36,6 +37,7 @@ use sqlparser::{ }; use std::collections::VecDeque; use std::fmt; +use std::str::FromStr; // Use `Parser::expected` instead, if possible macro_rules! parser_err { @@ -55,18 +57,25 @@ fn parse_file_type(s: &str) -> Result { /// DataFusion specific `EXPLAIN` /// -/// Syntax: +/// Supports both the legacy keyword form and, on dialects whose +/// [`Dialect::supports_explain_with_utility_options`] returns `true` +/// (PostgreSQL, DuckDB, etc.), the Postgres-style parenthesized option list: +/// /// ```sql +/// -- Legacy keyword form (any dialect) /// EXPLAIN [FORMAT format] statement +/// +/// -- Postgres-style option form (dialect-gated) +/// EXPLAIN (option [arg] [, ...]) statement /// ``` +/// +/// See [`ExplainStatementOptions`] for the list of supported options in the +/// parenthesized form. #[derive(Debug, Clone, PartialEq, Eq)] pub struct ExplainStatement { - /// `EXPLAIN ANALYZE ..` - pub analyze: bool, - /// `EXPLAIN .. VERBOSE ..` - pub verbose: bool, - /// `EXPLAIN .. FORMAT ` - pub format: Option, + /// Normalized options parsed from either the legacy keyword form or the + /// parenthesized option list. + pub options: ExplainStatementOptions, /// The statement to analyze. Note this is a DataFusion [`Statement`] (not a /// [`sqlparser::ast::Statement`] so that we can use `EXPLAIN`, `COPY`, and other /// DataFusion specific statements @@ -75,22 +84,47 @@ pub struct ExplainStatement { impl fmt::Display for ExplainStatement { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - let Self { - analyze, - verbose, - format, - statement, - } = self; + let Self { options, statement } = self; + + // If only the legacy-era fields are set, print the legacy keyword + // form so existing round-trip tests continue to pass. + let uses_parenthesized = options.analyze_level.is_some() + || options.analyze_categories.is_some() + || options.show_statistics.is_some(); write!(f, "EXPLAIN ")?; - if *analyze { - write!(f, "ANALYZE ")?; - } - if *verbose { - write!(f, "VERBOSE ")?; - } - if let Some(format) = format.as_ref() { - write!(f, "FORMAT {format} ")?; + if uses_parenthesized { + // Emit a parenthesized option list. + let mut parts: Vec = Vec::new(); + if options.analyze { + parts.push("ANALYZE".to_string()); + } + if options.verbose { + parts.push("VERBOSE".to_string()); + } + if let Some(format) = &options.format { + parts.push(format!("FORMAT {format}")); + } + if let Some(level) = options.analyze_level { + parts.push(format!("LEVEL {level}")); + } + if let Some(cats) = &options.analyze_categories { + parts.push(format!("METRICS '{cats}'")); + } + if let Some(stats) = options.show_statistics { + parts.push(format!("COSTS {}", if stats { "ON" } else { "OFF" })); + } + write!(f, "({}) ", parts.join(", "))?; + } else { + if options.analyze { + write!(f, "ANALYZE ")?; + } + if options.verbose { + write!(f, "VERBOSE ")?; + } + if let Some(format) = &options.format { + write!(f, "FORMAT {format} ")?; + } } write!(f, "{statement}") @@ -197,7 +231,7 @@ pub(crate) type LexOrdering = Vec; /// [ PARTITIONED BY ( | ) ] /// [ WITH ORDER () /// [ OPTIONS () ] -/// LOCATION +/// LOCATION | LOCATION ([, ...]) /// /// := ( , ...) /// @@ -215,8 +249,8 @@ pub struct CreateExternalTable { pub columns: Vec, /// File type (Parquet, NDJSON, CSV, etc) pub file_type: String, - /// Path to file - pub location: String, + /// Paths to files + pub locations: Vec, /// Partition Columns pub table_partition_cols: Vec, /// Ordered expressions @@ -255,7 +289,23 @@ impl fmt::Display for CreateExternalTable { } write!(f, ") ")?; } - write!(f, "LOCATION {}", self.location) + match self.locations.as_slice() { + [location] => write!( + f, + "LOCATION {}", + Value::SingleQuotedString(location.clone()) + ), + locations => { + write!(f, "LOCATION (")?; + for (idx, location) in locations.iter().enumerate() { + if idx > 0 { + write!(f, ", ")?; + } + write!(f, "{}", Value::SingleQuotedString(location.clone()))?; + } + write!(f, ")") + } + } } } @@ -325,6 +375,10 @@ fn ensure_not_set(field: &Option, name: &str) -> Result<(), DataFusionErro pub struct DFParser<'a> { pub parser: Parser<'a>, options: SqlParserOptions, + /// Whether the configured dialect supports Postgres-style + /// `EXPLAIN (option, ...)` utility-option syntax. Cached here because + /// sqlparser's [`Parser::dialect`] field is private. + supports_explain_with_utility_options: bool, } /// Same as `sqlparser` @@ -434,27 +488,38 @@ impl<'a, 'b> DFParserBuilder<'a, 'b> { .with_tokens_with_locations(tokens) .with_recursion_limit(self.recursion_limit), options: SqlParserOptions { - recursion_limit: self.recursion_limit, + recursion_limit: ConfigNonZeroUsize::try_new(self.recursion_limit)?, ..Default::default() }, + supports_explain_with_utility_options: self + .dialect + .supports_explain_with_utility_options(), }) } } -impl<'a> DFParser<'a> { - #[deprecated(since = "46.0.0", note = "DFParserBuilder")] - pub fn new(sql: &'a str) -> Result { - DFParserBuilder::new(sql).build() - } - - #[deprecated(since = "46.0.0", note = "DFParserBuilder")] - pub fn new_with_dialect( - sql: &'a str, - dialect: &'a dyn Dialect, - ) -> Result { - DFParserBuilder::new(sql).with_dialect(dialect).build() +/// Returns true when `tok` is the start of a query / parenthesized query +/// group. Used to disambiguate `EXPLAIN (SELECT ...)` (a parenthesized query) +/// from `EXPLAIN (ANALYZE) SELECT ...` (a Postgres-style option list). +fn token_starts_query(tok: &Token) -> bool { + match tok { + Token::LParen => true, + Token::Word(Word { keyword, .. }) => matches!( + keyword, + Keyword::SELECT + | Keyword::WITH + | Keyword::VALUES + | Keyword::TABLE + | Keyword::INSERT + | Keyword::UPDATE + | Keyword::DELETE + | Keyword::MERGE + ), + _ => false, } +} +impl<'a> DFParser<'a> { /// Parse a sql string into one or [`Statement`]s using the /// [`GenericDialect`]. pub fn parse_sql(sql: &'a str) -> Result, DataFusionError> { @@ -758,18 +823,46 @@ impl<'a> DFParser<'a> { } /// Parse a SQL `EXPLAIN` + /// + /// After the `EXPLAIN` keyword, if the dialect supports the Postgres-style + /// option list and the next non-whitespace token is `(`, we must + /// disambiguate between an option list (`EXPLAIN (ANALYZE) SELECT ...`) + /// and a parenthesized query (`EXPLAIN (SELECT ...)` or + /// `EXPLAIN (q1 EXCEPT q2) UNION ALL ...`). pub fn parse_explain(&mut self) -> Result { + if self.supports_explain_with_utility_options + && self.parser.peek_token().token == Token::LParen + && !token_starts_query(&self.parser.peek_nth_token(1).token) + { + let raw = self.parser.parse_utility_options()?; + let options = ExplainStatementOptions::from_utility_options(&raw)?; + let statement = self.parse_statement()?; + return Ok(Statement::Explain(ExplainStatement { + statement: Box::new(statement), + options, + })); + } + + // Legacy keyword form. let analyze = self.parser.parse_keyword(Keyword::ANALYZE); let verbose = self.parser.parse_keyword(Keyword::VERBOSE); - let format = self.parse_explain_format()?; + let format = self + .parse_explain_format()? + .map(|s| ExplainFormat::from_str(&s)) + .transpose()?; let statement = self.parse_statement()?; - Ok(Statement::Explain(ExplainStatement { - statement: Box::new(statement), + let options = ExplainStatementOptions { analyze, verbose, format, + ..Default::default() + }; + + Ok(Statement::Explain(ExplainStatement { + statement: Box::new(statement), + options, })) } @@ -1020,7 +1113,7 @@ impl<'a> DFParser<'a> { #[derive(Default)] struct Builder { file_type: Option, - location: Option, + locations: Option>, table_partition_cols: Option>, order_exprs: Vec, options: Option>, @@ -1044,8 +1137,8 @@ impl<'a> DFParser<'a> { builder.file_type = Some(self.parse_file_format()?); } Keyword::LOCATION => { - ensure_not_set(&builder.location, "LOCATION")?; - builder.location = Some(self.parser.parse_literal_string()?); + ensure_not_set(&builder.locations, "LOCATION")?; + builder.locations = Some(self.parse_locations()?); } Keyword::WITH => { if self.parser.parse_keyword(Keyword::ORDER) { @@ -1121,17 +1214,22 @@ impl<'a> DFParser<'a> { "Missing STORED AS clause in CREATE EXTERNAL TABLE statement".into(), )); } - if builder.location.is_none() { + if builder.locations.is_none() { return sql_err!(ParserError::ParserError( "Missing LOCATION clause in CREATE EXTERNAL TABLE statement".into(), )); } + let locations = builder.locations.unwrap(); + if locations.is_empty() { + return parser_err!("LOCATION requires at least one path"); + } + let create = CreateExternalTable { name: table_name, columns, file_type: builder.file_type.unwrap(), - location: builder.location.unwrap(), + locations, table_partition_cols: builder.table_partition_cols.unwrap_or(vec![]), order_exprs: builder.order_exprs, if_not_exists, @@ -1144,6 +1242,29 @@ impl<'a> DFParser<'a> { Ok(Statement::CreateExternalTable(create)) } + /// Parses one or more external table locations. + fn parse_locations(&mut self) -> Result, DataFusionError> { + if !self.parser.consume_token(&Token::LParen) { + return Ok(vec![self.parser.parse_literal_string()?]); + } + + let mut locations = vec![]; + loop { + locations.push(self.parser.parse_literal_string()?); + let comma = self.parser.consume_token(&Token::Comma); + if self.parser.consume_token(&Token::RParen) { + // Allow a trailing comma, even though it's not in standard + break; + } else if !comma { + return self.expected( + "',' or ')' after location definition", + &self.parser.peek_token(), + ); + } + } + Ok(locations) + } + /// Parses the set of valid formats fn parse_file_format(&mut self) -> Result { let token = self.parser.next_token(); @@ -1232,17 +1353,23 @@ mod tests { } } - #[test] - fn create_external_table() -> Result<(), DataFusionError> { - // positive case - let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv'"; - let display = None; - let name = ObjectName::from(vec![Ident::from("t")]); - let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), - columns: vec![make_column_def("c1", DataType::Int(display))], + fn make_create_external_table(location: &str) -> CreateExternalTable { + make_create_external_table_with_locations(&[location]) + } + + fn make_create_external_table_with_locations( + locations: &[&str], + ) -> CreateExternalTable { + let locations = locations + .iter() + .map(|location| location.to_string()) + .collect::>(); + + CreateExternalTable { + name: ObjectName::from(vec![Ident::from("t")]), + columns: vec![], file_type: "CSV".to_string(), - location: "foo.csv".into(), + locations, table_partition_cols: vec![], order_exprs: vec![], if_not_exists: false, @@ -1251,24 +1378,59 @@ mod tests { unbounded: false, options: vec![], constraints: vec![], + } + } + + #[test] + fn create_external_table() -> Result<(), DataFusionError> { + // positive case + let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv'"; + let display = None; + let expected = Statement::CreateExternalTable(CreateExternalTable { + columns: vec![make_column_def("c1", DataType::Int(display))], + ..make_create_external_table("foo.csv") + }); + expect_parse_ok(sql, expected)?; + + // positive case: literal comma remains part of a single path + let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo,bar.csv'"; + let expected = Statement::CreateExternalTable(CreateExternalTable { + columns: vec![make_column_def("c1", DataType::Int(display))], + ..make_create_external_table("foo,bar.csv") }); expect_parse_ok(sql, expected)?; + // positive case: multiple locations use an explicit list + let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION ('foo.csv', 'bar.csv')"; + let expected = Statement::CreateExternalTable(CreateExternalTable { + columns: vec![make_column_def("c1", DataType::Int(display))], + ..make_create_external_table_with_locations(&["foo.csv", "bar.csv"]) + }); + expect_parse_ok(sql, expected)?; + + assert_eq!( + Statement::CreateExternalTable(make_create_external_table("foo.csv")) + .to_string(), + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION 'foo.csv'" + ); + assert_eq!( + Statement::CreateExternalTable(make_create_external_table_with_locations(&[ + "foo.csv", "bar.csv" + ])) + .to_string(), + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION ('foo.csv', 'bar.csv')" + ); + assert_eq!( + Statement::CreateExternalTable(make_create_external_table("foo'bar.csv")) + .to_string(), + "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION 'foo''bar.csv'" + ); + // positive case: leading space let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv' "; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![make_column_def("c1", DataType::Int(None))], - file_type: "CSV".to_string(), - location: "foo.csv".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; @@ -1276,18 +1438,8 @@ mod tests { let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv' ;"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![make_column_def("c1", DataType::Int(None))], - file_type: "CSV".to_string(), - location: "foo.csv".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; @@ -1295,21 +1447,12 @@ mod tests { let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv' OPTIONS (format.delimiter '|')"; let display = None; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![make_column_def("c1", DataType::Int(display))], - file_type: "CSV".to_string(), - location: "foo.csv".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, options: vec![( "format.delimiter".into(), Value::SingleQuotedString("|".into()), )], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; @@ -1317,18 +1460,9 @@ mod tests { let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV PARTITIONED BY (p1, p2) LOCATION 'foo.csv'"; let display = None; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![make_column_def("c1", DataType::Int(display))], - file_type: "CSV".to_string(), - location: "foo.csv".into(), table_partition_cols: vec!["p1".to_string(), "p2".to_string()], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; @@ -1343,24 +1477,15 @@ mod tests { ('format.compression' 'XZ')", "XZ"), ("CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV LOCATION 'foo.csv' OPTIONS ('format.compression' 'ZSTD')", "ZSTD"), - ]; + ]; for (sql, compression) in sqls { let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![make_column_def("c1", DataType::Int(display))], - file_type: "CSV".to_string(), - location: "foo.csv".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, options: vec![( "format.compression".into(), Value::SingleQuotedString(compression.into()), )], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; } @@ -1368,72 +1493,33 @@ mod tests { // positive case: it is ok for parquet files not to have columns specified let sql = "CREATE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'foo.parquet'"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), - columns: vec![], file_type: "PARQUET".to_string(), - location: "foo.parquet".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.parquet") }); expect_parse_ok(sql, expected)?; // positive case: it is ok for parquet files to be other than upper case let sql = "CREATE EXTERNAL TABLE t STORED AS parqueT LOCATION 'foo.parquet'"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), - columns: vec![], file_type: "PARQUET".to_string(), - location: "foo.parquet".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.parquet") }); expect_parse_ok(sql, expected)?; // positive case: it is ok for avro files not to have columns specified let sql = "CREATE EXTERNAL TABLE t STORED AS AVRO LOCATION 'foo.avro'"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), - columns: vec![], file_type: "AVRO".to_string(), - location: "foo.avro".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.avro") }); expect_parse_ok(sql, expected)?; // positive case: it is ok for avro files not to have columns specified let sql = "CREATE EXTERNAL TABLE IF NOT EXISTS t STORED AS PARQUET LOCATION 'foo.parquet'"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), - columns: vec![], file_type: "PARQUET".to_string(), - location: "foo.parquet".into(), - table_partition_cols: vec![], - order_exprs: vec![], if_not_exists: true, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.parquet") }); expect_parse_ok(sql, expected)?; @@ -1441,39 +1527,21 @@ mod tests { let sql = "CREATE OR REPLACE EXTERNAL TABLE t STORED AS PARQUET LOCATION 'foo.parquet'"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), - columns: vec![], file_type: "PARQUET".to_string(), - location: "foo.parquet".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, or_replace: true, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.parquet") }); expect_parse_ok(sql, expected)?; // positive case: column definition allowed in 'partition by' clause let sql = "CREATE EXTERNAL TABLE t(c1 int) STORED AS CSV PARTITIONED BY (p1 int) LOCATION 'foo.csv'"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![ make_column_def("c1", DataType::Int(None)), make_column_def("p1", DataType::Int(None)), ], - file_type: "CSV".to_string(), - location: "foo.csv".into(), table_partition_cols: vec!["p1".to_string()], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; @@ -1495,39 +1563,21 @@ mod tests { let sql = "CREATE EXTERNAL TABLE t STORED AS x OPTIONS ('k1' 'v1') LOCATION 'blahblah'"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), - columns: vec![], file_type: "X".to_string(), - location: "blahblah".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, options: vec![("k1".into(), Value::SingleQuotedString("v1".into()))], - constraints: vec![], + ..make_create_external_table("blahblah") }); expect_parse_ok(sql, expected)?; // positive case: additional options (multiple entries) can be specified let sql = "CREATE EXTERNAL TABLE t STORED AS x OPTIONS ('k1' 'v1', k2 v2) LOCATION 'blahblah'"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), - columns: vec![], file_type: "X".to_string(), - location: "blahblah".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, options: vec![ ("k1".into(), Value::SingleQuotedString("v1".into())), ("k2".into(), Value::SingleQuotedString("v2".into())), ], - constraints: vec![], + ..make_create_external_table("blahblah") }); expect_parse_ok(sql, expected)?; @@ -1556,11 +1606,7 @@ mod tests { ]; for (sql, (asc, nulls_first)) in sqls.iter().zip(expected) { let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![make_column_def("c1", DataType::Int(None))], - file_type: "CSV".to_string(), - location: "foo.csv".into(), - table_partition_cols: vec![], order_exprs: vec![vec![OrderByExpr { expr: Identifier(Ident { value: "c1".to_owned(), @@ -1570,12 +1616,7 @@ mod tests { options: OrderByOptions { asc, nulls_first }, with_fill: None, }]], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; } @@ -1584,14 +1625,10 @@ mod tests { let sql = "CREATE EXTERNAL TABLE t(c1 int, c2 int) STORED AS CSV WITH ORDER (c1 ASC, c2 DESC NULLS FIRST) LOCATION 'foo.csv'"; let display = None; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![ make_column_def("c1", DataType::Int(display)), make_column_def("c2", DataType::Int(display)), ], - file_type: "CSV".to_string(), - location: "foo.csv".into(), - table_partition_cols: vec![], order_exprs: vec![vec![ OrderByExpr { expr: Identifier(Ident { @@ -1618,12 +1655,7 @@ mod tests { with_fill: None, }, ]], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; @@ -1631,14 +1663,10 @@ mod tests { let sql = "CREATE EXTERNAL TABLE t(c1 int, c2 int) STORED AS CSV WITH ORDER (c1 - c2 ASC) LOCATION 'foo.csv'"; let display = None; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![ make_column_def("c1", DataType::Int(display)), make_column_def("c2", DataType::Int(display)), ], - file_type: "CSV".to_string(), - location: "foo.csv".into(), - table_partition_cols: vec![], order_exprs: vec![vec![OrderByExpr { expr: Expr::BinaryOp { left: Box::new(Identifier(Ident { @@ -1659,12 +1687,7 @@ mod tests { }, with_fill: None, }]], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.csv") }); expect_parse_ok(sql, expected)?; @@ -1681,13 +1704,11 @@ mod tests { 'TRUNCATE' 'NO', 'format.has_header' 'true')"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![ make_column_def("c1", DataType::Int(None)), make_column_def("c2", DataType::Float(ExactNumberInfo::None)), ], file_type: "PARQUET".to_string(), - location: "foo.parquet".into(), table_partition_cols: vec!["c1".into()], order_exprs: vec![vec![OrderByExpr { expr: Expr::BinaryOp { @@ -1710,8 +1731,6 @@ mod tests { with_fill: None, }]], if_not_exists: true, - or_replace: false, - temporary: false, unbounded: true, options: vec![ ( @@ -1732,7 +1751,7 @@ mod tests { Value::SingleQuotedString("true".into()), ), ], - constraints: vec![], + ..make_create_external_table("foo.parquet") }); expect_parse_ok(sql, expected)?; @@ -1749,13 +1768,11 @@ mod tests { 'TRUNCATE' 'NO', 'format.has_header' 'true')"; let expected = Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![ make_column_def("c1", DataType::Int(None)), make_column_def("c2", DataType::Float(ExactNumberInfo::None)), ], file_type: "PARQUET".to_string(), - location: "foo.parquet".into(), table_partition_cols: vec!["c1".into()], order_exprs: vec![vec![OrderByExpr { expr: Expr::BinaryOp { @@ -1777,9 +1794,7 @@ mod tests { }, with_fill: None, }]], - if_not_exists: false, or_replace: true, - temporary: false, unbounded: true, options: vec![ ( @@ -1800,7 +1815,7 @@ mod tests { Value::SingleQuotedString("true".into()), ), ], - constraints: vec![], + ..make_create_external_table("foo.parquet") }); expect_parse_ok(sql, expected)?; @@ -1873,9 +1888,14 @@ mod tests { options: vec![], }); let expected = Statement::Explain(ExplainStatement { - analyze, - verbose, - format: None, + options: ExplainStatementOptions { + analyze, + verbose, + format: None, + analyze_level: None, + analyze_categories: None, + show_statistics: None, + }, statement: Box::new(expected_copy), }); assert_eq!(verified_stmt(sql), expected); @@ -2070,21 +2090,10 @@ mod tests { options: vec![], }), { - let name = ObjectName::from(vec![Ident::from("t")]); let display = None; Statement::CreateExternalTable(CreateExternalTable { - name: name.clone(), columns: vec![make_column_def("c1", DataType::Int(display))], - file_type: "CSV".to_string(), - location: "foo.csv".into(), - table_partition_cols: vec![], - order_exprs: vec![], - if_not_exists: false, - or_replace: false, - temporary: false, - unbounded: false, - options: vec![], - constraints: vec![], + ..make_create_external_table("foo.csv") }) }, { @@ -2203,4 +2212,164 @@ mod tests { "Expected: end of expression, found: bar", ) } + + // ------------------------------------------------------------------ + // Postgres-style `EXPLAIN (option, ...)` tests + // ------------------------------------------------------------------ + + fn parse_with_pg(sql: &str) -> Result { + let dialect = sqlparser::dialect::PostgreSqlDialect {}; + let mut statements = DFParser::parse_sql_with_dialect(sql, &dialect)?; + assert_eq!(statements.len(), 1, "Expected exactly one statement"); + Ok(statements.pop_front().unwrap()) + } + + fn parse_with_generic(sql: &str) -> Result { + let mut statements = DFParser::parse_sql(sql)?; + assert_eq!(statements.len(), 1, "Expected exactly one statement"); + Ok(statements.pop_front().unwrap()) + } + + #[test] + fn explain_legacy_keyword_form_postgres_dialect() { + // The legacy keyword form still works under PostgreSQL dialect. + let stmt = parse_with_pg("EXPLAIN ANALYZE VERBOSE SELECT 1").unwrap(); + let Statement::Explain(ExplainStatement { options, .. }) = stmt else { + panic!("Expected Statement::Explain"); + }; + assert!(options.analyze); + assert!(options.verbose); + assert!(options.format.is_none()); + assert!(options.analyze_level.is_none()); + } + + #[test] + fn explain_paren_form_on_generic_supports_utility_options() { + // sqlparser's GenericDialect also declares + // `supports_explain_with_utility_options = true`, so DataFusion's + // default parser accepts the parenthesized form too. + let stmt = parse_with_generic("EXPLAIN (FORMAT TREE) SELECT 1").unwrap(); + let Statement::Explain(ExplainStatement { options, .. }) = stmt else { + panic!("Expected Statement::Explain"); + }; + assert_eq!(options.format, Some(ExplainFormat::Tree)); + } + + #[test] + fn explain_paren_form_on_non_supporting_dialect_is_parse_error() { + // Dialects that do NOT declare support for utility options (e.g. + // Snowflake) must still error on the parenthesized form — proving + // the dialect gate itself works. + use sqlparser::dialect::SnowflakeDialect; + let dialect = SnowflakeDialect {}; + let res = + DFParser::parse_sql_with_dialect("EXPLAIN (FORMAT TREE) SELECT 1", &dialect); + assert!( + res.is_err(), + "expected parse error under non-supporting dialect" + ); + } + + #[test] + fn explain_paren_grouping_query_is_not_mistaken_for_options() { + // Historic DataFusion behavior allows parentheses around the + // query after EXPLAIN (e.g. `EXPLAIN (SELECT ...)` or + // `EXPLAIN (q1 EXCEPT q2) UNION ALL (q3 EXCEPT q4)`). The dialect + // gate for Postgres-style options must not swallow these. + for sql in [ + "EXPLAIN (SELECT 1)", + "EXPLAIN (WITH t AS (SELECT 1) SELECT * FROM t)", + "EXPLAIN (VALUES (1), (2))", + "EXPLAIN ((SELECT 1))", + ] { + let stmt = parse_with_pg(sql).unwrap_or_else(|e| { + panic!("{sql} failed under PG dialect: {e}"); + }); + let Statement::Explain(ExplainStatement { options, .. }) = stmt else { + panic!("Expected Statement::Explain for {sql}"); + }; + assert!(!options.analyze, "{sql} should not be ANALYZE"); + assert!(!options.verbose, "{sql} should not be VERBOSE"); + assert!(options.format.is_none(), "{sql} should have no FORMAT"); + } + } + + #[test] + fn explain_paren_form_analyze_verbose() { + let stmt = parse_with_pg("EXPLAIN (ANALYZE, VERBOSE) SELECT 1").unwrap(); + let Statement::Explain(ExplainStatement { options, .. }) = stmt else { + panic!("Expected Statement::Explain"); + }; + assert!(options.analyze); + assert!(options.verbose); + } + + #[test] + fn explain_paren_form_format_tree() { + let stmt = parse_with_pg("EXPLAIN (FORMAT tree) SELECT 1").unwrap(); + let Statement::Explain(ExplainStatement { options, .. }) = stmt else { + panic!("Expected Statement::Explain"); + }; + assert!(!options.analyze); + assert_eq!(options.format, Some(ExplainFormat::Tree)); + } + + #[test] + fn explain_paren_form_metrics_level() { + use datafusion_common::format::{ + ExplainAnalyzeCategories, MetricCategory, MetricType, + }; + let stmt = + parse_with_pg("EXPLAIN (ANALYZE, METRICS 'rows,bytes', LEVEL dev) SELECT 1") + .unwrap(); + let Statement::Explain(ExplainStatement { options, .. }) = stmt else { + panic!("Expected Statement::Explain"); + }; + assert!(options.analyze); + assert_eq!(options.analyze_level, Some(MetricType::Dev)); + assert_eq!( + options.analyze_categories, + Some(ExplainAnalyzeCategories::Only(vec![ + MetricCategory::Rows, + MetricCategory::Bytes, + ])) + ); + } + + #[test] + fn explain_paren_form_bool_spellings() { + let stmt = + parse_with_pg("EXPLAIN (ANALYZE ON, VERBOSE OFF, COSTS TRUE) SELECT 1") + .unwrap(); + let Statement::Explain(ExplainStatement { options, .. }) = stmt else { + panic!("Expected Statement::Explain"); + }; + assert!(options.analyze); + assert!(!options.verbose); + assert_eq!(options.show_statistics, Some(true)); + } + + #[test] + fn explain_paren_form_buffers_rejected() { + let err = parse_with_pg("EXPLAIN (BUFFERS) SELECT 1").unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("BUFFERS"), + "error should mention BUFFERS: {msg}" + ); + assert!( + msg.contains("not supported"), + "error should say not supported: {msg}" + ); + } + + #[test] + fn explain_paren_form_unknown_option_rejected() { + let err = parse_with_pg("EXPLAIN (ASDF) SELECT 1").unwrap_err(); + let msg = err.to_string(); + assert!( + msg.contains("unknown EXPLAIN option"), + "error should describe unknown option: {msg}" + ); + } } diff --git a/datafusion/sql/src/planner.rs b/datafusion/sql/src/planner.rs index 01215ae3434cf..3a696811be499 100644 --- a/datafusion/sql/src/planner.rs +++ b/datafusion/sql/src/planner.rs @@ -18,7 +18,7 @@ //! [`SqlToRel`]: SQL Query Planner (produces [`LogicalPlan`] from SQL AST) use std::collections::HashMap; use std::str::FromStr; -use std::sync::Arc; +use std::sync::{Arc, Mutex}; use std::vec; use crate::utils::make_decimal_type; @@ -32,10 +32,10 @@ use datafusion_common::{ DFSchemaRef, Diagnostic, SchemaError, field_not_found, internal_err, plan_datafusion_err, }; +use datafusion_expr::Expr; use datafusion_expr::logical_plan::{LogicalPlan, LogicalPlanBuilder}; pub use datafusion_expr::planner::ContextProvider; use datafusion_expr::utils::find_column_exprs; -use datafusion_expr::{Expr, col}; use sqlparser::ast::{ArrayElemTypeDef, ExactNumberInfo, TimezoneInfo}; use sqlparser::ast::{ColumnDef as SQLColumnDef, ColumnOption}; use sqlparser::ast::{DataType as SQLDataType, Ident, ObjectName, TableAlias}; @@ -455,6 +455,7 @@ pub struct SqlToRel<'a, S: ContextProvider> { pub(crate) context_provider: &'a S, pub(crate) options: ParserOptions, pub(crate) ident_normalizer: IdentNormalizer, + warnings: Mutex>, } impl<'a, S: ContextProvider> SqlToRel<'a, S> { @@ -477,9 +478,27 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { context_provider, options, ident_normalizer: IdentNormalizer::new(ident_normalize), + warnings: Mutex::new(vec![]), } } + pub(crate) fn add_warning(&self, warning: Diagnostic) { + self.warnings + .lock() + .expect("warning diagnostic lock poisoned") + .push(warning); + } + + /// Drain and return non-fatal warnings collected during SQL planning. + pub fn take_warnings(&self) -> Vec { + std::mem::take( + &mut self + .warnings + .lock() + .expect("warning diagnostic lock poisoned"), + ) + } + pub fn build_schema(&self, columns: Vec) -> Result { let mut fields = Vec::with_capacity(columns.len()); @@ -572,10 +591,10 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { idents.len() ) } else { - let fields = plan.schema().fields().clone(); + let columns = plan.schema().columns(); LogicalPlanBuilder::from(plan) - .project(fields.iter().zip(idents).map(|(field, ident)| { - col(field.name()).alias(self.ident_normalizer.normalize(ident)) + .project(columns.into_iter().zip(idents).map(|(col, ident)| { + Expr::Column(col).alias(self.ident_normalizer.normalize(ident)) }))? .build() } @@ -622,13 +641,13 @@ impl<'a, S: ContextProvider> SqlToRel<'a, S> { Diagnostic::new_error( format!( "column '{}' not found in '{}'", - &col.name, relation + col.name, relation ), col.spans().first(), ) } else { Diagnostic::new_error( - format!("column '{}' not found", &col.name), + format!("column '{}' not found", col.name), col.spans().first(), ) }; diff --git a/datafusion/sql/src/query.rs b/datafusion/sql/src/query.rs index 76124cbc7eb59..e2b9e4d2d5305 100644 --- a/datafusion/sql/src/query.rs +++ b/datafusion/sql/src/query.rs @@ -19,7 +19,6 @@ use std::sync::Arc; use crate::planner::{ContextProvider, PlannerContext, SqlToRel}; -use crate::stack::StackGuard; use datafusion_common::{Constraints, DFSchema, Result, not_impl_err}; use datafusion_expr::expr::{Sort, WildcardOptions}; @@ -81,11 +80,9 @@ impl SqlToRel<'_, S> { // The functions called from `set_expr_to_plan()` need more than 128KB // stack in debug builds as investigated in: // https://github.com/apache/datafusion/pull/13310#discussion_r1836813902 - let plan = { - // scope for dropping _guard - let _guard = StackGuard::new(256 * 1024); + let plan = crate::stack::maybe_grow(|| { self.set_expr_to_plan(other, planner_context) - }?; + })?; let oby_exprs = to_order_by_exprs(order_by)?; let order_by_rex = self.order_by_to_sort_expr( oby_exprs, diff --git a/datafusion/sql/src/relation/join.rs b/datafusion/sql/src/relation/join.rs index 3343890c6dc1d..475d9a5b38099 100644 --- a/datafusion/sql/src/relation/join.rs +++ b/datafusion/sql/src/relation/join.rs @@ -122,6 +122,7 @@ impl SqlToRel<'_, S> { JoinConstraint::On(sql_expr) => { let join_schema = left.schema().join(right.schema())?; // parse ON expression + self.warn_on_null_equality_predicate(&sql_expr); let expr = self.sql_to_expr(sql_expr, &join_schema, planner_context)?; LogicalPlanBuilder::from(left) .join_on(right, join_type, Some(expr))? diff --git a/datafusion/sql/src/resolve.rs b/datafusion/sql/src/resolve.rs index 955dbb86602a3..d1c172502ff11 100644 --- a/datafusion/sql/src/resolve.rs +++ b/datafusion/sql/src/resolve.rs @@ -20,9 +20,9 @@ use std::ops::ControlFlow; use datafusion_common::{DataFusionError, Result}; -use crate::TableReference; use crate::parser::{CopyToSource, CopyToStatement, Statement as DFStatement}; use crate::planner::object_name_to_table_reference; +use datafusion_common::TableReference; use sqlparser::ast::*; // following constants are used in `resolve_table_references` diff --git a/datafusion/sql/src/select.rs b/datafusion/sql/src/select.rs index b7f7d80e70815..bbd9d203eb124 100644 --- a/datafusion/sql/src/select.rs +++ b/datafusion/sql/src/select.rs @@ -25,16 +25,18 @@ use crate::utils::{ CheckColumnsMustReferenceAggregatePurpose, CheckColumnsSatisfyExprsPurpose, check_columns_satisfy_exprs, extract_aliases, rebase_expr, resolve_aliases_to_exprs, resolve_columns, resolve_positions_to_exprs, rewrite_recursive_unnest_bottom_up, - rewrite_recursive_unnests_bottom_up, + rewrite_recursive_unnests_bottom_up, substitute_top_level_alias, + substitute_top_level_aliases_in_sorts, }; use arrow::datatypes::DataType; use datafusion_common::error::DataFusionErrorBuilder; use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; use datafusion_common::{Column, DFSchema, DFSchemaRef, Result, not_impl_err, plan_err}; -use datafusion_common::{RecursionUnnestOption, UnnestOptions}; +use datafusion_common::{NullHandling, RecursionUnnestOption, UnnestOptions}; use datafusion_expr::ExprSchemable; use datafusion_expr::builder::get_struct_unnested_columns; +use datafusion_expr::expr::Unnest as UnnestExpr; use datafusion_expr::expr::{PlannedReplaceSelectItem, WildcardOptions}; use datafusion_expr::expr_rewriter::{ normalize_col, normalize_col_with_schemas_and_ambiguity_check, normalize_sorts, @@ -69,6 +71,24 @@ struct AggregatePlanResult { qualify_expr: Option, /// ORDER BY expressions rewritten to reference aggregate output columns order_by_exprs: Vec, + /// DISTINCT ON expressions rewritten to reference aggregate output columns + on_exprs: Vec, +} + +struct DistinctOnUnnestPlanResult { + plan: LogicalPlan, + select_exprs: Vec, + on_exprs: Vec, + order_by_exprs: Vec, +} + +struct RewrittenUnnestExprGroups { + plan: LogicalPlan, + expr_groups: Vec>, +} + +fn flatten_expr_groups(expr_groups: Vec>) -> Vec { + expr_groups.into_iter().flatten().collect() } impl SqlToRel<'_, S> { @@ -145,10 +165,43 @@ impl SqlToRel<'_, S> { // This alias map is resolved and looked up in both having exprs and group by exprs let alias_map = extract_aliases(&select_exprs); + // DISTINCT ON expressions are parsed alongside HAVING / QUALIFY so + // they participate in aggregate / window discovery and get rebased + // through the same pipeline. The SQL nodes are taken out of + // `select.distinct` so the later match on `Distinct::On` still fires + // but does not move the original Vec. + // + // Resolution precedence matches PostgreSQL and ORDER BY: SELECT + // aliases win over input columns. For example, + // SELECT DISTINCT ON (b) a AS b ... GROUP BY a + // resolves `b` to the alias for `a`, not to a same-named input + // column. + let on_exprs_sql: Vec = match &mut select.distinct { + Some(Distinct::On(exprs)) => std::mem::take(exprs), + _ => Vec::new(), + }; + let mut on_expr_schema = projected_plan.schema().as_ref().clone(); + on_expr_schema.merge(base_plan.schema()); + let on_exprs_pre_aggr: Vec = on_exprs_sql + .into_iter() + .map(|e| { + let expr = + self.sql_expr_to_logical_expr(e, &on_expr_schema, planner_context)?; + // PostgreSQL only substitutes an output alias when the whole + // ON expression is a bare identifier. `b` resolves to the + // alias; `b + 0` keeps `b` as the input column. + let expr = substitute_top_level_alias(expr, &alias_map); + let expr = normalize_col(expr, &projected_plan)?; + let (expr, _) = expr.infer_placeholder_types(&on_expr_schema)?; + Ok(expr) + }) + .collect::>>()?; + // Optionally the HAVING expression. let having_expr_opt = select .having .map::, _>(|having_expr| { + self.warn_on_null_equality_predicate(&having_expr); let having_expr = self.sql_expr_to_logical_expr( having_expr, &combined_schema, @@ -168,7 +221,10 @@ impl SqlToRel<'_, S> { // SELECT c1, MAX(c2) AS m FROM t GROUP BY c1 HAVING MAX(c2) > 10; // let having_expr = resolve_aliases_to_exprs(having_expr, &alias_map)?; - normalize_col(having_expr, &projected_plan) + let having_expr = normalize_col(having_expr, &projected_plan)?; + let (having_expr, _) = + having_expr.infer_placeholder_types(&combined_schema)?; + Ok(having_expr) }) .transpose()?; @@ -197,6 +253,8 @@ impl SqlToRel<'_, S> { base_plan.schema(), std::slice::from_ref(&group_by_expr), )?; + let (group_by_expr, _) = + group_by_expr.infer_placeholder_types(&combined_schema)?; Ok(group_by_expr) }) .collect::>>()? @@ -235,7 +293,10 @@ impl SqlToRel<'_, S> { // select row_number() over (PARTITION BY id) as rk from users qualify row_number() over (PARTITION BY id) > 1; // let qualify_expr = resolve_aliases_to_exprs(qualify_expr, &alias_map)?; - normalize_col(qualify_expr, &projected_plan) + let qualify_expr = normalize_col(qualify_expr, &projected_plan)?; + let (qualify_expr, _) = + qualify_expr.infer_placeholder_types(&combined_schema)?; + Ok(qualify_expr) }) .transpose()?; @@ -251,12 +312,15 @@ impl SqlToRel<'_, S> { // Find aggregates in ORDER BY let order_by_aggrs = find_aggregate_exprs(order_by_rex.iter().map(|s| &s.expr)); - // Combine: all aggregates from SELECT/HAVING/QUALIFY, plus ORDER BY aggregates - // that aren't already in SELECT/HAVING/QUALIFY + // Find aggregates in DISTINCT ON + let on_aggrs = find_aggregate_exprs(on_exprs_pre_aggr.iter()); + + // Combine: all aggregates from SELECT/HAVING/QUALIFY, plus ORDER BY + // and DISTINCT ON aggregates that aren't already covered. let mut aggr_exprs = select_having_qualify_aggrs; - for order_by_aggr in order_by_aggrs { - if !aggr_exprs.iter().any(|e| e == &order_by_aggr) { - aggr_exprs.push(order_by_aggr); + for extra_aggr in order_by_aggrs.into_iter().chain(on_aggrs) { + if !aggr_exprs.iter().any(|e| e == &extra_aggr) { + aggr_exprs.push(extra_aggr); } } @@ -267,6 +331,7 @@ impl SqlToRel<'_, S> { having_expr: having_expr_post_aggr, qualify_expr: qualify_expr_post_aggr, order_by_exprs: mut order_by_rex, + on_exprs: mut on_exprs_post_aggr, } = if !group_by_exprs.is_empty() || !aggr_exprs.is_empty() { self.aggregate( &base_plan, @@ -274,6 +339,7 @@ impl SqlToRel<'_, S> { having_expr_opt.as_ref(), qualify_expr_opt.as_ref(), &order_by_rex, + &on_exprs_pre_aggr, &group_by_exprs, &aggr_exprs, )? @@ -290,6 +356,7 @@ impl SqlToRel<'_, S> { having_expr: having_expr_opt, qualify_expr: qualify_expr_opt, order_by_exprs: order_by_rex, + on_exprs: on_exprs_pre_aggr, }, } }; @@ -304,12 +371,13 @@ impl SqlToRel<'_, S> { // All of the window expressions (deduplicated and rewritten to reference aggregates as // columns from input). Window functions may be sourced from the SELECT list, QUALIFY - // expression, or ORDER BY. + // expression, ORDER BY, or DISTINCT ON. let window_func_exprs = find_window_exprs( select_exprs_post_aggr .iter() .chain(qualify_expr_post_aggr.iter()) - .chain(order_by_rex.iter().map(|s| &s.expr)), + .chain(order_by_rex.iter().map(|s| &s.expr)) + .chain(on_exprs_post_aggr.iter()), ); // Process window functions after aggregation as they can reference @@ -336,6 +404,11 @@ impl SqlToRel<'_, S> { }) .collect::>>()?; + on_exprs_post_aggr = on_exprs_post_aggr + .iter() + .map(|expr| rebase_expr(expr, &window_func_exprs, &plan)) + .collect::>>()?; + plan }; @@ -377,39 +450,74 @@ impl SqlToRel<'_, S> { plan }; - // Try processing unnest expression or do the final projection - let plan = self.try_process_unnest(plan, select_exprs_post_aggr)?; - - // Process distinct clause + // Process distinct clause. For `DISTINCT ON` combined with + // aggregation, GROUP BY, or window functions we apply DistinctOn + // *before* the final projection so grouping columns and ORDER BY + // tie-breakers that aren't in the user SELECT stay in scope. + // DistinctOn provides the projection in that case (its select_expr + // list is wrapped in FIRST_VALUE during lowering). let plan = match select.distinct { - None => Ok(plan), - Some(Distinct::All) => Ok(plan), + None | Some(Distinct::All) => { + self.try_process_unnest(plan, select_exprs_post_aggr)? + } Some(Distinct::Distinct) => { - LogicalPlanBuilder::from(plan).distinct()?.build() + let plan = self.try_process_unnest(plan, select_exprs_post_aggr)?; + LogicalPlanBuilder::from(plan).distinct()?.build()? } - Some(Distinct::On(on_expr)) => { - if !aggr_exprs.is_empty() - || !group_by_exprs.is_empty() - || !window_func_exprs.is_empty() + Some(Distinct::On(_)) => { + if aggr_exprs.is_empty() + && group_by_exprs.is_empty() + && window_func_exprs.is_empty() { - return not_impl_err!( - "DISTINCT ON expressions with GROUP BY, aggregation or window functions are not supported " + // Fast path: no aggregation context. Fuse projection + // and deduplication into a single DistinctOn over + // `base_plan`. The sort attached to DistinctOn via + // `with_sort_expr` later normalizes against base_plan, + // so a bare ORDER BY alias (e.g. `ORDER BY x` where + // SELECT has `a AS x`) must be swapped back to the + // underlying input expression first. + order_by_rex = + substitute_top_level_aliases_in_sorts(order_by_rex, &alias_map); + LogicalPlanBuilder::from(base_plan) + .distinct_on(on_exprs_post_aggr, select_exprs, None)? + .build()? + } else { + // General path: DistinctOn layered over the post- + // aggregate / post-window plan (no extra Projection + // node — DistinctOn's lowering wraps each select_expr + // in FIRST_VALUE, which acts as the projection). + // + // The DistinctOn input has the post-aggregate raw + // column names (e.g. `max(t.c4)`), not the user-facing + // SELECT aliases (`agg2`). ORDER BY may reference + // those aliases — substitute them back to the + // underlying post-aggregate expression so they + // resolve against the DistinctOn input. + let select_alias_map = extract_aliases(&select_exprs_post_aggr); + order_by_rex = substitute_top_level_aliases_in_sorts( + order_by_rex, + &select_alias_map, ); - } - let on_expr = on_expr - .into_iter() - .map(|e| { - self.sql_expr_to_logical_expr(e, plan.schema(), planner_context) - }) - .collect::>>()?; + let DistinctOnUnnestPlanResult { + plan, + select_exprs: select_exprs_post_aggr, + on_exprs: on_exprs_post_aggr, + order_by_exprs: rewritten_order_by_rex, + } = self.try_process_distinct_on_unnest( + plan, + select_exprs_post_aggr, + on_exprs_post_aggr, + order_by_rex, + )?; + order_by_rex = rewritten_order_by_rex; - // Build the final plan - LogicalPlanBuilder::from(base_plan) - .distinct_on(on_expr, select_exprs, None)? - .build() + LogicalPlanBuilder::from(plan) + .distinct_on(on_exprs_post_aggr, select_exprs_post_aggr, None)? + .build()? + } } - }?; + }; // DISTRIBUTE BY let plan = if !select.distribute_by.is_empty() { @@ -441,98 +549,166 @@ impl SqlToRel<'_, S> { input: LogicalPlan, select_exprs: Vec, ) -> Result { + let RewrittenUnnestExprGroups { plan, expr_groups } = self + .rewrite_unnest_expr_groups( + input, + select_exprs.into_iter().map(|expr| vec![expr]).collect(), + )?; + + LogicalPlanBuilder::from(plan) + .project(flatten_expr_groups(expr_groups))? + .build() + } + + /// Rewrites SELECT-list UNNESTs while keeping hidden DISTINCT ON / ORDER + /// BY inputs available to the DistinctOn node. + fn try_process_distinct_on_unnest( + &self, + input: LogicalPlan, + select_exprs: Vec, + on_exprs: Vec, + order_by_exprs: Vec, + ) -> Result { + let select_len = select_exprs.len(); + let on_len = on_exprs.len(); + let mut expr_groups = select_exprs + .into_iter() + .map(|expr| vec![expr]) + .collect::>(); + expr_groups.extend(on_exprs.into_iter().map(|expr| vec![expr])); + expr_groups.extend( + order_by_exprs + .iter() + .map(|sort_expr| vec![sort_expr.expr.clone()]), + ); + + let RewrittenUnnestExprGroups { + plan, + mut expr_groups, + } = self.rewrite_unnest_expr_groups(input, expr_groups)?; + + let rewritten_select_exprs = + flatten_expr_groups(expr_groups.drain(..select_len).collect()); + let rewritten_on_exprs = expr_groups + .drain(..on_len) + .map(|exprs| self.expect_single_distinct_on_expr(exprs, "DISTINCT ON")) + .collect::>>()?; + let rewritten_order_by_exprs = order_by_exprs + .into_iter() + .zip(expr_groups) + .map(|(sort_expr, exprs)| { + Ok(sort_expr + .with_expr(self.expect_single_distinct_on_expr(exprs, "ORDER BY")?)) + }) + .collect::>>()?; + + Ok(DistinctOnUnnestPlanResult { + plan, + select_exprs: rewritten_select_exprs, + on_exprs: rewritten_on_exprs, + order_by_exprs: rewritten_order_by_exprs, + }) + } + + fn expect_single_distinct_on_expr( + &self, + exprs: Vec, + clause: &str, + ) -> Result { + if exprs.len() == 1 { + return Ok(exprs.into_iter().next().expect("len checked above")); + } + + not_impl_err!( + "{clause} expressions that expand to multiple columns are not supported with DISTINCT ON" + ) + } + + fn rewrite_unnest_expr_groups( + &self, + input: LogicalPlan, + expr_groups: Vec>, + ) -> Result { // Try process group by unnest let input = self.try_process_aggregate_unnest(input)?; let mut intermediate_plan = input; - let mut intermediate_select_exprs = select_exprs; - // Fast path: If there is are no unnests in the select_exprs, wrap the plan in a projection - if !intermediate_select_exprs - .iter() - .any(has_unnest_expr_recursively) - { - return LogicalPlanBuilder::from(intermediate_plan) - .project(intermediate_select_exprs)? - .build(); - } + let mut intermediate_expr_groups = expr_groups; - // Each expr in select_exprs can contains multiple unnest stage - // The transformation happen bottom up, one at a time for each iteration - // Only exhaust the loop if no more unnest transformation is found - for i in 0.. { + loop { let mut unnest_columns = IndexMap::new(); - // from which column used for projection, before the unnest happen - // including non unnest column and unnest column + // from which columns used for projection, before the unnest happen + // including non unnest columns and unnest columns let mut inner_projection_exprs = vec![]; + let mut outer_expr_groups = + Vec::with_capacity(intermediate_expr_groups.len()); + + for expr_group in &intermediate_expr_groups { + let mut outer_expr_group = vec![]; + for expr in expr_group { + let mut rewritten_exprs = rewrite_recursive_unnest_bottom_up( + &intermediate_plan, + &mut unnest_columns, + &mut inner_projection_exprs, + expr, + )?; - // expr returned here maybe different from the originals in inner_projection_exprs - // for example: - // - unnest(struct_col) will be transformed into struct_col.field1, struct_col.field2 - // - unnest(array_col) will be transformed into array_col.element - // - unnest(array_col) + 1 will be transformed into array_col.element +1 - let mut outer_projection_exprs = vec![]; - for expr in &intermediate_select_exprs { - let mut rewritten_exprs = rewrite_recursive_unnest_bottom_up( - &intermediate_plan, - &mut unnest_columns, - &mut inner_projection_exprs, - expr, - )?; + if let Some(columns) = + self.get_struct_unnest_columns(&intermediate_plan, expr)? + { + rewritten_exprs = rewritten_exprs + .into_iter() + .zip(columns) + .map(|(expr, column)| expr.alias(column.flat_name())) + .collect(); + } - if let Some(columns) = - self.get_struct_unnest_columns(&intermediate_plan, expr)? - { - rewritten_exprs = rewritten_exprs - .into_iter() - .zip(columns) - .map(|(expr, column)| expr.alias(column.flat_name())) - .collect(); + outer_expr_group.extend(rewritten_exprs); } - - outer_projection_exprs.extend(rewritten_exprs); + outer_expr_groups.push(outer_expr_group); } // No more unnest is possible if unnest_columns.is_empty() { - // The original expr does not contain any unnest - if i == 0 { - return LogicalPlanBuilder::from(intermediate_plan) - .project(intermediate_select_exprs)? - .build(); - } - break; - } else { - // Set preserve_nulls to false to ensure compatibility with DuckDB and PostgreSQL - let mut unnest_options = UnnestOptions::new().with_preserve_nulls(false); - let mut unnest_col_vec = vec![]; + return Ok(RewrittenUnnestExprGroups { + plan: intermediate_plan, + expr_groups: intermediate_expr_groups, + }); + } - for (col, maybe_list_unnest) in unnest_columns.into_iter() { - if let Some(list_unnest) = maybe_list_unnest { - unnest_options = list_unnest.into_iter().fold( - unnest_options, - |options, unnest_list| { - options.with_recursions(RecursionUnnestOption { - input_column: col.clone(), - output_column: unnest_list.output_column, - depth: unnest_list.depth, - }) - }, - ); - } - unnest_col_vec.push(col); + // The default SQL `UNNEST` matches DuckDB/PostgreSQL: drop both + // NULL and empty input lists. Outer-unnest (modelled as + // `Unnest { outer: true }`) overrides that and selects + // `NullHandling::PreserveAndExpandEmpty`. Mixing the two in a + // single SELECT is a planning error because `UnnestOptions` is + // per-`UnnestExec`, not per-column. + let null_handling = collect_unnest_null_handling(&intermediate_expr_groups)?; + let mut unnest_options = + UnnestOptions::new().with_null_handling(null_handling); + let mut unnest_col_vec = vec![]; + + for (col, maybe_list_unnest) in unnest_columns.into_iter() { + if let Some(list_unnest) = maybe_list_unnest { + unnest_options = list_unnest.into_iter().fold( + unnest_options, + |options, unnest_list| { + options.with_recursions(RecursionUnnestOption { + input_column: col.clone(), + output_column: unnest_list.output_column, + depth: unnest_list.depth, + }) + }, + ); } - let plan = LogicalPlanBuilder::from(intermediate_plan) - .project(inner_projection_exprs)? - .unnest_columns_with_options(unnest_col_vec, unnest_options)? - .build()?; - intermediate_plan = plan; - intermediate_select_exprs = outer_projection_exprs; + unnest_col_vec.push(col); } - } - LogicalPlanBuilder::from(intermediate_plan) - .project(intermediate_select_exprs)? - .build() + intermediate_plan = LogicalPlanBuilder::from(intermediate_plan) + .project(inner_projection_exprs)? + .unnest_columns_with_options(unnest_col_vec, unnest_options)? + .build()?; + intermediate_expr_groups = outer_expr_groups; + } } fn get_struct_unnest_columns( @@ -708,6 +884,7 @@ impl SqlToRel<'_, S> { Some(predicate_expr) => { let fallback_schemas = plan.fallback_normalize_schemas(); + self.warn_on_null_equality_predicate(&predicate_expr); let filter_expr = self.sql_to_expr(predicate_expr, plan.schema(), planner_context)?; @@ -1017,6 +1194,7 @@ impl SqlToRel<'_, S> { having_expr_opt: Option<&Expr>, qualify_expr_opt: Option<&Expr>, order_by_exprs: &[SortExpr], + on_exprs: &[Expr], group_by_exprs: &[Expr], aggr_exprs: &[Expr], ) -> Result { @@ -1183,12 +1361,28 @@ impl SqlToRel<'_, S> { ), )?; + // Rewrite the DISTINCT ON expressions to use the columns produced by + // the aggregation. Same shape as ORDER BY rewriting so a hidden + // grouping column or a raw aggregate expression in ON is resolved. + let on_exprs_post_aggr = on_exprs + .iter() + .map(|expr| rebase_expr(expr, &aggr_projection_exprs, input)) + .collect::>>()?; + check_columns_satisfy_exprs( + &all_valid_exprs, + &on_exprs_post_aggr, + CheckColumnsSatisfyExprsPurpose::Aggregate( + CheckColumnsMustReferenceAggregatePurpose::DistinctOn, + ), + )?; + Ok(AggregatePlanResult { plan, select_exprs: select_exprs_post_aggr, having_expr: having_expr_post_aggr, qualify_expr: qualify_expr_post_aggr, order_by_exprs: order_by_post_aggr, + on_exprs: on_exprs_post_aggr, }) } @@ -1275,3 +1469,45 @@ fn has_unnest_expr_recursively(expr: &Expr) -> bool { }); has_unnest } + +/// Walk `select_exprs`, observe every [`Expr::Unnest`] inside them, and +/// derive the [`NullHandling`] mode for the resulting [`UnnestOptions`]. +/// +/// * No unnest with `outer = true` → [`NullHandling::Drop`] (default SQL +/// `UNNEST(...)` semantics, matching DuckDB/PostgreSQL). +/// * Every unnest with `outer = true` → [`NullHandling::PreserveAndExpandEmpty`] +/// (outer-unnest semantics: `NULL` and empty input lists each produce a +/// single `NULL` output row). +/// * A mix of `outer = true` and `outer = false` in one SELECT → planning +/// error, because `UnnestOptions` applies per `Unnest` plan node, not +/// per output column. +fn collect_unnest_null_handling(expr_groups: &[Vec]) -> Result { + let mut saw_outer = false; + let mut saw_inner = false; + for group in expr_groups { + for expr in group { + expr.apply(|e| { + if let Expr::Unnest(UnnestExpr { outer, .. }) = e { + if *outer { + saw_outer = true; + } else { + saw_inner = true; + } + } + Ok(TreeNodeRecursion::Continue) + })?; + } + } + if saw_outer && saw_inner { + return plan_err!( + "Cannot mix `unnest(...)` with `unnest_outer(...)` in the same \ + SELECT — the unnest operator carries a single null-handling \ + mode. Split the query so each unnest projection uses one mode." + ); + } + Ok(if saw_outer { + NullHandling::PreserveAndExpandEmpty + } else { + NullHandling::Drop + }) +} diff --git a/datafusion/sql/src/set_expr.rs b/datafusion/sql/src/set_expr.rs index dc8e4f14d1ee8..51b11f3087095 100644 --- a/datafusion/sql/src/set_expr.rs +++ b/datafusion/sql/src/set_expr.rs @@ -25,70 +25,71 @@ use datafusion_expr::{LogicalPlan, LogicalPlanBuilder}; use sqlparser::ast::{SetExpr, SetOperator, SetQuantifier, Spanned}; impl SqlToRel<'_, S> { - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] pub(super) fn set_expr_to_plan( &self, set_expr: SetExpr, planner_context: &mut PlannerContext, ) -> Result { - let set_expr_span = Span::try_from_sqlparser_span(set_expr.span()); - match set_expr { - SetExpr::Select(s) => self.select_to_plan(*s, None, planner_context), - SetExpr::Values(v) => self.sql_values_to_plan(v, planner_context), - SetExpr::SetOperation { - op, - left, - right, - set_quantifier, - } => { - let left_span = Span::try_from_sqlparser_span(left.span()); - let right_span = Span::try_from_sqlparser_span(right.span()); - let left_plan = self.set_expr_to_plan(*left, planner_context); - // Store the left plan's schema so that the right side can - // alias duplicate expressions to match. Skip for BY NAME - // operations since those match columns by name, not position. - if let Ok(plan) = &left_plan - && plan.schema().fields().len() > 1 - && !matches!( - set_quantifier, - SetQuantifier::ByName - | SetQuantifier::AllByName - | SetQuantifier::DistinctByName - ) - { - planner_context - .set_set_expr_left_schema(Some(Arc::clone(plan.schema()))); - } - let right_plan = self.set_expr_to_plan(*right, planner_context); - planner_context.set_set_expr_left_schema(None); - let (left_plan, right_plan) = match (left_plan, right_plan) { - (Ok(left_plan), Ok(right_plan)) => (left_plan, right_plan), - (Err(left_err), Err(right_err)) => { - return Err(DataFusionError::Collection(vec![ - left_err, right_err, - ])); + crate::stack::maybe_grow(|| { + let set_expr_span = Span::try_from_sqlparser_span(set_expr.span()); + match set_expr { + SetExpr::Select(s) => self.select_to_plan(*s, None, planner_context), + SetExpr::Values(v) => self.sql_values_to_plan(v, planner_context), + SetExpr::SetOperation { + op, + left, + right, + set_quantifier, + } => { + let left_span = Span::try_from_sqlparser_span(left.span()); + let right_span = Span::try_from_sqlparser_span(right.span()); + let left_plan = self.set_expr_to_plan(*left, planner_context); + // Store the left plan's schema so that the right side can + // alias duplicate expressions to match. Skip for BY NAME + // operations since those match columns by name, not position. + if let Ok(plan) = &left_plan + && plan.schema().fields().len() > 1 + && !matches!( + set_quantifier, + SetQuantifier::ByName + | SetQuantifier::AllByName + | SetQuantifier::DistinctByName + ) + { + planner_context + .set_set_expr_left_schema(Some(Arc::clone(plan.schema()))); } - (Err(err), _) | (_, Err(err)) => { - return Err(err); + let right_plan = self.set_expr_to_plan(*right, planner_context); + planner_context.set_set_expr_left_schema(None); + let (left_plan, right_plan) = match (left_plan, right_plan) { + (Ok(left_plan), Ok(right_plan)) => (left_plan, right_plan), + (Err(left_err), Err(right_err)) => { + return Err(DataFusionError::Collection(vec![ + left_err, right_err, + ])); + } + (Err(err), _) | (_, Err(err)) => { + return Err(err); + } + }; + if !(set_quantifier == SetQuantifier::ByName + || set_quantifier == SetQuantifier::AllByName) + { + self.validate_set_expr_num_of_columns( + op, + left_span, + right_span, + &left_plan, + &right_plan, + set_expr_span, + )?; } - }; - if !(set_quantifier == SetQuantifier::ByName - || set_quantifier == SetQuantifier::AllByName) - { - self.validate_set_expr_num_of_columns( - op, - left_span, - right_span, - &left_plan, - &right_plan, - set_expr_span, - )?; + self.set_operation_to_plan(op, left_plan, right_plan, set_quantifier) } - self.set_operation_to_plan(op, left_plan, right_plan, set_quantifier) + SetExpr::Query(q) => self.query_to_plan(*q, planner_context), + _ => not_impl_err!("Query {set_expr} not implemented yet"), } - SetExpr::Query(q) => self.query_to_plan(*q, planner_context), - _ => not_impl_err!("Query {set_expr} not implemented yet"), - } + }) } pub(super) fn is_union_all(set_quantifier: SetQuantifier) -> Result { diff --git a/datafusion/sql/src/stack.rs b/datafusion/sql/src/stack.rs index b7d5eebdd7188..ed3bf1553ebfc 100644 --- a/datafusion/sql/src/stack.rs +++ b/datafusion/sql/src/stack.rs @@ -15,49 +15,43 @@ // specific language governing permissions and limitations // under the License. -pub use inner::StackGuard; - -/// A guard that sets the minimum stack size for the current thread to `min_stack_size` bytes. +/// The local red zone used by SQL recursive entry points. +/// +/// Some SQL planner and unparser recursion paths need more than `recursive`'s +/// default 128 KiB red zone in debug builds. Keep this value local to each +/// stack-growth checkpoint rather than mutating `recursive`'s process-global +/// minimum stack size. #[cfg(feature = "recursive_protection")] -mod inner { - /// Sets the stack size to `min_stack_size` bytes on call to `new()` and - /// resets to the previous value when this structure is dropped. - pub struct StackGuard { - previous_stack_size: usize, - } +pub(crate) const SQL_RECURSION_RED_ZONE: usize = 256 * 1024; - impl StackGuard { - /// Sets the stack size to `min_stack_size` bytes on call to `new()` and - /// resets to the previous value when this structure is dropped. - pub fn new(min_stack_size: usize) -> Self { - let previous_stack_size = recursive::get_minimum_stack_size(); - recursive::set_minimum_stack_size(min_stack_size); - Self { - previous_stack_size, - } - } - } - - impl Drop for StackGuard { - fn drop(&mut self) { - recursive::set_minimum_stack_size(self.previous_stack_size); - } - } +/// Runs `callback` on a stack with enough space for SQL recursive entry points. +#[cfg(feature = "recursive_protection")] +#[inline] +pub(crate) fn maybe_grow(callback: impl FnOnce() -> R) -> R { + stacker::maybe_grow( + SQL_RECURSION_RED_ZONE, + recursive::get_stack_allocation_size(), + callback, + ) } -/// A stub implementation of the stack guard when the recursive protection -/// feature is not enabled +/// Runs `callback` without stack growth when recursive protection is disabled. #[cfg(not(feature = "recursive_protection"))] -mod inner { - /// A stub implementation of the stack guard when the recursive protection - /// feature is not enabled that does nothing - pub struct StackGuard; +#[inline] +pub(crate) fn maybe_grow(callback: impl FnOnce() -> R) -> R { + callback() +} + +#[cfg(all(test, feature = "recursive_protection"))] +mod tests { + use super::*; + + #[test] + fn maybe_grow_does_not_mutate_recursive_minimum_stack_size() { + let before = recursive::get_minimum_stack_size(); + let observed = maybe_grow(recursive::get_minimum_stack_size); - impl StackGuard { - /// A stub implementation of the stack guard when the recursive protection - /// feature is not enabled - pub fn new(_min_stack_size: usize) -> Self { - Self - } + assert_eq!(observed, before); + assert_eq!(recursive::get_minimum_stack_size(), before); } } diff --git a/datafusion/sql/src/statement.rs b/datafusion/sql/src/statement.rs index 8c94610f7764c..fd5c34ff5d961 100644 --- a/datafusion/sql/src/statement.rs +++ b/datafusion/sql/src/statement.rs @@ -31,14 +31,18 @@ use crate::utils::normalize_ident; use arrow::datatypes::{Field, FieldRef, Fields}; use datafusion_common::error::_plan_err; +use datafusion_common::format::ExplainStatementOptions; use datafusion_common::parsers::CompressionTypeVariant; +use datafusion_common::tree_node::{Transformed, TreeNode, TreeNodeRecursion}; use datafusion_common::{ Column, Constraint, Constraints, DFSchema, DFSchemaRef, DataFusionError, Result, ScalarValue, SchemaError, SchemaReference, TableReference, ToDFSchema, exec_err, internal_err, not_impl_err, plan_datafusion_err, plan_err, schema_err, unqualified_field_not_found, }; -use datafusion_expr::dml::{CopyTo, InsertOp}; +use datafusion_expr::dml::{ + CopyTo, InsertOp, MergeIntoAction, MergeIntoClause, MergeIntoClauseKind, MergeIntoOp, +}; use datafusion_expr::expr_rewriter::normalize_col_with_schemas_and_ambiguity_check; use datafusion_expr::logical_plan::DdlStatement; use datafusion_expr::logical_plan::builder::project; @@ -52,7 +56,7 @@ use datafusion_expr::{ LogicalPlan, LogicalPlanBuilder, OperateFunctionArg, PlanType, Prepare, ResetVariable, SetVariable, SortExpr, Statement as PlanStatement, ToStringifiedPlan, TransactionAccessMode, TransactionConclusion, TransactionEnd, - TransactionIsolationLevel, TransactionStart, Volatility, WriteOp, cast, col, + TransactionIsolationLevel, TransactionStart, Volatility, WriteOp, cast, }; use sqlparser::ast::{ self, BeginTransactionKind, CheckConstraint, ForeignKeyConstraint, IndexColumn, @@ -227,12 +231,9 @@ impl SqlToRel<'_, S> { DFStatement::CreateExternalTable(s) => self.external_table_to_plan(s), DFStatement::Statement(s) => self.sql_statement_to_plan(*s), DFStatement::CopyTo(s) => self.copy_to_plan(s), - DFStatement::Explain(ExplainStatement { - verbose, - analyze, - format, - statement, - }) => self.explain_to_plan(verbose, analyze, format, *statement), + DFStatement::Explain(ExplainStatement { options, statement }) => { + self.explain_to_plan(options, *statement) + } DFStatement::Reset(statement) => self.reset_statement_to_plan(statement), } } @@ -280,12 +281,21 @@ impl SqlToRel<'_, S> { statement, analyze, format, - describe_alias: _, .. } => { - let format = format.map(|format| format.to_string()); + let format = format + .map(|format| ExplainFormat::from_str(&format.to_string())) + .transpose()?; let statement = DFStatement::Statement(statement); - self.explain_to_plan(verbose, analyze, format, statement) + let options = ExplainStatementOptions { + analyze, + verbose, + format, + analyze_level: None, + analyze_categories: None, + show_statistics: None, + }; + self.explain_to_plan(options, statement) } Statement::Query(query) => self.query_to_plan(*query, planner_context), Statement::ShowVariable { variable } => self.show_variable_to_plan(&variable), @@ -548,14 +558,14 @@ impl SqlToRel<'_, S> { input_schema.fields().len() ); } - let input_fields = input_schema.fields(); + let input_columns = input_schema.columns(); let project_exprs = schema .fields() .iter() - .zip(input_fields) - .map(|(field, input_field)| { + .zip(input_columns) + .map(|(field, input_column)| { cast( - col(input_field.name()), + Expr::Column(input_column), field.data_type().clone(), ) .alias(field.name()) @@ -1209,6 +1219,8 @@ impl SqlToRel<'_, S> { self.delete_to_plan(&table_name, selection, limit) } + Statement::Merge(merge) => self.merge_to_plan(merge), + Statement::StartTransaction { modes, begin: false, @@ -1460,7 +1472,7 @@ impl SqlToRel<'_, S> { function_body, }; - let statement = DdlStatement::CreateFunction(CreateFunction { + let statement = DdlStatement::CreateFunction(Box::new(CreateFunction { or_replace, temporary, name, @@ -1468,7 +1480,7 @@ impl SqlToRel<'_, S> { args, params, schema: DFSchemaRef::new(DFSchema::empty()), - }); + })); Ok(LogicalPlan::Ddl(statement)) } @@ -1797,7 +1809,7 @@ impl SqlToRel<'_, S> { name, columns, file_type, - location, + locations, table_partition_cols, if_not_exists, temporary, @@ -1846,19 +1858,29 @@ impl SqlToRel<'_, S> { let name = self.object_name_to_table_reference(name)?; let constraints = self.new_constraint_from_table_constraints(&all_constraints, &df_schema)?; + + let Some(location) = locations.first().cloned() else { + return plan_err!("CREATE EXTERNAL TABLE requires at least one location"); + }; + + // Keep the existing single-location builder API: seed it with the first + // location, then replace it with the complete list. Ok(LogicalPlan::Ddl(DdlStatement::CreateExternalTable( - PlanCreateExternalTable::builder(name, location, file_type, df_schema) - .with_partition_cols(table_partition_cols) - .with_if_not_exists(if_not_exists) - .with_or_replace(or_replace) - .with_temporary(temporary) - .with_definition(definition) - .with_order_exprs(ordered_exprs) - .with_unbounded(unbounded) - .with_options(options_map) - .with_constraints(constraints) - .with_column_defaults(column_defaults) - .build(), + Box::new( + PlanCreateExternalTable::builder(name, location, file_type, df_schema) + .with_locations(locations) + .with_partition_cols(table_partition_cols) + .with_if_not_exists(if_not_exists) + .with_or_replace(or_replace) + .with_temporary(temporary) + .with_definition(definition) + .with_order_exprs(ordered_exprs) + .with_unbounded(unbounded) + .with_options(options_map) + .with_constraints(constraints) + .with_column_defaults(column_defaults) + .build(), + ), ))) } @@ -2004,9 +2026,7 @@ impl SqlToRel<'_, S> { /// datafusion `EXPLAIN` statement. fn explain_to_plan( &self, - verbose: bool, - analyze: bool, - format: Option, + opts: ExplainStatementOptions, statement: DFStatement, ) -> Result { let plan = self.statement_to_plan(statement)?; @@ -2018,34 +2038,77 @@ impl SqlToRel<'_, S> { let schema = LogicalPlan::explain_schema(); let schema = schema.to_dfschema_ref()?; + let ExplainStatementOptions { + analyze, + verbose, + format, + analyze_level, + analyze_categories, + show_statistics, + } = opts; + + // Mutual exclusivity checks if verbose && format.is_some() { return plan_err!("EXPLAIN VERBOSE with FORMAT is not supported"); } + if !analyze { + if analyze_level.is_some() { + return plan_err!("EXPLAIN option LEVEL requires ANALYZE"); + } + if analyze_categories.is_some() { + return plan_err!("EXPLAIN option METRICS requires ANALYZE"); + } + } + if analyze && show_statistics.is_some() { + return plan_err!("EXPLAIN option COSTS cannot be combined with ANALYZE"); + } + + // Resolve the requested output format. + // + // Verbose mode only supports indent format, and for EXPLAIN ANALYZE + // only `Indent` and `PostgresJSON` are supported today — `Tree` and + // `Graphviz` require additional work to render with live metrics. + let options = self.context_provider.options(); + let format = if verbose { + ExplainFormat::Indent + } else if let Some(format) = format { + format + } else if analyze { + ExplainFormat::Indent + } else { + options.explain.format.clone() + }; if analyze { - if format.is_some() { - return plan_err!("EXPLAIN ANALYZE with FORMAT is not supported"); + match &format { + ExplainFormat::Indent => {} + ExplainFormat::PostgresJSON => { + // The pgjson renderer does not emit statistics yet, so + // reject the combination rather than silently ignoring it. + if options.explain.show_statistics { + return plan_err!( + "EXPLAIN ANALYZE with FORMAT pgjson does not support show_statistics" + ); + } + } + ExplainFormat::Tree | ExplainFormat::Graphviz => { + return plan_err!( + "EXPLAIN ANALYZE with FORMAT {format} is not supported" + ); + } } Ok(LogicalPlan::Analyze(Analyze { verbose, + format, input: plan, schema, + analyze_level, + analyze_categories, })) } else { let stringified_plans = vec![plan.to_stringified(PlanType::InitialLogicalPlan)]; - // default to configuration value - // verbose mode only supports indent format - let options = self.context_provider.options(); - let format = if verbose { - ExplainFormat::Indent - } else if let Some(format) = format { - ExplainFormat::from_str(&format)? - } else { - options.explain.format.clone() - }; - Ok(LogicalPlan::Explain(Explain { verbose, explain_format: format, @@ -2053,6 +2116,7 @@ impl SqlToRel<'_, S> { stringified_plans, schema, logical_optimization_succeeded: false, + show_statistics, })) } } @@ -2355,6 +2419,403 @@ impl SqlToRel<'_, S> { Ok(plan) } + fn merge_to_plan(&self, merge: ast::Merge) -> Result { + let ast::Merge { + table, + source, + on, + clauses, + into: _, + merge_token: _, + optimizer_hints, + output, + } = merge; + + if !optimizer_hints.is_empty() { + plan_err!("Optimizer hints not supported")?; + } + + if output.is_some() { + return not_impl_err!("MERGE OUTPUT clause is not supported"); + } + + if clauses.is_empty() { + return plan_err!("MERGE INTO requires at least one WHEN clause"); + } + + // 1. Resolve target table + let (target_table_name, target_alias) = match table { + TableFactor::Table { + name, + alias, + args, + with_hints, + version, + with_ordinality, + partitions, + json_path, + sample, + index_hints, + } => { + if alias + .as_ref() + .is_some_and(|alias| !alias.columns.is_empty()) + { + return not_impl_err!( + "MERGE target alias column lists are not supported" + ); + } + if args.is_some() + || !with_hints.is_empty() + || version.is_some() + || with_ordinality + || !partitions.is_empty() + || json_path.is_some() + || sample.is_some() + || !index_hints.is_empty() + { + return not_impl_err!( + "MERGE target table modifiers are not supported" + ); + } + (name, alias) + } + _ => plan_err!("Cannot MERGE INTO non-table relation!")?, + }; + let target_table_ref = self.object_name_to_table_reference(target_table_name)?; + let target_table_source = self + .context_provider + .get_table_source(target_table_ref.clone())?; + // Use alias as schema qualifier so `t.col` resolves when user writes + // `MERGE INTO target AS t`. Fall back to the table reference itself. + let target_qualifier = target_alias + .as_ref() + .map(|a| { + TableReference::bare(self.ident_normalizer.normalize(a.name.clone())) + }) + .unwrap_or_else(|| target_table_ref.clone()); + let target_schema = Arc::new(DFSchema::try_from_qualified_schema( + target_qualifier.clone(), + &target_table_source.schema(), + )?); + + // 2. Plan the source (USING clause) as a LogicalPlan + let mut planner_context = PlannerContext::new(); + let source_table_with_joins = TableWithJoins { + relation: source, + joins: vec![], + }; + let source_plan = + self.plan_from_tables(vec![source_table_with_joins], &mut planner_context)?; + + // 3. Build a combined schema for resolving expressions in ON and WHEN clauses + let combined_schema = + Arc::new(target_schema.as_ref().join(source_plan.schema())?); + + // 4. Convert the ON condition from sqlparser Expr to datafusion Expr + let on_expr = self.sql_to_expr(*on, &combined_schema, &mut planner_context)?; + + // 5. Convert each WHEN clause + let df_clauses = clauses + .into_iter() + .map(|clause| { + self.merge_clause_to_plan( + clause, + &combined_schema, + &target_schema, + &target_qualifier, + &mut planner_context, + ) + }) + .collect::>>()?; + + // 6. Build the MERGE operation. Column references to the target may be + // qualified with the SQL alias (`MERGE INTO target AS t ... t.col`). + // Canonicalize those to the real target table qualifier so the stored + // plan is independent of the alias: this lets the analyzer passes and + // proto deserialization rebuild the target schema from `table_name` + // alone, without carrying the alias as extra state. + let mut merge_op = MergeIntoOp { + on: on_expr, + clauses: df_clauses, + }; + if target_qualifier != target_table_ref { + // Target references in correlated subqueries are represented as + // `OuterReferenceColumn`s inside the embedded logical plan. The + // alias canonicalization below only rewrites top-level expression + // columns, so accepting such a subquery would leave the target + // alias in the public MERGE representation. Reject this case until + // the alias can be rewritten scope-safely inside subquery plans. + for expr in merge_op.exprs() { + if Self::has_outer_reference_to_qualifier(expr, &target_qualifier)? { + return not_impl_err!( + "MERGE subqueries correlated to target alias \ + '{target_qualifier}' are not supported" + ); + } + } + + // Canonicalizing target columns to `target_table_ref` is only safe + // when the source does not already use that qualifier. If it does + // (e.g. `MERGE INTO target AS t USING source AS target`), the two + // namespaces would collapse and later resolution could silently + // pick the source column for a target reference. Reject that + // collision rather than change the meaning of the condition. + if source_plan.schema().iter().any(|(qualifier, _)| { + qualifier.is_some_and(|q| q.resolved_eq(&target_table_ref)) + }) { + return plan_err!( + "MERGE source may not use the target table name '{target_table_ref}' \ + as a qualifier while the target is aliased as '{target_qualifier}'; \ + use a different source alias" + ); + } + let canonical = merge_op + .exprs() + .into_iter() + .cloned() + .map(|expr| { + Self::canonicalize_target_qualifier( + expr, + &target_qualifier, + &target_table_ref, + ) + }) + .collect::>>()?; + merge_op = merge_op.with_new_exprs(canonical)?; + } + + Ok(LogicalPlan::Dml(DmlStatement::new( + target_table_ref, + target_table_source, + WriteOp::MergeInto(Box::new(merge_op)), + Arc::new(source_plan), + ))) + } + + /// Rewrite every [`Expr::Column`] qualified with `from` to instead use + /// `to`, leaving all other columns untouched. Used to canonicalize MERGE + /// target-alias references to the real target table qualifier. + fn canonicalize_target_qualifier( + expr: Expr, + from: &TableReference, + to: &TableReference, + ) -> Result { + expr.transform(|expr| match expr { + Expr::Column(col) if col.relation.as_ref() == Some(from) => Ok( + Transformed::yes(Expr::Column(Column::new(Some(to.clone()), col.name))), + ), + other => Ok(Transformed::no(other)), + }) + .map(|transformed| transformed.data) + } + + /// Return true if an expression contains a subquery whose embedded plan + /// has an outer reference qualified by `qualifier`. + fn has_outer_reference_to_qualifier( + expr: &Expr, + qualifier: &TableReference, + ) -> Result { + let mut found = false; + expr.apply(|expr| { + let subquery = match expr { + Expr::Exists(exists) => Some(&exists.subquery), + Expr::InSubquery(in_subquery) => Some(&in_subquery.subquery), + Expr::SetComparison(set_comparison) => Some(&set_comparison.subquery), + Expr::ScalarSubquery(subquery) => Some(subquery), + _ => None, + }; + + if let Some(subquery) = subquery { + subquery.subquery.apply_with_subqueries(|plan| { + plan.apply_expressions(|expr| { + expr.apply(|expr| { + if let Expr::OuterReferenceColumn(_, column) = expr + && column.relation.as_ref() == Some(qualifier) + { + found = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } + }) + })?; + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + } + + Ok(if found { + TreeNodeRecursion::Stop + } else { + TreeNodeRecursion::Continue + }) + })?; + Ok(found) + } + + fn merge_target_column_name( + &self, + name: &ObjectName, + target_qualifier: &TableReference, + ) -> Result { + let part = name + .0 + .iter() + .last() + .ok_or_else(|| plan_datafusion_err!("Empty column name"))?; + let ident = part + .as_ident() + .cloned() + .ok_or_else(|| plan_datafusion_err!("Expected simple identifier"))?; + + if name.0.len() > 1 { + let qualifier = self.object_name_to_table_reference(ObjectName( + name.0[..name.0.len() - 1].to_vec(), + ))?; + if !qualifier.resolved_eq(target_qualifier) { + return plan_err!( + "MERGE assignment target '{name}' must reference target table \ + '{target_qualifier}'" + ); + } + } + + Ok(self.ident_normalizer.normalize(ident)) + } + + fn merge_clause_to_plan( + &self, + clause: ast::MergeClause, + combined_schema: &DFSchema, + target_schema: &DFSchema, + target_qualifier: &TableReference, + planner_context: &mut PlannerContext, + ) -> Result { + let kind = match clause.clause_kind { + ast::MergeClauseKind::Matched => MergeIntoClauseKind::Matched, + ast::MergeClauseKind::NotMatched => MergeIntoClauseKind::NotMatched, + ast::MergeClauseKind::NotMatchedByTarget => { + MergeIntoClauseKind::NotMatchedByTarget + } + ast::MergeClauseKind::NotMatchedBySource => { + MergeIntoClauseKind::NotMatchedBySource + } + }; + + let predicate = clause + .predicate + .map(|p| self.sql_to_expr(p, combined_schema, planner_context)) + .transpose()?; + + let action = match clause.action { + ast::MergeAction::Update(update_expr) => { + if update_expr.update_predicate.is_some() { + return not_impl_err!( + "MERGE UPDATE WHERE predicates are not supported" + ); + } + if update_expr.delete_predicate.is_some() { + return not_impl_err!( + "MERGE UPDATE DELETE WHERE predicates are not supported" + ); + } + let assignments = update_expr + .assignments + .into_iter() + .map(|assign| { + let col_name = match &assign.target { + AssignmentTarget::ColumnName(cols) => { + self.merge_target_column_name(cols, target_qualifier)? + } + _ => plan_err!("Tuples are not supported")?, + }; + // Validate column exists in target + target_schema.field_with_unqualified_name(&col_name)?; + let value = self.sql_to_expr( + assign.value, + combined_schema, + planner_context, + )?; + Ok((col_name, value)) + }) + .collect::>>()?; + let mut seen = HashSet::new(); + for (column, _) in &assignments { + if !seen.insert(column.as_str()) { + return plan_err!("Duplicate column '{column}' in MERGE UPDATE"); + } + } + MergeIntoAction::Update(assignments) + } + ast::MergeAction::Insert(insert_expr) => { + if insert_expr.insert_predicate.is_some() { + return not_impl_err!( + "MERGE INSERT WHERE predicates are not supported" + ); + } + let columns: Vec = insert_expr + .columns + .iter() + .map(|c| self.merge_target_column_name(c, target_qualifier)) + .collect::>>()?; + + // Validate: no duplicates, all columns exist in target schema + let mut seen = HashSet::new(); + for col in &columns { + if !seen.insert(col.as_str()) { + return plan_err!("Duplicate column '{col}' in MERGE INSERT"); + } + target_schema.field_with_unqualified_name(col)?; + } + + let num_target_cols = target_schema.fields().len(); + + let values = match insert_expr.kind { + ast::MergeInsertKind::Values(values) => { + if values.rows.len() != 1 { + return plan_err!( + "MERGE INSERT must have exactly one row of values" + ); + } + let row = values.rows.into_iter().next().unwrap().content; + let expected = if columns.is_empty() { + num_target_cols + } else { + columns.len() + }; + if row.len() != expected { + return plan_err!( + "MERGE INSERT has {expected} column(s) but {} value(s)", + row.len() + ); + } + row.into_iter() + .map(|v| { + self.sql_to_expr(v, combined_schema, planner_context) + }) + .collect::>>()? + } + ast::MergeInsertKind::Row => { + return not_impl_err!("MERGE INSERT ROW is not supported"); + } + }; + + MergeIntoAction::Insert { columns, values } + } + ast::MergeAction::Delete { .. } => MergeIntoAction::Delete, + }; + + Ok(MergeIntoClause { + kind, + predicate, + action, + }) + } + fn insert_to_plan( &self, table_name: ObjectName, @@ -2435,17 +2896,23 @@ impl SqlToRel<'_, S> { span: _, }) = val { - let name = - name.replace('$', "").parse::().map_err(|_| { - plan_datafusion_err!("Can't parse placeholder: {name}") - })? - 1; + let index = match name[1..].parse::().map_err(|_| { + plan_datafusion_err!("Can't parse placeholder: {name}") + })? { + 0 => { + return plan_err!( + "Invalid placeholder, zero is not a valid index: {name}" + ); + } + index => index - 1, + }; let field = fields.get(idx).ok_or_else(|| { plan_datafusion_err!( "Placeholder ${} refers to a non existent column", idx + 1 ) })?; - let _ = prepare_param_data_types.insert(name, Arc::clone(field)); + let _ = prepare_param_data_types.insert(index, Arc::clone(field)); } } } @@ -2574,17 +3041,35 @@ impl SqlToRel<'_, S> { "".to_string() }; + // Scalar / aggregate / window functions are resolved by joining + // parameters (IN rows aggregated per OUT row) with routines. + // Table functions (UDTFs) don't have parameter rows, so they are + // sourced directly from routines via a UNION branch. Restricting + // the JOIN to non-TABLE routines prevents same-named scalar+UDTF + // pairs (e.g. `generate_series`) from cross-joining. + let where_clause = where_clause.replace("p.function_name", "sc.function_name"); let query = format!( r#" SELECT DISTINCT - p.*, - r.function_type function_type, - r.description description, - r.syntax_example syntax_example -FROM - ( + sc.function_name, + sc.return_type, + sc.parameters, + sc.parameter_types, + sc.function_type, + sc.description, + sc.syntax_example +FROM ( + SELECT + p.function_name, + p.return_type, + p.parameters, + p.parameter_types, + r.function_type function_type, + r.description description, + r.syntax_example syntax_example + FROM ( SELECT - i.specific_name function_name, + o.specific_name function_name, o.data_type return_type, array_agg(i.parameter_name ORDER BY i.ordinal_position ASC) parameters, array_agg(i.data_type ORDER BY i.ordinal_position ASC) parameter_types @@ -2600,9 +3085,9 @@ FROM FROM information_schema.parameters WHERE - parameter_mode = 'IN' - ) i - JOIN + parameter_mode = 'OUT' + ) o + LEFT JOIN ( SELECT specific_catalog, @@ -2615,16 +3100,32 @@ FROM FROM information_schema.parameters WHERE - parameter_mode = 'OUT' - ) o + parameter_mode = 'IN' + ) i ON i.specific_catalog = o.specific_catalog AND i.specific_schema = o.specific_schema AND i.specific_name = o.specific_name AND i.rid = o.rid - GROUP BY 1, 2, i.rid + GROUP BY 1, 2, o.rid ) as p -JOIN information_schema.routines r -ON p.function_name = r.routine_name + JOIN information_schema.routines r + ON p.function_name = r.routine_name + AND r.function_type <> 'TABLE' + + UNION ALL + + SELECT + routine_name function_name, + data_type return_type, + array_agg(NULL) FILTER (WHERE FALSE) parameters, + array_agg(NULL) FILTER (WHERE FALSE) parameter_types, + function_type, + description, + syntax_example + FROM information_schema.routines + WHERE function_type = 'TABLE' + GROUP BY routine_name, data_type, function_type, description, syntax_example +) sc {where_clause} "# ); diff --git a/datafusion/sql/src/unparser/ast.rs b/datafusion/sql/src/unparser/ast.rs index 4b4e56c40cdc5..7418d0b5b7605 100644 --- a/datafusion/sql/src/unparser/ast.rs +++ b/datafusion/sql/src/unparser/ast.rs @@ -264,6 +264,9 @@ impl SelectBuilder { pub fn pop_from(&mut self) -> Option { self.from.pop() } + pub fn has_selection(&self) -> bool { + self.selection.is_some() + } pub fn lateral_views(&mut self, value: Vec) -> &mut Self { self.lateral_views = value; self @@ -483,6 +486,7 @@ pub struct RelationBuilder { enum TableFactorBuilder { Table(TableRelationBuilder), Derived(DerivedRelationBuilder), + NestedJoin(ast::TableWithJoins, Option), Unnest(UnnestRelationBuilder), Flatten(FlattenRelationBuilder), Empty, @@ -501,6 +505,15 @@ impl RelationBuilder { self } + pub fn nested_join( + &mut self, + value: ast::TableWithJoins, + alias: Option, + ) -> &mut Self { + self.relation = Some(TableFactorBuilder::NestedJoin(value, alias)); + self + } + pub fn unnest(&mut self, value: UnnestRelationBuilder) -> &mut Self { self.relation = Some(TableFactorBuilder::Unnest(value)); self @@ -524,6 +537,9 @@ impl RelationBuilder { Some(TableFactorBuilder::Derived(ref mut rel_builder)) => { rel_builder.alias = value; } + Some(TableFactorBuilder::NestedJoin(_, ref mut alias)) => { + *alias = value; + } Some(TableFactorBuilder::Unnest(ref mut rel_builder)) => { rel_builder.alias = value; } @@ -539,6 +555,12 @@ impl RelationBuilder { Ok(match self.relation { Some(TableFactorBuilder::Table(ref value)) => Some(value.build()?), Some(TableFactorBuilder::Derived(ref value)) => Some(value.build()?), + Some(TableFactorBuilder::NestedJoin(ref table_with_joins, ref alias)) => { + Some(ast::TableFactor::NestedJoin { + table_with_joins: Box::new(table_with_joins.clone()), + alias: alias.clone(), + }) + } Some(TableFactorBuilder::Unnest(ref value)) => Some(value.build()?), Some(TableFactorBuilder::Flatten(ref value)) => Some(value.build()?), Some(TableFactorBuilder::Empty) => None, diff --git a/datafusion/sql/src/unparser/dialect.rs b/datafusion/sql/src/unparser/dialect.rs index d9344622405fc..d7dad04014226 100644 --- a/datafusion/sql/src/unparser/dialect.rs +++ b/datafusion/sql/src/unparser/dialect.rs @@ -94,6 +94,11 @@ pub trait Dialect: Send + Sync { DateFieldExtractStyle::DatePart } + /// The style to use when unparsing DISTINCT FROM style expressions + fn distinct_from_style(&self) -> DistinctFromStyle { + DistinctFromStyle::FullText + } + /// The character length extraction style to use: `CharacterLengthStyle` fn character_length_style(&self) -> CharacterLengthStyle { CharacterLengthStyle::CharacterLength @@ -333,6 +338,15 @@ pub enum CharacterLengthStyle { CharacterLength, } +/// `DistinctFromStyle` to use for unparsing `IsDistinctFrom` and `IsNotDistinctFrom` operators +#[derive(Clone, Copy, PartialEq)] +pub enum DistinctFromStyle { + /// DBMS supports `IS (NOT) DISTINCT FROM` + FullText, + /// DBMS supports equivalent operations via `<=>` and `NOT <=>` + Spaceship, +} + pub struct DefaultDialect {} impl Dialect for DefaultDialect { @@ -385,6 +399,10 @@ impl Dialect for PostgreSqlDialect { ast::DataType::SmallInt(None) } + fn distinct_from_style(&self) -> DistinctFromStyle { + DistinctFromStyle::FullText + } + fn scalar_function_to_sql_overrides( &self, unparser: &Unparser, @@ -416,9 +434,11 @@ impl PostgreSqlDialect { }; Ok(Some(ast::Expr::AnyOp { - left: Box::new(unparser.expr_to_sql(needle)?), + // Recurse through the annotated entry point so the stack-growth + // protection engages on nested arguments; see issue #23056. + left: Box::new(unparser.expr_to_sql_with_nesting(needle)?), compare_op: BinaryOperator::Eq, - right: Box::new(unparser.expr_to_sql(haystack)?), + right: Box::new(unparser.expr_to_sql_with_nesting(haystack)?), is_some: false, })) } @@ -529,6 +549,10 @@ impl Dialect for DuckDBDialect { Ok(None) } + + fn distinct_from_style(&self) -> DistinctFromStyle { + DistinctFromStyle::FullText + } } pub struct MySqlDialect {} @@ -562,6 +586,10 @@ impl Dialect for MySqlDialect { DateFieldExtractStyle::Extract } + fn distinct_from_style(&self) -> DistinctFromStyle { + DistinctFromStyle::Spaceship + } + fn int64_cast_dtype(&self) -> ast::DataType { ast::DataType::Custom(ObjectName::from(vec![Ident::new("SIGNED")]), vec![]) } @@ -619,6 +647,10 @@ impl Dialect for SqliteDialect { CharacterLengthStyle::Length } + fn distinct_from_style(&self) -> DistinctFromStyle { + DistinctFromStyle::FullText + } + fn supports_column_alias_in_table_alias(&self) -> bool { false } diff --git a/datafusion/sql/src/unparser/expr.rs b/datafusion/sql/src/unparser/expr.rs index d7b1c6a3bb6de..9403e15406344 100644 --- a/datafusion/sql/src/unparser/expr.rs +++ b/datafusion/sql/src/unparser/expr.rs @@ -30,7 +30,7 @@ use std::sync::Arc; use std::vec; use super::Unparser; -use super::dialect::IntervalStyle; +use super::dialect::{DistinctFromStyle, IntervalStyle}; use arrow::array::{ ArrayRef, Date32Array, Date64Array, PrimitiveArray, types::{ @@ -94,16 +94,67 @@ const IS: &BinaryOperator = &BinaryOperator::BitwiseAnd; impl Unparser<'_> { pub fn expr_to_sql(&self, expr: &Expr) -> Result { - let mut root_expr = self.expr_to_sql_inner(expr)?; - if self.pretty { - root_expr = self.remove_unnecessary_nesting(root_expr, LOWEST, LOWEST); + // Unparsing recurses once per nesting level. The function-argument and + // dialect scalar-function-override paths cost more per level than the + // default `recursive` red zone, so without raising the minimum stack + // size the stack-growing trampoline engages too late and the OS stack + // overflows on deeply nested expressions (issue #23056). The size + // mirrors the planner's stack-growth usage in `query.rs`. + crate::stack::maybe_grow(|| self.expr_to_sql_with_nesting(expr)) + } + + /// Recursive entry point shared by the public [`Self::expr_to_sql`] and the + /// internal recursion sites (scalar-function arguments, arrays, maps, and + /// dialect scalar-function overrides). + /// + /// This is a stack-growth checkpoint. Internal recursion must call this + /// rather than the public [`Self::expr_to_sql`]: the public entry point + /// would re-enter the public stack-growth boundary on every level. + pub(crate) fn expr_to_sql_with_nesting(&self, expr: &Expr) -> Result { + crate::stack::maybe_grow(|| { + let mut root_expr = self.expr_to_sql_inner(expr)?; + if self.pretty { + root_expr = self.remove_unnecessary_nesting(root_expr, LOWEST, LOWEST); + } + Ok(root_expr) + }) + } + + fn distinct_from_to_sql( + &self, + left: ast::Expr, + right: ast::Expr, + is_distinct: bool, + ) -> Result { + match self.dialect.distinct_from_style() { + DistinctFromStyle::FullText => { + let expr = if is_distinct { + ast::Expr::IsDistinctFrom(Box::new(left), Box::new(right)) + } else { + ast::Expr::IsNotDistinctFrom(Box::new(left), Box::new(right)) + }; + Ok(ast::Expr::Nested(Box::new(expr))) + } + DistinctFromStyle::Spaceship => { + let expr = ast::Expr::Nested(Box::new(ast::Expr::BinaryOp { + left: Box::new(left), + right: Box::new(right), + op: BinaryOperator::Spaceship, + })); + if is_distinct { + Ok(ast::Expr::Nested(Box::new(ast::Expr::UnaryOp { + op: UnaryOperator::Not, + expr: Box::new(expr), + }))) + } else { + Ok(expr) + } + } } - Ok(root_expr) } - #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn expr_to_sql_inner(&self, expr: &Expr) -> Result { - match expr { + crate::stack::maybe_grow(|| match expr { Expr::InList(InList { expr, list, @@ -155,11 +206,7 @@ impl Unparser<'_> { }) => { let l = self.expr_to_sql_inner(left.as_ref())?; let r = self.expr_to_sql_inner(right.as_ref())?; - - Ok(ast::Expr::Nested(Box::new(ast::Expr::IsDistinctFrom( - Box::new(l), - Box::new(r), - )))) + self.distinct_from_to_sql(l, r, true) } Expr::BinaryExpr(BinaryExpr { left, @@ -168,11 +215,7 @@ impl Unparser<'_> { }) => { let l = self.expr_to_sql_inner(left.as_ref())?; let r = self.expr_to_sql_inner(right.as_ref())?; - - Ok(ast::Expr::Nested(Box::new(ast::Expr::IsNotDistinctFrom( - Box::new(l), - Box::new(r), - )))) + self.distinct_from_to_sql(l, r, false) } Expr::BinaryExpr(BinaryExpr { left, op, right }) => { let l = self.expr_to_sql_inner(left.as_ref())?; @@ -221,7 +264,7 @@ impl Unparser<'_> { } Expr::Cast(Cast { expr, field }) => Ok(self.cast_to_sql(expr, field)?), Expr::Literal(value, _) => Ok(self.scalar_to_sql(value)?), - Expr::Alias(Alias { expr, name: _, .. }) => self.expr_to_sql_inner(expr), + Expr::Alias(Alias { expr, .. }) => self.expr_to_sql_inner(expr), Expr::WindowFunction(window_fun) => { let WindowFunction { fun, @@ -359,20 +402,25 @@ impl Unparser<'_> { .. } = &agg.params; - let args = self.function_args_to_sql(args)?; + let args_to_use; + let within_group; + + // if this is a WITHIN GROUP aggregate, skip the prepended arg + if agg.func.supports_within_group_clause() && !order_by.is_empty() { + args_to_use = self.function_args_to_sql(&args[1..])?; + within_group = order_by + .iter() + .map(|sort_expr| self.sort_to_sql(sort_expr)) + .collect::>>()?; + } else { + args_to_use = self.function_args_to_sql(args)?; + within_group = Vec::new(); + } + let filter = match filter { Some(filter) => Some(Box::new(self.expr_to_sql_inner(filter)?)), None => None, }; - let within_group: Vec = - if agg.func.supports_within_group_clause() { - order_by - .iter() - .map(|sort_expr| self.sort_to_sql(sort_expr)) - .collect::>>()? - } else { - Vec::new() - }; Ok(ast::Expr::Function(Function { name: ObjectName::from(vec![Ident { value: func_name.to_string(), @@ -382,7 +430,7 @@ impl Unparser<'_> { args: ast::FunctionArguments::List(ast::FunctionArgumentList { duplicate_treatment: distinct .then_some(DuplicateTreatment::Distinct), - args, + args: args_to_use, clauses: vec![], }), filter, @@ -612,7 +660,7 @@ impl Unparser<'_> { Expr::LambdaVariable(l) => Ok(ast::Expr::Identifier( self.new_ident_quoted_if_needs(l.name.clone()), )), - } + }) } pub fn scalar_function_to_sql( @@ -660,7 +708,7 @@ impl Unparser<'_> { fn make_array_to_sql(&self, args: &[Expr]) -> Result { let args = args .iter() - .map(|e| self.expr_to_sql(e)) + .map(|e| self.expr_to_sql_with_nesting(e)) .collect::>>()?; Ok(ast::Expr::Array(Array { elem: args, @@ -687,8 +735,8 @@ impl Unparser<'_> { 2, "array_element must have exactly 2 arguments" ); - let array = self.expr_to_sql(&args[0])?; - let index = self.expr_to_sql(&args[1])?; + let array = self.expr_to_sql_with_nesting(&args[0])?; + let index = self.expr_to_sql_with_nesting(&args[1])?; Ok(ast::Expr::CompoundFieldAccess { root: Box::new(array), access_chain: vec![ast::AccessExpr::Subscript(Subscript::Index { index })], @@ -711,7 +759,7 @@ impl Unparser<'_> { Ok(ast::DictionaryField { key, - value: Box::new(self.expr_to_sql(&chunk[1])?), + value: Box::new(self.expr_to_sql_with_nesting(&chunk[1])?), }) }) .collect::>>()?; @@ -781,7 +829,8 @@ impl Unparser<'_> { fn map_to_sql(&self, args: &[Expr]) -> Result { assert_eq_or_internal_err!(args.len(), 2, "map must have exactly 2 arguments"); - let ast::Expr::Array(Array { elem: keys, .. }) = self.expr_to_sql(&args[0])? + let ast::Expr::Array(Array { elem: keys, .. }) = + self.expr_to_sql_with_nesting(&args[0])? else { return internal_err!( "map expects first argument to be an array, but received: {:?}", @@ -789,7 +838,8 @@ impl Unparser<'_> { ); }; - let ast::Expr::Array(Array { elem: values, .. }) = self.expr_to_sql(&args[1])? + let ast::Expr::Array(Array { elem: values, .. }) = + self.expr_to_sql_with_nesting(&args[1])? else { return internal_err!( "map expects second argument to be an array, but received: {:?}", @@ -923,7 +973,7 @@ impl Unparser<'_> { ) { Ok(ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Wildcard)) } else { - self.expr_to_sql(e) + self.expr_to_sql_with_nesting(e) .map(|e| ast::FunctionArg::Unnamed(ast::FunctionArgExpr::Expr(e))) } }) @@ -974,7 +1024,7 @@ impl Unparser<'_> { left_op: &BinaryOperator, right_op: &BinaryOperator, ) -> ast::Expr { - match expr { + crate::stack::maybe_grow(|| match expr { ast::Expr::Nested(nested) => { let surrounding_precedence = self .sql_op_precedence(left_op) @@ -1027,7 +1077,7 @@ impl Unparser<'_> { self.remove_unnecessary_nesting(*expr, left_op, IS), )), _ => expr, - } + }) } fn inner_precedence(&self, expr: &ast::Expr) -> u8 { @@ -1368,19 +1418,27 @@ impl Unparser<'_> { ScalarValue::Utf8(None) | ScalarValue::Utf8View(None) | ScalarValue::LargeUtf8(None) => Ok(ast::Expr::value(ast::Value::Null)), - ScalarValue::Binary(Some(_)) => not_impl_err!("Unsupported scalar: {v:?}"), - ScalarValue::Binary(None) => Ok(ast::Expr::value(ast::Value::Null)), - ScalarValue::BinaryView(Some(_)) => { - not_impl_err!("Unsupported scalar: {v:?}") - } - ScalarValue::BinaryView(None) => Ok(ast::Expr::value(ast::Value::Null)), - ScalarValue::FixedSizeBinary(..) => { - not_impl_err!("Unsupported scalar: {v:?}") - } - ScalarValue::LargeBinary(Some(_)) => { - not_impl_err!("Unsupported scalar: {v:?}") + ScalarValue::Binary(Some(bin)) + | ScalarValue::BinaryView(Some(bin)) + | ScalarValue::LargeBinary(Some(bin)) + | ScalarValue::FixedSizeBinary(_, Some(bin)) => { + let hex = bin + .iter() + .flat_map(|x| { + const HEX: [char; 16] = [ + '0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'a', 'b', + 'c', 'd', 'e', 'f', + ]; + let (hi, lo) = (((*x >> 4) & 0xfu8), (*x & 0xfu8)); + [HEX[hi as usize], HEX[lo as usize]] + }) + .collect::(); + Ok(ast::Expr::value(ast::Value::HexStringLiteral(hex))) } - ScalarValue::LargeBinary(None) => Ok(ast::Expr::value(ast::Value::Null)), + ScalarValue::Binary(None) + | ScalarValue::BinaryView(None) + | ScalarValue::FixedSizeBinary(_, None) + | ScalarValue::LargeBinary(None) => Ok(ast::Expr::value(ast::Value::Null)), ScalarValue::FixedSizeList(a) => self.scalar_value_list_to_sql(a.values()), ScalarValue::List(a) => self.scalar_value_list_to_sql(a.values()), ScalarValue::LargeList(a) => self.scalar_value_list_to_sql(a.values()), @@ -1900,7 +1958,7 @@ mod tests { use std::ops::{Add, Sub}; use std::{sync::Arc, vec}; - use crate::unparser::dialect::SqliteDialect; + use crate::unparser::dialect::{MySqlDialect, SqliteDialect}; use arrow::array::{LargeListArray, LargeListViewArray, ListArray, ListViewArray}; use arrow::datatypes::{DataType::Int8, Field, Int32Type, Schema, TimeUnit}; use ast::ObjectName; @@ -1908,11 +1966,12 @@ mod tests { use datafusion_common::{Spans, TableReference}; use datafusion_expr::expr::WildcardOptions; use datafusion_expr::{ - ColumnarValue, HigherOrderUDF, LambdaParametersProgress, ScalarFunctionArgs, - ScalarUDF, ScalarUDFImpl, Signature, ValueOrLambda, Volatility, WindowFrame, - WindowFunctionDefinition, case, cast, col, cube, exists, grouping_set, - interval_datetime_lit, interval_year_month_lit, lambda, lambda_var, lit, not, - not_exists, out_ref_col, placeholder, rollup, table_scan, try_cast, when, + ColumnarValue, HigherOrderUDF, HigherOrderUDFImpl, LambdaParametersProgress, + ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, ValueOrLambda, + Volatility, WindowFrame, WindowFunctionDefinition, case, cast, col, cube, exists, + grouping_set, interval_datetime_lit, interval_year_month_lit, lambda, lambda_var, + lit, not, not_exists, out_ref_col, placeholder, rollup, table_scan, try_cast, + when, }; use datafusion_expr::{ExprFunctionExt, interval_month_day_nano_lit}; use datafusion_functions::datetime::from_unixtime::FromUnixtimeFunc; @@ -1969,7 +2028,7 @@ mod tests { #[derive(Debug, Hash, Eq, PartialEq)] struct DummyHigherOrderUDF; - impl HigherOrderUDF for DummyHigherOrderUDF { + impl HigherOrderUDFImpl for DummyHigherOrderUDF { fn name(&self) -> &str { "dummy_higher_order_function" } @@ -2087,7 +2146,7 @@ mod tests { ), ( Expr::HigherOrderFunction(HigherOrderFunction::new( - Arc::new(DummyHigherOrderUDF), + Arc::new(HigherOrderUDF::new_from_impl(DummyHigherOrderUDF)), vec![col("a"), lambda(["v"], -lambda_var("v"))], )), r#"dummy_higher_order_function(a, (v) -> -v)"#, @@ -2393,6 +2452,7 @@ mod tests { name: "array_col".to_string(), spans: Spans::new(), })), + outer: false, }), r#"UNNEST("table".array_col)"#, ), @@ -3277,6 +3337,110 @@ mod tests { Ok(()) } + /// Regression test for https://github.com/apache/datafusion/issues/23056 + /// + /// Deeply-nested expressions whose unparse path routes through scalar + /// function arguments and dialect scalar-function overrides used to + /// overflow the OS stack even with `recursive_protection` enabled, + /// because the per-level stack cost of those paths exceeds the default + /// `recursive` red zone and the unparser installed no [`StackGuard`]. + /// + /// This test only asserts the protected behavior, so it is gated on the + /// `recursive_protection` feature. Without that feature the unparser is + /// not stack-safe by design and a deep enough expression will overflow. + #[cfg(feature = "recursive_protection")] + #[test] + fn test_deeply_nested_expr_does_not_overflow_stack() { + // Far deeper than the ~60 levels that overflow without protection, but + // bounded so the trampoline's heap stacks stay reasonable in debug. + const DEPTH: usize = 2_000; + + // Run on an explicit, realistically-sized thread stack. The work is + // performed on a spawned thread so an overflow (in the unfixed code) + // aborts the process and fails the test deterministically rather than + // depending on the harness thread's stack size. + let handle = std::thread::Builder::new() + .stack_size(2 * 1024 * 1024) + .spawn(|| { + // 1. Linear chain through a dialect scalar-function override: + // array_has(array_has(... array_has(col, 'x') ...), 'x'). + // PostgreSqlDialect unparses array_has via array_has_to_sql_any, + // which recurses back into the unparser for each argument. + let mut nested_fn: Expr = col("c"); + for _ in 0..DEPTH { + nested_fn = array_has(nested_fn, lit("x")); + } + let pg = PostgreSqlDialect {}; + Unparser::new(&pg) + .expr_to_sql(&nested_fn) + .expect("deeply nested scalar function should unparse"); + + // 2. Linear chain of plain binary operators, exercising the + // inner -> inner recursion on the default dialect. + let mut nested_binary: Expr = col("c"); + for _ in 0..DEPTH { + nested_binary = nested_binary + lit(1); + } + Unparser::default() + .expr_to_sql(&nested_binary) + .expect("deeply nested binary expression should unparse"); + + // 3. Same binary chain in pretty mode. Pretty mode runs + // `remove_unnecessary_nesting` at every level, which recurses + // alongside the unparse itself; this locks down that second + // recursion site fixed by this PR. + Unparser::default() + .with_pretty(true) + .expr_to_sql(&nested_binary) + .expect( + "deeply nested binary expression should unparse in pretty mode", + ); + }) + .unwrap(); + + // If the unparser overflows, the process aborts and this join is never + // reached; otherwise the spawned thread returns cleanly. + handle.join().expect("unparsing thread should not panic"); + } + + #[cfg(feature = "recursive_protection")] + #[test] + fn test_expr_to_sql_does_not_mutate_recursive_minimum_stack_size() -> Result<()> { + const DEFAULT_RECURSIVE_RED_ZONE: usize = 128 * 1024; + + let previous_minimum = recursive::get_minimum_stack_size(); + recursive::set_minimum_stack_size(DEFAULT_RECURSIVE_RED_ZONE); + + let observed_minimum = Arc::new(std::sync::atomic::AtomicUsize::new(usize::MAX)); + let dialect = DuckDBDialect::new().with_custom_scalar_overrides(vec![( + "dummy_udf", + Box::new({ + let observed_minimum = Arc::clone(&observed_minimum); + move |unparser: &Unparser, args: &[Expr]| { + observed_minimum.store( + recursive::get_minimum_stack_size(), + std::sync::atomic::Ordering::Relaxed, + ); + unparser.scalar_function_to_sql("dummy_udf", args).map(Some) + } + }) as ScalarFnToSqlHandler, + )]); + let expr = ScalarUDF::new_from_impl(DummyUDF::new()).call(vec![col("a")]); + + let result = Unparser::new(&dialect).expr_to_sql(&expr); + let final_minimum = recursive::get_minimum_stack_size(); + recursive::set_minimum_stack_size(previous_minimum); + + result?; + assert_eq!( + observed_minimum.load(std::sync::atomic::Ordering::Relaxed), + DEFAULT_RECURSIVE_RED_ZONE + ); + assert_eq!(final_minimum, DEFAULT_RECURSIVE_RED_ZONE); + + Ok(()) + } + #[test] fn test_window_func_support_window_frame() -> Result<()> { let default_dialect: Arc = @@ -3705,6 +3869,8 @@ mod tests { #[test] fn test_is_distinct_from() { + let mysql_unparser = Unparser::new(&MySqlDialect {}); + let expr = Expr::BinaryExpr(BinaryExpr::new( Box::new(col("c1")), Operator::IsDistinctFrom, @@ -3713,6 +3879,8 @@ mod tests { let sql = expr_to_sql(&expr).unwrap().to_string(); assert_eq!(sql, "(c1 IS DISTINCT FROM true)"); + let sql = mysql_unparser.expr_to_sql(&expr).unwrap().to_string(); + assert_eq!(sql, "(NOT (`c1` <=> true))"); let expr = Expr::BinaryExpr(BinaryExpr::new( Box::new(col("c1")), @@ -3722,5 +3890,75 @@ mod tests { let sql = expr_to_sql(&expr).unwrap().to_string(); assert_eq!(sql, "(c1 IS NOT DISTINCT FROM true)"); + let sql = mysql_unparser.expr_to_sql(&expr).unwrap().to_string(); + assert_eq!(sql, "(`c1` <=> true)"); + } + + #[test] + fn test_binary_literal() { + let value = vec![0xDEu8, 0xAD, 0xBE, 0xEF]; + let expected_hex = "X'deadbeef'"; + + assert_eq!( + expr_to_sql(&Expr::Literal( + ScalarValue::Binary(Some(value.clone())), + None + )) + .unwrap() + .to_string(), + expected_hex + ); + assert_eq!( + expr_to_sql(&Expr::Literal( + ScalarValue::BinaryView(Some(value.clone())), + None + )) + .unwrap() + .to_string(), + expected_hex + ); + assert_eq!( + expr_to_sql(&Expr::Literal( + ScalarValue::FixedSizeBinary(4, Some(value.clone())), + None + )) + .unwrap() + .to_string(), + expected_hex + ); + assert_eq!( + expr_to_sql(&Expr::Literal( + ScalarValue::LargeBinary(Some(value.clone())), + None + )) + .unwrap() + .to_string(), + expected_hex + ); + + assert_eq!( + expr_to_sql(&Expr::Literal(ScalarValue::Binary(None), None)) + .unwrap() + .to_string(), + "NULL" + ); + assert_eq!( + expr_to_sql(&Expr::Literal(ScalarValue::BinaryView(None), None)) + .unwrap() + .to_string(), + "NULL" + ); + assert_eq!( + expr_to_sql(&Expr::Literal(ScalarValue::FixedSizeBinary(1, None), None)) + .unwrap() + .to_string(), + "NULL" + ); + assert_eq!( + expr_to_sql(&Expr::Literal(ScalarValue::LargeBinary(None), None)) + .unwrap() + .to_string(), + "NULL" + ); } } diff --git a/datafusion/sql/src/unparser/plan.rs b/datafusion/sql/src/unparser/plan.rs index 6697b4ed748ae..f4b60176cfba9 100644 --- a/datafusion/sql/src/unparser/plan.rs +++ b/datafusion/sql/src/unparser/plan.rs @@ -42,15 +42,16 @@ use crate::unparser::{ }; use crate::utils::UNNEST_PLACEHOLDER; use datafusion_common::{ - Column, DataFusionError, Result, ScalarValue, TableReference, assert_or_internal_err, - internal_datafusion_err, internal_err, not_impl_err, + Column, DFSchema, DataFusionError, Result, ScalarValue, TableReference, + assert_or_internal_err, internal_datafusion_err, internal_err, not_impl_err, tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRecursion}, + utils::combine_limit, }; use datafusion_expr::expr::{OUTER_REFERENCE_COLUMN_PREFIX, UNNEST_COLUMN_PREFIX}; use datafusion_expr::{ - Aggregate, BinaryExpr, Distinct, Expr, JoinConstraint, JoinType, LogicalPlan, - LogicalPlanBuilder, Operator, Projection, SortExpr, TableScan, Unnest, - UserDefinedLogicalNode, Window, expr::Alias, + Aggregate, BinaryExpr, Distinct, Expr, FetchType, JoinConstraint, JoinType, + LogicalPlan, LogicalPlanBuilder, Operator, Projection, SkipType, Sort, SortExpr, + TableScan, Unnest, UserDefinedLogicalNode, Window, expr::Alias, }; use sqlparser::ast::{self, Ident, OrderByKind, SetExpr, TableAliasColumnDef}; use std::{sync::Arc, vec}; @@ -100,6 +101,69 @@ pub fn plan_to_sql(plan: &LogicalPlan) -> Result { unparser.plan_to_sql(plan) } +/// Aggregate-expression scope for one rendered SELECT block. +/// +/// When an aggregate's input is itself emitted as a derived subquery (a +/// projection sits between the aggregate and its relation), the input columns +/// are only reachable by that derived table's output names. Base-table +/// qualifiers like `t.col` name a relation that is out of scope above the +/// boundary, so emitting them produces SQL a strict engine rejects. +/// +/// Every clause that renders an aggregate expression (SELECT / GROUP BY / +/// HAVING / QUALIFY / ORDER BY) has to apply the same rule. Detect the +/// boundary once here and reuse it, so the clauses can't drift apart (which is +/// how earlier fixes left some clauses correct and others not). +struct UnparserAggScope<'a> { + agg: &'a Aggregate, + /// `agg.input` renders as a derived projection, so out-of-scope qualifiers + /// must be stripped from expressions in this scope. + input_is_derived_projection: bool, +} + +impl<'a> UnparserAggScope<'a> { + fn new(agg: &'a Aggregate) -> Self { + Self { + agg, + input_is_derived_projection: Unparser::contains_projection_before_relation( + agg.input.as_ref(), + ), + } + } + + /// Prepare a projected column or predicate that still references the + /// aggregate by its output columns: unproject it back onto the aggregate + /// (and `windows`) expressions, then normalize it for this scope. + fn prepare(&self, expr: Expr, windows: Option<&[&Window]>) -> Result { + self.normalize(unproject_agg_exprs(expr, self.agg, windows)?) + } + + /// Normalize an expression that is already in aggregate form (group / aggr + /// exprs, or an unprojected sort expr): strip the qualifiers that fall out + /// of scope once the input is a derived projection. No-op otherwise. + fn normalize(&self, expr: Expr) -> Result { + if self.input_is_derived_projection { + Unparser::strip_column_qualifiers_for_schema( + expr, + self.agg.input.schema().as_ref(), + ) + } else { + Ok(expr) + } + } + + /// Unproject a sort expression onto this aggregate, then normalize it so + /// ORDER BY uses the same scope as the other clauses. + fn prepare_sort_expr( + &self, + sort_expr: SortExpr, + input: &LogicalPlan, + ) -> Result { + let mut sort_expr = unproject_sort_expr(sort_expr, Some(self.agg), input)?; + sort_expr.expr = self.normalize(sort_expr.expr)?; + Ok(sort_expr) + } +} + impl Unparser<'_> { pub fn plan_to_sql(&self, plan: &LogicalPlan) -> Result { let mut plan = normalize_union_schema(plan)?; @@ -226,15 +290,29 @@ impl Unparser<'_> { Ok(SetExpr::Select(Box::new(select_builder.build()?))) } - /// Reconstructs a SELECT SQL statement from a logical plan by unprojecting column expressions - /// found in a [Projection] node. This requires scanning the plan tree for relevant Aggregate - /// and Window nodes and matching column expressions to the appropriate agg or window expressions. + /// Reconstructs a SELECT SQL statement from a logical plan by + /// unprojecting column expressions found in a [Projection] node. This + /// requires scanning the plan tree for relevant Aggregate and Window + /// nodes and matching column expressions to the appropriate agg or + /// window expressions. + /// + /// `fully_absorbed` reports whether the Projection arm was able to + /// absorb every `Sort`/`Limit` node between this Projection and the + /// Aggregate/Window into the current SELECT. When `false`, the + /// Aggregate/Window will end up in a derived subquery, so we fall + /// back to passthrough column references that resolve against that + /// subquery's output instead of unprojecting onto the original + /// aggregate expressions. + /// + /// Returns `true` if an Aggregate node was found and claimed for this + /// SELECT. fn reconstruct_select_statement( &self, plan: &LogicalPlan, p: &Projection, select: &mut SelectBuilder, - ) -> Result<()> { + fully_absorbed: bool, + ) -> Result { let mut exprs = p.expr.clone(); // If an Unnest node is found within the select, find and unproject the unnest column @@ -277,16 +355,32 @@ impl Unparser<'_> { .collect::>>()?; } - match ( - find_agg_node_within_select(plan, true), - find_window_nodes_within_select(plan, None, true), - ) { + // When some Sort/Limit nodes between this Projection and the + // Aggregate/Window couldn't be absorbed into the current SELECT, + // the Aggregate/Window will live inside a derived subquery. In + // that case we use the passthrough projection path — column refs + // resolve against the derived subquery's output columns instead + // of being unprojected onto the original aggregate/window + // expressions. + let agg = if fully_absorbed { + find_agg_node_within_select(plan, true) + } else { + None + }; + let window = if fully_absorbed { + find_window_nodes_within_select(plan, None, true) + } else { + None + }; + match (agg, window) { (Some(agg), window) => { let window_option = window.as_deref(); + let unparser_agg_scope = UnparserAggScope::new(agg); let items = exprs .into_iter() .map(|proj_expr| { - let unproj = unproject_agg_exprs(proj_expr, agg, window_option)?; + let unproj = + unparser_agg_scope.prepare(proj_expr, window_option)?; self.select_item_to_sql(&unproj) }) .collect::>>()?; @@ -295,10 +389,14 @@ impl Unparser<'_> { select.group_by(ast::GroupByExpr::Expressions( agg.group_expr .iter() - .map(|expr| self.expr_to_sql(expr)) + .cloned() + .map(|expr| { + self.expr_to_sql(&unparser_agg_scope.normalize(expr)?) + }) .collect::>>()?, vec![], )); + Ok(true) } (None, Some(window)) => { let items = exprs @@ -310,6 +408,7 @@ impl Unparser<'_> { .collect::>>()?; select.projection(items); + Ok(false) } _ => { let items = exprs @@ -328,9 +427,60 @@ impl Unparser<'_> { }) .collect::>>()?; select.projection(items); + Ok(false) } } - Ok(()) + } + + fn contains_projection_before_relation(plan: &LogicalPlan) -> bool { + match plan { + LogicalPlan::Projection(_) => true, + LogicalPlan::TableScan(_) + | LogicalPlan::Subquery(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Join(_) + | LogicalPlan::EmptyRelation(_) + | LogicalPlan::Values(_) => false, + _ => { + let inputs = plan.inputs(); + matches!( + inputs.as_slice(), + [input] if Self::contains_projection_before_relation(input) + ) + } + } + } + + fn contains_aggregate_before_relation(plan: &LogicalPlan) -> bool { + match plan { + LogicalPlan::Aggregate(_) => true, + LogicalPlan::TableScan(_) + | LogicalPlan::Subquery(_) + | LogicalPlan::SubqueryAlias(_) + | LogicalPlan::Join(_) + | LogicalPlan::EmptyRelation(_) + | LogicalPlan::Values(_) => false, + _ => { + let inputs = plan.inputs(); + matches!( + inputs.as_slice(), + [input] if Self::contains_aggregate_before_relation(input) + ) + } + } + } + + /// Unproject a sort expression; normalize it when the sort is above an + /// aggregate, otherwise just unproject (no scope to normalize against). + fn unproject_sort_expr_in_scope( + sort_expr: SortExpr, + agg: Option<&Aggregate>, + input: &LogicalPlan, + ) -> Result { + match agg { + Some(agg) => UnparserAggScope::new(agg).prepare_sort_expr(sort_expr, input), + None => unproject_sort_expr(sort_expr, None, input), + } } fn derive( @@ -438,7 +588,7 @@ impl Unparser<'_> { })); if !select.already_projected() { - self.reconstruct_select_statement(plan, p, select)?; + self.reconstruct_select_statement(plan, p, select, true)?; } if matches!( @@ -496,6 +646,9 @@ impl Unparser<'_> { window_expr .iter() .map(|expr| { + // No normalization: this agg branch is only reachable from a + // hand-built plan. SQL wraps windows in a projection, which + // reconstruct_select_statement handles (and normalizes). let expr = if let Some(agg) = agg { unproject_agg_exprs(expr.clone(), agg, None)? } else { @@ -540,8 +693,9 @@ impl Unparser<'_> { let input_schema = window.input.schema(); let mut alias_rewriter = TableAliasRewriter { - table_schema: input_schema.as_arrow(), + table_schema: input_schema.as_ref(), alias_name: TableReference::bare(input_alias), + rewrite_unqualified: true, }; let window_expr = window .window_expr @@ -552,6 +706,19 @@ impl Unparser<'_> { self.project_window_output(&window_expr, select, None) } + fn extract_join_input_table_scan_filters( + plan: &Arc, + table_scan_filters: &mut Vec, + ) -> Result> { + match try_transform_to_simple_table_scan_with_filters(plan)? { + Some((plan, filters)) => { + table_scan_filters.extend(filters); + Ok(Arc::new(plan)) + } + None => Ok(Arc::clone(plan)), + } + } + #[cfg_attr(feature = "recursive_protection", recursive::recursive)] fn select_to_sql_recursively( &self, @@ -680,8 +847,229 @@ impl Unparser<'_> { if self.dialect.unnest_as_lateral_flatten() { Self::collect_flatten_aliases(p.input.as_ref(), select); } - self.reconstruct_select_statement(plan, p, select)?; - self.select_to_sql_recursively(p.input.as_ref(), query, select, relation) + // Walk down through consecutive Sort/Limit nodes, greedily + // absorbing what can be folded into the SELECT we're + // building around the Aggregate. A single SQL SELECT can + // carry at most one `ORDER BY` (applied before `LIMIT`), + // so the safe shape between us and the Aggregate is + // `Limit* Sort?` (outer→inner). We stop at the first node + // that would violate this; that node becomes the + // subquery boundary, and recursion (seeing + // `already_projected = true`) wraps it in a derived + // relation. If we walk all the way to a non-Sort/non-Limit + // terminator, the entire chain folds into one SELECT. + // + // Stacked Sorts with nothing between them collapse to the + // outermost — the same simplification `EnforceSorting` + // applies on the physical side — but only when no Limit + // has been absorbed since the previous Sort, since the + // inner Sort would otherwise be determining which rows + // the Limit keeps. + // + // The fold is collected here without touching `query` + // (apart from non-literal direct Limits, which don't + // depend on projection form). Once we know whether every + // Sort/Limit was absorbed we can pick the right + // projection form and emit `ORDER BY` with or without + // unprojection. + let mut cur = p.input.as_ref(); + let mut absorbed_sort: Option<&Sort> = None; + let mut combined_skip: usize = 0; + let mut combined_fetch: Option = None; + let mut have_combined_limit = false; + let mut have_direct_limit = false; + let mut have_order_by = false; + loop { + match cur { + LogicalPlan::Limit(limit) => { + if have_order_by { + // Limit-below-Sort: `ORDER BY … LIMIT N` + // would apply the sort first, but the + // logical plan applies the Limit first. + break; + } + let skip_lit = limit.get_skip_type()?; + let fetch_lit = limit.get_fetch_type()?; + match (skip_lit, fetch_lit) { + (SkipType::Literal(s), FetchType::Literal(f)) => { + if have_direct_limit { + break; + } + if have_combined_limit { + // outer = already-accumulated; + // inner = this Limit. Same merge + // rule as the optimizer. + let (cs, cf) = combine_limit( + combined_skip, + combined_fetch, + s, + f, + ); + combined_skip = cs; + combined_fetch = cf; + } else { + combined_skip = s; + combined_fetch = f; + have_combined_limit = true; + } + } + _ => { + if have_combined_limit || have_direct_limit { + // Cannot safely merge a + // non-literal Limit with a prior + // one; let recursion handle it. + break; + } + let Some(query_ref) = query.as_mut() else { + return internal_err!( + "Limit operator only valid in a statement context." + ); + }; + if let Some(fetch) = &limit.fetch { + query_ref.limit(Some(self.expr_to_sql(fetch)?)); + } + if let Some(skip) = &limit.skip { + query_ref.offset(Some(ast::Offset { + rows: ast::OffsetRows::None, + value: self.expr_to_sql(skip)?, + })); + } + have_direct_limit = true; + } + } + cur = limit.input.as_ref(); + } + LogicalPlan::Sort(sort) if sort.fetch.is_some() => { + // `Sort { fetch }` is logically + // `Limit(fetch) -> Sort`. Try to absorb the + // virtual Limit first; only if that succeeds + // do we absorb the Sort. Otherwise we'd + // silently drop the fetch. + let fetch = sort.fetch.expect("guarded above"); + if have_order_by { + // The virtual Limit would sit below an + // already-absorbed outer Sort. + break; + } + if have_direct_limit { + // Cannot combine a literal fetch with a + // non-literal direct Limit; let the + // derived subquery preserve both. + break; + } + if have_combined_limit { + let (cs, cf) = combine_limit( + combined_skip, + combined_fetch, + 0, + Some(fetch), + ); + combined_skip = cs; + combined_fetch = cf; + } else { + combined_skip = 0; + combined_fetch = Some(fetch); + have_combined_limit = true; + } + // Now the Sort itself. We know + // `!have_order_by` from the check above. + absorbed_sort = Some(sort); + have_order_by = true; + cur = sort.input.as_ref(); + } + LogicalPlan::Sort(sort) => { + // Sort without `fetch`. + if have_order_by { + // Outer Sort already absorbed; the inner + // Sort is reordered by it and is + // conventionally dropped, matching + // `EnforceSorting` on the physical side. + cur = sort.input.as_ref(); + continue; + } + absorbed_sort = Some(sort); + have_order_by = true; + cur = sort.input.as_ref(); + } + _ => break, + } + } + + // `fully_absorbed` is the bottom-up algorithm's "walked + // all the way to the terminator without stopping": the + // Aggregate/Window will live in the same SELECT as this + // Projection, so we can unproject sort exprs and let + // `reconstruct_select_statement` claim it. + let fully_absorbed = + !matches!(cur, LogicalPlan::Limit(_) | LogicalPlan::Sort(_)); + let found_agg = + self.reconstruct_select_statement(plan, p, select, fully_absorbed)?; + + // Whether to bother emitting the absorbed clauses: only + // if there's an Aggregate either claimed in this SELECT + // or about to live in a derived subquery below us. If + // there's nothing aggregate-like to fold over, fall + // through and let the normal recursion handle the + // Projection's input. + let agg_below = + !fully_absorbed && find_agg_node_within_select(plan, true).is_some(); + if !(found_agg || agg_below) { + return self.select_to_sql_recursively( + p.input.as_ref(), + query, + select, + relation, + ); + } + + if let Some(sort) = absorbed_sort { + let Some(query_ref) = query.as_mut() else { + return internal_err!( + "Sort operator only valid in a statement context." + ); + }; + let sort_exprs: Vec = if fully_absorbed { + let agg = + find_agg_node_within_select(plan, select.already_projected()); + sort.expr + .iter() + .map(|sort_expr| { + Self::unproject_sort_expr_in_scope( + sort_expr.clone(), + agg, + sort.input.as_ref(), + ) + }) + .collect::>>()? + } else { + sort.expr.clone() + }; + query_ref.order_by(self.sorts_to_sql(&sort_exprs)?); + } + if have_combined_limit { + let Some(query_ref) = query.as_mut() else { + return internal_err!( + "Limit operator only valid in a statement context." + ); + }; + if let Some(fetch) = combined_fetch { + query_ref.limit(Some(ast::Expr::value(ast::Value::Number( + fetch.to_string(), + false, + )))); + } + if combined_skip > 0 { + query_ref.offset(Some(ast::Offset { + rows: ast::OffsetRows::None, + value: ast::Expr::value(ast::Value::Number( + combined_skip.to_string(), + false, + )), + })); + } + } + + self.select_to_sql_recursively(cur, query, select, relation) } LogicalPlan::Filter(filter) => { let window = find_window_nodes_within_select( @@ -697,13 +1085,14 @@ impl Unparser<'_> { let mut unprojected = unproject_window_exprs(filter.predicate.clone(), window)?; if let Some(agg) = agg { - unprojected = unproject_agg_exprs(unprojected, agg, None)?; + unprojected = + UnparserAggScope::new(agg).prepare(unprojected, None)?; } let filter_expr = self.expr_to_sql(&unprojected)?; select.qualify(Some(filter_expr)); } else if let Some(agg) = agg { - let unprojected = - unproject_agg_exprs(filter.predicate.clone(), agg, None)?; + let unprojected = UnparserAggScope::new(agg) + .prepare(filter.predicate.clone(), None)?; let filter_expr = self.expr_to_sql(&unprojected)?; select.having(Some(filter_expr)); } else { @@ -789,7 +1178,11 @@ impl Unparser<'_> { .expr .iter() .map(|sort_expr| { - unproject_sort_expr(sort_expr.clone(), agg, sort.input.as_ref()) + Self::unproject_sort_expr_in_scope( + sort_expr.clone(), + agg, + sort.input.as_ref(), + ) }) .collect::>>()?; @@ -805,23 +1198,38 @@ impl Unparser<'_> { LogicalPlan::Aggregate(agg) => { // Aggregation can be already handled in the projection case if !select.already_projected() { + let unparser_agg_scope = UnparserAggScope::new(agg); // The query returns aggregate and group expressions. If that weren't the case, // the aggregate would have been placed inside a projection, making the check above^ false let exprs: Vec<_> = agg .aggr_expr .iter() .chain(agg.group_expr.iter()) - .map(|expr| self.select_item_to_sql(expr)) + .cloned() + .map(|expr| { + self.select_item_to_sql(&unparser_agg_scope.normalize(expr)?) + }) .collect::>>()?; select.projection(exprs); select.group_by(ast::GroupByExpr::Expressions( agg.group_expr .iter() - .map(|expr| self.expr_to_sql(expr)) + .cloned() + .map(|expr| { + self.expr_to_sql(&unparser_agg_scope.normalize(expr)?) + }) .collect::>>()?, vec![], )); + } else if Self::contains_aggregate_before_relation(agg.input.as_ref()) { + return self.derive_with_dialect_alias( + "derived_aggregate", + agg.input.as_ref(), + relation, + false, + vec![], + ); } self.select_to_sql_recursively( @@ -900,14 +1308,15 @@ impl Unparser<'_> { // The outer projection plan will handle projecting the correct columns. let already_projected = select.already_projected(); - let left_plan = - match try_transform_to_simple_table_scan_with_filters(left_plan)? { - Some((plan, filters)) => { - table_scan_filters.extend(filters); - Arc::new(plan) - } - None => Arc::clone(left_plan), - }; + let left_plan = Self::extract_join_input_table_scan_filters( + left_plan, + &mut table_scan_filters, + )?; + let left_plan = if already_projected { + Self::unwrap_qualified_passthrough_join_projection(left_plan) + } else { + left_plan + }; self.select_to_sql_recursively( left_plan.as_ref(), @@ -923,23 +1332,28 @@ impl Unparser<'_> { None }; - let right_plan = - match try_transform_to_simple_table_scan_with_filters(right_plan)? { - Some((plan, filters)) => { - table_scan_filters.extend(filters); - Arc::new(plan) - } - None => Arc::clone(right_plan), - }; + let right_plan = Self::extract_join_input_table_scan_filters( + right_plan, + &mut table_scan_filters, + )?; let mut right_relation = RelationBuilder::default(); - - self.select_to_sql_recursively( - right_plan.as_ref(), - query, - select, - &mut right_relation, - )?; + if already_projected + && let Some(nested_relation) = self + .qualified_passthrough_join_projection_to_nested_relation( + right_plan.as_ref(), + query, + )? + { + right_relation = nested_relation; + } else { + self.select_to_sql_recursively( + right_plan.as_ref(), + query, + select, + &mut right_relation, + )?; + } let (join_filters, where_filters) = Self::split_join_on_and_where_filters( join.join_type, @@ -1633,7 +2047,7 @@ impl Unparser<'_> { let mut flatten = FlattenRelationBuilder::default(); flatten.input_expr(input_expr); - flatten.outer(unnest.options.preserve_nulls); + flatten.outer(unnest.options.preserve_nulls()); Ok(Some(flatten)) } @@ -1658,6 +2072,96 @@ impl Unparser<'_> { ) } + fn is_qualified_passthrough_projection(projection: &Projection) -> bool { + projection + .expr + .iter() + .all(|expr| matches!(expr, Expr::Column(column) if column.relation.is_some())) + } + + fn unwrap_qualified_passthrough_join_projection( + plan: Arc, + ) -> Arc { + if let LogicalPlan::Projection(projection) = plan.as_ref() + && matches!(projection.input.as_ref(), LogicalPlan::Join(_)) + && Self::is_qualified_passthrough_projection(projection) + { + Arc::clone(&projection.input) + } else { + plan + } + } + + fn qualified_passthrough_join_projection_to_nested_relation( + &self, + plan: &LogicalPlan, + query: &mut Option, + ) -> Result> { + let LogicalPlan::Projection(projection) = plan else { + return Ok(None); + }; + if !matches!(projection.input.as_ref(), LogicalPlan::Join(_)) + || !Self::is_qualified_passthrough_projection(projection) + { + return Ok(None); + } + + let original_query = query.clone(); + let mut nested_select = SelectBuilder::default(); + nested_select.push_from(TableWithJoinsBuilder::default()); + let mut nested_relation = RelationBuilder::default(); + self.select_to_sql_recursively( + projection.input.as_ref(), + query, + &mut nested_select, + &mut nested_relation, + )?; + if nested_select.has_selection() { + *query = original_query; + return Ok(None); + } + + let Some(mut nested_from) = nested_select.pop_from() else { + return internal_err!("Failed to build nested join relation"); + }; + nested_from.relation(nested_relation); + let Some(table_with_joins) = nested_from.build()? else { + return internal_err!("Failed to build nested join relation"); + }; + + let mut relation = RelationBuilder::default(); + relation.nested_join(table_with_joins, None); + Ok(Some(relation)) + } + + /// Strip the table qualifier from every column in an expression that must + /// resolve against an unnamed derived table's output columns rather than a + /// deeper table alias that is out of scope at this nesting level. + fn strip_column_qualifiers(expr: Expr) -> Result { + expr.transform(|e| match e { + Expr::Column(mut column) => { + column.relation = None; + Ok(Transformed::yes(Expr::Column(column))) + } + other => Ok(Transformed::no(other)), + }) + .data() + } + + fn strip_column_qualifiers_for_schema(expr: Expr, schema: &DFSchema) -> Result { + expr.transform(|e| match e { + Expr::Column(mut column) + if column.relation.is_some() + && schema.index_of_column(&column).is_ok() => + { + column.relation = None; + Ok(Transformed::yes(Expr::Column(column))) + } + other => Ok(Transformed::no(other)), + }) + .data() + } + /// Try to unparse a table scan with pushdown operations into a new subquery plan. /// If the table scan is without any pushdown operations, return None. fn unparse_table_scan_pushdown( @@ -1672,10 +2176,15 @@ impl Unparser<'_> { return Ok(None); } let table_schema = table_scan.source.schema(); + let filter_schema = DFSchema::try_from_qualified_schema( + table_scan.table_name.clone(), + table_schema.as_ref(), + )?; let mut filter_alias_rewriter = alias.as_ref().map(|alias_name| TableAliasRewriter { - table_schema: &table_schema, + table_schema: &filter_schema, alias_name: alias_name.clone(), + rewrite_unqualified: true, }); let mut builder = LogicalPlanBuilder::scan( @@ -1782,11 +2291,39 @@ impl Unparser<'_> { alias.clone(), already_projected, )? { + // The pushed-down scan alias is only in scope for the + // projection directly above the aliased table scan. `plan` + // is the result of pushing the alias further down: if it is + // itself a `Projection`, the input was another projection + // (e.g. common subexpression elimination stacked one), so + // this projection sits over a derived table rather than + // directly over the aliased scan, and the alias is out of + // scope here. Its qualified pass-through columns must then + // reference the derived table's output unqualified instead + // of being rebased to the alias. Build it directly so the + // unqualified columns are not re-normalized back to the + // alias. (Otherwise `plan` is the scan-derived plan and we + // fall through to rebase to the alias, correct one level + // above the scan.) + if alias.is_some() && matches!(plan, LogicalPlan::Projection(_)) { + let exprs = projection + .expr + .iter() + .cloned() + .map(Self::strip_column_qualifiers) + .collect::>>()?; + return Ok(Some(LogicalPlan::Projection(Projection::try_new( + exprs, + Arc::new(plan), + )?))); + } + let exprs = if alias.is_some() { let mut alias_rewriter = alias.as_ref().map(|alias_name| TableAliasRewriter { - table_schema: plan.schema().as_arrow(), + table_schema: plan.schema().as_ref(), alias_name: alias_name.clone(), + rewrite_unqualified: false, }); projection .expr diff --git a/datafusion/sql/src/unparser/rewrite.rs b/datafusion/sql/src/unparser/rewrite.rs index a6bfba4cca7af..6ee66f61938f0 100644 --- a/datafusion/sql/src/unparser/rewrite.rs +++ b/datafusion/sql/src/unparser/rewrite.rs @@ -17,10 +17,9 @@ use std::{collections::HashSet, sync::Arc}; -use arrow::datatypes::Schema; use datafusion_common::tree_node::TreeNodeContainer; use datafusion_common::{ - Column, HashMap, Result, TableReference, + Column, DFSchema, HashMap, Result, TableReference, tree_node::{Transformed, TransformedResult, TreeNode, TreeNodeRewriter}, }; use datafusion_expr::expr::{Alias, UNNEST_COLUMN_PREFIX}; @@ -502,20 +501,24 @@ fn find_projection(logical_plan: &LogicalPlan) -> Option<&Projection> { } /// A `TreeNodeRewriter` implementation that rewrites `Expr::Column` expressions by -/// replacing the column's name with an alias if the column exists in the provided schema. +/// replacing the column's qualifier with an alias if the column resolves to a +/// qualified field in the provided schema. /// /// This is typically used to apply table aliases in query plans, ensuring that /// the column references in the expressions use the correct table alias. /// /// # Fields /// -/// * `table_schema`: The schema (`SchemaRef`) representing the table structure -/// from which the columns are referenced. This is used to look up columns by their names. +/// * `table_schema`: The schema representing the table structure from which the +/// columns are referenced. This is used to look up columns by their names and qualifiers. /// * `alias_name`: The alias (`TableReference`) that will replace the table name /// in the column references when applicable. +/// * `rewrite_unqualified`: Whether columns that resolve to unqualified fields +/// in `table_schema` should also be rewritten to `alias_name`. pub struct TableAliasRewriter<'a> { - pub table_schema: &'a Schema, + pub table_schema: &'a DFSchema, pub alias_name: TableReference, + pub rewrite_unqualified: bool, } impl TreeNodeRewriter for TableAliasRewriter<'_> { @@ -524,12 +527,23 @@ impl TreeNodeRewriter for TableAliasRewriter<'_> { fn f_down(&mut self, expr: Expr) -> Result> { match expr { Expr::Column(column) => { - if let Ok(field) = self.table_schema.field_with_name(&column.name) { - let new_column = - Column::new(Some(self.alias_name.clone()), field.name().clone()); - Ok(Transformed::yes(Expr::Column(new_column))) - } else { - Ok(Transformed::no(Expr::Column(column))) + match self + .table_schema + .qualified_field_from_column(&column) + .or_else(|_| { + self.table_schema + .qualified_field_with_unqualified_name(&column.name) + }) { + Ok((qualifier, field)) + if qualifier.is_some() || self.rewrite_unqualified => + { + let new_column = Column::new( + Some(self.alias_name.clone()), + field.name().clone(), + ); + Ok(Transformed::yes(Expr::Column(new_column))) + } + Ok(_) | Err(_) => Ok(Transformed::no(Expr::Column(column))), } } _ => Ok(Transformed::no(expr)), @@ -540,7 +554,7 @@ impl TreeNodeRewriter for TableAliasRewriter<'_> { #[cfg(test)] mod tests { use super::*; - use arrow::datatypes::{DataType, Field}; + use arrow::datatypes::{DataType, Field, Schema}; use datafusion_expr::{LogicalPlanBuilder, col, table_scan}; // this is a regression test: when the outer projection has fewer expressions than diff --git a/datafusion/sql/src/unparser/utils.rs b/datafusion/sql/src/unparser/utils.rs index 732e030b335d8..949b49eb77be9 100644 --- a/datafusion/sql/src/unparser/utils.rs +++ b/datafusion/sql/src/unparser/utils.rs @@ -22,7 +22,7 @@ use super::{ rewrite::TableAliasRewriter, }; use datafusion_common::{ - Column, DataFusionError, Result, ScalarValue, TableReference, + Column, DFSchema, DataFusionError, Result, ScalarValue, TableReference, assert_eq_or_internal_err, internal_err, tree_node::{Transformed, TransformedResult, TreeNode}, }; @@ -389,11 +389,16 @@ pub(crate) fn try_transform_to_simple_table_scan_with_filters( } LogicalPlan::TableScan(table_scan) => { let table_schema = table_scan.source.schema(); + let filter_schema = DFSchema::try_from_qualified_schema( + table_scan.table_name.clone(), + table_schema.as_ref(), + )?; // optional rewriter if table has an alias let mut filter_alias_rewriter = table_alias.as_ref().map(|alias_name| TableAliasRewriter { - table_schema: &table_schema, + table_schema: &filter_schema, alias_name: alias_name.clone(), + rewrite_unqualified: true, }); // rewrite filters to use table alias if present @@ -448,7 +453,7 @@ pub(crate) fn date_part_to_sql( ) -> Result> { match (style, date_part_args.len()) { (DateFieldExtractStyle::Extract, 2) => { - let date_expr = unparser.expr_to_sql(&date_part_args[1])?; + let date_expr = unparser.expr_to_sql_with_nesting(&date_part_args[1])?; if let Expr::Literal(ScalarValue::Utf8(Some(field)), _) = &date_part_args[0] { let field = match field.to_lowercase().as_str() { "year" => ast::DateTimeField::Year, @@ -468,7 +473,7 @@ pub(crate) fn date_part_to_sql( } } (DateFieldExtractStyle::Strftime, 2) => { - let column = unparser.expr_to_sql(&date_part_args[1])?; + let column = unparser.expr_to_sql_with_nesting(&date_part_args[1])?; if let Expr::Literal(ScalarValue::Utf8(Some(field)), _) = &date_part_args[0] { let field = match field.to_lowercase().as_str() { diff --git a/datafusion/sql/src/utils.rs b/datafusion/sql/src/utils.rs index 1a76dd69f46c5..3b571eed279dd 100644 --- a/datafusion/sql/src/utils.rs +++ b/datafusion/sql/src/utils.rs @@ -35,7 +35,7 @@ use datafusion_expr::expr::{ }; use datafusion_expr::utils::{expr_as_column_expr, find_column_exprs}; use datafusion_expr::{ - ColumnUnnestList, Expr, ExprSchemable, LogicalPlan, col, expr_vec_fmt, + ColumnUnnestList, Expr, ExprSchemable, LogicalPlan, SortExpr, col, expr_vec_fmt, }; use indexmap::IndexMap; @@ -98,6 +98,7 @@ pub(crate) enum CheckColumnsMustReferenceAggregatePurpose { Having, Qualify, OrderBy, + DistinctOn, } #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -120,6 +121,9 @@ impl CheckColumnsSatisfyExprsPurpose { Self::Aggregate(CheckColumnsMustReferenceAggregatePurpose::OrderBy) => { "Column in ORDER BY must be in GROUP BY or an aggregate function" } + Self::Aggregate(CheckColumnsMustReferenceAggregatePurpose::DistinctOn) => { + "Column in DISTINCT ON must be in GROUP BY or an aggregate function" + } } } @@ -202,6 +206,42 @@ pub(crate) fn extract_aliases(exprs: &[Expr]) -> HashMap { .collect::>() } +/// If `expr` is a bare unqualified `Column` whose name matches a SELECT +/// alias, swap it for the alias's underlying expression. Nested occurrences +/// are left alone: PostgreSQL only resolves a top-level identifier as an +/// output alias in clauses like ORDER BY and DISTINCT ON. +pub(crate) fn substitute_top_level_alias( + expr: Expr, + aliases: &HashMap, +) -> Expr { + if let Expr::Column(col) = &expr + && col.relation.is_none() + && let Some(underlying) = aliases.get(&col.name) + { + return underlying.clone(); + } + + expr +} + +/// Applies [`substitute_top_level_alias`] to each sort expression. +pub(crate) fn substitute_top_level_aliases_in_sorts( + sort_exprs: Vec, + aliases: &HashMap, +) -> Vec { + if aliases.is_empty() { + return sort_exprs; + } + + sort_exprs + .into_iter() + .map(|sort_expr| { + sort_expr + .with_expr(substitute_top_level_alias(sort_expr.expr.clone(), aliases)) + }) + .collect() +} + /// Given an expression that's literal int encoding position, lookup the corresponding expression /// in the select_exprs list, if the index is within the bounds and it is indeed a position literal, /// otherwise, returns planning error. diff --git a/datafusion/sql/tests/cases/diagnostic.rs b/datafusion/sql/tests/cases/diagnostic.rs index 7a729739469d3..1f2cefdec0629 100644 --- a/datafusion/sql/tests/cases/diagnostic.rs +++ b/datafusion/sql/tests/cases/diagnostic.rs @@ -16,13 +16,19 @@ // under the License. use datafusion_functions::string; +use datafusion_functions_aggregate::sum::sum_udaf; use insta::assert_snapshot; -use std::{collections::HashMap, sync::Arc}; +use std::{collections::HashMap, ops::ControlFlow, sync::Arc}; +use datafusion_common::diagnostic::DiagnosticKind; use datafusion_common::{Diagnostic, Location, Result, Span}; use datafusion_sql::{ - parser::{DFParser, DFParserBuilder}, + parser::{DFParser, DFParserBuilder, Statement as DFStatement}, planner::{ParserOptions, SqlToRel}, + sqlparser::{ + ast::{Expr as SQLExpr, visit_expressions_mut}, + tokenizer::Span as SQLParserSpan, + }, }; use regex::Regex; @@ -39,7 +45,8 @@ fn do_query(sql: &'static str) -> Diagnostic { ..ParserOptions::default() }; let state = MockSessionState::default() - .with_scalar_function(Arc::new(string::concat().as_ref().clone())); + .with_scalar_function(Arc::new(string::concat().as_ref().clone())) + .with_aggregate_function(sum_udaf()); let context = MockContextProvider { state }; let sql_to_rel = SqlToRel::new_with_options(&context, options); match sql_to_rel.statement_to_plan(statement) { @@ -51,6 +58,41 @@ fn do_query(sql: &'static str) -> Diagnostic { } } +fn do_query_warnings(sql: &'static str) -> Vec { + let statement = DFParserBuilder::new(sql) + .build() + .expect("unable to create parser") + .parse_statement() + .expect("unable to parse query"); + do_statement_warnings(statement) +} + +fn do_statement_warnings(statement: DFStatement) -> Vec { + let options = ParserOptions { + collect_spans: true, + ..ParserOptions::default() + }; + let state = MockSessionState::default(); + let context = MockContextProvider { state }; + let sql_to_rel = SqlToRel::new_with_options(&context, options); + sql_to_rel + .statement_to_plan(statement) + .expect("expected planning to succeed"); + sql_to_rel.take_warnings() +} + +fn clear_value_spans(statement: &mut DFStatement) { + let DFStatement::Statement(statement) = statement else { + panic!("expected sqlparser statement"); + }; + let _ = visit_expressions_mut(statement.as_mut(), |expr| { + if let SQLExpr::Value(value) = expr { + value.span = SQLParserSpan::empty(); + } + ControlFlow::<()>::Continue(()) + }); +} + /// Given a query that contains tag delimited spans, returns a mapping from the /// span name to the [`Span`]. Tags are comments of the form `/*tag*/`. In case /// you want the same location to open two spans, or close open and open @@ -369,6 +411,56 @@ fn test_unary_op_plus_with_non_column() -> Result<()> { Ok(()) } +#[test] +fn test_unary_op_minus_with_column() -> Result<()> { + let query = "SELECT -/*whole*/first_name/*whole*/ FROM person"; + let spans = get_spans(query); + let diag = do_query(query); + assert_snapshot!(diag.message, @"- cannot be used with Utf8"); + assert_eq!(diag.span, Some(spans["whole"])); + assert_snapshot!(diag.notes[0].message, @"- can only be used with signed numeric types, intervals, and timestamps"); + assert_snapshot!(diag.helps[0].message, @"perhaps you need to cast person.first_name"); + Ok(()) +} + +#[test] +fn test_unary_op_minus_with_non_column() -> Result<()> { + let query = "SELECT -'a'"; + let diag = do_query(query); + assert_eq!(diag.message, "- cannot be used with Utf8"); + assert_snapshot!(diag.notes[0].message, @"- can only be used with signed numeric types, intervals, and timestamps"); + assert_eq!(diag.notes[0].span, None); + assert_snapshot!(diag.helps[0].message, @r#"perhaps you need to cast Utf8("a")"#); + assert_eq!(diag.helps[0].span, None); + assert_eq!(diag.span, None); + Ok(()) +} + +#[test] +fn test_unary_op_not_with_column() -> Result<()> { + let query = "SELECT NOT /*whole*/first_name/*whole*/ FROM person"; + let spans = get_spans(query); + let diag = do_query(query); + assert_snapshot!(diag.message, @"NOT cannot be used with Utf8"); + assert_eq!(diag.span, Some(spans["whole"])); + assert_snapshot!(diag.notes[0].message, @"NOT can only be used with boolean expressions"); + assert_snapshot!(diag.helps[0].message, @"perhaps you need to cast person.first_name"); + Ok(()) +} + +#[test] +fn test_unary_op_not_with_non_column() -> Result<()> { + let query = "SELECT NOT 'a'"; + let diag = do_query(query); + assert_eq!(diag.message, "NOT cannot be used with Utf8"); + assert_snapshot!(diag.notes[0].message, @"NOT can only be used with boolean expressions"); + assert_eq!(diag.notes[0].span, None); + assert_snapshot!(diag.helps[0].message, @r#"perhaps you need to cast Utf8("a")"#); + assert_eq!(diag.helps[0].span, None); + assert_eq!(diag.span, None); + Ok(()) +} + #[test] fn test_syntax_error() -> Result<()> { // create a table with a column of type varchar @@ -390,3 +482,231 @@ fn test_syntax_error() -> Result<()> { }, } } + +#[test] +fn test_eq_null_warning_in_where() -> Result<()> { + let query = "SELECT * FROM person WHERE /*cmp*/first_name = /*null*/NULL/*null+cmp*/"; + let spans = get_spans(query); + let warnings = do_query_warnings(query); + assert_eq!(warnings.len(), 1); + + let warning = &warnings[0]; + assert_eq!(warning.kind, DiagnosticKind::Warning); + assert_snapshot!( + warning.message, + @"comparison with NULL using `=` always evaluates to NULL" + ); + assert_eq!(warning.span, Some(spans["cmp"])); + assert_snapshot!( + warning.helps[0].message, + @"use `IS NULL` to check for NULL values" + ); + assert_eq!(warning.helps[0].span, Some(spans["null"])); + Ok(()) +} + +#[test] +fn test_null_eq_warning_in_where() -> Result<()> { + let query = "SELECT * FROM person WHERE /*cmp+null*/NULL/*null*/ = first_name/*cmp*/"; + let spans = get_spans(query); + let warnings = do_query_warnings(query); + assert_eq!(warnings.len(), 1); + assert_eq!(warnings[0].kind, DiagnosticKind::Warning); + assert_snapshot!( + warnings[0].message, + @"comparison with NULL using `=` always evaluates to NULL" + ); + assert_eq!(warnings[0].span, Some(spans["cmp"])); + assert_eq!(warnings[0].helps[0].span, Some(spans["null"])); + Ok(()) +} + +#[test] +fn test_not_eq_null_warning_in_where() -> Result<()> { + let query = + "SELECT * FROM person WHERE /*cmp*/first_name <> /*null*/NULL/*null+cmp*/"; + let spans = get_spans(query); + let warnings = do_query_warnings(query); + assert_eq!(warnings.len(), 1); + assert_eq!(warnings[0].kind, DiagnosticKind::Warning); + assert_snapshot!( + warnings[0].message, + @"comparison with NULL using `<>` always evaluates to NULL" + ); + assert_eq!(warnings[0].span, Some(spans["cmp"])); + assert_snapshot!( + warnings[0].helps[0].message, + @"use `IS NOT NULL` to check for non-NULL values" + ); + assert_eq!(warnings[0].helps[0].span, Some(spans["null"])); + Ok(()) +} + +#[test] +fn test_eq_null_warning_in_join_on() -> Result<()> { + let query = + "SELECT * FROM person a JOIN person b ON /*cmp*/a.id = /*null*/NULL/*null+cmp*/"; + let spans = get_spans(query); + let warnings = do_query_warnings(query); + assert_eq!(warnings.len(), 1); + assert_eq!(warnings[0].kind, DiagnosticKind::Warning); + assert_snapshot!( + warnings[0].message, + @"comparison with NULL using `=` always evaluates to NULL" + ); + assert_eq!(warnings[0].span, Some(spans["cmp"])); + Ok(()) +} + +#[test] +fn test_eq_null_warning_in_having() -> Result<()> { + let query = "SELECT first_name FROM person GROUP BY first_name HAVING /*cmp*/1 = /*null*/NULL/*null+cmp*/"; + let spans = get_spans(query); + let warnings = do_query_warnings(query); + assert_eq!(warnings.len(), 1); + assert_eq!(warnings[0].kind, DiagnosticKind::Warning); + assert_snapshot!( + warnings[0].message, + @"comparison with NULL using `=` always evaluates to NULL" + ); + assert_eq!(warnings[0].span, Some(spans["cmp"])); + Ok(()) +} + +#[test] +fn test_eq_null_warning_nested_in_case_predicate() -> Result<()> { + let query = "SELECT * FROM person WHERE CASE WHEN /*cmp*/first_name = /*null*/NULL/*null+cmp*/ THEN true ELSE false END"; + let spans = get_spans(query); + let warnings = do_query_warnings(query); + assert_eq!(warnings.len(), 1); + assert_eq!(warnings[0].kind, DiagnosticKind::Warning); + assert_eq!(warnings[0].span, Some(spans["cmp"])); + Ok(()) +} + +#[test] +fn test_eq_null_warning_under_is_null_predicate() -> Result<()> { + let query = "SELECT * FROM person WHERE (/*cmp*/first_name = /*null*/NULL/*null+cmp*/) IS NULL"; + let spans = get_spans(query); + let warnings = do_query_warnings(query); + assert_eq!(warnings.len(), 1); + assert_eq!(warnings[0].kind, DiagnosticKind::Warning); + assert_eq!(warnings[0].span, Some(spans["cmp"])); + assert_eq!(warnings[0].helps[0].span, Some(spans["null"])); + Ok(()) +} + +#[test] +fn test_eq_null_warning_without_null_span() -> Result<()> { + let query = "SELECT * FROM person WHERE first_name = NULL"; + let mut statement = DFParserBuilder::new(query) + .build() + .expect("unable to create parser") + .parse_statement() + .expect("unable to parse query"); + clear_value_spans(&mut statement); + + let warnings = do_statement_warnings(statement); + assert_eq!(warnings.len(), 1); + assert_eq!(warnings[0].kind, DiagnosticKind::Warning); + assert_snapshot!( + warnings[0].message, + @"comparison with NULL using `=` always evaluates to NULL" + ); + assert_eq!(warnings[0].helps[0].span, None); + Ok(()) +} + +#[test] +fn test_is_null_has_no_warning() -> Result<()> { + let warnings = do_query_warnings("SELECT * FROM person WHERE first_name IS NULL"); + assert!(warnings.is_empty(), "unexpected warnings: {warnings:?}"); + Ok(()) +} + +#[test] +fn test_eq_null_projection_has_no_warning() -> Result<()> { + let warnings = do_query_warnings("SELECT first_name = NULL FROM person"); + assert!(warnings.is_empty(), "unexpected warnings: {warnings:?}"); + Ok(()) +} + +#[test] +fn test_eq_null_projection_in_exists_has_no_warning() -> Result<()> { + let warnings = do_query_warnings( + "SELECT * FROM person WHERE EXISTS (SELECT first_name = NULL FROM person)", + ); + assert!(warnings.is_empty(), "unexpected warnings: {warnings:?}"); + Ok(()) +} + +#[test] +fn test_eq_null_warning_in_exists_subquery_where() -> Result<()> { + let query = "SELECT * FROM person WHERE EXISTS (SELECT 1 FROM person WHERE /*cmp*/first_name = /*null*/NULL/*null+cmp*/)"; + let spans = get_spans(query); + let warnings = do_query_warnings(query); + assert_eq!(warnings.len(), 1); + assert_eq!(warnings[0].kind, DiagnosticKind::Warning); + assert_snapshot!( + warnings[0].message, + @"comparison with NULL using `=` always evaluates to NULL" + ); + assert_eq!(warnings[0].span, Some(spans["cmp"])); + assert_eq!(warnings[0].helps[0].span, Some(spans["null"])); + Ok(()) +} + +#[test] +fn test_multiple_null_comparison_warnings() -> Result<()> { + let warnings = do_query_warnings( + "SELECT * FROM person WHERE first_name = NULL OR last_name <> NULL", + ); + assert_eq!(warnings.len(), 2); + assert!(warnings.iter().all(|w| w.kind == DiagnosticKind::Warning)); + assert_snapshot!( + warnings[0].message, + @"comparison with NULL using `=` always evaluates to NULL" + ); + assert_snapshot!( + warnings[1].message, + @"comparison with NULL using `<>` always evaluates to NULL" + ); + Ok(()) +} + +#[test] +fn test_nested_aggregate() -> Result<()> { + let query = "SELECT sum(sum(/*a*/age/*a*/)) FROM person"; + let spans = get_spans(query); + let diag = do_query(query); + assert_snapshot!(diag.message, @"Aggregate function calls cannot be nested"); + assert_eq!(diag.span, Some(spans["a"])); + assert_snapshot!( + diag.helps[0].message, + @"Compute 'sum(person.age)' in an inner query and aggregate its result" + ); + Ok(()) +} + +#[test] +fn test_window_function_inside_aggregate() -> Result<()> { + let query = "SELECT sum(sum(/*a*/age/*a*/) OVER ()) FROM person"; + let spans = get_spans(query); + let diag = do_query(query); + assert_snapshot!( + diag.message, + @"Aggregate function calls cannot contain window function calls" + ); + assert_eq!(diag.span, Some(spans["a"])); + Ok(()) +} + +#[test] +fn test_nested_window_function() -> Result<()> { + let query = "SELECT sum(sum(/*a*/age/*a*/) OVER ()) OVER () FROM person"; + let spans = get_spans(query); + let diag = do_query(query); + assert_snapshot!(diag.message, @"Window function calls cannot be nested"); + assert_eq!(diag.span, Some(spans["a"])); + Ok(()) +} diff --git a/datafusion/sql/tests/cases/plan_to_sql.rs b/datafusion/sql/tests/cases/plan_to_sql.rs index 62912c7ff86c9..d6c31570bf1b0 100644 --- a/datafusion/sql/tests/cases/plan_to_sql.rs +++ b/datafusion/sql/tests/cases/plan_to_sql.rs @@ -313,6 +313,12 @@ macro_rules! roundtrip_statement_with_dialect_helper { let state = MockSessionState::default() .with_aggregate_function(max_udaf()) .with_aggregate_function(min_udaf()) + .with_aggregate_function( + datafusion_functions_aggregate::approx_percentile_cont::approx_percentile_cont_udaf(), + ) + .with_aggregate_function( + datafusion_functions_aggregate::percentile_cont::percentile_cont_udaf(), + ) .with_expr_planner(Arc::new(CoreFunctionPlanner::default())) .with_expr_planner(Arc::new(NestedFunctionPlanner)) .with_expr_planner(Arc::new(FieldAccessPlanner)); @@ -2736,6 +2742,116 @@ fn test_unparse_inner_join_with_table_scan_projection() -> Result<()> { Ok(()) } +/// Build the three base table scans (`left_table`, `mid_table`, `right_table`) +/// shared by the nested passthrough-projection join unparsing tests. +fn nested_passthrough_join_tables() -> Result<(LogicalPlan, LogicalPlan, LogicalPlan)> { + let left_schema = Schema::new(vec![ + Field::new("left_id", DataType::Int32, false), + Field::new("mid_id", DataType::Int32, false), + ]); + let mid_schema = Schema::new(vec![ + Field::new("mid_id", DataType::Int32, false), + Field::new("right_id", DataType::Int32, false), + ]); + let right_schema = Schema::new(vec![ + Field::new("right_id", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ]); + + let left = table_scan(Some("left_table"), &left_schema, None)?.build()?; + let mid = table_scan(Some("mid_table"), &mid_schema, None)?.build()?; + let right = table_scan(Some("right_table"), &right_schema, None)?.build()?; + Ok((left, mid, right)) +} + +#[test] +fn test_unparse_projected_join_unwraps_right_nested_passthrough_projection() -> Result<()> +{ + let (left, mid, right) = nested_passthrough_join_tables()?; + + let nested_right = LogicalPlanBuilder::from(mid) + .join( + right, + datafusion_expr::JoinType::Inner, + (vec!["mid_table.right_id"], vec!["right_table.right_id"]), + None, + )? + .project(vec![ + col("mid_table.mid_id"), + col("mid_table.right_id"), + col("right_table.value"), + ])? + .build()?; + + let plan = LogicalPlanBuilder::from(left) + .join( + nested_right, + datafusion_expr::JoinType::Inner, + (vec!["left_table.mid_id"], vec!["mid_table.mid_id"]), + None, + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("right_table.value"), + ])? + .build()?; + + let sql = plan_to_sql(&plan)?; + assert_snapshot!( + sql, + @r#"SELECT left_table.left_id, mid_table.mid_id, right_table."value" FROM left_table INNER JOIN (mid_table INNER JOIN right_table ON mid_table.right_id = right_table.right_id) ON left_table.mid_id = mid_table.mid_id"# + ); + + Ok(()) +} + +#[test] +fn test_unparse_projected_join_unwraps_left_nested_passthrough_projection() -> Result<()> +{ + let (left, mid, right) = nested_passthrough_join_tables()?; + + // Left join input is a qualified passthrough `Projection(Join)`, and the + // outer join condition (`mid_table.right_id`) references an alias from + // inside it. The unparser must not wrap this in a derived table that would + // hide `mid_table.right_id` from the outer condition. + let nested_left = LogicalPlanBuilder::from(left) + .join( + mid, + datafusion_expr::JoinType::Inner, + (vec!["left_table.mid_id"], vec!["mid_table.mid_id"]), + None, + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("mid_table.right_id"), + ])? + .build()?; + + let plan = LogicalPlanBuilder::from(nested_left) + .join( + right, + datafusion_expr::JoinType::Inner, + (vec!["mid_table.right_id"], vec!["right_table.right_id"]), + None, + )? + .project(vec![ + col("left_table.left_id"), + col("mid_table.mid_id"), + col("right_table.value"), + ])? + .build()?; + + let sql = plan_to_sql(&plan)?; + assert_snapshot!( + sql, + @r#"SELECT left_table.left_id, mid_table.mid_id, right_table."value" FROM left_table INNER JOIN mid_table ON left_table.mid_id = mid_table.mid_id INNER JOIN right_table ON mid_table.right_id = right_table.right_id"# + ); + + Ok(()) +} + #[test] fn test_unparse_left_semi_join_with_table_scan_projection() -> Result<()> { let schema = Schema::new(vec![ @@ -3367,6 +3483,332 @@ fn roundtrip_subquery_aggregate_with_column_alias() -> Result<(), DataFusionErro Ok(()) } +/// Roundtrip: aggregate over a subquery projection with limit. +#[test] +fn roundtrip_aggregate_over_subquery() -> Result<(), DataFusionError> { + roundtrip_statement_with_dialect_helper!( + sql: r#"SELECT __agg_0 AS "min(j1_id)", __agg_1 AS "max(j1_id)" FROM (SELECT min(j1_rename) AS __agg_0, max(j1_rename) AS __agg_1 FROM (SELECT j1_id AS j1_rename FROM j1) AS bla LIMIT 20)"#, + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @r#"SELECT __agg_0 AS "min(j1_id)", __agg_1 AS "max(j1_id)" FROM (SELECT min(bla.j1_rename) AS __agg_0, max(bla.j1_rename) AS __agg_1 FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla LIMIT 20)"#, + ); + Ok(()) +} + +/// Projection → Limit → Aggregate (aliases inlined into Aggregate, no +/// intermediate Projection). Verifies the Limit is folded into the outer +/// SELECT rather than creating a spurious derived subquery. +#[test] +fn test_unparse_aggregate_over_subquery_no_inner_proj() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![ + max(col("bla.j1_rename")).alias("__agg_0"), + max(col("bla.j1_rename")).alias("__agg_1"), + ], + )? + .limit(0, Some(20))? + .project(vec![ + col("__agg_0").alias("max1(j1_id)"), + col("__agg_1").alias("max2(j1_id)"), + ])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT max(bla.j1_rename) AS "max1(j1_id)", max(bla.j1_rename) AS "max2(j1_id)" FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla LIMIT 20"#); + Ok(()) +} + +/// Projection → Aggregate (aliases inlined, no rename in outer Projection). +/// Verifies the aggregate aliases are preserved as output column names. +#[test] +fn test_unparse_aggregate_no_outer_rename() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![ + max(col("bla.j1_rename")).alias("__agg_0"), + max(col("bla.j1_rename")).alias("__agg_1"), + ], + )? + .project(vec![col("__agg_0"), col("__agg_1")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @"SELECT max(bla.j1_rename) AS __agg_0, max(bla.j1_rename) AS __agg_1 FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla"); + Ok(()) +} + +/// Projection → Sort → Aggregate (aliases inlined into Aggregate). +/// Verifies the Sort is folded into the outer SELECT rather than creating +/// a spurious derived subquery. +#[test] +fn test_unparse_aggregate_with_sort_no_inner_proj() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![max(col("bla.j1_rename")).alias("__agg_0")], + )? + .sort(vec![col("__agg_0").sort(true, true)])? + .project(vec![col("__agg_0").alias("max1(j1_id)")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT max(bla.j1_rename) AS "max1(j1_id)" FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla ORDER BY max(bla.j1_rename) ASC NULLS FIRST"#); + Ok(()) +} + +/// Projection → Limit → Sort → Aggregate (aliases inlined into Aggregate). +/// The Projection claims the Aggregate through the stacked Limit/Sort; +/// both clauses should fold into the outer SELECT instead of wrapping +/// the Sort in a derived subquery. +#[test] +fn test_unparse_aggregate_with_limit_sort_no_inner_proj() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![max(col("bla.j1_rename")).alias("__agg_0")], + )? + .sort(vec![col("__agg_0").sort(true, true)])? + .limit(0, Some(5))? + .project(vec![col("__agg_0").alias("max1(j1_id)")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT max(bla.j1_rename) AS "max1(j1_id)" FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla ORDER BY max(bla.j1_rename) ASC NULLS FIRST LIMIT 5"#); + Ok(()) +} + +/// Projection → Sort → Limit → Aggregate (aliases inlined into Aggregate). +/// The Sort sits above the Limit — the logical plan applies Limit first +/// and Sort second, which a single `ORDER BY … LIMIT` SELECT cannot +/// express (SQL applies the sort first). The outer Sort folds into the +/// outer SELECT using passthrough column references, while the Limit +/// (and the Aggregate it sits over) goes into a derived subquery. +#[test] +fn test_unparse_aggregate_with_sort_over_limit_no_inner_proj() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![max(col("bla.j1_rename")).alias("__agg_0")], + )? + .limit(0, Some(5))? + .sort(vec![col("__agg_0").sort(true, true)])? + .project(vec![col("__agg_0").alias("max1(j1_id)")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT __agg_0 AS "max1(j1_id)" FROM (SELECT max(bla.j1_rename) AS __agg_0 FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla LIMIT 5) ORDER BY __agg_0 ASC NULLS FIRST"#); + Ok(()) +} + +/// Projection → Limit(10) → Limit(5) → Aggregate. Two stacked Limits +/// merge via `combine_limit` (matching the optimizer's `PushDownLimit`): +/// outer fetch=10, inner fetch=5 → effective fetch=5. +#[test] +fn test_unparse_aggregate_with_repeated_limits_combines() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![max(col("bla.j1_rename")).alias("__agg_0")], + )? + .limit(0, Some(10))? + .limit(0, Some(5))? + .project(vec![col("__agg_0").alias("max1(j1_id)")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT max(bla.j1_rename) AS "max1(j1_id)" FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla LIMIT 5"#); + Ok(()) +} + +/// Projection → Limit(skip=2, fetch=10) → Limit(skip=3, fetch=20) +/// → Aggregate. Two stacked Limits merge via `combine_limit`: combined +/// skip = 3+2=5, combined fetch = min(10, 20-2) = 10. +#[test] +fn test_unparse_aggregate_with_repeated_limits_combines_offset() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![max(col("bla.j1_rename")).alias("__agg_0")], + )? + .limit(3, Some(20))? + .limit(2, Some(10))? + .project(vec![col("__agg_0").alias("max1(j1_id)")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT max(bla.j1_rename) AS "max1(j1_id)" FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla LIMIT 10 OFFSET 5"#); + Ok(()) +} + +/// Projection → Sort(DESC) → Sort(ASC) → Aggregate. Two stacked Sorts +/// fold into a single ORDER BY using the outermost (top) Sort's order; +/// the inner Sort is reordered by the outer one and is therefore +/// redundant. +#[test] +fn test_unparse_aggregate_with_repeated_sorts_keeps_outermost() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![max(col("bla.j1_rename")).alias("__agg_0")], + )? + .sort(vec![col("__agg_0").sort(false, false)])? + .sort(vec![col("__agg_0").sort(true, true)])? + .project(vec![col("__agg_0").alias("max1(j1_id)")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT max(bla.j1_rename) AS "max1(j1_id)" FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla ORDER BY max(bla.j1_rename) ASC NULLS FIRST"#); + Ok(()) +} + +/// Projection → Sort(ASC) → Limit(10) → Sort(DESC) → Aggregate. The +/// inner Sort determines which rows the Limit keeps and the outer Sort +/// re-orders the kept rows — a single SELECT cannot express that, so +/// the outer Sort folds into the outer SELECT (passthrough refs) and +/// the Limit + inner Sort + Aggregate go into a derived subquery. +#[test] +fn test_unparse_aggregate_with_sort_limit_sort_uses_derived_subquery() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![max(col("bla.j1_rename")).alias("__agg_0")], + )? + .sort(vec![col("__agg_0").sort(false, false)])? + .limit(0, Some(10))? + .sort(vec![col("__agg_0").sort(true, true)])? + .project(vec![col("__agg_0").alias("max1(j1_id)")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT __agg_0 AS "max1(j1_id)" FROM (SELECT max(bla.j1_rename) AS __agg_0 FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla ORDER BY max(bla.j1_rename) DESC NULLS LAST LIMIT 10) ORDER BY __agg_0 ASC NULLS FIRST"#); + Ok(()) +} + +/// Projection -> Limit(non-literal fetch) -> Sort { fetch = 5 } -> Aggregate. +/// The outer Limit is non-literal so it can't be combined with the inner +/// Sort's fetch=5. The walk must stop before absorbing the Sort so its +/// fetch survives as `LIMIT 5` in the derived subquery, while the +/// non-literal outer Limit applies on the outer SELECT. +#[test] +fn test_unparse_aggregate_with_non_literal_limit_over_sort_with_fetch() -> Result<()> { + let context = MockContextProvider { + state: MockSessionState::default(), + }; + let j1_schema = context + .get_table_source(TableReference::bare("j1"))? + .schema(); + + let scan = table_scan(Some("j1"), &j1_schema, None)?.build()?; + let plan = LogicalPlanBuilder::from(scan) + .project(vec![col("j1.j1_id").alias("j1_rename")])? + .alias("bla")? + .aggregate( + vec![] as Vec, + vec![max(col("bla.j1_rename")).alias("__agg_0")], + )? + .sort_with_limit(vec![col("__agg_0").sort(true, true)], Some(5))? + .limit_by_expr(None, Some(cast(lit(7_i64), DataType::Int32)))? + .project(vec![col("__agg_0").alias("max1(j1_id)")])? + .build()?; + + let sql = Unparser::default().plan_to_sql(&plan)?.to_string(); + insta::assert_snapshot!(sql, @r#"SELECT __agg_0 AS "max1(j1_id)" FROM (SELECT max(bla.j1_rename) AS __agg_0 FROM (SELECT j1.j1_id AS j1_rename FROM j1) AS bla ORDER BY max(bla.j1_rename) ASC NULLS FIRST LIMIT 5) LIMIT CAST(7 AS INTEGER)"#); + Ok(()) +} + /// Test that unparsing a manually constructed join with a subquery aggregate /// preserves the MAX aggregate function. /// @@ -3909,6 +4351,40 @@ fn snowflake_flatten_cross_join_unnest_table_column() -> Result<(), DataFusionEr Ok(()) } +#[test] +fn roundtrip_approx_percentile_cont_within_group() -> Result<(), DataFusionError> { + roundtrip_statement_with_dialect_helper!( + sql: "SELECT approx_percentile_cont(0.5) WITHIN GROUP (ORDER BY salary) FROM person", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT approx_percentile_cont(0.5) WITHIN GROUP (ORDER BY person.salary ASC NULLS LAST) FROM person", + ); + Ok(()) +} + +#[test] +fn roundtrip_percentile_cont_within_group() -> Result<(), DataFusionError> { + roundtrip_statement_with_dialect_helper!( + sql: "SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY salary) FROM person", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT percentile_cont(0.5) WITHIN GROUP (ORDER BY person.salary ASC NULLS LAST) FROM person", + ); + Ok(()) +} + +#[test] +fn roundtrip_approx_percentile_cont_within_group_with_centroids() +-> Result<(), DataFusionError> { + roundtrip_statement_with_dialect_helper!( + sql: "SELECT approx_percentile_cont(0.9, 200) WITHIN GROUP (ORDER BY salary * 2 DESC) FROM person", + parser_dialect: GenericDialect {}, + unparser_dialect: UnparserDefaultDialect {}, + expected: @"SELECT approx_percentile_cont(0.9, 200) WITHIN GROUP (ORDER BY (person.salary * 2) DESC NULLS FIRST) FROM person", + ); + Ok(()) +} + #[test] fn snowflake_flatten_multiple_unnest_cross_join() -> Result<(), DataFusionError> { // Realistic Snowflake pattern: diff --git a/datafusion/sql/tests/common/mod.rs b/datafusion/sql/tests/common/mod.rs index 71e864d2a733d..e7c819bbf64a6 100644 --- a/datafusion/sql/tests/common/mod.rs +++ b/datafusion/sql/tests/common/mod.rs @@ -56,7 +56,7 @@ impl Display for MockCsvType { #[derive(Default)] pub(crate) struct MockSessionState { scalar_functions: HashMap>, - higher_order_functions: HashMap>, + higher_order_functions: HashMap>, aggregate_functions: HashMap>, expr_planners: Vec>, type_planner: Option>, @@ -101,7 +101,7 @@ impl MockSessionState { pub fn with_higher_order_function( mut self, - higher_order_function: Arc, + higher_order_function: Arc, ) -> Self { self.higher_order_functions.insert( higher_order_function.name().to_string(), @@ -291,7 +291,7 @@ impl ContextProvider for MockContextProvider { self.state.scalar_functions.get(name).cloned() } - fn get_higher_order_meta(&self, name: &str) -> Option> { + fn get_higher_order_meta(&self, name: &str) -> Option> { self.state.higher_order_functions.get(name).cloned() } diff --git a/datafusion/sql/tests/sql_integration.rs b/datafusion/sql/tests/sql_integration.rs index 64763e33d93f7..08a95381b32c8 100644 --- a/datafusion/sql/tests/sql_integration.rs +++ b/datafusion/sql/tests/sql_integration.rs @@ -29,7 +29,7 @@ use common::MockContextProvider; use datafusion_common::{DFSchema, DataFusionError, Result, assert_contains}; use datafusion_expr::{ ColumnarValue, CreateIndex, DdlStatement, Expr, HigherOrderFunctionArgs, - HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDF, + HigherOrderReturnFieldArgs, HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, LambdaParametersProgress, ScalarFunctionArgs, ScalarUDF, ScalarUDFImpl, Signature, ValueOrLambda, Volatility, col, expr::{HigherOrderFunction, LambdaVariable, ScalarFunction}, @@ -100,7 +100,7 @@ fn parse_decimals_3() { assert_snapshot!( plan, @r" - Projection: Decimal128(Some(1),1,1) + Projection: Decimal128(0.1,1,1) EmptyRelation: rows=1 " ); @@ -114,7 +114,7 @@ fn parse_decimals_4() { assert_snapshot!( plan, @r" - Projection: Decimal128(Some(1),2,2) + Projection: Decimal128(0.01,2,2) EmptyRelation: rows=1 " ); @@ -128,7 +128,7 @@ fn parse_decimals_5() { assert_snapshot!( plan, @r" - Projection: Decimal128(Some(10),2,1) + Projection: Decimal128(1.0,2,1) EmptyRelation: rows=1 " ); @@ -142,7 +142,7 @@ fn parse_decimals_6() { assert_snapshot!( plan, @r" - Projection: Decimal128(Some(1001),4,2) + Projection: Decimal128(10.01,4,2) EmptyRelation: rows=1 " ); @@ -156,7 +156,7 @@ fn parse_decimals_7() { assert_snapshot!( plan, @r" - Projection: Decimal128(Some(1000000000000000000000),22,2) + Projection: Decimal128(10000000000000000000.00,22,2) EmptyRelation: rows=1 " ); @@ -184,7 +184,7 @@ fn parse_decimals_9() { assert_snapshot!( plan, @r" - Projection: Decimal128(Some(18446744073709551616),20,0) + Projection: Decimal128(18446744073709551616,20,0) EmptyRelation: rows=1 " ); @@ -725,7 +725,7 @@ fn plan_insert_no_target_columns() { )] #[case::non_existing_column( "INSERT INTO test_decimal (nonexistent, price) VALUES (1, 2), (4, 5)", - "Schema error: No field named nonexistent. \ + "Schema error: No field named nonexistent.\n\ Valid fields are id, price." )] #[case::target_column_count_mismatch( @@ -867,19 +867,27 @@ fn select_filter_cannot_use_alias() { #[test] fn select_neg_filter() { + // NOT requires a boolean expression; applying it to a Utf8 column is an error let sql = "SELECT id, first_name, last_name \ FROM person WHERE NOT state"; - let plan = logical_plan(sql).unwrap(); - assert_snapshot!( - plan, - @r" - Projection: person.id, person.first_name, person.last_name - Filter: NOT person.state - TableScan: person - " + let err = logical_plan(sql).unwrap_err(); + assert!( + err.to_string() + .contains("Unary operator 'NOT' requires a boolean expression"), + "unexpected error: {err}" ); } +#[test] +fn select_not_bool_filter() { + let sql = "SELECT order_id FROM orders WHERE NOT delivered"; + let plan = logical_plan(sql).unwrap(); + let expected = "Projection: orders.order_id\ + \n Filter: NOT orders.delivered\ + \n TableScan: orders"; + assert_eq!(expected, format!("{plan}")); +} + #[test] fn select_compound_filter() { let sql = "SELECT id, first_name, last_name \ @@ -1496,6 +1504,113 @@ fn select_aggregate_with_group_by_with_having_using_count_star_not_in_select() { ); } +/// Asserts that placeholder `id` (e.g. `"$1"`) was inferred as `expected` type +/// somewhere in `plan`. +fn assert_placeholder_type(plan: &LogicalPlan, id: &str, expected: DataType) { + let param_types = plan.get_parameter_types().unwrap(); + assert_eq!(param_types.get(id), Some(&Some(expected))); +} + +/// An expression containing a placeholder, written in both the SELECT list and +/// the GROUP BY, has to be recognised as one expression the way its literal +/// equivalent is. Otherwise the columns inside it read as ungrouped, because the +/// SELECT list has its placeholder types inferred and the grouping key does not. +#[test] +fn select_aggregate_with_group_by_placeholder_expression() { + let sql = "SELECT CASE WHEN age < $1 THEN 'young' ELSE 'old' END, count(*) + FROM person + GROUP BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Projection: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END, count(*) + Aggregate: groupBy=[[CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END]], aggr=[[count(*)]] + TableScan: person + "# + ); + assert_placeholder_type(&plan, "$1", DataType::Int32); +} + +/// The same, for a grouping expression repeated in HAVING. +#[test] +fn select_aggregate_with_having_placeholder_expression() { + let sql = "SELECT CASE WHEN age < $1 THEN 'young' ELSE 'old' END, count(*) + FROM person + GROUP BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END + HAVING CASE WHEN age < $1 THEN 'young' ELSE 'old' END = 'young'"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Projection: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END, count(*) + Filter: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END = Utf8("young") + Aggregate: groupBy=[[CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END]], aggr=[[count(*)]] + TableScan: person + "# + ); + assert_placeholder_type(&plan, "$1", DataType::Int32); +} + +/// The same, for a grouping expression repeated in ORDER BY. +#[test] +fn select_aggregate_with_order_by_placeholder_expression() { + let sql = "SELECT CASE WHEN age < $1 THEN 'young' ELSE 'old' END, count(*) + FROM person + GROUP BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END + ORDER BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Sort: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END ASC NULLS LAST + Projection: CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END, count(*) + Aggregate: groupBy=[[CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END]], aggr=[[count(*)]] + TableScan: person + "# + ); + assert_placeholder_type(&plan, "$1", DataType::Int32); +} + +/// The same, for a window expression repeated in QUALIFY. Here the two spellings +/// of the window expression collide by name instead, since they print alike but +/// do not compare equal. +#[test] +fn select_window_with_qualify_placeholder_expression() { + let sql = "SELECT first_name, + row_number() OVER (PARTITION BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END) + FROM person + QUALIFY row_number() OVER (PARTITION BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END) = 1"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + Projection: person.first_name, row_number() PARTITION BY [CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + Filter: row_number() PARTITION BY [CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING = Int64(1) + WindowAggr: windowExpr=[[row_number() PARTITION BY [CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] + TableScan: person + "# + ); + assert_placeholder_type(&plan, "$1", DataType::Int32); +} + +#[test] +fn select_distinct_on_with_order_by_placeholder_expression() { + let sql = + "SELECT DISTINCT ON (CASE WHEN age < $1 THEN 'young' ELSE 'old' END) first_name + FROM person + ORDER BY CASE WHEN age < $1 THEN 'young' ELSE 'old' END"; + let plan = logical_plan(sql).unwrap(); + assert_snapshot!( + plan, + @r#" + DistinctOn: on_expr=[[CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END]], select_expr=[[person.first_name]], sort_expr=[[CASE WHEN person.age < $1 THEN Utf8("young") ELSE Utf8("old") END ASC NULLS LAST]] + TableScan: person + "# + ); + assert_placeholder_type(&plan, "$1", DataType::Int32); +} + #[test] fn select_binary_expr() { let sql = "SELECT age + salary from person"; @@ -1681,7 +1796,10 @@ fn select_simple_aggregate_with_groupby_and_column_in_group_by_does_not_exist() assert_snapshot!( err.strip_backtrace(), - @r#"Schema error: No field named doesnotexist. Valid fields are "sum(person.age)", person.id, person.first_name, person.last_name, person.age, person.state, person.salary, person.birth_date, person."😀"."# + @r#" +Schema error: No field named doesnotexist. +Valid fields are "sum(person.age)", person.id, person.first_name, person.last_name, person.age, person.state, person.salary, person.birth_date, person."😀". +"# ); } @@ -1760,6 +1878,81 @@ fn select_simple_aggregate_with_groupby_position_out_of_range() { ); } +#[test] +fn select_nested_aggregate() { + // https://github.com/apache/datafusion/issues/23812 + let err = logical_plan("SELECT sum(sum(age)) FROM person") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(person.age)' is nested inside 'sum(sum(person.age))'" + ); + + let err = logical_plan("SELECT state, sum(count(age)) FROM person GROUP BY state") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'count(person.age)' is nested inside 'sum(count(person.age))'" + ); + + let err = + logical_plan("SELECT state FROM person GROUP BY state HAVING sum(sum(age)) > 0") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot be nested: 'sum(person.age)' is nested inside 'sum(sum(person.age))'" + ); +} + +#[test] +fn select_window_function_inside_aggregate() { + // https://github.com/apache/datafusion/issues/23812 + let err = logical_plan("SELECT sum(sum(age) OVER ()) FROM person") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Aggregate function calls cannot contain window function calls: 'sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING)'" + ); +} + +#[test] +fn select_nested_window_function() { + // https://github.com/apache/datafusion/issues/23812 + let err = logical_plan("SELECT sum(sum(age) OVER ()) OVER () FROM person") + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Window function calls cannot be nested: 'sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING' is nested inside 'sum(sum(person.age) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING'" + ); + + let err = logical_plan( + "SELECT rank() OVER (ORDER BY rank() OVER (ORDER BY age)) FROM person", + ) + .expect_err("query should have failed"); + assert_snapshot!( + err.strip_backtrace(), + @"Error during planning: Window function calls cannot be nested: 'rank() ORDER BY [person.age ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW' is nested inside 'rank() ORDER BY [rank() ORDER BY [person.age ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW'" + ); +} + +#[test] +fn select_aggregate_inside_window_function() { + // an aggregate as the argument of a window function is legal: the window + // function is evaluated on top of the aggregate + let plan = + logical_plan("SELECT state, sum(sum(age)) OVER () FROM person GROUP BY state") + .unwrap(); + assert_snapshot!( + plan, + @r" + Projection: person.state, sum(sum(person.age)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING + WindowAggr: windowExpr=[[sum(sum(person.age)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] + Aggregate: groupBy=[[person.state]], aggr=[[sum(person.age)]] + TableScan: person + " + ); +} + #[test] fn select_simple_aggregate_with_groupby_can_use_alias() { let plan = @@ -2267,6 +2460,29 @@ fn create_external_table_csv() { ); } +#[test] +fn create_external_table_multiple_locations() { + let sql = "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION ('foo.csv', 'bar.csv')"; + let plan = logical_plan(sql).unwrap(); + let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = plan else { + panic!("expected a CreateExternalTable plan"); + }; + assert_eq!( + cmd.locations, + vec!["foo.csv".to_string(), "bar.csv".to_string()] + ); +} + +#[test] +fn create_external_table_location_with_literal_comma() { + let sql = "CREATE EXTERNAL TABLE t STORED AS CSV LOCATION 'foo,bar.csv'"; + let plan = logical_plan(sql).unwrap(); + let LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) = plan else { + panic!("expected a CreateExternalTable plan"); + }; + assert_eq!(cmd.locations, vec!["foo,bar.csv".to_string()]); +} + #[test] fn create_external_table_with_pk() { let sql = "CREATE EXTERNAL TABLE t(c1 int, primary key(c1)) STORED AS CSV LOCATION 'foo.csv'"; @@ -3491,6 +3707,104 @@ fn select_groupby_orderby_aggregate_on_non_selected_column_original_issue() { ); } +#[test] +fn plan_merge_into_canonicalizes_qualifiers_and_preserves_quoted_columns() { + let plan = logical_plan( + "MERGE INTO person_quoted_cols AS t USING j2 AS s ON t.id = s.j2_id \ + WHEN MATCHED THEN UPDATE SET \"First Name\" = s.j2_string \ + WHEN NOT MATCHED THEN INSERT (id, \"Age\") VALUES (s.j2_id, 42)", + ) + .unwrap(); + let LogicalPlan::Dml(dml) = &plan else { + panic!("expected Dml, got {plan:?}"); + }; + let datafusion_expr::WriteOp::MergeInto(merge_op) = &dml.op else { + panic!("expected MergeInto, got {:?}", dml.op); + }; + + assert_eq!(merge_op.on.to_string(), "person_quoted_cols.id = s.j2_id"); + + let datafusion_expr::dml::MergeIntoAction::Update(assignments) = + &merge_op.clauses[0].action + else { + panic!("expected UPDATE"); + }; + assert_eq!(assignments[0].0, "First Name"); + assert_eq!(assignments[0].1.to_string(), "s.j2_string"); + + let datafusion_expr::dml::MergeIntoAction::Insert { columns, values } = + &merge_op.clauses[1].action + else { + panic!("expected INSERT"); + }; + assert_eq!(columns, &["id".to_string(), "Age".to_string()]); + assert_eq!(values[0].to_string(), "s.j2_id"); +} + +#[rstest] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = j2.j2_string WHERE false", + "MERGE UPDATE WHERE predicates are not supported" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = j2.j2_string DELETE WHERE false", + "MERGE UPDATE DELETE WHERE predicates are not supported" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN NOT MATCHED THEN INSERT (j1_id, j1_string) \ + VALUES (j2.j2_id, j2.j2_string) WHERE false", + "MERGE INSERT WHERE predicates are not supported" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN MATCHED THEN UPDATE SET j1_string = 'a', j1_string = 'b'", + "Duplicate column 'j1_string' in MERGE UPDATE" +)] +#[case( + "MERGE INTO j1 AS t USING j2 AS s ON t.j1_id = s.j2_id \ + WHEN MATCHED THEN UPDATE SET s.j1_string = s.j2_string", + "MERGE assignment target 's.j1_string' must reference target table 't'" +)] +#[case( + "MERGE INTO j1 AS t USING j2 AS s ON t.j1_id = s.j2_id \ + WHEN NOT MATCHED THEN INSERT (s.j1_id) VALUES (s.j2_id)", + "MERGE assignment target 's.j1_id' must reference target table 't'" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id", + "MERGE INTO requires at least one WHEN clause" +)] +#[case( + "MERGE INTO j1 USING j2 ON j1.j1_id = j2.j2_id \ + WHEN NOT MATCHED THEN INSERT (j1_id, J1_ID) VALUES (1, 2)", + "Duplicate column 'j1_id' in MERGE INSERT" +)] +#[case( + "MERGE INTO j1() USING j2 ON true WHEN MATCHED THEN DELETE", + "MERGE target table modifiers are not supported" +)] +#[case( + "MERGE INTO j1 PARTITION (p0) USING j2 ON true WHEN MATCHED THEN DELETE", + "MERGE target table modifiers are not supported" +)] +#[case( + "MERGE INTO j1 AS t(a) USING j2 ON true WHEN MATCHED THEN DELETE", + "MERGE target alias column lists are not supported" +)] +fn plan_merge_into_rejects_invalid_actions_and_structure( + #[case] sql: &str, + #[case] expected: &str, +) { + let err = logical_plan(sql).unwrap_err(); + assert!( + err.strip_backtrace().contains(expected), + "unexpected error: {err}" + ); +} + fn logical_plan(sql: &str) -> Result { logical_plan_with_options(sql, ParserOptions::default()) } @@ -3510,7 +3824,9 @@ fn logical_plan_with_options(sql: &str, options: ParserOptions) -> Result Result { let state = MockSessionState::default() .with_aggregate_function(sum_udaf()) - .with_higher_order_function(Arc::new(MockArrayReduce::new())) + .with_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( + MockArrayReduce::new(), + ))) .with_scalar_function(make_array_udf()) .with_expr_planner(Arc::new(CustomExprPlanner {})); // plan array literal let context = MockContextProvider { state }; @@ -5358,7 +5674,7 @@ fn test_progressive_lambda_parameters() { assert_eq!( expr, Expr::HigherOrderFunction(HigherOrderFunction::new( - Arc::new(MockArrayReduce::new()), + Arc::new(HigherOrderUDF::new_from_impl(MockArrayReduce::new())), vec![ Expr::ScalarFunction(ScalarFunction::new_udf( make_array_udf(), @@ -5402,7 +5718,7 @@ impl MockArrayReduce { } } -impl HigherOrderUDF for MockArrayReduce { +impl HigherOrderUDFImpl for MockArrayReduce { fn name(&self) -> &str { "array_reduce" } diff --git a/datafusion/sqllogictest/Cargo.toml b/datafusion/sqllogictest/Cargo.toml index e2ffe1415a1fb..13493d16c05e5 100644 --- a/datafusion/sqllogictest/Cargo.toml +++ b/datafusion/sqllogictest/Cargo.toml @@ -64,7 +64,7 @@ tempfile = { workspace = true } testcontainers-modules = { workspace = true, features = ["postgres"], optional = true } thiserror = "2.0.18" tokio = { workspace = true } -tokio-postgres = { version = "0.7.17", optional = true } +tokio-postgres = { version = "0.7.18", optional = true } [features] avro = ["datafusion/avro"] diff --git a/datafusion/sqllogictest/bin/sqllogictests.rs b/datafusion/sqllogictest/bin/sqllogictests.rs index 69ae3a2fa7dd3..da0beb0c29a28 100644 --- a/datafusion/sqllogictest/bin/sqllogictests.rs +++ b/datafusion/sqllogictest/bin/sqllogictests.rs @@ -61,6 +61,7 @@ const DATAFUSION_TESTING_TEST_DIRECTORY: &str = "../../datafusion-testing/data/" const PG_COMPAT_FILE_PREFIX: &str = "pg_compat_"; const TPCH_PREFIX: &str = "tpch"; const SQLITE_PREFIX: &str = "sqlite"; +const ENCRYPTED_PARQUET_FILE: &str = "encrypted_parquet.slt"; const ERRS_PER_FILE_LIMIT: usize = 10; const TIMING_DEBUG_SLOW_FILES_ENV: &str = "SLT_TIMING_DEBUG_SLOW_FILES"; @@ -452,7 +453,7 @@ async fn run_test_file_substrait_round_trip( let pb = mp.add(ProgressBar::new(count)); pb.set_style(mp_style); - pb.set_message(format!("{:?}", &relative_path)); + pb.set_message(format!("{relative_path:?}")); let mut runner = sqllogictest::Runner::new(|| async { Ok(DataFusionSubstraitRoundTrip::new( @@ -472,6 +473,10 @@ async fn run_test_file_substrait_round_trip( } #[cfg(not(feature = "substrait"))] +#[expect( + clippy::unused_async, + reason = "matches the substrait-enabled implementation" +)] async fn run_test_file_substrait_round_trip( _test_file: TestFile, _validator: Validator, @@ -507,7 +512,7 @@ async fn run_test_file( let pb = mp.add(ProgressBar::new(count)); pb.set_style(mp_style); - pb.set_message(format!("{:?}", &relative_path)); + pb.set_message(format!("{relative_path:?}")); // If DataFusion configuration has changed during test file runs, errors will be // pushed to this vec. @@ -626,7 +631,7 @@ async fn run_test_file_with_postgres( let pb = mp.add(ProgressBar::new(count)); pb.set_style(mp_style); - pb.set_message(format!("{:?}", &relative_path)); + pb.set_message(format!("{relative_path:?}")); let mut runner = sqllogictest::Runner::new(|| { Postgres::connect_with_tracked_sql( @@ -645,6 +650,10 @@ async fn run_test_file_with_postgres( } #[cfg(not(feature = "postgres"))] +#[expect( + clippy::unused_async, + reason = "matches the postgres-enabled implementation" +)] async fn run_test_file_with_postgres( _test_file: TestFile, _validator: Validator, @@ -681,7 +690,7 @@ async fn run_complete_file( let pb = mp.add(ProgressBar::new(count)); pb.set_style(mp_style); - pb.set_message(format!("{:?}", &relative_path)); + pb.set_message(format!("{relative_path:?}")); let config_change_errors = Arc::new(Mutex::new(Vec::new())); let mut runner = sqllogictest::Runner::new(|| async { @@ -737,7 +746,7 @@ async fn run_complete_file_with_postgres( let pb = mp.add(ProgressBar::new(count)); pb.set_style(mp_style); - pb.set_message(format!("{:?}", &relative_path)); + pb.set_message(format!("{relative_path:?}")); let mut runner = sqllogictest::Runner::new(|| { Postgres::connect_with_tracked_sql( @@ -770,6 +779,10 @@ async fn run_complete_file_with_postgres( } #[cfg(not(feature = "postgres"))] +#[expect( + clippy::unused_async, + reason = "matches the postgres-enabled implementation" +)] async fn run_complete_file_with_postgres( _test_file: TestFile, _validator: Validator, @@ -805,6 +818,10 @@ fn read_test_files(options: &Options) -> Result> { .filter(|f| f.is_slt_file()) .filter(|f| !f.relative_path_starts_with(TPCH_PREFIX) || options.include_tpch) .filter(|f| !f.relative_path_starts_with(SQLITE_PREFIX) || options.include_sqlite) + .filter(|f| { + !f.relative_path_starts_with(ENCRYPTED_PARQUET_FILE) + || cfg!(feature = "parquet_encryption") + }) .filter(|f| options.check_pg_compat_file(f.path.as_path())) .collect::>(); diff --git a/datafusion/sqllogictest/src/engines/conversion.rs b/datafusion/sqllogictest/src/engines/conversion.rs index 3e519042f4ee0..4447ea9f03bea 100644 --- a/datafusion/sqllogictest/src/engines/conversion.rs +++ b/datafusion/sqllogictest/src/engines/conversion.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use arrow::datatypes::{Decimal128Type, Decimal256Type, DecimalType, i256}; +use arrow::datatypes::DecimalType; use bigdecimal::BigDecimal; use half::f16; use std::str::FromStr; @@ -96,21 +96,15 @@ pub(crate) fn spark_f64_to_str(value: f64) -> String { } } -pub(crate) fn decimal_128_to_str(value: i128, scale: i8) -> String { +pub(crate) fn arrow_decimal_to_str( + value: T::Native, + scale: i8, +) -> String { let precision = u8::MAX; // does not matter + let value = T::format_decimal(value, precision, scale); big_decimal_to_str( - BigDecimal::from_str(&Decimal128Type::format_decimal(value, precision, scale)) - .unwrap(), - None, - ) -} - -pub(crate) fn decimal_256_to_str(value: i256, scale: i8) -> String { - let precision = u8::MAX; // does not matter - big_decimal_to_str( - BigDecimal::from_str(&Decimal256Type::format_decimal(value, precision, scale)) - .unwrap(), - None, + BigDecimal::from_str(&value).unwrap(), + Some(i64::from(scale)), ) } @@ -132,7 +126,10 @@ pub(crate) fn big_decimal_to_str(value: BigDecimal, round_digits: Option) - #[cfg(test)] mod tests { - use super::big_decimal_to_str; + use super::{arrow_decimal_to_str, big_decimal_to_str}; + use arrow::datatypes::{ + Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, i256, + }; use bigdecimal::{BigDecimal, num_bigint::BigInt}; macro_rules! assert_decimal_str_eq { @@ -196,4 +193,15 @@ mod tests { assert_decimal_str_eq!(10_i128.pow(13) + 11, 13, Some(13), "1.0000000000011"); } + + #[test] + fn test_arrow_decimal_to_str() { + assert_eq!(arrow_decimal_to_str::(12345, 2), "123.45"); + assert_eq!(arrow_decimal_to_str::(12345, 2), "123.45"); + assert_eq!(arrow_decimal_to_str::(12345, 2), "123.45"); + assert_eq!( + arrow_decimal_to_str::(i256::from(12345), 2), + "123.45" + ); + } } diff --git a/datafusion/sqllogictest/src/engines/datafusion_engine/normalize.rs b/datafusion/sqllogictest/src/engines/datafusion_engine/normalize.rs index 2c549422d6547..fa566a6a3d251 100644 --- a/datafusion/sqllogictest/src/engines/datafusion_engine/normalize.rs +++ b/datafusion/sqllogictest/src/engines/datafusion_engine/normalize.rs @@ -19,7 +19,9 @@ use super::super::conversion::*; use super::error::{DFSqlLogicTestError, Result}; use crate::engines::output::DFColumnType; use arrow::array::{Array, AsArray}; -use arrow::datatypes::{Fields, Schema}; +use arrow::datatypes::{ + Decimal32Type, Decimal64Type, Decimal128Type, Decimal256Type, Fields, Schema, +}; use arrow::util::display::ArrayFormatter; use arrow::{array, array::ArrayRef, datatypes::DataType, record_batch::RecordBatch}; use datafusion::common::internal_datafusion_err; @@ -209,13 +211,21 @@ pub fn cell_to_string(col: &ArrayRef, row: usize, is_spark_path: bool) -> Result Ok(f64_to_str(result)) } } + DataType::Decimal32(_, scale) => { + let value = get_row_value!(array::Decimal32Array, col, row); + Ok(arrow_decimal_to_str::(value, *scale)) + } + DataType::Decimal64(_, scale) => { + let value = get_row_value!(array::Decimal64Array, col, row); + Ok(arrow_decimal_to_str::(value, *scale)) + } DataType::Decimal128(_, scale) => { let value = get_row_value!(array::Decimal128Array, col, row); - Ok(decimal_128_to_str(value, *scale)) + Ok(arrow_decimal_to_str::(value, *scale)) } DataType::Decimal256(_, scale) => { let value = get_row_value!(array::Decimal256Array, col, row); - Ok(decimal_256_to_str(value, *scale)) + Ok(arrow_decimal_to_str::(value, *scale)) } DataType::LargeUtf8 => Ok(varchar_to_str(get_row_value!( array::LargeStringArray, @@ -268,9 +278,11 @@ pub fn convert_schema_to_types(columns: &Fields) -> Vec { | DataType::UInt16 | DataType::UInt32 | DataType::UInt64 => DFColumnType::Integer, - DataType::Float16 - | DataType::Float32 - | DataType::Float64 + DataType::Float16 | DataType::Float32 | DataType::Float64 => { + DFColumnType::Float + } + DataType::Decimal32(_, _) + | DataType::Decimal64(_, _) | DataType::Decimal128(_, _) | DataType::Decimal256(_, _) => DFColumnType::Float, DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View => { diff --git a/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs b/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs index c3f266dcd1b62..f085fb5708875 100644 --- a/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs +++ b/datafusion/sqllogictest/src/engines/postgres_engine/mod.rs @@ -75,8 +75,7 @@ impl Postgres { /// /// See https://docs.rs/tokio-postgres/latest/tokio_postgres/config/struct.Config.html#url for format pub async fn connect(relative_path: PathBuf, pb: ProgressBar) -> Result { - let uri = std::env::var("PG_URI") - .map_or_else(|_| PG_URI.to_string(), std::convert::identity); + let uri = std::env::var("PG_URI").unwrap_or_else(|_| PG_URI.to_string()); info!("Using postgres connection string: {uri}"); diff --git a/datafusion/sqllogictest/src/test_context.rs b/datafusion/sqllogictest/src/test_context.rs index 0edde71b939f4..d85ca2db76268 100644 --- a/datafusion/sqllogictest/src/test_context.rs +++ b/datafusion/sqllogictest/src/test_context.rs @@ -36,6 +36,7 @@ use arrow::record_batch::RecordBatch; use datafusion::catalog::{ CatalogProvider, MemoryCatalogProvider, MemorySchemaProvider, SchemaProvider, Session, }; +use datafusion::common::config::Dialect; use datafusion::common::{DataFusionError, Result, not_impl_err}; use datafusion::functions::math::abs; use datafusion::logical_expr::async_udf::{AsyncScalarUDF, AsyncScalarUDFImpl}; @@ -53,6 +54,8 @@ use datafusion::{ use datafusion_spark::SessionStateBuilderSpark; use crate::is_spark_path; +use range_partitioning::register_range_partitioned_table; + use async_trait::async_trait; use datafusion::common::cast::as_float64_array; use datafusion::execution::SessionStateBuilder; @@ -61,6 +64,8 @@ use log::info; use sqlparser::ast; use tempfile::TempDir; +mod range_partitioning; + /// Context for running tests pub struct TestContext { /// Context for running queries @@ -112,6 +117,9 @@ impl TestContext { if is_spark_path(relative_path) { state_builder = state_builder.with_spark_features(); + if let Some(config) = state_builder.config() { + config.options_mut().sql_parser.dialect = Dialect::Spark; + } } if matches!( @@ -134,15 +142,15 @@ impl TestContext { } "information_schema_table_types.slt" => { info!("Registering local temporary table"); - register_temp_table(test_ctx.session_ctx()).await; + register_temp_table(test_ctx.session_ctx()); } "information_schema_columns.slt" => { info!("Registering table with many types"); - register_table_with_many_types(test_ctx.session_ctx()).await; + register_table_with_many_types(test_ctx.session_ctx()); } "map.slt" => { info!("Registering table with map"); - register_table_with_map(test_ctx.session_ctx()).await; + register_table_with_map(test_ctx.session_ctx()); } "avro.slt" => { #[cfg(feature = "avro")] @@ -165,16 +173,25 @@ impl TestContext { test_ctx.ctx.register_udf(example_udf); register_partition_table(&mut test_ctx).await; info!("Registering table with many types"); - register_table_with_many_types(test_ctx.session_ctx()).await; + register_table_with_many_types(test_ctx.session_ctx()); + } + "range_partitioning.slt" => { + info!("Registering range partitioned table"); + register_range_partitioned_table(test_ctx.session_ctx()); } "metadata.slt" | "arrow_field.slt" => { info!("Registering metadata table tables"); - register_metadata_tables(test_ctx.session_ctx()).await; + register_metadata_tables(test_ctx.session_ctx()); + register_conflicting_metadata_tables(test_ctx.session_ctx()) } "union_function.slt" => { info!("Registering table with union column"); register_union_table(test_ctx.session_ctx()) } + "aggregate.slt" => { + info!("Registering table with union column for approx_distinct"); + register_approx_distinct_union_table(test_ctx.session_ctx()) + } "dictionary_struct.slt" => { info!("Registering table with dictionary-encoded struct column"); register_dictionary_struct_table(test_ctx.session_ctx()); @@ -354,7 +371,7 @@ pub async fn register_partition_table(test_ctx: &mut TestContext) { } // registers a LOCAL TEMPORARY table. -pub async fn register_temp_table(ctx: &SessionContext) { +pub fn register_temp_table(ctx: &SessionContext) { #[derive(Debug)] struct TestTable(TableType); @@ -386,7 +403,7 @@ pub async fn register_temp_table(ctx: &SessionContext) { .unwrap(); } -pub async fn register_table_with_many_types(ctx: &SessionContext) { +pub fn register_table_with_many_types(ctx: &SessionContext) { let catalog = MemoryCatalogProvider::new(); let schema = MemorySchemaProvider::new(); @@ -402,7 +419,7 @@ pub async fn register_table_with_many_types(ctx: &SessionContext) { .unwrap(); } -pub async fn register_table_with_map(ctx: &SessionContext) { +pub fn register_table_with_map(ctx: &SessionContext) { let key = Field::new("key", DataType::Int64, false); let value = Field::new("value", DataType::Int64, true); let map_field = @@ -452,7 +469,7 @@ fn table_with_many_types() -> Arc { } /// Registers a table_with_metadata that contains both field level and Table level metadata -pub async fn register_metadata_tables(ctx: &SessionContext) { +pub fn register_metadata_tables(ctx: &SessionContext) { let id = Field::new("id", DataType::Int32, true).with_metadata(HashMap::from([( String::from("metadata_key"), String::from("the id field"), @@ -556,14 +573,17 @@ fn register_union_table(ctx: &SessionContext) { ], ) .unwrap(), - ScalarBuffer::from(vec![3, 1, 3]), + ScalarBuffer::from(vec![3, 1, 3, 3, 1, 3]), None, vec![ - Arc::new(Int32Array::from(vec![1, 2, 3])), + Arc::new(Int32Array::from(vec![1, 2, 3, 1, 5, 3])), Arc::new(StringArray::from(vec![ Some("foo"), Some("bar"), Some("baz"), + Some("qux"), + Some("bar"), + Some("quux"), ])), ], ) @@ -581,6 +601,43 @@ fn register_union_table(ctx: &SessionContext) { ctx.register_batch("union_table", batch).unwrap(); } +fn register_approx_distinct_union_table(ctx: &SessionContext) { + let union = UnionArray::try_new( + UnionFields::try_new( + vec![0, 1], + vec![ + Field::new("i", DataType::Int32, true), + Field::new("s", DataType::Utf8, true), + ], + ) + .unwrap(), + ScalarBuffer::from(vec![0_i8, 0, 1, 1, 0, 0, 1, 0]), + Some(ScalarBuffer::from(vec![0, 1, 0, 1, 2, 3, 2, 4])), + vec![ + Arc::new(Int32Array::from(vec![ + Some(1), + Some(1), + None, + None, + Some(5), + ])), + Arc::new(StringArray::from(vec![Some("x"), Some("y"), None])), + ], + ) + .unwrap(); + + let schema = Schema::new(vec![ + Field::new("g", DataType::Int32, false), + Field::new("u", union.data_type().clone(), false), + ]); + + let g = Arc::new(Int32Array::from(vec![1, 1, 1, 2, 2, 3, 3, 4])); + let batch = RecordBatch::try_new(Arc::new(schema), vec![g, Arc::new(union)]).unwrap(); + + ctx.register_batch("approx_distinct_union_test", batch) + .unwrap(); +} + fn register_dictionary_struct_table(ctx: &SessionContext) { // Build deduplicated struct values: 3 unique structs let names = Arc::new(StringArray::from(vec!["Alice", "Bob", "Carol"])) as ArrayRef; @@ -712,3 +769,26 @@ fn register_async_abs_udf(ctx: &SessionContext) { let udf = AsyncScalarUDF::new(Arc::new(async_abs)); ctx.register_udf(udf.into_scalar_udf()); } + +fn register_conflicting_metadata_tables(ctx: &SessionContext) { + let schema_left = + Schema::new(vec![Field::new("a", DataType::Int32, false)]).with_metadata( + HashMap::from([(String::from("metadata_key"), String::from("left"))]), + ); + let data_left = + Arc::new(Int32Array::from(vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10])) as ArrayRef; + + let batch_left = + RecordBatch::try_new(Arc::new(schema_left), vec![Arc::new(data_left)]).unwrap(); + ctx.register_batch("larger_table", batch_left).unwrap(); + + let schema_right = + Schema::new(vec![Field::new("b", DataType::Int32, false)]).with_metadata( + HashMap::from([(String::from("metadata_key"), String::from("right"))]), + ); + let data_right = Arc::new(Int32Array::from(vec![1])) as ArrayRef; + + let batch_right = + RecordBatch::try_new(Arc::new(schema_right), vec![Arc::new(data_right)]).unwrap(); + ctx.register_batch("smaller_table", batch_right).unwrap(); +} diff --git a/datafusion/sqllogictest/src/test_context/range_partitioning.rs b/datafusion/sqllogictest/src/test_context/range_partitioning.rs new file mode 100644 index 0000000000000..3cde3939f0b7c --- /dev/null +++ b/datafusion/sqllogictest/src/test_context/range_partitioning.rs @@ -0,0 +1,294 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fs::{File, create_dir_all, remove_dir_all}; +use std::path::Path; +use std::sync::Arc; + +use arrow::array::{ArrayRef, Int32Array}; +use arrow::compute::SortOptions; +use arrow::datatypes::{DataType, Field, Schema, SchemaRef}; +use arrow::record_batch::RecordBatch; +use datafusion::catalog::streaming::StreamingTable; +use datafusion::common::{ScalarValue, SplitPoint}; +use datafusion::datasource::file_format::parquet::ParquetFormat; +use datafusion::datasource::listing::{ + ListingOptions, ListingTable, ListingTableConfig, ListingTableUrl, +}; +use datafusion::logical_expr::{Partitioning, RangePartitioning, col}; +use datafusion::parquet::arrow::ArrowWriter; +use datafusion::physical_expr::{ + Partitioning as PhysicalPartitioning, PhysicalSortExpr, + RangePartitioning as PhysicalRangePartitioning, expressions::col as physical_col, +}; +use datafusion::physical_plan::streaming::PartitionStream; +use datafusion::physical_plan::test::TestPartitionStream; +use datafusion::prelude::SessionContext; + +// ============================================================================== +// Range Partitioned Table (sqllogictest-only) +// ============================================================================== + +/// Registers a simple range-partitioned listing table for testing before +/// declaring such tables is supported via SQL. +pub(super) fn register_range_partitioned_table(ctx: &SessionContext) { + const RANGE_PARTITIONS: [&[(i32, i32, i32)]; 4] = [ + &[(1, 1, 10), (5, 2, 50)], + &[(10, 1, 100), (15, 2, 150)], + &[(20, 1, 200), (25, 2, 250)], + &[(30, 1, 300), (35, 2, 350)], + ]; + const SHIFTED_RANGE_PARTITIONS: [&[(i32, i32, i32)]; 4] = [ + &[(1, 1, 10), (5, 2, 50), (10, 1, 100)], + &[(15, 2, 150)], + &[(20, 1, 200), (25, 2, 250)], + &[(30, 1, 300), (35, 2, 350)], + ]; + const NARROW_RANGE_PARTITIONS: [&[(i32, i32, i32)]; 3] = [ + &[(1, 1, 10), (5, 2, 50)], + &[(10, 1, 100), (15, 2, 150)], + &[(20, 1, 200), (25, 2, 250), (30, 1, 300), (35, 2, 350)], + ]; + const SPARSE_RANGE_PARTITIONS: [&[(i32, i32, i32)]; 4] = [ + &[(5, 2, 50), (8, 3, 80)], + &[(10, 1, 100)], + &[(20, 1, 200)], + &[(30, 1, 300), (40, 4, 400)], + ]; + + let schema = Arc::new(Schema::new(vec![ + Field::new("range_key", DataType::Int32, false), + Field::new("non_range_key", DataType::Int32, false), + Field::new("value", DataType::Int32, false), + ])); + let output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + let range_table_dir = Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned"); + + register_parquet_listing_table( + ctx, + "range_partitioned", + &range_table_dir, + Arc::clone(&schema), + RANGE_PARTITIONS, + output_partitioning, + ); + + register_unbounded_range_stream_table( + ctx, + "unbounded_range_like", + Arc::clone(&schema), + [10, 20, 30], + RANGE_PARTITIONS.map(|rows| rows.to_vec()), + ); + register_unbounded_range_stream_table( + ctx, + "unbounded_range_like_shifted", + Arc::clone(&schema), + [15, 20, 30], + SHIFTED_RANGE_PARTITIONS.map(|rows| rows.to_vec()), + ); + + let shifted_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(15))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_parquet_listing_table( + ctx, + "range_partitioned_shifted", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_shifted"), + Arc::clone(&schema), + SHIFTED_RANGE_PARTITIONS, + shifted_output_partitioning, + ); + + // Same rows as `range_partitioned` but split into only three range + // partitions on `range_key`. Used to exercise the co-partition check when + // two Range inputs disagree on partition count. + let narrow_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_parquet_listing_table( + ctx, + "range_partitioned_narrow", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_narrow"), + Arc::clone(&schema), + NARROW_RANGE_PARTITIONS, + narrow_output_partitioning, + ); + + let sparse_output_partitioning = Partitioning::Range( + RangePartitioning::try_new( + vec![col("range_key").sort(true, true)], + vec![ + SplitPoint::new(vec![ScalarValue::Int32(Some(10))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(20))]), + SplitPoint::new(vec![ScalarValue::Int32(Some(30))]), + ], + ) + .expect("range partitioning should be valid"), + ); + + register_parquet_listing_table( + ctx, + "range_partitioned_sparse", + Path::new(env!("CARGO_MANIFEST_DIR")) + .join("test_files/scratch_range_partitioning/range_partitioned_sparse"), + Arc::clone(&schema), + SPARSE_RANGE_PARTITIONS, + sparse_output_partitioning, + ); +} + +fn register_parquet_listing_table( + ctx: &SessionContext, + name: &str, + table_dir: impl AsRef, + schema: SchemaRef, + partitions: impl IntoIterator, + output_partitioning: Partitioning, +) { + let table_dir = table_dir.as_ref(); + if table_dir.exists() { + remove_dir_all(table_dir).expect("test table dir should be removable"); + } + create_dir_all(table_dir).expect("test table dir should be created"); + for (idx, rows) in partitions.into_iter().enumerate() { + let batch = range_batch(Arc::clone(&schema), rows); + let file = File::create(table_dir.join(format!("part-{idx}.parquet"))) + .expect("test table parquet partition should be created"); + let mut writer = ArrowWriter::try_new(file, Arc::clone(&schema), None) + .expect("test table parquet writer should be created"); + writer + .write(&batch) + .expect("test table parquet partition should be written"); + writer + .close() + .expect("test table parquet writer should close"); + } + + let table_path = format!( + "{}/", + table_dir + .to_str() + .expect("test table path should be valid utf8") + ); + let table_url = + ListingTableUrl::parse(&table_path).expect("test table url should parse"); + let options = ListingOptions::new(Arc::new(ParquetFormat::default())) + .with_output_partitioning(Some(output_partitioning)); + let config = ListingTableConfig::new(table_url) + .with_listing_options(options) + .with_schema(schema); + let table = + ListingTable::try_new(config).expect("test listing table should be valid"); + + ctx.register_table(name, Arc::new(table)) + .expect("test listing table registration should succeed"); +} + +fn register_unbounded_range_stream_table( + ctx: &SessionContext, + name: &str, + schema: Arc, + split_points: [i32; 3], + partition_rows: [Vec<(i32, i32, i32)>; 4], +) { + let output_partitioning = PhysicalPartitioning::Range( + PhysicalRangePartitioning::try_new( + [PhysicalSortExpr { + expr: physical_col("range_key", &schema) + .expect("range key should exist in stream schema"), + options: SortOptions::default(), + }] + .into(), + split_points + .into_iter() + .map(|value| SplitPoint::new(vec![ScalarValue::Int32(Some(value))])) + .collect(), + ) + .expect("range partitioning should be valid"), + ); + let partitions = partition_rows + .into_iter() + .map(|rows| range_stream_partition(Arc::clone(&schema), &rows)) + .collect(); + + ctx.register_table( + name, + Arc::new( + StreamingTable::try_new(schema, partitions) + .expect("range stream table should be valid") + .with_infinite_table(true) + .with_output_partitioning(output_partitioning), + ), + ) + .expect("test stream table registration should succeed"); +} + +fn range_stream_partition( + schema: SchemaRef, + rows: &[(i32, i32, i32)], +) -> Arc { + Arc::new(TestPartitionStream::new_with_batches(vec![range_batch( + schema, rows, + )])) +} + +fn range_batch(schema: SchemaRef, rows: &[(i32, i32, i32)]) -> RecordBatch { + RecordBatch::try_new( + schema, + vec![ + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.0))) + as ArrayRef, + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.1))) + as ArrayRef, + Arc::new(Int32Array::from_iter_values(rows.iter().map(|row| row.2))) + as ArrayRef, + ], + ) + .expect("range batch should be valid") +} diff --git a/datafusion/sqllogictest/src/test_file.rs b/datafusion/sqllogictest/src/test_file.rs index 71dbfa6edc944..a7609bb8018ab 100644 --- a/datafusion/sqllogictest/src/test_file.rs +++ b/datafusion/sqllogictest/src/test_file.rs @@ -107,26 +107,38 @@ impl Ord for TestFile { /// $ cargo test --profile=ci --test sqllogictests -- --timing-summary top /// ... /// Per-file elapsed summary (deterministic): -/// 1. 3.568s aggregate.slt -/// 2. 3.464s joins.slt -/// 3. 3.336s imdb.slt -/// 4. 3.085s push_down_filter_regression.slt -/// 5. 2.926s aggregate_skip_partial.slt -/// 6. 2.399s window.slt -/// 7. 2.198s group_by.slt -/// 8. 1.281s clickbench.slt -/// 9. 1.058s datetime/timestamps.slt +/// 1. 5.437s nested_loop_join_spill.slt +/// 2. 3.471s push_down_filter_regression.slt +/// 3. 3.458s aggregate.slt +/// 4. 3.065s joins.slt +/// 5. 2.852s aggregate_skip_partial.slt +/// 6. 2.832s imdb.slt +/// 7. 2.453s window.slt +/// 8. 1.831s group_by.slt +/// 9. 1.282s clickbench.slt +/// 10. 1.055s datetime/timestamps.slt +/// 11. 0.994s array/array_has.slt +/// 12. 0.840s cte.slt +/// 13. 0.748s sort_pushdown.slt +/// 14. 0.714s push_down_filter_parquet.slt +/// 15. 0.668s projection_pushdown.slt /// ``` const TEST_PRIORITY_ENTRIES: &[&str] = &[ - "aggregate.slt", // longest-running files go first - "joins.slt", - "imdb.slt", + "nested_loop_join_spill.slt", // longest-running files go first "push_down_filter_regression.slt", + "aggregate.slt", + "joins.slt", "aggregate_skip_partial.slt", + "imdb.slt", "window.slt", "group_by.slt", "clickbench.slt", "datetime/timestamps.slt", + "array/array_has.slt", + "cte.slt", + "sort_pushdown.slt", + "push_down_filter_parquet.slt", + "projection_pushdown.slt", ]; /// Default priority for tests not in the priority map. Tests with lower diff --git a/datafusion/sqllogictest/test_files/aggregate.slt b/datafusion/sqllogictest/test_files/aggregate.slt index 25b69d16dd035..460d4cd2ffda3 100644 --- a/datafusion/sqllogictest/test_files/aggregate.slt +++ b/datafusion/sqllogictest/test_files/aggregate.slt @@ -212,203 +212,6 @@ WITHIN GROUP (ORDER BY c3) OVER (ROWS BETWEEN 4 PRECEDING AND CURRENT ROW) FROM aggregate_test_100 -# array agg can use order by -query ? -SELECT array_agg(c13 ORDER BY c13) -FROM - (SELECT * - FROM aggregate_test_100 - ORDER BY c13 - LIMIT 5) as t1 ----- -[0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8] - -# array agg can use order by with distinct -query ? -SELECT array_agg(DISTINCT c13 ORDER BY c13) -FROM - (SELECT * - FROM aggregate_test_100 - ORDER BY c13 - LIMIT 5) as t1 ----- -[0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8] - -query error Execution error: In an aggregate with DISTINCT, ORDER BY expressions must appear in argument list -SELECT array_agg(DISTINCT c13 ORDER BY c12) -FROM aggregate_test_100 - -query error Execution error: In an aggregate with DISTINCT, ORDER BY expressions must appear in argument list -SELECT array_agg(DISTINCT c13 ORDER BY c13, c12) -FROM aggregate_test_100 - -query ?? rowsort -with tbl as (SELECT * FROM (VALUES ('xxx', 'yyy'), ('xxx', 'yyy'), ('xxx2', 'yyy2')) AS t(x, y)) -select - array_agg(x order by x) as x_agg, - array_agg(y order by y) as y_agg -from tbl -group by all ----- -[xxx, xxx, xxx2] [yyy, yyy, yyy2] - -query ?? -SELECT - (SELECT array_agg(c12 ORDER BY c12) FROM aggregate_test_100), - (SELECT array_agg(c13 ORDER BY c13) FROM aggregate_test_100) ----- -[0.01479305307777301, 0.02182578039211991, 0.03968347085780355, 0.04429073092078406, 0.047343434291126085, 0.04893135681998029, 0.0494924465469434, 0.05573662213439634, 0.05636955101974106, 0.061029375346466685, 0.07260475960924484, 0.09465635123783445, 0.12357539988406441, 0.152498292971736, 0.16301110515739792, 0.1640882545084913, 0.1754261586710173, 0.17592486905979987, 0.17909035118828576, 0.18628859265874176, 0.19113293583306745, 0.2145232647388039, 0.21535402343780985, 0.24899794314659673, 0.2537253407987472, 0.2667177795079635, 0.27159190516490006, 0.2739938529235548, 0.28534428578703896, 0.2944158618048994, 0.296036538664718, 0.3051364088814128, 0.30585375151301186, 0.3114712539863804, 0.3231750610081745, 0.32869374687050157, 0.33639590659276175, 0.3600766362333053, 0.36936304600612724, 0.38870280983958583, 0.39144436569161134, 0.40342283197779727, 0.4094218353587008, 0.40975383525297016, 0.42073125331890115, 0.4273123318932347, 0.42950521730777025, 0.4830878559436823, 0.5081765563442366, 0.5437595540422571, 0.5590205548347534, 0.5593249815276734, 0.5603062368164834, 0.560333188635217, 0.5614503754617461, 0.565352842229935, 0.574210838214554, 0.5759450483859969, 0.5773498217058918, 0.5991138115095911, 0.6009475544728957, 0.6108938307533, 0.6316565296547284, 0.6404495093354053, 0.6405262429561641, 0.6425694115212065, 0.658671129040488, 0.6668423897406515, 0.6864391962767343, 0.7035635283169166, 0.7325106678655877, 0.7328050041291218, 0.7614304100703713, 0.7631239070049998, 0.7670021786149205, 0.7697753383420857, 0.7764360990307122, 0.7784918983501654, 0.7973920072996036, 0.819715865079681, 0.8506721053047003, 0.8813167497816289, 0.8824879447595726, 0.9185813970744787, 0.9231889896940375, 0.9237877978193884, 0.9255031346434324, 0.9293883502480845, 0.9294097332465232, 0.9463098243875633, 0.946325164889271, 0.9491397432856566, 0.9567595541247681, 0.9706712283358269, 0.9723580396501548, 0.9748360509016578, 0.9800193410444061, 0.980809631269599, 0.991517828651004, 0.9965400387585364] [0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8, 3BEOHQsMEFZ58VcNTOJYShTBpAPzbt, 4HX6feIvmNXBN7XGqgO4YVBkhu8GDI, 4JznSdBajNWhu4hRQwjV1FjTTxY68i, 52mKlRE3aHCBZtjECq6sY9OqVf8Dze, 56MZa5O1hVtX4c5sbnCfxuX5kDChqI, 6FPJlLAcaQ5uokyOWZ9HGdLZObFvOZ, 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW, 6oIXZuIPIqEoPBvFmbt2Nxy3tryGUE, 6x93sxYioWuq5c9Kkk8oTAAORM7cH0, 802bgTGl6Bk5TlkPYYTxp5JkKyaYUA, 8LIh0b6jmDGm87BmIyjdxNIpX4ugjD, 90gAtmGEeIqUTbo1ZrxCvWtsseukXC, 9UbObCsVkmYpJGcGrgfK90qOnwb2Lj, AFGCj7OWlEB5QfniEFgonMq90Tq5uH, ALuRhobVWbnQTTWZdSOk0iVe8oYFhW, Amn2K87Db5Es3dFQO9cw9cvpAM6h35, AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz, BJqx5WokrmrrezZA0dUbleMYkG5U2O, BPtQMxnuSPpxMExYV9YkDa6cAN7GP3, BsM5ZAYifRh5Lw3Y8X1r53I0cTJnfE, C2GT5KVyOPZpgKVl110TyZO0NcJ434, DuJNG8tufSqW0ZstHqWj3aGvFLMg4A, EcCuckwsF3gV1Ecgmh5v4KM8g1ozif, ErJFw6hzZ5fmI5r8bhE4JzlscnhKZU, F7NSTjWvQJyBburN7CXRUlbgp2dIrA, Fi4rJeTQq4eXj8Lxg3Hja5hBVTVV5u, H5j5ZHy1FGesOAHjkQEDYCucbpKWRu, HKSMQ9nTnwXCJIte1JrM1dtYnDtJ8g, IWl0G3ZlMNf7WT8yjIB49cx7MmYOmr, IZTkHMLvIKuiLjhDjYMmIHxh166we4, Ig1QcuKsjHXkproePdERo2w0mYzIqd, JHNgc2UCaiXOdmkxwDDyGhRlO0mnBQ, JN0VclewmjwYlSl8386MlWv5rEhWCz, JafwVLSVk5AVoXFuzclesQ000EE2k1, KJFcmTVjdkCMv94wYCtfHMFhzyRsmH, Ktb7GQ0N1DrxwkCkEUsTaIXk0xYinn, Ld2ej8NEv5zNcqU60FwpHeZKBhfpiV, LiEBxds3X0Uw0lxiYjDqrkAaAwoiIW, MXhhH1Var3OzzJCtI9VNyYvA0q8UyJ, MeSTAXq8gVxVjbEjgkvU9YLte0X9uE, NEhyk8uIx4kEULJGa8qIyFjjBcP2G6, O66j6PaYuZhEUtqV6fuU7TyjM2WxC5, OF7fQ37GzaZ5ikA2oMyvleKtgnLjXh, OPwBqCEK5PWTjWaiOyL45u2NLTaDWv, Oq6J4Rx6nde0YlhOIJkFsX2MsSvAQ0, Ow5PGpfTm4dXCfTDsXAOTatXRoAydR, QEHVvcP8gxI6EMJIrvcnIhgzPNjIvv, QJYm7YRA3YetcBHI5wkMZeLXVmfuNy, QYlaIAnJA6r8rlAb6f59wcxvcPcWFf, RilTlL1tKkPOUFuzmLydHAVZwv1OGl, Sfx0vxv1skzZWT1PqVdoRDdO6Sb6xH, TTQUwpMNSXZqVBKAFvXu7OlWvKXJKX, TtDKUZxzVxsq758G6AWPSYuZgVgbcl, VDhtJkYjAYPykCgOU9x3v7v3t4SO1a, VY0zXmXeksCT8BzvpzpPLbmU9Kp9Y4, Vp3gmWunM5A7wOC9YW2JroFqTWjvTi, WHmjWk2AY4c6m7DA4GitUx6nmb1yYS, XemNcT1xp61xcM1Qz3wZ1VECCnq06O, Z2sWcQr0qyCJRMHDpRy3aQr7PkHtkK, aDxBtor7Icd9C5hnTvvw5NrIre740e, akiiY5N0I44CMwEnBL6RTBk7BRkxEj, b3b9esRhTzFEawbs6XhpKnD9ojutHB, bgK1r6v3BCTh0aejJUhkA1Hn6idXGp, cBGc0kSm32ylBDnxogG727C0uhZEYZ, cq4WSAIFwx3wwTUS5bp1wCe71R6U5I, dVdvo6nUD5FgCgsbOZLds28RyGTpnx, e2Gh6Ov8XkXoFdJWhl0EjwEHlMDYyG, f9ALCzwDAKmdu7Rk2msJaB1wxe5IBX, fuyvs0w7WsKSlXqJ1e6HFSoLmx03AG, gTpyQnEODMcpsPnJMZC66gh33i3m0b, gpo8K5qtYePve6jyPt6xgJx4YOVjms, gxfHWUF8XgY2KdFxigxvNEXe2V2XMl, i6RQVXKUh7MzuGMDaNclUYnFUAireU, ioEncce3mPOXD2hWhpZpCPWGATG6GU, jQimhdepw3GKmioWUlVSWeBVRKFkY3, l7uwDoTepWwnAP0ufqtHJS3CRi7RfP, lqhzgLsXZ8JhtpeeUWWNbMz8PHI705, m6jD0LBIQWaMfenwRCTANI9eOdyyto, mhjME0zBHbrK6NMkytMTQzOssOa1gF, mzbkwXKrPeZnxg2Kn1LRF5hYSsmksS, nYVJnVicpGRqKZibHyBAmtmzBXAFfT, oHJMNvWuunsIMIWFnYG31RCfkOo2V7, oLZ21P2JEDooxV1pU31cIxQHEeeoLu, okOkcWflkNXIy4R8LzmySyY1EC3sYd, pLk3i59bZwd5KBZrI1FiweYTd5hteG, pTeu0WMjBRTaNRT15rLCuEh3tBJVc5, qnPOOmslCJaT45buUisMRnM0rc77EK, t6fQUjJejPcjc04wHvHTPe55S65B4V, ukOiFGGFnQJDHFgZxHMpvhD3zybF0M, ukyD7b0Efj7tNlFSRmzZ0IqkEzg2a8, waIGbOGl1PM6gnzZ4uuZt4E2yDWRHs, wwXqSGKLyBQyPkonlzBNYUJTCo4LRS, xipQ93429ksjNcXPX5326VSg1xJZcW, y7C453hRWd4E7ImjNDWlpexB8nUqjh, ydkwycaISlYSlEq3TlkS2m15I2pcp8] - -query ?? -SELECT - array_agg(c12 ORDER BY c12), - array_agg(c13 ORDER BY c13) -FROM aggregate_test_100 ----- -[0.01479305307777301, 0.02182578039211991, 0.03968347085780355, 0.04429073092078406, 0.047343434291126085, 0.04893135681998029, 0.0494924465469434, 0.05573662213439634, 0.05636955101974106, 0.061029375346466685, 0.07260475960924484, 0.09465635123783445, 0.12357539988406441, 0.152498292971736, 0.16301110515739792, 0.1640882545084913, 0.1754261586710173, 0.17592486905979987, 0.17909035118828576, 0.18628859265874176, 0.19113293583306745, 0.2145232647388039, 0.21535402343780985, 0.24899794314659673, 0.2537253407987472, 0.2667177795079635, 0.27159190516490006, 0.2739938529235548, 0.28534428578703896, 0.2944158618048994, 0.296036538664718, 0.3051364088814128, 0.30585375151301186, 0.3114712539863804, 0.3231750610081745, 0.32869374687050157, 0.33639590659276175, 0.3600766362333053, 0.36936304600612724, 0.38870280983958583, 0.39144436569161134, 0.40342283197779727, 0.4094218353587008, 0.40975383525297016, 0.42073125331890115, 0.4273123318932347, 0.42950521730777025, 0.4830878559436823, 0.5081765563442366, 0.5437595540422571, 0.5590205548347534, 0.5593249815276734, 0.5603062368164834, 0.560333188635217, 0.5614503754617461, 0.565352842229935, 0.574210838214554, 0.5759450483859969, 0.5773498217058918, 0.5991138115095911, 0.6009475544728957, 0.6108938307533, 0.6316565296547284, 0.6404495093354053, 0.6405262429561641, 0.6425694115212065, 0.658671129040488, 0.6668423897406515, 0.6864391962767343, 0.7035635283169166, 0.7325106678655877, 0.7328050041291218, 0.7614304100703713, 0.7631239070049998, 0.7670021786149205, 0.7697753383420857, 0.7764360990307122, 0.7784918983501654, 0.7973920072996036, 0.819715865079681, 0.8506721053047003, 0.8813167497816289, 0.8824879447595726, 0.9185813970744787, 0.9231889896940375, 0.9237877978193884, 0.9255031346434324, 0.9293883502480845, 0.9294097332465232, 0.9463098243875633, 0.946325164889271, 0.9491397432856566, 0.9567595541247681, 0.9706712283358269, 0.9723580396501548, 0.9748360509016578, 0.9800193410444061, 0.980809631269599, 0.991517828651004, 0.9965400387585364] [0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8, 3BEOHQsMEFZ58VcNTOJYShTBpAPzbt, 4HX6feIvmNXBN7XGqgO4YVBkhu8GDI, 4JznSdBajNWhu4hRQwjV1FjTTxY68i, 52mKlRE3aHCBZtjECq6sY9OqVf8Dze, 56MZa5O1hVtX4c5sbnCfxuX5kDChqI, 6FPJlLAcaQ5uokyOWZ9HGdLZObFvOZ, 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW, 6oIXZuIPIqEoPBvFmbt2Nxy3tryGUE, 6x93sxYioWuq5c9Kkk8oTAAORM7cH0, 802bgTGl6Bk5TlkPYYTxp5JkKyaYUA, 8LIh0b6jmDGm87BmIyjdxNIpX4ugjD, 90gAtmGEeIqUTbo1ZrxCvWtsseukXC, 9UbObCsVkmYpJGcGrgfK90qOnwb2Lj, AFGCj7OWlEB5QfniEFgonMq90Tq5uH, ALuRhobVWbnQTTWZdSOk0iVe8oYFhW, Amn2K87Db5Es3dFQO9cw9cvpAM6h35, AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz, BJqx5WokrmrrezZA0dUbleMYkG5U2O, BPtQMxnuSPpxMExYV9YkDa6cAN7GP3, BsM5ZAYifRh5Lw3Y8X1r53I0cTJnfE, C2GT5KVyOPZpgKVl110TyZO0NcJ434, DuJNG8tufSqW0ZstHqWj3aGvFLMg4A, EcCuckwsF3gV1Ecgmh5v4KM8g1ozif, ErJFw6hzZ5fmI5r8bhE4JzlscnhKZU, F7NSTjWvQJyBburN7CXRUlbgp2dIrA, Fi4rJeTQq4eXj8Lxg3Hja5hBVTVV5u, H5j5ZHy1FGesOAHjkQEDYCucbpKWRu, HKSMQ9nTnwXCJIte1JrM1dtYnDtJ8g, IWl0G3ZlMNf7WT8yjIB49cx7MmYOmr, IZTkHMLvIKuiLjhDjYMmIHxh166we4, Ig1QcuKsjHXkproePdERo2w0mYzIqd, JHNgc2UCaiXOdmkxwDDyGhRlO0mnBQ, JN0VclewmjwYlSl8386MlWv5rEhWCz, JafwVLSVk5AVoXFuzclesQ000EE2k1, KJFcmTVjdkCMv94wYCtfHMFhzyRsmH, Ktb7GQ0N1DrxwkCkEUsTaIXk0xYinn, Ld2ej8NEv5zNcqU60FwpHeZKBhfpiV, LiEBxds3X0Uw0lxiYjDqrkAaAwoiIW, MXhhH1Var3OzzJCtI9VNyYvA0q8UyJ, MeSTAXq8gVxVjbEjgkvU9YLte0X9uE, NEhyk8uIx4kEULJGa8qIyFjjBcP2G6, O66j6PaYuZhEUtqV6fuU7TyjM2WxC5, OF7fQ37GzaZ5ikA2oMyvleKtgnLjXh, OPwBqCEK5PWTjWaiOyL45u2NLTaDWv, Oq6J4Rx6nde0YlhOIJkFsX2MsSvAQ0, Ow5PGpfTm4dXCfTDsXAOTatXRoAydR, QEHVvcP8gxI6EMJIrvcnIhgzPNjIvv, QJYm7YRA3YetcBHI5wkMZeLXVmfuNy, QYlaIAnJA6r8rlAb6f59wcxvcPcWFf, RilTlL1tKkPOUFuzmLydHAVZwv1OGl, Sfx0vxv1skzZWT1PqVdoRDdO6Sb6xH, TTQUwpMNSXZqVBKAFvXu7OlWvKXJKX, TtDKUZxzVxsq758G6AWPSYuZgVgbcl, VDhtJkYjAYPykCgOU9x3v7v3t4SO1a, VY0zXmXeksCT8BzvpzpPLbmU9Kp9Y4, Vp3gmWunM5A7wOC9YW2JroFqTWjvTi, WHmjWk2AY4c6m7DA4GitUx6nmb1yYS, XemNcT1xp61xcM1Qz3wZ1VECCnq06O, Z2sWcQr0qyCJRMHDpRy3aQr7PkHtkK, aDxBtor7Icd9C5hnTvvw5NrIre740e, akiiY5N0I44CMwEnBL6RTBk7BRkxEj, b3b9esRhTzFEawbs6XhpKnD9ojutHB, bgK1r6v3BCTh0aejJUhkA1Hn6idXGp, cBGc0kSm32ylBDnxogG727C0uhZEYZ, cq4WSAIFwx3wwTUS5bp1wCe71R6U5I, dVdvo6nUD5FgCgsbOZLds28RyGTpnx, e2Gh6Ov8XkXoFdJWhl0EjwEHlMDYyG, f9ALCzwDAKmdu7Rk2msJaB1wxe5IBX, fuyvs0w7WsKSlXqJ1e6HFSoLmx03AG, gTpyQnEODMcpsPnJMZC66gh33i3m0b, gpo8K5qtYePve6jyPt6xgJx4YOVjms, gxfHWUF8XgY2KdFxigxvNEXe2V2XMl, i6RQVXKUh7MzuGMDaNclUYnFUAireU, ioEncce3mPOXD2hWhpZpCPWGATG6GU, jQimhdepw3GKmioWUlVSWeBVRKFkY3, l7uwDoTepWwnAP0ufqtHJS3CRi7RfP, lqhzgLsXZ8JhtpeeUWWNbMz8PHI705, m6jD0LBIQWaMfenwRCTANI9eOdyyto, mhjME0zBHbrK6NMkytMTQzOssOa1gF, mzbkwXKrPeZnxg2Kn1LRF5hYSsmksS, nYVJnVicpGRqKZibHyBAmtmzBXAFfT, oHJMNvWuunsIMIWFnYG31RCfkOo2V7, oLZ21P2JEDooxV1pU31cIxQHEeeoLu, okOkcWflkNXIy4R8LzmySyY1EC3sYd, pLk3i59bZwd5KBZrI1FiweYTd5hteG, pTeu0WMjBRTaNRT15rLCuEh3tBJVc5, qnPOOmslCJaT45buUisMRnM0rc77EK, t6fQUjJejPcjc04wHvHTPe55S65B4V, ukOiFGGFnQJDHFgZxHMpvhD3zybF0M, ukyD7b0Efj7tNlFSRmzZ0IqkEzg2a8, waIGbOGl1PM6gnzZ4uuZt4E2yDWRHs, wwXqSGKLyBQyPkonlzBNYUJTCo4LRS, xipQ93429ksjNcXPX5326VSg1xJZcW, y7C453hRWd4E7ImjNDWlpexB8nUqjh, ydkwycaISlYSlEq3TlkS2m15I2pcp8] - -query ?? rowsort -with tbl as (SELECT * FROM (VALUES ('xxx', 'yyy'), ('xxx', 'yyy'), ('xxx2', 'yyy2')) AS t(x, y)) -select - array_agg(distinct x order by x) as x_agg, - array_agg(distinct y order by y) as y_agg -from tbl -group by all ----- -[xxx, xxx2] [yyy, yyy2] - -query ?? -SELECT - (SELECT array_agg(DISTINCT c12 ORDER BY c12) FROM aggregate_test_100), - (SELECT array_agg(DISTINCT c13 ORDER BY c13) FROM aggregate_test_100) ----- -[0.01479305307777301, 0.02182578039211991, 0.03968347085780355, 0.04429073092078406, 0.047343434291126085, 0.04893135681998029, 0.0494924465469434, 0.05573662213439634, 0.05636955101974106, 0.061029375346466685, 0.07260475960924484, 0.09465635123783445, 0.12357539988406441, 0.152498292971736, 0.16301110515739792, 0.1640882545084913, 0.1754261586710173, 0.17592486905979987, 0.17909035118828576, 0.18628859265874176, 0.19113293583306745, 0.2145232647388039, 0.21535402343780985, 0.24899794314659673, 0.2537253407987472, 0.2667177795079635, 0.27159190516490006, 0.2739938529235548, 0.28534428578703896, 0.2944158618048994, 0.296036538664718, 0.3051364088814128, 0.30585375151301186, 0.3114712539863804, 0.3231750610081745, 0.32869374687050157, 0.33639590659276175, 0.3600766362333053, 0.36936304600612724, 0.38870280983958583, 0.39144436569161134, 0.40342283197779727, 0.4094218353587008, 0.40975383525297016, 0.42073125331890115, 0.4273123318932347, 0.42950521730777025, 0.4830878559436823, 0.5081765563442366, 0.5437595540422571, 0.5590205548347534, 0.5593249815276734, 0.5603062368164834, 0.560333188635217, 0.5614503754617461, 0.565352842229935, 0.574210838214554, 0.5759450483859969, 0.5773498217058918, 0.5991138115095911, 0.6009475544728957, 0.6108938307533, 0.6316565296547284, 0.6404495093354053, 0.6405262429561641, 0.6425694115212065, 0.658671129040488, 0.6668423897406515, 0.6864391962767343, 0.7035635283169166, 0.7325106678655877, 0.7328050041291218, 0.7614304100703713, 0.7631239070049998, 0.7670021786149205, 0.7697753383420857, 0.7764360990307122, 0.7784918983501654, 0.7973920072996036, 0.819715865079681, 0.8506721053047003, 0.8813167497816289, 0.8824879447595726, 0.9185813970744787, 0.9231889896940375, 0.9237877978193884, 0.9255031346434324, 0.9293883502480845, 0.9294097332465232, 0.9463098243875633, 0.946325164889271, 0.9491397432856566, 0.9567595541247681, 0.9706712283358269, 0.9723580396501548, 0.9748360509016578, 0.9800193410444061, 0.980809631269599, 0.991517828651004, 0.9965400387585364] [0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8, 3BEOHQsMEFZ58VcNTOJYShTBpAPzbt, 4HX6feIvmNXBN7XGqgO4YVBkhu8GDI, 4JznSdBajNWhu4hRQwjV1FjTTxY68i, 52mKlRE3aHCBZtjECq6sY9OqVf8Dze, 56MZa5O1hVtX4c5sbnCfxuX5kDChqI, 6FPJlLAcaQ5uokyOWZ9HGdLZObFvOZ, 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW, 6oIXZuIPIqEoPBvFmbt2Nxy3tryGUE, 6x93sxYioWuq5c9Kkk8oTAAORM7cH0, 802bgTGl6Bk5TlkPYYTxp5JkKyaYUA, 8LIh0b6jmDGm87BmIyjdxNIpX4ugjD, 90gAtmGEeIqUTbo1ZrxCvWtsseukXC, 9UbObCsVkmYpJGcGrgfK90qOnwb2Lj, AFGCj7OWlEB5QfniEFgonMq90Tq5uH, ALuRhobVWbnQTTWZdSOk0iVe8oYFhW, Amn2K87Db5Es3dFQO9cw9cvpAM6h35, AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz, BJqx5WokrmrrezZA0dUbleMYkG5U2O, BPtQMxnuSPpxMExYV9YkDa6cAN7GP3, BsM5ZAYifRh5Lw3Y8X1r53I0cTJnfE, C2GT5KVyOPZpgKVl110TyZO0NcJ434, DuJNG8tufSqW0ZstHqWj3aGvFLMg4A, EcCuckwsF3gV1Ecgmh5v4KM8g1ozif, ErJFw6hzZ5fmI5r8bhE4JzlscnhKZU, F7NSTjWvQJyBburN7CXRUlbgp2dIrA, Fi4rJeTQq4eXj8Lxg3Hja5hBVTVV5u, H5j5ZHy1FGesOAHjkQEDYCucbpKWRu, HKSMQ9nTnwXCJIte1JrM1dtYnDtJ8g, IWl0G3ZlMNf7WT8yjIB49cx7MmYOmr, IZTkHMLvIKuiLjhDjYMmIHxh166we4, Ig1QcuKsjHXkproePdERo2w0mYzIqd, JHNgc2UCaiXOdmkxwDDyGhRlO0mnBQ, JN0VclewmjwYlSl8386MlWv5rEhWCz, JafwVLSVk5AVoXFuzclesQ000EE2k1, KJFcmTVjdkCMv94wYCtfHMFhzyRsmH, Ktb7GQ0N1DrxwkCkEUsTaIXk0xYinn, Ld2ej8NEv5zNcqU60FwpHeZKBhfpiV, LiEBxds3X0Uw0lxiYjDqrkAaAwoiIW, MXhhH1Var3OzzJCtI9VNyYvA0q8UyJ, MeSTAXq8gVxVjbEjgkvU9YLte0X9uE, NEhyk8uIx4kEULJGa8qIyFjjBcP2G6, O66j6PaYuZhEUtqV6fuU7TyjM2WxC5, OF7fQ37GzaZ5ikA2oMyvleKtgnLjXh, OPwBqCEK5PWTjWaiOyL45u2NLTaDWv, Oq6J4Rx6nde0YlhOIJkFsX2MsSvAQ0, Ow5PGpfTm4dXCfTDsXAOTatXRoAydR, QEHVvcP8gxI6EMJIrvcnIhgzPNjIvv, QJYm7YRA3YetcBHI5wkMZeLXVmfuNy, QYlaIAnJA6r8rlAb6f59wcxvcPcWFf, RilTlL1tKkPOUFuzmLydHAVZwv1OGl, Sfx0vxv1skzZWT1PqVdoRDdO6Sb6xH, TTQUwpMNSXZqVBKAFvXu7OlWvKXJKX, TtDKUZxzVxsq758G6AWPSYuZgVgbcl, VDhtJkYjAYPykCgOU9x3v7v3t4SO1a, VY0zXmXeksCT8BzvpzpPLbmU9Kp9Y4, Vp3gmWunM5A7wOC9YW2JroFqTWjvTi, WHmjWk2AY4c6m7DA4GitUx6nmb1yYS, XemNcT1xp61xcM1Qz3wZ1VECCnq06O, Z2sWcQr0qyCJRMHDpRy3aQr7PkHtkK, aDxBtor7Icd9C5hnTvvw5NrIre740e, akiiY5N0I44CMwEnBL6RTBk7BRkxEj, b3b9esRhTzFEawbs6XhpKnD9ojutHB, bgK1r6v3BCTh0aejJUhkA1Hn6idXGp, cBGc0kSm32ylBDnxogG727C0uhZEYZ, cq4WSAIFwx3wwTUS5bp1wCe71R6U5I, dVdvo6nUD5FgCgsbOZLds28RyGTpnx, e2Gh6Ov8XkXoFdJWhl0EjwEHlMDYyG, f9ALCzwDAKmdu7Rk2msJaB1wxe5IBX, fuyvs0w7WsKSlXqJ1e6HFSoLmx03AG, gTpyQnEODMcpsPnJMZC66gh33i3m0b, gpo8K5qtYePve6jyPt6xgJx4YOVjms, gxfHWUF8XgY2KdFxigxvNEXe2V2XMl, i6RQVXKUh7MzuGMDaNclUYnFUAireU, ioEncce3mPOXD2hWhpZpCPWGATG6GU, jQimhdepw3GKmioWUlVSWeBVRKFkY3, l7uwDoTepWwnAP0ufqtHJS3CRi7RfP, lqhzgLsXZ8JhtpeeUWWNbMz8PHI705, m6jD0LBIQWaMfenwRCTANI9eOdyyto, mhjME0zBHbrK6NMkytMTQzOssOa1gF, mzbkwXKrPeZnxg2Kn1LRF5hYSsmksS, nYVJnVicpGRqKZibHyBAmtmzBXAFfT, oHJMNvWuunsIMIWFnYG31RCfkOo2V7, oLZ21P2JEDooxV1pU31cIxQHEeeoLu, okOkcWflkNXIy4R8LzmySyY1EC3sYd, pLk3i59bZwd5KBZrI1FiweYTd5hteG, pTeu0WMjBRTaNRT15rLCuEh3tBJVc5, qnPOOmslCJaT45buUisMRnM0rc77EK, t6fQUjJejPcjc04wHvHTPe55S65B4V, ukOiFGGFnQJDHFgZxHMpvhD3zybF0M, ukyD7b0Efj7tNlFSRmzZ0IqkEzg2a8, waIGbOGl1PM6gnzZ4uuZt4E2yDWRHs, wwXqSGKLyBQyPkonlzBNYUJTCo4LRS, xipQ93429ksjNcXPX5326VSg1xJZcW, y7C453hRWd4E7ImjNDWlpexB8nUqjh, ydkwycaISlYSlEq3TlkS2m15I2pcp8] - -query ?? -SELECT - array_agg(DISTINCT c12 ORDER BY c12), - array_agg(DISTINCT c13 ORDER BY c13) -FROM aggregate_test_100 ----- -[0.01479305307777301, 0.02182578039211991, 0.03968347085780355, 0.04429073092078406, 0.047343434291126085, 0.04893135681998029, 0.0494924465469434, 0.05573662213439634, 0.05636955101974106, 0.061029375346466685, 0.07260475960924484, 0.09465635123783445, 0.12357539988406441, 0.152498292971736, 0.16301110515739792, 0.1640882545084913, 0.1754261586710173, 0.17592486905979987, 0.17909035118828576, 0.18628859265874176, 0.19113293583306745, 0.2145232647388039, 0.21535402343780985, 0.24899794314659673, 0.2537253407987472, 0.2667177795079635, 0.27159190516490006, 0.2739938529235548, 0.28534428578703896, 0.2944158618048994, 0.296036538664718, 0.3051364088814128, 0.30585375151301186, 0.3114712539863804, 0.3231750610081745, 0.32869374687050157, 0.33639590659276175, 0.3600766362333053, 0.36936304600612724, 0.38870280983958583, 0.39144436569161134, 0.40342283197779727, 0.4094218353587008, 0.40975383525297016, 0.42073125331890115, 0.4273123318932347, 0.42950521730777025, 0.4830878559436823, 0.5081765563442366, 0.5437595540422571, 0.5590205548347534, 0.5593249815276734, 0.5603062368164834, 0.560333188635217, 0.5614503754617461, 0.565352842229935, 0.574210838214554, 0.5759450483859969, 0.5773498217058918, 0.5991138115095911, 0.6009475544728957, 0.6108938307533, 0.6316565296547284, 0.6404495093354053, 0.6405262429561641, 0.6425694115212065, 0.658671129040488, 0.6668423897406515, 0.6864391962767343, 0.7035635283169166, 0.7325106678655877, 0.7328050041291218, 0.7614304100703713, 0.7631239070049998, 0.7670021786149205, 0.7697753383420857, 0.7764360990307122, 0.7784918983501654, 0.7973920072996036, 0.819715865079681, 0.8506721053047003, 0.8813167497816289, 0.8824879447595726, 0.9185813970744787, 0.9231889896940375, 0.9237877978193884, 0.9255031346434324, 0.9293883502480845, 0.9294097332465232, 0.9463098243875633, 0.946325164889271, 0.9491397432856566, 0.9567595541247681, 0.9706712283358269, 0.9723580396501548, 0.9748360509016578, 0.9800193410444061, 0.980809631269599, 0.991517828651004, 0.9965400387585364] [0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8, 3BEOHQsMEFZ58VcNTOJYShTBpAPzbt, 4HX6feIvmNXBN7XGqgO4YVBkhu8GDI, 4JznSdBajNWhu4hRQwjV1FjTTxY68i, 52mKlRE3aHCBZtjECq6sY9OqVf8Dze, 56MZa5O1hVtX4c5sbnCfxuX5kDChqI, 6FPJlLAcaQ5uokyOWZ9HGdLZObFvOZ, 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW, 6oIXZuIPIqEoPBvFmbt2Nxy3tryGUE, 6x93sxYioWuq5c9Kkk8oTAAORM7cH0, 802bgTGl6Bk5TlkPYYTxp5JkKyaYUA, 8LIh0b6jmDGm87BmIyjdxNIpX4ugjD, 90gAtmGEeIqUTbo1ZrxCvWtsseukXC, 9UbObCsVkmYpJGcGrgfK90qOnwb2Lj, AFGCj7OWlEB5QfniEFgonMq90Tq5uH, ALuRhobVWbnQTTWZdSOk0iVe8oYFhW, Amn2K87Db5Es3dFQO9cw9cvpAM6h35, AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz, BJqx5WokrmrrezZA0dUbleMYkG5U2O, BPtQMxnuSPpxMExYV9YkDa6cAN7GP3, BsM5ZAYifRh5Lw3Y8X1r53I0cTJnfE, C2GT5KVyOPZpgKVl110TyZO0NcJ434, DuJNG8tufSqW0ZstHqWj3aGvFLMg4A, EcCuckwsF3gV1Ecgmh5v4KM8g1ozif, ErJFw6hzZ5fmI5r8bhE4JzlscnhKZU, F7NSTjWvQJyBburN7CXRUlbgp2dIrA, Fi4rJeTQq4eXj8Lxg3Hja5hBVTVV5u, H5j5ZHy1FGesOAHjkQEDYCucbpKWRu, HKSMQ9nTnwXCJIte1JrM1dtYnDtJ8g, IWl0G3ZlMNf7WT8yjIB49cx7MmYOmr, IZTkHMLvIKuiLjhDjYMmIHxh166we4, Ig1QcuKsjHXkproePdERo2w0mYzIqd, JHNgc2UCaiXOdmkxwDDyGhRlO0mnBQ, JN0VclewmjwYlSl8386MlWv5rEhWCz, JafwVLSVk5AVoXFuzclesQ000EE2k1, KJFcmTVjdkCMv94wYCtfHMFhzyRsmH, Ktb7GQ0N1DrxwkCkEUsTaIXk0xYinn, Ld2ej8NEv5zNcqU60FwpHeZKBhfpiV, LiEBxds3X0Uw0lxiYjDqrkAaAwoiIW, MXhhH1Var3OzzJCtI9VNyYvA0q8UyJ, MeSTAXq8gVxVjbEjgkvU9YLte0X9uE, NEhyk8uIx4kEULJGa8qIyFjjBcP2G6, O66j6PaYuZhEUtqV6fuU7TyjM2WxC5, OF7fQ37GzaZ5ikA2oMyvleKtgnLjXh, OPwBqCEK5PWTjWaiOyL45u2NLTaDWv, Oq6J4Rx6nde0YlhOIJkFsX2MsSvAQ0, Ow5PGpfTm4dXCfTDsXAOTatXRoAydR, QEHVvcP8gxI6EMJIrvcnIhgzPNjIvv, QJYm7YRA3YetcBHI5wkMZeLXVmfuNy, QYlaIAnJA6r8rlAb6f59wcxvcPcWFf, RilTlL1tKkPOUFuzmLydHAVZwv1OGl, Sfx0vxv1skzZWT1PqVdoRDdO6Sb6xH, TTQUwpMNSXZqVBKAFvXu7OlWvKXJKX, TtDKUZxzVxsq758G6AWPSYuZgVgbcl, VDhtJkYjAYPykCgOU9x3v7v3t4SO1a, VY0zXmXeksCT8BzvpzpPLbmU9Kp9Y4, Vp3gmWunM5A7wOC9YW2JroFqTWjvTi, WHmjWk2AY4c6m7DA4GitUx6nmb1yYS, XemNcT1xp61xcM1Qz3wZ1VECCnq06O, Z2sWcQr0qyCJRMHDpRy3aQr7PkHtkK, aDxBtor7Icd9C5hnTvvw5NrIre740e, akiiY5N0I44CMwEnBL6RTBk7BRkxEj, b3b9esRhTzFEawbs6XhpKnD9ojutHB, bgK1r6v3BCTh0aejJUhkA1Hn6idXGp, cBGc0kSm32ylBDnxogG727C0uhZEYZ, cq4WSAIFwx3wwTUS5bp1wCe71R6U5I, dVdvo6nUD5FgCgsbOZLds28RyGTpnx, e2Gh6Ov8XkXoFdJWhl0EjwEHlMDYyG, f9ALCzwDAKmdu7Rk2msJaB1wxe5IBX, fuyvs0w7WsKSlXqJ1e6HFSoLmx03AG, gTpyQnEODMcpsPnJMZC66gh33i3m0b, gpo8K5qtYePve6jyPt6xgJx4YOVjms, gxfHWUF8XgY2KdFxigxvNEXe2V2XMl, i6RQVXKUh7MzuGMDaNclUYnFUAireU, ioEncce3mPOXD2hWhpZpCPWGATG6GU, jQimhdepw3GKmioWUlVSWeBVRKFkY3, l7uwDoTepWwnAP0ufqtHJS3CRi7RfP, lqhzgLsXZ8JhtpeeUWWNbMz8PHI705, m6jD0LBIQWaMfenwRCTANI9eOdyyto, mhjME0zBHbrK6NMkytMTQzOssOa1gF, mzbkwXKrPeZnxg2Kn1LRF5hYSsmksS, nYVJnVicpGRqKZibHyBAmtmzBXAFfT, oHJMNvWuunsIMIWFnYG31RCfkOo2V7, oLZ21P2JEDooxV1pU31cIxQHEeeoLu, okOkcWflkNXIy4R8LzmySyY1EC3sYd, pLk3i59bZwd5KBZrI1FiweYTd5hteG, pTeu0WMjBRTaNRT15rLCuEh3tBJVc5, qnPOOmslCJaT45buUisMRnM0rc77EK, t6fQUjJejPcjc04wHvHTPe55S65B4V, ukOiFGGFnQJDHFgZxHMpvhD3zybF0M, ukyD7b0Efj7tNlFSRmzZ0IqkEzg2a8, waIGbOGl1PM6gnzZ4uuZt4E2yDWRHs, wwXqSGKLyBQyPkonlzBNYUJTCo4LRS, xipQ93429ksjNcXPX5326VSg1xJZcW, y7C453hRWd4E7ImjNDWlpexB8nUqjh, ydkwycaISlYSlEq3TlkS2m15I2pcp8] - -statement ok -CREATE EXTERNAL TABLE agg_order ( -c1 INT NOT NULL, -c2 INT NOT NULL, -c3 INT NOT NULL -) -STORED AS CSV -LOCATION '../core/tests/data/aggregate_agg_multi_order.csv' -OPTIONS ('format.has_header' 'true'); - -# test array_agg with order by multiple columns -query ? -select array_agg(c1 order by c2 desc, c3) from agg_order; ----- -[5, 6, 7, 8, 9, 1, 2, 3, 4, 10] - -query TT -explain select array_agg(c1 order by c2 desc, c3) from agg_order; ----- -logical_plan -01)Aggregate: groupBy=[[]], aggr=[[array_agg(agg_order.c1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]]] -02)--TableScan: agg_order projection=[c1, c2, c3] -physical_plan -01)AggregateExec: mode=Final, gby=[], aggr=[array_agg(agg_order.c1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] -02)--CoalescePartitionsExec -03)----AggregateExec: mode=Partial, gby=[], aggr=[array_agg(agg_order.c1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] -04)------SortExec: expr=[c2@1 DESC, c3@2 ASC NULLS LAST], preserve_partitioning=[true] -05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], file_type=csv, has_header=true - -# test array_agg_order with list data type -statement ok -CREATE TABLE array_agg_order_list_table AS VALUES - ('w', 2, [1,2,3], 10), - ('w', 1, [9,5,2], 20), - ('w', 1, [3,2,5], 30), - ('b', 2, [4,5,6], 20), - ('b', 1, [7,8,9], 30) -; - -query T? rowsort -select column1, array_agg(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; ----- -b [[7, 8, 9], [4, 5, 6]] -w [[3, 2, 5], [9, 5, 2], [1, 2, 3]] - -query T?? rowsort -select column1, first_value(column3 order by column2, column4 desc), last_value(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; ----- -b [7, 8, 9] [4, 5, 6] -w [3, 2, 5] [1, 2, 3] - -query T? rowsort -select column1, nth_value(column3, 2 order by column2, column4 desc) from array_agg_order_list_table group by column1; ----- -b [4, 5, 6] -w [9, 5, 2] - -query ? -select array_agg(DISTINCT column2 order by column2) from array_agg_order_list_table; ----- -[1, 2] - -query ? -select array_agg(DISTINCT column2 order by column2 desc) from array_agg_order_list_table; ----- -[2, 1] - -query ? -select array_agg(DISTINCT column2 + 1 order by column2 + 1 desc) from array_agg_order_list_table; ----- -[3, 2] - -query ? -select array_agg(DISTINCT column2 order by column2) from array_agg_order_list_table GROUP BY column1; ----- -[1, 2] -[1, 2] - -statement error In an aggregate with DISTINCT, ORDER BY expressions must appear in argument list -select array_agg(DISTINCT column2 order by column1) from array_agg_order_list_table; - -statement ok -drop table array_agg_order_list_table; - -# test array_agg_distinct with list data type -statement ok -CREATE TABLE array_agg_distinct_list_table AS VALUES - ('w', [0,1]), - ('w', [0,1]), - ('w', [1,0]), - ('b', [1,0]), - ('b', [1,0]), - ('b', [1,0]), - ('b', [0,1]), - (NULL, [0,1]), - ('b', NULL) -; - -# Apply array_sort to have deterministic result, higher dimension nested array also works but not for array sort, -# so they are covered in `datafusion/functions-aggregate/src/array_agg.rs` -query ?? -select array_sort(c1), array_sort(c2) from ( - select array_agg(distinct column1) as c1, array_agg(distinct column2) ignore nulls as c2 from array_agg_distinct_list_table -); ----- -[NULL, b, w] [[0, 1], [1, 0]] - -statement ok -drop table array_agg_distinct_list_table; - -# Test array_agg with DISTINCT and IGNORE NULLS (regression test for issue #19735) -query ? -SELECT array_sort(ARRAY_AGG(DISTINCT x IGNORE NULLS)) as result -FROM (VALUES (1), (2), (NULL), (2), (NULL), (1)) AS t(x); ----- -[1, 2] # Test that non-DISTINCT aggregates also preserve IGNORE NULLS when mixed with DISTINCT # This tests the two-phase aggregation rewrite in SingleDistinctToGroupBy @@ -456,75 +259,6 @@ FROM (VALUES ---- 2 [40, 30, 20, 10] -statement error This feature is not implemented: Calling array_agg: LIMIT not supported in function arguments: 1 -SELECT array_agg(c13 LIMIT 1) FROM aggregate_test_100 - - -# Test distinct aggregate function with merge batch -query II -with A as ( - select 1 as id, 2 as foo - UNION ALL - select 1, null - UNION ALL - select 1, null - UNION ALL - select 1, 3 - UNION ALL - select 1, 2 - ---- The order is non-deterministic, verify with length -) select array_length(array_agg(distinct a.foo)), sum(distinct 1) from A a group by a.id; ----- -3 1 - -# It has only AggregateExec with FinalPartitioned mode, so `merge_batch` is used -# If the plan is changed, whether the `merge_batch` is used should be verified to ensure the test coverage -query TT -explain with A as ( - select 1 as id, 2 as foo - UNION ALL - select 1, null - UNION ALL - select 1, null - UNION ALL - select 1, 3 - UNION ALL - select 1, 2 -) select array_length(array_agg(distinct a.foo)), sum(distinct 1) from A a group by a.id; ----- -logical_plan -01)Projection: array_length(array_agg(DISTINCT a.foo)), sum(DISTINCT Int64(1)) -02)--Aggregate: groupBy=[[a.id]], aggr=[[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))]] -03)----SubqueryAlias: a -04)------SubqueryAlias: a -05)--------Union -06)----------Projection: Int64(1) AS id, Int64(2) AS foo -07)------------EmptyRelation: rows=1 -08)----------Projection: Int64(1) AS id, Int64(NULL) AS foo -09)------------EmptyRelation: rows=1 -10)----------Projection: Int64(1) AS id, Int64(NULL) AS foo -11)------------EmptyRelation: rows=1 -12)----------Projection: Int64(1) AS id, Int64(3) AS foo -13)------------EmptyRelation: rows=1 -14)----------Projection: Int64(1) AS id, Int64(2) AS foo -15)------------EmptyRelation: rows=1 -physical_plan -01)ProjectionExec: expr=[array_length(array_agg(DISTINCT a.foo)@1) as array_length(array_agg(DISTINCT a.foo)), sum(DISTINCT Int64(1))@2 as sum(DISTINCT Int64(1))] -02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted -03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=5 -04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted -05)--------UnionExec -06)----------ProjectionExec: expr=[1 as id, 2 as foo] -07)------------PlaceholderRowExec -08)----------ProjectionExec: expr=[1 as id, NULL as foo] -09)------------PlaceholderRowExec -10)----------ProjectionExec: expr=[1 as id, NULL as foo] -11)------------PlaceholderRowExec -12)----------ProjectionExec: expr=[1 as id, 3 as foo] -13)------------PlaceholderRowExec -14)----------ProjectionExec: expr=[1 as id, 2 as foo] -15)------------PlaceholderRowExec - # FIX: custom absolute values # csv_query_avg_multi_batch @@ -1370,6 +1104,74 @@ ORDER BY tags, timestamp; statement ok DROP TABLE median_window_test; +# Regression: percentile_cont(DISTINCT ...) used to forward the extra +# percentile-argument column into the distinct-values buffer (which asserts a +# single input array), panicking on every distinct query. Plain aggregate: +statement ok +CREATE TABLE distinct_pct(id INT, x DOUBLE) AS VALUES + (1, 5), (2, 5), (3, 9); + +query R +SELECT percentile_cont(DISTINCT x, 0.5) FROM distinct_pct; +---- +7 + +# Regression: distinct sliding-window percentile must count value multiplicity +# on retract. Row 3's frame is {5, 9}; the row-1 `5` leaves the frame but the +# row-2 `5` remains, so the distinct set is still {5, 9} (median 7), not {9}. +query IR +SELECT id, percentile_cont(DISTINCT x, 0.5) + OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM distinct_pct +ORDER BY id; +---- +1 5 +2 5 +3 7 + +statement ok +DROP TABLE distinct_pct; + +# Regression: grouped percentile_cont(DISTINCT ...) forces two-phase +# (Partial + FinalPartitioned) aggregation, exercising the distinct +# accumulator's state()/merge_batch() paths. Duplicate values within a +# group must be de-duplicated across the per-partition merge. +# Group 1 distinct {1,5,9} -> median 5; group 2 distinct {3,7} -> median 5. +statement ok +CREATE TABLE grp_distinct_pct(g INT, x DOUBLE) AS VALUES + (1, 5), (1, 5), (1, 9), (1, 1), + (2, 7), (2, 7), (2, 3); + +query IR +SELECT g, percentile_cont(DISTINCT x, 0.5) FROM grp_distinct_pct GROUP BY g ORDER BY g; +---- +1 5 +2 5 + +statement ok +DROP TABLE grp_distinct_pct; + +# Regression: sliding-window percentile_cont(DISTINCT ...) over data with +# NULLs exercises the null_count() > 0 slow path in BOTH update_batch (a NULL +# row enters the frame) and retract_batch (a NULL row leaves the frame as the +# window slides). NULLs are ignored; distinct dedups the non-null values. +statement ok +CREATE TABLE distinct_pct_nulls(id INT, x DOUBLE) AS VALUES + (1, 5), (2, NULL), (3, 9), (4, 5); + +query IR +SELECT id, percentile_cont(DISTINCT x, 0.5) + OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM distinct_pct_nulls ORDER BY id; +---- +1 5 +2 5 +3 9 +4 7 + +statement ok +DROP TABLE distinct_pct_nulls; + query RT select approx_median(arrow_cast(col_f32, 'Float16')), arrow_typeof(approx_median(arrow_cast(col_f32, 'Float16'))) from median_table; ---- @@ -1398,6 +1200,11 @@ select approx_median(NULL), arrow_typeof(approx_median(NULL)) from median_table; ---- NULL Null +query ?T +select median(NULL), arrow_typeof(median(NULL)); +---- +NULL Null + # median decimal statement ok create table t(c decimal(10, 4)) as values (0.0001), (0.0002), (0.0003), (0.0004), (0.0005), (0.0006); @@ -1836,6 +1643,472 @@ SELECT approx_distinct(c14) AS a, approx_distinct(c15) AS b, approx_distinct(arr ---- 18 60 60 60 60 +# approx_distinct over Boolean: exact count via flag-pair accumulator (0..=2). +statement ok +CREATE TABLE approx_distinct_bool_test (g INT, b BOOLEAN) AS VALUES + (1, true), (1, true), (1, NULL), + (2, false), (2, false), + (3, true), (3, false), (3, NULL), (3, true), + (4, NULL), (4, NULL); + +query I +SELECT approx_distinct(b) FROM approx_distinct_bool_test WHERE g = 1; +---- +1 + +query I +SELECT approx_distinct(b) FROM approx_distinct_bool_test WHERE g = 2; +---- +1 + +query I +SELECT approx_distinct(b) FROM approx_distinct_bool_test WHERE g = 3; +---- +2 + +query I +SELECT approx_distinct(b) FROM approx_distinct_bool_test WHERE g = 4; +---- +0 + +query II +SELECT g, approx_distinct(b) FROM approx_distinct_bool_test GROUP BY g ORDER BY g; +---- +1 1 +2 1 +3 2 +4 0 + +query I +SELECT approx_distinct(b) FROM approx_distinct_bool_test; +---- +2 + +statement ok +DROP TABLE approx_distinct_bool_test; + +# Grouped approx_distinct uses a dedicated GroupsAccumulator (adaptive +# sparse -> dense HyperLogLog per group). Results are deterministic (the HLL uses +# a fixed hash seed); for these specific small inputs the 16384-register HLL +# estimates the true distinct count exactly. The key invariant is that the +# grouped path agrees with the scalar (no GROUP BY) path on the same data, which +# is checked explicitly below. +statement ok +CREATE TABLE approx_distinct_group_test (g INT, s VARCHAR, i INT) AS VALUES + (1, 'a', 10), (1, 'a', 10), (1, 'b', 20), + (2, 'c', 30), (2, 'd', 30), (2, 'c', 40), + (3, NULL, NULL), (3, NULL, NULL), + (4, 'e', 50); + +# Strings (Utf8): group 1 -> {a,b}=2, group 2 -> {c,d}=2, group 3 -> all null=0, group 4 -> {e}=1 +query II +SELECT g, approx_distinct(s) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + +# Utf8View takes the inline-view hashing path +query II +SELECT g, approx_distinct(arrow_cast(s, 'Utf8View')) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + +# BinaryView non-grouped +query I +SELECT approx_distinct(arrow_cast(arrow_cast(s, 'Binary'), 'BinaryView')) FROM approx_distinct_group_test WHERE g = 2; +---- +2 + + +# BinaryView grouped +query II +SELECT g, approx_distinct(arrow_cast(arrow_cast(s, 'Binary'), 'BinaryView')) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + + +# FixedSizeBinary non-grouped +query I +SELECT approx_distinct(arrow_cast(arrow_cast(s, 'Binary'), 'FixedSizeBinary(1)')) FROM approx_distinct_group_test WHERE g = 2; +---- +2 + +# FixedSizeBinary grouped +query II +SELECT g, approx_distinct(arrow_cast(arrow_cast(s, 'Binary'), 'FixedSizeBinary(1)')) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + + +# List +statement ok +CREATE TABLE approx_distinct_list_test (g INT, l INT[]) AS VALUES + (1, [1, 2]), (1, [1, 2]), (1, [3, 4]), + (2, [5, 6]), (2, NULL), + (3, NULL), (3, NULL), + (4, [7, 8]); + +# List non-grouped +query I +SELECT approx_distinct(l) FROM approx_distinct_list_test WHERE g = 1; +---- +2 + +# List grouped +# Group 1 -> {[1,2],[3,4]}=2, +# Group 2 -> {[5,6]}=1 (NULL excluded), +# Group 3 -> all null=0, +# Group 4 -> {[7,8]}=1 +query II +SELECT g, approx_distinct(l) FROM approx_distinct_list_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct lists across groups are still counted overall. +query I +SELECT approx_distinct(l) FROM approx_distinct_list_test; +---- +4 + +# LargeList non-grouped +query I +SELECT approx_distinct(arrow_cast(l, 'LargeList(Int32)')) FROM approx_distinct_list_test WHERE g = 1; +---- +2 + +# LargeList grouped +query II +SELECT g, approx_distinct(arrow_cast(l, 'LargeList(Int32)')) FROM approx_distinct_list_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct lists across groups are still counted overall. +query I +SELECT approx_distinct(arrow_cast(l, 'LargeList(Int32)')) FROM approx_distinct_list_test; +---- +4 + +# ListView non-grouped +query I +SELECT approx_distinct(arrow_cast(l, 'ListView(Int32)')) FROM approx_distinct_list_test WHERE g = 1; +---- +2 + +# ListView grouped +query II +SELECT g, approx_distinct(arrow_cast(l, 'ListView(Int32)')) FROM approx_distinct_list_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct lists across groups are still counted overall. +query I +SELECT approx_distinct(arrow_cast(l, 'ListView(Int32)')) FROM approx_distinct_list_test; +---- +4 + +# LargeListView non-grouped +query I +SELECT approx_distinct(arrow_cast(l, 'LargeListView(Int32)')) FROM approx_distinct_list_test WHERE g = 1; +---- +2 + +# LargeListView grouped +query II +SELECT g, approx_distinct(arrow_cast(l, 'LargeListView(Int32)')) FROM approx_distinct_list_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct lists across groups are still counted overall. +query I +SELECT approx_distinct(arrow_cast(l, 'LargeListView(Int32)')) FROM approx_distinct_list_test; +---- +4 + + +# FixedSizeList non-grouped +query I +SELECT approx_distinct(arrow_cast(l, 'FixedSizeList(2, Int32)')) FROM approx_distinct_list_test WHERE g = 1; +---- +2 + +# FixedSizeList grouped +query II +SELECT g, approx_distinct(arrow_cast(l, 'FixedSizeList(2, Int32)')) FROM approx_distinct_list_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct lists across groups are still counted overall. +query I +SELECT approx_distinct(arrow_cast(l, 'FixedSizeList(2, Int32)')) FROM approx_distinct_list_test; +---- +4 + +statement ok +DROP TABLE approx_distinct_list_test; + +# Map +statement ok +CREATE TABLE approx_distinct_map_test AS SELECT * FROM (VALUES + (1, MAP {'a': 1, 'b': 2}), (1, MAP {'a': 1, 'b': 2}), (1, MAP {'c': 3}), + (2, MAP {'d': 4}), (2, NULL), + (3, NULL), (3, NULL), + (4, MAP {'e': 5}) +) AS t(g, m); + +# Map non-grouped +query I +SELECT approx_distinct(m) FROM approx_distinct_map_test WHERE g = 1; +---- +2 + +# Map grouped +# Group 1 -> {{a:1,b:2},{c:3}}=2, +# Group 2 -> {{d:4}}=1 (NULL excluded), +# Group 3 -> all null=0, +# Group 4 -> {{e:5}}=1 +query II +SELECT g, approx_distinct(m) FROM approx_distinct_map_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct maps across groups are still counted overall. +query I +SELECT approx_distinct(m) FROM approx_distinct_map_test; +---- +4 + +statement ok +DROP TABLE approx_distinct_map_test; + +# Struct +statement ok +CREATE TABLE approx_distinct_struct_test AS SELECT * FROM (VALUES + (1, named_struct('a', 1, 'b', 2)), (1, named_struct('a', 1, 'b', 2)), (1, named_struct('a', 3, 'b', 3)), + (2, named_struct('a', 4, 'b', 4)), (2, NULL), + (3, NULL), (3, NULL), + (4, named_struct('a', 5, 'b', 5)) +) AS t(g, s); + +# Struct non-grouped +query I +SELECT approx_distinct(s) FROM approx_distinct_struct_test WHERE g = 1; +---- +2 + +# Struct grouped +# Group 1 -> {{a:1,b:2},{a:3,b:3}}=2, +# Group 2 -> {{a:4,b:4}}=1 (NULL excluded), +# Group 3 -> all null=0, +# Group 4 -> {{a:5,b:5}}=1 +query II +SELECT g, approx_distinct(s) FROM approx_distinct_struct_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct structs across groups are still counted overall. +query I +SELECT approx_distinct(s) FROM approx_distinct_struct_test; +---- +4 + +statement ok +DROP TABLE approx_distinct_struct_test; + +# Union +# `approx_distinct_union_test` (g INT, u UNION) is registered +# in test_context.rs because a union value cannot be constructed from SQL. + +# Union non-grouped +query I +SELECT approx_distinct(u) FROM approx_distinct_union_test WHERE g = 1; +---- +2 + +# Union grouped +# Group 1 -> {i:1, i:1, s:"x"}=2, +# Group 2 -> {s:"y"}=1 (NULL excluded), +# Group 3 -> all null=0, +# Group 4 -> {i:5}=1 +query II +SELECT g, approx_distinct(u) FROM approx_distinct_union_test GROUP BY g ORDER BY g; +---- +1 2 +2 1 +3 0 +4 1 + +# The non-group path must agree with the grouped path on +# the same data, and distinct union values across groups are still counted overall. +query I +SELECT approx_distinct(u) FROM approx_distinct_union_test; +---- +4 + +# Integers (Int32): group 1 -> {10,20}=2, group 2 -> {30,40}=2, group 3 -> 0, group 4 -> {50}=1 +query II +SELECT g, approx_distinct(i) FROM approx_distinct_group_test GROUP BY g ORDER BY g; +---- +1 2 +2 2 +3 0 +4 1 + +# Invariant: the scalar (no GROUP BY) path must agree with the grouped path on +# the same data. The grouped result for g = 2 above is 2, and so is the scalar +# result over only g = 2's rows. +query I +SELECT approx_distinct(s) FROM approx_distinct_group_test WHERE g = 2; +---- +2 + +query I +SELECT approx_distinct(i) FROM approx_distinct_group_test WHERE g = 2; +---- +2 + +statement ok +DROP TABLE approx_distinct_group_test; + +# Grouped approx_distinct that crosses the sparse -> dense promotion threshold: +# 2000 distinct values in group 0 and 2000 in group 1. The estimate should be +# within HyperLogLog's error margin (~0.8%) of the true cardinality. +statement ok +CREATE TABLE approx_distinct_dense_test AS + SELECT (v % 2) AS g, v AS i FROM generate_series(0, 3999) AS t(v); + +query B +SELECT min(c) > 1900 AND max(c) < 2100 FROM ( + SELECT g, approx_distinct(i) AS c FROM approx_distinct_dense_test GROUP BY g +); +---- +true + +statement ok +DROP TABLE approx_distinct_dense_test; + +# This test runs approx_distinct over decimal128 and decimal256 for the scalar and the grouped path. +statement ok +CREATE TABLE approx_distinct_decimal_test (g INT, dec128 DECIMAL(20, 2), dec256 DECIMAL(40, 2)) AS VALUES + (1, 12345678901234.56, 12345678901234567890123456.78), + (1, 98765432109876.54, 98765432109876543210987654.32), + (1, 98765432109876.54, 98765432109876543210987654.32), + (2, 55555555555555.55, 55555555555555555555555555.55), + (2, -0.0, -0.0), + (2, 0.0, 0.0); + +# Scalar path +query II +SELECT approx_distinct(dec128), approx_distinct(dec256) FROM approx_distinct_decimal_test; +---- +4 4 + +# Grouped path +query III +SELECT g, approx_distinct(dec128), approx_distinct(dec256) +FROM approx_distinct_decimal_test GROUP BY g ORDER BY g; +---- +1 2 2 +2 2 2 + +statement ok +DROP TABLE approx_distinct_decimal_test; + +# This test runs approx_distinct over all four Duration units for the scalar and the grouped path. +statement ok +CREATE TABLE approx_distinct_duration_test AS VALUES + (1, arrow_cast(1, 'Duration(Second)'), arrow_cast(1, 'Duration(Millisecond)'), arrow_cast(1, 'Duration(Microsecond)'), arrow_cast(1, 'Duration(Nanosecond)')), + (1, arrow_cast(2, 'Duration(Second)'), arrow_cast(2, 'Duration(Millisecond)'), arrow_cast(2, 'Duration(Microsecond)'), arrow_cast(2, 'Duration(Nanosecond)')), + (1, arrow_cast(2, 'Duration(Second)'), arrow_cast(2, 'Duration(Millisecond)'), arrow_cast(2, 'Duration(Microsecond)'), arrow_cast(2, 'Duration(Nanosecond)')), + (2, arrow_cast(3, 'Duration(Second)'), arrow_cast(3, 'Duration(Millisecond)'), arrow_cast(3, 'Duration(Microsecond)'), arrow_cast(3, 'Duration(Nanosecond)')), + (2, arrow_cast(0, 'Duration(Second)'), arrow_cast(0, 'Duration(Millisecond)'), arrow_cast(0, 'Duration(Microsecond)'), arrow_cast(0, 'Duration(Nanosecond)')), + (2, arrow_cast(0, 'Duration(Second)'), arrow_cast(0, 'Duration(Millisecond)'), arrow_cast(0, 'Duration(Microsecond)'), arrow_cast(0, 'Duration(Nanosecond)')); + +# Scalar path +query IIII +SELECT approx_distinct(column2), approx_distinct(column3), approx_distinct(column4), approx_distinct(column5) FROM approx_distinct_duration_test; +---- +4 4 4 4 + +# Grouped path +query IIIII +SELECT column1, approx_distinct(column2), approx_distinct(column3), approx_distinct(column4), approx_distinct(column5) +FROM approx_distinct_duration_test GROUP BY column1 ORDER BY column1; +---- +1 2 2 2 2 +2 2 2 2 2 + +statement ok +DROP TABLE approx_distinct_duration_test; + + +# This test runs approx_distinct over the intervals YearMonth, +# DayTime, MonthDayNano for the scalar and the grouped path. +statement ok +CREATE TABLE approx_distinct_interval_test (g INT, ym INTERVAL, dt INTERVAL, mdn INTERVAL) AS VALUES + (1, INTERVAL '1' MONTH, INTERVAL '1' DAY, INTERVAL '1' MONTH), + (1, INTERVAL '2' MONTH, INTERVAL '1 day 5 hours', INTERVAL '1 day 5 nanoseconds'), + (1, INTERVAL '2' MONTH, INTERVAL '1 day 5 hours', INTERVAL '1 day 5 nanoseconds'), + (2, INTERVAL '3' YEAR, INTERVAL '2' DAY, INTERVAL '2' DAY), + (2, INTERVAL '0' MONTH, INTERVAL '0' DAY, INTERVAL '0' DAY), + (2, INTERVAL '0' MONTH, INTERVAL '0' DAY, INTERVAL '0' DAY); + +# Scalar path +query III +SELECT approx_distinct(ym), approx_distinct(dt), approx_distinct(mdn) FROM approx_distinct_interval_test; +---- +4 4 4 + +# Grouped path +query IIII +SELECT g, approx_distinct(ym), approx_distinct(dt), approx_distinct(mdn) +FROM approx_distinct_interval_test GROUP BY g ORDER BY g; +---- +1 2 2 2 +2 2 2 2 + +statement ok +DROP TABLE approx_distinct_interval_test; + + ## This test executes the APPROX_PERCENTILE_CONT aggregation against the test ## data, asserting the estimated quantiles are ±5% their actual values. ## @@ -2383,7 +2656,6 @@ d 2.444444444444 25.444444444444 e 3 40.333333333333 - query TR SELECT c1, approx_percentile_cont(0.95) WITHIN GROUP (ORDER BY c3 DESC) AS c3_p95 FROM aggregate_test_100 GROUP BY 1 ORDER BY 1 ---- @@ -2651,23 +2923,6 @@ SELECT count(1 + 1) ---- 1 -# csv_query_array_agg -query ? -SELECT array_agg(c13) FROM (SELECT * FROM aggregate_test_100 ORDER BY c13 LIMIT 2) test ----- -[0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB] - -# csv_query_array_agg_empty -query ? -SELECT array_agg(c13) FROM (SELECT * FROM aggregate_test_100 LIMIT 0) test ----- -NULL - -# csv_query_array_agg_one -query ? -SELECT array_agg(c13) FROM (SELECT * FROM aggregate_test_100 ORDER BY c13 LIMIT 1) test ----- -[0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm] # csv_query_array_agg_with_overflow query IIRIII @@ -2720,12 +2975,6 @@ NULL 4 29 1.260869565217 123 -117 23 NULL 5 -194 -13.857142857143 118 -101 14 NULL NULL 781 7.81 125 -117 100 -# select with count to forces array_agg_distinct function, since single distinct expression is converted to group by by optimizer -# csv_query_array_agg_distinct -query ?I -SELECT array_sort(array_agg(distinct c2)), count(1) FROM aggregate_test_100 ----- -[1, 2, 3, 4, 5] 100 # aggregate_time_min_and_max query TT @@ -2888,7 +3137,6 @@ SELECT max(c1) FROM test; 3 - # count_basic statement ok create table t (c int) as values (1), (2), (null), (3), (null), (4), (5); @@ -4250,177 +4498,6 @@ SELECT percentile_cont(0.75) WITHIN GROUP (ORDER BY v DESC) FROM (VALUES (1), (2 ---- 2.75 -# array_agg_zero -query ? -SELECT ARRAY_AGG([]) ----- -[[]] - -# array_agg_one -query ? -SELECT ARRAY_AGG([1]) ----- -[[1]] - -# test array_agg with no row qualified -statement ok -create table t(a int, b float, c bigint) as values (1, 1.2, 2); - -# returns NULL, follows DuckDB's behaviour -query ? -select array_agg(a) from t where a > 2; ----- -NULL - -query ? -select array_agg(b) from t where b > 3.1; ----- -NULL - -query ? -select array_agg(c) from t where c > 3; ----- -NULL - -query ?I -select array_agg(c), count(1) from t where c > 3; ----- -NULL 0 - -# returns 0 rows if group by is applied, follows DuckDB's behaviour -query ? -select array_agg(a) from t where a > 3 group by a; ----- - -query ?I -select array_agg(a), count(1) from t where a > 3 group by a; ----- - -# returns NULL, follows DuckDB's behaviour -query ? -select array_agg(distinct a) from t where a > 3; ----- -NULL - -query ?I -select array_agg(distinct a), count(1) from t where a > 3; ----- -NULL 0 - -# returns 0 rows if group by is applied, follows DuckDB's behaviour -query ? -select array_agg(distinct a) from t where a > 3 group by a; ----- - -query ?I -select array_agg(distinct a), count(1) from t where a > 3 group by a; ----- - -# test order sensitive array agg -query ? -select array_agg(a order by a) from t where a > 3; ----- -NULL - -query ? -select array_agg(a order by a) from t where a > 3 group by a; ----- - -query ?I -select array_agg(a order by a), count(1) from t where a > 3 group by a; ----- - -statement ok -drop table t; - -# test with no values -statement ok -create table t(a int, b float, c bigint); - -query ? -select array_agg(a) from t; ----- -NULL - -query ? -select array_agg(b) from t; ----- -NULL - -query ? -select array_agg(c) from t; ----- -NULL - -query ?I -select array_agg(distinct a), count(1) from t; ----- -NULL 0 - -query ?I -select array_agg(distinct b), count(1) from t; ----- -NULL 0 - -query ?I -select array_agg(distinct b), count(1) from t; ----- -NULL 0 - -statement ok -drop table t; - - -# array_agg_i32 -statement ok -create table t (c1 int) as values (1), (2), (3), (4), (5); - -query ? -select array_agg(c1) from t; ----- -[1, 2, 3, 4, 5] - -statement ok -drop table t; - -# array_agg_nested -statement ok -create table t as values (make_array([1, 2, 3], [4, 5])), (make_array([6], [7, 8])), (make_array([9])); - -query ? -select array_agg(column1) from t; ----- -[[[1, 2, 3], [4, 5]], [[6], [7, 8]], [[9]]] - -statement ok -drop table t; - -# array_agg_ignore_nulls -statement ok -create table t as values (NULL, ''), (1, 'c'), (2, 'a'), (NULL, 'b'), (4, NULL), (NULL, NULL), (5, 'a'); - -query ? -select array_agg(column1) ignore nulls as c1 from t; ----- -[1, 2, 4, 5] - -query II -select count(*), array_length(array_agg(distinct column2) ignore nulls) from t; ----- -7 4 - -query ? -select array_agg(column2 order by column1) ignore nulls from t; ----- -[c, a, a, , b] - -query ? -select array_agg(DISTINCT column2 order by column2) ignore nulls from t; ----- -[, a, b, c] - -statement ok -drop table t; # variance_single_value query RRRR @@ -4435,7 +4512,6 @@ select var(sq.column1), var_pop(sq.column1), stddev(sq.column1), stddev_pop(sq.c 2 1 1.414213562373 1 - # aggregates on empty tables statement ok CREATE TABLE empty (column1 bigint, column2 int); @@ -5273,7 +5349,6 @@ DROP TABLE min_bool; ################# - ################# # min_max on strings/binary with null values and groups ################# @@ -5478,6 +5553,31 @@ SELECT id, MAX(value) FROM fixed_size_binary_views GROUP BY id ORDER BY id; 3 000101 4 NULL +# Group by a FixedSizeBinary column +# (exercises the FixedSizeBinary `GroupColumn` in `GroupValuesColumn`) +query ?I +SELECT value, COUNT(*) FROM fixed_size_binary_views GROUP BY value ORDER BY value; +---- +000101 2 +000102 1 +000103 3 +000104 2 +000109 1 +NULL 2 + +# Multi-column group by including a FixedSizeBinary column +query ?II +SELECT value, id, COUNT(*) FROM fixed_size_binary_views GROUP BY value, id ORDER BY value, id; +---- +000101 2 1 +000101 3 1 +000102 1 1 +000103 1 3 +000104 1 2 +000109 2 1 +NULL 1 1 +NULL 4 1 + statement ok DROP VIEW fixed_size_binary_views; @@ -5983,7 +6083,6 @@ ORDER BY tag 426172 426172 1 426172 426172 1 - statement ok drop table t_source; @@ -6197,6 +6296,103 @@ GROUP BY g ---- 0 0 +# first_value_with_group_by_and_nullable_filter +# Rows whose FILTER predicate evaluates to NULL must be excluded (#22666) +query II rowsort +SELECT g, first_value(a ORDER BY a) FILTER (WHERE b < 1) AS fv +FROM (VALUES (0, 10, CAST(NULL AS INT)), (0, 20, 2)) AS t(g, a, b) +GROUP BY g +---- +0 NULL + +# last_value_with_group_by_and_nullable_filter +query II rowsort +SELECT g, last_value(a ORDER BY a) FILTER (WHERE b < 1) AS lv +FROM (VALUES (0, 10, CAST(NULL AS INT)), (0, 20, 2)) AS t(g, a, b) +GROUP BY g +---- +0 NULL + +# first_last_value_with_group_by_and_mixed_filter_results +# Only rows whose FILTER predicate is TRUE participate: a = 10 (b = 1) and +# a = 20 (b = 0) in group 0. The NULL-predicate row (a = 5) and the +# FALSE-predicate row (a = 30) are excluded. No row passes the filter in +# group 1, so the aggregates return NULL there. +query III rowsort +SELECT g, + first_value(a ORDER BY a) FILTER (WHERE b < 2) AS fv, + last_value(a ORDER BY a) FILTER (WHERE b < 2) AS lv +FROM (VALUES (0, 5, CAST(NULL AS INT)), (0, 10, 1), (0, 30, 2), (0, 20, 0), + (1, 100, CAST(NULL AS INT)), (1, 50, 3)) AS t(g, a, b) +GROUP BY g +---- +0 10 20 +1 NULL NULL + +# first_last_value_with_group_by_filter_all_true_and_no_filter +# Behavior is unchanged when every row passes the FILTER or there is no FILTER +query IIIII rowsort +SELECT g, + first_value(a ORDER BY a) FILTER (WHERE a > 0) AS fv, + last_value(a ORDER BY a) FILTER (WHERE a > 0) AS lv, + first_value(a ORDER BY a) AS fv_no_filter, + last_value(a ORDER BY a) AS lv_no_filter +FROM (VALUES (0, 5, CAST(NULL AS INT)), (0, 10, 1), (0, 30, 2), (0, 20, 0)) AS t(g, a, b) +GROUP BY g +---- +0 5 30 5 30 + +# first_value_without_group_by_and_nullable_filter +query I rowsort +SELECT first_value(a ORDER BY a) FILTER (WHERE b < 1) AS fv +FROM (VALUES (10, CAST(NULL AS INT)), (20, 2)) AS t(a, b) +---- +NULL + +# first_value_window_function_no_regression +query II +SELECT a, first_value(a) OVER (ORDER BY a) AS fv +FROM (VALUES (10), (20), (5)) AS t(a) +ORDER BY a +---- +5 5 +10 5 +20 5 + +# query_with_untyped_null_filter +query I +SELECT count(*) FILTER (WHERE NULL) +---- +0 + +query I +SELECT count(1) FILTER (WHERE NULL) +---- +0 + +query I +SELECT sum(1) FILTER (WHERE NULL) +---- +NULL + +query R +SELECT avg(1) FILTER (WHERE NULL) +---- +NULL + +# window_aggregate_with_untyped_null_filter +query I +SELECT count(*) FILTER (WHERE NULL) OVER () +FROM (VALUES (1)) AS t(x) +---- +0 + +query I +SELECT sum(1) FILTER (WHERE NULL) OVER () +FROM (VALUES (1)) AS t(x) +---- +NULL + # query_with_and_without_filter query III rowsort SELECT @@ -6375,6 +6571,99 @@ c NULL 2 statement ok drop table dn; +# sum_interval +# Component-wise sum across all three Interval variants (matches PostgreSQL). + +# Basic Interval(MonthDayNano): the issue's repro. +# (0 mons, 0 days, 1s) + (12 mons, 0, 0) + (1 mon, 0, 0) = (13 mons, 0, 1s) +query T? +SELECT arrow_typeof(sum(v)), sum(v) FROM (VALUES + (interval '1 second'), + (interval '1 year'), + (interval '1 month')) t(v); +---- +Interval(MonthDayNano) 13 mons 1.000000000 secs + +# NULLs are skipped. +query ? +SELECT sum(v) FROM (VALUES + (interval '1 day'), + (NULL), + (interval '2 days')) t(v); +---- +3 days + +# Empty input → NULL. +query ? +SELECT sum(v) FROM (VALUES (interval '1 day')) t(v) WHERE 1 = 0; +---- +NULL + +# GROUP BY exercises the PrimitiveGroupsAccumulator path. +query I? rowsort +SELECT k, sum(v) FROM (VALUES + (1, interval '1 day'), + (1, interval '2 days'), + (2, interval '1 month')) t(k, v) +GROUP BY k; +---- +1 3 days +2 1 mons + +# Interval(YearMonth) via cast. +query T? +SELECT arrow_typeof(sum(v)), sum(v) FROM (VALUES + (arrow_cast('1 year', 'Interval(YearMonth)')), + (arrow_cast('6 months', 'Interval(YearMonth)'))) t(v); +---- +Interval(YearMonth) 1 years 6 mons + +# Interval(DayTime) via cast. +query T? +SELECT arrow_typeof(sum(v)), sum(v) FROM (VALUES + (arrow_cast('1 day', 'Interval(DayTime)')), + (arrow_cast('1 day', 'Interval(DayTime)'))) t(v); +---- +Interval(DayTime) 2 days + +# Sliding window sum on intervals. +query ?? +SELECT v, sum(v) OVER (ORDER BY v ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM (VALUES + (interval '1 day'), + (interval '2 days'), + (interval '3 days')) t(v); +---- +1 days 1 days +2 days 3 days +3 days 5 days + +# DISTINCT sum drops duplicates (DistinctSumAccumulator path). +query ? +SELECT sum(DISTINCT v) FROM (VALUES + (interval '1 day'), + (interval '1 day'), + (interval '2 days')) t(v); +---- +3 days + +# SUM(col + interval_lit) — exercises the simplify_expr_op_literal path. +query ? +SELECT sum(v + interval '1 day') FROM (VALUES + (interval '1 day'), + (interval '2 days'), + (interval '3 days')) t(v); +---- +9 days + +# Negative intervals: component-wise wrapping_add over signed i32/i64. +query ? +SELECT sum(v) FROM (VALUES + (interval '1 day'), + (interval '-3 days')) t(v); +---- +-2 days + # Prepare the table with dictionary values for testing statement ok CREATE TABLE value(x bigint) AS VALUES (1), (2), (3), (1), (3), (4), (5), (2); @@ -6759,7 +7048,6 @@ statement error select regr_sxy(NULL, 'bar'); - # regr_*() NULL results query RRIRRRRRR select regr_slope(1,1), regr_intercept(1,1), regr_count(1,1), regr_r2(1,1), regr_avgx(1,1), regr_avgy(1,1), regr_sxx(1,1), regr_syy(1,1), regr_sxy(1,1); @@ -6787,7 +7075,6 @@ select regr_slope(column2, column1), regr_intercept(column2, column1), regr_coun NULL NULL 3 NULL 1 4 0 8 0 - # regr_*() basic tests query RRIRRRRRR select @@ -6892,7 +7179,6 @@ b 3 0 2 1 2 6 2 18 6 c NULL NULL 1 NULL 1 10 0 0 0 - # regr_*() testing merge_batch() from RegrAccumulator's internal implementation statement ok set datafusion.execution.batch_size = 1; @@ -6952,7 +7238,6 @@ statement ok set datafusion.execution.batch_size = 8192; - # regr_*() testing retract_batch() from RegrAccumulator's internal implementation query RRIRRRRRR SELECT @@ -7422,13 +7707,11 @@ statement ok drop table distinct_count_large_binary_table; - ## Cleanup from distinct count tests statement ok drop table distinct_count_string_table; - # rule `aggregate_statistics` should not optimize MIN/MAX to wrong values on empty relation statement ok @@ -7862,6 +8145,23 @@ CREATE TABLE t1(v1 int); statement error DataFusion error: Error during planning: Aggregate functions are not allowed in the WHERE clause. Consider using HAVING instead SELECT v1 FROM t1 WHERE ((count(v1) % 1) << 1) > 0; +# issue: https://github.com/apache/datafusion/issues/11748 +query R +SELECT AVG(v1) FROM t1 GROUP BY false HAVING false; +---- + +query R +SELECT AVG(v1) FROM t1 GROUP BY false; +---- + +statement ok +INSERT INTO t1 VALUES (1), (2), (3); + +query R +SELECT AVG(v1) FROM t1 GROUP BY false; +---- +2 + statement ok DROP TABLE t1; @@ -8451,19 +8751,6 @@ VALUES ---- x 1 -query error Error during planning: WITHIN GROUP is only supported for ordered-set aggregate functions -SELECT array_agg(a_varchar) WITHIN GROUP (ORDER BY a_varchar) -FROM (VALUES ('a'), ('d'), ('c'), ('a')) t(a_varchar); - - -query error Error during planning: WITHIN GROUP is only supported for ordered-set aggregate functions -SELECT array_agg(DISTINCT a_varchar) WITHIN GROUP (ORDER BY a_varchar) -FROM (VALUES ('a'), ('d'), ('c'), ('a')) t(a_varchar); - - -query error Error during planning: ORDER BY and WITHIN GROUP clauses cannot be used together in the same aggregate function -SELECT array_agg(a_varchar order by a_varchar) WITHIN GROUP (ORDER BY a_varchar) -FROM (VALUES ('a'), ('d'), ('c'), ('a')) t(a_varchar); # distinct average statement ok @@ -8852,7 +9139,7 @@ ORDER BY g; # Config reset -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; @@ -9018,3 +9305,73 @@ SET datafusion.execution.target_partitions = 4; statement ok DROP TABLE hits_raw; + +# Nested aggregate function calls are rejected during planning +# issue: https://github.com/apache/datafusion/issues/23812 +statement ok +CREATE TABLE nested_agg_t AS VALUES (1, 10.0), (1, 20.0), (2, 30.0); + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested: 'sum\(nested_agg_t\.column2\)' is nested inside 'sum\(sum\(nested_agg_t\.column2\)\)' +SELECT column1, sum(sum(column2)) FROM nested_agg_t GROUP BY column1; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested: 'count\(nested_agg_t\.column2\)' is nested inside 'sum\(count\(nested_agg_t\.column2\)\)' +SELECT column1, sum(count(column2)) FROM nested_agg_t GROUP BY column1; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT sum(sum(column2)) FROM nested_agg_t; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT column1 FROM nested_agg_t GROUP BY column1 HAVING sum(sum(column2)) > 0; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT column1, sum(column2 + sum(column2)) FROM nested_agg_t GROUP BY column1; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT sum(column2) FILTER (WHERE sum(column2) > 0) FROM nested_agg_t; + +statement error DataFusion error: Error during planning: Aggregate function calls cannot be nested +SELECT array_agg(column2 ORDER BY sum(column2)) FROM nested_agg_t; + +# A window function nested inside an aggregate is rejected as well +statement error DataFusion error: Error during planning: Aggregate function calls cannot contain window function calls +SELECT sum(sum(column2) OVER ()) FROM nested_agg_t; + +# ... as are nested window functions +statement error DataFusion error: Error during planning: Window function calls cannot be nested +SELECT sum(sum(column2) OVER ()) OVER () FROM nested_agg_t; + +statement error DataFusion error: Error during planning: Window function calls cannot be nested +SELECT row_number() OVER (ORDER BY row_number() OVER ()) FROM nested_agg_t; + +# ... including a window call nested in `PARTITION BY` +statement error DataFusion error: Error during planning: Window function calls cannot be nested +SELECT row_number() OVER (PARTITION BY row_number() OVER ()) FROM nested_agg_t; + +# A scalar function applied to an aggregate is legal +query IR +SELECT column1, abs(sum(column2)) FROM nested_agg_t GROUP BY column1 ORDER BY column1; +---- +1 30 +2 30 + +# A window function applied to an aggregate is legal +query IR +SELECT column1, sum(sum(column2)) OVER () FROM nested_agg_t GROUP BY column1 ORDER BY column1; +---- +1 60 +2 60 + +# An aggregate over the result of an aggregate computed in a subquery is legal +query R +SELECT sum(s) FROM (SELECT sum(column2) AS s FROM nested_agg_t GROUP BY column1); +---- +60 + +# An aggregate over the result of a window function computed in a subquery is legal +query R +SELECT sum(s) FROM (SELECT sum(column2) OVER () AS s FROM nested_agg_t); +---- +180 + +statement ok +DROP TABLE nested_agg_t; diff --git a/datafusion/sqllogictest/test_files/aggregate_any_value.slt b/datafusion/sqllogictest/test_files/aggregate_any_value.slt new file mode 100644 index 0000000000000..3fe6f787d346d --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_any_value.slt @@ -0,0 +1,57 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +statement ok +CREATE TABLE any_value_test AS VALUES + (1, NULL, NULL), + (1, 10, 'first'), + (1, 20, 'second'), + (2, NULL, NULL), + (2, NULL, NULL), + (3, 30, 'third'); + +query B +SELECT any_value(column2) IN (10, 20) FROM any_value_test; +---- +true + +query IBB rowsort +SELECT + column1, + any_value(column2) IN (10, 20, 30), + any_value(column3) IN ('first', 'second', 'third') +FROM any_value_test +GROUP BY column1; +---- +1 true true +2 NULL NULL +3 true true + +query T +SELECT arrow_typeof(any_value(column3)) FROM any_value_test; +---- +Utf8 + +query I +SELECT any_value(column2) FROM any_value_test WHERE false; +---- +NULL + +query I +SELECT any_value(column2) FROM any_value_test WHERE column1 = 2; +---- +NULL diff --git a/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt new file mode 100644 index 0000000000000..3dbf880fd1fa9 --- /dev/null +++ b/datafusion/sqllogictest/test_files/aggregate_memory_spill.slt @@ -0,0 +1,235 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Memory-limited (spilling) grouped hash aggregation. +# +# High-cardinality GROUP BY under a tight memory limit: the aggregate spills to +# disk, re-groups the spilled state, and must still return the right answer. +# +# The group key is scrambled with `(v * 7) % 100000` because generate_series is +# sorted, which would take the streaming path that never spills. gcd(7, 100000) +# = 1, so it's a bijection over 1..100000. Still 100000 groups, just unsorted, +# so the hash table grows and spills. +# +# Each query aggregates over the grouped result, so the expected output is one +# row. sum(1..100000) = 5000050000, and every v lands in one group, so the +# per-group sums always add back to that total. + +# Single partition keeps the aggregation in one operator (no repartition). +statement ok +SET datafusion.execution.target_partitions = 1 + +statement ok +SET datafusion.execution.batch_size = 128 + +statement ok +SET datafusion.runtime.memory_limit = '1M' + +# --- Case A: single-column high-cardinality GROUP BY --- +query II +SELECT count(*), sum(total) +FROM ( + SELECT (v * 7) % 100000 AS k, sum(v) AS total + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- +100000 5000050000 + +# Assert spill happened, the `spill_count` metric must be > 0 +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(total) +FROM ( + SELECT (v * 7) % 100000 AS k, sum(v) AS total + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- + +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[spill_count=7,] + + +# --- Case B: multi-column GROUP BY (is_single() = false) --- +# Both keys are bijections of v, so each (a, b) pair is unique: 100000 groups. +query II +SELECT count(*), sum(total) +FROM ( + SELECT (v * 7) % 100000 AS a, (v * 13) % 100000 AS b, sum(v) AS total + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000, (v * 13) % 100000 +) +---- +100000 5000050000 + +# Assert spill happened, the `spill_count` metric must be > 0 +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(total) +FROM ( + SELECT (v * 7) % 100000 AS a, (v * 13) % 100000 AS b, sum(v) AS total + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000, (v * 13) % 100000 +) +---- + +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000), v@0 * 13 % 100000 as t.v * Int64(13) % Int64(100000)], aggr=[sum(t.v)], metrics=[spill_count=7,] + + +# --- Case C: DISTINCT aggregate under memory limit --- +# One distinct value per group, so each count(DISTINCT v) = 1. +query II +SELECT count(*), sum(d) +FROM ( + SELECT (v * 7) % 100000 AS k, count(DISTINCT v) AS d + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- +100000 100000 + +# Assert spill happened, the `spill_count` metric must be > 0 +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(d) +FROM ( + SELECT (v * 7) % 100000 AS k, count(DISTINCT v) AS d + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- + +04)------AggregateExec: mode=Single, gby=[group_alias_0@0 as group_alias_0], aggr=[count(alias1)], metrics=[spill_count=7,] + + +# --- Case D: multiple aggregates (sum/min/max) under memory limit --- +# Each group holds a single v, so min(v) = max(v) = v within the group. +query IIII +SELECT count(*), sum(s), min(mn), max(mx) +FROM ( + SELECT (v * 7) % 100000 AS k, sum(v) AS s, min(v) AS mn, max(v) AS mx + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- +100000 5000050000 1 100000 + +# Assert spill happened, the `spill_count` metric must be > 0 +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(s), min(mn), max(mx) +FROM ( + SELECT (v * 7) % 100000 AS k, sum(v) AS s, min(v) AS mn, max(v) AS mx + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- + +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v), min(t.v), max(t.v)], metrics=[spill_count=7,] + + +# --- Case E: avg() aggregate (Float64 output) under memory limit --- +# Each group holds a single v, so avg(v) = v within the group. +query IRR +SELECT count(*), min(a), max(a) +FROM ( + SELECT (v * 7) % 100000 AS k, avg(v) AS a + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- +100000 1 100000 + +# Assert spill happened, the `spill_count` metric must be > 0 +query TT +EXPLAIN ANALYZE +SELECT count(*), min(a), max(a) +FROM ( + SELECT (v * 7) % 100000 AS k, avg(v) AS a + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- + +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[avg(t.v)], metrics=[spill_count=7,] + + +# --- Case F: array_agg() aggregate (growable state) under memory limit --- +# Each group holds a single v, so array_length(array_agg(v)) = 1. +query II +SELECT count(*), sum(l) +FROM ( + SELECT (v * 7) % 100000 AS k, array_length(array_agg(v)) AS l + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- +100000 100000 + +# Assert spill happened, the `spill_count` metric must be > 0 +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(l) +FROM ( + SELECT (v * 7) % 100000 AS k, array_length(array_agg(v)) AS l + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- + +04)------AggregateExec: mode=Single, gby=[v@0 * 7 % 100000 as t.v * Int64(7) % Int64(100000)], aggr=[array_agg(t.v)], metrics=[spill_count=7,] + + +# --- Case G: partial/final aggregation under memory limit --- +statement ok +SET datafusion.execution.target_partitions = 4 + +query II +SELECT count(*), sum(total) +FROM ( + SELECT (v * 7) % 100000 AS k, sum(v) AS total + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- +100000 5000050000 + +# Assert spill happened in the final aggregation. +# In multi-partitions configuration, 'spilled_rows' is not deterministic, so assert +# the unit to be 'K' +query TT +EXPLAIN ANALYZE +SELECT count(*), sum(total) +FROM ( + SELECT (v * 7) % 100000 AS k, sum(v) AS total + FROM generate_series(1, 100000) AS t(v) + GROUP BY (v * 7) % 100000 +) +---- + +06)----------AggregateExec: mode=FinalPartitioned, gby=[t.v * Int64(7) % Int64(100000)@0 as t.v * Int64(7) % Int64(100000)], aggr=[sum(t.v)], metrics=[spilled_rows=K,] + + +# Restore settings to slt runner defaults +statement ok +RESET datafusion.runtime.memory_limit + +statement ok +RESET datafusion.execution.batch_size + +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema diff --git a/datafusion/sqllogictest/test_files/aggregate_repartition.slt b/datafusion/sqllogictest/test_files/aggregate_repartition.slt index 1f1e726811675..2302e161bfe72 100644 --- a/datafusion/sqllogictest/test_files/aggregate_repartition.slt +++ b/datafusion/sqllogictest/test_files/aggregate_repartition.slt @@ -131,7 +131,7 @@ physical_plan # Config reset -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok SET datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/aggregate_skip_partial.slt b/datafusion/sqllogictest/test_files/aggregate_skip_partial.slt index a10417f232409..2ed4c9921f3a7 100644 --- a/datafusion/sqllogictest/test_files/aggregate_skip_partial.slt +++ b/datafusion/sqllogictest/test_files/aggregate_skip_partial.slt @@ -220,7 +220,7 @@ e true false NULL statement ok reset datafusion.execution.batch_size; -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; @@ -636,6 +636,44 @@ FROM aggregate_test_100_null GROUP BY c2 ORDER BY c2; 4 3 5 6 +# Test variance and stddev with nullable fields and filters +query IRRR +SELECT c2, + var_samp(c11) FILTER (WHERE c3 > 0), + stddev_samp(c11) FILTER (WHERE c3 > 0), + stddev_pop(c11) FILTER (WHERE c3 > 0) +FROM aggregate_test_100_null GROUP BY c2 ORDER BY c2; +---- +1 0.085786074994 0.29289259976 0.276141791266 +2 0.070769227104 0.266024861816 0.24884347232 +3 0.087515365779 0.295829960922 0.270054571305 +4 0.038697229167 0.196716113134 0.182123731489 +5 0.081817232141 0.286037116719 0.23354832782 + +# Test corr with nullable fields and filters +query IR +SELECT c2, + corr(c3, c11) FILTER (WHERE c5 > 0) +FROM aggregate_test_100_null GROUP BY c2 ORDER BY c2; +---- +1 -0.38515658251 +2 0.414489249329 +3 -0.796447131429 +4 -0.938568248748 +5 -0.051058146743 + +# Test count distinct with nullable fields and filters +query II +SELECT c2, + count(distinct c3) FILTER (WHERE c11 > 0.5) +FROM aggregate_test_100_null GROUP BY c2 ORDER BY c2; +---- +1 10 +2 6 +3 3 +4 3 +5 6 + # Test median with nullable fields and filter query IRR SELECT c2, @@ -712,7 +750,7 @@ ORDER BY i; statement ok reset datafusion.execution.batch_size; -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; @@ -772,7 +810,7 @@ true false false false false true false NULL statement ok reset datafusion.execution.batch_size; -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/aggregates_topk.slt b/datafusion/sqllogictest/test_files/aggregates_topk.slt index 19ead8965ed01..e2d453068adf7 100644 --- a/datafusion/sqllogictest/test_files/aggregates_topk.slt +++ b/datafusion/sqllogictest/test_files/aggregates_topk.slt @@ -98,15 +98,15 @@ c 4 a 1 query TT -explain select trace_id, MAX(timestamp) from traces group by trace_id order by MAX(timestamp) desc limit 4; +explain select trace_id, MAX(timestamp) from traces group by trace_id order by MAX(timestamp) desc nulls last limit 4; ---- logical_plan -01)Sort: max(traces.timestamp) DESC NULLS FIRST, fetch=4 +01)Sort: max(traces.timestamp) DESC NULLS LAST, fetch=4 02)--Aggregate: groupBy=[[traces.trace_id]], aggr=[[max(traces.timestamp)]] 03)----TableScan: traces projection=[trace_id, timestamp] physical_plan -01)SortPreservingMergeExec: [max(traces.timestamp)@1 DESC], fetch=4 -02)--SortExec: TopK(fetch=4), expr=[max(traces.timestamp)@1 DESC], preserve_partitioning=[true] +01)SortPreservingMergeExec: [max(traces.timestamp)@1 DESC NULLS LAST], fetch=4 +02)--SortExec: TopK(fetch=4), expr=[max(traces.timestamp)@1 DESC NULLS LAST], preserve_partitioning=[true] 03)----AggregateExec: mode=FinalPartitioned, gby=[trace_id@0 as trace_id], aggr=[max(traces.timestamp)], lim=[4] 04)------RepartitionExec: partitioning=Hash([trace_id@0], 4), input_partitions=1 05)--------AggregateExec: mode=Partial, gby=[trace_id@0 as trace_id], aggr=[max(traces.timestamp)], lim=[4] @@ -218,17 +218,17 @@ x zebra z mango query TT -explain select category, max(val) max_val from string_topk group by category order by max_val desc limit 2; +explain select category, max(val) max_val from string_topk group by category order by max_val desc nulls last limit 2; ---- logical_plan -01)Sort: max_val DESC NULLS FIRST, fetch=2 +01)Sort: max_val DESC NULLS LAST, fetch=2 02)--Projection: string_topk.category, max(string_topk.val) AS max_val 03)----Aggregate: groupBy=[[string_topk.category]], aggr=[[max(string_topk.val)]] 04)------TableScan: string_topk projection=[category, val] physical_plan -01)SortPreservingMergeExec: [max_val@1 DESC], fetch=2 -02)--SortExec: TopK(fetch=2), expr=[max_val@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[category@0 as category, max(string_topk.val)@1 as max_val] +01)SortPreservingMergeExec: [max_val@1 DESC NULLS LAST], fetch=2 +02)--ProjectionExec: expr=[category@0 as category, max(string_topk.val)@1 as max_val] +03)----SortExec: TopK(fetch=2), expr=[max(string_topk.val)@1 DESC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[category@0 as category], aggr=[max(string_topk.val)], lim=[2] 05)--------RepartitionExec: partitioning=Hash([category@0], 4), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[category@0 as category], aggr=[max(string_topk.val)], lim=[2] @@ -241,19 +241,19 @@ x zebra z mango query TT -explain select category, max(val) max_val from string_topk_view group by category order by max_val desc limit 2; +explain select category, max(val) max_val from string_topk_view group by category order by max_val desc nulls last limit 2; ---- logical_plan -01)Sort: max_val DESC NULLS FIRST, fetch=2 +01)Sort: max_val DESC NULLS LAST, fetch=2 02)--Projection: string_topk_view.category, max(string_topk_view.val) AS max_val 03)----Aggregate: groupBy=[[string_topk_view.category]], aggr=[[max(string_topk_view.val)]] 04)------SubqueryAlias: string_topk_view 05)--------Projection: string_topk.category AS category, string_topk.val AS val 06)----------TableScan: string_topk projection=[category, val] physical_plan -01)SortPreservingMergeExec: [max_val@1 DESC], fetch=2 -02)--SortExec: TopK(fetch=2), expr=[max_val@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[category@0 as category, max(string_topk_view.val)@1 as max_val] +01)SortPreservingMergeExec: [max_val@1 DESC NULLS LAST], fetch=2 +02)--ProjectionExec: expr=[category@0 as category, max(string_topk_view.val)@1 as max_val] +03)----SortExec: TopK(fetch=2), expr=[max(string_topk_view.val)@1 DESC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[category@0 as category], aggr=[max(string_topk_view.val)], lim=[2] 05)--------RepartitionExec: partitioning=Hash([category@0], 4), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[category@0 as category], aggr=[max(string_topk_view.val)], lim=[2] @@ -268,11 +268,13 @@ NULL 0 0 c 1 2 # Regression tests for string max with ORDER BY ... LIMIT to ensure schema stability +# Note: the NULL group has an all-NULL trace_id, so its max is NULL and ranks +# first under DESC NULLS FIRST (previously the group was dropped: issue #23440) query TT select trace_id, max(trace_id) as max_trace from traces group by trace_id order by max_trace desc limit 2; ---- +NULL NULL c c -b b query TT explain select trace_id, max(trace_id) as max_trace from traces group by trace_id order by max_trace desc limit 2; @@ -284,11 +286,11 @@ logical_plan 04)------TableScan: traces projection=[trace_id] physical_plan 01)SortPreservingMergeExec: [max_trace@1 DESC], fetch=2 -02)--SortExec: TopK(fetch=2), expr=[max_trace@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[trace_id@0 as trace_id, max(traces.trace_id)@1 as max_trace] -04)------AggregateExec: mode=FinalPartitioned, gby=[trace_id@0 as trace_id], aggr=[max(traces.trace_id)], lim=[2] +02)--ProjectionExec: expr=[trace_id@0 as trace_id, max(traces.trace_id)@1 as max_trace] +03)----SortExec: TopK(fetch=2), expr=[max(traces.trace_id)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[trace_id@0 as trace_id], aggr=[max(traces.trace_id)] 05)--------RepartitionExec: partitioning=Hash([trace_id@0], 4), input_partitions=1 -06)----------AggregateExec: mode=Partial, gby=[trace_id@0 as trace_id], aggr=[max(traces.trace_id)], lim=[2] +06)----------AggregateExec: mode=Partial, gby=[trace_id@0 as trace_id], aggr=[max(traces.trace_id)] 07)------------DataSourceExec: partitions=1, partition_sizes=[1] @@ -303,15 +305,15 @@ AS SELECT FROM traces; query TT -explain select trace_id, MAX(timestamp) from traces_utf8view group by trace_id order by MAX(timestamp) desc limit 4; +explain select trace_id, MAX(timestamp) from traces_utf8view group by trace_id order by MAX(timestamp) desc nulls last limit 4; ---- logical_plan -01)Sort: max(traces_utf8view.timestamp) DESC NULLS FIRST, fetch=4 +01)Sort: max(traces_utf8view.timestamp) DESC NULLS LAST, fetch=4 02)--Aggregate: groupBy=[[traces_utf8view.trace_id]], aggr=[[max(traces_utf8view.timestamp)]] 03)----TableScan: traces_utf8view projection=[trace_id, timestamp] physical_plan -01)SortPreservingMergeExec: [max(traces_utf8view.timestamp)@1 DESC], fetch=4 -02)--SortExec: TopK(fetch=4), expr=[max(traces_utf8view.timestamp)@1 DESC], preserve_partitioning=[true] +01)SortPreservingMergeExec: [max(traces_utf8view.timestamp)@1 DESC NULLS LAST], fetch=4 +02)--SortExec: TopK(fetch=4), expr=[max(traces_utf8view.timestamp)@1 DESC NULLS LAST], preserve_partitioning=[true] 03)----AggregateExec: mode=FinalPartitioned, gby=[trace_id@0 as trace_id], aggr=[max(traces_utf8view.timestamp)], lim=[4] 04)------RepartitionExec: partitioning=Hash([trace_id@0], 4), input_partitions=1 05)--------AggregateExec: mode=Partial, gby=[trace_id@0 as trace_id], aggr=[max(traces_utf8view.timestamp)], lim=[4] @@ -329,15 +331,15 @@ AS SELECT FROM traces; query TT -explain select trace_id, MAX(timestamp) from traces_largeutf8 group by trace_id order by MAX(timestamp) desc limit 4; +explain select trace_id, MAX(timestamp) from traces_largeutf8 group by trace_id order by MAX(timestamp) desc nulls last limit 4; ---- logical_plan -01)Sort: max(traces_largeutf8.timestamp) DESC NULLS FIRST, fetch=4 +01)Sort: max(traces_largeutf8.timestamp) DESC NULLS LAST, fetch=4 02)--Aggregate: groupBy=[[traces_largeutf8.trace_id]], aggr=[[max(traces_largeutf8.timestamp)]] 03)----TableScan: traces_largeutf8 projection=[trace_id, timestamp] physical_plan -01)SortPreservingMergeExec: [max(traces_largeutf8.timestamp)@1 DESC], fetch=4 -02)--SortExec: TopK(fetch=4), expr=[max(traces_largeutf8.timestamp)@1 DESC], preserve_partitioning=[true] +01)SortPreservingMergeExec: [max(traces_largeutf8.timestamp)@1 DESC NULLS LAST], fetch=4 +02)--SortExec: TopK(fetch=4), expr=[max(traces_largeutf8.timestamp)@1 DESC NULLS LAST], preserve_partitioning=[true] 03)----AggregateExec: mode=FinalPartitioned, gby=[trace_id@0 as trace_id], aggr=[max(traces_largeutf8.timestamp)], lim=[4] 04)------RepartitionExec: partitioning=Hash([trace_id@0], 4), input_partitions=1 05)--------AggregateExec: mode=Partial, gby=[trace_id@0 as trace_id], aggr=[max(traces_largeutf8.timestamp)], lim=[4] @@ -456,6 +458,127 @@ select count(*) from (select category from values_table group by category order ---- 3 +# Test DISTINCT with NULLs and NULLS FIRST ordering (issue #22554) +statement ok +create table nullable_vals (v varchar) as values (NULL), (''), ('a'), ('b'); + +# Verify this regression test exercises the TopK aggregation path +query TT +explain select distinct v from nullable_vals order by v asc nulls first limit 1; +---- +logical_plan +01)Sort: nullable_vals.v ASC NULLS FIRST, fetch=1 +02)--Aggregate: groupBy=[[nullable_vals.v]], aggr=[[]] +03)----TableScan: nullable_vals projection=[v] +physical_plan +01)SortPreservingMergeExec: [v@0 ASC], fetch=1 +02)--SortExec: TopK(fetch=1), expr=[v@0 ASC], preserve_partitioning=[true] +03)----AggregateExec: mode=FinalPartitioned, gby=[v@0 as v], aggr=[], lim=[1] +04)------RepartitionExec: partitioning=Hash([v@0], 4), input_partitions=1 +05)--------AggregateExec: mode=Partial, gby=[v@0 as v], aggr=[], lim=[1] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] + +# NULLS FIRST: NULL should be the first row returned by LIMIT +query T +select distinct v from nullable_vals order by v asc nulls first limit 1; +---- +NULL + +query T +select distinct v from nullable_vals order by v asc nulls first limit 2; +---- +NULL +(empty) + +query T +select distinct v from nullable_vals order by v asc nulls first limit 3; +---- +NULL +(empty) +a + +# NULLS LAST: non-null values come first +query T +select distinct v from nullable_vals order by v asc nulls last limit 1; +---- +(empty) + +query T +select distinct v from nullable_vals order by v asc nulls last limit 4; +---- +(empty) +a +b +NULL + +# DESC NULLS FIRST: NULL comes first +query T +select distinct v from nullable_vals order by v desc nulls first limit 1; +---- +NULL + +# DESC NULLS LAST: NULL comes last +query T +select distinct v from nullable_vals order by v desc nulls last limit 1; +---- +b + +query T +select distinct v from nullable_vals order by v desc nulls last limit 4; +---- +b +a +(empty) +NULL + +# Test with integer column containing NULLs +statement ok +create table nullable_ints (v int) as values (NULL), (3), (1), (2); + +query I +select distinct v from nullable_ints order by v asc nulls first limit 1; +---- +NULL + +query I +select distinct v from nullable_ints order by v asc nulls first limit 3; +---- +NULL +1 +2 + +query I +select distinct v from nullable_ints order by v desc nulls last limit 2; +---- +3 +2 + +query I +select distinct v from nullable_ints order by v asc nulls last limit 4; +---- +1 +2 +3 +NULL + +# Test with all-NULL column +statement ok +create table all_nulls (v varchar) as values (NULL), (NULL); + +query T +select distinct v from all_nulls order by v asc nulls first limit 1; +---- +NULL + +statement ok +drop table nullable_vals; + +statement ok +drop table nullable_ints; + +statement ok +drop table all_nulls; + statement ok drop table values_table; @@ -464,3 +587,205 @@ drop table ids; statement ok drop table traces; + +####### +# Regression tests for all-NULL groups in TopK aggregation (issues #23440, #22190): +# a group whose aggregate inputs are all NULL must be emitted with a NULL +# aggregate value instead of disappearing from the result +####### +statement ok +CREATE TABLE t0 AS SELECT * FROM (VALUES ('gamma', CAST(NULL AS DOUBLE))) v(s, y); + +# MIN/MAX with NULLS FIRST must use regular aggregation because a group's +# aggregate can transition from NULL to non-NULL and worsen its rank. +query TT +explain select s, max(y) as max_y from t0 group by s order by max_y desc nulls first limit 3; +---- +logical_plan +01)Sort: max_y DESC NULLS FIRST, fetch=3 +02)--Projection: t0.s, max(t0.y) AS max_y +03)----Aggregate: groupBy=[[t0.s]], aggr=[[max(t0.y)]] +04)------TableScan: t0 projection=[s, y] +physical_plan +01)ProjectionExec: expr=[s@0 as s, max(t0.y)@1 as max_y] +02)--SortExec: TopK(fetch=3), expr=[max(t0.y)@1 DESC], preserve_partitioning=[false] +03)----AggregateExec: mode=SinglePartitioned, gby=[s@0 as s], aggr=[max(t0.y)] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +# issue #23440: single all-NULL group, MAX DESC NULLS FIRST LIMIT 3 +query R +SELECT max_y FROM (SELECT s, MAX(y) AS max_y FROM t0 GROUP BY s) ORDER BY max_y DESC NULLS FIRST LIMIT 3; +---- +NULL + +# issue #22190: single all-NULL group, MIN ASC NULLS LAST LIMIT 20 +query TT +EXPLAIN SELECT min_y FROM (SELECT s, MIN(y) AS min_y FROM t0 GROUP BY s) ORDER BY min_y ASC NULLS LAST LIMIT 20; +---- +logical_plan +01)Sort: min_y ASC NULLS LAST, fetch=20 +02)--Projection: min(t0.y) AS min_y +03)----Aggregate: groupBy=[[t0.s]], aggr=[[min(t0.y)]] +04)------TableScan: t0 projection=[s, y] +physical_plan +01)SortExec: TopK(fetch=20), expr=[min_y@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--ProjectionExec: expr=[min(t0.y)@1 as min_y] +03)----AggregateExec: mode=SinglePartitioned, gby=[s@0 as s], aggr=[min(t0.y)], lim=[20] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +query R +SELECT min_y FROM (SELECT s, MIN(y) AS min_y FROM t0 GROUP BY s) ORDER BY min_y ASC NULLS LAST LIMIT 20; +---- +NULL + +# one all-NULL group and one valued group, limit larger than the group count: +# both rows must be present +statement ok +CREATE TABLE topk_two_groups(s varchar, y bigint) AS VALUES +('a', CAST(NULL AS BIGINT)), +('b', 10), +('b', 20); + +query TI +select s, max_y from (select s, max(y) as max_y from topk_two_groups group by s) order by max_y desc nulls first limit 10; +---- +a NULL +b 20 + +# 5 all-NULL groups with LIMIT 2: exactly 2 rows survive +statement ok +CREATE TABLE topk_five_nulls(s varchar, y bigint) AS VALUES +('g1', CAST(NULL AS BIGINT)), +('g2', CAST(NULL AS BIGINT)), +('g3', CAST(NULL AS BIGINT)), +('g4', CAST(NULL AS BIGINT)), +('g5', CAST(NULL AS BIGINT)); + +query I +select max_y from (select s, max(y) as max_y from topk_five_nulls group by s) order by max_y desc nulls first limit 2; +---- +NULL +NULL + +# 2 all-NULL groups and 3 valued groups with LIMIT 4 +statement ok +CREATE TABLE topk_mixed(s varchar, y bigint) AS VALUES +('n1', CAST(NULL AS BIGINT)), +('n2', CAST(NULL AS BIGINT)), +('v1', 10), +('v2', 20), +('v3', 30); + +# DESC NULLS FIRST: both NULL groups rank before all values +query I +select max_y from (select s, max(y) as max_y from topk_mixed group by s) order by max_y desc nulls first limit 4; +---- +NULL +NULL +30 +20 + +# DESC NULLS LAST: NULL groups rank after all values +query I +select max_y from (select s, max(y) as max_y from topk_mixed group by s) order by max_y desc nulls last limit 4; +---- +30 +20 +10 +NULL + +# an all-NULL group that later produces a value losing to the current top-k +# must not be emitted with a NULL value +statement ok +CREATE TABLE topk_null_then_value(s varchar, y bigint) AS VALUES +('g1', CAST(NULL AS BIGINT)), +('g2', 10), +('g1', 5); + +query I +select max_y from (select s, max(y) as max_y from topk_null_then_value group by s) order by max_y desc nulls first limit 1; +---- +10 + +# single-row batches force NULL and non-NULL rows of the same group into +# different batches: a -> 7, b -> 5, c -> NULL +statement ok +set datafusion.execution.batch_size = 1; + +statement ok +CREATE TABLE topk_batches(s varchar, y bigint) AS VALUES +('a', CAST(NULL AS BIGINT)), +('b', 5), +('a', 3), +('c', CAST(NULL AS BIGINT)), +('b', CAST(NULL AS BIGINT)), +('a', 7); + +query TI +select s, max_y from (select s, max(y) as max_y from topk_batches group by s) order by max_y desc nulls first limit 3; +---- +c NULL +a 7 +b 5 + +# NULLS FIRST is not monotonic for MIN/MAX aggregation: a group initially +# registered as NULL can later become valued, so a bounded TopK cannot safely +# discard other NULL candidates. This must fall back to regular aggregation. +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +CREATE TABLE topk_null_backfill(s varchar, y bigint) AS VALUES +('a', CAST(NULL AS BIGINT)), +('b', CAST(NULL AS BIGINT)), +('c', CAST(NULL AS BIGINT)), +('a', 5); + +query I +select max_y from (select s, max(y) as max_y from topk_null_backfill group by s) order by max_y desc nulls first limit 2; +---- +NULL +NULL + +# An evicted valued group must not be re-registered and emitted as all-NULL. +statement ok +CREATE TABLE topk_evicted_then_null(s varchar, y bigint) AS VALUES +('a', 10), +('b', 20), +('a', CAST(NULL AS BIGINT)), +('c', CAST(NULL AS BIGINT)); + +query TI +select s, max_y from (select s, max(y) as max_y from topk_evicted_then_null group by s) order by max_y desc nulls first limit 1; +---- +c NULL + +statement ok +set datafusion.execution.batch_size = 8192; + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +drop table topk_batches; + +statement ok +drop table topk_evicted_then_null; + +statement ok +drop table topk_null_backfill; + +statement ok +drop table topk_null_then_value; + +statement ok +drop table topk_mixed; + +statement ok +drop table topk_five_nulls; + +statement ok +drop table topk_two_groups; + +statement ok +drop table t0; diff --git a/datafusion/sqllogictest/test_files/alias.slt b/datafusion/sqllogictest/test_files/alias.slt index 5339179db4c43..f19ce2a3b6e0b 100644 --- a/datafusion/sqllogictest/test_files/alias.slt +++ b/datafusion/sqllogictest/test_files/alias.slt @@ -57,3 +57,71 @@ drop table t1; statement count 0 drop table t2; + + +# Test table-aliasing a subquery with case sensitive columns +# (https://github.com/apache/datafusion/issues/22916) + +statement ok +create table t ("A" int, "B.C" int); + +query II +select * from (select * from t) t_(x, y); +---- + +query TT +explain select * from (select * from t) t_(x, y); +---- +logical_plan +01)SubqueryAlias: t_ +02)--Projection: t.A AS x, t.B.C AS y +03)----TableScan: t projection=[A, B.C] +physical_plan +01)ProjectionExec: expr=[A@0 as x, B.C@1 as y] +02)--DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain select * from (select * from t) t_(X, Y); +---- +logical_plan +01)SubqueryAlias: t_ +02)--Projection: t.A AS x, t.B.C AS y +03)----TableScan: t projection=[A, B.C] +physical_plan +01)ProjectionExec: expr=[A@0 as x, B.C@1 as y] +02)--DataSourceExec: partitions=1, partition_sizes=[0] + +query TT +explain select * from (select * from t) t_("X", "Y"); +---- +logical_plan +01)SubqueryAlias: t_ +02)--Projection: t.A AS X, t.B.C AS Y +03)----TableScan: t projection=[A, B.C] +physical_plan +01)ProjectionExec: expr=[A@0 as X, B.C@1 as Y] +02)--DataSourceExec: partitions=1, partition_sizes=[0] + +statement ok +insert into t values (1, 2); + +query II +select t_.x, t_.y +from (select "B.C", "A" from t) as t_(x, y); +---- +2 1 + +query I +select "x.y" +from (select "B.C" from t) as t_("x.y"); +---- +2 + +query I +select t_."x.y" +from (select "B.C" from t) as t_("x.y"); +---- +2 + +statement ok +drop table t; diff --git a/datafusion/sqllogictest/test_files/array/array_any_match.slt b/datafusion/sqllogictest/test_files/array/array_any_match.slt index 27f2a5339ef68..82133054e118a 100644 --- a/datafusion/sqllogictest/test_files/array/array_any_match.slt +++ b/datafusion/sqllogictest/test_files/array/array_any_match.slt @@ -103,6 +103,35 @@ SELECT list_any_match([1, 2, 3], x -> x > 2); ---- true +# null arg +query B +SELECT array_any_match(NULL, x -> x > 2); +---- +NULL + +# predicate can reference an outer column +query B +SELECT array_any_match(list, x -> x > number) FROM t; +---- +true +true +false + +# large list works +query B +SELECT array_any_match(arrow_cast([1, 2, 3], 'LargeList(Int32)'), x -> x > 2); +---- +true + +# other list representations are coerced during planning +query BBB +SELECT + array_any_match(arrow_cast([1, 2, 3], 'FixedSizeList(3, Int32)'), x -> x > 2), + array_any_match(arrow_cast([1, 2, 3], 'ListView(Int32)'), x -> x > 2), + array_any_match(arrow_cast([1, 2, 3], 'LargeListView(Int32)'), x -> x > 2); +---- +true true true + statement ok drop table t; diff --git a/datafusion/sqllogictest/test_files/array/array_any_value.slt b/datafusion/sqllogictest/test_files/array/array_any_value.slt index 6579e88ac7dba..c8976e8493261 100644 --- a/datafusion/sqllogictest/test_files/array/array_any_value.slt +++ b/datafusion/sqllogictest/test_files/array/array_any_value.slt @@ -145,6 +145,35 @@ select array_any_value(make_array(NULL, 1, 2, 3, 4, 5)), array_any_value(column1 1 41 1 51 +# array_any_value with empty (length-0) list elements +# A non-null but empty list must yield NULL, including a trailing empty element +# whose start offset equals the values length +statement ok +create table any_value_empty (id int, tags bigint[]) as values + (1, make_array(10)), + (2, cast(make_array() as bigint[])), + (3, make_array(20, 30)), + (4, cast(make_array() as bigint[])); + +query II +select id, array_any_value(tags) from any_value_empty order by id; +---- +1 10 +2 NULL +3 20 +4 NULL + +query II +select id, array_any_value(arrow_cast(tags, 'LargeList(Int64)')) from any_value_empty order by id; +---- +1 10 +2 NULL +3 20 +4 NULL + +statement ok +drop table any_value_empty; + # make_array with nulls query ??????? select make_array(make_array('a','b'), null), diff --git a/datafusion/sqllogictest/test_files/array/array_append.slt b/datafusion/sqllogictest/test_files/array/array_append.slt index 50949948c890e..0758a09a4925b 100644 --- a/datafusion/sqllogictest/test_files/array/array_append.slt +++ b/datafusion/sqllogictest/test_files/array/array_append.slt @@ -269,5 +269,67 @@ select array_append(column1, arrow_cast(make_array(1, 11, 111), 'FixedSizeList(3 [[1, 2, 3], [2, 9, 1], [7, 8, 9], [1, 2, 3], [1, 7, 4], [4, 5, 6], [1, 11, 111]] [[1, 2, 3], [11, 12, 13], [7, 8, 9]] [[4, 5, 6], [10, 11, 12], [4, 9, 8], [7, 8, 9], [10, 11, 12], [1, 8, 7], [1, 11, 111]] [[1, 2, 3], [11, 12, 13], [10, 11, 12]] +# the input list's inner field is carried through to the output, so that the +# returned type matches the one promised at planning time +query ?T +select + array_append(arrow_cast(column1, 'List(Int64, field: ''element'')'), 4), + arrow_typeof(array_append(arrow_cast(column1, 'List(Int64, field: ''element'')'), 4)) +from values (make_array(1, 2, 3)); +---- +[1, 2, 3, 4] List(Int64, field: 'element') + +query ?T +select + array_append(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 4), + arrow_typeof(array_append(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 4)) +from values (make_array(1, 2, 3)); +---- +[1, 2, 3, 4] LargeList(Int64, field: 'element') + +# nested value types go through a different kernel path +query ?T +select + array_append(arrow_cast(column1, 'List(List(Int64), field: ''element'')'), make_array(2)), + arrow_typeof(array_append(arrow_cast(column1, 'List(List(Int64), field: ''element'')'), make_array(2))) +from values (make_array(make_array(1))); +---- +[[1], [2]] List(List(Int64), field: 'element') + +# a non-nullable inner field stays non-nullable when the appended element cannot be null +query ??TT +select + array_append(column1, 4), + array_append(arrow_cast(column1, 'LargeList(non-null Int64)'), 4), + arrow_typeof(array_append(column1, 4)), + arrow_typeof(array_append(arrow_cast(column1, 'LargeList(non-null Int64)'), 4)) +from values + (arrow_cast(make_array(), 'List(non-null Int64)')), + (arrow_cast(NULL, 'List(non-null Int64)')), + (arrow_cast(make_array(1, 2, 3), 'List(non-null Int64)')) +; +---- +[4] [4] List(non-null Int64) LargeList(non-null Int64) +[4] [4] List(non-null Int64) LargeList(non-null Int64) +[1, 2, 3, 4] [1, 2, 3, 4] List(non-null Int64) LargeList(non-null Int64) + +# ... and is widened when the appended element is nullable, since the result +# genuinely contains a null element +query ?T +select + array_append(arrow_cast(make_array(1, 2), 'List(non-null Int64)'), NULL), + arrow_typeof(array_append(arrow_cast(make_array(1, 2), 'List(non-null Int64)'), NULL)); +---- +[1, 2, NULL] List(Int64) + +query ?T +select + array_append(arrow_cast(make_array(1, 2), 'List(non-null Int64)'), column1), + arrow_typeof(array_append(arrow_cast(make_array(1, 2), 'List(non-null Int64)'), column1)) +from values (3), (NULL); +---- +[1, 2, 3] List(Int64) +[1, 2, NULL] List(Int64) + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_concat.slt b/datafusion/sqllogictest/test_files/array/array_concat.slt index 168b307a1e636..5b7985ef6d194 100644 --- a/datafusion/sqllogictest/test_files/array/array_concat.slt +++ b/datafusion/sqllogictest/test_files/array/array_concat.slt @@ -419,5 +419,24 @@ select array_concat(make_array(column3), column1, column2) from arrays_values_v2 [NULL, 11, 12] [NULL] +# array_concat derives a fresh return type from the unified element types rather +# than cloning an input's, so its output field is always the default nullable +# `item` regardless of the inputs' inner fields +query ?T +select + array_concat(arrow_cast(column1, 'List(non-null Int64)'), make_array(3)), + arrow_typeof(array_concat(arrow_cast(column1, 'List(non-null Int64)'), make_array(3))) +from values (make_array(1, 2)); +---- +[1, 2, 3] List(Int64) + +query ?T +select + array_concat(arrow_cast(column1, 'List(Int64, field: ''element'')'), make_array(3)), + arrow_typeof(array_concat(arrow_cast(column1, 'List(Int64, field: ''element'')'), make_array(3))) +from values (make_array(1, 2)); +---- +[1, 2, 3] List(Int64) + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_distinct.slt b/datafusion/sqllogictest/test_files/array/array_distinct.slt index 88ffdf7f2ff78..2682413cac248 100644 --- a/datafusion/sqllogictest/test_files/array/array_distinct.slt +++ b/datafusion/sqllogictest/test_files/array/array_distinct.slt @@ -19,11 +19,10 @@ include ./init_data.slt.part ## array_distinct -#TODO: https://github.com/apache/datafusion/issues/7142 -#query ? -#select array_distinct(null); -#---- -#NULL +query ? +select array_distinct(null); +---- +NULL # test with empty row, the row that does not match the condition has row count 0 statement ok @@ -137,6 +136,12 @@ select array_compact(arrow_cast([NULL, NULL, NULL], 'List(Int64)')); ---- [] +# all nulls with untyped NULLs (List(Null) inner values are a NullArray) +query ? +select array_compact(make_array(NULL, NULL, NULL)); +---- +[] + # empty array query ? select array_compact([]); @@ -167,6 +172,17 @@ select array_compact([make_array(1, 2), NULL, make_array(3, 4)]); ---- [[1, 2], [3, 4]] +# nested array of all-null inner lists: outer elements are non-null, kept as-is +query ? +select array_compact(make_array(make_array(NULL, NULL, NULL), make_array(NULL, NULL))); +---- +[[NULL, NULL, NULL], [NULL, NULL]] + +query ? +select array_compact(make_array(make_array(NULL, NULL, NULL), make_array(NULL, NULL, NULL))); +---- +[[NULL, NULL, NULL], [NULL, NULL, NULL]] + # LargeList query ? select array_compact(arrow_cast([1, NULL, 2, NULL, 3], 'LargeList(Int64)')); @@ -210,5 +226,47 @@ select array_compact(arrow_cast(make_array(NULL, NULL, NULL), 'FixedSizeList(3, ---- [] +# Negative zero (-0.0) and positive zero (+0.0) compare equal under +# IEEE 754, but their bit patterns differ. array_distinct must normalize +# the sign so the canonical representative (+0.0) is used; otherwise +# group-by / dedup hashing on the raw bits keeps both as distinct +# elements. PostgreSQL / IEEE 754 expected output below. + +# array_distinct collapses +0.0 and -0.0 into a single element. +query ? +select array_distinct([0.0, -0.0]); +---- +[0.0] + +# General case with extra elements. +query ? +select array_distinct([0.0, -0.0, 0.0, 1.0, -0.0]); +---- +[0.0, 1.0] + +# array_length(array_distinct(...)) for {+0.0, -0.0, +0.0} must be 1. +query I +select array_length(array_distinct([0.0, -0.0, 0.0])); +---- +1 + +# Float32 list. +query ? +select array_distinct(arrow_cast([0.0, -0.0], 'List(Float32)')); +---- +[0.0] + +# LargeList(Float64). +query ? +select array_distinct(arrow_cast([0.0, -0.0], 'LargeList(Float64)')); +---- +[0.0] + +# FixedSizeList(Float64). +query ? +select array_distinct(arrow_cast([0.0, -0.0], 'FixedSizeList(2, Float64)')); +---- +[0.0] + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_element.slt b/datafusion/sqllogictest/test_files/array/array_element.slt index 7c960653edcf1..846d869ec667e 100644 --- a/datafusion/sqllogictest/test_files/array/array_element.slt +++ b/datafusion/sqllogictest/test_files/array/array_element.slt @@ -231,6 +231,22 @@ NULL NULL 55 +# array_element with null index from column +query I +select array_element(column1, column2) from slices where column2 is NULL; +---- +NULL + +query I +select array_element(arrow_cast(column1, 'LargeList(Int64)'), column2) from slices where column2 is NULL; +---- +NULL + +query I +select array_element(column1, column2) from fixed_slices where column2 is NULL; +---- +NULL + # array_element with columns and scalars query II select array_element(make_array(1, 2, 3, 4, 5), column2), array_element(column1, 3) from slices; diff --git a/datafusion/sqllogictest/test_files/array/array_empty.slt b/datafusion/sqllogictest/test_files/array/array_empty.slt index 62ac5f66b74c5..800f568934f2e 100644 --- a/datafusion/sqllogictest/test_files/array/array_empty.slt +++ b/datafusion/sqllogictest/test_files/array/array_empty.slt @@ -45,11 +45,10 @@ select empty(arrow_cast(make_array(), 'LargeList(Int64)')); ---- true -#TODO: https://github.com/apache/datafusion/issues/9158 -#query B -#select empty(arrow_cast(make_array(), 'FixedSizeList(0, Null)')); -#---- -#true +query B +select empty(arrow_cast(make_array(), 'FixedSizeList(0, Null)')); +---- +true # empty scalar function #3 query B @@ -69,10 +68,8 @@ false #TODO: https://github.com/apache/datafusion/issues/7142 # empty scalar function #4 -#query B -#select empty(NULL); -#---- -#NULL +query error array_empty does not support type Null +select empty(NULL); # empty scalar function #5 query B diff --git a/datafusion/sqllogictest/test_files/array/array_except.slt b/datafusion/sqllogictest/test_files/array/array_except.slt index a718723e58c38..1d41a5a79d15f 100644 --- a/datafusion/sqllogictest/test_files/array/array_except.slt +++ b/datafusion/sqllogictest/test_files/array/array_except.slt @@ -156,4 +156,52 @@ select array_except(arrow_cast([1, 2, 3, 4], 'FixedSizeList(4, Int64)'), arrow_c [1, 2] +# Negative zero (-0.0) and positive zero (+0.0) compare equal under +# IEEE 754, but their bit patterns differ. array_except must treat +# +0.0 and -0.0 as the same element when subtracting. PostgreSQL / +# IEEE 754 expected output below. + +# -0.0 in rhs removes +0.0 in lhs. +query ? +select array_except([0.0], [-0.0]); +---- +[] + +# Reverse direction. +query ? +select array_except([-0.0], [0.0]); +---- +[] + +# -0.0 in rhs also removes the +0.0 element from the lhs. +query ? +select array_except([0.0, -0.0], [-0.0]); +---- +[] + +# +0.0 in rhs also removes the -0.0 element from the lhs. +query ? +select array_except([0.0, -0.0], [0.0]); +---- +[] + +# More general case with extra unmatched element. +query ? +select array_except([0.0, -0.0, 1.0], [-0.0]); +---- +[1.0] + +# Float32 list. +query ? +select array_except(arrow_cast([0.0, -0.0], 'List(Float32)'), arrow_cast([0.0], 'List(Float32)')); +---- +[] + +# LargeList(Float64). +query ? +select array_except(arrow_cast([0.0, -0.0], 'LargeList(Float64)'), arrow_cast([-0.0], 'LargeList(Float64)')); +---- +[] + + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_filter.slt b/datafusion/sqllogictest/test_files/array/array_filter.slt index f22cfb219830c..b6d73fbe7d09d 100644 --- a/datafusion/sqllogictest/test_files/array/array_filter.slt +++ b/datafusion/sqllogictest/test_files/array/array_filter.slt @@ -120,6 +120,20 @@ SELECT array_filter(arrow_cast(list, 'ListView(Int32)'), v -> v > 2) from t; [4, 50] [7, 50] +# large list works +query ? +SELECT array_filter(arrow_cast([1, 2, 3, 4, 5], 'LargeList(Int32)'), v -> v > 2); +---- +[3, 4, 5] + +# FixedSizeList / LargeListView coercions during planning +query ?? +SELECT + array_filter(arrow_cast([1, 2, 3, 4], 'FixedSizeList(4, Int32)'), v -> v > 2), + array_filter(arrow_cast([1, 2, 3, 4], 'LargeListView(Int32)'), v -> v > 2); +---- +[3, 4] [3, 4] + # null array argument returns null query ? SELECT array_filter(arrow_cast(NULL, 'List(Int32)'), v -> v > 0); @@ -204,6 +218,12 @@ SELECT array_transform(array_filter(list, v -> v > 1), v -> v * 3) FROM with_nul [6] NULL +# null arg +query ? +SELECT array_filter(NULL, x -> x > 2); +---- +NULL + statement ok drop table t; diff --git a/datafusion/sqllogictest/test_files/array/array_first.slt b/datafusion/sqllogictest/test_files/array/array_first.slt new file mode 100644 index 0000000000000..d761c3a4d1f0a --- /dev/null +++ b/datafusion/sqllogictest/test_files/array/array_first.slt @@ -0,0 +1,127 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +############# +## array_first Tests +############# + +statement ok +set datafusion.sql_parser.dialect = databricks; + +statement ok +CREATE TABLE t (list array, number int) +AS VALUES +([1, 50], 10), +([4, 50], 40), +([7, 50], 60); + +# basic: returns the first element that matches the predicate +query I +SELECT array_first([1, 2, 3, 4], x -> x > 2); +---- +3 + +# no element matches returns null +query I +SELECT array_first([1, 2, 3], x -> x > 5); +---- +NULL + +# empty array returns null +query I +SELECT array_first(arrow_cast(make_array(), 'List(Int32)'), x -> x > 0); +---- +NULL + +# null array returns null +query I +SELECT array_first(arrow_cast(NULL, 'List(Int32)'), x -> x > 0); +---- +NULL + +# a predicate that returns null for an element is treated as not matching +query I +SELECT array_first([1, 2, NULL, 4], x -> x > 2); +---- +4 + +# the predicate may match a null element, which is returned as null +query I +SELECT array_first(arrow_cast([NULL, 2], 'List(Int32)'), x -> x IS NULL); +---- +NULL + +# predicate always returns null -> no match -> null +query I +SELECT array_first([1, 2, 3], x -> NULL::boolean); +---- +NULL + +# a predicate matching every element returns the first element +query I +SELECT array_first([10, 20, 30], x -> true); +---- +10 + +# string elements +query T +SELECT array_first(['a', 'bb', 'ccc'], x -> length(x) > 1); +---- +bb + +# multiple rows +query I +SELECT array_first(list, x -> x > 5) FROM t; +---- +50 +50 +7 + +# predicate can reference an outer column (last row has no match -> null) +query I +SELECT array_first(list, x -> x > number) FROM t; +---- +50 +50 +NULL + +# large list works +query I +SELECT array_first(arrow_cast([1, 2, 3, 4], 'LargeList(Int32)'), x -> x > 2); +---- +3 + +# other list representations are coerced during planning +query III +SELECT + array_first(arrow_cast([1, 2, 3, 4], 'FixedSizeList(4, Int32)'), x -> x > 2), + array_first(arrow_cast([1, 2, 3, 4], 'ListView(Int32)'), x -> x > 2), + array_first(arrow_cast([1, 2, 3, 4], 'LargeListView(Int32)'), x -> x > 2); +---- +3 3 3 + +# alias array_first/list_first work +query I +SELECT list_first([1, 2, 3, 4], x -> x > 2); +---- +3 + +statement ok +drop table t; + +statement ok +set datafusion.sql_parser.dialect = generic; diff --git a/datafusion/sqllogictest/test_files/array/array_has.slt b/datafusion/sqllogictest/test_files/array/array_has.slt index e343c1b1fae41..14bc331d8f2d9 100644 --- a/datafusion/sqllogictest/test_files/array/array_has.slt +++ b/datafusion/sqllogictest/test_files/array/array_has.slt @@ -41,13 +41,15 @@ select array_has([1, null, 2], 3), false false #TODO: array_has_all and array_has_any cannot handle NULL -#query BBBB -#select array_has_any([], null), -# array_has_any([1, 2, 3], null), -# array_has_all([], null), -# array_has_all([1, 2, 3], null); -#---- -#false false false false +query BB +select array_has_any([], null), + array_has_any([1, 2, 3], null); +---- +NULL NULL + +query error array_has does not support type 'Null' +select array_has_all([], null), + array_has_all([1, 2, 3], null); query BBBBBBBBBBBB select array_has(make_array(1,2), 1), @@ -894,4 +896,118 @@ statement ok DROP TABLE any_op_test; +# ------------------------------------------------------------------------- +# array_has with an array (column) needle -- one needle value per row, which +# goes through array_has_dispatch_for_array (the cases above use a scalar +# literal needle and take a different path). +# ------------------------------------------------------------------------- + +statement ok +create table array_has_int_needle (arr int[], needle int) as values + ([1, 2, 3], 2), -- found + ([4, 5, 6], 9), -- not found + (NULL, 5), -- null row + ([7, NULL, 9], NULL), -- null needle + ([7, NULL, 9], 7), -- element null skipped, found + ([0, NULL], 0), -- valid 0 matches + ([NULL, 5], 0), -- null-fill collision: a null slot must not match 0 + ([], 1), -- empty + ([NULL, NULL], 3); -- all null + +query B +select array_has(arr, needle) from array_has_int_needle; +---- +true +false +NULL +NULL +true +true +false +false +false + +# same over LargeList (i64) offsets +query B +select array_has(arrow_cast(arr, 'LargeList(Int32)'), needle) from array_has_int_needle; +---- +true +false +NULL +NULL +true +true +false +false +false + +statement ok +drop table array_has_int_needle; + +statement ok +create table array_has_str_needle (arr text[], needle text) as values + (['a', 'bb', 'ccc'], 'bb'), -- inline, found + (['short', 'tiny'], 'missing'), -- inline, not found + (['this_is_a_long_value_xyz'], 'this_is_a_long_value_xyz'), -- long, found + (['prefixAAAA_1111', 'prefixAAAA_2222'], 'prefixAAAA_2222'), -- long shared prefix + (['x', NULL, 'y'], 'y'), -- element null skipped + ([NULL], ''), -- null slot vs "" -> false + (NULL, 'q'); -- null row + +query B +select array_has(arr, needle) from array_has_str_needle; +---- +true +false +true +true +true +false +NULL + +# Utf8View exercises the view-aware fast path +query B +select array_has(arrow_cast(arr, 'List(Utf8View)'), arrow_cast(needle, 'Utf8View')) +from array_has_str_needle; +---- +true +false +true +true +true +false +NULL + +# LargeUtf8 elements +query B +select array_has(arrow_cast(arr, 'LargeList(LargeUtf8)'), arrow_cast(needle, 'LargeUtf8')) +from array_has_str_needle; +---- +true +false +true +true +true +false +NULL + +statement ok +drop table array_has_str_needle; + +# > ROW_CONVERSION_CHUNK_SIZE (512) rows with element nulls exercises the chunked +# element-null path. The needle equals an element that is always present, so all +# rows match; the second query shifts the needle out of range, so none do. +query I +select count(*) from generate_series(1, 2000) as t(v) +where array_has(make_array(v % 7, NULL, (v + 2) % 7), v % 7); +---- +2000 + +query I +select count(*) from generate_series(1, 2000) as t(v) +where array_has(make_array(v % 7, NULL, (v + 2) % 7), v % 7 + 100); +---- +0 + + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_index.slt b/datafusion/sqllogictest/test_files/array/array_index.slt index 9cd033418d24b..1d9e2989e2342 100644 --- a/datafusion/sqllogictest/test_files/array/array_index.slt +++ b/datafusion/sqllogictest/test_files/array/array_index.slt @@ -94,17 +94,23 @@ NULL NULL e [13, 14] NULL NULL [NULL, 18] NULL NULL -# TODO: support index as column # single index with columns #5 (index as column) -# query ? -# select make_array(1, 2, 3, 4, 5)[column2] from arrays_with_repeating_elements; -# ---- +query I +select make_array(1, 2, 3, 4, 5)[column2] from arrays_with_repeating_elements; +---- +2 +4 +NULL +NULL -# TODO: support argument and index as columns # single index with columns #6 (argument and index as columns) -# query I -# select column1[column2] from arrays_with_repeating_elements; -# ---- +query I +select column1[column2] from arrays_with_repeating_elements; +---- +2 +5 +7 +10 ## array[i:j] @@ -141,17 +147,17 @@ select arrow_cast([1, 2, 3], 'LargeList(Int64)')[1]; ---- 1 -# TODO: support multiple negative index # multiple index with columns #3 (negative index) -# query II -# select make_array(1, 2, 3)[-3:-1], make_array(1.0, 2.0, 3.0)[-3:-1], make_array('h', 'e', 'l', 'l', 'o')[-2:0]; -# ---- +query ??? +select make_array(1, 2, 3)[-3:-1], make_array(1.0, 2.0, 3.0)[-3:-1], make_array('h', 'e', 'l', 'l', 'o')[-2:0]; +---- +[1, 2, 3] [1.0, 2.0, 3.0] [] -# TODO: support complex index # multiple index with columns #4 (complex index) -# query III -# select make_array(1, 2, 3)[2 + 1 - 1:10], make_array(1.0, 2.0, 3.0)[2 | 2:10], make_array('h', 'e', 'l', 'l', 'o')[6 ^ 6:10]; -# ---- +query ??? +select make_array(1, 2, 3)[2 + 1 - 1:10], make_array(1.0, 2.0, 3.0)[(2 | 2):10], make_array('h', 'e', 'l', 'l', 'o')[6 ^ 6:10]; +---- +[2, 3] [2.0, 3.0] [h, e, l, l, o] # multiple index with columns #1 (positive index) query ??? @@ -177,36 +183,56 @@ NULL [13.3, 14.4, 15.5] [a, m, e, t] [[11, 12], [13, 14]] NULL [,] [[15, 16], [NULL, 18]] [16.6, 17.7, 18.8] NULL -# TODO: support negative index # multiple index with columns #3 (negative index) -# query ?RT -# select column1[-2:-4], column2[-3:-5], column3[-1:-4] from arrays; -# ---- -# [NULL, 2] 1.1 m +query ??? +select column1[-2:-4], column2[-3:-5], column3[-1:-4] from arrays; +---- +[] [] [] +[] [] [] +[] [] [] +[] [] [] +NULL [] [] +[] NULL [] +[] [] NULL -# TODO: support complex index # multiple index with columns #4 (complex index) -# query ?RT -# select column1[9 - 7:2 + 2], column2[1 * 0:2 * 3], column3[1 + 1 - 0:5 % 3] from arrays; -# ---- +query ??? +select column1[9 - 7:2 + 2], column2[1 * 0:2 * 3], column3[1 + 1 - 0:5 % 3] from arrays; +---- +[[3, NULL]] [1.1, 2.2, 3.3] [o] +[[5, 6]] [NULL, 5.5, 6.6] [p] +[[7, 8]] [7.7, 8.8, 9.9] [NULL] +[[9, 10]] [10.1, NULL, 12.2] [i] +NULL [13.3, 14.4, 15.5] [m] +[[13, 14]] NULL [] +[[NULL, 18]] [16.6, 17.7, 18.8] NULL -# TODO: support first index as column # multiple index with columns #5 (first index as column) -# query ? -# select make_array(1, 2, 3, 4, 5)[column2:4] from arrays_with_repeating_elements -# ---- +query ? +select make_array(1, 2, 3, 4, 5)[column2:4] from arrays_with_repeating_elements +---- +[2, 3, 4] +[4] +[] +[] -# TODO: support last index as column # multiple index with columns #6 (last index as column) -# query ?RT -# select make_array(1, 2, 3, 4, 5)[2:column3] from arrays_with_repeating_elements; -# ---- +query ? +select make_array(1, 2, 3, 4, 5)[2:column3] from arrays_with_repeating_elements; +---- +[2, 3, 4] +[2, 3, 4, 5] +[2, 3, 4, 5] +[2, 3, 4, 5] -# TODO: support argument and indices as column # multiple index with columns #7 (argument and indices as column) -# query ?RT -# select column1[column2:column3] from arrays_with_repeating_elements; -# ---- +query ? +select column1[column2:column3] from arrays_with_repeating_elements; +---- +[2, 1, 3] +[5, 6, 5, 5] +[7, 8, 7, 7] +[10] # array[i:j:k] @@ -222,12 +248,11 @@ select make_array(1, 2, 3)[0:0:2], make_array(1.0, 2.0, 3.0)[0:2:2], make_array( ---- [] [1.0] [h, l, o] -#TODO: sqlparser does not support negative index ## multiple index with columns #3 (negative index) -#query ??? -#select make_array(1, 2, 3)[-1:-2:-2], make_array(1.0, 2.0, 3.0)[-2:-3:-2], make_array('h', 'e', 'l', 'l', 'o')[-2:-4:-2]; -#---- -#[1] [2.0] [e, l] +query ??? +select make_array(1, 2, 3)[-1:-2:-2], make_array(1.0, 2.0, 3.0)[-2:-3:-2], make_array('h', 'e', 'l', 'l', 'o')[-2:-4:-2]; +---- +[3] [2.0] [l, e] # multiple index with columns #1 (positive index) query ??? diff --git a/datafusion/sqllogictest/test_files/array/array_length.slt b/datafusion/sqllogictest/test_files/array/array_length.slt index 1bb5382339854..7741d815bc234 100644 --- a/datafusion/sqllogictest/test_files/array/array_length.slt +++ b/datafusion/sqllogictest/test_files/array/array_length.slt @@ -159,21 +159,6 @@ select array_distance([2], [3]), list_distance([1], [2]), list_distance([1], [-2 query error select list_distance([1], [1, 2]); -query R -select array_distance([[1, 1]], [1, 2]); ----- -1 - -query R -select array_distance([[1, 1]], [[1, 2]]); ----- -1 - -query R -select array_distance([[1, 1]], [[1, 2]]); ----- -1 - query RR select array_distance([1, 1, 0, 0], [2, 2, 1, 1]), list_distance([1, 2, 3], [1, 2, 3]); ---- @@ -204,6 +189,40 @@ select list_distance([1, 2, 3], [1, 2, 3]) AS distance; ---- 0 +# array_distance with null outer arrays +query RR +select + array_distance(arrow_cast(NULL, 'List(Float64)'), [1, 2]), + array_distance([1, 2], arrow_cast(NULL, 'List(Float64)')); +---- +NULL NULL + +# invalid argument count and types +query error DataFusion error: Error during planning: Execution error: Function 'array_distance' user-defined coercion failed with: Execution error: array_distance function requires 2 arguments +select array_distance(); + +query error DataFusion error: Error during planning: Execution error: Function 'array_distance' user-defined coercion failed with: Execution error: array_distance function requires 2 arguments +select array_distance([1]); + +query error DataFusion error: Error during planning: Execution error: Function 'array_distance' user-defined coercion failed with: Execution error: array_distance function requires 2 arguments +select array_distance([1], [2], [3]); + +query error array_distance does not support type Int64 +select array_distance(1, [1]); + +query error array_distance does not support types +select array_distance([1], arrow_cast([1], 'LargeList(Float64)')); + +query error array_distance only supports one-dimensional arrays +select array_distance([[1, 1]], [1, 2]); + +query error array_distance only supports one-dimensional arrays +select array_distance([[1, 1]], [[1, 2]]); + +query error array_distance only supports one-dimensional arrays +select array_distance([[1, 2], [100, 100]], [[1, 4], [0, 0]]); + + # array_distance with columns query RRR select array_distance(column1, column2), array_distance(column1, column3), array_distance(column1, column4) from arrays_distance_table; diff --git a/datafusion/sqllogictest/test_files/array/array_pop.slt b/datafusion/sqllogictest/test_files/array/array_pop.slt index b830fa464a984..0b7ebf75f4b2b 100644 --- a/datafusion/sqllogictest/test_files/array/array_pop.slt +++ b/datafusion/sqllogictest/test_files/array/array_pop.slt @@ -22,10 +22,8 @@ include ./init_data.slt.part # array_pop_back scalar function with null #TODO: https://github.com/apache/datafusion/issues/7142 # follow clickhouse and duckdb -#query ? -#select array_pop_back(null); -#---- -#NULL +query error array_pop_back does not support type: Null +select array_pop_back(null); # array_pop_back scalar function #1 query ?? @@ -201,10 +199,8 @@ NULL #TODO:https://github.com/apache/datafusion/issues/7142 # array_pop_front scalar function with null # follow clickhouse and duckdb -#query ? -#select array_pop_front(null); -#---- -#NULL +query error array_pop_front does not support type: Null +select array_pop_front(null); # array_pop_front scalar function #1 query ?? @@ -322,5 +318,38 @@ select array_pop_front(arrow_cast([1, 2], 'LargeListView(Int64)')); ---- [2] +# maintains inner nullability +query ??TT +select + array_pop_front(column1), + array_pop_back(column1), + arrow_typeof(array_pop_front(column1)), + arrow_typeof(array_pop_back(column1)) +from values + (arrow_cast([], 'List(non-null Int32)')), + (arrow_cast(NULL, 'List(non-null Int32)')), + (arrow_cast([1, 3, 5, -5], 'List(non-null Int32)')) +; +---- +[] [] List(non-null Int32) List(non-null Int32) +NULL NULL List(non-null Int32) List(non-null Int32) +[3, 5, -5] [1, 3, 5] List(non-null Int32) List(non-null Int32) + +query ??TT +select + array_pop_front(column1), + array_pop_back(column1), + arrow_typeof(array_pop_front(column1)), + arrow_typeof(array_pop_back(column1)) +from values + (arrow_cast([], 'LargeList(non-null Int32)')), + (arrow_cast(NULL, 'LargeList(non-null Int32)')), + (arrow_cast([1, 3, 5, -5], 'LargeList(non-null Int32)')) +; +---- +[] [] LargeList(non-null Int32) LargeList(non-null Int32) +NULL NULL LargeList(non-null Int32) LargeList(non-null Int32) +[3, 5, -5] [1, 3, 5] LargeList(non-null Int32) LargeList(non-null Int32) + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_position.slt b/datafusion/sqllogictest/test_files/array/array_position.slt index 07e3d3143592c..e3dd830dfb77a 100644 --- a/datafusion/sqllogictest/test_files/array/array_position.slt +++ b/datafusion/sqllogictest/test_files/array/array_position.slt @@ -314,10 +314,8 @@ select array_positions([1, 2, 3, 4, 5], null); #TODO: https://github.com/apache/datafusion/issues/7142 # array_positions with NULL (follow PostgreSQL) -#query ? -#select array_positions(null, 1); -#---- -#NULL +query error array_positions does not support type 'Null' +select array_positions(null, 1); # array_positions scalar function #1 query ??? diff --git a/datafusion/sqllogictest/test_files/array/array_prepend.slt b/datafusion/sqllogictest/test_files/array/array_prepend.slt index 0782680ed2de9..bfb61ab4f9f91 100644 --- a/datafusion/sqllogictest/test_files/array/array_prepend.slt +++ b/datafusion/sqllogictest/test_files/array/array_prepend.slt @@ -57,7 +57,6 @@ select array_prepend(null, [[1,2,3]]); # DuckDB: [[]] # ClickHouse: [[]] -# TODO: We may also return [[]] query ? select array_prepend([], []); ---- @@ -274,5 +273,67 @@ select array_prepend(arrow_cast(make_array(1, 11, 111), 'FixedSizeList(3, Int64) [[1, 11, 111], [1, 2, 3], [2, 9, 1], [7, 8, 9], [1, 2, 3], [1, 7, 4], [4, 5, 6]] [[7, 8, 9], [1, 2, 3], [11, 12, 13]] [[1, 11, 111], [4, 5, 6], [10, 11, 12], [4, 9, 8], [7, 8, 9], [10, 11, 12], [1, 8, 7]] [[10, 11, 12], [1, 2, 3], [11, 12, 13]] +# the input list's inner field is carried through to the output, so that the +# returned type matches the one promised at planning time +query ?T +select + array_prepend(0, arrow_cast(column1, 'List(Int64, field: ''element'')')), + arrow_typeof(array_prepend(0, arrow_cast(column1, 'List(Int64, field: ''element'')'))) +from values (make_array(1, 2, 3)); +---- +[0, 1, 2, 3] List(Int64, field: 'element') + +query ?T +select + array_prepend(0, arrow_cast(column1, 'LargeList(Int64, field: ''element'')')), + arrow_typeof(array_prepend(0, arrow_cast(column1, 'LargeList(Int64, field: ''element'')'))) +from values (make_array(1, 2, 3)); +---- +[0, 1, 2, 3] LargeList(Int64, field: 'element') + +# nested value types go through a different kernel path +query ?T +select + array_prepend(make_array(0), arrow_cast(column1, 'List(List(Int64), field: ''element'')')), + arrow_typeof(array_prepend(make_array(0), arrow_cast(column1, 'List(List(Int64), field: ''element'')'))) +from values (make_array(make_array(1))); +---- +[[0], [1]] List(List(Int64), field: 'element') + +# a non-nullable inner field stays non-nullable when the prepended element cannot be null +query ??TT +select + array_prepend(0, column1), + array_prepend(0, arrow_cast(column1, 'LargeList(non-null Int64)')), + arrow_typeof(array_prepend(0, column1)), + arrow_typeof(array_prepend(0, arrow_cast(column1, 'LargeList(non-null Int64)'))) +from values + (arrow_cast(make_array(), 'List(non-null Int64)')), + (arrow_cast(NULL, 'List(non-null Int64)')), + (arrow_cast(make_array(1, 2, 3), 'List(non-null Int64)')) +; +---- +[0] [0] List(non-null Int64) LargeList(non-null Int64) +[0] [0] List(non-null Int64) LargeList(non-null Int64) +[0, 1, 2, 3] [0, 1, 2, 3] List(non-null Int64) LargeList(non-null Int64) + +# ... and is widened when the prepended element is nullable, since the result +# genuinely contains a null element +query ?T +select + array_prepend(NULL, arrow_cast(make_array(1, 2), 'List(non-null Int64)')), + arrow_typeof(array_prepend(NULL, arrow_cast(make_array(1, 2), 'List(non-null Int64)'))); +---- +[NULL, 1, 2] List(Int64) + +query ?T +select + array_prepend(column1, arrow_cast(make_array(1, 2), 'List(non-null Int64)')), + arrow_typeof(array_prepend(column1, arrow_cast(make_array(1, 2), 'List(non-null Int64)'))) +from values (0), (NULL); +---- +[0, 1, 2] List(Int64) +[NULL, 1, 2] List(Int64) + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_remove.slt b/datafusion/sqllogictest/test_files/array/array_remove.slt index c3ce7073eca83..3088042ff1400 100644 --- a/datafusion/sqllogictest/test_files/array/array_remove.slt +++ b/datafusion/sqllogictest/test_files/array/array_remove.slt @@ -63,13 +63,32 @@ select ---- [1, NULL, 3] [NULL, 2.2, 3.3] [NULL, bc] -#TODO: https://github.com/apache/datafusion/issues/7142 # follow PostgreSQL behavior -#query ? -#select -# array_remove(NULL, 1) -#---- -#NULL +# A NULL-typed array argument returns NULL, matching array_replace and SQL +# three-valued logic. +query ? +select array_remove(NULL, 1); +---- +NULL + +query ? +select array_remove_n(NULL, 1, 2); +---- +NULL + +query ? +select array_remove(column1, 1) from (values (NULL), (NULL), (NULL)); +---- +NULL +NULL +NULL + +query ? +select array_remove(column1, column2) from (values (NULL, 1), (NULL, 2), (NULL, 3)) as t(column1, column2); +---- +NULL +NULL +NULL query ?? select @@ -266,6 +285,13 @@ select array_remove_n(make_array(1, 2, 2, 1, 1), NULL, 2), ---- NULL [1, 1, 1] +# array_remove_n with null max scalar +query ?? +select array_remove_n(make_array(1, 2, 2, 1, 1), 2, NULL), + array_remove_n(arrow_cast(make_array(1, 2, 2, 1, 1), 'LargeList(Int64)'), 2, NULL); +---- +NULL NULL + # array_remove_n with null element scalar (LargeList) query ?? select array_remove_n(arrow_cast(make_array(1, 2, 2, 1, 1), 'LargeList(Int64)'), NULL, 2), @@ -273,12 +299,27 @@ select array_remove_n(arrow_cast(make_array(1, 2, 2, 1, 1), 'LargeList(Int64)'), ---- NULL [1, 1, 1] +# array_remove_n with null max scalar +query ?? +select array_remove_n(make_array(1, 2, 2, 1, 1), 2, NULL), + array_remove_n(make_array(1, 2, 2, 1, 1), 2, 2); +---- +NULL [1, 1, 1] + +# array_remove_n with null max scalar (LargeList) +query ?? +select array_remove_n(arrow_cast(make_array(1, 2, 2, 1, 1), 'LargeList(Int64)'), 2, NULL), + array_remove_n(arrow_cast(make_array(1, 2, 2, 1, 1), 'LargeList(Int64)'), 2, 2); +---- +NULL [1, 1, 1] + # array_remove_n with null element from column query ? select array_remove_n(column1, column2, column3) from (values (make_array(1, 2, 2, 1, 1), 2, 2), (make_array(3, 4, 4, 3, 3), null, 2), (make_array(5, 6, 6, 5, 5), 6, 1), + (make_array(7, 8, 8, 7, 7), 8, null), (null, 1, 1) ) as t(column1, column2, column3); ---- @@ -286,18 +327,21 @@ select array_remove_n(column1, column2, column3) from (values NULL [5, 6, 5, 5] NULL +NULL # array_remove_n with null element from column (LargeList) query ? select array_remove_n(column1, column2, column3) from (values (arrow_cast(make_array(1, 2, 2, 1, 1), 'LargeList(Int64)'), 2, 2), (arrow_cast(make_array(3, 4, 4, 3, 3), 'LargeList(Int64)'), null, 2), - (arrow_cast(make_array(5, 6, 6, 5, 5), 'LargeList(Int64)'), 6, 1) + (arrow_cast(make_array(5, 6, 6, 5, 5), 'LargeList(Int64)'), 6, 1), + (arrow_cast(make_array(7, 8, 8, 7, 7), 'LargeList(Int64)'), 8, null) ) as t(column1, column2, column3); ---- [1, 1, 1] NULL [5, 6, 5, 5] +NULL # array_remove_n scalar function #1 query ??? @@ -381,12 +425,11 @@ select array_remove_n(make_array([1, 2, 3], [4, 5, 6], [4, 5, 6], [10, 11, 12], ## array_remove_all (aliases: `list_removes`) -#TODO: https://github.com/apache/datafusion/issues/7142 # array_remove_all with NULL elements -#query ? -#select array_remove_all(NULL, 1); -#---- -#NULL +query ? +select array_remove_all(NULL, 1); +---- +NULL query ? select array_remove_all(make_array(1, 2, 2, 1, 1), NULL); @@ -537,4 +580,76 @@ select array_remove_all(make_array([1, 2, 3], [4, 5, 6], [4, 5, 6], [10, 11, 12] [[1, 2, 3], [4, 5, 6], [4, 5, 6], [10, 11, 12], [13, 14, 15], [10, 11, 12], [10, 11, 12], [19, 20, 21], [19, 20, 21], [19, 20, 21], [22, 23, 24]] [[28, 29, 30], [31, 32, 33], [34, 35, 36], [28, 29, 30], [31, 32, 33], [34, 35, 36], [28, 29, 30], [31, 32, 33], [34, 35, 36], [28, 29, 30]] +# array_remove scalar arguments over multiple input rows +query ??? +select + array_remove(column1, 2), + array_remove_n(column1, 2, 2), + array_remove_all(column1, 2) +from ( + values + (make_array(1, 2, 2, 3, 2, 1, 4)), + (make_array(42, 2, 55, 63, 2)) +) as t(column1); +---- +[1, 2, 3, 2, 1, 4] [1, 3, 2, 1, 4] [1, 3, 1, 4] +[42, 55, 63, 2] [42, 55, 63] [42, 55, 63] + +# array_remove with elements containing NULLs — scalar path preserves NULLs +query ??? +select + array_remove(column1, 2), + array_remove_n(column1, 2, 2), + array_remove_all(column1, 2) +from ( + values + (make_array(1, 2, NULL, 3, 2, NULL, 4)), + (make_array(42, 2, NULL, 63, 2)) +) as t(column1); +---- +[1, NULL, 3, 2, NULL, 4] [1, NULL, 3, NULL, 4] [1, NULL, 3, NULL, 4] +[42, NULL, 63, 2] [42, NULL, 63] [42, NULL, 63] + +# array_remove_n with n exceeding match count +query ? +select array_remove_n(make_array(1, 2, 2, 3), 2, 100); +---- +[1, 3] + +# array_remove_n with n=0 and n=-1 (no removal) +query ?? +select + array_remove_n(make_array(1, 2, 2, 3), 2, 0), + array_remove_n(make_array(1, 2, 2, 3), 2, -1); +---- +[1, 2, 2, 3] [1, 2, 2, 3] + +# array_remove on empty arrays +query ?? +select + array_remove(arrow_cast(make_array(), 'List(Int64)'), 1), + array_remove_all(arrow_cast(make_array(), 'List(Int64)'), 1); +---- +[] [] + +# array_remove needle not found — array unchanged +query ? +select array_remove_all(make_array(1, 2, 3, 4, 5), 99); +---- +[1, 2, 3, 4, 5] + +# array_remove all elements match +query ? +select array_remove_all(make_array(7, 7, 7, 7), 7); +---- +[] + +# LargeList scalar path edge cases +query ?? +select + array_remove_all(arrow_cast(make_array(1, 1, 1), 'LargeList(Int64)'), 1), + array_remove_n(arrow_cast(make_array(1, 1, 1), 'LargeList(Int64)'), 1, 2); +---- +[] [1] + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/array_repeat.slt b/datafusion/sqllogictest/test_files/array/array_repeat.slt index 8052f09cb32c7..5073c7d4c5822 100644 --- a/datafusion/sqllogictest/test_files/array/array_repeat.slt +++ b/datafusion/sqllogictest/test_files/array/array_repeat.slt @@ -43,6 +43,9 @@ select ---- [[1], [1], [1], [1], [1]] [[1.1, 2.2, 3.3], [1.1, 2.2, 3.3], [1.1, 2.2, 3.3]] [[NULL, NULL], [NULL, NULL], [NULL, NULL]] [[[1, 2], [3, 4]], [[1, 2], [3, 4]]] +query error DataFusion error: Execution error: array_repeat: requested length exceeds maximum array size +select array_repeat(1, 9223372036854775807); + query ???? select array_repeat(arrow_cast([1], 'LargeList(Int64)'), 5), @@ -76,6 +79,16 @@ Select ---- [] [] [] [] +# array_repeat returns an execution error on scalar output-size overflow +query error DataFusion error: Execution error: array_repeat: total repeated values overflowed usize +SELECT array_repeat(1, c) +FROM ( + VALUES + (9223372036854775807), + (9223372036854775807), + (9223372036854775807) +) AS t(c); + # array_repeat with columns #1 statement ok diff --git a/datafusion/sqllogictest/test_files/array/array_replace.slt b/datafusion/sqllogictest/test_files/array/array_replace.slt index 390ed4b946520..ce45e6440dddf 100644 --- a/datafusion/sqllogictest/test_files/array/array_replace.slt +++ b/datafusion/sqllogictest/test_files/array/array_replace.slt @@ -118,6 +118,37 @@ select array_replace(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)'), ---- [1, 2, 3, 4, 5] +# A NULL-typed array argument returns NULL, for both literal and multi-row +# column inputs. +query ? +select array_replace(NULL, 1, 2); +---- +NULL + +query ? +select array_replace_n(NULL, 1, 2, 3); +---- +NULL + +query ? +select array_replace_all(NULL, 1, 2); +---- +NULL + +query ? +select array_replace(column1, 1, 2) from (values (NULL), (NULL), (NULL)); +---- +NULL +NULL +NULL + +query ? +select array_replace(column1, column2, column3) from (values (NULL, 1, 2), (NULL, 3, 4), (NULL, 5, 6)) as t(column1, column2, column3); +---- +NULL +NULL +NULL + # array_replace scalar function with columns #1 query ? select array_replace(column1, column2, column3) from arrays_with_repeating_elements; @@ -212,6 +243,33 @@ from large_nested_arrays_with_repeating_elements; [[1, 2, 3], [4, 5, 6], [4, 5, 6], [10, 11, 12], [13, 14, 15], [10, 11, 12], [10, 11, 12], [28, 29, 30], [19, 20, 21], [28, 29, 30], [19, 20, 21], [22, 23, 24]] [[19, 20, 21], [19, 20, 21], [19, 20, 21], [22, 23, 24], [19, 20, 21], [25, 26, 27], [19, 20, 21], [22, 23, 24], [19, 20, 21], [19, 20, 21]] [[11, 12, 13], [19, 20, 21], [19, 20, 21], [22, 23, 24], [19, 20, 21], [25, 26, 27], [19, 20, 21], [22, 23, 24], [19, 20, 21], [19, 20, 21]] [[1, 2, 3], [4, 5, 6], [4, 5, 6], [10, 11, 12], [13, 14, 15], [10, 11, 12], [10, 11, 12], [19, 20, 21], [19, 20, 21], [37, 38, 39], [19, 20, 21], [22, 23, 24]] [[28, 29, 30], [31, 32, 33], [34, 35, 36], [28, 29, 30], [31, 32, 33], [34, 35, 36], [28, 29, 30], [31, 32, 33], [34, 35, 36], [28, 29, 30]] [[11, 12, 13], [31, 32, 33], [34, 35, 36], [28, 29, 30], [31, 32, 33], [34, 35, 36], [28, 29, 30], [31, 32, 33], [34, 35, 36], [28, 29, 30]] +# array_replace scalar arguments over multiple input rows +query ??? +select + array_replace(column1, 2, 9), + array_replace_n(column1, 2, 9, 2), + array_replace_all(column1, 2, 9) +from ( + values + (make_array(1, 2, 2, 3)), + (make_array(2, 4, 2)) +) as t(column1); +---- +[1, 9, 2, 3] [1, 9, 9, 3] [1, 9, 9, 3] +[9, 4, 2] [9, 4, 9] [9, 4, 9] + +# array_replace_n scalar max exceeding matches over multiple input rows +query ? +select array_replace_n(column1, 2, 9, 10) +from ( + values + (make_array(1, 2, 2, 3)), + (make_array(2, 4, 2)) +) as t(column1); +---- +[1, 9, 9, 3] +[9, 4, 9] + ## array_replace_n (aliases: `list_replace_n`) # array_replace_n scalar function #1 @@ -226,22 +284,35 @@ select ---- [1, 3, 3, 4] [1, 0, 0, 5, 4, 6, 7] [1, 2, 3] [1, 4, 4] [1, 4, 4] [0, 4, 0, 5] -query ???? +query ?????? select array_replace_n(arrow_cast(make_array(1, 2, 3, 4), 'LargeList(Int64)'), 2, 3, 2), array_replace_n(arrow_cast(make_array(1, 4, 4, 5, 4, 6, 7), 'LargeList(Int64)'), 4, 0, 2), array_replace_n(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)'), 4, 0, 3), - array_replace_n(arrow_cast(make_array(1, 4, 4), 'LargeList(Int64)'), 4, 0, 0); + array_replace_n(arrow_cast(make_array(1, 4, 4), 'LargeList(Int64)'), 4, 0, 0), + array_replace_n(arrow_cast(make_array(1, 4, 4), 'LargeList(Int64)'), 4, 0, -1), + array_replace_n(arrow_cast(make_array(1, 4, 1, 5), 'LargeList(Int64)'), 1, 0, 10); ---- -[1, 3, 3, 4] [1, 0, 0, 5, 4, 6, 7] [1, 2, 3] [1, 4, 4] +[1, 3, 3, 4] [1, 0, 0, 5, 4, 6, 7] [1, 2, 3] [1, 4, 4] [1, 4, 4] [0, 4, 0, 5] -query ??? +query ?????? select array_replace_n(arrow_cast(make_array(1, 2, 3, 4), 'FixedSizeList(4, Int64)'), 2, 3, 2), array_replace_n(arrow_cast(make_array(1, 4, 4, 5, 4, 6, 7), 'FixedSizeList(7, Int64)'), 4, 0, 2), - array_replace_n(arrow_cast(make_array(1, 2, 3), 'FixedSizeList(3, Int64)'), 4, 0, 3); + array_replace_n(arrow_cast(make_array(1, 2, 3), 'FixedSizeList(3, Int64)'), 4, 0, 3), + array_replace_n(arrow_cast(make_array(1, 4, 4), 'FixedSizeList(3, Int64)'), 4, 0, 0), + array_replace_n(arrow_cast(make_array(1, 4, 4), 'FixedSizeList(3, Int64)'), 4, 0, -1), + array_replace_n(arrow_cast(make_array(1, 4, 1, 5), 'FixedSizeList(4, Int64)'), 1, 0, 10); ---- -[1, 3, 3, 4] [1, 0, 0, 5, 4, 6, 7] [1, 2, 3] +[1, 3, 3, 4] [1, 0, 0, 5, 4, 6, 7] [1, 2, 3] [1, 4, 4] [1, 4, 4] [0, 4, 0, 5] + +# array_replace_n scalar max exceeding matches for empty arrays +query ?? +select + array_replace_n(arrow_cast(make_array(), 'List(Int64)'), 2, 9, 10), + array_replace_n(arrow_cast(make_array(), 'LargeList(Int64)'), 2, 9, 10); +---- +[] [] # array_replace_n scalar function #2 (element is list) query ?? @@ -316,12 +387,29 @@ select query ? select array_replace_n(make_array(1, 2, 3, 4, 5), NULL, NULL, NULL); ---- -[1, 2, 3, 4, 5] +NULL query ? select array_replace_n(arrow_cast(make_array(1, 2, 3, 4, 5), 'LargeList(Int64)'), NULL, NULL, NULL); ---- -[1, 2, 3, 4, 5] +NULL + +query ?? +select + array_replace_n(make_array(1, 2, 2), 2, 9, NULL), + array_replace_n(arrow_cast(make_array(1, 2, 2), 'LargeList(Int64)'), 2, 9, NULL); +---- +NULL NULL + +# array_replace_n with null max from column +query ? +select array_replace_n(column1, column2, column3, column4) from (values + (make_array(1, 2, 2), 2, 9, 2), + (make_array(3, 4, 4), 4, 8, null) +) as t(column1, column2, column3, column4); +---- +[1, 9, 9] +NULL # array_replace_n scalar function with columns #1 query ? @@ -657,6 +745,101 @@ select column1, column2, column3, column4, array_replace_n(column1, column2, col NULL 3 2 1 NULL [3, 1, 3] 3 NULL 1 [NULL, 1, 3] +query ??? +select + array_replace(make_array(3, NULL, NULL), NULL, 5), + array_replace_n(make_array(3, NULL, NULL), NULL, 5, 10), + array_replace_all(make_array(3, NULL, NULL), NULL, 5); +---- +[3, 5, NULL] [3, 5, 5] [3, 5, 5] + +# the input list's inner field is carried through to the output, so that the +# returned type matches the one promised at planning time +query ???TTT +select + array_replace(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9), + array_replace_n(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9, 1), + array_replace_all(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9), + arrow_typeof(array_replace(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9)), + arrow_typeof(array_replace_n(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9, 1)), + arrow_typeof(array_replace_all(arrow_cast(column1, 'List(Int64, field: ''element'')'), 2, 9)) +from values (make_array(1, 2, 2)); +---- +[1, 9, 2] [1, 9, 2] [1, 9, 9] List(Int64, field: 'element') List(Int64, field: 'element') List(Int64, field: 'element') + +query ???TTT +select + array_replace(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9), + array_replace_n(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9, 1), + array_replace_all(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9), + arrow_typeof(array_replace(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9)), + arrow_typeof(array_replace_n(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9, 1)), + arrow_typeof(array_replace_all(arrow_cast(column1, 'LargeList(Int64, field: ''element'')'), 2, 9)) +from values (make_array(1, 2, 2)); +---- +[1, 9, 2] [1, 9, 2] [1, 9, 9] LargeList(Int64, field: 'element') LargeList(Int64, field: 'element') LargeList(Int64, field: 'element') + +# nested from/to values fall back to the generic comparison path +query ?T +select + array_replace(arrow_cast(column1, 'List(List(Int64), field: ''element'')'), make_array(2), make_array(9)), + arrow_typeof(array_replace(arrow_cast(column1, 'List(List(Int64), field: ''element'')'), make_array(2), make_array(9))) +from values (make_array(make_array(1), make_array(2))); +---- +[[1], [9]] List(List(Int64), field: 'element') + +# a non-nullable inner field stays non-nullable when the replacement cannot be null +query ???TTT +select + array_replace(column1, 2, 9), + array_replace_n(column1, 2, 9, 1), + array_replace_all(arrow_cast(column1, 'LargeList(non-null Int64)'), 2, 9), + arrow_typeof(array_replace(column1, 2, 9)), + arrow_typeof(array_replace_n(column1, 2, 9, 1)), + arrow_typeof(array_replace_all(arrow_cast(column1, 'LargeList(non-null Int64)'), 2, 9)) +from values + (arrow_cast(make_array(), 'List(non-null Int64)')), + (arrow_cast(NULL, 'List(non-null Int64)')), + (arrow_cast(make_array(1, 2, 2), 'List(non-null Int64)')) +; +---- +[] [] [] List(non-null Int64) List(non-null Int64) LargeList(non-null Int64) +NULL NULL NULL List(non-null Int64) List(non-null Int64) LargeList(non-null Int64) +[1, 9, 2] [1, 9, 2] [1, 9, 9] List(non-null Int64) List(non-null Int64) LargeList(non-null Int64) + +# ... and is widened when the replacement is nullable, since the result +# genuinely contains a null element +query ???TTT +select + array_replace(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL), + array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL, 1), + array_replace_all(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL), + arrow_typeof(array_replace(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL)), + arrow_typeof(array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL, 1)), + arrow_typeof(array_replace_all(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL)) +from values (make_array(1, 2, 2)); +---- +[1, NULL, 2] [1, NULL, 2] [1, NULL, NULL] List(Int64) List(Int64) List(Int64) + +# a max of 0 short circuits without replacing anything, but must still return +# the promised (widened) type +query ?T +select + array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL, 0), + arrow_typeof(array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, NULL, 0)) +from values (make_array(1, 2, 2)); +---- +[1, 2, 2] List(Int64) + +# a NULL max yields a NULL row of the promised type +query ?T +select + array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, 9, NULL), + arrow_typeof(array_replace_n(arrow_cast(column1, 'List(non-null Int64)'), 2, 9, NULL)) +from values (make_array(1, 2, 2)); +---- +NULL List(non-null Int64) + statement ok diff --git a/datafusion/sqllogictest/test_files/array/array_resize.slt b/datafusion/sqllogictest/test_files/array/array_resize.slt index 91febb76ac00e..37f8f2c6935c6 100644 --- a/datafusion/sqllogictest/test_files/array/array_resize.slt +++ b/datafusion/sqllogictest/test_files/array/array_resize.slt @@ -64,6 +64,23 @@ select array_resize(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)'), 5, 4); query error select array_resize(make_array(1, 2, 3), -5, 2); +# array_resize with a very large size should error instead of panicking (capacity overflow) +query error DataFusion error: Execution error: array_resize: resulting array of 9223372036854775807 elements exceeds the maximum array size +select array_resize(make_array(1), 9223372036854775807, 0); + +query error DataFusion error: Execution error: array_resize: resulting array of 9223372036854775807 elements exceeds the maximum array size +select array_resize(arrow_cast(make_array(1), 'LargeList(Int64)'), 9223372036854775807, 0); + +# List size above i32::MAX must error via the offset-type guard instead of +# silently truncating the offsets or attempting a multi-GB allocation +query error DataFusion error: Execution error: array_resize: resulting array of 3000000000 elements exceeds the maximum array size +select array_resize(make_array(1), 3000000000, 0); + +# Non-primitive element types (e.g. Utf8) fall back to a conservative byte +# width; a size above that cap must error instead of attempting a huge allocation. +query error DataFusion error: Execution error: array_resize: resulting array of 600000000000000000 elements exceeds the maximum array size +select array_resize(arrow_cast(make_array('a'), 'LargeList(Utf8)'), 600000000000000000); + # array_resize scalar function #5 query ? select array_resize(make_array(1.1, 2.2, 3.3), 10, 9.9); @@ -75,6 +92,12 @@ select array_resize(arrow_cast(make_array(1.1, 2.2, 3.3), 'LargeList(Float64)'), ---- [1.1, 2.2, 3.3, 9.9, 9.9, 9.9, 9.9, 9.9, 9.9, 9.9] +# array_resize null size +query ?? +select array_resize(make_array(1, 2, 3), NULL), array_resize(arrow_cast(make_array(1, 2, 3), 'LargeList(Int64)'), NULL); +---- +NULL NULL + # array_resize scalar function #5 query ? select array_resize(column1, column2, column3) from arrays_values; @@ -84,7 +107,7 @@ select array_resize(column1, column2, column3) from arrays_values; [21, 22, 23, NULL, 25, 26, 27, 28, 29, 30, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3] [31, 32, 33, 34, 35, NULL, 37, 38, 39, 40, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4] NULL -[] +NULL [51, 52, NULL, 54, 55, 56, 57, 58, 59, 60, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL] [61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7] @@ -96,7 +119,7 @@ select array_resize(arrow_cast(column1, 'LargeList(Int64)'), column2, column3) f [21, 22, 23, NULL, 25, 26, 27, 28, 29, 30, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3] [31, 32, 33, 34, 35, NULL, 37, 38, 39, 40, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4, 4] NULL -[] +NULL [51, 52, NULL, 54, 55, 56, 57, 58, 59, 60, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL] [61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7, 7] @@ -139,7 +162,7 @@ select array_resize(column1, column2, column3) from array_resize_values; [21, 22, 23, 24, NULL, 26, 27, 28] [31, 32, 33, 34, 35, 36, NULL, 38, 39, 40, 4, 4] NULL -[] +NULL [51, 52, 53, 54, 55, NULL, 57, 58, 59, 60, NULL, NULL, NULL] [61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 7, 7, 7, 7, 7] @@ -152,7 +175,7 @@ select array_resize(arrow_cast(column1, 'LargeList(Int64)'), column2, column3) f [21, 22, 23, 24, NULL, 26, 27, 28] [31, 32, 33, 34, 35, 36, NULL, 38, 39, 40, 4, 4] NULL -[] +NULL [51, 52, 53, 54, 55, NULL, 57, 58, 59, 60, NULL, NULL, NULL] [61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 7, 7, 7, 7, 7] @@ -165,7 +188,7 @@ select array_resize(column1, column2, 9) from array_resize_values; [21, 22, 23, 24, NULL, 26, 27, 28] [31, 32, 33, 34, 35, 36, NULL, 38, 39, 40, 9, 9] NULL -[] +NULL [51, 52, 53, 54, 55, NULL, 57, 58, 59, 60, 9, 9, 9] [61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 9, 9, 9, 9, 9] diff --git a/datafusion/sqllogictest/test_files/array/array_slice.slt b/datafusion/sqllogictest/test_files/array/array_slice.slt index 14587a50b2266..76b81b28efc58 100644 --- a/datafusion/sqllogictest/test_files/array/array_slice.slt +++ b/datafusion/sqllogictest/test_files/array/array_slice.slt @@ -450,6 +450,37 @@ NULL NULL [1, 3, 5] +# maintains inner nullability +query ?T +select array_slice(column1, 2, 3), arrow_typeof(array_slice(column1, 2, 3)) +from values + (arrow_cast([], 'List(non-null Int32)')), + (arrow_cast(NULL, 'List(non-null Int32)')), + (arrow_cast([1, 3, 5, -5], 'List(non-null Int32)')) +; +---- +[] List(non-null Int32) +NULL List(non-null Int32) +[3, 5] List(non-null Int32) + +query ?T +select array_slice(column1, 2, 3), arrow_typeof(array_slice(column1, 2, 3)) +from values + (arrow_cast([], 'LargeList(non-null Int32)')), + (arrow_cast(NULL, 'LargeList(non-null Int32)')), + (arrow_cast([1, 3, 5, -5], 'LargeList(non-null Int32)')) +; +---- +[] LargeList(non-null Int32) +NULL LargeList(non-null Int32) +[3, 5] LargeList(non-null Int32) + +query ?T +select array_slice(column1, 2, 3, 2), arrow_typeof(array_slice(column1, 2, 3, 2)) +from values (arrow_cast([1, 3, 5, -5], 'List(non-null Int32)')); +---- +[3] List(non-null Int32) + # Testing with empty arguments should result in an error query error DataFusion error: Error during planning: 'array_slice' does not support zero arguments select array_slice(); diff --git a/datafusion/sqllogictest/test_files/array/array_transform.slt b/datafusion/sqllogictest/test_files/array/array_transform.slt index c8c43588c882c..5439d7441155b 100644 --- a/datafusion/sqllogictest/test_files/array/array_transform.slt +++ b/datafusion/sqllogictest/test_files/array/array_transform.slt @@ -393,6 +393,12 @@ physical_plan 02)--ProjectionExec: expr=[text@0 as text, list@1 as list, number@2 as number, CASE WHEN number@2 > 30 THEN array_transform(make_array(make_array(list@1)), (list) -> array_transform(list@3, (list) -> array_transform(list@4, (v) -> number@2 + v@5 + array_element(list@4, 1)))) ELSE array_transform(make_array(make_array(list@1)), (list) -> array_transform(list@3, (list) -> array_transform(list@4, (v) -> number@2 + array_element(list@4, 1)))) END as CASE WHEN t.number > Int64(30) THEN array_transform(make_array(make_array(t.list)),(list) -> array_transform(list,(list) -> array_transform(list,(v) -> t.number + v + list[Int64(1)]))) ELSE array_transform(make_array(make_array(t.list)),(list) -> array_transform(list,(list) -> array_transform(list,(v) -> t.number + list[Int64(1)]))) END] 03)----DataSourceExec: partitions=1, partition_sizes=[1] +# null arg +query ? +SELECT array_transform(NULL, x -> x * 2); +---- +NULL + query error select array_transform(); ---- diff --git a/datafusion/sqllogictest/test_files/array/array_union.slt b/datafusion/sqllogictest/test_files/array/array_union.slt index 6a0fdc546e7d7..edb90705af940 100644 --- a/datafusion/sqllogictest/test_files/array/array_union.slt +++ b/datafusion/sqllogictest/test_files/array/array_union.slt @@ -236,4 +236,64 @@ select array_except([1, 2], arrow_cast(null, 'List(Int64)')); NULL +# Negative zero (-0.0) and positive zero (+0.0) compare equal under +# IEEE 754, but their bit patterns differ. array_union and array_intersect +# must normalize the sign for dedup / matching; the canonical +# representative is +0.0. PostgreSQL / IEEE 754 expected output below. + +# array_union with +0.0 / -0.0 +query ? +select array_union([0.0], [-0.0]); +---- +[0.0] + +query ? +select array_union([0.0, 1.0], [-0.0]); +---- +[0.0, 1.0] + +query ? +select array_union([0.0, -0.0, 1.0], [-0.0, 1.0]); +---- +[0.0, 1.0] + +# Float32 list. +query ? +select array_union(arrow_cast([0.0], 'List(Float32)'), arrow_cast([-0.0], 'List(Float32)')); +---- +[0.0] + +# LargeList(Float64). +query ? +select array_union(arrow_cast([0.0], 'LargeList(Float64)'), arrow_cast([-0.0], 'LargeList(Float64)')); +---- +[0.0] + + +# array_intersect with +0.0 / -0.0 +# +0.0 in lhs matches -0.0 in rhs. +query ? +select array_intersect([0.0, 1.0], [-0.0]); +---- +[0.0] + +# Either +0.0 or -0.0 in lhs matches +0.0 in rhs (canonicalized to +0.0). +query ? +select array_intersect([0.0, -0.0], [0.0]); +---- +[0.0] + +# Same with -0.0 in rhs. +query ? +select array_intersect([0.0, -0.0], [-0.0]); +---- +[0.0] + +# Float32 list. +query ? +select array_intersect(arrow_cast([0.0], 'List(Float32)'), arrow_cast([-0.0], 'List(Float32)')); +---- +[0.0] + + include ./cleanup.slt.part diff --git a/datafusion/sqllogictest/test_files/array/cardinality.slt b/datafusion/sqllogictest/test_files/array/cardinality.slt index 52b1a2b5445d9..21e94b53b2768 100644 --- a/datafusion/sqllogictest/test_files/array/cardinality.slt +++ b/datafusion/sqllogictest/test_files/array/cardinality.slt @@ -51,6 +51,37 @@ select cardinality(arrow_cast([[1, 2], [3, 4], [5, 6]], 'FixedSizeList(3, List(I ---- 6 +# cardinality counts actual leaf elements in ragged nested arrays +query III +select cardinality([[1], [2, 3]]), + cardinality([[1, 2, 3], []]), + cardinality([[], [1, 2]]); +---- +3 3 2 + +query IIII +select cardinality(arrow_cast([[1], [2, 3]], 'ListView(List(Int64))')), + cardinality(arrow_cast([[1], [2, 3]], 'LargeListView(List(Int64))')), + cardinality(arrow_cast([[1, 2], [3, 4]], 'List(FixedSizeList(2, Int64))')), + cardinality(arrow_cast([[1], [2, 3]], 'LargeList(List(Int64))')); +---- +3 3 4 3 + +query III +select cardinality(arrow_cast([[[1]], [[2, 3], []]], 'List(ListView(List(Int64)))')), + cardinality(arrow_cast([[[1]], [[2, 3], []]], 'List(LargeListView(List(Int64)))')), + cardinality(arrow_cast([[[1, 2]], [[3, 4]]], 'List(List(FixedSizeList(2, Int64)))')); +---- +3 3 4 + +query IIII +select cardinality([[NULL], [1, 2]]), + cardinality([[[1]], [[2, 3], []]]), + cardinality(arrow_cast([[], [1, 2]], 'LargeList(List(Int64))')), + cardinality(make_array(NULL::int[], [1, 2])); +---- +3 3 2 2 + # cardinality scalar function #3 query II select cardinality(make_array()), cardinality(make_array(make_array())) @@ -67,12 +98,10 @@ select cardinality(arrow_cast(make_array(), 'LargeList(Int64)')), cardinality(ar ---- 0 0 -#TODO -#https://github.com/apache/datafusion/issues/9158 -#query II -#select cardinality(arrow_cast(make_array(), 'FixedSizeList(1, Null)')), cardinality(arrow_cast(make_array(make_array()), 'FixedSizeList(1, List(Int64))')) -#---- -#NULL 0 +query II +select cardinality(arrow_cast(make_array(null), 'FixedSizeList(1, Null)')), cardinality(arrow_cast(make_array(make_array()), 'FixedSizeList(1, List(Int64))')) +---- +1 0 # cardinality of NULL arrays should return NULL query II diff --git a/datafusion/sqllogictest/test_files/array/cleanup.slt.part b/datafusion/sqllogictest/test_files/array/cleanup.slt.part index a11a4770ec058..eff5d17acf37f 100644 --- a/datafusion/sqllogictest/test_files/array/cleanup.slt.part +++ b/datafusion/sqllogictest/test_files/array/cleanup.slt.part @@ -167,4 +167,3 @@ drop table large_arrays_values_without_nulls; statement ok drop table fixed_size_arrays_values_without_nulls; - diff --git a/datafusion/sqllogictest/test_files/array/init_data.slt.part b/datafusion/sqllogictest/test_files/array/init_data.slt.part index f5cc58fb2be58..bb8d76809f816 100644 --- a/datafusion/sqllogictest/test_files/array/init_data.slt.part +++ b/datafusion/sqllogictest/test_files/array/init_data.slt.part @@ -689,4 +689,3 @@ AS FROM arrays_distance_table ; - diff --git a/datafusion/sqllogictest/test_files/array_add.slt b/datafusion/sqllogictest/test_files/array_add.slt new file mode 100644 index 0000000000000..e13f6acd269cb --- /dev/null +++ b/datafusion/sqllogictest/test_files/array_add.slt @@ -0,0 +1,237 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +## array_add + +# Basic element-wise sum +query ? +select array_add([1.0, 2.0, 3.0], [10.0, 20.0, 30.0]); +---- +[11.0, 22.0, 33.0] + +# Negative components +query ? +select array_add([1.0, -2.0, 3.0], [-1.0, 2.0, -3.0]); +---- +[0.0, 0.0, 0.0] + +# Single-element arrays +query ? +select array_add([5.0], [7.0]); +---- +[12.0] + +# Bare NULL on left -> NULL row +query ? +select array_add(NULL, [1.0, 2.0]); +---- +NULL + +# Bare NULL on right -> NULL row +query ? +select array_add([1.0, 2.0], NULL); +---- +NULL + +# Both bare NULL -> NULL row +query ? +select array_add(NULL, NULL); +---- +NULL + +# NULL element on left propagates to that position only +query ? +select array_add([1.0, NULL, 3.0], [10.0, 20.0, 30.0]); +---- +[11.0, NULL, 33.0] + +# NULL element on right propagates to that position only +query ? +select array_add([1.0, 2.0, 3.0], [10.0, NULL, 30.0]); +---- +[11.0, NULL, 33.0] + +# NULL element on both sides at the same position +query ? +select array_add([1.0, NULL, 3.0], [10.0, NULL, 30.0]); +---- +[11.0, NULL, 33.0] + +# NULL elements at different positions both propagate +query ? +select array_add([1.0, NULL, 3.0], [NULL, 20.0, 30.0]); +---- +[NULL, NULL, 33.0] + +# Length mismatch is an exec error +query error array_add requires both list inputs to have the same length per row +select array_add([1.0, 2.0], [10.0, 20.0, 30.0]); + +# Empty arrays on both sides return empty array +query ? +select array_add(arrow_cast(make_array(), 'List(Float64)'), arrow_cast(make_array(), 'List(Float64)')); +---- +[] + +# Integer literals coerced to Float64 +query ? +select array_add([1, 2, 3], [10, 20, 30]); +---- +[11.0, 22.0, 33.0] + +# Mixed int + float literals coerced to Float64 +query ? +select array_add([1, 2, 3], [0.5, 0.5, 0.5]); +---- +[1.5, 2.5, 3.5] + +# LargeList input on both sides +query ? +select array_add( + arrow_cast([1.0, 2.0, 3.0], 'LargeList(Float64)'), + arrow_cast([10.0, 20.0, 30.0], 'LargeList(Float64)') +); +---- +[11.0, 22.0, 33.0] + +# Mixed List + LargeList -> both widened to LargeList +query ? +select array_add( + [1.0, 2.0, 3.0], + arrow_cast([10.0, 20.0, 30.0], 'LargeList(Float64)') +); +---- +[11.0, 22.0, 33.0] + +# FixedSizeList input (coerced to List) +query ? +select array_add( + arrow_cast([1.0, 2.0, 3.0], 'FixedSizeList(3, Float64)'), + arrow_cast([10.0, 20.0, 30.0], 'FixedSizeList(3, Float64)') +); +---- +[11.0, 22.0, 33.0] + +# Float32 inner type on one side +query ? +select array_add( + arrow_cast([1.0, 2.0, 3.0], 'List(Float32)'), + [10.0, 20.0, 30.0] +); +---- +[11.0, 22.0, 33.0] + +# Int64 inner type +query ? +select array_add( + arrow_cast([1, 2, 3], 'List(Int64)'), + arrow_cast([10, 20, 30], 'List(Int64)') +); +---- +[11.0, 22.0, 33.0] + +# Unsupported non-list input (plan error) +query error array_add does not support type +select array_add(1, [1.0, 2.0]); + +# Wrong arg count +query error array_add function requires 2 arguments, got 0 +select array_add(); + +query error array_add function requires 2 arguments, got 1 +select array_add([1.0, 2.0]); + +# Return type matches input variant +query ?T +select array_add([1.0, 2.0], [3.0, 4.0]), arrow_typeof(array_add([1.0, 2.0], [3.0, 4.0])); +---- +[4.0, 6.0] List(Float64) + +# Multi-row query: normal row, NULL row, element-NULL row, length-matched row +query ? +select array_add(a, b) from (values + (make_array(1.0, 2.0, 3.0), make_array(10.0, 20.0, 30.0)), + (NULL, make_array(1.0, 2.0, 3.0)), + (make_array(1.0, 2.0, 3.0), NULL), + (make_array(1.0, NULL, 3.0), make_array(10.0, 20.0, 30.0)) +) as t(a, b); +---- +[11.0, 22.0, 33.0] +NULL +NULL +[11.0, NULL, 33.0] + +# list_add alias +query ? +select list_add([1.0, 2.0], [3.0, 4.0]); +---- +[4.0, 6.0] + +# list_add alias multi-row +query ? +select list_add(a, b) from (values + (make_array(1.0, 2.0), make_array(10.0, 20.0)), + (NULL, make_array(1.0, 2.0)) +) as t(a, b); +---- +[11.0, 22.0] +NULL + +# Decimal element types are coerced to Float64 (lossy) like other array-math UDFs +query ? +select array_add( + arrow_cast([1, 2, 3], 'List(Decimal128(10, 2))'), + arrow_cast([10, 20, 30], 'List(Decimal128(10, 2))') +); +---- +[11.0, 22.0, 33.0] + +# Explicit cast to DOUBLE works as the documented opt-in +query ? +select array_add( + arrow_cast(arrow_cast([1, 2, 3], 'List(Decimal128(10, 2))'), 'List(Float64)'), + [10.0, 20.0, 30.0] +); +---- +[11.0, 22.0, 33.0] + +# Chained array_add: result of inner call feeds the outer call +query ? +select array_add(array_add([1.0, 2.0, 3.0], [10.0, 20.0, 30.0]), [100.0, 200.0, 300.0]); +---- +[111.0, 222.0, 333.0] + +# Chained array_add propagates element-level NULLs through both layers +query ? +select array_add( + array_add([1.0, NULL, 3.0], [10.0, 20.0, 30.0]), + [100.0, 200.0, NULL] +); +---- +[111.0, NULL, NULL] + +# Chained array_add over multiple rows +query ? +select array_add(array_add(a, b), c) from (values + (make_array(1.0, 2.0), make_array(10.0, 20.0), make_array(100.0, 200.0)), + (NULL, make_array(1.0, 2.0), make_array(3.0, 4.0)), + (make_array(1.0, 2.0), make_array(10.0, NULL), make_array(100.0, 200.0)) +) as t(a, b, c); +---- +[111.0, 222.0] +NULL +[111.0, NULL] \ No newline at end of file diff --git a/datafusion/sqllogictest/test_files/array_agg.slt b/datafusion/sqllogictest/test_files/array_agg.slt new file mode 100644 index 0000000000000..d5aaf8cab17c1 --- /dev/null +++ b/datafusion/sqllogictest/test_files/array_agg.slt @@ -0,0 +1,620 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +####### +# Tests for the array_agg aggregate function. +# +# Sliding (bounded) window frames, which exercise `retract_batch`, live in +# `array_agg_sliding_window.slt`. +####### + +####### +# Setup test data table +####### +statement ok +CREATE EXTERNAL TABLE aggregate_test_100 ( + c1 VARCHAR NOT NULL, + c2 TINYINT NOT NULL, + c3 SMALLINT NOT NULL, + c4 SMALLINT, + c5 INT, + c6 BIGINT NOT NULL, + c7 SMALLINT NOT NULL, + c8 INT NOT NULL, + c9 INT UNSIGNED NOT NULL, + c10 BIGINT UNSIGNED NOT NULL, + c11 FLOAT NOT NULL, + c12 DOUBLE NOT NULL, + c13 VARCHAR NOT NULL, + c14 DATE NOT NULL, + c15 TIMESTAMP NOT NULL, +) +STORED AS CSV +LOCATION '../../testing/data/csv/aggregate_test_100_with_dates.csv' +OPTIONS ('format.has_header' 'true'); + +####### +# Basic array_agg +####### + +# csv_query_array_agg +query ? +SELECT array_agg(c13) FROM (SELECT * FROM aggregate_test_100 ORDER BY c13 LIMIT 2) test +---- +[0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB] + +# csv_query_array_agg_empty +query ? +SELECT array_agg(c13) FROM (SELECT * FROM aggregate_test_100 LIMIT 0) test +---- +NULL + +# csv_query_array_agg_one +query ? +SELECT array_agg(c13) FROM (SELECT * FROM aggregate_test_100 ORDER BY c13 LIMIT 1) test +---- +[0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm] + +# array_agg_zero +query ? +SELECT ARRAY_AGG([]) +---- +[[]] + +# array_agg_one +query ? +SELECT ARRAY_AGG([1]) +---- +[[1]] + +# test array_agg with no row qualified +statement ok +create table t(a int, b float, c bigint) as values (1, 1.2, 2); + +# returns NULL, follows DuckDB's behaviour +query ? +select array_agg(a) from t where a > 2; +---- +NULL + +query ? +select array_agg(b) from t where b > 3.1; +---- +NULL + +query ? +select array_agg(c) from t where c > 3; +---- +NULL + +query ?I +select array_agg(c), count(1) from t where c > 3; +---- +NULL 0 + +# returns 0 rows if group by is applied, follows DuckDB's behaviour +query ? +select array_agg(a) from t where a > 3 group by a; +---- + +query ?I +select array_agg(a), count(1) from t where a > 3 group by a; +---- + +# returns NULL, follows DuckDB's behaviour +query ? +select array_agg(distinct a) from t where a > 3; +---- +NULL + +query ?I +select array_agg(distinct a), count(1) from t where a > 3; +---- +NULL 0 + +# returns 0 rows if group by is applied, follows DuckDB's behaviour +query ? +select array_agg(distinct a) from t where a > 3 group by a; +---- + +query ?I +select array_agg(distinct a), count(1) from t where a > 3 group by a; +---- + +# test order sensitive array agg +query ? +select array_agg(a order by a) from t where a > 3; +---- +NULL + +query ? +select array_agg(a order by a) from t where a > 3 group by a; +---- + +query ?I +select array_agg(a order by a), count(1) from t where a > 3 group by a; +---- + +statement ok +drop table t; + +# test with no values +statement ok +create table t(a int, b float, c bigint); + +query ? +select array_agg(a) from t; +---- +NULL + +query ? +select array_agg(b) from t; +---- +NULL + +query ? +select array_agg(c) from t; +---- +NULL + +query ?I +select array_agg(distinct a), count(1) from t; +---- +NULL 0 + +query ?I +select array_agg(distinct b), count(1) from t; +---- +NULL 0 + +query ?I +select array_agg(distinct b), count(1) from t; +---- +NULL 0 + +statement ok +drop table t; + + +# array_agg_i32 +statement ok +create table t (c1 int) as values (1), (2), (3), (4), (5); + +query ? +select array_agg(c1) from t; +---- +[1, 2, 3, 4, 5] + +statement ok +drop table t; + +# array_agg_nested +statement ok +create table t as values (make_array([1, 2, 3], [4, 5])), (make_array([6], [7, 8])), (make_array([9])); + +query ? +select array_agg(column1) from t; +---- +[[[1, 2, 3], [4, 5]], [[6], [7, 8]], [[9]]] + +statement ok +drop table t; + +# array_agg_ignore_nulls +statement ok +create table t as values (NULL, ''), (1, 'c'), (2, 'a'), (NULL, 'b'), (4, NULL), (NULL, NULL), (5, 'a'); + +query ? +select array_agg(column1) ignore nulls as c1 from t; +---- +[1, 2, 4, 5] + +query II +select count(*), array_length(array_agg(distinct column2) ignore nulls) from t; +---- +7 4 + +query ? +select array_agg(column2 order by column1) ignore nulls from t; +---- +[c, a, a, , b] + +query ? +select array_agg(DISTINCT column2 order by column2) ignore nulls from t; +---- +[, a, b, c] + +statement ok +drop table t; + +####### +# array_agg with ORDER BY +####### + +# array agg can use order by +query ? +SELECT array_agg(c13 ORDER BY c13) +FROM + (SELECT * + FROM aggregate_test_100 + ORDER BY c13 + LIMIT 5) as t1 +---- +[0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8] + +# array agg can use order by with distinct +query ? +SELECT array_agg(DISTINCT c13 ORDER BY c13) +FROM + (SELECT * + FROM aggregate_test_100 + ORDER BY c13 + LIMIT 5) as t1 +---- +[0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8] + +query error Execution error: In an aggregate with DISTINCT, ORDER BY expressions must appear in argument list +SELECT array_agg(DISTINCT c13 ORDER BY c12) +FROM aggregate_test_100 + +query error Execution error: In an aggregate with DISTINCT, ORDER BY expressions must appear in argument list +SELECT array_agg(DISTINCT c13 ORDER BY c13, c12) +FROM aggregate_test_100 + +query ?? rowsort +with tbl as (SELECT * FROM (VALUES ('xxx', 'yyy'), ('xxx', 'yyy'), ('xxx2', 'yyy2')) AS t(x, y)) +select + array_agg(x order by x) as x_agg, + array_agg(y order by y) as y_agg +from tbl +group by all +---- +[xxx, xxx, xxx2] [yyy, yyy, yyy2] + +query ?? +SELECT + (SELECT array_agg(c12 ORDER BY c12) FROM aggregate_test_100), + (SELECT array_agg(c13 ORDER BY c13) FROM aggregate_test_100) +---- +[0.01479305307777301, 0.02182578039211991, 0.03968347085780355, 0.04429073092078406, 0.047343434291126085, 0.04893135681998029, 0.0494924465469434, 0.05573662213439634, 0.05636955101974106, 0.061029375346466685, 0.07260475960924484, 0.09465635123783445, 0.12357539988406441, 0.152498292971736, 0.16301110515739792, 0.1640882545084913, 0.1754261586710173, 0.17592486905979987, 0.17909035118828576, 0.18628859265874176, 0.19113293583306745, 0.2145232647388039, 0.21535402343780985, 0.24899794314659673, 0.2537253407987472, 0.2667177795079635, 0.27159190516490006, 0.2739938529235548, 0.28534428578703896, 0.2944158618048994, 0.296036538664718, 0.3051364088814128, 0.30585375151301186, 0.3114712539863804, 0.3231750610081745, 0.32869374687050157, 0.33639590659276175, 0.3600766362333053, 0.36936304600612724, 0.38870280983958583, 0.39144436569161134, 0.40342283197779727, 0.4094218353587008, 0.40975383525297016, 0.42073125331890115, 0.4273123318932347, 0.42950521730777025, 0.4830878559436823, 0.5081765563442366, 0.5437595540422571, 0.5590205548347534, 0.5593249815276734, 0.5603062368164834, 0.560333188635217, 0.5614503754617461, 0.565352842229935, 0.574210838214554, 0.5759450483859969, 0.5773498217058918, 0.5991138115095911, 0.6009475544728957, 0.6108938307533, 0.6316565296547284, 0.6404495093354053, 0.6405262429561641, 0.6425694115212065, 0.658671129040488, 0.6668423897406515, 0.6864391962767343, 0.7035635283169166, 0.7325106678655877, 0.7328050041291218, 0.7614304100703713, 0.7631239070049998, 0.7670021786149205, 0.7697753383420857, 0.7764360990307122, 0.7784918983501654, 0.7973920072996036, 0.819715865079681, 0.8506721053047003, 0.8813167497816289, 0.8824879447595726, 0.9185813970744787, 0.9231889896940375, 0.9237877978193884, 0.9255031346434324, 0.9293883502480845, 0.9294097332465232, 0.9463098243875633, 0.946325164889271, 0.9491397432856566, 0.9567595541247681, 0.9706712283358269, 0.9723580396501548, 0.9748360509016578, 0.9800193410444061, 0.980809631269599, 0.991517828651004, 0.9965400387585364] [0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8, 3BEOHQsMEFZ58VcNTOJYShTBpAPzbt, 4HX6feIvmNXBN7XGqgO4YVBkhu8GDI, 4JznSdBajNWhu4hRQwjV1FjTTxY68i, 52mKlRE3aHCBZtjECq6sY9OqVf8Dze, 56MZa5O1hVtX4c5sbnCfxuX5kDChqI, 6FPJlLAcaQ5uokyOWZ9HGdLZObFvOZ, 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW, 6oIXZuIPIqEoPBvFmbt2Nxy3tryGUE, 6x93sxYioWuq5c9Kkk8oTAAORM7cH0, 802bgTGl6Bk5TlkPYYTxp5JkKyaYUA, 8LIh0b6jmDGm87BmIyjdxNIpX4ugjD, 90gAtmGEeIqUTbo1ZrxCvWtsseukXC, 9UbObCsVkmYpJGcGrgfK90qOnwb2Lj, AFGCj7OWlEB5QfniEFgonMq90Tq5uH, ALuRhobVWbnQTTWZdSOk0iVe8oYFhW, Amn2K87Db5Es3dFQO9cw9cvpAM6h35, AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz, BJqx5WokrmrrezZA0dUbleMYkG5U2O, BPtQMxnuSPpxMExYV9YkDa6cAN7GP3, BsM5ZAYifRh5Lw3Y8X1r53I0cTJnfE, C2GT5KVyOPZpgKVl110TyZO0NcJ434, DuJNG8tufSqW0ZstHqWj3aGvFLMg4A, EcCuckwsF3gV1Ecgmh5v4KM8g1ozif, ErJFw6hzZ5fmI5r8bhE4JzlscnhKZU, F7NSTjWvQJyBburN7CXRUlbgp2dIrA, Fi4rJeTQq4eXj8Lxg3Hja5hBVTVV5u, H5j5ZHy1FGesOAHjkQEDYCucbpKWRu, HKSMQ9nTnwXCJIte1JrM1dtYnDtJ8g, IWl0G3ZlMNf7WT8yjIB49cx7MmYOmr, IZTkHMLvIKuiLjhDjYMmIHxh166we4, Ig1QcuKsjHXkproePdERo2w0mYzIqd, JHNgc2UCaiXOdmkxwDDyGhRlO0mnBQ, JN0VclewmjwYlSl8386MlWv5rEhWCz, JafwVLSVk5AVoXFuzclesQ000EE2k1, KJFcmTVjdkCMv94wYCtfHMFhzyRsmH, Ktb7GQ0N1DrxwkCkEUsTaIXk0xYinn, Ld2ej8NEv5zNcqU60FwpHeZKBhfpiV, LiEBxds3X0Uw0lxiYjDqrkAaAwoiIW, MXhhH1Var3OzzJCtI9VNyYvA0q8UyJ, MeSTAXq8gVxVjbEjgkvU9YLte0X9uE, NEhyk8uIx4kEULJGa8qIyFjjBcP2G6, O66j6PaYuZhEUtqV6fuU7TyjM2WxC5, OF7fQ37GzaZ5ikA2oMyvleKtgnLjXh, OPwBqCEK5PWTjWaiOyL45u2NLTaDWv, Oq6J4Rx6nde0YlhOIJkFsX2MsSvAQ0, Ow5PGpfTm4dXCfTDsXAOTatXRoAydR, QEHVvcP8gxI6EMJIrvcnIhgzPNjIvv, QJYm7YRA3YetcBHI5wkMZeLXVmfuNy, QYlaIAnJA6r8rlAb6f59wcxvcPcWFf, RilTlL1tKkPOUFuzmLydHAVZwv1OGl, Sfx0vxv1skzZWT1PqVdoRDdO6Sb6xH, TTQUwpMNSXZqVBKAFvXu7OlWvKXJKX, TtDKUZxzVxsq758G6AWPSYuZgVgbcl, VDhtJkYjAYPykCgOU9x3v7v3t4SO1a, VY0zXmXeksCT8BzvpzpPLbmU9Kp9Y4, Vp3gmWunM5A7wOC9YW2JroFqTWjvTi, WHmjWk2AY4c6m7DA4GitUx6nmb1yYS, XemNcT1xp61xcM1Qz3wZ1VECCnq06O, Z2sWcQr0qyCJRMHDpRy3aQr7PkHtkK, aDxBtor7Icd9C5hnTvvw5NrIre740e, akiiY5N0I44CMwEnBL6RTBk7BRkxEj, b3b9esRhTzFEawbs6XhpKnD9ojutHB, bgK1r6v3BCTh0aejJUhkA1Hn6idXGp, cBGc0kSm32ylBDnxogG727C0uhZEYZ, cq4WSAIFwx3wwTUS5bp1wCe71R6U5I, dVdvo6nUD5FgCgsbOZLds28RyGTpnx, e2Gh6Ov8XkXoFdJWhl0EjwEHlMDYyG, f9ALCzwDAKmdu7Rk2msJaB1wxe5IBX, fuyvs0w7WsKSlXqJ1e6HFSoLmx03AG, gTpyQnEODMcpsPnJMZC66gh33i3m0b, gpo8K5qtYePve6jyPt6xgJx4YOVjms, gxfHWUF8XgY2KdFxigxvNEXe2V2XMl, i6RQVXKUh7MzuGMDaNclUYnFUAireU, ioEncce3mPOXD2hWhpZpCPWGATG6GU, jQimhdepw3GKmioWUlVSWeBVRKFkY3, l7uwDoTepWwnAP0ufqtHJS3CRi7RfP, lqhzgLsXZ8JhtpeeUWWNbMz8PHI705, m6jD0LBIQWaMfenwRCTANI9eOdyyto, mhjME0zBHbrK6NMkytMTQzOssOa1gF, mzbkwXKrPeZnxg2Kn1LRF5hYSsmksS, nYVJnVicpGRqKZibHyBAmtmzBXAFfT, oHJMNvWuunsIMIWFnYG31RCfkOo2V7, oLZ21P2JEDooxV1pU31cIxQHEeeoLu, okOkcWflkNXIy4R8LzmySyY1EC3sYd, pLk3i59bZwd5KBZrI1FiweYTd5hteG, pTeu0WMjBRTaNRT15rLCuEh3tBJVc5, qnPOOmslCJaT45buUisMRnM0rc77EK, t6fQUjJejPcjc04wHvHTPe55S65B4V, ukOiFGGFnQJDHFgZxHMpvhD3zybF0M, ukyD7b0Efj7tNlFSRmzZ0IqkEzg2a8, waIGbOGl1PM6gnzZ4uuZt4E2yDWRHs, wwXqSGKLyBQyPkonlzBNYUJTCo4LRS, xipQ93429ksjNcXPX5326VSg1xJZcW, y7C453hRWd4E7ImjNDWlpexB8nUqjh, ydkwycaISlYSlEq3TlkS2m15I2pcp8] + +query ?? +SELECT + array_agg(c12 ORDER BY c12), + array_agg(c13 ORDER BY c13) +FROM aggregate_test_100 +---- +[0.01479305307777301, 0.02182578039211991, 0.03968347085780355, 0.04429073092078406, 0.047343434291126085, 0.04893135681998029, 0.0494924465469434, 0.05573662213439634, 0.05636955101974106, 0.061029375346466685, 0.07260475960924484, 0.09465635123783445, 0.12357539988406441, 0.152498292971736, 0.16301110515739792, 0.1640882545084913, 0.1754261586710173, 0.17592486905979987, 0.17909035118828576, 0.18628859265874176, 0.19113293583306745, 0.2145232647388039, 0.21535402343780985, 0.24899794314659673, 0.2537253407987472, 0.2667177795079635, 0.27159190516490006, 0.2739938529235548, 0.28534428578703896, 0.2944158618048994, 0.296036538664718, 0.3051364088814128, 0.30585375151301186, 0.3114712539863804, 0.3231750610081745, 0.32869374687050157, 0.33639590659276175, 0.3600766362333053, 0.36936304600612724, 0.38870280983958583, 0.39144436569161134, 0.40342283197779727, 0.4094218353587008, 0.40975383525297016, 0.42073125331890115, 0.4273123318932347, 0.42950521730777025, 0.4830878559436823, 0.5081765563442366, 0.5437595540422571, 0.5590205548347534, 0.5593249815276734, 0.5603062368164834, 0.560333188635217, 0.5614503754617461, 0.565352842229935, 0.574210838214554, 0.5759450483859969, 0.5773498217058918, 0.5991138115095911, 0.6009475544728957, 0.6108938307533, 0.6316565296547284, 0.6404495093354053, 0.6405262429561641, 0.6425694115212065, 0.658671129040488, 0.6668423897406515, 0.6864391962767343, 0.7035635283169166, 0.7325106678655877, 0.7328050041291218, 0.7614304100703713, 0.7631239070049998, 0.7670021786149205, 0.7697753383420857, 0.7764360990307122, 0.7784918983501654, 0.7973920072996036, 0.819715865079681, 0.8506721053047003, 0.8813167497816289, 0.8824879447595726, 0.9185813970744787, 0.9231889896940375, 0.9237877978193884, 0.9255031346434324, 0.9293883502480845, 0.9294097332465232, 0.9463098243875633, 0.946325164889271, 0.9491397432856566, 0.9567595541247681, 0.9706712283358269, 0.9723580396501548, 0.9748360509016578, 0.9800193410444061, 0.980809631269599, 0.991517828651004, 0.9965400387585364] [0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8, 3BEOHQsMEFZ58VcNTOJYShTBpAPzbt, 4HX6feIvmNXBN7XGqgO4YVBkhu8GDI, 4JznSdBajNWhu4hRQwjV1FjTTxY68i, 52mKlRE3aHCBZtjECq6sY9OqVf8Dze, 56MZa5O1hVtX4c5sbnCfxuX5kDChqI, 6FPJlLAcaQ5uokyOWZ9HGdLZObFvOZ, 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW, 6oIXZuIPIqEoPBvFmbt2Nxy3tryGUE, 6x93sxYioWuq5c9Kkk8oTAAORM7cH0, 802bgTGl6Bk5TlkPYYTxp5JkKyaYUA, 8LIh0b6jmDGm87BmIyjdxNIpX4ugjD, 90gAtmGEeIqUTbo1ZrxCvWtsseukXC, 9UbObCsVkmYpJGcGrgfK90qOnwb2Lj, AFGCj7OWlEB5QfniEFgonMq90Tq5uH, ALuRhobVWbnQTTWZdSOk0iVe8oYFhW, Amn2K87Db5Es3dFQO9cw9cvpAM6h35, AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz, BJqx5WokrmrrezZA0dUbleMYkG5U2O, BPtQMxnuSPpxMExYV9YkDa6cAN7GP3, BsM5ZAYifRh5Lw3Y8X1r53I0cTJnfE, C2GT5KVyOPZpgKVl110TyZO0NcJ434, DuJNG8tufSqW0ZstHqWj3aGvFLMg4A, EcCuckwsF3gV1Ecgmh5v4KM8g1ozif, ErJFw6hzZ5fmI5r8bhE4JzlscnhKZU, F7NSTjWvQJyBburN7CXRUlbgp2dIrA, Fi4rJeTQq4eXj8Lxg3Hja5hBVTVV5u, H5j5ZHy1FGesOAHjkQEDYCucbpKWRu, HKSMQ9nTnwXCJIte1JrM1dtYnDtJ8g, IWl0G3ZlMNf7WT8yjIB49cx7MmYOmr, IZTkHMLvIKuiLjhDjYMmIHxh166we4, Ig1QcuKsjHXkproePdERo2w0mYzIqd, JHNgc2UCaiXOdmkxwDDyGhRlO0mnBQ, JN0VclewmjwYlSl8386MlWv5rEhWCz, JafwVLSVk5AVoXFuzclesQ000EE2k1, KJFcmTVjdkCMv94wYCtfHMFhzyRsmH, Ktb7GQ0N1DrxwkCkEUsTaIXk0xYinn, Ld2ej8NEv5zNcqU60FwpHeZKBhfpiV, LiEBxds3X0Uw0lxiYjDqrkAaAwoiIW, MXhhH1Var3OzzJCtI9VNyYvA0q8UyJ, MeSTAXq8gVxVjbEjgkvU9YLte0X9uE, NEhyk8uIx4kEULJGa8qIyFjjBcP2G6, O66j6PaYuZhEUtqV6fuU7TyjM2WxC5, OF7fQ37GzaZ5ikA2oMyvleKtgnLjXh, OPwBqCEK5PWTjWaiOyL45u2NLTaDWv, Oq6J4Rx6nde0YlhOIJkFsX2MsSvAQ0, Ow5PGpfTm4dXCfTDsXAOTatXRoAydR, QEHVvcP8gxI6EMJIrvcnIhgzPNjIvv, QJYm7YRA3YetcBHI5wkMZeLXVmfuNy, QYlaIAnJA6r8rlAb6f59wcxvcPcWFf, RilTlL1tKkPOUFuzmLydHAVZwv1OGl, Sfx0vxv1skzZWT1PqVdoRDdO6Sb6xH, TTQUwpMNSXZqVBKAFvXu7OlWvKXJKX, TtDKUZxzVxsq758G6AWPSYuZgVgbcl, VDhtJkYjAYPykCgOU9x3v7v3t4SO1a, VY0zXmXeksCT8BzvpzpPLbmU9Kp9Y4, Vp3gmWunM5A7wOC9YW2JroFqTWjvTi, WHmjWk2AY4c6m7DA4GitUx6nmb1yYS, XemNcT1xp61xcM1Qz3wZ1VECCnq06O, Z2sWcQr0qyCJRMHDpRy3aQr7PkHtkK, aDxBtor7Icd9C5hnTvvw5NrIre740e, akiiY5N0I44CMwEnBL6RTBk7BRkxEj, b3b9esRhTzFEawbs6XhpKnD9ojutHB, bgK1r6v3BCTh0aejJUhkA1Hn6idXGp, cBGc0kSm32ylBDnxogG727C0uhZEYZ, cq4WSAIFwx3wwTUS5bp1wCe71R6U5I, dVdvo6nUD5FgCgsbOZLds28RyGTpnx, e2Gh6Ov8XkXoFdJWhl0EjwEHlMDYyG, f9ALCzwDAKmdu7Rk2msJaB1wxe5IBX, fuyvs0w7WsKSlXqJ1e6HFSoLmx03AG, gTpyQnEODMcpsPnJMZC66gh33i3m0b, gpo8K5qtYePve6jyPt6xgJx4YOVjms, gxfHWUF8XgY2KdFxigxvNEXe2V2XMl, i6RQVXKUh7MzuGMDaNclUYnFUAireU, ioEncce3mPOXD2hWhpZpCPWGATG6GU, jQimhdepw3GKmioWUlVSWeBVRKFkY3, l7uwDoTepWwnAP0ufqtHJS3CRi7RfP, lqhzgLsXZ8JhtpeeUWWNbMz8PHI705, m6jD0LBIQWaMfenwRCTANI9eOdyyto, mhjME0zBHbrK6NMkytMTQzOssOa1gF, mzbkwXKrPeZnxg2Kn1LRF5hYSsmksS, nYVJnVicpGRqKZibHyBAmtmzBXAFfT, oHJMNvWuunsIMIWFnYG31RCfkOo2V7, oLZ21P2JEDooxV1pU31cIxQHEeeoLu, okOkcWflkNXIy4R8LzmySyY1EC3sYd, pLk3i59bZwd5KBZrI1FiweYTd5hteG, pTeu0WMjBRTaNRT15rLCuEh3tBJVc5, qnPOOmslCJaT45buUisMRnM0rc77EK, t6fQUjJejPcjc04wHvHTPe55S65B4V, ukOiFGGFnQJDHFgZxHMpvhD3zybF0M, ukyD7b0Efj7tNlFSRmzZ0IqkEzg2a8, waIGbOGl1PM6gnzZ4uuZt4E2yDWRHs, wwXqSGKLyBQyPkonlzBNYUJTCo4LRS, xipQ93429ksjNcXPX5326VSg1xJZcW, y7C453hRWd4E7ImjNDWlpexB8nUqjh, ydkwycaISlYSlEq3TlkS2m15I2pcp8] + +query ?? rowsort +with tbl as (SELECT * FROM (VALUES ('xxx', 'yyy'), ('xxx', 'yyy'), ('xxx2', 'yyy2')) AS t(x, y)) +select + array_agg(distinct x order by x) as x_agg, + array_agg(distinct y order by y) as y_agg +from tbl +group by all +---- +[xxx, xxx2] [yyy, yyy2] + +query ?? +SELECT + (SELECT array_agg(DISTINCT c12 ORDER BY c12) FROM aggregate_test_100), + (SELECT array_agg(DISTINCT c13 ORDER BY c13) FROM aggregate_test_100) +---- +[0.01479305307777301, 0.02182578039211991, 0.03968347085780355, 0.04429073092078406, 0.047343434291126085, 0.04893135681998029, 0.0494924465469434, 0.05573662213439634, 0.05636955101974106, 0.061029375346466685, 0.07260475960924484, 0.09465635123783445, 0.12357539988406441, 0.152498292971736, 0.16301110515739792, 0.1640882545084913, 0.1754261586710173, 0.17592486905979987, 0.17909035118828576, 0.18628859265874176, 0.19113293583306745, 0.2145232647388039, 0.21535402343780985, 0.24899794314659673, 0.2537253407987472, 0.2667177795079635, 0.27159190516490006, 0.2739938529235548, 0.28534428578703896, 0.2944158618048994, 0.296036538664718, 0.3051364088814128, 0.30585375151301186, 0.3114712539863804, 0.3231750610081745, 0.32869374687050157, 0.33639590659276175, 0.3600766362333053, 0.36936304600612724, 0.38870280983958583, 0.39144436569161134, 0.40342283197779727, 0.4094218353587008, 0.40975383525297016, 0.42073125331890115, 0.4273123318932347, 0.42950521730777025, 0.4830878559436823, 0.5081765563442366, 0.5437595540422571, 0.5590205548347534, 0.5593249815276734, 0.5603062368164834, 0.560333188635217, 0.5614503754617461, 0.565352842229935, 0.574210838214554, 0.5759450483859969, 0.5773498217058918, 0.5991138115095911, 0.6009475544728957, 0.6108938307533, 0.6316565296547284, 0.6404495093354053, 0.6405262429561641, 0.6425694115212065, 0.658671129040488, 0.6668423897406515, 0.6864391962767343, 0.7035635283169166, 0.7325106678655877, 0.7328050041291218, 0.7614304100703713, 0.7631239070049998, 0.7670021786149205, 0.7697753383420857, 0.7764360990307122, 0.7784918983501654, 0.7973920072996036, 0.819715865079681, 0.8506721053047003, 0.8813167497816289, 0.8824879447595726, 0.9185813970744787, 0.9231889896940375, 0.9237877978193884, 0.9255031346434324, 0.9293883502480845, 0.9294097332465232, 0.9463098243875633, 0.946325164889271, 0.9491397432856566, 0.9567595541247681, 0.9706712283358269, 0.9723580396501548, 0.9748360509016578, 0.9800193410444061, 0.980809631269599, 0.991517828651004, 0.9965400387585364] [0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8, 3BEOHQsMEFZ58VcNTOJYShTBpAPzbt, 4HX6feIvmNXBN7XGqgO4YVBkhu8GDI, 4JznSdBajNWhu4hRQwjV1FjTTxY68i, 52mKlRE3aHCBZtjECq6sY9OqVf8Dze, 56MZa5O1hVtX4c5sbnCfxuX5kDChqI, 6FPJlLAcaQ5uokyOWZ9HGdLZObFvOZ, 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW, 6oIXZuIPIqEoPBvFmbt2Nxy3tryGUE, 6x93sxYioWuq5c9Kkk8oTAAORM7cH0, 802bgTGl6Bk5TlkPYYTxp5JkKyaYUA, 8LIh0b6jmDGm87BmIyjdxNIpX4ugjD, 90gAtmGEeIqUTbo1ZrxCvWtsseukXC, 9UbObCsVkmYpJGcGrgfK90qOnwb2Lj, AFGCj7OWlEB5QfniEFgonMq90Tq5uH, ALuRhobVWbnQTTWZdSOk0iVe8oYFhW, Amn2K87Db5Es3dFQO9cw9cvpAM6h35, AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz, BJqx5WokrmrrezZA0dUbleMYkG5U2O, BPtQMxnuSPpxMExYV9YkDa6cAN7GP3, BsM5ZAYifRh5Lw3Y8X1r53I0cTJnfE, C2GT5KVyOPZpgKVl110TyZO0NcJ434, DuJNG8tufSqW0ZstHqWj3aGvFLMg4A, EcCuckwsF3gV1Ecgmh5v4KM8g1ozif, ErJFw6hzZ5fmI5r8bhE4JzlscnhKZU, F7NSTjWvQJyBburN7CXRUlbgp2dIrA, Fi4rJeTQq4eXj8Lxg3Hja5hBVTVV5u, H5j5ZHy1FGesOAHjkQEDYCucbpKWRu, HKSMQ9nTnwXCJIte1JrM1dtYnDtJ8g, IWl0G3ZlMNf7WT8yjIB49cx7MmYOmr, IZTkHMLvIKuiLjhDjYMmIHxh166we4, Ig1QcuKsjHXkproePdERo2w0mYzIqd, JHNgc2UCaiXOdmkxwDDyGhRlO0mnBQ, JN0VclewmjwYlSl8386MlWv5rEhWCz, JafwVLSVk5AVoXFuzclesQ000EE2k1, KJFcmTVjdkCMv94wYCtfHMFhzyRsmH, Ktb7GQ0N1DrxwkCkEUsTaIXk0xYinn, Ld2ej8NEv5zNcqU60FwpHeZKBhfpiV, LiEBxds3X0Uw0lxiYjDqrkAaAwoiIW, MXhhH1Var3OzzJCtI9VNyYvA0q8UyJ, MeSTAXq8gVxVjbEjgkvU9YLte0X9uE, NEhyk8uIx4kEULJGa8qIyFjjBcP2G6, O66j6PaYuZhEUtqV6fuU7TyjM2WxC5, OF7fQ37GzaZ5ikA2oMyvleKtgnLjXh, OPwBqCEK5PWTjWaiOyL45u2NLTaDWv, Oq6J4Rx6nde0YlhOIJkFsX2MsSvAQ0, Ow5PGpfTm4dXCfTDsXAOTatXRoAydR, QEHVvcP8gxI6EMJIrvcnIhgzPNjIvv, QJYm7YRA3YetcBHI5wkMZeLXVmfuNy, QYlaIAnJA6r8rlAb6f59wcxvcPcWFf, RilTlL1tKkPOUFuzmLydHAVZwv1OGl, Sfx0vxv1skzZWT1PqVdoRDdO6Sb6xH, TTQUwpMNSXZqVBKAFvXu7OlWvKXJKX, TtDKUZxzVxsq758G6AWPSYuZgVgbcl, VDhtJkYjAYPykCgOU9x3v7v3t4SO1a, VY0zXmXeksCT8BzvpzpPLbmU9Kp9Y4, Vp3gmWunM5A7wOC9YW2JroFqTWjvTi, WHmjWk2AY4c6m7DA4GitUx6nmb1yYS, XemNcT1xp61xcM1Qz3wZ1VECCnq06O, Z2sWcQr0qyCJRMHDpRy3aQr7PkHtkK, aDxBtor7Icd9C5hnTvvw5NrIre740e, akiiY5N0I44CMwEnBL6RTBk7BRkxEj, b3b9esRhTzFEawbs6XhpKnD9ojutHB, bgK1r6v3BCTh0aejJUhkA1Hn6idXGp, cBGc0kSm32ylBDnxogG727C0uhZEYZ, cq4WSAIFwx3wwTUS5bp1wCe71R6U5I, dVdvo6nUD5FgCgsbOZLds28RyGTpnx, e2Gh6Ov8XkXoFdJWhl0EjwEHlMDYyG, f9ALCzwDAKmdu7Rk2msJaB1wxe5IBX, fuyvs0w7WsKSlXqJ1e6HFSoLmx03AG, gTpyQnEODMcpsPnJMZC66gh33i3m0b, gpo8K5qtYePve6jyPt6xgJx4YOVjms, gxfHWUF8XgY2KdFxigxvNEXe2V2XMl, i6RQVXKUh7MzuGMDaNclUYnFUAireU, ioEncce3mPOXD2hWhpZpCPWGATG6GU, jQimhdepw3GKmioWUlVSWeBVRKFkY3, l7uwDoTepWwnAP0ufqtHJS3CRi7RfP, lqhzgLsXZ8JhtpeeUWWNbMz8PHI705, m6jD0LBIQWaMfenwRCTANI9eOdyyto, mhjME0zBHbrK6NMkytMTQzOssOa1gF, mzbkwXKrPeZnxg2Kn1LRF5hYSsmksS, nYVJnVicpGRqKZibHyBAmtmzBXAFfT, oHJMNvWuunsIMIWFnYG31RCfkOo2V7, oLZ21P2JEDooxV1pU31cIxQHEeeoLu, okOkcWflkNXIy4R8LzmySyY1EC3sYd, pLk3i59bZwd5KBZrI1FiweYTd5hteG, pTeu0WMjBRTaNRT15rLCuEh3tBJVc5, qnPOOmslCJaT45buUisMRnM0rc77EK, t6fQUjJejPcjc04wHvHTPe55S65B4V, ukOiFGGFnQJDHFgZxHMpvhD3zybF0M, ukyD7b0Efj7tNlFSRmzZ0IqkEzg2a8, waIGbOGl1PM6gnzZ4uuZt4E2yDWRHs, wwXqSGKLyBQyPkonlzBNYUJTCo4LRS, xipQ93429ksjNcXPX5326VSg1xJZcW, y7C453hRWd4E7ImjNDWlpexB8nUqjh, ydkwycaISlYSlEq3TlkS2m15I2pcp8] + +query ?? +SELECT + array_agg(DISTINCT c12 ORDER BY c12), + array_agg(DISTINCT c13 ORDER BY c13) +FROM aggregate_test_100 +---- +[0.01479305307777301, 0.02182578039211991, 0.03968347085780355, 0.04429073092078406, 0.047343434291126085, 0.04893135681998029, 0.0494924465469434, 0.05573662213439634, 0.05636955101974106, 0.061029375346466685, 0.07260475960924484, 0.09465635123783445, 0.12357539988406441, 0.152498292971736, 0.16301110515739792, 0.1640882545084913, 0.1754261586710173, 0.17592486905979987, 0.17909035118828576, 0.18628859265874176, 0.19113293583306745, 0.2145232647388039, 0.21535402343780985, 0.24899794314659673, 0.2537253407987472, 0.2667177795079635, 0.27159190516490006, 0.2739938529235548, 0.28534428578703896, 0.2944158618048994, 0.296036538664718, 0.3051364088814128, 0.30585375151301186, 0.3114712539863804, 0.3231750610081745, 0.32869374687050157, 0.33639590659276175, 0.3600766362333053, 0.36936304600612724, 0.38870280983958583, 0.39144436569161134, 0.40342283197779727, 0.4094218353587008, 0.40975383525297016, 0.42073125331890115, 0.4273123318932347, 0.42950521730777025, 0.4830878559436823, 0.5081765563442366, 0.5437595540422571, 0.5590205548347534, 0.5593249815276734, 0.5603062368164834, 0.560333188635217, 0.5614503754617461, 0.565352842229935, 0.574210838214554, 0.5759450483859969, 0.5773498217058918, 0.5991138115095911, 0.6009475544728957, 0.6108938307533, 0.6316565296547284, 0.6404495093354053, 0.6405262429561641, 0.6425694115212065, 0.658671129040488, 0.6668423897406515, 0.6864391962767343, 0.7035635283169166, 0.7325106678655877, 0.7328050041291218, 0.7614304100703713, 0.7631239070049998, 0.7670021786149205, 0.7697753383420857, 0.7764360990307122, 0.7784918983501654, 0.7973920072996036, 0.819715865079681, 0.8506721053047003, 0.8813167497816289, 0.8824879447595726, 0.9185813970744787, 0.9231889896940375, 0.9237877978193884, 0.9255031346434324, 0.9293883502480845, 0.9294097332465232, 0.9463098243875633, 0.946325164889271, 0.9491397432856566, 0.9567595541247681, 0.9706712283358269, 0.9723580396501548, 0.9748360509016578, 0.9800193410444061, 0.980809631269599, 0.991517828651004, 0.9965400387585364] [0VVIHzxWtNOFLtnhjHEKjXaJOSLJfm, 0keZ5G8BffGwgF2RwQD59TFzMStxCB, 0og6hSkhbX8AC1ktFS4kounvTzy8Vo, 1aOcrEGd0cOqZe2I5XBOm0nDcwtBZO, 2T3wSlHdEmASmO0xcXHnndkKEt6bz8, 3BEOHQsMEFZ58VcNTOJYShTBpAPzbt, 4HX6feIvmNXBN7XGqgO4YVBkhu8GDI, 4JznSdBajNWhu4hRQwjV1FjTTxY68i, 52mKlRE3aHCBZtjECq6sY9OqVf8Dze, 56MZa5O1hVtX4c5sbnCfxuX5kDChqI, 6FPJlLAcaQ5uokyOWZ9HGdLZObFvOZ, 6WfVFBVGJSQb7FhA7E0lBwdvjfZnSW, 6oIXZuIPIqEoPBvFmbt2Nxy3tryGUE, 6x93sxYioWuq5c9Kkk8oTAAORM7cH0, 802bgTGl6Bk5TlkPYYTxp5JkKyaYUA, 8LIh0b6jmDGm87BmIyjdxNIpX4ugjD, 90gAtmGEeIqUTbo1ZrxCvWtsseukXC, 9UbObCsVkmYpJGcGrgfK90qOnwb2Lj, AFGCj7OWlEB5QfniEFgonMq90Tq5uH, ALuRhobVWbnQTTWZdSOk0iVe8oYFhW, Amn2K87Db5Es3dFQO9cw9cvpAM6h35, AyYVExXK6AR2qUTxNZ7qRHQOVGMLcz, BJqx5WokrmrrezZA0dUbleMYkG5U2O, BPtQMxnuSPpxMExYV9YkDa6cAN7GP3, BsM5ZAYifRh5Lw3Y8X1r53I0cTJnfE, C2GT5KVyOPZpgKVl110TyZO0NcJ434, DuJNG8tufSqW0ZstHqWj3aGvFLMg4A, EcCuckwsF3gV1Ecgmh5v4KM8g1ozif, ErJFw6hzZ5fmI5r8bhE4JzlscnhKZU, F7NSTjWvQJyBburN7CXRUlbgp2dIrA, Fi4rJeTQq4eXj8Lxg3Hja5hBVTVV5u, H5j5ZHy1FGesOAHjkQEDYCucbpKWRu, HKSMQ9nTnwXCJIte1JrM1dtYnDtJ8g, IWl0G3ZlMNf7WT8yjIB49cx7MmYOmr, IZTkHMLvIKuiLjhDjYMmIHxh166we4, Ig1QcuKsjHXkproePdERo2w0mYzIqd, JHNgc2UCaiXOdmkxwDDyGhRlO0mnBQ, JN0VclewmjwYlSl8386MlWv5rEhWCz, JafwVLSVk5AVoXFuzclesQ000EE2k1, KJFcmTVjdkCMv94wYCtfHMFhzyRsmH, Ktb7GQ0N1DrxwkCkEUsTaIXk0xYinn, Ld2ej8NEv5zNcqU60FwpHeZKBhfpiV, LiEBxds3X0Uw0lxiYjDqrkAaAwoiIW, MXhhH1Var3OzzJCtI9VNyYvA0q8UyJ, MeSTAXq8gVxVjbEjgkvU9YLte0X9uE, NEhyk8uIx4kEULJGa8qIyFjjBcP2G6, O66j6PaYuZhEUtqV6fuU7TyjM2WxC5, OF7fQ37GzaZ5ikA2oMyvleKtgnLjXh, OPwBqCEK5PWTjWaiOyL45u2NLTaDWv, Oq6J4Rx6nde0YlhOIJkFsX2MsSvAQ0, Ow5PGpfTm4dXCfTDsXAOTatXRoAydR, QEHVvcP8gxI6EMJIrvcnIhgzPNjIvv, QJYm7YRA3YetcBHI5wkMZeLXVmfuNy, QYlaIAnJA6r8rlAb6f59wcxvcPcWFf, RilTlL1tKkPOUFuzmLydHAVZwv1OGl, Sfx0vxv1skzZWT1PqVdoRDdO6Sb6xH, TTQUwpMNSXZqVBKAFvXu7OlWvKXJKX, TtDKUZxzVxsq758G6AWPSYuZgVgbcl, VDhtJkYjAYPykCgOU9x3v7v3t4SO1a, VY0zXmXeksCT8BzvpzpPLbmU9Kp9Y4, Vp3gmWunM5A7wOC9YW2JroFqTWjvTi, WHmjWk2AY4c6m7DA4GitUx6nmb1yYS, XemNcT1xp61xcM1Qz3wZ1VECCnq06O, Z2sWcQr0qyCJRMHDpRy3aQr7PkHtkK, aDxBtor7Icd9C5hnTvvw5NrIre740e, akiiY5N0I44CMwEnBL6RTBk7BRkxEj, b3b9esRhTzFEawbs6XhpKnD9ojutHB, bgK1r6v3BCTh0aejJUhkA1Hn6idXGp, cBGc0kSm32ylBDnxogG727C0uhZEYZ, cq4WSAIFwx3wwTUS5bp1wCe71R6U5I, dVdvo6nUD5FgCgsbOZLds28RyGTpnx, e2Gh6Ov8XkXoFdJWhl0EjwEHlMDYyG, f9ALCzwDAKmdu7Rk2msJaB1wxe5IBX, fuyvs0w7WsKSlXqJ1e6HFSoLmx03AG, gTpyQnEODMcpsPnJMZC66gh33i3m0b, gpo8K5qtYePve6jyPt6xgJx4YOVjms, gxfHWUF8XgY2KdFxigxvNEXe2V2XMl, i6RQVXKUh7MzuGMDaNclUYnFUAireU, ioEncce3mPOXD2hWhpZpCPWGATG6GU, jQimhdepw3GKmioWUlVSWeBVRKFkY3, l7uwDoTepWwnAP0ufqtHJS3CRi7RfP, lqhzgLsXZ8JhtpeeUWWNbMz8PHI705, m6jD0LBIQWaMfenwRCTANI9eOdyyto, mhjME0zBHbrK6NMkytMTQzOssOa1gF, mzbkwXKrPeZnxg2Kn1LRF5hYSsmksS, nYVJnVicpGRqKZibHyBAmtmzBXAFfT, oHJMNvWuunsIMIWFnYG31RCfkOo2V7, oLZ21P2JEDooxV1pU31cIxQHEeeoLu, okOkcWflkNXIy4R8LzmySyY1EC3sYd, pLk3i59bZwd5KBZrI1FiweYTd5hteG, pTeu0WMjBRTaNRT15rLCuEh3tBJVc5, qnPOOmslCJaT45buUisMRnM0rc77EK, t6fQUjJejPcjc04wHvHTPe55S65B4V, ukOiFGGFnQJDHFgZxHMpvhD3zybF0M, ukyD7b0Efj7tNlFSRmzZ0IqkEzg2a8, waIGbOGl1PM6gnzZ4uuZt4E2yDWRHs, wwXqSGKLyBQyPkonlzBNYUJTCo4LRS, xipQ93429ksjNcXPX5326VSg1xJZcW, y7C453hRWd4E7ImjNDWlpexB8nUqjh, ydkwycaISlYSlEq3TlkS2m15I2pcp8] + +statement ok +CREATE EXTERNAL TABLE agg_order ( +c1 INT NOT NULL, +c2 INT NOT NULL, +c3 INT NOT NULL +) +STORED AS CSV +LOCATION '../core/tests/data/aggregate_agg_multi_order.csv' +OPTIONS ('format.has_header' 'true'); + +# test array_agg with order by multiple columns +query ? +select array_agg(c1 order by c2 desc, c3) from agg_order; +---- +[5, 6, 7, 8, 9, 1, 2, 3, 4, 10] + +query TT +explain select array_agg(c1 order by c2 desc, c3) from agg_order; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[array_agg(agg_order.c1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]]] +02)--TableScan: agg_order projection=[c1, c2, c3] +physical_plan +01)AggregateExec: mode=Final, gby=[], aggr=[array_agg(agg_order.c1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=Partial, gby=[], aggr=[array_agg(agg_order.c1) ORDER BY [agg_order.c2 DESC NULLS FIRST, agg_order.c3 ASC NULLS LAST]] +04)------SortExec: expr=[c2@1 DESC, c3@2 ASC NULLS LAST], preserve_partitioning=[true] +05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1, c2, c3], file_type=csv, has_header=true + +# Regression test: ARRAY_AGG with conflicting ASC/DESC ORDER BY in the same query. +# get_finer_aggregate_exprs_requirement picks ASC as the common requirement and +# reverses the DESC aggregate (is_reversed=true, ordering_req=[ASC]). +# The optimizer then sets is_input_pre_ordered=true on both. Without the fix, +# state() emits values reversed to DESC but ordering keys still in ASC order, +# causing merge_batch to pair each value with the wrong key (silent wrong results). +query TT +explain select array_agg(c1 order by c1), array_agg(c1 order by c1 desc) from agg_order; +---- +logical_plan +01)Aggregate: groupBy=[[]], aggr=[[array_agg(agg_order.c1) ORDER BY [agg_order.c1 ASC NULLS LAST], array_agg(agg_order.c1) ORDER BY [agg_order.c1 DESC NULLS FIRST]]] +02)--TableScan: agg_order projection=[c1] +physical_plan +01)AggregateExec: mode=Final, gby=[], aggr=[array_agg(agg_order.c1) ORDER BY [agg_order.c1 ASC NULLS LAST], array_agg(agg_order.c1) ORDER BY [agg_order.c1 DESC NULLS FIRST]] +02)--CoalescePartitionsExec +03)----AggregateExec: mode=Partial, gby=[], aggr=[array_agg(agg_order.c1) ORDER BY [agg_order.c1 ASC NULLS LAST], array_agg(agg_order.c1) ORDER BY [agg_order.c1 DESC NULLS FIRST]] +04)------SortExec: expr=[c1@0 ASC NULLS LAST], preserve_partitioning=[true] +05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_agg_multi_order.csv]]}, projection=[c1], file_type=csv, has_header=true + +query ?? +select array_agg(c1 order by c1), array_agg(c1 order by c1 desc) from agg_order; +---- +[1, 2, 3, 4, 5, 6, 7, 8, 9, 10] [10, 9, 8, 7, 6, 5, 4, 3, 2, 1] + +# test array_agg_order with list data type +statement ok +CREATE TABLE array_agg_order_list_table AS VALUES + ('w', 2, [1,2,3], 10), + ('w', 1, [9,5,2], 20), + ('w', 1, [3,2,5], 30), + ('b', 2, [4,5,6], 20), + ('b', 1, [7,8,9], 30) +; + +query T? rowsort +select column1, array_agg(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; +---- +b [[7, 8, 9], [4, 5, 6]] +w [[3, 2, 5], [9, 5, 2], [1, 2, 3]] + +query T?? rowsort +select column1, first_value(column3 order by column2, column4 desc), last_value(column3 order by column2, column4 desc) from array_agg_order_list_table group by column1; +---- +b [7, 8, 9] [4, 5, 6] +w [3, 2, 5] [1, 2, 3] + +query T? rowsort +select column1, nth_value(column3, 2 order by column2, column4 desc) from array_agg_order_list_table group by column1; +---- +b [4, 5, 6] +w [9, 5, 2] + +query ? +select array_agg(DISTINCT column2 order by column2) from array_agg_order_list_table; +---- +[1, 2] + +query ? +select array_agg(DISTINCT column2 order by column2 desc) from array_agg_order_list_table; +---- +[2, 1] + +query ? +select array_agg(DISTINCT column2 + 1 order by column2 + 1 desc) from array_agg_order_list_table; +---- +[3, 2] + +query ? +select array_agg(DISTINCT column2 order by column2) from array_agg_order_list_table GROUP BY column1; +---- +[1, 2] +[1, 2] + +statement error In an aggregate with DISTINCT, ORDER BY expressions must appear in argument list +select array_agg(DISTINCT column2 order by column1) from array_agg_order_list_table; + +statement ok +drop table array_agg_order_list_table; + +####### +# array_agg with DISTINCT +####### + +# select with count to forces array_agg_distinct function, since single distinct expression is converted to group by by optimizer +# csv_query_array_agg_distinct +query ?I +SELECT array_sort(array_agg(distinct c2)), count(1) FROM aggregate_test_100 +---- +[1, 2, 3, 4, 5] 100 + +# test array_agg_distinct with list data type +statement ok +CREATE TABLE array_agg_distinct_list_table AS VALUES + ('w', [0,1]), + ('w', [0,1]), + ('w', [1,0]), + ('b', [1,0]), + ('b', [1,0]), + ('b', [1,0]), + ('b', [0,1]), + (NULL, [0,1]), + ('b', NULL) +; + +# Apply array_sort to have deterministic result, higher dimension nested array also works but not for array sort, +# so they are covered in `datafusion/functions-aggregate/src/array_agg.rs` +query ?? +select array_sort(c1), array_sort(c2) from ( + select array_agg(distinct column1) as c1, array_agg(distinct column2) ignore nulls as c2 from array_agg_distinct_list_table +); +---- +[NULL, b, w] [[0, 1], [1, 0]] + +statement ok +drop table array_agg_distinct_list_table; + + +# Test array_agg with DISTINCT and IGNORE NULLS (regression test for issue #19735) +query ? +SELECT array_sort(ARRAY_AGG(DISTINCT x IGNORE NULLS)) as result +FROM (VALUES (1), (2), (NULL), (2), (NULL), (1)) AS t(x); +---- +[1, 2] + +# Test distinct aggregate function with merge batch +query II +with A as ( + select 1 as id, 2 as foo + UNION ALL + select 1, null + UNION ALL + select 1, null + UNION ALL + select 1, 3 + UNION ALL + select 1, 2 + ---- The order is non-deterministic, verify with length +) select array_length(array_agg(distinct a.foo)), sum(distinct 1) from A a group by a.id; +---- +3 1 + +# It has only AggregateExec with FinalPartitioned mode, so `merge_batch` is used +# If the plan is changed, whether the `merge_batch` is used should be verified to ensure the test coverage +query TT +explain with A as ( + select 1 as id, 2 as foo + UNION ALL + select 1, null + UNION ALL + select 1, null + UNION ALL + select 1, 3 + UNION ALL + select 1, 2 +) select array_length(array_agg(distinct a.foo)), sum(distinct 1) from A a group by a.id; +---- +logical_plan +01)Projection: array_length(array_agg(DISTINCT a.foo)), sum(DISTINCT Int64(1)) +02)--Aggregate: groupBy=[[a.id]], aggr=[[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))]] +03)----SubqueryAlias: a +04)------SubqueryAlias: a +05)--------Union +06)----------Projection: Int64(1) AS id, Int64(2) AS foo +07)------------EmptyRelation: rows=1 +08)----------Projection: Int64(1) AS id, Int64(NULL) AS foo +09)------------EmptyRelation: rows=1 +10)----------Projection: Int64(1) AS id, Int64(NULL) AS foo +11)------------EmptyRelation: rows=1 +12)----------Projection: Int64(1) AS id, Int64(3) AS foo +13)------------EmptyRelation: rows=1 +14)----------Projection: Int64(1) AS id, Int64(2) AS foo +15)------------EmptyRelation: rows=1 +physical_plan +01)ProjectionExec: expr=[array_length(array_agg(DISTINCT a.foo)@1) as array_length(array_agg(DISTINCT a.foo)), sum(DISTINCT Int64(1))@2 as sum(DISTINCT Int64(1))] +02)--AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted +03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=5 +04)------AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[array_agg(DISTINCT a.foo), sum(DISTINCT Int64(1))], ordering_mode=Sorted +05)--------UnionExec +06)----------ProjectionExec: expr=[1 as id, CAST(2 AS Int64) as foo] +07)------------PlaceholderRowExec +08)----------ProjectionExec: expr=[1 as id, NULL as foo] +09)------------PlaceholderRowExec +10)----------ProjectionExec: expr=[1 as id, NULL as foo] +11)------------PlaceholderRowExec +12)----------ProjectionExec: expr=[1 as id, CAST(3 AS Int64) as foo] +13)------------PlaceholderRowExec +14)----------ProjectionExec: expr=[1 as id, CAST(2 AS Int64) as foo] +15)------------PlaceholderRowExec + +####### +# Unsupported syntax +####### + +statement error This feature is not implemented: Calling array_agg: LIMIT not supported in function arguments: 1 +SELECT array_agg(c13 LIMIT 1) FROM aggregate_test_100 + +query error Error during planning: WITHIN GROUP is only supported for ordered-set aggregate functions +SELECT array_agg(a_varchar) WITHIN GROUP (ORDER BY a_varchar) +FROM (VALUES ('a'), ('d'), ('c'), ('a')) t(a_varchar); + + +query error Error during planning: WITHIN GROUP is only supported for ordered-set aggregate functions +SELECT array_agg(DISTINCT a_varchar) WITHIN GROUP (ORDER BY a_varchar) +FROM (VALUES ('a'), ('d'), ('c'), ('a')) t(a_varchar); + + +query error Error during planning: ORDER BY and WITHIN GROUP clauses cannot be used together in the same aggregate function +SELECT array_agg(a_varchar order by a_varchar) WITHIN GROUP (ORDER BY a_varchar) +FROM (VALUES ('a'), ('d'), ('c'), ('a')) t(a_varchar); + +# test array_agg_distinct with dictionary encoded data +statement ok +CREATE TABLE array_agg_distinct_dict_table AS VALUES + ('w', 1), + ('w', 1), + ('b', 2), + ('b', 1), + (NULL, 2) +; + +# Apply array_sort to have deterministic result +query ?? +select array_sort(c1), array_sort(c2) from ( + select array_agg(distinct arrow_cast(column1, 'Dictionary(Int32, Utf8)')) as c1, + array_agg(distinct arrow_cast(column2, 'Dictionary(Int8, Int64)')) ignore nulls as c2 + from array_agg_distinct_dict_table +); +---- +[NULL, b, w] [1, 2] + +# The element type of the returned list must stay dictionary encoded, otherwise the +# aggregate output does not match the schema it declared +query T +select arrow_typeof(array_agg(distinct arrow_cast(column1, 'Dictionary(Int32, Utf8)'))) +from array_agg_distinct_dict_table; +---- +List(Dictionary(Int32, Utf8)) + +# ... including when the dictionary is nested inside another type +query ? +select array_sort(c) from ( + select array_agg(distinct struct(arrow_cast(column1, 'Dictionary(Int32, Utf8)') as f)) as c + from array_agg_distinct_dict_table +); +---- +[{f: NULL}, {f: b}, {f: w}] + +# ... and when no rows are aggregated at all +query ? +select array_agg(distinct arrow_cast(column1, 'Dictionary(Int32, Utf8)')) +from array_agg_distinct_dict_table where column2 > 100; +---- +NULL + +query T +select arrow_typeof(array_agg(distinct arrow_cast(column1, 'Dictionary(Int32, Utf8)'))) +from array_agg_distinct_dict_table where column2 > 100; +---- +List(Dictionary(Int32, Utf8)) + +statement ok +drop table array_agg_distinct_dict_table; diff --git a/datafusion/sqllogictest/test_files/array_agg_sliding_window.slt b/datafusion/sqllogictest/test_files/array_agg_sliding_window.slt index 78d48513a6656..c828794b1dcb7 100644 --- a/datafusion/sqllogictest/test_files/array_agg_sliding_window.slt +++ b/datafusion/sqllogictest/test_files/array_agg_sliding_window.slt @@ -168,6 +168,233 @@ FROM t_nulls; [C] [C, E] +####### +# DISTINCT sliding window tests +# Validates retract_batch implementation on DistinctArrayAggAccumulator. +# DataFusion rejects `array_agg(... ORDER BY ...)` inside window functions, +# so we wrap with array_sort to make output deterministic +# (HashMap iteration order otherwise). +####### + +statement ok +CREATE TABLE t_dist(ts INT, val TEXT) AS VALUES + (1,'A'),(2,'A'),(3,'B'),(4,'C'),(5,'B'); + +# Duplicate stays in frame after partial retract. +# Frame contents per row (ts=1..5): +# [A] -> {A} +# [A,A] -> {A} (A appears twice, still distinct {A}) +# [A,A,B] -> {A,B} +# [A,B,C] -> {A,B,C} (one A retracted, one A remains) +# [B,C,B] -> {B,C} (last A retracted, B duplicate stays) +query ? +SELECT array_sort(array_agg(DISTINCT val) + OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND CURRENT ROW)) +FROM t_dist; +---- +[A] +[A] +[A, B] +[A, B, C] +[B, C] + +# Narrower ROWS frame +query ? +SELECT array_sort(array_agg(DISTINCT val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_dist; +---- +[A] +[A] +[A, B] +[B, C] +[B, C] + +# DESC window ORDER BY: frame walks input in reverse temporal order, so +# update/retract are called against the reversed row stream. Validates +# retract still tracks duplicates correctly when rows arrive in DESC order. +# Output rows are emitted in ts DESC order (ts=5,4,3,2,1). +# ts=5 (B): frame [B] -> {B} +# ts=4 (C): frame [B,C] -> {B,C} (1 preceding in DESC = ts=5) +# ts=3 (B): frame [C,B] -> {B,C} +# ts=2 (A): frame [B,A] -> {A,B} +# ts=1 (A): frame [A,A] -> {A} (duplicate A in frame) +query ? +SELECT array_sort(array_agg(DISTINCT val) + OVER (ORDER BY ts DESC ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_dist; +---- +[B] +[B, C] +[B, C] +[A, B] +[A] + +# RANGE frame with value gaps -> multi-row retract on shift +statement ok +CREATE TABLE t_dist_range(ts INT, val TEXT) AS VALUES + (1,'A'),(2,'A'),(3,'B'),(10,'A'),(11,'C'); + +query ? +SELECT array_sort(array_agg(DISTINCT val) + OVER (ORDER BY ts RANGE BETWEEN 2 PRECEDING AND CURRENT ROW)) +FROM t_dist_range; +---- +[A] +[A] +[A, B] +[A] +[A, C] + +# DISTINCT + IGNORE NULLS in sliding frame: nulls never enter state. +statement ok +CREATE TABLE t_dist_nulls(ts INT, val TEXT) AS VALUES + (1,'A'),(2,NULL),(3,'A'),(4,NULL),(5,'B'); + +query ? +SELECT array_sort(array_agg(DISTINCT val) IGNORE NULLS + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_dist_nulls; +---- +[A] +[A] +[A] +[A] +[B] + +# DISTINCT without IGNORE NULLS: NULL enters state with a refcount. +# Retract must remove the NULL key when its last occurrence leaves the frame. +# array_sort defaults to ASC NULLS FIRST, so a live NULL sorts ahead of A/B; +# rows with no live NULL have no NULL element. +# ts=1 (A): frame [A] -> {A} sorted [A] +# ts=2 (NULL): frame [A,NULL] -> {A,NULL} sorted [NULL, A] +# ts=3 (A): frame [NULL,A] -> {A,NULL} sorted [NULL, A] +# ts=4 (NULL): frame [A,NULL] -> {A,NULL} sorted [NULL, A] +# (the ts=2 NULL retracts but the ts=4 NULL is still present) +# ts=5 (B): frame [NULL,B] -> {B,NULL} sorted [NULL, B] +query ? +SELECT array_sort(array_agg(DISTINCT val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_dist_nulls; +---- +[A] +[NULL, A] +[NULL, A] +[NULL, A] +[NULL, B] + +# GROUPS frame with duplicated sort keys: rows tied on the ORDER BY column +# are batched into the same group, so a single shift can update or retract +# multiple rows at once. +statement ok +CREATE TABLE t_dist_groups(ts INT, val TEXT) AS VALUES + (1,'A'),(1,'A'),(2,'B'),(2,'C'),(3,'A'),(3,'D'); + +query ? +SELECT array_sort(array_agg(DISTINCT val) + OVER (ORDER BY ts GROUPS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_dist_groups; +---- +[A] +[A] +[A, B, C] +[A, B, C] +[A, B, C, D] +[A, B, C, D] + +# PARTITION BY: each partition retracts against its own state only. A leak of +# one partition's state into the next would surface as the next partition's +# first row carrying foreign values, or a retract hitting the +# `value not present in state` internal_err. 'A' lives only in grp 1, 'C' only +# in grp 2, 'B' in both — so leaked grp-1 state would make grp=2/ts=1 emit +# [A, B] instead of [B]. Rows emitted in (grp, ts) order. +# grp 1: ts=1 [A]->{A} ts=2 [A,A]->{A} ts=3 [A,B]->{A,B} +# grp 2: ts=1 [B]->{B} ts=2 [B,C]->{B,C} ts=3 [C,C]->{C} +statement ok +CREATE TABLE t_dist_parts(grp INT, ts INT, val TEXT) AS VALUES + (1,1,'A'),(1,2,'A'),(1,3,'B'), + (2,1,'B'),(2,2,'C'),(2,3,'C'); + +query ? +SELECT array_sort(array_agg(DISTINCT val) + OVER (PARTITION BY grp ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_dist_parts +ORDER BY grp, ts; +---- +[A] +[A] +[A, B] +[B] +[B, C] +[C] + +# Numeric element type: retract must hash and compare Int32 ScalarValues +# correctly (every sibling test uses Utf8). Mirrors the t_dist 2-PRECEDING walk. +# ts=1 [10] -> {10} +# ts=2 [10,10] -> {10} (duplicate, stays distinct {10}) +# ts=3 [10,10,20] -> {10,20} +# ts=4 [10,20,30] -> {10,20,30} (one 10 retracted, one 10 remains) +# ts=5 [20,30,20] -> {20,30} (last 10 retracted, 20 duplicate stays) +statement ok +CREATE TABLE t_dist_int(ts INT, val INT) AS VALUES + (1,10),(2,10),(3,20),(4,30),(5,20); + +query ? +SELECT array_sort(array_agg(DISTINCT val) + OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND CURRENT ROW)) +FROM t_dist_int; +---- +[10] +[10] +[10, 20] +[10, 20, 30] +[20, 30] + +# ORDER BY interaction — window context. +# DataFusion's planner rejects ANY aggregate-level ORDER BY inside a window +# function, so neither the valid (DISTINCT x ORDER BY x) nor the invalid +# (DISTINCT x ORDER BY y) form reaches the DISTINCT-arg-equality validator +# in window context. Both error at planning, but at the window-planner stage. +statement error Aggregate ORDER BY is not implemented for window functions +SELECT array_agg(DISTINCT val ORDER BY val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM t_dist; + +statement error Aggregate ORDER BY is not implemented for window functions +SELECT array_agg(DISTINCT val ORDER BY ts) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM t_dist; + +# ORDER BY interaction — non-window context (regression for the storage swap). +# The DISTINCT-arg-equality validator must still accept the matching case +# and reject the mismatched case after we changed the underlying state. +query ? +SELECT array_agg(DISTINCT val ORDER BY val) FROM t_dist; +---- +[A, B, C] + +statement error In an aggregate with DISTINCT, ORDER BY expressions must appear in argument list +SELECT array_agg(DISTINCT val ORDER BY ts) FROM t_dist; + +# Result cardinality bounded by frame cardinality (live-key proxy for state growth). +# Set up 100 rows over 50 cycling distinct values, then run a 2-row sliding frame. +# Since `evaluate` returns the live key set verbatim, max(result_length) == 2 +# proves keys are dropped as their last occurrence leaves the frame. A leaky +# retract would let the result balloon toward 50 (all distinct values seen) or +# error at runtime via the `value not present in state` internal_err!. +statement ok +CREATE TABLE t_dist_growth AS + SELECT i AS ts, ('v' || (i % 50)::TEXT) AS val FROM generate_series(1, 100) t(i); + +query I +SELECT max(cardinality(distinct_arr)) FROM ( + SELECT array_agg(DISTINCT val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS distinct_arr + FROM t_dist_growth +); +---- +2 + # Cleanup statement ok DROP TABLE t; @@ -182,4 +409,25 @@ statement ok DROP TABLE t_int; statement ok -DROP TABLE t_groups; \ No newline at end of file +DROP TABLE t_groups; + +statement ok +DROP TABLE t_dist; + +statement ok +DROP TABLE t_dist_range; + +statement ok +DROP TABLE t_dist_nulls; + +statement ok +DROP TABLE t_dist_groups; + +statement ok +DROP TABLE t_dist_growth; + +statement ok +DROP TABLE t_dist_parts; + +statement ok +DROP TABLE t_dist_int; diff --git a/datafusion/sqllogictest/test_files/array_avg.slt b/datafusion/sqllogictest/test_files/array_avg.slt new file mode 100644 index 0000000000000..ae00a38ae68d6 --- /dev/null +++ b/datafusion/sqllogictest/test_files/array_avg.slt @@ -0,0 +1,165 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +## array_avg + +# Basic case +query R +select array_avg([1.0, 2.0, 3.0]); +---- +2 + +# Single element +query R +select array_avg([5.0]); +---- +5 + +# Negative values +query R +select array_avg([-1.0, -2.0, -3.0]); +---- +-2 + +# Positive and negative cancel +query R +select array_avg([1.0, -1.0, 2.0, -2.0]); +---- +0 + +# Non-integer mean (sum / count) +query R +select array_avg([1.0, 2.0]); +---- +1.5 + +# Empty array returns NULL (matches PostgreSQL AVG, DuckDB list_avg, SQL Standard AVG-of-empty-set) +query R +select array_avg(arrow_cast(make_array(), 'List(Float64)')); +---- +NULL + +# Bare NULL input returns NULL row +query R +select array_avg(NULL); +---- +NULL + +# NULL elements are skipped from BOTH the sum and the count (SQL aggregate convention). +# avg([1, NULL, 3]) = (1 + 3) / 2 = 2 — not (1 + 3) / 3. +query R +select array_avg([1.0, NULL, 3.0]); +---- +2 + +# Single NULL among numeric: skip the NULL, divide by 1 +query R +select array_avg([NULL, 10.0]); +---- +10 + +# All-NULL array returns NULL row (matches SQL AVG over all-NULL) +query R +select array_avg(arrow_cast([NULL, NULL], 'List(Float64)')); +---- +NULL + +# LargeList support +query R +select array_avg(arrow_cast([1.0, 2.0, 3.0], 'LargeList(Float64)')); +---- +2 + +# FixedSizeList input (coerced to List) +query R +select array_avg(arrow_cast([1.0, 2.0, 3.0], 'FixedSizeList(3, Float64)')); +---- +2 + +# Float32 inner type (coerced to Float64) +query R +select array_avg(arrow_cast([1.0, 2.0, 3.0], 'List(Float32)')); +---- +2 + +# Int64 inner type (coerced to Float64) — integer mean returned as Float64 +query R +select array_avg(arrow_cast([1, 2, 3], 'List(Int64)')); +---- +2 + +# Integer literals (coerced to Float64) +query R +select array_avg([1, 2, 3]); +---- +2 + +# Integer mean that is NOT an integer (3 / 2 = 1.5) +query R +select array_avg([1, 2]); +---- +1.5 + +# Unsupported non-list input (plan error) +query error array_avg does not support type +select array_avg(1); + +# Multi-row query with mix of normal, partial-NULL, all-NULL elements, empty, NULL row +query R +select array_avg(column1) from (values + (make_array(1.0, 2.0, 3.0)), + (make_array(0.0)), + (make_array(1.0, NULL, 4.0)), + (arrow_cast(make_array(), 'List(Float64)')), + (NULL) +) as t(column1); +---- +2 +0 +2.5 +NULL +NULL + +# Wrong arity (zero args) +query error array_avg function requires 1 argument, got 0 +select array_avg(); + +# Wrong arity (two args) +query error array_avg function requires 1 argument, got 2 +select array_avg([1.0], [2.0]); + +# Return type is Float64 +query RT +select array_avg([1.0, 2.0, 3.0]), arrow_typeof(array_avg([1.0, 2.0, 3.0])); +---- +2 Float64 + +# list_avg alias produces the same result +query R +select list_avg([1.0, 2.0, 3.0]); +---- +2 + +# list_avg alias with NULL row propagates correctly +query R +select list_avg(column1) from (values + (make_array(1.0, 2.0)), + (NULL) +) as t(column1); +---- +1.5 +NULL diff --git a/datafusion/sqllogictest/test_files/array_product.slt b/datafusion/sqllogictest/test_files/array_product.slt new file mode 100644 index 0000000000000..ba60360d1c1a9 --- /dev/null +++ b/datafusion/sqllogictest/test_files/array_product.slt @@ -0,0 +1,145 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +## array_product + +# Basic product of three floats +query R +select array_product([1.0, 2.0, 3.0]); +---- +6 + +# Negative values: signs multiply +query R +select array_product([-2.0, 3.0]); +---- +-6 + +# Single element returns itself +query R +select array_product([5.0]); +---- +5 + +# Zero element produces zero (no short-circuit; we still multiply) +query R +select array_product([0.0, 3.0, 4.0]); +---- +0 + +# NULL elements inside the list are skipped (SQL aggregate convention) +query R +select array_product([2.0, NULL, 3.0]); +---- +6 + +# All-NULL elements: no data to reduce, returns NULL +query R +select array_product([CAST(NULL AS DOUBLE), CAST(NULL AS DOUBLE)]); +---- +NULL + +# Bare NULL input returns NULL +query R +select array_product(NULL); +---- +NULL + +# Empty array: no data to reduce, returns NULL +query R +select array_product(arrow_cast(make_array(), 'List(Float64)')); +---- +NULL + +# LargeList input +query R +select array_product(arrow_cast([2.0, 3.0, 4.0], 'LargeList(Float64)')); +---- +24 + +# FixedSizeList input (coerced to List) +query R +select array_product(arrow_cast([2.0, 3.0, 4.0], 'FixedSizeList(3, Float64)')); +---- +24 + +# Float32 inner type (coerced to Float64) +query R +select array_product(arrow_cast([2.0, 3.0, 4.0], 'List(Float32)')); +---- +24 + +# Int64 inner type (coerced to Float64) +query R +select array_product(arrow_cast([2, 3, 4], 'List(Int64)')); +---- +24 + +# Integer literals (coerced to Float64) +query R +select array_product([2, 3, 4]); +---- +24 + +# Unsupported non-list input (plan error) +query error array_product does not support type +select array_product(1); + +# No arguments error +query error array_product function requires 1 argument, got 0 +select array_product(); + +# Multi-row query: normal row, NULL row, empty list, all-NULL elements, +# element-NULL skip, single zero +query R +select array_product(column1) from (values + (make_array(2.0, 3.0, 4.0)), + (NULL), + (arrow_cast(make_array(), 'List(Float64)')), + (make_array(CAST(NULL AS DOUBLE), CAST(NULL AS DOUBLE))), + (make_array(CAST(2.0 AS DOUBLE), CAST(NULL AS DOUBLE), CAST(5.0 AS DOUBLE))), + (make_array(0.0, 7.0)) +) as t(column1); +---- +24 +NULL +NULL +NULL +10 +0 + +# Return type is always Float64 (scalar, not List) +query RT +select array_product([2.0, 3.0]), arrow_typeof(array_product([2.0, 3.0])); +---- +6 Float64 + +# list_product alias produces the same result +query R +select list_product([2.0, 3.0, 4.0]); +---- +24 + +# list_product alias multi-row +query R +select list_product(column1) from (values + (make_array(2.0, 3.0)), + (NULL) +) as t(column1); +---- +6 +NULL \ No newline at end of file diff --git a/datafusion/sqllogictest/test_files/array_scale.slt b/datafusion/sqllogictest/test_files/array_scale.slt new file mode 100644 index 0000000000000..15d6cd6d98f68 --- /dev/null +++ b/datafusion/sqllogictest/test_files/array_scale.slt @@ -0,0 +1,192 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +## array_scale + +# General case: scale vector by positive scalar +query ? +select array_scale([1.0, 2.0, 3.0], 2.0); +---- +[2.0, 4.0, 6.0] + +# Scale by 1 returns the same array +query ? +select array_scale([1.0, 2.0, 3.0], 1.0); +---- +[1.0, 2.0, 3.0] + +# Scale by 0 returns zeros +query ? +select array_scale([1.0, 2.0, 3.0], 0.0); +---- +[0.0, 0.0, 0.0] + +# Scale by negative scalar +query ? +select array_scale([1.0, 2.0, 3.0], -1.0); +---- +[-1.0, -2.0, -3.0] + +# Scale by fractional scalar +query ? +select array_scale([2.0, 4.0, 6.0], 0.5); +---- +[1.0, 2.0, 3.0] + +# Single-element array +query ? +select array_scale([5.0], 3.0); +---- +[15.0] + +# Bare NULL array returns NULL +query ? +select array_scale(NULL, 2.0); +---- +NULL + +# NULL scalar returns NULL row (whole-row null because the scalar applies uniformly) +query ? +select array_scale([1.0, 2.0, 3.0], NULL); +---- +NULL + +# Both NULL returns NULL +query ? +select array_scale(NULL, NULL); +---- +NULL + +# NULL element in array propagates only to that position +query ? +select array_scale([1.0, NULL, 3.0], 2.0); +---- +[2.0, NULL, 6.0] + +# All-NULL elements with valid scalar: each position remains NULL +query ? +select array_scale([NULL, NULL], 5.0); +---- +[NULL, NULL] + +# LargeList support +query ? +select array_scale(arrow_cast([1.0, 2.0, 3.0], 'LargeList(Float64)'), 2.0); +---- +[2.0, 4.0, 6.0] + +# FixedSizeList input (coerced to List) +query ? +select array_scale(arrow_cast([1.0, 2.0, 3.0], 'FixedSizeList(3, Float64)'), 2.0); +---- +[2.0, 4.0, 6.0] + +# Float32 inner type (coerced to Float64) +query ? +select array_scale(arrow_cast([1.0, 2.0, 3.0], 'List(Float32)'), 2.0); +---- +[2.0, 4.0, 6.0] + +# Int64 inner type (coerced to Float64) +query ? +select array_scale(arrow_cast([1, 2, 3], 'List(Int64)'), 2); +---- +[2.0, 4.0, 6.0] + +# Integer literals on both sides (coerced to Float64) +query ? +select array_scale([1, 2, 3], 2); +---- +[2.0, 4.0, 6.0] + +# Integer scalar with Float64 list +query ? +select array_scale([1.0, 2.0, 3.0], 3); +---- +[3.0, 6.0, 9.0] + +# Unsupported non-numeric scalar (plan error) +query error array_scale second argument must be numeric +select array_scale([1.0, 2.0, 3.0], 'foo'); + +# Unsupported non-list first argument (plan error) +query error array_scale first argument must be a list type +select array_scale(1.0, 2.0); + +# Multi-row query: constant scalar broadcast across rows +query ? +select array_scale(column1, 2.0) from (values + (make_array(1.0, 2.0, 3.0)), + (make_array(0.0, 0.0)), + (make_array(1.0, NULL, 3.0)), + (NULL) +) as t(column1); +---- +[2.0, 4.0, 6.0] +[0.0, 0.0] +[2.0, NULL, 6.0] +NULL + +# Multi-row query: scalar from a column (varies per row) +query ? +select array_scale(column1, column2) from (values + (make_array(1.0, 2.0, 3.0), 2.0), + (make_array(1.0, 2.0), 0.5), + (make_array(1.0, 2.0), arrow_cast(NULL, 'Float64')), + (NULL, 3.0) +) as t(column1, column2); +---- +[2.0, 4.0, 6.0] +[0.5, 1.0] +NULL +NULL + +# Empty array: array_scale of an empty array yields an empty array +query ? +select array_scale(arrow_cast(make_array(), 'List(Float64)'), 2.0); +---- +[] + +# Wrong arity (zero args) +query error array_scale function requires 2 arguments, got 0 +select array_scale(); + +# Wrong arity (one arg) +query error array_scale function requires 2 arguments, got 1 +select array_scale([1.0, 2.0]); + +# Return type matches input list shape: List(Float64) input yields List(Float64) output +query ?T +select array_scale([1.0, 2.0], 3.0), arrow_typeof(array_scale([1.0, 2.0], 3.0)); +---- +[3.0, 6.0] List(Float64) + +# list_scale alias produces the same result +query ? +select list_scale([1.0, 2.0, 3.0], 2.0); +---- +[2.0, 4.0, 6.0] + +# list_scale alias with NULL scalar propagates correctly +query ? +select list_scale(column1, column2) from (values + (make_array(1.0, 2.0), 2.0), + (make_array(1.0, 2.0), arrow_cast(NULL, 'Float64')) +) as t(column1, column2); +---- +[2.0, 4.0] +NULL diff --git a/datafusion/sqllogictest/test_files/array_subtract.slt b/datafusion/sqllogictest/test_files/array_subtract.slt new file mode 100644 index 0000000000000..4a680c93aae95 --- /dev/null +++ b/datafusion/sqllogictest/test_files/array_subtract.slt @@ -0,0 +1,237 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +## array_subtract + +# Basic element-wise difference +query ? +select array_subtract([10.0, 20.0, 30.0], [1.0, 2.0, 3.0]); +---- +[9.0, 18.0, 27.0] + +# Negative components +query ? +select array_subtract([1.0, -2.0, 3.0], [-1.0, 2.0, -3.0]); +---- +[2.0, -4.0, 6.0] + +# Single-element arrays +query ? +select array_subtract([7.0], [5.0]); +---- +[2.0] + +# Bare NULL on left -> NULL row +query ? +select array_subtract(NULL, [1.0, 2.0]); +---- +NULL + +# Bare NULL on right -> NULL row +query ? +select array_subtract([1.0, 2.0], NULL); +---- +NULL + +# Both bare NULL -> NULL row +query ? +select array_subtract(NULL, NULL); +---- +NULL + +# NULL element on left propagates to that position only +query ? +select array_subtract([10.0, NULL, 30.0], [1.0, 2.0, 3.0]); +---- +[9.0, NULL, 27.0] + +# NULL element on right propagates to that position only +query ? +select array_subtract([10.0, 20.0, 30.0], [1.0, NULL, 3.0]); +---- +[9.0, NULL, 27.0] + +# NULL element on both sides at the same position +query ? +select array_subtract([10.0, NULL, 30.0], [1.0, NULL, 3.0]); +---- +[9.0, NULL, 27.0] + +# NULL elements at different positions both propagate +query ? +select array_subtract([10.0, NULL, 30.0], [NULL, 2.0, 3.0]); +---- +[NULL, NULL, 27.0] + +# Length mismatch is an exec error +query error array_subtract requires both list inputs to have the same length per row +select array_subtract([1.0, 2.0], [10.0, 20.0, 30.0]); + +# Empty arrays on both sides return empty array +query ? +select array_subtract(arrow_cast(make_array(), 'List(Float64)'), arrow_cast(make_array(), 'List(Float64)')); +---- +[] + +# Integer literals coerced to Float64 +query ? +select array_subtract([10, 20, 30], [1, 2, 3]); +---- +[9.0, 18.0, 27.0] + +# Mixed int + float literals coerced to Float64 +query ? +select array_subtract([1, 2, 3], [0.5, 0.5, 0.5]); +---- +[0.5, 1.5, 2.5] + +# LargeList input on both sides +query ? +select array_subtract( + arrow_cast([10.0, 20.0, 30.0], 'LargeList(Float64)'), + arrow_cast([1.0, 2.0, 3.0], 'LargeList(Float64)') +); +---- +[9.0, 18.0, 27.0] + +# Mixed List + LargeList -> both widened to LargeList +query ? +select array_subtract( + [10.0, 20.0, 30.0], + arrow_cast([1.0, 2.0, 3.0], 'LargeList(Float64)') +); +---- +[9.0, 18.0, 27.0] + +# FixedSizeList input (coerced to List) +query ? +select array_subtract( + arrow_cast([10.0, 20.0, 30.0], 'FixedSizeList(3, Float64)'), + arrow_cast([1.0, 2.0, 3.0], 'FixedSizeList(3, Float64)') +); +---- +[9.0, 18.0, 27.0] + +# Float32 inner type on one side +query ? +select array_subtract( + arrow_cast([10.0, 20.0, 30.0], 'List(Float32)'), + [1.0, 2.0, 3.0] +); +---- +[9.0, 18.0, 27.0] + +# Int64 inner type +query ? +select array_subtract( + arrow_cast([10, 20, 30], 'List(Int64)'), + arrow_cast([1, 2, 3], 'List(Int64)') +); +---- +[9.0, 18.0, 27.0] + +# Unsupported non-list input (plan error) +query error array_subtract does not support type +select array_subtract(1, [1.0, 2.0]); + +# Wrong arg count +query error array_subtract function requires 2 arguments, got 0 +select array_subtract(); + +query error array_subtract function requires 2 arguments, got 1 +select array_subtract([1.0, 2.0]); + +# Return type matches input variant +query ?T +select array_subtract([1.0, 2.0], [3.0, 4.0]), arrow_typeof(array_subtract([1.0, 2.0], [3.0, 4.0])); +---- +[-2.0, -2.0] List(Float64) + +# Multi-row query: normal row, NULL row, element-NULL row, length-matched row +query ? +select array_subtract(a, b) from (values + (make_array(10.0, 20.0, 30.0), make_array(1.0, 2.0, 3.0)), + (NULL, make_array(1.0, 2.0, 3.0)), + (make_array(1.0, 2.0, 3.0), NULL), + (make_array(10.0, NULL, 30.0), make_array(1.0, 2.0, 3.0)) +) as t(a, b); +---- +[9.0, 18.0, 27.0] +NULL +NULL +[9.0, NULL, 27.0] + +# list_subtract alias +query ? +select list_subtract([3.0, 4.0], [1.0, 2.0]); +---- +[2.0, 2.0] + +# list_subtract alias multi-row +query ? +select list_subtract(a, b) from (values + (make_array(10.0, 20.0), make_array(1.0, 2.0)), + (NULL, make_array(1.0, 2.0)) +) as t(a, b); +---- +[9.0, 18.0] +NULL + +# Decimal element types are coerced to Float64 (lossy) like other array-math UDFs +query ? +select array_subtract( + arrow_cast([10, 20, 30], 'List(Decimal128(10, 2))'), + arrow_cast([1, 2, 3], 'List(Decimal128(10, 2))') +); +---- +[9.0, 18.0, 27.0] + +# Explicit cast to DOUBLE works as the documented opt-in +query ? +select array_subtract( + arrow_cast(arrow_cast([10, 20, 30], 'List(Decimal128(10, 2))'), 'List(Float64)'), + [1.0, 2.0, 3.0] +); +---- +[9.0, 18.0, 27.0] + +# Chained array_subtract: result of inner call feeds the outer call +query ? +select array_subtract(array_subtract([100.0, 200.0, 300.0], [10.0, 20.0, 30.0]), [1.0, 2.0, 3.0]); +---- +[89.0, 178.0, 267.0] + +# Chained array_subtract propagates element-level NULLs through both layers +query ? +select array_subtract( + array_subtract([100.0, NULL, 300.0], [10.0, 20.0, 30.0]), + [1.0, 2.0, NULL] +); +---- +[89.0, NULL, NULL] + +# Chained array_subtract over multiple rows +query ? +select array_subtract(array_subtract(a, b), c) from (values + (make_array(100.0, 200.0), make_array(10.0, 20.0), make_array(1.0, 2.0)), + (NULL, make_array(1.0, 2.0), make_array(3.0, 4.0)), + (make_array(100.0, 200.0), make_array(10.0, NULL), make_array(1.0, 2.0)) +) as t(a, b, c); +---- +[89.0, 178.0] +NULL +[89.0, NULL] \ No newline at end of file diff --git a/datafusion/sqllogictest/test_files/array_sum.slt b/datafusion/sqllogictest/test_files/array_sum.slt new file mode 100644 index 0000000000000..823d767c48489 --- /dev/null +++ b/datafusion/sqllogictest/test_files/array_sum.slt @@ -0,0 +1,152 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +## array_sum + +# Basic case +query R +select array_sum([1.0, 2.0, 3.0]); +---- +6 + +# Single element +query R +select array_sum([5.0]); +---- +5 + +# Negative values +query R +select array_sum([-1.0, -2.0, -3.0]); +---- +-6 + +# Positive and negative cancel +query R +select array_sum([1.0, -1.0, 2.0, -2.0]); +---- +0 + +# Empty array returns NULL (matches PostgreSQL, DuckDB list_sum, SQL Standard SUM-of-empty-set) +query R +select array_sum(arrow_cast(make_array(), 'List(Float64)')); +---- +NULL + +# Bare NULL input returns NULL row +query R +select array_sum(NULL); +---- +NULL + +# NULL elements are skipped (SQL aggregate convention) +query R +select array_sum([1.0, NULL, 3.0]); +---- +4 + +# Single NULL among numeric: skip the NULL +query R +select array_sum([NULL, 10.0]); +---- +10 + +# All-NULL array returns NULL row (matches SQL SUM over all-NULL) +query R +select array_sum(arrow_cast([NULL, NULL], 'List(Float64)')); +---- +NULL + +# LargeList support +query R +select array_sum(arrow_cast([1.0, 2.0, 3.0], 'LargeList(Float64)')); +---- +6 + +# FixedSizeList input (coerced to List) +query R +select array_sum(arrow_cast([1.0, 2.0, 3.0], 'FixedSizeList(3, Float64)')); +---- +6 + +# Float32 inner type (coerced to Float64) +query R +select array_sum(arrow_cast([1.0, 2.0, 3.0], 'List(Float32)')); +---- +6 + +# Int64 inner type (coerced to Float64) +query R +select array_sum(arrow_cast([1, 2, 3], 'List(Int64)')); +---- +6 + +# Integer literals (coerced to Float64) +query R +select array_sum([1, 2, 3]); +---- +6 + +# Unsupported non-list input (plan error) +query error array_sum does not support type +select array_sum(1); + +# Multi-row query with mix of normal, single-element, NULL elements, empty, NULL row +query R +select array_sum(column1) from (values + (make_array(1.0, 2.0, 3.0)), + (make_array(0.0)), + (make_array(1.0, NULL, 4.0)), + (arrow_cast(make_array(), 'List(Float64)')), + (NULL) +) as t(column1); +---- +6 +0 +5 +NULL +NULL + +# Wrong arity (zero args) +query error array_sum function requires 1 argument, got 0 +select array_sum(); + +# Wrong arity (two args) +query error array_sum function requires 1 argument, got 2 +select array_sum([1.0], [2.0]); + +# Return type is Float64 +query RT +select array_sum([1.0, 2.0, 3.0]), arrow_typeof(array_sum([1.0, 2.0, 3.0])); +---- +6 Float64 + +# list_sum alias produces the same result +query R +select list_sum([1.0, 2.0, 3.0]); +---- +6 + +# list_sum alias with NULL row propagates correctly +query R +select list_sum(column1) from (values + (make_array(1.0, 2.0)), + (NULL) +) as t(column1); +---- +3 +NULL diff --git a/datafusion/sqllogictest/test_files/arrow_typeof.slt b/datafusion/sqllogictest/test_files/arrow_typeof.slt index e00909ad5fc59..17fcb7fa36ed7 100644 --- a/datafusion/sqllogictest/test_files/arrow_typeof.slt +++ b/datafusion/sqllogictest/test_files/arrow_typeof.slt @@ -397,12 +397,21 @@ select arrow_cast(null, 'FixedSizeList(1, Int64)'); ---- NULL -#TODO: arrow-rs doesn't support it yet -#query ? -#select arrow_cast('1', 'FixedSizeList(1, Int64)'); -#---- -#[1] +query ? +select arrow_cast([], 'FixedSizeList(0, Null)'); +---- +[] + +query ? rowsort +select arrow_cast(a, 'FixedSizeList(0, Null)') from values ([]), (NULL) t(a); +---- +NULL +[] +query ? +select arrow_cast('1', 'FixedSizeList(1, Int64)'); +---- +[1] query ? select arrow_cast([1], 'FixedSizeList(1, Int64)'); diff --git a/datafusion/sqllogictest/test_files/binary.slt b/datafusion/sqllogictest/test_files/binary.slt index a57c31547f08d..91a9449343d2a 100644 --- a/datafusion/sqllogictest/test_files/binary.slt +++ b/datafusion/sqllogictest/test_files/binary.slt @@ -281,7 +281,7 @@ SELECT cast(binary as varchar) as str, character_length(binary) as binary_len, cast(largebinary as varchar) as large_str, - character_length(binary) as largebinary_len + character_length(largebinary) as largebinary_len from t; ---- Foo 3 Foo 3 @@ -298,6 +298,20 @@ SELECT character_length(X'20'); query error Encountered non UTF\-8 data: invalid utf\-8 sequence of 1 bytes from index 0 SELECT character_length(X'c328'); +# reverse function +query TTTT +SELECT + cast(binary as varchar) as str, + reverse(binary) as binary_reversed, + cast(largebinary as varchar) as large_str, + reverse(largebinary) as largebinary_reversed +from t; +---- +Foo ooF Foo ooF +NULL NULL NULL NULL +Bar raB Bar raB +FooBar raBooF FooBar raBooF + # regexp_replace query TTTT SELECT @@ -347,14 +361,20 @@ SELECT x'636166c3a9' || arrow_cast(x'68656c6c6f', 'FixedSizeBinary(5)'), arrow_t query ?T SELECT arrow_cast(x'6361', 'FixedSizeBinary(2)') || arrow_cast(x'68656c6c6f', 'FixedSizeBinary(5)'), arrow_typeof(arrow_cast(x'6361', 'FixedSizeBinary(2)') || arrow_cast(x'68656c6c6f', 'FixedSizeBinary(5)')); ---- -636168656c6c6f Binary +636168656c6c6f FixedSizeBinary(7) -# Byte pipe operator is forbidden for mixed binary and text -query error DataFusion error: Error during planning: Cannot infer common string type for string concat operation Binary || Utf8 +# Byte pipe operator is allowed for mixed binary and text +query T SELECT x'c3a9' || 'hello'; +---- +éhello -query error DataFusion error: Error during planning: Cannot infer common string type for string concat operation Utf8 || LargeBinary +query T SELECT 'hello' || arrow_cast(arrow_cast('hello', 'Binary'), 'LargeBinary'); +---- +hellohello -query error DataFusion error: Error during planning: Cannot infer common string type for string concat operation Utf8 || BinaryView +query T SELECT 'hello' || arrow_cast(arrow_cast('hello', 'Binary'), 'BinaryView'); +---- +hellohello diff --git a/datafusion/sqllogictest/test_files/case.slt b/datafusion/sqllogictest/test_files/case.slt index 3953878ceb666..f7ae380242942 100644 --- a/datafusion/sqllogictest/test_files/case.slt +++ b/datafusion/sqllogictest/test_files/case.slt @@ -41,6 +41,19 @@ NULL 6 7 +# CASE nullability remains consistent through type coercion +query I +SELECT count(endpoint) +FROM ( + SELECT CASE + WHEN a IS NOT NULL THEN CAST(a AS BIGINT) + ELSE CAST(0 AS BIGINT) + END AS endpoint + FROM foo +) +---- +6 + # column or explicit null query I SELECT CASE WHEN a > 2 THEN b ELSE null END FROM foo diff --git a/datafusion/sqllogictest/test_files/clickbench.slt b/datafusion/sqllogictest/test_files/clickbench.slt index 91463c9c2bff8..4a1ef833c91db 100644 --- a/datafusion/sqllogictest/test_files/clickbench.slt +++ b/datafusion/sqllogictest/test_files/clickbench.slt @@ -213,8 +213,8 @@ logical_plan 06)----------TableScan: hits_raw projection=[AdvEngineID], partial_filters=[hits_raw.AdvEngineID != Int16(0)] physical_plan 01)SortPreservingMergeExec: [count(*)@1 DESC] -02)--SortExec: expr=[count(*)@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[AdvEngineID@0 as AdvEngineID, count(Int64(1))@1 as count(*)] +02)--ProjectionExec: expr=[AdvEngineID@0 as AdvEngineID, count(Int64(1))@1 as count(*)] +03)----SortExec: expr=[count(Int64(1))@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[AdvEngineID@0 as AdvEngineID], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([AdvEngineID@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[AdvEngineID@0 as AdvEngineID], aggr=[count(Int64(1))] @@ -239,8 +239,8 @@ logical_plan 06)----------TableScan: hits_raw projection=[RegionID, UserID] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[u@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[RegionID@0 as RegionID, count(alias1)@1 as u] +02)--ProjectionExec: expr=[RegionID@0 as RegionID, count(alias1)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([RegionID@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID], aggr=[count(alias1)] @@ -269,8 +269,8 @@ logical_plan 05)--------TableScan: hits_raw projection=[RegionID, UserID, ResolutionWidth, AdvEngineID] physical_plan 01)SortPreservingMergeExec: [c@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[RegionID@0 as RegionID, sum(hits.AdvEngineID)@1 as sum(hits.AdvEngineID), count(Int64(1))@2 as c, avg(hits.ResolutionWidth)@3 as avg(hits.ResolutionWidth), count(DISTINCT hits.UserID)@4 as count(DISTINCT hits.UserID)] +02)--ProjectionExec: expr=[RegionID@0 as RegionID, sum(hits.AdvEngineID)@1 as sum(hits.AdvEngineID), count(Int64(1))@2 as c, avg(hits.ResolutionWidth)@3 as avg(hits.ResolutionWidth), count(DISTINCT hits.UserID)@4 as count(DISTINCT hits.UserID)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[RegionID@0 as RegionID], aggr=[sum(hits.AdvEngineID), count(Int64(1)), avg(hits.ResolutionWidth), count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([RegionID@0], 4), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[RegionID@0 as RegionID], aggr=[sum(hits.AdvEngineID), count(Int64(1)), avg(hits.ResolutionWidth), count(DISTINCT hits.UserID)] @@ -298,15 +298,15 @@ logical_plan 07)------------TableScan: hits_raw projection=[UserID, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[u@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[MobilePhoneModel@0 as MobilePhoneModel, count(alias1)@1 as u] +02)--ProjectionExec: expr=[MobilePhoneModel@0 as MobilePhoneModel, count(alias1)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([MobilePhoneModel@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[MobilePhoneModel@0 as MobilePhoneModel], aggr=[count(alias1)] 07)------------AggregateExec: mode=FinalPartitioned, gby=[MobilePhoneModel@0 as MobilePhoneModel, alias1@1 as alias1], aggr=[] 08)--------------RepartitionExec: partitioning=Hash([MobilePhoneModel@0, alias1@1], 4), input_partitions=4 09)----------------AggregateExec: mode=Partial, gby=[MobilePhoneModel@1 as MobilePhoneModel, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: MobilePhoneModel@1 != +10)------------------FilterExec: MobilePhoneModel@1 != 11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 12)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] @@ -328,15 +328,15 @@ logical_plan 07)------------TableScan: hits_raw projection=[UserID, MobilePhone, MobilePhoneModel], partial_filters=[hits_raw.MobilePhoneModel != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[u@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, count(alias1)@2 as u] +02)--ProjectionExec: expr=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, count(alias1)@2 as u] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@2 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([MobilePhone@0, MobilePhoneModel@1], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel], aggr=[count(alias1)] 07)------------AggregateExec: mode=FinalPartitioned, gby=[MobilePhone@0 as MobilePhone, MobilePhoneModel@1 as MobilePhoneModel, alias1@2 as alias1], aggr=[] 08)--------------RepartitionExec: partitioning=Hash([MobilePhone@0, MobilePhoneModel@1, alias1@2], 4), input_partitions=4 09)----------------AggregateExec: mode=Partial, gby=[MobilePhone@1 as MobilePhone, MobilePhoneModel@2 as MobilePhoneModel, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: MobilePhoneModel@2 != +10)------------------FilterExec: MobilePhoneModel@2 != 11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 12)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, MobilePhone, MobilePhoneModel], file_type=parquet, predicate=MobilePhoneModel@34 != , pruning_predicate=MobilePhoneModel_null_count@2 != row_count@3 AND (MobilePhoneModel_min@0 != OR != MobilePhoneModel_max@1), required_guarantees=[MobilePhoneModel not in ()] @@ -357,12 +357,12 @@ logical_plan 06)----------TableScan: hits_raw projection=[SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [c@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(Int64(1))@1 as c] +02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(Int64(1))@1 as c] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(Int64(1))] -07)------------FilterExec: SearchPhrase@0 != +07)------------FilterExec: SearchPhrase@0 != 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] @@ -384,15 +384,15 @@ logical_plan 07)------------TableScan: hits_raw projection=[UserID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [u@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[u@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(alias1)@1 as u] +02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, count(alias1)@1 as u] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@0 as SearchPhrase], aggr=[count(alias1)] 07)------------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase, alias1@1 as alias1], aggr=[] 08)--------------RepartitionExec: partitioning=Hash([SearchPhrase@0, alias1@1], 4), input_partitions=4 09)----------------AggregateExec: mode=Partial, gby=[SearchPhrase@1 as SearchPhrase, UserID@0 as alias1], aggr=[] -10)------------------FilterExec: SearchPhrase@1 != +10)------------------FilterExec: SearchPhrase@1 != 11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 12)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[UserID, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] @@ -413,12 +413,12 @@ logical_plan 06)----------TableScan: hits_raw projection=[SearchEngineID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [c@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[SearchEngineID@0 as SearchEngineID, SearchPhrase@1 as SearchPhrase, count(Int64(1))@2 as c] +02)--ProjectionExec: expr=[SearchEngineID@0 as SearchEngineID, SearchPhrase@1 as SearchPhrase, count(Int64(1))@2 as c] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[SearchEngineID@0 as SearchEngineID, SearchPhrase@1 as SearchPhrase], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([SearchEngineID@0, SearchPhrase@1], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[SearchEngineID@0 as SearchEngineID, SearchPhrase@1 as SearchPhrase], aggr=[count(Int64(1))] -07)------------FilterExec: SearchPhrase@1 != +07)------------FilterExec: SearchPhrase@1 != 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[SearchEngineID, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] @@ -438,8 +438,8 @@ logical_plan 05)--------TableScan: hits_raw projection=[UserID] physical_plan 01)SortPreservingMergeExec: [count(*)@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[count(*)@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[UserID@0 as UserID, count(Int64(1))@1 as count(*)] +02)--ProjectionExec: expr=[UserID@0 as UserID, count(Int64(1))@1 as count(*)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[UserID@0 as UserID], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([UserID@0], 4), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[UserID@0 as UserID], aggr=[count(Int64(1))] @@ -466,8 +466,8 @@ logical_plan 05)--------TableScan: hits_raw projection=[UserID, SearchPhrase] physical_plan 01)SortPreservingMergeExec: [count(*)@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[count(*)@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[UserID@0 as UserID, SearchPhrase@1 as SearchPhrase, count(Int64(1))@2 as count(*)] +02)--ProjectionExec: expr=[UserID@0 as UserID, SearchPhrase@1 as SearchPhrase, count(Int64(1))@2 as count(*)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[UserID@0 as UserID, SearchPhrase@1 as SearchPhrase], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([UserID@0, SearchPhrase@1], 4), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[UserID@0 as UserID, SearchPhrase@1 as SearchPhrase], aggr=[count(Int64(1))] @@ -521,8 +521,8 @@ logical_plan 05)--------TableScan: hits_raw projection=[EventTime, UserID, SearchPhrase] physical_plan 01)SortPreservingMergeExec: [count(*)@3 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[count(*)@3 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[UserID@0 as UserID, date_part(Utf8("MINUTE"),to_timestamp_seconds(hits.EventTime))@1 as m, SearchPhrase@2 as SearchPhrase, count(Int64(1))@3 as count(*)] +02)--ProjectionExec: expr=[UserID@0 as UserID, date_part(Utf8("MINUTE"),to_timestamp_seconds(hits.EventTime))@1 as m, SearchPhrase@2 as SearchPhrase, count(Int64(1))@3 as count(*)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@3 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[UserID@0 as UserID, date_part(Utf8("MINUTE"),to_timestamp_seconds(hits.EventTime))@1 as date_part(Utf8("MINUTE"),to_timestamp_seconds(hits.EventTime)), SearchPhrase@2 as SearchPhrase], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([UserID@0, date_part(Utf8("MINUTE"),to_timestamp_seconds(hits.EventTime))@1, SearchPhrase@2], 4), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[UserID@1 as UserID, date_part(MINUTE, to_timestamp_seconds(EventTime@0)) as date_part(Utf8("MINUTE"),to_timestamp_seconds(hits.EventTime)), SearchPhrase@2 as SearchPhrase], aggr=[count(Int64(1))] @@ -593,18 +593,18 @@ logical_plan 02)--Projection: hits.SearchPhrase, min(hits.URL), count(Int64(1)) AS count(*) AS c 03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[min(hits.URL), count(Int64(1))]] 04)------SubqueryAlias: hits -05)--------Filter: hits_raw.URL LIKE Utf8View("%google%") AND hits_raw.SearchPhrase != Utf8View("") -06)----------TableScan: hits_raw projection=[URL, SearchPhrase], partial_filters=[hits_raw.URL LIKE Utf8View("%google%"), hits_raw.SearchPhrase != Utf8View("")] +05)--------Filter: hits_raw.SearchPhrase != Utf8View("") AND hits_raw.URL LIKE Utf8View("%google%") +06)----------TableScan: hits_raw projection=[URL, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View(""), hits_raw.URL LIKE Utf8View("%google%")] physical_plan 01)SortPreservingMergeExec: [c@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, min(hits.URL)@1 as min(hits.URL), count(Int64(1))@2 as c] +02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, min(hits.URL)@1 as min(hits.URL), count(Int64(1))@2 as c] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[min(hits.URL), count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@1 as SearchPhrase], aggr=[min(hits.URL), count(Int64(1))] -07)------------FilterExec: URL@0 LIKE %google% AND SearchPhrase@1 != +07)------------FilterExec: SearchPhrase@1 != AND URL@0 LIKE %google% 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[URL, SearchPhrase], file_type=parquet, predicate=URL@13 LIKE %google% AND SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@4 != row_count@5 AND (SearchPhrase_min@2 != OR != SearchPhrase_max@3), required_guarantees=[SearchPhrase not in ()] +09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[URL, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND URL@13 LIKE %google%, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query TTI SELECT "SearchPhrase", MIN("URL"), COUNT(*) AS c FROM hits WHERE "URL" LIKE '%google%' AND "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; @@ -619,18 +619,18 @@ logical_plan 02)--Projection: hits.SearchPhrase, min(hits.URL), min(hits.Title), count(Int64(1)) AS count(*) AS c, count(DISTINCT hits.UserID) 03)----Aggregate: groupBy=[[hits.SearchPhrase]], aggr=[[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)]] 04)------SubqueryAlias: hits -05)--------Filter: hits_raw.Title LIKE Utf8View("%Google%") AND hits_raw.URL NOT LIKE Utf8View("%.google.%") AND hits_raw.SearchPhrase != Utf8View("") -06)----------TableScan: hits_raw projection=[Title, UserID, URL, SearchPhrase], partial_filters=[hits_raw.Title LIKE Utf8View("%Google%"), hits_raw.URL NOT LIKE Utf8View("%.google.%"), hits_raw.SearchPhrase != Utf8View("")] +05)--------Filter: hits_raw.SearchPhrase != Utf8View("") AND hits_raw.Title LIKE Utf8View("%Google%") AND hits_raw.URL NOT LIKE Utf8View("%.google.%") +06)----------TableScan: hits_raw projection=[Title, UserID, URL, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View(""), hits_raw.Title LIKE Utf8View("%Google%"), hits_raw.URL NOT LIKE Utf8View("%.google.%")] physical_plan 01)SortPreservingMergeExec: [c@3 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@3 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, min(hits.URL)@1 as min(hits.URL), min(hits.Title)@2 as min(hits.Title), count(Int64(1))@3 as c, count(DISTINCT hits.UserID)@4 as count(DISTINCT hits.UserID)] +02)--ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase, min(hits.URL)@1 as min(hits.URL), min(hits.Title)@2 as min(hits.Title), count(Int64(1))@3 as c, count(DISTINCT hits.UserID)@4 as count(DISTINCT hits.UserID)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@3 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[SearchPhrase@0 as SearchPhrase], aggr=[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)] 05)--------RepartitionExec: partitioning=Hash([SearchPhrase@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[SearchPhrase@3 as SearchPhrase], aggr=[min(hits.URL), min(hits.Title), count(Int64(1)), count(DISTINCT hits.UserID)] -07)------------FilterExec: Title@0 LIKE %Google% AND URL@2 NOT LIKE %.google.% AND SearchPhrase@3 != +07)------------FilterExec: SearchPhrase@3 != AND Title@0 LIKE %Google% AND URL@2 NOT LIKE %.google.% 08)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[Title, UserID, URL, SearchPhrase], file_type=parquet, predicate=Title@2 LIKE %Google% AND URL@13 NOT LIKE %.google.% AND SearchPhrase@39 != , pruning_predicate=SearchPhrase_null_count@6 != row_count@7 AND (SearchPhrase_min@4 != OR != SearchPhrase_max@5), required_guarantees=[SearchPhrase not in ()] +09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[Title, UserID, URL, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND Title@2 LIKE %Google% AND URL@13 NOT LIKE %.google.%, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query TTTII SELECT "SearchPhrase", MIN("URL"), MIN("Title"), COUNT(*) AS c, COUNT(DISTINCT "UserID") FROM hits WHERE "Title" LIKE '%Google%' AND "URL" NOT LIKE '%.google.%' AND "SearchPhrase" <> '' GROUP BY "SearchPhrase" ORDER BY c DESC LIMIT 10; @@ -652,7 +652,7 @@ physical_plan 03)----ProjectionExec: expr=[WatchID@0 as WatchID, JavaEnable@1 as JavaEnable, Title@2 as Title, GoodEvent@3 as GoodEvent, EventTime@4 as EventTime, CounterID@6 as CounterID, ClientIP@7 as ClientIP, RegionID@8 as RegionID, UserID@9 as UserID, CounterClass@10 as CounterClass, OS@11 as OS, UserAgent@12 as UserAgent, URL@13 as URL, Referer@14 as Referer, IsRefresh@15 as IsRefresh, RefererCategoryID@16 as RefererCategoryID, RefererRegionID@17 as RefererRegionID, URLCategoryID@18 as URLCategoryID, URLRegionID@19 as URLRegionID, ResolutionWidth@20 as ResolutionWidth, ResolutionHeight@21 as ResolutionHeight, ResolutionDepth@22 as ResolutionDepth, FlashMajor@23 as FlashMajor, FlashMinor@24 as FlashMinor, FlashMinor2@25 as FlashMinor2, NetMajor@26 as NetMajor, NetMinor@27 as NetMinor, UserAgentMajor@28 as UserAgentMajor, UserAgentMinor@29 as UserAgentMinor, CookieEnable@30 as CookieEnable, JavascriptEnable@31 as JavascriptEnable, IsMobile@32 as IsMobile, MobilePhone@33 as MobilePhone, MobilePhoneModel@34 as MobilePhoneModel, Params@35 as Params, IPNetworkID@36 as IPNetworkID, TraficSourceID@37 as TraficSourceID, SearchEngineID@38 as SearchEngineID, SearchPhrase@39 as SearchPhrase, AdvEngineID@40 as AdvEngineID, IsArtifical@41 as IsArtifical, WindowClientWidth@42 as WindowClientWidth, WindowClientHeight@43 as WindowClientHeight, ClientTimeZone@44 as ClientTimeZone, ClientEventTime@45 as ClientEventTime, SilverlightVersion1@46 as SilverlightVersion1, SilverlightVersion2@47 as SilverlightVersion2, SilverlightVersion3@48 as SilverlightVersion3, SilverlightVersion4@49 as SilverlightVersion4, PageCharset@50 as PageCharset, CodeVersion@51 as CodeVersion, IsLink@52 as IsLink, IsDownload@53 as IsDownload, IsNotBounce@54 as IsNotBounce, FUniqID@55 as FUniqID, OriginalURL@56 as OriginalURL, HID@57 as HID, IsOldCounter@58 as IsOldCounter, IsEvent@59 as IsEvent, IsParameter@60 as IsParameter, DontCountHits@61 as DontCountHits, WithHash@62 as WithHash, HitColor@63 as HitColor, LocalEventTime@64 as LocalEventTime, Age@65 as Age, Sex@66 as Sex, Income@67 as Income, Interests@68 as Interests, Robotness@69 as Robotness, RemoteIP@70 as RemoteIP, WindowName@71 as WindowName, OpenerName@72 as OpenerName, HistoryLength@73 as HistoryLength, BrowserLanguage@74 as BrowserLanguage, BrowserCountry@75 as BrowserCountry, SocialNetwork@76 as SocialNetwork, SocialAction@77 as SocialAction, HTTPError@78 as HTTPError, SendTiming@79 as SendTiming, DNSTiming@80 as DNSTiming, ConnectTiming@81 as ConnectTiming, ResponseStartTiming@82 as ResponseStartTiming, ResponseEndTiming@83 as ResponseEndTiming, FetchTiming@84 as FetchTiming, SocialSourceNetworkID@85 as SocialSourceNetworkID, SocialSourcePage@86 as SocialSourcePage, ParamPrice@87 as ParamPrice, ParamOrderID@88 as ParamOrderID, ParamCurrency@89 as ParamCurrency, ParamCurrencyID@90 as ParamCurrencyID, OpenstatServiceName@91 as OpenstatServiceName, OpenstatCampaignID@92 as OpenstatCampaignID, OpenstatAdID@93 as OpenstatAdID, OpenstatSourceID@94 as OpenstatSourceID, UTMSource@95 as UTMSource, UTMMedium@96 as UTMMedium, UTMCampaign@97 as UTMCampaign, UTMContent@98 as UTMContent, UTMTerm@99 as UTMTerm, FromTag@100 as FromTag, HasGCLID@101 as HasGCLID, RefererHash@102 as RefererHash, URLHash@103 as URLHash, CLID@104 as CLID, CAST(CAST(EventDate@5 AS Int32) AS Date32) as EventDate] 04)------FilterExec: URL@13 LIKE %google% 05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[WatchID, JavaEnable, Title, GoodEvent, EventTime, EventDate, CounterID, ClientIP, RegionID, UserID, CounterClass, OS, UserAgent, URL, Referer, IsRefresh, RefererCategoryID, RefererRegionID, URLCategoryID, URLRegionID, ResolutionWidth, ResolutionHeight, ResolutionDepth, FlashMajor, FlashMinor, FlashMinor2, NetMajor, NetMinor, UserAgentMajor, UserAgentMinor, CookieEnable, JavascriptEnable, IsMobile, MobilePhone, MobilePhoneModel, Params, IPNetworkID, TraficSourceID, SearchEngineID, SearchPhrase, AdvEngineID, IsArtifical, WindowClientWidth, WindowClientHeight, ClientTimeZone, ClientEventTime, SilverlightVersion1, SilverlightVersion2, SilverlightVersion3, SilverlightVersion4, PageCharset, CodeVersion, IsLink, IsDownload, IsNotBounce, FUniqID, OriginalURL, HID, IsOldCounter, IsEvent, IsParameter, DontCountHits, WithHash, HitColor, LocalEventTime, Age, Sex, Income, Interests, Robotness, RemoteIP, WindowName, OpenerName, HistoryLength, BrowserLanguage, BrowserCountry, SocialNetwork, SocialAction, HTTPError, SendTiming, DNSTiming, ConnectTiming, ResponseStartTiming, ResponseEndTiming, FetchTiming, SocialSourceNetworkID, SocialSourcePage, ParamPrice, ParamOrderID, ParamCurrency, ParamCurrencyID, OpenstatServiceName, OpenstatCampaignID, OpenstatAdID, OpenstatSourceID, UTMSource, UTMMedium, UTMCampaign, UTMContent, UTMTerm, FromTag, HasGCLID, RefererHash, URLHash, CLID], file_type=parquet, predicate=URL@13 LIKE %google% AND DynamicFilter [ empty ] +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[WatchID, JavaEnable, Title, GoodEvent, EventTime, EventDate, CounterID, ClientIP, RegionID, UserID, CounterClass, OS, UserAgent, URL, Referer, IsRefresh, RefererCategoryID, RefererRegionID, URLCategoryID, URLRegionID, ResolutionWidth, ResolutionHeight, ResolutionDepth, FlashMajor, FlashMinor, FlashMinor2, NetMajor, NetMinor, UserAgentMajor, UserAgentMinor, CookieEnable, JavascriptEnable, IsMobile, MobilePhone, MobilePhoneModel, Params, IPNetworkID, TraficSourceID, SearchEngineID, SearchPhrase, AdvEngineID, IsArtifical, WindowClientWidth, WindowClientHeight, ClientTimeZone, ClientEventTime, SilverlightVersion1, SilverlightVersion2, SilverlightVersion3, SilverlightVersion4, PageCharset, CodeVersion, IsLink, IsDownload, IsNotBounce, FUniqID, OriginalURL, HID, IsOldCounter, IsEvent, IsParameter, DontCountHits, WithHash, HitColor, LocalEventTime, Age, Sex, Income, Interests, Robotness, RemoteIP, WindowName, OpenerName, HistoryLength, BrowserLanguage, BrowserCountry, SocialNetwork, SocialAction, HTTPError, SendTiming, DNSTiming, ConnectTiming, ResponseStartTiming, ResponseEndTiming, FetchTiming, SocialSourceNetworkID, SocialSourcePage, ParamPrice, ParamOrderID, ParamCurrency, ParamCurrencyID, OpenstatServiceName, OpenstatCampaignID, OpenstatAdID, OpenstatSourceID, UTMSource, UTMMedium, UTMCampaign, UTMContent, UTMTerm, FromTag, HasGCLID, RefererHash, URLHash, CLID], file_type=parquet, predicate=URL@13 LIKE %google% AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible query IITIIIIIIIIITTIIIIIIIIIITIIITIIIITTIIITIIIIIIIIIITIIIIITIIIIIITIIIIIIIIIITTTTIIIIIIIITITTITTTTTTTTTTIIIID SELECT * FROM hits WHERE "URL" LIKE '%google%' ORDER BY "EventTime" LIMIT 10; @@ -672,10 +672,11 @@ logical_plan physical_plan 01)ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase] 02)--SortPreservingMergeExec: [EventTime@1 ASC NULLS LAST], fetch=10 -03)----SortExec: TopK(fetch=10), expr=[EventTime@1 ASC NULLS LAST], preserve_partitioning=[true] -04)------FilterExec: SearchPhrase@1 != , projection=[SearchPhrase@1, EventTime@0] -05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[EventTime, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND DynamicFilter [ empty ], pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] +03)----ProjectionExec: expr=[SearchPhrase@1 as SearchPhrase, EventTime@0 as EventTime] +04)------SortExec: TopK(fetch=10), expr=[EventTime@0 ASC NULLS LAST], preserve_partitioning=[true] +05)--------FilterExec: SearchPhrase@1 != +06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[EventTime, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query T SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "EventTime" LIMIT 10; @@ -693,9 +694,9 @@ logical_plan physical_plan 01)SortPreservingMergeExec: [SearchPhrase@0 ASC NULLS LAST], fetch=10 02)--SortExec: TopK(fetch=10), expr=[SearchPhrase@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----FilterExec: SearchPhrase@0 != +03)----FilterExec: SearchPhrase@0 != 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND DynamicFilter [ empty ], pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] +05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query T SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "SearchPhrase" LIMIT 10; @@ -715,10 +716,11 @@ logical_plan physical_plan 01)ProjectionExec: expr=[SearchPhrase@0 as SearchPhrase] 02)--SortPreservingMergeExec: [EventTime@1 ASC NULLS LAST, SearchPhrase@0 ASC NULLS LAST], fetch=10 -03)----SortExec: TopK(fetch=10), expr=[EventTime@1 ASC NULLS LAST, SearchPhrase@0 ASC NULLS LAST], preserve_partitioning=[true] -04)------FilterExec: SearchPhrase@1 != , projection=[SearchPhrase@1, EventTime@0] -05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[EventTime, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND DynamicFilter [ empty ], pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] +03)----ProjectionExec: expr=[SearchPhrase@1 as SearchPhrase, EventTime@0 as EventTime] +04)------SortExec: TopK(fetch=10), expr=[EventTime@0 ASC NULLS LAST, SearchPhrase@1 ASC NULLS LAST], preserve_partitioning=[true] +05)--------FilterExec: SearchPhrase@1 != +06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[EventTime, SearchPhrase], file_type=parquet, predicate=SearchPhrase@39 != AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=SearchPhrase_null_count@2 != row_count@3 AND (SearchPhrase_min@0 != OR != SearchPhrase_max@1), required_guarantees=[SearchPhrase not in ()] query T SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "EventTime", "SearchPhrase" LIMIT 10; @@ -726,58 +728,58 @@ SELECT "SearchPhrase" FROM hits WHERE "SearchPhrase" <> '' ORDER BY "EventTime", ## Q27 query TT -EXPLAIN SELECT "CounterID", AVG(length("URL")) AS l, COUNT(*) AS c FROM hits WHERE "URL" <> '' GROUP BY "CounterID" HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; +EXPLAIN SELECT "CounterID", AVG(octet_length("URL")) AS l, COUNT(*) AS c FROM hits WHERE "URL" <> '' GROUP BY "CounterID" HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; ---- logical_plan 01)Sort: l DESC NULLS FIRST, fetch=25 -02)--Projection: hits.CounterID, avg(length(hits.URL)) AS l, count(Int64(1)) AS count(*) AS c +02)--Projection: hits.CounterID, avg(octet_length(hits.URL)) AS l, count(Int64(1)) AS count(*) AS c 03)----Filter: count(Int64(1)) > Int64(100000) -04)------Aggregate: groupBy=[[hits.CounterID]], aggr=[[avg(CAST(character_length(hits.URL) AS length(hits.URL) AS Float64)), count(Int64(1))]] +04)------Aggregate: groupBy=[[hits.CounterID]], aggr=[[avg(CAST(octet_length(hits.URL) AS Float64)), count(Int64(1))]] 05)--------SubqueryAlias: hits 06)----------Filter: hits_raw.URL != Utf8View("") 07)------------TableScan: hits_raw projection=[CounterID, URL], partial_filters=[hits_raw.URL != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [l@1 DESC], fetch=25 -02)--SortExec: TopK(fetch=25), expr=[l@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[CounterID@0 as CounterID, avg(length(hits.URL))@1 as l, count(Int64(1))@2 as c] +02)--ProjectionExec: expr=[CounterID@0 as CounterID, avg(octet_length(hits.URL))@1 as l, count(Int64(1))@2 as c] +03)----SortExec: TopK(fetch=25), expr=[avg(octet_length(hits.URL))@1 DESC], preserve_partitioning=[true] 04)------FilterExec: count(Int64(1))@2 > 100000 -05)--------AggregateExec: mode=FinalPartitioned, gby=[CounterID@0 as CounterID], aggr=[avg(length(hits.URL)), count(Int64(1))] +05)--------AggregateExec: mode=FinalPartitioned, gby=[CounterID@0 as CounterID], aggr=[avg(octet_length(hits.URL)), count(Int64(1))] 06)----------RepartitionExec: partitioning=Hash([CounterID@0], 4), input_partitions=4 -07)------------AggregateExec: mode=Partial, gby=[CounterID@0 as CounterID], aggr=[avg(length(hits.URL)), count(Int64(1))] -08)--------------FilterExec: URL@1 != +07)------------AggregateExec: mode=Partial, gby=[CounterID@0 as CounterID], aggr=[avg(octet_length(hits.URL)), count(Int64(1))] +08)--------------FilterExec: URL@1 != 09)----------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 10)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[CounterID, URL], file_type=parquet, predicate=URL@13 != , pruning_predicate=URL_null_count@2 != row_count@3 AND (URL_min@0 != OR != URL_max@1), required_guarantees=[URL not in ()] query IRI -SELECT "CounterID", AVG(length("URL")) AS l, COUNT(*) AS c FROM hits WHERE "URL" <> '' GROUP BY "CounterID" HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; +SELECT "CounterID", AVG(octet_length("URL")) AS l, COUNT(*) AS c FROM hits WHERE "URL" <> '' GROUP BY "CounterID" HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; ---- ## Q28 query TT -EXPLAIN SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; +EXPLAIN SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(octet_length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; ---- logical_plan 01)Sort: l DESC NULLS FIRST, fetch=25 -02)--Projection: regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1")) AS k, avg(length(hits.Referer)) AS l, count(Int64(1)) AS count(*) AS c, min(hits.Referer) +02)--Projection: regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1")) AS k, avg(octet_length(hits.Referer)) AS l, count(Int64(1)) AS count(*) AS c, min(hits.Referer) 03)----Filter: count(Int64(1)) > Int64(100000) -04)------Aggregate: groupBy=[[regexp_replace(hits.Referer, Utf8View("^https?://(?:www\.)?([^/]+)/.*$"), Utf8View("\1")) AS regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))]], aggr=[[avg(CAST(character_length(hits.Referer) AS length(hits.Referer) AS Float64)), count(Int64(1)), min(hits.Referer)]] +04)------Aggregate: groupBy=[[regexp_replace(hits.Referer, Utf8View("^https?://(?:www\.)?([^/]+)/.*$"), Utf8View("\1")) AS regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))]], aggr=[[avg(CAST(octet_length(hits.Referer) AS Float64)), count(Int64(1)), min(hits.Referer)]] 05)--------SubqueryAlias: hits 06)----------Filter: hits_raw.Referer != Utf8View("") 07)------------TableScan: hits_raw projection=[Referer], partial_filters=[hits_raw.Referer != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [l@1 DESC], fetch=25 -02)--SortExec: TopK(fetch=25), expr=[l@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))@0 as k, avg(length(hits.Referer))@1 as l, count(Int64(1))@2 as c, min(hits.Referer)@3 as min(hits.Referer)] +02)--ProjectionExec: expr=[regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))@0 as k, avg(octet_length(hits.Referer))@1 as l, count(Int64(1))@2 as c, min(hits.Referer)@3 as min(hits.Referer)] +03)----SortExec: TopK(fetch=25), expr=[avg(octet_length(hits.Referer))@1 DESC], preserve_partitioning=[true] 04)------FilterExec: count(Int64(1))@2 > 100000 -05)--------AggregateExec: mode=FinalPartitioned, gby=[regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))@0 as regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))], aggr=[avg(length(hits.Referer)), count(Int64(1)), min(hits.Referer)] +05)--------AggregateExec: mode=FinalPartitioned, gby=[regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))@0 as regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))], aggr=[avg(octet_length(hits.Referer)), count(Int64(1)), min(hits.Referer)] 06)----------RepartitionExec: partitioning=Hash([regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))@0], 4), input_partitions=4 -07)------------AggregateExec: mode=Partial, gby=[regexp_replace(Referer@0, ^https?://(?:www\.)?([^/]+)/.*$, \1) as regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))], aggr=[avg(length(hits.Referer)), count(Int64(1)), min(hits.Referer)] -08)--------------FilterExec: Referer@0 != +07)------------AggregateExec: mode=Partial, gby=[regexp_replace(Referer@0, ^https?://(?:www\.)?([^/]+)/.*$, \1) as regexp_replace(hits.Referer,Utf8("^https?://(?:www\.)?([^/]+)/.*$"),Utf8("\1"))], aggr=[avg(octet_length(hits.Referer)), count(Int64(1)), min(hits.Referer)] +08)--------------FilterExec: Referer@0 != 09)----------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 10)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/clickbench_hits_10.parquet]]}, projection=[Referer], file_type=parquet, predicate=Referer@14 != , pruning_predicate=Referer_null_count@2 != row_count@3 AND (Referer_min@0 != OR != Referer_max@1), required_guarantees=[Referer not in ()] query TRIT -SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; +SELECT REGEXP_REPLACE("Referer", '^https?://(?:www\.)?([^/]+)/.*$', '\1') AS k, AVG(octet_length("Referer")) AS l, COUNT(*) AS c, MIN("Referer") FROM hits WHERE "Referer" <> '' GROUP BY k HAVING COUNT(*) > 100000 ORDER BY l DESC LIMIT 25; ---- ## Q29 @@ -815,8 +817,8 @@ logical_plan 07)------------TableScan: hits_raw projection=[ClientIP, IsRefresh, ResolutionWidth, SearchEngineID, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [c@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[SearchEngineID@0 as SearchEngineID, ClientIP@1 as ClientIP, count(Int64(1))@2 as c, sum(hits.IsRefresh)@3 as sum(hits.IsRefresh), avg(hits.ResolutionWidth)@4 as avg(hits.ResolutionWidth)] +02)--ProjectionExec: expr=[SearchEngineID@0 as SearchEngineID, ClientIP@1 as ClientIP, count(Int64(1))@2 as c, sum(hits.IsRefresh)@3 as sum(hits.IsRefresh), avg(hits.ResolutionWidth)@4 as avg(hits.ResolutionWidth)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[SearchEngineID@0 as SearchEngineID, ClientIP@1 as ClientIP], aggr=[count(Int64(1)), sum(hits.IsRefresh), avg(hits.ResolutionWidth)] 05)--------RepartitionExec: partitioning=Hash([SearchEngineID@0, ClientIP@1], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[SearchEngineID@3 as SearchEngineID, ClientIP@0 as ClientIP], aggr=[count(Int64(1)), sum(hits.IsRefresh), avg(hits.ResolutionWidth)] @@ -842,8 +844,8 @@ logical_plan 07)------------TableScan: hits_raw projection=[WatchID, ClientIP, IsRefresh, ResolutionWidth, SearchPhrase], partial_filters=[hits_raw.SearchPhrase != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [c@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[WatchID@0 as WatchID, ClientIP@1 as ClientIP, count(Int64(1))@2 as c, sum(hits.IsRefresh)@3 as sum(hits.IsRefresh), avg(hits.ResolutionWidth)@4 as avg(hits.ResolutionWidth)] +02)--ProjectionExec: expr=[WatchID@0 as WatchID, ClientIP@1 as ClientIP, count(Int64(1))@2 as c, sum(hits.IsRefresh)@3 as sum(hits.IsRefresh), avg(hits.ResolutionWidth)@4 as avg(hits.ResolutionWidth)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[WatchID@0 as WatchID, ClientIP@1 as ClientIP], aggr=[count(Int64(1)), sum(hits.IsRefresh), avg(hits.ResolutionWidth)] 05)--------RepartitionExec: partitioning=Hash([WatchID@0, ClientIP@1], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[WatchID@0 as WatchID, ClientIP@1 as ClientIP], aggr=[count(Int64(1)), sum(hits.IsRefresh), avg(hits.ResolutionWidth)] @@ -867,8 +869,8 @@ logical_plan 05)--------TableScan: hits_raw projection=[WatchID, ClientIP, IsRefresh, ResolutionWidth] physical_plan 01)SortPreservingMergeExec: [c@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[WatchID@0 as WatchID, ClientIP@1 as ClientIP, count(Int64(1))@2 as c, sum(hits.IsRefresh)@3 as sum(hits.IsRefresh), avg(hits.ResolutionWidth)@4 as avg(hits.ResolutionWidth)] +02)--ProjectionExec: expr=[WatchID@0 as WatchID, ClientIP@1 as ClientIP, count(Int64(1))@2 as c, sum(hits.IsRefresh)@3 as sum(hits.IsRefresh), avg(hits.ResolutionWidth)@4 as avg(hits.ResolutionWidth)] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[WatchID@0 as WatchID, ClientIP@1 as ClientIP], aggr=[count(Int64(1)), sum(hits.IsRefresh), avg(hits.ResolutionWidth)] 05)--------RepartitionExec: partitioning=Hash([WatchID@0, ClientIP@1], 4), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[WatchID@0 as WatchID, ClientIP@1 as ClientIP], aggr=[count(Int64(1)), sum(hits.IsRefresh), avg(hits.ResolutionWidth)] @@ -900,8 +902,8 @@ logical_plan 05)--------TableScan: hits_raw projection=[URL] physical_plan 01)SortPreservingMergeExec: [c@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[c@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[URL@0 as URL, count(Int64(1))@1 as c] +02)--ProjectionExec: expr=[URL@0 as URL, count(Int64(1))@1 as c] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[URL@0 as URL], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([URL@0], 4), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[URL@0 as URL], aggr=[count(Int64(1))] @@ -989,8 +991,8 @@ logical_plan 07)------------TableScan: hits_raw projection=[EventDate, CounterID, URL, IsRefresh, DontCountHits], partial_filters=[hits_raw.CounterID = Int32(62), hits_raw.EventDate >= UInt16(15887), hits_raw.EventDate <= UInt16(15917), hits_raw.DontCountHits = Int16(0), hits_raw.IsRefresh = Int16(0), hits_raw.URL != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [pageviews@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[pageviews@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[URL@0 as URL, count(Int64(1))@1 as pageviews] +02)--ProjectionExec: expr=[URL@0 as URL, count(Int64(1))@1 as pageviews] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[URL@0 as URL], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([URL@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[URL@0 as URL], aggr=[count(Int64(1))] @@ -1016,8 +1018,8 @@ logical_plan 07)------------TableScan: hits_raw projection=[Title, EventDate, CounterID, IsRefresh, DontCountHits], partial_filters=[hits_raw.CounterID = Int32(62), hits_raw.EventDate >= UInt16(15887), hits_raw.EventDate <= UInt16(15917), hits_raw.DontCountHits = Int16(0), hits_raw.IsRefresh = Int16(0), hits_raw.Title != Utf8View("")] physical_plan 01)SortPreservingMergeExec: [pageviews@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[pageviews@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[Title@0 as Title, count(Int64(1))@1 as pageviews] +02)--ProjectionExec: expr=[Title@0 as Title, count(Int64(1))@1 as pageviews] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[Title@0 as Title], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([Title@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[Title@0 as Title], aggr=[count(Int64(1))] @@ -1045,8 +1047,8 @@ logical_plan physical_plan 01)GlobalLimitExec: skip=1000, fetch=10 02)--SortPreservingMergeExec: [pageviews@1 DESC], fetch=1010 -03)----SortExec: TopK(fetch=1010), expr=[pageviews@1 DESC], preserve_partitioning=[true] -04)------ProjectionExec: expr=[URL@0 as URL, count(Int64(1))@1 as pageviews] +03)----ProjectionExec: expr=[URL@0 as URL, count(Int64(1))@1 as pageviews] +04)------SortExec: TopK(fetch=1010), expr=[count(Int64(1))@1 DESC], preserve_partitioning=[true] 05)--------AggregateExec: mode=FinalPartitioned, gby=[URL@0 as URL], aggr=[count(Int64(1))] 06)----------RepartitionExec: partitioning=Hash([URL@0], 4), input_partitions=4 07)------------AggregateExec: mode=Partial, gby=[URL@0 as URL], aggr=[count(Int64(1))] @@ -1074,8 +1076,8 @@ logical_plan physical_plan 01)GlobalLimitExec: skip=1000, fetch=10 02)--SortPreservingMergeExec: [pageviews@5 DESC], fetch=1010 -03)----SortExec: TopK(fetch=1010), expr=[pageviews@5 DESC], preserve_partitioning=[true] -04)------ProjectionExec: expr=[TraficSourceID@0 as TraficSourceID, SearchEngineID@1 as SearchEngineID, AdvEngineID@2 as AdvEngineID, CASE WHEN hits.SearchEngineID = Int64(0) AND hits.AdvEngineID = Int64(0) THEN hits.Referer ELSE Utf8("") END@3 as src, URL@4 as dst, count(Int64(1))@5 as pageviews] +03)----ProjectionExec: expr=[TraficSourceID@0 as TraficSourceID, SearchEngineID@1 as SearchEngineID, AdvEngineID@2 as AdvEngineID, CASE WHEN hits.SearchEngineID = Int64(0) AND hits.AdvEngineID = Int64(0) THEN hits.Referer ELSE Utf8("") END@3 as src, URL@4 as dst, count(Int64(1))@5 as pageviews] +04)------SortExec: TopK(fetch=1010), expr=[count(Int64(1))@5 DESC], preserve_partitioning=[true] 05)--------AggregateExec: mode=FinalPartitioned, gby=[TraficSourceID@0 as TraficSourceID, SearchEngineID@1 as SearchEngineID, AdvEngineID@2 as AdvEngineID, CASE WHEN hits.SearchEngineID = Int64(0) AND hits.AdvEngineID = Int64(0) THEN hits.Referer ELSE Utf8("") END@3 as CASE WHEN hits.SearchEngineID = Int64(0) AND hits.AdvEngineID = Int64(0) THEN hits.Referer ELSE Utf8("") END, URL@4 as URL], aggr=[count(Int64(1))] 06)----------RepartitionExec: partitioning=Hash([TraficSourceID@0, SearchEngineID@1, AdvEngineID@2, CASE WHEN hits.SearchEngineID = Int64(0) AND hits.AdvEngineID = Int64(0) THEN hits.Referer ELSE Utf8("") END@3, URL@4], 4), input_partitions=4 07)------------AggregateExec: mode=Partial, gby=[TraficSourceID@2 as TraficSourceID, SearchEngineID@3 as SearchEngineID, AdvEngineID@4 as AdvEngineID, CASE WHEN SearchEngineID@3 = 0 AND AdvEngineID@4 = 0 THEN Referer@1 ELSE END as CASE WHEN hits.SearchEngineID = Int64(0) AND hits.AdvEngineID = Int64(0) THEN hits.Referer ELSE Utf8("") END, URL@0 as URL], aggr=[count(Int64(1))] @@ -1103,8 +1105,8 @@ logical_plan physical_plan 01)GlobalLimitExec: skip=100, fetch=10 02)--SortPreservingMergeExec: [pageviews@2 DESC], fetch=110 -03)----SortExec: TopK(fetch=110), expr=[pageviews@2 DESC], preserve_partitioning=[true] -04)------ProjectionExec: expr=[URLHash@0 as URLHash, EventDate@1 as EventDate, count(Int64(1))@2 as pageviews] +03)----ProjectionExec: expr=[URLHash@0 as URLHash, EventDate@1 as EventDate, count(Int64(1))@2 as pageviews] +04)------SortExec: TopK(fetch=110), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 05)--------AggregateExec: mode=FinalPartitioned, gby=[URLHash@0 as URLHash, EventDate@1 as EventDate], aggr=[count(Int64(1))] 06)----------RepartitionExec: partitioning=Hash([URLHash@0, EventDate@1], 4), input_partitions=4 07)------------AggregateExec: mode=Partial, gby=[URLHash@0 as URLHash, EventDate@1 as EventDate], aggr=[count(Int64(1))] @@ -1133,8 +1135,8 @@ logical_plan physical_plan 01)GlobalLimitExec: skip=10000, fetch=10 02)--SortPreservingMergeExec: [pageviews@2 DESC], fetch=10010 -03)----SortExec: TopK(fetch=10010), expr=[pageviews@2 DESC], preserve_partitioning=[true] -04)------ProjectionExec: expr=[WindowClientWidth@0 as WindowClientWidth, WindowClientHeight@1 as WindowClientHeight, count(Int64(1))@2 as pageviews] +03)----ProjectionExec: expr=[WindowClientWidth@0 as WindowClientWidth, WindowClientHeight@1 as WindowClientHeight, count(Int64(1))@2 as pageviews] +04)------SortExec: TopK(fetch=10010), expr=[count(Int64(1))@2 DESC], preserve_partitioning=[true] 05)--------AggregateExec: mode=FinalPartitioned, gby=[WindowClientWidth@0 as WindowClientWidth, WindowClientHeight@1 as WindowClientHeight], aggr=[count(Int64(1))] 06)----------RepartitionExec: partitioning=Hash([WindowClientWidth@0, WindowClientHeight@1], 4), input_partitions=4 07)------------AggregateExec: mode=Partial, gby=[WindowClientWidth@0 as WindowClientWidth, WindowClientHeight@1 as WindowClientHeight], aggr=[count(Int64(1))] diff --git a/datafusion/sqllogictest/test_files/copy.slt b/datafusion/sqllogictest/test_files/copy.slt index 402ac8e8512bf..77977a6afcb11 100644 --- a/datafusion/sqllogictest/test_files/copy.slt +++ b/datafusion/sqllogictest/test_files/copy.slt @@ -33,7 +33,7 @@ COPY source_table TO 'test_files/scratch/copy/partitioned_table1/' STORED AS par # validate multiple partitioned parquet file output statement ok -CREATE EXTERNAL TABLE validate_partitioned_parquet STORED AS PARQUET +CREATE EXTERNAL TABLE validate_partitioned_parquet STORED AS PARQUET LOCATION 'test_files/scratch/copy/partitioned_table1/' PARTITIONED BY (col2); query IT @@ -44,7 +44,7 @@ select * from validate_partitioned_parquet order by col1, col2; # validate partition paths were actually generated statement ok -CREATE EXTERNAL TABLE validate_partitioned_parquet_bar STORED AS PARQUET +CREATE EXTERNAL TABLE validate_partitioned_parquet_bar STORED AS PARQUET LOCATION 'test_files/scratch/copy/partitioned_table1/col2=Bar'; query I @@ -61,7 +61,7 @@ OPTIONS ('format.compression' 'zstd(10)'); # validate multiple partitioned parquet file output statement ok -CREATE EXTERNAL TABLE validate_partitioned_parquet2 STORED AS PARQUET +CREATE EXTERNAL TABLE validate_partitioned_parquet2 STORED AS PARQUET LOCATION 'test_files/scratch/copy/partitioned_table2/' PARTITIONED BY (column2, column3); query ITT @@ -72,7 +72,7 @@ select * from validate_partitioned_parquet2 order by column1,column2,column3; 3 c z statement ok -CREATE EXTERNAL TABLE validate_partitioned_parquet_a_x STORED AS PARQUET +CREATE EXTERNAL TABLE validate_partitioned_parquet_a_x STORED AS PARQUET LOCATION 'test_files/scratch/copy/partitioned_table2/column2=a/column3=x'; query I @@ -89,7 +89,7 @@ OPTIONS ('format.compression' 'zstd(10)'); # validate multiple partitioned parquet file output statement ok -CREATE EXTERNAL TABLE validate_partitioned_parquet3 STORED AS PARQUET +CREATE EXTERNAL TABLE validate_partitioned_parquet3 STORED AS PARQUET LOCATION 'test_files/scratch/copy/partitioned_table3/' PARTITIONED BY (column1, column3); query TTT @@ -100,7 +100,7 @@ select column1, column2, column3 from validate_partitioned_parquet3 order by col 3 c z statement ok -CREATE EXTERNAL TABLE validate_partitioned_parquet_1_x STORED AS PARQUET +CREATE EXTERNAL TABLE validate_partitioned_parquet_1_x STORED AS PARQUET LOCATION 'test_files/scratch/copy/partitioned_table3/column1=1/column3=x'; query T @@ -143,7 +143,7 @@ select column1, column2, column3, column4, column5, column6, column7, column8, c statement ok -create table test ("'test'" varchar, "'test2'" varchar, "'test3'" varchar); +create table test ("'test'" varchar, "'test2'" varchar, "'test3'" varchar); # https://github.com/apache/datafusion/issues/9714 ## Until the partition by parsing uses ColumnDef, this test is meaningless since it becomes an overfit. Even in @@ -249,7 +249,7 @@ select * from validate_parquet; 2 Bar query I -copy (values (struct(timestamp '2021-01-01 01:00:01', 1)), (struct(timestamp '2022-01-01 01:00:01', 2)), +copy (values (struct(timestamp '2021-01-01 01:00:01', 1)), (struct(timestamp '2022-01-01 01:00:01', 2)), (struct(timestamp '2023-01-03 01:00:01', 3)), (struct(timestamp '2024-01-01 01:00:01', 4))) to 'test_files/scratch/copy/table_nested2/' STORED AS PARQUET; ---- @@ -267,15 +267,15 @@ select * from validate_parquet_nested2; {c0: 2024-01-01T01:00:01, c1: 4} query I -COPY -(values (struct ('foo', (struct ('foo', make_array(struct('a',1), struct('b',2))))), make_array(timestamp '2023-01-01 01:00:01',timestamp '2023-01-01 01:00:01')), -(struct('bar', (struct ('foo', make_array(struct('aa',10), struct('bb',20))))), make_array(timestamp '2024-01-01 01:00:01', timestamp '2024-01-01 01:00:01'))) +COPY +(values (struct ('foo', (struct ('foo', make_array(struct('a',1), struct('b',2))))), make_array(timestamp '2023-01-01 01:00:01',timestamp '2023-01-01 01:00:01')), +(struct('bar', (struct ('foo', make_array(struct('aa',10), struct('bb',20))))), make_array(timestamp '2024-01-01 01:00:01', timestamp '2024-01-01 01:00:01'))) to 'test_files/scratch/copy/table_nested/' STORED AS PARQUET; ---- 2 statement ok -CREATE EXTERNAL TABLE validate_parquet_nested STORED AS PARQUET +CREATE EXTERNAL TABLE validate_parquet_nested STORED AS PARQUET LOCATION 'test_files/scratch/copy/table_nested/'; query ?? @@ -285,14 +285,14 @@ select * from validate_parquet_nested; {c0: bar, c1: {c0: foo, c1: [{c0: aa, c1: 10}, {c0: bb, c1: 20}]}} [2024-01-01T01:00:01, 2024-01-01T01:00:01] query I -copy (values ([struct('foo', 1), struct('bar', 2)])) +copy (values ([struct('foo', 1), struct('bar', 2)])) to 'test_files/scratch/copy/array_of_struct/' STORED AS PARQUET; ---- 1 statement ok -CREATE EXTERNAL TABLE validate_array_of_struct +CREATE EXTERNAL TABLE validate_array_of_struct STORED AS PARQUET LOCATION 'test_files/scratch/copy/array_of_struct/'; query ? @@ -301,7 +301,7 @@ select * from validate_array_of_struct; [{c0: foo, c1: 1}, {c0: bar, c1: 2}] query I -copy (values (struct('foo', [1,2,3], struct('bar', [2,3,4])))) +copy (values (struct('foo', [1,2,3], struct('bar', [2,3,4])))) to 'test_files/scratch/copy/struct_with_array/' STORED AS PARQUET; ---- 1 @@ -326,6 +326,7 @@ OPTIONS ( 'format.compression::col1' 'zstd(5)', 'format.compression::col2' snappy, 'format.max_row_group_size' 12345, +'format.max_row_group_bytes' 2048, 'format.data_pagesize_limit' 1234, 'format.write_batch_size' 1234, 'format.writer_version' 2.0, @@ -577,8 +578,8 @@ select * from validate_arrow_file; # Copy from dict encoded values to single arrow file query I -COPY (values -('c', arrow_cast('foo', 'Dictionary(Int32, Utf8)')), ('d', arrow_cast('bar', 'Dictionary(Int32, Utf8)'))) +COPY (values +('c', arrow_cast('foo', 'Dictionary(Int32, Utf8)')), ('d', arrow_cast('bar', 'Dictionary(Int32, Utf8)'))) to 'test_files/scratch/copy/table_dict.arrow' STORED AS ARROW; ---- 2 diff --git a/datafusion/sqllogictest/test_files/create_external_table.slt b/datafusion/sqllogictest/test_files/create_external_table.slt index f56cff2a2a2f0..1d339f402501f 100644 --- a/datafusion/sqllogictest/test_files/create_external_table.slt +++ b/datafusion/sqllogictest/test_files/create_external_table.slt @@ -303,3 +303,42 @@ statement error DataFusion error: SQL error: ParserError\("'IF NOT EXISTS' canno CREATE OR REPLACE EXTERNAL TABLE IF NOT EXISTS t_conflict(c1 int) STORED AS CSV LOCATION 'foo.csv'; + +# Multiple listed locations are read together as a single table. +# Each partition-N.csv has 11 rows, so listing exactly two of them (rather than +# the whole directory) yields 22 rows. +statement ok +CREATE EXTERNAL TABLE multi_loc (c1 int, c2 bigint, c3 boolean) +STORED AS CSV +LOCATION ('../core/tests/data/partitioned_csv/partition-0.csv', '../core/tests/data/partitioned_csv/partition-1.csv') +OPTIONS ('format.has_header' 'false'); + +query I +SELECT count(*) FROM multi_loc; +---- +22 + +statement ok +DROP TABLE multi_loc; + +# Duplicate locations are rejected to avoid scanning the same data twice. +statement error Duplicate location +CREATE EXTERNAL TABLE multi_loc_duplicate (c1 int, c2 bigint, c3 boolean) +STORED AS CSV +LOCATION ('../core/tests/data/partitioned_csv/partition-0.csv', '../core/tests/data/partitioned_csv/partition-0.csv') +OPTIONS ('format.has_header' 'false'); + +# Whitespace around the list separators is ignored +statement ok +CREATE EXTERNAL TABLE multi_loc_ws (c1 int, c2 bigint, c3 boolean) +STORED AS CSV +LOCATION ( '../core/tests/data/partitioned_csv/partition-0.csv' , '../core/tests/data/partitioned_csv/partition-1.csv' ) +OPTIONS ('format.has_header' 'false'); + +query I +SELECT count(*) FROM multi_loc_ws; +---- +22 + +statement ok +DROP TABLE multi_loc_ws; diff --git a/datafusion/sqllogictest/test_files/csv_files.slt b/datafusion/sqllogictest/test_files/csv_files.slt index d980e802c83cb..af2c6d41af42e 100644 --- a/datafusion/sqllogictest/test_files/csv_files.slt +++ b/datafusion/sqllogictest/test_files/csv_files.slt @@ -376,7 +376,7 @@ id3 value3 # Reset repartition_file_min_size to default value statement ok -SET datafusion.optimizer.repartition_file_min_size = 10485760; +RESET datafusion.optimizer.repartition_file_min_size; statement ok drop table stored_table_with_cr_terminator; diff --git a/datafusion/sqllogictest/test_files/cte.slt b/datafusion/sqllogictest/test_files/cte.slt index d13e0d4f085e9..89110e9788914 100644 --- a/datafusion/sqllogictest/test_files/cte.slt +++ b/datafusion/sqllogictest/test_files/cte.slt @@ -171,7 +171,7 @@ logical_plan 07)--------TableScan: nodes projection=[id] physical_plan 01)RecursiveQueryExec: name=nodes, is_distinct=false -02)--ProjectionExec: expr=[1 as id] +02)--ProjectionExec: expr=[CAST(1 AS Int64) as id] 03)----PlaceholderRowExec 04)--CoalescePartitionsExec 05)----ProjectionExec: expr=[id@0 + 1 as id] @@ -179,6 +179,111 @@ physical_plan 07)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 08)----------WorkTableExec: name=nodes +# recursive CTE with a column-list alias (e.g. `t(n)`): the declared names must be +# applied to the static term so the recursive self-reference can resolve them +query I rowsort +WITH RECURSIVE t(n) AS ( + SELECT 1 + UNION ALL + SELECT n + 1 FROM t WHERE n < 10 +) +SELECT n FROM t +---- +1 +10 +2 +3 +4 +5 +6 +7 +8 +9 + +# recursive CTE with a multi-column column-list alias +query II rowsort +WITH RECURSIVE t(a, b) AS ( + SELECT 1, 2 + UNION ALL + SELECT a + 1, b * 2 FROM t WHERE a < 5 +) +SELECT a, b FROM t +---- +1 2 +2 4 +3 8 +4 16 +5 32 + +# recursive CTE with a column-list alias and UNION (DISTINCT) +query I rowsort +WITH RECURSIVE t(n) AS ( + SELECT 1 + UNION + SELECT n + 1 FROM t WHERE n < 5 +) +SELECT n FROM t +---- +1 +2 +3 +4 +5 + +# recursive CTE column-list alias arity mismatch is rejected cleanly (raised at +# the static term, rather than the old confusing "No field named ...") +query error DataFusion error: Error during planning: Source table contains 1 columns but only 2 names given as column alias +WITH RECURSIVE t(a, b) AS ( + SELECT 1 + UNION ALL + SELECT a + 1 FROM t WHERE a < 3 +) +SELECT * FROM t + +# explain a column-list-aliased recursive CTE: the declared name is applied to +# the static term, so there is no extra projection on top of RecursiveQuery +query TT +EXPLAIN WITH RECURSIVE t(n) AS ( + SELECT 1 + UNION ALL + SELECT n + 1 FROM t WHERE n < 10 +) +SELECT * FROM t +---- +logical_plan +01)SubqueryAlias: t +02)--RecursiveQuery: is_distinct=false +03)----Projection: Int64(1) AS n +04)------EmptyRelation: rows=1 +05)----Projection: t.n + Int64(1) +06)------Filter: t.n < Int64(10) +07)--------TableScan: t projection=[n] +physical_plan +01)RecursiveQueryExec: name=t, is_distinct=false +02)--ProjectionExec: expr=[CAST(1 AS Int64) as n] +03)----PlaceholderRowExec +04)--CoalescePartitionsExec +05)----ProjectionExec: expr=[n@0 + 1 as n] +06)------FilterExec: n@0 < 10 +07)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +08)----------WorkTableExec: name=t + +# recursive CTE with a quoted, case-sensitive column-list alias: `"N"` must be +# preserved (not lowercased) so the recursive self-reference resolves it +query I rowsort +WITH RECURSIVE t("N") AS ( + SELECT 1 + UNION ALL + SELECT "N" + 1 FROM t WHERE "N" < 5 +) +SELECT "N" FROM t +---- +1 +2 +3 +4 +5 + # simple deduplicating recursive CTE works query I WITH RECURSIVE nodes AS ( @@ -699,7 +804,7 @@ WITH RECURSIVE region_sales AS ( SELECT s.salesperson_id AS salesperson_id, SUM(s.sale_amount) AS amount, - SUM(0) as level + 0 as level FROM sales s GROUP BY @@ -842,12 +947,12 @@ logical_plan 03)----Projection: Int64(1) AS val 04)------EmptyRelation: rows=1 05)----Projection: Int64(2) AS val -06)------Cross Join: -07)--------Filter: recursive_cte.val < Int64(2) -08)----------TableScan: recursive_cte -09)--------SubqueryAlias: sub_cte -10)----------Projection: Int64(2) AS val -11)------------EmptyRelation: rows=1 +06)------Cross Join: +07)--------Projection: +08)----------Filter: recursive_cte.val < Int64(2) +09)------------TableScan: recursive_cte projection=[val] +10)--------SubqueryAlias: sub_cte +11)----------EmptyRelation: rows=1 physical_plan 01)RecursiveQueryExec: name=recursive_cte, is_distinct=false 02)--ProjectionExec: expr=[1 as val] @@ -855,11 +960,10 @@ physical_plan 04)--ProjectionExec: expr=[2 as val] 05)----CrossJoinExec 06)------CoalescePartitionsExec -07)--------FilterExec: val@0 < 2 +07)--------FilterExec: val@0 < 2, projection=[] 08)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 09)------------WorkTableExec: name=recursive_cte -10)------ProjectionExec: expr=[2 as val] -11)--------PlaceholderRowExec +10)------PlaceholderRowExec # Test issue: https://github.com/apache/datafusion/issues/9794 # Non-recursive term and recursive term have different types @@ -1079,7 +1183,7 @@ logical_plan 07)--------TableScan: numbers projection=[n] physical_plan 01)RecursiveQueryExec: name=numbers, is_distinct=false -02)--ProjectionExec: expr=[1 as n] +02)--ProjectionExec: expr=[CAST(1 AS Int64) as n] 03)----PlaceholderRowExec 04)--CoalescePartitionsExec 05)----ProjectionExec: expr=[n@0 + 1 as n] @@ -1104,7 +1208,7 @@ logical_plan 07)--------TableScan: numbers projection=[n] physical_plan 01)RecursiveQueryExec: name=numbers, is_distinct=false -02)--ProjectionExec: expr=[1 as n] +02)--ProjectionExec: expr=[CAST(1 AS Int64) as n] 03)----PlaceholderRowExec 04)--CoalescePartitionsExec 05)----ProjectionExec: expr=[n@0 + 1 as n] @@ -1161,7 +1265,7 @@ logical_plan physical_plan 01)GlobalLimitExec: skip=0, fetch=5 02)--RecursiveQueryExec: name=r, is_distinct=false -03)----ProjectionExec: expr=[0 as k, 0 as v] +03)----ProjectionExec: expr=[CAST(0 AS Int64) as k, CAST(0 AS Int64) as v] 04)------PlaceholderRowExec 05)----SortExec: TopK(fetch=1), expr=[v@1 ASC NULLS LAST], preserve_partitioning=[false] 06)------WorkTableExec: name=r @@ -1205,14 +1309,13 @@ EXPLAIN WITH RECURSIVE trans AS ( logical_plan 01)SubqueryAlias: trans 02)--RecursiveQuery: is_distinct=true -03)----Projection: closure.start, closure.end -04)------TableScan: closure -05)----Projection: l.start, r.end -06)------Inner Join: l.end = r.start -07)--------SubqueryAlias: l -08)----------TableScan: trans -09)--------SubqueryAlias: r -10)----------TableScan: closure +03)----TableScan: closure projection=[start, end] +04)----Projection: l.start, r.end +05)------Inner Join: l.end = r.start +06)--------SubqueryAlias: l +07)----------TableScan: trans projection=[start, end] +08)--------SubqueryAlias: r +09)----------TableScan: closure projection=[start, end] physical_plan 01)RecursiveQueryExec: name=trans, is_distinct=true 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/recursive_cte/closure.csv]]}, projection=[start, end], file_type=csv, has_header=true @@ -1300,6 +1403,186 @@ DROP TABLE cte_schema_reread; statement ok DROP TABLE cte_schema_records; +########## +## Recursive CTE nullability widening +## +## A recursive term can introduce NULLs that the static (anchor) term never +## produces. The recursive CTE output schema must therefore widen nullability +## across both terms, otherwise nullability-based optimizer simplifications +## (e.g. removing IS NULL / IS NOT NULL predicates) produce wrong results. +########## + +# recursive self-reference must use conservative nullability even when the +# anchor term uses non-null literals. Otherwise optimizer nullability-based +# simplification can remove this semantically required IS NOT NULL guard. +query II rowsort +WITH RECURSIVE t(a, b) AS ( + SELECT 0 AS a, 0 AS b + UNION ALL + SELECT b AS a, CAST(NULL AS INT) AS b FROM t WHERE a IS NOT NULL +) +SELECT * FROM t +---- +0 0 +0 NULL +NULL NULL + +# outer IS NOT NULL filters must see recursive output as nullable, not just the +# non-null anchor literal. +query II rowsort +WITH RECURSIVE t(a, b) AS ( + SELECT 0 AS a, 0 AS b + UNION ALL + SELECT b AS a, CAST(NULL AS INT) AS b FROM t WHERE a IS NOT NULL +) +SELECT * FROM t WHERE a IS NOT NULL +---- +0 0 +0 NULL + +# outer IS NULL filters must see recursive output as nullable, not just the +# non-null anchor literal. +query I +WITH RECURSIVE t(a) AS ( + SELECT 1 AS a + UNION ALL + SELECT CAST(NULL AS INT) AS a FROM t WHERE a IS NOT NULL +) +SELECT * FROM t WHERE a IS NULL +---- +NULL + +# deduplicating recursive CTE must preserve widened nullability for outer filters. +query I +WITH RECURSIVE t(a) AS ( + SELECT 1 AS a + UNION + SELECT CAST(NULL AS INT) AS a FROM t WHERE a IS NOT NULL +) +SELECT * FROM t WHERE a IS NULL +---- +NULL + +# recursive output nullability must be tracked per column, not just for the +# first column. +query II rowsort +WITH RECURSIVE t(a, b) AS ( + SELECT 0 AS a, 0 AS b + UNION ALL + SELECT b AS a, CAST(NULL AS INT) AS b FROM t WHERE a IS NOT NULL +) +SELECT * FROM t WHERE b IS NULL +---- +0 NULL +NULL NULL + +# recursive output nullability must survive recursive term type coercion. +query I +WITH RECURSIVE t(a) AS ( + SELECT 1::INT AS a + UNION ALL + SELECT CAST(NULL AS BIGINT) AS a FROM t WHERE a IS NOT NULL +) +SELECT * FROM t WHERE a IS NULL +---- +NULL + +# DESCRIBE should expose the widened recursive output nullability. +query TTT +DESCRIBE WITH RECURSIVE t(a) AS ( + SELECT 1 AS a + UNION ALL + SELECT CAST(NULL AS INT) AS a FROM t WHERE a IS NOT NULL +) +SELECT * FROM t +---- +a Int64 YES + +# recursive self-reference must not simplify away IS NULL guards when the +# anchor term is nullable and the recursive term is non-null. +query I rowsort +WITH RECURSIVE t(a) AS ( + SELECT CAST(NULL AS INT) AS a + UNION ALL + SELECT 1 AS a FROM t WHERE a IS NULL +) +SELECT * FROM t +---- +1 +NULL + +# widened recursive nullability must survive aggregate physical planning. +query III +WITH RECURSIVE t(a) AS ( + SELECT 1 AS a + UNION ALL + SELECT CAST(NULL AS INT) AS a FROM t WHERE a IS NOT NULL +) +SELECT COUNT(*), COUNT(a), SUM(a) FROM t +---- +2 1 1 + +# outer filters must still see widened nullability through a derived projection. +query I +WITH RECURSIVE t(a) AS ( + SELECT 1 AS a + UNION ALL + SELECT CAST(NULL AS INT) AS a FROM t WHERE a IS NOT NULL +) +SELECT x FROM (SELECT a AS x FROM t) WHERE x IS NULL +---- +NULL + +# per-column nullability widening must survive type coercion in multi-column +# recursive terms. +query II +WITH RECURSIVE t(a, b) AS ( + SELECT 1::INT AS a, 2::INT AS b + UNION ALL + SELECT a + 1 AS a, CAST(NULL AS BIGINT) AS b FROM t WHERE a < 2 +) +SELECT * FROM t WHERE b IS NULL +---- +2 NULL + +# join planning must preserve recursive output nullability for null-sensitive +# predicates above the recursive query. +query II +WITH RECURSIVE t(a) AS ( + SELECT 1 AS a + UNION ALL + SELECT CAST(NULL AS INT) AS a FROM t WHERE a IS NOT NULL +) +SELECT t.a, u.b FROM t LEFT JOIN (SELECT 1 AS b) u ON t.a = u.b WHERE u.b IS NULL +---- +NULL NULL + +# A recursive CTE must not inherit the static term's uniqueness / primary-key +# functional dependencies: the recursive term can append rows that duplicate +# the static term's "unique" keys, so an outer DISTINCT must not be optimized +# away based on a stale dependency. +statement ok +CREATE TABLE recursive_cte_pk(id INT NOT NULL PRIMARY KEY); + +statement ok +INSERT INTO recursive_cte_pk VALUES (2), (1); + +query I rowsort +SELECT DISTINCT id FROM ( + WITH RECURSIVE t(id) AS ( + SELECT id FROM recursive_cte_pk + UNION ALL + SELECT id - 1 FROM t WHERE id > 1 + ) + SELECT id FROM t +) +---- +1 +2 + +statement ok +DROP TABLE recursive_cte_pk; + statement count 0 set datafusion.execution.enable_recursive_ctes = false; @@ -1319,3 +1602,128 @@ RESET datafusion.execution.enable_recursive_ctes; statement ok RESET datafusion.sql_parser.enable_ident_normalization; + + +# Test projection optimization in recursive CTEs + +# https://github.com/apache/datafusion/issues/22249 +query I +with recursive t(k, v) as ( + select 1 k, 10 v + union all + select 2, 20 from t where k = 1 +) +select v +from t +order by 1; +---- +10 +20 + +# https://github.com/apache/datafusion/issues/22249 +query I +with recursive t(k, v) as ( + select 1 k, 10 v + union all + select 2, 20 from t where v = 10 +) +select v +from t +order by 1; +---- +10 +20 + +# Keep columns that are not selected by the outer query, but still affect +# recursive UNION distinctness. +query I +with recursive t(k, v) as ( + select 1 k, 10 v + union + select 2, 10 from t where v = 10 +) +select v +from t +order by 1; +---- +10 +10 + +statement ok +copy ( + select i as k, i as v1, i as v2 + from generate_series(1, 3) t(i) +) to 'test_files/scratch/cte/test.parquet'; + +statement ok +create external table test stored as parquet location 'test_files/scratch/cte/test.parquet'; + +# check that both the static and recursive terms are optimized +query TT +explain +with recursive r as ( + select k, v1 -- only needs to project k and v1 from table test + from test + union all + select k * 10, v1 + from r + where k < ( -- only needs to project k and v2 from table test + select v2 + from test + where k = 2 + ) +) +select * +from r +order by 1, 2; +---- +logical_plan +01)Sort: r.k ASC NULLS LAST, r.v1 ASC NULLS LAST +02)--SubqueryAlias: r +03)----RecursiveQuery: is_distinct=false +04)------TableScan: test projection=[k, v1] +05)------Projection: r.k * Int64(10), r.v1 +06)--------Filter: r.k < () +07)----------Subquery: +08)------------Projection: test.v2 +09)--------------Filter: test.k = Int64(2) +10)----------------TableScan: test projection=[k, v2], partial_filters=[test.k = Int64(2)] +11)----------TableScan: r projection=[k, v1] +physical_plan +01)ScalarSubqueryExec: subqueries=1 +02)--SortExec: expr=[k@0 ASC NULLS LAST, v1@1 ASC NULLS LAST], preserve_partitioning=[false] +03)----RecursiveQueryExec: name=r, is_distinct=false +04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/cte/test.parquet]]}, projection=[CAST(k@0 AS Int64) as k, CAST(v1@1 AS Int64) as v1], output_ordering=[k@0 ASC NULLS LAST], file_type=parquet +05)------CoalescePartitionsExec +06)--------ProjectionExec: expr=[k@0 * 10 as k, v1@1 as v1] +07)----------FilterExec: k@0 < scalar_subquery() +08)------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +09)--------------WorkTableExec: name=r +10)--FilterExec: k@0 = 2, projection=[v2@1] +11)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +12)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/cte/test.parquet]]}, projection=[k, v2], output_ordering=[k@0 ASC NULLS LAST], file_type=parquet, predicate=k@0 = 2, pruning_predicate=k_null_count@2 != row_count@3 AND k_min@0 <= 2 AND 2 <= k_max@1, required_guarantees=[k in (2)] + +query II +with recursive r as ( + select k, v1 + from test + union all + select k * 10, v1 + from r + where k < ( + select v2 + from test + where k = 2 + ) +) +select * +from r +order by 1, 2; +---- +1 1 +2 2 +3 3 +10 1 + +statement ok +drop table test; diff --git a/datafusion/sqllogictest/test_files/date_bin_errors.slt b/datafusion/sqllogictest/test_files/date_bin_errors.slt index b59201eb906f6..53cba506defd6 100644 --- a/datafusion/sqllogictest/test_files/date_bin_errors.slt +++ b/datafusion/sqllogictest/test_files/date_bin_errors.slt @@ -23,10 +23,24 @@ select date_bin(interval '1637426858 months', to_timestamp_millis(1040292460), t ---- NULL -# Negative timestamp with month interval - should return NULL instead of panicking +# Issue #22528: negative sub-second source with month interval. query P select date_bin(interval '1 month', to_timestamp_millis(-1040292460), timestamp '1984-01-07 00:00:00'); ---- +1969-12-07T00:00:00 + +# Array path should match the scalar path above. +query P +select date_bin(interval '1 month', c, timestamp '1984-01-07 00:00:00') +from values (to_timestamp_millis(-1040292460)) t(c); +---- +1969-12-07T00:00:00 + +# Array path should return NULL for per-row overflow. +query P +select date_bin(interval '1637426858 months', c, timestamp '1984-01-07 00:00:00') +from values (to_timestamp_millis(1040292460)) t(c); +---- NULL # Large stride causing overflow - should return NULL @@ -67,4 +81,37 @@ select date_bin( arrow_cast(-9223372036854775808, 'Timestamp(Nanosecond, None)') ); ---- -NULL \ No newline at end of file +NULL + +# compute_distance overflow: source at i64::MIN nanoseconds previously panicked +# inside compute_distance; it must return NULL through the SQL execution path +query P +select date_bin( + interval '3 nanoseconds', + arrow_cast(-9223372036854775808, 'Timestamp(Nanosecond, None)') +); +---- +NULL + +# Source timestamp scaling to nanoseconds overflows: should return NULL, not panic +query P +select date_bin( + interval '1 nanosecond', + arrow_cast(9223372036854775807, 'Timestamp(Second, None)'), + timestamp '1970-01-01 00:00:00' +); +---- +NULL + +# Source timestamp scaling to nanoseconds overflows in array path: should return NULL, not panic +query P +select date_bin( + interval '1 nanosecond', + ts, + timestamp '1970-01-01 00:00:00' +) +from ( + values (arrow_cast(9223372036854775807, 'Timestamp(Second, None)')) +) as t(ts); +---- +NULL diff --git a/datafusion/sqllogictest/test_files/datetime/arith_date_interval.slt b/datafusion/sqllogictest/test_files/datetime/arith_date_interval.slt index 01e1939996dfc..12fcc2bfd3464 100644 --- a/datafusion/sqllogictest/test_files/datetime/arith_date_interval.slt +++ b/datafusion/sqllogictest/test_files/datetime/arith_date_interval.slt @@ -47,3 +47,6 @@ SELECT arrow_cast('2020-01-01', 'Date64') + INTERVAL '999999' YEAR query error Arrow error: Compute error: Date arithmetic overflow SELECT arrow_cast('2020-01-01', 'Date64') - INTERVAL '999999' YEAR + +query error Arrow error: Compute error: Date arithmetic overflow +SELECT DATE '2262-04-10' + INTERVAL '999999999' DAY diff --git a/datafusion/sqllogictest/test_files/datetime/arith_time_interval.slt b/datafusion/sqllogictest/test_files/datetime/arith_time_interval.slt index 997eae9b1bd8b..1d2b0e15bb953 100644 --- a/datafusion/sqllogictest/test_files/datetime/arith_time_interval.slt +++ b/datafusion/sqllogictest/test_files/datetime/arith_time_interval.slt @@ -1,70 +1,143 @@ # postgresql behavior # # time + interval → time -# Add an interval to a time +# Add an interval to a time. The result is a `time` value that wraps within the +# 24-hour clock, matching PostgreSQL and DuckDB. # time '01:00' + interval '3 hours' → 04:00:00 -# -# note that while the above reflects what postgresql does -# in the case of datafusion/arrow that is not the case. The -# result will be an interval, not a time. +# time '22:00' + interval '3 hours' → 01:00:00 (wraps past midnight) -query ? +query D SELECT '01:00'::time + interval '3 hours' ---- -4 hours +04:00:00 query T SELECT arrow_typeof('01:00'::time + interval '3 hours') ---- -Interval(MonthDayNano) +Time64(ns) -query ? +query D SELECT '22:00'::time + interval '3 hours' ---- -25 hours +01:00:00 -query ? +query D SELECT interval '3 hours' + '22:00'::time ---- -25 hours +01:00:00 -query ? +# The result keeps the input time's unit, mirroring `timestamp + interval`, rather +# than widening to Time64(ns). +query D SELECT arrow_cast('22:00', 'Time32(Second)') + interval '3 hours' ---- -25 hours +01:00:00 + +query T +SELECT arrow_typeof(arrow_cast('22:00', 'Time32(Second)') + interval '3 hours') +---- +Time32(s) -query ? +query D SELECT arrow_cast('22:00', 'Time32(Millisecond)') + interval '3 hours' ---- -25 hours +01:00:00 -query ? +query T +SELECT arrow_typeof(arrow_cast('22:00', 'Time32(Millisecond)') + interval '3 hours') +---- +Time32(ms) + +query D SELECT arrow_cast('22:00', 'Time64(Microsecond)') + interval '3 hours' ---- -25 hours +01:00:00 -query ? +query T +SELECT arrow_typeof(arrow_cast('22:00', 'Time64(Microsecond)') + interval '3 hours') +---- +Time64(µs) + +query D SELECT arrow_cast('22:00', 'Time64(Nanosecond)') + interval '3 hours' ---- -25 hours +01:00:00 + +query T +SELECT arrow_typeof(arrow_cast('22:00', 'Time64(Nanosecond)') + interval '3 hours') +---- +Time64(ns) + +# The interval is applied at nanosecond precision and floored to the time's unit, exactly +# as for `timestamp(unit) ± interval`. Adding one nanosecond to a second-resolution time +# floors back to a no-op... +query D +SELECT arrow_cast('22:00', 'Time32(Second)') + interval '1 nanosecond' +---- +22:00:00 + +# ...but subtracting one nanosecond floors down a full second, matching +# `timestamp(s) - interval '1 nanosecond'` (= 09:59:59) rather than staying put. +query D +SELECT arrow_cast('10:00:00', 'Time32(Second)') - interval '1 nanosecond' +---- +09:59:59 + +query D +SELECT arrow_cast('12:00:00', 'Time32(Millisecond)') + interval '1 microsecond' +---- +12:00:00 + +query D +SELECT arrow_cast('12:00:00', 'Time64(Microsecond)') + interval '1 microsecond' +---- +12:00:00.000001 + +# Whole days and months in the interval do not affect a time-of-day (PostgreSQL). +query D +SELECT '10:00'::time + interval '1 day 2 hours' +---- +12:00:00 # postgresql behavior # # time - interval → time -# Subtract an interval from a time +# Subtract an interval from a time, wrapping within the 24-hour clock. # time '05:00' - interval '2 hours' → 03:00:00 +# time '02:00' - interval '3 hours' → 23:00:00 (wraps before midnight) -query ? +query D SELECT '05:00'::time - interval '2 hours' ---- -3 hours +03:00:00 query T SELECT arrow_typeof('05:00'::time - interval '2 hours') ---- -Interval(MonthDayNano) +Time64(ns) -query ? +query D SELECT '02:00'::time - interval '3 hours' ---- --1 hours +23:00:00 + +# Array inputs (not only scalars) exercise the columnar path, including nulls. +statement ok +CREATE TABLE time_vals(id INT, t TIME) AS VALUES (1, '01:00'::time), (2, '22:00'::time), (3, NULL); + +query D +SELECT t + interval '3 hours' FROM time_vals ORDER BY id +---- +04:00:00 +01:00:00 +NULL + +query D +SELECT t - interval '2 hours' FROM time_vals ORDER BY id +---- +23:00:00 +20:00:00 +NULL + +statement ok +DROP TABLE time_vals diff --git a/datafusion/sqllogictest/test_files/datetime/date_part.slt b/datafusion/sqllogictest/test_files/datetime/date_part.slt index 891319f9e2cd2..0a992b2d78a22 100644 --- a/datafusion/sqllogictest/test_files/datetime/date_part.slt +++ b/datafusion/sqllogictest/test_files/datetime/date_part.slt @@ -838,6 +838,40 @@ SELECT extract(millisecond from arrow_cast('23:32:50.123456789'::time, 'Time64(N ---- 50123 +# date32 and date64 + +statement ok +CREATE TABLE source_dt AS +with t as (values + ('1970-01-01'), + ('2020-06-02'), + ('2026-02-28'), + (NULL) +) +SELECT + arrow_cast(column1, 'Date32') as date32, + arrow_cast(column1, 'Date64') as date64, +FROM t; + +query IIIIIIIIII +SELECT date_part('year', date32), date_part('month', date32), date_part('week', date32), date_part('day', date32), date_part('hour', date32), date_part('minute', date32), date_part('second', date32), date_part('millisecond', date32), date_part('microsecond', date32), date_part('nanosecond', date32) FROM source_dt; +---- +1970 1 1 1 0 0 0 0 0 0 +2020 6 23 2 0 0 0 0 0 0 +2026 2 9 28 0 0 0 0 0 0 +NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL + +query IIIIIIIIII +SELECT date_part('year', date64), date_part('month', date64), date_part('week', date64), date_part('day', date64), date_part('hour', date64), date_part('minute', date64), date_part('second', date64), date_part('millisecond', date64), date_part('microsecond', date64), date_part('nanosecond', date64) FROM source_dt; +---- +1970 1 1 1 0 0 0 0 0 0 +2020 6 23 2 0 0 0 0 0 0 +2026 2 9 28 0 0 0 0 0 0 +NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL + +statement ok +drop table source_dt; + # just some floating point stuff happening in the result here query I SELECT date_part('microsecond', arrow_cast('23:32:50.123456789'::time, 'Time64(Nanosecond)')) diff --git a/datafusion/sqllogictest/test_files/datetime/dates.slt b/datafusion/sqllogictest/test_files/datetime/dates.slt index d2a7360b120c6..a6a5f480f72e2 100644 --- a/datafusion/sqllogictest/test_files/datetime/dates.slt +++ b/datafusion/sqllogictest/test_files/datetime/dates.slt @@ -139,6 +139,12 @@ SELECT to_date('01-14-2023 01:01:30+05:30', '%q', '%d-%m-%Y %H/%M/%S', '%+', '%m ---- 2023-01-13 +# Formatted pre-epoch datetimes retain their calendar date +query D +SELECT to_date('1969-12-31 12:00:00', '%Y-%m-%d %H:%M:%S'); +---- +1969-12-31 + statement error DataFusion error: Execution error: to_date function unsupported data type at index 1: List SELECT to_date('2022-08-03T14:38:50+05:30', make_array('%s', '%q', '%d-%m-%Y %H:%M:%S%#z', '%+')); @@ -298,6 +304,35 @@ SELECT to_date('2020-09-08 12/00/00+00:00', '%c', '%+') query error DataFusion error: Execution error: Error parsing timestamp from '2020\-09\-08 12/00/00\+00:00' using format '%q': trailing input SELECT to_date('2020-09-08 12/00/00+00:00', '%q') +# NULL string scalar inputs and all-NULL scalar formats return NULL +query DD +SELECT + to_date(NULL::VARCHAR, '%Y-%m-%d'), + to_date('2020-09-08', NULL::VARCHAR, NULL::VARCHAR) +---- +NULL NULL + +# NULL array formats are skipped for each row; rows with no usable format return NULL +query ID +SELECT id, to_date(value, format1, format2) +FROM ( + VALUES + (1, '2020-09-08', NULL::VARCHAR, '%Y-%m-%d'), + (2, '2020-09-08', NULL::VARCHAR, NULL::VARCHAR) +) AS t(id, value, format1, format2) +ORDER BY id +---- +1 2020-09-08 +2 NULL + +# Skipping NULL formats does not mask a later parse error. +query error DataFusion error: Execution error: Error parsing timestamp from '2020\-09\-08' using format '%q': trailing input +SELECT to_date('2020-09-08', NULL::VARCHAR, '%q') + +# Invalid format types are rejected before NULL input propagation. +query error DataFusion error: Execution error: to_date function unsupported data type at index 1: Int64 +SELECT to_date(NULL::VARCHAR, 12345) + statement ok create table ts_utf8_data(ts varchar(100), format varchar(100)) as values ('2020-09-08 12/00/00+00:00', '%Y-%m-%d %H/%M/%S%#z'), diff --git a/datafusion/sqllogictest/test_files/datetime/timestamps.slt b/datafusion/sqllogictest/test_files/datetime/timestamps.slt index d6e50f560aaf0..d73bc6eb06de8 100644 --- a/datafusion/sqllogictest/test_files/datetime/timestamps.slt +++ b/datafusion/sqllogictest/test_files/datetime/timestamps.slt @@ -518,6 +518,35 @@ SELECT COUNT(*) FROM ts_data_secs where ts > to_timestamp_seconds('2020-09-08 12 ---- 2 +# NULL string scalar inputs and all-NULL scalar formats return NULL +query PP +SELECT + to_timestamp(NULL::VARCHAR, '%Y-%m-%d'), + to_timestamp('2020-09-08', NULL::VARCHAR, NULL::VARCHAR) +---- +NULL NULL + +# NULL array formats are skipped for each row; rows with no usable format return NULL +query IP +SELECT id, to_timestamp(value, format1, format2) +FROM ( + VALUES + (1, '2020-09-08', NULL::VARCHAR, '%Y-%m-%d'), + (2, '2020-09-08', NULL::VARCHAR, NULL::VARCHAR) +) AS t(id, value, format1, format2) +ORDER BY id +---- +1 2020-09-08T00:00:00 +2 NULL + +# to_unixtime uses the same formatted string parsing path +query II +SELECT + to_unixtime(NULL::VARCHAR, '%Y-%m-%d'), + to_unixtime('2020-09-08', NULL::VARCHAR, NULL::VARCHAR) +---- +NULL NULL + # to_timestamp float inputs query PPP @@ -595,6 +624,12 @@ SELECT to_timestamp(arrow_cast(123456789.123456789, 'Decimal128(18,9)')) as c1, ---- 1973-11-29T21:33:09.123456784 1970-01-01T00:00:00.123456789 1970-01-01T00:00:00.123456789 +# Regression test for https://github.com/apache/datafusion/issues/22213 +query error .*overflows timestamp nanoseconds +SELECT to_timestamp( + arrow_cast('99999999999999999999999999999999999999', 'Decimal128(38,0)') +); + # from_unixtime @@ -1229,6 +1264,12 @@ SELECT DATE_BIN('5 month', '2022-01-01T00:00:00Z'); ---- 2021-09-01T00:00:00 +# test with utf8view +query P +SELECT DATE_BIN(arrow_cast('5 month', 'Utf8View'), '2022-01-01T00:00:00Z'); +---- +2021-09-01T00:00:00 + # month interval with default start time query P SELECT DATE_BIN('1 month', '2022-01-01 00:00:00Z'); @@ -1930,10 +1971,8 @@ SELECT '2000-01-01T00:00:00'::timestamp - '2010-01-01T00:00:00'::timestamp; -3653 days 0 hours 0 mins 0.000000000 secs # Interval - Timestamp => error -# statement error DataFusion error: Error during planning: Cannot coerce arithmetic expression Interval\(MonthDayNano\) \- Timestamp\(Nanosecond, None\) to valid types -# TODO: This query should raise error -# query P -# SELECT i - ts1 from FOO; +query error Cannot coerce arithmetic expression Interval\(MonthDayNano\) - Timestamp\(ns\) to valid types +SELECT i - ts1 from FOO; statement ok drop table foo; @@ -2474,6 +2513,51 @@ SELECT TIMESTAMPTZ '2020-01-01 00:00:00Z' = TIMESTAMP '2020-01-01' ---- true +query BBB +SELECT + arrow_cast(TIMESTAMP '2024-01-01 00:00:00', 'Timestamp(Second, None)') = + arrow_cast(TIMESTAMP '2024-01-01 00:00:00.123', 'Timestamp(Millisecond, None)'), + arrow_cast(TIMESTAMP '2024-01-01 00:00:00', 'Timestamp(Second, None)') = + arrow_cast(TIMESTAMP '2024-01-01 00:00:00.000', 'Timestamp(Millisecond, None)'), + arrow_cast(TIMESTAMP '2024-01-01 00:00:00', 'Timestamp(Second, None)') < + arrow_cast(TIMESTAMP '2024-01-01 00:00:00.123', 'Timestamp(Millisecond, None)') +---- +false true true + +query ? +SELECT + arrow_cast(TIMESTAMP '2024-01-01 00:00:00.123', 'Timestamp(Millisecond, None)') - + arrow_cast(TIMESTAMP '2024-01-01 00:00:00', 'Timestamp(Second, None)') +---- +0 days 0 hours 0 mins 0.123 secs + +query TP +SELECT arrow_typeof(ts), ts +FROM ( + SELECT arrow_cast(TIMESTAMP '2024-01-01 00:00:00', 'Timestamp(Second, None)') AS ts + UNION ALL + SELECT arrow_cast(TIMESTAMP '2024-01-01 00:00:00.123', 'Timestamp(Millisecond, None)') AS ts +) +ORDER BY ts +---- +Timestamp(ms) 2024-01-01T00:00:00 +Timestamp(ms) 2024-01-01T00:00:00.123 + +query TP +SELECT + arrow_typeof( + coalesce( + arrow_cast(NULL, 'Timestamp(Second, None)'), + arrow_cast(TIMESTAMP '2024-01-01 00:00:00.123', 'Timestamp(Millisecond, None)') + ) + ), + coalesce( + arrow_cast(NULL, 'Timestamp(Second, None)'), + arrow_cast(TIMESTAMP '2024-01-01 00:00:00.123', 'Timestamp(Millisecond, None)') + ) +---- +Timestamp(ms) 2024-01-01T00:00:00.123 + # verify timestamp cast with integer input query PPPPPP SELECT to_timestamp(null), to_timestamp(0), to_timestamp(1926632005), to_timestamp(1), to_timestamp(-1), to_timestamp(0-1) @@ -3445,6 +3529,28 @@ select to_time(time_str) from time_strings; statement ok drop table time_strings; +# Table input with multiple formats +# `%Q` is intentionally invalid; subsequent formats should still be tried. +query D rowsort +select to_time( + time_str, + '%Q', + '%H:%M:%S', + '%H-%M-%S', + '%H/%M/%S' +) from ( + values + ('12:30:45'), + ('14-25-30'), + ('09/05/01'), + (NULL) +) as formatted_time_strings(time_str); +---- +09:05:01 +12:30:45 +14:25:30 +NULL + # Error cases query error Error parsing 'not_a_time' as time @@ -3953,17 +4059,17 @@ true query ? select arrow_cast('2024-06-17T11:00:00', 'Timestamp(Nanosecond, Some("UTC"))') - arrow_cast('2024-06-17T12:00:00', 'Timestamp(Microsecond, Some("UTC"))'); ---- -0 days -1 hours 0 mins 0.000000 secs +0 days -1 hours 0 mins 0.000000000 secs query ? select arrow_cast('2024-06-17T13:00:00', 'Timestamp(Nanosecond, Some("+00:00"))') - arrow_cast('2024-06-17T12:00:00', 'Timestamp(Microsecond, Some("UTC"))'); ---- -0 days 1 hours 0 mins 0.000000 secs +0 days 1 hours 0 mins 0.000000000 secs query ? select arrow_cast('2024-06-17T13:00:00', 'Timestamp(Nanosecond, Some("UTC"))') - arrow_cast('2024-06-17T12:00:00', 'Timestamp(Microsecond, Some("+00:00"))'); ---- -0 days 1 hours 0 mins 0.000000 secs +0 days 1 hours 0 mins 0.000000000 secs # not supported: coercion across timezones query error @@ -5325,9 +5431,28 @@ SELECT to_timestamp(arrow_cast(-9223372036, 'Int64')); 1677-09-21T00:12:44 # Overflow error when value exceeds valid range -query error Arithmetic overflow +query error converted value exceeds the representable i64 range SELECT to_timestamp(arrow_cast(9223372037, 'Int64')); +# TRY_CAST returns NULL for timestamp/date casts that overflow +query P +SELECT TRY_CAST(arrow_cast(9223372037, 'Timestamp(s)') AS TIMESTAMP(9)); +---- +NULL + +query P +SELECT TRY_CAST(DATE '3000-01-01' AS TIMESTAMP(9)); +---- +NULL + +query P +SELECT TRY_CAST(ts AS TIMESTAMP(9)) AS ts +FROM ( + VALUES (arrow_cast(9223372037, 'Timestamp(s)')) +) t(ts); +---- +NULL + # Float truncation behavior query P SELECT to_timestamp_seconds(arrow_cast(-1.9, 'Float64')); @@ -5339,6 +5464,13 @@ SELECT to_timestamp_millis(arrow_cast(-1.9, 'Float64')); ---- 1969-12-31T23:59:59.999 +# Regression test for https://github.com/apache/datafusion/issues/22214 +query error .*out of range after truncating to year +SELECT date_trunc( + 'year', + arrow_cast(TIMESTAMP '1677-09-22 00:00:00', 'Timestamp(Nanosecond, None)') +); + ########## ## Common timestamp data diff --git a/datafusion/sqllogictest/test_files/ddl.slt b/datafusion/sqllogictest/test_files/ddl.slt index 82c30e9aba386..672bab553330d 100644 --- a/datafusion/sqllogictest/test_files/ddl.slt +++ b/datafusion/sqllogictest/test_files/ddl.slt @@ -200,10 +200,6 @@ SELECT foo_schema.bar.a FROM foo_schema.bar; ---- 1 -# TODO: Drop schema for cleanup, see #6027 -# statement ok -# DROP SCHEMA foo_schema; - ########## # Drop view error tests ########## @@ -654,7 +650,7 @@ LOCATION 'test_files/scratch/ddl/test_table'; query TTT DESCRIBE aggregate_table; ---- -id Int64 YES +id Int64 NO # Should insert into an empty table statement ok @@ -983,6 +979,20 @@ CREATE TABLE dup_src AS VALUES(1, 2); statement error DataFusion error: Schema error: Schema contains duplicate unqualified field name column1 CREATE TABLE dup_ctas AS SELECT * FROM dup_src LEFT JOIN dup_src y ON dup_src.column1 = y.column2; +statement ok +CREATE TABLE dup_ctas_with_schema(left_c1 bigint, right_c1 bigint) AS +SELECT dup_src.column1, right_src.column1 +FROM dup_src +CROSS JOIN (SELECT column2 AS column1 FROM dup_src) right_src; + +query II +SELECT left_c1, right_c1 FROM dup_ctas_with_schema; +---- +1 2 + +statement ok +DROP TABLE dup_ctas_with_schema; + statement error DataFusion error: Schema error: Schema contains duplicate unqualified field name column1 CREATE VIEW dup_view AS SELECT * FROM dup_src LEFT JOIN dup_src y ON dup_src.column1 = y.column2; diff --git a/datafusion/sqllogictest/test_files/decimal.slt b/datafusion/sqllogictest/test_files/decimal.slt index 5faf801c84652..4335ec06685f2 100644 --- a/datafusion/sqllogictest/test_files/decimal.slt +++ b/datafusion/sqllogictest/test_files/decimal.slt @@ -1046,17 +1046,17 @@ SELECT log(10, arrow_cast(1 , 'Decimal32(5, 1)')) query RT SELECT power(2::decimal(38, 0), 4), arrow_typeof(power(2::decimal(38, 0), 4)); ---- -16 Decimal128(38, 0) +16 Float64 query RT SELECT power(10000000000::decimal(38, 0), 2), arrow_typeof(power(10000000000::decimal(38, 0), 2)); ---- -100000000000000000000 Decimal128(38, 0) +100000000000000000000 Float64 query R SELECT power(2.5, 4) ---- -39 +39.0625 query R SELECT power(2.5, 1) @@ -1093,76 +1093,107 @@ SELECT power(2, 100000000000) ---- Infinity -# Negative exponent now works (fallback to f64) +# Negative exponent returns Float64 so fractional results are representable query RT SELECT power(2::decimal(38, 0), -5), arrow_typeof(power(2::decimal(38, 0), -5)); ---- -0 Decimal128(38, 0) +0.03125 Float64 + +query RT +SELECT power(CAST(2 AS DECIMAL(10, 0)), -3), arrow_typeof(power(CAST(2 AS DECIMAL(10, 0)), -3)); +---- +0.125 Float64 -# Negative exponent with scale preserves decimal places query RT SELECT power(4::decimal(38, 5), -1), arrow_typeof(power(4::decimal(38, 5), -1)); ---- -0.25 Decimal128(38, 5) +0.25 Float64 + +query IRT +SELECT exponent, power(2::decimal(10, 0), exponent), arrow_typeof(power(2::decimal(10, 0), exponent)) +FROM (VALUES (-3), (3)) AS t(exponent) +ORDER BY exponent; +---- +-3 0.125 Float64 +3 8 Float64 -# Expected to have `16 Decimal128(38, 0)` -# Due to type coericion, it becomes Float -> Float -> Float query RT SELECT power(2::decimal(38, 0), 4), arrow_typeof(power(2::decimal(38, 0), 4)); ---- -16 Decimal128(38, 0) +16 Float64 -# Arbitrary scale query RT SELECT power(2.5::decimal(38, 3), 4), arrow_typeof(power(2.5::decimal(38, 3), 4)); ---- -39.062 Decimal128(38, 3) +39.0625 Float64 + +# https://github.com/apache/datafusion/issues/22480 +query RT +SELECT power(2.5::decimal(20, 4), 10), arrow_typeof(power(2.5::decimal(20, 4), 10)); +---- +9536.7431640625 Float64 query RT SELECT power(2.5, 4.0), arrow_typeof(power(2.5, 4.0)); ---- -39 Decimal128(2, 1) +39.0625 Float64 -# Non-integer exponent now works (fallback to f64) query RT SELECT power(2.5, 4.2), arrow_typeof(power(2.5, 4.2)); ---- -46.9 Decimal128(2, 1) +46.9189232024 Float64 -query error Compute error: Cannot use non-finite exp: NaN -SELECT power(2::decimal(38, 0), arrow_cast('NaN','Float64')) +query RT +SELECT power(2::decimal(38, 0), arrow_cast('NaN','Float64')), + arrow_typeof(power(2::decimal(38, 0), arrow_cast('NaN','Float64'))); +---- +NaN Float64 -query error Compute error: Cannot use non-finite exp: inf -SELECT power(2::decimal(38, 0), arrow_cast('INF','Float64')) +query RT +SELECT power(2::decimal(38, 0), arrow_cast('INF','Float64')), + arrow_typeof(power(2::decimal(38, 0), arrow_cast('INF','Float64'))); +---- +Infinity Float64 -# Floating above u32::max now works (fallback to f64, returns infinity which is an error) -query error Arrow error: Arithmetic overflow: Result of 2\^5000000000.1 is not finite -SELECT power(2::decimal(38, 0), 5000000000.1) +# Result overflows finite Float64 range +query RT +SELECT power(2::decimal(38, 0), 5000000000.1), + arrow_typeof(power(2::decimal(38, 0), 5000000000.1)); +---- +Infinity Float64 -# Integer Above u32::max - still goes through integer path which fails -query error Arrow error: Arithmetic overflow: Unsupported exp value -SELECT power(2::decimal(38, 0), 5000000000) +# Integer above u32::max uses the Float64 decimal/int path +query RT +SELECT power(2::decimal(38, 0), 5000000000), + arrow_typeof(power(2::decimal(38, 0), 5000000000)); +---- +Infinity Float64 -query ?T +query RT SELECT power(arrow_cast(2, 'Decimal32(5, 0)'), 4), arrow_typeof(power(arrow_cast(2, 'Decimal32(5, 0)'), 4)); ---- -16 Decimal32(5, 0) +16 Float64 -query ?T +query RT SELECT power(arrow_cast(2, 'Decimal64(5, 0)'), 4), arrow_typeof(power(arrow_cast(2, 'Decimal64(5, 0)'), 4)); ---- -16 Decimal64(5, 0) +16 Float64 query RT SELECT power(2::decimal(76, 0), 4), arrow_typeof(power(2::decimal(76, 0), 4)); ---- -16 Decimal256(76, 0) +16 Float64 query R SELECT power(2.0, null) ---- NULL +query RT +SELECT power(2::decimal(38, 0), null), arrow_typeof(power(2::decimal(38, 0), null)); +---- +NULL Float64 + # Array variants of power function query RR rowsort SELECT distinct c1*100000, power(c1*100000, 2) from decimal_simple; @@ -1206,7 +1237,7 @@ select log(100000000000000000000000000000000000::decimal(38,0)) ---- 35 -# Result is decimal since argument is decimal regardless decimals-as-floats parsing +# Decimal x Int64 returns Float64 regardless of decimals-as-floats parsing query R SELECT power(10000000000::decimal(38, 0), 2); ---- @@ -1216,7 +1247,7 @@ query RT SELECT power(10000000000::decimal(38, 0), 2), arrow_typeof(power(10000000000::decimal(38, 0), 2)); ---- -100000000000000000000 Decimal128(38, 0) +100000000000000000000 Float64 query R SELECT power(2.5, 4.0) @@ -1260,3 +1291,52 @@ ORDER BY c1; statement ok DROP TABLE decimal_div_mismatch; + +# Regression tests: `avg` of a decimal column must accumulate its intermediate +# sum in a type wide enough not to overflow the input's native type. Each row +# count below is chosen so that the sum just exceeds the input's native maximum +# and would silently wrap if accumulated unwidened. + +# 21476 * 99999 = 2,147,578,524 > i32::MAX +query RT +select avg(d), arrow_typeof(avg(d)) +from ( + select arrow_cast(99999.0, 'Decimal32(5, 0)') as d + from generate_series(1, 21476) +) t; +---- +99999 Decimal32(9, 4) + +# 92235 * 99999999999999 ~= 9.22e18 > i64::MAX +query RT +select avg(d), arrow_typeof(avg(d)) +from ( + select arrow_cast('99999999999999', 'Decimal64(14, 0)') as d + from generate_series(1, 92235) +) t; +---- +99999999999999 Decimal64(18, 4) + +# 21476 * (10^34 - 1) ~= 2.15e38 > i128::MAX +query RT +select avg(d), arrow_typeof(avg(d)) +from ( + select arrow_cast('9999999999999999999999999999999999', 'Decimal128(34, 0)') as d + from generate_series(1, 21476) +) t; +---- +9999999999999999999999999999999999 Decimal128(38, 4) + +# Regression: `avg(DISTINCT ...)` must widen its intermediate sum the same way. +# The second distinct aggregate keeps `single_distinct_to_group_by` from +# rewriting the plan, so the distinct accumulator is the code under test. + +# sum(1..65536) = 2,147,516,416 > i32::MAX, avg = 32768.5 +query RTR +select avg(distinct d), arrow_typeof(avg(distinct d)), avg(distinct v) +from ( + select arrow_cast(v, 'Decimal32(9, 0)') as d, v + from generate_series(1, 65536) t(v) +) t; +---- +32768.5 Decimal32(9, 4) 32768.5 diff --git a/datafusion/sqllogictest/test_files/delete.slt b/datafusion/sqllogictest/test_files/delete.slt index 6131d6db3d5f7..1f33360824393 100644 --- a/datafusion/sqllogictest/test_files/delete.slt +++ b/datafusion/sqllogictest/test_files/delete.slt @@ -79,7 +79,7 @@ physical_plan # Deleting by columns that do not exist returns an error -query error DataFusion error: Schema error: No field named e. Valid fields are t1.a, t1.b, t1.c, t1.d. +query error DataFusion error: Schema error: No field named e\.\nValid fields are t1.a, t1.b, t1.c, t1.d. explain delete from t1 where e = 1; diff --git a/datafusion/sqllogictest/test_files/describe.slt b/datafusion/sqllogictest/test_files/describe.slt index 88347965c67a5..083a33657f0a3 100644 --- a/datafusion/sqllogictest/test_files/describe.slt +++ b/datafusion/sqllogictest/test_files/describe.slt @@ -142,3 +142,44 @@ name_count Int64 NO # Describing a statement that's not a query is not supported statement error Describing statements other than SELECT not supported DESCRIBE CREATE TABLE test_desc_table (id INT, name VARCHAR); + +########## +# VALUES nullability inference +# +# Inferred VALUES schemas should be non-nullable when no row contributes a +# NULL in that column position, and nullable when at least one does. +########## + +# All-non-null VALUES: every column is non-nullable. +query TTT rowsort +DESCRIBE SELECT * FROM (VALUES (1, 'a'), (2, 'b')) AS t(x, y); +---- +x Int64 NO +y Utf8 NO + +# Untyped NULL in one row makes that column nullable; sibling all-non-null +# columns remain non-nullable. +query TTT rowsort +DESCRIBE SELECT * FROM (VALUES (1, 'a'), (NULL, 'b')) AS t(x, y); +---- +x Int64 YES +y Utf8 NO + +# Typed NULL has the same effect as untyped NULL on column nullability. +query TTT +DESCRIBE SELECT * FROM (VALUES (1), (CAST(NULL AS BIGINT))) AS t(x); +---- +x Int64 YES + +# All-NULL column is nullable; the inferred type is Null. +query TTT +DESCRIBE SELECT * FROM (VALUES (NULL), (NULL)) AS t(x); +---- +x Null YES + +# A Null-typed value sourced from a non-nullable expression must still +# produce a nullable column: a DataType::Null field is always nullable. +query TTT +DESCRIBE SELECT * FROM (VALUES (arrow_cast(1, 'Null'))) AS t(x); +---- +x Null YES diff --git a/datafusion/sqllogictest/test_files/dictionary.slt b/datafusion/sqllogictest/test_files/dictionary.slt index 92e6c41835d75..f314254955824 100644 --- a/datafusion/sqllogictest/test_files/dictionary.slt +++ b/datafusion/sqllogictest/test_files/dictionary.slt @@ -80,12 +80,12 @@ SELECT * FROM m1; query TTT DESCRIBE m1; ---- -tag_id Dictionary(Int32, Utf8) YES -f1 Float64 YES -f2 Utf8 YES -f3 Utf8 YES -f4 Float64 YES -time Timestamp(ns) YES +tag_id Dictionary(Int32, Utf8) NO +f1 Float64 NO +f2 Utf8 NO +f3 Utf8 NO +f4 Float64 NO +time Timestamp(ns) NO # in list with dictionary input query BBB @@ -154,10 +154,10 @@ passive 1000 1000 2023-12-04T01:30:00 query TTT DESCRIBE m2; ---- -type Dictionary(Int32, Utf8) YES -tag_id Dictionary(Int32, Utf8) YES -f5 Float64 YES -time Timestamp(ns) YES +type Dictionary(Int32, Utf8) NO +tag_id Dictionary(Int32, Utf8) NO +f5 Float64 NO +time Timestamp(ns) NO query I select count(*) from m1 where tag_id = '1000' and time < '2024-01-03T14:46:35+01:00'; @@ -492,10 +492,10 @@ LOCATION 'test_files/scratch/dictionary/dict_hash_10.parquet'; query TTT DESCRIBE dict_hash_10; ---- -id Int64 YES -payload_hash Dictionary(Int32, Utf8) YES -metric Float64 YES -ts Timestamp(ns) YES +id Int64 NO +payload_hash Dictionary(Int32, Utf8) NO +metric Float64 NO +ts Timestamp(ns) NO query II SELECT COUNT(*), COUNT(DISTINCT payload_hash) @@ -515,3 +515,139 @@ DROP TABLE dict_hash_10; statement ok DROP TABLE dict_hash_src; + +statement ok +CREATE TABLE dict_large_utf8 AS +SELECT + arrow_cast(column1, 'Dictionary(Int32, LargeUtf8)') AS tag, + arrow_cast(column2, 'Float64') AS val +FROM (VALUES ('alpha', 1.0), ('beta', 2.0), ('alpha', 3.0), ('gamma', 4.0), ('beta', 5.0), (NULL, 6.0)); + +query TRI rowsort +SELECT tag, SUM(val), COUNT(*) FROM dict_large_utf8 GROUP BY tag; +---- +NULL 6 1 +alpha 4 2 +beta 7 2 +gamma 4 1 + +statement ok +DROP TABLE dict_large_utf8; + +# multiple dictionary columns as group keys + +statement ok +CREATE TABLE dict_multi_key AS +SELECT + arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS region, + arrow_cast(column2, 'Dictionary(Int16, Utf8)') AS status, + arrow_cast(column3, 'Dictionary(Int8, Utf8)') AS tier, + arrow_cast(column4, 'Float64') AS amount +FROM ( + VALUES + ('us', 'active', 'gold', 100.0), + ('eu', 'active', 'silver', 200.0), + ('us', 'active', 'gold', 300.0), + ('eu', 'inactive', 'gold', 400.0), + ('us', 'inactive', 'silver', 500.0), + ('eu', 'inactive', 'gold', 600.0), + ('us', 'active', 'silver', 150.0), + ('eu', 'active', 'silver', 250.0), + (NULL, 'active', 'gold', 700.0) +); + +query TTTRI rowsort +SELECT region, status, tier, SUM(amount), COUNT(*) FROM dict_multi_key GROUP BY region, status, tier; +---- +NULL active gold 700 1 +eu active silver 450 2 +eu inactive gold 1000 2 +us active gold 400 2 +us active silver 150 1 +us inactive silver 500 1 + +statement ok +DROP TABLE dict_multi_key; + +# mixed dict (2) and non-dict (2) group keys with nulls spread across all columns + +statement ok +CREATE TABLE dict_mixed_nulls AS +SELECT + arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS region, + arrow_cast(column2, 'Dictionary(Int16, Utf8)') AS status, + arrow_cast(column3, 'Utf8') AS tier, + arrow_cast(column4, 'Int32') AS category, + arrow_cast(column5, 'Float64') AS amount +FROM ( + VALUES + ('us', 'active', 'gold', 1, 100.0), + ('us', 'active', 'gold', 1, 200.0), + ('eu', 'active', 'silver', 2, 300.0), + (NULL, 'active', 'gold', 1, 400.0), + ('us', NULL, 'gold', 1, 500.0), + ('us', 'active', NULL, 1, 600.0), + ('us', 'active', 'gold', NULL, 700.0) +); + +query TTTIRI rowsort +SELECT region, status, tier, category, SUM(amount), COUNT(*) FROM dict_mixed_nulls GROUP BY region, status, tier, category; +---- +NULL active gold 1 400 1 +eu active silver 2 300 1 +us NULL gold 1 500 1 +us active NULL 1 600 1 +us active gold 1 300 2 +us active gold NULL 700 1 + +statement ok +DROP TABLE dict_mixed_nulls; + +########## +## Aggregation tests: COUNT(DISTINCT) on dictionary columns +########## + +statement ok +CREATE TABLE dict_count_distinct AS +SELECT + arrow_cast(column1, 'Dictionary(Int64, Utf8)') AS region, + arrow_cast(column2, 'Dictionary(Int8, Utf8)') AS sensor +FROM ( + VALUES + ('north', 's1'), ('north', 's2'), ('north', 's1'), ('north', 's3'), + ('south', 's3'), ('south', 's3'), ('south', 's4'), + ('east', 's1'), + (NULL, 's5'), + ('north', NULL) +); + +query TI rowsort +SELECT region, COUNT(DISTINCT sensor) FROM dict_count_distinct GROUP BY region; +---- +NULL 1 +east 1 +north 3 +south 2 + +statement ok +DROP TABLE dict_count_distinct; + +# same dictionary type but value order differs across batches so key ids refer to different strings; +# grouping must use the logical value, not the raw key id +query TI rowsort +WITH + first_batch AS ( + SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS region + FROM (VALUES ('west'), ('west'), ('west'), ('east'), (NULL)) AS t(column1) + ), + second_batch AS ( + SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS region + FROM (VALUES ('east'), ('east'), ('east'), ('west'), (NULL)) AS t(column1) + ) +SELECT region, count(*) +FROM (SELECT region FROM first_batch UNION ALL SELECT region FROM second_batch) +GROUP BY region; +---- +NULL 2 +east 4 +west 4 diff --git a/datafusion/sqllogictest/test_files/distinct_on.slt b/datafusion/sqllogictest/test_files/distinct_on.slt index 5b18915080f8f..0659b9c208f9c 100644 --- a/datafusion/sqllogictest/test_files/distinct_on.slt +++ b/datafusion/sqllogictest/test_files/distinct_on.slt @@ -195,3 +195,233 @@ RESET datafusion.explain.logical_plan_only; statement ok drop table t; + +# DISTINCT ON combined with GROUP BY + aggregation (issue #17256). +# ON references a grouping column; ORDER BY uses an aggregate alias. +query TII +SELECT DISTINCT ON (c1) c1, c3, max(c4) AS agg2 +FROM aggregate_test_100 GROUP BY c1, c3 ORDER BY c1, agg2; +---- +a 65 -28462 +b -60 -21739 +c 3 -30508 +d 102 -24558 +e -56 -31500 + +# DISTINCT ON referencing a SELECT alias for an aggregate. +query TI +SELECT DISTINCT ON (agg2) c1, max(c4) AS agg2 +FROM aggregate_test_100 GROUP BY c1 ORDER BY agg2; +---- +b 25286 +c 29106 +d 31106 +a 32064 +e 32514 + +# DISTINCT ON with a scalar function over a grouping column. +query TI +SELECT DISTINCT ON (upper(c1)) c1, sum(c3) FROM aggregate_test_100 +GROUP BY c1 ORDER BY upper(c1); +---- +a -385 +b -111 +c -28 +d 458 +e 847 + +# Hidden ORDER BY tie-breaker: c3 is in GROUP BY and ORDER BY but +# NOT in the SELECT list. PostgreSQL accepts this. +query TI +SELECT DISTINCT ON (c1) c1, max(c4) AS agg2 +FROM aggregate_test_100 GROUP BY c1, c3 ORDER BY c1, c3; +---- +a 11640 +b 19316 +c -30187 +d 5613 +e 13611 + +# Hidden DISTINCT ON key: ON references a grouping column that is +# NOT in the SELECT list. ORDER BY adds a deterministic tie-breaker. +query II +SELECT DISTINCT ON (c1) c2 % 2, count(*) AS n +FROM aggregate_test_100 GROUP BY c1, c2 % 2 ORDER BY c1, c2 % 2; +---- +0 7 +0 9 +0 11 +0 6 +0 12 + +# Raw aggregate expression in DISTINCT ON (not via an alias). +query TI +SELECT DISTINCT ON (sum(c3)) c1, sum(c3) AS total +FROM aggregate_test_100 GROUP BY c1 ORDER BY sum(c3); +---- +a -385 +b -111 +c -28 +d 458 +e 847 + +# DISTINCT ON with HAVING. +query TI +SELECT DISTINCT ON (c1) c1, count(*) AS cnt FROM aggregate_test_100 +GROUP BY c1 HAVING count(*) > 10 ORDER BY c1, cnt DESC; +---- +a 21 +b 19 +c 21 +d 18 +e 21 + +# DISTINCT ON combined with a window function over a unique ordering +# key, so the test is fully deterministic. +query II +WITH t(id, v) AS (VALUES (1, 10), (2, 20), (3, 10), (4, 30), (5, 20)) +SELECT DISTINCT ON (v) v, row_number() OVER (ORDER BY id) AS rn +FROM t ORDER BY v, rn; +---- +10 1 +20 2 +30 4 + +# Raw window expression in DISTINCT ON (not via an alias). Uses a +# unique ordering key so row_number is deterministic. +query II +WITH t(id, v) AS (VALUES (1, 10), (2, 20), (3, 30), (4, 40)) +SELECT DISTINCT ON (row_number() OVER (ORDER BY id)) id, v +FROM t +ORDER BY row_number() OVER (ORDER BY id); +---- +1 10 +2 20 +3 30 +4 40 + +# Qualified join columns with potential alias conflict. +query TII +WITH t1(k, v) AS (VALUES ('x', 1), ('x', 2), ('y', 3)), + t2(k, w) AS (VALUES ('x', 10), ('y', 20)) +SELECT DISTINCT ON (t1.k) t1.k, sum(t1.v) AS s, max(t2.w) AS mw +FROM t1 JOIN t2 ON t1.k = t2.k +GROUP BY t1.k ORDER BY t1.k; +---- +x 3 10 +y 3 20 + +# DISTINCT ON name conflicts with an input column of the same name. +# PostgreSQL resolves `b` to the SELECT alias `a AS b`, not the input +# column `t.b`. Groups should be keyed by `a`, not `t.b`. +query TI +WITH t(a, b) AS (VALUES ('x', 1), ('x', 2), ('y', 1)) +SELECT DISTINCT ON (b) a AS b, count(*) AS n +FROM t GROUP BY a ORDER BY b; +---- +x 2 +y 1 + +# A bare alias resolves to the SELECT expression, but inside a larger +# expression the same identifier refers to the input column. Postgres: +# ORDER BY a, b + 0 DESC +# uses `a` (post-aggregate) and `t.b + 0`, so for a=100 the row with +# t.b=2 wins (sum=2). DataFusion must not recursively swap `b` inside +# `b + 0` for the alias. +query II +WITH t(a, b) AS (VALUES (100, 1), (100, 2), (200, 1)) +SELECT DISTINCT ON (a) a AS b, sum(b) AS s +FROM t GROUP BY a, b ORDER BY a, b + 0 DESC; +---- +100 2 +200 1 + +# A nested ORDER BY expression over a SELECT alias is still rejected in +# the post-aggregate DISTINCT ON path. +query error No field named agg2 +WITH t(a, b) AS (VALUES (1, 10), (1, 20), (2, 30)) +SELECT DISTINCT ON (a) a, max(b) AS agg2 +FROM t GROUP BY a ORDER BY a, agg2 + 1; + +# DISTINCT ON after aggregation still needs the SELECT-list unnest +# rewrite, and bare ORDER BY aliases should keep working against the +# rewritten DistinctOn input. +query II +WITH t(a, b) AS (VALUES (1, 10), (1, 20), (2, 30)) +SELECT DISTINCT ON (a) a, unnest(array_agg(b)) AS b +FROM t GROUP BY a ORDER BY a, b DESC; +---- +1 20 +2 30 + +# DISTINCT ON after aggregation also needs to keep multi-column SELECT +# expansions from struct unnest working. +query III +WITH t(a, b) AS (VALUES (1, 10), (1, 20), (2, 30)) +SELECT DISTINCT ON (a) a, unnest(struct(max(b), min(b))) +FROM t GROUP BY a ORDER BY a; +---- +1 20 10 +2 30 30 + +# DISTINCT ON keys that expand to multiple columns are rejected. +query error DISTINCT ON expressions that expand to multiple columns are not supported with DISTINCT ON +WITH t(a, b) AS (VALUES (1, 10), (2, 20)) +SELECT DISTINCT ON (unnest(struct(max(b), min(b)))) a +FROM t GROUP BY a ORDER BY unnest(struct(max(b), min(b))), a; + +# ORDER BY tie-breakers that expand to multiple columns are rejected. +query error ORDER BY expressions that expand to multiple columns are not supported with DISTINCT ON +WITH t(a, b) AS (VALUES (1, 10), (2, 20)) +SELECT DISTINCT ON (a) a, max(b) +FROM t GROUP BY a ORDER BY a, unnest(struct(max(b), min(b))); + +# Fast path (no aggregation): a bare ORDER BY alias must resolve to +# its underlying SELECT expression so that the sort attached to +# DistinctOn normalizes against the base plan. +query T +WITH t(a, b) AS (VALUES ('x', 1), ('x', 2), ('y', 3)) +SELECT DISTINCT ON (x) a AS x FROM t ORDER BY x; +---- +x +y + +# EXPLAIN for the post-aggregation case. +statement ok +set datafusion.explain.logical_plan_only = true; + +query TT +explain SELECT DISTINCT ON (c1) c1, max(c4) AS agg2 +FROM aggregate_test_100 GROUP BY c1 ORDER BY c1, agg2; +---- +logical_plan +01)Projection: first_value(aggregate_test_100.c1) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST] AS c1, first_value(agg2) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST] AS agg2 +02)--Sort: aggregate_test_100.c1 ASC NULLS LAST +03)----Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[first_value(aggregate_test_100.c1) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST], first_value(max(aggregate_test_100.c4) AS agg2) ORDER BY [aggregate_test_100.c1 ASC NULLS LAST, max(aggregate_test_100.c4) ASC NULLS LAST]]] +04)------Aggregate: groupBy=[[aggregate_test_100.c1]], aggr=[[max(aggregate_test_100.c4)]] +05)--------TableScan: aggregate_test_100 projection=[c1, c4] + +statement ok +RESET datafusion.explain.logical_plan_only; + +# Ordinal ORDER BY still works in the post-aggregate DISTINCT ON path. +query TI +SELECT DISTINCT ON (c1) c1, max(c4) +FROM aggregate_test_100 GROUP BY c1 ORDER BY 1, 2; +---- +a 32064 +b 25286 +c 29106 +d 31106 +e 32514 + +# Synthetic repro for issue #17256. +query TIR +WITH t(a, b, c) AS ( + VALUES ('x', 1, 10.0), ('x', 1, 20.0), ('y', 2, 30.0) +) +SELECT DISTINCT ON (a) a, b, sum(c) AS total +FROM t GROUP BY a, b ORDER BY a, total DESC; +---- +x 1 30 +y 2 30 diff --git a/datafusion/sqllogictest/test_files/dynamic_filter_pushdown_config.slt b/datafusion/sqllogictest/test_files/dynamic_filter_pushdown_config.slt index e779ce2cbffb0..eec6e5ae179bc 100644 --- a/datafusion/sqllogictest/test_files/dynamic_filter_pushdown_config.slt +++ b/datafusion/sqllogictest/test_files/dynamic_filter_pushdown_config.slt @@ -90,7 +90,7 @@ logical_plan 02)--TableScan: test_parquet projection=[id, value, name] physical_plan 01)SortExec: TopK(fetch=3), expr=[value@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/test_data.parquet]]}, projection=[id, value, name], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[value@1 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/test_data.parquet]]}, projection=[id, value, name], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[value@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible statement ok set datafusion.explain.analyze_level = summary; @@ -104,7 +104,7 @@ Plan with Metrics 03)----ProjectionExec: expr=[id@0 as id, value@1 as v, value@1 + id@0 as name], metrics=[output_rows=10, ] 04)------FilterExec: value@1 > 3, metrics=[output_rows=10, , selectivity=100% (10/10)] 05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, metrics=[output_rows=10, ] -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/test_data.parquet]]}, projection=[id, value], file_type=parquet, predicate=value@1 > 3 AND DynamicFilter [ value@1 IS NULL OR value@1 > 800 ], pruning_predicate=value_null_count@1 != row_count@2 AND value_max@0 > 3 AND (value_null_count@1 > 0 OR value_null_count@1 != row_count@2 AND value_max@0 > 800), required_guarantees=[], metrics=[output_rows=10, elapsed_compute=, output_bytes=80.0 B, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched -> 1 fully matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=210, metadata_load_time=, scan_efficiency_ratio=18.31% (210/1.15 K)] +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/test_data.parquet]]}, projection=[id, value], file_type=parquet, predicate=value@1 > 3 AND DynamicFilter [ value@1 IS NULL OR value@1 > 800 ], dynamic_rg_pruning=eligible, pruning_predicate=value_null_count@1 != row_count@2 AND value_max@0 > 3 AND (value_null_count@1 > 0 OR value_null_count@1 != row_count@2 AND value_max@0 > 800), required_guarantees=[], metrics=[output_rows=10, elapsed_compute=, output_bytes=80.0 B, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched -> 1 fully matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=210, page_index_load_skipped=1, row_groups_pruned_dynamic_filter=0, metadata_load_time=, scan_efficiency_ratio=18.31% (210/1.15 K)] statement ok set datafusion.explain.analyze_level = dev; @@ -157,7 +157,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], projection=[id@2, data@3, info@1] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id, info], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Disable Join dynamic filter pushdown statement ok @@ -235,7 +235,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, id@0)], projection=[id@2, data@3, info@1] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id, info], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # RIGHT JOIN correctness: all right rows appear, unmatched left rows produce NULLs query ITT @@ -284,7 +284,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(id@0, id@0)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # LEFT SEMI JOIN (physical LeftSemi): reverse table roles so optimizer keeps LeftSemi # (right_parquet has 3 rows < left_parquet has 5 rows, so no swap occurs). @@ -304,7 +304,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(id@0, id@0)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id, info], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # LEFT SEMI (physical LeftSemi) correctness: only right rows with matching left ids query IT rowsort @@ -338,7 +338,7 @@ physical_plan 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet 03)--SortExec: expr=[data@1 DESC], preserve_partitioning=[false] 04)----FilterExec: DynamicFilter [ empty ] -05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[data@1 DESC], reverse_row_groups=true +05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[data@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible statement count 0 SET datafusion.execution.parquet.pushdown_filters = true; @@ -361,7 +361,7 @@ physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(id@0, id@0)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet 03)--SortExec: expr=[data@1 DESC], preserve_partitioning=[false] -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[data@1 DESC], reverse_row_groups=true +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[data@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible statement count 0 RESET datafusion.execution.parquet.pushdown_filters; @@ -374,16 +374,16 @@ FROM left_parquet l WHERE l.id NOT IN (SELECT r.id FROM right_parquet r); ---- logical_plan -01)LeftAnti Join: l.id = __correlated_sq_1.id +01)LeftAnti Join: l.id = __correlated_sq_1.id null_aware 02)--SubqueryAlias: l 03)----TableScan: left_parquet projection=[id, data] 04)--SubqueryAlias: __correlated_sq_1 05)----SubqueryAlias: r 06)------TableScan: right_parquet projection=[id] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)] +01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)], null_aware 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet # LEFT MARK JOIN: the OR prevents decorrelation to LeftSemi, so the optimizer # uses LeftMark. Self-generated dynamic filter pushes to the probe side. @@ -407,7 +407,7 @@ physical_plan 02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 03)----HashJoinExec: mode=CollectLeft, join_type=LeftMark, on=[(id@0, id@0)] 04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id, info], file_type=parquet -05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ] +05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # LEFT MARK correctness: all right rows match EXISTS, so all 3 appear query IT rowsort @@ -444,8 +444,8 @@ logical_plan physical_plan 01)SortExec: TopK(fetch=2), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] 02)--HashJoinExec: mode=CollectLeft, join_type=RightSemi, on=[(id@0, id@0)] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ] -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Correctness check query IT @@ -457,10 +457,9 @@ ORDER BY l.id LIMIT 2; 1 left1 3 left3 -# ANTI JOIN with TopK parent: TopK generates a dynamic filter on `id` (join -# key) that pushes through the LeftAnti join to both the preserved and -# non-preserved sides. The HashJoin pushes the self-generated filter to the -# right hand side of the LeftAnti join. +# ANTI JOIN with TopK parent: the TopK dynamic filter on `id` is pushed only +# to the preserved output side. Filtering the non-output side can create +# anti-join output. query TT EXPLAIN SELECT l.* FROM left_parquet l @@ -469,7 +468,7 @@ ORDER BY l.id LIMIT 2; ---- logical_plan 01)Sort: l.id ASC NULLS LAST, fetch=2 -02)--LeftAnti Join: l.id = __correlated_sq_1.id +02)--LeftAnti Join: l.id = __correlated_sq_1.id null_aware 03)----SubqueryAlias: l 04)------TableScan: left_parquet projection=[id, data] 05)----SubqueryAlias: __correlated_sq_1 @@ -477,9 +476,9 @@ logical_plan 07)--------TableScan: right_parquet projection=[id] physical_plan 01)SortExec: TopK(fetch=2), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ] -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ] +02)--HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)], null_aware +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id], file_type=parquet # Correctness check query IT @@ -491,6 +490,40 @@ ORDER BY l.id LIMIT 2; 2 left2 4 left4 +# A parent filter must remain when only an anti join's non-output side accepts +# pushdown; otherwise filtering that side creates incorrect anti-join rows. +statement ok +SET datafusion.optimizer.max_passes = 0; + +statement ok +SET datafusion.optimizer.join_reordering = false; + +statement ok +SET datafusion.execution.parquet.pushdown_filters = true; + +query I +SELECT count(*) +FROM join_left l LEFT ANTI JOIN right_parquet r USING (id) +WHERE false; +---- +0 + +query I +SELECT count(*) +FROM right_parquet r RIGHT ANTI JOIN join_left l USING (id) +WHERE false; +---- +0 + +statement ok +RESET datafusion.optimizer.max_passes; + +statement ok +RESET datafusion.optimizer.join_reordering; + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + # Test 3: Test independent control # Disable TopK, keep Join enabled @@ -516,7 +549,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], projection=[id@2, data@3, info@1] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id, info], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Enable TopK, disable Join statement ok @@ -588,7 +621,7 @@ physical_plan 02)--CoalescePartitionsExec 03)----AggregateExec: mode=Partial, gby=[], aggr=[max(agg_parquet.score)] 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/agg_data.parquet]]}, projection=[score], file_type=parquet, predicate=category@0 = alpha AND DynamicFilter [ empty ], pruning_predicate=category_null_count@2 != row_count@3 AND category_min@0 <= alpha AND alpha <= category_max@1, required_guarantees=[category in (alpha)] +05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/agg_data.parquet]]}, projection=[score], file_type=parquet, predicate=category@0 = alpha AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=category_null_count@2 != row_count@3 AND category_min@0 <= alpha AND alpha <= category_max@1, required_guarantees=[category in (alpha)] # Test 4b: COUNT + MAX — DynamicFilter should NOT appear here in mixed aggregates @@ -736,7 +769,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], projection=[id@2, data@3, info@1] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_right.parquet]]}, projection=[id, info], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_filter_pushdown_config/join_left.parquet]]}, projection=[id, data], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Test 6: Regression test for issue #20213 - dynamic filter applied to wrong table # when subquery join has same column names on both sides. diff --git a/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt new file mode 100644 index 0000000000000..c6700ebf0b97c --- /dev/null +++ b/datafusion/sqllogictest/test_files/dynamic_row_group_pruning.slt @@ -0,0 +1,289 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# End-to-end SLT for **dynamic row-group pruning** driven by a TopK +# `SortExec`'s `DynamicFilterPhysicalExpr`. +# +# Builds a 5-row-group parquet file with disjoint per-RG ranges of `v`: +# RG 0: 0..3, RG 1: 3..6, RG 2: 6..9, RG 3: 9..12, RG 4: 12..15 +# `ORDER BY v DESC LIMIT 3` fills the TopK heap from the row group with +# the largest values; the tightened threshold then proves every other +# row group unreachable. At each row-group boundary the runtime +# `RowGroupPruner` evaluates the current threshold against the next RGs' +# statistics, drops the ones it proves unwinnable, and rebuilds the +# decoder via `into_builder().with_row_groups(...)` to skip them. Each +# drop bumps `row_groups_pruned_dynamic_filter`. + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +set datafusion.explain.analyze_level = summary; + +statement ok +CREATE TABLE source_data AS VALUES +-- RG 0 + (0), (1), (2), +-- RG 1 + (3), (4), (5), +-- RG 2 + (6), (7), (8), +-- RG 3 + (9), (10), (11), +-- RG 4 + (12), (13), (14); + +statement ok +COPY (SELECT column1 as v FROM source_data) +TO 'test_files/scratch/dynamic_row_group_pruning/data.parquet' +STORED AS PARQUET +OPTIONS ( + 'format.max_row_group_size' '3' +); + +statement ok +drop table source_data; + +statement ok +CREATE EXTERNAL TABLE t +STORED AS PARQUET +LOCATION 'test_files/scratch/dynamic_row_group_pruning/data.parquet'; + +# Sanity: query returns the right rows. +query I +SELECT v FROM t ORDER BY v DESC LIMIT 3; +---- +14 +13 +12 + +# Plain `EXPLAIN` must surface the plan-time eligibility marker +# `dynamic_rg_pruning=eligible` on the `DataSourceExec` line: the +# predicate is dynamic, so the runtime row-group pruner will be +# consulted at each decoder-run boundary. This is the only knob users +# have for spotting the optimization without running the query. +query TT +explain select v from t order by v desc limit 3; +---- +logical_plan +01)Sort: t.v DESC NULLS FIRST, fetch=3 +02)--TableScan: t projection=[v] +physical_plan +01)SortExec: TopK(fetch=3), expr=[v@0 DESC], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_row_group_pruning/data.parquet]]}, projection=[v], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[v@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible + +# `EXPLAIN ANALYZE` must surface the runtime metric +# `row_groups_pruned_dynamic_filter` with a non-zero value. Five +# disjoint row groups, `LIMIT 3` fits inside the highest RG, so the +# pruner skips the other four. Note the exact `=4`: the data is small +# enough that the TopK heap fills in a single batch, and execution is +# single-threaded, so the count is deterministic. Time- and size-keyed +# fields are masked with ``. +query TT +explain analyze select v from t order by v desc limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[v@0 DESC], preserve_partitioning=[false], filter=[v@0 IS NULL OR v@0 > 12], metrics=[output_rows=3, elapsed_compute=, output_bytes=] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/dynamic_row_group_pruning/data.parquet]]}, projection=[v], file_type=parquet, predicate=DynamicFilter [ v@0 IS NULL OR v@0 > 12 ], sort_order_for_reorder=[v@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=v_null_count@0 > 0 OR v_null_count@0 != row_count@2 AND v_max@1 > 12, required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=5 total → 5 matched, row_groups_pruned_bloom_filter=5 total → 5 matched, page_index_pages_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, row_groups_pruned_dynamic_filter=4, metadata_load_time=, scan_efficiency_ratio=] + +statement ok +drop table t; + +# Config reset — without these the SLT runner flags the file for +# leaking session state into subsequent files. +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + +statement ok +RESET datafusion.explain.analyze_level; + +# Regression test for #24352: TopK dynamic filter + `pushdown_filters` must not +# re-read an already-delivered row group. The filter column (`search_phrase`) +# differs from the sort column (`event_time`), and one row group has an empty +# post-predicate selection that row-group statistics cannot see — its only small +# `event_time` (50) sits on the row where `search_phrase = ''`. arrow-rs finishes +# that RG without handing back a reader; without syncing `rg_plan` to the decoder +# frontier via `peek_next_row_group`, `rg_plan` trailed the decoder by one, so a +# later runtime prune rebuilt the decoder from a stale plan and re-read an +# already-delivered RG — duplicating rows and dropping the true top-k tail. +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +set datafusion.execution.target_partitions = 1; + +# Both dynamic-filter switches are on by default; set them explicitly so this +# test keeps exercising the prune/rebuild path even if the defaults change. +statement ok +set datafusion.optimizer.enable_dynamic_filter_pushdown = true; + +statement ok +set datafusion.optimizer.enable_topk_dynamic_filter_pushdown = true; + +statement ok +CREATE TABLE q26_src AS +SELECT + CAST(CASE + WHEN i < 2048 THEN i * 1000 + WHEN i < 4096 THEN (CASE WHEN i = 2048 THEN 50 ELSE 20000 + i END) + WHEN i < 6144 THEN 100 + (i - 4096) + ELSE 5000 + (i - 6144) + END AS BIGINT) AS event_time, + CASE WHEN i = 2048 THEN '' ELSE 'p' || CAST(i AS VARCHAR) END AS search_phrase +FROM generate_series(0, 8191) AS t(i); + +statement ok +COPY (SELECT * FROM q26_src) +TO 'test_files/scratch/dynamic_row_group_pruning/q26.parquet' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' '2048'); + +statement ok +drop table q26_src; + +statement ok +CREATE EXTERNAL TABLE q26 (event_time BIGINT NOT NULL, search_phrase VARCHAR NOT NULL) +STORED AS PARQUET +LOCATION 'test_files/scratch/dynamic_row_group_pruning/q26.parquet'; + +# Each search_phrase is unique, so any repeated value would be the same source +# row emitted twice. The result must be the 10 smallest-`event_time` non-empty +# phrases with no duplicates (matches DuckDB and pushdown-off DataFusion). +query T +SELECT search_phrase FROM q26 WHERE search_phrase <> '' ORDER BY event_time LIMIT 10; +---- +p0 +p4096 +p4097 +p4098 +p4099 +p4100 +p4101 +p4102 +p4103 +p4104 + +statement ok +drop table q26; + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# restore it explicitly rather than RESET (which would revert to the system +# default = num_cpus and leak modified config out of this file). +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +RESET datafusion.optimizer.enable_dynamic_filter_pushdown; + +statement ok +RESET datafusion.optimizer.enable_topk_dynamic_filter_pushdown; + +# Regression test for a scan where two pruning mechanisms are live at once: +# page-index pruning leaves an intra-row-group `RowSelection`, and a TopK +# dynamic filter prunes row groups at runtime. The property under test is that +# a `WHERE a >= 50 ORDER BY b ASC LIMIT 5` query returns the correct top-5 by +# `b` while the dynamic predicate prunes a row group during the application of +# multiple predicates. Layout (RG size 100): +# RG 0: b=1000..1099, a=100..199 (a>=50 keeps all) +# RG 1: b=2000..2099, a=0..99 (a>=50 keeps rows 50..99 — page-index prunes +# the first 5 pages, leaving `skip 50, select 50`) +# RG 2: b=3000..3099, a=100..199 (keeps all) +# RG 3: b=0..99, a=100..199 (keeps all) +# The correct top-5 by `b` (0..4) lives entirely in RG 3. +# `data_page_row_count_limit`/`write_batch_size` force multiple pages per RG so +# page-index pruning can produce the intra-RG selection. +# Tracking issue for the behavior change (keeping both mechanisms): +# https://github.com/apache/arrow-rs/issues/10624 / +# https://github.com/apache/datafusion/issues/24358 +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +CREATE TABLE rgsel_src AS +SELECT + CAST(CASE WHEN i / 100 = 1 THEN i % 100 ELSE 100 + (i % 100) END AS BIGINT) AS a, + CAST(CASE + WHEN i < 100 THEN 1000 + i + WHEN i < 200 THEN 2000 + (i - 100) + WHEN i < 300 THEN 3000 + (i - 200) + ELSE (i - 300) + END AS BIGINT) AS b +FROM generate_series(0, 399) AS t(i); + +statement ok +COPY (SELECT * FROM rgsel_src) +TO 'test_files/scratch/dynamic_row_group_pruning/rgsel.parquet' +STORED AS PARQUET +OPTIONS ( + 'format.max_row_group_size' '100', + 'format.data_page_row_count_limit' '10', + 'format.write_batch_size' '10' +); + +statement ok +drop table rgsel_src; + +statement ok +CREATE EXTERNAL TABLE rgsel (a BIGINT NOT NULL, b BIGINT NOT NULL) +STORED AS PARQUET +LOCATION 'test_files/scratch/dynamic_row_group_pruning/rgsel.parquet'; + +# The correct top-5 by `b` among rows with `a >= 50` is b = 0..4 (they live in +# RG 3, all of whose rows satisfy `a >= 50`). +query I +SELECT b FROM rgsel WHERE a >= 50 ORDER BY b ASC LIMIT 5; +---- +0 +1 +2 +3 +4 + +# The same query without filter pushdown never engages the runtime pruner, so +# its answer is the ground truth the pushdown path above must match. +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +query I +SELECT b FROM rgsel WHERE a >= 50 ORDER BY b ASC LIMIT 5; +---- +0 +1 +2 +3 +4 + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +drop table rgsel; + +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# restore it explicitly rather than RESET (which would revert to the system +# default = num_cpus and leak modified config out of this file). +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; diff --git a/datafusion/sqllogictest/test_files/eliminate_outer_join.slt b/datafusion/sqllogictest/test_files/eliminate_outer_join.slt index dff7692a4451e..afd491b0b64c8 100644 --- a/datafusion/sqllogictest/test_files/eliminate_outer_join.slt +++ b/datafusion/sqllogictest/test_files/eliminate_outer_join.slt @@ -370,6 +370,66 @@ select * from t1 left join t2 on t1.a = t2.x where (t2.y > 150) is unknown; 3 30 c NULL NULL NULL NULL 40 d NULL NULL NULL +# LEFT JOIN + WHERE NOT ((t2.y > 150) IS TRUE) -> stays LEFT JOIN +query TT +explain select * from t1 left join t2 on t1.a = t2.x where not ((t2.y > 150) is true); +---- +logical_plan +01)Filter: NOT t2.y > Int32(150) IS TRUE +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b, c] +04)----TableScan: t2 projection=[x, y, z] + +# Both the matched-with-low-y row AND the LEFT-padded NULL rows must +# survive. +query IITIIT rowsort +select * from t1 left join t2 on t1.a = t2.x where not ((t2.y > 150) is true); +---- +1 10 a 1 100 p +3 30 c NULL NULL NULL +NULL 40 d NULL NULL NULL + +# LEFT JOIN + WHERE NOT ((t2.y > 150) IS FALSE) -> stays LEFT JOIN +# NOT( IS FALSE) is TRUE when is TRUE OR NULL, so it accepts the +# LEFT-padded NULL rows and is not null-rejecting. +query TT +explain select * from t1 left join t2 on t1.a = t2.x where not ((t2.y > 150) is false); +---- +logical_plan +01)Filter: NOT t2.y > Int32(150) IS FALSE +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b, c] +04)----TableScan: t2 projection=[x, y, z] + +# The matched-with-high-y row AND the LEFT-padded NULL rows must survive. +query IITIIT rowsort +select * from t1 left join t2 on t1.a = t2.x where not ((t2.y > 150) is false); +---- +2 20 b 2 200 q +3 30 c NULL NULL NULL +NULL 40 d NULL NULL NULL + +# LEFT JOIN + WHERE NOT ((t2.y > 150) IS NOT UNKNOWN) -> stays LEFT JOIN +# NOT( IS NOT UNKNOWN) is equivalent to IS UNKNOWN: TRUE only when +# is NULL, so it accepts the LEFT-padded NULL rows and is not +# null-rejecting. +query TT +explain select * from t1 left join t2 on t1.a = t2.x where not ((t2.y > 150) is not unknown); +---- +logical_plan +01)Filter: NOT t2.y > Int32(150) IS NOT UNKNOWN +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b, c] +04)----TableScan: t2 projection=[x, y, z] + +# Only the LEFT-padded NULL rows (where t2.y > 150 evaluates to NULL) +# should survive. +query IITIIT rowsort +select * from t1 left join t2 on t1.a = t2.x where not ((t2.y > 150) is not unknown); +---- +3 30 c NULL NULL NULL +NULL 40 d NULL NULL NULL + ### ### FULL JOIN → LEFT / RIGHT conversion tests ### @@ -478,6 +538,326 @@ select * from t1 left join t2 on t1.a = t2.x where (t2.y > 150) is true and t2.z ---- 2 20 b 2 200 q +### +### Projection between Filter and Join +### + +# A filter on a volatile, projected expression can still be used for outer join +# elimination. +query TT +explain +select s.a +from ( + select t1.a, random() + cast(t2.y as double) as ry + from t1 left join t2 on t1.a = t2.x +) s +where s.ry > 150.0; +---- +logical_plan +01)SubqueryAlias: s +02)--Projection: t1.a +03)----Filter: ry > Float64(150) +04)------Projection: t1.a, random() + CAST(t2.y AS Float64) AS ry +05)--------Inner Join: t1.a = t2.x +06)----------TableScan: t1 projection=[a] +07)----------TableScan: t2 projection=[x, y] + +query I rowsort +select s.a +from ( + select t1.a, random() + cast(t2.y as double) as ry + from t1 left join t2 on t1.a = t2.x +) s +where s.ry > 150.0; +---- +2 + +# This query has the shape of TPC-DS Q49: `OptimizeProjections` results in +# placing a `Projection` node between the `Filter` and `Join`, but we can look +# through that node to convert the outer join. +statement ok +create table d(k int, flag int); + +statement ok +insert into d values (1, 1), (2, 1), (3, 0); + +query TT +explain +select t1.a, sum(coalesce(t2.y, 0)) as ret_sum +from t1 left join t2 on t1.a = t2.x, d +where t2.y > 150 + and t1.a = d.k + and d.flag = 1 +group by t1.a; +---- +logical_plan +01)Projection: t1.a, sum(coalesce(t2.y,Int64(0))) AS ret_sum +02)--Aggregate: groupBy=[[t1.a]], aggr=[[sum(CASE WHEN __common_expr_1 IS NOT NULL THEN __common_expr_1 ELSE Int64(0) END) AS sum(coalesce(t2.y,Int64(0)))]] +03)----Projection: CAST(t2.y AS Int64) AS __common_expr_1, t1.a +04)------Inner Join: t1.a = d.k +05)--------Projection: t1.a, t2.y +06)----------Inner Join: t1.a = t2.x +07)------------TableScan: t1 projection=[a] +08)------------Filter: t2.y > Int32(150) +09)--------------TableScan: t2 projection=[x, y] +10)--------Projection: d.k +11)----------Filter: d.flag = Int32(1) +12)------------TableScan: d projection=[k, flag] + +query II rowsort +select t1.a, sum(coalesce(t2.y, 0)) as ret_sum +from t1 left join t2 on t1.a = t2.x, d +where t2.y > 150 + and t1.a = d.k + and d.flag = 1 +group by t1.a; +---- +2 200 + +# A CTE can introduce a query boundary between the outer filter and the +# LEFT JOIN. +query TT +explain +with s as ( + select t1.a, t2.y + from t1 left join t2 on t1.a = t2.x +) +select s.a from s where s.y > 150; +---- +logical_plan +01)SubqueryAlias: s +02)--Projection: t1.a +03)----Inner Join: t1.a = t2.x +04)------TableScan: t1 projection=[a] +05)------Projection: t2.x +06)--------Filter: t2.y > Int32(150) +07)----------TableScan: t2 projection=[x, y] + +query I rowsort +with s as ( + select t1.a, t2.y + from t1 left join t2 on t1.a = t2.x +) +select s.a from s where s.y > 150; +---- +2 + +# https://github.com/apache/datafusion/issues/13232 +# Strict scalar functions over the nullable side of an outer join reject the +# NULL-padding rows when used in a null-rejecting filter. +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where abs(t2.y) > 5; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: abs(t2.y) > Int32(5) +06)--------TableScan: t2 projection=[x, y] + +query I rowsort +select t1.a +from t1 left join t2 on t1.a = t2.x +where abs(t2.y) > 5; +---- +1 +2 + +query TT +explain +select t1.a +from t1 inner join t2 on t1.a = t2.x +where abs(t2.y) > 5; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: abs(t2.y) > Int32(5) +06)--------TableScan: t2 projection=[x, y] + +query I rowsort +select t1.a +from t1 inner join t2 on t1.a = t2.x +where abs(t2.y) > 5; +---- +1 +2 + +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where abs(t2.y) is not null; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: abs(t2.y) IS NOT NULL +06)--------TableScan: t2 projection=[x, y] + +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where coalesce(t2.y, 0) > 5; +---- +logical_plan +01)Projection: t1.a +02)--Filter: CASE WHEN __common_expr_3 IS NOT NULL THEN __common_expr_3 ELSE Int64(0) END > Int64(5) +03)----Projection: CAST(t2.y AS Int64) AS __common_expr_3, t1.a +04)------Left Join: t1.a = t2.x +05)--------TableScan: t1 projection=[a] +06)--------TableScan: t2 projection=[x, y] + +### +### Strict math function matrix +### + +# Unary function on the nullable side of a LEFT JOIN -> INNER JOIN. +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where ceil(t2.y) > 150; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: ceil(CAST(t2.y AS Float64)) > Float64(150) +06)--------TableScan: t2 projection=[x, y] + +query I rowsort +select t1.a +from t1 left join t2 on t1.a = t2.x +where ceil(t2.y) > 150; +---- +2 + +# Unary function on the nullable side of a RIGHT JOIN -> INNER JOIN. +query TT +explain +select t2.x +from t1 right join t2 on t1.a = t2.x +where floor(t1.b) > 15; +---- +logical_plan +01)Projection: t2.x +02)--Inner Join: t1.a = t2.x +03)----Projection: t1.a +04)------Filter: CAST(t1.b AS Float64) >= Float64(16) +05)--------TableScan: t1 projection=[a, b] +06)----TableScan: t2 projection=[x] + +query I rowsort +select t2.x +from t1 right join t2 on t1.a = t2.x +where floor(t1.b) > 15; +---- +2 + +# Binary function with the nullable column as its first argument. +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where atan2(t2.y, 1) > 1; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: atan2(CAST(t2.y AS Float64), Float64(1)) > Float64(1) +06)--------TableScan: t2 projection=[x, y] + +query I rowsort +select t1.a +from t1 left join t2 on t1.a = t2.x +where atan2(t2.y, 1) > 1; +---- +1 +2 + +# Binary function with the nullable column as its second argument. +query TT +explain +select t1.a +from t1 left join t2 on t1.a = t2.x +where power(2, t2.y) > 100; +---- +logical_plan +01)Projection: t1.a +02)--Inner Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Projection: t2.x +05)------Filter: power(Float64(2), CAST(t2.y AS Float64)) > Float64(100) +06)--------TableScan: t2 projection=[x, y] + +query I rowsort +select t1.a +from t1 left join t2 on t1.a = t2.x +where power(2, t2.y) > 100; +---- +1 +2 + +# A strict function on only the right side of a FULL JOIN -> RIGHT JOIN. +query TT +explain +select t1.a, t2.y +from t1 full join t2 on t1.a = t2.x +where round(t2.y, -2) >= 100; +---- +logical_plan +01)Projection: t1.a, t2.y +02)--Right Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----Filter: round(t2.y, Int32(-2)) >= Int32(100) +05)------TableScan: t2 projection=[x, y] + +query II rowsort +select t1.a, t2.y +from t1 full join t2 on t1.a = t2.x +where round(t2.y, -2) >= 100; +---- +1 100 +2 200 +NULL 300 + +# A strict function on only the left side of a FULL JOIN -> LEFT JOIN. +query TT +explain +select t1.a, t1.b +from t1 full join t2 on t1.a = t2.x +where trunc(t1.b, -1) >= 10; +---- +logical_plan +01)Projection: t1.a, t1.b +02)--Left Join: t1.a = t2.x +03)----Filter: trunc(CAST(t1.b AS Float64), Int64(-1)) >= Float64(10) +04)------TableScan: t1 projection=[a, b] +05)----TableScan: t2 projection=[x] + +query II rowsort +select t1.a, t1.b +from t1 full join t2 on t1.a = t2.x +where trunc(t1.b, -1) >= 10; +---- +1 10 +2 20 +3 30 +NULL 40 + ### ### Cleanup ### @@ -490,3 +870,6 @@ drop table t1; statement ok drop table t2; + +statement ok +drop table d; diff --git a/datafusion/sqllogictest/test_files/errors.slt b/datafusion/sqllogictest/test_files/errors.slt index 20c1db5cb1511..ab934279c32ec 100644 --- a/datafusion/sqllogictest/test_files/errors.slt +++ b/datafusion/sqllogictest/test_files/errors.slt @@ -180,13 +180,13 @@ SELECT DISTINCT - 84 FROM tab0 AS cor0 WHERE NOT + 96 / + col1 <= NULL GROUP BY statement ok create table a(timestamp int, birthday int, ts int, tokens int, amp int, staamp int); -query error DataFusion error: Schema error: No field named timetamp\. Did you mean 'a\.timestamp'\?\. +query error DataFusion error: Schema error: No field named timetamp\. Did you mean 'a\.timestamp'\?\nValid fields are a\.timestamp, a\.birthday, a\.ts, a\.tokens, a\.amp, a\.staamp\. select timetamp from a; -query error DataFusion error: Schema error: No field named dadsada\. Valid fields are a\.timestamp, a\.birthday, a\.ts, a\.tokens, a\.amp, a\.staamp\. +query error DataFusion error: Schema error: No field named dadsada\.\nValid fields are a\.timestamp, a\.birthday, a\.ts, a\.tokens, a\.amp, a\.staamp\. select dadsada from a; -query error DataFusion error: Schema error: No field named ammp\. Did you mean 'a\.amp'\?\. +query error DataFusion error: Schema error: No field named ammp\. Did you mean 'a\.amp'\?\nValid fields are a\.timestamp, a\.birthday, a\.ts, a\.tokens, a\.amp, a\.staamp\. select ammp from a; statement ok diff --git a/datafusion/sqllogictest/test_files/explain.slt b/datafusion/sqllogictest/test_files/explain.slt index 67d2c1e7b516e..b6837002086ad 100644 --- a/datafusion/sqllogictest/test_files/explain.slt +++ b/datafusion/sqllogictest/test_files/explain.slt @@ -237,11 +237,10 @@ physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE physical_plan after FilterPushdown SAME TEXT AS ABOVE -physical_plan after EnforceDistribution SAME TEXT AS ABOVE +physical_plan after WindowTopN SAME TEXT AS ABOVE +physical_plan after EnsureRequirements SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE -physical_plan after EnforceSorting SAME TEXT AS ABOVE physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE -physical_plan after WindowTopN SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after OutputRequirements DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true physical_plan after LimitAggregation SAME TEXT AS ABOVE @@ -318,11 +317,10 @@ physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE physical_plan after FilterPushdown SAME TEXT AS ABOVE -physical_plan after EnforceDistribution SAME TEXT AS ABOVE +physical_plan after WindowTopN SAME TEXT AS ABOVE +physical_plan after EnsureRequirements SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE -physical_plan after EnforceSorting SAME TEXT AS ABOVE physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE -physical_plan after WindowTopN SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after OutputRequirements 01)GlobalLimitExec: skip=0, fetch=10, statistics=[Rows=Exact(8), Bytes=Absent, [(Col[0]: ScanBytes=Exact(32)),(Col[1]: ScanBytes=Inexact(24)),(Col[2]: ScanBytes=Exact(32)),(Col[3]: ScanBytes=Exact(32)),(Col[4]: ScanBytes=Exact(32)),(Col[5]: ScanBytes=Exact(64)),(Col[6]: ScanBytes=Exact(32)),(Col[7]: ScanBytes=Exact(64)),(Col[8]: ScanBytes=Inexact(88)),(Col[9]: ScanBytes=Inexact(49)),(Col[10]: ScanBytes=Exact(64))]] @@ -365,11 +363,10 @@ physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE physical_plan after FilterPushdown SAME TEXT AS ABOVE -physical_plan after EnforceDistribution SAME TEXT AS ABOVE +physical_plan after WindowTopN SAME TEXT AS ABOVE +physical_plan after EnsureRequirements SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE -physical_plan after EnforceSorting SAME TEXT AS ABOVE physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE -physical_plan after WindowTopN SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after OutputRequirements 01)GlobalLimitExec: skip=0, fetch=10 @@ -536,8 +533,13 @@ query error DataFusion error: Error during planning: EXPLAIN VERBOSE with FORMAT explain verbose format tree select * from values (1); # valid explain format -query error DataFusion error: Invalid or Unsupported Configuration: Invalid explain format. Expected 'indent', 'tree', 'pgjson' or 'graphviz'. Got 'xxx' +query error set datafusion.explain.format = "xxx"; +---- +DataFusion error: Error setting config datafusion.explain.format +caused by +Invalid or Unsupported Configuration: Invalid explain format. Expected 'indent', 'tree', 'pgjson' or 'graphviz'. Got 'xxx' + # verbose uses indent mode even when a different mode (e.g tree) is set @@ -614,11 +616,10 @@ physical_plan after aggregate_statistics SAME TEXT AS ABOVE physical_plan after join_selection SAME TEXT AS ABOVE physical_plan after LimitedDistinctAggregation SAME TEXT AS ABOVE physical_plan after FilterPushdown SAME TEXT AS ABOVE -physical_plan after EnforceDistribution SAME TEXT AS ABOVE +physical_plan after WindowTopN SAME TEXT AS ABOVE +physical_plan after EnsureRequirements SAME TEXT AS ABOVE physical_plan after CombinePartialFinalAggregate SAME TEXT AS ABOVE -physical_plan after EnforceSorting SAME TEXT AS ABOVE physical_plan after OptimizeAggregateOrder SAME TEXT AS ABOVE -physical_plan after WindowTopN SAME TEXT AS ABOVE physical_plan after ProjectionPushdown SAME TEXT AS ABOVE physical_plan after OutputRequirements DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/example.csv]]}, projection=[a, b, c], file_type=csv, has_header=true physical_plan after LimitAggregation SAME TEXT AS ABOVE @@ -682,7 +683,7 @@ logical_plan 11)--subgraph cluster_3 12)--{ 13)----graph[label="Detailed LogicalPlan"] -14)----4[shape=box label="Values: (Int64(1))\nSchema: [column1:Int64;N]"] +14)----4[shape=box label="Values: (Int64(1))\nSchema: [column1:Int64]"] 15)--} 16)} 17)// End DataFusion GraphViz Plan @@ -691,3 +692,150 @@ logical_plan statement ok drop table foo; + +# ------------------------------------------------------------------ +# Postgres-style `EXPLAIN (option, ...)` tests (dialect-gated). +# +# These require a dialect whose `supports_explain_with_utility_options()` +# returns true. DataFusion's default Generic dialect also declares this +# (mirroring sqlparser-rs 0.61.0), so the parenthesized form works there +# too. We set PostgreSQL explicitly for clarity. +# ------------------------------------------------------------------ + +statement ok +set datafusion.sql_parser.dialect = 'PostgreSQL'; + +# `EXPLAIN (FORMAT tree)` matches the legacy `EXPLAIN FORMAT tree` form. +query TT +EXPLAIN (FORMAT tree) SELECT 1; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ Int64(1): 1 │ +05)└─────────────┬─────────────┘ +06)┌─────────────┴─────────────┐ +07)│ PlaceholderRowExec │ +08)└───────────────────────────┘ + +# Unknown options are rejected with a clear error. +statement error DataFusion error: Error during planning: unknown EXPLAIN option: FOO +EXPLAIN (FOO) SELECT 1; + +# Postgres-only options return a "not supported" message pointing at METRICS. +statement error DataFusion error: This feature is not implemented: EXPLAIN option BUFFERS is not supported by DataFusion +EXPLAIN (BUFFERS) SELECT 1; + +statement error DataFusion error: This feature is not implemented: EXPLAIN option WAL is not supported by DataFusion +EXPLAIN (WAL) SELECT 1; + +# LEVEL / METRICS / TIMING / SUMMARY all require ANALYZE. +statement error DataFusion error: Error during planning: EXPLAIN option LEVEL requires ANALYZE +EXPLAIN (LEVEL dev) SELECT 1; + +statement error DataFusion error: Error during planning: EXPLAIN option METRICS requires ANALYZE +EXPLAIN (METRICS 'rows') SELECT 1; + +# COSTS and ANALYZE are mutually exclusive (COSTS only applies to plan-only +# EXPLAIN). +statement error DataFusion error: Error during planning: EXPLAIN option COSTS cannot be combined with ANALYZE +EXPLAIN (ANALYZE, COSTS ON) SELECT 1; + +# TIMING and SUMMARY are sugar for METRICS/LEVEL and likewise need ANALYZE. +statement error DataFusion error: Error during planning: EXPLAIN option METRICS requires ANALYZE +EXPLAIN (TIMING ON) SELECT 1; + +statement error DataFusion error: Error during planning: EXPLAIN option LEVEL requires ANALYZE +EXPLAIN (SUMMARY ON) SELECT 1; + +# VERBOSE is incompatible with any FORMAT, and ANALYZE only supports the +# `indent` and `pgjson` formats — `tree` and `graphviz` are rejected (these +# mappings come from the planner, not the parser). +statement error DataFusion error: Error during planning: EXPLAIN ANALYZE with FORMAT tree is not supported +EXPLAIN (ANALYZE, FORMAT tree) SELECT 1; + +statement error DataFusion error: Error during planning: EXPLAIN VERBOSE with FORMAT is not supported +EXPLAIN (VERBOSE, FORMAT tree) SELECT 1; + +# FORMAT argument can be a bare identifier (already tested) or a quoted +# string and produces the same plan either way. +query TT +EXPLAIN (FORMAT 'tree') SELECT 1; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ Int64(1): 1 │ +05)└─────────────┬─────────────┘ +06)┌─────────────┴─────────────┐ +07)│ PlaceholderRowExec │ +08)└───────────────────────────┘ + +# Bool option arguments accept bare/ON|OFF/TRUE|FALSE/1|0/=value forms. +# `ANALYZE OFF` is the same as a plain `EXPLAIN`. +query TT +EXPLAIN (ANALYZE OFF, FORMAT tree) SELECT 1; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ Int64(1): 1 │ +05)└─────────────┬─────────────┘ +06)┌─────────────┴─────────────┐ +07)│ PlaceholderRowExec │ +08)└───────────────────────────┘ + +# `COSTS OFF` overrides `datafusion.explain.show_statistics` per-statement +# (ANALYZE+COSTS is rejected above). +query TT +EXPLAIN (COSTS OFF) SELECT 1; +---- +logical_plan +01)Projection: Int64(1) +02)--EmptyRelation: rows=1 +physical_plan +01)ProjectionExec: expr=[1 as Int64(1)] +02)--PlaceholderRowExec + +# Bool argument forms: ON / TRUE / 1 all enable the option. The parenthesized +# form does not support `= value` for booleans (sqlparser's utility option +# grammar). Quoted-string booleans are accepted by the option parser. +statement ok +EXPLAIN (COSTS ON) SELECT 1; + +statement ok +EXPLAIN (COSTS TRUE) SELECT 1; + +statement ok +EXPLAIN (COSTS 1) SELECT 1; + +statement ok +EXPLAIN (COSTS 'true') SELECT 1; + +# Unrecognized argument for a boolean option. +statement error DataFusion error: Error during planning: expected boolean for EXPLAIN option costs, got 'maybe' +EXPLAIN (COSTS maybe) SELECT 1; + +# Unrecognized argument for a string/ident option. +statement error DataFusion error: Invalid or Unsupported Configuration: Invalid explain format\. Expected 'indent', 'tree', 'pgjson' or 'graphviz'\. Got 'bogus' +EXPLAIN (FORMAT bogus) SELECT 1; + +# Legacy keyword form still works on PostgreSQL dialect. +query TT +EXPLAIN FORMAT tree SELECT 1; +---- +physical_plan +01)┌───────────────────────────┐ +02)│ ProjectionExec │ +03)│ -------------------- │ +04)│ Int64(1): 1 │ +05)└─────────────┬─────────────┘ +06)┌─────────────┴─────────────┐ +07)│ PlaceholderRowExec │ +08)└───────────────────────────┘ + +statement ok +reset datafusion.sql_parser.dialect; diff --git a/datafusion/sqllogictest/test_files/explain_analyze.slt b/datafusion/sqllogictest/test_files/explain_analyze.slt index 7460148bab8f4..d64efe80ccae5 100644 --- a/datafusion/sqllogictest/test_files/explain_analyze.slt +++ b/datafusion/sqllogictest/test_files/explain_analyze.slt @@ -231,7 +231,7 @@ explain analyze select * from cat_tracking where species > 'M' AND s >= 50 order ---- Plan with Metrics 01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[] statement ok reset datafusion.explain.analyze_categories; @@ -247,7 +247,7 @@ explain analyze select * from cat_tracking where species > 'M' AND s >= 50 order ---- Plan with Metrics 01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, scan_efficiency_ratio=22.13% (521/2.35 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, scan_efficiency_ratio=21.75% (485/2.23 K)] statement ok reset datafusion.explain.analyze_categories; @@ -262,7 +262,7 @@ explain analyze select * from cat_tracking where species > 'M' AND s >= 50 order ---- Plan with Metrics 01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3, output_bytes=] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, scan_efficiency_ratio=] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, scan_efficiency_ratio=] statement ok reset datafusion.explain.analyze_categories; @@ -277,7 +277,7 @@ explain analyze select * from cat_tracking where species > 'M' AND s >= 50 order ---- Plan with Metrics 01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3, output_bytes=] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, scan_efficiency_ratio=] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, scan_efficiency_ratio=] statement ok reset datafusion.explain.analyze_categories; @@ -292,7 +292,7 @@ explain analyze select * from cat_tracking where species > 'M' AND s >= 50 order ---- Plan with Metrics 01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[elapsed_compute=] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[elapsed_compute=, metadata_load_time=] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[elapsed_compute=, metadata_load_time=] statement ok reset datafusion.explain.analyze_categories; @@ -300,6 +300,412 @@ reset datafusion.explain.analyze_categories; statement ok reset datafusion.explain.analyze_level; +# ------------------------------------------------ +# Test memory metrics display. +# ------------------------------------------------ + +statement ok +set datafusion.explain.analyze_level = dev; + +statement ok +set datafusion.explain.analyze_categories = 'bytes'; + +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.optimizer.repartition_joins = false; + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +statement ok +set datafusion.optimizer.enable_piecewise_merge_join = false; + +statement ok +set datafusion.optimizer.hash_join_inlist_pushdown_max_size = 0; + +statement ok +set datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values = 0; + +query TT +EXPLAIN ANALYZE +WITH t1 (k) AS ( + VALUES (1), (2), (3) +), t2 (k) AS ( + VALUES (1), (2), (3) +) +SELECT * +FROM t1 +JOIN t2 ON t1.k = t2.k; +---- +Plan with Metrics +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(k@0, k@0)], metrics=[output_bytes=128.0 KB, build_mem_used=44.0 B] +02)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +03)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +04)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +05)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +statement ok +set datafusion.execution.hash_join_buffering_capacity = 1024; + +query TT +EXPLAIN ANALYZE +WITH t1 (k) AS ( + VALUES (1), (2), (3) +), t2 (k) AS ( + VALUES (1), (2), (3) +) +SELECT * +FROM t1 +JOIN t2 ON t1.k = t2.k; +---- +Plan with Metrics +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(k@0, k@0)], metrics=[output_bytes=128.0 KB, build_mem_used=44.0 B] +02)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +03)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +04)--BufferExec: capacity=1024, metrics=[max_mem_used=128.0 B] +05)----ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +06)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +statement ok +reset datafusion.execution.hash_join_buffering_capacity; + +query TT +EXPLAIN ANALYZE +WITH t1 (k) AS ( + VALUES (1), (2) +), t2 (k) AS ( + VALUES (10), (20) +) +SELECT * +FROM t1 +CROSS JOIN t2; +---- +Plan with Metrics +01)CrossJoinExec, metrics=[output_bytes=96.0 B, build_mem_used=128.0 B] +02)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +03)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +04)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +05)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +query TT +EXPLAIN ANALYZE +WITH t1 (k) AS ( + VALUES (1), (2) +), t2 (k) AS ( + VALUES (2), (3) +) +SELECT * +FROM t1 +JOIN t2 ON t1.k < t2.k; +---- +Plan with Metrics +01)NestedLoopJoinExec: join_type=Inner, filter=k@0 < k@1, metrics=[output_bytes=128.0 KB, spilled_bytes=0.0 B, build_mem_used=128.0 B] +02)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +03)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +04)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +05)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +statement ok +set datafusion.optimizer.enable_piecewise_merge_join = true; + +query TT +EXPLAIN ANALYZE +WITH t1 (k) AS ( + VALUES (3), (4) +), t2 (k) AS ( + VALUES (1), (2) +) +SELECT * +FROM t1 +JOIN t2 ON t1.k > t2.k; +---- +Plan with Metrics +01)PiecewiseMergeJoin: operator=Gt, join_type=Inner, on=(k > k), metrics=[output_bytes=0.0 B, build_mem_used=144.0 B] +02)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=16.0 B] +03)----SortExec: expr=[column1@0 ASC], preserve_partitioning=[false], metrics=[output_bytes=16.0 B, spilled_bytes=0.0 B] +04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +05)--ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +06)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +statement ok +set datafusion.optimizer.enable_piecewise_merge_join = false; + +statement ok +set datafusion.execution.target_partitions = 2; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +statement ok +CREATE TABLE ea_smj_t1(a text, b int) AS VALUES ('Alice', 50), ('Alice', 100), ('Bob', 1); + +statement ok +CREATE TABLE ea_smj_t2(a text, b int) AS VALUES ('Alice', 2), ('Alice', 1); + +query TT +EXPLAIN ANALYZE +SELECT ea_smj_t1.a, ea_smj_t1.b, ea_smj_t2.a, ea_smj_t2.b +FROM ea_smj_t1 +JOIN ea_smj_t2 ON ea_smj_t1.a = ea_smj_t2.a + AND ea_smj_t2.b * 50 <= ea_smj_t1.b; +---- +Plan with Metrics +01)SortMergeJoinExec: join_type=Inner, on=[(a@0, a@0)], filter=CAST(b@1 AS Int64) * 50 <= CAST(b@0 AS Int64), metrics=[output_bytes=320.0 KB, spilled_bytes=0.0 B, peak_mem_used=432.0 B] +02)--SortExec: expr=[a@0 ASC], preserve_partitioning=[false], metrics=[output_bytes=60.0 B, spilled_bytes=0.0 B] +03)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +04)--SortExec: expr=[a@0 ASC], preserve_partitioning=[false], metrics=[output_bytes=40.0 B, spilled_bytes=0.0 B] +05)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +query TT +EXPLAIN ANALYZE +SELECT ea_smj_t1.a, ea_smj_t1.b +FROM ea_smj_t1 +LEFT SEMI JOIN ea_smj_t2 ON ea_smj_t1.a = ea_smj_t2.a; +---- +Plan with Metrics +01)SortMergeJoinExec: join_type=LeftSemi, on=[(a@0, a@0)], metrics=[output_bytes=160.0 KB, spilled_bytes=0.0 B, peak_mem_used=0.0 B] +02)--SortExec: expr=[a@0 ASC], preserve_partitioning=[false], metrics=[output_bytes=60.0 B, spilled_bytes=0.0 B] +03)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] +04)--SortExec: expr=[a@0 ASC], preserve_partitioning=[false], metrics=[output_bytes=32.0 B, spilled_bytes=0.0 B] +05)----DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +statement ok +DROP TABLE ea_smj_t1; + +statement ok +DROP TABLE ea_smj_t2; + +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.optimizer.repartition_joins = false; + +statement ok +set datafusion.execution.enable_migration_aggregate = false; + +query TT +EXPLAIN ANALYZE +WITH t (k) AS ( + VALUES (1), (2), (1), (3) +) +SELECT k, count(*) +FROM t +GROUP BY k; +---- +Plan with Metrics +01)ProjectionExec: expr=[k@0 as k, count(Int64(1))@1 as count(*)], metrics=[output_bytes=1056.0 B] +02)--AggregateExec: mode=Single, gby=[k@0 as k], aggr=[count(Int64(1))], metrics=[output_bytes=1056.0 B, spilled_bytes=0.0 B, peak_mem_used=9.2 KB] +03)----ProjectionExec: expr=[column1@0 as k], metrics=[output_bytes=32.0 B] +04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] + +statement ok +reset datafusion.execution.enable_migration_aggregate; + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +statement ok +reset datafusion.optimizer.hash_join_inlist_pushdown_max_size; + +statement ok +reset datafusion.optimizer.hash_join_inlist_pushdown_max_distinct_values; + +statement ok +reset datafusion.optimizer.repartition_joins; + +statement ok +reset datafusion.optimizer.enable_piecewise_merge_join; + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.explain.analyze_categories; + +statement ok +reset datafusion.explain.analyze_level; + +# ------------------------------------------------------------------ +# Same category/level filtering, but via the Postgres-style +# `EXPLAIN (ANALYZE, METRICS ..., LEVEL ...)` statement option list. +# +# Each block below mirrors one of the `set datafusion.explain.*` +# tests above so the equivalence between session config and +# per-statement overrides is exercised side-by-side. +# ------------------------------------------------------------------ + +statement ok +set datafusion.sql_parser.dialect = 'PostgreSQL'; + +# ---- (METRICS 'none', LEVEL summary) — plan only, no metrics ---- + +query TT +EXPLAIN (ANALYZE, METRICS 'none', LEVEL summary) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[] + +# ---- (METRICS 'rows', LEVEL summary) — row-count metrics only ---- + +query TT +EXPLAIN (ANALYZE, METRICS 'rows', LEVEL summary) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, scan_efficiency_ratio=21.75% (485/2.23 K)] + +# ---- Quoted-string METRICS with multiple categories ---- + +query TT +EXPLAIN (ANALYZE, METRICS 'rows,bytes', LEVEL summary) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3, output_bytes=] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, scan_efficiency_ratio=] + +# ---- (METRICS 'timing', LEVEL summary) — timing metrics only ---- + +query TT +EXPLAIN (ANALYZE, METRICS 'timing', LEVEL summary) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[elapsed_compute=] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[elapsed_compute=, metadata_load_time=] + +# ---- TIMING sugar: `METRICS 'rows,bytes', TIMING off` ↔ rows+bytes only ---- +# Equivalent to METRICS 'rows,bytes' since the sugar removes the timing +# category from the explicit METRICS selection. + +query TT +EXPLAIN (ANALYZE, METRICS 'rows,bytes', TIMING off, LEVEL summary) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3, output_bytes=] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, scan_efficiency_ratio=] + +# ---- TIMING sugar: `METRICS 'rows', TIMING on` ↔ rows + timing ---- + +query TT +EXPLAIN (ANALYZE, METRICS 'rows', TIMING on, LEVEL summary) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3, elapsed_compute=] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, metadata_load_time=, scan_efficiency_ratio=21.75% (485/2.23 K)] + +# ---- SUMMARY sugar: `SUMMARY on` ↔ `LEVEL summary` ---- +# Equivalent to METRICS 'rows', LEVEL summary above. + +query TT +EXPLAIN (ANALYZE, METRICS 'rows', SUMMARY on) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, scan_efficiency_ratio=21.75% (485/2.23 K)] + +# ---- Statement option overrides session config ---- +# Session says 'timing' but statement-level `METRICS 'rows'` wins. + +statement ok +set datafusion.explain.analyze_categories = 'timing'; + +query TT +EXPLAIN (ANALYZE, METRICS 'rows', LEVEL summary) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, scan_efficiency_ratio=21.75% (485/2.23 K)] + +# ---- pgjson format: structural golden with no metrics ---- + +query TT +EXPLAIN (ANALYZE, FORMAT PGJSON, METRICS 'none') SELECT * FROM (VALUES (1), (2), (3)) t(x); +---- +Plan with Metrics +01)[ +02)--{ +03)----"Plan": { +04)------"Node Type": "ProjectionExec", +05)------"Details": "ProjectionExec: expr=[column1@0 as x]", +06)------"Plans": [ +07)--------{ +08)----------"Node Type": "DataSourceExec", +09)----------"Details": "DataSourceExec: partitions=1, partition_sizes=[1]", +10)----------"Plans": [] +11)--------} +12)------] +13)----} +14)--} +15)] + +statement ok +reset datafusion.explain.analyze_categories; + +# ---- pgjson with METRICS 'rows': row-count surfaces as Actual Rows ---- + +query TT +EXPLAIN (ANALYZE, FORMAT PGJSON, METRICS 'rows') SELECT * FROM (VALUES (1), (2), (3)) t(x); +---- +Plan with Metrics +01)[ +02)--{ +03)----"Plan": { +04)------"Node Type": "ProjectionExec", +05)------"Details": "ProjectionExec: expr=[column1@0 as x]", +06)------"Actual Rows": 3, +07)------"Extras": { +08)--------"output_batches": 1 +09)------}, +10)------"Plans": [ +11)--------{ +12)----------"Node Type": "DataSourceExec", +13)----------"Details": "DataSourceExec: partitions=1, partition_sizes=[1]", +14)----------"Plans": [] +15)--------} +16)------] +17)----} +18)--} +19)] + +# ---- Argument syntax variants for METRICS ---- +# Bare identifier, `= value`, and quoted string forms should all parse +# to the same selection. + +query TT +EXPLAIN (ANALYZE, METRICS rows, LEVEL summary) select * from cat_tracking where species > 'M' AND s >= 50 order by species limit 3; +---- +Plan with Metrics +01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/explain_analyze/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, scan_efficiency_ratio=21.75% (485/2.23 K)] + +statement ok +reset datafusion.sql_parser.dialect; + +# ---- Reject formats that AnalyzeExec cannot render with live metrics ---- + +query error DataFusion error: Error during planning: EXPLAIN ANALYZE with FORMAT tree is not supported +explain analyze format tree select 1; + +query error DataFusion error: Error during planning: EXPLAIN ANALYZE with FORMAT graphviz is not supported +explain analyze format graphviz select 1; + +# ---- pgjson does not render statistics yet, so reject show_statistics ---- + +statement ok +set datafusion.explain.show_statistics = true; + +query error DataFusion error: Error during planning: EXPLAIN ANALYZE with FORMAT pgjson does not support show_statistics +explain analyze format pgjson select 1; + +statement ok +reset datafusion.explain.show_statistics; + # --- Teardown --- statement ok diff --git a/datafusion/sqllogictest/test_files/explain_tree.slt b/datafusion/sqllogictest/test_files/explain_tree.slt index 5bb4817be9644..4e0397bb41e2e 100644 --- a/datafusion/sqllogictest/test_files/explain_tree.slt +++ b/datafusion/sqllogictest/test_files/explain_tree.slt @@ -1100,6 +1100,28 @@ physical_plan 24)-----------------------------│ format: csv │ 25)-----------------------------└───────────────────────────┘ +# Query with null-aware anti join (NOT IN subquery). +query TT +explain select int_col from table1 where int_col not in (select int_col from table2); +---- +physical_plan +01)┌───────────────────────────┐ +02)│ HashJoinExec │ +03)│ -------------------- │ +04)│ join_type: LeftAnti │ +05)│ │ +06)│ null_aware ├──────────────┐ +07)│ │ │ +08)│ on: │ │ +09)│ (int_col = int_col) │ │ +10)└─────────────┬─────────────┘ │ +11)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ +12)│ DataSourceExec ││ DataSourceExec │ +13)│ -------------------- ││ -------------------- │ +14)│ files: 1 ││ files: 1 │ +15)│ format: csv ││ format: parquet │ +16)└───────────────────────────┘└───────────────────────────┘ + # Query with nested loop join. query TT explain select int_col from table1 where exists (select count(*) from table2); @@ -1577,7 +1599,7 @@ physical_plan 04)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ 05)│ ProjectionExec ││ CoalescePartitionsExec │ 06)│ -------------------- ││ │ -07)│ id: 1 ││ │ +07)│ id: CAST(1 AS Int64) ││ │ 08)└─────────────┬─────────────┘└─────────────┬─────────────┘ 09)┌─────────────┴─────────────┐┌─────────────┴─────────────┐ 10)│ PlaceholderRowExec ││ ProjectionExec │ diff --git a/datafusion/sqllogictest/test_files/expr.slt b/datafusion/sqllogictest/test_files/expr.slt index 51b7591b41199..32113890aadc0 100644 --- a/datafusion/sqllogictest/test_files/expr.slt +++ b/datafusion/sqllogictest/test_files/expr.slt @@ -67,7 +67,7 @@ statement error Parser error: Invalid timezone "Foo": failed to parse timezone SELECT arrow_cast('2021-01-02T03:04:00', 'Timestamp(Nanosecond, Some("Foo"))') # test_array_index -query III??IIIIII +query III??IIIIIIII SELECT ([5,4,3,2,1])[1], ([5,4,3,2,1])[2], @@ -80,11 +80,11 @@ SELECT -- out of bounds ([5,4,3,2,1])[0], ([5,4,3,2,1])[6], - -- ([5,4,3,2,1])[-1], -- TODO: wrong answer - -- ([5,4,3,2,1])[null], -- TODO: not supported + ([5,4,3,2,1])[-1], + ([5,4,3,2,1])[null], ([5,4,3,2,1])[100] ---- -5 4 1 [1, 2] [3, 4] 1 3 4 NULL NULL NULL +5 4 1 [1, 2] [3, 4] 1 3 4 NULL NULL 1 NULL NULL # test_array_literals query ????? @@ -330,7 +330,7 @@ SELECT ascii('222') 50 query I -SELECT ascii('0xa') +SELECT ascii('0xa') ---- 48 @@ -561,7 +561,7 @@ NULL query T SELECT ltrim(' zzzytest ') ---- -zzzytest +zzzytest query T SELECT ltrim('zzzytest', 'xyz') @@ -850,6 +850,15 @@ SELECT to_hex(0) ---- 0 +query T +SELECT to_hex(arrow_cast(a, 'Dictionary(Int32, Int64)')) +FROM (VALUES (0), (10), (255), (NULL)) AS t(a) +---- +0 +a +ff +NULL + # negative values (two's complement encoding) query T SELECT to_hex(-1) @@ -985,17 +994,16 @@ SELECT upper(NULL) ---- NULL -# TODO issue: https://github.com/apache/datafusion/issues/6596 -# query ?? -#SELECT -# CAST([1,2,3,4] AS INT[]) as a, -# CAST([1,2,3,4] AS NUMERIC(10,4)[]) as b -#---- -#[1, 2, 3, 4] [1.0000, 2.0000, 3.0000, 4.0000] +query ?? +SELECT + CAST([1,2,3,4] AS INT[]) as a, + CAST([1,2,3,4] AS NUMERIC(10,4)[]) as b +---- +[1, 2, 3, 4] [1.0000, 2.0000, 3.0000, 4.0000] # test_random_expression query BB -SELECT +SELECT random() BETWEEN 0.0 AND 1.0, random() = random() ---- @@ -1551,6 +1559,30 @@ SELECT md5(NULL); ---- NULL +# md5 string and binary array inputs +query BBBBBB +SELECT + md5(column1) = md5('tom'), + md5(arrow_cast(column1, 'LargeUtf8')) = md5('tom'), + md5(arrow_cast(column1, 'Utf8View')) = md5('tom'), + md5(arrow_cast(column1, 'Binary')) = md5('tom'), + md5(arrow_cast(column1, 'LargeBinary')) = md5('tom'), + md5(arrow_cast(column1, 'BinaryView')) = md5('tom') +FROM (VALUES ('tom'), (NULL)) AS t(column1); +---- +true true true true true true +NULL NULL NULL NULL NULL NULL + +# invalid argument count and type +query error DataFusion error: +SELECT md5(); + +query error DataFusion error: +SELECT md5('tom', 'extra'); + +query error DataFusion error: +SELECT md5(1); + query ? SELECT digest('','md5'); ---- @@ -1586,6 +1618,31 @@ SELECT sha224(NULL); ---- NULL +# sha224 string and binary array inputs +query BBBBBB +SELECT + sha224(column1) = sha224('tom'), + sha224(arrow_cast(column1, 'LargeUtf8')) = sha224('tom'), + sha224(arrow_cast(column1, 'Utf8View')) = sha224('tom'), + sha224(arrow_cast(column1, 'Binary')) = sha224('tom'), + sha224(arrow_cast(column1, 'LargeBinary')) = sha224('tom'), + sha224(arrow_cast(column1, 'BinaryView')) = sha224('tom') +FROM (VALUES ('tom'), (NULL), ('mot')) AS t(column1); +---- +true true true true true true +NULL NULL NULL NULL NULL NULL +false false false false false false + +# invalid argument count and type +query error DataFusion error: +SELECT sha224(); + +query error DataFusion error: +SELECT sha224('tom', 'extra'); + +query error DataFusion error: +SELECT sha224(1); + query ? SELECT digest(NULL,'sha224'); ---- @@ -1646,6 +1703,31 @@ SELECT sha384(NULL); ---- NULL +# sha384 string and binary array inputs +query BBBBBB +SELECT + sha384(column1) = sha384('tom'), + sha384(arrow_cast(column1, 'LargeUtf8')) = sha384('tom'), + sha384(arrow_cast(column1, 'Utf8View')) = sha384('tom'), + sha384(arrow_cast(column1, 'Binary')) = sha384('tom'), + sha384(arrow_cast(column1, 'LargeBinary')) = sha384('tom'), + sha384(arrow_cast(column1, 'BinaryView')) = sha384('tom') +FROM (VALUES ('tom'), (NULL), ('mot')) AS t(column1); +---- +true true true true true true +NULL NULL NULL NULL NULL NULL +false false false false false false + +# invalid argument count and type +query error DataFusion error: +SELECT sha384(); + +query error DataFusion error: +SELECT sha384('tom', 'extra'); + +query error DataFusion error: +SELECT sha384(1); + query ? SELECT digest(NULL,'sha384'); ---- @@ -1676,6 +1758,31 @@ SELECT sha512(NULL); ---- NULL +# sha512 string and binary array inputs +query BBBBBB +SELECT + sha512(column1) = sha512('tom'), + sha512(arrow_cast(column1, 'LargeUtf8')) = sha512('tom'), + sha512(arrow_cast(column1, 'Utf8View')) = sha512('tom'), + sha512(arrow_cast(column1, 'Binary')) = sha512('tom'), + sha512(arrow_cast(column1, 'LargeBinary')) = sha512('tom'), + sha512(arrow_cast(column1, 'BinaryView')) = sha512('tom') +FROM (VALUES ('tom'), (NULL), ('mot')) AS t(column1); +---- +true true true true true true +NULL NULL NULL NULL NULL NULL +false false false false false false + +# invalid argument count and type +query error DataFusion error: +SELECT sha512(); + +query error DataFusion error: +SELECT sha512('tom', 'extra'); + +query error DataFusion error: +SELECT sha512(1); + query ? SELECT digest(NULL,'sha512'); ---- @@ -1716,6 +1823,69 @@ SELECT digest('','blake3'); ---- af1349b9f5f9a1a6a0404dea36dcc9499bcb25c9adc112b7cc9a93cae41f3262 +# digest every supported algorithm over an array +query BBBBBBBB +SELECT + digest(column1, 'md5') = digest('tom', 'md5'), + digest(column1, 'sha224') = digest('tom', 'sha224'), + digest(column1, 'sha256') = digest('tom', 'sha256'), + digest(column1, 'sha384') = digest('tom', 'sha384'), + digest(column1, 'sha512') = digest('tom', 'sha512'), + digest(column1, 'blake2s') = digest('tom', 'blake2s'), + digest(column1, 'blake2b') = digest('tom', 'blake2b'), + digest(column1, 'blake3') = digest('tom', 'blake3') +FROM (VALUES ('tom'), (NULL), ('mot')) AS t(column1); +---- +true true true true true true true true +NULL NULL NULL NULL NULL NULL NULL NULL +false false false false false false false false + +# binary-view, large-utf8, and utf8-view array inputs +query BBBBBBBB +SELECT + digest(arrow_cast(column1, 'BinaryView'), 'md5') = digest('tom', 'md5'), + digest(arrow_cast(column1, 'BinaryView'), 'sha224') = digest('tom', 'sha224'), + digest(arrow_cast(column1, 'BinaryView'), 'sha256') = digest('tom', 'sha256'), + digest(arrow_cast(column1, 'BinaryView'), 'sha384') = digest('tom', 'sha384'), + digest(arrow_cast(column1, 'BinaryView'), 'sha512') = digest('tom', 'sha512'), + digest(arrow_cast(column1, 'BinaryView'), 'blake2s') = digest('tom', 'blake2s'), + digest(arrow_cast(column1, 'BinaryView'), 'blake2b') = digest('tom', 'blake2b'), + digest(arrow_cast(column1, 'BinaryView'), 'blake3') = digest('tom', 'blake3') +FROM (VALUES ('tom'), (NULL), ('mot')) AS t(column1); +---- +true true true true true true true true +NULL NULL NULL NULL NULL NULL NULL NULL +false false false false false false false false + +query BB +SELECT + digest(arrow_cast(column1, 'LargeUtf8'), 'md5') = digest('tom', 'md5'), + digest(arrow_cast(column1, 'Utf8View'), 'md5') = digest('tom', 'md5') +FROM (VALUES ('tom'), (NULL), ('mot')) AS t(column1); +---- +true true +NULL NULL +false false + +# invalid algorithm, dynamic algorithm, argument count, and argument type +query error There is no built-in digest algorithm named 'unknown' +SELECT digest('tom', 'unknown'); + +query error Digest using dynamically decided method is not yet supported +SELECT digest(column1, column2) FROM (VALUES ('tom', 'md5')) AS t(column1, column2); + +query error DataFusion error: +SELECT digest(); + +query error DataFusion error: +SELECT digest('tom'); + +query error DataFusion error: +SELECT digest('tom', 'md5', 'extra'); + +query error DataFusion error: +SELECT digest(1, 'md5'); + # vverify utf8view query ? SELECT sha224(arrow_cast('tom', 'Utf8View')); @@ -1978,15 +2148,15 @@ query B select column1 <=> column2 from (VALUES (1, 1), (2, 3), (NULL, NULL)) as t; ---- true -false +false true # Sanity test - comparing <=> with equivalent expression query B -SELECT - (column1 <=> column2) = +SELECT + (column1 <=> column2) = (IFNULL(column1, false) = IFNULL(column2, false)) AS comparison_result -FROM (VALUES +FROM (VALUES (1, 1), -- equal values (1, 2), -- different values (NULL, NULL), -- both NULL @@ -2280,20 +2450,20 @@ host3 3.3 # can have an aggregate function with an inner CASE WHEN query TR -select - t2.server_host as host, +select + t2.server_host as host, sum(( - case when t2.server_host is not null + case when t2.server_host is not null then t2.server_load2 end - )) + )) from ( - select + select struct(time,load1,load2,host)['c2'] as server_load2, struct(time,load1,load2,host)['c3'] as server_host from t1 - ) t2 - where server_host IS NOT NULL + ) t2 + where server_host IS NOT NULL group by server_host order by host; ---- host1 101 @@ -2302,19 +2472,19 @@ host3 303 # TODO: Issue tracked in https://github.com/apache/datafusion/issues/10364 query TR -select - t2.server['c3'] as host, +select + t2.server['c3'] as host, sum(( - case when t2.server['c3'] is not null + case when t2.server['c3'] is not null then t2.server['c2'] end - )) + )) from ( - select + select struct(time,load1,load2,host) as server from t1 - ) t2 - where t2.server['c3'] IS NOT NULL + ) t2 + where t2.server['c3'] IS NOT NULL group by t2.server['c3'] order by host; ---- host1 101 @@ -2323,22 +2493,22 @@ host3 303 # can have 2 projections with aggr(short_circuited), with different short-circuited expr query TRR -select - t2.server_host as host, +select + t2.server_host as host, sum(coalesce(server_load1)), sum(( - case when t2.server_host is not null + case when t2.server_host is not null then t2.server_load2 end - )) + )) from ( - select + select struct(time,load1,load2,host)['c1'] as server_load1, struct(time,load1,load2,host)['c2'] as server_load2, struct(time,load1,load2,host)['c3'] as server_host from t1 - ) t2 - where server_host IS NOT NULL + ) t2 + where server_host IS NOT NULL group by server_host order by host; ---- host1 1.1 101 @@ -2347,43 +2517,43 @@ host3 3.3 303 # TODO: Issue tracked in https://github.com/apache/datafusion/issues/10364 query error -select - t2.server['c3'] as host, +select + t2.server['c3'] as host, sum(coalesce(server['c1'])), sum(( - case when t2.server['c3'] is not null + case when t2.server['c3'] is not null then t2.server['c2'] end - )) + )) from ( - select + select struct(time,load1,load2,host) as server, from t1 - ) t2 - where server_host IS NOT NULL + ) t2 + where server_host IS NOT NULL group by server_host order by host; query TRR -select - t2.server_host as host, +select + t2.server_host as host, sum(( - case when t2.server_host is not null - then server_load1 + case when t2.server_host is not null + then server_load1 end - )), + )), sum(( - case when server_host is not null - then server_load2 + case when server_host is not null + then server_load2 end - )) + )) from ( - select + select struct(time,load1,load2,host)['c1'] as server_load1, struct(time,load1,load2,host)['c2'] as server_load2, struct(time,load1,load2,host)['c3'] as server_host from t1 - ) t2 - where server_host IS NOT NULL + ) t2 + where server_host IS NOT NULL group by server_host order by host; ---- host1 1.1 101 @@ -2392,24 +2562,24 @@ host3 3.3 303 # TODO: Issue tracked in https://github.com/apache/datafusion/issues/10364 query TRR -select - t2.server['c3'] as host, +select + t2.server['c3'] as host, sum(( - case when t2.server['c3'] is not null + case when t2.server['c3'] is not null then t2.server['c1'] end - )), + )), sum(( - case when t2.server['c3'] is not null + case when t2.server['c3'] is not null then t2.server['c2'] end - )) + )) from ( - select - struct(time,load1,load2,host) as server + select + struct(time,load1,load2,host) as server from t1 - ) t2 - where t2.server['c3'] IS NOT NULL + ) t2 + where t2.server['c3'] IS NOT NULL group by t2.server['c3'] order by host; ---- host1 1.1 101 @@ -2495,3 +2665,18 @@ false statement ok drop table t; + + +# Test numeric literals with underscore separators +# (https://github.com/apache/datafusion/issues/23877) + +statement ok +set datafusion.sql_parser.dialect = 'postgres' + +query IIIRI +select 1_000, 1_2_3_4, -1_2_3_4, 1_2.3_4, 0_0 +---- +1000 1234 -1234 12.34 0 + +statement ok +reset datafusion.sql_parser.dialect diff --git a/datafusion/sqllogictest/test_files/file_row_index.slt b/datafusion/sqllogictest/test_files/file_row_index.slt new file mode 100644 index 0000000000000..38822bebfdfd3 --- /dev/null +++ b/datafusion/sqllogictest/test_files/file_row_index.slt @@ -0,0 +1,171 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +statement ok +COPY (VALUES (10), (20), (30), (40), (50)) +TO 'test_files/scratch/file_row_index/parquet_table/data.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE parquet_table(column1 int) +STORED AS PARQUET +LOCATION 'test_files/scratch/file_row_index/parquet_table/'; + +query TT +EXPLAIN SELECT file_row_index(), column1 FROM parquet_table +---- +logical_plan +01)Projection: file_row_index(), parquet_table.column1 +02)--TableScan: parquet_table projection=[column1] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/file_row_index/parquet_table/data.parquet]]}, projection=[CAST(__datafusion_file_row_index@1 AS Int64) as file_row_index(), column1], file_type=parquet + + +query II +SELECT file_row_index(), column1 FROM parquet_table ORDER BY column1 +---- +0 10 +1 20 +2 30 +3 40 +4 50 + +query III +SELECT file_row_index(), file_row_index() + 1, column1 +FROM parquet_table +ORDER BY column1 +---- +0 1 10 +1 2 20 +2 3 30 +3 4 40 +4 5 50 + + +query II +SELECT file_row_index(), column1 +FROM parquet_table +WHERE file_row_index() > 2 +ORDER BY column1 +---- +3 40 +4 50 + +# Filter on file_row_index without having it in projection + +query TT +EXPLAIN SELECT column1 FROM parquet_table WHERE file_row_index() > 2 ORDER BY column1 +---- +logical_plan +01)Sort: parquet_table.column1 ASC NULLS LAST +02)--Projection: parquet_table.column1 +03)----Filter: __datafusion_extracted_1 > Int64(2) +04)------Projection: file_row_index() AS __datafusion_extracted_1, parquet_table.column1 +05)--------TableScan: parquet_table projection=[column1] +physical_plan +01)SortExec: expr=[column1@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--FilterExec: __datafusion_extracted_1@0 > 2, projection=[column1@1] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/file_row_index/parquet_table/data.parquet]]}, projection=[CAST(__datafusion_file_row_index@1 AS Int64) as __datafusion_extracted_1, column1], file_type=parquet + +query I +SELECT column1 FROM parquet_table WHERE file_row_index() > 2 ORDER BY column1 +---- +40 +50 + +# Filter on file_row_index without projecting it, while enabling filter pushdown + +statement ok +SET datafusion.execution.parquet.pushdown_filters = true; + +query TT +EXPLAIN SELECT column1 FROM parquet_table WHERE file_row_index() > 2 ORDER BY column1 +---- +logical_plan +01)Sort: parquet_table.column1 ASC NULLS LAST +02)--Projection: parquet_table.column1 +03)----Filter: __datafusion_extracted_1 > Int64(2) +04)------Projection: file_row_index() AS __datafusion_extracted_1, parquet_table.column1 +05)--------TableScan: parquet_table projection=[column1] +physical_plan +01)SortExec: expr=[column1@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--FilterExec: __datafusion_extracted_1@0 > 2, projection=[column1@1] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/file_row_index/parquet_table/data.parquet]]}, projection=[CAST(__datafusion_file_row_index@1 AS Int64) as __datafusion_extracted_1, column1], file_type=parquet + +query I +SELECT column1 FROM parquet_table WHERE file_row_index() > 2 ORDER BY column1 +---- +40 +50 + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + +# Without the rewrite in ParquetSource, `file_row_index()` errors because it +# depends on file-source context + +query error file_row_index\(\) is source dependent and cannot be evaluated directly +SELECT file_row_index() + +# Testing pushdown over a source that doesn't support `file_row_index()`. + +statement ok +COPY (VALUES (10), (20), (30), (40), (50)) +TO 'test_files/scratch/file_row_index/csv_table/data.csv' +STORED AS CSV; + +statement ok +CREATE EXTERNAL TABLE csv_table(column1 int) +STORED AS CSV +LOCATION 'test_files/scratch/file_row_index/csv_table/data.csv'; + +query error file_row_index\(\) is source dependent and cannot be evaluated directly +SELECT *, file_row_index() FROM csv_table; + +# Testing a table with two files. + +statement ok +COPY (VALUES (10), (20)) +TO 'test_files/scratch/file_row_index/parquet_two_files/part-1.parquet' +STORED AS PARQUET; + +statement ok +COPY (VALUES (30), (40)) +TO 'test_files/scratch/file_row_index/parquet_two_files/part-2.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE parquet_two_files(column1 int) +STORED AS PARQUET +LOCATION 'test_files/scratch/file_row_index/parquet_two_files/'; + +query II +SELECT file_row_index(), column1 +FROM parquet_two_files +WHERE file_row_index() = 1 +ORDER BY column1 +---- +1 20 +1 40 + +statement ok +DROP TABLE parquet_two_files; + +statement ok +DROP TABLE parquet_table; + +statement ok +DROP TABLE csv_table; diff --git a/datafusion/sqllogictest/test_files/first_last_nested.slt b/datafusion/sqllogictest/test_files/first_last_nested.slt new file mode 100644 index 0000000000000..b96b47f6ab9c7 --- /dev/null +++ b/datafusion/sqllogictest/test_files/first_last_nested.slt @@ -0,0 +1,80 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# SQL-level coverage for first_value / last_value over nested payloads +# (Struct, Map) through the grouped `FirstLastGroupsAccumulator` path. +# +# The accumulator is unit-tested directly in first_last.rs. These tests +# add the integration coverage the unit tests cannot: the SLT runner +# executes with target_partitions = 4, so a GROUP BY drives the two-phase +# aggregate (Partial -> FinalPartitioned) and the nested intermediate +# state produced by `state()` is round-tripped back through +# `merge_batch()` across the partition boundary. Struct and Map otherwise +# have no SQL-level first_value / last_value coverage (only List does, in +# array_agg.slt). + +######################################## +# Struct payload +######################################## + +statement ok +CREATE TABLE first_last_struct AS VALUES + (1, 1, named_struct('a', 10, 'b', 'x')), + (1, 2, named_struct('a', 20, 'b', 'y')), + (1, 3, named_struct('a', 30, 'b', 'z')), + (2, 1, named_struct('a', 40, 'b', 'p')), + (2, 2, named_struct('a', 50, 'b', 'q')); + +query I?? +select column1, first_value(column3 order by column2), last_value(column3 order by column2) +from first_last_struct group by column1 order by column1; +---- +1 {a: 10, b: x} {a: 30, b: z} +2 {a: 40, b: p} {a: 50, b: q} + +# Descending order flips first / last. +query I?? +select column1, first_value(column3 order by column2 desc), last_value(column3 order by column2 desc) +from first_last_struct group by column1 order by column1; +---- +1 {a: 30, b: z} {a: 10, b: x} +2 {a: 50, b: q} {a: 40, b: p} + +statement ok +drop table first_last_struct; + +######################################## +# Map payload +######################################## + +statement ok +CREATE TABLE first_last_map AS VALUES + (1, 1, MAP {'k1': 10, 'k2': 20}), + (1, 2, MAP {'k3': 30}), + (1, 3, MAP {'k4': 40, 'k5': 50}), + (2, 1, MAP {'k9': 99}), + (2, 2, MAP {'k8': 88, 'k7': 77}); + +query I?? +select column1, first_value(column3 order by column2), last_value(column3 order by column2) +from first_last_map group by column1 order by column1; +---- +1 {k1: 10, k2: 20} {k4: 40, k5: 50} +2 {k9: 99} {k8: 88, k7: 77} + +statement ok +drop table first_last_map; diff --git a/datafusion/sqllogictest/test_files/floor_preimage.slt b/datafusion/sqllogictest/test_files/floor_preimage.slt index 960b966ebbba0..b54e2d37ee563 100644 --- a/datafusion/sqllogictest/test_files/floor_preimage.slt +++ b/datafusion/sqllogictest/test_files/floor_preimage.slt @@ -149,7 +149,7 @@ query TT EXPLAIN SELECT * FROM test_data WHERE floor(decimal_val) = arrow_cast(100, 'Decimal128(10,2)'); ---- logical_plan -01)Filter: test_data.decimal_val >= Decimal128(Some(10000),10,2) AND test_data.decimal_val < Decimal128(Some(10100),10,2) +01)Filter: test_data.decimal_val >= Decimal128(100.00,10,2) AND test_data.decimal_val < Decimal128(101.00,10,2) 02)--TableScan: test_data projection=[id, float_val, int_val, decimal_val] # 4. Column on RHS - same transformation diff --git a/datafusion/sqllogictest/test_files/functional_dependencies.slt b/datafusion/sqllogictest/test_files/functional_dependencies.slt new file mode 100644 index 0000000000000..c49004190dc60 --- /dev/null +++ b/datafusion/sqllogictest/test_files/functional_dependencies.slt @@ -0,0 +1,310 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +# Tests for functional dependencies +# (`datafusion/common/src/functional_dependencies.rs`) +# +# A functional dependency records that one set of columns (the *determinant*) +# determines the values of the others. DataFusion derives them from PRIMARY +# KEY / UNIQUE constraints and from GROUP BY keys, and four optimizer rules +# consume them to remove redundant work, each tested here in a different section. +# +# NULL handling is (as always) important: +# +# * A PRIMARY KEY is unique AND not nullable. +# * A `UNIQUE` constraint permits *multiple NULL rows*, because NULLs +# compare distinct. +# +# It is important not to mix `UNIQUE` columns with `DISTINCT` or `GROUP BY`, +# which treat NULLs as equal and can produce wrong answers. +########## + +# These rules all run during logical optimization, so show only logical plans. +statement ok +set datafusion.explain.logical_plan_only = true; + +# Set target_partitions explicitly so query results are stable. +statement ok +set datafusion.execution.target_partitions = 4; + +########## +## Test tables +########## + +statement ok +CREATE TABLE t_pk (x INT, y INT, PRIMARY KEY (x)) AS VALUES (1, 10), (2, 20); + +statement ok +CREATE TABLE t_uniq (x INT UNIQUE, y INT) AS VALUES (NULL, 2), (NULL, 1), (1, 3); + +query II rowsort +SELECT x, y FROM t_uniq; +---- +1 3 +NULL 1 +NULL 2 + + +# 1.1 PRIMARY KEY: rows are unique; the DISTINCT is removed and no +# Aggregate appears in the plan. +query TT +EXPLAIN SELECT DISTINCT x FROM t_pk; +---- +logical_plan TableScan: t_pk projection=[x] + +# 1.2 Nullable UNIQUE: the DISTINCT must be KEPT. UNIQUE allows several NULL +# rows, but DISTINCT treats NULLs as equal and has to collapse them into one. +# +# BUG: the DISTINCT is removed and both NULL rows are returned. +# Expected: `1`, `NULL`. +# Issue: https://github.com/apache/datafusion/issues/23634 +query I +SELECT DISTINCT x FROM t_uniq ORDER BY x NULLS LAST; +---- +1 +NULL +NULL + +query TT +EXPLAIN SELECT DISTINCT x FROM t_uniq; +---- +logical_plan TableScan: t_uniq projection=[x] + +# 1.3 A PRIMARY KEY downgraded to a non-unique dependency by a LEFT JOIN +# so the DISTINCT must be KEPT. +# Fixed by: https://github.com/apache/datafusion/pull/23548 +statement ok +CREATE TABLE t_orders (x INT, amount INT) AS VALUES (1, 10), (1, 20), (2, 30); + +query I +SELECT DISTINCT p.x FROM t_pk p LEFT JOIN t_orders o ON p.x = o.x ORDER BY p.x; +---- +1 +2 + +query TT +EXPLAIN SELECT DISTINCT p.x FROM t_pk p LEFT JOIN t_orders o ON p.x = o.x; +---- +logical_plan +01)Aggregate: groupBy=[[p.x]], aggr=[[]] +02)--SubqueryAlias: p +03)----TableScan: t_pk projection=[x] + +statement ok +drop table t_orders; + +# 1.4 DISTINCT over a GROUP BY output. Grouping collapses the multiple NULL +# rows, (NULL included) and the DISTINCT can be removed. +query I +SELECT DISTINCT x FROM (SELECT x FROM t_uniq GROUP BY x) ORDER BY x NULLS LAST; +---- +1 +NULL + +query TT +EXPLAIN SELECT DISTINCT x FROM (SELECT x FROM t_uniq GROUP BY x); +---- +logical_plan +01)Aggregate: groupBy=[[t_uniq.x]], aggr=[[]] +02)--TableScan: t_uniq projection=[x] + + +# 2.1 PRIMARY KEY: `x` determines `y`, so `ORDER BY x, y` is equivalent to +# `ORDER BY x` and the `y` key is dropped from the plan. +query TT +EXPLAIN SELECT x, y FROM t_pk ORDER BY x, y; +---- +logical_plan +01)Sort: t_pk.x ASC NULLS LAST +02)--TableScan: t_pk projection=[x, y] + +# 2.2 Nullable UNIQUE: `x` does NOT determine `y` across the two NULL rows, +# so the `y` sort key must be kept. +# +# BUG: +# Expected: `1 3`, `NULL 1`, `NULL 2`. +# Issue: https://github.com/apache/datafusion/issues/23818 +query II +SELECT x, y FROM t_uniq ORDER BY x NULLS LAST, y; +---- +1 3 +NULL 2 +NULL 1 + +query TT +EXPLAIN SELECT x, y FROM t_uniq ORDER BY x NULLS LAST, y; +---- +logical_plan +01)Sort: t_uniq.x ASC NULLS LAST +02)--TableScan: t_uniq projection=[x, y] + +# 2.3 After `GROUP BY x` the `x` does determine `cnt`, so can drop `cnt` from sort +query TT +EXPLAIN SELECT x, cnt FROM (SELECT x, count(*) AS cnt FROM t_uniq GROUP BY x) ORDER BY x, cnt; +---- +logical_plan +01)Sort: t_uniq.x ASC NULLS LAST +02)--Projection: t_uniq.x, count(Int64(1)) AS cnt +03)----Aggregate: groupBy=[[t_uniq.x]], aggr=[[count(Int64(1))]] +04)------TableScan: t_uniq projection=[x] + + +# 3.1 PRIMARY KEY: `x` determines `y`, and `y` is not selected, so grouping +# by `x, y` is the same as grouping by `x`. +query TT +EXPLAIN SELECT x FROM t_pk GROUP BY x, y; +---- +logical_plan +01)Aggregate: groupBy=[[t_pk.x]], aggr=[[]] +02)--TableScan: t_pk projection=[x] + +# 3.2 Nullable UNIQUE: grouping by `x, y` is NOT the same as grouping by +# `x` -- two NULL rows differ in `y` and belong in separate groups. +# +# BUG: `y` is dropped from the GROUP BY and the two NULL groups are merged, +# so one row goes missing. +# Expected: `1`, `NULL`, `NULL` (three rows). +# Issue: https://github.com/apache/datafusion/issues/23819 +query I rowsort +SELECT x FROM t_uniq GROUP BY x, y; +---- +1 +NULL + +query TT +EXPLAIN SELECT x FROM t_uniq GROUP BY x, y; +---- +logical_plan +01)Aggregate: groupBy=[[t_uniq.x]], aggr=[[]] +02)--TableScan: t_uniq projection=[x] + +# 3.3 The same grouping, but with `y` selected so the parent needs it: no +# column can be dropped and the answer is right. +query II rowsort +SELECT x, y FROM t_uniq GROUP BY x, y; +---- +1 3 +NULL 1 +NULL 2 + +# 4.1 PRIMARY KEY: `x` determines `y`, so `y` has a single well-defined +# value per group and one row is returned per `x`. +query II rowsort +SELECT x, y FROM t_pk GROUP BY x; +---- +1 10 +2 20 + +query TT +EXPLAIN SELECT x, y FROM t_pk GROUP BY x; +---- +logical_plan +01)Aggregate: groupBy=[[t_pk.x, t_pk.y]], aggr=[[]] +02)--TableScan: t_pk projection=[x, y] + +# 4.2 Nullable UNIQUE: `x` does NOT determine `y`, so there is no +# well-defined `y` for the `x = NULL` group. +# +# BUG: `y` is appended to the GROUP BY anyway, so `GROUP BY x` returns TWO +# rows for `x = NULL`. +# Expected: one row per distinct `x` (or a planning error -- postgres +# rejects this query, and accepts the 4.1 PRIMARY KEY form). +# Issue: https://github.com/apache/datafusion/issues/23820 +query II rowsort +SELECT x, y FROM t_uniq GROUP BY x; +---- +1 3 +NULL 1 +NULL 2 + +query TT +EXPLAIN SELECT x, y FROM t_uniq GROUP BY x; +---- +logical_plan +01)Aggregate: groupBy=[[t_uniq.x, t_uniq.y]], aggr=[[]] +02)--TableScan: t_uniq projection=[x, y] + + +statement ok +CREATE TABLE t_null (x INT) AS VALUES (NULL), (NULL); + +statement ok +CREATE TABLE t_probe (z INT) AS VALUES (0), (2); + +# 5.1 Grouping by `g.x, g.cnt` must keep both columns: `g.x` alone does not +# determine `g.cnt` after NULL padding. +query II +SELECT g.x, count(*) AS c + FROM t_probe a + LEFT JOIN (SELECT x, count(*) AS cnt FROM t_null GROUP BY x) g + ON a.z = g.cnt + GROUP BY g.x, g.cnt + ORDER BY c; +---- +NULL 1 +NULL 1 + +query TT +EXPLAIN SELECT g.x, count(*) AS c + FROM t_probe a + LEFT JOIN (SELECT x, count(*) AS cnt FROM t_null GROUP BY x) g + ON a.z = g.cnt + GROUP BY g.x, g.cnt; +---- +logical_plan +01)Projection: g.x, count(Int64(1)) AS count(*) AS c +02)--Aggregate: groupBy=[[g.x, g.cnt]], aggr=[[count(Int64(1))]] +03)----Projection: g.x, g.cnt +04)------Left Join: CAST(a.z AS Int64) = g.cnt +05)--------SubqueryAlias: a +06)----------TableScan: t_probe projection=[z] +07)--------SubqueryAlias: g +08)----------Projection: t_null.x, count(Int64(1)) AS count(*) AS cnt +09)------------Aggregate: groupBy=[[t_null.x]], aggr=[[count(Int64(1))]] +10)--------------TableScan: t_null projection=[x] + +# 5.2 The ORDER BY variant: `g.x` is NULL for both rows, so the `g.cnt` +# tie-breaker is what orders them. +query II +SELECT g.x, g.cnt + FROM t_probe a + LEFT JOIN (SELECT x, count(*) AS cnt FROM t_null GROUP BY x) g + ON a.z = g.cnt + ORDER BY g.x, g.cnt; +---- +NULL 2 +NULL NULL + +statement ok +drop table t_null; + +statement ok +drop table t_probe; + +########## +## Cleanup +########## + +statement ok +drop table t_pk; + +statement ok +drop table t_uniq; + +statement ok +RESET datafusion.explain.logical_plan_only; diff --git a/datafusion/sqllogictest/test_files/functions.slt b/datafusion/sqllogictest/test_files/functions.slt index ea3cd6eb4bd33..008be05852c85 100644 --- a/datafusion/sqllogictest/test_files/functions.slt +++ b/datafusion/sqllogictest/test_files/functions.slt @@ -68,7 +68,7 @@ SELECT length('') ---- 0 -query I +query ? SELECT length(arrow_cast('', 'Dictionary(Int32, Utf8)')) ---- 0 @@ -83,7 +83,7 @@ SELECT length('josé') ---- 4 -query I +query ? SELECT length(arrow_cast('josé', 'Dictionary(Int32, Utf8)')) ---- 4 @@ -468,17 +468,35 @@ Utf8View query T SELECT arrow_typeof(upper(arrow_cast(arrow_cast('foo', 'Dictionary(Int32, Utf8)'), 'Dictionary(Int32, Utf8View)'))) ---- -Utf8View +Dictionary(Int32, Utf8View) + +statement ok +CREATE TABLE upper_dictionary_test AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS c1 FROM (VALUES +('foo'), +(NULL), +('Bar')); + +query TT +SELECT upper(c1), arrow_typeof(upper(c1)) FROM upper_dictionary_test +---- +FOO Dictionary(Int32, Utf8) +NULL Dictionary(Int32, Utf8) +BAR Dictionary(Int32, Utf8) + +statement ok +DROP TABLE upper_dictionary_test query T SELECT btrim(' foo ') ---- foo -query T -SELECT btrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')) +query TT +SELECT btrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')), + arrow_typeof(btrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)'))) ---- -foo +foo Dictionary(Int32, Utf8) query T SELECT initcap('foo') @@ -490,6 +508,106 @@ SELECT initcap(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- Foo +query TTTT +SELECT initcap(arrow_cast('foo BAR', 'Dictionary(Int32, LargeUtf8)')), + arrow_typeof(initcap(arrow_cast('foo BAR', 'Dictionary(Int32, LargeUtf8)'))), + initcap(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + )), + arrow_typeof(initcap(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + ))) +---- +Foo Bar Dictionary(Int32, LargeUtf8) Foo Bar Dictionary(Int32, Utf8View) + +query ?T +SELECT initcap(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(initcap(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +Foo Bar Dictionary(Int32, Dictionary(UInt32, Utf8)) + +statement ok +CREATE TABLE unicode_dictionary_test AS +SELECT column1 AS id, + arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS dict_col, + arrow_cast( + arrow_cast(column2, 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ) AS nested_dict_col +FROM (VALUES +(1, 'foo BAR'), +(2, 'éclair CAFÉ'), +(3, NULL)); + +query T?TT +SELECT initcap(dict_col), initcap(nested_dict_col), + arrow_typeof(initcap(dict_col)), + arrow_typeof(initcap(nested_dict_col)) +FROM unicode_dictionary_test +ORDER BY id +---- +Foo Bar Foo Bar Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) +Éclair Café Éclair Café Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) +NULL NULL Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) + +query TTTT +SELECT reverse(arrow_cast('foo BAR', 'Dictionary(Int32, LargeUtf8)')), + arrow_typeof(reverse(arrow_cast('foo BAR', 'Dictionary(Int32, LargeUtf8)'))), + reverse(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + )), + arrow_typeof(reverse(arrow_cast( + arrow_cast('foo BAR', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + ))) +---- +RAB oof Dictionary(Int32, LargeUtf8) RAB oof Dictionary(Int32, Utf8View) + +query T?TT +SELECT reverse(dict_col), reverse(nested_dict_col), + arrow_typeof(reverse(dict_col)), + arrow_typeof(reverse(nested_dict_col)) +FROM unicode_dictionary_test +ORDER BY id +---- +RAB oof RAB oof Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) +ÉFAC rialcé ÉFAC rialcé Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) +NULL NULL Dictionary(Int32, Utf8) Dictionary(Int32, Dictionary(UInt32, Utf8)) + +statement ok +DROP TABLE unicode_dictionary_test + +query ? +SELECT ascii(arrow_cast('é', 'Dictionary(Int32, Utf8)')) +---- +233 + +query T +SELECT arrow_typeof(ascii(arrow_cast('é', 'Dictionary(Int32, Utf8)'))) +---- +Dictionary(Int32, Int32) + +query ?T +SELECT ascii(arrow_cast( + arrow_cast('💯', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(ascii(arrow_cast( + arrow_cast('💯', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +128175 Dictionary(Int32, Dictionary(UInt32, Int32)) + query I SELECT instr('foobarbar', 'bar') ---- @@ -548,17 +666,35 @@ Utf8View query T SELECT arrow_typeof(lower(arrow_cast(arrow_cast('FOObar', 'Dictionary(Int32, Utf8)'), 'Dictionary(Int32, Utf8View)'))) ---- -Utf8View +Dictionary(Int32, Utf8View) + +statement ok +CREATE TABLE lower_dictionary_test AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS c1 FROM (VALUES +('FOO'), +(NULL), +('Bar')); + +query TT +SELECT lower(c1), arrow_typeof(lower(c1)) FROM lower_dictionary_test +---- +foo Dictionary(Int32, Utf8) +NULL Dictionary(Int32, Utf8) +bar Dictionary(Int32, Utf8) + +statement ok +DROP TABLE lower_dictionary_test query T SELECT ltrim(' foo') ---- foo -query T -SELECT ltrim(arrow_cast(' foo', 'Dictionary(Int32, Utf8)')) +query TT +SELECT ltrim(arrow_cast(' foo', 'Dictionary(Int32, Utf8)')), + arrow_typeof(ltrim(arrow_cast(' foo', 'Dictionary(Int32, Utf8)'))) ---- -foo +foo Dictionary(Int32, Utf8) query T SELECT md5('foo') @@ -576,20 +712,75 @@ SELECT rtrim(' foo ') ---- foo -query T -SELECT rtrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')) +query TT +SELECT rtrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')), + arrow_typeof(rtrim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)'))) ---- - foo + foo Dictionary(Int32, Utf8) query T SELECT trim(' foo ') ---- foo -query T -SELECT trim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')) +query TT +SELECT trim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)')), + arrow_typeof(trim(arrow_cast(' foo ', 'Dictionary(Int32, Utf8)'))) +---- +foo Dictionary(Int32, Utf8) + +query TTTT?T +SELECT btrim(arrow_cast(' foo ', 'Dictionary(Int32, LargeUtf8)')), + arrow_typeof(btrim(arrow_cast(' foo ', 'Dictionary(Int32, LargeUtf8)'))), + ltrim(arrow_cast( + arrow_cast(' bar', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + )), + arrow_typeof(ltrim(arrow_cast( + arrow_cast(' bar', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + ))), + rtrim(arrow_cast( + arrow_cast('baz ', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(rtrim(arrow_cast( + arrow_cast('baz ', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +foo Dictionary(Int32, LargeUtf8) bar Dictionary(Int32, Utf8View) baz Dictionary(Int32, Dictionary(UInt32, Utf8)) + +statement ok +CREATE TABLE trim_dictionary_test AS +SELECT arrow_cast(column1, 'Dictionary(Int32, Utf8)') AS both, + arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS leading, + arrow_cast(column3, 'Dictionary(Int32, Utf8)') AS trailing +FROM (VALUES +(' foo ', ' bar', 'baz '), +(NULL, NULL, NULL)); + +query TTTTTT +SELECT btrim(both), arrow_typeof(btrim(both)), + ltrim(leading), arrow_typeof(ltrim(leading)), + rtrim(trailing), arrow_typeof(rtrim(trailing)) +FROM trim_dictionary_test +---- +foo Dictionary(Int32, Utf8) bar Dictionary(Int32, Utf8) baz Dictionary(Int32, Utf8) +NULL Dictionary(Int32, Utf8) NULL Dictionary(Int32, Utf8) NULL Dictionary(Int32, Utf8) + +statement ok +DROP TABLE trim_dictionary_test + +query TTTTTT +SELECT btrim(arrow_cast('__foo__', 'Dictionary(Int32, Utf8)'), '_'), + arrow_typeof(btrim(arrow_cast('__foo__', 'Dictionary(Int32, Utf8)'), '_')), + ltrim(arrow_cast('__bar', 'Dictionary(Int32, Utf8)'), '_'), + arrow_typeof(ltrim(arrow_cast('__bar', 'Dictionary(Int32, Utf8)'), '_')), + rtrim(arrow_cast('baz__', 'Dictionary(Int32, Utf8)'), '_'), + arrow_typeof(rtrim(arrow_cast('baz__', 'Dictionary(Int32, Utf8)'), '_')) ---- -foo +foo Utf8 bar Utf8 baz Utf8 # Verify that trim, ltrim, and rtrim only strip spaces by default, # not other whitespace characters (tabs, newlines, etc.) @@ -605,31 +796,153 @@ SELECT bit_length('foo') ---- 24 -query I +query ? SELECT bit_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- 24 +query T +SELECT arrow_typeof(bit_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)'))) +---- +Dictionary(Int32, Int32) + +query ?T +SELECT bit_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(bit_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +16 Dictionary(Int32, Dictionary(UInt32, Int32)) + query I SELECT character_length('foo') ---- 3 -query I +query ? SELECT character_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- 3 +query ?T +SELECT character_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(character_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +1 Dictionary(Int32, Dictionary(UInt32, Int32)) + +query ?T?T +SELECT character_length(arrow_cast('foo', 'Dictionary(Int32, LargeUtf8)')), + arrow_typeof(character_length( + arrow_cast('foo', 'Dictionary(Int32, LargeUtf8)') + )), + character_length(arrow_cast( + arrow_cast('foo', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + )), + arrow_typeof(character_length(arrow_cast( + arrow_cast('foo', 'Dictionary(Int32, Utf8)'), + 'Dictionary(Int32, Utf8View)' + ))) +---- +3 Dictionary(Int32, Int64) 3 Dictionary(Int32, Int32) + query I SELECT octet_length('foo') ---- 3 -query I +query ? SELECT octet_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)')) ---- 3 +query T +SELECT arrow_typeof(octet_length(arrow_cast('foo', 'Dictionary(Int32, Utf8)'))) +---- +Dictionary(Int32, Int32) + +query ?T +SELECT octet_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + )), + arrow_typeof(octet_length(arrow_cast( + arrow_cast('é', 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ))) +---- +2 Dictionary(Int32, Dictionary(UInt32, Int32)) + +statement ok +CREATE TABLE string_length_dictionary_test AS +SELECT column1 AS id, + arrow_cast(column2, 'Dictionary(Int32, Utf8)') AS dict_col, + arrow_cast( + arrow_cast(column2, 'Dictionary(UInt32, Utf8)'), + 'Dictionary(Int32, Dictionary(UInt32, Utf8))' + ) AS nested_dict_col +FROM (VALUES +(1, 'foo'), +(2, 'é'), +(3, NULL)); + +query ??TT +SELECT bit_length(dict_col), bit_length(nested_dict_col), + arrow_typeof(bit_length(dict_col)), + arrow_typeof(bit_length(nested_dict_col)) +FROM string_length_dictionary_test +ORDER BY id +---- +24 24 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +16 16 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) + +query ??TT +SELECT octet_length(dict_col), octet_length(nested_dict_col), + arrow_typeof(octet_length(dict_col)), + arrow_typeof(octet_length(nested_dict_col)) +FROM string_length_dictionary_test +ORDER BY id +---- +3 3 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +2 2 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) + +query ??TT +SELECT character_length(dict_col), character_length(nested_dict_col), + arrow_typeof(character_length(dict_col)), + arrow_typeof(character_length(nested_dict_col)) +FROM string_length_dictionary_test +ORDER BY id +---- +3 3 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +1 1 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) + +query ??TT +SELECT ascii(dict_col), ascii(nested_dict_col), + arrow_typeof(ascii(dict_col)), + arrow_typeof(ascii(nested_dict_col)) +FROM string_length_dictionary_test +ORDER BY id +---- +102 102 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +233 233 Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) +NULL NULL Dictionary(Int32, Int32) Dictionary(Int32, Dictionary(UInt32, Int32)) + +statement ok +DROP TABLE string_length_dictionary_test + query I SELECT strpos('helloworld', 'world') ---- @@ -805,10 +1118,10 @@ SELECT overlay('abc' placing 'X' from 5 for 1) abcX # Start positions must be positive. -statement error negative substring length not allowed +statement error overlay start position must be at least 1: 0 SELECT overlay('abc' placing 'X' from 0 for 1) -statement error negative substring length not allowed +statement error overlay start position must be at least 1: -1 SELECT overlay('abc' placing 'X' from -1 for 1) # Negative count keeps the suffix from before the start position. diff --git a/datafusion/sqllogictest/test_files/group_by.slt b/datafusion/sqllogictest/test_files/group_by.slt index 8c055c25caeb2..38d1b7821451d 100644 --- a/datafusion/sqllogictest/test_files/group_by.slt +++ b/datafusion/sqllogictest/test_files/group_by.slt @@ -2942,8 +2942,8 @@ logical_plan 08)----------SubqueryAlias: e 09)------------TableScan: sales_global projection=[sn, ts, currency, amount] physical_plan -01)SortExec: expr=[sn@2 ASC NULLS LAST], preserve_partitioning=[false] -02)--ProjectionExec: expr=[zip_code@1 as zip_code, country@2 as country, sn@0 as sn, ts@3 as ts, currency@4 as currency, last_value(e.amount) ORDER BY [e.sn ASC NULLS LAST]@5 as last_rate] +01)ProjectionExec: expr=[zip_code@1 as zip_code, country@2 as country, sn@0 as sn, ts@3 as ts, currency@4 as currency, last_value(e.amount) ORDER BY [e.sn ASC NULLS LAST]@5 as last_rate] +02)--SortExec: expr=[sn@0 ASC NULLS LAST], preserve_partitioning=[false] 03)----AggregateExec: mode=Single, gby=[sn@2 as sn, zip_code@0 as zip_code, country@1 as country, ts@3 as ts, currency@4 as currency], aggr=[last_value(e.amount) ORDER BY [e.sn ASC NULLS LAST]] 04)------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(currency@2, currency@4)], filter=ts@0 >= ts@1, projection=[zip_code@4, country@5, sn@6, ts@7, currency@8, sn@0, amount@3] 05)--------DataSourceExec: partitions=1, partition_sizes=[1] @@ -2985,8 +2985,8 @@ logical_plan 04)------TableScan: sales_global projection=[country, ts, amount] physical_plan 01)SortPreservingMergeExec: [country@0 ASC NULLS LAST] -02)--SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[country@0 as country, first_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST]@1 as fv1, last_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST]@2 as fv2] +02)--ProjectionExec: expr=[country@0 as country, first_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST]@1 as fv1, last_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST]@2 as fv2] +03)----SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[country@0 as country], aggr=[first_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST], last_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST]] 05)--------RepartitionExec: partitioning=Hash([country@0], 8), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[country@0 as country], aggr=[first_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST], last_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST]] @@ -3019,8 +3019,8 @@ logical_plan 04)------TableScan: sales_global projection=[country, ts, amount] physical_plan 01)SortPreservingMergeExec: [country@0 ASC NULLS LAST] -02)--SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[country@0 as country, first_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST]@1 as fv1, last_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@2 as fv2] +02)--ProjectionExec: expr=[country@0 as country, first_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST]@1 as fv1, last_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]@2 as fv2] +03)----SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[country@0 as country], aggr=[first_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST], last_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]] 05)--------RepartitionExec: partitioning=Hash([country@0], 8), input_partitions=1 06)----------AggregateExec: mode=Partial, gby=[country@0 as country], aggr=[first_value(sales_global.amount) ORDER BY [sales_global.ts ASC NULLS LAST], last_value(sales_global.amount) ORDER BY [sales_global.ts DESC NULLS FIRST]] @@ -3181,8 +3181,8 @@ logical_plan 04)------TableScan: sales_global projection=[country, amount] physical_plan 01)SortPreservingMergeExec: [country@0 ASC NULLS LAST] -02)--SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[country@0 as country, array_agg(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@1 as array_agg1] +02)--ProjectionExec: expr=[country@0 as country, array_agg(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@1 as array_agg1] +03)----SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[country@0 as country], aggr=[array_agg(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]] 05)--------RepartitionExec: partitioning=Hash([country@0], 8), input_partitions=8 06)----------AggregateExec: mode=Partial, gby=[country@0 as country], aggr=[array_agg(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]] @@ -3216,8 +3216,8 @@ logical_plan 04)------TableScan: sales_global projection=[country, amount] physical_plan 01)SortPreservingMergeExec: [country@0 ASC NULLS LAST] -02)--SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[country@0 as country, array_agg(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@1 as amounts, first_value(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@2 as fv1, last_value(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@3 as fv2] +02)--ProjectionExec: expr=[country@0 as country, array_agg(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@1 as amounts, first_value(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST]@2 as fv1, last_value(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]@3 as fv2] +03)----SortExec: expr=[country@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[country@0 as country], aggr=[array_agg(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST], first_value(sales_global.amount) ORDER BY [sales_global.amount ASC NULLS LAST], last_value(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]] 05)--------RepartitionExec: partitioning=Hash([country@0], 8), input_partitions=8 06)----------AggregateExec: mode=Partial, gby=[country@0 as country], aggr=[array_agg(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST], last_value(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST], last_value(sales_global.amount) ORDER BY [sales_global.amount DESC NULLS FIRST]] @@ -3484,8 +3484,8 @@ logical_plan 09)------------TableScan: sales_global_with_pk projection=[sn, amount] physical_plan 01)SortPreservingMergeExec: [sn@0 ASC NULLS LAST] -02)--SortExec: expr=[sn@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[sn@0 as sn, sum(l.amount)@2 as sum(l.amount), amount@1 as amount] +02)--ProjectionExec: expr=[sn@0 as sn, sum(l.amount)@2 as sum(l.amount), amount@1 as amount] +03)----SortExec: expr=[sn@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[sn@0 as sn, amount@1 as amount], aggr=[sum(l.amount)] 05)--------RepartitionExec: partitioning=Hash([sn@0, amount@1], 8), input_partitions=8 06)----------AggregateExec: mode=Partial, gby=[sn@1 as sn, amount@2 as amount], aggr=[sum(l.amount)] @@ -3630,8 +3630,8 @@ logical_plan 08)--------------TableScan: sales_global_with_pk projection=[zip_code, country, sn, ts, currency, amount] physical_plan 01)SortPreservingMergeExec: [sn@2 ASC NULLS LAST] -02)--SortExec: expr=[sn@2 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[zip_code@1 as zip_code, country@2 as country, sn@0 as sn, ts@3 as ts, currency@4 as currency, amount@5 as amount, sum_amount@6 as sum_amount] +02)--ProjectionExec: expr=[zip_code@1 as zip_code, country@2 as country, sn@0 as sn, ts@3 as ts, currency@4 as currency, amount@5 as amount, sum_amount@6 as sum_amount] +03)----SortExec: expr=[sn@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[sn@0 as sn, zip_code@1 as zip_code, country@2 as country, ts@3 as ts, currency@4 as currency, amount@5 as amount, sum_amount@6 as sum_amount], aggr=[] 05)--------RepartitionExec: partitioning=Hash([sn@0, zip_code@1, country@2, ts@3, currency@4, amount@5, sum_amount@6], 8), input_partitions=8 06)----------AggregateExec: mode=Partial, gby=[sn@2 as sn, zip_code@0 as zip_code, country@1 as country, ts@3 as ts, currency@4 as currency, amount@5 as amount, sum_amount@6 as sum_amount], aggr=[] @@ -4327,8 +4327,8 @@ logical_plan 04)------TableScan: csv_with_timestamps projection=[ts] physical_plan 01)SortPreservingMergeExec: [months@0 DESC], fetch=5 -02)--SortExec: TopK(fetch=5), expr=[months@0 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[date_part(Utf8("MONTH"),csv_with_timestamps.ts)@0 as months] +02)--ProjectionExec: expr=[date_part(Utf8("MONTH"),csv_with_timestamps.ts)@0 as months] +03)----SortExec: TopK(fetch=5), expr=[date_part(Utf8("MONTH"),csv_with_timestamps.ts)@0 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[date_part(Utf8("MONTH"),csv_with_timestamps.ts)@0 as date_part(Utf8("MONTH"),csv_with_timestamps.ts)], aggr=[], lim=[5] 05)--------RepartitionExec: partitioning=Hash([date_part(Utf8("MONTH"),csv_with_timestamps.ts)@0], 8), input_partitions=8 06)----------AggregateExec: mode=Partial, gby=[date_part(MONTH, ts@0) as date_part(Utf8("MONTH"),csv_with_timestamps.ts)], aggr=[], lim=[5] @@ -4438,8 +4438,8 @@ logical_plan 05)--------TableScan: aggregate_test_100 projection=[c1, c2, c3, c4] physical_plan 01)SortPreservingMergeExec: [c1@0 ASC NULLS LAST] -02)--SortExec: expr=[c1@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[c1@0 as c1, count(alias1)@1 as count(DISTINCT aggregate_test_100.c2), min(alias1)@2 as min(DISTINCT aggregate_test_100.c2), sum(alias2)@3 as sum(aggregate_test_100.c3), max(alias3)@4 as max(aggregate_test_100.c4)] +02)--ProjectionExec: expr=[c1@0 as c1, count(alias1)@1 as count(DISTINCT aggregate_test_100.c2), min(alias1)@2 as min(DISTINCT aggregate_test_100.c2), sum(alias2)@3 as sum(aggregate_test_100.c3), max(alias3)@4 as max(aggregate_test_100.c4)] +03)----SortExec: expr=[c1@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[c1@0 as c1], aggr=[count(alias1), min(alias1), sum(alias2), max(alias3)] 05)--------RepartitionExec: partitioning=Hash([c1@0], 8), input_partitions=8 06)----------AggregateExec: mode=Partial, gby=[c1@0 as c1], aggr=[count(alias1), min(alias1), sum(alias2), max(alias3)] @@ -4608,9 +4608,9 @@ logical_plan physical_plan 01)SortPreservingMergeExec: [max(timestamp_table.t1)@1 DESC], fetch=4 02)--SortExec: TopK(fetch=4), expr=[max(timestamp_table.t1)@1 DESC], preserve_partitioning=[true] -03)----AggregateExec: mode=FinalPartitioned, gby=[c2@0 as c2], aggr=[max(timestamp_table.t1)], lim=[4] +03)----AggregateExec: mode=FinalPartitioned, gby=[c2@0 as c2], aggr=[max(timestamp_table.t1)] 04)------RepartitionExec: partitioning=Hash([c2@0], 8), input_partitions=8 -05)--------AggregateExec: mode=Partial, gby=[c2@1 as c2], aggr=[max(timestamp_table.t1)], lim=[4] +05)--------AggregateExec: mode=Partial, gby=[c2@1 as c2], aggr=[max(timestamp_table.t1)] 06)----------RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=4 07)------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/group_by/timestamp_table/0.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/group_by/timestamp_table/1.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/group_by/timestamp_table/2.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/group_by/timestamp_table/3.csv]]}, projection=[t1, c2], file_type=csv, has_header=true @@ -5634,10 +5634,150 @@ physical_plan statement ok reset datafusion.execution.batch_size; -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; statement count 0 drop table t; + +# DISTINCT must not be removed when a unique key is downgraded to a +# non-unique functional dependency by a join: `u.id` is a primary key, but +# after the LEFT JOIN each `u` row can occur once per matching order. +statement ok +CREATE TABLE users_with_pk (id INT, name VARCHAR, primary key(id)) AS VALUES + (1, 'alice'), + (2, 'bob'); + +statement ok +CREATE TABLE user_orders (user_id INT, amount INT) AS VALUES + (1, 10), + (1, 20), + (2, 30); + +query I +SELECT DISTINCT u.id + FROM users_with_pk u + LEFT JOIN user_orders o ON u.id = o.user_id + ORDER BY u.id; +---- +1 +2 + +# The DISTINCT must be planned as an Aggregate; it cannot be removed based +# on the (join-downgraded) primary key of `users_with_pk`. +query TT +EXPLAIN SELECT DISTINCT u.id + FROM users_with_pk u + LEFT JOIN user_orders o ON u.id = o.user_id; +---- +logical_plan +01)Aggregate: groupBy=[[u.id]], aggr=[[]] +02)--SubqueryAlias: u +03)----TableScan: users_with_pk projection=[id] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[] +02)--RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +drop table users_with_pk; + +statement ok +drop table user_orders; + +# Test multi group by int + Duration +statement ok +CREATE TABLE duration_group_test AS VALUES + (1, arrow_cast(5, 'Duration(Second)')), + (1, arrow_cast(5, 'Duration(Second)')), + (1, arrow_cast(7, 'Duration(Second)')), + (2, arrow_cast(5, 'Duration(Second)')); + +# Single Duration group key ({5s, 7s}, 5s x3) via the GroupValuesPrimitive path. +query ?I +SELECT column2, count(*) FROM duration_group_test GROUP BY column2 ORDER BY column2; +---- +0 days 0 hours 0 mins 5 secs 3 +0 days 0 hours 0 mins 7 secs 1 + +# Multi-column GROUP BY: a primitive key and a Duration key on the same path. +query I?I +SELECT column1, column2, count(*) +FROM duration_group_test GROUP BY column1, column2 ORDER BY column1, column2; +---- +1 0 days 0 hours 0 mins 5 secs 2 +1 0 days 0 hours 0 mins 7 secs 1 +2 0 days 0 hours 0 mins 5 secs 1 + +statement ok +DROP TABLE duration_group_test; + +# Test multi group by int + Float16 +statement ok +CREATE TABLE float16_group_test AS VALUES + (arrow_cast(1.5, 'Float16'), 1), + (arrow_cast(1.5, 'Float16'), 1), + (arrow_cast(2.5, 'Float16'), 2), + (arrow_cast(-0.0, 'Float16'), 3), + (arrow_cast(0.0, 'Float16'), 3); + +# Single Float16 group key ({1.5, 2.5, ±0.0}) via the GroupValuesPrimitive path. +query I +SELECT count(*) FROM (SELECT column1 FROM float16_group_test GROUP BY column1); +---- +3 + +# Multi-column GROUP BY: a primitive key and a Float16 key on the same path. +query I +SELECT count(*) FROM (SELECT column1, column2 FROM float16_group_test GROUP BY column1, column2); +---- +3 + +statement ok +DROP TABLE float16_group_test; + +# Test multi group by int + Interval +statement ok +CREATE TABLE interval_group_test AS VALUES + (1, INTERVAL '1' MONTH), + (1, INTERVAL '1' MONTH), + (1, INTERVAL '30' DAY), + (2, INTERVAL '1' MONTH); + +# Single Interval group key ({1 month, 30 days}) via the GroupValuesPrimitive path. +query I +SELECT count(*) FROM interval_group_test GROUP BY column2 ORDER BY count(*); +---- +1 +3 + +# Multi-column GROUP BY: a primitive key and an Interval key on the same path. +query II +SELECT column1, count(*) +FROM interval_group_test GROUP BY column1, column2 ORDER BY column1, count(*); +---- +1 1 +1 2 +2 1 + +statement ok +DROP TABLE interval_group_test; + +# Test multi group by int + Decimal256 +statement ok +create table decimal256_multi_group (k int, d decimal(50, 2)) as values + (1, 100.00), (1, 100.00), (1, 250.00), (2, 100.00), (2, NULL); + +query IRI +select k, d, count(*) from decimal256_multi_group group by k, d order by k, d; +---- +1 100 2 +1 250 1 +2 100 1 +2 NULL 1 + +statement ok +drop table decimal256_multi_group; diff --git a/datafusion/sqllogictest/test_files/grouping.slt b/datafusion/sqllogictest/test_files/grouping.slt index eac901b2a300f..7d893f94be74a 100644 --- a/datafusion/sqllogictest/test_files/grouping.slt +++ b/datafusion/sqllogictest/test_files/grouping.slt @@ -232,6 +232,12 @@ SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING ---- NULL +# grouping_sets_empty_input_avg: AVG returns NULL for the empty group +query R +SELECT AVG(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS(()) +---- +NULL + # grouping_sets_empty_input_count: COUNT returns 0 for the empty group, not a missing row query I SELECT COUNT(*) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS(()) @@ -255,3 +261,90 @@ query II SELECT SUM(v1), COUNT(*) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS((), (v1)) ---- NULL 0 + +# rollup_empty_input_outer_count: an outer COUNT(*) over ROLLUP is answered from the inner row-count statistics +query I +SELECT COUNT(*) FROM (SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY ROLLUP(v1)) +---- +1 + +# cube_empty_input_outer_count: an outer COUNT(*) over CUBE is answered from the inner row-count statistics +query I +SELECT COUNT(*) FROM (SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY CUBE(v1)) +---- +1 + +# grouping_sets_empty_input_outer_count: an outer COUNT(*) over GROUPING SETS is answered from the inner row-count statistics +query I +SELECT COUNT(*) FROM (SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS((), (v1))) +---- +1 + +# duplicate_empty_grouping_sets_empty_input: each empty grouping set emits its own grand-total row +query I +SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS((v1), (), ()) +---- +NULL +NULL + +# duplicate_empty_grouping_sets_empty_input_outer_count: the row-count statistics must match those two rows +query I +SELECT COUNT(*) FROM (SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS((v1), (), ())) +---- +2 + +# duplicate_empty_grouping_sets_empty_input_limit: LIMIT applies to the grand-total rows, so one of +# the two is returned +query I +SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS((v1), (), ()) LIMIT 1 +---- +NULL + +# duplicate_empty_grouping_sets_empty_input_limit_outer_count: the outer COUNT(*) is answered from +# the inner row-count statistics, which must agree with the rows the limit lets through +query I +SELECT COUNT(*) FROM (SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS((v1), (), ()) LIMIT 1) +---- +1 + +# duplicate_empty_grouping_sets_empty_input_limit_above_row_count_outer_count: a limit above the row +# count leaves both grand-total rows +query I +SELECT COUNT(*) FROM (SELECT SUM(v1) FROM generate_series(10) AS t1(v1) WHERE false GROUP BY GROUPING SETS((v1), (), ()) LIMIT 5) +---- +2 + +# An empty Hive-partitioned file has no rows and exact partition-column statistics, the +# combination an outer MIN/MAX needs to be answered from statistics. +statement ok +COPY (SELECT * FROM (VALUES (1)) v(a) WHERE false) +TO 'test_files/scratch/grouping/p=x/empty.parquet' STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE hive_partitioned_empty (a INT, p VARCHAR) +STORED AS PARQUET PARTITIONED BY (p) +LOCATION 'test_files/scratch/grouping/'; + +# rollup_empty_input_grand_total_row: the single row ROLLUP emits holds NULL in the grouping column +query T +SELECT p FROM hive_partitioned_empty GROUP BY ROLLUP(p) +---- +NULL + +# rollup_empty_input_outer_min_max: the only row is the NULL grand-total row, so MIN/MAX are NULL +query TT +SELECT MIN(p), MAX(p) FROM (SELECT p FROM hive_partitioned_empty GROUP BY ROLLUP(p)) +---- +NULL NULL + +# cube_empty_input_outer_min_max: the only row is the NULL grand-total row, so MIN/MAX are NULL +query TT +SELECT MIN(p), MAX(p) FROM (SELECT p FROM hive_partitioned_empty GROUP BY CUBE(p)) +---- +NULL NULL + +# group_by_empty_input_outer_min_max: a plain GROUP BY emits no rows, so MIN/MAX are NULL +query TT +SELECT MIN(p), MAX(p) FROM (SELECT p FROM hive_partitioned_empty GROUP BY p) +---- +NULL NULL diff --git a/datafusion/sqllogictest/test_files/ident_normalization.slt b/datafusion/sqllogictest/test_files/ident_normalization.slt index b1bdb1d882274..5de84c69bd82f 100644 --- a/datafusion/sqllogictest/test_files/ident_normalization.slt +++ b/datafusion/sqllogictest/test_files/ident_normalization.slt @@ -75,7 +75,7 @@ A Int64 NO # Expect error as 'a' is not a column -- "A" is and the identifiers # are not normalized -query error DataFusion error: Schema error: No field named a\. Valid fields are x\."A"\. +query error DataFusion error: Schema error: No field named a\. Did you mean 'x\."A"'\?\nColumn names are case sensitive\. You can use double quotes to refer to the x\."A" column or disable the datafusion\.sql_parser\.enable_ident_normalization configuration\.\nValid fields are x\."A"\. select a from x; # should work (note the uppercase 'A') diff --git a/datafusion/sqllogictest/test_files/identifiers.slt b/datafusion/sqllogictest/test_files/identifiers.slt index e5eec3bf7f2c0..a78eba04c4843 100644 --- a/datafusion/sqllogictest/test_files/identifiers.slt +++ b/datafusion/sqllogictest/test_files/identifiers.slt @@ -90,16 +90,16 @@ drop table case_insensitive_test statement ok CREATE TABLE test("Column1" string) AS VALUES ('content1'); -statement error DataFusion error: Schema error: No field named column1\. Did you mean 'test\.Column1'\?\. +statement error DataFusion error: Schema error: No field named column1\. Did you mean 'test\."Column1"'\?\nColumn names are case sensitive\. You can use double quotes to refer to the test\."Column1" column or disable the datafusion\.sql_parser\.enable_ident_normalization configuration\.\nValid fields are test\."Column1"\. SELECT COLumn1 from test -statement error DataFusion error: Schema error: No field named column1\. Did you mean 'test\.Column1'\?\. +statement error DataFusion error: Schema error: No field named column1\. Did you mean 'test\."Column1"'\?\nColumn names are case sensitive\. You can use double quotes to refer to the test\."Column1" column or disable the datafusion\.sql_parser\.enable_ident_normalization configuration\.\nValid fields are test\."Column1"\. SELECT Column1 from test -statement error DataFusion error: Schema error: No field named column1\. Did you mean 'test\.Column1'\?\. +statement error DataFusion error: Schema error: No field named column1\. Did you mean 'test\."Column1"'\?\nColumn names are case sensitive\. You can use double quotes to refer to the test\."Column1" column or disable the datafusion\.sql_parser\.enable_ident_normalization configuration\.\nValid fields are test\."Column1"\. SELECT column1 from test -statement error DataFusion error: Schema error: No field named column1\. Did you mean 'test\.Column1'\?\. +statement error DataFusion error: Schema error: No field named column1\. Did you mean 'test\."Column1"'\?\nColumn names are case sensitive\. You can use double quotes to refer to the test\."Column1" column or disable the datafusion\.sql_parser\.enable_ident_normalization configuration\.\nValid fields are test\."Column1"\. SELECT "column1" from test statement ok diff --git a/datafusion/sqllogictest/test_files/in_list.slt b/datafusion/sqllogictest/test_files/in_list.slt new file mode 100644 index 0000000000000..dbdad2056fbd8 --- /dev/null +++ b/datafusion/sqllogictest/test_files/in_list.slt @@ -0,0 +1,581 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +# IN List Tests +# +# This file focuses on the IN operator and its various specializations +# +# Note that "short" IN LISTS do not go through the InList implementation at all, +# instead they are rewritten into a series of OR expressions. See: +# https://github.com/apache/datafusion/blob/ed37b6c9555bc278130dc774ed833b8c0bd29bfa/datafusion/optimizer/src/simplify_expressions/inlist_simplifier.rs#L39-L88 +########## + + +# Tests for IN LIST integer specializations + + +statement ok +CREATE TABLE in_list_ints ( + label VARCHAR, + i8 TINYINT, + u8 TINYINT UNSIGNED, + i16 SMALLINT, + u16 SMALLINT UNSIGNED, + i32 INT, + u32 INT UNSIGNED, + i64 BIGINT, + u64 BIGINT UNSIGNED +) AS VALUES + ('min', -128, 0, -32768, 0, -2147483648, 0, -9223372036854775808, 0), + ('minus_one', -1, 1, -1, 1, -1, 1, -1, 1), + ('zero', 0, 0, 0, 0, 0, 0, 0, 0), + ('one', 1, 1, 1, 1, 1, 1, 1, 1), + ('eleven', 11, 11, 11, 11, 11, 11, 11, 11), + ('max', 127, 255, 32767, 65535, 2147483647, 4294967295, 9223372036854775807, 18446744073709551615); + +# Verify that the Arrow types of the columns are as expected. This is important because the IN LIST specializations are based on the column types. +query TTTTTTTT +SELECT + arrow_typeof(i8), + arrow_typeof(u8), + arrow_typeof(i16), + arrow_typeof(u16), + arrow_typeof(i32), + arrow_typeof(u32), + arrow_typeof(i64), + arrow_typeof(u64) +FROM in_list_ints +LIMIT 1 +---- +Int8 UInt8 Int16 UInt16 Int32 UInt32 Int64 UInt64 + +# Verify the data is as expected. +query TIIIIIIII +SELECT label, i8, u8, i16, u16, i32, u32, i64, u64 +FROM in_list_ints +ORDER BY label +---- +eleven 11 11 11 11 11 11 11 11 +max 127 255 32767 65535 2147483647 4294967295 9223372036854775807 18446744073709551615 +min -128 0 -32768 0 -2147483648 0 -9223372036854775808 0 +minus_one -1 1 -1 1 -1 1 -1 1 +one 1 1 1 1 1 1 1 1 +zero 0 0 0 0 0 0 0 0 + +# Empty IN lists are rejected by the SQL parser. +statement error .*Expected: an expression, found: \).* +SELECT 1 IN (); + +# Min for each type +query TBBBBBBBB +SELECT + label, + i8 IN (1, 2, 3, -128), + u8 IN (1, 2, 3, 0), + i16 IN (1, 2, 3, -32768), + u16 IN (1, 2, 3, 0), + i32 IN (1, 2, 3, -2147483648), + u32 IN (1, 2, 3, 0), + i64 IN (1, 2, 3, -9223372036854775808), + u64 IN (1, 2, 3, 0) +FROM in_list_ints +ORDER BY label +---- +eleven false false false false false false false false +max false false false false false false false false +min true true true true true true true true +minus_one false true false true false true false true +one true true true true true true true true +zero false true false true false true false true + +# Max for each type (use values that cover the entire input range +# e.g. values that have 1 non zero byte. 2 non zero bytes, etc) +query TBBBBBBBB +SELECT + label, + i8 IN (-64, -32, 32, 64, 127), + u8 IN (32, 64, 128, 200, 255), + i16 IN (3, 258, 4097, 16385, 32767), + u16 IN (3, 258, 4097, 16385, 65535), + i32 IN (3, 258, 66051, 16909060, 2147483647), + u32 IN (3, 258, 66051, 16909060, 4294967295), + i64 IN (3, 258, 66051, 16909060, 9223372036854775807), + u64 IN (3, 258, 66051, 16909060, 18446744073709551615) +FROM in_list_ints +ORDER BY label +---- +eleven false false false false false false false false +max true true true true true true true true +min false false false false false false false false +minus_one false false false false false false false false +one false false false false false false false false +zero false false false false false false false false + +# Twelve item IN list with no matches (also cover the entire range) +query TBBBBBBBB +SELECT + label, + i8 IN (-120, -64, -32, -16, -8, -4, -2, 2, 32, 64, 100, 126), + u8 IN (2, 3, 4, 8, 16, 32, 64, 100, 128, 150, 200, 254), + i16 IN (-30000, -16384, -1024, -257, 2, 257, 4097, 8192, 16385, 20000, 30000, 32000), + u16 IN (2, 3, 4, 8, 16, 32, 257, 4097, 16385, 32768, 60000, 65534), + i32 IN (-2000000000, -1000000000, -16711936, -65536, -1024, 2, 66051, 16909060, 305419896, 1076895760, 2000000000, 2147483646), + u32 IN (2, 3, 4, 8, 16, 66051, 16909060, 305419896, 1076895760, 2309737967, 4000000000, 4294967294), + i64 IN (-9000000000000000000, -72057594037927936, -1000000000000000000, -65536, -1024, 2, 72623859790382856, 81985529216486895, 819855292164868960, 1234605616436508552, 9000000000000000000, 9223372036854775806), + u64 IN (2, 3, 4, 8, 16, 72623859790382856, 81985529216486895, 819855292164868960, 1234605616436508552, 18000000000000000000, 18364758544493064720, 18446744073709551614) +FROM in_list_ints +ORDER BY label +---- +eleven false false false false false false false false +max false false false false false false false false +min false false false false false false false false +minus_one false false false false false false false false +one false false false false false false false false +zero false false false false false false false false + +# Twelve item IN list with matches, including 11. +query TBBBBBBBB +SELECT + label, + i8 IN (-120, -64, -32, -16, -8, -4, -2, -128, 1, 3, 5, 11), + u8 IN (2, 4, 8, 16, 32, 64, 128, 200, 0, 11, 250, 255), + i16 IN (-30000, -20000, -10000, -32768, -1024, -256, -128, 1, 2, 3, 5, 11), + u16 IN (2, 4, 8, 16, 32, 64, 128, 256, 0, 11, 60000, 65535), + i32 IN (-2000000000, -1000000000, -2147483648, -65536, -1024, -256, -128, 1, 2, 3, 5, 11), + u32 IN (2, 4, 8, 16, 32, 64, 128, 256, 0, 11, 4000000000, 4294967295), + i64 IN (-9000000000000000000, -9223372036854775808, -1000000000, -65536, -1024, -256, -128, 1, 2, 3, 5, 11), + u64 IN (2, 4, 8, 16, 32, 64, 128, 256, 0, 11, 18000000000000000000, 18446744073709551615) +FROM in_list_ints +ORDER BY label +---- +eleven true true true true true true true true +max false true false true false true false true +min true true true true true true true true +minus_one false false false false false false false false +one true false true false true false true false +zero false true false true false true false true + +# Seventeen item IN list (shorter lists have specialized implementation) +query TBB +SELECT + label, + i8 IN (-128, -120, -100, -80, -60, -40, -20, -10, -5, -3, -2, 2, 3, 5, 20, 40, 11), + u8 IN (2, 3, 4, 5, 6, 7, 8, 9, 10, 12, 13, 14, 16, 20, 0, 11, 255) +FROM in_list_ints +ORDER BY label +---- +eleven true true +max false true +min true true +minus_one false false +one false false +zero false true + +# Cleanup +statement ok +DROP TABLE in_list_ints; + +#### +## Integer Null Handling +#### + +# Table with nulls to test null handling for integer IN list specializations +statement ok +CREATE TABLE in_list_ints_nullable ( + label VARCHAR, + i8 TINYINT, + u8 TINYINT UNSIGNED, + i16 SMALLINT, + u16 SMALLINT UNSIGNED, + i32 INT, + u32 INT UNSIGNED, + i64 BIGINT, + u64 BIGINT UNSIGNED +) AS VALUES + ('match', 11, 11, 11, 11, 11, 11, 11, 11), + ('no_match', 7, 7, 7, 7, 7, 7, 7, 7), + ('nulls', NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL); + +# Null input values return NULL when the IN list has no nulls. +query TBBBBBBBB +SELECT + label, + i8 IN (3, 4, 5, 6, 11), + u8 IN (3, 4, 5, 6, 11), + i16 IN (3, 258, 4097, 16385, 11), + u16 IN (3, 258, 4097, 16385, 11), + i32 IN (3, 258, 66051, 16909060, 11), + u32 IN (3, 258, 66051, 16909060, 11), + i64 IN (3, 258, 66051, 16909060, 11), + u64 IN (3, 258, 66051, 16909060, 11) +FROM in_list_ints_nullable +ORDER BY label +---- +match true true true true true true true true +no_match false false false false false false false false +nulls NULL NULL NULL NULL NULL NULL NULL NULL + +# NOT IN without NULL list values returns false for matches and true for non-matches. +query TBBBBBBBB +SELECT + label, + i8 NOT IN (3, 4, 5, 6, 11), + u8 NOT IN (3, 4, 5, 6, 11), + i16 NOT IN (3, 258, 4097, 16385, 11), + u16 NOT IN (3, 258, 4097, 16385, 11), + i32 NOT IN (3, 258, 66051, 16909060, 11), + u32 NOT IN (3, 258, 66051, 16909060, 11), + i64 NOT IN (3, 258, 66051, 16909060, 11), + u64 NOT IN (3, 258, 66051, 16909060, 11) +FROM in_list_ints_nullable +ORDER BY label +---- +match false false false false false false false false +no_match true true true true true true true true +nulls NULL NULL NULL NULL NULL NULL NULL NULL + +# Null IN list values return true for matches and NULL for non-matches. +query TBBBBBBBB +SELECT + label, + i8 IN (NULL, 3, 4, 5, 11), + u8 IN (NULL, 3, 4, 5, 11), + i16 IN (NULL, 3, 258, 4097, 11), + u16 IN (NULL, 3, 258, 4097, 11), + i32 IN (NULL, 3, 258, 66051, 11), + u32 IN (NULL, 3, 258, 66051, 11), + i64 IN (NULL, 3, 258, 66051, 11), + u64 IN (NULL, 3, 258, 66051, 11) +FROM in_list_ints_nullable +ORDER BY label +---- +match true true true true true true true true +no_match NULL NULL NULL NULL NULL NULL NULL NULL +nulls NULL NULL NULL NULL NULL NULL NULL NULL + +# Null IN list values return false for matches and NULL for non-matches with NOT IN. +query TBBBBBBBB +SELECT + label, + i8 NOT IN (NULL, 3, 4, 5, 11), + u8 NOT IN (NULL, 3, 4, 5, 11), + i16 NOT IN (NULL, 3, 258, 4097, 11), + u16 NOT IN (NULL, 3, 258, 4097, 11), + i32 NOT IN (NULL, 3, 258, 66051, 11), + u32 NOT IN (NULL, 3, 258, 66051, 11), + i64 NOT IN (NULL, 3, 258, 66051, 11), + u64 NOT IN (NULL, 3, 258, 66051, 11) +FROM in_list_ints_nullable +ORDER BY label +---- +match false false false false false false false false +no_match NULL NULL NULL NULL NULL NULL NULL NULL +nulls NULL NULL NULL NULL NULL NULL NULL NULL + +# Cleanup +statement ok +DROP TABLE in_list_ints_nullable + +#### +## Float IN List Specializations +#### + +statement ok +CREATE TABLE in_list_floats AS +SELECT + label, + arrow_cast(value, 'Float16') AS f16, + arrow_cast(value, 'Float32') AS f32, + arrow_cast(value, 'Float64') AS f64 +FROM (VALUES + ('match', 11.0), + ('no_match', 7.0), + ('nulls', NULL) +) AS t(label, value); + +# Five element IN lists cover the specialized Float16/32/64 paths. +query TTBBB +SELECT + 'Float16', + label, + f16 IN (arrow_cast(3.0, 'Float16'), arrow_cast(4.0, 'Float16'), arrow_cast(5.0, 'Float16'), arrow_cast(6.0, 'Float16'), arrow_cast(11.0, 'Float16')), + f16 IN (arrow_cast(NULL, 'Float16'), arrow_cast(3.0, 'Float16'), arrow_cast(4.0, 'Float16'), arrow_cast(5.0, 'Float16'), arrow_cast(11.0, 'Float16')), + f16 IN (arrow_cast(3.0, 'Float16'), arrow_cast(4.0, 'Float16'), arrow_cast(5.0, 'Float16'), arrow_cast(6.0, 'Float16'), arrow_cast(8.0, 'Float16')) +FROM in_list_floats +UNION ALL +SELECT + 'Float32', + label, + f32 IN (arrow_cast(3.0, 'Float32'), arrow_cast(4.0, 'Float32'), arrow_cast(5.0, 'Float32'), arrow_cast(6.0, 'Float32'), arrow_cast(11.0, 'Float32')), + f32 IN (arrow_cast(NULL, 'Float32'), arrow_cast(3.0, 'Float32'), arrow_cast(4.0, 'Float32'), arrow_cast(5.0, 'Float32'), arrow_cast(11.0, 'Float32')), + f32 IN (arrow_cast(3.0, 'Float32'), arrow_cast(4.0, 'Float32'), arrow_cast(5.0, 'Float32'), arrow_cast(6.0, 'Float32'), arrow_cast(8.0, 'Float32')) +FROM in_list_floats +UNION ALL +SELECT + 'Float64', + label, + f64 IN (arrow_cast(3.0, 'Float64'), arrow_cast(4.0, 'Float64'), arrow_cast(5.0, 'Float64'), arrow_cast(6.0, 'Float64'), arrow_cast(11.0, 'Float64')), + f64 IN (arrow_cast(NULL, 'Float64'), arrow_cast(3.0, 'Float64'), arrow_cast(4.0, 'Float64'), arrow_cast(5.0, 'Float64'), arrow_cast(11.0, 'Float64')), + f64 IN (arrow_cast(3.0, 'Float64'), arrow_cast(4.0, 'Float64'), arrow_cast(5.0, 'Float64'), arrow_cast(6.0, 'Float64'), arrow_cast(8.0, 'Float64')) +FROM in_list_floats +ORDER BY 1, 2 +---- +Float16 match true true false +Float16 no_match false NULL false +Float16 nulls NULL NULL NULL +Float32 match true true false +Float32 no_match false NULL false +Float32 nulls NULL NULL NULL +Float64 match true true false +Float64 no_match false NULL false +Float64 nulls NULL NULL NULL + +# Nine element Float16 IN list (shorter lists have specialized code) +query TB +SELECT + label, + f16 IN (arrow_cast(1.0, 'Float16'), arrow_cast(2.0, 'Float16'), arrow_cast(3.0, 'Float16'), + arrow_cast(4.0, 'Float16'), arrow_cast(5.0, 'Float16'), arrow_cast(6.0, 'Float16'), + arrow_cast(8.0, 'Float16'), arrow_cast(9.0, 'Float16'), arrow_cast(11.0, 'Float16')) +FROM in_list_floats +ORDER BY label +---- +match true +no_match false +nulls NULL + +# Cleanup +statement ok +DROP TABLE in_list_floats + +#### +## Temporal IN List Specializations +#### + +statement ok +CREATE TABLE in_list_temporal AS +SELECT + label, + arrow_cast(value, 'Date32') AS d32, + arrow_cast(value, 'Date64') AS d64, + arrow_cast(arrow_cast(value, 'Int32'), 'Time32(Second)') AS t32s, + arrow_cast(value, 'Time64(Nanosecond)') AS t64ns, + arrow_cast(value, 'Timestamp(Nanosecond, None)') AS ts_ns, + arrow_cast(value, 'Timestamp(Second, Some("UTC"))') AS ts_s_utc, + arrow_cast(value, 'Duration(Second)') AS dur_s +FROM (VALUES + ('match', 11), + ('no_match', 7), + ('nulls', NULL) +) AS t(label, value); + +# Basic Temporal IN Lists +query TBBBBBBB +SELECT + label, + d32 IN (arrow_cast(3, 'Date32'), arrow_cast(4, 'Date32'), arrow_cast(5, 'Date32'), arrow_cast(11, 'Date32')), + d64 IN (arrow_cast(3, 'Date64'), arrow_cast(4, 'Date64'), arrow_cast(5, 'Date64'), arrow_cast(11, 'Date64')), + t32s IN (arrow_cast(arrow_cast(3, 'Int32'), 'Time32(Second)'), arrow_cast(arrow_cast(4, 'Int32'), 'Time32(Second)'), arrow_cast(arrow_cast(5, 'Int32'), 'Time32(Second)'), arrow_cast(arrow_cast(11, 'Int32'), 'Time32(Second)')), + t64ns IN (arrow_cast(3, 'Time64(Nanosecond)'), arrow_cast(4, 'Time64(Nanosecond)'), arrow_cast(5, 'Time64(Nanosecond)'), arrow_cast(11, 'Time64(Nanosecond)')), + ts_ns IN (arrow_cast(3, 'Timestamp(Nanosecond, None)'), arrow_cast(4, 'Timestamp(Nanosecond, None)'), arrow_cast(5, 'Timestamp(Nanosecond, None)'), arrow_cast(11, 'Timestamp(Nanosecond, None)')), + ts_s_utc IN (arrow_cast(3, 'Timestamp(Second, Some("UTC"))'), arrow_cast(4, 'Timestamp(Second, Some("UTC"))'), arrow_cast(5, 'Timestamp(Second, Some("UTC"))'), arrow_cast(11, 'Timestamp(Second, Some("UTC"))')), + dur_s IN (arrow_cast(3, 'Duration(Second)'), arrow_cast(4, 'Duration(Second)'), arrow_cast(5, 'Duration(Second)'), arrow_cast(11, 'Duration(Second)')) +FROM in_list_temporal +ORDER BY label +---- +match true true true true true true true +no_match false false false false false false false +nulls NULL NULL NULL NULL NULL NULL NULL + +# The same lists with NOT IN. +query TBBBBBBB +SELECT + label, + d32 NOT IN (arrow_cast(3, 'Date32'), arrow_cast(4, 'Date32'), arrow_cast(5, 'Date32'), arrow_cast(11, 'Date32')), + d64 NOT IN (arrow_cast(3, 'Date64'), arrow_cast(4, 'Date64'), arrow_cast(5, 'Date64'), arrow_cast(11, 'Date64')), + t32s NOT IN (arrow_cast(arrow_cast(3, 'Int32'), 'Time32(Second)'), arrow_cast(arrow_cast(4, 'Int32'), 'Time32(Second)'), arrow_cast(arrow_cast(5, 'Int32'), 'Time32(Second)'), arrow_cast(arrow_cast(11, 'Int32'), 'Time32(Second)')), + t64ns NOT IN (arrow_cast(3, 'Time64(Nanosecond)'), arrow_cast(4, 'Time64(Nanosecond)'), arrow_cast(5, 'Time64(Nanosecond)'), arrow_cast(11, 'Time64(Nanosecond)')), + ts_ns NOT IN (arrow_cast(3, 'Timestamp(Nanosecond, None)'), arrow_cast(4, 'Timestamp(Nanosecond, None)'), arrow_cast(5, 'Timestamp(Nanosecond, None)'), arrow_cast(11, 'Timestamp(Nanosecond, None)')), + ts_s_utc NOT IN (arrow_cast(3, 'Timestamp(Second, Some("UTC"))'), arrow_cast(4, 'Timestamp(Second, Some("UTC"))'), arrow_cast(5, 'Timestamp(Second, Some("UTC"))'), arrow_cast(11, 'Timestamp(Second, Some("UTC"))')), + dur_s NOT IN (arrow_cast(3, 'Duration(Second)'), arrow_cast(4, 'Duration(Second)'), arrow_cast(5, 'Duration(Second)'), arrow_cast(11, 'Duration(Second)')) +FROM in_list_temporal +ORDER BY label +---- +match false false false false false false false +no_match true true true true true true true +nulls NULL NULL NULL NULL NULL NULL NULL + +# Null IN list values return true for matches and NULL for non-matches. +query TBBBBBBB +SELECT + label, + d32 IN (NULL, arrow_cast(3, 'Date32'), arrow_cast(4, 'Date32'), arrow_cast(11, 'Date32')), + d64 IN (NULL, arrow_cast(3, 'Date64'), arrow_cast(4, 'Date64'), arrow_cast(11, 'Date64')), + t32s IN (NULL, arrow_cast(arrow_cast(3, 'Int32'), 'Time32(Second)'), arrow_cast(arrow_cast(4, 'Int32'), 'Time32(Second)'), arrow_cast(arrow_cast(11, 'Int32'), 'Time32(Second)')), + t64ns IN (NULL, arrow_cast(3, 'Time64(Nanosecond)'), arrow_cast(4, 'Time64(Nanosecond)'), arrow_cast(11, 'Time64(Nanosecond)')), + ts_ns IN (NULL, arrow_cast(3, 'Timestamp(Nanosecond, None)'), arrow_cast(4, 'Timestamp(Nanosecond, None)'), arrow_cast(11, 'Timestamp(Nanosecond, None)')), + ts_s_utc IN (NULL, arrow_cast(3, 'Timestamp(Second, Some("UTC"))'), arrow_cast(4, 'Timestamp(Second, Some("UTC"))'), arrow_cast(11, 'Timestamp(Second, Some("UTC"))')), + dur_s IN (NULL, arrow_cast(3, 'Duration(Second)'), arrow_cast(4, 'Duration(Second)'), arrow_cast(11, 'Duration(Second)')) +FROM in_list_temporal +ORDER BY label +---- +match true true true true true true true +no_match NULL NULL NULL NULL NULL NULL NULL +nulls NULL NULL NULL NULL NULL NULL NULL + +# A NULL in the list turns off some specializations +query TBBBBBBB +SELECT + label, + d32 IN (NULL, arrow_cast(11, 'Date32')), + d64 IN (NULL, arrow_cast(11, 'Date64')), + t32s IN (NULL, arrow_cast(arrow_cast(11, 'Int32'), 'Time32(Second)')), + t64ns IN (NULL, arrow_cast(11, 'Time64(Nanosecond)')), + ts_ns IN (NULL, arrow_cast(11, 'Timestamp(Nanosecond, None)')), + ts_s_utc IN (NULL, arrow_cast(11, 'Timestamp(Second, Some("UTC"))')), + dur_s IN (NULL, arrow_cast(11, 'Duration(Second)')) +FROM in_list_temporal +ORDER BY label +---- +match true true true true true true true +no_match NULL NULL NULL NULL NULL NULL NULL +nulls NULL NULL NULL NULL NULL NULL NULL + +# Cleanup +statement ok +DROP TABLE in_list_temporal + +#### +## Decimal128 IN List Specializations +#### + +statement ok +CREATE TABLE in_list_decimal AS +SELECT * FROM (VALUES + ('match', arrow_cast(11, 'Decimal128(10, 2)')), + ('no_match', arrow_cast(7, 'Decimal128(10, 2)')), + ('nulls', NULL) +) AS t(label, d128); + +query T +SELECT arrow_typeof(d128) FROM in_list_decimal LIMIT 1 +---- +Decimal128(10, 2) + +# Four non-null values and five non-null values (test different specializations) +query TBB +SELECT + label, + d128 IN (arrow_cast(3, 'Decimal128(10, 2)'), arrow_cast(4, 'Decimal128(10, 2)'), arrow_cast(5, 'Decimal128(10, 2)'), arrow_cast(11, 'Decimal128(10, 2)')), + d128 IN (arrow_cast(3, 'Decimal128(10, 2)'), arrow_cast(4, 'Decimal128(10, 2)'), arrow_cast(5, 'Decimal128(10, 2)'), arrow_cast(11, 'Decimal128(10, 2)'), arrow_cast(13, 'Decimal128(10, 2)')) +FROM in_list_decimal +ORDER BY label +---- +match true true +no_match false false +nulls NULL NULL + +# The same lists with NOT IN. +query TBB +SELECT + label, + d128 NOT IN (arrow_cast(3, 'Decimal128(10, 2)'), arrow_cast(4, 'Decimal128(10, 2)'), arrow_cast(5, 'Decimal128(10, 2)'), arrow_cast(11, 'Decimal128(10, 2)')), + d128 NOT IN (arrow_cast(3, 'Decimal128(10, 2)'), arrow_cast(4, 'Decimal128(10, 2)'), arrow_cast(5, 'Decimal128(10, 2)'), arrow_cast(11, 'Decimal128(10, 2)'), arrow_cast(13, 'Decimal128(10, 2)')) +FROM in_list_decimal +ORDER BY label +---- +match false false +no_match true true +nulls NULL NULL + +# Null IN list values, including short lists with a single non-null value. +query TBB +SELECT + label, + d128 IN (NULL, arrow_cast(3, 'Decimal128(10, 2)'), arrow_cast(4, 'Decimal128(10, 2)'), arrow_cast(11, 'Decimal128(10, 2)')), + d128 IN (NULL, arrow_cast(11, 'Decimal128(10, 2)')) +FROM in_list_decimal +ORDER BY label +---- +match true true +no_match NULL NULL +nulls NULL NULL + +# Cleanup +statement ok +DROP TABLE in_list_decimal + +#### +## Interval IN List Specializations +#### + +statement ok +CREATE TABLE in_list_interval AS +SELECT * FROM (VALUES + ('match', INTERVAL '11 months'), + ('no_match', INTERVAL '7 months'), + ('nulls', NULL) +) AS t(label, imdn); + +query T +SELECT arrow_typeof(imdn) FROM in_list_interval LIMIT 1 +---- +Interval(MonthDayNano) + +# Four non-null values and five non-null values (test different specializations) +query TBB +SELECT + label, + imdn IN (INTERVAL '3 months', INTERVAL '4 months', INTERVAL '5 months', INTERVAL '11 months'), + imdn IN (INTERVAL '3 months', INTERVAL '4 months', INTERVAL '5 months', INTERVAL '11 months', INTERVAL '13 months') +FROM in_list_interval +ORDER BY label +---- +match true true +no_match false false +nulls NULL NULL + +# The same lists with NOT IN. +query TBB +SELECT + label, + imdn NOT IN (INTERVAL '3 months', INTERVAL '4 months', INTERVAL '5 months', INTERVAL '11 months'), + imdn NOT IN (INTERVAL '3 months', INTERVAL '4 months', INTERVAL '5 months', INTERVAL '11 months', INTERVAL '13 months') +FROM in_list_interval +ORDER BY label +---- +match false false +no_match true true +nulls NULL NULL + +# Null IN list values, including short lists with a single non-null value. +query TBB +SELECT + label, + imdn IN (NULL, INTERVAL '3 months', INTERVAL '4 months', INTERVAL '11 months'), + imdn IN (NULL, INTERVAL '11 months') +FROM in_list_interval +ORDER BY label +---- +match true true +no_match NULL NULL +nulls NULL NULL + +# Cleanup +statement ok +DROP TABLE in_list_interval diff --git a/datafusion/sqllogictest/test_files/information_schema.slt b/datafusion/sqllogictest/test_files/information_schema.slt index 1aa9bc79e5bbe..573fb04b3451b 100644 --- a/datafusion/sqllogictest/test_files/information_schema.slt +++ b/datafusion/sqllogictest/test_files/information_schema.slt @@ -218,6 +218,8 @@ datafusion.execution.batch_size 8192 datafusion.execution.coalesce_batches true datafusion.execution.collect_statistics true datafusion.execution.enable_ansi_mode false +datafusion.execution.enable_file_stream_work_stealing true +datafusion.execution.enable_migration_aggregate true datafusion.execution.enable_recursive_ctes true datafusion.execution.enforce_batch_size_in_joins false datafusion.execution.hash_join_buffering_capacity 0 @@ -239,6 +241,10 @@ datafusion.execution.parquet.coerce_int96 NULL datafusion.execution.parquet.coerce_int96_tz NULL datafusion.execution.parquet.column_index_truncate_length 64 datafusion.execution.parquet.compression zstd(3) +datafusion.execution.parquet.content_defined_chunking.enabled false +datafusion.execution.parquet.content_defined_chunking.max_chunk_size 1048576 +datafusion.execution.parquet.content_defined_chunking.min_chunk_size 262144 +datafusion.execution.parquet.content_defined_chunking.norm_level 0 datafusion.execution.parquet.created_by datafusion datafusion.execution.parquet.data_page_row_count_limit 20000 datafusion.execution.parquet.data_pagesize_limit 1048576 @@ -247,7 +253,9 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 datafusion.execution.parquet.enable_page_index true datafusion.execution.parquet.encoding NULL datafusion.execution.parquet.force_filter_selections false +datafusion.execution.parquet.max_in_list_size 20 datafusion.execution.parquet.max_predicate_cache_size NULL +datafusion.execution.parquet.max_row_group_bytes NULL datafusion.execution.parquet.max_row_group_size 1048576 datafusion.execution.parquet.maximum_buffered_record_batches_per_stream 2 datafusion.execution.parquet.maximum_parallel_row_group_writers 1 @@ -260,7 +268,6 @@ datafusion.execution.parquet.skip_arrow_metadata false datafusion.execution.parquet.skip_metadata true datafusion.execution.parquet.statistics_enabled page datafusion.execution.parquet.statistics_truncate_length 64 -datafusion.execution.parquet.use_content_defined_chunking NULL datafusion.execution.parquet.write_batch_size 1024 datafusion.execution.parquet.writer_version 1.0 datafusion.execution.perfect_hash_join_min_key_density 0.15 @@ -303,6 +310,7 @@ datafusion.optimizer.enable_distinct_aggregation_soft_limit true datafusion.optimizer.enable_dynamic_filter_pushdown true datafusion.optimizer.enable_join_dynamic_filter_pushdown true datafusion.optimizer.enable_leaf_expression_pushdown true +datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery true datafusion.optimizer.enable_piecewise_merge_join false datafusion.optimizer.enable_round_robin_repartition true datafusion.optimizer.enable_sort_pushdown true @@ -325,7 +333,7 @@ datafusion.optimizer.prefer_existing_union false datafusion.optimizer.prefer_hash_join true datafusion.optimizer.preserve_file_partitions 0 datafusion.optimizer.repartition_aggregations true -datafusion.optimizer.repartition_file_min_size 10485760 +datafusion.optimizer.repartition_file_min_size 1048576 datafusion.optimizer.repartition_file_scans true datafusion.optimizer.repartition_joins true datafusion.optimizer.repartition_sorts true @@ -337,10 +345,12 @@ datafusion.optimizer.use_statistics_registry false datafusion.runtime.file_statistics_cache_limit 20M datafusion.runtime.list_files_cache_limit 1M datafusion.runtime.list_files_cache_ttl NULL +datafusion.runtime.max_spill_merge_fan_in 0 datafusion.runtime.max_temp_directory_size 100G datafusion.runtime.memory_limit unlimited datafusion.runtime.metadata_cache_limit 50M datafusion.runtime.temp_directory NULL +datafusion.spark.map_key_dedup_policy EXCEPTION datafusion.sql_parser.collect_spans false datafusion.sql_parser.default_null_ordering nulls_max datafusion.sql_parser.dialect generic @@ -366,15 +376,17 @@ datafusion.catalog.location NULL Location scanned to load tables for `default` s datafusion.catalog.newlines_in_values false Specifies whether newlines in (quoted) CSV values are supported. This is the default value for `format.newlines_in_values` for `CREATE EXTERNAL TABLE` if not specified explicitly in the statement. Parsing newlines in quoted values may be affected by execution behaviour such as parallel file scanning. Setting this to `true` ensures that newlines in values are parsed successfully, which may reduce performance. datafusion.execution.batch_size 8192 Default batch size while creating new batches, it's especially useful for buffer-in-memory batches since creating tiny batches would result in too much metadata memory consumption datafusion.execution.coalesce_batches true When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting -datafusion.execution.collect_statistics true Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Applies to the default `ListingTableProvider` in DataFusion. Defaults to true. +datafusion.execution.collect_statistics true Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Defaults to true. datafusion.execution.enable_ansi_mode false Whether to enable ANSI SQL mode. The flag is experimental and relevant only for DataFusion Spark built-in functions When `enable_ansi_mode` is set to `true`, the query engine follows ANSI SQL semantics for expressions, casting, and error handling. This means: - **Strict type coercion rules:** implicit casts between incompatible types are disallowed. - **Standard SQL arithmetic behavior:** operations such as division by zero, numeric overflow, or invalid casts raise runtime errors rather than returning `NULL` or adjusted values. - **Consistent ANSI behavior** for string concatenation, comparisons, and `NULL` handling. When `enable_ansi_mode` is `false` (the default), the engine uses a more permissive, non-ANSI mode designed for user convenience and backward compatibility. In this mode: - Implicit casts between types are allowed (e.g., string to integer when possible). - Arithmetic operations are more lenient — for example, `abs()` on the minimum representable integer value returns the input value instead of raising overflow. - Division by zero or invalid casts may return `NULL` instead of failing. # Default `false` — ANSI SQL mode is disabled by default. +datafusion.execution.enable_file_stream_work_stealing true When `true` (the default), DataFusion's built-in file scans dynamically rebalance files across partitions at query execution time: a partition that goes idle reads files (or byte-range morsels) originally assigned to a sibling partition, which keeps all partitions busy in a single process. Executors that depend on the plan-time partition assignment — such as Ballista and datafusion-distributed, which run each partition as an isolated task and never poll the siblings — should set this to `false` so each partition reads only its own file group and no runtime reassignment occurs. +datafusion.execution.enable_migration_aggregate true Temporary switch for aggregate stream implementations that are being migrated from `GroupedHashAggregateStream`. When set to true, DataFusion tries the migrated implementations when their preconditions are satisfied. When set to false, grouped aggregation falls back to `GroupedHashAggregateStream`. This option will be removed after the migration is finished. See for details. datafusion.execution.enable_recursive_ctes true Should DataFusion support recursive CTEs datafusion.execution.enforce_batch_size_in_joins false Should DataFusion enforce batch size in joins or not. By default, DataFusion will not enforce batch size in joins. Enforcing batch size in joins can reduce memory usage when joining large tables with a highly-selective join filter, but is also slightly slower. datafusion.execution.hash_join_buffering_capacity 0 How many bytes to buffer in the probe side of hash joins while the build side is concurrently being built. Without this, hash joins will wait until the full materialization of the build side before polling the probe side. This is useful in scenarios where the query is not completely CPU bounded, allowing to do some early work concurrently and reducing the latency of the query. Note that when hash join buffering is enabled, the probe side will start eagerly polling data, not giving time for the producer side of dynamic filters to produce any meaningful predicate. Queries with dynamic filters might see performance degradation. Disabled by default, set to a number greater than 0 for enabling it. datafusion.execution.keep_partition_by_columns false Should DataFusion keep the columns used for partition_by in the output RecordBatches datafusion.execution.listing_table_factory_infer_partitions true Should a `ListingTable` created through the `ListingTableFactory` infer table partitions from Hive compliant directories. Defaults to true (partition columns are inferred and will be represented in the table schema). datafusion.execution.listing_table_ignore_subdirectory true Should sub directories be ignored when scanning directories for data files. Defaults to true (ignores subdirectories), consistent with Hive. Note that this setting does not affect reading partitioned tables (e.g. `/table/year=2021/month=01/data.parquet`). -datafusion.execution.max_buffered_batches_per_output_file 2 This is the maximum number of RecordBatches buffered for each output file being worked. Higher values can potentially give faster write performance at the cost of higher peak memory consumption +datafusion.execution.max_buffered_batches_per_output_file 2 This is the maximum number of RecordBatches buffered for each output file being worked. Higher values can potentially give faster write performance at the cost of higher peak memory consumption. This budget is split evenly between two independent points in the write pipeline (see the demuxer diagram in #7791): how many files can be in flight from the demuxer to a writer task, and how many RecordBatches are buffered for a single file's writer. Must be at least 2 so each half gets at least 1 unit of buffering - 0 or 1 would leave one side with a zero-capacity channel and panic at write time. datafusion.execution.max_spill_file_size_bytes 134217728 Maximum size in bytes for individual spill files before rotating to a new file. When operators spill data to disk (e.g., RepartitionExec), they write multiple batches to the same file until this size limit is reached, then rotate to a new file. This reduces syscall overhead compared to one-file-per-batch while preventing files from growing too large. A larger value reduces file creation overhead but may hold more disk space. A smaller value creates more files but allows finer-grained space reclamation as files can be deleted once fully consumed. Now only `RepartitionExec` supports this spill file rotation feature, other spilling operators may create spill files larger than the limit. Default: 128 MB datafusion.execution.meta_fetch_concurrency 32 Number of files to read in parallel when inferring schema and statistics datafusion.execution.minimum_parallel_output_files 4 Guarantees a minimum level of output files running in parallel. RecordBatches will be distributed in round robin fashion to each parallel writer. Each writer is closed and a new file opened once soft_max_rows_per_output_file is reached. @@ -389,6 +401,10 @@ datafusion.execution.parquet.coerce_int96 NULL (reading) If true, parquet reader datafusion.execution.parquet.coerce_int96_tz NULL (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. datafusion.execution.parquet.column_index_truncate_length 64 (writing) Sets column index truncate length datafusion.execution.parquet.compression zstd(3) (writing) Sets default parquet compression codec. Valid values are: uncompressed, snappy, gzip(level), brotli(level), lz4, zstd(level), and lz4_raw. These values are not case sensitive. If NULL, uses default parquet writer setting Note that this default setting is not the same as the default parquet writer setting. +datafusion.execution.parquet.content_defined_chunking.enabled false (writing) EXPERIMENTAL: Enable content-defined chunking (CDC) when writing parquet files. When enabled, parallel writing is automatically disabled since the chunker state must persist across row groups. +datafusion.execution.parquet.content_defined_chunking.max_chunk_size 1048576 Maximum chunk size in bytes. A split is forced when the accumulated size exceeds this value. Default is 1 MiB. +datafusion.execution.parquet.content_defined_chunking.min_chunk_size 262144 Minimum chunk size in bytes. The rolling hash will not trigger a split until this many bytes have been accumulated. Default is 256 KiB. +datafusion.execution.parquet.content_defined_chunking.norm_level 0 Normalization level. Increasing this improves deduplication ratio but increases fragmentation. Recommended range is [-3, 3], default is 0. datafusion.execution.parquet.created_by datafusion (writing) Sets "created by" property datafusion.execution.parquet.data_page_row_count_limit 20000 (writing) Sets best effort maximum number of rows in data page datafusion.execution.parquet.data_pagesize_limit 1048576 (writing) Sets best effort maximum size of data page in bytes @@ -397,8 +413,10 @@ datafusion.execution.parquet.dictionary_page_size_limit 1048576 (writing) Sets b datafusion.execution.parquet.enable_page_index true (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. datafusion.execution.parquet.encoding NULL (writing) Sets default encoding for any column. Valid values are: plain, plain_dictionary, rle, bit_packed, delta_binary_packed, delta_length_byte_array, delta_byte_array, rle_dictionary, and byte_stream_split. These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.force_filter_selections false (reading) Force the use of RowSelections for filter results, when pushdown_filters is enabled. If false, the reader will automatically choose between a RowSelection and a Bitmap based on the number and pattern of selected rows. +datafusion.execution.parquet.max_in_list_size 20 Maximum number of values in an `IN (...)` list for which pruning will occur. Longer lists will not be used to prune files, row groups, or data pages. Higher values help in cases such as filtering on a list of ~25-100 identifiers, but also make the predicate more expensive to evaluate. Set to 0 to disable `IN (...)` list pruning entirely. Defaults to 20. datafusion.execution.parquet.max_predicate_cache_size NULL (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. -datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. +datafusion.execution.parquet.max_row_group_bytes NULL (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. +datafusion.execution.parquet.max_row_group_size 1048576 (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. datafusion.execution.parquet.maximum_buffered_record_batches_per_stream 2 (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. datafusion.execution.parquet.maximum_parallel_row_group_writers 1 (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. datafusion.execution.parquet.metadata_size_hint 524288 (reading) If specified, the parquet reader will try and fetch the last `size_hint` bytes of the parquet file optimistically. If not specified, two reads are required: One read to fetch the 8-byte parquet footer and another to fetch the metadata length encoded in the footer Default setting to 512 KiB, which should be sufficient for most parquet files, it can reduce one I/O operation per parquet file. If the metadata is larger than the hint, two reads will still be performed. @@ -410,7 +428,6 @@ datafusion.execution.parquet.skip_arrow_metadata false (writing) Skip encoding t datafusion.execution.parquet.skip_metadata true (reading) If true, the parquet reader skip the optional embedded metadata that may be in the file Schema. This setting can help avoid schema conflicts when querying multiple parquet files with schemas containing compatible types but different metadata datafusion.execution.parquet.statistics_enabled page (writing) Sets if statistics are enabled for any column Valid values are: "none", "chunk", and "page" These values are not case sensitive. If NULL, uses default parquet writer setting datafusion.execution.parquet.statistics_truncate_length 64 (writing) Sets statistics truncate length. If NULL, uses default parquet writer setting -datafusion.execution.parquet.use_content_defined_chunking NULL (writing) EXPERIMENTAL: Enable content-defined chunking (CDC) when writing parquet files. When `Some`, CDC is enabled with the given options; when `None` (the default), CDC is disabled. When CDC is enabled, parallel writing is automatically disabled since the chunker state must persist across row groups. datafusion.execution.parquet.write_batch_size 1024 (writing) Sets write_batch_size in rows datafusion.execution.parquet.writer_version 1.0 (writing) Sets parquet writer version valid values are "1.0" and "2.0" datafusion.execution.perfect_hash_join_min_key_density 0.15 The minimum required density of join keys on the build side to consider a perfect hash join (see `HashJoinExec` for more details). Density is calculated as: `(number of rows) / (max_key - min_key + 1)`. A perfect hash join may be used if the actual key density > this value. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. @@ -453,6 +470,7 @@ datafusion.optimizer.enable_distinct_aggregation_soft_limit true When set to tru datafusion.optimizer.enable_dynamic_filter_pushdown true When set to true attempts to push down dynamic filters generated by operators (TopK, Join & Aggregate) into the file scan phase. For example, for a query such as `SELECT * FROM t ORDER BY timestamp DESC LIMIT 10`, the optimizer will attempt to push down the current top 10 timestamps that the TopK operator references into the file scans. This means that if we already have 10 timestamps in the year 2025 any files that only have timestamps in the year 2024 can be skipped / pruned at various stages in the scan. The config will suppress `enable_join_dynamic_filter_pushdown`, `enable_topk_dynamic_filter_pushdown` & `enable_aggregate_dynamic_filter_pushdown` So if you disable `enable_topk_dynamic_filter_pushdown`, then enable `enable_dynamic_filter_pushdown`, the `enable_topk_dynamic_filter_pushdown` will be overridden. datafusion.optimizer.enable_join_dynamic_filter_pushdown true When set to true, the optimizer will attempt to push down Join dynamic filters into the file scan phase. datafusion.optimizer.enable_leaf_expression_pushdown true When set to true, the optimizer will extract leaf expressions (such as `get_field`) from filter/sort/join nodes into projections closer to the leaf table scans, and push those projections down towards the leaf nodes. +datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery true When set to true, uncorrelated scalar subqueries are left in the logical plan and executed by `ScalarSubqueryExec` during physical execution. When set to false, all scalar subqueries (including uncorrelated ones) are rewritten to left joins by the `ScalarSubqueryToJoin` optimizer rule. Note disabling this option is not recommended. It restores pre behavior, which silently produces incorrect results for multi-row subqueries and does not support scalar subqueries in ORDER BY / JOIN ON / aggregate-function arguments. This option is intended as a temporary escape hatch for distributed execution frameworks and is planned to be removed in a future DataFusion release. datafusion.optimizer.enable_piecewise_merge_join false When set to true, piecewise merge join is enabled. PiecewiseMergeJoin is currently experimental. Physical planner will opt for PiecewiseMergeJoin when there is only one range filter. datafusion.optimizer.enable_round_robin_repartition true When set to true, the physical plan optimizer will try to add round robin repartitioning to increase parallelism to leverage more CPU cores datafusion.optimizer.enable_sort_pushdown true Enable sort pushdown optimization. When enabled, attempts to push sort requirements down to data sources that can natively handle them (e.g., by reversing file/row group read order). Returns **inexact ordering**: Sort operator is kept for correctness, but optimized input enables early termination for TopK queries (ORDER BY ... LIMIT N), providing significant speedup. Memory: No additional overhead (only changes read order). Future: Will add option to detect perfectly sorted data and eliminate Sort completely. Default: true @@ -473,9 +491,9 @@ datafusion.optimizer.max_passes 3 Number of times that the optimizer will attemp datafusion.optimizer.prefer_existing_sort false When true, DataFusion will opportunistically remove sorts when the data is already sorted, (i.e. setting `preserve_order` to true on `RepartitionExec` and using `SortPreservingMergeExec`) When false, DataFusion will maximize plan parallelism using `RepartitionExec` even if this requires subsequently resorting data using a `SortExec`. datafusion.optimizer.prefer_existing_union false When set to true, the optimizer will not attempt to convert Union to Interleave datafusion.optimizer.prefer_hash_join true When set to true, the physical plan optimizer will prefer HashJoin over SortMergeJoin. HashJoin can work more efficiently than SortMergeJoin but consumes more memory -datafusion.optimizer.preserve_file_partitions 0 Minimum number of distinct partition values required to group files by their Hive partition column values (enabling Hash partitioning declaration). How the option is used: - preserve_file_partitions=0: Disable it. - preserve_file_partitions=1: Always enable it. - preserve_file_partitions=N, actual file partitions=M: Only enable when M >= N. This threshold preserves I/O parallelism when file partitioning is below it. Note: This may reduce parallelism, rooting from the I/O level, if the number of distinct partitions is less than the target_partitions. +datafusion.optimizer.preserve_file_partitions 0 Minimum number of distinct partition values required to group files by their Hive partition column values (enabling output partitioning declaration). How the option is used: - preserve_file_partitions=0: Disable it. - preserve_file_partitions=1: Always enable it. - preserve_file_partitions=N, actual file partitions=M: Only enable when M >= N. This threshold preserves I/O parallelism when file partitioning is below it. Note: This may reduce parallelism, rooting from the I/O level, if the number of distinct partitions is less than the target_partitions. datafusion.optimizer.repartition_aggregations true Should DataFusion repartition data using the aggregate keys to execute aggregates in parallel using the provided `target_partitions` level -datafusion.optimizer.repartition_file_min_size 10485760 Minimum total files size in bytes to perform file scan repartitioning. +datafusion.optimizer.repartition_file_min_size 1048576 Minimum total file size in bytes for file-group byte-range splitting to fire. Files (or merged file groups) smaller than this stay as one partition. Lower values produce more, smaller partitions — better at filling `target_partitions` worth of cores when files are modestly sized, at the cost of slightly more per-partition open / metadata-load overhead. datafusion.optimizer.repartition_file_scans true When set to `true`, datasource partitions will be repartitioned to achieve maximum parallelism. This applies to both in-memory partitions and FileSource's file groups (1 group is 1 partition). For FileSources, only Parquet and CSV formats are currently supported. If set to `true` for a FileSource, all files will be repartitioned evenly (i.e., a single large file might be partitioned into smaller chunks) for parallel scanning. If set to `false` for a FileSource, different files will be read in parallel, but repartitioning won't happen within a single file. If set to `true` for an in-memory source, all memtable's partitions will have their batches repartitioned evenly to the desired number of `target_partitions`. Repartitioning can change the total number of partitions and batches per partition, but does not slice the initial record tables provided to the MemTable on creation. datafusion.optimizer.repartition_joins true Should DataFusion repartition data using the join keys to execute joins in parallel using the provided `target_partitions` level datafusion.optimizer.repartition_sorts true Should DataFusion execute sorts in a per-partition fashion and merge afterwards instead of coalescing first and sorting globally. With this flag is enabled, plans in the form below ```text "SortExec: [a@0 ASC]", " CoalescePartitionsExec", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ``` would turn into the plan below which performs better in multithreaded environments ```text "SortPreservingMergeExec: [a@0 ASC]", " SortExec: [a@0 ASC]", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ``` @@ -487,13 +505,15 @@ datafusion.optimizer.use_statistics_registry false When set to true, the physica datafusion.runtime.file_statistics_cache_limit 20M Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.list_files_cache_limit 1M Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.list_files_cache_ttl NULL TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes. +datafusion.runtime.max_spill_merge_fan_in 0 Maximum number of spill files opened by one external merge pass. Use 0 for unlimited. Values below 2 still use 2 so a merge can make progress. datafusion.runtime.max_temp_directory_size 100G Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.memory_limit unlimited Maximum memory limit for query execution. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.metadata_cache_limit 50M Maximum memory to use for file metadata cache such as Parquet metadata. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. datafusion.runtime.temp_directory NULL The path to the temporary file directory. +datafusion.spark.map_key_dedup_policy EXCEPTION Policy for handling duplicate keys in Spark-compatible map-construction functions (`map_from_arrays`, `map_from_entries`, `str_to_map`). Mirrors Spark's [`spark.sql.mapKeyDedupPolicy`](https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961): - `EXCEPTION` (default): raise `[DUPLICATED_MAP_KEY]` at runtime on any duplicate key. - `LAST_WIN`: keep the last occurrence of each duplicate key. Values are case-insensitive. datafusion.sql_parser.collect_spans false When set to true, the source locations relative to the original SQL query (i.e. [`Span`](https://docs.rs/sqlparser/latest/sqlparser/tokenizer/struct.Span.html)) will be collected and recorded in the logical plan nodes. datafusion.sql_parser.default_null_ordering nulls_max Specifies the default null ordering for query results. There are 4 options: - `nulls_max`: Nulls appear last in ascending order. - `nulls_min`: Nulls appear first in ascending order. - `nulls_first`: Nulls always be first in any order. - `nulls_last`: Nulls always be last in any order. By default, `nulls_max` is used to follow Postgres's behavior. postgres rule: -datafusion.sql_parser.dialect generic Configure the SQL dialect used by DataFusion's parser; supported values include: Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, MsSQL, ClickHouse, BigQuery, Ansi, DuckDB and Databricks. +datafusion.sql_parser.dialect generic Configure the SQL dialect used by DataFusion's parser; supported values include: Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, MsSQL, ClickHouse, BigQuery, Ansi, DuckDB, Databricks, Spark. datafusion.sql_parser.enable_ident_normalization true When set to true, SQL parser will normalize ident (convert ident to lowercase when not quoted) datafusion.sql_parser.enable_options_value_normalization false When set to true, SQL parser will normalize options value (convert value to lowercase). Note that this option is ignored and will be removed in the future. All case-insensitive values are normalized automatically. datafusion.sql_parser.enable_subquery_sort_elimination true When set to true, DataFusion may remove `ORDER BY` clauses from subqueries or CTEs during SQL planning when their ordering cannot affect the result, such as when no `LIMIT` or other order-sensitive operator depends on them. Disable this option to preserve explicit subquery ordering in the planned query. @@ -561,8 +581,8 @@ CREATE OR REPLACE TABLE some_table AS VALUES (1,2),(3,4); query TTT rowsort DESCRIBE some_table ---- -column1 Int64 YES -column2 Int64 YES +column1 Int64 NO +column2 Int64 NO statement ok DROP TABLE public.some_table; @@ -575,8 +595,8 @@ CREATE OR REPLACE TABLE public.some_table AS VALUES (1,2),(3,4); query TTT rowsort DESCRIBE public.some_table ---- -column1 Int64 YES -column2 Int64 YES +column1 Int64 NO +column2 Int64 NO statement ok DROP TABLE public.some_table; @@ -589,8 +609,8 @@ CREATE OR REPLACE TABLE datafusion.public.some_table AS VALUES (1,2),(3,4); query TTT rowsort DESCRIBE datafusion.public.some_table ---- -column1 Int64 YES -column2 Int64 YES +column1 Int64 NO +column2 Int64 NO statement ok DROP TABLE datafusion.public.some_table; @@ -763,7 +783,7 @@ OPTIONS ('format.has_header' 'true'); query TTTT SHOW CREATE TABLE abc; ---- -datafusion public abc CREATE EXTERNAL TABLE abc STORED AS CSV LOCATION ../../testing/data/csv/aggregate_test_100.csv +datafusion public abc CREATE EXTERNAL TABLE abc STORED AS CSV LOCATION '../../testing/data/csv/aggregate_test_100.csv' # show_external_create_table_with_order statement ok @@ -776,7 +796,7 @@ OPTIONS ('format.has_header' 'true'); query TTTT SHOW CREATE TABLE abc_ordered; ---- -datafusion public abc_ordered CREATE EXTERNAL TABLE abc_ordered STORED AS CSV WITH ORDER (c1) LOCATION ../../testing/data/csv/aggregate_test_100.csv +datafusion public abc_ordered CREATE EXTERNAL TABLE abc_ordered STORED AS CSV WITH ORDER (c1) LOCATION '../../testing/data/csv/aggregate_test_100.csv' statement ok DROP TABLE abc_ordered; @@ -792,7 +812,7 @@ OPTIONS ('format.has_header' 'true'); query TTTT SHOW CREATE TABLE abc_multi_order; ---- -datafusion public abc_multi_order CREATE EXTERNAL TABLE abc_multi_order STORED AS CSV WITH ORDER (c1, c2 DESC) LOCATION ../../testing/data/csv/aggregate_test_100.csv +datafusion public abc_multi_order CREATE EXTERNAL TABLE abc_multi_order STORED AS CSV WITH ORDER (c1, c2 DESC) LOCATION '../../testing/data/csv/aggregate_test_100.csv' statement ok DROP TABLE abc_multi_order; @@ -808,7 +828,7 @@ OPTIONS ('format.has_header' 'true'); query TTTT SHOW CREATE TABLE abc_order_nulls; ---- -datafusion public abc_order_nulls CREATE EXTERNAL TABLE abc_order_nulls STORED AS CSV WITH ORDER (c1 NULLS LAST, c2 DESC NULLS FIRST) LOCATION ../../testing/data/csv/aggregate_test_100.csv +datafusion public abc_order_nulls CREATE EXTERNAL TABLE abc_order_nulls STORED AS CSV WITH ORDER (c1 NULLS LAST, c2 DESC NULLS FIRST) LOCATION '../../testing/data/csv/aggregate_test_100.csv' statement ok DROP TABLE abc_order_nulls; @@ -862,15 +882,6 @@ datafusion public string_agg 1 IN expression String NULL false 1 datafusion public string_agg 2 IN delimiter String NULL false 1 datafusion public string_agg 1 OUT NULL String NULL false 1 -# test variable length arguments -query TTTBI rowsort -select specific_name, data_type, parameter_mode, is_variadic, rid from information_schema.parameters where specific_name = 'concat'; ----- -concat Binary IN true 0 -concat String IN true 1 -concat String OUT false 0 -concat String OUT false 1 - # test ceorcion signature query TTITI rowsort select specific_name, data_type, ordinal_position, parameter_mode, rid from information_schema.parameters where specific_name = 'repeat'; @@ -888,6 +899,17 @@ date_trunc Time(ns) [precision, expression] [String, Time(ns)] SCALAR Truncates date_trunc Timestamp(ns) [precision, expression] [String, Timestamp(ns)] SCALAR Truncates a timestamp or time value to a specified precision. date_trunc(precision, expression) date_trunc Timestamp(ns, "+TZ") [precision, expression] [String, Timestamp(ns, "+TZ")] SCALAR Truncates a timestamp or time value to a specified precision. date_trunc(precision, expression) +# Table functions (UDTFs) appear in information_schema.routines with +# function_type = TABLE and data_type = TABLE. +# Note: built-in `generate_series` and `range` are registered as BOTH a +# scalar UDF and a UDTF, so this test filters to the TABLE rows to make +# a stable assertion. +query TTT rowsort +select routine_name, data_type, function_type from information_schema.routines where function_type = 'TABLE' order by routine_name; +---- +generate_series TABLE TABLE +range TABLE TABLE + statement ok show functions @@ -895,7 +917,7 @@ show functions statement ok reset datafusion.catalog.information_schema; -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/input_file_name.slt b/datafusion/sqllogictest/test_files/input_file_name.slt new file mode 100644 index 0000000000000..32110aa2d69af --- /dev/null +++ b/datafusion/sqllogictest/test_files/input_file_name.slt @@ -0,0 +1,127 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +## input_file_name() tests +########## + +statement ok +COPY (VALUES (10), (20), (30)) +TO 'test_files/scratch/input_file_name/csv/first.csv' +STORED AS CSV; + +statement ok +COPY (VALUES (40), (50), (60)) +TO 'test_files/scratch/input_file_name/csv/second.csv' +STORED AS CSV; + +statement ok +CREATE EXTERNAL TABLE csv_table(column1 int) +STORED AS CSV +LOCATION 'test_files/scratch/input_file_name/csv/'; + +query TI +SELECT + input_file_name(), + column1 +FROM csv_table +ORDER BY column1 +---- +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/csv/first.csv 10 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/csv/first.csv 20 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/csv/first.csv 30 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/csv/second.csv 40 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/csv/second.csv 50 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/csv/second.csv 60 + +query I +SELECT column1 +FROM csv_table +WHERE input_file_name() LIKE '%/first.csv' +ORDER BY column1; +---- +10 +20 +30 + +query TT +EXPLAIN SELECT column1 +FROM csv_table +WHERE input_file_name() LIKE '%/first.csv'; +---- +logical_plan +01)Projection: csv_table.column1 +02)--Filter: __datafusion_extracted_1 LIKE Utf8("%/first.csv") +03)----Projection: input_file_name() AS __datafusion_extracted_1, csv_table.column1 +04)------TableScan: csv_table projection=[column1] +physical_plan +01)FilterExec: __datafusion_extracted_1@0 LIKE %/first.csv, projection=[column1@1] +02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2 +03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/csv/first.csv], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/csv/second.csv]]}, projection=[input_file_name() as __datafusion_extracted_1, column1], file_type=csv, has_header=true + + +query error Execution error: input_file_name\(\) is source dependent and cannot be evaluated directly +SELECT input_file_name() FROM (VALUES (1)) v(x); + +statement ok +DROP TABLE csv_table; + +# Parquet tests as it has its own implementation + +statement ok +COPY (VALUES (10), (20), (30)) +TO 'test_files/scratch/input_file_name/parquet/first.parquet' +STORED AS PARQUET; + +statement ok +COPY (VALUES (40), (50), (60)) +TO 'test_files/scratch/input_file_name/parquet/second.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE pq_table(column1 int) +STORED AS PARQUET +LOCATION 'test_files/scratch/input_file_name/parquet/'; + +query I +SELECT column1 FROM pq_table +WHERE input_file_name() LIKE '%first.parquet' +ORDER BY column1 +---- +10 +20 +30 + +query TT +EXPLAIN SELECT column1 FROM pq_table +WHERE input_file_name() LIKE '%first.parquet' +ORDER BY column1 +---- +logical_plan +01)Sort: pq_table.column1 ASC NULLS LAST +02)--Projection: pq_table.column1 +03)----Filter: __datafusion_extracted_1 LIKE Utf8("%first.parquet") +04)------Projection: input_file_name() AS __datafusion_extracted_1, pq_table.column1 +05)--------TableScan: pq_table projection=[column1] +physical_plan +01)SortPreservingMergeExec: [column1@0 ASC NULLS LAST] +02)--SortExec: expr=[column1@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----FilterExec: __datafusion_extracted_1@0 LIKE %first.parquet, projection=[column1@1] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/parquet/first.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/input_file_name/parquet/second.parquet]]}, projection=[input_file_name() as __datafusion_extracted_1, column1], file_type=parquet + +statement ok +DROP TABLE pq_table; diff --git a/datafusion/sqllogictest/test_files/insert_to_external.slt b/datafusion/sqllogictest/test_files/insert_to_external.slt index 75476c0278c40..f8bd555612625 100644 --- a/datafusion/sqllogictest/test_files/insert_to_external.slt +++ b/datafusion/sqllogictest/test_files/insert_to_external.slt @@ -48,8 +48,8 @@ create table dictionary_encoded_values as values query TTT describe dictionary_encoded_values; ---- -column1 Utf8 YES -column2 Dictionary(Int32, Utf8) YES +column1 Utf8 NO +column2 Dictionary(Int32, Utf8) NO statement ok CREATE EXTERNAL TABLE dictionary_encoded_parquet_partitioned( @@ -128,8 +128,8 @@ logical_plan 03)----Values: (Int64(5), Int64(1)), (Int64(4), Int64(2)), (Int64(7), Int64(7)), (Int64(7), Int64(8)), (Int64(7), Int64(9))... physical_plan 01)DataSinkExec: sink=CsvSink(file_groups=[]) -02)--SortExec: expr=[a@0 ASC NULLS LAST, b@1 DESC], preserve_partitioning=[false] -03)----ProjectionExec: expr=[column1@0 as a, column2@1 as b] +02)--ProjectionExec: expr=[column1@0 as a, column2@1 as b] +03)----SortExec: expr=[column1@0 ASC NULLS LAST, column2@1 DESC], preserve_partitioning=[false] 04)------DataSourceExec: partitions=1, partition_sizes=[1] query I @@ -696,7 +696,7 @@ LOCATION 'test_files/scratch/insert_to_external/external_parquet_table_q7/'; # Config reset -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/insert_values_placeholders.slt b/datafusion/sqllogictest/test_files/insert_values_placeholders.slt new file mode 100644 index 0000000000000..a9cc0ba289344 --- /dev/null +++ b/datafusion/sqllogictest/test_files/insert_values_placeholders.slt @@ -0,0 +1,26 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +## INSERT VALUES placeholder tests +########## + +statement ok +CREATE TABLE placeholder_zero_insert(x BIGINT NULL); + +query error DataFusion error: Error during planning: Invalid placeholder, zero is not a valid index: \$0 +EXPLAIN INSERT INTO placeholder_zero_insert VALUES ($0); diff --git a/datafusion/sqllogictest/test_files/join.slt.part b/datafusion/sqllogictest/test_files/join.slt.part index b9d163d877596..00bea008fc2fc 100644 --- a/datafusion/sqllogictest/test_files/join.slt.part +++ b/datafusion/sqllogictest/test_files/join.slt.part @@ -94,7 +94,7 @@ statement ok set datafusion.execution.batch_size = 4096; # left semi with wrong where clause -query error DataFusion error: Schema error: No field named t2\.t2_id\. Did you mean 't1\.t1_id'\?\. +query error DataFusion error: Schema error: No field named t2\.t2_id\. Did you mean 't1\.t1_id'\?\nValid fields are t1\.t1_id, t1\.t1_name, t1\.t1_int\. SELECT t1.t1_id, t1.t1_name, t1.t1_int FROM t1 LEFT SEMI JOIN t2 ON t1.t1_id = t2.t2_id diff --git a/datafusion/sqllogictest/test_files/joins.slt b/datafusion/sqllogictest/test_files/joins.slt index e0be63fe71525..7a706836f44d6 100644 --- a/datafusion/sqllogictest/test_files/joins.slt +++ b/datafusion/sqllogictest/test_files/joins.slt @@ -1333,19 +1333,57 @@ inner join join_t2 on join_t1.t1_id = join_t2.t2_id ---- logical_plan 01)Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[]] -02)--Projection: join_t1.t1_id -03)----Inner Join: join_t1.t1_id = join_t2.t2_id -04)------TableScan: join_t1 projection=[t1_id] -05)------TableScan: join_t2 projection=[t2_id] +02)--LeftSemi Join: join_t1.t1_id = join_t2.t2_id +03)----TableScan: join_t1 projection=[t1_id] +04)----TableScan: join_t2 projection=[t2_id] physical_plan 01)AggregateExec: mode=FinalPartitioned, gby=[t1_id@0 as t1_id], aggr=[] 02)--RepartitionExec: partitioning=Hash([t1_id@0], 2), input_partitions=2 03)----AggregateExec: mode=Partial, gby=[t1_id@0 as t1_id], aggr=[] -04)------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(t1_id@0, t2_id@0)], projection=[t1_id@0] +04)------HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(t1_id@0, t2_id@0)] 05)--------DataSourceExec: partitions=1, partition_sizes=[1] 06)--------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 07)----------DataSourceExec: partitions=1, partition_sizes=[1] +statement ok +set datafusion.explain.logical_plan_only = true; + +# A single `count(DISTINCT col)` over a join whose other side is used only as an +# existence filter can be rewritten to a semi join. +query TT +EXPLAIN +select join_t1.t1_id, count(distinct join_t1.t1_int) +from join_t1 +inner join join_t2 on join_t1.t1_id = join_t2.t2_id +group by join_t1.t1_id +---- +logical_plan +01)Projection: join_t1.t1_id, count(alias1) AS count(DISTINCT join_t1.t1_int) +02)--Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(alias1)]] +03)----Aggregate: groupBy=[[join_t1.t1_id, join_t1.t1_int AS alias1]], aggr=[[]] +04)------LeftSemi Join: join_t1.t1_id = join_t2.t2_id +05)--------TableScan: join_t1 projection=[t1_id, t1_int] +06)--------TableScan: join_t2 projection=[t2_id] + +# A similar query with two DISTINCT aggregates is currently not rewritten +# TODO: https://github.com/apache/datafusion/issues/22644 +query TT +EXPLAIN +select join_t1.t1_id, count(distinct join_t1.t1_int), count(distinct join_t1.t1_name) +from join_t1 +inner join join_t2 on join_t1.t1_id = join_t2.t2_id +group by join_t1.t1_id +---- +logical_plan +01)Aggregate: groupBy=[[join_t1.t1_id]], aggr=[[count(DISTINCT join_t1.t1_int), count(DISTINCT join_t1.t1_name)]] +02)--Projection: join_t1.t1_id, join_t1.t1_name, join_t1.t1_int +03)----Inner Join: join_t1.t1_id = join_t2.t2_id +04)------TableScan: join_t1 projection=[t1_id, t1_name, t1_int] +05)------TableScan: join_t2 projection=[t2_id] + +statement ok +set datafusion.explain.logical_plan_only = false; + # Join on struct query TT explain select join_t3.s3, join_t4.s4 @@ -1411,10 +1449,9 @@ logical_plan 01)Projection: count(alias1) AS count(DISTINCT join_t1.t1_id) 02)--Aggregate: groupBy=[[]], aggr=[[count(alias1)]] 03)----Aggregate: groupBy=[[join_t1.t1_id AS alias1]], aggr=[[]] -04)------Projection: join_t1.t1_id -05)--------Inner Join: join_t1.t1_id = join_t2.t2_id -06)----------TableScan: join_t1 projection=[t1_id] -07)----------TableScan: join_t2 projection=[t2_id] +04)------LeftSemi Join: join_t1.t1_id = join_t2.t2_id +05)--------TableScan: join_t1 projection=[t1_id] +06)--------TableScan: join_t2 projection=[t2_id] physical_plan 01)ProjectionExec: expr=[count(alias1)@0 as count(DISTINCT join_t1.t1_id)] 02)--AggregateExec: mode=Final, gby=[], aggr=[count(alias1)] @@ -1423,7 +1460,7 @@ physical_plan 05)--------AggregateExec: mode=FinalPartitioned, gby=[alias1@0 as alias1], aggr=[] 06)----------RepartitionExec: partitioning=Hash([alias1@0], 2), input_partitions=2 07)------------AggregateExec: mode=Partial, gby=[t1_id@0 as alias1], aggr=[] -08)--------------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(t1_id@0, t2_id@0)], projection=[t1_id@0] +08)--------------HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(t1_id@0, t2_id@0)] 09)----------------DataSourceExec: partitions=1, partition_sizes=[1] 10)----------------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 11)------------------DataSourceExec: partitions=1, partition_sizes=[1] @@ -1952,7 +1989,7 @@ where join_t1.t1_id + 12 not in (select join_t2.t2_id + 1 from join_t2 where join_t1.t1_int > 0) ---- logical_plan -01)LeftAnti Join: CAST(join_t1.t1_id AS Int64) + Int64(12) = __correlated_sq_1.join_t2.t2_id + Int64(1) Filter: join_t1.t1_int > UInt32(0) +01)LeftAnti Join: CAST(join_t1.t1_id AS Int64) + Int64(12) = __correlated_sq_1.join_t2.t2_id + Int64(1) Filter: join_t1.t1_int > UInt32(0) null_aware 02)--TableScan: join_t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 04)----Projection: CAST(join_t2.t2_id AS Int64) + Int64(1) @@ -5148,7 +5185,7 @@ LEFT ANTI JOIN ( ) t2 ON t1.k = t2.k; ---- Plan with Metrics -01)HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(k@0, k@0)], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=1, array_map_created_count=0, build_input_batches=0, build_input_rows=0, input_batches=1, input_rows=2, build_mem_used=, build_time=, join_time=, avg_fanout=N/A (0/0), probe_hit_rate=0% (0/2)] +01)HashJoinExec: mode=CollectLeft, join_type=RightAnti, on=[(k@0, k@0)], metrics=[output_rows=2, elapsed_compute=, output_bytes=, output_batches=1, build_mem_used=, array_map_created_count=0, build_input_batches=0, build_input_rows=0, input_batches=1, input_rows=2, build_time=, join_time=, avg_fanout=N/A (0/0), probe_hit_rate=0% (0/2)] 02)--ProjectionExec: expr=[column1@0 as k], metrics=[output_rows=0, elapsed_compute=, output_bytes=, output_batches=0, expr_0_eval_time=] 03)----FilterExec: column1@0 != 1, metrics=[output_rows=0, elapsed_compute=, output_bytes=, output_batches=0, selectivity=0% (0/1)] 04)------DataSourceExec: partitions=1, partition_sizes=[1], metrics=[] @@ -5365,10 +5402,10 @@ LEFT JOIN issue_19067_right r ON l.join_key = r.join_key ORDER BY l.id; ---- physical_plan -01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--ProjectionExec: expr=[id@2 as id, join_key@3 as left_key, join_key@0 as right_key, value@1 as value] -03)----HashJoinExec: mode=CollectLeft, join_type=Right, on=[(join_key@0, join_key@1)] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[id@2 as id, join_key@3 as left_key, join_key@0 as right_key, value@1 as value] +02)--HashJoinExec: mode=CollectLeft, join_type=Right, on=[(join_key@0, join_key@1)] +03)----DataSourceExec: partitions=1, partition_sizes=[1] +04)----SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] 05)------DataSourceExec: partitions=1, partition_sizes=[1] statement count 0 @@ -5506,7 +5543,7 @@ statement count 0 DROP TABLE t2; statement ok -CREATE TABLE t1(a INT, b INT) AS VALUES +CREATE TABLE t1(a INT, b INT) AS VALUES (NULL, 1), (NULL, 2), (NULL, 3), (NULL, 4), (NULL, 5); statement ok @@ -5518,7 +5555,7 @@ CREATE TABLE t2(c INT) AS VALUES (1), (2); query II SELECT sub.a, sub.b FROM ( SELECT * FROM t1 ORDER BY b LIMIT 1 -) sub +) sub JOIN t2 ON sub.a = t2.c; ---- @@ -5527,3 +5564,660 @@ DROP TABLE t1; statement ok DROP TABLE t2; + +# Regression test for a LEFT JOIN with a non-equijoin predicate (forces +# NestedLoopJoinExec) and a multi-partition probe side. Previously the unmatched +# left rows could be emitted before all partitions finished probing, adding +# spurious NULL-padded rows. The result must include every left row exactly once. +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.execution.batch_size = 2; + +statement ok +CREATE TABLE nlj_left(id INT, v INT) AS VALUES (1, 4), (2, 72), (3, 41), (4, 98), (5, 91); + +statement ok +CREATE TABLE nlj_right(w INT) AS VALUES (49), (58), (83), (3), (76); + +query III +SELECT id, v, w FROM nlj_left LEFT JOIN nlj_right ON nlj_left.v < nlj_right.w ORDER BY id, w; +---- +1 4 49 +1 4 58 +1 4 76 +1 4 83 +2 72 76 +2 72 83 +3 41 49 +3 41 58 +3 41 76 +3 41 83 +4 98 NULL +5 91 NULL + +statement ok +DROP TABLE nlj_left; + +statement ok +DROP TABLE nlj_right; + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.execution.batch_size; + +########## +# Eliminate unused outer joins (`EliminateJoin` rule) +# +# An outer join whose non-preserved side is unreferenced above the join is +# removed entirely when it cannot duplicate the preserved side's rows: either +# the non-preserved side is unique on the join keys (e.g. PRIMARY KEY / +# UNIQUE constraint or GROUP BY), or the join's ancestors are +# duplicate-insensitive. Most cases below exercise the LEFT JOIN direction; +# RIGHT JOIN is symmetric and covered at the end of the section. +########## + +statement ok +CREATE TABLE elim_users (id INT primary key, name VARCHAR) AS VALUES + (1, 'alice'), + (2, 'bob'), + (4, 'dave'); + +statement ok +CREATE TABLE elim_orders (order_id INT, user_id INT, amount INT) AS VALUES + (1, 1, 100), + (2, 1, 200), + (3, 3, 50); + +# The right side is unique on the join key (primary key) and unused above the +# join: the LEFT JOIN is removed from the plan. +query TT +EXPLAIN SELECT order_id, amount FROM elim_orders LEFT JOIN elim_users ON user_id = id; +---- +logical_plan TableScan: elim_orders projection=[order_id, amount] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +# All orders are returned, including the one with no matching user. +query II rowsort +SELECT order_id, amount FROM elim_orders LEFT JOIN elim_users ON user_id = id; +---- +1 100 +2 200 +3 50 + +# A WHERE clause on left-side columns does not block the rewrite. +query TT +EXPLAIN SELECT order_id FROM elim_orders LEFT JOIN elim_users ON user_id = id WHERE amount > 100; +---- +logical_plan +01)Projection: elim_orders.order_id +02)--Filter: elim_orders.amount > Int32(100) +03)----TableScan: elim_orders projection=[order_id, amount] +physical_plan +01)FilterExec: amount@1 > 100, projection=[order_id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT order_id FROM elim_orders LEFT JOIN elim_users ON user_id = id WHERE amount > 100; +---- +2 + +# An extra join filter on right-side columns does not block the rewrite: for a +# left join it only decides whether a left row is matched or null-padded, and +# either way the row is emitted. +query TT +EXPLAIN SELECT order_id FROM elim_orders LEFT JOIN elim_users ON user_id = id AND name <> 'bob'; +---- +logical_plan TableScan: elim_orders projection=[order_id] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT order_id FROM elim_orders LEFT JOIN elim_users ON user_id = id AND name <> 'bob'; +---- +1 +2 +3 + +# count(*) is duplicate-sensitive, but the unique join key guarantees each +# order appears exactly once, so the join is still removed. +query TT +EXPLAIN SELECT count(*) FROM elim_orders LEFT JOIN elim_users ON user_id = id; +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +03)----TableScan: elim_orders projection=[] +physical_plan +01)ProjectionExec: expr=[3 as count(*)] +02)--PlaceholderRowExec + +query I +SELECT count(*) FROM elim_orders LEFT JOIN elim_users ON user_id = id; +---- +3 + +# A DISTINCT (or GROUP BY) right side is unique on its keys even without +# declared constraints, so the join is removed. +query TT +EXPLAIN SELECT o.order_id FROM elim_orders o LEFT JOIN (SELECT DISTINCT user_id FROM elim_orders) d ON o.user_id = d.user_id; +---- +logical_plan +01)SubqueryAlias: o +02)--TableScan: elim_orders projection=[order_id] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT o.order_id FROM elim_orders o LEFT JOIN (SELECT DISTINCT user_id FROM elim_orders) d ON o.user_id = d.user_id; +---- +1 +2 +3 + +# Negative case: the right side is referenced in the SELECT list, so the join +# must stay. +query TT +EXPLAIN SELECT order_id, name FROM elim_orders LEFT JOIN elim_users ON user_id = id; +---- +logical_plan +01)Projection: elim_orders.order_id, elim_users.name +02)--Left Join: elim_orders.user_id = elim_users.id +03)----TableScan: elim_orders projection=[order_id, user_id] +04)----TableScan: elim_users projection=[id, name] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Left, on=[(user_id@1, id@0)], projection=[order_id@0, name@3] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Negative case: the right side is not unique on the join key, so a left row +# may match several right rows; the join must stay. +query TT +EXPLAIN SELECT elim_users.id FROM elim_users LEFT JOIN elim_orders ON elim_users.id = elim_orders.user_id; +---- +logical_plan +01)Projection: elim_users.id +02)--Left Join: elim_users.id = elim_orders.user_id +03)----TableScan: elim_users projection=[id] +04)----TableScan: elim_orders projection=[user_id] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, user_id@0)], projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +# ... and the duplicates it produces are observable: user 1 has two orders. +query I rowsort +SELECT elim_users.id FROM elim_users LEFT JOIN elim_orders ON elim_users.id = elim_orders.user_id; +---- +1 +1 +2 +4 + +# The same non-unique right side under a DISTINCT: the join's ancestors are +# duplicate-insensitive, so the extra matches only affect row multiplicity +# and the join is removed even without uniqueness on the join key. +query TT +EXPLAIN SELECT DISTINCT name FROM elim_users LEFT JOIN elim_orders ON elim_users.id = elim_orders.user_id; +---- +logical_plan +01)Aggregate: groupBy=[[elim_users.name]], aggr=[[]] +02)--TableScan: elim_users projection=[name] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[name@0 as name], aggr=[] +02)--RepartitionExec: partitioning=Hash([name@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[name@0 as name], aggr=[] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +query T rowsort +SELECT DISTINCT name FROM elim_users LEFT JOIN elim_orders ON elim_users.id = elim_orders.user_id; +---- +alice +bob +dave + +# Negative case: count(*) observes row multiplicity and the right side is not +# unique on the join key, so the join must stay (user 1 has two orders). +query TT +EXPLAIN SELECT count(*) FROM elim_users LEFT JOIN elim_orders ON elim_users.id = elim_orders.user_id; +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +03)----Projection: +04)------Left Join: elim_users.id = elim_orders.user_id +05)--------TableScan: elim_users projection=[id] +06)--------TableScan: elim_orders projection=[user_id] +physical_plan +01)ProjectionExec: expr=[count(Int64(1))@0 as count(*)] +02)--AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] +03)----CoalescePartitionsExec +04)------AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] +05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +06)----------HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, user_id@0)], projection=[] +07)------------DataSourceExec: partitions=1, partition_sizes=[1] +08)------------DataSourceExec: partitions=1, partition_sizes=[1] + +query I +SELECT count(*) FROM elim_users LEFT JOIN elim_orders ON elim_users.id = elim_orders.user_id; +---- +4 + +# A left join with no equi-join keys at all (ON true) matches every left row +# with every right row. Under a duplicate-insensitive ancestor (DISTINCT) the +# multiplication is unobservable and the join is removed. +query TT +EXPLAIN SELECT DISTINCT order_id FROM elim_orders LEFT JOIN elim_users ON true; +---- +logical_plan +01)Aggregate: groupBy=[[elim_orders.order_id]], aggr=[[]] +02)--TableScan: elim_orders projection=[order_id] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[order_id@0 as order_id], aggr=[] +02)--RepartitionExec: partitioning=Hash([order_id@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[order_id@0 as order_id], aggr=[] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT DISTINCT order_id FROM elim_orders LEFT JOIN elim_users ON true; +---- +1 +2 +3 + +# Negative case: without the DISTINCT the multiplication is observable (each +# order is repeated once per user), so the join must stay. +query TT +EXPLAIN SELECT order_id FROM elim_orders LEFT JOIN elim_users ON true; +---- +logical_plan +01)Left Join: +02)--TableScan: elim_orders projection=[order_id] +03)--TableScan: elim_users projection=[] +physical_plan +01)NestedLoopJoinExec: join_type=Right +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT order_id FROM elim_orders LEFT JOIN elim_users ON true; +---- +1 +1 +1 +2 +2 +2 +3 +3 +3 + +# A LIMIT makes the row count observable, but the uniqueness path does not +# depend on duplicate-insensitivity: the unique (PK) right side is unused, so +# the join is removed even under a LIMIT. +query TT +EXPLAIN SELECT order_id FROM elim_orders LEFT JOIN elim_users ON user_id = id LIMIT 2; +---- +logical_plan +01)Limit: skip=0, fetch=2 +02)--TableScan: elim_orders projection=[order_id], fetch=2 +physical_plan DataSourceExec: partitions=1, partition_sizes=[1], fetch=2 + +query I rowsort +SELECT order_id FROM elim_orders LEFT JOIN elim_users ON user_id = id LIMIT 2; +---- +1 +2 + +# Negative case: a LIMIT between the join and a DISTINCT ancestor makes the +# row count observable, so the DISTINCT's duplicate-insensitivity does not +# reach the join; with a non-unique right side the join must stay. +query TT +EXPLAIN SELECT DISTINCT id FROM (SELECT elim_users.id FROM elim_users LEFT JOIN elim_orders ON elim_users.id = elim_orders.user_id LIMIT 5); +---- +logical_plan +01)Aggregate: groupBy=[[elim_users.id]], aggr=[[]] +02)--Projection: elim_users.id +03)----Limit: skip=0, fetch=5 +04)------Left Join: elim_users.id = elim_orders.user_id +05)--------Limit: skip=0, fetch=5 +06)----------TableScan: elim_users projection=[id], fetch=5 +07)--------TableScan: elim_orders projection=[user_id] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[id@0 as id], aggr=[] +02)--RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[id@0 as id], aggr=[] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)--------HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@0, user_id@0)], projection=[id@0], fetch=5 +06)----------DataSourceExec: partitions=1, partition_sizes=[1], fetch=5 +07)----------DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT DISTINCT id FROM (SELECT elim_users.id FROM elim_users LEFT JOIN elim_orders ON elim_users.id = elim_orders.user_id LIMIT 5); +---- +1 +2 +4 + +# LEFT JOIN LATERAL decorrelates into a plain left join, with equality +# predicates extracted as join keys: the unique (PK) lateral side is unused, +# so the join is removed. +query TT +EXPLAIN SELECT order_id FROM elim_orders LEFT JOIN LATERAL (SELECT * FROM elim_users WHERE id = user_id) AS u ON true; +---- +logical_plan TableScan: elim_orders projection=[order_id] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT order_id FROM elim_orders LEFT JOIN LATERAL (SELECT * FROM elim_users WHERE id = user_id) AS u ON true; +---- +1 +2 +3 + +# A non-equality lateral predicate becomes a join filter, which does not +# block removal under a duplicate-insensitive ancestor (DISTINCT). +query TT +EXPLAIN SELECT DISTINCT order_id FROM elim_orders LEFT JOIN LATERAL (SELECT * FROM elim_users WHERE id > user_id) AS u ON true; +---- +logical_plan +01)Aggregate: groupBy=[[elim_orders.order_id]], aggr=[[]] +02)--TableScan: elim_orders projection=[order_id] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[order_id@0 as order_id], aggr=[] +02)--RepartitionExec: partitioning=Hash([order_id@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[order_id@0 as order_id], aggr=[] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT DISTINCT order_id FROM elim_orders LEFT JOIN LATERAL (SELECT * FROM elim_users WHERE id > user_id) AS u ON true; +---- +1 +2 +3 + +# Negative case: without the DISTINCT a filter-only lateral can multiply left +# rows observably (each order matches every user with a greater id), so the +# join must stay. +query TT +EXPLAIN SELECT order_id FROM elim_orders LEFT JOIN LATERAL (SELECT * FROM elim_users WHERE id > user_id) AS u ON true; +---- +logical_plan +01)Projection: elim_orders.order_id +02)--Left Join: Filter: u.id > elim_orders.user_id +03)----TableScan: elim_orders projection=[order_id, user_id] +04)----SubqueryAlias: u +05)------TableScan: elim_users projection=[id] +physical_plan +01)NestedLoopJoinExec: join_type=Right, filter=id@1 > user_id@0, projection=[order_id@1] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT order_id FROM elim_orders LEFT JOIN LATERAL (SELECT * FROM elim_users WHERE id > user_id) AS u ON true; +---- +1 +1 +2 +2 +3 + +# RIGHT JOIN is symmetric: the join is removed when its *left* side is +# unreferenced above the join and cannot duplicate right rows. + +# The left side is unique on the join key (primary key) and unused above the +# join: the RIGHT JOIN is removed from the plan. +query TT +EXPLAIN SELECT order_id, amount FROM elim_users RIGHT JOIN elim_orders ON id = user_id; +---- +logical_plan TableScan: elim_orders projection=[order_id, amount] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +# All orders are returned, including the one with no matching user. +query II rowsort +SELECT order_id, amount FROM elim_users RIGHT JOIN elim_orders ON id = user_id; +---- +1 100 +2 200 +3 50 + +# An extra join filter on left-side columns does not block the rewrite: for a +# right join it only decides whether a right row is matched or null-padded, +# and either way the row is emitted. +query TT +EXPLAIN SELECT order_id FROM elim_users RIGHT JOIN elim_orders ON id = user_id AND name <> 'bob'; +---- +logical_plan TableScan: elim_orders projection=[order_id] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +query I rowsort +SELECT order_id FROM elim_users RIGHT JOIN elim_orders ON id = user_id AND name <> 'bob'; +---- +1 +2 +3 + +# count(*) is duplicate-sensitive, but the unique join key guarantees each +# order appears exactly once, so the join is still removed. +query TT +EXPLAIN SELECT count(*) FROM elim_users RIGHT JOIN elim_orders ON id = user_id; +---- +logical_plan +01)Projection: count(Int64(1)) AS count(*) +02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +03)----TableScan: elim_orders projection=[] +physical_plan +01)ProjectionExec: expr=[3 as count(*)] +02)--PlaceholderRowExec + +query I +SELECT count(*) FROM elim_users RIGHT JOIN elim_orders ON id = user_id; +---- +3 + +# The left side is not unique on the join key, but a DISTINCT ancestor makes +# the extra matches unobservable, so the join is removed even without +# uniqueness on the join key. +query TT +EXPLAIN SELECT DISTINCT name FROM elim_orders RIGHT JOIN elim_users ON elim_orders.user_id = elim_users.id; +---- +logical_plan +01)Aggregate: groupBy=[[elim_users.name]], aggr=[[]] +02)--TableScan: elim_users projection=[name] +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[name@0 as name], aggr=[] +02)--RepartitionExec: partitioning=Hash([name@0], 4), input_partitions=1 +03)----AggregateExec: mode=Partial, gby=[name@0 as name], aggr=[] +04)------DataSourceExec: partitions=1, partition_sizes=[1] + +query T rowsort +SELECT DISTINCT name FROM elim_orders RIGHT JOIN elim_users ON elim_orders.user_id = elim_users.id; +---- +alice +bob +dave + +# Negative case: the left side is referenced in the SELECT list, so the join +# must stay. +query TT +EXPLAIN SELECT order_id, name FROM elim_users RIGHT JOIN elim_orders ON id = user_id; +---- +logical_plan +01)Projection: elim_orders.order_id, elim_users.name +02)--Right Join: elim_users.id = elim_orders.user_id +03)----TableScan: elim_users projection=[id, name] +04)----TableScan: elim_orders projection=[order_id, user_id] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Left, on=[(user_id@1, id@0)], projection=[order_id@0, name@3] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Negative case: the left side is not unique on the join key, so a right row +# may match several left rows; the join must stay. +query TT +EXPLAIN SELECT elim_users.id FROM elim_orders RIGHT JOIN elim_users ON elim_orders.user_id = elim_users.id; +---- +logical_plan +01)Projection: elim_users.id +02)--Right Join: elim_orders.user_id = elim_users.id +03)----TableScan: elim_orders projection=[user_id] +04)----TableScan: elim_users projection=[id] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Right, on=[(user_id@0, id@0)], projection=[id@1] +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +# ... and the duplicates it produces are observable: user 1 has two orders. +query I rowsort +SELECT elim_users.id FROM elim_orders RIGHT JOIN elim_users ON elim_orders.user_id = elim_users.id; +---- +1 +1 +2 +4 + +statement ok +DROP TABLE elim_users; + +statement ok +DROP TABLE elim_orders; + +# A UNIQUE constraint, unlike PRIMARY KEY, permits NULLs — and per SQL +# semantics several NULLs may coexist in a UNIQUE column. Whether a nullable +# UNIQUE key proves uniqueness on the join keys therefore depends on the +# join's null semantics. +statement ok +CREATE TABLE elim_null_keys (id INT, k INT) AS VALUES + (1, 10), + (2, NULL), + (3, 30); + +statement ok +CREATE TABLE elim_null_lookup (ukey INT UNIQUE, payload INT) AS VALUES + (10, 100), + (NULL, 200), + (NULL, 300); + +# Under the default null semantics (`=`), NULL keys match nothing, so the +# nullable UNIQUE right side still yields at most one match per left row and +# the join is removed. +query TT +EXPLAIN SELECT id FROM elim_null_keys LEFT JOIN elim_null_lookup ON k = ukey; +---- +logical_plan TableScan: elim_null_keys projection=[id] +physical_plan DataSourceExec: partitions=1, partition_sizes=[1] + +# The NULL-keyed left row matches nothing and is emitted exactly once. +query I rowsort +SELECT id FROM elim_null_keys LEFT JOIN elim_null_lookup ON k = ukey; +---- +1 +2 +3 + +# Negative case: IS NOT DISTINCT FROM compares NULLs as equal, so both NULL +# rows in the UNIQUE column match a NULL left key; the right side is not +# unique under these semantics and the join must stay. +query TT +EXPLAIN SELECT id FROM elim_null_keys LEFT JOIN elim_null_lookup ON k IS NOT DISTINCT FROM ukey; +---- +logical_plan +01)Projection: elim_null_keys.id +02)--Left Join: elim_null_keys.k = elim_null_lookup.ukey +03)----TableScan: elim_null_keys projection=[id, k] +04)----TableScan: elim_null_lookup projection=[ukey] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=Right, on=[(ukey@0, k@1)], projection=[id@1], NullsEqual: true +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +# ... and the duplicates are observable: the NULL-keyed left row matches both +# NULL lookup rows. +query I rowsort +SELECT id FROM elim_null_keys LEFT JOIN elim_null_lookup ON k IS NOT DISTINCT FROM ukey; +---- +1 +2 +2 +3 + +statement ok +DROP TABLE elim_null_keys; + +statement ok +DROP TABLE elim_null_lookup; + +# Regression test: a `CollectLeft` `HashJoinExec` requires `SinglePartition` on its build +# (left) child, and the `CoalescePartitionsExec` that satisfies it must survive the +# sort-parallelization phase of `EnsureRequirements`. It used to be removed positionally +# (the traversal descends into the join because the *probe* side is linked to a coalesce), +# leaving a multi-partition build side that `SanityCheckPlan` rejects with +# "does not satisfy distribution requirements: SinglePartition". + +statement ok +set datafusion.execution.target_partitions = 8; + +# Keep the scan multi-partition as written, i.e. one partition per file. +statement ok +set datafusion.optimizer.repartition_file_scans = false; + +statement ok +CREATE TABLE collect_left_src (id INT, ts INT) AS VALUES (1, 10), (2, 20), (3, 30); + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/0.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/1.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/2.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY (SELECT * FROM collect_left_src) TO 'test_files/scratch/joins/collect_left/3.parquet' STORED AS PARQUET; +---- +3 + +statement ok +CREATE EXTERNAL TABLE collect_left STORED AS PARQUET LOCATION 'test_files/scratch/joins/collect_left/'; + +# The build side is the 4-partition scan; the probe side is the `DISTINCT ON` aggregate, +# whose `CoalescePartitionsExec` is what makes the traversal reach the join. +query I +SELECT a.id +FROM collect_left a +LEFT JOIN (SELECT DISTINCT ON (id) id, ts FROM collect_left ORDER BY id, ts) f + ON a.id = f.id +ORDER BY a.id; +---- +1 +1 +1 +1 +2 +2 +2 +2 +3 +3 +3 +3 + +statement ok +DROP TABLE collect_left; + +statement ok +DROP TABLE collect_left_src; + +statement ok +reset datafusion.optimizer.repartition_file_scans; + +statement ok +set datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/limit.slt b/datafusion/sqllogictest/test_files/limit.slt index fc62584dc3df1..58a655c02b2fc 100644 --- a/datafusion/sqllogictest/test_files/limit.slt +++ b/datafusion/sqllogictest/test_files/limit.slt @@ -748,7 +748,7 @@ explain select * from testSubQueryLimit as t1 join (select * from testSubQueryLi ---- logical_plan 01)Limit: skip=0, fetch=10 -02)--Cross Join: +02)--Cross Join: 03)----SubqueryAlias: t1 04)------Limit: skip=0, fetch=10 05)--------TableScan: testsubquerylimit projection=[a, b], fetch=10 @@ -773,7 +773,7 @@ explain select * from testSubQueryLimit as t1 join (select * from testSubQueryLi ---- logical_plan 01)Limit: skip=0, fetch=2 -02)--Cross Join: +02)--Cross Join: 03)----SubqueryAlias: t1 04)------Limit: skip=0, fetch=2 05)--------TableScan: testsubquerylimit projection=[a, b], fetch=2 @@ -868,7 +868,7 @@ physical_plan 01)ProjectionExec: expr=[1 as foo] 02)--SortPreservingMergeExec: [part_key@0 ASC NULLS LAST], fetch=1 03)----SortExec: TopK(fetch=1), expr=[part_key@0 ASC NULLS LAST], preserve_partitioning=[true] -04)------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit/test_limit_with_partitions/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit/test_limit_with_partitions/part-1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit/test_limit_with_partitions/part-2.parquet]]}, projection=[part_key], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[part_key@0 ASC NULLS LAST] +04)------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit/test_limit_with_partitions/part-0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit/test_limit_with_partitions/part-1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit/test_limit_with_partitions/part-2.parquet]]}, projection=[part_key], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[part_key@0 ASC NULLS LAST], dynamic_rg_pruning=eligible query I with selection as ( @@ -989,3 +989,58 @@ c-4 statement ok DROP TABLE t21176; + +# Regression test for https://github.com/apache/datafusion/issues/22489 +# An outer ORDER BY / OFFSET must not reduce an inner LIMIT when the two are +# separated by a sort on a *different* key. + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +CREATE TABLE t22489 (g INT, x INT, y INT) AS VALUES (1, 10, 4), (2, 20, 3), (3, 30, 2), (4, 40, 1); + +# Inner ORDER BY sx DESC LIMIT 4 keeps all four groups; the outer ORDER BY +# sy DESC OFFSET 1 then drops only the sy-max group (g=1), so g=2,3,4 remain. +query III +SELECT * FROM ( + SELECT g, SUM(x) AS sx, SUM(y) AS sy FROM t22489 GROUP BY g + ORDER BY sx DESC LIMIT 4 +) q +ORDER BY sy DESC +OFFSET 1; +---- +2 20 3 +3 30 2 +4 40 1 + +query TT +EXPLAIN +SELECT * FROM ( + SELECT g, SUM(x) AS sx, SUM(y) AS sy FROM t22489 GROUP BY g + ORDER BY sx DESC LIMIT 4 +) q +ORDER BY sy DESC +OFFSET 1; +---- +logical_plan +01)Limit: skip=1, fetch=None +02)--Sort: q.sy DESC NULLS FIRST +03)----SubqueryAlias: q +04)------Sort: sx DESC NULLS FIRST, fetch=4 +05)--------Projection: t22489.g, sum(t22489.x) AS sx, sum(t22489.y) AS sy +06)----------Aggregate: groupBy=[[t22489.g]], aggr=[[sum(CAST(t22489.x AS Int64)), sum(CAST(t22489.y AS Int64))]] +07)------------TableScan: t22489 projection=[g, x, y] +physical_plan +01)GlobalLimitExec: skip=1, fetch=None +02)--SortExec: expr=[sy@2 DESC], preserve_partitioning=[false] +03)----SortPreservingMergeExec: [sx@1 DESC], fetch=4 +04)------ProjectionExec: expr=[g@0 as g, sum(t22489.x)@1 as sx, sum(t22489.y)@2 as sy] +05)--------SortExec: TopK(fetch=4), expr=[sum(t22489.x)@1 DESC], preserve_partitioning=[true] +06)----------AggregateExec: mode=FinalPartitioned, gby=[g@0 as g], aggr=[sum(t22489.x), sum(t22489.y)] +07)------------RepartitionExec: partitioning=Hash([g@0], 4), input_partitions=1 +08)--------------AggregateExec: mode=Partial, gby=[g@0 as g], aggr=[sum(t22489.x), sum(t22489.y)] +09)----------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +DROP TABLE t22489; diff --git a/datafusion/sqllogictest/test_files/limit_pruning.slt b/datafusion/sqllogictest/test_files/limit_pruning.slt index 373e1636a2bb6..4ef0b5c74f3e7 100644 --- a/datafusion/sqllogictest/test_files/limit_pruning.slt +++ b/datafusion/sqllogictest/test_files/limit_pruning.slt @@ -63,7 +63,7 @@ set datafusion.explain.analyze_level = summary; query TT explain analyze select * from tracking_data where species > 'M' AND s >= 50 limit 3; ---- -Plan with Metrics DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit_pruning/data.parquet]]}, projection=[species, s], limit=3, file_type=parquet, predicate=species@0 > M AND s@1 >= 50, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50, required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=0 total → 0 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=2 total → 0 matched, bytes_scanned=, metadata_load_time=, scan_efficiency_ratio= (171/2.35 K)] +Plan with Metrics DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit_pruning/data.parquet]]}, projection=[species, s], limit=3, file_type=parquet, predicate=species@0 > M AND s@1 >= 50, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50, required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=0 total → 0 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=2 total → 0 matched, bytes_scanned=, metadata_load_time=, scan_efficiency_ratio= (159/2.23 K)] statement ok CREATE TABLE fully_matched_limit_source AS VALUES @@ -120,7 +120,7 @@ explain analyze select * from tracking_data where species > 'M' AND s >= 50 orde ---- Plan with Metrics 01)SortExec: TopK(fetch=3), expr=[species@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[species@0 < Nlpine Sheep], metrics=[output_rows=3, elapsed_compute=, output_bytes=] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit_pruning/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, metadata_load_time=, scan_efficiency_ratio= (521/2.35 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/limit_pruning/data.parquet]]}, projection=[species, s], file_type=parquet, predicate=species@0 > M AND s@1 >= 50 AND DynamicFilter [ species@0 < Nlpine Sheep ], sort_order_for_reorder=[species@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=species_null_count@1 != row_count@2 AND species_max@0 > M AND s_null_count@4 != row_count@2 AND s_max@3 >= 50 AND species_null_count@1 != row_count@2 AND species_min@5 < Nlpine Sheep, required_guarantees=[], metrics=[output_rows=3, elapsed_compute=, output_bytes=, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=4 total → 3 matched -> 1 fully matched, row_groups_pruned_bloom_filter=3 total → 3 matched, page_index_pages_pruned=2 total → 2 matched, page_index_pages_skipped_by_fully_matched=1, limit_pruned_row_groups=0 total → 0 matched, bytes_scanned=, row_groups_pruned_dynamic_filter=0, metadata_load_time=, scan_efficiency_ratio= (/)] statement ok drop table tracking_data; diff --git a/datafusion/sqllogictest/test_files/listing_table_statistics.slt b/datafusion/sqllogictest/test_files/listing_table_statistics.slt index 4b2aa0f563b22..3021ee5334f58 100644 --- a/datafusion/sqllogictest/test_files/listing_table_statistics.slt +++ b/datafusion/sqllogictest/test_files/listing_table_statistics.slt @@ -35,7 +35,7 @@ query TT explain format indent select * from t; ---- logical_plan TableScan: t projection=[int_col, str_col] -physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/listing_table_statistics/1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/listing_table_statistics/2.parquet]]}, projection=[int_col, str_col], file_type=parquet, statistics=[Rows=Exact(4), Bytes=Absent, [(Col[0]: Min=Exact(Int64(-1)) Max=Exact(Int64(3)) Null=Exact(0) ScanBytes=Exact(32)),(Col[1]: Min=Exact(Utf8View("a")) Max=Exact(Utf8View("d")) Null=Exact(0) ScanBytes=Inexact(100))]] +physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/listing_table_statistics/1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/listing_table_statistics/2.parquet]]}, projection=[int_col, str_col], file_type=parquet, statistics=[Rows=Exact(4), Bytes=Absent, [(Col[0]: Min=Exact(Int64(-1)) Max=Exact(Int64(3)) Null=Exact(0) ScanBytes=Exact(32)),(Col[1]: Min=Exact(Utf8View("a")) Max=Exact(Utf8View("d")) Null=Exact(0) ScanBytes=Inexact(88))]] statement ok drop table t; diff --git a/datafusion/sqllogictest/test_files/map.slt b/datafusion/sqllogictest/test_files/map.slt index 62e70e6080bab..9ec2d0b894535 100644 --- a/datafusion/sqllogictest/test_files/map.slt +++ b/datafusion/sqllogictest/test_files/map.slt @@ -579,7 +579,7 @@ SELECT MAP { 'a': 1, 'b': 2, 'c': 3 }['a']; # accessing map with non-string key in case expression query I -SELECT (CASE WHEN 1 > 0 THEN MAP {'x': 100} ELSE MAP {'y': 200} END)['x']; +SELECT (CASE WHEN 1 > 0 THEN MAP {'x': 100} ELSE MAP {'y': 200} END)['x']; ---- 100 @@ -642,6 +642,12 @@ select map_extract(MAP {1: 1, 2: 2, 3:3}, '1'), map_extract(MAP {1: 1, 2: 2, 3:3 ---- [1] [1] [1] [NULL] [1] +# null arg +query ? +select map_extract(NULL, 'a'); +---- +NULL + # map_extract with columns query ??? select map_extract(column1, 1), map_extract(column1, 5), map_extract(column1, 7) from map_array_table_1; @@ -901,3 +907,295 @@ NULL statement ok drop table tt; + +# mixed scalar/array inputs +query ? +SELECT map(['a','b'], [column1, column1 * 10]) FROM (VALUES (1), (2), (3)) t; +---- +{a: 1, b: 10} +{a: 2, b: 20} +{a: 3, b: 30} + +query ? +SELECT map([column1, column1 * 10], ['x','y']) FROM (VALUES (1), (2), (3)) t; +---- +{1: x, 10: y} +{2: x, 20: y} +{3: x, 30: y} + +# tests for DISTINCT / GROUP BY / aggregation on map columns +# https://github.com/apache/datafusion/issues/15428 + +# NOTE: the CAST(NULL AS BIGINT) in the VALUES list below predates the fix for +# https://github.com/apache/datafusion/issues/23474 and is no longer required. +# It is kept as-is to document the historical workaround; the un-cast form is +# exercised in the "map NULL value coercion in VALUES" section further below. +statement ok +CREATE TABLE map_distinct_table AS VALUES + (MAP {'k1': 1, 'k2': 2}, 'a', 1), + (MAP {'k1': 1, 'k2': 2}, 'a', 2), + (MAP {'k1': 1, 'k2': 2}, 'b', 3), + (MAP {'k1': 3}, 'a', 4), + (MAP {'k1': CAST(NULL AS BIGINT)}, 'b', 5); + +statement ok +INSERT INTO map_distinct_table VALUES (NULL, 'a', 6), (NULL, 'a', 7), (NULL, 'b', 8); + +# distinct on a map column collapses duplicate maps and duplicate NULLs +query ? rowsort +SELECT DISTINCT column1 FROM map_distinct_table; +---- +NULL +{k1: 1, k2: 2} +{k1: 3} +{k1: NULL} + +# exact reproducer from #15428: DISTINCT on a map column with LIMIT +query ? rowsort +SELECT DISTINCT column1 FROM map_distinct_table LIMIT 10; +---- +NULL +{k1: 1, k2: 2} +{k1: 3} +{k1: NULL} + +# distinct over a map column together with a scalar column +query ?T rowsort +SELECT DISTINCT column1, column2 FROM map_distinct_table; +---- +NULL a +NULL b +{k1: 1, k2: 2} a +{k1: 1, k2: 2} b +{k1: 3} a +{k1: NULL} b + +# group by a map column +query ?II rowsort +SELECT column1, COUNT(*), SUM(column3) FROM map_distinct_table GROUP BY column1; +---- +NULL 3 21 +{k1: 1, k2: 2} 3 6 +{k1: 3} 1 4 +{k1: NULL} 1 5 + +# group by a map column and a scalar column +query ?TI rowsort +SELECT column1, column2, COUNT(*) FROM map_distinct_table GROUP BY column1, column2; +---- +NULL a 2 +NULL b 1 +{k1: 1, k2: 2} a 2 +{k1: 1, k2: 2} b 1 +{k1: 3} a 1 +{k1: NULL} b 1 + +# empty maps compare equal under DISTINCT and are distinct from NULL +query ? rowsort +SELECT DISTINCT column1 FROM (VALUES (MAP {}), (MAP {}), (NULL)) t(column1); +---- +NULL +{} + +# HAVING clause with a map grouping key +query ?I rowsort +SELECT column1, COUNT(*) FROM map_distinct_table GROUP BY column1 HAVING COUNT(*) > 1; +---- +NULL 3 +{k1: 1, k2: 2} 3 + +# count and count distinct on a map column +query II +SELECT COUNT(column1), COUNT(DISTINCT column1) FROM map_distinct_table; +---- +5 3 + +# map column as input to an aggregate function +query T? +SELECT column2, array_agg(column1 ORDER BY column3) FROM map_distinct_table GROUP BY column2 ORDER BY column2; +---- +a [{k1: 1, k2: 2}, {k1: 1, k2: 2}, {k1: 3}, NULL, NULL] +b [{k1: 1, k2: 2}, {k1: NULL}, NULL] + +# UNION (distinct) on map columns +query ? +SELECT MAP {'a': 1} UNION SELECT MAP {'a': 1}; +---- +{a: 1} + +# unsorted maps are compared by entry order: maps with the same entries in a +# different order are treated as distinct values +query ? rowsort +SELECT DISTINCT column1 FROM (SELECT MAP {'k1': 1, 'k2': 2} AS column1 UNION ALL SELECT MAP {'k2': 2, 'k1': 1}); +---- +{k1: 1, k2: 2} +{k2: 2, k1: 1} + +statement ok +DROP TABLE map_distinct_table; + +# distinct / group by on map columns read from parquet +statement ok +CREATE EXTERNAL TABLE map_data +STORED AS PARQUET +LOCATION '../core/tests/data/parquet_map.parquet'; + +query I +SELECT COUNT(*) FROM (SELECT DISTINCT ints, strings FROM map_data); +---- +209 + +query TI rowsort +SELECT strings['method'] AS method, COUNT(*) FROM (SELECT DISTINCT strings FROM map_data) GROUP BY method; +---- +DELETE 24 +GET 27 +HEAD 33 +OPTION 29 +PATCH 30 +POST 41 +PUT 25 + +statement ok +DROP TABLE map_data; + +# map NULL value coercion in VALUES +# https://github.com/apache/datafusion/issues/23474 +# A bare NULL map value used to fail type unification across a VALUES list +# ("Inconsistent data type across values list") and required an explicit +# CAST(NULL AS ). The map value type now unifies with concrete value +# types following the same rules as scalar VALUES coercion. + +# concrete-typed row first, NULL-valued row second (the issue reproducer) +statement ok +CREATE TABLE map_null_concrete_first AS VALUES + (MAP {'k1': 1, 'k2': 2}), + (MAP {'k1': NULL}); + +# NULL must round-trip as NULL after coercion, not a default value +query ? rowsort +SELECT * FROM map_null_concrete_first; +---- +{k1: 1, k2: 2} +{k1: NULL} + +# the NULL value type is coerced to the concrete value type (Int64) +query T +SELECT arrow_typeof(column1) FROM map_null_concrete_first LIMIT 1; +---- +Map("entries": non-null Struct("key": non-null Utf8, "value": Int64), unsorted) + +statement ok +DROP TABLE map_null_concrete_first; + +# NULL-valued row first, concrete-typed row second (coercion is symmetric) +statement ok +CREATE TABLE map_null_first AS VALUES + (MAP {'k1': NULL}), + (MAP {'k1': 1, 'k2': 2}); + +query ? rowsort +SELECT * FROM map_null_first; +---- +{k1: 1, k2: 2} +{k1: NULL} + +statement ok +DROP TABLE map_null_first; + +# every row has a NULL value: succeeds and the value type stays Null +statement ok +CREATE TABLE map_all_null_values AS VALUES + (MAP {'k': NULL}), + (MAP {'k': NULL}); + +query ? rowsort +SELECT * FROM map_all_null_values; +---- +{k: NULL} +{k: NULL} + +query T +SELECT arrow_typeof(column1) FROM map_all_null_values LIMIT 1; +---- +Map("entries": non-null Struct("key": non-null Utf8, "value": Null), unsorted) + +statement ok +DROP TABLE map_all_null_values; + +# multiple keys where only one value is NULL +query ? rowsort +SELECT * FROM (VALUES + (MAP {'a': 1, 'b': NULL}), + (MAP {'a': 2, 'b': 3})) t(column1); +---- +{a: 1, b: NULL} +{a: 2, b: 3} + +# three rows with the NULL-valued row in the middle +query ? rowsort +SELECT * FROM (VALUES + (MAP {'k': 1}), + (MAP {'k': NULL}), + (MAP {'k': 2})) t(column1); +---- +{k: 1} +{k: 2} +{k: NULL} + +# numeric widening across a NULL-valued row follows the scalar rule +# (Int64 + Float64 -> Float64) +statement ok +CREATE TABLE map_null_widening AS VALUES + (MAP {'k': 1}), + (MAP {'k': NULL}), + (MAP {'k': 1.5}); + +query ? rowsort +SELECT * FROM map_null_widening; +---- +{k: 1.0} +{k: 1.5} +{k: NULL} + +query T +SELECT arrow_typeof(column1) FROM map_null_widening LIMIT 1; +---- +Map("entries": non-null Struct("key": non-null Utf8, "value": Float64), unsorted) + +statement ok +DROP TABLE map_null_widening; + +# incompatible concrete value types with a NULL-valued row in between still +# error; Int64/Utf8 follows the scalar VALUES rule (coerce to the numeric +# type, then fail to cast the non-numeric string) +query error Cast error: Cannot cast string 'hello' to value of Int64 type +SELECT * FROM (VALUES + (MAP {'k': 1}), + (MAP {'k': NULL}), + (MAP {'k': 'hello'})) t(column1); + +# NULL value type unification recurses into nested maps +query ? rowsort +SELECT * FROM (VALUES + (MAP {'outer': MAP {'inner': 1}}), + (MAP {'outer': MAP {'inner': NULL}})) t(column1); +---- +{outer: {inner: 1}} +{outer: {inner: NULL}} + +# INSERT INTO ... VALUES also accepts a NULL map value without a cast +statement ok +CREATE TABLE map_null_insert AS VALUES (MAP {'k1': 1, 'k2': 2}); + +statement ok +INSERT INTO map_null_insert VALUES (MAP {'k1': NULL}); + +query ? rowsort +SELECT * FROM map_null_insert; +---- +{k1: 1, k2: 2} +{k1: NULL} + +statement ok +DROP TABLE map_null_insert; diff --git a/datafusion/sqllogictest/test_files/math.slt b/datafusion/sqllogictest/test_files/math.slt index 475434883d315..b6bf51dd4799a 100644 --- a/datafusion/sqllogictest/test_files/math.slt +++ b/datafusion/sqllogictest/test_files/math.slt @@ -88,6 +88,102 @@ SELECT round(125.2345, -3), round(125.2345, -2), round(125.2345, -1), round(125. ---- 0 100 130 125 125 125.2 125.23 125.235 +# Round signed and unsigned integer scalar widths +query IIIIIIII +SELECT + round(arrow_cast('115', 'Int8'), -1), + round(arrow_cast('-115', 'Int16'), -1), + round(arrow_cast('115', 'Int32'), -1), + round(arrow_cast('-115', 'Int64'), -1), + round(arrow_cast('115', 'UInt8'), -1), + round(arrow_cast('115', 'UInt16'), -1), + round(arrow_cast('115', 'UInt32'), -1), + round(arrow_cast('115', 'UInt64'), -1); +---- +120 -120 120 -120 120 120 120 120 + +# Round signed and unsigned integer arrays, including null and oversized scales +query IIIIIIII +SELECT + round(arrow_cast(column1, 'Int8'), column2), + round(arrow_cast(column1, 'Int16'), column2), + round(arrow_cast(column1, 'Int32'), column2), + round(arrow_cast(column1, 'Int64'), column2), + round(arrow_cast(column1, 'UInt8'), column2), + round(arrow_cast(column1, 'UInt16'), column2), + round(arrow_cast(column1, 'UInt32'), column2), + round(arrow_cast(column1, 'UInt64'), column2) +FROM (VALUES ('115', -1), ('0', -1), (NULL, -20)) AS t(column1, column2); +---- +120 120 120 120 120 120 120 120 +0 0 0 0 0 0 0 0 +NULL NULL NULL NULL NULL NULL NULL NULL + +# Test columns without null +query I +SELECT + round(column1, column2) +FROM (VALUES (115, -20), (0, -1), (21, -1)) AS t(column1, column2); +---- +0 +0 +20 + +# Round all decimal widths as scalars +query RRRR +SELECT + round(arrow_cast('125.55', 'Decimal32(7,2)'), 1), + round(arrow_cast('-125.55', 'Decimal64(16,2)'), 1), + round(arrow_cast('125.55', 'Decimal128(30,2)'), 1), + round(arrow_cast('-125.55', 'Decimal256(40,2)'), 1); +---- +125.6 -125.6 125.6 -125.6 + +# Round all decimal widths as arrays with per-row decimal places +query RRRR +SELECT + round(arrow_cast(column1, 'Decimal32(7,2)'), column2), + round(arrow_cast(column1, 'Decimal64(16,2)'), column2), + round(arrow_cast(column1, 'Decimal128(30,2)'), column2), + round(arrow_cast(column1, 'Decimal256(40,2)'), column2) +FROM (VALUES ('125.55', 1), ('-125.55', 0), ('125.55', -1), (NULL, 1)) AS t(column1, column2); +---- +125.6 125.6 125.6 125.6 +-126 -126 -126 -126 +130 130 130 130 +NULL NULL NULL NULL + +# Float arrays with scalar and per-row decimal places +query RRRR +SELECT + round(arrow_cast(column1, 'Float32'), 1), + round(arrow_cast(column1, 'Float64'), 1), + round(arrow_cast(column1, 'Float32'), column2), + round(arrow_cast(column1, 'Float64'), column2) +FROM (VALUES ('125.55', 1), ('-125.55', 0), (NULL, -1)) AS t(column1, column2); +---- +125.6 125.6 125.6 125.6 +-125.6 -125.6 -126 -126 +NULL NULL NULL NULL + +# Null decimal places, invalid argument count/type, and out-of-range scale +query R +SELECT round(1.25, NULL); +---- +NULL + +query error DataFusion error: Error during planning: 'round' does not support zero arguments +SELECT round(); + +query error Error during planning: Internal error: Function 'round' failed to match any signature +SELECT round(1, 2, 3); + +query error Error during planning: Internal error: Function 'round' failed to match any signature +SELECT round('x'); + +query error round decimal_places 2147483648 is out of supported i32 range +SELECT round(1.25, 2147483648); + # atan2 query RRRRRRR SELECT atan2(2.0, 1.0), atan2(-2.0, 1.0), atan2(2.0, -1.0), atan2(-2.0, -1.0), atan2(NULL, 1.0), atan2(2.0, NULL), atan2(NULL, NULL); @@ -127,6 +223,92 @@ SELECT isnan(1::DECIMAL(10,2)), isnan(0::DECIMAL(10,2)), isnan(NULL::DECIMAL(10, ---- false false NULL false +# isnan: scalar values at the remaining numeric widths +query BBBBBBBBBB +SELECT + isnan(arrow_cast('NaN', 'Float16')), + isnan(arrow_cast('-1.5', 'Float16')), + isnan(arrow_cast('-128', 'Int8')), + isnan(arrow_cast('-32768', 'Int16')), + isnan(arrow_cast('-9223372036854775808', 'Int64')), + isnan(arrow_cast('65535', 'UInt16')), + isnan(arrow_cast('18446744073709551615', 'UInt64')), + isnan(arrow_cast('1.25', 'Decimal32(7,2)')), + isnan(arrow_cast('-12.34', 'Decimal64(16,2)')), + isnan(arrow_cast('0.00', 'Decimal256(40,2)')) +---- +true false false false false false false false false false + +# isnan: floating-point arrays, including infinities and nulls +query IBBB +SELECT id, + isnan(arrow_cast(v, 'Float16')), + isnan(arrow_cast(v, 'Float32')), + isnan(arrow_cast(v, 'Float64')) +FROM (VALUES + (1, 'NaN'), + (2, 'Infinity'), + (3, '-Infinity'), + (4, '0.0'), + (5, NULL) +) AS t(id, v) +ORDER BY id +---- +1 true true true +2 false false false +3 false false false +4 false false false +5 NULL NULL NULL + +# isnan: signed and unsigned integer arrays +query IBBBBBBBB +SELECT id, + isnan(arrow_cast(v, 'Int8')), + isnan(arrow_cast(v, 'Int16')), + isnan(arrow_cast(v, 'Int32')), + isnan(arrow_cast(v, 'Int64')), + isnan(arrow_cast(v, 'UInt8')), + isnan(arrow_cast(v, 'UInt16')), + isnan(arrow_cast(v, 'UInt32')), + isnan(arrow_cast(v, 'UInt64')) +FROM (VALUES (1, '0'), (2, '42'), (3, NULL)) AS t(id, v) +ORDER BY id +---- +1 false false false false false false false false +2 false false false false false false false false +3 NULL NULL NULL NULL NULL NULL NULL NULL + +# isnan: decimal arrays at every Arrow decimal width +query IBBBB +SELECT id, + isnan(arrow_cast(v, 'Decimal32(7,2)')), + isnan(arrow_cast(v, 'Decimal64(16,2)')), + isnan(arrow_cast(v, 'Decimal128(30,2)')), + isnan(arrow_cast(v, 'Decimal256(40,2)')) +FROM (VALUES (1, '0.00'), (2, '-12.34'), (3, NULL)) AS t(id, v) +ORDER BY id +---- +1 false false false false +2 false false false false +3 NULL NULL NULL NULL + +# isnan: an untyped all-null array +query B +SELECT isnan(v) FROM (VALUES (NULL), (NULL)) AS t(v) +---- +NULL +NULL + +# isnan: invalid argument count and type +statement error +SELECT isnan() + +statement error +SELECT isnan(1, 2) + +statement error +SELECT isnan('not numeric') + # iszero query BBBB SELECT iszero(1.0), iszero(0.0), iszero(-0.0), iszero(NULL) @@ -149,6 +331,86 @@ SELECT iszero(1::DECIMAL(10,2)), iszero(0::DECIMAL(10,2)), iszero(NULL::DECIMAL( ---- false true NULL false +# iszero: scalar boundary values at the remaining numeric widths +query BBBBBBBBBB +SELECT + iszero(arrow_cast(-0.0, 'Float16')), + iszero(arrow_cast('NaN', 'Float32')), + iszero(arrow_cast('-128', 'Int8')), + iszero(arrow_cast('-32768', 'Int16')), + iszero(arrow_cast('-9223372036854775808', 'Int64')), + iszero(arrow_cast('65535', 'UInt16')), + iszero(arrow_cast('18446744073709551615', 'UInt64')), + iszero(arrow_cast('0.00', 'Decimal32(7,2)')), + iszero(arrow_cast('-12.34', 'Decimal64(16,2)')), + iszero(arrow_cast('0.00', 'Decimal256(40,2)')) +---- +true false false false false false false true false true + +# iszero: signed integer arrays, including minimum values and nulls +query IBBBB +SELECT id, iszero(i8), iszero(i16), iszero(i32), iszero(i64) +FROM (VALUES + (1, 0::TINYINT, 0::SMALLINT, 0::INT, 0::BIGINT), + (2, arrow_cast('-128', 'Int8'), arrow_cast('-32768', 'Int16'), arrow_cast('-2147483648', 'Int32'), arrow_cast('-9223372036854775808', 'Int64')), + (3, NULL::TINYINT, NULL::SMALLINT, NULL::INT, NULL::BIGINT) +) AS t(id, i8, i16, i32, i64) +ORDER BY id +---- +1 true true true true +2 false false false false +3 NULL NULL NULL NULL + +# iszero: unsigned integer arrays +query IBBBB +SELECT id, iszero(u8), iszero(u16), iszero(u32), iszero(u64) +FROM (VALUES + (1, 0::TINYINT UNSIGNED, 0::SMALLINT UNSIGNED, 0::INT UNSIGNED, 0::BIGINT UNSIGNED), + (2, 255::TINYINT UNSIGNED, 65535::SMALLINT UNSIGNED, 4294967295::INT UNSIGNED, 4294967295::BIGINT UNSIGNED), + (3, NULL::TINYINT UNSIGNED, NULL::SMALLINT UNSIGNED, NULL::INT UNSIGNED, NULL::BIGINT UNSIGNED) +) AS t(id, u8, u16, u32, u64) +ORDER BY id +---- +1 true true true true +2 false false false false +3 NULL NULL NULL NULL + +# iszero: floating-point arrays, including signed zero, NaN, and nulls +query IBBB +SELECT id, + iszero(arrow_cast(v, 'Float16')), + iszero(arrow_cast(v, 'Float32')), + iszero(arrow_cast(v, 'Float64')) +FROM (VALUES (1, 0.0), (2, -0.0), (3, 'NaN'::DOUBLE), (4, -1.5), (5, NULL::DOUBLE)) AS t(id, v) +ORDER BY id +---- +1 true true true +2 true true true +3 false false false +4 false false false +5 NULL NULL NULL + +# iszero: decimal arrays at every Arrow decimal width +query IBBBB +SELECT id, + iszero(arrow_cast(v, 'Decimal32(7,2)')), + iszero(arrow_cast(v, 'Decimal64(16,2)')), + iszero(arrow_cast(v, 'Decimal128(30,2)')), + iszero(arrow_cast(v, 'Decimal256(40,2)')) +FROM (VALUES (1, '0.00'), (2, '-12.34'), (3, NULL)) AS t(id, v) +ORDER BY id +---- +1 true true true true +2 false false false false +3 NULL NULL NULL NULL + +# iszero: an untyped all-null array +query B +SELECT iszero(v) FROM (VALUES (NULL), (NULL)) AS t(v) +---- +NULL +NULL + # abs: empty argument statement error SELECT abs(); @@ -686,6 +948,38 @@ select gcd(-9223372036854775808, 0); query error DataFusion error: Arrow error: Compute error: Signed integer overflow in GCD\(0, \-9223372036854775808\) select gcd(0, -9223372036854775808); +# gcd decimal +query RT +select gcd(2::decimal(38, 0), 3::decimal(38, 0)), arrow_typeof(gcd(2::decimal(38, 0), 3::decimal(38, 0))); +---- +1 Decimal128(38, 0) + +query RT +select gcd(0::decimal(38, 0), 3::decimal(38, 0)), arrow_typeof(gcd(0::decimal(38, 0), 3::decimal(38, 0))); +---- +3 Decimal128(38, 0) + +query RT +select gcd(2, 3::decimal(38, 0)), arrow_typeof(gcd(2, 3::decimal(38, 0))); +---- +1 Decimal128(38, 0) + +query RR +select gcd(-15::decimal(38, 0), -3::decimal(38, 0)), gcd(-15::decimal(38, 0), 3::decimal(38, 0)); +---- +3 3 + +# non-whole number case +query RT +select gcd(15.3::decimal(38, 1), 2.9::decimal(38, 1)), arrow_typeof(gcd(15.3::decimal(38, 1), 2.9::decimal(38, 1))); +---- +0.1 Decimal128(38, 1) + +# both decimal arguments are coerced to widest - decimal(38, 5), return type is that as well +query RT +select gcd(15::decimal(30, 2), 3::decimal(38, 5)), arrow_typeof(gcd(15::decimal(30, 2), 3::decimal(38, 5))); +---- +3 Decimal128(38, 5) ## lcm @@ -727,6 +1021,28 @@ select lcm(1, -9223372036854775808); query error DataFusion error: Arrow error: Compute error: Signed integer overflow in LCM\(2, 9223372036854775803\) select lcm(2, 9223372036854775803); +# lcm decimal +query R +select lcm(2::decimal(38, 0), 3::decimal(38, 0)); +---- +6 + +query RT +select lcm(0::decimal(38, 0), 3::decimal(38, 0)), arrow_typeof(lcm(0::decimal(38, 0), 3::decimal(38, 0))); +---- +0 Decimal128(38, 0) + +query RT +select lcm(2, 3::decimal(38, 0)), arrow_typeof(lcm(2, 3::decimal(38, 0))); +---- +6 Decimal128(38, 0) + +# both decimal arguments are coerced to widest - decimal(38, 5), return type is that as well +query RT +select lcm(2::decimal(30, 2), 3::decimal(38, 5)), arrow_typeof(lcm(2::decimal(30, 2), 3::decimal(38, 5))); +---- +6 Decimal128(38, 5) + ## pow/power @@ -818,6 +1134,8 @@ from values 81 NULL +# There is no variant of `power` that accepts (Decimal, Decimal); type coercion +# casts both arguments to `Float64`, so the result is `Float64`. query RT rowsort select power(base::decimal(38, 0), exponent::decimal(38, 0)), @@ -830,19 +1148,19 @@ from values (2, 3), (3, 4) as t(base, exponent); ---- -0 Decimal128(38, 0) -1 Decimal128(38, 0) -4 Decimal128(38, 0) -625 Decimal128(38, 0) -8 Decimal128(38, 0) -81 Decimal128(38, 0) +0 Float64 +1 Float64 +4 Float64 +625 Float64 +8 Float64 +81 Float64 query RT select pow(2.5::decimal(2, 1), 4::bigint), arrow_typeof(pow(2.5::decimal(2, 1), 4::bigint)); ---- -39 Decimal128(2, 1) +39.0625 Float64 # factorial negative (PostgreSQL-compatible domain error) query error DataFusion error: Execution error: factorial of a negative number is undefined @@ -864,6 +1182,55 @@ logical_plan 02)--TableScan: aggregate_simple projection=[] physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/aggregate_simple.csv]]}, projection=[NULL as log(NULL,aggregate_simple.c2)], file_type=csv, has_header=true +# Simplification must preserve NULLs from nullable columns +query RRRRR rowsort +SELECT + log(a, 1), + log(a, a), + log(a, power(a, b)), + power(a, 0), + power(a, log(a, b)) +FROM (VALUES (NULL::double, 2.0::double), (2.0, 3.0)) AS t(a, b); +---- +0 1 3 1 3 +NULL NULL NULL NULL NULL + +# Nullable bases must remain in the optimized plan so they can propagate NULL +query TT +EXPLAIN SELECT + log(a, 1) AS l1, + log(a, a) AS la, + power(a, 0) AS p0, + power(a, log(a, b)) AS pl +FROM (VALUES (NULL::double, 2.0::double), (2.0, 3.0)) AS t(a, b); +---- +logical_plan +01)Projection: log(t.a, Float64(1)) AS l1, log(t.a, t.a) AS la, power(t.a, Float64(0)) AS p0, power(t.a, log(t.a, t.b)) AS pl +02)--SubqueryAlias: t +03)----Projection: column1 AS a, column2 AS b +04)------Values: (Float64(NULL) AS NULL, Float64(2)), (Float64(2), Float64(3)) +physical_plan +01)ProjectionExec: expr=[log(column1@0, 1) as l1, log(column1@0, column1@0) as la, power(column1@0, 0) as p0, power(column1@0, log(column1@0, column2@1)) as pl] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Non-nullable bases still use the existing simplifications +query TT +EXPLAIN SELECT + log(a, 1) AS l1, + log(a, a) AS la, + power(a, 0) AS p0, + power(a, log(a, b)) AS pl +FROM (VALUES (2.0::double, 3.0::double)) AS t(a, b); +---- +logical_plan +01)Projection: Float64(0) AS l1, Float64(1) AS la, Float64(1) AS p0, t.b AS pl +02)--SubqueryAlias: t +03)----Projection: column2 AS b +04)------Values: (Float64(2), Float64(3)) +physical_plan +01)ProjectionExec: expr=[0 as l1, 1 as la, 1 as p0, column2@1 as pl] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + # Float 16/32/64 for log query RT SELECT log(2.5, arrow_cast(10.9, 'Float16')), arrow_typeof(log(2.5, arrow_cast(10.9, 'Float16'))); @@ -897,6 +1264,59 @@ SELECT lcm(6, column1) FROM (VALUES (4), (9), (0)); 18 0 +query I +SELECT lcm(column1, column2) FROM (VALUES (0, 5), (3, 5), (25, 5), (-16, 5)); +---- +0 +15 +25 +80 + +query R +SELECT lcm(6, arrow_cast(column1, 'Decimal128(38,0)')) FROM (VALUES (4), (9), (0)); +---- +12 +18 +0 + +query R +SELECT lcm(arrow_cast(column1, 'Decimal32(7,0)'), arrow_cast(column2, 'Decimal32(7,0)')) FROM (VALUES (6, 4), (6, 9), (6, 0)); +---- +12 +18 +0 + +query R +SELECT lcm(arrow_cast(column1, 'Decimal64(16,0)'), arrow_cast(column2, 'Decimal64(16,0)')) FROM (VALUES (6, 4), (6, 9), (6, 0)); +---- +12 +18 +0 + +query R +SELECT lcm(arrow_cast(column1, 'Decimal128(38,0)'), arrow_cast(column2, 'Decimal128(38,0)')) FROM (VALUES (6, 4), (6, 9), (6, 0)); +---- +12 +18 +0 + +query R +SELECT lcm(arrow_cast(column1, 'Decimal256(40,0)'), arrow_cast(column2, 'Decimal256(40,0)')) FROM (VALUES (6, 4), (6, 9), (6, 0)); +---- +12 +18 +0 + +# invalid argument count and type +query error DataFusion error: +SELECT lcm(); + +query error DataFusion error: +SELECT lcm(1, 2, 3); + +query error DataFusion error: +SELECT lcm('x', 'y'); + # lcm array and scalar with nulls in the array query I SELECT lcm(column1, 5) FROM (VALUES (0), (NULL), (25)); @@ -940,6 +1360,84 @@ SELECT gcd(15, column1) FROM (VALUES (10), (25), (0)); 5 15 +query I +SELECT gcd(column1, column2) FROM (VALUES (8, 12), (18, 12), (0, 12), (-36, 12)); +---- +4 +6 +12 +12 + +query R +SELECT gcd(15, arrow_cast(column1, 'Decimal128(38,0)')) FROM (VALUES (10), (25), (0)); +---- +5 +5 +15 + +query R +SELECT gcd(arrow_cast(column1, 'Decimal128(38,0)'), arrow_cast(column2, 'Decimal128(38,0)')) FROM (VALUES (15, 10), (15, 25), (15, 0)); +---- +5 +5 +15 + +# gcd with the remaining decimal array widths +query R +SELECT gcd(arrow_cast(column1, 'Decimal32(7,0)'), arrow_cast(column2, 'Decimal32(7,0)')) FROM (VALUES (15, 10), (15, 25), (15, 0)); +---- +5 +5 +15 + +query R +SELECT gcd(arrow_cast(column1, 'Decimal64(16,0)'), arrow_cast(column2, 'Decimal64(16,0)')) FROM (VALUES (15, 10), (15, 25), (15, 0)); +---- +5 +5 +15 + +query R +SELECT gcd(arrow_cast(column1, 'Decimal256(40,0)'), arrow_cast(column2, 'Decimal256(40,0)')) FROM (VALUES (15, 10), (15, 25), (15, 0)); +---- +5 +5 +15 + +# gcd array with zero, minimum, and null scalars +query I +SELECT gcd(column1, 0) FROM (VALUES (1), (2), (0), (NULL)); +---- +1 +2 +0 +NULL + +query I +SELECT gcd(column1, -9223372036854775808) FROM (VALUES (1), (2), (NULL)); +---- +1 +2 +NULL + +query I +SELECT gcd(column1, NULL) FROM (VALUES (1), (2), (NULL)); +---- +NULL +NULL +NULL + +# invalid argument count and type +query error gcd function requires 2 arguments, got 0 +SELECT gcd(); + +query error gcd function requires 2 arguments, got 3 +SELECT gcd(1, 2, 3); + +query error Unsupported argument types Utf8 and Utf8 for function gcd +SELECT gcd('x', 'y'); + + # gcd array and scalar with nulls in the array query I SELECT gcd(column1, 12) FROM (VALUES (8), (NULL), (0), (-36)); diff --git a/datafusion/sqllogictest/test_files/merge_into.slt b/datafusion/sqllogictest/test_files/merge_into.slt new file mode 100644 index 0000000000000..f868bcbdc4862 --- /dev/null +++ b/datafusion/sqllogictest/test_files/merge_into.slt @@ -0,0 +1,248 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +## MERGE INTO Tests +## +## Note that MERGE INTO planning is supported, but the built-in MemTable does not +## (yet) support execution. These tests verify planning +########## + +statement ok +create table target(id int, val varchar, qty int); + +statement ok +insert into target values (1, 'foo', 100.0); + +statement ok +insert into target values (2, 'bar', 200.0); + +statement ok +insert into target values (3, 'baz', 300.0); + + +statement ok +create table source(id int, val varchar, is_active boolean); + +statement ok +insert into source values (2, 'xxxx', true); + +statement ok +insert into source values (4, 'yyyy', false); + + +########## +# Logical planning +########## + +query TT +explain merge into target using source on target.id = source.id +when matched then update set val = source.val +when not matched then insert (id, val) values (source.id, source.val); +---- +logical_plan +01)Dml: op=[MergeInto] table=[target] +02)--TableScan: source projection=[id, val, is_active] +physical_plan_error +01)MERGE INTO operation on table 'target' +02)caused by +03)This feature is not implemented: MERGE INTO not supported for Base table + +# Simple MATCHED DELETE +query TT +explain merge into target using source on target.id = source.id +when matched then delete; +---- +logical_plan +01)Dml: op=[MergeInto] table=[target] +02)--TableScan: source projection=[id, val, is_active] +physical_plan_error +01)MERGE INTO operation on table 'target' +02)caused by +03)This feature is not implemented: MERGE INTO not supported for Base table + +# Aliased target and source: alias is canonicalized to the table name +query TT +explain merge into target as t using source as s on t.id = s.id +when matched and s.is_active then update set val = s.val +when not matched by source then delete; +---- +logical_plan +01)Dml: op=[MergeInto] table=[target] +02)--SubqueryAlias: s +03)----TableScan: source projection=[id, val, is_active] +physical_plan_error +01)MERGE INTO operation on table 'target' +02)caused by +03)This feature is not implemented: MERGE INTO not supported for Base table + +# WHEN NOT MATCHED THEN DELETE is rejected by the parser (no target row exists); +query error DELETE is not allowed in a NOT MATCHED merge clause at Line: 2, Column: 23 +merge into target using source on target.id = source.id +when not matched then delete; + +# NOT MATCHED BY SOURCE +query TT +explain merge into target using source on target.id = source.id +when not matched by source then delete; +---- +logical_plan +01)Dml: op=[MergeInto] table=[target] +02)--TableScan: source projection=[id, val, is_active] +physical_plan_error +01)MERGE INTO operation on table 'target' +02)caused by +03)This feature is not implemented: MERGE INTO not supported for Base table + +# Subquery as the USING source +query TT +explain merge into target using (select id, max(val) as val from source group by id) as s +on target.id = s.id +when matched then update set val = s.val; +---- +logical_plan +01)Dml: op=[MergeInto] table=[target] +02)--SubqueryAlias: s +03)----Projection: source.id, max(source.val) AS val +04)------Aggregate: groupBy=[[source.id]], aggr=[[max(source.val)]] +05)--------TableScan: source projection=[id, val] +physical_plan_error +01)MERGE INTO operation on table 'target' +02)caused by +03)This feature is not implemented: MERGE INTO not supported for Base table + +# INSERT without an explicit column list requires values for all target columns +query TT +explain merge into target using source on target.id = source.id +when not matched then insert values (source.id, source.val, 0); +---- +logical_plan +01)Dml: op=[MergeInto] table=[target] +02)--TableScan: source projection=[id, val, is_active] +physical_plan_error +01)MERGE INTO operation on table 'target' +02)caused by +03)This feature is not implemented: MERGE INTO not supported for Base table + +# Execution fails: the default TableProvider does not implement merge_into +statement error +merge into target using source on target.id = source.id +when matched then delete; +---- +DataFusion error: MERGE INTO operation on table 'target' +caused by +This feature is not implemented: MERGE INTO not supported for Base table + + +########## +# Type coercion of ON / WHEN conditions +########## + +statement error DataFusion error: type_coercion\ncaused by\nError during planning: MERGE ON condition must be boolean type, but got Int64 +merge into target using source on 1 +when matched then delete; + +statement error DataFusion error: type_coercion\ncaused by\nError during planning: MERGE WHEN condition must be boolean type, but got Utf8 +merge into target using source on target.id = source.id +when matched and 'yes' then delete; + +########## +# Planning errors: invalid structure +########## + +statement error DataFusion error: Error during planning: MERGE INTO requires at least one WHEN clause +merge into target using source on target.id = source.id; + +statement error DataFusion error: Error during planning: Duplicate column 'val' in MERGE UPDATE +merge into target using source on target.id = source.id +when matched then update set val = source.val, val = 'x'; + +statement error DataFusion error: Error during planning: Duplicate column 'id' in MERGE INSERT +merge into target using source on target.id = source.id +when not matched then insert (id, ID) values (1, 2); + +statement error DataFusion error: Error during planning: MERGE INSERT has 2 column\(s\) but 1 value\(s\) +merge into target using source on target.id = source.id +when not matched then insert (id, val) values (source.id); + +statement error DataFusion error: Error during planning: MERGE INSERT has 3 column\(s\) but 2 value\(s\) +merge into target using source on target.id = source.id +when not matched then insert values (source.id, source.val); + +statement error DataFusion error: Error during planning: MERGE INSERT must have exactly one row of values +merge into target using source on target.id = source.id +when not matched then insert (id) values (1), (2); + +# Unknown column in UPDATE assignment +statement error DataFusion error: Schema error: No field named nonexistent. +merge into target using source on target.id = source.id +when matched then update set nonexistent = 1; + +# Unknown column in INSERT column list +statement error DataFusion error: Schema error: No field named nonexistent. +merge into target using source on target.id = source.id +when not matched then insert (nonexistent) values (1); + +# UPDATE assignment must reference the target table +statement error DataFusion error: Error during planning: MERGE assignment target 's.val' must reference target table 't' +merge into target as t using source as s on t.id = s.id +when matched then update set s.val = 'x'; + +########## +# Planning errors: qualifier and alias handling +########## + +# Source alias may not collide with the target table name when the target is aliased +statement error DataFusion error: Error during planning: MERGE source may not use the target table name 'target' as a qualifier while the target is aliased as 't'; use a different source alias +merge into target as t using source as target on t.id = target.id +when matched then delete; + +# Subqueries correlated to the target alias are not supported yet +statement error DataFusion error: This feature is not implemented: MERGE subqueries correlated to target alias 't' are not supported +merge into target as t using source as s +on exists (select 1 from source x where x.id = t.id) +when matched then delete; + +########## +# Planning errors: unsupported syntax +########## + +statement error DataFusion error: This feature is not implemented: MERGE target table modifiers are not supported +merge into target partition (p0) using source on target.id = source.id +when matched then delete; + +statement error DataFusion error: This feature is not implemented: MERGE target alias column lists are not supported +merge into target as t(a, b) using source as s on t.a = s.id +when matched then delete; + +statement error DataFusion error: This feature is not implemented: MERGE UPDATE WHERE predicates are not supported +merge into target using source on target.id = source.id +when matched then update set val = source.val where source.is_active; + +statement error DataFusion error: This feature is not implemented: MERGE INSERT WHERE predicates are not supported +merge into target using source on target.id = source.id +when not matched then insert (id) values (source.id) where source.is_active; + +statement error DataFusion error: This feature is not implemented: MERGE INSERT ROW is not supported +merge into target using source on target.id = source.id +when not matched then insert row; + +statement ok +drop table target; + +statement ok +drop table source; diff --git a/datafusion/sqllogictest/test_files/metadata.slt b/datafusion/sqllogictest/test_files/metadata.slt index 3e2a503e6b3fc..0fc74fa6cf602 100644 --- a/datafusion/sqllogictest/test_files/metadata.slt +++ b/datafusion/sqllogictest/test_files/metadata.slt @@ -520,3 +520,9 @@ NULL the id field statement ok drop table table_with_metadata; + +# Test that metadata on conflicting values raises an error. +# The larger_table has 10 values, smaller_tables 1 value and the fields of each table +# have conflicting metadata, same key different values See test:context.rs register_conflicting_metadata_tables +statement error DataFusion error: PhysicalOptimizer rule 'join_selection' failed\. Schema mismatch\.\ncaused by\nInternal error: Schema metadata mismatch: Expected original metadata: \{"metadata_key": "right"\}, got metadata: \{"metadata_key": "left"\} +select * from larger_table cross join smaller_table; diff --git a/datafusion/sqllogictest/test_files/monotonic_projection_test.slt b/datafusion/sqllogictest/test_files/monotonic_projection_test.slt index 7feefc169fcab..71e5fbc08e3eb 100644 --- a/datafusion/sqllogictest/test_files/monotonic_projection_test.slt +++ b/datafusion/sqllogictest/test_files/monotonic_projection_test.slt @@ -168,3 +168,150 @@ physical_plan 03)----ProjectionExec: expr=[CAST(a@0 + b@1 AS Int64) as sum_expr, a@0 as a, b@1 as b] 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true 05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b], output_ordering=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], file_type=csv, has_header=true + +# concat(a, b) is not lexicographically ordered just because a is ordered: +# "a" < "a0", but "a1" > "a01". The projected result still needs a sort. +query I +COPY ( + SELECT * FROM (VALUES ('a', '1'), ('a0', '1')) AS t(a, b) ORDER BY a +) TO 'test_files/scratch/monotonic_projection_test/concat_ordered.parquet'; +---- +2 + +statement ok +CREATE EXTERNAL TABLE concat_ordered (a VARCHAR, b VARCHAR) +STORED AS PARQUET +WITH ORDER (a) +WITH ORDER (b) +LOCATION 'test_files/scratch/monotonic_projection_test/concat_ordered.parquet'; + +query TT +EXPLAIN +SELECT concat(a, b) AS c +FROM concat_ordered +ORDER BY c; +---- +logical_plan +01)Sort: c ASC NULLS LAST +02)--Projection: concat(concat_ordered.a, concat_ordered.b) AS c +03)----TableScan: concat_ordered projection=[a, b] +physical_plan +01)SortExec: expr=[c@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/monotonic_projection_test/concat_ordered.parquet]]}, projection=[concat(a@0, b@1) as c], file_type=parquet + +query T +SELECT concat(a, b) AS c +FROM concat_ordered +ORDER BY c; +---- +a01 +a1 + +# An ordering on (c, a, b) does not imply an ordering on (a, b), even when +# FilterExec establishes c = concat(a, b). EnsureRequirements must retain the +# sort required by ORDER BY a, b. +query I +COPY ( + SELECT concat(a, b) AS c, a, b + FROM (VALUES ('a0', '1'), ('a', '1')) AS t(a, b) + ORDER BY c, a, b +) TO 'test_files/scratch/monotonic_projection_test/concat_equality_ordered.parquet'; +---- +2 + +statement ok +CREATE EXTERNAL TABLE concat_equality_ordered (c VARCHAR, a VARCHAR, b VARCHAR) +STORED AS PARQUET +WITH ORDER (c, a, b) +LOCATION 'test_files/scratch/monotonic_projection_test/concat_equality_ordered.parquet'; + +query TT +EXPLAIN +SELECT a, b +FROM concat_equality_ordered +WHERE c = concat(a, b) +ORDER BY a, b; +---- +logical_plan +01)Sort: concat_equality_ordered.a ASC NULLS LAST, concat_equality_ordered.b ASC NULLS LAST +02)--Projection: concat_equality_ordered.a, concat_equality_ordered.b +03)----Filter: concat_equality_ordered.c = concat(concat_equality_ordered.a, concat_equality_ordered.b) +04)------TableScan: concat_equality_ordered projection=[c, a, b], partial_filters=[concat_equality_ordered.c = concat(concat_equality_ordered.a, concat_equality_ordered.b)] +physical_plan +01)SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST] +02)--SortExec: expr=[a@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[true] +03)----FilterExec: c@0 = concat(a@1, b@2), projection=[a@1, b@2] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true +05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/monotonic_projection_test/concat_equality_ordered.parquet]]}, projection=[c, a, b], output_ordering=[c@0 ASC NULLS LAST, a@1 ASC NULLS LAST, b@2 ASC NULLS LAST], file_type=parquet, predicate=c@0 = concat(a@1, b@2) + +query TT +SELECT a, b +FROM concat_equality_ordered +WHERE c = concat(a, b) +ORDER BY a, b; +---- +a 1 +a0 1 + +# Test that precision-losing int-to-float casts do not invalidate suffix sort keys. +# +# When CAST(Int32 AS Float32) collapses distinct integer values (e.g., 16777216 and +# 16777217 both become 16777216.0), the suffix sort key (k) must still be sorted +# correctly. Before the fix, the optimizer incorrectly reused the pre-existing sort +# order and dropped the SortExec, producing wrong results. +# +# t1 is declared with a sort order, t2 is not — their results should be identical +# since CAST(v AS FLOAT) is not injective for 32-bit integers. +statement ok +CREATE EXTERNAL TABLE t1_int_float (k int, v int) +STORED AS CSV +WITH ORDER (v DESC, k DESC) +LOCATION '../core/tests/data/int_to_float_cast_precision.csv' +OPTIONS ('format.has_header' 'true'); + +statement ok +CREATE EXTERNAL TABLE t2_int_float (k int, v int) +STORED AS CSV +LOCATION '../core/tests/data/int_to_float_cast_precision.csv' +OPTIONS ('format.has_header' 'true'); + +# Both queries must return the same result: k=2 before k=1. +# (v_=16777216.0 for both rows; when tied on v_, DESC on k means k=2 comes first) +query IR +SELECT k, cast(v as float) v_ FROM t1_int_float ORDER BY v_ DESC, k DESC; +---- +2 16777216 +1 16777216 + +query IR +SELECT k, cast(v as float) v_ FROM t2_int_float ORDER BY v_ DESC, k DESC; +---- +2 16777216 +1 16777216 + +# Widening cast (Int32 -> Int64) is strictly 1-to-1, so the optimizer CAN +# legally reuse the pre-existing sort order and omit a SortExec. +statement ok +CREATE EXTERNAL TABLE t3_int_bigint (k int, v int) +STORED AS CSV +WITH ORDER (v DESC, k DESC) +LOCATION '../core/tests/data/int_to_float_cast_precision.csv' +OPTIONS ('format.has_header' 'true'); + +# CAST(Int32 AS BIGINT) is injective, so suffix key ordering is preserved. +query II +SELECT k, cast(v as bigint) v_ FROM t3_int_bigint ORDER BY v_ DESC, k DESC; +---- +1 16777217 +2 16777216 + +# Cleanup +statement ok +DROP TABLE t1_int_float; + +statement ok +DROP TABLE t2_int_float; + +statement ok +DROP TABLE t3_int_bigint; + diff --git a/datafusion/sqllogictest/test_files/negative_zero.slt b/datafusion/sqllogictest/test_files/negative_zero.slt new file mode 100644 index 0000000000000..8ea1122880e14 --- /dev/null +++ b/datafusion/sqllogictest/test_files/negative_zero.slt @@ -0,0 +1,231 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +## Negative Zero (-0.0) vs. Positive Zero (+0.0) Behavior +########## +# +# IEEE 754 specifies +0.0 == -0.0 (they compare equal). PostgreSQL follows +# this and treats them as the same value for DISTINCT, GROUP BY, UNION, +# INTERSECT, EXCEPT, and equality predicates. The bit patterns differ in +# the sign bit, so any code path that hashes / compares the raw bits (e.g. +# `f64::to_bits` or `f64::to_ne_bytes`) will treat them as distinct values +# and must be normalized before grouping / dedup. +# +# Note: the sqllogictest formatter renders both `-0.0` and `+0.0` as `0`, +# so the visible scalar values look identical in the expected output. The +# behavior is asserted via row counts and via auxiliary `1.0 / a` +# (`Infinity` vs `-Infinity`) columns that expose the sign. + +##### +## Equality and ordering predicates +##### + +# +0.0 == -0.0 is TRUE; +0.0 < -0.0 and +0.0 > -0.0 are both FALSE. +query BBB +SELECT 0.0 = -0.0 AS eq, 0.0 < -0.0 AS lt, 0.0 > -0.0 AS gt; +---- +true false false + +# 0.0 IS DISTINCT FROM -0.0 must be FALSE because the values are equal. +query B +SELECT 0.0 IS DISTINCT FROM -0.0 AS is_distinct; +---- +false + +##### +## SELECT DISTINCT with +0.0 / -0.0 (Float64) +##### + +# DISTINCT must collapse +0.0 and -0.0 into a single row. +query R rowsort +SELECT DISTINCT a +FROM (SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0); +---- +0 + +# Same query, with `1.0 / a` to expose the sign in the projection. The +# tuples `(+0.0, +Infinity)` and `(-0.0, -Infinity)` are not equal — the +# zero columns compare equal but `+Infinity != -Infinity` — so DISTINCT +# keeps both rows. PG returns the same two rows. +query RR rowsort +SELECT DISTINCT a, 1.0 / a AS inv +FROM (SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0); +---- +0 -Infinity +0 Infinity + +# COUNT(DISTINCT) over {+0.0, -0.0} must return 1. +query I +SELECT COUNT(DISTINCT a) +FROM (SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0); +---- +1 + +# GROUP BY must put +0.0 and -0.0 in the same group. +query RRI rowsort +SELECT a, 1.0 / a AS inv, COUNT(*) +FROM (SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0) +GROUP BY a; +---- +0 Infinity 3 + +# Multi-column DISTINCT; (+0.0, 1) and (-0.0, 1) must collapse. +query RI rowsort +SELECT DISTINCT a, b +FROM (SELECT 0.0 AS a, 1 AS b UNION ALL SELECT -0.0, 1 UNION ALL SELECT 0.0, 2); +---- +0 1 +0 2 + +##### +## SELECT DISTINCT with +0.0 / -0.0 (Float32 / REAL) +##### + +# DISTINCT for Float32: same collapse to a single row. +query R rowsort +SELECT DISTINCT a +FROM ( + SELECT arrow_cast(0.0, 'Float32') AS a + UNION ALL SELECT arrow_cast(-0.0, 'Float32') + UNION ALL SELECT arrow_cast(0.0, 'Float32') +); +---- +0 + +# COUNT(DISTINCT) for Float32: must be 1. +query I +SELECT COUNT(DISTINCT a) +FROM ( + SELECT arrow_cast(0.0, 'Float32') AS a + UNION ALL SELECT arrow_cast(-0.0, 'Float32') +); +---- +1 + +##### +## UNION (set semantics) with +0.0 / -0.0 +##### + +# UNION (DISTINCT) must collapse +0.0 / -0.0 into a single row. +query R rowsort +SELECT 0.0 AS a UNION SELECT -0.0 UNION SELECT 0.0; +---- +0 + +# UNION ALL preserves every input row regardless of sign — baseline. +query R rowsort +SELECT 0.0 AS a UNION ALL SELECT -0.0 UNION ALL SELECT 0.0; +---- +0 +0 +0 + +# UNION on Float32 must also collapse to a single row. +query R rowsort +SELECT arrow_cast(0.0, 'Float32') AS a +UNION +SELECT arrow_cast(-0.0, 'Float32'); +---- +0 + +##### +## INTERSECT with +0.0 / -0.0 +##### + +# INTERSECT treats +0.0 and -0.0 as equal — one matching row. +query R rowsort +SELECT 0.0 AS a INTERSECT SELECT -0.0; +---- +0 + +# INTERSECT ALL with multiplicities min(1,1) = 1. +query R rowsort +SELECT 0.0 AS a INTERSECT ALL SELECT -0.0; +---- +0 + +# INTERSECT for Float32: same matching behavior. +query R rowsort +SELECT arrow_cast(0.0, 'Float32') AS a +INTERSECT +SELECT arrow_cast(-0.0, 'Float32'); +---- +0 + +##### +## EXCEPT with +0.0 / -0.0 +##### + +# EXCEPT treats +0.0 and -0.0 as equal — zero rows after subtraction. +query R rowsort +SELECT 0.0 AS a EXCEPT SELECT -0.0; +---- + +# Reverse direction: also zero rows. +query R rowsort +SELECT -0.0 AS a EXCEPT SELECT 0.0; +---- + +# EXCEPT for Float32: zero rows. +query R rowsort +SELECT arrow_cast(0.0, 'Float32') AS a +EXCEPT +SELECT arrow_cast(-0.0, 'Float32'); +---- + +# EXCEPT ALL with matching multiplicities: zero rows. +query R rowsort +SELECT 0.0 AS a EXCEPT ALL SELECT -0.0; +---- + +##### +## INNER JOIN ON equality with +0.0 / -0.0 +##### + +# Equi-join on a = b matches +0.0 against -0.0. +query RR +SELECT t1.a, t2.b +FROM (SELECT 0.0 AS a) t1 +JOIN (SELECT -0.0 AS b) t2 ON t1.a = t2.b; +---- +0 0 + +# Sort-merge join must also match +0.0 against -0.0. SMJ builds equi-key +# matchers via `JoinKeyComparator`, which calls Arrow's `make_comparator` +# (IEEE 754 totalOrder); without normalization, +0.0 and -0.0 produce +# different orderings and miss the match. +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +query RR +SELECT t1.a, t2.b +FROM (SELECT 0.0 AS a) t1 +JOIN (SELECT -0.0 AS b) t2 ON t1.a = t2.b; +---- +0 0 + +# Float32 SMJ equi-join. +query RR +SELECT t1.a, t2.b +FROM (SELECT arrow_cast(0.0, 'Float32') AS a) t1 +JOIN (SELECT arrow_cast(-0.0, 'Float32') AS b) t2 ON t1.a = t2.b; +---- +0 0 + +statement ok +reset datafusion.optimizer.prefer_hash_join; diff --git a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt index 5907a85a9b923..bdb56cf22045a 100644 --- a/datafusion/sqllogictest/test_files/null_aware_anti_join.slt +++ b/datafusion/sqllogictest/test_files/null_aware_anti_join.slt @@ -53,12 +53,12 @@ query TT EXPLAIN SELECT * FROM outer_table WHERE id NOT IN (SELECT id FROM inner_table_no_null); ---- logical_plan -01)LeftAnti Join: outer_table.id = __correlated_sq_1.id +01)LeftAnti Join: outer_table.id = __correlated_sq_1.id null_aware 02)--TableScan: outer_table projection=[id, value] 03)--SubqueryAlias: __correlated_sq_1 04)----TableScan: inner_table_no_null projection=[id] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)] +01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)], null_aware 02)--DataSourceExec: partitions=1, partition_sizes=[1] 03)--DataSourceExec: partitions=1, partition_sizes=[1] @@ -70,6 +70,20 @@ query IT rowsort SELECT * FROM outer_table WHERE id NOT IN (SELECT id FROM inner_table_with_null); ---- +# Regression test + +statement ok +set datafusion.optimizer.filter_null_join_keys = true; + +# The subquery NULL must reach the join: every row's NOT IN is UNKNOWN or +# FALSE, so the result stays empty. +query IT rowsort +SELECT * FROM outer_table WHERE id NOT IN (SELECT id FROM inner_table_with_null); +---- + +statement ok +reset datafusion.optimizer.filter_null_join_keys; + # Verify the result is empty even though there are rows in outer_table # that don't match the non-NULL value (2) in the subquery. # This is correct null-aware behavior: if subquery contains NULL, result is unknown. @@ -193,12 +207,12 @@ query TT EXPLAIN SELECT * FROM outer_table WHERE id NOT IN (SELECT id FROM inner_table_with_null); ---- logical_plan -01)LeftAnti Join: outer_table.id = __correlated_sq_1.id +01)LeftAnti Join: outer_table.id = __correlated_sq_1.id null_aware 02)--TableScan: outer_table projection=[id, value] 03)--SubqueryAlias: __correlated_sq_1 04)----TableScan: inner_table_with_null projection=[id] physical_plan -01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)] +01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(id@0, id@0)], null_aware 02)--DataSourceExec: partitions=1, partition_sizes=[1] 03)--DataSourceExec: partitions=1, partition_sizes=[1] @@ -451,3 +465,113 @@ DROP TABLE customers_test; statement ok DROP TABLE all_null_banned; + +############# +## Test: dynamic filter pushdown must not drop inner (probe-side) NULLs. +## With join dynamic filter pushdown on, the build-side filter pushed to the probe scan would drop +## inner NULLs, but NOT IN three-valued logic needs them to collapse the result to zero rows. The +## in-memory VALUES scans above never apply the pushed filter, so this case needs a parquet scan. +############# + +statement ok +set datafusion.optimizer.enable_join_dynamic_filter_pushdown = true; + +# Row-level parquet filtering, so the pushed filter actually drops matching rows instead of only +# pruning row groups. Without this the single row group is read whole and the NULL never gets dropped. +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +statement ok +CREATE TABLE asa_outer(id INT) AS VALUES (1), (2), (3); + +statement ok +CREATE TABLE asa_inner(eid INT) AS VALUES (2), (NULL); + +query I +COPY asa_outer TO 'test_files/scratch/null_aware_anti_join/asa_outer.parquet' STORED AS PARQUET; +---- +3 + +query I +COPY asa_inner TO 'test_files/scratch/null_aware_anti_join/asa_inner.parquet' STORED AS PARQUET; +---- +2 + +statement ok +CREATE EXTERNAL TABLE asa_outer_parquet(id INT) +STORED AS PARQUET +LOCATION 'test_files/scratch/null_aware_anti_join/asa_outer.parquet'; + +statement ok +CREATE EXTERNAL TABLE asa_inner_parquet(eid INT) +STORED AS PARQUET +LOCATION 'test_files/scratch/null_aware_anti_join/asa_inner.parquet'; + +# Expected: zero rows. Before the fix the pushed dynamic filter dropped inner NULLs, so the join +# wrongly returned id = 1 and id = 3. +query I +SELECT id FROM asa_outer_parquet WHERE id NOT IN (SELECT eid FROM asa_inner_parquet) ORDER BY id; +---- + +statement ok +DROP TABLE asa_outer; + +statement ok +DROP TABLE asa_inner; + +statement ok +DROP TABLE asa_outer_parquet; + +statement ok +DROP TABLE asa_inner_parquet; + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + +statement ok +RESET datafusion.optimizer.enable_join_dynamic_filter_pushdown; + +############# +## Regression: null-aware NOT IN with an outer predicate on the join key +## +## `push_down_filter` used to infer the outer predicate `id > 5` onto the +## subquery side (as `eid > 5`), dropping the subquery's NULL row and wrongly +## returning outer rows. The subquery NULL must reach the join so that +## `NOT IN` stays UNKNOWN for every row. +############# + +statement ok +CREATE TABLE nai_outer(id INT) AS VALUES (3), (7); + +statement ok +CREATE TABLE nai_inner(id INT) AS VALUES (NULL); + +# Expected: zero rows (subquery contains NULL => NOT IN is UNKNOWN for all). +query I +SELECT id FROM nai_outer WHERE id > 5 AND id NOT IN (SELECT id FROM nai_inner) ORDER BY id; +---- + +# Same query under SortMergeJoin + multiple partitions: null-aware joins must +# be planned as a CollectLeft HashJoin, not a plain anti SortMergeJoin. +statement ok +SET datafusion.optimizer.prefer_hash_join = false; + +statement ok +SET datafusion.execution.target_partitions = 4; + +query I +SELECT id FROM nai_outer WHERE id NOT IN (SELECT id FROM nai_inner) ORDER BY id; +---- + +statement ok +SET datafusion.optimizer.prefer_hash_join = true; + +# The SLT runner sets target_partitions to 4, so restore that value explicitly. +statement ok +SET datafusion.execution.target_partitions = 4; + +statement ok +DROP TABLE nai_outer; + +statement ok +DROP TABLE nai_inner; diff --git a/datafusion/sqllogictest/test_files/operator.slt b/datafusion/sqllogictest/test_files/operator.slt index e50fa721c8850..926efe8fd56dc 100644 --- a/datafusion/sqllogictest/test_files/operator.slt +++ b/datafusion/sqllogictest/test_files/operator.slt @@ -287,7 +287,7 @@ EXPLAIN SELECT * FROM numeric_types WHERE int64 < 5 AND uint64 < 5 AND float64 < 5 AND decimal < 5; ---- physical_plan -01)FilterExec: int64@3 < 5 AND uint64@7 < 5 AND float64@9 < 5 AND decimal@10 < Some(500),5,2 +01)FilterExec: int64@3 < 5 AND uint64@7 < 5 AND float64@9 < 5 AND decimal@10 < 5.00 02)--DataSourceExec: partitions=1, partition_sizes=[1] ## < negative integer (expect no casts) @@ -296,7 +296,7 @@ EXPLAIN SELECT * FROM numeric_types WHERE int64 < -5 AND uint64 < -5 AND float64 < -5 AND decimal < -5; ---- physical_plan -01)FilterExec: int64@3 < -5 AND CAST(uint64@7 AS Decimal128(20, 0)) < Some(-5),20,0 AND float64@9 < -5 AND decimal@10 < Some(-500),5,2 +01)FilterExec: int64@3 < -5 AND CAST(uint64@7 AS Decimal128(20, 0)) < -5 AND float64@9 < -5 AND decimal@10 < -5.00 02)--DataSourceExec: partitions=1, partition_sizes=[1] ## < decimal (expect casts for integers to float) @@ -305,7 +305,7 @@ EXPLAIN SELECT * FROM numeric_types WHERE int64 < 5.1 AND uint64 < 5.1 AND float64 < 5.1 AND decimal < 5.1; ---- physical_plan -01)FilterExec: CAST(int64@3 AS Float64) < 5.1 AND CAST(uint64@7 AS Float64) < 5.1 AND float64@9 < 5.1 AND decimal@10 < Some(510),5,2 +01)FilterExec: CAST(int64@3 AS Float64) < 5.1 AND CAST(uint64@7 AS Float64) < 5.1 AND float64@9 < 5.1 AND decimal@10 < 5.10 02)--DataSourceExec: partitions=1, partition_sizes=[1] ## < negative decimal (expect casts for integers to float) @@ -314,7 +314,7 @@ EXPLAIN SELECT * FROM numeric_types WHERE int64 < -5.1 AND uint64 < -5.1 AND float64 < -5.1 AND decimal < -5.1; ---- physical_plan -01)FilterExec: CAST(int64@3 AS Float64) < -5.1 AND CAST(uint64@7 AS Float64) < -5.1 AND float64@9 < -5.1 AND decimal@10 < Some(-510),5,2 +01)FilterExec: CAST(int64@3 AS Float64) < -5.1 AND CAST(uint64@7 AS Float64) < -5.1 AND float64@9 < -5.1 AND decimal@10 < -5.10 02)--DataSourceExec: partitions=1, partition_sizes=[1] @@ -326,7 +326,7 @@ EXPLAIN SELECT * FROM numeric_types WHERE int64 = 5 AND uint64 = 5 AND float64 = 5 AND decimal = 5; ---- physical_plan -01)FilterExec: int64@3 = 5 AND uint64@7 = 5 AND float64@9 = 5 AND decimal@10 = Some(500),5,2 +01)FilterExec: int64@3 = 5 AND uint64@7 = 5 AND float64@9 = 5 AND decimal@10 = 5.00 02)--DataSourceExec: partitions=1, partition_sizes=[1] ## = negative integer (expect no casts) @@ -335,7 +335,7 @@ EXPLAIN SELECT * FROM numeric_types WHERE int64 = -5 AND uint64 = -5 AND float64 = -5 AND decimal = -5; ---- physical_plan -01)FilterExec: int64@3 = -5 AND CAST(uint64@7 AS Decimal128(20, 0)) = Some(-5),20,0 AND float64@9 = -5 AND decimal@10 = Some(-500),5,2 +01)FilterExec: int64@3 = -5 AND CAST(uint64@7 AS Decimal128(20, 0)) = -5 AND float64@9 = -5 AND decimal@10 = -5.00 02)--DataSourceExec: partitions=1, partition_sizes=[1] ## = decimal (expect casts for integers to float) @@ -344,7 +344,7 @@ EXPLAIN SELECT * FROM numeric_types WHERE int64 = 5.1 AND uint64 = 5.1 AND float64 = 5.1 AND decimal = 5.1; ---- physical_plan -01)FilterExec: CAST(int64@3 AS Float64) = 5.1 AND CAST(uint64@7 AS Float64) = 5.1 AND float64@9 = 5.1 AND decimal@10 = Some(510),5,2 +01)FilterExec: CAST(int64@3 AS Float64) = 5.1 AND CAST(uint64@7 AS Float64) = 5.1 AND float64@9 = 5.1 AND decimal@10 = 5.10 02)--DataSourceExec: partitions=1, partition_sizes=[1] ## = negative decimal (expect casts for integers to float) @@ -353,7 +353,7 @@ EXPLAIN SELECT * FROM numeric_types WHERE int64 = -5.1 AND uint64 = -5.1 AND float64 = -5.1 AND decimal = -5.1; ---- physical_plan -01)FilterExec: CAST(int64@3 AS Float64) = -5.1 AND CAST(uint64@7 AS Float64) = -5.1 AND float64@9 = -5.1 AND decimal@10 = Some(-510),5,2 +01)FilterExec: CAST(int64@3 AS Float64) = -5.1 AND CAST(uint64@7 AS Float64) = -5.1 AND float64@9 = -5.1 AND decimal@10 = -5.10 02)--DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt b/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt index da1e7de22bb7a..9df55512413f3 100644 --- a/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt +++ b/datafusion/sqllogictest/test_files/optimizer_group_by_constant.slt @@ -60,10 +60,9 @@ FROM test_table t group by 1, 2, 3 ---- logical_plan -01)Projection: Int64(123), Int64(456), Int64(789), count(Int64(1)), avg(t.c12) -02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1)), avg(t.c12)]] -03)----SubqueryAlias: t -04)------TableScan: test_table projection=[c12] +01)Aggregate: groupBy=[[Int64(123), Int64(456), Int64(789)]], aggr=[[count(Int64(1)), avg(t.c12)]] +02)--SubqueryAlias: t +03)----TableScan: test_table projection=[c12] query TT EXPLAIN @@ -72,8 +71,8 @@ FROM test_table t GROUP BY 1, 2 ---- logical_plan -01)Projection: Date32("2023-05-04") AS dt, Boolean(true) AS today_filter, count(Int64(1)) -02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] +01)Projection: to_date(Utf8("2023-05-04")) AS dt, date_part(Utf8("DAY"),now()) < Int64(1000) AS today_filter, count(Int64(1)) +02)--Aggregate: groupBy=[[Date32("2023-05-04") AS to_date(Utf8("2023-05-04")), Boolean(true) AS date_part(Utf8("DAY"),now()) < Int64(1000)]], aggr=[[count(Int64(1))]] 03)----SubqueryAlias: t 04)------TableScan: test_table projection=[] @@ -90,10 +89,9 @@ FROM test_table t GROUP BY 1 ---- logical_plan -01)Projection: Boolean(true) AS NOT date_part(Utf8("MONTH"),now()) BETWEEN Int64(50) AND Int64(60), count(Int64(1)) -02)--Aggregate: groupBy=[[]], aggr=[[count(Int64(1))]] -03)----SubqueryAlias: t -04)------TableScan: test_table projection=[] +01)Aggregate: groupBy=[[Boolean(true) AS NOT date_part(Utf8("MONTH"),now()) BETWEEN Int64(50) AND Int64(60)]], aggr=[[count(Int64(1))]] +02)--SubqueryAlias: t +03)----TableScan: test_table projection=[] query TT EXPLAIN @@ -119,7 +117,7 @@ logical_plan # Config reset -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/options.slt b/datafusion/sqllogictest/test_files/options.slt index 0d1583dbc0086..00d7664779a66 100644 --- a/datafusion/sqllogictest/test_files/options.slt +++ b/datafusion/sqllogictest/test_files/options.slt @@ -209,20 +209,20 @@ query RT select 123456789.0123456789012345678901234567890, arrow_typeof(123456789.0123456789012345678901234567890) ---- -123456789.012345678901 Decimal256(40, 31) +123456789.012345678901234567890123456789 Decimal256(40, 31) query RT select -123456789.0123456789012345678901234567890, arrow_typeof(-123456789.0123456789012345678901234567890) ---- --123456789.012345678901 Decimal256(40, 31) +-123456789.012345678901234567890123456789 Decimal256(40, 31) # max precision and scale of Decimal256 query RTRT select -1e-76, arrow_typeof(-1e-76), -1.234567e-70, arrow_typeof(-1.234567e-70) ---- -0 Decimal256(76, 76) 0 Decimal256(76, 76) +-0.0000000000000000000000000000000000000000000000000000000000000000000000000001 Decimal256(76, 76) -0.0000000000000000000000000000000000000000000000000000000000000000000001234567 Decimal256(76, 76) # Decimal256::MAX for nonnegative scale query RT @@ -243,13 +243,13 @@ query RTRT select 1e-38, arrow_typeof(1e-38), 1e-39, arrow_typeof(1e-39); ---- -0 Decimal128(38, 38) 0 Decimal256(39, 39) +0.00000000000000000000000000000000000001 Decimal128(38, 38) 0.000000000000000000000000000000000000001 Decimal256(39, 39) query RTRT select -1e-38, arrow_typeof(-1e-38), -1e-39, arrow_typeof(-1e-39); ---- -0 Decimal128(38, 38) 0 Decimal256(39, 39) +-0.00000000000000000000000000000000000001 Decimal128(38, 38) -0.000000000000000000000000000000000000001 Decimal256(39, 39) # unsupported precision query error Decimal precision 77 exceeds the maximum supported precision: 76 diff --git a/datafusion/sqllogictest/test_files/order.slt b/datafusion/sqllogictest/test_files/order.slt index ffd48d5996576..4b136d24b0751 100644 --- a/datafusion/sqllogictest/test_files/order.slt +++ b/datafusion/sqllogictest/test_files/order.slt @@ -272,8 +272,8 @@ logical_plan 04)------TableScan: aggregate_test_100 projection=[c2, c3] physical_plan 01)SortPreservingMergeExec: [c2@0 ASC NULLS LAST] -02)--SortExec: expr=[c2@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[c2@0 as c2, sum(aggregate_test_100.c3)@1 as total_sal] +02)--ProjectionExec: expr=[c2@0 as c2, sum(aggregate_test_100.c3)@1 as total_sal] +03)----SortExec: expr=[c2@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[c2@0 as c2], aggr=[sum(aggregate_test_100.c3)] 05)--------RepartitionExec: partitioning=Hash([c2@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[c2@0 as c2], aggr=[sum(aggregate_test_100.c3)] @@ -291,8 +291,8 @@ logical_plan 04)------TableScan: aggregate_test_100 projection=[c2, c3] physical_plan 01)SortPreservingMergeExec: [total_sal@1 ASC NULLS LAST, c2@0 ASC NULLS LAST] -02)--SortExec: expr=[total_sal@1 ASC NULLS LAST, c2@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[c2@0 as c2, sum(aggregate_test_100.c3)@1 as total_sal] +02)--ProjectionExec: expr=[c2@0 as c2, sum(aggregate_test_100.c3)@1 as total_sal] +03)----SortExec: expr=[sum(aggregate_test_100.c3)@1 ASC NULLS LAST, c2@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[c2@0 as c2], aggr=[sum(aggregate_test_100.c3)] 05)--------RepartitionExec: partitioning=Hash([c2@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[c2@0 as c2], aggr=[sum(aggregate_test_100.c3)] @@ -674,8 +674,8 @@ SELECT DISTINCT time as "first_seen" FROM t ORDER BY 1; statement ok drop table t; -# Create a table having 3 columns which are ordering equivalent by the source. In the next step, -# we will expect to observe the removed SortExec by propagating the orders across projection. +# Create a table with three independently ordered columns. Their sum is not +# necessarily ordered because integer addition can wrap. statement ok CREATE EXTERNAL TABLE multiple_ordered_table ( a0 INTEGER, @@ -702,9 +702,108 @@ logical_plan 03)----TableScan: multiple_ordered_table projection=[a, b, c] physical_plan 01)SortPreservingMergeExec: [result@0 ASC NULLS LAST] -02)--ProjectionExec: expr=[b@1 + a@0 + c@2 as result] -03)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true -04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST], [b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], file_type=csv, has_header=true +02)--SortExec: expr=[result@0 ASC NULLS LAST], preserve_partitioning=[true] +03)----ProjectionExec: expr=[b@1 + a@0 + c@2 as result] +04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true +05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a, b, c], output_orderings=[[a@0 ASC NULLS LAST], [b@1 ASC NULLS LAST], [c@2 ASC NULLS LAST]], file_type=csv, has_header=true + +statement ok +drop table multiple_ordered_table; + + +# Create a table having dependent sort order +statement ok +CREATE EXTERNAL TABLE multiple_ordered_table ( + a0 INTEGER, + a INTEGER, + b INTEGER, + c INTEGER, + d INTEGER +) +STORED AS CSV +WITH ORDER (a ASC, b ASC, c ASC) +LOCATION '../core/tests/data/window_2.csv' +OPTIONS ('format.has_header' 'true'); + +# Test without repartition so removal of sort is more apperant +statement ok +set datafusion.execution.target_partitions = 1; + +# A strictly order-preserving scalar function is one-to-one, so an ordering on +# its argument carries over to its result. `from_unixtime` reinterprets the +# input integer as a timestamp without changing the value, so the whole +# ordering is preserved and no SortExec is needed. +query TT +EXPLAIN SELECT from_unixtime(a) AS a_, from_unixtime(b) AS b_, from_unixtime(c) AS c_ +FROM multiple_ordered_table +ORDER BY a_, b_, c_; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST, b_ ASC NULLS LAST, c_ ASC NULLS LAST +02)--Projection: from_unixtime(CAST(multiple_ordered_table.a AS Int64)) AS a_, from_unixtime(CAST(multiple_ordered_table.b AS Int64)) AS b_, from_unixtime(CAST(multiple_ordered_table.c AS Int64)) AS c_ +03)----TableScan: multiple_ordered_table projection=[a, b, c] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[from_unixtime(CAST(a@1 AS Int64)) as a_, from_unixtime(CAST(b@2 AS Int64)) as b_, from_unixtime(CAST(c@3 AS Int64)) as c_], file_type=csv, has_header=true + +# Being one-to-one also justifies keeping the *suffix* sort keys: data sorted +# by [a, b] is also sorted by [from_unixtime(a), b], because rows with equal +# `a_` have equal `a`, within which `b` is already sorted. +query TT +EXPLAIN SELECT from_unixtime(a) AS a_, b +FROM multiple_ordered_table +ORDER BY a_, b; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST, multiple_ordered_table.b ASC NULLS LAST +02)--Projection: from_unixtime(CAST(multiple_ordered_table.a AS Int64)) AS a_, multiple_ordered_table.b +03)----TableScan: multiple_ordered_table projection=[a, b] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[from_unixtime(CAST(a@1 AS Int64)) as a_, b], file_type=csv, has_header=true + +# A widening CAST is one-to-one too: +query TT +EXPLAIN SELECT CAST(a AS BIGINT) AS a_, b +FROM multiple_ordered_table +ORDER BY a_, b; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST, multiple_ordered_table.b ASC NULLS LAST +02)--Projection: CAST(multiple_ordered_table.a AS Int64) AS a_, multiple_ordered_table.b +03)----TableScan: multiple_ordered_table projection=[a, b] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[CAST(a@1 AS Int64) as a_, b], file_type=csv, has_header=true + +# In contrast, a merely monotone (`preserves_lex_ordering`, but not strictly +# order-preserving) function such as floor() does NOT justify the suffix keys: +# in general floor() collapses distinct inputs into one output value, and `b` +# is not sorted within such a run, so a SortExec must remain. +query TT +EXPLAIN SELECT floor(CAST(a AS DOUBLE)) AS a_, b +FROM multiple_ordered_table +ORDER BY a_, b; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST, multiple_ordered_table.b ASC NULLS LAST +02)--Projection: floor(CAST(multiple_ordered_table.a AS Float64)) AS a_, multiple_ordered_table.b +03)----TableScan: multiple_ordered_table projection=[a, b] +physical_plan +01)SortExec: expr=[a_@0 ASC NULLS LAST, b@1 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[floor(CAST(a@1 AS Float64)) as a_, b], file_type=csv, has_header=true + +# Monotonicity alone is still enough when the expression is the *only* sort +# key, so here the SortExec is removed even though floor() is not strict: +query TT +EXPLAIN SELECT floor(CAST(a AS DOUBLE)) AS a_ +FROM multiple_ordered_table +ORDER BY a_; +---- +logical_plan +01)Sort: a_ ASC NULLS LAST +02)--Projection: floor(CAST(multiple_ordered_table.a AS Float64)) AS a_ +03)----TableScan: multiple_ordered_table projection=[a] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[floor(CAST(a@1 AS Float64)) as a_], file_type=csv, has_header=true + +# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# reset it explicitly. +statement ok +set datafusion.execution.target_partitions = 4; statement ok drop table multiple_ordered_table; @@ -1705,15 +1804,16 @@ EXPLAIN SELECT named_struct('sum', a + b) AS s FROM ordered ORDER BY s['sum']; ---- physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/data/composite_order.csv]]}, projection=[named_struct(sum, a@0 + b@1) as s], file_type=csv, has_header=true -# Wrapping a non-ordered column into a struct — SortExec required +# Wrapping a non-ordered column into a struct — SortExec required. # Reuses the `ordered` table above which has WITH ORDER (a + b). +# The simplifier resolves `get_field(named_struct(...), 'a')` so the sort key +# is not extracted into a separate scan projection column. query TT EXPLAIN SELECT named_struct('a', a, 'b', b) AS s FROM ordered ORDER BY s['a']; ---- physical_plan -01)ProjectionExec: expr=[s@0 as s] -02)--SortExec: expr=[__datafusion_extracted_1@1 ASC NULLS LAST], preserve_partitioning=[false] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/data/composite_order.csv]]}, projection=[named_struct(a, a@0, b, b@1) as s, get_field(named_struct(a, a@0, b, b@1), a) as __datafusion_extracted_1], file_type=csv, has_header=true +01)SortExec: expr=[get_field(s@0, a) ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/data/composite_order.csv]]}, projection=[named_struct(a, a@0, b, b@1) as s], file_type=csv, has_header=true # Simple column ordering tests using a table ordered by (a) statement ok @@ -1737,9 +1837,8 @@ query TT EXPLAIN SELECT named_struct('a', a, 'b', b) AS s FROM ordered_by_a ORDER BY s['b']; ---- physical_plan -01)ProjectionExec: expr=[s@0 as s] -02)--SortExec: expr=[__datafusion_extracted_1@1 ASC NULLS LAST], preserve_partitioning=[false] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/data/composite_order.csv]]}, projection=[named_struct(a, a@0, b, b@1) as s, get_field(named_struct(a, a@0, b, b@1), b) as __datafusion_extracted_1], file_type=csv, has_header=true +01)SortExec: expr=[get_field(s@0, b) ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/data/composite_order.csv]]}, projection=[named_struct(a, a@0, b, b@1) as s], file_type=csv, has_header=true # Mixed projection: top-level column alongside struct, order by struct field query TT @@ -1747,6 +1846,38 @@ EXPLAIN SELECT a, named_struct('a', a, 'b', b) AS s FROM ordered_by_a ORDER BY s ---- physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/data/composite_order.csv]]}, projection=[a, named_struct(a, a@0, b, b@1) as s], output_ordering=[a@0 ASC NULLS LAST], file_type=csv, has_header=true +query I +COPY ( + SELECT * FROM (VALUES (1, 1), (2, 3), (200, 10), (255, 10)) AS t(a, b) + ORDER BY a +) +TO 'test_files/scratch/order/uint8_overflow.csv' +OPTIONS ('format.has_header' 'false'); +---- +4 + +statement ok +CREATE EXTERNAL TABLE ordered_u8 ( + a TINYINT UNSIGNED NOT NULL, + b TINYINT UNSIGNED NOT NULL +) +STORED AS CSV +LOCATION 'test_files/scratch/order/uint8_overflow.csv' +OPTIONS ('format.has_header' 'false') +WITH ORDER (a ASC) +WITH ORDER (b ASC); + +query I +SELECT (a + b) AS result FROM ordered_u8 ORDER BY result ASC; +---- +2 +5 +9 +210 + +statement ok +DROP TABLE ordered_u8; + # Config reset statement ok reset datafusion.catalog.information_schema; @@ -1770,3 +1901,71 @@ reset datafusion.sql_parser.default_null_ordering; statement ok reset datafusion.sql_parser.dialect; + +# A global sort feeding a sink (CopyTo) must keep a leading key that is constant +# within each partition but differs across them ("a" is 2 on one union branch, +# 1 on the other). The merge above the union has to reorder rows across branches, +# so the physical plan must keep "a" in its ordering; dropping it (leaving only +# [b@1 ASC]) silently loses the global order under the sink. +statement ok +CREATE TABLE t2(b INT) AS VALUES (10), (20); + +query TT +EXPLAIN COPY ( + SELECT 2 AS a, b FROM t2 + UNION ALL + SELECT 1 AS a, b FROM t2 + ORDER BY a, b +) TO 'test_files/scratch/order/sort_key_sink.parquet'; +---- +logical_plan +01)CopyTo: format=parquet output_url=test_files/scratch/order/sort_key_sink.parquet options: () +02)--Sort: a ASC NULLS LAST, b ASC NULLS LAST +03)----Union +04)------Projection: Int64(2) AS a, t2.b +05)--------TableScan: t2 projection=[b] +06)------Projection: Int64(1) AS a, t2.b +07)--------TableScan: t2 projection=[b] +physical_plan +01)DataSinkExec: sink=ParquetSink(file_groups=[]) +02)--SortPreservingMergeExec: [a@0 ASC NULLS LAST, b@1 ASC NULLS LAST] +03)----UnionExec +04)------SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false] +05)--------ProjectionExec: expr=[2 as a, b@0 as b] +06)----------DataSourceExec: partitions=1, partition_sizes=[1] +07)------SortExec: expr=[b@1 ASC NULLS LAST], preserve_partitioning=[false] +08)--------ProjectionExec: expr=[1 as a, b@0 as b] +09)----------DataSourceExec: partitions=1, partition_sizes=[1] + +# Actually execute the COPY and verify the rows written to the file are in +# global (a, b) order, interleaving the two union branches. If "a" were +# dropped from the sort, the file would instead contain rows ordered only +# by "b" within each branch. +query I +COPY ( + SELECT 2 AS a, b FROM t2 + UNION ALL + SELECT 1 AS a, b FROM t2 + ORDER BY a, b +) TO 'test_files/scratch/order/sort_key_sink.parquet'; +---- +4 + +statement ok +CREATE EXTERNAL TABLE sort_key_sink STORED AS PARQUET +LOCATION 'test_files/scratch/order/sort_key_sink.parquet'; + +# Note: no ORDER BY here, so this checks the order rows were written in +query II +SELECT * FROM sort_key_sink; +---- +1 10 +1 20 +2 10 +2 20 + +statement ok +DROP TABLE sort_key_sink; + +statement ok +DROP TABLE t2; diff --git a/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt new file mode 100644 index 0000000000000..2c53c94144eb3 --- /dev/null +++ b/datafusion/sqllogictest/test_files/ordered_aggregate_spill.slt @@ -0,0 +1,253 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# End-to-end tests for ordered aggregation under finite memory. + +# Result set more than 100 lines will be hashed +hash-threshold 100 + +statement ok +SET datafusion.execution.target_partitions = 2 + +statement ok +SET datafusion.execution.batch_size = 128 + +statement ok +SET datafusion.optimizer.repartition_aggregations = true + +statement ok +SET datafusion.optimizer.prefer_existing_sort = true + +statement ok +SET datafusion.execution.enable_migration_aggregate = true + +statement ok +SET datafusion.runtime.memory_limit = '1M' + +# ================================================================================== +# Input is fully ordered by group keys (input order by (a,b), query is 'group by a,b') +# ================================================================================== + +# Fully ordered input uses the ordered partial and final streams without spill. +query TT +EXPLAIN ANALYZE +SELECT v1, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY v1 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,ordering_mode=Sorted, metrics=[spill_count=0,] +02)--RepartitionExec:preserve_order=true +03)----AggregateExec: mode=Partial,ordering_mode=Sorted, metrics=[spill_count=0,] + + +query II rowsort +SELECT v1, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY v1 +---- +40002 values hashing to 34c2b23730596cbd2489ed4a627c17d7 + +# The same fully ordered query cannot spill and reports OOM under tighter memory. +statement ok +SET datafusion.runtime.memory_limit = '1K' + +query error Resources exhausted +SELECT v1, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY v1 + +# ================================================================================== +# Input is partially ordered by group keys (input order by (a), query is 'group by a,b') +# +# Try different memory limits, ensure result is the same, but spill count differ + +# HACK: check `spilled_bytes=x KB` to ensure it has spilled. If it has not spilled, +# the it shows `spilled_bytes = 0B`. Should better check spill_count, but it's not +# stable due to ordered hash repartition, and `sqllogictest` don't support regex. +# ================================================================================== + +statement ok +SET datafusion.runtime.memory_limit = '2M' + +statement ok +SET datafusion.optimizer.enable_round_robin_repartition = false + +# Round 1: The input is partially ordered and does not spill with a 2 MB limit. +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,aggr=[sum(t1.v1 * Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] +02)--RepartitionExec:input_partitions=1, maintains_sort_order=true +03)----AggregateExec: mode=Partial,ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +# All rounds should have the same result hash +query III rowsort +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12 + +# Round 2: The same query spills five times with a 600 KB limit. +statement ok +SET datafusion.runtime.memory_limit = '600K' + +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,aggr=[sum(t1.v1 * Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spilled_bytes= KB,] +02)--RepartitionExec:input_partitions=1, maintains_sort_order=true +03)----AggregateExec: mode=Partial,ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +# All rounds should have the same result hash +query III rowsort +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12 + +# Round 3: The same query spills six times with a 500 KB limit. +statement ok +SET datafusion.runtime.memory_limit = '500K' + +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,aggr=[sum(t1.v1 * Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spilled_bytes= KB,] +02)--RepartitionExec:input_partitions=1, maintains_sort_order=true +03)----AggregateExec: mode=Partial,ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +# All rounds should have the same result hash +query III rowsort +SELECT round(v1, -4), v1 % 5000, sum(v1 * 2) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +45000 values hashing to e6ece4b4b86e6152a1c785e90ab5ba12 + +# Exercise the same spill path with a variable-width string aggregate state in +# the spilled payload. Keep one partial input partition so memory pressure is on +# the ordered aggregate rather than a repartition merge. +statement ok +SET datafusion.runtime.memory_limit = '600K' + +# Ensures final aggregate has spill_count > 0 +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, + sum(v1 * 2), min(CAST(v1 % 2 AS VARCHAR)) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=FinalPartitioned,aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spilled_bytes=KB,] +02)--RepartitionExec:input_partitions=1, maintains_sort_order=true +03)----AggregateExec: mode=Partial,aggr=[sum(t1.v1 * Int64(2)), min(t1.v1 % Int64(2))], ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +# ================================================================================== +# Single mode: with one partition the whole aggregation runs in a `Single` mode +# AggregateExec. min() keeps one intermediate state and avg() keeps two (sum + +# count), so both single- and multi-state accumulators are spilled and merged. +# ================================================================================== + +statement ok +SET datafusion.execution.target_partitions = 1 + +# Reference round: enough memory to aggregate without spilling. +statement ok +SET datafusion.runtime.memory_limit = '10M' + +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, min(v1 * 2), avg(v1) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=Single,aggr=[min(t1.v1 * Int64(2)), avg(t1.v1)], ordering_mode=PartiallySorted([0]), metrics=[spill_count=0,] + + +query IIIR rowsort +SELECT round(v1, -4), v1 % 5000, min(v1 * 2), avg(v1) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +60000 values hashing to 872df6cefd51f81820fc5c6e5d7480df + +# Spilling round: the same query under a 600 KB limit must spill. +statement ok +SET datafusion.runtime.memory_limit = '600K' + +query TT +EXPLAIN ANALYZE +SELECT round(v1, -4), v1 % 5000, min(v1 * 2), avg(v1) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +Plan with Metrics +01)AggregateExec: mode=Single,aggr=[min(t1.v1 * Int64(2)), avg(t1.v1)], ordering_mode=PartiallySorted([0]), metrics=[spilled_bytes= KB,] + + +# Same result hash as the no-spill round above +query IIIR rowsort +SELECT round(v1, -4), v1 % 5000, min(v1 * 2), avg(v1) +FROM generate_series(20000) AS t1(v1) +GROUP BY round(v1, -4), v1 % 5000 +---- +60000 values hashing to 872df6cefd51f81820fc5c6e5d7480df + +statement ok +RESET datafusion.runtime.memory_limit + +statement ok +RESET datafusion.optimizer.enable_round_robin_repartition + +statement ok +RESET datafusion.execution.enable_migration_aggregate + +statement ok +RESET datafusion.optimizer.prefer_existing_sort + +statement ok +RESET datafusion.optimizer.repartition_aggregations + +statement ok +RESET datafusion.execution.batch_size + +statement ok +SET datafusion.execution.target_partitions = 4 + +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema diff --git a/datafusion/sqllogictest/test_files/parquet_cdc.slt b/datafusion/sqllogictest/test_files/parquet_cdc.slt index f87f05af74a0c..bc9b3aeaeae07 100644 --- a/datafusion/sqllogictest/test_files/parquet_cdc.slt +++ b/datafusion/sqllogictest/test_files/parquet_cdc.slt @@ -28,14 +28,15 @@ CREATE TABLE cdc_source AS VALUES (5, 'eve', 500.99) # -# Test 1: Enable CDC with 'true' (uses default options) +# Test 1: Enable CDC with the explicit `content_defined_chunking.enabled` key +# (uses default chunking parameters). # query I COPY cdc_source TO 'test_files/scratch/parquet_cdc/enabled_true/' STORED AS PARQUET OPTIONS ( - 'format.use_content_defined_chunking' 'true' + 'format.content_defined_chunking.enabled' 'true' ) ---- 5 @@ -68,15 +69,14 @@ SELECT SUM(column3) FROM cdc_enabled_true_read 1502.49 # -# Test 2: Disable CDC with 'false' (same as default behavior) +# Test 2: CDC is disabled by default (no content_defined_chunking options set). +# It can also be turned off explicitly with +# `content_defined_chunking.enabled` = 'false'. # query I COPY cdc_source TO 'test_files/scratch/parquet_cdc/disabled_false/' STORED AS PARQUET -OPTIONS ( - 'format.use_content_defined_chunking' 'false' -) ---- 5 @@ -95,16 +95,17 @@ SELECT * FROM cdc_disabled_false_read 5 eve 500.99 # -# Test 3: Enable CDC with custom sub-field options +# Test 3: Enable CDC with custom chunking parameters # query I COPY cdc_source TO 'test_files/scratch/parquet_cdc/custom_chunks/' STORED AS PARQUET OPTIONS ( - 'format.use_content_defined_chunking.min_chunk_size' '1024', - 'format.use_content_defined_chunking.max_chunk_size' '4096', - 'format.use_content_defined_chunking.norm_level' '1' + 'format.content_defined_chunking.enabled' 'true', + 'format.content_defined_chunking.min_chunk_size' '1024', + 'format.content_defined_chunking.max_chunk_size' '4096', + 'format.content_defined_chunking.norm_level' '1' ) ---- 5 @@ -135,7 +136,7 @@ CREATE EXTERNAL TABLE cdc_external_write ( ) STORED AS PARQUET LOCATION 'test_files/scratch/parquet_cdc/external_table/' OPTIONS ( - 'format.use_content_defined_chunking' 'true' + 'format.content_defined_chunking.enabled' 'true' ) query I @@ -169,7 +170,7 @@ query I COPY cdc_large_source TO 'test_files/scratch/parquet_cdc/large/' STORED AS PARQUET OPTIONS ( - 'format.use_content_defined_chunking' 'true' + 'format.content_defined_chunking.enabled' 'true' ) ---- 1000 @@ -213,7 +214,7 @@ query I COPY cdc_types_source TO 'test_files/scratch/parquet_cdc/types/' STORED AS PARQUET OPTIONS ( - 'format.use_content_defined_chunking' 'true' + 'format.content_defined_chunking.enabled' 'true' ) ---- 3 diff --git a/datafusion/sqllogictest/test_files/parquet_cdc_config.slt b/datafusion/sqllogictest/test_files/parquet_cdc_config.slt new file mode 100644 index 0000000000000..2e2b3a5d2ca0b --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_cdc_config.slt @@ -0,0 +1,64 @@ +# Content-defined chunking (CDC) config resolution. +# +# CDC is a plain CdcOptions struct with an explicit `enabled` flag, so toggling +# `content_defined_chunking.enabled` is independent of the chunking parameters +# and of the order in which keys are set. There is no bare boolean form. + +statement ok +SET datafusion.catalog.information_schema = true + +# Disabled by default: enabled=false, parameters at their defaults. +query TT rowsort +SELECT name, value FROM information_schema.df_settings +WHERE name LIKE '%content_defined_chunking%' +---- +datafusion.execution.parquet.content_defined_chunking.enabled false +datafusion.execution.parquet.content_defined_chunking.max_chunk_size 1048576 +datafusion.execution.parquet.content_defined_chunking.min_chunk_size 262144 +datafusion.execution.parquet.content_defined_chunking.norm_level 0 + +# Setting a parameter does NOT enable CDC: `enabled` stays false. +statement ok +SET datafusion.execution.parquet.content_defined_chunking.min_chunk_size = 2048 + +query TT rowsort +SELECT name, value FROM information_schema.df_settings +WHERE name LIKE '%content_defined_chunking%' +---- +datafusion.execution.parquet.content_defined_chunking.enabled false +datafusion.execution.parquet.content_defined_chunking.max_chunk_size 1048576 +datafusion.execution.parquet.content_defined_chunking.min_chunk_size 2048 +datafusion.execution.parquet.content_defined_chunking.norm_level 0 + +# Enabling is explicit and independent of the parameters already set. +statement ok +SET datafusion.execution.parquet.content_defined_chunking.enabled = true + +query TT rowsort +SELECT name, value FROM information_schema.df_settings +WHERE name LIKE '%content_defined_chunking%' +---- +datafusion.execution.parquet.content_defined_chunking.enabled true +datafusion.execution.parquet.content_defined_chunking.max_chunk_size 1048576 +datafusion.execution.parquet.content_defined_chunking.min_chunk_size 2048 +datafusion.execution.parquet.content_defined_chunking.norm_level 0 + +# Disabling only flips the flag; the parameters are left untouched. +statement ok +SET datafusion.execution.parquet.content_defined_chunking.enabled = false + +query TT rowsort +SELECT name, value FROM information_schema.df_settings +WHERE name LIKE '%content_defined_chunking%' +---- +datafusion.execution.parquet.content_defined_chunking.enabled false +datafusion.execution.parquet.content_defined_chunking.max_chunk_size 1048576 +datafusion.execution.parquet.content_defined_chunking.min_chunk_size 2048 +datafusion.execution.parquet.content_defined_chunking.norm_level 0 + +# Restore defaults so the harness does not see modified configuration. +statement ok +SET datafusion.execution.parquet.content_defined_chunking.min_chunk_size = 262144 + +statement ok +SET datafusion.catalog.information_schema = false diff --git a/datafusion/sqllogictest/test_files/parquet_max_row_group_bytes.slt b/datafusion/sqllogictest/test_files/parquet_max_row_group_bytes.slt new file mode 100644 index 0000000000000..8de83329ae073 --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_max_row_group_bytes.slt @@ -0,0 +1,186 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# End-to-end tests for the `max_row_group_bytes` Parquet writer option: +# write Parquet files with the option set, then read them back to confirm +# the option is wired through from config to the writer. +# See datafusion/common/src/config.rs for the option definition. + +statement ok +CREATE TABLE source_table(id INT, name VARCHAR) AS VALUES +(1, 'one'), (2, 'two'), (3, 'three'), (4, 'four'), (5, 'five'); + +# Write with max_row_group_bytes set via COPY format options. +query I +COPY source_table +TO 'test_files/scratch/parquet_max_row_group_bytes/copy_options/' +STORED AS PARQUET +OPTIONS ('format.max_row_group_bytes' 1024); +---- +5 + +statement ok +CREATE EXTERNAL TABLE readback_copy_options +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_max_row_group_bytes/copy_options/'; + +query IT +SELECT id, name FROM readback_copy_options ORDER BY id; +---- +1 one +2 two +3 three +4 four +5 five + +# The option also applies when set via the session config (not just COPY OPTIONS). +statement ok +SET datafusion.execution.parquet.max_row_group_bytes = 2048; + +query I +COPY source_table +TO 'test_files/scratch/parquet_max_row_group_bytes/session_config/' +STORED AS PARQUET; +---- +5 + +statement ok +RESET datafusion.execution.parquet.max_row_group_bytes; + +statement ok +CREATE EXTERNAL TABLE readback_session_config +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_max_row_group_bytes/session_config/'; + +query IT +SELECT id, name FROM readback_session_config ORDER BY id; +---- +1 one +2 two +3 three +4 four +5 five + +# A zero byte limit is rejected with a clear configuration error. +query error DataFusion error: Invalid or Unsupported Configuration: max_row_group_bytes must be greater than 0 +COPY source_table +TO 'test_files/scratch/parquet_max_row_group_bytes/invalid/' +STORED AS PARQUET +OPTIONS ('format.max_row_group_bytes' 0); + +# ----------------------------------------------------------------------------- +# Row-group-count verification via EXPLAIN ANALYZE. +# +# `row_groups_pruned_statistics=N total` reports the number of row groups in the +# written file, so it lets us confirm that `max_row_group_bytes` actually +# changes how the writer splits row groups, and that combining it with +# `max_row_group_size` flushes on whichever limit is reached first. +# +# NOTE: byte-based flushing is currently honored only by the single-threaded +# Parquet writer (`AsyncArrowWriter`/`ArrowWriter`), which encodes inline and +# can therefore observe the in-progress row group's encoded size. The +# multi-threaded (parallel) writer decides row-group boundaries by row count +# only and ignores `max_row_group_bytes`, so these cases force the +# single-threaded path with `allow_single_file_parallelism = false`. Extending +# the parallel writer to honor the byte limit is a follow-up change. +# ----------------------------------------------------------------------------- + +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.execution.minimum_parallel_output_files = 1; + +statement ok +set datafusion.execution.batch_size = 1024; + +statement ok +set datafusion.execution.parquet.allow_single_file_parallelism = false; + +# Row-count limit only: 4096 rows with max_row_group_size = 1000 -> 5 row groups +# (four full groups of 1000 plus a remainder of 96). +statement ok +COPY (SELECT value AS id FROM range(0, 4096)) +TO 'test_files/scratch/parquet_max_row_group_bytes/rg_count_size_only/' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' 1000); + +statement ok +CREATE EXTERNAL TABLE rg_count_size_only +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_max_row_group_bytes/rg_count_size_only/'; + +query TT +EXPLAIN ANALYZE SELECT * FROM rg_count_size_only WHERE id >= 0; +---- +Plan with Metrics +01)FilterExec: id@0 >= 0 +02)--DataSourceExec: row_groups_pruned_statistics=5 total + +# Both limits set: the byte limit also flushes the sub-1000-row remainders that +# the row-count limit would otherwise carry into the next batch, so the file is +# split more finely -> 8 row groups (whichever limit is reached first). +statement ok +COPY (SELECT value AS id FROM range(0, 4096)) +TO 'test_files/scratch/parquet_max_row_group_bytes/rg_count_size_and_bytes/' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' 1000, 'format.max_row_group_bytes' 1); + +statement ok +CREATE EXTERNAL TABLE rg_count_size_and_bytes +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_max_row_group_bytes/rg_count_size_and_bytes/'; + +query TT +EXPLAIN ANALYZE SELECT * FROM rg_count_size_and_bytes WHERE id >= 0; +---- +Plan with Metrics +01)FilterExec: id@0 >= 0 +02)--DataSourceExec: row_groups_pruned_statistics=8 total + +# Byte limit drives alone: the row-count limit is far larger than the data, so +# only the byte limit splits. Each 1024-row batch fills a fresh (empty) row +# group, which is never split mid-batch -> 4 row groups. +statement ok +COPY (SELECT value AS id FROM range(0, 4096)) +TO 'test_files/scratch/parquet_max_row_group_bytes/rg_count_bytes_only/' +STORED AS PARQUET +OPTIONS ('format.max_row_group_size' 100000, 'format.max_row_group_bytes' 1); + +statement ok +CREATE EXTERNAL TABLE rg_count_bytes_only +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_max_row_group_bytes/rg_count_bytes_only/'; + +query TT +EXPLAIN ANALYZE SELECT * FROM rg_count_bytes_only WHERE id >= 0; +---- +Plan with Metrics +01)FilterExec: id@0 >= 0 +02)--DataSourceExec: row_groups_pruned_statistics=4 total + +statement ok +reset datafusion.execution.parquet.allow_single_file_parallelism; + +statement ok +reset datafusion.execution.batch_size; + +statement ok +reset datafusion.execution.minimum_parallel_output_files; + +statement ok +set datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/parquet_metadata_functions.slt b/datafusion/sqllogictest/test_files/parquet_metadata_functions.slt new file mode 100644 index 0000000000000..25a3c4eb4c6fa --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_metadata_functions.slt @@ -0,0 +1,84 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Test for Parquet scans with a mix of metadata functions + +statement ok +COPY (VALUES (10), (20), (30)) +TO 'test_files/scratch/parquet_metadata_functions/first.parquet' +STORED AS PARQUET; + +statement ok +COPY (VALUES (40), (50), (60)) +TO 'test_files/scratch/parquet_metadata_functions/second.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE test_table(column1 int) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_metadata_functions/'; + +query TII rowsort +SELECT input_file_name(), file_row_index(), column1 +FROM test_table +---- +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/first.parquet 0 10 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/first.parquet 1 20 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/first.parquet 2 30 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/second.parquet 0 40 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/second.parquet 1 50 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/second.parquet 2 60 + +query TT +EXPLAIN SELECT input_file_name(), file_row_index(), column1 +FROM test_table +---- +logical_plan +01)Projection: input_file_name(), file_row_index(), test_table.column1 +02)--TableScan: test_table projection=[column1] +physical_plan DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/first.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/second.parquet]]}, projection=[input_file_name() as input_file_name(), CAST(__datafusion_file_row_index@1 AS Int64) as file_row_index(), column1], file_type=parquet + + +# Make sure it also behaves consistently regardless of filter pushdown + +statement ok +SET datafusion.execution.parquet.pushdown_filters = false; + +query TII rowsort +SELECT input_file_name(), file_row_index(), column1 +FROM test_table +WHERE file_row_index() = 2 AND input_file_name() LIKE '%parquet'; +---- +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/first.parquet 2 30 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/second.parquet 2 60 + +statement ok +SET datafusion.execution.parquet.pushdown_filters = true; + +query TII rowsort +SELECT input_file_name(), file_row_index(), column1 +FROM test_table +WHERE file_row_index() = 2 AND input_file_name() LIKE '%parquet'; +---- +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/first.parquet 2 30 +WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_metadata_functions/second.parquet 2 60 + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + +statement ok +DROP TABLE test_table; diff --git a/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt new file mode 100644 index 0000000000000..d936a89beb9f7 --- /dev/null +++ b/datafusion/sqllogictest/test_files/parquet_nested_schema_pruning.slt @@ -0,0 +1,565 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +########## +# Nested projection pruning: a table whose declared nested type is narrower +# than the Parquet file's physical type reads only the declared leaves. +# +# This file covers both halves of that claim: the results are correct, and +# the scan really did read less. Each `explain analyze` below pins a literal +# bytes_scanned against a same-context baseline table declaring the file's +# own physical schema, so no cast is inserted and every leaf is read. A +# change that silently widens a clipped read shows up as a mismatch here. +########## + +# The file contains events: ARRAY> and +# s: STRUCT; the table below declares narrower nested types. +statement ok +COPY ( + SELECT id, events, s + FROM (VALUES + (1, [named_struct('x', 10, 'y', 'a1', 'pad_a', 'p', 'pad_b', 'q')], + named_struct('x', 100, 'y', 's1', 'pad', 'sp1')), + (2, [named_struct('x', 20, 'y', 'b1', 'pad_a', 'p', 'pad_b', 'q'), + named_struct('x', 21, 'y', 'b2', 'pad_a', 'p', 'pad_b', 'q')], + named_struct('x', 200, 'y', 's2', 'pad', 'sp2')), + (3, NULL, + NULL) + ) AS t(id, events, s) +) TO 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet' +STORED AS PARQUET; + +# Declared schema drops pad_a/pad_b from the list elements and pad from the +# struct, declares x as BIGINT (the file has INT), and adds a z column that +# does not exist in the file. +statement ok +CREATE EXTERNAL TABLE narrow ( + id INT, + events ARRAY>, + s STRUCT +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +query I?? +SELECT id, events, s FROM narrow ORDER BY id; +---- +1 [{x: 10, y: a1, z: NULL}] {x: 100, y: s1} +2 [{x: 20, y: b1, z: NULL}, {x: 21, y: b2, z: NULL}] {x: 200, y: s2} +3 NULL NULL + +# Struct-level nullability is preserved: row 3's struct is NULL, not a +# struct of NULLs. +query IBB +SELECT id, events IS NULL, s IS NULL FROM narrow ORDER BY id; +---- +1 false false +2 false false +3 true true + +query II +SELECT id, s['x'] FROM narrow ORDER BY id; +---- +1 100 +2 200 +3 NULL + +query II +SELECT id, e['x'] FROM (SELECT id, unnest(events) AS e FROM narrow) ORDER BY id, e['x']; +---- +1 10 +2 20 +2 21 + +# `full_schema` names every field the file has, so nothing can be clipped away +# and the scan always reads every leaf: a same-context baseline for the +# bytes_scanned comparison below. (A cast is still inserted — the declared +# leaf types differ from the file's, e.g. VARCHAR maps to Utf8View here while +# the file holds Utf8 — but it is not a *narrowing* one, so `clip_for_cast` +# keeps all the leaves.) +statement ok +CREATE EXTERNAL TABLE full_schema ( + id INT, + events ARRAY>, + s STRUCT +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +# bytes_scanned is a literal (not ) checked-in value: narrow +# reads fewer bytes than full_schema because the cast-clipped leaves drop +# pad_a, pad_b, and pad. A future change that widens the narrow read shows +# up here as a bytes_scanned mismatch. +query TT +explain analyze select events from narrow; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=172] + +query TT +explain analyze select events from full_schema; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=312] + +# Same for the top-level struct column: the clipped read drops `pad`. +query TT +explain analyze select s from narrow; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=146] + +query TT +explain analyze select s from full_schema; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=219] + +# `get_field` on a schema-narrowed struct becomes `get_field(CAST(s), 'x')`; +# the read clips to the cast target (every field the *narrow* schema +# declares), not further down to just `x`. The fair "nothing was clipped" +# baseline is therefore reading every physical leaf of `s` +# (`select s from full_schema` above), not the same `get_field` query against +# `full_schema` -- that one needs no cast at all and takes `get_field`'s own, +# more precise, single-leaf pushdown path. +query TT +explain analyze select s['x'] from narrow; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=146] + +# Mixed access -- the whole (narrowed) column and a subfield of it -- still +# reads only the narrow schema's leaves. +query TT +explain analyze select s, s['y'] from narrow; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=146] + +query TT +explain analyze select s, s['y'] from full_schema; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=219] + + +# `SELECT *` goes through the same clipped read as an explicit projection. +query I?? +SELECT * FROM narrow ORDER BY id; +---- +1 [{x: 10, y: a1, z: NULL}] {x: 100, y: s1} +2 [{x: 20, y: b1, z: NULL}, {x: 21, y: b2, z: NULL}] {x: 200, y: s2} +3 NULL NULL + +# Referencing the narrowed column as a whole *and* through a field access in +# the same query. +query I?T +SELECT id, s, s['y'] FROM narrow ORDER BY id; +---- +1 {x: 100, y: s1} s1 +2 {x: 200, y: s2} s2 +3 NULL NULL + +# Aggregating over a clipped nested column. +query IIT +SELECT count(*), sum(s['x']), string_agg(s['y'], ',' ORDER BY id) FROM narrow; +---- +3 300 s1,s2 + +# Filtering on a field of a clipped nested column. +query I? +SELECT id, s FROM narrow WHERE s['x'] = 200 ORDER BY id; +---- +2 {x: 200, y: s2} + +# A declared schema whose fields are in a different order from the file's: +# the values follow the declared order, not the physical one. +statement ok +CREATE EXTERNAL TABLE reordered ( + id INT, + events ARRAY>, + s STRUCT +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +query I?? +SELECT id, events, s FROM reordered ORDER BY id; +---- +1 [{y: a1, x: 10}] {y: s1, x: 100} +2 [{y: b1, x: 20}, {y: b2, x: 21}] {y: s2, x: 200} +3 NULL NULL + +statement ok +DROP TABLE reordered; + +# A declared struct sharing no field name with the file's is rejected rather +# than silently null-filled: `clip_for_cast` never sees a zero-overlap cast. +statement ok +CREATE EXTERNAL TABLE no_overlap ( + id INT, + s STRUCT +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +statement error DataFusion error: Execution error: Cannot cast column 's' +SELECT s FROM no_overlap; + +statement ok +DROP TABLE no_overlap; + +########## +# Struct nested inside a struct: both levels are clipped, and the reader's +# reconstruction of struct validity survives at both levels. +########## + +statement ok +COPY ( + SELECT id, n + FROM (VALUES + (1, named_struct('inner', named_struct('a', 1, 'pad_i', 'pi1'), 'c', 'c1', 'pad_o', 'po1')), + (2, named_struct('inner', named_struct('a', 2, 'pad_i', 'pi2'), 'c', 'c2', 'pad_o', 'po2')), + (3, NULL) + ) AS t(id, n) +) TO 'test_files/scratch/parquet_nested_schema_pruning/nested_struct.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE nested_narrow ( + id INT, + n STRUCT, c VARCHAR> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/nested_struct.parquet'; + +query I? +SELECT id, n FROM nested_narrow ORDER BY id; +---- +1 {inner: {a: 1}, c: c1} +2 {inner: {a: 2}, c: c2} +3 NULL + +query IBB +SELECT id, n IS NULL, n['inner'] IS NULL FROM nested_narrow ORDER BY id; +---- +1 false false +2 false false +3 true true + +statement ok +DROP TABLE nested_narrow; + +########## +# The exact shape reported in datafusion-comet#4859: a two-level +# `ARRAY>>>` column with a dropped +# struct sibling (`latency_parts`), a dropped map sibling (`feature_map`), a +# dropped nested-struct sibling (`diagnostics`), and dropped top-level +# sibling columns (`dimension_id`, `region_code`, `raw_payload`). +# Structurally the same ReadSchema/InputSchema pair as the issue (field names +# representative, not verbatim), which let Comet's production query read +# 1.35 TB where plain Spark, given the same pruned ReadSchema, read 30.9 GB. +# +# Every dropped sibling carries real data rather than NULLs, so the +# bytes_scanned gap below is attributable to the clip and not to NULL columns +# being cheap. +########## + +statement ok +COPY ( + SELECT id, is_flagged, dimension_id, region_code, events, raw_payload + FROM (VALUES + (1, true, 1001, 'us-east', [named_struct( + 'is_available', true, + 'event_time_ms', 10, + 'event_token', 'token-0', + 'latency_parts', named_struct('queue_time_ms', 5, 'retry_count', 1), + 'items', [named_struct('group_id', 1, 'entity_id', 101, 'metric_value', 1.5, + 'feature_map', MAP {'f1': 0.25}, + 'diagnostics', named_struct('module_id', 'm1', 'trace_id', 't1'), + 'pad', 'pad-0000'), + named_struct('group_id', 2, 'entity_id', 102, 'metric_value', 3.0, + 'feature_map', MAP {'f2': 0.5}, + 'diagnostics', named_struct('module_id', 'm2', 'trace_id', 't2'), + 'pad', 'pad-0001')])], + 'payload-0'), + (2, false, 1002, 'us-west', [named_struct( + 'is_available', false, + 'event_time_ms', 20, + 'event_token', 'token-1', + 'latency_parts', named_struct('queue_time_ms', 7, 'retry_count', 2), + 'items', [named_struct('group_id', 3, 'entity_id', 103, 'metric_value', 4.5, + 'feature_map', MAP {'f3': 0.75}, + 'diagnostics', named_struct('module_id', 'm3', 'trace_id', 't3'), + 'pad', 'pad-0002')])], + 'payload-1') + ) AS t(id, is_flagged, dimension_id, region_code, events, raw_payload) +) TO 'test_files/scratch/parquet_nested_schema_pruning/two_level.parquet' +STORED AS PARQUET; + +# Declares neither the dropped top-level columns nor, inside `events`, +# `event_token`/`latency_parts`, nor, inside `items`, the map, the nested +# struct, or the pad. +statement ok +CREATE EXTERNAL TABLE two_level_narrow ( + id INT, + is_flagged BOOLEAN, + events ARRAY> + >> +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/two_level.parquet'; + +# The file's own physical schema, so no cast is inserted: the same-context +# baseline for the bytes comparison. +statement ok +CREATE EXTERNAL TABLE two_level_full +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/two_level.parquet'; + +# Only the declared subfields survive, at *both* nesting levels: the printed +# structs are the emitted Arrow type. +query I? +SELECT id, events FROM two_level_narrow ORDER BY id; +---- +1 [{is_available: true, event_time_ms: 10, items: [{group_id: 1, entity_id: 101, metric_value: 1.5}, {group_id: 2, entity_id: 102, metric_value: 3.0}]}] +2 [{is_available: false, event_time_ms: 20, items: [{group_id: 3, entity_id: 103, metric_value: 4.5}]}] + +# Unnesting twice reaches the inner list's surviving leaves. +query III +SELECT id, i['group_id'], i['entity_id'] +FROM (SELECT id, unnest(e['items']) AS i + FROM (SELECT id, unnest(events) AS e FROM two_level_narrow)) +ORDER BY id, i['group_id']; +---- +1 1 101 +1 2 102 +2 3 103 + +query TT +explain analyze select events from two_level_narrow; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=2, bytes_scanned=381] + +query TT +explain analyze select events from two_level_full; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=2, bytes_scanned=1.05 K] + +statement ok +DROP TABLE two_level_narrow; + +statement ok +DROP TABLE two_level_full; + +########## +# A MAP column is never clipped (the runtime cast routes maps through Arrow's +# positional struct cast, which needs every child), but it must not stop a +# struct sibling from being clipped. The declared schema below omits the map +# entirely, leaving it in the file as an unprojected root. +########## + +statement ok +COPY ( + SELECT id, m, s + FROM (VALUES + (1, MAP {'k1': 1, 'k2': 2}, named_struct('x', 10, 'pad', 'p1')), + (2, MAP {'k1': 3}, named_struct('x', 20, 'pad', 'p2')) + ) AS t(id, m, s) +) TO 'test_files/scratch/parquet_nested_schema_pruning/with_map.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE map_sibling ( + id INT, + s STRUCT +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/with_map.parquet'; + +query I? +SELECT id, s FROM map_sibling ORDER BY id; +---- +1 {x: 10} +2 {x: 20} + +statement ok +DROP TABLE map_sibling; + +########## +# One table over two files, one physically narrow (no cast inserted) and one +# wide (clipped). Both must read correctly in the same scan. +########## + +statement ok +COPY ( + SELECT id, s + FROM (VALUES + (10, named_struct('x', 1000, 'y', 'w1', 'pad', 'wp1')) + ) AS t(id, s) +) TO 'test_files/scratch/parquet_nested_schema_pruning/mixed/wide.parquet' +STORED AS PARQUET; + +statement ok +COPY ( + SELECT id, s + FROM (VALUES + (20, named_struct('x', 2000, 'y', 'n1')) + ) AS t(id, s) +) TO 'test_files/scratch/parquet_nested_schema_pruning/mixed/narrow.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE mixed_files ( + id INT, + s STRUCT +) +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/mixed/'; + +query I? +SELECT id, s FROM mixed_files ORDER BY id; +---- +10 {x: 1000, y: w1} +20 {x: 2000, y: n1} + +statement ok +DROP TABLE mixed_files; + +########## +# A predicate on a primitive column with filter pushdown enabled, while the +# projected nested column is clipped: the clip and the row filter have to +# coexist on the same scan. +# +# The predicate is deliberately on `id` and not on a field of the clipped +# column: `WHERE s['x'] = ...` with pushdown enabled is silently dropped +# today (apache/datafusion#24109), which is a pre-existing row-filter bug +# rather than anything this feature does. +########## + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +query I? +SELECT id, s FROM narrow WHERE id >= 2 ORDER BY id; +---- +2 {x: 200, y: s2} +3 NULL + +query TT +explain analyze select s from narrow where id >= 2; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=2, bytes_scanned=219] + +query TT +explain analyze select s from full_schema where id >= 2; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=2, bytes_scanned=292] + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +########## +# Query-level casts. `ProjectionExec` is merged into the scan, so a `CAST` +# written in the query reaches the same read-plan analysis as an +# adapter-inserted one — including one column consumed through two *different* +# cast targets, which no single clipped read can serve. Clipping to one +# target's leaves would leave the other cast reading a struct that is missing +# the fields it names, which `cast_column` either null-fills (wrong results) +# or, for disjoint targets, rejects outright. +# +# `exact` infers its schema from the file, so no adapter cast is interposed +# and the casts below are the only ones the scan sees. +########## + +statement ok +CREATE EXTERNAL TABLE exact +STORED AS PARQUET +LOCATION 'test_files/scratch/parquet_nested_schema_pruning/wide.parquet'; + +# A single query-level cast is clipped like an adapter-inserted one. +query ? +SELECT CAST(s AS STRUCT) FROM exact ORDER BY id; +---- +{y: s1} +{y: s2} +NULL + +# Disjoint targets. +query ?? +SELECT CAST(s AS STRUCT) AS q0, + CAST(s AS STRUCT) AS q1 +FROM exact ORDER BY id; +---- +{x: 100} {pad: sp1} +{x: 200} {pad: sp2} +NULL NULL + +# Overlapping targets: q1 needs a leaf q0's clip would have dropped. +query ?? +SELECT CAST(s AS STRUCT) AS q0, + CAST(s AS STRUCT) AS q1 +FROM exact ORDER BY id; +---- +{x: 100} {x: 100, y: s1} +{x: 200} {x: 200, y: s2} +NULL NULL + +# Repeated identical targets still clip. +query ?? +SELECT CAST(s AS STRUCT) AS q0, + CAST(s AS STRUCT) AS q1 +FROM exact ORDER BY id; +---- +{x: 100} {x: 100} +{x: 200} {x: 200} +NULL NULL + +# A cast alongside a whole-column reference: the whole-column read wins. +query ?? +SELECT CAST(s AS STRUCT) AS q0, s +FROM exact ORDER BY id; +---- +{x: 100} {x: 100, y: s1, pad: sp1} +{x: 200} {x: 200, y: s2, pad: sp2} +NULL NULL + +# The conflicting-target fallback reads the whole column -- exactly what a +# scan with no clipping at all reads, and never more. These two must match: +# the first falls back, the second never clips in the first place. +query TT +explain analyze select CAST(s AS STRUCT) AS q0, CAST(s AS STRUCT) AS q1 from exact; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=219] + +query TT +explain analyze select s from exact; +---- +Plan with Metrics DataSourceExec: metrics=[output_rows=3, bytes_scanned=219] + +statement ok +DROP TABLE exact; + +# A query cast stacked on top of the adapter's cast for a narrowed table. +query ? +SELECT CAST(s AS STRUCT) FROM narrow ORDER BY id; +---- +{y: s1} +{y: s2} +NULL + +statement ok +DROP TABLE narrow; + +statement ok +DROP TABLE full_schema; diff --git a/datafusion/sqllogictest/test_files/parquet_statistics.slt b/datafusion/sqllogictest/test_files/parquet_statistics.slt index 1073f60a0fef2..9cf6b1e0381d1 100644 --- a/datafusion/sqllogictest/test_files/parquet_statistics.slt +++ b/datafusion/sqllogictest/test_files/parquet_statistics.slt @@ -59,7 +59,7 @@ query TT EXPLAIN SELECT * FROM test_table WHERE column1 = 1; ---- physical_plan -01)FilterExec: column1@0 = 1, statistics=[Rows=Inexact(2), Bytes=Inexact(10), [(Col[0]: Min=Exact(Int64(1)) Max=Exact(Int64(1)) Null=Inexact(0) Distinct=Exact(1) ScanBytes=Inexact(40))]] +01)FilterExec: column1@0 = 1, statistics=[Rows=Inexact(2), Bytes=Inexact(10), [(Col[0]: Min=Exact(Int64(1)) Max=Exact(Int64(1)) Null=Exact(0) Distinct=Exact(1) ScanBytes=Inexact(10))]] 02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2, statistics=[Rows=Inexact(5), Bytes=Inexact(40), [(Col[0]: Min=Inexact(Int64(1)) Max=Inexact(Int64(4)) Null=Inexact(0) ScanBytes=Inexact(40))]] 03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/test_table/0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/test_table/1.parquet]]}, projection=[column1], file_type=parquet, predicate=column1@0 = 1, pruning_predicate=column1_null_count@2 != row_count@3 AND column1_min@0 <= 1 AND 1 <= column1_max@1, required_guarantees=[column1 in (1)], statistics=[Rows=Inexact(5), Bytes=Inexact(40), [(Col[0]: Min=Inexact(Int64(1)) Max=Inexact(Int64(4)) Null=Inexact(0) ScanBytes=Inexact(40))]] @@ -84,7 +84,7 @@ query TT EXPLAIN SELECT * FROM test_table WHERE column1 = 1; ---- physical_plan -01)FilterExec: column1@0 = 1, statistics=[Rows=Inexact(2), Bytes=Inexact(10), [(Col[0]: Min=Exact(Int64(1)) Max=Exact(Int64(1)) Null=Inexact(0) Distinct=Exact(1) ScanBytes=Inexact(40))]] +01)FilterExec: column1@0 = 1, statistics=[Rows=Inexact(2), Bytes=Inexact(10), [(Col[0]: Min=Exact(Int64(1)) Max=Exact(Int64(1)) Null=Exact(0) Distinct=Exact(1) ScanBytes=Inexact(10))]] 02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2, statistics=[Rows=Inexact(5), Bytes=Inexact(40), [(Col[0]: Min=Inexact(Int64(1)) Max=Inexact(Int64(4)) Null=Inexact(0) ScanBytes=Inexact(40))]] 03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/test_table/0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/test_table/1.parquet]]}, projection=[column1], file_type=parquet, predicate=column1@0 = 1, pruning_predicate=column1_null_count@2 != row_count@3 AND column1_min@0 <= 1 AND 1 <= column1_max@1, required_guarantees=[column1 in (1)], statistics=[Rows=Inexact(5), Bytes=Inexact(40), [(Col[0]: Min=Inexact(Int64(1)) Max=Inexact(Int64(4)) Null=Inexact(0) ScanBytes=Inexact(40))]] @@ -109,7 +109,7 @@ query TT EXPLAIN SELECT * FROM test_table WHERE column1 = 1; ---- physical_plan -01)FilterExec: column1@0 = 1, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]: Min=Exact(Int64(1)) Max=Exact(Int64(1)) Distinct=Exact(1))]] +01)FilterExec: column1@0 = 1, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]: Min=Exact(Int64(1)) Max=Exact(Int64(1)) Null=Exact(0) Distinct=Inexact(1))]] 02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=2, statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:)]] 03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/test_table/0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/test_table/1.parquet]]}, projection=[column1], file_type=parquet, predicate=column1@0 = 1, pruning_predicate=column1_null_count@2 != row_count@3 AND column1_min@0 <= 1 AND 1 <= column1_max@1, required_guarantees=[column1 in (1)], statistics=[Rows=Absent, Bytes=Absent, [(Col[0]:)]] @@ -152,7 +152,7 @@ query TT EXPLAIN SELECT i8 FROM typed_table WHERE i8 = 2; ---- physical_plan -01)FilterExec: i8@0 = 2, statistics=[Rows=Inexact(1), Bytes=Inexact(1), [(Col[0]: Min=Exact(Int8(2)) Max=Exact(Int8(2)) Null=Inexact(0) Distinct=Exact(1) ScanBytes=Inexact(5))]] +01)FilterExec: i8@0 = 2, statistics=[Rows=Inexact(1), Bytes=Inexact(1), [(Col[0]: Min=Exact(Int8(2)) Max=Exact(Int8(2)) Null=Exact(0) Distinct=Exact(1) ScanBytes=Inexact(1))]] 02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, statistics=[Rows=Inexact(5), Bytes=Inexact(5), [(Col[0]: Min=Inexact(Int8(1)) Max=Inexact(Int8(5)) Null=Inexact(0) ScanBytes=Inexact(5))]] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/typed_table.parquet]]}, projection=[i8], file_type=parquet, predicate=i8@0 = 2, pruning_predicate=i8_null_count@2 != row_count@3 AND i8_min@0 <= 2 AND 2 <= i8_max@1, required_guarantees=[i8 in (2)], statistics=[Rows=Inexact(5), Bytes=Inexact(5), [(Col[0]: Min=Inexact(Int8(1)) Max=Inexact(Int8(5)) Null=Inexact(0) ScanBytes=Inexact(5))]] @@ -161,7 +161,7 @@ query TT EXPLAIN SELECT i64 FROM typed_table WHERE i64 = 2; ---- physical_plan -01)FilterExec: i64@0 = 2, statistics=[Rows=Inexact(1), Bytes=Inexact(8), [(Col[0]: Min=Exact(Int64(2)) Max=Exact(Int64(2)) Null=Inexact(0) Distinct=Exact(1) ScanBytes=Inexact(40))]] +01)FilterExec: i64@0 = 2, statistics=[Rows=Inexact(1), Bytes=Inexact(8), [(Col[0]: Min=Exact(Int64(2)) Max=Exact(Int64(2)) Null=Exact(0) Distinct=Exact(1) ScanBytes=Inexact(8))]] 02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, statistics=[Rows=Inexact(5), Bytes=Inexact(40), [(Col[0]: Min=Inexact(Int64(1)) Max=Inexact(Int64(5)) Null=Inexact(0) ScanBytes=Inexact(40))]] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/typed_table.parquet]]}, projection=[i64], file_type=parquet, predicate=i64@1 = 2, pruning_predicate=i64_null_count@2 != row_count@3 AND i64_min@0 <= 2 AND 2 <= i64_max@1, required_guarantees=[i64 in (2)], statistics=[Rows=Inexact(5), Bytes=Inexact(40), [(Col[0]: Min=Inexact(Int64(1)) Max=Inexact(Int64(5)) Null=Inexact(0) ScanBytes=Inexact(40))]] @@ -170,7 +170,7 @@ query TT EXPLAIN SELECT f32 FROM typed_table WHERE f32 = 2.5; ---- physical_plan -01)FilterExec: CAST(f32@0 AS Float64) = 2.5, statistics=[Rows=Inexact(1), Bytes=Inexact(1), [(Col[0]: Min=Exact(Float32(2.5)) Max=Exact(Float32(2.5)) Null=Inexact(0) Distinct=Exact(1) ScanBytes=Inexact(20))]] +01)FilterExec: CAST(f32@0 AS Float64) = 2.5, statistics=[Rows=Inexact(1), Bytes=Inexact(1), [(Col[0]: Min=Exact(Float32(2.5)) Max=Exact(Float32(2.5)) Null=Inexact(0) Distinct=Exact(1) ScanBytes=Inexact(1))]] 02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, statistics=[Rows=Inexact(5), Bytes=Inexact(20), [(Col[0]: Min=Inexact(Float32(1.5)) Max=Inexact(Float32(5.5)) Null=Inexact(0) ScanBytes=Inexact(20))]] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/typed_table.parquet]]}, projection=[f32], file_type=parquet, predicate=CAST(f32@2 AS Float64) = 2.5, pruning_predicate=f32_null_count@2 != row_count@3 AND CAST(f32_min@0 AS Float64) <= 2.5 AND 2.5 <= CAST(f32_max@1 AS Float64), required_guarantees=[], statistics=[Rows=Inexact(5), Bytes=Inexact(20), [(Col[0]: Min=Inexact(Float32(1.5)) Max=Inexact(Float32(5.5)) Null=Inexact(0) ScanBytes=Inexact(20))]] @@ -179,7 +179,7 @@ query TT EXPLAIN SELECT f64 FROM typed_table WHERE 2.5 = f64; ---- physical_plan -01)FilterExec: f64@0 = 2.5, statistics=[Rows=Inexact(1), Bytes=Inexact(1), [(Col[0]: Min=Exact(Float64(2.5)) Max=Exact(Float64(2.5)) Null=Inexact(0) Distinct=Exact(1) ScanBytes=Inexact(40))]] +01)FilterExec: f64@0 = 2.5, statistics=[Rows=Inexact(1), Bytes=Inexact(1), [(Col[0]: Min=Exact(Float64(2.5)) Max=Exact(Float64(2.5)) Null=Exact(0) Distinct=Exact(1) ScanBytes=Inexact(1))]] 02)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, statistics=[Rows=Inexact(5), Bytes=Inexact(40), [(Col[0]: Min=Inexact(Float64(1.5)) Max=Inexact(Float64(5.5)) Null=Inexact(0) ScanBytes=Inexact(40))]] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/parquet_statistics/typed_table.parquet]]}, projection=[f64], file_type=parquet, predicate=f64@3 = 2.5, pruning_predicate=f64_null_count@2 != row_count@3 AND f64_min@0 <= 2.5 AND 2.5 <= f64_max@1, required_guarantees=[f64 in (2.5)], statistics=[Rows=Inexact(5), Bytes=Inexact(40), [(Col[0]: Min=Inexact(Float64(1.5)) Max=Inexact(Float64(5.5)) Null=Inexact(0) ScanBytes=Inexact(40))]] diff --git a/datafusion/sqllogictest/test_files/pipe_operator.slt b/datafusion/sqllogictest/test_files/pipe_operator.slt index 406ddafc7bdea..4e2b867fd744d 100644 --- a/datafusion/sqllogictest/test_files/pipe_operator.slt +++ b/datafusion/sqllogictest/test_files/pipe_operator.slt @@ -15,11 +15,6 @@ # specific language governing permissions and limitations # under the License. -# BigQuery supports the pipe operator syntax -# TODO: Make the Generic dialect support the pipe operator syntax -statement ok -set datafusion.sql_parser.dialect = 'BigQuery'; - statement ok CREATE TABLE test( a INT, @@ -188,14 +183,10 @@ query TII |> AS produce_sales |> LEFT JOIN ( - SELECT "apples" AS item, 123 AS id + SELECT 'apples' AS item, 123 AS id ) AS produce_data ON produce_sales.item = produce_data.item |> SELECT produce_sales.item, sales, id; ---- apples 2 123 bananas 5 NULL - -# Config reset -statement ok -RESET datafusion.sql_parser.dialect; diff --git a/datafusion/sqllogictest/test_files/predicates.slt b/datafusion/sqllogictest/test_files/predicates.slt index d45c3e0b459b4..b4482a3af1beb 100644 --- a/datafusion/sqllogictest/test_files/predicates.slt +++ b/datafusion/sqllogictest/test_files/predicates.slt @@ -204,12 +204,171 @@ SELECT * FROM test WHERE column1 ~ 'z' ---- Bazzz +query T +SELECT * FROM test WHERE column1 ~ '^Bazzz$' +---- +Bazzz + +query T +SELECT * FROM test WHERE column1 ~ '^(foo|Bazzz)$' +---- +foo +Bazzz + +statement ok +CREATE TABLE test_regex_utf8view(s VARCHAR) AS VALUES ('foo'), ('Bazzz'); + +statement ok +set datafusion.explain.logical_plan_only = true + +# `~` anchored literal -> `= Utf8View(..)` +query TT +EXPLAIN SELECT * FROM test_regex_utf8view WHERE s ~ '^Bazzz$' +---- +logical_plan +01)Filter: test_regex_utf8view.s = Utf8View("Bazzz") +02)--TableScan: test_regex_utf8view projection=[s] + +# `~*` anchored literal -> `ILIKE Utf8View(..)` +query TT +EXPLAIN SELECT * FROM test_regex_utf8view WHERE s ~* '^bazzz$' +---- +logical_plan +01)Filter: test_regex_utf8view.s ILIKE Utf8View("bazzz") +02)--TableScan: test_regex_utf8view projection=[s] + +# `~` anchored alternation -> OR of `= Utf8View(..)` comparisons. +query TT +EXPLAIN SELECT * FROM test_regex_utf8view WHERE s ~ '^(foo|Bazzz)$' +---- +logical_plan +01)Filter: test_regex_utf8view.s = Utf8View("foo") OR test_regex_utf8view.s = Utf8View("Bazzz") +02)--TableScan: test_regex_utf8view projection=[s] + +# `~*` anchored alternation -> NOT simplified: it falls back to a regex match, +# because `IN`/`=` cannot express case-insensitive matching. +query TT +EXPLAIN SELECT * FROM test_regex_utf8view WHERE s ~* '^(foo|bazzz)$' +---- +logical_plan +01)Filter: test_regex_utf8view.s ~* Utf8View("^(foo|bazzz)$") +02)--TableScan: test_regex_utf8view projection=[s] + +# `!~` -> `!= Utf8View(..)` +query TT +EXPLAIN SELECT * FROM test_regex_utf8view WHERE s !~ '^Bazzz$' +---- +logical_plan +01)Filter: test_regex_utf8view.s != Utf8View("Bazzz") +02)--TableScan: test_regex_utf8view projection=[s] + +# `!~*` -> `NOT ILIKE Utf8View(..)` +query TT +EXPLAIN SELECT * FROM test_regex_utf8view WHERE s !~* '^bazzz$' +---- +logical_plan +01)Filter: test_regex_utf8view.s NOT ILIKE Utf8View("bazzz") +02)--TableScan: test_regex_utf8view projection=[s] + +# `!~` anchored alternation -> AND of `!= Utf8View(..)` comparisons. +query TT +EXPLAIN SELECT * FROM test_regex_utf8view WHERE s !~ '^(foo|Bazzz)$' +---- +logical_plan +01)Filter: test_regex_utf8view.s != Utf8View("foo") AND test_regex_utf8view.s != Utf8View("Bazzz") +02)--TableScan: test_regex_utf8view projection=[s] + +# `!~*` anchored alternation -> NOT simplified: it falls back to a regex match, +# same reason as the `~*` alternation above. +query TT +EXPLAIN SELECT * FROM test_regex_utf8view WHERE s !~* '^(foo|bazzz)$' +---- +logical_plan +01)Filter: test_regex_utf8view.s !~* Utf8View("^(foo|bazzz)$") +02)--TableScan: test_regex_utf8view projection=[s] + +statement ok +set datafusion.explain.logical_plan_only = false + +# Result assertions +query T +SELECT * FROM test_regex_utf8view WHERE s ~ '^Bazzz$' +---- +Bazzz + +query T +SELECT * FROM test_regex_utf8view WHERE s ~ '^(foo|Bazzz)$' +---- +foo +Bazzz + +# Case-insensitive anchored match over Utf8View: must be simplified to ILIKE +# (not a case-sensitive Eq) and must keep operand types as Utf8View. +query T +SELECT * FROM test_regex_utf8view WHERE s ~* '^bazzz$' +---- +Bazzz + +# Case-insensitive anchored alternation over Utf8View +query T rowsort +SELECT * FROM test_regex_utf8view WHERE s ~* '^(foo|bazzz)$' +---- +Bazzz +foo + +query T rowsort +SELECT * FROM test_regex_utf8view WHERE s !~ '^Bazzz$' +---- +foo + +query T rowsort +SELECT * FROM test_regex_utf8view WHERE s !~* '^bazzz$' +---- +foo + +# Both rows match the alternation, so the negated forms return nothing. +query T rowsort +SELECT * FROM test_regex_utf8view WHERE s !~ '^(foo|Bazzz)$' +---- + +query T rowsort +SELECT * FROM test_regex_utf8view WHERE s !~* '^(foo|bazzz)$' +---- + +statement ok +DROP TABLE test_regex_utf8view; + query T SELECT * FROM test WHERE column1 ~* 'z' ---- Bazzz ZZZZZ +query T +SELECT * FROM test WHERE column1 ~* '^barrr$' +---- +Barrr + +query T +SELECT * FROM test WHERE column1 ~* '^(barrr|bazzz)$' +---- +Barrr +Bazzz + +query T rowsort +SELECT * FROM test WHERE column1 !~ '^Bazzz$' +---- +Barrr +ZZZZZ +foo + +query T rowsort +SELECT * FROM test WHERE column1 !~* '^barrr$' +---- +Bazzz +ZZZZZ +foo + query T SELECT * FROM test WHERE column1 !~ 'z' ---- @@ -662,15 +821,15 @@ OR ---- logical_plan 01)Projection: lineitem.l_partkey -02)--Inner Join: lineitem.l_partkey = part.p_partkey Filter: part.p_brand = Utf8View("Brand#12") AND lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2) AND part.p_size <= Int32(15) -03)----Filter: lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) OR lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) OR lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2) -04)------TableScan: lineitem projection=[l_partkey, l_quantity], partial_filters=[lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) OR lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) OR lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2)] +02)--Inner Join: lineitem.l_partkey = part.p_partkey Filter: part.p_brand = Utf8View("Brand#12") AND lineitem.l_quantity >= Decimal128(1.00,15,2) AND lineitem.l_quantity <= Decimal128(11.00,15,2) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND lineitem.l_quantity >= Decimal128(10.00,15,2) AND lineitem.l_quantity <= Decimal128(20.00,15,2) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND lineitem.l_quantity >= Decimal128(20.00,15,2) AND lineitem.l_quantity <= Decimal128(30.00,15,2) AND part.p_size <= Int32(15) +03)----Filter: lineitem.l_quantity >= Decimal128(1.00,15,2) AND lineitem.l_quantity <= Decimal128(11.00,15,2) OR lineitem.l_quantity >= Decimal128(10.00,15,2) AND lineitem.l_quantity <= Decimal128(20.00,15,2) OR lineitem.l_quantity >= Decimal128(20.00,15,2) AND lineitem.l_quantity <= Decimal128(30.00,15,2) +04)------TableScan: lineitem projection=[l_partkey, l_quantity], partial_filters=[lineitem.l_quantity >= Decimal128(1.00,15,2) AND lineitem.l_quantity <= Decimal128(11.00,15,2) OR lineitem.l_quantity >= Decimal128(10.00,15,2) AND lineitem.l_quantity <= Decimal128(20.00,15,2) OR lineitem.l_quantity >= Decimal128(20.00,15,2) AND lineitem.l_quantity <= Decimal128(30.00,15,2)] 05)----Filter: part.p_size >= Int32(1) AND (part.p_brand = Utf8View("Brand#12") AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_size <= Int32(15)) 06)------TableScan: part projection=[p_partkey, p_brand, p_size], partial_filters=[part.p_size >= Int32(1), part.p_brand = Utf8View("Brand#12") AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_size <= Int32(15)] physical_plan -01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_partkey@0, p_partkey@0)], filter=p_brand@1 = Brand#12 AND l_quantity@0 >= Some(100),15,2 AND l_quantity@0 <= Some(1100),15,2 AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND l_quantity@0 >= Some(1000),15,2 AND l_quantity@0 <= Some(2000),15,2 AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND l_quantity@0 >= Some(2000),15,2 AND l_quantity@0 <= Some(3000),15,2 AND p_size@2 <= 15, projection=[l_partkey@0] +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_partkey@0, p_partkey@0)], filter=p_brand@1 = Brand#12 AND l_quantity@0 >= 1.00 AND l_quantity@0 <= 11.00 AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND l_quantity@0 >= 10.00 AND l_quantity@0 <= 20.00 AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND l_quantity@0 >= 20.00 AND l_quantity@0 <= 30.00 AND p_size@2 <= 15, projection=[l_partkey@0] 02)--RepartitionExec: partitioning=Hash([l_partkey@0], 4), input_partitions=4 -03)----FilterExec: l_quantity@1 >= Some(100),15,2 AND l_quantity@1 <= Some(1100),15,2 OR l_quantity@1 >= Some(1000),15,2 AND l_quantity@1 <= Some(2000),15,2 OR l_quantity@1 >= Some(2000),15,2 AND l_quantity@1 <= Some(3000),15,2 +03)----FilterExec: l_quantity@1 >= 1.00 AND l_quantity@1 <= 11.00 OR l_quantity@1 >= 10.00 AND l_quantity@1 <= 20.00 OR l_quantity@1 >= 20.00 AND l_quantity@1 <= 30.00 04)------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/tpch-csv/lineitem.csv]]}, projection=[l_partkey, l_quantity], file_type=csv, has_header=true 06)--RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 diff --git a/datafusion/sqllogictest/test_files/prepare.slt b/datafusion/sqllogictest/test_files/prepare.slt index 16e41834a3120..a3fe7cfb9010b 100644 --- a/datafusion/sqllogictest/test_files/prepare.slt +++ b/datafusion/sqllogictest/test_files/prepare.slt @@ -107,6 +107,86 @@ EXECUTE my_plan('j%'); statement ok DEALLOCATE my_plan +# Allow prepare $1 IN (subquery) +statement ok +PREPARE my_plan AS SELECT id FROM person WHERE $1 IN (SELECT age FROM person); + +query I rowsort +EXECUTE my_plan(20); +---- +1 + +query I rowsort +EXECUTE my_plan(99); +---- + +statement ok +DEALLOCATE my_plan + +# Allow prepare $1 NOT IN (subquery) +statement ok +PREPARE my_plan AS SELECT id FROM person WHERE $1 NOT IN (SELECT age FROM person); + +query I rowsort +EXECUTE my_plan(99); +---- +1 + +query I rowsort +EXECUTE my_plan(20); +---- + +statement ok +DEALLOCATE my_plan + +# Allow prepare $1 = ANY (subquery) +statement ok +PREPARE my_plan AS SELECT id FROM person WHERE $1 = ANY (SELECT age FROM person); + +query I rowsort +EXECUTE my_plan(20); +---- +1 + +query I rowsort +EXECUTE my_plan(99); +---- + +statement ok +DEALLOCATE my_plan + +# Allow prepare $1 <> ALL (subquery) +statement ok +PREPARE my_plan AS SELECT id FROM person WHERE $1 <> ALL (SELECT age FROM person); + +query I rowsort +EXECUTE my_plan(99); +---- +1 + +query I rowsort +EXECUTE my_plan(20); +---- + +statement ok +DEALLOCATE my_plan + +# Allow prepare $1 < ALL (subquery) +statement ok +PREPARE my_plan AS SELECT id FROM person WHERE $1 < ALL (SELECT age FROM person); + +query I rowsort +EXECUTE my_plan(10); +---- +1 + +query I rowsort +EXECUTE my_plan(50); +---- + +statement ok +DEALLOCATE my_plan + # Check for missing parameters statement ok PREPARE my_plan AS SELECT * FROM person WHERE id < $1; diff --git a/datafusion/sqllogictest/test_files/preserve_file_partitioning.slt b/datafusion/sqllogictest/test_files/preserve_file_partitioning.slt index 175d7d90cd8ed..e2dd22cc82bba 100644 --- a/datafusion/sqllogictest/test_files/preserve_file_partitioning.slt +++ b/datafusion/sqllogictest/test_files/preserve_file_partitioning.slt @@ -258,7 +258,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[f_dkey@0 as f_dkey, count(Int64(1))@1 as count(*), sum(fact_table.value)@2 as sum(fact_table.value)] 02)--AggregateExec: mode=SinglePartitioned, gby=[f_dkey@1 as f_dkey], aggr=[count(Int64(1)), sum(fact_table.value)] -03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], file_type=parquet +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_partitioning=Hash([f_dkey@1], 3), file_type=parquet # Verify results with optimization match results without optimization query TIR rowsort @@ -320,7 +320,7 @@ physical_plan 01)SortPreservingMergeExec: [f_dkey@0 ASC NULLS LAST] 02)--ProjectionExec: expr=[f_dkey@0 as f_dkey, count(Int64(1))@1 as count(*), avg(fact_table_ordered.value)@2 as avg(fact_table_ordered.value)] 03)----AggregateExec: mode=SinglePartitioned, gby=[f_dkey@1 as f_dkey], aggr=[count(Int64(1)), avg(fact_table_ordered.value)], ordering_mode=Sorted -04)------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_ordering=[f_dkey@1 ASC NULLS LAST], file_type=parquet +04)------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_ordering=[f_dkey@1 ASC NULLS LAST], output_partitioning=Hash([f_dkey@1], 3), file_type=parquet query TIR SELECT f_dkey, count(*), avg(value) FROM fact_table_ordered GROUP BY f_dkey ORDER BY f_dkey; @@ -367,7 +367,7 @@ physical_plan 08)--------------FilterExec: service@2 = log 09)----------------RepartitionExec: partitioning=RoundRobinBatch(3), input_partitions=1 10)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension/data.parquet]]}, projection=[d_dkey, env, service], file_type=parquet, predicate=service@2 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] -11)------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_ordering=[f_dkey@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +11)------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_ordering=[f_dkey@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify results without optimization query TTTIR rowsort @@ -418,7 +418,7 @@ physical_plan 06)----------FilterExec: service@2 = log 07)------------RepartitionExec: partitioning=RoundRobinBatch(3), input_partitions=1 08)--------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension/data.parquet]]}, projection=[d_dkey, env, service], file_type=parquet, predicate=service@2 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] -09)--------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_ordering=[f_dkey@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +09)--------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_ordering=[f_dkey@1 ASC NULLS LAST], output_partitioning=Hash([f_dkey@1], 3), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible query TTTIR rowsort SELECT f.f_dkey, MAX(d.env), MAX(d.service), count(*), sum(f.value) @@ -493,7 +493,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[f_dkey@2 as f_dkey, timestamp@0 as timestamp, value@1 as value, row_number() PARTITION BY [fact_table_ordered.f_dkey] ORDER BY [fact_table_ordered.timestamp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [fact_table_ordered.f_dkey] ORDER BY [fact_table_ordered.timestamp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [fact_table_ordered.f_dkey] ORDER BY [fact_table_ordered.timestamp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], file_type=parquet +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 3), file_type=parquet query TPRI rowsort SELECT f_dkey, timestamp, value, @@ -548,7 +548,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[f_dkey@0 as f_dkey, count(Int64(1))@1 as count(*), sum(high_cardinality_table.value)@2 as sum(high_cardinality_table.value)] 02)--AggregateExec: mode=SinglePartitioned, gby=[f_dkey@1 as f_dkey], aggr=[count(Int64(1)), sum(high_cardinality_table.value)] -03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=A/data.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=D/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=B/data.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=E/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], file_type=parquet +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=A/data.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=D/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=B/data.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=E/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/high_cardinality/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_partitioning=Hash([f_dkey@1], 3), file_type=parquet # Verify results with optimization match results without optimization query TIR rowsort @@ -643,7 +643,7 @@ physical_plan 05)--------RepartitionExec: partitioning=Hash([d_dkey@1], 3), input_partitions=3 06)----------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=C/data.parquet]]}, projection=[env, d_dkey], file_type=parquet 07)--------RepartitionExec: partitioning=Hash([f_dkey@1], 3), input_partitions=3 -08)----------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], file_type=parquet, predicate=DynamicFilter [ empty ] +08)----------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible query TTR rowsort SELECT f.f_dkey, d.env, sum(f.value) @@ -685,8 +685,8 @@ physical_plan 02)--RepartitionExec: partitioning=Hash([f_dkey@0, env@1], 3), input_partitions=3 03)----AggregateExec: mode=Partial, gby=[f_dkey@1 as f_dkey, env@2 as env], aggr=[sum(f.value)] 04)------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(d_dkey@1, f_dkey@1)], projection=[value@2, f_dkey@3, env@0] -05)--------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=C/data.parquet]]}, projection=[env, d_dkey], file_type=parquet -06)--------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], file_type=parquet +05)--------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/dimension_partitioned/d_dkey=C/data.parquet]]}, projection=[env, d_dkey], output_partitioning=Hash([d_dkey@1], 3), file_type=parquet +06)--------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[value, f_dkey], output_partitioning=Hash([f_dkey@1], 3), file_type=parquet query TTR rowsort SELECT f.f_dkey, d.env, sum(f.value) @@ -722,7 +722,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[f_dkey@0 as f_dkey, timestamp@1 as timestamp, count(Int64(1))@2 as count(*), avg(fact_table.value)@3 as avg(fact_table.value)] 02)--AggregateExec: mode=SinglePartitioned, gby=[f_dkey@2 as f_dkey, timestamp@0 as timestamp], aggr=[count(Int64(1)), avg(fact_table.value)] -03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], file_type=parquet +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/preserve_file_partitioning/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_partitioning=Hash([f_dkey@2], 3), file_type=parquet query TPIR rowsort SELECT f_dkey, timestamp, diff --git a/datafusion/sqllogictest/test_files/projection_pushdown.slt b/datafusion/sqllogictest/test_files/projection_pushdown.slt index 344aef1f92cf9..f59d9da0fe68c 100644 --- a/datafusion/sqllogictest/test_files/projection_pushdown.slt +++ b/datafusion/sqllogictest/test_files/projection_pushdown.slt @@ -412,8 +412,9 @@ logical_plan 02)--Projection: three_cols.col_a, three_cols.col_b, three_cols.col_c, three_cols.col_b AS col_b_dup 03)----TableScan: three_cols projection=[col_a, col_b, col_c] physical_plan -01)SortExec: expr=[col_a@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/three_cols.parquet]]}, projection=[col_a, col_b, col_c, col_b@1 as col_b_dup], file_type=parquet, sort_order_for_reorder=[col_a@0 ASC NULLS LAST] +01)ProjectionExec: expr=[col_a@0 as col_a, col_b@1 as col_b, col_c@2 as col_c, col_b@1 as col_b_dup] +02)--SortExec: expr=[col_a@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/three_cols.parquet]]}, projection=[col_a, col_b, col_c], file_type=parquet, sort_order_for_reorder=[col_a@0 ASC NULLS LAST] # Verify correctness query IIII @@ -444,7 +445,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id, s] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) as simple_struct.s[value]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) as simple_struct.s[value]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query II @@ -467,7 +468,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id, s] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) + 1 as simple_struct.s[value] + Int64(1)], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) + 1 as simple_struct.s[value] + Int64(1)], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query II @@ -490,7 +491,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id, s] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) as simple_struct.s[value], get_field(s@1, label) as simple_struct.s[label]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) as simple_struct.s[value], get_field(s@1, label) as simple_struct.s[label]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query IIT @@ -513,7 +514,7 @@ logical_plan 03)----TableScan: nested_struct projection=[id, nested] physical_plan 01)SortExec: TopK(fetch=2), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/nested.parquet]]}, projection=[id, get_field(nested@1, outer, inner) as nested_struct.nested[outer][inner]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/nested.parquet]]}, projection=[id, get_field(nested@1, outer, inner) as nested_struct.nested[outer][inner]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query II @@ -535,7 +536,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id, s] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, label) || _suffix as simple_struct.s[label] || Utf8("_suffix")], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, label) || _suffix as simple_struct.s[label] || Utf8("_suffix")], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query IT @@ -564,8 +565,8 @@ logical_plan 04)------Projection: get_field(simple_struct.s, Utf8("value")) AS __datafusion_extracted_1, simple_struct.id 05)--------TableScan: simple_struct projection=[id, s], partial_filters=[simple_struct.id > Int64(1)] physical_plan -01)SortExec: expr=[simple_struct.s[value]@1 ASC NULLS LAST], preserve_partitioning=[false] -02)--ProjectionExec: expr=[id@1 as id, __datafusion_extracted_1@0 as simple_struct.s[value]] +01)ProjectionExec: expr=[id@1 as id, __datafusion_extracted_1@0 as simple_struct.s[value]] +02)--SortExec: expr=[__datafusion_extracted_1@0 ASC NULLS LAST], preserve_partitioning=[false] 03)----FilterExec: id@1 > 1 04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=id@0 > 1, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 1, required_guarantees=[] @@ -592,10 +593,10 @@ logical_plan 04)------Projection: get_field(simple_struct.s, Utf8("value")) AS __datafusion_extracted_1, simple_struct.id 05)--------TableScan: simple_struct projection=[id, s], partial_filters=[simple_struct.id > Int64(1)] physical_plan -01)SortExec: TopK(fetch=2), expr=[simple_struct.s[value]@1 ASC NULLS LAST], preserve_partitioning=[false] -02)--ProjectionExec: expr=[id@1 as id, __datafusion_extracted_1@0 as simple_struct.s[value]] +01)ProjectionExec: expr=[id@1 as id, __datafusion_extracted_1@0 as simple_struct.s[value]] +02)--SortExec: TopK(fetch=2), expr=[__datafusion_extracted_1@0 ASC NULLS LAST], preserve_partitioning=[false] 03)----FilterExec: id@1 > 1 -04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=id@0 > 1 AND DynamicFilter [ empty ], pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 1, required_guarantees=[] +04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=id@0 > 1 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 1, required_guarantees=[] # Verify correctness query II @@ -621,7 +622,7 @@ physical_plan 01)SortExec: TopK(fetch=2), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] 02)--ProjectionExec: expr=[id@1 as id, __datafusion_extracted_1@0 + 1 as simple_struct.s[value] + Int64(1)] 03)----FilterExec: id@1 > 1 -04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=id@0 > 1 AND DynamicFilter [ empty ], pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 1, required_guarantees=[] +04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=id@0 > 1 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 1, required_guarantees=[] # Verify correctness query II @@ -713,7 +714,7 @@ logical_plan physical_plan 01)SortPreservingMergeExec: [id@0 ASC NULLS LAST], fetch=3 02)--SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part1.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part3.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part4.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part5.parquet]]}, projection=[id, get_field(s@1, value) as multi_struct.s[value]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part1.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part3.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part4.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part5.parquet]]}, projection=[id, get_field(s@1, value) as multi_struct.s[value]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query II @@ -737,7 +738,7 @@ logical_plan physical_plan 01)SortPreservingMergeExec: [id@0 ASC NULLS LAST], fetch=3 02)--SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part1.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part3.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part4.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part5.parquet]]}, projection=[id, get_field(s@1, value) + 1 as multi_struct.s[value] + Int64(1)], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part1.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part3.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part4.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part5.parquet]]}, projection=[id, get_field(s@1, value) + 1 as multi_struct.s[value] + Int64(1)], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query II @@ -762,8 +763,8 @@ logical_plan 05)--------TableScan: multi_struct projection=[id, s], partial_filters=[multi_struct.id > Int64(2)] physical_plan 01)SortPreservingMergeExec: [id@0 ASC NULLS LAST] -02)--SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[id@1 as id, __datafusion_extracted_1@0 as multi_struct.s[value]] +02)--ProjectionExec: expr=[id@1 as id, __datafusion_extracted_1@0 as multi_struct.s[value]] +03)----SortExec: expr=[id@1 ASC NULLS LAST], preserve_partitioning=[true] 04)------FilterExec: id@1 > 2 05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=3 06)----------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part1.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part3.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part4.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/multi/part5.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=id@0 > 2, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 2, required_guarantees=[] @@ -874,7 +875,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id, s] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) as simple_struct.s[value], get_field(s@1, value) + 10 as simple_struct.s[value] + Int64(10), get_field(s@1, label) as simple_struct.s[label]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) as simple_struct.s[value], get_field(s@1, value) + 10 as simple_struct.s[value] + Int64(10), get_field(s@1, label) as simple_struct.s[label]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query IIIT @@ -897,7 +898,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, 42 as constant], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, 42 as constant], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query II @@ -919,7 +920,7 @@ logical_plan 02)--TableScan: simple_struct projection=[id] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query I @@ -947,7 +948,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, id@0 + 100 as computed], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, id@0 + 100 as computed], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query II @@ -1039,7 +1040,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id, s] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) + id@0 as combined], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) + id@0 as combined], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query II @@ -1095,7 +1096,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id, s] physical_plan 01)SortExec: TopK(fetch=2), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, 42 as answer, get_field(s@1, label) as simple_struct.s[label]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, 42 as answer, get_field(s@1, label) as simple_struct.s[label]], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query IIT @@ -1118,7 +1119,7 @@ logical_plan 03)----TableScan: simple_struct projection=[id, s] physical_plan 01)SortExec: TopK(fetch=2), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) + 100 as simple_struct.s[value] + Int64(100), get_field(s@1, label) || _test as simple_struct.s[label] || Utf8("_test")], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) + 100 as simple_struct.s[value] + Int64(100), get_field(s@1, label) || _test as simple_struct.s[label] || Utf8("_test")], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Verify correctness query IIT @@ -1317,7 +1318,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[id@0 as id] 02)--SortExec: TopK(fetch=2), expr=[__datafusion_extracted_1@1 ASC NULLS LAST], preserve_partitioning=[false] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) as __datafusion_extracted_1], file_type=parquet, predicate=DynamicFilter [ empty ] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id, get_field(s@1, value) as __datafusion_extracted_1], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness query I @@ -1424,7 +1425,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(__datafusion_extracted_1@0, __datafusion_extracted_2 * Int64(10)@2)], projection=[id@1, id@3] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, id, get_field(s@1, level) * 10 as __datafusion_extracted_2 * Int64(10)], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, id, get_field(s@1, level) * 10 as __datafusion_extracted_2 * Int64(10)], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness - value = level * 10 # simple_struct: (1,100), (2,200), (3,150), (4,300), (5,250) @@ -1460,7 +1461,7 @@ physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)] 02)--FilterExec: __datafusion_extracted_1@0 > 150, projection=[id@1] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=get_field(s@1, value) > 150 -04)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ] +04)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness - id matches and value > 150 query II @@ -1500,7 +1501,7 @@ physical_plan 02)--FilterExec: __datafusion_extracted_1@0 > 100, projection=[id@1] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet, predicate=get_field(s@1, value) > 100 04)--FilterExec: __datafusion_extracted_2@0 > 3, projection=[id@1] -05)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, id], file_type=parquet, predicate=get_field(s@1, level) > 3 AND DynamicFilter [ empty ] +05)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, id], file_type=parquet, predicate=get_field(s@1, level) > 3 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness - id matches, value > 100, and level > 3 # Matching ids where value > 100: 2(200), 3(150), 4(300), 5(250) @@ -1536,7 +1537,7 @@ physical_plan 01)ProjectionExec: expr=[id@0 as id, __datafusion_extracted_1@1 as simple_struct.s[label], __datafusion_extracted_2@2 as join_right.s[role]] 02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@1, id@1)], projection=[id@1, __datafusion_extracted_1@0, __datafusion_extracted_2@2] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, label) as __datafusion_extracted_1, id], file_type=parquet -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, role) as __datafusion_extracted_2, id], file_type=parquet, predicate=DynamicFilter [ empty ] +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, role) as __datafusion_extracted_2, id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness query ITT @@ -1568,7 +1569,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[id], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness query II @@ -1607,7 +1608,7 @@ physical_plan 02)--HashJoinExec: mode=CollectLeft, join_type=Left, on=[(id@1, id@0)], projection=[id@1, __datafusion_extracted_2@0, __datafusion_extracted_3@3] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_2, id], file_type=parquet 04)----FilterExec: __datafusion_extracted_1@0 > 5, projection=[id@1, __datafusion_extracted_3@2] -05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_1, id, get_field(s@1, level) as __datafusion_extracted_3], file_type=parquet, predicate=get_field(s@1, level) > 5 AND DynamicFilter [ empty ] +05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_1, id, get_field(s@1, level) as __datafusion_extracted_3], file_type=parquet, predicate=get_field(s@1, level) > 5 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness - left join with level > 5 condition # Only join_right rows with level > 5 are matched: id=1 (level=10), id=4 (level=8) @@ -1688,8 +1689,9 @@ logical_plan 04)------Projection: get_field(simple_struct.s, Utf8("value")) AS __datafusion_extracted_1, get_field(simple_struct.s, Utf8("label")) AS __datafusion_extracted_2 05)--------TableScan: simple_struct projection=[s] physical_plan -01)SortExec: expr=[t.s[value]@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as t.s[value], get_field(s@1, label) as t.s[label]], file_type=parquet +01)ProjectionExec: expr=[__datafusion_extracted_1@0 as t.s[value], __datafusion_extracted_2@1 as t.s[label]] +02)--SortExec: expr=[__datafusion_extracted_1@0 ASC NULLS LAST], preserve_partitioning=[false] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, get_field(s@1, label) as __datafusion_extracted_2], file_type=parquet # Verify correctness query IT @@ -1816,13 +1818,14 @@ logical_plan 12)--------------TableScan: simple_struct projection=[id, s], partial_filters=[simple_struct.id > Int64(3)] physical_plan 01)SortPreservingMergeExec: [t.s[value]@0 ASC NULLS LAST] -02)--SortExec: expr=[t.s[value]@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[__datafusion_extracted_1@0 as t.s[value], __datafusion_extracted_2@1 as t.s[label]] -04)------UnionExec +02)--ProjectionExec: expr=[__datafusion_extracted_1@0 as t.s[value], __datafusion_extracted_2@1 as t.s[label]] +03)----UnionExec +04)------SortExec: expr=[__datafusion_extracted_1@0 ASC NULLS LAST], preserve_partitioning=[false] 05)--------FilterExec: id@2 <= 3, projection=[__datafusion_extracted_1@0, __datafusion_extracted_2@1] 06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, get_field(s@1, label) as __datafusion_extracted_2, id], file_type=parquet, predicate=id@0 <= 3, pruning_predicate=id_null_count@1 != row_count@2 AND id_min@0 <= 3, required_guarantees=[] -07)--------FilterExec: id@2 > 3, projection=[__datafusion_extracted_1@0, __datafusion_extracted_2@1] -08)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, get_field(s@1, label) as __datafusion_extracted_2, id], file_type=parquet, predicate=id@0 > 3, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 3, required_guarantees=[] +07)------SortExec: expr=[__datafusion_extracted_1@0 ASC NULLS LAST], preserve_partitioning=[false] +08)--------FilterExec: id@2 > 3, projection=[__datafusion_extracted_1@0, __datafusion_extracted_2@1] +09)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, get_field(s@1, label) as __datafusion_extracted_2, id], file_type=parquet, predicate=id@0 > 3, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 3, required_guarantees=[] # Verify correctness query IT @@ -1899,7 +1902,7 @@ physical_plan 01)ProjectionExec: expr=[__datafusion_extracted_3@0 as s.s[value], __datafusion_extracted_4@1 as j.s[role]] 02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@2, id@2)], filter=__datafusion_extracted_1@1 > __datafusion_extracted_2@0, projection=[__datafusion_extracted_3@4, __datafusion_extracted_4@1] 03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, get_field(s@1, role) as __datafusion_extracted_4, id], file_type=parquet -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, get_field(s@1, value) as __datafusion_extracted_3, id], file_type=parquet, predicate=DynamicFilter [ empty ] +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, get_field(s@1, value) as __datafusion_extracted_3, id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness - only admin roles match (ids 1 and 4) query II @@ -1935,7 +1938,7 @@ logical_plan physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@1, id@1)], filter=__datafusion_extracted_1@0 > __datafusion_extracted_2@1, projection=[id@1, id@3] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/simple.parquet]]}, projection=[get_field(s@1, value) as __datafusion_extracted_1, id], file_type=parquet -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, id], file_type=parquet, predicate=DynamicFilter [ empty ] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/join_right.parquet]]}, projection=[get_field(s@1, level) as __datafusion_extracted_2, id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify correctness - all rows match since value >> level for all ids # simple_struct: (1,100), (2,200), (3,150), (4,300), (5,250) @@ -2068,9 +2071,115 @@ SELECT s, id FROM simple_struct WHERE s['value'] > 100 AND id < 4; {value: 200, label: beta} 2 {value: 150, label: gamma} 3 +##################### +# Section 9: Join key extraction with pruned outputs +##################### + +statement ok +CREATE TABLE issue_22895_rt2 AS SELECT * FROM (VALUES + (named_struct('msg','user auth failed','sid','a'), 1, 'svc1'), + (named_struct('msg','login token','sid','b'), 2, 'svc2') +) v(attributes, id, name); + +query IT +SELECT a.id, b.name +FROM issue_22895_rt2 a JOIN issue_22895_rt2 b + ON a.attributes['sid'] = b.attributes['sid'] +WHERE a.attributes['msg'] LIKE '%auth%' +ORDER BY a.id, b.name; +---- +1 svc1 + +statement ok +CREATE TABLE issue_22895_rt AS SELECT * FROM (VALUES + (named_struct('uid','u1','t','t1'), TIMESTAMP '2026-06-08T10:00:00', 'a'), + (named_struct('uid','u2','t','t2'), TIMESTAMP '2026-06-08T11:00:00', 'b') +) v(attributes, start_timestamp, span_name); + +query P +SELECT r.start_timestamp +FROM issue_22895_rt r +JOIN (SELECT attributes['uid'] AS uid FROM issue_22895_rt) f + ON f.uid = r.attributes['uid'] +WHERE r.attributes['t'] IN (SELECT attributes['t'] FROM issue_22895_rt) +ORDER BY r.start_timestamp; +---- +2026-06-08T10:00:00 +2026-06-08T11:00:00 + # Config reset # The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok SET datafusion.execution.target_partitions = 4; + +##################### +# Section: volatile expressions are not duplicated by projection pushdown +# +# Regression test for #23220: a volatile expression (e.g. `random()`) aliased +# once in a subquery and referenced multiple times must be evaluated once and +# reused. Projection pushdown must not merge the outer projection into the file +# scan when doing so would inline and duplicate the volatile expression. +# Reproduces only against a file scan (not an in-memory table); if the volatile +# expression is duplicated, the two references diverge and `x = y` is false. +##################### + +statement ok +COPY (SELECT 1 AS id UNION ALL SELECT 2 UNION ALL SELECT 3) +TO 'test_files/scratch/projection_pushdown/volatile.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE volatile_scan STORED AS PARQUET +LOCATION 'test_files/scratch/projection_pushdown/volatile.parquet'; + +# The two references to the aliased `random()` value must be equal on every +# row: the expression is evaluated once and reused, not inlined twice. +query B rowsort +SELECT s.x = s.y +FROM (SELECT r AS x, r AS y FROM (SELECT random() AS r FROM volatile_scan) AS t) AS s; +---- +true +true +true + +##################### +# Section: expensive expressions are not re-inlined by projection pushdown +# +# A repeated expensive expression (e.g. `power(a, 2)`) is extracted by CSE into +# a single intermediate projection. Projection pushdown must keep it as one +# `ProjectionExec` above the scan (`power(a, 2)` computed once) rather than +# inlining it into the `DataSourceExec` projection and re-evaluating it at each +# reference site. +##################### + +statement ok +SET datafusion.execution.target_partitions = 1; + +statement ok +COPY (SELECT 1.0::double AS a, 2.0::double AS b, 3::bigint AS c + UNION ALL SELECT 4.0, 5.0, 6 + UNION ALL SELECT 7.0, 8.0, 9) +TO 'test_files/scratch/projection_pushdown/cse.parquet' +STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE cse_scan STORED AS PARQUET +LOCATION 'test_files/scratch/projection_pushdown/cse.parquet'; + +query TT +EXPLAIN SELECT power(a, 2) + b AS x, power(a, 2) - b AS y, power(a, 2) * c AS z +FROM cse_scan; +---- +logical_plan +01)Projection: __common_expr_1 + cse_scan.b AS x, __common_expr_1 - cse_scan.b AS y, __common_expr_1 * CAST(cse_scan.c AS Float64) AS z +02)--Projection: power(cse_scan.a, Float64(2)) AS __common_expr_1, cse_scan.b, cse_scan.c +03)----TableScan: cse_scan projection=[a, b, c] +physical_plan +01)ProjectionExec: expr=[__common_expr_1@0 + b@1 as x, __common_expr_1@0 - b@1 as y, __common_expr_1@0 * CAST(c@2 AS Float64) as z] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/projection_pushdown/cse.parquet]]}, projection=[power(a@0, 2) as __common_expr_1, b, c], file_type=parquet + +# Reset the config changed above (the SLT runner expects target_partitions = 4). +statement ok +SET datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt b/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt index b04b962a5df19..72d034067663e 100644 --- a/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt +++ b/datafusion/sqllogictest/test_files/push_down_filter_parquet.slt @@ -158,7 +158,7 @@ physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(k@0, k@0)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/small_table.parquet]]}, projection=[k], output_ordering=[k@0 ASC NULLS LAST], file_type=parquet 03)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/large_table.parquet]]}, projection=[k, v], output_ordering=[k@0 ASC NULLS LAST], file_type=parquet, predicate=v@1 >= 50 AND DynamicFilter [ empty ], pruning_predicate=v_null_count@1 != row_count@2 AND v_max@0 >= 50, required_guarantees=[] +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/large_table.parquet]]}, projection=[k, v], output_ordering=[k@0 ASC NULLS LAST], file_type=parquet, predicate=v@1 >= 50 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=v_null_count@1 != row_count@2 AND v_max@0 >= 50, required_guarantees=[] statement ok drop table small_table; @@ -206,7 +206,7 @@ EXPLAIN ANALYZE SELECT t FROM topk_pushdown ORDER BY t * t LIMIT 10; ---- Plan with Metrics 01)SortExec: TopK(fetch=10), expr=[t@0 * t@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[t@0 * t@0 < 1884329474306198481], metrics=[output_rows=10, output_batches=1, row_replacements=10] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_pushdown.parquet]]}, projection=[t], output_ordering=[t@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ t@0 * t@0 < 1884329474306198481 ], metrics=[output_rows=128, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=782 total → 782 matched, row_groups_pruned_bloom_filter=782 total → 782 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=128, pushdown_rows_pruned=99.87 K, predicate_cache_inner_records=128, predicate_cache_records=128, scan_efficiency_ratio=64.87% (258.7 K/398.8 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_pushdown.parquet]]}, projection=[t], output_ordering=[t@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ t@0 * t@0 < 1884329474306198481 ], dynamic_rg_pruning=eligible, metrics=[output_rows=128, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=782 total → 782 matched, row_groups_pruned_bloom_filter=782 total → 782 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=128, pushdown_rows_pruned=99.87 K, predicate_cache_inner_records=128, predicate_cache_records=128, scan_efficiency_ratio=64.87% (258.7 K/398.8 K)] statement ok reset datafusion.explain.analyze_categories; @@ -257,7 +257,7 @@ EXPLAIN SELECT * FROM topk_single_col ORDER BY b DESC LIMIT 1; ---- physical_plan 01)SortExec: TopK(fetch=1), expr=[b@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_single_col.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[b@1 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_single_col.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[b@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible statement ok set datafusion.explain.analyze_categories = 'rows'; @@ -268,7 +268,7 @@ EXPLAIN ANALYZE SELECT * FROM topk_single_col ORDER BY b DESC LIMIT 1; ---- Plan with Metrics 01)SortExec: TopK(fetch=1), expr=[b@1 DESC], preserve_partitioning=[false], filter=[b@1 IS NULL OR b@1 > bd], metrics=[output_rows=1, output_batches=1, row_replacements=1] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_single_col.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ b@1 IS NULL OR b@1 > bd ], sort_order_for_reorder=[b@1 DESC], reverse_row_groups=true, pruning_predicate=b_null_count@0 > 0 OR b_null_count@0 != row_count@2 AND b_max@1 > bd, required_guarantees=[], metrics=[output_rows=4, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=4, pushdown_rows_pruned=0, predicate_cache_inner_records=4, predicate_cache_records=4, scan_efficiency_ratio=22.37% (240/1.07 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_single_col.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ b@1 IS NULL OR b@1 > bd ], sort_order_for_reorder=[b@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=b_null_count@0 > 0 OR b_null_count@0 != row_count@2 AND b_max@1 > bd, required_guarantees=[], metrics=[output_rows=4, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=4, pushdown_rows_pruned=0, predicate_cache_inner_records=4, predicate_cache_records=4, scan_efficiency_ratio=21.62% (222/1.03 K)] statement ok reset datafusion.explain.analyze_categories; @@ -319,7 +319,7 @@ EXPLAIN ANALYZE SELECT * FROM topk_multi_col ORDER BY b ASC NULLS LAST, a DESC L ---- Plan with Metrics 01)SortExec: TopK(fetch=2), expr=[b@1 ASC NULLS LAST, a@0 DESC], preserve_partitioning=[false], filter=[b@1 < bb OR b@1 = bb AND (a@0 IS NULL OR a@0 > ac)], metrics=[output_rows=2, output_batches=1, row_replacements=2] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_multi_col.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ b@1 < bb OR b@1 = bb AND (a@0 IS NULL OR a@0 > ac) ], sort_order_for_reorder=[b@1 ASC NULLS LAST, a@0 DESC], pruning_predicate=b_null_count@1 != row_count@2 AND b_min@0 < bb OR b_null_count@1 != row_count@2 AND b_min@0 <= bb AND bb <= b_max@3 AND (a_null_count@4 > 0 OR a_null_count@4 != row_count@2 AND a_max@5 > ac), required_guarantees=[], metrics=[output_rows=4, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=4, pushdown_rows_pruned=0, predicate_cache_inner_records=8, predicate_cache_records=8, scan_efficiency_ratio=22.37% (240/1.07 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_multi_col.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ b@1 < bb OR b@1 = bb AND (a@0 IS NULL OR a@0 > ac) ], sort_order_for_reorder=[b@1 ASC NULLS LAST, a@0 DESC], dynamic_rg_pruning=eligible, pruning_predicate=b_null_count@1 != row_count@2 AND b_min@0 < bb OR b_null_count@1 != row_count@2 AND b_min@0 <= bb AND bb <= b_max@3 AND (a_null_count@4 > 0 OR a_null_count@4 != row_count@2 AND a_max@5 > ac), required_guarantees=[], metrics=[output_rows=4, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=4, pushdown_rows_pruned=0, predicate_cache_inner_records=8, predicate_cache_records=8, scan_efficiency_ratio=21.62% (222/1.03 K)] statement ok reset datafusion.explain.analyze_categories; @@ -388,8 +388,8 @@ FROM join_probe p INNER JOIN join_build AS build ---- Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, a@0), (b@1, b@1)], projection=[a@3, b@4, c@2, e@5], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=20.48% (214/1.04 K)] -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_probe.parquet]]}, projection=[a, b, e], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ], pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=22.78% (246/1.08 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=19.58% (196/1.00 K)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_probe.parquet]]}, projection=[a, b, e], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=22.05% (228/1.03 K)] statement ok reset datafusion.explain.analyze_categories; @@ -474,9 +474,9 @@ INNER JOIN nested_t3 ON nested_t2.c = nested_t3.d; Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(c@3, d@0)], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] 02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, b@0)], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nested_t1.parquet]]}, projection=[a, x], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=18.23% (144/790)] -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nested_t2.parquet]]}, projection=[b, c, y], file_type=parquet, predicate=DynamicFilter [ b@0 >= aa AND b@0 <= ab AND b@0 IN (SET) ([aa, ab]) ], pruning_predicate=b_null_count@1 != row_count@2 AND b_max@0 >= aa AND b_null_count@1 != row_count@2 AND b_min@3 <= ab AND (b_null_count@1 != row_count@2 AND b_min@3 <= aa AND aa <= b_max@0 OR b_null_count@1 != row_count@2 AND b_min@3 <= ab AND ab <= b_max@0), required_guarantees=[b in (aa, ab)], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=5 total → 5 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=3, predicate_cache_inner_records=5, predicate_cache_records=2, scan_efficiency_ratio=23.2% (252/1.09 K)] -05)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nested_t3.parquet]]}, projection=[d, z], file_type=parquet, predicate=DynamicFilter [ d@0 >= ca AND d@0 <= cb AND hash_lookup ], pruning_predicate=d_null_count@1 != row_count@2 AND d_max@0 >= ca AND d_null_count@1 != row_count@2 AND d_min@3 <= cb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=8 total → 8 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=6, predicate_cache_inner_records=8, predicate_cache_records=2, scan_efficiency_ratio=22.12% (184/832)] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nested_t1.parquet]]}, projection=[a, x], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=17.37% (132/760)] +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nested_t2.parquet]]}, projection=[b, c, y], file_type=parquet, predicate=DynamicFilter [ b@0 >= aa AND b@0 <= ab AND b@0 IN (SET) ([aa, ab]) ], dynamic_rg_pruning=eligible, pruning_predicate=b_null_count@1 != row_count@2 AND b_max@0 >= aa AND b_null_count@1 != row_count@2 AND b_min@3 <= ab AND (b_null_count@1 != row_count@2 AND b_min@3 <= aa AND aa <= b_max@0 OR b_null_count@1 != row_count@2 AND b_min@3 <= ab AND ab <= b_max@0), required_guarantees=[b in (aa, ab)], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=5 total → 5 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=3, predicate_cache_inner_records=5, predicate_cache_records=2, scan_efficiency_ratio=22.46% (234/1.04 K)] +05)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nested_t3.parquet]]}, projection=[d, z], file_type=parquet, predicate=DynamicFilter [ d@0 >= ca AND d@0 <= cb AND hash_lookup ], dynamic_rg_pruning=eligible, pruning_predicate=d_null_count@1 != row_count@2 AND d_max@0 >= ca AND d_null_count@1 != row_count@2 AND d_min@3 <= cb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=8 total → 8 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=6, predicate_cache_inner_records=8, predicate_cache_records=2, scan_efficiency_ratio=21.45% (172/802)] statement ok reset datafusion.explain.analyze_categories; @@ -541,7 +541,7 @@ physical_plan 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, d@0)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/parent_build.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=a@0 = aa, pruning_predicate=a_null_count@2 != row_count@3 AND a_min@0 <= aa AND aa <= a_max@1, required_guarantees=[a in (aa)] 03)--RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/parent_probe.parquet]]}, projection=[d, e, f], file_type=parquet, predicate=e@1 = ba AND d@0 = aa AND DynamicFilter [ empty ], pruning_predicate=e_null_count@2 != row_count@3 AND e_min@0 <= ba AND ba <= e_max@1 AND d_null_count@6 != row_count@3 AND d_min@4 <= aa AND aa <= d_max@5, required_guarantees=[d in (aa), e in (ba)] +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/parent_probe.parquet]]}, projection=[d, e, f], file_type=parquet, predicate=e@1 = ba AND d@0 = aa AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=e_null_count@2 != row_count@3 AND e_min@0 <= ba AND ba <= e_max@1 AND d_null_count@6 != row_count@3 AND d_min@4 <= aa AND aa <= d_max@5, required_guarantees=[d in (aa), e in (ba)] statement ok drop table parent_build; @@ -605,8 +605,8 @@ LIMIT 2; Plan with Metrics 01)SortExec: TopK(fetch=2), expr=[e@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[e@0 < bb], metrics=[output_rows=2, output_batches=1, row_replacements=2] 02)--HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, d@0)], projection=[e@2], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_join_build.parquet]]}, projection=[a], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=6.7% (70/1.04 K)] -04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_join_probe.parquet]]}, projection=[d, e], file_type=parquet, predicate=DynamicFilter [ d@0 >= aa AND d@0 <= ab AND d@0 IN (SET) ([aa, ab]) ] AND DynamicFilter [ e@1 < bb ], pruning_predicate=d_null_count@1 != row_count@2 AND d_max@0 >= aa AND d_null_count@1 != row_count@2 AND d_min@3 <= ab AND (d_null_count@1 != row_count@2 AND d_min@3 <= aa AND aa <= d_max@0 OR d_null_count@1 != row_count@2 AND d_min@3 <= ab AND ab <= d_max@0) AND e_null_count@5 != row_count@2 AND e_min@4 < bb, required_guarantees=[d in (aa, ab)], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=4 total → 4 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=15.37% (166/1.08 K)] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_join_build.parquet]]}, projection=[a], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=6.39% (64/1.00 K)] +04)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_join_probe.parquet]]}, projection=[d, e], file_type=parquet, predicate=DynamicFilter [ d@0 >= aa AND d@0 <= ab AND d@0 IN (SET) ([aa, ab]) ] AND DynamicFilter [ e@1 < bb ], dynamic_rg_pruning=eligible, pruning_predicate=d_null_count@1 != row_count@2 AND d_max@0 >= aa AND d_null_count@1 != row_count@2 AND d_min@3 <= ab AND (d_null_count@1 != row_count@2 AND d_min@3 <= aa AND aa <= d_max@0 OR d_null_count@1 != row_count@2 AND d_min@3 <= ab AND ab <= d_max@0) AND e_null_count@5 != row_count@2 AND e_min@4 < bb, required_guarantees=[d in (aa, ab)], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=4 total → 4 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=14.89% (154/1.03 K)] statement ok reset datafusion.explain.analyze_categories; @@ -654,8 +654,9 @@ query TT EXPLAIN ANALYZE SELECT b, a FROM topk_proj ORDER BY a LIMIT 2; ---- Plan with Metrics -01)SortExec: TopK(fetch=2), expr=[a@1 ASC NULLS LAST], preserve_partitioning=[false], filter=[a@1 < 2], metrics=[output_rows=2, output_batches=1, row_replacements=2] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_proj.parquet]]}, projection=[b, a], file_type=parquet, predicate=DynamicFilter [ a@0 < 2 ], sort_order_for_reorder=[a@1 ASC NULLS LAST], pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 2, required_guarantees=[], metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=13.72% (153/1.11 K)] +01)ProjectionExec: expr=[b@1 as b, a@0 as a], metrics=[output_rows=2, output_batches=1] +02)--SortExec: TopK(fetch=2), expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[a@0 < 2], metrics=[output_rows=2, output_batches=1, row_replacements=2] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_proj.parquet]]}, projection=[a, b], file_type=parquet, predicate=DynamicFilter [ a@0 < 2 ], sort_order_for_reorder=[a@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 2, required_guarantees=[], metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=13.21% (141/1.07 K)] # Case 2: prune — `SELECT a` — filter stays as `a < 2` on the scan. query TT @@ -663,7 +664,7 @@ EXPLAIN ANALYZE SELECT a FROM topk_proj ORDER BY a LIMIT 2; ---- Plan with Metrics 01)SortExec: TopK(fetch=2), expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[a@0 < 2], metrics=[output_rows=2, output_batches=1, row_replacements=2] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_proj.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ a@0 < 2 ], sort_order_for_reorder=[a@0 ASC NULLS LAST], pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 2, required_guarantees=[], metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=7.09% (79/1.11 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_proj.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ a@0 < 2 ], sort_order_for_reorder=[a@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 2, required_guarantees=[], metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=6.84% (73/1.07 K)] # Case 3: expression — `SELECT a+1 AS a_plus_1` — the TopK filter is on # `a_plus_1`, the scan predicate must read `a@0 + 1`. @@ -672,7 +673,7 @@ EXPLAIN ANALYZE SELECT a + 1 AS a_plus_1, b FROM topk_proj ORDER BY a_plus_1 LIM ---- Plan with Metrics 01)SortExec: TopK(fetch=2), expr=[a_plus_1@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[a_plus_1@0 < 3], metrics=[output_rows=2, output_batches=1, row_replacements=2] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_proj.parquet]]}, projection=[CAST(a@0 AS Int64) + 1 as a_plus_1, b], file_type=parquet, predicate=DynamicFilter [ CAST(a@0 AS Int64) + 1 < 3 ], metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=13.72% (153/1.11 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_proj.parquet]]}, projection=[CAST(a@0 AS Int64) + 1 as a_plus_1, b], file_type=parquet, predicate=DynamicFilter [ CAST(a@0 AS Int64) + 1 < 3 ], dynamic_rg_pruning=eligible, metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=13.21% (141/1.07 K)] # Case 4: alias shadowing — `SELECT a+1 AS a` — the projection renames # `a+1` to `a`, so the TopK's `a < 3` must still be rewritten to @@ -682,7 +683,7 @@ EXPLAIN ANALYZE SELECT a + 1 AS a, b FROM topk_proj ORDER BY a LIMIT 2; ---- Plan with Metrics 01)SortExec: TopK(fetch=2), expr=[a@0 ASC NULLS LAST], preserve_partitioning=[false], filter=[a@0 < 3], metrics=[output_rows=2, output_batches=1, row_replacements=2] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_proj.parquet]]}, projection=[CAST(a@0 AS Int64) + 1 as a, b], file_type=parquet, predicate=DynamicFilter [ CAST(a@0 AS Int64) + 1 < 3 ], sort_order_for_reorder=[a@0 ASC NULLS LAST], metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=13.72% (153/1.11 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/topk_proj.parquet]]}, projection=[CAST(a@0 AS Int64) + 1 as a, b], file_type=parquet, predicate=DynamicFilter [ CAST(a@0 AS Int64) + 1 < 3 ], sort_order_for_reorder=[a@0 ASC NULLS LAST], dynamic_rg_pruning=eligible, metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=13.21% (141/1.07 K)] statement ok reset datafusion.explain.analyze_categories; @@ -739,12 +740,12 @@ INNER JOIN ( ---- Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, a@0)], projection=[a@0, min_value@2], metrics=[output_rows=2, output_batches=2, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=2, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_agg_build.parquet]]}, projection=[a], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=15.32% (70/457)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_agg_build.parquet]]}, projection=[a], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=14.45% (64/443)] 03)--ProjectionExec: expr=[a@0 as a, min(join_agg_probe.value)@1 as min_value], metrics=[output_rows=2, output_batches=2] 04)----AggregateExec: mode=FinalPartitioned, gby=[a@0 as a], aggr=[min(join_agg_probe.value)], metrics=[output_rows=2, output_batches=2, spill_count=0, spilled_rows=0] 05)------RepartitionExec: partitioning=Hash([a@0], 4), input_partitions=1, metrics=[output_rows=2, output_batches=2, spill_count=0, spilled_rows=0] 06)--------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[min(join_agg_probe.value)], metrics=[output_rows=2, output_batches=1, spill_count=0, spilled_rows=0, skipped_aggregation_rows=0, reduction_factor=100% (2/2)] -07)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_agg_probe.parquet]]}, projection=[a, value], file_type=parquet, predicate=DynamicFilter [ a@0 >= h1 AND a@0 <= h2 AND a@0 IN (SET) ([h1, h2]) ], pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= h1 AND a_null_count@1 != row_count@2 AND a_min@3 <= h2 AND (a_null_count@1 != row_count@2 AND a_min@3 <= h1 AND h1 <= a_max@0 OR a_null_count@1 != row_count@2 AND a_min@3 <= h2 AND h2 <= a_max@0), required_guarantees=[a in (h1, h2)], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=4 total → 4 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=4, predicate_cache_records=2, scan_efficiency_ratio=19.81% (163/823)] +07)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/join_agg_probe.parquet]]}, projection=[a, value], file_type=parquet, predicate=DynamicFilter [ a@0 >= h1 AND a@0 <= h2 AND a@0 IN (SET) ([h1, h2]) ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= h1 AND a_null_count@1 != row_count@2 AND a_min@3 <= h2 AND (a_null_count@1 != row_count@2 AND a_min@3 <= h1 AND h1 <= a_max@0 OR a_null_count@1 != row_count@2 AND a_min@3 <= h2 AND h2 <= a_max@0), required_guarantees=[a in (h1, h2)], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=4 total → 4 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=4, predicate_cache_records=2, scan_efficiency_ratio=19.07% (151/792)] statement ok reset datafusion.explain.analyze_categories; @@ -807,7 +808,7 @@ ON nulls_build.a = nulls_probe.a AND nulls_build.b = nulls_probe.b; Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, a@0), (b@1, b@1)], metrics=[output_rows=1, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=3, input_batches=1, input_rows=1, avg_fanout=100% (1/1), probe_hit_rate=100% (1/1)] 02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nulls_build.parquet]]}, projection=[a, b], file_type=parquet, metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=18.6% (144/774)] -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nulls_probe.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= 1 AND b@1 <= 2 AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:1}, {c0:,c1:2}, {c0:ab,c1:}]) ], pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= 1 AND b_null_count@5 != row_count@2 AND b_min@6 <= 2, required_guarantees=[], metrics=[output_rows=1, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=1, pushdown_rows_pruned=3, predicate_cache_inner_records=8, predicate_cache_records=2, scan_efficiency_ratio=21.1% (237/1.12 K)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nulls_probe.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= 1 AND b@1 <= 2 AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:1}, {c0:,c1:2}, {c0:ab,c1:}]) ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= 1 AND b_null_count@5 != row_count@2 AND b_min@6 <= 2, required_guarantees=[], metrics=[output_rows=1, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=1, pushdown_rows_pruned=3, predicate_cache_inner_records=8, predicate_cache_records=2, scan_efficiency_ratio=20.18% (225/1.11 K)] statement ok reset datafusion.explain.analyze_categories; @@ -872,8 +873,8 @@ ON lj_build.a = lj_probe.a AND lj_build.b = lj_probe.b; ---- Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Left, on=[(a@0, a@0), (b@1, b@1)], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=2, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=20.48% (214/1.04 K)] -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_probe.parquet]]}, projection=[a, b, e], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ], pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=22.78% (246/1.08 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=19.58% (196/1.00 K)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_probe.parquet]]}, projection=[a, b, e], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=22.05% (228/1.03 K)] # LEFT SEMI JOIN: only matching build rows are returned; probe scan still # receives the dynamic filter. @@ -888,8 +889,8 @@ WHERE EXISTS ( ---- Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=LeftSemi, on=[(a@0, a@0), (b@1, b@1)], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=2, input_rows=4, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=20.48% (214/1.04 K)] -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_probe.parquet]]}, projection=[a, b], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ], pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=15.37% (166/1.08 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=19.58% (196/1.00 K)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/lj_probe.parquet]]}, projection=[a, b], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND struct(a@0, b@1) IN (SET) ([{c0:aa,c1:ba}, {c0:ab,c1:bb}]) ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=14.89% (154/1.03 K)] statement ok reset datafusion.explain.analyze_categories; @@ -958,8 +959,8 @@ FROM hl_probe p INNER JOIN hl_build AS build ---- Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, a@0), (b@1, b@1)], projection=[a@3, b@4, c@2, e@5], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/hl_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=20.48% (214/1.04 K)] -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/hl_probe.parquet]]}, projection=[a, b, e], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND hash_lookup ], pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=22.78% (246/1.08 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/hl_build.parquet]]}, projection=[a, b, c], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=19.58% (196/1.00 K)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/hl_probe.parquet]]}, projection=[a, b, e], file_type=parquet, predicate=DynamicFilter [ a@0 >= aa AND a@0 <= ab AND b@1 >= ba AND b@1 <= bb AND hash_lookup ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 >= aa AND a_null_count@1 != row_count@2 AND a_min@3 <= ab AND b_null_count@5 != row_count@2 AND b_max@4 >= ba AND b_null_count@5 != row_count@2 AND b_min@6 <= bb, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=22.05% (228/1.03 K)] statement ok drop table hl_build; @@ -1007,8 +1008,8 @@ FROM int_build b INNER JOIN int_probe p ---- Plan with Metrics 01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id1@0, id1@0), (id2@1, id2@1)], projection=[id1@0, id2@1, value@2, data@5], metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/int_build.parquet]]}, projection=[id1, id2, value], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=19.02% (222/1.17 K)] -03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/int_probe.parquet]]}, projection=[id1, id2, data], file_type=parquet, predicate=DynamicFilter [ id1@0 >= 1 AND id1@0 <= 2 AND id2@1 >= 10 AND id2@1 <= 20 AND hash_lookup ], pruning_predicate=id1_null_count@1 != row_count@2 AND id1_max@0 >= 1 AND id1_null_count@1 != row_count@2 AND id1_min@3 <= 2 AND id2_null_count@5 != row_count@2 AND id2_max@4 >= 10 AND id2_null_count@5 != row_count@2 AND id2_min@6 <= 20, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=21.43% (239/1.11 K)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/int_build.parquet]]}, projection=[id1, id2, value], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=18.23% (204/1.12 K)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/int_probe.parquet]]}, projection=[id1, id2, data], file_type=parquet, predicate=DynamicFilter [ id1@0 >= 1 AND id1@0 <= 2 AND id2@1 >= 10 AND id2@1 <= 20 AND hash_lookup ], dynamic_rg_pruning=eligible, pruning_predicate=id1_null_count@1 != row_count@2 AND id1_max@0 >= 1 AND id1_null_count@1 != row_count@2 AND id1_min@3 <= 2 AND id2_null_count@5 != row_count@2 AND id2_max@4 >= 10 AND id2_null_count@5 != row_count@2 AND id2_min@6 <= 20, required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=2, predicate_cache_inner_records=8, predicate_cache_records=4, scan_efficiency_ratio=20.67% (221/1.07 K)] statement ok reset datafusion.explain.analyze_categories; @@ -1023,6 +1024,309 @@ statement ok drop table int_probe; +######## +# Null-equal joins (IS NOT DISTINCT FROM, INTERSECT) keep dynamic filter pushdown. +# Min/max bounds and membership filters derived from the build side evaluate to NULL +# for a probe-side NULL key, so the pushed predicate carries an `IS NULL` disjunct that +# lets the probe NULL reach the join and null-match a build-side NULL. +######## + +statement ok +COPY (SELECT * FROM (VALUES (11), (22), (NULL)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/nej_probe.parquet' STORED AS PARQUET; + +statement ok +COPY (SELECT * FROM (VALUES (11), (NULL)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/nej_build.parquet' STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE nej_probe STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/nej_probe.parquet'; + +statement ok +CREATE EXTERNAL TABLE nej_build STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/nej_build.parquet'; + +# The probe-side NULL key must survive to match the build-side NULL +query II rowsort +SELECT nej_build.id, nej_probe.id FROM nej_build JOIN nej_probe ON nej_build.id IS NOT DISTINCT FROM nej_probe.id +---- +11 11 +NULL NULL + +# The populated filter shows the final shape: an IS NULL disjunct ahead of the +# bounds and membership checks, keeping the probe NULL alive for the join. +statement ok +set datafusion.explain.analyze_categories = 'rows'; + +query TT +EXPLAIN ANALYZE SELECT nej_build.id, nej_probe.id FROM nej_build JOIN nej_probe ON nej_build.id IS NOT DISTINCT FROM nej_probe.id +---- +Plan with Metrics +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], NullsEqual: true, metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=2, avg_fanout=100% (2/2), probe_hit_rate=100% (2/2)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nej_build.parquet]]}, projection=[id], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=12.92% (65/503)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nej_probe.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ id@0 IS NULL OR id@0 >= 11 AND id@0 <= 11 AND id@0 IN (SET) ([11, NULL]) ], dynamic_rg_pruning=eligible, pruning_predicate=id_null_count@0 > 0 OR id_null_count@0 != row_count@2 AND id_max@1 >= 11 AND id_null_count@0 != row_count@2 AND id_min@3 <= 11 AND (id_null_count@0 != row_count@2 AND id_min@3 <= 11 AND 11 <= id_max@1 OR id_null_count@0 != row_count@2 AND id_min@3 <= NULL AND NULL <= id_max@1), required_guarantees=[], metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=3 total → 3 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=2, pushdown_rows_pruned=1, predicate_cache_inner_records=3, predicate_cache_records=3, scan_efficiency_ratio=14.45% (74/512)] + +statement ok +reset datafusion.explain.analyze_categories; + +statement ok +drop table nej_build; + +statement ok +drop table nej_probe; + + +# Multi-key null-equal join: the IS NULL disjunct covers every nullable key, so a probe row with a +# NULL in either key still reaches the join and null-matches the build side. +statement ok +COPY (SELECT * FROM (VALUES (1, 10), (2, NULL), (NULL, 30)) v(a, b)) TO 'test_files/scratch/push_down_filter_parquet/mnej_probe.parquet' STORED AS PARQUET; + +statement ok +COPY (SELECT * FROM (VALUES (1, 10), (2, NULL)) v(a, b)) TO 'test_files/scratch/push_down_filter_parquet/mnej_build.parquet' STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE mnej_probe STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/mnej_probe.parquet'; + +statement ok +CREATE EXTERNAL TABLE mnej_build STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/mnej_build.parquet'; + +query IIII rowsort +SELECT mnej_build.a, mnej_build.b, mnej_probe.a, mnej_probe.b FROM mnej_build JOIN mnej_probe ON (mnej_build.a IS NOT DISTINCT FROM mnej_probe.a) AND (mnej_build.b IS NOT DISTINCT FROM mnej_probe.b) +---- +1 10 1 10 +2 NULL 2 NULL + +statement ok +set datafusion.explain.analyze_categories = 'rows'; + +# After execution the populated filter shows the applied predicate: an IS NULL disjunct +# per key ahead of the build-side membership check, because the build holds a NULL. +query TT +EXPLAIN ANALYZE SELECT mnej_build.a, mnej_build.b, mnej_probe.a, mnej_probe.b FROM mnej_build JOIN mnej_probe ON (mnej_build.a IS NOT DISTINCT FROM mnej_probe.a) AND (mnej_build.b IS NOT DISTINCT FROM mnej_probe.b) +---- +Plan with Metrics +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(a@0, a@0), (b@1, b@1)], NullsEqual: true, metrics=[output_rows=2, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=3, avg_fanout=100% (2/2), probe_hit_rate=66.67% (2/3)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/mnej_build.parquet]]}, projection=[a, b], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=16.42% (133/810)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/mnej_probe.parquet]]}, projection=[a, b], file_type=parquet, predicate=DynamicFilter [ a@0 IS NULL OR b@1 IS NULL OR a@0 >= 1 AND a@0 <= 2 AND b@1 >= 10 AND b@1 <= 10 AND struct(a@0, b@1) IN (SET) ([{c0:1,c1:10}, {c0:2,c1:}]) ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@0 > 0 OR b_null_count@1 > 0 OR a_null_count@0 != row_count@3 AND a_max@2 >= 1 AND a_null_count@0 != row_count@3 AND a_min@4 <= 2 AND b_null_count@1 != row_count@3 AND b_max@5 >= 10 AND b_null_count@1 != row_count@3 AND b_min@6 <= 10, required_guarantees=[], metrics=[output_rows=3, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=3, pushdown_rows_pruned=0, predicate_cache_inner_records=6, predicate_cache_records=6, scan_efficiency_ratio=18.16% (148/815)] + +statement ok +reset datafusion.explain.analyze_categories; + +statement ok +drop table mnej_build; + +statement ok +drop table mnej_probe; + + +# A NULL-free build has nothing for a probe NULL to null-match, so the pushed filter +# skips the IS NULL widening and keeps its full selectivity. +statement ok +COPY (SELECT * FROM (VALUES (11), (22)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/nnb_build.parquet' STORED AS PARQUET; + +statement ok +COPY (SELECT * FROM (VALUES (11), (33), (NULL)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/nnb_probe.parquet' STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE nnb_build STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/nnb_build.parquet'; + +statement ok +CREATE EXTERNAL TABLE nnb_probe STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/nnb_probe.parquet'; + +query II rowsort +SELECT nnb_build.id, nnb_probe.id FROM nnb_build JOIN nnb_probe ON nnb_build.id IS NOT DISTINCT FROM nnb_probe.id +---- +11 11 + +statement ok +set datafusion.explain.analyze_categories = 'rows'; + +# No IS NULL disjunct in the populated filter: the probe NULL can be pruned safely. +query TT +EXPLAIN ANALYZE SELECT nnb_build.id, nnb_probe.id FROM nnb_build JOIN nnb_probe ON nnb_build.id IS NOT DISTINCT FROM nnb_probe.id +---- +Plan with Metrics +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], NullsEqual: true, metrics=[output_rows=1, output_batches=1, array_map_created_count=1, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=1, avg_fanout=100% (1/1), probe_hit_rate=100% (1/1)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nnb_build.parquet]]}, projection=[id], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=13.71% (68/496)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nnb_probe.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ id@0 >= 11 AND id@0 <= 22 AND id@0 IN (SET) ([11, 22]) ], dynamic_rg_pruning=eligible, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 >= 11 AND id_null_count@1 != row_count@2 AND id_min@3 <= 22 AND (id_null_count@1 != row_count@2 AND id_min@3 <= 11 AND 11 <= id_max@0 OR id_null_count@1 != row_count@2 AND id_min@3 <= 22 AND 22 <= id_max@0), required_guarantees=[id in (11, 22)], metrics=[output_rows=1, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=3 total → 3 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=1, pushdown_rows_pruned=2, predicate_cache_inner_records=3, predicate_cache_records=1, scan_efficiency_ratio=14.45% (74/512)] + +statement ok +reset datafusion.explain.analyze_categories; + +statement ok +drop table nnb_build; + +statement ok +drop table nnb_probe; + + +# A probe key declared NOT NULL skips the disjunct even when the build holds a NULL: +# no probe row can be NULL, so there is nothing to keep. +statement ok +COPY (SELECT * FROM (VALUES (11), (NULL)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/nnp_build.parquet' STORED AS PARQUET; + +statement ok +COPY (SELECT * FROM (VALUES (11), (33)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/nnp_probe.parquet' STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE nnp_build STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/nnp_build.parquet'; + +statement ok +CREATE EXTERNAL TABLE nnp_probe (id BIGINT NOT NULL) STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/nnp_probe.parquet'; + +# The build NULL matches nothing here: the probe cannot produce a NULL. +query II rowsort +SELECT nnp_build.id, nnp_probe.id FROM nnp_build JOIN nnp_probe ON nnp_build.id IS NOT DISTINCT FROM nnp_probe.id +---- +11 11 + +statement ok +set datafusion.explain.analyze_categories = 'rows'; + +query TT +EXPLAIN ANALYZE SELECT nnp_build.id, nnp_probe.id FROM nnp_build JOIN nnp_probe ON nnp_build.id IS NOT DISTINCT FROM nnp_probe.id +---- +Plan with Metrics +01)HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(id@0, id@0)], NullsEqual: true, metrics=[output_rows=1, output_batches=1, array_map_created_count=0, build_input_batches=1, build_input_rows=2, input_batches=1, input_rows=1, avg_fanout=100% (1/1), probe_hit_rate=100% (1/1)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nnp_build.parquet]]}, projection=[id], file_type=parquet, metrics=[output_rows=2, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=0 total → 0 matched, page_index_rows_pruned=0 total → 0 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=0, pushdown_rows_pruned=0, predicate_cache_inner_records=0, predicate_cache_records=0, scan_efficiency_ratio=12.92% (65/503)] +03)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/nnp_probe.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ id@0 >= 11 AND id@0 <= 11 AND id@0 IN (SET) ([11, NULL]) ], dynamic_rg_pruning=eligible, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 >= 11 AND id_null_count@1 != row_count@2 AND id_min@3 <= 11 AND (id_null_count@1 != row_count@2 AND id_min@3 <= 11 AND 11 <= id_max@0 OR id_null_count@1 != row_count@2 AND id_min@3 <= NULL AND NULL <= id_max@0), required_guarantees=[id in (11, NULL)], metrics=[output_rows=1, output_batches=1, files_ranges_pruned_statistics=1 total → 1 matched, row_groups_pruned_statistics=1 total → 1 matched, row_groups_pruned_bloom_filter=1 total → 1 matched, page_index_pages_pruned=1 total → 1 matched, page_index_rows_pruned=2 total → 2 matched, limit_pruned_row_groups=0 total → 0 matched, batches_split=0, file_open_errors=0, file_scan_errors=0, files_opened=1, files_processed=1, num_predicate_creation_errors=0, predicate_evaluation_errors=0, pushdown_rows_matched=1, pushdown_rows_pruned=1, predicate_cache_inner_records=2, predicate_cache_records=1, scan_efficiency_ratio=13.71% (68/496)] + +statement ok +reset datafusion.explain.analyze_categories; + +statement ok +drop table nnp_build; + +statement ok +drop table nnp_probe; + + +# Partitioned mode: the per-partition CASE filter gets the same IS NULL widening, so a +# probe NULL routed to a pruning branch still reaches the join. +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.hash_join_single_partition_threshold = 0; + +statement ok +set datafusion.optimizer.hash_join_single_partition_threshold_rows = 0; + +# Two files per side so each scan starts with multiple partitions and the join +# runs real hash routing instead of collapsing to a single branch. +statement ok +COPY (SELECT * FROM (VALUES (11), (22)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/pnej_probe/1.parquet' STORED AS PARQUET; + +statement ok +COPY (SELECT * FROM (VALUES (33), (NULL)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/pnej_probe/2.parquet' STORED AS PARQUET; + +statement ok +COPY (SELECT * FROM (VALUES (11)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/pnej_build/1.parquet' STORED AS PARQUET; + +statement ok +COPY (SELECT * FROM (VALUES (NULL)) v(id)) TO 'test_files/scratch/push_down_filter_parquet/pnej_build/2.parquet' STORED AS PARQUET; + +statement ok +CREATE EXTERNAL TABLE pnej_probe STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/pnej_probe/'; + +statement ok +CREATE EXTERNAL TABLE pnej_build STORED AS PARQUET LOCATION 'test_files/scratch/push_down_filter_parquet/pnej_build/'; + +query TT +EXPLAIN SELECT pnej_build.id, pnej_probe.id FROM pnej_build JOIN pnej_probe ON pnej_build.id IS NOT DISTINCT FROM pnej_probe.id +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(id@0, id@0)], NullsEqual: true +02)--RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=2 +03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/pnej_build/1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/pnej_build/2.parquet]]}, projection=[id], file_type=parquet +04)--RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=2 +05)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/pnej_probe/1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_parquet/pnej_probe/2.parquet]]}, projection=[id], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + +query II rowsort +SELECT pnej_build.id, pnej_probe.id FROM pnej_build JOIN pnej_probe ON pnej_build.id IS NOT DISTINCT FROM pnej_probe.id +---- +11 11 +NULL NULL + +statement ok +drop table pnej_build; + +statement ok +drop table pnej_probe; + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +RESET datafusion.optimizer.hash_join_single_partition_threshold; + +statement ok +RESET datafusion.optimizer.hash_join_single_partition_threshold_rows; + + +######## +# Regression test for build-NULL + emptied-probe interaction in null-aware LeftAnti joins. +# +# `x NOT IN (subquery)` plans as a null-aware LeftAnti hash join where `x` is +# the build (left) side. The dynamic-filter pushdown derives a bounds/membership +# filter from the build keys and pushes it onto the probe scan. When the build +# contains a NULL key and the filter prunes every probe row, the probe looks +# empty to the join. A null-aware LeftAnti treats an empty probe as a genuinely- +# absent subquery, so it emits the build-side NULL as a matching row. That is +# wrong: `NULL NOT IN (non-empty set)` must be UNKNOWN, not TRUE. +# +# The fix: suppress dynamic-filter pushdown whenever the build key is nullable +# and the join is null-aware, so the probe is never artificially emptied. +######## + +statement ok +set datafusion.optimizer.enable_join_dynamic_filter_pushdown = true; + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +# Build side: `ao` has a nullable `id` column; the NULL row is the one that +# must NOT appear in the output. +query I +COPY (SELECT * FROM (VALUES (5), (NULL)) v(id)) +TO 'test_files/scratch/push_down_filter_parquet/ao_p.parquet' +STORED AS PARQUET; +---- +2 + +# Probe / subquery side: `i_disj` has two non-NULL values that don't match 5, +# and no NULLs. The subquery is non-empty, so `NULL NOT IN (...)` is UNKNOWN. +query I +COPY (SELECT * FROM (VALUES (2), (3)) v(eid)) +TO 'test_files/scratch/push_down_filter_parquet/i_disj_p.parquet' +STORED AS PARQUET; +---- +2 + +statement ok +CREATE EXTERNAL TABLE ao_p (id INT) STORED AS PARQUET +LOCATION 'test_files/scratch/push_down_filter_parquet/ao_p.parquet'; + +statement ok +CREATE EXTERNAL TABLE i_disj_p (eid INT) STORED AS PARQUET +LOCATION 'test_files/scratch/push_down_filter_parquet/i_disj_p.parquet'; + +# Must return only `5`. `NULL NOT IN (2, 3)` is UNKNOWN, so that row is dropped. +query I +SELECT id FROM ao_p WHERE id NOT IN (SELECT eid FROM i_disj_p) ORDER BY id; +---- +5 + +statement ok +drop table ao_p; + +statement ok +drop table i_disj_p; + +statement ok +RESET datafusion.optimizer.enable_join_dynamic_filter_pushdown; + +statement ok +RESET datafusion.execution.parquet.pushdown_filters; + + # Config reset statement ok RESET datafusion.explain.physical_plan_only; diff --git a/datafusion/sqllogictest/test_files/push_down_filter_regression.slt b/datafusion/sqllogictest/test_files/push_down_filter_regression.slt index 923a51afc8df9..57509fd0395b9 100644 --- a/datafusion/sqllogictest/test_files/push_down_filter_regression.slt +++ b/datafusion/sqllogictest/test_files/push_down_filter_regression.slt @@ -146,7 +146,7 @@ physical_plan 01)AggregateExec: mode=Final, gby=[], aggr=[max(agg_dyn_test.id)] 02)--CoalescePartitionsExec 03)----AggregateExec: mode=Partial, gby=[], aggr=[max(agg_dyn_test.id)] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-01/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-02/j5fUeSDQo22oPyPU.parquet], [WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-03/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-04/j5fUeSDQo22oPyPU.parquet]]}, projection=[id], file_type=parquet, predicate=id@0 > 1 AND DynamicFilter [ empty ], pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 1, required_guarantees=[] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-01/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-02/j5fUeSDQo22oPyPU.parquet], [WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-03/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-04/j5fUeSDQo22oPyPU.parquet]]}, projection=[id], file_type=parquet, predicate=id@0 > 1 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=id_null_count@1 != row_count@2 AND id_max@0 > 1, required_guarantees=[] query I select max(id) from agg_dyn_test where id > 1; @@ -161,7 +161,7 @@ physical_plan 01)AggregateExec: mode=Final, gby=[], aggr=[max(agg_dyn_test.id)] 02)--CoalescePartitionsExec 03)----AggregateExec: mode=Partial, gby=[], aggr=[max(agg_dyn_test.id)] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-01/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-02/j5fUeSDQo22oPyPU.parquet], [WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-03/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-04/j5fUeSDQo22oPyPU.parquet]]}, projection=[id], file_type=parquet, predicate=CAST(id@0 AS Int64) + 1 > 1 AND DynamicFilter [ empty ] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-01/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-02/j5fUeSDQo22oPyPU.parquet], [WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-03/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-04/j5fUeSDQo22oPyPU.parquet]]}, projection=[id], file_type=parquet, predicate=CAST(id@0 AS Int64) + 1 > 1 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Expect dynamic filter available inside data source query TT @@ -171,7 +171,7 @@ physical_plan 01)AggregateExec: mode=Final, gby=[], aggr=[max(agg_dyn_test.id), min(agg_dyn_test.id)] 02)--CoalescePartitionsExec 03)----AggregateExec: mode=Partial, gby=[], aggr=[max(agg_dyn_test.id), min(agg_dyn_test.id)] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-01/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-02/j5fUeSDQo22oPyPU.parquet], [WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-03/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-04/j5fUeSDQo22oPyPU.parquet]]}, projection=[id], file_type=parquet, predicate=id@0 < 10 AND DynamicFilter [ empty ], pruning_predicate=id_null_count@1 != row_count@2 AND id_min@0 < 10, required_guarantees=[] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-01/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-02/j5fUeSDQo22oPyPU.parquet], [WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-03/j5fUeSDQo22oPyPU.parquet, WORKSPACE_ROOT/datafusion/core/tests/data/test_statistics_per_partition/date=2025-03-04/j5fUeSDQo22oPyPU.parquet]]}, projection=[id], file_type=parquet, predicate=id@0 < 10 AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=id_null_count@1 != row_count@2 AND id_min@0 < 10, required_guarantees=[] # Dynamic filter should not be available for grouping sets query TT @@ -222,7 +222,7 @@ set datafusion.execution.collect_statistics = true; # execution (the order in which Partial aggregates publish dynamic filter # updates races against when the scan reads each partition). The original # Rust test only asserted matched < 4; the important invariant here is -# that the DynamicFilter text is correct. +# that dynamic filtering is applied and metrics are suppressed. statement ok set datafusion.explain.analyze_level = summary; @@ -236,7 +236,7 @@ Plan with Metrics 01)AggregateExec: mode=Final, gby=[], aggr=[max(agg_dyn_e2e.column1)], metrics=[] 02)--CoalescePartitionsExec, metrics=[] 03)----AggregateExec: mode=Partial, gby=[], aggr=[max(agg_dyn_e2e.column1)], metrics=[] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn/file_0.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn/file_1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn/file_2.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn/file_3.parquet]]}, projection=[column1], file_type=parquet, predicate=column1@0 > 1 AND DynamicFilter [ column1@0 > 4 ], pruning_predicate=column1_null_count@1 != row_count@2 AND column1_max@0 > 1 AND column1_null_count@1 != row_count@2 AND column1_max@0 > 4, required_guarantees=[], metrics=[] +04)------DataSourceExec: file_groups= projection=[column1], file_type=parquet, predicate=column1@0 > 1 AND DynamicFilter [ column1@0 > ], dynamic_rg_pruning=eligible, pruning_predicate=column1_null_count@1 != row_count@2 AND column1_max@0 > 1 AND column1_null_count@1 != row_count@2 AND column1_max@0 > , required_guarantees=[], metrics=[] statement ok reset datafusion.explain.analyze_categories; @@ -275,17 +275,25 @@ drop table agg_dyn_e2e; statement ok set datafusion.execution.target_partitions = 2; -# --- single-column fixture ([5, 1, 3, 8]) split across 2 files --- +# --- single-column fixture ([1, 8, 1, 8]) split across 2 files --- +# +# Every file shares the same per-file min (1) and max (8). This makes the +# DynamicFilter content deterministic under parallel execution: no matter the +# order in which the Partial aggregates publish their bounds, every partition +# contributes the same min/max, so any snapshot taken by `EXPLAIN ANALYZE` +# equals the fully converged filter. Using files with differing per-file +# extremes (e.g. min 1 vs 3) makes the snapshot race-dependent, which is what +# caused the flakiness reported in #22621. statement ok COPY ( - SELECT * FROM (VALUES (5), (1)) AS v(a) + SELECT * FROM (VALUES (1), (8)) AS v(a) ) TO 'test_files/scratch/push_down_filter_regression/agg_dyn_single/file_0.parquet' STORED AS PARQUET; statement ok COPY ( - SELECT * FROM (VALUES (3), (8)) AS v(a) + SELECT * FROM (VALUES (1), (8)) AS v(a) ) TO 'test_files/scratch/push_down_filter_regression/agg_dyn_single/file_1.parquet' STORED AS PARQUET; @@ -296,10 +304,11 @@ LOCATION 'test_files/scratch/push_down_filter_regression/agg_dyn_single/'; # Use `analyze_level = summary` + `analyze_categories = 'none'` so metrics # render empty; we only care that the `predicate=DynamicFilter [ ... ]` text -# matches. Pruning metrics here are subject to a parallel-execution race +# matches. The pruning *counts* are still subject to a parallel-execution race # (the order in which Partial aggregates publish filter updates vs. when the -# scan reads each partition), so the filter *content* is deterministic but -# the pruning counts are not. +# scan reads each partition), which is why metrics are suppressed. The filter +# *content* is kept deterministic by giving every file the same per-file +# min/max (see the fixture comment above). statement ok set datafusion.explain.analyze_level = summary; @@ -314,7 +323,7 @@ Plan with Metrics 01)AggregateExec: mode=Final, gby=[], aggr=[min(agg_dyn_single.a)], metrics=[] 02)--CoalescePartitionsExec, metrics=[] 03)----AggregateExec: mode=Partial, gby=[], aggr=[min(agg_dyn_single.a)], metrics=[] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_1.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ a@0 < 1 ], pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 1, required_guarantees=[], metrics=[] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_1.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ a@0 < 1 ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 1, required_guarantees=[], metrics=[] # MAX(a) -> DynamicFilter [ a > 8 ] query TT @@ -324,7 +333,7 @@ Plan with Metrics 01)AggregateExec: mode=Final, gby=[], aggr=[max(agg_dyn_single.a)], metrics=[] 02)--CoalescePartitionsExec, metrics=[] 03)----AggregateExec: mode=Partial, gby=[], aggr=[max(agg_dyn_single.a)], metrics=[] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_1.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ a@0 > 8 ], pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 > 8, required_guarantees=[], metrics=[] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_1.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ a@0 > 8 ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_max@0 > 8, required_guarantees=[], metrics=[] # MIN(a), MAX(a) -> DynamicFilter [ a < 1 OR a > 8 ] query TT @@ -334,7 +343,7 @@ Plan with Metrics 01)AggregateExec: mode=Final, gby=[], aggr=[min(agg_dyn_single.a), max(agg_dyn_single.a)], metrics=[] 02)--CoalescePartitionsExec, metrics=[] 03)----AggregateExec: mode=Partial, gby=[], aggr=[min(agg_dyn_single.a), max(agg_dyn_single.a)], metrics=[] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_1.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ a@0 < 1 OR a@0 > 8 ], pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 1 OR a_null_count@1 != row_count@2 AND a_max@3 > 8, required_guarantees=[], metrics=[] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_single/file_1.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ a@0 < 1 OR a@0 > 8 ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 1 OR a_null_count@1 != row_count@2 AND a_max@3 > 8, required_guarantees=[], metrics=[] # MIN(a+1) -> no dynamic filter (expression input is not a plain column) query TT @@ -350,16 +359,18 @@ statement ok drop table agg_dyn_single; # --- two-column fixture: MIN(a) + MAX(b) across columns --- +# Every file shares the same per-file min(a)=1 and max(b)=9 so the DynamicFilter +# content is deterministic regardless of publish order (see #22621). statement ok COPY ( - SELECT * FROM (VALUES (5, 7), (1, 2)) AS v(a, b) + SELECT * FROM (VALUES (1, 5), (4, 9)) AS v(a, b) ) TO 'test_files/scratch/push_down_filter_regression/agg_dyn_two_col/file_0.parquet' STORED AS PARQUET; statement ok COPY ( - SELECT * FROM (VALUES (3, 4), (8, 9)) AS v(a, b) + SELECT * FROM (VALUES (1, 6), (2, 9)) AS v(a, b) ) TO 'test_files/scratch/push_down_filter_regression/agg_dyn_two_col/file_1.parquet' STORED AS PARQUET; @@ -376,7 +387,7 @@ Plan with Metrics 01)AggregateExec: mode=Final, gby=[], aggr=[min(agg_dyn_two_col.a), max(agg_dyn_two_col.b)], metrics=[] 02)--CoalescePartitionsExec, metrics=[] 03)----AggregateExec: mode=Partial, gby=[], aggr=[min(agg_dyn_two_col.a), max(agg_dyn_two_col.b)], metrics=[] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_two_col/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_two_col/file_1.parquet]]}, projection=[a, b], file_type=parquet, predicate=DynamicFilter [ a@0 < 1 OR b@1 > 9 ], pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 1 OR b_null_count@4 != row_count@2 AND b_max@3 > 9, required_guarantees=[], metrics=[] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_two_col/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_two_col/file_1.parquet]]}, projection=[a, b], file_type=parquet, predicate=DynamicFilter [ a@0 < 1 OR b@1 > 9 ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 1 OR b_null_count@4 != row_count@2 AND b_max@3 > 9, required_guarantees=[], metrics=[] statement ok drop table agg_dyn_two_col; @@ -384,10 +395,12 @@ drop table agg_dyn_two_col; # --- mixed expressions: MIN(a), MAX(a), MAX(b), MIN(c+1) --- # Supported aggregates (MIN(a), MAX(a), MAX(b)) should drive a filter; # MIN(c+1) is unsupported and must not contribute. +# Every file shares the same per-file min(a)=1, max(a)=8 and max(b)=12 so the +# DynamicFilter content is deterministic regardless of publish order (see #22621). statement ok COPY ( - SELECT * FROM (VALUES (5, 10, 100), (1, 4, 70)) AS v(a, b, c) + SELECT * FROM (VALUES (1, 12, 100), (8, 4, 70)) AS v(a, b, c) ) TO 'test_files/scratch/push_down_filter_regression/agg_dyn_mixed/file_0.parquet' STORED AS PARQUET; @@ -410,7 +423,7 @@ Plan with Metrics 01)AggregateExec: mode=Final, gby=[], aggr=[min(agg_dyn_mixed.a), max(agg_dyn_mixed.a), max(agg_dyn_mixed.b), min(agg_dyn_mixed.c + Int64(1))], metrics=[] 02)--CoalescePartitionsExec, metrics=[] 03)----AggregateExec: mode=Partial, gby=[], aggr=[min(agg_dyn_mixed.a), max(agg_dyn_mixed.a), max(agg_dyn_mixed.b), min(agg_dyn_mixed.c + Int64(1))], metrics=[] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_mixed/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_mixed/file_1.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ a@0 < 1 OR a@0 > 8 OR b@1 > 12 ], pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 1 OR a_null_count@1 != row_count@2 AND a_max@3 > 8 OR b_null_count@5 != row_count@2 AND b_max@4 > 12, required_guarantees=[], metrics=[] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_mixed/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_mixed/file_1.parquet]]}, projection=[a, b, c], file_type=parquet, predicate=DynamicFilter [ a@0 < 1 OR a@0 > 8 OR b@1 > 12 ], dynamic_rg_pruning=eligible, pruning_predicate=a_null_count@1 != row_count@2 AND a_min@0 < 1 OR a_null_count@1 != row_count@2 AND a_max@3 > 8 OR b_null_count@5 != row_count@2 AND b_max@4 > 12, required_guarantees=[], metrics=[] statement ok drop table agg_dyn_mixed; @@ -442,7 +455,7 @@ Plan with Metrics 01)AggregateExec: mode=Final, gby=[], aggr=[min(agg_dyn_nulls.a)], metrics=[] 02)--CoalescePartitionsExec, metrics=[] 03)----AggregateExec: mode=Partial, gby=[], aggr=[min(agg_dyn_nulls.a)], metrics=[] -04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_nulls/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_nulls/file_1.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ true ], metrics=[] +04)------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_nulls/file_0.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_dyn_nulls/file_1.parquet]]}, projection=[a], file_type=parquet, predicate=DynamicFilter [ true ], dynamic_rg_pruning=eligible, metrics=[] statement ok reset datafusion.explain.analyze_categories; @@ -502,6 +515,44 @@ physical_plan 05)--------AggregateExec: mode=Partial, gby=[a@0 as a], aggr=[count(agg_filter_pushdown.b)] 06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_filter_pushdown.parquet]]}, projection=[a, b], file_type=parquet +# Mixed filters on an aggregate output and a grouping column must preserve their +# parent filter result order. The grouping-column filter can push below the +# aggregate, but the aggregate-output filter must remain above it. +# Disable logical optimizer passes for this regression so the logical filter +# pushdown rule does not split the mixed predicate before the physical +# `AggregateExec::gather_filters_for_pushdown` path sees it. +statement ok +set datafusion.optimizer.max_passes = 0; + +query TT +EXPLAIN SELECT a, b, cnt FROM ( + SELECT a, b, count(b) AS cnt + FROM agg_filter_pushdown + GROUP BY a, b +) q WHERE cnt = 2 AND b = 'foo'; +---- +physical_plan +01)FilterExec: cnt@2 = 2 +02)--ProjectionExec: expr=[a@0 as a, b@1 as b, count(agg_filter_pushdown.b)@2 as cnt] +03)----AggregateExec: mode=FinalPartitioned, gby=[a@0 as a, b@1 as b], aggr=[count(agg_filter_pushdown.b)] +04)------RepartitionExec: partitioning=Hash([a@0, b@1], 4), input_partitions=4 +05)--------AggregateExec: mode=Partial, gby=[a@0 as a, b@1 as b], aggr=[count(agg_filter_pushdown.b)] +06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/push_down_filter_regression/agg_filter_pushdown.parquet]]}, projection=[a, b], file_type=parquet, predicate=b@1 = CAST(foo AS Utf8View), pruning_predicate=b_null_count@2 != row_count@3 AND b_min@0 <= foo AND foo <= b_max@1, required_guarantees=[] + +# If the aggregate-output filter is incorrectly removed, this query returns 1. +query I +SELECT count(*) FROM ( + SELECT a, b, count(b) AS cnt + FROM agg_filter_pushdown + GROUP BY a, b +) q WHERE cnt = 2 AND b = 'foo'; +---- +0 + +statement ok +reset datafusion.optimizer.max_passes; + statement ok drop table agg_filter_pushdown; diff --git a/datafusion/sqllogictest/test_files/push_down_topk_through_join.slt b/datafusion/sqllogictest/test_files/push_down_topk_through_join.slt new file mode 100644 index 0000000000000..bdc04786f58f7 --- /dev/null +++ b/datafusion/sqllogictest/test_files/push_down_topk_through_join.slt @@ -0,0 +1,1127 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Tests for pushing a TopK (Sort with fetch) through an outer join. +# +# These queries exercise the scenarios handled by the PushDownTopKThroughJoin +# rule. That rule lands in a follow-up PR; the EXPLAIN plans below capture +# current behavior, so the follow-up's diff shows exactly which plans change. +# The query-result checks hold whether or not the rule is enabled. + +statement ok +set datafusion.execution.target_partitions = 1; + +statement ok +set datafusion.explain.logical_plan_only = true; + +statement ok +CREATE TABLE t1 (a INT, b INT, c VARCHAR) AS VALUES + (1, 10, 'one'), + (2, 20, 'two'), + (3, 30, 'three'), + (4, 40, 'four'), + (5, 50, 'five'); + +statement ok +CREATE TABLE t2 (x INT, y INT, z VARCHAR) AS VALUES + (1, 100, 'alpha'), + (2, 200, 'beta'), + (3, 300, 'gamma'), + (6, 600, 'delta'), + (7, 700, 'epsilon'); + +### +### Sort keys come entirely from the preserved side +### + +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC LIMIT 3; +---- +1 10 1 +2 20 2 +3 30 3 + +# RIGHT JOIN: the right input is the preserved side +query TT +EXPLAIN SELECT t1.a, t2.x, t2.y +FROM t1 RIGHT JOIN t2 ON t1.a = t2.x +ORDER BY t2.y ASC LIMIT 3; +---- +logical_plan +01)Sort: t2.y ASC NULLS LAST, fetch=3 +02)--Right Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----TableScan: t2 projection=[x, y] + +query III +SELECT t1.a, t2.x, t2.y +FROM t1 RIGHT JOIN t2 ON t1.a = t2.x +ORDER BY t2.y ASC LIMIT 3; +---- +1 1 100 +2 2 200 +3 3 300 + +### +### Cases where pushdown does not apply +### + +# INNER JOIN has no preserved side +query TT +EXPLAIN SELECT t1.a, t2.x +FROM t1 INNER JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Projection: t1.a, t2.x +02)--Sort: t1.b ASC NULLS LAST, fetch=3 +03)----Projection: t1.a, t2.x, t1.b +04)------Inner Join: t1.a = t2.x +05)--------TableScan: t1 projection=[a, b] +06)--------TableScan: t2 projection=[x] + +# LEFT JOIN sorted by a right-side (non-preserved) column +query TT +EXPLAIN SELECT t1.a, t2.x, t2.y +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t2.y ASC LIMIT 3; +---- +logical_plan +01)Sort: t2.y ASC NULLS LAST, fetch=3 +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----TableScan: t2 projection=[x, y] + +# FULL OUTER JOIN preserves neither side +query TT +EXPLAIN SELECT t1.a, t2.x +FROM t1 FULL OUTER JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Projection: t1.a, t2.x +02)--Sort: t1.b ASC NULLS LAST, fetch=3 +03)----Projection: t1.a, t2.x, t1.b +04)------Full Join: t1.a = t2.x +05)--------TableScan: t1 projection=[a, b] +06)--------TableScan: t2 projection=[x] + +# Non-equijoin filter in the ON clause only controls matching, not which +# preserved (left) rows appear, so all left rows are still emitted. +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x AND t1.b > t2.y +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--Projection: t1.a, t1.b, t2.x +03)----Left Join: t1.a = t2.x Filter: t1.b > t2.y +04)------TableScan: t1 projection=[a, b] +05)------TableScan: t2 projection=[x, y] + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x AND t1.b > t2.y +ORDER BY t1.b ASC LIMIT 3; +---- +1 10 NULL +2 20 NULL +3 30 NULL + +# Non-equijoin filter on the non-preserved side only +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x AND t2.y > 100 +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----Projection: t2.x +05)------Filter: t2.y > Int32(100) +06)--------TableScan: t2 projection=[x, y] + +# A preserved-side filter in the ON clause suppresses matches, but the rows +# still appear NULL-filled, so it does not change which rows are preserved. +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x AND t1.b > 20 +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--Left Join: t1.a = t2.x Filter: t1.b > Int32(20) +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x AND t1.b > 20 +ORDER BY t1.b ASC LIMIT 3; +---- +1 10 NULL +2 20 NULL +3 30 3 + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x AND t2.y > 100 +ORDER BY t1.b ASC LIMIT 3; +---- +1 10 NULL +2 20 2 +3 30 3 + +# Sort without LIMIT is not a TopK +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +### +### Preserved child already carries a Sort with a fetch +### + +# Inner Sort limits to 5 rows; the outer query takes 2. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 5) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY b ASC LIMIT 2; +---- +logical_plan +01)Sort: sub.b ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: sub +03)----Left Join: t1.a = t2.x +04)------SubqueryAlias: t1 +05)--------Sort: t1.b ASC NULLS LAST, fetch=5 +06)----------TableScan: t1 projection=[a, b] +07)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 5) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY b ASC LIMIT 2; +---- +1 10 1 +2 20 2 + +# Inner Sort limits to 2 rows; the outer query takes 5 (already tighter). +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 2) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY b ASC LIMIT 5; +---- +logical_plan +01)Sort: sub.b ASC NULLS LAST, fetch=5 +02)--SubqueryAlias: sub +03)----Left Join: t1.a = t2.x +04)------SubqueryAlias: t1 +05)--------Sort: t1.b ASC NULLS LAST, fetch=2 +06)----------TableScan: t1 projection=[a, b] +07)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 2) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY b ASC LIMIT 5; +---- +1 10 1 +2 20 2 + +### +### Semi/anti joins: not all preserved-side rows reach the output, so a +### pushed fetch could drop rows that would have survived the join filter +### + +query TT +EXPLAIN SELECT t1.a, t1.b +FROM t1 LEFT SEMI JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--LeftSemi Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +query TT +EXPLAIN SELECT t1.a, t1.b +FROM t1 LEFT ANTI JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--LeftAnti Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +query TT +EXPLAIN SELECT t2.x, t2.y +FROM t1 RIGHT SEMI JOIN t2 ON t1.a = t2.x +ORDER BY t2.y ASC LIMIT 3; +---- +logical_plan +01)Sort: t2.y ASC NULLS LAST, fetch=3 +02)--RightSemi Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----TableScan: t2 projection=[x, y] + +query TT +EXPLAIN SELECT t2.x, t2.y +FROM t1 RIGHT ANTI JOIN t2 ON t1.a = t2.x +ORDER BY t2.y ASC LIMIT 3; +---- +logical_plan +01)Sort: t2.y ASC NULLS LAST, fetch=3 +02)--RightAnti Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----TableScan: t2 projection=[x, y] + +### +### Multi-column sort and OFFSET +### + +# ORDER BY spans both sides (t1.b and t2.y), so the keys are not entirely +# from the preserved side. +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x, t2.y +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC, t2.y ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, t2.y ASC NULLS LAST, fetch=3 +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x, y] + +query IIII +SELECT t1.a, t1.b, t2.x, t2.y +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC, t2.y ASC LIMIT 3; +---- +1 10 1 100 +2 20 2 200 +3 30 3 300 + +# LIMIT with OFFSET: the eligible fetch is limit + offset (2 + 1 = 3). +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC LIMIT 2 OFFSET 1; +---- +logical_plan +01)Limit: skip=1, fetch=2 +02)--Sort: t1.b ASC NULLS LAST, fetch=3 +03)----Left Join: t1.a = t2.x +04)------TableScan: t1 projection=[a, b] +05)------TableScan: t2 projection=[x] + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC LIMIT 2 OFFSET 1; +---- +2 20 2 +3 30 3 + +### +### Resolving sort keys through a projection +### + +# ORDER BY references a projected expression (neg_b = -t1.b); resolution must +# map the alias back to the pre-projection expression. +query TT +EXPLAIN SELECT -t1.b AS neg_b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY neg_b ASC LIMIT 3; +---- +logical_plan +01)Sort: neg_b ASC NULLS LAST, fetch=3 +02)--Projection: (- t1.b) AS neg_b, t2.x +03)----Left Join: t1.a = t2.x +04)------TableScan: t1 projection=[a, b] +05)------TableScan: t2 projection=[x] + +# -b ascending means largest b first +query II +SELECT -t1.b AS neg_b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY neg_b ASC LIMIT 3; +---- +-50 NULL +-40 NULL +-30 3 + +# A non-deterministic sort expression (random()) cannot be duplicated. +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b + random() ASC LIMIT 3; +---- +logical_plan +01)Sort: CAST(t1.b AS Float64) + random() ASC NULLS LAST, fetch=3 +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +# Sort references a column that resolves to random() through the projection. +query TT +EXPLAIN SELECT rand_col, t2.x +FROM ( + SELECT random() AS rand_col, t1.a, t2.x + FROM t1 LEFT JOIN t2 ON t1.a = t2.x +) +ORDER BY rand_col ASC LIMIT 3; +---- +logical_plan +01)Sort: rand_col ASC NULLS LAST, fetch=3 +02)--Projection: random() AS rand_col, t2.x +03)----Left Join: t1.a = t2.x +04)------TableScan: t1 projection=[a] +05)------TableScan: t2 projection=[x] + +### +### SubqueryAlias edge cases +### + +# Preserved child is a SubqueryAlias over a TableScan with no inner Sort. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY b ASC LIMIT 2; +---- +logical_plan +01)Sort: sub.b ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: sub +03)----Left Join: t1.a = t2.x +04)------SubqueryAlias: t1 +05)--------TableScan: t1 projection=[a, b] +06)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY b ASC LIMIT 2; +---- +1 10 1 +2 20 2 + +# RIGHT JOIN; the preserved (right) child already limits to 10 rows. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.a, t2.x, t2.y + FROM t1 + RIGHT JOIN (SELECT * FROM t2 ORDER BY y ASC LIMIT 10) t2 + ON t1.a = t2.x +) sub +ORDER BY y ASC LIMIT 3; +---- +logical_plan +01)Sort: sub.y ASC NULLS LAST, fetch=3 +02)--SubqueryAlias: sub +03)----Right Join: t1.a = t2.x +04)------TableScan: t1 projection=[a] +05)------SubqueryAlias: t2 +06)--------Sort: t2.y ASC NULLS LAST, fetch=10 +07)----------TableScan: t2 projection=[x, y] + +query III +SELECT * FROM ( + SELECT t1.a, t2.x, t2.y + FROM t1 + RIGHT JOIN (SELECT * FROM t2 ORDER BY y ASC LIMIT 10) t2 + ON t1.a = t2.x +) sub +ORDER BY y ASC LIMIT 3; +---- +1 1 100 +2 2 200 +3 3 300 + +# Alias name (foo) differs from the table name; column resolution must follow +# the SubqueryAlias renaming. +query TT +EXPLAIN SELECT * FROM ( + SELECT foo.a, foo.b, t2.x + FROM (SELECT * FROM t1) foo + LEFT JOIN t2 ON foo.a = t2.x +) sub +ORDER BY b ASC LIMIT 3; +---- +logical_plan +01)Sort: sub.b ASC NULLS LAST, fetch=3 +02)--SubqueryAlias: sub +03)----Left Join: foo.a = t2.x +04)------SubqueryAlias: foo +05)--------TableScan: t1 projection=[a, b] +06)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT foo.a, foo.b, t2.x + FROM (SELECT * FROM t1) foo + LEFT JOIN t2 ON foo.a = t2.x +) sub +ORDER BY b ASC LIMIT 3; +---- +1 10 1 +2 20 2 +3 30 3 + +# ORDER BY a non-preserved-side column (t2.x) through a SubqueryAlias. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY x ASC LIMIT 3; +---- +logical_plan +01)Sort: sub.x ASC NULLS LAST, fetch=3 +02)--SubqueryAlias: sub +03)----Left Join: t1.a = t2.x +04)------SubqueryAlias: t1 +05)--------TableScan: t1 projection=[a, b] +06)------TableScan: t2 projection=[x] + +# INNER JOIN wrapped in a SubqueryAlias. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1) t1 + INNER JOIN t2 ON t1.a = t2.x +) sub +ORDER BY b ASC LIMIT 3; +---- +logical_plan +01)Sort: sub.b ASC NULLS LAST, fetch=3 +02)--SubqueryAlias: sub +03)----Inner Join: t1.a = t2.x +04)------SubqueryAlias: t1 +05)--------TableScan: t1 projection=[a, b] +06)------TableScan: t2 projection=[x] + +# Multiple sort columns, both from the preserved side, through a SubqueryAlias. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY a ASC, b ASC LIMIT 3; +---- +logical_plan +01)Sort: sub.a ASC NULLS LAST, sub.b ASC NULLS LAST, fetch=3 +02)--SubqueryAlias: sub +03)----Left Join: t1.a = t2.x +04)------SubqueryAlias: t1 +05)--------TableScan: t1 projection=[a, b] +06)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT t1.a, t1.b, t2.x + FROM (SELECT * FROM t1) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY a ASC, b ASC LIMIT 3; +---- +1 10 1 +2 20 2 +3 30 3 + +# A WHERE filter on the preserved side is pushed below the join by +# PushDownFilter before this scenario is considered. +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +WHERE t1.b > 10 +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--Left Join: t1.a = t2.x +03)----Filter: t1.b > Int32(10) +04)------TableScan: t1 projection=[a, b] +05)----TableScan: t2 projection=[x] + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +WHERE t1.b > 10 +ORDER BY t1.b ASC LIMIT 3; +---- +2 20 2 +3 30 3 +4 40 NULL + +### +### Descending order and explicit NULLS placement +### + +# DESC (NULLS FIRST by default) +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b DESC LIMIT 3; +---- +logical_plan +01)Sort: t1.b DESC NULLS FIRST, fetch=3 +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b DESC LIMIT 3; +---- +5 50 NULL +4 40 NULL +3 30 3 + +# ASC NULLS FIRST +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC NULLS FIRST LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS FIRST, fetch=3 +02)--Left Join: t1.a = t2.x +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 LEFT JOIN t2 ON t1.a = t2.x +ORDER BY t1.b ASC NULLS FIRST LIMIT 3; +---- +1 10 1 +2 20 2 +3 30 3 + +# DESC NULLS LAST on the preserved (right) side of a RIGHT JOIN +query TT +EXPLAIN SELECT t1.a, t2.x, t2.y +FROM t1 RIGHT JOIN t2 ON t1.a = t2.x +ORDER BY t2.y DESC NULLS LAST LIMIT 3; +---- +logical_plan +01)Sort: t2.y DESC NULLS LAST, fetch=3 +02)--Right Join: t1.a = t2.x +03)----TableScan: t1 projection=[a] +04)----TableScan: t2 projection=[x, y] + +query III +SELECT t1.a, t2.x, t2.y +FROM t1 RIGHT JOIN t2 ON t1.a = t2.x +ORDER BY t2.y DESC NULLS LAST LIMIT 3; +---- +NULL 7 700 +NULL 6 600 +3 3 300 + +### +### CROSS JOIN +### + +# Each left row appears |t2| times, so the top-N by left columns must come +# from the top-N left rows. +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 CROSS JOIN t2 +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--Cross Join: +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[x] + +query III +SELECT t1.a, t1.b, t2.x +FROM t1 CROSS JOIN t2 +ORDER BY t1.b ASC, t2.x ASC LIMIT 3; +---- +1 10 1 +1 10 2 +1 10 3 + +# CROSS JOIN sorted by right-side columns. +query TT +EXPLAIN SELECT t1.a, t2.x, t2.y +FROM t1 CROSS JOIN t2 +ORDER BY t2.y ASC LIMIT 3; +---- +logical_plan +01)Sort: t2.y ASC NULLS LAST, fetch=3 +02)--Cross Join: +03)----TableScan: t1 projection=[a] +04)----TableScan: t2 projection=[x, y] + +query III +SELECT t1.a, t2.x, t2.y +FROM t1 CROSS JOIN t2 +ORDER BY t2.y ASC, t1.a ASC LIMIT 3; +---- +1 1 100 +2 1 100 +3 1 100 + +# CROSS JOIN: ORDER BY spans both sides (t1.b + t2.y). +query TT +EXPLAIN SELECT t1.a, t1.b, t2.y +FROM t1 CROSS JOIN t2 +ORDER BY t1.b + t2.y ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b + t2.y ASC NULLS LAST, fetch=3 +02)--Cross Join: +03)----TableScan: t1 projection=[a, b] +04)----TableScan: t2 projection=[y] + +# INNER JOIN with only a non-equi filter: the filter can drop rows from either +# side, so a pushed fetch could select rows that get filtered out. +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x +FROM t1 INNER JOIN t2 ON t1.b > t2.y +ORDER BY t1.b ASC LIMIT 3; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=3 +02)--Projection: t1.a, t1.b, t2.x +03)----Inner Join: Filter: t1.b > t2.y +04)------TableScan: t1 projection=[a, b] +05)------TableScan: t2 projection=[x, y] + +### +### Multi-level outer joins +### + +# Chained LEFT JOINs share t1 as the preserved side. +statement ok +CREATE TABLE t3 (p INT, q INT) AS VALUES + (1, 1000), + (2, 2000), + (3, 3000); + +query TT +EXPLAIN SELECT t1.a, t1.b, t2.x, t3.p +FROM t1 +LEFT JOIN t2 ON t1.a = t2.x +LEFT JOIN t3 ON t1.a = t3.p +ORDER BY t1.b ASC LIMIT 2; +---- +logical_plan +01)Sort: t1.b ASC NULLS LAST, fetch=2 +02)--Left Join: t1.a = t3.p +03)----Left Join: t1.a = t2.x +04)------TableScan: t1 projection=[a, b] +05)------TableScan: t2 projection=[x] +06)----TableScan: t3 projection=[p] + +query IIII +SELECT t1.a, t1.b, t2.x, t3.p +FROM t1 +LEFT JOIN t2 ON t1.a = t2.x +LEFT JOIN t3 ON t1.a = t3.p +ORDER BY t1.b ASC LIMIT 2; +---- +1 10 1 1 +2 20 2 2 + +statement ok +DROP TABLE t3; + +### +### Tied sort keys +### + +# Three preserved-side rows tie on b=10; all tied rows still appear. +statement ok +CREATE TABLE t_tied (a INT, b INT) AS VALUES + (1, 10), + (2, 10), + (3, 10), + (4, 20), + (5, 30); + +statement ok +CREATE TABLE t_other (x INT) AS VALUES (1), (2), (3); + +query TT +EXPLAIN SELECT t_tied.a, t_tied.b, t_other.x +FROM t_tied LEFT JOIN t_other ON t_tied.a = t_other.x +ORDER BY t_tied.b ASC, t_tied.a ASC LIMIT 3; +---- +logical_plan +01)Sort: t_tied.b ASC NULLS LAST, t_tied.a ASC NULLS LAST, fetch=3 +02)--Left Join: t_tied.a = t_other.x +03)----TableScan: t_tied projection=[a, b] +04)----TableScan: t_other projection=[x] + +query III +SELECT t_tied.a, t_tied.b, t_other.x +FROM t_tied LEFT JOIN t_other ON t_tied.a = t_other.x +ORDER BY t_tied.b ASC, t_tied.a ASC LIMIT 3; +---- +1 10 1 +2 10 2 +3 10 3 + +statement ok +DROP TABLE t_tied; + +statement ok +DROP TABLE t_other; + +### +### Nested SubqueryAlias +### + +# Resolve the sort key through multiple alias layers. +query TT +EXPLAIN SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 2; +---- +logical_plan +01)Sort: outer_sub.b ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: outer_sub +03)----Left Join: inner_sub.a = t2.x +04)------SubqueryAlias: inner_sub +05)--------SubqueryAlias: inner_alias +06)----------TableScan: t1 projection=[a, b] +07)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 2; +---- +1 10 1 +2 20 2 + +# Existing inner Sort(fetch=5) sits behind two SubqueryAlias layers. +query TT +EXPLAIN SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 5) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 2; +---- +logical_plan +01)Sort: outer_sub.b ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: outer_sub +03)----Left Join: inner_sub.a = t2.x +04)------SubqueryAlias: inner_sub +05)--------SubqueryAlias: inner_alias +06)----------Sort: t1.b ASC NULLS LAST, fetch=5 +07)------------TableScan: t1 projection=[a, b] +08)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 5) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 2; +---- +1 10 1 +2 20 2 + +# Inner Sort already limits to 2 rows; the outer query takes 5. +query TT +EXPLAIN SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 2) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 5; +---- +logical_plan +01)Sort: outer_sub.b ASC NULLS LAST, fetch=5 +02)--SubqueryAlias: outer_sub +03)----Left Join: inner_sub.a = t2.x +04)------SubqueryAlias: inner_sub +05)--------SubqueryAlias: inner_alias +06)----------Sort: t1.b ASC NULLS LAST, fetch=2 +07)------------TableScan: t1 projection=[a, b] +08)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 2) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 5; +---- +1 10 1 +2 20 2 + +# Inner Sort orders by a (fetch=5); the outer query orders by a different +# column (b). +query TT +EXPLAIN SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1 ORDER BY a ASC LIMIT 5) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 2; +---- +logical_plan +01)Sort: outer_sub.b ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: outer_sub +03)----Left Join: inner_sub.a = t2.x +04)------SubqueryAlias: inner_sub +05)--------SubqueryAlias: inner_alias +06)----------Sort: t1.a ASC NULLS LAST, fetch=5 +07)------------TableScan: t1 projection=[a, b] +08)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1 ORDER BY a ASC LIMIT 5) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 2; +---- +1 10 1 +2 20 2 + +# Inner full sort (ORDER BY a, no fetch) under a different outer sort key. +query TT +EXPLAIN SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1 ORDER BY a ASC) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 2; +---- +logical_plan +01)Sort: outer_sub.b ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: outer_sub +03)----Left Join: inner_sub.a = t2.x +04)------SubqueryAlias: inner_sub +05)--------SubqueryAlias: inner_alias +06)----------TableScan: t1 projection=[a, b] +07)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT inner_sub.a, inner_sub.b, t2.x + FROM ( + SELECT * FROM (SELECT * FROM t1 ORDER BY a ASC) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY b ASC LIMIT 2; +---- +1 10 1 +2 20 2 + +# Inner Sort sits behind SubqueryAlias -> Projection(rename b -> renamed_b) -> +# SubqueryAlias; resolution must look through the Projection to find it. +query TT +EXPLAIN SELECT * FROM ( + SELECT inner_sub.a, inner_sub.renamed_b, t2.x + FROM ( + SELECT a, b AS renamed_b FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 5) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY renamed_b ASC LIMIT 2; +---- +logical_plan +01)Sort: outer_sub.renamed_b ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: outer_sub +03)----Left Join: inner_sub.a = t2.x +04)------SubqueryAlias: inner_sub +05)--------Projection: inner_alias.a, inner_alias.b AS renamed_b +06)----------SubqueryAlias: inner_alias +07)------------Sort: t1.b ASC NULLS LAST, fetch=5 +08)--------------TableScan: t1 projection=[a, b] +09)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT inner_sub.a, inner_sub.renamed_b, t2.x + FROM ( + SELECT a, b AS renamed_b FROM (SELECT * FROM t1 ORDER BY b ASC LIMIT 5) inner_alias + ) inner_sub + LEFT JOIN t2 ON inner_sub.a = t2.x +) outer_sub +ORDER BY renamed_b ASC LIMIT 2; +---- +1 10 1 +2 20 2 + +# Sort sits above a Projection that selects a column subset. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.a, t1.renamed_b, t2.x + FROM (SELECT a, b AS renamed_b FROM t1) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY renamed_b ASC LIMIT 2; +---- +logical_plan +01)Sort: sub.renamed_b ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: sub +03)----Left Join: t1.a = t2.x +04)------SubqueryAlias: t1 +05)--------Projection: t1.a, t1.b AS renamed_b +06)----------TableScan: t1 projection=[a, b] +07)------TableScan: t2 projection=[x] + +query III +SELECT * FROM ( + SELECT t1.a, t1.renamed_b, t2.x + FROM (SELECT a, b AS renamed_b FROM t1) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY renamed_b ASC LIMIT 2; +---- +1 10 1 +2 20 2 + +# random() is computed once in the Projection (as rand_col); ordering by the +# precomputed column does not re-evaluate it, unlike random() in the sort expr. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.rand_col, t2.x + FROM (SELECT random() AS rand_col, a FROM t1) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY rand_col ASC LIMIT 2; +---- +logical_plan +01)Sort: sub.rand_col ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: sub +03)----Projection: t1.rand_col, t2.x +04)------Left Join: t1.a = t2.x +05)--------SubqueryAlias: t1 +06)----------Projection: random() AS rand_col, t1.a +07)------------TableScan: t1 projection=[a] +08)--------TableScan: t2 projection=[x] + +# The outer ORDER BY column resolves to random() through the Projection, and an +# existing inner Sort is also on random() -- but they are independent random() +# invocations producing different orderings, so they must not be treated as the +# same expression. +query TT +EXPLAIN SELECT * FROM ( + SELECT t1.rand_col, t2.x + FROM ( + SELECT random() AS rand_col, a + FROM (SELECT a FROM t1 ORDER BY random() LIMIT 10) + ) t1 + LEFT JOIN t2 ON t1.a = t2.x +) sub +ORDER BY rand_col ASC LIMIT 2; +---- +logical_plan +01)Sort: sub.rand_col ASC NULLS LAST, fetch=2 +02)--SubqueryAlias: sub +03)----Projection: t1.rand_col, t2.x +04)------Left Join: t1.a = t2.x +05)--------SubqueryAlias: t1 +06)----------Projection: random() AS rand_col, t1.a +07)------------Sort: random() ASC NULLS LAST, fetch=10 +08)--------------TableScan: t1 projection=[a] +09)--------TableScan: t2 projection=[x] + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.explain.logical_plan_only; + +statement ok +DROP TABLE t1; + +statement ok +DROP TABLE t2; diff --git a/datafusion/sqllogictest/test_files/pwmj.slt b/datafusion/sqllogictest/test_files/pwmj.slt index 295eb94318ee5..9789c0e4e5392 100644 --- a/datafusion/sqllogictest/test_files/pwmj.slt +++ b/datafusion/sqllogictest/test_files/pwmj.slt @@ -39,9 +39,9 @@ query II SELECT t1.t1_id, t2.t2_id FROM join_t1 t1 JOIN join_t2 t2 - ON t1.t1_id > t2.t2_id -WHERE t1.t1_id > 10 - AND t2.t2_int > 1 + ON t1.t1_id > t2.t2_id +WHERE t1.t1_id > 10 + AND t2.t2_int > 1 ORDER BY 1; ---- 22 11 @@ -53,9 +53,9 @@ query IITI SELECT * FROM join_t1 t1 JOIN join_t2 t2 - ON t1.t1_id > t2.t2_id -WHERE t1.t1_id > 10 - AND t2.t2_int > 1 + ON t1.t1_id > t2.t2_id +WHERE t1.t1_id > 10 + AND t2.t2_int > 1 ORDER BY 1; ---- 22 11 z 3 @@ -67,9 +67,9 @@ EXPLAIN SELECT t1.t1_id, t2.t2_id FROM join_t1 t1 JOIN join_t2 t2 - ON t1.t1_id > t2.t2_id -WHERE t1.t1_id > 10 - AND t2.t2_int > 1 + ON t1.t1_id > t2.t2_id +WHERE t1.t1_id > 10 + AND t2.t2_int > 1 ORDER BY 1; ---- logical_plan @@ -326,8 +326,8 @@ logical_plan 06)------SubqueryAlias: t2 07)--------TableScan: null_join_t2 projection=[id] physical_plan -01)SortExec: expr=[left_id@0 ASC NULLS LAST, right_id@1 ASC NULLS LAST], preserve_partitioning=[false] -02)--ProjectionExec: expr=[id@0 as left_id, id@1 as right_id] +01)ProjectionExec: expr=[id@0 as left_id, id@1 as right_id] +02)--SortExec: expr=[id@0 ASC NULLS LAST, id@1 ASC NULLS LAST], preserve_partitioning=[false] 03)----NestedLoopJoinExec: join_type=Inner, filter=id@0 < id@0 + id@1 04)------DataSourceExec: partitions=1, partition_sizes=[1] 05)------DataSourceExec: partitions=1, partition_sizes=[1] @@ -339,8 +339,8 @@ JOIN null_join_t2 t2 ON t1.id < t2.id ORDER BY 1,2; ---- -1 3 -2 3 +1 3 +2 3 statement ok set datafusion.optimizer.enable_piecewise_merge_join = false; diff --git a/datafusion/sqllogictest/test_files/qualify.slt b/datafusion/sqllogictest/test_files/qualify.slt index ce58e3998cf57..b70f078327a42 100644 --- a/datafusion/sqllogictest/test_files/qualify.slt +++ b/datafusion/sqllogictest/test_files/qualify.slt @@ -39,8 +39,8 @@ CREATE TABLE users ( # Basic QUALIFY with ROW_NUMBER query ITI -SELECT id, name, ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) as rn -FROM users +SELECT id, name, ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) as rn +FROM users QUALIFY rn = 1 ORDER BY dept, id; ---- @@ -49,8 +49,8 @@ ORDER BY dept, id; # QUALIFY with RANK query ITI -SELECT id, name, RANK() OVER (ORDER BY salary DESC) as rank -FROM users +SELECT id, name, RANK() OVER (ORDER BY salary DESC) as rank +FROM users QUALIFY rank <= 3 ORDER BY rank, id; ---- @@ -60,8 +60,8 @@ ORDER BY rank, id; # QUALIFY with DENSE_RANK query ITI -SELECT id, name, DENSE_RANK() OVER (PARTITION BY dept ORDER BY age) as dense_rank -FROM users +SELECT id, name, DENSE_RANK() OVER (PARTITION BY dept ORDER BY age) as dense_rank +FROM users QUALIFY dense_rank <= 2 ORDER BY dept, dense_rank, id; ---- @@ -78,7 +78,7 @@ ORDER BY dept, dense_rank, id; query ITII SELECT id, name, ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) as rn, RANK() OVER (ORDER BY age) as age_rank -FROM users +FROM users QUALIFY rn <= 2 AND age_rank <= 5 ORDER BY dept, rn, id; ---- @@ -88,7 +88,7 @@ ORDER BY dept, rn, id; # QUALIFY with LAG function query ITRR SELECT id, name, salary, LAG(salary) OVER (PARTITION BY dept ORDER BY id) as prev_salary -FROM users +FROM users QUALIFY prev_salary IS NOT NULL AND salary > prev_salary ORDER BY dept, id; ---- @@ -99,7 +99,7 @@ ORDER BY dept, id; # QUALIFY with LEAD function query ITRR SELECT id, name, salary, LEAD(salary) OVER (PARTITION BY dept ORDER BY id) as next_salary -FROM users +FROM users QUALIFY next_salary IS NOT NULL AND salary < next_salary ORDER BY dept, id; ---- @@ -110,7 +110,7 @@ ORDER BY dept, id; # QUALIFY with NTILE query ITI SELECT id, name, NTILE(3) OVER (PARTITION BY dept ORDER BY salary DESC) as tile -FROM users +FROM users QUALIFY tile = 1 ORDER BY dept, id; ---- @@ -121,7 +121,7 @@ ORDER BY dept, id; # QUALIFY with PERCENT_RANK query ITR SELECT id, name, PERCENT_RANK() OVER (PARTITION BY dept ORDER BY salary) as pct_rank -FROM users +FROM users QUALIFY pct_rank >= 0.5 ORDER BY dept, pct_rank, id; ---- @@ -134,7 +134,7 @@ ORDER BY dept, pct_rank, id; # QUALIFY with CUME_DIST query ITR SELECT id, name, CUME_DIST() OVER (PARTITION BY dept ORDER BY age) as cume_dist -FROM users +FROM users QUALIFY cume_dist >= 0.75 ORDER BY dept, cume_dist, id; ---- @@ -145,11 +145,11 @@ ORDER BY dept, cume_dist, id; # QUALIFY with multiple window functions query ITIII -SELECT id, name, +SELECT id, name, ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) as rn, RANK() OVER (ORDER BY age) as age_rank, DENSE_RANK() OVER (PARTITION BY dept ORDER BY age) as dept_age_rank -FROM users +FROM users QUALIFY rn <= 2 AND age_rank <= 4 AND dept_age_rank <= 2 ORDER BY dept, rn, id; ---- @@ -158,9 +158,9 @@ ORDER BY dept, rn, id; # QUALIFY with arithmetic expressions query ITRI -SELECT id, name, salary, +SELECT id, name, salary, ROW_NUMBER() OVER (PARTITION BY dept ORDER BY salary DESC) as rn -FROM users +FROM users QUALIFY rn = 1 AND salary > 60000 ORDER BY dept, id; ---- @@ -169,9 +169,9 @@ ORDER BY dept, id; # QUALIFY with string functions query ITI -SELECT id, name, +SELECT id, name, ROW_NUMBER() OVER (PARTITION BY dept ORDER BY name) as rn -FROM users +FROM users QUALIFY rn = 1 ORDER BY dept, id; ---- @@ -181,7 +181,7 @@ ORDER BY dept, id; # window function with aggregate function query ITI SELECT id, name, COUNT(*) OVER (PARTITION BY dept) as cnt -FROM users +FROM users QUALIFY cnt > 4 ORDER BY dept, id; ---- @@ -198,7 +198,7 @@ FROM users WHERE salary > 5000 GROUP BY dept, salary HAVING SUM(salary) > 20000 -QUALIFY r > 60000 +QUALIFY r > 60000 ---- Marketing 70000 Marketing 70000 @@ -306,27 +306,27 @@ QUALIFY r > 60000 ---- logical_plan 01)Projection: users.dept, avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS r -02)--Filter: avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING > Decimal128(Some(60000000000),14,6) +02)--Filter: avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING > Decimal128(60000.000000,14,6) 03)----Projection: users.dept, avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING 04)------WindowAggr: windowExpr=[[avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] 05)--------Projection: users.dept, users.salary -06)----------Filter: sum(users.salary) > Decimal128(Some(2000000),20,2) +06)----------Filter: sum(users.salary) > Decimal128(20000.00,20,2) 07)------------Aggregate: groupBy=[[users.dept, users.salary]], aggr=[[sum(users.salary)]] -08)--------------Filter: users.salary > Decimal128(Some(500000),10,2) +08)--------------Filter: users.salary > Decimal128(5000.00,10,2) 09)----------------TableScan: users projection=[salary, dept] physical_plan 01)ProjectionExec: expr=[dept@0 as dept, avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as r] -02)--FilterExec: avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 > Some(60000000000),14,6 +02)--FilterExec: avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 > 60000.000000 03)----ProjectionExec: expr=[dept@0 as dept, avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] 04)------WindowAggExec: wdw=[avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "avg(users.salary) PARTITION BY [users.dept] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Decimal128(14, 6), nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] 05)--------SortExec: expr=[dept@0 ASC NULLS LAST], preserve_partitioning=[true] 06)----------RepartitionExec: partitioning=Hash([dept@0], 4), input_partitions=4 -07)------------FilterExec: sum(users.salary)@2 > Some(2000000),20,2, projection=[dept@0, salary@1] +07)------------FilterExec: sum(users.salary)@2 > 20000.00, projection=[dept@0, salary@1] 08)--------------AggregateExec: mode=FinalPartitioned, gby=[dept@0 as dept, salary@1 as salary], aggr=[sum(users.salary)] 09)----------------RepartitionExec: partitioning=Hash([dept@0, salary@1], 4), input_partitions=4 10)------------------AggregateExec: mode=Partial, gby=[dept@1 as dept, salary@0 as salary], aggr=[sum(users.salary)] 11)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -12)----------------------FilterExec: salary@0 > Some(500000),10,2 +12)----------------------FilterExec: salary@0 > 5000.00 13)------------------------DataSourceExec: partitions=1, partition_sizes=[1] # plan with aggregate function @@ -360,4 +360,4 @@ physical_plan # Clean up statement ok -DROP TABLE users; +DROP TABLE users; diff --git a/datafusion/sqllogictest/test_files/range_partitioning.slt b/datafusion/sqllogictest/test_files/range_partitioning.slt new file mode 100644 index 0000000000000..326856a352f36 --- /dev/null +++ b/datafusion/sqllogictest/test_files/range_partitioning.slt @@ -0,0 +1,1919 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# The sqllogictest harness registers range_partitioned(range_key, non_range_key, value) +# as a Parquet ListingTable with four declared range-partitioned file groups: +# +# partition 0: range_key in [..., 10), rows (1, 1, 10), (5, 2, 50) +# partition 1: range_key in [10, 20), rows (10, 1, 100), (15, 2, 150) +# partition 2: range_key in [20, 30), rows (20, 1, 200), (25, 2, 250) +# partition 3: range_key in [30, ...), rows (30, 1, 300), (35, 2, 350) + +statement ok +set datafusion.explain.physical_plan_only = true; + +statement ok +set datafusion.execution.collect_statistics = false; + +statement ok +set datafusion.optimizer.enable_join_dynamic_filter_pushdown = false; + +statement ok +set datafusion.optimizer.enable_topk_dynamic_filter_pushdown = false; + +statement ok +set datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown = false; + +statement ok +set datafusion.execution.parquet.pushdown_filters = false; + +########## +# TEST 1: Aggregate on Range Partition Column +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies the aggregate key and avoids repartitioning. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 2: Aggregate on Non-Range Column +# With subset threshold met and preserve-file disabled, grouping on a non-range +# key cannot reuse Range([range_key]) and requires hash repartitioning. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----AggregateExec: mode=Partial, gby=[non_range_key@0 as non_range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet + +query II +SELECT non_range_key, SUM(value) FROM range_partitioned GROUP BY non_range_key ORDER BY non_range_key; +---- +1 610 +2 800 + + +########## +# TEST 3: Aggregate Reuses Range Subset Partitioning +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies grouping by (range_key, non_range_key). +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key, non_range_key@1 as non_range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT range_key, non_range_key, SUM(value) FROM range_partitioned GROUP BY range_key, non_range_key ORDER BY range_key, non_range_key; +---- +1 1 10 +5 2 50 +10 1 100 +15 2 150 +20 1 200 +25 2 250 +30 1 300 +35 2 350 + + +########## +# TEST 4: Aggregate Preserves Range When Preserve File Threshold Met +# With preserve-file threshold 1 and 4 input partitions, Range is preserved +# even though target_partitions is 5. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 5: Aggregate Rehashes When Preserve File Threshold Not Met +# With preserve-file threshold 5 and only 4 input partitions, planning can +# repartition to increase parallelism. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 5; + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM range_partitioned GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=FinalPartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=5 +03)----AggregateExec: mode=Partial, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------RepartitionExec: partitioning=RoundRobinBatch(5), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +statement ok +set datafusion.optimizer.repartition_joins = true; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + + +########## +# TEST 6: Join on Range Partition Column +# A partitioned inner hash join requires co-partitioned KeyPartitioned inputs. +# Compatible Range layouts satisfy both the per-child key requirements and the +# cross-child layout requirement, so no Hash repartitioning is inserted. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 7: Incompatible Range Join Repartitions +# Both inputs are independently range partitioned on range_key, but their split +# points differ. The per-child key requirements can be satisfied by Range, but +# the co-partitioned layout requirement cannot, so Hash repartitioning repairs +# both sides. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 8: Non-Range Join Repartitions +# Range([range_key]) does not satisfy KeyPartitioned([non_range_key]), so +# planning inserts Hash repartitioning on the actual join key. +########## + +query TT +EXPLAIN SELECT l.non_range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(non_range_key@0, non_range_key@0)], projection=[non_range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet + +query III +SELECT l.non_range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.non_range_key = r.non_range_key +ORDER BY l.non_range_key, l.value, r.value; +---- +1 10 10 +1 10 100 +1 10 200 +1 10 300 +1 100 10 +1 100 100 +1 100 200 +1 100 300 +1 200 10 +1 200 100 +1 200 200 +1 200 300 +1 300 10 +1 300 100 +1 300 200 +1 300 300 +2 50 50 +2 50 150 +2 50 250 +2 50 350 +2 150 50 +2 150 150 +2 150 250 +2 150 350 +2 250 50 +2 250 150 +2 250 250 +2 250 350 +2 350 50 +2 350 150 +2 350 250 +2 350 350 + +########## +# TEST 9: Left-Side Range Hash Joins +# Compatible Range layouts satisfy left-side partitioned hash join +# requirements without Hash repartitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN (SELECT range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--FilterExec: value@1 <= 150 +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN (SELECT range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 NULL +25 250 NULL +30 300 NULL +35 350 NULL + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +LEFT SEMI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(range_key@0, range_key@0)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--FilterExec: value@1 <= 150, projection=[range_key@0] +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +LEFT SEMI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +LEFT ANTI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(range_key@0, range_key@0)] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--FilterExec: value@1 <= 150, projection=[range_key@0] +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +LEFT ANTI JOIN (SELECT range_key FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +20 200 +25 250 +30 300 +35 350 + +########## +# TEST 10: Left-Side Range Hash Joins With Incomplete Range Keys +# Range partitioning covers only range_key, so joins requiring additional +# or different keys are repaired with Hash repartitioning. +########## + +# Range([range_key]) is only a subset of the composite join key, so the +# co-partitioned hash join requirement is repaired with Hash repartitioning. +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN (SELECT range_key, non_range_key, value FROM range_partitioned WHERE value <= 150) r +ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0), (non_range_key@1, non_range_key@1)], projection=[range_key@0, non_range_key@1, value@2, value@5] +02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)----FilterExec: value@2 <= 150 +06)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +# Range([range_key]) does not satisfy a join keyed on non_range_key. +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(non_range_key@1, non_range_key@0)], projection=[range_key@0, non_range_key@1, value@2, value@4] +02)--RepartitionExec: partitioning=Hash([non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet + +########## +# TEST 11: Left-Side Range Hash Joins With Incompatible Range Layouts +# Different split points or partition counts do not satisfy the +# co-partitioned layout requirement. +########## + +# Different split points do not satisfy the co-partitioned layout requirement. +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +# Different partition counts do not satisfy the co-partitioned layout +# requirement. +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +LEFT JOIN range_partitioned_narrow r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Left, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=3 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20)], 3), file_type=parquet + +########## +# TEST 12: LeftMark Subqueries Over Range Hash Joins +# SQL IN subqueries decorrelate to LeftMark joins. These queries pin matched, +# unmatched, and NULL marker behavior over compatible Range inputs. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150); +---- +physical_plan +01)FilterExec: non_range_key@1 = 2 OR mark@3, projection=[range_key@0, value@2] +02)--HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(range_key@0, range_key@0)] +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)----FilterExec: value@1 <= 150, projection=[range_key@0] +05)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150) +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +query II +SELECT l.range_key, l.value +FROM range_partitioned l +WHERE l.non_range_key = 2 OR l.range_key IN ( + SELECT CASE WHEN value <= 150 THEN range_key ELSE NULL END + FROM range_partitioned) +ORDER BY l.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +########## +# TEST 13: Compatible Range Join Repartitions to Increase Parallelism +# Co-partitioning satisfaction does not prevent a repartition that increases +# parallelism. With target_partitions larger than the Range partition count, +# both sides are hash repartitioned. +########## + +statement ok +set datafusion.execution.target_partitions = 5; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([range_key@0], 5), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 14: Preserve File Partitions Preserves Range Join Inputs +# preserve_file_partitions preserves compatible Range inputs for partitioned +# joins even when target_partitions is higher than the input partition count. +########## + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +########## +# TEST 15: Nested Range Joins +# Compatible Range partitioning is preserved through the lower join, allowing +# the upper join to consume it without Hash repartitioning either input. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value, s.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +JOIN range_partitioned s ON r.range_key = s.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@2, range_key@0)], projection=[range_key@0, value@1, value@3, value@5] +02)--HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +05)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query IIII +SELECT l.range_key, l.value, r.value, s.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +JOIN range_partitioned s ON r.range_key = s.range_key +ORDER BY l.range_key; +---- +1 10 10 10 +5 50 50 50 +10 100 100 100 +15 150 150 150 +20 200 200 200 +25 250 250 250 +30 300 300 300 +35 350 350 350 + +########## +# TEST 16: Range Aggregates Feed Range Join +# Aggregates on range_key preserve reusable partitioning for the downstream +# partitioned join. +########## + +query TT +EXPLAIN WITH + l AS ( + SELECT range_key, SUM(value) AS l_sum + FROM range_partitioned + GROUP BY range_key + ), + r AS ( + SELECT range_key, SUM(value) AS r_sum + FROM range_partitioned + GROUP BY range_key + ) +SELECT l.range_key, l.l_sum, r.r_sum +FROM l JOIN r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, l_sum@1, r_sum@3] +02)--ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value)@1 as l_sum] +03)----AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +05)--ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value)@1 as r_sum] +06)----AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(range_partitioned.value)] +07)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +WITH + l AS ( + SELECT range_key, SUM(value) AS l_sum + FROM range_partitioned + GROUP BY range_key + ), + r AS ( + SELECT range_key, SUM(value) AS r_sum + FROM range_partitioned + GROUP BY range_key + ) +SELECT l.range_key, l.l_sum, r.r_sum +FROM l JOIN r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 17: Range Join Feeds Aggregate +# The join preserves compatible Range partitioning on range_key, allowing the +# aggregate above it to avoid Hash repartitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, SUM(l.value + r.value) +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +GROUP BY l.range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(l.value + r.value)] +02)--HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT l.range_key, SUM(l.value + r.value) +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +GROUP BY l.range_key +ORDER BY l.range_key; +---- +1 20 +5 100 +10 200 +15 300 +20 400 +25 500 +30 600 +35 700 + +########## +# TEST 18: Right Join on Range Partition Column +# Compatible Range inputs satisfy the join's partitioning requirements, so no +# Hash repartitioning is inserted. The left filter keeps its Range partitioning +# and the unmatched right rows above 150 are preserved. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--FilterExec: value@1 <= 150 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=value@2 <= 150, pruning_predicate=value_null_count@1 != row_count@2 AND value_min@0 <= 150, required_guarantees=[] +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +NULL 20 200 +NULL 25 250 +NULL 30 300 +NULL 35 350 + +########## +# TEST 19: Right Semi Join on Range Partition Column +# Compatible Range inputs avoid Hash repartitioning for RightSemi joins. +# Only right rows with a match on the filtered left side are returned. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT SEMI JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(range_key@0, range_key@0)] +02)--FilterExec: value@1 <= 150, projection=[range_key@0] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=value@2 <= 150, pruning_predicate=value_null_count@1 != row_count@2 AND value_min@0 <= 150, required_guarantees=[] +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT SEMI JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 + +########## +# TEST 20: Right Anti Join on Range Partition Column +# Compatible Range inputs avoid Hash repartitioning for RightAnti joins. +# Only right rows without a match on the filtered left side are returned. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT ANTI JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=RightAnti, on=[(range_key@0, range_key@0)] +02)--FilterExec: value@1 <= 150, projection=[range_key@0] +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=value@2 <= 150, pruning_predicate=value_null_count@1 != row_count@2 AND value_min@0 <= 150, required_guarantees=[] +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT r.range_key, r.value +FROM (SELECT range_key FROM range_partitioned WHERE value <= 150) l +RIGHT ANTI JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +20 200 +25 250 +30 300 +35 350 + +########## +# TEST 21: Incompatible Range Right Join Repartitions +# The split points of the two inputs differ, so the co-partitioned layout +# requirement cannot be satisfied and Hash repartitioning repairs both sides +# of the right join. Results stay correct on the repartitioned path. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----FilterExec: value@1 <= 150 +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=value@2 <= 150, pruning_predicate=value_null_count@1 != row_count@2 AND value_min@0 <= 150, required_guarantees=[] +05)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +06)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT range_key, value FROM range_partitioned WHERE value <= 150) l +RIGHT JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +NULL 20 200 +NULL 25 250 +NULL 30 300 +NULL 35 350 + +########## +# TEST 22: Composite-Key Right Join Repartitions +# Range([range_key]) does not satisfy a partitioned join on +# (range_key, non_range_key), so both sides repartition on the full key. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +query TT +EXPLAIN SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0), (non_range_key@1, non_range_key@1)], projection=[range_key@0, non_range_key@1, value@2, value@5] +02)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query IIII +SELECT l.range_key, l.non_range_key, l.value, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned r ON l.range_key = r.range_key AND l.non_range_key = r.non_range_key +ORDER BY l.range_key; +---- +1 1 10 10 +5 2 50 50 +10 1 100 100 +15 2 150 150 +20 1 200 200 +25 2 250 250 +30 1 300 300 +35 2 350 350 + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +########## +# TEST 23: Right Join with Mismatched Range Partition Counts Repartitions +# Both inputs are range partitioned on range_key, but declare a different number +# of partitions (four vs three). The per-child key requirements can be satisfied +# by Range, but the co-partitioned layout requirement cannot, so Hash +# repartitioning repairs both sides of the right join. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned_narrow r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(range_key@0, range_key@0)], projection=[value@1, range_key@2, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=3 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20)], 3), file_type=parquet + +query III +SELECT l.value, r.range_key, r.value +FROM range_partitioned l +RIGHT JOIN range_partitioned_narrow r ON l.range_key = r.range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +100 10 100 +150 15 150 +200 20 200 +250 25 250 +300 30 300 +350 35 350 + +########## +# TEST 24: Right Join on Non-Range Key Repartitions +# Both inputs expose Range([range_key]), but the join key is non_range_key. +# Range([range_key]) does not satisfy KeyPartitioned([non_range_key]), so +# planning inserts Hash repartitioning on the actual join key for the right join. +########## + +query TT +EXPLAIN SELECT l.value, r.range_key, r.value +FROM (SELECT non_range_key, value FROM range_partitioned WHERE range_key < 10) l +RIGHT JOIN range_partitioned r ON l.non_range_key = r.non_range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Right, on=[(non_range_key@0, non_range_key@1)], projection=[value@1, range_key@2, value@4] +02)--RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +03)----FilterExec: range_key@0 < 10, projection=[non_range_key@1, value@2] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=range_key@0 < 10, pruning_predicate=range_key_null_count@1 != row_count@2 AND range_key_min@0 < 10, required_guarantees=[] +05)--RepartitionExec: partitioning=Hash([non_range_key@1], 4), input_partitions=4 +06)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT l.value, r.range_key, r.value +FROM (SELECT non_range_key, value FROM range_partitioned WHERE range_key < 10) l +RIGHT JOIN range_partitioned r ON l.non_range_key = r.non_range_key +ORDER BY r.range_key; +---- +10 1 10 +50 5 50 +10 10 100 +50 15 150 +10 20 200 +50 25 250 +10 30 300 +50 35 350 + +########## +# TEST 25: Mark Join Marker Semantics +# Mark joins preserve matched, unmatched, and NULL-key marker behavior over +# range-partitioned inputs. +########## + +query TT +EXPLAIN SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150); +---- +physical_plan +01)FilterExec: non_range_key@1 = 2 OR mark@3, projection=[range_key@0, value@2] +02)--HashJoinExec: mode=Partitioned, join_type=LeftMark, on=[(range_key@0, range_key@0)] +03)----DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)----FilterExec: value@1 <= 150, projection=[range_key@0] +05)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +# Matched rows have mark=true and are returned; unmatched rows have +# mark=false and are only returned when non_range_key = 2. +query II +SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT range_key FROM range_partitioned WHERE value <= 150) +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +# NULL join keys on the build side never match: rows whose keys only "match" +# the NULL entries keep a non-true marker and are filtered out unless the +# non_range_key = 2 disjunct covers them. +query II +SELECT r.range_key, r.value +FROM range_partitioned r +WHERE r.non_range_key = 2 OR r.range_key IN ( + SELECT CASE WHEN value <= 150 THEN range_key ELSE NULL END FROM range_partitioned) +ORDER BY r.range_key; +---- +1 10 +5 50 +10 100 +15 150 +25 250 +35 350 + +########## +# TEST 26: Sort Merge Join Avoids Repartition for Compatible Range Inputs +# Compatible Range inputs satisfy SortMergeJoinExec's co-partitioned +# KeyPartitioned requirements. +########## + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] +03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, sort_order_for_reorder=[range_key@0 ASC] +05)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +06)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, sort_order_for_reorder=[range_key@0 ASC] + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 27: Sort Merge Join Repartitions Incompatible Range Inputs +# Different Range split points do not satisfy SortMergeJoinExec's +# co-partitioned KeyPartitioned requirements. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SortMergeJoinExec: join_type=Inner, on=[(range_key@0, range_key@0)] +03)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +06)----SortExec: expr=[range_key@0 ASC], preserve_partitioning=[true] +07)------RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +08)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +########## +# TEST 28: Symmetric Hash Join Avoids Repartition for Compatible Range Inputs +# Compatible Range streams satisfy SymmetricHashJoinExec's co-partitioned +# KeyPartitioned requirements. +########## + +statement ok +set datafusion.optimizer.prefer_hash_join = true; + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +03)----StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +04)----StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) + +query III rowsort +SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like r ON l.range_key = r.range_key; +---- +1 10 10 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 +5 50 50 + +########## +# TEST 29: Symmetric Hash Join Repartitions Incompatible Range Inputs +# Different Range split points do not satisfy SymmetricHashJoinExec's +# co-partitioned KeyPartitioned requirements. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, value@3 as value] +02)--SymmetricHashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)] +03)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +04)------StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4) +05)----RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +06)------StreamingTableExec: partition_sizes=4, projection=[range_key, value], infinite_source=true, output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4) + +query III rowsort +SELECT l.range_key, l.value, r.value +FROM unbounded_range_like l +FULL JOIN unbounded_range_like_shifted r ON l.range_key = r.range_key; +---- +1 10 10 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 +5 50 50 + +########## +# TEST 30: Full Outer Join on Range Partition Column +# Full partitioned hash joins also opt in to Range satisfying KeyPartitioned +# requirements, so compatible Range layouts avoid Hash repartitioning here too. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 31: Full Outer Join Incompatible Range Repartitions +# For Full joins, differing split points between the two Range-partitioned +# inputs still require Hash repartitioning to co-partition. +########## + +query TT +EXPLAIN SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_shifted r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--RepartitionExec: partitioning=Hash([range_key@0], 4), input_partitions=4 +05)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet + +query III +SELECT l.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_shifted r ON l.range_key = r.range_key +ORDER BY l.range_key; +---- +1 10 10 +5 50 50 +10 100 100 +15 150 150 +20 200 200 +25 250 250 +30 300 300 +35 350 350 + +########## +# TEST 32: Full Outer Join Produces Matched and Unmatched Rows +# `range_partitioned` and `range_partitioned_sparse` share the same Range +# split points/partition count but only partially overlapping range_key +# values, so this exercises matched rows, left-only unmatched rows (NULLs on +# the right), and right-only unmatched rows (NULLs on the left) while still +# avoiding Hash repartitioning. +########## + +query TT +EXPLAIN SELECT l.range_key, r.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Full, on=[(range_key@0, range_key@0)], projection=[range_key@0, range_key@2, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query IIII +SELECT l.range_key, r.range_key, l.value, r.value +FROM range_partitioned l +FULL JOIN range_partitioned_sparse r ON l.range_key = r.range_key +ORDER BY l.range_key, r.range_key; +---- +1 NULL 10 NULL +5 5 50 50 +10 10 100 100 +15 NULL 150 NULL +20 20 200 200 +25 NULL 250 NULL +30 30 300 300 +35 NULL 350 NULL +NULL 8 NULL 80 +NULL 40 NULL 400 + +statement ok +reset datafusion.optimizer.prefer_hash_join; + +statement ok +reset datafusion.optimizer.repartition_joins; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +########## +# TEST 33: Union of Range Partitioned Inputs +# Each input exposes the same Range partitioning on range_key, so the optimizer +# converts UnionExec to InterleaveExec to avoid redundant repartitioning. +########## + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned; +---- +physical_plan +01)InterleaveExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +ORDER BY range_key, value; +---- +1 10 +1 10 +5 50 +5 50 +10 100 +10 100 +15 150 +15 150 +20 200 +20 200 +25 250 +25 250 +30 300 +30 300 +35 350 +35 350 + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + + +########## +# TEST 34: Window on Range Partition Column +# Range([range_key]) colocates equal range_key values, so +# PARTITION BY range_key is satisfied without a hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 35: Unbounded-Frame Window on Range Partition Column +# The unbounded frame makes DataFusion use WindowAggExec instead of +# BoundedWindowAggExec, which likewise reuses Range partitioning without a +# hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] +02)--WindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 36: Window on Non-Range Column Rehashes +# Range([range_key]) does not colocate non_range_key values, so +# PARTITION BY non_range_key still requires a hash repartition. +########## + +query TT +EXPLAIN SELECT non_range_key, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[non_range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet + +query III +SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value) FROM range_partitioned ORDER BY non_range_key, value; +---- +1 10 10 +1 100 110 +1 200 310 +1 300 610 +2 50 50 +2 150 200 +2 250 450 +2 350 800 + + +########## +# TEST 37: Unbounded-Frame Window on Non-Range Column Rehashes +# The unbounded frame makes DataFusion use WindowAggExec; Range([range_key]) +# does not colocate non_range_key values, so PARTITION BY non_range_key +# still requires a hash repartition. +########## + +query TT +EXPLAIN SELECT non_range_key, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@2 as sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING] +02)--WindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(range_partitioned.value) PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING", data_type: Int64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(NULL)), end_bound: Following(UInt64(NULL)), is_causal: false }] +03)----SortExec: expr=[non_range_key@0 ASC NULLS LAST, value@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet + +query III +SELECT non_range_key, value, SUM(value) OVER (PARTITION BY non_range_key ORDER BY value ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) FROM range_partitioned ORDER BY non_range_key, value; +---- +1 10 610 +1 100 610 +1 200 610 +1 300 610 +2 50 800 +2 150 800 +2 250 800 +2 350 800 + + +########## +# TEST 38: Window Subset Satisfaction on Range Partition Column +# With the subset threshold met, Range([range_key]) satisfies +# PARTITION BY (range_key, non_range_key): equal composite keys share the +# same range_key, so they are already colocated. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, non_range_key@1 ASC NULLS LAST, value@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + + +########## +# TEST 39: Window Subset Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of PARTITION BY +# (range_key, non_range_key), so it should not satisfy the window key when +# subset satisfaction is disabled. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortExec: expr=[range_key@0 ASC NULLS LAST, non_range_key@1 ASC NULLS LAST, value@2 ASC NULLS LAST], preserve_partitioning=[true] +04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, SUM(value) OVER (PARTITION BY range_key, non_range_key ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 50 +10 100 +15 150 +20 200 +25 250 +30 300 +35 350 + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + + +########## +# TEST 40: Window Without Partition Keys Uses a Single Partition +# A window with no PARTITION BY requires a single partition; range +# partitioning is not applicable. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW] +02)--BoundedWindowAggExec: wdw=[sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(range_partitioned.value) ORDER BY [range_partitioned.value ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortPreservingMergeExec: [value@1 ASC NULLS LAST] +04)------SortExec: expr=[value@1 ASC NULLS LAST], preserve_partitioning=[true] +05)--------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, SUM(value) OVER (ORDER BY value) FROM range_partitioned ORDER BY range_key; +---- +1 10 +5 60 +10 160 +15 310 +20 510 +25 760 +30 1060 +35 1410 + + + +########## +# TEST 41: PartitionedTopK on Range Partition Column +# Exact Range([range_key]) satisfies the TopK partition key and avoids repartitioning. +########## + +statement ok +set datafusion.optimizer.enable_window_topn = true; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0], order=[value@1 DESC] +04)------DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query III +SELECT * FROM ( + SELECT range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY range_key; +---- +1 10 1 +5 50 1 +10 100 1 +15 150 1 +20 200 1 +25 250 1 +30 300 1 +35 350 1 + + +########## +# TEST 42: PartitionedTopK on Non-Range Column +# Partitioning on a non-range key cannot reuse Range([range_key]) and +# requires hash repartitioning. +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[non_range_key@0 as non_range_key, value@1 as value, row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[non_range_key@0], order=[value@1 DESC] +04)------RepartitionExec: partitioning=Hash([non_range_key@0], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[non_range_key, value], output_partitioning=UnknownPartitioning(4), file_type=parquet + +query III +SELECT * FROM ( + SELECT non_range_key, value, ROW_NUMBER() OVER (PARTITION BY non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY non_range_key; +---- +1 300 1 +2 350 1 + + +########## +# TEST 43: PartitionedTopK Reuses Range Subset Partitioning +# With subset threshold met and preserve-file disabled, Range([range_key]) +# satisfies partitioning by (range_key, non_range_key). +########## + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query IIII +SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1 +ORDER BY range_key, non_range_key; +---- +1 1 10 1 +5 2 50 1 +10 1 100 1 +15 2 150 1 +20 1 200 1 +25 2 250 1 +30 1 300 1 +35 2 350 1 + + +########## +# TEST 44: Range Subset PartitionedTopK Rehashes Below Subset Threshold +# Range([range_key]) is only a subset of PARTITION BY (range_key, non_range_key), +# so it should not satisfy the TopK partition key when subset satisfaction is +# disabled. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.subset_repartition_threshold = 5; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 0; + +query TT +EXPLAIN SELECT * FROM ( + SELECT range_key, non_range_key, value, ROW_NUMBER() OVER (PARTITION BY range_key, non_range_key ORDER BY value DESC) as rn + FROM range_partitioned +) WHERE rn <= 1; +---- +physical_plan +01)ProjectionExec: expr=[range_key@0 as range_key, non_range_key@1 as non_range_key, value@2 as value, row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [range_partitioned.range_key, range_partitioned.non_range_key] ORDER BY [range_partitioned.value DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----PartitionedTopKExec: fn=row_number, fetch=1, partition=[range_key@0, non_range_key@1], order=[value@2 DESC] +04)------RepartitionExec: partitioning=Hash([range_key@0, non_range_key@1], 4), input_partitions=4 +05)--------DataSourceExec: file_groups=, projection=[range_key, non_range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +statement ok +reset datafusion.optimizer.subset_repartition_threshold; + +statement ok +reset datafusion.explain.physical_plan_only; + +statement ok +reset datafusion.optimizer.enable_window_topn; + +########## +# TEST 45: Subset of Inputs Compatible Does Not Trigger InterleaveExec +# In a three-way union, two inputs share the same Range split points [10,20,30] +# while the third has a partially-overlapping but different set [15,20,30]. +# can_interleave requires ALL inputs to match, so UnionExec is kept. +########## + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted +ORDER BY range_key, value; +---- +1 10 +1 10 +1 10 +5 50 +5 50 +5 50 +10 100 +10 100 +10 100 +15 150 +15 150 +15 150 +20 200 +20 200 +20 200 +25 250 +25 250 +25 250 +30 300 +30 300 +30 300 +35 350 +35 350 +35 350 + +########## +# TEST 46: Incompatible Range Split Points Falls Back to UnionExec +# Two range-partitioned inputs with different split points cannot be interleaved, +# so the optimizer keeps UnionExec instead of converting to InterleaveExec. +########## + +statement ok +set datafusion.explain.physical_plan_only = true; + +query TT +EXPLAIN SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted; +---- +physical_plan +01)UnionExec +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(15), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, value FROM range_partitioned +UNION ALL +SELECT range_key, value FROM range_partitioned_shifted +ORDER BY range_key, value; +---- +1 10 +1 10 +5 50 +5 50 +10 100 +10 100 +15 150 +15 150 +20 200 +20 200 +25 250 +25 250 +30 300 +30 300 +35 350 +35 350 + +########## +# TEST 47: InterleaveExec Propagates Range Partitioning to Aggregate +# InterleaveExec outputs the same Range partitioning as its compatible inputs, +# allowing a downstream aggregate on range_key to run SinglePartitioned without +# a Hash repartition. +########## + +query TT +EXPLAIN SELECT range_key, SUM(value) FROM ( + SELECT range_key, value FROM range_partitioned + UNION ALL + SELECT range_key, value FROM range_partitioned +) GROUP BY range_key; +---- +physical_plan +01)AggregateExec: mode=SinglePartitioned, gby=[range_key@0 as range_key], aggr=[sum(value)] +02)--InterleaveExec +03)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet +04)----DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet + +query II +SELECT range_key, SUM(value) FROM ( + SELECT range_key, value FROM range_partitioned + UNION ALL + SELECT range_key, value FROM range_partitioned +) GROUP BY range_key ORDER BY range_key; +---- +1 20 +5 100 +10 200 +15 300 +20 400 +25 500 +30 600 +35 700 + +########## +# TEST 48: Hash Join Dynamic Filter Pushdown on Compatible Range Inputs +# The Parquet-backed probe accepts the partition-routed dynamic filter. Matching +# Range split points keep build filter i aligned with probe partition i. The +# build has rows only in partitions 0 and 2, so the runtime filter must route +# with a four-way CASE rather than collapse to a single filter. +########## + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.preserve_file_partitions = 1; + +statement ok +set datafusion.optimizer.enable_join_dynamic_filter_pushdown = true; + +statement ok +set datafusion.execution.parquet.pushdown_filters = true; + +query TT +EXPLAIN SELECT b.range_key, b.value, p.value +FROM ( + SELECT range_key, value + FROM range_partitioned + WHERE range_key IN (5, 20) +) b +JOIN range_partitioned p ON b.range_key = p.range_key; +---- +physical_plan +01)HashJoinExec: mode=Partitioned, join_type=Inner, on=[(range_key@0, range_key@0)], projection=[range_key@0, value@1, value@3] +02)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=range_key@0 = 5 OR range_key@0 = 20, pruning_predicate=range_key_null_count@2 != row_count@3 AND range_key_min@0 <= 5 AND 5 <= range_key_max@1 OR range_key_null_count@2 != row_count@3 AND range_key_min@0 <= 20 AND 20 <= range_key_max@1, required_guarantees=[range_key in (20, 5)] +03)--DataSourceExec: file_groups=, projection=[range_key, value], output_partitioning=Range([range_key@0 ASC], [(10), (20), (30)], 4), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible + +query TT +EXPLAIN ANALYZE SELECT b.range_key, b.value, p.value +FROM ( + SELECT range_key, value + FROM range_partitioned + WHERE range_key IN (5, 20) +) b +JOIN range_partitioned p ON b.range_key = p.range_key; +---- +Plan with Metrics +01)HashJoinExec: mode=Partitionedmetrics=[output_rows=2,] +02)--DataSourceExec: file_type=parquet, predicate=range_key@0 = 5 OR range_key@0 = 20metrics=[output_rows=2,] +03)--DataSourceExec: file_type=parquet, predicate=DynamicFilter [ CASE range_partition WHEN 0 THEN range_key@0 >= 5 AND range_key@0 <= 5 AND range_key@0 IN (SET) ([5]) WHEN 1 THEN false WHEN 2 THEN range_key@0 >= 20 AND range_key@0 <= 20 AND range_key@0 IN (SET) ([20]) ELSE false END ]metrics=[output_rows=2,] + +query III +SELECT b.range_key, b.value, p.value +FROM ( + SELECT range_key, value + FROM range_partitioned + WHERE range_key IN (5, 20) +) b +JOIN range_partitioned p ON b.range_key = p.range_key +ORDER BY b.range_key; +---- +5 50 50 +20 200 200 + +statement ok +reset datafusion.optimizer.preserve_file_partitions; + +statement ok +reset datafusion.execution.collect_statistics; + +statement ok +reset datafusion.optimizer.enable_join_dynamic_filter_pushdown; + +statement ok +reset datafusion.optimizer.enable_topk_dynamic_filter_pushdown; + +statement ok +reset datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown; + +statement ok +reset datafusion.execution.parquet.pushdown_filters; + +statement ok +reset datafusion.explain.physical_plan_only; diff --git a/datafusion/sqllogictest/test_files/references.slt b/datafusion/sqllogictest/test_files/references.slt index 0e72c5e5a29e9..3da1b385ee51b 100644 --- a/datafusion/sqllogictest/test_files/references.slt +++ b/datafusion/sqllogictest/test_files/references.slt @@ -66,7 +66,7 @@ CREATE TABLE test("f.c1" TEXT, "test.c2" INT, "...." INT) AS VALUES ('foobar', 2, 20), ('foobaz', 3, 30); -query error DataFusion error: Schema error: No field named f1\.c1\. Valid fields are test\."f\.c1", test\."test\.c2", test\."\.\.\.\."\. +query error DataFusion error: Schema error: No field named f1\.c1\. Did you mean 'test\."f\.c1"'\?\nValid fields are test\."f\.c1", test\."test\.c2", test\."\.\.\.\."\. SELECT f1.c1 FROM test; query T @@ -105,8 +105,8 @@ logical_plan 02)--Projection: test....., test..... AS c3 03)----TableScan: test projection=[....] physical_plan -01)SortExec: expr=[....@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--ProjectionExec: expr=[....@0 as ...., ....@0 as c3] +01)ProjectionExec: expr=[....@0 as ...., ....@0 as c3] +02)--SortExec: expr=[....@0 ASC NULLS LAST], preserve_partitioning=[false] 03)----DataSourceExec: partitions=1, partition_sizes=[1] diff --git a/datafusion/sqllogictest/test_files/regexp/regexp_count.slt b/datafusion/sqllogictest/test_files/regexp/regexp_count.slt index c87c194fa6b25..0b2b9e5e74559 100644 --- a/datafusion/sqllogictest/test_files/regexp/regexp_count.slt +++ b/datafusion/sqllogictest/test_files/regexp/regexp_count.slt @@ -76,6 +76,21 @@ SELECT regexp_count('abc', '', 5); ---- 0 +query I +SELECT regexp_count('', ''); +---- +1 + +query I +SELECT regexp_count('😀', '', 2); +---- +1 + +query I +SELECT regexp_count('abc', 'x*', 4); +---- +1 + statement error External error: query failed: DataFusion error: Arrow error: Compute error: regexp_count() requires start to be 1 based SELECT regexp_count('123123123123', '123', 0); @@ -91,7 +106,7 @@ SELECT regexp_count('123123123123', '123', 1, 'g'); query I SELECT regexp_count(str, '\w') from regexp_test_data; ---- -0 +NULL 3 3 3 @@ -107,7 +122,7 @@ SELECT regexp_count(str, '\w') from regexp_test_data; query I SELECT regexp_count(str, '\w{2}', start) from regexp_test_data; ---- -0 +NULL 1 1 1 @@ -123,7 +138,7 @@ SELECT regexp_count(str, '\w{2}', start) from regexp_test_data; query I SELECT regexp_count(str, 'ab', 1, 'i') from regexp_test_data; ---- -0 +NULL 1 1 1 @@ -140,7 +155,7 @@ SELECT regexp_count(str, 'ab', 1, 'i') from regexp_test_data; query I SELECT regexp_count(str, pattern) from regexp_test_data; ---- -0 +NULL 1 1 0 @@ -156,7 +171,7 @@ SELECT regexp_count(str, pattern) from regexp_test_data; query I SELECT regexp_count(str, pattern, start) from regexp_test_data; ---- -0 +NULL 1 1 0 @@ -172,35 +187,35 @@ SELECT regexp_count(str, pattern, start) from regexp_test_data; query I SELECT regexp_count(str, pattern, start, flags) from regexp_test_data; ---- -0 -1 -1 -1 -0 -0 -0 -0 -1 +NULL 1 1 1 +NULL +NULL +NULL +NULL +NULL +NULL +NULL +NULL # test type coercion query I SELECT regexp_count(arrow_cast(str, 'Utf8'), arrow_cast(pattern, 'LargeUtf8'), arrow_cast(start, 'Int32'), flags) from regexp_test_data; ---- -0 -1 -1 -1 -0 -0 -0 -0 -1 +NULL 1 1 1 +NULL +NULL +NULL +NULL +NULL +NULL +NULL +NULL # test string views @@ -211,7 +226,7 @@ SELECT arrow_cast(str, 'Utf8View') as str, arrow_cast(pattern, 'Utf8View') as pa query I SELECT regexp_count(str, '\w') from t_stringview; ---- -0 +NULL 3 3 3 @@ -227,7 +242,7 @@ SELECT regexp_count(str, '\w') from t_stringview; query I SELECT regexp_count(str, '\w{2}', start) from t_stringview; ---- -0 +NULL 1 1 1 @@ -243,7 +258,7 @@ SELECT regexp_count(str, '\w{2}', start) from t_stringview; query I SELECT regexp_count(str, 'ab', 1, 'i') from t_stringview; ---- -0 +NULL 1 1 1 @@ -260,7 +275,7 @@ SELECT regexp_count(str, 'ab', 1, 'i') from t_stringview; query I SELECT regexp_count(str, pattern) from t_stringview; ---- -0 +NULL 1 1 0 @@ -276,7 +291,7 @@ SELECT regexp_count(str, pattern) from t_stringview; query I SELECT regexp_count(str, pattern, start) from t_stringview; ---- -0 +NULL 1 1 0 @@ -292,57 +307,74 @@ SELECT regexp_count(str, pattern, start) from t_stringview; query I SELECT regexp_count(str, pattern, start, flags) from t_stringview; ---- -0 -1 -1 -1 -0 -0 -0 -0 -1 +NULL 1 1 1 +NULL +NULL +NULL +NULL +NULL +NULL +NULL +NULL # test type coercion query I SELECT regexp_count(arrow_cast(str, 'Utf8'), arrow_cast(pattern, 'LargeUtf8'), arrow_cast(start, 'Int32'), flags) from t_stringview; ---- -0 -1 -1 -1 -0 -0 -0 -0 -1 +NULL 1 1 1 +NULL +NULL +NULL +NULL +NULL +NULL +NULL +NULL -# NULL tests +# NULL tests: like PostgreSQL, a NULL in any argument produces a NULL result query I SELECT regexp_count(NULL, NULL); ---- -0 +NULL query I SELECT regexp_count(NULL, 'a'); ---- -0 +NULL query I SELECT regexp_count('a', NULL); ---- -0 +NULL query I SELECT regexp_count(NULL, NULL, NULL, NULL); ---- -0 +NULL + +query I +SELECT regexp_count('abc', 'b', NULL); +---- +NULL + +query I +SELECT regexp_count('abc', 'b', 1, NULL); +---- +NULL + +# NULL start position in one row of a column +query I +SELECT regexp_count(v, 'b', s) FROM (VALUES ('abc', 1), ('abcb', NULL)) AS t(v, s); +---- +1 +NULL statement ok CREATE TABLE empty_table (str varchar, pattern varchar, start int, flags varchar); @@ -357,10 +389,10 @@ INSERT INTO empty_table VALUES ('a', NULL, 1, 'i'), (NULL, 'a', 1, 'i'), (NULL, query I SELECT regexp_count(str, pattern, start, flags) from empty_table; ---- -0 -0 -0 -0 +NULL +NULL +NULL +NULL statement ok drop table t_stringview; diff --git a/datafusion/sqllogictest/test_files/regexp/regexp_instr.slt b/datafusion/sqllogictest/test_files/regexp/regexp_instr.slt index d4e98e6431678..bbe9693736442 100644 --- a/datafusion/sqllogictest/test_files/regexp/regexp_instr.slt +++ b/datafusion/sqllogictest/test_files/regexp/regexp_instr.slt @@ -23,6 +23,18 @@ SELECT regexp_instr('123123123123123', '(12)3'); ---- 1 +query IIIIIII +SELECT + regexp_instr('abc', ''), + regexp_instr('', ''), + regexp_instr('abc', '', 4), + regexp_instr('abc', '', 5), + regexp_instr('😀', '', 1, 2), + regexp_instr('abc', 'x*', 4), + regexp_instr(NULL, ''); +---- +1 1 4 0 2 4 NULL + query I SELECT regexp_instr('123123123123', '123', 1); ---- @@ -61,14 +73,15 @@ SELECT ---- 11 -statement error -External error: query failed: DataFusion error: Arrow error: Compute error: regexp_instr() requires start to be 1 based +statement error DataFusion error: Arrow error: Compute error: regexp_instr\(\) requires start to be 1-based SELECT regexp_instr('123123123123', '123', 0); -statement error -External error: query failed: DataFusion error: Arrow error: Compute error: regexp_instr() requires start to be 1 based +statement error DataFusion error: Arrow error: Compute error: regexp_instr\(\) requires start to be 1-based SELECT regexp_instr('123123123123', '123', -3); +statement error DataFusion error: Arrow error: Compute error: N must be 1 or greater +SELECT regexp_instr('abcabcabc', 'abc', 1, 0); + query I SELECT regexp_instr(str, pattern) FROM regexp_test_data; ---- @@ -161,6 +174,34 @@ SELECT regexp_instr('a', NULL); ---- NULL +# Like PostgreSQL, a NULL in any argument produces a NULL result +query I +SELECT regexp_instr('abc', 'b', NULL); +---- +NULL + +query I +SELECT regexp_instr('abc', 'b', 1, NULL); +---- +NULL + +query I +SELECT regexp_instr('abc', 'b', 1, 1, NULL); +---- +NULL + +query I +SELECT regexp_instr('abc', '(b)', 1, 1, 'i', NULL); +---- +NULL + +# NULL start position in one row of a column +query I +SELECT regexp_instr(v, 'b', s) FROM (VALUES ('abc', 1), ('abcb', NULL)) AS t(v, s); +---- +2 +NULL + query I SELECT regexp_instr('😀abcdef', 'abc'); ---- @@ -189,8 +230,27 @@ NULL NULL NULL +# The pattern column alternates between two regexes within a single batch, so +# the compiled regex for 'abc' must be looked up again from the regex cache +# after 'def' displaced it as the most recently used pattern +statement ok +CREATE TABLE t_alternating_pattern(str varchar, pattern varchar) AS VALUES + ('abcdef', 'abc'), + ('abcdef', 'def'), + ('abcdef', 'abc'); + +query I +SELECT regexp_instr(str, pattern) FROM t_alternating_pattern; +---- +1 +4 +1 + statement ok DROP TABLE t_stringview; statement ok DROP TABLE empty_table; + +statement ok +DROP TABLE t_alternating_pattern; diff --git a/datafusion/sqllogictest/test_files/regexp/regexp_like.slt b/datafusion/sqllogictest/test_files/regexp/regexp_like.slt index 22d5066d5f782..30fa913896bdf 100644 --- a/datafusion/sqllogictest/test_files/regexp/regexp_like.slt +++ b/datafusion/sqllogictest/test_files/regexp/regexp_like.slt @@ -168,6 +168,16 @@ SELECT 'foo\nbar\nbaz' ~ 'bar'; ---- true +query B +SELECT regexp_like(E'a\nb', '^b', 'm'); +---- +true + +query B +SELECT regexp_like(E'a\nb', '^b'); +---- +false + statement error Error during planning: Cannot infer common argument type for regex operation List(Field { name: "item", data_type: Int64, nullable: true, metadata: {} }) ~ List(Field { name: "item", data_type: Int64, nullable: true, metadata: {} }) select [1,2] ~ [3]; diff --git a/datafusion/sqllogictest/test_files/repartition_scan.slt b/datafusion/sqllogictest/test_files/repartition_scan.slt index 88eaf7118f8a5..aa5ef064ec67a 100644 --- a/datafusion/sqllogictest/test_files/repartition_scan.slt +++ b/datafusion/sqllogictest/test_files/repartition_scan.slt @@ -64,7 +64,7 @@ logical_plan 02)--TableScan: parquet_table projection=[column1], partial_filters=[parquet_table.column1 != Int32(42)] physical_plan 01)FilterExec: column1@0 != 42 -02)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:0..135], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:135..270], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:270..405], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:405..537]]}, projection=[column1], file_type=parquet, predicate=column1@0 != 42, pruning_predicate=column1_null_count@2 != row_count@3 AND (column1_min@0 != 42 OR 42 != column1_max@1), required_guarantees=[column1 not in (42)] +02)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:0..131], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:131..262], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:262..393], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:393..521]]}, projection=[column1], file_type=parquet, predicate=column1@0 != 42, pruning_predicate=column1_null_count@2 != row_count@3 AND (column1_min@0 != 42 OR 42 != column1_max@1), required_guarantees=[column1 not in (42)] # disable round robin repartitioning statement ok @@ -79,7 +79,7 @@ logical_plan 02)--TableScan: parquet_table projection=[column1], partial_filters=[parquet_table.column1 != Int32(42)] physical_plan 01)FilterExec: column1@0 != 42 -02)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:0..135], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:135..270], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:270..405], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:405..537]]}, projection=[column1], file_type=parquet, predicate=column1@0 != 42, pruning_predicate=column1_null_count@2 != row_count@3 AND (column1_min@0 != 42 OR 42 != column1_max@1), required_guarantees=[column1 not in (42)] +02)--DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:0..131], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:131..262], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:262..393], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:393..521]]}, projection=[column1], file_type=parquet, predicate=column1@0 != 42, pruning_predicate=column1_null_count@2 != row_count@3 AND (column1_min@0 != 42 OR 42 != column1_max@1), required_guarantees=[column1 not in (42)] # enable round robin repartitioning again statement ok @@ -103,7 +103,7 @@ physical_plan 01)SortPreservingMergeExec: [column1@0 ASC NULLS LAST] 02)--SortExec: expr=[column1@0 ASC NULLS LAST], preserve_partitioning=[true] 03)----FilterExec: column1@0 != 42 -04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/1.parquet:0..266], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/1.parquet:266..526, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:0..6], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:6..272], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:272..537]]}, projection=[column1], file_type=parquet, predicate=column1@0 != 42, pruning_predicate=column1_null_count@2 != row_count@3 AND (column1_min@0 != 42 OR 42 != column1_max@1), required_guarantees=[column1 not in (42)] +04)------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/1.parquet:0..258], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/1.parquet:258..510, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:0..6], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:6..264], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:264..521]]}, projection=[column1], file_type=parquet, predicate=column1@0 != 42, pruning_predicate=column1_null_count@2 != row_count@3 AND (column1_min@0 != 42 OR 42 != column1_max@1), required_guarantees=[column1 not in (42)] ## Read the files as though they are ordered @@ -138,7 +138,7 @@ logical_plan physical_plan 01)SortPreservingMergeExec: [column1@0 ASC NULLS LAST] 02)--FilterExec: column1@0 != 42 -03)----DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/1.parquet:0..263], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:0..268], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:268..537], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/1.parquet:263..526]]}, projection=[column1], output_ordering=[column1@0 ASC NULLS LAST], file_type=parquet, predicate=column1@0 != 42, pruning_predicate=column1_null_count@2 != row_count@3 AND (column1_min@0 != 42 OR 42 != column1_max@1), required_guarantees=[column1 not in (42)] +03)----DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/1.parquet:0..255], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:0..260], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/2.parquet:260..521], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_scan/parquet_table/1.parquet:255..510]]}, projection=[column1], output_ordering=[column1@0 ASC NULLS LAST], file_type=parquet, predicate=column1@0 != 42, pruning_predicate=column1_null_count@2 != row_count@3 AND (column1_min@0 != 42 OR 42 != column1_max@1), required_guarantees=[column1 not in (42)] # Cleanup statement ok diff --git a/datafusion/sqllogictest/test_files/repartition_subset_satisfaction.slt b/datafusion/sqllogictest/test_files/repartition_subset_satisfaction.slt index dbf31dec5e118..5371ca59beea1 100644 --- a/datafusion/sqllogictest/test_files/repartition_subset_satisfaction.slt +++ b/datafusion/sqllogictest/test_files/repartition_subset_satisfaction.slt @@ -164,7 +164,7 @@ physical_plan 03)----AggregateExec: mode=FinalPartitioned, gby=[f_dkey@0 as f_dkey, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)@1 as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)], aggr=[count(Int64(1)), avg(fact_table_ordered.value)], ordering_mode=Sorted 04)------RepartitionExec: partitioning=Hash([f_dkey@0, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)@1], 3), input_partitions=3, preserve_order=true, sort_exprs=f_dkey@0 ASC NULLS LAST, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)@1 ASC NULLS LAST 05)--------AggregateExec: mode=Partial, gby=[f_dkey@2 as f_dkey, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }, timestamp@0) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)], aggr=[count(Int64(1)), avg(fact_table_ordered.value)], ordering_mode=Sorted -06)----------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], file_type=parquet +06)----------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 3), file_type=parquet # Verify results without subset satisfaction query TPIR rowsort @@ -204,7 +204,7 @@ physical_plan 01)SortPreservingMergeExec: [f_dkey@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST] 02)--ProjectionExec: expr=[f_dkey@0 as f_dkey, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)@1 as time_bin, count(Int64(1))@2 as count(*), avg(fact_table_ordered.value)@3 as avg(fact_table_ordered.value)] 03)----AggregateExec: mode=SinglePartitioned, gby=[f_dkey@2 as f_dkey, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }, timestamp@0) as date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)], aggr=[count(Int64(1)), avg(fact_table_ordered.value)], ordering_mode=Sorted -04)------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], file_type=parquet +04)------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 3), file_type=parquet # Verify results match with subset satisfaction query TPIR rowsort @@ -251,7 +251,7 @@ physical_plan 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [fact_table_ordered.f_dkey, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)] ORDER BY [fact_table_ordered.timestamp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [fact_table_ordered.f_dkey, date_bin(IntervalMonthDayNano(\"IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }\"),fact_table_ordered.timestamp)] ORDER BY [fact_table_ordered.timestamp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 03)----SortExec: expr=[f_dkey@2 ASC NULLS LAST, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }, timestamp@0) ASC NULLS LAST, timestamp@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------RepartitionExec: partitioning=Hash([f_dkey@2, date_bin(IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }, timestamp@0)], 3), input_partitions=3 -05)--------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], file_type=parquet +05)--------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 3), file_type=parquet # Verify results without subset satisfaction query TPRI rowsort @@ -292,7 +292,7 @@ logical_plan physical_plan 01)ProjectionExec: expr=[f_dkey@2 as f_dkey, timestamp@0 as timestamp, value@1 as value, row_number() PARTITION BY [fact_table_ordered.f_dkey, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)] ORDER BY [fact_table_ordered.timestamp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [fact_table_ordered.f_dkey, date_bin(IntervalMonthDayNano("IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }"),fact_table_ordered.timestamp)] ORDER BY [fact_table_ordered.timestamp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [fact_table_ordered.f_dkey, date_bin(IntervalMonthDayNano(\"IntervalMonthDayNano { months: 0, days: 0, nanoseconds: 30000000000 }\"),fact_table_ordered.timestamp)] ORDER BY [fact_table_ordered.timestamp ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], file_type=parquet +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 3), file_type=parquet # Verify results match with subset satisfaction query TPRI rowsort @@ -367,8 +367,8 @@ logical_plan 15)--------------------TableScan: fact_table_ordered projection=[timestamp, value, f_dkey] physical_plan 01)SortPreservingMergeExec: [env@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST] -02)--SortExec: expr=[env@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[env@0 as env, time_bin@1 as time_bin, avg(a.max_bin_value)@2 as avg_max_value] +02)--ProjectionExec: expr=[env@0 as env, time_bin@1 as time_bin, avg(a.max_bin_value)@2 as avg_max_value] +03)----SortExec: expr=[env@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, time_bin@1 as time_bin], aggr=[avg(a.max_bin_value)] 05)--------RepartitionExec: partitioning=Hash([env@0, time_bin@1], 3), input_partitions=3 06)----------AggregateExec: mode=Partial, gby=[env@1 as env, time_bin@0 as time_bin], aggr=[avg(a.max_bin_value)] @@ -379,8 +379,8 @@ physical_plan 11)--------------------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(d_dkey@1, f_dkey@2)], projection=[f_dkey@4, env@0, timestamp@2, value@3] 12)----------------------CoalescePartitionsExec 13)------------------------FilterExec: service@1 = log, projection=[env@0, d_dkey@2] -14)--------------------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=A/data.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=D/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=C/data.parquet]]}, projection=[env, service, d_dkey], file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] -15)----------------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +14)--------------------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=A/data.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=D/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=C/data.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 3), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] +15)----------------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 3), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify results without subset satisfaction query TPR rowsort @@ -464,8 +464,8 @@ logical_plan 15)--------------------TableScan: fact_table_ordered projection=[timestamp, value, f_dkey] physical_plan 01)SortPreservingMergeExec: [env@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST] -02)--SortExec: expr=[env@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[env@0 as env, time_bin@1 as time_bin, avg(a.max_bin_value)@2 as avg_max_value] +02)--ProjectionExec: expr=[env@0 as env, time_bin@1 as time_bin, avg(a.max_bin_value)@2 as avg_max_value] +03)----SortExec: expr=[env@0 ASC NULLS LAST, time_bin@1 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[env@0 as env, time_bin@1 as time_bin], aggr=[avg(a.max_bin_value)] 05)--------RepartitionExec: partitioning=Hash([env@0, time_bin@1], 3), input_partitions=3 06)----------AggregateExec: mode=Partial, gby=[env@1 as env, time_bin@0 as time_bin], aggr=[avg(a.max_bin_value)] @@ -474,8 +474,8 @@ physical_plan 09)----------------HashJoinExec: mode=CollectLeft, join_type=Inner, on=[(d_dkey@1, f_dkey@2)], projection=[f_dkey@4, env@0, timestamp@2, value@3] 10)------------------CoalescePartitionsExec 11)--------------------FilterExec: service@1 = log, projection=[env@0, d_dkey@2] -12)----------------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=A/data.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=D/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=C/data.parquet]]}, projection=[env, service, d_dkey], file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] -13)------------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +12)----------------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=A/data.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=D/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/dimension/d_dkey=C/data.parquet]]}, projection=[env, service, d_dkey], output_partitioning=Hash([d_dkey@2], 3), file_type=parquet, predicate=service@1 = log, pruning_predicate=service_null_count@2 != row_count@3 AND service_min@0 <= log AND log <= service_max@1, required_guarantees=[service in (log)] +13)------------------DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=A/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=B/data.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/repartition_subset_satisfaction/fact/f_dkey=C/data.parquet]]}, projection=[timestamp, value, f_dkey], output_ordering=[f_dkey@2 ASC NULLS LAST, timestamp@0 ASC NULLS LAST], output_partitioning=Hash([f_dkey@2], 3), file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Verify results match with subset satisfaction query TPR rowsort @@ -517,7 +517,7 @@ prod 2023-01-01T09:12:30 197.7 # Config reset -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/scalar.slt b/datafusion/sqllogictest/test_files/scalar.slt index 89ae30e3c047b..7666b680e16a8 100644 --- a/datafusion/sqllogictest/test_files/scalar.slt +++ b/datafusion/sqllogictest/test_files/scalar.slt @@ -234,7 +234,26 @@ select round(atanh(a), 5), round(atanh(b), 5), round(atanh(c), 5) from small_flo query RRR rowsort select atan2(0, 1), atan2(1, 2), atan2(2, 2); ---- -0 0.4636476 0.7853982 +0 0.463647609001 0.785398163397 + +# atan2 returns Float32 only when both arguments are Float32; every other +# numeric combination (integers, Float64, mixed, NULL) is computed in Float64 +query TTTTTT +select + arrow_typeof(atan2(arrow_cast(1.0, 'Float32'), arrow_cast(1.0, 'Float32'))), + arrow_typeof(atan2(1, 1)), + arrow_typeof(atan2(arrow_cast(1.0, 'Float32'), arrow_cast(1.0, 'Float64'))), + arrow_typeof(atan2(arrow_cast(1.0, 'Float64'), arrow_cast(1.0, 'Float32'))), + arrow_typeof(atan2(null, null)), + arrow_typeof(atan2(null, 64)); +---- +Float32 Float64 Float64 Float64 Float64 Float64 + +# atan2 with integer inputs is computed in double precision +query B +select atan2(1, 1000000) = atan2(1.0, 1000000.0); +---- +true # atan2 scalar nulls query R rowsort @@ -777,6 +796,31 @@ select nanvl(null, null); ---- NULL +# nanvl evaluates in the common (widest) float type of its arguments. Mixing +# narrower floats widens losslessly (Float16 + Float32 -> Float32), while +# integers, decimals, and NULL are coerced to Float64. +query TTTTTTTT +select + arrow_typeof(nanvl(arrow_cast(1.0, 'Float16'), arrow_cast(2.0, 'Float16'))), + arrow_typeof(nanvl(arrow_cast(1.0, 'Float32'), arrow_cast(2.0, 'Float32'))), + arrow_typeof(nanvl(arrow_cast(1.0, 'Float16'), arrow_cast(2.0, 'Float32'))), + arrow_typeof(nanvl(arrow_cast(1.0, 'Float16'), arrow_cast(2.0, 'Float64'))), + arrow_typeof(nanvl(arrow_cast(1.0, 'Float32'), arrow_cast(2.0, 'Float64'))), + arrow_typeof(nanvl(1, 2)), + arrow_typeof(nanvl(1, arrow_cast(2.0, 'Float32'))), + arrow_typeof(nanvl(null, null)); +---- +Float16 Float32 Float32 Float64 Float64 Float64 Float64 Float64 + +# nanvl with an integer argument is computed in double precision, even when the +# other argument is Float32. +query BB +select + nanvl(16777217, 1) = nanvl(arrow_cast(16777217, 'Float64'), 1.0), + nanvl(16777217, arrow_cast(1.0, 'Float32')) = nanvl(arrow_cast(16777217, 'Float64'), 1.0); +---- +true true + # nanvl with columns (round is needed to normalize the outputs of different operating systems) query RRR rowsort select round(nanvl(asin(f + a), 2), 5), round(nanvl(asin(b + c), 3), 5), round(nanvl(asin(d + e), 4), 5) from small_floats; @@ -921,6 +965,61 @@ select round(a), round(b), round(c) from small_floats; 0 0 1 1 0 0 +# round int64 should preserve exact values above Float64 precision range +query TI +select arrow_typeof(round(arrow_cast(9007199254740993, 'Int64'))), + round(arrow_cast(9007199254740993, 'Int64')); +---- +Int64 9007199254740993 + +# round int64 with positive decimal_places should preserve exact values above Float64 precision range +query TI +select arrow_typeof(round(arrow_cast(9007199254740993, 'Int64'), 2)), + round(arrow_cast(9007199254740993, 'Int64'), 2); +---- +Int64 9007199254740993 + +# round int64 with negative decimal_places +query TI +select arrow_typeof(round(arrow_cast(125, 'Int64'), -1)), + round(arrow_cast(125, 'Int64'), -1); +---- +Int64 130 + +# round int64 with column decimal_places +query I +select round(v, dp) +from (values (arrow_cast(125, 'Int64'), 1), + (arrow_cast(125, 'Int64'), -1)) as t(v, dp); +---- +125 +130 + +# round int64 overflow with negative decimal_places +query error Overflow while rounding Int64 +select round(arrow_cast(9223372036854775807, 'Int64'), -1); + +# round uint64 should preserve exact values +query TI +select arrow_typeof(round(arrow_cast(18446744073709551615, 'UInt64'))), + round(arrow_cast(18446744073709551615, 'UInt64')); +---- +UInt64 18446744073709551615 + +# round uint64 with positive decimal_places should preserve exact values +query TI +select arrow_typeof(round(arrow_cast(18446744073709551615, 'UInt64'), 2)), + round(arrow_cast(18446744073709551615, 'UInt64'), 2); +---- +UInt64 18446744073709551615 + +# round int64 to place larger than the number itself +query TI +select arrow_typeof(round(arrow_cast(125, 'Int64'), -5)), + round(arrow_cast(125, 'Int64'), -5); +---- +Int64 0 + # round with too large # max Int32 is 2147483647 query error round decimal_places 2147483648 is out of supported i32 range @@ -1107,12 +1206,16 @@ NULL # sqrt with columns (round is needed to normalize the outputs of different operating systems) query RRR rowsort -select round(sqrt(a), 5), round(sqrt(b), 5), round(sqrt(c), 5) from signed_integers; +select round(sqrt(abs(a)), 5), round(sqrt(abs(b)), 5), round(sqrt(abs(c)), 5) from signed_integers; ---- -1.41421 NaN 11.09054 +1 10 23.81176 +1.41421 31.62278 11.09054 +1.73205 100 31.27299 2 NULL NULL -NaN 10 NaN -NaN 100 NaN + +# sqrt with negative column values should error +query error cannot take square root of a negative number +select round(sqrt(a), 5), round(sqrt(b), 5), round(sqrt(c), 5) from signed_integers; # sqrt scalar fraction query RR rowsort @@ -1128,10 +1231,12 @@ select sqrt(cast(10e8 as double)); # sqrt scalar negative -query R rowsort +query error cannot take square root of a negative number select sqrt(-1); ----- -NaN + +# sqrt scalar negative float8 +query error cannot take square root of a negative number +select sqrt((-1.0)::float8); ## tan @@ -1228,6 +1333,57 @@ from small_floats; 0.836 0.8 0.836 1 1 1 +# trunc with decimals +query RT +select trunc(arrow_cast(3.1415, 'Decimal128(10,4)')), arrow_typeof(trunc(arrow_cast(3.1415, 'Decimal128(10,4)'))); +---- +3 Decimal128(10, 4) + +# trunc with precision - decimals +query RRRRR rowsort +select + trunc(arrow_cast(4.267, 'Decimal32(8,3)'), 3), + trunc(arrow_cast(1.1234, 'Decimal64(18,6)'), 2), + trunc(arrow_cast(-1.1231, 'Decimal128(15,4)'), 6), + trunc(arrow_cast(1.2837284, 'Decimal256(35,7)'), 2), + trunc(arrow_cast(1.1, 'Decimal128(10,1)'), 0); +---- +4.267 1.12 -1.1231 1.28 1 + +# trunc with negative precision should truncate digits left of decimal - decimal types +query RT +select trunc(arrow_cast(12345.678, 'Decimal128(10,3)'), -3), + arrow_typeof(trunc(arrow_cast(12345.678, 'Decimal128(10,3)'), -3)); +---- +12000 Decimal128(10, 3) + +# trunc: coercion with a decimal argument and a non-int64 precision argument +query RT +select trunc(arrow_cast(1.2345678, 'Decimal128(20,14)'), arrow_cast(2, 'Int32')), + arrow_typeof(trunc(arrow_cast(1.2345678, 'Decimal128(20,14)'), arrow_cast(2, 'Int32'))); +---- +1.23 Decimal128(20, 14) + +# trunc with columns and precision - decimal128 +query RRR rowsort +select + trunc(arrow_cast(a, 'Decimal128(10,4)'), 0) as a0, + trunc(arrow_cast(b, 'Decimal128(10,4)'), 0) as b0, + trunc(arrow_cast(c, 'Decimal128(10,4)'), 0) as c0 +from small_floats; +---- +-1 NULL NULL +0 0 -1 +0 0 0 +0 0 1 + +# trunc issue #22512 +query R +select trunc(CAST(9007199254740993 AS DECIMAL(20,0))); +---- +9007199254740993 + + ## bitwise and # bitwise and with column and scalar @@ -1300,6 +1456,12 @@ NULL -32 statement ok set datafusion.sql_parser.dialect = postgresql; +# postgresql exponentiation uses caret +query R +select 2 ^ 3; +---- +8 + # postgresql bitwise xor with column and scalar query I rowsort select c # 856 from signed_integers; @@ -1753,7 +1915,7 @@ SELECT not(true), not(false) ---- false true -query error type_coercion\ncaused by\nError during planning: Cannot infer common argument type for comparison operation Int64 IS DISTINCT FROM Boolean +query error Error during planning: Unary operator 'NOT' requires a boolean expression, got Int64 SELECT not(1), not(0) query ?B @@ -1761,7 +1923,7 @@ SELECT null, not(null) ---- NULL NULL -query error type_coercion\ncaused by\nError during planning: Cannot infer common argument type for comparison operation Utf8 IS DISTINCT FROM Boolean +query error Error during planning: Unary operator 'NOT' requires a boolean expression, got Utf8 SELECT NOT('hi') # test_negative_expressions() @@ -1771,7 +1933,7 @@ SELECT null, -null ---- NULL NULL -query error type_coercion\ncaused by\nError during planning: Negation only supports numeric, interval and timestamp types +query error Error during planning: Unary operator '-' only supports signed numeric, interval and timestamp types SELECT -'100' query error DataFusion error: Error during planning: Unary operator '\+' only supports numeric, interval and timestamp types diff --git a/datafusion/sqllogictest/test_files/select.slt b/datafusion/sqllogictest/test_files/select.slt index 3e97dc4588655..4107921d2fda5 100644 --- a/datafusion/sqllogictest/test_files/select.slt +++ b/datafusion/sqllogictest/test_files/select.slt @@ -960,6 +960,12 @@ physical_plan 01)ProjectionExec: expr=[c1@0 >= 2 AND c1@0 <= 3 as select_between_data.c1 BETWEEN Int64(2) AND Int64(3)] 02)--DataSourceExec: partitions=1, partition_sizes=[1] +# regression test: full i64 BETWEEN bounds should not overflow +query I +SELECT * FROM (VALUES (1)) AS t(x) +WHERE x BETWEEN -9223372036854775808 AND 9223372036854775807 +---- +1 # TODO: query_get_indexed_field @@ -1181,7 +1187,7 @@ SELECT * FROM empty_table statement ok CREATE TABLE case_sensitive_table("INT32" int) AS VALUES (1), (2), (3), (4), (5); -statement error DataFusion error: Schema error: No field named int32\. Valid fields are case_sensitive_table\."INT32"\. +statement error DataFusion error: Schema error: No field named int32\. Did you mean 'case_sensitive_table\."INT32"'\?\nColumn names are case sensitive\. You can use double quotes to refer to the case_sensitive_table\."INT32" column or disable the datafusion\.sql_parser\.enable_ident_normalization configuration\.\nValid fields are case_sensitive_table\."INT32"\. select "int32" from case_sensitive_table query I @@ -1571,9 +1577,9 @@ physical_plan 03)----RepartitionExec: partitioning=Hash([c2@0], 2), input_partitions=2 04)------AggregateExec: mode=Partial, gby=[c2@0 as c2], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1 -06)----------ProjectionExec: expr=[c2@0 as c2] -07)------------SortExec: TopK(fetch=4), expr=[c1@1 ASC NULLS LAST, c2@0 ASC NULLS LAST], preserve_partitioning=[false] -08)--------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c2, c1], file_type=csv, has_header=true +06)----------ProjectionExec: expr=[c2@1 as c2] +07)------------SortExec: TopK(fetch=4), expr=[c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST], preserve_partitioning=[false] +08)--------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c2], file_type=csv, has_header=true # FilterExec can track equality of non-column expressions. # plan below shouldn't have a SortExec because given column 'a' is ordered. @@ -1585,16 +1591,15 @@ WHERE CAST(ROUND(b) as INT) = a ORDER BY CAST(ROUND(b) as INT); ---- logical_plan -01)Sort: CAST(round(CAST(annotated_data_finite2.b AS Float64)) AS Int32) ASC NULLS LAST -02)--Filter: CAST(round(CAST(annotated_data_finite2.b AS Float64)) AS Int32) = annotated_data_finite2.a -03)----TableScan: annotated_data_finite2 projection=[a0, a, b, c, d], partial_filters=[CAST(round(CAST(annotated_data_finite2.b AS Float64)) AS Int32) = annotated_data_finite2.a] +01)Sort: CAST(round(annotated_data_finite2.b) AS Int32) ASC NULLS LAST +02)--Filter: CAST(round(annotated_data_finite2.b) AS Int32) = annotated_data_finite2.a +03)----TableScan: annotated_data_finite2 projection=[a0, a, b, c, d], partial_filters=[CAST(round(annotated_data_finite2.b) AS Int32) = annotated_data_finite2.a] physical_plan -01)SortPreservingMergeExec: [CAST(round(CAST(b@2 AS Float64)) AS Int32) ASC NULLS LAST] -02)--FilterExec: CAST(round(CAST(b@2 AS Float64)) AS Int32) = a@1 +01)SortPreservingMergeExec: [round(b@2) ASC NULLS LAST] +02)--FilterExec: round(b@2) = a@1 03)----RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1, maintains_sort_order=true 04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC NULLS LAST, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], file_type=csv, has_header=true - statement ok drop table annotated_data_finite2; @@ -1823,7 +1828,7 @@ select a + b from (select 1 as a, 2 as b, 1 as "a + b"); 3 # Can't reference an output column by expression over projection. -query error DataFusion error: Schema error: No field named a\. Valid fields are "a \+ Int64\(1\)"\. +query error DataFusion error: Schema error: No field named a\.\nValid fields are "a \+ Int64\(1\)"\. select a + 1 from (select a+1 from (select 1 as a)); query I @@ -1861,7 +1866,7 @@ statement ok DROP TABLE test; # Can't reference an unqualified column by a qualified name -query error DataFusion error: Schema error: No field named t1\.v1\. Column names are case sensitive\. You can use double quotes to refer to the "t1\.v1" column or set the datafusion\.sql_parser\.enable_ident_normalization configuration\. Valid fields are "t1\.v1"\. +query error DataFusion error: Schema error: No field named t1\.v1\. Did you mean '"t1\.v1"'\?\nValid fields are "t1\.v1"\. SELECT t1.v1 FROM (SELECT 1 AS "t1.v1"); # Test issue: https://github.com/apache/datafusion/issues/14124 @@ -1963,7 +1968,7 @@ SELECT COUNT(*) FROM t0 AS tt0 WHERE (4==(3/0)); # Config reset -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; diff --git a/datafusion/sqllogictest/test_files/set_variable.slt b/datafusion/sqllogictest/test_files/set_variable.slt index 6f58e5fb3100b..b8db761e796fe 100644 --- a/datafusion/sqllogictest/test_files/set_variable.slt +++ b/datafusion/sqllogictest/test_files/set_variable.slt @@ -93,10 +93,10 @@ datafusion.execution.coalesce_batches false statement ok set datafusion.catalog.information_schema = true -statement error DataFusion error: Error parsing '1' as bool +statement error DataFusion error: Error setting config datafusion\.execution\.coalesce_batches\ncaused by\nError parsing '1' as bool SET datafusion.execution.coalesce_batches to 1 -statement error DataFusion error: Error parsing 'abc' as bool +statement error DataFusion error: Error setting config datafusion\.execution\.coalesce_batches\ncaused by\nError parsing 'abc' as bool SET datafusion.execution.coalesce_batches to abc # set u64 variable @@ -104,12 +104,12 @@ statement ok set datafusion.catalog.information_schema = true statement ok -SET datafusion.execution.batch_size to 0 +SET datafusion.execution.batch_size to 310104 query TT SHOW datafusion.execution.batch_size ---- -datafusion.execution.batch_size 0 +datafusion.execution.batch_size 310104 statement ok SET datafusion.execution.batch_size to '1' @@ -132,10 +132,10 @@ datafusion.execution.batch_size 2 statement ok set datafusion.catalog.information_schema = true -statement error DataFusion error: Error parsing '-1' as usize +statement error DataFusion error: Error setting config datafusion\.execution\.batch_size\ncaused by\nError parsing '-1' as usize SET datafusion.execution.batch_size to -1 -statement error DataFusion error: Error parsing 'abc' as usize +statement error DataFusion error: Error setting config datafusion\.execution\.batch_size\ncaused by\nError parsing 'abc' as usize SET datafusion.execution.batch_size to abc statement error External error: invalid digit found in string @@ -382,7 +382,7 @@ statement error DataFusion error: Invalid or Unsupported Configuration: Config v RESET datafusion.execution.batches_size # reset invalid variable - extra suffix on valid field -statement error DataFusion error: Invalid or Unsupported Configuration: Config field is a scalar usize and does not have nested field "bar" +statement error DataFusion error: Invalid or Unsupported Configuration: Config field batch_size is a scalar ConfigNonZeroUsize and does not have nested field "bar" RESET datafusion.execution.batch_size.bar ############################################# @@ -580,7 +580,7 @@ SHOW datafusion.format.date_format datafusion.format.date_format %Y-%m-%d # Invalid format option name -statement error DataFusion error: Invalid or Unsupported Configuration: Config value "unknown_option" not found on FormatOptions +statement error DataFusion error: Error setting config datafusion\.format\.unknown_option\ncaused by\nInvalid or Unsupported Configuration: Config value "unknown_option" not found on FormatOptions SET datafusion.format.unknown_option = true ############ @@ -611,6 +611,18 @@ SHOW datafusion.runtime.max_temp_directory_size ---- datafusion.runtime.max_temp_directory_size 10G +# Test SET and SHOW runtime.max_spill_merge_fan_in +statement ok +SET datafusion.runtime.max_spill_merge_fan_in = '16' + +query TT +SHOW datafusion.runtime.max_spill_merge_fan_in +---- +datafusion.runtime.max_spill_merge_fan_in 16 + +statement ok +RESET datafusion.runtime.max_spill_merge_fan_in + # Test SET and SHOW runtime.file_statistics_cache_limit statement ok SET datafusion.runtime.file_statistics_cache_limit = '42M' @@ -669,6 +681,7 @@ SELECT name FROM information_schema.df_settings WHERE name LIKE 'datafusion.runt datafusion.runtime.file_statistics_cache_limit datafusion.runtime.list_files_cache_limit datafusion.runtime.list_files_cache_ttl +datafusion.runtime.max_spill_merge_fan_in datafusion.runtime.max_temp_directory_size datafusion.runtime.memory_limit datafusion.runtime.metadata_cache_limit @@ -707,6 +720,64 @@ SET datafusion.runtime.list_files_cache_ttl = '1m18446744073709551555s' statement error DataFusion error: Error during planning: Duration has overflowed allowed maximum limit due to 'mins \* 60 \+ secs' when setting 'datafusion\.runtime\.list_files_cache_ttl' SET datafusion.runtime.list_files_cache_ttl = '1m18446744073709551556s' +# Set invalid value and ensures error +statement error DataFusion error: Error setting config datafusion\.execution\.batch_size\ncaused by\nInvalid or Unsupported Configuration: value must be greater than 0 +SET datafusion.execution.batch_size = 0 + +statement error DataFusion error: Error setting config datafusion\.execution\.meta_fetch_concurrency\ncaused by\nInvalid or Unsupported Configuration: value must be greater than 0 +SET datafusion.execution.meta_fetch_concurrency = 0 + +statement error +SET datafusion.execution.minimum_parallel_output_files = 0 +---- +DataFusion error: Error setting config datafusion.execution.minimum_parallel_output_files +caused by +Invalid or Unsupported Configuration: value must be greater than 0 + + +statement error +SET datafusion.execution.soft_max_rows_per_output_file = 0 +---- +DataFusion error: Error setting config datafusion.execution.soft_max_rows_per_output_file +caused by +Invalid or Unsupported Configuration: value must be greater than 0 + + +statement error +SET datafusion.execution.max_spill_file_size_bytes = 0 +---- +DataFusion error: Error setting config datafusion.execution.max_spill_file_size_bytes +caused by +Invalid or Unsupported Configuration: value must be greater than 0 + + +statement error +SET datafusion.sql_parser.recursion_limit = 0 +---- +DataFusion error: Error setting config datafusion.sql_parser.recursion_limit +caused by +Invalid or Unsupported Configuration: value must be greater than 0 + + +# max_buffered_batches_per_output_file is halved to size an internal channel +# capacity, so 0 and 1 both round down to a zero-capacity channel and must be +# rejected, not just 0. +statement error +SET datafusion.execution.max_buffered_batches_per_output_file = 0 +---- +DataFusion error: Error setting config datafusion.execution.max_buffered_batches_per_output_file +caused by +Invalid or Unsupported Configuration: value must be at least 2 + + +statement error +SET datafusion.execution.max_buffered_batches_per_output_file = 1 +---- +DataFusion error: Error setting config datafusion.execution.max_buffered_batches_per_output_file +caused by +Invalid or Unsupported Configuration: value must be at least 2 + + # Config reset statement ok RESET datafusion.catalog.create_default_catalog_and_schema diff --git a/datafusion/sqllogictest/test_files/simplify_expr.slt b/datafusion/sqllogictest/test_files/simplify_expr.slt index 58ec7a1b262c3..57dc440407dc0 100644 --- a/datafusion/sqllogictest/test_files/simplify_expr.slt +++ b/datafusion/sqllogictest/test_files/simplify_expr.slt @@ -146,3 +146,288 @@ logical_plan physical_plan 01)ProjectionExec: expr=[column1@0 = 1 as opt1, column1@0 = 2 AND column1@0 != 2 as noopt1, column1@0 = 4 as opt2, column1@0 != 5 AND column1@0 = 5 as noopt2] 02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Identity Date cast in a comparison predicate. +# `cast(d AS date)` where `d` is already Date32 is an identity cast and should +# fold away, so the predicate compares against the bare column `d`. This enables +# downstream pruning / filter pushdown that expects a bare-column comparison. +statement ok +create table dates(d date) as values (DATE '2024-01-01'), (DATE '2024-01-02'); + +query TT +explain select d from dates where cast(d as date) = DATE '2024-01-01'; +---- +logical_plan +01)Filter: dates.d = Date32("2024-01-01") +02)--TableScan: dates projection=[d] +physical_plan +01)FilterExec: d@0 = 2024-01-01 +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Identity Date cast inside an `IN` predicate. `IN` goes through a separate +# validation and rewrite path but relies on the same literal-cast helper, so the +# identity `cast(d AS date)` should likewise fold to a bare-column comparison. +query TT +explain select d from dates where cast(d as date) in (DATE '2024-01-01'); +---- +logical_plan +01)Filter: dates.d = Date32("2024-01-01") +02)--TableScan: dates projection=[d] +physical_plan +01)FilterExec: d@0 = 2024-01-01 +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +drop table dates; + +# ------------------------------------------------------------------------ +# Unwrapping Date32 <-> Date64 casts in comparison predicates. +# +# `Date32` counts whole days since the epoch; `Date64` counts milliseconds. +# Widening a `Date32` column up to `Date64` (`date32_col -> Date64`) is +# injective, so a comparison against a whole-day `Date64` literal can be +# rewritten onto the bare `Date32` column. Narrowing a `Date64` column down to +# `Date32` truncates the milliseconds to the day (many-to-one) and must NOT be +# rewritten: `CAST(date64 AS Date32) = ` matches any millisecond within +# that day. Arrow does not require `Date64` values to fall on a day boundary +# (arrow-rs#5288), so the table below intentionally stores sub-day `Date64` +# values (ids 2 and 4) to exercise that hazard. +# +# The `Date64` column is built from raw millisecond values with `arrow_cast`; +# `2025-01-01 00:00` = 1735689600000 ms (day 20089), `2025-01-01 12:00` adds +# 43200000 ms. `1969-12-31 00:00` = -86400000 ms (day -1); `1969-12-31 12:00` +# = -43200000 ms (a pre-epoch sub-day value). +statement ok +create table date_unwrap as +select + c.id, + arrow_cast(c.d32, 'Date32') as d32, + arrow_cast(c.d64ms, 'Date64') as d64 +from (values + (1, '2025-01-01', 1735689600000), + (2, '2025-01-01', 1735732800000), + (3, '1969-12-31', -86400000), + (4, '1969-12-31', -43200000), + (5, NULL, NULL) +) as c(id, d32, d64ms); + +query IDD +select id, d32, d64 from date_unwrap order by id; +---- +1 2025-01-01 2025-01-01T00:00:00 +2 2025-01-01 2025-01-01T12:00:00 +3 1969-12-31 1969-12-31T00:00:00 +4 1969-12-31 1969-12-31T12:00:00 +5 NULL NULL + +# --- Widening Date32 -> Date64: folds onto the bare column --------------- +# The plan for these widening queries is what changes when the optimization is +# enabled: the CAST moves off the column and onto the (whole-day) literal. +query TT +explain select id from date_unwrap where arrow_cast(d32, 'Date64') = arrow_cast(1735689600000, 'Date64'); +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: date_unwrap.d32 = Date32("2025-01-01") +03)----TableScan: date_unwrap projection=[id, d32] +physical_plan +01)FilterExec: d32@1 = 2025-01-01, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I +select id from date_unwrap where arrow_cast(d32, 'Date64') = arrow_cast(1735689600000, 'Date64') order by id; +---- +1 +2 + +# Range operators fold too (Date32 -> Date64 is monotonic). +query TT +explain select id from date_unwrap where arrow_cast(d32, 'Date64') < arrow_cast(1735689600000, 'Date64'); +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: date_unwrap.d32 < Date32("2025-01-01") +03)----TableScan: date_unwrap projection=[id, d32] +physical_plan +01)FilterExec: d32@1 < 2025-01-01, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query TT +explain select id from date_unwrap where arrow_cast(d32, 'Date64') >= arrow_cast(1735689600000, 'Date64'); +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: date_unwrap.d32 >= Date32("2025-01-01") +03)----TableScan: date_unwrap projection=[id, d32] +physical_plan +01)FilterExec: d32@1 >= 2025-01-01, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I +select id from date_unwrap where arrow_cast(d32, 'Date64') < arrow_cast(1735689600000, 'Date64') order by id; +---- +3 +4 + +query I +select id from date_unwrap where arrow_cast(d32, 'Date64') <= arrow_cast(1735689600000, 'Date64') order by id; +---- +1 +2 +3 +4 + +query I +select id from date_unwrap where arrow_cast(d32, 'Date64') > arrow_cast(1735689600000, 'Date64') order by id; +---- + +query I +select id from date_unwrap where arrow_cast(d32, 'Date64') >= arrow_cast(1735689600000, 'Date64') order by id; +---- +1 +2 + +# Reversed operands fold too: with the Date64 literal on the LEFT, logical +# simplification moves the bare column to the left and swaps the operator +# (`literal < CAST(col)` becomes `col > literal`). +query TT +explain select id from date_unwrap where arrow_cast(-86400000, 'Date64') < arrow_cast(d32, 'Date64'); +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: date_unwrap.d32 > Date32("1969-12-31") +03)----TableScan: date_unwrap projection=[id, d32] +physical_plan +01)FilterExec: d32@1 > 1969-12-31, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I +select id from date_unwrap where arrow_cast(-86400000, 'Date64') < arrow_cast(d32, 'Date64') order by id; +---- +1 +2 + +# IN-list widening also folds. +query TT +explain select id from date_unwrap where arrow_cast(d32, 'Date64') in (arrow_cast(1735689600000, 'Date64'), arrow_cast(-86400000, 'Date64')); +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: date_unwrap.d32 = Date32("2025-01-01") OR date_unwrap.d32 = Date32("1969-12-31") +03)----TableScan: date_unwrap projection=[id, d32] +physical_plan +01)FilterExec: d32@1 = 2025-01-01 OR d32@1 = 1969-12-31, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I +select id from date_unwrap where arrow_cast(d32, 'Date64') in (arrow_cast(1735689600000, 'Date64'), arrow_cast(-86400000, 'Date64')) order by id; +---- +1 +2 +3 +4 + +# A NON-whole-day literal is NOT foldable: a Date32-derived Date64 is always at +# midnight, so it can never equal a sub-day literal. The plan keeps the CAST and +# the query returns zero rows. +query TT +explain select id from date_unwrap where arrow_cast(d32, 'Date64') = arrow_cast(1735732800000, 'Date64'); +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: CAST(date_unwrap.d32 AS Date64) = Date64("2025-01-01") +03)----TableScan: date_unwrap projection=[id, d32] +physical_plan +01)FilterExec: CAST(d32@1 AS Date64) = 2025-01-01, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I +select id from date_unwrap where arrow_cast(d32, 'Date64') = arrow_cast(1735732800000, 'Date64') order by id; +---- + +# NULL comparison semantics are unchanged by the rewrite (three-valued logic: +# the NULL row yields NULL, not a dropped row). +query IB +select id, arrow_cast(d32, 'Date64') = arrow_cast(1735689600000, 'Date64') as eq from date_unwrap order by id; +---- +1 true +2 true +3 false +4 false +5 NULL + +# --- Narrowing Date64 -> Date32: must NOT fold (soundness) --------------- +# The plan for these queries is invariant: the CAST stays on the column. If it +# were unwrapped, the sub-day rows (ids 2 and 4) would be dropped. +query TT +explain select id from date_unwrap where cast(d64 as date) = DATE '2025-01-01'; +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: CAST(date_unwrap.d64 AS Date32) = Date32("2025-01-01") +03)----TableScan: date_unwrap projection=[id, d64] +physical_plan +01)FilterExec: CAST(d64@1 AS Date32) = 2025-01-01, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# id 2 is 2025-01-01 12:00 - it truncates to 2025-01-01 and MUST be returned. +query I +select id from date_unwrap where cast(d64 as date) = DATE '2025-01-01' order by id; +---- +1 +2 + +query TT +explain select id from date_unwrap where cast(d64 as date) < DATE '2025-01-01'; +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: CAST(date_unwrap.d64 AS Date32) < Date32("2025-01-01") +03)----TableScan: date_unwrap projection=[id, d64] +physical_plan +01)FilterExec: CAST(d64@1 AS Date32) < 2025-01-01, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# IN-list narrowing is guarded as well. +query TT +explain select id from date_unwrap where cast(d64 as date) in (DATE '2025-01-01'); +---- +logical_plan +01)Projection: date_unwrap.id +02)--Filter: CAST(date_unwrap.d64 AS Date32) = Date32("2025-01-01") +03)----TableScan: date_unwrap projection=[id, d64] +physical_plan +01)FilterExec: CAST(d64@1 AS Date32) = 2025-01-01, projection=[id@0] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +query I +select id from date_unwrap where cast(d64 as date) in (DATE '2025-01-01') order by id; +---- +1 +2 + +# Pre-epoch dates. Arrow's Date64 -> Date32 cast divides by 86_400_000 and +# truncates toward zero, so the pre-epoch sub-day value (id 4, -43200000 ms) +# truncates to day 0 (1970-01-01), not to 1969-12-31. This is arrow's runtime +# behavior; `scale_date_literal` only ever folds on exact whole-day multiples, +# so it can never disagree with the value the cast actually produces. +query ID +select id, cast(d64 as date) as truncated from date_unwrap where d64 is not null order by id; +---- +1 2025-01-01 +2 2025-01-01 +3 1969-12-31 +4 1970-01-01 + +query I +select id from date_unwrap where cast(d64 as date) = DATE '1969-12-31' order by id; +---- +3 + +query I +select id from date_unwrap where cast(d64 as date) = DATE '1970-01-01' order by id; +---- +4 + +statement ok +drop table date_unwrap; diff --git a/datafusion/sqllogictest/test_files/simplify_predicates.slt b/datafusion/sqllogictest/test_files/simplify_predicates.slt index c2a21ea7103c3..44fdedc9c8e1d 100644 --- a/datafusion/sqllogictest/test_files/simplify_predicates.slt +++ b/datafusion/sqllogictest/test_files/simplify_predicates.slt @@ -142,7 +142,7 @@ WHERE int_col > 5 AND float_col BETWEEN 1 AND 100; ---- logical_plan -01)Filter: test_data.str_col LIKE Utf8View("A%") AND test_data.float_col >= Float32(1) AND test_data.float_col <= Float32(100) AND test_data.int_col > Int32(10) +01)Filter: test_data.float_col >= Float32(1) AND test_data.float_col <= Float32(100) AND test_data.int_col > Int32(10) AND test_data.str_col LIKE Utf8View("A%") 02)--TableScan: test_data projection=[int_col, float_col, str_col, date_col, bool_col] statement ok diff --git a/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt b/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt new file mode 100644 index 0000000000000..69bb718bd8c1f --- /dev/null +++ b/datafusion/sqllogictest/test_files/sort_merge_join_spill.slt @@ -0,0 +1,308 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# End-to-end SortMergeJoinExec spilling tests. +# +# Each query runs as an unlimited-memory hash join for expected results, then as +# a memory-limited sort-merge join that must spill. + +hash-threshold 100 + +# Use multiple partitions so the planner can select SortMergeJoinExec. +statement ok +SET datafusion.execution.target_partitions = 2 + +statement ok +SET datafusion.execution.batch_size = 200 + +# Probe rows include one matching key; x=500 yields true, false, and NULL filters. +statement ok +CREATE VIEW probe AS +SELECT value AS k, 500 AS x FROM generate_series(1, 3); + +# Probe rows with no matching buffered key. +statement ok +CREATE VIEW probe_nomatch AS SELECT value AS k FROM generate_series(7, 9); + +# One 2,000-row key group with a 512-byte payload, split into 10 batches. +# Ordered generation avoids input sorts so only the join buffers the payload; +# x includes NULLs and values on both sides of 500. +statement ok +CREATE VIEW wide AS +SELECT 2 AS k, + value AS v, + CASE WHEN value % 10 = 0 THEN cast(NULL AS BIGINT) ELSE value % 1000 END AS x, + lpad(cast(value AS varchar), 512, 'x') AS p +FROM generate_series(1, 2000); + +# One narrow 20,000-row key group for the bitwise semi-join regression. Its +# 200-row input batches fit in 64 KB, while the complete group does not. +statement ok +CREATE VIEW bitwise_wide AS +SELECT 2 AS k, value AS v +FROM generate_series(1, 20000); + +# Keep output narrow while retaining the payload in the buffered input. + +query TT +EXPLAIN SELECT p.k, w.v, length(w.p) FROM probe p JOIN wide w ON p.k = w.k +---- +HashJoinExec + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p JOIN wide w ON p.k = w.k +---- +6000 values hashing to ae029ab21ba6942d04253c3eb1fafbee + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p LEFT JOIN wide w ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM wide w RIGHT JOIN probe p ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p FULL JOIN wide w ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +# Use the unlimited-memory hash join as the reference for filtered joins. +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +JOIN wide w ON p.k = w.k AND p.x < w.x +---- +2700 values hashing to 824832563a1e34fe419885d0b7cccc9d + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +LEFT JOIN wide w ON p.k = w.k AND p.x < w.x +---- +2706 values hashing to 38c8a5628a41b463e41d592c2401ca8d + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM wide w +RIGHT JOIN probe p ON p.k = w.k AND p.x < w.x +---- +2706 values hashing to 38c8a5628a41b463e41d592c2401ca8d + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +FULL JOIN wide w ON p.k = w.k AND p.x < w.x +---- +6006 values hashing to b6b875b2658ee19e8bca7cd6e993dee1 + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe_nomatch p FULL JOIN wide w ON p.k = w.k +---- +6009 values hashing to 126cd87356bc636448b65c5fb5f4bd2b + +# Only the first input batch satisfies this filtered semi join. +query I +SELECT pr.k +FROM probe pr +WHERE EXISTS ( + SELECT 1 + FROM bitwise_wide wi + WHERE pr.k = wi.k + AND wi.v <= pr.x - 300 +) +---- +2 + +# A 64 KB pool spills all 10 buffered batches; each result must match its +# unlimited-memory hash-join reference. + +statement ok +SET datafusion.optimizer.prefer_hash_join = false + +statement ok +SET datafusion.runtime.memory_limit = '64K' + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p JOIN wide w ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Inner, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p JOIN wide w ON p.k = w.k +---- +6000 values hashing to ae029ab21ba6942d04253c3eb1fafbee + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p LEFT JOIN wide w ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Left, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p LEFT JOIN wide w ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM wide w RIGHT JOIN probe p ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Right, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM wide w RIGHT JOIN probe p ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p FULL JOIN wide w ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Full, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p FULL JOIN wide w ON p.k = w.k +---- +6006 values hashing to 352109cc65a61f6224bb027dbee60df5 + +# Filtered spills cover true, false, and NULL masks; outer joins defer unmatched +# rows until the whole key group is restored. + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p +JOIN wide w ON p.k = w.k AND p.x < w.x +---- +Plan with Metrics +SortMergeJoinExec: join_type=Inner, on=[(k@0, k@0)], filter=x@0 < x@1, metrics=[output_rows=900,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +JOIN wide w ON p.k = w.k AND p.x < w.x +---- +2700 values hashing to 824832563a1e34fe419885d0b7cccc9d + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p +LEFT JOIN wide w ON p.k = w.k AND p.x < w.x +---- +Plan with Metrics +SortMergeJoinExec: join_type=Left, on=[(k@0, k@0)], filter=x@0 < x@1, metrics=[output_rows=902,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +LEFT JOIN wide w ON p.k = w.k AND p.x < w.x +---- +2706 values hashing to 38c8a5628a41b463e41d592c2401ca8d + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM wide w +RIGHT JOIN probe p ON p.k = w.k AND p.x < w.x +---- +Plan with Metrics +SortMergeJoinExec: join_type=Right, on=[(k@0, k@0)], filter=x@1 < x@0, metrics=[output_rows=902,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM wide w +RIGHT JOIN probe p ON p.k = w.k AND p.x < w.x +---- +2706 values hashing to 38c8a5628a41b463e41d592c2401ca8d + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe p +FULL JOIN wide w ON p.k = w.k AND p.x < w.x +---- +Plan with Metrics +SortMergeJoinExec: join_type=Full, on=[(k@0, k@0)], filter=x@0 < x@1, metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe p +FULL JOIN wide w ON p.k = w.k AND p.x < w.x +---- +6006 values hashing to b6b875b2658ee19e8bca7cd6e993dee1 + +# Full join restores all buffered batches to emit unmatched rows. + +query TT +EXPLAIN ANALYZE +SELECT p.k, w.v, length(w.p) FROM probe_nomatch p FULL JOIN wide w ON p.k = w.k +---- +Plan with Metrics +SortMergeJoinExec: join_type=Full, on=[(k@0, k@0)], metrics=[output_rows=2.00 K,spill_count=10, spilled_bytes=spilled_rows=2.00 K, peak_mem_used= + +query III rowsort +SELECT p.k, w.v, length(w.p) FROM probe_nomatch p FULL JOIN wide w ON p.k = w.k +---- +6009 values hashing to 126cd87356bc636448b65c5fb5f4bd2b + +# Let the required narrow input sorts merge within the constrained pool. +statement ok +SET datafusion.execution.sort_spill_reservation_bytes = 0 + +# Prove the correlated EXISTS uses the filtered bitwise LeftSemi stream and +# spills the complete multi-batch key group. +query TT +EXPLAIN ANALYZE +SELECT pr.k +FROM probe pr +WHERE EXISTS ( + SELECT 1 + FROM bitwise_wide wi + WHERE pr.k = wi.k + AND wi.v <= pr.x - 300 +) +---- +Plan with Metrics +SortMergeJoinExec: join_type=LeftSemi, on=[(k@0, k@0)], filter=v@1 <= x@0 - 300, metrics=[output_rows=1,spill_count=1, spilled_bytes= KB, spilled_rows= K, peak_mem_used= + +# The same query must retain the matching first slice after later overflows. +query I +SELECT pr.k +FROM probe pr +WHERE EXISTS ( + SELECT 1 + FROM bitwise_wide wi + WHERE pr.k = wi.k + AND wi.v <= pr.x - 300 +) +---- +2 + +statement ok +RESET datafusion.execution.sort_spill_reservation_bytes + +statement ok +RESET datafusion.runtime.memory_limit + +statement ok +RESET datafusion.optimizer.prefer_hash_join + +statement ok +RESET datafusion.execution.batch_size + +statement ok +SET datafusion.execution.target_partitions = 4 + +statement ok +RESET datafusion.catalog.create_default_catalog_and_schema diff --git a/datafusion/sqllogictest/test_files/sort_pushdown.slt b/datafusion/sqllogictest/test_files/sort_pushdown.slt index 540562eb3bc8d..f2442762f3fd2 100644 --- a/datafusion/sqllogictest/test_files/sort_pushdown.slt +++ b/datafusion/sqllogictest/test_files/sort_pushdown.slt @@ -43,7 +43,7 @@ logical_plan 02)--TableScan: sorted_parquet projection=[id, value, name] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_data.parquet]]}, projection=[id, value, name], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_data.parquet]]}, projection=[id, value, name], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible # Test 1.2: Verify results are correct query IIT @@ -74,7 +74,7 @@ logical_plan 02)--TableScan: sorted_parquet projection=[id, value, name] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_data.parquet]]}, projection=[id, value, name], output_ordering=[id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_data.parquet]]}, projection=[id, value, name], output_ordering=[id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Re-enable statement ok @@ -91,7 +91,7 @@ logical_plan physical_plan 01)GlobalLimitExec: skip=2, fetch=3 02)--SortExec: TopK(fetch=5), expr=[id@0 DESC], preserve_partitioning=[false] -03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_data.parquet]]}, projection=[id, value, name], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_data.parquet]]}, projection=[id, value, name], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query IIT SELECT * FROM sorted_parquet ORDER BY id DESC LIMIT 3 OFFSET 2; @@ -155,7 +155,7 @@ logical_plan 03)----TableScan: multi_rg_sorted projection=[id, category, value], partial_filters=[multi_rg_sorted.category = Utf8View("alpha") OR multi_rg_sorted.category = Utf8View("gamma")] physical_plan 01)SortExec: TopK(fetch=5), expr=[id@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/multi_rg_sorted.parquet]]}, projection=[id, category, value], file_type=parquet, predicate=(category@1 = alpha OR category@1 = gamma) AND DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, pruning_predicate=category_null_count@2 != row_count@3 AND category_min@0 <= alpha AND alpha <= category_max@1 OR category_null_count@2 != row_count@3 AND category_min@0 <= gamma AND gamma <= category_max@1, required_guarantees=[category in (alpha, gamma)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/multi_rg_sorted.parquet]]}, projection=[id, category, value], file_type=parquet, predicate=(category@1 = alpha OR category@1 = gamma) AND DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=category_null_count@2 != row_count@3 AND category_min@0 <= alpha AND alpha <= category_max@1 OR category_null_count@2 != row_count@3 AND category_min@0 <= gamma AND gamma <= category_max@1, required_guarantees=[category in (alpha, gamma)] # Verify the results are correct despite reverse scanning with row selection # Expected: gamma values (6, 5) then alpha values (2, 1), in DESC order by id @@ -272,7 +272,7 @@ logical_plan physical_plan 01)SortPreservingMergeExec: [id@0 DESC], fetch=3 02)--SortExec: TopK(fetch=3), expr=[id@0 DESC], preserve_partitioning=[true] -03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_multi/part1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_multi/part2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_multi/part3.parquet]]}, projection=[id, value, name], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true +03)----DataSourceExec: file_groups={3 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_multi/part1.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_multi/part2.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/sorted_multi/part3.parquet]]}, projection=[id, value, name], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible # Verify correctness with repartitioning and multiple files query IIT @@ -381,7 +381,7 @@ logical_plan 03)----TableScan: timeseries_parquet projection=[timeframe, period_end, value], partial_filters=[timeseries_parquet.timeframe = Utf8View("quarterly")] physical_plan 01)SortExec: TopK(fetch=2), expr=[period_end@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], file_type=parquet, predicate=timeframe@0 = quarterly AND DynamicFilter [ empty ], sort_order_for_reorder=[period_end@1 DESC], reverse_row_groups=true, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= quarterly AND quarterly <= timeframe_max@1, required_guarantees=[timeframe in (quarterly)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], file_type=parquet, predicate=timeframe@0 = quarterly AND DynamicFilter [ empty ], sort_order_for_reorder=[period_end@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= quarterly AND quarterly <= timeframe_max@1, required_guarantees=[timeframe in (quarterly)] # Test 2.2: Verify the results are correct query TIR @@ -440,7 +440,7 @@ logical_plan 02)--TableScan: timeseries_parquet projection=[timeframe, period_end, value] physical_plan 01)SortExec: TopK(fetch=3), expr=[timeframe@0 ASC NULLS LAST, period_end@1 DESC], preserve_partitioning=[false], sort_prefix=[timeframe@0 ASC NULLS LAST] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], output_ordering=[timeframe@0 ASC NULLS LAST, period_end@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], output_ordering=[timeframe@0 ASC NULLS LAST, period_end@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Test 2.7: Disable sort pushdown and verify filter still works statement ok @@ -458,7 +458,7 @@ logical_plan 03)----TableScan: timeseries_parquet projection=[timeframe, period_end, value], partial_filters=[timeseries_parquet.timeframe = Utf8View("quarterly")] physical_plan 01)SortExec: TopK(fetch=2), expr=[period_end@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], output_ordering=[timeframe@0 ASC NULLS LAST, period_end@1 ASC NULLS LAST], file_type=parquet, predicate=timeframe@0 = quarterly AND DynamicFilter [ empty ], pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= quarterly AND quarterly <= timeframe_max@1, required_guarantees=[timeframe in (quarterly)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], output_ordering=[timeframe@0 ASC NULLS LAST, period_end@1 ASC NULLS LAST], file_type=parquet, predicate=timeframe@0 = quarterly AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= quarterly AND quarterly <= timeframe_max@1, required_guarantees=[timeframe in (quarterly)] # Results should still be correct query TIR @@ -491,7 +491,7 @@ logical_plan 03)----TableScan: timeseries_parquet projection=[timeframe, period_end, value], partial_filters=[timeseries_parquet.timeframe = Utf8View("daily") OR timeseries_parquet.timeframe = Utf8View("weekly")] physical_plan 01)SortExec: TopK(fetch=3), expr=[period_end@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], file_type=parquet, predicate=(timeframe@0 = daily OR timeframe@0 = weekly) AND DynamicFilter [ empty ], sort_order_for_reorder=[period_end@1 DESC], reverse_row_groups=true, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= daily AND daily <= timeframe_max@1 OR timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= weekly AND weekly <= timeframe_max@1, required_guarantees=[timeframe in (daily, weekly)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], file_type=parquet, predicate=(timeframe@0 = daily OR timeframe@0 = weekly) AND DynamicFilter [ empty ], sort_order_for_reorder=[period_end@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= daily AND daily <= timeframe_max@1 OR timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= weekly AND weekly <= timeframe_max@1, required_guarantees=[timeframe in (daily, weekly)] # Test 2.9: Complex case - literal constant in sort expression itself # The literal 'constant' is ignored in sort analysis @@ -511,7 +511,7 @@ logical_plan 03)----TableScan: timeseries_parquet projection=[timeframe, period_end, value], partial_filters=[timeseries_parquet.timeframe = Utf8View("monthly")] physical_plan 01)SortExec: TopK(fetch=2), expr=[period_end@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], file_type=parquet, predicate=timeframe@0 = monthly AND DynamicFilter [ empty ], sort_order_for_reorder=[period_end@1 DESC], reverse_row_groups=true, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= monthly AND monthly <= timeframe_max@1, required_guarantees=[timeframe in (monthly)] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timeseries_sorted.parquet]]}, projection=[timeframe, period_end, value], file_type=parquet, predicate=timeframe@0 = monthly AND DynamicFilter [ empty ], sort_order_for_reorder=[period_end@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible, pruning_predicate=timeframe_null_count@2 != row_count@3 AND timeframe_min@0 <= monthly AND monthly <= timeframe_max@1, required_guarantees=[timeframe in (monthly)] # Verify results query TIR @@ -600,7 +600,7 @@ logical_plan 02)--TableScan: timestamp_parquet projection=[id, ts, volume, price] physical_plan 01)SortExec: TopK(fetch=3), expr=[ts@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timestamp_sorted.parquet]]}, projection=[id, ts, volume, price], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[ts@1 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timestamp_sorted.parquet]]}, projection=[id, ts, volume, price], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[ts@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible # Verify results query IPIR @@ -626,7 +626,7 @@ logical_plan 02)--TableScan: timestamp_parquet projection=[id, ts, volume, price] physical_plan 01)SortExec: TopK(fetch=3), expr=[date_trunc(day, ts@1) DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timestamp_sorted.parquet]]}, projection=[id, ts, volume, price], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[date_trunc(day, ts@1) DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/timestamp_sorted.parquet]]}, projection=[id, ts, volume, price], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[date_trunc(day, ts@1) DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible # Verify results (descending day) query IPIR @@ -686,7 +686,7 @@ logical_plan 02)--TableScan: multi_month_parquet projection=[id, ts, volume, price] physical_plan 01)SortExec: TopK(fetch=2), expr=[ts@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/multi_month_sorted.parquet]]}, projection=[id, ts, volume, price], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[ts@1 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/multi_month_sorted.parquet]]}, projection=[id, ts, volume, price], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[ts@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query IPIR SELECT * FROM multi_month_parquet @@ -712,7 +712,7 @@ logical_plan 02)--TableScan: multi_month_parquet projection=[id, ts, volume, price] physical_plan 01)SortExec: TopK(fetch=2), expr=[date_trunc(month, ts@1) DESC, ts@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/multi_month_sorted.parquet]]}, projection=[id, ts, volume, price], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[date_trunc(month, ts@1) DESC, ts@1 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/multi_month_sorted.parquet]]}, projection=[id, ts, volume, price], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[date_trunc(month, ts@1) DESC, ts@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query IPIR SELECT * FROM multi_month_parquet @@ -754,7 +754,7 @@ logical_plan 02)--TableScan: int_parquet projection=[id, small_val, big_val] physical_plan 01)SortExec: TopK(fetch=2), expr=[CAST(small_val@1 AS Int64) DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/int_sorted.parquet]]}, projection=[id, small_val, big_val], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[CAST(small_val@1 AS Int64) DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/int_sorted.parquet]]}, projection=[id, small_val, big_val], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[CAST(small_val@1 AS Int64) DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query III SELECT * FROM int_parquet @@ -796,7 +796,7 @@ logical_plan 02)--TableScan: float_parquet projection=[id, value] physical_plan 01)SortExec: TopK(fetch=3), expr=[ceil(value@1) DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/float_sorted.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[ceil(value@1) DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/float_sorted.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[ceil(value@1) DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query IR SELECT * FROM float_parquet @@ -839,7 +839,7 @@ logical_plan 02)--TableScan: signed_parquet projection=[id, value] physical_plan 01)SortExec: TopK(fetch=3), expr=[abs(value@1) DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/signed_sorted.parquet]]}, projection=[id, value], output_ordering=[value@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/signed_sorted.parquet]]}, projection=[id, value], output_ordering=[value@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Results should still be correct (no optimization applied) query IR @@ -1100,8 +1100,9 @@ CREATE EXTERNAL TABLE reversed_parquet(id INT, value INT) STORED AS PARQUET LOCATION 'test_files/scratch/sort_pushdown/reversed/'; -# Test 4.1: PushdownSort reorders files by min/max statistics so they are -# already in correct sort order → non-overlapping → no SortExec needed. +# Test 4.1: PushdownSort reorders files by min/max statistics; the +# post-sort file groups are non-overlapping, the inferred ordering +# re-validates, and the SortExec above can be eliminated. # (files reordered from [a_high, b_mid, c_low] to [c_low, b_mid, a_high]) query TT EXPLAIN SELECT * FROM reversed_parquet ORDER BY id ASC; @@ -1109,9 +1110,7 @@ EXPLAIN SELECT * FROM reversed_parquet ORDER BY id ASC; logical_plan 01)Sort: reversed_parquet.id ASC NULLS LAST 02)--TableScan: reversed_parquet projection=[id, value] -physical_plan -01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/c_low.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/a_high.parquet]]}, projection=[id, value], file_type=parquet, sort_order_for_reorder=[id@0 ASC NULLS LAST] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/c_low.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/a_high.parquet]]}, projection=[id, value], output_ordering=[id@0 ASC NULLS LAST], file_type=parquet # Test 4.2: Results must be correct query II @@ -1175,10 +1174,153 @@ SELECT * FROM overlap_parquet ORDER BY id ASC; 5 500 6 600 +# Test 5b: Safety case — no WITH ORDER, files written without ORDER BY (no +# sorting_columns metadata). Source has no way to declare per-file ordering, +# so even though min/max stats happen to be non-overlapping, the optimizer +# must NOT eliminate SortExec. +statement ok +CREATE TABLE no_decl_low(id INT, value INT) AS VALUES (1, 100), (3, 300), (2, 200); + +statement ok +CREATE TABLE no_decl_mid(id INT, value INT) AS VALUES (6, 600), (4, 400), (5, 500); + +statement ok +CREATE TABLE no_decl_high(id INT, value INT) AS VALUES (9, 900), (8, 800), (7, 700); + +# Write WITHOUT ORDER BY so each file lacks sorting_columns metadata. +query I +COPY no_decl_low TO 'test_files/scratch/sort_pushdown/no_decl/a_low.parquet'; +---- +3 + +query I +COPY no_decl_mid TO 'test_files/scratch/sort_pushdown/no_decl/b_mid.parquet'; +---- +3 + +query I +COPY no_decl_high TO 'test_files/scratch/sort_pushdown/no_decl/c_high.parquet'; +---- +3 + +statement ok +CREATE EXTERNAL TABLE no_decl_parquet(id INT, value INT) +STORED AS PARQUET +LOCATION 'test_files/scratch/sort_pushdown/no_decl/'; + +# Min/max stats per file happen to be non-overlapping (1-3, 4-6, 7-9) but the +# rows inside each file are NOT sorted by id. Without an ordering declaration +# (WITH ORDER or parquet sorting_columns), the optimizer cannot prove the +# output would be sorted — SortExec must stay. +query TT +EXPLAIN SELECT * FROM no_decl_parquet ORDER BY id ASC; +---- +logical_plan +01)Sort: no_decl_parquet.id ASC NULLS LAST +02)--TableScan: no_decl_parquet projection=[id, value] +physical_plan +01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/no_decl/a_low.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/no_decl/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/no_decl/c_high.parquet]]}, projection=[id, value], file_type=parquet, sort_order_for_reorder=[id@0 ASC NULLS LAST] + +# Results must still be correct (SortExec does the final sort) +query II +SELECT * FROM no_decl_parquet ORDER BY id ASC; +---- +1 100 +2 200 +3 300 +4 400 +5 500 +6 600 +7 700 +8 800 +9 900 + +# Cleanup Test 5b +statement ok +DROP TABLE no_decl_low; + +statement ok +DROP TABLE no_decl_mid; + +statement ok +DROP TABLE no_decl_high; + +statement ok +DROP TABLE no_decl_parquet; + +# Test 5c: NULL safety — files in **wrong** filesystem order so the +# Inexact branch fires; the previously-non-last file contains NULLs in +# the sort column. With NULLS LAST, NULLs inside a file sit after all +# non-null rows. If the next file's non-null values are smaller than +# the previous file's max, those values would land AFTER the NULLs in +# the concatenated stream — breaking the ordering. The fix must NOT +# upgrade to Exact here even though stats are non-overlapping. + +statement ok +CREATE TABLE null_safety_high(id INT, value INT) AS VALUES (4, 400), (5, 500), (6, 600); + +statement ok +CREATE TABLE null_safety_low_with_nulls(id INT, value INT) AS VALUES (1, 100), (2, 200), (3, 300), (NULL, 999); + +# Name files so alphabetical order is REVERSED relative to id order +# (a_high before b_low) — triggers the Inexact / re-validate path. +query I +COPY (SELECT * FROM null_safety_high ORDER BY id ASC NULLS LAST) +TO 'test_files/scratch/sort_pushdown/null_safety/a_high.parquet'; +---- +3 + +query I +COPY (SELECT * FROM null_safety_low_with_nulls ORDER BY id ASC NULLS LAST) +TO 'test_files/scratch/sort_pushdown/null_safety/b_low_nulls.parquet'; +---- +4 + +statement ok +CREATE EXTERNAL TABLE null_safety_parquet(id INT, value INT) +STORED AS PARQUET +LOCATION 'test_files/scratch/sort_pushdown/null_safety/' +WITH ORDER (id ASC NULLS LAST); + +# After Phase 2 reorder file_groups would be [b_low_nulls, a_high] and +# min/max would be non-overlapping — but b_low_nulls has NULLs in the +# sort column, so we must NOT upgrade to Exact. SortExec stays. +query TT +EXPLAIN SELECT * FROM null_safety_parquet ORDER BY id ASC NULLS LAST; +---- +logical_plan +01)Sort: null_safety_parquet.id ASC NULLS LAST +02)--TableScan: null_safety_parquet projection=[id, value] +physical_plan +01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/null_safety/b_low_nulls.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/null_safety/a_high.parquet]]}, projection=[id, value], file_type=parquet, sort_order_for_reorder=[id@0 ASC NULLS LAST] + +# Results must still be correct (SortExec does the final sort) +query II +SELECT * FROM null_safety_parquet ORDER BY id ASC NULLS LAST; +---- +1 100 +2 200 +3 300 +4 400 +5 500 +6 600 +NULL 999 + +statement ok +DROP TABLE null_safety_high; + +statement ok +DROP TABLE null_safety_low_with_nulls; + +statement ok +DROP TABLE null_safety_parquet; + # Test 6: WITH ORDER + reversed filesystem order # Same file setup as Test 4 but explicitly declaring ordering via WITH ORDER. -# Even with WITH ORDER, the optimizer should detect that inter-file order is wrong -# and keep SortExec. +# PushdownSort reorders files by min/max stats; after reorder the inter-file +# ordering re-validates and the SortExec above is eliminated. statement ok CREATE EXTERNAL TABLE reversed_with_order_parquet(id INT, value INT) @@ -1194,9 +1336,7 @@ EXPLAIN SELECT * FROM reversed_with_order_parquet ORDER BY id ASC; logical_plan 01)Sort: reversed_with_order_parquet.id ASC NULLS LAST 02)--TableScan: reversed_with_order_parquet projection=[id, value] -physical_plan -01)SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/c_low.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/a_high.parquet]]}, projection=[id, value], file_type=parquet, sort_order_for_reorder=[id@0 ASC NULLS LAST] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/c_low.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/reversed/a_high.parquet]]}, projection=[id, value], output_ordering=[id@0 ASC NULLS LAST], file_type=parquet # Test 6.2: Results must be correct query II @@ -1333,9 +1473,7 @@ EXPLAIN SELECT * FROM desc_reversed_parquet ORDER BY id DESC; logical_plan 01)Sort: desc_reversed_parquet.id DESC NULLS FIRST 02)--TableScan: desc_reversed_parquet projection=[id, value] -physical_plan -01)SortExec: expr=[id@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/desc_reversed/b_high.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/desc_reversed/a_low.parquet]]}, projection=[id, value], file_type=parquet, sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/desc_reversed/b_high.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/desc_reversed/a_low.parquet]]}, projection=[id, value], output_ordering=[id@0 DESC], file_type=parquet # Test 8.2: Results must be correct query II @@ -1348,6 +1486,78 @@ SELECT * FROM desc_reversed_parquet ORDER BY id DESC; 2 200 1 100 +# Test 8b: DESC with multiple row groups per file sharing a min value. +# Regression test for the Inexact→Exact upgrade: when SortExec is eliminated +# the files must be read in natural order. The opener's runtime row-group +# reorder (sort ASC-by-min then reverse) mis-orders two row groups in one file +# that share the same min — so the upgrade must NOT leave those hints active. +# +# File b_high is DESC-sorted [10,8,8,8] written with 2 rows per row group: +# RG0 = [10, 8] (min 8, max 10) +# RG1 = [ 8, 8] (min 8, max 8) +# Both row groups have min=8. Naively reordering RGs ASC-by-min then reversing +# yields [RG1, RG0] → 8,8,10,8 (wrong). Natural order [RG0, RG1] is correct. + +statement ok +CREATE TABLE rg_desc_high(id INT, value INT) AS VALUES (10, 100), (8, 801), (8, 802), (8, 803); + +statement ok +CREATE TABLE rg_desc_low(id INT, value INT) AS VALUES (3, 300), (2, 200), (1, 100); + +query I +COPY (SELECT * FROM rg_desc_high ORDER BY id DESC) +TO 'test_files/scratch/sort_pushdown/rg_desc/b_high.parquet' +OPTIONS ('format.max_row_group_size' '2'); +---- +4 + +query I +COPY (SELECT * FROM rg_desc_low ORDER BY id DESC) +TO 'test_files/scratch/sort_pushdown/rg_desc/a_low.parquet' +OPTIONS ('format.max_row_group_size' '2'); +---- +3 + +# Files named so filesystem order [a_low, b_high] is wrong for DESC → the +# Inexact path fires, stats reorder makes file groups [b_high, a_low] +# non-overlapping, and the upgrade eliminates SortExec. +statement ok +CREATE EXTERNAL TABLE rg_desc_parquet(id INT, value INT) +STORED AS PARQUET +LOCATION 'test_files/scratch/sort_pushdown/rg_desc/' +WITH ORDER (id DESC); + +# SortExec eliminated, files reordered, NO sort_order_for_reorder / +# reverse_row_groups (natural read is correct after the upgrade). +query TT +EXPLAIN SELECT id FROM rg_desc_parquet ORDER BY id DESC; +---- +logical_plan +01)Sort: rg_desc_parquet.id DESC NULLS FIRST +02)--TableScan: rg_desc_parquet projection=[id] +physical_plan DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/rg_desc/b_high.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/rg_desc/a_low.parquet]]}, projection=[id], output_ordering=[id@0 DESC], file_type=parquet + +# Results must be in DESC order — id=10 first. +query I +SELECT id FROM rg_desc_parquet ORDER BY id DESC; +---- +10 +8 +8 +8 +3 +2 +1 + +statement ok +DROP TABLE rg_desc_parquet; + +statement ok +DROP TABLE rg_desc_high; + +statement ok +DROP TABLE rg_desc_low; + # Test 9: Multi-column sort key validation # Files have (category, id) ordering. Files share a boundary value on category='B' # so column-level min/max statistics overlap on the primary key column. @@ -1778,7 +1988,7 @@ logical_plan 02)--TableScan: tb_overlap projection=[id, value] physical_plan 01)SortExec: TopK(fetch=5), expr=[id@0 DESC, value@1 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tb_overlap/file_z.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tb_overlap/file_y.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tb_overlap/file_x.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC, value@1 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tb_overlap/file_z.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tb_overlap/file_y.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tb_overlap/file_x.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC, value@1 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query II SELECT * FROM tb_overlap ORDER BY id DESC, value DESC LIMIT 5; @@ -1863,7 +2073,7 @@ logical_plan 02)--TableScan: tc_limit projection=[id, value] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tc_limit/file_c.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tc_limit/file_b.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tc_limit/file_a.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tc_limit/file_c.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tc_limit/file_b.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tc_limit/file_a.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query II SELECT * FROM tc_limit ORDER BY id DESC LIMIT 3; @@ -2218,7 +2428,7 @@ STORED AS PARQUET LOCATION 'test_files/scratch/sort_pushdown/tg_buffer/' WITH ORDER (id ASC); -# Test G.1: BufferExec appears between SPM and DataSourceExec +# Test G.1: SortExec eliminated; BufferExec replaces it between SPM and DataSourceExec query TT EXPLAIN SELECT * FROM tg_buffer ORDER BY id ASC; ---- @@ -2227,8 +2437,8 @@ logical_plan 02)--TableScan: tg_buffer projection=[id, value] physical_plan 01)SortPreservingMergeExec: [id@0 ASC NULLS LAST] -02)--SortExec: expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/a_high.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/c_low.parquet]]}, projection=[id, value], file_type=parquet, sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--BufferExec: capacity=1073741824 +03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/a_high.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/c_low.parquet]]}, projection=[id, value], output_ordering=[id@0 ASC NULLS LAST], file_type=parquet # Verify correctness query II @@ -2245,7 +2455,7 @@ SELECT * FROM tg_buffer ORDER BY id ASC; 9 900 10 1000 -# Test G.2: LIMIT query with BufferExec +# Test G.2: LIMIT query — SortExec eliminated, limit pushed to source; BufferExec stays query TT EXPLAIN SELECT * FROM tg_buffer ORDER BY id ASC LIMIT 3; ---- @@ -2254,8 +2464,8 @@ logical_plan 02)--TableScan: tg_buffer projection=[id, value] physical_plan 01)SortPreservingMergeExec: [id@0 ASC NULLS LAST], fetch=3 -02)--SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/a_high.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/c_low.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--BufferExec: capacity=1073741824 +03)----DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/b_mid.parquet, WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/a_high.parquet], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tg_buffer/c_low.parquet]]}, projection=[id, value], limit=3, output_ordering=[id@0 ASC NULLS LAST], file_type=parquet query II SELECT * FROM tg_buffer ORDER BY id ASC LIMIT 3; @@ -2334,7 +2544,7 @@ logical_plan 02)--TableScan: th_reorder projection=[id, value] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/th_reorder/data.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/th_reorder/data.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 ASC NULLS LAST], dynamic_rg_pruning=eligible # Results must be correct regardless of RG reorder. query II @@ -2353,7 +2563,7 @@ logical_plan 02)--TableScan: th_reorder projection=[id, value] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/th_reorder/data.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/th_reorder/data.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query II SELECT * FROM th_reorder ORDER BY id DESC LIMIT 3; @@ -2439,7 +2649,7 @@ logical_plan 02)--TableScan: tj_scrambled projection=[id, value] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tj_scrambled/data.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tj_scrambled/data.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible # Test J.2: Results must be correct query II @@ -2629,7 +2839,7 @@ logical_plan 02)--TableScan: tl_sorted projection=[id, value] physical_plan 01)SortExec: TopK(fetch=3), expr=[id@0 DESC, value@1 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tl_multikey/data.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC, value@1 ASC NULLS LAST], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/sort_pushdown/tl_multikey/data.parquet]]}, projection=[id, value], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[id@0 DESC, value@1 ASC NULLS LAST], reverse_row_groups=true, dynamic_rg_pruning=eligible query II SELECT id, value FROM tl_sorted ORDER BY id DESC, value ASC LIMIT 3; diff --git a/datafusion/sqllogictest/test_files/spark/aggregate/collect_window.slt b/datafusion/sqllogictest/test_files/spark/aggregate/collect_window.slt new file mode 100644 index 0000000000000..5661cb9432427 --- /dev/null +++ b/datafusion/sqllogictest/test_files/spark/aggregate/collect_window.slt @@ -0,0 +1,398 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +####### +# Tests for Spark-compat collect_list / collect_set as WINDOW functions. +# Spark semantics: +# - NULL inputs are skipped (Hive collect_list/collect_set behavior). +# - An empty frame (or one where all inputs were NULL) evaluates to [] +# rather than NULL (nullable = false in Spark's Collect aggregate). +# - collect_list preserves frame order; collect_set deduplicates. +# Validates that NullToEmptyListAccumulator forwards retract_batch +# so the wrapped ArrayAggAccumulator / DistinctArrayAggAccumulator can +# drive sliding window frames. +####### + +statement ok +CREATE TABLE t(ts INT, val TEXT) AS VALUES + (1, 'A'), (2, 'B'), (3, 'C'), (4, 'D'), (5, 'E'); + +# Unbounded preceding frame — accumulator only sees update_batch. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) +FROM t; +---- +[A] +[A, B] +[A, B, C] +[A, B, C, D] +[A, B, C, D, E] + +# Bounded sliding ROWS frame — requires retract_batch. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM t; +---- +[A] +[A, B] +[B, C] +[C, D] +[D, E] + +# Wider sliding window. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND CURRENT ROW) +FROM t; +---- +[A] +[A, B] +[A, B, C] +[B, C, D] +[C, D, E] + +# Centered sliding window with PRECEDING + FOLLOWING. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND 1 FOLLOWING) +FROM t; +---- +[A, B] +[A, B, C] +[B, C, D] +[C, D, E] +[D, E] + +# Unbounded both sides — every row sees the full input. +query ? +SELECT collect_list(val) + OVER (ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING) +FROM t; +---- +[A, B, C, D, E] +[A, B, C, D, E] +[A, B, C, D, E] +[A, B, C, D, E] +[A, B, C, D, E] + +# Single-row frame. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN CURRENT ROW AND CURRENT ROW) +FROM t; +---- +[A] +[B] +[C] +[D] +[E] + +# Empty leading frame on the first row — Spark returns []. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND 1 PRECEDING) +FROM t; +---- +[] +[A] +[A, B] +[B, C] +[C, D] + +# Empty trailing frame on the last row — Spark returns []. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN 1 FOLLOWING AND 2 FOLLOWING) +FROM t; +---- +[B, C] +[C, D] +[D, E] +[E] +[] + +####### +# NULL handling — Spark's collect_list skips NULL inputs. +####### + +statement ok +CREATE TABLE t_nulls(ts INT, val TEXT) AS VALUES + (1, 'A'), (2, NULL), (3, 'C'), (4, NULL), (5, 'E'); + +# NULLs filtered out of the materialized list, but the row still emits one entry. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM t_nulls; +---- +[A] +[A] +[C] +[C] +[E] + +# Wider frame. +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND CURRENT ROW) +FROM t_nulls; +---- +[A] +[A] +[A, C] +[C] +[C, E] + +# All-NULL frame collapses to []. +statement ok +CREATE TABLE t_allnull(ts INT, val TEXT) AS VALUES + (1, NULL), (2, NULL), (3, NULL); + +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW) +FROM t_allnull; +---- +[] +[] +[] + +####### +# PARTITION BY — each partition starts with fresh accumulator state. +####### + +statement ok +CREATE TABLE t_parts(grp INT, ts INT, val TEXT) AS VALUES + (1, 1, 'A'), (1, 2, 'B'), (1, 3, 'C'), + (2, 1, 'X'), (2, 2, 'Y'), (2, 3, 'Z'); + +query I? +SELECT grp, collect_list(val) + OVER (PARTITION BY grp ORDER BY ts + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM t_parts +ORDER BY grp, ts; +---- +1 [A] +1 [A, B] +1 [B, C] +2 [X] +2 [X, Y] +2 [Y, Z] + +####### +# RANGE frame with value gaps — exercises multi-row retract. +####### + +statement ok +CREATE TABLE t_range(ts INT, val TEXT) AS VALUES + (1, 'A'), (2, 'B'), (3, 'C'), (4, 'D'), (100, 'E'); + +query ? +SELECT collect_list(val) + OVER (ORDER BY ts RANGE BETWEEN 2 PRECEDING AND 2 FOLLOWING) +FROM t_range; +---- +[A, B, C] +[A, B, C, D] +[A, B, C, D] +[B, C, D] +[E] + +####### +# GROUPS frame — rows tied on ORDER BY are processed together. +####### + +statement ok +CREATE TABLE t_groups(ts INT, val TEXT) AS VALUES + (1, 'A'), (1, 'B'), (2, 'C'), (2, 'D'), (3, 'E'); + +query ? +SELECT collect_list(val) + OVER (ORDER BY ts GROUPS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM t_groups; +---- +[A, B] +[A, B] +[A, B, C, D] +[A, B, C, D] +[C, D, E] + +####### +# Integer-typed input — guards against type-specific regressions. +####### + +statement ok +CREATE TABLE t_int(ts INT, val INT) AS VALUES + (1, 10), (2, 20), (3, 30), (4, 40), (5, 50); + +query ? +SELECT collect_list(val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) +FROM t_int; +---- +[10] +[10, 20] +[20, 30] +[30, 40] +[40, 50] + +####### +# collect_set as a WINDOW function. +# array_sort wraps the result because the underlying HashMap iteration +# order is not deterministic. +####### + +statement ok +CREATE TABLE t_set(ts INT, val TEXT) AS VALUES + (1,'A'),(2,'A'),(3,'B'),(4,'C'),(5,'B'); + +# Sliding ROWS frame, 2 PRECEDING. +# Frame contents per row: +# [A] -> {A} +# [A,A] -> {A} +# [A,A,B] -> {A,B} +# [A,B,C] -> {A,B,C} +# [B,C,B] -> {B,C} +query ? +SELECT array_sort(collect_set(val) + OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND CURRENT ROW)) +FROM t_set; +---- +[A] +[A] +[A, B] +[A, B, C] +[B, C] + +# Narrower frame. +query ? +SELECT array_sort(collect_set(val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_set; +---- +[A] +[A] +[A, B] +[B, C] +[B, C] + +# Unbounded preceding — every distinct seen so far stays in. +query ? +SELECT array_sort(collect_set(val) + OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)) +FROM t_set; +---- +[A] +[A] +[A, B] +[A, B, C] +[A, B, C] + +# collect_set with NULLs — NULL never enters the set. +statement ok +CREATE TABLE t_set_nulls(ts INT, val TEXT) AS VALUES + (1,'A'),(2,NULL),(3,'A'),(4,NULL),(5,'B'); + +query ? +SELECT array_sort(collect_set(val) + OVER (ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_set_nulls; +---- +[A] +[A] +[A] +[A] +[B] + +# collect_set with PARTITION BY — partition isolation. +statement ok +CREATE TABLE t_set_parts(grp INT, ts INT, val TEXT) AS VALUES + (1, 1, 'A'), (1, 2, 'A'), (1, 3, 'B'), + (2, 1, 'B'), (2, 2, 'C'), (2, 3, 'C'); + +query I? +SELECT grp, array_sort(collect_set(val) + OVER (PARTITION BY grp ORDER BY ts + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW)) +FROM t_set_parts +ORDER BY grp, ts; +---- +1 [A] +1 [A] +1 [A, B] +2 [B] +2 [B, C] +2 [C] + +# Empty leading frame on the first row — Spark returns []. +query ? +SELECT array_sort(collect_set(val) + OVER (ORDER BY ts ROWS BETWEEN 2 PRECEDING AND 1 PRECEDING)) +FROM t_set; +---- +[] +[A] +[A] +[A, B] +[B, C] + +# All-NULL window — set is empty. +query ? +SELECT array_sort(collect_set(val) + OVER (ORDER BY ts ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW)) +FROM t_allnull; +---- +[] +[] +[] + +####### +# Cleanup +####### + +statement ok +DROP TABLE t; + +statement ok +DROP TABLE t_nulls; + +statement ok +DROP TABLE t_allnull; + +statement ok +DROP TABLE t_parts; + +statement ok +DROP TABLE t_range; + +statement ok +DROP TABLE t_groups; + +statement ok +DROP TABLE t_int; + +statement ok +DROP TABLE t_set; + +statement ok +DROP TABLE t_set_nulls; + +statement ok +DROP TABLE t_set_parts; diff --git a/datafusion/sqllogictest/test_files/spark/array/array_repeat.slt b/datafusion/sqllogictest/test_files/spark/array/array_repeat.slt index 923e349140976..d51767a264895 100644 --- a/datafusion/sqllogictest/test_files/spark/array/array_repeat.slt +++ b/datafusion/sqllogictest/test_files/spark/array/array_repeat.slt @@ -112,3 +112,10 @@ FROM VALUES [[123], [123]] [[], []] [[NULL], [NULL]] + + +# null count +query ? +select array_repeat('a', NULL); +---- +NULL diff --git a/datafusion/sqllogictest/test_files/spark/array/slice.slt b/datafusion/sqllogictest/test_files/spark/array/slice.slt index 6dfc1c0c6d0bf..f6fb431a0769b 100644 --- a/datafusion/sqllogictest/test_files/spark/array/slice.slt +++ b/datafusion/sqllogictest/test_files/spark/array/slice.slt @@ -137,3 +137,24 @@ query ? SELECT slice(slice(make_array(NULL), 1, 2), 1, 2) ---- [NULL] + +query ? +SELECT slice(make_array(1), -2, 2) +---- +[] + +query ? +SELECT slice(make_array(1, 2, 3, 4), -5, 2) +---- +[] + +query ? +SELECT slice(make_array(1), 3, 4) +---- +[] + +# the inner field name of the input list is preserved +query ?T +SELECT slice(array(1, 2, 3, 4), 2, 2), arrow_typeof(slice(array(1, 2, 3, 4), 2, 2)); +---- +[2, 3] List(Int64, field: 'element') diff --git a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt index 39dca512226b2..3ac5337cd7fd5 100644 --- a/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt +++ b/datafusion/sqllogictest/test_files/spark/bitmap/bitmap_count.slt @@ -68,6 +68,21 @@ SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int32, Binary)')) FROM (VALUES (X' 16 NULL +# The CAST to Dictionary below comes from the explicit arrow_cast. There must +# not be an additional outer CAST(... AS Binary) before bitmap_count. +query TT +EXPLAIN SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int32, Binary)')) +FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); +---- +logical_plan +01)Projection: bitmap_count(CAST(t.a AS Dictionary(Int32, Binary))) AS bitmap_count(arrow_cast(t.a,Utf8("Dictionary(Int32, Binary)"))) +02)--SubqueryAlias: t +03)----Projection: column1 AS a +04)------Values: (Binary("16,16")), (Binary("10,176")), (Binary("255,255")), (Binary(NULL)) +physical_plan +01)ProjectionExec: expr=[bitmap_count(CAST(column1@0 AS Dictionary(Int32, Binary))) as bitmap_count(arrow_cast(t.a,Utf8("Dictionary(Int32, Binary)")))] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + query I SELECT bitmap_count(arrow_cast(a, 'Dictionary(Int8, Binary)')) FROM (VALUES (X'1010'), (X'0AB0'), (X'FFFF'), (NULL)) AS t(a); ---- diff --git a/datafusion/sqllogictest/test_files/spark/collection/size.slt b/datafusion/sqllogictest/test_files/spark/collection/size.slt index 106760eebfe42..b9c445f4e6805 100644 --- a/datafusion/sqllogictest/test_files/spark/collection/size.slt +++ b/datafusion/sqllogictest/test_files/spark/collection/size.slt @@ -84,7 +84,7 @@ SELECT size(make_array(1, NULL, 3)); # NULL array returns -1 (Spark behavior) query I -SELECT size(NULL::int[]); +SELECT size(CAST(NULL AS ARRAY)); ---- -1 diff --git a/datafusion/sqllogictest/test_files/spark/datetime/monthname.slt b/datafusion/sqllogictest/test_files/spark/datetime/monthname.slt new file mode 100644 index 0000000000000..5927d79526a7b --- /dev/null +++ b/datafusion/sqllogictest/test_files/spark/datetime/monthname.slt @@ -0,0 +1,175 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# Scalar date input +query T +SELECT monthname('2024-03-15'::DATE); +---- +Mar + +# All 12 months +query T +SELECT monthname('2024-01-15'::DATE); +---- +Jan + +query T +SELECT monthname('2024-02-15'::DATE); +---- +Feb + +query T +SELECT monthname('2024-03-15'::DATE); +---- +Mar + +query T +SELECT monthname('2024-04-15'::DATE); +---- +Apr + +query T +SELECT monthname('2024-05-15'::DATE); +---- +May + +query T +SELECT monthname('2024-06-15'::DATE); +---- +Jun + +query T +SELECT monthname('2024-07-15'::DATE); +---- +Jul + +query T +SELECT monthname('2024-08-15'::DATE); +---- +Aug + +query T +SELECT monthname('2024-09-15'::DATE); +---- +Sep + +query T +SELECT monthname('2024-10-15'::DATE); +---- +Oct + +query T +SELECT monthname('2024-11-15'::DATE); +---- +Nov + +query T +SELECT monthname('2024-12-15'::DATE); +---- +Dec + +# NULL handling +query T +SELECT monthname(NULL::DATE); +---- +NULL + +# Array input +query T +SELECT monthname(d) FROM (VALUES ('2024-01-01'::DATE), ('2024-06-15'::DATE), ('2024-12-31'::DATE), (NULL::DATE)) AS t(d); +---- +Jan +Jun +Dec +NULL + +# Timestamp input: Spark coerces TIMESTAMP/TIMESTAMP_NTZ to DATE before evaluation +query T +SELECT monthname('2024-03-15 12:34:56'::TIMESTAMP); +---- +Mar + +query T +SELECT monthname('2024-07-04 00:00:00'::TIMESTAMP); +---- +Jul + +query T +SELECT monthname(NULL::TIMESTAMP); +---- +NULL + +# Timestamp array input +query T +SELECT monthname(ts) FROM (VALUES + ('2024-01-15 01:02:03'::TIMESTAMP), + ('2024-08-20 10:20:30'::TIMESTAMP), + ('2024-11-30 23:59:59'::TIMESTAMP), + (NULL::TIMESTAMP) +) AS t(ts); +---- +Jan +Aug +Nov +NULL + +# TIMESTAMP_NTZ (Timestamp without timezone) — explicit Microsecond precision +query T +SELECT monthname(arrow_cast('2024-04-10 09:15:00', 'Timestamp(Microsecond, None)')); +---- +Apr + +# TIMESTAMP_NTZ — explicit Millisecond precision +query T +SELECT monthname(arrow_cast('2024-09-05 18:45:30', 'Timestamp(Millisecond, None)')); +---- +Sep + +# TIMESTAMP_NTZ — explicit Second precision +query T +SELECT monthname(arrow_cast('2024-02-29 00:00:00', 'Timestamp(Second, None)')); +---- +Feb + +# TIMESTAMP_NTZ — NULL handling +query T +SELECT monthname(arrow_cast(NULL, 'Timestamp(Microsecond, None)')); +---- +NULL + +# TIMESTAMP with timezone (Spark TIMESTAMP / LTZ) — coerces to Date32 +query T +SELECT monthname(arrow_cast('2024-05-20 03:00:00', 'Timestamp(Nanosecond, Some("UTC"))')); +---- +May + +query T +SELECT monthname(arrow_cast('2024-10-31 23:59:59', 'Timestamp(Microsecond, Some("America/New_York"))')); +---- +Oct + +# Error: wrong argument type (string without cast) +statement error Function 'monthname' requires Date, but received String +SELECT monthname('not-a-date'); + +# Error: wrong argument type (integer) +statement error Function 'monthname' requires Date, but received Int64 +SELECT monthname(123); + +# Error: no arguments +statement error 'monthname' does not support zero arguments +SELECT monthname(); diff --git a/datafusion/sqllogictest/test_files/spark/datetime/next_day.slt b/datafusion/sqllogictest/test_files/spark/datetime/next_day.slt index 872d1f2b58eb6..74fc12e21e6d4 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/next_day.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/next_day.slt @@ -36,6 +36,12 @@ SELECT next_day('2015-07-27'::DATE, 'Sat'::string); ---- 2015-08-01 +# Whitespace-padded day names should be rejected (return NULL) per Spark behavior +query D +SELECT next_day('2015-01-14'::DATE, ' MO '::string); +---- +NULL + query error Failed to coerce arguments to satisfy a call to 'next_day' function SELECT next_day('2015-07-27'::DATE); @@ -79,3 +85,18 @@ FROM VALUES NULL NULL NULL + +# https://github.com/apache/datafusion/issues/23891 +# Far-future start dates whose next occurrence lands past chrono::NaiveDate::MAX +# (epoch day 95026236) must still return a value, matching Spark's integer +# arithmetic, rather than panicking. Cast the Date32 result to Int32 to assert +# the epoch day directly (these dates are past the printable range). +query I +SELECT arrow_cast(next_day(arrow_cast(95026236, 'Date32'), 'Mon'::string), 'Int32'); +---- +95026243 + +query I +SELECT arrow_cast(next_day(arrow_cast(95026230, 'Date32'), 'Tue'::string), 'Int32'); +---- +95026237 diff --git a/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt b/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt index b4f5444e8a2da..efa6b898c2a6a 100644 --- a/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt +++ b/datafusion/sqllogictest/test_files/spark/datetime/weekday.slt @@ -23,5 +23,109 @@ ## Original Query: SELECT weekday('2009-07-30'); ## PySpark 3.5.5 Result: {'weekday(2009-07-30)': 3, 'typeof(weekday(2009-07-30))': 'int', 'typeof(2009-07-30)': 'string'} -#query -#SELECT weekday('2009-07-30'::string); +# Spark `weekday` is 0-indexed with Monday = 0 .. Sunday = 6. +# 2009-07-30 is a Thursday -> 3. +query I +SELECT weekday('2009-07-30'::DATE); +---- +3 + +# All seven days of one week (2024-01-01 is a Monday). +query I +SELECT weekday('2024-01-01'::DATE); +---- +0 + +query I +SELECT weekday('2024-01-02'::DATE); +---- +1 + +query I +SELECT weekday('2024-01-03'::DATE); +---- +2 + +query I +SELECT weekday('2024-01-04'::DATE); +---- +3 + +query I +SELECT weekday('2024-01-05'::DATE); +---- +4 + +query I +SELECT weekday('2024-01-06'::DATE); +---- +5 + +query I +SELECT weekday('2024-01-07'::DATE); +---- +6 + +# NULL handling +query I +SELECT weekday(NULL::DATE); +---- +NULL + +# Array input (mix of weekdays and NULL) +query I +SELECT weekday(d) FROM (VALUES ('2024-01-01'::DATE), ('2024-01-06'::DATE), ('2024-01-07'::DATE), (NULL::DATE)) AS t(d); +---- +0 +5 +6 +NULL + +# Timestamp input: Spark coerces TIMESTAMP/TIMESTAMP_NTZ to DATE before evaluation +query I +SELECT weekday('2009-07-30 12:34:56'::TIMESTAMP); +---- +3 + +query I +SELECT weekday(NULL::TIMESTAMP); +---- +NULL + +# Timestamp array input +query I +SELECT weekday(ts) FROM (VALUES + ('2024-01-01 01:02:03'::TIMESTAMP), + ('2024-01-06 10:20:30'::TIMESTAMP), + ('2024-01-07 23:59:59'::TIMESTAMP), + (NULL::TIMESTAMP) +) AS t(ts); +---- +0 +5 +6 +NULL + +# TIMESTAMP_NTZ (Timestamp without timezone) — explicit Microsecond precision +query I +SELECT weekday(arrow_cast('2009-07-30 09:15:00', 'Timestamp(Microsecond, None)')); +---- +3 + +# TIMESTAMP with timezone (Spark TIMESTAMP / LTZ) — coerces to Date32 +query I +SELECT weekday(arrow_cast('2024-01-07 03:00:00', 'Timestamp(Nanosecond, Some("UTC"))')); +---- +6 + +# Error: wrong argument type (string without cast) +statement error Function 'weekday' requires Date, but received String +SELECT weekday('not-a-date'); + +# Error: wrong argument type (integer) +statement error Function 'weekday' requires Date, but received Int64 +SELECT weekday(123); + +# Error: no arguments +statement error 'weekday' does not support zero arguments +SELECT weekday(); diff --git a/datafusion/sqllogictest/test_files/spark/map/map_from_arrays.slt b/datafusion/sqllogictest/test_files/spark/map/map_from_arrays.slt index a26b0435c9291..7e501a31628e1 100644 --- a/datafusion/sqllogictest/test_files/spark/map/map_from_arrays.slt +++ b/datafusion/sqllogictest/test_files/spark/map/map_from_arrays.slt @@ -118,11 +118,25 @@ SELECT ---- {outer_key1: {inner_a: 1, inner_b: 2}, outer_key2: {inner_x: 10, inner_y: 20, inner_z: 30}} -# Test with duplicate keys -query ? +# Test with duplicate keys: raises DUPLICATED_MAP_KEY under Spark's default policy +query error DataFusion error: Execution error: \[DUPLICATED_MAP_KEY\] Duplicate map key true was found SELECT map_from_arrays(array(true, false, true), array('a', NULL, 'b')); ----- -{false: NULL, true: b} + +# Integer keys with a duplicate also raise DUPLICATED_MAP_KEY. +query error DataFusion error: Execution error: \[DUPLICATED_MAP_KEY\] Duplicate map key 1 was found +SELECT map_from_arrays(array(1, 2, 1), array('a', 'b', 'c')); + +# String keys with a duplicate also raise DUPLICATED_MAP_KEY. +query error DataFusion error: Execution error: \[DUPLICATED_MAP_KEY\] Duplicate map key k was found +SELECT map_from_arrays(array('k', 'k', 'k'), array(1, 2, 3)); + +# Multi-row: a clean row and a duplicate row still errors. +query error DataFusion error: Execution error: \[DUPLICATED_MAP_KEY\] Duplicate map key 1 was found +SELECT map_from_arrays(a, b) +FROM values + (array[1, 2], array['a', 'b']), + (array[1, 1], array['x', 'y']) +AS tab(a, b); # Tests with different list types query ? @@ -134,3 +148,40 @@ query ? SELECT map_from_arrays(arrow_cast(array('a', 'b', 'c'), 'FixedSizeList(3, Utf8)'), arrow_cast(array(1, 2, 3), 'LargeList(Int32)')); ---- {a: 1, b: 2, c: 3} + +# LAST_WIN policy: duplicates are allowed; later occurrences overwrite earlier ones. +statement ok +set datafusion.spark.map_key_dedup_policy = 'LAST_WIN'; + +query ? +SELECT map_from_arrays(array(1, 2, 1), array('a', 'b', 'c')); +---- +{1: c, 2: b} + +query ? +SELECT map_from_arrays(array('k', 'k', 'k'), array(1, 2, 3)); +---- +{k: 3} + +query ? +SELECT map_from_arrays(array(true, false, true), array('a', NULL, 'b')); +---- +{true: b, false: NULL} + +# Multi-row mix under LAST_WIN: clean, duplicate, empty and NULL rows all work. +query ? +SELECT map_from_arrays(a, b) +FROM values + (array[1, 2], array['a', 'b']), + (array[1, 1], array['x', 'y']), + (array[], array[]), + (NULL, NULL) +AS tab(a, b); +---- +{1: a, 2: b} +{1: y} +{} +NULL + +statement ok +set datafusion.spark.map_key_dedup_policy = 'EXCEPTION'; diff --git a/datafusion/sqllogictest/test_files/spark/map/map_from_entries.slt b/datafusion/sqllogictest/test_files/spark/map/map_from_entries.slt index 19b46886a027e..21f41f5ad976b 100644 --- a/datafusion/sqllogictest/test_files/spark/map/map_from_entries.slt +++ b/datafusion/sqllogictest/test_files/spark/map/map_from_entries.slt @@ -151,8 +151,8 @@ SELECT ---- {outer_key1: {inner_a: 1, inner_b: 2}, outer_key2: {inner_x: 10, inner_y: 20, inner_z: 30}} -# Test with duplicate keys -query ? +# Test with duplicate keys: raises DUPLICATED_MAP_KEY under Spark's default policy +query error DataFusion error: Execution error: \[DUPLICATED_MAP_KEY\] Duplicate map key true was found SELECT map_from_entries(array( struct(true, 'a'), struct(false, 'b'), @@ -160,5 +160,58 @@ SELECT map_from_entries(array( struct(false, cast(NULL as string)), struct(true, 'd') )); + +# Integer keys with a duplicate also raise DUPLICATED_MAP_KEY. +query error DataFusion error: Execution error: \[DUPLICATED_MAP_KEY\] Duplicate map key 1 was found +SELECT map_from_entries(array(struct(1, 'a'), struct(2, 'b'), struct(1, 'c'))); + +# String keys with triple occurrence also raise DUPLICATED_MAP_KEY. +query error DataFusion error: Execution error: \[DUPLICATED_MAP_KEY\] Duplicate map key k was found +SELECT map_from_entries(array(struct('k', 1), struct('k', 2), struct('k', 3))); + +# Multi-row: a clean row followed by a duplicate row still errors. +query error DataFusion error: Execution error: \[DUPLICATED_MAP_KEY\] Duplicate map key 1 was found +SELECT map_from_entries(data) +FROM values + (array[struct(1, 'a'), struct(2, 'b')]), + (array[struct(1, 'x'), struct(1, 'y')]) +AS tab(data); + +# LAST_WIN policy: duplicates are allowed; later occurrences overwrite earlier ones. +statement ok +set datafusion.spark.map_key_dedup_policy = 'LAST_WIN'; + +query ? +SELECT map_from_entries(array( + struct(true, 'a'), + struct(false, 'b'), + struct(true, 'c'), + struct(false, cast(NULL as string)), + struct(true, 'd') +)); ---- -{false: NULL, true: d} +{true: d, false: NULL} + +query ? +SELECT map_from_entries(array(struct(1, 'a'), struct(2, 'b'), struct(1, 'c'))); +---- +{1: c, 2: b} + +query ? +SELECT map_from_entries(array(struct('k', 1), struct('k', 2), struct('k', 3))); +---- +{k: 3} + +# Multi-row mix under LAST_WIN: clean row + duplicate row both succeed. +query ? +SELECT map_from_entries(data) +FROM values + (array[struct(1, 'a'), struct(2, 'b')]), + (array[struct(1, 'x'), struct(1, 'y')]) +AS tab(data); +---- +{1: a, 2: b} +{1: y} + +statement ok +set datafusion.spark.map_key_dedup_policy = 'EXCEPTION'; diff --git a/datafusion/sqllogictest/test_files/spark/map/str_to_map.slt b/datafusion/sqllogictest/test_files/spark/map/str_to_map.slt index 30d1672aef0ae..c1307468d7c6c 100644 --- a/datafusion/sqllogictest/test_files/spark/map/str_to_map.slt +++ b/datafusion/sqllogictest/test_files/spark/map/str_to_map.slt @@ -64,11 +64,25 @@ SELECT str_to_map('a=1&b=2&c=3', '&', '='); {a: 1, b: 2, c: 3} # Duplicate keys: EXCEPTION policy (Spark 3.0+ default) -# TODO: Add LAST_WIN policy tests when spark.sql.mapKeyDedupPolicy config is supported statement error Duplicate map key SELECT str_to_map('a:1,b:2,a:3'); +# Triple+ occurrences of the same key still raise DUPLICATED_MAP_KEY. +statement error +Duplicate map key 'a' +SELECT str_to_map('a:1,a:2,a:3'); + +# Duplicate where one occurrence is missing the kv_delim (value = NULL) still errors. +statement error +Duplicate map key 'a' +SELECT str_to_map('a,b:2,a:3'); + +# Multi-row input: a clean row followed by a duplicate row fails on the duplicate row. +statement error +Duplicate map key 'a' +SELECT str_to_map(col) FROM (VALUES ('a:1,b:2'), ('a:3,a:4')) AS t(col); + # Additional tests (DataFusion-specific) # NULL input returns NULL @@ -111,4 +125,43 @@ SELECT str_to_map(col1, col2, col3) FROM (VALUES ('a=1,b=2', ',', '='), ('x#9', ---- {a: 1, b: 2} {x: 9} -NULL \ No newline at end of file +NULL + +# LAST_WIN policy: duplicates are allowed; later occurrences overwrite earlier ones. +statement ok +set datafusion.spark.map_key_dedup_policy = 'LAST_WIN'; + +query ? +SELECT str_to_map('a:1,b:2,a:3'); +---- +{a: 3, b: 2} + +query ? +SELECT str_to_map('a:1,a:2,a:3'); +---- +{a: 3} + +# Missing kv_delim: the later occurrence overwrites the value at the key's +# first-seen position. +query ? +SELECT str_to_map('a:1,b:2,a'); +---- +{a: NULL, b: 2} + +# Multi-row: both clean and duplicate rows succeed under LAST_WIN. +query ? +SELECT str_to_map(col) FROM (VALUES ('a:1,b:2'), ('a:3,a:4')) AS t(col); +---- +{a: 1, b: 2} +{a: 4} + +statement ok +set datafusion.spark.map_key_dedup_policy = 'EXCEPTION'; + +# Invalid policy values are rejected at SET time with a clear message. +statement error +set datafusion.spark.map_key_dedup_policy = 'BOGUS'; +---- +DataFusion error: Error setting config datafusion.spark.map_key_dedup_policy +caused by +Invalid or Unsupported Configuration: Invalid MapKeyDedupPolicy: BOGUS. Expected one of: EXCEPTION, LAST_WIN diff --git a/datafusion/sqllogictest/test_files/spark/math/atan2.slt b/datafusion/sqllogictest/test_files/spark/math/atan2.slt index eb644854c402d..11e7a90202ddc 100644 --- a/datafusion/sqllogictest/test_files/spark/math/atan2.slt +++ b/datafusion/sqllogictest/test_files/spark/math/atan2.slt @@ -21,7 +21,151 @@ # For more information, please see: # https://github.com/apache/datafusion/issues/15914 -## Original Query: SELECT atan2(0, 0); -## PySpark 3.5.5 Result: {'ATAN2(0, 0)': 0.0, 'typeof(ATAN2(0, 0))': 'double', 'typeof(0)': 'int'} -#query -#SELECT atan2(0::int); +# standard angles in radians +query R +SELECT atan2(0, 0); +---- +0 + +query R +SELECT atan2(0, 1); +---- +0 + +# all four quadrants (atan2 is quadrant-aware via the signs of both arguments) +query R +SELECT atan2(1, 1); +---- +0.785398163397448 + +query R +SELECT atan2(1, -1); +---- +2.356194490192345 + +query R +SELECT atan2(-1, -1); +---- +-2.356194490192345 + +query R +SELECT atan2(-1, 1); +---- +-0.785398163397448 + +# on the axes +query R +SELECT atan2(1, 0); +---- +1.570796326794897 + +# negative x-axis: atan2 range extends to pi (atan only reaches +/- pi/2) +query R +SELECT atan2(0, -1); +---- +3.141592653589793 + +# NULL if either argument is NULL +query R +SELECT atan2(NULL::double, 1.0::double); +---- +NULL + +query R +SELECT atan2(1.0::double, NULL::double); +---- +NULL + +# NaN: any NaN input yields NaN (for atan2, NaN wins even over Infinity) +query R +SELECT atan2('NaN'::double, 1.0::double); +---- +NaN + +query R +SELECT atan2(1.0::double, 'NaN'::double); +---- +NaN + +query R +SELECT atan2('NaN'::double, 'NaN'::double); +---- +NaN + +query R +SELECT atan2('NaN'::double, 'Infinity'::double); +---- +NaN + +# NULL beats every special value (validity is checked before the value) +query R +SELECT atan2(NULL::double, 'Infinity'::double); +---- +NULL + +# both infinite: quadrant set by the signs (+/- pi/4, +/- 3pi/4) +query R +SELECT atan2('Infinity'::double, 'Infinity'::double); +---- +0.785398163397448 + +query R +SELECT atan2('-Infinity'::double, 'Infinity'::double); +---- +-0.785398163397448 + +query R +SELECT atan2('Infinity'::double, '-Infinity'::double); +---- +2.356194490192345 + +query R +SELECT atan2('-Infinity'::double, '-Infinity'::double); +---- +-2.356194490192345 + +# one infinite argument +query R +SELECT atan2('Infinity'::double, 1.0::double); +---- +1.570796326794897 + +query R +SELECT atan2('-Infinity'::double, 1.0::double); +---- +-1.570796326794897 + +query R +SELECT atan2(1.0::double, 'Infinity'::double); +---- +0 + +query R +SELECT atan2(1.0::double, '-Infinity'::double); +---- +3.141592653589793 + +query R +SELECT atan2(-1.0::double, '-Infinity'::double); +---- +-3.141592653589793 + +# signed zeros: -0 flips the sign on the negative x-axis (atan2(+0, -1) = pi above; atan2(-0, -1) = -pi) +query R +SELECT atan2(-0.0::double, -1.0::double); +---- +-3.141592653589793 + +# -0 in the first argument still returns 0 on the positive x-axis +query R +SELECT atan2(-0.0::double, 1.0::double); +---- +0 + +# array path, including a NULL row +query R +SELECT atan2(a, b) FROM (VALUES (0.0::double, 1.0::double), (1.0::double, 1.0::double), (NULL::double, 1.0::double)) AS t(a, b); +---- +0 +0.785398163397448 +NULL diff --git a/datafusion/sqllogictest/test_files/spark/math/hypot.slt b/datafusion/sqllogictest/test_files/spark/math/hypot.slt index 1349be0a95ee7..564b34add8b9f 100644 --- a/datafusion/sqllogictest/test_files/spark/math/hypot.slt +++ b/datafusion/sqllogictest/test_files/spark/math/hypot.slt @@ -21,7 +21,115 @@ # For more information, please see: # https://github.com/apache/datafusion/issues/15914 -## Original Query: SELECT hypot(3, 4); -## PySpark 3.5.5 Result: {'HYPOT(3, 4)': 5.0, 'typeof(HYPOT(3, 4))': 'double', 'typeof(3)': 'int', 'typeof(4)': 'int'} -#query -#SELECT hypot(3::int, 4::int); +# Scalar: classic Pythagorean triples (3-4-5, 5-12-13) +query R +SELECT hypot(3, 4); +---- +5 + +query R +SELECT hypot(5, 12); +---- +13 + +# Double inputs +query R +SELECT hypot(3.0::double, 4.0::double); +---- +5 + +# NULL if either argument is NULL +query R +SELECT hypot(NULL::double, 4.0::double); +---- +NULL + +query R +SELECT hypot(3.0::double, NULL::double); +---- +NULL + +# Array path, including a NULL row +query R +SELECT hypot(a, b) FROM (VALUES (3.0::double, 4.0::double), (6.0::double, 8.0::double), (NULL::double, 1.0::double)) AS t(a, b); +---- +5 +10 +NULL + +# Overflow-safe: naive sqrt(a*a + b*b) overflows to Infinity here; hypot stays finite (matches Spark's Math.hypot) +query B +SELECT hypot(3e200::double, 4e200::double) < 'Infinity'::double; +---- +true + +# any infinite input yields +Infinity, even when the other is NaN +query R +SELECT hypot('Infinity'::double, 4.0::double); +---- +Infinity + +query R +SELECT hypot(4.0::double, '-Infinity'::double); +---- +Infinity + +query R +SELECT hypot('Infinity'::double, 'NaN'::double); +---- +Infinity + +# NaN propagates when neither input is infinite +query R +SELECT hypot('NaN'::double, 4.0::double); +---- +NaN + +# signed zeros +query RRR +SELECT hypot(0.0::double, 0.0::double), hypot(-0.0::double, 0.0::double), hypot(3.0::double, -0.0::double); +---- +0 0 3 + +# NULL propagates even when the other input is Infinity +query R +SELECT hypot(NULL::double, 'Infinity'::double); +---- +NULL + +# negative inputs yield the positive magnitude +query RR +SELECT hypot(-3.0::double, -4.0::double), hypot(-3.0::double, 4.0::double); +---- +5 5 + +# Underflow-safe: naive sqrt(a*a + b*b) underflows to 0 for tiny inputs; hypot stays nonzero (matches Spark's Math.hypot) +query B +SELECT hypot(3e-200::double, 4e-200::double) > 0; +---- +true + +# Array path with special values (normal, +Infinity, NaN, NULL) +query R +SELECT hypot(a, b) FROM (VALUES + (3.0::double, 4.0::double), + ('Infinity'::double, 1.0::double), + ('NaN'::double, 1.0::double), + (NULL::double, 1.0::double)) AS t(a, b); +---- +5 +Infinity +NaN +NULL + +# both inputs NaN -> NaN +query R +SELECT hypot('NaN'::double, 'NaN'::double); +---- +NaN + +# both inputs infinite -> +Infinity +query R +SELECT hypot('Infinity'::double, '-Infinity'::double); +---- +Infinity \ No newline at end of file diff --git a/datafusion/sqllogictest/test_files/spark/math/pow.slt b/datafusion/sqllogictest/test_files/spark/math/pow.slt index 55b6f65b81235..17c3cfa18b0d3 100644 --- a/datafusion/sqllogictest/test_files/spark/math/pow.slt +++ b/datafusion/sqllogictest/test_files/spark/math/pow.slt @@ -22,6 +22,154 @@ # https://github.com/apache/datafusion/issues/15914 ## Original Query: SELECT pow(2, 3); -## PySpark 3.5.5 Result: {'pow(2, 3)': 8.0, 'typeof(pow(2, 3))': 'double', 'typeof(2)': 'int', 'typeof(3)': 'int'} -#query -#SELECT pow(2::int, 3::int); +## PySpark 3.5.5 Result: {'pow(2, 3)': 8.0, 'typeof(pow(2, 3))': 'double'} +## DataFusion: pow(int, int) returns int. Sqllogictest prints 8. +query R +SELECT pow(2::int, 3::int); +---- +8 + +## Spark returns Infinity for pow(0, negative) — see https://github.com/apache/datafusion/issues/22598 +## PostgreSQL / DataFusion default raises an error instead. +## PySpark 3.5.5: spark.sql("select pow(0, -1)").show() => Infinity + +query R +SELECT pow(0::double, -1::double); +---- +Infinity + +query R +SELECT power(0::double, -1::double); +---- +Infinity + +query R +SELECT pow(0.0, -1.0); +---- +Infinity + +# nulls +query R +SELECT pow(CAST(NULL AS DOUBLE), 1.0); +---- +NULL + +query R +SELECT pow(1.0, CAST(NULL AS DOUBLE)); +---- +NULL + +query R +SELECT pow(CAST(NULL AS DOUBLE), CAST(NULL AS DOUBLE)); +---- +NULL + +# nans +query R +SELECT pow(CAST('NaN' AS DOUBLE), 1.0); +---- +NaN + +query R +SELECT pow(1.0, CAST('NaN' AS DOUBLE)); +---- +1 + +query R +SELECT pow(CAST('NaN' AS DOUBLE), 0.0); +---- +1 + +# -0, +0 +query R +SELECT pow(0.0, 1.0); +---- +0 + +query R +SELECT pow(CAST('-0.0' AS DOUBLE), 1.0); +---- +0 + +query R +SELECT pow(0.0, -1.0); +---- +Infinity + +query R +SELECT pow(CAST('-0.0' AS DOUBLE), -1.0); +---- +Infinity + +# -inf, +inf +query R +SELECT pow(CAST('Infinity' AS DOUBLE), 1.0); +---- +Infinity + +query R +SELECT pow(CAST('Infinity' AS DOUBLE), -1.0); +---- +0 + +query R +SELECT pow(CAST('-Infinity' AS DOUBLE), 1.0); +---- +-Infinity + +query R +SELECT pow(CAST('-Infinity' AS DOUBLE), 2.0); +---- +Infinity + +query R +SELECT pow(2.0, CAST('Infinity' AS DOUBLE)); +---- +Infinity + +query R +SELECT pow(0.5, CAST('Infinity' AS DOUBLE)); +---- +0 + +query R +SELECT pow(2.0, CAST('-Infinity' AS DOUBLE)); +---- +0 + +query R +SELECT pow(0.5, CAST('-Infinity' AS DOUBLE)); +---- +Infinity + +# Test Array x Array +statement ok +CREATE TABLE t1(a DOUBLE, b DOUBLE) AS VALUES +(0.0, -1.0), +(2.0, 3.0), +(CAST(NULL AS DOUBLE), 1.0); + +query R +SELECT pow(a, b) FROM t1; +---- +Infinity +8 +NULL + +statement ok +DROP TABLE t1; + +# Test Scalar x Array +statement ok +CREATE TABLE t2(b DOUBLE) AS VALUES +(-1.0), +(2.0); + +query R +SELECT pow(0.0, b) FROM t2; +---- +Infinity +0 + +statement ok +DROP TABLE t2; diff --git a/datafusion/sqllogictest/test_files/spark/math/round.slt b/datafusion/sqllogictest/test_files/spark/math/round.slt index 91c5bdf0506f5..7fee15079d1d0 100644 --- a/datafusion/sqllogictest/test_files/spark/math/round.slt +++ b/datafusion/sqllogictest/test_files/spark/math/round.slt @@ -222,6 +222,18 @@ SELECT round(25::bigint, -1::int); ---- 30 +# round(bigint) should preserve exact values above Float64's exact integer range +query IT +SELECT round(arrow_cast(9007199254740993, 'Int64')), arrow_typeof(round(arrow_cast(9007199254740993, 'Int64'))); +---- +9007199254740993 Int64 + +# round(bigint, positive scale) should also preserve exact values above Float64's exact integer range +query IT +SELECT round(arrow_cast(9007199254740993, 'Int64'), 2::int), arrow_typeof(round(arrow_cast(9007199254740993, 'Int64'), 2::int)); +---- +9007199254740993 Int64 + # round(smallint, -1) query I SELECT round(25::smallint, -1::int); @@ -268,6 +280,18 @@ SELECT round(arrow_cast(25, 'UInt64'), -1::int); ---- 30 +# round(uint64) should preserve exact values above Float64's exact integer range +query IT +SELECT round(arrow_cast(18446744073709551615, 'UInt64')), arrow_typeof(round(arrow_cast(18446744073709551615, 'UInt64'))); +---- +18446744073709551615 UInt64 + +# round(uint64, positive scale) should also preserve exact values above Float64's exact integer range +query IT +SELECT round(arrow_cast(18446744073709551615, 'UInt64'), 2::int), arrow_typeof(round(arrow_cast(18446744073709551615, 'UInt64'), 2::int)); +---- +18446744073709551615 UInt64 + # round(uint32, positive scale) — no-op for integers query I SELECT round(arrow_cast(42, 'UInt32'), 2::int); @@ -281,52 +305,52 @@ SELECT round(arrow_cast(42, 'UInt32'), 2::int); # --- Decimal32 --- # round(decimal32, 0) — round to integer -query ? +query R SELECT round(arrow_cast(2.5, 'Decimal32(9, 1)'), 0::int); ---- -3.0 +3 -query ? +query R SELECT round(arrow_cast(-2.5, 'Decimal32(9, 1)'), 0::int); ---- --3.0 +-3 # round(decimal32, 2) -query ? +query R SELECT round(arrow_cast(2.345, 'Decimal32(9, 3)'), 2::int); ---- -2.350 +2.35 # round(decimal32) default scale = 0 -query ? +query R SELECT round(arrow_cast(3.5, 'Decimal32(9, 1)')); ---- -4.0 +4 # --- Decimal64 --- # round(decimal64, 0) — round to integer -query ? +query R SELECT round(arrow_cast(2.5, 'Decimal64(18, 1)'), 0::int); ---- -3.0 +3 -query ? +query R SELECT round(arrow_cast(-2.5, 'Decimal64(18, 1)'), 0::int); ---- --3.0 +-3 # round(decimal64, 2) -query ? +query R SELECT round(arrow_cast(2.345, 'Decimal64(18, 3)'), 2::int); ---- -2.350 +2.35 # round(decimal64) default scale = 0 -query ? +query R SELECT round(arrow_cast(3.5, 'Decimal64(18, 1)')); ---- -4.0 +4 # --- Decimal128 --- diff --git a/datafusion/sqllogictest/test_files/spark/string/concat.slt b/datafusion/sqllogictest/test_files/spark/string/concat.slt index df539a1c7a159..bd61ec29385ef 100644 --- a/datafusion/sqllogictest/test_files/spark/string/concat.slt +++ b/datafusion/sqllogictest/test_files/spark/string/concat.slt @@ -26,6 +26,7 @@ SELECT concat(arrow_cast('Spark', 'Utf8View'), arrow_cast('SQL', 'Utf8View')), a ---- SparkSQL Utf8View +# A major difference from the generic `concat` query T SELECT concat('Spark', 'SQL', NULL); ---- @@ -83,55 +84,14 @@ SELECT concat(arrow_cast('hello', 'Utf8View'), arrow_cast(' world', 'Binary')), ---- hello world Utf8View -# Test mixed types: Binary + Binary -query TT +# Test Binary + Binary +query ?T SELECT concat(arrow_cast('hello', 'Binary'), arrow_cast(' world', 'Binary')), arrow_typeof(concat(arrow_cast('hello', 'Binary'), arrow_cast(' world', 'Binary'))); ---- -hello world Utf8 - -# Test mixed types with ws: Binary + Binary -query TT -SELECT concat_ws('|', arrow_cast('hello', 'Binary'), arrow_cast('world', 'Binary')), arrow_typeof(concat_ws('|', arrow_cast('hello', 'Binary'), arrow_cast('world', 'Binary'))); ----- -hello|world Utf8 - -# Invalid UTF8 binaries for concatenation, scalar case -# 636166c3a9 = café , where c3a9 is a char é -# 68656c6c6f = hello -query error Execution error: invalid UTF-8 in binary literal -SELECT concat(x'636166c3', x'68656c6c6f'); - -query error Execution error: invalid UTF-8 in binary literal -SELECT concat(x'636166c3', arrow_cast(x'68656c6c6f', 'Utf8View')); - -statement ok -create table t as values (x'636166c3', x'68656c6c6f'); - -# Invalid UTF8 sequence for concatenation, array case -query error Arrow error: Invalid argument error: Invalid UTF8 sequence at string -SELECT concat(column1, column2) from t; +68656c6c6f20776f726c64 Binary -# Invalid UTF8 sequence for concatenation, array case -query error DataFusion error: Execution error: invalid UTF-8 in binary literal -SELECT concat(column1, arrow_cast(column2, 'Utf8View')) from t; - -statement ok -drop table t - -statement ok -create table t as values (x'636166c3', x'a968656c6c6f'); - -# Invalid UTF8 binaries make a valid UTF8 sequence after concatenation, array case -query T -SELECT concat(column1, column2) from t; ----- -caféhello - -statement ok -drop table t - -# Invalid UTF8 binaries make a valid UTF8 sequence after concatenation, scalar case -query T +# Test Binary + Binary, binary literals +query ? SELECT concat(x'636166c3', x'a968656c6c6f'); ---- -caféhello +636166c3a968656c6c6f diff --git a/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt b/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt index 62df636bba9ce..f6404cab8f3cd 100644 --- a/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt +++ b/datafusion/sqllogictest/test_files/spark/string/concat_ws.slt @@ -21,22 +21,367 @@ # For more information, please see: # https://github.com/apache/datafusion/issues/15914 -## Original Query: SELECT concat_ws(' ', 'Spark', 'SQL'); -## PySpark 3.5.5 Result: {'concat_ws( , Spark, SQL)': 'Spark SQL', 'typeof(concat_ws( , Spark, SQL))': 'string', 'typeof( )': 'string', 'typeof(Spark)': 'string', 'typeof(SQL)': 'string'} -#query -#SELECT concat_ws(' '::string, 'Spark'::string, 'SQL'::string); - -## Original Query: SELECT concat_ws('/', 'foo', null, 'bar'); -## PySpark 3.5.5 Result: {'concat_ws(/, foo, NULL, bar)': 'foo/bar', 'typeof(concat_ws(/, foo, NULL, bar))': 'string', 'typeof(/)': 'string', 'typeof(foo)': 'string', 'typeof(NULL)': 'void', 'typeof(bar)': 'string'} -#query -#SELECT concat_ws('/'::string, 'foo'::string, NULL::void, 'bar'::string); - -## Original Query: SELECT concat_ws('s'); -## PySpark 3.5.5 Result: {'concat_ws(s)': '', 'typeof(concat_ws(s))': 'string', 'typeof(s)': 'string'} -#query -#SELECT concat_ws('s'::string); - -## Original Query: SELECT concat_ws(null, 'Spark', 'SQL'); -## PySpark 3.5.5 Result: {'concat_ws(NULL, Spark, SQL)': None, 'typeof(concat_ws(NULL, Spark, SQL))': 'string', 'typeof(NULL)': 'void', 'typeof(Spark)': 'string', 'typeof(SQL)': 'string'} -#query -#SELECT concat_ws(NULL::void, 'Spark'::string, 'SQL'::string); +## ── Basic scalar usage ────────────────────────────────────── + +## Multiple string arguments +query T +SELECT concat_ws(',', 'a', 'b', 'c'); +---- +a,b,c + +## Space separator +query T +SELECT concat_ws(' ', 'Spark', 'SQL'); +---- +Spark SQL + +## Slash separator with null skipped +query T +SELECT concat_ws('/', 'foo', NULL, 'bar'); +---- +foo/bar + +## Single argument after separator +query T +SELECT concat_ws(',', 'a'); +---- +a + +## No arguments after separator → empty string +query T +SELECT concat_ws(','); +---- +(empty) + +## Null separator returns null +query T +SELECT concat_ws(NULL, 'a', 'b', 'c'); +---- +NULL + +## All null arguments → empty string +query T +SELECT concat_ws(',', CAST(NULL AS STRING), CAST(NULL AS STRING)); +---- +(empty) + +## ── Array arguments ───────────────────────────────────────── + +## Array argument +query T +SELECT concat_ws(',', array('a', 'b', 'c')); +---- +a,b,c + +## Array with nulls skipped +query T +SELECT concat_ws(',', array('a', NULL, 'c')); +---- +a,c + +## Multiple arrays +query T +SELECT concat_ws(',', array('a', 'b'), array('c', 'd')); +---- +a,b,c,d + +## Mixed scalar and array arguments +query T +SELECT concat_ws(',', 'x', array('a', 'b'), 'y'); +---- +x,a,b,y + +## Null array is skipped +query T +SELECT concat_ws(',', 'x', CAST(NULL AS ARRAY), 'y'); +---- +x,y + +## ── Edge cases ─────────────────────────────────────────────── + +## Separator column with no value arguments +query T +SELECT concat_ws(sep) AS result FROM VALUES (','), ('-') AS t(sep); +---- +(empty) +(empty) + +## Null separator in column with no value arguments +query T +SELECT concat_ws(sep) AS result FROM VALUES (CAST(NULL AS STRING)), (',') AS t(sep); +---- +NULL +(empty) + +## ── Column expressions ────────────────────────────────────── + +## concat_ws on columns +query T +SELECT concat_ws('-', a, b) AS result FROM VALUES ('hello', 'world'), ('foo', 'bar') AS t(a, b); +---- +hello-world +foo-bar + +## concat_ws with null in columns +query T +SELECT concat_ws(',', a, b) AS result FROM VALUES ('a', 'b'), ('c', CAST(NULL AS STRING)), (CAST(NULL AS STRING), 'd') AS t(a, b); +---- +a,b +c +d + +## Scalar-only arguments over multiple rows (broadcast test) +query T +SELECT concat_ws(',', 'a', 'b') AS result FROM VALUES (1), (2), (3) AS t(x); +---- +a,b +a,b +a,b + +## ── Additional edge cases ─────────────────────────────────── + +## Empty separator — values concatenated with nothing between +query T +SELECT concat_ws('', 'a', 'b', 'c'); +---- +abc + +## Empty-string values are NOT skipped (only NULLs are) +query T +SELECT concat_ws(',', '', 'a', '', 'b'); +---- +,a,,b + +## Multi-character separator +query T +SELECT concat_ws(' - ', 'a', 'b', 'c'); +---- +a - b - c + +## Utf8View separator +query TT +SELECT concat_ws(arrow_cast(',', 'Utf8View'), 'a', 'b'), arrow_typeof(concat_ws(arrow_cast(',', 'Utf8View'), 'a', 'b')); +---- +a,b Utf8 + +## LargeUtf8 separator +query TT +SELECT concat_ws(arrow_cast(',', 'LargeUtf8'), 'a', 'b'), arrow_typeof(concat_ws(arrow_cast(',', 'LargeUtf8'), 'a', 'b')); +---- +a,b Utf8 + +## Empty array → empty string +query T +SELECT concat_ws(',', array()); +---- +(empty) + +## Scalar + array + array mix +query T +SELECT concat_ws(',', array('a', 'b'), 'c', array('d', 'e')); +---- +a,b,c,d,e + +## All-NULL row mixed with non-NULL rows +query T +SELECT concat_ws(',', a, b, c) AS result FROM VALUES + ('a', 'b', 'c'), + (CAST(NULL AS STRING), 'b', 'c'), + ('a', CAST(NULL AS STRING), CAST(NULL AS STRING)), + (CAST(NULL AS STRING), CAST(NULL AS STRING), CAST(NULL AS STRING)) + AS t(a, b, c); +---- +a,b,c +b,c +a +(empty) + +## Separator from column (per-row separator), with NULL rows +query T +SELECT concat_ws(sep, a, b) AS result FROM VALUES + (',', 'a', 'b'), + (CAST(NULL AS STRING), 'a', 'b'), + ('|', 'x', 'y') + AS t(sep, a, b); +---- +a,b +NULL +x|y + +## ── Spark cross-checked extras ────────────────────────────── + +## Zero arguments → error (Spark: WRONG_NUM_ARGS) +query error +SELECT concat_ws(); + +## Numeric separator coerced to string +query T +SELECT concat_ws(1, 'a', 'b'); +---- +a1b + +## Only numeric separator, no values → empty string +query T +SELECT concat_ws(123); +---- +(empty) + +## Numeric values coerced to string +query T +SELECT concat_ws(',', 1, 2, 3); +---- +1,2,3 + +## Float values +query T +SELECT concat_ws(',', 1.5, 2.5); +---- +1.5,2.5 + +## Boolean values +query T +SELECT concat_ws(',', true, false); +---- +true,false + +## Mixed numeric and string +query T +SELECT concat_ws(',', CAST(1 AS BIGINT), 'a'); +---- +1,a + +## Date values +query T +SELECT concat_ws(',', DATE '2024-01-01', 'x'); +---- +2024-01-01,x + +## Multi-byte UTF-8 separator +query T +SELECT concat_ws('é', 'a', 'b'); +---- +aéb + +## Nested concat_ws +query T +SELECT concat_ws('|', concat_ws(',', 'a', 'b'), concat_ws(',', 'c', 'd')); +---- +a,b|c,d + +## All-NULL elements in array → empty string +query T +SELECT concat_ws(',', array(CAST(NULL AS STRING), CAST(NULL AS STRING), CAST(NULL AS STRING))); +---- +(empty) + +## Empty-string elements in arrays are NOT skipped +query T +SELECT concat_ws(',', array(''), array('')); +---- +, + +## Multiple arrays interleaved with scalars and NULLs +query T +SELECT concat_ws(',', array('a', 'b'), 'c', CAST(NULL AS STRING), array('d'), '', 'e'); +---- +a,b,c,d,,e + +## Long argument list (variadic) +query T +SELECT concat_ws('-', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'); +---- +a-b-c-d-e-f-g-h-i-j + +## Long array +query T +SELECT concat_ws(',', array('a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j')); +---- +a,b,c,d,e,f,g,h,i,j + +## All-empty arguments +query T +SELECT concat_ws('', '', '', ''); +---- +(empty) + +## Empty separator with all NULLs +query T +SELECT concat_ws('', CAST(NULL AS STRING), CAST(NULL AS STRING)); +---- +(empty) + +## Long string preserved (length sanity) +query I +SELECT length(concat_ws(',', repeat('x', 1000), repeat('y', 1000))); +---- +2001 + +## ── List variants (FixedSizeList / ListView) ──────────────── + +## FixedSizeList argument is expanded element-by-element +query T +SELECT concat_ws(',', arrow_cast(make_array('a', 'b', 'c'), 'FixedSizeList(3, Utf8)')); +---- +a,b,c + +## ListView argument is expanded element-by-element +query T +SELECT concat_ws(',', arrow_cast(make_array('a', 'b'), 'ListView(Utf8)')); +---- +a,b + +## LargeListView argument is expanded element-by-element +query T +SELECT concat_ws(',', arrow_cast(make_array('a', 'b'), 'LargeListView(Utf8)')); +---- +a,b + +## ── Binary coercion (Spark casts binary to its UTF-8 view) ── + +## Binary argument is coerced to its string representation +query T +SELECT concat_ws(',', X'4869'); +---- +Hi + +## ── Null-separator-over-rows shape ────────────────────────── + +## Null separator on a multi-row column yields NULL on every row +query T +SELECT concat_ws(NULL, v) AS result FROM VALUES ('a'), ('b'), ('c') AS t(v) ORDER BY v; +---- +NULL +NULL +NULL + +## ── Non-string list elements (planner-inserted element cast) ─ + +## Array of integers — elements must be cast to STRING +query T +SELECT concat_ws(',', array(1, 2, 3)); +---- +1,2,3 + +## Array of doubles +query T +SELECT concat_ws('-', array(1.5, 2.5, 3.5)); +---- +1.5-2.5-3.5 + +## Array of booleans +query T +SELECT concat_ws(',', array(true, false, true)); +---- +true,false,true + +## Mixed: string scalar + int array + string scalar +query T +SELECT concat_ws(',', 'x', array(1, 2), 'y'); +---- +x,1,2,y + +## ── Struct rejection ──────────────────────────────────────── + +## Struct argument is rejected (not coerced to string) +query error +SELECT concat_ws(',', named_struct('a', 1)); diff --git a/datafusion/sqllogictest/test_files/spark/string/elt.slt b/datafusion/sqllogictest/test_files/spark/string/elt.slt index 12917d17e1e47..9f0348324aadc 100644 --- a/datafusion/sqllogictest/test_files/spark/string/elt.slt +++ b/datafusion/sqllogictest/test_files/spark/string/elt.slt @@ -59,3 +59,143 @@ query T SELECT elt(1, 10, null) ---- 10 + +######################################## +# ANSI mode = false (default): invalid indices return NULL +######################################## + +# Index 0 -> NULL (Spark returns NULL when ANSI is off) +query T +SELECT elt(0::int, 'a', 'b'); +---- +NULL + +# Negative index -> NULL +query T +SELECT elt(-1::int, 'a', 'b'); +---- +NULL + +# Index far beyond the input list -> NULL +query T +SELECT elt(100::int, 'a', 'b', 'c'); +---- +NULL + +# NULL index -> NULL regardless of mode +query T +SELECT elt(NULL::int, 'a', 'b'); +---- +NULL + +# NULL value at the selected index -> NULL +query T +SELECT elt(2::int, 'a', NULL); +---- +NULL + +# Three-argument list, pick middle element +query T +SELECT elt(2::int, 'scala', 'java', 'python'); +---- +java + +# Three-argument list, pick last element +query T +SELECT elt(3::int, 'scala', 'java', 'python'); +---- +python + +# Mixed types get cast to string (Spark returns string) +query T +SELECT elt(2::int, 1, 2, 3); +---- +2 + +# Vectorized: mix of valid, out-of-range, and NULL indices in ANSI-off mode +statement ok +CREATE TABLE elt_rows(idx INT, a STRING, b STRING, c STRING) AS VALUES + (1, 'a1', 'b1', 'c1'), + (2, 'a2', 'b2', 'c2'), + (3, 'a3', 'b3', 'c3'), + (0, 'a4', 'b4', 'c4'), + (-1, 'a5', 'b5', 'c5'), + (4, 'a6', 'b6', 'c6'), + (NULL, 'a7', 'b7', 'c7'); + +query T +SELECT elt(idx, a, b, c) FROM elt_rows ORDER BY a; +---- +a1 +b2 +c3 +NULL +NULL +NULL +NULL + +statement ok +DROP TABLE elt_rows; + +######################################## +# ANSI mode = true: invalid indices raise ArrayIndexOutOfBoundsException +######################################## + +statement ok +set datafusion.execution.enable_ansi_mode = true; + +# Valid indices still work +query T +SELECT elt(1::int, 'scala', 'java'); +---- +scala + +query T +SELECT elt(2::int, 'scala', 'java'); +---- +java + +# NULL index still returns NULL (matches Spark: no error when index itself is NULL) +query T +SELECT elt(NULL::int, 'a', 'b'); +---- +NULL + +# NULL value at valid index still returns NULL (only invalid indices error) +query T +SELECT elt(1::int, NULL, 'b'); +---- +NULL + +# Out-of-range positive index errors +statement error DataFusion error: Execution error: The index 3 is out of bounds\. The array has 2 elements\. +SELECT elt(3::int, 'scala', 'java'); + +# Zero index errors +statement error DataFusion error: Execution error: The index 0 is out of bounds\. The array has 2 elements\. +SELECT elt(0::int, 'scala', 'java'); + +# Negative index errors +statement error DataFusion error: Execution error: The index -1 is out of bounds\. The array has 2 elements\. +SELECT elt(-1::int, 'scala', 'java'); + +# Large positive index errors +statement error DataFusion error: Execution error: The index 100 is out of bounds\. The array has 3 elements\. +SELECT elt(100::int, 'a', 'b', 'c'); + +# Vectorized: a batch that contains any invalid index errors in ANSI mode +statement ok +CREATE TABLE elt_ansi(idx INT, a STRING, b STRING) AS VALUES + (1, 'a1', 'b1'), + (2, 'a2', 'b2'), + (3, 'a3', 'b3'); + +statement error DataFusion error: Execution error: The index 3 is out of bounds\. The array has 2 elements\. +SELECT elt(idx, a, b) FROM elt_ansi; + +statement ok +DROP TABLE elt_ansi; + +# Reset ANSI mode +statement ok +set datafusion.execution.enable_ansi_mode = false; diff --git a/datafusion/sqllogictest/test_files/spark/string/quote.slt b/datafusion/sqllogictest/test_files/spark/string/quote.slt new file mode 100644 index 0000000000000..b5ef0f84e60d2 --- /dev/null +++ b/datafusion/sqllogictest/test_files/spark/string/quote.slt @@ -0,0 +1,161 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +query T +SELECT quote(arrow_cast(127, 'Int8')); +---- +'127' + +query T +SELECT quote(arrow_cast(-128, 'Int8')); +---- +'-128' + +query T +SELECT quote(arrow_cast(32767, 'Int16')); +---- +'32767' + +query T +SELECT quote(arrow_cast(-32768, 'Int16')); +---- +'-32768' + +query T +SELECT quote(arrow_cast(2147483647, 'Int32')); +---- +'2147483647' + +query T +SELECT quote(arrow_cast(-2147483648, 'Int32')); +---- +'-2147483648' + +query T +SELECT quote(arrow_cast(9223372036854775807, 'Int64')); +---- +'9223372036854775807' + +query T +SELECT quote(arrow_cast(-9223372036854775808, 'Int64')); +---- +'-9223372036854775808' + +query T +SELECT quote(arrow_cast(3.14, 'Float32')); +---- +'3.14' + +query T +SELECT quote(arrow_cast(2.718281828459045, 'Float64')); +---- +'2.718281828459045' + +query T +SELECT quote(arrow_cast(0, 'UInt8')); +---- +'0' + +query T +SELECT quote(arrow_cast(255, 'UInt8')); +---- +'255' + +query T +SELECT quote(arrow_cast(65535, 'UInt16')); +---- +'65535' + +query T +SELECT quote(arrow_cast(4294967295, 'UInt32')); +---- +'4294967295' + +query T +SELECT quote(arrow_cast(18446744073709551615, 'UInt64')); +---- +'18446744073709551615' + +query T +SELECT quote('special chars: !@#$%^&*()'); +---- +'special chars: !@#$%^&*()' + +query T +SELECT quote('tab\tseparated'); +---- +'tab\tseparated' + +query T +SELECT quote('carriage\rreturn'); +---- +'carriage\rreturn' + +query T +SELECT quote('backslash\\test'); +---- +'backslash\\test' + +query T +SELECT quote('quote\"inside\"'); +---- +'quote\"inside\"' + +query T +SELECT quote('mixed\nescape\tchars\r\n'); +---- +'mixed\nescape\tchars\r\n' + +query T +SELECT quote('unicode: 你好, 世界'); +---- +'unicode: 你好, 世界' + +query T +SELECT quote('emoji: 😀🎉❤️🚀'); +---- +'emoji: 😀🎉❤️🚀' + +query T +SELECT quote(arrow_cast('2024-01-15', 'Date32')); +---- +'2024-01-15' + +query T +SELECT quote(arrow_cast('2024-01-15T12:30:45', 'Timestamp(µs)')); +---- +'2024-01-15T12:30:45' + +query T +SELECT quote('special\n\t\r'); +---- +'special\n\t\r' + +query T +SELECT quote('a''b'); +---- +'a\'b' + +query T +SELECT quote('it''s a ''test'''); +---- +'it\'s a \'test\'' + +query T +SELECT quote(''''); +---- +'\'' diff --git a/datafusion/sqllogictest/test_files/statistics_registry.slt b/datafusion/sqllogictest/test_files/statistics_registry.slt index c856e779a0877..89258bec299c1 100644 --- a/datafusion/sqllogictest/test_files/statistics_registry.slt +++ b/datafusion/sqllogictest/test_files/statistics_registry.slt @@ -104,9 +104,9 @@ physical_plan 02)--RepartitionExec: partitioning=Hash([small_id@2], 4), input_partitions=1, maintains_sort_order=true 03)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@1)], projection=[region_id@1, order_id@2, small_id@4] 04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet -05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +05)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible 06)--RepartitionExec: partitioning=Hash([small_id@0], 4), input_partitions=1, maintains_sort_order=true -07)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +07)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/dim_small.parquet]]}, projection=[small_id, label], output_ordering=[small_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # -- With registry ----------------------------------------------------------- # Conservative estimate 100 > 50: dim_small correctly swapped to build side @@ -127,7 +127,7 @@ physical_plan 04)--RepartitionExec: partitioning=Hash([small_id@2], 4), input_partitions=1, maintains_sort_order=true 05)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(customer_id@0, customer_id@1)], projection=[region_id@1, order_id@2, small_id@4] 06)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/customers.parquet]]}, projection=[customer_id, region_id], output_ordering=[region_id@1 ASC NULLS LAST], file_type=parquet -07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ] +07)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/statistics_registry/orders.parquet]]}, projection=[order_id, customer_id, small_id], output_ordering=[order_id@0 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] AND DynamicFilter [ empty ], dynamic_rg_pruning=eligible # -- Verify results are identical regardless of join order -------------------- diff --git a/datafusion/sqllogictest/test_files/string/concat.slt b/datafusion/sqllogictest/test_files/string/concat.slt new file mode 100644 index 0000000000000..1749a57591bc6 --- /dev/null +++ b/datafusion/sqllogictest/test_files/string/concat.slt @@ -0,0 +1,138 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +# +# tests for concat and concat_ws +# + +# Test two Utf8View inputs: value and return type +query TT +SELECT concat(arrow_cast('Foo', 'Utf8View'), arrow_cast('Bar', 'Utf8View')), arrow_typeof(concat(arrow_cast('Foo', 'Utf8View'), arrow_cast('Bar', 'Utf8View'))); +---- +FooBar Utf8View + +query T +SELECT concat('Foo', 'Bar', NULL); +---- +FooBar + +query T +SELECT concat('', '1', '', '2'); +---- +12 + +query error does not support zero arguments +SELECT concat(); + +query T +SELECT concat(''); +---- +(empty) + +query T +SELECT concat(a, b, c) from (select 'a' a, 'b' b, 'c' c union all select null a, 'b', 'c') order by 1 nulls last; +---- +abc +bc + +# Test mixed types: Utf8View + Utf8 +query TT +SELECT concat(arrow_cast('hello', 'Utf8View'), ' world'), arrow_typeof(concat(arrow_cast('hello', 'Utf8View'), ' world')); +---- +hello world Utf8View + +# Test mixed string types +query TT +SELECT concat('a', arrow_cast('b', 'LargeUtf8')), arrow_typeof(concat('a', arrow_cast('b', 'LargeUtf8'))); +---- +ab LargeUtf8 + +# Test types mixed together +query TT +SELECT concat('a', arrow_cast('b', 'LargeUtf8'), arrow_cast('c', 'Utf8View')), arrow_typeof(concat('a', arrow_cast('b', 'LargeUtf8'), arrow_cast('c', 'Utf8View'))); +---- +abc Utf8View + +# Mixed Utf8 + Binary is allowed; binary is coerced to the widest string type +query TT +SELECT concat(arrow_cast('hello', 'Utf8'), arrow_cast(' world', 'Binary')), arrow_typeof(concat(arrow_cast('hello', 'Utf8'), arrow_cast(' world', 'Binary'))); +---- +hello world Utf8 + +# binary separator is allowed for string arguments +query TT +SELECT concat_ws(x'7c', 'hello', 'world'), arrow_typeof(concat_ws(x'7c', 'hello', 'world')); +---- +hello|world Utf8 + +# null separator +query T +SELECT concat_ws(NULL, 'hello', 'world'); +---- +NULL + +# Test Binary + Binary scalar concat +query ?T +SELECT concat(arrow_cast('hello', 'Binary'), arrow_cast(' world', 'Binary')), arrow_typeof(concat(arrow_cast('hello', 'Binary'), arrow_cast(' world', 'Binary'))); +---- +68656c6c6f20776f726c64 Binary + +# Test all binary types together: widened to BinaryView +query ?T +SELECT concat(arrow_cast('hello', 'Binary'), arrow_cast('there', 'BinaryView'), arrow_cast('world', 'LargeBinary')), arrow_typeof(concat(arrow_cast('hello', 'Binary'), arrow_cast('there', 'BinaryView'), arrow_cast('world', 'LargeBinary'))); +---- +68656c6c6f7468657265776f726c64 BinaryView + +# Test all binary types together with concat_ws: widened to BinaryView +query ?T +SELECT concat_ws(x'7c', arrow_cast('hello', 'Binary'), arrow_cast(' there', 'BinaryView'), arrow_cast(' world', 'LargeBinary')), arrow_typeof(concat_ws(x'7c', arrow_cast('hello', 'Binary'), arrow_cast(' there', 'BinaryView'), arrow_cast(' world', 'LargeBinary'))); +---- +68656c6c6f7c2074686572657c20776f726c64 BinaryView + +query TT +SELECT concat_ws('|', arrow_cast('hello', 'Utf8View'), 'world'), arrow_typeof(concat_ws('|', arrow_cast('hello', 'Utf8View'), ' world')); +---- +hello|world Utf8View + +query TT +SELECT concat_ws('|', arrow_cast('hello', 'Utf8View'), arrow_cast('there', 'LargeUtf8'), arrow_cast('world', 'Utf8')), arrow_typeof(concat_ws('|', arrow_cast('hello', 'Utf8View'), arrow_cast('there', 'LargeUtf8'), arrow_cast('world', 'Utf8'))); +---- +hello|there|world Utf8View + +# Test Binary + Binary scalar concat +query ?T +SELECT concat_ws(x'7c', arrow_cast('hello', 'Binary'), arrow_cast('world', 'Binary')), arrow_typeof(concat_ws(x'7c', arrow_cast('hello', 'Binary'), arrow_cast('world', 'Binary'))); +---- +68656c6c6f7c776f726c64 Binary + +statement ok +create table t as values (x'636166c3a9', x'68656c6c6f'); + +# Test binary + binary array concat +query ? +SELECT concat(column1, column2) from t; +---- +636166c3a968656c6c6f + +# Test binary + binary array concat_ws +query ? +SELECT concat_ws(x'7c', column1, column2) from t; +---- +636166c3a97c68656c6c6f + +statement ok +drop table t diff --git a/datafusion/sqllogictest/test_files/string/string_literal.slt b/datafusion/sqllogictest/test_files/string/string_literal.slt index 97f2a40c13fea..07aacaad9343b 100644 --- a/datafusion/sqllogictest/test_files/string/string_literal.slt +++ b/datafusion/sqllogictest/test_files/string/string_literal.slt @@ -341,6 +341,64 @@ SELECT lpad('x', 5, 'e' || chr(769)) = 'e' || chr(769) || 'e' || chr(769) || 'x' ---- true 5 +# lpad with string, length, and fill arrays in every string width +query BBB +SELECT + lpad(arrow_cast(column1, 'Utf8'), column2, arrow_cast(column3, 'Utf8')) IS NOT DISTINCT FROM column4, + lpad(arrow_cast(column1, 'LargeUtf8'), column2, arrow_cast(column3, 'LargeUtf8')) IS NOT DISTINCT FROM column4, + lpad(arrow_cast(column1, 'Utf8View'), column2, arrow_cast(column3, 'Utf8View')) IS NOT DISTINCT FROM column4 +FROM (VALUES + ('hi', 5, 'xy', 'xyxhi'), + ('abcdef', 3, 'z', 'abc'), + ('é', 4, '好', '好好好é'), + ('hi', 5, '', 'hi'), + (NULL, 5, 'x', NULL), + ('hi', NULL, 'x', NULL), + ('hi', 5, NULL, NULL) +) AS t(column1, column2, column3, column4); +---- +true true true +true true true +true true true +true true true +true true true +true true true +true true true + +# lpad array path with the default fill +query BBB +SELECT + lpad(arrow_cast(column1, 'Utf8'), column2) IS NOT DISTINCT FROM column3, + lpad(arrow_cast(column1, 'LargeUtf8'), column2) IS NOT DISTINCT FROM column3, + lpad(arrow_cast(column1, 'Utf8View'), column2) IS NOT DISTINCT FROM column3 +FROM (VALUES ('hi', 5, ' hi'), ('abcdef', 3, 'abc'), (NULL, 5, NULL)) AS t(column1, column2, column3); +---- +true true true +true true true +true true true + +# a large scalar target length skips the scalar fast path +query I +SELECT character_length(lpad('x', 16385, 'a')); +---- +16385 + +# invalid argument count/type and excessive target length +query error 'lpad' does not support zero arguments +SELECT lpad(); + +query error Failed to coerce arguments to satisfy a call to 'lpad' function: coercion from Utf8 to the signature +SELECT lpad('x'); + +query error Failed to coerce arguments to satisfy a call to 'lpad' function: coercion from Utf8, Int64, Utf8, Utf8 to the signature +SELECT lpad('x', 2, 'y', 'z'); + +query error Failed to coerce arguments to satisfy a call to 'lpad' function: coercion from Utf8, Utf8 to the signature +SELECT lpad('x', 'bad'); + +query error lpad requested length 2147483648 too large +SELECT lpad('x', 2147483648, 'y'); + query T SELECT regexp_replace('foobar', 'bar', 'xx', 'gi') ---- @@ -391,6 +449,10 @@ SELECT repeat(arrow_cast('foo', 'Dictionary(Int32, Utf8)'), 3) ---- foofoofoo +query error DataFusion error: Execution error: string size overflow on repeat, max size is 2147483647, but got \d+ +SELECT repeat(x, 9223372036854775807) +FROM (VALUES ('abc')) AS t(x); + query T SELECT arrow_typeof(repeat('foo', 3)) ---- @@ -426,6 +488,26 @@ SELECT replace(arrow_cast('foobar', 'LargeUtf8'), arrow_cast('bar', 'LargeUtf8') ---- foohello +# PostgreSQL compatibility: empty search string is a no-op (issue #22253) +query T +SELECT replace('abc', '', 'x') +---- +abc + +query T +SELECT replace(arrow_cast('abc', 'Dictionary(Int32, Utf8)'), '', 'x') +---- +abc + +query T +SELECT replace(arrow_cast('abc', 'Utf8View'), arrow_cast('', 'Utf8View'), arrow_cast('x', 'Utf8View')) +---- +abc + +query T +SELECT replace(arrow_cast('abc', 'LargeUtf8'), arrow_cast('', 'LargeUtf8'), arrow_cast('x', 'LargeUtf8')) +---- +abc query T SELECT reverse('abcde') @@ -462,6 +544,11 @@ SELECT reverse(arrow_cast('abcde', 'Dictionary(Int32, Utf8)')) ---- edcba +query T +SELECT arrow_typeof(reverse(arrow_cast('abcde', 'Dictionary(Int32, Utf8)'))) +---- +Dictionary(Int32, Utf8) + query T SELECT reverse('loẅks') ---- @@ -640,6 +727,64 @@ SELECT rpad('x', 5, 'e' || chr(769)) = 'x' || 'e' || chr(769) || 'e' || chr(769) ---- true 5 +# rpad with string, length, and fill arrays in every string width +query BBB +SELECT + rpad(arrow_cast(column1, 'Utf8'), column2, arrow_cast(column3, 'Utf8')) IS NOT DISTINCT FROM column4, + rpad(arrow_cast(column1, 'LargeUtf8'), column2, arrow_cast(column3, 'LargeUtf8')) IS NOT DISTINCT FROM column4, + rpad(arrow_cast(column1, 'Utf8View'), column2, arrow_cast(column3, 'Utf8View')) IS NOT DISTINCT FROM column4 +FROM (VALUES + ('hi', 5, 'xy', 'hixyx'), + ('abcdef', 3, 'z', 'abc'), + ('é', 4, '好', 'é好好好'), + ('hi', 5, '', 'hi'), + (NULL, 5, 'x', NULL), + ('hi', NULL, 'x', NULL), + ('hi', 5, NULL, NULL) +) AS t(column1, column2, column3, column4); +---- +true true true +true true true +true true true +true true true +true true true +true true true +true true true + +# rpad array path with the default fill +query BBB +SELECT + rpad(arrow_cast(column1, 'Utf8'), column2) IS NOT DISTINCT FROM column3, + rpad(arrow_cast(column1, 'LargeUtf8'), column2) IS NOT DISTINCT FROM column3, + rpad(arrow_cast(column1, 'Utf8View'), column2) IS NOT DISTINCT FROM column3 +FROM (VALUES ('hi', 5, 'hi '), ('abcdef', 3, 'abc'), (NULL, 5, NULL)) AS t(column1, column2, column3); +---- +true true true +true true true +true true true + +# a large scalar target length skips the scalar fast path +query I +SELECT character_length(rpad('x', 16385, 'a')); +---- +16385 + +# invalid argument count/type and excessive target length +query error 'rpad' does not support zero arguments +SELECT rpad(); + +query error Failed to coerce arguments to satisfy a call to 'rpad' function: coercion from Utf8 to the signature +SELECT rpad('x'); + +query error Failed to coerce arguments to satisfy a call to 'rpad' function: coercion from Utf8, Int64, Utf8, Utf8 to the signature +SELECT rpad('x', 2, 'y', 'z'); + +query error Failed to coerce arguments to satisfy a call to 'rpad' function: coercion from Utf8, Utf8 to the signature +SELECT rpad('x', 'bad'); + +query error rpad requested length 2147483648 too large +SELECT rpad('x', 2147483648, 'y'); + query I SELECT char_length('') ---- @@ -909,6 +1054,16 @@ SELECT find_in_set(arrow_cast('', 'Utf8View'), arrow_cast('a,b,c,d,a', 'Utf8View ---- 0 +# invalid scalar argument count and type +query error 'find_in_set' does not support zero arguments +SELECT find_in_set(); + +query error Failed to coerce arguments to satisfy a call to 'find_in_set' function +SELECT find_in_set('a'); + +query error Failed to coerce arguments to satisfy a call to 'find_in_set' function +SELECT find_in_set('a', 'a,b', 'extra'); + query T SELECT split_part('foo_bar', '_', 2) @@ -1855,7 +2010,7 @@ SELECT ---- 48 176 32 40 -query IIII +query ???? SELECT bit_length(arrow_cast('Andrew', 'Dictionary(Int32, Utf8)')), bit_length(arrow_cast('datafusion数据融合', 'Dictionary(Int32, Utf8)')), diff --git a/datafusion/sqllogictest/test_files/string/string_query.slt.part b/datafusion/sqllogictest/test_files/string/string_query.slt.part index 9e5b8f91e7d8e..dcddf06b557ed 100644 --- a/datafusion/sqllogictest/test_files/string/string_query.slt.part +++ b/datafusion/sqllogictest/test_files/string/string_query.slt.part @@ -645,10 +645,10 @@ drop table test_lowercase; query IIII SELECT - ASCII(ascii_1) as c1, - ASCII(ascii_2) as c2, - ASCII(unicode_1) as c3, - ASCII(unicode_2) as c4 + arrow_cast(ASCII(ascii_1), 'Int32') as c1, + arrow_cast(ASCII(ascii_2), 'Int32') as c2, + arrow_cast(ASCII(unicode_1), 'Int32') as c3, + arrow_cast(ASCII(unicode_2), 'Int32') as c4 FROM test_basic_operator; ---- 65 88 100 128293 @@ -884,10 +884,10 @@ Xiangpeng bar NULL bar NULL datafusion数据融合 Raphael baraphael NULL datafusionДатbarион NULL datafusionДатаФусион under_score under_score NULL un iść core NULL un iść core percent percent NULL pan Tadeusz ma iść w kąt NULL pan Tadeusz ma iść w kąt -(empty) (empty) NULL bar NULL (empty) -(empty) (empty) NULL bar NULL (empty) -% % NULL bar NULL (empty) -_ _ NULL bar NULL (empty) +(empty) (empty) NULL (empty) NULL (empty) +(empty) (empty) NULL (empty) NULL (empty) +% % NULL (empty) NULL (empty) +_ _ NULL (empty) NULL (empty) NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL NULL @@ -972,6 +972,82 @@ NULL NULL NULL NULL # Test FIND_IN_SET # -------------------------------------- +# array on the left and a literal on the right +query I +SELECT find_in_set(ascii_1, 'Andrew,Xiangpeng') FROM test_basic_operator +---- +1 +2 +0 +0 +0 +0 +0 +0 +0 +NULL +NULL + +# literal on the left and an array on the right +query I +SELECT find_in_set('🔥', unicode_2) FROM test_basic_operator +---- +1 +0 +0 +0 +0 +0 +0 +0 +0 +NULL +1 + +# arrays on both sides +query I +SELECT find_in_set(unicode_2, unicode_1) FROM test_basic_operator +---- +0 +1 +0 +0 +0 +1 +1 +1 +1 +NULL +NULL + +# Explicit casts are needed to exercise the LargeUtf8 scalar/array paths; +# otherwise string coercion chooses a different common physical type. +query II +SELECT + find_in_set(arrow_cast(ascii_1, 'LargeUtf8'), arrow_cast('Andrew,Xiangpeng', 'LargeUtf8')), + find_in_set(arrow_cast('🔥', 'LargeUtf8'), arrow_cast(unicode_2, 'LargeUtf8')) +FROM test_basic_operator +---- +1 1 +2 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +0 0 +NULL NULL +NULL 1 + +# null literals paired with arrays +query II +SELECT find_in_set(ascii_1, NULL), find_in_set(NULL, ascii_2) +FROM test_basic_operator +LIMIT 1 +---- +NULL NULL + query IIIIII SELECT FIND_IN_SET(ascii_1, 'a,b,c,d'), @@ -1253,8 +1329,8 @@ NULL NULL NULL NULL NULL NULL query II SELECT - CHARACTER_LENGTH(ascii_1), - CHARACTER_LENGTH(unicode_1) + arrow_cast(CHARACTER_LENGTH(ascii_1), 'Int64'), + arrow_cast(CHARACTER_LENGTH(unicode_1), 'Int64') FROM test_basic_operator ---- @@ -1275,7 +1351,12 @@ NULL NULL # -------------------------------------- query IIII -select bit_length(ascii_1), bit_length(ascii_2), bit_length(unicode_1), bit_length(unicode_2) from test_basic_operator; +select + arrow_cast(bit_length(ascii_1), 'Int64'), + arrow_cast(bit_length(ascii_2), 'Int64'), + arrow_cast(bit_length(unicode_1), 'Int64'), + arrow_cast(bit_length(unicode_2), 'Int64') +from test_basic_operator; ---- 48 8 144 32 72 72 176 176 @@ -1860,11 +1941,28 @@ SELECT left(ascii_1, 0), right(ascii_1, 0) FROM test_basic_operator NULL NULL NULL NULL -# left and right return Utf8View -query TT -SELECT arrow_typeof(left(ascii_1, 3)), arrow_typeof(right(ascii_1, 3)) FROM test_basic_operator LIMIT 1 +# left and right preserve the input string type +query TTTTTT +SELECT + arrow_typeof(left(arrow_cast(ascii_1, 'Utf8'), 3)), + arrow_typeof(right(arrow_cast(ascii_1, 'Utf8'), 3)), + arrow_typeof(left(arrow_cast(ascii_1, 'LargeUtf8'), 3)), + arrow_typeof(right(arrow_cast(ascii_1, 'LargeUtf8'), 3)), + arrow_typeof(left(arrow_cast(ascii_1, 'Utf8View'), 3)), + arrow_typeof(right(arrow_cast(ascii_1, 'Utf8View'), 3)) +FROM test_basic_operator LIMIT 1 +---- +Utf8 Utf8 LargeUtf8 LargeUtf8 Utf8View Utf8View + +# substr preserves the input string type +query TTT +SELECT + arrow_typeof(substr(arrow_cast(ascii_1, 'Utf8'), 1, 3)), + arrow_typeof(substr(arrow_cast(ascii_1, 'LargeUtf8'), 1, 3)), + arrow_typeof(substr(arrow_cast(ascii_1, 'Utf8View'), 1, 3)) +FROM test_basic_operator LIMIT 1 ---- -Utf8View Utf8View +Utf8 LargeUtf8 Utf8View # -------------------------------------- # Test repeat() against array inputs with various null patterns. The scalar diff --git a/datafusion/sqllogictest/test_files/struct.slt b/datafusion/sqllogictest/test_files/struct.slt index 5cf6e4817d475..a0e47a8691f34 100644 --- a/datafusion/sqllogictest/test_files/struct.slt +++ b/datafusion/sqllogictest/test_files/struct.slt @@ -126,6 +126,49 @@ physical_plan 01)ProjectionExec: expr=[struct(a@0, b@1, c@2) as struct(values.a,values.b,values.c)] 02)--DataSourceExec: partitions=1, partition_sizes=[1] +# get_field over an inline named_struct is resolved during logical +# simplification: the field access collapses to the underlying expression +# instead of materializing the intermediate struct. +query R +select get_field(named_struct('min', a, 'max', b), 'max') from values; +---- +1.1 +2.2 +3.3 + +query TT +explain select get_field(named_struct('min', a, 'max', b), 'max') from values; +---- +logical_plan +01)Projection: values.b AS named_struct(Utf8("min"),values.a,Utf8("max"),values.b)[max] +02)--TableScan: values projection=[b] +physical_plan +01)ProjectionExec: expr=[b@0 as named_struct(Utf8("min"),values.a,Utf8("max"),values.b)[max]] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# the same simplification applies to the positional struct() constructor, +# whose fields are named c0, c1, ... +query TT +explain select get_field(struct(a, b, c), 'c1') from values; +---- +logical_plan +01)Projection: values.b AS struct(values.a,values.b,values.c)[c1] +02)--TableScan: values projection=[b] +physical_plan +01)ProjectionExec: expr=[b@0 as struct(values.a,values.b,values.c)[c1]] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + +# nested constructors collapse all the way through +query TT +explain select named_struct('outer', named_struct('inner', a))['outer']['inner'] from values; +---- +logical_plan +01)Projection: values.a AS named_struct(Utf8("outer"),named_struct(Utf8("inner"),values.a))[outer][inner] +02)--TableScan: values projection=[a] +physical_plan +01)ProjectionExec: expr=[a@0 as named_struct(Utf8("outer"),named_struct(Utf8("inner"),values.a))[outer][inner]] +02)--DataSourceExec: partitions=1, partition_sizes=[1] + # error on 0 arguments query error select named_struct(); @@ -1671,3 +1714,22 @@ RESET datafusion.sql_parser.dialect; statement ok drop table t_agg_window; + +# extract_leaf_expressions regression +statement ok +create table leaf_base as select named_struct('status', 'active') as s, 1 as id; + +statement ok +create view leaf_view as select s, id, id + 1 as synth from leaf_base; + +query T?I +select s['status'], s, id from leaf_view where s['status'] is not null; +---- +active {status: active} 1 + +statement ok +drop view leaf_view; + +statement ok +drop table leaf_base; + diff --git a/datafusion/sqllogictest/test_files/subquery.slt b/datafusion/sqllogictest/test_files/subquery.slt index 25f124f217cbf..dcca13c4164c5 100644 --- a/datafusion/sqllogictest/test_files/subquery.slt +++ b/datafusion/sqllogictest/test_files/subquery.slt @@ -338,13 +338,13 @@ where c_acctbal < ( logical_plan 01)Sort: customer.c_custkey ASC NULLS LAST 02)--Projection: customer.c_custkey -03)----Inner Join: customer.c_custkey = __scalar_sq_1.o_custkey Filter: CAST(customer.c_acctbal AS Decimal128(25, 2)) < __scalar_sq_1.sum(orders.o_totalprice) +03)----LeftSemi Join: customer.c_custkey = __scalar_sq_1.o_custkey Filter: CAST(customer.c_acctbal AS Decimal128(25, 2)) < __scalar_sq_1.sum(orders.o_totalprice) 04)------TableScan: customer projection=[c_custkey, c_acctbal] 05)------SubqueryAlias: __scalar_sq_1 06)--------Projection: sum(orders.o_totalprice), orders.o_custkey 07)----------Aggregate: groupBy=[[orders.o_custkey]], aggr=[[sum(orders.o_totalprice)]] 08)------------Projection: orders.o_custkey, orders.o_totalprice -09)--------------Inner Join: orders.o_orderkey = __scalar_sq_2.l_orderkey Filter: CAST(orders.o_totalprice AS Decimal128(25, 2)) < __scalar_sq_2.price +09)--------------LeftSemi Join: orders.o_orderkey = __scalar_sq_2.l_orderkey Filter: CAST(orders.o_totalprice AS Decimal128(25, 2)) < __scalar_sq_2.price 10)----------------TableScan: orders projection=[o_orderkey, o_custkey, o_totalprice] 11)----------------SubqueryAlias: __scalar_sq_2 12)------------------Projection: sum(lineitem.l_extendedprice) AS price, lineitem.l_orderkey @@ -387,7 +387,7 @@ query TT explain SELECT t1_id, t1_name, t1_int FROM t1 WHERE EXISTS(SELECT t1_int FROM t1 WHERE t1.t1_id > t1.t1_int) ---- logical_plan -01)LeftSemi Join: +01)LeftSemi Join: 02)--TableScan: t1 projection=[t1_id, t1_name, t1_int] 03)--SubqueryAlias: __correlated_sq_1 04)----Projection: @@ -555,7 +555,7 @@ logical_plan 02)--TableScan: t0 projection=[t0_id, t0_name] 03)--SubqueryAlias: __correlated_sq_2 04)----Projection: t1.t1_name -05)------Inner Join: t1.t1_id = t2.t2_id +05)------LeftSemi Join: t1.t1_id = t2.t2_id 06)--------TableScan: t1 projection=[t1_id, t1_name] 07)--------TableScan: t2 projection=[t2_id] @@ -568,7 +568,7 @@ logical_plan 02)--TableScan: t0 projection=[t0_id, t0_name] 03)--SubqueryAlias: __correlated_sq_1 04)----Projection: t2.t2_name -05)------Inner Join: t1.t1_id = t2.t2_id +05)------RightSemi Join: t1.t1_id = t2.t2_id 06)--------TableScan: t1 projection=[t1_id] 07)--------SubqueryAlias: t2 08)----------TableScan: t2 projection=[t2_id, t2_name] @@ -606,7 +606,7 @@ query TT explain SELECT t1_id, t1_name FROM t1 WHERE EXISTS (SELECT NULL) ---- logical_plan -01)LeftSemi Join: +01)LeftSemi Join: 02)--TableScan: t1 projection=[t1_id, t1_name] 03)--SubqueryAlias: __correlated_sq_1 04)----EmptyRelation: rows=1 @@ -888,6 +888,68 @@ SELECT t1_id, (SELECT count(*) FROM t2 WHERE t2.t2_int = t1.t1_int) as cnt from 33 3 44 0 +#correlated_scalar_subquery_non_count_agg_empty_defaults +query III rowsort +SELECT + t1_id, + ( + SELECT regr_count(1.0, 1.0) + FROM t2 + WHERE t2.t2_int = t1.t1_int + ) AS r, + ( + SELECT approx_distinct(t2.t2_id) + FROM t2 + WHERE t2.t2_int = t1.t1_int + ) AS d +FROM t1 +---- +11 1 1 +22 0 0 +33 3 3 +44 0 0 + +query II rowsort +SELECT + t1_id, + ( + SELECT regr_count(1.0, 1.0) + approx_distinct(t2.t2_id) + FROM t2 + WHERE t2.t2_int = t1.t1_int + ) AS combined +FROM t1 +---- +11 2 +22 0 +33 6 +44 0 + +query I rowsort +SELECT t1_id +FROM t1 +WHERE + ( + SELECT approx_distinct(t2.t2_id) + FROM t2 + WHERE t2.t2_int = t1.t1_int + ) = 0 +---- +22 +44 + +query I rowsort +SELECT t1_id +FROM t1 +WHERE + ( + SELECT regr_count(1.0, 1.0) + FROM t2 + WHERE t2.t2_int = t1.t1_int + ) = 0 +---- +22 +44 + #correlated_scalar_subquery_count_agg_with_alias query TT explain SELECT t1_id, (SELECT count(*) as _cnt FROM t2 WHERE t2.t2_int = t1.t1_int) as cnt from t1 @@ -1255,6 +1317,149 @@ physical_plan 04)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 05)------DataSourceExec: partitions=1, partition_sizes=[2] +query TT +explain select t1_id from t1 +where t1_id > 40 or exists (select 1 from t2 where t2.t2_int = t1.t1_int) +---- +logical_plan +01)Projection: t1.t1_id +02)--Filter: t1.t1_id > Int32(40) OR __correlated_sq_1.mark +03)----Projection: t1.t1_id, __correlated_sq_1.mark +04)------LeftMark Join: t1.t1_int = __correlated_sq_1.t2_int +05)--------TableScan: t1 projection=[t1_id, t1_int] +06)--------SubqueryAlias: __correlated_sq_1 +07)----------TableScan: t2 projection=[t2_int] +physical_plan +01)FilterExec: t1_id@0 > 40 OR mark@1, projection=[t1_id@0] +02)--HashJoinExec: mode=CollectLeft, join_type=RightMark, on=[(t2_int@0, t1_int@1)], projection=[t1_id@0, mark@2] +03)----DataSourceExec: partitions=1, partition_sizes=[2] +04)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)------DataSourceExec: partitions=1, partition_sizes=[2] + +query I rowsort +select t1_id from t1 +where t1_id > 40 or exists (select 1 from t2 where t2.t2_int = t1.t1_int) +---- +11 +33 +44 + +query TT +explain select t1_id from t1 +where t1_id > 40 or not exists (select 1 from t2 where t2.t2_int > t1.t1_int) +---- +logical_plan +01)Projection: t1.t1_id +02)--Filter: t1.t1_id > Int32(40) OR NOT __correlated_sq_1.mark +03)----Projection: t1.t1_id, __correlated_sq_1.mark +04)------LeftMark Join: Filter: __correlated_sq_1.t2_int > t1.t1_int +05)--------TableScan: t1 projection=[t1_id, t1_int] +06)--------SubqueryAlias: __correlated_sq_1 +07)----------TableScan: t2 projection=[t2_int] +physical_plan +01)FilterExec: t1_id@0 > 40 OR NOT mark@1, projection=[t1_id@0] +02)--NestedLoopJoinExec: join_type=RightMark, filter=t2_int@1 > t1_int@0, projection=[t1_id@0, mark@2] +03)----DataSourceExec: partitions=1, partition_sizes=[2] +04)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)------DataSourceExec: partitions=1, partition_sizes=[2] + +query I rowsort +select t1_id from t1 +where t1_id > 40 or not exists (select 1 from t2 where t2.t2_int > t1.t1_int) +---- +33 +44 + +########## +# Regression for https://github.com/apache/datafusion/issues/23010: +# a projection that selects / reorders a subset of columns over a mark join. +# Schema-aware projection pushdown (driven by ColumnIndex / JoinSide) must keep +# the synthetic `mark` column (JoinSide::None) at the join output while pushing +# the child columns down. These lock the query results (which must stay stable +# across the refactor) and the current plan shape (the physical plan is expected +# to change once child pushdown is enabled for mark joins). Cover hash LeftMark, +# negated mark, and nested-loop mark. +########## + +query TT +EXPLAIN SELECT t1_name, t1_id FROM t1 +WHERE t1_id > 40 OR t1_id IN (SELECT t2_id FROM t2 WHERE t1_int > 0) +---- +logical_plan +01)Projection: t1.t1_name, t1.t1_id +02)--Filter: t1.t1_id > Int32(40) OR __correlated_sq_1.mark +03)----Projection: t1.t1_id, t1.t1_name, __correlated_sq_1.mark +04)------LeftMark Join: t1.t1_id = __correlated_sq_1.t2_id Filter: t1.t1_int > Int32(0) +05)--------TableScan: t1 projection=[t1_id, t1_name, t1_int] +06)--------SubqueryAlias: __correlated_sq_1 +07)----------TableScan: t2 projection=[t2_id] +physical_plan +01)FilterExec: t1_id@0 > 40 OR mark@2, projection=[t1_name@1, t1_id@0] +02)--HashJoinExec: mode=CollectLeft, join_type=RightMark, on=[(t2_id@0, t1_id@0)], filter=t1_int@0 > 0, projection=[t1_id@0, t1_name@1, mark@3] +03)----DataSourceExec: partitions=1, partition_sizes=[2] +04)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)------DataSourceExec: partitions=1, partition_sizes=[2] + +query TI rowsort +SELECT t1_name, t1_id FROM t1 +WHERE t1_id > 40 OR t1_id IN (SELECT t2_id FROM t2 WHERE t1_int > 0) +---- +a 11 +b 22 +d 44 + +query TT +EXPLAIN SELECT t1_int, t1_name FROM t1 +WHERE t1_id < 20 OR NOT EXISTS (SELECT 1 FROM t2 WHERE t1.t1_id = t2.t2_id) +---- +logical_plan +01)Projection: t1.t1_int, t1.t1_name +02)--Filter: t1.t1_id < Int32(20) OR NOT __correlated_sq_1.mark +03)----LeftMark Join: t1.t1_id = __correlated_sq_1.t2_id +04)------TableScan: t1 projection=[t1_id, t1_name, t1_int] +05)------SubqueryAlias: __correlated_sq_1 +06)--------TableScan: t2 projection=[t2_id] +physical_plan +01)FilterExec: t1_id@0 < 20 OR NOT mark@3, projection=[t1_int@2, t1_name@1] +02)--HashJoinExec: mode=CollectLeft, join_type=RightMark, on=[(t2_id@0, t1_id@0)] +03)----DataSourceExec: partitions=1, partition_sizes=[2] +04)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)------DataSourceExec: partitions=1, partition_sizes=[2] + +query IT rowsort +SELECT t1_int, t1_name FROM t1 +WHERE t1_id < 20 OR NOT EXISTS (SELECT 1 FROM t2 WHERE t1.t1_id = t2.t2_id) +---- +1 a +3 c + +query TT +EXPLAIN SELECT t1_name FROM t1 +WHERE t1_id > 40 OR EXISTS (SELECT 1 FROM t2 WHERE t1.t1_int > t2.t2_int) +---- +logical_plan +01)Projection: t1.t1_name +02)--Filter: t1.t1_id > Int32(40) OR __correlated_sq_1.mark +03)----Projection: t1.t1_id, t1.t1_name, __correlated_sq_1.mark +04)------LeftMark Join: Filter: t1.t1_int > __correlated_sq_1.t2_int +05)--------TableScan: t1 projection=[t1_id, t1_name, t1_int] +06)--------SubqueryAlias: __correlated_sq_1 +07)----------TableScan: t2 projection=[t2_int] +physical_plan +01)FilterExec: t1_id@0 > 40 OR mark@2, projection=[t1_name@1] +02)--NestedLoopJoinExec: join_type=RightMark, filter=t1_int@0 > t2_int@1, projection=[t1_id@0, t1_name@1, mark@3] +03)----DataSourceExec: partitions=1, partition_sizes=[2] +04)----RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +05)------DataSourceExec: partitions=1, partition_sizes=[2] + +query T rowsort +SELECT t1_name FROM t1 +WHERE t1_id > 40 OR EXISTS (SELECT 1 FROM t2 WHERE t1.t1_int > t2.t2_int) +---- +b +c +d + statement ok set datafusion.explain.logical_plan_only = true; @@ -1508,7 +1713,7 @@ query TT explain SELECT a FROM t1 WHERE EXISTS (SELECT count(*) FROM t2) ---- logical_plan -01)LeftSemi Join: +01)LeftSemi Join: 02)--TableScan: t1 projection=[a] 03)--SubqueryAlias: __correlated_sq_1 04)----EmptyRelation: rows=1 @@ -1525,7 +1730,7 @@ statement count 0 create table person(id int, last_name int, state int); query TT -explain SELECT id FROM person p WHERE EXISTS +explain SELECT id FROM person p WHERE EXISTS (SELECT * FROM person WHERE last_name = p.last_name AND state = p.state) ---- logical_plan @@ -1598,6 +1803,25 @@ logical_plan 21)----------Projection: column1 AS v 22)------------Values: (Int64(5)), (Int64(NULL)) +# same-table `= ANY` / `<> ALL` must plan without +# "duplicate unqualified field name mark". +statement ok +create table set_cmp_self(id int, age int) as values (1, 20), (2, 30), (3, 40); + +query I rowsort +select id from set_cmp_self where age = any(select age from set_cmp_self); +---- +1 +2 +3 + +query I +select id from set_cmp_self where age <> all(select age from set_cmp_self); +---- + +statement count 0 +drop table set_cmp_self; + # correlated_recursive_scalar_subquery_with_level_3_exists_subquery_referencing_level1_relation query TT explain select c_custkey from customer @@ -1613,7 +1837,7 @@ where c_acctbal < ( logical_plan 01)Sort: customer.c_custkey ASC NULLS LAST 02)--Projection: customer.c_custkey -03)----Inner Join: customer.c_custkey = __scalar_sq_2.o_custkey Filter: CAST(customer.c_acctbal AS Decimal128(25, 2)) < __scalar_sq_2.sum(orders.o_totalprice) +03)----LeftSemi Join: customer.c_custkey = __scalar_sq_2.o_custkey Filter: CAST(customer.c_acctbal AS Decimal128(25, 2)) < __scalar_sq_2.sum(orders.o_totalprice) 04)------TableScan: customer projection=[c_custkey, c_acctbal] 05)------SubqueryAlias: __scalar_sq_2 06)--------Projection: sum(orders.o_totalprice), orders.o_custkey @@ -1639,7 +1863,7 @@ where c_acctbal < ( logical_plan 01)Sort: customer.c_custkey ASC NULLS LAST 02)--Projection: customer.c_custkey -03)----Inner Join: customer.c_custkey = __scalar_sq_2.o_custkey Filter: CAST(customer.c_acctbal AS Decimal128(25, 2)) < __scalar_sq_2.sum(orders.o_totalprice) +03)----LeftSemi Join: customer.c_custkey = __scalar_sq_2.o_custkey Filter: CAST(customer.c_acctbal AS Decimal128(25, 2)) < __scalar_sq_2.sum(orders.o_totalprice) 04)------TableScan: customer projection=[c_custkey, c_acctbal] 05)------SubqueryAlias: __scalar_sq_2 06)--------Projection: sum(orders.o_totalprice), orders.o_custkey @@ -1684,7 +1908,7 @@ WHERE e1.salary > ( ---- logical_plan 01)Projection: e1.employee_name, e1.salary -02)--Inner Join: e1.dept_id = __scalar_sq_1.dept_id Filter: CAST(e1.salary AS Decimal128(38, 14)) > __scalar_sq_1.avg(e2.salary) +02)--LeftSemi Join: e1.dept_id = __scalar_sq_1.dept_id Filter: CAST(e1.salary AS Decimal128(38, 14)) > __scalar_sq_1.avg(e2.salary) 03)----SubqueryAlias: e1 04)------TableScan: employees projection=[employee_name, dept_id, salary] 05)----SubqueryAlias: __scalar_sq_1 @@ -2091,6 +2315,95 @@ SELECT (SELECT v FROM (SELECT 1 AS v UNION ALL SELECT 2) AS t ORDER BY v LIMIT 1 ---- 1 +############# +## End-to-end correctness coverage for the flag-off path. +## When `datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery` is false, +## uncorrelated scalar subqueries are rewritten to left joins by +## `ScalarSubqueryToJoin` instead of executed by `ScalarSubqueryExec`. This +## restores pre-PR-21240 behavior, which has two known shortcomings the +## physical-execution path was built to fix: multi-row subqueries silently +## return wrong results, and uncorrelated scalar subqueries do not work in +## ORDER BY / JOIN ON / aggregate-function arguments. Those cases are +## intentionally not covered here; the queries below are the ones where both +## paths agree. +############# + +statement ok +set datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery = false; + +# Scalar subquery returning exactly one row → success +query I +SELECT (SELECT v FROM sq_values LIMIT 1); +---- +1 + +# Scalar subquery returning exactly one row in WHERE → success +query I rowsort +SELECT x FROM sq_main WHERE x > (SELECT v FROM sq_values LIMIT 1); +---- +10 +20 + +# Scalar subquery returning zero rows → NULL +query I +SELECT (SELECT v FROM sq_empty); +---- +NULL + +# Scalar subquery returning zero rows in arithmetic → NULL propagation +query I +SELECT x + (SELECT v FROM sq_empty) FROM sq_main; +---- +NULL +NULL + +# Scalar subquery returning zero rows in WHERE comparison → no matching rows +query I +SELECT x FROM sq_main WHERE x > (SELECT v FROM sq_empty); +---- + +# Aggregated subquery always returns one row, even on empty input → success +query I +SELECT (SELECT count(*) FROM sq_empty); +---- +0 + +# Aggregated subquery on multi-row table → success +query I +SELECT (SELECT max(v) FROM sq_values); +---- +3 + +# HAVING clause with uncorrelated scalar subquery +query II rowsort +SELECT x, count(*) AS cnt FROM sq_main GROUP BY x +HAVING count(*) > (SELECT min(v) FROM sq_values); +---- + +# CASE WHEN with uncorrelated scalar subquery as condition +query T rowsort +SELECT CASE WHEN x > (SELECT min(v) FROM sq_values) + THEN 'big' ELSE 'small' END AS label +FROM sq_main; +---- +big +big + +# Doubly-nested constant subquery +query I +SELECT (SELECT (SELECT 42)); +---- +42 + +# NULL comparison semantics through subquery boundary +query B +SELECT 1 = (SELECT CAST(NULL AS INT)); +---- +NULL + +statement ok +RESET datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery; + statement count 0 DROP TABLE sq_values; @@ -2286,3 +2599,61 @@ DROP TABLE sq_count_customer; statement ok DROP TABLE sq_count_orders; + +# Regression test: `NOT IN` is a null-aware anti join. When the subquery yields a +# NULL the predicate is never TRUE, so the query must return zero rows. This must +# hold regardless of the chosen physical join operator. Previously, with +# prefer_hash_join = false and multiple partitions, the planner routed the +# null-aware anti join to SortMergeJoin (which is not null-aware) and returned +# wrong results; null-aware anti joins must use the CollectLeft HashJoin. + +statement ok +set datafusion.optimizer.prefer_hash_join = false; + +statement ok +CREATE TABLE nia_left(x INT) AS VALUES (1), (2), (3), (4); + +statement ok +CREATE TABLE nia_right_with_null(y INT) AS VALUES (2), (NULL); + +statement ok +CREATE TABLE nia_right_no_null(y INT) AS VALUES (2), (4); + +# Subquery contains a NULL -> NOT IN must return no rows. +query I +SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_with_null) ORDER BY x; +---- + +# The null-aware anti join must be planned as a CollectLeft HashJoinExec even with +# prefer_hash_join = false: SortMergeJoinExec is not null-aware and must not be used. +query TT +EXPLAIN SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_with_null); +---- +logical_plan +01)LeftAnti Join: nia_left.x = __correlated_sq_1.y null_aware +02)--TableScan: nia_left projection=[x] +03)--SubqueryAlias: __correlated_sq_1 +04)----TableScan: nia_right_with_null projection=[y] +physical_plan +01)HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(x@0, y@0)], null_aware +02)--DataSourceExec: partitions=1, partition_sizes=[1] +03)--DataSourceExec: partitions=1, partition_sizes=[1] + +# Subquery has no NULL -> NOT IN behaves like a normal anti join. +query I +SELECT x FROM nia_left WHERE x NOT IN (SELECT y FROM nia_right_no_null) ORDER BY x; +---- +1 +3 + +statement ok +DROP TABLE nia_left; + +statement ok +DROP TABLE nia_right_with_null; + +statement ok +DROP TABLE nia_right_no_null; + +statement ok +reset datafusion.optimizer.prefer_hash_join; diff --git a/datafusion/sqllogictest/test_files/subquery_sort.slt b/datafusion/sqllogictest/test_files/subquery_sort.slt index 6df93a3daabf6..080fd57c274d8 100644 --- a/datafusion/sqllogictest/test_files/subquery_sort.slt +++ b/datafusion/sqllogictest/test_files/subquery_sort.slt @@ -116,12 +116,11 @@ logical_plan 06)----------WindowAggr: windowExpr=[[rank() ORDER BY [sink_table.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] 07)------------TableScan: sink_table projection=[c1, c3, c9] physical_plan -01)ProjectionExec: expr=[c1@0 as c1, r@1 as r] -02)--SortExec: TopK(fetch=2), expr=[c1@0 ASC NULLS LAST, c3@2 ASC NULLS LAST, c9@3 ASC NULLS LAST], preserve_partitioning=[false] -03)----ProjectionExec: expr=[c1@0 as c1, rank() ORDER BY [sink_table.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as r, c3@1 as c3, c9@2 as c9] -04)------BoundedWindowAggExec: wdw=[rank() ORDER BY [sink_table.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() ORDER BY [sink_table.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -05)--------SortExec: expr=[c1@0 DESC], preserve_partitioning=[false] -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c3, c9], file_type=csv, has_header=true +01)ProjectionExec: expr=[c1@0 as c1, rank() ORDER BY [sink_table.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as r] +02)--SortExec: TopK(fetch=2), expr=[c1@0 ASC NULLS LAST, c3@1 ASC NULLS LAST, c9@2 ASC NULLS LAST], preserve_partitioning=[false] +03)----BoundedWindowAggExec: wdw=[rank() ORDER BY [sink_table.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() ORDER BY [sink_table.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------SortExec: expr=[c1@0 DESC], preserve_partitioning=[false] +05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c3, c9], file_type=csv, has_header=true #Test with utf8view for window function statement ok @@ -142,12 +141,11 @@ logical_plan 06)----------WindowAggr: windowExpr=[[rank() ORDER BY [sink_table_with_utf8view.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] 07)------------TableScan: sink_table_with_utf8view projection=[c1, c3, c9] physical_plan -01)ProjectionExec: expr=[c1@0 as c1, r@1 as r] -02)--SortExec: TopK(fetch=2), expr=[c1@0 ASC NULLS LAST, c3@2 ASC NULLS LAST, c9@3 ASC NULLS LAST], preserve_partitioning=[false] -03)----ProjectionExec: expr=[c1@0 as c1, rank() ORDER BY [sink_table_with_utf8view.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as r, c3@1 as c3, c9@2 as c9] -04)------BoundedWindowAggExec: wdw=[rank() ORDER BY [sink_table_with_utf8view.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() ORDER BY [sink_table_with_utf8view.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -05)--------SortExec: expr=[c1@0 DESC], preserve_partitioning=[false] -06)----------DataSourceExec: partitions=1, partition_sizes=[1] +01)ProjectionExec: expr=[c1@0 as c1, rank() ORDER BY [sink_table_with_utf8view.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as r] +02)--SortExec: TopK(fetch=2), expr=[c1@0 ASC NULLS LAST, c3@1 ASC NULLS LAST, c9@2 ASC NULLS LAST], preserve_partitioning=[false] +03)----BoundedWindowAggExec: wdw=[rank() ORDER BY [sink_table_with_utf8view.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() ORDER BY [sink_table_with_utf8view.c1 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------SortExec: expr=[c1@0 DESC], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] statement ok DROP TABLE sink_table_with_utf8view; diff --git a/datafusion/sqllogictest/test_files/table_functions.slt b/datafusion/sqllogictest/test_files/table_functions.slt index e1ab444d81044..e67d898d71475 100644 --- a/datafusion/sqllogictest/test_files/table_functions.slt +++ b/datafusion/sqllogictest/test_files/table_functions.slt @@ -197,6 +197,68 @@ SELECT * FROM generate_series(1, 2, 3, 4) statement error DataFusion error: Error during planning: Argument \#1 must be an INTEGER, TIMESTAMP, DATE or NULL, got Utf8 SELECT * FROM generate_series('foo', 'bar') +# Regression test for https://github.com/apache/datafusion/issues/22208 +# A step that would overflow i64 after the last reachable value must return the +# reachable values instead of panicking, matching PostgreSQL/DuckDB behavior. +query I +SELECT * FROM generate_series(9223372036854775806, 9223372036854775807, 2) +---- +9223372036854775806 + +# Same, in the descending direction +query I +SELECT * FROM generate_series(-9223372036854775806, -9223372036854775808, -2) +---- +-9223372036854775806 +-9223372036854775808 + +# Landing exactly on i64::MAX must include it +query I +SELECT * FROM generate_series(9223372036854775805, 9223372036854775807, 2) +---- +9223372036854775805 +9223372036854775807 + +# Same overflow behavior for `range` (end exclusive) +query I +SELECT * FROM range(9223372036854775806, 9223372036854775807, 2) +---- +9223372036854775806 + +# Regression test for https://github.com/apache/datafusion/issues/22193 +# Dates outside the nanosecond timestamp range must produce a clean planning +# error instead of panicking (debug) or silently wrapping (release). +statement error DataFusion error: Error during planning: First argument for generate_series is out of range of nanosecond timestamps +SELECT * FROM generate_series(DATE '0001-01-01', DATE '2000-01-01', INTERVAL '1' DAY) + +statement error DataFusion error: Error during planning: Second argument for generate_series is out of range of nanosecond timestamps +SELECT * FROM generate_series(DATE '2000-01-01', DATE '3000-01-01', INTERVAL '1' DAY) + +# Reaching the maximum representable date must not attempt to advance beyond it. +query P +SELECT * FROM generate_series(DATE '2262-04-11', DATE '2262-04-11', INTERVAL '1' DAY) +---- +2262-04-11T00:00:00 + +# Same for the maximum representable nanosecond timestamp. +query P +SELECT * FROM generate_series(TIMESTAMP '2262-04-11T23:47:16.854775807', TIMESTAMP '2262-04-11T23:47:16.854775807', INTERVAL '1' NANOSECOND) +---- +2262-04-11T23:47:16.854775807 + +# A timestamp step that exceeds the nanosecond range must terminate after the +# last reachable value instead of returning an overflow error. +query P +SELECT * FROM generate_series(TIMESTAMP '2262-04-11T23:47:16.854775806', TIMESTAMP '2262-04-11T23:47:16.854775807', INTERVAL '2' NANOSECOND) +---- +2262-04-11T23:47:16.854775806 + +# Same behavior for date series, which use the timestamp implementation. +query P +SELECT * FROM generate_series(DATE '2262-04-10', DATE '2262-04-11', INTERVAL '2' DAY) +---- +2262-04-10T00:00:00 + # UDF and UDTF `generate_series` can be used simultaneously query ? rowsort SELECT generate_series(1, t1.end) FROM generate_series(3, 5) as t1(end) diff --git a/datafusion/sqllogictest/test_files/topk.slt b/datafusion/sqllogictest/test_files/topk.slt index 8cab67dac0acb..180350a735b46 100644 --- a/datafusion/sqllogictest/test_files/topk.slt +++ b/datafusion/sqllogictest/test_files/topk.slt @@ -316,7 +316,7 @@ explain select number, letter, age from partial_sorted order by number desc, let ---- physical_plan 01)SortExec: TopK(fetch=3), expr=[number@0 DESC, letter@1 ASC NULLS LAST, age@2 DESC], preserve_partitioning=[false], sort_prefix=[number@0 DESC, letter@1 ASC NULLS LAST] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Explain variations of the above query with different orderings, and different sort prefixes. @@ -326,28 +326,28 @@ explain select number, letter, age from partial_sorted order by age desc limit 3 ---- physical_plan 01)SortExec: TopK(fetch=3), expr=[age@2 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[age@2 DESC], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[age@2 DESC], reverse_row_groups=true, dynamic_rg_pruning=eligible query TT explain select number, letter, age from partial_sorted order by number desc, letter desc limit 3; ---- physical_plan 01)SortExec: TopK(fetch=3), expr=[number@0 DESC, letter@1 DESC], preserve_partitioning=[false], sort_prefix=[number@0 DESC] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible query TT explain select number, letter, age from partial_sorted order by number asc limit 3; ---- physical_plan 01)SortExec: TopK(fetch=3), expr=[number@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[number@0 ASC NULLS LAST] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[number@0 ASC NULLS LAST], dynamic_rg_pruning=eligible query TT explain select number, letter, age from partial_sorted order by letter asc, number desc limit 3; ---- physical_plan 01)SortExec: TopK(fetch=3), expr=[letter@1 ASC NULLS LAST, number@0 DESC], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[letter@1 ASC NULLS LAST, number@0 DESC] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[letter@1 ASC NULLS LAST, number@0 DESC], dynamic_rg_pruning=eligible # Explicit NULLS ordering cases (reversing the order of the NULLS on the number and letter orderings) query TT @@ -355,14 +355,14 @@ explain select number, letter, age from partial_sorted order by number desc, let ---- physical_plan 01)SortExec: TopK(fetch=3), expr=[number@0 DESC, letter@1 ASC], preserve_partitioning=[false], sort_prefix=[number@0 DESC] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible query TT explain select number, letter, age from partial_sorted order by number desc NULLS LAST, letter asc limit 3; ---- physical_plan 01)SortExec: TopK(fetch=3), expr=[number@0 DESC NULLS LAST, letter@1 ASC NULLS LAST], preserve_partitioning=[false] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[number@0 DESC NULLS LAST, letter@1 ASC NULLS LAST], reverse_row_groups=true +02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], file_type=parquet, predicate=DynamicFilter [ empty ], sort_order_for_reorder=[number@0 DESC NULLS LAST, letter@1 ASC NULLS LAST], reverse_row_groups=true, dynamic_rg_pruning=eligible # Verify that the sort prefix is correctly computed on the normalized ordering (removing redundant aliased columns) @@ -370,20 +370,22 @@ query TT explain select number, letter, age, number as column4, letter as column5 from partial_sorted order by number desc, column4 desc, letter asc, column5 asc, age desc limit 3; ---- physical_plan -01)SortExec: TopK(fetch=3), expr=[number@0 DESC, letter@1 ASC NULLS LAST, age@2 DESC], preserve_partitioning=[false], sort_prefix=[number@0 DESC, letter@1 ASC NULLS LAST] -02)--DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age, number@0 as column4, letter@1 as column5], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ] +01)ProjectionExec: expr=[number@0 as number, letter@1 as letter, age@2 as age, number@0 as column4, letter@1 as column5] +02)--SortExec: TopK(fetch=3), expr=[number@0 DESC, letter@1 ASC NULLS LAST, age@2 DESC], preserve_partitioning=[false], sort_prefix=[number@0 DESC, letter@1 ASC NULLS LAST] +03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, letter, age], output_ordering=[number@0 DESC, letter@1 ASC NULLS LAST], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible -# Verify that the sort prefix is correctly computed over normalized, order-maintaining projections (number + 1, number, number + 1, age) +# `number + 1` is not order-maintaining (addition can overflow and wrap), so +# no sort prefix can be computed over the projected expression. query TT explain select number + 1 as number_plus, number, number + 1 as other_number_plus, age from partial_sorted order by number_plus desc, number desc, other_number_plus desc, age asc limit 3; ---- physical_plan 01)SortPreservingMergeExec: [number_plus@0 DESC, number@1 DESC, other_number_plus@2 DESC, age@3 ASC NULLS LAST], fetch=3 -02)--SortExec: TopK(fetch=3), expr=[number_plus@0 DESC, number@1 DESC, age@3 ASC NULLS LAST], preserve_partitioning=[true], sort_prefix=[number_plus@0 DESC, number@1 DESC] -03)----ProjectionExec: expr=[__common_expr_1@0 as number_plus, number@1 as number, __common_expr_1@0 as other_number_plus, age@2 as age] +02)--ProjectionExec: expr=[__common_expr_1@0 as number_plus, number@1 as number, __common_expr_1@0 as other_number_plus, age@2 as age] +03)----SortExec: TopK(fetch=3), expr=[__common_expr_1@0 DESC, number@1 DESC, age@2 ASC NULLS LAST], preserve_partitioning=[true] 04)------ProjectionExec: expr=[CAST(number@0 AS Int64) + 1 as __common_expr_1, number@0 as number, age@1 as age] 05)--------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1, maintains_sort_order=true -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, age], output_ordering=[number@0 DESC], file_type=parquet, predicate=DynamicFilter [ empty ] +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/topk/partial_sorted/1.parquet]]}, projection=[number, age], output_ordering=[number@0 DESC], file_type=parquet, predicate=DynamicFilter [ empty ], dynamic_rg_pruning=eligible # Cleanup statement ok diff --git a/datafusion/sqllogictest/test_files/tpch/create_tables.slt.part b/datafusion/sqllogictest/test_files/tpch/create_tables.slt.part index d6249cb579902..9488367e25569 100644 --- a/datafusion/sqllogictest/test_files/tpch/create_tables.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/create_tables.slt.part @@ -23,7 +23,7 @@ statement ok CREATE EXTERNAL TABLE IF NOT EXISTS supplier ( - s_suppkey BIGINT, + s_suppkey BIGINT PRIMARY KEY, s_name VARCHAR, s_address VARCHAR, s_nationkey BIGINT, @@ -35,7 +35,7 @@ CREATE EXTERNAL TABLE IF NOT EXISTS supplier ( statement ok CREATE EXTERNAL TABLE IF NOT EXISTS part ( - p_partkey BIGINT, + p_partkey BIGINT PRIMARY KEY, p_name VARCHAR, p_mfgr VARCHAR, p_brand VARCHAR, @@ -56,11 +56,12 @@ CREATE EXTERNAL TABLE IF NOT EXISTS partsupp ( ps_supplycost DECIMAL(15, 2), ps_comment VARCHAR, ps_rev VARCHAR, + PRIMARY KEY (ps_partkey, ps_suppkey), ) STORED AS CSV LOCATION 'test_files/tpch/data/partsupp.tbl' OPTIONS ('format.delimiter' '|', 'format.has_header' 'false'); statement ok CREATE EXTERNAL TABLE IF NOT EXISTS customer ( - c_custkey BIGINT, + c_custkey BIGINT PRIMARY KEY, c_name VARCHAR, c_address VARCHAR, c_nationkey BIGINT, @@ -73,7 +74,7 @@ CREATE EXTERNAL TABLE IF NOT EXISTS customer ( statement ok CREATE EXTERNAL TABLE IF NOT EXISTS orders ( - o_orderkey BIGINT, + o_orderkey BIGINT PRIMARY KEY, o_custkey BIGINT, o_orderstatus VARCHAR, o_totalprice DECIMAL(15, 2), @@ -104,11 +105,12 @@ CREATE EXTERNAL TABLE IF NOT EXISTS lineitem ( l_shipmode VARCHAR, l_comment VARCHAR, l_rev VARCHAR, + PRIMARY KEY (l_orderkey, l_linenumber), ) STORED AS CSV LOCATION 'test_files/tpch/data/lineitem.tbl' OPTIONS ('format.delimiter' '|', 'format.has_header' 'false'); statement ok CREATE EXTERNAL TABLE IF NOT EXISTS nation ( - n_nationkey BIGINT, + n_nationkey BIGINT PRIMARY KEY, n_name VARCHAR, n_regionkey BIGINT, n_comment VARCHAR, @@ -117,7 +119,7 @@ CREATE EXTERNAL TABLE IF NOT EXISTS nation ( statement ok CREATE EXTERNAL TABLE IF NOT EXISTS region ( - r_regionkey BIGINT, + r_regionkey BIGINT PRIMARY KEY, r_name VARCHAR, r_comment VARCHAR, r_rev VARCHAR, diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q1.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q1.slt.part index 10c229546b93b..b227f94553e2f 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q1.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q1.slt.part @@ -42,17 +42,17 @@ explain select logical_plan 01)Sort: lineitem.l_returnflag ASC NULLS LAST, lineitem.l_linestatus ASC NULLS LAST 02)--Projection: lineitem.l_returnflag, lineitem.l_linestatus, sum(lineitem.l_quantity) AS sum_qty, sum(lineitem.l_extendedprice) AS sum_base_price, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS sum_disc_price, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax) AS sum_charge, avg(lineitem.l_quantity) AS avg_qty, avg(lineitem.l_extendedprice) AS avg_price, avg(lineitem.l_discount) AS avg_disc, count(Int64(1)) AS count(*) AS count_order -03)----Aggregate: groupBy=[[lineitem.l_returnflag, lineitem.l_linestatus]], aggr=[[sum(lineitem.l_quantity), sum(lineitem.l_extendedprice), sum(__common_expr_1) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount), sum(__common_expr_1 * (Decimal128(Some(1),20,0) + lineitem.l_tax)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax), avg(lineitem.l_quantity), avg(lineitem.l_extendedprice), avg(lineitem.l_discount), count(Int64(1))]] -04)------Projection: lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount) AS __common_expr_1, lineitem.l_quantity, lineitem.l_extendedprice, lineitem.l_discount, lineitem.l_tax, lineitem.l_returnflag, lineitem.l_linestatus +03)----Aggregate: groupBy=[[lineitem.l_returnflag, lineitem.l_linestatus]], aggr=[[sum(lineitem.l_quantity), sum(lineitem.l_extendedprice), sum(__common_expr_1) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount), sum(__common_expr_1 * (Decimal128(1,20,0) + lineitem.l_tax)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax), avg(lineitem.l_quantity), avg(lineitem.l_extendedprice), avg(lineitem.l_discount), count(Int64(1))]] +04)------Projection: lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount) AS __common_expr_1, lineitem.l_quantity, lineitem.l_extendedprice, lineitem.l_discount, lineitem.l_tax, lineitem.l_returnflag, lineitem.l_linestatus 05)--------Filter: lineitem.l_shipdate <= Date32("1998-09-02") 06)----------TableScan: lineitem projection=[l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate], partial_filters=[lineitem.l_shipdate <= Date32("1998-09-02")] physical_plan 01)SortPreservingMergeExec: [l_returnflag@0 ASC NULLS LAST, l_linestatus@1 ASC NULLS LAST] -02)--SortExec: expr=[l_returnflag@0 ASC NULLS LAST, l_linestatus@1 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus, sum(lineitem.l_quantity)@2 as sum_qty, sum(lineitem.l_extendedprice)@3 as sum_base_price, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@4 as sum_disc_price, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax)@5 as sum_charge, avg(lineitem.l_quantity)@6 as avg_qty, avg(lineitem.l_extendedprice)@7 as avg_price, avg(lineitem.l_discount)@8 as avg_disc, count(Int64(1))@9 as count_order] -04)------AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[sum(lineitem.l_quantity), sum(lineitem.l_extendedprice), sum(__common_expr_1) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount), sum(__common_expr_1 * Some(1),20,0 + lineitem.l_tax) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax), avg(lineitem.l_quantity), avg(lineitem.l_extendedprice), avg(lineitem.l_discount), count(Int64(1))] +02)--ProjectionExec: expr=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus, sum(lineitem.l_quantity)@2 as sum_qty, sum(lineitem.l_extendedprice)@3 as sum_base_price, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@4 as sum_disc_price, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax)@5 as sum_charge, avg(lineitem.l_quantity)@6 as avg_qty, avg(lineitem.l_extendedprice)@7 as avg_price, avg(lineitem.l_discount)@8 as avg_disc, count(Int64(1))@9 as count_order] +03)----SortExec: expr=[l_returnflag@0 ASC NULLS LAST, l_linestatus@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[l_returnflag@0 as l_returnflag, l_linestatus@1 as l_linestatus], aggr=[sum(lineitem.l_quantity), sum(lineitem.l_extendedprice), sum(__common_expr_1) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount), sum(__common_expr_1 * 1 + lineitem.l_tax) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax), avg(lineitem.l_quantity), avg(lineitem.l_extendedprice), avg(lineitem.l_discount), count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([l_returnflag@0, l_linestatus@1], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[l_returnflag@5 as l_returnflag, l_linestatus@6 as l_linestatus], aggr=[sum(lineitem.l_quantity), sum(lineitem.l_extendedprice), sum(__common_expr_1) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount), sum(__common_expr_1 * Some(1),20,0 + lineitem.l_tax) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax), avg(lineitem.l_quantity), avg(lineitem.l_extendedprice), avg(lineitem.l_discount), count(Int64(1))] -07)------------ProjectionExec: expr=[l_extendedprice@0 * (Some(1),20,0 - l_discount@1) as __common_expr_1, l_quantity@2 as l_quantity, l_extendedprice@0 as l_extendedprice, l_discount@1 as l_discount, l_tax@3 as l_tax, l_returnflag@4 as l_returnflag, l_linestatus@5 as l_linestatus] +06)----------AggregateExec: mode=Partial, gby=[l_returnflag@5 as l_returnflag, l_linestatus@6 as l_linestatus], aggr=[sum(lineitem.l_quantity), sum(lineitem.l_extendedprice), sum(__common_expr_1) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount), sum(__common_expr_1 * 1 + lineitem.l_tax) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount * Int64(1) + lineitem.l_tax), avg(lineitem.l_quantity), avg(lineitem.l_extendedprice), avg(lineitem.l_discount), count(Int64(1))] +07)------------ProjectionExec: expr=[l_extendedprice@0 * (1 - l_discount@1) as __common_expr_1, l_quantity@2 as l_quantity, l_extendedprice@0 as l_extendedprice, l_discount@1 as l_discount, l_tax@3 as l_tax, l_returnflag@4 as l_returnflag, l_linestatus@5 as l_linestatus] 08)--------------FilterExec: l_shipdate@6 <= 1998-09-02, projection=[l_extendedprice@1, l_discount@2, l_quantity@0, l_tax@3, l_returnflag@4, l_linestatus@5] -09)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate], file_type=csv, has_header=false +09)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_quantity, l_extendedprice, l_discount, l_tax, l_returnflag, l_linestatus, l_shipdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q10.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q10.slt.part index 33d5e273a0d37..9b5db48e83ebc 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q10.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q10.slt.part @@ -53,7 +53,7 @@ limit 10; logical_plan 01)Sort: revenue DESC NULLS FIRST, fetch=10 02)--Projection: customer.c_custkey, customer.c_name, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS revenue, customer.c_acctbal, nation.n_name, customer.c_address, customer.c_phone, customer.c_comment -03)----Aggregate: groupBy=[[customer.c_custkey, customer.c_name, customer.c_acctbal, customer.c_phone, nation.n_name, customer.c_address, customer.c_comment]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] +03)----Aggregate: groupBy=[[customer.c_custkey, customer.c_name, customer.c_acctbal, customer.c_phone, nation.n_name, customer.c_address, customer.c_comment]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] 04)------Projection: customer.c_custkey, customer.c_name, customer.c_address, customer.c_phone, customer.c_acctbal, customer.c_comment, lineitem.l_extendedprice, lineitem.l_discount, nation.n_name 05)--------Inner Join: customer.c_nationkey = nation.n_nationkey 06)----------Projection: customer.c_custkey, customer.c_name, customer.c_address, customer.c_nationkey, customer.c_phone, customer.c_acctbal, customer.c_comment, lineitem.l_extendedprice, lineitem.l_discount @@ -70,23 +70,23 @@ logical_plan 17)----------TableScan: nation projection=[n_nationkey, n_name] physical_plan 01)SortPreservingMergeExec: [revenue@2 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[revenue@2 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[c_custkey@0 as c_custkey, c_name@1 as c_name, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@7 as revenue, c_acctbal@2 as c_acctbal, n_name@4 as n_name, c_address@5 as c_address, c_phone@3 as c_phone, c_comment@6 as c_comment] -04)------AggregateExec: mode=FinalPartitioned, gby=[c_custkey@0 as c_custkey, c_name@1 as c_name, c_acctbal@2 as c_acctbal, c_phone@3 as c_phone, n_name@4 as n_name, c_address@5 as c_address, c_comment@6 as c_comment], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +02)--ProjectionExec: expr=[c_custkey@0 as c_custkey, c_name@1 as c_name, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@7 as revenue, c_acctbal@2 as c_acctbal, n_name@4 as n_name, c_address@5 as c_address, c_phone@3 as c_phone, c_comment@6 as c_comment] +03)----SortExec: TopK(fetch=10), expr=[sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@7 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[c_custkey@0 as c_custkey, c_name@1 as c_name, c_acctbal@2 as c_acctbal, c_phone@3 as c_phone, n_name@4 as n_name, c_address@5 as c_address, c_comment@6 as c_comment], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 05)--------RepartitionExec: partitioning=Hash([c_custkey@0, c_name@1, c_acctbal@2, c_phone@3, n_name@4, c_address@5, c_comment@6], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[c_custkey@0 as c_custkey, c_name@1 as c_name, c_acctbal@4 as c_acctbal, c_phone@3 as c_phone, n_name@8 as n_name, c_address@2 as c_address, c_comment@5 as c_comment], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +06)----------AggregateExec: mode=Partial, gby=[c_custkey@0 as c_custkey, c_name@1 as c_name, c_acctbal@4 as c_acctbal, c_phone@3 as c_phone, n_name@8 as n_name, c_address@2 as c_address, c_comment@5 as c_comment], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 07)------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c_nationkey@3, n_nationkey@0)], projection=[c_custkey@0, c_name@1, c_address@2, c_phone@4, c_acctbal@5, c_comment@6, l_extendedprice@7, l_discount@8, n_name@10] 08)--------------RepartitionExec: partitioning=Hash([c_nationkey@3], 4), input_partitions=4 09)----------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(o_orderkey@7, l_orderkey@0)], projection=[c_custkey@0, c_name@1, c_address@2, c_nationkey@3, c_phone@4, c_acctbal@5, c_comment@6, l_extendedprice@9, l_discount@10] 10)------------------RepartitionExec: partitioning=Hash([o_orderkey@7], 4), input_partitions=4 11)--------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c_custkey@0, o_custkey@1)], projection=[c_custkey@0, c_name@1, c_address@2, c_nationkey@3, c_phone@4, c_acctbal@5, c_comment@6, o_orderkey@7] -12)----------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=1 -13)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_name, c_address, c_nationkey, c_phone, c_acctbal, c_comment], file_type=csv, has_header=false +12)----------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4 +13)------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_custkey, c_name, c_address, c_nationkey, c_phone, c_acctbal, c_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 14)----------------------RepartitionExec: partitioning=Hash([o_custkey@1], 4), input_partitions=4 15)------------------------FilterExec: o_orderdate@2 >= 1993-10-01 AND o_orderdate@2 < 1994-01-01, projection=[o_orderkey@0, o_custkey@1] -16)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_orderdate], file_type=csv, has_header=false +16)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_orderdate], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 17)------------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 18)--------------------FilterExec: l_returnflag@3 = R, projection=[l_orderkey@0, l_extendedprice@1, l_discount@2] -19)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_extendedprice, l_discount, l_returnflag], file_type=csv, has_header=false +19)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_extendedprice, l_discount, l_returnflag], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 20)--------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 -21)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false +21)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q11.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q11.slt.part index e8a224867df05..c1e0a638cc839 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q11.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q11.slt.part @@ -54,7 +54,7 @@ logical_plan 05)--------Projection: CAST(CAST(sum(partsupp.ps_supplycost * partsupp.ps_availqty) AS Float64) * Float64(0.0001) AS Decimal128(38, 15)) 06)----------Aggregate: groupBy=[[]], aggr=[[sum(partsupp.ps_supplycost * CAST(partsupp.ps_availqty AS Decimal128(10, 0)))]] 07)------------Projection: partsupp.ps_availqty, partsupp.ps_supplycost -08)--------------Inner Join: supplier.s_nationkey = nation.n_nationkey +08)--------------LeftSemi Join: supplier.s_nationkey = nation.n_nationkey 09)----------------Projection: partsupp.ps_availqty, partsupp.ps_supplycost, supplier.s_nationkey 10)------------------Inner Join: partsupp.ps_suppkey = supplier.s_suppkey 11)--------------------TableScan: partsupp projection=[ps_suppkey, ps_availqty, ps_supplycost] @@ -64,7 +64,7 @@ logical_plan 15)--------------------TableScan: nation projection=[n_nationkey, n_name], partial_filters=[nation.n_name = Utf8View("GERMANY")] 16)------Aggregate: groupBy=[[partsupp.ps_partkey]], aggr=[[sum(partsupp.ps_supplycost * CAST(partsupp.ps_availqty AS Decimal128(10, 0)))]] 17)--------Projection: partsupp.ps_partkey, partsupp.ps_availqty, partsupp.ps_supplycost -18)----------Inner Join: supplier.s_nationkey = nation.n_nationkey +18)----------LeftSemi Join: supplier.s_nationkey = nation.n_nationkey 19)------------Projection: partsupp.ps_partkey, partsupp.ps_availqty, partsupp.ps_supplycost, supplier.s_nationkey 20)--------------Inner Join: partsupp.ps_suppkey = supplier.s_suppkey 21)----------------TableScan: partsupp projection=[ps_partkey, ps_suppkey, ps_availqty, ps_supplycost] @@ -75,35 +75,35 @@ logical_plan physical_plan 01)ScalarSubqueryExec: subqueries=1 02)--SortPreservingMergeExec: [value@1 DESC], fetch=10 -03)----SortExec: TopK(fetch=10), expr=[value@1 DESC], preserve_partitioning=[true] -04)------ProjectionExec: expr=[ps_partkey@0 as ps_partkey, sum(partsupp.ps_supplycost * partsupp.ps_availqty)@1 as value] +03)----ProjectionExec: expr=[ps_partkey@0 as ps_partkey, sum(partsupp.ps_supplycost * partsupp.ps_availqty)@1 as value] +04)------SortExec: TopK(fetch=10), expr=[sum(partsupp.ps_supplycost * partsupp.ps_availqty)@1 DESC], preserve_partitioning=[true] 05)--------FilterExec: CAST(sum(partsupp.ps_supplycost * partsupp.ps_availqty)@1 AS Decimal128(38, 15)) > scalar_subquery() 06)----------AggregateExec: mode=FinalPartitioned, gby=[ps_partkey@0 as ps_partkey], aggr=[sum(partsupp.ps_supplycost * partsupp.ps_availqty)] 07)------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 08)--------------AggregateExec: mode=Partial, gby=[ps_partkey@0 as ps_partkey], aggr=[sum(partsupp.ps_supplycost * partsupp.ps_availqty)] -09)----------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@3, n_nationkey@0)], projection=[ps_partkey@0, ps_availqty@1, ps_supplycost@2] +09)----------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(s_nationkey@3, n_nationkey@0)], projection=[ps_partkey@0, ps_availqty@1, ps_supplycost@2] 10)------------------RepartitionExec: partitioning=Hash([s_nationkey@3], 4), input_partitions=4 11)--------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_suppkey@1, s_suppkey@0)], projection=[ps_partkey@0, ps_availqty@2, ps_supplycost@3, s_nationkey@5] 12)----------------------RepartitionExec: partitioning=Hash([ps_suppkey@1], 4), input_partitions=4 -13)------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_availqty, ps_supplycost], file_type=csv, has_header=false +13)------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_availqty, ps_supplycost], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false 14)----------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -15)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], file_type=csv, has_header=false +15)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 16)------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=4 17)--------------------FilterExec: n_name@1 = GERMANY, projection=[n_nationkey@0] 18)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -19)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false +19)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 20)--ProjectionExec: expr=[CAST(CAST(sum(partsupp.ps_supplycost * partsupp.ps_availqty)@0 AS Float64) * 0.0001 AS Decimal128(38, 15)) as sum(partsupp.ps_supplycost * partsupp.ps_availqty) * Float64(0.0001)] 21)----AggregateExec: mode=Final, gby=[], aggr=[sum(partsupp.ps_supplycost * partsupp.ps_availqty)] 22)------CoalescePartitionsExec 23)--------AggregateExec: mode=Partial, gby=[], aggr=[sum(partsupp.ps_supplycost * partsupp.ps_availqty)] -24)----------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@2, n_nationkey@0)], projection=[ps_availqty@0, ps_supplycost@1] +24)----------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(s_nationkey@2, n_nationkey@0)], projection=[ps_availqty@0, ps_supplycost@1] 25)------------RepartitionExec: partitioning=Hash([s_nationkey@2], 4), input_partitions=4 26)--------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_suppkey@0, s_suppkey@0)], projection=[ps_availqty@1, ps_supplycost@2, s_nationkey@4] 27)----------------RepartitionExec: partitioning=Hash([ps_suppkey@0], 4), input_partitions=4 -28)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_suppkey, ps_availqty, ps_supplycost], file_type=csv, has_header=false +28)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_suppkey, ps_availqty, ps_supplycost], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false 29)----------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -30)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], file_type=csv, has_header=false +30)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 31)------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=4 32)--------------FilterExec: n_name@1 = GERMANY, projection=[n_nationkey@0] 33)----------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -34)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false +34)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q12.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q12.slt.part index 84a6598cb992b..a9d579b6a590d 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q12.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q12.slt.part @@ -60,14 +60,14 @@ logical_plan 09)----------TableScan: orders projection=[o_orderkey, o_orderpriority] physical_plan 01)SortPreservingMergeExec: [l_shipmode@0 ASC NULLS LAST] -02)--SortExec: expr=[l_shipmode@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[l_shipmode@0 as l_shipmode, sum(CASE WHEN orders.o_orderpriority = Utf8("1-URGENT") OR orders.o_orderpriority = Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)@1 as high_line_count, sum(CASE WHEN orders.o_orderpriority != Utf8("1-URGENT") AND orders.o_orderpriority != Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)@2 as low_line_count] +02)--ProjectionExec: expr=[l_shipmode@0 as l_shipmode, sum(CASE WHEN orders.o_orderpriority = Utf8("1-URGENT") OR orders.o_orderpriority = Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)@1 as high_line_count, sum(CASE WHEN orders.o_orderpriority != Utf8("1-URGENT") AND orders.o_orderpriority != Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)@2 as low_line_count] +03)----SortExec: expr=[l_shipmode@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[l_shipmode@0 as l_shipmode], aggr=[sum(CASE WHEN orders.o_orderpriority = 1-URGENT OR orders.o_orderpriority = 2-HIGH THEN 1 ELSE 0 END) as sum(CASE WHEN orders.o_orderpriority = Utf8("1-URGENT") OR orders.o_orderpriority = Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END), sum(CASE WHEN orders.o_orderpriority != 1-URGENT AND orders.o_orderpriority != 2-HIGH THEN 1 ELSE 0 END) as sum(CASE WHEN orders.o_orderpriority != Utf8("1-URGENT") AND orders.o_orderpriority != Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)] 05)--------RepartitionExec: partitioning=Hash([l_shipmode@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[l_shipmode@0 as l_shipmode], aggr=[sum(CASE WHEN orders.o_orderpriority = 1-URGENT OR orders.o_orderpriority = 2-HIGH THEN 1 ELSE 0 END) as sum(CASE WHEN orders.o_orderpriority = Utf8("1-URGENT") OR orders.o_orderpriority = Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END), sum(CASE WHEN orders.o_orderpriority != 1-URGENT AND orders.o_orderpriority != 2-HIGH THEN 1 ELSE 0 END) as sum(CASE WHEN orders.o_orderpriority != Utf8("1-URGENT") AND orders.o_orderpriority != Utf8("2-HIGH") THEN Int64(1) ELSE Int64(0) END)] 07)------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_orderkey@0, o_orderkey@0)], projection=[l_shipmode@1, o_orderpriority@3] 08)--------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 09)----------------FilterExec: (l_shipmode@4 = MAIL OR l_shipmode@4 = SHIP) AND l_receiptdate@3 > l_commitdate@2 AND l_shipdate@1 < l_commitdate@2 AND l_receiptdate@3 >= 1994-01-01 AND l_receiptdate@3 < 1995-01-01, projection=[l_orderkey@0, l_shipmode@4] -10)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_shipdate, l_commitdate, l_receiptdate, l_shipmode], file_type=csv, has_header=false +10)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_shipdate, l_commitdate, l_receiptdate, l_shipmode], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 11)--------------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 -12)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_orderpriority], file_type=csv, has_header=false +12)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_orderpriority], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q13.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q13.slt.part index 94e0848bfcce1..9f9cbb3b6af68 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q13.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q13.slt.part @@ -54,16 +54,16 @@ logical_plan 12)--------------------TableScan: orders projection=[o_orderkey, o_custkey, o_comment], partial_filters=[orders.o_comment NOT LIKE Utf8View("%special%requests%")] physical_plan 01)SortPreservingMergeExec: [custdist@1 DESC, c_count@0 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[custdist@1 DESC, c_count@0 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[c_count@0 as c_count, count(Int64(1))@1 as custdist] +02)--ProjectionExec: expr=[c_count@0 as c_count, count(Int64(1))@1 as custdist] +03)----SortExec: TopK(fetch=10), expr=[count(Int64(1))@1 DESC, c_count@0 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[c_count@0 as c_count], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([c_count@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[c_count@0 as c_count], aggr=[count(Int64(1))] 07)------------ProjectionExec: expr=[count(orders.o_orderkey)@1 as c_count] 08)--------------AggregateExec: mode=SinglePartitioned, gby=[c_custkey@0 as c_custkey], aggr=[count(orders.o_orderkey)] 09)----------------HashJoinExec: mode=Partitioned, join_type=Left, on=[(c_custkey@0, o_custkey@1)], projection=[c_custkey@0, o_orderkey@1] -10)------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=1 -11)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey], file_type=csv, has_header=false +10)------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4 +11)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_custkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 12)------------------RepartitionExec: partitioning=Hash([o_custkey@1], 4), input_partitions=4 13)--------------------FilterExec: o_comment@2 NOT LIKE %special%requests%, projection=[o_orderkey@0, o_custkey@1] -14)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_comment], file_type=csv, has_header=false +14)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q14.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q14.slt.part index 198e6676f841f..68e7e3a329747 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q14.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q14.slt.part @@ -33,8 +33,8 @@ where ---- logical_plan 01)Projection: Float64(100) * CAST(sum(CASE WHEN part.p_type LIKE Utf8("PROMO%") THEN lineitem.l_extendedprice * Int64(1) - lineitem.l_discount ELSE Int64(0) END) AS Float64) / CAST(sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS Float64) AS promo_revenue -02)--Aggregate: groupBy=[[]], aggr=[[sum(CASE WHEN part.p_type LIKE Utf8View("PROMO%") THEN __common_expr_1 ELSE Decimal128(Some(0),38,4) END) AS sum(CASE WHEN part.p_type LIKE Utf8("PROMO%") THEN lineitem.l_extendedprice * Int64(1) - lineitem.l_discount ELSE Int64(0) END), sum(__common_expr_1) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] -03)----Projection: lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount) AS __common_expr_1, part.p_type +02)--Aggregate: groupBy=[[]], aggr=[[sum(CASE WHEN part.p_type LIKE Utf8View("PROMO%") THEN __common_expr_1 ELSE Decimal128(0.0000,38,4) END) AS sum(CASE WHEN part.p_type LIKE Utf8("PROMO%") THEN lineitem.l_extendedprice * Int64(1) - lineitem.l_discount ELSE Int64(0) END), sum(__common_expr_1) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] +03)----Projection: lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount) AS __common_expr_1, part.p_type 04)------Inner Join: lineitem.l_partkey = part.p_partkey 05)--------Projection: lineitem.l_partkey, lineitem.l_extendedprice, lineitem.l_discount 06)----------Filter: lineitem.l_shipdate >= Date32("1995-09-01") AND lineitem.l_shipdate < Date32("1995-10-01") @@ -42,13 +42,13 @@ logical_plan 08)--------TableScan: part projection=[p_partkey, p_type] physical_plan 01)ProjectionExec: expr=[100 * CAST(sum(CASE WHEN part.p_type LIKE Utf8("PROMO%") THEN lineitem.l_extendedprice * Int64(1) - lineitem.l_discount ELSE Int64(0) END)@0 AS Float64) / CAST(sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 AS Float64) as promo_revenue] -02)--AggregateExec: mode=Final, gby=[], aggr=[sum(CASE WHEN part.p_type LIKE PROMO% THEN __common_expr_1 ELSE Some(0),38,4 END) as sum(CASE WHEN part.p_type LIKE Utf8("PROMO%") THEN lineitem.l_extendedprice * Int64(1) - lineitem.l_discount ELSE Int64(0) END), sum(__common_expr_1) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +02)--AggregateExec: mode=Final, gby=[], aggr=[sum(CASE WHEN part.p_type LIKE PROMO% THEN __common_expr_1 ELSE 0.0000 END) as sum(CASE WHEN part.p_type LIKE Utf8("PROMO%") THEN lineitem.l_extendedprice * Int64(1) - lineitem.l_discount ELSE Int64(0) END), sum(__common_expr_1) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 03)----CoalescePartitionsExec -04)------AggregateExec: mode=Partial, gby=[], aggr=[sum(CASE WHEN part.p_type LIKE PROMO% THEN __common_expr_1 ELSE Some(0),38,4 END) as sum(CASE WHEN part.p_type LIKE Utf8("PROMO%") THEN lineitem.l_extendedprice * Int64(1) - lineitem.l_discount ELSE Int64(0) END), sum(__common_expr_1) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] -05)--------ProjectionExec: expr=[l_extendedprice@0 * (Some(1),20,0 - l_discount@1) as __common_expr_1, p_type@2 as p_type] +04)------AggregateExec: mode=Partial, gby=[], aggr=[sum(CASE WHEN part.p_type LIKE PROMO% THEN __common_expr_1 ELSE 0.0000 END) as sum(CASE WHEN part.p_type LIKE Utf8("PROMO%") THEN lineitem.l_extendedprice * Int64(1) - lineitem.l_discount ELSE Int64(0) END), sum(__common_expr_1) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +05)--------ProjectionExec: expr=[l_extendedprice@0 * (1 - l_discount@1) as __common_expr_1, p_type@2 as p_type] 06)----------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_partkey@0, p_partkey@0)], projection=[l_extendedprice@1, l_discount@2, p_type@4] 07)------------RepartitionExec: partitioning=Hash([l_partkey@0], 4), input_partitions=4 08)--------------FilterExec: l_shipdate@3 >= 1995-09-01 AND l_shipdate@3 < 1995-10-01, projection=[l_partkey@0, l_extendedprice@1, l_discount@2] -09)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_extendedprice, l_discount, l_shipdate], file_type=csv, has_header=false -10)------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=1 -11)--------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_type], file_type=csv, has_header=false +09)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_extendedprice, l_discount, l_shipdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false +10)------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 +11)--------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_type], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q15.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q15.slt.part index 388e473c00764..097b313cd69ae 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q15.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q15.slt.part @@ -61,11 +61,11 @@ logical_plan 09)--------------Aggregate: groupBy=[[]], aggr=[[max(revenue0.total_revenue)]] 10)----------------SubqueryAlias: revenue0 11)------------------Projection: sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS total_revenue -12)--------------------Aggregate: groupBy=[[lineitem.l_suppkey]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] +12)--------------------Aggregate: groupBy=[[lineitem.l_suppkey]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] 13)----------------------Projection: lineitem.l_suppkey, lineitem.l_extendedprice, lineitem.l_discount 14)------------------------Filter: lineitem.l_shipdate >= Date32("1996-01-01") AND lineitem.l_shipdate < Date32("1996-04-01") 15)--------------------------TableScan: lineitem projection=[l_suppkey, l_extendedprice, l_discount, l_shipdate], partial_filters=[lineitem.l_shipdate >= Date32("1996-01-01"), lineitem.l_shipdate < Date32("1996-04-01")] -16)------------Aggregate: groupBy=[[lineitem.l_suppkey]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] +16)------------Aggregate: groupBy=[[lineitem.l_suppkey]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] 17)--------------Projection: lineitem.l_suppkey, lineitem.l_extendedprice, lineitem.l_discount 18)----------------Filter: lineitem.l_shipdate >= Date32("1996-01-01") AND lineitem.l_shipdate < Date32("1996-04-01") 19)------------------TableScan: lineitem projection=[l_suppkey, l_extendedprice, l_discount, l_shipdate], partial_filters=[lineitem.l_shipdate >= Date32("1996-01-01"), lineitem.l_shipdate < Date32("1996-04-01")] @@ -75,20 +75,20 @@ physical_plan 03)----SortExec: expr=[s_suppkey@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_suppkey@0, supplier_no@0)], projection=[s_suppkey@0, s_name@1, s_address@2, s_phone@3, total_revenue@5] 05)--------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_name, s_address, s_phone], file_type=csv, has_header=false +06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_name, s_address, s_phone], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 07)--------ProjectionExec: expr=[l_suppkey@0 as supplier_no, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 as total_revenue] 08)----------FilterExec: sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 = scalar_subquery() -09)------------AggregateExec: mode=FinalPartitioned, gby=[l_suppkey@0 as l_suppkey], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +09)------------AggregateExec: mode=FinalPartitioned, gby=[l_suppkey@0 as l_suppkey], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 10)--------------RepartitionExec: partitioning=Hash([l_suppkey@0], 4), input_partitions=4 -11)----------------AggregateExec: mode=Partial, gby=[l_suppkey@0 as l_suppkey], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +11)----------------AggregateExec: mode=Partial, gby=[l_suppkey@0 as l_suppkey], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 12)------------------FilterExec: l_shipdate@3 >= 1996-01-01 AND l_shipdate@3 < 1996-04-01, projection=[l_suppkey@0, l_extendedprice@1, l_discount@2] -13)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_suppkey, l_extendedprice, l_discount, l_shipdate], file_type=csv, has_header=false +13)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_suppkey, l_extendedprice, l_discount, l_shipdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 14)--AggregateExec: mode=Final, gby=[], aggr=[max(revenue0.total_revenue)] 15)----CoalescePartitionsExec 16)------AggregateExec: mode=Partial, gby=[], aggr=[max(revenue0.total_revenue)] 17)--------ProjectionExec: expr=[sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 as total_revenue] -18)----------AggregateExec: mode=FinalPartitioned, gby=[l_suppkey@0 as l_suppkey], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +18)----------AggregateExec: mode=FinalPartitioned, gby=[l_suppkey@0 as l_suppkey], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 19)------------RepartitionExec: partitioning=Hash([l_suppkey@0], 4), input_partitions=4 -20)--------------AggregateExec: mode=Partial, gby=[l_suppkey@0 as l_suppkey], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +20)--------------AggregateExec: mode=Partial, gby=[l_suppkey@0 as l_suppkey], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 21)----------------FilterExec: l_shipdate@3 >= 1996-01-01 AND l_shipdate@3 < 1996-04-01, projection=[l_suppkey@0, l_extendedprice@1, l_discount@2] -22)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_suppkey, l_extendedprice, l_discount, l_shipdate], file_type=csv, has_header=false +22)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_suppkey, l_extendedprice, l_discount, l_shipdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part index b01110b567ca8..5902204e2f7a0 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q16.slt.part @@ -54,35 +54,34 @@ logical_plan 02)--Projection: part.p_brand, part.p_type, part.p_size, count(alias1) AS supplier_cnt 03)----Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size]], aggr=[[count(alias1)]] 04)------Aggregate: groupBy=[[part.p_brand, part.p_type, part.p_size, partsupp.ps_suppkey AS alias1]], aggr=[[]] -05)--------LeftAnti Join: partsupp.ps_suppkey = __correlated_sq_1.s_suppkey +05)--------LeftAnti Join: partsupp.ps_suppkey = __correlated_sq_1.s_suppkey null_aware 06)----------Projection: partsupp.ps_suppkey, part.p_brand, part.p_type, part.p_size 07)------------Inner Join: partsupp.ps_partkey = part.p_partkey 08)--------------TableScan: partsupp projection=[ps_partkey, ps_suppkey] -09)--------------Filter: part.p_brand != Utf8View("Brand#45") AND part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%") AND part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]) -10)----------------TableScan: part projection=[p_partkey, p_brand, p_type, p_size], partial_filters=[part.p_brand != Utf8View("Brand#45"), part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%"), part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)])] +09)--------------Filter: part.p_brand != Utf8View("Brand#45") AND part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]) AND part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%") +10)----------------TableScan: part projection=[p_partkey, p_brand, p_type, p_size], partial_filters=[part.p_brand != Utf8View("Brand#45"), part.p_size IN ([Int32(49), Int32(14), Int32(23), Int32(45), Int32(19), Int32(3), Int32(36), Int32(9)]), part.p_type NOT LIKE Utf8View("MEDIUM POLISHED%")] 11)----------SubqueryAlias: __correlated_sq_1 12)------------Projection: supplier.s_suppkey 13)--------------Filter: supplier.s_comment LIKE Utf8View("%Customer%Complaints%") 14)----------------TableScan: supplier projection=[s_suppkey, s_comment], partial_filters=[supplier.s_comment LIKE Utf8View("%Customer%Complaints%")] physical_plan 01)SortPreservingMergeExec: [supplier_cnt@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[supplier_cnt@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, count(alias1)@3 as supplier_cnt] +02)--ProjectionExec: expr=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, count(alias1)@3 as supplier_cnt] +03)----SortExec: TopK(fetch=10), expr=[count(alias1)@3 DESC, p_brand@0 ASC NULLS LAST, p_type@1 ASC NULLS LAST, p_size@2 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(alias1)] 05)--------RepartitionExec: partitioning=Hash([p_brand@0, p_type@1, p_size@2], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size], aggr=[count(alias1)] 07)------------AggregateExec: mode=FinalPartitioned, gby=[p_brand@0 as p_brand, p_type@1 as p_type, p_size@2 as p_size, alias1@3 as alias1], aggr=[] 08)--------------RepartitionExec: partitioning=Hash([p_brand@0, p_type@1, p_size@2, alias1@3], 4), input_partitions=4 09)----------------AggregateExec: mode=Partial, gby=[p_brand@1 as p_brand, p_type@2 as p_type, p_size@3 as p_size, ps_suppkey@0 as alias1], aggr=[] -10)------------------HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(ps_suppkey@0, s_suppkey@0)] +10)------------------HashJoinExec: mode=CollectLeft, join_type=LeftAnti, on=[(ps_suppkey@0, s_suppkey@0)], null_aware 11)--------------------CoalescePartitionsExec 12)----------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_partkey@0, p_partkey@0)], projection=[ps_suppkey@1, p_brand@3, p_type@4, p_size@5] 13)------------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 -14)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey], file_type=csv, has_header=false +14)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false 15)------------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 -16)--------------------------FilterExec: p_brand@1 != Brand#45 AND p_type@2 NOT LIKE MEDIUM POLISHED% AND p_size@3 IN (SET) ([49, 14, 23, 45, 19, 3, 36, 9]) -17)----------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -18)------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_brand, p_type, p_size], file_type=csv, has_header=false -19)--------------------FilterExec: s_comment@1 LIKE %Customer%Complaints%, projection=[s_suppkey@0] -20)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -21)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_comment], file_type=csv, has_header=false +16)--------------------------FilterExec: p_brand@1 != Brand#45 AND p_size@3 IN (SET) ([49, 14, 23, 45, 19, 3, 36, 9]) AND p_type@2 NOT LIKE MEDIUM POLISHED% +17)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_type, p_size], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +18)--------------------FilterExec: s_comment@1 LIKE %Customer%Complaints%, projection=[s_suppkey@0] +19)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +20)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q17.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q17.slt.part index 83294d61a1698..e678f8b440dd4 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q17.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q17.slt.part @@ -39,7 +39,7 @@ logical_plan 01)Projection: CAST(sum(lineitem.l_extendedprice) AS Float64) / Float64(7) AS avg_yearly 02)--Aggregate: groupBy=[[]], aggr=[[sum(lineitem.l_extendedprice)]] 03)----Projection: lineitem.l_extendedprice -04)------Inner Join: part.p_partkey = __scalar_sq_1.l_partkey Filter: CAST(lineitem.l_quantity AS Decimal128(30, 15)) < __scalar_sq_1.Float64(0.2) * avg(lineitem.l_quantity) +04)------LeftSemi Join: part.p_partkey = __scalar_sq_1.l_partkey Filter: CAST(lineitem.l_quantity AS Decimal128(30, 15)) < __scalar_sq_1.Float64(0.2) * avg(lineitem.l_quantity) 05)--------Projection: lineitem.l_quantity, lineitem.l_extendedprice, part.p_partkey 06)----------Inner Join: lineitem.l_partkey = part.p_partkey 07)------------TableScan: lineitem projection=[l_partkey, l_quantity, l_extendedprice] @@ -55,16 +55,15 @@ physical_plan 02)--AggregateExec: mode=Final, gby=[], aggr=[sum(lineitem.l_extendedprice)] 03)----CoalescePartitionsExec 04)------AggregateExec: mode=Partial, gby=[], aggr=[sum(lineitem.l_extendedprice)] -05)--------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(p_partkey@2, l_partkey@1)], filter=CAST(l_quantity@0 AS Decimal128(30, 15)) < Float64(0.2) * avg(lineitem.l_quantity)@1, projection=[l_extendedprice@1] +05)--------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(p_partkey@2, l_partkey@1)], filter=CAST(l_quantity@0 AS Decimal128(30, 15)) < Float64(0.2) * avg(lineitem.l_quantity)@1, projection=[l_extendedprice@1] 06)----------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_partkey@0, p_partkey@0)], projection=[l_quantity@1, l_extendedprice@2, p_partkey@3] 07)------------RepartitionExec: partitioning=Hash([l_partkey@0], 4), input_partitions=4 -08)--------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_quantity, l_extendedprice], file_type=csv, has_header=false +08)--------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_quantity, l_extendedprice], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 09)------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 10)--------------FilterExec: p_brand@1 = Brand#23 AND p_container@2 = MED BOX, projection=[p_partkey@0] -11)----------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -12)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_brand, p_container], file_type=csv, has_header=false -13)----------ProjectionExec: expr=[CAST(0.2 * CAST(avg(lineitem.l_quantity)@1 AS Float64) AS Decimal128(30, 15)) as Float64(0.2) * avg(lineitem.l_quantity), l_partkey@0 as l_partkey] -14)------------AggregateExec: mode=FinalPartitioned, gby=[l_partkey@0 as l_partkey], aggr=[avg(lineitem.l_quantity)] -15)--------------RepartitionExec: partitioning=Hash([l_partkey@0], 4), input_partitions=4 -16)----------------AggregateExec: mode=Partial, gby=[l_partkey@0 as l_partkey], aggr=[avg(lineitem.l_quantity)] -17)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_quantity], file_type=csv, has_header=false +11)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_container], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +12)----------ProjectionExec: expr=[CAST(0.2 * CAST(avg(lineitem.l_quantity)@1 AS Float64) AS Decimal128(30, 15)) as Float64(0.2) * avg(lineitem.l_quantity), l_partkey@0 as l_partkey] +13)------------AggregateExec: mode=FinalPartitioned, gby=[l_partkey@0 as l_partkey], aggr=[avg(lineitem.l_quantity)] +14)--------------RepartitionExec: partitioning=Hash([l_partkey@0], 4), input_partitions=4 +15)----------------AggregateExec: mode=Partial, gby=[l_partkey@0 as l_partkey], aggr=[avg(lineitem.l_quantity)] +16)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_quantity], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q18.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q18.slt.part index 617051d602bd6..3602aa1f4a8ed 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q18.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q18.slt.part @@ -63,7 +63,7 @@ logical_plan 10)----------TableScan: lineitem projection=[l_orderkey, l_quantity] 11)------SubqueryAlias: __correlated_sq_1 12)--------Projection: lineitem.l_orderkey -13)----------Filter: sum(lineitem.l_quantity) > Decimal128(Some(30000),25,2) +13)----------Filter: sum(lineitem.l_quantity) > Decimal128(300.00,25,2) 14)------------Aggregate: groupBy=[[lineitem.l_orderkey]], aggr=[[sum(lineitem.l_quantity)]] 15)--------------TableScan: lineitem projection=[l_orderkey, l_quantity] physical_plan @@ -74,14 +74,14 @@ physical_plan 05)--------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(o_orderkey@2, l_orderkey@0)], projection=[c_custkey@0, c_name@1, o_orderkey@2, o_totalprice@3, o_orderdate@4, l_quantity@6] 06)----------RepartitionExec: partitioning=Hash([o_orderkey@2], 4), input_partitions=4 07)------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c_custkey@0, o_custkey@1)], projection=[c_custkey@0, c_name@1, o_orderkey@2, o_totalprice@4, o_orderdate@5] -08)--------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=1 -09)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_name], file_type=csv, has_header=false +08)--------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4 +09)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_custkey, c_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 10)--------------RepartitionExec: partitioning=Hash([o_custkey@1], 4), input_partitions=4 -11)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_totalprice, o_orderdate], file_type=csv, has_header=false +11)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_totalprice, o_orderdate], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 12)----------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 -13)------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_quantity], file_type=csv, has_header=false -14)--------FilterExec: sum(lineitem.l_quantity)@1 > Some(30000),25,2, projection=[l_orderkey@0] +13)------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_quantity], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false +14)--------FilterExec: sum(lineitem.l_quantity)@1 > 300.00, projection=[l_orderkey@0] 15)----------AggregateExec: mode=FinalPartitioned, gby=[l_orderkey@0 as l_orderkey], aggr=[sum(lineitem.l_quantity)] 16)------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 17)--------------AggregateExec: mode=Partial, gby=[l_orderkey@0 as l_orderkey], aggr=[sum(lineitem.l_quantity)] -18)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_quantity], file_type=csv, has_header=false +18)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_quantity], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q19.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q19.slt.part index 9ac2aaa4a67fc..7ef36a72eca26 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q19.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q19.slt.part @@ -55,24 +55,23 @@ where ---- logical_plan 01)Projection: sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS revenue -02)--Aggregate: groupBy=[[]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] +02)--Aggregate: groupBy=[[]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] 03)----Projection: lineitem.l_extendedprice, lineitem.l_discount -04)------Inner Join: lineitem.l_partkey = part.p_partkey Filter: part.p_brand = Utf8View("Brand#12") AND part.p_container IN ([Utf8View("SM CASE"), Utf8View("SM BOX"), Utf8View("SM PACK"), Utf8View("SM PKG")]) AND lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_container IN ([Utf8View("MED BAG"), Utf8View("MED BOX"), Utf8View("MED PKG"), Utf8View("MED PACK")]) AND lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_container IN ([Utf8View("LG CASE"), Utf8View("LG BOX"), Utf8View("LG PACK"), Utf8View("LG PKG")]) AND lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2) AND part.p_size <= Int32(15) +04)------LeftSemi Join: lineitem.l_partkey = part.p_partkey Filter: part.p_brand = Utf8View("Brand#12") AND part.p_container IN ([Utf8View("SM CASE"), Utf8View("SM BOX"), Utf8View("SM PACK"), Utf8View("SM PKG")]) AND lineitem.l_quantity >= Decimal128(1.00,15,2) AND lineitem.l_quantity <= Decimal128(11.00,15,2) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_container IN ([Utf8View("MED BAG"), Utf8View("MED BOX"), Utf8View("MED PKG"), Utf8View("MED PACK")]) AND lineitem.l_quantity >= Decimal128(10.00,15,2) AND lineitem.l_quantity <= Decimal128(20.00,15,2) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_container IN ([Utf8View("LG CASE"), Utf8View("LG BOX"), Utf8View("LG PACK"), Utf8View("LG PKG")]) AND lineitem.l_quantity >= Decimal128(20.00,15,2) AND lineitem.l_quantity <= Decimal128(30.00,15,2) AND part.p_size <= Int32(15) 05)--------Projection: lineitem.l_partkey, lineitem.l_quantity, lineitem.l_extendedprice, lineitem.l_discount -06)----------Filter: (lineitem.l_shipmode = Utf8View("AIR") OR lineitem.l_shipmode = Utf8View("AIR REG")) AND lineitem.l_shipinstruct = Utf8View("DELIVER IN PERSON") AND (lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) OR lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) OR lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2)) -07)------------TableScan: lineitem projection=[l_partkey, l_quantity, l_extendedprice, l_discount, l_shipinstruct, l_shipmode], partial_filters=[lineitem.l_shipmode = Utf8View("AIR") OR lineitem.l_shipmode = Utf8View("AIR REG"), lineitem.l_shipinstruct = Utf8View("DELIVER IN PERSON"), lineitem.l_quantity >= Decimal128(Some(100),15,2) AND lineitem.l_quantity <= Decimal128(Some(1100),15,2) OR lineitem.l_quantity >= Decimal128(Some(1000),15,2) AND lineitem.l_quantity <= Decimal128(Some(2000),15,2) OR lineitem.l_quantity >= Decimal128(Some(2000),15,2) AND lineitem.l_quantity <= Decimal128(Some(3000),15,2)] +06)----------Filter: (lineitem.l_shipmode = Utf8View("AIR") OR lineitem.l_shipmode = Utf8View("AIR REG")) AND lineitem.l_shipinstruct = Utf8View("DELIVER IN PERSON") AND (lineitem.l_quantity >= Decimal128(1.00,15,2) AND lineitem.l_quantity <= Decimal128(11.00,15,2) OR lineitem.l_quantity >= Decimal128(10.00,15,2) AND lineitem.l_quantity <= Decimal128(20.00,15,2) OR lineitem.l_quantity >= Decimal128(20.00,15,2) AND lineitem.l_quantity <= Decimal128(30.00,15,2)) +07)------------TableScan: lineitem projection=[l_partkey, l_quantity, l_extendedprice, l_discount, l_shipinstruct, l_shipmode], partial_filters=[lineitem.l_shipmode = Utf8View("AIR") OR lineitem.l_shipmode = Utf8View("AIR REG"), lineitem.l_shipinstruct = Utf8View("DELIVER IN PERSON"), lineitem.l_quantity >= Decimal128(1.00,15,2) AND lineitem.l_quantity <= Decimal128(11.00,15,2) OR lineitem.l_quantity >= Decimal128(10.00,15,2) AND lineitem.l_quantity <= Decimal128(20.00,15,2) OR lineitem.l_quantity >= Decimal128(20.00,15,2) AND lineitem.l_quantity <= Decimal128(30.00,15,2)] 08)--------Filter: part.p_size >= Int32(1) AND (part.p_brand = Utf8View("Brand#12") AND part.p_container IN ([Utf8View("SM CASE"), Utf8View("SM BOX"), Utf8View("SM PACK"), Utf8View("SM PKG")]) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_container IN ([Utf8View("MED BAG"), Utf8View("MED BOX"), Utf8View("MED PKG"), Utf8View("MED PACK")]) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_container IN ([Utf8View("LG CASE"), Utf8View("LG BOX"), Utf8View("LG PACK"), Utf8View("LG PKG")]) AND part.p_size <= Int32(15)) 09)----------TableScan: part projection=[p_partkey, p_brand, p_size, p_container], partial_filters=[part.p_size >= Int32(1), part.p_brand = Utf8View("Brand#12") AND part.p_container IN ([Utf8View("SM CASE"), Utf8View("SM BOX"), Utf8View("SM PACK"), Utf8View("SM PKG")]) AND part.p_size <= Int32(5) OR part.p_brand = Utf8View("Brand#23") AND part.p_container IN ([Utf8View("MED BAG"), Utf8View("MED BOX"), Utf8View("MED PKG"), Utf8View("MED PACK")]) AND part.p_size <= Int32(10) OR part.p_brand = Utf8View("Brand#34") AND part.p_container IN ([Utf8View("LG CASE"), Utf8View("LG BOX"), Utf8View("LG PACK"), Utf8View("LG PKG")]) AND part.p_size <= Int32(15)] physical_plan 01)ProjectionExec: expr=[sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@0 as revenue] -02)--AggregateExec: mode=Final, gby=[], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +02)--AggregateExec: mode=Final, gby=[], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 03)----CoalescePartitionsExec -04)------AggregateExec: mode=Partial, gby=[], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] -05)--------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_partkey@0, p_partkey@0)], filter=p_brand@1 = Brand#12 AND p_container@3 IN (SET) ([SM CASE, SM BOX, SM PACK, SM PKG]) AND l_quantity@0 >= Some(100),15,2 AND l_quantity@0 <= Some(1100),15,2 AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND p_container@3 IN (SET) ([MED BAG, MED BOX, MED PKG, MED PACK]) AND l_quantity@0 >= Some(1000),15,2 AND l_quantity@0 <= Some(2000),15,2 AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND p_container@3 IN (SET) ([LG CASE, LG BOX, LG PACK, LG PKG]) AND l_quantity@0 >= Some(2000),15,2 AND l_quantity@0 <= Some(3000),15,2 AND p_size@2 <= 15, projection=[l_extendedprice@2, l_discount@3] +04)------AggregateExec: mode=Partial, gby=[], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +05)--------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(l_partkey@0, p_partkey@0)], filter=p_brand@1 = Brand#12 AND p_container@3 IN (SET) ([SM CASE, SM BOX, SM PACK, SM PKG]) AND l_quantity@0 >= 1.00 AND l_quantity@0 <= 11.00 AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND p_container@3 IN (SET) ([MED BAG, MED BOX, MED PKG, MED PACK]) AND l_quantity@0 >= 10.00 AND l_quantity@0 <= 20.00 AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND p_container@3 IN (SET) ([LG CASE, LG BOX, LG PACK, LG PKG]) AND l_quantity@0 >= 20.00 AND l_quantity@0 <= 30.00 AND p_size@2 <= 15, projection=[l_extendedprice@2, l_discount@3] 06)----------RepartitionExec: partitioning=Hash([l_partkey@0], 4), input_partitions=4 -07)------------FilterExec: (l_shipmode@5 = AIR OR l_shipmode@5 = AIR REG) AND l_shipinstruct@4 = DELIVER IN PERSON AND (l_quantity@1 >= Some(100),15,2 AND l_quantity@1 <= Some(1100),15,2 OR l_quantity@1 >= Some(1000),15,2 AND l_quantity@1 <= Some(2000),15,2 OR l_quantity@1 >= Some(2000),15,2 AND l_quantity@1 <= Some(3000),15,2), projection=[l_partkey@0, l_quantity@1, l_extendedprice@2, l_discount@3] -08)--------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_quantity, l_extendedprice, l_discount, l_shipinstruct, l_shipmode], file_type=csv, has_header=false +07)------------FilterExec: (l_shipmode@5 = AIR OR l_shipmode@5 = AIR REG) AND l_shipinstruct@4 = DELIVER IN PERSON AND (l_quantity@1 >= 1.00 AND l_quantity@1 <= 11.00 OR l_quantity@1 >= 10.00 AND l_quantity@1 <= 20.00 OR l_quantity@1 >= 20.00 AND l_quantity@1 <= 30.00), projection=[l_partkey@0, l_quantity@1, l_extendedprice@2, l_discount@3] +08)--------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_quantity, l_extendedprice, l_discount, l_shipinstruct, l_shipmode], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 09)----------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 10)------------FilterExec: p_size@2 >= 1 AND (p_brand@1 = Brand#12 AND p_container@3 IN (SET) ([SM CASE, SM BOX, SM PACK, SM PKG]) AND p_size@2 <= 5 OR p_brand@1 = Brand#23 AND p_container@3 IN (SET) ([MED BAG, MED BOX, MED PKG, MED PACK]) AND p_size@2 <= 10 OR p_brand@1 = Brand#34 AND p_container@3 IN (SET) ([LG CASE, LG BOX, LG PACK, LG PKG]) AND p_size@2 <= 15) -11)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -12)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_brand, p_size, p_container], file_type=csv, has_header=false +11)--------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_brand, p_size, p_container], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q2.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q2.slt.part index b1a15388270b3..b6fa1c4806bf4 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q2.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q2.slt.part @@ -65,9 +65,9 @@ limit 10; logical_plan 01)Sort: supplier.s_acctbal DESC NULLS FIRST, nation.n_name ASC NULLS LAST, supplier.s_name ASC NULLS LAST, part.p_partkey ASC NULLS LAST, fetch=10 02)--Projection: supplier.s_acctbal, supplier.s_name, nation.n_name, part.p_partkey, part.p_mfgr, supplier.s_address, supplier.s_phone, supplier.s_comment -03)----Inner Join: part.p_partkey = __scalar_sq_1.ps_partkey, partsupp.ps_supplycost = __scalar_sq_1.min(partsupp.ps_supplycost) +03)----LeftSemi Join: part.p_partkey = __scalar_sq_1.ps_partkey, partsupp.ps_supplycost = __scalar_sq_1.min(partsupp.ps_supplycost) 04)------Projection: part.p_partkey, part.p_mfgr, supplier.s_name, supplier.s_address, supplier.s_phone, supplier.s_acctbal, supplier.s_comment, partsupp.ps_supplycost, nation.n_name -05)--------Inner Join: nation.n_regionkey = region.r_regionkey +05)--------LeftSemi Join: nation.n_regionkey = region.r_regionkey 06)----------Projection: part.p_partkey, part.p_mfgr, supplier.s_name, supplier.s_address, supplier.s_phone, supplier.s_acctbal, supplier.s_comment, partsupp.ps_supplycost, nation.n_name, nation.n_regionkey 07)------------Inner Join: supplier.s_nationkey = nation.n_nationkey 08)--------------Projection: part.p_partkey, part.p_mfgr, supplier.s_name, supplier.s_address, supplier.s_nationkey, supplier.s_phone, supplier.s_acctbal, supplier.s_comment, partsupp.ps_supplycost @@ -87,7 +87,7 @@ logical_plan 22)--------Projection: min(partsupp.ps_supplycost), partsupp.ps_partkey 23)----------Aggregate: groupBy=[[partsupp.ps_partkey]], aggr=[[min(partsupp.ps_supplycost)]] 24)------------Projection: partsupp.ps_partkey, partsupp.ps_supplycost -25)--------------Inner Join: nation.n_regionkey = region.r_regionkey +25)--------------LeftSemi Join: nation.n_regionkey = region.r_regionkey 26)----------------Projection: partsupp.ps_partkey, partsupp.ps_supplycost, nation.n_regionkey 27)------------------Inner Join: supplier.s_nationkey = nation.n_nationkey 28)--------------------Projection: partsupp.ps_partkey, partsupp.ps_supplycost, supplier.s_nationkey @@ -101,9 +101,9 @@ logical_plan physical_plan 01)SortPreservingMergeExec: [s_acctbal@0 DESC, n_name@2 ASC NULLS LAST, s_name@1 ASC NULLS LAST, p_partkey@3 ASC NULLS LAST], fetch=10 02)--SortExec: TopK(fetch=10), expr=[s_acctbal@0 DESC, n_name@2 ASC NULLS LAST, s_name@1 ASC NULLS LAST, p_partkey@3 ASC NULLS LAST], preserve_partitioning=[true] -03)----HashJoinExec: mode=Partitioned, join_type=Inner, on=[(p_partkey@0, ps_partkey@1), (ps_supplycost@7, min(partsupp.ps_supplycost)@0)], projection=[s_acctbal@5, s_name@2, n_name@8, p_partkey@0, p_mfgr@1, s_address@3, s_phone@4, s_comment@6] +03)----HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(p_partkey@0, ps_partkey@1), (ps_supplycost@7, min(partsupp.ps_supplycost)@0)], projection=[s_acctbal@5, s_name@2, n_name@8, p_partkey@0, p_mfgr@1, s_address@3, s_phone@4, s_comment@6] 04)------RepartitionExec: partitioning=Hash([p_partkey@0, ps_supplycost@7], 4), input_partitions=4 -05)--------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(n_regionkey@9, r_regionkey@0)], projection=[p_partkey@0, p_mfgr@1, s_name@2, s_address@3, s_phone@4, s_acctbal@5, s_comment@6, ps_supplycost@7, n_name@8] +05)--------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(n_regionkey@9, r_regionkey@0)], projection=[p_partkey@0, p_mfgr@1, s_name@2, s_address@3, s_phone@4, s_acctbal@5, s_comment@6, ps_supplycost@7, n_name@8] 06)----------RepartitionExec: partitioning=Hash([n_regionkey@9], 4), input_partitions=4 07)------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@4, n_nationkey@0)], projection=[p_partkey@0, p_mfgr@1, s_name@2, s_address@3, s_phone@5, s_acctbal@6, s_comment@7, ps_supplycost@8, n_name@10, n_regionkey@11] 08)--------------RepartitionExec: partitioning=Hash([s_nationkey@4], 4), input_partitions=4 @@ -112,35 +112,34 @@ physical_plan 11)--------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(p_partkey@0, ps_partkey@0)], projection=[p_partkey@0, p_mfgr@1, ps_suppkey@3, ps_supplycost@4] 12)----------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 13)------------------------FilterExec: p_size@3 = 15 AND p_type@2 LIKE %BRASS, projection=[p_partkey@0, p_mfgr@1] -14)--------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -15)----------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_mfgr, p_type, p_size], file_type=csv, has_header=false -16)----------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 -17)------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_supplycost], file_type=csv, has_header=false -18)------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -19)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_name, s_address, s_nationkey, s_phone, s_acctbal, s_comment], file_type=csv, has_header=false -20)--------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 -21)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name, n_regionkey], file_type=csv, has_header=false -22)----------RepartitionExec: partitioning=Hash([r_regionkey@0], 4), input_partitions=4 -23)------------FilterExec: r_name@1 = EUROPE, projection=[r_regionkey@0] -24)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -25)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/region.tbl]]}, projection=[r_regionkey, r_name], file_type=csv, has_header=false -26)------RepartitionExec: partitioning=Hash([ps_partkey@1, min(partsupp.ps_supplycost)@0], 4), input_partitions=4 -27)--------ProjectionExec: expr=[min(partsupp.ps_supplycost)@1 as min(partsupp.ps_supplycost), ps_partkey@0 as ps_partkey] -28)----------AggregateExec: mode=FinalPartitioned, gby=[ps_partkey@0 as ps_partkey], aggr=[min(partsupp.ps_supplycost)] -29)------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 -30)--------------AggregateExec: mode=Partial, gby=[ps_partkey@0 as ps_partkey], aggr=[min(partsupp.ps_supplycost)] -31)----------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(n_regionkey@2, r_regionkey@0)], projection=[ps_partkey@0, ps_supplycost@1] -32)------------------RepartitionExec: partitioning=Hash([n_regionkey@2], 4), input_partitions=4 -33)--------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@2, n_nationkey@0)], projection=[ps_partkey@0, ps_supplycost@1, n_regionkey@4] -34)----------------------RepartitionExec: partitioning=Hash([s_nationkey@2], 4), input_partitions=4 -35)------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_suppkey@1, s_suppkey@0)], projection=[ps_partkey@0, ps_supplycost@2, s_nationkey@4] -36)--------------------------RepartitionExec: partitioning=Hash([ps_suppkey@1], 4), input_partitions=4 -37)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_supplycost], file_type=csv, has_header=false -38)--------------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -39)----------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], file_type=csv, has_header=false -40)----------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 -41)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_regionkey], file_type=csv, has_header=false -42)------------------RepartitionExec: partitioning=Hash([r_regionkey@0], 4), input_partitions=4 -43)--------------------FilterExec: r_name@1 = EUROPE, projection=[r_regionkey@0] -44)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -45)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/region.tbl]]}, projection=[r_regionkey, r_name], file_type=csv, has_header=false +14)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_mfgr, p_type, p_size], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +15)----------------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 +16)------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_supplycost], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false +17)------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 +18)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_name, s_address, s_nationkey, s_phone, s_acctbal, s_comment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +19)--------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 +20)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name, n_regionkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +21)----------RepartitionExec: partitioning=Hash([r_regionkey@0], 4), input_partitions=4 +22)------------FilterExec: r_name@1 = EUROPE, projection=[r_regionkey@0] +23)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +24)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/region.tbl]]}, projection=[r_regionkey, r_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +25)------RepartitionExec: partitioning=Hash([ps_partkey@1, min(partsupp.ps_supplycost)@0], 4), input_partitions=4 +26)--------ProjectionExec: expr=[min(partsupp.ps_supplycost)@1 as min(partsupp.ps_supplycost), ps_partkey@0 as ps_partkey] +27)----------AggregateExec: mode=FinalPartitioned, gby=[ps_partkey@0 as ps_partkey], aggr=[min(partsupp.ps_supplycost)] +28)------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 +29)--------------AggregateExec: mode=Partial, gby=[ps_partkey@0 as ps_partkey], aggr=[min(partsupp.ps_supplycost)] +30)----------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(n_regionkey@2, r_regionkey@0)], projection=[ps_partkey@0, ps_supplycost@1] +31)------------------RepartitionExec: partitioning=Hash([n_regionkey@2], 4), input_partitions=4 +32)--------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@2, n_nationkey@0)], projection=[ps_partkey@0, ps_supplycost@1, n_regionkey@4] +33)----------------------RepartitionExec: partitioning=Hash([s_nationkey@2], 4), input_partitions=4 +34)------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_suppkey@1, s_suppkey@0)], projection=[ps_partkey@0, ps_supplycost@2, s_nationkey@4] +35)--------------------------RepartitionExec: partitioning=Hash([ps_suppkey@1], 4), input_partitions=4 +36)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_supplycost], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false +37)--------------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 +38)----------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +39)----------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 +40)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_regionkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +41)------------------RepartitionExec: partitioning=Hash([r_regionkey@0], 4), input_partitions=4 +42)--------------------FilterExec: r_name@1 = EUROPE, projection=[r_regionkey@0] +43)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +44)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/region.tbl]]}, projection=[r_regionkey, r_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q20.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q20.slt.part index 426a1cbaa4e22..e038a7482d24f 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q20.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q20.slt.part @@ -60,14 +60,14 @@ logical_plan 02)--Projection: supplier.s_name, supplier.s_address 03)----LeftSemi Join: supplier.s_suppkey = __correlated_sq_2.ps_suppkey 04)------Projection: supplier.s_suppkey, supplier.s_name, supplier.s_address -05)--------Inner Join: supplier.s_nationkey = nation.n_nationkey +05)--------LeftSemi Join: supplier.s_nationkey = nation.n_nationkey 06)----------TableScan: supplier projection=[s_suppkey, s_name, s_address, s_nationkey] 07)----------Projection: nation.n_nationkey 08)------------Filter: nation.n_name = Utf8View("CANADA") 09)--------------TableScan: nation projection=[n_nationkey, n_name], partial_filters=[nation.n_name = Utf8View("CANADA")] 10)------SubqueryAlias: __correlated_sq_2 11)--------Projection: partsupp.ps_suppkey -12)----------Inner Join: partsupp.ps_partkey = __scalar_sq_3.l_partkey, partsupp.ps_suppkey = __scalar_sq_3.l_suppkey Filter: CAST(partsupp.ps_availqty AS Float64) > __scalar_sq_3.Float64(0.5) * sum(lineitem.l_quantity) +12)----------LeftSemi Join: partsupp.ps_partkey = __scalar_sq_3.l_partkey, partsupp.ps_suppkey = __scalar_sq_3.l_suppkey Filter: CAST(partsupp.ps_availqty AS Float64) > __scalar_sq_3.Float64(0.5) * sum(lineitem.l_quantity) 13)------------LeftSemi Join: partsupp.ps_partkey = __correlated_sq_1.p_partkey 14)--------------TableScan: partsupp projection=[ps_partkey, ps_suppkey, ps_availqty] 15)--------------SubqueryAlias: __correlated_sq_1 @@ -85,26 +85,25 @@ physical_plan 02)--SortExec: expr=[s_name@0 ASC NULLS LAST], preserve_partitioning=[true] 03)----HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(s_suppkey@0, ps_suppkey@0)], projection=[s_name@1, s_address@2] 04)------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=4 -05)--------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@3, n_nationkey@0)], projection=[s_suppkey@0, s_name@1, s_address@2] +05)--------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(s_nationkey@3, n_nationkey@0)], projection=[s_suppkey@0, s_name@1, s_address@2] 06)----------RepartitionExec: partitioning=Hash([s_nationkey@3], 4), input_partitions=1 -07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_name, s_address, s_nationkey], file_type=csv, has_header=false +07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_name, s_address, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 08)----------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=4 09)------------FilterExec: n_name@1 = CANADA, projection=[n_nationkey@0] 10)--------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -11)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false +11)----------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 12)------RepartitionExec: partitioning=Hash([ps_suppkey@0], 4), input_partitions=4 -13)--------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(ps_partkey@0, l_partkey@1), (ps_suppkey@1, l_suppkey@2)], filter=CAST(ps_availqty@0 AS Float64) > Float64(0.5) * sum(lineitem.l_quantity)@1, projection=[ps_suppkey@1] +13)--------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(ps_partkey@0, l_partkey@1), (ps_suppkey@1, l_suppkey@2)], filter=CAST(ps_availqty@0 AS Float64) > Float64(0.5) * sum(lineitem.l_quantity)@1, projection=[ps_suppkey@1] 14)----------RepartitionExec: partitioning=Hash([ps_partkey@0, ps_suppkey@1], 4), input_partitions=4 15)------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(ps_partkey@0, p_partkey@0)] 16)--------------RepartitionExec: partitioning=Hash([ps_partkey@0], 4), input_partitions=4 -17)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_availqty], file_type=csv, has_header=false +17)----------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_availqty], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false 18)--------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 19)----------------FilterExec: p_name@1 LIKE forest%, projection=[p_partkey@0] -20)------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -21)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_name], file_type=csv, has_header=false -22)----------ProjectionExec: expr=[0.5 * CAST(sum(lineitem.l_quantity)@2 AS Float64) as Float64(0.5) * sum(lineitem.l_quantity), l_partkey@0 as l_partkey, l_suppkey@1 as l_suppkey] -23)------------AggregateExec: mode=FinalPartitioned, gby=[l_partkey@0 as l_partkey, l_suppkey@1 as l_suppkey], aggr=[sum(lineitem.l_quantity)] -24)--------------RepartitionExec: partitioning=Hash([l_partkey@0, l_suppkey@1], 4), input_partitions=4 -25)----------------AggregateExec: mode=Partial, gby=[l_partkey@0 as l_partkey, l_suppkey@1 as l_suppkey], aggr=[sum(lineitem.l_quantity)] -26)------------------FilterExec: l_shipdate@3 >= 1994-01-01 AND l_shipdate@3 < 1995-01-01, projection=[l_partkey@0, l_suppkey@1, l_quantity@2] -27)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_suppkey, l_quantity, l_shipdate], file_type=csv, has_header=false +20)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +21)----------ProjectionExec: expr=[0.5 * CAST(sum(lineitem.l_quantity)@2 AS Float64) as Float64(0.5) * sum(lineitem.l_quantity), l_partkey@0 as l_partkey, l_suppkey@1 as l_suppkey] +22)------------AggregateExec: mode=FinalPartitioned, gby=[l_partkey@0 as l_partkey, l_suppkey@1 as l_suppkey], aggr=[sum(lineitem.l_quantity)] +23)--------------RepartitionExec: partitioning=Hash([l_partkey@0, l_suppkey@1], 4), input_partitions=4 +24)----------------AggregateExec: mode=Partial, gby=[l_partkey@0 as l_partkey, l_suppkey@1 as l_suppkey], aggr=[sum(lineitem.l_quantity)] +25)------------------FilterExec: l_shipdate@3 >= 1994-01-01 AND l_shipdate@3 < 1995-01-01, projection=[l_partkey@0, l_suppkey@1, l_quantity@2] +26)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_partkey, l_suppkey, l_quantity, l_shipdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q21.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q21.slt.part index 5e9192d677532..47e5d6d888dc5 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q21.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q21.slt.part @@ -65,59 +65,58 @@ logical_plan 05)--------LeftAnti Join: l1.l_orderkey = __correlated_sq_2.l_orderkey Filter: __correlated_sq_2.l_suppkey != l1.l_suppkey 06)----------LeftSemi Join: l1.l_orderkey = __correlated_sq_1.l_orderkey Filter: __correlated_sq_1.l_suppkey != l1.l_suppkey 07)------------Projection: supplier.s_name, l1.l_orderkey, l1.l_suppkey -08)--------------Inner Join: supplier.s_nationkey = nation.n_nationkey -09)----------------Projection: supplier.s_name, supplier.s_nationkey, l1.l_orderkey, l1.l_suppkey -10)------------------Inner Join: l1.l_orderkey = orders.o_orderkey -11)--------------------Projection: supplier.s_name, supplier.s_nationkey, l1.l_orderkey, l1.l_suppkey -12)----------------------Inner Join: supplier.s_suppkey = l1.l_suppkey -13)------------------------TableScan: supplier projection=[s_suppkey, s_name, s_nationkey] -14)------------------------SubqueryAlias: l1 -15)--------------------------Projection: lineitem.l_orderkey, lineitem.l_suppkey -16)----------------------------Filter: lineitem.l_receiptdate > lineitem.l_commitdate -17)------------------------------TableScan: lineitem projection=[l_orderkey, l_suppkey, l_commitdate, l_receiptdate], partial_filters=[lineitem.l_receiptdate > lineitem.l_commitdate] -18)--------------------Projection: orders.o_orderkey -19)----------------------Filter: orders.o_orderstatus = Utf8View("F") -20)------------------------TableScan: orders projection=[o_orderkey, o_orderstatus], partial_filters=[orders.o_orderstatus = Utf8View("F")] -21)----------------Projection: nation.n_nationkey -22)------------------Filter: nation.n_name = Utf8View("SAUDI ARABIA") -23)--------------------TableScan: nation projection=[n_nationkey, n_name], partial_filters=[nation.n_name = Utf8View("SAUDI ARABIA")] -24)------------SubqueryAlias: __correlated_sq_1 -25)--------------SubqueryAlias: l2 -26)----------------TableScan: lineitem projection=[l_orderkey, l_suppkey] -27)----------SubqueryAlias: __correlated_sq_2 -28)------------SubqueryAlias: l3 -29)--------------Projection: lineitem.l_orderkey, lineitem.l_suppkey -30)----------------Filter: lineitem.l_receiptdate > lineitem.l_commitdate -31)------------------TableScan: lineitem projection=[l_orderkey, l_suppkey, l_commitdate, l_receiptdate], partial_filters=[lineitem.l_receiptdate > lineitem.l_commitdate] +08)--------------LeftSemi Join: supplier.s_nationkey = nation.n_nationkey +09)----------------LeftSemi Join: l1.l_orderkey = orders.o_orderkey +10)------------------Projection: supplier.s_name, supplier.s_nationkey, l1.l_orderkey, l1.l_suppkey +11)--------------------Inner Join: supplier.s_suppkey = l1.l_suppkey +12)----------------------TableScan: supplier projection=[s_suppkey, s_name, s_nationkey] +13)----------------------SubqueryAlias: l1 +14)------------------------Projection: lineitem.l_orderkey, lineitem.l_suppkey +15)--------------------------Filter: lineitem.l_receiptdate > lineitem.l_commitdate +16)----------------------------TableScan: lineitem projection=[l_orderkey, l_suppkey, l_commitdate, l_receiptdate], partial_filters=[lineitem.l_receiptdate > lineitem.l_commitdate] +17)------------------Projection: orders.o_orderkey +18)--------------------Filter: orders.o_orderstatus = Utf8View("F") +19)----------------------TableScan: orders projection=[o_orderkey, o_orderstatus], partial_filters=[orders.o_orderstatus = Utf8View("F")] +20)----------------Projection: nation.n_nationkey +21)------------------Filter: nation.n_name = Utf8View("SAUDI ARABIA") +22)--------------------TableScan: nation projection=[n_nationkey, n_name], partial_filters=[nation.n_name = Utf8View("SAUDI ARABIA")] +23)------------SubqueryAlias: __correlated_sq_1 +24)--------------SubqueryAlias: l2 +25)----------------TableScan: lineitem projection=[l_orderkey, l_suppkey] +26)----------SubqueryAlias: __correlated_sq_2 +27)------------SubqueryAlias: l3 +28)--------------Projection: lineitem.l_orderkey, lineitem.l_suppkey +29)----------------Filter: lineitem.l_receiptdate > lineitem.l_commitdate +30)------------------TableScan: lineitem projection=[l_orderkey, l_suppkey, l_commitdate, l_receiptdate], partial_filters=[lineitem.l_receiptdate > lineitem.l_commitdate] physical_plan 01)SortPreservingMergeExec: [numwait@1 DESC, s_name@0 ASC NULLS LAST] -02)--SortExec: expr=[numwait@1 DESC, s_name@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[s_name@0 as s_name, count(Int64(1))@1 as numwait] +02)--ProjectionExec: expr=[s_name@0 as s_name, count(Int64(1))@1 as numwait] +03)----SortExec: expr=[count(Int64(1))@1 DESC, s_name@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[s_name@0 as s_name], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([s_name@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[s_name@0 as s_name], aggr=[count(Int64(1))] 07)------------HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(l_orderkey@1, l_orderkey@0)], filter=l_suppkey@1 != l_suppkey@0, projection=[s_name@0] 08)--------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(l_orderkey@1, l_orderkey@0)], filter=l_suppkey@1 != l_suppkey@0 09)----------------RepartitionExec: partitioning=Hash([l_orderkey@1], 4), input_partitions=4 -10)------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@1, n_nationkey@0)], projection=[s_name@0, l_orderkey@2, l_suppkey@3] +10)------------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(s_nationkey@1, n_nationkey@0)], projection=[s_name@0, l_orderkey@2, l_suppkey@3] 11)--------------------RepartitionExec: partitioning=Hash([s_nationkey@1], 4), input_partitions=4 -12)----------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_orderkey@2, o_orderkey@0)], projection=[s_name@0, s_nationkey@1, l_orderkey@2, l_suppkey@3] +12)----------------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(l_orderkey@2, o_orderkey@0)] 13)------------------------RepartitionExec: partitioning=Hash([l_orderkey@2], 4), input_partitions=4 14)--------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_suppkey@0, l_suppkey@1)], projection=[s_name@1, s_nationkey@2, l_orderkey@3, l_suppkey@4] 15)----------------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -16)------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_name, s_nationkey], file_type=csv, has_header=false +16)------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_name, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 17)----------------------------RepartitionExec: partitioning=Hash([l_suppkey@1], 4), input_partitions=4 18)------------------------------FilterExec: l_receiptdate@3 > l_commitdate@2, projection=[l_orderkey@0, l_suppkey@1] -19)--------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey, l_commitdate, l_receiptdate], file_type=csv, has_header=false +19)--------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey, l_commitdate, l_receiptdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 20)------------------------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 21)--------------------------FilterExec: o_orderstatus@1 = F, projection=[o_orderkey@0] -22)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_orderstatus], file_type=csv, has_header=false +22)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_orderstatus], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 23)--------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=4 24)----------------------FilterExec: n_name@1 = SAUDI ARABIA, projection=[n_nationkey@0] 25)------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -26)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false +26)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 27)----------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 -28)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey], file_type=csv, has_header=false +28)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 29)--------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 30)----------------FilterExec: l_receiptdate@3 > l_commitdate@2, projection=[l_orderkey@0, l_suppkey@1] -31)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey, l_commitdate, l_receiptdate], file_type=csv, has_header=false +31)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey, l_commitdate, l_receiptdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part index 3e9472e4a8867..d3f27021f1781 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q22.slt.part @@ -66,16 +66,16 @@ logical_plan 08)--------------Subquery: 09)----------------Aggregate: groupBy=[[]], aggr=[[avg(customer.c_acctbal)]] 10)------------------Projection: customer.c_acctbal -11)--------------------Filter: customer.c_acctbal > Decimal128(Some(0),15,2) AND substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")]) -12)----------------------TableScan: customer projection=[c_phone, c_acctbal], partial_filters=[customer.c_acctbal > Decimal128(Some(0),15,2), substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")])] +11)--------------------Filter: customer.c_acctbal > Decimal128(0.00,15,2) AND substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")]) +12)----------------------TableScan: customer projection=[c_phone, c_acctbal], partial_filters=[customer.c_acctbal > Decimal128(0.00,15,2), substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")])] 13)--------------TableScan: customer projection=[c_custkey, c_phone, c_acctbal], partial_filters=[substr(customer.c_phone, Int64(1), Int64(2)) IN ([Utf8View("13"), Utf8View("31"), Utf8View("23"), Utf8View("29"), Utf8View("30"), Utf8View("18"), Utf8View("17")])] 14)------------SubqueryAlias: __correlated_sq_1 15)--------------TableScan: orders projection=[o_custkey] physical_plan 01)ScalarSubqueryExec: subqueries=1 02)--SortPreservingMergeExec: [cntrycode@0 ASC NULLS LAST] -03)----SortExec: expr=[cntrycode@0 ASC NULLS LAST], preserve_partitioning=[true] -04)------ProjectionExec: expr=[cntrycode@0 as cntrycode, count(Int64(1))@1 as numcust, sum(custsale.c_acctbal)@2 as totacctbal] +03)----ProjectionExec: expr=[cntrycode@0 as cntrycode, count(Int64(1))@1 as numcust, sum(custsale.c_acctbal)@2 as totacctbal] +04)------SortExec: expr=[cntrycode@0 ASC NULLS LAST], preserve_partitioning=[true] 05)--------AggregateExec: mode=FinalPartitioned, gby=[cntrycode@0 as cntrycode], aggr=[count(Int64(1)), sum(custsale.c_acctbal)] 06)----------RepartitionExec: partitioning=Hash([cntrycode@0], 4), input_partitions=4 07)------------AggregateExec: mode=Partial, gby=[cntrycode@0 as cntrycode], aggr=[count(Int64(1)), sum(custsale.c_acctbal)] @@ -83,13 +83,11 @@ physical_plan 09)----------------HashJoinExec: mode=Partitioned, join_type=LeftAnti, on=[(c_custkey@0, o_custkey@0)], projection=[c_phone@1, c_acctbal@2] 10)------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4 11)--------------------FilterExec: substr(c_phone@1, 1, 2) IN (SET) ([13, 31, 23, 29, 30, 18, 17]) AND CAST(c_acctbal@2 AS Decimal128(19, 6)) > scalar_subquery() -12)----------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -13)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_phone, c_acctbal], file_type=csv, has_header=false -14)------------------RepartitionExec: partitioning=Hash([o_custkey@0], 4), input_partitions=4 -15)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_custkey], file_type=csv, has_header=false -16)--AggregateExec: mode=Final, gby=[], aggr=[avg(customer.c_acctbal)] -17)----CoalescePartitionsExec -18)------AggregateExec: mode=Partial, gby=[], aggr=[avg(customer.c_acctbal)] -19)--------FilterExec: c_acctbal@1 > Some(0),15,2 AND substr(c_phone@0, 1, 2) IN (SET) ([13, 31, 23, 29, 30, 18, 17]), projection=[c_acctbal@1] -20)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -21)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_phone, c_acctbal], file_type=csv, has_header=false +12)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_custkey, c_phone, c_acctbal], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +13)------------------RepartitionExec: partitioning=Hash([o_custkey@0], 4), input_partitions=4 +14)--------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_custkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +15)--AggregateExec: mode=Final, gby=[], aggr=[avg(customer.c_acctbal)] +16)----CoalescePartitionsExec +17)------AggregateExec: mode=Partial, gby=[], aggr=[avg(customer.c_acctbal)] +18)--------FilterExec: c_acctbal@1 > 0.00 AND substr(c_phone@0, 1, 2) IN (SET) ([13, 31, 23, 29, 30, 18, 17]), projection=[c_acctbal@1] +19)----------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_phone, c_acctbal], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q3.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q3.slt.part index ba56f10fab25f..724ef72ca324c 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q3.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q3.slt.part @@ -44,11 +44,11 @@ limit 10; logical_plan 01)Sort: revenue DESC NULLS FIRST, orders.o_orderdate ASC NULLS LAST, fetch=10 02)--Projection: lineitem.l_orderkey, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS revenue, orders.o_orderdate, orders.o_shippriority -03)----Aggregate: groupBy=[[lineitem.l_orderkey, orders.o_orderdate, orders.o_shippriority]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] +03)----Aggregate: groupBy=[[lineitem.l_orderkey, orders.o_orderdate, orders.o_shippriority]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] 04)------Projection: orders.o_orderdate, orders.o_shippriority, lineitem.l_orderkey, lineitem.l_extendedprice, lineitem.l_discount 05)--------Inner Join: orders.o_orderkey = lineitem.l_orderkey 06)----------Projection: orders.o_orderkey, orders.o_orderdate, orders.o_shippriority -07)------------Inner Join: customer.c_custkey = orders.o_custkey +07)------------RightSemi Join: customer.c_custkey = orders.o_custkey 08)--------------Projection: customer.c_custkey 09)----------------Filter: customer.c_mktsegment = Utf8View("BUILDING") 10)------------------TableScan: customer projection=[c_custkey, c_mktsegment], partial_filters=[customer.c_mktsegment = Utf8View("BUILDING")] @@ -59,19 +59,18 @@ logical_plan 15)--------------TableScan: lineitem projection=[l_orderkey, l_extendedprice, l_discount, l_shipdate], partial_filters=[lineitem.l_shipdate > Date32("1995-03-15")] physical_plan 01)SortPreservingMergeExec: [revenue@1 DESC, o_orderdate@2 ASC NULLS LAST], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[revenue@1 DESC, o_orderdate@2 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[l_orderkey@0 as l_orderkey, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@3 as revenue, o_orderdate@1 as o_orderdate, o_shippriority@2 as o_shippriority] -04)------AggregateExec: mode=SinglePartitioned, gby=[l_orderkey@2 as l_orderkey, o_orderdate@0 as o_orderdate, o_shippriority@1 as o_shippriority], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +02)--ProjectionExec: expr=[l_orderkey@0 as l_orderkey, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@3 as revenue, o_orderdate@1 as o_orderdate, o_shippriority@2 as o_shippriority] +03)----SortExec: TopK(fetch=10), expr=[sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@3 DESC, o_orderdate@1 ASC NULLS LAST], preserve_partitioning=[true] +04)------AggregateExec: mode=SinglePartitioned, gby=[l_orderkey@2 as l_orderkey, o_orderdate@0 as o_orderdate, o_shippriority@1 as o_shippriority], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 05)--------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(o_orderkey@0, l_orderkey@0)], projection=[o_orderdate@1, o_shippriority@2, l_orderkey@3, l_extendedprice@4, l_discount@5] 06)----------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 -07)------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c_custkey@0, o_custkey@1)], projection=[o_orderkey@1, o_orderdate@3, o_shippriority@4] +07)------------HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(c_custkey@0, o_custkey@1)], projection=[o_orderkey@0, o_orderdate@2, o_shippriority@3] 08)--------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4 09)----------------FilterExec: c_mktsegment@1 = BUILDING, projection=[c_custkey@0] -10)------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -11)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_mktsegment], file_type=csv, has_header=false -12)--------------RepartitionExec: partitioning=Hash([o_custkey@1], 4), input_partitions=4 -13)----------------FilterExec: o_orderdate@2 < 1995-03-15 -14)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_orderdate, o_shippriority], file_type=csv, has_header=false -15)----------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 -16)------------FilterExec: l_shipdate@3 > 1995-03-15, projection=[l_orderkey@0, l_extendedprice@1, l_discount@2] -17)--------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_extendedprice, l_discount, l_shipdate], file_type=csv, has_header=false +10)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_custkey, c_mktsegment], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +11)--------------RepartitionExec: partitioning=Hash([o_custkey@1], 4), input_partitions=4 +12)----------------FilterExec: o_orderdate@2 < 1995-03-15 +13)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_orderdate, o_shippriority], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +14)----------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 +15)------------FilterExec: l_shipdate@3 > 1995-03-15, projection=[l_orderkey@0, l_extendedprice@1, l_discount@2] +16)--------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_extendedprice, l_discount, l_shipdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q4.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q4.slt.part index 0007666f15365..470d7a6527a52 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q4.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q4.slt.part @@ -54,15 +54,15 @@ logical_plan 12)----------------TableScan: lineitem projection=[l_orderkey, l_commitdate, l_receiptdate], partial_filters=[lineitem.l_receiptdate > lineitem.l_commitdate] physical_plan 01)SortPreservingMergeExec: [o_orderpriority@0 ASC NULLS LAST] -02)--SortExec: expr=[o_orderpriority@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[o_orderpriority@0 as o_orderpriority, count(Int64(1))@1 as order_count] +02)--ProjectionExec: expr=[o_orderpriority@0 as o_orderpriority, count(Int64(1))@1 as order_count] +03)----SortExec: expr=[o_orderpriority@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[o_orderpriority@0 as o_orderpriority], aggr=[count(Int64(1))] 05)--------RepartitionExec: partitioning=Hash([o_orderpriority@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[o_orderpriority@0 as o_orderpriority], aggr=[count(Int64(1))] 07)------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(o_orderkey@0, l_orderkey@0)], projection=[o_orderpriority@1] 08)--------------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 09)----------------FilterExec: o_orderdate@1 >= 1993-07-01 AND o_orderdate@1 < 1993-10-01, projection=[o_orderkey@0, o_orderpriority@2] -10)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_orderdate, o_orderpriority], file_type=csv, has_header=false +10)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_orderdate, o_orderpriority], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 11)--------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 12)----------------FilterExec: l_receiptdate@2 > l_commitdate@1, projection=[l_orderkey@0] -13)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_commitdate, l_receiptdate], file_type=csv, has_header=false +13)------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_commitdate, l_receiptdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q5.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q5.slt.part index bda0586963159..0c4bdfda8daf0 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q5.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q5.slt.part @@ -45,9 +45,9 @@ order by logical_plan 01)Sort: revenue DESC NULLS FIRST 02)--Projection: nation.n_name, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount) AS revenue -03)----Aggregate: groupBy=[[nation.n_name]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] +03)----Aggregate: groupBy=[[nation.n_name]], aggr=[[sum(lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount)) AS sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)]] 04)------Projection: lineitem.l_extendedprice, lineitem.l_discount, nation.n_name -05)--------Inner Join: nation.n_regionkey = region.r_regionkey +05)--------LeftSemi Join: nation.n_regionkey = region.r_regionkey 06)----------Projection: lineitem.l_extendedprice, lineitem.l_discount, nation.n_name, nation.n_regionkey 07)------------Inner Join: supplier.s_nationkey = nation.n_nationkey 08)--------------Projection: lineitem.l_extendedprice, lineitem.l_discount, supplier.s_nationkey @@ -68,12 +68,12 @@ logical_plan 23)--------------TableScan: region projection=[r_regionkey, r_name], partial_filters=[region.r_name = Utf8View("ASIA")] physical_plan 01)SortPreservingMergeExec: [revenue@1 DESC] -02)--SortExec: expr=[revenue@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[n_name@0 as n_name, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 as revenue] -04)------AggregateExec: mode=FinalPartitioned, gby=[n_name@0 as n_name], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +02)--ProjectionExec: expr=[n_name@0 as n_name, sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 as revenue] +03)----SortExec: expr=[sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)@1 DESC], preserve_partitioning=[true] +04)------AggregateExec: mode=FinalPartitioned, gby=[n_name@0 as n_name], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] 05)--------RepartitionExec: partitioning=Hash([n_name@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[n_name@2 as n_name], aggr=[sum(lineitem.l_extendedprice * Some(1),20,0 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] -07)------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(n_regionkey@3, r_regionkey@0)], projection=[l_extendedprice@0, l_discount@1, n_name@2] +06)----------AggregateExec: mode=Partial, gby=[n_name@2 as n_name], aggr=[sum(lineitem.l_extendedprice * 1 - lineitem.l_discount) as sum(lineitem.l_extendedprice * Int64(1) - lineitem.l_discount)] +07)------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(n_regionkey@3, r_regionkey@0)], projection=[l_extendedprice@0, l_discount@1, n_name@2] 08)--------------RepartitionExec: partitioning=Hash([n_regionkey@3], 4), input_partitions=4 09)----------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@2, n_nationkey@0)], projection=[l_extendedprice@0, l_discount@1, n_name@4, n_regionkey@5] 10)------------------RepartitionExec: partitioning=Hash([s_nationkey@2], 4), input_partitions=4 @@ -82,18 +82,18 @@ physical_plan 13)------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(o_orderkey@1, l_orderkey@0)], projection=[c_nationkey@0, l_suppkey@3, l_extendedprice@4, l_discount@5] 14)--------------------------RepartitionExec: partitioning=Hash([o_orderkey@1], 4), input_partitions=4 15)----------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c_custkey@0, o_custkey@1)], projection=[c_nationkey@1, o_orderkey@2] -16)------------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=1 -17)--------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_nationkey], file_type=csv, has_header=false +16)------------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4 +17)--------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_custkey, c_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 18)------------------------------RepartitionExec: partitioning=Hash([o_custkey@1], 4), input_partitions=4 19)--------------------------------FilterExec: o_orderdate@2 >= 1994-01-01 AND o_orderdate@2 < 1995-01-01, projection=[o_orderkey@0, o_custkey@1] -20)----------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_orderdate], file_type=csv, has_header=false +20)----------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_orderdate], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 21)--------------------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 -22)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey, l_extendedprice, l_discount], file_type=csv, has_header=false +22)----------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey, l_extendedprice, l_discount], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 23)----------------------RepartitionExec: partitioning=Hash([s_suppkey@0, s_nationkey@1], 4), input_partitions=1 -24)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], file_type=csv, has_header=false +24)------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 25)------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 -26)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name, n_regionkey], file_type=csv, has_header=false +26)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name, n_regionkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 27)--------------RepartitionExec: partitioning=Hash([r_regionkey@0], 4), input_partitions=4 28)----------------FilterExec: r_name@1 = ASIA, projection=[r_regionkey@0] 29)------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -30)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/region.tbl]]}, projection=[r_regionkey, r_name], file_type=csv, has_header=false +30)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/region.tbl]]}, projection=[r_regionkey, r_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q6.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q6.slt.part index eb9063d691712..02a716557d039 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q6.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q6.slt.part @@ -31,12 +31,12 @@ logical_plan 01)Projection: sum(lineitem.l_extendedprice * lineitem.l_discount) AS revenue 02)--Aggregate: groupBy=[[]], aggr=[[sum(lineitem.l_extendedprice * lineitem.l_discount)]] 03)----Projection: lineitem.l_extendedprice, lineitem.l_discount -04)------Filter: lineitem.l_shipdate >= Date32("1994-01-01") AND lineitem.l_shipdate < Date32("1995-01-01") AND lineitem.l_discount >= Decimal128(Some(5),15,2) AND lineitem.l_discount <= Decimal128(Some(7),15,2) AND lineitem.l_quantity < Decimal128(Some(2400),15,2) -05)--------TableScan: lineitem projection=[l_quantity, l_extendedprice, l_discount, l_shipdate], partial_filters=[lineitem.l_shipdate >= Date32("1994-01-01"), lineitem.l_shipdate < Date32("1995-01-01"), lineitem.l_discount >= Decimal128(Some(5),15,2), lineitem.l_discount <= Decimal128(Some(7),15,2), lineitem.l_quantity < Decimal128(Some(2400),15,2)] +04)------Filter: lineitem.l_shipdate >= Date32("1994-01-01") AND lineitem.l_shipdate < Date32("1995-01-01") AND lineitem.l_discount >= Decimal128(0.05,15,2) AND lineitem.l_discount <= Decimal128(0.07,15,2) AND lineitem.l_quantity < Decimal128(24.00,15,2) +05)--------TableScan: lineitem projection=[l_quantity, l_extendedprice, l_discount, l_shipdate], partial_filters=[lineitem.l_shipdate >= Date32("1994-01-01"), lineitem.l_shipdate < Date32("1995-01-01"), lineitem.l_discount >= Decimal128(0.05,15,2), lineitem.l_discount <= Decimal128(0.07,15,2), lineitem.l_quantity < Decimal128(24.00,15,2)] physical_plan 01)ProjectionExec: expr=[sum(lineitem.l_extendedprice * lineitem.l_discount)@0 as revenue] 02)--AggregateExec: mode=Final, gby=[], aggr=[sum(lineitem.l_extendedprice * lineitem.l_discount)] 03)----CoalescePartitionsExec 04)------AggregateExec: mode=Partial, gby=[], aggr=[sum(lineitem.l_extendedprice * lineitem.l_discount)] -05)--------FilterExec: l_shipdate@3 >= 1994-01-01 AND l_shipdate@3 < 1995-01-01 AND l_discount@2 >= Some(5),15,2 AND l_discount@2 <= Some(7),15,2 AND l_quantity@0 < Some(2400),15,2, projection=[l_extendedprice@1, l_discount@2] -06)----------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_quantity, l_extendedprice, l_discount, l_shipdate], file_type=csv, has_header=false +05)--------FilterExec: l_shipdate@3 >= 1994-01-01 AND l_shipdate@3 < 1995-01-01 AND l_discount@2 >= 0.05 AND l_discount@2 <= 0.07 AND l_quantity@0 < 24.00, projection=[l_extendedprice@1, l_discount@2] +06)----------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_quantity, l_extendedprice, l_discount, l_shipdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q7.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q7.slt.part index 590a737703847..0db80ae202658 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q7.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q7.slt.part @@ -62,7 +62,7 @@ logical_plan 02)--Projection: shipping.supp_nation, shipping.cust_nation, shipping.l_year, sum(shipping.volume) AS revenue 03)----Aggregate: groupBy=[[shipping.supp_nation, shipping.cust_nation, shipping.l_year]], aggr=[[sum(shipping.volume)]] 04)------SubqueryAlias: shipping -05)--------Projection: n1.n_name AS supp_nation, n2.n_name AS cust_nation, date_part(Utf8("YEAR"), lineitem.l_shipdate) AS l_year, lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount) AS volume +05)--------Projection: n1.n_name AS supp_nation, n2.n_name AS cust_nation, date_part(Utf8("YEAR"), lineitem.l_shipdate) AS l_year, lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount) AS volume 06)----------Inner Join: customer.c_nationkey = n2.n_nationkey Filter: n1.n_name = Utf8View("FRANCE") AND n2.n_name = Utf8View("GERMANY") OR n1.n_name = Utf8View("GERMANY") AND n2.n_name = Utf8View("FRANCE") 07)------------Projection: lineitem.l_extendedprice, lineitem.l_discount, lineitem.l_shipdate, customer.c_nationkey, n1.n_name 08)--------------Inner Join: supplier.s_nationkey = n1.n_nationkey @@ -85,12 +85,12 @@ logical_plan 25)----------------TableScan: nation projection=[n_nationkey, n_name], partial_filters=[nation.n_name = Utf8View("GERMANY") OR nation.n_name = Utf8View("FRANCE")] physical_plan 01)SortPreservingMergeExec: [supp_nation@0 ASC NULLS LAST, cust_nation@1 ASC NULLS LAST, l_year@2 ASC NULLS LAST] -02)--SortExec: expr=[supp_nation@0 ASC NULLS LAST, cust_nation@1 ASC NULLS LAST, l_year@2 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[supp_nation@0 as supp_nation, cust_nation@1 as cust_nation, l_year@2 as l_year, sum(shipping.volume)@3 as revenue] +02)--ProjectionExec: expr=[supp_nation@0 as supp_nation, cust_nation@1 as cust_nation, l_year@2 as l_year, sum(shipping.volume)@3 as revenue] +03)----SortExec: expr=[supp_nation@0 ASC NULLS LAST, cust_nation@1 ASC NULLS LAST, l_year@2 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[supp_nation@0 as supp_nation, cust_nation@1 as cust_nation, l_year@2 as l_year], aggr=[sum(shipping.volume)] 05)--------RepartitionExec: partitioning=Hash([supp_nation@0, cust_nation@1, l_year@2], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[supp_nation@0 as supp_nation, cust_nation@1 as cust_nation, l_year@2 as l_year], aggr=[sum(shipping.volume)] -07)------------ProjectionExec: expr=[n_name@0 as supp_nation, n_name@1 as cust_nation, date_part(YEAR, l_shipdate@2) as l_year, l_extendedprice@3 * (Some(1),20,0 - l_discount@4) as volume] +07)------------ProjectionExec: expr=[n_name@0 as supp_nation, n_name@1 as cust_nation, date_part(YEAR, l_shipdate@2) as l_year, l_extendedprice@3 * (1 - l_discount@4) as volume] 08)--------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(c_nationkey@3, n_nationkey@0)], filter=n_name@0 = FRANCE AND n_name@1 = GERMANY OR n_name@0 = GERMANY AND n_name@1 = FRANCE, projection=[n_name@4, n_name@6, l_shipdate@2, l_extendedprice@0, l_discount@1] 09)----------------RepartitionExec: partitioning=Hash([c_nationkey@3], 4), input_partitions=4 10)------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@0, n_nationkey@0)], projection=[l_extendedprice@1, l_discount@2, l_shipdate@3, c_nationkey@4, n_name@6] @@ -101,19 +101,19 @@ physical_plan 15)----------------------------RepartitionExec: partitioning=Hash([l_orderkey@1], 4), input_partitions=4 16)------------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_suppkey@0, l_suppkey@1)], projection=[s_nationkey@1, l_orderkey@2, l_extendedprice@4, l_discount@5, l_shipdate@6] 17)--------------------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -18)----------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], file_type=csv, has_header=false +18)----------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 19)--------------------------------RepartitionExec: partitioning=Hash([l_suppkey@1], 4), input_partitions=4 20)----------------------------------FilterExec: l_shipdate@4 >= 1995-01-01 AND l_shipdate@4 <= 1996-12-31 -21)------------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey, l_extendedprice, l_discount, l_shipdate], file_type=csv, has_header=false +21)------------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_suppkey, l_extendedprice, l_discount, l_shipdate], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false 22)----------------------------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 -23)------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey], file_type=csv, has_header=false -24)------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=1 -25)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_nationkey], file_type=csv, has_header=false +23)------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +24)------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4 +25)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_custkey, c_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 26)--------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=4 27)----------------------FilterExec: n_name@1 = FRANCE OR n_name@1 = GERMANY 28)------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -29)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false +29)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false 30)----------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=4 31)------------------FilterExec: n_name@1 = GERMANY OR n_name@1 = FRANCE 32)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -33)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false +33)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q8.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q8.slt.part index 82de61c60b0a5..902413e9efb28 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q8.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q8.slt.part @@ -58,10 +58,10 @@ order by logical_plan 01)Sort: all_nations.o_year ASC NULLS LAST 02)--Projection: all_nations.o_year, CAST(CAST(sum(CASE WHEN all_nations.nation = Utf8("BRAZIL") THEN all_nations.volume ELSE Int64(0) END) AS Decimal128(12, 2)) / CAST(sum(all_nations.volume) AS Decimal128(12, 2)) AS Decimal128(15, 2)) AS mkt_share -03)----Aggregate: groupBy=[[all_nations.o_year]], aggr=[[sum(CASE WHEN all_nations.nation = Utf8View("BRAZIL") THEN all_nations.volume ELSE Decimal128(Some(0),38,4) END) AS sum(CASE WHEN all_nations.nation = Utf8("BRAZIL") THEN all_nations.volume ELSE Int64(0) END), sum(all_nations.volume)]] +03)----Aggregate: groupBy=[[all_nations.o_year]], aggr=[[sum(CASE WHEN all_nations.nation = Utf8View("BRAZIL") THEN all_nations.volume ELSE Decimal128(0.0000,38,4) END) AS sum(CASE WHEN all_nations.nation = Utf8("BRAZIL") THEN all_nations.volume ELSE Int64(0) END), sum(all_nations.volume)]] 04)------SubqueryAlias: all_nations -05)--------Projection: date_part(Utf8("YEAR"), orders.o_orderdate) AS o_year, lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount) AS volume, n2.n_name AS nation -06)----------Inner Join: n1.n_regionkey = region.r_regionkey +05)--------Projection: date_part(Utf8("YEAR"), orders.o_orderdate) AS o_year, lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount) AS volume, n2.n_name AS nation +06)----------LeftSemi Join: n1.n_regionkey = region.r_regionkey 07)------------Projection: lineitem.l_extendedprice, lineitem.l_discount, orders.o_orderdate, n1.n_regionkey, n2.n_name 08)--------------Inner Join: supplier.s_nationkey = n2.n_nationkey 09)----------------Projection: lineitem.l_extendedprice, lineitem.l_discount, supplier.s_nationkey, orders.o_orderdate, n1.n_regionkey @@ -73,7 +73,7 @@ logical_plan 15)----------------------------Projection: lineitem.l_orderkey, lineitem.l_extendedprice, lineitem.l_discount, supplier.s_nationkey 16)------------------------------Inner Join: lineitem.l_suppkey = supplier.s_suppkey 17)--------------------------------Projection: lineitem.l_orderkey, lineitem.l_suppkey, lineitem.l_extendedprice, lineitem.l_discount -18)----------------------------------Inner Join: part.p_partkey = lineitem.l_partkey +18)----------------------------------RightSemi Join: part.p_partkey = lineitem.l_partkey 19)------------------------------------Projection: part.p_partkey 20)--------------------------------------Filter: part.p_type = Utf8View("ECONOMY ANODIZED STEEL") 21)----------------------------------------TableScan: part projection=[p_partkey, p_type], partial_filters=[part.p_type = Utf8View("ECONOMY ANODIZED STEEL")] @@ -93,11 +93,11 @@ physical_plan 01)SortPreservingMergeExec: [o_year@0 ASC NULLS LAST] 02)--SortExec: expr=[o_year@0 ASC NULLS LAST], preserve_partitioning=[true] 03)----ProjectionExec: expr=[o_year@0 as o_year, CAST(CAST(sum(CASE WHEN all_nations.nation = Utf8("BRAZIL") THEN all_nations.volume ELSE Int64(0) END)@1 AS Decimal128(12, 2)) / CAST(sum(all_nations.volume)@2 AS Decimal128(12, 2)) AS Decimal128(15, 2)) as mkt_share] -04)------AggregateExec: mode=FinalPartitioned, gby=[o_year@0 as o_year], aggr=[sum(CASE WHEN all_nations.nation = BRAZIL THEN all_nations.volume ELSE Some(0),38,4 END) as sum(CASE WHEN all_nations.nation = Utf8("BRAZIL") THEN all_nations.volume ELSE Int64(0) END), sum(all_nations.volume)] +04)------AggregateExec: mode=FinalPartitioned, gby=[o_year@0 as o_year], aggr=[sum(CASE WHEN all_nations.nation = BRAZIL THEN all_nations.volume ELSE 0.0000 END) as sum(CASE WHEN all_nations.nation = Utf8("BRAZIL") THEN all_nations.volume ELSE Int64(0) END), sum(all_nations.volume)] 05)--------RepartitionExec: partitioning=Hash([o_year@0], 4), input_partitions=4 -06)----------AggregateExec: mode=Partial, gby=[o_year@0 as o_year], aggr=[sum(CASE WHEN all_nations.nation = BRAZIL THEN all_nations.volume ELSE Some(0),38,4 END) as sum(CASE WHEN all_nations.nation = Utf8("BRAZIL") THEN all_nations.volume ELSE Int64(0) END), sum(all_nations.volume)] -07)------------ProjectionExec: expr=[date_part(YEAR, o_orderdate@0) as o_year, l_extendedprice@1 * (Some(1),20,0 - l_discount@2) as volume, n_name@3 as nation] -08)--------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(n_regionkey@3, r_regionkey@0)], projection=[o_orderdate@2, l_extendedprice@0, l_discount@1, n_name@4] +06)----------AggregateExec: mode=Partial, gby=[o_year@0 as o_year], aggr=[sum(CASE WHEN all_nations.nation = BRAZIL THEN all_nations.volume ELSE 0.0000 END) as sum(CASE WHEN all_nations.nation = Utf8("BRAZIL") THEN all_nations.volume ELSE Int64(0) END), sum(all_nations.volume)] +07)------------ProjectionExec: expr=[date_part(YEAR, o_orderdate@0) as o_year, l_extendedprice@1 * (1 - l_discount@2) as volume, n_name@3 as nation] +08)--------------HashJoinExec: mode=Partitioned, join_type=LeftSemi, on=[(n_regionkey@3, r_regionkey@0)], projection=[o_orderdate@2, l_extendedprice@0, l_discount@1, n_name@4] 09)----------------RepartitionExec: partitioning=Hash([n_regionkey@3], 4), input_partitions=4 10)------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@2, n_nationkey@0)], projection=[l_extendedprice@0, l_discount@1, o_orderdate@3, n_regionkey@4, n_name@6] 11)--------------------RepartitionExec: partitioning=Hash([s_nationkey@2], 4), input_partitions=4 @@ -109,25 +109,24 @@ physical_plan 17)--------------------------------RepartitionExec: partitioning=Hash([l_orderkey@0], 4), input_partitions=4 18)----------------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_suppkey@1, s_suppkey@0)], projection=[l_orderkey@0, l_extendedprice@2, l_discount@3, s_nationkey@5] 19)------------------------------------RepartitionExec: partitioning=Hash([l_suppkey@1], 4), input_partitions=4 -20)--------------------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(p_partkey@0, l_partkey@1)], projection=[l_orderkey@1, l_suppkey@3, l_extendedprice@4, l_discount@5] +20)--------------------------------------HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(p_partkey@0, l_partkey@1)], projection=[l_orderkey@0, l_suppkey@2, l_extendedprice@3, l_discount@4] 21)----------------------------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 22)------------------------------------------FilterExec: p_type@1 = ECONOMY ANODIZED STEEL, projection=[p_partkey@0] -23)--------------------------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -24)----------------------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_type], file_type=csv, has_header=false -25)----------------------------------------RepartitionExec: partitioning=Hash([l_partkey@1], 4), input_partitions=4 -26)------------------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_partkey, l_suppkey, l_extendedprice, l_discount], file_type=csv, has_header=false -27)------------------------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -28)--------------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], file_type=csv, has_header=false -29)--------------------------------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 -30)----------------------------------FilterExec: o_orderdate@2 >= 1995-01-01 AND o_orderdate@2 <= 1996-12-31 -31)------------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_orderdate], file_type=csv, has_header=false -32)----------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=1 -33)------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl]]}, projection=[c_custkey, c_nationkey], file_type=csv, has_header=false -34)------------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 -35)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_regionkey], file_type=csv, has_header=false -36)--------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 -37)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false -38)----------------RepartitionExec: partitioning=Hash([r_regionkey@0], 4), input_partitions=4 -39)------------------FilterExec: r_name@1 = AMERICA, projection=[r_regionkey@0] -40)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -41)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/region.tbl]]}, projection=[r_regionkey, r_name], file_type=csv, has_header=false +23)--------------------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_type], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +24)----------------------------------------RepartitionExec: partitioning=Hash([l_partkey@1], 4), input_partitions=4 +25)------------------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_partkey, l_suppkey, l_extendedprice, l_discount], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false +26)------------------------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 +27)--------------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +28)--------------------------------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 +29)----------------------------------FilterExec: o_orderdate@2 >= 1995-01-01 AND o_orderdate@2 <= 1996-12-31 +30)------------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_custkey, o_orderdate], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +31)----------------------------RepartitionExec: partitioning=Hash([c_custkey@0], 4), input_partitions=4 +32)------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:0..606529], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:606529..1213058], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1213058..1819587], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/customer.tbl:1819587..2426114]]}, projection=[c_custkey, c_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +33)------------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 +34)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_regionkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +35)--------------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 +36)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +37)----------------RepartitionExec: partitioning=Hash([r_regionkey@0], 4), input_partitions=4 +38)------------------FilterExec: r_name@1 = AMERICA, projection=[r_regionkey@0] +39)--------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 +40)----------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/region.tbl]]}, projection=[r_regionkey, r_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/plans/q9.slt.part b/datafusion/sqllogictest/test_files/tpch/plans/q9.slt.part index 7a973490be479..29869bcddfeb1 100644 --- a/datafusion/sqllogictest/test_files/tpch/plans/q9.slt.part +++ b/datafusion/sqllogictest/test_files/tpch/plans/q9.slt.part @@ -56,7 +56,7 @@ logical_plan 02)--Projection: profit.nation, profit.o_year, sum(profit.amount) AS sum_profit 03)----Aggregate: groupBy=[[profit.nation, profit.o_year]], aggr=[[sum(profit.amount)]] 04)------SubqueryAlias: profit -05)--------Projection: nation.n_name AS nation, date_part(Utf8("YEAR"), orders.o_orderdate) AS o_year, lineitem.l_extendedprice * (Decimal128(Some(1),20,0) - lineitem.l_discount) - partsupp.ps_supplycost * lineitem.l_quantity AS amount +05)--------Projection: nation.n_name AS nation, date_part(Utf8("YEAR"), orders.o_orderdate) AS o_year, lineitem.l_extendedprice * (Decimal128(1,20,0) - lineitem.l_discount) - partsupp.ps_supplycost * lineitem.l_quantity AS amount 06)----------Inner Join: supplier.s_nationkey = nation.n_nationkey 07)------------Projection: lineitem.l_quantity, lineitem.l_extendedprice, lineitem.l_discount, supplier.s_nationkey, partsupp.ps_supplycost, orders.o_orderdate 08)--------------Inner Join: lineitem.l_orderkey = orders.o_orderkey @@ -64,24 +64,23 @@ logical_plan 10)------------------Inner Join: lineitem.l_suppkey = partsupp.ps_suppkey, lineitem.l_partkey = partsupp.ps_partkey 11)--------------------Projection: lineitem.l_orderkey, lineitem.l_partkey, lineitem.l_suppkey, lineitem.l_quantity, lineitem.l_extendedprice, lineitem.l_discount, supplier.s_nationkey 12)----------------------Inner Join: lineitem.l_suppkey = supplier.s_suppkey -13)------------------------Projection: lineitem.l_orderkey, lineitem.l_partkey, lineitem.l_suppkey, lineitem.l_quantity, lineitem.l_extendedprice, lineitem.l_discount -14)--------------------------Inner Join: part.p_partkey = lineitem.l_partkey -15)----------------------------Projection: part.p_partkey -16)------------------------------Filter: part.p_name LIKE Utf8View("%green%") -17)--------------------------------TableScan: part projection=[p_partkey, p_name], partial_filters=[part.p_name LIKE Utf8View("%green%")] -18)----------------------------TableScan: lineitem projection=[l_orderkey, l_partkey, l_suppkey, l_quantity, l_extendedprice, l_discount] -19)------------------------TableScan: supplier projection=[s_suppkey, s_nationkey] -20)--------------------TableScan: partsupp projection=[ps_partkey, ps_suppkey, ps_supplycost] -21)----------------TableScan: orders projection=[o_orderkey, o_orderdate] -22)------------TableScan: nation projection=[n_nationkey, n_name] +13)------------------------RightSemi Join: part.p_partkey = lineitem.l_partkey +14)--------------------------Projection: part.p_partkey +15)----------------------------Filter: part.p_name LIKE Utf8View("%green%") +16)------------------------------TableScan: part projection=[p_partkey, p_name], partial_filters=[part.p_name LIKE Utf8View("%green%")] +17)--------------------------TableScan: lineitem projection=[l_orderkey, l_partkey, l_suppkey, l_quantity, l_extendedprice, l_discount] +18)------------------------TableScan: supplier projection=[s_suppkey, s_nationkey] +19)--------------------TableScan: partsupp projection=[ps_partkey, ps_suppkey, ps_supplycost] +20)----------------TableScan: orders projection=[o_orderkey, o_orderdate] +21)------------TableScan: nation projection=[n_nationkey, n_name] physical_plan 01)SortPreservingMergeExec: [nation@0 ASC NULLS LAST, o_year@1 DESC], fetch=10 -02)--SortExec: TopK(fetch=10), expr=[nation@0 ASC NULLS LAST, o_year@1 DESC], preserve_partitioning=[true] -03)----ProjectionExec: expr=[nation@0 as nation, o_year@1 as o_year, sum(profit.amount)@2 as sum_profit] +02)--ProjectionExec: expr=[nation@0 as nation, o_year@1 as o_year, sum(profit.amount)@2 as sum_profit] +03)----SortExec: TopK(fetch=10), expr=[nation@0 ASC NULLS LAST, o_year@1 DESC], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[nation@0 as nation, o_year@1 as o_year], aggr=[sum(profit.amount)] 05)--------RepartitionExec: partitioning=Hash([nation@0, o_year@1], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[nation@0 as nation, o_year@1 as o_year], aggr=[sum(profit.amount)] -07)------------ProjectionExec: expr=[n_name@0 as nation, date_part(YEAR, o_orderdate@1) as o_year, l_extendedprice@2 * (Some(1),20,0 - l_discount@3) - ps_supplycost@4 * l_quantity@5 as amount] +07)------------ProjectionExec: expr=[n_name@0 as nation, date_part(YEAR, o_orderdate@1) as o_year, l_extendedprice@2 * (1 - l_discount@3) - ps_supplycost@4 * l_quantity@5 as amount] 08)--------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(s_nationkey@3, n_nationkey@0)], projection=[n_name@7, o_orderdate@5, l_extendedprice@1, l_discount@2, ps_supplycost@4, l_quantity@0] 09)----------------RepartitionExec: partitioning=Hash([s_nationkey@3], 4), input_partitions=4 10)------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_orderkey@0, o_orderkey@0)], projection=[l_quantity@1, l_extendedprice@2, l_discount@3, s_nationkey@4, ps_supplycost@5, o_orderdate@7] @@ -90,18 +89,17 @@ physical_plan 13)------------------------RepartitionExec: partitioning=Hash([l_suppkey@2, l_partkey@1], 4), input_partitions=4 14)--------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(l_suppkey@2, s_suppkey@0)], projection=[l_orderkey@0, l_partkey@1, l_suppkey@2, l_quantity@3, l_extendedprice@4, l_discount@5, s_nationkey@7] 15)----------------------------RepartitionExec: partitioning=Hash([l_suppkey@2], 4), input_partitions=4 -16)------------------------------HashJoinExec: mode=Partitioned, join_type=Inner, on=[(p_partkey@0, l_partkey@1)], projection=[l_orderkey@1, l_partkey@2, l_suppkey@3, l_quantity@4, l_extendedprice@5, l_discount@6] +16)------------------------------HashJoinExec: mode=Partitioned, join_type=RightSemi, on=[(p_partkey@0, l_partkey@1)] 17)--------------------------------RepartitionExec: partitioning=Hash([p_partkey@0], 4), input_partitions=4 18)----------------------------------FilterExec: p_name@1 LIKE %green%, projection=[p_partkey@0] -19)------------------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 -20)--------------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl]]}, projection=[p_partkey, p_name], file_type=csv, has_header=false -21)--------------------------------RepartitionExec: partitioning=Hash([l_partkey@1], 4), input_partitions=4 -22)----------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_partkey, l_suppkey, l_quantity, l_extendedprice, l_discount], file_type=csv, has_header=false -23)----------------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 -24)------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], file_type=csv, has_header=false -25)------------------------RepartitionExec: partitioning=Hash([ps_suppkey@1, ps_partkey@0], 4), input_partitions=4 -26)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_supplycost], file_type=csv, has_header=false -27)--------------------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 -28)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_orderdate], file_type=csv, has_header=false -29)----------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 -30)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], file_type=csv, has_header=false +19)------------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:0..597773], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:597773..1195546], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1195546..1793319], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/part.tbl:1793319..2391090]]}, projection=[p_partkey, p_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +20)--------------------------------RepartitionExec: partitioning=Hash([l_partkey@1], 4), input_partitions=4 +21)----------------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:0..18561749], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:18561749..37123498], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:37123498..55685247], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/lineitem.tbl:55685247..74246996]]}, projection=[l_orderkey, l_partkey, l_suppkey, l_quantity, l_extendedprice, l_discount], constraints=[PrimaryKey([0, 3])], file_type=csv, has_header=false +22)----------------------------RepartitionExec: partitioning=Hash([s_suppkey@0], 4), input_partitions=1 +23)------------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/supplier.tbl]]}, projection=[s_suppkey, s_nationkey], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +24)------------------------RepartitionExec: partitioning=Hash([ps_suppkey@1, ps_partkey@0], 4), input_partitions=4 +25)--------------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:0..2932049], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:2932049..5864098], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:5864098..8796147], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/partsupp.tbl:8796147..11728193]]}, projection=[ps_partkey, ps_suppkey, ps_supplycost], constraints=[PrimaryKey([0, 1])], file_type=csv, has_header=false +26)--------------------RepartitionExec: partitioning=Hash([o_orderkey@0], 4), input_partitions=4 +27)----------------------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:0..4223281], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:4223281..8446562], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:8446562..12669843], [WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/orders.tbl:12669843..16893122]]}, projection=[o_orderkey, o_orderdate], constraints=[PrimaryKey([0])], file_type=csv, has_header=false +28)----------------RepartitionExec: partitioning=Hash([n_nationkey@0], 4), input_partitions=1 +29)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/tpch/data/nation.tbl]]}, projection=[n_nationkey, n_name], constraints=[PrimaryKey([0])], file_type=csv, has_header=false diff --git a/datafusion/sqllogictest/test_files/tpch/tpch.slt b/datafusion/sqllogictest/test_files/tpch/tpch.slt index 764285784aa50..4a1cb4f9e02e2 100644 --- a/datafusion/sqllogictest/test_files/tpch/tpch.slt +++ b/datafusion/sqllogictest/test_files/tpch/tpch.slt @@ -21,6 +21,15 @@ include ./create_tables.slt.part include ./plans/q*.slt.part include ./answers/q*.slt.part +# test answers with uncorrelated scalar subqueries rewritten to joins +statement ok +set datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery = false; + +include ./answers/q*.slt.part + +statement ok +reset datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery; + # test answers with sort merge join statement ok set datafusion.optimizer.prefer_hash_join = false; @@ -31,4 +40,4 @@ include ./drop_tables.slt.part # Config reset statement ok -reset datafusion.optimizer.prefer_hash_join; \ No newline at end of file +reset datafusion.optimizer.prefer_hash_join; diff --git a/datafusion/sqllogictest/test_files/type_coercion.slt b/datafusion/sqllogictest/test_files/type_coercion.slt index 7039e66b38b15..6a56fc2407a94 100644 --- a/datafusion/sqllogictest/test_files/type_coercion.slt +++ b/datafusion/sqllogictest/test_files/type_coercion.slt @@ -301,4 +301,104 @@ query error does not support zero arguments SELECT * FROM (SELECT 1) WHERE CAST(STARTS_WITH() AS STRING) = 'x'; query error does not support zero arguments -SELECT * FROM (SELECT 1) WHERE TRY_CAST(STARTS_WITH() AS INT) = 1; \ No newline at end of file +SELECT * FROM (SELECT 1) WHERE TRY_CAST(STARTS_WITH() AS INT) = 1; + +################################################################### +## SIMILAR TO type coercion +## https://github.com/apache/datafusion/issues/22886 +## https://github.com/apache/datafusion/issues/23732 +################################################################### + +# NULL pattern is coerced to a typed NULL and evaluates to NULL instead of panicking +query B +SELECT 'a' SIMILAR TO NULL; +---- +NULL + +query B +SELECT NULL SIMILAR TO NULL; +---- +NULL + +query B +SELECT 'a' NOT SIMILAR TO NULL; +---- +NULL + +# operands of different string types are coerced to a common type +statement ok +CREATE TABLE t AS SELECT * FROM (VALUES ('user auth failed')) v(s); + +statement ok +CREATE TABLE p AS SELECT * FROM (VALUES ('(auth|login)')) v(pat); + +# Utf8View value with a non-scalar Utf8 pattern (issue repro) +query B +SELECT arrow_cast(t.s, 'Utf8View') SIMILAR TO p.pat FROM t CROSS JOIN p; +---- +true + +# LargeUtf8 value with a non-scalar Utf8 pattern +query B +SELECT arrow_cast(t.s, 'LargeUtf8') SIMILAR TO p.pat FROM t CROSS JOIN p; +---- +true + +# Dictionary value with a non-scalar Utf8 pattern must be unpacked before +# reaching the regex array kernel +query B +SELECT arrow_cast(t.s, 'Dictionary(Int32, Utf8)') SIMILAR TO p.pat FROM t CROSS JOIN p; +---- +true + +# non-scalar string-like patterns are coerced by the analyzer +query B +SELECT t.s SIMILAR TO arrow_cast(p.pat, 'Utf8View') FROM t CROSS JOIN p; +---- +true + +query B +SELECT t.s SIMILAR TO arrow_cast(p.pat, 'LargeUtf8') FROM t CROSS JOIN p; +---- +true + +query B +SELECT t.s NOT SIMILAR TO arrow_cast(p.pat, 'Utf8View') FROM t CROSS JOIN p; +---- +false + +query B +SELECT t.s SIMILAR TO arrow_cast(p.pat, 'Dictionary(Int32, Utf8)') FROM t CROSS JOIN p; +---- +true + +# NULL patterns (literal or Null-typed non-scalar) evaluate to NULL +query B +SELECT t.s SIMILAR TO NULL FROM t; +---- +NULL + +statement ok +CREATE TABLE pn AS SELECT NULL AS pat; + +query B +SELECT t.s SIMILAR TO pn.pat FROM t CROSS JOIN pn; +---- +NULL + +statement ok +DROP TABLE pn; + +statement ok +DROP TABLE t; + +statement ok +DROP TABLE p; + +# incompatible operand types are a planning error, not a panic +query error There isn't a common type to coerce Int64 and Utf8 in SIMILAR TO expression +SELECT 1 SIMILAR TO 'a'; + +# a non-string pattern is rejected by the analyzer +query error There isn't a common type to coerce Utf8 and Int64 in SIMILAR TO expression +SELECT 'a' SIMILAR TO 1; diff --git a/datafusion/sqllogictest/test_files/union.slt b/datafusion/sqllogictest/test_files/union.slt index a48ede604968b..d4776dd0c0ddb 100644 --- a/datafusion/sqllogictest/test_files/union.slt +++ b/datafusion/sqllogictest/test_files/union.slt @@ -21,11 +21,11 @@ statement ok CREATE TABLE t1( - id INT, + id INT, name TEXT ) as VALUES - (1, 'Alex'), - (2, 'Bob'), + (1, 'Alex'), + (2, 'Bob'), (3, 'Alice') ; @@ -34,20 +34,20 @@ CREATE TABLE t2( id TINYINT, name TEXT ) as VALUES - (1, 'Alex'), - (2, 'Bob'), + (1, 'Alex'), + (2, 'Bob'), (3, 'John') ; # union with EXCEPT(JOIN) query T rowsort -( +( SELECT name FROM t1 EXCEPT SELECT name FROM t2 -) +) UNION ALL -( +( SELECT name FROM t2 EXCEPT SELECT name FROM t1 @@ -58,13 +58,13 @@ John # union with type coercion query IT rowsort -( +( SELECT * FROM t1 EXCEPT SELECT * FROM t2 -) +) UNION ALL -( +( SELECT * FROM t2 EXCEPT SELECT * FROM t1 @@ -341,6 +341,14 @@ physical_plan 05)--------FilterExec: id@0 = 1 OR id@0 = 2 06)----------DataSourceExec: partitions=1, partition_sizes=[1] +# Regression: schema recomputation must preserve the unqualified UNION +# output labels while unions_to_filter is enabled. +query IT rowsort +SELECT id, name FROM t1 WHERE id = 1 UNION SELECT id, name FROM t1 WHERE id = 2 +---- +1 Alex +2 Bob + statement ok set datafusion.optimizer.enable_unions_to_filter = false; @@ -564,7 +572,7 @@ logical_plan physical_plan 01)CoalescePartitionsExec: fetch=3 02)--UnionExec -03)----ProjectionExec: expr=[count(Int64(1))@0 as cnt] +03)----ProjectionExec: expr=[CAST(count(Int64(1))@0 AS Int64) as cnt] 04)------GlobalLimitExec: skip=0, fetch=3 05)--------AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] 06)----------CoalescePartitionsExec @@ -576,7 +584,7 @@ physical_plan 12)----------------------FilterExec: c13@1 != C2GT5KVyOPZpgKVl110TyZO0NcJ434, projection=[c1@0] 13)------------------------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 14)--------------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c13], file_type=csv, has_header=true -15)----ProjectionExec: expr=[1 as cnt] +15)----ProjectionExec: expr=[CAST(1 AS Int64) as cnt] 16)------PlaceholderRowExec 17)----ProjectionExec: expr=[lead(b.c1,Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING@1 as cnt] 18)------GlobalLimitExec: skip=0, fetch=3 @@ -643,11 +651,11 @@ OPTIONS ('format.has_header' 'true'); query TT explain SELECT c1 FROM( -( +( SELECT c1 FROM t1 -) +) UNION ALL -( +( SELECT c1a FROM t2 )) ORDER BY c1 @@ -713,7 +721,7 @@ logical_plan 11)----------EmptyRelation: rows=1 physical_plan 01)UnionExec -02)--ProjectionExec: expr=[count(Int64(1))@1 as count, n@0 as n] +02)--ProjectionExec: expr=[count(Int64(1))@1 as count, CAST(n@0 AS Int64) as n] 03)----AggregateExec: mode=SinglePartitioned, gby=[n@0 as n], aggr=[count(Int64(1))], ordering_mode=Sorted 04)------ProjectionExec: expr=[5 as n] 05)--------PlaceholderRowExec @@ -822,8 +830,8 @@ DROP TABLE t4; # Test issue: https://github.com/apache/datafusion/issues/11742 query R rowsort -WITH - tt(v1) AS (VALUES (1::INT),(NULL::INT)) +WITH + tt(v1) AS (VALUES (1::INT),(NULL::INT)) SELECT NVL(v1, 0.5) FROM tt UNION ALL SELECT NULL WHERE FALSE; diff --git a/datafusion/sqllogictest/test_files/union_by_name.slt b/datafusion/sqllogictest/test_files/union_by_name.slt index 6a1608d5d1348..dbcaea778c0d9 100644 --- a/datafusion/sqllogictest/test_files/union_by_name.slt +++ b/datafusion/sqllogictest/test_files/union_by_name.slt @@ -124,7 +124,7 @@ NULL 5 # Ambiguous name -statement error DataFusion error: Schema error: No field named x. Valid fields are a, b. +statement error DataFusion error: Schema error: No field named x\.\nValid fields are a, b. SELECT x AS a FROM t1 UNION BY NAME SELECT x AS b FROM t1 ORDER BY x; query II diff --git a/datafusion/sqllogictest/test_files/union_function.slt b/datafusion/sqllogictest/test_files/union_function.slt index 74616490ab707..cb6f482dc9c72 100644 --- a/datafusion/sqllogictest/test_files/union_function.slt +++ b/datafusion/sqllogictest/test_files/union_function.slt @@ -28,6 +28,9 @@ select union_column, union_extract(union_column, 'int') from union_table; {int=1} 1 {string=bar} NULL {int=3} 3 +{int=1} 1 +{string=bar} NULL +{int=3} 3 query error DataFusion error: Execution error: field bool not found on union select union_extract(union_column, 'bool') from union_table; @@ -56,6 +59,9 @@ select union_column, union_tag(union_column) from union_table; {int=1} int {string=bar} string {int=3} int +{int=1} int +{string=bar} string +{int=3} int query error DataFusion error: Error during planning: 'union_tag' does not support zero arguments select union_tag() from union_table; @@ -65,3 +71,44 @@ select union_tag(union_column, 'int') from union_table; query error DataFusion error: Execution error: union_tag only support unions, got Utf8 select union_tag('int') from union_table; + +########## +## UNION Hashing Tests +########## + +query ?I +select union_column, count(*) +from union_table +group by union_column +order by union_column; +---- +{string=bar} 2 +{int=1} 2 +{int=3} 2 + +query ? +select distinct union_column +from union_table +order by union_column; +---- +{string=bar} +{int=1} +{int=3} + +query I +select count(distinct union_column) from union_table; +---- +3 + +query ?II +select + union_column, + count(*), + sum(union_extract(union_column, 'int')) +from union_table +group by union_column +order by union_column; +---- +{string=bar} 2 NULL +{int=1} 2 2 +{int=3} 2 6 diff --git a/datafusion/sqllogictest/test_files/unnest.slt b/datafusion/sqllogictest/test_files/unnest.slt index faeb5d59578e5..a3385b81d70d1 100644 --- a/datafusion/sqllogictest/test_files/unnest.slt +++ b/datafusion/sqllogictest/test_files/unnest.slt @@ -278,8 +278,8 @@ NULL NULL 17 NULL NULL 18 query IIII -select - unnest(column1), unnest(column2) + 2, +select + unnest(column1), unnest(column2) + 2, column3 * 10, unnest(array_remove(column1, 4)) from unnest_table; ---- @@ -903,7 +903,7 @@ query TT explain select * from unnest_table u, unnest(u.column1); ---- logical_plan -01)Cross Join: +01)Cross Join: 02)--SubqueryAlias: u 03)----TableScan: unnest_table projection=[column1, column2, column3, column4, column5] 04)--Subquery: @@ -1060,8 +1060,8 @@ logical_plan 04)------Projection: t.column1 AS __unnest_placeholder(t.column1), t.column2 05)--------TableScan: t projection=[column1, column2] physical_plan -01)SortExec: expr=[unnested@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--ProjectionExec: expr=[__unnest_placeholder(t.column1,depth=1)@0 as unnested, column2@1 as column2] +01)ProjectionExec: expr=[__unnest_placeholder(t.column1,depth=1)@0 as unnested, column2@1 as column2] +02)--SortExec: expr=[__unnest_placeholder(t.column1,depth=1)@0 ASC NULLS LAST], preserve_partitioning=[false] 03)----UnnestExec 04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/unnest/ordered_array.parquet]]}, projection=[column1@0 as __unnest_placeholder(t.column1), column2], output_ordering=[column2@1 ASC NULLS LAST], file_type=parquet @@ -1107,8 +1107,8 @@ logical_plan 05)--------Projection: struct(t.column1, t.column2, t.column3) AS __unnest_placeholder(struct(t.column1,t.column2,t.column3)) 06)----------TableScan: t projection=[column1, column2, column3] physical_plan -01)SortExec: expr=[struct(t.column1,t.column2,t.column3).c0@0 ASC NULLS LAST], preserve_partitioning=[false] -02)--ProjectionExec: expr=[__unnest_placeholder(struct(t.column1,t.column2,t.column3)).c0@0 as struct(t.column1,t.column2,t.column3).c0, __unnest_placeholder(struct(t.column1,t.column2,t.column3)).c1@1 as struct(t.column1,t.column2,t.column3).c1, __unnest_placeholder(struct(t.column1,t.column2,t.column3)).c2@2 as struct(t.column1,t.column2,t.column3).c2] +01)ProjectionExec: expr=[__unnest_placeholder(struct(t.column1,t.column2,t.column3)).c0@0 as struct(t.column1,t.column2,t.column3).c0, __unnest_placeholder(struct(t.column1,t.column2,t.column3)).c1@1 as struct(t.column1,t.column2,t.column3).c1, __unnest_placeholder(struct(t.column1,t.column2,t.column3)).c2@2 as struct(t.column1,t.column2,t.column3).c2] +02)--SortExec: expr=[__unnest_placeholder(struct(t.column1,t.column2,t.column3)).c0@0 ASC NULLS LAST], preserve_partitioning=[false] 03)----UnnestExec 04)------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/unnest/ordered_tuples.parquet]]}, projection=[struct(column1@0, column2@1, column3@2) as __unnest_placeholder(struct(t.column1,t.column2,t.column3))], file_type=parquet @@ -1419,3 +1419,230 @@ FROM ( statement ok DROP TABLE unused_unnest_pruning; + +## Regression: pushing a leaf-extracted projection (containing get_field, +## which has MoveTowardsLeafNodes placement) through an `Unnest` used to +## trip `Assertion failed: expr.is_empty(): Unnest` inside +## `PushDownLeafProjections`. The optimizer must not try to pushdown these +## projections through an `Unnest` and should produce a valid plan. + +statement ok +CREATE TABLE struct_and_list_table +AS VALUES + (struct(1, 2), [10, 20, 30]), + (struct(3, 4), [40, 50]); + +query I +SELECT sum(get_field(s, 'c0')) +FROM (SELECT s, unnest(arr) + FROM (SELECT column1 AS s, column2 AS arr + FROM struct_and_list_table)); +---- +9 + +statement ok +DROP TABLE struct_and_list_table; + +## Regression: get_field directly references the struct produced by unnest. +## This covers the case where the leaf-extracted expression depends on the +## unnested column itself rather than a sibling input column below the Unnest. + +statement ok +CREATE TABLE list_struct_table +AS VALUES + ([struct(1, 'a'), struct(2, 'b')]), + ([struct(3, 'c')]); + +query IT +SELECT get_field(unnest(column1), 'c0'), get_field(unnest(column1), 'c1') +FROM list_struct_table; +---- +1 a +2 b +3 c + +statement ok +DROP TABLE list_struct_table; + +#################################### +# `unnest_outer` Tests +# +# `unnest_outer(col)` is the outer-unnest peer to `unnest(col)`. Rows whose +# input list is `NULL` or empty produce a single output row containing +# `NULL`; rows with values are exploded element-by-element the same way as +# plain `unnest`. +# +# Column types on the tables below are inferred from the `VALUES` rows. +# DataFusion's SQL parser does not accept PostgreSQL `TYPE[]` array-column +# syntax inside `CREATE TABLE ... AS VALUES`, so tables are declared as +# CTAS over a `VALUES` subquery with aliased column names. +#################################### + +## unnest vs unnest_outer on an integer list + +statement ok +CREATE TABLE int_lists AS +SELECT column1 AS id, column2 AS xs FROM (VALUES + (1, [10, 20, 30]), + (4, [40]), + (5, [NULL, 50]), + (2, arrow_cast(make_array(), 'List(Int64)')), + (3, NULL) +); + +## Plain `unnest`: drops both NULL and empty input rows. +## Inner NULL elements survive. +query II +SELECT id, unnest(xs) AS x FROM int_lists ORDER BY id, x; +---- +1 10 +1 20 +1 30 +4 40 +5 50 +5 NULL + +## `unnest_outer`: NULL and empty input lists each produce one NULL row. +query II +SELECT id, unnest_outer(xs) AS x FROM int_lists ORDER BY id, x; +---- +1 10 +1 20 +1 30 +2 NULL +3 NULL +4 40 +5 50 +5 NULL + +## String list with inner NULLs: inner NULL elements must survive (they are +## not the same as "empty"), while NULL and empty input lists become a +## single NULL output row. + +statement ok +CREATE TABLE str_lists AS +SELECT column1 AS id, column2 AS tags FROM (VALUES + ('A', ['x', 'y']), + ('B', ['p', NULL, 'q']), + ('C', arrow_cast(make_array(), 'List(Utf8)')), + ('D', NULL) +); + +query TT +SELECT id, unnest_outer(tags) AS tag FROM str_lists ORDER BY id, tag; +---- +A x +A y +B p +B q +B NULL +C NULL +D NULL + +## Mixed list lengths — verify row-wise expansion. + +statement ok +CREATE TABLE varied_lists AS +SELECT column1 AS id, column2 AS xs FROM (VALUES + (1, [1, 2, 3, 4]), + (2, [5]), + (3, arrow_cast(make_array(), 'List(Int64)')), + (4, NULL) +); + +query II +SELECT id, unnest_outer(xs) AS x FROM varied_lists ORDER BY id, x; +---- +1 1 +1 2 +1 3 +1 4 +2 5 +3 NULL +4 NULL + +## Aliased output column. +query II +SELECT id, unnest_outer(xs) AS unwrapped FROM int_lists WHERE id = 2; +---- +2 NULL + +## Mixing `unnest` and `unnest_outer` in one SELECT is a planning error. +## `UnnestOptions` is per-`UnnestExec`, so we refuse to silently pick one mode. + +statement error DataFusion error: Error during planning: Cannot mix `unnest\(\.\.\.\)` with `unnest_outer\(\.\.\.\)` in the same SELECT +SELECT unnest(xs), unnest_outer(xs) FROM int_lists; + +## Chained `unnest` → `unnest_outer` via subquery. +## `unnest(xs)` (inner) drops NULL and empty outer rows, then +## `unnest_outer(ys)` (outer) preserves NULL and empty sub-lists from the +## inner unnest. + +statement ok +CREATE TABLE nested_lists AS +SELECT column1 AS id, column2 AS xs FROM (VALUES + (100, [[1, 2, 3], NULL, [4, 5]]), + (200, [[7], arrow_cast(make_array(), 'List(Int64)')]), + (300, NULL) +); + +query II +SELECT id, unnest_outer(ys) AS y +FROM (SELECT id, unnest(xs) AS ys FROM nested_lists) +ORDER BY id, y; +---- +100 1 +100 2 +100 3 +100 4 +100 5 +100 NULL +200 7 +200 NULL + +statement ok +DROP TABLE nested_lists; + +## `unnest_outer` agrees with `unnest` when no NULL or empty rows exist. + +statement ok +CREATE TABLE dense_lists AS +SELECT column1 AS id, column2 AS xs FROM (VALUES + (1, [10, 20]), + (2, [30]), + (3, [40, 50, 60]) +); + +query II +SELECT id, unnest(xs) AS x FROM dense_lists ORDER BY id, x; +---- +1 10 +1 20 +2 30 +3 40 +3 50 +3 60 + +query II +SELECT id, unnest_outer(xs) AS x FROM dense_lists ORDER BY id, x; +---- +1 10 +1 20 +2 30 +3 40 +3 50 +3 60 + +## Cleanup + +statement ok +DROP TABLE int_lists; + +statement ok +DROP TABLE str_lists; + +statement ok +DROP TABLE varied_lists; + +statement ok +DROP TABLE dense_lists; diff --git a/datafusion/sqllogictest/test_files/window.slt b/datafusion/sqllogictest/test_files/window.slt index 1c614f6a22c1e..6374cbf4f4b80 100644 --- a/datafusion/sqllogictest/test_files/window.slt +++ b/datafusion/sqllogictest/test_files/window.slt @@ -272,8 +272,8 @@ logical_plan 16)------------------EmptyRelation: rows=1 physical_plan 01)SortPreservingMergeExec: [b@0 ASC NULLS LAST] -02)--SortExec: expr=[b@0 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[b@0 as b, max(d.a)@1 as max_a] +02)--ProjectionExec: expr=[b@0 as b, max(d.a)@1 as max_a] +03)----SortExec: expr=[b@0 ASC NULLS LAST], preserve_partitioning=[true] 04)------AggregateExec: mode=FinalPartitioned, gby=[b@0 as b], aggr=[max(d.a)] 05)--------RepartitionExec: partitioning=Hash([b@0], 4), input_partitions=4 06)----------AggregateExec: mode=Partial, gby=[b@1 as b], aggr=[max(d.a)], ordering_mode=Sorted @@ -1215,6 +1215,10 @@ NULL 3917 -1114 -1114 15673 15673 +statement error Execution error: The second argument of nth_value must not be i64::MIN +SELECT nth_value(x, -9223372036854775808) OVER (ORDER BY x) +FROM (VALUES (1)) AS t(x); + @@ -2257,8 +2261,9 @@ physical_plan 06)----------ProjectionExec: expr=[c2@1 as c2, c8@2 as c8, c9@3 as c9, c1_alias@4 as c1_alias, sum(t1.c9) PARTITION BY [t1.c1, t1.c2] ORDER BY [t1.c9 ASC NULLS LAST, t1.c8 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING@5 as sum(t1.c9) PARTITION BY [t1.c1, t1.c2] ORDER BY [t1.c9 ASC NULLS LAST, t1.c8 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING, sum(t1.c9) PARTITION BY [t1.c1, t1.c2] ORDER BY [t1.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING@6 as sum(t1.c9) PARTITION BY [t1.c1, t1.c2] ORDER BY [t1.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING] 07)------------BoundedWindowAggExec: wdw=[sum(t1.c9) PARTITION BY [t1.c1, t1.c2] ORDER BY [t1.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING: Field { "sum(t1.c9) PARTITION BY [t1.c1, t1.c2] ORDER BY [t1.c9 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING": nullable UInt64 }, frame: ROWS BETWEEN 1 PRECEDING AND 5 FOLLOWING], mode=[Sorted] 08)--------------WindowAggExec: wdw=[sum(t1.c9) PARTITION BY [t1.c1, t1.c2] ORDER BY [t1.c9 ASC NULLS LAST, t1.c8 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING: Ok(Field { name: "sum(t1.c9) PARTITION BY [t1.c1, t1.c2] ORDER BY [t1.c9 ASC NULLS LAST, t1.c8 ASC NULLS LAST] ROWS BETWEEN 1 PRECEDING AND UNBOUNDED FOLLOWING", data_type: UInt64, nullable: true }), frame: WindowFrame { units: Rows, start_bound: Preceding(UInt64(1)), end_bound: Following(UInt64(NULL)), is_causal: false }] -09)----------------SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, c9@3 ASC NULLS LAST, c8@2 ASC NULLS LAST], preserve_partitioning=[false] -10)------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c2, c8, c9, c1@0 as c1_alias], file_type=csv, has_header=true +09)----------------ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, c8@2 as c8, c9@3 as c9, c1@0 as c1_alias] +10)------------------SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, c9@3 ASC NULLS LAST, c8@2 ASC NULLS LAST], preserve_partitioning=[false] +11)--------------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c1, c2, c8, c9], file_type=csv, has_header=true query IIIII SELECT c9, @@ -2404,8 +2409,8 @@ logical_plan 03)----WindowAggr: windowExpr=[[row_number() ORDER BY [aggregate_test_100.c9 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] 04)------TableScan: aggregate_test_100 projection=[c9] physical_plan -01)SortExec: TopK(fetch=5), expr=[rn1@1 DESC], preserve_partitioning=[false] -02)--ProjectionExec: expr=[c9@0 as c9, row_number() ORDER BY [aggregate_test_100.c9 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@1 as rn1] +01)ProjectionExec: expr=[c9@0 as c9, row_number() ORDER BY [aggregate_test_100.c9 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@1 as rn1] +02)--SortExec: TopK(fetch=5), expr=[row_number() ORDER BY [aggregate_test_100.c9 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@1 DESC], preserve_partitioning=[false] 03)----BoundedWindowAggExec: wdw=[row_number() ORDER BY [aggregate_test_100.c9 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() ORDER BY [aggregate_test_100.c9 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 04)------SortExec: expr=[c9@0 DESC], preserve_partitioning=[false] 05)--------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c9], file_type=csv, has_header=true @@ -5481,7 +5486,7 @@ order by c1, c2, rank; query TT explain select c1, c2, rank1, rank2 from ( - select c1, c2, rank() over (partition by c1 order by c2) as rank1, + select c1, c2, rank() over (partition by c1 order by c2) as rank1, rank() over (partition by c2, c1 order by c1) as rank2 from t1 ) @@ -5497,8 +5502,8 @@ logical_plan 06)----------TableScan: t1 projection=[c1, c2] physical_plan 01)SortPreservingMergeExec: [c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, rank1@2 ASC NULLS LAST, rank2@3 ASC NULLS LAST] -02)--SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, rank1@2 ASC NULLS LAST, rank2@3 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rank1, rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rank2] +02)--ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rank1, rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rank2] +03)----SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 ASC NULLS LAST, rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 ASC NULLS LAST], preserve_partitioning=[true] 04)------BoundedWindowAggExec: wdw=[rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 05)--------SortExec: expr=[c2@1 ASC NULLS LAST, c1@0 ASC NULLS LAST], preserve_partitioning=[true] 06)----------RepartitionExec: partitioning=Hash([c2@1, c1@0], 2), input_partitions=2 @@ -5512,7 +5517,7 @@ physical_plan query IIII select c1, c2, rank1, rank2 from ( - select c1, c2, rank() over (partition by c1 order by c2) as rank1, + select c1, c2, rank() over (partition by c1 order by c2) as rank1, rank() over (partition by c2, c1 order by c1) as rank2 from t1 ) @@ -5529,7 +5534,7 @@ order by c1, c2, rank1, rank2; query TT explain select c1, c2, rank1, rank2 from ( - select c1, c2, rank() over (partition by c1 order by c2) as rank1, + select c1, c2, rank() over (partition by c1 order by c2) as rank1, rank() over (partition by c2, c1 order by c1) as rank2 from t1 ) @@ -5545,8 +5550,8 @@ logical_plan 06)----------TableScan: t1 projection=[c1, c2] physical_plan 01)SortPreservingMergeExec: [c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, rank1@2 ASC NULLS LAST, rank2@3 ASC NULLS LAST] -02)--SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, rank1@2 ASC NULLS LAST, rank2@3 ASC NULLS LAST], preserve_partitioning=[true] -03)----ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rank1, rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rank2] +02)--ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rank1, rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rank2] +03)----SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 ASC NULLS LAST, rank() PARTITION BY [t1.c1] ORDER BY [t1.c2 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 ASC NULLS LAST, rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 ASC NULLS LAST], preserve_partitioning=[true] 04)------BoundedWindowAggExec: wdw=[rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [t1.c2, t1.c1] ORDER BY [t1.c1 ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 05)--------SortExec: expr=[c2@1 ASC NULLS LAST, c1@0 ASC NULLS LAST], preserve_partitioning=[true] 06)----------RepartitionExec: partitioning=Hash([c2@1, c1@0], 2), input_partitions=2 @@ -5559,7 +5564,7 @@ physical_plan query IIII select c1, c2, rank1, rank2 from ( - select c1, c2, rank() over (partition by c1 order by c2) as rank1, + select c1, c2, rank() over (partition by c1 order by c2) as rank1, rank() over (partition by c2, c1 order by c1) as rank2 from t1 ) @@ -5955,6 +5960,63 @@ physical_plan 07)------------DataSourceExec: partitions=2, partition_sizes=[5, 4] +# SUM(DISTINCT) over sliding frames must skip NULLs and return NULL +# for frames containing no non-null values. +statement ok +CREATE TABLE table_distinct_sum_nulls(ts INT, v BIGINT) AS VALUES + (1, NULL), (2, 3), (3, NULL), (4, NULL), (5, 5); + +query II +SELECT + ts, + SUM(DISTINCT v) OVER ( + ORDER BY ts + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ) AS s +FROM table_distinct_sum_nulls; +---- +1 NULL +2 3 +3 3 +4 NULL +5 5 + + +# SUM(DISTINCT) over sliding (bounded) window frames is only implemented +# for Int64. Other SUM-supported input types must fail with a clear +# capability error instead of an accumulator-internal one. +statement ok +CREATE TABLE table_distinct_sum_types(ts INT, f DOUBLE, d DECIMAL(10, 2)) AS VALUES + (1, 1.5, 1.50), (2, 2.5, 2.50), (3, 1.5, 1.50); + +query error DataFusion error: This feature is not implemented: SUM\(DISTINCT\) over sliding window frames is only supported for Int64, got Float64 +SELECT SUM(DISTINCT f) OVER ( + ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW +) FROM table_distinct_sum_types; + +query error DataFusion error: This feature is not implemented: SUM\(DISTINCT\) over sliding window frames is only supported for Int64, got Decimal128\(10, 2\) +SELECT SUM(DISTINCT d) OVER ( + ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW +) FROM table_distinct_sum_types; + +query error DataFusion error: This feature is not implemented: SUM\(DISTINCT\) over sliding window frames is only supported for Int64, got UInt64 +SELECT SUM(DISTINCT arrow_cast(ts, 'UInt64')) OVER ( + ORDER BY ts ROWS BETWEEN 1 PRECEDING AND CURRENT ROW +) FROM table_distinct_sum_types; + +# Unbounded frames take the regular distinct-sum path and keep +# supporting all SUM input types. +query R +SELECT SUM(DISTINCT f) OVER (ORDER BY ts) FROM table_distinct_sum_types; +---- +1.5 +4 +4 + +statement ok +DROP TABLE table_distinct_sum_types; + + # FILTER clause with window functions # Verify FILTER clause with non-aggregate window functions fails with a clear message @@ -6023,8 +6085,8 @@ physical_plan 03)----BoundedWindowAggExec: wdw=[sum(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, sum(test.c2) FILTER (WHERE test.c2 >= Int64(2) AND test.c2 < Int64(4) AND test.c1 > Int64(0)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "sum(test.c2) FILTER (WHERE test.c2 >= Int64(2) AND test.c2 < Int64(4) AND test.c1 > Int64(0)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, count(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "count(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": Int64 }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, array_agg(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "array_agg(test.c2) FILTER (WHERE test.c2 >= Int64(2)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable List(Int64) }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, array_agg(test.c2) FILTER (WHERE test.c2 >= Int64(2) AND test.c2 < Int64(4) AND test.c1 > Int64(0)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "array_agg(test.c2) FILTER (WHERE test.c2 >= Int64(2) AND test.c2 < Int64(4) AND test.c1 > Int64(0)) ORDER BY [test.c1 ASC NULLS LAST, test.c2 ASC NULLS LAST] ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": nullable List(Int64) }, frame: ROWS BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] 04)------SortPreservingMergeExec: [c1@2 ASC NULLS LAST, c2@3 ASC NULLS LAST], fetch=5 05)--------SortExec: TopK(fetch=5), expr=[c1@2 ASC NULLS LAST, c2@3 ASC NULLS LAST], preserve_partitioning=[true] -06)----------DataSourceExec: file_groups={4 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-0.csv], [WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-1.csv], [WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-2.csv], [WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-3.csv]]}, projection=[c2@1 >= 2 as __common_expr_1, c2@1 >= 2 AND c2@1 < 4 AND c1@0 > 0 as __common_expr_2, c1, c2], file_type=csv, has_header=false - +06)----------ProjectionExec: expr=[__common_expr_3@0 as __common_expr_1, __common_expr_3@0 AND c2@2 < 4 AND c1@1 > 0 as __common_expr_2, c1@1 as c1, c2@2 as c2] +07)------------DataSourceExec: file_groups={2 groups: [[WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-0.csv, WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-1.csv], [WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-2.csv, WORKSPACE_ROOT/datafusion/core/tests/data/partitioned_csv/partition-3.csv]]}, projection=[c2@1 >= 2 as __common_expr_3, c1, c2], file_type=csv, has_header=false # FILTER filters out some rows query IIIII?? @@ -6567,6 +6629,23 @@ FROM ( 2 1 3 1 +# A RANGE frame can transition from non-empty to empty and back to non-empty +# when the ORDER BY values contain gaps. The sliding accumulator must discard +# the state from the previous non-empty frame. +query IIIII +SELECT k, + SUM(v) OVER w, + COUNT(v) OVER w, + MIN(v) OVER w, + MAX(v) OVER w +FROM (VALUES (0, 100), (10, 90), (30, 10), (40, 20)) AS t(k, v) +WINDOW w AS (ORDER BY k RANGE BETWEEN 10 PRECEDING AND 5 PRECEDING); +---- +0 NULL 0 NULL NULL +10 100 1 100 100 +30 NULL 0 NULL NULL +40 10 1 10 10 + # AVG over a sliding window must yield NULL when the frame has no non-NULL # values — including frames that became empty via `retract_batch`. Covers # Float64, Decimal, and the narrow-frame retract-to-empty case. @@ -6600,6 +6679,142 @@ ORDER BY i; 3 1 4 NULL +# Covariance/correlation sliding-window regression test. Verifies correct +# results across row removals and a NULL-gap empty-frame transition. +query IRRR +SELECT + column1, + covar_pop(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + covar_samp(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + corr(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ) +FROM ( + VALUES + (1, 10.0, 5.0), + (2, NULL, NULL), + (3, NULL, NULL), + (4, 30.0, 10.0), + (5, 40.0, 20.0), + (6, 50.0, 10.0) +); +---- +1 0 NULL NULL +2 0 NULL NULL +3 NULL NULL NULL +4 0 NULL NULL +5 25 50 1 +6 -25 -50 -1 + +# Multi-row covariance/correlation sliding-window regression test. Verifies +# correct accumulation when valid rows enter the frame after a reset. +query IRRR +SELECT + column1, + covar_pop(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 2 PRECEDING AND CURRENT ROW + ), + covar_samp(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 2 PRECEDING AND CURRENT ROW + ), + corr(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 2 PRECEDING AND CURRENT ROW + ) +FROM ( + VALUES + (1, 10.0, 5.0), + (2, NULL, NULL), + (3, NULL, NULL), + (4, 30.0, 10.0), + (5, 40.0, 20.0), + (6, 50.0, 10.0) +); +---- +1 0 NULL NULL +2 0 NULL NULL +3 0 NULL NULL +4 0 NULL NULL +5 25 50 1 +6 0 0 0 + +# Covariance/correlation sliding-window regression test. Rows with NULL in +# either input column must not contribute to the aggregate state. +query IRRR +SELECT + column1, + covar_pop(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 3 PRECEDING AND CURRENT ROW + ), + covar_samp(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 3 PRECEDING AND CURRENT ROW + ), + corr(column2, column3) OVER ( + ORDER BY column1 + ROWS BETWEEN 3 PRECEDING AND CURRENT ROW + ) +FROM ( + VALUES + (1, 10.0, 5.0), + (2, 20.0, NULL), + (3, NULL, 15.0), + (4, 30.0, 10.0), + (5, 40.0, 20.0) +); +---- +1 0 NULL NULL +2 0 NULL NULL +3 0 NULL NULL +4 25 50 1 +5 25 50 1 + +# Variance/stddev sliding-window regression test. Verifies that retracting +# the last valid row resets the aggregate state. +query IRRRR +SELECT + column1, + var_pop(column2) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + var_samp(column2) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + stddev_pop(column2) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ), + stddev_samp(column2) OVER ( + ORDER BY column1 + ROWS BETWEEN 1 PRECEDING AND CURRENT ROW + ) +FROM ( + VALUES + (1, 10.0), + (2, NULL), + (3, NULL), + (4, 30.0), + (5, 40.0) +); +---- +1 0 NULL 0 NULL +2 0 NULL 0 NULL +3 NULL NULL NULL NULL +4 0 NULL 0 NULL +5 25 50 5 7.071067811865 + # Decimal variant — the integer-division path would otherwise panic on an # empty frame. query IR @@ -6619,7 +6834,7 @@ ORDER BY i; statement ok reset datafusion.execution.batch_size; -# The SLT runner sets `target_partitions` to 4 instead of using the default, so +# The SLT runner sets `target_partitions` to 4 instead of using the default, so # reset it explicitly. statement ok set datafusion.execution.target_partitions = 4; @@ -6644,3 +6859,18 @@ DROP TABLE issue_20194_t1; statement ok DROP TABLE issue_20194_t2; + +# Sliding-window over a frame whose non-NULL values have all been retracted should yield NULL. +query IIIIRR +SELECT id, x, + MIN(x) OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS min_x, + MAX(x) OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS max_x, + percentile_cont(x, 0.5) OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS percentile_x, + median(x) OVER (ORDER BY id ROWS BETWEEN 1 PRECEDING AND CURRENT ROW) AS median_x, +FROM (VALUES (1, 3), (2, NULL), (3, NULL), (4, 7)) t(id, x) +ORDER BY id +---- +1 3 3 3 3 3 +2 NULL 3 3 3 3 +3 NULL NULL NULL NULL NULL +4 7 7 7 7 7 diff --git a/datafusion/sqllogictest/test_files/window_topn.slt b/datafusion/sqllogictest/test_files/window_topn.slt index bf9ce26b35537..44cb31153b004 100644 --- a/datafusion/sqllogictest/test_files/window_topn.slt +++ b/datafusion/sqllogictest/test_files/window_topn.slt @@ -64,8 +64,9 @@ logical_plan physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] # Test 3: rn < 4 should give same results (fetch=3) query III rowsort @@ -131,8 +132,9 @@ logical_plan physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] # Test 7: Filter on data column (not window output) — should NOT optimize query TT @@ -164,8 +166,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] statement ok SET datafusion.explain.physical_plan_only = false; @@ -233,22 +236,32 @@ physical_plan 30)│ CURRENT ROW │ 31)└─────────────┬─────────────┘ 32)┌─────────────┴─────────────┐ -33)│ PartitionedTopKExec │ +33)│ RepartitionExec │ 34)│ -------------------- │ -35)│ fetch: 3 │ -36)│ │ -37)│ order: │ -38)│ [val@2 ASC NULLS LAST] │ -39)│ │ -40)│ partition: [pk@1] │ -41)└─────────────┬─────────────┘ -42)┌─────────────┴─────────────┐ -43)│ DataSourceExec │ -44)│ -------------------- │ -45)│ bytes: 480 │ -46)│ format: memory │ -47)│ rows: 1 │ -48)└───────────────────────────┘ +35)│ partition_count(in->out): │ +36)│ 1 -> 4 │ +37)│ │ +38)│ partitioning_scheme: │ +39)│ Hash([pk@1], 4) │ +40)└─────────────┬─────────────┘ +41)┌─────────────┴─────────────┐ +42)│ PartitionedTopKExec │ +43)│ -------------------- │ +44)│ fetch: 3 │ +45)│ fn: row_number │ +46)│ │ +47)│ order: │ +48)│ [val@2 ASC NULLS LAST] │ +49)│ │ +50)│ partition: [pk@1] │ +51)└─────────────┬─────────────┘ +52)┌─────────────┴─────────────┐ +53)│ DataSourceExec │ +54)│ -------------------- │ +55)│ bytes: 480 │ +56)│ format: memory │ +57)│ rows: 1 │ +58)└───────────────────────────┘ statement ok SET datafusion.explain.format = indent; @@ -308,9 +321,9 @@ EXPLAIN SELECT * FROM ( ---- physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 as rnk] -02)--FilterExec: rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@4 <= 3 -03)----BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] 05)--------DataSourceExec: partitions=1, partition_sizes=[1] # Test 14: Filter on rn AND rnk — compound predicate should NOT optimize @@ -360,8 +373,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk, window_topn_t.id] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk, window_topn_t.id] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk, window_topn_t.id] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=3, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1, id@0], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] statement ok SET datafusion.explain.physical_plan_only = false; @@ -391,8 +405,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.id] ORDER BY [window_topn_t.id ASC NULLS LAST, window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=3, partition=[id@0], order=[val@2 ASC NULLS LAST] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([id@0], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[id@0], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] # Test 19: Overlapping keys correctness (each id is unique, so rn=1 for all) statement ok @@ -426,8 +441,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk ASC NULLS LAST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 DESC] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 DESC] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] # Test 21: Correctness for PARTITION BY pk ORDER BY pk, val DESC statement ok @@ -460,8 +476,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.pk DESC NULLS FIRST, window_topn_t.val DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 DESC] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 DESC] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] statement ok SET datafusion.explain.physical_plan_only = false; @@ -494,8 +511,9 @@ QUALIFY rn <= 3; physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] # Test 30: QUALIFY with < operator statement ok @@ -522,9 +540,9 @@ QUALIFY rnk <= 3; ---- physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rnk] -02)--FilterExec: rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 <= 3 -03)----BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false] +02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_t.pk] ORDER BY [window_topn_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] 05)--------DataSourceExec: partitions=1, partition_sizes=[1] statement ok @@ -601,8 +619,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val ASC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=2, partition=[pk@1], order=[val@2 ASC] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@1], order=[val@2 ASC] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] query TT EXPLAIN SELECT * FROM ( @@ -612,8 +631,9 @@ EXPLAIN SELECT * FROM ( physical_plan 01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rn] 02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [window_topn_nulls.pk] ORDER BY [window_topn_nulls.val DESC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] -03)----PartitionedTopKExec: fetch=2, partition=[pk@1], order=[val@2 DESC NULLS LAST] -04)------DataSourceExec: partitions=1, partition_sizes=[1] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=row_number, fetch=2, partition=[pk@1], order=[val@2 DESC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] statement ok SET datafusion.explain.physical_plan_only = false; @@ -621,6 +641,478 @@ SET datafusion.explain.physical_plan_only = false; statement ok DROP TABLE window_topn_nulls; +############################################################################### +# RANK() tests +############################################################################### +# +# RANK semantics differ from ROW_NUMBER in that ties at the boundary are +# retained (`WHERE rk <= K` may keep more than K rows per partition). The +# tests below exercise both the boundary-Equal case (incoming row tied +# with current K-th-best) and the boundary-unchanged-after-eviction case +# (PartitionedTopKRank: heap evicts a tied row → push to per-partition +# `ties` Vec). + +# Table designed to produce ties at and around the rank-K boundary +statement ok +CREATE TABLE window_topn_rank_t (id INT, pk INT, val INT) AS VALUES + -- pk=1: ties at rank 2 (val=20 thrice), val=30 jumps to rank 5 + (1, 1, 10), + (2, 1, 20), + (3, 1, 20), + (4, 1, 20), + (5, 1, 30), + -- pk=2: distinct values, no ties + (6, 2, 5), + (7, 2, 15), + (8, 2, 25), + -- pk=3: 100 then four 200s — exercises the boundary-unchanged-with-eviction + -- case from the design doc's worked example (heap fills with three 200s, + -- the fourth ties, then 100 evicts a 200 but new boundary is still 200, + -- so the evicted 200 must move to ties) + (9, 3, 100), + (10, 3, 200), + (11, 3, 200), + (12, 3, 200), + (13, 3, 200), + (14, 3, 300); + +# Test R1: Basic RANK correctness with ties at the boundary. +# Expected per partition (RANK ASC, rk <= 3): +# pk=1: 10 (rk=1), 20×3 (rk=2 each) → 4 rows +# pk=2: 5 (rk=1), 15 (rk=2), 25 (rk=3) → 3 rows +# pk=3: 100 (rk=1), 200×4 (rk=2 each) → 5 rows +# Total: 12 rows kept, val=30 (pk=1, rk=5) and val=300 (pk=3, rk=6) dropped. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t +) WHERE rk <= 3; +---- +1 1 10 +10 3 200 +11 3 200 +12 3 200 +13 3 200 +2 1 20 +3 1 20 +4 1 20 +6 2 5 +7 2 15 +8 2 25 +9 3 100 + +# Test R2: EXPLAIN shows PartitionedTopKExec with fn=rank +query TT +EXPLAIN SELECT * FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t +) WHERE rk <= 3; +---- +logical_plan +01)Projection: window_topn_rank_t.id, window_topn_rank_t.pk, window_topn_rank_t.val, rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk +02)--Filter: rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(3) +03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: window_topn_rank_t projection=[id, pk, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk] +02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_t.pk] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Test R3: rk < 4 should give the same results (fetch = K-1 = 3) +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t +) WHERE rk < 4; +---- +1 1 10 +10 3 200 +11 3 200 +12 3 200 +13 3 200 +2 1 20 +3 1 20 +4 1 20 +6 2 5 +7 2 15 +8 2 25 +9 3 100 + +# Test R4: Flipped predicate `3 >= rk` should also trigger optimization +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t +) WHERE 3 >= rk; +---- +1 1 10 +10 3 200 +11 3 200 +12 3 200 +13 3 200 +2 1 20 +3 1 20 +4 1 20 +6 2 5 +7 2 15 +8 2 25 +9 3 100 + +# Test R5: Flipped strict `4 > rk` should also trigger optimization (fetch=3) +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_t +) WHERE 4 > rk; +---- +1 1 10 +10 3 200 +11 3 200 +12 3 200 +13 3 200 +2 1 20 +3 1 20 +4 1 20 +6 2 5 +7 2 15 +8 2 25 +9 3 100 + +# Test R6: RANK without PARTITION BY — should NOT trigger the optimization +# (global top-K with ties; SortExec with fetch handles this without our rule). +# Use window_topn_rank_t (still alive); window_topn_t was dropped earlier. +query II rowsort +SELECT id, val FROM ( + SELECT *, RANK() OVER (ORDER BY val) as rk FROM window_topn_rank_t +) WHERE rk <= 3; +---- +1 10 +6 5 +7 15 + +# Test R7: RANK with multi-column PARTITION BY +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk, id ORDER BY val) as rk FROM window_topn_rank_t +) WHERE rk <= 1; +---- +1 1 10 +10 3 200 +11 3 200 +12 3 200 +13 3 200 +14 3 300 +2 1 20 +3 1 20 +4 1 20 +5 1 30 +6 2 5 +7 2 15 +8 2 25 +9 3 100 + +# Test R8: Verify multi-column partition plan still uses fn=rank +query TT +EXPLAIN SELECT * FROM ( + SELECT *, RANK() OVER (PARTITION BY pk, id ORDER BY val) as rk FROM window_topn_rank_t +) WHERE rk <= 1; +---- +logical_plan +01)Projection: window_topn_rank_t.id, window_topn_rank_t.pk, window_topn_rank_t.val, rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk +02)--Filter: rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: window_topn_rank_t projection=[id, pk, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk] +02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_t.pk, window_topn_rank_t.id] ORDER BY [window_topn_rank_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----RepartitionExec: partitioning=Hash([pk@1, id@0], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=rank, fetch=1, partition=[pk@1, id@0], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Test R9: RANK with DESC ordering +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val DESC) as rk FROM window_topn_rank_t +) WHERE rk <= 1; +---- +14 3 300 +5 1 30 +8 2 25 + +# Test R10: Mixed window functions — RANK + ROW_NUMBER in the same query. +# Filter is on the RANK column; rule should still fire (matches by col_idx). +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, + ROW_NUMBER() OVER (PARTITION BY pk ORDER BY val) as rn, + RANK() OVER (PARTITION BY pk ORDER BY val) as rk + FROM window_topn_rank_t +) WHERE rk <= 1; +---- +1 1 10 +6 2 5 +9 3 100 + +# Test R11: QUALIFY form (parser desugars to the same plan) +query IIII rowsort +SELECT id, pk, val, + RANK() OVER (PARTITION BY pk ORDER BY val) as rk +FROM window_topn_rank_t +QUALIFY rk <= 1; +---- +1 1 10 1 +6 2 5 1 +9 3 100 1 + +statement ok +DROP TABLE window_topn_rank_t; + +############################################################################### +# RANK() — equality predicate (negative: rule supports only =/>) +############################################################################### +# +# `extract_window_limit` matches only `<, <=, >, >=`. Equality predicates +# `rk = N` are NOT optimized by this rule (regardless of N). DuckDB +# special-cases `rk = 1` as equivalent to `rk <= 1`; we don't. The two +# tests below pin current behavior so that an accidental rule extension +# (or regression) shows up. + +statement ok +CREATE TABLE window_topn_rank_eq_t (id INT, pk INT, val INT) AS VALUES + (1, 1, 10), (2, 1, 20), (3, 1, 30), + (4, 2, 5), (5, 2, 15), (6, 2, 25); + +# Test R12: `rk = 1` — correct results, but plan should still contain +# FilterExec + SortExec (rule did NOT fire). +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t +) WHERE rk = 1; +---- +1 1 10 +4 2 5 + +query TT +EXPLAIN SELECT * FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_eq_t +) WHERE rk = 1; +---- +logical_plan +01)Projection: window_topn_rank_eq_t.id, window_topn_rank_eq_t.pk, window_topn_rank_eq_t.val, rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk +02)--Filter: rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW = UInt64(1) +03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: window_topn_rank_eq_t projection=[id, pk, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk] +02)--FilterExec: rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 = 1 +03)----BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_eq_t.pk] ORDER BY [window_topn_rank_eq_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +04)------SortExec: expr=[pk@1 ASC NULLS LAST, val@2 ASC NULLS LAST], preserve_partitioning=[false] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +DROP TABLE window_topn_rank_eq_t; + +############################################################################### +# RANK() — dense-ties boundary preservation +############################################################################### +# +# Heap fills with K=3 rows tied at the same value, then a strictly-better +# row arrives. The heap evicts one of the tied rows, but the new +# K-th-best is still tied with the evicted row (boundary unchanged). +# PartitionedTopKRank must push the evicted row into `ties` rather than +# discarding it. Without that branch, a `rk <= 3` query loses the +# evicted tied row. + +statement ok +CREATE TABLE window_topn_rank_dense_t (id INT, pk INT, val INT) AS VALUES + -- ten rows with the same val + one strictly-better row + (1, 1, 10), (2, 1, 10), (3, 1, 10), (4, 1, 10), (5, 1, 10), + (6, 1, 10), (7, 1, 10), (8, 1, 10), (9, 1, 10), (10, 1, 10), + (11, 1, 5); + +# Test R14: With `rk <= 3`, every row should be retained: +# - val=5 → rk=1 +# - val=10 (×10) → rk=2 each +# Total 11 rows. If the boundary-unchanged-eviction branch ever drops a +# tied row, this query would return fewer than 11. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_dense_t +) WHERE rk <= 3; +---- +1 1 10 +10 1 10 +11 1 5 +2 1 10 +3 1 10 +4 1 10 +5 1 10 +6 1 10 +7 1 10 +8 1 10 +9 1 10 + +# Test R15: rule fired (no FilterExec/SortExec) +query TT +EXPLAIN SELECT * FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val) as rk FROM window_topn_rank_dense_t +) WHERE rk <= 3; +---- +logical_plan +01)Projection: window_topn_rank_dense_t.id, window_topn_rank_dense_t.pk, window_topn_rank_dense_t.val, rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk +02)--Filter: rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(3) +03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: window_topn_rank_dense_t projection=[id, pk, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk] +02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_dense_t.pk] ORDER BY [window_topn_rank_dense_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=rank, fetch=3, partition=[pk@1], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +DROP TABLE window_topn_rank_dense_t; + +############################################################################### +# RANK() — NULL handling in ORDER BY +############################################################################### +# +# RANK ASSIGNMENTS WITH NULLS: +# ORDER BY val ASC NULLS LAST → non-NULLs ranked first, NULLs at the end +# ORDER BY val DESC NULLS LAST → same shape, different non-NULL order +# ORDER BY val ASC NULLS FIRST → NULLs all tie at rank 1 +# ORDER BY val DESC NULLS FIRST → NULLs all tie at rank 1 +# +# Multiple NULLs in the same partition all share the same rank (they're +# tied under the encoded ORDER BY). + +statement ok +CREATE TABLE window_topn_rank_null_t (id INT, pk INT, val INT) AS VALUES + -- pk=1: distinct vals plus one NULL → ASC NULLS LAST → 1,2,3,NULL ranks 1,2,3,4 + (1, 1, 1), (2, 1, 2), (3, 1, 3), (4, 1, NULL), + -- pk=2: one non-NULL plus two NULLs → ASC NULLS LAST → 5,NULL,NULL ranks 1,2,2 + (5, 2, 5), (6, 2, NULL), (7, 2, NULL); + +# Test R16: ASC NULLS LAST, rk <= 4 covers everything in pk=1, only rk≤2 +# in pk=2 (since both NULLs tie at rank 2 and there's no rank 3 or 4). +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS LAST) as rk FROM window_topn_rank_null_t +) WHERE rk <= 4; +---- +1 1 1 +2 1 2 +3 1 3 +4 1 NULL +5 2 5 +6 2 NULL +7 2 NULL + +# Test R17: ASC NULLS LAST, rk <= 2 — pk=1's NULL (rk=4) drops out; +# pk=2's NULLs (rk=2 each) are retained. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS LAST) as rk FROM window_topn_rank_null_t +) WHERE rk <= 2; +---- +1 1 1 +2 1 2 +5 2 5 +6 2 NULL +7 2 NULL + +# Test R18: rule fires for NULLS LAST configuration +query TT +EXPLAIN SELECT * FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS LAST) as rk FROM window_topn_rank_null_t +) WHERE rk <= 2; +---- +logical_plan +01)Projection: window_topn_rank_null_t.id, window_topn_rank_null_t.pk, window_topn_rank_null_t.val, rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rk +02)--Filter: rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(2) +03)----WindowAggr: windowExpr=[[rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: window_topn_rank_null_t projection=[id, pk, val] +physical_plan +01)ProjectionExec: expr=[id@0 as id, pk@1 as pk, val@2 as val, rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@3 as rk] +02)--BoundedWindowAggExec: wdw=[rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "rank() PARTITION BY [window_topn_rank_null_t.pk] ORDER BY [window_topn_rank_null_t.val ASC NULLS LAST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----RepartitionExec: partitioning=Hash([pk@1], 4), input_partitions=1, maintains_sort_order=true +04)------PartitionedTopKExec: fn=rank, fetch=2, partition=[pk@1], order=[val@2 ASC NULLS LAST] +05)--------DataSourceExec: partitions=1, partition_sizes=[1] + +# Test R19: DESC NULLS LAST — pk=1: 3,2,1,NULL ranks 1,2,3,4; pk=2: 5,NULL,NULL ranks 1,2,2. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val DESC NULLS LAST) as rk FROM window_topn_rank_null_t +) WHERE rk <= 4; +---- +1 1 1 +2 1 2 +3 1 3 +4 1 NULL +5 2 5 +6 2 NULL +7 2 NULL + +# Test R20: ASC NULLS FIRST — pk=1: NULL,1,2,3 ranks 1,2,3,4; +# pk=2: NULL,NULL,5 ranks 1,1,3. With rk <= 2, pk=2's NULLs are kept, +# pk=1 keeps NULL and val=1. +query III rowsort +SELECT id, pk, val FROM ( + SELECT *, RANK() OVER (PARTITION BY pk ORDER BY val ASC NULLS FIRST) as rk FROM window_topn_rank_null_t +) WHERE rk <= 2; +---- +1 1 1 +4 1 NULL +6 2 NULL +7 2 NULL + +statement ok +DROP TABLE window_topn_rank_null_t; + # Reset config to default (false) statement ok SET datafusion.optimizer.enable_window_topn = false; + +statement ok +create table t(c1 int, c2 int) as values (1, 2), (3, 4); + +statement ok +set datafusion.execution.target_partitions = 5; + +statement ok +set datafusion.optimizer.repartition_windows = false; + +statement ok +set datafusion.execution.batch_size = 1; + +statement ok +set datafusion.optimizer.enable_window_topn = true; + +query TT +EXPLAIN SELECT * FROM ( + SELECT c1, c2, ROW_NUMBER() OVER (PARTITION BY c1 ORDER BY c2 DESC) as rn + FROM t +) WHERE rn <= 1; +---- +logical_plan +01)Projection: t.c1, t.c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW AS rn +02)--Filter: row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW <= UInt64(1) +03)----WindowAggr: windowExpr=[[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW]] +04)------TableScan: t projection=[c1, c2] +physical_plan +01)ProjectionExec: expr=[c1@0 as c1, c2@1 as c2, row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW@2 as rn] +02)--BoundedWindowAggExec: wdw=[row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW: Field { "row_number() PARTITION BY [t.c1] ORDER BY [t.c2 DESC NULLS FIRST] RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW": UInt64 }, frame: RANGE BETWEEN UNBOUNDED PRECEDING AND CURRENT ROW], mode=[Sorted] +03)----SortPreservingMergeExec: [c1@0 ASC NULLS LAST, c2@1 DESC] +04)------SortExec: expr=[c1@0 ASC NULLS LAST, c2@1 DESC], preserve_partitioning=[true] +05)--------PartitionedTopKExec: fn=row_number, fetch=1, partition=[c1@0], order=[c2@1 DESC] +06)----------RepartitionExec: partitioning=Hash([c1@0], 5), input_partitions=1 +07)------------DataSourceExec: partitions=1, partition_sizes=[1] + +statement ok +set datafusion.execution.target_partitions = 4; + +statement ok +set datafusion.optimizer.repartition_windows = true; + +statement ok +set datafusion.execution.batch_size = 8192; + +statement ok +set datafusion.optimizer.enable_window_topn = false; diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs index dae6c625ef55b..be084f360358a 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/field_reference.rs @@ -21,7 +21,7 @@ use datafusion::logical_expr::Expr; use std::sync::Arc; use substrait::proto::expression::FieldReference; use substrait::proto::expression::field_reference::ReferenceType::DirectReference; -use substrait::proto::expression::field_reference::RootType; +use substrait::proto::expression::field_reference::{LambdaParameterReference, RootType}; use substrait::proto::expression::reference_segment::ReferenceType::StructField; pub async fn from_field_reference( @@ -56,9 +56,9 @@ pub(crate) fn from_substrait_field_reference( Some(RootType::Expression(_)) => not_impl_err!( "Expression root type in field reference is not supported" ), - Some(RootType::LambdaParameterReference(_)) => not_impl_err!( - "Lambda parameter reference in field reference is not yet supported" - ), + Some(RootType::LambdaParameterReference( + LambdaParameterReference { steps_out }, + )) => consumer.lambda_variable(*steps_out as usize, field_idx), } } _ => not_impl_err!( @@ -85,3 +85,85 @@ fn resolve_outer_reference( let col = Column::from((qualifier, field)); Ok(Expr::OuterReferenceColumn(Arc::clone(field), col)) } + +#[cfg(test)] +mod tests { + use datafusion::{ + common::{DFSchema, assert_contains}, + prelude::SessionContext, + }; + use substrait::proto::{ + Type, + expression::{ + FieldReference, ReferenceSegment, + field_reference::{self, LambdaParameterReference, RootType}, + reference_segment::{ReferenceType, StructField}, + }, + r#type::{I64, Kind}, + }; + + use crate::{ + extensions::Extensions, + logical_plan::consumer::{ + DefaultSubstraitConsumer, SubstraitConsumer, from_field_reference, + }, + }; + + #[tokio::test] + async fn test_lambda_variable_invalid_steps_out() { + let lambda_field_ref = lambda_field_ref(0, 99); + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = + from_field_reference(&consumer, &lambda_field_ref, DFSchema::empty_ref()) + .await + .unwrap_err(); + + assert_contains!(err.to_string(), "No lambda at 99 steps out, got only 0"); + } + + #[tokio::test] + async fn test_lambda_variable_invalid_field_idx() { + let lambda_field_ref = lambda_field_ref(1, 0); + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + let _names = consumer + .push_lambda_parameters( + &[Type { + kind: Some(Kind::I64(I64::default())), + }], + DFSchema::empty_ref(), + ) + .unwrap(); + + let err = + from_field_reference(&consumer, &lambda_field_ref, DFSchema::empty_ref()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "At lambda 0 steps out, no field at index 1, got only 1" + ); + } + + fn lambda_field_ref(field: i32, steps_out: u32) -> FieldReference { + FieldReference { + reference_type: Some(field_reference::ReferenceType::DirectReference( + ReferenceSegment { + reference_type: Some(ReferenceType::StructField(Box::new( + StructField { field, child: None }, + ))), + }, + )), + root_type: Some(RootType::LambdaParameterReference( + LambdaParameterReference { steps_out }, + )), + } + } +} diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs b/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs new file mode 100644 index 0000000000000..c4554dea8770d --- /dev/null +++ b/datafusion/substrait/src/logical_plan/consumer/expr/lambda.rs @@ -0,0 +1,102 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::{ + common::{DFSchema, substrait_err}, + prelude::{Expr, lambda}, +}; +use substrait::proto; + +use crate::logical_plan::consumer::SubstraitConsumer; + +pub async fn from_lambda( + consumer: &impl SubstraitConsumer, + expr: &proto::expression::Lambda, + input_schema: &DFSchema, +) -> datafusion::common::Result { + let Some(parameters) = expr.parameters.as_ref() else { + return substrait_err!("Lambda expression without parameters is not allowed"); + }; + + let names = consumer.push_lambda_parameters(¶meters.types, input_schema)?; + + let Some(body) = expr.body.as_ref() else { + return substrait_err!("Lambda expression without body is not allowed"); + }; + + let body = consumer.consume_expression(body, input_schema).await?; + + consumer.pop_lambda_parameters(); + + Ok(lambda(names, body)) +} + +#[cfg(test)] +mod tests { + use datafusion::{ + common::{DFSchema, assert_contains}, + prelude::SessionContext, + }; + use substrait::proto::{self, Expression, r#type::Struct}; + + use crate::{ + extensions::Extensions, + logical_plan::consumer::{DefaultSubstraitConsumer, from_lambda}, + }; + + #[tokio::test] + async fn test_lambda_without_body() { + let lambda = proto::expression::Lambda { + parameters: Some(Struct::default()), + body: None, + }; + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = from_lambda(&consumer, &lambda, DFSchema::empty_ref()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "Lambda expression without body is not allowed" + ); + } + + #[tokio::test] + async fn test_lambda_without_parameters() { + let lambda = proto::expression::Lambda { + parameters: None, + body: Some(Box::new(Expression::default())), + }; + + let extensions = Extensions::default(); + let session_state = SessionContext::new().state(); + let consumer = DefaultSubstraitConsumer::new(&extensions, &session_state); + + let err = from_lambda(&consumer, &lambda, DFSchema::empty_ref()) + .await + .unwrap_err(); + + assert_contains!( + err.to_string(), + "Lambda expression without parameters is not allowed" + ); + } +} diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs b/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs index 295456e95f9f3..2fcc11f4e417d 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/mod.rs @@ -20,6 +20,7 @@ mod cast; mod field_reference; mod function_arguments; mod if_then; +mod lambda; mod literal; mod nested; mod scalar_function; @@ -32,6 +33,7 @@ pub use cast::*; pub use field_reference::*; pub use function_arguments::*; pub use if_then::*; +pub use lambda::*; pub use literal::*; pub use nested::*; pub use scalar_function::*; @@ -95,8 +97,11 @@ pub async fn from_substrait_rex( RexType::DynamicParameter(expr) => { consumer.consume_dynamic_parameter(expr, input_schema).await } - RexType::Lambda(_) | RexType::LambdaInvocation(_) => { - not_impl_err!("Lambda expressions are not yet supported") + RexType::Lambda(lambda) => { + consumer.consume_lambda(lambda.as_ref(), input_schema).await + } + RexType::LambdaInvocation(_) => { + not_impl_err!("Lambda invocations are not supported") } }, None => substrait_err!("Expression must set rex_type: {expression:?}"), diff --git a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs index 1a0fb3f55f609..47a944504c510 100644 --- a/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/consumer/expr/scalar_function.rs @@ -30,7 +30,6 @@ pub async fn from_scalar_function( f: &ScalarFunction, input_schema: &DFSchema, ) -> Result { - //TODO: handle higher order functions, as they are also encoded as scalar functions let Some(fn_signature) = consumer .get_extensions() .functions @@ -45,6 +44,20 @@ pub async fn from_scalar_function( let fn_name = substrait_fun_name(fn_signature); let args = from_substrait_func_args(consumer, &f.arguments, input_schema).await?; + let higher_order_func = consumer + .get_function_registry() + .higher_order_function(fn_name) + .or_else(|e| { + if let Some(alt_name) = substrait_to_df_name(fn_name) { + consumer + .get_function_registry() + .higher_order_function(alt_name) + .or(Err(e)) + } else { + Err(e) + } + }); + let udf_func = consumer.get_function_registry().udf(fn_name).or_else(|e| { if let Some(alt_name) = substrait_to_df_name(fn_name) { consumer.get_function_registry().udf(alt_name).or(Err(e)) @@ -53,9 +66,14 @@ pub async fn from_scalar_function( } }); - // try to first match the requested function into registered udfs, then built-in ops + // try to first match the requested function into registered higher-order functions, then udfs, built-in ops // and finally built-in expressions - if let Ok(func) = udf_func { + if let Ok(func) = higher_order_func { + Ok(Expr::HigherOrderFunction(expr::HigherOrderFunction::new( + func.to_owned(), + args, + ))) + } else if let Ok(func) = udf_func { Ok(Expr::ScalarFunction(expr::ScalarFunction::new_udf( func.to_owned(), args, @@ -70,7 +88,7 @@ pub async fn from_scalar_function( // In those cases we build a balanced tree of BinaryExprs arg_list_to_binary_op_tree(op, args) } else if let Some(builder) = BuiltinExprBuilder::try_from_name(fn_name) { - builder.build(consumer, f, args).await + builder.build(consumer, f, args) } else { not_impl_err!("Unsupported function name: {fn_name:?}") } @@ -188,34 +206,32 @@ impl BuiltinExprBuilder { } } - pub async fn build( + pub fn build( self, consumer: &impl SubstraitConsumer, f: &ScalarFunction, args: Vec, ) -> Result { match self.expr_name.as_str() { - "like" => Self::build_like_expr(false, false, f, args).await, - "ilike" => Self::build_like_expr(true, false, f, args).await, - "like_match" => Self::build_like_expr(false, false, f, args).await, - "like_imatch" => Self::build_like_expr(true, false, f, args).await, - "like_not_match" => Self::build_like_expr(false, true, f, args).await, - "like_not_imatch" => Self::build_like_expr(true, true, f, args).await, + "like" => Self::build_like_expr(false, false, f, args), + "ilike" => Self::build_like_expr(true, false, f, args), + "like_match" => Self::build_like_expr(false, false, f, args), + "like_imatch" => Self::build_like_expr(true, false, f, args), + "like_not_match" => Self::build_like_expr(false, true, f, args), + "like_not_imatch" => Self::build_like_expr(true, true, f, args), "not" | "negative" | "negate" | "is_null" | "is_not_null" | "is_true" | "is_false" | "is_not_true" | "is_not_false" | "is_unknown" - | "is_not_unknown" => Self::build_unary_expr(&self.expr_name, args).await, - "and_not" | "xor" => Self::build_binary_expr(&self.expr_name, args).await, - "between" => Self::build_between_expr(&self.expr_name, args).await, - "logb" => { - Self::build_custom_handling_expr(consumer, &self.expr_name, args).await - } + | "is_not_unknown" => Self::build_unary_expr(&self.expr_name, args), + "and_not" | "xor" => Self::build_binary_expr(&self.expr_name, args), + "between" => Self::build_between_expr(&self.expr_name, args), + "logb" => Self::build_custom_handling_expr(consumer, &self.expr_name, args), _ => { not_impl_err!("Unsupported builtin expression: {}", self.expr_name) } } } - async fn build_unary_expr(fn_name: &str, args: Vec) -> Result { + fn build_unary_expr(fn_name: &str, args: Vec) -> Result { let [arg] = match args.try_into() { Ok(args_arr) => args_arr, Err(_) => return substrait_err!("Expected one argument for {fn_name} expr"), @@ -239,7 +255,7 @@ impl BuiltinExprBuilder { Ok(expr) } - async fn build_like_expr( + fn build_like_expr( case_insensitive: bool, negated: bool, f: &ScalarFunction, @@ -288,7 +304,7 @@ impl BuiltinExprBuilder { })) } - async fn build_binary_expr(fn_name: &str, args: Vec) -> Result { + fn build_binary_expr(fn_name: &str, args: Vec) -> Result { let [a, b] = match args.try_into() { Ok(args_arr) => args_arr, Err(_) => { @@ -312,7 +328,7 @@ impl BuiltinExprBuilder { Self::build_and_not_expr(or_expr, and_expr) } - async fn build_between_expr(fn_name: &str, args: Vec) -> Result { + fn build_between_expr(fn_name: &str, args: Vec) -> Result { let [expression, low, high] = match args.try_into() { Ok(args_arr) => args_arr, Err(_) => { @@ -329,18 +345,18 @@ impl BuiltinExprBuilder { } //This handles any functions that require custom handling - async fn build_custom_handling_expr( + fn build_custom_handling_expr( consumer: &impl SubstraitConsumer, fn_name: &str, args: Vec, ) -> Result { match fn_name { - "logb" => Self::build_logb_expr(consumer, args).await, + "logb" => Self::build_logb_expr(consumer, args), _ => not_impl_err!("Unsupported custom handled expression: {}", fn_name), } } - async fn build_logb_expr( + fn build_logb_expr( consumer: &impl SubstraitConsumer, args: Vec, ) -> Result { diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs index ac7d2479c397a..413ee4b537c29 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/aggregate_rel.rs @@ -109,11 +109,17 @@ pub async fn from_aggregate_rel( aggr_exprs.push(std::sync::Arc::unwrap_or_clone(agg_func?)); } - // Ensure that all expressions have a unique name + // Ensure that all expressions have a unique name. Both grouping and + // aggregate expressions become fields in the aggregate's output schema, + // so they share a single namespace. let mut name_tracker = NameTracker::new(); let group_exprs = group_exprs - .iter() - .map(|e| name_tracker.get_uniquely_named_expr(e.clone())) + .into_iter() + .map(|e| name_tracker.get_uniquely_named_expr(e)) + .collect::, _>>()?; + let aggr_exprs = aggr_exprs + .into_iter() + .map(|e| name_tracker.get_uniquely_named_expr(e)) .collect::, _>>()?; input.aggregate(group_exprs, aggr_exprs)?.build() diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs index 0a4048650fa2b..5aea6c809b701 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/project_rel.rs @@ -20,6 +20,7 @@ use crate::logical_plan::consumer::utils::NameTracker; use async_recursion::async_recursion; use datafusion::common::{Column, not_impl_err}; use datafusion::logical_expr::builder::project; +use datafusion::logical_expr::utils::find_window_exprs; use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; use std::collections::HashSet; use std::sync::Arc; @@ -57,13 +58,9 @@ pub async fn from_project_rel( let e = consumer .consume_expression(expr, input.clone().schema()) .await?; - // if the expression is WindowFunction, wrap in a Window relation - if let Expr::WindowFunction(_) = &e { - // Adding the same expression here and in the project below - // works because the project's builder uses columnize_expr(..) - // to transform it into a column reference - window_exprs.insert(e.clone()); - } + // The project's builder uses columnize_expr(..) to transform + // nested window expressions into column references. + window_exprs.extend(find_window_exprs([&e])); explicit_exprs.push(name_tracker.get_uniquely_named_expr(e)?); } diff --git a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs index 832110e11131c..78951a3aff549 100644 --- a/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/consumer/rel/read_rel.rs @@ -148,19 +148,48 @@ pub async fn from_read_rel( let values = if !vt.expressions.is_empty() { let mut exprs = vec![]; for row in &vt.expressions { + if row.fields.len() != substrait_schema.fields().len() { + return substrait_err!( + "Field count mismatch: expected {} fields but found {} in virtual table row", + substrait_schema.fields().len(), + row.fields.len() + ); + } + let mut row_exprs = vec![]; + let mut name_idx = 0; for expression in &row.fields { - let expr = consumer - .consume_expression(expression, &substrait_schema) - .await?; + // Top-level names are provided through schema + // Each expression consumes at least one name, and Literals may consume additional names. + name_idx += 1; + let expr = match expression.rex_type.as_ref() { + Some(substrait::proto::expression::RexType::Literal(lit)) => { + // Values literals need 'named_struct.names' so nested struct fields keep their names from the ReadRel base schema. + // This is important for nested struct fields to retain their names. + Expr::Literal( + from_substrait_literal( + consumer, + lit, + &named_struct.names, + &mut name_idx, + )?, + None, + ) + } + _ => { + consumer + .consume_expression(expression, &substrait_schema) + .await? + } + }; row_exprs.push(expr); } - // For expressions, validate against top-level schema fields, not nested names - if row_exprs.len() != substrait_schema.fields().len() { + + if name_idx != named_struct.names.len() { return substrait_err!( - "Field count mismatch: expected {} fields but found {} in virtual table row", - substrait_schema.fields().len(), - row_exprs.len() + "Names list must match exactly to nested schema, but found {} uses for {} names", + name_idx, + named_struct.names.len() ); } exprs.push(row_exprs); diff --git a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs index 65bc53ce0834e..bbd80b4cff001 100644 --- a/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs +++ b/datafusion/substrait/src/logical_plan/consumer/substrait_consumer.rs @@ -23,21 +23,27 @@ use super::{ from_substrait_rex, from_window_function, }; use crate::extensions::Extensions; +use crate::logical_plan::consumer::{ + field_from_substrait_type_without_names, from_lambda, +}; use async_trait::async_trait; -use datafusion::arrow::datatypes::DataType; +use datafusion::arrow::datatypes::{DataType, FieldRef}; use datafusion::catalog::TableProvider; +use datafusion::common::datatype::FieldExt; use datafusion::common::{ DFSchema, ScalarValue, TableReference, not_impl_err, substrait_err, }; use datafusion::execution::{FunctionRegistry, SessionState}; +use datafusion::logical_expr::expr::LambdaVariable; use datafusion::logical_expr::{Expr, Extension, LogicalPlan}; +use std::collections::VecDeque; use std::sync::{Arc, RwLock}; -use substrait::proto; use substrait::proto::expression as substrait_expression; use substrait::proto::expression::{ Enum, FieldReference, IfThen, Literal, MultiOrList, Nested, ScalarFunction, SingularOrList, SwitchExpression, WindowFunction, }; +use substrait::proto::{self, Type}; use substrait::proto::{ AggregateRel, ConsistentPartitionWindowRel, CrossRel, DynamicParameter, ExchangeRel, Expression, ExtensionLeafRel, ExtensionMultiRel, ExtensionSingleRel, FetchRel, @@ -62,17 +68,19 @@ use substrait::proto::{ /// # use datafusion::logical_expr::{Expr, LogicalPlan, LogicalPlanBuilder}; /// # use std::sync::Arc; /// # use substrait::proto; -/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel}; +/// # use substrait::proto::{ExtensionLeafRel, FilterRel, ProjectRel, Type}; /// # use datafusion::arrow::datatypes::DataType; /// # use datafusion::logical_expr::expr::ScalarFunction; /// # use datafusion_substrait::extensions::Extensions; /// # use datafusion_substrait::logical_plan::consumer::{ -/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer +/// # from_project_rel, from_substrait_rel, from_substrait_rex, SubstraitConsumer, DefaultSubstraitLambdaConsumer /// # }; /// /// struct CustomSubstraitConsumer { /// extensions: Arc, /// state: Arc, +/// // You can reuse existing consumer code related to lambdas +/// lambda_consumer: DefaultSubstraitLambdaConsumer, /// } /// /// #[async_trait] @@ -95,6 +103,30 @@ use substrait::proto::{ /// self.state.as_ref() /// } /// +/// fn push_lambda_parameters( +/// &self, +/// lambda_parameters: &[Type], +/// input_schema: &DFSchema, +/// ) -> datafusion::common::Result> { +/// self.lambda_consumer.push_lambda_parameters( +/// self, +/// lambda_parameters, +/// input_schema, +/// ) +/// } +/// +/// fn pop_lambda_parameters(&self) { +/// self.lambda_consumer.pop_lambda_parameters(); +/// } +/// +/// fn lambda_variable( +/// &self, +/// steps_out: usize, +/// field_idx: usize, +/// ) -> datafusion::common::Result { +/// self.lambda_consumer.lambda_variable(steps_out, field_idx) +/// } +/// /// // You can reuse existing consumer code to assist in handling advanced extensions /// async fn consume_project(&self, rel: &ProjectRel) -> Result { /// let df_plan = from_project_rel(self, rel).await?; @@ -384,6 +416,14 @@ pub trait SubstraitConsumer: Send + Sync + Sized { )) } + async fn consume_lambda( + &self, + expr: &proto::expression::Lambda, + input_schema: &DFSchema, + ) -> datafusion::common::Result { + from_lambda(self, expr, input_schema).await + } + // Outer Schema Stack // These methods manage a stack of outer schemas for correlated subquery support. // When entering a subquery, the enclosing query's schema is pushed onto the stack. @@ -481,6 +521,35 @@ pub trait SubstraitConsumer: Send + Sync + Sized { }; substrait_err!("Missing handler for user-defined literals {}", type_ref) } + + // Lambda related methods + + /// Push the given lambda parameters onto the stack when entering a lambda and + /// returns the names they got assigned + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaConsumer] and forward this method to it + fn push_lambda_parameters( + &self, + _lambda_parameters: &[Type], + _input_schema: &DFSchema, + ) -> datafusion::common::Result> { + not_impl_err!("SubstraitConsumer::push_lambda_parameters") + } + + /// Pop lambda parameters from the stack when leaving a lambda. + fn pop_lambda_parameters(&self) {} + + /// Returns an expression corresponding to the lambda variable with the given field_idx within the lambda it originates from, + /// at the lambda `step_outs` of the current scope + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaConsumer] and forward this method to it + fn lambda_variable( + &self, + _steps_out: usize, + _field_idx: usize, + ) -> datafusion::common::Result { + not_impl_err!("SubstraitConsumer::lambda_variable") + } } /// Default SubstraitConsumer for converting standard Substrait without user-defined extensions. @@ -490,6 +559,7 @@ pub struct DefaultSubstraitConsumer<'a> { pub(super) extensions: &'a Extensions, pub(super) state: &'a SessionState, outer_schemas: RwLock>>, + lambda_consumer: DefaultSubstraitLambdaConsumer, } impl<'a> DefaultSubstraitConsumer<'a> { @@ -498,6 +568,7 @@ impl<'a> DefaultSubstraitConsumer<'a> { extensions, state, outer_schemas: RwLock::new(Vec::new()), + lambda_consumer: DefaultSubstraitLambdaConsumer::new(), } } } @@ -594,6 +665,140 @@ impl SubstraitConsumer for DefaultSubstraitConsumer<'_> { let plan = plan.with_exprs_and_inputs(plan.expressions(), inputs)?; Ok(LogicalPlan::Extension(Extension { node: plan })) } + + fn push_lambda_parameters( + &self, + lambda_parameters: &[Type], + input_schema: &DFSchema, + ) -> datafusion::common::Result> { + self.lambda_consumer + .push_lambda_parameters(self, lambda_parameters, input_schema) + } + + fn pop_lambda_parameters(&self) { + self.lambda_consumer.pop_lambda_parameters() + } + + fn lambda_variable( + &self, + steps_out: usize, + field_idx: usize, + ) -> datafusion::common::Result { + self.lambda_consumer.lambda_variable(steps_out, field_idx) + } +} + +/// Default implementation of lambda related methods of the [SubstraitConsumer] trait +/// +/// Can be embedded into a custom [SubstraitConsumer] to implement them +pub struct DefaultSubstraitLambdaConsumer { + inner: RwLock, +} + +struct DefaultSubstraitLambdaConsumerInner { + /// Parameters of the lambdas currently in scope, ordered from innermost + /// to outermost. Index 0 is the lambda being consumed; higher indices + /// are enclosing lambdas, matching the `steps_out` value used by + /// [`DefaultSubstraitLambdaConsumer::lambda_variable`] and `LambdaParameterReference`. + lambda_parameters: VecDeque>, + next_lambda_parameter: usize, +} + +impl Default for DefaultSubstraitLambdaConsumer { + fn default() -> Self { + Self::new() + } +} + +impl DefaultSubstraitLambdaConsumer { + pub fn new() -> Self { + Self { + inner: RwLock::new(DefaultSubstraitLambdaConsumerInner { + lambda_parameters: VecDeque::new(), + next_lambda_parameter: 0, + }), + } + } + + pub fn push_lambda_parameters( + &self, + consumer: &impl SubstraitConsumer, + lambda_parameters: &[Type], + input_schema: &DFSchema, + ) -> datafusion::common::Result> { + let mut inner = self.inner.write().unwrap(); + + let lambda_parameters = lambda_parameters + .iter() + .map(|ty| { + let (assigned_number, default_name) = + next_lambda_parameter_name(inner.next_lambda_parameter, input_schema); + + inner.next_lambda_parameter = assigned_number + 1; + + Ok(field_from_substrait_type_without_names(consumer, ty)? + .renamed(&default_name)) + }) + .collect::>>()?; + + let names = lambda_parameters.iter().map(|f| f.name().clone()).collect(); + + inner.lambda_parameters.push_front(lambda_parameters); + + Ok(names) + } + + pub fn pop_lambda_parameters(&self) { + self.inner.write().unwrap().lambda_parameters.pop_front(); + } + + pub fn lambda_variable( + &self, + steps_out: usize, + field_idx: usize, + ) -> datafusion::common::Result { + let lambda_parameters = &self.inner.read().unwrap().lambda_parameters; + + let Some(lambda_parameters) = lambda_parameters.get(steps_out) else { + return substrait_err!( + "No lambda at {steps_out} steps out, got only {}", + lambda_parameters.len() + ); + }; + + let Some(var) = lambda_parameters.get(field_idx) else { + return substrait_err!( + "At lambda {steps_out} steps out, no field at index {field_idx}, got only {}", + lambda_parameters.len() + ); + }; + + Ok(Expr::LambdaVariable(LambdaVariable::new( + var.name().clone(), + Some(Arc::clone(var)), + ))) + } +} + +/// Returns the next available lambda parameter name and the index it was assigned. +/// +/// Names follow the pattern `pN` where `N` starts at `next_lambda_parameter`. If `pN` +/// conflicts with an existing column name in `input_schema`, `N` is incremented until +/// a free name is found. +fn next_lambda_parameter_name( + mut next_lambda_parameter: usize, + input_schema: &DFSchema, +) -> (usize, String) { + loop { + let name = format!("p{next_lambda_parameter}"); + + // avoid conflicts with column names + if !input_schema.has_column_with_unqualified_name(&name) { + return (next_lambda_parameter, name); + } + + next_lambda_parameter += 1; + } } #[cfg(test)] diff --git a/datafusion/substrait/src/logical_plan/consumer/utils.rs b/datafusion/substrait/src/logical_plan/consumer/utils.rs index c654cc070938d..824c79452d86e 100644 --- a/datafusion/substrait/src/logical_plan/consumer/utils.rs +++ b/datafusion/substrait/src/logical_plan/consumer/utils.rs @@ -18,12 +18,11 @@ use crate::logical_plan::consumer::SubstraitConsumer; use datafusion::arrow::datatypes::{DataType, Field, Schema, TimeUnit, UnionFields}; use datafusion::common::{ - DFSchema, DFSchemaRef, exec_err, not_impl_err, substrait_datafusion_err, - substrait_err, + DFSchema, DFSchemaRef, TableReference, exec_err, not_impl_err, + substrait_datafusion_err, substrait_err, }; use datafusion::logical_expr::expr::Sort; use datafusion::logical_expr::{Cast, Expr, ExprSchemable}; -use datafusion::sql::TableReference; use std::collections::HashSet; use std::sync::Arc; use substrait::proto::SortField; @@ -570,12 +569,11 @@ pub(crate) mod tests { use crate::extensions::Extensions; use crate::logical_plan::consumer::DefaultSubstraitConsumer; use datafusion::arrow::datatypes::{DataType, Field, Fields, Schema}; - use datafusion::common::DFSchema; + use datafusion::common::{DFSchema, TableReference}; use datafusion::error::Result; use datafusion::execution::SessionState; use datafusion::logical_expr::{Expr, col}; use datafusion::prelude::SessionContext; - use datafusion::sql::TableReference; use std::collections::HashMap; use std::sync::{Arc, LazyLock}; diff --git a/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs b/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs new file mode 100644 index 0000000000000..0d32dab2ccadd --- /dev/null +++ b/datafusion/substrait/src/logical_plan/producer/expr/lambda.rs @@ -0,0 +1,48 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::{common::DFSchemaRef, logical_expr::expr::Lambda}; +use substrait::proto::{ + Expression, + expression::RexType, + r#type::{Nullability, Struct}, +}; + +use crate::logical_plan::producer::SubstraitProducer; + +pub fn from_lambda( + producer: &mut impl SubstraitProducer, + lambda: &Lambda, + schema: &DFSchemaRef, +) -> Result { + Ok(Expression { + rex_type: Some(RexType::Lambda(Box::new( + substrait::proto::expression::Lambda { + parameters: Some(Struct { + nullability: Nullability::Required as i32, + type_variation_reference: 0, + types: lambda + .params + .iter() + .map(|p| producer.lambda_parameter_type(p)) + .collect::>()?, + }), + body: Some(Box::new(producer.handle_expr(&lambda.body, schema)?)), + }, + ))), + }) +} diff --git a/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs b/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs new file mode 100644 index 0000000000000..3d7f06e2332a1 --- /dev/null +++ b/datafusion/substrait/src/logical_plan/producer/expr/lambda_variable.rs @@ -0,0 +1,49 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::logical_expr::expr::LambdaVariable; +use substrait::proto::{ + Expression, + expression::{ + FieldReference, ReferenceSegment, RexType, + field_reference::{LambdaParameterReference, ReferenceType, RootType}, + reference_segment::{self, StructField}, + }, +}; + +use crate::logical_plan::producer::SubstraitProducer; + +pub fn from_lambda_variable( + producer: &mut impl SubstraitProducer, + lambda_variable: &LambdaVariable, + _schema: &datafusion::common::DFSchema, +) -> Result { + let (steps_out, field) = producer.lambda_variable(&lambda_variable.name)?; + + Ok(Expression { + rex_type: Some(RexType::Selection(Box::new(FieldReference { + reference_type: Some(ReferenceType::DirectReference(ReferenceSegment { + reference_type: Some(reference_segment::ReferenceType::StructField( + Box::new(StructField { field, child: None }), + )), + })), + root_type: Some(RootType::LambdaParameterReference( + LambdaParameterReference { steps_out }, + )), + }))), + }) +} diff --git a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs index 6e053f0d90a96..c728af2f1458d 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/mod.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/mod.rs @@ -19,6 +19,8 @@ mod aggregate_function; mod cast; mod field_reference; mod if_then; +mod lambda; +mod lambda_variable; mod literal; mod placeholder; mod scalar_function; @@ -30,6 +32,8 @@ pub use aggregate_function::*; pub use cast::*; pub use field_reference::*; pub use if_then::*; +pub use lambda::*; +pub use lambda_variable::*; pub use literal::*; pub use placeholder::*; pub use scalar_function::*; @@ -154,10 +158,8 @@ pub fn to_substrait_rex( Expr::HigherOrderFunction(expr) => { producer.handle_higher_order_function(expr, schema) } - Expr::Lambda(expr) => not_impl_err!("Cannot convert {expr:?} to Substrait"), - Expr::LambdaVariable(expr) => { - not_impl_err!("Cannot convert {expr:?} to Substrait") - } + Expr::Lambda(expr) => producer.handle_lambda(expr, schema), + Expr::LambdaVariable(expr) => producer.handle_lambda_variable(expr, schema), } } diff --git a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs index e36d5128cd293..75720395aae7c 100644 --- a/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs +++ b/datafusion/substrait/src/logical_plan/producer/expr/scalar_function.rs @@ -15,19 +15,35 @@ // specific language governing permissions and limitations // under the License. -use crate::logical_plan::producer::{SubstraitProducer, to_substrait_literal_expr}; -use datafusion::common::{DFSchemaRef, ScalarValue, not_impl_err}; -use datafusion::logical_expr::{Between, BinaryExpr, Expr, Like, Operator, expr}; +use crate::logical_plan::producer::{ + SubstraitProducer, to_substrait_literal_expr, to_substrait_type, +}; +use datafusion::arrow::datatypes::DataType; +use datafusion::common::datatype::FieldExt; +use datafusion::common::{ + DFSchemaRef, ScalarValue, internal_datafusion_err, not_impl_err, substrait_err, +}; +use datafusion::logical_expr::{ + Between, BinaryExpr, Expr, ExprSchemable, Like, Operator, expr, +}; use substrait::proto::expression::{RexType, ScalarFunction}; use substrait::proto::function_argument::ArgType; -use substrait::proto::{Expression, FunctionArgument}; +use substrait::proto::{Expression, FunctionArgument, Type}; pub fn from_scalar_function( producer: &mut impl SubstraitProducer, fun: &expr::ScalarFunction, schema: &DFSchemaRef, ) -> datafusion::common::Result { - from_function(producer, fun.name(), &fun.args, schema) + let (_, output_field) = Expr::ScalarFunction(fun.clone()).to_field(schema)?; + from_function( + producer, + fun.name(), + &fun.args, + output_field.data_type(), + output_field.is_nullable(), + schema, + ) } pub fn from_higher_order_function( @@ -35,13 +51,94 @@ pub fn from_higher_order_function( fun: &expr::HigherOrderFunction, schema: &DFSchemaRef, ) -> datafusion::common::Result { - from_function(producer, fun.name(), &fun.args, schema) + let mut lambda_parameters = fun.lambda_parameters(schema)?.into_iter(); + + let num_lambdas = fun + .args + .iter() + .filter(|arg| matches!(arg, Expr::Lambda(_))) + .count(); + + if lambda_parameters.len() != num_lambdas { + return substrait_err!( + "{} returned {} lambdas but {num_lambdas} expected", + fun.name(), + lambda_parameters.len() + ); + } + + let arguments = fun + .args + .iter() + .map(|arg| { + let arg = match arg { + Expr::Lambda(l) => { + let lambda_parameters = + lambda_parameters.next().ok_or_else(|| { + internal_datafusion_err!( + "lambda_parameters len should have been checked above" + ) + })?; + + if l.params.len() > lambda_parameters.len() { + return substrait_err!( + "Lambda defined {} parameters ({}) but function {} supports only {}", + l.params.len(), + l.params.join(","), + fun.name(), + lambda_parameters.len() + ) + } + + let named_lambda_parameters = + std::iter::zip(&l.params, lambda_parameters) + .map(|(name, parameter)| parameter.renamed(name)) + .collect(); + + producer.push_lambda_parameters(named_lambda_parameters)?; + + let arg = producer.handle_lambda(l, schema); + + producer.pop_lambda_parameters()?; + + arg + } + _ => producer.handle_expr(arg, schema), + }?; + + Ok(FunctionArgument { + arg_type: Some(ArgType::Value(arg)), + }) + }) + .collect::>()?; + + let function_anchor = producer.register_function(fun.name().to_string()); + + let (_, output_field) = Expr::HigherOrderFunction(fun.clone()).to_field(schema)?; + let output_type = to_substrait_type( + producer, + output_field.data_type(), + output_field.is_nullable(), + )?; + + #[expect(deprecated)] + Ok(Expression { + rex_type: Some(RexType::ScalarFunction(ScalarFunction { + function_reference: function_anchor, + arguments, + output_type: Some(output_type), + options: vec![], + args: vec![], + })), + }) } fn from_function( producer: &mut impl SubstraitProducer, name: &str, args: &[Expr], + output_type: &DataType, + output_nullability: bool, schema: &DFSchemaRef, ) -> datafusion::common::Result { let mut arguments: Vec = vec![]; @@ -52,6 +149,7 @@ fn from_function( } let arguments = custom_argument_handler(name, arguments); + let output_type = to_substrait_type(producer, output_type, output_nullability)?; let function_anchor = producer.register_function(name.to_string()); #[expect(deprecated)] @@ -59,7 +157,7 @@ fn from_function( rex_type: Some(RexType::ScalarFunction(ScalarFunction { function_reference: function_anchor, arguments, - output_type: None, + output_type: Some(output_type), options: vec![], args: vec![], })), @@ -103,7 +201,13 @@ pub fn from_unary_expr( Expr::Negative(arg) => ("negate", arg), expr => not_impl_err!("Unsupported expression: {expr:?}")?, }; - to_substrait_unary_scalar_fn(producer, fn_name, arg, schema) + let (_, output_field) = expr.to_field(schema)?; + let output_type = to_substrait_type( + producer, + output_field.data_type(), + output_field.is_nullable(), + )?; + to_substrait_unary_scalar_fn(producer, fn_name, arg, schema, &output_type) } pub fn from_binary_expr( @@ -114,7 +218,19 @@ pub fn from_binary_expr( let BinaryExpr { left, op, right } = expr; let l = producer.handle_expr(left, schema)?; let r = producer.handle_expr(right, schema)?; - Ok(make_binary_op_scalar_func(producer, &l, &r, *op)) + let (_, output_field) = Expr::BinaryExpr(expr.clone()).to_field(schema)?; + let output_type = to_substrait_type( + producer, + output_field.data_type(), + output_field.is_nullable(), + )?; + Ok(make_binary_op_scalar_func( + producer, + &l, + &r, + *op, + &output_type, + )) } pub fn from_like( @@ -209,6 +325,7 @@ fn to_substrait_unary_scalar_fn( fn_name: &str, arg: &Expr, schema: &DFSchemaRef, + output_type: &Type, ) -> datafusion::common::Result { let function_anchor = producer.register_function(fn_name.to_string()); let substrait_expr = producer.handle_expr(arg, schema)?; @@ -219,7 +336,7 @@ fn to_substrait_unary_scalar_fn( arguments: vec![FunctionArgument { arg_type: Some(ArgType::Value(substrait_expr)), }], - output_type: None, + output_type: Some(output_type.clone()), options: vec![], ..Default::default() })), @@ -232,6 +349,7 @@ pub fn make_binary_op_scalar_func( lhs: &Expression, rhs: &Expression, op: Operator, + output_type: &Type, ) -> Expression { let function_anchor = producer.register_function(operator_to_name(op).to_string()); #[expect(deprecated)] @@ -246,7 +364,7 @@ pub fn make_binary_op_scalar_func( arg_type: Some(ArgType::Value(rhs.clone())), }, ], - output_type: None, + output_type: Some(output_type.clone()), args: vec![], options: vec![], })), @@ -264,57 +382,21 @@ pub fn from_between( low, high, } = between; - if *negated { - // `expr NOT BETWEEN low AND high` can be translated into (expr < low OR high < expr) - let substrait_expr = producer.handle_expr(expr.as_ref(), schema)?; - let substrait_low = producer.handle_expr(low.as_ref(), schema)?; - let substrait_high = producer.handle_expr(high.as_ref(), schema)?; - - let l_expr = make_binary_op_scalar_func( - producer, - &substrait_expr, - &substrait_low, - Operator::Lt, - ); - let r_expr = make_binary_op_scalar_func( - producer, - &substrait_high, - &substrait_expr, - Operator::Lt, - ); - Ok(make_binary_op_scalar_func( - producer, - &l_expr, - &r_expr, - Operator::Or, - )) + let expr = if *negated { + // `expr NOT BETWEEN low AND high` can be translated into (expr < low OR high < expr) + Expr::or( + Expr::lt(*expr.clone(), *low.clone()), + Expr::lt(*high.clone(), *expr.clone()), + ) } else { // `expr BETWEEN low AND high` can be translated into (low <= expr AND expr <= high) - let substrait_expr = producer.handle_expr(expr.as_ref(), schema)?; - let substrait_low = producer.handle_expr(low.as_ref(), schema)?; - let substrait_high = producer.handle_expr(high.as_ref(), schema)?; - - let l_expr = make_binary_op_scalar_func( - producer, - &substrait_low, - &substrait_expr, - Operator::LtEq, - ); - let r_expr = make_binary_op_scalar_func( - producer, - &substrait_expr, - &substrait_high, - Operator::LtEq, - ); - - Ok(make_binary_op_scalar_func( - producer, - &l_expr, - &r_expr, - Operator::And, - )) - } + Expr::and( + Expr::lt_eq(*low.clone(), *expr.clone()), + Expr::lt_eq(*expr.clone(), *high.clone()), + ) + }; + producer.handle_expr(&expr, schema) } pub fn operator_to_name(op: Operator) -> &'static str { @@ -364,3 +446,37 @@ pub fn operator_to_name(op: Operator) -> &'static str { Operator::Colon => "colon", } } + +#[cfg(test)] +mod tests { + use crate::logical_plan::producer::{ + DefaultSubstraitProducer, SubstraitProducer, to_substrait_type, + }; + use datafusion::arrow::datatypes::DataType; + use datafusion::common::{DFSchema, DFSchemaRef}; + use datafusion::execution::SessionStateBuilder; + use datafusion::prelude::lit; + use substrait::proto::Expression; + use substrait::proto::expression::{RexType, ScalarFunction}; + + #[tokio::test] + async fn binary_expr_output_type() -> datafusion::common::Result<()> { + let state = SessionStateBuilder::default().build(); + let empty_schema = DFSchemaRef::new(DFSchema::empty()); + let mut producer = DefaultSubstraitProducer::new(&state); + + let expr = lit(1i64) + lit(2i64); + let substrait_expr = producer.handle_expr(&expr, &empty_schema)?; + if let Expression { + rex_type: Some(RexType::ScalarFunction(ScalarFunction { output_type, .. })), + } = substrait_expr + { + let expected_type = + to_substrait_type(&mut producer, &DataType::Int64, false)?; + assert_eq!(output_type, Some(expected_type)); + Ok(()) + } else { + panic!("Substrait ScalarFunction expected") + } + } +} diff --git a/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs index 50c4b3da86cbe..1b9e91c7c475a 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/exchange_rel.rs @@ -32,6 +32,11 @@ pub fn from_repartition( let partition_count = match repartition.partitioning_scheme { Partitioning::RoundRobinBatch(num) => num, Partitioning::Hash(_, num) => num, + Partitioning::Range(_) => { + // TODO: Support range repartitioning in Substrait exchange output. + // Tracked by https://github.com/apache/datafusion/issues/22788 + return not_impl_err!("Substrait does not support Range repartitioning"); + } Partitioning::DistributeBy(_) => { return not_impl_err!( "Physical plan does not support DistributeBy partitioning" @@ -50,6 +55,11 @@ pub fn from_repartition( .collect::>>()?; ExchangeKind::ScatterByFields(ScatterFields { fields }) } + Partitioning::Range(_) => { + // TODO: Support range repartitioning in Substrait exchange output. + // Tracked by https://github.com/apache/datafusion/issues/22788 + return not_impl_err!("Substrait does not support Range repartitioning"); + } Partitioning::DistributeBy(_) => { return not_impl_err!( "Physical plan does not support DistributeBy partitioning" diff --git a/datafusion/substrait/src/logical_plan/producer/rel/join.rs b/datafusion/substrait/src/logical_plan/producer/rel/join.rs index cbf5593ffc86c..9094774780e10 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/join.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/join.rs @@ -15,59 +15,38 @@ // specific language governing permissions and limitations // under the License. -use crate::logical_plan::producer::{SubstraitProducer, make_binary_op_scalar_func}; -use datafusion::common::{ - DFSchemaRef, JoinConstraint, JoinType, NullEquality, not_impl_err, -}; +use crate::logical_plan::producer::SubstraitProducer; +use datafusion::common::{JoinConstraint, JoinType, NullEquality, not_impl_err}; +use datafusion::logical_expr::utils::conjunction; use datafusion::logical_expr::{Expr, Join, Operator}; +use datafusion::prelude::binary_expr; use std::sync::Arc; use substrait::proto::rel::RelType; -use substrait::proto::{Expression, JoinRel, Rel, join_rel}; +use substrait::proto::{JoinRel, Rel, join_rel}; pub fn from_join( producer: &mut impl SubstraitProducer, join: &Join, ) -> datafusion::common::Result> { - let left = producer.handle_plan(join.left.as_ref())?; - let right = producer.handle_plan(join.right.as_ref())?; - let join_type = to_substrait_jointype(join.join_type); - // we only support basic joins so return an error for anything not yet supported + // only ON constraints are supported right now match join.join_constraint { JoinConstraint::On => {} JoinConstraint::Using => return not_impl_err!("join constraint: `using`"), } - let in_join_schema = Arc::new(join.left.schema().join(join.right.schema())?); - - // convert filter if present - let join_filter = match &join.filter { - Some(filter) => Some(producer.handle_expr(filter, &in_join_schema)?), - None => None, - }; - // map the left and right columns to binary expressions in the form `l = r` - // build a single expression for the ON condition, such as `l.a = r.a AND l.b = r.b` - let eq_op = match join.null_equality { - NullEquality::NullEqualsNothing => Operator::Eq, - NullEquality::NullEqualsNull => Operator::IsNotDistinctFrom, - }; - let join_on = to_substrait_join_expr(producer, &join.on, eq_op, &in_join_schema)?; + let left = producer.handle_plan(join.left.as_ref())?; + let right = producer.handle_plan(join.right.as_ref())?; + let join_type = to_substrait_jointype(join.join_type); - // create conjunction between `join_on` and `join_filter` to embed all join conditions, - // whether equal or non-equal in a single expression - let join_expr = match &join_on { - Some(on_expr) => match &join_filter { - Some(filter) => Some(Box::new(make_binary_op_scalar_func( - producer, - on_expr, - filter, - Operator::And, - ))), - None => join_on.map(Box::new), // the join expression will only contain `join_on` if filter doesn't exist - }, - None => match &join_filter { - Some(_) => join_filter.map(Box::new), // the join expression will only contain `join_filter` if the `on` condition doesn't exist - None => None, - }, + let join_expr = + to_substrait_join_expr(join.on.clone(), join.null_equality, join.filter.clone()); + let join_expression = match join_expr { + Some(expr) => { + let in_join_schema = Arc::new(join.left.schema().join(join.right.schema())?); + let expression = producer.handle_expr(&expr, &in_join_schema)?; + Some(Box::new(expression)) + } + None => None, }; Ok(Box::new(Rel { @@ -76,7 +55,7 @@ pub fn from_join( left: Some(left), right: Some(right), r#type: join_type as i32, - expression: join_expr, + expression: join_expression, post_join_filter: None, advanced_extension: None, }))), @@ -84,25 +63,20 @@ pub fn from_join( } fn to_substrait_join_expr( - producer: &mut impl SubstraitProducer, - join_conditions: &Vec<(Expr, Expr)>, - eq_op: Operator, - join_schema: &DFSchemaRef, -) -> datafusion::common::Result> { - // Only support AND conjunction for each binary expression in join conditions - let mut exprs: Vec = vec![]; - for (left, right) in join_conditions { - let l = producer.handle_expr(left, join_schema)?; - let r = producer.handle_expr(right, join_schema)?; - // AND with existing expression - exprs.push(make_binary_op_scalar_func(producer, &l, &r, eq_op)); - } - - let join_expr: Option = - exprs.into_iter().reduce(|acc: Expression, e: Expression| { - make_binary_op_scalar_func(producer, &acc, &e, Operator::And) - }); - Ok(join_expr) + join_on: Vec<(Expr, Expr)>, + null_equality: NullEquality, + join_filter: Option, +) -> Option { + // Combine join on and filter conditions into a single Boolean expression (#7611) + let eq_op = match null_equality { + NullEquality::NullEqualsNothing => Operator::Eq, + NullEquality::NullEqualsNull => Operator::IsNotDistinctFrom, + }; + let all_conditions = join_on + .into_iter() + .map(|(left, right)| binary_expr(left, eq_op, right)) + .chain(join_filter); + conjunction(all_conditions) } fn to_substrait_jointype(join_type: JoinType) -> join_rel::JoinType { @@ -119,3 +93,85 @@ fn to_substrait_jointype(join_type: JoinType) -> join_rel::JoinType { JoinType::RightSemi => join_rel::JoinType::RightSemi, } } + +#[cfg(test)] +mod tests { + use crate::logical_plan::producer::{ + DefaultSubstraitProducer, SubstraitProducer, to_substrait_type, + }; + use datafusion::arrow::datatypes::{DataType, Field, Schema}; + use datafusion::common::{JoinConstraint, JoinType, NullEquality}; + use datafusion::execution::SessionStateBuilder; + use datafusion::logical_expr::utils::conjunction; + use datafusion::logical_expr::{Join, col, table_scan}; + use std::sync::Arc; + use substrait::proto::expression::{RexType, ScalarFunction}; + use substrait::proto::rel::RelType; + use substrait::proto::{Expression, JoinRel, Rel, join_rel}; + + #[test] + fn test_from_join() -> datafusion::common::Result<()> { + let state = SessionStateBuilder::default().build(); + let mut producer = DefaultSubstraitProducer::new(&state); + + let schema = Schema::new(vec![ + Field::new("a", DataType::Int32, false), + Field::new("b", DataType::Int32, false), + Field::new("c", DataType::Int32, false), + ]); + let left_scan = table_scan(Some("t1"), &schema, None)?.build()?; + let right_scan = table_scan(Some("t2"), &schema, None)?.build()?; + let join = Join::try_new( + Arc::new(left_scan.clone()), + Arc::new(right_scan.clone()), + vec![(col("t1.a"), col("t2.a")), (col("t1.b"), col("t2.b"))], + Some(col("t1.c").gt(col("t2.c"))), + JoinType::Inner, + JoinConstraint::On, + NullEquality::NullEqualsNothing, + false, + )?; + let join_expr = producer.handle_join(&join)?; + + let in_join_schema = Arc::new(join.left.schema().join(join.right.schema())?); + let expected_join_expr = conjunction(vec![ + // Join on + col("t1.a").eq(col("t2.a")), + col("t1.b").eq(col("t2.b")), + // Join filter + col("t1.c").gt(col("t2.c")), + ]) + .unwrap(); + let expected_join_expression = + producer.handle_expr(&expected_join_expr, &in_join_schema)?; + + assert_eq!( + join_expr, + Box::new(Rel { + rel_type: Some(RelType::Join(Box::new(JoinRel { + common: None, + left: Some(producer.handle_plan(&left_scan)?), + right: Some(producer.handle_plan(&right_scan)?), + r#type: join_rel::JoinType::Inner as i32, + expression: Some(Box::new(expected_join_expression.clone())), + post_join_filter: None, + advanced_extension: None, + }))) + }) + ); + + // Check that the join_expression has the expected output_type + if let Expression { + rex_type: Some(RexType::ScalarFunction(ScalarFunction { output_type, .. })), + } = expected_join_expression + { + let expected_type = + to_substrait_type(&mut producer, &DataType::Boolean, false)?; + assert_eq!(output_type, Some(expected_type)); + } else { + panic!("Substrait ScalarFunction expected") + } + + Ok(()) + } +} diff --git a/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs b/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs index 8dfbb36d3767d..900273bf8e6d7 100644 --- a/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs +++ b/datafusion/substrait/src/logical_plan/producer/rel/read_rel.rs @@ -15,55 +15,19 @@ // specific language governing permissions and limitations // under the License. -use crate::logical_plan::producer::{ - SubstraitProducer, to_substrait_literal, to_substrait_named_struct, -}; +use crate::logical_plan::producer::{SubstraitProducer, to_substrait_named_struct}; use datafusion::common::{DFSchema, ToDFSchema, substrait_datafusion_err}; use datafusion::logical_expr::utils::conjunction; use datafusion::logical_expr::{EmptyRelation, Expr, TableScan, Values}; use datafusion::scalar::ScalarValue; use std::sync::Arc; use substrait::proto::expression::MaskExpression; -use substrait::proto::expression::literal::Struct as LiteralStruct; use substrait::proto::expression::mask_expression::{StructItem, StructSelect}; use substrait::proto::expression::nested::Struct as NestedStruct; use substrait::proto::read_rel::{NamedTable, ReadType, VirtualTable}; use substrait::proto::rel::RelType; use substrait::proto::{ReadRel, Rel}; -/// Converts rows of literal expressions into Substrait literal structs. -/// -/// Each row is expected to contain only `Expr::Literal` or `Expr::Alias` wrapping literals. -/// Aliases are unwrapped and the underlying literal is converted. -fn convert_literal_rows( - producer: &mut impl SubstraitProducer, - rows: &[Vec], -) -> datafusion::common::Result> { - rows.iter() - .map(|row| { - let fields = row - .iter() - .map(|expr| match expr { - Expr::Literal(sv, _) => to_substrait_literal(producer, sv), - Expr::Alias(alias) => match alias.expr.as_ref() { - // The schema gives us the names, so we can skip aliases - Expr::Literal(sv, _) => to_substrait_literal(producer, sv), - _ => Err(substrait_datafusion_err!( - "Only literal types can be aliased in Virtual Tables, got: {}", - alias.expr.variant_name() - )), - }, - _ => Err(substrait_datafusion_err!( - "Only literal types and aliases are supported in Virtual Tables, got: {}", - expr.variant_name() - )), - }) - .collect::>()?; - Ok(LiteralStruct { fields }) - }) - .collect() -} - /// Converts rows of arbitrary expressions into Substrait nested structs. /// /// Validates that each row has the expected schema length and converts each expression @@ -163,6 +127,7 @@ pub fn from_empty_relation( let base_schema = to_substrait_named_struct(producer, &e.schema)?; let read_type = if e.produce_one_row { + let empty_schema = Arc::new(DFSchema::empty()); // Create one row with default scalar values for each field in the schema. // For example, an Int32 field gets Int32(NULL), a Utf8 field gets Utf8(NULL), etc. // This represents the "phantom row" that provides a context for evaluating @@ -173,25 +138,16 @@ pub fn from_empty_relation( .iter() .map(|f| { let scalar = ScalarValue::try_from(f.data_type())?; - to_substrait_literal(producer, &scalar) + producer.handle_expr(&Expr::Literal(scalar, None), &empty_schema) }) .collect::>()?; ReadType::VirtualTable(VirtualTable { - // Use deprecated 'values' field instead of 'expressions' because the consumer's - // nested expression support (RexType::Nested) is not yet implemented. - // The 'values' field uses literal::Struct which the consumer can properly - // deserialize with field name preservation. - #[expect(deprecated)] - values: vec![LiteralStruct { fields }], - expressions: vec![], + expressions: vec![NestedStruct { fields }], + ..Default::default() }) } else { - ReadType::VirtualTable(VirtualTable { - #[expect(deprecated)] - values: vec![], - expressions: vec![], - }) + ReadType::VirtualTable(VirtualTable::default()) }; Ok(Box::new(Rel { rel_type: Some(RelType::Read(Box::new(ReadRel { @@ -212,23 +168,8 @@ pub fn from_values( ) -> datafusion::common::Result> { let schema_len = v.schema.fields().len(); let empty_schema = Arc::new(DFSchema::empty()); - - let use_literals = v.values.iter().all(|row| { - row.iter().all(|expr| match expr { - Expr::Literal(_, _) => true, - Expr::Alias(alias) => matches!(alias.expr.as_ref(), Expr::Literal(_, _)), - _ => false, - }) - }); - - let (values, expressions) = if use_literals { - let values = convert_literal_rows(producer, &v.values)?; - (values, vec![]) - } else { - let expressions = - convert_expression_rows(producer, &v.values, schema_len, &empty_schema)?; - (vec![], expressions) - }; + let expressions = + convert_expression_rows(producer, &v.values, schema_len, &empty_schema)?; Ok(Box::new(Rel { rel_type: Some(RelType::Read(Box::new(ReadRel { common: None, @@ -237,10 +178,9 @@ pub fn from_values( best_effort_filter: None, projection: None, advanced_extension: None, - #[expect(deprecated)] read_type: Some(ReadType::VirtualTable(VirtualTable { - values, expressions, + ..Default::default() })), }))), })) diff --git a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs index 4228c32435897..6d54d32cad3db 100644 --- a/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs +++ b/datafusion/substrait/src/logical_plan/producer/substrait_producer.rs @@ -20,18 +20,23 @@ use crate::logical_plan::producer::{ from_aggregate, from_aggregate_function, from_alias, from_between, from_binary_expr, from_case, from_cast, from_column, from_distinct, from_empty_relation, from_exists, from_filter, from_higher_order_function, from_in_list, from_in_subquery, from_join, - from_like, from_limit, from_literal, from_placeholder, from_projection, - from_repartition, from_scalar_function, from_scalar_subquery, from_set_comparison, - from_sort, from_subquery_alias, from_table_scan, from_try_cast, from_unary_expr, - from_union, from_values, from_window, from_window_function, to_substrait_rel, - to_substrait_rex, + from_lambda, from_lambda_variable, from_like, from_limit, from_literal, + from_placeholder, from_projection, from_repartition, from_scalar_function, + from_scalar_subquery, from_set_comparison, from_sort, from_subquery_alias, + from_table_scan, from_try_cast, from_unary_expr, from_union, from_values, + from_window, from_window_function, to_substrait_rel, to_substrait_rex, + to_substrait_type_from_field, +}; +use datafusion::arrow::datatypes::FieldRef; +use datafusion::common::{ + Column, DFSchemaRef, HashMap, ScalarValue, not_impl_err, substrait_err, }; -use datafusion::common::{Column, DFSchemaRef, ScalarValue, substrait_err}; use datafusion::execution::SessionState; use datafusion::execution::registry::SerializerRegistry; use datafusion::logical_expr::Subquery; use datafusion::logical_expr::expr::{ - Alias, Exists, InList, InSubquery, Placeholder, SetComparison, WindowFunction, + Alias, Exists, InList, InSubquery, Lambda, LambdaVariable, Placeholder, + SetComparison, WindowFunction, }; use datafusion::logical_expr::{ Aggregate, Between, BinaryExpr, Case, Cast, Distinct, EmptyRelation, Expr, Extension, @@ -57,16 +62,19 @@ use substrait::proto::{ /// # use std::sync::Arc; /// # use substrait::proto::{Expression, Rel}; /// # use substrait::proto::rel::RelType; +/// # use datafusion::arrow::datatypes::FieldRef; /// # use datafusion::common::DFSchemaRef; /// # use datafusion::error::Result; /// # use datafusion::execution::SessionState; /// # use datafusion::logical_expr::{Between, Extension, Projection}; /// # use datafusion_substrait::extensions::Extensions; -/// # use datafusion_substrait::logical_plan::producer::{from_projection, SubstraitProducer}; +/// # use datafusion_substrait::logical_plan::producer::{from_projection, SubstraitProducer, DefaultSubstraitLambdaProducer, lambda_parameters_map}; /// /// struct CustomSubstraitProducer { /// extensions: Extensions, /// state: Arc, +/// // You can reuse existing producer code related to lambdas +/// lambda_producer: DefaultSubstraitLambdaProducer, /// } /// /// impl SubstraitProducer for CustomSubstraitProducer { @@ -83,6 +91,33 @@ use substrait::proto::{ /// self.extensions /// } /// +/// fn push_lambda_parameters( +/// &mut self, +/// lambda_parameters: Vec, +/// ) -> datafusion::common::Result<()> { +/// let lambda_parameters_map = lambda_parameters_map(self, lambda_parameters)?; +/// +/// self.lambda_producer +/// .push_lambda_parameters(lambda_parameters_map); +/// +/// Ok(()) +/// } +/// +/// fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { +/// self.lambda_producer.pop_lambda_parameters() +/// } +/// +/// fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { +/// self.lambda_producer.lambda_variable(name) +/// } +/// +/// fn lambda_parameter_type( +/// &self, +/// name: &str, +/// ) -> datafusion::common::Result { +/// self.lambda_producer.lambda_parameter_type(name) +/// } +/// /// // You can set additional metadata on the Rels you produce /// fn handle_projection(&mut self, plan: &Projection) -> Result> { /// let mut rel = from_projection(self, plan)?; @@ -405,11 +440,65 @@ pub trait SubstraitProducer: Send + Sync + Sized { ) -> datafusion::common::Result { from_placeholder(self, placeholder) } + + fn handle_lambda( + &mut self, + lambda: &Lambda, + schema: &DFSchemaRef, + ) -> datafusion::common::Result { + from_lambda(self, lambda, schema) + } + + fn handle_lambda_variable( + &mut self, + lambda_variable: &LambdaVariable, + schema: &DFSchemaRef, + ) -> datafusion::common::Result { + from_lambda_variable(self, lambda_variable, schema) + } + + // Lambda related methods + + /// Push the given `lambda_parameters` into this producer so they can be referenced by lambda variables + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn push_lambda_parameters( + &mut self, + _lambda_parameters: Vec, + ) -> datafusion::common::Result<()> { + not_impl_err!("SubstraitProducer::push_lambda_parameters") + } + + /// Pop the last pushed `lambda_parameters` so that it unshadow any previously shadowed lambda parameter + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + not_impl_err!("SubstraitProducer::pop_lambda_parameters") + } + + /// Get the (`steps_out`, `field_idx`) of the lambda variable with the given `name`. `steps_out` refers to the number + /// of lambda boundaries to traverse (0 = current lambda), and `field_idx` refers to the index within the lambda parameters + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn lambda_variable(&self, _name: &str) -> datafusion::common::Result<(u32, i32)> { + not_impl_err!("SubstraitProducer::lambda_variable") + } + + /// Get the type of the lambda parameter with the given `name` + /// + /// Note for custom implementations it's possible to embed a [DefaultSubstraitLambdaProducer] and forward this method to it + fn lambda_parameter_type( + &self, + _name: &str, + ) -> datafusion::common::Result { + not_impl_err!("SubstraitProducer::lambda_parameter_type") + } } pub struct DefaultSubstraitProducer<'a> { extensions: Extensions, serializer_registry: &'a dyn SerializerRegistry, + lambda_producer: DefaultSubstraitLambdaProducer, } impl<'a> DefaultSubstraitProducer<'a> { @@ -417,6 +506,7 @@ impl<'a> DefaultSubstraitProducer<'a> { DefaultSubstraitProducer { extensions: Extensions::default(), serializer_registry: state.serializer_registry().as_ref(), + lambda_producer: DefaultSubstraitLambdaProducer::new(), } } } @@ -471,4 +561,109 @@ impl SubstraitProducer for DefaultSubstraitProducer<'_> { rel_type: Some(rel_type), })) } + + fn push_lambda_parameters( + &mut self, + lambda_parameters: Vec, + ) -> datafusion::common::Result<()> { + let lambda_parameters_map = lambda_parameters_map(self, lambda_parameters)?; + + self.lambda_producer + .push_lambda_parameters(lambda_parameters_map); + + Ok(()) + } + + fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + self.lambda_producer.pop_lambda_parameters() + } + + fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { + self.lambda_producer.lambda_variable(name) + } + + fn lambda_parameter_type( + &self, + name: &str, + ) -> datafusion::common::Result { + self.lambda_producer.lambda_parameter_type(name) + } +} + +/// Default implementation of lambda related methods of the [SubstraitProducer] trait +/// +/// Can be embedded into a custom [SubstraitProducer] to implement them +pub struct DefaultSubstraitLambdaProducer { + lambdas_variables: Vec>, +} + +impl Default for DefaultSubstraitLambdaProducer { + fn default() -> Self { + Self::new() + } +} + +impl DefaultSubstraitLambdaProducer { + pub fn new() -> Self { + Self { + lambdas_variables: Vec::new(), + } + } + + /// Note you can construct the `lambda_parameters` argument using [lambda_parameters_map] + pub fn push_lambda_parameters( + &mut self, + lambda_parameters: HashMap, + ) { + self.lambdas_variables.push(lambda_parameters); + } + + pub fn pop_lambda_parameters(&mut self) -> datafusion::common::Result<()> { + match self.lambdas_variables.pop() { + Some(_) => Ok(()), + None => substrait_err!("no lambda_parameters to pop"), + } + } + + pub fn lambda_variable(&self, name: &str) -> datafusion::common::Result<(u32, i32)> { + for (steps_out, lambda_parameters) in + self.lambdas_variables.iter().rev().enumerate() + { + if let Some((field_idx, _type)) = lambda_parameters.get(name) { + return Ok((steps_out as u32, *field_idx as i32)); + } + } + + substrait_err!("unknown lambda variable {name}") + } + + pub fn lambda_parameter_type( + &self, + name: &str, + ) -> datafusion::common::Result { + for lambda_parameters in self.lambdas_variables.iter().rev() { + if let Some((_field_idx, type_)) = lambda_parameters.get(name) { + return Ok(type_.clone()); + } + } + + substrait_err!("unknown lambda variable {name}") + } +} + +/// Produces a map of lambda parameters as expected by [DefaultSubstraitLambdaProducer::push_lambda_parameters] +pub fn lambda_parameters_map( + producer: &mut impl SubstraitProducer, + lambda_parameters: Vec, +) -> datafusion::common::Result> { + lambda_parameters + .into_iter() + .enumerate() + .map(|(field_idx, field)| { + Ok(( + field.name().clone(), + (field_idx, to_substrait_type_from_field(producer, &field)?), + )) + }) + .collect::>() } diff --git a/datafusion/substrait/src/physical_plan/producer.rs b/datafusion/substrait/src/physical_plan/producer.rs index 17ca99ceff6e4..21282b9e8b48d 100644 --- a/datafusion/substrait/src/physical_plan/producer.rs +++ b/datafusion/substrait/src/physical_plan/producer.rs @@ -74,13 +74,9 @@ pub fn to_substrait_rel( let mut types = vec![]; for field in file_config.file_schema().fields.iter() { - match to_substrait_type(field.data_type(), field.is_nullable()) { - Ok(t) => { - names.push(field.name().clone()); - types.push(t); - } - Err(e) => return Err(e), - } + let t = to_substrait_type(field.data_type(), field.is_nullable())?; + names.push(field.name().clone()); + types.push(t); } let type_info = Struct { diff --git a/datafusion/substrait/src/serializer.rs b/datafusion/substrait/src/serializer.rs index ee71bc3121afe..bcc9f5cf50eac 100644 --- a/datafusion/substrait/src/serializer.rs +++ b/datafusion/substrait/src/serializer.rs @@ -70,12 +70,12 @@ pub async fn deserialize(path: impl AsRef) -> Result> { let mut file = OpenOptions::new().read(true).open(path).await?; file.read_to_end(&mut protobuf_in).await?; - deserialize_bytes(protobuf_in).await + deserialize_bytes(&protobuf_in) } /// Deserializes a plan from the bytes. -pub async fn deserialize_bytes(proto_bytes: Vec) -> Result> { - Ok(Box::new(Message::decode(&*proto_bytes).map_err(|e| { +pub fn deserialize_bytes(proto_bytes: &[u8]) -> Result> { + Ok(Box::new(Message::decode(proto_bytes).map_err(|e| { DataFusionError::Substrait(format!("Failed to decode plan: {e}")) })?)) } diff --git a/datafusion/substrait/tests/cases/consumer_integration.rs b/datafusion/substrait/tests/cases/consumer_integration.rs index b5d9f36620c67..1f30a753772cb 100644 --- a/datafusion/substrait/tests/cases/consumer_integration.rs +++ b/datafusion/substrait/tests/cases/consumer_integration.rs @@ -207,7 +207,7 @@ mod tests { @r#" Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE * LINEITEM.L_DISCOUNT) AS REVENUE]] Projection: LINEITEM.L_EXTENDEDPRICE * LINEITEM.L_DISCOUNT - Filter: LINEITEM.L_SHIPDATE >= CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-01-01") AS Date32) AND LINEITEM.L_DISCOUNT >= Decimal128(Some(5),3,2) AND LINEITEM.L_DISCOUNT <= Decimal128(Some(7),3,2) AND LINEITEM.L_QUANTITY < CAST(Int32(24) AS Decimal128(15, 2)) + Filter: LINEITEM.L_SHIPDATE >= CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-01-01") AS Date32) AND LINEITEM.L_DISCOUNT >= Decimal128(0.05,3,2) AND LINEITEM.L_DISCOUNT <= Decimal128(0.07,3,2) AND LINEITEM.L_QUANTITY < CAST(Int32(24) AS Decimal128(15, 2)) TableScan: LINEITEM "# ); @@ -273,7 +273,7 @@ mod tests { Sort: sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY) DESC NULLS FIRST Filter: sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY) > () Subquery: - Projection: sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY) * Decimal128(Some(1000000),11,10) + Projection: sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY) * Decimal128(0.0001000000,11,10) Aggregate: groupBy=[[]], aggr=[[sum(PARTSUPP.PS_SUPPLYCOST * PARTSUPP.PS_AVAILQTY)]] Projection: PARTSUPP.PS_SUPPLYCOST * CAST(PARTSUPP.PS_AVAILQTY AS Decimal128(19, 0)) Filter: PARTSUPP.PS_SUPPKEY = SUPPLIER.S_SUPPKEY AND SUPPLIER.S_NATIONKEY = NATION.N_NATIONKEY AND NATION.N_NAME = Utf8("JAPAN") @@ -340,9 +340,9 @@ mod tests { assert_snapshot!( plan_str, @r#" - Projection: Decimal128(Some(10000),5,2) * sum(CASE WHEN PART.P_TYPE LIKE Utf8("PROMO%") THEN LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT ELSE Decimal128(Some(0),19,4) END) / sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS PROMO_REVENUE - Aggregate: groupBy=[[]], aggr=[[sum(CASE WHEN PART.P_TYPE LIKE Utf8("PROMO%") THEN LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT ELSE Decimal128(Some(0),19,4) END), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT)]] - Projection: CASE WHEN PART.P_TYPE LIKE CAST(Utf8("PROMO%") AS Utf8) THEN LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) ELSE Decimal128(Some(0),19,4) END, LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) + Projection: Decimal128(100.00,5,2) * sum(CASE WHEN PART.P_TYPE LIKE Utf8("PROMO%") THEN LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT ELSE Decimal128(0.0000,19,4) END) / sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT) AS PROMO_REVENUE + Aggregate: groupBy=[[]], aggr=[[sum(CASE WHEN PART.P_TYPE LIKE Utf8("PROMO%") THEN LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT ELSE Decimal128(0.0000,19,4) END), sum(LINEITEM.L_EXTENDEDPRICE * Int32(1) - LINEITEM.L_DISCOUNT)]] + Projection: CASE WHEN PART.P_TYPE LIKE CAST(Utf8("PROMO%") AS Utf8) THEN LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) ELSE Decimal128(0.0000,19,4) END, LINEITEM.L_EXTENDEDPRICE * (CAST(Int32(1) AS Decimal128(15, 2)) - LINEITEM.L_DISCOUNT) Filter: LINEITEM.L_PARTKEY = PART.P_PARTKEY AND LINEITEM.L_SHIPDATE >= Date32("1995-09-01") AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-10-01") AS Date32) Cross Join: TableScan: LINEITEM @@ -389,12 +389,12 @@ mod tests { assert_snapshot!( plan_str, @r#" - Projection: sum(LINEITEM.L_EXTENDEDPRICE) / Decimal128(Some(70),2,1) AS AVG_YEARLY + Projection: sum(LINEITEM.L_EXTENDEDPRICE) / Decimal128(7.0,2,1) AS AVG_YEARLY Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_EXTENDEDPRICE)]] Projection: LINEITEM.L_EXTENDEDPRICE Filter: PART.P_PARTKEY = LINEITEM.L_PARTKEY AND PART.P_BRAND = Utf8("Brand#23") AND PART.P_CONTAINER = Utf8("MED BOX") AND LINEITEM.L_QUANTITY < () Subquery: - Projection: Decimal128(Some(2),2,1) * avg(LINEITEM.L_QUANTITY) + Projection: Decimal128(0.2,2,1) * avg(LINEITEM.L_QUANTITY) Aggregate: groupBy=[[]], aggr=[[avg(LINEITEM.L_QUANTITY)]] Projection: LINEITEM.L_QUANTITY Filter: LINEITEM.L_PARTKEY = outer_ref(PART.P_PARTKEY) @@ -468,7 +468,7 @@ mod tests { Filter: PART.P_NAME LIKE CAST(Utf8("forest%") AS Utf8) TableScan: PART Subquery: - Projection: Decimal128(Some(5),2,1) * sum(LINEITEM.L_QUANTITY) + Projection: Decimal128(0.5,2,1) * sum(LINEITEM.L_QUANTITY) Aggregate: groupBy=[[]], aggr=[[sum(LINEITEM.L_QUANTITY)]] Projection: LINEITEM.L_QUANTITY Filter: LINEITEM.L_PARTKEY = outer_ref(PARTSUPP.PS_PARTKEY) AND LINEITEM.L_SUPPKEY = outer_ref(PARTSUPP.PS_SUPPKEY) AND LINEITEM.L_SHIPDATE >= CAST(Utf8("1994-01-01") AS Date32) AND LINEITEM.L_SHIPDATE < CAST(Utf8("1995-01-01") AS Date32) @@ -526,7 +526,7 @@ mod tests { Subquery: Aggregate: groupBy=[[]], aggr=[[avg(CUSTOMER.C_ACCTBAL)]] Projection: CUSTOMER.C_ACCTBAL - Filter: CUSTOMER.C_ACCTBAL > Decimal128(Some(0),3,2) AND (substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("13") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("31") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("23") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("29") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("30") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("18") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("17") AS Utf8)) + Filter: CUSTOMER.C_ACCTBAL > Decimal128(0.00,3,2) AND (substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("13") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("31") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("23") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("29") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("30") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("18") AS Utf8) OR substr(CUSTOMER.C_PHONE, Int32(1), Int32(2)) = CAST(Utf8("17") AS Utf8)) TableScan: CUSTOMER Subquery: Filter: ORDERS.O_CUSTKEY = outer_ref(CUSTOMER.C_CUSTKEY) diff --git a/datafusion/substrait/tests/cases/logical_plans.rs b/datafusion/substrait/tests/cases/logical_plans.rs index 663a372fe2e4f..522381de6efdf 100644 --- a/datafusion/substrait/tests/cases/logical_plans.rs +++ b/datafusion/substrait/tests/cases/logical_plans.rs @@ -19,6 +19,7 @@ #[cfg(test)] mod tests { + use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; use crate::utils::test::{add_plan_schemas_to_ctx, read_json}; use datafusion::common::test_util::format_batches; use std::collections::HashSet; @@ -90,6 +91,31 @@ mod tests { Ok(()) } + #[tokio::test] + async fn nested_window_function_in_expression() -> Result<()> { + // The Substrait Project expression represents: + // SELECT 1 + count(*) OVER () FROM DATA + let proto_plan = read_json( + "tests/testdata/test_plans/nested_window_expression.substrait.json", + ); + let ctx = add_plan_schemas_to_ctx(SessionContext::new(), &proto_plan)?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + assert_snapshot!( + plan, + @r" + Projection: Int64(1) + count(Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING AS EXPR$0 + WindowAggr: windowExpr=[[count(Int64(1)) ROWS BETWEEN UNBOUNDED PRECEDING AND UNBOUNDED FOLLOWING]] + TableScan: DATA + " + ); + + // Trigger execution to ensure the nested window is physically plannable + DataFrame::new(ctx.state(), plan).show().await?; + + Ok(()) + } + #[tokio::test] async fn double_window_function() -> Result<()> { // Confirms a WindowExpr can be repeated in the same project. @@ -293,4 +319,25 @@ mod tests { Ok(()) } + + #[tokio::test] + async fn higher_order_function() -> Result<()> { + let proto_plan = + read_json("tests/testdata/test_plans/higher_order_function.json"); + // ctx already contains the queried table + let ctx = higher_order_function_ctx().await?; + let plan = from_substrait_plan(&ctx.state(), &proto_plan).await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_concat(array_transform2(p0, (p3, p4) -> p3 * p2 * p4), array_transform2(p0, (p5, p6) -> p5 * p2 * p6))) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_concat(array_transform2(v,(v, j) -> v * i * j),array_transform2(v,(v, j) -> v * i * j))) + TableScan: data3 + " + ); + + // Trigger execution to ensure plan validity + DataFrame::new(ctx.state(), plan).show().await?; + Ok(()) + } } diff --git a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs index 1d65256d76420..f084d3170edcc 100644 --- a/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs +++ b/datafusion/substrait/tests/cases/roundtrip_logical_plan.rs @@ -17,8 +17,13 @@ use crate::utils::test::read_json; use datafusion::arrow::array::ArrayRef; +use datafusion::config::Dialect; use datafusion::functions_nested::map::map; -use datafusion::logical_expr::LogicalPlanBuilder; +use datafusion::logical_expr::{ + ColumnarValue, HigherOrderFunctionArgs, HigherOrderReturnFieldArgs, + HigherOrderSignature, HigherOrderUDF, HigherOrderUDFImpl, LambdaParametersProgress, + LogicalPlanBuilder, ValueOrLambda, +}; use datafusion::physical_plan::Accumulator; use datafusion::scalar::ScalarValue; use datafusion_substrait::logical_plan::{ @@ -27,7 +32,9 @@ use datafusion_substrait::logical_plan::{ use std::cmp::Ordering; use std::mem::size_of_val; -use datafusion::arrow::datatypes::{DataType, Field, IntervalUnit, Schema, TimeUnit}; +use datafusion::arrow::datatypes::{ + DataType, Field, FieldRef, IntervalUnit, Schema, TimeUnit, +}; use datafusion::common::tree_node::Transformed; use datafusion::common::{DFSchema, DFSchemaRef, Spans, not_impl_err, plan_err}; use datafusion::error::Result; @@ -1116,6 +1123,27 @@ async fn aggregate_identical_grouping_expressions() -> Result<()> { Ok(()) } +#[tokio::test] +async fn aggregate_identical_measures() -> Result<()> { + // Two identical aggregate measures share the same schema_name; without + // NameTracker dedup over measures, building the Aggregate's output + // DFSchema fails with "Schema contains duplicate unqualified field name". + let proto_plan = read_json( + "tests/testdata/test_plans/aggregate_identical_measures.substrait.json", + ); + + let plan = generate_plan_from_substrait(proto_plan).await?; + assert_snapshot!( + plan, + @r" + Projection: __common_expr_1 AS sum_a_1, __common_expr_1 AS sum(data.a)__temp__0 AS sum_a_2 + Aggregate: groupBy=[[]], aggr=[[sum(data.a) AS __common_expr_1]] + TableScan: data projection=[a] + " + ); + Ok(()) +} + #[tokio::test] async fn simple_intersect_consume() -> Result<()> { let proto_plan = read_json("tests/testdata/test_plans/intersect.substrait.json"); @@ -1922,6 +1950,217 @@ async fn roundtrip_placeholder_typed_utf8() -> Result<()> { Ok(()) } +#[tokio::test] +async fn roundtrip_array_transform_higher_order_function() -> Result<()> { + let ctx = higher_order_function_ctx().await?; + + // simple + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], p0 -> p0 * 2) from data3", + ctx.clone(), + ) + .await?; + + // dont use the parameter + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], p0 -> 3) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters using both + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> p0 * p2) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters only last + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> 2 * p2) from data3", + ctx.clone(), + ) + .await?; + + // multiple parameters use none + roundtrip_with_ctx( + "SELECT array_transform2([data3.p1], (p0, p2) -> 3) from data3", + ctx.clone(), + ) + .await?; + + // nested without variable shadowing + roundtrip_with_ctx("SELECT array_transform2([[data3.p1]], p0 -> array_transform2(p0, p2 -> p2 * 2)) from data3", ctx.clone()) + .await?; + + // nested with multiple parameters without variable shadowing + roundtrip_with_ctx("SELECT array_transform2([[data3.p1]], (p0, p2) -> array_transform2(p0, (p3, p4) -> p2 * p3 * p4)) from data3", ctx.clone()) + .await?; + + // since substrait doesn't encode lambda parameters names, they got generated, non-conflicting names during consumption + // testing name shadowing requires to assert against the generated plan and check the correct parameter usage instead of round tripping + + // nested with variable shadowing. + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2([[data3.p1]], v -> array_transform2(v, v -> v * 2)) from data3", + true, + true, + &ctx, + ) + .await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0) -> array_transform2(p0, (p2) -> p2 * Int64(2))) AS array_transform2(make_array(make_array(data3.p1)),(v) -> array_transform2(v,(v) -> v * Int64(2))) + TableScan: data3 projection=[p1] + " + ); + + // nested with variable shadowing with multiple parameters + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2([[data3.p1]], (v, i) -> array_transform2(v, (v, i) -> v * i)) from data3", + true, + true, + &ctx, + ) + .await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0, p2) -> array_transform2(p0, (p3, p4) -> p3 * p4)) AS array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_transform2(v,(v, i) -> v * i)) + TableScan: data3 projection=[p1] + " + ); + + // nested with variable shadowing and later reuse of the shadowed var after exiting the shadowing expression + let plan = generate_plan_from_sql_with_ctx( + "SELECT array_transform2( + [[data3.p1]], + v -> array_concat( + -- when entering this expression, inner v is pushed into the producer and shadows outer v, but after exiting this, + -- it should be removed and unshadow the outer v, so that it can be used in the next expression + array_transform2(v, v -> v * 2), + array_transform2(v, v -> v * 2) + ) + ) from data3", + true, + true, + &ctx, + ) + .await?; + + assert_snapshot!( + plan, + @" + Projection: array_transform2(make_array(make_array(data3.p1)), (p0) -> array_concat(array_transform2(p0, (p2) -> p2 * Int64(2)), array_transform2(p0, (p3) -> p3 * Int64(2)))) AS array_transform2(make_array(make_array(data3.p1)),(v) -> array_concat(array_transform2(v,(v) -> v * Int64(2)),array_transform2(v,(v) -> v * Int64(2)))) + TableScan: data3 projection=[p1] + " + ); + + Ok(()) +} + +pub(crate) async fn higher_order_function_ctx() -> Result { + let ctx = create_context_with_dialect(Some(Dialect::Databricks)).await?; + + ctx.register_higher_order_function(Arc::new(HigherOrderUDF::new_from_impl( + ArrayTransform::new(), + ))); + + let data3_fields = vec![ + Field::new("p1", DataType::Int64, true), // lambda parameters should not conflict with this column + ]; + let data3 = Schema::new(data3_fields); + let mut data3_options = CsvReadOptions::new(); + data3_options.schema = Some(&data3); + data3_options.has_header = false; + ctx.register_csv("data3", "tests/testdata/empty.csv", data3_options) + .await?; + + Ok(ctx) +} + +// todo use core array_transform when it supports multiple lambda parameters +#[derive(Debug, PartialEq, Eq, Hash)] +struct ArrayTransform { + signature: HigherOrderSignature, +} + +impl ArrayTransform { + fn new() -> Self { + Self { + signature: HigherOrderSignature::variadic_any(Volatility::Immutable), + } + } +} + +impl HigherOrderUDFImpl for ArrayTransform { + fn name(&self) -> &str { + "array_transform2" + } + + fn aliases(&self) -> &[String] { + &[] + } + + fn signature(&self) -> &HigherOrderSignature { + &self.signature + } + + fn lambda_parameters( + &self, + _step: usize, + fields: &[ValueOrLambda>], + ) -> Result { + let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(_)] = fields else { + unreachable!() + }; + + let field = match list.data_type() { + DataType::List(field) => field, + _ => unreachable!(), + }; + + Ok(LambdaParametersProgress::Complete(vec![vec![ + Arc::clone(field), + Arc::new(Field::new("", DataType::Int64, true)), + ]])) + } + + fn return_field_from_args( + &self, + args: HigherOrderReturnFieldArgs, + ) -> Result { + let [ValueOrLambda::Value(list), ValueOrLambda::Lambda(lambda)] = args.arg_fields + else { + unreachable!() + }; + + let field = Arc::new(Field::new( + Field::LIST_FIELD_DEFAULT_NAME, + lambda.data_type().clone(), + lambda.is_nullable(), + )); + + let return_type = match list.data_type() { + DataType::List(_) => DataType::List(field), + _ => unreachable!(), + }; + + Ok(Arc::new(Field::new("", return_type, list.is_nullable()))) + } + + fn invoke_with_args(&self, args: HigherOrderFunctionArgs) -> Result { + // this function is only tested with roundtrip_with_ctx, which only prints the output + // and generate_plan_from_sql_with_ctx which doesn't execute nothing, so the output doesn't matter + Ok(ColumnarValue::Scalar(ScalarValue::new_default( + args.return_type(), + )?)) + } +} + fn check_post_join_filters(rel: &Rel) -> Result<()> { // search for target_rel and field value in proto match &rel.rel_type { @@ -1985,7 +2224,7 @@ fn check_post_join_filters(rel: &Rel) -> Result<()> { } } -async fn verify_post_join_filter_value(proto: Box) -> Result<()> { +fn verify_post_join_filter_value(proto: &Plan) -> Result<()> { for relation in &proto.relations { match relation.rel_type.as_ref() { Some(rt) => match rt { @@ -2024,10 +2263,7 @@ fn count_read_filters(rel: &Rel, filter_count: &mut u32) -> Result<()> { } } -async fn assert_read_filter_count( - proto: Box, - expected_filter_count: u32, -) -> Result<()> { +fn assert_read_filter_count(proto: &Plan, expected_filter_count: u32) -> Result<()> { let mut filter_count: u32 = 0; for relation in &proto.relations { match relation.rel_type.as_ref() { @@ -2063,6 +2299,15 @@ async fn generate_plan_from_sql( optimized: bool, ) -> Result { let ctx = create_context().await?; + generate_plan_from_sql_with_ctx(sql, assert_schema, optimized, &ctx).await +} + +async fn generate_plan_from_sql_with_ctx( + sql: &str, + assert_schema: bool, + optimized: bool, + ctx: &SessionContext, +) -> Result { let df: DataFrame = ctx.sql(sql).await?; let plan = if optimized { @@ -2396,7 +2641,7 @@ async fn roundtrip_verify_post_join_filter(sql: &str) -> Result<()> { let proto = roundtrip_with_ctx(sql, ctx).await?; // verify that the join filters are None - verify_post_join_filter_value(proto).await + verify_post_join_filter_value(&proto) } async fn roundtrip_verify_read_filter_count( @@ -2407,7 +2652,7 @@ async fn roundtrip_verify_read_filter_count( let proto = roundtrip_with_ctx(sql, ctx).await?; // verify that filter counts in read relations are as expected - assert_read_filter_count(proto, expected_filter_count).await + assert_read_filter_count(&proto, expected_filter_count) } async fn roundtrip_all_types(sql: &str) -> Result<()> { @@ -2416,8 +2661,18 @@ async fn roundtrip_all_types(sql: &str) -> Result<()> { } async fn create_context() -> Result { + create_context_with_dialect(None).await +} + +async fn create_context_with_dialect(dialect: Option) -> Result { + let mut session_config = SessionConfig::default(); + + if let Some(dialect) = dialect { + session_config.options_mut().sql_parser.dialect = dialect; + } + let mut state = SessionStateBuilder::new() - .with_config(SessionConfig::default()) + .with_config(session_config) .with_runtime_env(Arc::new(RuntimeEnv::default())) .with_default_features() .with_serializer_registry(Arc::new(MockSerializerRegistry)) diff --git a/datafusion/substrait/tests/cases/serialize.rs b/datafusion/substrait/tests/cases/serialize.rs index 2d7257fad3394..1981ef66db377 100644 --- a/datafusion/substrait/tests/cases/serialize.rs +++ b/datafusion/substrait/tests/cases/serialize.rs @@ -28,9 +28,16 @@ mod tests { use insta::assert_snapshot; use std::fs; + use substrait::proto::expression::field_reference::{ReferenceType, RootType}; + use substrait::proto::expression::reference_segment; + use substrait::proto::expression::{ReferenceSegment, RexType}; + use substrait::proto::function_argument::ArgType; use substrait::proto::plan_rel::RelType; use substrait::proto::rel_common::{Emit, EmitKind}; - use substrait::proto::{RelCommon, rel}; + use substrait::proto::r#type::{I64, Kind as TypeKind, List, Nullability, Struct}; + use substrait::proto::{Expression, RelCommon, Type, rel}; + + use crate::cases::roundtrip_logical_plan::higher_order_function_ctx; #[tokio::test] async fn serialize_to_file() -> Result<()> { @@ -196,6 +203,101 @@ mod tests { panic!("plan did not match expected structure") } + #[tokio::test] + async fn higher_order_function() -> Result<()> { + let ctx = higher_order_function_ctx().await?; + let df = ctx + .sql( + "SELECT array_transform2( + [[data3.p1]], + (v, i) -> array_concat( + -- when entering this expression, inner v is pushed into the producer and shadows outer v, but after exiting this, + -- it should be removed and unshadow the outer v, so that it can be used in the next expression + array_transform2(v, (v, j) -> v * i * j), + array_transform2(v, (v, j) -> v * i * j) + ) + ) from data3" + ) + .await?; + let datafusion_plan = df.into_optimized_plan()?; + let plan = to_substrait_plan(&datafusion_plan, &ctx.state())? + .as_ref() + .clone(); + + let relation = plan.relations.first().unwrap().rel_type.as_ref(); + let root_rel = match relation { + Some(RelType::Root(root)) => root.input.as_ref().unwrap(), + _ => panic!("expected Root"), + }; + + let Some(rel::RelType::Project(p)) = root_rel.rel_type.as_ref() else { + panic!("expected Project at top of plan") + }; + + let mut params = vec![]; + let mut lambda_refs = vec![]; + + collect_lambda_ref(&p.expressions[0], &mut params, &mut lambda_refs); + + let nullable_i64 = Type { + kind: Some(TypeKind::I64(I64 { + type_variation_reference: 0, + nullability: Nullability::Nullable as i32, + })), + }; + + let inner_lambda_struct = Struct { + // v, j + types: vec![nullable_i64.clone(); 2], + type_variation_reference: 0, + nullability: Nullability::Required as i32, + }; + + assert_eq!( + params, + vec![ + Struct { + types: vec![ + // v + Type { + kind: Some(TypeKind::List(Box::new(List { + r#type: Some(Box::new(nullable_i64.clone())), + type_variation_reference: 0, + nullability: Nullability::Nullable as i32 + }))) + }, + // i + nullable_i64, + ], + type_variation_reference: 0, + nullability: Nullability::Required as i32, + }, + inner_lambda_struct.clone(), + inner_lambda_struct, + ] + ); + + assert_eq!( + lambda_refs, + vec![ + // first inner array_transform2 argument: outer v + (0, 0), + // first inner lambda body: v * i * j + (0, 0), + (1, 1), + (0, 1), + // second inner array_transform2 argument: outer v + (0, 0), + // second inner lambda body: v * i * j + (0, 0), + (1, 1), + (0, 1), + ] + ); + + Ok(()) + } + fn assert_emit(rel_common: Option<&RelCommon>, output_mapping: Vec) { assert_eq!( rel_common.unwrap().emit_kind.clone(), @@ -211,4 +313,54 @@ mod tests { .await?; Ok(ctx) } + + // Recursively walks a expression tree depth-first, collecting in visit order: + // - `params`: the parameter struct of each Lambda encountered + // - `lambda_refs`: every field reference whose root is a LambdaParameterReference, + // recorded as (steps_out, field_index) so tests can assert which enclosing + // lambda each reference resolves to and which parameter within it. + fn collect_lambda_ref( + expr: &Expression, + params: &mut Vec, + lambda_refs: &mut Vec<(u32, i32)>, + ) { + if let Some(rex_type) = &expr.rex_type { + match rex_type { + RexType::Selection(field_reference) => { + if let ( + Some(ReferenceType::DirectReference(ReferenceSegment { + reference_type: + Some(reference_segment::ReferenceType::StructField( + struct_field, + )), + })), + Some(RootType::LambdaParameterReference(lambda_param_ref)), + ) = (&field_reference.reference_type, &field_reference.root_type) + { + lambda_refs.push((lambda_param_ref.steps_out, struct_field.field)) + } + } + RexType::ScalarFunction(scalar_function) => { + for arg in &scalar_function.arguments { + match &arg.arg_type { + Some(ArgType::Value(value)) => { + collect_lambda_ref(value, params, lambda_refs) + } + _ => unreachable!(), + } + } + } + RexType::Lambda(lambda) => { + if let Some(parameters) = &lambda.parameters { + params.push(parameters.clone()); + } + if let Some(body) = &lambda.body { + collect_lambda_ref(body, params, lambda_refs); + } + } + RexType::Literal(_literal) => {} + _ => unreachable!(), + } + } + } } diff --git a/datafusion/substrait/tests/testdata/test_plans/aggregate_identical_measures.substrait.json b/datafusion/substrait/tests/testdata/test_plans/aggregate_identical_measures.substrait.json new file mode 100644 index 0000000000000..620d55e93ee1e --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/aggregate_identical_measures.substrait.json @@ -0,0 +1,103 @@ +{ + "extensionUris": [{ + "extensionUriAnchor": 1, + "uri": "/functions_arithmetic.yaml" + }], + "extensions": [{ + "extensionFunction": { + "extensionUriReference": 1, + "functionAnchor": 0, + "name": "sum:i64" + } + }], + "relations": [{ + "root": { + "input": { + "aggregate": { + "common": { + "direct": {} + }, + "input": { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": ["a"], + "struct": { + "types": [{ + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": ["data"] + } + } + }, + "groupings": [{ + "groupingExpressions": [] + }], + "measures": [ + { + "measure": { + "functionReference": 0, + "phase": "AGGREGATION_PHASE_INITIAL_TO_RESULT", + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "invocation": "AGGREGATION_INVOCATION_ALL", + "arguments": [{ + "value": { + "selection": { + "directReference": { + "structField": { + "field": 0 + } + }, + "rootReference": {} + } + } + }] + } + }, + { + "measure": { + "functionReference": 0, + "phase": "AGGREGATION_PHASE_INITIAL_TO_RESULT", + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "invocation": "AGGREGATION_INVOCATION_ALL", + "arguments": [{ + "value": { + "selection": { + "directReference": { + "structField": { + "field": 0 + } + }, + "rootReference": {} + } + } + }] + } + } + ] + } + }, + "names": ["sum_a_1", "sum_a_2"] + } + }], + "version": { + "minorNumber": 54, + "producer": "manual" + } +} diff --git a/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json b/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json new file mode 100644 index 0000000000000..da613b2573447 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/higher_order_function.json @@ -0,0 +1,438 @@ +{ + "version": { + "minorNumber": 85, + "producer": "datafusion" + }, + "extensions": [ + { + "extensionFunction": { + "extensionUrnReference": 2, + "functionAnchor": 2, + "name": "array_transform2" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 2, + "name": "make_array" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 2, + "functionAnchor": 3, + "name": "array_concat" + } + }, + { + "extensionFunction": { + "extensionUrnReference": 1, + "functionAnchor": 1, + "name": "multiply" + } + } + ], + "relations": [ + { + "root": { + "input": { + "project": { + "common": { + "emit": { + "outputMapping": [ + 1 + ] + } + }, + "input": { + "read": { + "baseSchema": { + "names": [ + "p1" + ], + "struct": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + } + }, + "projection": { + "select": { + "structItems": [ + {} + ] + } + }, + "namedTable": { + "names": [ + "data3" + ] + } + } + }, + "expressions": [ + { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "scalarFunction": { + "arguments": [ + { + "value": { + "scalarFunction": { + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "rootReference": {} + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 3, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": { + "stepsOut": 1 + } + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": {} + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "scalarFunction": { + "functionReference": 2, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "lambda": { + "parameters": { + "types": [ + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "nullability": "NULLABILITY_REQUIRED" + }, + "body": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "scalarFunction": { + "functionReference": 1, + "arguments": [ + { + "value": { + "selection": { + "directReference": { + "structField": {} + }, + "lambdaParameterReference": {} + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": { + "stepsOut": 1 + } + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + }, + { + "value": { + "selection": { + "directReference": { + "structField": { + "field": 1 + } + }, + "lambdaParameterReference": {} + } + } + } + ], + "outputType": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + } + } + } + ], + "outputType": { + "list": { + "type": { + "list": { + "type": { + "i64": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullability": "NULLABILITY_NULLABLE" + } + } + } + } + ] + } + }, + "names": [ + "array_transform2(make_array(make_array(data3.p1)),(v, i) -> array_concat(array_transform2(v,(v, j) -> v * i * j),array_transform2(v,(v, j) -> v * i * j)))" + ] + } + } + ], + "extensionUrns": [ + { + "extensionUrnAnchor": 1, + "urn": "extension:io.substrait:functions_arithmetic" + }, + { + "extensionUrnAnchor": 2, + "urn": "extension:io.substrait:functions_list" + } + ] +} diff --git a/datafusion/substrait/tests/testdata/test_plans/join_with_expression_key.json b/datafusion/substrait/tests/testdata/test_plans/join_with_expression_key.json index 73fa06eea5f05..8a81a9a0c780f 100644 --- a/datafusion/substrait/tests/testdata/test_plans/join_with_expression_key.json +++ b/datafusion/substrait/tests/testdata/test_plans/join_with_expression_key.json @@ -100,29 +100,52 @@ } }, "virtualTable": { - "values": [{ - "fields": [{ - "string": "aaa", - "nullable": true - }, { - "string": "host-a", - "nullable": true - }, { - "i64": "128", - "nullable": true - }] - }, { - "fields": [{ - "string": "bbb", - "nullable": true - }, { - "string": "host-b", - "nullable": true - }, { - "i64": "256", - "nullable": true - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "string": "aaa", + "nullable": true + } + }, + { + "literal": { + "string": "host-a", + "nullable": true + } + }, + { + "literal": { + "i64": "128", + "nullable": true + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "bbb", + "nullable": true + } + }, + { + "literal": { + "string": "host-b", + "nullable": true + } + }, + { + "literal": { + "i64": "256", + "nullable": true + } + } + ] + } + ] } } }, @@ -293,23 +316,40 @@ } }, "virtualTable": { - "values": [{ - "fields": [{ - "string": "host-a", - "nullable": true - }, { - "i64": "107", - "nullable": true - }] - }, { - "fields": [{ - "string": "host-b", - "nullable": true - }, { - "i64": "214", - "nullable": true - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "string": "host-a", + "nullable": true + } + }, + { + "literal": { + "i64": "107", + "nullable": true + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "host-b", + "nullable": true + } + }, + { + "literal": { + "i64": "214", + "nullable": true + } + } + ] + } + ] } } }, @@ -365,29 +405,52 @@ } }, "virtualTable": { - "values": [{ - "fields": [{ - "string": "aaa", - "nullable": true - }, { - "string": "host-a", - "nullable": true - }, { - "i64": "128", - "nullable": true - }] - }, { - "fields": [{ - "string": "bbb", - "nullable": true - }, { - "string": "host-b", - "nullable": true - }, { - "i64": "256", - "nullable": true - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "string": "aaa", + "nullable": true + } + }, + { + "literal": { + "string": "host-a", + "nullable": true + } + }, + { + "literal": { + "i64": "128", + "nullable": true + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "bbb", + "nullable": true + } + }, + { + "literal": { + "string": "host-b", + "nullable": true + } + }, + { + "literal": { + "i64": "256", + "nullable": true + } + } + ] + } + ] } } }, diff --git a/datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct.json b/datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct.json index 642256c562995..13c1e5899db0b 100644 --- a/datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct.json +++ b/datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct.json @@ -24,13 +24,147 @@ } }, "virtualTable": { - "values": [ - { "fields": [{ "string": "1", "nullable": false }, { "string": "a", "nullable": true }, { "string": "c1", "nullable": false }] }, - { "fields": [{ "string": "2", "nullable": false }, { "string": "b", "nullable": true }, { "string": "c2", "nullable": false }] }, - { "fields": [{ "string": "3", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c3", "nullable": false }] }, - { "fields": [{ "string": "4", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c4", "nullable": false }] }, - { "fields": [{ "string": "5", "nullable": false }, { "string": "e", "nullable": true }, { "string": "c5", "nullable": false }] }, - { "fields": [{ "string": "6", "nullable": false }, { "string": "f", "nullable": true }, { "string": "c6", "nullable": false }] } + "expressions": [ + { + "fields": [ + { + "literal": { + "string": "1", + "nullable": false + } + }, + { + "literal": { + "string": "a", + "nullable": true + } + }, + { + "literal": { + "string": "c1", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "2", + "nullable": false + } + }, + { + "literal": { + "string": "b", + "nullable": true + } + }, + { + "literal": { + "string": "c2", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "3", + "nullable": false + } + }, + { + "literal": { + "null": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullable": true + } + }, + { + "literal": { + "string": "c3", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "4", + "nullable": false + } + }, + { + "literal": { + "null": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullable": true + } + }, + { + "literal": { + "string": "c4", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "5", + "nullable": false + } + }, + { + "literal": { + "string": "e", + "nullable": true + } + }, + { + "literal": { + "string": "c5", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "6", + "nullable": false + } + }, + { + "literal": { + "string": "f", + "nullable": true + } + }, + { + "literal": { + "string": "c6", + "nullable": false + } + } + ] + } ] } } @@ -50,13 +184,147 @@ } }, "virtualTable": { - "values": [ - { "fields": [{ "string": "1", "nullable": false }, { "string": "a", "nullable": true }, { "string": "c1", "nullable": false }] }, - { "fields": [{ "string": "2", "nullable": false }, { "string": "b", "nullable": true }, { "string": "c2", "nullable": false }] }, - { "fields": [{ "string": "3", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c3", "nullable": false }] }, - { "fields": [{ "string": "4", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c4", "nullable": false }] }, - { "fields": [{ "string": "5", "nullable": false }, { "string": "e", "nullable": true }, { "string": "c5", "nullable": false }] }, - { "fields": [{ "string": "6", "nullable": false }, { "string": "f", "nullable": true }, { "string": "c6", "nullable": false }] } + "expressions": [ + { + "fields": [ + { + "literal": { + "string": "1", + "nullable": false + } + }, + { + "literal": { + "string": "a", + "nullable": true + } + }, + { + "literal": { + "string": "c1", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "2", + "nullable": false + } + }, + { + "literal": { + "string": "b", + "nullable": true + } + }, + { + "literal": { + "string": "c2", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "3", + "nullable": false + } + }, + { + "literal": { + "null": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullable": true + } + }, + { + "literal": { + "string": "c3", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "4", + "nullable": false + } + }, + { + "literal": { + "null": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullable": true + } + }, + { + "literal": { + "string": "c4", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "5", + "nullable": false + } + }, + { + "literal": { + "string": "e", + "nullable": true + } + }, + { + "literal": { + "string": "c5", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "6", + "nullable": false + } + }, + { + "literal": { + "string": "f", + "nullable": true + } + }, + { + "literal": { + "string": "c6", + "nullable": false + } + } + ] + } ] } } diff --git a/datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct_left.json b/datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct_left.json index f16672947e1ee..481bba44d839b 100644 --- a/datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct_left.json +++ b/datafusion/substrait/tests/testdata/test_plans/mixed_join_equal_and_indistinct_left.json @@ -24,13 +24,147 @@ } }, "virtualTable": { - "values": [ - { "fields": [{ "string": "1", "nullable": false }, { "string": "a", "nullable": true }, { "string": "c1", "nullable": false }] }, - { "fields": [{ "string": "2", "nullable": false }, { "string": "b", "nullable": true }, { "string": "c2", "nullable": false }] }, - { "fields": [{ "string": "3", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c3", "nullable": false }] }, - { "fields": [{ "string": "4", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c4", "nullable": false }] }, - { "fields": [{ "string": "5", "nullable": false }, { "string": "e", "nullable": true }, { "string": "c5", "nullable": false }] }, - { "fields": [{ "string": "6", "nullable": false }, { "string": "f", "nullable": true }, { "string": "c6", "nullable": false }] } + "expressions": [ + { + "fields": [ + { + "literal": { + "string": "1", + "nullable": false + } + }, + { + "literal": { + "string": "a", + "nullable": true + } + }, + { + "literal": { + "string": "c1", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "2", + "nullable": false + } + }, + { + "literal": { + "string": "b", + "nullable": true + } + }, + { + "literal": { + "string": "c2", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "3", + "nullable": false + } + }, + { + "literal": { + "null": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullable": true + } + }, + { + "literal": { + "string": "c3", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "4", + "nullable": false + } + }, + { + "literal": { + "null": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullable": true + } + }, + { + "literal": { + "string": "c4", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "5", + "nullable": false + } + }, + { + "literal": { + "string": "e", + "nullable": true + } + }, + { + "literal": { + "string": "c5", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "6", + "nullable": false + } + }, + { + "literal": { + "string": "f", + "nullable": true + } + }, + { + "literal": { + "string": "c6", + "nullable": false + } + } + ] + } ] } } @@ -50,13 +184,147 @@ } }, "virtualTable": { - "values": [ - { "fields": [{ "string": "1", "nullable": false }, { "string": "a", "nullable": true }, { "string": "c1", "nullable": false }] }, - { "fields": [{ "string": "2", "nullable": false }, { "string": "b", "nullable": true }, { "string": "c2", "nullable": false }] }, - { "fields": [{ "string": "3", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c3", "nullable": false }] }, - { "fields": [{ "string": "4", "nullable": false }, { "null": { "string": { "nullability": "NULLABILITY_NULLABLE" } }, "nullable": true }, { "string": "c4", "nullable": false }] }, - { "fields": [{ "string": "5", "nullable": false }, { "string": "e", "nullable": true }, { "string": "c5", "nullable": false }] }, - { "fields": [{ "string": "6", "nullable": false }, { "string": "f", "nullable": true }, { "string": "c6", "nullable": false }] } + "expressions": [ + { + "fields": [ + { + "literal": { + "string": "1", + "nullable": false + } + }, + { + "literal": { + "string": "a", + "nullable": true + } + }, + { + "literal": { + "string": "c1", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "2", + "nullable": false + } + }, + { + "literal": { + "string": "b", + "nullable": true + } + }, + { + "literal": { + "string": "c2", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "3", + "nullable": false + } + }, + { + "literal": { + "null": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullable": true + } + }, + { + "literal": { + "string": "c3", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "4", + "nullable": false + } + }, + { + "literal": { + "null": { + "string": { + "nullability": "NULLABILITY_NULLABLE" + } + }, + "nullable": true + } + }, + { + "literal": { + "string": "c4", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "5", + "nullable": false + } + }, + { + "literal": { + "string": "e", + "nullable": true + } + }, + { + "literal": { + "string": "c5", + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "string": "6", + "nullable": false + } + }, + { + "literal": { + "string": "f", + "nullable": true + } + }, + { + "literal": { + "string": "c6", + "nullable": false + } + } + ] + } ] } } diff --git a/datafusion/substrait/tests/testdata/test_plans/multiple_joins.json b/datafusion/substrait/tests/testdata/test_plans/multiple_joins.json index e88cce648da7c..15c0313b43b54 100644 --- a/datafusion/substrait/tests/testdata/test_plans/multiple_joins.json +++ b/datafusion/substrait/tests/testdata/test_plans/multiple_joins.json @@ -72,19 +72,30 @@ } }, "virtualTable": { - "values": [{ - "fields": [{ - "i64": "1", - "nullable": true, - "typeVariationReference": 0 - }] - }, { - "fields": [{ - "i64": "2", - "nullable": true, - "typeVariationReference": 0 - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "i64": "1", + "nullable": true, + "typeVariationReference": 0 + } + } + ] + }, + { + "fields": [ + { + "literal": { + "i64": "2", + "nullable": true, + "typeVariationReference": 0 + } + } + ] + } + ] } } }, @@ -153,27 +164,44 @@ } }, "virtualTable": { - "values": [{ - "fields": [{ - "i64": "1", - "nullable": true, - "typeVariationReference": 0 - }, { - "string": "info", - "nullable": true, - "typeVariationReference": 0 - }] - }, { - "fields": [{ - "i64": "2", - "nullable": true, - "typeVariationReference": 0 - }, { - "string": "low", - "nullable": true, - "typeVariationReference": 0 - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "i64": "1", + "nullable": true, + "typeVariationReference": 0 + } + }, + { + "literal": { + "string": "info", + "nullable": true, + "typeVariationReference": 0 + } + } + ] + }, + { + "fields": [ + { + "literal": { + "i64": "2", + "nullable": true, + "typeVariationReference": 0 + } + }, + { + "literal": { + "string": "low", + "nullable": true, + "typeVariationReference": 0 + } + } + ] + } + ] } } }, @@ -272,19 +300,30 @@ } }, "virtualTable": { - "values": [{ - "fields": [{ - "i64": "1", - "nullable": true, - "typeVariationReference": 0 - }] - }, { - "fields": [{ - "i64": "2", - "nullable": true, - "typeVariationReference": 0 - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "i64": "1", + "nullable": true, + "typeVariationReference": 0 + } + } + ] + }, + { + "fields": [ + { + "literal": { + "i64": "2", + "nullable": true, + "typeVariationReference": 0 + } + } + ] + } + ] } } }, @@ -389,19 +428,30 @@ } }, "virtualTable": { - "values": [{ - "fields": [{ - "i64": "1", - "nullable": true, - "typeVariationReference": 0 - }] - }, { - "fields": [{ - "i64": "2", - "nullable": true, - "typeVariationReference": 0 - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "i64": "1", + "nullable": true, + "typeVariationReference": 0 + } + } + ] + }, + { + "fields": [ + { + "literal": { + "i64": "2", + "nullable": true, + "typeVariationReference": 0 + } + } + ] + } + ] } } }, diff --git a/datafusion/substrait/tests/testdata/test_plans/nested_window_expression.substrait.json b/datafusion/substrait/tests/testdata/test_plans/nested_window_expression.substrait.json new file mode 100644 index 0000000000000..f4dc73a9ca672 --- /dev/null +++ b/datafusion/substrait/tests/testdata/test_plans/nested_window_expression.substrait.json @@ -0,0 +1,131 @@ +{ + "extensionUris": [ + { + "extensionUriAnchor": 1, + "uri": "/functions_arithmetic.yaml" + }, + { + "extensionUriAnchor": 2, + "uri": "/functions_aggregate_generic.yaml" + } + ], + "extensions": [ + { + "extensionFunction": { + "extensionUriReference": 1, + "functionAnchor": 0, + "name": "add:i64_i64" + } + }, + { + "extensionFunction": { + "extensionUriReference": 2, + "functionAnchor": 1, + "name": "count:any" + } + } + ], + "relations": [ + { + "root": { + "input": { + "project": { + "common": { + "emit": { + "outputMapping": [ + 1 + ] + } + }, + "input": { + "read": { + "common": { + "direct": {} + }, + "baseSchema": { + "names": [ + "A" + ], + "struct": { + "types": [ + { + "i64": { + "typeVariationReference": 0, + "nullability": "NULLABILITY_NULLABLE" + } + } + ], + "typeVariationReference": 0, + "nullability": "NULLABILITY_REQUIRED" + } + }, + "namedTable": { + "names": [ + "DATA" + ] + } + } + }, + "expressions": [ + { + "scalarFunction": { + "functionReference": 0, + "args": [], + "outputType": { + "i64": { + "typeVariationReference": 0, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "arguments": [ + { + "value": { + "literal": { + "i64": 1, + "nullable": false, + "typeVariationReference": 0 + } + } + }, + { + "value": { + "windowFunction": { + "functionReference": 1, + "partitions": [], + "sorts": [], + "upperBound": { + "unbounded": {} + }, + "lowerBound": { + "unbounded": {} + }, + "phase": "AGGREGATION_PHASE_INITIAL_TO_RESULT", + "outputType": { + "i64": { + "typeVariationReference": 0, + "nullability": "NULLABILITY_NULLABLE" + } + }, + "args": [], + "arguments": [], + "invocation": "AGGREGATION_INVOCATION_ALL", + "options": [], + "boundsType": "BOUNDS_TYPE_ROWS" + } + } + } + ], + "options": [] + } + } + ] + } + }, + "names": [ + "EXPR$0" + ] + } + } + ], + "expectedTypeUrls": [] +} diff --git a/datafusion/substrait/tests/testdata/test_plans/non_nullable_lists.substrait.json b/datafusion/substrait/tests/testdata/test_plans/non_nullable_lists.substrait.json index e1c5574f8bec2..e29c000ee669b 100644 --- a/datafusion/substrait/tests/testdata/test_plans/non_nullable_lists.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/non_nullable_lists.substrait.json @@ -34,26 +34,28 @@ } }, "virtualTable": { - "values": [ + "expressions": [ { "fields": [ { - "list": { - "values": [ - { - "i32": 1, - "nullable": false, - "typeVariationReference": 0 - }, - { - "i32": 2, - "nullable": false, - "typeVariationReference": 0 - } - ] - }, - "nullable": false, - "typeVariationReference": 0 + "literal": { + "list": { + "values": [ + { + "i32": 1, + "nullable": false, + "typeVariationReference": 0 + }, + { + "i32": 2, + "nullable": false, + "typeVariationReference": 0 + } + ] + }, + "nullable": false, + "typeVariationReference": 0 + } } ] } diff --git a/datafusion/substrait/tests/testdata/test_plans/scalar_fn_logb_expr.substrait.json b/datafusion/substrait/tests/testdata/test_plans/scalar_fn_logb_expr.substrait.json index eeaf5a3dd8476..d5209a683d633 100644 --- a/datafusion/substrait/tests/testdata/test_plans/scalar_fn_logb_expr.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/scalar_fn_logb_expr.substrait.json @@ -85,23 +85,40 @@ "direct": {} }, "virtualTable": { - "values": [{ - "fields": [{ - "fp32": 1.0, - "nullable": false - }, { - "fp32": 10.0, - "nullable": false - }] - }, { - "fields": [{ - "fp32": 100.0, - "nullable": false - }, { - "fp32": 10.0, - "nullable": false - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "fp32": 1.0, + "nullable": false + } + }, + { + "literal": { + "fp32": 10.0, + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "fp32": 100.0, + "nullable": false + } + }, + { + "literal": { + "fp32": 10.0, + "nullable": false + } + } + ] + } + ] } } } diff --git a/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_between_expr.substrait.json b/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_between_expr.substrait.json index 6749a301b17df..f609d26138ad8 100644 --- a/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_between_expr.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_between_expr.substrait.json @@ -106,29 +106,52 @@ "direct": {} }, "virtualTable": { - "values": [{ - "fields": [{ - "i8": 2, - "nullable": false - }, { - "i8": 1, - "nullable": false - }, { - "i8": 3, - "nullable": false - }] - }, { - "fields": [{ - "i8": 4, - "nullable": false - }, { - "i8": 1, - "nullable": false - }, { - "i8": 2, - "nullable": false - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "i8": 2, + "nullable": false + } + }, + { + "literal": { + "i8": 1, + "nullable": false + } + }, + { + "literal": { + "i8": 3, + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "i8": 4, + "nullable": false + } + }, + { + "literal": { + "i8": 1, + "nullable": false + } + }, + { + "literal": { + "i8": 2, + "nullable": false + } + } + ] + } + ] } } } diff --git a/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_built_in_binary_expr_and_not.substrait.json b/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_built_in_binary_expr_and_not.substrait.json index 8365b1edfe250..5d91342257825 100644 --- a/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_built_in_binary_expr_and_not.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_built_in_binary_expr_and_not.substrait.json @@ -85,39 +85,72 @@ "direct": {} }, "virtualTable": { - "values": [{ - "fields": [{ - "boolean": true, - "nullable": false - }, { - "boolean": true, - "nullable": false - }] - }, { - "fields": [{ - "boolean": true, - "nullable": false - }, { - "boolean": false, - "nullable": false - }] - }, { - "fields": [{ - "boolean": false, - "nullable": false - }, { - "boolean": true, - "nullable": false - }] - }, { - "fields": [{ - "boolean": false, - "nullable": false - }, { - "boolean": false, - "nullable": false - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "boolean": true, + "nullable": false + } + }, + { + "literal": { + "boolean": true, + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "boolean": true, + "nullable": false + } + }, + { + "literal": { + "boolean": false, + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "boolean": false, + "nullable": false + } + }, + { + "literal": { + "boolean": true, + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "boolean": false, + "nullable": false + } + }, + { + "literal": { + "boolean": false, + "nullable": false + } + } + ] + } + ] } } } diff --git a/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_built_in_binary_expr_xor.substrait.json b/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_built_in_binary_expr_xor.substrait.json index cfd760de890c0..2514c0afc9448 100644 --- a/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_built_in_binary_expr_xor.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/scalar_fn_to_built_in_binary_expr_xor.substrait.json @@ -85,39 +85,72 @@ "direct": {} }, "virtualTable": { - "values": [{ - "fields": [{ - "boolean": true, - "nullable": false - }, { - "boolean": true, - "nullable": false - }] - }, { - "fields": [{ - "boolean": true, - "nullable": false - }, { - "boolean": false, - "nullable": false - }] - }, { - "fields": [{ - "boolean": false, - "nullable": false - }, { - "boolean": true, - "nullable": false - }] - }, { - "fields": [{ - "boolean": false, - "nullable": false - }, { - "boolean": false, - "nullable": false - }] - }] + "expressions": [ + { + "fields": [ + { + "literal": { + "boolean": true, + "nullable": false + } + }, + { + "literal": { + "boolean": true, + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "boolean": true, + "nullable": false + } + }, + { + "literal": { + "boolean": false, + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "boolean": false, + "nullable": false + } + }, + { + "literal": { + "boolean": true, + "nullable": false + } + } + ] + }, + { + "fields": [ + { + "literal": { + "boolean": false, + "nullable": false + } + }, + { + "literal": { + "boolean": false, + "nullable": false + } + } + ] + } + ] } } } diff --git a/datafusion/substrait/tests/testdata/test_plans/select_count_from_select_1.substrait.json b/datafusion/substrait/tests/testdata/test_plans/select_count_from_select_1.substrait.json index e9f6795880185..b0d4ba4813bcf 100644 --- a/datafusion/substrait/tests/testdata/test_plans/select_count_from_select_1.substrait.json +++ b/datafusion/substrait/tests/testdata/test_plans/select_count_from_select_1.substrait.json @@ -43,12 +43,14 @@ } }, "virtualTable": { - "values": [ + "expressions": [ { "fields": [ { - "i64": "0", - "nullable": false + "literal": { + "i64": "0", + "nullable": false + } } ] } diff --git a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json index f107ac473a987..34b0d22f00c4c 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package-lock.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package-lock.json @@ -15,7 +15,7 @@ "copy-webpack-plugin": "14.0.0", "webpack": "5.105.0", "webpack-cli": "5.1.4", - "webpack-dev-server": "5.2.4" + "webpack-dev-server": "6.0.0" } }, "../pkg": { @@ -391,21 +391,20 @@ "dev": true }, "node_modules/@types/express": { - "version": "4.17.25", - "resolved": "https://registry.npmjs.org/@types/express/-/express-4.17.25.tgz", - "integrity": "sha512-dVd04UKsfpINUnK0yBoYHDF3xu7xVH4BuDotC/xGuycx4CgbP48X/KF/586bcObxT0HENHXEU8Nqtu6NR+eKhw==", + "version": "5.0.6", + "resolved": "https://registry.npmjs.org/@types/express/-/express-5.0.6.tgz", + "integrity": "sha512-sKYVuV7Sv9fbPIt/442koC7+IIwK5olP1KWeD88e/idgoJqDm3JV/YUiPwkoKK92ylff2MGxSz1CSjsXelx0YA==", "dev": true, "dependencies": { "@types/body-parser": "*", - "@types/express-serve-static-core": "^4.17.33", - "@types/qs": "*", - "@types/serve-static": "^1" + "@types/express-serve-static-core": "^5.0.0", + "@types/serve-static": "^2" } }, "node_modules/@types/express-serve-static-core": { - "version": "4.17.36", - "resolved": "https://registry.npmjs.org/@types/express-serve-static-core/-/express-serve-static-core-4.17.36.tgz", - "integrity": "sha512-zbivROJ0ZqLAtMzgzIUC4oNqDG9iF0lSsAqpOD9kbs5xcIM3dTiyuHvBc7R8MtWBp3AAWGaovJa+wzWPjLYW7Q==", + "version": "5.1.2", + "resolved": "https://registry.npmjs.org/@types/express-serve-static-core/-/express-serve-static-core-5.1.2.tgz", + "integrity": "sha512-d3KvEXBSo/lOAMc2u6fkyDHBvetBHeqD7wm/AcXfLpSOQwlmG9D/aQ0SFswVjv05p7ullQS7Mjohj6/VdbZuTg==", "dev": true, "dependencies": { "@types/node": "*", @@ -415,20 +414,11 @@ } }, "node_modules/@types/http-errors": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/@types/http-errors/-/http-errors-2.0.2.tgz", - "integrity": "sha512-lPG6KlZs88gef6aD85z3HNkztpj7w2R7HmR3gygjfXCQmsLloWNARFkMuzKiiY8FGdh1XDpgBdrSf4aKDiA7Kg==", + "version": "2.0.5", + "resolved": "https://registry.npmjs.org/@types/http-errors/-/http-errors-2.0.5.tgz", + "integrity": "sha512-r8Tayk8HJnX0FztbZN7oVqGccWgw98T/0neJphO91KkmOzug1KkofZURD4UaD5uH8AqcFLfdPErnBod0u71/qg==", "dev": true }, - "node_modules/@types/http-proxy": { - "version": "1.17.12", - "resolved": "https://registry.npmjs.org/@types/http-proxy/-/http-proxy-1.17.12.tgz", - "integrity": "sha512-kQtujO08dVtQ2wXAuSFfk9ASy3sug4+ogFR8Kd8UgP8PEuc1/G/8yjYRmp//PcDNJEUKOza/MrQu15bouEUCiw==", - "dev": true, - "dependencies": { - "@types/node": "*" - } - }, "node_modules/@types/json-schema": { "version": "7.0.15", "resolved": "https://registry.npmjs.org/@types/json-schema/-/json-schema-7.0.15.tgz", @@ -459,13 +449,6 @@ "integrity": "sha512-EEhsLsD6UsDM1yFhAvy0Cjr6VwmpMWqFBCb9w07wVugF7w9nfajxLuVmngTIpgS6svCnm6Vaw+MZhoDCKnOfsw==", "dev": true }, - "node_modules/@types/retry": { - "version": "0.12.2", - "resolved": "https://registry.npmjs.org/@types/retry/-/retry-0.12.2.tgz", - "integrity": "sha512-XISRgDJ2Tc5q4TRqvgJtzsRkFYNJzZrhTdtMoGVBttwzzQJkPnS3WWTFc7kuDRoPtPakl+T+OfdEUjYJj7Jbow==", - "dev": true, - "license": "MIT" - }, "node_modules/@types/send": { "version": "0.17.1", "resolved": "https://registry.npmjs.org/@types/send/-/send-0.17.1.tgz", @@ -487,24 +470,12 @@ } }, "node_modules/@types/serve-static": { - "version": "1.15.7", - "resolved": "https://registry.npmjs.org/@types/serve-static/-/serve-static-1.15.7.tgz", - "integrity": "sha512-W8Ym+h8nhuRwaKPaDw34QUkwsGi6Rc4yYqvKFo5rm2FUEhCFbzVWrxXUxuKK8TASjWsysJY0nsmNCGhCOIsrOw==", + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/@types/serve-static/-/serve-static-2.2.0.tgz", + "integrity": "sha512-8mam4H1NHLtu7nmtalF7eyBH14QyOASmcxHhSfEoRyr0nP/YdoesEtU+uSRvMe96TW/HPTtkoKqQLl53N7UXMQ==", "dev": true, - "license": "MIT", "dependencies": { "@types/http-errors": "*", - "@types/node": "*", - "@types/send": "*" - } - }, - "node_modules/@types/sockjs": { - "version": "0.3.36", - "resolved": "https://registry.npmjs.org/@types/sockjs/-/sockjs-0.3.36.tgz", - "integrity": "sha512-MK9V6NzAS1+Ud7JV9lJLFqW85VbC9dq3LmwZCuBe4wBDgKC0Kj/jd8Xl+nSviU+Qc3+m7umHHyHg//2KSa0a0Q==", - "dev": true, - "license": "MIT", - "dependencies": { "@types/node": "*" } }, @@ -721,18 +692,43 @@ "dev": true }, "node_modules/accepts": { - "version": "1.3.8", - "resolved": "https://registry.npmjs.org/accepts/-/accepts-1.3.8.tgz", - "integrity": "sha512-PYAthTa2m2VKxuvSD3DPC/Gy+U+sOA1LAuT8mkmRuvw+NACSaeXEQ+NHcVF7rONl6qcaxV3Uuemwawk+7+SJLw==", + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/accepts/-/accepts-2.0.0.tgz", + "integrity": "sha512-5cvg6CtKwfgdmVqY1WIiXKc3Q1bkRqGLi+2W/6ao+6Y7gu/RCwRuAhGEzh5B4KlszSuTLgZYuqFqo5bImjNKng==", "dev": true, "dependencies": { - "mime-types": "~2.1.34", - "negotiator": "0.6.3" + "mime-types": "^3.0.0", + "negotiator": "^1.0.0" }, "engines": { "node": ">= 0.6" } }, + "node_modules/accepts/node_modules/mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true, + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/accepts/node_modules/mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", + "dev": true, + "dependencies": { + "mime-db": "^1.54.0" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, "node_modules/acorn": { "version": "8.15.0", "resolved": "https://registry.npmjs.org/acorn/-/acorn-8.15.0.tgz", @@ -814,27 +810,6 @@ "ansi-html": "bin/ansi-html" } }, - "node_modules/anymatch": { - "version": "3.1.3", - "resolved": "https://registry.npmjs.org/anymatch/-/anymatch-3.1.3.tgz", - "integrity": "sha512-KMReFUr0B4t+D+OBkjR3KYqvocp2XaSzO55UcB6mgQMd3KbcE+mWTyvVV7D/zsdEbNnV6acZUutkiHQXvTr1Rw==", - "dev": true, - "license": "ISC", - "dependencies": { - "normalize-path": "^3.0.0", - "picomatch": "^2.0.4" - }, - "engines": { - "node": ">= 8" - } - }, - "node_modules/array-flatten": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/array-flatten/-/array-flatten-1.1.1.tgz", - "integrity": "sha512-PCVAQswWemu6UdxsDFFX/+gVeYqKAod3D3UVm91jHwynguOwAvYPhx8nNlM++NqRcK6CxxpUafjmhIdKiHibqg==", - "dev": true, - "license": "MIT" - }, "node_modules/asn1js": { "version": "3.0.10", "resolved": "https://registry.npmjs.org/asn1js/-/asn1js-3.0.10.tgz", @@ -867,114 +842,44 @@ "node_modules/batch": { "version": "0.6.1", "resolved": "https://registry.npmjs.org/batch/-/batch-0.6.1.tgz", - "integrity": "sha1-3DQxT05nkxgJP8dgJyUl+UvyXBY=", + "integrity": "sha512-x+VAiMRL6UPkx+kudNvxTl6hB2XNNCG2r+7wixVfIYwu/2HKRXimwQyaumLjMveWvT2Hkd/cAJw+QBMfJ/EKVw==", "dev": true }, - "node_modules/binary-extensions": { - "version": "2.3.0", - "resolved": "https://registry.npmjs.org/binary-extensions/-/binary-extensions-2.3.0.tgz", - "integrity": "sha512-Ceh+7ox5qe7LJuLHoY0feh3pHuUDHAcRUeyL2VYghZwfpkNIy/+8Ocg0a3UuSoYzavmylwuLWQOf3hl0jjMMIw==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">=8" - }, - "funding": { - "url": "https://github.com/sponsors/sindresorhus" - } - }, "node_modules/body-parser": { - "version": "1.20.5", - "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.5.tgz", - "integrity": "sha512-3grm+/2tUOvu2cjJkvsIxrv/wVpfXQW4PsQHYm7yk4vfpu7Ekl6nEsYBoJUL6qDwZUx8wUhQ8tR2qz+ad9c9OA==", - "dev": true, - "dependencies": { - "bytes": "~3.1.2", - "content-type": "~1.0.5", - "debug": "2.6.9", - "depd": "2.0.0", - "destroy": "~1.2.0", - "http-errors": "~2.0.1", - "iconv-lite": "~0.4.24", - "on-finished": "~2.4.1", - "qs": "~6.15.1", - "raw-body": "~2.5.3", - "type-is": "~1.6.18", - "unpipe": "~1.0.0" - }, - "engines": { - "node": ">= 0.8", - "npm": "1.2.8000 || >= 1.4.16" - } - }, - "node_modules/body-parser/node_modules/debug": { - "version": "2.6.9", - "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", - "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", - "dev": true, - "license": "MIT", - "dependencies": { - "ms": "2.0.0" - } - }, - "node_modules/body-parser/node_modules/depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.8" - } - }, - "node_modules/body-parser/node_modules/http-errors": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", - "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-2.3.0.tgz", + "integrity": "sha512-2cGmJupaNgg+QUwVLAucDuWuoMZ6EX9iHDRswZ5lsNYEmwPaRknMPCLZz07yTzVq/83p4o/wzbDZbBrTvGGTIw==", "dev": true, "dependencies": { - "depd": "~2.0.0", - "inherits": "~2.0.4", - "setprototypeof": "~1.2.0", - "statuses": "~2.0.2", - "toidentifier": "~1.0.1" + "bytes": "^3.1.2", + "content-type": "^2.0.0", + "debug": "^4.4.3", + "http-errors": "^2.0.1", + "iconv-lite": "^0.7.2", + "on-finished": "^2.4.1", + "qs": "^6.15.2", + "raw-body": "^3.0.2", + "type-is": "^2.1.0" }, "engines": { - "node": ">= 0.8" + "node": ">=18" }, "funding": { "type": "opencollective", "url": "https://opencollective.com/express" } }, - "node_modules/body-parser/node_modules/inherits": { - "version": "2.0.4", - "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", - "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", - "dev": true - }, - "node_modules/body-parser/node_modules/qs": { - "version": "6.15.2", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.15.2.tgz", - "integrity": "sha512-Rzq0KEyX/w/tEybncDgdkZrJgVUsUMk3xjh3t5bv3S1HTAtg+uOYt72+ZfwiQwKdysThkTBdL/rTi6HDmX9Ddw==", + "node_modules/body-parser/node_modules/content-type": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-2.0.0.tgz", + "integrity": "sha512-j/O/d7GcZCyNl7/hwZAb606rzqkyvaDctLmckbxLzHvFBzTJHuGEdodATcP3yIRoDrLHkIATJuvzbFlp/ki2cQ==", "dev": true, - "dependencies": { - "side-channel": "^1.1.0" - }, "engines": { - "node": ">=0.6" + "node": ">=18" }, "funding": { - "url": "https://github.com/sponsors/ljharb" - } - }, - "node_modules/body-parser/node_modules/statuses": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", - "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", - "dev": true, - "engines": { - "node": ">= 0.8" + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/bonjour-service": { @@ -1044,7 +949,6 @@ "resolved": "https://registry.npmjs.org/bundle-name/-/bundle-name-4.1.0.tgz", "integrity": "sha512-tjwM5exMg6BGRI+kNmTntNsvdZS1X8BFYS6tnJ2hdH0kVxM6/eVZ2xy+FqStSWvYmtfFMDLIxurorHwDKfDz5Q==", "dev": true, - "license": "MIT", "dependencies": { "run-applescript": "^7.0.0" }, @@ -1078,7 +982,6 @@ "resolved": "https://registry.npmjs.org/call-bind-apply-helpers/-/call-bind-apply-helpers-1.0.2.tgz", "integrity": "sha512-Sp1ablJ0ivDkSzjcaJdxEunN5/XvksFJ2sMBFfq6x0ryhQV/2b/KwFe21cMpmHtPOSij8K99/wSfoEuTObmuMQ==", "dev": true, - "license": "MIT", "dependencies": { "es-errors": "^1.3.0", "function-bind": "^1.1.2" @@ -1092,7 +995,6 @@ "resolved": "https://registry.npmjs.org/call-bound/-/call-bound-1.0.4.tgz", "integrity": "sha512-+ys997U96po4Kx/ABpBCqhA9EuxJaQWDQg7295H4hBphv3IZg0boBKuwYpt4YXp6MZ5AmZQnU/tyMTlRpaSejg==", "dev": true, - "license": "MIT", "dependencies": { "call-bind-apply-helpers": "^1.0.2", "get-intrinsic": "^1.3.0" @@ -1125,28 +1027,18 @@ ] }, "node_modules/chokidar": { - "version": "3.6.0", - "resolved": "https://registry.npmjs.org/chokidar/-/chokidar-3.6.0.tgz", - "integrity": "sha512-7VT13fmjotKpGipCW9JEQAusEPE+Ei8nl6/g4FBAmIm0GOOLMua9NDDo/DWp0ZAxCr3cPq5ZpBqmPAQgDda2Pw==", + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/chokidar/-/chokidar-5.0.0.tgz", + "integrity": "sha512-TQMmc3w+5AxjpL8iIiwebF73dRDF4fBIieAqGn9RGCWaEVwQ6Fb2cGe31Yns0RRIzii5goJ1Y7xbMwo1TxMplw==", "dev": true, - "license": "MIT", "dependencies": { - "anymatch": "~3.1.2", - "braces": "~3.0.2", - "glob-parent": "~5.1.2", - "is-binary-path": "~2.1.0", - "is-glob": "~4.0.1", - "normalize-path": "~3.0.0", - "readdirp": "~3.6.0" + "readdirp": "^5.0.0" }, "engines": { - "node": ">= 8.10.0" + "node": ">= 20.19.0" }, "funding": { "url": "https://paulmillr.com/funding/" - }, - "optionalDependencies": { - "fsevents": "~2.3.2" } }, "node_modules/chrome-trace-event": { @@ -1265,65 +1157,44 @@ } }, "node_modules/content-disposition": { - "version": "0.5.4", - "resolved": "https://registry.npmjs.org/content-disposition/-/content-disposition-0.5.4.tgz", - "integrity": "sha512-FveZTNuGw04cxlAiWbzi6zTAL/lhehaWbTtgluJh4/E95DqMwTmha3KZN1aAWA8cFIhHzMZUvLevkw5Rqk+tSQ==", + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/content-disposition/-/content-disposition-1.1.0.tgz", + "integrity": "sha512-5jRCH9Z/+DRP7rkvY83B+yGIGX96OYdJmzngqnw2SBSxqCFPd0w2km3s5iawpGX8krnwSGmF0FW5Nhr0Hfai3g==", "dev": true, - "license": "MIT", - "dependencies": { - "safe-buffer": "5.2.1" - }, "engines": { - "node": ">= 0.6" + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, - "node_modules/content-disposition/node_modules/safe-buffer": { - "version": "5.2.1", - "resolved": "https://registry.npmjs.org/safe-buffer/-/safe-buffer-5.2.1.tgz", - "integrity": "sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ==", - "dev": true, - "funding": [ - { - "type": "github", - "url": "https://github.com/sponsors/feross" - }, - { - "type": "patreon", - "url": "https://www.patreon.com/feross" - }, - { - "type": "consulting", - "url": "https://feross.org/support" - } - ], - "license": "MIT" - }, "node_modules/content-type": { "version": "1.0.5", "resolved": "https://registry.npmjs.org/content-type/-/content-type-1.0.5.tgz", "integrity": "sha512-nTjqfcBFEipKdXCv4YDQWCfmcLZKm81ldF0pAopTvyrFGVbcR6P/VAAd5G7N+0tTr8QqiU0tFadD6FK4NtJwOA==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.6" } }, "node_modules/cookie": { - "version": "0.7.1", - "resolved": "https://registry.npmjs.org/cookie/-/cookie-0.7.1.tgz", - "integrity": "sha512-6DnInpx7SJ2AK3+CTUE/ZM0vWTUboZCegxhC2xiIydHR9jNuTAASBrfEpHhiGOZw/nX51bHt6YQl8jsGo4y/0w==", + "version": "0.7.2", + "resolved": "https://registry.npmjs.org/cookie/-/cookie-0.7.2.tgz", + "integrity": "sha512-yki5XnKuf750l50uGTllt6kKILY4nQ1eNIQatoXEByZ5dWgnKqbnqmTrBE5B4N7lrMJKQ2ytWMiTO2o0v6Ew/w==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.6" } }, "node_modules/cookie-signature": { - "version": "1.0.6", - "resolved": "https://registry.npmjs.org/cookie-signature/-/cookie-signature-1.0.6.tgz", - "integrity": "sha512-QADzlaHc8icV8I7vbaJXJwod9HWYp8uCqf1xa4OfNu1T7JVxQIrUgOWtHdNDtPiywmFbiS12VjotIXLrKM3orQ==", + "version": "1.2.2", + "resolved": "https://registry.npmjs.org/cookie-signature/-/cookie-signature-1.2.2.tgz", + "integrity": "sha512-D76uU73ulSXrD1UXF4KE2TMxVVwhsnCgfAyTg9k8P6KGZjlXKrOLe4dJQKI3Bxi5wjesZoFXJWElNWBjPZMbhg==", "dev": true, - "license": "MIT" + "engines": { + "node": ">=6.6.0" + } }, "node_modules/copy-webpack-plugin": { "version": "14.0.0", @@ -1370,12 +1241,6 @@ "node": ">=20.0.0" } }, - "node_modules/core-util-is": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/core-util-is/-/core-util-is-1.0.2.tgz", - "integrity": "sha1-tf1UIgqivFq1eqtxQMlAdUUDwac=", - "dev": true - }, "node_modules/cross-spawn": { "version": "7.0.3", "resolved": "https://registry.npmjs.org/cross-spawn/-/cross-spawn-7.0.3.tgz", @@ -1395,27 +1260,33 @@ "link": true }, "node_modules/debug": { - "version": "4.1.1", - "resolved": "https://registry.npmjs.org/debug/-/debug-4.1.1.tgz", - "integrity": "sha512-pYAIzeRo8J6KPEaJ0VWOh5Pzkbw/RetuzehGM7QRRX5he4fPHx2rdKMB256ehJCkX+XRQm16eZLqLNS8RSZXZw==", - "deprecated": "Debug versions >=3.2.0 <3.2.7 || >=4 <4.3.1 have a low-severity ReDos regression when used in a Node.js environment. It is recommended you upgrade to 3.2.7 or 4.3.1. (https://github.com/visionmedia/debug/issues/797)", + "version": "4.4.3", + "resolved": "https://registry.npmjs.org/debug/-/debug-4.4.3.tgz", + "integrity": "sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==", "dev": true, "dependencies": { - "ms": "^2.1.1" + "ms": "^2.1.3" + }, + "engines": { + "node": ">=6.0" + }, + "peerDependenciesMeta": { + "supports-color": { + "optional": true + } } }, "node_modules/debug/node_modules/ms": { - "version": "2.1.2", - "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.2.tgz", - "integrity": "sha512-sGkPx+VjMtmA6MX27oA4FBFELFCZZ4S4XqeGOXCv68tT+jb3vk/RyaKWP0PTKyWtmLSM0b+adUTEvbs1PEaH2w==", + "version": "2.1.3", + "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.3.tgz", + "integrity": "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==", "dev": true }, "node_modules/default-browser": { - "version": "5.2.1", - "resolved": "https://registry.npmjs.org/default-browser/-/default-browser-5.2.1.tgz", - "integrity": "sha512-WY/3TUME0x3KPYdRRxEJJvXRHV4PyPoUsxtZa78lwItwRQRHhd2U9xOscaT/YTf8uCXIAjeJOFBVEh/7FtD8Xg==", + "version": "5.5.0", + "resolved": "https://registry.npmjs.org/default-browser/-/default-browser-5.5.0.tgz", + "integrity": "sha512-H9LMLr5zwIbSxrmvikGuI/5KGhZ8E2zH3stkMgM5LpOWDutGM2JZaj460Udnf1a+946zc7YBgrqEWwbk7zHvGw==", "dev": true, - "license": "MIT", "dependencies": { "bundle-name": "^4.1.0", "default-browser-id": "^5.0.0" @@ -1428,11 +1299,10 @@ } }, "node_modules/default-browser-id": { - "version": "5.0.0", - "resolved": "https://registry.npmjs.org/default-browser-id/-/default-browser-id-5.0.0.tgz", - "integrity": "sha512-A6p/pu/6fyBcA1TRz/GqWYPViplrftcW2gZC9q79ngNCKAeR/X3gcEdXQHl4KNXV+3wgIJ1CPkJQ3IHM6lcsyA==", + "version": "5.0.1", + "resolved": "https://registry.npmjs.org/default-browser-id/-/default-browser-id-5.0.1.tgz", + "integrity": "sha512-x1VCxdX4t+8wVfd1so/9w+vQ4vx7lKd2Qp5tDRutErwmR85OgmfX7RlLRMWafRMY7hbEiXIbudNrjOAPa/hL8Q==", "dev": true, - "license": "MIT", "engines": { "node": ">=18" }, @@ -1445,7 +1315,6 @@ "resolved": "https://registry.npmjs.org/define-lazy-prop/-/define-lazy-prop-3.0.0.tgz", "integrity": "sha512-N+MeXYoqr3pOgn8xfyRPREN7gHakLYjhsHhWGT3fWAiL4IkAt0iDw14QiiEm2bE30c5XX5q0FtAA3CK5f9/BUg==", "dev": true, - "license": "MIT", "engines": { "node": ">=12" }, @@ -1454,31 +1323,14 @@ } }, "node_modules/depd": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/depd/-/depd-1.1.2.tgz", - "integrity": "sha1-m81S4UwJd2PnSbJ0xDRu0uVgtak=", - "dev": true, - "engines": { - "node": ">= 0.6" - } - }, - "node_modules/destroy": { - "version": "1.2.0", - "resolved": "https://registry.npmjs.org/destroy/-/destroy-1.2.0.tgz", - "integrity": "sha512-2sJGJTaXIIaR1w4iJSNoN0hnMY7Gpc/n8D4qSCJw8QqFWXf7cuAgnEHxBpweaVcPevC2l3KpjYCx3NypQQgaJg==", + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", + "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", "dev": true, - "license": "MIT", "engines": { - "node": ">= 0.8", - "npm": "1.2.8000 || >= 1.4.16" + "node": ">= 0.8" } }, - "node_modules/detect-node": { - "version": "2.0.4", - "resolved": "https://registry.npmjs.org/detect-node/-/detect-node-2.0.4.tgz", - "integrity": "sha512-ZIzRpLJrOj7jjP2miAtgqIfmzbxa4ZOr5jJc601zklsfEx9oTzmmj2nVpIPRpNlRTIh8lc1kyViIY7BWSGNmKw==", - "dev": true - }, "node_modules/dns-packet": { "version": "5.6.1", "resolved": "https://registry.npmjs.org/dns-packet/-/dns-packet-5.6.1.tgz", @@ -1497,7 +1349,6 @@ "resolved": "https://registry.npmjs.org/dunder-proto/-/dunder-proto-1.0.1.tgz", "integrity": "sha512-KIN/nDJBQRcXw0MLVhZE9iQHmG68qAVIBg9CqmUYjmQIhgij9U5MFvrqkUL5FbtyyzZuOeOt0zdeRe4UY7ct+A==", "dev": true, - "license": "MIT", "dependencies": { "call-bind-apply-helpers": "^1.0.1", "es-errors": "^1.3.0", @@ -1511,8 +1362,7 @@ "version": "1.1.1", "resolved": "https://registry.npmjs.org/ee-first/-/ee-first-1.1.1.tgz", "integrity": "sha512-WMwm9LhRUo+WUaRN+vRuETqG89IgZphVSNkdFgeb6sS/E4OrDIN7t48CAewSHXc6C8lefD8KKfr5vY61brQlow==", - "dev": true, - "license": "MIT" + "dev": true }, "node_modules/electron-to-chromium": { "version": "1.5.286", @@ -1525,7 +1375,6 @@ "resolved": "https://registry.npmjs.org/encodeurl/-/encodeurl-2.0.0.tgz", "integrity": "sha512-Q0n9HRi4m6JuGIV1eFlmvJB7ZEVxu93IrMyiMsGC0lrMJMWzRgx6WGquyfQgZVb31vhGgXnfmPNNXmxnOkRBrg==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.8" } @@ -1560,7 +1409,6 @@ "resolved": "https://registry.npmjs.org/es-define-property/-/es-define-property-1.0.1.tgz", "integrity": "sha512-e3nRfgfUZ4rNGL232gUgX06QNyyez04KdjFrF+LTRoOXmrOgFKDg4BCdsjW8EnT69eqdYGmRpJwiPVYNrCaW3g==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.4" } @@ -1570,7 +1418,6 @@ "resolved": "https://registry.npmjs.org/es-errors/-/es-errors-1.3.0.tgz", "integrity": "sha512-Zf5H2Kxt2xjTvbJvP2ZWLEICxA6j+hAmMzIlypy4xcBg1vKVnx89Wy0GbS+kf5cwCVFFzdCFh2XSCFNULS6csw==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.4" } @@ -1582,11 +1429,10 @@ "dev": true }, "node_modules/es-object-atoms": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/es-object-atoms/-/es-object-atoms-1.1.1.tgz", - "integrity": "sha512-FGgH2h8zKNim9ljj7dankFPcICIK9Cp5bm+c2gQSYePhpaG5+esrLODihIorn+Pe6FGJzWhXQotPv73jTaldXA==", + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/es-object-atoms/-/es-object-atoms-1.1.2.tgz", + "integrity": "sha512-HWcBoN6NileqtSydK2FqHbS/LoDd2pqrnQHLyJzBj4kOp/ky2MWMN694xOfkK8/SnUsW2DH7EfyVlydKCsm1Zw==", "dev": true, - "license": "MIT", "dependencies": { "es-errors": "^1.3.0" }, @@ -1606,7 +1452,7 @@ "node_modules/escape-html": { "version": "1.0.3", "resolved": "https://registry.npmjs.org/escape-html/-/escape-html-1.0.3.tgz", - "integrity": "sha1-Aljq5NPQwJdN4cFpGI7wBR0dGYg=", + "integrity": "sha512-NiSupZ4OeuGwr68lGIeym/ksIZMJodUGOSCZ/FSnTxcrekbvqrgdUxlJOMpijaKZVjAJrWrGs/6Jy8OMuyj9ow==", "dev": true }, "node_modules/eslint-scope": { @@ -1657,17 +1503,10 @@ "resolved": "https://registry.npmjs.org/etag/-/etag-1.8.1.tgz", "integrity": "sha512-aIL5Fx7mawVa300al2BnEE4iNvo1qETxLrPI/o05L7z6go7fCw1J6EQmbK4FmJ2AS7kgVF/KEZWufBfdClMcPg==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.6" } }, - "node_modules/eventemitter3": { - "version": "4.0.7", - "resolved": "https://registry.npmjs.org/eventemitter3/-/eventemitter3-4.0.7.tgz", - "integrity": "sha512-8guHBZCwKnFhYdHr2ysuRWErTwhoN2X8XELRlrRwpmfeY2jjuUN4taQMsULKUVo1K4DvZl+0pgfyoysHxvmvEw==", - "dev": true - }, "node_modules/events": { "version": "3.3.0", "resolved": "https://registry.npmjs.org/events/-/events-3.3.0.tgz", @@ -1678,100 +1517,71 @@ } }, "node_modules/express": { - "version": "4.22.1", - "resolved": "https://registry.npmjs.org/express/-/express-4.22.1.tgz", - "integrity": "sha512-F2X8g9P1X7uCPZMA3MVf9wcTqlyNp7IhH5qPCI0izhaOIYXaW9L535tGA3qmjRzpH+bZczqq7hVKxTR4NWnu+g==", - "dev": true, - "dependencies": { - "accepts": "~1.3.8", - "array-flatten": "1.1.1", - "body-parser": "~1.20.3", - "content-disposition": "~0.5.4", - "content-type": "~1.0.4", - "cookie": "~0.7.1", - "cookie-signature": "~1.0.6", - "debug": "2.6.9", - "depd": "2.0.0", - "encodeurl": "~2.0.0", - "escape-html": "~1.0.3", - "etag": "~1.8.1", - "finalhandler": "~1.3.1", - "fresh": "~0.5.2", - "http-errors": "~2.0.0", - "merge-descriptors": "1.0.3", - "methods": "~1.1.2", - "on-finished": "~2.4.1", - "parseurl": "~1.3.3", - "path-to-regexp": "~0.1.12", - "proxy-addr": "~2.0.7", - "qs": "~6.14.0", - "range-parser": "~1.2.1", - "safe-buffer": "5.2.1", - "send": "~0.19.0", - "serve-static": "~1.16.2", - "setprototypeof": "1.2.0", - "statuses": "~2.0.1", - "type-is": "~1.6.18", - "utils-merge": "1.0.1", - "vary": "~1.1.2" + "version": "5.2.1", + "resolved": "https://registry.npmjs.org/express/-/express-5.2.1.tgz", + "integrity": "sha512-hIS4idWWai69NezIdRt2xFVofaF4j+6INOpJlVOLDO8zXGpUVEVzIYk12UUi2JzjEzWL3IOAxcTubgz9Po0yXw==", + "dev": true, + "dependencies": { + "accepts": "^2.0.0", + "body-parser": "^2.2.1", + "content-disposition": "^1.0.0", + "content-type": "^1.0.5", + "cookie": "^0.7.1", + "cookie-signature": "^1.2.1", + "debug": "^4.4.0", + "depd": "^2.0.0", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "etag": "^1.8.1", + "finalhandler": "^2.1.0", + "fresh": "^2.0.0", + "http-errors": "^2.0.0", + "merge-descriptors": "^2.0.0", + "mime-types": "^3.0.0", + "on-finished": "^2.4.1", + "once": "^1.4.0", + "parseurl": "^1.3.3", + "proxy-addr": "^2.0.7", + "qs": "^6.14.0", + "range-parser": "^1.2.1", + "router": "^2.2.0", + "send": "^1.1.0", + "serve-static": "^2.2.0", + "statuses": "^2.0.1", + "type-is": "^2.0.1", + "vary": "^1.1.2" }, "engines": { - "node": ">= 0.10.0" + "node": ">= 18" }, "funding": { "type": "opencollective", "url": "https://opencollective.com/express" } }, - "node_modules/express/node_modules/debug": { - "version": "2.6.9", - "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", - "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", + "node_modules/express/node_modules/mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", "dev": true, - "license": "MIT", - "dependencies": { - "ms": "2.0.0" + "engines": { + "node": ">= 0.6" } }, - "node_modules/express/node_modules/depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", + "node_modules/express/node_modules/mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", "dev": true, - "license": "MIT", + "dependencies": { + "mime-db": "^1.54.0" + }, "engines": { - "node": ">= 0.8" - } - }, - "node_modules/express/node_modules/safe-buffer": { - "version": "5.2.1", - "resolved": "https://registry.npmjs.org/safe-buffer/-/safe-buffer-5.2.1.tgz", - "integrity": "sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ==", - "dev": true, - "funding": [ - { - "type": "github", - "url": "https://github.com/sponsors/feross" - }, - { - "type": "patreon", - "url": "https://www.patreon.com/feross" - }, - { - "type": "consulting", - "url": "https://feross.org/support" - } - ], - "license": "MIT" - }, - "node_modules/express/node_modules/statuses": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", - "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.8" + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/fast-deep-equal": { @@ -1781,9 +1591,9 @@ "dev": true }, "node_modules/fast-uri": { - "version": "3.1.0", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.0.tgz", - "integrity": "sha512-iPeeDKJSWf4IEOasVVrknXpaBV0IApz/gp7S2bb7Z4Lljbl2MGJRqInZiUrQwV16cpzw/D3S5j5Julj/gT52AA==", + "version": "3.1.5", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", + "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", "dev": true, "funding": [ { @@ -1805,18 +1615,6 @@ "node": ">= 4.9.1" } }, - "node_modules/faye-websocket": { - "version": "0.11.4", - "resolved": "https://registry.npmjs.org/faye-websocket/-/faye-websocket-0.11.4.tgz", - "integrity": "sha512-CzbClwlXAuiRQAlUyfqPgvPoNKTckTPGfwZV4ZdAhVcP2lh9KUxJg2b5GkE7XbjKQ3YJnQ9z6D9ntLAlB+tP8g==", - "dev": true, - "dependencies": { - "websocket-driver": ">=0.5.1" - }, - "engines": { - "node": ">=0.8.0" - } - }, "node_modules/fill-range": { "version": "7.1.1", "resolved": "https://registry.npmjs.org/fill-range/-/fill-range-7.1.1.tgz", @@ -1830,42 +1628,24 @@ } }, "node_modules/finalhandler": { - "version": "1.3.1", - "resolved": "https://registry.npmjs.org/finalhandler/-/finalhandler-1.3.1.tgz", - "integrity": "sha512-6BN9trH7bp3qvnrRyzsBz+g3lZxTNZTbVO2EV1CS0WIcDbawYVdYvGflME/9QP0h0pYlCDBCTjYa9nZzMDpyxQ==", + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/finalhandler/-/finalhandler-2.1.1.tgz", + "integrity": "sha512-S8KoZgRZN+a5rNwqTxlZZePjT/4cnm0ROV70LedRHZ0p8u9fRID0hJUZQpkKLzro8LfmC8sx23bY6tVNxv8pQA==", "dev": true, - "license": "MIT", "dependencies": { - "debug": "2.6.9", - "encodeurl": "~2.0.0", - "escape-html": "~1.0.3", - "on-finished": "2.4.1", - "parseurl": "~1.3.3", - "statuses": "2.0.1", - "unpipe": "~1.0.0" + "debug": "^4.4.0", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "on-finished": "^2.4.1", + "parseurl": "^1.3.3", + "statuses": "^2.0.1" }, "engines": { - "node": ">= 0.8" - } - }, - "node_modules/finalhandler/node_modules/debug": { - "version": "2.6.9", - "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", - "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", - "dev": true, - "license": "MIT", - "dependencies": { - "ms": "2.0.0" - } - }, - "node_modules/finalhandler/node_modules/statuses": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", - "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.8" + "node": ">= 18.0.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/find-up": { @@ -1881,59 +1661,22 @@ "node": ">=8" } }, - "node_modules/follow-redirects": { - "version": "1.16.0", - "resolved": "https://registry.npmjs.org/follow-redirects/-/follow-redirects-1.16.0.tgz", - "integrity": "sha512-y5rN/uOsadFT/JfYwhxRS5R7Qce+g3zG97+JrtFZlC9klX/W5hD7iiLzScI4nZqUS7DNUdhPgw4xI8W2LuXlUw==", - "dev": true, - "funding": [ - { - "type": "individual", - "url": "https://github.com/sponsors/RubenVerborgh" - } - ], - "engines": { - "node": ">=4.0" - }, - "peerDependenciesMeta": { - "debug": { - "optional": true - } - } - }, "node_modules/forwarded": { "version": "0.2.0", "resolved": "https://registry.npmjs.org/forwarded/-/forwarded-0.2.0.tgz", "integrity": "sha512-buRG0fpBtRHSTCOASe6hD258tEubFoRLb4ZNA6NxMVHNw2gOcwHo9wyablzMzOA5z9xA9L1KNjk/Nt6MT9aYow==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.6" } }, "node_modules/fresh": { - "version": "0.5.2", - "resolved": "https://registry.npmjs.org/fresh/-/fresh-0.5.2.tgz", - "integrity": "sha512-zJ2mQYM18rEFOudeV4GShTGIQ7RbzA7ozbU9I/XBpm7kqgMywgmylMwXHxZJmkVoYkna9d2pVXVXPdYTP9ej8Q==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.6" - } - }, - "node_modules/fsevents": { - "version": "2.3.3", - "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", - "integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==", + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/fresh/-/fresh-2.0.0.tgz", + "integrity": "sha512-Rx/WycZ60HOaqLKAi6cHRKKI7zxWbJ31MhntmtwMoaTeF7XFH9hhBp8vITaMidfljRQ6eYWCKkaTK+ykVJHP2A==", "dev": true, - "hasInstallScript": true, - "license": "MIT", - "optional": true, - "os": [ - "darwin" - ], "engines": { - "node": "^8.16.0 || ^10.6.0 || >=11.0.0" + "node": ">= 0.8" } }, "node_modules/function-bind": { @@ -1950,7 +1693,6 @@ "resolved": "https://registry.npmjs.org/get-intrinsic/-/get-intrinsic-1.3.0.tgz", "integrity": "sha512-9fSjSaos/fRIVIp+xSJlE6lfwhES7LNtKaCBIamHsjr2na1BiABJPo0mOjjz8GJDURarmCPGqaiVg5mfjb98CQ==", "dev": true, - "license": "MIT", "dependencies": { "call-bind-apply-helpers": "^1.0.2", "es-define-property": "^1.0.1", @@ -1975,7 +1717,6 @@ "resolved": "https://registry.npmjs.org/get-proto/-/get-proto-1.0.1.tgz", "integrity": "sha512-sTSfBjoXBp89JvIKIefqw7U2CCebsc74kiY6awiGogKtoSGbgjYE/G/+l9sF3MWFPNc9IcoOC4ODfKHfxFmp0g==", "dev": true, - "license": "MIT", "dependencies": { "dunder-proto": "^1.0.1", "es-object-atoms": "^1.0.0" @@ -1984,18 +1725,6 @@ "node": ">= 0.4" } }, - "node_modules/glob-parent": { - "version": "5.1.2", - "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-5.1.2.tgz", - "integrity": "sha512-AOIgSQCepiJYwP3ARnGx+5VnTu2HBYdzbGP45eLw1vr3zB3vZLeyed1sC9hnbcOc9/SrMyM5RPQrkGz4aS9Zow==", - "dev": true, - "dependencies": { - "is-glob": "^4.0.1" - }, - "engines": { - "node": ">= 6" - } - }, "node_modules/glob-to-regexp": { "version": "0.4.1", "resolved": "https://registry.npmjs.org/glob-to-regexp/-/glob-to-regexp-0.4.1.tgz", @@ -2007,7 +1736,6 @@ "resolved": "https://registry.npmjs.org/gopd/-/gopd-1.2.0.tgz", "integrity": "sha512-ZUKRh6/kUFoAiTAtTYPZJ3hw9wNxx+BIBOijnlG9PnrJsCcSjs1wyyD6vJpaYtgnzDrKYRSqf3OO6Rfa93xsRg==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.4" }, @@ -2021,12 +1749,6 @@ "integrity": "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==", "dev": true }, - "node_modules/handle-thing": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/handle-thing/-/handle-thing-2.0.1.tgz", - "integrity": "sha512-9Qn4yBxelxoh2Ow62nP+Ka/kMnOXRi8BXnRaUwezLNhqelnN49xKz4F/dPP8OYLxLxq6JDtZb2i9XznUQbNPTg==", - "dev": true - }, "node_modules/has": { "version": "1.0.3", "resolved": "https://registry.npmjs.org/has/-/has-1.0.3.tgz", @@ -2053,7 +1775,6 @@ "resolved": "https://registry.npmjs.org/has-symbols/-/has-symbols-1.1.0.tgz", "integrity": "sha512-1cDNdwJ2Jaohmb3sg4OmKaMBwuC48sYni5HUw2DvsC8LjGTLK9h+eb1X6RyuOHe4hT0ULCW68iomhjUoKUqlPQ==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.4" }, @@ -2062,11 +1783,10 @@ } }, "node_modules/hasown": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.2.tgz", - "integrity": "sha512-0hJU9SCPvmMzIBdZFqNPXWa6dqh7WdH0cII9y+CyS8rG3nL48Bclra9HmKhVVUHyPWNH5Y7xDwAB7bfgSjkUMQ==", + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.4.tgz", + "integrity": "sha512-T2UbfbBEF32wiepXIsMlTW9+dDYC6wMh/t/vYA4tuOMKqWz/n3vr1NFSxQiyP+zk2mXsoMA/i/7qV6LKut1t1A==", "dev": true, - "license": "MIT", "dependencies": { "function-bind": "^1.1.2" }, @@ -2074,133 +1794,71 @@ "node": ">= 0.4" } }, - "node_modules/hpack.js": { - "version": "2.1.6", - "resolved": "https://registry.npmjs.org/hpack.js/-/hpack.js-2.1.6.tgz", - "integrity": "sha1-h3dMCUnlE/QuhFdbPEVoH63ioLI=", - "dev": true, - "dependencies": { - "inherits": "^2.0.1", - "obuf": "^1.0.0", - "readable-stream": "^2.0.1", - "wbuf": "^1.1.0" - } - }, - "node_modules/http-deceiver": { - "version": "1.2.7", - "resolved": "https://registry.npmjs.org/http-deceiver/-/http-deceiver-1.2.7.tgz", - "integrity": "sha1-+nFolEq5pRnTN8sL7HKE3D5yPYc=", - "dev": true - }, "node_modules/http-errors": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.0.tgz", - "integrity": "sha512-FtwrG/euBzaEjYeRqOgly7G0qviiXoJWnvEH2Z1plBdXgbyjv34pHTSb9zoeHMyDy33+DWy5Wt9Wo+TURtOYSQ==", + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", + "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", "dev": true, - "license": "MIT", "dependencies": { - "depd": "2.0.0", - "inherits": "2.0.4", - "setprototypeof": "1.2.0", - "statuses": "2.0.1", - "toidentifier": "1.0.1" + "depd": "~2.0.0", + "inherits": "~2.0.4", + "setprototypeof": "~1.2.0", + "statuses": "~2.0.2", + "toidentifier": "~1.0.1" }, "engines": { "node": ">= 0.8" - } - }, - "node_modules/http-errors/node_modules/depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.8" - } - }, - "node_modules/http-errors/node_modules/inherits": { - "version": "2.0.4", - "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", - "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", - "dev": true, - "license": "ISC" - }, - "node_modules/http-errors/node_modules/statuses": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", - "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.8" - } - }, - "node_modules/http-parser-js": { - "version": "0.5.8", - "resolved": "https://registry.npmjs.org/http-parser-js/-/http-parser-js-0.5.8.tgz", - "integrity": "sha512-SGeBX54F94Wgu5RH3X5jsDtf4eHyRogWX1XGT3b4HuW3tQPM4AaBzoUji/4AAJNXCEOWZ5O0DgZmJw1947gD5Q==", - "dev": true - }, - "node_modules/http-proxy": { - "version": "1.18.1", - "resolved": "https://registry.npmjs.org/http-proxy/-/http-proxy-1.18.1.tgz", - "integrity": "sha512-7mz/721AbnJwIVbnaSv1Cz3Am0ZLT/UBwkC92VlxhXv/k/BBQfM2fXElQNC27BVGr0uwUpplYPQM9LnaBMR5NQ==", - "dev": true, - "dependencies": { - "eventemitter3": "^4.0.0", - "follow-redirects": "^1.0.0", - "requires-port": "^1.0.0" }, - "engines": { - "node": ">=8.0.0" + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/http-proxy-middleware": { - "version": "2.0.9", - "resolved": "https://registry.npmjs.org/http-proxy-middleware/-/http-proxy-middleware-2.0.9.tgz", - "integrity": "sha512-c1IyJYLYppU574+YI7R4QyX2ystMtVXZwIdzazUIPIJsHuWNd+mho2j+bKoHftndicGj9yh+xjd+l0yj7VeT1Q==", + "version": "4.2.0", + "resolved": "https://registry.npmjs.org/http-proxy-middleware/-/http-proxy-middleware-4.2.0.tgz", + "integrity": "sha512-ZA+oNOoM+GLoFTIzhkJptVQov73Srep2LBqhF8hG8CIPKO3nam1jonXVQ/QUH8RbwsmaaVz2SOJdzBNBHNtKbw==", "dev": true, - "license": "MIT", "dependencies": { - "@types/http-proxy": "^1.17.8", - "http-proxy": "^1.18.1", - "is-glob": "^4.0.1", - "is-plain-obj": "^3.0.0", - "micromatch": "^4.0.2" + "debug": "^4.4.3", + "httpxy": "^0.5.4", + "is-glob": "^4.0.3", + "is-plain-obj": "^4.1.0", + "micromatch": "^4.0.8" }, "engines": { - "node": ">=12.0.0" - }, - "peerDependencies": { - "@types/express": "^4.17.13" - }, - "peerDependenciesMeta": { - "@types/express": { - "optional": true - } + "node": "^22.15.0 || ^24.0.0 || >=26.0.0" } }, + "node_modules/httpxy": { + "version": "0.5.5", + "resolved": "https://registry.npmjs.org/httpxy/-/httpxy-0.5.5.tgz", + "integrity": "sha512-uDjmnPyp1q4Sgzf3w+J/Fc6UqcCEj0x4Wjp7OqK5dGhNeDgpyrAmnS6ey8QWrX3SWDon2DMKf9sBa5X9+CVyMA==", + "dev": true + }, "node_modules/hyperdyperid": { "version": "1.2.0", "resolved": "https://registry.npmjs.org/hyperdyperid/-/hyperdyperid-1.2.0.tgz", "integrity": "sha512-Y93lCzHYgGWdrJ66yIktxiaGULYc6oGiABxhcO5AufBeOyoIdZF7bIfLaOrbM0iGIOXQQgxxRrFEnb+Y6w1n4A==", "dev": true, - "license": "MIT", "engines": { "node": ">=10.18" } }, "node_modules/iconv-lite": { - "version": "0.4.24", - "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz", - "integrity": "sha512-v3MXnZAcvnywkTUEZomIActle7RXXeedOR31wwl7VlyoXO4Qi9arvSenNQWne1TcRwhCL1HwLI21bEqdpj8/rA==", + "version": "0.7.3", + "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.7.3.tgz", + "integrity": "sha512-IKXpvIzjnC9XTAUbVBcMfGS0EPaIXtW6v+zr+RRp+hqULEpo0owZax6wyRwPOJbWbzjYspQwusTsfVr0ifh4uQ==", "dev": true, "dependencies": { - "safer-buffer": ">= 2.1.2 < 3" + "safer-buffer": ">= 2.1.2 < 3.0.0" }, "engines": { "node": ">=0.10.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/import-local": { @@ -2223,9 +1881,9 @@ } }, "node_modules/inherits": { - "version": "2.0.3", - "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.3.tgz", - "integrity": "sha1-Yzwsg+PaQqUC9SRmAiSA9CCCYd4=", + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", + "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", "dev": true }, "node_modules/interpret": { @@ -2238,27 +1896,14 @@ } }, "node_modules/ipaddr.js": { - "version": "2.1.0", - "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-2.1.0.tgz", - "integrity": "sha512-LlbxQ7xKzfBusov6UMi4MFpEg0m+mAm9xyNGEduwXMEDuf4WfzB/RZwMVYEd7IKGvh4IUkEXYxtAVu9T3OelJQ==", + "version": "2.4.0", + "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-2.4.0.tgz", + "integrity": "sha512-9VGk3HGanVE6JoZXHiCpnGy5X0jYDnN4EA4lntFPj+1vIWlFhIylq2CrrCOJH9EAhc5CYhq18F2Av2tgoAPsYQ==", "dev": true, "engines": { "node": ">= 10" } }, - "node_modules/is-binary-path": { - "version": "2.1.0", - "resolved": "https://registry.npmjs.org/is-binary-path/-/is-binary-path-2.1.0.tgz", - "integrity": "sha512-ZMERYes6pDydyuGidse7OsHxtbI7WVeUEozgR/g7rd0xUimYNlvZRE/K2MgZTjWy725IfelLeVcEM97mmtRGXw==", - "dev": true, - "license": "MIT", - "dependencies": { - "binary-extensions": "^2.0.0" - }, - "engines": { - "node": ">=8" - } - }, "node_modules/is-core-module": { "version": "2.13.0", "resolved": "https://registry.npmjs.org/is-core-module/-/is-core-module-2.13.0.tgz", @@ -2276,7 +1921,6 @@ "resolved": "https://registry.npmjs.org/is-docker/-/is-docker-3.0.0.tgz", "integrity": "sha512-eljcgEDlEns/7AXFosB5K/2nCM4P7FQPkGc/DWLy5rmFEWvZayGrik1d9/QIY5nJ4f9YsVvBkA6kJpHn9rISdQ==", "dev": true, - "license": "MIT", "bin": { "is-docker": "cli.js" }, @@ -2308,12 +1952,23 @@ "node": ">=0.10.0" } }, + "node_modules/is-in-ssh": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/is-in-ssh/-/is-in-ssh-1.0.0.tgz", + "integrity": "sha512-jYa6Q9rH90kR1vKB6NM7qqd1mge3Fx4Dhw5TVlK1MUBqhEOuCagrEHMevNuCcbECmXZ0ThXkRm+Ymr51HwEPAw==", + "dev": true, + "engines": { + "node": ">=20" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } + }, "node_modules/is-inside-container": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/is-inside-container/-/is-inside-container-1.0.0.tgz", "integrity": "sha512-KIYLCCJghfHZxqjYBE7rEy0OBuTd5xCHS7tHVgvCLkx7StIoaxwNW3hCALgEUjFfeRk+MG/Qxmp/vtETEF3tRA==", "dev": true, - "license": "MIT", "dependencies": { "is-docker": "^3.0.0" }, @@ -2328,11 +1983,10 @@ } }, "node_modules/is-network-error": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/is-network-error/-/is-network-error-1.1.0.tgz", - "integrity": "sha512-tUdRRAnhT+OtCZR/LxZelH/C7QtjtFrTu5tXCA8pl55eTUElUHT+GPYV8MBMBvea/j+NxQqVt3LbWMRir7Gx9g==", + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/is-network-error/-/is-network-error-1.3.2.tgz", + "integrity": "sha512-PhBY86zaxNZUuWP6h13Vu5oFe0XY6/UlKzQnYFELzGVHygP3MxmvTfYSG7GN3aIab/iWudSMgjSnG9Dq+nHrgA==", "dev": true, - "license": "MIT", "engines": { "node": ">=16" }, @@ -2350,12 +2004,12 @@ } }, "node_modules/is-plain-obj": { - "version": "3.0.0", - "resolved": "https://registry.npmjs.org/is-plain-obj/-/is-plain-obj-3.0.0.tgz", - "integrity": "sha512-gwsOE28k+23GP1B6vFl1oVh/WOzmawBrKwo5Ev6wMKzPkaXaCDIQKzLnvsA42DRlbVTWorkgTKIviAKCWkfUwA==", + "version": "4.1.0", + "resolved": "https://registry.npmjs.org/is-plain-obj/-/is-plain-obj-4.1.0.tgz", + "integrity": "sha512-+Pgi+vMuUNkJyExiMBt5IlFoMyKnr5zhJ4Uspz58WOhBF5QoIZkFyNHIbBAtHwzVAgk5RtndVNsDRN61/mmDqg==", "dev": true, "engines": { - "node": ">=10" + "node": ">=12" }, "funding": { "url": "https://github.com/sponsors/sindresorhus" @@ -2373,12 +2027,17 @@ "node": ">=0.10.0" } }, + "node_modules/is-promise": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/is-promise/-/is-promise-4.0.0.tgz", + "integrity": "sha512-hvpoI6korhJMnej285dSg6nu1+e6uxs7zG3BYAm5byqDsgJNWwxzM6z6iZiAgQR4TJ30JmBTOwqZUw3WlyH3AQ==", + "dev": true + }, "node_modules/is-wsl": { - "version": "3.1.0", - "resolved": "https://registry.npmjs.org/is-wsl/-/is-wsl-3.1.0.tgz", - "integrity": "sha512-UcVfVfaK4Sc4m7X3dUSoHoozQGBEFeDC+zVo06t98xe8CzHSZZBekNXH+tu0NalHolcJ/QAGqS46Hef7QXBIMw==", + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/is-wsl/-/is-wsl-3.1.1.tgz", + "integrity": "sha512-e6rvdUCiQCAuumZslxRJWR/Doq4VpPR82kqclvcS0efgt430SlGIk05vdCN58+VrzgtIcfNODjozVielycD4Sw==", "dev": true, - "license": "MIT", "dependencies": { "is-inside-container": "^1.0.0" }, @@ -2389,12 +2048,6 @@ "url": "https://github.com/sponsors/sindresorhus" } }, - "node_modules/isarray": { - "version": "1.0.0", - "resolved": "https://registry.npmjs.org/isarray/-/isarray-1.0.0.tgz", - "integrity": "sha1-u5NdSFgsuhaMBoNJV6VKPgcSTxE=", - "dev": true - }, "node_modules/isexe": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/isexe/-/isexe-2.0.0.tgz", @@ -2446,14 +2099,13 @@ } }, "node_modules/launch-editor": { - "version": "2.10.0", - "resolved": "https://registry.npmjs.org/launch-editor/-/launch-editor-2.10.0.tgz", - "integrity": "sha512-D7dBRJo/qcGX9xlvt/6wUYzQxjh5G1RvZPgPv8vi4KRU99DVQL/oW7tnVOCCTm2HGeo3C5HvGE5Yrh6UBoZ0vA==", + "version": "2.14.1", + "resolved": "https://registry.npmjs.org/launch-editor/-/launch-editor-2.14.1.tgz", + "integrity": "sha512-QWBrQsMpH7gPr965dsKD/3cKWiNoTjpATQf++Xq63N6sKRGMwlVXz41O1IZTMfZQgBctD/K5Zt06+/I6pP6+HA==", "dev": true, - "license": "MIT", "dependencies": { - "picocolors": "^1.0.0", - "shell-quote": "^1.8.1" + "picocolors": "^1.1.1", + "shell-quote": "^1.8.4" } }, "node_modules/loader-runner": { @@ -2486,39 +2138,50 @@ "resolved": "https://registry.npmjs.org/math-intrinsics/-/math-intrinsics-1.1.0.tgz", "integrity": "sha512-/IXtbwEk5HTPyEwyKX6hGkYXxM9nbj64B+ilVJnC/R6B0pH5G4V3b0pVbL7DBj4tkhBAppbQUlf6F6Xl9LHu1g==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.4" } }, "node_modules/media-typer": { - "version": "0.3.0", - "resolved": "https://registry.npmjs.org/media-typer/-/media-typer-0.3.0.tgz", - "integrity": "sha512-dq+qelQ9akHpcOl/gUVRTxVIOkAJ1wR3QAvb4RsVjS8oVoFjDGTc679wJYmUmknUF5HwMLOgb5O+a3KxfWapPQ==", + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/media-typer/-/media-typer-1.1.1.tgz", + "integrity": "sha512-yz3xRaG20c6/BOzvYoDaGtPmGscs7YivItZEEqe6GbwNfHuxu9YNmvnEkMzKldAGY4/80pRcQRZSEnhquk9XuQ==", "dev": true, - "license": "MIT", "engines": { - "node": ">= 0.6" + "node": ">= 0.8" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/memfs": { - "version": "4.17.2", - "resolved": "https://registry.npmjs.org/memfs/-/memfs-4.17.2.tgz", - "integrity": "sha512-NgYhCOWgovOXSzvYgUW0LQ7Qy72rWQMGGFJDoWg4G30RHd3z77VbYdtJ4fembJXBy8pMIUA31XNAupobOQlwdg==", - "dev": true, - "license": "Apache-2.0", - "dependencies": { - "@jsonjoy.com/json-pack": "^1.0.3", - "@jsonjoy.com/util": "^1.3.0", - "tree-dump": "^1.0.1", + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/memfs/-/memfs-4.64.0.tgz", + "integrity": "sha512-Kw72fgY7Wn+sD8KmtNWSafl1dz0UvAsE/PHs3YVfLiaZuA3HxNm9sRLqAu0ATiBGJvME1PxZXbBZPv5GycDeAw==", + "dev": true, + "dependencies": { + "@jsonjoy.com/fs-core": "4.64.0", + "@jsonjoy.com/fs-fsa": "4.64.0", + "@jsonjoy.com/fs-node": "4.64.0", + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-to-fsa": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "@jsonjoy.com/fs-print": "4.64.0", + "@jsonjoy.com/fs-snapshot": "4.64.0", + "@jsonjoy.com/json-pack": "^1.11.0", + "@jsonjoy.com/util": "^1.9.0", + "glob-to-regex.js": "^1.0.1", + "thingies": "^2.5.0", + "tree-dump": "^1.0.3", "tslib": "^2.0.0" }, - "engines": { - "node": ">= 4.0.0" - }, "funding": { "type": "github", "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" } }, "node_modules/memfs/node_modules/@jsonjoy.com/base64": { @@ -2526,7 +2189,6 @@ "resolved": "https://registry.npmjs.org/@jsonjoy.com/base64/-/base64-1.1.2.tgz", "integrity": "sha512-q6XAnWQDIMA3+FTiOYajoYqySkO+JSat0ytXGSuRdq9uXE7o92gzuQwQM14xaCRlBLGq3v5miDGC4vkVTn54xA==", "dev": true, - "license": "Apache-2.0", "engines": { "node": ">=10.0" }, @@ -2538,18 +2200,11 @@ "tslib": "2" } }, - "node_modules/memfs/node_modules/@jsonjoy.com/json-pack": { - "version": "1.2.0", - "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pack/-/json-pack-1.2.0.tgz", - "integrity": "sha512-io1zEbbYcElht3tdlqEOFxZ0dMTYrHz9iMf0gqn1pPjZFTCgM5R4R5IMA20Chb2UPYYsxjzs8CgZ7Nb5n2K2rA==", + "node_modules/memfs/node_modules/@jsonjoy.com/buffers": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/buffers/-/buffers-17.67.0.tgz", + "integrity": "sha512-tfExRpYxBvi32vPs9ZHaTjSP4fHAfzSmcahOfNxtvGHcyJel+aibkPlGeBB+7AoC6hL7lXIE++8okecBxx7lcw==", "dev": true, - "license": "Apache-2.0", - "dependencies": { - "@jsonjoy.com/base64": "^1.1.1", - "@jsonjoy.com/util": "^1.1.2", - "hyperdyperid": "^1.2.0", - "thingies": "^1.20.0" - }, "engines": { "node": ">=10.0" }, @@ -2561,12 +2216,11 @@ "tslib": "2" } }, - "node_modules/memfs/node_modules/@jsonjoy.com/util": { - "version": "1.6.0", - "resolved": "https://registry.npmjs.org/@jsonjoy.com/util/-/util-1.6.0.tgz", - "integrity": "sha512-sw/RMbehRhN68WRtcKCpQOPfnH6lLP4GJfqzi3iYej8tnzpZUDr6UkZYJjcjjC0FWEJOJbyM3PTIwxucUmDG2A==", + "node_modules/memfs/node_modules/@jsonjoy.com/codegen": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/codegen/-/codegen-1.0.0.tgz", + "integrity": "sha512-E8Oy+08cmCf0EK/NMxpaJZmOxPqM+6iSe2S4nlSBrPZOORoDJILxtbSUEDKQyTamm/BVAhIGllOBNU79/dwf0g==", "dev": true, - "license": "Apache-2.0", "engines": { "node": ">=10.0" }, @@ -2578,25 +2232,38 @@ "tslib": "2" } }, - "node_modules/memfs/node_modules/thingies": { - "version": "1.21.0", - "resolved": "https://registry.npmjs.org/thingies/-/thingies-1.21.0.tgz", - "integrity": "sha512-hsqsJsFMsV+aD4s3CWKk85ep/3I9XzYV/IXaSouJMYIoDlgyi11cBhsqYe9/geRfB0YIikBQg6raRaM+nIMP9g==", + "node_modules/memfs/node_modules/@jsonjoy.com/fs-core": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-core/-/fs-core-4.64.0.tgz", + "integrity": "sha512-zs2TAq7Six5jgMuoMNjpspAvOP3mhtgq/k1UyQodEzCtQi/N83y2/y+zcvnZSGp/Rxq96DBN+bValOBQAyn/ew==", "dev": true, - "license": "Unlicense", + "dependencies": { + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "thingies": "^2.5.0" + }, "engines": { - "node": ">=10.18" + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" }, "peerDependencies": { - "tslib": "^2" + "tslib": "2" } }, - "node_modules/memfs/node_modules/tree-dump": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/tree-dump/-/tree-dump-1.0.3.tgz", - "integrity": "sha512-il+Cv80yVHFBwokQSfd4bldvr1Md951DpgAGfmhydt04L+YzHgubm2tQ7zueWDcGENKHq0ZvGFR/hjvNXilHEg==", + "node_modules/memfs/node_modules/@jsonjoy.com/fs-fsa": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-fsa/-/fs-fsa-4.64.0.tgz", + "integrity": "sha512-nMWOVbkLFyEgmXZih3wyvxA9XpgyyqyfrINMHvEFqhi7uqfRl7c9ERJt6yX7vgMPrB9Uo+OJO+Spa0cFzPD01w==", "dev": true, - "license": "Apache-2.0", + "dependencies": { + "@jsonjoy.com/fs-core": "4.64.0", + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "thingies": "^2.5.0" + }, "engines": { "node": ">=10.0" }, @@ -2608,45 +2275,402 @@ "tslib": "2" } }, - "node_modules/memfs/node_modules/tslib": { - "version": "2.8.1", - "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", - "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", - "dev": true, - "license": "0BSD" - }, - "node_modules/merge-descriptors": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/merge-descriptors/-/merge-descriptors-1.0.3.tgz", - "integrity": "sha512-gaNvAS7TZ897/rVaZ0nMtAyxNyi/pdbjbAwUpFQpN70GqnVfOiXpeUUMKRBmzXaSQ8DdTX4/0ms62r2K+hE6mQ==", + "node_modules/memfs/node_modules/@jsonjoy.com/fs-node": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-node/-/fs-node-4.64.0.tgz", + "integrity": "sha512-dO+NNkODbUli4uV42bcNrrLvq5rE7SNpdZ5TNd0dtbLsAaNK3MDiIC9lUi+brboGoIjW6vd2fB1qao60nrk5xA==", "dev": true, - "license": "MIT", - "funding": { - "url": "https://github.com/sponsors/sindresorhus" - } + "dependencies": { + "@jsonjoy.com/fs-core": "4.64.0", + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "@jsonjoy.com/fs-print": "4.64.0", + "@jsonjoy.com/fs-snapshot": "4.64.0", + "glob-to-regex.js": "^1.0.0", + "thingies": "^2.5.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } }, - "node_modules/merge-stream": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/merge-stream/-/merge-stream-2.0.0.tgz", - "integrity": "sha512-abv/qOcuPfk3URPfDzmZU1LKmuw8kT+0nIHvKrKgFrwifol/doWcdA4ZqsWQ8ENrFKkd67Mfpo/LovbIUsbt3w==", + "node_modules/memfs/node_modules/@jsonjoy.com/fs-node-builtins": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-node-builtins/-/fs-node-builtins-4.64.0.tgz", + "integrity": "sha512-/o7WRFhUWaM/fOrslwLZGnzn4RmRILykn+lAL+mNObqqRNw+CQSiij6hpCeZ+C7buhdoVo7go/OYqzaSUfDYmA==", + "dev": true, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-node-to-fsa": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-node-to-fsa/-/fs-node-to-fsa-4.64.0.tgz", + "integrity": "sha512-WDD9WVs0hb7UAEKTgZW2f66WDrbj7gIIWwpP3spbLyXa0rghtUaFTB8L4gdR3ZCWwiKIsj38/CNijpVmpnuPUw==", + "dev": true, + "dependencies": { + "@jsonjoy.com/fs-fsa": "4.64.0", + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-node-utils": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-node-utils/-/fs-node-utils-4.64.0.tgz", + "integrity": "sha512-k5Indsx9hWW9xSF7Y6oSKKwtCUNhzZxadub3owhIlitc+iMRVlPPdX2duTKQWBL3qNWpXya8jykgaaWpheeS4w==", + "dev": true, + "dependencies": { + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "glob-to-regex.js": "^1.0.1" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-print": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-print/-/fs-print-4.64.0.tgz", + "integrity": "sha512-PHZFccchvkhWrwPWHjmVAhbC3vSHCtyZvlZfJJ3ho2bnzl450hXri6/8e6pbkWdH+SkmLXNml0sV8e5HDAfxKw==", + "dev": true, + "dependencies": { + "@jsonjoy.com/fs-node-utils": "4.64.0", + "tree-dump": "^1.1.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-snapshot": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-snapshot/-/fs-snapshot-4.64.0.tgz", + "integrity": "sha512-oM7UDeL83q6NBzzsfKAsYKXKVXlykKFqqOLh4xZZKAzzROTlInkPbc6LTDGThEOnPiFiUzA7tYziHG9xavd76Q==", + "dev": true, + "dependencies": { + "@jsonjoy.com/buffers": "^17.65.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "@jsonjoy.com/json-pack": "^17.65.0", + "@jsonjoy.com/util": "^17.65.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-snapshot/node_modules/@jsonjoy.com/base64": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/base64/-/base64-17.67.0.tgz", + "integrity": "sha512-5SEsJGsm15aP8TQGkDfJvz9axgPwAEm98S5DxOuYe8e1EbfajcDmgeXXzccEjh+mLnjqEKrkBdjHWS5vFNwDdw==", + "dev": true, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-snapshot/node_modules/@jsonjoy.com/codegen": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/codegen/-/codegen-17.67.0.tgz", + "integrity": "sha512-idnkUplROpdBOV0HMcwhsCUS5TRUi9poagdGs70A6S4ux9+/aPuKbh8+UYRTLYQHtXvAdNfQWXDqZEx5k4Dj2Q==", + "dev": true, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-snapshot/node_modules/@jsonjoy.com/json-pack": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pack/-/json-pack-17.67.0.tgz", + "integrity": "sha512-t0ejURcGaZsn1ClbJ/3kFqSOjlryd92eQY465IYrezsXmPcfHPE/av4twRSxf6WE+TkZgLY+71vCZbiIiFKA/w==", + "dev": true, + "dependencies": { + "@jsonjoy.com/base64": "17.67.0", + "@jsonjoy.com/buffers": "17.67.0", + "@jsonjoy.com/codegen": "17.67.0", + "@jsonjoy.com/json-pointer": "17.67.0", + "@jsonjoy.com/util": "17.67.0", + "hyperdyperid": "^1.2.0", + "thingies": "^2.5.0", + "tree-dump": "^1.1.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-snapshot/node_modules/@jsonjoy.com/json-pointer": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pointer/-/json-pointer-17.67.0.tgz", + "integrity": "sha512-+iqOFInH+QZGmSuaybBUNdh7yvNrXvqR+h3wjXm0N/3JK1EyyFAeGJvqnmQL61d1ARLlk/wJdFKSL+LHJ1eaUA==", + "dev": true, + "dependencies": { + "@jsonjoy.com/util": "17.67.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/fs-snapshot/node_modules/@jsonjoy.com/util": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/util/-/util-17.67.0.tgz", + "integrity": "sha512-6+8xBaz1rLSohlGh68D1pdw3AwDi9xydm8QNlAFkvnavCJYSze+pxoW2VKP8p308jtlMRLs5NTHfPlZLd4w7ew==", + "dev": true, + "dependencies": { + "@jsonjoy.com/buffers": "17.67.0", + "@jsonjoy.com/codegen": "17.67.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/json-pack": { + "version": "1.21.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pack/-/json-pack-1.21.0.tgz", + "integrity": "sha512-+AKG+R2cfZMShzrF2uQw34v3zbeDYUqnQ+jg7ORic3BGtfw9p/+N6RJbq/kkV8JmYZaINknaEQ2m0/f693ZPpg==", + "dev": true, + "dependencies": { + "@jsonjoy.com/base64": "^1.1.2", + "@jsonjoy.com/buffers": "^1.2.0", + "@jsonjoy.com/codegen": "^1.0.0", + "@jsonjoy.com/json-pointer": "^1.0.2", + "@jsonjoy.com/util": "^1.9.0", + "hyperdyperid": "^1.2.0", + "thingies": "^2.5.0", + "tree-dump": "^1.1.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/json-pack/node_modules/@jsonjoy.com/buffers": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/buffers/-/buffers-1.2.1.tgz", + "integrity": "sha512-12cdlDwX4RUM3QxmUbVJWqZ/mrK6dFQH4Zxq6+r1YXKXYBNgZXndx2qbCJwh3+WWkCSn67IjnlG3XYTvmvYtgA==", + "dev": true, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/json-pointer": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pointer/-/json-pointer-1.0.2.tgz", + "integrity": "sha512-Fsn6wM2zlDzY1U+v4Nc8bo3bVqgfNTGcn6dMgs6FjrEnt4ZCe60o6ByKRjOGlI2gow0aE/Q41QOigdTqkyK5fg==", + "dev": true, + "dependencies": { + "@jsonjoy.com/codegen": "^1.0.0", + "@jsonjoy.com/util": "^1.9.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/util": { + "version": "1.9.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/util/-/util-1.9.0.tgz", + "integrity": "sha512-pLuQo+VPRnN8hfPqUTLTHk126wuYdXVxE6aDmjSeV4NCAgyxWbiOIeNJVtID3h1Vzpoi9m4jXezf73I6LgabgQ==", + "dev": true, + "dependencies": { + "@jsonjoy.com/buffers": "^1.0.0", + "@jsonjoy.com/codegen": "^1.0.0" + }, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/@jsonjoy.com/util/node_modules/@jsonjoy.com/buffers": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/buffers/-/buffers-1.2.1.tgz", + "integrity": "sha512-12cdlDwX4RUM3QxmUbVJWqZ/mrK6dFQH4Zxq6+r1YXKXYBNgZXndx2qbCJwh3+WWkCSn67IjnlG3XYTvmvYtgA==", + "dev": true, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/glob-to-regex.js": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/glob-to-regex.js/-/glob-to-regex.js-1.2.0.tgz", + "integrity": "sha512-QMwlOQKU/IzqMUOAZWubUOT8Qft+Y0KQWnX9nK3ch0CJg0tTp4TvGZsTfudYKv2NzoQSyPcnA6TYeIQ3jGichQ==", + "dev": true, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/thingies": { + "version": "2.6.0", + "resolved": "https://registry.npmjs.org/thingies/-/thingies-2.6.0.tgz", + "integrity": "sha512-rMHRjmlFLM1R96UYPvpmnc3LYtdFrT33JIB7L9hetGue1qAPfn1N2LJeEjxUSidu1Iku+haLZXDuEXUHNGO/lg==", + "dev": true, + "engines": { + "node": ">=10.18" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "^2" + } + }, + "node_modules/memfs/node_modules/tree-dump": { + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/tree-dump/-/tree-dump-1.1.0.tgz", + "integrity": "sha512-rMuvhU4MCDbcbnleZTFezWsaZXRFemSqAM+7jPnzUl1fo9w3YEKOxAeui0fz3OI4EU4hf23iyA7uQRVko+UaBA==", + "dev": true, + "engines": { + "node": ">=10.0" + }, + "funding": { + "type": "github", + "url": "https://github.com/sponsors/streamich" + }, + "peerDependencies": { + "tslib": "2" + } + }, + "node_modules/memfs/node_modules/tslib": { + "version": "2.8.1", + "resolved": "https://registry.npmjs.org/tslib/-/tslib-2.8.1.tgz", + "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", "dev": true }, - "node_modules/methods": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/methods/-/methods-1.1.2.tgz", - "integrity": "sha512-iclAHeNqNm68zFtnZ0e+1L2yUIdvzNoauKU4WBA3VvH/vPFieF7qfRlwUZU+DA9P9bPXIS90ulxoUoCH23sV2w==", + "node_modules/merge-descriptors": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/merge-descriptors/-/merge-descriptors-2.0.0.tgz", + "integrity": "sha512-Snk314V5ayFLhp3fkUREub6WtjBfPdCPY1Ln8/8munuLuiYhsABgBVWsozAG+MWMbVEvcdcpbi9R7ww22l9Q3g==", "dev": true, - "license": "MIT", "engines": { - "node": ">= 0.6" + "node": ">=18" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" } }, + "node_modules/merge-stream": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/merge-stream/-/merge-stream-2.0.0.tgz", + "integrity": "sha512-abv/qOcuPfk3URPfDzmZU1LKmuw8kT+0nIHvKrKgFrwifol/doWcdA4ZqsWQ8ENrFKkd67Mfpo/LovbIUsbt3w==", + "dev": true + }, "node_modules/micromatch": { "version": "4.0.8", "resolved": "https://registry.npmjs.org/micromatch/-/micromatch-4.0.8.tgz", "integrity": "sha512-PXwfBhYu0hBCPw8Dn0E+WDYb7af3dSLVWKi3HGv84IdF4TyFoC0ysxFd0Goxw7nSv4T/PzEJQxsYsEiFCKo2BA==", "dev": true, - "license": "MIT", "dependencies": { "braces": "^3.0.3", "picomatch": "^2.3.1" @@ -2655,19 +2679,6 @@ "node": ">=8.6" } }, - "node_modules/mime": { - "version": "1.6.0", - "resolved": "https://registry.npmjs.org/mime/-/mime-1.6.0.tgz", - "integrity": "sha512-x0Vn8spI+wuJ1O6S7gnbaQg8Pxh4NNHb7KSINmEWKiPE4RKOplvijn+NkmYmmRgP68mc70j2EbeTFRsrswaQeg==", - "dev": true, - "license": "MIT", - "bin": { - "mime": "cli.js" - }, - "engines": { - "node": ">=4" - } - }, "node_modules/mime-db": { "version": "1.52.0", "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.52.0.tgz", @@ -2689,12 +2700,6 @@ "node": ">= 0.6" } }, - "node_modules/minimalistic-assert": { - "version": "1.0.1", - "resolved": "https://registry.npmjs.org/minimalistic-assert/-/minimalistic-assert-1.0.1.tgz", - "integrity": "sha512-UtJcAD4yEaGtjPezWuO9wC4nwUnVH/8/Im3yEHQP4b67cXlD/Qr9hdITCU1xDbSEXg2XKNaP8jsReV7vQd00/A==", - "dev": true - }, "node_modules/ms": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz", @@ -2716,9 +2721,9 @@ } }, "node_modules/negotiator": { - "version": "0.6.3", - "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-0.6.3.tgz", - "integrity": "sha512-+EUsqGPLsM+j/zdChZjsnX51g4XrHFOIXwfnCVPGlQk/k5giakcKsuxCObBRu6DSm9opw/O6slWbJdghQM4bBg==", + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-1.0.0.tgz", + "integrity": "sha512-8Ofs/AUQh8MaEcrlq5xOX0CQ9ypTF5dl78mjlMNfOK08fzpgTHQRQPBxcPlEtIw0yRpws+Zo/3r+5WRby7u3Gg==", "dev": true, "engines": { "node": ">= 0.6" @@ -2750,7 +2755,6 @@ "resolved": "https://registry.npmjs.org/object-inspect/-/object-inspect-1.13.4.tgz", "integrity": "sha512-W67iLl4J2EXEGTbfeHCffrjDfitvLANg0UlX3wFUUSTx92KXRFegMHUVgSqE+wvhAbi4WqjGg9czysTV2Epbew==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.4" }, @@ -2758,18 +2762,11 @@ "url": "https://github.com/sponsors/ljharb" } }, - "node_modules/obuf": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/obuf/-/obuf-1.1.2.tgz", - "integrity": "sha512-PX1wu0AmAdPqOL1mWhqmlOd8kOIZQwGZw6rh7uby9fTc5lhaOWFLX3I6R1hrF9k3zUY40e6igsLGkDXK92LJNg==", - "dev": true - }, "node_modules/on-finished": { "version": "2.4.1", "resolved": "https://registry.npmjs.org/on-finished/-/on-finished-2.4.1.tgz", "integrity": "sha512-oVlzkg3ENAhCk2zdv7IJwd/QUD4z2RxRwpkcGY8psCVcCYZNq4wYnVWALHM+brtuJjePWiYF/ClmuDr8Ch5+kg==", "dev": true, - "license": "MIT", "dependencies": { "ee-first": "1.1.1" }, @@ -2786,20 +2783,30 @@ "node": ">= 0.8" } }, + "node_modules/once": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz", + "integrity": "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w==", + "dev": true, + "dependencies": { + "wrappy": "1" + } + }, "node_modules/open": { - "version": "10.1.2", - "resolved": "https://registry.npmjs.org/open/-/open-10.1.2.tgz", - "integrity": "sha512-cxN6aIDPz6rm8hbebcP7vrQNhvRcveZoJU72Y7vskh4oIm+BZwBECnx5nTmrlres1Qapvx27Qo1Auukpf8PKXw==", + "version": "11.0.0", + "resolved": "https://registry.npmjs.org/open/-/open-11.0.0.tgz", + "integrity": "sha512-smsWv2LzFjP03xmvFoJ331ss6h+jixfA4UUV/Bsiyuu4YJPfN+FIQGOIiv4w9/+MoHkfkJ22UIaQWRVFRfH6Vw==", "dev": true, - "license": "MIT", "dependencies": { - "default-browser": "^5.2.1", + "default-browser": "^5.4.0", "define-lazy-prop": "^3.0.0", + "is-in-ssh": "^1.0.0", "is-inside-container": "^1.0.0", - "is-wsl": "^3.1.0" + "powershell-utils": "^0.1.0", + "wsl-utils": "^0.3.0" }, "engines": { - "node": ">=18" + "node": ">=20" }, "funding": { "url": "https://github.com/sponsors/sindresorhus" @@ -2833,18 +2840,15 @@ } }, "node_modules/p-retry": { - "version": "6.2.1", - "resolved": "https://registry.npmjs.org/p-retry/-/p-retry-6.2.1.tgz", - "integrity": "sha512-hEt02O4hUct5wtwg4H4KcWgDdm+l1bOaEy/hWzd8xtXB9BqxTWBBhb+2ImAtH4Cv4rPjV76xN3Zumqk3k3AhhQ==", + "version": "8.0.0", + "resolved": "https://registry.npmjs.org/p-retry/-/p-retry-8.0.0.tgz", + "integrity": "sha512-kFVqH1HxOHp8LupNsOys7bSV09VYTRLxarH/mokO4Rqhk6wGi70E0jh4VzvVGXfEVNggHoHLAMWsQqHyU1Ey9A==", "dev": true, - "license": "MIT", "dependencies": { - "@types/retry": "0.12.2", - "is-network-error": "^1.0.0", - "retry": "^0.13.1" + "is-network-error": "^1.3.0" }, "engines": { - "node": ">=16.17" + "node": ">=22" }, "funding": { "url": "https://github.com/sponsors/sindresorhus" @@ -2893,11 +2897,14 @@ "dev": true }, "node_modules/path-to-regexp": { - "version": "0.1.12", - "resolved": "https://registry.npmjs.org/path-to-regexp/-/path-to-regexp-0.1.12.tgz", - "integrity": "sha512-RA1GjUVMnvYFxuqovrEqZoxxW5NUZqbwKtYz/Tt7nXerk0LbLblQmrsgdeOxV5SFHf0UDggjS/bSeOZwt1pmEQ==", + "version": "8.4.2", + "resolved": "https://registry.npmjs.org/path-to-regexp/-/path-to-regexp-8.4.2.tgz", + "integrity": "sha512-qRcuIdP69NPm4qbACK+aDogI5CBDMi1jKe0ry5rSQJz8JVLsC7jV8XpiJjGRLLol3N+R5ihGYcrPLTno6pAdBA==", "dev": true, - "license": "MIT" + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } }, "node_modules/picocolors": { "version": "1.1.1", @@ -2952,18 +2959,23 @@ "integrity": "sha512-oJFu94HQb+KVduSUQL7wnpmqnfmLsOA/nAh6b6EH0wCEoK0/mPeXU6c3wKDV83MkOuHPRHtSXKKU99IBazS/2w==", "dev": true }, - "node_modules/process-nextick-args": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/process-nextick-args/-/process-nextick-args-2.0.0.tgz", - "integrity": "sha512-MtEC1TqN0EU5nephaJ4rAtThHtC86dNN9qCuEhtshvpVBkAW5ZO7BASN9REnF9eoXGcRub+pFuKEpOHE+HbEMw==", - "dev": true + "node_modules/powershell-utils": { + "version": "0.1.0", + "resolved": "https://registry.npmjs.org/powershell-utils/-/powershell-utils-0.1.0.tgz", + "integrity": "sha512-dM0jVuXJPsDN6DvRpea484tCUaMiXWjuCn++HGTqUWzGDjv5tZkEZldAJ/UMlqRYGFrD/etByo4/xOuC/snX2A==", + "dev": true, + "engines": { + "node": ">=20" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } }, "node_modules/proxy-addr": { "version": "2.0.7", "resolved": "https://registry.npmjs.org/proxy-addr/-/proxy-addr-2.0.7.tgz", "integrity": "sha512-llQsMLSUDUPT44jdrU/O37qlnifitDP+ZwrmmZcoSKyLKvtZxpyV0n2/bD/N4tBAAZ/gJEdZU7KMraoK1+XYAg==", "dev": true, - "license": "MIT", "dependencies": { "forwarded": "0.2.0", "ipaddr.js": "1.9.1" @@ -2977,7 +2989,6 @@ "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-1.9.1.tgz", "integrity": "sha512-0KI/607xoxSToH7GjN1FfSbLoU0+btTicjsQSWQlh/hZykN8KpmMf7uYwPW3R+akZ6R/w18ZlXSHBYXiYUPO3g==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.10" } @@ -3007,12 +3018,13 @@ } }, "node_modules/qs": { - "version": "6.14.2", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.14.2.tgz", - "integrity": "sha512-V/yCWTTF7VJ9hIh18Ugr2zhJMP01MY7c5kh4J870L7imm6/DIzBsNLTXzMwUA3yZ5b/KBqLx8Kp3uRvd7xSe3Q==", + "version": "6.15.3", + "resolved": "https://registry.npmjs.org/qs/-/qs-6.15.3.tgz", + "integrity": "sha512-O9gl3zCl5h5blw1KGUzQKhA5oUXSl8rwUIM5o0S3nCXMliSvy5Dzx7/DJcI+SwgICv+IneSZwhBh1oSyEHA71A==", "dev": true, "dependencies": { - "side-channel": "^1.1.0" + "es-define-property": "^1.0.1", + "side-channel": "^1.1.1" }, "engines": { "node": ">=0.6" @@ -3022,100 +3034,44 @@ } }, "node_modules/range-parser": { - "version": "1.2.1", - "resolved": "https://registry.npmjs.org/range-parser/-/range-parser-1.2.1.tgz", - "integrity": "sha512-Hrgsx+orqoygnmhFbKaHE6c296J+HTAQXoxEF6gNupROmmGJRoyzfG3ccAveqCBrwr/2yxQ5BVd/GTl5agOwSg==", + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/range-parser/-/range-parser-1.3.0.tgz", + "integrity": "sha512-hek2mFQpPuI4E1BBKrSto+BU3e3x4xuarsbiwr3+lf7p44juvFMV0XFWQAP3xUyqXA4RrXLIoaSUGbSt056ZMw==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.6" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/raw-body": { - "version": "2.5.3", - "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.3.tgz", - "integrity": "sha512-s4VSOf6yN0rvbRZGxs8Om5CWj6seneMwK3oDb4lWDH0UPhWcxwOWw5+qk24bxq87szX1ydrwylIOp2uG1ojUpA==", + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-3.0.2.tgz", + "integrity": "sha512-K5zQjDllxWkf7Z5xJdV0/B0WTNqx6vxG70zJE4N0kBs4LovmEYWJzQGxC9bS9RAKu3bgM40lrd5zoLJ12MQ5BA==", "dev": true, "dependencies": { "bytes": "~3.1.2", "http-errors": "~2.0.1", - "iconv-lite": "~0.4.24", + "iconv-lite": "~0.7.0", "unpipe": "~1.0.0" }, "engines": { - "node": ">= 0.8" - } - }, - "node_modules/raw-body/node_modules/depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", - "dev": true, - "engines": { - "node": ">= 0.8" + "node": ">= 0.10" } }, - "node_modules/raw-body/node_modules/http-errors": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", - "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", + "node_modules/readdirp": { + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/readdirp/-/readdirp-5.0.0.tgz", + "integrity": "sha512-9u/XQ1pvrQtYyMpZe7DXKv2p5CNvyVwzUB6uhLAnQwHMSgKMBR62lc7AHljaeteeHXn11XTAaLLUVZYVZyuRBQ==", "dev": true, - "dependencies": { - "depd": "~2.0.0", - "inherits": "~2.0.4", - "setprototypeof": "~1.2.0", - "statuses": "~2.0.2", - "toidentifier": "~1.0.1" - }, "engines": { - "node": ">= 0.8" + "node": ">= 20.19.0" }, "funding": { - "type": "opencollective", - "url": "https://opencollective.com/express" - } - }, - "node_modules/raw-body/node_modules/inherits": { - "version": "2.0.4", - "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", - "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", - "dev": true - }, - "node_modules/raw-body/node_modules/statuses": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", - "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", - "dev": true, - "engines": { - "node": ">= 0.8" - } - }, - "node_modules/readable-stream": { - "version": "2.3.6", - "resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-2.3.6.tgz", - "integrity": "sha512-tQtKA9WIAhBF3+VLAseyMqZeBjW0AHJoxOtYqSUZNJxauErmLbVm2FW1y+J/YA9dUrAC39ITejlZWhVIwawkKw==", - "dev": true, - "dependencies": { - "core-util-is": "~1.0.0", - "inherits": "~2.0.3", - "isarray": "~1.0.0", - "process-nextick-args": "~2.0.0", - "safe-buffer": "~5.1.1", - "string_decoder": "~1.1.1", - "util-deprecate": "~1.0.1" - } - }, - "node_modules/readdirp": { - "version": "3.6.0", - "resolved": "https://registry.npmjs.org/readdirp/-/readdirp-3.6.0.tgz", - "integrity": "sha512-hOS089on8RduqdbhvQ5Z37A0ESjsqz6qnRcffsMU3495FuTdqSm+7bhJ29JvIOsBDEEnan5DPu9t3To9VRlMzA==", - "dev": true, - "license": "MIT", - "dependencies": { - "picomatch": "^2.2.1" - }, - "engines": { - "node": ">=8.10.0" + "type": "individual", + "url": "https://paulmillr.com/funding/" } }, "node_modules/rechoir": { @@ -3145,12 +3101,6 @@ "node": ">=0.10.0" } }, - "node_modules/requires-port": { - "version": "1.0.0", - "resolved": "https://registry.npmjs.org/requires-port/-/requires-port-1.0.0.tgz", - "integrity": "sha512-KigOCHcocU3XODJxsu8i/j8T9tzT4adHiecwORRQ0ZZFcp7ahwXuRU1m+yuO90C5ZUyGeGfocHDI14M3L3yDAQ==", - "dev": true - }, "node_modules/resolve": { "version": "1.22.6", "resolved": "https://registry.npmjs.org/resolve/-/resolve-1.22.6.tgz", @@ -3189,22 +3139,27 @@ "node": ">=8" } }, - "node_modules/retry": { - "version": "0.13.1", - "resolved": "https://registry.npmjs.org/retry/-/retry-0.13.1.tgz", - "integrity": "sha512-XQBQ3I8W1Cge0Seh+6gjj03LbmRFWuoszgK9ooCpwYIrhhoO80pfq4cUkU5DkknwfOfFteRwlZ56PYOGYyFWdg==", + "node_modules/router": { + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/router/-/router-2.2.0.tgz", + "integrity": "sha512-nLTrUKm2UyiL7rlhapu/Zl45FwNgkZGaCpZbIHajDYgwlJCOzLSk+cIPAnsEqV955GjILJnKbdQC1nVPz+gAYQ==", "dev": true, - "license": "MIT", + "dependencies": { + "debug": "^4.4.0", + "depd": "^2.0.0", + "is-promise": "^4.0.0", + "parseurl": "^1.3.3", + "path-to-regexp": "^8.0.0" + }, "engines": { - "node": ">= 4" + "node": ">= 18" } }, "node_modules/run-applescript": { - "version": "7.0.0", - "resolved": "https://registry.npmjs.org/run-applescript/-/run-applescript-7.0.0.tgz", - "integrity": "sha512-9by4Ij99JUr/MCFBUkDKLWK3G9HVXmabKz9U5MlIAIuvuzkiOicRYs8XJLxX+xahD+mLiiCYDqF9dKAgtzKP1A==", + "version": "7.1.0", + "resolved": "https://registry.npmjs.org/run-applescript/-/run-applescript-7.1.0.tgz", + "integrity": "sha512-DPe5pVFaAsinSaV6QjQ6gdiedWDcRCbUuiQfQa2wmWV7+xC9bGulGI8+TdRmoFkAPaBXk8CrAbnlY2ISniJ47Q==", "dev": true, - "license": "MIT", "engines": { "node": ">=18" }, @@ -3212,12 +3167,6 @@ "url": "https://github.com/sponsors/sindresorhus" } }, - "node_modules/safe-buffer": { - "version": "5.1.2", - "resolved": "https://registry.npmjs.org/safe-buffer/-/safe-buffer-5.1.2.tgz", - "integrity": "sha512-Gd2UZBJDkXlY7GbJxfsE8/nvKkUEU1G38c1siN6QP6a9PT9MmHB8GnpscSmMJSoF8LOIrt8ud/wPtojys4G6+g==", - "dev": true - }, "node_modules/safer-buffer": { "version": "2.1.2", "resolved": "https://registry.npmjs.org/safer-buffer/-/safer-buffer-2.1.2.tgz", @@ -3243,12 +3192,6 @@ "url": "https://opencollective.com/webpack" } }, - "node_modules/select-hose": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/select-hose/-/select-hose-2.0.0.tgz", - "integrity": "sha1-Yl2GWPhlr0Psliv8N2o3NZpJlMo=", - "dev": true - }, "node_modules/selfsigned": { "version": "5.5.0", "resolved": "https://registry.npmjs.org/selfsigned/-/selfsigned-5.5.0.tgz", @@ -3263,100 +3206,95 @@ } }, "node_modules/send": { - "version": "0.19.0", - "resolved": "https://registry.npmjs.org/send/-/send-0.19.0.tgz", - "integrity": "sha512-dW41u5VfLXu8SJh5bwRmyYUbAoSB3c9uQh6L8h/KtsFREPWpbX1lrljJo186Jc4nmci/sGUZ9a0a0J2zgfq2hw==", - "dev": true, - "license": "MIT", - "dependencies": { - "debug": "2.6.9", - "depd": "2.0.0", - "destroy": "1.2.0", - "encodeurl": "~1.0.2", - "escape-html": "~1.0.3", - "etag": "~1.8.1", - "fresh": "0.5.2", - "http-errors": "2.0.0", - "mime": "1.6.0", - "ms": "2.1.3", - "on-finished": "2.4.1", - "range-parser": "~1.2.1", - "statuses": "2.0.1" + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/send/-/send-1.2.1.tgz", + "integrity": "sha512-1gnZf7DFcoIcajTjTwjwuDjzuz4PPcY2StKPlsGAQ1+YH20IRVrBaXSWmdjowTJ6u8Rc01PoYOGHXfP1mYcZNQ==", + "dev": true, + "dependencies": { + "debug": "^4.4.3", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "etag": "^1.8.1", + "fresh": "^2.0.0", + "http-errors": "^2.0.1", + "mime-types": "^3.0.2", + "ms": "^2.1.3", + "on-finished": "^2.4.1", + "range-parser": "^1.2.1", + "statuses": "^2.0.2" }, "engines": { - "node": ">= 0.8.0" - } - }, - "node_modules/send/node_modules/debug": { - "version": "2.6.9", - "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", - "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", - "dev": true, - "license": "MIT", - "dependencies": { - "ms": "2.0.0" + "node": ">= 18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, - "node_modules/send/node_modules/debug/node_modules/ms": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz", - "integrity": "sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A==", - "dev": true, - "license": "MIT" - }, - "node_modules/send/node_modules/depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", + "node_modules/send/node_modules/mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", "dev": true, - "license": "MIT", "engines": { - "node": ">= 0.8" + "node": ">= 0.6" } }, - "node_modules/send/node_modules/encodeurl": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/encodeurl/-/encodeurl-1.0.2.tgz", - "integrity": "sha512-TPJXq8JqFaVYm2CWmPvnP2Iyo4ZSM7/QKcSmuMLDObfpH5fi7RUGmd/rTDf+rut/saiDiQEeVTNgAmJEdAOx0w==", + "node_modules/send/node_modules/mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", "dev": true, - "license": "MIT", + "dependencies": { + "mime-db": "^1.54.0" + }, "engines": { - "node": ">= 0.8" + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/send/node_modules/ms": { "version": "2.1.3", "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.3.tgz", "integrity": "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==", - "dev": true, - "license": "MIT" - }, - "node_modules/send/node_modules/statuses": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", - "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.8" - } + "dev": true }, "node_modules/serve-index": { - "version": "1.9.1", - "resolved": "https://registry.npmjs.org/serve-index/-/serve-index-1.9.1.tgz", - "integrity": "sha1-03aNabHn2C5c4FD/9bRTvqEqkjk=", + "version": "1.9.2", + "resolved": "https://registry.npmjs.org/serve-index/-/serve-index-1.9.2.tgz", + "integrity": "sha512-KDj11HScOaLmrPxl70KYNW1PksP4Nb/CLL2yvC+Qd2kHMPEEpfc4Re2e4FOay+bC/+XQl/7zAcWON3JVo5v3KQ==", "dev": true, "dependencies": { - "accepts": "~1.3.4", + "accepts": "~1.3.8", "batch": "0.6.1", "debug": "2.6.9", "escape-html": "~1.0.3", - "http-errors": "~1.6.2", - "mime-types": "~2.1.17", - "parseurl": "~1.3.2" + "http-errors": "~1.8.0", + "mime-types": "~2.1.35", + "parseurl": "~1.3.3" }, "engines": { "node": ">= 0.8.0" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/serve-index/node_modules/accepts": { + "version": "1.3.8", + "resolved": "https://registry.npmjs.org/accepts/-/accepts-1.3.8.tgz", + "integrity": "sha512-PYAthTa2m2VKxuvSD3DPC/Gy+U+sOA1LAuT8mkmRuvw+NACSaeXEQ+NHcVF7rONl6qcaxV3Uuemwawk+7+SJLw==", + "dev": true, + "dependencies": { + "mime-types": "~2.1.34", + "negotiator": "0.6.3" + }, + "engines": { + "node": ">= 0.6" } }, "node_modules/serve-index/node_modules/debug": { @@ -3368,49 +3306,73 @@ "ms": "2.0.0" } }, + "node_modules/serve-index/node_modules/depd": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/depd/-/depd-1.1.2.tgz", + "integrity": "sha512-7emPTl6Dpo6JRXOXjLRxck+FlLRX5847cLKEn00PLAgc3g2hTZZgr+e4c2v6QpSmLeFP3n5yUo7ft6avBK/5jQ==", + "dev": true, + "engines": { + "node": ">= 0.6" + } + }, "node_modules/serve-index/node_modules/http-errors": { - "version": "1.6.3", - "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-1.6.3.tgz", - "integrity": "sha1-i1VoC7S+KDoLW/TqLjhYC+HZMg0=", + "version": "1.8.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-1.8.1.tgz", + "integrity": "sha512-Kpk9Sm7NmI+RHhnj6OIWDI1d6fIoFAtFt9RLaTMRlg/8w49juAStsrBgp0Dp4OdxdVbRIeKhtCUvoi/RuAhO4g==", "dev": true, "dependencies": { "depd": "~1.1.2", - "inherits": "2.0.3", - "setprototypeof": "1.1.0", - "statuses": ">= 1.4.0 < 2" + "inherits": "2.0.4", + "setprototypeof": "1.2.0", + "statuses": ">= 1.5.0 < 2", + "toidentifier": "1.0.1" }, "engines": { "node": ">= 0.6" } }, - "node_modules/serve-index/node_modules/setprototypeof": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/setprototypeof/-/setprototypeof-1.1.0.tgz", - "integrity": "sha512-BvE/TwpZX4FXExxOxZyRGQQv651MSwmWKZGqvmPcRIjDqWub67kTKuIMx43cZZrS/cBBzwBcNDWoFxt2XEFIpQ==", - "dev": true + "node_modules/serve-index/node_modules/negotiator": { + "version": "0.6.3", + "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-0.6.3.tgz", + "integrity": "sha512-+EUsqGPLsM+j/zdChZjsnX51g4XrHFOIXwfnCVPGlQk/k5giakcKsuxCObBRu6DSm9opw/O6slWbJdghQM4bBg==", + "dev": true, + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/serve-index/node_modules/statuses": { + "version": "1.5.0", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-1.5.0.tgz", + "integrity": "sha512-OpZ3zP+jT1PI7I8nemJX4AKmAX070ZkYPVWV/AaKTJl+tXCTGyVdC1a4SL8RUQYEwk/f34ZX8UTykN68FwrqAA==", + "dev": true, + "engines": { + "node": ">= 0.6" + } }, "node_modules/serve-static": { - "version": "1.16.2", - "resolved": "https://registry.npmjs.org/serve-static/-/serve-static-1.16.2.tgz", - "integrity": "sha512-VqpjJZKadQB/PEbEwvFdO43Ax5dFBZ2UECszz8bQ7pi7wt//PWe1P6MN7eCnjsatYtBT6EuiClbjSWP2WrIoTw==", + "version": "2.2.1", + "resolved": "https://registry.npmjs.org/serve-static/-/serve-static-2.2.1.tgz", + "integrity": "sha512-xRXBn0pPqQTVQiC8wyQrKs2MOlX24zQ0POGaj0kultvoOCstBQM5yvOhAVSUwOMjQtTvsPWoNCHfPGwaaQJhTw==", "dev": true, - "license": "MIT", "dependencies": { - "encodeurl": "~2.0.0", - "escape-html": "~1.0.3", - "parseurl": "~1.3.3", - "send": "0.19.0" + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "parseurl": "^1.3.3", + "send": "^1.2.0" }, "engines": { - "node": ">= 0.8.0" + "node": ">= 18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" } }, "node_modules/setprototypeof": { "version": "1.2.0", "resolved": "https://registry.npmjs.org/setprototypeof/-/setprototypeof-1.2.0.tgz", "integrity": "sha512-E5LDX7Wrp85Kil5bhZv46j8jOeboKq5JMmYM3gVGdGH8xFpPWXUMsNrlODCrkoxMEeNi/XZIwuRvY4XNwYMJpw==", - "dev": true, - "license": "ISC" + "dev": true }, "node_modules/shallow-clone": { "version": "3.0.1", @@ -3446,11 +3408,10 @@ } }, "node_modules/shell-quote": { - "version": "1.8.3", - "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.8.3.tgz", - "integrity": "sha512-ObmnIF4hXNg1BqhnHmgbDETF8dLPCggZWBjkQfhZpbszZnYur5DUljTcCHii5LC3J5E0yeO/1LIMyH+UvHQgyw==", + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.10.0.tgz", + "integrity": "sha512-w1aiOKwKuRgtwAReIIj89puqg+I7GvX4IbLrvmhXbzQsj1+Zwi4VO3+fa6ZF91TWSjIxoEkKnMeHcLEODK5ZXA==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.4" }, @@ -3459,15 +3420,14 @@ } }, "node_modules/side-channel": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/side-channel/-/side-channel-1.1.0.tgz", - "integrity": "sha512-ZX99e6tRweoUXqR+VBrslhda51Nh5MTQwou5tnUDgbtyM0dBgmhEDtWGP/xbKn6hqfPRHujUNwz5fy/wbbhnpw==", + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/side-channel/-/side-channel-1.1.1.tgz", + "integrity": "sha512-6x6dK6zJdpTzF4sQeNYxwtvBzf6Eg4GtlesS94HOvTudUeyK2WXAaIfmDgsyslYrRBeFIlsi54AYsFGUuhmvrQ==", "dev": true, - "license": "MIT", "dependencies": { "es-errors": "^1.3.0", - "object-inspect": "^1.13.3", - "side-channel-list": "^1.0.0", + "object-inspect": "^1.13.4", + "side-channel-list": "^1.0.1", "side-channel-map": "^1.0.1", "side-channel-weakmap": "^1.0.2" }, @@ -3479,14 +3439,13 @@ } }, "node_modules/side-channel-list": { - "version": "1.0.0", - "resolved": "https://registry.npmjs.org/side-channel-list/-/side-channel-list-1.0.0.tgz", - "integrity": "sha512-FCLHtRD/gnpCiCHEiJLOwdmFP+wzCmDEkc9y7NsYxeF4u7Btsn1ZuwgwJGxImImHicJArLP4R0yX4c2KCrMrTA==", + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/side-channel-list/-/side-channel-list-1.0.1.tgz", + "integrity": "sha512-mjn/0bi/oUURjc5Xl7IaWi/OJJJumuoJFQJfDDyO46+hBWsfaVM65TBHq2eoZBhzl9EchxOijpkbRC8SVBQU0w==", "dev": true, - "license": "MIT", "dependencies": { "es-errors": "^1.3.0", - "object-inspect": "^1.13.3" + "object-inspect": "^1.13.4" }, "engines": { "node": ">= 0.4" @@ -3500,7 +3459,6 @@ "resolved": "https://registry.npmjs.org/side-channel-map/-/side-channel-map-1.0.1.tgz", "integrity": "sha512-VCjCNfgMsby3tTdo02nbjtM/ewra6jPHmpThenkTYh8pG9ucZ/1P8So4u4FGBek/BjpOVsDCMoLA/iuBKIFXRA==", "dev": true, - "license": "MIT", "dependencies": { "call-bound": "^1.0.2", "es-errors": "^1.3.0", @@ -3519,7 +3477,6 @@ "resolved": "https://registry.npmjs.org/side-channel-weakmap/-/side-channel-weakmap-1.0.2.tgz", "integrity": "sha512-WPS/HvHQTYnHisLo9McqBHOJk2FkHO/tlpvldyrnem4aeQp4hai3gythswg6p01oSoTl58rcpiFAjF2br2Ak2A==", "dev": true, - "license": "MIT", "dependencies": { "call-bound": "^1.0.2", "es-errors": "^1.3.0", @@ -3534,17 +3491,6 @@ "url": "https://github.com/sponsors/ljharb" } }, - "node_modules/sockjs": { - "version": "0.3.24", - "resolved": "https://registry.npmjs.org/sockjs/-/sockjs-0.3.24.tgz", - "integrity": "sha512-GJgLTZ7vYb/JtPSSZ10hsOYIvEYsjbNU+zPdIHcUaWVNUEPivzxku31865sSSud0Da0W4lEeOPlmw93zLQchuQ==", - "dev": true, - "dependencies": { - "faye-websocket": "^0.11.3", - "uuid": "^8.3.2", - "websocket-driver": "^0.7.4" - } - }, "node_modules/source-map": { "version": "0.6.1", "resolved": "https://registry.npmjs.org/source-map/-/source-map-0.6.1.tgz", @@ -3564,66 +3510,13 @@ "source-map": "^0.6.0" } }, - "node_modules/spdy": { - "version": "4.0.2", - "resolved": "https://registry.npmjs.org/spdy/-/spdy-4.0.2.tgz", - "integrity": "sha512-r46gZQZQV+Kl9oItvl1JZZqJKGr+oEkB08A6BzkiR7593/7IbtuncXHd2YoYeTsG4157ZssMu9KYvUHLcjcDoA==", - "dev": true, - "dependencies": { - "debug": "^4.1.0", - "handle-thing": "^2.0.0", - "http-deceiver": "^1.2.7", - "select-hose": "^2.0.0", - "spdy-transport": "^3.0.0" - }, - "engines": { - "node": ">=6.0.0" - } - }, - "node_modules/spdy-transport": { - "version": "3.0.0", - "resolved": "https://registry.npmjs.org/spdy-transport/-/spdy-transport-3.0.0.tgz", - "integrity": "sha512-hsLVFE5SjA6TCisWeJXFKniGGOpBgMLmerfO2aCyCU5s7nJ/rpAepqmFifv/GCbSbueEeAJJnmSQ2rKC/g8Fcw==", - "dev": true, - "dependencies": { - "debug": "^4.1.0", - "detect-node": "^2.0.4", - "hpack.js": "^2.1.6", - "obuf": "^1.1.2", - "readable-stream": "^3.0.6", - "wbuf": "^1.7.3" - } - }, - "node_modules/spdy-transport/node_modules/readable-stream": { - "version": "3.6.0", - "resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-3.6.0.tgz", - "integrity": "sha512-BViHy7LKeTz4oNnkcLJ+lVSL6vpiFeX6/d3oSH8zCW7UxP2onchk+vTGB143xuFjHS3deTgkKoXXymXqymiIdA==", - "dev": true, - "dependencies": { - "inherits": "^2.0.3", - "string_decoder": "^1.1.1", - "util-deprecate": "^1.0.1" - }, - "engines": { - "node": ">= 6" - } - }, "node_modules/statuses": { - "version": "1.5.0", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-1.5.0.tgz", - "integrity": "sha1-Fhx9rBd2Wf2YEfQ3cfqZOBR4Yow=", + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", + "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", "dev": true, "engines": { - "node": ">= 0.6" - } - }, - "node_modules/string_decoder": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/string_decoder/-/string_decoder-1.1.1.tgz", - "integrity": "sha512-n/ShnvDi6FHbbVfviro+WojiFzv+s8MPMHBczVePfUpDJLwoLT0ht1l4YwBCbi8pJAveEEdnkHyPyTP/mzRfwg==", - "dev": true, - "dependencies": { - "safe-buffer": "~5.1.0" + "node": ">= 0.8" } }, "node_modules/supports-color": { @@ -3786,7 +3679,6 @@ "resolved": "https://registry.npmjs.org/toidentifier/-/toidentifier-1.0.1.tgz", "integrity": "sha512-o5sSPKEkg/DIQNmH43V0/uerLrpzVedkUh8tGNvaeXpfpuwjKenlSox/2O/BTlZUtEe+JG7s5YhEz608PlAHRA==", "dev": true, - "license": "MIT", "engines": { "node": ">=0.6" } @@ -3810,25 +3702,66 @@ } }, "node_modules/type-is": { - "version": "1.6.18", - "resolved": "https://registry.npmjs.org/type-is/-/type-is-1.6.18.tgz", - "integrity": "sha512-TkRKr9sUTxEH8MdfuCSP7VizJyzRNMjj2J2do2Jr3Kym598JVdEksuzPQCnlFPW4ky9Q+iA+ma9BGm06XQBy8g==", + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/type-is/-/type-is-2.1.0.tgz", + "integrity": "sha512-faYHw0anBbc/kWF3zFTEnxSFOAGUX9GFbOBthvDdLsIlEoWOFOtS0zgCiQYwIskL9iGXZL3kAXD8OoZ4GmMATA==", "dev": true, - "license": "MIT", "dependencies": { - "media-typer": "0.3.0", - "mime-types": "~2.1.24" + "content-type": "^2.0.0", + "media-typer": "^1.1.0", + "mime-types": "^3.0.0" + }, + "engines": { + "node": ">= 18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/type-is/node_modules/content-type": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-2.0.0.tgz", + "integrity": "sha512-j/O/d7GcZCyNl7/hwZAb606rzqkyvaDctLmckbxLzHvFBzTJHuGEdodATcP3yIRoDrLHkIATJuvzbFlp/ki2cQ==", + "dev": true, + "engines": { + "node": ">=18" }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, + "node_modules/type-is/node_modules/mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true, "engines": { "node": ">= 0.6" } }, + "node_modules/type-is/node_modules/mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", + "dev": true, + "dependencies": { + "mime-db": "^1.54.0" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, "node_modules/unpipe": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/unpipe/-/unpipe-1.0.0.tgz", "integrity": "sha512-pjy2bYhSsufwWlKwPc+l3cN7+wuJlK6uz0YdJEOlQDbl6jo/YlPi4mb8agUkVC8BF7V8NuzeyPNqRksA3hztKQ==", "dev": true, - "license": "MIT", "engines": { "node": ">= 0.8" } @@ -3863,31 +3796,6 @@ "browserslist": ">= 4.21.0" } }, - "node_modules/util-deprecate": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz", - "integrity": "sha1-RQ1Nyfpw3nMnYvvS1KKJgUGaDM8=", - "dev": true - }, - "node_modules/utils-merge": { - "version": "1.0.1", - "resolved": "https://registry.npmjs.org/utils-merge/-/utils-merge-1.0.1.tgz", - "integrity": "sha512-pMZTvIkT1d+TFGvDOqodOclx0QWkkgi6Tdoa8gC8ffGAAqz9pzPTZWAybbsHHoED/ztMtkv/VoYTYyShUn81hA==", - "dev": true, - "license": "MIT", - "engines": { - "node": ">= 0.4.0" - } - }, - "node_modules/uuid": { - "version": "8.3.2", - "resolved": "https://registry.npmjs.org/uuid/-/uuid-8.3.2.tgz", - "integrity": "sha512-+NYs2QeMWy+GWFOEm9xnn6HCDp0l7QBD7ml8zLUmJ+93Q5NF0NocErnwkTkXVFNiX3/fpC6afS8Dhb/gz7R7eg==", - "dev": true, - "bin": { - "uuid": "dist/bin/uuid" - } - }, "node_modules/vary": { "version": "1.1.2", "resolved": "https://registry.npmjs.org/vary/-/vary-1.1.2.tgz", @@ -3910,15 +3818,6 @@ "node": ">=10.13.0" } }, - "node_modules/wbuf": { - "version": "1.7.3", - "resolved": "https://registry.npmjs.org/wbuf/-/wbuf-1.7.3.tgz", - "integrity": "sha512-O84QOnr0icsbFGLS0O3bI5FswxzRr8/gHwWkDlQFskhSPryQXvrTMxjxGP4+iWYoauLoBvfDpkrOauZ+0iZpDA==", - "dev": true, - "dependencies": { - "minimalistic-assert": "^1.0.0" - } - }, "node_modules/webpack": { "version": "5.105.0", "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.105.0.tgz", @@ -4022,28 +3921,25 @@ } }, "node_modules/webpack-dev-middleware": { - "version": "7.4.2", - "resolved": "https://registry.npmjs.org/webpack-dev-middleware/-/webpack-dev-middleware-7.4.2.tgz", - "integrity": "sha512-xOO8n6eggxnwYpy1NlzUKpvrjfJTvae5/D6WOK0S2LSo7vjmo5gCM1DbLUmFqrMTJP+W/0YZNctm7jasWvLuBA==", + "version": "8.0.4", + "resolved": "https://registry.npmjs.org/webpack-dev-middleware/-/webpack-dev-middleware-8.0.4.tgz", + "integrity": "sha512-9dFzIvIfbdnkOlRjXDHEmEKlY/KPsELNIyKWdoNfK4WaHN9Db+JyVG0gi4/APUPX2UVhnCZ6jp7x0EyM7yTq1Q==", "dev": true, - "license": "MIT", "dependencies": { - "colorette": "^2.0.10", - "memfs": "^4.6.0", - "mime-types": "^2.1.31", - "on-finished": "^2.4.1", + "memfs": "^4.56.10", + "mime-types": "^3.0.2", "range-parser": "^1.2.1", - "schema-utils": "^4.0.0" + "schema-utils": "^4.3.3" }, "engines": { - "node": ">= 18.12.0" + "node": ">= 20.9.0" }, "funding": { "type": "opencollective", "url": "https://opencollective.com/webpack" }, "peerDependencies": { - "webpack": "^5.0.0" + "webpack": "^5.101.0" }, "peerDependenciesMeta": { "webpack": { @@ -4051,53 +3947,75 @@ } } }, + "node_modules/webpack-dev-middleware/node_modules/mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true, + "engines": { + "node": ">= 0.6" + } + }, + "node_modules/webpack-dev-middleware/node_modules/mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", + "dev": true, + "dependencies": { + "mime-db": "^1.54.0" + }, + "engines": { + "node": ">=18" + }, + "funding": { + "type": "opencollective", + "url": "https://opencollective.com/express" + } + }, "node_modules/webpack-dev-server": { - "version": "5.2.4", - "resolved": "https://registry.npmjs.org/webpack-dev-server/-/webpack-dev-server-5.2.4.tgz", - "integrity": "sha512-GqDPGZN9bRqKBTkp4aWkobDDHMsrXKoGSdOH56smIri8qR0JG8gfL8/v/f/OZR3/OKXjG8uwJbFVhKm/FNU/UA==", + "version": "6.0.0", + "resolved": "https://registry.npmjs.org/webpack-dev-server/-/webpack-dev-server-6.0.0.tgz", + "integrity": "sha512-q9SD4ItOGhZLeU6EGT10caDZdHjF50Pz1DtkRZZOPsfluMXOkacWKKOtSBSLVkPqKiF67eFUC0rI88U/tSFPEw==", "dev": true, "dependencies": { "@types/bonjour": "^3.5.13", "@types/connect-history-api-fallback": "^1.5.4", - "@types/express": "^4.17.25", - "@types/express-serve-static-core": "^4.17.21", + "@types/express": "^5.0.6", + "@types/express-serve-static-core": "^5.1.1", "@types/serve-index": "^1.9.4", - "@types/serve-static": "^1.15.5", - "@types/sockjs": "^0.3.36", - "@types/ws": "^8.5.10", + "@types/serve-static": "^2.2.0", + "@types/ws": "^8.18.1", "ansi-html-community": "^0.0.8", - "bonjour-service": "^1.2.1", - "chokidar": "^3.6.0", - "colorette": "^2.0.10", + "bonjour-service": "^1.3.0", + "chokidar": "^5.0.0", "compression": "^1.8.1", "connect-history-api-fallback": "^2.0.0", - "express": "^4.22.1", - "graceful-fs": "^4.2.6", - "http-proxy-middleware": "^2.0.9", - "ipaddr.js": "^2.1.0", - "launch-editor": "^2.6.1", - "open": "^10.0.3", - "p-retry": "^6.2.0", - "schema-utils": "^4.2.0", + "express": "^5.2.1", + "graceful-fs": "^4.2.11", + "http-proxy-middleware": "^4.1.1", + "ipaddr.js": "^2.3.0", + "launch-editor": "^2.14.1", + "open": "^11.0.0", + "p-retry": "^8.0.0", + "schema-utils": "^4.3.3", "selfsigned": "^5.5.0", - "serve-index": "^1.9.1", - "sockjs": "^0.3.24", - "spdy": "^4.0.2", - "webpack-dev-middleware": "^7.4.2", - "ws": "^8.18.0" + "serve-index": "^1.9.2", + "tinyglobby": "^0.2.15", + "webpack-dev-middleware": "^8.0.3", + "ws": "^8.20.0" }, "bin": { "webpack-dev-server": "bin/webpack-dev-server.js" }, "engines": { - "node": ">= 18.12.0" + "node": ">= 22.15.0" }, "funding": { "type": "opencollective", "url": "https://opencollective.com/webpack" }, "peerDependencies": { - "webpack": "^5.0.0" + "webpack": "^5.101.0" }, "peerDependenciesMeta": { "webpack": { @@ -4130,29 +4048,6 @@ "node": ">=10.13.0" } }, - "node_modules/websocket-driver": { - "version": "0.7.4", - "resolved": "https://registry.npmjs.org/websocket-driver/-/websocket-driver-0.7.4.tgz", - "integrity": "sha512-b17KeDIQVjvb0ssuSDF2cYXSg2iztliJ4B9WdsuB6J952qCPKmnVq4DyW5motImXHDC1cBT/1UezrJVsKw5zjg==", - "dev": true, - "dependencies": { - "http-parser-js": ">=0.5.1", - "safe-buffer": ">=5.1.0", - "websocket-extensions": ">=0.1.1" - }, - "engines": { - "node": ">=0.8.0" - } - }, - "node_modules/websocket-extensions": { - "version": "0.1.4", - "resolved": "https://registry.npmjs.org/websocket-extensions/-/websocket-extensions-0.1.4.tgz", - "integrity": "sha512-OqedPIGOfsDlo31UNwYbCFMSaO9m9G/0faIHj5/dZFDMFqPTcx6UwqyOy3COEaEOg/9VsGIpdqn62W5KhoKSpg==", - "dev": true, - "engines": { - "node": ">=0.8.0" - } - }, "node_modules/which": { "version": "2.0.2", "resolved": "https://registry.npmjs.org/which/-/which-2.0.2.tgz", @@ -4174,12 +4069,17 @@ "integrity": "sha512-CC1bOL87PIWSBhDcTrdeLo6eGT7mCFtrg0uIJtqJUFyK+eJnzl8A1niH56uu7KMa5XFrtiV+AQuHO3n7DsHnLQ==", "dev": true }, + "node_modules/wrappy": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz", + "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==", + "dev": true + }, "node_modules/ws": { - "version": "8.18.2", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.18.2.tgz", - "integrity": "sha512-DMricUmwGZUVr++AEAe2uiVM7UoO9MAVZMDu05UQOaUII0lp+zOzLLU4Xqh/JvTqklB1T4uELaaPBKyjE1r4fQ==", + "version": "8.21.1", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.1.tgz", + "integrity": "sha512-+0NTnW77fFN/DjQi6k/Sq/Yvk4Sgajw7urW8V+asjXnRgDs9gyGkdb7EzgfhA4goXsRIZKE28fzIXBHEzhuiWw==", "dev": true, - "license": "MIT", "engines": { "node": ">=10.0.0" }, @@ -4195,6 +4095,22 @@ "optional": true } } + }, + "node_modules/wsl-utils": { + "version": "0.3.1", + "resolved": "https://registry.npmjs.org/wsl-utils/-/wsl-utils-0.3.1.tgz", + "integrity": "sha512-g/eziiSUNBSsdDJtCLB8bdYEUMj4jR7AGeUo96p/3dTafgjHhpF4RiCFPiRILwjQoDXx5MqkBr4fwWtR3Ky4Wg==", + "dev": true, + "dependencies": { + "is-wsl": "^3.1.0", + "powershell-utils": "^0.1.0" + }, + "engines": { + "node": ">=20" + }, + "funding": { + "url": "https://github.com/sponsors/sindresorhus" + } } }, "dependencies": { @@ -4573,21 +4489,20 @@ "dev": true }, "@types/express": { - "version": "4.17.25", - "resolved": "https://registry.npmjs.org/@types/express/-/express-4.17.25.tgz", - "integrity": "sha512-dVd04UKsfpINUnK0yBoYHDF3xu7xVH4BuDotC/xGuycx4CgbP48X/KF/586bcObxT0HENHXEU8Nqtu6NR+eKhw==", + "version": "5.0.6", + "resolved": "https://registry.npmjs.org/@types/express/-/express-5.0.6.tgz", + "integrity": "sha512-sKYVuV7Sv9fbPIt/442koC7+IIwK5olP1KWeD88e/idgoJqDm3JV/YUiPwkoKK92ylff2MGxSz1CSjsXelx0YA==", "dev": true, "requires": { "@types/body-parser": "*", - "@types/express-serve-static-core": "^4.17.33", - "@types/qs": "*", - "@types/serve-static": "^1" + "@types/express-serve-static-core": "^5.0.0", + "@types/serve-static": "^2" } }, "@types/express-serve-static-core": { - "version": "4.17.36", - "resolved": "https://registry.npmjs.org/@types/express-serve-static-core/-/express-serve-static-core-4.17.36.tgz", - "integrity": "sha512-zbivROJ0ZqLAtMzgzIUC4oNqDG9iF0lSsAqpOD9kbs5xcIM3dTiyuHvBc7R8MtWBp3AAWGaovJa+wzWPjLYW7Q==", + "version": "5.1.2", + "resolved": "https://registry.npmjs.org/@types/express-serve-static-core/-/express-serve-static-core-5.1.2.tgz", + "integrity": "sha512-d3KvEXBSo/lOAMc2u6fkyDHBvetBHeqD7wm/AcXfLpSOQwlmG9D/aQ0SFswVjv05p7ullQS7Mjohj6/VdbZuTg==", "dev": true, "requires": { "@types/node": "*", @@ -4597,20 +4512,11 @@ } }, "@types/http-errors": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/@types/http-errors/-/http-errors-2.0.2.tgz", - "integrity": "sha512-lPG6KlZs88gef6aD85z3HNkztpj7w2R7HmR3gygjfXCQmsLloWNARFkMuzKiiY8FGdh1XDpgBdrSf4aKDiA7Kg==", + "version": "2.0.5", + "resolved": "https://registry.npmjs.org/@types/http-errors/-/http-errors-2.0.5.tgz", + "integrity": "sha512-r8Tayk8HJnX0FztbZN7oVqGccWgw98T/0neJphO91KkmOzug1KkofZURD4UaD5uH8AqcFLfdPErnBod0u71/qg==", "dev": true }, - "@types/http-proxy": { - "version": "1.17.12", - "resolved": "https://registry.npmjs.org/@types/http-proxy/-/http-proxy-1.17.12.tgz", - "integrity": "sha512-kQtujO08dVtQ2wXAuSFfk9ASy3sug4+ogFR8Kd8UgP8PEuc1/G/8yjYRmp//PcDNJEUKOza/MrQu15bouEUCiw==", - "dev": true, - "requires": { - "@types/node": "*" - } - }, "@types/json-schema": { "version": "7.0.15", "resolved": "https://registry.npmjs.org/@types/json-schema/-/json-schema-7.0.15.tgz", @@ -4641,12 +4547,6 @@ "integrity": "sha512-EEhsLsD6UsDM1yFhAvy0Cjr6VwmpMWqFBCb9w07wVugF7w9nfajxLuVmngTIpgS6svCnm6Vaw+MZhoDCKnOfsw==", "dev": true }, - "@types/retry": { - "version": "0.12.2", - "resolved": "https://registry.npmjs.org/@types/retry/-/retry-0.12.2.tgz", - "integrity": "sha512-XISRgDJ2Tc5q4TRqvgJtzsRkFYNJzZrhTdtMoGVBttwzzQJkPnS3WWTFc7kuDRoPtPakl+T+OfdEUjYJj7Jbow==", - "dev": true - }, "@types/send": { "version": "0.17.1", "resolved": "https://registry.npmjs.org/@types/send/-/send-0.17.1.tgz", @@ -4667,22 +4567,12 @@ } }, "@types/serve-static": { - "version": "1.15.7", - "resolved": "https://registry.npmjs.org/@types/serve-static/-/serve-static-1.15.7.tgz", - "integrity": "sha512-W8Ym+h8nhuRwaKPaDw34QUkwsGi6Rc4yYqvKFo5rm2FUEhCFbzVWrxXUxuKK8TASjWsysJY0nsmNCGhCOIsrOw==", + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/@types/serve-static/-/serve-static-2.2.0.tgz", + "integrity": "sha512-8mam4H1NHLtu7nmtalF7eyBH14QyOASmcxHhSfEoRyr0nP/YdoesEtU+uSRvMe96TW/HPTtkoKqQLl53N7UXMQ==", "dev": true, "requires": { "@types/http-errors": "*", - "@types/node": "*", - "@types/send": "*" - } - }, - "@types/sockjs": { - "version": "0.3.36", - "resolved": "https://registry.npmjs.org/@types/sockjs/-/sockjs-0.3.36.tgz", - "integrity": "sha512-MK9V6NzAS1+Ud7JV9lJLFqW85VbC9dq3LmwZCuBe4wBDgKC0Kj/jd8Xl+nSviU+Qc3+m7umHHyHg//2KSa0a0Q==", - "dev": true, - "requires": { "@types/node": "*" } }, @@ -4875,13 +4765,30 @@ "dev": true }, "accepts": { - "version": "1.3.8", - "resolved": "https://registry.npmjs.org/accepts/-/accepts-1.3.8.tgz", - "integrity": "sha512-PYAthTa2m2VKxuvSD3DPC/Gy+U+sOA1LAuT8mkmRuvw+NACSaeXEQ+NHcVF7rONl6qcaxV3Uuemwawk+7+SJLw==", + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/accepts/-/accepts-2.0.0.tgz", + "integrity": "sha512-5cvg6CtKwfgdmVqY1WIiXKc3Q1bkRqGLi+2W/6ao+6Y7gu/RCwRuAhGEzh5B4KlszSuTLgZYuqFqo5bImjNKng==", "dev": true, "requires": { - "mime-types": "~2.1.34", - "negotiator": "0.6.3" + "mime-types": "^3.0.0", + "negotiator": "^1.0.0" + }, + "dependencies": { + "mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true + }, + "mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", + "dev": true, + "requires": { + "mime-db": "^1.54.0" + } + } } }, "acorn": { @@ -4933,22 +4840,6 @@ "integrity": "sha512-1APHAyr3+PCamwNw3bXCPp4HFLONZt/yIH0sZp0/469KWNTEy+qN5jQ3GVX6DMZ1UXAi34yVwtTeaG/HpBuuzw==", "dev": true }, - "anymatch": { - "version": "3.1.3", - "resolved": "https://registry.npmjs.org/anymatch/-/anymatch-3.1.3.tgz", - "integrity": "sha512-KMReFUr0B4t+D+OBkjR3KYqvocp2XaSzO55UcB6mgQMd3KbcE+mWTyvVV7D/zsdEbNnV6acZUutkiHQXvTr1Rw==", - "dev": true, - "requires": { - "normalize-path": "^3.0.0", - "picomatch": "^2.0.4" - } - }, - "array-flatten": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/array-flatten/-/array-flatten-1.1.1.tgz", - "integrity": "sha512-PCVAQswWemu6UdxsDFFX/+gVeYqKAod3D3UVm91jHwynguOwAvYPhx8nNlM++NqRcK6CxxpUafjmhIdKiHibqg==", - "dev": true - }, "asn1js": { "version": "3.0.10", "resolved": "https://registry.npmjs.org/asn1js/-/asn1js-3.0.10.tgz", @@ -4977,82 +4868,30 @@ "batch": { "version": "0.6.1", "resolved": "https://registry.npmjs.org/batch/-/batch-0.6.1.tgz", - "integrity": "sha1-3DQxT05nkxgJP8dgJyUl+UvyXBY=", - "dev": true - }, - "binary-extensions": { - "version": "2.3.0", - "resolved": "https://registry.npmjs.org/binary-extensions/-/binary-extensions-2.3.0.tgz", - "integrity": "sha512-Ceh+7ox5qe7LJuLHoY0feh3pHuUDHAcRUeyL2VYghZwfpkNIy/+8Ocg0a3UuSoYzavmylwuLWQOf3hl0jjMMIw==", + "integrity": "sha512-x+VAiMRL6UPkx+kudNvxTl6hB2XNNCG2r+7wixVfIYwu/2HKRXimwQyaumLjMveWvT2Hkd/cAJw+QBMfJ/EKVw==", "dev": true }, "body-parser": { - "version": "1.20.5", - "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-1.20.5.tgz", - "integrity": "sha512-3grm+/2tUOvu2cjJkvsIxrv/wVpfXQW4PsQHYm7yk4vfpu7Ekl6nEsYBoJUL6qDwZUx8wUhQ8tR2qz+ad9c9OA==", + "version": "2.3.0", + "resolved": "https://registry.npmjs.org/body-parser/-/body-parser-2.3.0.tgz", + "integrity": "sha512-2cGmJupaNgg+QUwVLAucDuWuoMZ6EX9iHDRswZ5lsNYEmwPaRknMPCLZz07yTzVq/83p4o/wzbDZbBrTvGGTIw==", "dev": true, "requires": { - "bytes": "~3.1.2", - "content-type": "~1.0.5", - "debug": "2.6.9", - "depd": "2.0.0", - "destroy": "~1.2.0", - "http-errors": "~2.0.1", - "iconv-lite": "~0.4.24", - "on-finished": "~2.4.1", - "qs": "~6.15.1", - "raw-body": "~2.5.3", - "type-is": "~1.6.18", - "unpipe": "~1.0.0" + "bytes": "^3.1.2", + "content-type": "^2.0.0", + "debug": "^4.4.3", + "http-errors": "^2.0.1", + "iconv-lite": "^0.7.2", + "on-finished": "^2.4.1", + "qs": "^6.15.2", + "raw-body": "^3.0.2", + "type-is": "^2.1.0" }, "dependencies": { - "debug": { - "version": "2.6.9", - "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", - "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", - "dev": true, - "requires": { - "ms": "2.0.0" - } - }, - "depd": { + "content-type": { "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", - "dev": true - }, - "http-errors": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", - "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", - "dev": true, - "requires": { - "depd": "~2.0.0", - "inherits": "~2.0.4", - "setprototypeof": "~1.2.0", - "statuses": "~2.0.2", - "toidentifier": "~1.0.1" - } - }, - "inherits": { - "version": "2.0.4", - "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", - "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", - "dev": true - }, - "qs": { - "version": "6.15.2", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.15.2.tgz", - "integrity": "sha512-Rzq0KEyX/w/tEybncDgdkZrJgVUsUMk3xjh3t5bv3S1HTAtg+uOYt72+ZfwiQwKdysThkTBdL/rTi6HDmX9Ddw==", - "dev": true, - "requires": { - "side-channel": "^1.1.0" - } - }, - "statuses": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", - "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-2.0.0.tgz", + "integrity": "sha512-j/O/d7GcZCyNl7/hwZAb606rzqkyvaDctLmckbxLzHvFBzTJHuGEdodATcP3yIRoDrLHkIATJuvzbFlp/ki2cQ==", "dev": true } } @@ -5143,19 +4982,12 @@ "dev": true }, "chokidar": { - "version": "3.6.0", - "resolved": "https://registry.npmjs.org/chokidar/-/chokidar-3.6.0.tgz", - "integrity": "sha512-7VT13fmjotKpGipCW9JEQAusEPE+Ei8nl6/g4FBAmIm0GOOLMua9NDDo/DWp0ZAxCr3cPq5ZpBqmPAQgDda2Pw==", + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/chokidar/-/chokidar-5.0.0.tgz", + "integrity": "sha512-TQMmc3w+5AxjpL8iIiwebF73dRDF4fBIieAqGn9RGCWaEVwQ6Fb2cGe31Yns0RRIzii5goJ1Y7xbMwo1TxMplw==", "dev": true, "requires": { - "anymatch": "~3.1.2", - "braces": "~3.0.2", - "fsevents": "~2.3.2", - "glob-parent": "~5.1.2", - "is-binary-path": "~2.1.0", - "is-glob": "~4.0.1", - "normalize-path": "~3.0.0", - "readdirp": "~3.6.0" + "readdirp": "^5.0.0" } }, "chrome-trace-event": { @@ -5244,21 +5076,10 @@ "dev": true }, "content-disposition": { - "version": "0.5.4", - "resolved": "https://registry.npmjs.org/content-disposition/-/content-disposition-0.5.4.tgz", - "integrity": "sha512-FveZTNuGw04cxlAiWbzi6zTAL/lhehaWbTtgluJh4/E95DqMwTmha3KZN1aAWA8cFIhHzMZUvLevkw5Rqk+tSQ==", - "dev": true, - "requires": { - "safe-buffer": "5.2.1" - }, - "dependencies": { - "safe-buffer": { - "version": "5.2.1", - "resolved": "https://registry.npmjs.org/safe-buffer/-/safe-buffer-5.2.1.tgz", - "integrity": "sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ==", - "dev": true - } - } + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/content-disposition/-/content-disposition-1.1.0.tgz", + "integrity": "sha512-5jRCH9Z/+DRP7rkvY83B+yGIGX96OYdJmzngqnw2SBSxqCFPd0w2km3s5iawpGX8krnwSGmF0FW5Nhr0Hfai3g==", + "dev": true }, "content-type": { "version": "1.0.5", @@ -5267,15 +5088,15 @@ "dev": true }, "cookie": { - "version": "0.7.1", - "resolved": "https://registry.npmjs.org/cookie/-/cookie-0.7.1.tgz", - "integrity": "sha512-6DnInpx7SJ2AK3+CTUE/ZM0vWTUboZCegxhC2xiIydHR9jNuTAASBrfEpHhiGOZw/nX51bHt6YQl8jsGo4y/0w==", + "version": "0.7.2", + "resolved": "https://registry.npmjs.org/cookie/-/cookie-0.7.2.tgz", + "integrity": "sha512-yki5XnKuf750l50uGTllt6kKILY4nQ1eNIQatoXEByZ5dWgnKqbnqmTrBE5B4N7lrMJKQ2ytWMiTO2o0v6Ew/w==", "dev": true }, "cookie-signature": { - "version": "1.0.6", - "resolved": "https://registry.npmjs.org/cookie-signature/-/cookie-signature-1.0.6.tgz", - "integrity": "sha512-QADzlaHc8icV8I7vbaJXJwod9HWYp8uCqf1xa4OfNu1T7JVxQIrUgOWtHdNDtPiywmFbiS12VjotIXLrKM3orQ==", + "version": "1.2.2", + "resolved": "https://registry.npmjs.org/cookie-signature/-/cookie-signature-1.2.2.tgz", + "integrity": "sha512-D76uU73ulSXrD1UXF4KE2TMxVVwhsnCgfAyTg9k8P6KGZjlXKrOLe4dJQKI3Bxi5wjesZoFXJWElNWBjPZMbhg==", "dev": true }, "copy-webpack-plugin": { @@ -5308,12 +5129,6 @@ } } }, - "core-util-is": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/core-util-is/-/core-util-is-1.0.2.tgz", - "integrity": "sha1-tf1UIgqivFq1eqtxQMlAdUUDwac=", - "dev": true - }, "cross-spawn": { "version": "7.0.3", "resolved": "https://registry.npmjs.org/cross-spawn/-/cross-spawn-7.0.3.tgz", @@ -5329,26 +5144,26 @@ "version": "file:../pkg" }, "debug": { - "version": "4.1.1", - "resolved": "https://registry.npmjs.org/debug/-/debug-4.1.1.tgz", - "integrity": "sha512-pYAIzeRo8J6KPEaJ0VWOh5Pzkbw/RetuzehGM7QRRX5he4fPHx2rdKMB256ehJCkX+XRQm16eZLqLNS8RSZXZw==", + "version": "4.4.3", + "resolved": "https://registry.npmjs.org/debug/-/debug-4.4.3.tgz", + "integrity": "sha512-RGwwWnwQvkVfavKVt22FGLw+xYSdzARwm0ru6DhTVA3umU5hZc28V3kO4stgYryrTlLpuvgI9GiijltAjNbcqA==", "dev": true, "requires": { - "ms": "^2.1.1" + "ms": "^2.1.3" }, "dependencies": { "ms": { - "version": "2.1.2", - "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.2.tgz", - "integrity": "sha512-sGkPx+VjMtmA6MX27oA4FBFELFCZZ4S4XqeGOXCv68tT+jb3vk/RyaKWP0PTKyWtmLSM0b+adUTEvbs1PEaH2w==", + "version": "2.1.3", + "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.3.tgz", + "integrity": "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==", "dev": true } } }, "default-browser": { - "version": "5.2.1", - "resolved": "https://registry.npmjs.org/default-browser/-/default-browser-5.2.1.tgz", - "integrity": "sha512-WY/3TUME0x3KPYdRRxEJJvXRHV4PyPoUsxtZa78lwItwRQRHhd2U9xOscaT/YTf8uCXIAjeJOFBVEh/7FtD8Xg==", + "version": "5.5.0", + "resolved": "https://registry.npmjs.org/default-browser/-/default-browser-5.5.0.tgz", + "integrity": "sha512-H9LMLr5zwIbSxrmvikGuI/5KGhZ8E2zH3stkMgM5LpOWDutGM2JZaj460Udnf1a+946zc7YBgrqEWwbk7zHvGw==", "dev": true, "requires": { "bundle-name": "^4.1.0", @@ -5356,9 +5171,9 @@ } }, "default-browser-id": { - "version": "5.0.0", - "resolved": "https://registry.npmjs.org/default-browser-id/-/default-browser-id-5.0.0.tgz", - "integrity": "sha512-A6p/pu/6fyBcA1TRz/GqWYPViplrftcW2gZC9q79ngNCKAeR/X3gcEdXQHl4KNXV+3wgIJ1CPkJQ3IHM6lcsyA==", + "version": "5.0.1", + "resolved": "https://registry.npmjs.org/default-browser-id/-/default-browser-id-5.0.1.tgz", + "integrity": "sha512-x1VCxdX4t+8wVfd1so/9w+vQ4vx7lKd2Qp5tDRutErwmR85OgmfX7RlLRMWafRMY7hbEiXIbudNrjOAPa/hL8Q==", "dev": true }, "define-lazy-prop": { @@ -5368,21 +5183,9 @@ "dev": true }, "depd": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/depd/-/depd-1.1.2.tgz", - "integrity": "sha1-m81S4UwJd2PnSbJ0xDRu0uVgtak=", - "dev": true - }, - "destroy": { - "version": "1.2.0", - "resolved": "https://registry.npmjs.org/destroy/-/destroy-1.2.0.tgz", - "integrity": "sha512-2sJGJTaXIIaR1w4iJSNoN0hnMY7Gpc/n8D4qSCJw8QqFWXf7cuAgnEHxBpweaVcPevC2l3KpjYCx3NypQQgaJg==", - "dev": true - }, - "detect-node": { - "version": "2.0.4", - "resolved": "https://registry.npmjs.org/detect-node/-/detect-node-2.0.4.tgz", - "integrity": "sha512-ZIzRpLJrOj7jjP2miAtgqIfmzbxa4ZOr5jJc601zklsfEx9oTzmmj2nVpIPRpNlRTIh8lc1kyViIY7BWSGNmKw==", + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", + "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", "dev": true }, "dns-packet": { @@ -5458,9 +5261,9 @@ "dev": true }, "es-object-atoms": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/es-object-atoms/-/es-object-atoms-1.1.1.tgz", - "integrity": "sha512-FGgH2h8zKNim9ljj7dankFPcICIK9Cp5bm+c2gQSYePhpaG5+esrLODihIorn+Pe6FGJzWhXQotPv73jTaldXA==", + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/es-object-atoms/-/es-object-atoms-1.1.2.tgz", + "integrity": "sha512-HWcBoN6NileqtSydK2FqHbS/LoDd2pqrnQHLyJzBj4kOp/ky2MWMN694xOfkK8/SnUsW2DH7EfyVlydKCsm1Zw==", "dev": true, "requires": { "es-errors": "^1.3.0" @@ -5475,7 +5278,7 @@ "escape-html": { "version": "1.0.3", "resolved": "https://registry.npmjs.org/escape-html/-/escape-html-1.0.3.tgz", - "integrity": "sha1-Aljq5NPQwJdN4cFpGI7wBR0dGYg=", + "integrity": "sha512-NiSupZ4OeuGwr68lGIeym/ksIZMJodUGOSCZ/FSnTxcrekbvqrgdUxlJOMpijaKZVjAJrWrGs/6Jy8OMuyj9ow==", "dev": true }, "eslint-scope": { @@ -5517,12 +5320,6 @@ "integrity": "sha512-aIL5Fx7mawVa300al2BnEE4iNvo1qETxLrPI/o05L7z6go7fCw1J6EQmbK4FmJ2AS7kgVF/KEZWufBfdClMcPg==", "dev": true }, - "eventemitter3": { - "version": "4.0.7", - "resolved": "https://registry.npmjs.org/eventemitter3/-/eventemitter3-4.0.7.tgz", - "integrity": "sha512-8guHBZCwKnFhYdHr2ysuRWErTwhoN2X8XELRlrRwpmfeY2jjuUN4taQMsULKUVo1K4DvZl+0pgfyoysHxvmvEw==", - "dev": true - }, "events": { "version": "3.3.0", "resolved": "https://registry.npmjs.org/events/-/events-3.3.0.tgz", @@ -5530,70 +5327,55 @@ "dev": true }, "express": { - "version": "4.22.1", - "resolved": "https://registry.npmjs.org/express/-/express-4.22.1.tgz", - "integrity": "sha512-F2X8g9P1X7uCPZMA3MVf9wcTqlyNp7IhH5qPCI0izhaOIYXaW9L535tGA3qmjRzpH+bZczqq7hVKxTR4NWnu+g==", - "dev": true, - "requires": { - "accepts": "~1.3.8", - "array-flatten": "1.1.1", - "body-parser": "~1.20.3", - "content-disposition": "~0.5.4", - "content-type": "~1.0.4", - "cookie": "~0.7.1", - "cookie-signature": "~1.0.6", - "debug": "2.6.9", - "depd": "2.0.0", - "encodeurl": "~2.0.0", - "escape-html": "~1.0.3", - "etag": "~1.8.1", - "finalhandler": "~1.3.1", - "fresh": "~0.5.2", - "http-errors": "~2.0.0", - "merge-descriptors": "1.0.3", - "methods": "~1.1.2", - "on-finished": "~2.4.1", - "parseurl": "~1.3.3", - "path-to-regexp": "~0.1.12", - "proxy-addr": "~2.0.7", - "qs": "~6.14.0", - "range-parser": "~1.2.1", - "safe-buffer": "5.2.1", - "send": "~0.19.0", - "serve-static": "~1.16.2", - "setprototypeof": "1.2.0", - "statuses": "~2.0.1", - "type-is": "~1.6.18", - "utils-merge": "1.0.1", - "vary": "~1.1.2" - }, - "dependencies": { - "debug": { - "version": "2.6.9", - "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", - "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", + "version": "5.2.1", + "resolved": "https://registry.npmjs.org/express/-/express-5.2.1.tgz", + "integrity": "sha512-hIS4idWWai69NezIdRt2xFVofaF4j+6INOpJlVOLDO8zXGpUVEVzIYk12UUi2JzjEzWL3IOAxcTubgz9Po0yXw==", + "dev": true, + "requires": { + "accepts": "^2.0.0", + "body-parser": "^2.2.1", + "content-disposition": "^1.0.0", + "content-type": "^1.0.5", + "cookie": "^0.7.1", + "cookie-signature": "^1.2.1", + "debug": "^4.4.0", + "depd": "^2.0.0", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "etag": "^1.8.1", + "finalhandler": "^2.1.0", + "fresh": "^2.0.0", + "http-errors": "^2.0.0", + "merge-descriptors": "^2.0.0", + "mime-types": "^3.0.0", + "on-finished": "^2.4.1", + "once": "^1.4.0", + "parseurl": "^1.3.3", + "proxy-addr": "^2.0.7", + "qs": "^6.14.0", + "range-parser": "^1.2.1", + "router": "^2.2.0", + "send": "^1.1.0", + "serve-static": "^2.2.0", + "statuses": "^2.0.1", + "type-is": "^2.0.1", + "vary": "^1.1.2" + }, + "dependencies": { + "mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true + }, + "mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", "dev": true, "requires": { - "ms": "2.0.0" + "mime-db": "^1.54.0" } - }, - "depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", - "dev": true - }, - "safe-buffer": { - "version": "5.2.1", - "resolved": "https://registry.npmjs.org/safe-buffer/-/safe-buffer-5.2.1.tgz", - "integrity": "sha512-rp3So07KcdmmKbGvgaNxQSJr7bGVSVk5S9Eq1F+ppbRo70+YeaDxkw5Dd8NPN+GD6bjnYm2VuPuCXmpuYvmCXQ==", - "dev": true - }, - "statuses": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", - "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", - "dev": true } } }, @@ -5604,9 +5386,9 @@ "dev": true }, "fast-uri": { - "version": "3.1.0", - "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.0.tgz", - "integrity": "sha512-iPeeDKJSWf4IEOasVVrknXpaBV0IApz/gp7S2bb7Z4Lljbl2MGJRqInZiUrQwV16cpzw/D3S5j5Julj/gT52AA==", + "version": "3.1.5", + "resolved": "https://registry.npmjs.org/fast-uri/-/fast-uri-3.1.5.tgz", + "integrity": "sha512-gHwA1O9LDIcKunMKhObS/HimwtehO1nPUECKAu5TpKgaO19fcWEl4bliWe1jWxVFvIXztJjjQ4L8XQ1EU9f7Jw==", "dev": true }, "fastest-levenshtein": { @@ -5615,15 +5397,6 @@ "integrity": "sha512-eRnCtTTtGZFpQCwhJiUOuxPQWRXVKYDn0b2PeHfXL6/Zi53SLAzAHfVhVWK2AryC/WH05kGfxhFIPvTF0SXQzg==", "dev": true }, - "faye-websocket": { - "version": "0.11.4", - "resolved": "https://registry.npmjs.org/faye-websocket/-/faye-websocket-0.11.4.tgz", - "integrity": "sha512-CzbClwlXAuiRQAlUyfqPgvPoNKTckTPGfwZV4ZdAhVcP2lh9KUxJg2b5GkE7XbjKQ3YJnQ9z6D9ntLAlB+tP8g==", - "dev": true, - "requires": { - "websocket-driver": ">=0.5.1" - } - }, "fill-range": { "version": "7.1.1", "resolved": "https://registry.npmjs.org/fill-range/-/fill-range-7.1.1.tgz", @@ -5634,35 +5407,17 @@ } }, "finalhandler": { - "version": "1.3.1", - "resolved": "https://registry.npmjs.org/finalhandler/-/finalhandler-1.3.1.tgz", - "integrity": "sha512-6BN9trH7bp3qvnrRyzsBz+g3lZxTNZTbVO2EV1CS0WIcDbawYVdYvGflME/9QP0h0pYlCDBCTjYa9nZzMDpyxQ==", + "version": "2.1.1", + "resolved": "https://registry.npmjs.org/finalhandler/-/finalhandler-2.1.1.tgz", + "integrity": "sha512-S8KoZgRZN+a5rNwqTxlZZePjT/4cnm0ROV70LedRHZ0p8u9fRID0hJUZQpkKLzro8LfmC8sx23bY6tVNxv8pQA==", "dev": true, "requires": { - "debug": "2.6.9", - "encodeurl": "~2.0.0", - "escape-html": "~1.0.3", - "on-finished": "2.4.1", - "parseurl": "~1.3.3", - "statuses": "2.0.1", - "unpipe": "~1.0.0" - }, - "dependencies": { - "debug": { - "version": "2.6.9", - "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", - "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", - "dev": true, - "requires": { - "ms": "2.0.0" - } - }, - "statuses": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", - "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", - "dev": true - } + "debug": "^4.4.0", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "on-finished": "^2.4.1", + "parseurl": "^1.3.3", + "statuses": "^2.0.1" } }, "find-up": { @@ -5675,12 +5430,6 @@ "path-exists": "^4.0.0" } }, - "follow-redirects": { - "version": "1.16.0", - "resolved": "https://registry.npmjs.org/follow-redirects/-/follow-redirects-1.16.0.tgz", - "integrity": "sha512-y5rN/uOsadFT/JfYwhxRS5R7Qce+g3zG97+JrtFZlC9klX/W5hD7iiLzScI4nZqUS7DNUdhPgw4xI8W2LuXlUw==", - "dev": true - }, "forwarded": { "version": "0.2.0", "resolved": "https://registry.npmjs.org/forwarded/-/forwarded-0.2.0.tgz", @@ -5688,18 +5437,11 @@ "dev": true }, "fresh": { - "version": "0.5.2", - "resolved": "https://registry.npmjs.org/fresh/-/fresh-0.5.2.tgz", - "integrity": "sha512-zJ2mQYM18rEFOudeV4GShTGIQ7RbzA7ozbU9I/XBpm7kqgMywgmylMwXHxZJmkVoYkna9d2pVXVXPdYTP9ej8Q==", + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/fresh/-/fresh-2.0.0.tgz", + "integrity": "sha512-Rx/WycZ60HOaqLKAi6cHRKKI7zxWbJ31MhntmtwMoaTeF7XFH9hhBp8vITaMidfljRQ6eYWCKkaTK+ykVJHP2A==", "dev": true }, - "fsevents": { - "version": "2.3.3", - "resolved": "https://registry.npmjs.org/fsevents/-/fsevents-2.3.3.tgz", - "integrity": "sha512-5xoDfX+fL7faATnagmWPpbFtwh/R77WmMMqqHGS65C3vvB0YHrgF+B1YmZ3441tMj5n63k0212XNoJwzlhffQw==", - "dev": true, - "optional": true - }, "function-bind": { "version": "1.1.2", "resolved": "https://registry.npmjs.org/function-bind/-/function-bind-1.1.2.tgz", @@ -5734,15 +5476,6 @@ "es-object-atoms": "^1.0.0" } }, - "glob-parent": { - "version": "5.1.2", - "resolved": "https://registry.npmjs.org/glob-parent/-/glob-parent-5.1.2.tgz", - "integrity": "sha512-AOIgSQCepiJYwP3ARnGx+5VnTu2HBYdzbGP45eLw1vr3zB3vZLeyed1sC9hnbcOc9/SrMyM5RPQrkGz4aS9Zow==", - "dev": true, - "requires": { - "is-glob": "^4.0.1" - } - }, "glob-to-regexp": { "version": "0.4.1", "resolved": "https://registry.npmjs.org/glob-to-regexp/-/glob-to-regexp-0.4.1.tgz", @@ -5761,12 +5494,6 @@ "integrity": "sha512-RbJ5/jmFcNNCcDV5o9eTnBLJ/HszWV0P73bc+Ff4nS/rJj+YaS6IGyiOL0VoBYX+l1Wrl3k63h/KrH+nhJ0XvQ==", "dev": true }, - "handle-thing": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/handle-thing/-/handle-thing-2.0.1.tgz", - "integrity": "sha512-9Qn4yBxelxoh2Ow62nP+Ka/kMnOXRi8BXnRaUwezLNhqelnN49xKz4F/dPP8OYLxLxq6JDtZb2i9XznUQbNPTg==", - "dev": true - }, "has": { "version": "1.0.3", "resolved": "https://registry.npmjs.org/has/-/has-1.0.3.tgz", @@ -5789,95 +5516,46 @@ "dev": true }, "hasown": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.2.tgz", - "integrity": "sha512-0hJU9SCPvmMzIBdZFqNPXWa6dqh7WdH0cII9y+CyS8rG3nL48Bclra9HmKhVVUHyPWNH5Y7xDwAB7bfgSjkUMQ==", + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/hasown/-/hasown-2.0.4.tgz", + "integrity": "sha512-T2UbfbBEF32wiepXIsMlTW9+dDYC6wMh/t/vYA4tuOMKqWz/n3vr1NFSxQiyP+zk2mXsoMA/i/7qV6LKut1t1A==", "dev": true, "requires": { "function-bind": "^1.1.2" } }, - "hpack.js": { - "version": "2.1.6", - "resolved": "https://registry.npmjs.org/hpack.js/-/hpack.js-2.1.6.tgz", - "integrity": "sha1-h3dMCUnlE/QuhFdbPEVoH63ioLI=", - "dev": true, - "requires": { - "inherits": "^2.0.1", - "obuf": "^1.0.0", - "readable-stream": "^2.0.1", - "wbuf": "^1.1.0" - } - }, - "http-deceiver": { - "version": "1.2.7", - "resolved": "https://registry.npmjs.org/http-deceiver/-/http-deceiver-1.2.7.tgz", - "integrity": "sha1-+nFolEq5pRnTN8sL7HKE3D5yPYc=", - "dev": true - }, "http-errors": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.0.tgz", - "integrity": "sha512-FtwrG/euBzaEjYeRqOgly7G0qviiXoJWnvEH2Z1plBdXgbyjv34pHTSb9zoeHMyDy33+DWy5Wt9Wo+TURtOYSQ==", - "dev": true, - "requires": { - "depd": "2.0.0", - "inherits": "2.0.4", - "setprototypeof": "1.2.0", - "statuses": "2.0.1", - "toidentifier": "1.0.1" - }, - "dependencies": { - "depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", - "dev": true - }, - "inherits": { - "version": "2.0.4", - "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", - "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", - "dev": true - }, - "statuses": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", - "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", - "dev": true - } - } - }, - "http-parser-js": { - "version": "0.5.8", - "resolved": "https://registry.npmjs.org/http-parser-js/-/http-parser-js-0.5.8.tgz", - "integrity": "sha512-SGeBX54F94Wgu5RH3X5jsDtf4eHyRogWX1XGT3b4HuW3tQPM4AaBzoUji/4AAJNXCEOWZ5O0DgZmJw1947gD5Q==", - "dev": true - }, - "http-proxy": { - "version": "1.18.1", - "resolved": "https://registry.npmjs.org/http-proxy/-/http-proxy-1.18.1.tgz", - "integrity": "sha512-7mz/721AbnJwIVbnaSv1Cz3Am0ZLT/UBwkC92VlxhXv/k/BBQfM2fXElQNC27BVGr0uwUpplYPQM9LnaBMR5NQ==", + "version": "2.0.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", + "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", "dev": true, "requires": { - "eventemitter3": "^4.0.0", - "follow-redirects": "^1.0.0", - "requires-port": "^1.0.0" + "depd": "~2.0.0", + "inherits": "~2.0.4", + "setprototypeof": "~1.2.0", + "statuses": "~2.0.2", + "toidentifier": "~1.0.1" } }, "http-proxy-middleware": { - "version": "2.0.9", - "resolved": "https://registry.npmjs.org/http-proxy-middleware/-/http-proxy-middleware-2.0.9.tgz", - "integrity": "sha512-c1IyJYLYppU574+YI7R4QyX2ystMtVXZwIdzazUIPIJsHuWNd+mho2j+bKoHftndicGj9yh+xjd+l0yj7VeT1Q==", + "version": "4.2.0", + "resolved": "https://registry.npmjs.org/http-proxy-middleware/-/http-proxy-middleware-4.2.0.tgz", + "integrity": "sha512-ZA+oNOoM+GLoFTIzhkJptVQov73Srep2LBqhF8hG8CIPKO3nam1jonXVQ/QUH8RbwsmaaVz2SOJdzBNBHNtKbw==", "dev": true, "requires": { - "@types/http-proxy": "^1.17.8", - "http-proxy": "^1.18.1", - "is-glob": "^4.0.1", - "is-plain-obj": "^3.0.0", - "micromatch": "^4.0.2" + "debug": "^4.4.3", + "httpxy": "^0.5.4", + "is-glob": "^4.0.3", + "is-plain-obj": "^4.1.0", + "micromatch": "^4.0.8" } }, + "httpxy": { + "version": "0.5.5", + "resolved": "https://registry.npmjs.org/httpxy/-/httpxy-0.5.5.tgz", + "integrity": "sha512-uDjmnPyp1q4Sgzf3w+J/Fc6UqcCEj0x4Wjp7OqK5dGhNeDgpyrAmnS6ey8QWrX3SWDon2DMKf9sBa5X9+CVyMA==", + "dev": true + }, "hyperdyperid": { "version": "1.2.0", "resolved": "https://registry.npmjs.org/hyperdyperid/-/hyperdyperid-1.2.0.tgz", @@ -5885,12 +5563,12 @@ "dev": true }, "iconv-lite": { - "version": "0.4.24", - "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.4.24.tgz", - "integrity": "sha512-v3MXnZAcvnywkTUEZomIActle7RXXeedOR31wwl7VlyoXO4Qi9arvSenNQWne1TcRwhCL1HwLI21bEqdpj8/rA==", + "version": "0.7.3", + "resolved": "https://registry.npmjs.org/iconv-lite/-/iconv-lite-0.7.3.tgz", + "integrity": "sha512-IKXpvIzjnC9XTAUbVBcMfGS0EPaIXtW6v+zr+RRp+hqULEpo0owZax6wyRwPOJbWbzjYspQwusTsfVr0ifh4uQ==", "dev": true, "requires": { - "safer-buffer": ">= 2.1.2 < 3" + "safer-buffer": ">= 2.1.2 < 3.0.0" } }, "import-local": { @@ -5904,9 +5582,9 @@ } }, "inherits": { - "version": "2.0.3", - "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.3.tgz", - "integrity": "sha1-Yzwsg+PaQqUC9SRmAiSA9CCCYd4=", + "version": "2.0.4", + "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", + "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", "dev": true }, "interpret": { @@ -5916,20 +5594,11 @@ "dev": true }, "ipaddr.js": { - "version": "2.1.0", - "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-2.1.0.tgz", - "integrity": "sha512-LlbxQ7xKzfBusov6UMi4MFpEg0m+mAm9xyNGEduwXMEDuf4WfzB/RZwMVYEd7IKGvh4IUkEXYxtAVu9T3OelJQ==", + "version": "2.4.0", + "resolved": "https://registry.npmjs.org/ipaddr.js/-/ipaddr.js-2.4.0.tgz", + "integrity": "sha512-9VGk3HGanVE6JoZXHiCpnGy5X0jYDnN4EA4lntFPj+1vIWlFhIylq2CrrCOJH9EAhc5CYhq18F2Av2tgoAPsYQ==", "dev": true }, - "is-binary-path": { - "version": "2.1.0", - "resolved": "https://registry.npmjs.org/is-binary-path/-/is-binary-path-2.1.0.tgz", - "integrity": "sha512-ZMERYes6pDydyuGidse7OsHxtbI7WVeUEozgR/g7rd0xUimYNlvZRE/K2MgZTjWy725IfelLeVcEM97mmtRGXw==", - "dev": true, - "requires": { - "binary-extensions": "^2.0.0" - } - }, "is-core-module": { "version": "2.13.0", "resolved": "https://registry.npmjs.org/is-core-module/-/is-core-module-2.13.0.tgz", @@ -5960,6 +5629,12 @@ "is-extglob": "^2.1.1" } }, + "is-in-ssh": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/is-in-ssh/-/is-in-ssh-1.0.0.tgz", + "integrity": "sha512-jYa6Q9rH90kR1vKB6NM7qqd1mge3Fx4Dhw5TVlK1MUBqhEOuCagrEHMevNuCcbECmXZ0ThXkRm+Ymr51HwEPAw==", + "dev": true + }, "is-inside-container": { "version": "1.0.0", "resolved": "https://registry.npmjs.org/is-inside-container/-/is-inside-container-1.0.0.tgz", @@ -5970,9 +5645,9 @@ } }, "is-network-error": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/is-network-error/-/is-network-error-1.1.0.tgz", - "integrity": "sha512-tUdRRAnhT+OtCZR/LxZelH/C7QtjtFrTu5tXCA8pl55eTUElUHT+GPYV8MBMBvea/j+NxQqVt3LbWMRir7Gx9g==", + "version": "1.3.2", + "resolved": "https://registry.npmjs.org/is-network-error/-/is-network-error-1.3.2.tgz", + "integrity": "sha512-PhBY86zaxNZUuWP6h13Vu5oFe0XY6/UlKzQnYFELzGVHygP3MxmvTfYSG7GN3aIab/iWudSMgjSnG9Dq+nHrgA==", "dev": true }, "is-number": { @@ -5982,9 +5657,9 @@ "dev": true }, "is-plain-obj": { - "version": "3.0.0", - "resolved": "https://registry.npmjs.org/is-plain-obj/-/is-plain-obj-3.0.0.tgz", - "integrity": "sha512-gwsOE28k+23GP1B6vFl1oVh/WOzmawBrKwo5Ev6wMKzPkaXaCDIQKzLnvsA42DRlbVTWorkgTKIviAKCWkfUwA==", + "version": "4.1.0", + "resolved": "https://registry.npmjs.org/is-plain-obj/-/is-plain-obj-4.1.0.tgz", + "integrity": "sha512-+Pgi+vMuUNkJyExiMBt5IlFoMyKnr5zhJ4Uspz58WOhBF5QoIZkFyNHIbBAtHwzVAgk5RtndVNsDRN61/mmDqg==", "dev": true }, "is-plain-object": { @@ -5996,21 +5671,21 @@ "isobject": "^3.0.1" } }, + "is-promise": { + "version": "4.0.0", + "resolved": "https://registry.npmjs.org/is-promise/-/is-promise-4.0.0.tgz", + "integrity": "sha512-hvpoI6korhJMnej285dSg6nu1+e6uxs7zG3BYAm5byqDsgJNWwxzM6z6iZiAgQR4TJ30JmBTOwqZUw3WlyH3AQ==", + "dev": true + }, "is-wsl": { - "version": "3.1.0", - "resolved": "https://registry.npmjs.org/is-wsl/-/is-wsl-3.1.0.tgz", - "integrity": "sha512-UcVfVfaK4Sc4m7X3dUSoHoozQGBEFeDC+zVo06t98xe8CzHSZZBekNXH+tu0NalHolcJ/QAGqS46Hef7QXBIMw==", + "version": "3.1.1", + "resolved": "https://registry.npmjs.org/is-wsl/-/is-wsl-3.1.1.tgz", + "integrity": "sha512-e6rvdUCiQCAuumZslxRJWR/Doq4VpPR82kqclvcS0efgt430SlGIk05vdCN58+VrzgtIcfNODjozVielycD4Sw==", "dev": true, "requires": { "is-inside-container": "^1.0.0" } }, - "isarray": { - "version": "1.0.0", - "resolved": "https://registry.npmjs.org/isarray/-/isarray-1.0.0.tgz", - "integrity": "sha1-u5NdSFgsuhaMBoNJV6VKPgcSTxE=", - "dev": true - }, "isexe": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/isexe/-/isexe-2.0.0.tgz", @@ -6053,13 +5728,13 @@ "dev": true }, "launch-editor": { - "version": "2.10.0", - "resolved": "https://registry.npmjs.org/launch-editor/-/launch-editor-2.10.0.tgz", - "integrity": "sha512-D7dBRJo/qcGX9xlvt/6wUYzQxjh5G1RvZPgPv8vi4KRU99DVQL/oW7tnVOCCTm2HGeo3C5HvGE5Yrh6UBoZ0vA==", + "version": "2.14.1", + "resolved": "https://registry.npmjs.org/launch-editor/-/launch-editor-2.14.1.tgz", + "integrity": "sha512-QWBrQsMpH7gPr965dsKD/3cKWiNoTjpATQf++Xq63N6sKRGMwlVXz41O1IZTMfZQgBctD/K5Zt06+/I6pP6+HA==", "dev": true, "requires": { - "picocolors": "^1.0.0", - "shell-quote": "^1.8.1" + "picocolors": "^1.1.1", + "shell-quote": "^1.8.4" } }, "loader-runner": { @@ -6084,20 +5759,30 @@ "dev": true }, "media-typer": { - "version": "0.3.0", - "resolved": "https://registry.npmjs.org/media-typer/-/media-typer-0.3.0.tgz", - "integrity": "sha512-dq+qelQ9akHpcOl/gUVRTxVIOkAJ1wR3QAvb4RsVjS8oVoFjDGTc679wJYmUmknUF5HwMLOgb5O+a3KxfWapPQ==", + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/media-typer/-/media-typer-1.1.1.tgz", + "integrity": "sha512-yz3xRaG20c6/BOzvYoDaGtPmGscs7YivItZEEqe6GbwNfHuxu9YNmvnEkMzKldAGY4/80pRcQRZSEnhquk9XuQ==", "dev": true }, "memfs": { - "version": "4.17.2", - "resolved": "https://registry.npmjs.org/memfs/-/memfs-4.17.2.tgz", - "integrity": "sha512-NgYhCOWgovOXSzvYgUW0LQ7Qy72rWQMGGFJDoWg4G30RHd3z77VbYdtJ4fembJXBy8pMIUA31XNAupobOQlwdg==", - "dev": true, - "requires": { - "@jsonjoy.com/json-pack": "^1.0.3", - "@jsonjoy.com/util": "^1.3.0", - "tree-dump": "^1.0.1", + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/memfs/-/memfs-4.64.0.tgz", + "integrity": "sha512-Kw72fgY7Wn+sD8KmtNWSafl1dz0UvAsE/PHs3YVfLiaZuA3HxNm9sRLqAu0ATiBGJvME1PxZXbBZPv5GycDeAw==", + "dev": true, + "requires": { + "@jsonjoy.com/fs-core": "4.64.0", + "@jsonjoy.com/fs-fsa": "4.64.0", + "@jsonjoy.com/fs-node": "4.64.0", + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-to-fsa": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "@jsonjoy.com/fs-print": "4.64.0", + "@jsonjoy.com/fs-snapshot": "4.64.0", + "@jsonjoy.com/json-pack": "^1.11.0", + "@jsonjoy.com/util": "^1.9.0", + "glob-to-regex.js": "^1.0.1", + "thingies": "^2.5.0", + "tree-dump": "^1.0.3", "tslib": "^2.0.0" }, "dependencies": { @@ -6108,36 +5793,231 @@ "dev": true, "requires": {} }, + "@jsonjoy.com/buffers": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/buffers/-/buffers-17.67.0.tgz", + "integrity": "sha512-tfExRpYxBvi32vPs9ZHaTjSP4fHAfzSmcahOfNxtvGHcyJel+aibkPlGeBB+7AoC6hL7lXIE++8okecBxx7lcw==", + "dev": true, + "requires": {} + }, + "@jsonjoy.com/codegen": { + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/codegen/-/codegen-1.0.0.tgz", + "integrity": "sha512-E8Oy+08cmCf0EK/NMxpaJZmOxPqM+6iSe2S4nlSBrPZOORoDJILxtbSUEDKQyTamm/BVAhIGllOBNU79/dwf0g==", + "dev": true, + "requires": {} + }, + "@jsonjoy.com/fs-core": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-core/-/fs-core-4.64.0.tgz", + "integrity": "sha512-zs2TAq7Six5jgMuoMNjpspAvOP3mhtgq/k1UyQodEzCtQi/N83y2/y+zcvnZSGp/Rxq96DBN+bValOBQAyn/ew==", + "dev": true, + "requires": { + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "thingies": "^2.5.0" + } + }, + "@jsonjoy.com/fs-fsa": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-fsa/-/fs-fsa-4.64.0.tgz", + "integrity": "sha512-nMWOVbkLFyEgmXZih3wyvxA9XpgyyqyfrINMHvEFqhi7uqfRl7c9ERJt6yX7vgMPrB9Uo+OJO+Spa0cFzPD01w==", + "dev": true, + "requires": { + "@jsonjoy.com/fs-core": "4.64.0", + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "thingies": "^2.5.0" + } + }, + "@jsonjoy.com/fs-node": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-node/-/fs-node-4.64.0.tgz", + "integrity": "sha512-dO+NNkODbUli4uV42bcNrrLvq5rE7SNpdZ5TNd0dtbLsAaNK3MDiIC9lUi+brboGoIjW6vd2fB1qao60nrk5xA==", + "dev": true, + "requires": { + "@jsonjoy.com/fs-core": "4.64.0", + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "@jsonjoy.com/fs-print": "4.64.0", + "@jsonjoy.com/fs-snapshot": "4.64.0", + "glob-to-regex.js": "^1.0.0", + "thingies": "^2.5.0" + } + }, + "@jsonjoy.com/fs-node-builtins": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-node-builtins/-/fs-node-builtins-4.64.0.tgz", + "integrity": "sha512-/o7WRFhUWaM/fOrslwLZGnzn4RmRILykn+lAL+mNObqqRNw+CQSiij6hpCeZ+C7buhdoVo7go/OYqzaSUfDYmA==", + "dev": true, + "requires": {} + }, + "@jsonjoy.com/fs-node-to-fsa": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-node-to-fsa/-/fs-node-to-fsa-4.64.0.tgz", + "integrity": "sha512-WDD9WVs0hb7UAEKTgZW2f66WDrbj7gIIWwpP3spbLyXa0rghtUaFTB8L4gdR3ZCWwiKIsj38/CNijpVmpnuPUw==", + "dev": true, + "requires": { + "@jsonjoy.com/fs-fsa": "4.64.0", + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "@jsonjoy.com/fs-node-utils": "4.64.0" + } + }, + "@jsonjoy.com/fs-node-utils": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-node-utils/-/fs-node-utils-4.64.0.tgz", + "integrity": "sha512-k5Indsx9hWW9xSF7Y6oSKKwtCUNhzZxadub3owhIlitc+iMRVlPPdX2duTKQWBL3qNWpXya8jykgaaWpheeS4w==", + "dev": true, + "requires": { + "@jsonjoy.com/fs-node-builtins": "4.64.0", + "glob-to-regex.js": "^1.0.1" + } + }, + "@jsonjoy.com/fs-print": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-print/-/fs-print-4.64.0.tgz", + "integrity": "sha512-PHZFccchvkhWrwPWHjmVAhbC3vSHCtyZvlZfJJ3ho2bnzl450hXri6/8e6pbkWdH+SkmLXNml0sV8e5HDAfxKw==", + "dev": true, + "requires": { + "@jsonjoy.com/fs-node-utils": "4.64.0", + "tree-dump": "^1.1.0" + } + }, + "@jsonjoy.com/fs-snapshot": { + "version": "4.64.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/fs-snapshot/-/fs-snapshot-4.64.0.tgz", + "integrity": "sha512-oM7UDeL83q6NBzzsfKAsYKXKVXlykKFqqOLh4xZZKAzzROTlInkPbc6LTDGThEOnPiFiUzA7tYziHG9xavd76Q==", + "dev": true, + "requires": { + "@jsonjoy.com/buffers": "^17.65.0", + "@jsonjoy.com/fs-node-utils": "4.64.0", + "@jsonjoy.com/json-pack": "^17.65.0", + "@jsonjoy.com/util": "^17.65.0" + }, + "dependencies": { + "@jsonjoy.com/base64": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/base64/-/base64-17.67.0.tgz", + "integrity": "sha512-5SEsJGsm15aP8TQGkDfJvz9axgPwAEm98S5DxOuYe8e1EbfajcDmgeXXzccEjh+mLnjqEKrkBdjHWS5vFNwDdw==", + "dev": true, + "requires": {} + }, + "@jsonjoy.com/codegen": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/codegen/-/codegen-17.67.0.tgz", + "integrity": "sha512-idnkUplROpdBOV0HMcwhsCUS5TRUi9poagdGs70A6S4ux9+/aPuKbh8+UYRTLYQHtXvAdNfQWXDqZEx5k4Dj2Q==", + "dev": true, + "requires": {} + }, + "@jsonjoy.com/json-pack": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pack/-/json-pack-17.67.0.tgz", + "integrity": "sha512-t0ejURcGaZsn1ClbJ/3kFqSOjlryd92eQY465IYrezsXmPcfHPE/av4twRSxf6WE+TkZgLY+71vCZbiIiFKA/w==", + "dev": true, + "requires": { + "@jsonjoy.com/base64": "17.67.0", + "@jsonjoy.com/buffers": "17.67.0", + "@jsonjoy.com/codegen": "17.67.0", + "@jsonjoy.com/json-pointer": "17.67.0", + "@jsonjoy.com/util": "17.67.0", + "hyperdyperid": "^1.2.0", + "thingies": "^2.5.0", + "tree-dump": "^1.1.0" + } + }, + "@jsonjoy.com/json-pointer": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pointer/-/json-pointer-17.67.0.tgz", + "integrity": "sha512-+iqOFInH+QZGmSuaybBUNdh7yvNrXvqR+h3wjXm0N/3JK1EyyFAeGJvqnmQL61d1ARLlk/wJdFKSL+LHJ1eaUA==", + "dev": true, + "requires": { + "@jsonjoy.com/util": "17.67.0" + } + }, + "@jsonjoy.com/util": { + "version": "17.67.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/util/-/util-17.67.0.tgz", + "integrity": "sha512-6+8xBaz1rLSohlGh68D1pdw3AwDi9xydm8QNlAFkvnavCJYSze+pxoW2VKP8p308jtlMRLs5NTHfPlZLd4w7ew==", + "dev": true, + "requires": { + "@jsonjoy.com/buffers": "17.67.0", + "@jsonjoy.com/codegen": "17.67.0" + } + } + } + }, "@jsonjoy.com/json-pack": { - "version": "1.2.0", - "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pack/-/json-pack-1.2.0.tgz", - "integrity": "sha512-io1zEbbYcElht3tdlqEOFxZ0dMTYrHz9iMf0gqn1pPjZFTCgM5R4R5IMA20Chb2UPYYsxjzs8CgZ7Nb5n2K2rA==", + "version": "1.21.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pack/-/json-pack-1.21.0.tgz", + "integrity": "sha512-+AKG+R2cfZMShzrF2uQw34v3zbeDYUqnQ+jg7ORic3BGtfw9p/+N6RJbq/kkV8JmYZaINknaEQ2m0/f693ZPpg==", "dev": true, "requires": { - "@jsonjoy.com/base64": "^1.1.1", - "@jsonjoy.com/util": "^1.1.2", + "@jsonjoy.com/base64": "^1.1.2", + "@jsonjoy.com/buffers": "^1.2.0", + "@jsonjoy.com/codegen": "^1.0.0", + "@jsonjoy.com/json-pointer": "^1.0.2", + "@jsonjoy.com/util": "^1.9.0", "hyperdyperid": "^1.2.0", - "thingies": "^1.20.0" + "thingies": "^2.5.0", + "tree-dump": "^1.1.0" + }, + "dependencies": { + "@jsonjoy.com/buffers": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/buffers/-/buffers-1.2.1.tgz", + "integrity": "sha512-12cdlDwX4RUM3QxmUbVJWqZ/mrK6dFQH4Zxq6+r1YXKXYBNgZXndx2qbCJwh3+WWkCSn67IjnlG3XYTvmvYtgA==", + "dev": true, + "requires": {} + } + } + }, + "@jsonjoy.com/json-pointer": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/json-pointer/-/json-pointer-1.0.2.tgz", + "integrity": "sha512-Fsn6wM2zlDzY1U+v4Nc8bo3bVqgfNTGcn6dMgs6FjrEnt4ZCe60o6ByKRjOGlI2gow0aE/Q41QOigdTqkyK5fg==", + "dev": true, + "requires": { + "@jsonjoy.com/codegen": "^1.0.0", + "@jsonjoy.com/util": "^1.9.0" } }, "@jsonjoy.com/util": { - "version": "1.6.0", - "resolved": "https://registry.npmjs.org/@jsonjoy.com/util/-/util-1.6.0.tgz", - "integrity": "sha512-sw/RMbehRhN68WRtcKCpQOPfnH6lLP4GJfqzi3iYej8tnzpZUDr6UkZYJjcjjC0FWEJOJbyM3PTIwxucUmDG2A==", + "version": "1.9.0", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/util/-/util-1.9.0.tgz", + "integrity": "sha512-pLuQo+VPRnN8hfPqUTLTHk126wuYdXVxE6aDmjSeV4NCAgyxWbiOIeNJVtID3h1Vzpoi9m4jXezf73I6LgabgQ==", + "dev": true, + "requires": { + "@jsonjoy.com/buffers": "^1.0.0", + "@jsonjoy.com/codegen": "^1.0.0" + }, + "dependencies": { + "@jsonjoy.com/buffers": { + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/@jsonjoy.com/buffers/-/buffers-1.2.1.tgz", + "integrity": "sha512-12cdlDwX4RUM3QxmUbVJWqZ/mrK6dFQH4Zxq6+r1YXKXYBNgZXndx2qbCJwh3+WWkCSn67IjnlG3XYTvmvYtgA==", + "dev": true, + "requires": {} + } + } + }, + "glob-to-regex.js": { + "version": "1.2.0", + "resolved": "https://registry.npmjs.org/glob-to-regex.js/-/glob-to-regex.js-1.2.0.tgz", + "integrity": "sha512-QMwlOQKU/IzqMUOAZWubUOT8Qft+Y0KQWnX9nK3ch0CJg0tTp4TvGZsTfudYKv2NzoQSyPcnA6TYeIQ3jGichQ==", "dev": true, "requires": {} }, "thingies": { - "version": "1.21.0", - "resolved": "https://registry.npmjs.org/thingies/-/thingies-1.21.0.tgz", - "integrity": "sha512-hsqsJsFMsV+aD4s3CWKk85ep/3I9XzYV/IXaSouJMYIoDlgyi11cBhsqYe9/geRfB0YIikBQg6raRaM+nIMP9g==", + "version": "2.6.0", + "resolved": "https://registry.npmjs.org/thingies/-/thingies-2.6.0.tgz", + "integrity": "sha512-rMHRjmlFLM1R96UYPvpmnc3LYtdFrT33JIB7L9hetGue1qAPfn1N2LJeEjxUSidu1Iku+haLZXDuEXUHNGO/lg==", "dev": true, "requires": {} }, "tree-dump": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/tree-dump/-/tree-dump-1.0.3.tgz", - "integrity": "sha512-il+Cv80yVHFBwokQSfd4bldvr1Md951DpgAGfmhydt04L+YzHgubm2tQ7zueWDcGENKHq0ZvGFR/hjvNXilHEg==", + "version": "1.1.0", + "resolved": "https://registry.npmjs.org/tree-dump/-/tree-dump-1.1.0.tgz", + "integrity": "sha512-rMuvhU4MCDbcbnleZTFezWsaZXRFemSqAM+7jPnzUl1fo9w3YEKOxAeui0fz3OI4EU4hf23iyA7uQRVko+UaBA==", "dev": true, "requires": {} }, @@ -6150,9 +6030,9 @@ } }, "merge-descriptors": { - "version": "1.0.3", - "resolved": "https://registry.npmjs.org/merge-descriptors/-/merge-descriptors-1.0.3.tgz", - "integrity": "sha512-gaNvAS7TZ897/rVaZ0nMtAyxNyi/pdbjbAwUpFQpN70GqnVfOiXpeUUMKRBmzXaSQ8DdTX4/0ms62r2K+hE6mQ==", + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/merge-descriptors/-/merge-descriptors-2.0.0.tgz", + "integrity": "sha512-Snk314V5ayFLhp3fkUREub6WtjBfPdCPY1Ln8/8munuLuiYhsABgBVWsozAG+MWMbVEvcdcpbi9R7ww22l9Q3g==", "dev": true }, "merge-stream": { @@ -6161,12 +6041,6 @@ "integrity": "sha512-abv/qOcuPfk3URPfDzmZU1LKmuw8kT+0nIHvKrKgFrwifol/doWcdA4ZqsWQ8ENrFKkd67Mfpo/LovbIUsbt3w==", "dev": true }, - "methods": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/methods/-/methods-1.1.2.tgz", - "integrity": "sha512-iclAHeNqNm68zFtnZ0e+1L2yUIdvzNoauKU4WBA3VvH/vPFieF7qfRlwUZU+DA9P9bPXIS90ulxoUoCH23sV2w==", - "dev": true - }, "micromatch": { "version": "4.0.8", "resolved": "https://registry.npmjs.org/micromatch/-/micromatch-4.0.8.tgz", @@ -6177,12 +6051,6 @@ "picomatch": "^2.3.1" } }, - "mime": { - "version": "1.6.0", - "resolved": "https://registry.npmjs.org/mime/-/mime-1.6.0.tgz", - "integrity": "sha512-x0Vn8spI+wuJ1O6S7gnbaQg8Pxh4NNHb7KSINmEWKiPE4RKOplvijn+NkmYmmRgP68mc70j2EbeTFRsrswaQeg==", - "dev": true - }, "mime-db": { "version": "1.52.0", "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.52.0.tgz", @@ -6198,12 +6066,6 @@ "mime-db": "1.52.0" } }, - "minimalistic-assert": { - "version": "1.0.1", - "resolved": "https://registry.npmjs.org/minimalistic-assert/-/minimalistic-assert-1.0.1.tgz", - "integrity": "sha512-UtJcAD4yEaGtjPezWuO9wC4nwUnVH/8/Im3yEHQP4b67cXlD/Qr9hdITCU1xDbSEXg2XKNaP8jsReV7vQd00/A==", - "dev": true - }, "ms": { "version": "2.0.0", "resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz", @@ -6221,9 +6083,9 @@ } }, "negotiator": { - "version": "0.6.3", - "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-0.6.3.tgz", - "integrity": "sha512-+EUsqGPLsM+j/zdChZjsnX51g4XrHFOIXwfnCVPGlQk/k5giakcKsuxCObBRu6DSm9opw/O6slWbJdghQM4bBg==", + "version": "1.0.0", + "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-1.0.0.tgz", + "integrity": "sha512-8Ofs/AUQh8MaEcrlq5xOX0CQ9ypTF5dl78mjlMNfOK08fzpgTHQRQPBxcPlEtIw0yRpws+Zo/3r+5WRby7u3Gg==", "dev": true }, "neo-async": { @@ -6250,12 +6112,6 @@ "integrity": "sha512-W67iLl4J2EXEGTbfeHCffrjDfitvLANg0UlX3wFUUSTx92KXRFegMHUVgSqE+wvhAbi4WqjGg9czysTV2Epbew==", "dev": true }, - "obuf": { - "version": "1.1.2", - "resolved": "https://registry.npmjs.org/obuf/-/obuf-1.1.2.tgz", - "integrity": "sha512-PX1wu0AmAdPqOL1mWhqmlOd8kOIZQwGZw6rh7uby9fTc5lhaOWFLX3I6R1hrF9k3zUY40e6igsLGkDXK92LJNg==", - "dev": true - }, "on-finished": { "version": "2.4.1", "resolved": "https://registry.npmjs.org/on-finished/-/on-finished-2.4.1.tgz", @@ -6271,16 +6127,27 @@ "integrity": "sha512-737ZY3yNnXy37FHkQxPzt4UZ2UWPWiCZWLvFZ4fu5cueciegX0zGPnrlY6bwRg4FdQOe9YU8MkmJwGhoMybl8A==", "dev": true }, + "once": { + "version": "1.4.0", + "resolved": "https://registry.npmjs.org/once/-/once-1.4.0.tgz", + "integrity": "sha512-lNaJgI+2Q5URQBkccEKHTQOPaXdUxnZZElQTZY0MFUAuaEqe1E+Nyvgdz/aIyNi6Z9MzO5dv1H8n58/GELp3+w==", + "dev": true, + "requires": { + "wrappy": "1" + } + }, "open": { - "version": "10.1.2", - "resolved": "https://registry.npmjs.org/open/-/open-10.1.2.tgz", - "integrity": "sha512-cxN6aIDPz6rm8hbebcP7vrQNhvRcveZoJU72Y7vskh4oIm+BZwBECnx5nTmrlres1Qapvx27Qo1Auukpf8PKXw==", + "version": "11.0.0", + "resolved": "https://registry.npmjs.org/open/-/open-11.0.0.tgz", + "integrity": "sha512-smsWv2LzFjP03xmvFoJ331ss6h+jixfA4UUV/Bsiyuu4YJPfN+FIQGOIiv4w9/+MoHkfkJ22UIaQWRVFRfH6Vw==", "dev": true, "requires": { - "default-browser": "^5.2.1", + "default-browser": "^5.4.0", "define-lazy-prop": "^3.0.0", + "is-in-ssh": "^1.0.0", "is-inside-container": "^1.0.0", - "is-wsl": "^3.1.0" + "powershell-utils": "^0.1.0", + "wsl-utils": "^0.3.0" } }, "p-locate": { @@ -6304,14 +6171,12 @@ } }, "p-retry": { - "version": "6.2.1", - "resolved": "https://registry.npmjs.org/p-retry/-/p-retry-6.2.1.tgz", - "integrity": "sha512-hEt02O4hUct5wtwg4H4KcWgDdm+l1bOaEy/hWzd8xtXB9BqxTWBBhb+2ImAtH4Cv4rPjV76xN3Zumqk3k3AhhQ==", + "version": "8.0.0", + "resolved": "https://registry.npmjs.org/p-retry/-/p-retry-8.0.0.tgz", + "integrity": "sha512-kFVqH1HxOHp8LupNsOys7bSV09VYTRLxarH/mokO4Rqhk6wGi70E0jh4VzvVGXfEVNggHoHLAMWsQqHyU1Ey9A==", "dev": true, "requires": { - "@types/retry": "0.12.2", - "is-network-error": "^1.0.0", - "retry": "^0.13.1" + "is-network-error": "^1.3.0" } }, "p-try": { @@ -6345,9 +6210,9 @@ "dev": true }, "path-to-regexp": { - "version": "0.1.12", - "resolved": "https://registry.npmjs.org/path-to-regexp/-/path-to-regexp-0.1.12.tgz", - "integrity": "sha512-RA1GjUVMnvYFxuqovrEqZoxxW5NUZqbwKtYz/Tt7nXerk0LbLblQmrsgdeOxV5SFHf0UDggjS/bSeOZwt1pmEQ==", + "version": "8.4.2", + "resolved": "https://registry.npmjs.org/path-to-regexp/-/path-to-regexp-8.4.2.tgz", + "integrity": "sha512-qRcuIdP69NPm4qbACK+aDogI5CBDMi1jKe0ry5rSQJz8JVLsC7jV8XpiJjGRLLol3N+R5ihGYcrPLTno6pAdBA==", "dev": true }, "picocolors": { @@ -6393,10 +6258,10 @@ } } }, - "process-nextick-args": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/process-nextick-args/-/process-nextick-args-2.0.0.tgz", - "integrity": "sha512-MtEC1TqN0EU5nephaJ4rAtThHtC86dNN9qCuEhtshvpVBkAW5ZO7BASN9REnF9eoXGcRub+pFuKEpOHE+HbEMw==", + "powershell-utils": { + "version": "0.1.0", + "resolved": "https://registry.npmjs.org/powershell-utils/-/powershell-utils-0.1.0.tgz", + "integrity": "sha512-dM0jVuXJPsDN6DvRpea484tCUaMiXWjuCn++HGTqUWzGDjv5tZkEZldAJ/UMlqRYGFrD/etByo4/xOuC/snX2A==", "dev": true }, "proxy-addr": { @@ -6441,88 +6306,38 @@ "dev": true }, "qs": { - "version": "6.14.2", - "resolved": "https://registry.npmjs.org/qs/-/qs-6.14.2.tgz", - "integrity": "sha512-V/yCWTTF7VJ9hIh18Ugr2zhJMP01MY7c5kh4J870L7imm6/DIzBsNLTXzMwUA3yZ5b/KBqLx8Kp3uRvd7xSe3Q==", + "version": "6.15.3", + "resolved": "https://registry.npmjs.org/qs/-/qs-6.15.3.tgz", + "integrity": "sha512-O9gl3zCl5h5blw1KGUzQKhA5oUXSl8rwUIM5o0S3nCXMliSvy5Dzx7/DJcI+SwgICv+IneSZwhBh1oSyEHA71A==", "dev": true, "requires": { - "side-channel": "^1.1.0" + "es-define-property": "^1.0.1", + "side-channel": "^1.1.1" } }, "range-parser": { - "version": "1.2.1", - "resolved": "https://registry.npmjs.org/range-parser/-/range-parser-1.2.1.tgz", - "integrity": "sha512-Hrgsx+orqoygnmhFbKaHE6c296J+HTAQXoxEF6gNupROmmGJRoyzfG3ccAveqCBrwr/2yxQ5BVd/GTl5agOwSg==", + "version": "1.3.0", + "resolved": "https://registry.npmjs.org/range-parser/-/range-parser-1.3.0.tgz", + "integrity": "sha512-hek2mFQpPuI4E1BBKrSto+BU3e3x4xuarsbiwr3+lf7p44juvFMV0XFWQAP3xUyqXA4RrXLIoaSUGbSt056ZMw==", "dev": true }, "raw-body": { - "version": "2.5.3", - "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-2.5.3.tgz", - "integrity": "sha512-s4VSOf6yN0rvbRZGxs8Om5CWj6seneMwK3oDb4lWDH0UPhWcxwOWw5+qk24bxq87szX1ydrwylIOp2uG1ojUpA==", + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/raw-body/-/raw-body-3.0.2.tgz", + "integrity": "sha512-K5zQjDllxWkf7Z5xJdV0/B0WTNqx6vxG70zJE4N0kBs4LovmEYWJzQGxC9bS9RAKu3bgM40lrd5zoLJ12MQ5BA==", "dev": true, "requires": { "bytes": "~3.1.2", "http-errors": "~2.0.1", - "iconv-lite": "~0.4.24", + "iconv-lite": "~0.7.0", "unpipe": "~1.0.0" - }, - "dependencies": { - "depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", - "dev": true - }, - "http-errors": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-2.0.1.tgz", - "integrity": "sha512-4FbRdAX+bSdmo4AUFuS0WNiPz8NgFt+r8ThgNWmlrjQjt1Q7ZR9+zTlce2859x4KSXrwIsaeTqDoKQmtP8pLmQ==", - "dev": true, - "requires": { - "depd": "~2.0.0", - "inherits": "~2.0.4", - "setprototypeof": "~1.2.0", - "statuses": "~2.0.2", - "toidentifier": "~1.0.1" - } - }, - "inherits": { - "version": "2.0.4", - "resolved": "https://registry.npmjs.org/inherits/-/inherits-2.0.4.tgz", - "integrity": "sha512-k/vGaX4/Yla3WzyMCvTQOXYeIHvqOKtnqBduzTHpzpQZzAskKMhZ2K+EnBiSM9zGSoIFeMpXKxa4dYeZIQqewQ==", - "dev": true - }, - "statuses": { - "version": "2.0.2", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", - "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", - "dev": true - } - } - }, - "readable-stream": { - "version": "2.3.6", - "resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-2.3.6.tgz", - "integrity": "sha512-tQtKA9WIAhBF3+VLAseyMqZeBjW0AHJoxOtYqSUZNJxauErmLbVm2FW1y+J/YA9dUrAC39ITejlZWhVIwawkKw==", - "dev": true, - "requires": { - "core-util-is": "~1.0.0", - "inherits": "~2.0.3", - "isarray": "~1.0.0", - "process-nextick-args": "~2.0.0", - "safe-buffer": "~5.1.1", - "string_decoder": "~1.1.1", - "util-deprecate": "~1.0.1" } }, "readdirp": { - "version": "3.6.0", - "resolved": "https://registry.npmjs.org/readdirp/-/readdirp-3.6.0.tgz", - "integrity": "sha512-hOS089on8RduqdbhvQ5Z37A0ESjsqz6qnRcffsMU3495FuTdqSm+7bhJ29JvIOsBDEEnan5DPu9t3To9VRlMzA==", - "dev": true, - "requires": { - "picomatch": "^2.2.1" - } + "version": "5.0.0", + "resolved": "https://registry.npmjs.org/readdirp/-/readdirp-5.0.0.tgz", + "integrity": "sha512-9u/XQ1pvrQtYyMpZe7DXKv2p5CNvyVwzUB6uhLAnQwHMSgKMBR62lc7AHljaeteeHXn11XTAaLLUVZYVZyuRBQ==", + "dev": true }, "rechoir": { "version": "0.8.0", @@ -6545,12 +6360,6 @@ "integrity": "sha512-Xf0nWe6RseziFMu+Ap9biiUbmplq6S9/p+7w7YXP/JBHhrUDDUhwa+vANyubuqfZWTveU//DYVGsDG7RKL/vEw==", "dev": true }, - "requires-port": { - "version": "1.0.0", - "resolved": "https://registry.npmjs.org/requires-port/-/requires-port-1.0.0.tgz", - "integrity": "sha512-KigOCHcocU3XODJxsu8i/j8T9tzT4adHiecwORRQ0ZZFcp7ahwXuRU1m+yuO90C5ZUyGeGfocHDI14M3L3yDAQ==", - "dev": true - }, "resolve": { "version": "1.22.6", "resolved": "https://registry.npmjs.org/resolve/-/resolve-1.22.6.tgz", @@ -6577,22 +6386,23 @@ "integrity": "sha512-qYg9KP24dD5qka9J47d0aVky0N+b4fTU89LN9iDnjB5waksiC49rvMB0PrUJQGoTmH50XPiqOvAjDfaijGxYZw==", "dev": true }, - "retry": { - "version": "0.13.1", - "resolved": "https://registry.npmjs.org/retry/-/retry-0.13.1.tgz", - "integrity": "sha512-XQBQ3I8W1Cge0Seh+6gjj03LbmRFWuoszgK9ooCpwYIrhhoO80pfq4cUkU5DkknwfOfFteRwlZ56PYOGYyFWdg==", - "dev": true + "router": { + "version": "2.2.0", + "resolved": "https://registry.npmjs.org/router/-/router-2.2.0.tgz", + "integrity": "sha512-nLTrUKm2UyiL7rlhapu/Zl45FwNgkZGaCpZbIHajDYgwlJCOzLSk+cIPAnsEqV955GjILJnKbdQC1nVPz+gAYQ==", + "dev": true, + "requires": { + "debug": "^4.4.0", + "depd": "^2.0.0", + "is-promise": "^4.0.0", + "parseurl": "^1.3.3", + "path-to-regexp": "^8.0.0" + } }, "run-applescript": { - "version": "7.0.0", - "resolved": "https://registry.npmjs.org/run-applescript/-/run-applescript-7.0.0.tgz", - "integrity": "sha512-9by4Ij99JUr/MCFBUkDKLWK3G9HVXmabKz9U5MlIAIuvuzkiOicRYs8XJLxX+xahD+mLiiCYDqF9dKAgtzKP1A==", - "dev": true - }, - "safe-buffer": { - "version": "5.1.2", - "resolved": "https://registry.npmjs.org/safe-buffer/-/safe-buffer-5.1.2.tgz", - "integrity": "sha512-Gd2UZBJDkXlY7GbJxfsE8/nvKkUEU1G38c1siN6QP6a9PT9MmHB8GnpscSmMJSoF8LOIrt8ud/wPtojys4G6+g==", + "version": "7.1.0", + "resolved": "https://registry.npmjs.org/run-applescript/-/run-applescript-7.1.0.tgz", + "integrity": "sha512-DPe5pVFaAsinSaV6QjQ6gdiedWDcRCbUuiQfQa2wmWV7+xC9bGulGI8+TdRmoFkAPaBXk8CrAbnlY2ISniJ47Q==", "dev": true }, "safer-buffer": { @@ -6613,12 +6423,6 @@ "ajv-keywords": "^5.1.0" } }, - "select-hose": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/select-hose/-/select-hose-2.0.0.tgz", - "integrity": "sha1-Yl2GWPhlr0Psliv8N2o3NZpJlMo=", - "dev": true - }, "selfsigned": { "version": "5.5.0", "resolved": "https://registry.npmjs.org/selfsigned/-/selfsigned-5.5.0.tgz", @@ -6630,84 +6434,72 @@ } }, "send": { - "version": "0.19.0", - "resolved": "https://registry.npmjs.org/send/-/send-0.19.0.tgz", - "integrity": "sha512-dW41u5VfLXu8SJh5bwRmyYUbAoSB3c9uQh6L8h/KtsFREPWpbX1lrljJo186Jc4nmci/sGUZ9a0a0J2zgfq2hw==", - "dev": true, - "requires": { - "debug": "2.6.9", - "depd": "2.0.0", - "destroy": "1.2.0", - "encodeurl": "~1.0.2", - "escape-html": "~1.0.3", - "etag": "~1.8.1", - "fresh": "0.5.2", - "http-errors": "2.0.0", - "mime": "1.6.0", - "ms": "2.1.3", - "on-finished": "2.4.1", - "range-parser": "~1.2.1", - "statuses": "2.0.1" + "version": "1.2.1", + "resolved": "https://registry.npmjs.org/send/-/send-1.2.1.tgz", + "integrity": "sha512-1gnZf7DFcoIcajTjTwjwuDjzuz4PPcY2StKPlsGAQ1+YH20IRVrBaXSWmdjowTJ6u8Rc01PoYOGHXfP1mYcZNQ==", + "dev": true, + "requires": { + "debug": "^4.4.3", + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "etag": "^1.8.1", + "fresh": "^2.0.0", + "http-errors": "^2.0.1", + "mime-types": "^3.0.2", + "ms": "^2.1.3", + "on-finished": "^2.4.1", + "range-parser": "^1.2.1", + "statuses": "^2.0.2" }, "dependencies": { - "debug": { - "version": "2.6.9", - "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", - "integrity": "sha512-bC7ElrdJaJnPbAP+1EotYvqZsb3ecl5wi6Bfi6BJTUcNowp6cvspg0jXznRTKDjm/E7AdgFBVeAPVMNcKGsHMA==", + "mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true + }, + "mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", "dev": true, "requires": { - "ms": "2.0.0" - }, - "dependencies": { - "ms": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/ms/-/ms-2.0.0.tgz", - "integrity": "sha512-Tpp60P6IUJDTuOq/5Z8cdskzJujfwqfOTkrwIwj7IRISpnkJnT6SyJ4PCPnGMoFjC9ddhal5KVIYtAt97ix05A==", - "dev": true - } + "mime-db": "^1.54.0" } }, - "depd": { - "version": "2.0.0", - "resolved": "https://registry.npmjs.org/depd/-/depd-2.0.0.tgz", - "integrity": "sha512-g7nH6P6dyDioJogAAGprGpCtVImJhpPk/roCzdb3fIh61/s/nPsfR6onyMwkCAR/OlC3yBC0lESvUoQEAssIrw==", - "dev": true - }, - "encodeurl": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/encodeurl/-/encodeurl-1.0.2.tgz", - "integrity": "sha512-TPJXq8JqFaVYm2CWmPvnP2Iyo4ZSM7/QKcSmuMLDObfpH5fi7RUGmd/rTDf+rut/saiDiQEeVTNgAmJEdAOx0w==", - "dev": true - }, "ms": { "version": "2.1.3", "resolved": "https://registry.npmjs.org/ms/-/ms-2.1.3.tgz", "integrity": "sha512-6FlzubTLZG3J2a/NVCAleEhjzq5oxgHyaCU9yYXvcLsvoVaHJq/s5xXI6/XXP6tz7R9xAOtHnSO/tXtF3WRTlA==", "dev": true - }, - "statuses": { - "version": "2.0.1", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.1.tgz", - "integrity": "sha512-RwNA9Z/7PrK06rYLIzFMlaF+l73iwpzsqRIFgbMLbTcLD6cOao82TaWefPXQvB2fOC4AjuYSEndS7N/mTCbkdQ==", - "dev": true } } }, "serve-index": { - "version": "1.9.1", - "resolved": "https://registry.npmjs.org/serve-index/-/serve-index-1.9.1.tgz", - "integrity": "sha1-03aNabHn2C5c4FD/9bRTvqEqkjk=", + "version": "1.9.2", + "resolved": "https://registry.npmjs.org/serve-index/-/serve-index-1.9.2.tgz", + "integrity": "sha512-KDj11HScOaLmrPxl70KYNW1PksP4Nb/CLL2yvC+Qd2kHMPEEpfc4Re2e4FOay+bC/+XQl/7zAcWON3JVo5v3KQ==", "dev": true, "requires": { - "accepts": "~1.3.4", + "accepts": "~1.3.8", "batch": "0.6.1", "debug": "2.6.9", "escape-html": "~1.0.3", - "http-errors": "~1.6.2", - "mime-types": "~2.1.17", - "parseurl": "~1.3.2" + "http-errors": "~1.8.0", + "mime-types": "~2.1.35", + "parseurl": "~1.3.3" }, "dependencies": { + "accepts": { + "version": "1.3.8", + "resolved": "https://registry.npmjs.org/accepts/-/accepts-1.3.8.tgz", + "integrity": "sha512-PYAthTa2m2VKxuvSD3DPC/Gy+U+sOA1LAuT8mkmRuvw+NACSaeXEQ+NHcVF7rONl6qcaxV3Uuemwawk+7+SJLw==", + "dev": true, + "requires": { + "mime-types": "~2.1.34", + "negotiator": "0.6.3" + } + }, "debug": { "version": "2.6.9", "resolved": "https://registry.npmjs.org/debug/-/debug-2.6.9.tgz", @@ -6717,36 +6509,49 @@ "ms": "2.0.0" } }, + "depd": { + "version": "1.1.2", + "resolved": "https://registry.npmjs.org/depd/-/depd-1.1.2.tgz", + "integrity": "sha512-7emPTl6Dpo6JRXOXjLRxck+FlLRX5847cLKEn00PLAgc3g2hTZZgr+e4c2v6QpSmLeFP3n5yUo7ft6avBK/5jQ==", + "dev": true + }, "http-errors": { - "version": "1.6.3", - "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-1.6.3.tgz", - "integrity": "sha1-i1VoC7S+KDoLW/TqLjhYC+HZMg0=", + "version": "1.8.1", + "resolved": "https://registry.npmjs.org/http-errors/-/http-errors-1.8.1.tgz", + "integrity": "sha512-Kpk9Sm7NmI+RHhnj6OIWDI1d6fIoFAtFt9RLaTMRlg/8w49juAStsrBgp0Dp4OdxdVbRIeKhtCUvoi/RuAhO4g==", "dev": true, "requires": { "depd": "~1.1.2", - "inherits": "2.0.3", - "setprototypeof": "1.1.0", - "statuses": ">= 1.4.0 < 2" + "inherits": "2.0.4", + "setprototypeof": "1.2.0", + "statuses": ">= 1.5.0 < 2", + "toidentifier": "1.0.1" } }, - "setprototypeof": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/setprototypeof/-/setprototypeof-1.1.0.tgz", - "integrity": "sha512-BvE/TwpZX4FXExxOxZyRGQQv651MSwmWKZGqvmPcRIjDqWub67kTKuIMx43cZZrS/cBBzwBcNDWoFxt2XEFIpQ==", + "negotiator": { + "version": "0.6.3", + "resolved": "https://registry.npmjs.org/negotiator/-/negotiator-0.6.3.tgz", + "integrity": "sha512-+EUsqGPLsM+j/zdChZjsnX51g4XrHFOIXwfnCVPGlQk/k5giakcKsuxCObBRu6DSm9opw/O6slWbJdghQM4bBg==", + "dev": true + }, + "statuses": { + "version": "1.5.0", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-1.5.0.tgz", + "integrity": "sha512-OpZ3zP+jT1PI7I8nemJX4AKmAX070ZkYPVWV/AaKTJl+tXCTGyVdC1a4SL8RUQYEwk/f34ZX8UTykN68FwrqAA==", "dev": true } } }, "serve-static": { - "version": "1.16.2", - "resolved": "https://registry.npmjs.org/serve-static/-/serve-static-1.16.2.tgz", - "integrity": "sha512-VqpjJZKadQB/PEbEwvFdO43Ax5dFBZ2UECszz8bQ7pi7wt//PWe1P6MN7eCnjsatYtBT6EuiClbjSWP2WrIoTw==", + "version": "2.2.1", + "resolved": "https://registry.npmjs.org/serve-static/-/serve-static-2.2.1.tgz", + "integrity": "sha512-xRXBn0pPqQTVQiC8wyQrKs2MOlX24zQ0POGaj0kultvoOCstBQM5yvOhAVSUwOMjQtTvsPWoNCHfPGwaaQJhTw==", "dev": true, "requires": { - "encodeurl": "~2.0.0", - "escape-html": "~1.0.3", - "parseurl": "~1.3.3", - "send": "0.19.0" + "encodeurl": "^2.0.0", + "escape-html": "^1.0.3", + "parseurl": "^1.3.3", + "send": "^1.2.0" } }, "setprototypeof": { @@ -6780,32 +6585,32 @@ "dev": true }, "shell-quote": { - "version": "1.8.3", - "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.8.3.tgz", - "integrity": "sha512-ObmnIF4hXNg1BqhnHmgbDETF8dLPCggZWBjkQfhZpbszZnYur5DUljTcCHii5LC3J5E0yeO/1LIMyH+UvHQgyw==", + "version": "1.10.0", + "resolved": "https://registry.npmjs.org/shell-quote/-/shell-quote-1.10.0.tgz", + "integrity": "sha512-w1aiOKwKuRgtwAReIIj89puqg+I7GvX4IbLrvmhXbzQsj1+Zwi4VO3+fa6ZF91TWSjIxoEkKnMeHcLEODK5ZXA==", "dev": true }, "side-channel": { - "version": "1.1.0", - "resolved": "https://registry.npmjs.org/side-channel/-/side-channel-1.1.0.tgz", - "integrity": "sha512-ZX99e6tRweoUXqR+VBrslhda51Nh5MTQwou5tnUDgbtyM0dBgmhEDtWGP/xbKn6hqfPRHujUNwz5fy/wbbhnpw==", + "version": "1.1.1", + "resolved": "https://registry.npmjs.org/side-channel/-/side-channel-1.1.1.tgz", + "integrity": "sha512-6x6dK6zJdpTzF4sQeNYxwtvBzf6Eg4GtlesS94HOvTudUeyK2WXAaIfmDgsyslYrRBeFIlsi54AYsFGUuhmvrQ==", "dev": true, "requires": { "es-errors": "^1.3.0", - "object-inspect": "^1.13.3", - "side-channel-list": "^1.0.0", + "object-inspect": "^1.13.4", + "side-channel-list": "^1.0.1", "side-channel-map": "^1.0.1", "side-channel-weakmap": "^1.0.2" } }, "side-channel-list": { - "version": "1.0.0", - "resolved": "https://registry.npmjs.org/side-channel-list/-/side-channel-list-1.0.0.tgz", - "integrity": "sha512-FCLHtRD/gnpCiCHEiJLOwdmFP+wzCmDEkc9y7NsYxeF4u7Btsn1ZuwgwJGxImImHicJArLP4R0yX4c2KCrMrTA==", + "version": "1.0.1", + "resolved": "https://registry.npmjs.org/side-channel-list/-/side-channel-list-1.0.1.tgz", + "integrity": "sha512-mjn/0bi/oUURjc5Xl7IaWi/OJJJumuoJFQJfDDyO46+hBWsfaVM65TBHq2eoZBhzl9EchxOijpkbRC8SVBQU0w==", "dev": true, "requires": { "es-errors": "^1.3.0", - "object-inspect": "^1.13.3" + "object-inspect": "^1.13.4" } }, "side-channel-map": { @@ -6833,17 +6638,6 @@ "side-channel-map": "^1.0.1" } }, - "sockjs": { - "version": "0.3.24", - "resolved": "https://registry.npmjs.org/sockjs/-/sockjs-0.3.24.tgz", - "integrity": "sha512-GJgLTZ7vYb/JtPSSZ10hsOYIvEYsjbNU+zPdIHcUaWVNUEPivzxku31865sSSud0Da0W4lEeOPlmw93zLQchuQ==", - "dev": true, - "requires": { - "faye-websocket": "^0.11.3", - "uuid": "^8.3.2", - "websocket-driver": "^0.7.4" - } - }, "source-map": { "version": "0.6.1", "resolved": "https://registry.npmjs.org/source-map/-/source-map-0.6.1.tgz", @@ -6860,61 +6654,12 @@ "source-map": "^0.6.0" } }, - "spdy": { - "version": "4.0.2", - "resolved": "https://registry.npmjs.org/spdy/-/spdy-4.0.2.tgz", - "integrity": "sha512-r46gZQZQV+Kl9oItvl1JZZqJKGr+oEkB08A6BzkiR7593/7IbtuncXHd2YoYeTsG4157ZssMu9KYvUHLcjcDoA==", - "dev": true, - "requires": { - "debug": "^4.1.0", - "handle-thing": "^2.0.0", - "http-deceiver": "^1.2.7", - "select-hose": "^2.0.0", - "spdy-transport": "^3.0.0" - } - }, - "spdy-transport": { - "version": "3.0.0", - "resolved": "https://registry.npmjs.org/spdy-transport/-/spdy-transport-3.0.0.tgz", - "integrity": "sha512-hsLVFE5SjA6TCisWeJXFKniGGOpBgMLmerfO2aCyCU5s7nJ/rpAepqmFifv/GCbSbueEeAJJnmSQ2rKC/g8Fcw==", - "dev": true, - "requires": { - "debug": "^4.1.0", - "detect-node": "^2.0.4", - "hpack.js": "^2.1.6", - "obuf": "^1.1.2", - "readable-stream": "^3.0.6", - "wbuf": "^1.7.3" - }, - "dependencies": { - "readable-stream": { - "version": "3.6.0", - "resolved": "https://registry.npmjs.org/readable-stream/-/readable-stream-3.6.0.tgz", - "integrity": "sha512-BViHy7LKeTz4oNnkcLJ+lVSL6vpiFeX6/d3oSH8zCW7UxP2onchk+vTGB143xuFjHS3deTgkKoXXymXqymiIdA==", - "dev": true, - "requires": { - "inherits": "^2.0.3", - "string_decoder": "^1.1.1", - "util-deprecate": "^1.0.1" - } - } - } - }, "statuses": { - "version": "1.5.0", - "resolved": "https://registry.npmjs.org/statuses/-/statuses-1.5.0.tgz", - "integrity": "sha1-Fhx9rBd2Wf2YEfQ3cfqZOBR4Yow=", + "version": "2.0.2", + "resolved": "https://registry.npmjs.org/statuses/-/statuses-2.0.2.tgz", + "integrity": "sha512-DvEy55V3DB7uknRo+4iOGT5fP1slR8wQohVdknigZPMpMstaKJQWhwiYBACJE3Ul2pTnATihhBYnRhZQHGBiRw==", "dev": true }, - "string_decoder": { - "version": "1.1.1", - "resolved": "https://registry.npmjs.org/string_decoder/-/string_decoder-1.1.1.tgz", - "integrity": "sha512-n/ShnvDi6FHbbVfviro+WojiFzv+s8MPMHBczVePfUpDJLwoLT0ht1l4YwBCbi8pJAveEEdnkHyPyTP/mzRfwg==", - "dev": true, - "requires": { - "safe-buffer": "~5.1.0" - } - }, "supports-color": { "version": "8.1.1", "resolved": "https://registry.npmjs.org/supports-color/-/supports-color-8.1.1.tgz", @@ -7022,13 +6767,37 @@ } }, "type-is": { - "version": "1.6.18", - "resolved": "https://registry.npmjs.org/type-is/-/type-is-1.6.18.tgz", - "integrity": "sha512-TkRKr9sUTxEH8MdfuCSP7VizJyzRNMjj2J2do2Jr3Kym598JVdEksuzPQCnlFPW4ky9Q+iA+ma9BGm06XQBy8g==", + "version": "2.1.0", + "resolved": "https://registry.npmjs.org/type-is/-/type-is-2.1.0.tgz", + "integrity": "sha512-faYHw0anBbc/kWF3zFTEnxSFOAGUX9GFbOBthvDdLsIlEoWOFOtS0zgCiQYwIskL9iGXZL3kAXD8OoZ4GmMATA==", "dev": true, "requires": { - "media-typer": "0.3.0", - "mime-types": "~2.1.24" + "content-type": "^2.0.0", + "media-typer": "^1.1.0", + "mime-types": "^3.0.0" + }, + "dependencies": { + "content-type": { + "version": "2.0.0", + "resolved": "https://registry.npmjs.org/content-type/-/content-type-2.0.0.tgz", + "integrity": "sha512-j/O/d7GcZCyNl7/hwZAb606rzqkyvaDctLmckbxLzHvFBzTJHuGEdodATcP3yIRoDrLHkIATJuvzbFlp/ki2cQ==", + "dev": true + }, + "mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true + }, + "mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", + "dev": true, + "requires": { + "mime-db": "^1.54.0" + } + } } }, "unpipe": { @@ -7047,24 +6816,6 @@ "picocolors": "^1.1.1" } }, - "util-deprecate": { - "version": "1.0.2", - "resolved": "https://registry.npmjs.org/util-deprecate/-/util-deprecate-1.0.2.tgz", - "integrity": "sha1-RQ1Nyfpw3nMnYvvS1KKJgUGaDM8=", - "dev": true - }, - "utils-merge": { - "version": "1.0.1", - "resolved": "https://registry.npmjs.org/utils-merge/-/utils-merge-1.0.1.tgz", - "integrity": "sha512-pMZTvIkT1d+TFGvDOqodOclx0QWkkgi6Tdoa8gC8ffGAAqz9pzPTZWAybbsHHoED/ztMtkv/VoYTYyShUn81hA==", - "dev": true - }, - "uuid": { - "version": "8.3.2", - "resolved": "https://registry.npmjs.org/uuid/-/uuid-8.3.2.tgz", - "integrity": "sha512-+NYs2QeMWy+GWFOEm9xnn6HCDp0l7QBD7ml8zLUmJ+93Q5NF0NocErnwkTkXVFNiX3/fpC6afS8Dhb/gz7R7eg==", - "dev": true - }, "vary": { "version": "1.1.2", "resolved": "https://registry.npmjs.org/vary/-/vary-1.1.2.tgz", @@ -7081,15 +6832,6 @@ "graceful-fs": "^4.1.2" } }, - "wbuf": { - "version": "1.7.3", - "resolved": "https://registry.npmjs.org/wbuf/-/wbuf-1.7.3.tgz", - "integrity": "sha512-O84QOnr0icsbFGLS0O3bI5FswxzRr8/gHwWkDlQFskhSPryQXvrTMxjxGP4+iWYoauLoBvfDpkrOauZ+0iZpDA==", - "dev": true, - "requires": { - "minimalistic-assert": "^1.0.0" - } - }, "webpack": { "version": "5.105.0", "resolved": "https://registry.npmjs.org/webpack/-/webpack-5.105.0.tgz", @@ -7153,53 +6895,65 @@ } }, "webpack-dev-middleware": { - "version": "7.4.2", - "resolved": "https://registry.npmjs.org/webpack-dev-middleware/-/webpack-dev-middleware-7.4.2.tgz", - "integrity": "sha512-xOO8n6eggxnwYpy1NlzUKpvrjfJTvae5/D6WOK0S2LSo7vjmo5gCM1DbLUmFqrMTJP+W/0YZNctm7jasWvLuBA==", + "version": "8.0.4", + "resolved": "https://registry.npmjs.org/webpack-dev-middleware/-/webpack-dev-middleware-8.0.4.tgz", + "integrity": "sha512-9dFzIvIfbdnkOlRjXDHEmEKlY/KPsELNIyKWdoNfK4WaHN9Db+JyVG0gi4/APUPX2UVhnCZ6jp7x0EyM7yTq1Q==", "dev": true, "requires": { - "colorette": "^2.0.10", - "memfs": "^4.6.0", - "mime-types": "^2.1.31", - "on-finished": "^2.4.1", + "memfs": "^4.56.10", + "mime-types": "^3.0.2", "range-parser": "^1.2.1", - "schema-utils": "^4.0.0" + "schema-utils": "^4.3.3" + }, + "dependencies": { + "mime-db": { + "version": "1.54.0", + "resolved": "https://registry.npmjs.org/mime-db/-/mime-db-1.54.0.tgz", + "integrity": "sha512-aU5EJuIN2WDemCcAp2vFBfp/m4EAhWJnUNSSw0ixs7/kXbd6Pg64EmwJkNdFhB8aWt1sH2CTXrLxo/iAGV3oPQ==", + "dev": true + }, + "mime-types": { + "version": "3.0.2", + "resolved": "https://registry.npmjs.org/mime-types/-/mime-types-3.0.2.tgz", + "integrity": "sha512-Lbgzdk0h4juoQ9fCKXW4by0UJqj+nOOrI9MJ1sSj4nI8aI2eo1qmvQEie4VD1glsS250n15LsWsYtCugiStS5A==", + "dev": true, + "requires": { + "mime-db": "^1.54.0" + } + } } }, "webpack-dev-server": { - "version": "5.2.4", - "resolved": "https://registry.npmjs.org/webpack-dev-server/-/webpack-dev-server-5.2.4.tgz", - "integrity": "sha512-GqDPGZN9bRqKBTkp4aWkobDDHMsrXKoGSdOH56smIri8qR0JG8gfL8/v/f/OZR3/OKXjG8uwJbFVhKm/FNU/UA==", + "version": "6.0.0", + "resolved": "https://registry.npmjs.org/webpack-dev-server/-/webpack-dev-server-6.0.0.tgz", + "integrity": "sha512-q9SD4ItOGhZLeU6EGT10caDZdHjF50Pz1DtkRZZOPsfluMXOkacWKKOtSBSLVkPqKiF67eFUC0rI88U/tSFPEw==", "dev": true, "requires": { "@types/bonjour": "^3.5.13", "@types/connect-history-api-fallback": "^1.5.4", - "@types/express": "^4.17.25", - "@types/express-serve-static-core": "^4.17.21", + "@types/express": "^5.0.6", + "@types/express-serve-static-core": "^5.1.1", "@types/serve-index": "^1.9.4", - "@types/serve-static": "^1.15.5", - "@types/sockjs": "^0.3.36", - "@types/ws": "^8.5.10", + "@types/serve-static": "^2.2.0", + "@types/ws": "^8.18.1", "ansi-html-community": "^0.0.8", - "bonjour-service": "^1.2.1", - "chokidar": "^3.6.0", - "colorette": "^2.0.10", + "bonjour-service": "^1.3.0", + "chokidar": "^5.0.0", "compression": "^1.8.1", "connect-history-api-fallback": "^2.0.0", - "express": "^4.22.1", - "graceful-fs": "^4.2.6", - "http-proxy-middleware": "^2.0.9", - "ipaddr.js": "^2.1.0", - "launch-editor": "^2.6.1", - "open": "^10.0.3", - "p-retry": "^6.2.0", - "schema-utils": "^4.2.0", + "express": "^5.2.1", + "graceful-fs": "^4.2.11", + "http-proxy-middleware": "^4.1.1", + "ipaddr.js": "^2.3.0", + "launch-editor": "^2.14.1", + "open": "^11.0.0", + "p-retry": "^8.0.0", + "schema-utils": "^4.3.3", "selfsigned": "^5.5.0", - "serve-index": "^1.9.1", - "sockjs": "^0.3.24", - "spdy": "^4.0.2", - "webpack-dev-middleware": "^7.4.2", - "ws": "^8.18.0" + "serve-index": "^1.9.2", + "tinyglobby": "^0.2.15", + "webpack-dev-middleware": "^8.0.3", + "ws": "^8.20.0" } }, "webpack-merge": { @@ -7218,23 +6972,6 @@ "integrity": "sha512-yd1RBzSGanHkitROoPFd6qsrxt+oFhg/129YzheDGqeustzX0vTZJZsSsQjVQC4yzBQ56K55XU8gaNCtIzOnTg==", "dev": true }, - "websocket-driver": { - "version": "0.7.4", - "resolved": "https://registry.npmjs.org/websocket-driver/-/websocket-driver-0.7.4.tgz", - "integrity": "sha512-b17KeDIQVjvb0ssuSDF2cYXSg2iztliJ4B9WdsuB6J952qCPKmnVq4DyW5motImXHDC1cBT/1UezrJVsKw5zjg==", - "dev": true, - "requires": { - "http-parser-js": ">=0.5.1", - "safe-buffer": ">=5.1.0", - "websocket-extensions": ">=0.1.1" - } - }, - "websocket-extensions": { - "version": "0.1.4", - "resolved": "https://registry.npmjs.org/websocket-extensions/-/websocket-extensions-0.1.4.tgz", - "integrity": "sha512-OqedPIGOfsDlo31UNwYbCFMSaO9m9G/0faIHj5/dZFDMFqPTcx6UwqyOy3COEaEOg/9VsGIpdqn62W5KhoKSpg==", - "dev": true - }, "which": { "version": "2.0.2", "resolved": "https://registry.npmjs.org/which/-/which-2.0.2.tgz", @@ -7250,12 +6987,28 @@ "integrity": "sha512-CC1bOL87PIWSBhDcTrdeLo6eGT7mCFtrg0uIJtqJUFyK+eJnzl8A1niH56uu7KMa5XFrtiV+AQuHO3n7DsHnLQ==", "dev": true }, + "wrappy": { + "version": "1.0.2", + "resolved": "https://registry.npmjs.org/wrappy/-/wrappy-1.0.2.tgz", + "integrity": "sha512-l4Sp/DRseor9wL6EvV2+TuQn63dMkPjZ/sp9XkghTEbV9KlPS1xUsZ3u7/IQO4wxtcFB4bgpQPRcR3QCvezPcQ==", + "dev": true + }, "ws": { - "version": "8.18.2", - "resolved": "https://registry.npmjs.org/ws/-/ws-8.18.2.tgz", - "integrity": "sha512-DMricUmwGZUVr++AEAe2uiVM7UoO9MAVZMDu05UQOaUII0lp+zOzLLU4Xqh/JvTqklB1T4uELaaPBKyjE1r4fQ==", + "version": "8.21.1", + "resolved": "https://registry.npmjs.org/ws/-/ws-8.21.1.tgz", + "integrity": "sha512-+0NTnW77fFN/DjQi6k/Sq/Yvk4Sgajw7urW8V+asjXnRgDs9gyGkdb7EzgfhA4goXsRIZKE28fzIXBHEzhuiWw==", "dev": true, "requires": {} + }, + "wsl-utils": { + "version": "0.3.1", + "resolved": "https://registry.npmjs.org/wsl-utils/-/wsl-utils-0.3.1.tgz", + "integrity": "sha512-g/eziiSUNBSsdDJtCLB8bdYEUMj4jR7AGeUo96p/3dTafgjHhpF4RiCFPiRILwjQoDXx5MqkBr4fwWtR3Ky4Wg==", + "dev": true, + "requires": { + "is-wsl": "^3.1.0", + "powershell-utils": "^0.1.0" + } } } } diff --git a/datafusion/wasmtest/datafusion-wasm-app/package.json b/datafusion/wasmtest/datafusion-wasm-app/package.json index 428460cb39486..e9e98f49495f9 100644 --- a/datafusion/wasmtest/datafusion-wasm-app/package.json +++ b/datafusion/wasmtest/datafusion-wasm-app/package.json @@ -29,7 +29,7 @@ "devDependencies": { "webpack": "5.105.0", "webpack-cli": "5.1.4", - "webpack-dev-server": "5.2.4", + "webpack-dev-server": "6.0.0", "copy-webpack-plugin": "14.0.0" } } diff --git a/datafusion/wasmtest/src/lib.rs b/datafusion/wasmtest/src/lib.rs index f545ccf19306a..d8da5d4b4f323 100644 --- a/datafusion/wasmtest/src/lib.rs +++ b/datafusion/wasmtest/src/lib.rs @@ -206,6 +206,34 @@ mod test { let _ = collect(physical_plan, task_ctx).await.unwrap(); } + #[wasm_bindgen_test(unsupported = tokio::test)] + async fn test_create_table_as_select() { + let ctx = get_ctx(); + ctx.sql("CREATE TABLE t AS SELECT 1 AS a, 'x' AS b") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let result = ctx + .sql("SELECT * FROM t") + .await + .unwrap() + .collect() + .await + .unwrap(); + + assert_eq!( + batches_to_string(&result), + "+---+---+\n\ + | a | b |\n\ + +---+---+\n\ + | 1 | x |\n\ + +---+---+" + ); + } + #[wasm_bindgen_test(unsupported = tokio::test)] async fn test_parquet_write() { let (schema, batch) = create_test_data(); diff --git a/dev/changelog/54.0.0.md b/dev/changelog/54.0.0.md new file mode 100644 index 0000000000000..4bae126539c84 --- /dev/null +++ b/dev/changelog/54.0.0.md @@ -0,0 +1,929 @@ + + +# Apache DataFusion 54.0.0 Changelog + +This release consists of 740 commits from 139 contributors. See credits at the end of this changelog for more information. + +See the [upgrade guide](https://datafusion.apache.org/library-user-guide/upgrading.html) for information on how to upgrade from previous versions. + +**Breaking changes:** + +- Add `ExecutionPlan::apply_expressions()` [#20337](https://github.com/apache/datafusion/pull/20337) (LiaCastaneda) +- Add `Field` to `Expr::Cast` -- allow logical expressions to express a cast to an extension type [#18136](https://github.com/apache/datafusion/pull/18136) (paleolimbot) +- feat: parse `JsonAccess` as a binary operator, add `Operator::Colon` [#20628](https://github.com/apache/datafusion/pull/20628) (Samyak2) +- Wrap Arc to Statistics for `partition_statistics` API [#20570](https://github.com/apache/datafusion/pull/20570) (xudong963) +- Replace ahash with foldhash for faster hashing in datafusion-common [#20958](https://github.com/apache/datafusion/pull/20958) (Dandandan) +- fix: `arrays_zip/list_zip` allow single array argument [#21047](https://github.com/apache/datafusion/pull/21047) (hsiang-c) +- Remove file prefetching from FileStream [#20916](https://github.com/apache/datafusion/pull/20916) (Dandandan) +- Remove as_any from scalar UDF trait definition [#20812](https://github.com/apache/datafusion/pull/20812) (timsaucer) +- Provide session to the udtf call [#20222](https://github.com/apache/datafusion/pull/20222) (askalt) +- chore: remove as_any from aggregate and window functions [#21209](https://github.com/apache/datafusion/pull/21209) (timsaucer) +- chore: remove as_any from ExecutionPlan [#21263](https://github.com/apache/datafusion/pull/21263) (timsaucer) +- fix: Prefer numeric in type coercion for comparisons [#20426](https://github.com/apache/datafusion/pull/20426) (neilconway) +- refactor(pruning): remove column param from PruningStatistics::row_counts [#21369](https://github.com/apache/datafusion/pull/21369) (adriangb) +- Remove CastColumnExpr and custom_file_casts example; unify on field-aware CastExpr [#21563](https://github.com/apache/datafusion/pull/21563) (kosiew) +- perf: Optimize NULL handling in `StringViewArrayBuilder` [#21538](https://github.com/apache/datafusion/pull/21538) (neilconway) +- Remove `as_any` on the `PhysicalExpr` trait [#21573](https://github.com/apache/datafusion/pull/21573) (timsaucer) +- Remove trait function `as_any` from datafusion-datasource [#21576](https://github.com/apache/datafusion/pull/21576) (timsaucer) +- feat: change approx percentile/median UDFs to return floats [#21074](https://github.com/apache/datafusion/pull/21074) (theirix) +- chore: Rename concat-specific string builders, make pub(crate) [#21695](https://github.com/apache/datafusion/pull/21695) (neilconway) +- perf: Implement physical execution of uncorrelated scalar subqueries [#21240](https://github.com/apache/datafusion/pull/21240) (neilconway) +- Add lambda support and array_transform udf [#21679](https://github.com/apache/datafusion/pull/21679) (gstvg) +- perf: strength reduce hash partition modulo (up to 1.16x faster) [#21900](https://github.com/apache/datafusion/pull/21900) (Dandandan) +- feat: Improve InListExpr types, flatten dict haystacks and validate in try_new_from_array [#21402](https://github.com/apache/datafusion/pull/21402) (buraksenn) +- feat: type-keyed extensions map for PartitionedFile [#21993](https://github.com/apache/datafusion/pull/21993) (adriangb) +- Add support for lambda column capture [#21323](https://github.com/apache/datafusion/pull/21323) (gstvg) +- feat: Add Protobuf support for Explain node [#21994](https://github.com/apache/datafusion/pull/21994) (danielhumanmod) +- deprecate: mark Statistics V2 framework (PR #14699) as deprecated [#22071](https://github.com/apache/datafusion/pull/22071) (alamb) +- feat: impl Any for MemoryPool [#21803](https://github.com/apache/datafusion/pull/21803) (haohuaijin) +- Add metrics to `FFI_ExecutionPlan` [#22136](https://github.com/apache/datafusion/pull/22136) (mailmindlin) +- fix(aggregate): show aliased expr in explain [#21739](https://github.com/apache/datafusion/pull/21739) (kumarUjjawal) +- proto: serialize dynamic filters on Sort, Aggregate, HashJoin plan nodes [#22011](https://github.com/apache/datafusion/pull/22011) (jayshrivastava) +- Add exact HigherOrderSignature [#22326](https://github.com/apache/datafusion/pull/22326) (LiaCastaneda) +- Add a memory bound FileStatisticsCache for the Listing Table [#20047](https://github.com/apache/datafusion/pull/20047) (mkleen) +- Add configurable UNION DISTINCT to FILTER rewrite optimization [#21075](https://github.com/apache/datafusion/pull/21075) (xiedeyantu) +- minor: make HigherOrderSignature less error-prone [#22106](https://github.com/apache/datafusion/pull/22106) (gstvg) +- Expose `ExecutionPlan` statistics across the FFI boundary [#22157](https://github.com/apache/datafusion/pull/22157) (mailmindlin) +- feat: optional timezone for coerce_int96 [#22318](https://github.com/apache/datafusion/pull/22318) (andygrove) + +**Performance related:** + +- perf: Optimize `array_to_string` to avoid a copy [#20639](https://github.com/apache/datafusion/pull/20639) (neilconway) +- perf: Apply logical regexp optimizations to Utf8View and LargeUtf8 inputs [#20581](https://github.com/apache/datafusion/pull/20581) (petern48) +- perf: Optimize `array_concat` using `MutableArrayData` [#20620](https://github.com/apache/datafusion/pull/20620) (neilconway) +- perf: Optimize `to_char` to allocate less, fix NULL handling [#20635](https://github.com/apache/datafusion/pull/20635) (neilconway) +- Eliminate deterministic group by keys with deterministic transformations [#20706](https://github.com/apache/datafusion/pull/20706) (Dandandan) +- perf: short-circuit and collect_bool for IN list with column references [#20694](https://github.com/apache/datafusion/pull/20694) (zhangxffff) +- perf: sort replace free()->try_grow() pattern with try_resize() to reduce memory pool interactions [#20729](https://github.com/apache/datafusion/pull/20729) (mbutrovich) +- perf: Optimize set operations to avoid RowConverter deserialization overhead [#20623](https://github.com/apache/datafusion/pull/20623) (neilconway) +- perf: Use batched row conversion for `array_has_any`, `array_has_all` [#20588](https://github.com/apache/datafusion/pull/20588) (neilconway) +- perf: Optimize array set ops on sliced arrays [#20693](https://github.com/apache/datafusion/pull/20693) (neilconway) +- perf: Optimize comparison on nested types [#20716](https://github.com/apache/datafusion/pull/20716) (neilconway) +- perf: Optimize `array_positions()` for scalar needle [#20770](https://github.com/apache/datafusion/pull/20770) (neilconway) +- perf: Optimize `approx_distinct()` for string, binary inputs [#21037](https://github.com/apache/datafusion/pull/21037) (neilconway) +- perf: Optimize `approx_distinct` for inline Utf8View [#21064](https://github.com/apache/datafusion/pull/21064) (neilconway) +- perf: Optimize `strpos()` for scalar needle, plus optimize UTF-8 codepath [#20754](https://github.com/apache/datafusion/pull/20754) (neilconway) +- perf: Optimize `lpad()`, `rpad()` for scalar args [#20657](https://github.com/apache/datafusion/pull/20657) (neilconway) +- perf: add in-place fast path for ScalarValue::add [#20959](https://github.com/apache/datafusion/pull/20959) (kumarUjjawal) +- perf: Optimize `array_sort()` [#21083](https://github.com/apache/datafusion/pull/21083) (neilconway) +- Super fast extended tests and improved planning speed linux [#21084](https://github.com/apache/datafusion/pull/21084) (blaginin) +- Add a builder to `SimplifyContext` to avoid allocating default values [#21092](https://github.com/apache/datafusion/pull/21092) (AdamGS) +- Avoid creating new RecordBatches to simplify expressions [#20534](https://github.com/apache/datafusion/pull/20534) (alamb) +- perf: optimize scatter with type-specific specialization [#20498](https://github.com/apache/datafusion/pull/20498) (CuteChuanChuan) +- perf: Optimize `array_min`, `array_max` for arrays of primitive types [#21101](https://github.com/apache/datafusion/pull/21101) (neilconway) +- perf: optimize map validation for common key types [#20805](https://github.com/apache/datafusion/pull/20805) (lyne7-sc) +- perf: specialized SemiAntiSortMergeJoinStream [#20806](https://github.com/apache/datafusion/pull/20806) (mbutrovich) +- Improvement: keep order-preserving repartitions for streaming aggregates [#21107](https://github.com/apache/datafusion/pull/21107) (xudong963) +- perf: Add support for `GroupsAccumulator` to `string_agg` [#21154](https://github.com/apache/datafusion/pull/21154) (neilconway) +- perf: Optimize `split_part`, support `Utf8View` [#21119](https://github.com/apache/datafusion/pull/21119) (neilconway) +- perf: sort-merge join (SMJ) batch deferred filtering and move mark joins to bitwise stream. Near-unique LEFT and FULL SMJ 20-50x faster [#21184](https://github.com/apache/datafusion/pull/21184) (mbutrovich) +- perf: Optimize `string_to_array` for scalar args [#21131](https://github.com/apache/datafusion/pull/21131) (neilconway) +- Misc minor optimizations to query optimizer performance [#21128](https://github.com/apache/datafusion/pull/21128) (AdamGS) +- ensure dynamic filters are correctly pushed down through aggregations [#21059](https://github.com/apache/datafusion/pull/21059) (jayshrivastava) +- perf: Merge Precision in-place [#21219](https://github.com/apache/datafusion/pull/21219) (AdamGS) +- feat: support GroupsAccumulator for first_value and last_value with string/binary types [#21090](https://github.com/apache/datafusion/pull/21090) (UBarney) +- perf: Optimize `split_part` for scalar args [#21238](https://github.com/apache/datafusion/pull/21238) (neilconway) +- perf: optimize object store requests when reading JSON [#20823](https://github.com/apache/datafusion/pull/20823) (ariel-miculas) +- perf: Optimize `split_part` for `Utf8View` [#21420](https://github.com/apache/datafusion/pull/21420) (neilconway) +- Eliminate outer joins with empty relations via null-padded projection [#21321](https://github.com/apache/datafusion/pull/21321) (SubhamSinghal) +- Optimize `regexp_replace` by stripping trailing .\* from anchored patterns. 2.4x improvement (ClickBench Q28) [#21379](https://github.com/apache/datafusion/pull/21379) (Dandandan) +- perf: use DynComparator in sort-merge join (SMJ), microbenchmark queries up to 12% faster, TPC-H overall ~5% faster [#21484](https://github.com/apache/datafusion/pull/21484) (mbutrovich) +- perf: Optimize NULL handling in `substr` [#21519](https://github.com/apache/datafusion/pull/21519) (neilconway) +- perf: replace SMJ's join_filter_not_matched_map HashMap with Vec [#21517](https://github.com/apache/datafusion/pull/21517) (mbutrovich) +- perf: Optimize NULL handling in `find_in_set` [#21464](https://github.com/apache/datafusion/pull/21464) (neilconway) +- perf: Optimize NULL handling in `lcm`, `gcd` [#21468](https://github.com/apache/datafusion/pull/21468) (neilconway) +- perf: Optimize NULL handling in `arrays_zip` [#21475](https://github.com/apache/datafusion/pull/21475) (neilconway) +- perf: Optimize NULL handling in `array_remove` [#21532](https://github.com/apache/datafusion/pull/21532) (neilconway) +- perf: Optimize NULL handling in `array_slice` [#21482](https://github.com/apache/datafusion/pull/21482) (neilconway) +- perf: Optimize NULL handling in some datetime functions [#21477](https://github.com/apache/datafusion/pull/21477) (neilconway) +- perf: Optimize NULL handling in `array_has` [#21471](https://github.com/apache/datafusion/pull/21471) (neilconway) +- perf: Optimize `Utf8View` string concat [#21535](https://github.com/apache/datafusion/pull/21535) (neilconway) +- Conditionally build page pruning predicates [#21480](https://github.com/apache/datafusion/pull/21480) (fpetkovski) +- perf: add fast path for uniform fill values in `array_resize` [#20617](https://github.com/apache/datafusion/pull/20617) (lyne7-sc) +- perf : Optimize count distinct using bitmaps instead of hashsets for smaller datatypes [#21456](https://github.com/apache/datafusion/pull/21456) (coderfender) +- perf: Optimize `left`, `right` to reduce copying [#21442](https://github.com/apache/datafusion/pull/21442) (neilconway) +- perf: Optimize `substr` for Utf8, LargeUtf8 [#21366](https://github.com/apache/datafusion/pull/21366) (neilconway) +- feat: Optimize ORDER BY by Pruning Functionally Redundant Sort Keys [#21362](https://github.com/apache/datafusion/pull/21362) (xiedeyantu) +- perf: Optimize logical optimizer's `OptimizeProjections` pass [#21726](https://github.com/apache/datafusion/pull/21726) (neilconway) +- perf: Optimize `DFSchema::qualified_name` [#21722](https://github.com/apache/datafusion/pull/21722) (neilconway) +- perf: Tweak vec capacity in `project_statistics` [#21734](https://github.com/apache/datafusion/pull/21734) (neilconway) +- perf: Reduce `Box` and `Arc` allocation churn during tree rewriting [#21749](https://github.com/apache/datafusion/pull/21749) (neilconway) +- perf: Implement groups accumulator count distinct primitive types [#21561](https://github.com/apache/datafusion/pull/21561) (coderfender) +- perf: Optimize approx count distinct using bitmaps instead of HLL for smaller int datatypes [#21453](https://github.com/apache/datafusion/pull/21453) (coderfender) +- perf: Optimize `lower`, `upper` for sliced arrays [#21814](https://github.com/apache/datafusion/pull/21814) (neilconway) +- perf: Add bulk NULL-aware string builders, use in `lower` and `upper` [#21789](https://github.com/apache/datafusion/pull/21789) (neilconway) +- perf: Use bulk-NULL builder in `uuid` [#21845](https://github.com/apache/datafusion/pull/21845) (neilconway) +- Skip map_expressions rebuild for Extension nodes with empty expressions [#21701](https://github.com/apache/datafusion/pull/21701) (zhuqi-lucas) +- Refactor InListExpr into static-filter modules [#21649](https://github.com/apache/datafusion/pull/21649) (geoffreyclaude) +- perf: Use bulk-NULL string builder in `initcap` [#21863](https://github.com/apache/datafusion/pull/21863) (neilconway) +- perf: Use bulk-NULL builder in `chr` [#21847](https://github.com/apache/datafusion/pull/21847) (neilconway) +- perf: implement convert_to_state for SparkAvg [#21548](https://github.com/apache/datafusion/pull/21548) (azhangd) +- perf: optimise `first_value`, `last_value` aggregate function [#21383](https://github.com/apache/datafusion/pull/21383) (theirix) +- perf(spark): use 256-entry byte-pair table in hex encoding [#21836](https://github.com/apache/datafusion/pull/21836) (Scolliq) +- perf: Optimize `substr_index` to use bulk-NULL string builder [#21877](https://github.com/apache/datafusion/pull/21877) (neilconway) +- perf: Use bulk-NULL builder in `replace` [#21849](https://github.com/apache/datafusion/pull/21849) (neilconway) +- Add SQL based benchmarking harness, port tpch to use framework [#21707](https://github.com/apache/datafusion/pull/21707) (Omega359) +- perf: Add `BulkNullStringArrayBuilder` trait, use in `repeat` [#21854](https://github.com/apache/datafusion/pull/21854) (neilconway) +- perf: optimize retract_batch for `median` and `percentile_cont` [#21894](https://github.com/apache/datafusion/pull/21894) (lyne7-sc) +- perf: Optimize `reverse` using bulk-NULL string builders [#21991](https://github.com/apache/datafusion/pull/21991) (neilconway) +- perf: Optimize `lower`, `upper` for ASCII inputs [#21980](https://github.com/apache/datafusion/pull/21980) (neilconway) +- perf: Cast entire Date32 array to Date64 on 1st failure [#21948](https://github.com/apache/datafusion/pull/21948) (huymq1710) +- perf: Use `NullBuffer::union_many` [#22070](https://github.com/apache/datafusion/pull/22070) (neilconway) +- perf: improve Int64 `generate_series` and `range` performance [#21891](https://github.com/apache/datafusion/pull/21891) (lyne7-sc) +- perf: batch contiguous extend calls in `array_replace` [#22119](https://github.com/apache/datafusion/pull/22119) (lyne7-sc) +- perf: Add `append_with` to string builders, use in `replace` [#22029](https://github.com/apache/datafusion/pull/22029) (neilconway) +- perf: reuse mask in `truncate_list_nulls` and avoid counting all true bits [#22158](https://github.com/apache/datafusion/pull/22158) (rluvaton) +- Skip RowFilter and page pruning for fully matched row groups [#21637](https://github.com/apache/datafusion/pull/21637) (xudong963) +- perf: bypass values.value(i) for inline strings in ArrowBytesViewMap [#22172](https://github.com/apache/datafusion/pull/22172) (RyanJamesStewart) +- perf: Elimiate SortExec on generate_series() [#22238](https://github.com/apache/datafusion/pull/22238) (2010YOUY01) +- perf: coalesce batches before sending to distributor channels in RepartitionExec [#22010](https://github.com/apache/datafusion/pull/22010) (gabotechs) +- Resolve MIN/MAX from Parquet metadata for Single-mode aggregates and CAST projections [#21651](https://github.com/apache/datafusion/pull/21651) (Dandandan) +- Compact more aggressively in TopK based upon memory usage [#20381](https://github.com/apache/datafusion/pull/20381) (cetra3) + +**Implemented enhancements:** + +- feat: support nanosecond date_part [#20674](https://github.com/apache/datafusion/pull/20674) (mhilton) +- feat: Support Spark `array_contains` builtin function [#20685](https://github.com/apache/datafusion/pull/20685) (comphead) +- feat: Integrate CastColumnExpr into PhysicalExprAdapter [#20269](https://github.com/apache/datafusion/pull/20269) (kumarUjjawal) +- feat: `partition_statistics()` for HashJoinExec [#20711](https://github.com/apache/datafusion/pull/20711) (jonathanc-n) +- feat: make DefaultLogicalExtensionCodec support serialisation of buil… [#20638](https://github.com/apache/datafusion/pull/20638) (Acfboy) +- feat: correct struct column names for `arrays_zip` return type [#20886](https://github.com/apache/datafusion/pull/20886) (comphead) +- feat: Reduce allocations for aggregating `Statistics` [#20768](https://github.com/apache/datafusion/pull/20768) (jonathanc-n) +- feat: add `custom_string_literal_override` to unparser Dialect trait [#20590](https://github.com/apache/datafusion/pull/20590) (goldmedal) +- feat: Extract NDV (distinct_count) statistics from Parquet metadata [#19957](https://github.com/apache/datafusion/pull/19957) (asolimando) +- feat: support repartitioning of FFI execution plans [#20449](https://github.com/apache/datafusion/pull/20449) (timsaucer) +- feat: create a datafusion-example for in-memory file format [#20394](https://github.com/apache/datafusion/pull/20394) (kumarUjjawal) +- feat: implement PhysicalOptimizerRule in FFI crate [#20451](https://github.com/apache/datafusion/pull/20451) (timsaucer) +- feat(metric): Add output skewness metric to detect skewed plans easier [#21211](https://github.com/apache/datafusion/pull/21211) (2010YOUY01) +- feat: add sort pushdown benchmark and SLT tests [#21213](https://github.com/apache/datafusion/pull/21213) (zhuqi-lucas) +- feat(sql): unparse array_has as ANY for Postgres [#20654](https://github.com/apache/datafusion/pull/20654) (vimeh) +- feat: feature-gate `sqllogictests` datafusion-substrait behind optional 'substrait' feature [#21268](https://github.com/apache/datafusion/pull/21268) (zhuqi-lucas) +- feat: generate reversed-name data for sort pushdown benchmark [#21266](https://github.com/apache/datafusion/pull/21266) (zhuqi-lucas) +- feat: Complete basic `LATERAL JOIN` functionality [#21202](https://github.com/apache/datafusion/pull/21202) (neilconway) +- feat: Use NDV for equality filter selectivity calculation [#20789](https://github.com/apache/datafusion/pull/20789) (jonathanc-n) +- feat: make BatchPartitioner::partition_iter public [#21341](https://github.com/apache/datafusion/pull/21341) (hcrosse) +- feat: spark compatible float to timestamp cast with ANSI support [#21212](https://github.com/apache/datafusion/pull/21212) (coderfender) +- feat(spark): Adds spark round function [#21062](https://github.com/apache/datafusion/pull/21062) (SubhamSinghal) +- feat: make DataFrame::create_physical_plan take &self instead of self [#20562](https://github.com/apache/datafusion/pull/20562) (xanderbailey) +- feat: add support for parquet content defined chunking options [#21110](https://github.com/apache/datafusion/pull/21110) (kszucs) +- feat: sort file groups by statistics during sort pushdown (Sort pushdown phase 2) [#21182](https://github.com/apache/datafusion/pull/21182) (zhuqi-lucas) +- feat: Set NDV to Exact(1) for numeric equality filter predicates [#21077](https://github.com/apache/datafusion/pull/21077) (asolimando) +- feat: make sort pushdown BufferExec capacity configurable, default 1GB [#21426](https://github.com/apache/datafusion/pull/21426) (zhuqi-lucas) +- feat: Propagate orderings through struct-producing projections [#21218](https://github.com/apache/datafusion/pull/21218) (rkrishn7) +- feat: add cast_to_type UDF for type-based casting [#21322](https://github.com/apache/datafusion/pull/21322) (adriangb) +- feat: Add pluggable StatisticsRegistry for operator-level statistics propagation [#21483](https://github.com/apache/datafusion/pull/21483) (asolimando) +- feat: Add Hash trait to Aggregate enums [#21569](https://github.com/apache/datafusion/pull/21569) (rluvaton) +- feat(substrait): support Placeholder <-> DynamicParameter in Substrait producer/consumer [#20977](https://github.com/apache/datafusion/pull/20977) (bvolpato) +- feat: add `with_metadata` scalar UDF to attach Arrow field metadata [#21509](https://github.com/apache/datafusion/pull/21509) (adriangb) +- feat: Additional Canonical Extension Types [#21291](https://github.com/apache/datafusion/pull/21291) (tschwarzinger) +- feat: Add memory-limited execution for NestedLoopJoinExec [#21448](https://github.com/apache/datafusion/pull/21448) (viirya) +- feat(stats): cap NDV at row count in statistics estimation [#21081](https://github.com/apache/datafusion/pull/21081) (asolimando) +- feat: support `array_compact` builtin function [#21522](https://github.com/apache/datafusion/pull/21522) (comphead) +- feat: add a config to disable subquery_sort_elimination [#21614](https://github.com/apache/datafusion/pull/21614) (haohuaijin) +- feat: extend single ndv optimization to non-arithmetic supporting types for equality predicates [#21473](https://github.com/apache/datafusion/pull/21473) (buraksenn) +- feat: extend interval analysis support for temporal types [#21520](https://github.com/apache/datafusion/pull/21520) (buraksenn) +- feat: add sort_pushdown_inexact benchmark for RG reorder [#21674](https://github.com/apache/datafusion/pull/21674) (zhuqi-lucas) +- feat: support '>', '<', '>=', '<=', '<>' in all operator [#21416](https://github.com/apache/datafusion/pull/21416) (buraksenn) +- feat: Add support for `LEFT JOIN LATERAL` [#21352](https://github.com/apache/datafusion/pull/21352) (neilconway) +- feat: Expose used `MemoryPool` details in `ResourcesExhausted` error messages [#20387](https://github.com/apache/datafusion/pull/20387) (erenavsarogullari) +- feat: estimate cardinality for semi and anti-joins using distinct counts [#20904](https://github.com/apache/datafusion/pull/20904) (buraksenn) +- feat: support `ListView` and `LargeListView` in `ScalarValue` [#21669](https://github.com/apache/datafusion/pull/21669) (Jefffrey) +- feat: add cosine_distance scalar function [#21542](https://github.com/apache/datafusion/pull/21542) (crm26) +- feat: remove `__unnest_placeholder` from struct unnest projection [#21725](https://github.com/apache/datafusion/pull/21725) (akoshchiy) +- feat(unparser): Keep inner join `Filter → TableScan` predicates to `WHERE` instead of moving to `JOIN ON` [#21694](https://github.com/apache/datafusion/pull/21694) (sgrebnov) +- feat: minor lambda perf improvements [#21896](https://github.com/apache/datafusion/pull/21896) (comphead) +- feat: automatically cast `ListView` to `List` for UDFs [#21855](https://github.com/apache/datafusion/pull/21855) (Jefffrey) +- feat: support binary arguments for StringConcat operator [#21883](https://github.com/apache/datafusion/pull/21883) (theirix) +- feat: add inner_product scalar function [#21861](https://github.com/apache/datafusion/pull/21861) (crm26) +- feat: Support RIGHT/FULL joins in NLJ memory-limited execution [#21833](https://github.com/apache/datafusion/pull/21833) (viirya) +- feat: Improved multiple column aggregation performance by using bitmasks rather than `Vec` [#21886](https://github.com/apache/datafusion/pull/21886) (huymq1710) +- feat: Making From conversions fallible with `TryFrom` [#21985](https://github.com/apache/datafusion/pull/21985) (Soham-Bhattacharjee-work) +- feat: support spark compatible floor function [#21933](https://github.com/apache/datafusion/pull/21933) (athlcode) +- feat: fix NTILE distribution logic [#22051](https://github.com/apache/datafusion/pull/22051) (comphead) +- feat: implement retract_batch for array_agg sliding window support [#22015](https://github.com/apache/datafusion/pull/22015) (SubhamSinghal) +- feat: Upgrade to sqlparser-rs 0.62.0 [#22069](https://github.com/apache/datafusion/pull/22069) (andygrove) +- feat: fix windows frame positive/neg overflows [#22140](https://github.com/apache/datafusion/pull/22140) (comphead) +- feat: fix AVG sliding windows wrong results with NULLs [#22139](https://github.com/apache/datafusion/pull/22139) (comphead) +- feat: fix windows decimal casting frame [#22174](https://github.com/apache/datafusion/pull/22174) (comphead) +- feat: eliminate GlobalLimitExec when input statistics prove limit is already satisfied [#22150](https://github.com/apache/datafusion/pull/22150) (xiedeyantu) +- feat: globally reorder files and row groups by statistics for TopK queries [#21956](https://github.com/apache/datafusion/pull/21956) (zhuqi-lucas) +- feat: Restore nullability when consuming substrait fields [#22105](https://github.com/apache/datafusion/pull/22105) (neilconway) +- feat: add array_normalize scalar function [#22013](https://github.com/apache/datafusion/pull/22013) (crm26) +- feat: add Spark-compatible xxhash64 function [#21967](https://github.com/apache/datafusion/pull/21967) (andygrove) + +**Fixed bugs:** + +- fix: make the `sql` feature truly optional [#20625](https://github.com/apache/datafusion/pull/20625) (linhr) +- fix: use try_shrink instead of shrink in try_resize [#20424](https://github.com/apache/datafusion/pull/20424) (ariel-miculas) +- fix: Provide more generic API for the capacity limit parsing [#20372](https://github.com/apache/datafusion/pull/20372) (erenavsarogullari) +- fix: Fix bug in `array_has` scalar path with sliced arrays [#20677](https://github.com/apache/datafusion/pull/20677) (neilconway) +- fix: `HashJoin` panic with String dictionary keys (don't flatten keys) [#20505](https://github.com/apache/datafusion/pull/20505) (alamb) +- fix: Return `probe_side.len()` for RightMark/Anti count(\*) queries [#20710](https://github.com/apache/datafusion/pull/20710) (jonathanc-n) +- fix: preserve None projection semantics across FFI boundary in ForeignTableProvider::scan [#20393](https://github.com/apache/datafusion/pull/20393) (Kontinuation) +- fix(spark): handle divide-by-zero in Spark `mod`/`pmod` with ANSI mode support [#20461](https://github.com/apache/datafusion/pull/20461) (davidlghellin) +- fix: sqllogictest cannot convert to Substrait [#19739](https://github.com/apache/datafusion/pull/19739) (kumarUjjawal) +- fix: interval analysis error when have two filterexec that inner filter proves zero selectivity [#20743](https://github.com/apache/datafusion/pull/20743) (haohuaijin) +- fix: SanityCheckPlan error with window functions and NVL filter [#20231](https://github.com/apache/datafusion/pull/20231) (EeshanBembi) +- fix: Avoid unnecessary type casts in `concat_ws` [#20436](https://github.com/apache/datafusion/pull/20436) (neilconway) +- fix: Remove `!=0` check from `supports_collect_by_thresholds` [#20730](https://github.com/apache/datafusion/pull/20730) (jonathanc-n) +- fix: do not recompute hash join exec properties if not required [#20900](https://github.com/apache/datafusion/pull/20900) (askalt) +- fix: Optimize `!~ '.*'` case to `col IS NULL AND Boolean(NULL)` instead of `Eq ""` [#20702](https://github.com/apache/datafusion/pull/20702) (petern48) +- fix: Track metrics in hash joins with empty build sides [#20810](https://github.com/apache/datafusion/pull/20810) (nuno-faria) +- fix: dfbench respects DATAFUSION_RUNTIME_MEMORY_LIMIT env var [#20631](https://github.com/apache/datafusion/pull/20631) (adriangb) +- fix(spark): return input string for PATH/FILE on schemeless URLs in `parse_url` [#20506](https://github.com/apache/datafusion/pull/20506) (davidlghellin) +- fix: InList Dictionary filter pushdown type mismatch [#20962](https://github.com/apache/datafusion/pull/20962) (erratic-pattern) +- fix: Run release verification with `--profile=ci` [#20987](https://github.com/apache/datafusion/pull/20987) (alamb) +- fix: move overflow guard before dense ratio in hash join to prevent overflows [#20998](https://github.com/apache/datafusion/pull/20998) (buraksenn) +- fix: improve GroupOrdering docs [#20994](https://github.com/apache/datafusion/pull/20994) (alamb) +- fix: update clickbench expected plan for NDV-aware optimization [#21050](https://github.com/apache/datafusion/pull/21050) (asolimando) +- fix: use datafusion_expr instead of datafusion crate in spark [#21043](https://github.com/apache/datafusion/pull/21043) (davidlghellin) +- Fix CTE reference resolution slt tests [#21049](https://github.com/apache/datafusion/pull/21049) (jonahgao) +- fix: validate wrapped negation during type coercion [#20965](https://github.com/apache/datafusion/pull/20965) (myandpr) +- fix(sql): handle GROUP BY ALL with aliased aggregates [#20943](https://github.com/apache/datafusion/pull/20943) (kumarUjjawal) +- fix: string_to_array('', delim) returns empty array for PostgreSQL compatibility [#21104](https://github.com/apache/datafusion/pull/21104) (dd-david-levin) +- Fix push_down_filter for children with non-empty fetch fields [#21057](https://github.com/apache/datafusion/pull/21057) (shivbhatia10) +- fix(stats): widen sum_value integer arithmetic to SUM-compatible types [#20865](https://github.com/apache/datafusion/pull/20865) (kumarUjjawal) +- fix: skip empty metadata in intersect_metadata_for_union to prevent s… [#21127](https://github.com/apache/datafusion/pull/21127) (RafaelHerrero) +- fix: Df int timestamp cast fix failing CI [#21163](https://github.com/apache/datafusion/pull/21163) (coderfender) +- fix(unparser): Fix BigQuery timestamp literal format in SQL unparsing [#21103](https://github.com/apache/datafusion/pull/21103) (sgrebnov) +- fix: propagate errors for unsupported table function arguments instead of silently dropping them [#21135](https://github.com/apache/datafusion/pull/21135) (buraksenn) +- fix: Fix `main` compilation failure [#21242](https://github.com/apache/datafusion/pull/21242) (2010YOUY01) +- fix: Revert "Fix/support duplicate column names #6543 (#21126)" [#21254](https://github.com/apache/datafusion/pull/21254) (mbutrovich) +- fix: Fix three bugs in query decorrelation [#21208](https://github.com/apache/datafusion/pull/21208) (neilconway) +- fix: date overflow panic [#21233](https://github.com/apache/datafusion/pull/21233) (haohuaijin) +- fix: `SELECT * EXCLUDE(...)` silently returns empty rows when all columns are excluded [#21259](https://github.com/apache/datafusion/pull/21259) (xiedeyantu) +- fix(unparser): use to_rfc3339 for default TIMESTAMPTZ formatting [#21295](https://github.com/apache/datafusion/pull/21295) (sgrebnov) +- fix: use spill writer's schema instead of the first batch schema for spill files [#21293](https://github.com/apache/datafusion/pull/21293) (gruuya) +- fix: binary string concat [#20787](https://github.com/apache/datafusion/pull/20787) (theirix) +- fix(sql): fix a bug when planning semi- or antijoins [#20990](https://github.com/apache/datafusion/pull/20990) (aalexandrov) +- fix(datasource): keep stats absent when collect_stats is false [#21149](https://github.com/apache/datafusion/pull/21149) (kumarUjjawal) +- fix: preserve source field metadata in TryCast expressions [#21390](https://github.com/apache/datafusion/pull/21390) (adriangb) +- fix: skips projection pruning for whole subtree [#20545](https://github.com/apache/datafusion/pull/20545) (Acfboy) +- fix: preserve subquery structure when unparsing SubqueryAlias over Ag… [#21099](https://github.com/apache/datafusion/pull/21099) (yonatan-sevenai) +- fix: FilterExec should drop projection when apply projection pushdown [#21460](https://github.com/apache/datafusion/pull/21460) (haohuaijin) +- fix: preserve duplicate GROUPING SETS rows [#21058](https://github.com/apache/datafusion/pull/21058) (xiedeyantu) +- fix: apply the left side schema on the right side in set expressions [#21052](https://github.com/apache/datafusion/pull/21052) (gruuya) +- fix: Use codepoints in `lpad`, `rpad`, `translate` [#21405](https://github.com/apache/datafusion/pull/21405) (neilconway) +- fix: PostgreSQL dialect can not support tinyint type [#21445](https://github.com/apache/datafusion/pull/21445) (xiedeyantu) +- fix: DataFusion benchmark panicked: failed to cast '2013-07-01' to UInt16 [#21498](https://github.com/apache/datafusion/pull/21498) (xiedeyantu) +- fix(sql): return planner error for malformed typed literals [#21454](https://github.com/apache/datafusion/pull/21454) (officialasishkumar) +- fix: Preserve quoted mixed-case identifiers in the `pivot_unpivot` example [#21432](https://github.com/apache/datafusion/pull/21432) (niebayes) +- fix(spark): array_repeat returns repeated NULLs instead of NULL when element is NULL [#21558](https://github.com/apache/datafusion/pull/21558) (buraksenn) +- fix: grouping with alias [#21438](https://github.com/apache/datafusion/pull/21438) (timsaucer) +- fix(spark): mod/pmod returns NULL instead of NaN for float division by zero [#21557](https://github.com/apache/datafusion/pull/21557) (buraksenn) +- fix: LazyMemoryExec should produce independent streams per execute() [#21565](https://github.com/apache/datafusion/pull/21565) (viirya) +- fix: json scan performance on local files [#21478](https://github.com/apache/datafusion/pull/21478) (ariel-miculas) +- fix(benchmarks): correct TPC-H benchmark SQL [#21615](https://github.com/apache/datafusion/pull/21615) (kumarUjjawal) +- fix: suppress nondeterministic metrics in agg_dyn_e2e sqllogictest [#21657](https://github.com/apache/datafusion/pull/21657) (mbutrovich) +- fix: Fix compilation error on `main` [#21664](https://github.com/apache/datafusion/pull/21664) (2010YOUY01) +- fix: `median` retract logic for sliding window frames [#21300](https://github.com/apache/datafusion/pull/21300) (lyne7-sc) +- fix: Fix Spark `slice` function `Null` type to `GenericListArray` casting issue [#20469](https://github.com/apache/datafusion/pull/20469) (erenavsarogullari) +- fix: Remove nested async block causing Stacked Borrows violation in PushDecoderStreamState [#21663](https://github.com/apache/datafusion/pull/21663) (mbutrovich) +- fix: impl `handle_child_pushdown_result` for `SortExec` [#21527](https://github.com/apache/datafusion/pull/21527) (haohuaijin) +- fix: SortMergeJoin full outer join incorrectly matches rows when filter evaluates to NULL [#21660](https://github.com/apache/datafusion/pull/21660) (mbutrovich) +- fix: try again to fix Miri in ParquetOpener [#21680](https://github.com/apache/datafusion/pull/21680) (mbutrovich) +- fix: `optimize_projections` failure after mark joins created by `EXISTS OR EXISTS` [#21265](https://github.com/apache/datafusion/pull/21265) (buraksenn) +- fix: import from `datafusion_expr` in `make_valid_utf8` [#21687](https://github.com/apache/datafusion/pull/21687) (hcrosse) +- fix: linearized operands in physical binaryexpr protobuf to avoid recursion limit [#21031](https://github.com/apache/datafusion/pull/21031) (haohuaijin) +- fix: remove unnecessary `as_any()` to fix compilation error [#21693](https://github.com/apache/datafusion/pull/21693) (Jefffrey) +- fix: Prevent CLI crash on wide tables [#21721](https://github.com/apache/datafusion/pull/21721) (Geethapranay1) +- fix(unparser): make `BigQueryDialect` more robust [#21296](https://github.com/apache/datafusion/pull/21296) (sgrebnov) +- fix: insert placeholder type inference showing wrong type when there is function wrapped placeholder (unknown type) [#20744](https://github.com/apache/datafusion/pull/20744) (buraksenn) +- fix: array_concat widens container variant for mixed List/LargeList inputs [#21704](https://github.com/apache/datafusion/pull/21704) (hcrosse) +- fix: Fix local `datafusion-cli` test failure [#21761](https://github.com/apache/datafusion/pull/21761) (2010YOUY01) +- fix: Validate spill read schema [#21738](https://github.com/apache/datafusion/pull/21738) (2010YOUY01) +- fix: improve sort pushdown benchmark data and add DESC LIMIT queries [#21711](https://github.com/apache/datafusion/pull/21711) (zhuqi-lucas) +- fix: rebind RecursiveQueryExec batches to the declared output schema [#21770](https://github.com/apache/datafusion/pull/21770) (adriangb) +- fix: Enable `arrow-ipc/zstd` in `datasource-arrow` to make `test_spill_compression` pass in every config [#21504](https://github.com/apache/datafusion/pull/21504) (AdamGS) +- fix: Do not highlight the CLI hint directly [#21858](https://github.com/apache/datafusion/pull/21858) (nuno-faria) +- fix: fix elapsed_compute metric in ParquetSink to report encoding time only [#21825](https://github.com/apache/datafusion/pull/21825) (fred1268) +- fix: grouping separator for float and decimal [#20268](https://github.com/apache/datafusion/pull/20268) (Druva-D) +- fix: Fix `.gitignore` in `benchmarks/` [#21954](https://github.com/apache/datafusion/pull/21954) (2010YOUY01) +- fix(proto): correctly serialize FilterExec empty projection [#21885](https://github.com/apache/datafusion/pull/21885) (Adez017) +- fix: Make conversion from FileDecryptionProperties to ConfigFileDecryptionProperties fallible [#21603](https://github.com/apache/datafusion/pull/21603) (adamreeve) +- fix: Avoid unnecessary input repartitioning with `ScalarSubqueryExec` [#21986](https://github.com/apache/datafusion/pull/21986) (neilconway) +- fix: error on CREATE EXTERNAL TABLE with no files and no explicit schema [#21965](https://github.com/apache/datafusion/pull/21965) (adriangb) +- fix: `median` returns Float64 for integer inputs to avoid truncation [#21988](https://github.com/apache/datafusion/pull/21988) (CuteChuanChuan) +- fix: Correct the number of pruned/matched Parquet pages [#22031](https://github.com/apache/datafusion/pull/22031) (nuno-faria) +- fix: use datafusion_expr instead of datafusion crate [#22052](https://github.com/apache/datafusion/pull/22052) (hsiang-c) +- fix(spark): align parse_url empty FILE path [#21969](https://github.com/apache/datafusion/pull/21969) (kumarUjjawal) +- fix: drop input plan early in `CoalescePartitionsExec` [#22017](https://github.com/apache/datafusion/pull/22017) (Samyak2) +- fix: track join_arrays memory in reservation after SMJ spill [#21962](https://github.com/apache/datafusion/pull/21962) (SubhamSinghal) +- fix: Avoid `overlay` panic on valid Unicode input, Postgres compatibility [#22046](https://github.com/apache/datafusion/pull/22046) (neilconway) +- fix: Panic in Spark's `format_string` for illegal characters [#22077](https://github.com/apache/datafusion/pull/22077) (neilconway) +- fix: Incorrect behavior for `FILTER` on NULLs [#22068](https://github.com/apache/datafusion/pull/22068) (neilconway) +- fix: coerce operand types in Interval mul/div/intersect/union/contains [#22027](https://github.com/apache/datafusion/pull/22027) (adriangb) +- fix(bench): avoid OOM in `array_replace` bench [#22120](https://github.com/apache/datafusion/pull/22120) (kumarUjjawal) +- fix: Nested self-referential CASE chains should not cause exponential hashing work during physical planning. [#22175](https://github.com/apache/datafusion/pull/22175) (avantgardnerio) +- fix: preserve Inexact precision in Statistics [#22146](https://github.com/apache/datafusion/pull/22146) (timsaucer) +- fix: Handle EXECUTE without statement name [#22204](https://github.com/apache/datafusion/pull/22204) (Dandandan) +- fix(sql): reject duplicate unqualified names in CTAS, CREATE VIEW, and SELECT INTO [#22290](https://github.com/apache/datafusion/pull/22290) (kumarUjjawal) +- fix: reduce memory allocation overhead during partial aggregation ear… [#22165](https://github.com/apache/datafusion/pull/22165) (ariel-miculas) +- fix: Fix bug with structurally equal correlated subqueries [#22313](https://github.com/apache/datafusion/pull/22313) (neilconway) +- fix: return error instead of capacity overflow panic in generate_series [#22323](https://github.com/apache/datafusion/pull/22323) (sweb) +- fix: simplifier on leaf nodes returns null [#22368](https://github.com/apache/datafusion/pull/22368) (timsaucer) + +**Documentation updates:** + +- Update DataFusion meetups page on docs [#20629](https://github.com/apache/datafusion/pull/20629) (alamb) +- docs: Update `datafusion-cli` doc for `top-memory-consumers` config [#20390](https://github.com/apache/datafusion/pull/20390) (erenavsarogullari) +- [main] Update version to 52.2.0 [#20573](https://github.com/apache/datafusion/pull/20573) (alamb) +- Update releases links with releases in 2025-2026 [#20630](https://github.com/apache/datafusion/pull/20630) (alamb) +- doc: Add more context to `Precision` [#20713](https://github.com/apache/datafusion/pull/20713) (jonathanc-n) +- Minor: Add comment explaining rationale to avoid dependencies on functions [#20667](https://github.com/apache/datafusion/pull/20667) (alamb) +- Hash join buffering on probe side [#19761](https://github.com/apache/datafusion/pull/19761) (gabotechs) +- Copy limits before repartitions [#20736](https://github.com/apache/datafusion/pull/20736) (avantgardnerio) +- Allow SQL `TypePlanner` to plan SQL types as extension types [#20676](https://github.com/apache/datafusion/pull/20676) (paleolimbot) +- doc: Add documentation for pushing limit into plan [#20271](https://github.com/apache/datafusion/pull/20271) (2010YOUY01) +- [main] Bump to 52.3.0 and changelog (#20790) [#20849](https://github.com/apache/datafusion/pull/20849) (alamb) +- refactor: Improve `SessionContext::parse_duration` API [#20816](https://github.com/apache/datafusion/pull/20816) (erenavsarogullari) +- docs: in release email, be specific about changelog location [#20975](https://github.com/apache/datafusion/pull/20975) (kevinjqliu) +- optimizer: Add configuration to disable join reordering [#21072](https://github.com/apache/datafusion/pull/21072) (2010YOUY01) +- docs: Improve getting started and testing guides for humans and agents [#20970](https://github.com/apache/datafusion/pull/20970) (alamb) +- docs: clarify NULL handling for array_remove functions (#21014) [#21018](https://github.com/apache/datafusion/pull/21018) (Xavrir) +- chore: Add `substr()` benchmarks, refactor [#20803](https://github.com/apache/datafusion/pull/20803) (neilconway) +- docs: Document the TableProvider evaluation order for filter, limit and projection [#21091](https://github.com/apache/datafusion/pull/21091) (alamb) +- Add `arrow_try_cast` UDF [#21130](https://github.com/apache/datafusion/pull/21130) (adriangb) +- docs: Add explicit fmt and clippy commands to AGENTS.md [#21171](https://github.com/apache/datafusion/pull/21171) (zhuqi-lucas) +- docs: add KalamDB to known users [#21181](https://github.com/apache/datafusion/pull/21181) (jamals86) +- [main] Update version to 53.0.0 and bring changelog [#21189](https://github.com/apache/datafusion/pull/21189) (alamb) +- Migrate Avro reader to arrow-avro and remove internal conversion code [#17861](https://github.com/apache/datafusion/pull/17861) (getChan) +- Add metric category filtering for EXPLAIN ANALYZE [#21160](https://github.com/apache/datafusion/pull/21160) (adriangb) +- docs: Add `RESET` Command Documentation [#21245](https://github.com/apache/datafusion/pull/21245) (erenavsarogullari) +- chore: fix upgrade guide link for object_store release notes [#21283](https://github.com/apache/datafusion/pull/21283) (haohuaijin) +- doc: Add documentation explaining the behavior of `null` values ​​in struct comparisons [#21226](https://github.com/apache/datafusion/pull/21226) (xiedeyantu) +- [docs] Add weekly sync details to contributor communication guide [#21298](https://github.com/apache/datafusion/pull/21298) (alamb) +- [docs] add sql example to timestamp/datetime docs for time zone [#21082](https://github.com/apache/datafusion/pull/21082) (buraksenn) +- Update documentation with recent blogs and events [#21462](https://github.com/apache/datafusion/pull/21462) (alamb) +- Update 53 upgrade guide to note release, other changes [#21449](https://github.com/apache/datafusion/pull/21449) (alamb) +- docs: Incorporate writing table provider blog post to user documentation [#21398](https://github.com/apache/datafusion/pull/21398) (buraksenn) +- remove as_any from TableProvider, SchemaProvider, CatalogProvider, and CatalogProviderList [#21346](https://github.com/apache/datafusion/pull/21346) (timsaucer) +- port 52.5.0 changelog to main [#21553](https://github.com/apache/datafusion/pull/21553) (alamb) +- Add `arrow_field(expr)` scalar UDF [#21389](https://github.com/apache/datafusion/pull/21389) (adriangb) +- Reorder `cargo publish` commands by dependency [#21552](https://github.com/apache/datafusion/pull/21552) (alamb) +- chore(deps): update jinja2 requirement from <4,>=3.1 to >=3.1.6,<4 in /docs [#21606](https://github.com/apache/datafusion/pull/21606) (dependabot[bot]) +- chore(deps): update pydata-sphinx-theme requirement from <1,>=0.16 to >=0.17.0,<1 in /docs [#21609](https://github.com/apache/datafusion/pull/21609) (dependabot[bot]) +- Add release management page to the documentation [#21001](https://github.com/apache/datafusion/pull/21001) (alamb) +- Perf: Window topn optimisation [#21479](https://github.com/apache/datafusion/pull/21479) (SubhamSinghal) +- chore(deps): update setuptools requirement from <83,>=82 to >=82.0.1,<83 in /docs [#21607](https://github.com/apache/datafusion/pull/21607) (dependabot[bot]) +- chore(deps): update maturin requirement from <2,>=1.11 to >=1.13.1,<2 in /docs [#21608](https://github.com/apache/datafusion/pull/21608) (dependabot[bot]) +- docs: Update `map_extract` examples [#21360](https://github.com/apache/datafusion/pull/21360) (nuno-faria) +- docs: add April 2026 readings and meetup links [#21644](https://github.com/apache/datafusion/pull/21644) (alamb) +- chore: backport version from `branch-53`, update some dependencies [#21708](https://github.com/apache/datafusion/pull/21708) (comphead) +- chore: add `array_remove_*` NULL handling changes to `Upgrade Guide` [#21769](https://github.com/apache/datafusion/pull/21769) (comphead) +- docs: fix some comments on query_planning example [#21783](https://github.com/apache/datafusion/pull/21783) (jotare) +- docs: fix typos in documentation [#21875](https://github.com/apache/datafusion/pull/21875) (jx2lee) +- docs: refresh CLI usage output in the user guide [#21874](https://github.com/apache/datafusion/pull/21874) (jx2lee) +- docs: clarify ExecutionProps and TaskContext docs [#21872](https://github.com/apache/datafusion/pull/21872) (alamb) +- chore: add internal markdown link check [#21831](https://github.com/apache/datafusion/pull/21831) (Geethapranay1) +- Update documentation for PhysicalExpr::evaluate_bounds [#21879](https://github.com/apache/datafusion/pull/21879) (alamb) +- chore(deps): update pydata-sphinx-theme requirement from <1,>=0.17.0 to >=0.17.1,<1 in /docs [#21889](https://github.com/apache/datafusion/pull/21889) (dependabot[bot]) +- docs(optimizer): add generated optimizer rules reference [#21824](https://github.com/apache/datafusion/pull/21824) (kumarUjjawal) +- add any_match higher-order function [#21903](https://github.com/apache/datafusion/pull/21903) (LiaCastaneda) +- docs: update commiter list [#21978](https://github.com/apache/datafusion/pull/21978) (coderfender) +- chore: update PMC/committer list [#21989](https://github.com/apache/datafusion/pull/21989) (comphead) +- Support '0' value for parse_capacity_limit() [#22014](https://github.com/apache/datafusion/pull/22014) (mkleen) +- docs: add llms.txt ecosystem hub at site root [#22003](https://github.com/apache/datafusion/pull/22003) (timsaucer) +- chore(deps): update maturin requirement from <2,>=1.13.1 to >=1.13.3,<2 in /docs [#22127](https://github.com/apache/datafusion/pull/22127) (dependabot[bot]) +- fix `date_part('isodow')` [#22116](https://github.com/apache/datafusion/pull/22116) (sdf-jkl) +- docs: updating arrays_zip output field naming [#22133](https://github.com/apache/datafusion/pull/22133) (timsaucer) +- Add rand() alias for random() [#22147](https://github.com/apache/datafusion/pull/22147) (xiedeyantu) +- chore: Update Rust toolchain to 1.95 [#22177](https://github.com/apache/datafusion/pull/22177) (Dandandan) +- docs: add DataFusion Java to subproject listings [#22149](https://github.com/apache/datafusion/pull/22149) (andygrove) +- Fix: deadlink in "Concepts, Reading, Events" page to DataFusion blog [#22325](https://github.com/apache/datafusion/pull/22325) (JarroVGIT) +- fixing factorial negative values [#22278](https://github.com/apache/datafusion/pull/22278) (raushanprabhakar1) +- docs(optimizer): Fix PushDownFilter doc typos. [#22320](https://github.com/apache/datafusion/pull/22320) (JSOD11) +- minor: add higher-order function methods to SessionContext [#21950](https://github.com/apache/datafusion/pull/21950) (gstvg) +- Add higher-order functions changes to upgrade guide [#22107](https://github.com/apache/datafusion/pull/22107) (gstvg) +- chore(deps): update myst-parser requirement from <6,>=5 to >=5.1.0,<6 in /docs [#22378](https://github.com/apache/datafusion/pull/22378) (dependabot[bot]) +- feat(functions-nested): add array_filter higher-order function [#21895](https://github.com/apache/datafusion/pull/21895) (ologlogn) +- Add SQL as a category in breaking API change policy [#22179](https://github.com/apache/datafusion/pull/22179) (alamb) +- [branch-54] Bump to version 54.0.0 [#22396](https://github.com/apache/datafusion/pull/22396) (mbutrovich) +- [branch-54] Revert "Add `ExecutionPlan::apply_expressions()` (#20337)" (#22437) [#22445](https://github.com/apache/datafusion/pull/22445) (alamb) +- [branch-54] Gate new ScalarSubqueryExec node behind session property (#22530) [#22690](https://github.com/apache/datafusion/pull/22690) (LiaCastaneda) + +**Other:** + +- Add metrics for parquet sink [#20307](https://github.com/apache/datafusion/pull/20307) (xudong963) +- Extend dynamic filter to joins that preserve probe side ON [#20447](https://github.com/apache/datafusion/pull/20447) (helgikrs) +- Improve sqllogicteset speed by creating only a single large file rather than 2 [#20586](https://github.com/apache/datafusion/pull/20586) (Tim-53) +- cli: Fix datafusion-cli hint edge cases [#20609](https://github.com/apache/datafusion/pull/20609) (comphead) +- Speedup sqllogictests by running long running tests first [#20576](https://github.com/apache/datafusion/pull/20576) (alamb) +- Fix custom metric display [#20643](https://github.com/apache/datafusion/pull/20643) (gabotechs) +- refactor: Set expected runtime config in error message when the used disk space during the spilling process has exceeded the allocation limit [#20375](https://github.com/apache/datafusion/pull/20375) (erenavsarogullari) +- more families for the CI [#20663](https://github.com/apache/datafusion/pull/20663) (blaginin) +- CI: Add CodeQL workflow for GitHub Actions security scanning [#20636](https://github.com/apache/datafusion/pull/20636) (kevinjqliu) +- chore(deps): bump astral-sh/setup-uv from 7.3.0 to 7.3.1 [#20660](https://github.com/apache/datafusion/pull/20660) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.68.8 to 2.68.16 [#20661](https://github.com/apache/datafusion/pull/20661) (dependabot[bot]) +- Improve formatting of datatypes [#20605](https://github.com/apache/datafusion/pull/20605) (emilk) +- Add explain plans for ClickBench queries [#20666](https://github.com/apache/datafusion/pull/20666) (alamb) +- Add files_processed and files_scanned metrics to FileStreamMetrics [#20592](https://github.com/apache/datafusion/pull/20592) (adriangb) +- Speedup push_down_filter_regression.slt by using uncompressed parquet [#20652](https://github.com/apache/datafusion/pull/20652) (alamb) +- Implement cardinality_effect for window execs and UnionExec [#20321](https://github.com/apache/datafusion/pull/20321) (getChan) +- ci: Harden labeler workflow, remove unnecessary checkout from pull_request_target job [#20637](https://github.com/apache/datafusion/pull/20637) (kevinjqliu) +- Add tests for sqllogictest prioritization [#20656](https://github.com/apache/datafusion/pull/20656) (alamb) +- correct parquet leaf index mapping when schema contains struct cols [#20698](https://github.com/apache/datafusion/pull/20698) (friendlymatthew) +- Reattach parquet metadata cache after deserializing in datafusion-proto [#20574](https://github.com/apache/datafusion/pull/20574) (nathanb9) +- Wire up with_new_state with DataSource [#20718](https://github.com/apache/datafusion/pull/20718) (gabotechs) +- chore: Enable `assigning_clones` clippy lint [#20670](https://github.com/apache/datafusion/pull/20670) (neilconway) +- FFI_TableOptions are using default values only [#20721](https://github.com/apache/datafusion/pull/20721) (timsaucer) +- Improve documentation for `AggregateUdfImpl::simplify` and `WindowUDFImpl::simplify` [#20712](https://github.com/apache/datafusion/pull/20712) (alamb) +- Fix test that's broken on Windows due to naive path handling [#20692](https://github.com/apache/datafusion/pull/20692) (Rafferty97) +- Fix DELETE/UPDATE filter extraction when predicates are pushed down into TableScan [#19884](https://github.com/apache/datafusion/pull/19884) (kosiew) +- use linker optimization for extended sqllogictests [#20740](https://github.com/apache/datafusion/pull/20740) (blaginin) +- Push even local limits past windows [#20752](https://github.com/apache/datafusion/pull/20752) (avantgardnerio) +- Add case-heavy LEFT JOIN benchmark and debug timing/logging for PushDownFilter hot paths [#20664](https://github.com/apache/datafusion/pull/20664) (kosiew) +- Fix repartition from dropping data when spilling [#20672](https://github.com/apache/datafusion/pull/20672) (xanderbailey) +- test: Add `datafusion-cli` `fair` and `unbounded` memory-pool test coverage [#20565](https://github.com/apache/datafusion/pull/20565) (erenavsarogullari) +- ser/de fetch in FilterExec [#20738](https://github.com/apache/datafusion/pull/20738) (haohuaijin) +- Add tests for simplifying multiple aggregate expressions [#20723](https://github.com/apache/datafusion/pull/20723) (alamb) +- Update reverse UDF to emit utf8view when input is utf8view [#20604](https://github.com/apache/datafusion/pull/20604) (Omega359) +- Make lower and upper emit Utf8View for Utf8View input [#20616](https://github.com/apache/datafusion/pull/20616) (kumarUjjawal) +- Fix FilterExec converting Absent column stats to Exact(NULL) [#20391](https://github.com/apache/datafusion/pull/20391) (fwojciec) +- Clean up date_part preimage implementation [#20350](https://github.com/apache/datafusion/pull/20350) (sdf-jkl) +- Make Physical CastExpr Field-aware and unify cast semantics across physical expressions [#20814](https://github.com/apache/datafusion/pull/20814) (kosiew) +- Pass ConfigOptions to scalar UDFs via FFI [#20454](https://github.com/apache/datafusion/pull/20454) (timsaucer) +- [datafusion-cli] Replace mutex with AtomicU64 for stream duration tracking in instrumentedObjectStore [#20802](https://github.com/apache/datafusion/pull/20802) (buraksenn) +- Make translate emit Utf8View for Utf8View input [#20624](https://github.com/apache/datafusion/pull/20624) (shivaaang) +- Allow filters on struct fields to be pushed down into Parquet scan [#20822](https://github.com/apache/datafusion/pull/20822) (friendlymatthew) +- Used constant with mapping instead of write! to display scalar value bytes [#20719](https://github.com/apache/datafusion/pull/20719) (buraksenn) +- chore(deps): bump taiki-e/install-action from 2.68.16 to 2.68.25 [#20842](https://github.com/apache/datafusion/pull/20842) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.32.5 to 4.32.6 [#20843](https://github.com/apache/datafusion/pull/20843) (dependabot[bot]) +- chore: Ignore RUSTSEC-2024-0421 [#20850](https://github.com/apache/datafusion/pull/20850) (comphead) +- chore(deps): bump quinn-proto from 0.11.13 to 0.11.14 [#20859](https://github.com/apache/datafusion/pull/20859) (dependabot[bot]) +- Use `ParquetPushDecoder` in `ParquetOpener` [#20839](https://github.com/apache/datafusion/pull/20839) (Dandandan) +- [Minor] Remove redundant ProjectionExec nodes in sort-based plans [#20780](https://github.com/apache/datafusion/pull/20780) (Dandandan) +- impl ser/de for preserve_order in RepartitionExec [#20798](https://github.com/apache/datafusion/pull/20798) (haohuaijin) +- Fix FileStream scanning_total to include sync next-file open time [#20627](https://github.com/apache/datafusion/pull/20627) (RatulDawar) +- chore: Ignore RUSTSEC-2024-0014 [#20862](https://github.com/apache/datafusion/pull/20862) (comphead) +- chore: clean up dependencies [#20861](https://github.com/apache/datafusion/pull/20861) (comphead) +- Add benchmark for struct field filter pushdown in Parquet [#20829](https://github.com/apache/datafusion/pull/20829) (friendlymatthew) +- Add Null Type Coercions for Placeholders [#20543](https://github.com/apache/datafusion/pull/20543) (cetra3) +- Minor: Deprecate unused `PartitionedFileStream` [#20869](https://github.com/apache/datafusion/pull/20869) (alamb) +- chore(deps): bump substrait from 0.62 to 0.63.0 [#20876](https://github.com/apache/datafusion/pull/20876) (benbellick) +- [Minor] propagate distinct_count as inexact through unions [#20846](https://github.com/apache/datafusion/pull/20846) (buraksenn) +- try to remove redundant alias in expression rewriter and select [#20867](https://github.com/apache/datafusion/pull/20867) (buraksenn) +- Fix duplicate group keys after hash aggregation spill (#20724) [#20858](https://github.com/apache/datafusion/pull/20858) (gboucher90) +- Include .proto files in datafusion-proto-common distribution [#20921](https://github.com/apache/datafusion/pull/20921) (haohuaijin) +- Check sqllogictests for any dangling config settings (#17914) [#20838](https://github.com/apache/datafusion/pull/20838) (cj-zhukov) +- Add support for ListView in unnest [#20760](https://github.com/apache/datafusion/pull/20760) (brancz) +- Project only accessed struct leaves in Parquet row filter pushdown [#20854](https://github.com/apache/datafusion/pull/20854) (friendlymatthew) +- minor: Move PreparedAccessPlan to same module as ParquetAccessPlan [#20929](https://github.com/apache/datafusion/pull/20929) (alamb) +- chore(deps): bump pyjwt from 2.11.0 to 2.12.0 [#20938](https://github.com/apache/datafusion/pull/20938) (dependabot[bot]) +- Rewrite `SUM(expr + scalar)` --> `SUM(expr) + scalar*COUNT(expr)` [#20749](https://github.com/apache/datafusion/pull/20749) (alamb) +- Add AGENTS.md / CLAUDE.md [#20939](https://github.com/apache/datafusion/pull/20939) (Dandandan) +- Support `columns_sorted` in row_filters [#20497](https://github.com/apache/datafusion/pull/20497) (sdf-jkl) +- Add --simulate-latency / SIMULATE_LATENCY option to dfbench / ./bench.sh [#20954](https://github.com/apache/datafusion/pull/20954) (Dandandan) +- Minor: make signatures of `SessionContext::register_*` methods consistent [#20873](https://github.com/apache/datafusion/pull/20873) (alexandreyc) +- test: add reproducer for Dictionary InList pushdown type mismatch (#2… [#20960](https://github.com/apache/datafusion/pull/20960) (erratic-pattern) +- Extract shared `ParquetReadPlan` for leaf column resolution [#20913](https://github.com/apache/datafusion/pull/20913) (friendlymatthew) +- chore: Remove usage of `paste` crate [#20946](https://github.com/apache/datafusion/pull/20946) (coderfender) +- Use exact distinct_count from statistics if exists for `COUNT(DISTINCT column))` calculations [#20845](https://github.com/apache/datafusion/pull/20845) (buraksenn) +- thin-ci [#20972](https://github.com/apache/datafusion/pull/20972) (blaginin) +- chore(deps): bump lz4_flex from 0.12.0 to 0.12.1 [#20973](https://github.com/apache/datafusion/pull/20973) (dependabot[bot]) +- Fix decimal log precision for non-power values [#20433](https://github.com/apache/datafusion/pull/20433) (kumarUjjawal) +- chore(deps): bump Swatinem/rust-cache from 2.8.2 to 2.9.1 [#20979](https://github.com/apache/datafusion/pull/20979) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.68.25 to 2.68.34 [#20983](https://github.com/apache/datafusion/pull/20983) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.32.6 to 4.33.0 [#20982](https://github.com/apache/datafusion/pull/20982) (dependabot[bot]) +- chore(deps): bump astral-sh/setup-uv from 7.3.1 to 7.6.0 [#20981](https://github.com/apache/datafusion/pull/20981) (dependabot[bot]) +- chore(deps): bump runs-on/action from 2.0.3 to 2.1.0 [#20980](https://github.com/apache/datafusion/pull/20980) (dependabot[bot]) +- [Minor] Update Cargo.lock, Fix Tokio minor breaking change [#20978](https://github.com/apache/datafusion/pull/20978) (Dandandan) +- chore(deps): Revert "chore(deps): bump runs-on/action from 2.0.3 to 2.1.0 (#20980)" [#21002](https://github.com/apache/datafusion/pull/21002) (mbutrovich) +- bug: fix `array_remove_*` with NULLS [#21013](https://github.com/apache/datafusion/pull/21013) (comphead) +- Simplify logic for memory pressure partial emit from ordered group by [#20559](https://github.com/apache/datafusion/pull/20559) (alamb) +- Fix memory reservation starvation in sort-merge [#20642](https://github.com/apache/datafusion/pull/20642) (xudong963) +- infra: automatically delete branch on pr merge [#21033](https://github.com/apache/datafusion/pull/21033) (kevinjqliu) +- Add support for nested lists in substrait consumer [#20953](https://github.com/apache/datafusion/pull/20953) (alexanderbianchi) +- build: update Rust toolchain version to 1.94.0 [#21045](https://github.com/apache/datafusion/pull/21045) (dariocurr) +- chore: Cleanup fully-qualified ScalarFunctionArgs [#20804](https://github.com/apache/datafusion/pull/20804) (neilconway) +- Support '>', '<', '>=', '<=', '<>' in any operator [#20830](https://github.com/apache/datafusion/pull/20830) (buraksenn) +- keep fetch when merge FilterExec in FilterPushdown [#21070](https://github.com/apache/datafusion/pull/21070) (haohuaijin) +- Fix Subtraction overflow in `max_distinct_count` when hash join has a pushed-down limit [#20799](https://github.com/apache/datafusion/pull/20799) (KARTIK64-rgb) +- Restore Sort unparser guard for correct ORDER BY placement [#20658](https://github.com/apache/datafusion/pull/20658) (krinart) +- chore(deps): bump rustls-webpki from 0.103.9 to 0.103.10 [#21089](https://github.com/apache/datafusion/pull/21089) (dependabot[bot]) +- chore: Remove duplicate imports in test code [#21061](https://github.com/apache/datafusion/pull/21061) (neilconway) +- test: update sqllogictest expectation for negation type coercion [#21102](https://github.com/apache/datafusion/pull/21102) (myandpr) +- fix[physical-expr-adapter]: support casting structs nested inside complex types [#20907](https://github.com/apache/datafusion/pull/20907) (asubiotto) +- Fix index panic in unparser with mismatched stacked projections [#21094](https://github.com/apache/datafusion/pull/21094) (friendlymatthew) +- chore: Fix all sqllogictest dangling configs [#21108](https://github.com/apache/datafusion/pull/21108) (2010YOUY01) +- Preserve SPM when parent maintains input order [#21097](https://github.com/apache/datafusion/pull/21097) (rkrishn7) +- chore: update testcontainers and astral-tokio-tar for cargo audit [#21114](https://github.com/apache/datafusion/pull/21114) (getChan) +- Spark soundex function implementation [#20725](https://github.com/apache/datafusion/pull/20725) (kazantsev-maksim) +- chore(deps): bump env_logger from 0.11.9 to 0.11.10 in the all-other-cargo-deps group across 1 directory [#21136](https://github.com/apache/datafusion/pull/21136) (dependabot[bot]) +- Fix `elapsed_compute` metric for Parquet DataSourceExec [#20767](https://github.com/apache/datafusion/pull/20767) (ernestprovo23) +- chore(deps): bump taiki-e/install-action from 2.68.34 to 2.69.7 [#21133](https://github.com/apache/datafusion/pull/21133) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.33.0 to 4.34.1 [#21132](https://github.com/apache/datafusion/pull/21132) (dependabot[bot]) +- Update to arrow/parquet `58.1.0` [#21044](https://github.com/apache/datafusion/pull/21044) (alamb) +- Simplify sqllogictest timing summary to boolean flag and remove top-N modes [#20598](https://github.com/apache/datafusion/pull/20598) (kosiew) +- Substrait join consumer should not merge nullability of join keys [#21121](https://github.com/apache/datafusion/pull/21121) (hareshkh) +- Enable debug assertions in CI. [#20832](https://github.com/apache/datafusion/pull/20832) (stuhood) +- chore(deps): bump requests from 2.32.5 to 2.33.0 [#21153](https://github.com/apache/datafusion/pull/21153) (dependabot[bot]) +- feat : support spark compatible int to timestamp cast [#20555](https://github.com/apache/datafusion/pull/20555) (coderfender) +- [Minor]: support window functions in order by expressions [#20963](https://github.com/apache/datafusion/pull/20963) (buraksenn) +- chore: Optimize schema rewriter usages [#21158](https://github.com/apache/datafusion/pull/21158) (comphead) +- Add benchmarks for Parquet struct leaf-level projection pruning [#21180](https://github.com/apache/datafusion/pull/21180) (friendlymatthew) +- chore: re-export projection in datafusion::datasource [#21185](https://github.com/apache/datafusion/pull/21185) (rluvaton) +- test: add SMJ benchmarks from #21184 [#21188](https://github.com/apache/datafusion/pull/21188) (mbutrovich) +- Fix sort merge interleave overflow [#20922](https://github.com/apache/datafusion/pull/20922) (xudong963) +- Reduce parquet struct projection benchmark data volume [#21187](https://github.com/apache/datafusion/pull/21187) (friendlymatthew) +- Minor: compute qualify window expressions only when QUALIFY clause is present [#21173](https://github.com/apache/datafusion/pull/21173) (buraksenn) +- fix[physical-plan/aggregates]: fix grouping by Ree [#21195](https://github.com/apache/datafusion/pull/21195) (asubiotto) +- [main] add 52.4.0 changelog [#21053](https://github.com/apache/datafusion/pull/21053) (alamb) +- Use leaf level `ProjectionMask` for parquet projections [#20925](https://github.com/apache/datafusion/pull/20925) (friendlymatthew) +- test: scale remaining sort-merge join (SMJ) benchmark queries [#21200](https://github.com/apache/datafusion/pull/21200) (mbutrovich) +- Fix: MemTable LIMIT ignored with reordered projections [#21177](https://github.com/apache/datafusion/pull/21177) (RamakrishnaChilaka) +- No cargo test for `sort_mem_validation` [#21222](https://github.com/apache/datafusion/pull/21222) (blaginin) +- Fix/support duplicate column names #6543 [#21126](https://github.com/apache/datafusion/pull/21126) (RafaelHerrero) +- Use spot instances for extended tests [#21221](https://github.com/apache/datafusion/pull/21221) (blaginin) +- chore: Cleanup Cargo profiles [#21214](https://github.com/apache/datafusion/pull/21214) (neilconway) +- chore(benchmark): Fix/update compile profile benchmark [#21223](https://github.com/apache/datafusion/pull/21223) (2010YOUY01) +- Basic Extension Type Registry Implementation [#20312](https://github.com/apache/datafusion/pull/20312) (tschwarzinger) +- chore(deps): bump serialize-javascript, terser-webpack-plugin and copy-webpack-plugin in /datafusion/wasmtest/datafusion-wasm-app [#21235](https://github.com/apache/datafusion/pull/21235) (dependabot[bot]) +- chore(deps-dev): bump node-forge from 1.3.2 to 1.4.0 in /datafusion/wasmtest/datafusion-wasm-app [#21225](https://github.com/apache/datafusion/pull/21225) (dependabot[bot]) +- chore(deps): bump cryptography from 46.0.5 to 46.0.6 [#21224](https://github.com/apache/datafusion/pull/21224) (dependabot[bot]) +- Fix FilterExec tree render missing fetch display [#21230](https://github.com/apache/datafusion/pull/21230) (zhuqi-lucas) +- ci: use ubuntu-slim runner for lightweight CI jobs [#21252](https://github.com/apache/datafusion/pull/21252) (CuteChuanChuan) +- kill `check_run_id` and `pr_number` from extended tests [#21228](https://github.com/apache/datafusion/pull/21228) (blaginin) +- [Minor] add non topk benchmarks for utf8/utf8view string aggregates [#21073](https://github.com/apache/datafusion/pull/21073) (buraksenn) +- ci: Add datafusion/sql as a folder to trigger extended tests for on changes [#21255](https://github.com/apache/datafusion/pull/21255) (mbutrovich) +- Misc minor optimization in the Physical Optimizer [#21216](https://github.com/apache/datafusion/pull/21216) (AdamGS) +- chore: Replace `TryInto` impl by `TryFrom` [#21203](https://github.com/apache/datafusion/pull/21203) (Tpt) +- Refactor parquet datasource into an explicit state machine [#21190](https://github.com/apache/datafusion/pull/21190) (alamb) +- Add flat vs. struct field projection benchmarks [#21257](https://github.com/apache/datafusion/pull/21257) (friendlymatthew) +- Refactor: expose predicate constant inference from physical-expr [#21167](https://github.com/apache/datafusion/pull/21167) (xudong963) +- Add end-to-end Parquet tests for List and LargeList struct schema evolution [#20840](https://github.com/apache/datafusion/pull/20840) (kosiew) +- chore(deps): bump taiki-e/install-action from 2.69.7 to 2.70.3 [#21271](https://github.com/apache/datafusion/pull/21271) (dependabot[bot]) +- chore(deps): bump rustyline from 17.0.2 to 18.0.0 [#21276](https://github.com/apache/datafusion/pull/21276) (dependabot[bot]) +- chore(deps): bump ctor from 0.6.3 to 0.8.0 [#21282](https://github.com/apache/datafusion/pull/21282) (dependabot[bot]) +- chore(deps): bump snmalloc-rs from 0.3.8 to 0.7.4 [#21280](https://github.com/apache/datafusion/pull/21280) (dependabot[bot]) +- chore(deps): bump sha1 from 0.10.6 to 0.11.0 [#21277](https://github.com/apache/datafusion/pull/21277) (dependabot[bot]) +- chore(deps): bump astral-sh/setup-uv from 7.6.0 to 8.0.0 [#21272](https://github.com/apache/datafusion/pull/21272) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.34.1 to 4.35.1 [#21273](https://github.com/apache/datafusion/pull/21273) (dependabot[bot]) +- chore(deps): bump pygments from 2.19.2 to 2.20.0 [#21256](https://github.com/apache/datafusion/pull/21256) (dependabot[bot]) +- feat(memory_pool): add `TrackConsumersPool::metrics()` to expose cons… [#21147](https://github.com/apache/datafusion/pull/21147) (bert-beyondloops) +- Update repeat UDF to emit utf8view when input is utf8view [#20645](https://github.com/apache/datafusion/pull/20645) (Omega359) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 7 updates [#21274](https://github.com/apache/datafusion/pull/21274) (dependabot[bot]) +- chore(deps): bump runs-on/action from 2.0.3 to 2.1.0 [#21134](https://github.com/apache/datafusion/pull/21134) (dependabot[bot]) +- chore: add `.claude/settings.local.json` to `.gitignore` [#21312](https://github.com/apache/datafusion/pull/21312) (jonahgao) +- Add `FileStreamBuilder` for creating FileStreams [#21261](https://github.com/apache/datafusion/pull/21261) (alamb) +- refactor: Split Parquet BloomFilter CPU and IO into separate states [#21285](https://github.com/apache/datafusion/pull/21285) (alamb) +- chore(deps): bump object_store from 0.13.1 to 0.13.2 [#21275](https://github.com/apache/datafusion/pull/21275) (dependabot[bot]) +- Merge queue: make dev checks required + add .asf.yaml validation [#21239](https://github.com/apache/datafusion/pull/21239) (blaginin) +- Adds INList and Between expr to skip outer join [#21303](https://github.com/apache/datafusion/pull/21303) (SubhamSinghal) +- No merge group for rust.yml yet [#21343](https://github.com/apache/datafusion/pull/21343) (blaginin) +- Disallow order by within ordered-set aggregate functions argument lists [#20421](https://github.com/apache/datafusion/pull/20421) (cj-zhukov) +- chore: Fix clippy and CI [#21287](https://github.com/apache/datafusion/pull/21287) (comphead) +- Split FileStreamMetrics into its own module [#21340](https://github.com/apache/datafusion/pull/21340) (alamb) +- Skip probe-side consumption when hash join build side is empty [#21068](https://github.com/apache/datafusion/pull/21068) (kosiew) +- Use ParquetMetaDataPushDecoder instead of ParquetMetaDataReader [#21357](https://github.com/apache/datafusion/pull/21357) (Dandandan) +- Eliminate redundant `ProjectionExec`s [#21333](https://github.com/apache/datafusion/pull/21333) (Dandandan) +- Minor: add tests for regexp_replace and capture groups [#21413](https://github.com/apache/datafusion/pull/21413) (alamb) +- bench: add benchmarks for first_value, last_value [#21409](https://github.com/apache/datafusion/pull/21409) (theirix) +- chore(deps): bump the all-other-cargo-deps group with 4 updates [#21435](https://github.com/apache/datafusion/pull/21435) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.70.3 to 2.74.0 [#21434](https://github.com/apache/datafusion/pull/21434) (dependabot[bot]) +- test: Add `datafusion.format.*` configs test coverage [#21355](https://github.com/apache/datafusion/pull/21355) (erenavsarogullari) +- Estimate aggregate output rows using existing NDV statistics [#20926](https://github.com/apache/datafusion/pull/20926) (buraksenn) +- Follow-up: remove interleave panic recovery after Arrow 58.1.0 [#21436](https://github.com/apache/datafusion/pull/21436) (xudong963) +- writing table to parquet followed by read and schema check [#21444](https://github.com/apache/datafusion/pull/21444) (Rich-T-kid) +- chore(deps): bump cryptography from 46.0.6 to 46.0.7 [#21489](https://github.com/apache/datafusion/pull/21489) (dependabot[bot]) +- Preserve logical cast field semantics during physical lowering with field-aware CastExpr [#20836](https://github.com/apache/datafusion/pull/20836) (kosiew) +- Add more regexp_replace test coverage [#21485](https://github.com/apache/datafusion/pull/21485) (alamb) +- Introduce Morselizer API, rewrite `ParquetOpener` to `ParquetMorselizer` [#21327](https://github.com/apache/datafusion/pull/21327) (alamb) +- chore: create benches small ints for count_distinct [#21521](https://github.com/apache/datafusion/pull/21521) (coderfender) +- refactor: extract sort pushdown logic from FileScanConfig into separate module [#21457](https://github.com/apache/datafusion/pull/21457) (zhuqi-lucas) +- chore: Add array_slice tests for overlapping nulls across inputs [#21540](https://github.com/apache/datafusion/pull/21540) (neilconway) +- Migrate PhysicalExprAdapter to unified CastExpr and remove CastColumnExpr usage [#21493](https://github.com/apache/datafusion/pull/21493) (kosiew) +- Unify cast handling by removing `CastColumnExpr` branches in pruning and ordering equivalence [#21545](https://github.com/apache/datafusion/pull/21545) (kosiew) +- [datafusion-spark] Add Spark-compatible ceil function [#20593](https://github.com/apache/datafusion/pull/20593) (shivbhatia10) +- sql: render PostgreSQL array literals as ARRAY[...] in unparser [#21513](https://github.com/apache/datafusion/pull/21513) (xiedeyantu) +- physical_optimizer: preserve_file_partitions when num file groups < target_partitions [#21533](https://github.com/apache/datafusion/pull/21533) (jayshrivastava) +- EliminateOuterJoin with Like, IsTrue, IsFalse, IsNotUnknown [#21549](https://github.com/apache/datafusion/pull/21549) (SubhamSinghal) +- chore(deps): bump hashbrown from 0.16.1 to 0.17.0 [#21611](https://github.com/apache/datafusion/pull/21611) (dependabot[bot]) +- chore(deps): bump ctor from 0.8.0 to 0.10.0 [#21612](https://github.com/apache/datafusion/pull/21612) (dependabot[bot]) +- Rewrite FileStream in terms of Morsel API [#21342](https://github.com/apache/datafusion/pull/21342) (alamb) +- Consolidate special case `regexp_match` logic [#21486](https://github.com/apache/datafusion/pull/21486) (alamb) +- chore(deps): bump taiki-e/install-action from 2.74.0 to 2.75.10 [#21605](https://github.com/apache/datafusion/pull/21605) (dependabot[bot]) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 3 updates [#21610](https://github.com/apache/datafusion/pull/21610) (dependabot[bot]) +- bench: first_last remove noisy benchmarks, add update_batch [#21487](https://github.com/apache/datafusion/pull/21487) (theirix) +- chore: Fix `typo` problems [#21495](https://github.com/apache/datafusion/pull/21495) (erenavsarogullari) +- chore(deps-dev): bump follow-redirects from 1.15.6 to 1.16.0 in /datafusion/wasmtest/datafusion-wasm-app [#21601](https://github.com/apache/datafusion/pull/21601) (dependabot[bot]) +- bench: Scale sort benchmarks to 1M rows to exercise merge path [#21630](https://github.com/apache/datafusion/pull/21630) (mbutrovich) +- Port filter_pushdown.rs async tests to sqllogictest [#21620](https://github.com/apache/datafusion/pull/21620) (adriangb) +- chore: fix cargo audit and dependencies check on main [#21655](https://github.com/apache/datafusion/pull/21655) (alamb) +- Spark make_valid_utf8 function implementation [#20633](https://github.com/apache/datafusion/pull/20633) (kazantsev-maksim) +- chore(deps): update tokio from 1.51 to 1.52 [#21670](https://github.com/apache/datafusion/pull/21670) (ahmed-mez) +- Use ListArray nullability instead of offsets for `array_element`, `array_any_value`. [#21672](https://github.com/apache/datafusion/pull/21672) (tabac) +- chore: breakdown `array.slt` into smaller files [#21658](https://github.com/apache/datafusion/pull/21658) (comphead) +- chore: Add more tests with `GROUP BY` to test spark `collect_set` [#21659](https://github.com/apache/datafusion/pull/21659) (comphead) +- Add strategy-focused InList benchmarks [#21648](https://github.com/apache/datafusion/pull/21648) (geoffreyclaude) +- Fix massive spill files for StringView/BinaryView columns II [#21633](https://github.com/apache/datafusion/pull/21633) (adriangb) +- chore: Backport 53.1.0 changelog [#21686](https://github.com/apache/datafusion/pull/21686) (comphead) +- refactor: Introduce SpillState enum for memory-limited NLJ execution [#21636](https://github.com/apache/datafusion/pull/21636) (viirya) +- Support Date32/Date64 in unwrap_cast optimization [#21665](https://github.com/apache/datafusion/pull/21665) (Dandandan) +- feat[expr-common]: add REE arithmetic coercion for numeric and decimal [#21179](https://github.com/apache/datafusion/pull/21179) (asubiotto) +- Make `test_display_pg_json` pass regardless of build setup and dependencies [#21502](https://github.com/apache/datafusion/pull/21502) (AdamGS) +- refactor: Share left-side spill file across partitions on OOM fallback [#21699](https://github.com/apache/datafusion/pull/21699) (viirya) +- Spark is_valid_utf8 function implementation [#21627](https://github.com/apache/datafusion/pull/21627) (kazantsev-maksim) +- chore: use bench array helpers from Arrow bench_util [#21544](https://github.com/apache/datafusion/pull/21544) (theirix) +- chore: add count distinct group benchmarks [#21575](https://github.com/apache/datafusion/pull/21575) (coderfender) +- minor: More comments to `read_spill_as_stream` [#21713](https://github.com/apache/datafusion/pull/21713) (2010YOUY01) +- Dynamic work scheduling in FileStream [#21351](https://github.com/apache/datafusion/pull/21351) (alamb) +- chore: Update Release instructions [#21705](https://github.com/apache/datafusion/pull/21705) (comphead) +- test: add tests for spill file sizes to verify View GC [#21750](https://github.com/apache/datafusion/pull/21750) (RatulDawar) +- chore(deps): bump astral-sh/setup-uv from 8.0.0 to 8.1.0 [#21759](https://github.com/apache/datafusion/pull/21759) (dependabot[bot]) +- chore(deps): bump aws-config from 1.8.15 to 1.8.16 in the all-other-cargo-deps group [#21760](https://github.com/apache/datafusion/pull/21760) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.35.1 to 4.35.2 [#21758](https://github.com/apache/datafusion/pull/21758) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.75.10 to 2.75.18 [#21757](https://github.com/apache/datafusion/pull/21757) (dependabot[bot]) +- Snowflake Unparser dialect and UNNEST support [#21593](https://github.com/apache/datafusion/pull/21593) (yonatan-sevenai) +- Skip files outside partition structure in hive-partitioned listing tables [#21756](https://github.com/apache/datafusion/pull/21756) (zhuqi-lucas) +- Handle canceled partitioned hash join dynamic filters lazily [#21666](https://github.com/apache/datafusion/pull/21666) (adriangb) +- Improve ergonomics for ExecutionPlanMetricsSet and MetricsSet [#21762](https://github.com/apache/datafusion/pull/21762) (gabotechs) +- [Minor]: unify ANY/ALL planning and align ANY NULL semantics with PG [#21743](https://github.com/apache/datafusion/pull/21743) (buraksenn) +- [Minor]: fix security audit because of rustls-webpki version [#21785](https://github.com/apache/datafusion/pull/21785) (buraksenn) +- refactor: Simplify NLJ re-scans with `ReplayableStreamSource` [#21742](https://github.com/apache/datafusion/pull/21742) (2010YOUY01) +- ci: permit stale workflow to delete cache [#21772](https://github.com/apache/datafusion/pull/21772) (Jefffrey) +- Unparser drops ORDER BY alias when flattening Projection through SubqueryAlias [#21491](https://github.com/apache/datafusion/pull/21491) (yonatan-sevenai) +- chore: re-enable `add_months` overflow test [#21774](https://github.com/apache/datafusion/pull/21774) (Jefffrey) +- chore: add aggregation test for listview types [#21776](https://github.com/apache/datafusion/pull/21776) (Jefffrey) +- chore: re-enable `array_union` nested null array edge case test [#21773](https://github.com/apache/datafusion/pull/21773) (Jefffrey) +- Fix: allow coercion from Binary and LargeBinary into BinaryView [#21800](https://github.com/apache/datafusion/pull/21800) (bert-beyondloops) +- chore: leave specialised bench helpers [#21810](https://github.com/apache/datafusion/pull/21810) (theirix) +- Add quote style and trimming to csv writier [#20813](https://github.com/apache/datafusion/pull/20813) (xanderbailey) +- chore(deps): bump picomatch from 2.3.1 to 2.3.2 in /datafusion/wasmtest/datafusion-wasm-app [#21164](https://github.com/apache/datafusion/pull/21164) (dependabot[bot]) +- perf(substr_index): speed up scalar and Utf8View [#21754](https://github.com/apache/datafusion/pull/21754) (kumarUjjawal) +- Fix PushdownSort dropping LIMIT when eliminating SortExec [#21744](https://github.com/apache/datafusion/pull/21744) (sgrebnov) +- chore: use Arc::unwrap_or_clone in more places [#21823](https://github.com/apache/datafusion/pull/21823) (Dandandan) +- build: explicitly set `publish = false` for internal crates [#21869](https://github.com/apache/datafusion/pull/21869) (rluvaton) +- chore: bump API limit for stale workflow [#21867](https://github.com/apache/datafusion/pull/21867) (Jefffrey) +- chore: bump `sha` & `md-5` to `0.11.0` [#21840](https://github.com/apache/datafusion/pull/21840) (Jefffrey) +- feat : ABI upgrade from abi_stabby to stabby since abi_stable is no longer maintained [#21030](https://github.com/apache/datafusion/pull/21030) (coderfender) +- Add protobuf serialization/deserialization support for `EmptyTable` scans [#20844](https://github.com/apache/datafusion/pull/20844) (OlegWock) +- Support Dictionary Arrays in MIN/MAX Aggregates [#21315](https://github.com/apache/datafusion/pull/21315) (kosiew) +- Fix some GH action permission issues identified by CodeQL [#21838](https://github.com/apache/datafusion/pull/21838) (Jefffrey) +- Add support for nested types to nullif. [#21764](https://github.com/apache/datafusion/pull/21764) (tabac) +- chore(deps): bump taiki-e/install-action from 2.75.18 to 2.75.23 [#21887](https://github.com/apache/datafusion/pull/21887) (dependabot[bot]) +- chore(deps): bump libloading from 0.8.9 to 0.9.0 [#21890](https://github.com/apache/datafusion/pull/21890) (dependabot[bot]) +- refactor `array_remove` benchmarks & add nested benches [#21834](https://github.com/apache/datafusion/pull/21834) (Jefffrey) +- Update `astral-tokio-tar` to appease cargo_audit [#21902](https://github.com/apache/datafusion/pull/21902) (alamb) +- Remove unnecessary Mutex in SharedMemoryReservation [#21899](https://github.com/apache/datafusion/pull/21899) (gabotechs) +- ci: add breaking change detector [#21499](https://github.com/apache/datafusion/pull/21499) (rluvaton) +- Fix GH action permissions in `rust.yml` and `docs.yaml` workflows [#21884](https://github.com/apache/datafusion/pull/21884) (Jefffrey) +- chore: fix `iff` typos [#21904](https://github.com/apache/datafusion/pull/21904) (comphead) +- Deduplicate InList primitive static filters [#21932](https://github.com/apache/datafusion/pull/21932) (geoffreyclaude) +- Fix nesting of permissions block in docs workflow [#21930](https://github.com/apache/datafusion/pull/21930) (Jefffrey) +- dependencies check are now required to merge ci [#21940](https://github.com/apache/datafusion/pull/21940) (blaginin) +- build: allow posting comments on PRs made from forks and fix missing protobuf [#21913](https://github.com/apache/datafusion/pull/21913) (rluvaton) +- Use shared statistics merge for union stats [#21430](https://github.com/apache/datafusion/pull/21430) (kumarUjjawal) +- Add ClickBench URL pushdown benchmark [#21945](https://github.com/apache/datafusion/pull/21945) (xudong963) +- test(sqllogictest): stabilize parquet output_rows_skew with WITH ORDER [#21898](https://github.com/apache/datafusion/pull/21898) (RatulDawar) +- Skip unnecessary plan rebuild in adjust_input_keys_ordering for non-join plans [#21947](https://github.com/apache/datafusion/pull/21947) (zhuqi-lucas) +- Adding Use of arrow's has_true() / has_false() [#21806](https://github.com/apache/datafusion/pull/21806) (raushanprabhakar1) +- feat[expr-common]: support regex and LIKE coercion on REE and Dict value types that require an extra coercion step [#21924](https://github.com/apache/datafusion/pull/21924) (asubiotto) +- feat[expr-common]: support REE in coalesce [#21919](https://github.com/apache/datafusion/pull/21919) (asubiotto) +- proto: serialize and dedupe dynamic filters v2 [#21807](https://github.com/apache/datafusion/pull/21807) (jayshrivastava) +- chore: fix `datafusion-spark` substring [#21963](https://github.com/apache/datafusion/pull/21963) (comphead) +- Respect DATA_DIR location for sql benchmarks [#21961](https://github.com/apache/datafusion/pull/21961) (Omega359) +- ci: use base repository branch for breaking change detector [#22006](https://github.com/apache/datafusion/pull/22006) (rluvaton) +- bench: add to_char_array_date32 [#22007](https://github.com/apache/datafusion/pull/22007) (huymq1710) +- ci: add `auto detected api change` label on breaking change detecting in the CI [#21953](https://github.com/apache/datafusion/pull/21953) (rluvaton) +- Fix fully matched row groups with null counts [#21907](https://github.com/apache/datafusion/pull/21907) (xudong963) +- functions: Add dict support for get field [#21115](https://github.com/apache/datafusion/pull/21115) (brancz) +- fix(physical-plan): set column byte_size to 0 in FilterExec zero-row interval stats [#21999](https://github.com/apache/datafusion/pull/21999) (buraksenn) +- Explicitly declare spill codec dependency in `physical-plan` [#21917](https://github.com/apache/datafusion/pull/21917) (kosiew) +- Add benchmark_runner for sql_benchmarks with help and list commands [#22001](https://github.com/apache/datafusion/pull/22001) (Omega359) +- chore: `datafusion-spark` substring to support Binary types [#21979](https://github.com/apache/datafusion/pull/21979) (comphead) +- Add reusable plan-time schema alignment helper and apply to RecursiveQueryExec [#21912](https://github.com/apache/datafusion/pull/21912) (kosiew) +- Upgrade to arrow-rs / parquet / avro 58.2.0 [#21812](https://github.com/apache/datafusion/pull/21812) (alamb) +- kill `linux-build-lib` from extended tests [#21227](https://github.com/apache/datafusion/pull/21227) (blaginin) +- chore: Rust checks are required + merge queue [#21941](https://github.com/apache/datafusion/pull/21941) (blaginin) +- Add wide-schema benchmark suite for measuring per-file metadata overhead [#21970](https://github.com/apache/datafusion/pull/21970) (adriangb) +- chore(deps): bump ctor from 0.10.1 to 1.0.1 [#22023](https://github.com/apache/datafusion/pull/22023) (dependabot[bot]) +- ci: narrow macOS test scope to datafusion-ffi, run benchmarks on amd64 [#22048](https://github.com/apache/datafusion/pull/22048) (blaginin) +- chore(deps): bump github/codeql-action from 4.35.2 to 4.35.3 [#22019](https://github.com/apache/datafusion/pull/22019) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.74.0 to 2.77.0 [#22018](https://github.com/apache/datafusion/pull/22018) (dependabot[bot]) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 2 updates [#22022](https://github.com/apache/datafusion/pull/22022) (dependabot[bot]) +- Allow benchmark allocator features together [#21905](https://github.com/apache/datafusion/pull/21905) (xudong963) +- Rich t kid/introduce dict benchmarks [#21860](https://github.com/apache/datafusion/pull/21860) (Rich-T-kid) +- Add benchmarks for dictionary path of new_group_values [#22004](https://github.com/apache/datafusion/pull/22004) (Rich-T-kid) +- Support `IS (NOT) DISTINCT FROM` in Unparser [#22054](https://github.com/apache/datafusion/pull/22054) (cetra3) +- chore: Fix broken build with `--benches --all-features` [#22081](https://github.com/apache/datafusion/pull/22081) (neilconway) +- Chore: Fix TPC-DS schema/query (fixes q30 run) [#22086](https://github.com/apache/datafusion/pull/22086) (Dandandan) +- chore(deps): (fix CI) bump taiki-e/install-action from 2.77.0 to 2.77.6 [#22110](https://github.com/apache/datafusion/pull/22110) (gstvg) +- Prevent empty grouping sets from being eliminated on empty input [#22039](https://github.com/apache/datafusion/pull/22039) (xiedeyantu) +- Consolidate and document SQL AST shims [#22094](https://github.com/apache/datafusion/pull/22094) (alamb) +- Support distinct-from predicates in Parquet pruning [#22084](https://github.com/apache/datafusion/pull/22084) (Dandandan) +- minor: Track Parquet rows and pages matched when the page index is skipped [#22085](https://github.com/apache/datafusion/pull/22085) (nuno-faria) +- Update to `arrow` / `parquet` from 58.2.0 --> 58.3.0 [#22066](https://github.com/apache/datafusion/pull/22066) (alamb) +- Add sqllogictest coverage for unused UNNEST pruning edge cases [#22074](https://github.com/apache/datafusion/pull/22074) (kosiew) +- chore(deps): bump actions/labeler from 6.0.1 to 6.1.0 [#22124](https://github.com/apache/datafusion/pull/22124) (dependabot[bot]) +- chore(deps): bump the all-other-cargo-deps group with 5 updates [#22128](https://github.com/apache/datafusion/pull/22128) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.35.3 to 4.35.4 [#22122](https://github.com/apache/datafusion/pull/22122) (dependabot[bot]) +- mem: Cleanup resources of done streams immediately [#22064](https://github.com/apache/datafusion/pull/22064) (EmilyMatt) +- Propagate field metadata through NTH_VALUE, FIRST_VALUE, and LAST_VALUE window functions [#22112](https://github.com/apache/datafusion/pull/22112) (paleolimbot) +- Minor: Disallow async function in lambdas [#22097](https://github.com/apache/datafusion/pull/22097) (gstvg) +- chore(deps): bump runs-on/action from 2.1.0 to 2.1.2 [#22123](https://github.com/apache/datafusion/pull/22123) (dependabot[bot]) +- bench: remove stale `array_expression` benchmark [#22143](https://github.com/apache/datafusion/pull/22143) (kumarUjjawal) +- Add resolve_lambda_variables helper to Expr and LogicalPlan [#22101](https://github.com/apache/datafusion/pull/22101) (gstvg) +- Fix panic on deep compound identifiers [#22186](https://github.com/apache/datafusion/pull/22186) (Dandandan) +- Refactor scalar min/max dispatch into function-based helpers [#22062](https://github.com/apache/datafusion/pull/22062) (kosiew) +- fix missing window expressions when unparsing plans without outer projections [#21801](https://github.com/apache/datafusion/pull/21801) (nathanb9) +- chore(deps): bump urllib3 from 2.6.3 to 2.7.0 [#22109](https://github.com/apache/datafusion/pull/22109) (dependabot[bot]) +- Call take arrays once per repartitioned input batch [#22159](https://github.com/apache/datafusion/pull/22159) (gene-bordegaray) +- Refactor parquet row filter setup [#22191](https://github.com/apache/datafusion/pull/22191) (xudong963) +- fix date_bin overflows subtracting extreme nanosecond timestamp origin [#22251](https://github.com/apache/datafusion/pull/22251) (xiedeyantu) +- fix date_trunc overflows converting extreme non-ns timestamps to nanoseconds [#22262](https://github.com/apache/datafusion/pull/22262) (xiedeyantu) +- Extract parquet push decoder module [#22289](https://github.com/apache/datafusion/pull/22289) (xudong963) +- Track spill read-back memory in SMJ [#22103](https://github.com/apache/datafusion/pull/22103) (SubhamSinghal) +- refactor(parquet-datasource): split opener.rs into an opener/ module [#22346](https://github.com/apache/datafusion/pull/22346) (adriangb) +- refactor(parquet-datasource): split sink and schema_coercion out of file_format.rs [#22347](https://github.com/apache/datafusion/pull/22347) (adriangb) +- fixing negative power to zero [#22277](https://github.com/apache/datafusion/pull/22277) (raushanprabhakar1) +- refactor(parquet-datasource): split bloom_filter out of row_group_filter.rs [#22348](https://github.com/apache/datafusion/pull/22348) (adriangb) +- Revert "[Minor]: unify ANY/ALL planning and align ANY NULL semantics with PG (#21743)" [#22345](https://github.com/apache/datafusion/pull/22345) (alamb) +- Fix pruning predicate for `LIKE` expressions with escape sequences [#22375](https://github.com/apache/datafusion/pull/22375) (masonh22) +- Fix: lead/lag extreme offsets handling [#22243](https://github.com/apache/datafusion/pull/22243) (Dandandan) +- chore(deps): fix CI, bump astral-tokio-tar [#22382](https://github.com/apache/datafusion/pull/22382) (gstvg) +- chore: Replace stray old-style string builder in `substr` [#22183](https://github.com/apache/datafusion/pull/22183) (neilconway) +- chore(deps): bump taiki-e/install-action from 2.77.6 to 2.79.2 [#22377](https://github.com/apache/datafusion/pull/22377) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.35.4 to 4.35.5 [#22376](https://github.com/apache/datafusion/pull/22376) (dependabot[bot]) +- chore(deps-dev): bump webpack-dev-server from 5.2.1 to 5.2.4 in /datafusion/wasmtest/datafusion-wasm-app [#22349](https://github.com/apache/datafusion/pull/22349) (dependabot[bot]) +- chore(deps): bump idna from 3.11 to 3.15 [#22381](https://github.com/apache/datafusion/pull/22381) (dependabot[bot]) +- Refactor Spark `format_string` numeric `%c` conversion dispatch [#22166](https://github.com/apache/datafusion/pull/22166) (kosiew) +- chore(deps): bump sysinfo from 0.38.4 to 0.39.2 [#22380](https://github.com/apache/datafusion/pull/22380) (dependabot[bot]) +- fix regexp_count should count empty-pattern matches [#22311](https://github.com/apache/datafusion/pull/22311) (xiedeyantu) +- Actually preserve predicate execution order in PushDownFilter [#21643](https://github.com/apache/datafusion/pull/21643) (joroKr21) +- chore(deps): bump qs and body-parser in /datafusion/wasmtest/datafusion-wasm-app [#22321](https://github.com/apache/datafusion/pull/22321) (dependabot[bot]) +- [branch-54] add changelog [#22402](https://github.com/apache/datafusion/pull/22402) (mbutrovich) +- [branch-54]: Backport 22404. Fix Spark slice function on negative OOB [#22443](https://github.com/apache/datafusion/pull/22443) (comphead) +- [branch-54] Cherry-pick #22493: restore SortExec elimination after stats-based file reorder [#22501](https://github.com/apache/datafusion/pull/22501) (zhuqi-lucas) +- [branch-54] Fix: compact view buffers in ScalarValue::compact for all container types (#21934) [#22446](https://github.com/apache/datafusion/pull/22446) (alamb) +- [branch-54] Support transparent ExecutionPlan downcasts [#22565](https://github.com/apache/datafusion/pull/22565) (geoffreyclaude) +- [branch-54] Fix TopK DISTINCT aggregation preserving NULLs (#22571) [#22634](https://github.com/apache/datafusion/pull/22634) (alamb) +- [branch-54] refactor: wrap HigherOrderUDFImpl in a concrete HigherOrderUDF struct (#22593) [#22635](https://github.com/apache/datafusion/pull/22635) (alamb) +- [branch-54] chore: Cleanup and refactor `build_join` in `ScalarSubqueryToJoin` (#… [#22693](https://github.com/apache/datafusion/pull/22693) (LiaCastaneda) +- [branch-54] fix: clear handled OFFSET before child recursion in LimitPushdown (#22525) [#22631](https://github.com/apache/datafusion/pull/22631) (alamb) +- [branch-54] refactor: give parquet CDC options an explicit `enabled` flag (backport #22632) [#22648](https://github.com/apache/datafusion/pull/22648) (kszucs) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 70 Neil Conway + 68 dependabot[bot] + 56 Andrew Lamb + 26 Burak Şen + 26 Oleks V + 21 Daniël Heres + 20 Adrian Garcia Badaracco + 18 Kumar Ujjawal + 17 Matt Butrovich + 16 kosiew + 15 Qi Zhu + 15 Tim Saucer + 14 Zhen Chen + 13 Dmitrii Blaginin + 13 Jeffrey Vo + 13 Yongting You + 13 xudong.w + 12 Huaijin + 11 Bhargava Vadlamani + 10 Eren Avsarogullari + 10 Matthew Kim + 9 gstvg + 8 Raz Luvaton + 8 Subham Singhal + 8 theirix + 6 Adam Gutglick + 6 Gabriel + 6 Jonathan Chen + 5 Alessandro Solimando + 5 Alfonso Subiotto Marqués + 5 Bruce Ritchie + 5 Liang-Chi Hsieh + 5 Lía Adriana + 5 Nuno Faria + 5 Sergei Grebnov + 4 Andy Grove + 4 Ariel Miculas-Trif + 4 Geoffrey Claude + 4 Jayant Shrivastava + 4 Kevin Liu + 4 lyne + 3 Brent Gardner + 3 David López + 3 Dewey Dunnington + 3 Harrison Crosse + 3 Huy Mac + 3 Kazantsev Maksim + 3 Konstantin Tarasov + 3 Namgung Chan + 3 Peter L + 3 RIchard Baah + 3 Ratul Dawar + 3 Raushan Prabhakar + 3 Xander + 3 Yonatan Striem Amit + 3 Yu-Chuan Hung + 3 crm26 + 2 Acfboy + 2 Adam Curtis + 2 Albert Skalt + 2 Anastasios Bakogiannis + 2 Bert Vermeiren + 2 Frederic Branczyk + 2 Geethapranay1 + 2 Jonah Gao + 2 Krisztián Szűcs + 2 Liam Feehery + 2 Marko Grujic + 2 Michael Kleen + 2 Peter Nguyen + 2 Rafael Herrero + 2 Rohan Krishnaswamy + 2 Samyak Sarnayak + 2 Sergey Zhukov + 2 Shiv Bhatia + 2 Tobias Schwarzinger + 2 hsiang-c + 2 jj.lee + 2 linfeng + 2 yaommen + 1 Adam Reeve + 1 Ahmed Mezghani + 1 Alex Zhang + 1 Alexander Alexandrov + 1 Alexander Rafferty + 1 Alexandre Crayssac + 1 Andrey Koshchiy + 1 Asish Kumar + 1 Ben Bellick + 1 Bruno Volpato + 1 Daniel Tu + 1 Druva + 1 EeshanBembi + 1 Emil Ernerfeldt + 1 Emily Matheys + 1 Ernest Provo + 1 Filip Petkovski + 1 Filip Wojciechowski + 1 Florian Müller + 1 Fred Thomas + 1 Gene Bordegaray + 1 Georgi Krastev + 1 Guillaume Boucher + 1 Haresh Khanna + 1 Helgi Kristvin Sigurbjarnarson + 1 Heran Lin + 1 Jamal Saad + 1 Jarro van Ginkel + 1 Jax Liu + 1 Joan Antoni RE + 1 Justin O'Dwyer + 1 Kartik Gupta + 1 Krishna Sudarshan J + 1 Kristin Cowalcijk + 1 Lavkesh Lahngir + 1 Martin Hilton + 1 Mason + 1 Nathan + 1 Oleh + 1 Ramakrishna Chilaka + 1 Rizky Mirzaviandy Priambodo + 1 RyanStewart + 1 Shivaang + 1 Soham Bhattacharjee + 1 Stu Hood + 1 Thomas Tanon + 1 Tim-53 + 1 UBarney + 1 Viktor Yershov + 1 Vinay Mehta + 1 Zhang Xiaofeng + 1 aditya singh rathore + 1 alexanderbianchi + 1 blaginin + 1 dario curreri + 1 dd-david-levin + 1 gabriel + 1 nathan + 1 niebayes +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. diff --git a/dev/changelog/54.1.0.md b/dev/changelog/54.1.0.md new file mode 100644 index 0000000000000..b45f42c9b1ece --- /dev/null +++ b/dev/changelog/54.1.0.md @@ -0,0 +1,67 @@ + + +# Apache DataFusion 54.1.0 Changelog + +This release consists of 19 commits from 9 contributors. See credits at the end of this changelog for more information. + +See the [upgrade guide](https://datafusion.apache.org/library-user-guide/upgrading.html) for information on how to upgrade from previous versions. + +**Documentation updates:** + +- [branch-54] Add datafusion.execution.enable_file_stream_work_stealing config [#23296](https://github.com/apache/datafusion/pull/23296) (andygrove) + +**Other:** + +- [branch-54] fix: preserve null_aware on logical JoinNode proto round-trip (backport #22104) [#22785](https://github.com/apache/datafusion/pull/22785) (mithuncy) +- [branch-54]: backport #22811 (bugfix: changed return type of spark's width_bucket to i64) [#23087](https://github.com/apache/datafusion/pull/23087) (mbutrovich) +- [branch-54] backport #22857 (Skip loading Parquet page index when row-group statistics already prove it cannot prune) [#23088](https://github.com/apache/datafusion/pull/23088) (mbutrovich) +- [branch-54] backport #23192 `array_compact` handle edge case with NULLs [#23196](https://github.com/apache/datafusion/pull/23196) (comphead) +- [branch-54] fix: Avoid panicing when stats are not available for a file group split (backport #23277) [#23340](https://github.com/apache/datafusion/pull/23340) (mkleen) +- [branch-54] fix: `approx_distinct` over-counts for utf8view (backport #22815, adapted) [#23576](https://github.com/apache/datafusion/pull/23576) (mbutrovich) +- [branch-54] fix: isolate anonymous file statistics cache (backport #22950, adapted) [#23573](https://github.com/apache/datafusion/pull/23573) (mbutrovich) +- [branch-54] fix: `= ANY (SELECT ...)` / `<> ALL (SELECT ...)` schema error (backport #22915) [#23575](https://github.com/apache/datafusion/pull/23575) (mbutrovich) +- [branch-54] fix: NestedLoopJoinExec emits spurious unmatched-left rows with multiple probe partitions (backport #22791) [#23577](https://github.com/apache/datafusion/pull/23577) (mbutrovich) +- [branch-54] fix: regex simplification of anchored patterns produces wrong results (backport #22727) [#23578](https://github.com/apache/datafusion/pull/23578) (mbutrovich) +- [branch-54] fix: Correctly compute nullability in recursive CTE schemas (backport #22552) [#23579](https://github.com/apache/datafusion/pull/23579) (mbutrovich) +- [branch-54] fix: handle `IS TRUE` correctly in `EliminateOuterJoin` (backport #22444) [#23580](https://github.com/apache/datafusion/pull/23580) (mbutrovich) +- [branch-54] fix: preserve no-filter SMJ matches across pending outer batches (backport #23049) [#23574](https://github.com/apache/datafusion/pull/23574) (mbutrovich) +- [branch-54] perf: avoid intermediate slice allocation in Spark slice function (backport #23481) [#23582](https://github.com/apache/datafusion/pull/23582) (mbutrovich) +- [branch-54] fix: don't duplicate volatile expressions when pushing projection into file scan (backport #23395, adapted) [#23585](https://github.com/apache/datafusion/pull/23585) (fordN) +- [branch-54] chore: fix cargo audit [#23607](https://github.com/apache/datafusion/pull/23607) (alamb) +- [branch-54] Resolve lost wakeup in SpillPoolReader with multiple conc… [#23654](https://github.com/apache/datafusion/pull/23654) (pepijnve) +- [branch-54] Handle nulls in type coercion of higher-order UDFs, map_extract, spark array_repeat (backport #23071) [#23629](https://github.com/apache/datafusion/pull/23629) (gstvg) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 11 Matt Butrovich + 1 Andrew Lamb + 1 Andy Grove + 1 Ford + 1 Michael Kleen + 1 Mithun Chicklore Yogendra + 1 Oleks V + 1 Pepijn Van Eeckhoudt + 1 gstvg +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. diff --git a/dev/changelog/55.0.0.md b/dev/changelog/55.0.0.md new file mode 100644 index 0000000000000..30ee1d369878f --- /dev/null +++ b/dev/changelog/55.0.0.md @@ -0,0 +1,1099 @@ + + +# Apache DataFusion 55.0.0 Changelog + +This release consists of 877 commits from 175 contributors. See credits at the end of this changelog for more information. + +See the [upgrade guide](https://datafusion.apache.org/library-user-guide/upgrading.html) for information on how to upgrade from previous versions. + +**Breaking changes:** + +- fix: preserve null_aware on logical JoinNode proto round-trip [#22104](https://github.com/apache/datafusion/pull/22104) (mithuncy) +- PushdownFilter optimizations [#21668](https://github.com/apache/datafusion/pull/21668) (joroKr21) +- proto: add proto converter reference to PhysicalExtensionCodec trait [#21055](https://github.com/apache/datafusion/pull/21055) (jayshrivastava) +- fix ^ evaluates as bitwise XOR instead of exponentiation [#22314](https://github.com/apache/datafusion/pull/22314) (xiedeyantu) +- Add EnsureRequirements: merged EnforceDistribution + EnforceSorting with idempotent pushdown_sorts [#21976](https://github.com/apache/datafusion/pull/21976) (zhuqi-lucas) +- feat(physical-expr): DynamicFilterTracker for cheap dynamic-filter change detection [#22460](https://github.com/apache/datafusion/pull/22460) (adriangb) +- Add minimal APIs / hooks for granular statistics collection in TableProvider implementations [#22300](https://github.com/apache/datafusion/pull/22300) (adriangb) +- Add lambda substrait support [#21193](https://github.com/apache/datafusion/pull/21193) (gstvg) +- minor: add `Any` to `QueryPlanner` trait [#22241](https://github.com/apache/datafusion/pull/22241) (milenkovicm) +- Add Physical `Partitioning::Range` enum variant [#22207](https://github.com/apache/datafusion/pull/22207) (gene-bordegaray) +- refactor: cache schema_without_virtual_columns and remove TableSchema::with_virtual_columns [#22600](https://github.com/apache/datafusion/pull/22600) (mbutrovich) +- feat(sql): Postgres-style `EXPLAIN (...)` option list [#21768](https://github.com/apache/datafusion/pull/21768) (adriangb) +- refactor: wrap HigherOrderUDFImpl in a concrete HigherOrderUDF struct [#22593](https://github.com/apache/datafusion/pull/22593) (LiaCastaneda) +- feat: add pgjson format support for EXPLAIN ANALYZE [#21767](https://github.com/apache/datafusion/pull/21767) (adriangb) +- Gate new ScalarSubqueryExec node behind session property [#22530](https://github.com/apache/datafusion/pull/22530) (LiaCastaneda) +- fix: Correctly compute nullability in recursive CTE schemas [#22552](https://github.com/apache/datafusion/pull/22552) (neilconway) +- Allow specifying an arrow schema for PartitionedFile [#22360](https://github.com/apache/datafusion/pull/22360) (fpetkovski) +- refactor: give parquet CDC options an explicit `enabled` flag [#22632](https://github.com/apache/datafusion/pull/22632) (kszucs) +- Add optimize_with_context to FFI_PhysicalOptimizerRule [#22584](https://github.com/apache/datafusion/pull/22584) (nathanb9) +- perf(logical-plan): box CreateExternalTable / CreateFunction in DdlStatement (-45% LogicalPlan size) [#22733](https://github.com/apache/datafusion/pull/22733) (zhuqi-lucas) +- feat: add max_row_group_bytes option to ParquetOptions [#22649](https://github.com/apache/datafusion/pull/22649) (Satyr09) +- feat: Add Spark SQL parser dialect config [#22529](https://github.com/apache/datafusion/pull/22529) (kumarUjjawal) +- refactor: Split hash aggregation logic into separated streams [#22729](https://github.com/apache/datafusion/pull/22729) (2010YOUY01) +- Add logical range partitioning representation [#22777](https://github.com/apache/datafusion/pull/22777) (gene-bordegaray) +- refactor: centralize SQL dialect metadata [#22840](https://github.com/apache/datafusion/pull/22840) (kumarUjjawal) +- Revert custom allocator auditing of MemoryPool tracking in SLTs [#22860](https://github.com/apache/datafusion/pull/22860) (avantgardnerio) +- fix: Correct output-count stats for partitioned partial aggs [#22780](https://github.com/apache/datafusion/pull/22780) (neilconway) +- fix: preserve async UDF return field metadata [#22663](https://github.com/apache/datafusion/pull/22663) (Kontinuation) +- fix: preserve Spark next_day whitespace validation [#22720](https://github.com/apache/datafusion/pull/22720) (xfocus3) +- FFI: plumb `placement` for `FFI_ScalarUDF` [#22608](https://github.com/apache/datafusion/pull/22608) (Amogh-2404) +- refactor: remove `opt_filter` in `GroupsAccumulator::merge_batch` [#22816](https://github.com/apache/datafusion/pull/22816) (haohuaijin) +- feat: decimal support for gcd and lcm [#22655](https://github.com/apache/datafusion/pull/22655) (theirix) +- refactor: Update SortMergeJoin to use async spill abstractions [#22230](https://github.com/apache/datafusion/pull/22230) (pantShrey) +- Add MERGE INTO types to datafusion-expr [#20763](https://github.com/apache/datafusion/pull/20763) (wirybeaver) +- Remove redundant `collect_stat` and `target_partitions` on `ListingOptions` [#22969](https://github.com/apache/datafusion/pull/22969) (gabotechs) +- fix: Omit NULL values from build side of hash joins [#22893](https://github.com/apache/datafusion/pull/22893) (neilconway) +- refactor: Simplify `approx_distinct` (-200 LoC) [#22921](https://github.com/apache/datafusion/pull/22921) (2010YOUY01) +- Introduce generic memory-limiting cache for parquet metadata [#22613](https://github.com/apache/datafusion/pull/22613) (mkleen) +- Add StatisticsContext parameter to partition_statistics [#21815](https://github.com/apache/datafusion/pull/21815) (asolimando) +- feat(parquet): intra-file early stopping via statistics + dynamic filters [#22450](https://github.com/apache/datafusion/pull/22450) (zhuqi-lucas) +- feat: logical plan protobuf representation for range repartitioning [#23030](https://github.com/apache/datafusion/pull/23030) (saadtajwar) +- perf: optimize object store requests when reading CSV [#22962](https://github.com/apache/datafusion/pull/22962) (saadtajwar) +- Group scan time expression rewrite functionality for UDFs in new module in `datafusion-physical-expr-adapter` [#23125](https://github.com/apache/datafusion/pull/23125) (AdamGS) +- [physical-plan]: remove deprecated UnionExec::new [#23100](https://github.com/apache/datafusion/pull/23100) (mgkz0) +- chore(datasource): remove deprecated `create_writer` free function (Closes #23080 — partial) [#23129](https://github.com/apache/datafusion/pull/23129) (Dodothereal) +- chore(catalog): remove deprecated ViewTable try_new (Closes #23080 - partial) [#23131](https://github.com/apache/datafusion/pull/23131) (Dodothereal) +- Add `ListingOptions::output_partitioning` and `FileScanConfig::output_partitioning` for pre-defined file partitioning [#22657](https://github.com/apache/datafusion/pull/22657) (gene-bordegaray) +- chore(parquet): remove deprecated schema-coercion helpers (Closes #23080 - partial) [#23132](https://github.com/apache/datafusion/pull/23132) (Dodothereal) +- chore(expr): remove deprecated Filter::try_new_with_having (Closes #23080 - partial) [#23150](https://github.com/apache/datafusion/pull/23150) (Dodothereal) +- chore(common): remove deprecated DFSchema::check_arrow_schema_type_compatible (Closes #23080 - partial) [#23151](https://github.com/apache/datafusion/pull/23151) (Dodothereal) +- chore(catalog-listing): remove deprecated split_files free fn (Closes #23080) [#23152](https://github.com/apache/datafusion/pull/23152) (Dodothereal) +- chore(sql): remove deprecated DFParser constructors (Closes #23080 - partial) [#23142](https://github.com/apache/datafusion/pull/23142) (Dodothereal) +- chore(common): remove deprecated DFSchema type-check method (Closes #23080 - partial) [#23144](https://github.com/apache/datafusion/pull/23144) (Dodothereal) +- chore(expr): remove deprecated Filter::try_new_with_having (Closes #23080 - partial) [#23145](https://github.com/apache/datafusion/pull/23145) (Dodothereal) +- chore(expr-common): remove deprecated Signature::get_possible_types (Closes #23080 - partial) [#23147](https://github.com/apache/datafusion/pull/23147) (Dodothereal) +- [sql]: remove old deprecated `DFParser::new` and `DFParser::new_with_dialect` [#23101](https://github.com/apache/datafusion/pull/23101) (mgkz0) +- chore(common): remove deprecated equivalent_names_and_types (Closes #23080) [#23153](https://github.com/apache/datafusion/pull/23153) (Dodothereal) +- chore(expr-common): remove deprecated Signature get_possible_types (Closes #23080 - partial) [#23135](https://github.com/apache/datafusion/pull/23135) (Dodothereal) +- [execution] Remove deprecated disk manager configuration API [#23139](https://github.com/apache/datafusion/pull/23139) (mgkz0) +- [physical-plan]: remove deprecated spill_record_batch_by_size [#23029](https://github.com/apache/datafusion/pull/23029) (alamb) +- fix: Fix peak memory display in `EXPLAIN ANALYZE` for multiple operators [#23140](https://github.com/apache/datafusion/pull/23140) (2010YOUY01) +- chore(datasource): remove deprecated add_row_stats (Closes #23080 - partial) [#23134](https://github.com/apache/datafusion/pull/23134) (Dodothereal) +- feat: introduce pluggable SpillFile trait and TempFileFactory for custom spill backends [#21882](https://github.com/apache/datafusion/pull/21882) (pantShrey) +- Add `Distribution::HashPartitioned` to `Distribution::KeyPartitioned` API bridge [#23259](https://github.com/apache/datafusion/pull/23259) (gene-bordegaray) +- feat: add datafusion.execution.enable_file_stream_work_stealing config [#23294](https://github.com/apache/datafusion/pull/23294) (andygrove) +- refactor: make file-statistics cache keys schema-aware [#23201](https://github.com/apache/datafusion/pull/23201) (Phoenix500526) +- perf: preserve dictionary encoding for lower/upper to avoid materializing low-cardinality columns [#22905](https://github.com/apache/datafusion/pull/22905) (lyne7-sc) +- Remove unstable public methods for `DynamicFilterPhysicalExpr` after proto migration [#23423](https://github.com/apache/datafusion/pull/23423) (jayshrivastava) +- refactor: remove redundant partitioned_by_file_group file scan field [#23189](https://github.com/apache/datafusion/pull/23189) (Phoenix500526) +- Add protobuf support for lambdas [#22362](https://github.com/apache/datafusion/pull/22362) (gstvg) +- Support co-partitioned range inner equi joins [#23184](https://github.com/apache/datafusion/pull/23184) (gene-bordegaray) +- refactor: Migrate ScalarSubqueryExpr to self-serialization proto pattern [#23130](https://github.com/apache/datafusion/pull/23130) (mattp5657) +- refactor(physical-plan): externalize statistics traversal into StatisticsContext [#23051](https://github.com/apache/datafusion/pull/23051) (asolimando) +- perf: Extend WindowTopN to support RANK [#22885](https://github.com/apache/datafusion/pull/22885) (SubhamSinghal) +- refactor: make join projection pushdown schema-aware via ColumnIndex/… [#23185](https://github.com/apache/datafusion/pull/23185) (Phoenix500526) +- ci: reintroduce code coverage reporting with cargo-llvm-cov [#23336](https://github.com/apache/datafusion/pull/23336) (buraksenn) +- fix: align dictionary coercion across typed signatures [#23549](https://github.com/apache/datafusion/pull/23549) (lyne7-sc) +- fix: preserve EmptyExec and PlaceholderRowExec partition count across proto round-trip [#23643](https://github.com/apache/datafusion/pull/23643) (andygrove) +- Resolve lost wakeup in SpillPoolReader with multiple concurrent SpillPoolWriters [#23522](https://github.com/apache/datafusion/pull/23522) (pepijnve) +- chore: deprecate record_batch macro in favor of upstream one [#23295](https://github.com/apache/datafusion/pull/23295) (buraksenn) +- fix: `time ± interval` returns a wrapped `time` instead of an interval [#23279](https://github.com/apache/datafusion/pull/23279) (vismaytiwari) +- Add ExecutionPlan try_to_proto / try_from_proto hooks + ProjectionExec reference [#23495](https://github.com/apache/datafusion/pull/23495) (adriangb) +- feat: Support multiple external table locations [#22695](https://github.com/apache/datafusion/pull/22695) (kumarUjjawal) +- refactor: pass `PhysicalPlanningContext` explicitly through planner traits [#23649](https://github.com/apache/datafusion/pull/23649) (timsaucer) +- feat(proto): thread expr encode/decode context into try_encode_expr / try_decode_expr [#23733](https://github.com/apache/datafusion/pull/23733) (adriangb) +- refactor(proto): migrate FilterExec serde [#23708](https://github.com/apache/datafusion/pull/23708) (Phoenix500526) +- refactor(proto): migrate single-child plans [#23710](https://github.com/apache/datafusion/pull/23710) (Phoenix500526) +- refactor(proto): migrate sort merge join serde [#23712](https://github.com/apache/datafusion/pull/23712) (Phoenix500526) +- feat: Range Partitioning FFI [#23520](https://github.com/apache/datafusion/pull/23520) (saadtajwar) +- Bump MSRV from `1.88.0` to `1.94.0` [#23632](https://github.com/apache/datafusion/pull/23632) (Jefffrey) +- refactor: move catalog traits to session crate [#23703](https://github.com/apache/datafusion/pull/23703) (timsaucer) +- refactor(proto): migrate SortExec and SortPreservingMergeExec serde [#23794](https://github.com/apache/datafusion/pull/23794) (buraksenn) +- refactor(proto): migrate UnnestExec serde [#23739](https://github.com/apache/datafusion/pull/23739) (Phoenix500526) +- refactor(proto): migrate GlobalLimitExec and LocalLimitExec serde [#23791](https://github.com/apache/datafusion/pull/23791) (buraksenn) +- refactor(proto): migrate RepartitionExec serde [#23792](https://github.com/apache/datafusion/pull/23792) (buraksenn) +- refactor(proto): migrate CrossJoinExec and NestedLoopJoinExec serde [#23834](https://github.com/apache/datafusion/pull/23834) (buraksenn) +- refactor(proto): migrate UnionExec and InterleaveExec serde [#23782](https://github.com/apache/datafusion/pull/23782) (buraksenn) +- refactor(proto): migrate symmetric hash join serde [#23736](https://github.com/apache/datafusion/pull/23736) (Phoenix500526) +- feat: migrate EmptyExec and PlaceholderRowExec to ExecutionPlan proto hooks [#23784](https://github.com/apache/datafusion/pull/23784) (847850277) +- Remove `GroupsAccumulator::supports_convert_to_state` and require `convert_to_state` [#23489](https://github.com/apache/datafusion/pull/23489) (lyne7-sc) +- chore: Enable `unused_async` lint, make some functions sync [#23679](https://github.com/apache/datafusion/pull/23679) (neilconway) +- refactor(proto): migrate HashJoinExec serde [#23853](https://github.com/apache/datafusion/pull/23853) (buraksenn) +- refactor(proto): migrate AsyncFuncExec to self-serializing proto [#23825](https://github.com/apache/datafusion/pull/23825) (mattp5657) +- refactor(proto): migrate window serde [#23780](https://github.com/apache/datafusion/pull/23780) (Phoenix500526) +- Migrate ExplainExec and AnalyzeExec protobuf serde [#23742](https://github.com/apache/datafusion/pull/23742) (Phoenix500526) +- refactor(proto): migrate aggregate exec serde [#23779](https://github.com/apache/datafusion/pull/23779) (Phoenix500526) +- refactor(proto): remove legacy scan field [#23445](https://github.com/apache/datafusion/pull/23445) (Phoenix500526) +- `ScalarUdfImpl::strictly_order_preserving`: Allow expression to report whether they keep the same ordering of the input [#23807](https://github.com/apache/datafusion/pull/23807) (rluvaton) +- FFI: forward ScalarUDF preserves_lex_ordering [#23069](https://github.com/apache/datafusion/pull/23069) (Amogh-2404) +- perf(functions-aggregate): optimize sliding window MIN/MAX using monotonic deques (#23826) [#23827](https://github.com/apache/datafusion/pull/23827) (pavan51) +- chore(deps): bump syn from 2.0.119 to 3.0.2 [#23945](https://github.com/apache/datafusion/pull/23945) (dependabot[bot]) +- refactor(proto): migrate scalar subquery serde [#23915](https://github.com/apache/datafusion/pull/23915) (Phoenix500526) +- refactor: mark the ExecutionPlan proto dispatch traits as non-public API [#24001](https://github.com/apache/datafusion/pull/24001) (adriangb) +- feat: add GroupColumn support for Duration in multi-column GROUP BY [#23783](https://github.com/apache/datafusion/pull/23783) (tohuya6) +- refactor: move planning APIs to session crate [#23842](https://github.com/apache/datafusion/pull/23842) (timsaucer) +- feat: Add support for `unnest_outer` function for arrays. [#22100](https://github.com/apache/datafusion/pull/22100) (athlcode) +- perf: track `BoundedWindowAggExec` Linear-mode watermark once per stream [#24033](https://github.com/apache/datafusion/pull/24033) (neilconway) +- fix(ffi): preserve aggregate null-handling support [#23908](https://github.com/apache/datafusion/pull/23908) (Amogh-2404) +- refactor: unify `ParquetFileReader` and `CachedParquetFileReader` [#24036](https://github.com/apache/datafusion/pull/24036) (alamb) +- feat(pruning): expose pruning predicate IN-list rewrite size cap as a config option [#24074](https://github.com/apache/datafusion/pull/24074) (zhuqi-lucas) +- refactor join-key equality filtering [#23843](https://github.com/apache/datafusion/pull/23843) (shehab-ali) +- Proto: migrate file sink serialization [#23781](https://github.com/apache/datafusion/pull/23781) (Phoenix500526) +- refactor(pruning): deprecate PruningPredicate::try_new [#24129](https://github.com/apache/datafusion/pull/24129) (goutamadwant) +- refactor: move lambda variable scope into Physical Planning Context [#23989](https://github.com/apache/datafusion/pull/23989) (sweb) +- feat: Implement FFI_QueryPlanner [#24028](https://github.com/apache/datafusion/pull/24028) (timsaucer) +- add ExecutionPlan::dynamic_expressions_produced() method [#24068](https://github.com/apache/datafusion/pull/24068) (jayshrivastava) +- fix(proto): preserve HashJoinExec fetch across serialization [#24165](https://github.com/apache/datafusion/pull/24165) (adriangb) +- refactor(proto): migrate CsvSource serde [#24177](https://github.com/apache/datafusion/pull/24177) (buraksenn) +- refactor(proto): migrate ParquetSource serde [#24169](https://github.com/apache/datafusion/pull/24169) (buraksenn) +- refactor(proto): migrate JsonSource serde [#24178](https://github.com/apache/datafusion/pull/24178) (buraksenn) +- Proto: migrate MemorySourceConfig to per-source try_to_proto / try_from_proto hooks [#24187](https://github.com/apache/datafusion/pull/24187) (adriangb) +- refactor(proto): migrate AvroSource serde [#24190](https://github.com/apache/datafusion/pull/24190) (buraksenn) +- refactor(proto): migrate ArrowSource serde [#24189](https://github.com/apache/datafusion/pull/24189) (buraksenn) +- chore(proto): deprecate `AsyncFuncExec::async_exprs`, which only existed for proto serialization [#24168](https://github.com/apache/datafusion/pull/24168) (adriangb) +- Reapply "Add ExecutionPlan::apply_expressions() (apache#20337)" (apache#22437) [#24018](https://github.com/apache/datafusion/pull/24018) (jayshrivastava) +- fix(proto): serialize Global/LocalLimitExec required_ordering [#24183](https://github.com/apache/datafusion/pull/24183) (buraksenn) +- fix(proto): preserve AggregateExec schema and reversed state [#24207](https://github.com/apache/datafusion/pull/24207) (buraksenn) +- perf: remove per-row String allocations from the Spark url functions [#23884](https://github.com/apache/datafusion/pull/23884) (andygrove) +- Expose accumulator state to allow prefix scanning [#24035](https://github.com/apache/datafusion/pull/24035) (avantgardnerio) +- fix(lambda): only push referenced params into the merged batch [#24162](https://github.com/apache/datafusion/pull/24162) (LiaCastaneda) +- Enable dynamic filters for range-partitioned joins [#23854](https://github.com/apache/datafusion/pull/23854) (peterxcli) +- chore(proto): remove never-released deprecated PhysicalPlanNodeExt scaffolding [#24269](https://github.com/apache/datafusion/pull/24269) (adriangb) +- Restore the From / TryFrom proto conversions dropped since 54.1.0 [#24205](https://github.com/apache/datafusion/pull/24205) (adriangb) +- FFI: plumb with_updated_config for FFI_ScalarUDF [#22797](https://github.com/apache/datafusion/pull/22797) (Amogh-2404) +- fix(physical-plan): CTAS panic on wasm32-unknown-unknown [#24275](https://github.com/apache/datafusion/pull/24275) (kentkwu) +- fix: ensure new_list respects data_type argument [#24029](https://github.com/apache/datafusion/pull/24029) (Ruchirtripathi) + +**Performance related:** + +- Optimize logical optimizer: skip map_subqueries + in-place rewriting [#22298](https://github.com/apache/datafusion/pull/22298) (adriangb) +- perf: collapse chained projections in a single optimizer pass; reduce memory usage / recursion [#22389](https://github.com/apache/datafusion/pull/22389) (Dandandan) +- Fix: compact view buffers in ScalarValue::compact for all container t… [#21934](https://github.com/apache/datafusion/pull/21934) (bert-beyondloops) +- perf: Optimize `translate` to use new bulk-NULL string builders [#22171](https://github.com/apache/datafusion/pull/22171) (neilconway) +- perf: Optimize `overlay` with new string builder [#22182](https://github.com/apache/datafusion/pull/22182) (neilconway) +- Optimize metric label cloning [#22406](https://github.com/apache/datafusion/pull/22406) (xudong963) +- perf: optimize `array_replace` for scalar needle [#22387](https://github.com/apache/datafusion/pull/22387) (lyne7-sc) +- perf: optimize array_remove for scalar needle [#22390](https://github.com/apache/datafusion/pull/22390) (lyne7-sc) +- perf: Optimize `split_part` using bulk-NULL string builders [#22283](https://github.com/apache/datafusion/pull/22283) (neilconway) +- perf: hoist split_vec_min_alloc to datafusion-common and shrink the emitted prefix [#22416](https://github.com/apache/datafusion/pull/22416) (RyanJamesStewart) +- perf(physical-optimizer): skip ensure_distribution rebuild when children are unchanged [#22521](https://github.com/apache/datafusion/pull/22521) (zhuqi-lucas) +- perf: Handle intermediate `Projection` nodes in `EliminateOuterJoin` [#22534](https://github.com/apache/datafusion/pull/22534) (neilconway) +- perf: array-free fast paths for `ScalarValue::cast_to` [#22576](https://github.com/apache/datafusion/pull/22576) (alamb) +- perf(optimizer): EliminateCrossJoin fast-path for join-free plans [#22612](https://github.com/apache/datafusion/pull/22612) (zhuqi-lucas) +- perf: optimize date subtraction to avoid intermediate array allocation [#22591](https://github.com/apache/datafusion/pull/22591) (lyne7-sc) +- perf: optimize arrays_zip perfect list zips [#22285](https://github.com/apache/datafusion/pull/22285) (puneetdixit200) +- perf: Reorder predicates in conjuncts via simple heuristic [#22343](https://github.com/apache/datafusion/pull/22343) (neilconway) +- perf: avoid unnecessary large allocations [#22558](https://github.com/apache/datafusion/pull/22558) (ariel-miculas) +- perf: Optimize semi-, anti-join index alignment [#22794](https://github.com/apache/datafusion/pull/22794) (neilconway) +- perf: improve approx_distinct performance 100x when there are fewer distinct values with many groups [#22768](https://github.com/apache/datafusion/pull/22768) (haohuaijin) +- perf: fast-path inline strings in ByteViewGroupValueBuilder::vectorized_append [#21794](https://github.com/apache/datafusion/pull/21794) (EeshanBembi) +- perf: Convert inner joins to semi joins when equivalent [#22652](https://github.com/apache/datafusion/pull/22652) (neilconway) +- refactor: use raw view access in do_append_val_inner and consolidate duplicated logic [#22907](https://github.com/apache/datafusion/pull/22907) (EeshanBembi) +- perf: avoid possibly expensive string formatting if no error is encountered [#23157](https://github.com/apache/datafusion/pull/23157) (tschwarzinger) +- Perf: cache primitive sort key in SortPreservingMerge to drop per-comparison bounds checks [#23162](https://github.com/apache/datafusion/pull/23162) (Dandandan) +- IN LIST: add UInt16 bitmap filter [#23012](https://github.com/apache/datafusion/pull/23012) (geoffreyclaude) +- perf: coalesce single-column sort runs to cut merge fan-in [#23202](https://github.com/apache/datafusion/pull/23202) (Dandandan) +- perf: share encoder/reservation across PartitionedTopKExec partition … [#23096](https://github.com/apache/datafusion/pull/23096) (SubhamSinghal) +- Optimize Int8 and Int16 integer IN filters [#23299](https://github.com/apache/datafusion/pull/23299) (alamb) +- feat: Implement state conversion for remaining group accumulators [#23275](https://github.com/apache/datafusion/pull/23275) (lyne7-sc) +- perf: optimize encode in datafusion-functions [#23456](https://github.com/apache/datafusion/pull/23456) (andygrove) +- perf: optimize ascii in datafusion-functions [#23462](https://github.com/apache/datafusion/pull/23462) (andygrove) +- perf: optimize nanvl in datafusion-functions [#23458](https://github.com/apache/datafusion/pull/23458) (andygrove) +- perf: avoid intermediate slice allocation in Spark slice function [#23481](https://github.com/apache/datafusion/pull/23481) (andygrove) +- perf: optimize make_date in datafusion-functions [#23470](https://github.com/apache/datafusion/pull/23470) (andygrove) +- perf: speedup `date_part` isodow by using `DayOfWeekMonday1` [#23491](https://github.com/apache/datafusion/pull/23491) (theirix) +- perf: optimisation for date_part with seconds [#23444](https://github.com/apache/datafusion/pull/23444) (theirix) +- perf: optimize `round` expression [#23471](https://github.com/apache/datafusion/pull/23471) (andygrove) +- perf: optimize `string_trim` [#23541](https://github.com/apache/datafusion/pull/23541) (andygrove) +- perf: optimize `date_trunc` [#23542](https://github.com/apache/datafusion/pull/23542) (andygrove) +- perf: Optimize array_has() for array needle [#23337](https://github.com/apache/datafusion/pull/23337) (freakyzoidberg) +- perf: optimize `trunc` for scalar precision case (10x faster) [#23593](https://github.com/apache/datafusion/pull/23593) (andygrove) +- perf: optimize `upper` (6% faster) [#23588](https://github.com/apache/datafusion/pull/23588) (andygrove) +- perf: preallocate memory in `pad` [#23586](https://github.com/apache/datafusion/pull/23586) (theirix) +- perf: optimize `replace` (2x faster) [#23589](https://github.com/apache/datafusion/pull/23589) (andygrove) +- perf: optimize `regexp_match` for literal pattern usage (20% faster) [#23547](https://github.com/apache/datafusion/pull/23547) (andygrove) +- perf: avoid per-row copy in Spark hex byte encoding [#23473](https://github.com/apache/datafusion/pull/23473) (andygrove) +- perf: optimize `get_field` [#23537](https://github.com/apache/datafusion/pull/23537) (andygrove) +- perf: optimize `regexp_instr` (40% faster) [#23540](https://github.com/apache/datafusion/pull/23540) (andygrove) +- perf: don't re-inline CSE'd expensive expressions in projection pushdown [#23459](https://github.com/apache/datafusion/pull/23459) (fordN) +- perf: optimize left_right in datafusion-functions [#23762](https://github.com/apache/datafusion/pull/23762) (andygrove) +- perf: preserve dictionary encoding for `bit_length`, `octet_length`, and `ascii` [#23743](https://github.com/apache/datafusion/pull/23743) (lyne7-sc) +- perf: optimize LEAD/LAG IGNORE NULLS evaluation [#23711](https://github.com/apache/datafusion/pull/23711) (xudong963) +- feat: add OR pre-selection short-circuit [#22979](https://github.com/apache/datafusion/pull/22979) (kumarUjjawal) +- perf: optimize `find_in_set` (up to 24x faster) [#23460](https://github.com/apache/datafusion/pull/23460) (andygrove) +- refactor: share hex encoding across datafusion-common, functions, and spark [#23766](https://github.com/apache/datafusion/pull/23766) (andygrove) +- perf: avoid per-row String allocation in Spark bin and char [#23881](https://github.com/apache/datafusion/pull/23881) (andygrove) +- IN LIST: add branchless filter for small primitive lists [#23014](https://github.com/apache/datafusion/pull/23014) (geoffreyclaude) +- perf: optimize `array_empty` udf [#23923](https://github.com/apache/datafusion/pull/23923) (rluvaton) +- feat(physical-plan): generic Rows-backed GroupColumn keeps mixed schemas on the column-wise path [#23523](https://github.com/apache/datafusion/pull/23523) (zhuqi-lucas) +- perf: `array_agg()` performance improvements [#23716](https://github.com/apache/datafusion/pull/23716) (fred1268) +- perf: Optimize hashing, null-free fast path for `percentile_cont`, `median` [#23954](https://github.com/apache/datafusion/pull/23954) (neilconway) +- perf: null-free fast path for COUNT(DISTINCT) primitive accumulator [#23956](https://github.com/apache/datafusion/pull/23956) (viirya) +- perf: precompile formats in to_time [#23964](https://github.com/apache/datafusion/pull/23964) (lyne7-sc) +- perf: Replace SipHash with foldhash in `BoundedWindowAggExec` [#23984](https://github.com/apache/datafusion/pull/23984) (neilconway) +- perf: preserve dictionary encoding for `character_length`, `initcap`, and `reverse` [#23930](https://github.com/apache/datafusion/pull/23930) (lyne7-sc) +- perf: skip re-slicing window partition batches with nothing to prune [#24047](https://github.com/apache/datafusion/pull/24047) (neilconway) +- perf: gather Linear-mode window input more efficiently [#24034](https://github.com/apache/datafusion/pull/24034) (neilconway) +- perf: use Vec in ArrowBytesMap [#24071](https://github.com/apache/datafusion/pull/24071) (Punisheroot) +- perf: preallocate RowsGroupColumn buffers in take_n [#24070](https://github.com/apache/datafusion/pull/24070) (saadtajwar) +- feat: add GroupColumn support for Decimal256 in multi-column GROUP BY [#23849](https://github.com/apache/datafusion/pull/23849) (tohuya6) +- perf: preserve dictionary encoding for `btrim`, `ltrim`, and `rtrim` [#24100](https://github.com/apache/datafusion/pull/24100) (lyne7-sc) +- Skip page index load (and `ParquetMetaData` clone) when the file has no page index [#24150](https://github.com/apache/datafusion/pull/24150) (alamb) +- perf: optimize char -> byte offset mapping in `regexp_count` [#24153](https://github.com/apache/datafusion/pull/24153) (neilconway) +- perf: skip evaluating fully calculated window partitions [#24127](https://github.com/apache/datafusion/pull/24127) (neilconway) +- perf: prune window state only for partitions that made progress [#24148](https://github.com/apache/datafusion/pull/24148) (neilconway) + +**Implemented enhancements:** + +- feat: fix `slice` function on OOB ranges [#22404](https://github.com/apache/datafusion/pull/22404) (comphead) +- feat: Analyze `VALUES` for nullability [#22089](https://github.com/apache/datafusion/pull/22089) (neilconway) +- feat: Add Spark-compatible `monthname` function to datafusion-spark [#21639](https://github.com/apache/datafusion/pull/21639) (JeelRajodiya) +- feat: Improve display of `Decimal` values [#22500](https://github.com/apache/datafusion/pull/22500) (neilconway) +- feat(catalog): expose InformationSchemataBuilder as public API [#22499](https://github.com/apache/datafusion/pull/22499) (zfarrell) +- feat: add array_scale scalar function [#22466](https://github.com/apache/datafusion/pull/22466) (crm26) +- feat: adds array_add function [#22459](https://github.com/apache/datafusion/pull/22459) (SubhamSinghal) +- feat: add TableSchemaBuilder and store partition columns as Fields [#22496](https://github.com/apache/datafusion/pull/22496) (adriangb) +- feat: lower repartition_file_min_size default from 10 MiB to 1 MiB [#22439](https://github.com/apache/datafusion/pull/22439) (adriangb) +- feat: Plumb Parquet virtual columns (row_number) through TableSchema and ParquetOpener [#22026](https://github.com/apache/datafusion/pull/22026) (mbutrovich) +- feat: add SparkPow UDF returning Infinity for pow(0, negative) [#22605](https://github.com/apache/datafusion/pull/22605) (Brijesh-Thakkar) +- feat: add array_subtract scalar function [#22556](https://github.com/apache/datafusion/pull/22556) (SubhamSinghal) +- feat: support Boolean in approx_distinct [#22707](https://github.com/apache/datafusion/pull/22707) (JeelRajodiya) +- feat: implement retract_batch for array_agg(DISTINCT) sliding window [#22719](https://github.com/apache/datafusion/pull/22719) (SubhamSinghal) +- feat: add DataFrame fill_nan [#22702](https://github.com/apache/datafusion/pull/22702) (Nagato-Yuzuru) +- feat: add array_sum scalar function [#22542](https://github.com/apache/datafusion/pull/22542) (crm26) +- feat: Support IEEE 754 negative zero semantics [#22835](https://github.com/apache/datafusion/pull/22835) (comphead) +- feat: Add From> trait for Precision enum [#22792](https://github.com/apache/datafusion/pull/22792) (devanbenz) +- feat: implement Spark-compatible weekday function [#22740](https://github.com/apache/datafusion/pull/22740) (sjhddh) +- feat(spark): add `concat_ws` with array support [#20928](https://github.com/apache/datafusion/pull/20928) (davidlghellin) +- feat: support reading from stdin in datafusion-cli [#22839](https://github.com/apache/datafusion/pull/22839) (huan233usc) +- feat(unparser): support binary literals [#23001](https://github.com/apache/datafusion/pull/23001) (zyuiop) +- feat: warn on NULL equality predicates [#22948](https://github.com/apache/datafusion/pull/22948) (ametel01) +- feat: support file-level parquet row selections [#22940](https://github.com/apache/datafusion/pull/22940) (haohuaijin) +- feat: support mixed binary and string types for concat UDFs [#22244](https://github.com/apache/datafusion/pull/22244) (theirix) +- feat(unparser): support DISTINCT FROM operators in the MySQL dialect [#22999](https://github.com/apache/datafusion/pull/22999) (zyuiop) +- feat: Add new `input_file_name` UDF for file-backed scans [#22978](https://github.com/apache/datafusion/pull/22978) (AdamGS) +- feat: add array_avg scalar function [#23168](https://github.com/apache/datafusion/pull/23168) (crm26) +- feat: Support Decimal type in `approx_distinct` [#23190](https://github.com/apache/datafusion/pull/23190) (mkleen) +- feat: Support interval type in approx_distinct [#23234](https://github.com/apache/datafusion/pull/23234) (mkleen) +- feat: Re-spill sort stream if unable to reserve for 2 streams [#22945](https://github.com/apache/datafusion/pull/22945) (EmilyMatt) +- feat: Expose cache hits in statistics_cache function [#23253](https://github.com/apache/datafusion/pull/23253) (mkleen) +- feat: Eagerly drop last finished stream in `FusedStreams` [#23283](https://github.com/apache/datafusion/pull/23283) (rluvaton) +- feat: cap spill merge fan-in [#23066](https://github.com/apache/datafusion/pull/23066) (yinli-systems) +- feat: Support duration type in approx_distinct [#23291](https://github.com/apache/datafusion/pull/23291) (mkleen) +- feat: Allow datafusion-ffi to opt out of proto parquet [#22951](https://github.com/apache/datafusion/pull/22951) (Xuanwo) +- feat: Support BinaryView type in approx_distinct [#23333](https://github.com/apache/datafusion/pull/23333) (mkleen) +- feat: support decimals in trunc UDF [#23320](https://github.com/apache/datafusion/pull/23320) (theirix) +- feat: add strictness metadata for scalar UDF null propagation and use it in outer join elimination [#23148](https://github.com/apache/datafusion/pull/23148) (lyne7-sc) +- feat: physical execution for range partitioning [#23231](https://github.com/apache/datafusion/pull/23231) (saadtajwar) +- feat: Support FixedSizedBinary type for approx_distinct [#23417](https://github.com/apache/datafusion/pull/23417) (mkleen) +- feat: Support List/ListView types in approx_distinct [#23443](https://github.com/apache/datafusion/pull/23443) (mkleen) +- feat: add array_first higher-order array function [#23267](https://github.com/apache/datafusion/pull/23267) (EdsonPetry) +- feat: Expose cache hits in list_files_cache function [#23439](https://github.com/apache/datafusion/pull/23439) (mkleen) +- feat: Support Map type in approx_distinct [#23526](https://github.com/apache/datafusion/pull/23526) (mkleen) +- feat: allow Partitioning::Range to satisfy window Distribution::KeyPartitioned requirements [#23416](https://github.com/apache/datafusion/pull/23416) (mithuncy) +- feat: benchmark_runner, improve `--list`, optional `DATA_DIR` [#23354](https://github.com/apache/datafusion/pull/23354) (Omega359) +- feat: Support Struct type in approx_distinct [#23663](https://github.com/apache/datafusion/pull/23663) (mkleen) +- feat: allow Full joins to reuse range co-partitioning in HashJoinExec [#23583](https://github.com/apache/datafusion/pull/23583) (mattp5657) +- feat: support co-partitioned range right-side equi hash joins [#23484](https://github.com/apache/datafusion/pull/23484) (gmhelmold) +- feat: complete range repartition physical planning [#23617](https://github.com/apache/datafusion/pull/23617) (saadtajwar) +- feat: Support Union type in approx_distinct [#23714](https://github.com/apache/datafusion/pull/23714) (mkleen) +- feat: add validating non-Arrow TDigest constructor and accessors [#23737](https://github.com/apache/datafusion/pull/23737) (adriangb) +- feat: add Spark-compatible hypot function [#23774](https://github.com/apache/datafusion/pull/23774) (KarpagamKarthikeyan) +- feat: add BuildHasher variants for hash_utils [#21820](https://github.com/apache/datafusion/pull/21820) (xudong963) +- feat: support `ansi` for `elt` [#23928](https://github.com/apache/datafusion/pull/23928) (comphead) +- feat: centralizing higher-order list lambda evaluation helpers [#23911](https://github.com/apache/datafusion/pull/23911) (saadtajwar) +- feat: add GroupColumn support for Float16 in multi-column GROUP BY [#23785](https://github.com/apache/datafusion/pull/23785) (tohuya6) +- feat: switch VirtualTable producer to use expressions field instead of deprecated values [#23672](https://github.com/apache/datafusion/pull/23672) (eliot1480) +- feat: add GroupColumn support for Interval in multi-column GROUP BY [#23786](https://github.com/apache/datafusion/pull/23786) (tohuya6) +- feat: drop generator on error to free memory faster [#23967](https://github.com/apache/datafusion/pull/23967) (rluvaton) +- feat: eliminate LEFT/RIGHT JOINs with redundant sides [#23566](https://github.com/apache/datafusion/pull/23566) (simonvandel) +- feat(parquet): multi-column lexicographic stats reorder for TopK sort pushdown [#23888](https://github.com/apache/datafusion/pull/23888) (zhuqi-lucas) +- feat: add Spark-compatible atan2 function [#23962](https://github.com/apache/datafusion/pull/23962) (KarpagamKarthikeyan) +- feat: Calculate non-distinct `sum` from column statistics when available [#23863](https://github.com/apache/datafusion/pull/23863) (AdamGS) +- feat: prune unread Parquet leaves when a nested column is cast to a narrower type [#24090](https://github.com/apache/datafusion/pull/24090) (mbutrovich) +- feat: Add SQL planner, physical planner, and TableProvider hook for MERGE INTO [#22988](https://github.com/apache/datafusion/pull/22988) (wirybeaver) +- feat(dataframe): add f16 support to dataframe! macro [#24234](https://github.com/apache/datafusion/pull/24234) (cj-zhukov) + +**Fixed bugs:** + +- fix: indentation for markdown block comments in docstrings [#22409](https://github.com/apache/datafusion/pull/22409) (ariel-miculas) +- fix(unparser): fold Limit/Sort into outer SELECT when Projection claims Aggregate through them [#21375](https://github.com/apache/datafusion/pull/21375) (yonatan-sevenai) +- fix(substrait): dedupe names of aggregate measures, not just groupings [#22453](https://github.com/apache/datafusion/pull/22453) (LiaCastaneda) +- fix: `Operator::returns_null_on_null()` should include string concat (`||`) [#22458](https://github.com/apache/datafusion/pull/22458) (neilconway) +- fix: custom_datasource example ignores projection pushdown in execute() [#22417](https://github.com/apache/datafusion/pull/22417) (kumarUjjawal) +- fix: handle `IS TRUE` correctly in `EliminateOuterJoin` [#22444](https://github.com/apache/datafusion/pull/22444) (neilconway) +- fix: avoid panic in TableSchema::with_table_partition_cols on shared Arc [#22372](https://github.com/apache/datafusion/pull/22372) (adriangb) +- fix: avoid panic in date_bin compute_distance near i64::MIN [#22408](https://github.com/apache/datafusion/pull/22408) (SAY-5) +- fix: make array null argument handling follow SQL semantics [#22508](https://github.com/apache/datafusion/pull/22508) (kumarUjjawal) +- fix: Set Substrait output types for expressions [#20597](https://github.com/apache/datafusion/pull/20597) (wlhjason) +- fix: clear handled OFFSET before child recursion in LimitPushdown [#22525](https://github.com/apache/datafusion/pull/22525) (kumarUjjawal) +- fix: Avoid precision loss for `atan2` with integer args [#22516](https://github.com/apache/datafusion/pull/22516) (neilconway) +- fix: LIKE 'prefix%' pruning fails on Utf8View and LargeUtf8 columns [#22562](https://github.com/apache/datafusion/pull/22562) (lyne7-sc) +- fix: widen `power(decimal, float)` to Float64, fix bugs [#22482](https://github.com/apache/datafusion/pull/22482) (neilconway) +- fix: reborrow metadata values when intersecting union metadata [#22491](https://github.com/apache/datafusion/pull/22491) (officialasishkumar) +- fix: Correct join cardinality estimation for semi and anti joins with disjoint column ranges [#22674](https://github.com/apache/datafusion/pull/22674) (neilconway) +- fix: Projection stats Absent for columns referenced >1 time [#22679](https://github.com/apache/datafusion/pull/22679) (neilconway) +- fix(substrait): plan nested projected window expressions [#22630](https://github.com/apache/datafusion/pull/22630) (bvolpato) +- fix: render binary columns as hex in DataFrame::describe() [#21728](https://github.com/apache/datafusion/pull/21728) (diegoQuinas) +- fix: wrong precision in a decimal256 log test [#22578](https://github.com/apache/datafusion/pull/22578) (theirix) +- fix: Avoid panic decoding invalid parquet writer version from proto [#22467](https://github.com/apache/datafusion/pull/22467) (fallintoplace) +- fix: make PushDownLeafProjections work with unnest [#22620](https://github.com/apache/datafusion/pull/22620) (pabadrubio) +- fix: correct cross join byte size statistics [#22700](https://github.com/apache/datafusion/pull/22700) (neilconway) +- fix: Improve consistency of per-column stats on `FilterExec` output [#22718](https://github.com/apache/datafusion/pull/22718) (neilconway) +- fix: Correct computation of selectivity for multi-key joins [#22725](https://github.com/apache/datafusion/pull/22725) (neilconway) +- fix: replace with empty search string should be a no-op [#22497](https://github.com/apache/datafusion/pull/22497) (Amogh-2404) +- fix: Remove `power(decimal, int)` code path [#22651](https://github.com/apache/datafusion/pull/22651) (neilconway) +- fix: avoid extraneous casts for equivalent nested types [#20945](https://github.com/apache/datafusion/pull/20945) (feichai0017) +- fix: handle NULLs in sliding SUM(DISTINCT) window frames [#22755](https://github.com/apache/datafusion/pull/22755) (kumarUjjawal) +- fix: Scale semi/anti-join column stats by estimated row count [#22762](https://github.com/apache/datafusion/pull/22762) (neilconway) +- fix: preserve timestamp precision when coercing mixed time units [#22759](https://github.com/apache/datafusion/pull/22759) (fengys1996) +- fix: make skip_partial_aggregation_probe_ratio_threshold match the docs [#22752](https://github.com/apache/datafusion/pull/22752) (haohuaijin) +- fix: NestedLoopJoinExec emits spurious unmatched-left rows with multiple probe partitions [#22791](https://github.com/apache/datafusion/pull/22791) (nathanb9) +- fix: Optimize projections in recursive CTEs [#22476](https://github.com/apache/datafusion/pull/22476) (nuno-faria) +- fix: Coerce aggregate FILTER predicates to boolean [#22774](https://github.com/apache/datafusion/pull/22774) (pchintar) +- fix: approx_distinct over-counts for utf8view [#22815](https://github.com/apache/datafusion/pull/22815) (haohuaijin) +- fix: regex simplification of anchored patterns produces wrong results [#22727](https://github.com/apache/datafusion/pull/22727) (lyne7-sc) +- fix: add backtrace for `assert_*_or_internal_err` helpers [#18910](https://github.com/apache/datafusion/pull/18910) (rluvaton) +- fix: map() fails when keys are literals and values are column expressions [#22784](https://github.com/apache/datafusion/pull/22784) (nathanb9) +- fix: Avoid incorrectly rounding large integers in `nanvl` [#22575](https://github.com/apache/datafusion/pull/22575) (neilconway) +- fix: Enable sliding window execution for covar_pop, covar_samp, and corr [#22764](https://github.com/apache/datafusion/pull/22764) (pchintar) +- fix: handle `date_bin` negative subsecond and overflow cases [#22610](https://github.com/apache/datafusion/pull/22610) (kumarUjjawal) +- fix: TRY_CAST returns NULL for timestamp/date overflow [#22897](https://github.com/apache/datafusion/pull/22897) (fengys1996) +- fix: count shared buffers once in hash join build-side memory accounting [#22862](https://github.com/apache/datafusion/pull/22862) (jordepic) +- fix(topk): call attempt_early_completion when filter rejects entire batch [#22852](https://github.com/apache/datafusion/pull/22852) (ajegou) +- fix: Disable join dynamic filters for null-equal joins [#22965](https://github.com/apache/datafusion/pull/22965) (neilconway) +- fix: ProjectionPushdown internal error on NestedLoopJoin mark joins [#22902](https://github.com/apache/datafusion/pull/22902) (lyne7-sc) +- fix: parquet limit pruning for row group selections [#22942](https://github.com/apache/datafusion/pull/22942) (haohuaijin) +- fix: isolate anonymous file statistics cache [#22950](https://github.com/apache/datafusion/pull/22950) (kumarUjjawal) +- fix: Parquet bloom filter pruning can incorrectly filter decimals encoded as FIXED_LEN_BYTE_ARRAY [#22995](https://github.com/apache/datafusion/pull/22995) (lyne7-sc) +- fix: Consider column names' case when aliasing tables [#22917](https://github.com/apache/datafusion/pull/22917) (nuno-faria) +- fix: prevent unparser stack overflow on deeply nested expressions [#23058](https://github.com/apache/datafusion/pull/23058) (adriangb) +- fix: block timestamp precision narrowing unwrap [#22837](https://github.com/apache/datafusion/pull/22837) (discord9) +- fix: preserve no-filter SMJ matches across pending outer batches [#23049](https://github.com/apache/datafusion/pull/23049) (neilconway) +- fix(proto): honor ExecutionPlan downcast_delegate during serialization [#23154](https://github.com/apache/datafusion/pull/23154) (geoffreyclaude) +- fix: add assert to `HashJoinExec::swap_inputs` [#23078](https://github.com/apache/datafusion/pull/23078) (haohuaijin) +- fix: preserve empty projection when ser/de `HashJoinExec` and `NestedLoopJoinExec` [#23082](https://github.com/apache/datafusion/pull/23082) (haohuaijin) +- fix: `array_compact` handle edge case with NULLs [#23192](https://github.com/apache/datafusion/pull/23192) (comphead) +- fix(spark): return error from ELT coerce_types when fewer than 2 args [#23164](https://github.com/apache/datafusion/pull/23164) (davidlghellin) +- fix: Preserve integer values in round() for large Int64 and UInt64 inputs [#22697](https://github.com/apache/datafusion/pull/22697) (pchintar) +- fix: surface BufferExec input panics instead of silently truncating output [#23243](https://github.com/apache/datafusion/pull/23243) (Tristan1900) +- fix: apply recursive CTE column-list aliases to the static term [#23098](https://github.com/apache/datafusion/pull/23098) (tomsanbear) +- fix: unparse columns of stacked pushdown projections unqualified [#23176](https://github.com/apache/datafusion/pull/23176) (Phoenix500526) +- fix(sort): record output_batches, output_bytes and end_time for when not using merge sort [#22878](https://github.com/apache/datafusion/pull/22878) (rluvaton) +- fix: Handle decimal columns consistently in SLT tests [#23161](https://github.com/apache/datafusion/pull/23161) (AdamGS) +- fix: avoid panic parsing non-ASCII runtime config values [#23316](https://github.com/apache/datafusion/pull/23316) (ByteBaker) +- fix: avoid global SQL stack guard mutation in unparser [#23284](https://github.com/apache/datafusion/pull/23284) (ametel01) +- fix: Avoid panicing when stats are not available for a file group split [#23277](https://github.com/apache/datafusion/pull/23277) (mkleen) +- fix: gate debug-only assertions in physical planner test test_optimization_invariant_checker [#23323](https://github.com/apache/datafusion/pull/23323) (buraksenn) +- fix: Reject out-of-range `ArrayMap` probe keys on 32-bit targets [#22911](https://github.com/apache/datafusion/pull/22911) (neilconway) +- fix: return execution error instead of capacity overflow panic in array_resize [#23306](https://github.com/apache/datafusion/pull/23306) (buraksenn) +- fix: cardinality returns incorrect results for ragged nested arrays [#23271](https://github.com/apache/datafusion/pull/23271) (lyne7-sc) +- fix: cast `[]` to `FixedSizeList(0, _)` [#23381](https://github.com/apache/datafusion/pull/23381) (Jefffrey) +- fix: don't duplicate volatile expressions when pushing projection into file scan [#23395](https://github.com/apache/datafusion/pull/23395) (fordN) +- fix: fix typo on doc [#23457](https://github.com/apache/datafusion/pull/23457) (Rich-T-kid) +- fix: Batch size limit in re-spill compounds [#23286](https://github.com/apache/datafusion/pull/23286) (EmilyMatt) +- fix: ensure a maximum of `buffer_len` RecordBatches are cached in `spawn_buffered` [#23560](https://github.com/apache/datafusion/pull/23560) (ariel-miculas) +- fix: close the markdown block in docstring [#23562](https://github.com/apache/datafusion/pull/23562) (ariel-miculas) +- fix: preserve range partitioning through joins [#23584](https://github.com/apache/datafusion/pull/23584) (EdsonPetry) +- fix: optimize_projections failure with struct-field join keys [#22903](https://github.com/apache/datafusion/pull/22903) (kumarUjjawal) +- fix: Handle potential overflow in internal state for `avg(decimal)` [#22714](https://github.com/apache/datafusion/pull/22714) (AdamGS) +- fix: support type coercion for MAP literals with NULL values in VALUES lists [#23521](https://github.com/apache/datafusion/pull/23521) (PG1204) +- fix: handle interleaved HashJoin projections in sort pushdown [#23591](https://github.com/apache/datafusion/pull/23591) (xudong963) +- fix: do not remove DISTINCT when a unique key was downgraded by a join [#23548](https://github.com/apache/datafusion/pull/23548) (simonvandel) +- fix: preserve aggregate scope when unparsing [#23327](https://github.com/apache/datafusion/pull/23327) (Phoenix500526) +- fix: keep null-aware anti-join NULLs in the pushed dynamic filter [#23104](https://github.com/apache/datafusion/pull/23104) (mdashti) +- fix: handle null date and timestamp format arguments [#23641](https://github.com/apache/datafusion/pull/23641) (lyne7-sc) +- fix: prevent LEAD/LAG IGNORE NULLS panic without null bitmap [#23706](https://github.com/apache/datafusion/pull/23706) (xudong963) +- fix: Preserve metadata when a cross-join is swapped [#23605](https://github.com/apache/datafusion/pull/23605) (mkleen) +- fix: coerce SIMILAR TO operands to a common string type [#23704](https://github.com/apache/datafusion/pull/23704) (u70b3) +- fix: Capture global ORDER BY requirement under ScalarSubqueryExec root [#23677](https://github.com/apache/datafusion/pull/23677) (sgrebnov) +- fix: avoid overflow in join cardinality estimation [#23788](https://github.com/apache/datafusion/pull/23788) (xudong963) +- fix: unwrap identity Date cast in comparison unwrapping [#23727](https://github.com/apache/datafusion/pull/23727) (adriangb) +- fix: reject nested aggregate functions (e.g. `sum(sum(x))`) during logical planning [#23813](https://github.com/apache/datafusion/pull/23813) (adriangb) +- fix: array_any_value returns NULL for empty list elements [#23775](https://github.com/apache/datafusion/pull/23775) (bjchambers) +- fix: fixed decode buffer size estimate for BinaryViewArray [#23765](https://github.com/apache/datafusion/pull/23765) (liningpan) +- fix: grouped first_value/last_value FILTER excludes NULL predicate rows [#23707](https://github.com/apache/datafusion/pull/23707) (u70b3) +- fix: NOT IN with NULL subquery returns wrong results under SortMergeJoin [#22810](https://github.com/apache/datafusion/pull/22810) (nathanb9) +- fix: align physical CASE nullability through casts [#23844](https://github.com/apache/datafusion/pull/23844) (friendlymatthew) +- fix: Handle null-aware joins correctly in `FilterNullJoinKeys` when its enabled [#23848](https://github.com/apache/datafusion/pull/23848) (AdamGS) +- fix: don't infer join predicates for null-aware joins in push_down_filter [#23901](https://github.com/apache/datafusion/pull/23901) (viirya) +- fix: skip dynamic filter pushdown for null-aware anti joins with a nullable build key [#23173](https://github.com/apache/datafusion/pull/23173) (mdashti) +- fix: Handle `input_file_name()` pushdown into `ParquetSource` with filter pushdown enabled [#23638](https://github.com/apache/datafusion/pull/23638) (AdamGS) +- fix: exclude precision-losing integer-to-float conversions from CastExpr::check_bigger_cast (#23808) [#23809](https://github.com/apache/datafusion/pull/23809) (getChan) +- fix: eliminate group by constant empty input [#22132](https://github.com/apache/datafusion/pull/22132) (HairstonE) +- fix: sliding window `min()` returns wrong value for all-NULL windows [#23874](https://github.com/apache/datafusion/pull/23874) (neilconway) +- fix: correct percentile_cont(DISTINCT) accumulation and sliding-window retract [#23913](https://github.com/apache/datafusion/pull/23913) (viirya) +- fix: support parentheses for negative decimal formatting [#23718](https://github.com/apache/datafusion/pull/23718) (wangzhigang1999) +- fix: last value accumulator merge indexing [#23905](https://github.com/apache/datafusion/pull/23905) (peterxcli) +- fix: accept LargeUtf8 and Utf8View patterns in SIMILAR TO planning [#23735](https://github.com/apache/datafusion/pull/23735) (u70b3) +- fix: preserve aggregate filter pushdown order [#22926](https://github.com/apache/datafusion/pull/22926) (discord9) +- fix(common): preserve an exact zero through filter selectivity estimation [#23936](https://github.com/apache/datafusion/pull/23936) (asolimando) +- fix: preserve dictionary-value nulls in scalar regex operators [#23966](https://github.com/apache/datafusion/pull/23966) (discord9) +- fix(datasource): avoid over-conservative transformation of num_rows statistics in file scan config [#23670](https://github.com/apache/datafusion/pull/23670) (tschwarzinger) +- fix: keep a CoalescePartitionsExec required by a SinglePartition child [#23948](https://github.com/apache/datafusion/pull/23948) (adriangb) +- fix: TopK aggregation drops groups whose MIN/MAX value is NULL [#23684](https://github.com/apache/datafusion/pull/23684) (u70b3) +- fix(sql): preserve source qualifiers in CTAS with explicit schema [#23879](https://github.com/apache/datafusion/pull/23879) (lyne7-sc) +- fix: reject nested arrays in array_distance [#23995](https://github.com/apache/datafusion/pull/23995) (2010YOUY01) +- fix: Improve error message for metadata conflict in schema [#23952](https://github.com/apache/datafusion/pull/23952) (mkleen) +- fix: handle empty patterns in regexp_instr [#24054](https://github.com/apache/datafusion/pull/24054) (iamhaseebn) +- fix: prevent incorrect results when pushing filters through anti joins [#24045](https://github.com/apache/datafusion/pull/24045) (buraksenn) +- fix: UnionExec now conforms each batch to the union's declared schema [#23861](https://github.com/apache/datafusion/pull/23861) (dariocurr) +- fix: do not derive ordering for arithmetic that can overflow [#23910](https://github.com/apache/datafusion/pull/23910) (buraksenn) +- fix: preserve projection field metadata during physical planning [#23981](https://github.com/apache/datafusion/pull/23981) (subotac) +- fix(proto): preserve empty projection when ser/de MemoryScanExec [#24087](https://github.com/apache/datafusion/pull/24087) (buraksenn) +- fix: Fix nullability of logical `InSubquery` expression [#23429](https://github.com/apache/datafusion/pull/23429) (AdamGS) +- fix: keep every spilled slice of a sort-merge join inner key group [#24056](https://github.com/apache/datafusion/pull/24056) (buraksenn) +- fix: reduce peak memory usage when round robin tiebreaker is disabled [#23606](https://github.com/apache/datafusion/pull/23606) (ariel-miculas) +- fix: preserve total_byte_size in calculate_total_byte_size when num_r… [#24027](https://github.com/apache/datafusion/pull/24027) (bert-beyondloops) +- fix: box aws-config loading future avoid clippy warning [#24175](https://github.com/apache/datafusion/pull/24175) (neilconway) +- fix: Correctly process numeric literals with underscores [#24046](https://github.com/apache/datafusion/pull/24046) (nuno-faria) +- fix: typo for the builder error type [#24052](https://github.com/apache/datafusion/pull/24052) (JosephLenton) +- fix: support untyped NULL input for median [#24104](https://github.com/apache/datafusion/pull/24104) (Sigma-Ma) +- fix(proto): prevent logical plan serialization stack overflow [#24124](https://github.com/apache/datafusion/pull/24124) (mithuncy) +- fix: prevent next_day panic on far-future start dates [#24194](https://github.com/apache/datafusion/pull/24194) (viirya) +- fix: return error instead of panic when decoding ParquetScan/AvroScan without features [#24198](https://github.com/apache/datafusion/pull/24198) (nam2ee) +- fix: re-enable null-equal join dynamic filters with an IS NULL predicate [#23106](https://github.com/apache/datafusion/pull/23106) (mdashti) +- fix: generate_series overflow panics at i64 boundary and out-of-range dates [#23723](https://github.com/apache/datafusion/pull/23723) (u70b3) +- fix: clear stale sliding aggregate state for empty RANGE frames [#24185](https://github.com/apache/datafusion/pull/24185) (lyne7-sc) +- fix(parquet): remap sorting columns for partitioned writes [#24211](https://github.com/apache/datafusion/pull/24211) (xudong963) +- fix: reject max_buffered_batches_per_output_file values below 2 [#24204](https://github.com/apache/datafusion/pull/24204) (DevShiba) +- fix: avoid buffering unbounded repartition output indefinitely [#24193](https://github.com/apache/datafusion/pull/24193) (goutamadwant) +- fix: infer placeholder types in GROUP BY, HAVING, QUALIFY and ORDER BY (fix for #24042) [#24043](https://github.com/apache/datafusion/pull/24043) (Braedon-Wooding-Displayr) +- fix: Propagate NULLs in `regexp_count`, `regexp_instr` [#24239](https://github.com/apache/datafusion/pull/24239) (neilconway) +- fix: preserve NULL semantics in `log` and `power` simplification [#24247](https://github.com/apache/datafusion/pull/24247) (lyne7-sc) + +**Documentation updates:** + +- Revert "Add `ExecutionPlan::apply_expressions()` (#20337)" [#22437](https://github.com/apache/datafusion/pull/22437) (alamb) +- docs: add agent skill for datafusion-ffi crate patterns [#22327](https://github.com/apache/datafusion/pull/22327) (timsaucer) +- docs: clarify difference between try_cast_literal_to_type and ScalarValue::cast_to [#22592](https://github.com/apache/datafusion/pull/22592) (alamb) +- added support for MapFromEntries [#21720](https://github.com/apache/datafusion/pull/21720) (athlcode) +- chore: update Rust toolchain to 1.96.0 [#22611](https://github.com/apache/datafusion/pull/22611) (Dandandan) +- chore(deps): update pydata-sphinx-theme requirement from <1,>=0.17.1 to >=0.18.0,<1 in /docs [#22540](https://github.com/apache/datafusion/pull/22540) (dependabot[bot]) +- Track allocator-level memory vs MemoryPool during SLTs to prevent OOMs [#22626](https://github.com/apache/datafusion/pull/22626) (avantgardnerio) +- Add `array_product` UDF [#22703](https://github.com/apache/datafusion/pull/22703) (SubhamSinghal) +- docs: revise OptimizerRule trait method descriptions [#22582](https://github.com/apache/datafusion/pull/22582) (jiengup) +- docs: add Boston DataFusion meetup [#22722](https://github.com/apache/datafusion/pull/22722) (alamb) +- Add example for PartitionedFile schema [#22809](https://github.com/apache/datafusion/pull/22809) (fpetkovski) +- [main] Update version and changelog to 54.0.0 [#22855](https://github.com/apache/datafusion/pull/22855) (alamb) +- docs: link release tracking issue to release management page [#22822](https://github.com/apache/datafusion/pull/22822) (alamb) +- chore: Define backport criteria [#22766](https://github.com/apache/datafusion/pull/22766) (comphead) +- docs: Update/improve `SELECT` reference [#22672](https://github.com/apache/datafusion/pull/22672) (neilconway) +- docs: link to 2026 Q3-Q4 roadmap discussion [#22884](https://github.com/apache/datafusion/pull/22884) (alamb) +- refactor(hash-aggr): Migrate the partial aggregation skip optimization to the new hash aggregation impl [#22899](https://github.com/apache/datafusion/pull/22899) (2010YOUY01) +- Add `file_row_index` UDF to query file-level row indexes from Parquet files [#22604](https://github.com/apache/datafusion/pull/22604) (AdamGS) +- chore(deps): update maturin requirement from <2,>=1.13.3 to >=1.14.0,<2 in /docs [#22974](https://github.com/apache/datafusion/pull/22974) (dependabot[bot]) +- chore(deps): update pydata-sphinx-theme requirement from <1,>=0.18.0 to >=0.19.0,<1 in /docs [#22972](https://github.com/apache/datafusion/pull/22972) (dependabot[bot]) +- docs: clarify stdin store buffers on construction, not first use [#23060](https://github.com/apache/datafusion/pull/23060) (huan233usc) +- Docs: Add `PartialSortExec` documentation [#23092](https://github.com/apache/datafusion/pull/23092) (alamb) +- docs: Add Shanghai Apache DataFusion Meetup to events page [#23025](https://github.com/apache/datafusion/pull/23025) (alamb) +- chore(deps): update maturin requirement from <2,>=1.14.0 to >=1.14.1,<2 in /docs [#23117](https://github.com/apache/datafusion/pull/23117) (dependabot[bot]) +- Add Hotdata to the "known users" list in introduction.md [#23004](https://github.com/apache/datafusion/pull/23004) (zfarrell) +- doc: More comments on GroupedHashAggregateStream refactor [#23200](https://github.com/apache/datafusion/pull/23200) (2010YOUY01) +- docs: show struct-returning aggregate window metadata pattern [#23248](https://github.com/apache/datafusion/pull/23248) (ametel01) +- Align DataFrame::fill_null column argument with fill_nan [#22904](https://github.com/apache/datafusion/pull/22904) (Nagato-Yuzuru) +- v54 upgrade guide: Remove unreleased-note [#23331](https://github.com/apache/datafusion/pull/23331) (simonvandel) +- docs: document ClickBench setup details [#23315](https://github.com/apache/datafusion/pull/23315) (ByteBaker) +- docs: add DataFusion Ballista to related subproject [#23377](https://github.com/apache/datafusion/pull/23377) (coderfender) +- chore(deps): update setuptools requirement from <83,>=82.0.1 to >=83.0.0,<84 in /docs [#23361](https://github.com/apache/datafusion/pull/23361) (dependabot[bot]) +- [codex] chore: update Rust toolchain to 1.96.1 [#23379](https://github.com/apache/datafusion/pull/23379) (alamb) +- docs: add infino to known users [#23383](https://github.com/apache/datafusion/pull/23383) (savannahar68) +- Update Rust toolchain to 1.97.0 [#23430](https://github.com/apache/datafusion/pull/23430) (Dandandan) +- chore(deps): update pydata-sphinx-theme requirement from <1,>=0.19.0 to >=0.20.0,<1 in /docs [#23551](https://github.com/apache/datafusion/pull/23551) (dependabot[bot]) +- doc: More comments to aggregate planning overview [#23525](https://github.com/apache/datafusion/pull/23525) (2010YOUY01) +- docs: add partitioned ClickBench SQL example [#23637](https://github.com/apache/datafusion/pull/23637) (ByteBaker) +- docs: Update committer and PMC list [#23621](https://github.com/apache/datafusion/pull/23621) (alamb) +- chore: Fix duplicated word typos in comments [#23662](https://github.com/apache/datafusion/pull/23662) (jackylee-ch) +- Add any_value aggregate function [#23043](https://github.com/apache/datafusion/pull/23043) (yinli-systems) +- docs: update Polygon.io reference to Massive.com [#23734](https://github.com/apache/datafusion/pull/23734) (xudong963) +- chore: Update version 54.1.0, add changelog (#23689) [#23764](https://github.com/apache/datafusion/pull/23764) (mbutrovich) +- docs: add Supermetal to known users [#23790](https://github.com/apache/datafusion/pull/23790) (kumarUjjawal) +- chore: remove Github filter `status:success` for `pending PR` shield [#23846](https://github.com/apache/datafusion/pull/23846) (comphead) +- Add codecov badge to README [#23860](https://github.com/apache/datafusion/pull/23860) (Jefffrey) +- docs: add datapress to known users list [#23919](https://github.com/apache/datafusion/pull/23919) (jeroenflvr) +- test: add regression coverage and docs for NULL format handling [#23669](https://github.com/apache/datafusion/pull/23669) (U0001F3A2) +- docs: Fixes incorrect type name in `UserDefinedLogicalNode` comment [#23992](https://github.com/apache/datafusion/pull/23992) (vikrantmehta123) +- docs: Add more documentation about `PartialSortExec` operator [#24048](https://github.com/apache/datafusion/pull/24048) (alamb) +- Docs: Update PR template to ask for user-visible rationale [#24053](https://github.com/apache/datafusion/pull/24053) (alamb) +- docs: document all fields and methods of `DFParquetMetadata` [#24037](https://github.com/apache/datafusion/pull/24037) (alamb) +- Fix syntax examples of some functions [#23212](https://github.com/apache/datafusion/pull/23212) (Viicos) +- chore(deps): Update to arrow/parquet 59.2.0 [#24030](https://github.com/apache/datafusion/pull/24030) (alamb) +- Fix duplicated words in documentation [#24176](https://github.com/apache/datafusion/pull/24176) (latent-9) +- chore: fix some scalar function docs [#24134](https://github.com/apache/datafusion/pull/24134) (Jefffrey) +- Docs: Add community showcase to the docs page [#24217](https://github.com/apache/datafusion/pull/24217) (alamb) +- docs: explain Parquet content-defined chunking [#24155](https://github.com/apache/datafusion/pull/24155) (goutamadwant) +- docs: add IceGate to the list of featured data platforms [#24240](https://github.com/apache/datafusion/pull/24240) (frisbeeman) +- Docs: Add PR review guide [#24051](https://github.com/apache/datafusion/pull/24051) (alamb) +- [branch-55] Update additional references to version number [#24295](https://github.com/apache/datafusion/pull/24295) (timsaucer) +- [branch-55] Backport of refactor(physical-plan): Simplify `ExecutionPlan` API with `replace_children` [#24296](https://github.com/apache/datafusion/pull/24296) (JSOD11) + +**Other:** + +- chore: protect branch-53 and branch-54 [#22403](https://github.com/apache/datafusion/pull/22403) (mbutrovich) +- refactor(parquet-datasource): extract DecoderProjection from build_stream [#22398](https://github.com/apache/datafusion/pull/22398) (adriangb) +- Fix: Infer placeholder type from subquery [#22436](https://github.com/apache/datafusion/pull/22436) (HairstonE) +- Split proto serialization to encapsulate private state (#21835) [#21929](https://github.com/apache/datafusion/pull/21929) (adriangb) +- test: add more tests and docs for heap size estimation [#22358](https://github.com/apache/datafusion/pull/22358) (mkleen) +- chore: Cleanup and refactor `build_join` in `ScalarSubqueryToJoin` [#22316](https://github.com/apache/datafusion/pull/22316) (neilconway) +- Fix missing field `partitioned_by_file_group` in serialization [#22365](https://github.com/apache/datafusion/pull/22365) (marc-pydantic) +- chore: Add existence (semi / anti ) benchmarks for hashjoinexec [#21821](https://github.com/apache/datafusion/pull/21821) (coderfender) +- chore(deps): bump qs and express in /datafusion/wasmtest/datafusion-wasm-app [#22469](https://github.com/apache/datafusion/pull/22469) (dependabot[bot]) +- chore: Disallow `reserve()` in clippy to prevent panics [#22386](https://github.com/apache/datafusion/pull/22386) (2010YOUY01) +- Support DISTINCT ON with aggregation and windows [#22169](https://github.com/apache/datafusion/pull/22169) (kumarUjjawal) +- Benchmark multi-column GROUP BY performance [#22322](https://github.com/apache/datafusion/pull/22322) (nathanb9) +- fix(sort-pushdown): restore SortExec elimination after stats-based file reorder [#22493](https://github.com/apache/datafusion/pull/22493) (zhuqi-lucas) +- fix array_repeat capacity overflow on constant scalar with large count [#22305](https://github.com/apache/datafusion/pull/22305) (xiedeyantu) +- fix sqrt(-1.0::float8) should error, not return NaN [#22308](https://github.com/apache/datafusion/pull/22308) (xiedeyantu) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 9 updates [#22470](https://github.com/apache/datafusion/pull/22470) (dependabot[bot]) +- Port LikeExpr to use try_to_proto / try_from_proto [#22471](https://github.com/apache/datafusion/pull/22471) (jx2lee) +- chore(deps-dev): bump fast-uri from 3.1.0 to 3.1.2 in /datafusion/wasmtest/datafusion-wasm-app [#22083](https://github.com/apache/datafusion/pull/22083) (dependabot[bot]) +- refactor: port InListExpr to use try_to_proto/try_from_proto hooks [#22503](https://github.com/apache/datafusion/pull/22503) (kkrainov) +- refactor(physical-expr): add proto ctx expr helpers and adopt in InList/Like [#22513](https://github.com/apache/datafusion/pull/22513) (adriangb) +- minor: Make `union_with_mix_of_presorted_and_explicitly_resorted_inputs_impl` cross platform [#22478](https://github.com/apache/datafusion/pull/22478) (nuno-faria) +- refactor: add `try_to_proto` to `HashTableLookupExpr` [#22451](https://github.com/apache/datafusion/pull/22451) (AnuragRaut08) +- Simplify get_field over inline struct constructors [#22239](https://github.com/apache/datafusion/pull/22239) (adriangb) +- Add regression coverage for DATE interval overflow [#22519](https://github.com/apache/datafusion/pull/22519) (puneetdixit200) +- Make DiskManager max_temp_directory_size dynamically adjustable [#22246](https://github.com/apache/datafusion/pull/22246) (Bukhtawar) +- chore(deps): bump taiki-e/install-action from 2.79.2 to 2.79.8 [#22537](https://github.com/apache/datafusion/pull/22537) (dependabot[bot]) +- chore(deps): bump actions/stale from 10.2.0 to 10.3.0 [#22536](https://github.com/apache/datafusion/pull/22536) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.35.5 to 4.36.0 [#22535](https://github.com/apache/datafusion/pull/22535) (dependabot[bot]) +- chore(deps): bump log from 0.4.29 to 0.4.30 in the all-other-cargo-deps group [#22539](https://github.com/apache/datafusion/pull/22539) (dependabot[bot]) +- test: add test that validate partial reduce with different number of state fields [#21175](https://github.com/apache/datafusion/pull/21175) (rluvaton) +- chore: fix two comment typos [#22524](https://github.com/apache/datafusion/pull/22524) (mvanhorn) +- port `NegativeExpr` to use the `try_to_proto` / `try_from_proto` hooks [#22483](https://github.com/apache/datafusion/pull/22483) (kevinhongzl) +- chore: update sqllogictest priority list with latest timing summary (8s --> 6s) [#22549](https://github.com/apache/datafusion/pull/22549) (alamb) +- Support transparent ExecutionPlan downcasts [#22559](https://github.com/apache/datafusion/pull/22559) (geoffreyclaude) +- Return None for cardinality overflow [#22309](https://github.com/apache/datafusion/pull/22309) (jx2lee) +- Fix correlated subquery empty defaults for regr_count and approx_distinct [#22319](https://github.com/apache/datafusion/pull/22319) (nathanb9) +- refactor: port HashExpr proto hooks [#22502](https://github.com/apache/datafusion/pull/22502) (nanookclaw) +- Port CastExpr to proto hooks [#22569](https://github.com/apache/datafusion/pull/22569) (feichai0017) +- ci(breaking-change-detector): don't use `maintain-one-comment` and instead do it manually [#22568](https://github.com/apache/datafusion/pull/22568) (rluvaton) +- refactor(physical-expr-common): add proto helpers for the recurring shapes in #22418, port already-migrated exprs [#22596](https://github.com/apache/datafusion/pull/22596) (adriangb) +- Port NotExpr proto hooks [#22463](https://github.com/apache/datafusion/pull/22463) (Herrtian) +- Migrate UnKnownColumn proto hooks [#22464](https://github.com/apache/datafusion/pull/22464) (koopatroopa787) +- refactor: Port IsNotNullExpr proto serialization hooks [#22532](https://github.com/apache/datafusion/pull/22532) (chakkk309) +- Optimize Parquet metadata row-group level statistics collection [#22462](https://github.com/apache/datafusion/pull/22462) (AdamGS) +- refactor: Port IsNullExpr proto serialization hooks [#22509](https://github.com/apache/datafusion/pull/22509) (chakkk309) +- chore: Fix typos in comments [#22625](https://github.com/apache/datafusion/pull/22625) (neilconway) +- Fix TopK DISTINCT aggregation preserving NULLs [#22571](https://github.com/apache/datafusion/pull/22571) (kumarUjjawal) +- Add range partitioning sqllogictest fixture [#22607](https://github.com/apache/datafusion/pull/22607) (gene-bordegaray) +- fix(physical-plan): make HashJoinExec dynamic filter pushdown idempotent [#22523](https://github.com/apache/datafusion/pull/22523) (wirybeaver) +- minor: Improve error message for invalid column expression in `SELECT` statement [#22486](https://github.com/apache/datafusion/pull/22486) (2010YOUY01) +- fix(physical-optimizer): make OutputRequirements idempotent [#22522](https://github.com/apache/datafusion/pull/22522) (wirybeaver) +- fix(array_agg): reverse ordering_values in state() when accumulator is reversed [#22597](https://github.com/apache/datafusion/pull/22597) (ologlogn) +- chore: Add primary key constraints for TPC-H, TPC-DS [#22646](https://github.com/apache/datafusion/pull/22646) (neilconway) +- test: cover regexp_like multiline flag [#22284](https://github.com/apache/datafusion/pull/22284) (nanookclaw) +- test: make push_down_filter_regression dynamic filter content deterministic (#22621) [#22643](https://github.com/apache/datafusion/pull/22643) (diegoQuinas) +- Revert addition of benchmark_runner for sql_benchmarks [#22624](https://github.com/apache/datafusion/pull/22624) (Omega359) +- fix array_repeat scalar path overflows total repeated-value count [#22274](https://github.com/apache/datafusion/pull/22274) (xiedeyantu) +- Refactor Spark `format_string` integer conversion dispatch [#22388](https://github.com/apache/datafusion/pull/22388) (kosiew) +- chore: Make sqllogictest pass with default features [#22619](https://github.com/apache/datafusion/pull/22619) (AdamGS) +- refactor: Port TryCastExpr proto serialization hooks [#22550](https://github.com/apache/datafusion/pull/22550) (chakkk309) +- fix nth_value window function negates i64::MIN [#22304](https://github.com/apache/datafusion/pull/22304) (xiedeyantu) +- sqllogictest: account before alloc to avoid panic-after-alloc hazards [#22742](https://github.com/apache/datafusion/pull/22742) (avantgardnerio) +- chore(deps): bump taiki-e/install-action from 2.79.8 to 2.81.3 [#22745](https://github.com/apache/datafusion/pull/22745) (dependabot[bot]) +- chore(deps): bump github/codeql-action from 4.36.0 to 4.36.1 [#22746](https://github.com/apache/datafusion/pull/22746) (dependabot[bot]) +- chore(deps): bump actions/checkout from 6.0.2 to 6.0.3 [#22748](https://github.com/apache/datafusion/pull/22748) (dependabot[bot]) +- chore(deps): bump astral-sh/setup-uv from 8.1.0 to 8.2.0 [#22747](https://github.com/apache/datafusion/pull/22747) (dependabot[bot]) +- fix date_bin overflows scaling extreme Timestamp(Second) source [#22315](https://github.com/apache/datafusion/pull/22315) (xiedeyantu) +- test: benchmarks and SLT tests for push-down TopK through join [#22760](https://github.com/apache/datafusion/pull/22760) (adriangb) +- Refactor hash join build-report lifecycle into `BuildReportHandle` [#22623](https://github.com/apache/datafusion/pull/22623) (kosiew) +- Mark BufferExec and AnalyzeExec as eager [#22711](https://github.com/apache/datafusion/pull/22711) (geoffreyclaude) +- feat(physical-expr): port Literal to try_to_proto / try_from_proto hooks [#22636](https://github.com/apache/datafusion/pull/22636) (koopatroopa787) +- Add clickbench SQL benchmark [#22633](https://github.com/apache/datafusion/pull/22633) (Omega359) +- Add imdb SQL benchmark [#22680](https://github.com/apache/datafusion/pull/22680) (Omega359) +- Add partitioning compatibility API [#22590](https://github.com/apache/datafusion/pull/22590) (gene-bordegaray) +- Add h2o SQL benchmark [#22660](https://github.com/apache/datafusion/pull/22660) (Omega359) +- chore(deps): bump the all-other-cargo-deps group with 6 updates [#22750](https://github.com/apache/datafusion/pull/22750) (dependabot[bot]) +- test: make ensure_requirements tests deterministic [#22789](https://github.com/apache/datafusion/pull/22789) (kumarUjjawal) +- Spark quote function implementation [#22642](https://github.com/apache/datafusion/pull/22642) (kazantsev-maksim) +- coerce Union vs scalar in comparisons [#22825](https://github.com/apache/datafusion/pull/22825) (friendlymatthew) +- bench: add predicate_eval SQL micro-benchmark suite for conjunctive filter evaluation [#22704](https://github.com/apache/datafusion/pull/22704) (adriangb) +- minor: More comments to `AggregateMode::PartialReduce` [#22800](https://github.com/apache/datafusion/pull/22800) (2010YOUY01) +- bench: make wide_schema honor DATA_DIR like the other sql_benchmarks [#22836](https://github.com/apache/datafusion/pull/22836) (adriangb) +- refactor: Port CaseExpr proto serialization hooks [#22838](https://github.com/apache/datafusion/pull/22838) (chakkk309) +- chore(deps): bump github/codeql-action from 4.36.1 to 4.36.2 [#22842](https://github.com/apache/datafusion/pull/22842) (dependabot[bot]) +- Add tpcds SQL benchmark [#22801](https://github.com/apache/datafusion/pull/22801) (Omega359) +- chore(deps): bump taiki-e/install-action from 2.81.3 to 2.81.8 [#22841](https://github.com/apache/datafusion/pull/22841) (dependabot[bot]) +- Add hj SQL benchmark [#22802](https://github.com/apache/datafusion/pull/22802) (Omega359) +- chore(deps): bump the all-other-cargo-deps group with 3 updates [#22844](https://github.com/apache/datafusion/pull/22844) (dependabot[bot]) +- add clickbench sorted SQL benchmark [#22807](https://github.com/apache/datafusion/pull/22807) (Omega359) +- Add nlj SQL benchmark [#22805](https://github.com/apache/datafusion/pull/22805) (Omega359) +- Add clickbench extended SQL benchmark [#22804](https://github.com/apache/datafusion/pull/22804) (Omega359) +- Add smj SQL benchmark [#22803](https://github.com/apache/datafusion/pull/22803) (Omega359) +- chore(deps-dev): bump shell-quote from 1.8.3 to 1.8.4 in /datafusion/wasmtest/datafusion-wasm-app [#22856](https://github.com/apache/datafusion/pull/22856) (dependabot[bot]) +- refactor(hash-aggr): Forward port the soft limit optimization to the new hash aggregation impl [#22824](https://github.com/apache/datafusion/pull/22824) (2010YOUY01) +- refactor(physical-plan): extract make_group_column factory + eager init at try_new + tighten Time variants [#22751](https://github.com/apache/datafusion/pull/22751) (zhuqi-lucas) +- minor: handle NULL array input in array_remove and array_replace [#22790](https://github.com/apache/datafusion/pull/22790) (lyne7-sc) +- Add sort tpch SQL benchmark [#22814](https://github.com/apache/datafusion/pull/22814) (Omega359) +- chore: Update to arrow/parquet 59.0.0 [#22744](https://github.com/apache/datafusion/pull/22744) (alamb) +- Upgrade minimal tokio-postgres version to address security advisory [#22937](https://github.com/apache/datafusion/pull/22937) (AdamGS) +- Clearly gate sliding SUM(DISTINCT) type support [#22866](https://github.com/apache/datafusion/pull/22866) (kumarUjjawal) +- refactor: introduce ProbeEnd state in NestedLoopJoinExec [#22865](https://github.com/apache/datafusion/pull/22865) (nathanb9) +- refactor: Simplify heap size estimation for types that own no heap allocations [#22918](https://github.com/apache/datafusion/pull/22918) (mkleen) +- refactor(hash-aggr): Migrate existing tests on `GroupsHashAggregateStream` [#22953](https://github.com/apache/datafusion/pull/22953) (2010YOUY01) +- Include `null_aware` status in the relevant Join node display implementations [#22913](https://github.com/apache/datafusion/pull/22913) (AdamGS) +- chore(deps): bump pyjwt from 2.12.0 to 2.13.0 [#22966](https://github.com/apache/datafusion/pull/22966) (dependabot[bot]) +- ci: Setup valid `Cargo.lock` for `depcheck` to unblock CI [#22933](https://github.com/apache/datafusion/pull/22933) (AdamGS) +- chore(deps-dev): bump launch-editor from 2.10.0 to 2.14.1 in /datafusion/wasmtest/datafusion-wasm-app [#22970](https://github.com/apache/datafusion/pull/22970) (dependabot[bot]) +- chore(deps): bump cryptography from 46.0.7 to 48.0.1 [#22968](https://github.com/apache/datafusion/pull/22968) (dependabot[bot]) +- refactor: Simplify heap size estimation for arrays [#22954](https://github.com/apache/datafusion/pull/22954) (mkleen) +- Remove orphaned `snowflake_flatten_validation.sql` script [#22938](https://github.com/apache/datafusion/pull/22938) (AdamGS) +- chore(deps): bump insta-cmd from 0.6.0 to 0.7.0 [#22976](https://github.com/apache/datafusion/pull/22976) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.81.8 to 2.81.11 [#22973](https://github.com/apache/datafusion/pull/22973) (dependabot[bot]) +- chore(deps): bump prost-build from 0.14.3 to 0.14.4 [#22843](https://github.com/apache/datafusion/pull/22843) (dependabot[bot]) +- Add `.gitignore` for `proto-models` [#22977](https://github.com/apache/datafusion/pull/22977) (Jefffrey) +- Fix leaf expression reconciliation [#22971](https://github.com/apache/datafusion/pull/22971) (cetra3) +- Make LogicalPlan::Unnest expression/rebuild contracts consistent [#22783](https://github.com/apache/datafusion/pull/22783) (nathanb9) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 6 updates [#22975](https://github.com/apache/datafusion/pull/22975) (dependabot[bot]) +- Refactor outer join null-rejection analysis to track join sides directly [#22870](https://github.com/apache/datafusion/pull/22870) (kosiew) +- chore: attach Diagnostic to unary operator type errors [#21288](https://github.com/apache/datafusion/pull/21288) (hcrosse) +- refactor: make scalar distance u64 and overflow aware [#22892](https://github.com/apache/datafusion/pull/22892) (sweb) +- bugfix: changed return type of spark's width_bucket to i64 [#22811](https://github.com/apache/datafusion/pull/22811) (aguilaredu) +- chore(deps-dev): bump webpack-dev-server from 5.2.4 to 5.2.5 in /datafusion/wasmtest/datafusion-wasm-app [#23009](https://github.com/apache/datafusion/pull/23009) (dependabot[bot]) +- Add sorted TopK TPC-H benchmark target [#23003](https://github.com/apache/datafusion/pull/23003) (geoffreyclaude) +- test: correct feature gating of two datafusion-common tests [#23044](https://github.com/apache/datafusion/pull/23044) (Phoenix500526) +- test: gate hash-dependent approx_distinct tests behind not(force_hash_collisions) [#23053](https://github.com/apache/datafusion/pull/23053) (Phoenix500526) +- Skip loading Parquet page index when row-group statistics already prove it cannot prune [#22857](https://github.com/apache/datafusion/pull/22857) (RatulDawar) +- Return errors on string builder offset overflow in `replace` and `initcap` [#22990](https://github.com/apache/datafusion/pull/22990) (kosiew) +- minor: reuse ColumnarValue::into_array in map's expand_if_scalar and avoid uncessary clones [#22984](https://github.com/apache/datafusion/pull/22984) (nathanb9) +- refactor: add `try_to_proto` / `try_from_proto` to `DynamicFilterPhysicalExpr` [#22452](https://github.com/apache/datafusion/pull/22452) (AnuragRaut08) +- minor: Validate `batch_size` configuration when setting it [#23054](https://github.com/apache/datafusion/pull/23054) (2010YOUY01) +- Fix shared TopK early exit with shared prefix threshold [#22991](https://github.com/apache/datafusion/pull/22991) (geoffreyclaude) +- bench: add correlated-proxy case to the predicate_eval suite [#22919](https://github.com/apache/datafusion/pull/22919) (adriangb) +- refactor: name build-row and matchable-map presence checks in hash join [#23024](https://github.com/apache/datafusion/pull/23024) (Phoenix500526) +- IN LIST: clean up generic static filtering [#21927](https://github.com/apache/datafusion/pull/21927) (geoffreyclaude) +- test: drive stdin store reuse through get_or_create [#23061](https://github.com/apache/datafusion/pull/23061) (huan233usc) +- test: Move default cache tests to default cache file [#23040](https://github.com/apache/datafusion/pull/23040) (mkleen) +- Optimize Parquet row-filter struct schema pruning [#22960](https://github.com/apache/datafusion/pull/22960) (shehab-ali) +- Fix DuckDB unparse for optimized join projections [#23002](https://github.com/apache/datafusion/pull/23002) (goutamadwant) +- chore: gate `internal_datafusion_err` import behind the `proto` feature [#23075](https://github.com/apache/datafusion/pull/23075) (Phoenix500526) +- Perf: avoid redundant comparison in SortPreservingMerge round-robin tie-breaker; optimize inner loop [#23107](https://github.com/apache/datafusion/pull/23107) (Dandandan) +- Move Parquet `input_file_name()` tests to `input_file_name.slt` [#23123](https://github.com/apache/datafusion/pull/23123) (AdamGS) +- chore(deps): bump actions/checkout from 6.0.3 to 7.0.0 [#23115](https://github.com/apache/datafusion/pull/23115) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.81.11 to 2.82.2 [#23114](https://github.com/apache/datafusion/pull/23114) (dependabot[bot]) +- [sql]: remove deprecated TableReference re-exports [#23102](https://github.com/apache/datafusion/pull/23102) (mgkz0) +- chore: `cargo update -p quinn` to resolve security audit issue [#23122](https://github.com/apache/datafusion/pull/23122) (Jefffrey) +- refactor(hash-aggr): Use `EmitTo` to output [#23055](https://github.com/apache/datafusion/pull/23055) (2010YOUY01) +- refactor: centralize TopK heap boundary handling [#23091](https://github.com/apache/datafusion/pull/23091) (kumarUjjawal) +- IN LIST: add UInt8 bitmap filter [#23011](https://github.com/apache/datafusion/pull/23011) (geoffreyclaude) +- chore(physical-plan): remove deprecated RowIndex struct (Closes #23080 - partial) [#23143](https://github.com/apache/datafusion/pull/23143) (Dodothereal) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 5 updates [#23118](https://github.com/apache/datafusion/pull/23118) (dependabot[bot]) +- Fix projection functional dependency remapping [#23028](https://github.com/apache/datafusion/pull/23028) (hhhizzz) +- Migrate case conversion and substr_index to fallible string view builder APIs [#23074](https://github.com/apache/datafusion/pull/23074) (kosiew) +- chore: use `Vec` instead of `OffsetBuilder` [#23195](https://github.com/apache/datafusion/pull/23195) (comphead) +- Fix final hash aggregate output regression by materializing once [#23182](https://github.com/apache/datafusion/pull/23182) (hhhizzz) +- Add regression coverage for quoted dotted column aliases [#23155](https://github.com/apache/datafusion/pull/23155) (kosiew) +- feat(functions-aggregate): support sum(interval) [#23177](https://github.com/apache/datafusion/pull/23177) (SubhamSinghal) +- chore(deps): bump itertools from 0.14.0 to 0.15.0 [#23119](https://github.com/apache/datafusion/pull/23119) (dependabot[bot]) +- refactor: centralize join-input table-scan filter extraction before u… [#23166](https://github.com/apache/datafusion/pull/23166) (Phoenix500526) +- refactor: factor distinct-from unparsing into a shared helper [#23163](https://github.com/apache/datafusion/pull/23163) (Phoenix500526) +- IN LIST: unify bitmap filter implementations [#23035](https://github.com/apache/datafusion/pull/23035) (geoffreyclaude) +- Avoid repeated `EmitTo::First` in partial hash aggregate output [#23250](https://github.com/apache/datafusion/pull/23250) (hhhizzz) +- Fix metrics for repartition when `preserve_order=true` [#20924](https://github.com/apache/datafusion/pull/20924) (xanderbailey) +- chore(deps): bump taiki-e/install-action from 2.82.2 to 2.82.6 [#23254](https://github.com/apache/datafusion/pull/23254) (dependabot[bot]) +- chore(deps): bump the all-other-cargo-deps group with 5 updates [#23256](https://github.com/apache/datafusion/pull/23256) (dependabot[bot]) +- refactor: `make_map_batch` array handling [#23228](https://github.com/apache/datafusion/pull/23228) (nathanb9) +- bench(hj): Add missing Q16–Q23 to benchmarks [#23257](https://github.com/apache/datafusion/pull/23257) (LiaCastaneda) +- Restrict trigger push branch for GitHub Workflow [#23278](https://github.com/apache/datafusion/pull/23278) (apupier) +- Aggregations Support `Partitioning::Range` [#23239](https://github.com/apache/datafusion/pull/23239) (gene-bordegaray) +- refactor(hash-aggr): Migrate ordered partial/final aggregation [#23181](https://github.com/apache/datafusion/pull/23181) (2010YOUY01) +- Fix CI failure by Ignore quick-xml audit advisories [#23298](https://github.com/apache/datafusion/pull/23298) (alamb) +- refactor(hash-aggr): Migrate partial-reduce hash aggregation [#23233](https://github.com/apache/datafusion/pull/23233) (2010YOUY01) +- fix(`EnsureRequirements`): remap sort requirement through `ProjectionExec` on pushdown [#23199](https://github.com/apache/datafusion/pull/23199) (Jeadie) +- chore(deps): bump cmov from 0.5.3 to 0.5.4 [#23300](https://github.com/apache/datafusion/pull/23300) (dependabot[bot]) +- Minor: Make `BloomFilterStatistics` and `RowGroupAccessPlanFilter::prune_by_bloom_filters` public [#23302](https://github.com/apache/datafusion/pull/23302) (xudong963) +- Add basic sql benchmark runner for running sql benchmarks [#23052](https://github.com/apache/datafusion/pull/23052) (Omega359) +- spark: support `collect_list` `collect_set` for `windows` execution [#23281](https://github.com/apache/datafusion/pull/23281) (comphead) +- chore: extend pre commit instructions for AI agents [#23313](https://github.com/apache/datafusion/pull/23313) (comphead) +- chore: add Cargo http options to handle download errors [#23314](https://github.com/apache/datafusion/pull/23314) (comphead) +- Fix inexact partitioned TopK sort pushdown [#23301](https://github.com/apache/datafusion/pull/23301) (xudong963) +- Add IN list sqllogictest test (and integer type coverage) [#23305](https://github.com/apache/datafusion/pull/23305) (alamb) +- bench: add array_has array-needle benchmarks [#23335](https://github.com/apache/datafusion/pull/23335) (freakyzoidberg) +- Add regression tests for hash-join dynamic filter expression policy [#23319](https://github.com/apache/datafusion/pull/23319) (kosiew) +- chore: update crossbeam-epoch to 0.9.20 [#23358](https://github.com/apache/datafusion/pull/23358) (Phoenix500526) +- chore(docs): resolve some docs typos [#23347](https://github.com/apache/datafusion/pull/23347) (devanbenz) +- IN LIST: add Float16 bitmap filter [#23311](https://github.com/apache/datafusion/pull/23311) (geoffreyclaude) +- chore(deps): bump astral-sh/setup-uv from 8.2.0 to 8.3.1 [#23366](https://github.com/apache/datafusion/pull/23366) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.82.6 to 2.82.10 [#23365](https://github.com/apache/datafusion/pull/23365) (dependabot[bot]) +- chore(deps): bump actions/checkout from 6.0.3 to 7.0.0 [#23367](https://github.com/apache/datafusion/pull/23367) (dependabot[bot]) +- minor: rename aggregate stream modules to match contents [#23372](https://github.com/apache/datafusion/pull/23372) (alamb) +- test: cover float IN list predicates [#23373](https://github.com/apache/datafusion/pull/23373) (alamb) +- test: Add coverage for `NOT IN` predicates [#23378](https://github.com/apache/datafusion/pull/23378) (alamb) +- chore(deps): bump runs-on/action from 2.1.2 to 2.2.0 [#23363](https://github.com/apache/datafusion/pull/23363) (dependabot[bot]) +- chore: Update to arrow/parquet 59.1.0 [#23312](https://github.com/apache/datafusion/pull/23312) (alamb) +- Fix:22477 any all schema error [#22915](https://github.com/apache/datafusion/pull/22915) (HairstonE) +- Push sort requirements through simple projections [#23288](https://github.com/apache/datafusion/pull/23288) (aectaan) +- refactor(hash-aggr): Simplify aggregate hash table with tempated functions [#23324](https://github.com/apache/datafusion/pull/23324) (2010YOUY01) +- refactor: centralizing shared-allocation accounting for Arc DFHeapSize impls [#23349](https://github.com/apache/datafusion/pull/23349) (saadtajwar) +- Fix union equivalence schema rewrite with stale constants [#23375](https://github.com/apache/datafusion/pull/23375) (xudong963) +- Fix memory size accounting for grouped `median` and `avg` [#23357](https://github.com/apache/datafusion/pull/23357) (lyne7-sc) +- chore(spm): extract initialize all parititions helper [#23419](https://github.com/apache/datafusion/pull/23419) (rluvaton) +- refactor: extract parquet projection read plan into its own module [#23396](https://github.com/apache/datafusion/pull/23396) (adriangb) +- Perf: Add short circuit for primitive vectorized equal_to [#23343](https://github.com/apache/datafusion/pull/23343) (Rich-T-kid) +- refactor: centralize date_bin per-row mapping [#23034](https://github.com/apache/datafusion/pull/23034) (kumarUjjawal) +- chore: cleanup some TODO items in sqllogictests [#23382](https://github.com/apache/datafusion/pull/23382) (Jefffrey) +- bench: add date_part benchmark [#23350](https://github.com/apache/datafusion/pull/23350) (theirix) +- refactor: de-duplicate parquet read plan construction [#23426](https://github.com/apache/datafusion/pull/23426) (adriangb) +- chore(deps): bump soupsieve from 2.8.3 to 2.8.4 [#23432](https://github.com/apache/datafusion/pull/23432) (dependabot[bot]) +- Use `concat_elements_dyn` from `arrow-rs` [#23211](https://github.com/apache/datafusion/pull/23211) (pepijnve) +- perf(physical-plan): fold PlanProperties fast-path into with_new_children_if_necessary (PR 1 of #22555) [#23332](https://github.com/apache/datafusion/pull/23332) (zhuqi-lucas) +- Minor: Fix docs for JoinSet [#23448](https://github.com/apache/datafusion/pull/23448) (alamb) +- test: add Poll::Pending spill stream coverage for async spill re-entry paths [#23353](https://github.com/apache/datafusion/pull/23353) (pantShrey) +- Test: add more aggregation focused dictionary sql logic test [#23280](https://github.com/apache/datafusion/pull/23280) (Rich-T-kid) +- minor: Remove `.gitignore` item for datafusion-examples [#23409](https://github.com/apache/datafusion/pull/23409) (2010YOUY01) +- minor: remove local file commited by mistake [#23476](https://github.com/apache/datafusion/pull/23476) (2010YOUY01) +- bench: add sort benchmarks for various data profile [#23346](https://github.com/apache/datafusion/pull/23346) (rluvaton) +- refactor(hash-aggr): Migrate single mode hash aggregation [#23408](https://github.com/apache/datafusion/pull/23408) (2010YOUY01) +- test: add sqllogictest coverage for DISTINCT / GROUP BY / aggregation on map columns [#23406](https://github.com/apache/datafusion/pull/23406) (PG1204) +- chore: use new `OffsetBuffer::subtract` helper [#23424](https://github.com/apache/datafusion/pull/23424) (rluvaton) +- Decode Hive partition values in listing tables [#23226](https://github.com/apache/datafusion/pull/23226) (yinli-systems) +- ci: Use `install-action` instead of `cargo install` to speed up CI [#23477](https://github.com/apache/datafusion/pull/23477) (2010YOUY01) +- Add minimal genarator-like stream implementation [#23530](https://github.com/apache/datafusion/pull/23530) (pepijnve) +- chore(deps): bump actions/stale from 10.3.0 to 10.4.0 [#23557](https://github.com/apache/datafusion/pull/23557) (dependabot[bot]) +- chore(deps): bump actions/labeler from 6.1.0 to 6.2.0 [#23556](https://github.com/apache/datafusion/pull/23556) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.82.10 to 2.83.2 [#23555](https://github.com/apache/datafusion/pull/23555) (dependabot[bot]) +- chore(deps): bump astral-sh/setup-uv from 8.3.1 to 8.3.2 [#23554](https://github.com/apache/datafusion/pull/23554) (dependabot[bot]) +- chore(deps): bump actions/setup-node from 6 to 7 [#23550](https://github.com/apache/datafusion/pull/23550) (dependabot[bot]) +- perf(physical-expr): cache remapped expression in DynamicFilterPhysicalExpr::current() [#23532](https://github.com/apache/datafusion/pull/23532) (zhuqi-lucas) +- Preserve string slice function return types [#23330](https://github.com/apache/datafusion/pull/23330) (xudong963) +- Allow Range partitioned inputs to PartitionedTopK [#23355](https://github.com/apache/datafusion/pull/23355) (stuhood) +- chore: Simplifying `SortPreservingMergeStream` to use generators instead of state machine [#23407](https://github.com/apache/datafusion/pull/23407) (rluvaton) +- Enforce co-partitioning for sort merge and symmetric hash joins [#23480](https://github.com/apache/datafusion/pull/23480) (gene-bordegaray) +- Infer placeholder type from ANY/ALL subquery, unit tests [#22545](https://github.com/apache/datafusion/pull/22545) (HairstonE) +- Use `octet_length` for ClickBench Q27/Q28 byte-length semantics [#23475](https://github.com/apache/datafusion/pull/23475) (kosiew) +- chore: group codeql action dependabot updates [#23561](https://github.com/apache/datafusion/pull/23561) (Jefffrey) +- chore(deps): bump the codeql-actions group with 2 updates [#23610](https://github.com/apache/datafusion/pull/23610) (dependabot[bot]) +- minor: validate config `recursion_limit` when setting it [#23592](https://github.com/apache/datafusion/pull/23592) (2010YOUY01) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 10 updates [#23613](https://github.com/apache/datafusion/pull/23613) (dependabot[bot]) +- Fix within group aggregates with unparser [#22195](https://github.com/apache/datafusion/pull/22195) (cetra3) +- bench(sort): fix sort axis benchmark run on single partition [#23614](https://github.com/apache/datafusion/pull/23614) (rluvaton) +- minor: validate config `max_spill_file_size_bytes` when setting it [#23594](https://github.com/apache/datafusion/pull/23594) (2010YOUY01) +- minor: validate config `soft_max_rows_per_output_file` when setting it [#23597](https://github.com/apache/datafusion/pull/23597) (2010YOUY01) +- chore(deps-dev): bump websocket-driver from 0.7.4 to 0.7.5 in /datafusion/wasmtest/datafusion-wasm-app [#23625](https://github.com/apache/datafusion/pull/23625) (dependabot[bot]) +- chore(deps): bump serde_with from 3.18.0 to 3.21.0 [#23624](https://github.com/apache/datafusion/pull/23624) (dependabot[bot]) +- minor: validate config `minimum_parallel_output_files` when setting it [#23596](https://github.com/apache/datafusion/pull/23596) (2010YOUY01) +- minor: validate config `meta_fetch_concurrency` when setting it [#23595](https://github.com/apache/datafusion/pull/23595) (2010YOUY01) +- try parallel ci [#23618](https://github.com/apache/datafusion/pull/23618) (blaginin) +- allow interleaveExec to support Range partioning [#23623](https://github.com/apache/datafusion/pull/23623) (Rich-T-kid) +- Support UDTFs in information_schema.routines / SHOW FUNCTIONS [#23438](https://github.com/apache/datafusion/pull/23438) (zhuqi-lucas) +- Handle nulls in type coercion of higher-order UDFs, map_extract, spark array_repeat [#23071](https://github.com/apache/datafusion/pull/23071) (gstvg) +- minor(CI): Use `install-action` to speed up ci [#23661](https://github.com/apache/datafusion/pull/23661) (2010YOUY01) +- bench: add FixedSizeBinary coverage to multi_group_by benchmark [#23650](https://github.com/apache/datafusion/pull/23650) (alamb) +- test: Fix malformed `regexp_instr` error tests and add slt coverage [#23620](https://github.com/apache/datafusion/pull/23620) (alamb) +- Mark null-propagating math functions as strict [#23527](https://github.com/apache/datafusion/pull/23527) (lyne7-sc) +- Fix ordering for UNION ALL over heterogeneous constants [#23528](https://github.com/apache/datafusion/pull/23528) (vadimpiven) +- chore: downsize `sql_planner_extended` `logical_plan_optimize` sample size to 5 [#23659](https://github.com/apache/datafusion/pull/23659) (Jefffrey) +- test: add advanced dictionary test [#23483](https://github.com/apache/datafusion/pull/23483) (Rich-T-kid) +- bench: parquet scan with a table schema narrower than a nested column [#23397](https://github.com/apache/datafusion/pull/23397) (adriangb) +- Cap SortPreservingMerge statistics by fetch [#23359](https://github.com/apache/datafusion/pull/23359) (discord9) +- `array_agg()` add tests and benchmarks [#23740](https://github.com/apache/datafusion/pull/23740) (fred1268) +- test: More `slt` tests for `iszero` function [#23713](https://github.com/apache/datafusion/pull/23713) (2010YOUY01) +- feat(physical-plan): Allow co-partitioned Partitioning::Range inputs for left-side hash joins [#23487](https://github.com/apache/datafusion/pull/23487) (JSOD11) +- chore(deps-dev): bump webpack-dev-server from 5.2.5 to 5.2.6 in /datafusion/wasmtest/datafusion-wasm-app [#23768](https://github.com/apache/datafusion/pull/23768) (dependabot[bot]) +- chore(deps): bump actions/checkout from 7.0.0 to 7.0.1 [#23746](https://github.com/apache/datafusion/pull/23746) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.82.6 to 2.84.0 [#23748](https://github.com/apache/datafusion/pull/23748) (dependabot[bot]) +- chore(deps): bump actions/labeler from 6.2.0 to 7.0.0 [#23749](https://github.com/apache/datafusion/pull/23749) (dependabot[bot]) +- chore(deps): bump codecov/codecov-action from 5.5.5 to 7.0.0 [#23750](https://github.com/apache/datafusion/pull/23750) (dependabot[bot]) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 15 updates [#23771](https://github.com/apache/datafusion/pull/23771) (dependabot[bot]) +- chore: fix `SlidingDistinctCountAccumulator::size()` to include budget for distinct values [#23399](https://github.com/apache/datafusion/pull/23399) (comphead) +- test: improve `md5` function SQL test coverage [#23757](https://github.com/apache/datafusion/pull/23757) (2010YOUY01) +- chore(deps-dev): bump fast-uri from 3.1.2 to 3.1.4 in /datafusion/wasmtest/datafusion-wasm-app [#23778](https://github.com/apache/datafusion/pull/23778) (dependabot[bot]) +- chore(deps-dev): bump shell-quote from 1.8.4 to 1.10.0 in /datafusion/wasmtest/datafusion-wasm-app [#23769](https://github.com/apache/datafusion/pull/23769) (dependabot[bot]) +- test: improve `isnan` function SQL test coverage [#23754](https://github.com/apache/datafusion/pull/23754) (2010YOUY01) +- chore(deps): bump the codeql-actions group with 2 updates [#23745](https://github.com/apache/datafusion/pull/23745) (dependabot[bot]) +- test: improve `digest` function SQL test coverage [#23756](https://github.com/apache/datafusion/pull/23756) (2010YOUY01) +- test: improve `sha` function SQL test coverage [#23758](https://github.com/apache/datafusion/pull/23758) (2010YOUY01) +- test: improve `lcm` function SQL test coverage [#23755](https://github.com/apache/datafusion/pull/23755) (2010YOUY01) +- allow range to satisfy key distribution generally [#23680](https://github.com/apache/datafusion/pull/23680) (gene-bordegaray) +- fix: do not treat concat as preserving lexicographical ordering [#23804](https://github.com/apache/datafusion/pull/23804) (buraksenn) +- refactor(proto): delegate deprecated ProjectionExec serde shims to new hooks [#23731](https://github.com/apache/datafusion/pull/23731) (adriangb) +- Add setter for `TaskContext::task_id` [#23837](https://github.com/apache/datafusion/pull/23837) (pepijnve) +- refactor(hash-aggr): Support spilling for ordered aggregation [#23657](https://github.com/apache/datafusion/pull/23657) (2010YOUY01) +- Unwrap widening Date32 -> Date64 casts in comparison predicates [#23729](https://github.com/apache/datafusion/pull/23729) (adriangb) +- test (slt): add memory-limited aggregation sqllogictests [#23838](https://github.com/apache/datafusion/pull/23838) (naman-modi) +- chore(deps-dev): bump ws from 8.18.2 to 8.21.1 in /datafusion/wasmtest/datafusion-wasm-app [#23866](https://github.com/apache/datafusion/pull/23866) (dependabot[bot]) +- chore(deps-dev): bump http-proxy-middleware from 2.0.9 to 2.0.10 in /datafusion/wasmtest/datafusion-wasm-app [#23865](https://github.com/apache/datafusion/pull/23865) (dependabot[bot]) +- test: add functional_dependencies.slt covering functional dependency driven optimizations [#23821](https://github.com/apache/datafusion/pull/23821) (alamb) +- chore(deps-dev): bump webpack-dev-server from 5.2.6 to 6.0.0 in /datafusion/wasmtest/datafusion-wasm-app [#23868](https://github.com/apache/datafusion/pull/23868) (dependabot[bot]) +- refactor(unparser): centralize aggregate-scope rendering in the SQL unparser [#23789](https://github.com/apache/datafusion/pull/23789) (naman-modi) +- Add FixedSizeList support for recursive struct schema adaptation [#22980](https://github.com/apache/datafusion/pull/22980) (kosiew) +- test: cover `array_agg(DISTINCT)` on dictionaries and bounded `retract_batch` memory [#23873](https://github.com/apache/datafusion/pull/23873) (alamb) +- chore: adjust `size` accounting for `min_max` [#23899](https://github.com/apache/datafusion/pull/23899) (comphead) +- Add ObjectStore-backed TempFileFactor / spill example [#23170](https://github.com/apache/datafusion/pull/23170) (alamb) +- Various `ScalarValue` numeric method fixes & refactors (especially decimal) [#23631](https://github.com/apache/datafusion/pull/23631) (Jefffrey) +- chore: simplify SortPreservingMergeStream to be as textbook-like as possible [#23702](https://github.com/apache/datafusion/pull/23702) (rluvaton) +- chore: Squelch "unused code" warning [#23924](https://github.com/apache/datafusion/pull/23924) (neilconway) +- Add name filter to metrics [#23719](https://github.com/apache/datafusion/pull/23719) (gabotechs) +- chore(deps): bump the codeql-actions group with 2 updates [#23938](https://github.com/apache/datafusion/pull/23938) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.84.0 to 2.85.2 [#23941](https://github.com/apache/datafusion/pull/23941) (dependabot[bot]) +- chore(deps): bump actions/stale from 10.4.0 to 11.0.0 [#23942](https://github.com/apache/datafusion/pull/23942) (dependabot[bot]) +- chore(deps): bump base64 from 0.22.1 to 0.23.0 [#23944](https://github.com/apache/datafusion/pull/23944) (dependabot[bot]) +- chore(deps): bump astral-sh/setup-uv from 8.3.2 to 9.0.0 [#23939](https://github.com/apache/datafusion/pull/23939) (dependabot[bot]) +- chore: refactor SortMergeJoin bitwise stream to generators and simplify to be textbook like as possible [#23761](https://github.com/apache/datafusion/pull/23761) (rluvaton) +- refactor: address review feedback on percentile_cont(DISTINCT) accumulator [#23946](https://github.com/apache/datafusion/pull/23946) (viirya) +- chore: remove unused `header` file [#23958](https://github.com/apache/datafusion/pull/23958) (Jefffrey) +- Fill in missing utf8view support in function type coercion [#23916](https://github.com/apache/datafusion/pull/23916) (Jefffrey) +- chore: refactor `VarianceAccumulator`, add tests and benchmark [#23977](https://github.com/apache/datafusion/pull/23977) (neilconway) +- minor(test): cover partially ordered aggregate spilling [#23947](https://github.com/apache/datafusion/pull/23947) (buraksenn) +- chore(ordered-partial-aggregate): move `OrderedPartialAggregateStream` to generators for readability [#23951](https://github.com/apache/datafusion/pull/23951) (rluvaton) +- test: improve `round` sqllogictest coverage [#23973](https://github.com/apache/datafusion/pull/23973) (2010YOUY01) +- test: improve `gcd` sqllogictest coverage [#23972](https://github.com/apache/datafusion/pull/23972) (2010YOUY01) +- test: improve `rpad` sqllogictest coverage [#23968](https://github.com/apache/datafusion/pull/23968) (2010YOUY01) +- test: improve `lpad` sqllogictest coverage [#23969](https://github.com/apache/datafusion/pull/23969) (2010YOUY01) +- Add benchmarks for hashjoin candidate equality filtering [#23980](https://github.com/apache/datafusion/pull/23980) (shehab-ali) +- Optimize Spark hex null handling [#23688](https://github.com/apache/datafusion/pull/23688) (floze-the-genius) +- fix(proto): prevent duplicate partition statistics on roundtrip [#23999](https://github.com/apache/datafusion/pull/23999) (buraksenn) +- Report peak MemoryPool reservation per query in benchmarks [#23985](https://github.com/apache/datafusion/pull/23985) (adriangb) +- refactor(proto): move PartitionedFile / FileGroup serde into datafusion-datasource [#24006](https://github.com/apache/datafusion/pull/24006) (adriangb) +- chore: refactor `MaterializingSortMergeJoinStream` into generators and simplify code to be textbook like as possible [#23976](https://github.com/apache/datafusion/pull/23976) (rluvaton) +- refactor(proto): put Partitioning / sort-expression serde on the types [#24003](https://github.com/apache/datafusion/pull/24003) (adriangb) +- bench: use seedable rng for reproducibility [#23653](https://github.com/apache/datafusion/pull/23653) (theirix) +- test: Fix data_pagesize_limit extraction in parquet writer props roundtrip test [#23664](https://github.com/apache/datafusion/pull/23664) (jackylee-ch) +- bench: extend BoundedWindowAggExec many-partitions benchmark [#24032](https://github.com/apache/datafusion/pull/24032) (neilconway) +- refactor: move arrow integer hex dispatch to datafusion-common [#23917](https://github.com/apache/datafusion/pull/23917) (buraksenn) +- test: add IN list slt coverage for temporal, Decimal128 and Interval types [#23875](https://github.com/apache/datafusion/pull/23875) (alamb) +- chore: rows_to_array cleanup for expecting single field [#24040](https://github.com/apache/datafusion/pull/24040) (saadtajwar) +- minor: Add `slt` test for nullable window retract [#24025](https://github.com/apache/datafusion/pull/24025) (2010YOUY01) +- WindowTopN dense_rank benchmark [#24050](https://github.com/apache/datafusion/pull/24050) (SubhamSinghal) +- minor(fix): correct to_date results for formatted pre-epoch datetimes [#24049](https://github.com/apache/datafusion/pull/24049) (buraksenn) +- minor(test): strengthen sort-merge join spilling coverage [#23988](https://github.com/apache/datafusion/pull/23988) (buraksenn) +- refactor(hash-aggr): Support spilling for single mode aggregation [#23965](https://github.com/apache/datafusion/pull/23965) (2010YOUY01) +- test: improve `find_in_set` sqllogictest coverage [#23970](https://github.com/apache/datafusion/pull/23970) (2010YOUY01) +- IN LIST: isolate branchless filter implementation [#23907](https://github.com/apache/datafusion/pull/23907) (geoffreyclaude) +- bench: add nested-type (List/Struct/Map) cases to first_value/last_value benchmark [#24075](https://github.com/apache/datafusion/pull/24075) (zhuqi-lucas) +- chore(deps): bump taiki-e/install-action from 2.85.2 to 2.85.6 [#24081](https://github.com/apache/datafusion/pull/24081) (dependabot[bot]) +- chore(deps): bump the codeql-actions group with 2 updates [#24080](https://github.com/apache/datafusion/pull/24080) (dependabot[bot]) +- bench: add ArrowBytesMap benchmarks [#24078](https://github.com/apache/datafusion/pull/24078) (Punisheroot) +- chore(deps): bump cryptography from 48.0.1 to 50.0.0 [#24091](https://github.com/apache/datafusion/pull/24091) (dependabot[bot]) +- fix(physical-plan): preserve Exact(0) in FilterExec for null_count, distinct_count and total_byte_size upon empty input [#24000](https://github.com/apache/datafusion/pull/24000) (asolimando) +- chore: cleanup `OrderedPartialAggregateStream` more [#24012](https://github.com/apache/datafusion/pull/24012) (rluvaton) +- chore: apply workspace lints to all crates [#24076](https://github.com/apache/datafusion/pull/24076) (emilk) +- feat(functions-aggregate): support nested types (List, Struct, Map) in first_value / last_value GroupsAccumulator [#23628](https://github.com/apache/datafusion/pull/23628) (zhuqi-lucas) +- Add config-matrix tests in enforce_distribution.rs for range-satisfaction settings [#23627](https://github.com/apache/datafusion/pull/23627) (blinding-pixels) +- chore(deps-dev): bump fast-uri from 3.1.4 to 3.1.5 in /datafusion/wasmtest/datafusion-wasm-app [#24092](https://github.com/apache/datafusion/pull/24092) (dependabot[bot]) +- Add support for running sql benchmarks with command line arguments [#23772](https://github.com/apache/datafusion/pull/23772) (Omega359) +- refactor(hash-aggr): Support spilling for `partial` and `final` mode aggregation [#24061](https://github.com/apache/datafusion/pull/24061) (2010YOUY01) +- Proto: add DataSink serialization hook [#23752](https://github.com/apache/datafusion/pull/23752) (Phoenix500526) +- bench: multi-conjunct shared-prefix struct row-filter pushdown [#23524](https://github.com/apache/datafusion/pull/23524) (SubhamSinghal) +- Preserve grouping ID during aggregate CSE [#24144](https://github.com/apache/datafusion/pull/24144) (notfilippo) +- Add DataSource/FileSource proto hooks and FileScanConfig serde [#23683](https://github.com/apache/datafusion/pull/23683) (kumarUjjawal) +- tests: add SLT test coverage for `MERGE INTO` [#24174](https://github.com/apache/datafusion/pull/24174) (alamb) +- chore: add runendencoded & listview types to dfschema equality methods [#24138](https://github.com/apache/datafusion/pull/24138) (Jefffrey) +- refactor(proto): destructure plan and proto structs in aggregate and window serde hooks [#24166](https://github.com/apache/datafusion/pull/24166) (adriangb) +- chore(deps): bump the all-other-cargo-deps group across 1 directory with 10 updates [#24163](https://github.com/apache/datafusion/pull/24163) (dependabot[bot]) +- test: add UnionArray hashing SQL coverage. [#24199](https://github.com/apache/datafusion/pull/24199) (VaibhaveS) +- fix(physical-plan): count empty grouping sets in the aggregate row estimate for an empty input [#24039](https://github.com/apache/datafusion/pull/24039) (asolimando) +- test(proto): add missing physical plan round-trip coverage [#24172](https://github.com/apache/datafusion/pull/24172) (adriangb) +- test(proto): split roundtrip_physical_plan.rs by plan category [#24223](https://github.com/apache/datafusion/pull/24223) (adriangb) +- physical-plan: coerce UNION/INTERLEAVE schema mismatches at plan time [#24094](https://github.com/apache/datafusion/pull/24094) (dariocurr) +- Parquet row filter struct access tree [#23217](https://github.com/apache/datafusion/pull/23217) (SubhamSinghal) +- Fix aggregate accumulator capacity accounting [#24099](https://github.com/apache/datafusion/pull/24099) (kosiew) +- refactor: Refactor numeric sign and padding in Spark format_string [#24115](https://github.com/apache/datafusion/pull/24115) (JSOD11) +- chore(deps): bump the all-other-cargo-deps group with 4 updates [#24254](https://github.com/apache/datafusion/pull/24254) (dependabot[bot]) +- chore(deps): bump taiki-e/install-action from 2.85.6 to 2.85.10 [#24253](https://github.com/apache/datafusion/pull/24253) (dependabot[bot]) +- chore(deps): bump runs-on/action from 2.2.0 to 2.3.0 [#24252](https://github.com/apache/datafusion/pull/24252) (dependabot[bot]) +- chore(deps): bump Swatinem/rust-cache from 2.9.1 to 2.9.2 [#24251](https://github.com/apache/datafusion/pull/24251) (dependabot[bot]) +- Add FixedSizeBinary support for MultiGroupBy [#23646](https://github.com/apache/datafusion/pull/23646) (maxburke) +- refactor: make apply_expression_roots more ergonomic [#24226](https://github.com/apache/datafusion/pull/24226) (jayshrivastava) +- chore(deps): bump toml from 0.9.12+spec-1.1.0 to 1.1.3+spec-1.1.0 [#24256](https://github.com/apache/datafusion/pull/24256) (dependabot[bot]) +- refactor: moving WindowTopN before EnsureRequirements [#24191](https://github.com/apache/datafusion/pull/24191) (saadtajwar) +- chore(deps): bump the codeql-actions group with 2 updates [#24250](https://github.com/apache/datafusion/pull/24250) (dependabot[bot]) +- [branch-55] Prepare for 55 release - version number, changelog [#24292](https://github.com/apache/datafusion/pull/24292) (timsaucer) +- [branch-55] Update changelog [#24314](https://github.com/apache/datafusion/pull/24314) (timsaucer) +- [branch-55] fix: correct list field inner type in array functions (#24345) [#24367](https://github.com/apache/datafusion/pull/24367) (timsaucer) +- [branch-55] fix wrong TopK results from re-reading already-delivered row groups (#24352) [#24368](https://github.com/apache/datafusion/pull/24368) (zhuqi-lucas) +- [branch-55]: don't runtime-prune row groups while a page-index RowSelection is live (#24355) [#24374](https://github.com/apache/datafusion/pull/24374) (zhuqi-lucas) +- [branch-55] fix: preserve the input list's inner field in array_append/prepend/replace - #24365 [#24377](https://github.com/apache/datafusion/pull/24377) (timsaucer) + +## Credits + +Thank you to everyone who contributed to this release. Here is a breakdown of commits (PRs merged) per contributor. + +``` + 80 dependabot[bot] + 46 Neil Conway + 42 Yongting You + 41 Adrian Garcia Badaracco + 37 Andrew Lamb + 27 Burak Şen + 24 Phoenix + 24 linfeng + 22 Andy Grove + 19 Kumar Ujjawal + 19 Michael Kleen + 18 Raz Luvaton + 17 Qi Zhu + 16 Adam Gutglick + 14 Bruce Ritchie + 14 Giorgio Maria Federico Birnthaler + 13 Jeffrey Vo + 12 Geoffrey Claude + 12 Oleks V + 12 xudong.w + 10 Gene Bordegaray + 10 Saad Tajwar + 10 Subham Singhal + 10 Tim Saucer + 10 kosiew + 9 Nathan + 9 theirix + 8 Huaijin + 6 Daniël Heres + 6 Zhen Chen + 5 Alessandro Solimando + 5 Amogh Ramesh + 5 Ariel Miculas-Trif + 5 Jayant Shrivastava + 5 Liang-Chi Hsieh + 5 Lía Adriana + 5 Matt Butrovich + 5 RIchard Baah + 5 kid + 4 Brent Gardner + 4 Goutam Adwant + 4 H + 4 Megakaizo + 4 Nuno Faria + 4 Pepijn Van Eeckhoudt + 4 Sean Kenneth Doherty + 4 Varun + 4 Xuanyi Li + 4 chakkk309 + 4 discord9 + 3 Alex Metelli + 3 ByteBaker + 3 Huang Qiwei + 3 Justin O'Dwyer + 3 Matthew Patton + 3 Mithun Chicklore Yogendra + 3 Moe + 3 Shehab Ali + 3 Simon Vandel Sillesen + 3 Xin Huang + 3 Yin Li + 3 crm26 + 3 gstvg + 3 pantShrey + 3 pchintar + 2 Anurag Tryambak Raut + 2 Bert Vermeiren + 2 Bhargava Vadlamani + 2 David López + 2 Diego Perez Giordán + 2 Edson Petry + 2 EeshanBembi + 2 Emily Matheys + 2 Filip Petkovski + 2 Florian Müller + 2 Ford + 2 Fred Thomas + 2 Gabriel + 2 Guocheng(Eric) Song + 2 JS + 2 Kanishk Sachan + 2 Karpagam Balasubramaniam + 2 Krishna Sudarshan J + 2 Louis Vialar + 2 Matthew Kim + 2 Nagato Yuzuru + 2 Naman Modi + 2 Peter L + 2 Peter Lee + 2 Pierre Lacave + 2 Prateek Ganigi + 2 Puneet Dixit + 2 Tobias Schwarzinger + 2 WeblWabl + 2 Zac Farrell + 2 Zeel Rajodiya + 2 dario curreri + 2 fys + 2 jackylee + 2 jj.lee + 2 nanookclaw + 1 7. Sun + 1 Ahmed EL. + 1 Asish Kumar + 1 Aurélien Pupier + 1 Ben Chambers + 1 Braedon Wooding + 1 Brijesh Thakkar + 1 Bruno Volpato + 1 Bukhtawar Khan + 1 Daipayan Mukherjee + 1 DevShiba + 1 Dmitrii Blaginin + 1 Eduardo Aguilar + 1 Egor Markov + 1 Emil Ernerfeldt + 1 Evgeniy Mineev + 1 Filippo + 1 Floze + 1 Georgi Krastev + 1 Gunther Xing + 1 Gustavo Schneiter + 1 Harrison Crosse + 1 Haseeb Nazir + 1 Jack Eadie + 1 Jason Wong + 1 Jordan Epstein + 1 Joseph Lenton + 1 Kazantsev Maksim + 1 Kent Wu + 1 Kristin Cowalcijk + 1 Krisztián Szűcs + 1 Lavkesh Lahngir + 1 Lining Pan + 1 Ma Zhengxuan + 1 Marc Brinkmann + 1 Marko Milenković + 1 Matt Van Horn + 1 Max Burke + 1 Minh Vu + 1 Nam2ee + 1 Namgung Chan + 1 Nathan Bezualem + 1 Pablo Abad Rubio + 1 Pavan51 + 1 Ratul Dawar + 1 Recoordinate + 1 Ruchir Tripathi + 1 RyanStewart + 1 Sai Asish Y + 1 Savan Nahar + 1 Sergei Grebnov + 1 Sergey Zhukov + 1 Stu Hood + 1 Thomas Santerre + 1 Tian Teng + 1 Vadim Piven + 1 VaibhaveS + 1 Victorien + 1 Vikrant Mehta + 1 Vismay + 1 Wenqi Mou + 1 Xander + 1 Xuanwo + 1 Yonatan Striem Amit + 1 Zhen-Lun (Kevin) Hong + 1 ajegou + 1 blinding-pixels + 1 eliot1480 + 1 jeroenflvr + 1 kkrainov + 1 subotac + 1 yoongbok lee + 1 zhengpeng + 1 zhigang +``` + +Thank you also to everyone who contributed in other ways such as filing issues, reviewing PRs, and providing feedback on this release. diff --git a/dev/depcheck/.gitignore b/dev/depcheck/.gitignore deleted file mode 100644 index 03314f77b5aa4..0000000000000 --- a/dev/depcheck/.gitignore +++ /dev/null @@ -1 +0,0 @@ -Cargo.lock diff --git a/dev/depcheck/Cargo.lock b/dev/depcheck/Cargo.lock new file mode 100644 index 0000000000000..3018c79c5a827 --- /dev/null +++ b/dev/depcheck/Cargo.lock @@ -0,0 +1,4167 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "annotate-snippets" +version = "0.12.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f211a51805bc641f3ad5b7664c77d2547af685cc33b4cd8d31964027a46f13f1" +dependencies = [ + "anstyle", + "memchr", + "unicode-width", +] + +[[package]] +name = "anstream" +version = "0.6.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43d5b281e737544384e969a5ccad3f1cdd24b48086a0fc1b2a5262a26b8f4f4a" +dependencies = [ + "anstyle", + "anstyle-parse 0.2.7", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse 1.0.0", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e7644824f0aa2c7b9384579234ef10eb7efb6a0deb83f9630a49594dd9c15c2" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[package]] +name = "arc-swap" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a3a1fd6f75306b68087b831f025c712524bcb19aad54e557b1129cfa0a2b207" +dependencies = [ + "rustversion", +] + +[[package]] +name = "arrayref" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76a2e8124351fda1ef8aaaa3bbd7ebbcb486bbcd4225aca0aa0d84bb2db8fecb" + +[[package]] +name = "arrayvec" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c02d123df017efcdfbd739ef81735b36c5ba83ec3c59c80a9d7ecc718f92e50" + +[[package]] +name = "autocfg" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2032f911046de80f0a198e0901378627c33f59ea0ac00e363d481118bd70a53" + +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "bitmaps" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "031043d04099746d8db04daf1fa424b2bc8bd69d92b25962dcde24da39ab64a2" +dependencies = [ + "typenum", +] + +[[package]] +name = "blake3" +version = "1.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0aa83c34e62843d924f905e0f5c866eb1dd6545fc4d719e803d9ba6030371fce" +dependencies = [ + "arrayref", + "arrayvec", + "cc", + "cfg-if", + "constant_time_eq", + "cpufeatures 0.3.0", +] + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "block2" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdeb9d870516001442e364c5220d3574d2da8dc765554b4a617230d33fa58ef5" +dependencies = [ + "objc2", +] + +[[package]] +name = "bstr" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63044e1ae8e69f3b5a92c736ca6269b8d12fa7efe39bf34ddb06d102cf0e2cab" +dependencies = [ + "memchr", + "regex-automata", + "serde", +] + +[[package]] +name = "bumpalo" +version = "3.20.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72f5acc6cb2ba439de613abc23857ec3d78374d8ed5ac84e9d11336e87da8649" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" + +[[package]] +name = "cargo" +version = "0.92.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89389877f508bae1d45a48b4e76cb0dac5e41a6ac5103752e44d29be0c69c394" +dependencies = [ + "annotate-snippets", + "anstream 0.6.21", + "anstyle", + "anyhow", + "base64", + "blake3", + "cargo-credential", + "cargo-credential-libsecret", + "cargo-credential-macos-keychain", + "cargo-credential-wincred", + "cargo-platform", + "cargo-util", + "cargo-util-schemas", + "clap", + "clap_complete", + "color-print", + "crates-io", + "curl", + "curl-sys", + "filetime", + "flate2", + "git2", + "git2-curl", + "gix", + "glob", + "hex", + "hmac", + "home", + "http-auth", + "ignore", + "im-rc", + "indexmap", + "itertools", + "jiff", + "jobserver", + "lazycell", + "libc", + "libgit2-sys", + "memchr", + "opener", + "os_info", + "pasetors", + "pathdiff", + "rand", + "regex", + "rusqlite", + "rustc-hash", + "rustc-stable-hash", + "rustfix", + "same-file", + "semver", + "serde", + "serde-untagged", + "serde_ignored", + "serde_json", + "sha1", + "shell-escape", + "supports-hyperlinks", + "supports-unicode", + "tar", + "tempfile", + "thiserror", + "time", + "toml", + "toml_edit", + "tracing", + "tracing-chrome", + "tracing-subscriber", + "unicase", + "unicode-width", + "unicode-xid", + "url", + "walkdir", + "windows-sys 0.60.2", + "winnow 0.7.15", +] + +[[package]] +name = "cargo-credential" +version = "0.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e36f089041deadf16226478a7737a833864fbda09408c7af237b9d615eeb6d69" +dependencies = [ + "anyhow", + "libc", + "serde", + "serde_json", + "thiserror", + "time", + "windows-sys 0.60.2", +] + +[[package]] +name = "cargo-credential-libsecret" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90161b8b1b98a28f0fbdfccafb6adcf2b0be948a4fad3acc31461abf5447debe" +dependencies = [ + "anyhow", + "cargo-credential", + "libloading", +] + +[[package]] +name = "cargo-credential-macos-keychain" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e95b9c2431165b30ea111f2933ed6799bfa9a66c9503046064cf8f001960ea1b" +dependencies = [ + "cargo-credential", + "security-framework", +] + +[[package]] +name = "cargo-credential-wincred" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c35397b066a83f2e036fb23fca2fb400bfa65e8e8453c21e0b1690cf8250e414" +dependencies = [ + "cargo-credential", + "windows-sys 0.60.2", +] + +[[package]] +name = "cargo-platform" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87a0c0e6148f11f01f32650a2ea02d532b2ad4e81d8bd41e6e565b5adc5e6082" +dependencies = [ + "serde", + "serde_core", +] + +[[package]] +name = "cargo-util" +version = "0.2.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f97c9ef0f8af69bfcecfe4c17a414d7bb978fe794bc1a38952e27b5c5d87492d" +dependencies = [ + "anyhow", + "core-foundation", + "filetime", + "hex", + "ignore", + "jobserver", + "libc", + "miow", + "same-file", + "sha2", + "shell-escape", + "tempfile", + "tracing", + "walkdir", + "windows-sys 0.60.2", +] + +[[package]] +name = "cargo-util-schemas" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "549c00f5bb23fdaf26135d747d7530563402a101f1887a5a1916afe2c09cf229" +dependencies = [ + "semver", + "serde", + "serde-untagged", + "serde-value", + "thiserror", + "toml", + "unicode-xid", + "url", +] + +[[package]] +name = "cc" +version = "1.2.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dad887fd958be91b5098c0248def011f4523ab786cd411be668777e55063501f" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex 2.0.1", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "clap" +version = "4.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ddb117e43bbf7dacf0a4190fef4d345b9bad68dfc649cb349e7d17d28428e51" +dependencies = [ + "clap_builder", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream 1.0.0", + "anstyle", + "clap_lex", + "strsim", + "terminal_size", +] + +[[package]] +name = "clap_complete" +version = "4.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0a7a9bfdb35811f9e59832f0f05975114d2251b415fb534108e6f34060fd772" +dependencies = [ + "clap", + "clap_lex", + "is_executable", + "shlex 1.3.0", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "clru" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "197fd99cb113a8d5d9b6376f3aa817f32c1078f2343b714fff7d2ca44fdf67d5" +dependencies = [ + "hashbrown 0.16.1", +] + +[[package]] +name = "color-print" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3aa954171903797d5623e047d9ab69d91b493657917bdfb8c2c80ecaf9cdb6f4" +dependencies = [ + "color-print-proc-macro", +] + +[[package]] +name = "color-print-proc-macro" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "692186b5ebe54007e45a59aea47ece9eb4108e141326c304cdc91699a7118a22" +dependencies = [ + "nom", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "const-oid" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + +[[package]] +name = "constant_time_eq" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d52eff69cd5e647efe296129160853a42795992097e8af39800e1060caeea9b" + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "cpufeatures" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b2a41393f66f16b0823bb79094d54ac5fbd34ab292ddafb9a0456ac9f87d201" +dependencies = [ + "libc", +] + +[[package]] +name = "crates-io" +version = "0.40.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "574ce0b8170c097cf174097b84bff181956ad2ab2bbe092ab58d1c08d9f1f417" +dependencies = [ + "curl", + "percent-encoding", + "serde", + "serde_json", + "thiserror", + "url", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "subtle", + "zeroize", +] + +[[package]] +name = "crypto-common" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1bfb12502f3fc46cca1bb51ac28df9d618d813cdc3d2f25b9fe775a34af26bb3" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "ct-codecs" +version = "1.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49fb0c6640b4507ebd99ff67677009e381ba5eee1d14df78de4a3d16eb123c39" + +[[package]] +name = "curl" +version = "0.4.50" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a45ee8994e5307cb4c60cfc1c20bf7263ffb771ddc135c9f768a14bcbc15b09" +dependencies = [ + "curl-sys", + "libc", + "openssl-probe", + "openssl-sys", + "schannel", + "socket2", + "windows-sys 0.61.2", +] + +[[package]] +name = "curl-sys" +version = "0.4.89+curl-8.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d680779285438f2d0927485973ab45b212ea990bddb80de8a55a1e3c1d9ba22" +dependencies = [ + "cc", + "libc", + "libnghttp2-sys", + "libz-sys", + "openssl-sys", + "pkg-config", + "vcpkg", + "windows-sys 0.61.2", +] + +[[package]] +name = "dashmap" +version = "6.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6361d5c062261c78a176addb82d4c821ae42bed6089de0e12603cd25de2059c" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "depcheck" +version = "0.0.0" +dependencies = [ + "cargo", +] + +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "const-oid", + "pem-rfc7468", + "zeroize", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "powerfmt", + "serde_core", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "const-oid", + "crypto-common", + "subtle", +] + +[[package]] +name = "dispatch2" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e0e367e4e7da84520dedcac1901e4da967309406d1e51017ae1abfb97adbd38" +dependencies = [ + "bitflags", + "objc2", +] + +[[package]] +name = "displaydoc" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ac70aa55017e108007fbaf5aa0f54b021c98f92ff8af59d42eda9da96e3dd4f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + +[[package]] +name = "ecdsa" +version = "0.16.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" +dependencies = [ + "der", + "digest", + "elliptic-curve", + "rfc6979", + "signature", + "spki", +] + +[[package]] +name = "ed25519-compact" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f5c0284a5d4b1a2fae017a9fe55fd7d01699711f1b572493f16593e173ea2801" +dependencies = [ + "getrandom 0.4.2", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "elliptic-curve" +version = "0.13.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest", + "ff", + "generic-array", + "group", + "hkdf", + "pem-rfc7468", + "pkcs8", + "rand_core 0.6.4", + "sec1", + "subtle", + "zeroize", +] + +[[package]] +name = "encoding_rs" +version = "0.8.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "erased-serde" +version = "0.4.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2add8a07dd6a8d93ff627029c51de145e12686fbc36ecb298ac22e74cf02dec" +dependencies = [ + "serde", + "serde_core", + "typeid", +] + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "fallible-iterator" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2acce4a10f12dc2fb14a218589d4f1f62ef011b2d0cc4b3cb1bba8e94da14649" + +[[package]] +name = "fallible-streaming-iterator" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7360491ce676a36bf9bb3c56c1aa791658183a54d2744120f27285738d90465a" + +[[package]] +name = "faster-hex" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7223ae2d2f179b803433d9c830478527e92b8117eab39460edae7f1614d9fb73" +dependencies = [ + "heapless", + "serde", +] + +[[package]] +name = "fastrand" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" + +[[package]] +name = "ff" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" +dependencies = [ + "rand_core 0.6.4", + "subtle", +] + +[[package]] +name = "fiat-crypto" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64cd1e32ddd350061ae6edb1b082d7c54915b5c672c389143b9a63403a109f24" + +[[package]] +name = "filetime" +version = "0.2.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" +dependencies = [ + "cfg-if", + "libc", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "miniz_oxide", + "zlib-rs", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "generic-array" +version = "0.14.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bb6743198531e02858aeaea5398fcc883e71851fcbcb5a2f773e2fb6cb1edf2" +dependencies = [ + "typenum", + "version_check", + "zeroize", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "libc", + "r-efi 5.3.0", + "wasip2", +] + +[[package]] +name = "getrandom" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 6.0.0", + "wasip2", + "wasip3", + "wasm-bindgen", +] + +[[package]] +name = "git2" +version = "0.20.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b88256088d75a56f8ecfa070513a775dd9107f6530ef14919dac831af9cfe2b" +dependencies = [ + "bitflags", + "libc", + "libgit2-sys", + "log", + "openssl-probe", + "openssl-sys", + "url", +] + +[[package]] +name = "git2-curl" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be8dcabbc09ece4d30a9aa983d5804203b7e2f8054a171f792deff59b56d31fa" +dependencies = [ + "curl", + "git2", + "log", + "url", +] + +[[package]] +name = "gix" +version = "0.73.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "514c29cc879bdc0286b0cbc205585a49b252809eb86c69df4ce4f855ee75f635" +dependencies = [ + "gix-actor", + "gix-attributes", + "gix-command", + "gix-commitgraph", + "gix-config", + "gix-credentials", + "gix-date", + "gix-diff", + "gix-dir", + "gix-discover", + "gix-features", + "gix-filter", + "gix-fs", + "gix-glob", + "gix-hash", + "gix-hashtable", + "gix-ignore", + "gix-index", + "gix-lock", + "gix-negotiate", + "gix-object", + "gix-odb", + "gix-pack", + "gix-path", + "gix-pathspec", + "gix-prompt", + "gix-protocol", + "gix-ref", + "gix-refspec", + "gix-revision", + "gix-revwalk", + "gix-sec", + "gix-shallow", + "gix-status", + "gix-submodule", + "gix-tempfile", + "gix-trace", + "gix-transport", + "gix-traverse", + "gix-url", + "gix-utils", + "gix-validate", + "gix-worktree", + "once_cell", + "prodash", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-actor" +version = "0.35.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "987a51a7e66db6ef4dc030418eb2a42af6b913a79edd8670766122d8af3ba59e" +dependencies = [ + "bstr", + "gix-date", + "gix-utils", + "itoa", + "thiserror", + "winnow 0.7.15", +] + +[[package]] +name = "gix-attributes" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45442188216d08a5959af195f659cb1f244a50d7d2d0c3873633b1cd7135f638" +dependencies = [ + "bstr", + "gix-glob", + "gix-path", + "gix-quote", + "gix-trace", + "kstring", + "smallvec", + "thiserror", + "unicode-bom", +] + +[[package]] +name = "gix-bitmap" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d982fc7ef0608e669851d0d2a6141dae74c60d5a27e8daa451f2a4857bbf41e2" +dependencies = [ + "thiserror", +] + +[[package]] +name = "gix-chunk" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c356b3825677cb6ff579551bb8311a81821e184453cbd105e2fc5311b288eeb" +dependencies = [ + "thiserror", +] + +[[package]] +name = "gix-command" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "46f9c425730a654835351e6da8c3c69ba1804f8b8d4e96d027254151138d5c64" +dependencies = [ + "bstr", + "gix-path", + "gix-quote", + "gix-trace", + "shell-words", +] + +[[package]] +name = "gix-commitgraph" +version = "0.29.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6bb23121e952f43a5b07e3e80890336cb847297467a410475036242732980d06" +dependencies = [ + "bstr", + "gix-chunk", + "gix-hash", + "memmap2", + "thiserror", +] + +[[package]] +name = "gix-config" +version = "0.46.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5dfb898c5b695fd4acfc3c0ab638525a65545d47706064dcf7b5ead6cdb136c0" +dependencies = [ + "bstr", + "gix-config-value", + "gix-features", + "gix-glob", + "gix-path", + "gix-ref", + "gix-sec", + "memchr", + "once_cell", + "smallvec", + "thiserror", + "unicode-bom", + "winnow 0.7.15", +] + +[[package]] +name = "gix-config-value" +version = "0.15.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c489abb061c74b0c3ad790e24a606ef968cebab48ec673d6a891ece7d5aef64" +dependencies = [ + "bitflags", + "bstr", + "gix-path", + "libc", + "thiserror", +] + +[[package]] +name = "gix-credentials" +version = "0.30.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0039dd3ac606dd80b16353a41b61fc237ca5cb8b612f67a9f880adfad4be4e05" +dependencies = [ + "bstr", + "gix-command", + "gix-config-value", + "gix-date", + "gix-path", + "gix-prompt", + "gix-sec", + "gix-trace", + "gix-url", + "thiserror", +] + +[[package]] +name = "gix-date" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "661245d045aa7c16ba4244daaabd823c562c3e45f1f25b816be2c57ee09f2171" +dependencies = [ + "bstr", + "itoa", + "jiff", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-diff" +version = "0.53.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de854852010d44a317f30c92d67a983e691c9478c8a3fb4117c1f48626bcdea8" +dependencies = [ + "bstr", + "gix-attributes", + "gix-command", + "gix-filter", + "gix-fs", + "gix-hash", + "gix-index", + "gix-object", + "gix-path", + "gix-pathspec", + "gix-tempfile", + "gix-trace", + "gix-traverse", + "gix-worktree", + "imara-diff", + "thiserror", +] + +[[package]] +name = "gix-dir" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dad34e4f373f94902df1ba1d2a1df3a1b29eacd15e316ac5972d842e31422dd7" +dependencies = [ + "bstr", + "gix-discover", + "gix-fs", + "gix-ignore", + "gix-index", + "gix-object", + "gix-path", + "gix-pathspec", + "gix-trace", + "gix-utils", + "gix-worktree", + "thiserror", +] + +[[package]] +name = "gix-discover" +version = "0.41.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffb180c91ca1a2cf53e828bb63d8d8f8fa7526f49b83b33d7f46cbeb5d79d30a" +dependencies = [ + "bstr", + "dunce", + "gix-fs", + "gix-hash", + "gix-path", + "gix-ref", + "gix-sec", + "thiserror", +] + +[[package]] +name = "gix-features" +version = "0.43.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd1543cd9b8abcbcebaa1a666a5c168ee2cda4dea50d3961ee0e6d1c42f81e5b" +dependencies = [ + "bytes", + "crc32fast", + "crossbeam-channel", + "flate2", + "gix-path", + "gix-trace", + "gix-utils", + "libc", + "once_cell", + "parking_lot", + "prodash", + "thiserror", + "walkdir", +] + +[[package]] +name = "gix-filter" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa6571a3927e7ab10f64279a088e0dae08e8da05547771796d7389bbe28ad9ff" +dependencies = [ + "bstr", + "encoding_rs", + "gix-attributes", + "gix-command", + "gix-hash", + "gix-object", + "gix-packetline-blocking", + "gix-path", + "gix-quote", + "gix-trace", + "gix-utils", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-fs" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a4d90307d064fa7230e0f87b03231be28f8ba63b913fc15346f489519d0c304" +dependencies = [ + "bstr", + "fastrand", + "gix-features", + "gix-path", + "gix-utils", + "thiserror", +] + +[[package]] +name = "gix-glob" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b947db8366823e7a750c254f6bb29e27e17f27e457bf336ba79b32423db62cd5" +dependencies = [ + "bitflags", + "bstr", + "gix-features", + "gix-path", +] + +[[package]] +name = "gix-hash" +version = "0.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "251fad79796a731a2a7664d9ea95ee29a9e99474de2769e152238d4fdb69d50e" +dependencies = [ + "faster-hex", + "gix-features", + "sha1-checked", + "thiserror", +] + +[[package]] +name = "gix-hashtable" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c35300b54896153e55d53f4180460931ccd69b7e8d2f6b9d6401122cdedc4f07" +dependencies = [ + "gix-hash", + "hashbrown 0.15.5", + "parking_lot", +] + +[[package]] +name = "gix-ignore" +version = "0.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "564d6fddf46e2c981f571b23d6ad40cb08bddcaf6fc7458b1d49727ad23c2870" +dependencies = [ + "bstr", + "gix-glob", + "gix-path", + "gix-trace", + "unicode-bom", +] + +[[package]] +name = "gix-index" +version = "0.41.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af39fde3ce4ce11371d9ce826f2936ec347318f2d1972fe98c2e7134e267e25" +dependencies = [ + "bitflags", + "bstr", + "filetime", + "fnv", + "gix-bitmap", + "gix-features", + "gix-fs", + "gix-hash", + "gix-lock", + "gix-object", + "gix-traverse", + "gix-utils", + "gix-validate", + "hashbrown 0.15.5", + "itoa", + "libc", + "memmap2", + "rustix", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-lock" +version = "18.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9fa71da90365668a621e184eb5b979904471af1b3b09b943a84bc50e8ad42ed" +dependencies = [ + "gix-tempfile", + "gix-utils", + "thiserror", +] + +[[package]] +name = "gix-negotiate" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d58d4c9118885233be971e0d7a589f5cfb1a8bd6cb6e2ecfb0fc6b1b293c83b" +dependencies = [ + "bitflags", + "gix-commitgraph", + "gix-date", + "gix-hash", + "gix-object", + "gix-revwalk", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-object" +version = "0.50.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d69ce108ab67b65fbd4fb7e1331502429d78baeb2eee10008bdef55765397c07" +dependencies = [ + "bstr", + "gix-actor", + "gix-date", + "gix-features", + "gix-hash", + "gix-hashtable", + "gix-path", + "gix-utils", + "gix-validate", + "itoa", + "smallvec", + "thiserror", + "winnow 0.7.15", +] + +[[package]] +name = "gix-odb" +version = "0.70.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c9d7af10fda9df0bb4f7f9bd507963560b3c66cb15a5b825caf752e0eb109ac" +dependencies = [ + "arc-swap", + "gix-date", + "gix-features", + "gix-fs", + "gix-hash", + "gix-hashtable", + "gix-object", + "gix-pack", + "gix-path", + "gix-quote", + "parking_lot", + "tempfile", + "thiserror", +] + +[[package]] +name = "gix-pack" +version = "0.60.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d8571df89bfca5abb49c3e3372393f7af7e6f8b8dbe2b96303593cef5b263019" +dependencies = [ + "clru", + "gix-chunk", + "gix-features", + "gix-hash", + "gix-hashtable", + "gix-object", + "gix-path", + "gix-tempfile", + "memmap2", + "parking_lot", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-packetline" +version = "0.19.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64286a8b5148e76ab80932e72762dd27ccf6169dd7a134b027c8a262a8262fcf" +dependencies = [ + "bstr", + "faster-hex", + "gix-trace", + "thiserror", +] + +[[package]] +name = "gix-packetline-blocking" +version = "0.19.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89c59c3ad41e68cb38547d849e9ef5ccfc0d00f282244ba1441ae856be54d001" +dependencies = [ + "bstr", + "faster-hex", + "gix-trace", + "thiserror", +] + +[[package]] +name = "gix-path" +version = "0.10.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cb06c3e4f8eed6e24fd915fa93145e28a511f4ea0e768bae16673e05ed3f366" +dependencies = [ + "bstr", + "gix-trace", + "gix-validate", + "thiserror", +] + +[[package]] +name = "gix-pathspec" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "daedead611c9bd1f3640dc90a9012b45f790201788af4d659f28d94071da7fba" +dependencies = [ + "bitflags", + "bstr", + "gix-attributes", + "gix-config-value", + "gix-glob", + "gix-path", + "thiserror", +] + +[[package]] +name = "gix-prompt" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "868e6516dfa16fdcbc5f8c935167d085f2ae65ccd4c9476a4319579d12a69d8d" +dependencies = [ + "gix-command", + "gix-config-value", + "parking_lot", + "rustix", + "thiserror", +] + +[[package]] +name = "gix-protocol" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12b4b807c47ffcf7c1e5b8119585368a56449f3493da93b931e1d4239364e922" +dependencies = [ + "bstr", + "gix-credentials", + "gix-date", + "gix-features", + "gix-hash", + "gix-lock", + "gix-negotiate", + "gix-object", + "gix-ref", + "gix-refspec", + "gix-revwalk", + "gix-shallow", + "gix-trace", + "gix-transport", + "gix-utils", + "maybe-async", + "thiserror", + "winnow 0.7.15", +] + +[[package]] +name = "gix-quote" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96fc2ff2ec8cc0c92807f02eab1f00eb02619fc2810d13dc42679492fcc36757" +dependencies = [ + "bstr", + "gix-utils", + "thiserror", +] + +[[package]] +name = "gix-ref" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b966f578079a42f4a51413b17bce476544cca1cf605753466669082f94721758" +dependencies = [ + "gix-actor", + "gix-features", + "gix-fs", + "gix-hash", + "gix-lock", + "gix-object", + "gix-path", + "gix-tempfile", + "gix-utils", + "gix-validate", + "memmap2", + "thiserror", + "winnow 0.7.15", +] + +[[package]] +name = "gix-refspec" +version = "0.31.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d29cae1ae31108826e7156a5e60bffacab405f4413f5bc0375e19772cce0055" +dependencies = [ + "bstr", + "gix-hash", + "gix-revision", + "gix-validate", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-revision" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f651f2b1742f760bb8161d6743229206e962b73d9c33c41f4e4aefa6586cbd3d" +dependencies = [ + "bstr", + "gix-commitgraph", + "gix-date", + "gix-hash", + "gix-object", + "gix-revwalk", + "thiserror", +] + +[[package]] +name = "gix-revwalk" +version = "0.21.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06e74f91709729e099af6721bd0fa7d62f243f2005085152301ca5cdd86ec02c" +dependencies = [ + "gix-commitgraph", + "gix-date", + "gix-hash", + "gix-hashtable", + "gix-object", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-sec" +version = "0.12.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea9962ed6d9114f7f100efe038752f41283c225bb507a2888903ac593dffa6be" +dependencies = [ + "bitflags", + "gix-path", + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "gix-shallow" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d936745103243ae4c510f19e0760ce73fb0f08096588fdbe0f0d7fb7ce8944b7" +dependencies = [ + "bstr", + "gix-hash", + "gix-lock", + "thiserror", +] + +[[package]] +name = "gix-status" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4afff9b34eeececa8bdc32b42fb318434b6b1391d9f8d45fe455af08dc2d35" +dependencies = [ + "bstr", + "filetime", + "gix-diff", + "gix-dir", + "gix-features", + "gix-filter", + "gix-fs", + "gix-hash", + "gix-index", + "gix-object", + "gix-path", + "gix-pathspec", + "gix-worktree", + "portable-atomic", + "thiserror", +] + +[[package]] +name = "gix-submodule" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "657cc5dd43cbc7a14d9c5aaf02cfbe9c2a15d077cded3f304adb30ef78852d3e" +dependencies = [ + "bstr", + "gix-config", + "gix-path", + "gix-pathspec", + "gix-refspec", + "gix-url", + "thiserror", +] + +[[package]] +name = "gix-tempfile" +version = "18.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "666c0041bcdedf5fa05e9bef663c897debab24b7dc1741605742412d1d47da57" +dependencies = [ + "dashmap", + "gix-fs", + "libc", + "once_cell", + "parking_lot", + "tempfile", +] + +[[package]] +name = "gix-trace" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44dc45eae785c0eb14173e0f152e6e224dcf4d45b6a6999a3aed22af541ad678" + +[[package]] +name = "gix-transport" +version = "0.48.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12f7cc0179fc89d53c54e1f9ce51229494864ab4bf136132d69db1b011741ca3" +dependencies = [ + "base64", + "bstr", + "curl", + "gix-command", + "gix-credentials", + "gix-features", + "gix-packetline", + "gix-quote", + "gix-sec", + "gix-url", + "thiserror", +] + +[[package]] +name = "gix-traverse" +version = "0.47.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7cdc82509d792ba0ad815f86f6b469c7afe10f94362e96c4494525a6601bdd5" +dependencies = [ + "bitflags", + "gix-commitgraph", + "gix-date", + "gix-hash", + "gix-hashtable", + "gix-object", + "gix-revwalk", + "smallvec", + "thiserror", +] + +[[package]] +name = "gix-url" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b76a9d266254ad287ffd44467cd88e7868799b08f4d52e02d942b93e514d16f" +dependencies = [ + "bstr", + "gix-features", + "gix-path", + "percent-encoding", + "thiserror", + "url", +] + +[[package]] +name = "gix-utils" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66c50966184123caf580ffa64e28031a878597f1c7fceb8fe19566c38eb1b771" +dependencies = [ + "bstr", + "fastrand", + "unicode-normalization", +] + +[[package]] +name = "gix-validate" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b1e63a5b516e970a594f870ed4571a8fdcb8a344e7bd407a20db8bd61dbfde4" +dependencies = [ + "bstr", + "thiserror", +] + +[[package]] +name = "gix-worktree" +version = "0.42.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55f625ac9126c19bef06dbc6d2703cdd7987e21e35b497bb265ac37d383877b1" +dependencies = [ + "bstr", + "gix-attributes", + "gix-features", + "gix-fs", + "gix-glob", + "gix-hash", + "gix-ignore", + "gix-index", + "gix-object", + "gix-path", + "gix-validate", +] + +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "globset" +version = "0.4.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52dfc19153a48bde0cbd630453615c8151bce3a5adfac7a0aebfbf0a1e1f57e3" +dependencies = [ + "aho-corasick", + "bstr", + "log", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core 0.6.4", + "subtle", +] + +[[package]] +name = "hash32" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47d60b12902ba28e2730cd37e95b8c9223af2808df9e902d4df49588d1470606" +dependencies = [ + "byteorder", +] + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.1.5", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.2.0", +] + +[[package]] +name = "hashbrown" +version = "0.17.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" + +[[package]] +name = "hashlink" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7382cf6263419f2d8df38c55d7da83da5c18aef87fc7a7fc1fb1e344edfe14c1" +dependencies = [ + "hashbrown 0.15.5", +] + +[[package]] +name = "heapless" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" +dependencies = [ + "hash32", + "stable_deref_trait", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "hkdf" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7" +dependencies = [ + "hmac", +] + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "http-auth" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "150fa4a9462ef926824cf4519c84ed652ca8f4fbae34cb8af045b5cbcaf98822" +dependencies = [ + "memchr", +] + +[[package]] +name = "icu_collections" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2984d1cd16c883d7935b9e07e44071dca8d917fd52ecc02c04d5fa0b5a3f191c" +dependencies = [ + "displaydoc", + "potential_utf", + "utf8_iter", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92219b62b3e2b4d88ac5119f8904c10f8f61bf7e95b640d25ba3075e6cac2c29" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c56e5ee99d6e3d33bd91c5d85458b6005a22140021cc324cea84dd0e72cff3b4" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3be0ae77ea334f4da67c12f149704f19f81d1adf7c51cf482943e84a2bad38" + +[[package]] +name = "icu_properties" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bee3b67d0ea5c2cca5003417989af8996f8604e34fb9ddf96208a033901e70de" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e2bbb201e0c04f7b4b3e14382af113e17ba4f63e2c9d2ee626b720cbce54a14" + +[[package]] +name = "icu_provider" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "139c4cf31c8b5f33d7e199446eff9c1e02decfc2f0eec2c8d71f65befa45b421" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb68373c0d6620ef8105e855e7745e18b0d00d3bdb07fb532e434244cdb9a714" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "ignore" +version = "0.4.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b915661dd01db3f05050265b2477bcc6527b3792388e2749b41623cc592be67d" +dependencies = [ + "crossbeam-deque", + "globset", + "log", + "memchr", + "regex-automata", + "same-file", + "walkdir", + "winapi-util", +] + +[[package]] +name = "im-rc" +version = "15.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af1955a75fa080c677d3972822ec4bad316169ab1cfc6c257a942c2265dbe5fe" +dependencies = [ + "bitmaps", + "rand_core 0.6.4", + "rand_xoshiro", + "sized-chunks", + "typenum", + "version_check", +] + +[[package]] +name = "imara-diff" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "17d34b7d42178945f775e84bc4c36dde7c1c6cdfea656d3354d009056f2bb3d2" +dependencies = [ + "hashbrown 0.15.5", +] + +[[package]] +name = "indexmap" +version = "2.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" +dependencies = [ + "equivalent", + "hashbrown 0.17.1", + "serde", + "serde_core", +] + +[[package]] +name = "is_executable" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baabb8b4867b26294d818bf3f651a454b6901431711abb96e296245888d6e8c4" +dependencies = [ + "windows-sys 0.60.2", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "jiff" +version = "0.2.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4603d3033e49e2b0e31229fcab20a5d40089c607d975cd9c80551dc69eed9102" +dependencies = [ + "jiff-static", + "jiff-tzdb-platform", + "log", + "portable-atomic", + "portable-atomic-util", + "serde_core", + "windows-link", +] + +[[package]] +name = "jiff-static" +version = "0.2.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "782d32378dddf207193ac91cefb848ad41abb58195c95168e1291227a0832b47" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "jiff-tzdb" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c900ef84826f1338a557697dc8fc601df9ca9af4ac137c7fb61d4c6f2dfd3076" + +[[package]] +name = "jiff-tzdb-platform" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "875a5a69ac2bab1a891711cf5eccbec1ce0341ea805560dcd90b7a2e925132e8" +dependencies = [ + "jiff-tzdb", +] + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom 0.3.4", + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.100" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2025f20d7a4fa7785846e7b63d10a76d3f1cee98ee5cb79ea59703f95e42162" +dependencies = [ + "cfg-if", + "wasm-bindgen", +] + +[[package]] +name = "kstring" +version = "2.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "558bf9508a558512042d3095138b1f7b8fe90c5467d94f9f1da28b3731c5dbd1" +dependencies = [ + "static_assertions", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "lazycell" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830d08ce1d1d941e6b30645f1a0eb5643013d835ce3779a5fc208261dbe10f55" + +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libgit2-sys" +version = "0.18.5+1.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "005d6ae6eac1912906073e069f7db60b1fa98e052a68227824afe3e3a1c59ca2" +dependencies = [ + "cc", + "libc", + "libssh2-sys", + "libz-sys", + "openssl-sys", + "pkg-config", +] + +[[package]] +name = "libloading" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7c4b02199fee7c5d21a5ae7d8cfa79a6ef5bb2fc834d6e9058e89c825efdc55" +dependencies = [ + "cfg-if", + "windows-link", +] + +[[package]] +name = "libnghttp2-sys" +version = "0.1.13+1.68.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "492e00167f1418c15648144f42bbfc63099806ecee9bf8d09a6353d6b4856b3c" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "libsqlite3-sys" +version = "0.35.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "133c182a6a2c87864fe97778797e46c7e999672690dc9fa3ee8e241aa4a9c13f" +dependencies = [ + "cc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "libssh2-sys" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "220e4f05ad4a218192533b300327f5150e809b54c4ec83b5a1d91833601811b9" +dependencies = [ + "cc", + "libc", + "libz-sys", + "openssl-sys", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "libz-sys" +version = "1.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85bc9657773828b90eeb625adff10eeac83cc21bbfd8e23a03eaa8a33c9e28d9" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "953f07c43838f8e6f9758cab68bf5bed85465e7587ebe0b823f1bcd81978ad3a" + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "maybe-async" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "746873a384ad60adc5db74471dfaba74bd278afbdcfd81db93fafcdfc8b5ca0c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "memmap2" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714098028fe011992e1c3962653c96b2d578c4b4bce9036e15ff220319b1e0e3" +dependencies = [ + "libc", +] + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "miow" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "536bfad37a309d62069485248eeaba1e8d9853aaf951caaeaed0585a95346f08" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "nix" +version = "0.31.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf20d2fde8ff38632c426f1165ed7436270b44f199fc55284c38276f9db47c3d" +dependencies = [ + "bitflags", + "cfg-if", + "cfg_aliases", + "libc", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "normpath" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9985ef7269fa99f3b12437bb698381da2428743ab90f20393f399fa14cab21a" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "num-conv" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "objc2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a12a8ed07aefc768292f076dc3ac8c48f3781c8f2d5851dd3d98950e8c5a89f" +dependencies = [ + "objc2-encode", +] + +[[package]] +name = "objc2-cloud-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73ad74d880bb43877038da939b7427bba67e9dd42004a18b809ba7d87cee241c" +dependencies = [ + "bitflags", + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-core-data" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b402a653efbb5e82ce4df10683b6b28027616a2715e90009947d50b8dd298fa" +dependencies = [ + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-core-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a180dd8642fa45cdb7dd721cd4c11b1cadd4929ce112ebd8b9f5803cc79d536" +dependencies = [ + "bitflags", + "dispatch2", + "objc2", +] + +[[package]] +name = "objc2-core-graphics" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e022c9d066895efa1345f8e33e584b9f958da2fd4cd116792e15e07e4720a807" +dependencies = [ + "bitflags", + "dispatch2", + "objc2", + "objc2-core-foundation", + "objc2-io-surface", +] + +[[package]] +name = "objc2-core-image" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5d563b38d2b97209f8e861173de434bd0214cf020e3423a52624cd1d989f006" +dependencies = [ + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-core-location" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca347214e24bc973fc025fd0d36ebb179ff30536ed1f80252706db19ee452009" +dependencies = [ + "objc2", + "objc2-foundation", +] + +[[package]] +name = "objc2-core-text" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cde0dfb48d25d2b4862161a4d5fcc0e3c24367869ad306b0c9ec0073bfed92d" +dependencies = [ + "bitflags", + "objc2", + "objc2-core-foundation", + "objc2-core-graphics", +] + +[[package]] +name = "objc2-encode" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ef25abbcd74fb2609453eb695bd2f860d389e457f67dc17cafc8b8cbc89d0c33" + +[[package]] +name = "objc2-foundation" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3e0adef53c21f888deb4fa59fc59f7eb17404926ee8a6f59f5df0fd7f9f3272" +dependencies = [ + "bitflags", + "block2", + "libc", + "objc2", + "objc2-core-foundation", +] + +[[package]] +name = "objc2-io-surface" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "180788110936d59bab6bd83b6060ffdfffb3b922ba1396b312ae795e1de9d81d" +dependencies = [ + "bitflags", + "objc2", + "objc2-core-foundation", +] + +[[package]] +name = "objc2-quartz-core" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96c1358452b371bf9f104e21ec536d37a650eb10f7ee379fff67d2e08d537f1f" +dependencies = [ + "bitflags", + "objc2", + "objc2-core-foundation", + "objc2-foundation", +] + +[[package]] +name = "objc2-ui-kit" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d87d638e33c06f577498cbcc50491496a3ed4246998a7fbba7ccb98b1e7eab22" +dependencies = [ + "bitflags", + "block2", + "objc2", + "objc2-cloud-kit", + "objc2-core-data", + "objc2-core-foundation", + "objc2-core-graphics", + "objc2-core-image", + "objc2-core-location", + "objc2-core-text", + "objc2-foundation", + "objc2-quartz-core", + "objc2-user-notifications", +] + +[[package]] +name = "objc2-user-notifications" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9df9128cbbfef73cda168416ccf7f837b62737d748333bfe9ab71c245d76613e" +dependencies = [ + "objc2", + "objc2-foundation", +] + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "opener" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2b03ff07a220d0d0ec9a1f0f238951b7967a5a2e96aefcd21a117b1083415e9" +dependencies = [ + "bstr", + "normpath", + "windows-sys 0.61.2", +] + +[[package]] +name = "openssl-probe" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" + +[[package]] +name = "openssl-sys" +version = "0.9.116" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f28a22dc7140cda5f096e5e7724a6962ca81a7f8bfd2979f9b18c11af56318c4" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "ordered-float" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "orion" +version = "0.17.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e6758747fd1ce1efaf2bd43219ac4aa9e28263b236b2b6a1e486bcd06820707" +dependencies = [ + "fiat-crypto", + "subtle", +] + +[[package]] +name = "os_info" +version = "3.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cf20a545b305cf1da722b236b5155c9bb35f1d5ceb28c048bd96ca842f41b5b" +dependencies = [ + "android_system_properties", + "log", + "nix", + "objc2", + "objc2-foundation", + "objc2-ui-kit", + "windows-sys 0.61.2", +] + +[[package]] +name = "p384" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe42f1670a52a47d448f14b6a5c61dd78fce51856e68edaa38f7ae3a46b8d6b6" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "pasetors" +version = "0.7.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e838401fb2873bad417e6a03179014c748746f67311cb7317ab14fc0881fa9f0" +dependencies = [ + "ct-codecs", + "ed25519-compact", + "getrandom 0.4.2", + "orion", + "p384", + "rand_core 0.6.4", + "regex", + "serde", + "serde_derive", + "serde_json", + "sha2", + "subtle", + "time", + "zeroize", +] + +[[package]] +name = "pathdiff" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df94ce210e5bc13cb6651479fa48d14f601d9858cfe0467f43ae157023b938d3" + +[[package]] +name = "pem-rfc7468" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412" +dependencies = [ + "base64ct", +] + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pkcs8" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" +dependencies = [ + "der", + "spki", +] + +[[package]] +name = "pkg-config" +version = "0.3.33" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portable-atomic-util" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a106d1259c23fac8e543272398ae0e3c0b8d33c88ed73d0cc71b0f1d902618" +dependencies = [ + "portable-atomic", +] + +[[package]] +name = "potential_utf" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0103b1cef7ec0cf76490e969665504990193874ea05c85ff9bab8b911d0a0564" +dependencies = [ + "zerovec", +] + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "primeorder" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" +dependencies = [ + "elliptic-curve", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "prodash" +version = "30.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6efc566849d3d9d737c5cb06cc50e48950ebe3d3f9d70631490fff3a07b139" +dependencies = [ + "parking_lot", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44c5af06bb1b7d3216d91932aed5265164bf384dc89cd6ba05cf59a35f5f76ea" +dependencies = [ + "rand_chacha", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_xoshiro" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6f97cdb2a36ed4183de61b2f824cc45c9f1037f28afe0a322e9fff4c108b5aaa" +dependencies = [ + "rand_core 0.6.4", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "regex" +version = "1.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1292b7759ae1cb9ec195452d1390a074f0cd8541ab7a5a8c31cd6db45d4a6ba" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6f6ff9a378485b298a5286656da665ba74413d36db0979633275d2e708145d4" + +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac", + "subtle", +] + +[[package]] +name = "rusqlite" +version = "0.37.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "165ca6e57b20e1351573e3729b958bc62f0e48025386970b6e4d29e7a7e71f3f" +dependencies = [ + "bitflags", + "fallible-iterator", + "fallible-streaming-iterator", + "hashlink", + "libsqlite3-sys", + "smallvec", +] + +[[package]] +name = "rustc-hash" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94300abf3f1ae2e2b8ffb7b58043de3d399c73fa6f4b73826402a5c457614dbe" + +[[package]] +name = "rustc-stable-hash" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "781442f29170c5c93b7185ad559492601acdc71d5bb0706f5868094f45cfcd08" + +[[package]] +name = "rustfix" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "864792a841a1d785ba91b8d2a75e1936b40bc517020c3c2958ac403b92e4f00a" +dependencies = [ + "serde", + "serde_json", + "thiserror", + "tracing", +] + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "sec1" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "semver" +version = "1.0.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a7852d02fc848982e0c167ef163aaff9cd91dc640ba85e263cb1ce46fae51cd" +dependencies = [ + "serde", + "serde_core", +] + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde-untagged" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f9faf48a4a2d2693be24c6289dbe26552776eb7737074e6722891fadbe6c5058" +dependencies = [ + "erased-serde", + "serde", + "serde_core", + "typeid", +] + +[[package]] +name = "serde-value" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3a1a3341211875ef120e117ea7fd5228530ae7e7036a779fdc9117be6b3282c" +dependencies = [ + "ordered-float", + "serde", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_ignored" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115dffd5f3853e06e746965a20dcbae6ee747ae30b543d91b0e089668bb07798" +dependencies = [ + "serde", + "serde_core", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_spanned" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6662b5879511e06e8999a8a235d848113e942c9124f211511b16466ee2995f26" +dependencies = [ + "serde_core", +] + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest", +] + +[[package]] +name = "sha1-checked" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89f599ac0c323ebb1c6082821a54962b839832b03984598375bff3975b804423" +dependencies = [ + "digest", + "sha1", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shell-escape" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45bb67a18fa91266cc7807181f62f9178a6873bfad7dc788c42e6430db40184f" + +[[package]] +name = "shell-words" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc6fe69c597f9c37bfeeeeeb33da3530379845f10be461a66d16d03eca2ded77" + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "digest", + "rand_core 0.6.4", +] + +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" + +[[package]] +name = "sized-chunks" +version = "0.6.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16d69225bde7a69b235da73377861095455d298f2b970996eec25ddbb42b3d1e" +dependencies = [ + "bitmaps", + "typenum", +] + +[[package]] +name = "smallvec" +version = "1.15.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ed6a63f02c8539c91a8685a86f4099661ba3da017932f6ebbea6de3f0fa7c90" + +[[package]] +name = "socket2" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52d1cfed4120b4d927bf7c0f86d2087a4a7d6027c906d9f9d525a80573b9be51" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "supports-hyperlinks" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e396b6523b11ccb83120b115a0b7366de372751aa6edf19844dfb13a6af97e91" + +[[package]] +name = "supports-unicode" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7401a30af6cb5818bb64852270bb722533397edcfc7344954a38f420819ece2" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tar" +version = "0.4.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6221d9a6003c78398e3b239969f352578258df48c8eb051caadae0015bc840" +dependencies = [ + "filetime", + "libc", +] + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.2", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "terminal_size" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "230a1b821ccbd75b185820a1f1ff7b14d21da1e442e22c0863ea5f08771a8874" +dependencies = [ + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "time" +version = "0.3.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" + +[[package]] +name = "time-macros" +version = "0.2.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215" +dependencies = [ + "num-conv", + "time-core", +] + +[[package]] +name = "tinystr" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8323304221c2a851516f22236c5722a72eaa19749016521d6dff0824447d96d" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "toml" +version = "0.9.12+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf92845e79fc2e2def6a5d828f0801e29a2f8acc037becc5ab08595c7d5e9863" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow 0.7.15", +] + +[[package]] +name = "toml_datetime" +version = "0.7.5+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92e1cfed4a3038bc5a127e35a2d360f145e1f4b971b551a2ba5fd7aedf7e1347" +dependencies = [ + "serde_core", +] + +[[package]] +name = "toml_edit" +version = "0.23.10+spec-1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84c8b9f757e028cee9fa244aea147aab2a9ec09d5325a9b01e0a49730c2b5269" +dependencies = [ + "indexmap", + "serde_core", + "serde_spanned", + "toml_datetime", + "toml_parser", + "toml_writer", + "winnow 0.7.15", +] + +[[package]] +name = "toml_parser" +version = "1.1.2+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2abe9b86193656635d2411dc43050282ca48aa31c2451210f4202550afb7526" +dependencies = [ + "winnow 1.0.3", +] + +[[package]] +name = "toml_writer" +version = "1.1.1+spec-1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-chrome" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bf0a738ed5d6450a9fb96e86a23ad808de2b727fd1394585da5cdd6788ffe724" +dependencies = [ + "serde_json", + "tracing-core", + "tracing-subscriber", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", +] + +[[package]] +name = "typeid" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc7d623258602320d5c55d1bc22793b57daff0ec7efc270ea7d55ce1d5f5471c" + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicase" +version = "2.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbc4bc3a9f746d862c45cb89d705aa10f187bb96c76001afab07a0d35ce60142" + +[[package]] +name = "unicode-bom" +version = "2.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7eec5d1121208364f6793f7d2e222bf75a915c19557537745b195b253dd64217" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization" +version = "0.1.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fd4f6878c9cb28d874b009da9e8d183b5abc80117c40bbd187a1fde336be6e8" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.4+wasi-0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b67efb37e106e55ce722a510d6b5f9c17f083e5fc79afc2badeb12cc313d9487" +dependencies = [ + "wit-bindgen 0.57.1", +] + +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen 0.51.0", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.123" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a254a4b10c19a76f09a27640e7ffbf9bc30bf67e16a3bf28aaefa4920fe81563" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.123" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24a40fc75b0ec6f3746ceb10d36f53a93dcd68a93b11b6445983945d79eba0dc" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.123" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "908f34bd9b9ce3d4caf07b72dfab63d61504d156856c6bd3cd87fa350cf3985b" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.123" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7acbf7616c27b194bbb550bf77ed0c2c3e5b7fd1260a93082b95fb7f47959b92" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap", + "wasm-encoder", + "wasmparser", +] + +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags", + "hashbrown 0.15.5", + "indexmap", + "semver", +] + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm", + "windows_aarch64_msvc", + "windows_i686_gnu", + "windows_i686_gnullvm", + "windows_i686_msvc", + "windows_x86_64_gnu", + "windows_x86_64_gnullvm", + "windows_x86_64_msvc", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "winnow" +version = "0.7.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df79d97927682d2fd8adb29682d1140b343be4ac0f08fd68b7765d9c059d3945" +dependencies = [ + "memchr", +] + +[[package]] +name = "winnow" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" + +[[package]] +name = "wit-bindgen" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + +[[package]] +name = "wit-bindgen" +version = "0.57.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ebf944e87a7c253233ad6766e082e3cd714b5d03812acc24c318f549614536e" + +[[package]] +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap", + "prettyplease", + "syn", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags", + "indexmap", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] + +[[package]] +name = "writeable" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ffae5123b2d3fc086436f8834ae3ab053a283cfac8fe0a0b8eaae044768a4c4" + +[[package]] +name = "yoke" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "709fe23a0424b6a435d82152b1bd3fdfb0833487d5fa90d05d42762a9891fef5" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "de844c262c8848816172cef550288e7dc6c7b7814b4ee56b3e1553f275f1858e" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ce1022995ff5ff5d841ad7d994facc23098cd40152f2c1d11cd607c6f530653f" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.52" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ae7f38b72ec2a254e2b87ef277cf2cd4fb97cbebf944faa6f33354da0867930" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11532158c46691caf0f2593ea8358fed6bbf68a0315e80aae9bd41fbade684a1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" + +[[package]] +name = "zerotrie" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f9152d31db0792fa83f70fb2f83148effb5c1f5b8c7686c3459e361d9bc20bf" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90f911cbc359ab6af17377d242225f4d75119aec87ea711a880987b18cd7b239" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "625dc425cab0dca6dc3c3319506e6593dcb08a9f387ea3b284dbd52a92c40555" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zlib-rs" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3be3d40e40a133f9c916ee3f9f4fa2d9d63435b5fbe1bfc6d9dae0aa0ada1513" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/dev/release/README.md b/dev/release/README.md index 2ca495cbb135f..5b57fbc448ed9 100644 --- a/dev/release/README.md +++ b/dev/release/README.md @@ -98,6 +98,15 @@ running: cargo check -p datafusion ``` +Within the user documentation there are references to the current version number. +Update these to the current version. At the time of this writing we need to manually +update the following files + +- `docs/source/download.md` +- `docs/source/user-guide/configs.md` +- `docs/source/user-guide/crate-configuration.md` +- `docs/source/user-guide/example-usage.md` + Then commit the changes and create a PR targeting the release branch `branch-N`. ```shell diff --git a/dev/release/rat_exclude_files.txt b/dev/release/rat_exclude_files.txt index f5ce368df724e..77da7db87e409 100644 --- a/dev/release/rat_exclude_files.txt +++ b/dev/release/rat_exclude_files.txt @@ -52,9 +52,9 @@ Cargo.lock .history parquet-testing/* *rat.txt -datafusion/proto/src/generated/datafusion_proto_common.rs -datafusion/proto/src/generated/pbjson.rs -datafusion/proto/src/generated/prost.rs +datafusion/proto-models/src/generated/datafusion_proto_common.rs +datafusion/proto-models/src/generated/pbjson.rs +datafusion/proto-models/src/generated/prost.rs datafusion/proto-common/src/generated/pbjson.rs datafusion/proto-common/src/generated/prost.rs .github/ISSUE_TEMPLATE/bug_report.yml diff --git a/dev/rust_lint.sh b/dev/rust_lint.sh index 43d29bd88166d..73cab9c7f70bd 100755 --- a/dev/rust_lint.sh +++ b/dev/rust_lint.sh @@ -106,6 +106,7 @@ declare -a WRITE_STEPS=( ) declare -a READONLY_STEPS=( + "ci/scripts/check_no_cargo_install_in_workflows.sh|false" "ci/scripts/rust_docs.sh|false" ) diff --git a/docs/pyproject.toml b/docs/pyproject.toml index 548c5fd858a59..c09415e8e8c86 100644 --- a/docs/pyproject.toml +++ b/docs/pyproject.toml @@ -5,9 +5,9 @@ requires-python = ">=3.11" dependencies = [ "sphinx>=9,<10", "sphinx-reredirects>=1.1,<2", - "pydata-sphinx-theme>=0.17.1,<1", + "pydata-sphinx-theme>=0.20.0,<1", "myst-parser>=5.1.0,<6", - "maturin>=1.13.3,<2", + "maturin>=1.14.1,<2", "jinja2>=3.1.6,<4", - "setuptools>=82.0.1,<83", + "setuptools>=83.0.0,<84", ] diff --git a/docs/source/contributor-guide/development_environment.md b/docs/source/contributor-guide/development_environment.md index faffa29c9cf71..2e4e00726580b 100644 --- a/docs/source/contributor-guide/development_environment.md +++ b/docs/source/contributor-guide/development_environment.md @@ -108,7 +108,7 @@ DataFusion is written in Rust and it uses a standard rust toolkit: - `rustup update stable` DataFusion generally uses the latest stable release of Rust, though it may lag when new Rust toolchains release - See which toolchain is currently pinned in the [`rust-toolchain.toml`](https://github.com/apache/datafusion/blob/main/rust-toolchain.toml) file - - This can cause issues such as not having the rust-analyzer component installed for the specified toolchain, in which case just install it manually, e.g. `rustup component add --toolchain 1.95.0 rust-analyzer` + - This can cause issues such as not having the rust-analyzer component installed for the specified toolchain, in which case just install it manually, e.g. `rustup component add --toolchain 1.97.0 rust-analyzer` - `cargo build` - `cargo fmt` to format the code - etc. diff --git a/docs/source/contributor-guide/governance.md b/docs/source/contributor-guide/governance.md index 52c212a7c0b1b..c0208c9de1476 100644 --- a/docs/source/contributor-guide/governance.md +++ b/docs/source/contributor-guide/governance.md @@ -43,8 +43,8 @@ DataFusion is currently governed by the following individuals The following table can be updated by running the following script: ```bash -python 3 docs/scripts/update_committer_list.py -prettier -w docs/scripts/update_committer_list.py +python3 docs/scripts/update_committer_list.py +ci/scripts/doc_prettier_check.sh --write --allow-dirty ``` Notes: @@ -71,6 +71,7 @@ Notes: | Jeffrey Vo | jeffreyvo | [Jefffrey](https://github.com/Jefffrey) | | PMC | | Jonah Gao | jonah | [jonahgao](https://github.com/jonahgao) | | PMC | | Kun Liu | liukun | [liukun4515](https://github.com/liukun4515) | | PMC | +| Matt Butrovich | mbutrovich | [mbutrovich](https://github.com/mbutrovich) | Apple | PMC | | Marko Milenković | milenkovicm | [milenkovicm](https://github.com/milenkovicm) | | PMC | | Mehmet Ozan Kabak | ozankabak | [ozankabak](https://github.com/ozankabak) | Synnada, Inc | PMC | | Tim Saucer | timsaucer | [timsaucer](https://github.com/timsaucer) | | PMC | @@ -95,14 +96,14 @@ Notes: | Siew Kam Onn | kosiew | [kosiew](https://github.com/kosiew) | | Committer | | Kumar Ujjawal | kumarujjawal | [kumarUjjawal](https://github.com/kumarUjjawal) | | Committer | | Lewis Zhang | linwei | [lewiszlw](https://github.com/lewiszlw) | diit.cn | Committer | -| Matt Butrovich | mbutrovich | [mbutrovich](https://github.com/mbutrovich) | Apple | Committer | | Metehan Yildirim | mete | [metegenez](https://github.com/metegenez) | | Committer | -| Martin Tzvetanov Grigorov | mgrigorov | | | Committer | +| Martin Tzvetanov Grigorov | mgrigorov | [martin-g](https://github.com/martin-g) | | Committer | | Wang Mingming | mingmwang | [mingmwang](https://github.com/mingmwang) | | Committer | | Michael Ward | mjward | [Michael-J-Ward ](https://github.com/Michael-J-Ward) | | Committer | | Marco Neumann | mneumann | [crepererum](https://github.com/crepererum) | InfluxData | Committer | +| Neil Conway | neilc | [neilconway](https://github.com/neilconway) | | Committer | | Zhong Yanghong | nju_yaho | [yahoNanJing](https://github.com/yahoNanJing) | | Committer | -| Nuno Faria | nunofaria | | | Committer | +| Nuno Faria | nunofaria | [nuno-faria](https://github.com/nuno-faria) | | Committer | | Paddy Horan | paddyhoran | [paddyhoran](https://github.com/paddyhoran) | Assured Allies | Committer | | Parth Chandra | parthc | [parthchandra](https://github.com/parthchandra) | Apple | Committer | | Rémi Dettai | rdettai | [rdettai](https://github.com/rdettai) | | Committer | diff --git a/docs/source/contributor-guide/index.md b/docs/source/contributor-guide/index.md index 6ec1efa4d99fa..6f1a0f1c19907 100644 --- a/docs/source/contributor-guide/index.md +++ b/docs/source/contributor-guide/index.md @@ -101,6 +101,11 @@ If you are concerned that a larger design will be lost in a string of small PRs, Note all commits in a PR are squashed when merged to the `main` branch so there is one commit per PR after merge. +For larger PRs, it is often helpful to leave a review on your own PR with +comments calling out important changes or specific important choices. These +annotations can help reviewers quickly find areas they should focus on, thus +speeding up review. + ## Release Management and Backports Contributor-facing guidance for release branches, patch releases, and backports @@ -135,22 +140,8 @@ do take priority over the conventional commit approach, allowing maintainers to ## Reviewing Pull Requests -Some helpful links: - -- [PRs Waiting for Review] on GitHub -- [Approved PRs Waiting for Merge] on GitHub - -[prs waiting for review]: https://github.com/apache/datafusion/pulls?q=is%3Apr+is%3Aopen+-review%3Aapproved+-is%3Adraft+ -[approved prs waiting for merge]: https://github.com/apache/datafusion/pulls?q=is%3Apr+is%3Aopen+review%3Aapproved+-is%3Adraft - -When reviewing PRs, our primary goal is to improve DataFusion and its community together. PR feedback should be constructive with the aim to help improve the code as well as the understanding of the contributor. - -Please ensure any issues you raise contains a rationale and suggested alternative -- it is frustrating to be told "don't do it this way" without any clear reason or alternate provided. - -Some things to specifically check: - -1. Is the feature or fix covered sufficiently with tests (see the [Testing](testing.md) section)? -2. Is the code clear, and fits the style of the existing codebase? +See the [Reviewing Pull Requests](pr_review.md) guide for what we look for +when reviewing PRs and how to prepare your own for review. ## Performance Improvements diff --git a/docs/source/contributor-guide/pr_review.md b/docs/source/contributor-guide/pr_review.md new file mode 100644 index 0000000000000..1154ae0e96b04 --- /dev/null +++ b/docs/source/contributor-guide/pr_review.md @@ -0,0 +1,240 @@ + + +# Reviewing Pull Requests + +When reviewing PRs, our primary goal is to improve DataFusion and its community +together. PR feedback should be constructive and help improve the code as well +as the understanding of the contributor. + +Review bandwidth is currently our most limited resource, and reviews from the +broader community are both welcomed and encouraged. Reviewing PRs is a great way +to learn the codebase, and you do not need to be a committer to leave valuable +review feedback. In fact, one of the best ways to become a committer is to +thoughtfully review other PRs. + +Please ensure any comments you leave contain a rationale and suggested +alternative -- it is frustrating to be told "don't do it this way" without any +clear reason or alternative provided. + +The criteria in this guide are also a useful checklist when preparing your own +PR for review. + +## PR Review Mechanics + +Some helpful links: + +- [PRs Waiting for Review] on GitHub +- [Approved PRs Waiting for Merge] on GitHub + +[prs waiting for review]: https://github.com/apache/datafusion/pulls?q=is%3Apr+is%3Aopen+-review%3Aapproved+-is%3Adraft+ +[approved prs waiting for merge]: https://github.com/apache/datafusion/pulls?q=is%3Apr+is%3Aopen+review%3Aapproved+-is%3Adraft + +The overall PR lifecycle (CI triggering, approval, the 24-hour rule for +"major" PRs, and merging) is described in the +[Pull Request Overview](index.md#pull-request-overview) section of the +contributor guide. + +Practical tips: + +1. Check out the changes locally to explore them in your IDE or with an + agent, e.g. `gh pr checkout ` using the [GitHub CLI]. +2. There is normally no need to rerun locally any tests that CI has already run. +3. Leave comments on specific lines of the diff where possible, so the + discussion has context. +4. If you review a PR but don't feel confident approving it, leaving comments + is still valuable: a partial review (e.g. "I reviewed the tests and they + look good") helps the next reviewer focus their time. +5. Anything that does not need to block the current PR can be noted as a + potential follow-up (ideally by filing an issue), keeping the PR focused + and quick to merge. + +[github cli]: https://cli.github.com/ + +## Review the PR Description + +The PR description is often what users and contributors will find when they have +a question about the intention behind a change, or when the code itself is not +clear. The PR description also becomes the extended commit message. + +Check that the description: + +1. Concisely describes the **problem being solved from the user's point of + view**. + +2. Follows the [PR template], and answers the template's questions. + +3. Accurately describes the content of the PR, including any relevant context or + background. + Great descriptions have a high signal-to-noise ratio, summarizing + important implementation changes without repeating technical minutiae that + are already present in the code itself. + +4. Explicitly calls out any user-facing or API changes (see + [Review the Code](#review-the-code) below). + +[pr template]: https://github.com/apache/datafusion/blob/main/.github/pull_request_template.md + +## Review the Code Comments + +The goal of code comments is to help future readers of the code understand what +is not obvious from reading the code itself. Great comments make the code easier +to reason about for readers with the expected background, and help future +maintainers. + +Some practical guidelines for reviewing comments: + +1. The code has adequate comments focused on the **rationale** for any + non-obvious change (the "why"), not a restatement of what the code does + (the "what"), which is typically clear from reading the code itself. +2. Comments do not narrate irrelevant internal implementation details or the + history of how the change was developed (this is common in LLM-assisted + code, e.g. "// changed to use a HashMap" or "// this handles the case + mentioned above"). Such comments become irrelevant as soon as the PR merges. +3. When comments refer to other structs, functions, or modules, they should use + [rustdoc intra-doc links] (e.g. `` [`SessionContext`] ``) rather than plain + text names, so that `cargo doc` link checking ensures the references stay + valid as the code evolves. +4. New public APIs have doc comments, including examples where appropriate + (doc examples are also tested by CI, so they double as test coverage). +5. When documenting modules, functions, or fields, start with simple examples + and intuitive explanations, and optionally add formal, math-like + definitions when necessary. This makes the implementation easier to reason + about. +6. When something is confusing on first read, treat that as a good + opportunity to improve the comments. + +[rustdoc intra-doc links]: https://doc.rust-lang.org/rustdoc/write-documentation/linking-to-items-by-name.html + +## Review the Test Coverage + +Check that the feature or fix is covered sufficiently with tests (see the +[Testing](testing.md) guide for more details): the PR should include tests for +any new functionality, and a bug fix should include a test that reproduces the +reported problem. + +Guidelines for evaluating tests: + +1. Prefer `sqllogictest` (`.slt`) tests or DataFrame API tests where + possible, as they exercise **user-visible behavior** and are less coupled + to internal implementation details than unit tests. +2. Verify tests cover edge cases and common failure scenarios, not just the + common successful path. However, it is NOT necessary to test every possible + error path, especially if it is difficult to trigger or unlikely to occur in + practice. +3. Verify test coverage of changed code using the `codecov` check on the PR, + or by running [`cargo llvm-cov`] locally for an HTML report. Use judgment + about any uncovered lines -- the goal is confidence in the change, not + slavishly hitting some coverage number. +4. Avoid tests with lots of repeated boilerplate: when many tests share + near-identical setup, it is hard to understand what is different + (and thus what is actually being tested) between them. Make the _difference_ + between cases obvious. +5. Check that tests assert on specific expected values or plans (e.g. via + `insta` snapshots or `.slt` expected output) rather than merely checking + "no error occurred". +6. Verify tests actually cover the bug ("Ablation Testing"): For bug fixes, revert + the fix locally and check that the new test fails without it (i.e. the test + actually reproduces the bug or covers the new feature). + +[`cargo llvm-cov`]: https://github.com/taiki-e/cargo-llvm-cov + +## Review the Code + +Check that: + +1. The code is clear and fits the style of the existing codebase. +2. New functions and tests are placed near similar functions and + tests. For example, helper functions should be defined close to where they are used, + and new tests should be placed in the same module as the code they test. + SLT tests should be placed in an existing .slt file with related functionality, + unless the new tests are large enough to justify their own file. +3. New APIs are consistent with existing public APIs and patterns; where a + similar mechanism already exists, the PR should extend it rather than + introduce a parallel one. +4. Any changes to the public API follow the [API health policy]. +5. The change is appropriately scoped: unrelated refactoring, formatting + churn, or drive-by changes make review longer and are better as separate + PRs. +6. New errors are actionable, mention the offending item, and use + the right error variant (e.g. `plan_err!` for user-triggerable errors vs + `internal_err!` for invariant violations). + +[api health policy]: api-health.md + +## Review the Performance + +Performance is a key feature of DataFusion. See [Performance Improvements](index.md#performance-improvements) +for the project policy: an improvement should be "enough" to justify any +added code complexity, and performance PRs should come with benchmark +results. + +When reviewing: + +1. Find any relevant existing benchmarks and run them against `main`: + the [system-level SQL benchmarks] are run with `bench.sh` (see the + [benchmarks README]), and microbenchmarks (e.g. in + `datafusion/functions/benches`) are run with `cargo bench`. +2. Be aware that benchmarking on a machine where other + work is being done will make results hard to reproduce. Prefer a quiet, + dedicated machine and repeated runs. +3. If the PR claims a performance improvement, check that the reported + results are reproducible and that the benchmark exercises the changed + code path. + +[system-level sql benchmarks]: https://github.com/apache/datafusion/tree/main/benchmarks +[benchmarks readme]: https://github.com/apache/datafusion/blob/main/benchmarks/README.md + +## Best Practices for Reviewers + +Here are some suggested best practices to follow when reviewing PRs. + +### Review Tone: Thank Contributors and Praise Good Work Specifically + +Open reviews by thanking the author by name, and when a PR is well done, say +specifically what makes it good -- positive feedback encourages people to keep +contributing and helps them understand what is valued in the project. + +### State Approval Conditions Explicitly + +If you are not ready to approve, list concretely what you would need to see +before approving (e.g. "benchmark results and an upgrade guide entry") so +the author has a clear path to merge. + +### Defer Non-Blocking Work to Follow-On Issues + +Explicitly defer non-critical suggestions to a follow-on PR and file +(or ask the author to file) issues for them, so good PRs merge quickly +without scope creep. + +Similarly, when a PR mixes refactoring with behavior changes or fixes a narrow +problem with a broad mechanism, ask for it to be split or scoped down rather +than reviewing it as-is. + +### Narrate What You Verified When Approving + +Rather than a bare "LGTM", say what you actually checked ("traced the state +transitions by hand", "confirmed the hasher change cannot affect ordering") +so it is clear what was verified and what was not. + +### Invite Additional Committers on Core Changes + +For changes to core, widely shared code, leave the PR open for other +committers to look at and cc those who know the area, even after you have +approved. diff --git a/docs/source/contributor-guide/release_management.md b/docs/source/contributor-guide/release_management.md index 0515204a5ecbc..7053d0f994559 100644 --- a/docs/source/contributor-guide/release_management.md +++ b/docs/source/contributor-guide/release_management.md @@ -44,10 +44,11 @@ Changes reach a release branch in one of two ways: - (Most common) Fix the issue on `main` and then backport the merged change to the release branch - Fix the issue on the release branch and then forward-port the change to `main` -Releases are coordinated in a GitHub issue, such as the -[release issue for 50.3.0]. If you think a fix should be included in a patch -release, discuss it on the relevant tracking issue first. You can also open the -backport PR first and then link it from the tracking issue. +Releases are coordinated using GitHub issues. Each planned release is listed in +the [DataFusion Releases tracking issue], and each release is coordinated in a +dedicated issue, such as the [release issue for 50.3.0]. If you think a fix +should be included in a patch release, discuss it on the relevant tracking issue +or open a backport PR and link it there. To prepare for a new release series, maintainers: @@ -59,6 +60,81 @@ To prepare for a new release series, maintainers: - Create release candidate artifacts from the release branch - After approval, publish to crates.io, ASF distribution servers, and Git tags +## Backport Criteria + +A release branch is a stabilization branch for an imminent or recent patch +release. The bar for landing a change on a release branch is therefore +_higher_ than the bar for landing on `main`, not lower. These criteria define +what is eligible for backport; the [Backport Workflow](#backport-workflow) +below describes the mechanics. + +DataFusion follows Cargo SemVer, with breaking changes allowed at major +version boundaries — see the [API health policy] for the full framing of +public Rust and SQL API stability. Patch releases (`x.y.z`, `z ≥ 1`) carry +fixes only and never introduce new features or breaking changes. + +### Eligible for backport + +- **Security fixes.** Fixes for known or reported security issues should be + backported to every actively maintained release branch. +- **Correctness fixes.** Fixes for queries that produce incorrect results, + panics, data loss, or crashes. If the fix itself changes user-visible SQL + semantics to make a wrong result right, follow [Behavior changes] below. +- **Stability and regression fixes.** Fixes for regressions introduced in the + current release line, hangs, deadlocks, memory leaks, or other availability + issues. +- **Build, CI, and test fixes** required to keep the branch buildable and + releasable. +- **Documentation fixes** for behavior already in the release. Documentation + for behavior that exists only on `main` does not belong on a release branch. + +### Not recommended for backport + +- **New features**, including new SQL functions, new optimizer rules, new + configuration options, new public APIs, and new file-format support. Land + on `main` and ship in the next major release. +- **Breaking API changes** of any kind, Rust or SQL. DataFusion makes + breaking changes only at major version boundaries — see [API health policy]. +- **Refactors and cleanup** that do not fix a bug, even if they are correct. +- **Performance improvements** that are not also correctness or stability + fixes. Land on `main`. +- **Dependency upgrades**, except when the upgrade itself is the security or + correctness fix and there is no narrower alternative. + +### Behavior changes + +A "behavior change" is any fix that alters user-visible results: SQL +semantics (values, ordering, types, null handling), error messages that +downstream users may rely on, plan output, or default configuration values. + +Behavior-changing fixes need extra scrutiny on a release branch because +users upgrading between patch versions do not expect their queries to start +returning different results. When proposing one for backport, state on the +release tracking issue _why_ the change should ship in this patch release +rather than wait for the next major. The previous and new behavior should +already be documented on the original issue or PR — link to that rather +than restating it. + +If in doubt, default to "land on `main`, ship in the next major." + +### Who decides + +The release manager for the active release line is the final reviewer of +what goes into the patch release. They coordinate via the release tracking +issue (for example, the [release issue for 50.3.0]). Anyone may propose a +backport by opening a backport PR and linking it from the tracking issue; +inclusion is the release manager's call. + +### Active release branches + +DataFusion does not maintain Long-Term Support branches. In general only the +most recent `branch-NN` is actively maintained for backports, but if you need +fixes in older releases, we are open to discussion. + +Security fixes are an exception: a maintainer may choose to backport a +critical security fix to an older branch even after it would otherwise be +closed. Discuss on the dev list or in a tracking issue before doing so. + ## Backport Workflow The usual workflow is: @@ -117,7 +193,10 @@ This PR: [`main` branch]: https://github.com/apache/datafusion/tree/main [`branch-50`]: https://github.com/apache/datafusion/tree/branch-50 [the release process readme in `dev/release`]: https://github.com/apache/datafusion/blob/main/dev/release/README.md +[datafusion releases tracking issue]: https://github.com/apache/datafusion/issues/19783 [release issue for 50.3.0]: https://github.com/apache/datafusion/issues/18072 [example backport pr]: https://github.com/apache/datafusion/pull/18131 [additional backport pr example]: https://github.com/apache/datafusion/pull/20792 [testing documentation]: testing.md +[api health policy]: api-health.md +[behavior changes]: #behavior-changes diff --git a/docs/source/contributor-guide/roadmap.md b/docs/source/contributor-guide/roadmap.md index bfaf398d3f549..903b2c7b7ef38 100644 --- a/docs/source/contributor-guide/roadmap.md +++ b/docs/source/contributor-guide/roadmap.md @@ -52,12 +52,16 @@ any single organization or coordinating committee. We typically discuss our roadmap using GitHub issues, approximately quarterly, and invite you to join the discussion. +The current roadmap discussion is +[DataFusion 2026 Q3-Q4 Roadmap Discussion](https://github.com/apache/datafusion/issues/22882). + For more information: 1. [Search for issues labeled `roadmap`](https://github.com/apache/datafusion/issues?q=is%3Aissue%20%20%20roadmap) -2. [DataFusion Road Map: Q1 2026](https://github.com/apache/datafusion/issues/18494) -3. [DataFusion Road Map: Q3-Q4 2025](https://github.com/apache/datafusion/issues/15878) -4. [2024 Q4 / 2025 Q1 Roadmap](https://github.com/apache/datafusion/issues/13274) +2. [DataFusion 2026 Q3-Q4 Roadmap Discussion](https://github.com/apache/datafusion/issues/22882) +3. [DataFusion Road Map: Q1 2026](https://github.com/apache/datafusion/issues/18494) +4. [DataFusion Road Map: Q3-Q4 2025](https://github.com/apache/datafusion/issues/15878) +5. [2024 Q4 / 2025 Q1 Roadmap](https://github.com/apache/datafusion/issues/13274) ## Improvement Proposals diff --git a/docs/source/download.md b/docs/source/download.md index ed8fc06440f0c..34296262071c8 100644 --- a/docs/source/download.md +++ b/docs/source/download.md @@ -26,7 +26,7 @@ For example: ```toml [dependencies] -datafusion = "53.0.0" +datafusion = "55.0.0" ``` While DataFusion is distributed via [crates.io] as a convenience, the diff --git a/docs/source/index.rst b/docs/source/index.rst index b939be86a0e25..ea6ebb74c08b1 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -126,6 +126,7 @@ To get started, see user-guide/sql/index user-guide/configs user-guide/explain-usage + user-guide/parquet-content-defined-chunking user-guide/metrics user-guide/faq @@ -158,6 +159,7 @@ To get started, see :caption: Contributor Guide contributor-guide/index + contributor-guide/pr_review contributor-guide/communication contributor-guide/development_environment contributor-guide/architecture diff --git a/docs/source/library-user-guide/building-logical-plans.md b/docs/source/library-user-guide/building-logical-plans.md index 9dc0fcbf31578..6efd97879ac4d 100644 --- a/docs/source/library-user-guide/building-logical-plans.md +++ b/docs/source/library-user-guide/building-logical-plans.md @@ -86,7 +86,7 @@ Filter: person.id > Int32(500) [id:Int32;N, name:Utf8;N] DataFusion logical plans can be created using the [LogicalPlanBuilder] struct. There is also a [DataFrame] API which is a higher-level API that delegates to [LogicalPlanBuilder]. -There are several functions that can can be used to create a new builder, such as +There are several functions that can be used to create a new builder, such as - `empty` - create an empty plan with no fields - `values` - create a plan from a set of literal values diff --git a/docs/source/library-user-guide/custom-table-providers.md b/docs/source/library-user-guide/custom-table-providers.md index 81b2d131e65c3..c094f8bf7eb1b 100644 --- a/docs/source/library-user-guide/custom-table-providers.md +++ b/docs/source/library-user-guide/custom-table-providers.md @@ -247,14 +247,22 @@ impl ExecutionPlan for MyExecPlan { vec![] // Leaf node -- no children } - fn with_new_children( + fn replace_children( self: Arc, children: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { assert!(children.is_empty()); Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute)) + } + fn execute( &self, partition: usize, @@ -655,7 +663,7 @@ and reading files that cannot possibly match the query. # use datafusion::execution::context::TaskContext; # use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; # use datafusion::physical_expr::EquivalenceProperties; -# use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr, PlanProperties}; +# use datafusion::physical_plan::{DisplayAs, DisplayFormatType, ExecutionPlan, Partitioning, PhysicalExpr, PlanProperties, ChildrenPropertiesMode, ReplaceChildrenOptions}; # use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; # /// A table provider backed by date-partitioned directories. @@ -764,9 +772,17 @@ impl DatePartitionedTable { # fn name(&self) -> &str { "DatePartitionedExec" } # fn properties(&self) -> &Arc { &self.properties } # fn children(&self) -> Vec<&Arc> { vec![] } -# fn with_new_children(self: Arc, _: Vec>) -> Result> { Ok(self) } +# fn replace_children(self: Arc, _: Vec>, _: ReplaceChildrenOptions) -> Result> { Ok(self) } +# +# fn with_new_children( +# self: Arc, +# children: Vec>, +# ) -> Result> { +# self.replace_children(children, ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute)) +# } +# # fn execute(&self, _: usize, _: Arc) -> Result { todo!() } -# fn apply_expressions(&self, _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } +# fn apply_expressions(&self, _f: &mut dyn FnMut(&Arc) -> Result) -> Result { Ok(TreeNodeRecursion::Continue) } # } ``` @@ -801,10 +817,8 @@ use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{ -# DisplayAs, DisplayFormatType, - ExecutionPlan, Partitioning, -# PhysicalExpr, - PlanProperties, +# DisplayAs, DisplayFormatType, PhysicalExpr, + ChildrenPropertiesMode, ReplaceChildrenOptions, ExecutionPlan, Partitioning, PlanProperties, }; use futures::stream; @@ -874,13 +888,21 @@ impl ExecutionPlan for CountingExec { fn properties(&self) -> &Arc { &self.properties } fn children(&self) -> Vec<&Arc> { vec![] } - fn with_new_children( + fn replace_children( self: Arc, - _children: Vec>, + _: Vec>, + _: ReplaceChildrenOptions, ) -> Result> { Ok(self) } + fn with_new_children( + self: Arc, + children: Vec>, + ) -> Result> { + self.replace_children(children, ReplaceChildrenOptions::new(ChildrenPropertiesMode::Recompute)) + } + fn execute( &self, partition: usize, @@ -913,7 +935,7 @@ impl ExecutionPlan for CountingExec { # fn apply_expressions( # &self, -# _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, +# _f: &mut dyn FnMut(&Arc) -> Result, # ) -> Result { # Ok(TreeNodeRecursion::Continue) # } diff --git a/docs/source/library-user-guide/functions/adding-udfs.md b/docs/source/library-user-guide/functions/adding-udfs.md index 0221e2e5adeb0..c3a40557a006d 100644 --- a/docs/source/library-user-guide/functions/adding-udfs.md +++ b/docs/source/library-user-guide/functions/adding-udfs.md @@ -397,9 +397,9 @@ impl AsyncUpper { pub fn new() -> Self { Self { signature: Signature::new( - TypeSignature::Coercible(vec![Coercion::Exact { - desired_type: TypeSignatureClass::Native(logical_string()), - }]), + TypeSignature::Coercible(vec![Coercion::new_exact( + TypeSignatureClass::Native(logical_string()), + )]), Volatility::Volatile, ), } @@ -497,9 +497,9 @@ We can now transfer the async UDF into the normal scalar using `into_scalar_udf` # pub fn new() -> Self { # Self { # signature: Signature::new( -# TypeSignature::Coercible(vec![Coercion::Exact { -# desired_type: TypeSignatureClass::Native(logical_string()), -# }]), +# TypeSignature::Coercible(vec![Coercion::new_exact( +# TypeSignatureClass::Native(logical_string()), +# )]), # Volatility::Volatile, # ), # } @@ -1229,6 +1229,36 @@ The `create_udaf` has six arguments to check: - The fifth argument is the function implementation. This is the function that we defined above. - The sixth argument is the description of the state, which will by passed between execution stages. +### Returning multiple values from an Aggregate UDF + +An aggregate UDF can return a `DataType::Struct` when one aggregate result needs +to carry multiple values. This is useful for time-windowing extensions that +need to return metadata such as the window start, window end, and the aggregate +value together. + +Pass the relevant input columns to the aggregate so the accumulator has enough +information to update and merge state normally in multi-stage aggregate plans. +For example, rows can be grouped into time buckets with the built-in `date_bin` +function, while a struct-returning aggregate computes the value and carries +metadata about each bucket: + +```sql +SELECT + augmented_avg(time, value)['window_start'] AS window_start, + augmented_avg(time, value)['window_end'] AS window_end, + augmented_avg(time, value)['window_duration'] AS window_duration, + augmented_avg(time, value)['avg_value'] AS avg_value +FROM t +GROUP BY date_bin(INTERVAL '30 seconds', time) +ORDER BY window_start; +``` + +In this pattern `date_bin(...)` assigns rows to a time bucket, while +`augmented_avg(time, value)` is a normal aggregate UDF whose accumulator stores +mergeable state such as `window_start`, `window_end`, `sum`, and `count`. +The aggregate's `evaluate` method returns a `ScalarValue::Struct`, and callers +can project individual fields from that struct. + ```rust # use datafusion::arrow::array::ArrayRef; diff --git a/docs/source/library-user-guide/upgrading/54.0.0.md b/docs/source/library-user-guide/upgrading/54.0.0.md index 0c3bf20a91ed5..f8e7ac93c08d8 100644 --- a/docs/source/library-user-guide/upgrading/54.0.0.md +++ b/docs/source/library-user-guide/upgrading/54.0.0.md @@ -21,10 +21,6 @@ ## DataFusion 54.0.0 -**Note:** DataFusion `54.0.0` has not been released yet. The information provided -in this section pertains to features and changes that have already been merged -to the main branch and are awaiting release in this version. - ### `AggregateFunctionExpr::human_display()` now returns `Option<&str>` `datafusion_physical_expr::aggregate::AggregateFunctionExpr::human_display()` @@ -169,73 +165,6 @@ where string types are preferred (`UNION`, `CASE THEN/ELSE`, `NVL2`). string-preferring behavior - Crates that call `get_coerce_type_for_case_expression` -### `ExecutionPlan::apply_expressions` is now a required method - -`apply_expressions` has been added as a **required** method on the `ExecutionPlan` trait (no default implementation). The same applies to the `FileSource` and `DataSource` traits. Any custom implementation of these traits must now implement `apply_expressions`. - -**Who is affected:** - -- Users who implement custom `ExecutionPlan` nodes -- Users who implement custom `FileSource` or `DataSource` sources - -**Migration guide:** - -Add `apply_expressions` to your implementation. Call `f` on each top-level `PhysicalExpr` your node owns, using `visit_sibling` to correctly propagate `TreeNodeRecursion`: - -**Node with no expressions:** - -```rust,ignore -fn apply_expressions( - &self, - _f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, -) -> Result { - Ok(TreeNodeRecursion::Continue) -} -``` - -**Node with a single expression:** - -```rust,ignore -fn apply_expressions( - &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, -) -> Result { - f(self.predicate.as_ref()) -} -``` - -**Node with multiple expressions:** - -```rust,ignore -fn apply_expressions( - &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, -) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - for expr in &self.expressions { - tnr = tnr.visit_sibling(|| f(expr.as_ref()))?; - } - Ok(tnr) -} -``` - -**Node whose only expressions are in `output_ordering()` (e.g. a synthetic test node with no owned expression fields):** - -```rust,ignore -fn apply_expressions( - &self, - f: &mut dyn FnMut(&dyn PhysicalExpr) -> Result, -) -> Result { - let mut tnr = TreeNodeRecursion::Continue; - if let Some(ordering) = self.cache.output_ordering() { - for sort_expr in ordering { - tnr = tnr.visit_sibling(|| f(sort_expr.expr.as_ref()))?; - } - } - Ok(tnr) -} -``` - ### `ExecutionPlan::partition_statistics` now returns `Arc` `ExecutionPlan::partition_statistics` now returns `Result>` instead of `Result`. This avoids cloning `Statistics` when it is shared across multiple consumers. @@ -406,6 +335,44 @@ This produces two user-visible changes: `ScalarSubqueryExpr` expression. Code that walks or transforms `LogicalPlan` / `ExecutionPlan` trees, as well as `EXPLAIN` output, may need updating. +### Filter predicate evaluation order may differ from query text + +The logical optimizer now reorders filters so that cheap predicates (most binary +comparisons, `IS NULL`, `Between`, `InList`, etc.) evaluate before expensive +ones (`LIKE`, regex, scalar function calls, subqueries). For example, +`WHERE col LIKE '%foo%' AND col2 = 5` may evaluate `col2 = 5` before +`col LIKE '%foo%'`. + +**Evaluation order has never been guaranteed to match the order written in the +query.** The SQL standard explicitly allows implementations to evaluate operands +in any order; major engines (PostgreSQL, SQL Server, Oracle, MySQL) document the +same. Queries should not rely on left-to-right evaluation or short-circuit +semantics for `AND` or `OR`. Previous versions of DataFusion already reordered +predicates (e.g., as part of expression simplification or predicate pushdown); +the new reordering pass just increases the scenarios where the optimizer will +change predicate evaluation order. + +**Fallible-predicate patterns are particularly affected.** For example: + +```sql +WHERE s ~ '^[0-9]+$' AND CAST(s AS INT) > 0 +``` + +The intent is likely to filter non-numeric strings before the `CAST` runs, +but this depends on evaluation-order behavior the SQL standard does not +guaranteed. The new reorder makes this kind of pattern more likely to fail +at runtime if the optimizer moves the `CAST` ahead of the regex. To force +conditional evaluation, rewrite using `CASE`, which has standardized +short-circuit semantics: + +```sql +WHERE CASE WHEN s ~ '^[0-9]+$' THEN CAST(s AS INT) > 0 ELSE false END +``` + +Volatile expressions (`random()`, `now()`, etc.) are exempt — their position +in the conjunct list is preserved so the number of times they evaluate per +query does not change. + ### `datafusion-proto`: expression deserialization now takes a `TaskContext` `Serializeable::from_bytes_with_registry` is renamed to `from_bytes_with_ctx` @@ -961,3 +928,8 @@ match register_function { RegisterFunction::Table(name, table) => {}, } ``` + +### New `Dialect::Spark` variant + +The `Dialect` enum in `datafusion_common::config` now includes a `Spark` variant. +If you match exhaustively on `Dialect`, add a `Dialect::Spark` arm. diff --git a/docs/source/library-user-guide/upgrading/55.0.0.md b/docs/source/library-user-guide/upgrading/55.0.0.md new file mode 100644 index 0000000000000..d64f287ea0b52 --- /dev/null +++ b/docs/source/library-user-guide/upgrading/55.0.0.md @@ -0,0 +1,1293 @@ + + +# Upgrade Guides + +## DataFusion 55.0.0 + +**Note:** DataFusion `55.0.0` has not been released yet. The information provided +in this section pertains to features and changes that have already been merged +to the main branch and are awaiting release in this version. + +### `DataFrame::fill_null` now borrows its arguments + +`DataFrame::fill_null` previously took its arguments by value: + +```rust,ignore +// Before +pub fn fill_null( + &self, + value: ScalarValue, + columns: Vec, +) -> Result +``` + +It now borrows them, matching the signature of the newly added +`DataFrame::fill_nan`: + +```rust,ignore +// After +pub fn fill_null( + &self, + value: &ScalarValue, + columns: &[&str], +) -> Result +``` + +This lets callers pass a borrowed `ScalarValue` and slice literals (or +`&str` column names) without first allocating owned `String`s. + +**Migration guide:** + +Borrow the value and pass a slice of `&str` instead of an owned `Vec`: + +```rust,ignore +// Before +let df = df.fill_null(ScalarValue::from(0), vec!["a".to_owned(), "c".to_owned()])?; +let df = df.fill_null(ScalarValue::from(0), vec![])?; + +// After +let df = df.fill_null(&ScalarValue::from(0), &["a", "c"])?; +let df = df.fill_null(&ScalarValue::from(0), &[])?; +``` + +### `FileScanConfig::partitioned_by_file_group` removed + +`FileScanConfig::partitioned_by_file_group` and +`FileScanConfigBuilder::with_partitioned_by_file_group(...)` have been removed. +Use `FileScanConfig::output_partitioning` and +`FileScanConfigBuilder::with_output_partitioning(...)` instead. +The corresponding +`datafusion_proto::protobuf::FileScanExecConf::partitioned_by_file_group` +field has also been removed. + +**Who is affected:** + +- Users who accessed `FileScanConfig::partitioned_by_file_group` directly. +- Users who called + `FileScanConfigBuilder::with_partitioned_by_file_group(true)`. +- Users who constructed or accessed + `datafusion_proto::protobuf::FileScanExecConf::partitioned_by_file_group`. + +**Migration guide:** + +If your file groups are organized by table partition column values, declare hash +output partitioning over those partition columns: + +```rust,ignore +use datafusion_datasource::file_scan_config::{ + FileScanConfigBuilder, output_partitioning_from_partition_fields, +}; + +let output_partitioning = output_partitioning_from_partition_fields( + source.table_schema().table_schema(), + source.table_schema().table_partition_cols(), + file_groups.len(), +); + +let config = FileScanConfigBuilder::new(object_store_url, source) + .with_file_groups(file_groups) + .with_output_partitioning(output_partitioning) + .build(); +``` + +`output_partitioning_from_partition_fields` returns +`Some(Partitioning::Hash(...))` when partition columns are present and `None` +otherwise. If you construct the partitioning manually, pass +`Some(Partitioning::Hash(partition_exprs, partition_count))` to +`with_output_partitioning(...)`. + +When constructing `FileScanExecConf`, omit `partitioned_by_file_group` and set +`output_partitioning` instead. + +### User `SpillFile` traits instead of [`RefCountedTempFile`] + +Spill file APIs now use the `datafusion_execution::SpillFile` trait instead of +the concrete [`RefCountedTempFile`] type. [`DiskManager::create_tmp_file`] now +returns `Arc`. +This change was introduced in [PR #21882], which adds pluggable spill file +backends via `SpillFile` and `TempFileFactory`. + +If your code matched on [`DiskManagerMode`], add a `DiskManagerMode::Custom(_)` +arm. + +If your code wrote directly to a [`RefCountedTempFile`] or called +[`RefCountedTempFile::update_disk_usage`], open a spill writer instead: + +```diff +- temp_file.inner().as_file().write_all(bytes)?; +- temp_file.update_disk_usage()?; ++ temp_file.open_writer()?.write_all(bytes)?; +``` + +Use `temp_file.size()` instead of [`RefCountedTempFile::current_disk_usage`]. + +[`diskmanager::create_tmp_file`]: https://docs.rs/datafusion-execution/latest/datafusion_execution/disk_manager/struct.DiskManager.html#method.create_tmp_file +[`diskmanagermode`]: https://docs.rs/datafusion-execution/latest/datafusion_execution/disk_manager/enum.DiskManagerMode.html +[`pr #21882`]: https://github.com/apache/datafusion/pull/21882 +[`refcountedtempfile`]: https://docs.rs/datafusion-execution/latest/datafusion_execution/disk_manager/struct.RefCountedTempFile.html +[`refcountedtempfile::current_disk_usage`]: https://docs.rs/datafusion-execution/latest/datafusion_execution/disk_manager/struct.RefCountedTempFile.html#method.current_disk_usage +[`refcountedtempfile::update_disk_usage`]: https://docs.rs/datafusion-execution/latest/datafusion_execution/disk_manager/struct.RefCountedTempFile.html#method.update_disk_usage + +### `Dialect::AVAILABLE` replaced by `Dialect::available()` + +`datafusion_common::config::Dialect::AVAILABLE` has been removed. Use +`Dialect::available()` instead. + +### `spill_record_batch_by_size` removed + +`datafusion_physical_plan::spill::spill_record_batch_by_size` has been removed. +This function was deprecated in DataFusion `46.0.0`. + +Use `datafusion_physical_plan::spill::SpillManager::spill_record_batch_by_size` +instead. + +### `CreateExternalTable` supports multiple locations + +`CREATE EXTERNAL TABLE` now accepts multiple paths in a single `LOCATION` +clause, which are read together as one table: + +```sql +CREATE EXTERNAL TABLE hits +STORED AS PARQUET +LOCATION ('file_1.parquet', 'file_2.parquet'); +``` + +To support this, the `location` field of both +`datafusion_expr::CreateExternalTable` and +`datafusion_sql::parser::CreateExternalTable` changed from a `String` to a +`Vec` named `locations`: + +```rust +// Before (54.0.0) +let location: String = create_external_table.location; + +// After (55.0.0) +let locations: Vec = create_external_table.locations; +``` + +The `CreateExternalTable::builder(name, location, file_type, schema)` +constructor is unchanged and still takes a single location; use the new +`CreateExternalTableBuilder::with_locations(Vec)` to set more than one. +All listed locations must resolve to the same schema and reside on the same +object store. A plain string literal remains a single location, so paths that +contain commas continue to work, for example `LOCATION 'path/with,comma.csv'`. + +### Decimal scalar formatting uses human-readable values + +Decimal scalar literals in `EXPLAIN` output, expression display strings, and +auto-generated column names now format the decimal value using its scale while +still showing the precision and scale. For example, a `Decimal128` literal with +stored value `1`, precision `1`, and scale `1` is now rendered as +`Decimal128(0.1,1,1)` instead of `Decimal128(Some(1),1,1)`. When formatting a +`ScalarValue` directly, it now appears as `0.1` instead of `Some(1),1,1`. + +`NULL` decimal literals were previously shown as `Decimal128(None,10,2)`; they +will now appear as `Decimal128(NULL,10,2)`. + +Query result values already used human-readable decimal formatting and are +unchanged. + +### `Coercion` supports dictionary encoding preservation + +`datafusion_expr_common::signature::Coercion` now supports optional dictionary +encoding preservation. Typed coercions materialize dictionary inputs by +default, including both `TypeSignatureClass::Native(...)` and broader classes +such as `Integer`, `Numeric`, and `Binary`. When preservation is enabled, +DataFusion instead coerces dictionary inputs to +`Dictionary(original_key_type, coerced_value_type)` instead of materializing them +to the coerced value type. + +User-defined functions can opt in by setting dictionary encoding preservation on +the relevant coercion: + +```rust +Coercion::new_exact(TypeSignatureClass::Native(logical_string())) + .with_encoding_preservation(EncodingPreservation::dictionary()) +``` + +This changes the coerced argument type passed to the function. If a function +derives its return type from that coerced argument type, code that checks exact +result types may need to update its expectations or add an explicit cast to +materialize the result. + +This changes the previous behavior of typed non-native classes such as +`Integer` and `Binary`, which retained the physical dictionary type by default. +UDFs relying on that behavior must now explicitly enable dictionary +preservation. `TypeSignatureClass::Any` is unaffected. + +### `GroupsAccumulator::merge_batch` no longer takes `opt_filter` + +The `opt_filter` argument has been removed from +`datafusion_expr_common::groups_accumulator::GroupsAccumulator::merge_batch`: + +```diff + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], +- opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()>; +``` + +Aggregate `FILTER` clauses only apply to raw input rows during the partial +(update) phase, so by the time intermediate states are merged there is nothing +left to filter per row. In practice `opt_filter` was always `None` here, so +removing it makes the API self-explanatory and impossible to misuse. + +**Who is affected:** + +- Anyone with a custom `GroupsAccumulator` implementation. +- Anyone calling `merge_batch` directly. + +**Migration guide:** + +Drop the `opt_filter` argument from your `merge_batch` signature and from any +call sites: + +```diff + fn merge_batch( + &mut self, + values: &[ArrayRef], + group_indices: &[usize], +- opt_filter: Option<&BooleanArray>, + total_num_groups: usize, + ) -> Result<()> { + // ... + } +``` + +```diff +- acc.merge_batch(values, group_indices, None, total_num_groups)?; ++ acc.merge_batch(values, group_indices, total_num_groups)?; +``` + +If your implementation previously inspected `opt_filter` (for example asserting +it was `None`), that code can simply be deleted. + +See [issue #22775](https://github.com/apache/datafusion/issues/22775) for details. + +### `GroupsAccumulator::convert_to_state` is now required + +`datafusion_expr_common::groups_accumulator::GroupsAccumulator::convert_to_state` +no longer provides a default implementation, and the +`GroupsAccumulator::supports_convert_to_state` capability method has been +removed. All `GroupsAccumulator` implementations must now support converting +input batches directly to intermediate aggregate state. + +**Who is affected:** + +- Users with custom `GroupsAccumulator` implementations. +- FFI providers and consumers that use `FFI_GroupsAccumulator`. + +**Migration guide:** + +Custom `GroupsAccumulator` implementations must now provide their own +`convert_to_state` implementation. + +Delete `supports_convert_to_state` implementations because `convert_to_state` +is now required: + +```diff +- fn supports_convert_to_state(&self) -> bool { +- true +- } +``` + +The `supports_convert_to_state` field has also been removed from +`datafusion_ffi::udaf::groups_accumulator::FFI_GroupsAccumulator`, changing its +ABI layout. Rebuild both FFI providers and consumers against DataFusion 55, and +do not exchange this struct with libraries built against older major versions. + +See [issue #23081](https://github.com/apache/datafusion/issues/23081) for details. + +### `is_dynamic_physical_expr` is deprecated + +`datafusion_physical_expr_common::physical_expr::is_dynamic_physical_expr` is +deprecated. It was a thin wrapper over `snapshot_generation(expr) != 0` used to +ask "does this predicate contain a dynamic filter?". + +Prefer asking the question directly against the concrete type. For a one-off +check, downcast to `DynamicFilterPhysicalExpr`: + +```rust +use datafusion_physical_expr::expressions::DynamicFilterPhysicalExpr; +use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion}; + +let mut is_dynamic = false; +predicate.apply(|e| { + if e.downcast_ref::().is_some() { + is_dynamic = true; + Ok(TreeNodeRecursion::Stop) + } else { + Ok(TreeNodeRecursion::Continue) + } +})?; +``` + +If you also need to know whether the dynamic filters can still change (and to be +notified when they do), use the new `DynamicFilterTracking` / +`DynamicFilterTracker` API in `datafusion_physical_expr`: + +```rust +use datafusion_physical_expr::DynamicFilterTracking; + +let tracking = DynamicFilterTracking::classify(&predicate); +if tracking.contains_dynamic_filter() { + // worth re-evaluating the predicate at runtime +} +``` + +### `PruningPredicate::try_new` is deprecated + +`datafusion_pruning::PruningPredicate::try_new` is deprecated. Use +`PruningPredicateBuilder` instead. The deprecated constructor remains available +in DataFusion 55 and preserves its existing behavior. + +```rust +// Before +let predicate = PruningPredicate::try_new(expr, schema)?; + +// After +let predicate = PruningPredicateBuilder::new() + .with_file_schema(schema) + .try_build(expr)?; +``` + +### `FilePruner::try_new` no longer builds a pruner for static predicates without statistics + +`datafusion_pruning::FilePruner::try_new` now returns `None` when the predicate +is purely static _and_ the file carries no usable column statistics, because +such a pruner can never prune anything beyond what planning already did. +Previously it returned `Some` whenever a statistics struct was present (the +"is this worth pruning?" decision lived in the Parquet opener). Files with column +statistics, and predicates that carry a dynamic filter, are unaffected. + +### `QueryPlanner` adds `Any` as a supertrait + +To enable downcasting of `dyn QueryPlanner` to concrete query planner types (via +`is::()` / `downcast_ref::()`), the `QueryPlanner` trait now has `Any` +as a supertrait: + +```diff +- pub trait QueryPlanner: Debug ++ pub trait QueryPlanner: Any + Debug +``` + +### `ExecutionPlan::partition_statistics` deprecated in favor of `statistics_from_inputs` + +`ExecutionPlan::partition_statistics` is deprecated. Statistics computation is +now split into two parts: + +- `StatisticsContext` owns the bottom-up plan-tree traversal and a per-walk + cache of memoized child statistics. Call `StatisticsContext::compute` to + obtain statistics for a plan. +- `ExecutionPlan::statistics_from_inputs` computes a node's statistics from its + children's already-resolved statistics, which the context passes in. The node + does not traverse the tree itself. + +Existing implementations of `partition_statistics` continue to work unchanged. +The default `statistics_from_inputs` delegates to the deprecated method, so no +migration is required until the deprecated method is removed. + +> **Warning:** The delegation is **one-way**: the default `statistics_from_inputs` +> calls `partition_statistics`, but the default `partition_statistics` does +> **not** call `statistics_from_inputs` — it returns `Statistics::new_unknown`. +> Nodes that override only `statistics_from_inputs` will silently return +> `Statistics::new_unknown` to any caller still using the deprecated +> `partition_statistics`. + +**Who is affected:** + +- Users who implement custom `ExecutionPlan` nodes (recommended to migrate) +- Users who call `partition_statistics` directly (recommended to switch to `StatisticsContext::compute`) + +**Migration guide:** + +For **implementations**, override `statistics_from_inputs` instead of +`partition_statistics`, plus `child_stats_requests` to declare which children to +resolve. Child statistics then arrive pre-computed in `input_stats` (one entry per +child, in `children()` order), so the node only expresses its local propagation +logic. Leaf nodes, and nodes that derive their statistics without reading children, +need neither override (the default `child_stats_requests` skips every child). + +```rust,ignore +// Before: +fn partition_statistics(&self, partition: Option) -> Result> { + let child_stats = self.input.partition_statistics(partition)?; + // ... transform child_stats ... +} + +// After: declare the child to resolve, then compute from its statistics. +fn child_stats_requests(&self, partition: Option) -> Vec { + vec![ChildStats::At(partition)] +} + +fn statistics_from_inputs( + &self, + input_stats: &[Arc], + args: &StatisticsArgs, +) -> Result> { + let child_stats = Arc::clone(&input_stats[0]); + // ... transform child_stats ... +} +``` + +> **Important:** the default `child_stats_requests` skips every child, so a node that +> reads `input_stats` must override it to declare the children it uses, or those slots +> are filled with `Statistics::new_unknown` placeholders. Request a child with +> `ChildStats::At(partition)` (`None` = overall) and omit one with `ChildStats::Skip`. +> For example, a partition-merging operator requests `ChildStats::At(None)`, and a +> broadcast join requests its build side at `None`. + +For **callers**, walk a plan through `StatisticsContext::compute`. The cache is +created with the context: + +```rust,ignore +use datafusion_physical_plan::{StatisticsArgs, StatisticsContext}; + +// Before: +let stats = plan.partition_statistics(None)?; + +// After: +let stats = StatisticsContext::new().compute(plan.as_ref(), &StatisticsArgs::new())?; +``` + +### `DdlStatement::CreateExternalTable` and `CreateFunction` are now boxed + +The two largest variants of `datafusion_expr::DdlStatement` are now +`Box`ed: + +```rust,ignore +// Before +pub enum DdlStatement { + CreateExternalTable(CreateExternalTable), + // ... + CreateFunction(CreateFunction), + // ... +} + +// After +pub enum DdlStatement { + CreateExternalTable(Box), + // ... + CreateFunction(Box), + // ... +} +``` + +`CreateExternalTable` is 312 bytes and `CreateFunction` is 288 bytes, so +without boxing they forced the entire `LogicalPlan` enum to 320 bytes +even on SELECT-only query paths that never instantiate them. Boxing +shrinks `LogicalPlan` from 320 → 176 bytes (−45%), making every +`mem::take` / `mem::swap` / `Arc` store on the planning +hot path move a smaller payload. + +**Who is affected:** + +- Users who construct `DdlStatement::CreateExternalTable(...)` or + `DdlStatement::CreateFunction(...)` from an owned struct. +- Users who pattern-match these variants and destructure the inner + struct in the same pattern (e.g. + `DdlStatement::CreateExternalTable(CreateExternalTable { name, .. })`). +- Code that consumes the inner struct out of these variants (e.g. to + pass `CreateExternalTable` by value to another function). + +**Migration guide:** + +When constructing the variants, wrap the inner struct in `Box::new`: + +```rust,ignore +// Before +let stmt = DdlStatement::CreateFunction(CreateFunction { name, args, .. }); + +// After +let stmt = DdlStatement::CreateFunction(Box::new(CreateFunction { + name, + args, + .. +})); +``` + +When pattern-matching, bind the boxed value and either access fields +through it (Rust auto-derefs the `Box`) or destructure via `.as_ref()`: + +```rust,ignore +// Before +match ddl { + DdlStatement::CreateExternalTable(CreateExternalTable { + name, location, .. + }) => { /* use name, location */ } +} + +// After — access fields through the box +match ddl { + DdlStatement::CreateExternalTable(ce) => { + let name = &ce.name; + let location = &ce.location; + /* ... */ + } +} + +// After — destructure the dereferenced struct +match ddl { + DdlStatement::CreateExternalTable(ce) => { + let CreateExternalTable { name, location, .. } = ce.as_ref(); + /* ... */ + } +} +``` + +When you need an owned `CreateExternalTable` / `CreateFunction` out of +the variant, dereference the box with `*`: + +```rust,ignore +// Before +match plan { + LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) => Ok(cmd), + _ => { /* ... */ } +} + +// After +match plan { + LogicalPlan::Ddl(DdlStatement::CreateExternalTable(cmd)) => Ok(*cmd), + _ => { /* ... */ } +} +``` + +See [PR #22733](https://github.com/apache/datafusion/pull/22733) for +details, including the per-variant size breakdown and benchmark +results. + +### `ExecutionPlan::with_new_children` and `ExecutionPlan::with_new_children_and_same_properties` deprecated + +`with_new_children` and `with_new_children_and_same_properties` have been +deprecated. These methods are used to replace the child plans of an +`ExecutionPlan` while leaving the plan otherwise identical. + +`with_new_children_if_necessary` has also been deprecated in favor of +`replace_children_if_necessary` for consistency in naming. + +As noted [here](https://github.com/apache/datafusion/pull/23332#discussion_r3554897693), +while the addition of `with_new_children_and_same_properties` has the benefit +of skipping potentially expensive computation in the case that replacement children +have the same properties as the original children, it widens the API surface area +of `ExecutionPlan` in a way that could be confusing for users. + +Thus, to rectify this, we unify these methods by introducing `replace_children`. +`replace_children` solves this problem by taking `ReplaceChildrenOptions`, +which includes a `ChildrenPropertiesMode`. The mode has two variants, +`Keep` and `Recompute`, which tell `replace_children` whether plan +properties can be reused or need to be recomputed. + +This method is called from `replace_children_if_necessary`, which is the +standard entry point that should be used for replacing the children of a node. + +**Migration guide:** + +To migrate from `with_new_children` and `with_new_children_and_same_properties` +to `replace_children`, it is recommended to implement `replace_children` with +a `match` statement matching on the `ChildrenPropertiesMode`. In the case that +the properties match the children, `ChildrenPropertiesMode::Keep`, +follow the body of `with_new_children_and_same_properties`. In the case that +the properties do not match the children, `ChildrenPropertiesMode::Recompute`, +follow the body of `with_new_children`. + +For example, take a look at the implementation for `FilterExec`: + +``` + fn replace_children( + self: Arc, + mut children: Vec>, + options: ReplaceChildrenOptions, + ) -> Result> { + validate_child_count!(self, children); + match options.children_properties { + ChildrenPropertiesMode::Keep => Ok(Arc::new(Self { + input: children.swap_remove(0), + metrics: ExecutionPlanMetricsSet::new(), + ..Self::clone(&*self) + })), + ChildrenPropertiesMode::Recompute => { + let new_input = children.swap_remove(0); + FilterExecBuilder::from(&*self) + .with_input(new_input) + .build() + .map(|e| Arc::new(e) as _) + } + } + } +``` + +In the case that the options indicate the properties are the same, we can simply +swap the children without having to recompute the properties. In the other case, +we create a new node from scratch. + +To ensure that this works correctly, it is recommended that users also look +through their codebase and ensure that they use `replace_children_if_necessary` +for these changes — `replace_children_if_necessary` should be preferred over +manual use of `replace_children`, since `replace_children_if_necessary` will +call `replace_children` with the correct options filled in. + +See [PR #23903](https://github.com/apache/datafusion/pull/23903) for details. + +### `ListingOptions::target_partitions` and `collect_stat` removed + +The `target_partitions` and `collect_stat` fields on +`datafusion_catalog_listing::ListingOptions`, their builder methods +(`with_target_partitions`, `with_collect_stat`), and the +`with_session_config_options` helper have been removed. + +`ListingTable` now reads both values directly from the active `SessionConfig` +at scan time instead of from a copy snapshotted onto the table at construction +time. + +**Who is affected:** + +- Code that set `target_partitions` / `collect_stat` per table via + `ListingOptions`, or read those public fields. +- Code that relied on a `ListingTable` freezing these values at construction + time independently of the session config. The table now always reflects the + current `SessionConfig`. + +**Migration guide:** + +Configure these on the `SessionConfig` instead: + +```rust,ignore +// Before +let options = ListingOptions::new(format) + .with_target_partitions(8) + .with_collect_stat(true); + +// After +let config = SessionConfig::new() + .with_target_partitions(8) + .with_collect_statistics(true); +``` + +See [PR #22969](https://github.com/apache/datafusion/pull/22969) for details. + +### Spark map functions now reject duplicate keys by default + +The Spark-compatibility map-construction functions (`map_from_arrays`, +`map_from_entries`, `str_to_map`) now raise `[DUPLICATED_MAP_KEY]` at runtime +when constructing a map that contains duplicate keys. This matches the default +of Spark's [`spark.sql.mapKeyDedupPolicy`](https://github.com/apache/spark/blob/v4.0.0/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4502-L4511). + +A new config option, `datafusion.spark.map_key_dedup_policy`, controls the +behavior: + +- `EXCEPTION` (default): raise on any duplicate key. +- `LAST_WIN`: keep the last occurrence of each duplicate key. The key stays at + its first-seen position with the value from its last occurrence (matching + Spark's `ArrayBasedMapBuilder`). + +**Who is affected:** + +- Queries calling `map_from_arrays` or `str_to_map` on data that contains + duplicate keys. Previously these functions either tolerated duplicates + silently or raised a non-configurable error. + +**Migration guide:** + +To restore lenient duplicate-key handling, set the policy to `LAST_WIN`: + +```sql +SET datafusion.spark.map_key_dedup_policy = 'LAST_WIN'; +``` + +See [PR #21720](https://github.com/apache/datafusion/pull/21720) for details. + +### Unify LRU memory-limiting caches into one generic cache + +The caches `DefaultFileMetadataCache`, `DefaultListFilesCache` and `DefaultFileStatisticsCache` +are merged into one generic implementation `DefaultCache`. The corresponding traits are now +type aliases: + +```diff +- pub trait FileStatisticsCache: CacheAccessor +- pub trait ListFilesCache: CacheAccessor +- pub trait FileMetadataCache: CacheAccessor ++ pub type FileStatisticsCache = dyn Cache; ++ pub type ListFilesCache = dyn Cache; ++ pub type FileMetadataCache = dyn Cache; +``` + +**Who is affected:** + +- Users who introduced their own implementation of `FileMetadataCache`, `ListFilesCache` or `FileStatisticsCache`. + +**Migration guide:** + +Implement the newly introduced types for your custom cache implementation. + +See [PR #22613](https://github.com/apache/datafusion/pull/22613) for details. + +### `CachedFileMetadata` now validates file schema + +The file-statistics cache remains keyed by `TableScopedPath`, but +`CachedFileMetadata` now stores a `SchemaFingerprint` of the `file_schema` used +to compute the cached statistics. Cache hits are valid only when both the file +metadata and schema fingerprint match. + +**Who is affected:** + +- Users constructing `CachedFileMetadata` values directly. + +**Migration guide:** + +- Pass `Arc::new(SchemaFingerprint::from_schema(file_schema))` to + `CachedFileMetadata::new`. +- Pass the current schema fingerprint to `CachedFileMetadata::is_valid_for`. + +See [PR #23201](https://github.com/apache/datafusion/pull/23201) for details. + +### `EmptyExecNode` and `PlaceholderRowExecNode` gained a `partitions` field + +The generated protobuf structs `EmptyExecNode` and `PlaceholderRowExecNode` +encoded only a schema, so the partition count set by `EmptyExec::with_partitions` +was silently dropped when a physical plan was serialized and deserialized: a plan +that reported `n` partitions before encoding reported `1` after. Both messages now +carry a `partitions` field that round-trips the count. + +**Who is affected:** + +- Users constructing `EmptyExecNode` or `PlaceholderRowExecNode` with an + exhaustive struct literal. + +**Migration guide:** + +Set the new field, or fill it from `Default`: + +```rust,ignore +// Before +EmptyExecNode { schema: Some(schema) } + +// After +EmptyExecNode { schema: Some(schema), partitions: 4 } +// or +EmptyExecNode { schema: Some(schema), ..Default::default() } +``` + +The wire format stays compatible in both directions. Plans encoded before this +field existed decode as a single partition, the previous default, and plans +encoded after it add a field that older readers ignore. + +See [PR #23643](https://github.com/apache/datafusion/pull/23643) for details. + +### `time ± interval` now returns a `time` instead of an `interval` + +Adding or subtracting an `interval` to/from a `time` value now returns a `time` +that wraps within the 24-hour clock, matching PostgreSQL and DuckDB. Previously +DataFusion returned an `interval`. + +```sql +-- 55.0.0 onwards: returns a time +SELECT time '23:30:00' + interval '2 hours'; +-- 01:30:00 +``` + +Only the sub-day portion of the interval affects the result; whole days and +months are ignored, as in PostgreSQL. The result keeps the input time's unit +(mirroring `timestamp + interval`), and any interval precision finer than that +unit is truncated -- so `time(s) + interval '1 nanosecond'` is a no-op. + +See [PR #23279](https://github.com/apache/datafusion/pull/23279) for details. + +### Physical-planning state moved to an explicit `PhysicalPlanningContext` + +The `subquery_indexes` and `subquery_results` public fields on +`datafusion_expr::execution_props::ExecutionProps` have been removed. They were +added in `54.0.0` as the channel through which the physical planner passed +uncorrelated scalar-subquery state to functions that create physical +`Arc` values from logical `Expr` values. + +The `lambda_variable_qualifier` public field and the +`with_qualified_lambda_variables` method on `ExecutionProps` have been removed +for the same reason: they carried the qualifiers of the lambda variables in +scope while `create_physical_expr` descended into a lambda body. + +That state is now carried by a dedicated +`datafusion_expr::physical_planning_context::PhysicalPlanningContext` passed explicitly +through functions and planner traits. Unlike `ExecutionProps`, which applies +throughout the planning of an entire query, this context is scoped to the +logical plan subtree currently being converted. This removes the need for the +physical planner to clone and mutate a `SessionState`, is a prerequisite for +letting the planner take `&dyn Session`, and lets `ExtensionPlanner` +implementations create physical +expressions containing scalar subqueries against the same subquery state as the +rest of the plan. + +The following functions take a new trailing +`planning_ctx: &PhysicalPlanningContext` parameter: + +- `datafusion_physical_expr::create_physical_expr` / `create_physical_exprs` +- `datafusion_physical_expr::create_physical_sort_expr` / + `create_physical_sort_exprs` / `create_physical_partitioning` +- `datafusion::physical_planner::create_window_expr` / + `create_window_expr_with_name` +- `datafusion_physical_expr::aggregate::LoweredAggregateBuilder::new` + +The planner traits changed accordingly: + +- `PhysicalPlanner::create_physical_expr` takes + `planning_ctx: &PhysicalPlanningContext` +- `ExtensionPlanner::plan_extension` and `plan_table_scan` receive + `planning_ctx: &PhysicalPlanningContext` and should forward it to + `PhysicalPlanner::create_physical_expr` when creating physical expressions + +Convenience methods such as `SessionContext::create_physical_expr` and +`SessionState::create_physical_expr` are unchanged. + +**Who is affected:** + +- Code calling the functions above: pass + `&PhysicalPlanningContext::default()` unless you are creating physical + expressions as part of a physical plan that contains uncorrelated scalar + subqueries. +- Custom `PhysicalPlanner` or `ExtensionPlanner` implementations: add the new + parameter and forward it. +- Code that read or wrote `execution_props.subquery_indexes` / + `execution_props.subquery_results`: build a `PhysicalPlanningContext` instead. +- Code that read `execution_props.lambda_variable_qualifier` or called + `ExecutionProps::with_qualified_lambda_variables`: remove that usage. Callers + that only plan a `HigherOrderFunction` are not affected -- + `create_physical_expr` populates the lambda qualifiers itself as it descends + into lambda bodies. Code that needs to read or extend the lambda + scope should use the equivalents on `PhysicalPlanningContext`: + `PhysicalPlanningContext::lambda_variable_qualifier` and + `PhysicalPlanningContext::with_qualified_lambda_variables`. + +**Migration guide:** + +When creating a physical expression outside of physical planning, pass an empty +context: + +```rust,ignore +use datafusion_expr::physical_planning_context::PhysicalPlanningContext; +use datafusion_physical_expr::create_physical_expr; + +// Before +let phys = create_physical_expr(&expr, &schema, &props)?; + +// After +let phys = create_physical_expr( + &expr, + &schema, + &props, + &PhysicalPlanningContext::default(), +)?; +``` + +For `ExtensionPlanner` implementations, accept and forward the context: + +```rust,ignore +async fn plan_extension( + &self, + planner: &dyn PhysicalPlanner, + node: &dyn UserDefinedLogicalNode, + logical_inputs: &[&LogicalPlan], + physical_inputs: &[Arc], + session: &dyn Session, + planning_ctx: &PhysicalPlanningContext, // new parameter +) -> Result>> { + for expr in node.expressions() { + // Forward the context so scalar subqueries in this node's + // expressions resolve against the plan's subquery state + planner.create_physical_expr(&expr, node.schema(), session, planning_ctx)?; + } + // ... +} +``` + +See [PR #23649](https://github.com/apache/datafusion/pull/23649) and +[PR #23989](https://github.com/apache/datafusion/pull/23989) for details. + +### Catalog, planner, and optimizer contracts moved to `datafusion-session` + +The catalog, planner, and physical optimizer contract traits now live in the +`datafusion-session` crate. This makes them available through `Session` without +downcasting to `SessionState`, including across the FFI boundary. + +The moved catalog traits are `CatalogProviderList`, `CatalogProvider`, +`SchemaProvider`, `TableProvider`, `TableProviderFactory`, and +`TableFunctionImpl`. The related `TableFunction` struct also moved. The +`datafusion-catalog` crate re-exports these items from their new location, so +paths such as `datafusion::catalog::TableProvider` and +`datafusion_catalog::CatalogProvider` continue to work unchanged. + +The moved planning and optimization traits are `QueryPlanner`, +`PhysicalPlanner`, `ExtensionPlanner`, `PhysicalOptimizerRule`, and +`PhysicalOptimizerContext`. Their previous paths also continue to work through +re-exports: + +- `datafusion::execution::context::QueryPlanner` +- `datafusion::physical_planner::{PhysicalPlanner, ExtensionPlanner}` +- `datafusion_physical_optimizer::{PhysicalOptimizerRule, PhysicalOptimizerContext}` + +The session argument for methods on `QueryPlanner`, `PhysicalPlanner`, and +`ExtensionPlanner` changed from `&SessionState` to `&dyn Session`. Custom planner +implementations should update their signatures. Planner code should use methods +on `Session` instead of downcasting it to `SessionState`. + +The `Session` trait now requires a `catalog_list` method that returns the +catalogs registered with the session: + +```rust +fn catalog_list(&self) -> Arc; +``` + +Custom `Session` implementations must add this method. Implementations that do +not expose a catalog can return the new `EmptyCatalogProviderList`: + +```rust +use std::sync::Arc; +use datafusion_session::{CatalogProviderList, EmptyCatalogProviderList}; + +fn catalog_list(&self) -> Arc { + Arc::new(EmptyCatalogProviderList) +} +``` + +`Session` gains a `query_planner` method alongside `optimize`, +`physical_optimizers`, and `statistics_registry`. All four have default +implementations, so existing `Session` implementations that do not perform +physical planning require no changes: `query_planner` defaults to the new +`UnsupportedQueryPlanner`, `optimize` returns the plan unchanged, +`physical_optimizers` returns no rules, and `statistics_registry` returns +`None`. + +A custom session that drives planning through `DefaultQueryPlanner` or +`DefaultPhysicalPlanner` must override these methods to expose its planning and +optimization behavior; the defaults will otherwise produce unoptimized plans or +fail to plan at all. The simplest approach is to delegate to a `SessionState`: + +```rust +use std::sync::Arc; +use datafusion_session::{PhysicalOptimizerRule, QueryPlanner}; + +fn query_planner(&self) -> Arc { + self.inner.query_planner() +} + +fn optimize(&self, plan: &LogicalPlan) -> Result { + self.inner.optimize(plan) +} + +fn physical_optimizers(&self) -> &[Arc] { + self.inner.physical_optimizers() +} +``` + +`ForeignSession::create_physical_plan` runs the complete planning pipeline in the +library that owns the session. `ForeignSession::query_planner`, `optimize`, and +`physical_optimizers` forward to the owning session across the FFI boundary. A +foreign query planner can also be installed on a session through the new +`datafusion_ffi::query_planner::FFI_QueryPlanner`; see that module's +documentation for how plans and extension codecs cross the boundary. + +See [PR #23703](https://github.com/apache/datafusion/pull/23703) for details on +the catalog changes. + +### `FFI_LogicalExtensionCodec::task_ctx_provider` is now private + +The `task_ctx_provider` field on +`datafusion_ffi::proto::logical_extension_codec::FFI_LogicalExtensionCodec` was +`pub` and is now crate-private, matching `FFI_PhysicalExtensionCodec`. + +**Who is affected:** + +- Code that read or cloned `FFI_LogicalExtensionCodec::task_ctx_provider` + directly. Pass the task context provider to `FFI_LogicalExtensionCodec::new` + instead, and keep your own copy if you need it elsewhere. + +### Unused `async` removed from several public functions + +Public functions that were declared `async` but never awaited anything are now +synchronous: + +- `CsvFormat::read_to_delimited_chunks_from_stream` (in + `datafusion_datasource_csv`, re-exported as + `datafusion::datasource::file_format::csv::CsvFormat`) +- `datafusion_substrait::serializer::deserialize_bytes`, which now also borrows + its input as `&[u8]` instead of taking an owned `Vec` +- `datafusion::test_util::parquet::TestParquetFile::create_scan` + +**Migration guide:** + +Remove `.await` from call sites; the compiler flags each one, since `.await` +on a non-future value does not compile: + +```rust,ignore +// Before +let stream = csv_format + .read_to_delimited_chunks_from_stream(input) + .await; +let plan = deserialize_bytes(proto_bytes).await?; + +// After +let stream = csv_format.read_to_delimited_chunks_from_stream(input); +let plan = deserialize_bytes(&proto_bytes)?; +``` + +### `MovingMin` and `MovingMax` changed to `pub(crate)` + +`MovingMin` and `MovingMax` in `datafusion_functions_aggregate::min_max` have been changed from `pub` to `pub(crate)` visibility as they are internal helper data structures for DataFusion's sliding window aggregators. + +**Who is affected:** + +- Code that directly imported `MovingMin` or `MovingMax` from `datafusion_functions_aggregate`. Standard SQL window functions (`MIN(...) OVER (...)` / `MAX(...) OVER (...)`) are unaffected. + +See [PR #23827](https://github.com/apache/datafusion/pull/23827) for details. + +### `ExecutionPlan::apply_expressions` is now a required method + +`apply_expressions` has been added as a **required** method on the `ExecutionPlan`, `FileSource`, and `DataSource` traits. Any custom implementation of +these traits must now implement `apply_expressions`. See docs on `ExecutionPlan::apply_expressions` for migration details. + +### `WindowExpr::evaluate_stateful` now takes a `WindowEvalContext` + +`WindowExpr::evaluate_stateful` (and the provided +`AggregateWindowExpr::aggregate_evaluate_stateful` method) take a new +`WindowEvalContext` argument carrying stream-level information that is shared +by all partitions: + +```rust,ignore +// Before +fn evaluate_stateful( + &self, + partition_batches: &PartitionBatches, + window_agg_state: &mut PartitionWindowAggStates, +) -> Result<()> + +// After +fn evaluate_stateful( + &self, + partition_batches: &PartitionBatches, + window_agg_state: &mut PartitionWindowAggStates, + eval_ctx: &WindowEvalContext<'_>, +) -> Result<()> +``` + +`WindowEvalContext` currently carries the most recent input row, which +previously lived in each partition's `PartitionBatchState` (see the next +section). The struct is `#[non_exhaustive]` so that fields can be added +without further signature changes: construct it with +`WindowEvalContext::default()` and set fields through its builder methods. + +**Who is affected:** + +- Implementations of the `WindowExpr` trait that override `evaluate_stateful` + must add the new parameter. +- Callers of `evaluate_stateful` or `aggregate_evaluate_stateful` must pass a + context. + +**Migration guide:** + +```rust,ignore +use datafusion_physical_expr::window::WindowEvalContext; + +// Before +window_expr.evaluate_stateful(&partition_batches, &mut window_agg_state)?; + +// After +let eval_ctx = WindowEvalContext::default() + .with_most_recent_row(most_recent_row.as_ref()); +window_expr.evaluate_stateful( + &partition_batches, + &mut window_agg_state, + &eval_ctx, +)?; +``` + +Pass `WindowEvalContext::default()` when no most-recent-row watermark is +available (for example, when the input is sorted by the partition keys and +partition ends are detected directly). + +### `PartitionBatchState::most_recent_row` removed + +The `most_recent_row` field and the `set_most_recent_row` method have been +removed from `datafusion_expr::window_state::PartitionBatchState`. The most +recent input row is a property of the whole input stream rather than +per-partition state: every partition observed the same value. It is now +tracked once by the operator driving the evaluation and passed to window +expressions through the new `WindowEvalContext` argument of +`WindowExpr::evaluate_stateful` described above. + +**Who is affected:** + +- Code that read `PartitionBatchState::most_recent_row` or called + `set_most_recent_row`, such as custom streaming window operators. + +**Migration guide:** + +Track the most recent input row once per stream (for example, a one-row +slice of the last non-empty input batch) and pass it to window expressions +via `WindowEvalContext::with_most_recent_row` instead of copying it into +each partition's state. + +### `MSRV` updated to 1.94.0 + +The Minimum Supported Rust Version (MSRV) has been updated to [`1.94.0`]. + +[`1.94.0`]: https://releases.rs/docs/1.94.0/ + +### `CachedParquetFileReader` removed; `ParquetFileReader` fields are now private + +`CachedParquetFileReader` duplicated `ParquetFileReader` and has been removed; +`ParquetFileReader`'s fields are also now private, with +`file_metrics()` and `partitioned_file()` accessors added for the two that +were previously public. + +**Who is affected:** + +- Code that names the `CachedParquetFileReader` type. +- Code that constructs a `ParquetFileReader` directly via a struct literal, or + reads/writes its fields. + +**Migration guide:** + +`ParquetFileReader::new` is no longer public; build a reader through +`ParquetFileReaderFactory::create_reader` (via `DefaultParquetFileReaderFactory` +or `CachedParquetFileReaderFactory`) instead of constructing one directly: + +```rust,ignore +// Before +let inner = ParquetObjectReader::new(Arc::clone(&store), location).with_file_size(size); +let reader = CachedParquetFileReader::new( + file_metrics, + store, + inner, + partitioned_file, + metadata_cache, + metadata_size_hint, +); + +// After +let reader = CachedParquetFileReaderFactory::new(store, metadata_cache) + .create_reader(partition_index, partitioned_file, metadata_size_hint, &metrics)?; +``` + +Replace field access with the new accessor methods: + +```rust,ignore +// Before +let bytes_scanned = reader.file_metrics.bytes_scanned.value(); +let location = &reader.partitioned_file.object_meta.location; + +// After +let bytes_scanned = reader.file_metrics().bytes_scanned.value(); +let location = &reader.partitioned_file().object_meta.location; +``` + +### `array_distance` scalar function now rejects multidimensional arrays + +`array_distance` only supports one-dimensional arrays. Previously, when given +multidimensional arrays, it computed the distance using only the first +subarray and ignored the remaining subarrays. For example: + +```sql +SELECT array_distance( + [[1, 2], [100, 100]], + [[1, 4], [0, 0]] +); +``` + +Previously, this query returned `2.0`, the distance between `[1, 2]` and +`[1, 4]`. It now returns a planning error stating that `array_distance` only +supports one-dimensional arrays. + +### `ParquetObjectReader` / `ParquetObjectWriter` deprecated upstream + +The [`parquet` crate] deprecated [`ParquetObjectReader`] +and [`ParquetObjectWriter`] in favor of implementing +[`AsyncFileReader`] directly (see the example on the [`AsyncFileReader`] trait and +[`parquet/examples/object_store.rs`] in `arrow-rs`) or passing an +[`BufWriter`] straight to [`AsyncArrowWriter`]. + +**Who is affected:** + +- Custom [`ParquetFileReaderFactory`] implementations that construct a + [`ParquetObjectReader`] directly and now see a deprecation warning after + upgrading the `parquet` dependency. + +**Migration guide:** + +If your [`AsyncFileReader`] implementation exists mainly to read from an +[`ObjectStore`] and track metrics, consider using DataFusion's +[`ParquetFileReader`] instead of wrapping a [`ParquetObjectReader`]: + +```rust,ignore +// Before +let inner = ParquetObjectReader::new(store, location).with_file_size(size); +Ok(Box::new(MyReader { inner, file_metrics, partitioned_file })) + +// After +Ok(Box::new(ParquetFileReader { + file_metrics, + store, + metadata_size_hint, + partitioned_file, +})) +``` + +If you need custom behavior (I/O coalescing, byte caching, a dedicated I/O +runtime), implement `AsyncFileReader` directly against your `ObjectStore`, +following the pattern in [`parquet/examples/object_store.rs`] + +See [PR #24030](https://github.com/apache/datafusion/pull/24030) for details. + +[`parquet` crate]: https://crates.io/crates/parquet +[`parquetobjectreader`]: https://docs.rs/parquet/59.1.0/parquet/arrow/async_reader/struct.ParquetObjectReader.html +[`parquetobjectwriter`]: https://docs.rs/parquet/59.1.0/parquet/arrow/async_writer/struct.ParquetObjectWriter.html +[`parquetfilereader`]: https://docs.rs/datafusion/latest/datafusion/datasource/physical_plan/parquet/struct.ParquetFileReader.html +[`parquetfilereaderfactory`]: https://docs.rs/datafusion/latest/datafusion/datasource/physical_plan/parquet/trait.ParquetFileReaderFactory.html +[`asyncfilereader`]: https://docs.rs/parquet/59.1.0/parquet/arrow/async_reader/trait.AsyncFileReader.html +[`objectstore`]: https://docs.rs/object_store/latest/object_store/trait.ObjectStore.html +[`bufwriter`]: https://docs.rs/tokio/latest/tokio/io/struct.BufWriter.html +[`asyncarrowwriter`]: https://docs.rs/parquet/59.1.0/parquet/arrow/async_writer/struct.AsyncArrowWriter.html +[`parquet/examples/object_store.rs`]: https://github.com/apache/arrow-rs/blob/main/parquet/examples/object_store.rs + +### `datafusion-proto`: parquet options conversions are fallible + +`protobuf::ParquetOptions` and `protobuf::TableParquetOptions` validate +`writer_version` when converting into their `datafusion-common` counterparts, so +those conversions are `TryFrom` rather than `From`. + +Every other `From` / `TryFrom` conversion between DataFusion types and +`datafusion_proto::protobuf` messages is unchanged. Several impls moved to the +crate that owns their DataFusion type, but trait impls are global, so +`X::try_from(&proto)` and `proto.try_into()` still resolve with no import +changes. + +**Migration guide:** + +```rust,ignore +// Before +let opts = ParquetOptions::from(&proto_opts); +let table_opts = TableParquetOptions::from(&proto_table_opts); + +// After +let opts = ParquetOptions::try_from(&proto_opts)?; +let table_opts = TableParquetOptions::try_from(&proto_table_opts)?; +``` + +See [issue #24019](https://github.com/apache/datafusion/issues/24019) for details. diff --git a/docs/source/library-user-guide/upgrading/index.rst b/docs/source/library-user-guide/upgrading/index.rst index 1ed5eca2a5d2a..51c7f1413172b 100644 --- a/docs/source/library-user-guide/upgrading/index.rst +++ b/docs/source/library-user-guide/upgrading/index.rst @@ -21,6 +21,7 @@ Upgrade Guides .. toctree:: :maxdepth: 1 + DataFusion 55.0.0 <55.0.0> DataFusion 54.0.0 <54.0.0> DataFusion 53.0.0 <53.0.0> DataFusion 52.0.0 <52.0.0> diff --git a/docs/source/library-user-guide/working-with-exprs.md b/docs/source/library-user-guide/working-with-exprs.md index 472ab2481360e..2f15fa90610d9 100644 --- a/docs/source/library-user-guide/working-with-exprs.md +++ b/docs/source/library-user-guide/working-with-exprs.md @@ -167,7 +167,7 @@ In DataFusion, an `OptimizerRule` is a trait that supports rewriting `Expr`s tha We'll call our rule `AddOneInliner` and implement the `OptimizerRule` trait. The `OptimizerRule` trait has two methods: - `name` - returns the name of the rule -- `try_optimize` - takes a `LogicalPlan` and returns an `Option`. If the rule is able to optimize the plan, it returns `Some(LogicalPlan)` with the optimized plan. If the rule is not able to optimize the plan, it returns `None`. +- `rewrite` - takes a `LogicalPlan` and `&dyn OptimizerConfig`, and returns a `Result>`. If the rule is able to optimize the plan, it returns `Transformed::yes` with the optimized plan. If the rule is not able to optimize the plan, it returns `Transformed::no`. ```rust use std::sync::Arc; diff --git a/docs/source/user-guide/cli/datasources.md b/docs/source/user-guide/cli/datasources.md index 6b1a4887a8a0f..59a6b0aa43284 100644 --- a/docs/source/user-guide/cli/datasources.md +++ b/docs/source/user-guide/cli/datasources.md @@ -132,6 +132,30 @@ select count(*) from hits; 1 row in set. Query took 0.344 seconds. ``` +## Reading from standard input + +On Unix-like systems you can pipe data into the CLI and query it by pointing the +`LOCATION` at the `/dev/stdin` pseudo-file: + +```console +$ cat hits.csv | datafusion-cli -c " +CREATE EXTERNAL TABLE hits STORED AS CSV LOCATION '/dev/stdin' OPTIONS ('format.has_header' 'true'); +SELECT count(*) FROM hits;" +``` + +This works for CSV, JSON, and Parquet. Because standard input is not seekable +(and Parquet stores its metadata at the end of the file), the CLI buffers the +entire input into memory before querying it, so the data must fit in memory. +Standard input is read only once: the buffered contents are reused for any +further tables backed by `/dev/stdin` in the same session. Those tables must +declare the same `STORED AS` format as the first one; a differing format is +rejected with an error. + +The SQL must be passed with `-c`/`--command` or `-f`/`--file` so that standard +input is free to carry the data. In the interactive shell (and when SQL is +piped to the CLI without `-c`/`-f`) standard input carries the SQL itself, and +`LOCATION '/dev/stdin'` returns an error. + **Why Wildcards Are Not Supported** Although wildcards (e.g., _.parquet or \*\*/_.parquet) may work for local diff --git a/docs/source/user-guide/cli/functions.md b/docs/source/user-guide/cli/functions.md index ea353d5c8dcc8..baf054ef5a12c 100644 --- a/docs/source/user-guide/cli/functions.md +++ b/docs/source/user-guide/cli/functions.md @@ -168,6 +168,7 @@ The columns of the returned table are: | num_rows | Utf8 | Number of rows in the table | | num_columns | UInt64 | Number of columns in the table | | table_size_bytes | Utf8 | Size of the table, in bytes | +| hits | UInt64 | Number of times the cached file statistics has been accessed | | statistics_size_bytes | UInt64 | Size of the cached statistics in memory | ## `list_files_cache` @@ -200,13 +201,15 @@ location 's3://overturemaps-us-west-2/release/2025-12-17.0/theme=base/type=infra ``` The columns of the returned table are: -| column_name | data_type | Description | -| ------------------- | ------------ | ----------------------------------------------------------------------------------------- | -| table | Utf8 | Name of the table | -| path | Utf8 | File path relative to the object store / filesystem root | -| metadata_size_bytes | UInt64 | Size of the cached metadata in memory (not its thrift encoded form) | -| expires_in | Duration(ms) | Last modified time of the file | -| metadata_list | List(Struct) | List of metadatas, one for each file under the path. | + +| column_name | data_type | Description | +| ------------------- | ------------ | ------------------------------------------------------------------- | +| table | Utf8 | Name of the table | +| path | Utf8 | File path relative to the object store / filesystem root | +| metadata_size_bytes | UInt64 | Size of the cached metadata in memory (not its thrift encoded form) | +| expires_in | Duration(ms) | Last modified time of the file | +| hits | UInt64 | Number of times the cached metadata has been accessed | +| metadata_list | List(Struct) | List of metadatas, one for each file under the path. | A metadata struct in the metadata_list contains the following fields: diff --git a/docs/source/user-guide/concepts-readings-events.md b/docs/source/user-guide/concepts-readings-events.md index 712f54a046123..c5d2a6486f2c2 100644 --- a/docs/source/user-guide/concepts-readings-events.md +++ b/docs/source/user-guide/concepts-readings-events.md @@ -200,9 +200,20 @@ This is a list of DataFusion related blog posts, articles, and other resources. - **2025-02-02** [Apache DataFusion Ballista 43.0.0 Released](https://datafusion.apache.org/blog/2025/02/02/datafusion-ballista-43.0.0) - **2025-01-17** [Apache DataFusion Comet 0.5.0 Release](https://datafusion.apache.org/blog/2025/01/17/datafusion-comet-0.5.0) +# 🎥 Community Showcase + +The [DataFusion Community Showcase](https://github.com/apache/datafusion/issues/22963) is a +regular virtual event where community members share what they are building with DataFusion. + +- **2026-08-06** [Vol. 3: ASAPQuery (Milind Srivastava) & Streamling (Yaroslav Tkachenko)](https://www.youtube.com/watch?v=0-BIHyzODH8) +- **2026-07-23** [Vol. 2: DataFusion Comet (Jordan Epstein) & DataFusion Ballista (Phillip LeBlanc)](https://www.youtube.com/watch?v=G8In--2RUwI) +- **2026-07-09** [Vol. 1: SedonaDB (Dewey Dunnington) & Xarray-SQL (Alex Merose)](https://www.youtube.com/watch?v=5o-4hL8vGPw) + # 🌎 Community Events +- **2026-09-03** [Boston Apache DataFusion Meetup](https://github.com/apache/datafusion/discussions/21541) - [RSVP](https://luma.com/yexgqifv) - **2026-07-22** [Denver Apache DataFusion Meetup](https://github.com/apache/datafusion/discussions/18428) - [RSVP](https://luma.com/jsu6faie) +- **2026-06-28** [Shanghai Apache DataFusion Meetup](https://github.com/apache/datafusion/discussions/16334) - [RSVP](https://luma.com/7xrhm9rx), [LinkedIn](https://www.linkedin.com/posts/ruihang-xia_we-are-going-to-have-a-apache-datafusion-share-7473348653169160194-NcmY) - **2026-05-12** [New York City Apache DataFusion Meetup](https://github.com/apache/datafusion/discussions/20030) - [RSVP](https://luma.com/adhshv92) - **2026-05-11** [San Francisco Apache DataFusion Meetup](https://github.com/apache/datafusion/discussions/21638) - [RSVP](https://luma.com/k3ointcl) - **2026-04-23** [Seattle Apache DataFusion Meetup](https://github.com/apache/datafusion/discussions/13500) - [RSVP](https://luma.com/hxshbp0m) diff --git a/docs/source/user-guide/configs.md b/docs/source/user-guide/configs.md index 576137bda29d1..e02ada03fc413 100644 --- a/docs/source/user-guide/configs.md +++ b/docs/source/user-guide/configs.md @@ -77,7 +77,7 @@ The following configuration settings are available: | datafusion.execution.perfect_hash_join_small_build_threshold | 1024 | A perfect hash join (see `HashJoinExec` for more details) will be considered if the range of keys (max - min) on the build side is < this threshold. This provides a fast path for joins with very small key ranges, bypassing the density check. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. | | datafusion.execution.perfect_hash_join_min_key_density | 0.15 | The minimum required density of join keys on the build side to consider a perfect hash join (see `HashJoinExec` for more details). Density is calculated as: `(number of rows) / (max_key - min_key + 1)`. A perfect hash join may be used if the actual key density > this value. Currently only supports cases where build_side.num_rows() < u32::MAX. Support for build_side.num_rows() >= u32::MAX will be added in the future. | | datafusion.execution.coalesce_batches | true | When set to true, record batches will be examined between each operator and small batches will be coalesced into larger batches. This is helpful when there are highly selective filters or joins that could produce tiny output batches. The target batch size is determined by the configuration setting | -| datafusion.execution.collect_statistics | true | Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Applies to the default `ListingTableProvider` in DataFusion. Defaults to true. | +| datafusion.execution.collect_statistics | true | Should DataFusion collect statistics when first creating a table. Has no effect after the table is created. Defaults to true. | | datafusion.execution.target_partitions | 0 | Number of partitions for query execution. Increasing partitions can increase concurrency. Defaults to the number of CPU cores on the system | | datafusion.execution.time_zone | NULL | The default time zone Some functions, e.g. `now` return timestamps in this time zone | | datafusion.execution.parquet.enable_page_index | true | (reading) If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce the I/O and number of rows decoded. | @@ -93,6 +93,7 @@ The following configuration settings are available: | datafusion.execution.parquet.coerce_int96_tz | NULL | (reading) Optional timezone applied to INT96 columns when `coerce_int96` is set. When `Some`, INT96 columns coerce to `Timestamp(, Some())` instead of the default `Timestamp(, None)`. Spark and other systems write INT96 values as UTC-adjusted instants, so callers that need the resulting Arrow type to be timezone-aware (e.g. for Spark `TimestampType` semantics) should set this to `"UTC"`. No effect when `coerce_int96` is `None`. | | datafusion.execution.parquet.bloom_filter_on_read | true | (reading) Use any available bloom filters when reading parquet files | | datafusion.execution.parquet.max_predicate_cache_size | NULL | (reading) The maximum predicate cache size, in bytes. When `pushdown_filters` is enabled, sets the maximum memory used to cache the results of predicate evaluation between filter evaluation and output generation. Decreasing this value will reduce memory usage, but may increase IO and CPU usage. None means use the default parquet reader setting. 0 means no caching. | +| datafusion.execution.parquet.max_in_list_size | 20 | Maximum number of values in an `IN (...)` list for which pruning will occur. Longer lists will not be used to prune files, row groups, or data pages. Higher values help in cases such as filtering on a list of ~25-100 identifiers, but also make the predicate more expensive to evaluate. Set to 0 to disable `IN (...)` list pruning entirely. Defaults to 20. | | datafusion.execution.parquet.data_pagesize_limit | 1048576 | (writing) Sets best effort maximum size of data page in bytes | | datafusion.execution.parquet.write_batch_size | 1024 | (writing) Sets write_batch_size in rows | | datafusion.execution.parquet.writer_version | 1.0 | (writing) Sets parquet writer version valid values are "1.0" and "2.0" | @@ -101,8 +102,9 @@ The following configuration settings are available: | datafusion.execution.parquet.dictionary_enabled | true | (writing) Sets if dictionary encoding is enabled. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.dictionary_page_size_limit | 1048576 | (writing) Sets best effort maximum dictionary page size, in bytes | | datafusion.execution.parquet.statistics_enabled | page | (writing) Sets if statistics are enabled for any column Valid values are: "none", "chunk", and "page" These values are not case sensitive. If NULL, uses default parquet writer setting | -| datafusion.execution.parquet.max_row_group_size | 1048576 | (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. | -| datafusion.execution.parquet.created_by | datafusion version 53.1.0 | (writing) Sets "created by" property | +| datafusion.execution.parquet.max_row_group_size | 1048576 | (writing) Target maximum number of rows in each row group (defaults to 1M rows). Writing larger row groups requires more memory to write, but can get better compression and be faster to read. When `max_row_group_bytes` is also set, the writer flushes a row group when either limit is reached, whichever comes first. | +| datafusion.execution.parquet.max_row_group_bytes | NULL | (writing) Target maximum size of each row group in bytes. When set, the writer flushes whenever either this limit or `max_row_group_size` is reached, whichever comes first. Useful for bounding writer memory on wide schemas where a row-count limit can map to very different byte sizes. Matches the behavior of `parquet.block.size` in parquet-mr. If `None` (the default), only the row-count limit applies. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores this limit. | +| datafusion.execution.parquet.created_by | datafusion version 55.0.0 | (writing) Sets "created by" property | | datafusion.execution.parquet.column_index_truncate_length | 64 | (writing) Sets column index truncate length | | datafusion.execution.parquet.statistics_truncate_length | 64 | (writing) Sets statistics truncate length. If NULL, uses default parquet writer setting | | datafusion.execution.parquet.data_page_row_count_limit | 20000 | (writing) Sets best effort maximum number of rows in data page | @@ -113,9 +115,13 @@ The following configuration settings are available: | datafusion.execution.parquet.allow_single_file_parallelism | true | (writing) Controls whether DataFusion will attempt to speed up writing parquet files by serializing them in parallel. Each column in each row group in each output file are serialized in parallel leveraging a maximum possible core count of n_files*n_row_groups*n_columns. | | datafusion.execution.parquet.maximum_parallel_row_group_writers | 1 | (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. | | datafusion.execution.parquet.maximum_buffered_record_batches_per_stream | 2 | (writing) By default parallel parquet writer is tuned for minimum memory usage in a streaming execution plan. You may see a performance benefit when writing large parquet files by increasing maximum_parallel_row_group_writers and maximum_buffered_record_batches_per_stream if your system has idle cores and can tolerate additional memory usage. Boosting these values is likely worthwhile when writing out already in-memory data, such as from a cached data frame. | -| datafusion.execution.parquet.use_content_defined_chunking | NULL | (writing) EXPERIMENTAL: Enable content-defined chunking (CDC) when writing parquet files. When `Some`, CDC is enabled with the given options; when `None` (the default), CDC is disabled. When CDC is enabled, parallel writing is automatically disabled since the chunker state must persist across row groups. | +| datafusion.execution.parquet.content_defined_chunking.enabled | false | (writing) EXPERIMENTAL: Enable content-defined chunking (CDC) when writing parquet files. When enabled, parallel writing is automatically disabled since the chunker state must persist across row groups. | +| datafusion.execution.parquet.content_defined_chunking.min_chunk_size | 262144 | Minimum chunk size in bytes. The rolling hash will not trigger a split until this many bytes have been accumulated. Default is 256 KiB. | +| datafusion.execution.parquet.content_defined_chunking.max_chunk_size | 1048576 | Maximum chunk size in bytes. A split is forced when the accumulated size exceeds this value. Default is 1 MiB. | +| datafusion.execution.parquet.content_defined_chunking.norm_level | 0 | Normalization level. Increasing this improves deduplication ratio but increases fragmentation. Recommended range is [-3, 3], default is 0. | | datafusion.execution.planning_concurrency | 0 | Fan-out during initial physical planning. This is mostly use to plan `UNION` children in parallel. Defaults to the number of CPU cores on the system | | datafusion.execution.skip_physical_aggregate_schema_check | false | When set to true, skips verifying that the schema produced by planning the input of `LogicalPlan::Aggregate` exactly matches the schema of the input plan. When set to false, if the schema does not match exactly (including nullability and metadata), a planning error will be raised. This is used to workaround bugs in the planner that are now caught by the new schema verification step. | +| datafusion.execution.enable_migration_aggregate | true | Temporary switch for aggregate stream implementations that are being migrated from `GroupedHashAggregateStream`. When set to true, DataFusion tries the migrated implementations when their preconditions are satisfied. When set to false, grouped aggregation falls back to `GroupedHashAggregateStream`. This option will be removed after the migration is finished. See for details. | | datafusion.execution.spill_compression | uncompressed | Sets the compression codec used when spilling data to disk. Since datafusion writes spill files using the Arrow IPC Stream format, only codecs supported by the Arrow IPC Stream Writer are allowed. Valid values are: uncompressed, lz4_frame, zstd. Note: lz4_frame offers faster (de)compression, but typically results in larger spill files. In contrast, zstd achieves higher compression ratios at the cost of slower (de)compression speed. | | datafusion.execution.sort_spill_reservation_bytes | 10485760 | Specifies the reserved memory for each spillable sort operation to facilitate an in-memory merge. When a sort operation spills to disk, the in-memory data must be sorted and merged before being written to a file. This setting reserves a specific amount of memory for that in-memory sort/merge process. Note: This setting is irrelevant if the sort operation cannot spill (i.e., if there's no `DiskManager` configured). | | datafusion.execution.sort_in_place_threshold_bytes | 1048576 | When sorting, below what size should data be concatenated and sorted in a single RecordBatch rather than sorted in batches and merged. | @@ -124,12 +130,13 @@ The following configuration settings are available: | datafusion.execution.meta_fetch_concurrency | 32 | Number of files to read in parallel when inferring schema and statistics | | datafusion.execution.minimum_parallel_output_files | 4 | Guarantees a minimum level of output files running in parallel. RecordBatches will be distributed in round robin fashion to each parallel writer. Each writer is closed and a new file opened once soft_max_rows_per_output_file is reached. | | datafusion.execution.soft_max_rows_per_output_file | 50000000 | Target number of rows in output files when writing multiple. This is a soft max, so it can be exceeded slightly. There also will be one file smaller than the limit if the total number of rows written is not roughly divisible by the soft max | -| datafusion.execution.max_buffered_batches_per_output_file | 2 | This is the maximum number of RecordBatches buffered for each output file being worked. Higher values can potentially give faster write performance at the cost of higher peak memory consumption | +| datafusion.execution.max_buffered_batches_per_output_file | 2 | This is the maximum number of RecordBatches buffered for each output file being worked. Higher values can potentially give faster write performance at the cost of higher peak memory consumption. This budget is split evenly between two independent points in the write pipeline (see the demuxer diagram in #7791): how many files can be in flight from the demuxer to a writer task, and how many RecordBatches are buffered for a single file's writer. Must be at least 2 so each half gets at least 1 unit of buffering - 0 or 1 would leave one side with a zero-capacity channel and panic at write time. | | datafusion.execution.listing_table_ignore_subdirectory | true | Should sub directories be ignored when scanning directories for data files. Defaults to true (ignores subdirectories), consistent with Hive. Note that this setting does not affect reading partitioned tables (e.g. `/table/year=2021/month=01/data.parquet`). | | datafusion.execution.listing_table_factory_infer_partitions | true | Should a `ListingTable` created through the `ListingTableFactory` infer table partitions from Hive compliant directories. Defaults to true (partition columns are inferred and will be represented in the table schema). | | datafusion.execution.enable_recursive_ctes | true | Should DataFusion support recursive CTEs | | datafusion.execution.split_file_groups_by_statistics | false | Attempt to eliminate sorts by packing & sorting files with non-overlapping statistics into the same file groups. Currently experimental | | datafusion.execution.keep_partition_by_columns | false | Should DataFusion keep the columns used for partition_by in the output RecordBatches | +| datafusion.execution.enable_file_stream_work_stealing | true | When `true` (the default), DataFusion's built-in file scans dynamically rebalance files across partitions at query execution time: a partition that goes idle reads files (or byte-range morsels) originally assigned to a sibling partition, which keeps all partitions busy in a single process. Executors that depend on the plan-time partition assignment — such as Ballista and datafusion-distributed, which run each partition as an isolated task and never poll the siblings — should set this to `false` so each partition reads only its own file group and no runtime reassignment occurs. | | datafusion.execution.skip_partial_aggregation_probe_ratio_threshold | 0.8 | Aggregation ratio (number of distinct groups / number of input rows) threshold for skipping partial aggregation. If the value is greater then partial aggregation will skip aggregation for further input | | datafusion.execution.skip_partial_aggregation_probe_rows_threshold | 100000 | Number of input rows partial aggregation partition should process, before aggregation ratio check and trying to switch to skipping aggregation mode | | datafusion.execution.use_row_number_estimates_to_optimize_partitioning | false | Should DataFusion use row number estimates at the input to decide whether increasing parallelism is beneficial or not. By default, only exact row numbers (not estimates) are used for this decision. Setting this flag to `true` will likely produce better plans. if the source of statistics is accurate. We plan to make this the default in the future. | @@ -144,16 +151,17 @@ The following configuration settings are available: | datafusion.optimizer.enable_window_topn | false | When set to true, the optimizer will replace Filter(rn<=K) → Window(ROW_NUMBER) → Sort patterns with a PartitionedTopKExec that maintains per-partition heaps, avoiding a full sort of the input. When the window partition key has low cardinality, enabling this optimization can improve performance. However, for high cardinality keys, it may cause regressions in both memory usage and runtime. | | datafusion.optimizer.enable_topk_repartition | true | When set to true, the optimizer will push TopK (Sort with fetch) below hash repartition when the partition key is a prefix of the sort key, reducing data volume before the shuffle. | | datafusion.optimizer.enable_topk_dynamic_filter_pushdown | true | When set to true, the optimizer will attempt to push down TopK dynamic filters into the file scan phase. | +| datafusion.optimizer.enable_physical_uncorrelated_scalar_subquery | true | When set to true, uncorrelated scalar subqueries are left in the logical plan and executed by `ScalarSubqueryExec` during physical execution. When set to false, all scalar subqueries (including uncorrelated ones) are rewritten to left joins by the `ScalarSubqueryToJoin` optimizer rule. Note disabling this option is not recommended. It restores pre behavior, which silently produces incorrect results for multi-row subqueries and does not support scalar subqueries in ORDER BY / JOIN ON / aggregate-function arguments. This option is intended as a temporary escape hatch for distributed execution frameworks and is planned to be removed in a future DataFusion release. | | datafusion.optimizer.enable_join_dynamic_filter_pushdown | true | When set to true, the optimizer will attempt to push down Join dynamic filters into the file scan phase. | | datafusion.optimizer.enable_aggregate_dynamic_filter_pushdown | true | When set to true, the optimizer will attempt to push down Aggregate dynamic filters into the file scan phase. | | datafusion.optimizer.enable_dynamic_filter_pushdown | true | When set to true attempts to push down dynamic filters generated by operators (TopK, Join & Aggregate) into the file scan phase. For example, for a query such as `SELECT * FROM t ORDER BY timestamp DESC LIMIT 10`, the optimizer will attempt to push down the current top 10 timestamps that the TopK operator references into the file scans. This means that if we already have 10 timestamps in the year 2025 any files that only have timestamps in the year 2024 can be skipped / pruned at various stages in the scan. The config will suppress `enable_join_dynamic_filter_pushdown`, `enable_topk_dynamic_filter_pushdown` & `enable_aggregate_dynamic_filter_pushdown` So if you disable `enable_topk_dynamic_filter_pushdown`, then enable `enable_dynamic_filter_pushdown`, the `enable_topk_dynamic_filter_pushdown` will be overridden. | | datafusion.optimizer.filter_null_join_keys | false | When set to true, the optimizer will insert filters before a join between a nullable and non-nullable column to filter out nulls on the nullable side. This filter can add additional overhead when the file format does not fully support predicate push down. | | datafusion.optimizer.repartition_aggregations | true | Should DataFusion repartition data using the aggregate keys to execute aggregates in parallel using the provided `target_partitions` level | -| datafusion.optimizer.repartition_file_min_size | 10485760 | Minimum total files size in bytes to perform file scan repartitioning. | +| datafusion.optimizer.repartition_file_min_size | 1048576 | Minimum total file size in bytes for file-group byte-range splitting to fire. Files (or merged file groups) smaller than this stay as one partition. Lower values produce more, smaller partitions — better at filling `target_partitions` worth of cores when files are modestly sized, at the cost of slightly more per-partition open / metadata-load overhead. | | datafusion.optimizer.repartition_joins | true | Should DataFusion repartition data using the join keys to execute joins in parallel using the provided `target_partitions` level | | datafusion.optimizer.allow_symmetric_joins_without_pruning | true | Should DataFusion allow symmetric hash joins for unbounded data sources even when its inputs do not have any ordering or filtering If the flag is not enabled, the SymmetricHashJoin operator will be unable to prune its internal buffers, resulting in certain join types - such as Full, Left, LeftAnti, LeftSemi, Right, RightAnti, and RightSemi - being produced only at the end of the execution. This is not typical in stream processing. Additionally, without proper design for long runner execution, all types of joins may encounter out-of-memory errors. | | datafusion.optimizer.repartition_file_scans | true | When set to `true`, datasource partitions will be repartitioned to achieve maximum parallelism. This applies to both in-memory partitions and FileSource's file groups (1 group is 1 partition). For FileSources, only Parquet and CSV formats are currently supported. If set to `true` for a FileSource, all files will be repartitioned evenly (i.e., a single large file might be partitioned into smaller chunks) for parallel scanning. If set to `false` for a FileSource, different files will be read in parallel, but repartitioning won't happen within a single file. If set to `true` for an in-memory source, all memtable's partitions will have their batches repartitioned evenly to the desired number of `target_partitions`. Repartitioning can change the total number of partitions and batches per partition, but does not slice the initial record tables provided to the MemTable on creation. | -| datafusion.optimizer.preserve_file_partitions | 0 | Minimum number of distinct partition values required to group files by their Hive partition column values (enabling Hash partitioning declaration). How the option is used: - preserve_file_partitions=0: Disable it. - preserve_file_partitions=1: Always enable it. - preserve_file_partitions=N, actual file partitions=M: Only enable when M >= N. This threshold preserves I/O parallelism when file partitioning is below it. Note: This may reduce parallelism, rooting from the I/O level, if the number of distinct partitions is less than the target_partitions. | +| datafusion.optimizer.preserve_file_partitions | 0 | Minimum number of distinct partition values required to group files by their Hive partition column values (enabling output partitioning declaration). How the option is used: - preserve_file_partitions=0: Disable it. - preserve_file_partitions=1: Always enable it. - preserve_file_partitions=N, actual file partitions=M: Only enable when M >= N. This threshold preserves I/O parallelism when file partitioning is below it. Note: This may reduce parallelism, rooting from the I/O level, if the number of distinct partitions is less than the target_partitions. | | datafusion.optimizer.repartition_windows | true | Should DataFusion repartition data using the partitions keys to execute window functions in parallel using the provided `target_partitions` level | | datafusion.optimizer.repartition_sorts | true | Should DataFusion execute sorts in a per-partition fashion and merge afterwards instead of coalescing first and sorting globally. With this flag is enabled, plans in the form below `text "SortExec: [a@0 ASC]", " CoalescePartitionsExec", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ` would turn into the plan below which performs better in multithreaded environments `text "SortPreservingMergeExec: [a@0 ASC]", " SortExec: [a@0 ASC]", " RepartitionExec: partitioning=RoundRobinBatch(8), input_partitions=1", ` | | datafusion.optimizer.subset_repartition_threshold | 4 | Partition count threshold for subset satisfaction optimization. When the current partition count is >= this threshold, DataFusion will skip repartitioning if the required partitioning expression is a subset of the current partition expression such as Hash(a) satisfies Hash(a, b). When the current partition count is < this threshold, DataFusion will repartition to increase parallelism even when subset satisfaction applies. Set to 0 to always repartition (disable subset satisfaction optimization). Set to a high value to always use subset satisfaction. Example (subset_repartition_threshold = 4): `text Hash([a]) satisfies Hash([a, b]) because (Hash([a, b]) is subset of Hash([a]) If current partitions (3) < threshold (4), repartition: AggregateExec: mode=FinalPartitioned, gby=[a, b], aggr=[SUM(x)] RepartitionExec: partitioning=Hash([a, b], 8), input_partitions=3 AggregateExec: mode=Partial, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 3) If current partitions (8) >= threshold (4), use subset satisfaction: AggregateExec: mode=SinglePartitioned, gby=[a, b], aggr=[SUM(x)] DataSourceExec: file_groups={...}, output_partitioning=Hash([a], 8) ` | @@ -187,7 +195,7 @@ The following configuration settings are available: | datafusion.sql_parser.parse_float_as_decimal | false | When set to true, SQL parser will parse float as decimal type | | datafusion.sql_parser.enable_ident_normalization | true | When set to true, SQL parser will normalize ident (convert ident to lowercase when not quoted) | | datafusion.sql_parser.enable_options_value_normalization | false | When set to true, SQL parser will normalize options value (convert value to lowercase). Note that this option is ignored and will be removed in the future. All case-insensitive values are normalized automatically. | -| datafusion.sql_parser.dialect | generic | Configure the SQL dialect used by DataFusion's parser; supported values include: Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, MsSQL, ClickHouse, BigQuery, Ansi, DuckDB and Databricks. | +| datafusion.sql_parser.dialect | generic | Configure the SQL dialect used by DataFusion's parser; supported values include: Generic, MySQL, PostgreSQL, Hive, SQLite, Snowflake, Redshift, MsSQL, ClickHouse, BigQuery, Ansi, DuckDB, Databricks, Spark. | | datafusion.sql_parser.support_varchar_with_length | true | If true, permit lengths for `VARCHAR` such as `VARCHAR(20)`, but ignore the length. If false, error if a `VARCHAR` with a length is specified. The Arrow type system does not have a notion of maximum string length and thus DataFusion can not enforce such limits. | | datafusion.sql_parser.map_string_types_to_utf8view | true | If true, string types (VARCHAR, CHAR, Text, and String) are mapped to `Utf8View` during SQL planning. If false, they are mapped to `Utf8`. Default is true. | | datafusion.sql_parser.collect_spans | false | When set to true, the source locations relative to the original SQL query (i.e. [`Span`](https://docs.rs/sqlparser/latest/sqlparser/tokenizer/struct.Span.html)) will be collected and recorded in the logical plan nodes. | @@ -203,6 +211,7 @@ The following configuration settings are available: | datafusion.format.time_format | %H:%M:%S%.f | Time format for time arrays | | datafusion.format.duration_format | pretty | Duration format. Can be either `"pretty"` or `"ISO8601"` | | datafusion.format.types_info | false | Show types in visual representation batches | +| datafusion.spark.map_key_dedup_policy | EXCEPTION | Policy for handling duplicate keys in Spark-compatible map-construction functions (`map_from_arrays`, `map_from_entries`, `str_to_map`). Mirrors Spark's [`spark.sql.mapKeyDedupPolicy`](https://github.com/apache/spark/blob/cf3a34e19dfcf70e2d679217ff1ba21302212472/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala#L4961): - `EXCEPTION` (default): raise `[DUPLICATED_MAP_KEY]` at runtime on any duplicate key. - `LAST_WIN`: keep the last occurrence of each duplicate key. Values are case-insensitive. | You can also reset configuration options to default settings via SQL using the `RESET` command. For example, to set and reset `datafusion.execution.batch_size`: @@ -236,6 +245,7 @@ The following runtime configuration settings are available: | datafusion.runtime.file_statistics_cache_limit | 20M | Maximum memory to use for file statistics cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | | datafusion.runtime.list_files_cache_limit | 1M | Maximum memory to use for list files cache. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | | datafusion.runtime.list_files_cache_ttl | NULL | TTL (time-to-live) of the entries in the list file cache. Supports units m (minutes), and s (seconds). Example: '2m' for 2 minutes. | +| datafusion.runtime.max_spill_merge_fan_in | 0 | Maximum number of spill files opened by one external merge pass. Use 0 for unlimited. Values below 2 still use 2 so a merge can make progress. | | datafusion.runtime.max_temp_directory_size | 100G | Maximum temporary file directory size. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | | datafusion.runtime.memory_limit | NULL | Maximum memory limit for query execution. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | | datafusion.runtime.metadata_cache_limit | 50M | Maximum memory to use for file metadata cache such as Parquet metadata. Supports suffixes K (kilobytes), M (megabytes), and G (gigabytes) or '0' for 0. Example: '2G' for 2 gigabytes. | diff --git a/docs/source/user-guide/crate-configuration.md b/docs/source/user-guide/crate-configuration.md index 92c0f37807c72..09c65107e58c8 100644 --- a/docs/source/user-guide/crate-configuration.md +++ b/docs/source/user-guide/crate-configuration.md @@ -156,7 +156,7 @@ By default, Datafusion returns errors as a plain text message. You can enable mo such as backtraces by enabling the `backtrace` feature to your `Cargo.toml` file like this: ```toml -datafusion = { version = "53.0.0", features = ["backtrace"]} +datafusion = { version = "55.0.0", features = ["backtrace"]} ``` Set environment [variables](https://doc.rust-lang.org/std/backtrace/index.html#environment-variables) diff --git a/docs/source/user-guide/example-usage.md b/docs/source/user-guide/example-usage.md index fd755715eec91..6f2419b9cd182 100644 --- a/docs/source/user-guide/example-usage.md +++ b/docs/source/user-guide/example-usage.md @@ -29,7 +29,7 @@ Find latest available Datafusion version on [DataFusion's crates.io] page. Add the dependency to your `Cargo.toml` file: ```toml -datafusion = "53.0.0" +datafusion = "55.0.0" tokio = { version = "1.0", features = ["rt-multi-thread"] } ``` diff --git a/docs/source/user-guide/explain-usage.md b/docs/source/user-guide/explain-usage.md index 9e06acbce4bd6..bc9dace297068 100644 --- a/docs/source/user-guide/explain-usage.md +++ b/docs/source/user-guide/explain-usage.md @@ -169,7 +169,7 @@ debugging to see why and when DataFusion added and removed operators from a plan During execution, DataFusion operators collect detailed metrics. You can access them programmatically via [`ExecutionPlan::metrics`] as well as with the -`EXPLAIN ANALYZE` command. For example here is the same query query as +`EXPLAIN ANALYZE` command. For example here is the same query as above but with `EXPLAIN ANALYZE` (note the output is edited for clarity) [`executionplan::metrics`]: https://docs.rs/datafusion/latest/datafusion/physical_plan/trait.ExecutionPlan.html#method.metrics @@ -241,6 +241,55 @@ When predicate pushdown is enabled, `DataSourceExec` with `ParquetSource` gains - `row_pushdown_eval_time`: time spent evaluating row-level filters - `page_index_eval_time`: time required to evaluate the page index filters +## Postgres-style `EXPLAIN (...)` options + +In addition to the legacy keyword form (`EXPLAIN ANALYZE VERBOSE FORMAT tree SELECT ...`), +DataFusion accepts a Postgres-style option list on dialects whose +[`supports_explain_with_utility_options`](https://docs.rs/sqlparser/latest/sqlparser/dialect/trait.Dialect.html#method.supports_explain_with_utility_options) +returns `true`. This includes the default `GenericDialect`, `PostgreSqlDialect`, and +`DuckDbDialect`, among others. + +```sql +EXPLAIN (ANALYZE, VERBOSE, METRICS 'rows,bytes', LEVEL dev) +SELECT ... ; +``` + +The recognized options are: + +| Option | Argument | Effect | +| --------- | ----------------- | ------------------------------------------------------------------------------------------------------------------------------------ | +| `ANALYZE` | boolean, optional | Execute the plan and collect metrics. Defaults to `TRUE` when bare. Equivalent to the `ANALYZE` keyword. | +| `VERBOSE` | boolean, optional | Show per-partition metrics and additional detail. Equivalent to the `VERBOSE` keyword. | +| `FORMAT` | identifier/string | One of `indent`, `tree`, `pgjson`, `graphviz`. Equivalent to the `FORMAT ` clause. | +| `METRICS` | string | Filter `ANALYZE` metrics by category. Accepts `'all'`, `'none'`, or any comma-separated subset of `rows,bytes,timing,uncategorized`. | +| `LEVEL` | identifier/string | `summary` or `dev`. Controls metric verbosity for `ANALYZE`. | +| `TIMING` | boolean | Sugar over `METRICS`: toggles inclusion of the `timing` category. | +| `SUMMARY` | boolean | Sugar over `LEVEL`: `TRUE` → `summary`, `FALSE` → `dev`. | +| `COSTS` | boolean | Include statistics in plain `EXPLAIN` output (equivalent to `SET datafusion.explain.show_statistics`). Not valid with `ANALYZE`. | + +Boolean arguments can be written bare (`ANALYZE` → `true`), as `TRUE`/`FALSE`, +`ON`/`OFF`, or `0`/`1`. + +When combined with `ANALYZE`, `FORMAT` supports `indent` (the default) and +`pgjson`; `tree` and `graphviz` are rejected. The `pgjson` form emits the +physical plan with live metrics, which is handy for plan visualizers — see +[`pgjson` format with `ANALYZE`](sql/explain.md#pgjson-format-with-analyze): + +```sql +EXPLAIN (ANALYZE, FORMAT pgjson) SELECT ...; +``` + +The statement-level options take precedence over session config, so you can leave +the session defaults alone and override just for the current query: + +```sql +EXPLAIN (ANALYZE, LEVEL dev, METRICS 'rows,bytes') SELECT ...; +``` + +Postgres options that DataFusion does not model (`BUFFERS`, `WAL`, `SETTINGS`, +`GENERIC_PLAN`, `MEMORY`) return a clear error rather than being silently +accepted — use `METRICS` to filter what appears in the output. + ## Partitions and Execution DataFusion determines the optimal number of cores to use as part of query @@ -316,7 +365,7 @@ For this query, let's again read the plan from the bottom to the top: - `gby=[UserID@0 as UserID]`: Represents `GROUP BY` in the [physical plan] and groups together the same values of `UserID`. - `aggr=[count(*)]`: Applies the `COUNT` aggregate on all rows for each group. - `RepartitionExec` - - `partitioning=Hash([UserID@0], 10)`: Divides the input into into 10 (new) output partitions based on the value of `hash(UserID)`. You can read more about this in the [partitioning] documentation. + - `partitioning=Hash([UserID@0], 10)`: Divides the input into 10 (new) output partitions based on the value of `hash(UserID)`. You can read more about this in the [partitioning] documentation. - `input_partitions=10`: Number of input partitions. - `CoalesceBatchesExec` - `target_batch_size=8192`: Combines smaller batches in to larger batches. In this case approximately 8192 rows in each batch. diff --git a/docs/source/user-guide/introduction.md b/docs/source/user-guide/introduction.md index b89457c66d919..1d9b618a012e6 100644 --- a/docs/source/user-guide/introduction.md +++ b/docs/source/user-guide/introduction.md @@ -103,6 +103,7 @@ Here are some active projects using DataFusion: - [Comet](https://github.com/apache/datafusion-comet) Apache Spark native query execution plugin - [Cube Store] Cube’s universal semantic layer platform is the next evolution of OLAP technology for AI, BI, spreadsheets, and embedded analytics - [datafusion-dft](https://github.com/datafusion-contrib/datafusion-dft) Batteries included CLI, TUI, and server implementations for DataFusion. +- [datapress](https://docs.datap-rs.org) An opinionated small and fast data server on parquet and delta tables. - [dbt Fusion engine](https://github.com/dbt-labs/dbt-fusion) The dbt Fusion engine, written in Rust, designed for speed and correctness with a native SQL understanding across DWH SQL dialects. - [delta-rs] Native Rust implementation of Delta Lake - [EDB Postgres Lakehouse] built with [Seafowl] @@ -112,7 +113,9 @@ Here are some active projects using DataFusion: - [GreptimeDB] Open Source & Cloud Native Distributed Time Series Database - [hiop](https://hiop.io) Serverless Data Logistic Platform - [HoraeDB] Distributed Time-Series Database +- [Hotdata](https://www.hotdata.dev) On-demand databases for AI agents with a unified query engine for vector, OLAP, and full-text search. - [Iceberg-rust](https://github.com/apache/iceberg-rust) Rust implementation of Apache Iceberg +- [IceGate](https://icegate.tech) Observability data lake engine for metrics, traces, logs, and events, built on Apache Iceberg with OpenTelemetry ingestion - [InfluxDB] Time Series Database - [Kamu] Planet-scale streaming data pipeline - [Kubeflow Trainer](https://github.com/kubeflow/trainer) Kubernetes-native project designed for @@ -122,7 +125,7 @@ Here are some active projects using DataFusion: - [OpenObserve] Distributed cloud native observability platform - [ParadeDB](https://github.com/paradedb/paradedb) PostgreSQL for Search & Analytics - [Parseable] Log storage and observability platform -- [Polygon.io](https://polygon.io/) Stock Market API +- [Massive.com](https://massive.com/) Stock Market API - [qv] Quickly view your data - [R2 Query Engine](https://blog.cloudflare.com/r2-sql-deep-dive/) Cloudflare's distributed engine for querying data in Iceberg Catalogs - [rerun.io](https://rerun.io/) Visualize and query robotics logs and transform them into training data. @@ -132,12 +135,14 @@ Here are some active projects using DataFusion: - [SedonaDB](https://github.com/apache/sedona-db) A single-node analytical database engine with geospatial as a first-class citizen - [Sleeper](https://github.com/gchq/sleeper) Serverless, cloud-native, log-structured merge tree based, scalable key-value store - [Spice.ai] Building blocks for data-driven AI applications +- [Supermetal](https://supermetal.io/) is a change data capture (CDC) platform that synchronizes data between databases, data warehouses, and lakehouses - [Synnada] Streaming-first framework for data products - [VegaFusion] Server-side acceleration for the [Vega](https://vega.github.io/) visualization grammar - [Vortex] An extensible, state of the art columnar file format - [Telemetry](https://telemetry.sh/) Structured logging made easy - [Xorq](https://github.com/xorq-labs/xorq/) Xorq is a multi-engine batch transformation framework built on Ibis, DataFusion and Arrow - [KalamDB](https://github.com/jamals86/KalamDB) SQL-first realtime state database for AI agents, chat products, and multi-tenant SaaS. +- [Infino](https://github.com/infino-ai/infino) Fast retrieval engine for SQL, full-text search, and vector search over Parquet on object storage Here are some less active projects that used DataFusion: diff --git a/docs/source/user-guide/parquet-content-defined-chunking.md b/docs/source/user-guide/parquet-content-defined-chunking.md new file mode 100644 index 0000000000000..d456200e69990 --- /dev/null +++ b/docs/source/user-guide/parquet-content-defined-chunking.md @@ -0,0 +1,163 @@ + + +# Parquet Content-Defined Chunking + +Content-defined chunking (CDC) is an experimental Parquet writer feature that +makes data page boundaries depend on column values rather than fixed row or byte +counts. This makes unchanged regions more likely to produce identical pages when +closely related versions of a dataset are written with the same settings. + +CDC is useful when the resulting files are stored or transferred through a +content-addressable or block-deduplicating system. Such a system can reuse the +identical pages instead of storing or transferring them again. For example, a +small insertion near the beginning of a dataset can change one page while later +page boundaries converge back to those of the previous version. + +CDC does not itself deduplicate data or provide a page store. On a conventional +filesystem or object store, each Parquet file is still stored in full. The output +is a normal Parquet file and requires no CDC-specific reader support. + +## When to enable CDC + +Consider CDC when all of the following apply: + +- You regularly write similar versions of the same dataset. +- Your storage or transfer layer detects and reuses duplicate byte ranges. +- Reducing storage or network transfer is more important than maximizing write + parallelism for an individual file. + +Leave CDC disabled for ordinary Parquet output unless you have measured a benefit +in the system that stores or transfers the files. CDC is disabled by default. + +When CDC is enabled, DataFusion uses the sequential Arrow writer for each output +file because the chunker's state must persist across row groups. This can reduce +write throughput compared with DataFusion's parallel writer path. Writing +different output files can still proceed concurrently. + +CDC operates independently for each output file. When `COPY` targets a +directory, DataFusion distributes input RecordBatches in round-robin order across +parallel output files; `datafusion.execution.minimum_parallel_output_files` +defaults to four. If batching or file assignment changes between dataset +versions, unchanged rows can move between files and reduce deduplication. For +the best results, keep the input order and output file layout stable. Use a +filename target for single-file output, or partition by stable keys when +multiple files are required. See +[Configuration Settings](configs.md#setting-configuration-options). + +## Enable CDC with SQL + +Set CDC for one [`COPY`](sql/dml.md#copy) operation with Parquet format options. +The filename target in this example selects single-file output: + +```sql +COPY ( + SELECT + value AS id, + CONCAT('event-', CAST(value AS VARCHAR)) AS event + FROM generate_series(1, 100000) +) TO 'cdc-output.parquet' +STORED AS PARQUET +OPTIONS ( + 'format.content_defined_chunking.enabled' 'true' +); +``` + +The default chunking parameters are a good starting point. The next example +specifies those defaults explicitly for one write; it does not change their +values: + +```sql +COPY source_table TO 'cdc-output.parquet' +STORED AS PARQUET +OPTIONS ( + 'format.content_defined_chunking.enabled' 'true', + 'format.content_defined_chunking.min_chunk_size' '262144', + 'format.content_defined_chunking.max_chunk_size' '1048576', + 'format.content_defined_chunking.norm_level' '0' +); +``` + +Change these values only after measuring with representative data. + +You can instead enable CDC for subsequent Parquet writes in the session: + +```sql +SET datafusion.execution.parquet.content_defined_chunking.enabled = true; +``` + +The corresponding environment variable is +`DATAFUSION_EXECUTION_PARQUET_CONTENT_DEFINED_CHUNKING_ENABLED`. See +[Configuration Settings](configs.md#setting-configuration-options) for all ways +to set session options. + +## Enable CDC with the Rust API + +Pass [`TableParquetOptions`] to [`DataFrame::write_parquet`]: + +```rust +use datafusion::config::{ParquetCdcOptions, TableParquetOptions}; +use datafusion::dataframe::DataFrameWriteOptions; +use datafusion::error::Result; +use datafusion::prelude::SessionContext; + +#[tokio::main] +async fn main() -> Result<()> { + let ctx = SessionContext::new(); + let df = ctx + .sql("SELECT value AS id FROM generate_series(1, 100000)") + .await?; + + let mut parquet_options = TableParquetOptions::default(); + parquet_options.global.content_defined_chunking = ParquetCdcOptions::enabled(); + + df.write_parquet( + "cdc-output.parquet", + DataFrameWriteOptions::new().with_single_file_output(true), + Some(parquet_options), + ) + .await?; + + Ok(()) +} +``` + +Set the fields of `ParquetCdcOptions` directly to use non-default chunk sizes or +normalization. + +## Tuning + +| Option | Default | Effect | +| ---------------- | ------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `min_chunk_size` | 256 KiB | Minimum logical size before the rolling hash can select a boundary. | +| `max_chunk_size` | 1 MiB | Maximum logical size before the writer forces a boundary. It must be greater than `min_chunk_size`. | +| `norm_level` | `0` | Controls how aggressively boundaries are selected. Higher values can improve deduplication but create more small pages; recommended range is `-3` through `3`. | + +Chunk sizes are measured from logical column data before encoding and +compression. Definition and repetition levels for nested data also count toward +the size. + +Use the same CDC, encoding, compression, and schema settings when comparing +dataset versions. Changing writer settings can change the page bytes and reduce +deduplication even when the logical data is unchanged. Measure the deduplication +ratio, output size, network transfer, and write time with representative data +before changing the defaults. + +[`dataframe::write_parquet`]: https://docs.rs/datafusion/latest/datafusion/dataframe/struct.DataFrame.html#method.write_parquet +[`tableparquetoptions`]: https://docs.rs/datafusion/latest/datafusion/common/config/struct.TableParquetOptions.html diff --git a/docs/source/user-guide/sql/aggregate_functions.md b/docs/source/user-guide/sql/aggregate_functions.md index ba9c6ae12477b..c681ccb28e1ee 100644 --- a/docs/source/user-guide/sql/aggregate_functions.md +++ b/docs/source/user-guide/sql/aggregate_functions.md @@ -80,6 +80,7 @@ SELECT SUM(x) WITHIN GROUP (ORDER BY x) FROM t; ## General Functions +- [any_value](#any_value) - [array_agg](#array_agg) - [avg](#avg) - [bit_and](#bit_and) @@ -105,6 +106,29 @@ SELECT SUM(x) WITHIN GROUP (ORDER BY x) FROM t; - [var_samp](#var_samp) - [var_sample](#var_sample) +### `any_value` + +Returns an arbitrary non-null value from a group, or NULL if the group contains only NULL values. + +```sql +any_value(expression) +``` + +#### Arguments + +- **expression**: The expression to operate on. Can be a constant, column, or function, and any combination of operators. + +#### Example + +```sql +> SELECT any_value(column_name) FROM table_name; ++------------------------+ +| any_value(column_name) | ++------------------------+ +| arbitrary_value | ++------------------------+ +``` + ### `array_agg` Returns an array created from the expression elements. If ordering is required, elements are inserted in the specified order. diff --git a/docs/source/user-guide/sql/ddl.md b/docs/source/user-guide/sql/ddl.md index 3a5c934ae8156..0d76775bcc1c6 100644 --- a/docs/source/user-guide/sql/ddl.md +++ b/docs/source/user-guide/sql/ddl.md @@ -82,6 +82,21 @@ For a comprehensive list of format-specific options that can be specified in the a path to a file or directory of partitioned files locally or on an object store. +Multiple locations can be supplied as a parenthesized list of string literals, +in which case the files are read together as one table: + +```sql +CREATE EXTERNAL TABLE hits +STORED AS PARQUET +LOCATION ( + 's3://clickhouse-public-datasets/hits_compatible/athena_partitioned/hits_1.parquet', + 's3://clickhouse-public-datasets/hits_compatible/athena_partitioned/hits_2.parquet' +); +``` + +All listed locations must reside on the same object store and resolve to data +with the same schema. + ### Example: Parquet Parquet data sources can be registered by executing a `CREATE EXTERNAL TABLE` SQL statement such as the following. It is not necessary to diff --git a/docs/source/user-guide/sql/explain.md b/docs/source/user-guide/sql/explain.md index 23101632625b1..e7be47b35001c 100644 --- a/docs/source/user-guide/sql/explain.md +++ b/docs/source/user-guide/sql/explain.md @@ -227,8 +227,9 @@ Elapsed 0.010 seconds. ## `EXPLAIN ANALYZE` -Shows the execution plan and metrics of a statement. Note that `EXPLAIN ANALYZE` -only supports the `indent` format. +Shows the execution plan and metrics of a statement. `EXPLAIN ANALYZE` supports +the `indent` format (the default) and the [`pgjson`](#pgjson-format-with-analyze) +format; the `tree` and `graphviz` formats are not supported with `ANALYZE`. ```sql EXPLAIN ANALYZE SELECT SUM(x) FROM table GROUP BY b; @@ -251,4 +252,46 @@ By default `EXPLAIN ANALYZE` shows the aggregated metrics from all partitions fo You can also set `datafusion.explain.analyze_level` from the [configuration value] to control the detail level for the metrics displayed. +### `pgjson` format with `ANALYZE` + +`EXPLAIN ANALYZE` can also emit the physical plan and its live execution +metrics in the [`pgjson`](#pgjson-format) format, so the analyzed plan can be +loaded into PostgreSQL plan visualizers such as [dalibo]. Each node reports its +`Actual Rows` and `Actual Total Time` (compute time, in milliseconds) using the +PostgreSQL key names, with any remaining DataFusion metrics under `Extras`. + +The format can be requested with either the keyword form +(`EXPLAIN ANALYZE FORMAT pgjson ...`) or, more idiomatically, the PostgreSQL +[option-list form](../explain-usage.md), which also lets you combine it with the +`METRICS` and `LEVEL` knobs in a single statement: + +```sql +> CREATE TABLE t(x int, b int) AS VALUES (1, 2), (2, 3); +> EXPLAIN (ANALYZE, FORMAT pgjson, METRICS 'rows') SELECT x FROM t WHERE b > 2; ++-------------------+---------------------------------------------------------------------------+ +| plan_type | plan | ++-------------------+---------------------------------------------------------------------------+ +| Plan with Metrics | [ | +| | { | +| | "Plan": { | +| | "Node Type": "FilterExec", | +| | "Details": "FilterExec: b@1 > 2, projection=[x@0]", | +| | "Actual Rows": 1, | +| | "Extras": { | +| | "output_batches": 1, | +| | "selectivity": "50% (1/2)" | +| | }, | +| | "Plans": [ | +| | { | +| | "Node Type": "DataSourceExec", | +| | "Details": "DataSourceExec: partitions=1, partition_sizes=[1]", | +| | "Plans": [] | +| | } | +| | ] | +| | } | +| | } | +| | ] | ++-------------------+---------------------------------------------------------------------------+ +``` + [configuration value]: ../configs.md diff --git a/docs/source/user-guide/sql/format_options.md b/docs/source/user-guide/sql/format_options.md index 46d251c18ed74..719fd5bd6b1ee 100644 --- a/docs/source/user-guide/sql/format_options.md +++ b/docs/source/user-guide/sql/format_options.md @@ -142,6 +142,7 @@ The following options are available when reading or writing Parquet files. If an | BLOOM_FILTER_FPP | Yes | Sets bloom filter false positive probability (global or per column). | `'bloom_filter_fpp'` or `'bloom_filter_fpp::col'` | None | | BLOOM_FILTER_NDV | Yes | Sets bloom filter number of distinct values (global or per column). | `'bloom_filter_ndv'` or `'bloom_filter_ndv::col'` | None | | MAX_ROW_GROUP_SIZE | No | Sets the maximum number of rows per row group. Larger groups require more memory but can improve compression and scan efficiency. | `'max_row_group_size'` | 1048576 | +| MAX_ROW_GROUP_BYTES | No | Sets the maximum size of each row group in bytes. When both this and `MAX_ROW_GROUP_SIZE` are set, the row group flushes whenever either limit is reached. Mirrors `parquet.block.size` from parquet-mr. Currently only honored when `allow_single_file_parallelism` is `false`; by default the parallel file writer ignores it. | `'max_row_group_bytes'` | None | | ENABLE_PAGE_INDEX | No | If true, reads the Parquet data page level metadata (the Page Index), if present, to reduce I/O and decoding. | `'enable_page_index'` | true | | PRUNING | No | If true, enables row group pruning based on min/max statistics. | `'pruning'` | true | | SKIP_METADATA | No | If true, skips optional embedded metadata in the file schema. | `'skip_metadata'` | true | @@ -163,6 +164,10 @@ The following options are available when reading or writing Parquet files. If an | ALLOW_SINGLE_FILE_PARALLELISM | No | Enables parallel serialization of columns in a single file. | `'allow_single_file_parallelism'` | true | | MAXIMUM_PARALLEL_ROW_GROUP_WRITERS | No | Maximum number of parallel row group writers. | `'maximum_parallel_row_group_writers'` | 1 | | MAXIMUM_BUFFERED_RECORD_BATCHES_PER_STREAM | No | Maximum number of buffered record batches per stream. | `'maximum_buffered_record_batches_per_stream'` | 2 | +| CONTENT_DEFINED_CHUNKING_ENABLED | No | Enables experimental content-defined chunking when writing Parquet files. Enabling it uses the sequential writer for each output file so chunker state persists across row groups. See [Parquet Content-Defined Chunking](../parquet-content-defined-chunking.md). | `'content_defined_chunking.enabled'` | false | +| CONTENT_DEFINED_CHUNKING_MIN_CHUNK_SIZE | No | Minimum logical size in bytes before the rolling hash can select a chunk boundary. | `'content_defined_chunking.min_chunk_size'` | 262144 (256 KiB) | +| CONTENT_DEFINED_CHUNKING_MAX_CHUNK_SIZE | No | Maximum logical size in bytes before the writer forces a chunk boundary. Must be greater than `content_defined_chunking.min_chunk_size`. | `'content_defined_chunking.max_chunk_size'` | 1048576 (1 MiB) | +| CONTENT_DEFINED_CHUNKING_NORM_LEVEL | No | Controls how aggressively chunk boundaries are selected. Higher values can improve deduplication but increase fragmentation. The recommended range is `-3` through `3`. | `'content_defined_chunking.norm_level'` | 0 | | KEY_VALUE_METADATA | No (Key is specific) | Adds custom key-value pairs to the file metadata. Use the format `'metadata::your_key_name' 'your_value'`. Multiple entries allowed. | `'metadata::key_name'` | None | **Example:** diff --git a/docs/source/user-guide/sql/scalar_functions.md b/docs/source/user-guide/sql/scalar_functions.md index 6bf61391eb10e..1bfec4ce43599 100644 --- a/docs/source/user-guide/sql/scalar_functions.md +++ b/docs/source/user-guide/sql/scalar_functions.md @@ -200,7 +200,7 @@ atan(numeric_expression) #### Example ```sql - > SELECT atan(1); +> SELECT atan(1); +-----------+ | atan(1) | +-----------+ @@ -227,11 +227,11 @@ atan2(expression_y, expression_x) ```sql > SELECT atan2(1, 1); -+------------+ -| atan2(1,1) | -+------------+ -| 0.7853982 | -+------------+ ++--------------------+ +| atan2(1,1) | ++--------------------+ +| 0.7853981633974483 | ++--------------------+ ``` ### `atanh` @@ -249,7 +249,7 @@ atanh(numeric_expression) #### Example ```sql - > SELECT atanh(0.5); +> SELECT atanh(0.5); +-------------+ | atanh(0.5) | +-------------+ @@ -387,7 +387,7 @@ degrees(numeric_expression) #### Example ```sql - > SELECT degrees(pi()); +> SELECT degrees(pi()); +------------+ | degrees(0) | +------------+ @@ -913,12 +913,12 @@ tanh(numeric_expression) #### Example ```sql - > SELECT tanh(20); - +----------+ - | tanh(20) | - +----------+ - | 1.0 | - +----------+ +> SELECT tanh(20); ++----------+ +| tanh(20) | ++----------+ +| 1.0 | ++----------+ ``` ### `trunc` @@ -2195,6 +2195,15 @@ encode(expression, format) Apache DataFusion uses a [PCRE-like](https://en.wikibooks.org/wiki/Regular_Expressions/Perl-Compatible_Regular_Expressions) regular expression [syntax](https://docs.rs/regex/latest/regex/#syntax) (minus support for several features including look-around and backreferences). + +The following flags are optionally supported in functions: + +- **i**: case-insensitive: letters match both upper and lower case +- **m**: multi-line mode: `^` and `$` match begin/end of line +- **s**: allow `.` to match `\n` +- **R**: enables CRLF mode: when multi-line mode is enabled, `\r\n` is used +- **U**: swap the meaning of `x*` and `x*?` + The following regular expression functions are supported: - [regexp_count](#regexp_count) @@ -2208,20 +2217,15 @@ The following regular expression functions are supported: Returns the number of matches that a [regular expression](https://docs.rs/regex/latest/regex/#syntax) has in a string. ```sql -regexp_count(str, regexp[, start, flags]) +regexp_count(str, regexp[, start[, flags]]) ``` #### Arguments - **str**: String expression to operate on. Can be a constant, column, or function, and any combination of operators. - **regexp**: Regular expression to operate on. Can be a constant, column, or function, and any combination of operators. -- **start**: - **start**: Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function. -- **flags**: Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: - - **i**: case-insensitive: letters match both upper and lower case - - **m**: multi-line mode: ^ and $ match begin/end of line - - **s**: allow . to match \n - - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used - - **U**: swap the meaning of x* and x*? +- **start**: Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function. +- **flags**: Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags. #### Example @@ -2246,14 +2250,9 @@ regexp_instr(str, regexp[, start[, N[, flags[, subexpr]]]]) - **str**: String expression to operate on. Can be a constant, column, or function, and any combination of operators. - **regexp**: Regular expression to operate on. Can be a constant, column, or function, and any combination of operators. -- **start**: - **start**: Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function. Defaults to 1 -- **N**: - **N**: Optional The N-th occurrence of pattern to find. Defaults to 1 (first match). Can be a constant, column, or function. -- **flags**: Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: - - **i**: case-insensitive: letters match both upper and lower case - - **m**: multi-line mode: ^ and $ match begin/end of line - - **s**: allow . to match \n - - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used - - **U**: swap the meaning of x* and x*? +- **start**: Optional start position (the first position is 1) to search for the regular expression. Can be a constant, column, or function. Defaults to 1 +- **N**: Optional The N-th occurrence of pattern to find. Defaults to 1 (first match). Can be a constant, column, or function. +- **flags**: Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags. - **subexpr**: Optional Specifies which capture group (subexpression) to return the position for. Defaults to 0, which returns the position of the entire match. #### Example @@ -2279,12 +2278,7 @@ regexp_like(str, regexp[, flags]) - **str**: String expression to operate on. Can be a constant, column, or function, and any combination of operators. - **regexp**: Regular expression to operate on. Can be a constant, column, or function, and any combination of operators. -- **flags**: Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: - - **i**: case-insensitive: letters match both upper and lower case - - **m**: multi-line mode: ^ and $ match begin/end of line - - **s**: allow . to match \n - - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used - - **U**: swap the meaning of x* and x*? +- **flags**: Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags. #### Example @@ -2318,12 +2312,7 @@ regexp_match(str, regexp[, flags]) - **str**: String expression to operate on. Can be a constant, column, or function, and any combination of operators. - **regexp**: Regular expression to match against. Can be a constant, column, or function. -- **flags**: Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: - - **i**: case-insensitive: letters match both upper and lower case - - **m**: multi-line mode: ^ and $ match begin/end of line - - **s**: allow . to match \n - - **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used - - **U**: swap the meaning of x* and x*? +- **flags**: Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags. #### Example @@ -2358,13 +2347,7 @@ regexp_replace(str, regexp, replacement[, flags]) - **regexp**: Regular expression to match against. Can be a constant, column, or function. - **replacement**: Replacement string expression to operate on. Can be a constant, column, or function, and any combination of operators. -- **flags**: Optional regular expression flags that control the behavior of the regular expression. The following flags are supported: -- **g**: (global) Search globally and don't return after the first match -- **i**: case-insensitive: letters match both upper and lower case -- **m**: multi-line mode: ^ and $ match begin/end of line -- **s**: allow . to match \n -- **R**: enables CRLF mode: when multi-line mode is enabled, \r\n is used -- **U**: swap the meaning of x* and x*? +- **flags**: Optional regular expression flags that control the behavior of the regular expression. Refer to the flags reference above for supported flags. #### Example @@ -2420,8 +2403,6 @@ The `current_date()` return value is determined at query time and will return th ```sql current_date() - (optional) SET datafusion.execution.time_zone = '+00:00'; - SELECT current_date(); ``` #### Example @@ -2458,8 +2439,6 @@ The session time zone can be set using the statement 'SET datafusion.execution.t ```sql current_time() - (optional) SET datafusion.execution.time_zone = '+00:00'; - SELECT current_time(); ``` #### Example @@ -2493,14 +2472,14 @@ Calculates time intervals and returns the start of the interval nearest to the s For example, if you "bin" or "window" data into 15 minute intervals, an input timestamp of `2023-01-01T18:18:18Z` will be updated to the start time of the 15 minute bin it is in: `2023-01-01T18:15:00Z`. ```sql -date_bin(interval, expression, origin-timestamp) +date_bin(interval, expression[, origin_timestamp]) ``` #### Arguments - **interval**: Bin interval. - **expression**: Time expression to operate on. Can be a constant, column, or function. -- **origin-timestamp**: Optional. Starting point used to determine bin boundaries. If not specified defaults 1970-01-01T00:00:00Z (the UNIX epoch in UTC). The following intervals are supported: +- **origin_timestamp**: Optional. Starting point used to determine bin boundaries. If not specified defaults 1970-01-01T00:00:00Z (the UNIX epoch in UTC). The following intervals are supported: - nanoseconds - microseconds @@ -2813,7 +2792,6 @@ to_char(expression, format) - **expression**: Expression to operate on. Can be a constant, column, or function that results in a date, time, timestamp or duration. - **format**: A [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) string to use to convert the expression. -- **day**: Day to use when making the date. Can be a constant, column or function, and any combination of arithmetic operators. #### Example @@ -2843,7 +2821,7 @@ Returns the corresponding date. Note: `to_date` returns Date32, which represents its values as the number of days since unix epoch(`1970-01-01`) stored as signed 32 bit value. The largest supported date value is `9999-12-31`. ```sql -to_date('2017-05-31', '%Y-%m-%d') +to_date(expression[, format1, ..., format_n]) ``` #### Arguments @@ -2851,7 +2829,7 @@ to_date('2017-05-31', '%Y-%m-%d') - **expression**: String expression to operate on. Can be a constant, column, or function, and any combination of operators. - **format_n**: Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully parse the expression - an error will be returned. + an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. #### Example @@ -2944,7 +2922,7 @@ Returns the corresponding time. Note: `to_time` returns Time64(Nanosecond), which represents the time of day in nanoseconds since midnight. ```sql -to_time('12:30:45', '%H:%M:%S') +to_time(expression[, format1, ..., format_n]) ``` #### Arguments @@ -3006,7 +2984,8 @@ to_timestamp(expression[, ..., format_n]) - **format_n**: Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully - parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is + parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. + Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. #### Example @@ -3050,7 +3029,8 @@ to_timestamp_micros(expression[, ..., format_n]) - **format_n**: Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully - parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is + parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. + Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. #### Example @@ -3094,7 +3074,8 @@ to_timestamp_millis(expression[, ..., format_n]) - **format_n**: Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully - parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is + parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. + Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. #### Example @@ -3137,7 +3118,8 @@ to_timestamp_nanos(expression[, ..., format_n]) - **format_n**: Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully - parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is + parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. + Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. #### Example @@ -3181,7 +3163,8 @@ to_timestamp_seconds(expression[, ..., format_n]) - **format_n**: Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully - parse the expression an error will be returned. Note: parsing of named timezones (e.g. 'America/New_York') using %Z is + parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. + Note: parsing of named timezones (e.g. 'America/New_York') using %Z is only supported at the end of the string preceded by a space. #### Example @@ -3218,7 +3201,7 @@ to_unixtime(expression[, ..., format_n]) #### Arguments - **expression**: Expression to operate on. Can be a constant, column, or function, and any combination of arithmetic operators. -- **format_n**: Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully parse the expression an error will be returned. +- **format_n**: Optional [Chrono format](https://docs.rs/chrono/latest/chrono/format/strftime/index.html) strings to use to parse the expression. Formats will be tried in the order they appear with the first successful one being returned. If none of the formats successfully parse the expression an error will be returned. NULL formats are skipped. If every format is NULL the result is NULL. #### Example @@ -3244,9 +3227,11 @@ _Alias of [current_date](#current_date)._ ## Array Functions - [any_match](#any_match) +- [array_add](#array_add) - [array_any_match](#array_any_match) - [array_any_value](#array_any_value) - [array_append](#array_append) +- [array_avg](#array_avg) - [array_cat](#array_cat) - [array_compact](#array_compact) - [array_concat](#array_concat) @@ -3259,6 +3244,7 @@ _Alias of [current_date](#current_date)._ - [array_except](#array_except) - [array_extract](#array_extract) - [array_filter](#array_filter) +- [array_first](#array_first) - [array_has](#array_has) - [array_has_all](#array_has_all) - [array_has_any](#array_has_any) @@ -3275,6 +3261,7 @@ _Alias of [current_date](#current_date)._ - [array_position](#array_position) - [array_positions](#array_positions) - [array_prepend](#array_prepend) +- [array_product](#array_product) - [array_push_back](#array_push_back) - [array_push_front](#array_push_front) - [array_remove](#array_remove) @@ -3286,8 +3273,11 @@ _Alias of [current_date](#current_date)._ - [array_replace_n](#array_replace_n) - [array_resize](#array_resize) - [array_reverse](#array_reverse) +- [array_scale](#array_scale) - [array_slice](#array_slice) - [array_sort](#array_sort) +- [array_subtract](#array_subtract) +- [array_sum](#array_sum) - [array_to_string](#array_to_string) - [array_transform](#array_transform) - [array_union](#array_union) @@ -3300,9 +3290,11 @@ _Alias of [current_date](#current_date)._ - [flatten](#flatten) - [generate_series](#generate_series) - [inner_product](#inner_product) +- [list_add](#list_add) - [list_any_match](#list_any_match) - [list_any_value](#list_any_value) - [list_append](#list_append) +- [list_avg](#list_avg) - [list_cat](#list_cat) - [list_compact](#list_compact) - [list_concat](#list_concat) @@ -3315,6 +3307,7 @@ _Alias of [current_date](#current_date)._ - [list_except](#list_except) - [list_extract](#list_extract) - [list_filter](#list_filter) +- [list_first](#list_first) - [list_has](#list_has) - [list_has_all](#list_has_all) - [list_has_any](#list_has_any) @@ -3330,6 +3323,7 @@ _Alias of [current_date](#current_date)._ - [list_position](#list_position) - [list_positions](#list_positions) - [list_prepend](#list_prepend) +- [list_product](#list_product) - [list_push_back](#list_push_back) - [list_push_front](#list_push_front) - [list_remove](#list_remove) @@ -3341,8 +3335,11 @@ _Alias of [current_date](#current_date)._ - [list_replace_n](#list_replace_n) - [list_resize](#list_resize) - [list_reverse](#list_reverse) +- [list_scale](#list_scale) - [list_slice](#list_slice) - [list_sort](#list_sort) +- [list_subtract](#list_subtract) +- [list_sum](#list_sum) - [list_to_string](#list_to_string) - [list_transform](#list_transform) - [list_union](#list_union) @@ -3357,6 +3354,34 @@ _Alias of [current_date](#current_date)._ _Alias of [array_any_match](#array_any_match)._ +### `array_add` + +Returns the element-wise sum of two numeric arrays of equal length, computed as `array1[i] + array2[i]` per position. NULL is propagated per element: if either input element at position `i` is NULL, the corresponding output element is NULL (positions are preserved). Returns NULL if either entire input array is NULL. Errors if the per-row lengths differ. Returns an empty array if both inputs are empty. + +```sql +array_add(array1, array2) +``` + +#### Arguments + +- **array1**: Array expression. Can be a constant, column, or function, and any combination of array operators. +- **array2**: Array expression. Can be a constant, column, or function, and any combination of array operators. + +#### Example + +```sql +> select array_add([1.0, 2.0, 3.0], [10.0, 20.0, 30.0]); ++---------------------------------------------------------+ +| array_add(List([1.0,2.0,3.0]),List([10.0,20.0,30.0])) | ++---------------------------------------------------------+ +| [11.0, 22.0, 33.0] | ++---------------------------------------------------------+ +``` + +#### Aliases + +- list_add + ### `array_any_match` Returns whether any elements of an array match the given predicate. Returns true if one or more elements match, false if none match (including empty arrays), and null if the predicate returns null for some elements and false for all others. @@ -3388,7 +3413,7 @@ any_match(array, predicate) ### `array_any_value` -Returns the first non-null element in the array. +Returns the first non-null element in the array. Returns NULL if the array is empty or NULL. ```sql array_any_value(array) @@ -3443,6 +3468,33 @@ array_append(array, element) - array_push_back - list_push_back +### `array_avg` + +Returns the arithmetic mean (sum divided by count) of the elements of the input array. NULL elements are skipped (per SQL aggregate convention) and excluded from the count. Returns NULL if the input row is NULL, every element is NULL, or the array is empty. + +```sql +array_avg(array) +``` + +#### Arguments + +- **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. + +#### Example + +```sql +> select array_avg([1.0, 2.0, 3.0]); ++----------------------------+ +| array_avg(List([1.0,2.0,3.0])) | ++----------------------------+ +| 2.0 | ++----------------------------+ +``` + +#### Aliases + +- list_avg + ### `array_cat` _Alias of [array_concat](#array_concat)._ @@ -3537,7 +3589,7 @@ array_dims(array) ### `array_distance` -Returns the Euclidean distance between two input arrays of equal length. +Returns the Euclidean distance between two one-dimensional input arrays of equal length. ```sql array_distance(array1, array2) @@ -3690,6 +3742,34 @@ array_filter(array, x -> x > 2) - list_filter +### `array_first` + +Returns the first element of an array that satisfies the given predicate. Returns null if the array is empty or no element matches. A predicate that returns null for an element is treated as not matching. + +```sql +array_first(array, predicate) +``` + +#### Arguments + +- **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. +- **predicate**: Lambda predicate that returns a boolean. The first element for which it returns true is returned. + +#### Example + +```sql +> select array_first([1, 2, 3, 4], x -> x > 2); ++----------------------------------------+ +| array_first([1,2,3,4],x -> x > 2) | ++----------------------------------------+ +| 3 | ++----------------------------------------+ +``` + +#### Aliases + +- list_first + ### `array_has` Returns true if the array contains the element. @@ -3722,26 +3802,26 @@ array_has(array, element) ### `array_has_all` -Returns true if all elements of sub-array exist in array. +Returns true if all elements of sub_array exist in array. ```sql -array_has_all(array, sub-array) +array_has_all(array, sub_array) ``` #### Arguments - **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. -- **sub-array**: Array expression. Can be a constant, column, or function, and any combination of array operators. +- **sub_array**: Array expression. Can be a constant, column, or function, and any combination of array operators. #### Example ```sql > select array_has_all([1, 2, 3, 4], [2, 3]); -+--------------------------------------------+ ++---------------------------------------------+ | array_has_all(List([1,2,3,4]), List([2,3])) | -+--------------------------------------------+ -| true | -+--------------------------------------------+ ++---------------------------------------------+ +| true | ++---------------------------------------------+ ``` #### Aliases @@ -3824,13 +3904,13 @@ _Alias of [array_to_string](#array_to_string)._ Returns the length of the array dimension. ```sql -array_length(array, dimension) +array_length(array[, dimension]) ``` #### Arguments - **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. -- **dimension**: Array dimension. +- **dimension**: Array dimension. Default is 1 #### Example @@ -3902,13 +3982,12 @@ array_min(array) Returns the number of dimensions of the array. ```sql -array_ndims(array, element) +array_ndims(array) ``` #### Arguments - **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. -- **element**: Array element. #### Example @@ -4011,15 +4090,14 @@ array_pop_front(array) Returns the position of the first occurrence of the specified element in the array, or NULL if not found. Comparisons are done using `IS DISTINCT FROM` semantics, so NULL is considered to match NULL. ```sql -array_position(array, element) -array_position(array, element, index) +array_position(array, element[, index]) ``` #### Arguments - **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. - **element**: Element to search for in the array. -- **index**: Index at which to start searching (1-indexed). +- **index**: Index at which to start searching (1-indexed). Defaults to searching from the start #### Example @@ -4102,6 +4180,33 @@ array_prepend(element, array) - array_push_front - list_push_front +### `array_product` + +Returns the product of the elements in the input numeric array. NULL elements inside the array are skipped (matching SQL aggregate convention). Returns NULL if the input is NULL, every element is NULL, or the array is empty. The result is always returned as `Float64`. + +```sql +array_product(array) +``` + +#### Arguments + +- **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. + +#### Example + +```sql +> select array_product([1.0, 2.0, 3.0]); ++------------------------------------+ +| array_product(List([1.0,2.0,3.0])) | ++------------------------------------+ +| 6.0 | ++------------------------------------+ +``` + +#### Aliases + +- list_product + ### `array_push_back` _Alias of [array_append](#array_append)._ @@ -4340,17 +4445,17 @@ array_replace_n(array, from, to, max) ### `array_resize` -Resizes the list to contain size elements. Initializes new elements with value or empty if value is not set. +Resizes the list to contain size elements. ```sql -array_resize(array, size, value) +array_resize(array, size[, value]) ``` #### Arguments - **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. - **size**: New size of given array. -- **value**: Defines new elements' value or empty if value is not set. +- **value**: If expanding the array, defines the values to fill in. Defaults to null. #### Example @@ -4394,12 +4499,40 @@ array_reverse(array) - list_reverse +### `array_scale` + +Returns a new array with each element of the input array multiplied by a scalar value, computed as `array[i] * scalar`. Returns NULL if the input row is NULL or the scalar is NULL. If a NULL element appears in the input array at position `i`, the result element at position `i` is NULL. Returns an empty array for an empty input array. + +```sql +array_scale(array, scalar) +``` + +#### Arguments + +- **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. +- **scalar**: Numeric scalar to multiply each element by. Can be a constant or column expression. + +#### Example + +```sql +> select array_scale([1.0, 2.0, 3.0], 2.0); ++----------------------------------+ +| array_scale(List([1.0,2.0,3.0]),Float64(2.0)) | ++----------------------------------+ +| [2.0, 4.0, 6.0] | ++----------------------------------+ +``` + +#### Aliases + +- list_scale + ### `array_slice` Returns a slice of the array based on 1-indexed start and end positions. ```sql -array_slice(array, begin, end) +array_slice(array, begin, end[, stride]) ``` #### Arguments @@ -4429,14 +4562,14 @@ array_slice(array, begin, end) Sort array. ```sql -array_sort(array, desc, nulls_first) +array_sort(array[, order[, nulls_order]]) ``` #### Arguments - **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. -- **desc**: Whether to sort in ascending (`ASC`) or descending (`DESC`) order. The default is `ASC`. -- **nulls_first**: Whether to sort nulls first (`NULLS FIRST`) or last (`NULLS LAST`). The default is `NULLS FIRST`. +- **order**: Whether to sort in ascending (`ASC`) or descending (`DESC`) order. The default is `ASC`. +- **nulls_order**: Whether to sort nulls first (`NULLS FIRST`) or last (`NULLS LAST`). The default is `NULLS FIRST`. #### Example @@ -4447,12 +4580,73 @@ array_sort(array, desc, nulls_first) +-----------------------------+ | [1, 2, 3] | +-----------------------------+ +> select array_sort([3, 1, NULL, 2], 'desc', 'nulls last'); ++--------------------------------------------------+ +| array_sort(List(3,1,NULL,2),'desc','nulls last') | ++--------------------------------------------------+ +| [3, 2, 1, NULL] | ++--------------------------------------------------+ ``` #### Aliases - list_sort +### `array_subtract` + +Returns the element-wise difference of two numeric arrays of equal length, computed as `array1[i] - array2[i]` per position. NULL is propagated per element: if either input element at position `i` is NULL, the corresponding output element is NULL (positions are preserved). Returns NULL if either entire input array is NULL. Errors if the per-row lengths differ. Returns an empty array if both inputs are empty. + +```sql +array_subtract(array1, array2) +``` + +#### Arguments + +- **array1**: Array expression. Can be a constant, column, or function, and any combination of array operators. +- **array2**: Array expression. Can be a constant, column, or function, and any combination of array operators. + +#### Example + +```sql +> select array_subtract([10.0, 20.0, 30.0], [1.0, 2.0, 3.0]); ++--------------------------------------------------------------+ +| array_subtract(List([10.0,20.0,30.0]),List([1.0,2.0,3.0])) | ++--------------------------------------------------------------+ +| [9.0, 18.0, 27.0] | ++--------------------------------------------------------------+ +``` + +#### Aliases + +- list_subtract + +### `array_sum` + +Returns the sum of the elements of the input array, computed as `array[0] + array[1] + ...`. NULL elements are skipped (per SQL aggregate convention). Returns NULL if the input row is NULL, every element is NULL, or the array is empty. + +```sql +array_sum(array) +``` + +#### Arguments + +- **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. + +#### Example + +```sql +> select array_sum([1.0, 2.0, 3.0]); ++----------------------------+ +| array_sum(List([1.0,2.0,3.0])) | ++----------------------------+ +| 6.0 | ++----------------------------+ +``` + +#### Aliases + +- list_sum + ### `array_to_string` Converts each element to its text representation. @@ -4489,13 +4683,13 @@ array_to_string(array, delimiter[, null_string]) transforms the values of an array ```sql -array_transform(array, x -> x*2) +array_transform(array, lambda) ``` #### Arguments - **array**: Array expression. Can be a constant, column, or function, and any combination of array operators. -- **lambda**: Lambda +- **lambda**: The lambda function used to transform each value of the array. #### Example @@ -4745,6 +4939,10 @@ inner_product(array1, array2) - dot_product +### `list_add` + +_Alias of [array_add](#array_add)._ + ### `list_any_match` _Alias of [array_any_match](#array_any_match)._ @@ -4757,6 +4955,10 @@ _Alias of [array_any_value](#array_any_value)._ _Alias of [array_append](#array_append)._ +### `list_avg` + +_Alias of [array_avg](#array_avg)._ + ### `list_cat` _Alias of [array_concat](#array_concat)._ @@ -4805,6 +5007,10 @@ _Alias of [array_element](#array_element)._ _Alias of [array_filter](#array_filter)._ +### `list_first` + +_Alias of [array_first](#array_first)._ + ### `list_has` _Alias of [array_has](#array_has)._ @@ -4865,6 +5071,10 @@ _Alias of [array_positions](#array_positions)._ _Alias of [array_prepend](#array_prepend)._ +### `list_product` + +_Alias of [array_product](#array_product)._ + ### `list_push_back` _Alias of [array_append](#array_append)._ @@ -4909,6 +5119,10 @@ _Alias of [array_resize](#array_resize)._ _Alias of [array_reverse](#array_reverse)._ +### `list_scale` + +_Alias of [array_scale](#array_scale)._ + ### `list_slice` _Alias of [array_slice](#array_slice)._ @@ -4917,6 +5131,14 @@ _Alias of [array_slice](#array_slice)._ _Alias of [array_sort](#array_sort)._ +### `list_subtract` + +_Alias of [array_subtract](#array_subtract)._ + +### `list_sum` + +_Alias of [array_sum](#array_sum)._ + ### `list_to_string` _Alias of [array_to_string](#array_to_string)._ @@ -4938,7 +5160,7 @@ _Alias of [arrays_zip](#arrays_zip)._ Returns an array using the specified input expressions. ```sql -make_array(expression1[, ..., expression_n]) +make_array([expression1, ..., expression_n]) ``` #### Arguments @@ -5157,7 +5379,7 @@ The `make_map` function creates a map from two lists: one for keys and one for v ```sql map(key, value) -map(key: value) +map {key: value} make_map(['key1', 'key2'], ['value1', 'value2']) ``` @@ -5534,7 +5756,9 @@ union_tag(union_expression) - [arrow_try_cast](#arrow_try_cast) - [arrow_typeof](#arrow_typeof) - [cast_to_type](#cast_to_type) +- [file_row_index](#file_row_index) - [get_field](#get_field) +- [input_file_name](#input_file_name) - [try_cast_to_type](#try_cast_to_type) - [version](#version) - [with_metadata](#with_metadata) @@ -5717,6 +5941,27 @@ cast_to_type(expression, reference) +-----+ ``` +### `file_row_index` + +Returns the zero-based row offset within the source file +that produced the current row. + +The value is scoped to one file, so rows from different files in the same scan +can have the same row index. This function is intended to be rewritten at +file-scan time. If the input file is not known (for example, if this function +is evaluated outside a file scan, or was not pushed down into one), direct +evaluation returns an error. + +```sql +file_row_index() +``` + +#### Example + +```sql +SELECT file_row_index() FROM t; +``` + ### `get_field` Returns a field within a map or a struct with the given key. @@ -5769,6 +6014,26 @@ get_field(expression, field_name[, field_name2, ...]) +--------+ ``` +### `input_file_name` + +Returns the path of the input file that produced the current row. + +Note: file paths/URIs may be sensitive metadata depending on your environment. + +This function is intended to be rewritten at file-scan time (when the file is +known). If the input file is not known (for example, if this function is +evaluated outside a file scan, or was not pushed down into one), direct evaluation returns an error. + +```sql +input_file_name() +``` + +#### Example + +```sql +SELECT input_file_name() FROM t; +``` + ### `try_cast_to_type` Casts the first argument to the data type of the second argument, returning NULL if the cast fails. Only the type of the second argument is used; its value is ignored. diff --git a/docs/source/user-guide/sql/select.md b/docs/source/user-guide/sql/select.md index 3564884b041ad..af442de6597c1 100644 --- a/docs/source/user-guide/sql/select.md +++ b/docs/source/user-guide/sql/select.md @@ -19,76 +19,343 @@ # SELECT syntax -The queries in DataFusion scan data from tables and return 0 or more rows. -Please be aware that column names in queries are made lower-case, but not on the inferred schema. Accordingly, if you -want to query against a capitalized field, make sure to use double quotes. Please see this -[example](https://datafusion.apache.org/user-guide/example-usage.html) for clarification. -In this documentation we describe the SQL syntax in DataFusion. - -DataFusion supports the following syntax for queries: - - -[ [WITH](#with-clause) with_query [, ...] ]
-[SELECT](#select-clause) [ ALL | DISTINCT ] select_expr [, ...]
-[ [FROM](#from-clause) from_item [, ...] ]
-[ [JOIN](#join-clause) join_item [, ...] ]
-[ [WHERE](#where-clause) condition ]
-[ [GROUP BY](#group-by-clause) grouping_element [, ...] ]
-[ [HAVING](#having-clause) condition]
-[ [QUALIFY](#qualify-clause) condition]
-[ [UNION](#union-clause) [ ALL | select ]
-[ [ORDER BY](#order-by-clause) expression [ ASC | DESC ][, ...] ]
-[ [LIMIT](#limit-clause) count ]
-[ [EXCLUDE | EXCEPT](#exclude-and-except-clause) ]
-[Pipe operators](#pipe-operators)
- -
+Queries in DataFusion scan data from tables, subqueries, table functions, or +literal values and return zero or more rows. DataFusion supports the following +general form for `SELECT` queries. Optional clauses can be omitted. The linked +sections describe each clause in more detail. + +
+ +Unquoted identifiers are normalized to lower case in SQL queries, but inferred +schema field names are not changed. If a field name contains capital letters or +other characters that require quoting, reference it with double quotes. See this +[example](https://datafusion.apache.org/user-guide/example-usage.html) for +clarification. ## WITH clause -A with clause allows to give names for queries and reference them by name. +```text +WITH [RECURSIVE] cte_name [(column_name [, ...])] AS (query) [, ...] +``` + +A `WITH` clause defines common table expressions (CTEs) that can be referenced +by name in the rest of the query. + +Examples: ```sql WITH x AS (SELECT a, MAX(b) AS b FROM t GROUP BY a) SELECT a, b FROM x; ``` +CTEs can also rename their output columns: + +```sql +WITH x(key, total) AS ( + SELECT a, SUM(b) FROM t GROUP BY a +) +SELECT key, total FROM x; +``` + +DataFusion supports `WITH RECURSIVE` for recursive CTEs. Recursive CTE support +is controlled by the `datafusion.execution.enable_recursive_ctes` configuration +setting, which is enabled by default. + +```sql +WITH RECURSIVE numbers AS ( + SELECT 1 AS n + UNION ALL + SELECT n + 1 FROM numbers WHERE n < 3 +) +SELECT n FROM numbers; +``` + ## SELECT clause -Example: +```text +SELECT [ALL | DISTINCT | DISTINCT ON (expression [, ...])] + select_item [, ...] + [INTO table_name] +``` + +The `SELECT` list can contain column references, arbitrary expressions, scalar +functions, aggregate functions, window functions, scalar subqueries, and +wildcards. + +Examples: ```sql -SELECT a, b, a + b FROM table +SELECT a, b, a + b AS sum_ab FROM table_name; ``` -The `DISTINCT` quantifier can be added to make the query return all distinct rows. -By default `ALL` will be used, which returns all the rows. +Aliases can be written with or without `AS`: ```sql -SELECT DISTINCT person, age FROM employees +SELECT a AS key, b value FROM table_name; +``` + +`SELECT` can be used without a `FROM` clause when the selected expressions do +not need input rows: + +```sql +SELECT 1 + 2 AS three; +``` + +`SELECT *` requires a `FROM` clause. + +### DISTINCT + +```text +SELECT DISTINCT select_item [, ...] +SELECT DISTINCT ON (expression [, ...]) select_item [, ...] +``` + +By default, `SELECT` uses `ALL` semantics and returns every row. The `DISTINCT` +quantifier removes duplicate rows from the query result. + +Examples: + +```sql +SELECT DISTINCT person, age FROM employees; +``` + +DataFusion also supports PostgreSQL-style `DISTINCT ON`, which keeps one row for +each distinct value of the listed expressions. Use `ORDER BY` to choose which +row is kept for each group. When `ORDER BY` is present, the initial `ORDER BY` +expressions must match the `DISTINCT ON` expressions. + +If multiple rows have the same `DISTINCT ON` values and the `ORDER BY` clause +does not fully order those rows, the row that is kept is not specified. Add +additional `ORDER BY` expressions to make the choice deterministic. + +```sql +SELECT DISTINCT ON (customer_id) customer_id, order_id, order_date +FROM orders +ORDER BY customer_id, order_date DESC; +``` + +### Wildcards + +```text +* +table_alias.* +* EXCLUDE column_name +* EXCLUDE (column_name [, ...]) +* EXCEPT column_name +* EXCEPT (column_name [, ...]) +* REPLACE (expression AS column_name [, ...]) +``` + +Use `*` to select all columns, or `table_alias.*` to select all columns from a +specific input. + +Examples: + +```sql +SELECT * FROM orders; +SELECT o.* FROM orders AS o; +``` + +Wildcard projections support `EXCLUDE` and `EXCEPT` to omit columns. Both +accept either a single column name or a parenthesized list of column names. + +```sql +SELECT * EXCLUDE customer_id FROM orders; +SELECT * EXCLUDE (customer_id, internal_note) FROM orders; +SELECT * EXCEPT customer_id FROM orders; +SELECT * EXCEPT (customer_id, internal_note) FROM orders; +SELECT o.* EXCLUDE (internal_note) FROM orders AS o; +``` + +Every name in an `EXCLUDE` or `EXCEPT` list must refer to an existing column. +The list must not name the same column more than once, and the wildcard must +not expand to zero columns. + +Wildcard projections also support `REPLACE`, which keeps the original column +name but substitutes a new expression for that column. + +```sql +SELECT * REPLACE (price * 2 AS price) FROM products; +SELECT p.* REPLACE (price * 2 AS price, product_id + 1000 AS product_id) +FROM products AS p; +``` + +`RENAME` and wildcard aliases such as `* AS alias` are not supported. + +### SELECT INTO + +```text +SELECT select_item [, ...] INTO table_name FROM ... +``` + +`SELECT ... INTO table_name` creates an in-memory table from the query result. +It is similar to [`CREATE TABLE ... AS SELECT`](ddl.md#create-table). + +```sql +SELECT customer_id, SUM(amount) AS total +INTO customer_totals +FROM orders +GROUP BY customer_id; ``` ## FROM clause -Example: +```text +FROM from_item [, ...] + +from_item: + table_name [[AS] alias [(column_alias [, ...])]] +| (query) [[AS] alias [(column_alias [, ...])]] +| VALUES (expression [, ...]) [, ...] [[AS] alias [(column_alias [, ...])]] +| table_function(argument [, ...]) [[AS] alias [(column_alias [, ...])]] +| UNNEST(expression) [[AS] alias [(column_alias [, ...])]] +``` + +The `FROM` clause specifies the input relations for the query. Supported inputs +include tables, CTEs, derived tables, `VALUES`, table functions, and `UNNEST`. + +Examples: + +```sql +SELECT t.a FROM table_name AS t; +``` + +Table aliases can include column aliases: + +```sql +SELECT x, y +FROM some_table AS t(x, y); +``` + +Subqueries can be used in the `FROM` clause: + +```sql +SELECT q.a +FROM (SELECT a FROM table_name WHERE a > 10) AS q; +``` + +`VALUES` can be used as a table expression: + +```sql +SELECT * +FROM VALUES (1, 'a'), (2, 'b') AS t(id, label); +``` + +Table functions such as `range` and `generate_series` can be used in `FROM`: ```sql -SELECT t.a FROM table AS t +SELECT value FROM range(0, 3); ``` +`UNNEST` expands a list, array, or similar nested value into one row for each +element. It can be used in the `SELECT` list to expand a value in each input +row, or as an input relation in `FROM`. When used in `FROM`, it can be given a +table alias and column alias. + +```sql +SELECT * FROM UNNEST([1, 2, 3]) AS u(value); +``` + +To expand a column for each input row, use `UNNEST` in the `SELECT` list: + +```sql +SELECT id, UNNEST(items) FROM orders; +``` + +`UNNEST` in the `FROM` clause cannot yet reference columns from preceding `FROM` +items (implicit lateral references such as `FROM orders AS t, UNNEST(t.items)` +are not currently supported). + +### `unnest_outer` + +`unnest_outer(col)` is the outer-unnest peer to `UNNEST(col)`. The two differ +only in how `NULL` and empty input lists are handled: + +| Form | `NULL` input list | Empty input list | +| ------------------- | ----------------- | ---------------- | +| `UNNEST(col)` | dropped | dropped | +| `unnest_outer(col)` | one `NULL` row | one `NULL` row | + +```sql +SELECT id, unnest_outer(tags) AS tag FROM rows; +``` + +An input row with an empty `tags` array or `NULL` `tags` produces one output +row whose `tag` is `NULL`, instead of being dropped. This is analogous to the +outer variant offered by other engines (Spark `explode_outer`, Hive `EXPLODE OUTER`, Snowflake `FLATTEN(OUTER => true)`). + +`unnest_outer` cannot be mixed with `unnest` in the same `SELECT` — the +unnest plan node carries a single null-handling mode for all its output +columns, so a mix would be ambiguous. The planner returns an error in that +case. + ## WHERE clause -Example: +```text +WHERE condition +``` + +The `WHERE` clause filters input rows before grouping, aggregation, and window +processing. ```sql -SELECT a FROM table WHERE a > 10 +SELECT a FROM table_name WHERE a > 10; ``` ## JOIN clause -DataFusion supports `INNER JOIN`, `LEFT OUTER JOIN`, `RIGHT OUTER JOIN`, `FULL OUTER JOIN`, `NATURAL JOIN`, `CROSS JOIN`, `LEFT SEMI JOIN`, `RIGHT SEMI JOIN`, `LEFT ANTI JOIN`, `RIGHT ANTI JOIN`, `LATERAL JOIN`, and `LEFT JOIN LATERAL`. +```text +from_item [join_type] JOIN from_item [join_condition] +from_item CROSS JOIN from_item +from_item NATURAL JOIN from_item +from_item [join_type] JOIN LATERAL (query) AS alias [join_condition] +from_item, LATERAL (query) AS alias + +join_type: + INNER +| LEFT [OUTER] +| RIGHT [OUTER] +| FULL [OUTER] +| LEFT SEMI +| RIGHT SEMI +| LEFT ANTI +| RIGHT ANTI + +join_condition: + ON condition +| USING (column_name [, ...]) +``` + +Joins are written inside the `FROM` clause between input relations. -The following examples are based on this table: +Join conditions can use `ON` or `USING`. + +Examples: + +```sql +SELECT * +FROM orders AS o +JOIN customers AS c ON o.customer_id = c.id; + +SELECT * +FROM orders +JOIN customers USING (customer_id); +``` + +The join examples below use this table: ```sql select * from x; @@ -104,7 +371,7 @@ select * from x; The keywords `JOIN` or `INNER JOIN` define a join that only shows rows where there is a match in both tables. ```sql -SELECT * FROM x INNER JOIN x y ON x.column_1 = y.column_1; +SELECT * FROM x INNER JOIN x AS y ON x.column_1 = y.column_1; +----------+----------+----------+----------+ | column_1 | column_2 | column_1 | column_2 | +----------+----------+----------+----------+ @@ -112,13 +379,20 @@ SELECT * FROM x INNER JOIN x y ON x.column_1 = y.column_1; +----------+----------+----------+----------+ ``` +The same behavior can also be written by listing both inputs in the `FROM` +clause and putting the join condition in the `WHERE` clause: + +```sql +SELECT * FROM x, x AS y WHERE x.column_1 = y.column_1; +``` + ### LEFT OUTER JOIN The keywords `LEFT JOIN` or `LEFT OUTER JOIN` define a join that includes all rows from the left table even if there is not a match in the right table. When there is no match, null values are produced for the right side of the join. ```sql -SELECT * FROM x LEFT JOIN x y ON x.column_1 = y.column_2; +SELECT * FROM x LEFT JOIN x AS y ON x.column_1 = y.column_2; +----------+----------+----------+----------+ | column_1 | column_2 | column_1 | column_2 | +----------+----------+----------+----------+ @@ -132,7 +406,7 @@ The keywords `RIGHT JOIN` or `RIGHT OUTER JOIN` define a join that includes all is not a match in the left table. When there is no match, null values are produced for the left side of the join. ```sql -SELECT * FROM x RIGHT JOIN x y ON x.column_1 = y.column_2; +SELECT * FROM x RIGHT JOIN x AS y ON x.column_1 = y.column_2; +----------+----------+----------+----------+ | column_1 | column_2 | column_1 | column_2 | +----------+----------+----------+----------+ @@ -147,7 +421,7 @@ The keywords `FULL JOIN` or `FULL OUTER JOIN` define a join that is effectively either side of the join where there is not a match. ```sql -SELECT * FROM x FULL OUTER JOIN x y ON x.column_1 = y.column_2; +SELECT * FROM x FULL OUTER JOIN x AS y ON x.column_1 = y.column_2; +----------+----------+----------+----------+ | column_1 | column_2 | column_1 | column_2 | +----------+----------+----------+----------+ @@ -162,7 +436,7 @@ A `NATURAL JOIN` defines an inner join based on common column names found betwee column names are found, it behaves like a `CROSS JOIN`. ```sql -SELECT * FROM x NATURAL JOIN x y; +SELECT * FROM x NATURAL JOIN x AS y; +----------+----------+ | column_1 | column_2 | +----------+----------+ @@ -176,7 +450,7 @@ A `CROSS JOIN` produces a cartesian product that matches every row in the left s right side of the join. ```sql -SELECT * FROM x CROSS JOIN x y; +SELECT * FROM x CROSS JOIN x AS y; +----------+----------+----------+----------+ | column_1 | column_2 | column_1 | column_2 | +----------+----------+----------+----------+ @@ -190,7 +464,7 @@ The `LEFT SEMI JOIN` returns all rows from the left table that have at least one projects only the columns from the left table. ```sql -SELECT * FROM x LEFT SEMI JOIN x y ON x.column_1 = y.column_1; +SELECT * FROM x LEFT SEMI JOIN x AS y ON x.column_1 = y.column_1; +----------+----------+ | column_1 | column_2 | +----------+----------+ @@ -204,7 +478,7 @@ The `RIGHT SEMI JOIN` returns all rows from the right table that have at least o only projects the columns from the right table. ```sql -SELECT * FROM x RIGHT SEMI JOIN x y ON x.column_1 = y.column_1; +SELECT * FROM x RIGHT SEMI JOIN x AS y ON x.column_1 = y.column_1; +----------+----------+ | column_1 | column_2 | +----------+----------+ @@ -215,10 +489,10 @@ SELECT * FROM x RIGHT SEMI JOIN x y ON x.column_1 = y.column_1; ### LEFT ANTI JOIN The `LEFT ANTI JOIN` returns all rows from the left table that do not have any matching row in the right table, projecting -only the left table’s columns. +only the left table's columns. ```sql -SELECT * FROM x LEFT ANTI JOIN x y ON x.column_1 = y.column_1; +SELECT * FROM x LEFT ANTI JOIN x AS y ON x.column_1 = y.column_1; +----------+----------+ | column_1 | column_2 | +----------+----------+ @@ -228,10 +502,10 @@ SELECT * FROM x LEFT ANTI JOIN x y ON x.column_1 = y.column_1; ### RIGHT ANTI JOIN The `RIGHT ANTI JOIN` returns all rows from the right table that do not have any matching row in the left table, projecting -only the right table’s columns. +only the right table's columns. ```sql -SELECT * FROM x RIGHT ANTI JOIN x y ON x.column_1 = y.column_1; +SELECT * FROM x RIGHT ANTI JOIN x AS y ON x.column_1 = y.column_1; +----------+----------+ | column_1 | column_2 | +----------+----------+ @@ -367,43 +641,131 @@ The following patterns are not yet supported: - Outer references in the `SELECT` list of the lateral subquery (e.g., `LATERAL (SELECT outer.col + 1)`). - `HAVING` in lateral subqueries. +- `FULL OUTER JOIN LATERAL`, `RIGHT JOIN LATERAL`, `RIGHT SEMI JOIN LATERAL`, and `RIGHT ANTI JOIN LATERAL`. ## GROUP BY clause -Example: +```text +GROUP BY ALL +GROUP BY grouping_element [, ...] + +grouping_element: + expression + ordinal_position + ROLLUP(expression [, ...]) + CUBE(expression [, ...]) + GROUPING SETS ((grouping_element [, ...]) [, ...]) +``` + +The `GROUP BY` clause groups rows before aggregate expressions are evaluated. +Grouping elements can be expressions, output aliases, or ordinal positions in +the `SELECT` list. + +Examples: ```sql -SELECT a, b, MAX(c) FROM table GROUP BY a, b +SELECT a, b, MAX(c) FROM table_name GROUP BY a, b; +SELECT a AS key, COUNT(*) FROM table_name GROUP BY key; +SELECT a, b, COUNT(*) FROM table_name GROUP BY 1, 2; ``` -Some aggregation functions accept optional ordering requirement, such as `ARRAY_AGG`. If a requirement is given, -aggregation is calculated in the order of the requirement. +`GROUP BY ALL` groups by every non-aggregate expression in the `SELECT` list. + +```sql +SELECT a, b, SUM(c) FROM table_name GROUP BY ALL; +``` + +Grouping sets allow a single query to compute aggregates for multiple grouping +levels. `ROLLUP(a, b)` computes aggregate rows grouped by `(a, b)`, then by +`a`, then over all input rows. `CUBE(a, b)` computes aggregate rows for all +combinations of `a` and `b`. `GROUPING SETS` lets you list the grouping levels +explicitly. + +```sql +SELECT a, b, SUM(c) FROM table_name GROUP BY ROLLUP(a, b); +SELECT a, b, SUM(c) FROM table_name GROUP BY CUBE(a, b); +SELECT a, b, SUM(c) +FROM table_name +GROUP BY GROUPING SETS ((a), (a, b), ()); +``` -Example: +Some aggregate functions accept an optional ordering requirement, such as +`ARRAY_AGG`. If an ordering requirement is given, aggregation is calculated in +that order. ```sql -SELECT a, b, ARRAY_AGG(c, ORDER BY d) FROM table GROUP BY a, b +SELECT a, b, ARRAY_AGG(c ORDER BY d) FROM table_name GROUP BY a, b; ``` ## HAVING clause -Example: +```text +HAVING condition +``` + +The `HAVING` clause filters groups after aggregation. It can reference grouping +expressions, aggregate expressions, and aliases from the `SELECT` list. ```sql -SELECT a, b, MAX(c) FROM table GROUP BY a, b HAVING MAX(c) > 10 +SELECT a, b, MAX(c) AS max_c +FROM table_name +GROUP BY a, b +HAVING max_c > 10; +``` + +## WINDOW clause + +```text +WINDOW window_name AS (window_definition) [, ...] +``` + +The `WINDOW` clause defines named window specifications that can be referenced +from window functions. See [Window Functions](window_functions.md) for the full +window-function reference. + +```sql +SELECT + depname, + empno, + salary, + AVG(salary) OVER w AS avg_salary +FROM empsalary +WINDOW w AS (PARTITION BY depname ORDER BY salary DESC); ``` ## QUALIFY clause -Example: +```text +QUALIFY condition +``` + +The `QUALIFY` clause filters rows after window functions are evaluated. A query +with `QUALIFY` must contain a window function in either the `SELECT` list or the +`QUALIFY` expression. `QUALIFY` can reference aliases from the `SELECT` list. ```sql -SELECT ROW_NUMBER() OVER (PARTITION BY region) AS rk FROM table QUALIFY rk > 1; +SELECT ROW_NUMBER() OVER (PARTITION BY region ORDER BY sales DESC) AS rk +FROM table_name +QUALIFY rk <= 3; +``` + +## Set operations + +```text +query UNION [ALL | DISTINCT] [BY NAME] query +query INTERSECT [ALL | DISTINCT] query +query EXCEPT [ALL | DISTINCT] query ``` -## UNION clause +Set operations combine the results of two queries into a single result. They +operate on whole rows rather than on individual columns, and the input queries +must produce compatible columns. Except for `UNION ... BY NAME` variants, +inputs must have the same number of output columns. + +`UNION` returns rows from both inputs and removes duplicates by default. +`UNION DISTINCT` is equivalent to `UNION`; `UNION ALL` preserves duplicates. -Example: +Examples: ```sql SELECT @@ -419,55 +781,108 @@ SELECT FROM table2 ``` +`INTERSECT` returns rows that appear in both inputs. `EXCEPT` returns rows from +the left input that do not appear in the right input. Both support `ALL` and +`DISTINCT`. + +```sql +SELECT a FROM table1 +INTERSECT +SELECT a FROM table2; + +SELECT a FROM table1 +EXCEPT ALL +SELECT a FROM table2; +``` + +`UNION BY NAME` matches columns by name instead of by position. `UNION ALL BY NAME` preserves duplicates, and `UNION DISTINCT BY NAME` removes duplicates. + +```sql +SELECT a, b FROM table1 +UNION BY NAME +SELECT b, a FROM table2; +``` + +Set operations can be followed by `ORDER BY`, `LIMIT`, and `OFFSET` clauses, +which apply to the combined result. + ## ORDER BY clause -Orders the results by the referenced expression. By default it uses ascending order (`ASC`). -This order can be changed to descending by adding `DESC` after the order-by expressions. +```text +ORDER BY order_expression [ASC | DESC] [NULLS FIRST | NULLS LAST] [, ...] +``` + +`ORDER BY` sorts query results. Each `order_expression` can be an expression, a +`SELECT` alias, or an ordinal position. The default direction is ascending +(`ASC`). + +If multiple rows have the same values for every `ORDER BY` expression, their +relative order is not specified. Add additional `ORDER BY` expressions to break +ties when the exact row order matters. Examples: ```sql -SELECT age, person FROM table ORDER BY age; -SELECT age, person FROM table ORDER BY age DESC; -SELECT age, person FROM table ORDER BY age, person DESC; +SELECT age, person FROM table_name ORDER BY age; +SELECT age, person FROM table_name ORDER BY age DESC; +SELECT age AS years, person FROM table_name ORDER BY years; +SELECT age, person FROM table_name ORDER BY 1, person DESC; ``` -## LIMIT clause +Use `NULLS FIRST` or `NULLS LAST` to control where null values sort: -Limits the number of rows to be a maximum of `count` rows. `count` should be a non-negative integer. +```sql +SELECT age, person FROM table_name ORDER BY age DESC NULLS LAST; +``` -Example: +With the DuckDB dialect, DataFusion supports `ORDER BY ALL`, which orders by +every column in the `SELECT` list from left to right. All selected items must +be column references; ordering by computed expressions such as `a + b` is not +supported: ```sql -SELECT age, person FROM table -LIMIT 10 +SET datafusion.sql_parser.dialect = 'DuckDB'; +SELECT address, zip FROM addresses ORDER BY ALL DESC; +``` + +## LIMIT and OFFSET clauses + +```text +[LIMIT count] +[OFFSET count] ``` -## EXCLUDE and EXCEPT clause +`LIMIT` restricts the number of rows returned. `OFFSET` skips rows before +returning results. The count expressions must be constant expressions that +evaluate to non-negative integers or `NULL`; column references are not allowed. +`NULL` has no effect. -Excluded named columns from query results. +Without an `ORDER BY` clause, `LIMIT` and `OFFSET` operate on an unspecified row +order, so the returned rows are not guaranteed to be deterministic. -Example selecting all columns except for `age` and `person`: +Examples: ```sql -SELECT * EXCEPT(age, person) -FROM table; +SELECT age, person FROM table_name LIMIT 10; +SELECT age, person FROM table_name OFFSET 20; +SELECT age, person FROM table_name LIMIT 10 OFFSET 20; +SELECT age, person FROM table_name OFFSET 20 LIMIT 10; ``` +DataFusion also accepts MySQL-style `LIMIT offset, count`: + ```sql -SELECT * EXCLUDE(age, person) -FROM table; +SELECT age, person FROM table_name LIMIT 20, 10; ``` ## Pipe operators -Some SQL dialects (e.g. BigQuery) support the pipe operator `|>`. -The SQL dialect can be set like this: - -```sql -set datafusion.sql_parser.dialect = 'BigQuery'; +```text +query |> pipe_operator [|> pipe_operator ...] ``` +DataFusion supports BigQuery-style pipe operators (`|>`). + DataFusion currently supports the following pipe operators: - [WHERE](#pipe_where) diff --git a/header b/header deleted file mode 100644 index 70665d1a26295..0000000000000 --- a/header +++ /dev/null @@ -1,16 +0,0 @@ -Licensed to the Apache Software Foundation (ASF) under one -or more contributor license agreements. See the NOTICE file -distributed with this work for additional information -regarding copyright ownership. The ASF licenses this file -to you under the Apache License, Version 2.0 (the -"License"); you may not use this file except in compliance -with the License. You may obtain a copy of the License at - -http://www.apache.org/licenses/LICENSE-2.0 - -Unless required by applicable law or agreed to in writing, software -distributed under the License is distributed on an "AS IS" BASIS, -WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -See the License for the specific language governing permissions and -limitations under the License. - diff --git a/licenserc.toml b/licenserc.toml index 105d969ea56e6..a1e01a5fd0ace 100644 --- a/licenserc.toml +++ b/licenserc.toml @@ -27,4 +27,5 @@ excludes = [ # generated code "datafusion/proto/src/generated/", "datafusion/proto-common/src/generated/", + "datafusion/proto-models/src/generated/", ] diff --git a/rust-toolchain.toml b/rust-toolchain.toml index 5df661d61cd6f..5639a821f5b98 100644 --- a/rust-toolchain.toml +++ b/rust-toolchain.toml @@ -19,5 +19,5 @@ # to compile this workspace and run CI jobs. [toolchain] -channel = "1.95.0" +channel = "1.97.0" components = ["rustfmt", "clippy"] diff --git a/snowflake_flatten_validation.sql b/snowflake_flatten_validation.sql deleted file mode 100644 index cae6f5ea59e77..0000000000000 --- a/snowflake_flatten_validation.sql +++ /dev/null @@ -1,219 +0,0 @@ --- ============================================================================ --- Snowflake LATERAL FLATTEN validation queries --- --- Run this file against a real Snowflake instance to verify that the --- Unparser-generated SQL is syntactically and semantically correct. --- --- Each section shows: --- 1. The DataFusion input (SQL parsed by the planner) --- 2. The Snowflake SQL produced by the Unparser --- --- NOTE: The Unparser emits array literals as [1, 2, 3] (DataFusion syntax). --- Snowflake requires ARRAY_CONSTRUCT(1, 2, 3). The queries below use --- ARRAY_CONSTRUCT so they can run directly on Snowflake. The exact Unparser --- output is shown in the "Unparser output:" comment above each query. --- ============================================================================ - --- ---------------------------------------------------------------------------- --- Setup: create and seed test tables --- ---------------------------------------------------------------------------- - -CREATE OR REPLACE TABLE source ( - items ARRAY -); - -INSERT INTO source SELECT PARSE_JSON('[1, 2, 3]'); -INSERT INTO source SELECT PARSE_JSON('["a", "b"]'); -INSERT INTO source SELECT NULL; - -CREATE OR REPLACE TABLE unnest_table ( - array_col ARRAY -); - -INSERT INTO unnest_table SELECT PARSE_JSON('[10, 20, 30]'); -INSERT INTO unnest_table SELECT PARSE_JSON('[40, 50]'); -INSERT INTO unnest_table SELECT NULL; - -CREATE OR REPLACE TABLE multi_array_table ( - column_a ARRAY, - column_b ARRAY -); - -INSERT INTO multi_array_table SELECT PARSE_JSON('[1, 2, 3]'), PARSE_JSON('["x", "y"]'); -INSERT INTO multi_array_table SELECT PARSE_JSON('[4]'), PARSE_JSON('["z"]'); - --- ============================================================================ --- Roundtrip tests: SQL parsed → plan → Snowflake SQL --- ============================================================================ - --- -------------------------------------------------------------------------- --- Test: snowflake_unnest_to_lateral_flatten_simple --- DataFusion input: SELECT * FROM UNNEST([1,2,3]) --- Unparser output: SELECT "_unnest_1"."VALUE" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS "_unnest_1" --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" -FROM LATERAL FLATTEN(INPUT => ARRAY_CONSTRUCT(1, 2, 3)) AS "_unnest_1"; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_implicit_from --- DataFusion input: SELECT UNNEST([1,2,3]) --- Unparser output: SELECT "_unnest_1"."VALUE" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS "_unnest_1" --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" -FROM LATERAL FLATTEN(INPUT => ARRAY_CONSTRUCT(1, 2, 3)) AS "_unnest_1"; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_string_array --- DataFusion input: SELECT * FROM UNNEST(['a','b','c']) --- Unparser output: SELECT "_unnest_1"."VALUE" FROM LATERAL FLATTEN(INPUT => ['a', 'b', 'c']) AS "_unnest_1" --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" -FROM LATERAL FLATTEN(INPUT => ARRAY_CONSTRUCT('a', 'b', 'c')) AS "_unnest_1"; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_select_unnest_with_alias --- DataFusion input: SELECT UNNEST([1,2,3]) as c1 --- Unparser output: SELECT "_unnest_1"."VALUE" AS "c1" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS "_unnest_1" --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" AS "c1" -FROM LATERAL FLATTEN(INPUT => ARRAY_CONSTRUCT(1, 2, 3)) AS "_unnest_1"; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_from_unnest_with_table_alias --- DataFusion input: SELECT * FROM UNNEST([1,2,3]) AS t1 (c1) --- Unparser output: SELECT "t1"."VALUE" FROM LATERAL FLATTEN(INPUT => [1, 2, 3]) AS "t1" --- -------------------------------------------------------------------------- -SELECT "t1"."VALUE" -FROM LATERAL FLATTEN(INPUT => ARRAY_CONSTRUCT(1, 2, 3)) AS "t1"; - --- ============================================================================ --- Plan-built tests: LogicalPlan → Snowflake SQL --- These use a table called "source" with an ARRAY column "items". --- ============================================================================ - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_limit_between_projection_and_unnest --- Plan: Projection → Limit → Unnest → Projection → TableScan --- Unparser output: SELECT "_unnest_1"."VALUE" AS "item" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" LIMIT 5 --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" AS "item" -FROM "source" -CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" -LIMIT 5; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_sort_between_projection_and_unnest --- Plan: Projection → Sort → Unnest → Projection → TableScan --- Unparser output: SELECT "_unnest_1"."VALUE" AS "item" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" ORDER BY "_unnest_1"."VALUE" ASC NULLS FIRST --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" AS "item" -FROM "source" -CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" -ORDER BY "_unnest_1"."VALUE" ASC NULLS FIRST; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_limit_between_projection_and_unnest_with_subquery_alias --- Plan: Projection → Limit → Unnest → SubqueryAlias → Projection → TableScan --- Unparser output: SELECT "_unnest_1"."VALUE" AS "item" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" LIMIT 10 --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" AS "item" -FROM "source" -CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" -LIMIT 10; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_composed_expression_wrapping_unnest --- Plan: Projection(CAST(placeholder AS Int64)) → Unnest → Projection → TableScan --- Unparser output: SELECT CAST("_unnest_1"."VALUE" AS BIGINT) AS "item_id" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" --- -------------------------------------------------------------------------- -SELECT CAST("_unnest_1"."VALUE" AS BIGINT) AS "item_id" -FROM "source" -CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1"; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_composed_expression_with_limit --- Plan: Projection(CAST) → Limit → Unnest → Projection → TableScan --- Unparser output: SELECT CAST("_unnest_1"."VALUE" AS BIGINT) AS "item_id" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" LIMIT 5 --- -------------------------------------------------------------------------- -SELECT CAST("_unnest_1"."VALUE" AS BIGINT) AS "item_id" -FROM "source" -CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" -LIMIT 5; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_multi_expression_projection --- Plan: Projection([CAST AS Int64, CAST AS Utf8]) → Unnest → Projection → TableScan --- Unparser output: SELECT CAST("_unnest_1"."VALUE" AS BIGINT) AS "a", CAST("_unnest_1"."VALUE" AS VARCHAR) AS "b" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" --- -------------------------------------------------------------------------- -SELECT CAST("_unnest_1"."VALUE" AS BIGINT) AS "a", - CAST("_unnest_1"."VALUE" AS VARCHAR) AS "b" -FROM "source" -CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1"; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_multi_expression_with_limit --- Plan: Projection([CAST, CAST]) → Limit → Unnest → Projection → TableScan --- Unparser output: SELECT CAST("_unnest_1"."VALUE" AS BIGINT) AS "a", CAST("_unnest_1"."VALUE" AS VARCHAR) AS "b" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" LIMIT 10 --- -------------------------------------------------------------------------- -SELECT CAST("_unnest_1"."VALUE" AS BIGINT) AS "a", - CAST("_unnest_1"."VALUE" AS VARCHAR) AS "b" -FROM "source" -CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" -LIMIT 10; - --- -------------------------------------------------------------------------- --- Test: snowflake_unnest_through_subquery_alias --- Plan: Projection → Unnest → SubqueryAlias → Projection → TableScan --- Unparser output: SELECT "_unnest_1"."VALUE" AS "item" FROM "source" CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1" --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" AS "item" -FROM "source" -CROSS JOIN LATERAL FLATTEN(INPUT => "source"."items", OUTER => true) AS "_unnest_1"; - --- ============================================================================ --- Roundtrip tests with table columns --- ============================================================================ - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_unnest_from_subselect --- DataFusion input: SELECT UNNEST(array_col) FROM (SELECT array_col FROM unnest_table WHERE array_col IS NOT NULL LIMIT 3) --- Unparser output: SELECT "_unnest_1"."VALUE" FROM (SELECT "unnest_table"."array_col" FROM "unnest_table" WHERE "unnest_table"."array_col" IS NOT NULL LIMIT 3) CROSS JOIN LATERAL FLATTEN(INPUT => "unnest_table"."array_col") AS "_unnest_1" --- -------------------------------------------------------------------------- -SELECT "_unnest_1"."VALUE" -FROM ( - SELECT "unnest_table"."array_col" - FROM "unnest_table" - WHERE "unnest_table"."array_col" IS NOT NULL - LIMIT 3 -) CROSS JOIN LATERAL FLATTEN(INPUT => "unnest_table"."array_col") AS "_unnest_1"; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_cross_join_unnest_table_column --- DataFusion input: SELECT * FROM multi_array_table CROSS JOIN UNNEST(column_a) AS a (a) --- Unparser output: SELECT "multi_array_table"."column_a", "multi_array_table"."column_b", "a"."VALUE" FROM "multi_array_table" CROSS JOIN LATERAL FLATTEN(INPUT => "multi_array_table"."column_a") AS "a" --- -------------------------------------------------------------------------- -SELECT "multi_array_table"."column_a", - "multi_array_table"."column_b", - "a"."VALUE" -FROM "multi_array_table" -CROSS JOIN LATERAL FLATTEN(INPUT => "multi_array_table"."column_a") AS "a"; - --- -------------------------------------------------------------------------- --- Test: snowflake_flatten_multiple_unnest_cross_join --- DataFusion input: SELECT a.a, b.b FROM multi_array_table --- CROSS JOIN UNNEST(column_a) AS a (a) --- CROSS JOIN UNNEST(column_b) AS b (b) --- Unparser output: SELECT "a"."VALUE", "b"."VALUE" FROM "multi_array_table" CROSS JOIN LATERAL FLATTEN(INPUT => "multi_array_table"."column_a") AS "a" CROSS JOIN LATERAL FLATTEN(INPUT => "multi_array_table"."column_b") AS "b" --- -------------------------------------------------------------------------- -SELECT "a"."VALUE", - "b"."VALUE" -FROM "multi_array_table" -CROSS JOIN LATERAL FLATTEN(INPUT => "multi_array_table"."column_a") AS "a" -CROSS JOIN LATERAL FLATTEN(INPUT => "multi_array_table"."column_b") AS "b"; - --- ============================================================================ --- Cleanup --- ============================================================================ --- DROP TABLE IF EXISTS source; --- DROP TABLE IF EXISTS unnest_table; --- DROP TABLE IF EXISTS multi_array_table; diff --git a/test-utils/src/lib.rs b/test-utils/src/lib.rs index be2bc0712afbd..55717c717c4af 100644 --- a/test-utils/src/lib.rs +++ b/test-utils/src/lib.rs @@ -19,6 +19,7 @@ use arrow::datatypes::Schema; use arrow::record_batch::RecordBatch; use datafusion_common::cast::as_int32_array; +use datafusion_common::{Constraint, Constraints}; use rand::prelude::StdRng; use rand::{Rng, SeedableRng}; @@ -113,6 +114,7 @@ pub fn stagger_batch_with_seed(batch: RecordBatch, seed: u64) -> Vec Self { + self.constraints = constraints; + self + } +} + +fn primary_key(schema: &Schema, column_names: &[&str]) -> Constraint { + let indices = column_names + .iter() + .map(|column_name| { + schema.index_of(column_name).unwrap_or_else(|_| { + panic!("primary key column '{column_name}' not found in schema") + }) + }) + .collect(); + + Constraint::PrimaryKey(indices) } diff --git a/test-utils/src/tpcds.rs b/test-utils/src/tpcds.rs index 28992eb043036..af1f727531d75 100644 --- a/test-utils/src/tpcds.rs +++ b/test-utils/src/tpcds.rs @@ -15,12 +15,18 @@ // specific language governing permissions and limitations // under the License. -use crate::TableDef; +use crate::{TableDef, primary_key}; use arrow::datatypes::{DataType, Field, Schema}; +use datafusion_common::Constraints; pub fn tpcds_schemas() -> Vec { + let def = |name, schema: Schema| { + let constraints = tpcds_constraints(name, &schema); + TableDef::new(name, schema).with_constraints(constraints) + }; + vec![ - TableDef::new( + def( "catalog_sales", Schema::new(vec![ Field::new("cs_sold_date_sk", DataType::Int32, false), @@ -63,7 +69,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("cs_net_profit", DataType::Decimal128(7, 2), false), ]), ), - TableDef::new( + def( "catalog_returns", Schema::new(vec![ Field::new("cr_returned_date_sk", DataType::Int32, false), @@ -95,7 +101,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("cr_net_loss", DataType::Decimal128(7, 2), false), ]), ), - TableDef::new( + def( "inventory", Schema::new(vec![ Field::new("inv_date_sk", DataType::Int32, false), @@ -104,7 +110,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("inv_quantity_on_hand", DataType::Int32, false), ]), ), - TableDef::new( + def( "store_sales", Schema::new(vec![ Field::new("ss_sold_date_sk", DataType::Int32, false), @@ -132,7 +138,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("ss_net_profit", DataType::Decimal128(7, 2), false), ]), ), - TableDef::new( + def( "store_returns", Schema::new(vec![ Field::new("sr_returned_date_sk", DataType::Int32, false), @@ -157,7 +163,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("sr_net_loss", DataType::Decimal128(7, 2), false), ]), ), - TableDef::new( + def( "web_sales", Schema::new(vec![ Field::new("ws_sold_date_sk", DataType::Int32, false), @@ -200,7 +206,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("ws_net_profit", DataType::Decimal128(7, 2), false), ]), ), - TableDef::new( + def( "web_returns", Schema::new(vec![ Field::new("wr_returned_date_sk", DataType::Int32, false), @@ -229,7 +235,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("wr_net_loss", DataType::Decimal128(7, 2), false), ]), ), - TableDef::new( + def( "call_center", Schema::new(vec![ Field::new("cc_call_center_sk", DataType::Int32, false), @@ -265,7 +271,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("cc_tax_percentage", DataType::Decimal128(5, 2), false), ]), ), - TableDef::new( + def( "catalog_page", Schema::new(vec![ Field::new("cp_catalog_page_sk", DataType::Int32, false), @@ -279,7 +285,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("cp_type", DataType::Utf8, false), ]), ), - TableDef::new( + def( "customer", Schema::new(vec![ Field::new("c_customer_sk", DataType::Int32, false), @@ -302,7 +308,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("c_last_review_date_sk", DataType::Int32, false), ]), ), - TableDef::new( + def( "customer_address", Schema::new(vec![ Field::new("ca_address_sk", DataType::Int32, false), @@ -320,7 +326,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("ca_location_type", DataType::Utf8, false), ]), ), - TableDef::new( + def( "customer_demographics", Schema::new(vec![ Field::new("cd_demo_sk", DataType::Int32, false), @@ -334,7 +340,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("cd_dep_college_count", DataType::Int32, false), ]), ), - TableDef::new( + def( "date_dim", Schema::new(vec![ Field::new("d_date_sk", DataType::Int32, false), @@ -367,7 +373,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("d_current_year", DataType::Utf8, false), ]), ), - TableDef::new( + def( "household_demographics", Schema::new(vec![ Field::new("hd_demo_sk", DataType::Int32, false), @@ -377,7 +383,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("hd_vehicle_count", DataType::Int32, false), ]), ), - TableDef::new( + def( "income_band", Schema::new(vec![ Field::new("ib_income_band_sk", DataType::Int32, false), @@ -385,7 +391,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("ib_upper_bound", DataType::Int32, false), ]), ), - TableDef::new( + def( "item", Schema::new(vec![ Field::new("i_item_sk", DataType::Int32, false), @@ -412,7 +418,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("i_product_name", DataType::Utf8, false), ]), ), - TableDef::new( + def( "promotion", Schema::new(vec![ Field::new("p_promo_sk", DataType::Int32, false), @@ -436,7 +442,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("p_discount_active", DataType::Utf8, false), ]), ), - TableDef::new( + def( "reason", Schema::new(vec![ Field::new("r_reason_sk", DataType::Int32, false), @@ -444,7 +450,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("r_reason_desc", DataType::Utf8, false), ]), ), - TableDef::new( + def( "ship_mode", //), Schema::new(vec![ @@ -456,7 +462,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("sm_contract", DataType::Utf8, false), ]), ), - TableDef::new( + def( "store", Schema::new(vec![ Field::new("s_store_sk", DataType::Int32, false), @@ -490,7 +496,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("s_tax_precentage", DataType::Decimal128(5, 2), false), ]), ), - TableDef::new( + def( "time_dim", Schema::new(vec![ Field::new("t_time_sk", DataType::Int32, false), @@ -505,7 +511,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("t_meal_time", DataType::Utf8, false), ]), ), - TableDef::new( + def( "warehouse", //), Schema::new(vec![ @@ -525,7 +531,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("w_gmt_offset", DataType::Decimal128(5, 2), false), ]), ), - TableDef::new( + def( "web_page", Schema::new(vec![ Field::new("wp_web_page_sk", DataType::Int32, false), @@ -544,7 +550,7 @@ pub fn tpcds_schemas() -> Vec { Field::new("wp_max_ad_count", DataType::Int32, false), ]), ), - TableDef::new( + def( "web_site", Schema::new(vec![ Field::new("web_site_sk", DataType::Int32, false), @@ -577,3 +583,43 @@ pub fn tpcds_schemas() -> Vec { ), ] } + +static TPCDS_PRIMARY_KEYS: &[(&str, &[&str])] = &[ + ("call_center", &["cc_call_center_sk"]), + ("catalog_page", &["cp_catalog_page_sk"]), + ("catalog_returns", &["cr_item_sk", "cr_order_number"]), + ("catalog_sales", &["cs_item_sk", "cs_order_number"]), + ("customer", &["c_customer_sk"]), + ("customer_address", &["ca_address_sk"]), + ("customer_demographics", &["cd_demo_sk"]), + ("date_dim", &["d_date_sk"]), + ("household_demographics", &["hd_demo_sk"]), + ("income_band", &["ib_income_band_sk"]), + ( + "inventory", + &["inv_date_sk", "inv_item_sk", "inv_warehouse_sk"], + ), + ("item", &["i_item_sk"]), + ("promotion", &["p_promo_sk"]), + ("reason", &["r_reason_sk"]), + ("ship_mode", &["sm_ship_mode_sk"]), + ("store", &["s_store_sk"]), + ("store_returns", &["sr_item_sk", "sr_ticket_number"]), + ("store_sales", &["ss_item_sk", "ss_ticket_number"]), + ("time_dim", &["t_time_sk"]), + ("warehouse", &["w_warehouse_sk"]), + ("web_page", &["wp_web_page_sk"]), + ("web_returns", &["wr_item_sk", "wr_order_number"]), + ("web_sales", &["ws_item_sk", "ws_order_number"]), + ("web_site", &["web_site_sk"]), +]; + +fn tpcds_constraints(table: &str, schema: &Schema) -> Constraints { + let columns = TPCDS_PRIMARY_KEYS + .iter() + .find(|(name, _)| *name == table) + .map(|(_, columns)| *columns) + .unwrap_or_else(|| unimplemented!("unknown TPC-DS table: {table}")); + + Constraints::new_unverified(vec![primary_key(schema, columns)]) +} diff --git a/test-utils/src/tpch.rs b/test-utils/src/tpch.rs index 636221f71e519..3836a5ebab159 100644 --- a/test-utils/src/tpch.rs +++ b/test-utils/src/tpch.rs @@ -15,8 +15,9 @@ // specific language governing permissions and limitations // under the License. -use crate::TableDef; +use crate::{TableDef, primary_key}; use arrow::datatypes::{DataType, Field, Schema}; +use datafusion_common::Constraints; /// Schemas for the TPCH tables pub fn tpch_schemas() -> Vec { @@ -105,14 +106,41 @@ pub fn tpch_schemas() -> Vec { Field::new("r_comment", DataType::Utf8, false), ]); + let def = |name, schema: Schema| { + let constraints = tpch_constraints(name, &schema); + TableDef::new(name, schema).with_constraints(constraints) + }; + vec![ - TableDef::new("lineitem", lineitem_schema), - TableDef::new("orders", orders_schema), - TableDef::new("part", part_schema), - TableDef::new("supplier", supplier_schema), - TableDef::new("partsupp", partsupp_schema), - TableDef::new("customer", customer_schema), - TableDef::new("nation", nation_schema), - TableDef::new("region", region_schema), + def("lineitem", lineitem_schema), + def("orders", orders_schema), + def("part", part_schema), + def("supplier", supplier_schema), + def("partsupp", partsupp_schema), + def("customer", customer_schema), + def("nation", nation_schema), + def("region", region_schema), ] } + +/// Primary-key columns for each TPC-H table. +static TPCH_PRIMARY_KEYS: &[(&str, &[&str])] = &[ + ("region", &["r_regionkey"]), + ("nation", &["n_nationkey"]), + ("part", &["p_partkey"]), + ("supplier", &["s_suppkey"]), + ("partsupp", &["ps_partkey", "ps_suppkey"]), + ("customer", &["c_custkey"]), + ("orders", &["o_orderkey"]), + ("lineitem", &["l_orderkey", "l_linenumber"]), +]; + +fn tpch_constraints(table: &str, schema: &Schema) -> Constraints { + let columns = TPCH_PRIMARY_KEYS + .iter() + .find(|(name, _)| *name == table) + .map(|(_, columns)| *columns) + .unwrap_or_else(|| unimplemented!("unknown TPC-H table: {table}")); + + Constraints::new_unverified(vec![primary_key(schema, columns)]) +} diff --git a/uv.lock b/uv.lock index f86b732dfd6d5..85fcce1e9db48 100644 --- a/uv.lock +++ b/uv.lock @@ -240,61 +240,58 @@ wheels = [ [[package]] name = "cryptography" -version = "46.0.7" +version = "50.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/47/93/ac8f3d5ff04d54bc814e961a43ae5b0b146154c89c61b47bb07557679b18/cryptography-46.0.7.tar.gz", hash = "sha256:e4cfd68c5f3e0bfdad0d38e023239b96a2fe84146481852dffbcca442c245aa5", size = 750652, upload-time = "2026-04-08T01:57:54.692Z" } +sdist = { url = "https://files.pythonhosted.org/packages/de/41/6cbdcf9142d00fe82836fbb51e503e58088575cf7a0fe1dbff6695bf0840/cryptography-50.0.0.tar.gz", hash = "sha256:eeac2acb5a20ed25e0ad6d1df9891a520b78b404266b6d11778f25d5d691a6c9", size = 880201, upload-time = "2026-07-31T14:25:10.11Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/0b/5d/4a8f770695d73be252331e60e526291e3df0c9b27556a90a6b47bccca4c2/cryptography-46.0.7-cp311-abi3-macosx_10_9_universal2.whl", hash = "sha256:ea42cbe97209df307fdc3b155f1b6fa2577c0defa8f1f7d3be7d31d189108ad4", size = 7179869, upload-time = "2026-04-08T01:56:17.157Z" }, - { url = "https://files.pythonhosted.org/packages/5f/45/6d80dc379b0bbc1f9d1e429f42e4cb9e1d319c7a8201beffd967c516ea01/cryptography-46.0.7-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:b36a4695e29fe69215d75960b22577197aca3f7a25b9cf9d165dcfe9d80bc325", size = 4275492, upload-time = "2026-04-08T01:56:19.36Z" }, - { url = "https://files.pythonhosted.org/packages/4a/9a/1765afe9f572e239c3469f2cb429f3ba7b31878c893b246b4b2994ffe2fe/cryptography-46.0.7-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5ad9ef796328c5e3c4ceed237a183f5d41d21150f972455a9d926593a1dcb308", size = 4426670, upload-time = "2026-04-08T01:56:21.415Z" }, - { url = "https://files.pythonhosted.org/packages/8f/3e/af9246aaf23cd4ee060699adab1e47ced3f5f7e7a8ffdd339f817b446462/cryptography-46.0.7-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:73510b83623e080a2c35c62c15298096e2a5dc8d51c3b4e1740211839d0dea77", size = 4280275, upload-time = "2026-04-08T01:56:23.539Z" }, - { url = "https://files.pythonhosted.org/packages/0f/54/6bbbfc5efe86f9d71041827b793c24811a017c6ac0fd12883e4caa86b8ed/cryptography-46.0.7-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:cbd5fb06b62bd0721e1170273d3f4d5a277044c47ca27ee257025146c34cbdd1", size = 4928402, upload-time = "2026-04-08T01:56:25.624Z" }, - { url = "https://files.pythonhosted.org/packages/2d/cf/054b9d8220f81509939599c8bdbc0c408dbd2bdd41688616a20731371fe0/cryptography-46.0.7-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:420b1e4109cc95f0e5700eed79908cef9268265c773d3a66f7af1eef53d409ef", size = 4459985, upload-time = "2026-04-08T01:56:27.309Z" }, - { url = "https://files.pythonhosted.org/packages/f9/46/4e4e9c6040fb01c7467d47217d2f882daddeb8828f7df800cb806d8a2288/cryptography-46.0.7-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:24402210aa54baae71d99441d15bb5a1919c195398a87b563df84468160a65de", size = 3990652, upload-time = "2026-04-08T01:56:29.095Z" }, - { url = "https://files.pythonhosted.org/packages/36/5f/313586c3be5a2fbe87e4c9a254207b860155a8e1f3cca99f9910008e7d08/cryptography-46.0.7-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:8a469028a86f12eb7d2fe97162d0634026d92a21f3ae0ac87ed1c4a447886c83", size = 4279805, upload-time = "2026-04-08T01:56:30.928Z" }, - { url = "https://files.pythonhosted.org/packages/69/33/60dfc4595f334a2082749673386a4d05e4f0cf4df8248e63b2c3437585f2/cryptography-46.0.7-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:9694078c5d44c157ef3162e3bf3946510b857df5a3955458381d1c7cfc143ddb", size = 4892883, upload-time = "2026-04-08T01:56:32.614Z" }, - { url = "https://files.pythonhosted.org/packages/c7/0b/333ddab4270c4f5b972f980adef4faa66951a4aaf646ca067af597f15563/cryptography-46.0.7-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:42a1e5f98abb6391717978baf9f90dc28a743b7d9be7f0751a6f56a75d14065b", size = 4459756, upload-time = "2026-04-08T01:56:34.306Z" }, - { url = "https://files.pythonhosted.org/packages/d2/14/633913398b43b75f1234834170947957c6b623d1701ffc7a9600da907e89/cryptography-46.0.7-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:91bbcb08347344f810cbe49065914fe048949648f6bd5c2519f34619142bbe85", size = 4410244, upload-time = "2026-04-08T01:56:35.977Z" }, - { url = "https://files.pythonhosted.org/packages/10/f2/19ceb3b3dc14009373432af0c13f46aa08e3ce334ec6eff13492e1812ccd/cryptography-46.0.7-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:5d1c02a14ceb9148cc7816249f64f623fbfee39e8c03b3650d842ad3f34d637e", size = 4674868, upload-time = "2026-04-08T01:56:38.034Z" }, - { url = "https://files.pythonhosted.org/packages/1a/bb/a5c213c19ee94b15dfccc48f363738633a493812687f5567addbcbba9f6f/cryptography-46.0.7-cp311-abi3-win32.whl", hash = "sha256:d23c8ca48e44ee015cd0a54aeccdf9f09004eba9fc96f38c911011d9ff1bd457", size = 3026504, upload-time = "2026-04-08T01:56:39.666Z" }, - { url = "https://files.pythonhosted.org/packages/2b/02/7788f9fefa1d060ca68717c3901ae7fffa21ee087a90b7f23c7a603c32ae/cryptography-46.0.7-cp311-abi3-win_amd64.whl", hash = "sha256:397655da831414d165029da9bc483bed2fe0e75dde6a1523ec2fe63f3c46046b", size = 3488363, upload-time = "2026-04-08T01:56:41.893Z" }, - { url = "https://files.pythonhosted.org/packages/7b/56/15619b210e689c5403bb0540e4cb7dbf11a6bf42e483b7644e471a2812b3/cryptography-46.0.7-cp314-cp314t-macosx_10_9_universal2.whl", hash = "sha256:d151173275e1728cf7839aaa80c34fe550c04ddb27b34f48c232193df8db5842", size = 7119671, upload-time = "2026-04-08T01:56:44Z" }, - { url = "https://files.pythonhosted.org/packages/74/66/e3ce040721b0b5599e175ba91ab08884c75928fbeb74597dd10ef13505d2/cryptography-46.0.7-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:db0f493b9181c7820c8134437eb8b0b4792085d37dbb24da050476ccb664e59c", size = 4268551, upload-time = "2026-04-08T01:56:46.071Z" }, - { url = "https://files.pythonhosted.org/packages/03/11/5e395f961d6868269835dee1bafec6a1ac176505a167f68b7d8818431068/cryptography-46.0.7-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ebd6daf519b9f189f85c479427bbd6e9c9037862cf8fe89ee35503bd209ed902", size = 4408887, upload-time = "2026-04-08T01:56:47.718Z" }, - { url = "https://files.pythonhosted.org/packages/40/53/8ed1cf4c3b9c8e611e7122fb56f1c32d09e1fff0f1d77e78d9ff7c82653e/cryptography-46.0.7-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:b7b412817be92117ec5ed95f880defe9cf18a832e8cafacf0a22337dc1981b4d", size = 4271354, upload-time = "2026-04-08T01:56:49.312Z" }, - { url = "https://files.pythonhosted.org/packages/50/46/cf71e26025c2e767c5609162c866a78e8a2915bbcfa408b7ca495c6140c4/cryptography-46.0.7-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:fbfd0e5f273877695cb93baf14b185f4878128b250cc9f8e617ea0c025dfb022", size = 4905845, upload-time = "2026-04-08T01:56:50.916Z" }, - { url = "https://files.pythonhosted.org/packages/c0/ea/01276740375bac6249d0a971ebdf6b4dc9ead0ee0a34ef3b5a88c1a9b0d4/cryptography-46.0.7-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:ffca7aa1d00cf7d6469b988c581598f2259e46215e0140af408966a24cf086ce", size = 4444641, upload-time = "2026-04-08T01:56:52.882Z" }, - { url = "https://files.pythonhosted.org/packages/3d/4c/7d258f169ae71230f25d9f3d06caabcff8c3baf0978e2b7d65e0acac3827/cryptography-46.0.7-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:60627cf07e0d9274338521205899337c5d18249db56865f943cbe753aa96f40f", size = 3967749, upload-time = "2026-04-08T01:56:54.597Z" }, - { url = "https://files.pythonhosted.org/packages/b5/2a/2ea0767cad19e71b3530e4cad9605d0b5e338b6a1e72c37c9c1ceb86c333/cryptography-46.0.7-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:80406c3065e2c55d7f49a9550fe0c49b3f12e5bfff5dedb727e319e1afb9bf99", size = 4270942, upload-time = "2026-04-08T01:56:56.416Z" }, - { url = "https://files.pythonhosted.org/packages/41/3d/fe14df95a83319af25717677e956567a105bb6ab25641acaa093db79975d/cryptography-46.0.7-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:c5b1ccd1239f48b7151a65bc6dd54bcfcc15e028c8ac126d3fada09db0e07ef1", size = 4871079, upload-time = "2026-04-08T01:56:58.31Z" }, - { url = "https://files.pythonhosted.org/packages/9c/59/4a479e0f36f8f378d397f4eab4c850b4ffb79a2f0d58704b8fa0703ddc11/cryptography-46.0.7-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:d5f7520159cd9c2154eb61eb67548ca05c5774d39e9c2c4339fd793fe7d097b2", size = 4443999, upload-time = "2026-04-08T01:57:00.508Z" }, - { url = "https://files.pythonhosted.org/packages/28/17/b59a741645822ec6d04732b43c5d35e4ef58be7bfa84a81e5ae6f05a1d33/cryptography-46.0.7-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:fcd8eac50d9138c1d7fc53a653ba60a2bee81a505f9f8850b6b2888555a45d0e", size = 4399191, upload-time = "2026-04-08T01:57:02.654Z" }, - { url = "https://files.pythonhosted.org/packages/59/6a/bb2e166d6d0e0955f1e9ff70f10ec4b2824c9cfcdb4da772c7dd69cc7d80/cryptography-46.0.7-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:65814c60f8cc400c63131584e3e1fad01235edba2614b61fbfbfa954082db0ee", size = 4655782, upload-time = "2026-04-08T01:57:04.592Z" }, - { url = "https://files.pythonhosted.org/packages/95/b6/3da51d48415bcb63b00dc17c2eff3a651b7c4fed484308d0f19b30e8cb2c/cryptography-46.0.7-cp314-cp314t-win32.whl", hash = "sha256:fdd1736fed309b4300346f88f74cd120c27c56852c3838cab416e7a166f67298", size = 3002227, upload-time = "2026-04-08T01:57:06.91Z" }, - { url = "https://files.pythonhosted.org/packages/32/a8/9f0e4ed57ec9cebe506e58db11ae472972ecb0c659e4d52bbaee80ca340a/cryptography-46.0.7-cp314-cp314t-win_amd64.whl", hash = "sha256:e06acf3c99be55aa3b516397fe42f5855597f430add9c17fa46bf2e0fb34c9bb", size = 3475332, upload-time = "2026-04-08T01:57:08.807Z" }, - { url = "https://files.pythonhosted.org/packages/a7/7f/cd42fc3614386bc0c12f0cb3c4ae1fc2bbca5c9662dfed031514911d513d/cryptography-46.0.7-cp38-abi3-macosx_10_9_universal2.whl", hash = "sha256:462ad5cb1c148a22b2e3bcc5ad52504dff325d17daf5df8d88c17dda1f75f2a4", size = 7165618, upload-time = "2026-04-08T01:57:10.645Z" }, - { url = "https://files.pythonhosted.org/packages/a5/d0/36a49f0262d2319139d2829f773f1b97ef8aef7f97e6e5bd21455e5a8fb5/cryptography-46.0.7-cp38-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:84d4cced91f0f159a7ddacad249cc077e63195c36aac40b4150e7a57e84fffe7", size = 4270628, upload-time = "2026-04-08T01:57:12.885Z" }, - { url = "https://files.pythonhosted.org/packages/8a/6c/1a42450f464dda6ffbe578a911f773e54dd48c10f9895a23a7e88b3e7db5/cryptography-46.0.7-cp38-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:128c5edfe5e5938b86b03941e94fac9ee793a94452ad1365c9fc3f4f62216832", size = 4415405, upload-time = "2026-04-08T01:57:14.923Z" }, - { url = "https://files.pythonhosted.org/packages/9a/92/4ed714dbe93a066dc1f4b4581a464d2d7dbec9046f7c8b7016f5286329e2/cryptography-46.0.7-cp38-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:5e51be372b26ef4ba3de3c167cd3d1022934bc838ae9eaad7e644986d2a3d163", size = 4272715, upload-time = "2026-04-08T01:57:16.638Z" }, - { url = "https://files.pythonhosted.org/packages/b7/e6/a26b84096eddd51494bba19111f8fffe976f6a09f132706f8f1bf03f51f7/cryptography-46.0.7-cp38-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:cdf1a610ef82abb396451862739e3fc93b071c844399e15b90726ef7470eeaf2", size = 4918400, upload-time = "2026-04-08T01:57:19.021Z" }, - { url = "https://files.pythonhosted.org/packages/c7/08/ffd537b605568a148543ac3c2b239708ae0bd635064bab41359252ef88ed/cryptography-46.0.7-cp38-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:1d25aee46d0c6f1a501adcddb2d2fee4b979381346a78558ed13e50aa8a59067", size = 4450634, upload-time = "2026-04-08T01:57:21.185Z" }, - { url = "https://files.pythonhosted.org/packages/16/01/0cd51dd86ab5b9befe0d031e276510491976c3a80e9f6e31810cce46c4ad/cryptography-46.0.7-cp38-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:cdfbe22376065ffcf8be74dc9a909f032df19bc58a699456a21712d6e5eabfd0", size = 3985233, upload-time = "2026-04-08T01:57:22.862Z" }, - { url = "https://files.pythonhosted.org/packages/92/49/819d6ed3a7d9349c2939f81b500a738cb733ab62fbecdbc1e38e83d45e12/cryptography-46.0.7-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:abad9dac36cbf55de6eb49badd4016806b3165d396f64925bf2999bcb67837ba", size = 4271955, upload-time = "2026-04-08T01:57:24.814Z" }, - { url = "https://files.pythonhosted.org/packages/80/07/ad9b3c56ebb95ed2473d46df0847357e01583f4c52a85754d1a55e29e4d0/cryptography-46.0.7-cp38-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:935ce7e3cfdb53e3536119a542b839bb94ec1ad081013e9ab9b7cfd478b05006", size = 4879888, upload-time = "2026-04-08T01:57:26.88Z" }, - { url = "https://files.pythonhosted.org/packages/b8/c7/201d3d58f30c4c2bdbe9b03844c291feb77c20511cc3586daf7edc12a47b/cryptography-46.0.7-cp38-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:35719dc79d4730d30f1c2b6474bd6acda36ae2dfae1e3c16f2051f215df33ce0", size = 4449961, upload-time = "2026-04-08T01:57:29.068Z" }, - { url = "https://files.pythonhosted.org/packages/a5/ef/649750cbf96f3033c3c976e112265c33906f8e462291a33d77f90356548c/cryptography-46.0.7-cp38-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:7bbc6ccf49d05ac8f7d7b5e2e2c33830d4fe2061def88210a126d130d7f71a85", size = 4401696, upload-time = "2026-04-08T01:57:31.029Z" }, - { url = "https://files.pythonhosted.org/packages/41/52/a8908dcb1a389a459a29008c29966c1d552588d4ae6d43f3a1a4512e0ebe/cryptography-46.0.7-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:a1529d614f44b863a7b480c6d000fe93b59acee9c82ffa027cfadc77521a9f5e", size = 4664256, upload-time = "2026-04-08T01:57:33.144Z" }, - { url = "https://files.pythonhosted.org/packages/4b/fa/f0ab06238e899cc3fb332623f337a7364f36f4bb3f2534c2bb95a35b132c/cryptography-46.0.7-cp38-abi3-win32.whl", hash = "sha256:f247c8c1a1fb45e12586afbb436ef21ff1e80670b2861a90353d9b025583d246", size = 3013001, upload-time = "2026-04-08T01:57:34.933Z" }, - { url = "https://files.pythonhosted.org/packages/d2/f1/00ce3bde3ca542d1acd8f8cfa38e446840945aa6363f9b74746394b14127/cryptography-46.0.7-cp38-abi3-win_amd64.whl", hash = "sha256:506c4ff91eff4f82bdac7633318a526b1d1309fc07ca76a3ad182cb5b686d6d3", size = 3472985, upload-time = "2026-04-08T01:57:36.714Z" }, - { url = "https://files.pythonhosted.org/packages/63/0c/dca8abb64e7ca4f6b2978769f6fea5ad06686a190cec381f0a796fdcaaba/cryptography-46.0.7-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:fc9ab8856ae6cf7c9358430e49b368f3108f050031442eaeb6b9d87e4dcf4e4f", size = 3476879, upload-time = "2026-04-08T01:57:38.664Z" }, - { url = "https://files.pythonhosted.org/packages/3a/ea/075aac6a84b7c271578d81a2f9968acb6e273002408729f2ddff517fed4a/cryptography-46.0.7-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:d3b99c535a9de0adced13d159c5a9cf65c325601aa30f4be08afd680643e9c15", size = 4219700, upload-time = "2026-04-08T01:57:40.625Z" }, - { url = "https://files.pythonhosted.org/packages/6c/7b/1c55db7242b5e5612b29fc7a630e91ee7a6e3c8e7bf5406d22e206875fbd/cryptography-46.0.7-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:d02c738dacda7dc2a74d1b2b3177042009d5cab7c7079db74afc19e56ca1b455", size = 4385982, upload-time = "2026-04-08T01:57:42.725Z" }, - { url = "https://files.pythonhosted.org/packages/cb/da/9870eec4b69c63ef5925bf7d8342b7e13bc2ee3d47791461c4e49ca212f4/cryptography-46.0.7-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:04959522f938493042d595a736e7dbdff6eb6cc2339c11465b3ff89343b65f65", size = 4219115, upload-time = "2026-04-08T01:57:44.939Z" }, - { url = "https://files.pythonhosted.org/packages/f4/72/05aa5832b82dd341969e9a734d1812a6aadb088d9eb6f0430fc337cc5a8f/cryptography-46.0.7-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:3986ac1dee6def53797289999eabe84798ad7817f3e97779b5061a95b0ee4968", size = 4385479, upload-time = "2026-04-08T01:57:46.86Z" }, - { url = "https://files.pythonhosted.org/packages/20/2a/1b016902351a523aa2bd446b50a5bc1175d7a7d1cf90fe2ef904f9b84ebc/cryptography-46.0.7-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:258514877e15963bd43b558917bc9f54cf7cf866c38aa576ebf47a77ddbc43a4", size = 3412829, upload-time = "2026-04-08T01:57:48.874Z" }, + { url = "https://files.pythonhosted.org/packages/c5/5c/59086b4aac5e879d38ddbcf74e4be7ade89cebc3eb199a55da998c3bb46a/cryptography-50.0.0-cp311-abi3-macosx_11_0_arm64.whl", hash = "sha256:031e2d5dd4bb9caa3ca9c82e5a197fd8ae680232cee62603d1a813f3f07e3d03", size = 4001252, upload-time = "2026-07-31T14:23:33.331Z" }, + { url = "https://files.pythonhosted.org/packages/57/ef/8f2df13c7216bcad3e1c74e07f6e193d93e998e114f524a53877c9af27ad/cryptography-50.0.0-cp311-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:fd9192b7b70c573d7f214eb1ae35e00d359f6f5e4b27c7e21e30de1fc6204645", size = 4719554, upload-time = "2026-07-31T14:23:35.611Z" }, + { url = "https://files.pythonhosted.org/packages/d9/41/029086c34d91052fc3b88bcc8056f709a7c915c7a23b235a54eb800b1c97/cryptography-50.0.0-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:06a32a980526a6ab9a4b9bf8f7385800791e2bb960903cb6b530e4817509a3b7", size = 4702130, upload-time = "2026-07-31T14:23:37.635Z" }, + { url = "https://files.pythonhosted.org/packages/7d/ff/b6ce0954962e7f7b969f850a883744197bb3910bdfd7b6da162eab7d9f68/cryptography-50.0.0-cp311-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:a1b30560f2acc95aa8b2e06e716a13dbfc97314747b80d9707e307f77b40d6b3", size = 4725244, upload-time = "2026-07-31T14:23:39.471Z" }, + { url = "https://files.pythonhosted.org/packages/06/1e/63a1027cb7fec360a182208e1b7767d5aa1fe57be3d6aa856e69a321edc0/cryptography-50.0.0-cp311-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:8d89f3976b10b4ce31118de72329025f70d2c6ead14a8217c5514dd2c6d5a78f", size = 5342265, upload-time = "2026-07-31T14:23:41.286Z" }, + { url = "https://files.pythonhosted.org/packages/6b/72/a1116d683a6d7ece94590013882515de087edf9ef0e6292aae615a44df73/cryptography-50.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b42a28c1844fd9de8f3f7d540e36b66f3a9c83fceac7170ebc7a6a19edd9dcae", size = 4734609, upload-time = "2026-07-31T14:23:43.139Z" }, + { url = "https://files.pythonhosted.org/packages/15/37/36a9c479bbe49acea2636c7fd3360d20f7b7e079c300352011c44850b181/cryptography-50.0.0-cp311-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:900131fafd8aead39ac7dd3a7e833be754c17a95cfd91221636949fe4eb0aa8a", size = 4356517, upload-time = "2026-07-31T14:23:44.939Z" }, + { url = "https://files.pythonhosted.org/packages/32/98/8a151d64367204cbc63ec65d37502f1d9c53cf4bfc6ec3c532614dbec60d/cryptography-50.0.0-cp311-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:07949c449a1abcf60d1ee6e88956d89404c7df3c8258f46589e912988e551987", size = 4724529, upload-time = "2026-07-31T14:23:46.93Z" }, + { url = "https://files.pythonhosted.org/packages/22/f6/ec13b470172126464a86bf54d2294a46d29837fc51ba3e45d4047946fb5e/cryptography-50.0.0-cp311-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:f89831ef99dd7dd169ab06d63a831adb9e20a87aac6d380266bbda5823349169", size = 5299852, upload-time = "2026-07-31T14:23:48.851Z" }, + { url = "https://files.pythonhosted.org/packages/da/3a/f05e32c99d440c9bb891ea0e36c9091891e36be5a9a87ab2ee6ea20729f6/cryptography-50.0.0-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:82148ec5bddac30b51a5b3c1945075f896fa022cb93f8e4a01e9f6ee95292c5f", size = 4734462, upload-time = "2026-07-31T14:23:50.861Z" }, + { url = "https://files.pythonhosted.org/packages/ca/dc/bd72b26be8953f80625f63151efd38eee71c76ca6cf591c08ff34615a79e/cryptography-50.0.0-cp311-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:1489e263a8048bb8b6a8bac662eb2d402ea5d2b7b4699b72f385f1e2772db105", size = 4852708, upload-time = "2026-07-31T14:23:52.715Z" }, + { url = "https://files.pythonhosted.org/packages/27/20/c930314a2ab476d15dec966ec87e2e9637bb02b06106b12c0396c57bb603/cryptography-50.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7cec5b856506da6defb290f30c9ee687d5f5e8cb0bd3f6459dde43b0b4fa40ef", size = 5004179, upload-time = "2026-07-31T14:23:54.887Z" }, + { url = "https://files.pythonhosted.org/packages/32/2e/c9db68a0c4bfa28e310707527c0ee3a2bd254104d2e02e68f368e197aa4c/cryptography-50.0.0-cp311-abi3-win_amd64.whl", hash = "sha256:bd1c592e4d5974f0d08d4888e432157adba757c66da0246918e43677fafa2d30", size = 3840395, upload-time = "2026-07-31T14:23:56.677Z" }, + { url = "https://files.pythonhosted.org/packages/c3/fb/951032a3bf22a5697c83183fb6294a4843772947a70e616c57b3ff5f522e/cryptography-50.0.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:49e7d93abdbd2990caced757e5fade25302f719c3c8fb6e6fff2dde98999fc41", size = 3989258, upload-time = "2026-07-31T14:23:58.881Z" }, + { url = "https://files.pythonhosted.org/packages/d4/67/91eb047e69c5e845f2f14b8a2e4a1aab0f283cb885531e9e22c8adb176bc/cryptography-50.0.0-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:19736989797678c6af1e55cd49055cdbcb55d8f6b5583ac5335f933aba9101dc", size = 4700648, upload-time = "2026-07-31T14:24:00.702Z" }, + { url = "https://files.pythonhosted.org/packages/30/82/85f0f7425c856b9f96459411eb12e74ef72df9caf6f8f15bf23a33ff131f/cryptography-50.0.0-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:80b63928fa35083b33966ce1efb70e5b9607181e49dcd1c22c8c005e319f667f", size = 4682442, upload-time = "2026-07-31T14:24:02.538Z" }, + { url = "https://files.pythonhosted.org/packages/1a/28/b555a365adff1cca2fbe7b9e487d68a40de6bc67ff2cb587473eb43de0e7/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:d58c3db7cd6eed54e6c06744db55456b65ebd7492ddeae9c1e93cfca7aa857d3", size = 4707596, upload-time = "2026-07-31T14:24:04.394Z" }, + { url = "https://files.pythonhosted.org/packages/72/d8/f52538140cc719df62a01cf87d1c7142318d235817109d6f4054d7c352d6/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_ppc64le.whl", hash = "sha256:df2a58a472f332225671c35b0a830208b86d004f82baa8530fa3782c85646533", size = 5314552, upload-time = "2026-07-31T14:24:06.31Z" }, + { url = "https://files.pythonhosted.org/packages/38/14/6120e5bd7c5aa022ad15424ba4d5c5269d0d9448ed4d55e492ea91e3c1c4/cryptography-50.0.0-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:11b74db56cdbe3cdee6e3f6982ecb70334fa10dce99ed58bf7894aaaa3b2a037", size = 4717113, upload-time = "2026-07-31T14:24:08.349Z" }, + { url = "https://files.pythonhosted.org/packages/fa/71/190bf38c3ee2e0f8efc9860ae100c9df4169742eef274b91e7aa1cb133b9/cryptography-50.0.0-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:f59e38625469987d7ef6d495323c55e7db6c212eaf6112267e0d3b565a2e9c9f", size = 4338580, upload-time = "2026-07-31T14:24:10.227Z" }, + { url = "https://files.pythonhosted.org/packages/3a/63/504ccfbbe61fd8aa983f7f146399cdf034c72c2fc55f5b2dfdcdcdb20c99/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_aarch64.whl", hash = "sha256:ecfed7367f965a0328cfbdd70da860f15441f002f613185668c6e6ebf5a0ac11", size = 4707038, upload-time = "2026-07-31T14:24:12.169Z" }, + { url = "https://files.pythonhosted.org/packages/01/77/2cf79bbfc4d12ca106437a6e170d6aaa01a373e93093118aaaef0e801bd4/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_ppc64le.whl", hash = "sha256:9aa87839c383bdbab6ef865787a1fb877af8dd03464c4400322726feaaadfc6d", size = 5273110, upload-time = "2026-07-31T14:24:14.38Z" }, + { url = "https://files.pythonhosted.org/packages/e5/45/8aae2972c520145377ea3559a605a899bebe227bf070b33cdb445929a9b9/cryptography-50.0.0-cp314-cp314t-manylinux_2_34_x86_64.whl", hash = "sha256:6ba6a53445bd3cfa809ef3ef5f1589aa6ba08784a1d962bf47d0940e871dab1c", size = 4716439, upload-time = "2026-07-31T14:24:16.415Z" }, + { url = "https://files.pythonhosted.org/packages/7b/20/4fe50b619a48c2525cc46e2dbc1ac490708d704be5d467bdaac6dc955682/cryptography-50.0.0-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:3f5735ffe4996d28b809371756219f5354864902a3b9e7c0b9ee87041209fc9c", size = 4837383, upload-time = "2026-07-31T14:24:18.553Z" }, + { url = "https://files.pythonhosted.org/packages/92/91/3a31366e183343d3703f8995c095f5734676bd6938118047e50fcf279eb4/cryptography-50.0.0-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:1b4a266766514614f8aa60416e71f2fc6e575d36e7bdc90f644fadb2f4b75b95", size = 4985772, upload-time = "2026-07-31T14:24:20.385Z" }, + { url = "https://files.pythonhosted.org/packages/74/9a/02ffe35b2853d121689871eb5dce862092562b3a1ed5cc98f1aaed441506/cryptography-50.0.0-cp314-cp314t-win_amd64.whl", hash = "sha256:12b9c6996425c76ea6c457ace4f3073e715b8c545add07cd1a8f3a4f90691269", size = 3816291, upload-time = "2026-07-31T14:24:22.125Z" }, + { url = "https://files.pythonhosted.org/packages/03/37/73d005be173aff344af30e9fd2a576575cb2391a7101d9cd3842e1fa8cce/cryptography-50.0.0-cp39-abi3-macosx_11_0_arm64.whl", hash = "sha256:ccdc4a71a4dabae05de219404f9f4abc38e3b58422177ff93d0da05967dafa07", size = 4036009, upload-time = "2026-07-31T14:24:24.122Z" }, + { url = "https://files.pythonhosted.org/packages/ff/c6/7a6202a534e32103a285b7834a120869557fe198d51d7cfe59754c8bda9c/cryptography-50.0.0-cp39-abi3-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:910e1d2668e7de9648f2bcee30e180db2a6b15c30f887d7c4c93ddf96e3992e3", size = 4745252, upload-time = "2026-07-31T14:24:26.118Z" }, + { url = "https://files.pythonhosted.org/packages/85/4f/0fa8c2f4428198f15d9ff8d63400e27afbf94ce833f6108da1eb3753f945/cryptography-50.0.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a91296cb61e8df6f86d0c19cc4068228da256bf59bf86049fbd821084565327f", size = 4728939, upload-time = "2026-07-31T14:24:27.994Z" }, + { url = "https://files.pythonhosted.org/packages/d1/63/54dd723490ba2dc09b299682c10b38db38f159728bcaae8c591b8af2f22d/cryptography-50.0.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:e722f16708d854fe924790e051061f6704a472c3bac347b6fd88033ea8dd0dc5", size = 4748483, upload-time = "2026-07-31T14:24:30.254Z" }, + { url = "https://files.pythonhosted.org/packages/1d/dd/7c77d26285cc7f6991efce64a0f5b4f9383bfa5dd8c5033003eaf7db4cdb/cryptography-50.0.0-cp39-abi3-manylinux_2_28_ppc64le.whl", hash = "sha256:d764dcf130c428ef66786f866dd750f53182bc608813489915e9fc106bb0c82f", size = 5367599, upload-time = "2026-07-31T14:24:32.457Z" }, + { url = "https://files.pythonhosted.org/packages/46/c9/f60aed34c013f317f92817b6c171c2d22a78270fa41109bd4b08af26b194/cryptography-50.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:105110f43a471dbd0060b9c9516cb8a6a79233631a04cc2ba16f28323ac6e025", size = 4762647, upload-time = "2026-07-31T14:24:34.599Z" }, + { url = "https://files.pythonhosted.org/packages/be/f3/f9a0173b139372c3a48ed98154b45cc6b9de17c789d5ab552e621c293609/cryptography-50.0.0-cp39-abi3-manylinux_2_31_armv7l.whl", hash = "sha256:828743d939e9629bc267b8e2d08d8bb67cd4319c771a33d4b18b22dd8fb7440a", size = 4385197, upload-time = "2026-07-31T14:24:36.647Z" }, + { url = "https://files.pythonhosted.org/packages/d8/36/83bb81f6e569bc38e1e4a7bc80f29b46bb9601920bc455fc8e888f5d5742/cryptography-50.0.0-cp39-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:2a8183b489dc1f7f80f135780fadc1108f14b31b8a40411c7a5b17425f65f28b", size = 4748095, upload-time = "2026-07-31T14:24:39.493Z" }, + { url = "https://files.pythonhosted.org/packages/6b/16/d3008eff98c764979865834c3d386d4fd041b5f52e7f34fc29ac1a5eb515/cryptography-50.0.0-cp39-abi3-manylinux_2_34_ppc64le.whl", hash = "sha256:6e7d61120573a7f2cd94cc095f9e81f6967c61ccdf194285aa143ecec8e0b708", size = 5325948, upload-time = "2026-07-31T14:24:41.556Z" }, + { url = "https://files.pythonhosted.org/packages/9c/f8/d97f9603efda3888187bfdb893f26c41be4735c10631d05d284ee6b047c4/cryptography-50.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:37fdb0d0111f1e2ff07139dfb79f1b49531f8e213c46f1163dd7642979b58c47", size = 4762400, upload-time = "2026-07-31T14:24:43.636Z" }, + { url = "https://files.pythonhosted.org/packages/64/a2/4615c8f7d81a00b1d6e6afe19f694e1543582349fb5f4076f6cb5dc36485/cryptography-50.0.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:c87f62a3d3b9888ed0fdde100ec06aa61ca9cd44bad9057d1dff9a516b5f5bb9", size = 4878208, upload-time = "2026-07-31T14:24:45.522Z" }, + { url = "https://files.pythonhosted.org/packages/d2/1a/efcfb02f91407149a0dacffffab791f7e19bf6385f63b3666dc8b5e5c9c8/cryptography-50.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:65c2c3add92b45fd0709db8594536aea39c2a67af0e27ffcf049c498501140b7", size = 5037050, upload-time = "2026-07-31T14:24:47.697Z" }, + { url = "https://files.pythonhosted.org/packages/57/30/4a22984d4f1bdfb8c054f07a92bc176b97a3134cc1d6c4b3bffb1f3688b4/cryptography-50.0.0-cp39-abi3-win_amd64.whl", hash = "sha256:d24fead1d4d076e1bfb006dcec392074a3cd8d7b4fc8a595aa64073b2b7a96ba", size = 3874135, upload-time = "2026-07-31T14:24:50.085Z" }, + { url = "https://files.pythonhosted.org/packages/9d/3e/e54cde8c01631a5a8226ccd617eab9e57fd5cfdad90f1a9e6bb570794631/cryptography-50.0.0-pp311-pypy311_pp73-macosx_11_0_arm64.whl", hash = "sha256:5e34edd123674534acd70147f0ca331eaa2c74e6325fb2028c886aa26ba0b68c", size = 3963170, upload-time = "2026-07-31T14:24:51.968Z" }, + { url = "https://files.pythonhosted.org/packages/01/b6/0b9e125e90f3d2dcf599a218a899cda7326a3158cfa258723f0b398b08f6/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_28_aarch64.whl", hash = "sha256:8eb5e1172eb569ea8a872796576e6a67c276351728b6455d5beb01242b027c6a", size = 4692441, upload-time = "2026-07-31T14:24:53.743Z" }, + { url = "https://files.pythonhosted.org/packages/53/c9/a5151588710785a96d7bc4de27d4cd62f263bbbcb203cfe29df537eb6505/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_28_x86_64.whl", hash = "sha256:910d11e1a385c654bf738bf3e6b8e6ed5de0f5610fcae2be9e5b398d8081d20e", size = 4699810, upload-time = "2026-07-31T14:24:55.746Z" }, + { url = "https://files.pythonhosted.org/packages/c7/1a/15b92b25eb6ce3089cd49377ae990a0f3ad485a510f968aed1f19dbdcdf2/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_34_aarch64.whl", hash = "sha256:62598a8a57f815db4c6259a4e97d857dab56697e7de8e8ab02352ab74da1995d", size = 4691924, upload-time = "2026-07-31T14:24:58.082Z" }, + { url = "https://files.pythonhosted.org/packages/62/15/219075012ab13e8905f3cd572204f4acb4b111df787104346b9bc0cea789/cryptography-50.0.0-pp311-pypy311_pp73-manylinux_2_34_x86_64.whl", hash = "sha256:07479a1cb08219ab719147e742e76090c9c773321959bb94946fffdd397a6437", size = 4699593, upload-time = "2026-07-31T14:24:59.951Z" }, + { url = "https://files.pythonhosted.org/packages/8e/b5/c2c5fce26f0ee40d21bafe7f191d29a34b35a65ac4fe8a1191d1983612e9/cryptography-50.0.0-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:c99c003e088647b8a5b7c145d6f78c335f6348332b62e142d411c4b63d1460b9", size = 3813796, upload-time = "2026-07-31T14:25:02.298Z" }, ] [[package]] @@ -349,10 +346,10 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "jinja2", specifier = ">=3.1.6,<4" }, - { name = "maturin", specifier = ">=1.13.3,<2" }, - { name = "myst-parser", specifier = ">=5,<6" }, - { name = "pydata-sphinx-theme", specifier = ">=0.17.1,<1" }, - { name = "setuptools", specifier = ">=82.0.1,<83" }, + { name = "maturin", specifier = ">=1.14.1,<2" }, + { name = "myst-parser", specifier = ">=5.1.0,<6" }, + { name = "pydata-sphinx-theme", specifier = ">=0.20.0,<1" }, + { name = "setuptools", specifier = ">=83.0.0,<84" }, { name = "sphinx", specifier = ">=9,<10" }, { name = "sphinx-reredirects", specifier = ">=1.1,<2" }, ] @@ -465,14 +462,14 @@ wheels = [ [[package]] name = "markdown-it-py" -version = "4.0.0" +version = "4.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "mdurl" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } +sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/94/54/e7d793b573f298e1c9013b8c4dade17d481164aa517d1d7148619c2cedbf/markdown_it_py-4.0.0-py3-none-any.whl", hash = "sha256:87327c59b172c5011896038353a81343b6754500a08cd7a4973bb48c6d578147", size = 87321, upload-time = "2025-08-11T12:57:51.923Z" }, + { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, ] [[package]] @@ -551,35 +548,35 @@ wheels = [ [[package]] name = "maturin" -version = "1.13.3" +version = "1.14.1" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/9c/1c/612d23d33ec21b9ae7ece7b3f0dd5f9dfd57b4009e9d2938165869ebd6ae/maturin-1.13.3.tar.gz", hash = "sha256:771e1e9e71a278e56db01552e0d1acfd1464259f9575b6e72842f893cd299079", size = 357934, upload-time = "2026-05-11T07:43:39.027Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/b3/addd877f871fb1860d46d3a4f206ecb10b946c85846805e6367631926fd3/maturin-1.14.1.tar.gz", hash = "sha256:9d6577a62cd08e0ceba7a0db06fb098e0c9b1b3429bad747a4f3a18215a1b3df", size = 369637, upload-time = "2026-06-19T05:19:49.774Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/71/66/18c2aaac0b2a5dea9f1db5984ce83b905ad205cfc7c02d0091e707c0c2e7/maturin-1.13.3-py3-none-linux_armv6l.whl", hash = "sha256:3cc13929ca82aefa4adbf0f2c35419369796213c6fb0eb24e914945f50ef5d8c", size = 10190971, upload-time = "2026-05-11T07:43:10.431Z" }, - { url = "https://files.pythonhosted.org/packages/bc/71/26a988d092e4fd6a9523d46d44400a46cad7cdf3fd206ce702240c748aee/maturin-1.13.3-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:53b08bd075649ce96513ad9abf241a43cb685ed6e9e7790f8dbc2d66e95d8323", size = 19716714, upload-time = "2026-05-11T07:43:36.911Z" }, - { url = "https://files.pythonhosted.org/packages/82/5c/f3fd0e184255d9fc7e272c62af3dfa84c617b2577ef83af9ce615f5279cc/maturin-1.13.3-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:4cd478e6e4c56251e48ed079b8efd55b30bc5c09cf695a1bdafaeb582ee735a0", size = 10194726, upload-time = "2026-05-11T07:43:07.05Z" }, - { url = "https://files.pythonhosted.org/packages/a9/e1/f4edb69fb647b77c4769a9bfd4d6fb62961e653d164bc277ecdffac3ab61/maturin-1.13.3-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:a2675e25f313034ae6f57388cf14818f87d8961c4a96795287f3e155f59beb11", size = 10172781, upload-time = "2026-05-11T07:43:40.796Z" }, - { url = "https://files.pythonhosted.org/packages/c7/7d/a1be934690cdcc3c6609769ceaad322ab7501c2ee5bafcac1b14d609e403/maturin-1.13.3-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:4667ef609ab446c1b5e0bfe4f9fb99699ab6d8548433f8d1a684256e0b67217f", size = 10682670, upload-time = "2026-05-11T07:43:13.132Z" }, - { url = "https://files.pythonhosted.org/packages/18/f5/372ae19b72ce8f6e37e5864ae4dc5b252ee9fce0619ccc3aa366aa3a7f97/maturin-1.13.3-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:3db93337ed97e60ffc878aa8b493cd7ae44d3a5e1a37256db3a4491f57565018", size = 10060363, upload-time = "2026-05-11T07:43:21.107Z" }, - { url = "https://files.pythonhosted.org/packages/cb/5b/c68340cca09368af0df80965dfabed4234205a492a93da00793c7b9aae20/maturin-1.13.3-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:1cc0a110b224ca90406b668a3e3c1f5a515062e59e26292f6dbaf5fd4909c6f3", size = 10017551, upload-time = "2026-05-11T07:43:33.916Z" }, - { url = "https://files.pythonhosted.org/packages/28/1e/f90fb2b000bad9e6d850cd5afb88b2f1e2a279cfb4de02ea40078484690e/maturin-1.13.3-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:c00ea6428dea17bf616fe93770837634454b28c2de1a876e42ef8036c616079a", size = 13301712, upload-time = "2026-05-11T07:43:26.492Z" }, - { url = "https://files.pythonhosted.org/packages/be/58/1670f68a8f04ccd7b90df11047bd9a046585310e84e1967cc9849cd1c5a3/maturin-1.13.3-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:49fd6ab08da28098ccf37afca24cdba72376ba9c1eedf9dd25ff82ed771961ff", size = 10946765, upload-time = "2026-05-11T07:43:16.135Z" }, - { url = "https://files.pythonhosted.org/packages/4b/ac/00c955c2ef134817b1a7bdaa76b0309e9c5291eb17d9ff88069eecd08bc2/maturin-1.13.3-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:b6741d7bf4af97da937528fd1e523c6ab54f53d9a21870fa735d6e67fd88e273", size = 10388661, upload-time = "2026-05-11T07:43:18.727Z" }, - { url = "https://files.pythonhosted.org/packages/97/c6/cbf8a51dde19c19aeba0d9b075095a2effb9b31fd312b1aae3ac79f8aea2/maturin-1.13.3-py3-none-win32.whl", hash = "sha256:0ef257e692cc756c87af5bea95ddfe7d3ac49d3376a7a87f728d63f06e7b6f8b", size = 8901838, upload-time = "2026-05-11T07:43:23.76Z" }, - { url = "https://files.pythonhosted.org/packages/a1/ff/c6a50a59dc8313097d43ac5f4d74df6a500c8cb62b0dc9e054f53e203a48/maturin-1.13.3-py3-none-win_amd64.whl", hash = "sha256:def4a435ea9d2ee93b18ba579dc8c9cf898889a66f312cd379b5e374ec3e3ad6", size = 10340801, upload-time = "2026-05-11T07:43:29.239Z" }, - { url = "https://files.pythonhosted.org/packages/6c/93/e32e79333f0902ba292b996f504f5f06be59587f7d02ab8d5ed1e3066445/maturin-1.13.3-py3-none-win_arm64.whl", hash = "sha256:2389fe92d017cea9d94e521fa0175314a4c52f79a1057b901fbc9f8686ef7d0b", size = 9706562, upload-time = "2026-05-11T07:43:31.743Z" }, + { url = "https://files.pythonhosted.org/packages/f4/f0/97c5a5bd9c71653a066c0976a484eaaae50b9369557838a4176b7b0bdaa5/maturin-1.14.1-py3-none-linux_armv6l.whl", hash = "sha256:522292398945442cdafa9daeb2271b2340fbde57027b818f923f88eab04174f8", size = 10207496, upload-time = "2026-06-19T05:19:09.321Z" }, + { url = "https://files.pythonhosted.org/packages/fe/83/294bca639b0e052f1e2f65199b3db258780c7d4e31408b934c9c974a1379/maturin-1.14.1-py3-none-macosx_10_12_x86_64.macosx_11_0_arm64.macosx_10_12_universal2.whl", hash = "sha256:ffe5ad71f21d1e6603c4dd75f7fee34adf5ed5ebcebb692886549888ebb329ed", size = 19680113, upload-time = "2026-06-19T05:19:13.43Z" }, + { url = "https://files.pythonhosted.org/packages/43/b6/79c881410a3b1c187f7eb3d407aecae646c6a4433d630d72200359015e83/maturin-1.14.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:f3306078070c1508fd715b9116070cbcaff5959024272a9f1e6f5cb29768b86c", size = 10169205, upload-time = "2026-06-19T05:19:16.615Z" }, + { url = "https://files.pythonhosted.org/packages/93/9d/44b6f26dcb7f7a04c5501ac2dbb6ca1490150682baa525ca5860504f9eab/maturin-1.14.1-py3-none-manylinux_2_12_i686.manylinux2010_i686.musllinux_1_1_i686.whl", hash = "sha256:cd457cd88961156e26379e1155bd287cc0ec1c8b2f1582b0660fb31b87c8842d", size = 10188098, upload-time = "2026-06-19T05:19:19.736Z" }, + { url = "https://files.pythonhosted.org/packages/1a/bd/9c0d5d6983905ce2c9edaa073a7e89355a9cf7f396988e05d32f1c37785d/maturin-1.14.1-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.musllinux_1_1_x86_64.whl", hash = "sha256:dfc54ae32e6fcb18302193ab9a30b0b25eefffba994ae13238974805533ef75e", size = 10627576, upload-time = "2026-06-19T05:19:22.713Z" }, + { url = "https://files.pythonhosted.org/packages/e5/33/b096412bd6a7cb399652b260666f901adf88a687181a6dbd6a3f89f0a94e/maturin-1.14.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.musllinux_1_1_aarch64.whl", hash = "sha256:a131d912b5267e640bc96d70f4914e10590aed64082ec9abacba7cea52004224", size = 10085181, upload-time = "2026-06-19T05:19:25.69Z" }, + { url = "https://files.pythonhosted.org/packages/56/8d/08c3bf469c38a23c9e6c877e338193001eb604d010fedc08341974e38528/maturin-1.14.1-py3-none-manylinux_2_17_armv7l.manylinux2014_armv7l.musllinux_1_1_armv7l.whl", hash = "sha256:be18fc568fb76884c0205456336892a75105ec398e6b667cd777c6268bd06d69", size = 10026363, upload-time = "2026-06-19T05:19:28.904Z" }, + { url = "https://files.pythonhosted.org/packages/3a/a4/c4d1a92839f8745ab4aab988a7db884a79d6d710bd3b286fcf9316dece1a/maturin-1.14.1-py3-none-manylinux_2_17_ppc64le.manylinux2014_ppc64le.musllinux_1_1_ppc64le.whl", hash = "sha256:994a0c8ba3ad8a92b3a9ee1b02645d200d610216b15cff5102b0fe65e8e08666", size = 13321347, upload-time = "2026-06-19T05:19:32.411Z" }, + { url = "https://files.pythonhosted.org/packages/b3/fa/170f04624d03fd07d2a8b1b67de83a127af93aef9eaa425839553347297b/maturin-1.14.1-py3-none-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:be80866363e605d137991b491a741a84cde9ae350183c4c85f49690ca9aaaa65", size = 10877609, upload-time = "2026-06-19T05:19:35.448Z" }, + { url = "https://files.pythonhosted.org/packages/61/ad/1ae2e1d0ded282bf2c55ac13f0811d87deb425e200ae64a15785675dede9/maturin-1.14.1-py3-none-manylinux_2_31_riscv64.musllinux_1_1_riscv64.whl", hash = "sha256:5282dffd4b539d2be245f4e5b1a5ab6bc1033b58f4a4872f5833f9d43c954aa4", size = 10417316, upload-time = "2026-06-19T05:19:38.28Z" }, + { url = "https://files.pythonhosted.org/packages/fb/27/bf677183920718da49cd7982d6a3ffc440aad8919329f571d189f81b7bdf/maturin-1.14.1-py3-none-win32.whl", hash = "sha256:1a04de0a20188f95c721b5702eed18140bdcccb28c386797093eca3f62f4d4e0", size = 8931293, upload-time = "2026-06-19T05:19:41.183Z" }, + { url = "https://files.pythonhosted.org/packages/63/4b/585adeb9167b08d3cdff0032a938b0e72655c92003df4f52c3f696a1bcc2/maturin-1.14.1-py3-none-win_amd64.whl", hash = "sha256:3c9f94640ecc4895e94abaf834a0684430032c865b2748a36c12461fd9252fdd", size = 10314067, upload-time = "2026-06-19T05:19:44.389Z" }, + { url = "https://files.pythonhosted.org/packages/51/d4/dac8c0720ae246be1700afb6fbdbbea20fe35b13f6570b2f70faa005df77/maturin-1.14.1-py3-none-win_arm64.whl", hash = "sha256:15cea8fcb3ba47dd636f50092bb34baea8b04ac777392f23e6bf8a9a61efb894", size = 9718943, upload-time = "2026-06-19T05:19:47.49Z" }, ] [[package]] name = "mdit-py-plugins" -version = "0.5.0" +version = "0.6.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "markdown-it-py" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/b2/fd/a756d36c0bfba5f6e39a1cdbdbfdd448dc02692467d83816dff4592a1ebc/mdit_py_plugins-0.5.0.tar.gz", hash = "sha256:f4918cb50119f50446560513a8e311d574ff6aaed72606ddae6d35716fe809c6", size = 44655, upload-time = "2025-08-11T07:25:49.083Z" } +sdist = { url = "https://files.pythonhosted.org/packages/59/fc/f8d0863f8862f25602c0404d75568e89fb6b4109804645e5cdfb1be5cf56/mdit_py_plugins-0.6.1.tar.gz", hash = "sha256:a2bca0f039f39dbd35fb74ae1b5f998608c437463371f0ff7f49a19a17a114d0", size = 56114, upload-time = "2026-05-13T09:03:38.91Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/fb/86/dd6e5db36df29e76c7a7699123569a4a18c1623ce68d826ed96c62643cae/mdit_py_plugins-0.5.0-py3-none-any.whl", hash = "sha256:07a08422fc1936a5d26d146759e9155ea466e842f5ab2f7d2266dd084c8dab1f", size = 57205, upload-time = "2025-08-11T07:25:47.597Z" }, + { url = "https://files.pythonhosted.org/packages/a5/69/6da5581c6a7fede7dc261bf4e67d6adca4196f176b43288b55b3db395b6e/mdit_py_plugins-0.6.1-py3-none-any.whl", hash = "sha256:214c82fb2ac524472ab6a5bcab1de80f73b50443e187f401bfd77efbc7c6481d", size = 66663, upload-time = "2026-05-13T09:03:37.76Z" }, ] [[package]] @@ -593,7 +590,7 @@ wheels = [ [[package]] name = "myst-parser" -version = "5.0.0" +version = "5.1.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "docutils" }, @@ -604,9 +601,9 @@ dependencies = [ { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/33/fa/7b45eef11b7971f0beb29d27b7bfe0d747d063aa29e170d9edd004733c8a/myst_parser-5.0.0.tar.gz", hash = "sha256:f6f231452c56e8baa662cc352c548158f6a16fcbd6e3800fc594978002b94f3a", size = 98535, upload-time = "2026-01-15T09:08:18.036Z" } +sdist = { url = "https://files.pythonhosted.org/packages/21/dc/603751677fff302f34396e206b610f556a59d7fe58b9a2145f54e96b48e8/myst_parser-5.1.0.tar.gz", hash = "sha256:ab69322dc6719dcc7f296479dbb70181b66df6ed315064f92dbc85c0e1bf2f02", size = 101182, upload-time = "2026-05-13T09:38:19.361Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d3/ac/686789b9145413f1a61878c407210e41bfdb097976864e0913078b24098c/myst_parser-5.0.0-py3-none-any.whl", hash = "sha256:ab31e516024918296e169139072b81592336f2fef55b8986aa31c9f04b5f7211", size = 84533, upload-time = "2026-01-15T09:08:16.788Z" }, + { url = "https://files.pythonhosted.org/packages/09/dc/f3dfb7488b770f3f67e6545085bf2abea5172e88f57b8ad25ef860ca704c/myst_parser-5.1.0-py3-none-any.whl", hash = "sha256:9c91c52b3cdb4d94a6506e4fab4e2f296c7623a0da0dcbe6de1565c3dad67a8a", size = 85817, upload-time = "2026-05-13T09:38:17.904Z" }, ] [[package]] @@ -758,21 +755,22 @@ wheels = [ [[package]] name = "pydata-sphinx-theme" -version = "0.17.1" +version = "0.20.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "accessible-pygments" }, { name = "babel" }, { name = "beautifulsoup4" }, { name = "docutils" }, + { name = "jinja2" }, { name = "pygments" }, + { name = "requests" }, { name = "sphinx", version = "9.0.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.12'" }, { name = "sphinx", version = "9.1.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.12'" }, - { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/ce/f7/c74c7100a7f4c0f77b5dcacb7dfdb8fee774fb70e487dd97acba2b930774/pydata_sphinx_theme-0.17.1.tar.gz", hash = "sha256:2cfc1d926c753c77039b7ee53f0ccebcbee5e81f0db61432b01cbb10ad7fd0af", size = 4991415, upload-time = "2026-04-21T13:00:34.263Z" } +sdist = { url = "https://files.pythonhosted.org/packages/ae/8e/add936feaaa9dade7d5b87c6852566d85e518b38ad32224dea32ef958984/pydata_sphinx_theme-0.20.0.tar.gz", hash = "sha256:0da172d41e19a66de875f4002f7054b385372ec65763852193791e658d50bb4a", size = 5004756, upload-time = "2026-07-09T09:09:14.693Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/e2/bc/2cb8c78300ce1ace4eeac3b3522218cea2c2053bfa6b4e32cc972a477f9a/pydata_sphinx_theme-0.17.1-py3-none-any.whl", hash = "sha256:320b022d7808bdf5920d9a28e573f27aace9b23e1af6ca103eecc752411df492", size = 6823346, upload-time = "2026-04-21T13:00:31.978Z" }, + { url = "https://files.pythonhosted.org/packages/80/08/28e2194ed1c3c3a3e86e0600e2dcc7f21dcd670bd31f3efab2e3e2cf0cbd/pydata_sphinx_theme-0.20.0-py3-none-any.whl", hash = "sha256:56744483c9d72c783e075de716ab95d486108b69605df7528078090b73f11f69", size = 6201166, upload-time = "2026-07-09T09:09:12.899Z" }, ] [[package]] @@ -802,11 +800,11 @@ wheels = [ [[package]] name = "pyjwt" -version = "2.12.0" +version = "2.13.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/a8/10/e8192be5f38f3e8e7e046716de4cae33d56fd5ae08927a823bb916be36c1/pyjwt-2.12.0.tar.gz", hash = "sha256:2f62390b667cd8257de560b850bb5a883102a388829274147f1d724453f8fb02", size = 102511, upload-time = "2026-03-12T17:15:30.831Z" } +sdist = { url = "https://files.pythonhosted.org/packages/3b/81/58d0ac84e1ef3a3843791d6954d94c0b33d526c75eeb1efbce9d0a4c4077/pyjwt-2.13.0.tar.gz", hash = "sha256:41571c89ca91598c79e8ef18a2d07367d4810fbbd6f637794879baf1b7703423", size = 107515, upload-time = "2026-05-21T19:54:36.618Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/15/70/70f895f404d363d291dcf62c12c85fdd47619ad9674ac0f53364d035925a/pyjwt-2.12.0-py3-none-any.whl", hash = "sha256:9bb459d1bdd0387967d287f5656bf7ec2b9a26645d1961628cda1764e087fd6e", size = 29700, upload-time = "2026-03-12T17:15:29.257Z" }, + { url = "https://files.pythonhosted.org/packages/a3/5e/ecf12fdb62546d64385c158514e9b2b671f7832108ef2ecd2020ce0af2d1/pyjwt-2.13.0-py3-none-any.whl", hash = "sha256:66adcc2aff09b3f1bbd95fc1e1577df8ac8723c978552fd43304c8a290ac5728", size = 31274, upload-time = "2026-05-21T19:54:35.362Z" }, ] [package.optional-dependencies] @@ -943,11 +941,11 @@ wheels = [ [[package]] name = "setuptools" -version = "82.0.1" +version = "83.0.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/4f/db/cfac1baf10650ab4d1c111714410d2fbb77ac5a616db26775db562c8fab2/setuptools-82.0.1.tar.gz", hash = "sha256:7d872682c5d01cfde07da7bccc7b65469d3dca203318515ada1de5eda35efbf9", size = 1152316, upload-time = "2026-03-09T12:47:17.221Z" } +sdist = { url = "https://files.pythonhosted.org/packages/34/26/f5d29e25ffdb535afef2d35cdb55b325298f96debd670da4c325e08d70f4/setuptools-83.0.0.tar.gz", hash = "sha256:025bccbbf0fa05b6192bc64ae1e7b16e001fd6d6d4d5de03c97b1c1ade523bef", size = 1154254, upload-time = "2026-07-04T15:31:22.699Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/9d/76/f789f7a86709c6b087c5a2f52f911838cad707cc613162401badc665acfe/setuptools-82.0.1-py3-none-any.whl", hash = "sha256:a59e362652f08dcd477c78bb6e7bd9d80a7995bc73ce773050228a348ce2e5bb", size = 1006223, upload-time = "2026-03-09T12:47:15.026Z" }, + { url = "https://files.pythonhosted.org/packages/5d/40/e1e72872c6354b306daef1703549e8e83b4d43cfea356311bf722a043752/setuptools-83.0.0-py3-none-any.whl", hash = "sha256:29b23c360f22f414dc7336bb39178cc7bcbf6021ed2733cde173f09dba19abb3", size = 1008090, upload-time = "2026-07-04T15:31:20.885Z" }, ] [[package]] @@ -970,11 +968,11 @@ wheels = [ [[package]] name = "soupsieve" -version = "2.8.3" +version = "2.8.4" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/7b/ae/2d9c981590ed9999a0d91755b47fc74f74de286b0f5cee14c9269041e6c4/soupsieve-2.8.3.tar.gz", hash = "sha256:3267f1eeea4251fb42728b6dfb746edc9acaffc4a45b27e19450b676586e8349", size = 118627, upload-time = "2026-01-20T04:27:02.457Z" } +sdist = { url = "https://files.pythonhosted.org/packages/47/2c/0a5f6f8ee0d5589e48c7640213ed5175d52cf540a06725b628cc1a45d6ce/soupsieve-2.8.4.tar.gz", hash = "sha256:e121fd02e975c695e4e9e8774a5ee35d74714b59307868dcc5319ad2d9e3328e", size = 121110, upload-time = "2026-05-24T13:55:57.154Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/46/2c/1462b1d0a634697ae9e55b3cecdcb64788e8b7d63f54d923fcd0bb140aed/soupsieve-2.8.3-py3-none-any.whl", hash = "sha256:ed64f2ba4eebeab06cc4962affce381647455978ffc1e36bb79a545b91f45a95", size = 37016, upload-time = "2026-01-20T04:27:01.012Z" }, + { url = "https://files.pythonhosted.org/packages/5e/f5/0c41cb68dcae6b7de4fac4188a3a9589e21fb31df21ea3a2e888db95e6c9/soupsieve-2.8.4-py3-none-any.whl", hash = "sha256:e7e6b0769c8f51ed59acab6e994b00621096cfb1c640a7509295987388fbaf65", size = 37304, upload-time = "2026-05-24T13:55:55.406Z" }, ] [[package]] @@ -985,23 +983,23 @@ resolution-markers = [ "python_full_version < '3.12'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version < '3.12'" }, - { name = "babel", marker = "python_full_version < '3.12'" }, - { name = "colorama", marker = "python_full_version < '3.12' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version < '3.12'" }, - { name = "imagesize", marker = "python_full_version < '3.12'" }, - { name = "jinja2", marker = "python_full_version < '3.12'" }, - { name = "packaging", marker = "python_full_version < '3.12'" }, - { name = "pygments", marker = "python_full_version < '3.12'" }, - { name = "requests", marker = "python_full_version < '3.12'" }, - { name = "roman-numerals", marker = "python_full_version < '3.12'" }, - { name = "snowballstemmer", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version < '3.12'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version < '3.12'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/42/50/a8c6ccc36d5eacdfd7913ddccd15a9cee03ecafc5ee2bc40e1f168d85022/sphinx-9.0.4.tar.gz", hash = "sha256:594ef59d042972abbc581d8baa577404abe4e6c3b04ef61bd7fc2acbd51f3fa3", size = 8710502, upload-time = "2025-12-04T07:45:27.343Z" } wheels = [ @@ -1016,23 +1014,23 @@ resolution-markers = [ "python_full_version >= '3.12'", ] dependencies = [ - { name = "alabaster", marker = "python_full_version >= '3.12'" }, - { name = "babel", marker = "python_full_version >= '3.12'" }, - { name = "colorama", marker = "python_full_version >= '3.12' and sys_platform == 'win32'" }, - { name = "docutils", marker = "python_full_version >= '3.12'" }, - { name = "imagesize", marker = "python_full_version >= '3.12'" }, - { name = "jinja2", marker = "python_full_version >= '3.12'" }, - { name = "packaging", marker = "python_full_version >= '3.12'" }, - { name = "pygments", marker = "python_full_version >= '3.12'" }, - { name = "requests", marker = "python_full_version >= '3.12'" }, - { name = "roman-numerals", marker = "python_full_version >= '3.12'" }, - { name = "snowballstemmer", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-applehelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-devhelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-htmlhelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-jsmath", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-qthelp", marker = "python_full_version >= '3.12'" }, - { name = "sphinxcontrib-serializinghtml", marker = "python_full_version >= '3.12'" }, + { name = "alabaster" }, + { name = "babel" }, + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "docutils" }, + { name = "imagesize" }, + { name = "jinja2" }, + { name = "packaging" }, + { name = "pygments" }, + { name = "requests" }, + { name = "roman-numerals" }, + { name = "snowballstemmer" }, + { name = "sphinxcontrib-applehelp" }, + { name = "sphinxcontrib-devhelp" }, + { name = "sphinxcontrib-htmlhelp" }, + { name = "sphinxcontrib-jsmath" }, + { name = "sphinxcontrib-qthelp" }, + { name = "sphinxcontrib-serializinghtml" }, ] sdist = { url = "https://files.pythonhosted.org/packages/cd/bd/f08eb0f4eed5c83f1ba2a3bd18f7745a2b1525fad70660a1c00224ec468a/sphinx-9.1.0.tar.gz", hash = "sha256:7741722357dd75f8190766926071fed3bdc211c74dd2d7d4df5404da95930ddb", size = 8718324, upload-time = "2025-12-31T15:09:27.646Z" } wheels = [
[ WITH cte [, ...] ]
+SELECT select_item [, ...]
+[ INTO table_name ]
+[ FROM from_item [, ...] ]
+[ JOIN join_item ... ]
+[ WHERE condition ]
+[ GROUP BY grouping_element [, ...] | GROUP BY ALL ]
+[ HAVING condition ]
+[ WINDOW window_name AS (window_definition) [, ...] ]
+[ QUALIFY condition ]
+[ { UNION | INTERSECT | EXCEPT } query ] [...]
+[ ORDER BY order_expression [, ...] ]
+[ LIMIT count ] [ OFFSET count ]
+[ |> pipe_operator ... ]